@symbolica/agentica 0.3.1 → 0.4.0

package/dist/src/client-session-manager/client-session-manager.d.ts

@@ -20,9 +20,9 @@ export interface InvokeAgentParams {
     parentIid?: string;
 }
 export declare class ClientSessionManager {
-    private websocket;
-    private tasks;
-    private sendQueue;
+    private sessionManagers;
+    private smLocks;
+    private uidToSm;
     private uidIidRecvQueue;
     private uidIidException;
     private matchIid;
@@ -46,15 +46,17 @@ export declare class ClientSessionManager {
     setEndpoints(baseHttp: string, baseWs: string): void;
     private registerSession;
     private buildHeaders;
-    private wsBackgroundTaskWriter;
-    private wsBackgroundTaskReader;
-    private _ensureWebSocketConnection;
+    private smWriter;
+    private smReader;
+    private _failSmFutures;
+    private _cleanupSessionManager;
+    private _ensureSmConnection;
     private _initAgentState;
     newAgent(cmar: CreateAgentRequest, parentLogger?: ScopedLogger): Promise<string>;
     invokeAgent(params: InvokeAgentParams, parentLogger?: ScopedLogger): Promise<AgentInvocationHandle>;
     echo(cancelSignal: AbortSignal, uid: string, iid?: string): AsyncGenerator<Chunk, void, unknown>;
     logs(uid: string, iid?: string, params?: Record<string, any>): Promise<Array<Record<string, any>>>;
-    fetchUsage(uid: string, iid: string, lastTotal?: Usage): Promise<{
+    fetchUsage(uid: string, iid: string, _lastTotal?: Usage): Promise<{
         usage: Usage;
         newTotal: Usage;
     }>;
@@ -69,5 +71,5 @@ declare global {
     var __AGENTICA_AGENTIC_MANAGER: ClientSessionManager | undefined;
 }
 export declare const __AGENTIC_MANAGER: ClientSessionManager;
-export declare const __AGENTIC_PROTO_VER = "typescript/0.3.1";
+export declare const __AGENTIC_PROTO_VER = "typescript/0.4.0";
 //# sourceMappingURL=client-session-manager.d.ts.map

package/dist/src/client-session-manager/client-session-manager.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"client-session-manager.d.ts","sourceRoot":"","sources":["../../../src/client-session-manager/client-session-manager.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EACR,kBAAkB,EAKlB,cAAc,EACjB,MAAM,SAAS,CAAC;AAEjB,OAAO,EAAE,KAAK,EAAE,KAAK,EAAY,MAAM,kBAAkB,CAAC;AAE1D,OAAO,EAAE,KAAK,YAAY,EAA+B,MAAM,gBAAgB,CAAC;AA6DhF,MAAM,WAAW,0BAA0B;IAEvC,MAAM,CAAC,EAAE,MAAM,GAAG,IAAI,CAAC;IAEvB,iBAAiB,CAAC,EAAE,OAAO,CAAC;CAC/B;AAGD,MAAM,WAAW,qBAAqB;IAClC,YAAY,EAAE,CAAC,IAAI,EAAE,UAAU,KAAK,OAAO,CAAC,IAAI,CAAC,CAAC;IAClD,YAAY,EAAE,MAAM,OAAO,CAAC,UAAU,CAAC,CAAC;IACxC,SAAS,EAAE,OAAO,CAAC,IAAI,CAAC,CAAC;IACzB,GAAG,EAAE,MAAM,CAAC;CACf;AAGD,MAAM,WAAW,iBAAiB;IAC9B,GAAG,EAAE,MAAM,CAAC;IACZ,iBAAiB,EAAE,UAAU,CAAC;IAC9B,QAAQ,EAAE,MAAM,GAAG,cAAc,CAAC;IAClC,SAAS,EAAE,OAAO,CAAC;IACnB,SAAS,CAAC,EAAE,MAAM,CAAC;IACnB,SAAS,CAAC,EAAE,MAAM,CAAC;CACtB;AAcD,qBAAa,oBAAoB;IAC7B,OAAO,CAAC,SAAS,CAA8B;IAC/C,OAAO,CAAC,KAAK,CAA2C;IACxD,OAAO,CAAC,SAAS,CAA8C;IAC/D,OAAO,CAAC,eAAe,CAA8C;IACrE,OAAO,CAAC,eAAe,CAGrB;IACF,OAAO,CAAC,QAAQ,CAAoF;IACpG,OAAO,CAAC,QAAQ,CAAsB;IACtC,OAAO,CAAC,SAAS,CAAc;IAE/B,OAAO,CAAC,QAAQ,CAAuB;IACvC,OAAO,CAAC,MAAM,CAAuB;IACrC,OAAO,CAAC,MAAM,CAAuB;IACrC,OAAO,CAAC,QAAQ,CAAe;IAC/B,OAAO,CAAC,WAAW,CAAwC;IAC3D,OAAO,CAAC,SAAS,CAAkB;IACnC,OAAO,CAAC,UAAU,CAAkB;IACpC,OAAO,CAAC,MAAM,CAAe;IAC7B,OAAO,CAAC,WAAW,CAAsD;IAGzE,OAAO,CAAC,eAAe,CAAS;IAGhC,OAAO,CAAC,qBAAqB,CAAa;IAC1C,OAAO,CAAC,eAAe,CAA+C;gBAE1D,MAAM,GAAE,0BAA+B;IAyCnD,OAAO,CAAC,aAAa;IAWrB,OAAO,CAAC,eAAe;IAQvB,YAAY,CAAC,QAAQ,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,IAAI;YAsBtC,eAAe;IAwC7B,OAAO,CAAC,YAAY;YAYN,sBAAsB;YA0DtB,sBAAsB;YAsJtB,0BAA0B;IA0FxC,OAAO,CAAC,eAAe;IAejB,QAAQ,CAAC,IAAI,EAAE,kBAAkB,EAAE,YAAY,CAAC,EAAE,YAAY,GAAG,OAAO,CAAC,MAAM,CAAC;IA8FhF,WAAW,CAAC,MAAM,EAAE,iBAAiB,EAAE,YAAY,CAAC,EAAE,YAAY,GAAG,OAAO,CAAC,qBAAqB,CAAC;IAmGlG,IAAI,CAAC,YAAY,EAAE,WAAW,EAAE,GAAG,EAAE,MAAM,EAAE,GAAG,CAAC,EAAE,MAAM,GAAG,cAAc,CAAC,KAAK,EAAE,IAAI,EAAE,OAAO,CAAC;IA+KjG,IAAI,CAAC,GAAG,EAAE,MAAM,EAAE,GAAG,CAAC,EAAE,MAAM,EAAE,MAAM,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,GAAG,CAAC,GAAG,OAAO,CAAC,KAAK,CAAC,MAAM,CAAC,MAAM,EAAE,GAAG,CAAC,CAAC,CAAC;IAuClG,UAAU,CAAC,GAAG,EAAE,MAAM,EAAE,GAAG,EAAE,MAAM,EAAE,SAAS,CAAC,EAAE,KAAK,GAAG,OAAO,CAAC;QAAE,KAAK,EAAE,KAAK,CAAC;QAAC,QAAQ,EAAE,KAAK,CAAA;KAAE,CAAC;IAyBzG,WAAW,CAAC,GAAG,EAAE,MAAM,GAAG,OAAO;IAOjC,YAAY,CAAC,GAAG,EAAE,MAAM,GAAG,MAAM,GAAG,SAAS;IAOvC,YAAY,CAAC,GAAG,EAAE,MAAM,GAAG,OAAO,CAAC,IAAI,CAAC;IAoCxC,UAAU,CAAC,GAAG,EAAE,MAAM,GAAG,OAAO,CAAC,IAAI,CAAC;IAoCtC,IAAI,IAAI,OAAO,CAAC,IAAI,CAAC;IAoErB,KAAK,IAAI,OAAO,CAAC,IAAI,CAAC;CAG/B;AAED,OAAO,CAAC,MAAM,CAAC;IACX,IAAI,0BAA0B,EAAE,oBAAoB,GAAG,SAAS,CAAC;CACpE;AAGD,eAAO,MAAM,iBAAiB,EAAE,oBAE6C,CAAC;AAE9E,eAAO,MAAM,mBAAmB,qBAA0B,CAAC"}
+{"version":3,"file":"client-session-manager.d.ts","sourceRoot":"","sources":["../../../src/client-session-manager/client-session-manager.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EACR,kBAAkB,EAKlB,cAAc,EACjB,MAAM,SAAS,CAAC;AAEjB,OAAO,EAAE,KAAK,EAAE,KAAK,EAAY,MAAM,kBAAkB,CAAC;AAE1D,OAAO,EAAE,KAAK,YAAY,EAA+B,MAAM,gBAAgB,CAAC;AAsEhF,MAAM,WAAW,0BAA0B;IAEvC,MAAM,CAAC,EAAE,MAAM,GAAG,IAAI,CAAC;IAEvB,iBAAiB,CAAC,EAAE,OAAO,CAAC;CAC/B;AAGD,MAAM,WAAW,qBAAqB;IAClC,YAAY,EAAE,CAAC,IAAI,EAAE,UAAU,KAAK,OAAO,CAAC,IAAI,CAAC,CAAC;IAClD,YAAY,EAAE,MAAM,OAAO,CAAC,UAAU,CAAC,CAAC;IACxC,SAAS,EAAE,OAAO,CAAC,IAAI,CAAC,CAAC;IACzB,GAAG,EAAE,MAAM,CAAC;CACf;AAGD,MAAM,WAAW,iBAAiB;IAC9B,GAAG,EAAE,MAAM,CAAC;IACZ,iBAAiB,EAAE,UAAU,CAAC;IAC9B,QAAQ,EAAE,MAAM,GAAG,cAAc,CAAC;IAClC,SAAS,EAAE,OAAO,CAAC;IACnB,SAAS,CAAC,EAAE,MAAM,CAAC;IACnB,SAAS,CAAC,EAAE,MAAM,CAAC;CACtB;AAcD,qBAAa,oBAAoB;IAE7B,OAAO,CAAC,eAAe,CAAoD;IAC3E,OAAO,CAAC,OAAO,CAAyC;IACxD,OAAO,CAAC,OAAO,CAAkC;IAEjD,OAAO,CAAC,eAAe,CAA8C;IACrE,OAAO,CAAC,eAAe,CAGrB;IACF,OAAO,CAAC,QAAQ,CAAoF;IACpG,OAAO,CAAC,QAAQ,CAAsB;IACtC,OAAO,CAAC,SAAS,CAAc;IAE/B,OAAO,CAAC,QAAQ,CAAuB;IACvC,OAAO,CAAC,MAAM,CAAuB;IACrC,OAAO,CAAC,MAAM,CAAuB;IACrC,OAAO,CAAC,QAAQ,CAAe;IAC/B,OAAO,CAAC,WAAW,CAAwC;IAC3D,OAAO,CAAC,SAAS,CAAkB;IACnC,OAAO,CAAC,UAAU,CAAkB;IACpC,OAAO,CAAC,MAAM,CAAe;IAC7B,OAAO,CAAC,WAAW,CAAsD;IAGzE,OAAO,CAAC,eAAe,CAAS;IAGhC,OAAO,CAAC,qBAAqB,CAAa;IAC1C,OAAO,CAAC,eAAe,CAA+C;gBAE1D,MAAM,GAAE,0BAA+B;IAyCnD,OAAO,CAAC,aAAa;IAWrB,OAAO,CAAC,eAAe;IAQvB,YAAY,CAAC,QAAQ,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,IAAI;YAsBtC,eAAe;IAwC7B,OAAO,CAAC,YAAY;YAYN,QAAQ;YA+CR,QAAQ;IAsItB,OAAO,CAAC,cAAc;YA+BR,sBAAsB;YAkDtB,mBAAmB;IAqIjC,OAAO,CAAC,eAAe;IAejB,QAAQ,CAAC,IAAI,EAAE,kBAAkB,EAAE,YAAY,CAAC,EAAE,YAAY,GAAG,OAAO,CAAC,MAAM,CAAC;IAiHhF,WAAW,CAAC,MAAM,EAAE,iBAAiB,EAAE,YAAY,CAAC,EAAE,YAAY,GAAG,OAAO,CAAC,qBAAqB,CAAC;IA2HlG,IAAI,CAAC,YAAY,EAAE,WAAW,EAAE,GAAG,EAAE,MAAM,EAAE,GAAG,CAAC,EAAE,MAAM,GAAG,cAAc,CAAC,KAAK,EAAE,IAAI,EAAE,OAAO,CAAC;IA+KjG,IAAI,CAAC,GAAG,EAAE,MAAM,EAAE,GAAG,CAAC,EAAE,MAAM,EAAE,MAAM,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,GAAG,CAAC,GAAG,OAAO,CAAC,KAAK,CAAC,MAAM,CAAC,MAAM,EAAE,GAAG,CAAC,CAAC,CAAC;IAuClG,UAAU,CAAC,GAAG,EAAE,MAAM,EAAE,GAAG,EAAE,MAAM,EAAE,UAAU,CAAC,EAAE,KAAK,GAAG,OAAO,CAAC;QAAE,KAAK,EAAE,KAAK,CAAC;QAAC,QAAQ,EAAE,KAAK,CAAA;KAAE,CAAC;IAqB1G,WAAW,CAAC,GAAG,EAAE,MAAM,GAAG,OAAO;IAejC,YAAY,CAAC,GAAG,EAAE,MAAM,GAAG,MAAM,GAAG,SAAS;IAOvC,YAAY,CAAC,GAAG,EAAE,MAAM,GAAG,OAAO,CAAC,IAAI,CAAC;IAoCxC,UAAU,CAAC,GAAG,EAAE,MAAM,GAAG,OAAO,CAAC,IAAI,CAAC;IAqCtC,IAAI,IAAI,OAAO,CAAC,IAAI,CAAC;IA8CrB,KAAK,IAAI,OAAO,CAAC,IAAI,CAAC;CAG/B;AAED,OAAO,CAAC,MAAM,CAAC;IACX,IAAI,0BAA0B,EAAE,oBAAoB,GAAG,SAAS,CAAC;CACpE;AAGD,eAAO,MAAM,iBAAiB,EAAE,oBAE6C,CAAC;AAE9E,eAAO,MAAM,mBAAmB,qBAA0B,CAAC"}

package/dist/src/client-session-manager/client-session-manager.js

@@ -27,9 +27,9 @@ import { injectTraceContext } from "../otel-config.js";
 import { CustomLogFW } from "../otel-logging.js";
 import { version } from "../version.js";
 export class ClientSessionManager {
-    websocket = null;
-    tasks = null;
-    sendQueue = null;
+    sessionManagers = new Map();
+    smLocks = new Map();
+    uidToSm = new Map();
     uidIidRecvQueue;
     uidIidException;
     matchIid;
@@ -136,29 +136,28 @@ export class ClientSessionManager {
         headers['X-Client-Session-ID'] = this.clientSessionId;
         return headers;
     }
-    async wsBackgroundTaskWriter(control) {
+    async smWriter(smId, control) {
         await waitForTracing;
-        const wsLogger = this.logger.withScope('ws-writer');
+        const wsLogger = this.logger.withScope(`ws-writer-${smId.slice(0, 8)}`);
         wsLogger.debug('Writer task started');
-        const ws = this.websocket;
         try {
             while (!control.shouldStop) {
+                const smc = this.sessionManagers.get(smId);
+                if (!smc) {
+                    wsLogger.debug('Session manager connection removed, exiting writer');
+                    break;
+                }
                 try {
-                    const sendQueue = this.sendQueue;
-                    if (!sendQueue) {
-                        wsLogger.debug('Send queue removed, exiting writer');
-                        break;
-                    }
-                    const msg = await sendQueue.get();
+                    const msg = await smc.sendQueue.get();
                     const msgUid = msg.uid;
                     const msgLogger = msgUid ? (this.uidToLogger.get(msgUid) ?? wsLogger) : wsLogger;
-                    if (ws.readyState !== 1) {
+                    if (smc.websocket.readyState !== 1) {
                         const error = new WebSocketConnectionError(`cannot send as websocket is not open.`);
                         enrichError(error, { uid: msgUid, sessionId: this.clientSessionId });
                         throw error;
                     }
                     msgLogger.debug(`Sending ${msg.type} message`);
-                    ws.send(encodeMessage(msg));
+                    smc.websocket.send(encodeMessage(msg));
                 }
                 catch (error) {
                     if (error instanceof Error && error.message === 'Queue closed') {
@@ -172,23 +171,20 @@ export class ClientSessionManager {
         }
         catch (error) {
             wsLogger.error('Writer task failed', error);
-            for (const [matchId, handlers] of this.matchIid.entries()) {
-                handlers.reject(new WebSocketConnectionError(`WebSocket writer failed: ${error.message}`));
-                this.matchIid.delete(matchId);
-            }
-            for (const [_, exceptionMap] of this.uidIidException.entries()) {
-                for (const [_iid, handlers] of exceptionMap.entries()) {
-                    handlers.reject(new WebSocketConnectionError(`WebSocket writer failed: ${error.message}`));
-                }
-            }
+            await this._cleanupSessionManager(smId);
             throw error;
         }
     }
-    async wsBackgroundTaskReader(control) {
+    async smReader(smId, control) {
         await waitForTracing;
-        const wsLogger = this.logger.withScope('ws-reader');
+        const wsLogger = this.logger.withScope(`ws-reader-${smId.slice(0, 8)}`);
         wsLogger.debug('Reader task started');
-        const ws = this.websocket;
+        const smc = this.sessionManagers.get(smId);
+        if (!smc) {
+            wsLogger.warn('Session manager connection not found');
+            return;
+        }
+        const ws = smc.websocket;
         ws.onmessage = (event) => {
             if (control.shouldStop)
                 return;
@@ -289,45 +285,99 @@ export class ClientSessionManager {
             ws.onclose = (event) => {
                 wsLogger.debug(`WebSocket closed (code=${event.code}, reason=${event.reason || 'none'})`);
                 control.shouldStop = true;
-                for (const [matchId, handlers] of this.matchIid.entries()) {
-                    handlers.reject(new WebSocketConnectionError(`WebSocket closed: ${event.reason || 'Connection closed'}`));
-                    this.matchIid.delete(matchId);
-                }
-                for (const [_, exceptionMap] of this.uidIidException.entries()) {
-                    for (const [_iid, handlers] of exceptionMap.entries()) {
-                        handlers.reject(new WebSocketConnectionError(`WebSocket closed: ${event.reason || 'Connection closed'}`));
-                    }
-                    exceptionMap.clear();
-                }
+                this._failSmFutures(smId, `WebSocket closed: ${event.reason || 'Connection closed'}`);
                 resolve();
             };
             ws.onerror = (_error) => {
                 wsLogger.error('WebSocket error');
                 control.shouldStop = true;
-                for (const [matchId, handlers] of this.matchIid.entries()) {
-                    handlers.reject(new WebSocketConnectionError('WebSocket error occurred'));
-                    this.matchIid.delete(matchId);
-                }
-                for (const [_, exceptionMap] of this.uidIidException.entries()) {
-                    for (const [_iid, handlers] of exceptionMap.entries()) {
-                        handlers.reject(new WebSocketConnectionError('WebSocket error occurred'));
-                    }
-                    exceptionMap.clear();
-                }
+                this._failSmFutures(smId, 'WebSocket error occurred');
                 resolve();
             };
         });
     }
-    async _ensureWebSocketConnection() {
-        if (this.websocket && this.websocket.readyState === 1) {
+    _failSmFutures(smId, reason) {
+        const affectedUids = [];
+        for (const [uid, sid] of this.uidToSm.entries()) {
+            if (sid === smId) {
+                affectedUids.push(uid);
+            }
+        }
+        for (const [matchId, handlers] of this.matchIid.entries()) {
+            handlers.reject(new WebSocketConnectionError(reason));
+            this.matchIid.delete(matchId);
+        }
+        for (const uid of affectedUids) {
+            const exceptionMap = this.uidIidException.get(uid);
+            if (exceptionMap) {
+                for (const [_iid, handlers] of exceptionMap.entries()) {
+                    handlers.reject(new WebSocketConnectionError(reason));
+                }
+                exceptionMap.clear();
+            }
+        }
+    }
+    async _cleanupSessionManager(smId) {
+        const smc = this.sessionManagers.get(smId);
+        if (!smc) {
             return;
         }
+        this.logger.debug(`Cleaning up session manager ${smId.slice(0, 8)}`);
+        this.sessionManagers.delete(smId);
+        this.smLocks.delete(smId);
+        const [readerControl, writerControl] = smc.tasks;
+        readerControl.shouldStop = true;
+        writerControl.shouldStop = true;
+        smc.sendQueue.close();
+        try {
+            if (smc.websocket.readyState === 0 || smc.websocket.readyState === 1) {
+                smc.websocket.close(1000, 'Normal closure');
+            }
+        }
+        catch (error) {
+            this.logger.warn(`Failed to close WebSocket for session manager ${smId.slice(0, 8)}`, error);
+        }
+        await Promise.allSettled([readerControl.promise, writerControl.promise]);
+        this._failSmFutures(smId, 'Session manager connection closed');
+        const affectedUids = [];
+        for (const [uid, sid] of this.uidToSm.entries()) {
+            if (sid === smId) {
+                affectedUids.push(uid);
+            }
+        }
+        for (const uid of affectedUids) {
+            this.uidToSm.delete(uid);
+        }
+        this.logger.debug(`Cleaned up session manager ${smId.slice(0, 8)}`);
+    }
+    async _ensureSmConnection(smId) {
+        const existing = this.sessionManagers.get(smId);
+        if (existing && existing.websocket.readyState === 1) {
+            return;
+        }
+        const existingLock = this.smLocks.get(smId);
+        if (existingLock) {
+            await existingLock;
+            return;
+        }
+        let resolveLock;
+        const lockPromise = new Promise((resolve) => {
+            resolveLock = resolve;
+        });
+        this.smLocks.set(smId, lockPromise);
         await waitForTracing;
-        const span = this.logger.startSpan('sdk.websocket_connection');
+        const span = this.logger.startSpan('sdk.sm_connection');
+        span.setAttribute('session_manager.id', smId);
         try {
-            this.sendQueue = new Queue();
+            const doubleCheck = this.sessionManagers.get(smId);
+            if (doubleCheck && doubleCheck.websocket.readyState === 1) {
+                resolveLock();
+                this.smLocks.delete(smId);
+                return;
+            }
+            const sendQueue = new Queue();
             const websocket_uri = `${this.baseWs}/socket`;
-            this.logger.debug(`Connecting WebSocket`);
+            this.logger.debug(`Connecting WebSocket for session manager ${smId.slice(0, 8)}`);
             span.setAttribute('websocket.uri', websocket_uri);
             span.setAttribute('websocket.state', 'connecting');
             const headers = {};
@@ -347,7 +397,6 @@ export class ClientSessionManager {
             ws.binaryType = 'arraybuffer';
             return new Promise((resolve, reject) => {
                 ws.onopen = () => {
-                    this.websocket = ws;
                     try {
                         ws._socket?.unref?.();
                     }
@@ -361,23 +410,32 @@ export class ClientSessionManager {
                         shouldStop: false,
                         promise: Promise.resolve(),
                     };
-                    readerControl.promise = this.wsBackgroundTaskReader(readerControl).catch((error) => {
-                        this.logger.error('Reader background task crashed', error);
+                    const smc = {
+                        websocket: ws,
+                        sendQueue,
+                        tasks: [readerControl, writerControl],
+                    };
+                    this.sessionManagers.set(smId, smc);
+                    readerControl.promise = this.smReader(smId, readerControl).catch((error) => {
+                        this.logger.error(`Reader task crashed for session manager ${smId.slice(0, 8)}`, error);
                     });
-                    writerControl.promise = this.wsBackgroundTaskWriter(writerControl).catch((error) => {
-                        this.logger.error('Writer background task crashed', error);
+                    writerControl.promise = this.smWriter(smId, writerControl).catch((error) => {
+                        this.logger.error(`Writer task crashed for session manager ${smId.slice(0, 8)}`, error);
                     });
-                    this.tasks = [readerControl, writerControl];
-                    this.logger.info(`WebSocket connected`);
+                    this.logger.info(`WebSocket connected for session manager ${smId.slice(0, 8)}`);
                     span.setAttribute('websocket.state', 'connected');
                     span.end();
+                    resolveLock();
+                    this.smLocks.delete(smId);
                     resolve();
                 };
                 ws.onerror = (error) => {
-                    this.logger.error(`WebSocket connection failed:\n${error.toString()}`);
+                    this.logger.error(`WebSocket connection failed for session manager ${smId.slice(0, 8)}:\n${error.toString()}`);
                     span.setAttribute('websocket.state', 'error');
                     span.recordException(error);
                     span.end();
+                    resolveLock();
+                    this.smLocks.delete(smId);
                     reject(new WebSocketConnectionError('WebSocket connection failed'));
                 };
             });
@@ -386,6 +444,8 @@ export class ClientSessionManager {
             span.setAttribute('websocket.state', 'error');
             span.recordException(error);
             span.end();
+            resolveLock();
+            this.smLocks.delete(smId);
             throw error;
         }
     }
@@ -457,10 +517,22 @@ export class ClientSessionManager {
                     logger.warn(message);
                 }
             }
-            const uid = await response.text();
+            const responseData = await response.json();
+            const uid = responseData.uid;
+            const smId = responseData.session_manager_id;
+            if (!uid) {
+                logger.error(`Server response missing "uid" field: ${JSON.stringify(responseData)}`);
+                throw new ConnectionError('Invalid server response: missing "uid" field');
+            }
+            if (!smId) {
+                logger.error(`Server response missing "session_manager_id" field: ${JSON.stringify(responseData)}`);
+                throw new ConnectionError('Invalid server response: missing "session_manager_id" field');
+            }
             logger.info('Created agent');
             span.setAttribute('agent.uid', uid);
-            await this._ensureWebSocketConnection();
+            span.setAttribute('agent.session_manager_id', smId);
+            this.uidToSm.set(uid, smId);
+            await this._ensureSmConnection(smId);
             this.knownUids.add(uid);
             this._initAgentState(uid, parentLogger);
             return uid;
@@ -493,7 +565,20 @@ export class ClientSessionManager {
                 enrichError(error, { uid, sessionId: this.clientSessionId });
                 throw error;
             }
+            const smId = this.uidToSm.get(uid);
+            if (!smId) {
+                const error = new WebSocketConnectionError(`No session manager found for agent ${uid}`);
+                enrichError(error, { uid, sessionId: this.clientSessionId });
+                throw error;
+            }
+            const smc = this.sessionManagers.get(smId);
+            if (!smc || smc.websocket.readyState !== 1) {
+                const error = new WebSocketConnectionError(`No active connection for session manager ${smId}`);
+                enrichError(error, { uid, sessionId: this.clientSessionId });
+                throw error;
+            }
             span.setAttribute('agent.uid', uid);
+            span.setAttribute('agent.session_manager_id', smId);
             span.setAttribute('agent.task_desc', typeof taskDesc === 'string' ? taskDesc : taskDesc.template);
             span.setAttribute('invocation.streaming', streaming);
             const matchId = this.idIssuer();
@@ -506,11 +591,11 @@ export class ClientSessionManager {
                 prompt: taskDesc,
                 streaming: streaming,
             };
-            invLogger.debug(`Invoking with match_id=${matchId}`);
+            invLogger.debug(`Invoking with match_id=${matchId} on session manager ${smId.slice(0, 8)}`);
             const iidPromise = new Promise((resolve, reject) => {
                 this.matchIid.set(matchId, { resolve, reject });
             });
-            this.sendQueue.put(msg);
+            smc.sendQueue.put(msg);
             const iid = await iidPromise;
             this.iidToUid.set(iid, uid);
             invLogger.info(`Invocation created with iid=${iid}`);
@@ -527,8 +612,13 @@ export class ClientSessionManager {
                 return { resolvers: { resolve, reject, promise }, promise };
             })();
             this.uidIidException.get(uid).set(iid, resolvers);
+            const capturedSmId = smId;
             return {
                 send_message: async (data) => {
+                    const currentSmc = this.sessionManagers.get(capturedSmId);
+                    if (!currentSmc || currentSmc.websocket.readyState !== 1) {
+                        throw new WebSocketConnectionError(`Session manager connection closed for ${capturedSmId}`);
+                    }
                     const dataMsg = {
                         type: 'data',
                         uid,
@@ -536,7 +626,7 @@ export class ClientSessionManager {
                         data,
                         timestamp: new Date().toISOString(),
                     };
-                    this.sendQueue.put(dataMsg);
+                    currentSmc.sendQueue.put(dataMsg);
                     invLogger.debug(`Queued ${data.length} bytes to send`);
                 },
                 recv_message: async () => {
@@ -728,23 +818,27 @@ export class ClientSessionManager {
         }
         return response.json();
     }
-    async fetchUsage(uid, iid, lastTotal) {
+    async fetchUsage(uid, iid, _lastTotal) {
         let total = Usage.zero();
         const logs = await this.logs(uid, iid, { type: 'sm_inference_usage' });
         for (const log of logs) {
             const logUsage = log.usage;
             if (logUsage) {
-                total = total.add(Usage.fromCompletions(logUsage, total));
+                total = total.add(Usage.fromCompletions(logUsage));
             }
         }
-        let usage = total;
-        if (lastTotal) {
-            usage = total.sub(lastTotal.replace({ outputTokens: 0 }));
-        }
-        return { usage, newTotal: total };
+        return { usage: total, newTotal: total };
     }
     agentExists(uid) {
-        return this.knownUids.has(uid) && this.websocket !== null && this.websocket.readyState === 1;
+        if (!this.knownUids.has(uid)) {
+            return false;
+        }
+        const smId = this.uidToSm.get(uid);
+        if (!smId) {
+            return false;
+        }
+        const smc = this.sessionManagers.get(smId);
+        return smc !== undefined && smc.websocket.readyState === 1;
     }
     getUidForIid(iid) {
         return this.iidToUid.get(iid);
@@ -796,6 +890,7 @@ export class ClientSessionManager {
         this.uidIidException.delete(uid);
         this.knownUids.delete(uid);
         this.uidToLogger.delete(uid);
+        this.uidToSm.delete(uid);
         logger.debug(`Closed agent ${uid.slice(0, 8)}`);
         await this.destroyAgent(uid);
         this.unrefConnection();
@@ -807,32 +902,14 @@ export class ClientSessionManager {
         const span = this.logger.startSpan('stop');
         try {
             const agentCount = this.knownUids.size;
-            this.logger.debug(`CSM Stop: stopping with ${agentCount} agent(s)`);
+            const smCount = this.sessionManagers.size;
+            this.logger.debug(`CSM Stop: stopping with ${agentCount} agent(s) on ${smCount} session manager(s)`);
             span.setAttribute('agent_count', agentCount);
+            span.setAttribute('session_manager_count', smCount);
             const uids = Array.from(this.knownUids);
             await Promise.allSettled(uids.map((uid) => this.closeAgent(uid)));
-            if (this.sendQueue) {
-                this.sendQueue.close();
-                this.sendQueue = null;
-            }
-            if (this.websocket) {
-                try {
-                    if (this.websocket.readyState === 0 || this.websocket.readyState === 1) {
-                        this.websocket.close(1000, 'Normal closure');
-                    }
-                }
-                catch (error) {
-                    this.logger.warn('Failed to close WebSocket', error);
-                }
-                this.websocket = null;
-            }
-            if (this.tasks) {
-                const [readerControl, writerControl] = this.tasks;
-                readerControl.shouldStop = true;
-                writerControl.shouldStop = true;
-                await Promise.allSettled([readerControl.promise, writerControl.promise]);
-                this.tasks = null;
-            }
+            const smIds = Array.from(this.sessionManagers.keys());
+            await Promise.allSettled(smIds.map((smId) => this._cleanupSessionManager(smId)));
             this.logger.debug('All resources cleaned up');
             this.isStopped = true;
             if (this.invocationGuard) {