kage-research/level2-test.ts

/**
 * Level 2 Test: Concurrency
 * 
 * Tests:
 * 1. Run 2 shadows in parallel
 * 2. Hit concurrency limit (max=1, try to create 2nd)
 */

import { Agent, type AgentTool, type AgentMessage, type AgentEvent } from "@mariozechner/pi-agent-core";
import { registerBuiltInApiProviders } from "@mariozechner/pi-ai";
import type { Model } from "@mariozechner/pi-ai";
import * as fs from "fs";
import * as path from "path";
import { exec } from "child_process";

// ============== CONFIG ==============
const OPENROUTER_API_KEY = process.env.OPENROUTER_API_KEY || "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";
process.env.OPENROUTER_API_KEY = OPENROUTER_API_KEY;

const model: Model<"openai-responses"> = {
  id: "stepfun/step-3.5-flash:free",
  name: "Step-3.5 Flash (Free)",
  api: "openai-responses",
  provider: "openrouter",
  baseUrl: "https://openrouter.ai/api/v1",
  reasoning: false,
  input: ["text"],
  output: ["text"],
  cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
  contextWindow: 128000,
  maxTokens: 8192,
};

// ============== SIMPLE TOOLS ==============
function createTools(cwd: string = process.cwd()): AgentTool[] {
  return [
    {
      name: "bash",
      label: "Run Command",
      description: "Run a shell command",
      parameters: {
        type: "object",
        properties: {
          command: { type: "string", description: "Command to run" },
        },
        required: ["command"],
      } as const,
      execute: async (toolCallId: string, params: { command: string }) => {
        return new Promise((resolve) => {
          exec(params.command, { cwd }, (error, stdout, stderr) => {
            if (error) {
              resolve({
                content: [{ type: "text", text: stderr || error.message }],
                details: { command: params.command, exitCode: error.code },
                isError: true,
              });
            } else {
              resolve({
                content: [{ type: "text", text: stdout }],
                details: { command: params.command, exitCode: 0 },
              });
            }
          });
        });
      },
    },
  ];
}

// ============== SHADOW CLASS ==============
class Shadow {
  public readonly id: string;
  public readonly agent: Agent;
  public readonly worktreePath: string;
  public status: "idle" | "running" | "completed" | "error" = "idle";

  constructor(id: string, worktreePath: string, systemPrompt: string) {
    this.id = id;
    this.worktreePath = worktreePath;

    this.agent = new Agent({
      initialState: {
        systemPrompt,
        model: model,
        tools: createTools(worktreePath) as any,
        messages: [],
      },
      convertToLlm: (messages: AgentMessage[]) => {
        return messages
          .filter((m) => m.role === "user" || m.role === "assistant" || m.role === "toolResult")
          .map((m) => ({ role: m.role, content: m.content }));
      },
    });

    this.agent.subscribe((event) => {
      if (event.type === "agent_start") this.status = "running";
      if (event.type === "agent_end") this.status = "completed";
    });
  }

  async prompt(message: string) {
    return this.agent.prompt(message);
  }

  abort() {
    this.agent.abort();
  }
}

// ============== SHADOW MANAGER ==============
class ShadowManager {
  private shadows: Map<string, Shadow> = new Map();
  private maxConcurrent: number;

  constructor(maxConcurrent: number) {
    this.maxConcurrent = maxConcurrent;
  }

  get activeCount(): number {
    return Array.from(this.shadows.values()).filter(s => s.status === "running").length;
  }

  get totalCount(): number {
    return this.shadows.size;
  }

  createShadow(id: string, worktreePath: string, systemPrompt?: string): Shadow {
    // Check BOTH running and total count
    if (this.activeCount >= this.maxConcurrent || this.totalCount >= this.maxConcurrent) {
      throw new Error(`Max concurrent (${this.maxConcurrent}) reached! Current: ${this.activeCount} running, ${this.totalCount} total`);
    }

    const shadow = new Shadow(id, worktreePath, systemPrompt || "You are a helpful assistant.");
    this.shadows.set(id, shadow);
    return shadow;
  }

  getShadow(id: string): Shadow | undefined {
    return this.shadows.get(id);
  }

  terminateShadow(id: string) {
    const shadow = this.shadows.get(id);
    if (shadow) {
      shadow.abort();
      this.shadows.delete(id);
    }
  }

  getStats() {
    return {
      active: this.activeCount,
      maxConcurrent: this.maxConcurrent,
      totalShadows: this.shadows.size,
    };
  }
}

// ============== TEST 1: MULTIPLE SHADOWS ==============
async function testMultipleShadows() {
  console.log("\n" + "=".repeat(50));
  console.log("TEST 1: Multiple Shadows (2 in parallel)");
  console.log("=".repeat(50));

  const manager = new ShadowManager(2); // Allow 2 concurrent

  // Create 2 shadows
  const shadow1 = manager.createShadow("shadow-1", "/tmp");
  const shadow2 = manager.createShadow("shadow-2", "/tmp");

  console.log(`Created 2 shadows`);
  console.log(`Stats:`, manager.getStats());

  // Run both in parallel
  console.log("\n🚀 Running both shadows in parallel...\n");

  const [result1, result2] = await Promise.all([
    shadow1.prompt("Say 'Hello from Shadow 1'"),
    shadow2.prompt("Say 'Hello from Shadow 2'"),
  ]);

  console.log("\n✅ Both shadows completed!");
  console.log(`Stats:`, manager.getStats());

  // Cleanup
  manager.terminateShadow("shadow-1");
  manager.terminateShadow("shadow-2");
}

// ============== TEST 2: CONCURRENCY LIMIT ==============
async function testConcurrencyLimit() {
  console.log("\n" + "=".repeat(50));
  console.log("TEST 2: Concurrency Limit (max=1, create 2nd)");
  console.log("=".repeat(50));

  const manager = new ShadowManager(1); // Only allow 1 concurrent!

  // Create first shadow - should work
  const shadow1 = manager.createShadow("shadow-1", "/tmp");
  console.log(`Created shadow-1:`, manager.getStats());

  // Try to create second shadow - should fail!
  console.log("\n🔴 Trying to create shadow-2 (should fail)...");
  try {
    manager.createShadow("shadow-2", "/tmp");
    console.log("❌ ERROR: Should have thrown!");
  } catch (error: any) {
    console.log(`✅ Correctly rejected: ${error.message}`);
  }

  console.log(`\nStats:`, manager.getStats());

  // Cleanup
  manager.terminateShadow("shadow-1");
}

// ============== MAIN ==============
async function main() {
  console.log("🧪 Level 2 Concurrency Tests\n");

  registerBuiltInApiProviders();

  await testMultipleShadows();
  await testConcurrencyLimit();

  console.log("\n✅ All tests complete!");
}

main().catch(console.error);