wave-code 0.19.3 → 0.19.5

package/dist/stdio/agentBridge.js

@@ -2,105 +2,124 @@
  * AgentBridge — wraps the SDK Agent and translates between the JSON-RPC-like
  * stdio protocol and Agent method calls / callbacks.
  *
+ * Multi-tenant: maintains a Map<sessionId, {agent, storedConfig}> so a single
+ * `wave --stdio` process can host multiple sessions. Session-scoped requests
+ * carry `sessionId` on the JSON-RPC envelope for routing; global requests
+ * (listSessions/searchFiles/auth/plugins) don't require it but may use it
+ * for workdir fallback.
+ *
  * Responsibilities:
- * - Route incoming requests to the appropriate Agent method
- * - Translate AgentCallbacks into outgoing notifications
+ * - Route incoming requests to the appropriate Agent method (by sessionId)
+ * - Translate AgentCallbacks into outgoing notifications (with sessionId)
  * - Implement the canUseTool permission flow over the stdio protocol
  * - Handle config updates by destroying and recreating the Agent
  */
 import { Agent, listSessions, searchFiles, PromptHistoryManager, AuthService, PluginCore, } from "wave-agent-sdk";
 import { INTERNAL_ERROR as PROTOCOL_INTERNAL_ERROR, METHOD_NOT_FOUND as PROTOCOL_METHOD_NOT_FOUND, } from "./protocol.js";
