test(agent-manager): add session ID verification and improve agent state persistence

2026-07-29 12:01:31 +00:00 · 2025-10-26 08:29:52 +01:00
parent 604e359437
commit 1d8e86985b
4 changed files with 264 additions and 95 deletions
--- a/packages/server/agents.json
+++ b/packages/server/agents.json
@@ -1,14 +1,134 @@
 [
  {
-    "id": "fc818184-5676-4f24-a92a-77e9e15c515d",
-    "title": "Agent fc818184",
+    "id": "299ab613-74d7-478b-a990-f95b044606b4",
+    "title": "Agent 299ab613",
+    "sessionId": "019a1f5d-1184-772c-acb4-0b61ad47bb18",
+    "options": {
+      "type": "claude",
+      "sessionId": null
+    },
+    "createdAt": "2025-10-26T07:13:03.607Z",
+    "lastActivityAt": "2025-10-26T07:13:07.664Z",
+    "cwd": "/var/folders/xl/kkk9drfd3ms_t8x7rmy4z6900000gn/T/acp-test-c5m2nU"
+  },
+  {
+    "id": "c6132199-a93b-428e-ac6b-6e8067855117",
+    "title": "Agent c6132199",
+    "sessionId": "019a1f5d-5a27-712a-a003-c33d85514452",
+    "options": {
+      "type": "claude",
+      "sessionId": null
+    },
+    "createdAt": "2025-10-26T07:13:22.279Z",
+    "lastActivityAt": "2025-10-26T07:13:34.996Z",
+    "cwd": "/var/folders/xl/kkk9drfd3ms_t8x7rmy4z6900000gn/T/acp-test-YqAd2n"
+  },
+  {
+    "id": "01cb507d-5c45-402a-9010-278077d72511",
+    "title": "Agent 01cb507d",
+    "sessionId": "019a1f5e-63ae-7318-a927-ace1f5177b40",
+    "options": {
+      "type": "claude",
+      "sessionId": null
+    },
+    "createdAt": "2025-10-26T07:14:29.299Z",
+    "lastActivityAt": "2025-10-26T07:14:34.243Z",
+    "cwd": "/var/folders/xl/kkk9drfd3ms_t8x7rmy4z6900000gn/T/acp-test-YqAd2n"
+  },
+  {
+    "id": "3647cfd5-7691-469d-a5b4-dfddf4950327",
+    "title": "Agent 3647cfd5",
+    "sessionId": "019a1f5f-f3e1-72b9-bca0-6a58fc8bd52a",
+    "options": {
+      "type": "claude",
+      "sessionId": null
+    },
+    "createdAt": "2025-10-26T07:16:12.630Z",
+    "lastActivityAt": "2025-10-26T07:16:26.477Z",
+    "cwd": "/var/folders/xl/kkk9drfd3ms_t8x7rmy4z6900000gn/T/acp-test-SHg6iy"
+  },
+  {
+    "id": "52623113-d211-4ee0-bdfd-329e8dd8378b",
+    "title": "Agent 52623113",
+    "sessionId": "019a1f60-2ce3-756c-8850-b310052bb04e",
+    "options": {
+      "type": "claude",
+      "sessionId": null
+    },
+    "createdAt": "2025-10-26T07:16:27.236Z",
+    "lastActivityAt": "2025-10-26T07:16:30.640Z",
+    "cwd": "/var/folders/xl/kkk9drfd3ms_t8x7rmy4z6900000gn/T/acp-test-SHg6iy"
+  },
+  {
+    "id": "9c675fcd-50c0-4228-a8d5-544286655713",
+    "title": "Agent 9c675fcd",
+    "sessionId": "019a1f60-b291-739b-8f8c-d4615aaf8e80",
+    "options": {
+      "type": "claude",
+      "sessionId": null
+    },
+    "createdAt": "2025-10-26T07:17:01.370Z",
+    "lastActivityAt": "2025-10-26T07:17:16.318Z",
+    "cwd": "/var/folders/xl/kkk9drfd3ms_t8x7rmy4z6900000gn/T/acp-test-0sCjpk"
+  },
+  {
+    "id": "9b7dc1ab-2465-4772-a821-760dd505b8d7",
+    "title": "Agent 9b7dc1ab",
+    "sessionId": "019a1f60-ef6b-76ab-9bce-e25cd24556cf",
+    "options": {
+      "type": "claude",
+      "sessionId": null
+    },
+    "createdAt": "2025-10-26T07:17:17.000Z",
+    "lastActivityAt": "2025-10-26T07:17:22.332Z",
+    "cwd": "/var/folders/xl/kkk9drfd3ms_t8x7rmy4z6900000gn/T/acp-test-0sCjpk"
+  },
+  {
+    "id": "c83eea1e-f5c8-435a-bc02-716542d60e0e",
+    "title": "Agent c83eea1e",
    "sessionId": null,
    "options": {
      "type": "claude",
      "sessionId": null
    },
-    "createdAt": "2025-10-26T06:59:13.289Z",
-    "lastActivityAt": "2025-10-26T06:59:22.379Z",
-    "cwd": "/var/folders/xl/kkk9drfd3ms_t8x7rmy4z6900000gn/T/acp-test-wm5mun"
+    "createdAt": "2025-10-26T07:28:27.901Z",
+    "lastActivityAt": "2025-10-26T07:28:41.570Z",
+    "cwd": "/var/folders/xl/kkk9drfd3ms_t8x7rmy4z6900000gn/T/acp-test-QcKr0I"
+  },
+  {
+    "id": "354b099b-ffb9-4c42-9b82-2590550e7299",
+    "title": "Agent 354b099b",
+    "sessionId": null,
+    "options": {
+      "type": "claude",
+      "sessionId": null
+    },
+    "createdAt": "2025-10-26T07:28:42.609Z",
+    "lastActivityAt": "2025-10-26T07:28:52.065Z",
+    "cwd": "/var/folders/xl/kkk9drfd3ms_t8x7rmy4z6900000gn/T/acp-test-QcKr0I"
+  },
+  {
+    "id": "28b888fa-7cbc-48a7-973a-964e8a5f29f5",
+    "title": "Agent 28b888fa",
+    "sessionId": "03e4b827-033e-4212-b3ca-cfee515cf1f0",
+    "options": {
+      "type": "claude",
+      "sessionId": "03e4b827-033e-4212-b3ca-cfee515cf1f0"
+    },
+    "createdAt": "2025-10-26T07:29:17.693Z",
+    "lastActivityAt": "2025-10-26T07:29:31.675Z",
+    "cwd": "/var/folders/xl/kkk9drfd3ms_t8x7rmy4z6900000gn/T/acp-test-Gi719a"
+  },
+  {
+    "id": "b535f9b4-230f-40a9-8db8-30b791be1df4",
+    "title": "Agent b535f9b4",
+    "sessionId": "06c94e04-8cc6-4219-8bf5-d0ef7eb9e46a",
+    "options": {
+      "type": "claude",
+      "sessionId": "06c94e04-8cc6-4219-8bf5-d0ef7eb9e46a"
+    },
+    "createdAt": "2025-10-26T07:29:32.156Z",
+    "lastActivityAt": "2025-10-26T07:29:40.533Z",
+    "cwd": "/var/folders/xl/kkk9drfd3ms_t8x7rmy4z6900000gn/T/acp-test-Gi719a"
  }
 ]
