feat(webapp): query boundary pinned end-to-end and a capped query retry - #4549
feat(webapp): query boundary pinned end-to-end and a capped query retry#4549kathiekiwi wants to merge 21 commits into
Conversation
The read-only guard was enforced by the TRQL grammar and a parser test, but nothing proved the route itself refuses a write; a route test now drives api.v1.query with a signed environment JWT and asserts nothing reaches ClickHouse. readonly=1 is no longer overridable by a caller's clickhouseSettings. run_query gives up after three consecutive failures so a broken query can't burn a whole agent turn. TRI-11165
|
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
WalkthroughThe webapp now rejects mutating query statements and enforces 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 2📝 Generate docstrings 💡
🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
…65' into feat/query-safety-tri-11165
@trigger.dev/build
trigger.dev
@trigger.dev/core
@trigger.dev/python
@trigger.dev/react-hooks
@trigger.dev/redis-worker
@trigger.dev/rsc
@trigger.dev/schema-to-json
@trigger.dev/sdk
commit: |
…65' into feat/query-safety-tri-11165
…x/watch-mode-keepalive-tri-13065
| vi.mock("~/db.server", () => { | ||
| const client = { | ||
| runtimeEnvironment: { | ||
| findFirst: mocks.runtimeEnvironmentFindFirst, | ||
| findMany: async () => [], | ||
| }, | ||
| revokedApiKey: { findMany: async () => [], findFirst: async () => null }, | ||
| project: { findMany: async () => [] }, | ||
| customerQuery: { findFirst: async () => null, create: mocks.customerQueryCreate }, | ||
| }; | ||
| return { prisma: client, $replica: client }; | ||
| }); | ||
| vi.mock("~/env.server", () => ({ | ||
| env: { | ||
| SESSION_SECRET: "test-session-secret", | ||
| QUERY_CLICKHOUSE_MAX_EXECUTION_TIME: "30", | ||
| QUERY_CLICKHOUSE_MAX_MEMORY_USAGE: 1000000, | ||
| QUERY_CLICKHOUSE_MAX_AST_ELEMENTS: 50000, | ||
| QUERY_CLICKHOUSE_MAX_EXPANDED_AST_ELEMENTS: 500000, | ||
| QUERY_CLICKHOUSE_MAX_BYTES_BEFORE_EXTERNAL_GROUP_BY: 1000000, | ||
| QUERY_CLICKHOUSE_MAX_RETURNED_ROWS: 1000, | ||
| }, | ||
| })); | ||
| vi.mock("~/services/clickhouse/clickhouseFactoryInstance.server", () => ({ | ||
| clickhouseFactory: { | ||
| getClickhouseForOrganization: async () => ({ | ||
| reader: { queryWithStats: mocks.queryWithStats }, | ||
| }), | ||
| }, | ||
| })); | ||
| vi.mock("~/services/platform.v3.server", () => ({ getLimit: async () => 30 })); | ||
| vi.mock("~/services/queryConcurrencyLimiter.server", () => ({ | ||
| queryConcurrencyLimiter: { | ||
| acquire: async () => ({ success: true }), | ||
| release: async () => {}, | ||
| }, | ||
| DEFAULT_ORG_CONCURRENCY_LIMIT: 10, | ||
| GLOBAL_CONCURRENCY_LIMIT: 100, | ||
| })); | ||
| vi.mock("~/services/logger.server", () => ({ | ||
| logger: { debug: vi.fn(), error: vi.fn(), warn: vi.fn(), info: vi.fn() }, | ||
| })); | ||
| vi.mock("~/v3/services/worker/workerGroupTokenService.server", () => ({ | ||
| WorkerGroupTokenService: class {}, | ||
| })); | ||
| vi.mock("~/v3/services/common.server", () => ({ ServiceValidationError: class extends Error {} })); | ||
| vi.mock("@internal/run-engine", () => ({ EngineServiceValidationError: class extends Error {} })); |
There was a problem hiding this comment.
🟡 New tests rely on mocking, which the repository's test guidance forbids
The new test suites stub out the database, environment config, ClickHouse factory, logger and concurrency limiter with mock objects (vi.mock("~/db.server", ...) at apps/webapp/test/queryRouteReadOnly.test.ts:37-83), whereas the repository's testing guidance requires real containers instead of mocks.
Impact: The tests can keep passing while the real database, environment and ClickHouse wiring drifts, so the read-only guarantee they claim to pin is only proven against stand-ins.
Which rule is violated and where
AGENTS.md ("Testing") states: "We use vitest exclusively. Never mock anything - use testcontainers instead." Both new files rely on mocking: apps/webapp/test/queryRouteReadOnly.test.ts:31-83 mocks ~/db.server, ~/env.server, the ClickHouse factory instance, the platform limits service, the concurrency limiter and the logger; internal-packages/dashboard-agent/src/tool-query-retry-cap.test.ts:12-26 hand-builds a fake DashboardAgentApiClient. The webapp suite has an existing test/ layout with container-based e2e configs (apps/webapp/test/README.md) that these route-level assertions could use instead.
Was this helpful? React with 👍 or 👎 to provide feedback.
| if (!result.ok) { | ||
| // Only SQL errors count toward the cap; transport errors are transient. | ||
| if (result.kind === "query") { | ||
| consecutiveQueryFailures++; | ||
| if (consecutiveQueryFailures >= MAX_CONSECUTIVE_QUERY_FAILURES) { | ||
| return { | ||
| error: `${result.error} That is ${consecutiveQueryFailures} queries in a row that failed. Stop querying and answer the user with what you already have.`, | ||
| }; | ||
| } | ||
| } | ||
| return { error: result.error }; | ||
| } | ||
| consecutiveQueryFailures = 0; |
There was a problem hiding this comment.
🔍 Chart-query validation failures bypass the retry cap
render_view validates chart queries through validateChartQuery, which also posts to the query endpoint (internal-packages/dashboard-agent/src/tool-api-client.ts), but a failure there returns an error prompting the model to "Fix the query ... and render the chart again" (internal-packages/dashboard-agent/src/tool-api.ts:407-412) without touching consecutiveQueryFailures. So the failure mode the cap is meant to prevent — a model burning the turn's step budget rewriting a query it can't fix — is still reachable via repeated render_view calls.
Was this helpful? React with 👍 or 👎 to provide feedback.
What & why
The agent's query tool is read-only, but that was true by three separate facts and only one of them had a test. This proves the boundary holds by contract rather than by prompt, and stops a broken query from burning a whole agent turn.
Two small guards, the rest is tests. TRI-11165.
Stack
Stacked on #4548 (watch-mode keepalive). Merge that first.
What's inside
apps/webapp/test/queryRouteReadOnly.test.ts) that drivesapi.v1.querywith a real signed environment JWT: a multi-statement write and a mutating statement are both refused before anything reaches ClickHouse, and a plain read passes so the seam stays live.readonly=1made non-overridable inqueryService.server.ts— callerclickhouseSettingswere spread after the defaults and could clear it.run_querytool (internal-packages/dashboard-agent/src/tool-api.ts): three consecutive failures returns a terminal "stop and answer with what you have".Key decisions
readonly=1, and the org/project/env scoping is injected server-side from the credential. The request body can't widen scope or turn a read into a write.Testing
queryRouteReadOnly.test.ts— write statements → 400, ClickHouse never called; a read passes.tool-query-retry-cap.test.ts— terminal at the third consecutive failure, counter resets on a success.