diff --git a/src/commands/config-command.ts b/src/commands/config-command.ts index 28bb0dfa..611b553c 100644 --- a/src/commands/config-command.ts +++ b/src/commands/config-command.ts @@ -160,10 +160,10 @@ export async function handleConfigCommand(args: string[]): Promise { try { // Start server - const { server, wss } = await startServer({ port, dev: options.dev }); + const { server, wss, cleanup } = await startServer({ port, dev: options.dev }); // Setup graceful shutdown - setupGracefulShutdown(server, wss); + setupGracefulShutdown(server, wss, cleanup); const url = `http://localhost:${port}`; diff --git a/src/web-server/index.ts b/src/web-server/index.ts index fd64547c..61dd2941 100644 --- a/src/web-server/index.ts +++ b/src/web-server/index.ts @@ -13,6 +13,7 @@ import { WebSocketServer } from 'ws'; import { setupWebSocket } from './websocket'; import { createSessionMiddleware, authMiddleware } from './middleware/auth-middleware'; import { startAutoSyncWatcher, stopAutoSyncWatcher } from '../cliproxy/sync'; +import { shutdownUsageAggregator } from './usage/aggregator'; export interface ServerOptions { port: number; @@ -107,6 +108,7 @@ export async function startServer(options: ServerOptions): Promise { wsCleanup(); stopAutoSyncWatcher().catch(() => {}); + shutdownUsageAggregator(); }; // Start listening