retrace-sdk 0.1.3

package/dist/transport.js

@@ -0,0 +1,124 @@
+import WebSocket from "ws";
+import { getConfig } from "./config.js";
+export class WSTransport {
+    ws = null;
+    connected = false;
+    backoff = 1000;
+    queue = [];
+    connect() {
+        const cfg = getConfig();
+        const url = `${cfg.wsUrl}/ws/v1/stream`;
+        this.ws = new WebSocket(url);
+        this.ws.on("open", () => {
+            this.ws.send(JSON.stringify({ type: "auth", api_key: cfg.apiKey }));
+        });
+        this.ws.on("message", (raw) => {
+            const msg = JSON.parse(raw.toString());
+            if (msg.type === "auth_ok") {
+                this.connected = true;
+                this.backoff = 1000;
+                this.flushQueue();
+            }
+            else if (msg.type === "ping") {
+                this.ws?.send(JSON.stringify({ type: "pong" }));
+            }
+        });
+        this.ws.on("close", () => {
+            this.connected = false;
+            this.ws = null;
+            setTimeout(() => this.reconnect(), this.backoff);
+            this.backoff = Math.min(this.backoff * 2, 30000);
+        });
+        this.ws.on("error", () => {
+            this.ws?.close();
+        });
+    }
+    reconnect() {
+        if (!this.connected && !this.ws)
+            this.connect();
+    }
+    flushQueue() {
+        while (this.queue.length && this.connected) {
+            this.ws?.send(this.queue.shift());
+        }
+    }
+    send(eventType, data) {
+        const msg = JSON.stringify({ type: eventType, data });
+        if (this.connected && this.ws?.readyState === WebSocket.OPEN) {
+            this.ws.send(msg);
+        }
+        else {
+            this.queue.push(msg);
+            if (!this.ws)
+                this.connect();
+        }
+    }
+    close() {
+        if (this.ws) {
+            this.ws.close();
+            this.ws = null;
+        }
+        this.connected = false;
+    }
+}
+export class HTTPTransport {
+    traceData = null;
+    spans = [];
+    send(eventType, data) {
+        if (eventType === "trace_started") {
+            this.traceData = data;
+        }
+        else if (eventType === "span_started" || eventType === "span_ended") {
+            this.spans.push({ ...data, _event: eventType });
+        }
+        else if (eventType === "trace_ended") {
+            if (this.traceData)
+                Object.assign(this.traceData, data);
+            this.flush();
+        }
+    }
+    flush() {
+        if (!this.traceData)
+            return;
+        const cfg = getConfig();
+        const url = `${cfg.baseUrl}/api/v1/traces`;
+        const body = { ...this.traceData, spans: this.buildSpans() };
+        fetch(url, {
+            method: "POST",
+            headers: { "x-retrace-key": cfg.apiKey, "Content-Type": "application/json" },
+            body: JSON.stringify(body),
+        }).catch(() => { });
+        this.traceData = null;
+        this.spans = [];
+    }
+    buildSpans() {
+        const merged = new Map();
+        for (const ev of this.spans) {
+            const { _event, ...rest } = ev;
+            const id = rest.id;
+            if (_event === "span_started") {
+                merged.set(id, rest);
+            }
+            else if (_event === "span_ended" && merged.has(id)) {
+                Object.assign(merged.get(id), rest);
+            }
+        }
+        return [...merged.values()];
+    }
+    close() {
+        this.flush();
+    }
+}
+export function createTransport(mode = "auto") {
+    if (mode === "http")
+        return new HTTPTransport();
+    if (mode === "ws")
+        return new WSTransport();
+    // Auto: try WS
+    try {
+        return new WSTransport();
+    }
+    catch {
+        return new HTTPTransport();
+    }
+}

package/dist/utils.d.ts

@@ -0,0 +1,4 @@
+export declare function genId(): string;
+export declare function nowIso(): string;
+export declare function utcNow(): Date;
+export declare function truncateJson(obj: unknown, maxBytes?: number): unknown;

package/dist/utils.js

