Initial commit: kage-research project files
This commit is contained in:
shokollm
385
level3.ts
Normal file
385
level3.ts
Normal file
@@ -0,0 +1,385 @@
|
||||
/**
|
||||
* Level 3: Checkpoint/Recovery + Task Tracking
|
||||
*
|
||||
* Features:
|
||||
* 1. Task status (pending/running/completed/failed)
|
||||
* 2. Error tracking (why it failed)
|
||||
* 3. Retry mechanism with backoff
|
||||
* 4. Checkpoint/recovery
|
||||
*/
|
||||
|
||||
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,
|
||||
};
|
||||
|
||||
const WORKSPACE = "/tmp/shadows-level3";
|
||||
|
||||
// ============== TASK STATUS ==============
|
||||
type TaskStatus = "pending" | "running" | "completed" | "failed" | "retrying";
|
||||
|
||||
interface TaskError {
|
||||
message: string;
|
||||
tool?: string;
|
||||
timestamp: number;
|
||||
attempt: number;
|
||||
}
|
||||
|
||||
interface Task {
|
||||
id: string;
|
||||
message: string;
|
||||
status: TaskStatus;
|
||||
createdAt: number;
|
||||
startedAt?: number;
|
||||
completedAt?: number;
|
||||
error?: TaskError;
|
||||
attempts: number;
|
||||
maxRetries: number;
|
||||
retryDelay: number; // ms
|
||||
result?: string;
|
||||
}
|
||||
|
||||
interface Checkpoint {
|
||||
tasks: Task[];
|
||||
shadows: { id: string; taskId: string; state: any }[];
|
||||
savedAt: number;
|
||||
}
|
||||
|
||||
// ============== TASK MANAGER ==============
|
||||
class TaskManager {
|
||||
private tasks: Map<string, Task> = new Map();
|
||||
private maxRetries = 3;
|
||||
private retryDelay = 5000; // 5 seconds base
|
||||
private checkpointDir: string;
|
||||
|
||||
constructor(checkpointDir: string) {
|
||||
this.checkpointDir = checkpointDir;
|
||||
if (!fs.existsSync(checkpointDir)) {
|
||||
fs.mkdirSync(checkpointDir, { recursive: true });
|
||||
}
|
||||
}
|
||||
|
||||
// Create a new task
|
||||
createTask(id: string, message: string): Task {
|
||||
const task: Task = {
|
||||
id,
|
||||
message,
|
||||
status: "pending",
|
||||
createdAt: Date.now(),
|
||||
attempts: 0,
|
||||
maxRetries: this.maxRetries,
|
||||
retryDelay: this.retryDelay,
|
||||
};
|
||||
this.tasks.set(id, task);
|
||||
this.saveCheckpoint();
|
||||
return task;
|
||||
}
|
||||
|
||||
// Get next pending task
|
||||
getNextPending(): Task | undefined {
|
||||
for (const task of this.tasks.values()) {
|
||||
if (task.status === "pending" || task.status === "retrying") {
|
||||
return task;
|
||||
}
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
// Start a task
|
||||
startTask(id: string): Task | undefined {
|
||||
const task = this.tasks.get(id);
|
||||
if (!task) return undefined;
|
||||
|
||||
task.status = "running";
|
||||
task.startedAt = Date.now();
|
||||
task.attempts++;
|
||||
this.saveCheckpoint();
|
||||
return task;
|
||||
}
|
||||
|
||||
// Complete a task
|
||||
completeTask(id: string, result: string): Task | undefined {
|
||||
const task = this.tasks.get(id);
|
||||
if (!task) return undefined;
|
||||
|
||||
task.status = "completed";
|
||||
task.completedAt = Date.now();
|
||||
task.result = result;
|
||||
this.saveCheckpoint();
|
||||
return task;
|
||||
}
|
||||
|
||||
// Fail a task
|
||||
failTask(id: string, error: string, tool?: string): Task | undefined {
|
||||
const task = this.tasks.get(id);
|
||||
if (!task) return undefined;
|
||||
|
||||
task.error = {
|
||||
message: error,
|
||||
tool,
|
||||
timestamp: Date.now(),
|
||||
attempt: task.attempts,
|
||||
};
|
||||
|
||||
// Check if we can retry
|
||||
if (task.attempts < task.maxRetries) {
|
||||
task.status = "retrying";
|
||||
// Exponential backoff: 5s, 10s, 20s...
|
||||
task.retryDelay = task.retryDelay * 2;
|
||||
} else {
|
||||
task.status = "failed";
|
||||
}
|
||||
|
||||
this.saveCheckpoint();
|
||||
return task;
|
||||
}
|
||||
|
||||
// Get task by ID
|
||||
getTask(id: string): Task | undefined {
|
||||
return this.tasks.get(id);
|
||||
}
|
||||
|
||||
// List all tasks
|
||||
listTasks(): Task[] {
|
||||
return Array.from(this.tasks.values());
|
||||
}
|
||||
|
||||
// Save checkpoint to disk
|
||||
saveCheckpoint() {
|
||||
const checkpoint: Checkpoint = {
|
||||
tasks: this.listTasks(),
|
||||
shadows: [],
|
||||
savedAt: Date.now(),
|
||||
};
|
||||
fs.writeFileSync(
|
||||
path.join(this.checkpointDir, "checkpoint.json"),
|
||||
JSON.stringify(checkpoint, null, 2)
|
||||
);
|
||||
}
|
||||
|
||||
// Load checkpoint from disk
|
||||
loadCheckpoint(): boolean {
|
||||
const checkpointPath = path.join(this.checkpointDir, "checkpoint.json");
|
||||
if (!fs.existsSync(checkpointPath)) return false;
|
||||
|
||||
try {
|
||||
const data = fs.readFileSync(checkpointPath, "utf-8");
|
||||
const checkpoint: Checkpoint = JSON.parse(data);
|
||||
|
||||
// Restore tasks
|
||||
for (const task of checkpoint.tasks) {
|
||||
this.tasks.set(task.id, task);
|
||||
}
|
||||
return true;
|
||||
} catch (e) {
|
||||
console.error("Failed to load checkpoint:", e);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// Get stats
|
||||
getStats() {
|
||||
const tasks = this.listTasks();
|
||||
return {
|
||||
total: tasks.length,
|
||||
pending: tasks.filter(t => t.status === "pending").length,
|
||||
running: tasks.filter(t => t.status === "running").length,
|
||||
completed: tasks.filter(t => t.status === "completed").length,
|
||||
failed: tasks.filter(t => t.status === "failed").length,
|
||||
retrying: tasks.filter(t => t.status === "retrying").length,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// ============== SHADOW ==============
|
||||
class Shadow {
|
||||
public id: string;
|
||||
public status: "idle" | "running" = "idle";
|
||||
private agent: Agent;
|
||||
|
||||
constructor(id: string, worktreePath: string, systemPrompt: string, tools: AgentTool[]) {
|
||||
this.id = id;
|
||||
this.agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt,
|
||||
model,
|
||||
tools: 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 = "running";
|
||||
if (event.type === "agent_end") this.status = "idle";
|
||||
});
|
||||
}
|
||||
|
||||
async run(message: string): Promise<string> {
|
||||
const events: string[] = [];
|
||||
|
||||
this.agent.subscribe((event) => {
|
||||
events.push(event.type);
|
||||
|
||||
// Log tool errors
|
||||
if (event.type === "tool_execution_end" && (event as any).isError) {
|
||||
console.log(` ⚠️ Tool error in ${event.toolName}`);
|
||||
}
|
||||
});
|
||||
|
||||
await this.agent.prompt(message);
|
||||
|
||||
// Get last assistant message
|
||||
const lastMsg = this.agent.state.messages.filter(m => m.role === "assistant").pop();
|
||||
return lastMsg ? JSON.stringify(lastMsg.content) : "No response";
|
||||
}
|
||||
|
||||
abort() {
|
||||
this.agent.abort();
|
||||
}
|
||||
}
|
||||
|
||||
// ============== EXECUTOR ==============
|
||||
class Executor {
|
||||
private shadow: Shadow;
|
||||
private taskManager: TaskManager;
|
||||
private isRunning = false;
|
||||
|
||||
constructor(taskManager: TaskManager, worktreePath: string) {
|
||||
this.taskManager = taskManager;
|
||||
this.shadow = new Shadow(
|
||||
"executor-1",
|
||||
worktreePath,
|
||||
"You are a helpful coding assistant. Use the bash tool to run commands.",
|
||||
[]
|
||||
);
|
||||
}
|
||||
|
||||
async run(): Promise<void> {
|
||||
this.isRunning = true;
|
||||
|
||||
while (this.isRunning) {
|
||||
// Get next pending task
|
||||
const task = this.taskManager.getNextPending();
|
||||
|
||||
if (!task) {
|
||||
console.log("😴 No pending tasks, waiting...");
|
||||
await this.sleep(3000);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Start the task
|
||||
this.taskManager.startTask(task.id);
|
||||
console.log(`\n▶️ Running task ${task.id}: "${task.message.substring(0, 50)}..."`);
|
||||
console.log(` Attempt ${task.attempts}/${task.maxRetries}`);
|
||||
|
||||
try {
|
||||
// Run the task
|
||||
const result = await this.shadow.run(task.message);
|
||||
|
||||
// Success
|
||||
this.taskManager.completeTask(task.id, result);
|
||||
console.log(`✅ Task ${task.id} completed!`);
|
||||
|
||||
} catch (error: any) {
|
||||
// Failed
|
||||
this.taskManager.failTask(task.id, error.message);
|
||||
console.log(`❌ Task ${task.id} failed: ${error.message}`);
|
||||
|
||||
// Check if will retry
|
||||
const updatedTask = this.taskManager.getTask(task.id);
|
||||
if (updatedTask?.status === "retrying") {
|
||||
console.log(` 🔄 Will retry in ${updatedTask.retryDelay}ms...`);
|
||||
await this.sleep(updatedTask.retryDelay);
|
||||
}
|
||||
}
|
||||
|
||||
// Show stats
|
||||
console.log(`\n📊 Stats:`, this.taskManager.getStats());
|
||||
}
|
||||
}
|
||||
|
||||
stop() {
|
||||
this.isRunning = false;
|
||||
this.shadow.abort();
|
||||
}
|
||||
|
||||
private sleep(ms: number) {
|
||||
return new Promise(resolve => setTimeout(resolve, ms));
|
||||
}
|
||||
}
|
||||
|
||||
// ============== MAIN ==============
|
||||
async function main() {
|
||||
console.log("🧪 Level 3: Checkpoint/Recovery + Task Tracking\n");
|
||||
|
||||
registerBuiltInApiProviders();
|
||||
|
||||
// Create task manager with checkpoint directory
|
||||
const taskManager = new TaskManager(WORKSPACE);
|
||||
|
||||
// Check for existing checkpoint
|
||||
const loaded = taskManager.loadCheckpoint();
|
||||
if (loaded) {
|
||||
console.log("📂 Loaded checkpoint, existing tasks:", taskManager.getStats());
|
||||
}
|
||||
|
||||
// Create some test tasks
|
||||
console.log("📝 Creating test tasks...");
|
||||
taskManager.createTask("task-1", "Say hello and run 'echo Hello from Task 1'");
|
||||
taskManager.createTask("task-2", "Say hi and run 'echo Hello from Task 2'");
|
||||
taskManager.createTask("task-3", "Run 'date' to get current time");
|
||||
|
||||
console.log("📊 Initial stats:", taskManager.getStats());
|
||||
|
||||
// Create executor and run
|
||||
const executor = new Executor(taskManager, "/tmp");
|
||||
|
||||
// Run for a bit then stop (for demo)
|
||||
const runPromise = executor.run();
|
||||
|
||||
// Let it run for 60 seconds then stop
|
||||
await new Promise(resolve => setTimeout(resolve, 60000));
|
||||
executor.stop();
|
||||
|
||||
await runPromise;
|
||||
|
||||
console.log("\n✅ Demo complete!");
|
||||
console.log("📊 Final stats:", taskManager.getStats());
|
||||
|
||||
// Show failed tasks with error details
|
||||
const tasks = taskManager.listTasks();
|
||||
const failed = tasks.filter(t => t.status === "failed");
|
||||
if (failed.length > 0) {
|
||||
console.log("\n❌ Failed tasks:");
|
||||
failed.forEach(t => {
|
||||
console.log(` - ${t.id}: ${t.error?.message} (attempt ${t.error?.attempt})`);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
main().catch(console.error);
|
||||
Reference in New Issue
Block a user