refactor: add user turn transcript recorder

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<void> | null = null;
-      let runtimeUserTranscriptPersistencePromise: Promise<void> | 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<UserTurnTranscriptRecorder> =
+        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) => {

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<UserTurnTranscriptPersistenceTarget | undefined>);
+
+export type UserTurnTranscriptRecorder = {
+  readonly message: PersistedUserTurnMessage | undefined;
+  markRuntimePersistencePending: (pending: Promise<void>) => void;
+  markRuntimePersisted: (message?: PersistedUserTurnMessage) => void;
+  markBlocked: () => void;
+  hasPersisted: () => boolean;
+  isBlocked: () => boolean;
+  hasRuntimePersistencePending: () => boolean;
+  waitForRuntimePersistence: () => Promise<void>;
+  persistApproved: (params?: {
+    updateMode?: UserTurnTranscriptUpdateMode;
+  }) => Promise<UserTurnTranscriptPersistResult | undefined>;
+  persistFallback: (params?: {
+    updateMode?: UserTurnTranscriptUpdateMode;
+  }) => Promise<UserTurnTranscriptPersistResult | undefined>;
+};
+
+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<UserTurnTranscriptPersistResult | undefined> {
   const message = resolvePersistedUserTurnMessage(params);
   if (!message) {
     return undefined;
@@ -485,3 +519,123 @@ export async function tryPersistInlineUserTurnTranscript(
     return undefined;
   }
 }
+
+async function resolveUserTurnTranscriptTarget(
+  target: UserTurnTranscriptTargetResolver,
+): Promise<UserTurnTranscriptPersistenceTarget | undefined> {
+  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<void> | undefined;
+  let selfPersistencePromise: Promise<UserTurnTranscriptPersistResult | undefined> | 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<UserTurnTranscriptPersistResult | undefined> => {
+    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,
+      }),
+  };
+}