refactor: single base connection, abstract acp further for claude

jonathanlab · jonathanlab · commit 40b8ed433dca · 2026-01-23T14:00:16.000+01:00
diff --git a/packages/agent/src/adapters/acp-connection.ts b/packages/agent/src/adapters/acp-connection.ts
@@ -6,7 +6,7 @@ import {
   createTappedWritableStream,
   type StreamPair,
 } from "@/utils/streams.js";
-import { ClaudeAcpAgent } from "./claude/agent.js";
+import { ClaudeAcpAgent } from "./claude/claude-agent.js";
 
 export type AgentAdapter = "claude";
 
@@ -92,7 +92,7 @@ export function createAcpConnection(
       logger.info("Cleaning up ACP connection");
 
       if (agent) {
-        await agent.closeAllSessions();
+        await agent.closeSession();
       }
 
       // Then close the streams to properly terminate the ACP connection
diff --git a/packages/agent/src/adapters/base-acp-agent.ts b/packages/agent/src/adapters/base-acp-agent.ts
@@ -16,16 +16,18 @@ import type {
   WriteTextFileResponse,
 } from "@agentclientprotocol/sdk";
 import { Logger } from "@/utils/logger.js";
-import type { SessionState } from "./types.js";
 
-export interface BaseSession extends SessionState {
-  abortController: AbortController;
+export interface BaseSession {
+  notificationHistory: SessionNotification[];
+  cancelled: boolean;
   interruptReason?: string;
+  abortController: AbortController;
 }
 
 export abstract class BaseAcpAgent implements Agent {
   abstract readonly adapterName: string;
-  protected sessions: { [key: string]: BaseSession } = {};
+  protected session: BaseSession | null = null;
+  protected sessionId: string | null = null;
   client: AgentSideConnection;
   protected logger: Logger;
   protected fileContentCache: { [key: string]: string } = {};
@@ -38,26 +40,49 @@ export abstract class BaseAcpAgent implements Agent {
   abstract initialize(request: InitializeRequest): Promise<InitializeResponse>;
   abstract newSession(params: NewSessionRequest): Promise<NewSessionResponse>;
   abstract prompt(params: PromptRequest): Promise<PromptResponse>;
-  abstract cancel(params: CancelNotification): Promise<void>;
+  protected abstract interruptSession(): Promise<void>;
+
+  async cancel(params: CancelNotification): Promise<void> {
+    if (this.sessionId !== params.sessionId || !this.session) {
+      throw new Error("Session not found");
+    }
+    this.session.cancelled = true;
+    const meta = params._meta as { interruptReason?: string } | undefined;
+    if (meta?.interruptReason) {
+      this.session.interruptReason = meta.interruptReason;
+    }
+    await this.interruptSession();
+  }
 
-  async closeAllSessions(): Promise<void> {
-    for (const [sessionId, session] of Object.entries(this.sessions)) {
-      try {
-        await this.cancel({ sessionId });
-        session.abortController.abort();
-        this.logger.info("Closed session", { sessionId });
-      } catch (err) {
-        this.logger.warn("Failed to close session", { sessionId, error: err });
-      }
+  async closeSession(): Promise<void> {
+    if (!this.session || !this.sessionId) {
+      return;
+    }
+    try {
+      await this.cancel({ sessionId: this.sessionId });
+      this.session.abortController.abort();
+      this.logger.info("Closed session", { sessionId: this.sessionId });
+    } catch (err) {
+      this.logger.warn("Failed to close session", {
+        sessionId: this.sessionId,
+        error: err,
+      });
     }
-    this.sessions = {};
+    this.session = null;
+    this.sessionId = null;
+  }
+
+  hasSession(sessionId: string): boolean {
+    return this.sessionId === sessionId && this.session !== null;
   }
 
   appendNotification(
     sessionId: string,
     notification: SessionNotification,
   ): void {
-    this.sessions[sessionId]?.notificationHistory.push(notification);
+    if (this.sessionId === sessionId && this.session) {
+      this.session.notificationHistory.push(notification);
+    }
   }
 
   async readTextFile(
diff --git a/packages/agent/src/adapters/claude/claude-agent.ts b/packages/agent/src/adapters/claude/claude-agent.ts
@@ -5,7 +5,6 @@ import {
   type AgentSideConnection,
   type AuthenticateRequest,
   type AvailableCommand,
-  type CancelNotification,
   type ClientCapabilities,
   type InitializeRequest,
   type InitializeResponse,
@@ -147,7 +146,7 @@ async function getAvailableSlashCommands(
 
 export class ClaudeAcpAgent extends BaseAcpAgent {
   readonly adapterName = "claude";
-  declare sessions: { [key: string]: Session };
+  declare session: Session | null;
   toolUseCache: ToolUseCache;
   backgroundTerminals: { [key: string]: BackgroundTerminal } = {};
   clientCapabilities?: ClientCapabilities;
@@ -175,7 +174,8 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
       notificationHistory: [],
       abortController,
     };
-    this.sessions[sessionId] = session;
+    this.session = session;
+    this.sessionId = sessionId;
     return session;
   }
 
@@ -239,12 +239,11 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
     }, 0);
   }
 
-  private getSession(sessionId: string): Session {
-    const session = this.sessions[sessionId];
-    if (!session) {
+  private getSession(): Session {
+    if (!this.session) {
       throw new Error("Session not found");
     }
-    return session;
+    return this.session;
   }
 
   private registerPersistence(
@@ -413,7 +412,7 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
   }
 
   async prompt(params: PromptRequest): Promise<PromptResponse> {
-    const session = this.getSession(params.sessionId);
+    const session = this.getSession();
     session.cancelled = false;
     session.interruptReason = undefined;
 
@@ -498,25 +497,22 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
     throw new Error("Session did not end in result");
   }
 
-  async cancel(params: CancelNotification): Promise<void> {
-    const session = this.getSession(params.sessionId);
-    session.cancelled = true;
-    const meta = params._meta as { interruptReason?: string } | undefined;
-    if (meta?.interruptReason) {
-      session.interruptReason = meta.interruptReason;
+  protected async interruptSession(): Promise<void> {
+    if (!this.session) {
+      return;
     }
-    await session.query.interrupt();
+    await this.session.query.interrupt();
   }
 
   async setSessionModel(params: SetSessionModelRequest) {
-    const session = this.getSession(params.sessionId);
+    const session = this.getSession();
     await session.query.setModel(params.modelId);
   }
 
   async setSessionMode(
     params: SetSessionModeRequest,
   ): Promise<SetSessionModeResponse> {
-    const session = this.getSession(params.sessionId);
+    const session = this.getSession();
 
     switch (params.modeId) {
       case "default":
@@ -545,14 +541,14 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
 
   canUseTool(sessionId: string): CanUseTool {
     return async (toolName, toolInput, { suggestions, toolUseID }) => {
-      const session = this.sessions[sessionId];
-      if (!session) {
+      if (this.sessionId !== sessionId || !this.session) {
         return {
           behavior: "deny",
           message: "Session not found",
           interrupt: true,
         };
       }
+      const session = this.session;
 
       const context = {
         session,
@@ -606,7 +602,7 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
     this.logger.info("[RESUME] Resuming session", { params });
     const { sessionId } = params;
 
-    if (this.sessions[sessionId]) {
+    if (this.sessionId === sessionId && this.session) {
       return {};
     }
 
diff --git a/packages/agent/src/adapters/claude/mcp-server.ts b/packages/agent/src/adapters/claude/mcp-server.ts
@@ -7,7 +7,7 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import * as diff from "diff";
 import { z } from "zod";
 import { Logger } from "@/utils/logger.js";
-import type { ClaudeAcpAgent } from "./agent.js";
+import type { ClaudeAcpAgent } from "./claude-agent.js";
 import { extractLinesWithByteLimit, sleep, unreachable } from "./utils.js";
 
 export const SYSTEM_REMINDER = `
@@ -111,7 +111,7 @@ Usage:
       },
       async (input) => {
         try {
-          if (!agent.sessions[sessionId]) {
+          if (!agent.hasSession(sessionId)) {
             return sessionNotFound();
           }
 
@@ -194,7 +194,7 @@ Usage:
       },
       async (input) => {
         try {
-          if (!agent.sessions[sessionId]) {
+          if (!agent.hasSession(sessionId)) {
             return sessionNotFound();
           }
           await agent.writeTextFile({
@@ -251,7 +251,7 @@ Usage:
       },
       async (input) => {
         try {
-          if (!agent.sessions[sessionId]) {
+          if (!agent.hasSession(sessionId)) {
             return sessionNotFound();
           }
 
@@ -334,7 +334,7 @@ Output: Create directory 'foo'`),
         },
       },
       async (input, extra) => {
-        if (!agent.sessions[sessionId]) {
+        if (!agent.hasSession(sessionId)) {
           return sessionNotFound();
         }
 
diff --git a/packages/agent/src/adapters/claude/session-helpers.ts b/packages/agent/src/adapters/claude/session-helpers.ts
@@ -9,7 +9,7 @@ import type {
   PermissionMode,
 } from "@anthropic-ai/claude-agent-sdk";
 import type { Logger } from "@/utils/logger.js";
-import type { ClaudeAcpAgent } from "./agent.js";
+import type { ClaudeAcpAgent } from "./claude-agent.js";
 import { createPostToolUseHook } from "./hooks.js";
 import { createMcpServer, toolNames } from "./mcp-server.js";
 import { clearStatsigCache } from "./plan-utils.js";
diff --git a/packages/agent/src/adapters/claude/types.ts b/packages/agent/src/adapters/claude/types.ts
@@ -1,5 +1,4 @@
 import type {
-  SessionNotification,
   TerminalHandle,
   TerminalOutputResponse,
 } from "@agentclientprotocol/sdk";
@@ -10,18 +9,15 @@ import type {
   SDKUserMessage,
 } from "@anthropic-ai/claude-agent-sdk";
 import type { Pushable } from "@/utils/streams.js";
+import type { BaseSession } from "../base-acp-agent.js";
 
-export type Session = {
+export type Session = BaseSession & {
   query: Query;
   input: Pushable<SDKUserMessage>;
-  cancelled: boolean;
   permissionMode: PermissionMode;
-  notificationHistory: SessionNotification[];
   sdkSessionId?: string;
   lastPlanFilePath?: string;
   lastPlanContent?: string;
-  abortController: AbortController;
-  interruptReason?: string;
 };
 
 export type ToolUseCache = {
diff --git a/packages/agent/src/adapters/types.ts b/packages/agent/src/adapters/types.ts
