From 8a1b7710d776231fd10dfe580c92c8139da0fd78 Mon Sep 17 00:00:00 2001 From: Shakker Date: Tue, 26 May 2026 13:47:08 +0100 Subject: [PATCH] refactor: add user turn transcript recorder --- src/gateway/server-methods/chat.ts | 99 +++++++-------- src/sessions/user-turn-transcript.ts | 172 +++++++++++++++++++++++++-- 2 files changed, 214 insertions(+), 57 deletions(-) diff --git a/src/gateway/server-methods/chat.ts b/src/gateway/server-methods/chat.ts index afd1536b7dbc..d5ca3111a58f 100644 --- a/src/gateway/server-methods/chat.ts +++ b/src/gateway/server-methods/chat.ts @@ -52,8 +52,9 @@ import { normalizeInputProvenance, type InputProvenance } from "../../sessions/i import { resolveSendPolicy } from "../../sessions/send-policy.js"; import { parseAgentSessionKey } from "../../sessions/session-key-utils.js"; import { - persistInlineUserTurnTranscript, + createUserTurnTranscriptRecorder, type UserTurnInput, + type UserTurnTranscriptRecorder, } from "../../sessions/user-turn-transcript.js"; import { uniqueStrings } from "../../shared/string-normalization.js"; import { deliveryContextFromSession } from "../../utils/delivery-context.shared.js"; @@ -2698,56 +2699,48 @@ export const chatHandlers: GatewayRequestHandlers = { }; const deliveredReplies: Array<{ payload: ReplyPayload; kind: "block" | "final" }> = []; let appendedWebchatAgentMedia = false; - let userTranscriptUpdatePromise: Promise | null = null; - let runtimeUserTranscriptPersistencePromise: Promise | null = null; let agentRunStarted = false; - let agentUserMessagePersisted = false; - let beforeAgentRunBlocked = false; - const persistGatewayUserTurnTranscript = async () => { - const runtimePersistence = runtimeUserTranscriptPersistencePromise; - if (runtimePersistence) { - await runtimePersistence.catch((error) => { - context.logGateway.warn( - `runtime user transcript persistence failed before fallback: ${formatForLog(error)}`, - ); - }); - if (agentUserMessagePersisted) { - return; - } - } - if (userTranscriptUpdatePromise) { - await userTranscriptUpdatePromise; - return; - } - userTranscriptUpdatePromise = (async () => { - await measureDiagnosticsTimelineSpan( - "gateway.chat_send.persist_user_transcript", - async () => { + const userTurnRecorderPromise: Promise = + userTurnInputPromise.then((input) => + createUserTurnTranscriptRecorder({ + input, + target: () => { const { storePath: latestStorePath, entry: latestEntry } = loadSessionEntry(sessionKey); const resolvedSessionId = latestEntry?.sessionId ?? backingSessionId; if (!resolvedSessionId) { - return; + return undefined; } - const userTurnInput = await userTurnInputPromise; - await persistInlineUserTurnTranscript({ + return { sessionId: resolvedSessionId, sessionKey, sessionEntry: latestEntry ?? entry, storePath: latestStorePath, agentId, config: cfg, - input: userTurnInput, - }); + }; }, - { - phase: "agent-turn", - config: cfg, - attributes: chatSendTraceAttributes, + errorContext: "gateway chat user turn transcript", + onPersistenceError: (error) => { + context.logGateway.warn( + `gateway user transcript persistence failed: ${formatForLog(error)}`, + ); }, - ); - })(); - await userTranscriptUpdatePromise; + }), + ); + const persistGatewayUserTurnTranscript = async () => { + const userTurnRecorder = await userTurnRecorderPromise; + await measureDiagnosticsTimelineSpan( + "gateway.chat_send.persist_user_transcript", + async () => { + await userTurnRecorder.persistFallback(); + }, + { + phase: "agent-turn", + config: cfg, + attributes: chatSendTraceAttributes, + }, + ); }; const appendWebchatAgentMediaTranscriptIfNeeded = async (payload: ReplyPayload) => { if (!agentRunStarted || appendedWebchatAgentMedia || !isMediaBearingPayload(payload)) { @@ -2847,7 +2840,8 @@ export const chatHandlers: GatewayRequestHandlers = { }, deliver: async (payload, info) => { if (getReplyPayloadMetadata(payload)?.beforeAgentRunBlocked === true) { - beforeAgentRunBlocked = true; + const userTurnRecorder = await userTurnRecorderPromise; + userTurnRecorder.markBlocked(); } switch (info.kind) { case "block": @@ -2906,11 +2900,15 @@ export const chatHandlers: GatewayRequestHandlers = { } } }, - onUserMessagePersisted: () => { - agentUserMessagePersisted = true; + onUserMessagePersisted: (message) => { + return userTurnRecorderPromise.then((userTurnRecorder) => { + userTurnRecorder.markRuntimePersisted(message); + }); }, onUserMessagePersistencePending: (pending) => { - runtimeUserTranscriptPersistencePromise = pending; + void userTurnRecorderPromise.then((userTurnRecorder) => { + userTurnRecorder.markRuntimePersistencePending(pending); + }); }, onModelSelected: (modelSelection) => { updateChatRunProvider(context.chatAbortControllers, { @@ -2924,7 +2922,10 @@ export const chatHandlers: GatewayRequestHandlers = { }, }, }); - beforeAgentRunBlocked = dispatchResult.beforeAgentRunBlocked === true; + if (dispatchResult.beforeAgentRunBlocked === true) { + const userTurnRecorder = await userTurnRecorderPromise; + userTurnRecorder.markBlocked(); + } return dispatchResult; }, { @@ -2947,20 +2948,21 @@ export const chatHandlers: GatewayRequestHandlers = { .map((payload) => payload.text?.trim()) .filter((text): text is string => Boolean(text)) .join(" | ") || undefined; + const userTurnRecorder = await userTurnRecorderPromise; if ( agentRunStarted && returnedAgentErrorPayloads.length > 0 && - !agentUserMessagePersisted && - !beforeAgentRunBlocked + !userTurnRecorder.hasPersisted() && + !userTurnRecorder.isBlocked() ) { await persistGatewayUserTurnTranscript(); } if ( agentRunStarted && returnedAgentErrorPayloads.length === 0 && - !agentUserMessagePersisted && - !beforeAgentRunBlocked && - runtimeUserTranscriptPersistencePromise + !userTurnRecorder.hasPersisted() && + !userTurnRecorder.isBlocked() && + userTurnRecorder.hasRuntimePersistencePending() ) { await persistGatewayUserTurnTranscript(); } @@ -3507,8 +3509,9 @@ export const chatHandlers: GatewayRequestHandlers = { ); }) .catch(async (err) => { + const userTurnRecorder = await userTurnRecorderPromise; const emitAfterError = - agentUserMessagePersisted || beforeAgentRunBlocked + userTurnRecorder.hasPersisted() || userTurnRecorder.isBlocked() ? Promise.resolve() : persistGatewayUserTurnTranscript(); await emitAfterError.catch((transcriptErr) => { diff --git a/src/sessions/user-turn-transcript.ts b/src/sessions/user-turn-transcript.ts index 4493bb903b37..9d1e6f59f16f 100644 --- a/src/sessions/user-turn-transcript.ts +++ b/src/sessions/user-turn-transcript.ts @@ -74,6 +74,46 @@ export type UserTurnTranscriptPersistenceTarget = Omit< "input" | "message" | "updateMode" >; +export type UserTurnTranscriptPersistResult = { + sessionFile: string; + sessionEntry: SessionEntry | undefined; + messageId: string; + message: PersistedUserTurnMessage; +}; + +export type UserTurnTranscriptTargetResolver = + | UserTurnTranscriptPersistenceTarget + | (() => + | UserTurnTranscriptPersistenceTarget + | undefined + | Promise); + +export type UserTurnTranscriptRecorder = { + readonly message: PersistedUserTurnMessage | undefined; + markRuntimePersistencePending: (pending: Promise) => void; + markRuntimePersisted: (message?: PersistedUserTurnMessage) => void; + markBlocked: () => void; + hasPersisted: () => boolean; + isBlocked: () => boolean; + hasRuntimePersistencePending: () => boolean; + waitForRuntimePersistence: () => Promise; + persistApproved: (params?: { + updateMode?: UserTurnTranscriptUpdateMode; + }) => Promise; + persistFallback: (params?: { + updateMode?: UserTurnTranscriptUpdateMode; + }) => Promise; +}; + +export type CreateUserTurnTranscriptRecorderParams = { + input?: UserTurnInput; + message?: PersistedUserTurnMessage; + target: UserTurnTranscriptTargetResolver; + updateMode?: UserTurnTranscriptUpdateMode; + errorContext?: string; + onPersistenceError?: (error: unknown) => void; +}; + type InlineUserTurnTranscriptSource = | { message: PersistedUserTurnMessage; @@ -397,15 +437,9 @@ export async function appendInlineUserTurnTranscriptMessage( }); } -export async function persistUserTurnTranscript(params: PersistUserTurnTranscriptParams): Promise< - | { - sessionFile: string; - sessionEntry: SessionEntry | undefined; - messageId: string; - message: PersistedUserTurnMessage; - } - | undefined -> { +export async function persistUserTurnTranscript( + params: PersistUserTurnTranscriptParams, +): Promise { const message = resolvePersistedUserTurnMessage(params); if (!message) { return undefined; @@ -485,3 +519,123 @@ export async function tryPersistInlineUserTurnTranscript( return undefined; } } + +async function resolveUserTurnTranscriptTarget( + target: UserTurnTranscriptTargetResolver, +): Promise { + return typeof target === "function" ? await target() : target; +} + +export function createUserTurnTranscriptRecorder( + params: CreateUserTurnTranscriptRecorderParams, +): UserTurnTranscriptRecorder { + const message = resolvePersistedUserTurnMessage(params); + let blocked = false; + let persisted = false; + let persistedResult: UserTurnTranscriptPersistResult | undefined; + let runtimePersistencePromise: Promise | undefined; + let selfPersistencePromise: Promise | undefined; + + const handlePersistenceError = (error: unknown) => { + if (params.onPersistenceError) { + params.onPersistenceError(error); + return; + } + logVerbose( + `failed to persist ${params.errorContext ?? "user turn transcript"}: ${String(error)}`, + ); + }; + + const waitForRuntimePersistence = async () => { + if (!runtimePersistencePromise) { + return; + } + try { + await runtimePersistencePromise; + } catch (error) { + handlePersistenceError(error); + } + }; + + const persistPrepared = async (options: { + waitForRuntime: boolean; + skipWhenBlocked: boolean; + updateMode?: UserTurnTranscriptUpdateMode; + }): Promise => { + if (persisted) { + return persistedResult; + } + if (options.skipWhenBlocked && blocked) { + return undefined; + } + if (!message) { + return undefined; + } + if (options.waitForRuntime) { + await waitForRuntimePersistence(); + if (persisted) { + return persistedResult; + } + } + if (selfPersistencePromise) { + return await selfPersistencePromise; + } + selfPersistencePromise = (async () => { + const target = await resolveUserTurnTranscriptTarget(params.target); + if (!target) { + return undefined; + } + const result = await persistUserTurnTranscript({ + ...target, + message, + updateMode: options.updateMode ?? params.updateMode ?? "inline", + }); + if (result) { + persisted = true; + persistedResult = result; + } + return result; + })(); + try { + return await selfPersistencePromise; + } catch (error) { + handlePersistenceError(error); + throw error; + } + }; + + return { + message, + markRuntimePersistencePending: (pending) => { + runtimePersistencePromise = pending; + }, + markRuntimePersisted: (persistedMessage) => { + persisted = true; + if (persistedMessage && persistedResult) { + persistedResult = { + ...persistedResult, + message: persistedMessage, + }; + } + }, + markBlocked: () => { + blocked = true; + }, + hasPersisted: () => persisted, + isBlocked: () => blocked, + hasRuntimePersistencePending: () => runtimePersistencePromise !== undefined, + waitForRuntimePersistence, + persistApproved: async (options) => + await persistPrepared({ + waitForRuntime: false, + skipWhenBlocked: true, + updateMode: options?.updateMode, + }), + persistFallback: async (options) => + await persistPrepared({ + waitForRuntime: true, + skipWhenBlocked: true, + updateMode: options?.updateMode, + }), + }; +}