From 19d86fcf984faa4008c932dc32b56ad1fa19e72e Mon Sep 17 00:00:00 2001 From: Nicholas Klos Date: Wed, 21 Oct 2026 12:31:19 -0500 Subject: [PATCH] Add querypilot trace and hooks: audit, guard and diagnose agent SQL tool calls trace import reads Claude Code transcripts into the audit log, pairing each query tool call with its result (SQL, outcome, failure class, latency; never result rows). hooks install logs calls live, and with --guard a PreToolUse hook validates SQL against schemas cached from the agent's own schema calls, denies unknown columns with did-you-mean feedback, catches epoch-ms arithmetic on DATE columns, and enforces a row cap. trace doctor reports failure classes, retry chains, latency, a guard replay, and tool/prompt/data fixes, and exports recovered chains for eval replay. Palantir Foundry's SQL tool is mapped by default; any tool maps by name and SQL field. Includes the first redacted report over the Foundry history. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 14 + README.md | 12 + docs/trace-doctor-foundry-2026-10-21.md | 173 ++++ docs/trace.md | 99 ++ src/querypilot/audit/types.py | 5 + src/querypilot/cli.py | 13 + src/querypilot/trace/__init__.py | 1 + src/querypilot/trace/classify.py | 173 ++++ src/querypilot/trace/cli.py | 288 ++++++ src/querypilot/trace/doctor.py | 939 ++++++++++++++++++ src/querypilot/trace/guard.py | 288 ++++++ src/querypilot/trace/hooks.py | 417 ++++++++ src/querypilot/trace/schema_cache.py | 232 +++++ src/querypilot/trace/tools.py | 79 ++ src/querypilot/trace/transcript.py | 341 +++++++ .../claude_code_foundry_transcript.jsonl | 22 + tests/test_trace.py | 499 ++++++++++ 17 files changed, 3595 insertions(+) create mode 100644 docs/trace-doctor-foundry-2026-10-21.md create mode 100644 docs/trace.md create mode 100644 src/querypilot/trace/__init__.py create mode 100644 src/querypilot/trace/classify.py create mode 100644 src/querypilot/trace/cli.py create mode 100644 src/querypilot/trace/doctor.py create mode 100644 src/querypilot/trace/guard.py create mode 100644 src/querypilot/trace/hooks.py create mode 100644 src/querypilot/trace/schema_cache.py create mode 100644 src/querypilot/trace/tools.py create mode 100644 src/querypilot/trace/transcript.py create mode 100644 tests/fixtures/claude_code_foundry_transcript.jsonl create mode 100644 tests/test_trace.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 73d7ca2..b0683cd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,20 @@ follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ### Added +- **`querypilot trace` and `querypilot hooks`.** Observe and improve the SQL agents + run through MCP and connector tools that QueryPilot cannot proxy, starting with + Palantir Foundry's `run_sql_query_on_foundry_dataset` in Claude Code. Any other + query tool is mapped by name and SQL input field (`--tool NAME=FIELD[:DIALECT]`). + `trace import` reads Claude Code transcripts into the JSONL audit log (SQL, + outcome, failure class, latency, session; never result rows). `hooks install` + adds PostToolUse/PostToolUseFailure hooks that log calls live. With `--guard`, + a PreToolUse hook validates SQL against a schema cache built from the agent's + schema calls, denies unknown columns with "did you mean" feedback, catches + epoch-millisecond arithmetic on DATE columns, and enforces `--max-rows`. + `trace doctor` reports failure classes, retry chains, latency and stalls, and a + guard replay. It tags each finding as a tool, prompt or data fix, and exports + recovered retry chains for `eval replay`. `QueryAuditRecord` gains `tool_name`, + `tool_call_id` and `details`. See `docs/trace.md`. - **`querypilot proxy` (prototype).** A Postgres wire-protocol proxy that puts the validator and access policy in front of a real database, so any client (psql, drivers, ORMs, MCP database servers, BI tools) is protected unchanged. It diff --git a/README.md b/README.md index da26f4a..7dd67d6 100644 --- a/README.md +++ b/README.md @@ -135,6 +135,18 @@ Every statement then gets the following: Driver handshakes, `SET`/`SHOW` of formatting settings, and catalog queries for `psql \d` and ORM reflection pass through an allowlist. Catalog queries that could read data, such as `pg_stats` or `query_to_xml(...)`, are rejected. The policy file is a normal `querypilot.yaml`, and its `database_url` is used when you omit `--upstream`. Off loopback, set `QUERYPILOT_PROXY_PASSWORD`. Protocol coverage, known gaps (no TLS termination yet, no server-side cursors), and the reasons it fails closed are in [docs/proxy.md](docs/proxy.md). +## Watch the SQL your agent already runs: `querypilot trace` and `hooks` + +Some query tools can't sit behind QueryPilot, for example Palantir Foundry's MCP SQL tool. For those, QueryPilot reads what the agent did instead. `trace import` pairs every query tool call in Claude Code transcripts with its result. `hooks install` logs new calls live. `trace doctor` reports the failure rate by class, retry chains, latency and stalls, and says for each problem whether it is best fixed in the tool, the prompt, or the data: + +```bash +querypilot trace doctor # reads ~/.claude/projects +querypilot trace doctor --output report.md # redacted, shareable +querypilot hooks install --guard --max-rows 200 +``` + +With `--guard`, a PreToolUse hook validates each query against schemas cached from the agent's own schema calls. It denies unknown columns with a "did you mean" reply before the query reaches the warehouse, blocks date arithmetic on DATE columns stored as epoch milliseconds, and enforces a row cap. Foundry's tool is mapped by default; map any other with `--tool NAME=SQL_FIELD[:DIALECT]`. Recovered retry chains export to `eval replay` as a regression suite. See [docs/trace.md](docs/trace.md) and the [first Foundry report](docs/trace-doctor-foundry-2026-10-21.md). + ## Better SQL from the model When QueryPilot writes the SQL itself (`ask`, `ask_database`, evals), it uses the techniques with the largest measured gains in the text-to-SQL literature. The sources and numbers are in [docs/research/developer-needs.md](docs/research/developer-needs.md). diff --git a/docs/trace-doctor-foundry-2026-10-21.md b/docs/trace-doctor-foundry-2026-10-21.md new file mode 100644 index 0000000..4d52554 --- /dev/null +++ b/docs/trace-doctor-foundry-2026-10-21.md @@ -0,0 +1,173 @@ +# Trace doctor: Foundry SQL in Claude Code (2026-07-17 to 2026-10-21) + +Tools: `mcp__palantir-mcp__run_sql_query_on_foundry_dataset`. Period: 2026-07-17 to 2026-10-21. + +- **5,163** query calls in **71** sessions (224 transcripts, including subagents) +- **609** failed: **11.8%** of the 5,158 that ran; 5 denied by the user or a permission rule +- Latency p50 **3.8s**, p90 **11.6s**, p99 80.8s, max 901s; **84** calls over 60s (6 over 5 min) + +## Failures by class + +| Class | Calls | % of calls | % of failures | Sessions | Fix first in | +|---|---:|---:|---:|---:|---| +| Unresolved column | 191 | 3.7% | 31% | 38 | tool | +| Result too large for context | 165 | 3.2% | 27% | 9 | tool | +| Upstream / network (no SQL error returned) | 107 | 2.1% | 18% | 22 | tool | +| Unknown table / dataset | 57 | 1.1% | 9% | 14 | tool | +| Data type mismatch | 47 | 0.9% | 8% | 19 | data | +| Syntax / parse | 20 | 0.4% | 3% | 8 | prompt | +| Other SQL analysis error | 13 | 0.3% | 2% | 9 | prompt | +| Invalid tool call | 8 | 0.2% | 1% | 7 | prompt | +| Timeout | 1 | 0.0% | 0% | 1 | tool | + +## Retry chains + +A chain starts at a failed query and follows the next queries on the same datasets in the same transcript until one succeeds (recovered) or the agent moves on (abandoned). + +- 382 chains, 609 failed attempts; 260 recovered, 122 abandoned +- Failed attempts per chain: 1: 300, 2: 45, 3: 13, 4: 8, 5+: 16 (max 15) +- Time spent on failed attempts: 134 min; median time from first failure to recovery 17s +- Agent fetched a schema mid-chain in 40 chains +- First failure class: Unresolved column 152, Upstream / network (no SQL error returned) 74, Result too large for context 42, Data type mismatch 41, Unknown table / dataset 40, Syntax / parse 15, Other SQL analysis error 10, Invalid tool call 7, Timeout 1 +- 91 successful queries were exact repeats of an earlier success in the session + +## Latency + +| Outcome | p50 (s) | +|---|---:| +| Succeeded | 3.8 | +| Unresolved column | 3.2 | +| Result too large for context | 22.1 | +| Upstream / network (no SQL error returned) | 2.5 | +| Unknown table / dataset | 2.1 | +| Data type mismatch | 3.7 | +| Syntax / parse | 0.6 | +| Other SQL analysis error | 3.6 | +| Invalid tool call | 1.3 | +| Timeout | 901.3 | + +## Guard replay + +The schema cache built from this history holds 924 datasets. The guard could check columns for 3,487 queries (every dataset cached at the time). + +- Failures it would have stopped before execution: Unresolved column 98, Data type mismatch 24, Upstream / network (no SQL error returned) 13, Syntax / parse 6, Result too large for context 4. Outside the column and type classes, these queries failed for another reason first but also used a column the dataset lacks or did date arithmetic on a DATE column. +- Successful queries it would have blocked by mistake: 38 of 3,099 checked. These are mostly datasets rebuilt with new columns after their schema was cached; live, the agent resends the query once and the hook drops the stale entry. + +Sample guard replies (redacted): + +- Column 'product' does not exist; did you mean Product_Name? +- Column 'category_name' does not exist; did you mean Customer_Name? +- Column 'rank' does not exist. + +## Findings + +### 1. [tool fix] Validate columns before the query runs (QueryPilot guard) (191 calls) + +191 queries referenced a column the dataset does not have. Replaying the history, the guard (schema cache built from earlier schema calls and SELECT * results) would have stopped 98 of them before they reached the warehouse, with a 'did you mean' reply, and wrongly blocked 38 of 3099 successful queries it could check. The rest touched datasets with no cached schema. + +- Column 'product' does not exist; did you mean Product_Name? +- Column 'category_name' does not exist; did you mean Customer_Name? +- Column 'rank' does not exist. + +**Fix:** querypilot hooks install --guard (optionally --require-schema, which makes the agent fetch a dataset's schema before its first query). + +### 2. [tool fix] Cap result size before it floods the context window (165 calls) + +165 results exceeded Claude Code's output limit. 141 had neither a LIMIT nor a rowLimit at or below the cap, and 122 used SELECT *. rowLimit on those calls: <= 50: 13, <= 100: 5, <= 200: 6, <= 500: 3, <= 1000: 137, not set: 1. Wide text columns overflow even at 100 rows. + +**Fix:** querypilot hooks install --guard --max-rows 200 enforces a cap; the prompt fix is to select only the needed columns and aggregate in SQL instead of reading raw rows. + +### 3. [tool fix] Transport and server errors reach the agent as opaque failures (107 calls) + +107 calls failed without a SQL error: network fetch failures, Foundry internal errors, and HTTP errors whose body the MCP tool dropped (a stack trace instead of a message). + +**Fix:** Fix in the tool layer: retry with backoff on network and 5xx errors and surface Foundry's error body. The agent cannot fix these by rewriting SQL. + +### 4. [prompt fix] Tell the agent the real names of columns it keeps guessing (92 calls) + +The same wrong column name failed repeatedly, often across sessions, so the agent is not carrying the schema forward. Recurring guesses: + +- `termination_reason` x20 in 1 session (dataset_A, dataset_B); closest real columns: terminated_date, job_change_reason, pension_id +- `product` x9 in 3 sessions (dataset_E, dataset_F); closest real columns: product_id, Product_Name, periodicity +- `run_seconds` x5 in 1 session (dataset_G); closest real columns: raw_json, branch_name, build_status +- `evidence_grade` x4 in 1 session (dataset_H); closest real columns: tenure_years, hr_role_group, department +- `posting_key` x4 in 1 session (dataset_I); closest real columns: cell_key, state, title +- `record_id` x4 in 3 sessions (dataset_F, dataset_J); closest real columns: ERP_ID, Company_ID, Deal_ID +- `pct_days_used` x3 in 1 session (dataset_M, dataset_N); closest real columns: Repeat_Users, Churned_Users, Play_Count +- `name` x3 in 3 sessions (dataset_E, dataset_J); closest real columns: NAME, status_name, ACTIVE + +**Fix:** Add the hints below to the agent's instructions (CLAUDE.md, or `hints:` in querypilot.yaml); when a guess recurs across many sessions, the data fix is a column description or an alias column in a view. + +### 5. [tool fix] Identical queries re-run in the same session (91 calls) + +91 successful queries repeated a query that had already succeeded earlier in the same session. + +**Fix:** Cache results per session in the tool, or tell the agent to keep earlier results instead of re-querying. + +### 6. [tool fix] Long-running queries stall the agent (84 calls) + +84 calls took longer than 60s (6 over 5 minutes); 14 of those failed anyway. 1 calls failed with a timeout. + +**Fix:** Cap rows and require a filter on the partition/date column for large datasets; for queries the agent repeats, materialize the result as a dataset or view (data fix). + +### 7. [tool fix] Catch wrong or inaccessible dataset references (57 calls) + +57 queries named a dataset that was not found or not readable (Foundry answers 403 for both a wrong RID and missing access). The guard caught 0 (malformed RIDs). + +**Fix:** Keep a registry of the datasets the agent uses (RID, path, one-line purpose) in its instructions, and enable --require-schema so a typo fails at the cheap schema call. + +### 8. [prompt fix] Fetch the schema before the first query, not after it fails (40 calls) + +In 40 retry chains the agent fetched a dataset schema only after a query failed. + +**Fix:** Instruct: 'Before querying a Foundry dataset for the first time in a session, call get_foundry_dataset_schema and use its exact column names.' The guard's --require-schema enforces this. + +### 9. [prompt fix] Spark SQL dialect mistakes (33 calls) + +33 queries failed to parse or failed Spark analysis (SCALAR_SUBQUERY_IS_IN_GROUP_BY_OR_AGGREGATE_FUNCTION x2, UNSUPPORTED_GENERATOR x2, INVALID_USAGE_OF_STAR_OR_REGEX x2, UNRESOLVED_ROUTINE x1, DIVIDE_BY_ZERO x1). + +**Fix:** State the dialect in the agent's instructions: 'Foundry SQL is Spark SQL: quote dataset RIDs in backticks, every non-aggregated column must be in GROUP BY, use try_cast/try_divide.' + +### 10. [data fix] Date columns read as epoch milliseconds (30 calls) + +47 queries failed on a data type mismatch; 30 of them treated a DATE or TIMESTAMP column as epoch milliseconds (`col / 1000`, `col >= 1767225600000`). 327 successful results displayed a DATE/TIMESTAMP column as 13-digit numbers, which is where the agent learns the wrong type. + +- `Report_Date` (DATE, dataset_F) x3 +- `hired_date` (DATE, dataset_A) x3 +- `Date` (DATE, dataset_T) x3 +- `Close_Date` (DATE, dataset_K) x2 +- `Activity_Date` (DATE, dataset_U) x2 +- `start_date` (DATE, dataset_V) x2 +- `Report_Date` (DATE, dataset_W) x1 +- `start_date` (DATE, dataset_X) x1 + +**Fix:** Data: add column descriptions that state the type (suggested `descriptions:` below), or publish views with explicit DATE columns. Tool: the PostToolUse hook annotates results whose date columns are rendered as epoch ms, and the guard blocks `date_col / 1000`. + +## Suggested prompt hints + +```yaml +hints: + - "In dataset_A, dataset_B, there is no `termination_reason` column; use `terminated_date`." + - "In dataset_E, dataset_F, there is no `product` column; use `product_id`." + - "In dataset_E, dataset_J, there is no `name` column; use `NAME`." + - "Foundry SQL results show DATE/TIMESTAMP values as epoch-millisecond numbers, but the columns are real dates: never divide them by 1000 or compare them with 13-digit numbers; use DATE 'YYYY-MM-DD' literals and date functions." + - "Foundry SQL is Spark SQL: quote dataset RIDs and paths in backticks, list every non-aggregated column in GROUP BY, and prefer try_cast / try_divide." +``` + +## Suggested column descriptions + +```yaml +descriptions: + "dataset_A.hired_date": "DATE. Not epoch milliseconds, although query results display it as a 13-digit number; filter with DATE literals (e.g. hired_date >= DATE '2026-01-01') or date functions." + "dataset_AA.Created_At": "DATE. Not epoch milliseconds, although query results display it as a 13-digit number; filter with DATE literals (e.g. Created_At >= DATE '2026-01-01') or date functions." + "dataset_AB.date": "DATE. Not epoch milliseconds, although query results display it as a 13-digit number; filter with DATE literals (e.g. date >= DATE '2026-01-01') or date functions." + "dataset_AC.start_date": "DATE. Not epoch milliseconds, although query results display it as a 13-digit number; filter with DATE literals (e.g. start_date >= DATE '2026-01-01') or date functions." + "dataset_F.Report_Date": "DATE. Not epoch milliseconds, although query results display it as a 13-digit number; filter with DATE literals (e.g. Report_Date >= DATE '2026-01-01') or date functions." + "dataset_K.Close_Date": "DATE. Not epoch milliseconds, although query results display it as a 13-digit number; filter with DATE literals (e.g. Close_Date >= DATE '2026-01-01') or date functions." + "dataset_T.Date": "DATE. Not epoch milliseconds, although query results display it as a 13-digit number; filter with DATE literals (e.g. Date >= DATE '2026-01-01') or date functions." + "dataset_U.Activity_Date": "DATE. Not epoch milliseconds, although query results display it as a 13-digit number; filter with DATE literals (e.g. Activity_Date >= DATE '2026-01-01') or date functions." + "dataset_W.Report_Date": "DATE. Not epoch milliseconds, although query results display it as a 13-digit number; filter with DATE literals (e.g. Report_Date >= DATE '2026-01-01') or date functions." + "dataset_X.start_date": "DATE. Not epoch milliseconds, although query results display it as a 13-digit number; filter with DATE literals (e.g. start_date >= DATE '2026-01-01') or date functions." + "dataset_Y.release_date": "DATE. Not epoch milliseconds, although query results display it as a 13-digit number; filter with DATE literals (e.g. release_date >= DATE '2026-01-01') or date functions." + "dataset_Z.transaction_datetime": "TIMESTAMP. Not epoch milliseconds, although query results display it as a 13-digit number; filter with DATE literals (e.g. transaction_datetime >= DATE '2026-01-01') or date functions." +``` diff --git a/docs/trace.md b/docs/trace.md new file mode 100644 index 0000000..00ab9a4 --- /dev/null +++ b/docs/trace.md @@ -0,0 +1,99 @@ +# Design note: `querypilot trace` and `querypilot hooks` + +Status: first version. Claude Code transcripts and hooks; Palantir Foundry's SQL tool is mapped by default, and any other MCP or connector query tool can be mapped by name. + +## Why + +Agents often run SQL through a tool that QueryPilot cannot sit in front of. Palantir Foundry's MCP server (`mcp__palantir-mcp__run_sql_query_on_foundry_dataset`) is one: Foundry does not speak the Postgres protocol, so `querypilot proxy` cannot intercept it. What QueryPilot can see is the agent's side of the conversation: Claude Code writes every `tool_use` and `tool_result` to a JSONL transcript, and its hooks run before and after each tool call. `trace` reads the first, `hooks` uses the second, and `trace doctor` says where each recurring problem is best fixed: + +- **tool**: QueryPilot's guard, or the tool itself (validate before running, cap rows, retry transport errors); +- **prompt**: a hint the agent should get up front (the dialect's rules, the real name of a column it keeps guessing); +- **data**: the dataset should describe or model itself better (a column description, a view with a proper DATE column). + +## Commands + +```bash +# Read past sessions (default ~/.claude/projects) into the audit log and schema cache +querypilot trace import + +# Diagnose: terminal summary, or a shareable markdown/JSON report (redacted by default) +querypilot trace doctor +querypilot trace doctor ~/.claude/projects --output report.md + +# Export recovered retry chains (failed query -> the SQL that worked) as an eval replay suite +querypilot trace doctor --export-replay chains.jsonl +querypilot eval replay --audit-jsonl chains.jsonl --fixture-db sqlite:///mirror.db --output suite.yaml + +# Log every query tool call live (PostToolUse + PostToolUseFailure) +querypilot hooks install # ~/.claude/settings.json; --scope project|local, or --settings PATH +querypilot hooks install --guard --max-rows 200 --require-schema +querypilot hooks uninstall +``` + +Other tools are mapped with `--tool NAME=SQL_FIELD[:DIALECT[:ROW_LIMIT_FIELD]]`, where `NAME` may be a glob, for example `--tool 'mcp__warehouse__*=query:snowflake'`. Foundry's SQL tool (`sqlQuery`, Spark SQL, `rowLimit`) and its schema tool are always included. + +## What gets recorded + +`trace import` and the hooks write `QueryAuditRecord`s (operation `tool_call`) to `~/.querypilot/trace/audit.jsonl` (`QUERYPILOT_HOME` moves it). Each record holds the SQL, the outcome and failure class, the latency, the session and transcript, the result size in characters and rows, and the result's column header. Result rows are never written. An error message is kept (first 500 characters) because the failure class and the "did you mean" suggestions come from it. `trace import` also keeps the user's last prompt before the call (first 500 characters) as `question`, which is what the replay export needs; all of this stays on the local machine. + +Calls are paired by `tool_use_id`, so a call seen in both a transcript and the live log, or in a resumed session's copied history, counts once. Latency from a transcript is the gap between the `tool_use` and `tool_result` records, which includes any permission prompt; from the hooks it is Claude Code's `duration_ms`. + +## Failure classes + +Foundry reports query failures as ordinary text (`Sql Query Failed ... [UNRESOLVED_COLUMN...]`), not as tool errors, so `trace/classify.py` reads the text. The classes are: unresolved column, unknown table/dataset, data type mismatch, result too large for the context window, timeout, upstream/network, syntax, other SQL analysis error, and invalid tool call. Calls the user or a permission rule refused are counted as denied, not as failures, and calls the guard blocked are counted as blocked. + +Two Foundry details matter. The MCP server appends a "403 Forbidden ... the dataset RID is incorrect" suggestion to every HTTP failure, so the HTTP status decides first: 5xx is upstream, 403/404 is unknown dataset (Foundry answers 403 for both a wrong RID and missing access). Second, a bare `Failed to fetch 400` followed by a stack trace means the tool dropped Foundry's error body, so it counts as upstream: the agent cannot act on it. + +## The guard + +`querypilot hooks install --guard` adds a PreToolUse hook. For each query it: + +1. parses the SQL in the tool's dialect (Spark for Foundry); if sqlglot cannot parse it, the call is allowed, because the engine's parser is the authority; +2. looks up every referenced dataset (RID or path) in the schema cache, which is filled from `get_foundry_dataset_schema` results and from successful `SELECT *` queries on a single dataset; +3. when all of them are cached, runs `SQLValidator` against those schemas and denies unknown or ambiguous columns, with `repair_hints`' "did you mean" text, and denies arithmetic that treats a DATE/TIMESTAMP column as epoch milliseconds (`Report_Date / 1000`, `Report_Date >= 1767225600000`); +4. with `--max-rows N`, denies queries that have neither `LIMIT <= N`, a tool row limit `<= N`, nor an aggregate-only projection; +5. with `--require-schema`, denies queries on datasets the agent has not fetched a schema for, which turns a wrong RID into a cheap schema-call failure. + +A dataset that is not cached is never judged. Datasets get rebuilt with new columns, so a cached schema can go stale. When the guard denies a query for its columns, the message tells the agent it may resend the identical query, and the hook allows that resend. If the query then succeeds, the PostToolUse hook drops the stale cache entry. The denial costs one turn, not the task. + +With `--guard`, the PostToolUse hook also adds context when a result shows a DATE/TIMESTAMP column as a 13-digit number. Foundry renders dates as epoch milliseconds, and in the history below that is where the agent learned to write `Report_Date / 1000`. + +Hooks never fail a tool call: every exception is caught and logged to `hook-errors.log` next to the audit log, and the hook prints nothing. + +## `trace doctor` + +The doctor reports: + +- the failure rate overall and by class, with the layer to fix first; +- **retry chains**: a failed query, then the following queries on the same datasets in the same transcript, up to the first success (recovered) or until the agent moves to other datasets (abandoned). It reports attempts per chain, time spent on failed attempts, chains where the agent fetched a schema only after failing, and exact repeats of queries that had already succeeded; +- **latency**: p50/p90/p99, stalls (`--stall-seconds`, default 60), and p50 by outcome; +- **guard replay**: the history replayed in time order, building the schema cache the way the hooks would, to count the failures the guard would have stopped and the successful queries it would have blocked; +- **findings**, each tagged tool, prompt or data, plus ready-to-paste `hints:` and `descriptions:` blocks for `querypilot.yaml` or the agent's instructions. + +Markdown and JSON reports are redacted by default: dataset RIDs and paths become `dataset_A`, `dataset_B`, ..., string literals become `'?'`, and result rows are never shown. Column names are kept. `--no-redact` shows dataset names, for local use. + +## First run: the Foundry history + +[trace-doctor-foundry-2026-10-21.md](trace-doctor-foundry-2026-10-21.md) is `trace doctor` over three months of Claude Code sessions that used Foundry. A quick scan done before this work estimated the same picture with looser regexes; the doctor's numbers line up: + +| | Quick scan | `trace doctor` | Why they differ | +|---|---:|---:|---| +| Foundry SQL calls | 5,157 | 5,163 | sessions kept running between the two counts | +| Sessions | 228 | 71 sessions, 224 transcripts | the scan counted transcript files, which include subagents | +| Failed | 12% | 11.8% (609) | | +| Unresolved column | 231 | 191 | the scan's `cannot resolve` pattern also matched type-mismatch errors | +| Result too large | 165 | 165 | | +| Unknown table | 80 | 57 | 5xx responses carrying Foundry's "403 ... RID is incorrect" suggestion are upstream errors | +| Data type mismatch | 34 | 47 | adds `CAST_INVALID_INPUT` and unparseable values | +| Network + Foundry internal | 42 + 21 | 107 upstream | adds 37 bare `Failed to fetch 400` errors with no body | +| Latency p50 / p90 | 3.8s / 11.7s | 3.8s / 11.6s | | +| Calls over 60s | 84 | 84 | | + +Replaying the guard over that history, it would have stopped 98 of the 191 unresolved-column failures and 24 of the 47 type mismatches before they reached Foundry. It would have blocked 38 of 3,099 successful queries it could check, mostly on datasets rebuilt with new columns after their schema was cached; with the resend rule each of those costs one turn. + +## Known gaps + +- Only Claude Code's transcript format is read. Other agents need their own reader that produces `ToolCall`s. +- The schema-tool parser understands Foundry's resource tree only. Other schema tools need a parser (`SchemaToolSpec.format`). +- Retry chains are a heuristic: they link queries that share a dataset in one transcript, so two unrelated questions about one dataset can join a chain. +- The guard cannot see Spark struct fields or backticked names that contain dots, so it skips those queries instead of guessing. diff --git a/src/querypilot/audit/types.py b/src/querypilot/audit/types.py index 9838595..7cf0f6e 100644 --- a/src/querypilot/audit/types.py +++ b/src/querypilot/audit/types.py @@ -33,6 +33,11 @@ class QueryAuditRecord(BaseModel): session_id: str | None = None app_name: str | None = None trace_id: str | None = None + # Set for agent tool calls captured by `querypilot trace` and the Claude Code hooks. + tool_name: str | None = None + tool_call_id: str | None = None + # Tool-call context: sequence, failure class, result size, guard decision. + details: dict = Field(default_factory=dict) @classmethod def create( diff --git a/src/querypilot/cli.py b/src/querypilot/cli.py index 1cb4ef0..2ae13c0 100644 --- a/src/querypilot/cli.py +++ b/src/querypilot/cli.py @@ -134,6 +134,11 @@ def main(argv: list[str] | None = None) -> int: ) proxy_parser.add_argument("--verbose", "-v", action="store_true", help="Log every connection and statement.") + from querypilot.trace.cli import add_hooks_parser, add_trace_parser + + add_trace_parser(subparsers) + add_hooks_parser(subparsers) + eval_parser = subparsers.add_parser("eval", help="Eval-driven SQL reliability harness.") eval_subparsers = eval_parser.add_subparsers(dest="eval_command", required=True) @@ -204,6 +209,14 @@ def main(argv: list[str] | None = None) -> int: from querypilot.proxy.cli import run_proxy return run_proxy(args) + if args.command == "trace": + from querypilot.trace.cli import run_trace + + return run_trace(args) + if args.command == "hooks": + from querypilot.trace.cli import run_hooks + + return run_hooks(args) if args.command == "eval": if args.eval_command == "run": return _eval_run(args) diff --git a/src/querypilot/trace/__init__.py b/src/querypilot/trace/__init__.py new file mode 100644 index 0000000..abb7de9 --- /dev/null +++ b/src/querypilot/trace/__init__.py @@ -0,0 +1 @@ +"""Observe the SQL agents run through MCP and connector tools: transcripts, hooks, guard, doctor.""" diff --git a/src/querypilot/trace/classify.py b/src/querypilot/trace/classify.py new file mode 100644 index 0000000..26565e6 --- /dev/null +++ b/src/querypilot/trace/classify.py @@ -0,0 +1,173 @@ +"""Classify what a query tool returned: success, failure class, or not run. + +Foundry's SQL tool reports query failures as ordinary text ("Sql Query +Failed ... [UNRESOLVED_COLUMN...]"), not as tool errors, so classification reads +the text. Patterns cover Spark SQL error classes, Foundry HTTP errors, MCP +transport failures and Claude Code's own messages (oversized results, denied +calls); other engines' messages fall through to the generic patterns. +""" +from __future__ import annotations + +import re +from dataclasses import dataclass + +UNRESOLVED_COLUMN = "unresolved_column" +UNKNOWN_TABLE = "unknown_table" +TYPE_MISMATCH = "type_mismatch" +RESULT_TOO_LARGE = "result_too_large" +TIMEOUT = "timeout" +UPSTREAM = "upstream" +SYNTAX = "syntax" +SQL_ANALYSIS = "sql_analysis" +TOOL_INPUT = "tool_input" +OTHER = "other" + +CLASS_LABELS = { + UNRESOLVED_COLUMN: "Unresolved column", + UNKNOWN_TABLE: "Unknown table / dataset", + TYPE_MISMATCH: "Data type mismatch", + RESULT_TOO_LARGE: "Result too large for context", + TIMEOUT: "Timeout", + UPSTREAM: "Upstream / network (no SQL error returned)", + SYNTAX: "Syntax / parse", + SQL_ANALYSIS: "Other SQL analysis error", + TOOL_INPUT: "Invalid tool call", + OTHER: "Other", +} + +# Call outcomes. +OK = "ok" +ERROR = "error" +# The user, a permission rule or another hook refused the call; it never ran. +DENIED = "denied" +# The QueryPilot guard refused the call before it ran. +BLOCKED = "blocked" +# No tool_result was recorded (interrupted session or still running). +NO_RESULT = "no_result" + +GUARD_PREFIX = "QueryPilot guard" + +_DENIED = re.compile( + r"doesn't want to proceed with this tool use|Permission for this action was denied|" + r"denied by (?:the )?(?:user|hook|PreToolUse)|tool use was rejected", + re.IGNORECASE, +) +_TOO_LARGE = re.compile(r"exceeds maximum allowed tokens|response too large|result is too large", re.IGNORECASE) +_SPILLED = re.compile(r"^|Output too large \(", re.IGNORECASE) +# Only the start of the text: result rows can contain any of these words. +_FAILED_SQL = re.compile( + r"^\s*(?:Sql Query Failed|Error\b|MCP error||Failed to |fetch failed)|" + r"Failed to execute SQL", + re.IGNORECASE, +) + +# Ordered: specific error classes before generic transport text, because Foundry +# wraps some engine errors ("Failed writing blob data ... [CAST_INVALID_INPUT]"). +_RULES: list[tuple[str, re.Pattern[str]]] = [ + (TOOL_INPUT, re.compile(r"Input validation error|No such tool available|Invalid arguments for tool", re.IGNORECASE)), + ( + UNRESOLVED_COLUMN, + re.compile( + r"UNRESOLVED_COLUMN|CANNOT_RESOLVE_STAR_EXPAND|UNRESOLVED_FIELD|AMBIGUOUS_REFERENCE|" + r"no such column|column \"?[\w.]+\"? does not exist|unknown column|invalid identifier", + re.IGNORECASE, + ), + ), + ( + UNKNOWN_TABLE, + re.compile( + r"TABLE_OR_VIEW_NOT_FOUND|DatasetNotFound|referenced dataset, schema, or table could not be found|" + r"dataset RID is incorrect|ReadQueryInputsPermissionDenied|no such table|" + r"relation \"?[\w.]+\"? does not exist|table '[^']+' doesn't exist", + re.IGNORECASE, + ), + ), + ( + TYPE_MISMATCH, + re.compile( + r"DATATYPE_MISMATCH|CAST_INVALID_INPUT|INCOMPATIBLE_DATA_TYPE|" + r"value used by the query could not be parsed|cannot be cast to|operator does not exist", + re.IGNORECASE, + ), + ), + (SYNTAX, re.compile(r"QueryParseError|PARSE_SYNTAX_ERROR|mismatched input|syntax error|cannot be parsed", re.IGNORECASE)), + ( + TIMEOUT, + re.compile(r"timed? ?out|timeout|deadline exceeded|statement cancel|canceled because it wasn't finished", re.IGNORECASE), + ), + ( + UPSTREAM, + re.compile( + r"fetch failed|Failed to fetch 5\d\d|INTERNAL_ERROR|internal error|\"type\": ?\"NETWORK\"|" + r"ECONNRESET|ETIMEDOUT|socket hang up|Service Unavailable|Bad Gateway|errorCode\":\"INTERNAL|" + r"Refer to the server with this errorInstanceId", + re.IGNORECASE, + ), + ), + ( + SQL_ANALYSIS, + re.compile( + r"\[[A-Z][A-Z_]+(?:\.[A-Z_]+)?\]|failed analysis|SqlAnalyzerException|Failed to fetch 4\d\d", re.IGNORECASE + ), + ), +] + + +_HTTP_STATUS = re.compile(r"Failed to fetch (\d{3})\b") + + +@dataclass(frozen=True) +class Outcome: + status: str + error_class: str | None = None + row_count: int | None = None + # Claude Code saved the result to a file and showed the model a preview. + spilled: bool = False + + +def classify_output(text: str, *, is_error: bool = False) -> Outcome: + head = text[:4000] + # Claude Code may wrap a hook's deny reason in its own text. + if GUARD_PREFIX in head[:500]: + return Outcome(BLOCKED) + if _DENIED.search(head[:400]): + return Outcome(DENIED) + if _TOO_LARGE.search(head[:400]): + return Outcome(ERROR, RESULT_TOO_LARGE) + if not is_error and not _FAILED_SQL.search(head[:300]): + return Outcome(OK, row_count=markdown_row_count(text), spilled=bool(_SPILLED.search(head[:200]))) + return Outcome(ERROR, error_class(head)) + + +def error_class(text: str) -> str: + # Foundry HTTP errors: the MCP server appends a "403 Forbidden ... RID is + # incorrect" suggestion to every failure, so the status code decides first. + http = _HTTP_STATUS.search(text[:200]) + if http: + status = int(http.group(1)) + if status >= 500: + return UPSTREAM + if status in (403, 404): + return UNKNOWN_TABLE + if "{" not in text[:400]: + # A bare "Failed to fetch 400" plus a stack trace: the tool dropped the error body. + return UPSTREAM + for name, pattern in _RULES: + if pattern.search(text): + return name + return OTHER + + +def markdown_row_count(text: str) -> int | None: + """Rows in a markdown table result (header and separator excluded).""" + lines = [line for line in text.splitlines() if line.startswith("|")] + if len(lines) < 2 or not set(lines[1].replace("|", "").strip()) <= {"-", ":", " "}: + return None + return len(lines) - 2 + + +def markdown_header(text: str) -> list[str] | None: + lines = [line for line in text.splitlines() if line.startswith("|")] + if len(lines) < 2 or not set(lines[1].replace("|", "").strip()) <= {"-", ":", " "}: + return None + return [cell.strip() for cell in lines[0].strip().strip("|").split("|")] diff --git a/src/querypilot/trace/cli.py b/src/querypilot/trace/cli.py new file mode 100644 index 0000000..05fa175 --- /dev/null +++ b/src/querypilot/trace/cli.py @@ -0,0 +1,288 @@ +from __future__ import annotations + +import argparse +import logging +import sys +from pathlib import Path + +from querypilot.trace.hooks import SCOPES, default_audit_log + +DEFAULT_TRANSCRIPTS = "~/.claude/projects" + + +def add_trace_parser(subparsers) -> None: + trace_parser = subparsers.add_parser( + "trace", help="Import and diagnose the SQL an agent ran through MCP or connector tools." + ) + trace_sub = trace_parser.add_subparsers(dest="trace_command", required=True) + + import_parser = trace_sub.add_parser( + "import", help="Read Claude Code transcripts into the QueryPilot audit log and schema cache." + ) + import_parser.add_argument( + "paths", nargs="*", default=[DEFAULT_TRANSCRIPTS], help="Transcript .jsonl files or directories." + ) + import_parser.add_argument( + "--output", default=None, help=f"Audit JSONL to append to (default: {_display(default_audit_log())})." + ) + import_parser.add_argument("--cache", default=None, help="Schema cache file (default: ~/.querypilot/trace).") + import_parser.add_argument("--no-cache", action="store_true", help="Do not update the schema cache.") + _add_tool_arg(import_parser) + + doctor_parser = trace_sub.add_parser( + "doctor", help="Report failure classes, retry chains, latency, and where each problem is best fixed." + ) + doctor_parser.add_argument( + "paths", + nargs="*", + help=f"Transcripts, directories or audit logs (default: {DEFAULT_TRANSCRIPTS}).", + ) + doctor_parser.add_argument("--format", choices=("terminal", "md", "json"), default=None) + doctor_parser.add_argument("--output", default=None, help="Write the report here (.md or .json).") + doctor_parser.add_argument( + "--no-redact", + action="store_true", + help="Show dataset names and literals in examples (reports are redacted by default).", + ) + doctor_parser.add_argument("--title", default="QueryPilot trace doctor", help="Markdown report title.") + doctor_parser.add_argument("--stall-seconds", type=float, default=60.0, help="Calls slower than this stall.") + doctor_parser.add_argument( + "--max-rows", type=int, default=200, help="Row cap assumed when judging oversized results." + ) + doctor_parser.add_argument( + "--export-replay", + default=None, + metavar="PATH", + help="Write recovered retry chains as audit JSONL for `querypilot eval replay`.", + ) + doctor_parser.add_argument( + "--suite", default=None, metavar="PATH", help="Also build the replay suite (needs --fixture-db)." + ) + doctor_parser.add_argument("--fixture-db", default=None, help="Database URL for the exported suite's cases.") + _add_tool_arg(doctor_parser) + + +def add_hooks_parser(subparsers) -> None: + hooks_parser = subparsers.add_parser( + "hooks", help="Install Claude Code hooks that log (and optionally guard) query tool calls." + ) + hooks_sub = hooks_parser.add_subparsers(dest="hooks_command", required=True) + + install = hooks_sub.add_parser("install", help="Add QueryPilot hooks to a Claude Code settings file.") + _add_scope_args(install) + _add_hook_options(install) + install.add_argument("--dry-run", action="store_true", help="Print the settings instead of writing them.") + + uninstall = hooks_sub.add_parser("uninstall", help="Remove QueryPilot hooks from a settings file.") + _add_scope_args(uninstall) + + run = hooks_sub.add_parser("run", help="Handle one hook event (Claude Code calls this).") + run.add_argument("event", choices=("pre", "post", "post-failure")) + _add_hook_options(run) + + +def run_trace(args: argparse.Namespace) -> int: + logging.getLogger("sqlglot").setLevel(logging.ERROR) + if args.trace_command == "import": + return _trace_import(args) + if args.trace_command == "doctor": + return _trace_doctor(args) + raise SystemExit(f"Unknown trace subcommand: {args.trace_command}") + + +def run_hooks(args: argparse.Namespace) -> int: + from querypilot.trace.hooks import ( + install_hooks, + main_hook, + options_from_args, + settings_path, + uninstall_hooks, + ) + + if args.hooks_command == "run": + return main_hook(args.event, options_from_args(args)) + path = Path(args.settings).expanduser() if args.settings else settings_path(args.scope) + if args.hooks_command == "uninstall": + removed = uninstall_hooks(path) + print(f"Removed {removed} QueryPilot hook(s) from {path}" if removed else f"No QueryPilot hooks in {path}") + return 0 + try: + options = options_from_args(args) + except ValueError as exc: + raise SystemExit(str(exc)) from exc + if args.dry_run: + import json + import tempfile + + with tempfile.TemporaryDirectory() as tmp: + preview = Path(tmp) / "settings.json" + if path.is_file(): + preview.write_text(path.read_text(encoding="utf-8"), encoding="utf-8") + print(json.dumps(install_hooks(preview, options), indent=2)) + return 0 + install_hooks(path, options) + mode = "logging + guard" if options.guard else "logging" + print(f"Installed QueryPilot hooks ({mode}) in {path}") + print(f" audit log: {options.audit_log}") + print(" tools: " + ", ".join(options.registry.tool_names())) + print("Restart Claude Code (or run /hooks) to load them. Remove with: querypilot hooks uninstall") + return 0 + + +def _trace_import(args: argparse.Namespace) -> int: + from querypilot.audit.sinks import JSONLAuditSink + from querypilot.trace.schema_cache import SchemaCache + from querypilot.trace.transcript import ( + QUERY, + SCHEMA, + load_calls, + read_audit_log, + to_audit_record, + ) + + registry = _registry(args) + output = Path(args.output).expanduser() if args.output else default_audit_log() + calls = load_calls(args.paths, registry) + known = {call.call_id for call in read_audit_log(output, registry)} if output.is_file() else set() + sink = JSONLAuditSink(output) + written = 0 + for call in calls: + if call.call_id in known: + continue + sink.write(to_audit_record(call)) + written += 1 + + cache_note = "" + if not args.no_cache: + cache = SchemaCache.load(args.cache) + for call in calls: + if call.kind == SCHEMA and call.dataset is not None: + cache.add(call.dataset) + elif call.kind == QUERY and call.header and call.sql: + cache.learn_from_select_star(call.sql, call.header, call.dialect) + cache_note = f"; schema cache: {len(cache)} datasets at {cache.save(args.cache)}" + queries = sum(call.kind == QUERY for call in calls) + sessions = len({call.session_id for call in calls if call.kind == QUERY}) + print( + f"Read {queries} query calls ({sessions} sessions) and {len(calls) - queries} schema calls; " + f"wrote {written} new records to {output}{cache_note}" + ) + return 0 + + +def _trace_doctor(args: argparse.Namespace) -> int: + from querypilot.trace.doctor import ( + diagnose, + render_markdown, + render_terminal, + replay_records, + ) + from querypilot.trace.transcript import load_calls + + paths = args.paths or [DEFAULT_TRANSCRIPTS] + calls = load_calls(paths, _registry(args)) + if not any(call.kind == "query" for call in calls): + print(f"No query tool calls found in {', '.join(paths)}.", file=sys.stderr) + return 1 + report = diagnose(calls, stall_seconds=args.stall_seconds, max_rows=args.max_rows, redact=not args.no_redact) + + fmt = args.format or _infer_format(args.output) or "terminal" + if fmt == "md": + text = render_markdown(report, title=args.title) + elif fmt == "json": + text = report.to_json() + else: + text = render_terminal(report) + if args.output: + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + Path(args.output).write_text(text, encoding="utf-8") + print(render_terminal(report)) + print(f"\nWrote {fmt} report to {args.output}") + else: + print(text) + + if args.export_replay or args.suite: + records = replay_records(calls) + target = Path(args.export_replay or Path(args.suite).with_suffix(".replay.jsonl")) + target.parent.mkdir(parents=True, exist_ok=True) + with target.open("w", encoding="utf-8") as handle: + for record in records: + handle.write(record.model_dump_json() + "\n") + print(f"Wrote {len(records)} recovered retry chains to {target}") + if args.suite: + if not args.fixture_db: + raise SystemExit("--suite needs --fixture-db (the database the cases run against).") + from querypilot.evals.loader import write_suite + from querypilot.evals.replay import replay_from_jsonl + + suite = replay_from_jsonl( + target, fixture_db=args.fixture_db, suite_name="trace_regressions", extra_tags=["trace"] + ) + print(f"Wrote {len(suite.cases)} cases to {write_suite(suite, args.suite)}") + return 0 + + +def _registry(args): + from querypilot.trace.tools import ToolRegistry + + try: + return ToolRegistry.from_args(args.tool) + except ValueError as exc: + raise SystemExit(str(exc)) from exc + + +def _add_tool_arg(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--tool", + action="append", + default=[], + metavar="NAME=SQL_FIELD[:DIALECT[:ROW_LIMIT_FIELD]]", + help=( + "Another query tool to track, e.g. mcp__warehouse__run_query=query:snowflake " + "(repeatable; NAME may be a glob). Foundry's SQL tool is always tracked." + ), + ) + + +def _add_scope_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--scope", + choices=SCOPES, + default="user", + help="user: ~/.claude/settings.json (default); project: .claude/settings.json; " + "local: .claude/settings.local.json.", + ) + parser.add_argument("--settings", default=None, help="Explicit settings file (overrides --scope).") + + +def _add_hook_options(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--audit-log", default=None, help=f"Audit JSONL (default: {_display(default_audit_log())}).") + parser.add_argument("--cache", default=None, help="Schema cache file (default: ~/.querypilot/trace).") + parser.add_argument( + "--guard", + action="store_true", + help="Validate SQL against cached schemas before it runs and deny unknown columns/tables.", + ) + parser.add_argument( + "--max-rows", type=int, default=None, help="With --guard: deny queries without a row cap at or below this." + ) + parser.add_argument( + "--require-schema", + action="store_true", + help="With --guard: deny queries on datasets whose schema the agent has not fetched yet.", + ) + _add_tool_arg(parser) + + +def _infer_format(path: str | None) -> str | None: + if not path: + return None + suffix = Path(path).suffix.lower() + return {".md": "md", ".markdown": "md", ".json": "json"}.get(suffix) + + +def _display(path: Path) -> str: + try: + return "~/" + str(path.relative_to(Path.home())) + except ValueError: + return str(path) diff --git a/src/querypilot/trace/doctor.py b/src/querypilot/trace/doctor.py new file mode 100644 index 0000000..af977cc --- /dev/null +++ b/src/querypilot/trace/doctor.py @@ -0,0 +1,939 @@ +"""Diagnose an agent's query-tool history and say where each problem is best fixed. + +Every finding is tagged with the layer that should change: + +- **tool**: QueryPilot's guard or the tool itself (validate before running, + cap rows, retry transport errors, render dates readably); +- **prompt**: a hint the agent should get up front (dialect rules, the real + name of a column it keeps guessing); +- **data**: the dataset should describe or model itself better (a column + description, a view with a proper DATE column). + +The guard replay walks the history in time order, building the schema cache +the way the hooks would, and asks what the guard would have done with each +query: how many failures it would have stopped before they reached the +warehouse, and how many successful queries it would have blocked by mistake. +""" +from __future__ import annotations + +import difflib +import json +import re +import statistics +from collections import Counter, defaultdict +from dataclasses import asdict, dataclass, field +from datetime import UTC, datetime + +import sqlglot +from sqlglot import exp + +from querypilot.audit.types import QueryAuditRecord +from querypilot.trace.classify import ( + BLOCKED, + CLASS_LABELS, + DENIED, + ERROR, + NO_RESULT, + OK, + RESULT_TOO_LARGE, + SQL_ANALYSIS, + SYNTAX, + TIMEOUT, + TOOL_INPUT, + TYPE_MISMATCH, + UNKNOWN_TABLE, + UNRESOLVED_COLUMN, + UPSTREAM, +) +from querypilot.trace.guard import guard_sql +from querypilot.trace.schema_cache import ( + FOUNDRY_RID, + SchemaCache, + table_ref_name, + table_refs, +) +from querypilot.trace.transcript import QUERY, SCHEMA, ToolCall + +TOOL_FIX = "tool" +PROMPT_FIX = "prompt" +DATA_FIX = "data" + +# Where each failure class is best fixed first. +PRIMARY_FIX = { + UNRESOLVED_COLUMN: TOOL_FIX, + UNKNOWN_TABLE: TOOL_FIX, + TYPE_MISMATCH: DATA_FIX, + RESULT_TOO_LARGE: TOOL_FIX, + TIMEOUT: TOOL_FIX, + UPSTREAM: TOOL_FIX, + SYNTAX: PROMPT_FIX, + SQL_ANALYSIS: PROMPT_FIX, + TOOL_INPUT: PROMPT_FIX, +} + +_UNRESOLVED = re.compile( + r"with name (`[^`]+`(?:\.`[^`]+`)*) cannot be resolved\.(?: Did you mean one of the following\? \[([^\]]*)\])?" +) +_MISMATCH = re.compile(r'Cannot resolve "(.+?)" due to data type mismatch.*?\("([A-Z]+)" and "([A-Z]+)"\)', re.DOTALL) +_BACKTICKED = re.compile(r"`([^`]+)`") + + +@dataclass +class ClassRow: + error_class: str + label: str + count: int + share_of_calls: float + share_of_failures: float + fix: str + sessions: int + + +@dataclass +class ChainStats: + chains: int = 0 + recovered: int = 0 + abandoned: int = 0 + failed_attempts: int = 0 + max_attempts: int = 0 + attempts_histogram: dict[str, int] = field(default_factory=dict) + schema_fetched_in_chain: int = 0 + wasted_seconds: float = 0.0 + recovery_seconds_p50: float | None = None + by_first_class: dict[str, int] = field(default_factory=dict) + duplicate_successes: int = 0 + + +@dataclass +class LatencyStats: + measured: int = 0 + p50: float | None = None + p90: float | None = None + p99: float | None = None + max: float | None = None + stall_seconds: float = 60.0 + stalls: int = 0 + stalls_failed: int = 0 + over_300s: int = 0 + no_result: int = 0 + by_class_p50: dict[str, float] = field(default_factory=dict) + + +@dataclass +class GuardReplay: + checked: int = 0 + caught: dict[str, int] = field(default_factory=dict) + false_blocks: int = 0 + successes_checked: int = 0 + false_block_examples: list[str] = field(default_factory=list) + caught_examples: list[str] = field(default_factory=list) + cached_datasets: int = 0 + limit_missing_in_too_large: int = 0 + select_star_in_too_large: int = 0 + too_large_row_limits: dict[str, int] = field(default_factory=dict) + + +@dataclass +class Finding: + fix: str + title: str + impact: int + detail: str + suggestion: str + examples: list[str] = field(default_factory=list) + + +@dataclass +class DoctorReport: + tools: list[str] + calls: int + sessions: int + threads: int + ran: int + failed: int + denied: int + blocked: int + no_result: int + failure_rate: float + first_call: str | None + last_call: str | None + classes: list[ClassRow] + chains: ChainStats + latency: LatencyStats + guard: GuardReplay + findings: list[Finding] + hints: list[str] + descriptions: dict[str, str] + + def to_json(self) -> str: + return json.dumps(asdict(self), indent=2, default=str) + + +# --------------------------------------------------------------------------- redaction + + +class Redactor: + """Stable labels for datasets and '?' for literals, so reports can be shared.""" + + def __init__(self, enabled: bool = True, cache: SchemaCache | None = None) -> None: + self.enabled = enabled + self.cache = cache + self._labels: dict[str, str] = {} + + def dataset(self, ref: str) -> str: + ref = ref.strip("`") + if not self.enabled: + hit = self.cache.lookup(ref) if self.cache else None + return hit.display_name if hit else ref + hit = self.cache.lookup(ref) if self.cache else None + key = hit.key if hit else ref + if key not in self._labels: + self._labels[key] = f"dataset_{_letters(len(self._labels))}" + return self._labels[key] + + def sql(self, sql: str | None, dialect: str = "spark", limit: int = 240) -> str: + if not sql: + return "" + if not self.enabled: + return _one_line(sql)[:limit] + try: + expression = sqlglot.parse_one(sql, read=dialect) + for table in table_refs(expression): + table.set("this", exp.to_identifier(self.dataset(table_ref_name(table)))) + table.set("db", None) + table.set("catalog", None) + for literal in list(expression.find_all(exp.Literal)): + if literal.is_string: + literal.replace(exp.Literal.string("?")) + text = expression.sql(dialect=dialect) + except Exception: # noqa: BLE001 - fall back to text redaction + text = self.text(sql) + return _one_line(text)[:limit] + + def guard_text(self, text: str) -> str: + """A guard reply with dataset names removed; its quoted words are column names, not data.""" + if not self.enabled: + return _one_line(text)[:240] + text = re.sub(r"\b[A-Za-z_]\w*\.(?=[A-Za-z_])", "", text) + text = re.sub( + r"did you mean ([^?]+)\?", + lambda m: "did you mean " + ", ".join(dict.fromkeys(m.group(1).split(", "))) + "?", + text, + ) + text = FOUNDRY_RID.sub(lambda m: self.dataset(m.group(0)), text) + return _one_line(re.sub(r"/[\w .-]+(?:/[\w .-]+)+", "", text))[:240] + + def text(self, text: str, limit: int = 240) -> str: + if not self.enabled: + return _one_line(text)[:limit] + text = FOUNDRY_RID.sub(lambda m: self.dataset(m.group(0)), text) + text = re.sub(r"(?", text) + return _one_line(text)[:limit] + + +def _plural(n: int, word: str) -> str: + return f"{n} {word}" if n == 1 else f"{n} {word}s" + + +def _letters(n: int) -> str: + out = "" + n += 1 + while n: + n, rem = divmod(n - 1, 26) + out = chr(65 + rem) + out + return out + + +def _one_line(text: str) -> str: + return re.sub(r"\s+", " ", text).strip() + + +# --------------------------------------------------------------------------- analysis + + +def diagnose( + calls: list[ToolCall], + *, + stall_seconds: float = 60.0, + max_rows: int = 200, + redact: bool = True, +) -> DoctorReport: + queries = [c for c in calls if c.kind == QUERY] + ran = [c for c in queries if c.ran and c.status != NO_RESULT] + failed = [c for c in ran if c.status == ERROR] + cache = SchemaCache() + guard = _replay_guard(calls, cache, max_rows=max_rows) + redactor = Redactor(redact, cache) + guard.caught_examples = [redactor.guard_text(e) for e in guard.caught_examples] + guard.false_block_examples = [redactor.guard_text(e) for e in guard.false_block_examples] + + classes = _class_rows(failed, len(ran)) + chains, _ = retry_chains(calls) + latency = _latency(queries, stall_seconds) + findings, hints, descriptions = _findings(calls, failed, classes, chains, latency, guard, cache, redactor) + starts = [c.started_at for c in queries if c.started_at] + return DoctorReport( + tools=sorted({c.tool for c in queries}), + calls=len(queries), + sessions=len({c.session_id for c in queries}), + threads=len({c.thread for c in queries}), + ran=len(ran), + failed=len(failed), + denied=sum(c.status == DENIED for c in queries), + blocked=sum(c.status == BLOCKED for c in queries), + no_result=sum(c.status == NO_RESULT for c in queries), + failure_rate=len(failed) / len(ran) if ran else 0.0, + first_call=min(starts).date().isoformat() if starts else None, + last_call=max(starts).date().isoformat() if starts else None, + classes=classes, + chains=chains, + latency=latency, + guard=guard, + findings=findings, + hints=hints, + descriptions=descriptions, + ) + + +def _class_rows(failed: list[ToolCall], ran: int) -> list[ClassRow]: + counts = Counter(c.error_class for c in failed) + sessions: dict[str, set] = defaultdict(set) + for call in failed: + sessions[call.error_class].add(call.session_id) + return [ + ClassRow( + error_class=name, + label=CLASS_LABELS.get(name, name), + count=count, + share_of_calls=count / ran if ran else 0.0, + share_of_failures=count / len(failed) if failed else 0.0, + fix=PRIMARY_FIX.get(name, PROMPT_FIX), + sessions=len(sessions[name]), + ) + for name, count in counts.most_common() + ] + + +def query_tables(sql: str | None, dialect: str) -> set[str]: + if not sql: + return set() + try: + expression = sqlglot.parse_one(sql, read=dialect) + return {table_ref_name(t).lower() for t in table_refs(expression)} + except Exception: # noqa: BLE001 - fall back to backticked names + return {m.lower() for m in _BACKTICKED.findall(sql) if m.startswith(("ri.", "/"))} + + +def retry_chains(calls: list[ToolCall]) -> tuple[ChainStats, list[tuple[list[ToolCall], ToolCall | None]]]: + """Failed query -> follow-up queries on the same tables, until one succeeds. + + Returns the stats and each chain as (failed calls, the call that recovered or None). + """ + stats = ChainStats() + found: list[tuple[list[ToolCall], ToolCall | None]] = [] + by_thread: dict[str | None, list[ToolCall]] = defaultdict(list) + for call in calls: + if call.ran and call.status != NO_RESULT: + by_thread[call.thread].append(call) + histogram: Counter = Counter() + recovery_times: list[float] = [] + first_classes: Counter = Counter() + + for thread_calls in by_thread.values(): + thread_calls.sort(key=lambda c: (c.started_at.timestamp() if c.started_at else 0.0, c.sequence)) + seen_ok: Counter = Counter() + chain: list[ToolCall] | None = None + tables: set[str] = set() + schema_seen = False + + def close(recovered_by: ToolCall | None) -> None: + nonlocal chain, schema_seen + if not chain: + return + found.append((chain, recovered_by)) + stats.chains += 1 + stats.failed_attempts += len(chain) + stats.max_attempts = max(stats.max_attempts, len(chain)) + histogram[str(len(chain)) if len(chain) < 5 else "5+"] += 1 + stats.wasted_seconds += sum(c.latency_s or 0 for c in chain) + first_classes[chain[0].error_class] += 1 + if schema_seen: + stats.schema_fetched_in_chain += 1 + if recovered_by is not None: + stats.recovered += 1 + end = recovered_by.ended_at or recovered_by.started_at + if end and chain[0].started_at: + recovery_times.append((end - chain[0].started_at).total_seconds()) + else: + stats.abandoned += 1 + chain, schema_seen = None, False + + for call in thread_calls: + if call.kind == SCHEMA: + if chain: + schema_seen = True + continue + call_tables = query_tables(call.sql, call.dialect) + if chain is not None and call_tables and tables and not (call_tables & tables): + close(None) + if call.status == ERROR: + if chain is None: + chain, tables = [], set(call_tables) + chain.append(call) + tables |= call_tables + elif call.status == OK: + close(call) + key = _normalized(call.sql) + if key: + seen_ok[key] += 1 + if seen_ok[key] > 1: + stats.duplicate_successes += 1 + close(None) + + stats.attempts_histogram = dict(sorted(histogram.items())) + stats.recovery_seconds_p50 = statistics.median(recovery_times) if recovery_times else None + stats.by_first_class = dict(first_classes.most_common()) + return stats, found + + +def _latency(queries: list[ToolCall], stall_seconds: float) -> LatencyStats: + measured = [c for c in queries if c.ran and c.status != NO_RESULT and c.latency_s is not None] + values = sorted(c.latency_s for c in measured) + stats = LatencyStats(measured=len(values), stall_seconds=stall_seconds) + stats.no_result = sum(c.status == NO_RESULT for c in queries) + if not values: + return stats + stats.p50 = _percentile(values, 50) + stats.p90 = _percentile(values, 90) + stats.p99 = _percentile(values, 99) + stats.max = values[-1] + stalls = [c for c in measured if c.latency_s > stall_seconds] + stats.stalls = len(stalls) + stats.stalls_failed = sum(c.status == ERROR for c in stalls) + stats.over_300s = sum(c.latency_s > 300 for c in measured) + by_class: dict[str, list[float]] = defaultdict(list) + for call in measured: + by_class[call.error_class or "ok"].append(call.latency_s) + stats.by_class_p50 = {k: round(statistics.median(v), 1) for k, v in sorted(by_class.items())} + return stats + + +def _percentile(values: list[float], pct: float) -> float: + if len(values) == 1: + return values[0] + rank = (len(values) - 1) * pct / 100 + low = int(rank) + high = min(low + 1, len(values) - 1) + return values[low] + (values[high] - values[low]) * (rank - low) + + +def _replay_guard(calls: list[ToolCall], cache: SchemaCache, *, max_rows: int) -> GuardReplay: + replay = GuardReplay() + caught: Counter = Counter() + row_limits: Counter = Counter() + for call in calls: + if call.kind == SCHEMA: + if call.dataset is not None: + cache.add(call.dataset) + continue + if not call.sql or not call.ran or call.status == NO_RESULT: + continue + decision = guard_sql(call.sql, cache, dialect=call.dialect) + if decision.checked_columns: + replay.checked += 1 + if call.status == OK: + replay.successes_checked += decision.checked_columns + if decision.denied: + # Live, the agent resends the query, it succeeds, and the hook evicts the stale entry. + replay.false_blocks += 1 + if len(replay.false_block_examples) < 8: + replay.false_block_examples.append(_first_line(decision.reason)) + cache.evict(decision.tables) + if call.header: + cache.learn_from_select_star(call.sql, call.header, call.dialect) + elif decision.denied: + caught[call.error_class] += 1 + if len(replay.caught_examples) < 5: + replay.caught_examples.append(_first_line(decision.reason)) + if call.error_class == RESULT_TOO_LARGE: + limited = guard_sql(call.sql, SchemaCache(), dialect=call.dialect, max_rows=max_rows, + tool_row_limit=_int(call.params.get("rowLimit"))) + replay.limit_missing_in_too_large += limited.rule == "limit" + replay.select_star_in_too_large += bool(re.search(r"select\s+(?:\w+\.)?\*", call.sql, re.IGNORECASE)) + row_limits[_bucket(_int(call.params.get("rowLimit")))] += 1 + replay.caught = dict(caught.most_common()) + replay.cached_datasets = len(cache) + replay.too_large_row_limits = {k: row_limits[k] for k in _BUCKETS if k in row_limits} + return replay + + +_BUCKETS = ("<= 50", "<= 100", "<= 200", "<= 500", "<= 1000", "> 1000", "not set") + + +def _bucket(value: int | None) -> str: + if value is None: + return "not set" + for bound in (50, 100, 200, 500, 1000): + if value <= bound: + return f"<= {bound}" + return "> 1000" + + +def _int(value) -> int | None: + try: + return int(value) + except (TypeError, ValueError): + return None + + +def _first_line(text: str | None) -> str: + """The guard's first problem, without the column listing that follows it.""" + lines = [line for line in (text or "").splitlines() if line.startswith("- ")] + first = lines[0][2:] if lines else (text or "") + return first.split(" Columns of ")[0] + + +def _normalized(sql: str | None) -> str | None: + return re.sub(r"\s+", " ", sql.strip().lower()) if sql else None + + +# --------------------------------------------------------------------------- findings + + +def _findings( + calls: list[ToolCall], + failed: list[ToolCall], + classes: list[ClassRow], + chains: ChainStats, + latency: LatencyStats, + guard: GuardReplay, + cache: SchemaCache, + redactor: Redactor, +) -> tuple[list[Finding], list[str], dict[str, str]]: + counts = {row.error_class: row.count for row in classes} + findings: list[Finding] = [] + hints: list[str] = [] + descriptions: dict[str, str] = {} + + # Unresolved columns: guard first, then recurring guesses as prompt/data fixes. + if counts.get(UNRESOLVED_COLUMN): + n = counts[UNRESOLVED_COLUMN] + caught = guard.caught.get(UNRESOLVED_COLUMN, 0) + findings.append( + Finding( + TOOL_FIX, + "Validate columns before the query runs (QueryPilot guard)", + n, + f"{n} queries referenced a column the dataset does not have. Replaying the history, " + f"the guard (schema cache built from earlier schema calls and SELECT * results) would have " + f"stopped {caught} of them before they reached the warehouse, with a 'did you mean' reply, " + f"and wrongly blocked {guard.false_blocks} of {guard.successes_checked} successful queries it " + "could check. The rest touched datasets with no cached schema.", + "querypilot hooks install --guard (optionally --require-schema, which makes the agent fetch " + "a dataset's schema before its first query).", + guard.caught_examples[:3], + ) + ) + guesses = _column_guesses(failed) + recurring = [(wrong, data) for wrong, data in guesses if data["count"] >= 2] + if recurring: + lines = [] + for wrong, data in recurring[:8]: + best = data["best"] + tables = ", ".join(sorted({redactor.dataset(t) for t in data["tables"]})[:2]) or "its dataset" + lines.append( + f"`{wrong}` x{data['count']} in {_plural(len(data['sessions']), 'session')} " + f"({tables}); closest real columns: {', '.join(data['closest'][:3]) or 'none listed'}" + ) + if best: + hints.append(f"In {tables}, there is no `{wrong}` column; use `{best}`.") + findings.append( + Finding( + PROMPT_FIX, + "Tell the agent the real names of columns it keeps guessing", + sum(d["count"] for _, d in recurring), + "The same wrong column name failed repeatedly, often across sessions, so the agent is " + "not carrying the schema forward. Recurring guesses:", + "Add the hints below to the agent's instructions (CLAUDE.md, or `hints:` in " + "querypilot.yaml); when a guess recurs across many sessions, the data fix is a column " + "description or an alias column in a view.", + lines, + ) + ) + schema_after = chains.schema_fetched_in_chain + if schema_after: + findings.append( + Finding( + PROMPT_FIX, + "Fetch the schema before the first query, not after it fails", + schema_after, + f"In {schema_after} retry chains the agent fetched a dataset schema only after a query " + "failed.", + "Instruct: 'Before querying a Foundry dataset for the first time in a session, call " + "get_foundry_dataset_schema and use its exact column names.' The guard's --require-schema " + "enforces this.", + ) + ) + + if counts.get(UNKNOWN_TABLE): + n = counts[UNKNOWN_TABLE] + findings.append( + Finding( + TOOL_FIX, + "Catch wrong or inaccessible dataset references", + n, + f"{n} queries named a dataset that was not found or not readable (Foundry answers 403 for " + f"both a wrong RID and missing access). The guard caught {guard.caught.get(UNKNOWN_TABLE, 0)} " + "(malformed RIDs).", + "Keep a registry of the datasets the agent uses (RID, path, one-line purpose) in its " + "instructions, and enable --require-schema so a typo fails at the cheap schema call.", + ) + ) + + if counts.get(TYPE_MISMATCH): + n = counts[TYPE_MISMATCH] + epoch = _epoch_mismatches(failed, cache) + epoch_total = sum(epoch.values()) + rendered = _epoch_rendered(calls, cache) + examples = [ + f"`{col}` ({kind}, {redactor.dataset(table)}) x{count}" for (col, kind, table), count in epoch.most_common(8) + ] + for col, kind, table in list(epoch)[:12]: + descriptions[f"{redactor.dataset(table)}.{col}"] = ( + f"{kind}. Not epoch milliseconds, although query results display it as a 13-digit number; " + f"filter with DATE literals (e.g. {col} >= DATE '2026-01-01') or date functions." + ) + if epoch_total: + hints.append( + "Foundry SQL results show DATE/TIMESTAMP values as epoch-millisecond numbers, but the " + "columns are real dates: never divide them by 1000 or compare them with 13-digit numbers; " + "use DATE 'YYYY-MM-DD' literals and date functions." + ) + findings.append( + Finding( + DATA_FIX, + "Date columns read as epoch milliseconds", + epoch_total or n, + f"{n} queries failed on a data type mismatch; {epoch_total} of them treated a DATE or " + f"TIMESTAMP column as epoch milliseconds (`col / 1000`, `col >= 1767225600000`). " + f"{rendered} successful results displayed a DATE/TIMESTAMP column as 13-digit numbers, " + "which is where the agent learns the wrong type.", + "Data: add column descriptions that state the type (suggested `descriptions:` below), or " + "publish views with explicit DATE columns. Tool: the PostToolUse hook annotates results whose " + "date columns are rendered as epoch ms, and the guard blocks `date_col / 1000`.", + examples, + ) + ) + + if counts.get(RESULT_TOO_LARGE): + n = counts[RESULT_TOO_LARGE] + findings.append( + Finding( + TOOL_FIX, + "Cap result size before it floods the context window", + n, + f"{n} results exceeded Claude Code's output limit. {guard.limit_missing_in_too_large} had " + f"neither a LIMIT nor a rowLimit at or below the cap, and {guard.select_star_in_too_large} " + f"used SELECT *. rowLimit on those calls: " + + ", ".join(f"{k}: {v}" for k, v in guard.too_large_row_limits.items()) + + ". Wide text columns overflow even at 100 rows.", + "querypilot hooks install --guard --max-rows 200 enforces a cap; the prompt fix is to " + "select only the needed columns and aggregate in SQL instead of reading raw rows.", + ) + ) + + if latency.stalls or counts.get(TIMEOUT): + findings.append( + Finding( + TOOL_FIX, + "Long-running queries stall the agent", + latency.stalls, + f"{latency.stalls} calls took longer than {latency.stall_seconds:.0f}s " + f"({latency.over_300s} over 5 minutes); {latency.stalls_failed} of those failed anyway. " + f"{counts.get(TIMEOUT, 0)} calls failed with a timeout.", + "Cap rows and require a filter on the partition/date column for large datasets; for queries " + "the agent repeats, materialize the result as a dataset or view (data fix).", + ) + ) + + if counts.get(UPSTREAM): + n = counts[UPSTREAM] + findings.append( + Finding( + TOOL_FIX, + "Transport and server errors reach the agent as opaque failures", + n, + f"{n} calls failed without a SQL error: network fetch failures, Foundry internal errors, and " + "HTTP errors whose body the MCP tool dropped (a stack trace instead of a message).", + "Fix in the tool layer: retry with backoff on network and 5xx errors and surface Foundry's " + "error body. The agent cannot fix these by rewriting SQL.", + ) + ) + + syntax_like = counts.get(SYNTAX, 0) + counts.get(SQL_ANALYSIS, 0) + if syntax_like: + codes = Counter( + m.group(1) for c in failed if c.error_class in (SYNTAX, SQL_ANALYSIS) + for m in [re.search(r"\[([A-Z][A-Z_]+)(?:\.[A-Z_]+)?\]", c.output)] if m + ) + findings.append( + Finding( + PROMPT_FIX, + "Spark SQL dialect mistakes", + syntax_like, + f"{syntax_like} queries failed to parse or failed Spark analysis " + + (f"({', '.join(f'{k} x{v}' for k, v in codes.most_common(5))})." if codes else "."), + "State the dialect in the agent's instructions: 'Foundry SQL is Spark SQL: quote dataset " + "RIDs in backticks, every non-aggregated column must be in GROUP BY, use try_cast/try_divide.'", + ) + ) + hints.append( + "Foundry SQL is Spark SQL: quote dataset RIDs and paths in backticks, list every " + "non-aggregated column in GROUP BY, and prefer try_cast / try_divide." + ) + + if chains.duplicate_successes: + findings.append( + Finding( + TOOL_FIX, + "Identical queries re-run in the same session", + chains.duplicate_successes, + f"{chains.duplicate_successes} successful queries repeated a query that had already succeeded " + "earlier in the same session.", + "Cache results per session in the tool, or tell the agent to keep earlier results instead of " + "re-querying.", + ) + ) + + order = {TOOL_FIX: 0, PROMPT_FIX: 1, DATA_FIX: 2} + findings.sort(key=lambda f: (-f.impact, order[f.fix])) + return findings, list(dict.fromkeys(hints)), descriptions + + +def _column_guesses(failed: list[ToolCall]) -> list[tuple[str, dict]]: + out: dict[str, dict] = {} + for call in failed: + if call.error_class != UNRESOLVED_COLUMN: + continue + match = _UNRESOLVED.search(call.output) + if not match: + continue + wrong = match.group(1).replace("`", "").split(".")[-1] + suggestions = [s.strip("` ").split(".")[-1] for s in (match.group(2) or "").split(",") if s.strip()] + entry = out.setdefault( + wrong.lower(), {"count": 0, "sessions": set(), "suggestions": [], "tables": set(), "best": None} + ) + entry["count"] += 1 + entry["sessions"].add(call.session_id) + entry["tables"] |= query_tables(call.sql, call.dialect) + for s in suggestions: + if s not in entry["suggestions"]: + entry["suggestions"].append(s) + for wrong, entry in out.items(): + ranked = sorted( + entry["suggestions"], key=lambda s: -difflib.SequenceMatcher(None, wrong, s.lower()).ratio() + ) + entry["closest"] = ranked + if ranked and difflib.SequenceMatcher(None, wrong, ranked[0].lower()).ratio() >= 0.6: + entry["best"] = ranked[0] + return sorted(out.items(), key=lambda item: -item[1]["count"]) + + +def _epoch_mismatches(failed: list[ToolCall], cache: SchemaCache) -> Counter: + # Only real dataset columns; Spark also reports CTE aliases such as `min`. + known = {c.name.lower() for d in cache.datasets.values() for c in d.columns} + found: Counter = Counter() + for call in failed: + if call.error_class != TYPE_MISMATCH: + continue + match = _MISMATCH.search(call.output) + if not match: + continue + expression, left, right = match.groups() + types = {left, right} + if not types & {"DATE", "TIMESTAMP"} or not types & {"INT", "BIGINT", "LONG", "DOUBLE"}: + continue + if not re.search(r"\b1000\b|86400000|\b\d{12,}\b", expression): + continue + column = re.match(r"\(?\s*([A-Za-z_]\w*)", expression) + if column and column.group(1).lower() in known: + table = next(iter(sorted(query_tables(call.sql, call.dialect))), "?") + found[(column.group(1), "DATE" if "DATE" in types else "TIMESTAMP", table)] += 1 + return found + + +def _epoch_rendered(calls: list[ToolCall], cache: SchemaCache) -> int: + """Successful results whose DATE/TIMESTAMP columns show up as 13-digit numbers.""" + date_columns = { + c.name.lower() for d in cache.datasets.values() for c in d.columns if c.type.upper().startswith(("DATE", "TIMESTAMP")) + } + hits = 0 + for call in calls: + if call.kind != QUERY or call.status != OK or not call.header: + continue + idx = [i for i, name in enumerate(call.header) if name.lower() in date_columns] + if not idx: + continue + rows = [line for line in call.output.splitlines() if line.startswith("|")][2:4] + for row in rows: + cells = [cell.strip() for cell in row.strip().strip("|").split("|")] + if any(i < len(cells) and re.fullmatch(r"-?\d{12,14}", cells[i]) for i in idx): + hits += 1 + break + return hits + + +# --------------------------------------------------------------------------- export + + +def replay_records(calls: list[ToolCall]) -> list[QueryAuditRecord]: + """Recovered retry chains as `eval replay` input: the question and the SQL that worked. + + Each record is an ``ask`` audit record whose ``rewritten_sql`` is the query + that finally succeeded, so ``querypilot eval replay`` turns it into a case + with that query as gold SQL. The failed attempts ride along in ``details``. + """ + _, chains = retry_chains(calls) + records = [] + for failed, recovered in chains: + if recovered is None or not recovered.sql: + continue + question = failed[0].question or recovered.question + if not question: + continue + records.append( + QueryAuditRecord( + timestamp=failed[0].started_at or datetime.now(UTC), + operation="ask", + question=question, + sql=recovered.sql, + rewritten_sql=recovered.sql, + valid=True, + executed=True, + row_count=recovered.row_count if recovered.row_count else 1, + session_id=recovered.session_id, + app_name=f"trace-{failed[0].error_class}", + tool_name=recovered.tool, + tool_call_id=recovered.call_id, + details={ + "failed_sql": [c.sql for c in failed], + "failure_classes": [c.error_class for c in failed], + "dialect": recovered.dialect, + }, + ) + ) + return records + + +# --------------------------------------------------------------------------- rendering + + +def render_markdown(report: DoctorReport, *, title: str = "QueryPilot trace doctor") -> str: + r = report + lines = [f"# {title}", ""] + period = f"{r.first_call} to {r.last_call}" if r.first_call else "n/a" + lines += [ + f"Tools: {', '.join(f'`{t}`' for t in r.tools) or 'none'}. Period: {period}.", + "", + f"- **{r.calls:,}** query calls in **{r.sessions}** sessions ({r.threads} transcripts, including subagents)", + f"- **{r.failed:,}** failed: **{r.failure_rate:.1%}** of the {r.ran:,} that ran" + + (f"; {r.denied} denied by the user or a permission rule" if r.denied else "") + + (f"; {r.blocked} blocked by the QueryPilot guard" if r.blocked else "") + + (f"; {r.no_result} never returned" if r.no_result else ""), + ] + lat = r.latency + if lat.p50 is not None: + lines.append( + f"- Latency p50 **{lat.p50:.1f}s**, p90 **{lat.p90:.1f}s**, p99 {lat.p99:.1f}s, max {lat.max:.0f}s; " + f"**{lat.stalls}** calls over {lat.stall_seconds:.0f}s ({lat.over_300s} over 5 min)" + ) + lines += ["", "## Failures by class", "", "| Class | Calls | % of calls | % of failures | Sessions | Fix first in |", + "|---|---:|---:|---:|---:|---|"] + for row in r.classes: + lines.append( + f"| {row.label} | {row.count} | {row.share_of_calls:.1%} | {row.share_of_failures:.0%} | " + f"{row.sessions} | {row.fix} |" + ) + c = r.chains + lines += [ + "", + "## Retry chains", + "", + ("A chain starts at a failed query and follows the next queries on the same datasets in the same " + "transcript until one succeeds (recovered) or the agent moves on (abandoned)."), + "", + f"- {c.chains} chains, {c.failed_attempts} failed attempts; {c.recovered} recovered, {c.abandoned} abandoned", + "- Failed attempts per chain: " + + ", ".join(f"{k}: {v}" for k, v in c.attempts_histogram.items()) + + f" (max {c.max_attempts})", + f"- Time spent on failed attempts: {c.wasted_seconds / 60:.0f} min" + + (f"; median time from first failure to recovery {c.recovery_seconds_p50:.0f}s" if c.recovery_seconds_p50 else ""), + f"- Agent fetched a schema mid-chain in {c.schema_fetched_in_chain} chains", + "- First failure class: " + ", ".join(f"{CLASS_LABELS.get(k, k)} {v}" for k, v in c.by_first_class.items()), + f"- {c.duplicate_successes} successful queries were exact repeats of an earlier success in the session", + ] + lines += ["", "## Latency", "", "| Outcome | p50 (s) |", "|---|---:|"] + order = ["ok", *(row.error_class for row in r.classes)] + for name, value in sorted(lat.by_class_p50.items(), key=lambda kv: order.index(kv[0]) if kv[0] in order else 99): + lines.append(f"| {CLASS_LABELS.get(name, 'Succeeded' if name == 'ok' else name)} | {value} |") + g = r.guard + lines += [ + "", + "## Guard replay", + "", + (f"The schema cache built from this history holds {g.cached_datasets} datasets. The guard could check " + f"columns for {g.checked:,} queries (every dataset cached at the time)."), + "", + "- Failures it would have stopped before execution: " + + (", ".join(f"{CLASS_LABELS.get(k, k)} {v}" for k, v in g.caught.items()) or "none") + + ". Outside the column and type classes, these queries failed for another reason first but " + "also used a column the dataset lacks or did date arithmetic on a DATE column.", + (f"- Successful queries it would have blocked by mistake: {g.false_blocks} of {g.successes_checked:,} " + "checked. These are mostly datasets rebuilt with new columns after their schema was cached; " + "live, the agent resends the query once and the hook drops the stale entry."), + ] + if g.caught_examples: + lines += ["", "Sample guard replies (redacted):", ""] + lines += [f"- {e}" for e in g.caught_examples[:3]] + lines += ["", "## Findings", ""] + for i, f in enumerate(r.findings, 1): + lines += [f"### {i}. [{f.fix} fix] {f.title} ({f.impact} calls)", "", f.detail, ""] + lines += [f"- {e}" for e in f.examples] + if f.examples: + lines.append("") + lines += [f"**Fix:** {f.suggestion}", ""] + if r.hints: + lines += ["## Suggested prompt hints", "", "```yaml", "hints:"] + lines += [f" - {json.dumps(h)}" for h in r.hints] + lines += ["```", ""] + if r.descriptions: + lines += ["## Suggested column descriptions", "", "```yaml", "descriptions:"] + lines += [f" {json.dumps(k)}: {json.dumps(v)}" for k, v in sorted(r.descriptions.items())] + lines += ["```", ""] + return "\n".join(lines).rstrip() + "\n" + + +def render_terminal(report: DoctorReport) -> str: + r = report + lat = r.latency + out = [ + f"{r.calls:,} query calls, {r.sessions} sessions; {r.failed:,} failed ({r.failure_rate:.1%})", + ] + if lat.p50 is not None: + out.append(f"latency p50 {lat.p50:.1f}s p90 {lat.p90:.1f}s >{lat.stall_seconds:.0f}s: {lat.stalls}") + out.append("") + for row in r.classes: + out.append(f" {row.count:>5} {row.share_of_calls:>6.1%} {row.label:<44} fix: {row.fix}") + c = r.chains + out += [ + "", + (f"retry chains: {c.chains} ({c.recovered} recovered, {c.abandoned} abandoned), " + f"{c.failed_attempts} failed attempts, {c.wasted_seconds / 60:.0f} min"), + (f"guard replay: would stop {sum(r.guard.caught.values())} failures, " + f"{r.guard.false_blocks} false blocks of {r.guard.successes_checked:,} checked successes"), + "", + "findings:", + ] + out += [f" [{f.fix}] {f.title} ({f.impact})" for f in r.findings] + return "\n".join(out) diff --git a/src/querypilot/trace/guard.py b/src/querypilot/trace/guard.py new file mode 100644 index 0000000..3704918 --- /dev/null +++ b/src/querypilot/trace/guard.py @@ -0,0 +1,288 @@ +"""Check an agent's SQL before the query tool runs it. + +The guard validates column and table references with ``SQLValidator`` against +the schema cache and returns "did you mean" feedback, so the agent fixes a bad +column without a round trip to the warehouse. It only judges queries whose +datasets are all cached; anything it cannot parse or resolve is allowed +(fail open), because a guard that blocks valid SQL costs more than it saves. +It can also require a row cap (a LIMIT, or the tool's own row-limit input). +""" +from __future__ import annotations + +import re +from dataclasses import dataclass, field + +import sqlglot +from sqlglot import exp + +from querypilot.core.config import QueryPilotConfig, SafetyPolicy +from querypilot.core.types import ColumnSchema, DatabaseSchema, TableSchema +from querypilot.trace.classify import GUARD_PREFIX +from querypilot.trace.schema_cache import ( + FOUNDRY_RID, + CachedDataset, + SchemaCache, + table_ref_name, + table_refs, +) +from querypilot.validation.feedback import repair_hints +from querypilot.validation.validator import SQLValidator + +ALLOW = "allow" +DENY = "deny" + +_RESOLUTION_ERRORS = ("Unknown column:", "Ambiguous column:", "Unknown table qualifier:") + + +@dataclass +class GuardDecision: + action: str = ALLOW + reason: str | None = None + # "schema" (unknown column/table), "limit" (no row cap), "schema_required". + rule: str | None = None + # SQL with a LIMIT added, for callers that rewrite instead of denying. + limited_sql: str | None = None + # True when every referenced dataset was in the cache, so columns were checked. + checked_columns: bool = False + tables: list[str] = field(default_factory=list) + + @property + def denied(self) -> bool: + return self.action == DENY + + +def guard_sql( + sql: str, + cache: SchemaCache, + *, + dialect: str = "spark", + max_rows: int | None = None, + tool_row_limit: int | None = None, + require_schema: bool = False, + schema_tool_hint: str = "the schema tool", +) -> GuardDecision: + stripped = sql.strip().rstrip(";").strip() + try: + expression = sqlglot.parse_one(stripped, read=dialect) + except (sqlglot.errors.SqlglotError, ValueError): + # The engine's parser is the authority on syntax; ours may lag its dialect. + return GuardDecision() + if expression is None: + return GuardDecision() + + refs = table_refs(expression) + names = list(dict.fromkeys(table_ref_name(ref) for ref in refs)) + decision = GuardDecision(tables=names) + + # Transaction-pinned refs (ri.foundry.main.transaction...@branch.ri...) are valid; only + # check names that claim to be a bare dataset RID. + malformed = [ + name for name in names if name.startswith("ri.foundry.main.dataset.") and not FOUNDRY_RID.fullmatch(name) + ] + if malformed: + decision.action, decision.rule = DENY, "schema" + decision.reason = _message( + [f"`{name}` is not a valid dataset RID (expected ri.foundry.main.dataset.)." for name in malformed] + + _close_rids(malformed, cache) + ) + return decision + + datasets = {name: cache.lookup(name) for name in names} + missing = [name for name, dataset in datasets.items() if dataset is None] + if missing and require_schema: + decision.action, decision.rule = DENY, "schema_required" + decision.reason = _message( + [ + ("Fetch the schema before querying: call " + f"{schema_tool_hint} for {', '.join(f'`{name}`' for name in missing)}, " + "then write the query with the real column names.") + ] + ) + return decision + + if names and not missing: + decision.checked_columns = True + problems = _column_problems(expression, refs, datasets, dialect) or _epoch_misuse(expression, datasets) + if problems: + decision.action, decision.rule = DENY, "schema" + decision.reason = _message(problems) + return decision + + if max_rows is not None and not _row_capped(expression, max_rows, tool_row_limit): + decision.limited_sql = _with_limit(stripped, expression, max_rows, dialect) + decision.action, decision.rule = DENY, "limit" + decision.reason = _message( + [ + (f"Queries must return at most {max_rows} rows: add `LIMIT {max_rows}` " + "(or a smaller one), or aggregate. Large results overflow the context window.") + ] + ) + return decision + + +def _column_problems( + expression: exp.Expression, + refs: list[exp.Table], + datasets: dict[str, CachedDataset | None], + dialect: str, +) -> list[str]: + aliases = _aliases(datasets) + renamed = expression.copy() + for table in table_refs(renamed): + alias = aliases[table_ref_name(table)] + table.set("this", exp.to_identifier(alias)) + table.set("db", None) + table.set("catalog", None) + schema = DatabaseSchema( + dialect=dialect, + tables=[ + TableSchema( + name=aliases[name], + columns=[ColumnSchema(name=c.name, type=c.type, comment=c.description or None) for c in ds.columns], + ) + for name, ds in datasets.items() + if ds is not None + ], + ) + config = QueryPilotConfig( + dialect=dialect, + max_rows=1_000_000, + safety_policy=SafetyPolicy(reject_cartesian_joins=False, blocked_functions=[], warn_on_select_star=False), + ) + renamed_sql = renamed.sql(dialect=dialect) + try: + result = SQLValidator(config).validate(renamed_sql, schema) + except Exception: # noqa: BLE001 - fail open on validator crashes + return [] + # Spark resolves `struct_col.field` and backticked names containing dots; the + # validator reads both as table.column, so those references are not judged. + value_aliases = {alias.alias.lower() for alias in expression.find_all(exp.Alias) if alias.alias} + dotted = any("." in ident.name for col in expression.find_all(exp.Column) for ident in col.parts) + if dotted: + return [] + errors = [ + error + for error in result.errors + if error.startswith(_RESOLUTION_ERRORS) + and not ( + error.startswith("Unknown table qualifier:") + and error.split(":", 1)[1].strip().lower() in value_aliases | _column_names(datasets) + ) + ] + if not errors: + return [] + hints = repair_hints(errors, schema, renamed_sql) or errors + legend = [ + f"{aliases[name]} = `{name}`" + (f" ({ds.path})" if ds and ds.path and ds.path != name else "") + for name, ds in datasets.items() + if aliases[name] != name + ] + if legend: + hints.append("Dataset names above: " + "; ".join(legend) + ".") + return hints + + +def _column_names(datasets: dict[str, CachedDataset | None]) -> set[str]: + return {column.name.lower() for dataset in datasets.values() if dataset for column in dataset.columns} + + +_DATE_TYPES = ("DATE", "TIMESTAMP") +_ARITHMETIC_OR_COMPARISON = (exp.Div, exp.Mul, exp.Add, exp.Sub, exp.GT, exp.GTE, exp.LT, exp.LTE, exp.EQ, exp.NEQ) + + +def _epoch_misuse(expression: exp.Expression, datasets: dict[str, CachedDataset | None]) -> list[str]: + """DATE/TIMESTAMP columns used as epoch milliseconds (``Report_Date / 1000``).""" + types: dict[str, str] = {} + for dataset in datasets.values(): + for column in dataset.columns if dataset else []: + types.setdefault(column.name.lower(), column.type.upper()) + problems: list[str] = [] + for node in expression.find_all(*_ARITHMETIC_OR_COMPARISON): + left, right = node.left, node.right + for column, other in ((left, right), (right, left)): + if not isinstance(column, exp.Column): + continue + column_type = types.get(column.name.lower(), "") + if not column_type.startswith(_DATE_TYPES) or not _epoch_like(other): + continue + message = ( + f"`{column.name}` is a {column_type} column, not epoch milliseconds: " + f"`{node.sql()}` fails with a type mismatch. Compare it with a literal like " + f"DATE '2026-01-01' or use date functions (year(), date_trunc(), datediff())." + ) + if message not in problems: + problems.append(message) + return problems + + +def _epoch_like(node: exp.Expression) -> bool: + numbers = [node] if isinstance(node, exp.Literal) else list(node.find_all(exp.Literal)) + for literal in numbers: + if not literal.is_number: + continue + try: + value = abs(float(literal.this)) + except ValueError: + continue + # 1000 (ms <-> s), 86400000 (ms per day), or a 13-digit epoch-ms timestamp. + if value in (1000, 86_400_000) or value >= 1e11: + return True + return False + + +def _aliases(datasets: dict[str, CachedDataset | None]) -> dict[str, str]: + out: dict[str, str] = {} + used: set[str] = set() + for name, dataset in datasets.items(): + base = dataset.display_name if dataset else name + alias = re.sub(r"\W+", "_", base).strip("_") or "dataset" + if alias[0].isdigit(): + alias = f"d_{alias}" + candidate, n = alias, 2 + while candidate.lower() in used: + candidate, n = f"{alias}_{n}", n + 1 + used.add(candidate.lower()) + out[name] = candidate + return out + + +def _close_rids(names: list[str], cache: SchemaCache) -> list[str]: + import difflib + + known = [key for key in cache.datasets if key.startswith("ri.")] + out = [] + for name in names: + close = difflib.get_close_matches(name, known, n=2, cutoff=0.8) + if close: + out.append(f"Did you mean {', '.join(f'`{c}`' for c in close)}?") + return out + + +def _row_capped(expression: exp.Expression, max_rows: int, tool_row_limit: int | None) -> bool: + if tool_row_limit is not None and 0 < tool_row_limit <= max_rows: + return True + limit = expression.args.get("limit") + value = limit.args.get("expression") if isinstance(limit, exp.Limit) else None + if isinstance(value, exp.Literal) and value.is_number: + try: + return 0 <= int(value.this) <= max_rows + except ValueError: + return False + # An aggregate without GROUP BY returns one row. + if isinstance(expression, exp.Select) and not expression.args.get("group"): + projections = expression.expressions + return bool(projections) and all(p.find(exp.AggFunc) is not None for p in projections) + return False + + +def _with_limit(stripped: str, expression: exp.Expression, max_rows: int, dialect: str) -> str: + if expression.args.get("limit") is None and "--" not in stripped: + return f"{stripped}\nLIMIT {max_rows}" + capped = expression.copy() + capped.set("limit", exp.Limit(expression=exp.Literal.number(max_rows))) + return capped.sql(dialect=dialect) + + +def _message(lines: list[str]) -> str: + body = "\n".join(f"- {line}" for line in lines) + return f"{GUARD_PREFIX} blocked this query before it ran:\n{body}\nFix the SQL and call the tool again." diff --git a/src/querypilot/trace/hooks.py b/src/querypilot/trace/hooks.py new file mode 100644 index 0000000..93d71c8 --- /dev/null +++ b/src/querypilot/trace/hooks.py @@ -0,0 +1,417 @@ +"""Claude Code hooks that log query-tool calls live and, opt-in, guard them. + +``querypilot hooks install`` adds PostToolUse and PostToolUseFailure hooks (and +a PreToolUse hook with ``--guard``) to a Claude Code settings file. Each hook +runs ``python -m querypilot hooks run ``, which reads the hook payload +on stdin and: + +- **post**: appends one audit record per query call (SQL, outcome, failure + class, latency, result size; never result rows) to the JSONL audit log, and + adds schema-tool results to the schema cache; +- **pre** (guard): validates the SQL against the schema cache and denies the + call with "did you mean" feedback, or when it has no row cap. A query that + was denied for its columns goes through if the agent sends it again + unchanged, so a stale cache entry costs one turn, not the task. + +A hook must never break the tool call: every error is swallowed (and logged to +``hook-errors.log`` next to the audit log), and the hook prints nothing. +""" +from __future__ import annotations + +import hashlib +import json +import os +import re +import shlex +import sys +import tempfile +import time +import traceback +from dataclasses import dataclass, field +from datetime import UTC, datetime +from pathlib import Path + +from querypilot.audit.sinks import JSONLAuditSink +from querypilot.trace.classify import BLOCKED, OK, markdown_header +from querypilot.trace.schema_cache import SchemaCache, querypilot_home +from querypilot.trace.tools import ToolRegistry, parse_tool_spec +from querypilot.trace.transcript import ( + QUERY, + SCHEMA, + ToolCall, + result_text, + small_value, + to_audit_record, +) + +HOOK_MARKER = "querypilot hooks run" +SCOPES = ("user", "project", "local") +_STATE_TTL_SECONDS = 6 * 3600 + + +def default_audit_log() -> Path: + return querypilot_home() / "trace" / "audit.jsonl" + + +def settings_path(scope: str, project_dir: str | Path | None = None) -> Path: + if scope == "user": + return Path.home() / ".claude" / "settings.json" + base = Path(project_dir or os.getcwd()) + return base / ".claude" / ("settings.json" if scope == "project" else "settings.local.json") + + +@dataclass +class HookOptions: + audit_log: Path = field(default_factory=default_audit_log) + cache_path: Path | None = None + guard: bool = False + max_rows: int | None = None + require_schema: bool = False + tools: list[str] = field(default_factory=list) + + @property + def registry(self) -> ToolRegistry: + return ToolRegistry.from_args(self.tools) + + def cli_args(self) -> list[str]: + args = ["--audit-log", str(self.audit_log)] + if self.cache_path: + args += ["--cache", str(self.cache_path)] + if self.guard: + args.append("--guard") + if self.max_rows is not None: + args += ["--max-rows", str(self.max_rows)] + if self.require_schema: + args.append("--require-schema") + for tool in self.tools: + args += ["--tool", tool] + return args + + @property + def state_path(self) -> Path: + return self.audit_log.parent / "hook-state.json" + + +# --------------------------------------------------------------------------- install + + +def hook_command(event: str, options: HookOptions, python: str | None = None) -> str: + parts = [python or sys.executable, "-m", "querypilot", "hooks", "run", event, *options.cli_args()] + return " ".join(shlex.quote(part) for part in parts) + + +def tool_matcher(names: list[str]) -> str: + """Claude Code treats a matcher of [A-Za-z0-9_|-] as an exact list, anything else as a regex.""" + if all(re.fullmatch(r"[\w-]+", name) for name in names): + return "|".join(names) + alternatives = [re.sub(r"([.+?^${}()\[\]\\])", r"\\\1", name).replace("*", ".*") for name in names] + return "^(?:" + "|".join(alternatives) + ")$" + + +def install_hooks(path: Path, options: HookOptions, *, python: str | None = None) -> dict: + settings = _read_settings(path) + hooks = _remove_ours(settings.get("hooks") or {}) + registry = options.registry + all_tools = tool_matcher(registry.tool_names()) + query_tools = tool_matcher([spec.name for spec in registry.query_tools]) + events = [("PostToolUse", "post", all_tools), ("PostToolUseFailure", "post-failure", all_tools)] + if options.guard: + events.insert(0, ("PreToolUse", "pre", query_tools)) + for event, name, matcher in events: + hooks.setdefault(event, []).append( + { + "matcher": matcher, + "hooks": [{"type": "command", "command": hook_command(name, options, python), "timeout": 30}], + } + ) + settings["hooks"] = hooks + _write_settings(path, settings) + return settings + + +def uninstall_hooks(path: Path) -> int: + settings = _read_settings(path) + before = json.dumps(settings.get("hooks") or {}) + hooks = _remove_ours(settings.get("hooks") or {}) + removed = before.count(HOOK_MARKER) - json.dumps(hooks).count(HOOK_MARKER) + if hooks: + settings["hooks"] = hooks + else: + settings.pop("hooks", None) + if removed: + _write_settings(path, settings) + return removed + + +def _remove_ours(hooks: dict) -> dict: + cleaned: dict = {} + for event, groups in hooks.items(): + kept_groups = [] + for group in groups if isinstance(groups, list) else []: + entries = [h for h in group.get("hooks", []) if HOOK_MARKER not in str(h.get("command", ""))] + if entries: + kept_groups.append({**group, "hooks": entries}) + if kept_groups: + cleaned[event] = kept_groups + return cleaned + + +def _read_settings(path: Path) -> dict: + if not path.is_file(): + return {} + text = path.read_text(encoding="utf-8").strip() + if not text: + return {} + data = json.loads(text) + if not isinstance(data, dict): + raise TypeError(f"{path} does not hold a JSON object") + return data + + +def _write_settings(path: Path, settings: dict) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(settings, indent=2) + "\n", encoding="utf-8") + + +# --------------------------------------------------------------------------- run + + +def run_hook(event: str, payload: dict, options: HookOptions) -> dict | None: + """Handle one hook event; return the JSON to print, or None for no output.""" + tool_name = str(payload.get("tool_name") or "") + registry = options.registry + spec = registry.query_spec(tool_name) + is_schema = spec is None and registry.schema_spec(tool_name) is not None + if spec is None and not is_schema: + return None + tool_input = payload.get("tool_input") if isinstance(payload.get("tool_input"), dict) else {} + call_id = str(payload.get("tool_use_id") or "") + state = _State.load(options.state_path) + + if event == "pre": + if spec is None or not options.guard: + return None + state.starts[call_id] = time.time() + output = _pre(payload, tool_input, spec, options, state) + state.save(options.state_path) + return output + + call = ToolCall( + call_id=call_id or f"hook-{time.time_ns()}", + tool=tool_name, + kind=QUERY if spec else SCHEMA, + session_id=payload.get("session_id"), + thread=_thread(payload), + sql=tool_input.get(spec.sql_field) if spec and isinstance(tool_input.get(spec.sql_field), str) else None, + dialect=spec.dialect if spec else "spark", + params={ + k: v for k, v in tool_input.items() if (spec is None or k != spec.sql_field) and small_value(v) + }, + ) + now = datetime.now(UTC) + started = state.starts.pop(call_id, None) + duration = payload.get("duration_ms") + if isinstance(duration, int | float): + call.duration_ms = int(duration) + elif started: + call.duration_ms = round((time.time() - started) * 1000) + call.started_at = datetime.fromtimestamp(started, UTC) if started else now + if event == "post-failure": + text, is_error = str(payload.get("error") or "tool call failed"), True + else: + response = payload.get("tool_response") + text = result_text(response.get("content") if isinstance(response, dict) and "content" in response else response) + is_error = bool(isinstance(response, dict) and response.get("isError")) + call.apply_result(text, is_error=is_error, ended_at=now) + + record = to_audit_record(call) + record.details["source"] = "hook" + record.details["duration_ms"] = call.duration_ms + JSONLAuditSink(options.audit_log).write(record) + + output = _learn(call, tool_input, text, options) if (options.guard or is_schema) else None + if state.dirty: + state.save(options.state_path) + return output + + +def _pre(payload: dict, tool_input: dict, spec, options: HookOptions, state: _State) -> dict | None: + from querypilot.trace.guard import guard_sql + + sql = tool_input.get(spec.sql_field) + if not isinstance(sql, str) or not sql.strip(): + return None + row_limit = tool_input.get(spec.row_limit_field) if spec.row_limit_field else None + decision = guard_sql( + sql, + SchemaCache.load(options.cache_path), + dialect=spec.dialect, + max_rows=options.max_rows, + tool_row_limit=int(row_limit) if isinstance(row_limit, int | float | str) and str(row_limit).isdigit() else None, + require_schema=options.require_schema, + schema_tool_hint="the dataset schema tool", + ) + if not decision.denied: + return None + session = str(payload.get("session_id") or "") + fingerprint = hashlib.sha256(f"{session}\n{sql.strip()}".encode()).hexdigest()[:16] + # The agent resent a query denied for its columns: trust it, the cache may be stale. + if decision.rule == "schema" and fingerprint in state.denied: + state.denied.pop(fingerprint, None) + return None + state.denied[fingerprint] = time.time() + reason = decision.reason or "Blocked by the QueryPilot guard." + if decision.rule == "schema": + reason += " If you are sure the query is right (the dataset may have changed), send it again unchanged." + + call = ToolCall( + call_id=str(payload.get("tool_use_id") or f"hook-{time.time_ns()}"), + tool=str(payload.get("tool_name")), + kind=QUERY, + session_id=payload.get("session_id"), + thread=_thread(payload), + started_at=datetime.now(UTC), + sql=sql, + dialect=spec.dialect, + output=reason, + status=BLOCKED, + ) + record = to_audit_record(call) + record.valid = False + record.details.update({"source": "hook", "guard_rule": decision.rule}) + JSONLAuditSink(options.audit_log).write(record) + return { + "hookSpecificOutput": { + "hookEventName": "PreToolUse", + "permissionDecision": "deny", + "permissionDecisionReason": reason, + } + } + + +def _learn(call: ToolCall, tool_input: dict, text: str, options: HookOptions) -> dict | None: + from querypilot.trace.guard import guard_sql + + cache = SchemaCache.load(options.cache_path) + changed = False + if call.kind == SCHEMA: + changed = cache.learn_from_schema_result(text, tool_input) is not None + elif call.status == OK and call.sql: + # A query that succeeded although the cache says it should fail proves the entry stale. + decision = guard_sql(call.sql, cache, dialect=call.dialect) + if decision.denied and decision.rule == "schema": + changed = bool(cache.evict(decision.tables)) + header = markdown_header(text[:20000]) + if header and cache.learn_from_select_star(call.sql, header, call.dialect): + changed = True + if changed: + cache.save(options.cache_path) + note = _epoch_note(call, text, cache) if call.kind == QUERY and call.status == OK else None + if note: + return {"hookSpecificOutput": {"hookEventName": "PostToolUse", "additionalContext": note}} + return None + + +def _epoch_note(call: ToolCall, text: str, cache: SchemaCache) -> str | None: + """Foundry renders DATE/TIMESTAMP values as epoch ms; say so before the agent does arithmetic on them.""" + from querypilot.trace.doctor import query_tables + + header = call.header or [] + if not header: + return None + types = {} + for ref in query_tables(call.sql, call.dialect): + dataset = cache.lookup(ref) + for column in dataset.columns if dataset else []: + types[column.name.lower()] = column.type.upper() + rows = [line for line in text.splitlines() if line.startswith("|")][2:5] + flagged = [] + for i, name in enumerate(header): + column_type = types.get(name.lower(), "") + if not column_type.startswith(("DATE", "TIMESTAMP")): + continue + for row in rows: + cells = [cell.strip() for cell in row.strip().strip("|").split("|")] + if i < len(cells) and re.fullmatch(r"-?\d{12,14}", cells[i]): + flagged.append(f"{name} ({column_type})") + break + if not flagged: + return None + return ( + f"QueryPilot: {', '.join(flagged)} are date columns displayed as epoch milliseconds. " + "In SQL they are dates: filter with DATE 'YYYY-MM-DD' literals or date functions, " + "never divide by 1000 or compare with 13-digit numbers." + ) + + +def _thread(payload: dict) -> str | None: + transcript = payload.get("transcript_path") + agent = payload.get("agent_id") + if agent: + return f"{transcript}#{agent}" + return transcript + + +class _State: + """Small per-machine state: call start times and recently denied queries.""" + + def __init__(self, starts: dict | None = None, denied: dict | None = None) -> None: + self.starts: dict[str, float] = starts or {} + self.denied: dict[str, float] = denied or {} + self._snapshot = json.dumps([self.starts, self.denied], sort_keys=True) + + @property + def dirty(self) -> bool: + return json.dumps([self.starts, self.denied], sort_keys=True) != self._snapshot + + @classmethod + def load(cls, path: Path) -> _State: + try: + raw = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return cls() + cutoff = time.time() - _STATE_TTL_SECONDS + fresh = lambda d: {k: v for k, v in (d or {}).items() if isinstance(v, int | float) and v > cutoff} + return cls(fresh(raw.get("starts")), fresh(raw.get("denied"))) + + def save(self, path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp = tempfile.mkstemp(dir=path.parent, prefix=".hook-state-", suffix=".json") + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump({"starts": self.starts, "denied": self.denied}, handle) + os.replace(tmp, path) + + +def main_hook(event: str, options: HookOptions, stdin=None, stdout=None) -> int: + stdin = stdin or sys.stdin + stdout = stdout or sys.stdout + try: + import logging + + logging.getLogger("sqlglot").setLevel(logging.ERROR) + payload = json.loads(stdin.read() or "{}") + output = run_hook(event, payload, options) if isinstance(payload, dict) else None + if output: + stdout.write(json.dumps(output)) + except Exception: # noqa: BLE001 - a logging hook must never fail the tool call + try: + log = options.audit_log.parent / "hook-errors.log" + log.parent.mkdir(parents=True, exist_ok=True) + with log.open("a", encoding="utf-8") as handle: + handle.write(f"{datetime.now(UTC).isoformat()} {event}\n{traceback.format_exc()}\n") + except OSError: + pass + return 0 + + +def options_from_args(args) -> HookOptions: + for value in args.tool or []: + parse_tool_spec(value) # fail fast on a malformed --tool + return HookOptions( + audit_log=Path(args.audit_log).expanduser() if args.audit_log else default_audit_log(), + cache_path=Path(args.cache).expanduser() if args.cache else None, + guard=bool(getattr(args, "guard", False)), + max_rows=getattr(args, "max_rows", None), + require_schema=bool(getattr(args, "require_schema", False)), + tools=list(args.tool or []), + ) diff --git a/src/querypilot/trace/schema_cache.py b/src/querypilot/trace/schema_cache.py new file mode 100644 index 0000000..68b0cd4 --- /dev/null +++ b/src/querypilot/trace/schema_cache.py @@ -0,0 +1,232 @@ +"""A local cache of dataset schemas, learned from the agent's own tool calls. + +Entries come from schema tool results (``get_foundry_dataset_schema``) and from +successful ``SELECT *`` queries on a single dataset, whose result header lists +every column. The guard validates SQL against this cache, so it only judges +queries whose datasets it has seen. +""" +from __future__ import annotations + +import json +import os +import re +import tempfile +from dataclasses import asdict, dataclass, field +from datetime import UTC, datetime +from pathlib import Path + +import sqlglot +from sqlglot import exp + +FOUNDRY_RID = re.compile(r"ri\.foundry\.main\.dataset\.[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}") + +SOURCE_SCHEMA_TOOL = "schema_tool" +SOURCE_SELECT_STAR = "select_star" + + +def querypilot_home() -> Path: + return Path(os.getenv("QUERYPILOT_HOME") or Path.home() / ".querypilot") + + +def default_cache_path() -> Path: + return querypilot_home() / "trace" / "schema-cache.json" + + +@dataclass +class CachedColumn: + name: str + type: str = "UNKNOWN" + description: str = "" + + +@dataclass +class CachedDataset: + key: str + path: str | None = None + columns: list[CachedColumn] = field(default_factory=list) + source: str = SOURCE_SCHEMA_TOOL + updated_at: str = field(default_factory=lambda: datetime.now(UTC).isoformat()) + + @property + def display_name(self) -> str: + if self.path: + return self.path.rstrip("/").rsplit("/", 1)[-1] or self.key + return self.key + + +class SchemaCache: + def __init__(self, datasets: dict[str, CachedDataset] | None = None) -> None: + self.datasets: dict[str, CachedDataset] = dict(datasets or {}) + self._paths = {d.path.lower(): d.key for d in self.datasets.values() if d.path} + + def __len__(self) -> int: + return len(self.datasets) + + def lookup(self, ref: str) -> CachedDataset | None: + ref = ref.strip().strip("`\"") + hit = self.datasets.get(ref) + if hit is None and ref.lower() in self._paths: + hit = self.datasets.get(self._paths[ref.lower()]) + return hit + + def add(self, dataset: CachedDataset) -> bool: + """Store a dataset; a schema-tool entry is never replaced by a SELECT * one.""" + current = self.datasets.get(dataset.key) + if ( + current is not None + and current.source == SOURCE_SCHEMA_TOOL + and dataset.source != SOURCE_SCHEMA_TOOL + ): + return False + if current is not None and dataset.path is None: + dataset.path = current.path + self.datasets[dataset.key] = dataset + if dataset.path: + self._paths[dataset.path.lower()] = dataset.key + return True + + def evict(self, refs: list[str]) -> list[str]: + """Forget datasets whose cached schema proved wrong (e.g. rebuilt with new columns).""" + evicted = [] + for ref in refs: + dataset = self.lookup(ref) + if dataset is None: + continue + self.datasets.pop(dataset.key, None) + if dataset.path: + self._paths.pop(dataset.path.lower(), None) + evicted.append(dataset.key) + return evicted + + def learn_from_schema_result(self, text: str, tool_input: dict | None = None) -> CachedDataset | None: + dataset = parse_foundry_schema(text) + if dataset is None: + return None + if dataset.path is None and tool_input and tool_input.get("datasetPath"): + dataset.path = str(tool_input["datasetPath"]) + self.add(dataset) + return dataset + + def learn_from_select_star(self, sql: str, header: list[str] | None, dialect: str) -> CachedDataset | None: + """``SELECT * FROM one_dataset`` returns every column; record their names.""" + if not header: + return None + ref = select_star_source(sql, dialect) + if ref is None: + return None + current = self.lookup(ref) + dataset = CachedDataset( + key=current.key if current else ref, + path=current.path if current else (ref if ref.startswith("/") else None), + columns=[CachedColumn(name) for name in header if name], + source=SOURCE_SELECT_STAR, + ) + return dataset if self.add(dataset) else None + + # ------------------------------------------------------------------ storage + + @classmethod + def load(cls, path: str | Path | None = None) -> SchemaCache: + target = Path(path) if path else default_cache_path() + if not target.is_file(): + return cls() + try: + payload = json.loads(target.read_text(encoding="utf-8")) + except (OSError, ValueError): + return cls() + datasets = {} + for key, raw in (payload.get("datasets") or {}).items(): + columns = [CachedColumn(**column) for column in raw.get("columns", [])] + datasets[key] = CachedDataset( + key=key, + path=raw.get("path"), + columns=columns, + source=raw.get("source", SOURCE_SCHEMA_TOOL), + updated_at=raw.get("updated_at", ""), + ) + return cls(datasets) + + def save(self, path: str | Path | None = None) -> Path: + target = Path(path) if path else default_cache_path() + target.parent.mkdir(parents=True, exist_ok=True) + payload = { + "version": 1, + "datasets": { + key: {k: v for k, v in asdict(dataset).items() if k != "key"} + for key, dataset in sorted(self.datasets.items()) + }, + } + # Hooks for parallel tool calls can save at once; replace atomically. + fd, tmp = tempfile.mkstemp(dir=target.parent, prefix=".schema-cache-", suffix=".json") + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=1) + os.replace(tmp, target) + return target + + +def parse_foundry_schema(text: str) -> CachedDataset | None: + """Parse a ``get_foundry_dataset_schema`` result (a JSON resource tree).""" + start = text.find("{") + if start == -1: + return None + try: + payload = json.loads(text[start:]) + except ValueError: + return None + if not isinstance(payload, dict): + return None + properties = payload.get("properties") or {} + rid = properties.get("rid") + if not rid: + return None + columns: list[CachedColumn] = [] + for child in payload.get("children") or []: + if not isinstance(child, dict) or child.get("type") != "schema": + continue + for column in child.get("children") or []: + props = column.get("properties") if isinstance(column, dict) else None + if props and props.get("name"): + columns.append( + CachedColumn( + name=str(props["name"]), + type=str(props.get("type") or "UNKNOWN"), + description=str(props.get("description") or ""), + ) + ) + if not columns: + return None + return CachedDataset(key=str(rid), path=properties.get("path"), columns=columns) + + +def table_refs(expression: exp.Expression) -> list[exp.Table]: + """Physical table references, excluding CTE names.""" + ctes = {cte.alias_or_name.lower() for cte in expression.find_all(exp.CTE)} + refs = [] + for table in expression.find_all(exp.Table): + if not table.name: + continue + if not table.args.get("db") and table.name.lower() in ctes: + continue + refs.append(table) + return refs + + +def table_ref_name(table: exp.Table) -> str: + parts = [part.name for part in table.parts] + return ".".join(parts) + + +def select_star_source(sql: str, dialect: str) -> str | None: + try: + expression = sqlglot.parse_one(sql, read=dialect) + except (sqlglot.errors.SqlglotError, ValueError): + return None + if not isinstance(expression, exp.Select) or expression.args.get("joins"): + return None + if len(expression.expressions) != 1 or not isinstance(expression.expressions[0], exp.Star): + return None + source = expression.args.get("from") or expression.args.get("from_") + table = source.this if source is not None else None + if not isinstance(table, exp.Table) or expression.find(exp.CTE) is not None: + return None + return table_ref_name(table) diff --git a/src/querypilot/trace/tools.py b/src/querypilot/trace/tools.py new file mode 100644 index 0000000..b56a5c4 --- /dev/null +++ b/src/querypilot/trace/tools.py @@ -0,0 +1,79 @@ +"""Which agent tools run SQL, and where their input keeps it. + +A query tool is mapped by its tool name (exact, or a glob such as +``mcp__warehouse__*``) and the input field that holds the SQL. Schema tools are +the ones whose results describe a dataset's columns; their results feed the +schema cache the guard validates against. +""" +from __future__ import annotations + +import fnmatch +from dataclasses import dataclass + +FOUNDRY_SQL_TOOL = "mcp__palantir-mcp__run_sql_query_on_foundry_dataset" +FOUNDRY_SCHEMA_TOOL = "mcp__palantir-mcp__get_foundry_dataset_schema" + + +@dataclass(frozen=True) +class QueryToolSpec: + name: str + sql_field: str + dialect: str = "spark" + # Input field that caps returned rows, when the tool has one. + row_limit_field: str | None = None + + def matches(self, tool_name: str) -> bool: + return fnmatch.fnmatchcase(tool_name, self.name) + + +@dataclass(frozen=True) +class SchemaToolSpec: + name: str + # Parser for the tool's result; "foundry" is the only built-in format. + format: str = "foundry" + + def matches(self, tool_name: str) -> bool: + return fnmatch.fnmatchcase(tool_name, self.name) + + +DEFAULT_QUERY_TOOLS: tuple[QueryToolSpec, ...] = ( + QueryToolSpec(FOUNDRY_SQL_TOOL, "sqlQuery", "spark", row_limit_field="rowLimit"), +) +DEFAULT_SCHEMA_TOOLS: tuple[SchemaToolSpec, ...] = (SchemaToolSpec(FOUNDRY_SCHEMA_TOOL),) + + +@dataclass(frozen=True) +class ToolRegistry: + query_tools: tuple[QueryToolSpec, ...] = DEFAULT_QUERY_TOOLS + schema_tools: tuple[SchemaToolSpec, ...] = DEFAULT_SCHEMA_TOOLS + + def query_spec(self, tool_name: str) -> QueryToolSpec | None: + return next((spec for spec in self.query_tools if spec.matches(tool_name)), None) + + def schema_spec(self, tool_name: str) -> SchemaToolSpec | None: + return next((spec for spec in self.schema_tools if spec.matches(tool_name)), None) + + def tool_names(self) -> list[str]: + return [spec.name for spec in (*self.query_tools, *self.schema_tools)] + + @classmethod + def from_args(cls, tools: list[str] | None, *, include_defaults: bool = True) -> ToolRegistry: + extra = tuple(parse_tool_spec(value) for value in tools or []) + base = DEFAULT_QUERY_TOOLS if include_defaults else () + return cls(query_tools=(*extra, *base), schema_tools=DEFAULT_SCHEMA_TOOLS) + + +def parse_tool_spec(value: str) -> QueryToolSpec: + """Parse ``NAME=SQL_FIELD[:DIALECT[:ROW_LIMIT_FIELD]]``. + + ``mcp__warehouse__run_query=query:snowflake`` maps the ``query`` input of + that tool as Snowflake SQL. + """ + name, sep, rest = value.partition("=") + if not sep or not name.strip() or not rest.strip(): + raise ValueError(f"Expected TOOL_NAME=SQL_FIELD[:DIALECT[:ROW_LIMIT_FIELD]], got {value!r}") + parts = rest.split(":") + sql_field = parts[0].strip() + dialect = parts[1].strip() if len(parts) > 1 and parts[1].strip() else "spark" + row_limit_field = parts[2].strip() if len(parts) > 2 and parts[2].strip() else None + return QueryToolSpec(name.strip(), sql_field, dialect, row_limit_field) diff --git a/src/querypilot/trace/transcript.py b/src/querypilot/trace/transcript.py new file mode 100644 index 0000000..0bd4df7 --- /dev/null +++ b/src/querypilot/trace/transcript.py @@ -0,0 +1,341 @@ +"""Read agent query-tool calls out of Claude Code transcripts and audit logs. + +Claude Code writes one JSONL transcript per session under +``~/.claude/projects//``. Each assistant ``tool_use`` block is answered +by a ``tool_result`` block with the same id in a later user record; pairing +them gives the SQL, the outcome and the latency (the gap between the two +records' timestamps, which includes any permission prompt). Resumed sessions +copy earlier history into a new file, so calls are de-duplicated by id. +""" +from __future__ import annotations + +import json +from collections.abc import Iterable, Iterator +from dataclasses import dataclass, field +from datetime import datetime +from pathlib import Path + +from querypilot.audit.types import QueryAuditRecord +from querypilot.trace.classify import ( + BLOCKED, + DENIED, + NO_RESULT, + OK, + classify_output, + markdown_header, +) +from querypilot.trace.schema_cache import CachedDataset, parse_foundry_schema +from querypilot.trace.tools import ToolRegistry + +QUERY = "query" +SCHEMA = "schema" +TOOL_CALL_OPERATION = "tool_call" + +_KEPT_OUTPUT = 2000 +_KEPT_QUESTION = 500 + + +@dataclass +class ToolCall: + call_id: str + tool: str + kind: str + session_id: str | None = None + # Transcript file (one per session or subagent); retry chains never cross it. + thread: str | None = None + sequence: int = 0 + started_at: datetime | None = None + ended_at: datetime | None = None + duration_ms: int | None = None + sql: str | None = None + dialect: str = "spark" + params: dict = field(default_factory=dict) + # Result or error text, truncated; held in memory only, never written by import. + output: str = "" + result_chars: int | None = None + is_error: bool = False + status: str = NO_RESULT + error_class: str | None = None + row_count: int | None = None + spilled: bool = False + header: list[str] | None = None + dataset: CachedDataset | None = None + question: str | None = None + + @property + def latency_s(self) -> float | None: + if self.duration_ms is not None: + return self.duration_ms / 1000 + if self.started_at and self.ended_at: + return max((self.ended_at - self.started_at).total_seconds(), 0.0) + return None + + @property + def ran(self) -> bool: + return self.status not in (DENIED, BLOCKED) + + def apply_result(self, text: str, *, is_error: bool, ended_at: datetime | None) -> None: + self.ended_at = ended_at + self.is_error = is_error + self.result_chars = len(text) + self.output = text[:_KEPT_OUTPUT] + if self.kind == SCHEMA: + self.dataset = None if is_error else parse_foundry_schema(text) + self.status = OK if self.dataset else ("error" if is_error else OK) + return + outcome = classify_output(text, is_error=is_error) + self.status, self.error_class = outcome.status, outcome.error_class + self.row_count, self.spilled = outcome.row_count, outcome.spilled + if outcome.status == OK: + self.header = markdown_header(text[:20000]) + + +def iter_transcript_files(paths: Iterable[str | Path]) -> Iterator[Path]: + for raw in paths: + path = Path(raw).expanduser() + if path.is_dir(): + yield from sorted(path.rglob("*.jsonl")) + elif path.is_file(): + yield path + + +def is_audit_log(path: Path) -> bool: + try: + with path.open(encoding="utf-8", errors="replace") as handle: + first = handle.readline() + return '"audit_id"' in first and '"operation"' in first + except OSError: + return False + + +def load_calls(paths: Iterable[str | Path], registry: ToolRegistry | None = None) -> list[ToolCall]: + """Tool calls from transcripts and/or QueryPilot audit logs, de-duplicated by id.""" + registry = registry or ToolRegistry() + calls: dict[str, ToolCall] = {} + for path in iter_transcript_files(paths): + found = read_audit_log(path, registry) if is_audit_log(path) else read_transcript(path, registry) + for call in found: + current = calls.get(call.call_id) + # Prefer the copy that saw the result; hooks measure latency exactly. + if current is None or (current.status == NO_RESULT and call.status != NO_RESULT): + calls[call.call_id] = call + elif call.duration_ms is not None and current.duration_ms is None: + current.duration_ms = call.duration_ms + ordered = sorted(calls.values(), key=lambda c: (c.started_at is None, c.started_at or datetime.min, c.sequence)) + return ordered + + +def read_transcript(path: Path, registry: ToolRegistry) -> list[ToolCall]: + pending: dict[str, ToolCall] = {} + out: list[ToolCall] = [] + question: str | None = None + sequence = 0 + try: + handle = path.open(encoding="utf-8", errors="replace") + except OSError: + return [] + with handle: + for line in handle: + if '"tool_use"' not in line and '"tool_result"' not in line and '"user"' not in line: + continue + try: + record = json.loads(line) + except ValueError: + continue + message = record.get("message") if isinstance(record, dict) else None + if not isinstance(message, dict): + continue + content = message.get("content") + timestamp = _parse_time(record.get("timestamp")) + # Skill bodies, reminders and compaction summaries arrive as user records too. + injected = bool(record.get("isMeta") or record.get("isCompactSummary")) + if message.get("role") == "user" and isinstance(content, str): + question = (None if injected else _prompt_text(content)) or question + continue + if not isinstance(content, list): + continue + for block in content: + if not isinstance(block, dict): + continue + kind = block.get("type") + if kind == "text" and message.get("role") == "user" and not injected: + question = _prompt_text(block.get("text")) or question + elif kind == "tool_use": + call = _call_from_tool_use(block, registry) + if call is None: + continue + sequence += 1 + call.sequence = sequence + call.session_id = record.get("sessionId") or path.stem + call.thread = str(path) + call.started_at = timestamp + call.question = question + pending[call.call_id] = call + out.append(call) + elif kind == "tool_result": + call = pending.pop(block.get("tool_use_id"), None) + if call is not None: + call.apply_result( + result_text(block.get("content")), + is_error=bool(block.get("is_error")), + ended_at=timestamp, + ) + return out + + +def read_audit_log(path: Path, registry: ToolRegistry) -> list[ToolCall]: + out = [] + with path.open(encoding="utf-8", errors="replace") as handle: + for line in handle: + line = line.strip() + if not line: + continue + try: + record = QueryAuditRecord.model_validate_json(line) + except ValueError: + continue + if record.operation == TOOL_CALL_OPERATION and record.tool_name and record.tool_call_id: + out.append(call_from_record(record, registry)) + return out + + +def _call_from_tool_use(block: dict, registry: ToolRegistry) -> ToolCall | None: + name = block.get("name") or "" + tool_input = block.get("input") if isinstance(block.get("input"), dict) else {} + spec = registry.query_spec(name) + if spec is not None: + sql = tool_input.get(spec.sql_field) + return ToolCall( + call_id=str(block.get("id")), + tool=name, + kind=QUERY, + sql=sql if isinstance(sql, str) else None, + dialect=spec.dialect, + params={k: v for k, v in tool_input.items() if k != spec.sql_field and small_value(v)}, + ) + if registry.schema_spec(name) is not None: + return ToolCall( + call_id=str(block.get("id")), + tool=name, + kind=SCHEMA, + params={k: v for k, v in tool_input.items() if small_value(v)}, + ) + return None + + +def result_text(content) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + return "\n".join( + str(item.get("text", "")) for item in content if isinstance(item, dict) and item.get("type") == "text" + ) + if isinstance(content, dict): + return json.dumps(content) + return "" if content is None else str(content) + + +def to_audit_record(call: ToolCall) -> QueryAuditRecord: + """The audit record `trace import` writes: SQL, outcome and timing, never result rows.""" + error = None + if call.status not in (OK, NO_RESULT): + error = call.output[:500] + return QueryAuditRecord( + timestamp=call.started_at or datetime.now().astimezone(), + operation=TOOL_CALL_OPERATION, + question=call.question, + sql=call.sql, + executed=call.status == OK, + row_count=call.row_count, + execution_time_ms=round(call.latency_s * 1000) if call.latency_s is not None else None, + error=error, + session_id=call.session_id, + tool_name=call.tool, + tool_call_id=call.call_id, + details={ + "kind": call.kind, + "status": call.status, + "error_class": call.error_class, + "thread": call.thread, + "sequence": call.sequence, + "ended_at": call.ended_at.isoformat() if call.ended_at else None, + "result_chars": call.result_chars, + "spilled": call.spilled, + "params": call.params, + "dialect": call.dialect, + "header": call.header, + "dataset": _dataset_payload(call.dataset), + }, + ) + + +def call_from_record(record: QueryAuditRecord, registry: ToolRegistry) -> ToolCall: + details = record.details or {} + spec = registry.query_spec(record.tool_name or "") + dataset = None + if details.get("dataset"): + from querypilot.trace.schema_cache import CachedColumn + + raw = details["dataset"] + dataset = CachedDataset( + key=raw["key"], path=raw.get("path"), columns=[CachedColumn(**c) for c in raw.get("columns", [])] + ) + call = ToolCall( + call_id=record.tool_call_id or record.audit_id, + tool=record.tool_name or "", + kind=details.get("kind") or (QUERY if spec else SCHEMA), + session_id=record.session_id, + thread=details.get("thread") or record.session_id, + sequence=int(details.get("sequence") or 0), + started_at=record.timestamp, + ended_at=_parse_time(details.get("ended_at")), + duration_ms=details.get("duration_ms"), + sql=record.sql, + dialect=details.get("dialect") or (spec.dialect if spec else "spark"), + params=details.get("params") or {}, + output=record.error or "", + result_chars=details.get("result_chars"), + status=details.get("status") or (OK if record.executed else "error"), + error_class=details.get("error_class"), + row_count=record.row_count, + spilled=bool(details.get("spilled")), + header=details.get("header"), + dataset=dataset, + question=record.question, + ) + if call.duration_ms is None and record.execution_time_ms is not None and call.ended_at is None: + call.duration_ms = record.execution_time_ms + return call + + +def _dataset_payload(dataset: CachedDataset | None) -> dict | None: + if dataset is None: + return None + return { + "key": dataset.key, + "path": dataset.path, + "columns": [{"name": c.name, "type": c.type, "description": c.description} for c in dataset.columns], + } + + +def _prompt_text(text) -> str | None: + if not isinstance(text, str): + return None + text = text.strip() + # Slash-command wrappers, hook output and system reminders are not the user's question. + if not text or text.startswith(("<", "[Request interrupted")): + return None + return text[:_KEPT_QUESTION] + + +def small_value(value) -> bool: + return isinstance(value, (int, float, bool)) or (isinstance(value, str) and len(value) <= 200) + + +def _parse_time(value) -> datetime | None: + if not isinstance(value, str) or not value: + return None + try: + return datetime.fromisoformat(value) + except ValueError: + return None diff --git a/tests/fixtures/claude_code_foundry_transcript.jsonl b/tests/fixtures/claude_code_foundry_transcript.jsonl new file mode 100644 index 0000000..41802d8 --- /dev/null +++ b/tests/fixtures/claude_code_foundry_transcript.jsonl @@ -0,0 +1,22 @@ +{"type": "user", "sessionId": "sess-1", "timestamp": "2026-10-01T10:00:00.000Z", "message": {"role": "user", "content": "Which products had the most monthly active users last month?"}} +{"type": "assistant", "sessionId": "sess-1", "timestamp": "2026-10-01T10:00:01.000Z", "message": {"role": "assistant", "content": [{"type": "text", "text": "Querying."}, {"type": "tool_use", "id": "toolu_01", "name": "Bash", "input": {"command": "ls"}}]}} +{"type": "user", "sessionId": "sess-1", "timestamp": "2026-10-01T10:00:02.000Z", "message": {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "is_error": false, "content": "ok"}]}} +{"type": "assistant", "sessionId": "sess-1", "timestamp": "2026-10-01T10:00:03.000Z", "message": {"role": "assistant", "content": [{"type": "text", "text": "Querying."}, {"type": "tool_use", "id": "toolu_02", "name": "mcp__palantir-mcp__run_sql_query_on_foundry_dataset", "input": {"sqlQuery": "SELECT Product, SUM(MAU) AS mau FROM `ri.foundry.main.dataset.11111111-2222-3333-4444-555555555555` GROUP BY Product", "rowLimit": 100}}]}} +{"type": "user", "sessionId": "sess-1", "timestamp": "2026-10-01T10:00:07.000Z", "message": {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_02", "is_error": false, "content": [{"type": "text", "text": "Sql Query Failed\nError Message: Failed to execute SQL query: [UNRESOLVED_COLUMN.WITH_SUGGESTION] A column or function parameter with name `Product` cannot be resolved. Did you mean one of the following? [`MAU`, `Product_Name`, `Report_Date`].; line 1 pos 7"}]}]}} +{"type": "assistant", "sessionId": "sess-1", "timestamp": "2026-10-01T10:00:08.000Z", "message": {"role": "assistant", "content": [{"type": "text", "text": "Querying."}, {"type": "tool_use", "id": "toolu_03", "name": "mcp__palantir-mcp__get_foundry_dataset_schema", "input": {"datasetRid": "ri.foundry.main.dataset.11111111-2222-3333-4444-555555555555"}}]}} +{"type": "user", "sessionId": "sess-1", "timestamp": "2026-10-01T10:00:09.000Z", "message": {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_03", "is_error": false, "content": [{"type": "text", "text": "{\"type\": \"dataset\", \"properties\": {\"rid\": \"ri.foundry.main.dataset.11111111-2222-3333-4444-555555555555\", \"path\": \"/Demo/Usage/monthly_usage\", \"branch\": \"master\"}, \"children\": [{\"type\": \"schema\", \"children\": [{\"type\": \"column\", \"properties\": {\"name\": \"Customer_Name\", \"type\": \"STRING\", \"nullable\": \"true\", \"description\": \"\"}}, {\"type\": \"column\", \"properties\": {\"name\": \"Product_Name\", \"type\": \"STRING\", \"nullable\": \"true\", \"description\": \"\"}}, {\"type\": \"column\", \"properties\": {\"name\": \"Report_Date\", \"type\": \"DATE\", \"nullable\": \"true\", \"description\": \"\"}}, {\"type\": \"column\", \"properties\": {\"name\": \"MAU\", \"type\": \"LONG\", \"nullable\": \"true\", \"description\": \"\"}}]}]}"}]}]}} +{"type": "assistant", "sessionId": "sess-1", "timestamp": "2026-10-01T10:00:10.000Z", "message": {"role": "assistant", "content": [{"type": "text", "text": "Querying."}, {"type": "tool_use", "id": "toolu_04", "name": "mcp__palantir-mcp__run_sql_query_on_foundry_dataset", "input": {"sqlQuery": "SELECT Product_Name, SUM(MAU) AS mau FROM `ri.foundry.main.dataset.11111111-2222-3333-4444-555555555555` WHERE Report_Date >= 1790812800000 GROUP BY Product_Name", "rowLimit": 100}}]}} +{"type": "user", "sessionId": "sess-1", "timestamp": "2026-10-01T10:00:13.000Z", "message": {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_04", "is_error": false, "content": [{"type": "text", "text": "Sql Query Failed\nError Message: Failed to execute SQL query: [DATATYPE_MISMATCH.BINARY_OP_DIFF_TYPES] Cannot resolve \"(Report_Date >= 1790812800000)\" due to data type mismatch: the left and right operands of the binary operator have incompatible types (\"DATE\" and \"BIGINT\").; line 1 pos 70"}]}]}} +{"type": "assistant", "sessionId": "sess-1", "timestamp": "2026-10-01T10:00:14.000Z", "message": {"role": "assistant", "content": [{"type": "text", "text": "Querying."}, {"type": "tool_use", "id": "toolu_05", "name": "mcp__palantir-mcp__run_sql_query_on_foundry_dataset", "input": {"sqlQuery": "SELECT Product_Name, SUM(MAU) AS mau FROM `ri.foundry.main.dataset.11111111-2222-3333-4444-555555555555` WHERE Report_Date >= DATE '2026-09-01' GROUP BY Product_Name ORDER BY mau DESC LIMIT 10", "rowLimit": 100}}]}} +{"type": "user", "sessionId": "sess-1", "timestamp": "2026-10-01T10:00:18.000Z", "message": {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_05", "is_error": false, "content": [{"type": "text", "text": "|Product_Name|mau|\n|---|---|\n|Alpha|120|\n|Beta|80|\n"}]}]}} +{"type": "user", "sessionId": "sess-1", "timestamp": "2026-10-01T10:00:30.000Z", "message": {"role": "user", "content": "Show me some raw rows."}} +{"type": "user", "sessionId": "sess-1", "isMeta": true, "timestamp": "2026-10-01T10:00:29.000Z", "message": {"role": "user", "content": [{"type": "text", "text": "Base directory for this skill: /x\n# Skill body"}]}} +{"type": "assistant", "sessionId": "sess-1", "timestamp": "2026-10-01T10:00:31.000Z", "message": {"role": "assistant", "content": [{"type": "text", "text": "Querying."}, {"type": "tool_use", "id": "toolu_06", "name": "mcp__palantir-mcp__run_sql_query_on_foundry_dataset", "input": {"sqlQuery": "SELECT * FROM `ri.foundry.main.dataset.11111111-2222-3333-4444-555555555555` LIMIT 5", "rowLimit": 5}}]}} +{"type": "user", "sessionId": "sess-1", "timestamp": "2026-10-01T10:00:33.000Z", "message": {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_06", "is_error": false, "content": [{"type": "text", "text": "|Customer_Name|Product_Name|Report_Date|MAU|\n|---|---|---|---|\n|Acme|Alpha|1788220800000|12|\n|Globex|Beta|1788220800000|7|\n"}]}]}} +{"type": "assistant", "sessionId": "sess-1", "timestamp": "2026-10-01T10:00:40.000Z", "message": {"role": "assistant", "content": [{"type": "text", "text": "Querying."}, {"type": "tool_use", "id": "toolu_07", "name": "mcp__palantir-mcp__run_sql_query_on_foundry_dataset", "input": {"sqlQuery": "SELECT * FROM `ri.foundry.main.dataset.11111111-2222-3333-4444-555555555555`", "rowLimit": 1000}}]}} +{"type": "user", "sessionId": "sess-1", "timestamp": "2026-10-01T10:02:20.000Z", "message": {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_07", "is_error": false, "content": "Error: result (912,345 characters across 1,001 lines) exceeds maximum allowed tokens. Output has been saved to /tmp/x.txt"}]}} +{"type": "assistant", "sessionId": "sess-1", "timestamp": "2026-10-01T10:02:30.000Z", "message": {"role": "assistant", "content": [{"type": "text", "text": "Querying."}, {"type": "tool_use", "id": "toolu_08", "name": "mcp__palantir-mcp__run_sql_query_on_foundry_dataset", "input": {"sqlQuery": "SELECT COUNT(*) FROM `ri.foundry.main.dataset.11111111-2222-3333-4444-555555555555`"}}]}} +{"type": "user", "sessionId": "sess-1", "timestamp": "2026-10-01T10:02:32.000Z", "message": {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_08", "is_error": true, "content": "Error: fetch failed\n{}"}]}} +{"type": "assistant", "sessionId": "sess-1", "timestamp": "2026-10-01T10:02:40.000Z", "message": {"role": "assistant", "content": [{"type": "text", "text": "Querying."}, {"type": "tool_use", "id": "toolu_09", "name": "mcp__palantir-mcp__run_sql_query_on_foundry_dataset", "input": {"sqlQuery": "SELECT COUNT(*) FROM `ri.foundry.main.dataset.11111111-2222-3333-4444-555555555555`"}}]}} +{"type": "user", "sessionId": "sess-1", "timestamp": "2026-10-01T10:02:43.000Z", "message": {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_09", "is_error": false, "content": [{"type": "text", "text": "|count(1)|\n|---|\n|42|\n"}]}]}} +{"type": "assistant", "sessionId": "sess-1", "timestamp": "2026-10-01T10:02:50.000Z", "message": {"role": "assistant", "content": [{"type": "text", "text": "Querying."}, {"type": "tool_use", "id": "toolu_10", "name": "mcp__palantir-mcp__run_sql_query_on_foundry_dataset", "input": {"sqlQuery": "SELECT Customer_Name FROM `ri.foundry.main.dataset.11111111-2222-3333-4444-555555555555` LIMIT 3"}}]}} diff --git a/tests/test_trace.py b/tests/test_trace.py new file mode 100644 index 0000000..2c1e338 --- /dev/null +++ b/tests/test_trace.py @@ -0,0 +1,499 @@ +from __future__ import annotations + +import io +import json +import subprocess +import sys +from pathlib import Path + +import pytest + +from querypilot.cli import main +from querypilot.evals.replay import replay_from_jsonl +from querypilot.trace import classify +from querypilot.trace.doctor import Redactor, diagnose, render_markdown, replay_records +from querypilot.trace.guard import guard_sql +from querypilot.trace.hooks import ( + HookOptions, + install_hooks, + main_hook, + run_hook, + tool_matcher, + uninstall_hooks, +) +from querypilot.trace.schema_cache import ( + CachedColumn, + CachedDataset, + SchemaCache, + parse_foundry_schema, +) +from querypilot.trace.tools import ( + FOUNDRY_SCHEMA_TOOL, + FOUNDRY_SQL_TOOL, + ToolRegistry, + parse_tool_spec, +) +from querypilot.trace.transcript import QUERY, SCHEMA, load_calls, result_text + +FIXTURE = Path(__file__).parent / "fixtures" / "claude_code_foundry_transcript.jsonl" +RID = "ri.foundry.main.dataset.11111111-2222-3333-4444-555555555555" + + +@pytest.fixture(autouse=True) +def _isolated_home(tmp_path, monkeypatch): + monkeypatch.setenv("QUERYPILOT_HOME", str(tmp_path / "qp-home")) + + +def _usage_cache() -> SchemaCache: + return SchemaCache( + { + RID: CachedDataset( + key=RID, + path="/Demo/Usage/monthly_usage", + columns=[ + CachedColumn("Customer_Name", "STRING"), + CachedColumn("Product_Name", "STRING"), + CachedColumn("Report_Date", "DATE"), + CachedColumn("MAU", "LONG"), + ], + ) + } + ) + + +# --------------------------------------------------------------------------- classify + + +@pytest.mark.parametrize( + ("text", "is_error", "status", "error_class"), + [ + ("|a|\n|---|\n|1|\n", False, classify.OK, None), + ( + "Sql Query Failed\nError Message: Failed to execute SQL query: [UNRESOLVED_COLUMN.WITH_SUGGESTION] x", + False, + classify.ERROR, + classify.UNRESOLVED_COLUMN, + ), + ("Sql Query Failed\nError Message: [DATATYPE_MISMATCH.BINARY_OP_DIFF_TYPES] y", False, "error", "type_mismatch"), + ("Error: result (1 characters) exceeds maximum allowed tokens.", False, "error", "result_too_large"), + ("Error: fetch failed\n{}", True, "error", "upstream"), + ("Error: Failed to fetch 500 \n**Suggestion:** A 403 Forbidden ... RID is incorrect", True, "error", "upstream"), + ("Error: Failed to fetch 403 \n**Suggestion:** dataset RID is incorrect", True, "error", "unknown_table"), + ('Error: Failed to fetch 400 \n{"errorName":"QueryParseError"}', True, "error", "syntax"), + ("Error: Failed to fetch 400 . Error: Failed to fetch 400\nat Object.fetch (cli.cjs:1:2)", True, "error", "upstream"), + ("Sql Query Failed\nError Message: [MISSING_GROUP_BY] z", False, "error", "sql_analysis"), + ("The user doesn't want to proceed with this tool use.", True, classify.DENIED, None), + ("QueryPilot guard blocked this query before it ran:\n- x", True, classify.BLOCKED, None), + ], +) +def test_classify_output(text, is_error, status, error_class) -> None: + outcome = classify.classify_output(text, is_error=is_error) + assert outcome.status == status + assert outcome.error_class == error_class + + +def test_result_rows_containing_error_words_are_still_success() -> None: + text = "|status|note|\n|---|---|\n|Error|Sql Query Failed upstream|\n|[ACTIVE]|timeout|\n" + outcome = classify.classify_output(text) + assert outcome.status == classify.OK + assert outcome.row_count == 2 + assert classify.markdown_header(text) == ["status", "note"] + + +# --------------------------------------------------------------------------- tools + + +def test_tool_specs_map_any_tool_by_name_and_field() -> None: + spec = parse_tool_spec("mcp__warehouse__*=query:snowflake:max_rows") + registry = ToolRegistry.from_args(["mcp__warehouse__*=query:snowflake:max_rows"]) + assert (spec.sql_field, spec.dialect, spec.row_limit_field) == ("query", "snowflake", "max_rows") + assert registry.query_spec("mcp__warehouse__run_query") == spec + assert registry.query_spec(FOUNDRY_SQL_TOOL).sql_field == "sqlQuery" + assert registry.schema_spec(FOUNDRY_SCHEMA_TOOL) is not None + assert registry.query_spec("Bash") is None + with pytest.raises(ValueError): + parse_tool_spec("no-field") + + +def test_tool_matcher_uses_exact_lists_when_possible() -> None: + assert tool_matcher([FOUNDRY_SQL_TOOL, FOUNDRY_SCHEMA_TOOL]) == f"{FOUNDRY_SQL_TOOL}|{FOUNDRY_SCHEMA_TOOL}" + assert tool_matcher(["mcp__wh__*", "mcp__a.b__run"]) == r"^(?:mcp__wh__.*|mcp__a\.b__run)$" + + +# --------------------------------------------------------------------------- transcripts + + +def test_load_calls_pairs_tool_use_with_result() -> None: + calls = load_calls([FIXTURE]) + queries = [c for c in calls if c.kind == QUERY] + schemas = [c for c in calls if c.kind == SCHEMA] + + assert len(queries) == 8 + assert len(schemas) == 1 and schemas[0].dataset.key == RID + by_id = {c.call_id: c for c in calls} + first = by_id["toolu_02"] + assert first.status == "error" and first.error_class == "unresolved_column" + assert first.latency_s == 4.0 + assert first.params == {"rowLimit": 100} + assert first.question == "Which products had the most monthly active users last month?" + assert by_id["toolu_05"].row_count == 2 + assert by_id["toolu_07"].error_class == "result_too_large" and by_id["toolu_07"].latency_s == 100.0 + assert by_id["toolu_08"].error_class == "upstream" + assert by_id["toolu_10"].status == "no_result" + assert by_id["toolu_06"].question == "Show me some raw rows." + + +def test_resumed_transcripts_are_deduplicated(tmp_path) -> None: + copy = tmp_path / "resumed.jsonl" + copy.write_text(FIXTURE.read_text()) + assert len(load_calls([FIXTURE, copy])) == len(load_calls([FIXTURE])) + + +def test_result_text_handles_mcp_shapes() -> None: + assert result_text("plain") == "plain" + assert result_text([{"type": "text", "text": "a"}, {"type": "image"}, {"type": "text", "text": "b"}]) == "a\nb" + assert result_text(None) == "" + + +# --------------------------------------------------------------------------- schema cache + + +def test_schema_cache_round_trip_and_select_star(tmp_path) -> None: + calls = load_calls([FIXTURE]) + schema_call = next(c for c in calls if c.kind == SCHEMA) + assert [c.name for c in schema_call.dataset.columns] == ["Customer_Name", "Product_Name", "Report_Date", "MAU"] + + cache = SchemaCache() + learned = cache.learn_from_select_star("SELECT * FROM `/Demo/other` LIMIT 5", ["a", "b"], "spark") + assert learned is not None and cache.lookup("/demo/OTHER").columns[1].name == "b" + assert cache.learn_from_select_star("SELECT a FROM `/Demo/other`", ["a"], "spark") is None + assert cache.learn_from_select_star("SELECT * FROM x JOIN y ON x.id = y.id", ["id"], "spark") is None + + cache.add(schema_call.dataset) + path = cache.save(tmp_path / "cache.json") + loaded = SchemaCache.load(path) + assert loaded.lookup(RID).path == "/Demo/Usage/monthly_usage" + assert loaded.lookup("/Demo/Usage/monthly_usage").key == RID + assert loaded.evict([RID]) == [RID] and loaded.lookup("/Demo/Usage/monthly_usage") is None + + +def test_select_star_never_overwrites_schema_tool_entry() -> None: + cache = _usage_cache() + assert cache.learn_from_select_star(f"SELECT * FROM `{RID}`", ["only_one"], "spark") is None + assert len(cache.lookup(RID).columns) == 4 + + +def test_parse_foundry_schema_rejects_non_schema_text() -> None: + assert parse_foundry_schema("no json here") is None + assert parse_foundry_schema('{"properties": {}}') is None + + +# --------------------------------------------------------------------------- guard + + +def test_guard_denies_unknown_column_with_did_you_mean() -> None: + decision = guard_sql(f"SELECT Product, SUM(MAU) FROM `{RID}` GROUP BY Product", _usage_cache()) + assert decision.denied and decision.rule == "schema" and decision.checked_columns + assert "did you mean monthly_usage.Product_Name" in decision.reason + assert f"monthly_usage = `{RID}`" in decision.reason + + +def test_guard_allows_valid_sql_and_uncached_datasets() -> None: + cache = _usage_cache() + sql = ( + f"WITH t AS (SELECT Product_Name AS p, MAU FROM `{RID}`) " + "SELECT p, SUM(MAU) AS m FROM t GROUP BY p ORDER BY m DESC LIMIT 5" + ) + assert not guard_sql(sql, cache).denied + other = guard_sql("SELECT nope FROM `ri.foundry.main.dataset.99999999-2222-3333-4444-555555555555`", cache) + assert not other.denied and not other.checked_columns + assert not guard_sql("SELECT FROM WHERE (((", cache).denied + + +def test_guard_require_schema_and_malformed_rid() -> None: + cache = _usage_cache() + unknown = "ri.foundry.main.dataset.99999999-2222-3333-4444-555555555555" + decision = guard_sql(f"SELECT a FROM `{unknown}`", cache, require_schema=True) + assert decision.denied and decision.rule == "schema_required" and "Fetch the schema" in decision.reason + bad = guard_sql("SELECT a FROM `ri.foundry.main.dataset.11111111-2222`", cache) + assert bad.denied and "not a valid dataset RID" in bad.reason + pinned = "ri.foundry.main.transaction.0000-1@master.ri.foundry.main.dataset.11111111-2222-3333-4444-555555555555" + assert not guard_sql(f"SELECT * FROM `{pinned}`", cache).denied + + +def test_guard_epoch_milliseconds_on_date_columns() -> None: + cache = _usage_cache() + for sql in ( + f"SELECT Report_Date / 1000 FROM `{RID}`", + f"SELECT MAU FROM `{RID}` WHERE Report_Date >= 1790812800000", + f"SELECT MAU FROM `{RID}` WHERE Report_Date >= unix_timestamp('2026-06-01') * 1000", + ): + decision = guard_sql(sql, cache) + assert decision.denied, sql + assert "not epoch milliseconds" in decision.reason + assert not guard_sql(f"SELECT MAU / 1000 FROM `{RID}` WHERE Report_Date >= DATE '2026-01-01'", cache).denied + + +def test_guard_row_cap() -> None: + cache = _usage_cache() + decision = guard_sql(f"SELECT * FROM `{RID}`", cache, max_rows=200) + assert decision.denied and decision.rule == "limit" + assert decision.limited_sql.endswith("LIMIT 200") + assert guard_sql(f"SELECT * FROM `{RID}` LIMIT 5000", cache, max_rows=200).denied + assert not guard_sql(f"SELECT * FROM `{RID}` LIMIT 50", cache, max_rows=200).denied + assert not guard_sql(f"SELECT * FROM `{RID}`", cache, max_rows=200, tool_row_limit=100).denied + assert not guard_sql(f"SELECT COUNT(*), MAX(MAU) FROM `{RID}`", cache, max_rows=200).denied + + +def test_guard_ignores_struct_fields_and_dotted_names() -> None: + cache = _usage_cache() + struct = ( + f"WITH z AS (SELECT explode(arrays_zip(Product_Name, MAU)) zz FROM `{RID}`) " + "SELECT zz.Product_Name FROM z" + ) + assert not guard_sql(struct, cache).denied + assert not guard_sql(f"SELECT `vendor.entityid` FROM `{RID}`", cache).denied + + +# --------------------------------------------------------------------------- doctor + + +def test_diagnose_sample_transcript() -> None: + report = diagnose(load_calls([FIXTURE])) + + assert (report.calls, report.ran, report.failed, report.no_result) == (8, 7, 4, 1) + assert report.failure_rate == pytest.approx(4 / 7) + assert {row.error_class: row.count for row in report.classes} == { + "unresolved_column": 1, + "type_mismatch": 1, + "result_too_large": 1, + "upstream": 1, + } + assert {row.error_class: row.fix for row in report.classes}["type_mismatch"] == "data" + assert report.chains.chains == 2 and report.chains.recovered == 2 + assert report.chains.schema_fetched_in_chain == 1 + assert report.latency.stalls == 1 and report.latency.p50 == 3.0 + assert report.guard.caught == {"type_mismatch": 1} + assert report.guard.false_blocks == 0 + fixes = {f.title: f.fix for f in report.findings} + assert fixes["Date columns read as epoch milliseconds"] == "data" + assert fixes["Validate columns before the query runs (QueryPilot guard)"] == "tool" + assert any("epoch" in hint for hint in report.hints) + assert any(key.endswith(".Report_Date") for key in report.descriptions) + + +def test_markdown_report_is_redacted() -> None: + text = render_markdown(diagnose(load_calls([FIXTURE]))) + assert "## Failures by class" in text and "| Unresolved column | 1 |" in text + assert RID not in text and "/Demo/Usage" not in text and "monthly_usage" not in text + assert "Acme" not in text and "Globex" not in text + + +def test_redactor_replaces_literals_and_datasets() -> None: + redactor = Redactor(True) + out = redactor.sql(f"SELECT a FROM `{RID}` WHERE name = 'Alice'") + assert "Alice" not in out and RID not in out and "dataset_A" in out + assert redactor.text(f"`{RID}` failed for '/Org/secret/path'") == "`dataset_A` failed for '?'" + + +def test_replay_export_feeds_eval_replay(tmp_path) -> None: + records = replay_records(load_calls([FIXTURE])) + assert [r.question for r in records] == [ + "Which products had the most monthly active users last month?", + "Show me some raw rows.", + ] + assert "DATE '2026-09-01'" in records[0].rewritten_sql + assert records[0].details["failure_classes"] == ["unresolved_column", "type_mismatch"] + + path = tmp_path / "replay.jsonl" + path.write_text("".join(r.model_dump_json() + "\n" for r in records)) + suite = replay_from_jsonl(path, fixture_db="sqlite:///mirror.db") + assert len(suite.cases) == 2 + assert "app:trace-unresolved_column" in suite.cases[0].tags + + +# --------------------------------------------------------------------------- hooks + + +def test_install_and_uninstall_preserve_other_hooks(tmp_path) -> None: + settings = tmp_path / "settings.json" + other = {"matcher": "Bash", "hooks": [{"type": "command", "command": "echo hi"}]} + settings.write_text(json.dumps({"model": "x", "hooks": {"PreToolUse": [other]}})) + options = HookOptions(audit_log=tmp_path / "audit.jsonl", guard=True, max_rows=200) + + install_hooks(settings, options, python="/usr/bin/python3") + install_hooks(settings, options, python="/usr/bin/python3") # idempotent + data = json.loads(settings.read_text()) + assert data["model"] == "x" + pre = data["hooks"]["PreToolUse"] + assert pre[0] == other and len(pre) == 2 + assert pre[1]["matcher"] == FOUNDRY_SQL_TOOL + command = pre[1]["hooks"][0]["command"] + assert command.startswith("/usr/bin/python3 -m querypilot hooks run pre") + assert "--guard" in command and "--max-rows 200" in command + assert FOUNDRY_SCHEMA_TOOL in data["hooks"]["PostToolUse"][0]["matcher"] + assert "PostToolUseFailure" in data["hooks"] + + assert uninstall_hooks(settings) == 3 + data = json.loads(settings.read_text()) + assert data["hooks"] == {"PreToolUse": [other]} + + +def test_logging_only_install_has_no_pre_hook(tmp_path) -> None: + settings = tmp_path / "settings.json" + install_hooks(settings, HookOptions(audit_log=tmp_path / "a.jsonl")) + assert set(json.loads(settings.read_text())["hooks"]) == {"PostToolUse", "PostToolUseFailure"} + + +def _hook_payloads(path: Path): + """Turn a transcript into the PreToolUse/PostToolUse payloads Claude Code would send.""" + uses = {} + for line in path.read_text().splitlines(): + record = json.loads(line) + content = record["message"]["content"] + if not isinstance(content, list): + continue + for block in content: + base = {"session_id": record["sessionId"], "transcript_path": str(path)} + if block.get("type") == "tool_use": + uses[block["id"]] = block + yield "pre", {**base, "hook_event_name": "PreToolUse", "tool_name": block["name"], + "tool_input": block["input"], "tool_use_id": block["id"]} + elif block.get("type") == "tool_result": + use = uses[block["tool_use_id"]] + payload = {**base, "tool_name": use["name"], "tool_input": use["input"], + "tool_use_id": use["id"], "duration_ms": 1500} + if block.get("is_error"): + yield "post-failure", {**payload, "hook_event_name": "PostToolUseFailure", + "error": result_text(block["content"])} + else: + yield "post", {**payload, "hook_event_name": "PostToolUse", "tool_response": block["content"]} + + +def test_hooks_replaying_sample_transcript_log_every_query_call(tmp_path) -> None: + options = HookOptions(audit_log=tmp_path / "audit.jsonl", cache_path=tmp_path / "cache.json") + outputs = [run_hook(event, payload, options) for event, payload in _hook_payloads(FIXTURE)] + + assert all(out is None for out in outputs) # logging only: never changes a call + hooked = {c.call_id: c for c in load_calls([options.audit_log])} + direct = {c.call_id: c for c in load_calls([FIXTURE]) if c.status != "no_result"} + assert set(hooked) == set(direct) + for call_id, call in direct.items(): + assert (hooked[call_id].status, hooked[call_id].error_class) == (call.status, call.error_class) + assert hooked[call_id].latency_s == 1.5 + assert SchemaCache.load(options.cache_path).lookup(RID) is not None + + report = diagnose(list(hooked.values())) + assert report.failed == 4 and report.chains.recovered == 2 + lines = options.audit_log.read_text() + assert "Alpha" not in lines and "Acme" not in lines # result rows are never logged + + +def test_guard_hook_denies_then_allows_identical_resend(tmp_path) -> None: + options = HookOptions(audit_log=tmp_path / "audit.jsonl", cache_path=tmp_path / "cache.json", guard=True) + _usage_cache().save(options.cache_path) + payload = { + "session_id": "s1", + "tool_name": FOUNDRY_SQL_TOOL, + "tool_use_id": "t1", + "tool_input": {"sqlQuery": f"SELECT Product FROM `{RID}` LIMIT 5", "rowLimit": 5}, + } + first = run_hook("pre", payload, options) + decision = first["hookSpecificOutput"] + assert decision["permissionDecision"] == "deny" + assert "Product_Name" in decision["permissionDecisionReason"] + assert "send it again unchanged" in decision["permissionDecisionReason"] + blocked = load_calls([options.audit_log]) + assert [c.status for c in blocked] == ["blocked"] + + assert run_hook("pre", {**payload, "tool_use_id": "t2"}, options) is None + # It succeeded, so the cached schema was stale: the post hook evicts it. + run_hook("post", {**payload, "tool_use_id": "t2", "tool_response": [{"type": "text", "text": "|Product|\n|---|\n|A|\n"}]}, options) + assert SchemaCache.load(options.cache_path).lookup(RID) is None + + +def test_guard_hook_annotates_epoch_rendered_dates(tmp_path) -> None: + options = HookOptions(audit_log=tmp_path / "audit.jsonl", cache_path=tmp_path / "cache.json", guard=True) + _usage_cache().save(options.cache_path) + out = run_hook( + "post", + { + "session_id": "s1", + "tool_name": FOUNDRY_SQL_TOOL, + "tool_use_id": "t1", + "tool_input": {"sqlQuery": f"SELECT Product_Name, Report_Date FROM `{RID}` LIMIT 2"}, + "tool_response": [{"type": "text", "text": "|Product_Name|Report_Date|\n|---|---|\n|A|1788220800000|\n"}], + }, + options, + ) + assert "Report_Date (DATE)" in out["hookSpecificOutput"]["additionalContext"] + + +def test_hook_ignores_other_tools_and_never_raises(tmp_path) -> None: + options = HookOptions(audit_log=tmp_path / "audit.jsonl", guard=True) + assert run_hook("pre", {"tool_name": "Bash", "tool_input": {"command": "ls"}}, options) is None + out = io.StringIO() + assert main_hook("post", options, stdin=io.StringIO("not json"), stdout=out) == 0 + assert out.getvalue() == "" + assert "Traceback" in (tmp_path / "hook-errors.log").read_text() + + +def test_hook_subprocess_round_trip(tmp_path) -> None: + audit = tmp_path / "audit.jsonl" + cache = tmp_path / "cache.json" + _usage_cache().save(cache) + payload = { + "session_id": "s1", + "tool_name": FOUNDRY_SQL_TOOL, + "tool_use_id": "t1", + "tool_input": {"sqlQuery": f"SELECT Report_Date / 1000 FROM `{RID}` LIMIT 1"}, + } + proc = subprocess.run( + [sys.executable, "-m", "querypilot", "hooks", "run", "pre", "--guard", + "--audit-log", str(audit), "--cache", str(cache)], + input=json.dumps(payload), capture_output=True, text=True, check=True, timeout=60, + ) + out = json.loads(proc.stdout) + assert out["hookSpecificOutput"]["permissionDecision"] == "deny" + assert "not epoch milliseconds" in out["hookSpecificOutput"]["permissionDecisionReason"] + + +# --------------------------------------------------------------------------- CLI + + +def test_cli_trace_import_and_doctor(tmp_path, capsys) -> None: + audit = tmp_path / "audit.jsonl" + cache = tmp_path / "cache.json" + assert main(["trace", "import", str(FIXTURE), "--output", str(audit), "--cache", str(cache)]) == 0 + assert main(["trace", "import", str(FIXTURE), "--output", str(audit), "--cache", str(cache)]) == 0 + out = capsys.readouterr().out + assert "wrote 9 new records" in out and "wrote 0 new records" in out + assert SchemaCache.load(cache).lookup(RID) is not None + assert "Alpha" not in audit.read_text() + + report = tmp_path / "report.md" + suite = tmp_path / "suite.yaml" + code = main( + ["trace", "doctor", str(audit), "--output", str(report), "--suite", str(suite), + "--fixture-db", "sqlite:///mirror.db"] + ) + assert code == 0 + assert "| Data type mismatch | 1 |" in report.read_text() + assert "Wrote 2 cases" in capsys.readouterr().out + assert suite.is_file() + + +def test_cli_hooks_install_dry_run(tmp_path, capsys) -> None: + settings = tmp_path / "settings.json" + assert main(["hooks", "install", "--settings", str(settings), "--guard", "--dry-run"]) == 0 + assert not settings.exists() + assert "querypilot hooks run pre" in capsys.readouterr().out + assert main(["hooks", "install", "--settings", str(settings), "--tool", "mcp__wh__run=query:snowflake"]) == 0 + assert "mcp__wh__run" in settings.read_text() + assert main(["hooks", "uninstall", "--settings", str(settings)]) == 0 + + +def test_guard_text_keeps_column_names_and_drops_dataset_names() -> None: + redactor = Redactor(True) + text = "Column 'product' does not exist; did you mean usage.Product_Name, other.Product_Name?" + assert redactor.guard_text(text) == "Column 'product' does not exist; did you mean Product_Name?"