elasticdash-sdk 0.2.0

package/src/interceptors/workflow-ai.ts

@@ -0,0 +1,419 @@
+import { getCaptureContext } from '../capture/recorder.js'
+import { rawDateNow } from './side-effects.js'
+import { getHttpRunContext, getHttpFrozenEvent, getHttpPromptMock, getHttpUserPromptMock, getHttpAIMock, pushTelemetryEvent, tryAutoInitHttpContext, getObservabilityContext } from './telemetry-push.js'
+import { resolveAIMock, resolvePromptMock, resolveUserPromptMock } from '../internals/mock-resolver.js'
+import { consumeCapturedLLMRequest } from './ai-interceptor.js'
+import { getEdReplayContext, replayCall } from '../ci/replay.js'
+import type { WorkflowEvent } from '../capture/event.js'
+
+/**
+ * Tracks whether we are currently inside a `wrapAI` execution.
+ * When active, `ai-interceptor` should skip recording to avoid duplicates.
+ * Uses a counter (not boolean) to handle nested wrapAI calls correctly.
+ */
+const AI_WRAPPER_KEY = '__elasticdash_ai_wrapper_depth__'
+const g = globalThis as Record<string, unknown>
+if (g[AI_WRAPPER_KEY] == null) g[AI_WRAPPER_KEY] = 0
+
+function enterAIWrapper(): void {
+  g[AI_WRAPPER_KEY] = ((g[AI_WRAPPER_KEY] as number) ?? 0) + 1
+}
+
+function leaveAIWrapper(): void {
+  g[AI_WRAPPER_KEY] = Math.max(0, ((g[AI_WRAPPER_KEY] as number) ?? 0) - 1)
+}
+
+type UsageInfo = { inputTokens?: number; outputTokens?: number; totalTokens?: number }
+
+function extractUsage(output: unknown): UsageInfo | undefined {
+  if (!output || typeof output !== 'object') return undefined
+  const o = output as Record<string, unknown>
+  if (o.usage && typeof o.usage === 'object') {
+    const u = o.usage as Record<string, number>
+    // Anthropic SDK: input_tokens / output_tokens
+    if (u.input_tokens != null || u.output_tokens != null) {
+      return {
+        inputTokens: u.input_tokens,
+        outputTokens: u.output_tokens,
+        totalTokens: (u.input_tokens ?? 0) + (u.output_tokens ?? 0),
+      }
+    }
+    // OpenAI SDK: prompt_tokens / completion_tokens
+    if (u.prompt_tokens != null || u.completion_tokens != null) {
+      return {
+        inputTokens: u.prompt_tokens,
+        outputTokens: u.completion_tokens,
+        totalTokens: u.total_tokens,
+      }
+    }
+    // Vercel AI SDK: { inputTokens, outputTokens }
+    if (u.inputTokens != null || u.outputTokens != null) {
+      return {
+        inputTokens: u.inputTokens,
+        outputTokens: u.outputTokens,
+        totalTokens: (u.inputTokens ?? 0) + (u.outputTokens ?? 0),
+      }
+    }
+  }
+  // Gemini SDK: usageMetadata
+  if (o.usageMetadata && typeof o.usageMetadata === 'object') {
+    const u = o.usageMetadata as Record<string, number>
+    return {
+      inputTokens: u.promptTokenCount,
+      outputTokens: u.candidatesTokenCount,
+      totalTokens: u.totalTokenCount,
+    }
+  }
+  // Flat token fields (e.g. { tokens: N } or { outputTokens: N, inputTokens: N })
+  if (typeof o.tokens === 'number' || typeof o.outputTokens === 'number') {
+    return {
+      inputTokens: typeof o.inputTokens === 'number' ? o.inputTokens : undefined,
+      outputTokens: typeof o.outputTokens === 'number' ? o.outputTokens : (typeof o.tokens === 'number' ? o.tokens : undefined),
+      totalTokens: (typeof o.inputTokens === 'number' ? o.inputTokens : 0) + (typeof o.outputTokens === 'number' ? o.outputTokens : (typeof o.tokens === 'number' ? o.tokens : 0)),
+    }
+  }
+  return undefined
+}
+
+/**
+ * After callFn returns, consume any captured LLM request from ai-interceptor
+ * and merge it into the event input so the recorded event contains the actual
+ * HTTP payload sent to the LLM (messages, model, parameters).
+ * Also returns captured usage when the app's output doesn't include it.
+ */
+function enrichFromLLMCapture(
+  input: unknown,
+  appUsage: UsageInfo | undefined,
+  fallbackModel?: string,
+  fallbackProvider?: string,
+): { input: unknown; usage: UsageInfo | undefined } {
+  const captured = consumeCapturedLLMRequest()
+  if (captured) {
+    const enrichedInput = (input && typeof input === 'object')
+      ? { ...(input as Record<string, unknown>), llmRequest: captured.body, promptSnippet: captured.promptSnippet }
+      : { originalInput: input, llmRequest: captured.body, promptSnippet: captured.promptSnippet }
+    const usage = appUsage ?? captured.usage
+    return { input: enrichedInput, usage }
+  }
+
+  // Interceptor capture failed (e.g. OpenAI SDK uses native fetch not globalThis.fetch).
+  // Fall back to model/provider from wrapAI options or function arguments.
+  const resolvedModel = fallbackModel
+    || (input && typeof input === 'object' ? (input as Record<string, unknown>).model as string | undefined : undefined)
+  const resolvedProvider = fallbackProvider
+    || (input && typeof input === 'object' ? (input as Record<string, unknown>).provider as string | undefined : undefined)
+
+  if (resolvedModel || resolvedProvider) {
+    const enrichedInput = (input && typeof input === 'object')
+      ? { ...(input as Record<string, unknown>), llmRequest: { model: resolvedModel, provider: resolvedProvider } }
+      : { originalInput: input, llmRequest: { model: resolvedModel, provider: resolvedProvider } }
+    return { input: enrichedInput, usage: appUsage }
+  }
+
+  return { input, usage: appUsage }
+}
+
+function isReadableStream(v: unknown): v is ReadableStream<Uint8Array> {
+  return (
+    typeof v === 'object' &&
+    v !== null &&
+    typeof (v as ReadableStream).getReader === 'function' &&
+    typeof (v as ReadableStream).tee === 'function'
+  )
+}
+
+function isAsyncIterable(v: unknown): v is AsyncIterable<unknown> {
+  return typeof v === 'object' && v !== null && Symbol.asyncIterator in (v as object)
+}
+
+async function bufferReadableStream(stream: ReadableStream<Uint8Array>): Promise<string> {
+  const decoder = new TextDecoder()
+  const reader = stream.getReader()
+  let raw = ''
+  try {
+    for (;;) {
+      const { done, value } = await reader.read()
+      if (done) break
+      raw += decoder.decode(value, { stream: true })
+    }
+  } finally {
+    reader.releaseLock()
+  }
+  return raw
+}
+
+function reconstructStream(raw: string): ReadableStream<Uint8Array> {
+  const encoder = new TextEncoder()
+  return new ReadableStream<Uint8Array>({
+    start(ctrl) {
+      ctrl.enqueue(encoder.encode(raw))
+      ctrl.close()
+    },
+  })
+}
+
+/** Wraps an AsyncIterable so chunks are collected while the caller iterates */
+function wrapAsyncIterable<T>(
+  source: AsyncIterable<T>,
+  onComplete: (chunks: T[]) => void,
+): AsyncIterable<T> {
+  return {
+    [Symbol.asyncIterator]() {
+      const iter = source[Symbol.asyncIterator]()
+      const collected: T[] = []
+      return {
+        async next() {
+          const result = await iter.next()
+          if (!result.done) {
+            collected.push(result.value)
+          } else {
+            onComplete(collected)
+          }
+          return result
+        },
+        async return(value?: unknown) {
+          onComplete(collected)
+          return iter.return ? iter.return(value) : { done: true as const, value: undefined }
+        },
+      }
+    },
+  }
+}
+
+export function wrapAI<Args extends unknown[], R>(
+  modelName: string,
+  callFn: (...args: Args) => Promise<R>,
+  _options?: { model?: string; provider?: string },
+): (...args: Args) => Promise<R> {
+  return async (...args: Args): Promise<R> => {
+    // Phase 3 fixture replay: short-circuit if ed-runner replay is active
+    const edReplay = getEdReplayContext()
+    if (edReplay) {
+      const input = args.length === 1 ? args[0] : args
+      const { output } = replayCall(edReplay, 'ai_call', modelName, input)
+      return output as R
+    }
+
+    await tryAutoInitHttpContext()
+    const ctx = getCaptureContext()
+    const httpCtx = getHttpRunContext()
+    const obsCtx = getObservabilityContext()
+
+    if (!ctx && !httpCtx && !obsCtx) return callFn(...args)
+
+    enterAIWrapper()
+    try {
+    const start = rawDateNow()
+
+    // Observability-only mode: record and push, no mocks/replay
+    if (!ctx && !httpCtx && obsCtx) {
+      const id = obsCtx.nextId()
+      const input = args.length === 1 ? args[0] : args
+      try {
+        const output = await callFn(...args)
+        const enriched = enrichFromLLMCapture(input, extractUsage(output), _options?.model, _options?.provider)
+
+        if (isReadableStream(output)) {
+          const [streamForCaller, streamForRecorder] = output.tee()
+          bufferReadableStream(streamForRecorder).then((rawText) => {
+            const durationMs = rawDateNow() - start
+            pushTelemetryEvent({ id, type: 'ai', name: modelName, input: enriched.input, output: null, streamed: true, streamRaw: rawText, timestamp: start, durationMs, ...(enriched.usage ? { usage: enriched.usage } : {}) })
+          }).catch(() => {
+            const durationMs = rawDateNow() - start
+            pushTelemetryEvent({ id, type: 'ai', name: modelName, input: enriched.input, output: null, streamed: true, streamRaw: '', timestamp: start, durationMs, ...(enriched.usage ? { usage: enriched.usage } : {}) })
+          })
+          return streamForCaller as unknown as R
+        }
+
+        if (isAsyncIterable(output)) {
+          return wrapAsyncIterable(output, (chunks) => {
+            const durationMs = rawDateNow() - start
+            const rawText = chunks.map((c) => (typeof c === 'string' ? c : JSON.stringify(c))).join('')
+            pushTelemetryEvent({ id, type: 'ai', name: modelName, input: enriched.input, output: null, streamed: true, streamRaw: rawText, timestamp: start, durationMs, ...(enriched.usage ? { usage: enriched.usage } : {}) })
+          }) as unknown as R
+        }
+
+        const durationMs = rawDateNow() - start
+        const event: WorkflowEvent = {
+          id, type: 'ai', name: modelName, input: enriched.input, output,
+          timestamp: start, durationMs,
+          ...(enriched.usage ? { usage: enriched.usage } : {}),
+        }
+        pushTelemetryEvent(event)
+        return output
+      } catch (e) {
+        const enriched = enrichFromLLMCapture(input, undefined, _options?.model, _options?.provider)
+        const durationMs = rawDateNow() - start
+        pushTelemetryEvent({
+          id, type: 'ai', name: modelName, input: enriched.input,
+          output: { error: String(e) }, timestamp: start, durationMs,
+        })
+        throw e
+      }
+    }
+
+    if (ctx) {
+      const { recorder, replay } = ctx
+      const id = recorder.nextId()
+
+      if (replay.shouldReplay(id)) {
+        const historical = replay.getRecordedEvent(id)
+        if (historical?.streamed === true) {
+          const raw = typeof historical.streamRaw === 'string' ? historical.streamRaw : ''
+          return reconstructStream(raw) as unknown as R
+        }
+        return replay.getRecordedResult(id) as R
+      }
+
+      // Check AI mock (output mock — skip real call, return recorded result)
+      const aiMock = resolveAIMock(modelName)
+      if (aiMock.mocked) {
+        const input = args.length === 1 ? args[0] : args
+        const event: WorkflowEvent = {
+          id, type: 'ai', name: modelName, input,
+          output: aiMock.result, timestamp: start, durationMs: 0,
+        }
+        recorder.record(event)
+        if (httpCtx) pushTelemetryEvent(event)
+        return aiMock.result as R
+      }
+
+      // Check prompt mocks (system + user prompt replacement — call real LLM with modified prompts)
+      const rawInput = args.length === 1 ? args[0] : args
+      const promptModifiedInput = resolvePromptMock(rawInput)
+      const userPromptModifiedInput = resolveUserPromptMock(promptModifiedInput !== undefined ? promptModifiedInput : rawInput)
+      const modifiedInput = userPromptModifiedInput !== undefined ? userPromptModifiedInput : promptModifiedInput
+      const effectiveArgs: Args = modifiedInput !== undefined ? [modifiedInput] as unknown as Args : args
+      const input = modifiedInput !== undefined ? modifiedInput : rawInput
+
+      try {
+        const output = await callFn(...effectiveArgs)
+
+        if (isReadableStream(output)) {
+          const [streamForCaller, streamForRecorder] = output.tee()
+          bufferReadableStream(streamForRecorder).then((rawText) => {
+            const durationMs = rawDateNow() - start
+            const event: WorkflowEvent = { id, type: 'ai', name: modelName, input, output: null, streamed: true, streamRaw: rawText, timestamp: start, durationMs }
+            recorder.record(event)
+            if (httpCtx) pushTelemetryEvent(event)
+          }).catch(() => {
+            const durationMs = rawDateNow() - start
+            const event: WorkflowEvent = { id, type: 'ai', name: modelName, input, output: null, streamed: true, streamRaw: '', timestamp: start, durationMs }
+            recorder.record(event)
+            if (httpCtx) pushTelemetryEvent(event)
+          })
+          return streamForCaller as unknown as R
+        }
+
+        if (isAsyncIterable(output)) {
+          return wrapAsyncIterable(output, (chunks) => {
+            const durationMs = rawDateNow() - start
+            const rawText = chunks.map((c) => (typeof c === 'string' ? c : JSON.stringify(c))).join('')
+            const event: WorkflowEvent = { id, type: 'ai', name: modelName, input, output: null, streamed: true, streamRaw: rawText, timestamp: start, durationMs }
+            recorder.record(event)
+            if (httpCtx) pushTelemetryEvent(event)
+          }) as unknown as R
+        }
+
+        const durationMs = rawDateNow() - start
+        const usage = extractUsage(output)
+        const event: WorkflowEvent = {
+          id, type: 'ai', name: modelName, input, output,
+          timestamp: start, durationMs,
+          ...(usage ? { usage } : {}),
+        }
+        recorder.record(event)
+        if (httpCtx) pushTelemetryEvent(event)
+        return output
+      } catch (e) {
+        const durationMs = rawDateNow() - start
+        const event: WorkflowEvent = {
+          id, type: 'ai', name: modelName, input,
+          output: { error: String(e) }, timestamp: start, durationMs,
+        }
+        recorder.record(event)
+        if (httpCtx) pushTelemetryEvent(event)
+        throw e
+      }
+    }
+
+    // HTTP mode only (no capture context)
+    const id = httpCtx!.nextId()
+
+    // Replay frozen step: push historical event so dashboard trace stays complete
+    const frozen = getHttpFrozenEvent(id)
+    if (frozen) {
+      pushTelemetryEvent(frozen)
+      if (frozen.streamed === true) {
+        const raw = typeof frozen.streamRaw === 'string' ? frozen.streamRaw : ''
+        return reconstructStream(raw) as unknown as R
+      }
+      return frozen.output as R
+    }
+
+    // Check AI output mock (skip real call, return mocked result)
+    const aiMock = getHttpAIMock(modelName)
+    if (aiMock.mocked) {
+      const input = args.length === 1 ? args[0] : args
+      const event: WorkflowEvent = {
+        id, type: 'ai', name: modelName, input,
+        output: aiMock.result, timestamp: start, durationMs: 0,
+      }
+      pushTelemetryEvent(event)
+      return aiMock.result as R
+    }
+
+    // Check prompt mocks (system + user prompt replacement in HTTP mode)
+    const rawHttpInput = args.length === 1 ? args[0] : args
+    const httpSystemModified = getHttpPromptMock(rawHttpInput)
+    const httpUserModified = getHttpUserPromptMock(httpSystemModified !== undefined ? httpSystemModified : rawHttpInput)
+    const httpModifiedInput = httpUserModified !== undefined ? httpUserModified : httpSystemModified
+    const httpEffectiveArgs: Args = httpModifiedInput !== undefined ? [httpModifiedInput] as unknown as Args : args
+    const input = httpModifiedInput !== undefined ? httpModifiedInput : rawHttpInput
+
+    try {
+      const output = await callFn(...httpEffectiveArgs)
+      const enriched = enrichFromLLMCapture(input, extractUsage(output), _options?.model, _options?.provider)
+
+      if (isReadableStream(output)) {
+        const [streamForCaller, streamForRecorder] = output.tee()
+        bufferReadableStream(streamForRecorder).then((rawText) => {
+          const durationMs = rawDateNow() - start
+          pushTelemetryEvent({ id, type: 'ai', name: modelName, input: enriched.input, output: null, streamed: true, streamRaw: rawText, timestamp: start, durationMs, ...(enriched.usage ? { usage: enriched.usage } : {}) })
+        }).catch(() => {
+          const durationMs = rawDateNow() - start
+          pushTelemetryEvent({ id, type: 'ai', name: modelName, input: enriched.input, output: null, streamed: true, streamRaw: '', timestamp: start, durationMs, ...(enriched.usage ? { usage: enriched.usage } : {}) })
+        })
+        return streamForCaller as unknown as R
+      }
+
+      if (isAsyncIterable(output)) {
+        return wrapAsyncIterable(output, (chunks) => {
+          const durationMs = rawDateNow() - start
+          const rawText = chunks.map((c) => (typeof c === 'string' ? c : JSON.stringify(c))).join('')
+          pushTelemetryEvent({ id, type: 'ai', name: modelName, input: enriched.input, output: null, streamed: true, streamRaw: rawText, timestamp: start, durationMs, ...(enriched.usage ? { usage: enriched.usage } : {}) })
+        }) as unknown as R
+      }
+
+      const durationMs = rawDateNow() - start
+      const event: WorkflowEvent = {
+        id, type: 'ai', name: modelName, input: enriched.input, output,
+        timestamp: start, durationMs,
+        ...(enriched.usage ? { usage: enriched.usage } : {}),
+      }
+      pushTelemetryEvent(event)
+      return output
+    } catch (e) {
+      const enriched = enrichFromLLMCapture(input, undefined, _options?.model, _options?.provider)
+      const durationMs = rawDateNow() - start
+      pushTelemetryEvent({
+        id, type: 'ai', name: modelName, input: enriched.input,
+        output: { error: String(e) }, timestamp: start, durationMs,
+      })
+      throw e
+    }
+    } finally {
+      leaveAIWrapper()
+    }
+  }
+}

