elasticdash-sdk 0.2.0

package/src/interceptors/tool.ts

@@ -0,0 +1,287 @@
+import { getCaptureContext } from '../capture/recorder.js'
+import { getCurrentTrace } from '../trace-adapter/context.js'
+import { rawDateNow } from './side-effects.js'
+import { getHttpRunContext, getHttpFrozenEvent, getHttpToolMock, pushTelemetryEvent, tryAutoInitHttpContext, getObservabilityContext } from './telemetry-push.js'
+import { getEdReplayContext, replayCall } from '../ci/replay.js'
+import { debugLog } from '../utils/debug.js'
+
+const TOOL_WRAPPER_ACTIVE_KEY = '__elasticdash_tool_wrapper_active__'
+
+function toTraceArgs(input: unknown): Record<string, unknown> | undefined {
+  if (input && typeof input === 'object' && !Array.isArray(input)) {
+    return input as Record<string, unknown>
+  }
+  if (input === undefined) return undefined
+  return { value: input }
+}
+
+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 wrapTool<Args extends unknown[], R>(
+  name: string,
+  fn: (...args: Args) => Promise<R>,
+): (...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, 'tool_call', name, input)
+      return output as R
+    }
+
+    await tryAutoInitHttpContext()
+    const ctx = getCaptureContext()
+    const httpCtx = getHttpRunContext()
+    const obsCtx = getObservabilityContext()
+    debugLog(`[elasticdash] Tool called: ${name}`, { args })
+    debugLog(`[elasticdash] Current capture context:`, ctx ? { hasRecorder: !!ctx.recorder, hasReplay: !!ctx.replay } : null)
+    debugLog(`[elasticdash] Current HTTP context:`, httpCtx ? { hasHttpContext: !!httpCtx } : null)
+    if (!ctx && !httpCtx && !obsCtx) return fn(...args)
+
+    // Observability-only mode: record and push, no mocks/replay
+    if (!ctx && !httpCtx && obsCtx) {
+      debugLog(`[elasticdash] wrapTool obs-only path: ${name}`)
+      const id = obsCtx.nextId()
+      const input = args.length === 1 ? args[0] : args
+      const start = rawDateNow()
+      try {
+        const output = await fn(...args)
+        const durationMs = rawDateNow() - start
+        if (isReadableStream(output)) {
+          const [streamForCaller, streamForRecorder] = output.tee()
+          bufferReadableStream(streamForRecorder).then((rawText) => {
+            const durationMs = rawDateNow() - start
+            pushTelemetryEvent({ id, type: 'tool', name, input, output: null, streamed: true, streamRaw: rawText, timestamp: start, durationMs })
+          }).catch(() => {
+            const durationMs = rawDateNow() - start
+            pushTelemetryEvent({ id, type: 'tool', name, input, output: null, streamed: true, streamRaw: '', timestamp: start, durationMs })
+          })
+          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: 'tool', name, input, output: null, streamed: true, streamRaw: rawText, timestamp: start, durationMs })
+          }) as unknown as R
+        }
+        pushTelemetryEvent({ id, type: 'tool', name, input, output, timestamp: start, durationMs })
+        return output
+      } catch (e) {
+        const durationMs = rawDateNow() - start
+        pushTelemetryEvent({ id, type: 'tool', name, input, output: { error: String(e) }, timestamp: start, durationMs })
+        throw e
+      }
+    }
+
+    // HTTP mode only (no capture context) — execute, record timing, push telemetry
+    if (!ctx) {
+      const id = httpCtx!.nextId()
+      const input = args.length === 1 ? args[0] : args
+
+      // 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 tool output mock (skip real call, return mocked result)
+      const toolMock = getHttpToolMock(name)
+      if (toolMock.mocked) {
+        const event = { id, type: 'tool' as const, name, input, output: toolMock.result, timestamp: rawDateNow(), durationMs: 0 }
+        pushTelemetryEvent(event)
+        return toolMock.result as R
+      }
+
+      const start = rawDateNow()
+      try {
+        const output = await fn(...args)
+        if (isReadableStream(output)) {
+          const [streamForCaller, streamForRecorder] = output.tee()
+          bufferReadableStream(streamForRecorder).then((rawText) => {
+            const durationMs = rawDateNow() - start
+            pushTelemetryEvent({ id, type: 'tool', name, input, output: null, streamed: true, streamRaw: rawText, timestamp: start, durationMs })
+          }).catch(() => {
+            const durationMs = rawDateNow() - start
+            pushTelemetryEvent({ id, type: 'tool', name, input, output: null, streamed: true, streamRaw: '', timestamp: start, durationMs })
+          })
+          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: 'tool', name, input, output: null, streamed: true, streamRaw: rawText, timestamp: start, durationMs })
+          }) as unknown as R
+        }
+        const durationMs = rawDateNow() - start
+        pushTelemetryEvent({ id, type: 'tool', name, input, output, timestamp: start, durationMs })
+        return output
+      } catch (e) {
+        const durationMs = rawDateNow() - start
+        pushTelemetryEvent({ id, type: 'tool', name, input, output: { error: String(e) }, timestamp: start, durationMs })
+        throw e
+      }
+    }
+
+    const trace = getCurrentTrace()
+    const { recorder, replay } = ctx
+    const id = recorder.nextId()
+    const input = args.length === 1 ? args[0] : args
+
+    if (replay.shouldReplay(id)) {
+      const historical = replay.getRecordedEvent(id)
+      if (historical) recorder.record(historical)
+
+      if (historical?.streamed === true) {
+        const raw = typeof historical.streamRaw === 'string' ? historical.streamRaw : ''
+        const stream = reconstructStream(raw) as unknown as R
+        if (trace && typeof trace.recordToolCall === 'function') {
+          trace.recordToolCall({ name, args: toTraceArgs(input), result: stream, workflowEventId: id })
+        }
+        return stream
+      }
+
+      const replayed = replay.getRecordedResult(id) as R
+      if (trace && typeof trace.recordToolCall === 'function') {
+        trace.recordToolCall({ name, args: toTraceArgs(input), result: replayed, workflowEventId: id })
+      }
+      return replayed
+    }
+
+    const g = globalThis as Record<string, unknown>
+    const prev = g[TOOL_WRAPPER_ACTIVE_KEY]
+    g[TOOL_WRAPPER_ACTIVE_KEY] = true
+    const start = rawDateNow()
+
+    try {
+      const output = await fn(...args)
+
+      if (isReadableStream(output)) {
+        const [streamForCaller, streamForRecorder] = output.tee()
+        bufferReadableStream(streamForRecorder).then((rawText) => {
+          const durationMs = rawDateNow() - start
+          const event = { id, type: 'tool' as const, name, input, output: null, streamed: true, streamRaw: rawText, timestamp: start, durationMs }
+          recorder.record(event)
+          if (httpCtx) pushTelemetryEvent(event)
+        }).catch(() => {
+          const durationMs = rawDateNow() - start
+          const event = { id, type: 'tool' as const, name, input, output: null, streamed: true, streamRaw: '', timestamp: start, durationMs }
+          recorder.record(event)
+          if (httpCtx) pushTelemetryEvent(event)
+        })
+        const result = streamForCaller as unknown as R
+        if (trace && typeof trace.recordToolCall === 'function') {
+          trace.recordToolCall({ name, args: toTraceArgs(input), result, workflowEventId: id })
+        }
+        return result
+      }
+
+      if (isAsyncIterable(output)) {
+        const wrapped = wrapAsyncIterable(output, (chunks) => {
+          const durationMs = rawDateNow() - start
+          const rawText = chunks.map((c) => (typeof c === 'string' ? c : JSON.stringify(c))).join('')
+          const event = { id, type: 'tool' as const, name, input, output: null, streamed: true, streamRaw: rawText, timestamp: start, durationMs }
+          recorder.record(event)
+          if (httpCtx) pushTelemetryEvent(event)
+        }) as unknown as R
+        if (trace && typeof trace.recordToolCall === 'function') {
+          trace.recordToolCall({ name, args: toTraceArgs(input), result: wrapped, workflowEventId: id })
+        }
+        return wrapped
+      }
+
+      const durationMs = rawDateNow() - start
+      const successEvent = { id, type: 'tool' as const, name, input, output, timestamp: start, durationMs }
+      recorder.record(successEvent)
+      if (httpCtx) pushTelemetryEvent(successEvent)
+      if (trace && typeof trace.recordToolCall === 'function') {
+        trace.recordToolCall({ name, args: toTraceArgs(input), result: output, workflowEventId: id, durationMs })
+      }
+      return output
+    } catch (e) {
+      const durationMs = rawDateNow() - start
+      const errorEvent = { id, type: 'tool' as const, name, input, output: { error: String(e) }, timestamp: start, durationMs }
+      recorder.record(errorEvent)
+      if (httpCtx) pushTelemetryEvent(errorEvent)
+      if (trace && typeof trace.recordToolCall === 'function') {
+        trace.recordToolCall({ name, args: toTraceArgs(input), result: { error: String(e) }, workflowEventId: id, durationMs })
+      }
+      throw e
+    } finally {
+      if (prev === undefined) delete g[TOOL_WRAPPER_ACTIVE_KEY]
+      else g[TOOL_WRAPPER_ACTIVE_KEY] = prev
+    }
+  }
+}