From 27b04fb9108bce90ec10389dd84dc52aa9ced189 Mon Sep 17 00:00:00 2001 From: Katia Bulatova Date: Mon, 10 Aug 2026 09:56:28 +0000 Subject: [PATCH 01/11] fix(sdk): watch-mode chat subscriptions survive quiet windows (TRI-13065) --- .changeset/watch-mode-keepalive.md | 5 + packages/trigger-sdk/src/v3/chat.test.ts | 127 +++++++++++++++++- packages/trigger-sdk/src/v3/chat.ts | 8 +- .../test/chat-transport-events.test.ts | 13 +- 4 files changed, 147 insertions(+), 6 deletions(-) create mode 100644 .changeset/watch-mode-keepalive.md diff --git a/.changeset/watch-mode-keepalive.md b/.changeset/watch-mode-keepalive.md new file mode 100644 index 00000000000..49e8fafb6cf --- /dev/null +++ b/.changeset/watch-mode-keepalive.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/sdk": patch +--- + +Watch-mode chat subscriptions now stay connected across quiet periods. diff --git a/packages/trigger-sdk/src/v3/chat.test.ts b/packages/trigger-sdk/src/v3/chat.test.ts index 21aa004c9a9..42da9c83951 100644 --- a/packages/trigger-sdk/src/v3/chat.test.ts +++ b/packages/trigger-sdk/src/v3/chat.test.ts @@ -1218,6 +1218,122 @@ describe("TriggerChatTransport", () => { }); }); + describe("watch mode across long-poll window boundaries", () => { + function settled(response: Response): Response { + const headers = new Headers(response.headers); + headers.set("X-Session-Settled", "true"); + return new Response(response.body, { status: 200, headers }); + } + + it("resubscribes after a completed turn and receives a later wake", async () => { + let subscribeCount = 0; + global.fetch = vi.fn().mockImplementation(async (url: string | URL) => { + const urlStr = typeof url === "string" ? url : url.toString(); + if (isSessionOutSubscribeUrl(urlStr)) { + subscribeCount++; + // Window 1: a turn completes, then the body EOFs with no + // settled header — the quiet long-poll boundary. + return subscribeCount === 1 + ? defaultSseResponse([ + { type: "text-delta", id: "p1", delta: "turn1" }, + { type: "trigger:turn-complete" }, + ]) + : settled(defaultSseResponse([{ type: "text-delta", id: "p2", delta: "wake" }])); + } + throw new Error(`Unexpected URL: ${urlStr}`); + }); + + const transport = new TriggerChatTransport({ + task: "my-chat-task", + accessToken: () => "pat", + watch: true, + sessions: { "chat-watch-eof": { publicAccessToken: "p", isStreaming: true } }, + }); + + const stream = await transport.reconnectToStream({ chatId: "chat-watch-eof" }); + const chunks = await drainChunks(stream!); + + expect(subscribeCount).toBe(2); + expect(chunks).toEqual([ + { type: "text-delta", id: "p1", delta: "turn1" }, + { type: "text-delta", id: "p2", delta: "wake" }, + ]); + }); + + it("stops when the server says the session settled", async () => { + let subscribeCount = 0; + global.fetch = vi.fn().mockImplementation(async (url: string | URL) => { + const urlStr = typeof url === "string" ? url : url.toString(); + if (isSessionOutSubscribeUrl(urlStr)) { + subscribeCount++; + return settled( + defaultSseResponse([ + { type: "text-delta", id: "p1", delta: "last" }, + { type: "trigger:turn-complete" }, + ]) + ); + } + throw new Error(`Unexpected URL: ${urlStr}`); + }); + + const transport = new TriggerChatTransport({ + task: "my-chat-task", + accessToken: () => "pat", + watch: true, + sessions: { "chat-watch-settled": { publicAccessToken: "p", isStreaming: true } }, + }); + + const stream = await transport.reconnectToStream({ chatId: "chat-watch-settled" }); + const chunks = await drainChunks(stream!); + + expect(subscribeCount).toBe(1); + expect(chunks).toHaveLength(1); + expect(transport.getSession("chat-watch-settled")?.isStreaming).toBe(false); + }); + + it("stops promptly when aborted during backoff", async () => { + vi.useFakeTimers(); + try { + let subscribeCount = 0; + global.fetch = vi.fn().mockImplementation(async (url: string | URL) => { + const urlStr = typeof url === "string" ? url : url.toString(); + if (isSessionStreamAppendUrl(urlStr)) return defaultAppendResponse(); + if (isSessionOutSubscribeUrl(urlStr)) { + subscribeCount++; + // Every window is quiet: EOF with no records, never settled. + return defaultSseResponse([]); + } + throw new Error(`Unexpected URL: ${urlStr}`); + }); + + const abortController = new AbortController(); + const transport = new TriggerChatTransport({ + task: "my-chat-task", + accessToken: () => "pat", + watch: true, + sessions: { "chat-watch-abort": { publicAccessToken: "p", isStreaming: true } }, + }); + + const stream = await transport.reconnectToStream({ + chatId: "chat-watch-abort", + abortSignal: abortController.signal, + }); + const drained = drainChunks(stream!); + await vi.advanceTimersByTimeAsync(10_000); + // The budget doesn't apply in watch mode, so it is still reconnecting. + expect(subscribeCount).toBeGreaterThan(6); + + const countAtAbort = subscribeCount; + abortController.abort(); + await drained; + await vi.advanceTimersByTimeAsync(10_000); + expect(subscribeCount).toBe(countAtAbort); + } finally { + vi.useRealTimers(); + } + }); + }); + describe("multi-tab coordination", () => { it("isReadOnly defaults to false when multiTab is disabled", () => { const transport = new TriggerChatTransport({ @@ -1488,9 +1604,18 @@ describe("TriggerChatTransport", () => { { type: "text-delta", id: "p2", delta: "Again" }, { type: "trigger:turn-complete" }, ]; + let subscribeCount = 0; global.fetch = vi.fn().mockImplementation(async (url: string | URL) => { const urlStr = typeof url === "string" ? url : url.toString(); - if (isSessionOutSubscribeUrl(urlStr)) return defaultSseResponse(turn1); + if (isSessionOutSubscribeUrl(urlStr)) { + subscribeCount++; + if (subscribeCount === 1) return defaultSseResponse(turn1); + // Watch mode reconnects past the body EOF; settle so the drain ends. + const response = defaultSseResponse([]); + const headers = new Headers(response.headers); + headers.set("X-Session-Settled", "true"); + return new Response(response.body, { status: 200, headers }); + } throw new Error(`Unexpected URL: ${urlStr}`); }); diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index 137222b1bae..07e89fc6d9a 100644 --- a/packages/trigger-sdk/src/v3/chat.ts +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -1780,11 +1780,13 @@ export class TriggerChatTransport implements ChatTransport { let eofResubscribes = 0; const resumeAfterEof = async () => { + // Watch mode is a standing subscription: it outlives turn-complete + // (which clears `isStreaming`) and idle windows EOF by design, so the + // give-up budget doesn't apply. Only abort or a settled session ends it. while ( - state.isStreaming && + (this.watchMode || (state.isStreaming && eofResubscribes < MAX_EOF_RESUBSCRIBES)) && !currentSubscription?.sessionSettled && - !combinedSignal.aborted && - eofResubscribes < MAX_EOF_RESUBSCRIBES + !combinedSignal.aborted ) { eofResubscribes++; // Sleep, but wake immediately on abort — otherwise a stop lands diff --git a/packages/trigger-sdk/test/chat-transport-events.test.ts b/packages/trigger-sdk/test/chat-transport-events.test.ts index ded1895dfd8..39f4e53d722 100644 --- a/packages/trigger-sdk/test/chat-transport-events.test.ts +++ b/packages/trigger-sdk/test/chat-transport-events.test.ts @@ -211,11 +211,20 @@ describe("transport stream events", () => { ``, ].join("\n"); + let subscribes = 0; const { transport, events } = makeTransport({ watch: true, sessions: { c1: { publicAccessToken: "tok_test", isStreaming: true } }, - fetch: async (_url, _init, ctx) => - ctx.endpoint === "in" ? jsonOk() : sseResponse(TWO_TURNS), + fetch: async (_url, _init, ctx) => { + if (ctx.endpoint === "in") return jsonOk(); + if (subscribes++ > 0) { + // Watch mode reconnects past the body EOF; settle so the read ends. + const settled = sseResponse(""); + settled.headers.set("X-Session-Settled", "true"); + return settled; + } + return sseResponse(TWO_TURNS); + }, }); const stream = await transport.reconnectToStream({ chatId: "c1" }); From 574a84d659a03d3be00048d925fb292689b7898e Mon Sep 17 00:00:00 2001 From: Katia Bulatova Date: Mon, 10 Aug 2026 10:43:20 +0000 Subject: [PATCH 02/11] fix(sdk): stopping a turn is an owning-transport action, not an abort side effect (TRI-13070) --- .changeset/watch-mode-keepalive.md | 2 + packages/trigger-sdk/src/v3/chat.test.ts | 88 ++++++++++++++++++++++++ packages/trigger-sdk/src/v3/chat.ts | 21 +++++- 3 files changed, 108 insertions(+), 3 deletions(-) diff --git a/.changeset/watch-mode-keepalive.md b/.changeset/watch-mode-keepalive.md index 49e8fafb6cf..e50a04f2a8c 100644 --- a/.changeset/watch-mode-keepalive.md +++ b/.changeset/watch-mode-keepalive.md @@ -3,3 +3,5 @@ --- Watch-mode chat subscriptions now stay connected across quiet periods. + +Read-only chat subscriptions no longer stop a turn when they disconnect. diff --git a/packages/trigger-sdk/src/v3/chat.test.ts b/packages/trigger-sdk/src/v3/chat.test.ts index 42da9c83951..f9851f52ade 100644 --- a/packages/trigger-sdk/src/v3/chat.test.ts +++ b/packages/trigger-sdk/src/v3/chat.test.ts @@ -1334,6 +1334,94 @@ describe("TriggerChatTransport", () => { }); }); + describe("reconnectToStream stop-on-abort ownership (TRI-13070)", () => { + // A quiet stream: EOF, no records, never settled — the subscription + // stays alive (watch mode) so an abort mid-flight exercises the stop path. + function quietWatchTransport(): { + transport: TriggerChatTransport; + appends: () => number; + } { + let appendCount = 0; + global.fetch = vi.fn().mockImplementation(async (url: string | URL) => { + const urlStr = typeof url === "string" ? url : url.toString(); + if (isSessionStreamAppendUrl(urlStr)) { + appendCount++; + return defaultAppendResponse(); + } + if (isSessionOutSubscribeUrl(urlStr)) return defaultSseResponse([]); + throw new Error(`Unexpected URL: ${urlStr}`); + }); + const transport = new TriggerChatTransport({ + task: "my-chat-task", + accessToken: () => "pat", + watch: true, + sessions: { "chat-own": { publicAccessToken: "p", isStreaming: true } }, + }); + return { transport, appends: () => appendCount }; + } + + it("passive subscriber aborting writes no stop chunk to .in", async () => { + vi.useFakeTimers(); + try { + const { transport, appends } = quietWatchTransport(); + const abort = new AbortController(); + const stream = await transport.reconnectToStream({ + chatId: "chat-own", + abortSignal: abort.signal, + }); + const drained = drainChunks(stream!); + await vi.advanceTimersByTimeAsync(1_000); + abort.abort(); + await drained; + await vi.advanceTimersByTimeAsync(1_000); + expect(appends()).toBe(0); + } finally { + vi.useRealTimers(); + } + }); + + it("owning subscriber with stopOnAbort:true sends a stop chunk on abort", async () => { + vi.useFakeTimers(); + try { + const { transport, appends } = quietWatchTransport(); + const abort = new AbortController(); + const stream = await transport.reconnectToStream({ + chatId: "chat-own", + abortSignal: abort.signal, + stopOnAbort: true, + }); + const drained = drainChunks(stream!); + await vi.advanceTimersByTimeAsync(1_000); + abort.abort(); + await drained; + await vi.advanceTimersByTimeAsync(1_000); + expect(appends()).toBe(1); + } finally { + vi.useRealTimers(); + } + }); + + it("abortSignal presence alone (stopOnAbort unset) sends no stop", async () => { + vi.useFakeTimers(); + try { + const { transport, appends } = quietWatchTransport(); + const abort = new AbortController(); + const stream = await transport.reconnectToStream({ + chatId: "chat-own", + abortSignal: abort.signal, + }); + const drained = drainChunks(stream!); + await vi.advanceTimersByTimeAsync(1_000); + abort.abort(); + await drained; + await vi.advanceTimersByTimeAsync(1_000); + expect(appends()).toBe(0); + } finally { + vi.useRealTimers(); + } + }); + }); + describe("multi-tab coordination", () => { it("isReadOnly defaults to false when multiTab is disabled", () => { const transport = new TriggerChatTransport({ diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index 07e89fc6d9a..34cc2ad2320 100644 --- a/packages/trigger-sdk/src/v3/chat.ts +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -873,7 +873,11 @@ export class TriggerChatTransport implements ChatTransport { state.isStreaming = true; this.notifySessionChange(chatId, state); - return this.subscribeToSessionStream(state, abortSignal, chatId, { sinceInSeq: inSeq }); + // Owning turn: aborting this live send stops the turn the user drives. + return this.subscribeToSessionStream(state, abortSignal, chatId, { + sinceInSeq: inSeq, + sendStopOnAbort: true, + }); }; /** @@ -1146,6 +1150,13 @@ export class TriggerChatTransport implements ChatTransport { options: { chatId: string; abortSignal?: AbortSignal | undefined; + /** + * Whether aborting this subscription sends `{kind:"stop"}` on `.in`. + * A subscription ending is not session ownership — a passive/watch + * reader unmounting must never stop a turn it doesn't drive. Only + * pass `true` from a caller that owns the live turn. @default false + */ + stopOnAbort?: boolean; } & ChatRequestOptions ): Promise | null> => { const state = this.sessions.get(options.chatId); @@ -1163,7 +1174,7 @@ export class TriggerChatTransport implements ChatTransport { return this.subscribeToSessionStream(state, abortSignal, options.chatId, { resumed: true, - sendStopOnAbort: !!options.abortSignal, + sendStopOnAbort: options.stopOnAbort ?? false, // Reconnect-on-reload opts into the server's settled-peek shortcut // so the SSE doesn't hang for 60s when no turn is in flight. Active // send-a-message paths must keep wait=60 to avoid racing the @@ -1266,7 +1277,11 @@ export class TriggerChatTransport implements ChatTransport { state.isStreaming = true; this.notifySessionChange(chatId, state); - return this.subscribeToSessionStream(state, undefined, chatId, { sinceInSeq: inSeq }); + // Owning action: aborting this send stops the turn the user drives. + return this.subscribeToSessionStream(state, undefined, chatId, { + sinceInSeq: inSeq, + sendStopOnAbort: true, + }); }; // ------------------------------------------------------------------------- From fef7b438f78830c19972dab99988110458a231e2 Mon Sep 17 00:00:00 2001 From: Katia Bulatova Date: Mon, 10 Aug 2026 14:46:29 +0000 Subject: [PATCH 03/11] fix(sdk): surface an error when a watch-mode turn is truncated by reconnect exhaustion --- .changeset/chat-truncated-turn-error.md | 5 +++++ packages/trigger-sdk/src/v3/chat.ts | 14 ++++++++++++++ 2 files changed, 19 insertions(+) create mode 100644 .changeset/chat-truncated-turn-error.md diff --git a/.changeset/chat-truncated-turn-error.md b/.changeset/chat-truncated-turn-error.md new file mode 100644 index 00000000000..f0805d00b8b --- /dev/null +++ b/.changeset/chat-truncated-turn-error.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/sdk": patch +--- + +Chat in the browser now shows an error when a reply is cut off by a lost connection that can't be re-established, instead of presenting the partial reply as if it had finished. diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index 34cc2ad2320..fb4b2df3f28 100644 --- a/packages/trigger-sdk/src/v3/chat.ts +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -1821,6 +1821,20 @@ export class TriggerChatTransport implements ChatTransport { if (opened) return opened; } + // A settled session or an abort ends the turn cleanly. Exhausting the + // resubscribe budget while the turn is still streaming means it was cut + // off — surface an error so the UI doesn't read a truncated reply as + // complete. The caller's catch emits stream-error and errors the stream. + if ( + state.isStreaming && + !currentSubscription?.sessionSettled && + !combinedSignal.aborted + ) { + throw new Error( + "Chat stream ended before the turn completed (reconnect budget exhausted)." + ); + } + // Settled close, or the turn is gone — tell the UI instead of // leaving it spinning on a stream nobody will finish. if (state.isStreaming) { From a95973c88d3189a45a01da9bba0f9858978c49ec Mon Sep 17 00:00:00 2001 From: Katia Bulatova Date: Mon, 10 Aug 2026 14:47:03 +0000 Subject: [PATCH 04/11] fix(sdk): jitter the chat reconnect backoff --- packages/trigger-sdk/src/v3/chat.ts | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index fb4b2df3f28..2ab7ede39df 100644 --- a/packages/trigger-sdk/src/v3/chat.ts +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -1813,7 +1813,10 @@ export class TriggerChatTransport implements ChatTransport { combinedSignal.removeEventListener("abort", done); resolve(); }; - timer = setTimeout(done, Math.min(100 * 2 ** (eofResubscribes - 1), 5_000)); + // Jitter the backoff so many clients reconnecting after the same + // dropped window don't resubscribe in lockstep. + const backoff = Math.min(100 * 2 ** (eofResubscribes - 1), 5_000); + timer = setTimeout(done, backoff * (0.5 + Math.random() * 0.5)); combinedSignal.addEventListener("abort", done); }); if (combinedSignal.aborted) break; From bdf01b005667ba5f7760c4705ca313c98c02a2a2 Mon Sep 17 00:00:00 2001 From: Katia Bulatova Date: Mon, 10 Aug 2026 18:56:35 +0000 Subject: [PATCH 05/11] fix(sdk): keep watch-mode chat streams open across turns and clear state on give-up - reconnect no longer peek-settles in watch mode, so a settled peek between turns can't close the standing subscription before the next turn. - the returned stream now aborts its resubscribe loop when the reader is cancelled, instead of leaking it. - clear and persist isStreaming before the budget-exhaustion throw so a reload doesn't reopen a doomed subscription. --- .changeset/chat-truncated-turn-error.md | 2 ++ packages/trigger-sdk/src/v3/chat.ts | 16 ++++++++++++++-- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/.changeset/chat-truncated-turn-error.md b/.changeset/chat-truncated-turn-error.md index f0805d00b8b..959231a8fc8 100644 --- a/.changeset/chat-truncated-turn-error.md +++ b/.changeset/chat-truncated-turn-error.md @@ -3,3 +3,5 @@ --- Chat in the browser now shows an error when a reply is cut off by a lost connection that can't be re-established, instead of presenting the partial reply as if it had finished. + +Watch-mode viewers now keep receiving later turns instead of the stream closing after the first turn completes. diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index 2ab7ede39df..9dbeab2aeb3 100644 --- a/packages/trigger-sdk/src/v3/chat.ts +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -1178,8 +1178,10 @@ export class TriggerChatTransport implements ChatTransport { // Reconnect-on-reload opts into the server's settled-peek shortcut // so the SSE doesn't hang for 60s when no turn is in flight. Active // send-a-message paths must keep wait=60 to avoid racing the - // freshly-triggered turn's first chunk. - peekSettled: true, + // freshly-triggered turn's first chunk. Watch mode must NOT peek: a + // settled peek between turns sets sessionSettled and closes the + // standing subscription, so the viewer never sees the next turn. + peekSettled: !this.watchMode, }); }; @@ -1833,6 +1835,11 @@ export class TriggerChatTransport implements ChatTransport { !currentSubscription?.sessionSettled && !combinedSignal.aborted ) { + // Clear + persist before throwing so the surfaced error leaves + // consistent state — otherwise a reload sees isStreaming: true + // and reopens a doomed subscription. + state.isStreaming = false; + this.notifySessionChange(chatId, state); throw new Error( "Chat stream ended before the turn completed (reconnect budget exhausted)." ); @@ -2039,6 +2046,11 @@ export class TriggerChatTransport implements ChatTransport { this.coordinator?.release(chatId); } }, + // A consumer that stops reading without aborting (drops the reader) + // would otherwise leave the resubscribe loop running forever. + cancel() { + internalAbort.abort(); + }, }); } } From e5ede6f62cd0a8fa0e63776ba3361108429f454e Mon Sep 17 00:00:00 2001 From: Katia Bulatova Date: Mon, 10 Aug 2026 18:56:44 +0000 Subject: [PATCH 06/11] test(sdk): cover watch-mode keepalive, reader-cancel, and budget-exhaustion error --- packages/trigger-sdk/src/v3/chat.test.ts | 100 ++++++++++++++++++++++- 1 file changed, 98 insertions(+), 2 deletions(-) diff --git a/packages/trigger-sdk/src/v3/chat.test.ts b/packages/trigger-sdk/src/v3/chat.test.ts index f9851f52ade..09296fc7f22 100644 --- a/packages/trigger-sdk/src/v3/chat.test.ts +++ b/packages/trigger-sdk/src/v3/chat.test.ts @@ -1176,7 +1176,7 @@ describe("TriggerChatTransport", () => { expect(transport.getSession("chat-slow")?.isStreaming).toBe(false); }); - it("gives up after a bounded number of resubscribes", async () => { + it("surfaces an error after the resubscribe budget is exhausted", async () => { // Fake timers so the 100ms..1.6s backoffs don't cost real seconds. vi.useFakeTimers(); try { @@ -1205,12 +1205,18 @@ describe("TriggerChatTransport", () => { messages: [createUserMessage("hi")], abortSignal: undefined, }); + // A cut-off turn surfaces an error rather than reading as complete. + // Attach the rejection assertion before advancing timers so the + // rejection is never unhandled. const drained = drainChunks(stream); + const rejects = expect(drained).rejects.toThrow(/reconnect budget exhausted/i); await vi.advanceTimersByTimeAsync(10_000); - await drained; + await rejects; // One initial connect plus the five-attempt resubscribe budget. expect(subscribeCount).toBe(6); + // State is cleared before the throw, so a reload won't reopen a + // doomed subscription. expect(transport.getSession("chat-empty")?.isStreaming).toBe(false); } finally { vi.useRealTimers(); @@ -1260,6 +1266,96 @@ describe("TriggerChatTransport", () => { ]); }); + it("does not peek-settle an idle resubscribe, so the next turn is delivered", async () => { + // Watch mode must NOT send X-Peek-Settled between turns: a settled peek + // while no turn is in flight closes the standing subscription and the + // viewer never sees turn 2. This mock plays the server's peek shortcut — + // a peek request with nothing in flight settles — to prove the transport + // long-polls instead. + const subscribeHeaders: Headers[] = []; + global.fetch = vi.fn().mockImplementation(async (url: string | URL, init?: RequestInit) => { + const urlStr = typeof url === "string" ? url : url.toString(); + if (isSessionOutSubscribeUrl(urlStr)) { + subscribeHeaders.push(new Headers(init?.headers)); + const n = subscribeHeaders.length; + if (n === 1) { + // Turn 1 completes, then the body EOFs (no settled header). + return defaultSseResponse([ + { type: "text-delta", id: "p1", delta: "turn1" }, + { type: "trigger:turn-complete" }, + ]); + } + if (n === 2) { + // Idle resubscribe. If it peeked, the server settles and the + // subscription would close before turn 2; a long-poll delivers it. + if (init && new Headers(init.headers).get("X-Peek-Settled")) { + return settled(defaultSseResponse([])); + } + return defaultSseResponse([ + { type: "text-delta", id: "p2", delta: "turn2" }, + { type: "trigger:turn-complete" }, + ]); + } + // Turn 2 done — end the watch cleanly. + return settled(defaultSseResponse([])); + } + throw new Error(`Unexpected URL: ${urlStr}`); + }); + + const transport = new TriggerChatTransport({ + task: "my-chat-task", + accessToken: () => "pat", + watch: true, + sessions: { "chat-watch-turn2": { publicAccessToken: "p", isStreaming: true } }, + }); + + const stream = await transport.reconnectToStream({ chatId: "chat-watch-turn2" }); + const chunks = await drainChunks(stream!); + + expect(subscribeHeaders[1]?.get("X-Peek-Settled")).toBeNull(); + expect(chunks).toEqual([ + { type: "text-delta", id: "p1", delta: "turn1" }, + { type: "text-delta", id: "p2", delta: "turn2" }, + ]); + }); + + it("cancelling the reader stops the resubscribe loop", async () => { + // A consumer that stops reading without aborting must not leak the + // resubscribe loop — the stream's cancel() aborts it. + vi.useFakeTimers(); + try { + let subscribeCount = 0; + global.fetch = vi.fn().mockImplementation(async (url: string | URL) => { + const urlStr = typeof url === "string" ? url : url.toString(); + if (isSessionOutSubscribeUrl(urlStr)) { + subscribeCount++; + // Quiet: EOF, no records, never settled — watch keeps resubscribing. + return defaultSseResponse([]); + } + throw new Error(`Unexpected URL: ${urlStr}`); + }); + + const transport = new TriggerChatTransport({ + task: "my-chat-task", + accessToken: () => "pat", + watch: true, + sessions: { "chat-watch-cancel": { publicAccessToken: "p", isStreaming: true } }, + }); + + const stream = await transport.reconnectToStream({ chatId: "chat-watch-cancel" }); + const reader = stream!.getReader(); + await vi.advanceTimersByTimeAsync(10_000); + expect(subscribeCount).toBeGreaterThan(1); + + const countAtCancel = subscribeCount; + await reader.cancel(); + await vi.advanceTimersByTimeAsync(10_000); + expect(subscribeCount).toBe(countAtCancel); + } finally { + vi.useRealTimers(); + } + }); + it("stops when the server says the session settled", async () => { let subscribeCount = 0; global.fetch = vi.fn().mockImplementation(async (url: string | URL) => { From c40e5f08c21692a0181ebf09a507bee2a094833b Mon Sep 17 00:00:00 2001 From: Katia Bulatova Date: Mon, 10 Aug 2026 19:06:28 +0000 Subject: [PATCH 07/11] fix(sdk): guard controller.close() so a clean reader cancel emits no stream-error A consumer cancelling the watch stream aborts the resubscribe loop, which reaches controller.close() on an already-closed controller. The resulting 'Invalid state' throw was surfaced as a bogus stream-error on every clean watch-viewer unmount. Wrap the remaining bare close sites to match the existing pattern. --- packages/trigger-sdk/src/v3/chat.test.ts | 8 +++++++- packages/trigger-sdk/src/v3/chat.ts | 18 +++++++++++++++--- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/packages/trigger-sdk/src/v3/chat.test.ts b/packages/trigger-sdk/src/v3/chat.test.ts index 09296fc7f22..9b0bfd598a8 100644 --- a/packages/trigger-sdk/src/v3/chat.test.ts +++ b/packages/trigger-sdk/src/v3/chat.test.ts @@ -1,6 +1,6 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; import type { UIMessage, UIMessageChunk } from "ai"; -import { TriggerChatTransport, createChatTransport } from "./chat.js"; +import { TriggerChatTransport, createChatTransport, type ChatTransportEvent } from "./chat.js"; // ─────────────────────────────────────────────────────────────────────────── // Test helpers @@ -1335,10 +1335,12 @@ describe("TriggerChatTransport", () => { throw new Error(`Unexpected URL: ${urlStr}`); }); + const events: ChatTransportEvent[] = []; const transport = new TriggerChatTransport({ task: "my-chat-task", accessToken: () => "pat", watch: true, + onEvent: (e) => events.push(e), sessions: { "chat-watch-cancel": { publicAccessToken: "p", isStreaming: true } }, }); @@ -1351,6 +1353,10 @@ describe("TriggerChatTransport", () => { await reader.cancel(); await vi.advanceTimersByTimeAsync(10_000); expect(subscribeCount).toBe(countAtCancel); + // A clean cancel must not surface a spurious stream-error (an + // unguarded controller.close() after cancel would throw "Invalid + // state" and leak it onto the telemetry channel). + expect(events.some((e) => e.type === "stream-error")).toBe(false); } finally { vi.useRealTimers(); } diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index 9dbeab2aeb3..05242b96faa 100644 --- a/packages/trigger-sdk/src/v3/chat.ts +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -1864,7 +1864,11 @@ export class TriggerChatTransport implements ChatTransport { const opened = (await openWithAuthRetry()) ?? (await resumeAfterEof()); if (opened === null) { - controller.close(); + try { + controller.close(); + } catch { + /* already closed by a consumer cancel */ + } return; } reader = opened.reader; @@ -1894,7 +1898,11 @@ export class TriggerChatTransport implements ChatTransport { if (next.done) { const resumed = await resumeAfterEof(); if (resumed === null) { - controller.close(); + try { + controller.close(); + } catch { + /* already closed by a consumer cancel */ + } return; } reader = resumed.reader; @@ -1907,7 +1915,11 @@ export class TriggerChatTransport implements ChatTransport { if (combinedSignal.aborted) { internalAbort.abort(); await reader.cancel(); - controller.close(); + try { + controller.close(); + } catch { + /* already closed by a consumer cancel */ + } return; } From d9dc61eddd88242f44bba680034a7eb59b0fda1b Mon Sep 17 00:00:00 2001 From: Katia Bulatova Date: Tue, 11 Aug 2026 13:14:08 +0000 Subject: [PATCH 08/11] chore(changeset): consolidate the chat keepalive changesets into one --- .changeset/chat-truncated-turn-error.md | 7 ------- .changeset/watch-mode-keepalive.md | 4 +--- 2 files changed, 1 insertion(+), 10 deletions(-) delete mode 100644 .changeset/chat-truncated-turn-error.md diff --git a/.changeset/chat-truncated-turn-error.md b/.changeset/chat-truncated-turn-error.md deleted file mode 100644 index 959231a8fc8..00000000000 --- a/.changeset/chat-truncated-turn-error.md +++ /dev/null @@ -1,7 +0,0 @@ ---- -"@trigger.dev/sdk": patch ---- - -Chat in the browser now shows an error when a reply is cut off by a lost connection that can't be re-established, instead of presenting the partial reply as if it had finished. - -Watch-mode viewers now keep receiving later turns instead of the stream closing after the first turn completes. diff --git a/.changeset/watch-mode-keepalive.md b/.changeset/watch-mode-keepalive.md index e50a04f2a8c..ba6c0441724 100644 --- a/.changeset/watch-mode-keepalive.md +++ b/.changeset/watch-mode-keepalive.md @@ -2,6 +2,4 @@ "@trigger.dev/sdk": patch --- -Watch-mode chat subscriptions now stay connected across quiet periods. - -Read-only chat subscriptions no longer stop a turn when they disconnect. +Watch-mode chat streams now survive quiet windows and keep delivering later turns, and a reply cut off by a lost connection now shows an error instead of appearing finished. From b1814dc70f48e64a72b056838ed807f9c0c74595 Mon Sep 17 00:00:00 2001 From: Katia Bulatova Date: Tue, 11 Aug 2026 22:04:18 +0000 Subject: [PATCH 09/11] fix(sdk): keep the successor stream's abort controller when an aborted stream tears down --- .changeset/chat-stream-supersede-race.md | 5 + packages/trigger-sdk/src/v3/chat.test.ts | 135 +++++++++++++++++++++++ packages/trigger-sdk/src/v3/chat.ts | 9 +- 3 files changed, 147 insertions(+), 2 deletions(-) create mode 100644 .changeset/chat-stream-supersede-race.md diff --git a/.changeset/chat-stream-supersede-race.md b/.changeset/chat-stream-supersede-race.md new file mode 100644 index 00000000000..bc36a92d4de --- /dev/null +++ b/.changeset/chat-stream-supersede-race.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/sdk": patch +--- + +Fixed a race where quickly restarting a chat stream could break stop and reconnect for the new stream. diff --git a/packages/trigger-sdk/src/v3/chat.test.ts b/packages/trigger-sdk/src/v3/chat.test.ts index 9b0bfd598a8..a0c3be9e346 100644 --- a/packages/trigger-sdk/src/v3/chat.test.ts +++ b/packages/trigger-sdk/src/v3/chat.test.ts @@ -132,6 +132,35 @@ function defaultSseResponse( }); } +/** + * An SSE response whose body stays open until the request signal aborts. + * Models a live subscription sitting on a quiet server. + */ +function openSseResponse(signal?: AbortSignal | null): Response { + const body = new ReadableStream({ + start(controller) { + const onAbort = () => { + const err = new Error("aborted"); + err.name = "AbortError"; + try { + controller.error(err); + } catch { + /* already errored */ + } + }; + if (signal?.aborted) onAbort(); + else signal?.addEventListener("abort", onAbort, { once: true }); + }, + }); + return new Response(body, { + status: 200, + headers: { + "content-type": "text/event-stream", + "X-Stream-Version": "v2", + }, + }); +} + function authError(status = 401): Response { return new Response(JSON.stringify({ error: "Unauthorized", name: "TriggerApiError", status }), { status, @@ -1524,6 +1553,112 @@ describe("TriggerChatTransport", () => { }); }); + describe("superseded stream teardown", () => { + it("keeps the successor's controller registered when the aborted stream tears down", async () => { + vi.useFakeTimers(); + try { + let appendCount = 0; + global.fetch = vi.fn().mockImplementation(async (url: string | URL) => { + const urlStr = typeof url === "string" ? url : url.toString(); + if (isSessionStreamAppendUrl(urlStr)) { + appendCount++; + return defaultAppendResponse(); + } + // Quiet stream: EOF, no records, never settled — watch keeps it open. + if (isSessionOutSubscribeUrl(urlStr)) return defaultSseResponse([]); + throw new Error(`Unexpected URL: ${urlStr}`); + }); + + const transport = new TriggerChatTransport({ + task: "my-chat-task", + accessToken: () => "pat", + watch: true, + sessions: { "chat-race": { publicAccessToken: "p", isStreaming: true } }, + }); + + const send = () => + transport.sendMessages({ + trigger: "submit-message" as const, + chatId: "chat-race", + messageId: undefined, + messages: [createUserMessage("hi")], + abortSignal: undefined, + }); + + const first = drainChunks(await send()); + await vi.advanceTimersByTimeAsync(1_000); + + // Supersede: the new stream registers its controller synchronously, + // the aborted one tears down a microtask later. + const second = await send(); + let secondClosed = false; + const secondDrain = drainChunks(second).then(() => { + secondClosed = true; + }); + await first; + await vi.advanceTimersByTimeAsync(1_000); + + // stopGeneration posts the stop chunk either way — only the + // closing assertion proves it found the successor to abort. + appendCount = 0; + expect(await transport.stopGeneration("chat-race")).toBe(true); + await vi.advanceTimersByTimeAsync(1_000); + expect(appendCount).toBe(1); + expect(secondClosed).toBe(true); + + transport.dispose(); + await secondDrain; + } finally { + vi.useRealTimers(); + } + }); + + it("keeps the tab claim the successor took (multi-tab)", async () => { + vi.useFakeTimers(); + try { + global.fetch = vi.fn().mockImplementation(async (url: string | URL, init?: RequestInit) => { + const urlStr = typeof url === "string" ? url : url.toString(); + if (isSessionStreamAppendUrl(urlStr)) return defaultAppendResponse(); + // Open SSE that only ends when the subscription is aborted, so + // the superseded stream tears down while the successor is live. + if (isSessionOutSubscribeUrl(urlStr)) return openSseResponse(init?.signal); + throw new Error(`Unexpected URL: ${urlStr}`); + }); + + const transport = new TriggerChatTransport({ + task: "my-chat-task", + accessToken: () => "pat", + multiTab: true, + sessions: { "chat-race-tab": { publicAccessToken: "p", isStreaming: true } }, + }); + + const send = () => + transport.sendMessages({ + trigger: "submit-message" as const, + chatId: "chat-race-tab", + messageId: undefined, + messages: [createUserMessage("hi")], + abortSignal: undefined, + }); + + const first = drainChunks(await send()); + await vi.advanceTimersByTimeAsync(1_000); + const secondDrain = drainChunks(await send()); + await first; + await vi.advanceTimersByTimeAsync(1_000); + + // The superseded stream must not release the claim its successor + // holds — otherwise this tab flips to read-only mid-turn. + expect(transport.hasClaim("chat-race-tab")).toBe(true); + + transport.dispose(); + await secondDrain; + } finally { + vi.useRealTimers(); + } + }); + }); + describe("multi-tab coordination", () => { it("isReadOnly defaults to false when multiTab is disabled", () => { const transport = new TriggerChatTransport({ diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index 05242b96faa..4b0a9ef2268 100644 --- a/packages/trigger-sdk/src/v3/chat.ts +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -2054,8 +2054,13 @@ export class TriggerChatTransport implements ChatTransport { controller.error(error); } finally { teardownWakeListeners(); - this.activeStreams.delete(chatId); - this.coordinator?.release(chatId); + // Only clear the entry (and drop the tab claim) if it is still + // ours — a superseding send registers its controller before this + // teardown runs, and owns the claim from then on. + if (this.activeStreams.get(chatId) === internalAbort) { + this.activeStreams.delete(chatId); + this.coordinator?.release(chatId); + } } }, // A consumer that stops reading without aborting (drops the reader) From a026776f3fbda88d69e3bed74e5d837c20689d35 Mon Sep 17 00:00:00 2001 From: Katia Bulatova Date: Tue, 11 Aug 2026 22:36:47 +0000 Subject: [PATCH 10/11] fix(sdk): release the tab claim when the user stops generation --- .changeset/chat-stream-supersede-race.md | 2 +- packages/trigger-sdk/src/v3/chat.test.ts | 43 ++++++++++++++++++++++++ packages/trigger-sdk/src/v3/chat.ts | 5 +++ 3 files changed, 49 insertions(+), 1 deletion(-) diff --git a/.changeset/chat-stream-supersede-race.md b/.changeset/chat-stream-supersede-race.md index bc36a92d4de..7e2b807d8c3 100644 --- a/.changeset/chat-stream-supersede-race.md +++ b/.changeset/chat-stream-supersede-race.md @@ -2,4 +2,4 @@ "@trigger.dev/sdk": patch --- -Fixed a race where quickly restarting a chat stream could break stop and reconnect for the new stream. +Fixed a race where quickly restarting a chat stream could break stop and reconnect for the new stream. Stopping a chat now also hands it back to your other tabs instead of leaving them read-only. diff --git a/packages/trigger-sdk/src/v3/chat.test.ts b/packages/trigger-sdk/src/v3/chat.test.ts index a0c3be9e346..71566457dd8 100644 --- a/packages/trigger-sdk/src/v3/chat.test.ts +++ b/packages/trigger-sdk/src/v3/chat.test.ts @@ -1657,6 +1657,49 @@ describe("TriggerChatTransport", () => { vi.useRealTimers(); } }); + + it("releases the tab claim when the user stops generation (multi-tab)", async () => { + vi.useFakeTimers(); + try { + global.fetch = vi.fn().mockImplementation(async (url: string | URL, init?: RequestInit) => { + const urlStr = typeof url === "string" ? url : url.toString(); + if (isSessionStreamAppendUrl(urlStr)) return defaultAppendResponse(); + if (isSessionOutSubscribeUrl(urlStr)) return openSseResponse(init?.signal); + throw new Error(`Unexpected URL: ${urlStr}`); + }); + + const transport = new TriggerChatTransport({ + task: "my-chat-task", + accessToken: () => "pat", + multiTab: true, + sessions: { "chat-stop-tab": { publicAccessToken: "p", isStreaming: true } }, + }); + + const drain = drainChunks( + await transport.sendMessages({ + trigger: "submit-message" as const, + chatId: "chat-stop-tab", + messageId: undefined, + messages: [createUserMessage("hi")], + abortSignal: undefined, + }) + ); + await vi.advanceTimersByTimeAsync(1_000); + expect(transport.hasClaim("chat-stop-tab")).toBe(true); + + expect(await transport.stopGeneration("chat-stop-tab")).toBe(true); + await vi.advanceTimersByTimeAsync(1_000); + + // The turn ends here with no successor stream, so the claim must be + // freed or other tabs stay read-only until this one closes. + expect(transport.hasClaim("chat-stop-tab")).toBe(false); + + transport.dispose(); + await drain; + } finally { + vi.useRealTimers(); + } + }); }); describe("multi-tab coordination", () => { diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index 4b0a9ef2268..b7075e6f71e 100644 --- a/packages/trigger-sdk/src/v3/chat.ts +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -1218,6 +1218,11 @@ export class TriggerChatTransport implements ChatTransport { activeStream.abort(); this.activeStreams.delete(chatId); } + // Release here, not in the stream teardown: that only releases while it + // still owns the map entry, and we just deleted it. Unlike a supersede, + // no successor stream follows a stop, so the claim would never be freed + // and other tabs would stay read-only until this one closes. + this.coordinator?.release(chatId); // The turn won't reach its turn-complete on this client (we just // aborted the reader), so clear the streaming flag here and persist — From e2948f1a24a0b352d9ac73ff11fe15c37974829c Mon Sep 17 00:00:00 2001 From: Katia Bulatova Date: Wed, 12 Aug 2026 10:27:57 +0200 Subject: [PATCH 11/11] feat(webapp): query boundary pinned end-to-end and a capped query retry (#4549) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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](https://linear.app/triggerdotdev/issue/TRI-11165). ## Stack Stacked on **#4548** (watch-mode keepalive). Merge that first. ## What's inside - **A route-level read-only test** (`apps/webapp/test/queryRouteReadOnly.test.ts`) that drives `api.v1.query` with 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=1` made non-overridable** in `queryService.server.ts` — caller `clickhouseSettings` were spread after the defaults and could clear it. - **A per-turn query-retry cap** in the agent's `run_query` tool (`internal-packages/dashboard-agent/src/tool-api.ts`): three consecutive failures returns a terminal "stop and answer with what you have". ## Key decisions - **The read-only guarantee is grammar-level, not filtered.** TRQL has no write statements — they don't parse — ClickHouse runs with `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. - **The deny test runs through the route, not the parser.** A parser-only test would stay green if a refactor routed agent SQL around the compiler; driving the real route with a signed JWT pins the boundary end-to-end, and the deliberate positive read keeps the assertion honest. - **The retry cap lives per turn, not in the prompt.** A failed query hands the model the database error to fix, and usually it does — but the only other limit was the turn's 10 steps, so one query the model couldn't fix could eat the whole turn and leave the user with no answer. The tool set is built per turn, so the counter caps consecutive failures; a success resets it. The retry instruction rides the error text, so the prompt prefix is unchanged. ## 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. - The load-bearing guards were control-broken first (readonly override re-enabled; cap removed) and the tests went red. --- .server-changes/agent-message-quota.md | 6 + .server-changes/agent-watch-plan-limits.md | 6 + .../query-boundary-and-retry-cap.md | 6 + .server-changes/query-busy-retry.md | 6 + .../dashboard-agent/AgentUpgradeGate.tsx | 11 +- .../dashboard-agent/DashboardAgentChat.tsx | 33 +- .../dashboard-agent/DashboardAgentDraft.tsx | 61 +- .../dashboard-agent/DashboardAgentPanel.tsx | 13 + .../dashboard-agent/message-quota.test.ts | 40 +- .../dashboard-agent/message-quota.ts | 22 + .../dashboard-agent/useAgentMessageQuota.ts | 47 +- apps/webapp/app/routes/api.v1.query.ts | 12 +- ...aram.env.$envParam.dashboard-agent.in.$.ts | 28 + ...jectParam.env.$envParam.dashboard-agent.ts | 31 +- ...dashboardAgentInvestigationSweep.server.ts | 60 + .../services/dashboardAgentQuota.server.ts | 101 ++ .../dashboardAgentWatchErrorStatus.server.ts | 1 + .../dashboardAgentWatchLimits.server.ts | 59 + .../services/dashboardAgentWatches.server.ts | 42 + .../webapp/app/services/platform.v3.server.ts | 31 + .../app/services/queryService.server.ts | 1 + .../test/dashboardAgentDurableResume.test.ts | 347 +++++ .../dashboardAgentInvestigationPoison.test.ts | 166 ++ apps/webapp/test/dashboardAgentQuota.test.ts | 184 +++ .../dashboardAgentTenantIsolation.test.ts | 240 +++ .../dashboardAgentWatchLimitStatus.test.ts | 163 ++ .../test/dashboardAgentWatchLimits.test.ts | 411 +++++ apps/webapp/test/queryRouteReadOnly.test.ts | 181 +++ docs/v3-openapi.yaml | 10 + .../drizzle/0004_stale_corsair.sql | 8 + .../drizzle/0005_ambitious_mordo.sql | 2 + .../drizzle/meta/0004_snapshot.json | 1344 ++++++++++++++++ .../drizzle/meta/0005_snapshot.json | 1357 +++++++++++++++++ .../drizzle/meta/_journal.json | 14 + .../dashboard-agent-db/src/queries.ts | 73 +- .../dashboard-agent-db/src/schema.ts | 23 + .../dashboard-agent-db/src/watch-queries.ts | 16 + .../dashboard-agent/src/tool-api-client.ts | 16 +- .../src/tool-api-transport.test.ts | 39 + .../dashboard-agent/src/tool-api.ts | 24 +- .../src/tool-query-retry-cap.test.ts | 122 ++ 41 files changed, 5293 insertions(+), 64 deletions(-) create mode 100644 .server-changes/agent-message-quota.md create mode 100644 .server-changes/agent-watch-plan-limits.md create mode 100644 .server-changes/query-boundary-and-retry-cap.md create mode 100644 .server-changes/query-busy-retry.md create mode 100644 apps/webapp/app/services/dashboardAgentQuota.server.ts create mode 100644 apps/webapp/app/services/dashboardAgentWatchLimits.server.ts create mode 100644 apps/webapp/test/dashboardAgentDurableResume.test.ts create mode 100644 apps/webapp/test/dashboardAgentInvestigationPoison.test.ts create mode 100644 apps/webapp/test/dashboardAgentQuota.test.ts create mode 100644 apps/webapp/test/dashboardAgentTenantIsolation.test.ts create mode 100644 apps/webapp/test/dashboardAgentWatchLimitStatus.test.ts create mode 100644 apps/webapp/test/dashboardAgentWatchLimits.test.ts create mode 100644 apps/webapp/test/queryRouteReadOnly.test.ts create mode 100644 internal-packages/dashboard-agent-db/drizzle/0004_stale_corsair.sql create mode 100644 internal-packages/dashboard-agent-db/drizzle/0005_ambitious_mordo.sql create mode 100644 internal-packages/dashboard-agent-db/drizzle/meta/0004_snapshot.json create mode 100644 internal-packages/dashboard-agent-db/drizzle/meta/0005_snapshot.json create mode 100644 internal-packages/dashboard-agent/src/tool-query-retry-cap.test.ts diff --git a/.server-changes/agent-message-quota.md b/.server-changes/agent-message-quota.md new file mode 100644 index 00000000000..296b082de79 --- /dev/null +++ b/.server-changes/agent-message-quota.md @@ -0,0 +1,6 @@ +--- +area: webapp +type: feature +--- + +The dashboard agent now comes with a monthly message allowance. A message that fails to send doesn't count against it. diff --git a/.server-changes/agent-watch-plan-limits.md b/.server-changes/agent-watch-plan-limits.md new file mode 100644 index 00000000000..d323e7194f0 --- /dev/null +++ b/.server-changes/agent-watch-plan-limits.md @@ -0,0 +1,6 @@ +--- +area: webapp +type: feature +--- + +Watches now respect your plan's limits: free plans can run a limited number of watches at once and for a shorter window, with a prompt to upgrade for more. diff --git a/.server-changes/query-boundary-and-retry-cap.md b/.server-changes/query-boundary-and-retry-cap.md new file mode 100644 index 00000000000..b1191d2746c --- /dev/null +++ b/.server-changes/query-boundary-and-retry-cap.md @@ -0,0 +1,6 @@ +--- +area: webapp +type: improvement +--- + +Queries stay read-only, and the agent now stops after a few failed queries in a row and answers with what it found instead of spending the whole reply retrying. diff --git a/.server-changes/query-busy-retry.md b/.server-changes/query-busy-retry.md new file mode 100644 index 00000000000..c7e8d11d37b --- /dev/null +++ b/.server-changes/query-busy-retry.md @@ -0,0 +1,6 @@ +--- +area: webapp +type: fix +--- + +When queries are queued up and one is turned away, you now get a clear "try again shortly" instead of an error that looks like a problem with the query itself. diff --git a/apps/webapp/app/components/dashboard-agent/AgentUpgradeGate.tsx b/apps/webapp/app/components/dashboard-agent/AgentUpgradeGate.tsx index 4b795edde9e..498a0431579 100644 --- a/apps/webapp/app/components/dashboard-agent/AgentUpgradeGate.tsx +++ b/apps/webapp/app/components/dashboard-agent/AgentUpgradeGate.tsx @@ -1,9 +1,10 @@ import { Link } from "@remix-run/react"; +import { AgentMonoLogo } from "~/components/primitives/AgentDotMatrix"; import { LinkButton } from "~/components/primitives/Buttons"; import { useOrganization } from "~/hooks/useOrganizations"; -import { cn } from "~/utils/cn"; import { v3BillingPath } from "~/utils/pathBuilder"; -import { AgentIcon, AGENT_ICON_ACCENT_CLASS, ASK_AGENT_LABEL } from "./agent-identity"; +import { ASK_AGENT_LABEL } from "./agent-identity"; +import { messageQuotaReachedCopy } from "./message-quota"; // Matches the composer's outer geometry so the replacement lands in the same place. const SLOT = "flex shrink-0 flex-col bg-background-bright px-3 pb-3 pt-1"; @@ -22,14 +23,12 @@ export function AgentUpgradeBlock({ {context}
- + Upgrade to unlock {ASK_AGENT_LABEL}
-

