Initial commit: kage-research project files
This commit is contained in:
386
level3c.ts
Normal file
386
level3c.ts
Normal file
@@ -0,0 +1,386 @@
|
||||
/**
|
||||
* Level 3c: Queue System with Worker Pool
|
||||
*
|
||||
* Features:
|
||||
* 1. Task queue - register many tasks
|
||||
* 2. Worker pool - max concurrent agents
|
||||
* 3. Auto-pull - workers take next task when free
|
||||
* 4. Priority support - high/normal/low
|
||||
* 5. Backpressure - reject when queue full
|
||||
*/
|
||||
|
||||
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 { 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,
|
||||
};
|
||||
|
||||
// ============== TASK ==============
|
||||
type TaskPriority = "high" | "normal" | "low";
|
||||
|
||||
interface QueuedTask {
|
||||
id: string;
|
||||
message: string;
|
||||
priority: TaskPriority;
|
||||
createdAt: number;
|
||||
status: "queued" | "running" | "completed" | "failed";
|
||||
}
|
||||
|
||||
// ============== QUEUE ==============
|
||||
class TaskQueue {
|
||||
private queue: QueuedTask[] = [];
|
||||
private maxSize: number;
|
||||
|
||||
constructor(maxSize: number = 100) {
|
||||
this.maxSize = maxSize;
|
||||
}
|
||||
|
||||
// Add task to queue
|
||||
enqueue(task: QueuedTask): boolean {
|
||||
if (this.queue.length >= this.maxSize) {
|
||||
return false; // Queue full - backpressure
|
||||
}
|
||||
|
||||
// Insert based on priority
|
||||
let insertIndex = this.queue.length;
|
||||
const priorityOrder = { high: 0, normal: 1, low: 2 };
|
||||
|
||||
for (let i = 0; i < this.queue.length; i++) {
|
||||
if (priorityOrder[task.priority] < priorityOrder[this.queue[i].priority]) {
|
||||
insertIndex = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
this.queue.splice(insertIndex, 0, task);
|
||||
return true;
|
||||
}
|
||||
|
||||
// Get next task (highest priority first)
|
||||
dequeue(): QueuedTask | undefined {
|
||||
const task = this.queue.shift();
|
||||
if (task) {
|
||||
task.status = "running";
|
||||
}
|
||||
return task;
|
||||
}
|
||||
|
||||
// Peek at next task without removing
|
||||
peek(): QueuedTask | undefined {
|
||||
return this.queue[0];
|
||||
}
|
||||
|
||||
// Get queue size
|
||||
size(): number {
|
||||
return this.queue.length;
|
||||
}
|
||||
|
||||
// Get all queued tasks
|
||||
getAll(): QueuedTask[] {
|
||||
return [...this.queue];
|
||||
}
|
||||
|
||||
// Update task status
|
||||
updateStatus(id: string, status: "completed" | "failed") {
|
||||
const task = this.queue.find(t => t.id === id);
|
||||
if (task) {
|
||||
task.status = status;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ============== WORKER (SHADOW) ==============
|
||||
class Worker {
|
||||
public id: string;
|
||||
public status: "idle" | "busy" = "idle";
|
||||
private agent: Agent;
|
||||
|
||||
constructor(id: string, worktreePath: string) {
|
||||
this.id = id;
|
||||
this.agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful coding assistant.",
|
||||
model: model,
|
||||
tools: [] 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 = "busy";
|
||||
if (event.type === "agent_end") this.status = "idle";
|
||||
});
|
||||
}
|
||||
|
||||
async run(task: QueuedTask): Promise<string> {
|
||||
this.status = "busy";
|
||||
await this.agent.prompt(task.message);
|
||||
this.status = "idle";
|
||||
return "completed";
|
||||
}
|
||||
|
||||
abort() {
|
||||
this.agent.abort();
|
||||
}
|
||||
}
|
||||
|
||||
// ============== QUEUE SYSTEM ==============
|
||||
class QueueSystem {
|
||||
private queue: TaskQueue;
|
||||
private workers: Worker[] = [];
|
||||
private maxWorkers: number;
|
||||
private maxQueueSize: number;
|
||||
private running = false;
|
||||
|
||||
constructor(maxWorkers: number = 2, maxQueueSize: number = 100) {
|
||||
this.maxWorkers = maxWorkers;
|
||||
this.maxQueueSize = maxQueueSize;
|
||||
this.queue = new TaskQueue(maxQueueSize);
|
||||
}
|
||||
|
||||
// Submit a task to queue
|
||||
submit(message: string, priority: TaskPriority = "normal"): boolean {
|
||||
const task: QueuedTask = {
|
||||
id: `task-${Date.now()}-${Math.random().toString(36).substr(2, 5)}`,
|
||||
message,
|
||||
priority,
|
||||
createdAt: Date.now(),
|
||||
status: "queued",
|
||||
};
|
||||
|
||||
const success = this.queue.enqueue(task);
|
||||
if (success) {
|
||||
console.log(`📥 Queued: ${task.id} (priority: ${priority}, queue: ${this.queue.size()})`);
|
||||
this.dispatch(); // Try to dispatch immediately
|
||||
} else {
|
||||
console.log(`❌ Queue full! Rejected: ${task.id}`);
|
||||
}
|
||||
return success;
|
||||
}
|
||||
|
||||
// Dispatch task to available worker
|
||||
private dispatch() {
|
||||
// Find idle workers
|
||||
const idleWorkers = this.workers.filter(w => w.status === "idle");
|
||||
|
||||
// Get next task
|
||||
const task = this.queue.peek();
|
||||
|
||||
if (!task || idleWorkers.length === 0) {
|
||||
return; // No task or no workers
|
||||
}
|
||||
|
||||
// Assign task to first idle worker
|
||||
const worker = idleWorkers[0];
|
||||
this.queue.dequeue(); // Remove from queue
|
||||
|
||||
console.log(`▶️ Dispatching ${task.id} to worker ${worker.id}`);
|
||||
|
||||
// Run task
|
||||
worker.run(task).then(() => {
|
||||
task.status = "completed";
|
||||
console.log(`✅ Completed: ${task.id}`);
|
||||
|
||||
// Check for more tasks
|
||||
this.dispatch();
|
||||
}).catch((error) => {
|
||||
task.status = "failed";
|
||||
console.log(`❌ Failed: ${task.id} - ${error.message}`);
|
||||
|
||||
// Check for more tasks
|
||||
this.dispatch();
|
||||
});
|
||||
}
|
||||
|
||||
// Add a worker
|
||||
addWorker(id: string, worktreePath: string): Worker {
|
||||
const worker = new Worker(id, worktreePath);
|
||||
this.workers.push(worker);
|
||||
console.log(`👷 Added worker: ${id} (total: ${this.workers.length})`);
|
||||
|
||||
// Try to dispatch
|
||||
this.dispatch();
|
||||
|
||||
return worker;
|
||||
}
|
||||
|
||||
// Remove a worker
|
||||
removeWorker(id: string) {
|
||||
const worker = this.workers.find(w => w.id === id);
|
||||
if (worker) {
|
||||
worker.abort();
|
||||
this.workers = this.workers.filter(w => w.id !== id);
|
||||
console.log(`👷 Removed worker: ${id}`);
|
||||
}
|
||||
}
|
||||
|
||||
// Get stats
|
||||
getStats() {
|
||||
return {
|
||||
queueSize: this.queue.size(),
|
||||
maxQueueSize: this.maxQueueSize,
|
||||
workers: this.workers.length,
|
||||
idleWorkers: this.workers.filter(w => w.status === "idle").length,
|
||||
busyWorkers: this.workers.filter(w => w.status === "busy").length,
|
||||
maxWorkers: this.maxWorkers,
|
||||
queuedTasks: this.queue.getAll().map(t => ({
|
||||
id: t.id,
|
||||
priority: t.priority,
|
||||
status: t.status,
|
||||
})),
|
||||
};
|
||||
}
|
||||
|
||||
// Start the queue processor
|
||||
start() {
|
||||
this.running = true;
|
||||
console.log(`🚀 Queue system started (max ${this.maxWorkers} workers)`);
|
||||
}
|
||||
|
||||
// Stop the queue processor
|
||||
stop() {
|
||||
this.running = false;
|
||||
this.workers.forEach(w => w.abort());
|
||||
console.log("🛑 Queue system stopped");
|
||||
}
|
||||
}
|
||||
|
||||
// ============== TESTS ==============
|
||||
async function testSequential() {
|
||||
console.log("\n" + "=" .repeat(50));
|
||||
console.log("TEST 1: Sequential (1 worker, multiple tasks)");
|
||||
console.log("=" .repeat(50));
|
||||
|
||||
const queue = new QueueSystem(1, 50);
|
||||
queue.start();
|
||||
queue.addWorker("worker-1", "/tmp");
|
||||
|
||||
// Submit 3 tasks
|
||||
queue.submit("Say 'Task 1'", "normal");
|
||||
queue.submit("Say 'Task 2'", "normal");
|
||||
queue.submit("Say 'Task 3'", "normal");
|
||||
|
||||
// Wait for tasks to complete
|
||||
await new Promise(resolve => setTimeout(resolve, 30000));
|
||||
|
||||
console.log("\n📊 Stats:", queue.getStats());
|
||||
queue.stop();
|
||||
}
|
||||
|
||||
async function testParallel() {
|
||||
console.log("\n" + "=" .repeat(50));
|
||||
console.log("TEST 2: Parallel (2 workers, multiple tasks)");
|
||||
console.log("=" .repeat(50));
|
||||
|
||||
const queue = new QueueSystem(2, 50);
|
||||
queue.start();
|
||||
queue.addWorker("worker-1", "/tmp");
|
||||
queue.addWorker("worker-2", "/tmp");
|
||||
|
||||
// Submit 4 tasks
|
||||
queue.submit("Say 'Task A'", "normal");
|
||||
queue.submit("Say 'Task B'", "normal");
|
||||
queue.submit("Say 'Task C'", "normal");
|
||||
queue.submit("Say 'Task D'", "normal");
|
||||
|
||||
// Wait for tasks to complete
|
||||
await new Promise(resolve => setTimeout(resolve, 30000));
|
||||
|
||||
console.log("\n📊 Stats:", queue.getStats());
|
||||
queue.stop();
|
||||
}
|
||||
|
||||
async function testPriority() {
|
||||
console.log("\n" + "=" .repeat(50));
|
||||
console.log("TEST 3: Priority (high priority first)");
|
||||
console.log("=" .repeat(50));
|
||||
|
||||
const queue = new QueueSystem(1, 50);
|
||||
queue.start();
|
||||
queue.addWorker("worker-1", "/tmp");
|
||||
|
||||
// Submit in random order with different priorities
|
||||
queue.submit("Say 'Normal 1'", "normal");
|
||||
queue.submit("Say 'Low'", "low");
|
||||
queue.submit("Say 'High 1'", "high");
|
||||
queue.submit("Say 'Normal 2'", "normal");
|
||||
queue.submit("Say 'High 2'", "high");
|
||||
|
||||
console.log("\n📊 Queue order:", queue.getStats().queuedTasks.map(t => `${t.priority}:${t.id.slice(-3)}`));
|
||||
|
||||
// Wait for tasks to complete
|
||||
await new Promise(resolve => setTimeout(resolve, 40000));
|
||||
|
||||
console.log("\n📊 Stats:", queue.getStats());
|
||||
queue.stop();
|
||||
}
|
||||
|
||||
async function testBackpressure() {
|
||||
console.log("\n" + "=" .repeat(50));
|
||||
console.log("TEST 4: Backpressure (queue full)");
|
||||
console.log("=" .repeat(50));
|
||||
|
||||
// Very small queue (3 max)
|
||||
const queue = new QueueSystem(1, 3);
|
||||
queue.start();
|
||||
queue.addWorker("worker-1", "/tmp");
|
||||
|
||||
// Submit 5 tasks (should reject 2)
|
||||
const results = [];
|
||||
results.push(queue.submit("Task 1", "normal"));
|
||||
results.push(queue.submit("Task 2", "normal"));
|
||||
results.push(queue.submit("Task 3", "normal"));
|
||||
results.push(queue.submit("Task 4", "normal")); // Should fail
|
||||
results.push(queue.submit("Task 5", "normal")); // Should fail
|
||||
|
||||
console.log("\n📊 Submit results:", results.map((r, i) => `Task ${i+1}: ${r ? '✅' : '❌'}`).join(", "));
|
||||
console.log("\n📊 Stats:", queue.getStats());
|
||||
|
||||
// Wait a bit then cleanup
|
||||
await new Promise(resolve => setTimeout(resolve, 5000));
|
||||
queue.stop();
|
||||
}
|
||||
|
||||
// ============== MAIN ==============
|
||||
async function main() {
|
||||
console.log("🧪 Level 3c: Queue System with Worker Pool\n");
|
||||
|
||||
registerBuiltInApiProviders();
|
||||
|
||||
await testSequential();
|
||||
await new Promise(resolve => setTimeout(resolve, 3000));
|
||||
|
||||
await testParallel();
|
||||
await new Promise(resolve => setTimeout(resolve, 3000));
|
||||
|
||||
await testPriority();
|
||||
await new Promise(resolve => setTimeout(resolve, 3000));
|
||||
|
||||
await testBackpressure();
|
||||
|
||||
console.log("\n✅ All tests complete!");
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
Reference in New Issue
Block a user