--- a/packages/server/src/server/acp/agent-manager.test.ts
+++ b/packages/server/src/server/acp/agent-manager.test.ts
@@ -106,6 +106,12 @@ describe("AgentManager", () => {

      expect(status).toBe("completed");

+      const agentBeforeKill = manager.listAgents().find((a) => a.id === agentId);
+      const claudeSessionId = manager.getClaudeSessionId(agentId);
+      console.log("Before kill - ACP Session ID:", agentBeforeKill?.sessionId);
+      console.log("Before kill - Claude Session ID:", claudeSessionId);
+      console.log("Working directory:", tmpDir);
+
      await manager.killAgent(agentId);

      const newManager = new AgentManager();
@@ -113,6 +119,10 @@ describe("AgentManager", () => {

      const agents = newManager.listAgents();
      const loadedAgent = agents.find((a) => a.id === agentId);
+      const loadedClaudeSessionId = newManager.getClaudeSessionId(agentId);
+
+      console.log("After reload - ACP Session ID:", loadedAgent?.sessionId);
+      console.log("After reload - Claude Session ID:", loadedClaudeSessionId);

      expect(loadedAgent).toBeDefined();
      expect(loadedAgent?.id).toBe(agentId);
--- a/packages/server/src/server/acp/agent-manager.ts
+++ b/packages/server/src/server/acp/agent-manager.ts
@@ -29,7 +29,11 @@ import type {
  AgentRuntime,
  ManagedAgentState,
 } from "./types.js";
-import { AgentPersistence, type AgentOptions } from "./agent-persistence.js";
+import {
+  AgentPersistence,
+  type AgentOptions,
+  type PersistedAgent,
+} from "./agent-persistence.js";

 interface PendingPermission {
  requestId: string;
@@ -56,6 +60,14 @@ interface ManagedAgent {
  state: ManagedAgentState;
 }

+type UpdateAgentCallback = (
+  agentId: string,
+  updateFn: (
+    agent: ManagedAgent
+  ) => boolean | void | Promise<boolean | void>,
+  options?: { sessionId: string | null }
+) => Promise<void>;
+
 /**
 * Get the status from an agent's state
 */
@@ -87,7 +99,7 @@ class ACPClient implements Client {
      agentId: string,
      params: RequestPermissionRequest
    ) => Promise<RequestPermissionResponse>,
-    private persistence: AgentPersistence
+    private updateAgent: UpdateAgentCallback
  ) {}

  async requestPermission(
@@ -100,21 +112,28 @@ class ACPClient implements Client {
  }

  async sessionUpdate(params: SessionNotification): Promise<void> {
+    if (params.update.sessionUpdate === "agent_message_chunk") {
+      console.log(`[Agent ${this.agentId}] Message chunk update:`, JSON.stringify(params, null, 2));
+    }
+
    // Check if this update contains a Claude session ID
    const claudeSessionId = params._meta?.claudeSessionId as string | undefined;
-    if (claudeSessionId && this.persistence) {
-      const persisted = await this.persistence.load();
-      const persistedAgent = persisted.find((a) => a.id === this.agentId);
-      if (persistedAgent && persistedAgent.options.type === "claude") {
-        if (
-          persistedAgent.options.sessionId === null ||
-          persistedAgent.options.sessionId !== claudeSessionId
-        ) {
-          // Update the Claude session ID
-          persistedAgent.options.sessionId = claudeSessionId;
-          await this.persistence.upsert(persistedAgent);
+    if (claudeSessionId) {
+      await this.updateAgent(this.agentId, (agent) => {
+        if (agent.options.type !== "claude") {
+          return false;
        }
-      }
+        if (agent.options.sessionId === claudeSessionId) {
+          return false;
+        }
+
+        agent.options = {
+          ...agent.options,
+          sessionId: claudeSessionId,
+        };
+
+        return true;
+      });
    }
    this.onUpdate(this.agentId, params);
  }
@@ -252,15 +271,7 @@ export class AgentManager {

    this.agents.set(agentId, agent);

-    await this.persistence.upsert({
-      id: agentId,
-      title: agent.title || `Agent ${agentId.slice(0, 8)}`,
-      sessionId: null,
-      options: agent.options,
-      createdAt: createdAt.toISOString(),
-      lastActivityAt: agent.lastActivityAt.toISOString(),
-      cwd: agent.cwd,
-    });
+    await this.updateAgent(agentId, () => undefined);

    this.notifySubscribers(agentId);

@@ -510,17 +521,8 @@ export class AgentManager {
    const runtime = this.getRuntime(agent);

    // Persist current state before killing
-    // Clear sessionId so next start creates a new session
    if (runtime) {
-      await this.persistence.upsert({
-        id: agent.id,
-        title: agent.title || `Agent ${agent.id.slice(0, 8)}`,
-        sessionId: null,
-        options: { ...agent.options, sessionId: null },
-        createdAt: agent.createdAt.toISOString(),
-        lastActivityAt: agent.lastActivityAt.toISOString(),
-        cwd: agent.cwd,
-      });
+      await this.updateAgent(agentId, () => true);
    }

    agent.state = { type: "killed" };
@@ -862,6 +864,63 @@ export class AgentManager {
    return null;
  }

+  private getPersistableSessionId(agent: ManagedAgent): string | null {
+    const { state } = agent;
+
+    switch (state.type) {
+      case "ready":
+      case "processing":
+      case "completed":
+        return state.runtime.sessionId;
+      case "initializing":
+        return state.runtime?.sessionId ?? state.persistedSessionId ?? null;
+      case "failed":
+        return state.runtime?.sessionId ?? null;
+      case "uninitialized":
+        return state.persistedSessionId;
+      case "killed":
+      default:
+        return null;
+    }
+  }
+
+  private serializeAgent(agent: ManagedAgent): PersistedAgent {
+    return {
+      id: agent.id,
+      title: agent.title || `Agent ${agent.id.slice(0, 8)}`,
+      sessionId: this.getPersistableSessionId(agent),
+      options: agent.options,
+      createdAt: agent.createdAt.toISOString(),
+      lastActivityAt: agent.lastActivityAt.toISOString(),
+      cwd: agent.cwd,
+    };
+  }
+
+  private async updateAgent(
+    agentId: string,
+    updateFn: (
+      agent: ManagedAgent
+    ) => boolean | void | Promise<boolean | void>,
+    options?: { sessionId: string | null }
+  ): Promise<void> {
+    const agent = this.agents.get(agentId);
+    if (!agent) {
+      throw new Error(`Agent ${agentId} not found`);
+    }
+
+    const shouldPersist = await updateFn(agent);
+    if (shouldPersist === false) {
+      return;
+    }
+
+    const persistedAgent = this.serializeAgent(agent);
+    if (options) {
+      persistedAgent.sessionId = options.sessionId;
+    }
+
+    await this.persistence.upsert(persistedAgent);
+  }
+
  /**
   * Start runtime for an agent (spawn process, create connection, initialize session)
   */
@@ -905,7 +964,7 @@ export class AgentManager {
      (id, params) => {
        return this.handlePermissionRequest(id, params);
      },
-      this.persistence
+      (id, updateFn, options) => this.updateAgent(id, updateFn, options)
    );
    const connection = new ClientSideConnection(() => client, stream);

@@ -1003,37 +1062,30 @@ export class AgentManager {
      const claudeSessionId = sessionResponse._meta?.claudeSessionId as
        | string
        | undefined;
-      if (
-        claudeSessionId &&
-        agent.options.type === "claude" &&
-        agent.options.sessionId !== claudeSessionId
-      ) {
-        agent.options = {
-          ...agent.options,
-          sessionId: claudeSessionId,
-        };
-      }
-
      console.log(
        `[Agent ${agentId}] Session ${
          mode === "new" ? "created" : "loaded"
        }: ACP=${effectiveSessionId}, Claude=${claudeSessionId || "N/A"}`
      );

-      await this.persistence.upsert({
-        id: agentId,
-        title: agent.title || `Agent ${agentId.slice(0, 8)}`,
-        sessionId: runtime.sessionId,
-        options: agent.options,
-        createdAt: agent.createdAt.toISOString(),
-        lastActivityAt: agent.lastActivityAt.toISOString(),
-        cwd: agent.cwd,
-      });
+      await this.updateAgent(agentId, (managedAgent) => {
+        if (
+          claudeSessionId &&
+          managedAgent.options.type === "claude" &&
+          managedAgent.options.sessionId !== claudeSessionId
+        ) {
+          managedAgent.options = {
+            ...managedAgent.options,
+            sessionId: claudeSessionId,
+          };
+        }

-      agent.state = {
-        type: "ready",
-        runtime,
-      };
+        managedAgent.state = {
+          type: "ready",
+          runtime,
+        };
+        return true;
+      });

      this.notifySubscribers(agentId);

@@ -1223,15 +1275,12 @@ export class AgentManager {
   * Set the title for an agent
   */
  async setAgentTitle(agentId: string, title: string): Promise<void> {
-    const agent = this.agents.get(agentId);
-    if (!agent) {
-      throw new Error(`Agent ${agentId} not found`);
-    }
+    await this.updateAgent(agentId, (managedAgent) => {
+      managedAgent.title = title;
+      return true;
+    });

-    agent.title = title;
    console.log(`[Agent ${agentId}] Title set to: "${title}"`);
-
-    await this.persistence.updateTitle(agentId, title);
  }

  /**
@@ -1245,6 +1294,20 @@ export class AgentManager {
    return agent.title ?? null;
  }

+  /**
+   * Get Claude session ID for an agent
+   */
+  getClaudeSessionId(agentId: string): string | null {
+    const agent = this.agents.get(agentId);
+    if (!agent) {
+      throw new Error(`Agent ${agentId} not found`);
+    }
+    if (agent.options.type === "claude") {
+      return agent.options.sessionId;
+    }
+    return null;
+  }
+
  /**
   * Mark that title generation has been triggered for this agent
   */
@@ -1470,15 +1533,7 @@ export class AgentManager {

          // Persist current state if agent has a session
          if (runtime) {
-            await this.persistence.upsert({
-              id: agent.id,
-              title: agent.title || `Agent ${agent.id.slice(0, 8)}`,
-              sessionId: runtime.sessionId,
-              options: agent.options,
-              createdAt: agent.createdAt.toISOString(),
-              lastActivityAt: agent.lastActivityAt.toISOString(),
-              cwd: agent.cwd,
-            });
+            await this.updateAgent(agent.id, () => undefined);
            console.log(`[Agent ${agent.id}] State persisted`);

            // Send graceful termination signal
--- a/packages/server/src/server/acp/agent-persistence.ts
+++ b/packages/server/src/server/acp/agent-persistence.ts
@@ -132,20 +132,4 @@ export class AgentPersistence {
      console.log(`[AgentPersistence] Removed agent ${agentId}`);
    }
  }
-
-  /**
-   * Update agent title
-   */
-  async updateTitle(agentId: string, title: string): Promise<void> {
-    const agents = await this.load();
-    const agent = agents.find((a) => a.id === agentId);
-
-    if (agent) {
-      agent.title = title;
-      await this.save(agents);
-      console.log(
-        `[AgentPersistence] Updated title for agent ${agentId}: "${title}"`
-      );
-    }
-  }
 }