elasticdash-test 0.1.14 → 0.1.16

package/src/internals/mock-resolver.ts

@@ -1,10 +1,11 @@
 /**
- * Runtime tool mock resolution for module-imported tools.
+ * Runtime mock resolution for module-imported tools and AI calls.
  *
- * Tools that are statically imported (not accessed via globalThis) cannot be
- * intercepted by the worker's proxy-based mocking. Instead, each tool function
- * calls `resolveMock` at its entry point. The worker writes the mock config to
- * `__ELASTICDASH_TOOL_MOCKS__` before the workflow runs and clears it after.
+ * 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 {
@@ -13,6 +14,17 @@ interface ToolMockEntry {
   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 }
@@ -90,3 +102,117 @@ export function resolveMock(toolName: string): MockResult {
 
   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
+}
+
+/**
+ * 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, string>` keyed by the original
+ * system prompt text, set by the worker/dashboard before the workflow runs.
+ */
+export function resolvePromptMock(input: unknown): unknown | undefined {
+  const g = globalThis as Record<string, unknown>
+  const mocks = g['__ELASTICDASH_PROMPT_MOCKS__'] as Record<string, string> | undefined
+  if (!mocks || Object.keys(mocks).length === 0) return undefined
+  const systemPrompt = extractSystemPrompt(input)
+  if (systemPrompt === undefined) return undefined
+  const newSystemPrompt = mocks[systemPrompt]
+  if (newSystemPrompt === undefined) return undefined
+  return replaceSystemPrompt(input, newSystemPrompt)
+}
+
+/**
+ * Resolves whether the current call to `modelName` should be mocked.
+ * Mirrors `resolveMock` but reads `__ELASTICDASH_AI_MOCKS__` and
+ * `__ELASTICDASH_AI_CALL_COUNTERS__`.
+ */
+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 }
+}

package/src/workflow-runner-worker.ts

@@ -79,6 +79,17 @@ interface ToolMockConfig {
   [toolName: string]: ToolMockEntry
 }
 
+/** Per-model AI mock configuration (mirrors ToolMockConfig) */
+interface AIMockEntry {
+  mode: 'live' | 'mock-all' | 'mock-specific'
+  callIndices?: number[]
+  mockData?: Record<number, unknown>
+}
+
+interface AIMockConfig {
+  [modelName: string]: AIMockEntry
+}
+
 /** Minimal inline tool-wrapping — records each tool call to the trace. */
 async function loadAndWrapTools(
   toolsModulePath: string,
@@ -210,6 +221,10 @@ async function main() {
     agentState?: AgentState
     /** Optional tool mock configuration from the dashboard UI */
     toolMockConfig?: ToolMockConfig
+    /** Optional AI/LLM mock configuration from the dashboard UI */
+    aiMockConfig?: AIMockConfig
+    /** Optional prompt mock config: replacement system prompts keyed by original system prompt text */
+    promptMockConfig?: Record<string, string>
   }
   try {
     payload = JSON.parse(raw)
@@ -219,7 +234,7 @@ async function main() {
     return
   }
 
-  const { workflowsModulePath, toolsModulePath, workflowName, args, input, replayMode = false, checkpoint = 0, history = [], agentState, toolMockConfig } = payload
+  const { workflowsModulePath, toolsModulePath, workflowName, args, input, replayMode = false, checkpoint = 0, history = [], agentState, toolMockConfig, aiMockConfig, promptMockConfig } = payload
 
   const { context, finalise } = startTraceSession()
   setCurrentTrace(context.trace)
@@ -258,9 +273,12 @@ async function main() {
   let workflowError: Error | undefined
 
   try {
-    // Write mock config to globals so module-imported tools can read it via resolveMock()
+    // Write mock configs to globals so module-imported tools/AI can read them via resolveMock/resolveAIMock()
     ;(globalThis as any).__ELASTICDASH_TOOL_MOCKS__ = toolMockConfig ?? {}
     ;(globalThis as any).__ELASTICDASH_TOOL_CALL_COUNTERS__ = {}
+    ;(globalThis as any).__ELASTICDASH_AI_MOCKS__ = aiMockConfig ?? {}
+    ;(globalThis as any).__ELASTICDASH_AI_CALL_COUNTERS__ = {}
+    ;(globalThis as any).__ELASTICDASH_PROMPT_MOCKS__ = promptMockConfig ?? {}
 
     await installDBAutoInterceptor()
     installAIInterceptor()
@@ -321,6 +339,9 @@ async function main() {
     // Clear mock globals
     delete (globalThis as any).__ELASTICDASH_TOOL_MOCKS__
     delete (globalThis as any).__ELASTICDASH_TOOL_CALL_COUNTERS__
+    delete (globalThis as any).__ELASTICDASH_AI_MOCKS__
+    delete (globalThis as any).__ELASTICDASH_AI_CALL_COUNTERS__
+    delete (globalThis as any).__ELASTICDASH_PROMPT_MOCKS__
   }
 
   await recorder.flush()