refactor: use generated types for session store RPC handlers

SteveSandersonMS · Copilot · SteveSandersonMS · commit a0daf7cff0aa · 2026-03-24T12:43:28.000Z
- Replace hand-written onRequest registrations with generated
  registerClientApiHandlers() and a routing SessionStoreHandler
- Re-export SessionStoreHandler and ClientApiHandlers from index.ts
- Fix listSessions concurrency: replace singleton listSessionsStore
  with a Map keyed by incrementing ID
- Add cleanup of sessionStoreConfigs and listSessionsStores on
  stop() and forceStop()

Co-authored-by: Copilot &lt;223556219+Copilot@users.noreply.github.com&gt;
diff --git a/nodejs/src/client.ts b/nodejs/src/client.ts
@@ -23,7 +23,8 @@ import {
     StreamMessageReader,
     StreamMessageWriter,
 } from "vscode-jsonrpc/node.js";
-import { createServerRpc } from "./generated/rpc.js";
+import { createServerRpc, registerClientApiHandlers } from "./generated/rpc.js";
+import type { SessionStoreHandler } from "./generated/rpc.js";
 import { getSdkProtocolVersion } from "./sdkProtocolVersion.js";
 import { CopilotSession, NO_RESULT_PERMISSION_V2_ERROR } from "./session.js";
 import { getTraceContext } from "./telemetry.js";
