keystone-cli 0.7.2 → 1.0.0

package/src/runner/optimization-runner.ts

@@ -1,10 +1,8 @@
-import { readFileSync, writeFileSync } from 'node:fs';
 import { dirname } from 'node:path';
-import { stringify } from 'yaml';
-import { parseAgent, resolveAgentPath } from '../parser/agent-parser';
+import { resolveAgentPath } from '../parser/agent-parser';
 import type { LlmStep, Step, Workflow } from '../parser/schema';
-import { extractJson } from '../utils/json-parser';
-import { getAdapter } from './llm-adapter';
+import { TIMEOUTS } from '../utils/constants';
+import { ConsoleLogger, type Logger } from '../utils/logger';
 import { executeLlmStep } from './llm-executor';
 import { WorkflowRunner } from './workflow-runner';
 
@@ -13,6 +11,7 @@ export interface OptimizationOptions {
   targetStepId: string;
   inputs?: Record<string, unknown>;
   iterations?: number;
+  logger?: Logger;
 }
 
 export class OptimizationRunner {
@@ -21,6 +20,8 @@ export class OptimizationRunner {
   private targetStepId: string;
   private iterations: number;
   private inputs: Record<string, unknown>;
+  private logger: Logger;
+  private secrets: Record<string, string>;
 
   constructor(workflow: Workflow, options: OptimizationOptions) {
     this.workflow = workflow;
@@ -28,6 +29,18 @@ export class OptimizationRunner {
     this.targetStepId = options.targetStepId;
     this.iterations = options.iterations || 5;
     this.inputs = options.inputs || {};
+    this.logger = options.logger || new ConsoleLogger();
+    this.secrets = OptimizationRunner.extractSecrets(Bun.env);
+  }
+
+  private static extractSecrets(env: Record<string, string | undefined>): Record<string, string> {
+    const secrets: Record<string, string> = {};
+    for (const [key, value] of Object.entries(env)) {
+      if (typeof value === 'string') {
+        secrets[key] = value;
+      }
+    }
+    return secrets;
   }
 
   public async optimize(): Promise<{ bestPrompt: string; bestScore: number }> {
@@ -40,8 +53,8 @@ export class OptimizationRunner {
       throw new Error(`Target step "${this.targetStepId}" not found or is not an LLM/Shell step`);
     }
 
-    console.log(`\n🚀 Optimizing step: ${this.targetStepId} (${targetStep.type})`);
-    console.log(`📊 Iterations: ${this.iterations}`);
+    this.logger.log(`\n🚀 Optimizing step: ${this.targetStepId} (${targetStep.type})`);
+    this.logger.log(`📊 Iterations: ${this.iterations}`);
 
     let bestPrompt =
       targetStep.type === 'llm'
@@ -52,8 +65,8 @@ export class OptimizationRunner {
     let currentPrompt = bestPrompt;
 
     for (let i = 1; i <= this.iterations; i++) {
-      console.log(`\n--- Iteration ${i}/${this.iterations} ---`);
-      console.log(
+      this.logger.log(`\n--- Iteration ${i}/${this.iterations} ---`);
+      this.logger.log(
         `Current Prompt: ${currentPrompt.substring(0, 100)}${currentPrompt.length > 100 ? '...' : ''}`
       );
 
@@ -74,18 +87,19 @@ export class OptimizationRunner {
       const runner = new WorkflowRunner(modifiedWorkflow, {
         inputs: this.inputs,
         workflowDir: dirname(this.workflowPath),
+        logger: this.logger,
       });
 
       const outputs = await runner.run();
 
       // 2. Evaluate the result
       const score = await this.evaluate(outputs);
-      console.log(`Score: ${score}/100`);
+      this.logger.log(`Score: ${score}/100`);
 
       if (score > bestScore) {
         bestScore = score;
         bestPrompt = currentPrompt;
-        console.log(`✨ New best score: ${bestScore}`);
+        this.logger.log(`✨ New best score: ${bestScore}`);
       }
 
       // 3. Suggest next prompt (if not last iteration)
@@ -103,15 +117,18 @@ export class OptimizationRunner {
     if (!evalConfig) return 0;
 
     if (evalConfig.scorer === 'script') {
-      // Note: getAdapter already imported at top level
-      const { executeStep } = await import('./step-executor');
+      if (!evalConfig.allowInsecure) {
+        throw new Error(
+          'Eval script execution is disabled by default. Set eval.allowInsecure: true to enable.'
+        );
+      }
+      const allowSecrets = evalConfig.allowSecrets === true;
 
       // Create a context with outputs available
       const context = {
         inputs: this.inputs,
         steps: {},
-        // biome-ignore lint/suspicious/noExplicitAny: environment access
-        secrets: Bun.env as any,
+        secrets: allowSecrets ? this.secrets : {},
         env: this.workflow.env,
         outputs, // Direct access
         output: outputs, // For convenience
@@ -136,7 +153,9 @@ export class OptimizationRunner {
       // Note: OptimizationRunner should probably import executeStep
       const { SafeSandbox } = await import('../utils/sandbox');
       try {
-        const result = await SafeSandbox.execute(scriptStep.run, context, { timeout: 5000 });
+        const result = await SafeSandbox.execute(scriptStep.run, context, {
+          timeout: TIMEOUTS.DEFAULT_SCRIPT_TIMEOUT_MS,
+        });
         if (typeof result === 'object' && result !== null && 'stdout' in result) {
           // biome-ignore lint/suspicious/noExplicitAny: result typing
           const match = (result as any).stdout.match(/\d+/);
@@ -150,14 +169,14 @@ export class OptimizationRunner {
           if (match) return Number.parseInt(match[0], 10);
         }
       } catch (e) {
-        console.error('Eval script failed:', e);
+        this.logger.error(`Eval script failed: ${String(e)}`);
       }
       return 0;
     }
 
     // LLM Scorer
     if (!evalConfig.agent || !evalConfig.prompt) {
-      console.warn('Skipping LLM evaluation: agent or prompt missing');
+      this.logger.warn('Skipping LLM evaluation: agent or prompt missing');
       return 0;
     }
 
@@ -168,7 +187,7 @@ export class OptimizationRunner {
       prompt: `${evalConfig.prompt}\n\nOutputs to evaluate:\n${JSON.stringify(outputs, null, 2)}`,
       needs: [],
       maxIterations: 10,
-      schema: {
+      outputSchema: {
         type: 'object',
         properties: {
           score: { type: 'number', minimum: 0, maximum: 100 },
@@ -182,8 +201,7 @@ export class OptimizationRunner {
     const context = {
       inputs: this.inputs,
       steps: {},
-      // biome-ignore lint/suspicious/noExplicitAny: environment access
-      secrets: Bun.env as any,
+      secrets: this.secrets,
       env: this.workflow.env,
     };
 
@@ -194,7 +212,7 @@ export class OptimizationRunner {
       async () => {
         throw new Error('Tools not supported in eval');
       },
-      console
+      this.logger
     );
 
     if (result.status === 'success' && result.output && typeof result.output === 'object') {
@@ -244,8 +262,7 @@ Return ONLY the new prompt text.`,
     const context = {
       inputs: this.inputs,
       steps: {},
-      // biome-ignore lint/suspicious/noExplicitAny: environment access
-      secrets: Bun.env as any,
+      secrets: this.secrets,
       env: this.workflow.env,
     };
 
@@ -257,7 +274,7 @@ Return ONLY the new prompt text.`,
         async () => {
           throw new Error('Tools not supported in meta-opt');
         },
-        console,
+        this.logger,
         undefined,
         dirname(this.workflowPath) // Pass workflowDir to resolve agent
       );
@@ -265,7 +282,7 @@ Return ONLY the new prompt text.`,
         return result.output.trim();
       }
     } catch (e) {
-      console.warn(`  ⚠️ Meta-optimizer failed: ${e instanceof Error ? e.message : String(e)}`);
+      this.logger.warn(`  ⚠️ Meta-optimizer failed: ${e instanceof Error ? e.message : String(e)}`);
       // Adding a dummy mutation for testing purposes if env var is set
       if (Bun.env.TEST_OPTIMIZER) {
         return `${currentPrompt}!`;
@@ -276,28 +293,28 @@ Return ONLY the new prompt text.`,
   }
 
   private async saveBestPrompt(prompt: string): Promise<void> {
-    console.log(`\n💾 Saving best prompt to ${this.workflowPath}`);
+    this.logger.log(`\n💾 Saving best prompt to ${this.workflowPath}`);
 
     // We need to be careful here. The prompt might be in the workflow YAML directly,
     // or it might be in an agent file.
 
     const targetStep = this.workflow.steps.find((s) => s.id === this.targetStepId);
 
-    console.log(`--- BEST PROMPT/RUN ---\n${prompt}\n-----------------------`);
+    this.logger.log(`--- BEST PROMPT/RUN ---\n${prompt}\n-----------------------`);
 
     if (targetStep?.type === 'llm') {
       const agentPath = resolveAgentPath((targetStep as LlmStep).agent, dirname(this.workflowPath));
       try {
         // For MVP, we just logged it. Automatic replacement in arbitrary files is risky without robust parsing.
         // But we can try to warn/notify.
-        console.log(
+        this.logger.log(
           `To apply this optimization, update the 'systemPrompt' or instruction in: ${agentPath}`
         );
       } catch (e) {
-        console.warn(`Could not locate agent file: ${e}`);
+        this.logger.warn(`Could not locate agent file: ${String(e)}`);
       }
     } else {
-      console.log(
+      this.logger.log(
         `To apply this optimization, update the 'run' command for step '${this.targetStepId}' in ${this.workflowPath}`
       );
     }

package/src/runner/reflexion.test.ts

@@ -46,7 +46,7 @@ describe('WorkflowRunner Reflexion', () => {
 
     // biome-ignore lint/suspicious/noExplicitAny: Accessing private property for testing
     const db = (runner as any).db;
-    await db.createRun(runner.getRunId(), workflow.name, {});
+    await db.createRun(runner.runId, workflow.name, {});
 
     const spy = jest.spyOn(StepExecutor, 'executeStep');
 

package/src/runner/resource-pool.test.ts

@@ -0,0 +1,113 @@
+import { describe, expect, it, mock } from 'bun:test';
+import { ConsoleLogger } from '../utils/logger';
+import { ResourcePoolManager } from './resource-pool';
+
+describe('ResourcePoolManager', () => {
+  const logger = new ConsoleLogger();
+
+  it('should respect pool limits', async () => {
+    const manager = new ResourcePoolManager(logger, {
+      pools: { test: 2 },
+    });
+
+    let activeCount = 0;
+    const run = async () => {
+      const release = await manager.acquire('test');
+      activeCount++;
+      expect(activeCount).toBeLessThanOrEqual(2);
+      await Bun.sleep(50);
+      activeCount--;
+      release();
+    };
+
+    await Promise.all([run(), run(), run(), run()]);
+  });
+
+  it('should use default limit for unknown pools', async () => {
+    const manager = new ResourcePoolManager(logger, {
+      defaultLimit: 3,
+    });
+
+    let activeCount = 0;
+    const run = async () => {
+      const release = await manager.acquire('unknown');
+      activeCount++;
+      expect(activeCount).toBeLessThanOrEqual(3);
+      await Bun.sleep(50);
+      activeCount--;
+      release();
+    };
+
+    await Promise.all([run(), run(), run(), run(), run()]);
+  });
+
+  it('should handle cancellation via AbortSignal', async () => {
+    const manager = new ResourcePoolManager(logger, {
+      pools: { test: 1 },
+    });
+
+    // Acquire the only slot
+    const release1 = await manager.acquire('test');
+
+    const controller = new AbortController();
+    const pendingAcquisition = manager.acquire('test', { signal: controller.signal });
+
+    // Cancel after a bit
+    setTimeout(() => controller.abort(), 10);
+
+    await expect(pendingAcquisition).rejects.toThrow('Acquisition aborted');
+    release1();
+  });
+
+  it('should respect priority in queue', async () => {
+    const manager = new ResourcePoolManager(logger, {
+      pools: { test: 1 },
+    });
+
+    const results: number[] = [];
+    const release1 = await manager.acquire('test');
+
+    const p1 = manager.acquire('test', { priority: 1 }).then((r) => {
+      results.push(1);
+      r();
+    });
+    const p2 = manager.acquire('test', { priority: 10 }).then((r) => {
+      results.push(10);
+      r();
+    });
+    const p3 = manager.acquire('test', { priority: 5 }).then((r) => {
+      results.push(5);
+      r();
+    });
+
+    release1();
+    await Promise.all([p1, p2, p3]);
+
+    // Priority 10 should run first, then 5, then 1
+    expect(results).toEqual([10, 5, 1]);
+  });
+
+  it('should provide metrics', async () => {
+    const manager = new ResourcePoolManager(logger, {
+      pools: { test: 2 },
+    });
+
+    const release1 = await manager.acquire('test');
+    const release2 = await manager.acquire('test');
+
+    const metrics = manager.getMetrics('test');
+    expect(metrics?.active).toBe(2);
+    expect(metrics?.queued).toBe(0);
+
+    const p3 = manager.acquire('test');
+    const metrics2 = manager.getMetrics('test');
+    expect(metrics2?.queued).toBe(1);
+
+    release1();
+    release2();
+    await p3.then((r) => r());
+
+    const metrics3 = manager.getMetrics('test');
+    expect(metrics3?.totalAcquired).toBe(3);
+  });
+});

package/src/runner/resource-pool.ts

@@ -0,0 +1,164 @@
+import type { Logger } from '../utils/logger.ts';
+
+export type ReleaseFunction = () => void;
+
+export interface PoolMetrics {
+  name: string;
+  limit: number;
+  active: number;
+  queued: number;
+  totalAcquired: number;
+  totalWaitTimeMs: number;
+}
+
+interface QueuedRequest {
+  resolve: (release: ReleaseFunction) => void;
+  reject: (error: Error) => void;
+  signal?: AbortSignal;
+  priority: number;
+  timestamp: number;
+}
+
+export class ResourcePoolManager {
+  private pools = new Map<
+    string,
+    {
+      limit: number;
+      active: number;
+      queue: QueuedRequest[];
+      totalAcquired: number;
+      totalWaitTimeMs: number;
+    }
+  >();
+  private globalLimit: number;
+
+  constructor(
+    private logger: Logger,
+    options: { defaultLimit?: number; pools?: Record<string, number> } = {}
+  ) {
+    this.globalLimit = options.defaultLimit || 10;
+    if (options.pools) {
+      for (const [name, limit] of Object.entries(options.pools)) {
+        this.pools.set(name, { limit, active: 0, queue: [], totalAcquired: 0, totalWaitTimeMs: 0 });
+      }
+    }
+  }
+
+  /**
+   * Acquire a resource from a pool.
+   * If the pool doesn't exist, it uses the global limit.
+   */
+  async acquire(
+    poolName: string,
+    options: { priority?: number; signal?: AbortSignal } = {}
+  ): Promise<ReleaseFunction> {
+    let pool = this.pools.get(poolName);
+    if (!pool) {
+      // Create a pool for this name if it doesn't exist, using global limit
+      pool = {
+        limit: this.globalLimit,
+        active: 0,
+        queue: [],
+        totalAcquired: 0,
+        totalWaitTimeMs: 0,
+      };
+      this.pools.set(poolName, pool);
+    }
+
+    if (pool.active < pool.limit && pool.queue.length === 0) {
+      pool.active++;
+      pool.totalAcquired++;
+      return this.createReleaseFn(poolName);
+    }
+
+    // Queue the request
+    const timestamp = Date.now();
+    const poolRef = pool;
+    return new Promise<ReleaseFunction>((resolve, reject) => {
+      const request: QueuedRequest = {
+        resolve,
+        reject,
+        signal: options.signal,
+        priority: options.priority || 0,
+        timestamp,
+      };
+
+      // Add to queue and sort by priority (desc) then timestamp (asc)
+      poolRef.queue.push(request);
+      poolRef.queue.sort((a, b) => {
+        if (b.priority !== a.priority) return b.priority - a.priority;
+        return a.timestamp - b.timestamp;
+      });
+
+      // Handle abort signal
+      if (options.signal) {
+        options.signal.addEventListener(
+          'abort',
+          () => {
+            const index = poolRef.queue.indexOf(request);
+            if (index !== -1) {
+              poolRef.queue.splice(index, 1);
+              reject(new Error('Acquisition aborted'));
+            }
+          },
+          { once: true }
+        );
+      }
+    });
+  }
+
+  private createReleaseFn(poolName: string): ReleaseFunction {
+    let released = false;
+    return () => {
+      if (released) return;
+      released = true;
+      this.release(poolName);
+    };
+  }
+
+  private release(poolName: string) {
+    const pool = this.pools.get(poolName);
+    if (!pool) return;
+
+    pool.active--;
+
+    // Process queue
+    while (pool.active < pool.limit && pool.queue.length > 0) {
+      const request = pool.queue.shift();
+      if (!request) break;
+
+      // Skip if signal already aborted
+      if (request.signal?.aborted) continue;
+
+      pool.active++;
+      pool.totalAcquired++;
+      pool.totalWaitTimeMs += Date.now() - request.timestamp;
+      request.resolve(this.createReleaseFn(poolName));
+    }
+  }
+
+  getMetrics(poolName: string): PoolMetrics | undefined {
+    const pool = this.pools.get(poolName);
+    if (!pool) return undefined;
+
+    return {
+      name: poolName,
+      limit: pool.limit,
+      active: pool.active,
+      queued: pool.queue.length,
+      totalAcquired: pool.totalAcquired,
+      totalWaitTimeMs: pool.totalWaitTimeMs,
+    };
+  }
+
+  getAllMetrics(): PoolMetrics[] {
+    return Array.from(this.pools.entries()).map(([name, pool]) => ({
+      name,
+      limit: pool.limit,
+      active: pool.active,
+      queued: pool.queue.length,
+      totalAcquired: pool.totalAcquired,
+      totalWaitTimeMs: pool.totalWaitTimeMs,
+    }));
+  }
+}