part of: ENG-3645 - flush OTEL spans on SIGTERM/SIGINT so killed OpenCode runs keep child spans

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,13 +57,18 @@ describe("LangfusePlugin", () => {
   const originalEnv = { ...process.env };
 
   beforeEach(() => {
+    __resetSignalFlushForTests();
     mockForceFlush.mockClear();
     mockStart.mockClear();
     mockShutdown.mockClear();
     mockLog.mockClear();
     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();
+    });
+  });
 });

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<unknown>;
+}
+interface ShutdownTarget {
+  shutdown: () => Promise<unknown>;
+}
+
+// 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<void> {
+  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");
+      }
     },
   };
 };