fix: graceful ctrl+c shutdown and await session-end memory writes

This commit is contained in:
William Valentin
2026-02-18 11:30:47 -08:00
parent 55cde541ea
commit 21232748b9
6 changed files with 103 additions and 30 deletions
+8 -3
View File
@@ -564,11 +564,14 @@ export class GatewayServer {
}
}
// Close all WebSocket connections first
// Close all WebSocket connections first.
// Await disconnects so end-of-session summary/memory writes complete before process exit.
const disconnects: Array<Promise<void>> = [];
for (const [ws, connectionId] of this.connectionMap) {
this.sessionBridge.disconnect(connectionId);
disconnects.push(this.sessionBridge.disconnect(connectionId));
ws.close(1001, 'Server shutting down');
}
await Promise.allSettled(disconnects);
this.connectionMap.clear();
this.connectionStateMap.clear();
@@ -628,7 +631,9 @@ export class GatewayServer {
});
ws.on('close', () => {
this.sessionBridge.disconnect(connectionId);
void this.sessionBridge.disconnect(connectionId).catch((error) => {
console.warn('Session disconnect failed:', error);
});
this.connectionMap.delete(ws);
this.connectionRateMap.delete(connectionId);
this.connectionStateMap.delete(connectionId);
+5 -6
View File
@@ -96,10 +96,10 @@ describe('SessionBridge', () => {
expect(bridge.getSessionId('conn-1')).toBe('ws:conn-1');
});
it('disconnect removes client but preserves session', () => {
it('disconnect removes client but preserves session', async () => {
const bridge = createBridge();
bridge.connect('conn-1');
bridge.disconnect('conn-1');
await bridge.disconnect('conn-1');
expect(bridge.connectionCount).toBe(0);
expect(bridge.getAgent('conn-1')).toBeUndefined();
});
@@ -148,8 +148,7 @@ describe('SessionBridge', () => {
} as unknown as SessionBridgeConfig['config'],
});
bridge.connect('conn-end-summary');
bridge.disconnect('conn-end-summary');
await new Promise(resolve => setTimeout(resolve, 0));
await bridge.disconnect('conn-end-summary');
expect(memoryStore.write).toHaveBeenCalled();
});
@@ -242,13 +241,13 @@ describe('SessionBridge', () => {
expect(sessions).toEqual([{ sessionId: 'ws:conn-1', connections: 2 }]);
});
it('shared sessions keep agent alive when one client disconnects', () => {
it('shared sessions keep agent alive when one client disconnects', async () => {
const bridge = createBridge();
bridge.connect('conn-1');
bridge.connect('conn-2');
bridge.switchSession('conn-2', 'ws:conn-1');
bridge.disconnect('conn-1');
await bridge.disconnect('conn-1');
// conn-2 still has the session
expect(bridge.getAgent('conn-2')).toBeDefined();
expect(bridge.connectionCount).toBe(1);
+11 -9
View File
@@ -53,7 +53,7 @@ export class SessionBridge {
}
/** Remove a WS connection. Does NOT destroy the session (persists in SQLite). */
disconnect(connectionId: string): void {
async disconnect(connectionId: string): Promise<void> {
const client = this.clients.get(connectionId);
if (client) {
// Only remove the agent if no other clients share the session
@@ -73,15 +73,17 @@ export class SessionBridge {
writeToMemory: summaryConfig.write_to_memory,
memoryNamespace: summaryConfig.memory_namespace,
};
void summarizeSessionOnEnd({
agent,
sessionId: client.sessionId,
history,
config: mappedConfig,
memoryStore: this.config.memoryStore,
}).catch((error) => {
try {
await summarizeSessionOnEnd({
agent,
sessionId: client.sessionId,
history,
config: mappedConfig,
memoryStore: this.config.memoryStore,
});
} catch (error) {
console.warn('Session end summary failed:', error);
});
}
}
this.agents.delete(client.sessionId);
}