fix: make serve shutdown reliable

.all-contributorsrc

@@ -76,6 +76,15 @@
       "contributions": [
         "code"
       ]
+    },
+    {
+      "login": "lmingde",
+      "name": "Mingde",
+      "avatar_url": "https://avatars.githubusercontent.com/u/38581341?v=4",
+      "profile": "https://github.com/lmingde",
+      "contributions": [
+        "bug"
+      ]
     }
   ],
   "contributorsPerLine": 7,

CONTRIBUTING.md

@@ -28,9 +28,11 @@ pnpm lint
 
 ```text
 packages/
-  core/     - Shared types, database schema, pricing data
-  cli/      - CLI tool for parsing logs, querying data, sync
-  web/      - SvelteKit web dashboard (SPA)
+  core/     - Shared types, database schema, pricing data, utilities
+  cli/      - Published CLI, parsers, local API server, sync, PM2 helpers
+  web/      - Local SvelteKit dashboard bundled into the CLI
+  widget/   - Electron tray/menu-bar widget
+  site/     - Official website, docs, accounts, uploads, leaderboard
 ```
 
 ## Submitting Changes

README.md

@@ -191,6 +191,9 @@ Thanks to these wonderful people ([emoji key](https://allcontributors.org/en/ref
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/jlxyfll"><img src="https://avatars.githubusercontent.com/u/16436887?v=4?s=100" width="100px;" alt="jlxyfll"/><br /><sub><b>jlxyfll</b></sub></a><br /><a href="https://github.com/juliantanx/aiusage/commits?author=jlxyfll" title="Code">💻</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/Ed-Bg"><img src="https://avatars.githubusercontent.com/u/68063080?v=4?s=100" width="100px;" alt="Harlan Hu"/><br /><sub><b>Harlan Hu</b></sub></a><br /><a href="https://github.com/juliantanx/aiusage/commits?author=Ed-Bg" title="Code">💻</a></td>
     </tr>
+    <tr>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/lmingde"><img src="https://avatars.githubusercontent.com/u/38581341?v=4?s=100" width="100px;" alt="Mingde"/><br /><sub><b>Mingde</b></sub></a><br /><a href="https://github.com/juliantanx/aiusage/issues?q=author%3Almingde" title="Bug reports">🐛</a></td>
+    </tr>
   </tbody>
 </table>
 

README_zh.md

@@ -191,6 +191,9 @@ pnpm dev
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/jlxyfll"><img src="https://avatars.githubusercontent.com/u/16436887?v=4?s=100" width="100px;" alt="jlxyfll"/><br /><sub><b>jlxyfll</b></sub></a><br /><a href="https://github.com/juliantanx/aiusage/commits?author=jlxyfll" title="Code">💻</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/Ed-Bg"><img src="https://avatars.githubusercontent.com/u/68063080?v=4?s=100" width="100px;" alt="Harlan Hu"/><br /><sub><b>Harlan Hu</b></sub></a><br /><a href="https://github.com/juliantanx/aiusage/commits?author=Ed-Bg" title="Code">💻</a></td>
     </tr>
+    <tr>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/lmingde"><img src="https://avatars.githubusercontent.com/u/38581341?v=4?s=100" width="100px;" alt="Mingde"/><br /><sub><b>Mingde</b></sub></a><br /><a href="https://github.com/juliantanx/aiusage/issues?q=author%3Almingde" title="Bug reports">🐛</a></td>
+    </tr>
   </tbody>
 </table>
 

packages/cli/src/commands/serve-shutdown.ts

@@ -0,0 +1,57 @@
+import type http from 'node:http'
+
+export const FORCE_SHUTDOWN_TIMEOUT_MS = 1_000
+
+interface GracefulShutdownOptions {
+  server: http.Server
+  cleanup: () => void
+  stopRuntime: () => void
+  exit?: (code: number) => void
+  log?: (message: string) => void
+  timeoutMs?: number
+}
+
+export function createGracefulShutdownHandler(options: GracefulShutdownOptions): () => void {
+  const exit = options.exit ?? ((code: number) => process.exit(code))
+  const log = options.log ?? ((message: string) => console.log(message))
+  const timeoutMs = options.timeoutMs ?? FORCE_SHUTDOWN_TIMEOUT_MS
+  let shuttingDown = false
+  let exited = false
+  let forceTimer: ReturnType<typeof setTimeout> | undefined
+
+  const exitOnce = () => {
+    if (exited) return
+    exited = true
+    if (forceTimer) clearTimeout(forceTimer)
+    exit(0)
+  }
+
+  return () => {
+    if (shuttingDown) {
+      // A second signal means the caller no longer wants to wait for graceful
+      // shutdown. Do not call server.close() again: each call adds another
+      // close listener while active connections are still draining.
+      options.server.closeAllConnections()
+      exitOnce()
+      return
+    }
+    shuttingDown = true
+
+    log('\nShutting down...')
+    options.cleanup()
+    options.stopRuntime()
+
+    forceTimer = setTimeout(() => {
+      options.server.closeAllConnections()
+      exitOnce()
+    }, timeoutMs)
+    forceTimer.unref()
+
+    try {
+      options.server.close(exitOnce)
+      options.server.closeIdleConnections()
+    } catch {
+      exitOnce()
+    }
+  }
+}

packages/cli/src/commands/serve.ts

@@ -15,6 +15,7 @@ import { getSyncTarget } from '../sync/target.js'
 import { RuntimeSettingsController } from '../runtime/settings-controller.js'
 import { AsyncTaskQueue } from '../db/write-queue.js'
 import { fetchExchangeRate, CACHE_TTL_MS } from '@aiusage/core'
+import { createGracefulShutdownHandler } from './serve-shutdown.js'
 import type Database from 'better-sqlite3'
 
 export interface ServeOptions {
@@ -202,20 +203,11 @@ export function serve(options: ServeOptions): void {
     try { unlinkSync(PORT_FILE) } catch {}
   }
 
-  process.on('SIGINT', () => {
-    console.log('\nShutting down...')
-    cleanup()
-    runtimeSettings.stop()
-    server.close(() => {
-      process.exit(0)
-    })
-  })
-
-  process.on('SIGTERM', () => {
-    cleanup()
-    runtimeSettings.stop()
-    server.close(() => {
-      process.exit(0)
-    })
+  const shutdown = createGracefulShutdownHandler({
+    server,
+    cleanup,
+    stopRuntime: () => runtimeSettings.stop(),
   })
+  process.once('SIGINT', shutdown)
+  process.once('SIGTERM', shutdown)
 }

packages/cli/tests/commands/serve-shutdown.test.ts

@@ -0,0 +1,96 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+import { createGracefulShutdownHandler } from '../../src/commands/serve-shutdown.js'
+
+function createServer() {
+  return {
+    close: vi.fn(),
+    closeIdleConnections: vi.fn(),
+    closeAllConnections: vi.fn(),
+  }
+}
+
+describe('createGracefulShutdownHandler', () => {
+  beforeEach(() => {
+    vi.useFakeTimers()
+  })
+
+  afterEach(() => {
+    vi.useRealTimers()
+    vi.restoreAllMocks()
+  })
+
+  it('cleans up and exits after the server closes', () => {
+    const server = createServer()
+    server.close.mockImplementation((callback: () => void) => {
+      callback()
+      return server
+    })
+    const cleanup = vi.fn()
+    const stopRuntime = vi.fn()
+    const exit = vi.fn()
+    const log = vi.fn()
+    const shutdown = createGracefulShutdownHandler({
+      server: server as any,
+      cleanup,
+      stopRuntime,
+      exit,
+      log,
+    })
+
+    shutdown()
+
+    expect(log).toHaveBeenCalledWith('\nShutting down...')
+    expect(cleanup).toHaveBeenCalledOnce()
+    expect(stopRuntime).toHaveBeenCalledOnce()
+    expect(server.close).toHaveBeenCalledOnce()
+    expect(server.closeIdleConnections).toHaveBeenCalledOnce()
+    expect(server.closeAllConnections).not.toHaveBeenCalled()
+    expect(exit).toHaveBeenCalledWith(0)
+  })
+
+  it('force-closes active connections when graceful shutdown times out', () => {
+    const server = createServer()
+    server.close.mockReturnValue(server)
+    const exit = vi.fn()
+    const shutdown = createGracefulShutdownHandler({
+      server: server as any,
+      cleanup: vi.fn(),
+      stopRuntime: vi.fn(),
+      exit,
+      log: vi.fn(),
+      timeoutMs: 100,
+    })
+
+    shutdown()
+    vi.advanceTimersByTime(99)
+    expect(exit).not.toHaveBeenCalled()
+
+    vi.advanceTimersByTime(1)
+    expect(server.closeAllConnections).toHaveBeenCalledOnce()
+    expect(exit).toHaveBeenCalledWith(0)
+  })
+
+  it('force-exits on a second signal without calling server.close again', () => {
+    const server = createServer()
+    server.close.mockReturnValue(server)
+    const cleanup = vi.fn()
+    const stopRuntime = vi.fn()
+    const exit = vi.fn()
+    const shutdown = createGracefulShutdownHandler({
+      server: server as any,
+      cleanup,
+      stopRuntime,
+      exit,
+      log: vi.fn(),
+    })
+
+    shutdown()
+    shutdown()
+
+    expect(cleanup).toHaveBeenCalledOnce()
+    expect(stopRuntime).toHaveBeenCalledOnce()
+    expect(server.close).toHaveBeenCalledOnce()
+    expect(server.closeAllConnections).toHaveBeenCalledOnce()
+    expect(exit).toHaveBeenCalledOnce()
+  })
+})