From cad1af1e69fda9cb300705b30a82ee46bb534d59 Mon Sep 17 00:00:00 2001
From: Julian <1041324235@qq.com>
Date: Tue, 1 Sep 2026 09:49:40 +0800
Subject: [PATCH] fix: make serve shutdown reliable
---
.all-contributorsrc | 9 ++
CONTRIBUTING.md | 8 +-
README.md | 3 +
README_zh.md | 3 +
packages/cli/src/commands/serve-shutdown.ts | 57 +++++++++++
packages/cli/src/commands/serve.ts | 22 ++---
.../cli/tests/commands/serve-shutdown.test.ts | 96 +++++++++++++++++++
7 files changed, 180 insertions(+), 18 deletions(-)
create mode 100644 packages/cli/src/commands/serve-shutdown.ts
create mode 100644 packages/cli/tests/commands/serve-shutdown.test.ts
diff --git a/.all-contributorsrc b/.all-contributorsrc
index eec7a9cd..0da604d2 100644
--- a/.all-contributorsrc
+++ b/.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,
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 7e0f4d47..74f5623e 100644
--- a/CONTRIBUTING.md
+++ b/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
diff --git a/README.md b/README.md
index 090574e0..73cede1b 100644
--- a/README.md
+++ b/README.md
@@ -191,6 +191,9 @@ Thanks to these wonderful people ([emoji key](https://allcontributors.org/en/ref
jlxyfll
💻 |
Harlan Hu
💻 |
+
+ 
Mingde
🐛 |
+
diff --git a/README_zh.md b/README_zh.md
index 5c150b3d..323a16c3 100644
--- a/README_zh.md
+++ b/README_zh.md
@@ -191,6 +191,9 @@ pnpm dev
jlxyfll
💻 |
Harlan Hu
💻 |
+
+
Mingde
🐛 |
+
diff --git a/packages/cli/src/commands/serve-shutdown.ts b/packages/cli/src/commands/serve-shutdown.ts
new file mode 100644
index 00000000..00ff96b6
--- /dev/null
+++ b/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 | 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()
+ }
+ }
+}
diff --git a/packages/cli/src/commands/serve.ts b/packages/cli/src/commands/serve.ts
index 96b8cd0e..b3059c14 100644
--- a/packages/cli/src/commands/serve.ts
+++ b/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)
}
diff --git a/packages/cli/tests/commands/serve-shutdown.test.ts b/packages/cli/tests/commands/serve-shutdown.test.ts
new file mode 100644
index 00000000..67b72de0
--- /dev/null
+++ b/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()
+ })
+})