fix(agent-core-v2): rebind stale agent runtime to local on first use after restore

.changeset/fix-stale-runtime-binding-restore.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+Fix restored sessions permanently losing tool and subagent access when their previous runtime no longer exists.

packages/agent-core-v2/src/agent/runtimeBinding/agentRuntime.ts

@@ -4,7 +4,7 @@ import { Emitter, type Event } from '#/_base/event';
 import type { IDisposable } from '#/_base/di/lifecycle';
 import { LifecycleScope } from '#/app/scopes';
 import type { Runtime, RuntimeBinding, RuntimeCapability, RuntimeLease } from '#/runtime/runtime';
-import { runtimeStatusAllows, type RuntimeGenerationSnapshot } from '#/runtime/runtimeRegistry';
+import { RuntimeError, runtimeStatusAllows, type RuntimeGenerationSnapshot } from '#/runtime/runtimeRegistry';
 import {
   IRuntimeResolver,
   IWorkspaceInstanceManager,
@@ -76,7 +76,13 @@ export class AgentRuntimeService implements IAgentRuntimeService {
   }
 
   inspect(): Runtime {
-    return this.resolver.inspect(this.binding.current);
+    try {
+      return this.resolver.inspect(this.binding.current);
+    } catch (error) {
+      if (!(error instanceof RuntimeError) || error.code !== 'runtime.not_found') throw error;
+      this.heal(error);
+      return this.resolver.inspect(this.binding.current);
+    }
   }
 
   isAvailable(required: readonly RuntimeCapability[] = []): boolean {
@@ -89,7 +95,13 @@ export class AgentRuntimeService implements IAgentRuntimeService {
   }
 
   acquire(required: readonly RuntimeCapability[] = []): RuntimeLease {
-    return this.resolver.acquire(this.binding.current, required);
+    try {
+      return this.resolver.acquire(this.binding.current, required);
+    } catch (error) {
+      if (!(error instanceof RuntimeError) || error.code !== 'runtime.not_found') throw error;
+      this.heal(error);
+      return this.resolver.acquire(this.binding.current, required);
+    }
   }
 
   dispose(): void {
@@ -104,6 +116,17 @@ export class AgentRuntimeService implements IAgentRuntimeService {
     this.changeEmitter.fire();
   }
 
+  private heal(error: RuntimeError): void {
+    const current = this.binding.current;
+    if (current.runtimeId === 'local') throw error;
+    try {
+      this.binding.set({ workspaceId: current.workspaceId, runtimeId: 'local' });
+    } catch (fallbackError) {
+      error.cause = fallbackError;
+      throw error;
+    }
+  }
+
   private bindRegistry(): void {
     this.registrySubscription?.dispose();
     const binding = this.binding.current;

packages/agent-core-v2/test/agent/runtimeBinding/runtimeBindingService.test.ts

@@ -1,7 +1,9 @@
 import { describe, expect, it } from 'vitest';
 
 import { Emitter } from '#/_base/event';
+import type { Event2 } from '#/app/event/event2';
 import { AgentRuntimeService, snapshotAgentRuntimeBinding } from '#/agent/runtimeBinding/agentRuntime';
+import { RuntimeSetBinding, runtimeBindingKey } from '#/agent/runtimeBinding/runtimeBindingOps';
 import { AgentRuntimeBindingService, agentRuntimeBindingKey } from '#/agent/runtimeBinding/runtimeBindingService';
 import { AgentStateService } from '#/agent/state/agentStateService';
 import { FakeRuntime } from '#/runtime/fakeRuntime';
@@ -53,10 +55,25 @@ function setup() {
     sessionScope: 'sessions/session',
     cwd: '/workspace',
   });
+  let restoreHook: ((context: undefined, next: () => Promise<void>) => Promise<void>) | undefined;
+  const dispatched: Event2[] = [];
   const dispatcher = {
     _serviceBrand: undefined,
-    dispatch: () => Promise.resolve(),
-    hooks: { onDidRestore: { register: () => ({ dispose: () => {} }) } },
+    dispatch: (event: Event2) => {
+      dispatched.push(event);
+      return Promise.resolve();
+    },
+    hooks: {
+      onDidRestore: {
+        register: (
+          _id: string,
+          hook: (context: undefined, next: () => Promise<void>) => Promise<void>,
+        ) => {
+          restoreHook = hook;
+          return { dispose: () => {} };
+        },
+      },
+    },
   } as unknown as IEventDispatcher;
   const binding = new AgentRuntimeBindingService(
     {
@@ -77,6 +94,11 @@ function setup() {
     onDidChange: workspaceChanges.event,
     get: () => ({ runtimes: registry }),
   } as unknown as IWorkspaceInstanceManager;
+  const restore = async (replayed: RuntimeBinding): Promise<void> => {
+    state.set(runtimeBindingKey, replayed);
+    if (restoreHook === undefined) throw new Error('restore hook was not registered');
+    await restoreHook(undefined, async () => {});
+  };
   return {
     registry,
     resolver,
@@ -85,6 +107,8 @@ function setup() {
     local,
     remote,
     localRegistration,
+    dispatched,
+    restore,
     workspaceChanges,
     agentRuntime: new AgentRuntimeService(binding, resolver, workspaces),
   };
@@ -230,4 +254,74 @@ describe('AgentRuntimeBindingService', () => {
     local.setStatus('ready');
     expect(changes).toHaveLength(1);
   });
+
+  it('heals a stale restored binding through acquire and persists the fallback', async () => {
+    const { binding, dispatched, restore, agentRuntime } = setup();
+    await restore({ workspaceId: 'workspace', runtimeId: 'acp:session_gone' });
+
+    const firstLease = agentRuntime.acquire();
+    expect(firstLease.runtime.identity.runtimeId).toBe('local');
+    expect(binding.current).toEqual({ workspaceId: 'workspace', runtimeId: 'local' });
+    const secondLease = agentRuntime.acquire();
+    expect(secondLease.runtime.identity.runtimeId).toBe('local');
+    expect(dispatched).toEqual([
+      expect.objectContaining({
+        type: RuntimeSetBinding.type,
+        workspaceId: 'workspace',
+        runtimeId: 'local',
+      }),
+    ]);
+    firstLease.dispose();
+    secondLease.dispose();
+  });
+
+  it('heals a stale restored binding through inspect', async () => {
+    const { binding, dispatched, restore, agentRuntime } = setup();
+    await restore({ workspaceId: 'workspace', runtimeId: 'acp:session_gone' });
+
+    expect(agentRuntime.inspect().identity.runtimeId).toBe('local');
+    expect(binding.current).toEqual({ workspaceId: 'workspace', runtimeId: 'local' });
+    expect(dispatched).toEqual([
+      expect.objectContaining({ type: RuntimeSetBinding.type, runtimeId: 'local' }),
+    ]);
+  });
+
+  it('heals a stale restored binding through availability and persists the fallback', async () => {
+    const { binding, dispatched, restore, agentRuntime } = setup();
+    await restore({ workspaceId: 'workspace', runtimeId: 'acp:session_gone' });
+
+    expect(agentRuntime.isAvailable(['fs'])).toBe(true);
+    expect(binding.current).toEqual({ workspaceId: 'workspace', runtimeId: 'local' });
+    expect(dispatched).toEqual([
+      expect.objectContaining({ type: RuntimeSetBinding.type, runtimeId: 'local' }),
+    ]);
+    const lease = agentRuntime.acquire();
+    expect(lease.runtime.identity.runtimeId).toBe('local');
+    lease.dispose();
+  });
+
+  it('preserves the stale binding and original error when the local fallback is missing', async () => {
+    const { binding, localRegistration, restore, agentRuntime } = setup();
+    const stale = { workspaceId: 'workspace', runtimeId: 'acp:session_gone' };
+    await localRegistration.remove();
+    await restore(stale);
+
+    let thrown: unknown;
+    try {
+      agentRuntime.acquire();
+    } catch (error) {
+      thrown = error;
+    }
+    expect(thrown).toEqual(
+      expect.objectContaining<Partial<RuntimeError>>({
+        code: 'runtime.not_found',
+        message: 'runtime acp:session_gone does not exist in workspace workspace',
+        cause: expect.objectContaining<Partial<RuntimeError>>({
+          code: 'runtime.not_found',
+          message: 'runtime local does not exist in workspace workspace',
+        }),
+      }),
+    );
+    expect(binding.current).toEqual(stale);
+  });
 });