refactor(agents): require live edit diff state (#121722)

src/agents/agent-tool-handler-state.test-helpers.ts

@@ -13,6 +13,7 @@ export function createBaseToolHandlerState() {
     toolMetas: [] as Array<{ toolName?: string; meta?: string; asyncStarted?: boolean }>,
     acceptedSessionSpawns: [],
     toolSummaryById: new Set<string>(),
+    liveEditDiffStateById: new Map(),
     itemActiveIds: new Set<string>(),
     itemStartedCount: 0,
     itemCompletedCount: 0,

src/agents/embedded-agent-subscribe.handlers.lifecycle.test.ts

@@ -66,6 +66,7 @@ function createContext(
     },
     state: {
       lastAssistant: lastAssistant as EmbeddedAgentSubscribeContext["state"]["lastAssistant"],
+      liveEditDiffStateById: new Map(),
       pendingCompactionRetry: 0,
       pendingToolMediaUrls: [],
       pendingToolMediaTrustByUrl: new Map(),

src/agents/embedded-agent-subscribe.handlers.lifecycle.ts

@@ -67,7 +67,7 @@ export function handleAgentEnd(
   ctx: EmbeddedAgentSubscribeContext,
   evt?: Extract<AgentSessionEvent, { type: "agent_end" }>,
 ): void | Promise<void> {
-  ctx.state.liveEditDiffStateById?.clear();
+  ctx.state.liveEditDiffStateById.clear();
   type BeforeTerminalDeliveryDecision = void | { suppressTerminalDelivery?: boolean };
   const lastAssistant = ctx.state.lastAssistant;
   const isError = isAssistantMessage(lastAssistant) && lastAssistant.stopReason === "error";

src/agents/embedded-agent-subscribe.handlers.tools.media.test.ts

@@ -29,6 +29,7 @@ function createMockContext(overrides?: {
       toolMetaById: new Map(),
       toolMetas: [],
       toolSummaryById: new Set(),
+      liveEditDiffStateById: new Map(),
       itemActiveIds: new Set(),
       itemStartedCount: 0,
       itemCompletedCount: 0,

src/agents/embedded-agent-subscribe.handlers.tools.test.ts

@@ -178,6 +178,7 @@ function createTestContext(): {
       toolMetas: [],
       acceptedSessionSpawns: [],
       toolSummaryById: new Set<string>(),
+      liveEditDiffStateById: new Map(),
       itemActiveIds: new Set<string>(),
       itemStartedCount: 0,
       itemCompletedCount: 0,

src/agents/embedded-agent-subscribe.handlers.tools.ts

@@ -1087,7 +1087,7 @@ export function handleToolExecutionStart(
   },
 ): void | Promise<void> {
   const startToolName = normalizeToolName(evt.toolName);
-  ctx.state.liveEditDiffStateById?.delete(evt.toolCallId);
+  ctx.state.liveEditDiffStateById.delete(evt.toolCallId);
   const askUserPromptReservation =
     startToolName === "ask_user" && ctx.params.onToolResult
       ? buildAskUserPromptPayload(evt.toolCallId, ctx.params.sessionKey, ctx.params.runId, evt.args)
@@ -1491,7 +1491,7 @@ export async function handleToolExecutionEnd(
   const toolName = normalizeToolName(rawToolName);
   const hideFromChannelProgress = evt.hideFromChannelProgress === true;
   const toolCallId = evt.toolCallId;
-  ctx.state.liveEditDiffStateById?.delete(toolCallId);
+  ctx.state.liveEditDiffStateById.delete(toolCallId);
   if (toolName === "ask_user") {
     cancelAskUserPromptDelivery(toolCallId, ctx.params.sessionKey, ctx.params.runId);
   }

src/agents/embedded-agent-subscribe.handlers.types.ts

@@ -342,6 +342,7 @@ type ToolHandlerState = Pick<
   | "acceptedSessionSpawns"
   | "toolSummaryById"
   | "execLiveUpdateStateById"
+  | "liveEditDiffStateById"
   | "itemActiveIds"
   | "itemStartedCount"
   | "itemCompletedCount"
@@ -368,9 +369,7 @@ type ToolHandlerState = Pick<
   | "deterministicApprovalPromptSent"
   | "toolExecutionSinceLastBlockReply"
   | "assistantMessageIndex"
-> & {
-  liveEditDiffStateById?: EmbeddedAgentSubscribeState["liveEditDiffStateById"];
-};
+>;
 
 export type ToolHandlerContext = {
   params: ToolHandlerParams;