@traceai/codex-sdk 0.1.0 → 0.1.1

package/src/instrumentation.ts

@@ -5,6 +5,7 @@
  */
 import { TracerProvider, diag, trace } from "@opentelemetry/api";
 import { FITracer, TraceConfigOptions } from "@traceai/fi-core";
+import { patchCodex } from "./codexWrapper";
 import { ContentPolicy } from "./spans";
 import { TurnTracingContext, patchThread } from "./threadWrapper";
 import { CodexModuleLike, ThreadLike } from "./types";
@@ -69,9 +70,20 @@ export function resolveContentPolicy(traceConfig?: TraceConfigOptions): ContentP
 export class CodexSDKInstrumentation {
   private tracerProvider?: TracerProvider;
   private readonly traceConfig?: TraceConfigOptions;
-  /** Undo functions for the class patch and wrapCodex() factories (never one per thread). */
+  /** Undo functions for class patches. */
   private readonly undo: Array<() => void> = [];
-  /** False after unpatch(): instance patches stay in place but pass events through untraced. */
+  /** Client undo functions hold only WeakRefs; collected clients release their bookkeeping. */
+  private readonly clientUndo = new Set<() => void>();
+  private readonly clientFinalizer = new FinalizationRegistry<() => void>((undo) => {
+    // Deactivate the owner too: saved factories must remain inert after collection.
+    undo();
+    this.clientUndo.delete(undo);
+  });
+  /** Thread undo functions hold only WeakRefs; collected threads release their bookkeeping. */
+  private readonly threadUndo = new Set<() => void>();
+  private readonly threadFinalizer = new FinalizationRegistry<() => void>((undo) => {
+    this.threadUndo.delete(undo);
+  });
   private enabled = true;
 
   constructor(config: CodexSDKInstrumentationConfig = {}) {
@@ -106,13 +118,16 @@ export class CodexSDKInstrumentation {
   }
 
   /**
-   * Patch one Thread instance. Returns the same instance. No reference to the
-   * thread is kept: after `unpatch()` the patched method passes events through
-   * untraced instead of being removed.
+   * Patch one Thread instance. Returns the same instance. Only a weak reference
+   * is kept, so `unpatch()` can restore live threads without retaining dropped ones.
    */
   wrapThread<T extends object>(thread: T): T {
     this.enabled = true;
-    patchThread(thread as ThreadLike, () => this.tracingContext());
+    const undo = patchThread(thread as ThreadLike, () => this.tracingContext());
+    if (undo) {
+      this.threadUndo.add(undo);
+      this.threadFinalizer.register(thread, undo, undo);
+    }
     return thread;
   }
 
@@ -120,27 +135,10 @@ export class CodexSDKInstrumentation {
   wrapCodex<T extends object>(codex: T): T {
     if (!codex || typeof codex !== "object") return codex;
     this.enabled = true;
-    const target = codex as Record<string, unknown>;
-    // Held weakly: a wrapped client the app drops can be collected before unpatch().
-    const ref = new WeakRef(target);
-    for (const method of ["startThread", "resumeThread"] as const) {
-      const original = target[method] as ((...args: unknown[]) => object) | undefined;
-      if (typeof original !== "function") continue;
-      const hadOwn = Object.prototype.hasOwnProperty.call(codex, method);
-      const self = this;
-      Object.defineProperty(codex, method, {
-        value: function tracedThreadFactory(this: unknown, ...args: unknown[]) {
-          return self.wrapThread(original.apply(this, args));
-        },
-        writable: true,
-        configurable: true,
-      });
-      this.undo.push(() => {
-        const live = ref.deref();
-        if (!live) return;
-        if (hadOwn) Object.defineProperty(live, method, { value: original, writable: true, configurable: true });
-        else delete live[method];
-      });
+    const undo = patchCodex(codex as Record<string | symbol, unknown>, (thread) => this.wrapThread(thread));
+    if (undo) {
+      this.clientUndo.add(undo);
+      this.clientFinalizer.register(codex, undo, undo);
     }
     return codex;
   }
@@ -148,6 +146,20 @@ export class CodexSDKInstrumentation {
   /** Remove every patch this instrumentation applied. */
   unpatch(): void {
     this.enabled = false;
+    for (const undo of this.threadUndo) {
+      this.threadFinalizer.unregister(undo);
+      undo();
+    }
+    this.threadUndo.clear();
+    for (const undo of this.clientUndo) {
+      this.clientFinalizer.unregister(undo);
+      try {
+        undo();
+      } catch (error) {
+        diag.debug(`${INSTRUMENTATION_NAME}: unpatch failed: ${error}`);
+      }
+    }
+    this.clientUndo.clear();
     while (this.undo.length) {
       const fn = this.undo.pop();
       try {

package/src/spans.ts

@@ -515,7 +515,7 @@ export class TurnTracer {
     if (turnUsage) {
       if (this.threadId) usageBaselines.set(this.threadId, cumulative);
     } else {
-      // The total went down (a reset, or an SDK that reports per-turn usage): use it as
+      // The total went down (a reset within the supported cumulative-usage SDKs): use it as
       // reported, and keep the larger baseline so a later turn is not counted twice.
       turnUsage = cumulative;
       attributes[C.USAGE_NON_MONOTONIC] = true;

package/src/threadWrapper.ts

@@ -16,6 +16,7 @@ import { ContentPolicy, TurnTracer, clockMs } from "./spans";
 import { InputLike, ThreadEventLike, ThreadLike, TurnOptionsLike } from "./types";
 
 const PATCHED = Symbol.for("@traceai/codex-sdk/patched");
+const PATCH_RECORD = Symbol.for("@traceai/codex-sdk/patch-record");
 const LOG = "@traceai/codex-sdk";
 
 export interface TurnTracingContext {
@@ -26,16 +27,37 @@ export interface TurnTracingContext {
 type GetContext = () => TurnTracingContext | undefined;
 
 interface PatchRecord {
-  target: ThreadLike;
   method: "runStreamedInternal" | "runStreamed";
-  original: unknown;
-  hadOwn: boolean;
+  originalDescriptor: PropertyDescriptor | undefined;
+  wrapper: unknown;
+  owner: symbol;
 }
 
 export function isPatched(target: unknown): boolean {
   return !!target && (target as Record<symbol, unknown>)[PATCHED] === true;
 }
 
+/** The undo closure retains neither the target nor its possibly bound original method. */
+function weakUndo(ref: WeakRef<ThreadLike>, owner: symbol): () => void {
+  return () => {
+    const target = ref.deref();
+    if (!target) return;
+    const record = Object.getOwnPropertyDescriptor(target, PATCH_RECORD)?.value as PatchRecord | undefined;
+    if (record?.owner !== owner) return;
+    try {
+      // Do not clobber a method that the application replaced after our patch.
+      if (Object.getOwnPropertyDescriptor(target, record.method)?.value === record.wrapper) {
+        if (record.originalDescriptor) Object.defineProperty(target, record.method, record.originalDescriptor);
+        else delete (target as Record<string, unknown>)[record.method];
+      }
+      delete (target as Record<symbol, unknown>)[PATCHED];
+      delete (target as Record<symbol, unknown>)[PATCH_RECORD];
+    } catch (error) {
+      diag.debug(`${LOG}: unpatch failed: ${error}`);
+    }
+  };
+}
+
 /**
  * Patch a Thread prototype or a single Thread instance. Returns an undo
  * function, or `undefined` when the target was already patched or has no
@@ -44,10 +66,11 @@ export function isPatched(target: unknown): boolean {
 export function patchThread(target: ThreadLike, getContext: GetContext): (() => void) | undefined {
   if (!target || typeof target !== "object" || isPatched(target)) return undefined;
 
-  let record: PatchRecord | undefined;
+  const owner = Symbol();
+  let record: PatchRecord;
   if (typeof target.runStreamedInternal === "function") {
     const original = target.runStreamedInternal;
-    record = { target, method: "runStreamedInternal", original, hadOwn: Object.prototype.hasOwnProperty.call(target, "runStreamedInternal") };
+    const originalDescriptor = Object.getOwnPropertyDescriptor(target, "runStreamedInternal");
     const traced = function tracedRunStreamedInternal(this: ThreadLike, input: InputLike, turnOptions?: TurnOptionsLike) {
       const startTimeMs = clockMs();
       const parentContext = otelContext.active();
@@ -55,10 +78,11 @@ export function patchThread(target: ThreadLike, getContext: GetContext): (() =>
       const generator = original.apply(this, arguments as unknown as [InputLike, TurnOptionsLike?]);
       return traceEvents(generator, this, input, turnOptions, parentContext, startTimeMs, getContext);
     };
+    record = { method: "runStreamedInternal", originalDescriptor, wrapper: traced, owner };
     defineMethod(target, "runStreamedInternal", traced);
   } else if (typeof target.runStreamed === "function") {
     const original = target.runStreamed;
-    record = { target, method: "runStreamed", original, hadOwn: Object.prototype.hasOwnProperty.call(target, "runStreamed") };
+    const originalDescriptor = Object.getOwnPropertyDescriptor(target, "runStreamed");
     diag.warn(`${LOG}: Thread.runStreamedInternal not found; tracing runStreamed() only, run() may be untraced`);
     const traced = async function tracedRunStreamed(this: ThreadLike, input: InputLike, turnOptions?: TurnOptionsLike) {
       const startTimeMs = clockMs();
@@ -71,22 +95,19 @@ export function patchThread(target: ThreadLike, getContext: GetContext): (() =>
         events: traceEvents(result.events, this, input, turnOptions, parentContext, startTimeMs, getContext),
       };
     };
+    record = { method: "runStreamed", originalDescriptor, wrapper: traced, owner };
     defineMethod(target, "runStreamed", traced);
   } else {
     return undefined;
   }
 
-  Object.defineProperty(target, PATCHED, { value: true, configurable: true });
-  const done = record;
-  return () => {
-    try {
-      if (done.hadOwn) defineMethod(done.target, done.method, done.original);
-      else delete (done.target as Record<string, unknown>)[done.method];
-      delete (done.target as Record<symbol, unknown>)[PATCHED];
-    } catch (error) {
-      diag.debug(`${LOG}: unpatch failed: ${error}`);
-    }
-  };
+  // Keep the public marker boolean so older installed copies also recognize an
+  // active patch. Store ownership separately, only on the weakly tracked target.
+  Object.defineProperties(target, {
+    [PATCHED]: { value: true, configurable: true },
+    [PATCH_RECORD]: { value: record, configurable: true },
+  });
+  return weakUndo(new WeakRef(target), owner);
 }
 
 function defineMethod(target: object, name: string, value: unknown): void {

package/src/version.ts

@@ -1 +1 @@
-export const VERSION = "0.1.0";
+export const VERSION = "0.1.1"; // x-release-please-version