keystone-cli 2.0.0 → 2.1.0

package/src/runner/executors/memory-executor.ts

@@ -43,45 +43,52 @@ export async function executeMemoryStep(
 
   const memoryDb = memoryDbFromOptions || new MemoryDb('.keystone/memory.db', dimension);
 
-  // Helper to get embedding using AI SDK
-  const getEmbedding = async (text: string): Promise<number[]> => {
-    const model = await getEmbeddingModel(modelName);
-    const result = await embed({
-      model,
-      value: text,
-      abortSignal,
-    });
-    return result.embedding;
-  };
+  try {
+    // Helper to get embedding using AI SDK
+    const getEmbedding = async (text: string): Promise<number[]> => {
+      const model = await getEmbeddingModel(modelName);
+      const result = await embed({
+        model,
+        value: text,
+        abortSignal,
+      });
+      return result.embedding;
+    };
 
-  switch (step.op) {
-    case 'store': {
-      const text = ExpressionEvaluator.evaluateString(step.text || '', context);
-      if (!text) throw new Error('Text is required for memory store operation');
+    switch (step.op) {
+      case 'store': {
+        const text = ExpressionEvaluator.evaluateString(step.text || '', context);
+        if (!text) throw new Error('Text is required for memory store operation');
 
-      const embedding = await getEmbedding(text);
-      const metadata = step.metadata || {};
-      const id = await memoryDb.store(text, embedding, metadata as Record<string, unknown>);
+        const embedding = await getEmbedding(text);
+        const metadata = step.metadata || {};
+        const id = await memoryDb.store(text, embedding, metadata as Record<string, unknown>);
 
-      return {
-        output: { id, status: 'stored' },
-        status: 'success',
-      };
-    }
-    case 'search': {
-      const query = ExpressionEvaluator.evaluateString(step.query || '', context);
-      if (!query) throw new Error('Query is required for memory search operation');
+        return {
+          output: { id, status: 'stored' },
+          status: 'success',
+        };
+      }
+      case 'search': {
+        const query = ExpressionEvaluator.evaluateString(step.query || '', context);
+        if (!query) throw new Error('Query is required for memory search operation');
 
-      const embedding = await getEmbedding(query);
-      const limit = step.limit || 5;
-      const results = await memoryDb.search(embedding, limit);
+        const embedding = await getEmbedding(query);
+        const limit = step.limit || 5;
+        const results = await memoryDb.search(embedding, limit);
 
-      return {
-        output: results,
-        status: 'success',
-      };
+        return {
+          output: results,
+          status: 'success',
+        };
+      }
+      default:
+        throw new Error(`Unknown memory operation: ${(step as any).op}`);
+    }
+  } finally {
+    // Only close if we created it ourselves
+    if (!memoryDbFromOptions) {
+      memoryDb.close();
     }
-    default:
-      throw new Error(`Unknown memory operation: ${(step as any).op}`);
   }
 }

package/src/runner/executors/shell-executor.ts

@@ -26,6 +26,7 @@
 import type { ExpressionContext } from '../../expression/evaluator.ts';
 import { ExpressionEvaluator } from '../../expression/evaluator.ts';
 import type { ShellStep } from '../../parser/schema.ts';
+import { ConfigLoader } from '../../utils/config-loader.ts';
 import { LIMITS } from '../../utils/constants.ts';
 import { filterSensitiveEnv } from '../../utils/env-filter.ts';
 import { ConsoleLogger, type Logger } from '../../utils/logger.ts';
@@ -146,6 +147,15 @@ function formatShellResult(result: ShellResult, logger: Logger): StepResult {
  */
 export function escapeShellArg(arg: unknown): string {
   const value = arg === null || arg === undefined ? '' : String(arg);
+
+  // Windows escaping (cmd.exe)
+  if (process.platform === 'win32') {
+    // Replace " with "" and wrap in double quotes
+    // This is the standard way to escape arguments for CRT-based programs in cmd
+    return `"${value.replace(/"/g, '""')}"`;
+  }
+
+  // POSIX escaping (sh)
   // Replace single quotes with '\'' (end quote, escaped quote, start quote)
   return `'${value.replace(/'/g, "'\\''")}'`;
 }
@@ -158,7 +168,7 @@ export interface ShellResult {
   stderrTruncated?: boolean;
 }
 
-const TRUNCATED_SUFFIX = '... [truncated output]';
+import { TRUNCATED_SUFFIX, createOutputLimiter } from '../../utils/stream-utils.ts';
 
 async function readStreamWithLimit(
   stream: ReadableStream<Uint8Array> | null | undefined,
@@ -173,65 +183,40 @@ async function readStreamWithLimit(
   }
 
   const reader = stream.getReader();
-  const decoder = new TextDecoder();
-  let text = '';
-  let bytesRead = 0;
+  const limiter = createOutputLimiter(maxBytes);
 
   while (true) {
     const { value, done } = await reader.read();
     if (done) break;
     if (!value) continue;
 
-    if (bytesRead + value.byteLength > maxBytes) {
-      const allowed = maxBytes - bytesRead;
-      if (allowed > 0) {
-        text += decoder.decode(value.slice(0, allowed), { stream: true });
-      }
-      text += decoder.decode();
+    limiter.append(Buffer.from(value));
+
+    if (limiter.truncated) {
       try {
         await reader.cancel();
       } catch {}
-      return { text: `${text}${TRUNCATED_SUFFIX}`, truncated: true };
+      break;
     }
-
-    bytesRead += value.byteLength;
-    text += decoder.decode(value, { stream: true });
   }
 
-  text += decoder.decode();
-  return { text, truncated: false };
+  return { text: limiter.finalize(), truncated: limiter.truncated };
 }
 
 // Whitelist of allowed characters for secure shell command execution
-// Allows: Alphanumeric, space, and common safe punctuation (_ . / : @ , + - = ' " ! ~)
-// Blocks: Newlines (\n, \r), Pipes, redirects, subshells, variables ($), etc.
-const SAFE_SHELL_CHARS = /^[a-zA-Z0-9 _./:@,+=~'"!-]+$/;
+// Allows: Alphanumeric, space, and common safe punctuation (_ . / : @ , + - =)
+// Blocks: Quotes, Newlines, Pipes, redirects, subshells, variables, backslashes, etc.
+const SAFE_SHELL_CHARS = /^[a-zA-Z0-9 _./:@,+=~"'-]+$/;
 
 export function detectShellInjectionRisk(rawCommand: string): boolean {
-  // We scan the command to handle single quotes correctly.
-  // Characters inside single quotes are considered escaped/literal and safe from shell injection.
-  let inSingleQuote = false;
-
-  for (let i = 0; i < rawCommand.length; i++) {
-    const char = rawCommand[i];
-
-    if (char === "'") {
-      inSingleQuote = !inSingleQuote;
-      continue;
-    }
-
-    // Outside single quotes, we enforce the strict whitelist
-    if (!inSingleQuote) {
-      if (!SAFE_SHELL_CHARS.test(char)) {
-        return true;
-      }
-    }
-    // Inside single quotes, everything is treated as a literal string by the shell,
-    // so we don't need to block special characters.
-  }
-
-  // If we ended with an unclosed single quote, it's a syntax risk
-  return inSingleQuote;
+  // We can safely ignore anything inside single quotes because our escape()
+  // function (which is the recommended way to interpolate) uses single quotes
+  // and correctly escapes nested single quotes as '\''.
+  // This regex matches '...' including correctly escaped internal single quotes.
+  const quotedRegex = /'([^']|'\\'')*'/g;
+  const stripped = rawCommand.replace(quotedRegex, "'QUOTED_STR'");
+
+  return !SAFE_SHELL_CHARS.test(stripped);
 }
 
 /**
@@ -256,13 +241,59 @@ export async function executeShell(
   if (!step.allowInsecure) {
     if (detectShellInjectionRisk(command)) {
       throw new Error(
-        `Security Error: Command execution blocked.\nCommand: "${command.substring(0, 100)}${
+        `Security Error: Command execution blocked to prevent potential shell injection.\nCommand: "${command.substring(0, 100)}${
           command.length > 100 ? '...' : ''
-        }"\nReason: Contains characters not in the strict whitelist (alphanumeric, whitespace, and _./:@,+=~-).\nThis protects against shell injection attacks.\nFix: either simplify your command or set 'allowInsecure: true' in your step definition if you trust the input.`
+        }"\nReason: Contains characters not in the strict whitelist (alphanumeric, whitespace, and _./:@,+=~-).\nThis protects against chaining malicious commands (e.g. '; rm -rf /'). It does NOT evaluate if the command itself is destructive.\nFix: either simplify your command or set 'allowInsecure: true' in your step definition if you trust the input.`
+      );
+    }
+
+    // Additional Check: Prevent Directory Traversal in Binary Path
+    // Even if it passes the whitelist, we don't want to allow 'cat ../../../etc/passwd'
+    // or executing '../../../../bin/malice'.
+    // We check for '..' characters which might indicate directory traversal.
+    if (command.includes('..') && (command.includes('/') || command.includes('\\'))) {
+      throw new Error(
+        `Security Error: Command blocked due to potential directory traversal ('..').\nCommand: "${command.substring(0, 100)}"\nTo allow relative paths outside the current directory, set 'allowInsecure: true'.`
       );
     }
   }
 
+  // Security Check: Enforce Denylist (e.g. rm, mkfs, etc.)
+  // We check this even if allowInsecure is true, because these are explicitly banned by policy.
+  const config = ConfigLoader.load();
+  if (config.engines?.denylist && config.engines.denylist.length > 0) {
+    // Robust parsing to get the command binary
+    // This handles:
+    // 1. Chained commands (e.g. "echo foo; rm -rf /")
+    // 2. Pre-command modifiers (e.g. "watch rm") - though difficult to do perfectly without a full shell parser,
+    //    we can check for common dangerous patterns or just strictly check tokens.
+    //
+    // Strategy: Tokenize by shell delimiters (;, |, &, &&, ||, ``, $()) and check the first word of each segment.
+
+    // Split by command separators
+    const segments = command.split(/[;|&]|\$\(|\`|\r?\n/);
+
+    for (const segment of segments) {
+      if (!segment.trim()) continue;
+
+      // Get the first token of the segment
+      const tokens = segment.trim().split(/\s+/);
+      let bin = tokens[0];
+
+      // Handle path prefixes (e.g. /bin/rm -> rm)
+      if (bin.includes('/')) {
+        const parts = bin.split(/[/\\]/);
+        bin = parts[parts.length - 1];
+      }
+
+      if (config.engines.denylist.includes(bin)) {
+        throw new Error(
+          `Security Error: Command "${bin}" is in the denylist and cannot be executed.`
+        );
+      }
+    }
+  }
+
   // Evaluate environment variables
   const env: Record<string, string> = context.env ? { ...context.env } : {};
   if (step.env) {
@@ -300,10 +331,14 @@ export async function executeShell(
     let stderrTruncated = false;
     const maxOutputBytes = LIMITS.MAX_PROCESS_OUTPUT_BYTES;
 
-    // Use 'sh -c' for everything to ensure consistent argument parsing
+    // Use 'sh -c' (POSIX) or 'cmd.exe /d /s /c' (Windows)
     // Security is guaranteed by the strict whitelist check above for allowInsecure: false
     // which prevents injection of metacharacters, quotes, escapes, etc.
-    const proc = Bun.spawn(['sh', '-c', command], {
+    const isWindows = process.platform === 'win32';
+    const shellCommand = isWindows ? 'cmd.exe' : 'sh';
+    const shellArgs = isWindows ? ['/d', '/s', '/c'] : ['-c'];
+
+    const proc = Bun.spawn([shellCommand, ...shellArgs, command], {
       cwd: cwd || process.cwd(),
       env: mergedEnv,
       stdout: 'pipe',
@@ -323,9 +358,15 @@ export async function executeShell(
     const stdoutPromise = readStreamWithLimit(proc.stdout, maxOutputBytes);
     const stderrPromise = readStreamWithLimit(proc.stderr, maxOutputBytes);
 
-    // Wait for exit
-    exitCode = await proc.exited;
-    const [stdoutResult, stderrResult] = await Promise.all([stdoutPromise, stderrPromise]);
+    // Wait for exit and streams simultaneously to prevent deadlocks
+    // (If the pipe fills up, the process blocks on write. If we await exit first, we never drain the pipe -> Deadlock)
+    const [exitResult, stdoutResult, stderrResult] = await Promise.all([
+      proc.exited,
+      stdoutPromise,
+      stderrPromise,
+    ]);
+
+    exitCode = exitResult;
 
     stdoutString = stdoutResult.text;
     stderrString = stderrResult.text;
@@ -406,8 +447,11 @@ export async function executeShellArgs(
     const stdoutPromise = readStreamWithLimit(proc.stdout, maxOutputBytes);
     const stderrPromise = readStreamWithLimit(proc.stderr, maxOutputBytes);
 
-    const exitCode = await proc.exited;
-    const [stdoutResult, stderrResult] = await Promise.all([stdoutPromise, stderrPromise]);
+    const [exitCode, stdoutResult, stderrResult] = await Promise.all([
+      proc.exited,
+      stdoutPromise,
+      stderrPromise,
+    ]);
 
     if (abortSignal) {
       abortSignal.removeEventListener('abort', abortHandler);

package/src/runner/executors/subworkflow-executor.ts

@@ -35,6 +35,8 @@ export async function executeSubWorkflow(
     parentDepth: number;
     parentOptions: any;
     abortSignal?: AbortSignal;
+    stepExecutionId?: string;
+    parentDb?: any; // WorkflowDb
   }
 ): Promise<StepResult> {
   if (options.abortSignal?.aborted) {
@@ -57,6 +59,7 @@ export async function executeSubWorkflow(
     ...options.parentOptions,
     inputs,
     dbPath: options.parentDbPath,
+    db: options.parentDb, // Reuse existing DB connection
     logger: options.parentLogger,
     mcpManager: options.parentMcpManager,
     workflowDir: subWorkflowDir,
@@ -64,6 +67,19 @@ export async function executeSubWorkflow(
     signal: options.abortSignal,
   });
 
+  // Track sub-workflow run ID in parent step metadata for rollback safety
+  if (options.stepExecutionId && options.parentDb) {
+    try {
+      await options.parentDb.updateStepMetadata(options.stepExecutionId, {
+        __subRunId: subRunner.runId,
+      });
+    } catch (error) {
+      options.parentLogger.warn(
+        `Failed to store sub-workflow run ID in metadata: ${error instanceof Error ? error.message : String(error)}`
+      );
+    }
+  }
+
   try {
     const output = await subRunner.run();
 

package/src/runner/executors/types.ts

@@ -46,7 +46,8 @@ export interface StepExecutorOptions {
   executeWorkflowFn?: (
     step: WorkflowStep,
     context: ExpressionContext,
-    abortSignal?: AbortSignal
+    abortSignal?: AbortSignal,
+    stepExecutionId?: string
   ) => Promise<StepResult>;
   mcpManager?: MCPManager;
   db?: WorkflowDb;
@@ -62,6 +63,7 @@ export interface StepExecutorOptions {
   debug?: boolean;
   allowInsecure?: boolean;
   emitEvent?: (event: WorkflowEvent) => void;
+  depth?: number;
 
   executeStep?: (step: Step, context: ExpressionContext) => Promise<StepResult>; // To avoid circular dependency
   executeLlmStep?: typeof executeLlmStep;

package/src/runner/executors/verification_fixes.test.ts

@@ -0,0 +1,46 @@
+import { describe, expect, test } from 'bun:test';
+import { validateRemoteUrl } from '../mcp-client';
+import { executeShell } from './shell-executor';
+
+describe('Verification Fixes', () => {
+  describe('SSRF Protection (mcp-client)', () => {
+    test('validateRemoteUrl should throw on 127.0.0.1', async () => {
+      expect(validateRemoteUrl('https://127.0.0.1')).rejects.toThrow('SSRF Protection');
+    });
+
+    test('validateRemoteUrl should throw on localhost', async () => {
+      expect(validateRemoteUrl('https://localhost')).rejects.toThrow('SSRF Protection');
+    });
+
+    test('validateRemoteUrl should throw on metadata IP', async () => {
+      expect(validateRemoteUrl('https://169.254.169.254')).rejects.toThrow('SSRF Protection');
+    });
+  });
+
+  describe('Shell Path Traversal (shell-executor)', () => {
+    const mockContext = { env: {}, steps: {}, inputs: {}, envOverrides: {}, secrets: {} };
+
+    test('should block command with ".." and "/" in secure mode', async () => {
+      const step = {
+        id: 'test',
+        type: 'shell' as const,
+        run: 'cat ../secret.txt',
+        allowInsecure: false,
+      };
+      // It should throw BEFORE spawning
+      // The error message I added was "Directory Traversal" or similar
+      // Let's check the implementation: "Command blocked due to potential directory traversal"
+      await expect(executeShell(step, mockContext)).rejects.toThrow('Command blocked');
+    });
+
+    test('should block absolute path with ".." in secure mode', async () => {
+      const step = {
+        id: 'test',
+        type: 'shell' as const,
+        run: '/bin/ls ../',
+        allowInsecure: false,
+      };
+      await expect(executeShell(step, mockContext)).rejects.toThrow('Command blocked');
+    });
+  });
+});

package/src/runner/join-scheduling.test.ts

@@ -123,7 +123,7 @@ describe('Join Scheduling & Resume', () => {
   });
 
   it('should resume and retry a step that previously exhausted retries', async () => {
-    const dbPath = 'test-resume-retry.db';
+    const dbPath = `test-resume-retry-${Date.now()}.db`;
     if (existsSync(dbPath)) rmSync(dbPath);
 
     const counterFile = `/tmp/keystone-test-resume-${Date.now()}.txt`;
@@ -177,6 +177,7 @@ describe('Join Scheduling & Resume', () => {
     // Verify it failed twice (initial + 1 retry)
     let val = await Bun.file(counterFile).text();
     expect(val.trim()).toBe('2');
+    await runner1.stop();
 
     // Now resume. It should try again (Run 3) and succeed.
     const runner2 = new WorkflowRunner(workflow, {

package/src/runner/llm-adapter.integration.test.ts

@@ -54,10 +54,13 @@ describe('LLM Adapter (AI SDK)', () => {
         model_mappings: {},
       } as any);
 
-      // With shared setupLlmMocks, we expect 'mock' provider
+      // Mock the provider to return a callable function that returns a mock model
+      const mockProvider = (modelId: string) => mockLanguageModel;
+      spyOn(DynamicProviderRegistry, 'getProvider').mockResolvedValue(() => mockProvider);
+
       const model = (await getModel('model-name')) as any;
-      expect(model.modelId).toBe('mock-model');
-      expect(model.provider).toBe('mock');
+      expect(model.modelId).toBe('test-model');
+      expect(model.provider).toBe('test-provider');
     });
 
     it('should handle auth token retrieval for standard providers', async () => {
@@ -73,11 +76,13 @@ describe('LLM Adapter (AI SDK)', () => {
         model_mappings: {},
       } as any);
 
+      // Mock the provider to return a callable function
+      const mockProvider = (modelId: string) => mockLanguageModel;
+      spyOn(DynamicProviderRegistry, 'getProvider').mockResolvedValue(() => mockProvider);
       spyOn(ConfigLoader, 'getSecret').mockReturnValue('fake-token');
 
       const model = (await getModel('gpt-4')) as any;
-      // With global mock, we mostly check it didn't throw and loaded the 'mock' provider
-      expect(model.provider).toBe('mock');
+      expect(model.provider).toBe('test-provider');
       expect(ConfigLoader.getSecret).toHaveBeenCalledWith('OPENAI_API_KEY');
     });
   });

package/src/runner/llm-adapter.ts

@@ -55,7 +55,7 @@ class LocalEmbeddingModel {
         return Array.from(output.data) as number[];
       })
     );
-    return { embeddings };
+    return { embeddings, warnings: [] };
   }
 
   /**
@@ -72,23 +72,30 @@ export type { LanguageModel, EmbeddingModel } from 'ai';
 
 const userRequire = createRequire(join(process.cwd(), 'package.json'));
 
-// Lazy-loaded global require to avoid blocking import time
+// Lazy-loaded global require
 let globalRequire: NodeRequire | undefined;
-let globalRequireResolved = false;
-
-function getGlobalRequire(): NodeRequire | undefined {
-  if (globalRequireResolved) {
+let globalRequirePromise: Promise<NodeRequire | undefined> | null = null;
+
+async function getGlobalRequire(): Promise<NodeRequire | undefined> {
+  if (globalRequire) return globalRequire;
+  if (globalRequirePromise) return globalRequirePromise;
+
+  globalRequirePromise = (async () => {
+    try {
+      const { exec } = await import('node:child_process');
+      const { promisify } = await import('node:util');
+      const execAsync = promisify(exec);
+
+      const { stdout } = await execAsync('npm root -g', { encoding: 'utf-8', timeout: 5000 });
+      const globalRoot = stdout.trim();
+      globalRequire = createRequire(join(globalRoot, 'package.json'));
+    } catch {
+      // Global npm root not found or command failed
+    }
     return globalRequire;
-  }
-  globalRequireResolved = true;
-  try {
-    const globalRoot = execSync('npm root -g', { encoding: 'utf-8' }).trim();
-    globalRequire = createRequire(join(globalRoot, 'package.json'));
-  } catch {
-    // Global npm root not found - this is expected in some environments (e.g., containers, CI)
-    // Global package resolution will be disabled silently
-  }
-  return globalRequire;
+  })();
+
+  return globalRequirePromise;
 }
 
 // Compatibility types for Keystone
@@ -157,14 +164,15 @@ export class DynamicProviderRegistry {
         let pkg: any;
         try {
           // Try local project first
-          pkg = await import(config.package);
+          const localPath = userRequire.resolve(config.package);
+          pkg = await import(localPath);
         } catch {
           try {
             const pkgPath = userRequire.resolve(config.package);
             pkg = await import(pkgPath);
           } catch {
             // Try global if local resolution fails
-            const globalReq = getGlobalRequire();
+            const globalReq = await getGlobalRequire();
             if (globalReq) {
               try {
                 const globalPkgPath = globalReq.resolve(config.package);
@@ -218,8 +226,22 @@ export class DynamicProviderRegistry {
           return pkg.default;
         }
 
+        // Check for standard generic export names
+        if (typeof pkg.createProvider === 'function') {
+          DynamicProviderRegistry.loadedProviders.set(providerName, pkg.createProvider);
+          return pkg.createProvider;
+        }
+        if (typeof pkg.provider === 'function') {
+          DynamicProviderRegistry.loadedProviders.set(providerName, pkg.provider);
+          return pkg.provider;
+        }
+
         const firstFn = Object.values(pkg).find((v) => typeof v === 'function');
         if (firstFn) {
+          // Warn about loose resolution only if we really had to fall back this far
+          new ConsoleLogger().warn(
+            `[Keystone] Warning: Provider '${providerName}' resolution fell back to the first exported function found in '${config.package}'. This may be unstable.`
+          );
           DynamicProviderRegistry.loadedProviders.set(providerName, firstFn as any);
           return firstFn as any;
         }
@@ -298,10 +320,22 @@ async function prepareProvider(
 }
 
 export async function getModel(model: string): Promise<LanguageModel> {
+  const configValues = ConfigLoader.load();
+  if (configValues.logging?.suppress_ai_sdk_warnings) {
+    process.env.AI_SDK_LOG_WARNINGS = 'false';
+  }
+
   const { provider, resolvedModel } = await prepareProvider(model);
 
   // AI SDK convention: provider(modelId)
   if (typeof provider === 'function') {
+    // Prefer explicit .chat() or .chatModel() if available to ensure correct protocol (Chat vs Completion)
+    if (typeof (provider as any).chat === 'function') {
+      return (provider as any).chat(resolvedModel);
+    }
+    if (typeof (provider as any).chatModel === 'function') {
+      return (provider as any).chatModel(resolvedModel);
+    }
     return (provider as any)(resolvedModel);
   }
 
@@ -321,6 +355,11 @@ export async function getModel(model: string): Promise<LanguageModel> {
 }
 
 export async function getEmbeddingModel(model: string): Promise<EmbeddingModel> {
+  const configValues = ConfigLoader.load();
+  if (configValues.logging?.suppress_ai_sdk_warnings) {
+    process.env.AI_SDK_LOG_WARNINGS = 'false';
+  }
+
   // 1. Check for local fallback
   if (model === 'local' || model === 'keystone-local') {
     return new LocalEmbeddingModel();

package/src/runner/llm-clarification.test.ts

@@ -14,6 +14,7 @@ import * as agentParser from '../parser/agent-parser';
 import type { Config } from '../parser/config-schema';
 import type { Agent, LlmStep, Step } from '../parser/schema';
 import { ConfigLoader } from '../utils/config-loader';
+import * as llmAdapter from './llm-adapter';
 import type { LLMMessage } from './llm-adapter';
 import type { StepResult } from './step-executor';
 
@@ -28,10 +29,11 @@ let currentChatFn: (messages: any[], options?: any) => Promise<MockLLMResponse>;
 describe('LLM Clarification', () => {
   let resolveAgentPathSpy: ReturnType<typeof spyOn>;
   let parseAgentSpy: ReturnType<typeof spyOn>;
+  let getModelSpy: ReturnType<typeof spyOn>;
 
   beforeAll(async () => {
     setupLlmMocks();
-    mockGetModel.mockResolvedValue(createUnifiedMockModel());
+    getModelSpy = spyOn(llmAdapter, 'getModel').mockResolvedValue(createUnifiedMockModel() as any);
     const module = await import('./executors/llm-executor.ts');
     executeLlmStep = module.executeLlmStep;
   });
@@ -64,6 +66,7 @@ describe('LLM Clarification', () => {
     ConfigLoader.clear();
     resolveAgentPathSpy.mockRestore();
     parseAgentSpy.mockRestore();
+    getModelSpy?.mockClear();
     resetLlmMocks();
   });