@@ -0,0 +1,21 @@
+import { randomUUID } from "crypto";
+export function genId() {
+    return randomUUID();
+}
+export function nowIso() {
+    return new Date().toISOString();
+}
+export function utcNow() {
+    return new Date();
+}
+export function truncateJson(obj, maxBytes = 10240) {
+    try {
+        const s = JSON.stringify(obj);
+        if (Buffer.byteLength(s) <= maxBytes)
+            return obj;
+        return JSON.parse(Buffer.from(s).subarray(0, maxBytes).toString());
+    }
+    catch {
+        return String(obj).slice(0, maxBytes);
+    }
+}

package/package.json

@@ -0,0 +1,36 @@
+{
+  "name": "retrace-sdk",
+  "version": "0.1.3",
+  "description": "The execution replay engine for AI agents. Record, replay, fork, and share agent executions.",
+  "type": "module",
+  "main": "dist/index.js",
+  "types": "dist/index.d.ts",
+  "exports": {
+    ".": {
+      "import": "./dist/index.js",
+      "types": "./dist/index.d.ts"
+    }
+  },
+  "scripts": {
+    "dev": "tsc --watch",
+    "build": "tsc",
+    "lint": "eslint src/",
+    "test": "vitest run"
+  },
+  "dependencies": {
+    "ws": "8.20.0"
+  },
+  "peerDependencies": {
+    "@google/genai": ">=1.52.0"
+  },
+  "peerDependenciesMeta": {
+    "@google/genai": {
+      "optional": true
+    }
+  },
+  "devDependencies": {
+    "typescript": "6.0.3",
+    "@types/node": "22.15.3",
+    "@types/ws": "8.18.1"
+  }
+}

package/src/config.test.ts

@@ -0,0 +1,16 @@
+import { describe, it, expect } from "vitest";
+import { configure, getConfig } from "./config.js";
+
+describe("SDK config", () => {
+  it("has defaults", () => {
+    const c = getConfig();
+    expect(c.baseUrl).toBe("http://localhost:3001");
+    expect(c.enabled).toBe(true);
+  });
+  it("configures custom values", () => {
+    configure({ apiKey: "rt_live_test", baseUrl: "http://custom:3001" });
+    const c = getConfig();
+    expect(c.apiKey).toBe("rt_live_test");
+    expect(c.baseUrl).toBe("http://custom:3001");
+  });
+});

package/src/config.ts

@@ -0,0 +1,31 @@
+export interface Config {
+  apiKey: string;
+  baseUrl: string;
+  wsUrl: string;
+  projectId: string | undefined;
+  enabled: boolean;
+}
+
+const config: Config = {
+  apiKey: process.env.RETRACE_API_KEY || "",
+  baseUrl: process.env.RETRACE_BASE_URL || "http://localhost:3001",
+  wsUrl: "",
+  projectId: process.env.RETRACE_PROJECT_ID || undefined,
+  enabled: !["false", "0"].includes((process.env.RETRACE_ENABLED || "true").toLowerCase()),
+};
+config.wsUrl = config.baseUrl.replace("https://", "wss://").replace("http://", "ws://");
+
+export function configure(opts: Partial<Config>): Config {
+  if (opts.apiKey && !opts.apiKey.startsWith("rt_live_")) {
+    console.warn("[retrace] API key does not start with 'rt_live_'. This may be invalid.");
+  }
+  Object.assign(config, opts);
+  if (opts.baseUrl && !opts.wsUrl) {
+    config.wsUrl = config.baseUrl.replace("https://", "wss://").replace("http://", "ws://");
+  }
+  return config;
+}
+
+export function getConfig(): Config {
+  return config;
+}

package/src/index.ts

@@ -0,0 +1,6 @@
+export { configure, getConfig } from "./config.js";
+export { record, trace, TraceRecorder } from "./recorder.js";
+export { SpanBuilder, TraceBuilder } from "./trace.js";
+export type { SpanData, TraceData } from "./trace.js";
+export { SpanType, TraceStatus } from "./trace.js";
+export { installGeminiInterceptor, uninstallGeminiInterceptor } from "./interceptors/gemini.js";

package/src/interceptors/gemini.ts