+import { readFileSync } from "fs";
+import path from "path";
+import { fileURLToPath } from "url";
+const __dirname = path.dirname(fileURLToPath(import.meta.url));
+const CLI_VERSION = (() => {
+    try {
+        const pkg = JSON.parse(readFileSync(path.resolve(__dirname, "../../package.json"), "utf-8"));
+        return pkg.version ?? "";
+    }
+    catch {
+        return "";
+    }
+})();
 export class AgentBridge {
     constructor(options) {
+        this.sessions = new Map();
         this.pendingPermissions = new Map();
         this.permissionCounter = 0;
-        this.storedConfig = {};
         this.emit = options.emit;
     }
     // ── Public API ────────────────────────────────────────────────
-    async handleRequest(method, params) {
+    async handleRequest(method, params, sessionId) {
         const p = (params ?? {});
         switch (method) {
             // ── Lifecycle ──
             case "initialize":
                 return this.initialize(p);
             case "destroy":
-                return this.destroy();
+                return this.destroy(sessionId);
             case "restoreSession":
-                return this.restoreSession(p.sessionId);
+                return this.restoreSession(p.sessionId, sessionId);
             case "listSessions":
-                return this.listSessions(p.workdir);
+                return this.listSessions(p.workdir, sessionId);
             case "getSessionInfo":
-                return this.getSessionInfo();
+                return this.getSessionInfo(sessionId);
             case "updateConfig":
-                return this.updateConfig(p);
+                return this.updateConfig(p, sessionId);
             // ── Messages ──
             case "sendMessage":
-                return this.sendMessage(p);
+                return this.sendMessage(p, sessionId);
             case "bang":
-                return this.bang(p.command);
+                return this.bang(p.command, sessionId);
             case "abortMessage":
-                return this.abortMessage();
+                return this.abortMessage(sessionId);
             case "clearMessages":
-                return this.clearMessages();
+                return this.clearMessages(sessionId);
             case "rewindToMessage":
-                return this.rewindToMessage(p.messageId);
+                return this.rewindToMessage(p.messageId, sessionId);
             case "deleteQueuedMessage":
-                return this.deleteQueuedMessage(p.index);
+                return this.deleteQueuedMessage(p.index, sessionId);
             case "getMessages":
-                return this.getMessages();
+                return this.getMessages(sessionId);
             case "getFullMessageThread":
-                return this.getFullMessageThread();
+                return this.getFullMessageThread(sessionId);
             // ── Permissions ──
             case "setPermissionMode":
-                return this.setPermissionMode(p.mode);
+                return this.setPermissionMode(p.mode, sessionId);
             case "getPermissionMode":
-                return this.getPermissionMode();
+                return this.getPermissionMode(sessionId);
             // ── MCP ──
             case "getMcpServers":
-                return this.getMcpServers();
+                return this.getMcpServers(sessionId);
             case "connectMcpServer":
-                return this.connectMcpServer(p.serverName);
+                return this.connectMcpServer(p.serverName, sessionId);
             case "disconnectMcpServer":
-                return this.disconnectMcpServer(p.serverName);
+                return this.disconnectMcpServer(p.serverName, sessionId);
             // ── Commands ──
             case "getSlashCommands":
-                return this.getSlashCommands();
-            // ── File / History ──
+                return this.getSlashCommands(sessionId);
+            // ── File / History (global — no session required) ──
             case "searchFiles":
-                return this.searchFiles(p);
+                return this.searchFiles(p, sessionId);
             case "getPromptHistory":
-                return this.getPromptHistory(p.workdir);
+                return this.getPromptHistory(p.workdir, sessionId);
             case "searchPromptHistory":
-                return this.searchPromptHistory(p.query, p.workdir);
-            // ── Auth ──
+                return this.searchPromptHistory(p.query, p.workdir, sessionId);
+            // ── Auth (global — no session required) ──
             case "getAuthStatus":
                 return this.getAuthStatus();
             case "login":
                 return this.login(p.serverUrl);
             case "logout":
                 return this.logout();
-            // ── Plugins ──
+            // ── Plugins (global — no session required) ──
             case "listPlugins":
-                return this.listPlugins(p.workdir);
+                return this.listPlugins(p.workdir, sessionId);
             case "installPlugin":
-                return this.installPlugin(p.pluginId, p.scope, p.workdir);
+                return this.installPlugin(p.pluginId, p.scope, p.workdir, sessionId);
             case "uninstallPlugin":
-                return this.uninstallPlugin(p.pluginId, p.workdir);
+                return this.uninstallPlugin(p.pluginId, p.workdir, sessionId);
             case "enablePlugin":
-                return this.enablePlugin(p.pluginId, p.scope, p.workdir);
+                return this.enablePlugin(p.pluginId, p.scope, p.workdir, sessionId);
             case "disablePlugin":
-                return this.disablePlugin(p.pluginId, p.scope, p.workdir);
+                return this.disablePlugin(p.pluginId, p.scope, p.workdir, sessionId);
             case "updatePlugin":
-                return this.updatePlugin(p.pluginId, p.workdir);
+                return this.updatePlugin(p.pluginId, p.workdir, sessionId);
             case "listMarketplaces":
-                return this.listMarketplaces(p.workdir);
+                return this.listMarketplaces(p.workdir, sessionId);
             case "addMarketplace":
-                return this.addMarketplace(p.input, p.scope, p.workdir);
+                return this.addMarketplace(p.input, p.scope, p.workdir, sessionId);
             case "removeMarketplace":
-                return this.removeMarketplace(p.name, p.scope, p.workdir);
+                return this.removeMarketplace(p.name, p.scope, p.workdir, sessionId);
             case "updateMarketplace":
-                return this.updateMarketplace(p.name, p.workdir);
+                return this.updateMarketplace(p.name, p.workdir, sessionId);
             default:
                 throw new RpcError(PROTOCOL_METHOD_NOT_FOUND, `Method not found: ${method}`);
         }
@@ -108,6 +127,7 @@ export class AgentBridge {
     handleNotification(method, params) {
         if (method === "permissionResponse") {
             const p = params;
+            // requestId is process-level unique; lookup doesn't need sessionId
             const resolve = this.pendingPermissions.get(p.requestId);
             if (resolve) {
                 this.pendingPermissions.delete(p.requestId);
@@ -117,9 +137,8 @@ export class AgentBridge {
     }
     // ── Lifecycle ─────────────────────────────────────────────────
     async initialize(params) {
-        // Merge with stored config (CLI defaults can be overridden by client)
-        this.storedConfig = { ...this.storedConfig, ...params };
-        const callbacks = this.createCallbacks();
+        const ctx = {};
+        const callbacks = this.createCallbacks(ctx);
         const options = {
             callbacks,
             workdir: params.workdir,
@@ -136,168 +155,207 @@ export class AgentBridge {
             disallowedTools: params.disallowedTools,
             plugins: params.pluginDirs?.map((path) => ({ type: "local", path })),
             mcpServers: params.mcpServers,
-            canUseTool: (context) => this.canUseTool(context),
+            canUseTool: (context) => this.canUseTool(context, ctx),
         };
-        this.agent = await Agent.create(options);
+        const agent = await Agent.create(options);
+        ctx.agent = agent;
+        ctx.registeredSessionId = agent.sessionId;
+        this.sessions.set(agent.sessionId, {
+            agent,
+            storedConfig: { ...params },
+        });
+        if (params.clientVersion) {
+            console.debug(`[agentBridge] clientVersion=${params.clientVersion} serverVersion=${CLI_VERSION}`);
+        }
         return {
-            sessionId: this.agent.sessionId,
-            workingDirectory: this.agent.workingDirectory,
-            permissionMode: this.agent.getPermissionMode(),
-            latestTotalTokens: this.agent.latestTotalTokens,
+            sessionId: agent.sessionId,
+            workingDirectory: agent.workingDirectory,
+            permissionMode: agent.getPermissionMode(),
+            latestTotalTokens: agent.latestTotalTokens,
+            serverVersion: CLI_VERSION,
         };
     }
-    async destroy() {
-        if (this.agent) {
-            await this.agent.destroy();
-            this.agent = undefined;
+    async destroy(sessionId) {
+        if (sessionId) {
+            const entry = this.sessions.get(sessionId);
+            if (entry) {
+                await entry.agent.destroy();
+                this.sessions.delete(sessionId);
+            }
         }
         return null;
     }
-    async restoreSession(sessionId) {
-        this.requireAgent();
-        await this.agent.restoreSession(sessionId);
+    async restoreSession(restoreId, sessionId) {
+        const entry = this.requireSession(sessionId);
+        await entry.agent.restoreSession(restoreId);
         return null;
     }
-    async listSessions(workdir) {
-        const sessions = await listSessions(workdir || this.agent?.workingDirectory || process.cwd());
+    async listSessions(workdir, sessionId) {
+        const sessions = await listSessions(workdir || this.getSessionWorkdir(sessionId) || process.cwd());
         return { sessions };
     }
-    getSessionInfo() {
-        this.requireAgent();
+    getSessionInfo(sessionId) {
+        const entry = this.requireSession(sessionId);
         return {
-            sessionId: this.agent.sessionId,
-            workingDirectory: this.agent.workingDirectory,
-            latestTotalTokens: this.agent.latestTotalTokens,
-            permissionMode: this.agent.getPermissionMode(),
-            availableTools: this.agent.getAvailableToolNames(),
+            sessionId: entry.agent.sessionId,
+            workingDirectory: entry.agent.workingDirectory,
+            latestTotalTokens: entry.agent.latestTotalTokens,
+            permissionMode: entry.agent.getPermissionMode(),
+            availableTools: entry.agent.getAvailableToolNames(),
         };
     }
-    async updateConfig(params) {
-        this.requireAgent();
-        const currentSessionId = this.agent.sessionId;
+    async updateConfig(params, sessionId) {
+        const entry = this.requireSession(sessionId);
+        const currentSessionId = entry.agent.sessionId;
         // Merge new config into stored config
-        this.storedConfig = { ...this.storedConfig, ...params };
-        // Destroy and recreate
-        await this.agent.destroy();
-        this.agent = undefined;
-        await this.initialize({
-            ...this.storedConfig,
+        entry.storedConfig = { ...entry.storedConfig, ...params };
+        // Destroy and recreate within the same session slot
+        await entry.agent.destroy();
+        this.sessions.delete(currentSessionId);
+        const ctx = {};
+        const callbacks = this.createCallbacks(ctx);
+        const options = {
+            callbacks,
+            workdir: entry.storedConfig.workdir,
             restoreSessionId: currentSessionId,
+            apiKey: entry.storedConfig.apiKey,
+            baseURL: entry.storedConfig.baseURL,
+            defaultHeaders: entry.storedConfig.defaultHeaders,
+            model: entry.storedConfig.model,
+            fastModel: entry.storedConfig.fastModel,
+            language: entry.storedConfig.language,
+            permissionMode: entry.storedConfig.permissionMode,
+            tools: entry.storedConfig.tools,
+            allowedTools: entry.storedConfig.allowedTools,
+            disallowedTools: entry.storedConfig.disallowedTools,
+            plugins: entry.storedConfig.pluginDirs?.map((p) => ({
+                type: "local",
+                path: p,
+            })),
+            mcpServers: entry.storedConfig.mcpServers,
+            canUseTool: (context) => this.canUseTool(context, ctx),
+        };
+        const agent = await Agent.create(options);
+        ctx.agent = agent;
+        ctx.registeredSessionId = agent.sessionId;
+        this.sessions.set(agent.sessionId, {
+            agent,
+            storedConfig: { ...entry.storedConfig },
         });
-        return { sessionId: this.agent.sessionId };
+        return { sessionId: agent.sessionId };
     }
     // ── Messages ──────────────────────────────────────────────────
-    async sendMessage(params) {
-        this.requireAgent();
+    async sendMessage(params, sessionId) {
+        const entry = this.requireSession(sessionId);
         if (params.force) {
-            this.agent.abortMessage();
+            entry.agent.abortMessage();
         }
         // Save prompt to history (mirrors VSCE chatSession.ts:236-242)
         try {
-            await PromptHistoryManager.addEntry(params.text, this.agent.sessionId, {}, this.agent.workingDirectory);
+            await PromptHistoryManager.addEntry(params.text, entry.agent.sessionId, {}, entry.agent.workingDirectory);
         }
         catch {
             // Best-effort; don't block message sending on history save failure
         }
-        await this.agent.sendMessage(params.text, params.images);
+        await entry.agent.sendMessage(params.text, params.images);
         return null;
     }
-    async bang(command) {
-        this.requireAgent();
-        await this.agent.bang(command);
+    async bang(command, sessionId) {
+        const entry = this.requireSession(sessionId);
+        await entry.agent.bang(command);
         return null;
     }
-    async abortMessage() {
-        this.requireAgent();
-        this.agent.abortMessage();
+    async abortMessage(sessionId) {
+        const entry = this.requireSession(sessionId);
+        entry.agent.abortMessage();
         return null;
     }
-    async clearMessages() {
-        this.requireAgent();
-        this.agent.clearMessages();
+    async clearMessages(sessionId) {
+        const entry = this.requireSession(sessionId);
+        entry.agent.clearMessages();
         return null;
     }
-    async rewindToMessage(messageId) {
-        this.requireAgent();
-        const { messages } = await this.agent.getFullMessageThread();
+    async rewindToMessage(messageId, sessionId) {
+        const entry = this.requireSession(sessionId);
+        const { messages } = await entry.agent.getFullMessageThread();
         const index = messages.findIndex((m) => m.id === messageId);
         if (index === -1) {
             throw new RpcError(PROTOCOL_INTERNAL_ERROR, `Message not found: ${messageId}`);
         }
         const message = messages[index];
         const textBlock = message.blocks.find((b) => b.type === "text");
-        await this.agent.truncateHistory(index);
+        await entry.agent.truncateHistory(index);
         return { inputContent: textBlock?.content || "" };
     }
-    deleteQueuedMessage(index) {
-        this.requireAgent();
-        this.agent.removeQueuedMessage(index);
+    deleteQueuedMessage(index, sessionId) {
+        const entry = this.requireSession(sessionId);
+        entry.agent.removeQueuedMessage(index);
         return null;
     }
-    getMessages() {
-        this.requireAgent();
-        return { messages: this.agent.messages };
+    getMessages(sessionId) {
+        const entry = this.requireSession(sessionId);
+        return { messages: entry.agent.messages };
     }
-    async getFullMessageThread() {
-        this.requireAgent();
-        return this.agent.getFullMessageThread();
+    async getFullMessageThread(sessionId) {
+        const entry = this.requireSession(sessionId);
+        return entry.agent.getFullMessageThread();
     }
     // ── Permissions ───────────────────────────────────────────────
-    async setPermissionMode(mode) {
-        this.requireAgent();
-        await this.agent.setPermissionMode(mode);
+    async setPermissionMode(mode, sessionId) {
+        const entry = this.requireSession(sessionId);
+        await entry.agent.setPermissionMode(mode);
         return null;
     }
-    getPermissionMode() {
-        this.requireAgent();
-        return { mode: this.agent.getPermissionMode() };
+    getPermissionMode(sessionId) {
+        const entry = this.requireSession(sessionId);
+        return { mode: entry.agent.getPermissionMode() };
     }
     // ── MCP ───────────────────────────────────────────────────────
-    getMcpServers() {
-        this.requireAgent();
-        return { servers: this.agent.getMcpServers() };
+    getMcpServers(sessionId) {
+        const entry = this.requireSession(sessionId);
+        return { servers: entry.agent.getMcpServers() };
     }
-    async connectMcpServer(serverName) {
-        this.requireAgent();
-        const success = await this.agent.connectMcpServer(serverName);
+    async connectMcpServer(serverName, sessionId) {
+        const entry = this.requireSession(sessionId);
+        const success = await entry.agent.connectMcpServer(serverName);
         return { success };
     }
-    async disconnectMcpServer(serverName) {
-        this.requireAgent();
-        const success = await this.agent.disconnectMcpServer(serverName);
+    async disconnectMcpServer(serverName, sessionId) {
+        const entry = this.requireSession(sessionId);
+        const success = await entry.agent.disconnectMcpServer(serverName);
         return { success };
     }
     // ── Commands ──────────────────────────────────────────────────
-    getSlashCommands() {
-        this.requireAgent();
-        return { commands: this.agent.getSlashCommands() };
+    getSlashCommands(sessionId) {
+        const entry = this.requireSession(sessionId);
+        return { commands: entry.agent.getSlashCommands() };
     }
-    // ── File / History ────────────────────────────────────────────
-    async searchFiles(params) {
+    // ── File / History (global) ───────────────────────────────────
+    async searchFiles(params, sessionId) {
         const files = await searchFiles(params.query, {
             maxResults: params.maxResults,
-            workingDirectory: params.workdir || this.agent?.workingDirectory || process.cwd(),
+            workingDirectory: params.workdir || this.getSessionWorkdir(sessionId) || process.cwd(),
         });
         return { files };
     }
-    async getPromptHistory(workdir) {
+    async getPromptHistory(workdir, sessionId) {
         const history = await PromptHistoryManager.getHistory({
-            workdir: workdir || this.agent?.workingDirectory,
+            workdir: workdir || this.getSessionWorkdir(sessionId),
         });
         return { history };
     }
-    async searchPromptHistory(query, workdir) {
+    async searchPromptHistory(query, workdir, sessionId) {
         const history = await PromptHistoryManager.searchHistory(query, {
-            workdir: workdir || this.agent?.workingDirectory,
+            workdir: workdir || this.getSessionWorkdir(sessionId),
         });
         return { history };
     }
     // ── canUseTool flow ───────────────────────────────────────────
-    canUseTool(context) {
+    canUseTool(context, ctx) {
         const requestId = `perm_${++this.permissionCounter}`;
         return new Promise((resolve) => {
             this.pendingPermissions.set(requestId, resolve);
-            this.emit("permissionRequest", { requestId, context });
+            this.emit("permissionRequest", { requestId, context }, ctx.registeredSessionId);
             // 5-minute timeout → auto-deny
             setTimeout(() => {
                 if (this.pendingPermissions.has(requestId)) {
@@ -310,12 +368,13 @@ export class AgentBridge {
             }, 5 * 60 * 1000);
         });
     }
-    // ── Auth ─────────────────────────────────────────────────────
+    // ── Auth (global) ────────────────────────────────────────────
     async getAuthStatus() {
         const authService = AuthService.getInstance();
         return {
             isAuthenticated: authService.isSSOAuthenticated(),
             user: authService.getAuthUser(),
+            serverUrl: authService.getServerUrl(),
         };
     }
     async login(serverUrl) {
@@ -333,17 +392,17 @@ export class AgentBridge {
         await authService.clearAuth();
         return null;
     }
-    // ── Plugins ──────────────────────────────────────────────────
-    getPluginCore(workdir) {
-        const resolvedWorkdir = workdir || this.agent?.workingDirectory || process.cwd();
+    // ── Plugins (global) ─────────────────────────────────────────
+    getPluginCore(workdir, sessionId) {
+        const resolvedWorkdir = workdir || this.getSessionWorkdir(sessionId) || process.cwd();
         if (!this.pluginCore || this.pluginCoreWorkdir !== resolvedWorkdir) {
             this.pluginCore = new PluginCore(resolvedWorkdir);
             this.pluginCoreWorkdir = resolvedWorkdir;
         }
         return this.pluginCore;
     }
-    async listPlugins(workdir) {
-        const core = this.getPluginCore(workdir);
+    async listPlugins(workdir, sessionId) {
+        const core = this.getPluginCore(workdir, sessionId);
         const { plugins, mergedEnabled } = await core.listPlugins();
         return {
             plugins: plugins.map((p) => {
@@ -361,117 +420,135 @@ export class AgentBridge {
             }),
         };
     }
-    async installPlugin(pluginId, scope, workdir) {
-        return this.getPluginCore(workdir).installPlugin(pluginId, scope);
+    async installPlugin(pluginId, scope, workdir, sessionId) {
+        return this.getPluginCore(workdir, sessionId).installPlugin(pluginId, scope);
     }
-    async uninstallPlugin(pluginId, workdir) {
-        await this.getPluginCore(workdir).uninstallPlugin(pluginId);
+    async uninstallPlugin(pluginId, workdir, sessionId) {
+        await this.getPluginCore(workdir, sessionId).uninstallPlugin(pluginId);
         return null;
     }
-    async enablePlugin(pluginId, scope, workdir) {
-        return this.getPluginCore(workdir).enablePlugin(pluginId, scope);
+    async enablePlugin(pluginId, scope, workdir, sessionId) {
+        return this.getPluginCore(workdir, sessionId).enablePlugin(pluginId, scope);
     }
-    async disablePlugin(pluginId, scope, workdir) {
-        return this.getPluginCore(workdir).disablePlugin(pluginId, scope);
+    async disablePlugin(pluginId, scope, workdir, sessionId) {
+        return this.getPluginCore(workdir, sessionId).disablePlugin(pluginId, scope);
     }
-    async updatePlugin(pluginId, workdir) {
-        return this.getPluginCore(workdir).updatePlugin(pluginId);
+    async updatePlugin(pluginId, workdir, sessionId) {
+        return this.getPluginCore(workdir, sessionId).updatePlugin(pluginId);
     }
-    async listMarketplaces(workdir) {
-        return this.getPluginCore(workdir).listMarketplaces();
+    async listMarketplaces(workdir, sessionId) {
+        return this.getPluginCore(workdir, sessionId).listMarketplaces();
     }
-    async addMarketplace(input, scope, workdir) {
-        return this.getPluginCore(workdir).addMarketplace(input, scope);
+    async addMarketplace(input, scope, workdir, sessionId) {
+        return this.getPluginCore(workdir, sessionId).addMarketplace(input, scope);
     }
-    async removeMarketplace(name, scope, workdir) {
-        await this.getPluginCore(workdir).removeMarketplace(name, scope);
+    async removeMarketplace(name, scope, workdir, sessionId) {
+        await this.getPluginCore(workdir, sessionId).removeMarketplace(name, scope);
         return null;
     }
-    async updateMarketplace(name, workdir) {
-        await this.getPluginCore(workdir).updateMarketplace(name);
+    async updateMarketplace(name, workdir, sessionId) {
+        await this.getPluginCore(workdir, sessionId).updateMarketplace(name);
         return null;
     }
     // ── Callbacks → Notifications ─────────────────────────────────
-    createCallbacks() {
+    createCallbacks(ctx) {
         return {
             onMessagesChange: (messages) => {
-                this.emit("messagesChange", { messages });
+                this.emit("messagesChange", { messages }, ctx.registeredSessionId);
             },
             onUserMessageAdded: () => {
-                const msg = this.findLastUserMessage();
+                const msg = this.findLastUserMessage(ctx.agent);
                 if (msg)
-                    this.emit("userMessageAdded", { message: msg });
+                    this.emit("userMessageAdded", { message: msg }, ctx.registeredSessionId);
             },
             onAssistantMessageAdded: (messageId) => {
-                const msg = this.agent?.messages.find((m) => m.id === messageId);
+                const msg = ctx.agent?.messages.find((m) => m.id === messageId);
                 if (msg)
-                    this.emit("assistantMessageAdded", { message: msg });
+                    this.emit("assistantMessageAdded", { message: msg }, ctx.registeredSessionId);
             },
             onAssistantContentUpdated: (params) => {
-                this.emit("assistantContentUpdated", params);
+                this.emit("assistantContentUpdated", params, ctx.registeredSessionId);
             },
             onAssistantReasoningUpdated: (params) => {
-                this.emit("assistantReasoningUpdated", params);
+                this.emit("assistantReasoningUpdated", params, ctx.registeredSessionId);
             },
             onToolBlockUpdated: (params) => {
-                this.emit("toolBlockUpdated", params);
+                this.emit("toolBlockUpdated", params, ctx.registeredSessionId);
             },
             onErrorBlockAdded: (error) => {
-                this.emit("errorBlockAdded", { error });
+                this.emit("errorBlockAdded", { error }, ctx.registeredSessionId);
             },
             onLoadingChange: (loading) => {
                 this.emit("loadingChange", {
                     loading,
-                    latestTotalTokens: this.agent?.latestTotalTokens,
-                });
+                    latestTotalTokens: ctx.agent?.latestTotalTokens,
+                }, ctx.registeredSessionId);
             },
             onCommandRunningChange: (running) => {
-                this.emit("commandRunningChange", { running });
+                this.emit("commandRunningChange", { running }, ctx.registeredSessionId);
             },
             onQueuedMessagesChange: (messages) => {
-                this.emit("queuedMessagesChange", { messages });
+                this.emit("queuedMessagesChange", { messages }, ctx.registeredSessionId);
             },
             onTasksChange: (tasks) => {
-                this.emit("tasksChange", { tasks });
+                this.emit("tasksChange", { tasks }, ctx.registeredSessionId);
             },
-            onSessionIdChange: (sessionId) => {
-                this.emit("sessionIdChange", { sessionId });
+            onSessionIdChange: (newSessionId) => {
+                const oldSessionId = ctx.registeredSessionId;
+                // Emit with the OLD sessionId so the client's router can deliver it
+                this.emit("sessionIdChange", { sessionId: newSessionId }, oldSessionId);
+                // Update the sessions Map key atomically (single-threaded, no await)
+                if (oldSessionId && oldSessionId !== newSessionId) {
+                    const entry = this.sessions.get(oldSessionId);
+                    if (entry) {
+                        this.sessions.delete(oldSessionId);
+                        this.sessions.set(newSessionId, entry);
+                    }
+                }
+                ctx.registeredSessionId = newSessionId;
             },
             onPermissionModeChange: (mode) => {
-                this.emit("permissionModeChange", { mode });
+                this.emit("permissionModeChange", { mode }, ctx.registeredSessionId);
             },
             onMcpServersChange: (servers) => {
-                this.emit("mcpServersChange", { servers });
+                this.emit("mcpServersChange", { servers }, ctx.registeredSessionId);
             },
             onAddBangMessage: () => {
-                this.emit("bangMessageAdded", {});
+                this.emit("bangMessageAdded", {}, ctx.registeredSessionId);
             },
             onUpdateBangMessage: () => {
-                this.emit("bangMessageUpdated", {});
+                this.emit("bangMessageUpdated", {}, ctx.registeredSessionId);
             },
             onCompleteBangMessage: () => {
-                this.emit("bangMessageCompleted", {});
+                this.emit("bangMessageCompleted", {}, ctx.registeredSessionId);
             },
             onNotificationMessageAdded: (params) => {
-                const msg = this.agent?.messages.find((m) => m.role === "user" &&
+                const msg = ctx.agent?.messages.find((m) => m.role === "user" &&
                     m.blocks.some((b) => b.type === "task_notification" &&
                         b.taskId === params.taskId));
-                this.emit("notificationMessageAdded", {
-                    ...params,
-                    message: msg,
-                });
+                this.emit("notificationMessageAdded", { ...params, message: msg }, ctx.registeredSessionId);
             },
         };
     }
-    findLastUserMessage() {
-        const userMessages = this.agent?.messages.filter((m) => m.role === "user") ?? [];
+    findLastUserMessage(agent) {
+        const userMessages = agent?.messages.filter((m) => m.role === "user") ?? [];
         return userMessages[userMessages.length - 1];
     }
     // ── Utils ─────────────────────────────────────────────────────
-    requireAgent() {
-        if (!this.agent) {
-            throw new RpcError(PROTOCOL_INTERNAL_ERROR, "Agent not initialized");
+    requireSession(sessionId) {
+        if (!sessionId) {
+            throw new RpcError(PROTOCOL_INTERNAL_ERROR, "sessionId is required for this request");
         }
+        const entry = this.sessions.get(sessionId);
+        if (!entry) {
+            throw new RpcError(PROTOCOL_INTERNAL_ERROR, `Session not found: ${sessionId}`);
+        }
+        return entry;
+    }
+    getSessionWorkdir(sessionId) {
+        if (!sessionId)
+            return undefined;
+        return this.sessions.get(sessionId)?.agent.workingDirectory;
     }
 }
 // ── Error class for protocol errors ─────────────────────────────