- You've used all {limit} messages included on the Free plan. Your chats stay here to read. -

+

{messageQuotaReachedCopy(limit)}

Upgrade diff --git a/apps/webapp/app/components/dashboard-agent/DashboardAgentChat.tsx b/apps/webapp/app/components/dashboard-agent/DashboardAgentChat.tsx index 19e65a570ca..aaecd89e404 100644 --- a/apps/webapp/app/components/dashboard-agent/DashboardAgentChat.tsx +++ b/apps/webapp/app/components/dashboard-agent/DashboardAgentChat.tsx @@ -17,6 +17,7 @@ import { DashboardAgentContextBanner } from "./DashboardAgentContextBanner"; import { DashboardAgentHero } from "./DashboardAgentHero"; import { DashboardAgentMessages, type TurnActivity } from "./DashboardAgentMessages"; import { MESSAGE_TOO_LARGE_ERROR } from "./message-limits"; +import { FREE_PLAN_MESSAGE_LIMIT, parseQuotaReachedResponse } from "./message-quota"; import { createTranscriptOrder, orderTranscript } from "./message-order"; import { navigateDestination } from "./navigate-target"; import { pendingNavigateIntents, pendingWatchIntents } from "./pending-intents"; @@ -102,6 +103,9 @@ export function DashboardAgentChat({ onActivityChange?: (chatId: string, activity: TurnActivity | null) => void; }) { const [input, setInput] = useState(""); + // Set when the server refuses a send over the cap, so the block shows at once rather than + // waiting for the next quota poll. + const [quotaReached, setQuotaReached] = useState<{ limit: number } | null>(null); const navigate = useNavigate(); const location = useLocation(); const toast = useToast(); @@ -128,6 +132,18 @@ export function DashboardAgentChat({ .catch(() => null)) as { error?: string } | null; throw new Error(data?.error ?? MESSAGE_TOO_LARGE_ERROR); } + // Over the message cap: show the upgrade block instead of a generic turn error. + if (res.status === 403) { + const data = (await res + .clone() + .json() + .catch(() => null)) as { error?: string; limit?: number } | null; + const reached = parseQuotaReachedResponse(res.status, data); + if (reached) { + setQuotaReached(reached); + throw new Error("You've reached your message limit."); + } + } return res; }, clientData, @@ -185,9 +201,12 @@ export function DashboardAgentChat({ const orderRef = useRef(createTranscriptOrder(initialMessages)); const messages = orderTranscript(rawMessages, orderRef.current); - // Counted here, not in the panel, so it includes the turn just sent. - const quota = useAgentMessageQuota({ actionPath, chatId, messages }); - const atMessageCap = quota.kind === "reached"; + // Read here, not in the panel, so it re-reads as each turn settles. + const quota = useAgentMessageQuota({ actionPath, chatId, status }); + // Either the poll saw the cap, or a send was just refused over it. + const atMessageCap = quota.kind === "reached" || quotaReached !== null; + const messageCapLimit = + quotaReached?.limit ?? (quota.kind === "reached" ? quota.limit : FREE_PLAN_MESSAGE_LIMIT); const isStreaming = status === "streaming"; // From status, not the last part: the indicator must stay up through silent tool calls. @@ -252,6 +271,8 @@ export function DashboardAgentChat({ }, [sendRequest, submit, canSend]); const retry = useCallback(() => { + // Over the cap, a retry only earns another 403 — same guard as `submit`. + if (atMessageCap) return; // A watch's consent record is a user message nobody typed, so retry never treats it as one. const action = retryAction( messages.filter((m) => !(m.role === "user" && isWatchRequestMessageId(m.id))) @@ -264,7 +285,7 @@ export function DashboardAgentChat({ return; } void sendMessage({ text: action.text, messageId: action.messageId }); - }, [messages, sendMessage, regenerate, clearError]); + }, [messages, sendMessage, regenerate, clearError, atMessageCap]); const resolveUri = useTriggerUriResolver(actionPath); @@ -414,9 +435,9 @@ export function DashboardAgentChat({ /> )} {watchCard ?
{watchCard}
: null} - {quota.kind === "reached" ? ( + {atMessageCap ? ( void; projectSlug: string; @@ -24,6 +26,7 @@ export function DashboardAgentDraft({ pageContext?: AgentPageContext; promotedPrompt?: SuggestedPrompt; watchCard?: React.ReactNode; + capReached?: { limit: number } | null; }) { const [input, setInput] = useState(""); @@ -43,12 +46,14 @@ export function DashboardAgentDraft({ const submit = useCallback( (text: string) => { + // Suggested prompts reach here via the hero, bypassing the composer's cap guard. + if (capReached) return; const trimmed = text.trim(); if (!trimmed) return; setInput(""); onSubmit(trimmed); }, - [onSubmit] + [onSubmit, capReached] ); return ( @@ -57,25 +62,41 @@ export function DashboardAgentDraft({ pageContext={pageContext} promoted={promotedPrompt} composer={ -
- {watchCard} - submit(input)} - onStop={() => {}} - isStreaming={false} - placeholderSuggestion={watchCard ? undefined : placeholderSuggestion} - context={ - - } - /> -
+ capReached ? ( +
+ {watchCard} + + } + /> +
+ ) : ( +
+ {watchCard} + submit(input)} + onStop={() => {}} + isStreaming={false} + placeholderSuggestion={watchCard ? undefined : placeholderSuggestion} + context={ + + } + /> +
+ ) } /> ); diff --git a/apps/webapp/app/components/dashboard-agent/DashboardAgentPanel.tsx b/apps/webapp/app/components/dashboard-agent/DashboardAgentPanel.tsx index 894a80e8e61..2a85325c1d6 100644 --- a/apps/webapp/app/components/dashboard-agent/DashboardAgentPanel.tsx +++ b/apps/webapp/app/components/dashboard-agent/DashboardAgentPanel.tsx @@ -24,6 +24,7 @@ import { writeLastChat, } from "./last-chat-storage"; import { DashboardAgentDraft } from "./DashboardAgentDraft"; +import { parseQuotaReachedResponse } from "./message-quota"; import { WatchCard } from "./WatchCard"; import { watchDraftFor } from "./watch-card"; import { NO_WATCH_CARD, watchCardReducer } from "./watch-card-state"; @@ -114,6 +115,8 @@ export function DashboardAgentPanel({ // Until the list has arrived, the page load's server count is the better answer. const [chatsLoaded, setChatsLoaded] = useState(false); const [active, setActive] = useState(null); + // A refused `create` over the cap: the draft shows the upgrade block instead of a raw toast. + const [capReached, setCapReached] = useState<{ limit: number } | null>(null); // Starts true so an `openWith` request waits for the restore instead of racing it. const [loading, setLoading] = useState( () => readLastChat(storageKey)?.path === location.pathname @@ -260,14 +263,22 @@ export function DashboardAgentPanel({ publicAccessToken?: string; headStarted?: boolean; error?: string; + limit?: number; }; if (seq !== openChatRequestSeq.current) return; if (!res.ok || !data.chatId || !data.publicAccessToken) { + const reached = parseQuotaReachedResponse(res.status, data); + if (reached) { + setCapReached(reached); + setActive(null); + return; + } console.error(`Dashboard agent: failed to create chat (${res.status})`, data.error); toast.error(data.error ?? "We couldn't start that chat. Try again in a moment."); setActive(null); return; } + setCapReached(null); setActive({ chatId: data.chatId, organizationId: organization.id, @@ -309,6 +320,7 @@ export function DashboardAgentPanel({ panelOrg.current = organization.id; claimChatSlot(); setActive(null); + setCapReached(null); setLoading(false); setChats([]); setChatsLoaded(false); @@ -634,6 +646,7 @@ export function DashboardAgentPanel({ pageContext={pageContext} promotedPrompt={promotedPrompt} watchCard={watchCardElement} + capReached={capReached} /> )} diff --git a/apps/webapp/app/components/dashboard-agent/message-quota.test.ts b/apps/webapp/app/components/dashboard-agent/message-quota.test.ts index 5a252db7ac1..3b47395e99a 100644 --- a/apps/webapp/app/components/dashboard-agent/message-quota.test.ts +++ b/apps/webapp/app/components/dashboard-agent/message-quota.test.ts @@ -1,5 +1,12 @@ import { describe, expect, it } from "vitest"; -import { countUserMessages, FREE_PLAN_MESSAGE_LIMIT, resolveMessageQuota } from "./message-quota"; +import { + countUserMessages, + FREE_PLAN_MESSAGE_LIMIT, + MESSAGE_QUOTA_REACHED_ERROR, + messageQuotaReachedCopy, + parseQuotaReachedResponse, + resolveMessageQuota, +} from "./message-quota"; describe("resolveMessageQuota", () => { it("caps a Free plan at the limit", () => { @@ -42,6 +49,37 @@ describe("resolveMessageQuota", () => { }); }); +describe("parseQuotaReachedResponse", () => { + it("maps a create/in 403 cap body to the limit", () => { + // Both the create path and the `in` transport refuse with this exact body. + expect( + parseQuotaReachedResponse(403, { error: MESSAGE_QUOTA_REACHED_ERROR, limit: 20 }) + ).toEqual({ limit: 20 }); + }); + + it("falls back to the free limit when the body omits it", () => { + expect(parseQuotaReachedResponse(403, { error: MESSAGE_QUOTA_REACHED_ERROR })).toEqual({ + limit: FREE_PLAN_MESSAGE_LIMIT, + }); + }); + + it("ignores other errors and non-403 statuses so they surface normally", () => { + expect(parseQuotaReachedResponse(403, { error: "something_else" })).toBeNull(); + expect(parseQuotaReachedResponse(500, { error: MESSAGE_QUOTA_REACHED_ERROR })).toBeNull(); + expect(parseQuotaReachedResponse(403, null)).toBeNull(); + }); +}); + +describe("messageQuotaReachedCopy", () => { + it("is a friendly sentence naming the limit, never the raw code", () => { + const copy = messageQuotaReachedCopy(20); + expect(copy).toContain("all 20 messages"); + expect(copy).toContain("Free plan"); + // Control break: if the mapping leaked the server code, this fails. + expect(copy).not.toContain(MESSAGE_QUOTA_REACHED_ERROR); + }); +}); + describe("countUserMessages", () => { it("counts only what the user sent", () => { expect( diff --git a/apps/webapp/app/components/dashboard-agent/message-quota.ts b/apps/webapp/app/components/dashboard-agent/message-quota.ts index f65481c8705..c603a836a2c 100644 --- a/apps/webapp/app/components/dashboard-agent/message-quota.ts +++ b/apps/webapp/app/components/dashboard-agent/message-quota.ts @@ -27,6 +27,28 @@ export function resolveMessageQuota({ : { kind: "within", used, limit, remaining }; } +// The server code both the create and `in` paths refuse with. The client owns the copy, +// so this code must never reach the UI as text. +export const MESSAGE_QUOTA_REACHED_ERROR = "message_quota_reached"; + +// Maps a 403 refusal body to the cap signal, or null for any other error. Both paths use +// this so a `message_quota_reached` code routes to the upgrade block, never a raw toast. +export function parseQuotaReachedResponse( + status: number, + data: { error?: string; limit?: number } | null | undefined +): { limit: number } | null { + if (status === 403 && data?.error === MESSAGE_QUOTA_REACHED_ERROR) { + return { limit: data.limit ?? FREE_PLAN_MESSAGE_LIMIT }; + } + return null; +} + +// The upgrade block's sentence. Pure so the copy is asserted directly, and so the raw +// server code can never be what the user reads. +export function messageQuotaReachedCopy(limit: number): string { + return `You've used all ${limit} messages included on the Free plan. Your chats stay here to read.`; +} + // A watch's consent record is a user message the person never typed, so it is // excluded here exactly as the stored count excludes it. export function countUserMessages(messages: { role: string; id?: string }[]): number { diff --git a/apps/webapp/app/components/dashboard-agent/useAgentMessageQuota.ts b/apps/webapp/app/components/dashboard-agent/useAgentMessageQuota.ts index da6cc3d83ba..958d8dfcc14 100644 --- a/apps/webapp/app/components/dashboard-agent/useAgentMessageQuota.ts +++ b/apps/webapp/app/components/dashboard-agent/useAgentMessageQuota.ts @@ -1,46 +1,55 @@ -import type { UIMessage } from "@ai-sdk/react"; -import { useEffect, useState } from "react"; -import { countUserMessages, resolveMessageQuota, type MessageQuota } from "./message-quota"; +import { useEffect, useRef, useState } from "react"; +import { useCurrentPlan } from "~/routes/_app.orgs.$organizationSlug/route"; +import { resolveMessageQuota, type MessageQuota } from "./message-quota"; -// Always undefined until billing supplies plan detection, which means no cap. +// Gated on billing PRESENCE, not the plan value: no subscription means billing isn't wired +// up (self-hosted), so there is no cap and no upgrade UI. A wired-up, non-paying plan is free. function useIsFreePlan(): boolean | undefined { - return undefined; + const subscription = useCurrentPlan()?.v3Subscription; + if (!subscription) return undefined; + return subscription.isPaying === false; } -// Counted in two halves: the server aggregates other chats, this chat's own count -// comes from the live transcript so the message just sent counts immediately. +// `used` is the server's per-period count for the org. Re-read once a turn settles — the +// server increment happens mid-turn in the `.in` proxy, so reading on optimistic append +// would lag the count by one message and show the cap a message late. export function useAgentMessageQuota({ actionPath, chatId, - messages, + status, }: { actionPath: string; chatId: string; - messages: UIMessage[]; + status: string; }): MessageQuota { const isFreePlan = useIsFreePlan(); - const [usedElsewhere, setUsedElsewhere] = useState(undefined); + const [used, setUsed] = useState(undefined); + + // Bumped each time the status leaves streaming/submitted, which drives the re-read. + const [settleTick, setSettleTick] = useState(0); + const prevStatus = useRef(status); + useEffect(() => { + const wasInFlight = prevStatus.current === "streaming" || prevStatus.current === "submitted"; + const nowSettled = status === "ready" || status === "error"; + prevStatus.current = status; + if (wasInFlight && nowSettled) setSettleTick((tick) => tick + 1); + }, [status]); useEffect(() => { if (isFreePlan !== true) return; const controller = new AbortController(); void (async () => { try { - const res = await fetch(`${actionPath}?quota=1&chatId=${encodeURIComponent(chatId)}`, { - signal: controller.signal, - }); + const res = await fetch(`${actionPath}?quota=1`, { signal: controller.signal }); if (!res.ok) return; const data = (await res.json()) as { used?: number }; - if (typeof data.used === "number") setUsedElsewhere(data.used); + if (typeof data.used === "number") setUsed(data.used); } catch { // Leave the count unknown, which means no cap. See `resolveMessageQuota`. } })(); return () => controller.abort(); - }, [isFreePlan, actionPath, chatId]); + }, [isFreePlan, actionPath, chatId, settleTick]); - return resolveMessageQuota({ - isFreePlan, - used: usedElsewhere === undefined ? undefined : usedElsewhere + countUserMessages(messages), - }); + return resolveMessageQuota({ isFreePlan, used }); } diff --git a/apps/webapp/app/routes/api.v1.query.ts b/apps/webapp/app/routes/api.v1.query.ts index adad14859db..5ff7f56cfec 100644 --- a/apps/webapp/app/routes/api.v1.query.ts +++ b/apps/webapp/app/routes/api.v1.query.ts @@ -2,7 +2,11 @@ import { json } from "@remix-run/server-runtime"; import { QueryError } from "@internal/clickhouse"; import { z } from "zod"; import { createActionApiRoute, everyResource } from "~/services/routeBuilders/apiBuilder.server"; -import { executeQuery, type QueryScope } from "~/services/queryService.server"; +import { + executeQuery, + isQueryConcurrencyRejection, + type QueryScope, +} from "~/services/queryService.server"; import { logger } from "~/services/logger.server"; import { rowsToCSV } from "~/utils/dataExport"; import { detectQueryTables } from "~/v3/detectQueryTables"; @@ -78,6 +82,12 @@ const { action, loader } = createActionApiRoute( }); if (!queryResult.success) { + // A concurrency rejection is "too busy", not a bad query: 429 so callers retry it + // instead of rewriting a query that was fine. + if (isQueryConcurrencyRejection(queryResult.error)) { + return json({ error: queryResult.error.message }, { status: 429 }); + } + // QueryError surfaces customer SQL problems (invalid syntax, // unsupported construct). Returned to the caller as 400; system // handles it gracefully, no alert needed. diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.in.$.ts b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.in.$.ts index 5939e836db1..655da1d9aa7 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.in.$.ts +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.in.$.ts @@ -7,6 +7,7 @@ import { MESSAGE_TOO_LARGE_CODE, MESSAGE_TOO_LARGE_ERROR, } from "~/components/dashboard-agent/message-limits"; +import { MESSAGE_QUOTA_REACHED_ERROR } from "~/components/dashboard-agent/message-quota"; import { findProjectBySlug } from "~/models/project.server"; import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server"; import { @@ -15,6 +16,12 @@ import { resolveDashboardAgentRepoSnapshot, } from "~/services/dashboardAgent.server"; import { dashboardAgentEnvironmentAddress } from "~/services/dashboardAgentEnvironmentAddress.server"; +import { dashboardAgentDb } from "~/services/dashboardAgentDb.server"; +import { + agentTurnCountsAgainstQuota, + recordAgentMessageSent, + resolveAgentMessageQuota, +} from "~/services/dashboardAgentQuota.server"; import { logger } from "~/services/logger.server"; import { requireUser } from "~/services/session.server"; import { readBoundedBodyText } from "~/utils/boundedRequestBody.server"; @@ -115,6 +122,9 @@ export async function action({ request, params }: ActionFunctionArgs) { parsed = undefined; } + // Hoisted so it is visible after the fetch: quota is charged only once the send succeeds. + let countsAgainstQuota = false; + if (parsed) { // Actions are placed by the server only, and this proxy is the one path a browser // can reach `.in` through. @@ -127,6 +137,17 @@ export async function action({ request, params }: ActionFunctionArgs) { return tooLarge(); } + // Only a real user message consumes quota; action turns were refused above. + countsAgainstQuota = agentTurnCountsAgainstQuota(parsed); + if (countsAgainstQuota) { + const quota = await resolveAgentMessageQuota(dashboardAgentDb, { + organizationId: project.organizationId, + }); + if (quota?.reached) { + return json({ error: MESSAGE_QUOTA_REACHED_ERROR, limit: quota.limit }, { status: 403 }); + } + } + let userActorToken: string; try { userActorToken = await mintDashboardAgentUserActorToken(user.id, { @@ -165,6 +186,13 @@ export async function action({ request, params }: ActionFunctionArgs) { try { const upstream = await fetch(upstreamUrl, { method: "POST", headers, body }); const text = await upstream.text(); + // Charge quota only for a delivered message: a non-2xx upstream (or a throw below) + // must not burn a send that never reached the agent. + if (countsAgainstQuota && upstream.ok) { + await recordAgentMessageSent(dashboardAgentDb, { + organizationId: project.organizationId, + }); + } return new Response(text, { status: upstream.status, headers: { "content-type": upstream.headers.get("content-type") ?? "application/json" }, diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.ts b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.ts index d808c567548..60e8c2999a4 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.ts +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.ts @@ -3,7 +3,7 @@ import { chatExists, countUnreadWatchWakes, countChatsWithUnreadWork, - countUserMessages, + getAgentMessageUsage, createChat, getChatMessages, getSession, @@ -28,6 +28,7 @@ import { MESSAGE_TOO_LARGE_CODE, MESSAGE_TOO_LARGE_ERROR, } from "~/components/dashboard-agent/message-limits"; +import { MESSAGE_QUOTA_REACHED_ERROR } from "~/components/dashboard-agent/message-quota"; import { MAX_URIS_PER_RESOLVE_REQUEST } from "~/components/dashboard-agent/resolve-uris"; import { $replica } from "~/db.server"; import { env } from "~/env.server"; @@ -53,6 +54,11 @@ import { dashboardAgentEnvironmentAddress } from "~/services/dashboardAgentEnvir import { watchErrorStatus } from "~/services/dashboardAgentWatchErrorStatus.server"; import { startDashboardAgentHeadStart } from "~/services/dashboardAgentHeadStart.server"; import { dashboardAgentDb } from "~/services/dashboardAgentDb.server"; +import { + currentAgentMessagePeriod, + recordAgentMessageSent, + resolveAgentMessageQuota, +} from "~/services/dashboardAgentQuota.server"; import { logger } from "~/services/logger.server"; import { resolveTriggerUri } from "~/services/resolveTriggerUri.server"; import { requireUser } from "~/services/session.server"; @@ -151,13 +157,11 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { const project = await findProjectBySlug(organizationSlug, projectParam, userId); if (!project) return json({ error: "Project not found" }, { status: 404 }); - // The open chat is excluded and counted from the live transcript instead, so an - // unpersisted turn still counts against the cap. + // The per-period counter, org-wide: a deleted chat can't lower it within the period. if (searchParams.get("quota") === "1") { - const used = await countUserMessages(dashboardAgentDb, { + const used = await getAgentMessageUsage(dashboardAgentDb, { organizationId: project.organizationId, - userId, - excludeChatId: searchParams.get("chatId") ?? undefined, + period: currentAgentMessagePeriod(), }); return json({ used }); } @@ -291,6 +295,13 @@ export const action = async ({ request, params }: ActionFunctionArgs) => { return messageTooLarge(); } + const quota = await resolveAgentMessageQuota(dashboardAgentDb, { + organizationId: project.organizationId, + }); + if (quota?.reached) { + return json({ error: MESSAGE_QUOTA_REACHED_ERROR, limit: quota.limit }, { status: 403 }); + } + let clientData: Record | undefined; try { clientData = parsed.data.clientData @@ -388,6 +399,14 @@ export const action = async ({ request, params }: ActionFunctionArgs) => { throw error; } + // Only the head start dispatches the first message here; a cold start sends it through + // the `in` proxy, which counts it there. Counting both would double-count. + if (headStarted) { + await recordAgentMessageSent(dashboardAgentDb, { + organizationId: project.organizationId, + }); + } + let publicAccessToken: string; try { publicAccessToken = await mintDashboardAgentToken(chatId); diff --git a/apps/webapp/app/services/dashboardAgentInvestigationSweep.server.ts b/apps/webapp/app/services/dashboardAgentInvestigationSweep.server.ts index b10ee6b06af..38853ea0ddd 100644 --- a/apps/webapp/app/services/dashboardAgentInvestigationSweep.server.ts +++ b/apps/webapp/app/services/dashboardAgentInvestigationSweep.server.ts @@ -5,8 +5,11 @@ import { listStaleOpenInvestigations, + recordInvestigationSweepAttempt, settleInvestigationAndCloseCard, + settleInvestigationAsInconclusive, type Investigation, + type SettledInvestigation, type SettledInvestigationCard, } from "@internal/dashboard-agent-db"; import { UNSETTLED_INVESTIGATION_NOTE } from "@internal/dashboard-agent-contracts"; @@ -22,6 +25,13 @@ export const INVESTIGATION_STALE_MS = 30 * 60 * 1000; /** Per-run cap. Oldest first, so the rest land next run. */ const SWEEP_BATCH_LIMIT = 100; +/** + * After this many failed settle attempts a row is force-abandoned: settled `inconclusive` + * WITHOUT the closing card, so a card that never renders leaves the queue instead of + * looping forever. The rare stuck spinner is the price of not starving every other row. + */ +export const MAX_SWEEP_ATTEMPTS = 5; + export type InvestigationSweepResult = { /** Stale `in_progress` rows seen. */ stale: number; @@ -30,6 +40,8 @@ export type InvestigationSweepResult = { closed: number; /** A turn (or another sweep) settled it first. */ alreadySettled: number; + /** Rows past the attempt cap, force-settled without a card so they leave the queue. */ + abandoned: number; failed: number; }; @@ -46,6 +58,10 @@ export type InvestigationSweepDeps = { chatId: string; note: string; }) => Promise; + /** Record a failed settle out-of-band; returns the new attempt count, or null if gone. */ + recordAttempt?: (params: { id: string }) => Promise; + /** Force a poison row terminal without the failing render path. */ + forceAbandon?: (params: { id: string; note: string }) => Promise; }; /** @@ -61,12 +77,17 @@ export async function sweepDashboardAgentInvestigations( deps.listStale ?? ((params) => listStaleOpenInvestigations(dashboardAgentDb, params)); const settleAndClose = deps.settleAndClose ?? ((params) => settleInvestigationAndCloseCard(dashboardAgentDb, params)); + const recordAttempt = + deps.recordAttempt ?? ((params) => recordInvestigationSweepAttempt(dashboardAgentDb, params)); + const forceAbandon = + deps.forceAbandon ?? ((params) => settleInvestigationAsInconclusive(dashboardAgentDb, params)); const result: InvestigationSweepResult = { stale: 0, settled: 0, closed: 0, alreadySettled: 0, + abandoned: 0, failed: 0, }; @@ -93,10 +114,49 @@ export async function sweepDashboardAgentInvestigations( result.settled++; if (outcome.closed) result.closed++; } catch (error) { + // The settle rolled back, so the row is still `in_progress`. Record the attempt in + // its own write — this rotates the row to the back of the sweep order (see + // `listStaleOpenInvestigations`) so it can't pin the head and starve newer rows. + let attempts: number | null = null; + try { + attempts = await recordAttempt({ id: investigation.id }); + } catch (recordError) { + logger.error("Dashboard agent investigation sweep: failed to record a sweep attempt", { + investigationId: investigation.id, + chatId: investigation.chatId, + error: recordError, + }); + } + + // Past the cap the card will never render; force it terminal without the render + // path so it leaves the queue instead of looping forever. + if (attempts !== null && attempts >= MAX_SWEEP_ATTEMPTS) { + try { + await forceAbandon({ id: investigation.id, note: UNSETTLED_INVESTIGATION_NOTE }); + result.abandoned++; + logger.warn( + "Dashboard agent investigation sweep: abandoned a card past the attempt cap", + { + investigationId: investigation.id, + chatId: investigation.chatId, + attempts, + } + ); + continue; + } catch (abandonError) { + logger.error("Dashboard agent investigation sweep: failed to abandon a poison card", { + investigationId: investigation.id, + chatId: investigation.chatId, + error: abandonError, + }); + } + } + result.failed++; logger.error("Dashboard agent investigation sweep: failed to settle an investigation", { investigationId: investigation.id, chatId: investigation.chatId, + attempts, error, }); } diff --git a/apps/webapp/app/services/dashboardAgentQuota.server.ts b/apps/webapp/app/services/dashboardAgentQuota.server.ts new file mode 100644 index 00000000000..4bfca8808bd --- /dev/null +++ b/apps/webapp/app/services/dashboardAgentQuota.server.ts @@ -0,0 +1,101 @@ +import type { Limits } from "@trigger.dev/platform"; +import { + getAgentMessageUsage, + incrementAgentMessageUsage, + type DashboardAgentDb, +} from "@internal/dashboard-agent-db"; +import { getCachedLimit } from "./platform.v3.server"; +import { logger } from "./logger.server"; + +// The repo's unlimited sentinel. Never Infinity: it serializes to null in the limit cache. +export const UNLIMITED_AGENT_MESSAGES = 100_000_000; + +// Filled by cloud billing (TRI-12863 P0). Absent until then, and always on self-hosted, +// so the fallback applies and the cap is effectively off. +const AGENT_MESSAGE_LIMIT_KEY = "agentMessages" as keyof Limits; + +/** The billing period the counter is scoped to: a UTC calendar month, "YYYY-MM". */ +export function currentAgentMessagePeriod(now: Date = new Date()): string { + return `${now.getUTCFullYear()}-${String(now.getUTCMonth() + 1).padStart(2, "0")}`; +} + +/** Pure so the send routes and, later, the MCP path share one rule. */ +export function checkAgentMessageQuota({ used, limit }: { used: number; limit: number }): { + reached: boolean; +} { + return { reached: used >= limit }; +} + +export type AgentMessageQuota = { reached: boolean; used: number; limit: number }; + +/** + * The period counter and the cached plan limit for one org. Fails open: an absent limit + * (self-hosted, or before the cloud side ships) resolves to the unlimited sentinel, and a + * counter read that throws returns `undefined` — either way there is no cap. + */ +export async function resolveAgentMessageQuota( + db: DashboardAgentDb, + params: { + organizationId: string; + now?: Date; + readLimit?: (organizationId: string) => Promise; + } +): Promise { + const readLimit = + params.readLimit ?? + (async (organizationId: string) => { + const cached = await getCachedLimit( + organizationId, + AGENT_MESSAGE_LIMIT_KEY, + UNLIMITED_AGENT_MESSAGES + ); + // A cache error leaves `val` empty; fall open to unlimited. + return cached.val ?? UNLIMITED_AGENT_MESSAGES; + }); + try { + const [limit, used] = await Promise.all([ + readLimit(params.organizationId), + getAgentMessageUsage(db, { + organizationId: params.organizationId, + period: currentAgentMessagePeriod(params.now), + }), + ]); + return { ...checkAgentMessageQuota({ used, limit }), used, limit }; + } catch (error) { + logger.error("Failed to resolve dashboard agent message quota", { + organizationId: params.organizationId, + error, + }); + return undefined; + } +} + +/** Record one sent user message. Swallows errors: the cap is a nudge, never a send blocker. */ +export async function recordAgentMessageSent( + db: DashboardAgentDb, + params: { organizationId: string; now?: Date } +): Promise { + try { + await incrementAgentMessageUsage(db, { + organizationId: params.organizationId, + period: currentAgentMessagePeriod(params.now), + }); + } catch (error) { + logger.error("Failed to record a dashboard agent message against the quota", { + organizationId: params.organizationId, + error, + }); + } +} + +/** + * Whether an agent turn consumes quota. Only a genuine new user message counts: the transport + * tags it `trigger: "submit-message"`. A retry/regenerate re-runs the agent from its own history + * without a new message (`trigger: "regenerate-message"`), and a wake is `"action"` — neither is + * something the user typed, so neither counts. + */ +export function agentTurnCountsAgainstQuota( + turn: { kind?: string; payload?: { trigger?: string } } | undefined +): boolean { + return turn?.kind === "message" && turn.payload?.trigger === "submit-message"; +} diff --git a/apps/webapp/app/services/dashboardAgentWatchErrorStatus.server.ts b/apps/webapp/app/services/dashboardAgentWatchErrorStatus.server.ts index de32c3595cf..9e17c567a96 100644 --- a/apps/webapp/app/services/dashboardAgentWatchErrorStatus.server.ts +++ b/apps/webapp/app/services/dashboardAgentWatchErrorStatus.server.ts @@ -6,6 +6,7 @@ import type { SubmitWatchErrorCode } from "./dashboardAgentWatches.server"; */ const STATUS_BY_CODE: Record = { limit_reached: 409, + watch_limit_reached: 409, duplicate: 409, request_conflict: 409, invalid_target: 404, diff --git a/apps/webapp/app/services/dashboardAgentWatchLimits.server.ts b/apps/webapp/app/services/dashboardAgentWatchLimits.server.ts new file mode 100644 index 00000000000..4c96196ebdd --- /dev/null +++ b/apps/webapp/app/services/dashboardAgentWatchLimits.server.ts @@ -0,0 +1,59 @@ +import type { Limits } from "@trigger.dev/platform"; +import { WATCH_MAX_HOURS } from "@internal/dashboard-agent-contracts"; +import { getCachedLimitAllowingZero, isBillingConfigured } from "./platform.v3.server"; + +// The unlimited sentinel, matching the message quota (TRI-12863 P1). Never Infinity: it +// serializes to null in the limit cache. +export const UNLIMITED_WATCH_LIMIT = 100_000_000; + +// Filled by cloud billing (TRI-12863 P0). Absent until then, and always on self-hosted, so +// the fallback applies and the plan floor is off. +const WATCH_MAX_HOURS_LIMIT_KEY = "agentWatchMaxHours" as keyof Limits; +const WATCH_COUNT_LIMIT_KEY = "agentWatchers" as keyof Limits; + +export type WatchPlanLimits = { + /** Longest window one watch may run for, in hours. */ + maxHours: number; + /** How many active watches the org may run at once. */ + watchers: number; +}; + +async function readLimit(organizationId: string, key: keyof Limits): Promise { + // A plan of 0 means zero, not absent: an org with watches switched off must not read as + // unlimited. Only a missing limit falls open. + const cached = await getCachedLimitAllowingZero(organizationId, key, UNLIMITED_WATCH_LIMIT); + // A cache error leaves `val` empty; fall open to unlimited. + return cached.val ?? UNLIMITED_WATCH_LIMIT; +} + +/** + * The org's plan floors for watches. Fails open: an absent limit (self-hosted, or before the + * cloud side ships) resolves to the unlimited sentinel, so neither floor bites. `read` is the + * plan-limit seam: tests pass their own reader instead of the cached platform one. + */ +export async function resolveWatchPlanLimits( + organizationId: string, + read: (organizationId: string, key: keyof Limits) => Promise = readLimit +): Promise { + const [maxHours, watchers] = await Promise.all([ + read(organizationId, WATCH_MAX_HOURS_LIMIT_KEY), + read(organizationId, WATCH_COUNT_LIMIT_KEY), + ]); + return { maxHours, watchers }; +} + +/** + * The window ceiling actually in force: the plan floor under the code ceiling. A plan that + * allows 100 hours still caps at {@link WATCH_MAX_HOURS}. + */ +export function effectiveWatchMaxHours(planMaxHours: number): number { + return Math.min(planMaxHours, WATCH_MAX_HOURS); +} + +/** + * A watch-limit refusal, plus an upgrade nudge when billing is present. Self-hosted never + * hits this (fails open above), and the nudge is gated so a stray refusal stays quiet there. + */ +export function watchLimitHint(base: string, billingConfigured = isBillingConfigured()): string { + return billingConfigured ? `${base} Upgrade your plan for more.` : base; +} diff --git a/apps/webapp/app/services/dashboardAgentWatches.server.ts b/apps/webapp/app/services/dashboardAgentWatches.server.ts index 10d0d7bf9d2..06973d58e54 100644 --- a/apps/webapp/app/services/dashboardAgentWatches.server.ts +++ b/apps/webapp/app/services/dashboardAgentWatches.server.ts @@ -10,6 +10,7 @@ import { cancelWatch, chatExists, claimWatchSubmission, + countActiveWatchesForOrg, createChat, createWatch, generateWatchId, @@ -68,6 +69,12 @@ import { import { watchCreationCheckDeps } from "~/services/dashboardAgentWatchChecks.server"; import { normalizeErrorFingerprint } from "~/services/dashboardAgentWatchErrorChecks"; import { subscribeUserToWatchAlerts } from "~/services/dashboardAgentWatchAlerts.server"; +import { + effectiveWatchMaxHours, + resolveWatchPlanLimits, + watchLimitHint, + type WatchPlanLimits, +} from "~/services/dashboardAgentWatchLimits.server"; import { mintDashboardAgentWatchBatchToken, mintDashboardAgentWatchToken, @@ -165,6 +172,7 @@ export async function authorizeWatchEnvironmentById(params: { export type CreateWatchErrorCode = | "limit_reached" + | "watch_limit_reached" | "duplicate" | "invalid_target" | "chat_not_found" @@ -273,6 +281,12 @@ export async function createDashboardAgentWatch(params: { scheduleTick?: typeof scheduleWatchTick; /** Skip the real trigger-config gate when a tick scheduler is injected. */ configured?: () => boolean; + /** Plan floors on window and count. Fails open to unlimited when absent. */ + resolveLimits?: (organizationId: string) => Promise; + /** Org-wide active-watch count, for the watcher-count floor. */ + countActiveWatches?: (organizationId: string) => Promise; + /** Gates the upgrade nudge, so self-hosted stays quiet. */ + billingConfigured?: () => boolean; }; }): Promise { const { environment, userId, chatId } = params; @@ -284,6 +298,11 @@ export async function createDashboardAgentWatch(params: { const buildCheckDeps = params.deps?.checkDeps ?? watchCreationCheckDeps; const scheduleTick = params.deps?.scheduleTick ?? scheduleWatchTick; const isDashboardAgentConfigured = params.deps?.configured ?? isDashboardAgentConfiguredDefault; + const resolveLimits = params.deps?.resolveLimits ?? resolveWatchPlanLimits; + const countActiveWatches = + params.deps?.countActiveWatches ?? + ((organizationId: string) => countActiveWatchesForOrg(dashboardAgentDb, { organizationId })); + const hint = (base: string) => watchLimitHint(base, params.deps?.billingConfigured?.()); const checkDeps = buildCheckDeps(environment, now); if (!isDashboardAgentConfigured()) { @@ -331,6 +350,29 @@ export async function createDashboardAgentWatch(params: { return { ok: true, watching: false, identity, immediate }; } + // Both floors are read only now the immediate check didn't answer: a one-shot creates no + // row, so a plan floor must not turn an answerable question into an upgrade nudge. Plan + // floors sit below the code ceilings (min(plan, ceiling)) and fail open: an absent limit + // resolves to unlimited, so neither bites on self-hosted. + const planLimits = await resolveLimits(environment.organizationId); + if (spec.maxHours > effectiveWatchMaxHours(planLimits.maxHours)) { + return { + ok: false, + code: "watch_limit_reached", + error: hint("That watch window is longer than your plan allows."), + }; + } + + // The per-chat cap of 3 still applies independently, in `createWatch`. + const activeCount = await countActiveWatches(environment.organizationId); + if (activeCount >= planLimits.watchers) { + return { + ok: false, + code: "watch_limit_reached", + error: hint("You've reached the number of active watches your plan allows."), + }; + } + const expiresAt = new Date(now.getTime() + spec.maxHours * 60 * 60 * 1000); const created = await createWatch(dashboardAgentDb, { diff --git a/apps/webapp/app/services/platform.v3.server.ts b/apps/webapp/app/services/platform.v3.server.ts index f8b524e3cdd..187ef3b063e 100644 --- a/apps/webapp/app/services/platform.v3.server.ts +++ b/apps/webapp/app/services/platform.v3.server.ts @@ -482,6 +482,37 @@ export async function getCachedLimit(orgId: string, limit: keyof Limits, fallbac }); } +/** + * Reads one plan limit, treating 0 as zero rather than absent: only a missing limit falls back. + * {@link getLimit} keeps its `!result` fallback, which its callers depend on. + */ +export function limitValueAllowingZero( + limits: Limits | undefined, + limit: keyof Limits, + fallback: number +): number { + const result = limits?.[limit]; + + if (result === undefined || result === null) return fallback; + if (typeof result === "number") return result; + if (typeof result === "object" && "number" in result) return result.number; + return fallback; +} + +/** + * Like {@link getCachedLimit}, but a plan value of 0 means zero. Cached under its own key so it + * never crosses with {@link getCachedLimit}. + */ +export async function getCachedLimitAllowingZero( + orgId: string, + limit: keyof Limits, + fallback: number +) { + return platformCache.limits.swr(`${orgId}:${limit}:allow-zero`, async () => + limitValueAllowingZero(await getLimits(orgId), limit, fallback) + ); +} + export async function customerPortalUrl(orgId: string, orgSlug: string) { if (!client) return undefined; diff --git a/apps/webapp/app/services/queryService.server.ts b/apps/webapp/app/services/queryService.server.ts index 986172775cb..f4bf4b940a4 100644 --- a/apps/webapp/app/services/queryService.server.ts +++ b/apps/webapp/app/services/queryService.server.ts @@ -397,6 +397,7 @@ export async function executeQuery( ...getDefaultClickhouseSettings(), ...queryCacheSettings, ...baseOptions.clickhouseSettings, // Allow caller overrides if needed + readonly: "1", // Not overridable: every query through here is read-only. }, querySettings: { maxRows: env.QUERY_CLICKHOUSE_MAX_RETURNED_ROWS, diff --git a/apps/webapp/test/dashboardAgentDurableResume.test.ts b/apps/webapp/test/dashboardAgentDurableResume.test.ts new file mode 100644 index 00000000000..8fd75d52e0a --- /dev/null +++ b/apps/webapp/test/dashboardAgentDurableResume.test.ts @@ -0,0 +1,347 @@ +import { + appendChatMessageOnceByChatId, + createChat, + createDashboardAgentDb, + getChatMessages, + getSession, + persistMessages, + persistTurn, + type DashboardAgentDb, + type DashboardAgentDbClient, +} from "@internal/dashboard-agent-db"; +import { postgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { readdirSync, readFileSync } from "node:fs"; +import path from "node:path"; +import { afterEach, describe, expect } from "vitest"; + +/** + * Durability of a chat.agent turn across a crash and a resume, against a real table + * (TRI-11166). + * + * The primitive gives chat.agent durability by snapshotting the transcript and replaying it + * on the next boot. These tests pin the store seam that replay lands on: the completing turn + * re-sends its whole snapshot, so the store has to fold that replay into exactly one row per + * message — no double-appended turn, no lost mid-turn message — and reconstruct the session + * cursor a refreshed client resumes from. + * + * What is NOT covered here, because it lives inside the closed chat.agent primitive package + * (object-store snapshot write, S2 `.in`/`.out` replay, `.out` trimming, OOM restart): the + * transport-level replay and the snapshot URL's own auth. The client-side reconnect / Last- + * Event-ID replay is covered in packages/trigger-sdk/src/v3/chat.test.ts. These tests are the + * store-level backstop those depend on. See the PR body for the residual follow-ups. + */ + +let agentDb: DashboardAgentDb; +let agentDbClient: DashboardAgentDbClient | undefined; + +const MIGRATIONS = path.resolve(__dirname, "../../../internal-packages/dashboard-agent-db/drizzle"); + +async function applyAgentSchema(prisma: PrismaClient) { + for (const name of readdirSync(MIGRATIONS) + .filter((file) => file.endsWith(".sql")) + .sort()) { + const sql = readFileSync(path.join(MIGRATIONS, name), "utf8"); + for (const statement of sql.split("--> statement-breakpoint")) { + const trimmed = statement.trim(); + if (trimmed.length > 0) await prisma.$executeRawUnsafe(trimmed); + } + } +} + +const ORG = "org_resume"; +const USER = "user_resume"; + +async function boot(prisma: PrismaClient, connectionUri: string, chatId: string) { + await applyAgentSchema(prisma); + agentDbClient = createDashboardAgentDb(connectionUri, { max: 4 }); + agentDb = agentDbClient.db; + await createChat(agentDb, { id: chatId, organizationId: ORG, userId: USER }); +} + +afterEach(async () => { + await agentDbClient?.close(); + agentDbClient = undefined; +}); + +function textMessage(id: string, role: "user" | "assistant" = "assistant", text = id) { + return { id, role, parts: [{ type: "text", text }] }; +} + +/** A tool part, so a mid-flight call and its completed result share an id but differ in body. */ +function toolMessage(id: string, state: "input-available" | "output-available") { + return { + id, + role: "assistant" as const, + parts: [{ type: "tool-get_query_schema", state, toolCallId: `${id}_call`, input: {} }], + }; +} + +async function transcript(chatId: string): Promise<{ id: string }[]> { + return (await getChatMessages(agentDb, { chatId, organizationId: ORG, userId: USER })) as { + id: string; + }[]; +} + +/** The allocator, where a wasted/duplicated slot is observable. */ +async function nextPosition(prisma: PrismaClient, chatId: string): Promise { + const rows = await prisma.$queryRawUnsafe<{ next_message_position: number }[]>( + `select next_message_position from trigger_dashboard_agent.chats where id = $1`, + chatId + ); + return rows[0]!.next_message_position; +} + +async function rowCount(prisma: PrismaClient, chatId: string): Promise { + const rows = await prisma.$queryRawUnsafe<{ count: bigint }[]>( + `select count(*)::int as count from trigger_dashboard_agent.chat_messages where chat_id = $1`, + chatId + ); + return Number(rows[0]!.count); +} + +describe("a streamed-then-resumed turn is not double-appended", () => { + postgresTest( + "re-delivering the completing turn finalises in place and appends nothing", + async ({ prisma, postgresContainer }) => { + const chatId = "chat_no_double"; + await boot(prisma, postgresContainer.getConnectionUri(), chatId); + + // The turn started: onTurnStart stored the user turn and the tool call mid-flight. + await persistMessages(agentDb, { + chatId, + messages: [textMessage("u1", "user"), toolMessage("a1", "input-available")], + }); + expect(await rowCount(prisma, chatId)).toBe(2); + + const completing = { + chatId, + messages: [textMessage("u1", "user"), toolMessage("a1", "output-available")], + finalizeMessageIds: ["a1"], + session: { publicAccessToken: "pat", lastEventId: "7", runId: "run" }, + }; + + // The turn completes, replaying its whole snapshot. `a1` is finalised, not re-added. + await persistTurn(agentDb, completing); + // The resume: the same completed turn is delivered again (client reconnected and the + // host re-persisted). It must converge — no second `a1`, no extra row of any kind. + await persistTurn(agentDb, completing); + + expect((await transcript(chatId)).map((m) => m.id)).toEqual(["u1", "a1"]); + expect(await rowCount(prisma, chatId)).toBe(2); + // Only u1 and a1 ever reserved a slot (allocator starts at 1); the finalisation and the + // replay reserve none, so the next free position is still 3. + expect(await nextPosition(prisma, chatId)).toBe(3); + // And `a1` is the completed body the user saw, not the mid-flight call. + const stored = (await transcript(chatId))[1] as unknown as { + parts: { state: string }[]; + }; + expect(stored.parts[0]!.state).toBe("output-available"); + }, + 30_000 + ); +}); + +describe("a crash mid-turn is reconstructed by the next boot's replay", () => { + postgresTest( + "the resumed turn keeps the mid-turn append, finalises its own message, and rebuilds the session cursor", + async ({ prisma, postgresContainer }) => { + const chatId = "chat_crash_resume"; + await boot(prisma, postgresContainer.getConnectionUri(), chatId); + + // Turn in flight: the snapshot it started from, stored before the model finished. + const snapshot = [textMessage("u1", "user"), toolMessage("a1", "input-available")]; + await persistMessages(agentDb, { chatId, messages: snapshot }); + + // A wake lands mid-turn, off its own lane — the message the old replace-the-array + // write used to lose. + await appendChatMessageOnceByChatId(agentDb, { + chatId, + message: textMessage("wake:w1"), + }); + + // Before the crash there is no session row to resume from. + expect(await getSession(agentDb, { chatId, organizationId: ORG, userId: USER })).toBeNull(); + + // Boot after the crash: replay the whole transcript, finalise the turn's own message, + // and write the session the client resumes from — all in one persistTurn. + await persistTurn(agentDb, { + chatId, + messages: [ + textMessage("u1", "user"), + toolMessage("a1", "output-available"), + textMessage("a2"), + ], + finalizeMessageIds: ["a1", "a2"], + session: { publicAccessToken: "pat_resumed", lastEventId: "99", runId: "run_resumed" }, + }); + + // Nothing was lost and the wake sits where it happened: after the snapshot, before the + // reply the turn went on to produce. + expect((await transcript(chatId)).map((m) => m.id)).toEqual(["u1", "a1", "wake:w1", "a2"]); + + const session = await getSession(agentDb, { chatId, organizationId: ORG, userId: USER }); + expect(session).toMatchObject({ + publicAccessToken: "pat_resumed", + lastEventId: "99", + runId: "run_resumed", + }); + }, + 30_000 + ); +}); + +describe("the session cursor a refreshed client resumes from", () => { + postgresTest( + "getSession returns the last persisted cursor, and a later turn advances it", + async ({ prisma, postgresContainer }) => { + const chatId = "chat_cursor"; + await boot(prisma, postgresContainer.getConnectionUri(), chatId); + + await persistTurn(agentDb, { + chatId, + messages: [textMessage("u1", "user"), textMessage("a1")], + session: { publicAccessToken: "pat1", lastEventId: "10", runId: "run1" }, + }); + // A mid-stream refresh reads exactly this cursor and resumes .out from it. + expect( + (await getSession(agentDb, { chatId, organizationId: ORG, userId: USER }))?.lastEventId + ).toBe("10"); + + // The next turn overwrites the cursor — a stale value is replaced, never appended. + await persistTurn(agentDb, { + chatId, + messages: [textMessage("u1", "user"), textMessage("a1"), textMessage("a2")], + session: { publicAccessToken: "pat2", lastEventId: "25", runId: "run2" }, + }); + const session = await getSession(agentDb, { chatId, organizationId: ORG, userId: USER }); + expect(session).toMatchObject({ + publicAccessToken: "pat2", + lastEventId: "25", + runId: "run2", + }); + }, + 30_000 + ); +}); + +describe("a failed snapshot write leaves the next boot a clean replay", () => { + postgresTest( + "a persistTurn that throws mid-write rolls back what it already wrote, and the retry replays with no loss", + async ({ prisma, postgresContainer }) => { + const chatId = "chat_write_fail"; + await boot(prisma, postgresContainer.getConnectionUri(), chatId); + + // A durable first turn, its tool call still mid-flight, and the session cursor it left. + await persistTurn(agentDb, { + chatId, + messages: [textMessage("u1", "user"), toolMessage("a1", "input-available")], + session: { publicAccessToken: "pat1", lastEventId: "1", runId: "run1" }, + }); + const positionBefore = await nextPosition(prisma, chatId); + + // Tear the next turn at the INSERT itself, so the failure lands after `a1` is finalised + // in place and after the slots are reserved no matter how the store orders its up-front + // validation. A row planted directly at the position the allocator is about to hand out + // makes that insert violate `chat_messages_chat_position_key`. Scaffolding, not part of + // the transcript under test — removed once the tear has fired. + await prisma.$executeRawUnsafe( + `insert into trigger_dashboard_agent.chat_messages (chat_id, message_id, position, role, message) + values ($1, 'planted_collision', $2, 'assistant', '{}'::jsonb)`, + chatId, + positionBefore + ); + + // The driver names the failing statement, so the rejection itself pins where the tear fired. + await expect( + persistTurn(agentDb, { + chatId, + messages: [ + textMessage("u1", "user"), + toolMessage("a1", "output-available"), + textMessage("a2"), + ], + finalizeMessageIds: ["a1"], + session: { publicAccessToken: "pat_torn", lastEventId: "2", runId: "run_torn" }, + }) + ).rejects.toThrow(/Failed query: insert into .*chat_messages/); + + await prisma.$executeRawUnsafe( + `delete from trigger_dashboard_agent.chat_messages where chat_id = $1 and message_id = 'planted_collision'`, + chatId + ); + + // The whole turn rolled back. The in-place rewrite the store had already applied is undone: + // `a1` is the mid-flight call again, not the finalised body the torn turn wrote. + expect((await transcript(chatId)).map((m) => m.id)).toEqual(["u1", "a1"]); + const tornA1 = (await transcript(chatId))[1] as unknown as { parts: { state: string }[] }; + expect(tornA1.parts[0]!.state).toBe("input-available"); + expect(await rowCount(prisma, chatId)).toBe(2); + // The slot it reserved for `a2` came back too, so the retry doesn't leave a gap. + expect(await nextPosition(prisma, chatId)).toBe(positionBefore); + // The cursor is still the first turn's: the failed turn never got as far as writing one. + expect( + await getSession(agentDb, { chatId, organizationId: ORG, userId: USER }) + ).toMatchObject({ publicAccessToken: "pat1", lastEventId: "1" }); + + // The retry — a clean replay of the same turn — lands everything exactly once. + await persistTurn(agentDb, { + chatId, + messages: [ + textMessage("u1", "user"), + toolMessage("a1", "output-available"), + textMessage("a2"), + ], + finalizeMessageIds: ["a1"], + session: { publicAccessToken: "pat2", lastEventId: "2", runId: "run2" }, + }); + expect((await transcript(chatId)).map((m) => m.id)).toEqual(["u1", "a1", "a2"]); + const retriedA1 = (await transcript(chatId))[1] as unknown as { parts: { state: string }[] }; + expect(retriedA1.parts[0]!.state).toBe("output-available"); + // One new row, one new slot: the rolled-back reservation was not double-counted. + expect(await nextPosition(prisma, chatId)).toBe(positionBefore + 1); + expect( + await getSession(agentDb, { chatId, organizationId: ORG, userId: USER }) + ).toMatchObject({ publicAccessToken: "pat2", lastEventId: "2" }); + }, + 30_000 + ); +}); + +describe("an OOM restart replays the turn cleanly", () => { + postgresTest( + "a restarted turn that re-sends its snapshot loses no data and doubles nothing", + async ({ prisma, postgresContainer }) => { + // The store seam an OOM restart lands on: the primitive restarts the run, replays `.in`, + // and re-persists. `.out` trimming and the OOM restart itself are inside the primitive + // (not reachable here) — this pins that a re-run's re-sent snapshot is idempotent. + const chatId = "chat_oom_restart"; + await boot(prisma, postgresContainer.getConnectionUri(), chatId); + + const firstAttempt = [textMessage("u1", "user"), toolMessage("a1", "input-available")]; + await persistMessages(agentDb, { chatId, messages: firstAttempt }); + const positionAfterFirst = await nextPosition(prisma, chatId); + + // The run OOMs and restarts. It replays the same input, produces the same ids, and + // finalises the turn it now completes. + const restarted = { + chatId, + messages: [ + textMessage("u1", "user"), + toolMessage("a1", "output-available"), + textMessage("a2"), + ], + finalizeMessageIds: ["a1", "a2"], + session: { publicAccessToken: "pat", lastEventId: "5", runId: "run_restarted" }, + }; + await persistTurn(agentDb, restarted); + // A second restart delivering the same turn again still converges. + await persistTurn(agentDb, restarted); + + expect((await transcript(chatId)).map((m) => m.id)).toEqual(["u1", "a1", "a2"]); + // The replayed u1/a1 reserved no new slots; only a2 was genuinely new. + expect(await nextPosition(prisma, chatId)).toBe(positionAfterFirst + 1); + }, + 30_000 + ); +}); diff --git a/apps/webapp/test/dashboardAgentInvestigationPoison.test.ts b/apps/webapp/test/dashboardAgentInvestigationPoison.test.ts new file mode 100644 index 00000000000..52b87f35c17 --- /dev/null +++ b/apps/webapp/test/dashboardAgentInvestigationPoison.test.ts @@ -0,0 +1,166 @@ +import { + createChat, + createDashboardAgentDb, + getInvestigation, + settleInvestigationAndCloseCard, + upsertInvestigationRevision, + type DashboardAgentDb, + type DashboardAgentDbClient, +} from "@internal/dashboard-agent-db"; +import { + investigationStateSchema, + type InvestigationState, +} from "@internal/dashboard-agent-contracts"; +import { postgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { readdirSync, readFileSync } from "node:fs"; +import path from "node:path"; +import { afterEach, describe, expect, vi } from "vitest"; + +const ctx = vi.hoisted(() => ({ + agentDb: undefined as unknown as DashboardAgentDb, +})); + +vi.mock("~/services/dashboardAgentDb.server", () => ({ + get dashboardAgentDb() { + return ctx.agentDb; + }, +})); + +const { sweepDashboardAgentInvestigations, INVESTIGATION_STALE_MS, MAX_SWEEP_ATTEMPTS } = + await import("~/services/dashboardAgentInvestigationSweep.server"); + +async function applyAgentSchema(prisma: PrismaClient) { + const folder = path.resolve(__dirname, "../../../internal-packages/dashboard-agent-db/drizzle"); + const migrations = readdirSync(folder) + .filter((file) => file.endsWith(".sql")) + .sort(); + for (const name of migrations) { + const sql = readFileSync(path.join(folder, name), "utf8"); + for (const statement of sql.split("--> statement-breakpoint")) { + const trimmed = statement.trim(); + if (trimmed.length > 0) await prisma.$executeRawUnsafe(trimmed); + } + } +} + +let agentDbClient: DashboardAgentDbClient | undefined; +let prismaForRaw: PrismaClient | undefined; + +async function boot(prisma: PrismaClient, connectionUri: string) { + await applyAgentSchema(prisma); + agentDbClient = createDashboardAgentDb(connectionUri, { max: 2 }); + ctx.agentDb = agentDbClient.db; + prismaForRaw = prisma; +} + +afterEach(async () => { + await agentDbClient?.close(); + agentDbClient = undefined; +}); + +const ORG = "org_poison"; +const USER = "user_poison"; + +function openState(): InvestigationState { + return investigationStateSchema.parse({ + outcome: "in_progress", + severity: "warn", + confidence: "medium", + title: "a stuck card", + headline: "Still checking.", + progress: "Reading spans", + checkNext: [], + hypotheses: [], + evidence: [], + }); +} + +async function seedInvestigation(chatId: string, ageMs: number): Promise { + await createChat(ctx.agentDb, { id: chatId, organizationId: ORG, userId: USER }); + const created = await upsertInvestigationRevision(ctx.agentDb, { + chatId, + projectRef: "proj", + environmentRef: "env", + state: openState(), + }); + if (!created.ok) throw new Error("fixture investigation not created"); + await prismaForRaw!.$executeRawUnsafe( + `update trigger_dashboard_agent.investigations + set updated_at = now() - ($2 || ' milliseconds')::interval where id = $1`, + created.id, + String(ageMs) + ); + return created.id; +} + +async function outcomeOf(id: string): Promise { + const row = await getInvestigation(ctx.agentDb, { id }); + return row ? (row.state as { outcome?: string }).outcome : undefined; +} + +const STALE_AGE_MS = INVESTIGATION_STALE_MS + 60_000; +const OLDER_AGE_MS = STALE_AGE_MS + 60_000; + +describe("the investigation sweep with a poison row", () => { + postgresTest( + "a row that always fails to settle cannot pin the head and starve a newer row", + async ({ prisma, postgresContainer }) => { + await boot(prisma, postgresContainer.getConnectionUri()); + + // Poison sorts first (older `updated_at`); renderable is newer. + const poisonId = await seedInvestigation("chat_poison", OLDER_AGE_MS); + const renderableId = await seedInvestigation("chat_ok", STALE_AGE_MS); + + // Only the poison row's settle throws; the renderable one goes through the real path. + const settleAndClose = (params: { id: string; chatId: string; note: string }) => { + if (params.id === poisonId) throw new Error("state isn't renderable"); + return settleInvestigationAndCloseCard(ctx.agentDb, params); + }; + + // limit 1 forces head contention: without backoff the poison row would win every run. + // A failed run throws so the job retries, but the attempt is recorded before it does. + await expect( + sweepDashboardAgentInvestigations({ limit: 1, settleAndClose }) + ).rejects.toThrow(); + expect(await outcomeOf(poisonId)).toBe("in_progress"); + expect(await outcomeOf(renderableId)).toBe("in_progress"); + + // Next run: the poison row now sorts behind the never-attempted renderable one, + // so the newer row is picked and settled despite the poison row still being stale. + const second = await sweepDashboardAgentInvestigations({ limit: 1, settleAndClose }); + expect(second).toMatchObject({ stale: 1, settled: 1, failed: 0 }); + expect(await outcomeOf(renderableId)).toBe("inconclusive"); + expect(await outcomeOf(poisonId)).toBe("in_progress"); + }, + 30_000 + ); + + postgresTest( + "after the attempt cap the poison row is abandoned and leaves the queue", + async ({ prisma, postgresContainer }) => { + await boot(prisma, postgresContainer.getConnectionUri()); + const poisonId = await seedInvestigation("chat_poison", STALE_AGE_MS); + + const settleAndClose = () => { + throw new Error("state isn't renderable"); + }; + + // The first MAX_SWEEP_ATTEMPTS-1 runs record a failed attempt and throw; the row stays stale. + for (let i = 1; i < MAX_SWEEP_ATTEMPTS; i++) { + await expect(sweepDashboardAgentInvestigations({ settleAndClose })).rejects.toThrow(); + expect(await outcomeOf(poisonId)).toBe("in_progress"); + } + + // The capped run force-settles the row without the render path, so it leaves the queue. + const capped = await sweepDashboardAgentInvestigations({ settleAndClose }); + expect(capped).toMatchObject({ stale: 1, abandoned: 1, failed: 0 }); + expect(await outcomeOf(poisonId)).toBe("inconclusive"); + + // Nothing stale remains, so the poison row is no longer swept. + const after = await sweepDashboardAgentInvestigations({ settleAndClose }); + expect(after).toMatchObject({ stale: 0 }); + }, + 30_000 + ); +}); diff --git a/apps/webapp/test/dashboardAgentQuota.test.ts b/apps/webapp/test/dashboardAgentQuota.test.ts new file mode 100644 index 00000000000..21d5e46ffe7 --- /dev/null +++ b/apps/webapp/test/dashboardAgentQuota.test.ts @@ -0,0 +1,184 @@ +import { + createChat, + createDashboardAgentDb, + getAgentMessageUsage, + incrementAgentMessageUsage, + softDeleteChat, + type DashboardAgentDb, + type DashboardAgentDbClient, +} from "@internal/dashboard-agent-db"; +import { postgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { readdirSync, readFileSync } from "node:fs"; +import path from "node:path"; +import { afterEach, describe, expect, it } from "vitest"; +import { + agentTurnCountsAgainstQuota, + checkAgentMessageQuota, + currentAgentMessagePeriod, + resolveAgentMessageQuota, + UNLIMITED_AGENT_MESSAGES, +} from "~/services/dashboardAgentQuota.server"; + +/** + * Server-side agent message quota (TRI-12863): a per-(org, period) counter that a deleted chat + * can't lower, a pure at/over/under rule, and a resolver that fails open when the limit is + * absent (self-hosted) or the counter read throws. + */ + +const MIGRATIONS = path.resolve(__dirname, "../../../internal-packages/dashboard-agent-db/drizzle"); + +async function applyAgentSchema(prisma: PrismaClient) { + for (const name of readdirSync(MIGRATIONS) + .filter((file) => file.endsWith(".sql")) + .sort()) { + const sql = readFileSync(path.join(MIGRATIONS, name), "utf8"); + for (const statement of sql.split("--> statement-breakpoint")) { + const trimmed = statement.trim(); + if (trimmed.length > 0) await prisma.$executeRawUnsafe(trimmed); + } + } +} + +const ORG = "org_quota"; +const USER = "user_quota"; + +let agentDbClient: DashboardAgentDbClient | undefined; + +async function boot(prisma: PrismaClient, connectionUri: string): Promise { + await applyAgentSchema(prisma); + agentDbClient = createDashboardAgentDb(connectionUri, { max: 4 }); + return agentDbClient.db; +} + +afterEach(async () => { + await agentDbClient?.close(); + agentDbClient = undefined; +}); + +describe("checkAgentMessageQuota", () => { + it("is not reached under the limit", () => { + expect(checkAgentMessageQuota({ used: 5, limit: 20 })).toEqual({ reached: false }); + }); + + it("is reached at the limit", () => { + // Control break: `>=`. Flip to `>` and this fails. + expect(checkAgentMessageQuota({ used: 20, limit: 20 })).toEqual({ reached: true }); + }); + + it("is reached over the limit", () => { + expect(checkAgentMessageQuota({ used: 21, limit: 20 })).toEqual({ reached: true }); + }); + + it("is never reached against the unlimited sentinel", () => { + expect(checkAgentMessageQuota({ used: 10_000, limit: UNLIMITED_AGENT_MESSAGES })).toEqual({ + reached: false, + }); + }); +}); + +describe("agentTurnCountsAgainstQuota", () => { + it("counts a genuine new user message (submit-message)", () => { + expect( + agentTurnCountsAgainstQuota({ kind: "message", payload: { trigger: "submit-message" } }) + ).toBe(true); + }); + + it("does not count a retry/regenerate", () => { + // Control break: a regenerate re-runs from history with no new message, so it must not + // burn quota. Widen the rule back to `!== "action"` and this fails. + expect( + agentTurnCountsAgainstQuota({ kind: "message", payload: { trigger: "regenerate-message" } }) + ).toBe(false); + }); + + it("does not count a wake (action turn)", () => { + expect(agentTurnCountsAgainstQuota({ kind: "message", payload: { trigger: "action" } })).toBe( + false + ); + }); + + it("does not count a non-message turn or a missing body", () => { + expect(agentTurnCountsAgainstQuota({ kind: "action" })).toBe(false); + expect(agentTurnCountsAgainstQuota(undefined)).toBe(false); + }); +}); + +describe("currentAgentMessagePeriod", () => { + it("is a zero-padded UTC calendar month", () => { + expect(currentAgentMessagePeriod(new Date(Date.UTC(2026, 7, 9)))).toBe("2026-08"); + expect(currentAgentMessagePeriod(new Date(Date.UTC(2026, 0, 1)))).toBe("2026-01"); + }); +}); + +describe("the per-(org, period) counter", () => { + postgresTest( + "accumulates and a deleted chat cannot free quota within the period", + async ({ prisma, postgresContainer }) => { + const db = await boot(prisma, postgresContainer.getConnectionUri()); + const period = "2026-08"; + + // The create path and then an append: two messages, same period. + expect(await incrementAgentMessageUsage(db, { organizationId: ORG, period })).toBe(1); + expect(await incrementAgentMessageUsage(db, { organizationId: ORG, period })).toBe(2); + expect(await getAgentMessageUsage(db, { organizationId: ORG, period })).toBe(2); + + // Deleting a chat must not move the counter: it is not joined to chats. + await createChat(db, { id: "chat_del", organizationId: ORG, userId: USER }); + await softDeleteChat(db, { chatId: "chat_del", userId: USER, organizationId: ORG }); + expect(await getAgentMessageUsage(db, { organizationId: ORG, period })).toBe(2); + + // The next period and other orgs start fresh. + expect(await getAgentMessageUsage(db, { organizationId: ORG, period: "2026-09" })).toBe(0); + expect(await getAgentMessageUsage(db, { organizationId: "org_other", period })).toBe(0); + } + ); +}); + +describe("resolveAgentMessageQuota", () => { + postgresTest( + "reports reached over the limit, and never reached when unlimited", + async ({ prisma, postgresContainer }) => { + const db = await boot(prisma, postgresContainer.getConnectionUri()); + const now = new Date(); + const period = currentAgentMessagePeriod(now); + for (let i = 0; i < 3; i++) { + await incrementAgentMessageUsage(db, { organizationId: ORG, period }); + } + + expect( + await resolveAgentMessageQuota(db, { organizationId: ORG, now, readLimit: async () => 3 }) + ).toEqual({ + reached: true, + used: 3, + limit: 3, + }); + expect( + await resolveAgentMessageQuota(db, { organizationId: ORG, now, readLimit: async () => 20 }) + ).toEqual({ reached: false, used: 3, limit: 20 }); + + // Self-hosted: the limit is absent, so the fallback (unlimited sentinel) applies and there + // is no cap — no extra branching, it falls out of the fallback. + const selfHosted = await resolveAgentMessageQuota(db, { + organizationId: ORG, + now, + readLimit: async () => UNLIMITED_AGENT_MESSAGES, + }); + expect(selfHosted?.reached).toBe(false); + } + ); + + it("fails open when the counter read throws", async () => { + const throwingDb = { + select: () => { + throw new Error("db down"); + }, + } as unknown as DashboardAgentDb; + + const result = await resolveAgentMessageQuota(throwingDb, { + organizationId: ORG, + readLimit: async () => 5, + }); + expect(result).toBeUndefined(); + }); +}); diff --git a/apps/webapp/test/dashboardAgentTenantIsolation.test.ts b/apps/webapp/test/dashboardAgentTenantIsolation.test.ts new file mode 100644 index 00000000000..3aeb1bc16e7 --- /dev/null +++ b/apps/webapp/test/dashboardAgentTenantIsolation.test.ts @@ -0,0 +1,240 @@ +import { + appendChatMessageOnce, + chatExists, + countUserMessages, + createChat, + createDashboardAgentDb, + getChatMessages, + getSession, + listChats, + persistTurn, + type DashboardAgentDb, + type DashboardAgentDbClient, +} from "@internal/dashboard-agent-db"; +import { postgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { readdirSync, readFileSync } from "node:fs"; +import path from "node:path"; +import { afterEach, describe, expect } from "vitest"; + +/** + * Cross-tenant isolation for the chat store, against a real table (TRI-11166). + * + * The 2026-06-10 chat.agent audit flagged a cross-tenant read: a chat/session belongs to + * one (org, user) pair, and every read that hands back its transcript or its session token + * has to be scoped by that pair. A chatId from another tenant must read as not-found — never + * as another tenant's transcript, and never as another tenant's public access token, which + * is the credential a resumed session boots from. + * + * The store's own queries are the floor: the resource route scopes on project.organizationId + * above this, but a bug there would still be caught here because these queries refuse a + * foreign (org, user) outright rather than trusting the caller. + */ + +let agentDb: DashboardAgentDb; +let agentDbClient: DashboardAgentDbClient | undefined; + +const MIGRATIONS = path.resolve(__dirname, "../../../internal-packages/dashboard-agent-db/drizzle"); + +/** Replays every migration in order, so a new migration can't leave the suite on a stale schema. */ +async function applyAgentSchema(prisma: PrismaClient) { + for (const name of readdirSync(MIGRATIONS) + .filter((file) => file.endsWith(".sql")) + .sort()) { + const sql = readFileSync(path.join(MIGRATIONS, name), "utf8"); + for (const statement of sql.split("--> statement-breakpoint")) { + const trimmed = statement.trim(); + if (trimmed.length > 0) await prisma.$executeRawUnsafe(trimmed); + } + } +} + +// Org A owns the chat. Org B and a same-org other user are the foreign tenants. +const ORG_A = "org_a"; +const USER_A = "user_a"; +const ORG_B = "org_b"; +const USER_B = "user_b"; +const CHAT = "chat_owned_by_a"; + +async function boot(prisma: PrismaClient, connectionUri: string) { + await applyAgentSchema(prisma); + agentDbClient = createDashboardAgentDb(connectionUri, { max: 4 }); + agentDb = agentDbClient.db; +} + +afterEach(async () => { + await agentDbClient?.close(); + agentDbClient = undefined; +}); + +function textMessage(id: string, role: "user" | "assistant" = "assistant") { + return { id, role, parts: [{ type: "text", text: id }] }; +} + +/** Seed a chat under org A with a transcript and a live session (its PAT is the credential). */ +async function seedOwnedChat() { + await createChat(agentDb, { id: CHAT, organizationId: ORG_A, userId: USER_A }); + await persistTurn(agentDb, { + chatId: CHAT, + messages: [textMessage("u1", "user"), textMessage("a1")], + session: { publicAccessToken: "pat_secret_of_a", lastEventId: "42", runId: "run_a" }, + }); +} + +const foreignScopes = [ + { name: "another org", organizationId: ORG_B, userId: USER_B }, + // Same org, different user: a member of A's org still isn't the chat's owner. + { name: "another user in the same org", organizationId: ORG_A, userId: USER_B }, + // Right user id, wrong org: the id alone must not carry across a tenant boundary. + { name: "the owner's user id under another org", organizationId: ORG_B, userId: USER_A }, +]; + +describe("getChatMessages is scoped to the owning (org, user)", () => { + postgresTest( + "the owner reads the transcript; every foreign tenant reads not-found", + async ({ prisma, postgresContainer }) => { + await boot(prisma, postgresContainer.getConnectionUri()); + await seedOwnedChat(); + + const owned = await getChatMessages(agentDb, { + chatId: CHAT, + organizationId: ORG_A, + userId: USER_A, + }); + expect((owned as { id: string }[]).map((m) => m.id)).toEqual(["u1", "a1"]); + + for (const scope of foreignScopes) { + // null is not-found. It must never be [] (a visible-but-empty chat) and never A's rows. + const seen = await getChatMessages(agentDb, { + chatId: CHAT, + organizationId: scope.organizationId, + userId: scope.userId, + }); + expect(seen, scope.name).toBeNull(); + } + }, + 30_000 + ); +}); + +describe("getSession never hands a foreign tenant the owner's access token", () => { + postgresTest( + "the owner gets the session; every foreign tenant gets null", + async ({ prisma, postgresContainer }) => { + await boot(prisma, postgresContainer.getConnectionUri()); + await seedOwnedChat(); + + const owned = await getSession(agentDb, { + chatId: CHAT, + organizationId: ORG_A, + userId: USER_A, + }); + expect(owned?.publicAccessToken).toBe("pat_secret_of_a"); + + for (const scope of foreignScopes) { + const seen = await getSession(agentDb, { + chatId: CHAT, + organizationId: scope.organizationId, + userId: scope.userId, + }); + // A leaked session row would carry A's PAT — the resume credential. Refuse outright. + expect(seen, scope.name).toBeNull(); + } + }, + 30_000 + ); +}); + +describe("chatExists is the owner check the action routes gate on", () => { + postgresTest( + "true for the owner, false for every foreign tenant", + async ({ prisma, postgresContainer }) => { + await boot(prisma, postgresContainer.getConnectionUri()); + await seedOwnedChat(); + + expect( + await chatExists(agentDb, { chatId: CHAT, organizationId: ORG_A, userId: USER_A }) + ).toBe(true); + for (const scope of foreignScopes) { + expect( + await chatExists(agentDb, { + chatId: CHAT, + organizationId: scope.organizationId, + userId: scope.userId, + }), + scope.name + ).toBe(false); + } + }, + 30_000 + ); +}); + +describe("listChats and countUserMessages never surface another tenant's chat", () => { + postgresTest( + "a foreign tenant lists nothing and counts nothing of the owner's", + async ({ prisma, postgresContainer }) => { + await boot(prisma, postgresContainer.getConnectionUri()); + await seedOwnedChat(); + + const ownedList = await listChats(agentDb, { organizationId: ORG_A, userId: USER_A }); + expect(ownedList.map((c) => c.id)).toEqual([CHAT]); + expect(await countUserMessages(agentDb, { organizationId: ORG_A, userId: USER_A })).toBe(1); + + for (const scope of foreignScopes) { + const list = await listChats(agentDb, { + organizationId: scope.organizationId, + userId: scope.userId, + }); + expect(list, scope.name).toEqual([]); + expect( + await countUserMessages(agentDb, { + organizationId: scope.organizationId, + userId: scope.userId, + }), + scope.name + ).toBe(0); + } + }, + 30_000 + ); +}); + +describe("a foreign org cannot append to another tenant's chat", () => { + postgresTest( + "appendChatMessageOnce with a foreign org writes nothing and leaves the transcript intact", + async ({ prisma, postgresContainer }) => { + await boot(prisma, postgresContainer.getConnectionUri()); + await seedOwnedChat(); + + const before = await getChatMessages(agentDb, { + chatId: CHAT, + organizationId: ORG_A, + userId: USER_A, + }); + + // A chat id from another org appends nothing when the org is verified. + const wroteForeignOrg = await appendChatMessageOnce(agentDb, { + chatId: CHAT, + userId: USER_A, + organizationId: ORG_B, + message: { id: "intruder", role: "assistant" }, + }); + expect(wroteForeignOrg).toBe(false); + + // And a foreign user, same org, is refused too. + const wroteForeignUser = await appendChatMessageOnce(agentDb, { + chatId: CHAT, + userId: USER_B, + organizationId: ORG_A, + message: { id: "intruder2", role: "assistant" }, + }); + expect(wroteForeignUser).toBe(false); + + expect( + await getChatMessages(agentDb, { chatId: CHAT, organizationId: ORG_A, userId: USER_A }) + ).toEqual(before); + }, + 30_000 + ); +}); diff --git a/apps/webapp/test/dashboardAgentWatchLimitStatus.test.ts b/apps/webapp/test/dashboardAgentWatchLimitStatus.test.ts new file mode 100644 index 00000000000..a2ad2f3fd5a --- /dev/null +++ b/apps/webapp/test/dashboardAgentWatchLimitStatus.test.ts @@ -0,0 +1,163 @@ +import { + createDashboardAgentDb, + type DashboardAgentDb, + type DashboardAgentDbClient, +} from "@internal/dashboard-agent-db"; +import { postgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { readdirSync, readFileSync } from "node:fs"; +import path from "node:path"; +import { afterEach, describe, expect, vi } from "vitest"; +import type * as WatchLimitsModule from "~/services/dashboardAgentWatchLimits.server"; + +// A plan-limit refusal (`watch_limit_reached`) is a 409, not a 500. The card submit's status +// ladder must map it the same way the MCP route does, or a full org sees an "unexpected error". + +const ctx = vi.hoisted(() => ({ + prisma: undefined as unknown as PrismaClient, + agentDb: undefined as unknown as DashboardAgentDb, + userId: "", +})); + +vi.mock("~/db.server", () => { + const proxy = new Proxy( + {}, + { get: (_target, prop) => (ctx.prisma as unknown as Record)[prop as string] } + ); + return { prisma: proxy, $replica: proxy, sqlDatabaseSchema: undefined }; +}); + +vi.mock("~/services/session.server", () => ({ + requireUser: async () => ({ id: ctx.userId, admin: false, isImpersonating: false }), +})); + +vi.mock("~/v3/canAccessDashboardAgent.server", () => ({ + canAccessDashboardAgent: async () => true, +})); + +vi.mock("~/services/dashboardAgentDb.server", () => ({ + get dashboardAgentDb() { + return ctx.agentDb; + }, +})); + +// The only stub: the plan floor billing would resolve. A 1-hour window makes a 2-hour watch +// exceed the plan, so the real submit path returns `watch_limit_reached`. Everything else runs. +vi.mock("~/services/dashboardAgentWatchLimits.server", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + resolveWatchPlanLimits: async () => ({ + maxHours: 1, + watchers: actual.UNLIMITED_WATCH_LIMIT, + }), + }; +}); + +process.env.SESSION_SECRET = "test-session-secret-for-watch-limit-status"; +// Unset, watch creation stops at `not_configured` (501) before the plan floor is read. +process.env.DASHBOARD_AGENT_SECRET_KEY = "test-dashboard-agent-secret"; + +const { action } = + await import("~/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent"); + +/** Replays every migration in order, so a new migration can't leave the suite on a stale schema. */ +async function applyAgentSchema(prisma: PrismaClient) { + const folder = path.resolve(__dirname, "../../../internal-packages/dashboard-agent-db/drizzle"); + for (const name of readdirSync(folder) + .filter((file) => file.endsWith(".sql")) + .sort()) { + const sql = readFileSync(path.join(folder, name), "utf8"); + for (const statement of sql.split("--> statement-breakpoint")) { + const trimmed = statement.trim(); + if (trimmed.length > 0) await prisma.$executeRawUnsafe(trimmed); + } + } +} + +let agentDbClient: DashboardAgentDbClient | undefined; + +async function seed(prisma: PrismaClient) { + const slug = `limit_status_${Math.random().toString(36).slice(2, 10)}`; + const user = await prisma.user.create({ + data: { email: `${slug}@example.com`, authenticationMethod: "MAGIC_LINK" }, + }); + const organization = await prisma.organization.create({ data: { title: slug, slug } }); + await prisma.orgMember.create({ + data: { organizationId: organization.id, userId: user.id, role: "ADMIN" }, + }); + const project = await prisma.project.create({ + data: { name: slug, slug, organizationId: organization.id, externalRef: `proj_${slug}` }, + }); + await prisma.runtimeEnvironment.create({ + data: { + slug: "prod", + type: "PRODUCTION", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_prod_${slug}`, + pkApiKey: `pk_prod_${slug}`, + shortcode: `p${slug.slice(0, 6)}`, + }, + }); + ctx.userId = user.id; + return { user, organization, project }; +} + +// error_recurrence resolves its target with no run/queue read, so the plan floor is the only +// thing standing between a valid submit and a created watch. +const DRAFT = JSON.stringify({ + spec: { + kind: "error_recurrence", + fingerprint: "a1b2c3", + checkEveryMinutes: 5, + maxHours: 2, + note: "ping me if it happens again", + }, + followUp: { investigateOnAttention: false, notifyExternally: false }, +}); + +function submitRequest(slug: string, body: Record) { + const form = new URLSearchParams(body); + return action({ + request: new Request( + `https://app.trigger.dev/resources/orgs/${slug}/projects/${slug}/env/prod/dashboard-agent`, + { + method: "POST", + headers: { "content-type": "application/x-www-form-urlencoded" }, + body: form.toString(), + } + ), + params: { organizationSlug: slug, projectParam: slug, envParam: "prod" }, + context: {}, + } as never) as Promise; +} + +afterEach(async () => { + await agentDbClient?.close(); + agentDbClient = undefined; +}); + +describe("the watch card submit's status for a plan-limit refusal", () => { + postgresTest( + "answers 409, not 500, when the window is longer than the plan allows", + async ({ prisma, postgresContainer }) => { + ctx.prisma = prisma; + await applyAgentSchema(prisma); + agentDbClient = createDashboardAgentDb(postgresContainer.getConnectionUri(), { max: 4 }); + ctx.agentDb = agentDbClient.db; + + const seeded = await seed(prisma); + + const response = await submitRequest(seeded.organization.slug, { + intent: "watch-create", + draft: DRAFT, + clientRequestId: "wreq_limit_1", + }); + + expect(response.status).toBe(409); + expect(await response.json()).toMatchObject({ code: "watch_limit_reached" }); + }, + 30_000 + ); +}); diff --git a/apps/webapp/test/dashboardAgentWatchLimits.test.ts b/apps/webapp/test/dashboardAgentWatchLimits.test.ts new file mode 100644 index 00000000000..7bd8f4ee223 --- /dev/null +++ b/apps/webapp/test/dashboardAgentWatchLimits.test.ts @@ -0,0 +1,411 @@ +import { + countActiveWatchesForOrg, + createChat, + createDashboardAgentDb, + listActiveWatchesForChat, + type DashboardAgentDb, + type DashboardAgentDbClient, +} from "@internal/dashboard-agent-db"; +import type { WatchSpec } from "@internal/dashboard-agent-contracts"; +import type * as TriggerSdk from "@trigger.dev/sdk"; +import { postgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { readdirSync, readFileSync } from "node:fs"; +import path from "node:path"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import type { WatchCheckDeps, WatchRunRow } from "~/services/dashboardAgentWatchChecks"; +import type { WatchPlanLimits } from "~/services/dashboardAgentWatchLimits.server"; + +vi.setConfig({ testTimeout: 60_000 }); + +const ctx = vi.hoisted(() => ({ + prisma: undefined as unknown as PrismaClient, + agentDb: undefined as unknown as DashboardAgentDb, +})); + +vi.mock("~/db.server", () => { + const proxy = new Proxy( + {}, + { get: (_target, prop) => (ctx.prisma as unknown as Record)[prop as string] } + ); + return { prisma: proxy, $replica: proxy, sqlDatabaseSchema: undefined }; +}); + +vi.mock("~/services/dashboardAgentDb.server", () => ({ + get dashboardAgentDb() { + return ctx.agentDb; + }, +})); + +vi.mock("@trigger.dev/sdk", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + TriggerClient: class { + tasks = { trigger: async () => ({ id: "run_test" }) }; + }, + }; +}); + +process.env.SESSION_SECRET = "test-session-secret-for-watch-limits"; + +const { createDashboardAgentWatch } = await import("~/services/dashboardAgentWatches.server"); +const { effectiveWatchMaxHours, resolveWatchPlanLimits, watchLimitHint, UNLIMITED_WATCH_LIMIT } = + await import("~/services/dashboardAgentWatchLimits.server"); +const { limitValueAllowingZero } = await import("~/services/platform.v3.server"); + +async function applyAgentSchema(prisma: PrismaClient) { + const folder = path.resolve(__dirname, "../../../internal-packages/dashboard-agent-db/drizzle"); + const migrations = readdirSync(folder) + .filter((file) => file.endsWith(".sql")) + .sort(); + for (const name of migrations) { + const sql = readFileSync(path.join(folder, name), "utf8"); + for (const statement of sql.split("--> statement-breakpoint")) { + const trimmed = statement.trim(); + if (trimmed.length > 0) await prisma.$executeRawUnsafe(trimmed); + } + } +} + +let agentDbClient: DashboardAgentDbClient | undefined; + +async function boot(prisma: PrismaClient, connectionUri: string) { + ctx.prisma = prisma; + await applyAgentSchema(prisma); + agentDbClient = createDashboardAgentDb(connectionUri, { max: 4 }); + ctx.agentDb = agentDbClient.db; +} + +async function seed(prisma: PrismaClient, slugBase: string) { + const slug = `${slugBase}_${Math.random().toString(36).slice(2, 10)}`; + const user = await prisma.user.create({ + data: { email: `${slug}@example.com`, authenticationMethod: "MAGIC_LINK" }, + }); + const organization = await prisma.organization.create({ data: { title: slug, slug } }); + await prisma.orgMember.create({ + data: { organizationId: organization.id, userId: user.id, role: "ADMIN" }, + }); + const project = await prisma.project.create({ + data: { name: slug, slug, organizationId: organization.id, externalRef: `proj_${slug}` }, + }); + const environment = await prisma.runtimeEnvironment.create({ + data: { + slug: "prod", + type: "PRODUCTION", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_prod_${slug}`, + pkApiKey: `pk_prod_${slug}`, + shortcode: `p${slug.slice(0, 6)}`, + }, + }); + return { user, organization, project, environment }; +} + +type Seeded = Awaited>; + +function authenticated(seeded: Seeded) { + return { + id: seeded.environment.id, + organizationId: seeded.organization.id, + projectId: seeded.project.id, + slug: "prod", + type: "PRODUCTION", + project: { id: seeded.project.id, externalRef: seeded.project.externalRef }, + organization: { id: seeded.organization.id, slug: seeded.organization.slug }, + } as any; +} + +async function seedChat(seeded: Seeded, chatId: string) { + await createChat(ctx.agentDb, { + id: chatId, + organizationId: seeded.organization.id, + userId: seeded.user.id, + }); + return chatId; +} + +function runRow(overrides: Partial = {}): WatchRunRow { + return { + friendlyId: "run_1", + status: "PENDING", + queue: "task/my-task", + createdAt: new Date(), + queuedAt: null, + startedAt: null, + completedAt: null, + delayUntil: null, + ...overrides, + }; +} + +function fakeCheckDeps(overrides: Partial = {}): WatchCheckDeps { + return { + readRun: async () => runRow(), + queueExists: async () => true, + readQueueDepth: async () => ({ depth: 7, source: "live_queue", current: true }), + readQueueOldestAge: async () => ({ ageMs: 30_000, source: "live_queue", current: true }), + readErrorRecurrence: async () => null, + readHealth: async () => ({ trustworthy: true, severity: "warn" }), + ...overrides, + }; +} + +const UNLIMITED: WatchPlanLimits = { + maxHours: UNLIMITED_WATCH_LIMIT, + watchers: UNLIMITED_WATCH_LIMIT, +}; + +function runStart(runId: string, maxHours = 2): WatchSpec { + return { kind: "run_start", runId, checkEveryMinutes: 1, maxHours, note: "tell me" }; +} + +function create(args: { + seeded: Seeded; + spec: WatchSpec; + chatId: string; + limits?: WatchPlanLimits; + billingConfigured?: boolean; + countActiveWatches?: (organizationId: string) => Promise; + checkDeps?: Partial; +}) { + return createDashboardAgentWatch({ + environment: authenticated(args.seeded), + userId: args.seeded.user.id, + chatId: args.chatId, + spec: args.spec, + deps: { + configured: () => true, + checkDeps: () => fakeCheckDeps(args.checkDeps), + scheduleTick: async () => {}, + resolveLimits: async () => args.limits ?? UNLIMITED, + ...(args.countActiveWatches ? { countActiveWatches: args.countActiveWatches } : {}), + ...(args.billingConfigured === undefined + ? {} + : { billingConfigured: () => args.billingConfigured! }), + }, + }); +} + +afterEach(async () => { + await agentDbClient?.close(); + agentDbClient = undefined; +}); + +describe("watch plan limits (pure)", () => { + it("caps the window ceiling at the code ceiling of 24 hours", () => { + expect(effectiveWatchMaxHours(100)).toBe(24); + expect(effectiveWatchMaxHours(1)).toBe(1); + expect(effectiveWatchMaxHours(0.5)).toBe(0.5); + }); + + it("reads a plan limit of zero as zero, not as an absent limit", async () => { + // The read the cached platform limit performs: a plan that switched watches off must not + // fall back to the unlimited sentinel. + expect( + limitValueAllowingZero( + { agentWatchMaxHours: 0 } as never, + "agentWatchMaxHours" as never, + UNLIMITED_WATCH_LIMIT + ) + ).toBe(0); + expect( + limitValueAllowingZero(undefined, "agentWatchMaxHours" as never, UNLIMITED_WATCH_LIMIT) + ).toBe(UNLIMITED_WATCH_LIMIT); + + expect(await resolveWatchPlanLimits("org_1", async () => 0)).toEqual({ + maxHours: 0, + watchers: 0, + }); + }); + + it("adds an upgrade nudge only when billing is configured", () => { + expect(watchLimitHint("too long.", true)).toBe("too long. Upgrade your plan for more."); + expect(watchLimitHint("too long.", false)).toBe("too long."); + }); +}); + +describe("createDashboardAgentWatch plan enforcement", () => { + postgresTest( + "refuses a window longer than the plan allows", + async ({ prisma, postgresContainer }) => { + await boot(prisma, postgresContainer.getConnectionUri()); + const seeded = await seed(prisma, "window"); + await seedChat(seeded, "chat_1"); + + const result = await create({ + seeded, + chatId: "chat_1", + spec: runStart("run_1", 2), + limits: { maxHours: 1, watchers: UNLIMITED_WATCH_LIMIT }, + billingConfigured: true, + }); + + expect(result).toMatchObject({ ok: false, code: "watch_limit_reached" }); + if (result.ok) return; + expect(result.error).toContain("Upgrade your plan"); + expect(await listActiveWatchesForChat(ctx.agentDb, { chatId: "chat_1" })).toHaveLength(0); + } + ); + + postgresTest( + "creates a watch whose window is within the plan", + async ({ prisma, postgresContainer }) => { + await boot(prisma, postgresContainer.getConnectionUri()); + const seeded = await seed(prisma, "within"); + await seedChat(seeded, "chat_1"); + + const result = await create({ + seeded, + chatId: "chat_1", + spec: runStart("run_1", 1), + limits: { maxHours: 1, watchers: UNLIMITED_WATCH_LIMIT }, + }); + + expect(result.ok).toBe(true); + expect(await listActiveWatchesForChat(ctx.agentDb, { chatId: "chat_1" })).toHaveLength(1); + } + ); + + postgresTest( + "refuses once the org is at its watcher count, counting active watches for real", + async ({ prisma, postgresContainer }) => { + await boot(prisma, postgresContainer.getConnectionUri()); + const seeded = await seed(prisma, "count"); + await seedChat(seeded, "chat_1"); + await seedChat(seeded, "chat_2"); + + const limits: WatchPlanLimits = { maxHours: UNLIMITED_WATCH_LIMIT, watchers: 1 }; + + const first = await create({ seeded, chatId: "chat_1", spec: runStart("run_1"), limits }); + expect(first.ok).toBe(true); + expect( + await countActiveWatchesForOrg(ctx.agentDb, { organizationId: seeded.organization.id }) + ).toBe(1); + + const second = await create({ seeded, chatId: "chat_2", spec: runStart("run_2"), limits }); + expect(second).toMatchObject({ ok: false, code: "watch_limit_reached" }); + expect(await listActiveWatchesForChat(ctx.agentDb, { chatId: "chat_2" })).toHaveLength(0); + } + ); + + postgresTest( + "fails open: an absent limit resolves to unlimited and a 2h watch is created", + async ({ prisma, postgresContainer }) => { + await boot(prisma, postgresContainer.getConnectionUri()); + const seeded = await seed(prisma, "failopen"); + await seedChat(seeded, "chat_1"); + + const result = await create({ + seeded, + chatId: "chat_1", + spec: runStart("run_1", 2), + limits: UNLIMITED, + }); + + expect(result.ok).toBe(true); + expect(await listActiveWatchesForChat(ctx.agentDb, { chatId: "chat_1" })).toHaveLength(1); + } + ); + + postgresTest( + "leaves no upgrade nudge on a refusal when billing is unconfigured", + async ({ prisma, postgresContainer }) => { + await boot(prisma, postgresContainer.getConnectionUri()); + const seeded = await seed(prisma, "selfhosted"); + await seedChat(seeded, "chat_1"); + + const result = await create({ + seeded, + chatId: "chat_1", + spec: runStart("run_1", 2), + limits: { maxHours: 1, watchers: UNLIMITED_WATCH_LIMIT }, + billingConfigured: false, + }); + + expect(result).toMatchObject({ ok: false, code: "watch_limit_reached" }); + if (result.ok) return; + expect(result.error).not.toContain("Upgrade"); + } + ); + + postgresTest( + "a plan window of zero hours refuses every watch", + async ({ prisma, postgresContainer }) => { + await boot(prisma, postgresContainer.getConnectionUri()); + const seeded = await seed(prisma, "zerohours"); + await seedChat(seeded, "chat_1"); + + const result = await create({ + seeded, + chatId: "chat_1", + spec: runStart("run_1", 1), + limits: { maxHours: 0, watchers: UNLIMITED_WATCH_LIMIT }, + }); + + expect(result).toMatchObject({ ok: false, code: "watch_limit_reached" }); + expect(await listActiveWatchesForChat(ctx.agentDb, { chatId: "chat_1" })).toHaveLength(0); + } + ); + + postgresTest( + "a plan of zero watchers refuses creation", + async ({ prisma, postgresContainer }) => { + await boot(prisma, postgresContainer.getConnectionUri()); + const seeded = await seed(prisma, "zerowatchers"); + await seedChat(seeded, "chat_1"); + + const result = await create({ + seeded, + chatId: "chat_1", + spec: runStart("run_1", 1), + limits: { maxHours: UNLIMITED_WATCH_LIMIT, watchers: 0 }, + }); + + expect(result).toMatchObject({ ok: false, code: "watch_limit_reached" }); + expect(await listActiveWatchesForChat(ctx.agentDb, { chatId: "chat_1" })).toHaveLength(0); + } + ); + + postgresTest( + "answers a condition that already happened, instead of refusing the window", + async ({ prisma, postgresContainer }) => { + await boot(prisma, postgresContainer.getConnectionUri()); + const seeded = await seed(prisma, "instant"); + await seedChat(seeded, "chat_1"); + + const result = await create({ + seeded, + chatId: "chat_1", + spec: runStart("run_1", 2), + limits: { maxHours: 1, watchers: 0 }, + billingConfigured: true, + checkDeps: { + readRun: async () => runRow({ status: "EXECUTING", startedAt: new Date() }), + }, + }); + + expect(result).toMatchObject({ ok: true, watching: false }); + expect(await listActiveWatchesForChat(ctx.agentDb, { chatId: "chat_1" })).toHaveLength(0); + } + ); + + postgresTest( + "min semantics: a plan of 100 hours still permits only up to the 24h ceiling", + async ({ prisma, postgresContainer }) => { + await boot(prisma, postgresContainer.getConnectionUri()); + const seeded = await seed(prisma, "minsem"); + await seedChat(seeded, "chat_1"); + + const created = await create({ + seeded, + chatId: "chat_1", + spec: runStart("run_1", 24), + limits: { maxHours: 100, watchers: UNLIMITED_WATCH_LIMIT }, + }); + expect(created.ok).toBe(true); + expect(await listActiveWatchesForChat(ctx.agentDb, { chatId: "chat_1" })).toHaveLength(1); + } + ); +}); diff --git a/apps/webapp/test/queryRouteReadOnly.test.ts b/apps/webapp/test/queryRouteReadOnly.test.ts new file mode 100644 index 00000000000..5cd08d65e6a --- /dev/null +++ b/apps/webapp/test/queryRouteReadOnly.test.ts @@ -0,0 +1,181 @@ +import { generateJWT } from "@trigger.dev/core/v3/jwt"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +/** + * The query API is read-only, and the grammar is what enforces it. A parser test alone would + * stay green if the route ever compiled agent SQL somewhere else, so these drive the real route + * with a real signed environment JWT and stub only the ClickHouse client. A write must be + * refused before anything reaches ClickHouse. + */ + +const ENVIRONMENT_ID = "env_1234"; +const API_KEY = "tr_dev_abcdefghijklmnop"; + +const environment = { + id: ENVIRONMENT_ID, + type: "DEVELOPMENT", + slug: "dev", + branchName: null, + apiKey: API_KEY, + organizationId: "org_1", + projectId: "proj_1", + archivedAt: null, + concurrencyLimitBurstFactor: { toNumber: () => 1 }, + maximumConcurrencyLimit: 10, + project: { id: "proj_1", externalRef: "proj_ref", deletedAt: null }, + organization: { id: "org_1" }, + orgMember: null, + parentEnvironment: null, +}; + +const mocks = vi.hoisted(() => ({ + runtimeEnvironmentFindFirst: vi.fn(), + queryWithStats: vi.fn(), + customerQueryCreate: vi.fn(), + concurrencyAcquire: vi.fn(), +})); + +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: mocks.concurrencyAcquire, + 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 {} })); + +import { action } from "~/routes/api.v1.query"; +import { executeQuery } from "~/services/queryService.server"; + +/** The claims the env-JWT exchange mints (api.v1.projects.$projectRef.$env.jwt.ts). */ +function mintEnvJwt(scopes: string[]) { + return generateJWT({ + secretKey: API_KEY, + payload: { + sub: ENVIRONMENT_ID, + pub: true, + scopes, + act: { sub: "usr_1", client: "dashboard-agent" }, + }, + expirationTime: "1h", + }); +} + +async function runQuery(query: string): Promise<{ status: number; body: any }> { + const jwt = await mintEnvJwt(["read:query"]); + const response = await action({ + request: new Request("https://api.trigger.dev/api/v1/query", { + method: "POST", + headers: { Authorization: `Bearer ${jwt}`, "Content-Type": "application/json" }, + body: JSON.stringify({ query }), + }), + params: {}, + context: {}, + } as any); + return { status: response.status, body: await response.json() }; +} + +describe("the query API route", () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.runtimeEnvironmentFindFirst.mockResolvedValue(environment); + mocks.customerQueryCreate.mockResolvedValue({ id: "cq_1" }); + mocks.concurrencyAcquire.mockResolvedValue({ success: true }); + mocks.queryWithStats.mockReturnValue(async () => [null, { rows: [], stats: {} }]); + }); + + // Pins the seam the two refusals assert against: a read really does reach ClickHouse here, + // so `not.toHaveBeenCalled()` below means refused, not unreachable. + it("runs a read against ClickHouse", async () => { + const result = await runQuery("SELECT count() FROM runs"); + + expect(result.status).toBe(200); + expect(mocks.queryWithStats).toHaveBeenCalled(); + }); + + it("refuses a write smuggled in as a second statement", async () => { + const result = await runQuery("SELECT 1 FROM runs; DROP TABLE runs"); + + expect(result.status).toBe(400); + expect(mocks.queryWithStats).not.toHaveBeenCalled(); + }); + + it("refuses a mutating statement", async () => { + const result = await runQuery("INSERT INTO runs (task_identifier) VALUES ('x')"); + + expect(result.status).toBe(400); + expect(mocks.queryWithStats).not.toHaveBeenCalled(); + }); + + // A busy service is not a bad query: 400 would tell a caller to rewrite a query that was fine. + it("answers a concurrency rejection with 429", async () => { + mocks.concurrencyAcquire.mockResolvedValue({ success: false, reason: "key_limit" }); + + const result = await runQuery("SELECT count() FROM runs"); + + expect(result.status).toBe(429); + expect(result.body.error).toContain("try again later"); + }); +}); + +describe("the query service", () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.concurrencyAcquire.mockResolvedValue({ success: true }); + mocks.queryWithStats.mockReturnValue(async () => [null, { rows: [], stats: {} }]); + }); + + it("keeps ClickHouse read-only when a caller overrides the settings", async () => { + await executeQuery({ + name: "test-query", + query: "SELECT count() FROM runs", + scope: "environment", + organizationId: "org_1", + projectId: "proj_1", + environmentId: ENVIRONMENT_ID, + clickhouseSettings: { readonly: "0" }, + } as any); + + expect(mocks.queryWithStats).toHaveBeenCalled(); + expect(mocks.queryWithStats.mock.calls[0][0].settings.readonly).toBe("1"); + }); +}); diff --git a/docs/v3-openapi.yaml b/docs/v3-openapi.yaml index bb42f7af2ca..a97ae70307e 100644 --- a/docs/v3-openapi.yaml +++ b/docs/v3-openapi.yaml @@ -1203,6 +1203,16 @@ paths: description: Error message describing the query error "401": description: Unauthorized - API key is missing or invalid + "429": + description: Query service is busy or rate limited - retry shortly + content: + application/json: + schema: + type: object + properties: + error: + type: string + description: Error message describing why the query was turned away "500": description: Internal server error during query execution tags: diff --git a/internal-packages/dashboard-agent-db/drizzle/0004_stale_corsair.sql b/internal-packages/dashboard-agent-db/drizzle/0004_stale_corsair.sql new file mode 100644 index 00000000000..d581b4fec3f --- /dev/null +++ b/internal-packages/dashboard-agent-db/drizzle/0004_stale_corsair.sql @@ -0,0 +1,8 @@ +CREATE TABLE "trigger_dashboard_agent"."agent_message_usage" ( + "organization_id" text NOT NULL, + "period" text NOT NULL, + "count" integer DEFAULT 0 NOT NULL, + "created_at" timestamp with time zone DEFAULT now() NOT NULL, + "updated_at" timestamp with time zone DEFAULT now() NOT NULL, + CONSTRAINT "agent_message_usage_organization_id_period_pk" PRIMARY KEY("organization_id","period") +); diff --git a/internal-packages/dashboard-agent-db/drizzle/0005_ambitious_mordo.sql b/internal-packages/dashboard-agent-db/drizzle/0005_ambitious_mordo.sql new file mode 100644 index 00000000000..3d56343452d --- /dev/null +++ b/internal-packages/dashboard-agent-db/drizzle/0005_ambitious_mordo.sql @@ -0,0 +1,2 @@ +ALTER TABLE "trigger_dashboard_agent"."investigations" ADD COLUMN "sweep_attempts" integer DEFAULT 0 NOT NULL;--> statement-breakpoint +ALTER TABLE "trigger_dashboard_agent"."investigations" ADD COLUMN "last_sweep_attempt_at" timestamp with time zone; \ No newline at end of file diff --git a/internal-packages/dashboard-agent-db/drizzle/meta/0004_snapshot.json b/internal-packages/dashboard-agent-db/drizzle/meta/0004_snapshot.json new file mode 100644 index 00000000000..5ee40ae24f0 --- /dev/null +++ b/internal-packages/dashboard-agent-db/drizzle/meta/0004_snapshot.json @@ -0,0 +1,1344 @@ +{ + "id": "f7cbfef4-7fc8-4deb-8da2-59248b242a60", + "prevId": "efb6f8b8-af9f-4ba7-9e38-bafd1f430b28", + "version": "7", + "dialect": "postgresql", + "tables": { + "trigger_dashboard_agent.agent_message_usage": { + "name": "agent_message_usage", + "schema": "trigger_dashboard_agent", + "columns": { + "organization_id": { + "name": "organization_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "period": { + "name": "period", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "count": { + "name": "count", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "agent_message_usage_organization_id_period_pk": { + "name": "agent_message_usage_organization_id_period_pk", + "columns": ["organization_id", "period"] + } + }, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.chat_messages": { + "name": "chat_messages", + "schema": "trigger_dashboard_agent", + "columns": { + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "message_id": { + "name": "message_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "position": { + "name": "position", + "type": "integer", + "primaryKey": false, + "notNull": true + }, + "role": { + "name": "role", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "message": { + "name": "message", + "type": "jsonb", + "primaryKey": false, + "notNull": true + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "chat_messages_chat_user_role_idx": { + "name": "chat_messages_chat_user_role_idx", + "columns": [ + { + "expression": "chat_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "message_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"chat_messages\".\"role\" = 'user'", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": { + "chat_messages_chat_id_message_id_pk": { + "name": "chat_messages_chat_id_message_id_pk", + "columns": ["chat_id", "message_id"] + } + }, + "uniqueConstraints": { + "chat_messages_chat_position_key": { + "name": "chat_messages_chat_position_key", + "nullsNotDistinct": false, + "columns": ["chat_id", "position"] + } + }, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.chat_sessions": { + "name": "chat_sessions", + "schema": "trigger_dashboard_agent", + "columns": { + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": true, + "notNull": true + }, + "public_access_token": { + "name": "public_access_token", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "last_event_id": { + "name": "last_event_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.chat_turn_evals": { + "name": "chat_turn_evals", + "schema": "trigger_dashboard_agent", + "columns": { + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "turn": { + "name": "turn", + "type": "integer", + "primaryKey": false, + "notNull": true + }, + "organization_id": { + "name": "organization_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "agent_run_id": { + "name": "agent_run_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "eval_run_id": { + "name": "eval_run_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "project_ref": { + "name": "project_ref", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "environment": { + "name": "environment", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "current_page": { + "name": "current_page", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "prompt_slug": { + "name": "prompt_slug", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "prompt_version": { + "name": "prompt_version", + "type": "integer", + "primaryKey": false, + "notNull": false + }, + "tools_used": { + "name": "tools_used", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'[]'::jsonb" + }, + "tool_error": { + "name": "tool_error", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "judge_model": { + "name": "judge_model", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "score_grounded": { + "name": "score_grounded", + "type": "smallint", + "primaryKey": false, + "notNull": false + }, + "score_answered": { + "name": "score_answered", + "type": "smallint", + "primaryKey": false, + "notNull": false + }, + "score_concise": { + "name": "score_concise", + "type": "smallint", + "primaryKey": false, + "notNull": false + }, + "passed": { + "name": "passed", + "type": "boolean", + "primaryKey": false, + "notNull": false + }, + "intent_category": { + "name": "intent_category", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "outcome": { + "name": "outcome", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "sentiment": { + "name": "sentiment", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "capability_gap": { + "name": "capability_gap", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "docs_gap": { + "name": "docs_gap", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "support_opportunity": { + "name": "support_opportunity", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "feature_request": { + "name": "feature_request", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "topics": { + "name": "topics", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'[]'::jsonb" + }, + "signals": { + "name": "signals", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'[]'::jsonb" + }, + "summary": { + "name": "summary", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "user_text": { + "name": "user_text", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "judge": { + "name": "judge", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "chat_turn_evals_org_created_idx": { + "name": "chat_turn_evals_org_created_idx", + "columns": [ + { + "expression": "organization_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "created_at", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "chat_turn_evals_created_idx": { + "name": "chat_turn_evals_created_idx", + "columns": [ + { + "expression": "created_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "chat_turn_evals_org_opps_idx": { + "name": "chat_turn_evals_org_opps_idx", + "columns": [ + { + "expression": "organization_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "created_at", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"chat_turn_evals\".\"capability_gap\" or \"trigger_dashboard_agent\".\"chat_turn_evals\".\"docs_gap\" or \"trigger_dashboard_agent\".\"chat_turn_evals\".\"support_opportunity\" or \"trigger_dashboard_agent\".\"chat_turn_evals\".\"feature_request\"", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": { + "chat_turn_evals_chat_id_turn_pk": { + "name": "chat_turn_evals_chat_id_turn_pk", + "columns": ["chat_id", "turn"] + } + }, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.chats": { + "name": "chats", + "schema": "trigger_dashboard_agent", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true + }, + "organization_id": { + "name": "organization_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "title": { + "name": "title", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'New chat'" + }, + "metadata": { + "name": "metadata", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'{}'::jsonb" + }, + "messages": { + "name": "messages", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'[]'::jsonb" + }, + "pinned_at": { + "name": "pinned_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "last_read_at": { + "name": "last_read_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "deleted_at": { + "name": "deleted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "last_message_at": { + "name": "last_message_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "next_message_position": { + "name": "next_message_position", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 1 + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "chats_org_user_last_msg_idx": { + "name": "chats_org_user_last_msg_idx", + "columns": [ + { + "expression": "organization_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "user_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "last_message_at", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"chats\".\"deleted_at\" is null", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.investigations": { + "name": "investigations", + "schema": "trigger_dashboard_agent", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true + }, + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "project_ref": { + "name": "project_ref", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "environment_ref": { + "name": "environment_ref", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "revision": { + "name": "revision", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "state": { + "name": "state", + "type": "jsonb", + "primaryKey": false, + "notNull": true + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "investigations_chat_idx": { + "name": "investigations_chat_idx", + "columns": [ + { + "expression": "chat_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "investigations_open_updated_idx": { + "name": "investigations_open_updated_idx", + "columns": [ + { + "expression": "updated_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"investigations\".\"state\"->>'outcome' = 'in_progress'", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.watch_batches": { + "name": "watch_batches", + "schema": "trigger_dashboard_agent", + "columns": { + "environment_id": { + "name": "environment_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "cadence_minutes": { + "name": "cadence_minutes", + "type": "integer", + "primaryKey": false, + "notNull": true + }, + "epoch": { + "name": "epoch", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "generation": { + "name": "generation", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'running'" + }, + "armed_at": { + "name": "armed_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "last_tick_at": { + "name": "last_tick_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "watch_batches_environment_id_cadence_minutes_pk": { + "name": "watch_batches_environment_id_cadence_minutes_pk", + "columns": ["environment_id", "cadence_minutes"] + } + }, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.watch_submissions": { + "name": "watch_submissions", + "schema": "trigger_dashboard_agent", + "columns": { + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "client_request_id": { + "name": "client_request_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "organization_id": { + "name": "organization_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "project_id": { + "name": "project_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "environment_id": { + "name": "environment_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "draft_hash": { + "name": "draft_hash", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "draft": { + "name": "draft", + "type": "jsonb", + "primaryKey": false, + "notNull": true + }, + "state": { + "name": "state", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'pending'" + }, + "watch_id": { + "name": "watch_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "unavailable": { + "name": "unavailable", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "external_notification_status": { + "name": "external_notification_status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'not_requested'" + }, + "external_notification_reason": { + "name": "external_notification_reason", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "immediate_result": { + "name": "immediate_result", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "refusal_code": { + "name": "refusal_code", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "refusal_error": { + "name": "refusal_error", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "refusal_existing_id": { + "name": "refusal_existing_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "watch_submissions_created_idx": { + "name": "watch_submissions_created_idx", + "columns": [ + { + "expression": "created_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": { + "watch_submissions_chat_id_client_request_id_pk": { + "name": "watch_submissions_chat_id_client_request_id_pk", + "columns": ["chat_id", "client_request_id"] + } + }, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.watches": { + "name": "watches", + "schema": "trigger_dashboard_agent", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true + }, + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "identity": { + "name": "identity", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "spec": { + "name": "spec", + "type": "jsonb", + "primaryKey": false, + "notNull": true + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'active'" + }, + "delivery_status": { + "name": "delivery_status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'not_required'" + }, + "cancel_reason": { + "name": "cancel_reason", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "resolution": { + "name": "resolution", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "observed_outcome": { + "name": "observed_outcome", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "investigate_on_attention": { + "name": "investigate_on_attention", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "organization_id": { + "name": "organization_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "project_id": { + "name": "project_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "project_ref": { + "name": "project_ref", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "environment_id": { + "name": "environment_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "expires_at": { + "name": "expires_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true + }, + "last_checked_at": { + "name": "last_checked_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "last_attempted_at": { + "name": "last_attempted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "fired_at": { + "name": "fired_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "delivery_claimed_at": { + "name": "delivery_claimed_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "delivery_claim_id": { + "name": "delivery_claim_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "delivered_at": { + "name": "delivered_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "cancelled_at": { + "name": "cancelled_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "last_result": { + "name": "last_result", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "tick_count": { + "name": "tick_count", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "alert_dispatch_key": { + "name": "alert_dispatch_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "retention_at": { + "name": "retention_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false, + "generated": { + "as": "greatest(delivered_at, cancelled_at, fired_at, last_checked_at, created_at)", + "type": "stored" + } + }, + "cadence_minutes": { + "name": "cadence_minutes", + "type": "integer", + "primaryKey": false, + "notNull": false, + "generated": { + "as": "((spec ->> 'checkEveryMinutes')::int)", + "type": "stored" + } + } + }, + "indexes": { + "watches_chat_idx": { + "name": "watches_chat_idx", + "columns": [ + { + "expression": "chat_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_chat_active_identity_key": { + "name": "watches_chat_active_identity_key", + "columns": [ + { + "expression": "chat_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "project_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "environment_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "identity", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "where": "\"trigger_dashboard_agent\".\"watches\".\"status\" = 'active'", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_status_expires_idx": { + "name": "watches_status_expires_idx", + "columns": [ + { + "expression": "status", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "expires_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_pending_delivery_idx": { + "name": "watches_pending_delivery_idx", + "columns": [ + { + "expression": "fired_at", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "last_checked_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"delivery_status\" in ('pending', 'delivering')", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_org_user_wake_idx": { + "name": "watches_org_user_wake_idx", + "columns": [ + { + "expression": "organization_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "user_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "coalesce(\"fired_at\", \"last_checked_at\") desc", + "asc": true, + "isExpression": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"delivery_status\" = 'delivered' and \"trigger_dashboard_agent\".\"watches\".\"status\" in ('fired', 'expired')", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_org_user_active_idx": { + "name": "watches_org_user_active_idx", + "columns": [ + { + "expression": "organization_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "user_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"status\" = 'active'", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_active_env_cadence_idx": { + "name": "watches_active_env_cadence_idx", + "columns": [ + { + "expression": "environment_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "cadence_minutes", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "coalesce(\"last_attempted_at\", \"last_checked_at\", \"created_at\")", + "asc": true, + "isExpression": true, + "nulls": "last" + }, + { + "expression": "expires_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"status\" = 'active'", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_env_cadence_delivery_idx": { + "name": "watches_env_cadence_delivery_idx", + "columns": [ + { + "expression": "environment_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "cadence_minutes", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "delivery_status", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "coalesce(\"fired_at\", \"last_checked_at\")", + "asc": true, + "isExpression": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"status\" in ('fired', 'expired') and \"trigger_dashboard_agent\".\"watches\".\"delivery_status\" in ('pending', 'delivering')", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_retention_idx": { + "name": "watches_retention_idx", + "columns": [ + { + "expression": "retention_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"status\" in ('fired', 'expired', 'cancelled') and \"trigger_dashboard_agent\".\"watches\".\"delivery_status\" in ('not_required', 'delivered')", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + } + }, + "enums": {}, + "schemas": { + "trigger_dashboard_agent": "trigger_dashboard_agent" + }, + "sequences": {}, + "roles": {}, + "policies": {}, + "views": {}, + "_meta": { + "columns": {}, + "schemas": {}, + "tables": {} + } +} diff --git a/internal-packages/dashboard-agent-db/drizzle/meta/0005_snapshot.json b/internal-packages/dashboard-agent-db/drizzle/meta/0005_snapshot.json new file mode 100644 index 00000000000..b00ae150c57 --- /dev/null +++ b/internal-packages/dashboard-agent-db/drizzle/meta/0005_snapshot.json @@ -0,0 +1,1357 @@ +{ + "id": "9f0a4739-19ca-4a15-82dd-25598116feb9", + "prevId": "f7cbfef4-7fc8-4deb-8da2-59248b242a60", + "version": "7", + "dialect": "postgresql", + "tables": { + "trigger_dashboard_agent.agent_message_usage": { + "name": "agent_message_usage", + "schema": "trigger_dashboard_agent", + "columns": { + "organization_id": { + "name": "organization_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "period": { + "name": "period", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "count": { + "name": "count", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "agent_message_usage_organization_id_period_pk": { + "name": "agent_message_usage_organization_id_period_pk", + "columns": ["organization_id", "period"] + } + }, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.chat_messages": { + "name": "chat_messages", + "schema": "trigger_dashboard_agent", + "columns": { + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "message_id": { + "name": "message_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "position": { + "name": "position", + "type": "integer", + "primaryKey": false, + "notNull": true + }, + "role": { + "name": "role", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "message": { + "name": "message", + "type": "jsonb", + "primaryKey": false, + "notNull": true + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "chat_messages_chat_user_role_idx": { + "name": "chat_messages_chat_user_role_idx", + "columns": [ + { + "expression": "chat_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "message_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"chat_messages\".\"role\" = 'user'", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": { + "chat_messages_chat_id_message_id_pk": { + "name": "chat_messages_chat_id_message_id_pk", + "columns": ["chat_id", "message_id"] + } + }, + "uniqueConstraints": { + "chat_messages_chat_position_key": { + "name": "chat_messages_chat_position_key", + "nullsNotDistinct": false, + "columns": ["chat_id", "position"] + } + }, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.chat_sessions": { + "name": "chat_sessions", + "schema": "trigger_dashboard_agent", + "columns": { + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": true, + "notNull": true + }, + "public_access_token": { + "name": "public_access_token", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "last_event_id": { + "name": "last_event_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.chat_turn_evals": { + "name": "chat_turn_evals", + "schema": "trigger_dashboard_agent", + "columns": { + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "turn": { + "name": "turn", + "type": "integer", + "primaryKey": false, + "notNull": true + }, + "organization_id": { + "name": "organization_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "agent_run_id": { + "name": "agent_run_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "eval_run_id": { + "name": "eval_run_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "project_ref": { + "name": "project_ref", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "environment": { + "name": "environment", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "current_page": { + "name": "current_page", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "prompt_slug": { + "name": "prompt_slug", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "prompt_version": { + "name": "prompt_version", + "type": "integer", + "primaryKey": false, + "notNull": false + }, + "tools_used": { + "name": "tools_used", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'[]'::jsonb" + }, + "tool_error": { + "name": "tool_error", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "judge_model": { + "name": "judge_model", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "score_grounded": { + "name": "score_grounded", + "type": "smallint", + "primaryKey": false, + "notNull": false + }, + "score_answered": { + "name": "score_answered", + "type": "smallint", + "primaryKey": false, + "notNull": false + }, + "score_concise": { + "name": "score_concise", + "type": "smallint", + "primaryKey": false, + "notNull": false + }, + "passed": { + "name": "passed", + "type": "boolean", + "primaryKey": false, + "notNull": false + }, + "intent_category": { + "name": "intent_category", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "outcome": { + "name": "outcome", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "sentiment": { + "name": "sentiment", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "capability_gap": { + "name": "capability_gap", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "docs_gap": { + "name": "docs_gap", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "support_opportunity": { + "name": "support_opportunity", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "feature_request": { + "name": "feature_request", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "topics": { + "name": "topics", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'[]'::jsonb" + }, + "signals": { + "name": "signals", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'[]'::jsonb" + }, + "summary": { + "name": "summary", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "user_text": { + "name": "user_text", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "judge": { + "name": "judge", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "chat_turn_evals_org_created_idx": { + "name": "chat_turn_evals_org_created_idx", + "columns": [ + { + "expression": "organization_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "created_at", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "chat_turn_evals_created_idx": { + "name": "chat_turn_evals_created_idx", + "columns": [ + { + "expression": "created_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "chat_turn_evals_org_opps_idx": { + "name": "chat_turn_evals_org_opps_idx", + "columns": [ + { + "expression": "organization_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "created_at", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"chat_turn_evals\".\"capability_gap\" or \"trigger_dashboard_agent\".\"chat_turn_evals\".\"docs_gap\" or \"trigger_dashboard_agent\".\"chat_turn_evals\".\"support_opportunity\" or \"trigger_dashboard_agent\".\"chat_turn_evals\".\"feature_request\"", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": { + "chat_turn_evals_chat_id_turn_pk": { + "name": "chat_turn_evals_chat_id_turn_pk", + "columns": ["chat_id", "turn"] + } + }, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.chats": { + "name": "chats", + "schema": "trigger_dashboard_agent", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true + }, + "organization_id": { + "name": "organization_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "title": { + "name": "title", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'New chat'" + }, + "metadata": { + "name": "metadata", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'{}'::jsonb" + }, + "messages": { + "name": "messages", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'[]'::jsonb" + }, + "pinned_at": { + "name": "pinned_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "last_read_at": { + "name": "last_read_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "deleted_at": { + "name": "deleted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "last_message_at": { + "name": "last_message_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "next_message_position": { + "name": "next_message_position", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 1 + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "chats_org_user_last_msg_idx": { + "name": "chats_org_user_last_msg_idx", + "columns": [ + { + "expression": "organization_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "user_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "last_message_at", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"chats\".\"deleted_at\" is null", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.investigations": { + "name": "investigations", + "schema": "trigger_dashboard_agent", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true + }, + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "project_ref": { + "name": "project_ref", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "environment_ref": { + "name": "environment_ref", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "revision": { + "name": "revision", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "state": { + "name": "state", + "type": "jsonb", + "primaryKey": false, + "notNull": true + }, + "sweep_attempts": { + "name": "sweep_attempts", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "last_sweep_attempt_at": { + "name": "last_sweep_attempt_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "investigations_chat_idx": { + "name": "investigations_chat_idx", + "columns": [ + { + "expression": "chat_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "investigations_open_updated_idx": { + "name": "investigations_open_updated_idx", + "columns": [ + { + "expression": "updated_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"investigations\".\"state\"->>'outcome' = 'in_progress'", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.watch_batches": { + "name": "watch_batches", + "schema": "trigger_dashboard_agent", + "columns": { + "environment_id": { + "name": "environment_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "cadence_minutes": { + "name": "cadence_minutes", + "type": "integer", + "primaryKey": false, + "notNull": true + }, + "epoch": { + "name": "epoch", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "generation": { + "name": "generation", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'running'" + }, + "armed_at": { + "name": "armed_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "last_tick_at": { + "name": "last_tick_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "watch_batches_environment_id_cadence_minutes_pk": { + "name": "watch_batches_environment_id_cadence_minutes_pk", + "columns": ["environment_id", "cadence_minutes"] + } + }, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.watch_submissions": { + "name": "watch_submissions", + "schema": "trigger_dashboard_agent", + "columns": { + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "client_request_id": { + "name": "client_request_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "organization_id": { + "name": "organization_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "project_id": { + "name": "project_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "environment_id": { + "name": "environment_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "draft_hash": { + "name": "draft_hash", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "draft": { + "name": "draft", + "type": "jsonb", + "primaryKey": false, + "notNull": true + }, + "state": { + "name": "state", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'pending'" + }, + "watch_id": { + "name": "watch_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "unavailable": { + "name": "unavailable", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "external_notification_status": { + "name": "external_notification_status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'not_requested'" + }, + "external_notification_reason": { + "name": "external_notification_reason", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "immediate_result": { + "name": "immediate_result", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "refusal_code": { + "name": "refusal_code", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "refusal_error": { + "name": "refusal_error", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "refusal_existing_id": { + "name": "refusal_existing_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "watch_submissions_created_idx": { + "name": "watch_submissions_created_idx", + "columns": [ + { + "expression": "created_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": { + "watch_submissions_chat_id_client_request_id_pk": { + "name": "watch_submissions_chat_id_client_request_id_pk", + "columns": ["chat_id", "client_request_id"] + } + }, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.watches": { + "name": "watches", + "schema": "trigger_dashboard_agent", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true + }, + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "identity": { + "name": "identity", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "spec": { + "name": "spec", + "type": "jsonb", + "primaryKey": false, + "notNull": true + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'active'" + }, + "delivery_status": { + "name": "delivery_status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'not_required'" + }, + "cancel_reason": { + "name": "cancel_reason", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "resolution": { + "name": "resolution", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "observed_outcome": { + "name": "observed_outcome", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "investigate_on_attention": { + "name": "investigate_on_attention", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "organization_id": { + "name": "organization_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "project_id": { + "name": "project_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "project_ref": { + "name": "project_ref", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "environment_id": { + "name": "environment_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "expires_at": { + "name": "expires_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true + }, + "last_checked_at": { + "name": "last_checked_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "last_attempted_at": { + "name": "last_attempted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "fired_at": { + "name": "fired_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "delivery_claimed_at": { + "name": "delivery_claimed_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "delivery_claim_id": { + "name": "delivery_claim_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "delivered_at": { + "name": "delivered_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "cancelled_at": { + "name": "cancelled_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "last_result": { + "name": "last_result", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "tick_count": { + "name": "tick_count", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "alert_dispatch_key": { + "name": "alert_dispatch_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "retention_at": { + "name": "retention_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false, + "generated": { + "as": "greatest(delivered_at, cancelled_at, fired_at, last_checked_at, created_at)", + "type": "stored" + } + }, + "cadence_minutes": { + "name": "cadence_minutes", + "type": "integer", + "primaryKey": false, + "notNull": false, + "generated": { + "as": "((spec ->> 'checkEveryMinutes')::int)", + "type": "stored" + } + } + }, + "indexes": { + "watches_chat_idx": { + "name": "watches_chat_idx", + "columns": [ + { + "expression": "chat_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_chat_active_identity_key": { + "name": "watches_chat_active_identity_key", + "columns": [ + { + "expression": "chat_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "project_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "environment_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "identity", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "where": "\"trigger_dashboard_agent\".\"watches\".\"status\" = 'active'", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_status_expires_idx": { + "name": "watches_status_expires_idx", + "columns": [ + { + "expression": "status", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "expires_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_pending_delivery_idx": { + "name": "watches_pending_delivery_idx", + "columns": [ + { + "expression": "fired_at", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "last_checked_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"delivery_status\" in ('pending', 'delivering')", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_org_user_wake_idx": { + "name": "watches_org_user_wake_idx", + "columns": [ + { + "expression": "organization_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "user_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "coalesce(\"fired_at\", \"last_checked_at\") desc", + "asc": true, + "isExpression": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"delivery_status\" = 'delivered' and \"trigger_dashboard_agent\".\"watches\".\"status\" in ('fired', 'expired')", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_org_user_active_idx": { + "name": "watches_org_user_active_idx", + "columns": [ + { + "expression": "organization_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "user_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"status\" = 'active'", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_active_env_cadence_idx": { + "name": "watches_active_env_cadence_idx", + "columns": [ + { + "expression": "environment_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "cadence_minutes", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "coalesce(\"last_attempted_at\", \"last_checked_at\", \"created_at\")", + "asc": true, + "isExpression": true, + "nulls": "last" + }, + { + "expression": "expires_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"status\" = 'active'", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_env_cadence_delivery_idx": { + "name": "watches_env_cadence_delivery_idx", + "columns": [ + { + "expression": "environment_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "cadence_minutes", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "delivery_status", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "coalesce(\"fired_at\", \"last_checked_at\")", + "asc": true, + "isExpression": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"status\" in ('fired', 'expired') and \"trigger_dashboard_agent\".\"watches\".\"delivery_status\" in ('pending', 'delivering')", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_retention_idx": { + "name": "watches_retention_idx", + "columns": [ + { + "expression": "retention_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"status\" in ('fired', 'expired', 'cancelled') and \"trigger_dashboard_agent\".\"watches\".\"delivery_status\" in ('not_required', 'delivered')", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + } + }, + "enums": {}, + "schemas": { + "trigger_dashboard_agent": "trigger_dashboard_agent" + }, + "sequences": {}, + "roles": {}, + "policies": {}, + "views": {}, + "_meta": { + "columns": {}, + "schemas": {}, + "tables": {} + } +} diff --git a/internal-packages/dashboard-agent-db/drizzle/meta/_journal.json b/internal-packages/dashboard-agent-db/drizzle/meta/_journal.json index 9efe7bc1a10..213320fa640 100644 --- a/internal-packages/dashboard-agent-db/drizzle/meta/_journal.json +++ b/internal-packages/dashboard-agent-db/drizzle/meta/_journal.json @@ -29,6 +29,20 @@ "when": 1786264383741, "tag": "0003_backfill_chat_last_read_at", "breakpoints": true + }, + { + "idx": 4, + "version": "7", + "when": 1786359241538, + "tag": "0004_stale_corsair", + "breakpoints": true + }, + { + "idx": 5, + "version": "7", + "when": 1786376934874, + "tag": "0005_ambitious_mordo", + "breakpoints": true } ] } diff --git a/internal-packages/dashboard-agent-db/src/queries.ts b/internal-packages/dashboard-agent-db/src/queries.ts index 40c2a10f55f..fc1ce8f7686 100644 --- a/internal-packages/dashboard-agent-db/src/queries.ts +++ b/internal-packages/dashboard-agent-db/src/queries.ts @@ -8,6 +8,7 @@ import type { DashboardAgentDb } from "./client.js"; import { generateInvestigationId } from "./ids.js"; import { lockChatForWatches, type DashboardAgentDbOrTx } from "./internal.js"; import { + agentMessageUsage, chatMessages, chats, chatSessions, @@ -126,6 +127,45 @@ export async function countUserMessages( return rows[0]?.count ?? 0; } +/** + * The message count for one org in one billing period. Reads the standalone counter, + * never the chat rows, so a deleted chat can't lower it within the period. `period` is + * a UTC calendar month, "YYYY-MM"; the caller chooses it. + */ +export async function getAgentMessageUsage( + db: DashboardAgentDb, + params: { organizationId: string; period: string } +): Promise { + const rows = await db + .select({ count: agentMessageUsage.count }) + .from(agentMessageUsage) + .where( + and( + eq(agentMessageUsage.organizationId, params.organizationId), + eq(agentMessageUsage.period, params.period) + ) + ) + .limit(1); + return rows[0]?.count ?? 0; +} + +/** Bump the counter by one, creating the period row on first use. Returns the new count. */ +export async function incrementAgentMessageUsage( + db: DashboardAgentDb, + params: { organizationId: string; period: string; by?: number } +): Promise { + const by = params.by ?? 1; + const rows = await db + .insert(agentMessageUsage) + .values({ organizationId: params.organizationId, period: params.period, count: by }) + .onConflictDoUpdate({ + target: [agentMessageUsage.organizationId, agentMessageUsage.period], + set: { count: sql`${agentMessageUsage.count} + ${by}`, updatedAt: sql`now()` }, + }) + .returning({ count: agentMessageUsage.count }); + return rows[0]?.count ?? by; +} + /** * Chats whose transcript moved on after their owner last looked. A watch wake is one way * that happens; an answer that landed while the panel was closed is another, and the panel @@ -1045,6 +1085,10 @@ export async function listChatIdsWithOpenInvestigations( /** * Sweep for investigations nothing else settles. `olderThan` is on `updated_at`, * which every revision bumps, so a card a live turn is writing to stays out. + * + * Order is `last_sweep_attempt_at` nulls first, then `updated_at`: a never-attempted + * row is always seen before one a prior sweep already failed on, so a row that can't + * settle rotates to the back instead of pinning the head and starving newer rows. */ export async function listStaleOpenInvestigations( db: DashboardAgentDb, @@ -1062,12 +1106,39 @@ export async function listStaleOpenInvestigations( sql`${investigations.updatedAt} <= ${params.olderThan.toISOString()}::timestamptz` ) ) - .orderBy(investigations.updatedAt) + .orderBy(sql`${investigations.lastSweepAttemptAt} asc nulls first`, investigations.updatedAt) .limit(params.limit ?? 100); return rows.map((row) => row.investigation); } +/** + * Record a failed stale-sweep settle on its own, committed outside the settle tx that + * rolled back. Bumps the attempt count and stamps `last_sweep_attempt_at` — which does + * NOT touch `updated_at`, so the row still reads as stale, only later in the order. + * Returns the new count, or null when the row is no longer `in_progress`. + */ +export async function recordInvestigationSweepAttempt( + db: DashboardAgentDbOrTx, + params: { id: string } +): Promise { + const rows = await db + .update(investigations) + .set({ + sweepAttempts: sql`${investigations.sweepAttempts} + 1`, + lastSweepAttemptAt: sql`now()`, + }) + .where( + and( + eq(investigations.id, params.id), + sql`${investigations.state}->>'outcome' = 'in_progress'` + ) + ) + .returning({ sweepAttempts: investigations.sweepAttempts }); + + return rows[0]?.sweepAttempts ?? null; +} + /** What the settle wrote, which is what the closing card has to render. */ export type SettledInvestigation = { id: string; revision: number; state: unknown }; diff --git a/internal-packages/dashboard-agent-db/src/schema.ts b/internal-packages/dashboard-agent-db/src/schema.ts index 0943f838d1f..d080d759915 100644 --- a/internal-packages/dashboard-agent-db/src/schema.ts +++ b/internal-packages/dashboard-agent-db/src/schema.ts @@ -170,6 +170,10 @@ export const investigations = dashboardAgentSchema.table( // Monotonic; bumped by a single atomic UPDATE. revision: integer("revision").notNull().default(0), state: jsonb("state").$type().notNull(), + // Failed stale-sweep settle attempts. Bumped outside the rolled-back settle tx so a + // row that can't render rotates to the back of the sweep order instead of pinning it. + sweepAttempts: integer("sweep_attempts").notNull().default(0), + lastSweepAttemptAt: timestamp("last_sweep_attempt_at", { withTimezone: true }), createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(), updatedAt: timestamp("updated_at", { withTimezone: true }).notNull().defaultNow(), }, @@ -183,6 +187,23 @@ export const investigations = dashboardAgentSchema.table( ] ); +/** + * Per-(org, period) message counter. Deliberately not joined to chats: deleting a chat + * must not free quota inside the period. `period` is a UTC calendar month, "YYYY-MM". + * Org id is a main-DB id with no FK. + */ +export const agentMessageUsage = dashboardAgentSchema.table( + "agent_message_usage", + { + organizationId: text("organization_id").notNull(), + period: text("period").notNull(), + count: integer("count").notNull().default(0), + createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(), + updatedAt: timestamp("updated_at", { withTimezone: true }).notNull().defaultNow(), + }, + (t) => [primaryKey({ columns: [t.organizationId, t.period] })] +); + export type Chat = typeof chats.$inferSelect; export type NewChat = typeof chats.$inferInsert; export type ChatMessage = typeof chatMessages.$inferSelect; @@ -193,3 +214,5 @@ export type ChatTurnEval = typeof chatTurnEvals.$inferSelect; export type NewChatTurnEval = typeof chatTurnEvals.$inferInsert; export type Investigation = typeof investigations.$inferSelect; export type NewInvestigation = typeof investigations.$inferInsert; +export type AgentMessageUsage = typeof agentMessageUsage.$inferSelect; +export type NewAgentMessageUsage = typeof agentMessageUsage.$inferInsert; diff --git a/internal-packages/dashboard-agent-db/src/watch-queries.ts b/internal-packages/dashboard-agent-db/src/watch-queries.ts index 1f8164f6689..2755adeb803 100644 --- a/internal-packages/dashboard-agent-db/src/watch-queries.ts +++ b/internal-packages/dashboard-agent-db/src/watch-queries.ts @@ -494,6 +494,22 @@ export async function countUnreadWatchWakes( return rows[0]?.count ?? 0; } +/** + * How many active watches an org has, across all its chats and users. The plan-limit floor + * is org-wide, so this is org-scoped only; a chat deletion cancels its watches, so `active` + * is the whole count. + */ +export async function countActiveWatchesForOrg( + db: DashboardAgentDb, + params: { organizationId: string } +): Promise { + const rows = await db + .select({ count: sql`count(*)::int` }) + .from(watches) + .where(and(eq(watches.status, "active"), eq(watches.organizationId, params.organizationId))); + return rows[0]?.count ?? 0; +} + /** * Whether this user has a watch that can still wake them here. Covered by * `watches_org_user_active_idx`; a chat deletion cancels its watches, so `active` is enough. diff --git a/internal-packages/dashboard-agent/src/tool-api-client.ts b/internal-packages/dashboard-agent/src/tool-api-client.ts index 2392fe6d8c8..2458f31c792 100644 --- a/internal-packages/dashboard-agent/src/tool-api-client.ts +++ b/internal-packages/dashboard-agent/src/tool-api-client.ts @@ -42,11 +42,12 @@ const GET_TIMEOUT_MS = 10_000; const JWT_TIMEOUT_MS = 10_000; const QUERY_TIMEOUT_MS = 30_000; -// "query" is the server rejecting the TRQL, "transport" is the request breaking. Chart +// "query" is the server rejecting the TRQL, "transport" is the request breaking, "busy" is +// the server too loaded or rate limited to answer — the same query may work shortly. Chart // validation only fails a render on "query". export type QueryPostResult = | { ok: true; rows: Array> } - | { ok: false; kind: "query" | "transport"; error: string }; + | { ok: false; kind: "query" | "transport" | "busy"; error: string }; export const NO_AUTH = { error: "No delegated access is available for this turn." } as const; @@ -215,6 +216,15 @@ export function createApiClient(ctx: ApiClientContext): DashboardAgentApiClient // The route returns 400 with { error } for invalid TRQL. const data = (await res.json().catch(() => ({}))) as { results?: unknown; error?: string }; if (!res.ok) { + // 429 is the concurrency rejection and the rate limiter: nothing is wrong with the + // query, so it is not a query error. + if (res.status === 429) { + return { + ok: false, + kind: "busy", + error: `${data.error ?? "The query service is busy right now."} You can retry the same query shortly.`, + }; + } return { ok: false, kind: res.status >= 500 ? "transport" : "query", @@ -234,7 +244,7 @@ export function createApiClient(ctx: ApiClientContext): DashboardAgentApiClient ): Promise { const result = await postQuery(query, period); if (isEnvUnavailable(result) || result.ok) return null; - if (result.kind === "transport") { + if (result.kind === "transport" || result.kind === "busy") { logger.warn("Skipped chart query validation", { error: result.error }); return null; } diff --git a/internal-packages/dashboard-agent/src/tool-api-transport.test.ts b/internal-packages/dashboard-agent/src/tool-api-transport.test.ts index 137c43409cf..6c609886811 100644 --- a/internal-packages/dashboard-agent/src/tool-api-transport.test.ts +++ b/internal-packages/dashboard-agent/src/tool-api-transport.test.ts @@ -88,6 +88,45 @@ describe("a broken request reads as a broken request, never as an answer", () => expect(JSON.stringify(result)).not.toContain("isn't locked to a deployed version"); }); + it("classifies a busy query route as busy, not as a bad query", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async (url: string) => + url.endsWith("/jwt") + ? Response.json({ token: "jwt" }) + : Response.json( + { error: "We're experiencing a lot of queries at the moment." }, + { status: 429 } + ) + ) + ); + + const result = await createApiClient(CTX).postQuery("SELECT 1", undefined); + + expect(result).toMatchObject({ ok: false, kind: "busy" }); + expect((result as { error: string }).error).toContain("retry the same query shortly"); + }); + + // The other half of the same invariant: only 429 is busy, so a rejected query still counts. + it("classifies a rejected query as a query error", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async (url: string) => + url.endsWith("/jwt") + ? Response.json({ token: "jwt" }) + : Response.json({ error: "Unknown expression identifier 'createdAt'." }, { status: 400 }) + ) + ); + + const result = await createApiClient(CTX).postQuery("SELECT createdAt FROM runs", undefined); + + expect(result).toMatchObject({ + ok: false, + kind: "query", + error: "Unknown expression identifier 'createdAt'.", + }); + }); + it("still reports a real 404 as the answer it is", async () => { vi.stubGlobal( "fetch", diff --git a/internal-packages/dashboard-agent/src/tool-api.ts b/internal-packages/dashboard-agent/src/tool-api.ts index b043acad12b..e0376015080 100644 --- a/internal-packages/dashboard-agent/src/tool-api.ts +++ b/internal-packages/dashboard-agent/src/tool-api.ts @@ -186,6 +186,9 @@ export function withLiveState(metrics: unknown, queueType: "task" | "custom", li }; } +/** Failed `run_query` calls in a row before the tool tells the model to stop and answer. */ +export const MAX_CONSECUTIVE_QUERY_FAILURES = 3; + export function buildApiTools(args: { ctx: DashboardAgentToolContext; client: DashboardAgentApiClient; @@ -195,6 +198,12 @@ export function buildApiTools(args: { const { userActorToken, projectRef, environmentName, environmentBranch } = ctx; const { origin, hasAuth, envApiGet, postQuery, validateChartQuery } = client; + // A failed query hands the model the database error to fix, and it usually does. When it + // doesn't, the only other limit is the turn's 10 steps, so one broken query can eat the + // whole turn and leave the user with no answer at all. This tool set is built per turn, + // so the counter caps consecutive failures within one turn. + let consecutiveQueryFailures = 0; + return { list_projects: tool({ ...listProjectsSchema, @@ -350,7 +359,20 @@ export function buildApiTools(args: { execute: async ({ query, period }) => { const result = await postQuery(query, period); if (isEnvUnavailable(result)) return envUnavailableError(result, "query"); - if (!result.ok) return { error: result.error }; + if (!result.ok) { + // Only SQL errors count toward the cap; transport and busy errors are transient, + // and the same query may work on a retry. + 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; const cap = 200; const rows = result.rows; return { rows: rows.slice(0, cap), rowCount: rows.length, truncated: rows.length > cap }; diff --git a/internal-packages/dashboard-agent/src/tool-query-retry-cap.test.ts b/internal-packages/dashboard-agent/src/tool-query-retry-cap.test.ts new file mode 100644 index 00000000000..0cf3b2f0a92 --- /dev/null +++ b/internal-packages/dashboard-agent/src/tool-query-retry-cap.test.ts @@ -0,0 +1,122 @@ +import { describe, expect, it, vi } from "vitest"; +import { buildApiTools, MAX_CONSECUTIVE_QUERY_FAILURES } from "./tool-api"; +import type { DashboardAgentApiClient } from "./tool-api-client"; + +/** + * A failed query hands the model the database error to fix. Without a cap, the only other + * limit is the turn's step budget, so a model that keeps rewriting the same broken query + * burns the whole turn and the user gets no answer. After three failures in a row the tool + * tells it to stop and answer. + */ + +function queryTool(postQuery: DashboardAgentApiClient["postQuery"]) { + const client = { + origin: "https://api.example.com", + hasAuth: true, + envApiGet: async () => ({ ok: false as const, status: 500 }), + postQuery, + validateChartQuery: async () => null, + } as unknown as DashboardAgentApiClient; + const tools = buildApiTools({ + ctx: { userActorToken: "uat", apiOrigin: client.origin }, + client, + renderInvestigations: (() => []) as any, + }); + return (query: string) => (tools.run_query as any).execute({ query }, {} as any); +} + +const failure = { + ok: false as const, + kind: "query" as const, + error: "Unknown expression identifier 'createdAt'.", +}; +const transportFailure = { + ok: false as const, + kind: "transport" as const, + error: "The environment is temporarily unavailable.", +}; +const busyFailure = { + ok: false as const, + kind: "busy" as const, + error: "We're experiencing a lot of queries at the moment. You can retry the same query shortly.", +}; +const success = { ok: true as const, rows: [{ n: 1 }] }; + +describe("run_query's consecutive-failure cap", () => { + it("keeps handing back the plain error until the cap", async () => { + const run = queryTool(async () => failure); + + for (let attempt = 1; attempt < MAX_CONSECUTIVE_QUERY_FAILURES; attempt++) { + const result = await run("SELECT createdAt FROM runs"); + expect(result.error).toBe(failure.error); + } + }); + + it("tells the model to stop and answer at the cap", async () => { + const run = queryTool(async () => failure); + + let result: { error: string } = { error: "" }; + for (let attempt = 0; attempt < MAX_CONSECUTIVE_QUERY_FAILURES; attempt++) { + result = await run("SELECT createdAt FROM runs"); + } + + expect(result.error).toContain(failure.error); + expect(result.error).toContain("answer the user with what you already have"); + }); + + it("counts consecutive failures only, so a good query clears the count", async () => { + const postQuery = vi + .fn() + .mockResolvedValueOnce(failure) + .mockResolvedValueOnce(failure) + .mockResolvedValueOnce(success) + .mockResolvedValue(failure); + const run = queryTool(postQuery as any); + + await run("bad"); + await run("bad"); + await run("good"); + const result = await run("bad"); + + expect(result.error).toBe(failure.error); + }); + + it("does not count transport errors toward the cap", async () => { + const run = queryTool(async () => transportFailure); + + let result: { error: string } = { error: "" }; + for (let attempt = 0; attempt < MAX_CONSECUTIVE_QUERY_FAILURES + 2; attempt++) { + result = await run("SELECT createdAt FROM runs"); + } + + expect(result.error).toBe(transportFailure.error); + expect(result.error).not.toContain("answer the user with what you already have"); + }); + + // A "too busy" rejection says nothing about the query, so spending the cap on it would + // stop the model over a queue that clears in seconds. + it("does not count busy rejections toward the cap", async () => { + const run = queryTool(async () => busyFailure); + + let result: { error: string } = { error: "" }; + for (let attempt = 0; attempt < MAX_CONSECUTIVE_QUERY_FAILURES + 2; attempt++) { + result = await run("SELECT createdAt FROM runs"); + } + + expect(result.error).toBe(busyFailure.error); + expect(result.error).not.toContain("answer the user with what you already have"); + }); + + it("still caps real SQL errors that follow busy rejections", async () => { + const postQuery = vi.fn().mockResolvedValueOnce(busyFailure).mockResolvedValue(failure); + const run = queryTool(postQuery as any); + + await run("busy"); + let result: { error: string } = { error: "" }; + for (let attempt = 0; attempt < MAX_CONSECUTIVE_QUERY_FAILURES; attempt++) { + result = await run("SELECT createdAt FROM runs"); + } + + expect(result.error).toContain("answer the user with what you already have"); + }); +});