elasticdash-sdk 0.2.0

package/src/observability.ts

@@ -0,0 +1,501 @@
+import { randomUUID } from 'node:crypto'
+import { TelemetryBatcher } from './telemetry-batcher.js'
+import {
+  setObservabilityContext,
+  getObservabilityContext,
+  clearObservabilityContext,
+  pushTelemetryEvent,
+} from './interceptors/telemetry-push.js'
+import type { ObservabilityContext } from './interceptors/telemetry-push.js'
+import { installAIInterceptor } from './interceptors/ai-interceptor.js'
+import { interceptFetch, getOriginalFetch } from './interceptors/http.js'
+import { installDBAutoInterceptor } from './interceptors/db-auto.js'
+import { connectToBackend, disconnectFromBackend } from './socket-connector.js'
+import { executeTrigger } from './trigger-executor.js'
+import { readFileSync, existsSync } from 'node:fs'
+import { join } from 'node:path'
+import { scanTools, scanWorkflows } from './execution/tool-runner.js'
+import { debugLog } from './utils/debug.js'
+import { isTraceCaptureEnabled, maybeCaptureTrace, maybeCaptureTraceSync } from './ci/trace-writer.js'
+import { detectGitInfo } from './ci/git-info.js'
+import type { WorkflowEvent } from './capture/event.js'
+
+export interface ObservabilityOptions {
+  serverUrl?: string
+  apiKey?: string
+  sessionId?: string
+  batchIntervalMs?: number
+  maxBatchSize?: number
+  heartbeatIntervalMs?: number
+  sampleRate?: number
+  redactKeys?: string[]
+}
+
+export interface ObservabilityHandle {
+  sessionId: string
+  shutdown: () => Promise<void>
+}
+
+let heartbeatTimer: ReturnType<typeof setInterval> | null = null
+let shutdownRegistered = false
+
+/**
+ * Extracts workflow names from the `workflows` object in elasticdash.config.ts
+ * by reading the source text. Returns an empty array if the file does not exist
+ * or cannot be parsed.
+ */
+/**
+ * Resolves the default workflow name used as traceId prefix.
+ *
+ * Filters ed_workflows.ts exports to exclude known SDK utility functions
+ * (edStartTrace, edEndTrace, setElasticDashModule) that are not actual
+ * workflow handlers. If exactly one candidate remains, use it. Otherwise
+ * fall back to 'unknown-workflow'.
+ */
+function resolveDefaultWorkflowName(cwd: string, workflows: ReturnType<typeof scanWorkflows>): string {
+  // Filter out SDK utility functions that users commonly export from ed_workflows.ts
+  const UTILITY_PREFIXES = ['edStartTrace', 'edEndTrace', 'setElasticDashModule', 'setElasticDash']
+  const candidates = workflows.filter(w => !UTILITY_PREFIXES.some(p => w.name === p))
+
+  debugLog(`[elasticdash] scanWorkflows found ${workflows.length} exports: [${workflows.map(w => w.name).join(', ')}]`)
+  debugLog(`[elasticdash] After filtering utilities: ${candidates.length} candidates: [${candidates.map(w => w.name).join(', ')}]`)
+
+  if (candidates.length === 1) {
+    debugLog(`[elasticdash] Resolved workflow name: ${candidates[0].name}`)
+    return candidates[0].name
+  }
+
+  // Multiple candidates — try reading elasticdash.config.ts for workflow keys
+  if (candidates.length > 1) {
+    const configNames = readConfigWorkflowNames(cwd)
+    if (configNames.length >= 1) {
+      debugLog(`[elasticdash] Config workflow names: [${configNames.join(', ')}] — using first: ${configNames[0]}`)
+      return configNames[0]
+    }
+    // No config — use first candidate
+    debugLog(`[elasticdash] Multiple workflow candidates, using first: ${candidates[0].name}`)
+    return candidates[0].name
+  }
+
+  // No candidates after filtering — use first raw export if available
+  if (workflows.length > 0) {
+    debugLog(`[elasticdash] No workflow candidates after filtering, using first export: ${workflows[0].name}`)
+    return workflows[0].name
+  }
+
+  debugLog(`[elasticdash] No workflows found — using 'unknown-workflow'`)
+  return 'unknown-workflow'
+}
+
+/**
+ * Reads workflow keys from the `workflows` object in elasticdash.config.ts.
+ * Uses dynamic import to parse the config properly.
+ * Returns an empty array if the file does not exist or has no workflows.
+ */
+function readConfigWorkflowNames(cwd: string): string[] {
+  for (const filename of ['elasticdash.config.ts', 'elasticdash.config.js']) {
+    const configPath = join(cwd, filename)
+    if (!existsSync(configPath)) continue
+    try {
+      const src = readFileSync(configPath, 'utf8')
+      // Match top-level keys inside `workflows: { ... }` using a line-based approach:
+      // workflow keys appear as `identifier: {` at the start of a line (with indentation)
+      // within the workflows block.
+      const wfMatch = src.match(/workflows\s*:\s*\{/)
+      if (!wfMatch || wfMatch.index === undefined) continue
+
+      // Find the workflows block by brace-balancing (skip strings)
+      const startIdx = wfMatch.index + wfMatch[0].length
+      let depth = 1
+      let inString: string | null = null
+      let endIdx = startIdx
+      for (let i = startIdx; i < src.length && depth > 0; i++) {
+        const ch = src[i]
+        if (inString) {
+          if (ch === '\\') { i++; continue }
+          if (ch === inString) inString = null
+          continue
+        }
+        if (ch === "'" || ch === '"' || ch === '`') { inString = ch; continue }
+        if (ch === '{') depth++
+        else if (ch === '}') depth--
+        if (depth > 0) endIdx = i
+      }
+
+      // Extract keys at depth 0 within the block — only `identifier: {` with `mode:` inside
+      const block = src.slice(startIdx, endIdx + 1)
+      const keys: string[] = []
+      // Simple approach: find lines matching `word: {` and check they contain `mode:`
+      for (const m of block.matchAll(/^\s+(\w+)\s*:\s*\{/gm)) {
+        // Verify this is a workflow config (has a `mode:` property) not a nested object
+        const keyStart = m.index! + m[0].length
+        // Scan ahead to the matching } to check for `mode:`
+        let d = 1
+        let inStr: string | null = null
+        let slice = ''
+        for (let i = keyStart; i < block.length && d > 0; i++) {
+          const c = block[i]
+          if (inStr) {
+            if (c === '\\') { i++; continue }
+            if (c === inStr) inStr = null
+            continue
+          }
+          if (c === "'" || c === '"' || c === '`') { inStr = c; continue }
+          if (c === '{') d++
+          else if (c === '}') d--
+          if (d > 0) slice += c
+        }
+        if (/\bmode\s*:/.test(slice)) {
+          keys.push(m[1])
+        }
+      }
+      if (keys.length > 0) return keys
+    } catch { /* config unreadable — fall through */ }
+  }
+  return []
+}
+
+/**
+ * Initialise active observability mode. All `wrapTool` / `wrapAI` calls will
+ * automatically record and stream trace events to the configured backend.
+ *
+ * Reads from environment variables as fallback:
+ * - `ELASTICDASH_API_URL` — backend API URL (required)
+ * - `ELASTICDASH_API_KEY` — project auth token
+ * - `ELASTICDASH_SESSION_ID` — session identifier (auto-generated if omitted)
+ */
+export function initObservability(options?: ObservabilityOptions): ObservabilityHandle {
+  const serverUrl = options?.serverUrl ?? process.env.ELASTICDASH_API_URL ?? ''
+  if (!serverUrl) {
+    throw new Error('[elasticdash] initObservability: serverUrl is required (set ELASTICDASH_API_URL or pass serverUrl option)')
+  }
+
+  const apiKey = options?.apiKey ?? process.env.ELASTICDASH_API_KEY
+  const sessionId = options?.sessionId ?? process.env.ELASTICDASH_SESSION_ID ?? randomUUID()
+  const sampleRate = options?.sampleRate ?? 1.0
+  const redactKeys = options?.redactKeys ?? []
+  const heartbeatIntervalMs = options?.heartbeatIntervalMs ?? 30_000
+
+  // Discover workflow name for traceId prefix.
+  // Priority: elasticdash.config.ts workflow keys > single ed_workflows.ts export > 'unknown-workflow'
+  const cwd = process.cwd()
+  const workflows = scanWorkflows(cwd)
+  const defaultWorkflowName = resolveDefaultWorkflowName(cwd, workflows)
+
+  // Detect git info for session metadata (branch, commit, CI provider)
+  // If detection fails, values are omitted — never store 'unknown' placeholders
+  let gitInfo: ReturnType<typeof detectGitInfo> = {}
+  try { gitInfo = detectGitInfo() } catch { /* not in a git repo — skip */ }
+  const sessionMetadata: Record<string, unknown> = {}
+  if (gitInfo.branch && gitInfo.branch !== 'unknown') sessionMetadata.branch = gitInfo.branch
+  if (gitInfo.commit && gitInfo.commit !== 'unknown') sessionMetadata.commit = gitInfo.commit
+  if (gitInfo.ciProvider && gitInfo.ciProvider !== 'local') sessionMetadata.ciProvider = gitInfo.ciProvider
+
+  const batcher = new TelemetryBatcher({
+    serverUrl,
+    apiKey,
+    sessionId,
+    metadata: Object.keys(sessionMetadata).length > 0 ? sessionMetadata : undefined,
+    batchIntervalMs: options?.batchIntervalMs,
+    maxBatchSize: options?.maxBatchSize,
+    redactKeys,
+    onTrigger: async (trigger) => {
+      await executeTrigger(serverUrl, apiKey, trigger)
+    },
+  })
+
+  let counter = 0
+  const traceId = `${defaultWorkflowName}::${Date.now()}::${randomUUID().slice(0, 8)}`
+  const ctx: ObservabilityContext = {
+    sessionId,
+    serverUrl,
+    apiKey,
+    batcher,
+    nextId: () => ++counter,
+    sampleRate,
+    redactKeys,
+    traceId,
+    defaultWorkflowName,
+  }
+
+  setObservabilityContext(ctx)
+  installAIInterceptor()
+  interceptFetch()
+  installDBAutoInterceptor().catch(() => {})
+  connectToBackend({ serverUrl, apiKey, sessionId })
+
+  // Heartbeat
+  heartbeatTimer = setInterval(() => {
+    batcher.enqueue({
+      id: ctx.nextId(),
+      type: 'side_effect',
+      name: '__heartbeat__',
+      input: { sessionId },
+      output: { uptime: process.uptime() },
+      timestamp: Date.now(),
+      durationMs: 0,
+      schemaVersion: 1,
+    })
+  }, heartbeatIntervalMs)
+  if (heartbeatTimer && typeof heartbeatTimer === 'object' && 'unref' in heartbeatTimer) {
+    heartbeatTimer.unref()
+  }
+
+  // Register process exit hooks (once)
+  if (!shutdownRegistered) {
+    shutdownRegistered = true
+    const onExit = () => { shutdownObservability().catch(() => {}) }
+    process.once('beforeExit', onExit)
+    process.once('SIGTERM', onExit)
+    process.once('SIGINT', onExit)
+
+    // Sync exit handler: last resort to flush any pending trace capture.
+    // process.on('exit') only allows synchronous code — the async handlers above
+    // may not complete before the process terminates (e.g. SIGINT on a dev server).
+    if (isTraceCaptureEnabled()) {
+      process.once('exit', () => {
+        const exitCtx = getObservabilityContext()
+        if (exitCtx?.eventCollector && exitCtx.eventCollector.length > 0) {
+          maybeCaptureTraceSync(exitCtx.eventCollector, exitCtx.traceId)
+        }
+      })
+    }
+  }
+
+  debugLog(`[elasticdash] Observability initialised — sessionId=${sessionId} workflow=${defaultWorkflowName} server=${serverUrl}`)
+
+  // Push workflow/tool catalog to backend (fire-and-forget)
+  pushCatalog(serverUrl, apiKey, workflows).catch(() => {})
+
+  return {
+    sessionId,
+    shutdown: () => shutdownObservability(),
+  }
+}
+
+/**
+ * POST the workflow/tool catalog to the backend so it registers the real
+ * function names from ed_workflows.ts and ed_tools.ts.
+ */
+async function pushCatalog(serverUrl: string, apiKey?: string, workflows: ReturnType<typeof scanWorkflows> = []): Promise<void> {
+  const cwd = process.cwd()
+  const tools = scanTools(cwd)
+
+  if (workflows.length === 0 && tools.length === 0) return
+
+  const body: Record<string, unknown> = {}
+  if (workflows.length > 0) body.workflows = workflows.map(w => ({ name: w.name }))
+  if (tools.length > 0) body.tools = tools.map(t => ({ name: t.name }))
+
+  const headers: Record<string, string> = { 'Content-Type': 'application/json' }
+  if (apiKey) headers['Authorization'] = `Bearer ${apiKey}`
+
+  try {
+    const res = await getOriginalFetch()(`${serverUrl.replace(/\/$/, '')}/api/observability/catalog`, {
+      method: 'POST',
+      headers,
+      body: JSON.stringify(body),
+    })
+    debugLog(`[elasticdash] Catalog pushed: ${workflows.length} workflows, ${tools.length} tools (status ${res.status})`)
+  } catch (err) {
+    debugLog(`[elasticdash] Catalog push failed: ${err instanceof Error ? err.message : String(err)}`)
+  }
+}
+
+/**
+ * Gracefully shut down observability: flush remaining events, send session_end,
+ * stop heartbeat, and clear context.
+ */
+export async function shutdownObservability(): Promise<void> {
+  const ctx = getObservabilityContext()
+  if (!ctx) return
+
+  // Flush any pending trace capture before shutdown
+  if (ctx.eventCollector && ctx.eventCollector.length > 0) {
+    const events = ctx.eventCollector
+    ctx.eventCollector = undefined
+    await maybeCaptureTrace(events, ctx.traceId)
+  }
+
+  // Stop heartbeat
+  if (heartbeatTimer) {
+    clearInterval(heartbeatTimer)
+    heartbeatTimer = null
+  }
+
+  // Send session_end event
+  ctx.batcher.enqueue({
+    id: ctx.nextId(),
+    type: 'side_effect',
+    name: '__session_end__',
+    input: { sessionId: ctx.sessionId },
+    output: { uptime: process.uptime() },
+    timestamp: Date.now(),
+    durationMs: 0,
+    schemaVersion: 1,
+  })
+
+  // Disconnect socket
+  await disconnectFromBackend()
+
+  // Flush and shut down the batcher
+  await ctx.batcher.shutdown()
+
+  clearObservabilityContext()
+  debugLog(`[elasticdash] Observability shut down — sessionId=${ctx.sessionId}`)
+}
+
+/**
+ * Start a new trace for the current request/operation.
+ *
+ * Call this at the start of each request handler to group all tool/AI events
+ * under one trace, identified by the workflow name.
+ *
+ * The generated traceId encodes the workflow name and timestamp for easy
+ * parsing by the backend: `{workflowName}::{timestamp}::{shortId}`
+ *
+ * @param workflowName — the workflow/function being executed (e.g. 'chatStreamHandler')
+ * @returns The new traceId
+ *
+ * @example
+ * ```ts
+ * // In a route handler:
+ * startTrace('chatStreamHandler')
+ * // All subsequent wrapTool/wrapAI calls get traceId = "chatStreamHandler::1712851200000::a1b2c3d4"
+ * ```
+ */
+export function startTrace(workflowName?: string): string {
+  const ctx = getObservabilityContext()
+  if (!ctx) throw new Error('[elasticdash] startTrace: observability not initialised')
+
+  // Flush any pending trace capture from a previous startTrace() that had no endTrace()
+  if (ctx.eventCollector && ctx.eventCollector.length > 0) {
+    debugLog(`[elasticdash] startTrace: flushing ${ctx.eventCollector.length} events from previous trace`)
+    const events = ctx.eventCollector
+    ctx.eventCollector = undefined
+    maybeCaptureTrace(events, ctx.traceId).catch((err) => {
+      debugLog(`[elasticdash] Failed to capture trace: ${err instanceof Error ? err.message : String(err)}`)
+    })
+  }
+
+  const resolvedName = workflowName || ctx.defaultWorkflowName
+  const timestamp = Date.now()
+  const shortId = randomUUID().slice(0, 8)
+  ctx.traceId = `${resolvedName}::${timestamp}::${shortId}`
+
+  // When ELASTICDASH_CAPTURE_TRACE=1, start collecting events for disk trace
+  const captureEnabled = isTraceCaptureEnabled()
+  debugLog(`[elasticdash] startTrace: ${resolvedName}, capture=${captureEnabled}`)
+  if (captureEnabled) {
+    ctx.eventCollector = []
+  }
+
+  // Emit a workflow event so the backend can discover the workflow name
+  // (same as wrapWorkflow, but without wrapping — for streaming handlers)
+  pushTelemetryEvent({
+    id: ctx.nextId(),
+    type: 'workflow',
+    name: resolvedName,
+    input: null,
+    output: null,
+    timestamp,
+    durationMs: 0,
+    schemaVersion: 1,
+  })
+
+  return ctx.traceId
+}
+
+/**
+ * End the current trace. Resets the traceId so subsequent events (heartbeats,
+ * new requests) are not attributed to the completed workflow.
+ *
+ * Call this when a workflow/request handler finishes — especially for streaming
+ * handlers where the trace would otherwise persist indefinitely.
+ *
+ * @example
+ * ```ts
+ * startTrace('chatStreamHandler')
+ * // ... workflow logic ...
+ * endTrace()
+ * ```
+ */
+export function endTrace(): void {
+  const ctx = getObservabilityContext()
+  debugLog(`[elasticdash] endTrace called, obsCtx=${!!ctx}, eventCollector=${ctx?.eventCollector?.length ?? 'none'}`)
+  if (!ctx) return
+
+  // Flush collected events to disk if trace capture is enabled
+  if (ctx.eventCollector && ctx.eventCollector.length > 0) {
+    debugLog(`[elasticdash] endTrace: flushing ${ctx.eventCollector.length} events to disk for traceId=${ctx.traceId}`)
+    const events = ctx.eventCollector
+    ctx.eventCollector = undefined
+    maybeCaptureTrace(events, ctx.traceId).catch((err) => {
+      debugLog(`[elasticdash] Failed to capture trace: ${err instanceof Error ? err.message : String(err)}`)
+    })
+  }
+
+  const timestamp = Date.now()
+  const shortId = randomUUID().slice(0, 8)
+  ctx.traceId = `${ctx.defaultWorkflowName}::${timestamp}::${shortId}`
+}
+
+/**
+ * Wrap a workflow function so that every invocation automatically starts a new
+ * trace with the workflow name. All tool/AI calls within the function execution
+ * are grouped under that trace.
+ *
+ * Use this in `ed_workflows.ts` to automatically tag each workflow execution:
+ *
+ * @example
+ * ```ts
+ * export const chatStreamHandler = wrapWorkflow('chatStreamHandler', async (input) => {
+ *   const result = await fetchUser(input.userId)
+ *   const reply = await generateReply(result)
+ *   return reply
+ * })
+ * ```
+ */
+export function wrapWorkflow<Args extends unknown[], R>(
+  name: string,
+  fn: (...args: Args) => Promise<R>,
+): (...args: Args) => Promise<R> {
+  return async (...args: Args): Promise<R> => {
+    const ctx = getObservabilityContext()
+    if (ctx) {
+      startTrace(name)
+      const collectedEvents: import('./capture/event.js').WorkflowEvent[] = []
+      ctx.eventCollector = collectedEvents
+      const start = Date.now()
+      let result: R
+      let error: unknown = undefined
+      try {
+        result = await fn(...args)
+      } catch (err) {
+        error = err
+        result = undefined as R
+      }
+      // Emit a single workflow event regardless of success or failure.
+      // Inner tool/AI events already record their own errors individually —
+      // the workflow event should reflect the workflow's own output, not
+      // escalate inner step failures to the workflow level.
+      pushTelemetryEvent({
+        id: ctx.nextId(),
+        type: 'workflow',
+        name,
+        input: args.length === 1 ? args[0] : args,
+        output: error !== undefined ? null : result,
+        timestamp: start,
+        durationMs: Date.now() - start,
+        schemaVersion: 1,
+      })
+      ctx.eventCollector = undefined
+      // Persist trace to disk if ELASTICDASH_CAPTURE_TRACE=1
+      await maybeCaptureTrace(collectedEvents, ctx.traceId)
+      endTrace()
+      if (error !== undefined) throw error
+      return result
+    }
+    return fn(...args)
+  }
+}