retrace-sdk 0.11.0 → 0.11.2

package/dist/config.d.ts

@@ -11,6 +11,10 @@ export interface Config {
      *  request/response only (recommended for short-lived scripts and serverless — it never
      *  holds an open socket and always surfaces upload errors); "ws" forces WebSocket. */
     transport: "auto" | "ws" | "http";
+    /** Called with a STRUCTURED signal when the server signals credits_exhausted | rate_limited |
+     *  halt | error. Branch on `signal.code`; use `signal.retryable`/`signal.fatal` to decide
+     *  behavior. Defaults to a throttled console warning so signals are never silently dropped. */
+    onError?: (signal: import("./errors.js").RetraceServerSignal) => void;
 }
 export declare function configure(opts: Partial<Config>): Config;
 export declare function requireApiKey(): string;

package/dist/config.js

@@ -17,6 +17,13 @@ export function configure(opts) {
     if (opts.baseUrl && !opts.wsUrl) {
         config.wsUrl = config.baseUrl.replace("https://", "wss://").replace("http://", "ws://");
     }
+    // Eagerly install provider interceptors so clients constructed AFTER configure() are patched.
+    // @google/genai binds generateContent as an own instance property, and our accessor only wraps
+    // instances built after install — so install must precede client construction. Fire-and-forget;
+    // the dynamic import resolves before the first awaited LLM call in any real async flow.
+    if (config.enabled) {
+        void import("./interceptors/install.js").then((m) => m.ensureInterceptorsInstalled()).catch(() => { });
+    }
     return config;
 }
 export function requireApiKey() {

package/dist/errors.d.ts

@@ -14,3 +14,25 @@ export declare class RetraceRateLimitError extends RetraceError {
     retryAfter: number;
     constructor(retryAfter: number);
 }
+/**
+ * Structured server-originated signal handed to `onError`. Actionable WITHOUT string-matching:
+ * branch on `code`, decide retry from `retryable`, decide whether recording is still alive from
+ * `fatal`. (A raw message alone forces the user to string-match — this type exists to avoid that.)
+ */
+export type RetraceSignalCode = "credits_exhausted" | "rate_limited" | "halt" | "error";
+export interface RetraceServerSignal {
+    /** Machine-readable category — branch on THIS, never on `message`. */
+    code: RetraceSignalCode;
+    /** Human-readable detail from the server. */
+    message: string;
+    /** Will retrying / backing off plausibly succeed? rate_limited=true; credits/halt/error=false. */
+    retryable: boolean;
+    /** Did this STOP recording? halt=true (transport closed); others leave recording alive. */
+    fatal: boolean;
+}
+/**
+ * Map a raw server frame to a structured signal. Single source of truth for category + retryable +
+ * fatal, shared by the WS dispatch. Kept here (not inline in the dispatch) so TS and Python classify
+ * identically and the CI gate can assert the mapping.
+ */
+export declare function classifyServerSignal(rawType: string, message: string): RetraceServerSignal;

package/dist/errors.js

@@ -14,3 +14,21 @@ export class RetraceRateLimitError extends RetraceError {
     retryAfter;
     constructor(retryAfter) { super(`Rate limited. Retry after ${retryAfter}s`); this.name = "RetraceRateLimitError"; this.retryAfter = retryAfter; }
 }
+/**
+ * Map a raw server frame to a structured signal. Single source of truth for category + retryable +
+ * fatal, shared by the WS dispatch. Kept here (not inline in the dispatch) so TS and Python classify
+ * identically and the CI gate can assert the mapping.
+ */
+export function classifyServerSignal(rawType, message) {
+    if (rawType === "halt") {
+        return { code: "halt", message: message || "Guardrail triggered", retryable: false, fatal: true };
+    }
+    // rawType === "error" (or anything else carrying an error string)
+    if (message?.includes("limit reached")) {
+        return { code: "credits_exhausted", message, retryable: false, fatal: false };
+    }
+    if (message?.includes("Rate limit")) {
+        return { code: "rate_limited", message, retryable: true, fatal: false };
+    }
+    return { code: "error", message: message || "Server error", retryable: false, fatal: false };
+}

package/dist/index.d.ts

@@ -2,6 +2,7 @@ export { configure, getConfig } from "./config.js";
 export { init, getActiveRecorder, shutdown } from "./init.js";
 export type { InitOptions } from "./init.js";
 export { record, trace, TraceRecorder } from "./recorder.js";
+export { stream } from "./stream.js";
 export { SpanBuilder, TraceBuilder } from "./trace.js";
 export type { SpanData, TraceData } from "./trace.js";
 export { SpanType, TraceStatus } from "./trace.js";
@@ -13,7 +14,7 @@ export { registerResumable, handleResume } from "./resume.js";
 export type { ResumeCommand } from "./resume.js";
 export { isReplaying, consumeCassetteEntry, handleReplay } from "./replay.js";
 export type { CassetteEntry, ReplayCommand } from "./replay.js";
-export { setTraceContext, clearTraceContext, getTraceparent, injectTraceparent, parseTraceparent } from "./traceparent.js";
+export { setTraceContext, clearTraceContext, getTraceparent, injectTraceparent, parseTraceparent, withTraceContext } from "./traceparent.js";
 export { markGolden } from "./golden.js";
 export { createLangChainHandler } from "./adapters/langchain.js";
 export { retraceOnStepFinish, recordVercelStep } from "./adapters/vercel-ai.js";

package/dist/index.js

@@ -1,6 +1,7 @@
 export { configure, getConfig } from "./config.js";
 export { init, getActiveRecorder, shutdown } from "./init.js";
 export { record, trace, TraceRecorder } from "./recorder.js";
+export { stream } from "./stream.js";
 export { SpanBuilder, TraceBuilder } from "./trace.js";
 export { SpanType, TraceStatus } from "./trace.js";
 export { installGeminiInterceptor, uninstallGeminiInterceptor } from "./interceptors/gemini.js";
@@ -9,10 +10,16 @@ export { installAnthropicInterceptor, uninstallAnthropicInterceptor } from "./in
 export { RetraceError, RetraceAuthError, RetraceCreditsExhaustedError, RetraceConnectionError, RetraceRateLimitError } from "./errors.js";
 export { registerResumable, handleResume } from "./resume.js";
 export { isReplaying, consumeCassetteEntry, handleReplay } from "./replay.js";
-export { setTraceContext, clearTraceContext, getTraceparent, injectTraceparent, parseTraceparent } from "./traceparent.js";
+export { setTraceContext, clearTraceContext, getTraceparent, injectTraceparent, parseTraceparent, withTraceContext } from "./traceparent.js";
 export { markGolden } from "./golden.js";
 // Framework adapters (5B) — drop-in instrumentation for LangChain/LangGraph + Vercel AI SDK.
 export { createLangChainHandler } from "./adapters/langchain.js";
 export { retraceOnStepFinish, recordVercelStep } from "./adapters/vercel-ai.js";
+// Patch provider SDKs at import (fire-and-forget; NO top-level await → CJS/bundler-safe). The Gemini
+// interceptor patches RETROACTIVE prototype methods (generateContentInternal/...Stream), so capture
+// works regardless of when the user constructs their client — including module-level clients built
+// in the same tick. This matches the Python SDK: no ordering contract, no ready() footgun.
+import { ensureInterceptorsInstalled } from "./interceptors/install.js";
+void ensureInterceptorsInstalled();
 // v0.5.0
 // trigger

package/dist/init.js

@@ -1,5 +1,6 @@
 import { configure, getConfig, requireApiKey } from "./config.js";
-import { TraceRecorder, flushSharedTransport } from "./recorder.js";
+import { TraceRecorder, drainSharedTransportOnExit } from "./recorder.js";
+import { onProcessExit } from "./transport.js";
 import { TraceStatus } from "./trace.js";
 let ambient = null;
 let exitHooked = false;
@@ -46,28 +47,34 @@ export function init(opts = {}) {
     ambient.start(traceName); // installs the provider interceptors against the ambient recorder
     if (!exitHooked && typeof process !== "undefined") {
         exitHooked = true;
-        const finish = (status) => {
+        const finish = (status, terminatedEarly = false) => {
             const rec = ambient;
             ambient = null;
             try {
-                rec?.end(undefined, status);
+                rec?.end(undefined, status, { terminatedEarly });
             }
             catch { /* best effort on shutdown */ }
         };
-        // On signal-triggered exits, process.exit() would otherwise kill the process before the
-        // final trace_ended is delivered over the network. End the trace, then await a transport
-        // drain (capped by a hard timeout so a hung network can't block shutdown) before exiting.
-        const finishAndExit = (status, code) => {
-            finish(status);
+        // Finish the ambient trace as a pre-exit hook: registerProcessExitFlush (in recorder.ts) runs
+        // this BEFORE draining the transport, so the final trace_ended is in the buffer for the
+        // HTTP one-shot — and signal ownership (sole-listener-flush-then-exit vs user-owns-exit) is
+        // handled there in one place, not duplicated here.
+        //
+        // Only a graceful exit (event loop emptied = the program finished its work) produces a CLEAN
+        // terminal. Signal/uncaught exits interrupted the run mid-flight, so the synthesized terminal
+        // is marked terminated_early — otherwise we'd manufacture a clean-looking terminal for a
+        // truncated run and defeat the replay-guard's no-terminal rule.
+        onProcessExit((reason) => finish(reason === "uncaught" ? TraceStatus.FAILED : TraceStatus.COMPLETED, reason !== "graceful"));
+        // uncaughtException is not covered by registerProcessExitFlush (it's status-specific and must
+        // exit non-zero): finish FAILED + terminated_early, drain best-effort within a hard cap, then exit.
+        process.once("uncaughtException", (err) => {
+            console.error(err);
+            finish(TraceStatus.FAILED, true);
             void Promise.race([
-                flushSharedTransport().catch(() => { }),
+                drainSharedTransportOnExit(1500).catch(() => { }),
                 new Promise((r) => setTimeout(r, 3000)),
-            ]).then(() => process.exit(code));
-        };
-        process.once("beforeExit", () => finish(TraceStatus.COMPLETED));
-        process.once("SIGINT", () => finishAndExit(TraceStatus.COMPLETED, 130));
-        process.once("SIGTERM", () => finishAndExit(TraceStatus.COMPLETED, 143));
-        process.once("uncaughtException", (err) => { console.error(err); finishAndExit(TraceStatus.FAILED, 1); });
+            ]).then(() => process.exit(1));
+        });
     }
     return ambient;
 }