@@ -0,0 +1,86 @@
+import { SpanData, SpanType } from "../trace.js";
+import { genId, nowIso, truncateJson } from "../utils.js";
+
+const PRICING: Record<string, [number, number]> = {
+  "gemini-3.1-pro-preview": [2.0, 12.0],
+  "gemini-2.5-pro": [1.25, 10.0],
+  "gemini-2.5-flash": [0.15, 0.60],
+  "gemini-2.0-flash": [0.10, 0.40],
+};
+
+function calcCost(model: string, inputTokens: number, outputTokens: number): number {
+  const p = PRICING[model] || [0, 0];
+  return (inputTokens * p[0] + outputTokens * p[1]) / 1_000_000;
+}
+
+ 
+let originalGenerateContent: ((...args: unknown[]) => unknown) | null = null;
+let installed = false;
+let onSpanCallback: ((span: SpanData) => void) | null = null;
+
+export function installGeminiInterceptor(onSpan: (span: SpanData) => void) {
+  if (installed) { onSpanCallback = onSpan; return; }
+  onSpanCallback = onSpan;
+
+  import("@google/genai").then((genaiMod) => {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    const mod = genaiMod as Record<string, any>;
+    const modelsProto = mod?.Models?.prototype || mod?.default?.Models?.prototype;
+    if (!modelsProto?.generateContent) return;
+
+    originalGenerateContent = modelsProto.generateContent;
+
+    modelsProto.generateContent = async function (...args: unknown[]) {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const opts = (args[0] as Record<string, any>) || {};
+      const model = (opts.model as string) || "unknown";
+      const contents = opts.contents;
+      const spanId = genId();
+      const startedAt = nowIso();
+      const startMs = Date.now();
+
+      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 as any;
+        const inputTokens = res?.usageMetadata?.promptTokenCount || 0;
+        const outputTokens = res?.usageMetadata?.candidatesTokenCount || 0;
+
+        const span: SpanData = {
+          id: spanId, trace_id: "", parent_id: null,
+          span_type: SpanType.LLM_CALL, name: "retrace.ai.generate", model,
+          input: truncateJson(contents), output: truncateJson(res?.text || ""),
+          input_tokens: inputTokens, output_tokens: outputTokens,
+          cost: calcCost(model, inputTokens, outputTokens),
+          duration_ms: durationMs, started_at: startedAt, ended_at: nowIso(),
+        };
+        onSpanCallback?.(span);
+        return result;
+      } catch (err) {
+        const span: SpanData = {
+          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);
+        throw err;
+      }
+    };
+    installed = true;
+  }).catch(() => { /* @google/genai not installed — skip */ });
+}
+
+export function uninstallGeminiInterceptor() {
+  if (!installed || !originalGenerateContent) return;
+  import("@google/genai").then((genaiMod) => {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    const mod = genaiMod as Record<string, any>;
+    const modelsProto = mod?.Models?.prototype || mod?.default?.Models?.prototype;
+    if (modelsProto) modelsProto.generateContent = originalGenerateContent;
+  }).catch(() => {});
+  installed = false;
+  onSpanCallback = null;
+}

package/src/interceptors/index.ts

@@ -0,0 +1 @@
+export { installGeminiInterceptor, uninstallGeminiInterceptor } from "./gemini.js";

package/src/recorder.ts

