elasticdash-sdk 0.2.0

package/src/internals/mock-resolver.ts

@@ -0,0 +1,492 @@
+/**
+ * Runtime mock resolution for module-imported tools and AI calls.
+ *
+ * Tools/AI calls that are statically imported (not accessed via globalThis)
+ * cannot be intercepted by the worker's proxy-based mocking. Instead, each
+ * wrapped function calls `resolveMock` / `resolveAIMock` at its entry point.
+ * The worker writes the mock config to `__ELASTICDASH_TOOL_MOCKS__` /
+ * `__ELASTICDASH_AI_MOCKS__` before the workflow runs and clears it after.
+ */
+
+interface ToolMockEntry {
+  mode: 'live' | 'mock-all' | 'mock-specific'
+  callIndices?: number[]
+  mockData?: Record<number, unknown>
+}
+
+/** Per-model AI mock configuration (mirrors ToolMockConfig) */
+export interface AIMockEntry {
+  mode: 'live' | 'mock-all' | 'mock-specific'
+  callIndices?: number[]
+  mockData?: Record<number, unknown>
+}
+
+export interface AIMockConfig {
+  [modelName: string]: AIMockEntry
+}
+
+type MockResult =
+  | { mocked: true; result: unknown }
+  | { mocked: false }
+
+/**
+ * Recursively normalises a mock result fetched from the trace.
+ *
+ * The trace recorder stores tool outputs as-is. When a tool returns a JSON
+ * string (rather than a parsed object), the stored value is that string, and
+ * inner string values may themselves be JSON-quoted. This function unwraps
+ * every layer so the calling code receives the same shape as the real output.
+ *
+ * - string  → try JSON.parse, then recurse on the result
+ * - array   → recurse on every element
+ * - object  → recurse on every value
+ * - other   → return unchanged
+ */
+export function normaliseMockResult(value: unknown): unknown {
+  if (typeof value === 'string') {
+    try { return normaliseMockResult(JSON.parse(value)) } catch { return value }
+  }
+  if (Array.isArray(value)) {
+    return value.map(normaliseMockResult)
+  }
+  if (value !== null && typeof value === 'object') {
+    const out: Record<string, unknown> = {}
+    for (const [k, v] of Object.entries(value as Record<string, unknown>)) {
+      out[k] = normaliseMockResult(v)
+    }
+    return out
+  }
+  return value
+}
+
+/**
+ * Resolves whether the current call to `toolName` should be mocked.
+ *
+ * - Returns `{ mocked: false }` when no mock config is active or mode is `'live'`.
+ * - Increments the per-tool call counter on every non-live invocation so that
+ *   `mock-specific` indices remain accurate even when some calls run live.
+ * - Safe to call in production: no-op when the mock globals are absent.
+ */
+export function resolveMock(toolName: string): MockResult {
+  const g = globalThis as Record<string, unknown>
+
+  const mocks = g['__ELASTICDASH_TOOL_MOCKS__'] as Record<string, ToolMockEntry> | undefined
+  if (!mocks) return { mocked: false }
+
+  const entry = mocks[toolName]
+  if (!entry || entry.mode === 'live') return { mocked: false }
+
+  // Initialise counters map if not yet present
+  if (!g['__ELASTICDASH_TOOL_CALL_COUNTERS__']) {
+    g['__ELASTICDASH_TOOL_CALL_COUNTERS__'] = {} as Record<string, number>
+  }
+  const counters = g['__ELASTICDASH_TOOL_CALL_COUNTERS__'] as Record<string, number>
+  counters[toolName] = (counters[toolName] ?? 0) + 1
+  const callNumber = counters[toolName]
+
+  if (entry.mode === 'mock-all') {
+    const data = entry.mockData ?? {}
+    const raw = data[callNumber] !== undefined ? data[callNumber] : data[0]
+    return { mocked: true, result: normaliseMockResult(raw) }
+  }
+
+  if (entry.mode === 'mock-specific') {
+    const indices = entry.callIndices ?? []
+    if (indices.includes(callNumber)) {
+      const data = entry.mockData ?? {}
+      return { mocked: true, result: normaliseMockResult(data[callNumber]) }
+    }
+    // Counter already incremented; this specific call runs live
+    return { mocked: false }
+  }
+
+  return { mocked: false }
+}
+
+/**
+ * Extracts the system prompt string from an LLM call input.
+ * Handles:
+ *  - Anthropic style:     `{ system: "...", messages: [...] }`
+ *  - OpenAI style:        `{ messages: [{ role: "system", content: "..." }, ...] }`
+ *  - Plain message array: `[{ role: "system", content: "..." }, ...]`
+ *  - Separate field:      `{ systemPrompt: "...", messages: [...] }` (custom wrapAI callers)
+ */
+export function extractSystemPrompt(input: unknown): string | undefined {
+  if (!input || typeof input !== 'object') return undefined
+  const o = input as Record<string, unknown>
+  if (typeof o.system === 'string') return o.system
+  if (typeof o.systemPrompt === 'string' && o.systemPrompt.length > 0) return o.systemPrompt
+  const msgs = Array.isArray(o.messages) ? o.messages : (Array.isArray(input) ? input as unknown[] : null)
+  if (msgs) {
+    for (const m of msgs) {
+      if (m && typeof m === 'object') {
+        const msg = m as Record<string, unknown>
+        if (msg.role === 'system' && typeof msg.content === 'string') return msg.content
+      }
+    }
+  }
+  return undefined
+}
+
+/**
+ * Returns a shallow copy of `input` with the system prompt replaced by `newSystemPrompt`.
+ */
+export function replaceSystemPrompt(input: unknown, newSystemPrompt: string): unknown {
+  if (!input || typeof input !== 'object') return input
+  const o = input as Record<string, unknown>
+  if (typeof o.system === 'string') return { ...o, system: newSystemPrompt }
+  if (typeof o.systemPrompt === 'string') return { ...o, systemPrompt: newSystemPrompt }
+  if (Array.isArray(input)) {
+    return (input as unknown[]).map(m => {
+      if (m && typeof m === 'object') {
+        const msg = m as Record<string, unknown>
+        if (msg.role === 'system' && typeof msg.content === 'string') return { ...msg, content: newSystemPrompt }
+      }
+      return m
+    })
+  }
+  if (Array.isArray(o.messages)) {
+    return {
+      ...o,
+      messages: (o.messages as unknown[]).map(m => {
+        if (m && typeof m === 'object') {
+          const msg = m as Record<string, unknown>
+          if (msg.role === 'system' && typeof msg.content === 'string') return { ...msg, content: newSystemPrompt }
+        }
+        return m
+      }),
+    }
+  }
+  return input
+}
+
+/**
+ * System prompt mock entry — mirrors user prompt mock entry.
+ */
+export interface SystemPromptMockEntry {
+  mode: 'live' | 'replace-all' | 'replace-specific'
+  replacement: string
+  callIndices?: number[]
+}
+
+export interface SystemPromptMockConfig {
+  [originalText: string]: SystemPromptMockEntry
+}
+
+/**
+ * If a prompt mock is configured for the system prompt found in `input`, returns
+ * a copy of `input` with the system prompt replaced. Otherwise returns `undefined`.
+ *
+ * `__ELASTICDASH_PROMPT_MOCKS__` is `Record<string, SystemPromptMockEntry>` keyed by
+ * the original system prompt text. Supports replace-all and replace-specific modes
+ * with per-text call counting.
+ *
+ * For backward compatibility, also accepts the legacy `Record<string, string>` format
+ * (treated as replace-all).
+ */
+export function resolvePromptMock(input: unknown): unknown | undefined {
+  const g = globalThis as Record<string, unknown>
+  const mocks = g['__ELASTICDASH_PROMPT_MOCKS__'] as Record<string, string | SystemPromptMockEntry> | undefined
+  if (!mocks || Object.keys(mocks).length === 0) return undefined
+  const systemPrompt = extractSystemPrompt(input)
+  if (systemPrompt === undefined) return undefined
+
+  const match = lookupMockEntry(mocks, systemPrompt)
+  if (!match) return undefined
+  const entry = match.entry
+
+  // Legacy format: plain string replacement (treat as replace-all)
+  if (typeof entry === 'string') {
+    return replaceSystemPrompt(input, entry)
+  }
+
+  if (entry.mode === 'live') return undefined
+
+  // Initialise per-text call counters
+  if (!g['__ELASTICDASH_PROMPT_CALL_COUNTERS__']) {
+    g['__ELASTICDASH_PROMPT_CALL_COUNTERS__'] = {} as Record<string, number>
+  }
+  const counters = g['__ELASTICDASH_PROMPT_CALL_COUNTERS__'] as Record<string, number>
+  counters[match.key] = (counters[match.key] ?? 0) + 1
+  const callNumber = counters[match.key]
+
+  if (entry.mode === 'replace-all') {
+    return replaceSystemPrompt(input, entry.replacement)
+  }
+
+  if (entry.mode === 'replace-specific') {
+    const indices = entry.callIndices ?? []
+    if (indices.includes(callNumber)) {
+      return replaceSystemPrompt(input, entry.replacement)
+    }
+  }
+
+  return undefined
+}
+
+/**
+ * Strips one layer of JSON string quoting if present.
+ * `'"hello"'` → `'hello'`, `'hello'` → `'hello'`
+ */
+function stripJsonQuotes(s: string): string {
+  if (s.length >= 2 && s.startsWith('"') && s.endsWith('"')) {
+    try {
+      const parsed = JSON.parse(s)
+      if (typeof parsed === 'string') return parsed
+    } catch { /* not valid JSON, return as-is */ }
+  }
+  return s
+}
+
+/**
+ * Look up a mock entry by key, falling back to stripped-quotes variants
+ * to handle mismatches between trace-recorded text and runtime text.
+ */
+export function lookupMockEntry<T>(mocks: Record<string, T>, text: string): { key: string; entry: T } | undefined {
+  if (mocks[text]) return { key: text, entry: mocks[text] }
+  // Try stripping quotes from the runtime text
+  const stripped = stripJsonQuotes(text)
+  if (stripped !== text && mocks[stripped]) return { key: stripped, entry: mocks[stripped] }
+  // Try adding quotes to the runtime text (mock key might have quotes)
+  const quoted = JSON.stringify(text)
+  if (quoted !== text && mocks[quoted]) return { key: quoted, entry: mocks[quoted] }
+  return undefined
+}
+
+/**
+ * Extracts all user-role message content strings from an LLM call input.
+ * Handles OpenAI-style `{ messages: [...] }` and plain message arrays.
+ */
+export function extractUserPrompts(input: unknown): string[] {
+  if (!input || typeof input !== 'object') return []
+  const o = input as Record<string, unknown>
+  const msgs = Array.isArray(o.messages) ? o.messages : (Array.isArray(input) ? input as unknown[] : null)
+  if (!msgs) return []
+  const results: string[] = []
+  for (const m of msgs) {
+    if (m && typeof m === 'object') {
+      const msg = m as Record<string, unknown>
+      if (msg.role === 'user' && typeof msg.content === 'string') results.push(msg.content)
+    }
+  }
+  return results
+}
+
+/**
+ * Returns a shallow copy of `input` with all user-role messages whose content
+ * matches `originalText` replaced with `newText`.
+ */
+export function replaceUserPrompt(input: unknown, originalText: string, newText: string): unknown {
+  if (!input || typeof input !== 'object') return input
+  const o = input as Record<string, unknown>
+
+  const replaceInMsgs = (msgs: unknown[]): unknown[] =>
+    msgs.map(m => {
+      if (m && typeof m === 'object') {
+        const msg = m as Record<string, unknown>
+        if (msg.role === 'user' && typeof msg.content === 'string' && msg.content === originalText) {
+          return { ...msg, content: newText }
+        }
+      }
+      return m
+    })
+
+  if (Array.isArray(input)) return replaceInMsgs(input as unknown[])
+  if (Array.isArray(o.messages)) return { ...o, messages: replaceInMsgs(o.messages as unknown[]) }
+  return input
+}
+
+/**
+ * User prompt mock entry — mirrors tool mock but for user-role messages.
+ */
+export interface UserPromptMockEntry {
+  mode: 'live' | 'replace-all' | 'replace-specific'
+  replacement: string
+  callIndices?: number[]
+}
+
+export interface UserPromptMockConfig {
+  [originalText: string]: UserPromptMockEntry
+}
+
+/**
+ * If user prompt mocks are configured, applies all matching replacements to `input`.
+ * Returns modified input or `undefined` if no replacements applied.
+ *
+ * `__ELASTICDASH_USER_PROMPT_MOCKS__` is `Record<string, UserPromptMockEntry>` keyed
+ * by original user message text.
+ *
+ * Call counting is per unique text — "the x-th LLM call where this text appears".
+ */
+export function resolveUserPromptMock(input: unknown): unknown | undefined {
+  const g = globalThis as Record<string, unknown>
+  const mocks = g['__ELASTICDASH_USER_PROMPT_MOCKS__'] as UserPromptMockConfig | undefined
+  if (!mocks || Object.keys(mocks).length === 0) return undefined
+
+  const userPrompts = extractUserPrompts(input)
+  if (userPrompts.length === 0) return undefined
+
+  // Initialise per-text call counters
+  if (!g['__ELASTICDASH_USER_PROMPT_CALL_COUNTERS__']) {
+    g['__ELASTICDASH_USER_PROMPT_CALL_COUNTERS__'] = {} as Record<string, number>
+  }
+  const counters = g['__ELASTICDASH_USER_PROMPT_CALL_COUNTERS__'] as Record<string, number>
+
+  // Determine which texts in this call have active mock entries
+  const uniqueTexts = Array.from(new Set(userPrompts))
+  let modified = false
+  let result = input
+
+  for (const text of uniqueTexts) {
+    const match = lookupMockEntry(mocks, text)
+    if (!match || match.entry.mode === 'live') continue
+    const entry = match.entry
+
+    // Increment counter for this text (this is the x-th call where this text appears)
+    counters[match.key] = (counters[match.key] ?? 0) + 1
+    const callNumber = counters[match.key]
+
+    if (entry.mode === 'replace-all') {
+      result = replaceUserPrompt(result, text, entry.replacement)
+      modified = true
+    } else if (entry.mode === 'replace-specific') {
+      const indices = entry.callIndices ?? []
+      if (indices.includes(callNumber)) {
+        result = replaceUserPrompt(result, text, entry.replacement)
+        modified = true
+      }
+    }
+  }
+
+  return modified ? result : undefined
+}
+
+/**
+ * Resolves whether the current call to `modelName` should be mocked.
+ * Mirrors `resolveMock` but reads `__ELASTICDASH_AI_MOCKS__` and
+ * `__ELASTICDASH_AI_CALL_COUNTERS__`.
+ */
+/**
+ * Standalone utility: applies all matching user prompt mock replacements to
+ * a messages array (or LLM input object). Works without ALS — takes the mock
+ * config and call counters as explicit parameters.
+ *
+ * Returns the modified input, or the original input unchanged if no mocks matched.
+ */
+export function applyUserPromptMocks(
+  input: unknown,
+  userPromptMocks: Record<string, UserPromptMockEntry>,
+  callCounters?: Record<string, number>,
+): unknown {
+  if (!userPromptMocks || Object.keys(userPromptMocks).length === 0) return input
+
+  const userPrompts = extractUserPrompts(input)
+  if (userPrompts.length === 0) return input
+
+  const counters = callCounters ?? {}
+  const uniqueTexts = Array.from(new Set(userPrompts))
+  let modified = false
+  let result = input
+
+  for (const text of uniqueTexts) {
+    const match = lookupMockEntry(userPromptMocks, text)
+    if (!match || match.entry.mode === 'live') continue
+    const entry = match.entry
+
+    counters[match.key] = (counters[match.key] ?? 0) + 1
+    const callNumber = counters[match.key]
+
+    if (entry.mode === 'replace-all') {
+      result = replaceUserPrompt(result, text, entry.replacement)
+      modified = true
+    } else if (entry.mode === 'replace-specific') {
+      const indices = entry.callIndices ?? []
+      if (indices.includes(callNumber)) {
+        result = replaceUserPrompt(result, text, entry.replacement)
+        modified = true
+      }
+    }
+  }
+
+  return modified ? result : input
+}
+
+/**
+ * Standalone utility: applies system prompt mock replacement to an LLM input
+ * object. Works without ALS — takes the mock config and call counters as
+ * explicit parameters.
+ *
+ * Returns the modified input, or the original input unchanged if no mock matched.
+ */
+export function applySystemPromptMocks(
+  input: unknown,
+  promptMocks: Record<string, string | SystemPromptMockEntry>,
+  callCounters?: Record<string, number>,
+): unknown {
+  if (!promptMocks || Object.keys(promptMocks).length === 0) return input
+
+  const systemPrompt = extractSystemPrompt(input)
+  if (systemPrompt === undefined) return input
+
+  const match = lookupMockEntry(promptMocks, systemPrompt)
+  if (!match) return input
+  const entry = match.entry
+
+  // Legacy format: plain string replacement (treat as replace-all)
+  if (typeof entry === 'string') {
+    return replaceSystemPrompt(input, entry)
+  }
+
+  if (entry.mode === 'live') return input
+
+  const counters = callCounters ?? {}
+  counters[match.key] = (counters[match.key] ?? 0) + 1
+  const callNumber = counters[match.key]
+
+  if (entry.mode === 'replace-all') {
+    return replaceSystemPrompt(input, entry.replacement)
+  }
+
+  if (entry.mode === 'replace-specific') {
+    const indices = entry.callIndices ?? []
+    if (indices.includes(callNumber)) {
+      return replaceSystemPrompt(input, entry.replacement)
+    }
+  }
+
+  return input
+}
+
+export function resolveAIMock(modelName: string): MockResult {
+  const g = globalThis as Record<string, unknown>
+
+  const mocks = g['__ELASTICDASH_AI_MOCKS__'] as Record<string, AIMockEntry> | undefined
+  if (!mocks) return { mocked: false }
+
+  const entry = mocks[modelName]
+  if (!entry || entry.mode === 'live') return { mocked: false }
+
+  if (!g['__ELASTICDASH_AI_CALL_COUNTERS__']) {
+    g['__ELASTICDASH_AI_CALL_COUNTERS__'] = {} as Record<string, number>
+  }
+  const counters = g['__ELASTICDASH_AI_CALL_COUNTERS__'] as Record<string, number>
+  counters[modelName] = (counters[modelName] ?? 0) + 1
+  const callNumber = counters[modelName]
+
+  if (entry.mode === 'mock-all') {
+    const data = entry.mockData ?? {}
+    const raw = data[callNumber] !== undefined ? data[callNumber] : data[0]
+    return { mocked: true, result: normaliseMockResult(raw) }
+  }
+
+  if (entry.mode === 'mock-specific') {
+    const indices = entry.callIndices ?? []
+    if (indices.includes(callNumber)) {
+      const data = entry.mockData ?? {}
+      return { mocked: true, result: normaliseMockResult(data[callNumber]) }
+    }
+    return { mocked: false }
+  }
+
+  return { mocked: false }
+}