package/dist/interceptors/_dispatch.d.ts

@@ -0,0 +1,25 @@
+import type { SpanData } from "../trace.js";
+export type SpanCallback = (span: SpanData) => void;
+/** A recorder-side sink for the two-phase streaming-span lifecycle (open at invocation, finalize
+ *  once at clean-drain / break / error / trace-end / exit). Routed like SpanCallback. */
+export interface OpenSpanSink {
+    registerOpenSpan(spanId: string, finalize: (reason: "complete" | "partial") => void): void;
+    unregisterOpenSpan(spanId: string): void;
+}
+/** Stable interceptor callback — routes an intercepted span to the recorder active in this context. */
+export declare function dispatchInterceptedSpan(span: SpanData): void;
+/** Capture the span sink active in THIS context (synchronously, at invocation) so a deferred
+ *  finalizer can emit to the right recorder even when later called from a context where the ALS
+ *  store is absent (the AFC layer's .return(), trace-end, exit-flush). */
+export declare function captureActiveSpanEmit(): SpanCallback | null;
+/** Register an open streaming span's finalizer with the active recorder (two-phase, model (b)). */
+export declare function dispatchRegisterOpenSpan(spanId: string, finalize: (reason: "complete" | "partial") => void): void;
+/** Drop an open span's finalizer once it has been finalized in-band. */
+export declare function dispatchUnregisterOpenSpan(spanId: string): void;
+/** Run `fn` with `cb` as the active intercepted-span handler for its async context (fully isolated). */
+export declare function runWithActiveRecorder<T>(cb: SpanCallback, fn: () => T, sink?: OpenSpanSink): T;
+/** Set the imperative-API fallback handler. Returns the PREVIOUS handler so callers can restore
+ * it (e.g. a nested trace() must not wipe the ambient init() fallback). Pass null to clear. */
+export declare function setActiveRecorderFallback(cb: SpanCallback | null, sink?: OpenSpanSink | null): SpanCallback | null;
+/** The current imperative-API open-span sink (so a nested trace can save/restore it). */
+export declare function currentFallbackSink(): OpenSpanSink | null;