@@ -0,0 +1,136 @@
+import { getConfig } from "./config.js";
+import { SpanBuilder, SpanData, SpanType, TraceBuilder, TraceStatus } from "./trace.js";
+import { createTransport, Transport } from "./transport.js";
+import { installGeminiInterceptor } from "./interceptors/gemini.js";
+
+export interface RecordOptions {
+  name?: string;
+  input?: unknown;
+  metadata?: Record<string, unknown>;
+}
+
+export class TraceRecorder {
+  private builder: TraceBuilder;
+  private transport: Transport;
+  private interceptorsInstalled = false;
+  output: unknown = undefined;
+
+  constructor(opts?: RecordOptions) {
+    this.builder = new TraceBuilder();
+    this.transport = createTransport();
+    const cfg = getConfig();
+    if (cfg.projectId) this.builder.setProjectId(cfg.projectId);
+    if (opts?.metadata) this.builder.setMetadata(opts.metadata);
+    if (opts?.name || opts?.input) {
+      this.builder.start(opts.name, opts.input);
+    }
+  }
+
+  get traceId() { return this.builder.id; }
+
+  start(name?: string, input?: unknown): this {
+    this.builder.start(name, input);
+    this.installInterceptors();
+    this.transport.send("trace_started", this.builder.toDict() as unknown as Record<string, unknown>);
+    return this;
+  }
+
+  end(output?: unknown, status: TraceStatus = TraceStatus.COMPLETED) {
+    if (output !== undefined) this.output = output;
+    const data = this.builder.end(this.output, status);
+    this.transport.send("trace_ended", {
+      id: data.id,
+      ended_at: data.ended_at!,
+      output: data.output,
+      status: data.status,
+      total_tokens: data.total_tokens,
+      total_cost: data.total_cost,
+    });
+    this.transport.close();
+  }
+
+  addSpan(span: SpanData) {
+    span.trace_id = this.builder.id;
+    this.builder.addSpan(span);
+    this.transport.send("span_started", span as unknown as Record<string, unknown>);
+    if (span.ended_at) {
+      this.transport.send("span_ended", {
+        id: span.id,
+        ended_at: span.ended_at,
+        output: span.output,
+        output_tokens: span.output_tokens,
+        cost: span.cost,
+        error: span.error,
+      });
+    }
+  }
+
+  startSpan(name: string, spanType: SpanType = SpanType.LLM_CALL, input?: unknown, model?: string, parentId?: string): SpanBuilder {
+    const sb = new SpanBuilder(name, spanType).start();
+    sb.setTraceId(this.builder.id);
+    if (input !== undefined) sb.setInput(input);
+    if (model) sb.setModel(model);
+    if (parentId) sb.setParentId(parentId);
+    this.transport.send("span_started", sb.toData() as unknown as Record<string, unknown>);
+    return sb;
+  }
+
+  endSpan(spanBuilder: SpanBuilder, output?: unknown, error?: string) {
+    const span = spanBuilder.end(output, error);
+    span.trace_id = this.builder.id;
+    this.builder.addSpan(span);
+    this.transport.send("span_ended", {
+      id: span.id,
+      ended_at: span.ended_at,
+      output: span.output,
+      output_tokens: span.output_tokens,
+      cost: span.cost,
+      error: span.error,
+    });
+  }
+
+  private installInterceptors() {
+    if (this.interceptorsInstalled) return;
+    installGeminiInterceptor((span) => this.addSpan(span));
+    this.interceptorsInstalled = true;
+  }
+}
+
+export function record(opts?: RecordOptions): TraceRecorder {
+  const cfg = getConfig();
+  if (!cfg.enabled) {
+    // Return a no-op recorder
+    return new TraceRecorder(opts);
+  }
+  return new TraceRecorder(opts);
+}
+
+export function trace<T>(fn: (...args: unknown[]) => T, opts?: RecordOptions): (...args: unknown[]) => T {
+  const cfg = getConfig();
+  return (...args: unknown[]): T => {
+    if (!cfg.enabled) return fn(...args);
+
+    const recorder = new TraceRecorder({
+      name: opts?.name || fn.name || "anonymous",
+      input: opts?.input ?? args,
+      metadata: opts?.metadata,
+    });
+    recorder.start(opts?.name || fn.name || "anonymous", opts?.input ?? args);
+
+    try {
+      const result = fn(...args);
+      // Handle async functions
+      if (result && typeof (result as { then?: unknown }).then === "function") {
+        return (result as unknown as Promise<unknown>).then(
+          (resolved) => { recorder.end(resolved, TraceStatus.COMPLETED); return resolved; },
+          (err) => { recorder.end(undefined, TraceStatus.FAILED); throw err; }
+        ) as unknown as T;
+      }
+      recorder.end(result, TraceStatus.COMPLETED);
+      return result;
+    } catch (err) {
+      recorder.end(undefined, TraceStatus.FAILED);
+      throw err;
+    }
+  };
+}

package/src/trace.ts

