diff --git a/src/agents/agent-tool-handler-state.test-helpers.ts b/src/agents/agent-tool-handler-state.test-helpers.ts
index d7701a920b6a..773abed05214 100644
--- a/src/agents/agent-tool-handler-state.test-helpers.ts
+++ b/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,
diff --git a/src/agents/embedded-agent-subscribe.handlers.lifecycle.test.ts b/src/agents/embedded-agent-subscribe.handlers.lifecycle.test.ts
index 034a75018f9b..36d09eb351c7 100644
--- a/src/agents/embedded-agent-subscribe.handlers.lifecycle.test.ts
+++ b/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(),
diff --git a/src/agents/embedded-agent-subscribe.handlers.lifecycle.ts b/src/agents/embedded-agent-subscribe.handlers.lifecycle.ts
index 9be5106f88c1..999cf84b1e76 100644
--- a/src/agents/embedded-agent-subscribe.handlers.lifecycle.ts
+++ b/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";
diff --git a/src/agents/embedded-agent-subscribe.handlers.tools.media.test.ts b/src/agents/embedded-agent-subscribe.handlers.tools.media.test.ts
index e7bb9dfbd2c4..a0e6a44e246d 100644
--- a/src/agents/embedded-agent-subscribe.handlers.tools.media.test.ts
+++ b/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,
diff --git a/src/agents/embedded-agent-subscribe.handlers.tools.test.ts b/src/agents/embedded-agent-subscribe.handlers.tools.test.ts
index add9b5d61b35..6446fa62ea9b 100644
--- a/src/agents/embedded-agent-subscribe.handlers.tools.test.ts
+++ b/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,
diff --git a/src/agents/embedded-agent-subscribe.handlers.tools.ts b/src/agents/embedded-agent-subscribe.handlers.tools.ts
index 76f223d90c0f..e1c8f482e7d9 100644
--- a/src/agents/embedded-agent-subscribe.handlers.tools.ts
+++ b/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);
   }
diff --git a/src/agents/embedded-agent-subscribe.handlers.types.ts b/src/agents/embedded-agent-subscribe.handlers.types.ts
index e519fb22423a..8785e8b96acf 100644
--- a/src/agents/embedded-agent-subscribe.handlers.types.ts
+++ b/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;