386 lines
10 KiB
TypeScript
386 lines
10 KiB
TypeScript
/**
|
||
* 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);
|