test(cli): cover rejected telemetry shutdown

packages/kilo-telemetry/src/__tests__/telemetry-shutdown.test.ts

@@ -3,16 +3,19 @@
 // top-level `Telemetry` import there resolve the real PostHog into the module
 // cache before the mock is set, making the test rely on bun:test's cache
 // invalidation timing rather than testing the shutdown path directly.
-import { describe, test, expect, mock } from "bun:test"
+import { beforeEach, describe, test, expect, mock } from "bun:test"
+
+const timeout = "Timeout while shutting down PostHog. Some events may not have been sent."
 
 mock.module("posthog-node", () => ({
   PostHog: class {
     async flush() {
       flushCalls += 1
-      await new Promise((resolve) => setTimeout(resolve, 60_000))
+      throw new Error("flush should not be called")
     }
     async shutdown(timeoutMs?: number) {
       shutdownCalls.push(timeoutMs)
+      throw timeout
     }
     optIn() {}
     optOut() {}
@@ -25,6 +28,11 @@ let flushCalls = 0
 const shutdownCalls: Array<number | undefined> = []
 
 describe("Telemetry.shutdown timeout (#9788)", () => {
+  beforeEach(() => {
+    flushCalls = 0
+    shutdownCalls.length = 0
+  })
+
   test("passes timeoutMs through to PostHog.shutdown and skips unbounded explicit flush()", async () => {
     // Reproduces the CLI exit hang reported in #9788: when the PostHog endpoint
     // is unreachable (offline, firewall, DNS adblock resolving the host to
@@ -36,12 +44,9 @@ describe("Telemetry.shutdown timeout (#9788)", () => {
     const { Telemetry } = await import("../telemetry.js")
     const { Client } = await import("../client.js")
     Client.init()
-    const start = Date.now()
-    await Telemetry.shutdown(50)
-    const elapsed = Date.now() - start
+    await expect(Telemetry.shutdown(50)).rejects.toBe(timeout)
 
     expect(flushCalls).toBe(0)
     expect(shutdownCalls).toEqual([50])
-    expect(elapsed).toBeLessThan(1000)
   })
 })

packages/opencode/test/kilocode/cli-shutdown.test.ts

@@ -0,0 +1,127 @@
+import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test"
+
+const calls: string[] = []
+const timeouts: Array<number | undefined> = []
+let err: unknown
+let exit: string | number | null | undefined
+
+mock.module("@opencode-ai/core/global", () => ({
+  Global: { Path: { data: "/tmp/kilo-test" } },
+}))
+
+mock.module("@opencode-ai/core/installation/version", () => ({
+  InstallationBuildKind: "release",
+  InstallationVersion: "test",
+}))
+
+mock.module("@kilocode/kilo-telemetry", () => ({
+  Telemetry: {
+    async init() {},
+    async updateIdentity() {},
+    trackCliStart() {},
+    trackCliExit(code?: number) {
+      calls.push(`track:${code ?? "undefined"}`)
+    },
+    async shutdown(timeout?: number) {
+      calls.push("telemetry")
+      timeouts.push(timeout)
+      if (err) throw err
+    },
+  },
+}))
+
+mock.module("@kilocode/kilo-gateway", () => ({
+  ENV_FEATURE: "KILO_FEATURE",
+  ENV_VERSION: "KILO_VERSION",
+  async migrateLegacyKiloAuth() {},
+}))
+
+mock.module("@/effect/app-runtime", () => ({
+  AppRuntime: { runPromise: (value: unknown) => value },
+}))
+
+mock.module("@/config/config", () => ({
+  Config: { Service: { use: () => ({ experimental: {} }) } },
+}))
+
+mock.module("@/auth", () => ({
+  Auth: { Service: { use: () => undefined } },
+}))
+
+mock.module("@/project/instance-runtime", () => ({
+  InstanceRuntime: {
+    async disposeAllInstances() {
+      calls.push("dispose")
+    },
+  },
+}))
+
+mock.module("@/kilocode/session-export", () => ({
+  SessionExport: {
+    async shutdown() {
+      calls.push("session")
+    },
+  },
+}))
+
+mock.module("@/kilocode/help-command", () => ({
+  createHelpCommand: () => ({ command: "help", handler() {} }),
+}))
+
+for (const path of [
+  "@/kilocode/cli/cmd/console",
+  "@/kilocode/cli/cmd/roll-call",
+  "@/kilocode/cli/cmd/profile",
+  "@/kilocode/cli/cmd/daemon",
+  "@/kilocode/cli/dev-setup",
+  "@/cli/cmd/remote",
+  "@/cli/cmd/config",
+]) {
+  mock.module(path, () => ({
+    KiloConsoleCommand: { command: "console", handler() {} },
+    RollCallCommand: { command: "roll-call", handler() {} },
+    ProfileCommand: { command: "profile", handler() {} },
+    DaemonCommand: { command: "daemon", handler() {} },
+    DevSetupCommand: { command: "dev-setup", handler() {} },
+    DevAliasCommand: { command: "dev-alias", handler() {} },
+    RemoteCommand: { command: "remote", handler() {} },
+    ConfigCommand: { command: "config", handler() {} },
+  }))
+}
+
+describe("KiloCli.shutdown", () => {
+  beforeEach(() => {
+    calls.length = 0
+    timeouts.length = 0
+    err = undefined
+    exit = process.exitCode
+    process.exitCode = undefined
+  })
+
+  afterEach(() => {
+    process.exitCode = exit
+  })
+
+  test("keeps telemetry shutdown timeout best-effort and still disposes instances", async () => {
+    err = "Timeout while shutting down PostHog. Some events may not have been sent."
+    process.exitCode = 0
+    const { KiloCli } = await import("../../src/kilocode/cli/setup")
+
+    await expect(KiloCli.shutdown()).resolves.toBeUndefined()
+
+    expect(timeouts).toEqual([2000])
+    expect(calls).toEqual(["track:0", "session", "telemetry", "dispose"])
+    expect(process.exitCode).toBe(0)
+  })
+
+  test("preserves failing command exit status", async () => {
+    process.exitCode = 1
+    const { KiloCli } = await import("../../src/kilocode/cli/setup")
+
+    await KiloCli.shutdown()
+
+    expect(timeouts).toEqual([2000])
+    expect(calls).toEqual(["track:1", "session", "telemetry", "dispose"])
+    expect(process.exitCode).toBe(1)
+  })
+})