diff --git a/.changeset/chat-stream-supersede-race.md b/.changeset/chat-stream-supersede-race.md new file mode 100644 index 00000000000..7e2b807d8c3 --- /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. Stopping a chat now also hands it back to your other tabs instead of leaving them read-only. diff --git a/.changeset/watch-mode-keepalive.md b/.changeset/watch-mode-keepalive.md new file mode 100644 index 00000000000..ba6c0441724 --- /dev/null +++ b/.changeset/watch-mode-keepalive.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/sdk": patch +--- + +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. 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"); + }); +}); diff --git a/packages/trigger-sdk/src/v3/chat.test.ts b/packages/trigger-sdk/src/v3/chat.test.ts index 21aa004c9a9..71566457dd8 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 @@ -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, @@ -1176,7 +1205,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 +1234,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(); @@ -1218,6 +1253,455 @@ 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("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 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 } }, + }); + + 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); + // 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(); + } + }); + + 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("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("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(); + } + }); + + 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", () => { it("isReadOnly defaults to false when multiTab is disabled", () => { const transport = new TriggerChatTransport({ @@ -1488,9 +1972,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..b7075e6f71e 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,12 +1174,14 @@ 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 - // 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, }); }; @@ -1205,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 — @@ -1266,7 +1284,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, + }); }; // ------------------------------------------------------------------------- @@ -1780,11 +1802,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 @@ -1796,7 +1820,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; @@ -1804,6 +1831,25 @@ 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 + ) { + // 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)." + ); + } + // Settled close, or the turn is gone — tell the UI instead of // leaving it spinning on a stream nobody will finish. if (state.isStreaming) { @@ -1823,7 +1869,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; @@ -1853,7 +1903,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; @@ -1866,7 +1920,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; } @@ -2001,10 +2059,20 @@ 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) + // would otherwise leave the resubscribe loop running forever. + cancel() { + internalAbort.abort(); + }, }); } } 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" });