From ee8ddb38499349771499fd798296d63b5c864683 Mon Sep 17 00:00:00 2001 From: Jakob Wennberg <149234542+jakobwennberg@users.noreply.github.com> Date: Sun, 26 Jul 2026 17:48:18 +0200 Subject: [PATCH] fix(assistant): stop cross-user conversation access, bricked threads and lost sessions (#1209) * fix(assistant): stop cross-user conversation access, bricked threads and lost sessions Hotfix batch (PR1 of the assistant UI makeover, dev_docs/assistant_redesign_plan.md section 7). No visual change; each of these is wrong today regardless of which design lands, and three are unrecoverable per incident. /api/agent/invoke never checked who owns a resumed conversation_id. RLS on agent_conversations/agent_messages is company-scoped, not user-scoped (20260517204000), so a member could post a colleague's conversation id, have their history loaded into the prompt and read it back, while their own turns were appended to that thread. The conversations list route filters on user_id for exactly this reason. Also pins company and intent: resuming a thread from another company would mix ledgers, and resuming under a different intent would swap the tool whitelist under history the model has already seen. A turn persists the assistant message carrying tool_use blocks before the tools run, and their results only after the batch finishes. Dying in between (client disconnect terminating the function, a deploy, a slow tool) left history ending on an unanswered tool_use, which the Messages API rejects on replay: every later turn 400s, and agent_messages is append-only for the BFL trail, so nothing could repair it. History is now patched on read by synthesizing is_error tool_results, leaving the stored trail untouched. check_and_increment_agent_quota is SECURITY DEFINER in public with a caller-chosen p_user_id, so any authenticated user could drain a colleague's minute/day budget and lock them out of every agent endpoint. A plain REVOKE would break the limiter (all three callers use the user's RLS client) and, as it fails open, silently remove the spend cap: the function now refuses to act for anyone but the caller, while service-role connections keep passing an explicit id. The single reject route re-read status and then wrote unguarded, so losing the race with commit's atomic pending -> committing claim stamped `rejected` over an operation that had already posted a verifikat, invisible to the committing-state recovery sweep. Guarded on status like bulk-reject already is; a lost race is now a 409. The sheet's Escape handler listened on window with no defaultPrevented or target check while the sheet is deliberately non-modal, so pressing Esc to dismiss the reject-reason Select inside an approval card, the command palette or any dialog unmounted the sheet and discarded the conversation, the streaming turn and the un-actioned proposal. It now yields to open overlays and to focus outside the sheet. Verified: 9526 unit tests pass, lint clean on touched files, guards pass, and the new pg-real test proves the quota guard against real Postgres (attacker raises 42501, victim counters stay at 0). The four unrelated pg-real failures on this machine reproduce identically with these changes stashed. Co-Authored-By: Claude Opus 5 * fix(assistant): close anon path on the quota RPC, order the ownership check ahead of writes Review follow-ups on the hotfix batch. The caller guard used auth.uid() alone, which is NULL for the `anon` role just as it is for backend roles, so an unauthenticated caller holding the public anon key (it ships in the browser bundle) could still pick any p_user_id and drain that user's quota. The guard now keys on the request role: anon and authenticated may only ever spend their own quota, backend roles keep passing an explicit id. The default PUBLIC execute grant is revoked as a second layer, with execute granted only to authenticated and service_role. Covered by a new pg test for the anon path. The ownership check ran after the onboarding.intake stamp, so a request that was about to be rejected could still write intake_completed_at. It now sits directly after the capability gate, ahead of every side effect and ahead of the company and profile reads, which also makes a rejected request cheaper. The tool-result repair matched ids anywhere in the history, but the API needs results in the message IMMEDIATELY after the tool_use. A result persisted after an intervening turn (two turns racing on one conversation) left a shape that still 400s. The repair is now positional, and orphaned or late-duplicate tool_results are dropped, since an unmatched tool_result is rejected just as an unanswered tool_use is. The Escape guard matched the Radix popper wrapper, which stays mounted when a popper is force-mounted; it now requires data-state="open" so a closed popper cannot block Escape for the rest of the session. Both new route errors are Swedish, per the user-facing error rule. Co-Authored-By: Claude Opus 5 --------- Co-authored-by: Claude Opus 5 --- app/api/agent/__tests__/invoke.test.ts | 224 ++++++++++++++++++ app/api/agent/invoke/route.ts | 40 +++- .../[id]/reject/__tests__/route.test.ts | 20 +- .../pending-operations/[id]/reject/route.ts | 22 +- components/agent/AgentSheet.tsx | 30 ++- .../__tests__/run-turn-dangling-tools.test.ts | 169 +++++++++++++ lib/agent/chat/run-turn.ts | 105 +++++++- ...726090000_agent_quota_rpc_caller_guard.sql | 100 ++++++++ tests/pg/agent-quota-caller-guard.pg.test.ts | 128 ++++++++++ 9 files changed, 830 insertions(+), 8 deletions(-) create mode 100644 app/api/agent/__tests__/invoke.test.ts create mode 100644 lib/agent/chat/__tests__/run-turn-dangling-tools.test.ts create mode 100644 supabase/migrations/20260726090000_agent_quota_rpc_caller_guard.sql create mode 100644 tests/pg/agent-quota-caller-guard.pg.test.ts diff --git a/app/api/agent/__tests__/invoke.test.ts b/app/api/agent/__tests__/invoke.test.ts new file mode 100644 index 00000000..a0bb1133 --- /dev/null +++ b/app/api/agent/__tests__/invoke.test.ts @@ -0,0 +1,224 @@ +/** + * Tests for POST /api/agent/invoke, focused on conversation ownership. + * + * RLS on agent_conversations/agent_messages is company-scoped, not user-scoped + * (migration 20260517204000), so the route itself has to prove that a resumed + * conversation_id belongs to the caller. Without that check, a member could + * post a colleague's conversation id and have their history loaded into the + * prompt (and their own turns appended to it). + */ +import { describe, it, expect, vi, beforeEach } from 'vitest' +import { NextResponse } from 'next/server' +import { createQueuedMockSupabase, createMockRequest, parseJsonResponse } from '@/tests/helpers' + +const { supabase, enqueue, reset } = createQueuedMockSupabase() + +const requireAuthMock = vi.fn() +vi.mock('@/lib/auth/require-auth', () => ({ + requireAuth: (...args: unknown[]) => requireAuthMock(...args), +})) + +vi.mock('@/lib/company/context', () => ({ + getActiveCompanyId: vi.fn().mockResolvedValue('company-1'), + requireCompanyId: vi.fn().mockResolvedValue('company-1'), +})) + +const checkRateMock = vi.fn() +vi.mock('@/lib/rate-limits/agent', () => ({ + checkAgentRateLimit: (...args: unknown[]) => checkRateMock(...args), + agentRateLimitResponseBody: () => ({ error: 'För många förfrågningar.' }), +})) + +vi.mock('@/lib/sandbox/guard', () => ({ + guardSandbox: vi.fn().mockResolvedValue(null), +})) + +vi.mock('@/lib/entitlements/has-capability', () => ({ + requireCapability: vi.fn().mockResolvedValue(null), +})) + +vi.mock('@/lib/entitlements/keys', () => ({ + CAPABILITY: { ai: 'ai' }, +})) + +vi.mock('@/lib/init', () => ({ + ensureInitialized: vi.fn(), +})) + +const getIntentMock = vi.fn() +vi.mock('@/lib/agent/intents/registry', () => ({ + getIntent: (...args: unknown[]) => getIntentMock(...args), +})) + +const runChatTurnMock = vi.fn() +vi.mock('@/lib/agent/chat/run-turn', () => ({ + runChatTurn: (...args: unknown[]) => runChatTurnMock(...args), + friendlyModelError: () => 'Något gick fel hos assistenten.', +})) + +import { POST } from '../invoke/route' + +const CONVERSATION_ID = '11111111-1111-4111-8111-111111111111' + +function body(overrides: Record = {}) { + return { + intent_id: 'general.help', + user_message: 'Hur gick juli?', + ...overrides, + } +} + +/** + * Queue the reads the route performs before it resolves the conversation. + * + * Ownership is validated ahead of every side effect and of the company/profile + * reads, so a rejected request costs exactly one membership read plus the + * conversation lookup. + */ +function enqueuePreamble() { + enqueue({ data: { role: 'owner' } }) // company_members +} + +/** The company + profile reads that only happen once a request is accepted. */ +function enqueueAcceptedTail() { + enqueue({ data: { name: 'Nordvik Bygg AB' } }) // companies + enqueue({ data: { full_name: 'Johan Nordvik' } }) // profiles +} + +beforeEach(() => { + vi.clearAllMocks() + reset() + requireAuthMock.mockResolvedValue({ user: { id: 'user-1' }, supabase, error: null }) + checkRateMock.mockResolvedValue({ ok: true }) + getIntentMock.mockReturnValue({ id: 'general.help', sheetTitle: 'Assistenten' }) + runChatTurnMock.mockResolvedValue(undefined) +}) + +describe('POST /api/agent/invoke', () => { + it('returns 401 when not authenticated', async () => { + requireAuthMock.mockResolvedValue({ + user: null, + supabase, + error: NextResponse.json({ error: 'Unauthorized' }, { status: 401 }), + }) + const res = await POST(createMockRequest('/api/agent/invoke', { method: 'POST', body: body() })) + const { status } = await parseJsonResponse(res) + expect(status).toBe(401) + }) + + it('returns 400 on an invalid body', async () => { + const res = await POST(createMockRequest('/api/agent/invoke', { method: 'POST', body: { intent_id: '' } })) + const { status } = await parseJsonResponse(res) + expect(status).toBe(400) + }) + + it('returns 400 for an unknown intent', async () => { + getIntentMock.mockReturnValue(undefined) + const res = await POST(createMockRequest('/api/agent/invoke', { method: 'POST', body: body() })) + const { status } = await parseJsonResponse<{ error: string }>(res) + expect(status).toBe(400) + }) + + it('returns 403 when the user is not a member of the company', async () => { + enqueue({ data: null }) // company_members: no row + const res = await POST(createMockRequest('/api/agent/invoke', { method: 'POST', body: body() })) + const { status } = await parseJsonResponse(res) + expect(status).toBe(403) + }) + + it('returns 404 when resuming a conversation owned by another user', async () => { + enqueuePreamble() + enqueue({ + data: { + id: CONVERSATION_ID, + user_id: 'user-2', // someone else in the same company + company_id: 'company-1', + intent_id: 'general.help', + }, + }) + + const res = await POST( + createMockRequest('/api/agent/invoke', { method: 'POST', body: body({ conversation_id: CONVERSATION_ID }) }), + ) + + const { status, body: json } = await parseJsonResponse<{ error: string }>(res) + expect(status).toBe(404) + expect(json.error).toBe('Konversationen hittades inte.') + expect(runChatTurnMock).not.toHaveBeenCalled() + }) + + it('returns 404 when resuming a conversation from another company', async () => { + enqueuePreamble() + enqueue({ + data: { + id: CONVERSATION_ID, + user_id: 'user-1', // same user... + company_id: 'company-2', // ...but a company other than the active one + intent_id: 'general.help', + }, + }) + + const res = await POST( + createMockRequest('/api/agent/invoke', { method: 'POST', body: body({ conversation_id: CONVERSATION_ID }) }), + ) + + const { status } = await parseJsonResponse(res) + expect(status).toBe(404) + expect(runChatTurnMock).not.toHaveBeenCalled() + }) + + it('returns 404 when the conversation does not exist', async () => { + enqueuePreamble() + enqueue({ data: null }) + + const res = await POST( + createMockRequest('/api/agent/invoke', { method: 'POST', body: body({ conversation_id: CONVERSATION_ID }) }), + ) + + const { status } = await parseJsonResponse(res) + expect(status).toBe(404) + expect(runChatTurnMock).not.toHaveBeenCalled() + }) + + it('returns 400 when the conversation belongs to a different intent', async () => { + enqueuePreamble() + enqueue({ + data: { + id: CONVERSATION_ID, + user_id: 'user-1', + company_id: 'company-1', + intent_id: 'vat.review', // tool loadout differs from general.help + }, + }) + + const res = await POST( + createMockRequest('/api/agent/invoke', { method: 'POST', body: body({ conversation_id: CONVERSATION_ID }) }), + ) + + const { status } = await parseJsonResponse(res) + expect(status).toBe(400) + expect(runChatTurnMock).not.toHaveBeenCalled() + }) + + it('runs the turn when the caller owns the conversation', async () => { + enqueuePreamble() + enqueue({ + data: { + id: CONVERSATION_ID, + user_id: 'user-1', + company_id: 'company-1', + intent_id: 'general.help', + }, + }) + enqueueAcceptedTail() + + const res = await POST( + createMockRequest('/api/agent/invoke', { method: 'POST', body: body({ conversation_id: CONVERSATION_ID }) }), + ) + + expect(res.status).toBe(200) + // The route streams NDJSON; draining it is enough to know the turn ran. + await res.text() + expect(runChatTurnMock).toHaveBeenCalledTimes(1) + }) +}) diff --git a/app/api/agent/invoke/route.ts b/app/api/agent/invoke/route.ts index a590aa04..3732ff10 100644 --- a/app/api/agent/invoke/route.ts +++ b/app/api/agent/invoke/route.ts @@ -117,6 +117,43 @@ export async function POST(request: Request) { const capBlocked = await requireCapability(supabase, companyId, CAPABILITY.ai) if (capBlocked) return capBlocked + // Resolve the conversation BEFORE any side effect below (the onboarding + // intake stamp): a request that is about to be rejected must not write. + let conversationId = body.conversation_id ?? null + if (conversationId) { + // A resumed conversation id comes straight from the client, so ownership + // has to be proven here. RLS on agent_conversations/agent_messages is + // COMPANY-scoped (migration 20260517204000), not user-scoped, so RLS alone + // would happily load a colleague's thread into the prompt and append this + // user's turns to it. The conversations list route filters on user_id for + // exactly this reason; the same rule applies to the turn itself. + // + // The company check matters too: a user who belongs to several companies + // must not resume a thread from company B while the turn runs with company + // A's ledger, tools and staged operations. + const { data: conv } = await supabase + .from('agent_conversations') + .select('id, user_id, company_id, intent_id') + .eq('id', conversationId) + .maybeSingle() + + if (!conv || conv.user_id !== user.id || conv.company_id !== companyId) { + // Same response for "doesn't exist" and "isn't yours": a 403 here would + // confirm that someone else's conversation id is real. + return NextResponse.json({ error: 'Konversationen hittades inte.' }, { status: 404 }) + } + + // The intent decides the tool loadout and the system prompt. Letting a + // resumed thread switch intent mid-conversation would swap the tool + // whitelist under history the model has already been shown. + if (conv.intent_id !== body.intent_id) { + return NextResponse.json( + { error: 'Konversationen hör till ett annat sammanhang.' }, + { status: 400 }, + ) + } + } + // onboarding.intake completion signal: once the user has actually // engaged (typed a real reply, not the auto-fired greeting prompt that // mounts the chat), stamp intake_completed_at on the profile so re-entry @@ -149,8 +186,7 @@ export async function POST(request: Request) { const companyName = company?.name ?? '' const firstName = profile?.full_name?.split(' ')[0] ?? null - // Resolve / create the conversation row. - let conversationId = body.conversation_id ?? null + // Create the conversation row when this is a fresh thread. if (!conversationId) { const { data: newConv, error: convErr } = await supabase .from('agent_conversations') diff --git a/app/api/pending-operations/[id]/reject/__tests__/route.test.ts b/app/api/pending-operations/[id]/reject/__tests__/route.test.ts index 6f15ae8d..a38b7f99 100644 --- a/app/api/pending-operations/[id]/reject/__tests__/route.test.ts +++ b/app/api/pending-operations/[id]/reject/__tests__/route.test.ts @@ -87,7 +87,7 @@ describe('POST /api/pending-operations/:id/reject', () => { it('rejects successfully', async () => { enqueueMany([ { data: { id: 'op-1', status: 'pending' } }, // fetch op - { data: null, error: null }, // update status + { data: [{ id: 'op-1' }], error: null }, // guarded update returns the row ]) const request = createMockRequest('/api/pending-operations/op-1/reject', { method: 'POST' }) @@ -97,4 +97,22 @@ describe('POST /api/pending-operations/:id/reject', () => { expect(status).toBe(200) expect(body.data.status).toBe('rejected') }) + + it('returns 409 when commit claims the row between the status read and the write', async () => { + // The status re-check above the write is a read, and commit claims rows + // atomically (pending -> committing). Losing that race must not stamp + // `rejected` over an operation that has since posted a verifikat, so the + // UPDATE is guarded on status and matches zero rows here. + enqueueMany([ + { data: { id: 'op-1', status: 'pending' } }, // fetch op: still pending + { data: [], error: null }, // guarded update matches nothing + ]) + + const request = createMockRequest('/api/pending-operations/op-1/reject', { method: 'POST' }) + const response = await POST(request, routeParams) + const { status, body } = await parseJsonResponse<{ error: string }>(response) + + expect(status).toBe(409) + expect(body.error).toContain('hann behandlas') + }) }) diff --git a/app/api/pending-operations/[id]/reject/route.ts b/app/api/pending-operations/[id]/reject/route.ts index 10b2043c..16965814 100644 --- a/app/api/pending-operations/[id]/reject/route.ts +++ b/app/api/pending-operations/[id]/reject/route.ts @@ -83,7 +83,13 @@ export const POST = withRouteContext<{ params: Promise<{ id: string }> }>( ) } - const { error: updateError } = await supabase + // The status check above is a read, and commit claims rows atomically + // (pending -> committing). Without a status guard on the write, a reject + // that loses that race stamps `rejected` over an operation that has since + // posted a verifikat: a booked entry whose row says rejected, which the + // committing-state recovery sweep will never find. Guard the UPDATE the + // same way bulk-reject does and report a lost race as 409. + const { data: updated, error: updateError } = await supabase .from('pending_operations') .update({ status: 'rejected', @@ -92,11 +98,25 @@ export const POST = withRouteContext<{ params: Promise<{ id: string }> }>( ...(rejectionReason ? { rejection_reason: rejectionReason } : {}), }) .eq('id', id) + .eq('company_id', companyId) + .eq('status', 'pending') + .select('id') if (updateError) { return NextResponse.json({ error: getUserErrorMessage(updateError) }, { status: 500 }) } + if (!updated || updated.length === 0) { + return NextResponse.json( + { + error: + 'Åtgärden hann behandlas av någon annan och kunde inte avvisas. ' + + 'Ladda om sidan och kontrollera resultatet.', + }, + { status: 409 }, + ) + } + return NextResponse.json({ data: { id, status: 'rejected' } }) }, { requireWrite: true }, diff --git a/components/agent/AgentSheet.tsx b/components/agent/AgentSheet.tsx index cf4d5081..6259c955 100644 --- a/components/agent/AgentSheet.tsx +++ b/components/agent/AgentSheet.tsx @@ -1,6 +1,6 @@ 'use client' -import { useEffect, useState } from 'react' +import { useEffect, useRef, useState } from 'react' import { X, Expand, Shrink, PanelRightClose, Eraser, History, ChevronLeft, Loader2 } from 'lucide-react' import AgentChat, { normalizeStoredMessages, type ChatMessage } from './AgentChat' import AgentAvatar from './AgentAvatar' @@ -68,12 +68,39 @@ export default function AgentSheet({ const sheetTitle = intentToTitle(intentId, agentName) const displayTitle = loaded ? (loaded.title ?? intentToTitle(loaded.intentId, agentName)) : sheetTitle const activeConversationId = loaded?.id ?? conversationId + const sheetRef = useRef(null) // Esc: back out of the session list first, otherwise close. Never while // collapsed (the sheet is hidden off-screen, so Esc belongs elsewhere). + // + // The sheet is deliberately non-modal, so this listener sits on window while + // the rest of the page stays interactive: it must therefore only claim the + // key when nothing nearer the user wants it. Closing the sheet discards the + // whole in-memory conversation, so an Esc meant for a dropdown inside an + // approval card, the command palette, or any dialog used to destroy the + // session outright. Three guards, cheapest first: + // - defaultPrevented: a Radix popover/dialog that handled Esc marks it. + // - an open overlay anywhere on the page (Radix marks these on the body + // and on the overlay elements themselves) means the key isn't ours. + // - focus sitting outside the sheet means the user is working elsewhere. useEffect(() => { const onKey = (e: KeyboardEvent) => { if (collapsed || e.key !== 'Escape') return + if (e.defaultPrevented) return + + if (typeof document !== 'undefined') { + // Match on data-state="open", not on the popper wrapper itself: a + // force-mounted popper stays in the DOM while closed, and keying off + // the wrapper alone would then block Escape for the rest of the session. + const overlayOpen = document.querySelector( + '[data-radix-popper-content-wrapper] [data-state="open"], [role="dialog"][data-state="open"], [role="alertdialog"][data-state="open"], [role="listbox"][data-state="open"], [data-radix-menu-content][data-state="open"], [data-radix-select-content][data-state="open"]', + ) + if (overlayOpen) return + + const active = document.activeElement + if (active && sheetRef.current && !sheetRef.current.contains(active)) return + } + if (view === 'list') setView('chat') else onClose() } @@ -135,6 +162,7 @@ export default function AgentSheet({ return (
[{ type: 'text', text }] + +describe('repairDanglingToolUse', () => { + it('leaves a well-formed conversation untouched', () => { + const messages = [ + { role: 'user' as const, content: textBlock('Boka om Circle K') }, + { + role: 'assistant' as const, + content: [ + { type: 'tool_use', id: 'tu_1', name: 'gnubok_query_journal', input: {} }, + ], + }, + { + role: 'user' as const, + content: [{ type: 'tool_result', tool_use_id: 'tu_1', content: 'A-207' }], + }, + { role: 'assistant' as const, content: textBlock('Klart att godkänna.') }, + ] + + expect(repairDanglingToolUse(messages)).toEqual(messages) + }) + + it('synthesizes an error tool_result for an unanswered tool_use', () => { + const messages = [ + { role: 'user' as const, content: textBlock('Gör klart juli') }, + { + role: 'assistant' as const, + content: [{ type: 'tool_use', id: 'tu_dead', name: 'gnubok_list_transactions', input: {} }], + }, + ] + + const repaired = repairDanglingToolUse(messages) + + expect(repaired).toHaveLength(3) + expect(repaired[2]).toEqual({ + role: 'user', + content: [ + { + type: 'tool_result', + tool_use_id: 'tu_dead', + content: expect.stringContaining('Avbröts'), + is_error: true, + }, + ], + }) + }) + + it('repairs only the unanswered call in a mixed batch', () => { + const messages = [ + { + role: 'assistant' as const, + content: [ + { type: 'tool_use', id: 'tu_ok', name: 'a', input: {} }, + { type: 'tool_use', id: 'tu_dead', name: 'b', input: {} }, + ], + }, + { + role: 'user' as const, + content: [{ type: 'tool_result', tool_use_id: 'tu_ok', content: 'fine' }], + }, + ] + + const repaired = repairDanglingToolUse(messages) + const synthesized = repaired.filter( + (m) => + Array.isArray(m.content) && + m.content.some((b: { is_error?: boolean }) => b.is_error === true), + ) + + expect(synthesized).toHaveLength(1) + expect(synthesized[0]!.content).toHaveLength(1) + expect(synthesized[0]!.content[0].tool_use_id).toBe('tu_dead') + }) + + it('inserts the stub immediately after the assistant turn that opened it', () => { + const messages = [ + { + role: 'assistant' as const, + content: [{ type: 'tool_use', id: 'tu_dead', name: 'a', input: {} }], + }, + { role: 'user' as const, content: textBlock('är du kvar?') }, + ] + + const repaired = repairDanglingToolUse(messages) + + // The stub must sit between the tool_use and the next user turn, otherwise + // the API still sees an unanswered tool_use followed by a user message. + expect(repaired.map((m) => m.role)).toEqual(['assistant', 'user', 'user']) + expect(repaired[1]!.content[0].tool_use_id).toBe('tu_dead') + expect(repaired[2]!.content[0].text).toBe('är du kvar?') + }) + + it('repairs a result that arrived after an intervening turn', () => { + // Two turns racing on one conversation can persist a user message between + // the tool_use and its results. The id IS answered somewhere, but not in + // the message immediately following the tool_use, which the API still + // rejects: the stub has to go in the gap and the late copy has to go. + const messages = [ + { + role: 'assistant' as const, + content: [{ type: 'tool_use', id: 'tu_late', name: 'a', input: {} }], + }, + { role: 'user' as const, content: textBlock('hallå?') }, + { + role: 'user' as const, + content: [{ type: 'tool_result', tool_use_id: 'tu_late', content: 'sent too late' }], + }, + ] + + const repaired = repairDanglingToolUse(messages) + + // Stub sits directly after the tool_use... + expect(repaired[1]!.content[0]).toMatchObject({ + type: 'tool_result', + tool_use_id: 'tu_late', + is_error: true, + }) + // ...the intervening user text survives... + expect(repaired[2]!.content[0].text).toBe('hallå?') + // ...and the orphaned late result is gone, along with its emptied message. + const lateResults = repaired + .flatMap((m) => (Array.isArray(m.content) ? m.content : [])) + .filter((b: { content?: string }) => b?.content === 'sent too late') + expect(lateResults).toHaveLength(0) + }) + + it('drops a tool_result whose tool_use never preceded it', () => { + const messages = [ + { role: 'user' as const, content: textBlock('hej') }, + { + role: 'user' as const, + content: [{ type: 'tool_result', tool_use_id: 'tu_ghost', content: 'orphan' }], + }, + ] + + const repaired = repairDanglingToolUse(messages) + + // An unmatched tool_result is rejected by the API just as an unanswered + // tool_use is, so the emptied message is dropped entirely. + expect(repaired).toHaveLength(1) + expect(repaired[0]!.content[0].text).toBe('hej') + }) + + it('tolerates string and non-array content without throwing', () => { + const messages = [ + { role: 'user' as const, content: 'plain string content' as unknown as [] }, + { role: 'assistant' as const, content: null as unknown as [] }, + ] + + expect(() => repairDanglingToolUse(messages)).not.toThrow() + expect(repairDanglingToolUse(messages)).toHaveLength(2) + }) +}) diff --git a/lib/agent/chat/run-turn.ts b/lib/agent/chat/run-turn.ts index 8eada84c..ded60b39 100644 --- a/lib/agent/chat/run-turn.ts +++ b/lib/agent/chat/run-turn.ts @@ -605,12 +605,111 @@ async function loadConversationMessages( .order('created_at', { ascending: true }) // role='tool' messages were written as user messages on the Anthropic side. - return (data ?? []).map((m: { role: string; content: ContentBlock }) => { + const messages = (data ?? []).map((m: { role: string; content: ContentBlock }) => { if (m.role === 'assistant') { - return { role: 'assistant', content: m.content as ContentBlock } + return { role: 'assistant' as const, content: m.content as ContentBlock } } - return { role: 'user', content: m.content as ContentBlock } + return { role: 'user' as const, content: m.content as ContentBlock } }) + + return repairDanglingToolUse(messages) +} + +/** + * Synthesize `tool_result` blocks for any `tool_use` the stored history never + * answered. + * + * The assistant message carrying `tool_use` blocks is persisted before the + * tools run, and their results only after the whole batch finishes. If the + * process dies in between (client disconnect terminating the function, a + * deploy, a tool that outlives the request), the stored conversation ends on an + * unanswered `tool_use`. The Messages API rejects that shape on replay, so + * every later turn 400s: and because agent_messages is append-only by design + * (no UPDATE/DELETE policies, BFL audit trail), nothing can repair the row. + * The conversation is bricked forever. + * + * Repairing on read keeps the stored trail untouched and the thread usable. + * The synthesized result is flagged as an error so the model treats it as a + * failed call rather than silently inventing an outcome from it. + */ +export function repairDanglingToolUse( + messages: { role: 'user' | 'assistant'; content: ContentBlock }[], +): { role: 'user' | 'assistant'; content: ContentBlock }[] { + const toolResultIds = (content: ContentBlock): Set => { + const ids = new Set() + if (!Array.isArray(content)) return ids + for (const block of content) { + if (block?.type === 'tool_result' && typeof block.tool_use_id === 'string') { + ids.add(block.tool_use_id) + } + } + return ids + } + + // The API requires results in the message IMMEDIATELY following the tool_use, + // so position matters, not just presence: a result that landed after an + // intervening turn (two turns racing on one conversation) is still an invalid + // shape. Walk pairwise, and treat only same-position results as answers. + const out: { role: 'user' | 'assistant'; content: ContentBlock }[] = [] + const satisfied = new Set() + + for (let i = 0; i < messages.length; i++) { + const m = messages[i]! + out.push(m) + if (m.role !== 'assistant' || !Array.isArray(m.content)) continue + + const pending = m.content + .filter((block: ContentBlock) => block?.type === 'tool_use' && typeof block.id === 'string') + .map((block: ContentBlock) => block.id as string) + if (pending.length === 0) continue + + const answeredHere = toolResultIds(messages[i + 1]?.content) + const missing = pending.filter((id) => !answeredHere.has(id)) + for (const id of pending) { + if (answeredHere.has(id)) satisfied.add(id) + } + + if (missing.length > 0) { + for (const id of missing) satisfied.add(id) + out.push({ + role: 'user', + content: missing.map((id) => ({ + type: 'tool_result' as const, + tool_use_id: id, + content: 'Avbröts innan verktyget hann svara. Kör om det om du behöver resultatet.', + is_error: true, + })) as ContentBlock, + }) + } + } + + // Drop any tool_result that is now orphaned: either a late duplicate of one + // we just stubbed, or a result whose tool_use never immediately preceded it. + // An unmatched tool_result is rejected by the API just as an unanswered + // tool_use is, so leaving it in would defeat the repair. + return out + .map((m, idx) => { + if (!Array.isArray(m.content)) return m + const prev = out[idx - 1] + const openedByPrev = + prev?.role === 'assistant' && Array.isArray(prev.content) + ? new Set( + prev.content + .filter( + (b: ContentBlock) => b?.type === 'tool_use' && typeof b.id === 'string', + ) + .map((b: ContentBlock) => b.id as string), + ) + : new Set() + + const kept = m.content.filter((block: ContentBlock) => { + if (block?.type !== 'tool_result') return true + return openedByPrev.has(block.tool_use_id) + }) + if (kept.length === m.content.length) return m + return { ...m, content: kept as ContentBlock } + }) + .filter((m) => !Array.isArray(m.content) || m.content.length > 0) } async function persistMessage( diff --git a/supabase/migrations/20260726090000_agent_quota_rpc_caller_guard.sql b/supabase/migrations/20260726090000_agent_quota_rpc_caller_guard.sql new file mode 100644 index 00000000..32035d24 --- /dev/null +++ b/supabase/migrations/20260726090000_agent_quota_rpc_caller_guard.sql @@ -0,0 +1,100 @@ +-- Harden check_and_increment_agent_quota against caller-chosen p_user_id. +-- +-- The function is SECURITY DEFINER and lives in `public`, so PostgREST exposes +-- it at /rest/v1/rpc/check_and_increment_agent_quota with the default +-- EXECUTE-to-PUBLIC grant. p_user_id is a plain argument, and user ids are +-- discoverable through company_members, so any caller could burn down a +-- targeted user's minute/day budget and lock them out of every agent endpoint +-- for the rest of the day. +-- +-- A plain REVOKE FROM PUBLIC alone is not the fix: all three callers (agent +-- invoke, onboarding stream, composer) call this with the user's own RLS +-- client, i.e. as `authenticated`, so revoking everything would break the +-- limiter and, because it fails open on error, silently remove the spend cap it +-- exists to enforce. +-- +-- So: two layers. +-- 1) Grants. Only `authenticated` and `service_role` may execute it at all. +-- `anon` is explicitly excluded: the anon key ships in the browser bundle, +-- so an unauthenticated caller could otherwise reach this RPC directly. +-- 2) A caller guard in the body. An end-user role may only ever spend its own +-- quota. Roles that are not PostgREST end users (service_role, and the +-- migration/superuser role used by cron, jobs and tests) keep passing an +-- explicit user id, which they need in order to act on someone's behalf. +-- +-- auth.uid() alone is NOT sufficient as the guard: it is NULL for `anon` as +-- well as for backend roles, so an unauthenticated caller would pass it. +-- +-- Function body is otherwise byte-identical to 20260526140000. + +CREATE OR REPLACE FUNCTION public.check_and_increment_agent_quota( + p_user_id uuid, + p_minute_max integer, + p_day_max integer +) RETURNS jsonb +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = public +AS $$ +DECLARE + v_minute_key text := to_char(now() AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI'); + v_day_key text := to_char(now() AT TIME ZONE 'Europe/Stockholm', 'YYYY-MM-DD'); + v_minute_count integer; + v_day_count integer; + v_role text := coalesce( + nullif(current_setting('request.jwt.claim.role', true), ''), + nullif(current_setting('request.jwt.claims', true), '')::jsonb ->> 'role', + '' + ); +BEGIN + -- Caller guard: anything arriving through PostgREST as an end user (anon or + -- authenticated) may only spend the quota of the user it is authenticated as. + IF v_role IN ('anon', 'authenticated') THEN + IF auth.uid() IS NULL OR p_user_id IS DISTINCT FROM auth.uid() THEN + RAISE EXCEPTION 'check_and_increment_agent_quota: p_user_id must be the calling user' + USING ERRCODE = '42501'; + END IF; + END IF; + + -- 1) Minute window — burst guard. + INSERT INTO public.agent_rate_counters (user_id, window_kind, window_key, count) + VALUES (p_user_id, 'minute', v_minute_key, 1) + ON CONFLICT (user_id, window_kind, window_key) + DO UPDATE SET count = agent_rate_counters.count + 1, updated_at = now() + RETURNING count INTO v_minute_count; + + IF v_minute_count > p_minute_max THEN + UPDATE public.agent_rate_counters SET count = count - 1 + WHERE user_id = p_user_id AND window_kind = 'minute' AND window_key = v_minute_key; + RETURN jsonb_build_object('ok', false, 'scope', 'minute', 'retry_after_sec', 60); + END IF; + + -- 2) Day window — slow-drip backstop (only checked once minute passes). + INSERT INTO public.agent_rate_counters (user_id, window_kind, window_key, count) + VALUES (p_user_id, 'day', v_day_key, 1) + ON CONFLICT (user_id, window_kind, window_key) + DO UPDATE SET count = agent_rate_counters.count + 1, updated_at = now() + RETURNING count INTO v_day_count; + + IF v_day_count > p_day_max THEN + -- Roll both counters back: the request didn't go through. + UPDATE public.agent_rate_counters SET count = count - 1 + WHERE user_id = p_user_id AND window_kind = 'day' AND window_key = v_day_key; + UPDATE public.agent_rate_counters SET count = count - 1 + WHERE user_id = p_user_id AND window_kind = 'minute' AND window_key = v_minute_key; + RETURN jsonb_build_object('ok', false, 'scope', 'day', 'retry_after_sec', 3600); + END IF; + + RETURN jsonb_build_object('ok', true); +END; +$$; + +-- Close the default PUBLIC grant and hand execute rights only to the roles that +-- legitimately call this: the three agent endpoints (as `authenticated`) and +-- backend/service contexts. +REVOKE ALL ON FUNCTION public.check_and_increment_agent_quota(uuid, integer, integer) FROM PUBLIC; +REVOKE ALL ON FUNCTION public.check_and_increment_agent_quota(uuid, integer, integer) FROM anon; +GRANT EXECUTE ON FUNCTION public.check_and_increment_agent_quota(uuid, integer, integer) TO authenticated; +GRANT EXECUTE ON FUNCTION public.check_and_increment_agent_quota(uuid, integer, integer) TO service_role; + +NOTIFY pgrst, 'reload schema'; diff --git a/tests/pg/agent-quota-caller-guard.pg.test.ts b/tests/pg/agent-quota-caller-guard.pg.test.ts new file mode 100644 index 00000000..d7825831 --- /dev/null +++ b/tests/pg/agent-quota-caller-guard.pg.test.ts @@ -0,0 +1,128 @@ +import { describe, expect, it } from 'vitest' +import type { PoolClient } from 'pg' +import { getClient, getPool } from './setup' +import { insertAuthUser } from './fixtures' + +/** + * Caller guard on check_and_increment_agent_quota (migration 20260726090000). + * + * The function is SECURITY DEFINER in `public`, so PostgREST exposes it to any + * authenticated user, and p_user_id is a plain argument. Without a guard, one + * user could spend another user's minute/day budget and lock them out of every + * agent endpoint. This test pins the guard against a real Postgres, because a + * mocked Supabase client cannot see a PL/pgSQL condition at all. + * + * Locks in: + * - An authenticated caller may spend their OWN quota. + * - An authenticated caller spending someone else's quota raises 42501 and + * leaves the victim's counters untouched. + * - A service-role/superuser connection (auth.uid() IS NULL) may still pass + * an explicit user id: cron, tests and backend jobs rely on that. + */ + +async function asRole( + role: 'authenticated' | 'anon', + userId: string | null, + fn: (client: PoolClient) => Promise, +): Promise { + const client = await getClient() + try { + await client.query('BEGIN') + await client.query(`SELECT set_config('request.jwt.claims', $1, true)`, [ + JSON.stringify(userId ? { sub: userId, role } : { role }), + ]) + await client.query(`SELECT set_config('request.jwt.claim.sub', $1, true)`, [userId ?? '']) + await client.query(`SELECT set_config('request.jwt.claim.role', $1, true)`, [role]) + await client.query(`SET LOCAL ROLE ${role}`) + const result = await fn(client) + await client.query('COMMIT') + return result + } catch (error) { + await client.query('ROLLBACK').catch(() => {}) + throw error + } finally { + client.release() + } +} + +function asAuthenticatedUser( + userId: string, + fn: (client: PoolClient) => Promise, +): Promise { + return asRole('authenticated', userId, fn) +} + +async function counterFor(userId: string): Promise { + const res = await getPool().query<{ count: number }>( + `SELECT COALESCE(SUM(count), 0)::int AS count + FROM public.agent_rate_counters + WHERE user_id = $1`, + [userId], + ) + return res.rows[0]?.count ?? 0 +} + +describe('check_and_increment_agent_quota caller guard.pg', () => { + it('lets an authenticated user spend their own quota', async () => { + const userId = await insertAuthUser() + + const result = await asAuthenticatedUser(userId, async (client) => { + const res = await client.query<{ result: { ok: boolean } }>( + `SELECT public.check_and_increment_agent_quota($1::uuid, 30, 1000) AS result`, + [userId], + ) + return res.rows[0]!.result + }) + + expect(result.ok).toBe(true) + expect(await counterFor(userId)).toBeGreaterThan(0) + }) + + it('refuses to spend another user quota and leaves their counters untouched', async () => { + const attacker = await insertAuthUser() + const victim = await insertAuthUser() + + await expect( + asAuthenticatedUser(attacker, async (client) => { + await client.query( + `SELECT public.check_and_increment_agent_quota($1::uuid, 30, 1000) AS result`, + [victim], + ) + }), + ).rejects.toMatchObject({ code: '42501' }) + + expect(await counterFor(victim)).toBe(0) + }) + + it('refuses an unauthenticated anon caller entirely', async () => { + // The anon key ships in the browser bundle, so this RPC is reachable + // without a session. auth.uid() is NULL for anon exactly as it is for + // backend roles, so a uid-only guard would have let this through. + const victim = await insertAuthUser() + + await expect( + asRole('anon', null, async (client) => { + await client.query( + `SELECT public.check_and_increment_agent_quota($1::uuid, 30, 1000) AS result`, + [victim], + ) + }), + ).rejects.toMatchObject({ code: expect.stringMatching(/^42501$/) }) + + expect(await counterFor(victim)).toBe(0) + }) + + it('still allows a service-role connection to pass an explicit user id', async () => { + const userId = await insertAuthUser() + + // The shared pool connects as the migration/superuser role, so auth.uid() + // is NULL here: the same shape cron jobs and backend scripts run under. + const res = await getPool().query<{ result: { ok: boolean } }>( + `SELECT public.check_and_increment_agent_quota($1::uuid, 30, 1000) AS result`, + [userId], + ) + + expect(res.rows[0]!.result.ok).toBe(true) + expect(await counterFor(userId)).toBeGreaterThan(0) + }) +})