package/src/internals/conditional-recorder.ts

@@ -0,0 +1,63 @@
+/**
+ * Conditional recorder utility for worker contexts
+ * 
+ * This module provides a safe wrapper around recordToolCall that only executes
+ * in worker environments. In non-worker contexts, it's a no-op to avoid unnecessary
+ * imports and function calls.
+ */
+
+// Check if running in an Elasticdash worker context
+const isWorkerContext = (): boolean => {
+  if (typeof globalThis === 'undefined') {
+    return false;
+  }
+  return (globalThis as any).__ELASTICDASH_WORKER__ === true;
+};
+
+let recordToolCallFn: ((name: string, input: any, output: any) => void) | null = null;
+
+// Lazy load recordToolCall only if in worker context
+const getRecordToolCall = async (): Promise<((name: string, input: any, output: any) => void) | null> => {
+  if (!isWorkerContext()) {
+    return null;
+  }
+
+  if (recordToolCallFn !== null) {
+    return recordToolCallFn;
+  }
+
+  try {
+    const { recordToolCall } = await import('../tracing.js');
+    recordToolCallFn = recordToolCall;
+    return recordToolCallFn;
+  } catch (err) {
+    console.warn('Failed to load recordToolCall from tracing module:', err);
+    return null;
+  }
+};
+
+/**
+ * Safely record a tool call only if in worker context
+ * 
+ * @param name - Name of the tool/function
+ * @param input - Input parameters to the tool
+ * @param output - Output/result from the tool
+ */
+export const safeRecordToolCall = async (
+  name: string,
+  input: any,
+  output: any
+): Promise<void> => {
+  const recorder = await getRecordToolCall();
+  if (recorder) {
+    recorder(name, input, output);
+  }
+};
+
+/**
+ * Synchronous version - checks if we're in worker context without importing
+ * Use this if you want to avoid async/await in your tool functions
+ */
+export const isWorker = (): boolean => {
+  return isWorkerContext();
+};