@@ -0,0 +1,135 @@
+import { genId, nowIso, utcNow } from "./utils.js";
+
+export enum SpanType {
+  LLM_CALL = "llm_call",
+  TOOL_CALL = "tool_call",
+  TOOL_RESULT = "tool_result",
+  REASONING = "reasoning",
+  ACTION = "action",
+  ERROR = "error",
+  FORK_POINT = "fork_point",
+}
+
+export enum TraceStatus {
+  RUNNING = "running",
+  COMPLETED = "completed",
+  FAILED = "failed",
+}
+
+export interface SpanData {
+  id: string;
+  trace_id: string;
+  parent_id: string | null;
+  span_type: string;
+  name: string;
+  model?: string;
+  input?: unknown;
+  output?: unknown;
+  input_tokens?: number;
+  output_tokens?: number;
+  cost?: number;
+  duration_ms?: number;
+  metadata?: Record<string, unknown>;
+  started_at: string;
+  ended_at?: string;
+  error?: string;
+}
+
+export interface TraceData {
+  id: string;
+  name?: string;
+  input?: unknown;
+  output?: unknown;
+  status: string;
+  total_tokens: number;
+  total_cost: number;
+  total_duration_ms: number;
+  metadata?: Record<string, unknown>;
+  started_at: string;
+  ended_at?: string;
+  spans?: SpanData[];
+  project_id?: string;
+}
+
+export class SpanBuilder {
+  private data: Partial<SpanData> & { id: string; span_type: string; name: string };
+  private _startTime?: Date;
+
+  constructor(name: string, spanType: SpanType) {
+    this.data = { id: genId(), span_type: spanType, name, trace_id: "", parent_id: null, started_at: nowIso() };
+  }
+
+  setModel(model: string) { this.data.model = model; return this; }
+  setInput(input: unknown) { this.data.input = input; return this; }
+  setOutput(output: unknown) { this.data.output = output; return this; }
+  setParentId(id: string) { this.data.parent_id = id; return this; }
+  setTraceId(id: string) { this.data.trace_id = id; return this; }
+  setMetadata(m: Record<string, unknown>) { this.data.metadata = m; return this; }
+
+  start(): this {
+    this._startTime = utcNow();
+    this.data.started_at = this._startTime.toISOString();
+    return this;
+  }
+
+  end(output?: unknown, error?: string): SpanData {
+    const now = utcNow();
+    this.data.ended_at = now.toISOString();
+    if (output !== undefined) this.data.output = output;
+    if (error) this.data.error = error;
+    if (this._startTime) {
+      this.data.duration_ms = now.getTime() - this._startTime.getTime();
+    }
+    return this.data as SpanData;
+  }
+
+  get id() { return this.data.id; }
+  toData(): SpanData { return this.data as SpanData; }
+}
+
+export class TraceBuilder {
+  private data: TraceData;
+  private _startTime?: Date;
+
+  constructor() {
+    this.data = {
+      id: genId(),
+      status: TraceStatus.RUNNING,
+      total_tokens: 0,
+      total_cost: 0,
+      total_duration_ms: 0,
+      started_at: nowIso(),
+      spans: [],
+    };
+  }
+
+  start(name?: string, input?: unknown): this {
+    this._startTime = utcNow();
+    this.data.started_at = this._startTime.toISOString();
+    if (name) this.data.name = name;
+    if (input !== undefined) this.data.input = input;
+    return this;
+  }
+
+  end(output?: unknown, status: TraceStatus = TraceStatus.COMPLETED): TraceData {
+    const now = utcNow();
+    this.data.ended_at = now.toISOString();
+    this.data.status = status;
+    if (output !== undefined) this.data.output = output;
+    if (this._startTime) {
+      this.data.total_duration_ms = now.getTime() - this._startTime.getTime();
+    }
+    return this.data;
+  }
+
+  addSpan(span: SpanData) {
+    this.data.spans!.push(span);
+    this.data.total_tokens += (span.input_tokens || 0) + (span.output_tokens || 0);
+    this.data.total_cost += span.cost || 0;
+  }
+
+  get id() { return this.data.id; }
+  setProjectId(id: string) { this.data.project_id = id; }
+  setMetadata(m: Record<string, unknown>) { this.data.metadata = m; }
+  toDict(): TraceData { return this.data; }
+}