fix: graceful ctrl+c shutdown and await session-end memory writes
This commit is contained in:
+30
-10
@@ -227,15 +227,29 @@ export function registerTuiCommand(program: Command): void {
|
||||
},
|
||||
});
|
||||
|
||||
const cleanup = () => {
|
||||
void lifecycle.shutdown();
|
||||
sessionStore.close();
|
||||
let cleanupPromise: Promise<void> | null = null;
|
||||
const cleanup = async () => {
|
||||
if (!cleanupPromise) {
|
||||
cleanupPromise = (async () => {
|
||||
await lifecycle.shutdown();
|
||||
sessionStore.close();
|
||||
})();
|
||||
}
|
||||
return cleanupPromise;
|
||||
};
|
||||
|
||||
process.on('SIGINT', () => {
|
||||
cleanup();
|
||||
process.exit(0);
|
||||
});
|
||||
const signalHandler = (signal: NodeJS.Signals) => {
|
||||
console.log(`\nReceived ${signal}; shutting down TUI...`);
|
||||
void cleanup()
|
||||
.then(() => process.exit(0))
|
||||
.catch((error) => {
|
||||
console.error('TUI shutdown failed:', error);
|
||||
process.exit(1);
|
||||
});
|
||||
};
|
||||
|
||||
process.on('SIGINT', signalHandler);
|
||||
process.on('SIGTERM', signalHandler);
|
||||
|
||||
const transferSessionToTarget = (target: string): string => {
|
||||
const normalizedTarget = target.trim().toLowerCase();
|
||||
@@ -274,7 +288,9 @@ export function registerTuiCommand(program: Command): void {
|
||||
modelProviderConfigs,
|
||||
contextThresholdPct: config.compaction.threshold_pct,
|
||||
onTransfer: transferSessionToTarget,
|
||||
onExit: cleanup,
|
||||
onExit: () => {
|
||||
void cleanup();
|
||||
},
|
||||
});
|
||||
} else {
|
||||
let switchingToFullscreen = false;
|
||||
@@ -316,12 +332,16 @@ export function registerTuiCommand(program: Command): void {
|
||||
modelProviderConfigs,
|
||||
contextThresholdPct: config.compaction.threshold_pct,
|
||||
onTransfer: transferSessionToTarget,
|
||||
onExit: cleanup,
|
||||
onExit: () => {
|
||||
void cleanup();
|
||||
},
|
||||
});
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
cleanup();
|
||||
await cleanup();
|
||||
process.off('SIGINT', signalHandler);
|
||||
process.off('SIGTERM', signalHandler);
|
||||
});
|
||||
}
|
||||
|
||||
+34
-2
@@ -540,8 +540,40 @@ export async function startServices(deps: {
|
||||
});
|
||||
|
||||
// Signal handlers
|
||||
const signalHandler = () => {
|
||||
lifecycle.shutdown().then(() => process.exit(0));
|
||||
let shutdownPromise: Promise<void> | null = null;
|
||||
let forceExitArmed = false;
|
||||
const forceExitTimeoutMs = 15_000;
|
||||
|
||||
const signalHandler = (signal: NodeJS.Signals) => {
|
||||
if (shutdownPromise) {
|
||||
if (forceExitArmed) {
|
||||
console.warn(`Second ${signal} received; forcing exit now.`);
|
||||
process.exit(130);
|
||||
}
|
||||
forceExitArmed = true;
|
||||
console.warn(`Shutdown already in progress. Send ${signal} again to force exit.`);
|
||||
return;
|
||||
}
|
||||
|
||||
console.log(`Received ${signal}; shutting down gracefully...`);
|
||||
shutdownPromise = lifecycle.shutdown();
|
||||
const forceTimer = setTimeout(() => {
|
||||
forceExitArmed = true;
|
||||
console.error(`Graceful shutdown timed out after ${forceExitTimeoutMs / 1000}s; forcing exit.`);
|
||||
process.exit(130);
|
||||
}, forceExitTimeoutMs);
|
||||
forceTimer.unref();
|
||||
|
||||
shutdownPromise
|
||||
.then(() => {
|
||||
clearTimeout(forceTimer);
|
||||
process.exit(0);
|
||||
})
|
||||
.catch((error) => {
|
||||
clearTimeout(forceTimer);
|
||||
console.error('Graceful shutdown failed:', error);
|
||||
process.exit(1);
|
||||
});
|
||||
};
|
||||
|
||||
process.on('SIGINT', signalHandler);
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user