fix: Preserve session state on idle-kill instead of destroying it (#1459)

Author: charlesvien

apps/code/src/main/services/agent/service.test.ts

@@ -285,6 +285,7 @@ describe("AgentService", () => {
         lastActivityAt: Date.now(),
         config: {},
         promptPending: false,
+        inFlightMcpToolCalls: new Map(),
         ...overrides,
       });
     }
@@ -375,6 +376,34 @@ describe("AgentService", () => {
       expect(getIdleTimeouts(service).has("run-1")).toBe(true);
     });
 
+    it("reschedules when inFlightMcpToolCalls is non-empty at timeout", () => {
+      const toolCalls = new Map([["tool-1", "some-mcp-tool"]]);
+      injectSession(service, "run-1", { inFlightMcpToolCalls: toolCalls });
+      service.recordActivity("run-1");
+
+      vi.advanceTimersByTime(15 * 60 * 1000);
+
+      expect(service.emit).not.toHaveBeenCalledWith(
+        "session-idle-killed",
+        expect.anything(),
+      );
+      expect(getIdleTimeouts(service).has("run-1")).toBe(true);
+    });
+
+    it("kills session when inFlightMcpToolCalls is empty", () => {
+      injectSession(service, "run-1", {
+        inFlightMcpToolCalls: new Map(),
+      });
+      service.recordActivity("run-1");
+
+      vi.advanceTimersByTime(15 * 60 * 1000);
+
+      expect(service.emit).toHaveBeenCalledWith(
+        "session-idle-killed",
+        expect.objectContaining({ taskRunId: "run-1" }),
+      );
+    });
+
     it("checkIdleDeadlines kills expired sessions on resume", () => {
       injectSession(service, "run-1");
       service.recordActivity("run-1");

apps/code/src/main/services/agent/service.ts

@@ -365,7 +365,7 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
    */
   public hasActiveSessions(): boolean {
     for (const session of this.sessions.values()) {
-      if (session.promptPending) {
+      if (session.promptPending || session.inFlightMcpToolCalls.size > 0) {
         log.info("Active session found", { sessionId: session.taskRunId });
         return true;
       }
@@ -391,7 +391,7 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
   private killIdleSession(taskRunId: string): void {
     const session = this.sessions.get(taskRunId);
     if (!session) return;
-    if (session.promptPending) {
+    if (session.promptPending || session.inFlightMcpToolCalls.size > 0) {
       this.recordActivity(taskRunId);
       return;
     }

apps/code/src/renderer/features/sessions/service/service.ts

@@ -127,7 +127,7 @@ export class SessionService {
         onData: (event: { taskRunId: string }) => {
           const { taskRunId } = event;
           log.info("Session idle-killed by main process", { taskRunId });
-          this.teardownSession(taskRunId);
+          this.handleIdleKill(taskRunId);
         },
         onError: (err: unknown) => {
           log.debug("Idle-killed subscription error", { error: err });
@@ -481,6 +481,24 @@ export class SessionService {
     removePersistedConfigOptions(taskRunId);
   }
 
+  /**
+   * Handle an idle-kill from the main process without destroying session state.
+   * The main process already cleaned up the agent, so we only need to
+   * unsubscribe from the channel and mark the session as errored.
+   * Preserves events, logUrl, configOptions and adapter so that Retry
+   * can reconnect with full context via unstable_resumeSession.
+   */
+  private handleIdleKill(taskRunId: string): void {
+    this.unsubscribeFromChannel(taskRunId);
+    sessionStoreSetters.updateSession(taskRunId, {
+      status: "error",
+      errorMessage:
+        "Session disconnected due to inactivity. Click Retry to reconnect.",
+      isPromptPending: false,
+      promptStartedAt: null,
+    });
+  }
+
   private setErrorSession(
     taskId: string,
     taskRunId: string,