Initial commit: kage-research project files
This commit is contained in:
shokollm
229
level2-test.ts
Normal file
229
level2-test.ts
Normal file
@@ -0,0 +1,229 @@
|
||||
/**
|
||||
* 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);
|
||||
Reference in New Issue
Block a user