package/dist/interceptors/_dispatch.js

@@ -0,0 +1,57 @@
+/**
+ * Context-isolated routing for auto-instrumented spans (mirrors the Python `_dispatch.py`
+ * ContextVar dispatcher).
+ *
+ * The interceptors (openai/anthropic/gemini) are patched globally and invoke ONE stable
+ * dispatcher. The target recorder for the current async context lives in an AsyncLocalStorage
+ * store, so concurrent traces on a server each resolve their own recorder — instead of a
+ * module-global, last-writer-wins callback that cross-routed intercepted spans to whichever
+ * trace was created most recently.
+ */
+import { AsyncLocalStorage } from "async_hooks";
+const activeRecorder = new AsyncLocalStorage();
+const activeOpenSink = new AsyncLocalStorage();
+// Fallback for the imperative record()/start()/end() API used outside a runWithActiveRecorder
+// scope. Last-writer-wins (documented limitation for purely imperative concurrent traces).
+let fallbackCb = null;
+let fallbackSink = null;
+/** Stable interceptor callback — routes an intercepted span to the recorder active in this context. */
+export function dispatchInterceptedSpan(span) {
+    const cb = activeRecorder.getStore() ?? fallbackCb;
+    cb?.(span);
+}
+/** Capture the span sink active in THIS context (synchronously, at invocation) so a deferred
+ *  finalizer can emit to the right recorder even when later called from a context where the ALS
+ *  store is absent (the AFC layer's .return(), trace-end, exit-flush). */
+export function captureActiveSpanEmit() {
+    return activeRecorder.getStore() ?? fallbackCb;
+}
+/** Register an open streaming span's finalizer with the active recorder (two-phase, model (b)). */
+export function dispatchRegisterOpenSpan(spanId, finalize) {
+    const sink = activeOpenSink.getStore() ?? fallbackSink;
+    sink?.registerOpenSpan(spanId, finalize);
+}
+/** Drop an open span's finalizer once it has been finalized in-band. */
+export function dispatchUnregisterOpenSpan(spanId) {
+    const sink = activeOpenSink.getStore() ?? fallbackSink;
+    sink?.unregisterOpenSpan(spanId);
+}
+/** Run `fn` with `cb` as the active intercepted-span handler for its async context (fully isolated). */
+export function runWithActiveRecorder(cb, fn, sink) {
+    if (sink)
+        return activeRecorder.run(cb, () => activeOpenSink.run(sink, fn));
+    return activeRecorder.run(cb, fn);
+}
+/** Set the imperative-API fallback handler. Returns the PREVIOUS handler so callers can restore
+ * it (e.g. a nested trace() must not wipe the ambient init() fallback). Pass null to clear. */
+export function setActiveRecorderFallback(cb, sink) {
+    const prev = fallbackCb;
+    fallbackCb = cb;
+    if (sink !== undefined)
+        fallbackSink = sink;
+    return prev;
+}
+/** The current imperative-API open-span sink (so a nested trace can save/restore it). */
+export function currentFallbackSink() {
+    return fallbackSink;
+}

