From 60a0942cbdd29bcd8bc826f933ccc1d5db051d83 Mon Sep 17 00:00:00 2001 From: Tanner Linsley Date: Mon, 5 Oct 2026 01:22:51 -0600 Subject: [PATCH 1/2] perf(chat): reserve usage and authorize admission in fewer round trips --- docs/tanchat-send-latency.md | 10 ++ .../0041_tanchat_usage_reservation.sql | 86 +++++++++++ drizzle/migrations/meta/_journal.json | 7 + harness-tests/core/run-usage-admin.test.ts | 145 ------------------ .../pending-runtime/run-usage-admin.test.ts | 59 +++++++ scripts/verify-tanchat-schema.ts | 5 + src/chat/conversation-identity.server.ts | 36 ++++- src/chat/server/conversation.ts | 42 +++-- src/chat/server/run-usage.ts | 125 +++------------ 9 files changed, 252 insertions(+), 263 deletions(-) create mode 100644 drizzle/migrations/0041_tanchat_usage_reservation.sql delete mode 100644 harness-tests/core/run-usage-admin.test.ts diff --git a/docs/tanchat-send-latency.md b/docs/tanchat-send-latency.md index f06b36c47..9f9372251 100644 --- a/docs/tanchat-send-latency.md +++ b/docs/tanchat-send-latency.md @@ -51,3 +51,13 @@ Live Worker tracing on the next short reply measured 1,398 ms of wall time and 5 The HTTP send path now overlaps workflow ownership, lifecycle metadata, and run context reads after resolving the authenticated conversation identity. All three reads settle before the request database context is released. Authorization, lifecycle rejection, workflow ownership rejection, and post-preparation checks remain in place. Structured chat_send_timing, chat_admission_timing, and chat_run_timing logs record cumulative milliseconds within each operation. They separate HTTP authentication and metadata, admission guards and preparation, model credentials and selection, usage reservation, enrichment, context preparation, first visible text, its durable save, and response completion. They do not log message text, credentials, or account identity. Admission and run logs use the message receipt ID for correlation. The first-text save is not a measurement of browser paint or stream delivery. + +## Measured database round trips, October 5 + +The first instrumented live run measured 301 ms of HTTP authentication and metadata, 694 ms inside admission, and 1,319 ms from run start to model context. The two admission guards accounted for 555 ms. Usage reservation took 916 ms after credentials, and enrichment took another 171 ms. First text arrived 2,895 ms after model context, consistent with the recorded 3.01-second model call. The browser observed the reply after 5.935 seconds, including browser control and observation overhead. + +Reservation now executes the existing protocol inside one PostgreSQL function call. It keeps the same receipt and day advisory lock keys, fresh receipt and role checks, allowance ordering, and atomic receipts, counters, and funded reservations. The function is SECURITY INVOKER and VOLATILE, so queries inside it obtain fresh snapshots after waiting for locks. A real database regression revokes admin access while reservation waits on the day lock and verifies rejection without a receipt. + +Ordinary admission guards now read membership, ownership, bot lifecycle, and thread lifecycle together. Both checks before and after preparation remain fresh. Retry readiness still receives a separate lifecycle check after its asynchronous readiness wait. + +[PostgreSQL function volatility](https://www.postgresql.org/docs/current/xfunc-volatility.html) explains the fresh snapshot behavior. These changes reduce application overhead without changing the model or its response quality. They do not promise near-zero provider latency. diff --git a/drizzle/migrations/0041_tanchat_usage_reservation.sql b/drizzle/migrations/0041_tanchat_usage_reservation.sql new file mode 100644 index 000000000..27a01380c --- /dev/null +++ b/drizzle/migrations/0041_tanchat_usage_reservation.sql @@ -0,0 +1,86 @@ +-- Run the existing allowance protocol next to its data, in one network call. +-- VOLATILE gives each query a fresh snapshot after the advisory locks resolve. +CREATE FUNCTION public.reserve_chat_run_usage(request jsonb, receipt_lock_key text) +RETURNS TABLE(outcome text, receipt_day text) +LANGUAGE plpgsql VOLATILE SECURITY INVOKER SET search_path = public +AS $$ +DECLARE + conversation_id_value text := request->>'conversationId'; + run_id_value text := request->>'runId'; + workspace_id_value text := request->>'workspaceId'; + user_id_value uuid := (request->>'userId')::uuid; + bot_id_value text := request->>'botId'; + scheduled_value boolean := (request->>'scheduled')::boolean; + day_value text := request->>'day'; + unlimited_value boolean := (request->>'unlimited')::boolean; + spend jsonb := request->'spend'; + reservation bigint := (spend->>'reservationMicros')::bigint; + receipt chat_run_usage_receipts%ROWTYPE; + user_turns bigint; + global_turns bigint; + scheduled_turns bigint; + user_spend bigint; + global_spend bigint; + included_admin boolean := false; +BEGIN + PERFORM pg_advisory_xact_lock(hashtextextended(receipt_lock_key, 0)); + PERFORM pg_advisory_xact_lock(hashtextextended('chat-run-usage:' || day_value, 0)); + SELECT * INTO receipt FROM chat_run_usage_receipts + WHERE conversation_id=conversation_id_value AND run_id=run_id_value; + IF FOUND THEN + IF receipt.workspace_id <> workspace_id_value OR receipt.user_id <> user_id_value + OR receipt.bot_id <> bot_id_value OR receipt.scheduled <> scheduled_value THEN + RETURN QUERY SELECT 'conflict'::text, receipt.day; + ELSE + RETURN QUERY SELECT 'duplicate'::text, receipt.day; + END IF; + RETURN; + END IF; + SELECT + COALESCE((SELECT turns FROM chat_daily_usage WHERE user_id=user_id_value::text AND day=day_value),0), + COALESCE((SELECT turns FROM chat_daily_usage WHERE user_id='__global' AND day=day_value),0), + COALESCE((SELECT turns FROM chat_scheduled_daily_usage WHERE user_id=user_id_value AND day=day_value),0), + COALESCE((SELECT SUM(billed_micros) FROM chat_funded_spend WHERE user_id=user_id_value AND day=day_value),0), + COALESCE((SELECT SUM(billed_micros) FROM chat_funded_spend WHERE day=day_value),0) + INTO user_turns,global_turns,scheduled_turns,user_spend,global_spend; + IF NOT (request->>'lift')::boolean AND user_turns >= (request->>'dailyLimit')::integer THEN + RETURN QUERY SELECT 'user'::text, day_value; RETURN; + END IF; + IF NOT (request->>'lift')::boolean AND global_turns >= 300 THEN + RETURN QUERY SELECT 'global'::text, day_value; RETURN; + END IF; + IF NOT unlimited_value AND scheduled_value AND scheduled_turns >= (request->>'scheduledLimit')::integer THEN + RETURN QUERY SELECT 'scheduled'::text, day_value; RETURN; + END IF; + IF reservation IS NOT NULL AND NOT unlimited_value THEN + SELECT EXISTS ( + SELECT 1 FROM users u + LEFT JOIN role_assignments a ON a.user_id=u.id + LEFT JOIN roles r ON r.id=a.role_id + WHERE u.id=user_id_value AND + ('admin'=ANY(u.capabilities::text::text[]) OR 'admin'=ANY(r.capabilities::text::text[])) + ) INTO included_admin; + IF NOT included_admin AND user_spend + reservation > (spend->>'userCapMicros')::bigint THEN + RETURN QUERY SELECT 'user-spend'::text, day_value; RETURN; + END IF; + IF NOT included_admin AND global_spend + reservation > (spend->>'globalCapMicros')::bigint THEN + RETURN QUERY SELECT 'global-spend'::text, day_value; RETURN; + END IF; + END IF; + INSERT INTO chat_run_usage_receipts(conversation_id,run_id,workspace_id,user_id,bot_id,scheduled,day,created_at,attempt_id) + VALUES(conversation_id_value,run_id_value,workspace_id_value,user_id_value,bot_id_value,scheduled_value,day_value, + (request->>'now')::bigint,(request->>'attemptId')::uuid); + INSERT INTO chat_daily_usage(user_id,day,turns) + VALUES(user_id_value::text,day_value,1),('__global',day_value,1) + ON CONFLICT(user_id,day) DO UPDATE SET turns=chat_daily_usage.turns+1; + IF reservation IS NOT NULL THEN + INSERT INTO chat_funded_spend(conversation_id,run_id,user_id,day,reserved_micros,billed_micros) + VALUES(conversation_id_value,run_id_value,user_id_value,day_value,reservation,reservation); + END IF; + IF scheduled_value THEN + INSERT INTO chat_scheduled_daily_usage(user_id,day,turns) VALUES(user_id_value,day_value,1) + ON CONFLICT(user_id,day) DO UPDATE SET turns=chat_scheduled_daily_usage.turns+1; + END IF; + RETURN QUERY SELECT 'reserved'::text, day_value; +END; +$$; diff --git a/drizzle/migrations/meta/_journal.json b/drizzle/migrations/meta/_journal.json index b1b5c8f9d..90638cf44 100644 --- a/drizzle/migrations/meta/_journal.json +++ b/drizzle/migrations/meta/_journal.json @@ -288,6 +288,13 @@ "when": 1790963800008, "tag": "0040_tanchat_waitlist", "breakpoints": true + }, + { + "idx": 41, + "version": "7", + "when": 1791183973089, + "tag": "0041_tanchat_usage_reservation", + "breakpoints": true } ] } diff --git a/harness-tests/core/run-usage-admin.test.ts b/harness-tests/core/run-usage-admin.test.ts deleted file mode 100644 index d230eb8bb..000000000 --- a/harness-tests/core/run-usage-admin.test.ts +++ /dev/null @@ -1,145 +0,0 @@ -import { beforeEach, expect, it, vi } from 'vitest' -import type { SQL } from 'drizzle-orm' -import { PgDialect } from 'drizzle-orm/pg-core' - -const mocks = vi.hoisted(() => ({ - execute: vi.fn(), - transaction: vi.fn(), - capabilities: vi.fn(), -})) -vi.mock('~/db/client', () => ({ db: mocks })) -vi.mock('~/auth/repositories.server', () => ({ - DrizzleCapabilitiesRepository: class { - getEffectiveCapabilities = mocks.capabilities - }, -})) -import { - reserveRunUsage, - settleFundedSpend, - type ReserveRunUsageInput, -} from '../../src/chat/server/run-usage' - -const dialect = new PgDialect() -const queries: { sql: string; params: unknown[] }[] = [] -const usage = { - user_turns: 1, - global_turns: 1, - scheduled_turns: 0, - user_spend: 2_000_000, - global_spend: 20_000_000, - user_is_admin: false, -} -const input: ReserveRunUsageInput = { - identity: { - workspaceId: 'workspace', - botId: 'bot', - conversationId: 'room', - userId: '877728d5-9673-4b29-b330-6b110d447c53', - }, - runId: 'run', - policy: { dailyTurns: 20 }, - localDevelopment: false, - scheduled: false, - fundedSpend: { - userCapMicros: 2_000_000, - globalCapMicros: 20_000_000, - reservationMicros: 250_000, - }, -} - -beforeEach(() => { - queries.length = 0 - Object.assign(usage, { - user_turns: 1, - global_turns: 1, - scheduled_turns: 0, - user_spend: 2_000_000, - global_spend: 20_000_000, - user_is_admin: false, - }) - mocks.capabilities.mockReset() - mocks.capabilities.mockImplementation(async () => - usage.user_is_admin ? ['admin'] : [], - ) - mocks.execute.mockReset() - mocks.execute.mockImplementation(async (statement: SQL) => { - const query = dialect.sqlToQuery(statement) - queries.push(query) - if (query.sql.includes('AS user_turns')) return [usage] - if (query.sql.includes('SELECT day FROM chat_funded_spend')) - return [{ day: '2026-10-04' }] - return [] - }) - mocks.transaction.mockImplementation((work) => - work({ execute: mocks.execute }), - ) -}) - -it('admits a stored admin above both spending caps while retaining billing and turn receipts', async () => { - usage.user_is_admin = true - await expect(reserveRunUsage(input)).resolves.toMatchObject({ - status: 'reserved', - }) - expect(mocks.capabilities).toHaveBeenCalledWith( - input.identity.userId, - expect.objectContaining({ execute: mocks.execute }), - ) - expect( - queries.some((query) => - query.sql.includes('INSERT INTO chat_run_usage_receipts'), - ), - ).toBe(true) - expect( - queries.some( - (query) => - query.sql.includes('INSERT INTO chat_funded_spend') && - query.params.includes(250_000), - ), - ).toBe(true) - expect( - queries.filter((query) => - query.sql.includes('INSERT INTO chat_daily_usage'), - ), - ).toHaveLength(1) - await settleFundedSpend('room', 'run', 0.01, false) - expect( - queries.some( - (query) => - query.sql.includes('UPDATE chat_funded_spend') && - query.params.includes(10_000), - ), - ).toBe(true) -}) - -it('continues rejecting a regular account at its included spending cap', async () => { - await expect(reserveRunUsage(input)).rejects.toMatchObject({ - code: 'user-spend', - }) - expect(queries.some((query) => query.sql.includes('INSERT INTO'))).toBe(false) -}) - -it('continues rejecting a regular account at the shared included spending cap', async () => { - usage.user_spend = 0 - await expect(reserveRunUsage(input)).rejects.toMatchObject({ - code: 'global-spend', - }) -}) - -it.each([ - ['user_turns', 20, 'user'], - ['global_turns', 300, 'global'], - ['scheduled_turns', 20, 'scheduled'], -] as const)('retains %s limits for admins', async (field, value, code) => { - usage.user_is_admin = true - usage[field] = value - await expect( - reserveRunUsage({ ...input, scheduled: true }), - ).rejects.toMatchObject({ code }) -}) - -it('avoids resolving admin capabilities when no funded cap applies', async () => { - await expect( - reserveRunUsage({ ...input, fundedSpend: undefined }), - ).resolves.toMatchObject({ status: 'reserved' }) - expect(mocks.capabilities).not.toHaveBeenCalled() -}) diff --git a/harness-tests/pending-runtime/run-usage-admin.test.ts b/harness-tests/pending-runtime/run-usage-admin.test.ts index 166949614..8fa759864 100644 --- a/harness-tests/pending-runtime/run-usage-admin.test.ts +++ b/harness-tests/pending-runtime/run-usage-admin.test.ts @@ -124,3 +124,62 @@ it('increments both daily counters once for concurrent reservations and duplicat await h.db`SELECT turns FROM chat_daily_usage WHERE day='2026-10-04' AND turns=4`, ).toHaveLength(2) }) + +it('rechecks role revocation after waiting for the usage lock', async () => { + const h = await conversationHarness() + await h.db`UPDATE users SET capabilities=ARRAY['admin']::capability[] WHERE id=${input.identity.userId}` + await h.db`INSERT INTO chat_funded_spend(conversation_id,run_id,user_id,day,reserved_micros,billed_micros) VALUES('main-conversation','prior',${input.identity.userId},'2026-10-04',20000000,20000000)` + const result = await h.db.begin(async (tx) => { + await tx.unsafe( + "SELECT pg_advisory_xact_lock(hashtextextended('chat-run-usage:2026-10-04',0))", + ) + const pending = reserveRunUsage(input).then( + (value) => ({ value }), + (error: unknown) => ({ error }), + ) + await expect + .poll(async () => { + const [row] = await tx.unsafe<{ waiting: boolean }[]>( + "SELECT EXISTS(SELECT 1 FROM pg_locks WHERE locktype='advisory' AND NOT granted) AS waiting", + ) + return row.waiting + }) + .toBe(true) + await tx.unsafe( + "UPDATE users SET capabilities='{}'::capability[] WHERE id=$1", + [input.identity.userId], + ) + return { pending } + }) + expect(await result.pending).toMatchObject({ error: { code: 'user-spend' } }) + expect( + await h.db`SELECT run_id FROM chat_run_usage_receipts WHERE run_id='admin-run'`, + ).toHaveLength(0) +}) + +it('rejects a regular account at the shared included spending cap', async () => { + const h = await conversationHarness() + await h.db`INSERT INTO users(id) VALUES('00000000-0000-4000-8000-000000000002')` + await h.db`INSERT INTO chat_funded_spend(conversation_id,run_id,user_id,day,reserved_micros,billed_micros) VALUES('main-conversation','prior','00000000-0000-4000-8000-000000000002','2026-10-04',20000000,20000000)` + await expect(reserveRunUsage(input)).rejects.toMatchObject({ + code: 'global-spend', + }) +}) + +it.each(['user', 'global', 'scheduled'] as const)( + 'retains the %s turn limit for admins', + async (code) => { + const h = await conversationHarness() + await h.db`UPDATE users SET capabilities=ARRAY['admin']::capability[] WHERE id=${input.identity.userId}` + if (code === 'scheduled') + await h.db`INSERT INTO chat_scheduled_daily_usage(user_id,day,turns) VALUES(${input.identity.userId},'2026-10-04',20)` + else + await h.db`INSERT INTO chat_daily_usage(user_id,day,turns) VALUES(${code === 'user' ? input.identity.userId : '__global'},'2026-10-04',${code === 'user' ? 20 : 300})` + await expect( + reserveRunUsage({ ...input, scheduled: true }), + ).rejects.toMatchObject({ code }) + expect( + await h.db`SELECT run_id FROM chat_run_usage_receipts WHERE run_id='admin-run'`, + ).toHaveLength(0) + }, +) diff --git a/scripts/verify-tanchat-schema.ts b/scripts/verify-tanchat-schema.ts index 710a8f5a3..e8be39eea 100644 --- a/scripts/verify-tanchat-schema.ts +++ b/scripts/verify-tanchat-schema.ts @@ -38,6 +38,11 @@ try { missing.push(`${table.name}.${column.name} must be NOT NULL`) } } + const [reservation] = await sql<{ ready: boolean }[]>` + SELECT to_regprocedure('public.reserve_chat_run_usage(jsonb,text)') IS NOT NULL AS ready + ` + if (!reservation.ready) + missing.push('reserve_chat_run_usage function missing') if (missing.length) throw new Error(`Chat schema is not ready:\n${missing.join('\n')}`) console.log( diff --git a/src/chat/conversation-identity.server.ts b/src/chat/conversation-identity.server.ts index 53e271f91..525bd187d 100644 --- a/src/chat/conversation-identity.server.ts +++ b/src/chat/conversation-identity.server.ts @@ -5,6 +5,7 @@ import { chatBots, chatConversations, chatConversationMains, + chatConversationThreads, chatMemberships, } from '~/db/schema' @@ -25,7 +26,7 @@ export class ConversationIdentityError extends Error { } /** Membership and ownership authorize access, never a caller-supplied ID alone. */ -export async function resolveConversationIdentity( +export async function resolveConversationAccess( input: z.infer, ) { const parsed = identityInput.safeParse(input) @@ -50,6 +51,10 @@ export async function resolveConversationIdentity( botId: chatConversations.botId, userId: chatConversations.userId, workspaceId: chatBots.workspaceId, + botArchivedAt: chatBots.archivedAt, + botDeletedAt: chatBots.deletedAt, + threadId: chatConversationThreads.conversationId, + threadArchivedAt: chatConversationThreads.archivedAt, }) .from(chatConversations) .innerJoin(chatBots, eq(chatBots.id, chatConversations.botId)) @@ -64,10 +69,37 @@ export async function resolveConversationIdentity( eq(chatConversationMains.userId, chatConversations.userId), ), ) + .leftJoin( + chatConversationThreads, + eq(chatConversationThreads.conversationId, chatConversations.id), + ) .where(and(...conditions)) .limit(2) if (rows.length !== 1) throw new ConversationIdentityError() - return rows[0] + const row = rows[0] + return { + identity: { + conversationId: row.conversationId, + botId: row.botId, + userId: row.userId, + workspaceId: row.workspaceId, + }, + lifecycle: { + bot: { + archived_at: row.botArchivedAt?.getTime() ?? null, + deleted_at: row.botDeletedAt?.getTime() ?? null, + }, + thread: row.threadId + ? { archived_at: row.threadArchivedAt?.getTime() ?? null } + : null, + }, + } +} + +export async function resolveConversationIdentity( + input: z.infer, +) { + return (await resolveConversationAccess(input)).identity } export type ConversationIdentity = Awaited< diff --git a/src/chat/server/conversation.ts b/src/chat/server/conversation.ts index 64279bcab..c011228f1 100644 --- a/src/chat/server/conversation.ts +++ b/src/chat/server/conversation.ts @@ -218,6 +218,7 @@ import { markConversationRead } from './bot-activity' import { ConversationIdentityError, resolveConversationIdentity, + resolveConversationAccess, type ConversationIdentity, } from '../conversation-identity.server' import { browserExecutionEnabled } from './browser-execution' @@ -2074,6 +2075,9 @@ export class Conversation extends DurableObject { }) } private async authorizeIdentity(identity: ActivityIdentity) { + return (await this.authorizeAccess(identity)).identity + } + private async authorizeAccess(identity: ActivityIdentity) { const old = this.state.identity if ( old && @@ -2085,10 +2089,11 @@ export class Conversation extends DurableObject { old.conversationId !== identity.conversationId)) ) throw new Error('Conversation identity cannot change.') - const resolved = await resolveConversationIdentity({ + const access = await resolveConversationAccess({ ...identity, conversationId: identity.conversationId ?? old?.conversationId, }) + const resolved = access.identity if ( !this.ctx.id.equals( this.env.CONVERSATIONS.idFromName(resolved.conversationId), @@ -2097,7 +2102,7 @@ export class Conversation extends DurableObject { throw new Error('This identity belongs to another conversation.') // Check again after storage authorization, before changing durable state. this.setIdentity(resolved) - return resolved + return access } private setIdentity(identity: ActivityIdentity) { const old = this.state.identity @@ -2291,15 +2296,17 @@ export class Conversation extends DurableObject { ) { throw new Error('This conversation copy is not ready yet.') } - if (input) - await this.authorizeIdentity({ - conversationId: input.conversationId, - botId: input.bot.id, - workspaceId: input.bot.workspace_id, - userId: input.userId, - }) - else if (this.state.identity) - await this.authorizeIdentity(this.state.identity) + const suppliedIdentity = input + ? { + conversationId: input.conversationId, + botId: input.bot.id, + workspaceId: input.bot.workspace_id, + userId: input.userId, + } + : this.state.identity + const access = suppliedIdentity + ? await this.authorizeAccess(suppliedIdentity) + : undefined const identity = this.state.identity if (this.state.copyOrigin?.retryAttemptId && !allowInactive) { const ready = await conversationRetryReady({ @@ -2314,11 +2321,14 @@ export class Conversation extends DurableObject { 'This retry is still preparing its request and files. Resume preparation before sending.', ) } - if (identity && !allowInactive) { - const { bot: metadata, thread } = await readConversationLifecycle( - identity.botId, - identity.conversationId ?? '', - ) + if (identity && access && !allowInactive) { + // Retry readiness can yield after authorization, so read lifecycle again. + const { bot: metadata, thread } = this.state.copyOrigin?.retryAttemptId + ? await readConversationLifecycle( + identity.botId, + identity.conversationId ?? '', + ) + : access.lifecycle if (metadata?.archived_at != null || metadata?.deleted_at != null) throw new ConversationInactiveError( 'Restore this conversation before continuing.', diff --git a/src/chat/server/run-usage.ts b/src/chat/server/run-usage.ts index 6e57567f9..a6c467105 100644 --- a/src/chat/server/run-usage.ts +++ b/src/chat/server/run-usage.ts @@ -1,6 +1,4 @@ import { db } from '~/db/client' -import { isAdmin } from '~/db/types' -import { DrizzleCapabilitiesRepository } from '~/auth/repositories.server' import { sql } from 'drizzle-orm' import { z } from 'zod' import { conversationRunIdentitySchema } from '../core/conversation-runs' @@ -117,105 +115,32 @@ export async function reserveRunUsage( if (identity.userId === '__global') throw new RunUsageConflictError() const day = new Date(now).toISOString().split('T')[0], attemptId = crypto.randomUUID() - return db.transaction(async (tx) => { - await tx.execute( - sql`SELECT pg_advisory_xact_lock(hashtextextended(${JSON.stringify(['chat-run-receipt', identity.conversationId, runId])},0))`, - ) - await tx.execute( - sql`SELECT pg_advisory_xact_lock(hashtextextended(${'chat-run-usage:' + day},0))`, - ) - const [usage] = await tx.execute<{ - receipt: { - workspace_id: string - user_id: string - bot_id: string - scheduled: boolean - day: string - } | null - user_turns: number - global_turns: number - scheduled_turns: number - user_spend: number - global_spend: number - }>(sql`SELECT - (SELECT jsonb_build_object('workspace_id',workspace_id,'user_id',user_id, - 'bot_id',bot_id,'scheduled',scheduled,'day',day) - FROM chat_run_usage_receipts WHERE conversation_id=${identity.conversationId} AND run_id=${runId}) AS receipt, - COALESCE((SELECT turns FROM chat_daily_usage WHERE user_id=${identity.userId} AND day=${day}),0)::double precision AS user_turns, - COALESCE((SELECT turns FROM chat_daily_usage WHERE user_id='__global' AND day=${day}),0)::double precision AS global_turns, - COALESCE((SELECT turns FROM chat_scheduled_daily_usage WHERE user_id=${identity.userId}::uuid AND day=${day}),0)::double precision AS scheduled_turns, - COALESCE((SELECT SUM(billed_micros) FROM chat_funded_spend WHERE user_id=${identity.userId}::uuid AND day=${day}),0)::double precision AS user_spend, - COALESCE((SELECT SUM(billed_micros) FROM chat_funded_spend WHERE day=${day}),0)::double precision AS global_spend`) - const receipt = usage.receipt - if (receipt) { - if ( - receipt.workspace_id !== identity.workspaceId || - receipt.user_id !== identity.userId || - receipt.bot_id !== identity.botId || - receipt.scheduled !== scheduled - ) - throw new RunUsageConflictError() - return { status: 'duplicate' as const, day: receipt.day } - } - if (!lift && usage.user_turns >= dailyLimit) - throw new RunUsageAllowanceError('user') - if (!lift && usage.global_turns >= 300) - throw new RunUsageAllowanceError('global') - if (!unlimited && scheduled && usage.scheduled_turns >= scheduledLimit) - throw new RunUsageAllowanceError('scheduled') - // Use the same effective capabilities as authentication, including roles. - const includedAdmin = - spend && !unlimited - ? isAdmin( - await new DrizzleCapabilitiesRepository().getEffectiveCapabilities( - identity.userId, - tx, - ), - ) - : false - if ( - !unlimited && - !includedAdmin && - spend && - usage.user_spend + spend.reservationMicros > spend.userCapMicros - ) - throw new RunUsageAllowanceError('user-spend') - if ( - !unlimited && - !includedAdmin && - spend && - usage.global_spend + spend.reservationMicros > spend.globalCapMicros - ) - throw new RunUsageAllowanceError('global-spend') - // Publish the receipt, funded reservation, and counters atomically in one - // statement while holding the existing receipt and day locks. - await tx.execute(sql`WITH receipt AS ( - INSERT INTO chat_run_usage_receipts(conversation_id,run_id,workspace_id,user_id,bot_id,scheduled,day,created_at,attempt_id) - VALUES(${identity.conversationId},${runId},${identity.workspaceId},${identity.userId}::uuid,${identity.botId},${scheduled},${day},${now},${attemptId}::uuid) - RETURNING conversation_id - ), turns AS ( - INSERT INTO chat_daily_usage(user_id,day,turns) - SELECT counters.user_id,${day},1 FROM receipt CROSS JOIN (VALUES (${identity.userId}),('__global')) counters(user_id) - ON CONFLICT(user_id,day) DO UPDATE SET turns=chat_daily_usage.turns+1 - )${ - spend - ? sql`, funded AS ( - INSERT INTO chat_funded_spend(conversation_id,run_id,user_id,day,reserved_micros,billed_micros) - SELECT conversation_id,${runId},${identity.userId}::uuid,${day},${spend.reservationMicros},${spend.reservationMicros} FROM receipt - )` - : sql`` - }${ - scheduled - ? sql`, scheduled AS ( - INSERT INTO chat_scheduled_daily_usage(user_id,day,turns) - SELECT ${identity.userId}::uuid,${day},1 FROM receipt - ON CONFLICT(user_id,day) DO UPDATE SET turns=chat_scheduled_daily_usage.turns+1 - )` - : sql`` - } - SELECT conversation_id FROM receipt`) - return { status: 'reserved' as const, day } + const request = JSON.stringify({ + ...identity, + runId, + scheduled, + day, + now, + attemptId, + dailyLimit, + scheduledLimit, + lift, + unlimited, + spend, }) + const [result] = await db.execute<{ + outcome: + | 'reserved' + | 'duplicate' + | 'conflict' + | ConstructorParameters[0] + receipt_day: string + }>(sql`SELECT * FROM public.reserve_chat_run_usage(${request}::jsonb, + ${JSON.stringify(['chat-run-receipt', identity.conversationId, runId])})`) + if (result.outcome === 'conflict') throw new RunUsageConflictError() + if (result.outcome !== 'reserved' && result.outcome !== 'duplicate') + throw new RunUsageAllowanceError(result.outcome) + return { status: result.outcome, day: result.receipt_day } } export async function fundedSpendReport( From ed980e3c1c44199c29151af46df23a0f869c7498 Mon Sep 17 00:00:00 2001 From: Tanner Linsley Date: Mon, 5 Oct 2026 10:23:46 -0600 Subject: [PATCH 2/2] fix(chat): keep transcript stable when sending --- .../core/conversation-events.test.ts | 95 ++++++++++++++++++- src/chat/components/ConversationSession.tsx | 28 +++--- src/chat/components/ConversationTurn.tsx | 31 +----- src/chat/components/VirtualMessages.tsx | 19 ---- src/chat/core/conversation-snapshot.ts | 26 +++++ src/chat/server/conversation-events.ts | 43 +++++++-- 6 files changed, 172 insertions(+), 70 deletions(-) create mode 100644 src/chat/core/conversation-snapshot.ts diff --git a/harness-tests/core/conversation-events.test.ts b/harness-tests/core/conversation-events.test.ts index 339559f43..595032b2f 100644 --- a/harness-tests/core/conversation-events.test.ts +++ b/harness-tests/core/conversation-events.test.ts @@ -1,3 +1,4 @@ +import { retainArchivedMessages } from '../../src/chat/core/conversation-snapshot' import { describe, expect, it } from 'vitest' import { StreamProcessor } from '@tanstack/ai' import { readMessageAttachments } from '../../src/chat/core/message-attachments' @@ -105,6 +106,7 @@ describe('conversation AG-UI projection', () => { role: 'user', parts: [{ type: 'text', content: 'A copied message' }], metadata: { gumInherited: true, obsolete: true }, + createdAt: new Date('2026-10-01T12:00:00Z'), }, ], } @@ -113,8 +115,14 @@ describe('conversation AG-UI projection', () => { const processor = new StreamProcessor() for (const event of conversationEvents(undefined, before, 'chat')) processor.processChunk(event) - for (const event of conversationEvents(before, after, 'chat')) + for (const event of conversationEvents(before, after, 'chat')) { processor.processChunk(event) + expect(processor.getMessages()).toHaveLength(1) + expect(processor.getMessages()[0].parts).toEqual(before.messages[0].parts) + expect(processor.getMessages()[0].createdAt).toEqual( + before.messages[0].createdAt, + ) + } expect(processor.getMessages()).toHaveLength(1) expect(processor.getMessages()[0]).toMatchObject({ metadata: { gumInherited: true }, @@ -238,6 +246,89 @@ describe('conversation AG-UI projection', () => { }) } }) + it('keeps displayed history when completed turns leave the live server window', () => { + const before: ProjectableConversation = { + ...initial(), + messages: [ + { + id: 'old', + role: 'user', + parts: [{ type: 'text', content: 'Earlier' }], + }, + { + id: 'current', + role: 'user', + parts: [{ type: 'text', content: 'Current' }], + }, + ], + } + const after: ProjectableConversation = { + ...before, + archivedTurns: 1, + messages: [ + before.messages[1], + { id: 'new', role: 'user', parts: [{ type: 'text', content: 'New' }] }, + ], + } + const processor = new StreamProcessor() + for (const event of conversationEvents(undefined, before, 'chat')) + processor.processChunk(event) + for (const event of conversationEvents(before, after, 'chat')) { + processor.processChunk(event) + expect(processor.getMessages()[0]).toMatchObject(before.messages[0]) + } + expect(processor.getMessages().map((message) => message.id)).toEqual([ + 'old', + 'current', + 'new', + ]) + }) + it('keeps archived rows across an authoritative metadata snapshot and clears on reset', () => { + const archived: ProjectableConversation = { + ...initial(), + archivedTurns: 1, + messages: [ + { + id: 'current', + role: 'user', + parts: [{ type: 'text', content: 'Current' }], + metadata: { obsolete: true }, + }, + ], + } + const displayed = [ + { + id: 'earlier', + role: 'user', + parts: [{ type: 'text', content: 'Earlier' }], + }, + ...archived.messages, + ] satisfies ProjectableConversation['messages'] + const next = structuredClone(archived) + next.messages[0].metadata = undefined + const processor = new StreamProcessor() + for (const event of conversationEvents( + undefined, + { ...archived, messages: displayed }, + 'chat', + )) + processor.processChunk(event) + for (const event of conversationEvents(archived, next, 'chat')) { + processor.processChunk( + retainArchivedMessages(event, processor.getMessages()), + ) + expect(processor.getMessages().map((message) => message.id)).toEqual([ + 'earlier', + 'current', + ]) + } + expect(processor.getMessages()[1].metadata).not.toHaveProperty('obsolete') + for (const event of conversationEvents(next, initial(), 'chat')) + processor.processChunk( + retainArchivedMessages(event, processor.getMessages()), + ) + expect(processor.getMessages()).toEqual([]) + }) it('publishes only appended text between saves and clears transcript on reset', () => { const old = { ...initial(), @@ -256,7 +347,7 @@ describe('conversation AG-UI projection', () => { messageId: 'a', delta: ' world', }) - expect(conversationEvents(next, initial(), 'chat')[0]).toEqual({ + expect(conversationEvents(next, initial(), 'chat')[0]).toMatchObject({ type: 'MESSAGES_SNAPSHOT', messages: [], }) diff --git a/src/chat/components/ConversationSession.tsx b/src/chat/components/ConversationSession.tsx index c4270ffc7..48edf94d8 100644 --- a/src/chat/components/ConversationSession.tsx +++ b/src/chat/components/ConversationSession.tsx @@ -1,3 +1,4 @@ +import { retainArchivedMessages } from '../core/conversation-snapshot' import { workspacePreviewsEnabled } from '../core/workspace-panels' import { useQueryClient } from '@tanstack/react-query' import { projectSavedEvent } from '../core/project-events' @@ -5,7 +6,6 @@ import { ProjectsPanel } from './ProjectsPanel' import { isPersonalAssistant } from '../core/bot-workspace' import './assistant-chat.css' import { DraftSyncNotice } from './DraftSyncNotice' -import { motion } from 'motion/react' import { useDebugDetails } from './useDebugDetails' import { Button } from './ui/Button' import { conversationActionRequest } from '../core/conversation-actions' @@ -461,6 +461,7 @@ export function ConversationSession({ const forkMessage = onFork ? (messageId: string) => onFork(messageId, destination.conversationId) : undefined + const displayedMessages = useRef(initialHistory.messages) const connection = useMemo(() => { const durable = durableStreamConnection({ sendUrl: apiUrl(`${path}/send`), @@ -468,8 +469,13 @@ export function ConversationSession({ initialOffset: initialHistory.streamOffset, emitSnapshotOnSubscribe: false, }) + const subscribe = + durable.subscribe as SubscribeConnectionAdapter['subscribe'] return { - subscribe: durable.subscribe as SubscribeConnectionAdapter['subscribe'], + subscribe: async function* (signal) { + for await (const chunk of subscribe(signal)) + yield retainArchivedMessages(chunk, displayedMessages.current) + }, send: async (messages, data) => { const message = messages.at(-1) if (!message || !('parts' in message)) @@ -525,6 +531,7 @@ export function ConversationSession({ } }, }) + displayedMessages.current = chat.messages const h = { ...metadata, messages: chat.messages } const emptyThreadState = useRef(false) emptyThreadState.current = @@ -1121,8 +1128,7 @@ export function ConversationSession({ origin={h.copyOrigin} /> )} - {thread && ( - @@ -1145,7 +1144,7 @@ export function ConversationSession({ {thread.source.text} {thread.source.truncated &&

