diff --git a/src/index.test.ts b/src/index.test.ts index b7a048e..bfd74ef 100644 --- a/src/index.test.ts +++ b/src/index.test.ts @@ -1,5 +1,17 @@ -import { describe, it, expect, beforeEach, afterEach, mock } from "bun:test"; -import { LangfusePlugin } from "./index"; +import { + describe, + it, + expect, + beforeEach, + afterEach, + afterAll, + mock, +} from "bun:test"; +import { + LangfusePlugin, + flushAndShutdown, + __resetSignalFlushForTests, +} from "./index"; const mockForceFlush = mock(() => Promise.resolve()); const mockStart = mock(() => {}); @@ -45,6 +57,7 @@ describe("LangfusePlugin", () => { const originalEnv = { ...process.env }; beforeEach(() => { + __resetSignalFlushForTests(); mockForceFlush.mockClear(); mockStart.mockClear(); mockShutdown.mockClear(); @@ -52,6 +65,10 @@ describe("LangfusePlugin", () => { capturedNodeSDKOptions = {}; }); + afterAll(() => { + __resetSignalFlushForTests(); + }); + afterEach(() => { process.env = { ...originalEnv }; }); @@ -247,4 +264,67 @@ describe("LangfusePlugin", () => { }); }); }); + + describe("shutdown flush (ENG-3645)", () => { + it("registers SIGTERM and SIGINT handlers on init", async () => { + setupEnv(); + const termBefore = process.listeners("SIGTERM").length; + const intBefore = process.listeners("SIGINT").length; + + await LangfusePlugin(mockPluginInput()); + + expect(process.listeners("SIGTERM").length).toBe(termBefore + 1); + expect(process.listeners("SIGINT").length).toBe(intBefore + 1); + }); + + it("flushes batched spans, then shuts the exporter down, on shutdown", async () => { + setupEnv(); + await LangfusePlugin(mockPluginInput()); + + await flushAndShutdown("SIGTERM"); + + expect(mockForceFlush).toHaveBeenCalled(); + expect(mockShutdown).toHaveBeenCalled(); + }); + + it("flushes and shuts down at most once across repeated triggers", async () => { + setupEnv(); + await LangfusePlugin(mockPluginInput()); + + await flushAndShutdown("SIGTERM"); + await flushAndShutdown("SIGINT"); + await flushAndShutdown("server.instance.disposed"); + + expect(mockForceFlush).toHaveBeenCalledTimes(1); + expect(mockShutdown).toHaveBeenCalledTimes(1); + }); + + it("registers nothing and flushes nothing when credentials are missing", async () => { + delete process.env.LANGFUSE_PUBLIC_KEY; + delete process.env.LANGFUSE_SECRET_KEY; + const termBefore = process.listeners("SIGTERM").length; + + await LangfusePlugin(mockPluginInput()); + + expect(process.listeners("SIGTERM").length).toBe(termBefore); + // Must be safe to call with nothing registered. + await flushAndShutdown("SIGTERM"); + expect(mockForceFlush).not.toHaveBeenCalled(); + }); + + it("does not terminate the process (OpenCode still owns the grace window)", async () => { + setupEnv(); + await LangfusePlugin(mockPluginInput()); + const exitSpy = mock((() => undefined) as (code?: number) => never); + const originalExit = process.exit; + process.exit = exitSpy as unknown as typeof process.exit; + try { + await flushAndShutdown("SIGTERM"); + } finally { + process.exit = originalExit; + } + + expect(exitSpy).not.toHaveBeenCalled(); + }); + }); }); diff --git a/src/index.ts b/src/index.ts index 5ad70a9..a0f725d 100644 --- a/src/index.ts +++ b/src/index.ts @@ -60,6 +60,100 @@ class ParentAwareContextManager extends AsyncLocalStorageContextManager { } } +type LogFn = (level: "info" | "warn" | "error", message: string) => void; + +/** Minimal shape of the OTEL pieces the shutdown routine drives — kept local so + * the routine is trivially unit-testable and independent of SDK internals. */ +interface FlushTarget { + forceFlush: () => Promise; +} +interface ShutdownTarget { + shutdown: () => Promise; +} + +// ENG-3645: the OpenCode readwrite subprocess is torn down by SIGTERM (pod +// soft-timeout) or by client-disconnect cancellation. OpenTelemetry batches +// spans and only exports them on the graceful `session.idle` / +// `server.instance.disposed` lifecycle events — neither of which fires on a +// kill — so every killed run silently loses its child spans. These module-level +// refs plus a single guarded shutdown routine let a process-signal handler +// flush the batch within the OPENCODE_KILL_GRACE_S window the pod already +// reserves after SIGTERM. +let activeProcessor: FlushTarget | undefined; +let activeSdk: ShutdownTarget | undefined; +let activeLog: LogFn | undefined; +let shuttingDown = false; +let signalHandlersInstalled = false; + +const SHUTDOWN_SIGNALS = ["SIGTERM", "SIGINT"] as const; +const signalListeners = new Map< + (typeof SHUTDOWN_SIGNALS)[number], + () => void +>(); + +/** + * Flush any batched spans and shut the exporter down exactly once, whether the + * trigger was a process signal or the `server.instance.disposed` lifecycle + * event. Never throws — a failed flush/shutdown must not mask the original + * teardown. Deliberately does NOT call `process.exit()`: OpenCode owns the + * process and still needs the same grace window to flush its own stdout NDJSON + * (the source of the salvaged diff); exiting here would truncate it. + */ +export async function flushAndShutdown(reason: string): Promise { + if (shuttingDown) return; + shuttingDown = true; + + const safeLog: LogFn = (level, message) => { + try { + activeLog?.(level, message); + } catch { + // The log transport may already be closing during teardown; ignore. + } + }; + + safeLog("info", `Flushing OTEL spans before ${reason} shutdown`); + try { + await activeProcessor?.forceFlush(); + } catch (err) { + safeLog("warn", `forceFlush on ${reason} failed: ${String(err)}`); + } + try { + await activeSdk?.shutdown(); + } catch (err) { + safeLog("warn", `shutdown on ${reason} failed: ${String(err)}`); + } +} + +/** + * Register SIGTERM/SIGINT handlers once per process. Guarded so repeated plugin + * initialisations (or a second plugin instance) cannot stack duplicate + * listeners or double-flush. + */ +function installSignalHandlers(): void { + if (signalHandlersInstalled) return; + signalHandlersInstalled = true; + for (const signal of SHUTDOWN_SIGNALS) { + const listener = () => { + void flushAndShutdown(signal); + }; + signalListeners.set(signal, listener); + process.on(signal, listener); + } +} + +/** Test-only: reset the module-level shutdown state between test cases. */ +export function __resetSignalFlushForTests(): void { + for (const [signal, listener] of signalListeners) { + process.removeListener(signal, listener); + } + signalListeners.clear(); + activeProcessor = undefined; + activeSdk = undefined; + activeLog = undefined; + shuttingDown = false; + signalHandlersInstalled = false; +} + export const LangfusePlugin: Plugin = async ({ client }) => { const publicKey = process.env.LANGFUSE_PUBLIC_KEY; const secretKey = process.env.LANGFUSE_SECRET_KEY; @@ -136,6 +230,13 @@ export const LangfusePlugin: Plugin = async ({ client }) => { log("info", `OTEL tracing initialized → ${baseUrl}`); } + // ENG-3645: expose this run's pipeline to the process-signal shutdown path and + // register the handlers (once) so a killed readwrite still exports its spans. + activeProcessor = processor; + activeSdk = sdk; + activeLog = log; + installSignalHandlers(); + return { config: async (config) => { if (!config.experimental?.openTelemetry) { @@ -151,7 +252,9 @@ export const LangfusePlugin: Plugin = async ({ client }) => { await processor.forceFlush(); } - if (event.type === "server.instance.disposed") await sdk.shutdown(); + if (event.type === "server.instance.disposed") { + await flushAndShutdown("server.instance.disposed"); + } }, }; };