package/dist/interceptors/gemini.d.ts

@@ -1,3 +1,3 @@
 import { SpanData } from "../trace.js";
-export declare function installGeminiInterceptor(onSpan: (span: SpanData) => void): void;
+export declare function installGeminiInterceptor(onSpan: (span: SpanData) => void): Promise<void>;
 export declare function uninstallGeminiInterceptor(): void;

package/dist/interceptors/gemini.js

@@ -1,5 +1,6 @@
 import { SpanType } from "../trace.js";
 import { genId, nowIso, truncateJson } from "../utils.js";
+import { dispatchRegisterOpenSpan, dispatchUnregisterOpenSpan, captureActiveSpanEmit } from "./_dispatch.js";
 import { emitGeminiToolCalls, emitGeminiToolResults, resetToolResultDedup, extractToolSchemas, extractSamplingParams } from "./tool-spans.js";
 const PRICING = {
     "gemini-3.1-flash-lite": [0.10, 0.40],
@@ -17,93 +18,233 @@ function calcCost(model, inputTokens, outputTokens) {
     const p = PRICING[model] || [0, 0];
     return (inputTokens * p[0] + outputTokens * p[1]) / 1_000_000;
 }
-let originalGenerateContent = null;
-let installed = false;
 let onSpanCallback = null;
-export function installGeminiInterceptor(onSpan) {
-    if (installed) {
-        onSpanCallback = onSpan;
-        resetToolResultDedup();
-        return;
-    }
-    onSpanCallback = onSpan;
-    resetToolResultDedup();
-    import("@google/genai").then((genaiMod) => {
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+let patchedProto = null;
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+let origGenerate = null;
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+let origStream = null;
+let installPromise = null;
+// Wrap a single `generateContent` implementation (the per-instance bound method). Returns a
+// function with identical signature that records a span around the original call.
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+function wrapGenerate(original) {
+    return async function (...args) {
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        const mod = genaiMod;
-        const modelsProto = mod?.Models?.prototype || mod?.default?.Models?.prototype;
-        if (!modelsProto?.generateContent)
-            return;
-        originalGenerateContent = modelsProto.generateContent;
-        modelsProto.generateContent = async function (...args) {
+        const opts = args[0] || {};
+        const model = opts.model || "unknown";
+        const contents = opts.contents;
+        const toolSchemas = extractToolSchemas("gemini", opts.config?.tools);
+        const sampling = extractSamplingParams("gemini", opts);
+        const spanMeta = { ...(toolSchemas ? { tool_schemas: toolSchemas } : {}), ...(sampling ? { sampling } : {}) };
+        const spanId = genId();
+        const startedAt = nowIso();
+        const startMs = Date.now();
+        const { isReplaying, consumeCassetteEntry } = await import("../replay.js");
+        if (isReplaying()) {
+            const entry = consumeCassetteEntry("retrace.ai.generate", "llm_call");
+            if (entry) {
+                return { text: entry.output || "", usageMetadata: { promptTokenCount: 0, candidatesTokenCount: 0 }, candidates: [] };
+            }
+        }
+        try {
+            const result = await original.apply(this, args);
+            const durationMs = Date.now() - startMs;
             // eslint-disable-next-line @typescript-eslint/no-explicit-any
-            const opts = args[0] || {};
-            const model = opts.model || "unknown";
-            const contents = opts.contents;
-            const toolSchemas = extractToolSchemas("gemini", opts.config?.tools);
-            const sampling = extractSamplingParams("gemini", opts);
-            const spanMeta = { ...(toolSchemas ? { tool_schemas: toolSchemas } : {}), ...(sampling ? { sampling } : {}) };
-            const spanId = genId();
-            const startedAt = nowIso();
-            const startMs = Date.now();
-            // Replay mode — return mocked response from cassette
-            const { isReplaying, consumeCassetteEntry } = await import("../replay.js");
-            if (isReplaying()) {
-                const entry = consumeCassetteEntry("retrace.ai.generate", "llm_call");
-                if (entry) {
-                    return { text: entry.output || "", usageMetadata: { promptTokenCount: 0, candidatesTokenCount: 0 }, candidates: [] };
-                }
+            const res = result;
+            const inputTokens = res?.usageMetadata?.promptTokenCount || 0;
+            const outputTokens = res?.usageMetadata?.candidatesTokenCount || 0;
+            const fnNames = res?.candidates?.[0]?.content?.parts?.filter((p) => p.functionCall).map((p) => p.functionCall.name).join(", ");
+            const outputText = res?.text ?? (fnNames ? `[function_call: ${fnNames}]` : "");
+            const span = {
+                id: spanId, trace_id: "", parent_id: null,
+                span_type: SpanType.LLM_CALL, name: "retrace.ai.generate", model,
+                input: truncateJson(contents), output: truncateJson(outputText),
+                input_tokens: inputTokens, output_tokens: outputTokens,
+                cost: calcCost(model, inputTokens, outputTokens),
+                duration_ms: durationMs, started_at: startedAt, ended_at: nowIso(),
+                ...(Object.keys(spanMeta).length ? { metadata: spanMeta } : {}),
+            };
+            onSpanCallback?.(span);
+            if (onSpanCallback) {
+                emitGeminiToolResults(contents, onSpanCallback);
+                emitGeminiToolCalls(res?.candidates, spanId, model, onSpanCallback);
+            }
+            return result;
+        }
+        catch (err) {
+            onSpanCallback?.({
+                id: spanId, trace_id: "", parent_id: null,
+                span_type: SpanType.LLM_CALL, name: "retrace.ai.generate", model,
+                input: truncateJson(contents), started_at: startedAt, ended_at: nowIso(),
+                duration_ms: Date.now() - startMs,
+                error: err instanceof Error ? err.message : String(err),
+            });
+            throw err;
+        }
+    };
+}
+// Wrap a single `generateContentStream` implementation. Accumulates text + usage across chunks and
+// emits the span when the stream completes OR is abandoned/errors (via finally).
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+function wrapStream(original) {
+    return async function (...args) {
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        const opts = args[0] || {};
+        const model = opts.model || "unknown";
+        const contents = opts.contents;
+        const toolSchemas = extractToolSchemas("gemini", opts.config?.tools);
+        const sampling = extractSamplingParams("gemini", opts);
+        const spanMeta = { streaming: true, ...(toolSchemas ? { tool_schemas: toolSchemas } : {}), ...(sampling ? { sampling } : {}) };
+        const spanId = genId();
+        const startedAt = nowIso();
+        const startMs = Date.now();
+        const { isReplaying, consumeCassetteEntry } = await import("../replay.js");
+        if (isReplaying()) {
+            const entry = consumeCassetteEntry("retrace.ai.generate", "llm_call");
+            if (entry) {
+                const text = entry.output || "";
+                async function* mockStream() { yield { text, usageMetadata: { promptTokenCount: 0, candidatesTokenCount: 0 } }; }
+                return mockStream();
+            }
+        }
+        let iterable;
+        try {
+            iterable = await original.apply(this, args);
+        }
+        catch (err) {
+            onSpanCallback?.({
+                id: spanId, trace_id: "", parent_id: null,
+                span_type: SpanType.LLM_CALL, name: "retrace.ai.generate", model,
+                input: truncateJson(contents), started_at: startedAt, ended_at: nowIso(),
+                duration_ms: Date.now() - startMs,
+                error: err instanceof Error ? err.message : String(err),
+            });
+            throw err;
+        }
+        const chunks = [];
+        let inputTokens = 0;
+        let outputTokens = 0;
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        let lastCandidates;
+        let streamError;
+        let emitted = false;
+        let sawFunctionCall = false;
+        // Bind the emit target NOW (caller's context) so finalize routes to the right recorder even
+        // when called later from the AFC's .return(), trace-end, or exit-flush.
+        const boundEmit = captureActiveSpanEmit() ?? onSpanCallback;
+        // Streaming-A, model (b): the span is OPENED at invocation (registered with the active
+        // recorder, synchronously in the CALLER's async context — registering lazily inside the
+        // generator fails because the AFC layer pulls it in a different context where the ALS sink is
+        // absent) and EMITTED EXACTLY ONCE at finalization — never silent. `capture_complete` is true
+        // ONLY on an observed clean full-drain; false for early-break, error, or a finalize forced by
+        // trace-end/exit (the AFC layer abandons this generator on a normal full consume — which is
+        // why we no longer rely on observing `done`, and dropped the read-ahead hack). A
+        // capture_complete:false span is NOT byte-replayable.
+        const finalize = (reason) => {
+            if (emitted)
+                return;
+            emitted = true;
+            dispatchUnregisterOpenSpan(spanId);
+            // capture_complete:true means fully-captured AND byte-replay-eligible — clean observed drain
+            // AND no function call (function-call output isn't captured as text and the AFC path may
+            // re-issue). Everything else is partial.
+            const complete = reason === "complete" && !sawFunctionCall;
+            const fnNames = lastCandidates?.[0]?.content?.parts
+                ?.filter((p) => p.functionCall)
+                .map((p) => p.functionCall.name).join(", ");
+            const outputText = chunks.length ? chunks.join("") : (fnNames ? `[function_call: ${fnNames}]` : "");
+            const span = {
+                id: spanId, trace_id: "", parent_id: null,
+                span_type: SpanType.LLM_CALL, name: "retrace.ai.generate", model,
+                input: truncateJson(contents), output: truncateJson(outputText),
+                input_tokens: inputTokens, output_tokens: outputTokens,
+                cost: calcCost(model, inputTokens, outputTokens),
+                duration_ms: Date.now() - startMs, started_at: startedAt, ended_at: nowIso(),
+                ...(streamError ? { error: streamError } : {}),
+                metadata: { ...spanMeta, capture_complete: complete },
+            };
+            boundEmit?.(span);
+            if (boundEmit) {
+                emitGeminiToolResults(contents, boundEmit);
+                emitGeminiToolCalls(lastCandidates, spanId, model, boundEmit);
             }
+        };
+        // Register in the caller's context so trace-end / exit-flush finalizes us partial if the
+        // consumer (AFC) abandons the generator mid-drain without ever reaching `done` / `.return()`.
+        dispatchRegisterOpenSpan(spanId, finalize);
+        async function* wrapped() {
             try {
-                const result = await originalGenerateContent.apply(this, args);
-                const durationMs = Date.now() - startMs;
-                // eslint-disable-next-line @typescript-eslint/no-explicit-any
-                const res = result;
-                const inputTokens = res?.usageMetadata?.promptTokenCount || 0;
-                const outputTokens = res?.usageMetadata?.candidatesTokenCount || 0;
-                const outputText = res?.text ?? (res?.candidates?.[0]?.content?.parts?.filter((p) => p.functionCall).map((p) => p.functionCall.name).join(", ") ? `[function_call: ${res.candidates[0].content.parts.filter((p) => p.functionCall).map((p) => p.functionCall.name).join(", ")}]` : "");
-                const span = {
-                    id: spanId, trace_id: "", parent_id: null,
-                    span_type: SpanType.LLM_CALL, name: "retrace.ai.generate", model,
-                    input: truncateJson(contents), output: truncateJson(outputText),
-                    input_tokens: inputTokens, output_tokens: outputTokens,
-                    cost: calcCost(model, inputTokens, outputTokens),
-                    duration_ms: durationMs, started_at: startedAt, ended_at: nowIso(),
-                    ...(Object.keys(spanMeta).length ? { metadata: spanMeta } : {}),
-                };
-                onSpanCallback?.(span);
-                // Auto-capture tool usage (functionCall parts in response, functionResponse in input).
-                if (onSpanCallback) {
-                    emitGeminiToolResults(contents, onSpanCallback);
-                    emitGeminiToolCalls(res?.candidates, spanId, model, onSpanCallback);
+                for await (const chunk of iterable) { // eslint-disable-line @typescript-eslint/no-explicit-any
+                    if (typeof chunk?.text === "string")
+                        chunks.push(chunk.text);
+                    if (chunk?.candidates)
+                        lastCandidates = chunk.candidates;
+                    // A function-call stream is NOT byte-replay-eligible: its output isn't plain text (so
+                    // chunks.join misses it) and the AFC layer may re-issue. Mark it so finalize never
+                    // stamps capture_complete:true even on a clean drain.
+                    if (chunk?.candidates?.[0]?.content?.parts?.some((p) => p.functionCall)) {
+                        sawFunctionCall = true;
+                    }
+                    if (chunk?.usageMetadata) {
+                        inputTokens = chunk.usageMetadata.promptTokenCount || inputTokens;
+                        outputTokens = chunk.usageMetadata.candidatesTokenCount || outputTokens;
+                    }
+                    yield chunk;
                 }
-                return result;
+                finalize("complete"); // observed clean full-drain (e.g. via retrace.stream helper)
             }
             catch (err) {
-                const span = {
-                    id: spanId, trace_id: "", parent_id: null,
-                    span_type: SpanType.LLM_CALL, name: "retrace.ai.generate", model,
-                    input: truncateJson(contents), started_at: startedAt, ended_at: nowIso(),
-                    duration_ms: Date.now() - startMs,
-                    error: err instanceof Error ? err.message : String(err),
-                };
-                onSpanCallback?.(span);
+                streamError = err instanceof Error ? err.message : String(err);
+                finalize("partial");
                 throw err;
             }
-        };
-        installed = true;
-    }).catch(() => { });
+            finally {
+                finalize("partial"); // early-break (consumer .return()) / no clean drain observed
+            }
+        }
+        return wrapped();
+    };
 }
-export function uninstallGeminiInterceptor() {
-    if (!installed || !originalGenerateContent)
-        return;
-    import("@google/genai").then((genaiMod) => {
+// @google/genai binds the PUBLIC `generateContent`/`generateContentStream` as own bound instance
+// properties (not on the prototype), so patching the prototype's public method is a no-op. However,
+// the public methods delegate to `generateContentInternal` / `generateContentStreamInternal`, which
+// ARE regular methods on `Models.prototype`. Patching those is RETROACTIVE to every instance
+// regardless of construction order (mirroring the Python SDK's class-method patch) — so no
+// install-before-construction requirement, no race, and no `ready()` escape hatch is needed.
+export function installGeminiInterceptor(onSpan) {
+    onSpanCallback = onSpan;
+    resetToolResultDedup();
+    if (installPromise)
+        return installPromise; // synchronous dedupe — prevents the double-wrap race
+    installPromise = import("@google/genai").then((genaiMod) => {
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
         const mod = genaiMod;
-        const modelsProto = mod?.Models?.prototype || mod?.default?.Models?.prototype;
-        if (modelsProto)
-            modelsProto.generateContent = originalGenerateContent;
+        const Models = mod?.Models || mod?.default?.Models;
+        const proto = Models?.prototype;
+        if (!proto || typeof proto.generateContentInternal !== "function")
+            return;
+        patchedProto = proto;
+        origGenerate = proto.generateContentInternal;
+        proto.generateContentInternal = wrapGenerate(origGenerate);
+        if (typeof proto.generateContentStreamInternal === "function") {
+            origStream = proto.generateContentStreamInternal;
+            proto.generateContentStreamInternal = wrapStream(origStream);
+        }
     }).catch(() => { });
-    installed = false;
+    return installPromise;
+}
+export function uninstallGeminiInterceptor() {
+    if (patchedProto) {
+        if (origGenerate)
+            patchedProto.generateContentInternal = origGenerate;
+        if (origStream)
+            patchedProto.generateContentStreamInternal = origStream;
+    }
+    installPromise = null;
     onSpanCallback = null;
+    patchedProto = null;
+    origGenerate = null;
+    origStream = null;
 }

package/dist/interceptors/install.d.ts

@@ -0,0 +1 @@
+export declare function ensureInterceptorsInstalled(): Promise<void>;

package/dist/interceptors/install.js

@@ -0,0 +1,21 @@
+import { dispatchInterceptedSpan } from "./_dispatch.js";
+import { installGeminiInterceptor } from "./gemini.js";
+import { installOpenAIInterceptor } from "./openai.js";
+import { installAnthropicInterceptor } from "./anthropic.js";
+/**
+ * Install ALL provider interceptors against the single stable dispatcher. Idempotent and memoized.
+ * Called at import time and from configure()/init(). The Gemini interceptor patches retroactive
+ * prototype methods, so the (async) install landing slightly after import is fine — capture works
+ * regardless of when the user constructs their client.
+ */
+let _installPromise = null;
+export function ensureInterceptorsInstalled() {
+    if (_installPromise)
+        return _installPromise;
+    _installPromise = Promise.all([
+        Promise.resolve(installGeminiInterceptor(dispatchInterceptedSpan)),
+        Promise.resolve(installOpenAIInterceptor(dispatchInterceptedSpan)),
+        Promise.resolve(installAnthropicInterceptor(dispatchInterceptedSpan)),
+    ]).then(() => { });
+    return _installPromise;
+}