@@ -179,8 +180,12 @@ export class CopilotClient {
     private negotiatedProtocolVersion: number | null = null;
     /** Per-session store configs, keyed by sessionId. Used to route sessionStore.* RPC requests. */
     private sessionStoreConfigs: Map<string, SessionStoreConfig> = new Map();
-    /** Temporary store used during listSessions() calls with a session store. */
-    private listSessionsStore: SessionStoreConfig | null = null;
+    /**
+     * Stores used during listSessions() calls. Keyed by an incrementing ID to support concurrent calls.
+     * The latest registered store is used when the server calls sessionStore.list.
+     */
+    private listSessionsStores: Map<number, SessionStoreConfig> = new Map();
+    private listSessionsStoreNextId = 0;
 
     /**
      * Typed server-scoped RPC methods.
@@ -403,6 +408,8 @@ export class CopilotClient {
             }
         }
         this.sessions.clear();
+        this.sessionStoreConfigs.clear();
+        this.listSessionsStores.clear();
 
         // Close connection
         if (this.connection) {
@@ -487,6 +494,8 @@ export class CopilotClient {
 
         // Clear sessions immediately without trying to destroy them
         this.sessions.clear();
+        this.sessionStoreConfigs.clear();
+        this.listSessionsStores.clear();
 
         // Force close connection
         if (this.connection) {
@@ -1005,9 +1014,11 @@ export class CopilotClient {
                   ? { filter: filterOrOptions as SessionListFilter }
                   : undefined;
 
-        // Temporarily register the store for the duration of this RPC call
+        // Register the store for the duration of this RPC call (supports concurrent calls)
+        let storeId: number | undefined;
         if (options?.sessionStore) {
-            this.listSessionsStore = options.sessionStore;
+            storeId = this.listSessionsStoreNextId++;
+            this.listSessionsStores.set(storeId, options.sessionStore);
         }
 
         try {
@@ -1037,8 +1048,8 @@ export class CopilotClient {
                 context: s.context,
             }));
         } finally {
-            if (options?.sessionStore) {
-                this.listSessionsStore = null;
+            if (storeId !== undefined) {
+                this.listSessionsStores.delete(storeId);
             }
         }
     }
@@ -1505,78 +1516,63 @@ export class CopilotClient {
             }): Promise<{ output?: unknown }> => await this.handleHooksInvoke(params)
         );
 
-        // Register session store RPC handlers (server→client callbacks for event persistence)
-        this.connection.onRequest(
-            "sessionStore.load",
-            async (params: { sessionId: string }): Promise<{ events: SessionEvent[] }> => {
-                const store = this.sessionStoreConfigs.get(params.sessionId);
-                if (!store) {
-                    throw new Error(`No session store registered for session ${params.sessionId}`);
-                }
-                const events = await store.onLoad(params.sessionId);
-                return { events };
-            }
-        );
+        // Register session store RPC handlers via generated registration function.
+        // The handler routes each call to the appropriate SessionStoreConfig based on sessionId.
+        registerClientApiHandlers(this.connection, {
+            sessionStore: this.createSessionStoreHandler(),
+        });
 
-        this.connection.onRequest(
-            "sessionStore.append",
-            async (params: { sessionId: string; events: SessionEvent[] }): Promise<void> => {
-                const store = this.sessionStoreConfigs.get(params.sessionId);
-                if (!store) {
-                    throw new Error(`No session store registered for session ${params.sessionId}`);
-                }
-                await store.onAppend(params.events, params.sessionId);
-            }
-        );
+        this.connection.onClose(() => {
+            this.state = "disconnected";
+        });
 
-        this.connection.onRequest(
-            "sessionStore.truncate",
-            async (params: {
-                sessionId: string;
-                upToEventId: string;
-            }): Promise<{ eventsRemoved: number; eventsKept: number }> => {
-                const store = this.sessionStoreConfigs.get(params.sessionId);
-                if (!store) {
-                    throw new Error(`No session store registered for session ${params.sessionId}`);
-                }
-                return await store.onTruncate(params.upToEventId, params.sessionId);
+        this.connection.onError((_error) => {
+            this.state = "disconnected";
+        });
+    }
+
+    /**
+     * Create a {@link SessionStoreHandler} that routes each RPC call to the
+     * appropriate {@link SessionStoreConfig} registered for the given session.
+     */
+    private createSessionStoreHandler(): SessionStoreHandler {
+        const getStore = (sessionId: string): SessionStoreConfig => {
+            const store = this.sessionStoreConfigs.get(sessionId);
+            if (!store) {
+                throw new Error(`No session store registered for session ${sessionId}`);
             }
-        );
+            return store;
+        };
 
-        this.connection.onRequest(
-            "sessionStore.list",
-            async (): Promise<{
-                sessions: Array<{ sessionId: string; mtime: string; birthtime: string }>;
-            }> => {
+        return {
+            load: async (params) => {
+                const store = getStore(params.sessionId);
+                const events = await store.onLoad(params.sessionId);
+                // Events are opaque on the wire but typed in the SDK consumer API
+                return { events: events as Record<string, unknown>[] };
+            },
+            append: async (params) => {
+                const store = getStore(params.sessionId);
+                await store.onAppend(params.events as SessionEvent[], params.sessionId);
+            },
+            truncate: async (params) => {
+                const store = getStore(params.sessionId);
+                return store.onTruncate(params.upToEventId, params.sessionId);
+            },
+            list: async () => {
                 // Use the most recently registered store for listing.
-                // In practice, listSessions() registers a temporary store before the RPC call.
-                const store = this.listSessionsStore;
+                const store = Array.from(this.listSessionsStores.values()).pop();
                 if (!store) {
                     throw new Error("No session store registered for listing");
                 }
                 const sessions = await store.onListSessions();
                 return { sessions };
-            }
-        );
-
-        this.connection.onRequest(
-            "sessionStore.delete",
-            async (params: { sessionId: string }): Promise<void> => {
-                const store = this.sessionStoreConfigs.get(params.sessionId);
-                if (!store) {
-                    throw new Error(`No session store registered for session ${params.sessionId}`);
-                }
+            },
+            delete: async (params) => {
+                const store = getStore(params.sessionId);
                 await store.onDelete(params.sessionId);
-            }
-        );
-
-        this.connection.onClose(() => {
-            this.state = "disconnected";
-        });
-
-        this.connection.onError((_error) => {
-            this.state = "disconnected";
-        });
+            },
+        };
     }
 
     private handleSessionEventNotification(notification: unknown): void {
diff --git a/nodejs/src/index.ts b/nodejs/src/index.ts
@@ -44,6 +44,8 @@ export type {
     SessionListFilter,
     SessionMetadata,
     SessionStoreConfig,
+    SessionStoreHandler,
+    ClientApiHandlers,
     SystemMessageAppendConfig,
     SystemMessageConfig,
     SystemMessageReplaceConfig,
diff --git a/nodejs/src/types.ts b/nodejs/src/types.ts
@@ -10,6 +10,19 @@
 import type { SessionEvent as GeneratedSessionEvent } from "./generated/session-events.js";
 export type SessionEvent = GeneratedSessionEvent;
 
+// Re-export generated client API types
+export type {
+    SessionStoreHandler,
+    SessionStoreLoadParams,
+    SessionStoreLoadResult,
+    SessionStoreAppendParams,
+    SessionStoreTruncateParams,
+    SessionStoreTruncateResult,
+    SessionStoreListResult,
+    SessionStoreDeleteParams,
+    ClientApiHandlers,
+} from "./generated/rpc.js";
+
 /**
  * Options for creating a CopilotClient
  */
@@ -1022,6 +1035,14 @@ export interface SessionContext {
  * the server routes all event persistence through these callbacks over RPC instead
  * of using its default file-based storage.
  */
+/**
+ * Configuration for a custom session store backend.
+ *
+ * The `descriptor` identifies the storage backend. The callback methods
+ * correspond to the generated {@link SessionStoreHandler} interface but
+ * use a consumer-friendly signature where `sessionId` is a plain argument
+ * rather than embedded in a params object.
+ */
 export interface SessionStoreConfig {
     /**
      * Opaque descriptor identifying this storage backend.