Source excerpt

} - + )} {(!!h.archivedTurns || (!!search.message && !liveTarget)) && ( navigation.onSelectMessage(undefined)} > {turns.map((turn, index, turns) => ( @@ -1441,7 +1439,7 @@ export function ConversationSession({ ))} )} - +
diff --git a/src/chat/components/ConversationTurn.tsx b/src/chat/components/ConversationTurn.tsx index a5c1990c4..cbf9ed789 100644 --- a/src/chat/components/ConversationTurn.tsx +++ b/src/chat/components/ConversationTurn.tsx @@ -1,4 +1,3 @@ -import { motion } from 'motion/react' import type { UIMessage, ToolResultPart, ToolCallPart } from '@tanstack/ai' import { useEffect, useRef, useState, type ReactNode } from 'react' import { @@ -296,18 +295,7 @@ export function ConversationTurn({ )} {turn.prompt && ( - )} - + )}
{interim.map((m) => ( @@ -497,18 +485,7 @@ export function ConversationTurn({ finalParts.some( (part) => part.type === 'text' && part.content.trim(), ) && ( - {final.id !== responseEnd?.id && renderMessageThreads?.(final.id, { canCreate: false })} - + )} {showResponseActions && (
void onJumpToBottom?: () => void jumpToBottomContainer?: HTMLDivElement | null - followRequest?: string contentVersion?: string | number messageScope?: string }) { @@ -114,23 +112,6 @@ export function VirtualMessages({ } }, [rows.length, scrollElement, scrollMargin, virtualizer]) - const followedRequest = useRef(undefined) - useLayoutEffect(() => { - if ( - !followRequest || - followedRequest.current === followRequest || - !scrollElement || - scrollMargin === null || - !rows.length - ) - return - followedRequest.current = followRequest - pendingTarget.current = undefined - pendingFocus.current?.cancel() - pendingFocus.current = null - virtualizer.scrollToEnd() - }, [followRequest, scrollElement, scrollMargin, rows.length, virtualizer]) - const alignTarget = () => { if (!pendingTarget.current || !scrollElement) return if (pendingFocus.current && !pendingFocus.current.active()) { diff --git a/src/chat/core/conversation-snapshot.ts b/src/chat/core/conversation-snapshot.ts new file mode 100644 index 000000000..076e4d17a --- /dev/null +++ b/src/chat/core/conversation-snapshot.ts @@ -0,0 +1,26 @@ +import type { StreamChunk, UIMessage } from '@tanstack/ai' + +// The server window omits archived turns, but an open transcript keeps them. +// A reset has no boundary and still replaces the entire displayed history. +export function retainArchivedMessages( + chunk: StreamChunk, + displayed: UIMessage[], +): StreamChunk { + if (chunk.type !== 'MESSAGES_SNAPSHOT') return chunk + const boundary = chunk.rawEvent?.retainedHistoryBefore + if (typeof boundary !== 'string') return chunk + const index = displayed.findIndex((message) => message.id === boundary) + if (index < 1) return chunk + return { + ...chunk, + messages: [ + ...displayed.slice(0, index).map((message) => ({ + ...message, + content: message.parts + .flatMap((part) => (part.type === 'text' ? [part.content] : [])) + .join(''), + })), + ...chunk.messages, + ], + } +} diff --git a/src/chat/server/conversation-events.ts b/src/chat/server/conversation-events.ts index b2beca0da..349f0c32e 100644 --- a/src/chat/server/conversation-events.ts +++ b/src/chat/server/conversation-events.ts @@ -6,6 +6,7 @@ export interface ProjectableConversation { messages: UIMessage[] status: string activeRun: string | null + archivedTurns?: number error?: string } @@ -31,19 +32,34 @@ export function conversationEvents( !previous || previous.messages.some((m) => { const message = nextById.get(m.id) - // AG-UI metadata deltas merge keys. Replay is required when authoritative - // metadata changes so removed fields cannot remain on an existing message. + // Replace authoritative metadata atomically so removed fields cannot remain + // on an existing message, without clearing and replaying the transcript. return ( - !message || - JSON.stringify(m.metadata) !== JSON.stringify(message.metadata) + (!message && + (next.archivedTurns ?? 0) <= (previous.archivedTurns ?? 0)) || + (message && + JSON.stringify(m.metadata) !== JSON.stringify(message.metadata)) ) }) if (reset) { - // Rebuild history using the same events as live delivery, including tools. - events.push({ type: EventType.MESSAGES_SNAPSHOT, messages: [] }) + // The native processor preserves UI parts and timestamps in snapshots. + events.push({ + type: EventType.MESSAGES_SNAPSHOT, + rawEvent: { + retainedHistoryBefore: + (next.archivedTurns ?? 0) > 0 ? next.messages[0]?.id : undefined, + }, + messages: next.messages.map((message) => ({ + ...message, + metadata: message.metadata ?? {}, + content: message.parts + .flatMap((part) => (part.type === 'text' ? [part.content] : [])) + .join(''), + })), + }) } const before = new Map((reset ? [] : previous.messages).map((m) => [m.id, m])) - for (const message of next.messages) { + for (const message of reset ? [] : next.messages) { const old = before.get(message.id) let textChanged = false if (!old) @@ -163,6 +179,19 @@ export function conversationEvents( }, }) if (previous?.status === 'running' && next.status !== 'running') { + // Error events need an existing assistant anchor after a snapshot clears + // transient processor state. No text is replayed into that message. + if (next.status === 'error') { + const assistant = [...next.messages] + .reverse() + .find((message) => message.role === 'assistant') + if (assistant) + events.push({ + type: EventType.TEXT_MESSAGE_START, + messageId: assistant.id, + role: assistant.role, + }) + } events.push( next.status === 'error' ? {