450 lines
14 KiB
TypeScript
450 lines
14 KiB
TypeScript
/**
|
|
* Level 2: Shadow + Shadow Manager + Tool Registry
|
|
*
|
|
* This adds:
|
|
* 1. Shadow class with context isolation
|
|
* 2. Shadow Manager for spawning/terminating
|
|
* 3. Tool registry (read, write, edit, bash, grep, find, ls)
|
|
* 4. Basic concurrency control
|
|
*/
|
|
|
|
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, spawn } 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;
|
|
|
|
// Model config (using free stepfun model)
|
|
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,
|
|
};
|
|
|
|
// ============== TOOL REGISTRY ==============
|
|
function createTools(cwd: string = process.cwd()): AgentTool[] {
|
|
return [
|
|
{
|
|
name: "read",
|
|
label: "Read File",
|
|
description: "Read the contents of a file",
|
|
parameters: {
|
|
type: "object",
|
|
properties: {
|
|
path: { type: "string", description: "Path to the file to read" },
|
|
},
|
|
required: ["path"],
|
|
} as const,
|
|
execute: async (toolCallId: string, params: { path: string }) => {
|
|
const fullPath = path.resolve(cwd, params.path);
|
|
try {
|
|
if (!fs.existsSync(fullPath)) {
|
|
throw new Error(`File not found: ${fullPath}`);
|
|
}
|
|
const content = fs.readFileSync(fullPath, "utf-8");
|
|
return {
|
|
content: [{ type: "text", text: content }],
|
|
details: { path: fullPath, lines: content.split("\n").length },
|
|
};
|
|
} catch (error: any) {
|
|
throw new Error(`Failed to read file: ${error.message}`);
|
|
}
|
|
},
|
|
},
|
|
{
|
|
name: "write",
|
|
label: "Write File",
|
|
description: "Write content to a file (creates or overwrites)",
|
|
parameters: {
|
|
type: "object",
|
|
properties: {
|
|
path: { type: "string", description: "Path to the file to write" },
|
|
content: { type: "string", description: "Content to write" },
|
|
},
|
|
required: ["path", "content"],
|
|
} as const,
|
|
execute: async (toolCallId: string, params: { path: string; content: string }) => {
|
|
const fullPath = path.resolve(cwd, params.path);
|
|
try {
|
|
// Ensure directory exists
|
|
const dir = path.dirname(fullPath);
|
|
if (!fs.existsSync(dir)) {
|
|
fs.mkdirSync(dir, { recursive: true });
|
|
}
|
|
fs.writeFileSync(fullPath, params.content, "utf-8");
|
|
return {
|
|
content: [{ type: "text", text: `Written ${params.content.length} bytes to ${fullPath}` }],
|
|
details: { path: fullPath, bytes: params.content.length },
|
|
};
|
|
} catch (error: any) {
|
|
throw new Error(`Failed to write file: ${error.message}`);
|
|
}
|
|
},
|
|
},
|
|
{
|
|
name: "edit",
|
|
label: "Edit File",
|
|
description: "Edit a file by replacing specific text",
|
|
parameters: {
|
|
type: "object",
|
|
properties: {
|
|
path: { type: "string", description: "Path to the file to edit" },
|
|
find: { type: "string", description: "Text to find" },
|
|
replace: { type: "string", description: "Text to replace with" },
|
|
},
|
|
required: ["path", "find"],
|
|
} as const,
|
|
execute: async (toolCallId: string, params: { path: string; find: string; replace?: string }) => {
|
|
const fullPath = path.resolve(cwd, params.path);
|
|
try {
|
|
if (!fs.existsSync(fullPath)) {
|
|
throw new Error(`File not found: ${fullPath}`);
|
|
}
|
|
let content = fs.readFileSync(fullPath, "utf-8");
|
|
const newContent = params.replace !== undefined
|
|
? content.replace(params.find, params.replace)
|
|
: content.replace(params.find, "");
|
|
|
|
if (content === newContent) {
|
|
throw new Error(`Text not found: "${params.find}"`);
|
|
}
|
|
|
|
fs.writeFileSync(fullPath, newContent, "utf-8");
|
|
return {
|
|
content: [{ type: "text", text: `Edited ${fullPath}` }],
|
|
details: { path: fullPath },
|
|
};
|
|
} catch (error: any) {
|
|
throw new Error(`Failed to edit file: ${error.message}`);
|
|
}
|
|
},
|
|
},
|
|
{
|
|
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 },
|
|
});
|
|
}
|
|
});
|
|
});
|
|
},
|
|
},
|
|
{
|
|
name: "grep",
|
|
label: "Search Text",
|
|
description: "Search for text in files",
|
|
parameters: {
|
|
type: "object",
|
|
properties: {
|
|
pattern: { type: "string", description: "Pattern to search for" },
|
|
path: { type: "string", description: "Path to search in (file or directory)" },
|
|
},
|
|
required: ["pattern"],
|
|
} as const,
|
|
execute: async (toolCallId: string, params: { pattern: string; path?: string }) => {
|
|
const searchPath = params.path || cwd;
|
|
return new Promise((resolve) => {
|
|
exec(`grep -r "${params.pattern}" ${searchPath} --line-number 2>/dev/null || true`, { cwd }, (error, stdout) => {
|
|
resolve({
|
|
content: [{ type: "text", text: stdout || `No matches found for "${params.pattern}"` }],
|
|
details: { pattern: params.pattern, path: searchPath },
|
|
});
|
|
});
|
|
});
|
|
},
|
|
},
|
|
{
|
|
name: "ls",
|
|
label: "List Files",
|
|
description: "List files in a directory",
|
|
parameters: {
|
|
type: "object",
|
|
properties: {
|
|
path: { type: "string", description: "Directory to list" },
|
|
},
|
|
} as const,
|
|
execute: async (toolCallId: string, params: { path?: string }) => {
|
|
const listPath = params.path ? path.resolve(cwd, params.path) : cwd;
|
|
try {
|
|
const files = fs.readdirSync(listPath);
|
|
return {
|
|
content: [{ type: "text", text: files.join("\n") }],
|
|
details: { path: listPath, count: files.length },
|
|
};
|
|
} catch (error: any) {
|
|
throw new Error(`Failed to list: ${error.message}`);
|
|
}
|
|
},
|
|
},
|
|
];
|
|
}
|
|
|
|
// ============== SHADOW CLASS ==============
|
|
interface ShadowConfig {
|
|
id: string;
|
|
systemPrompt: string;
|
|
worktreePath: string;
|
|
modelId?: string;
|
|
}
|
|
|
|
interface ShadowState {
|
|
id: string;
|
|
status: "idle" | "running" | "completed" | "error";
|
|
createdAt: Date;
|
|
worktreePath: string;
|
|
}
|
|
|
|
class Shadow {
|
|
public readonly id: string;
|
|
public readonly agent: Agent;
|
|
public readonly worktreePath: string;
|
|
public state: ShadowState;
|
|
|
|
private eventCallback?: (event: AgentEvent) => void;
|
|
|
|
constructor(config: ShadowConfig) {
|
|
this.id = config.id;
|
|
this.worktreePath = config.worktreePath;
|
|
this.state = {
|
|
id: config.id,
|
|
status: "idle",
|
|
createdAt: new Date(),
|
|
worktreePath: config.worktreePath,
|
|
};
|
|
|
|
// Create Pi Agent with isolated context
|
|
this.agent = new Agent({
|
|
initialState: {
|
|
systemPrompt: config.systemPrompt,
|
|
model: model,
|
|
tools: createTools(config.worktreePath) as any,
|
|
messages: [],
|
|
},
|
|
convertToLlm: (messages: AgentMessage[]) => {
|
|
// ISOLATION: Filter to only this shadow's messages
|
|
// We add a special role suffix to identify messages from this shadow
|
|
return messages
|
|
.filter((m) => {
|
|
// Keep messages that either:
|
|
// 1. Have no shadowId (legacy) OR
|
|
// 2. Have matching shadowId
|
|
const msg = m as any;
|
|
return !msg._shadowId || msg._shadowId === this.id;
|
|
})
|
|
.filter((m) => m.role === "user" || m.role === "assistant" || m.role === "toolResult")
|
|
.map((m) => ({
|
|
role: m.role,
|
|
content: m.content,
|
|
}));
|
|
},
|
|
});
|
|
|
|
// Subscribe to events
|
|
this.agent.subscribe((event) => {
|
|
// Track state changes
|
|
if (event.type === "agent_start") {
|
|
this.state.status = "running";
|
|
} else if (event.type === "agent_end") {
|
|
this.state.status = "completed";
|
|
} else if (event.type === "tool_execution_start") {
|
|
// Tool running
|
|
} else if (event.type === "tool_execution_end" && (event as any).isError) {
|
|
this.state.status = "error";
|
|
}
|
|
|
|
// Forward events
|
|
this.eventCallback?.(event);
|
|
});
|
|
}
|
|
|
|
onEvent(callback: (event: AgentEvent) => void) {
|
|
this.eventCallback = callback;
|
|
}
|
|
|
|
async prompt(message: string): Promise<AgentEvent[]> {
|
|
const events: AgentEvent[] = [];
|
|
|
|
// Tag message with shadow ID for isolation
|
|
const shadowMessage: AgentMessage = {
|
|
role: "user",
|
|
content: [{ type: "text", text: message }],
|
|
timestamp: Date.now(),
|
|
_shadowId: this.id, // Our custom field for isolation
|
|
};
|
|
|
|
return this.agent.prompt(shadowMessage);
|
|
}
|
|
|
|
abort() {
|
|
this.agent.abort();
|
|
}
|
|
|
|
reset() {
|
|
this.agent.reset();
|
|
this.state.status = "idle";
|
|
}
|
|
}
|
|
|
|
// ============== SHADOW MANAGER ==============
|
|
interface ShadowManagerConfig {
|
|
maxConcurrent?: number;
|
|
defaultSystemPrompt?: string;
|
|
}
|
|
|
|
class ShadowManager {
|
|
private shadows: Map<string, Shadow> = new Map();
|
|
private maxConcurrent: number;
|
|
private defaultSystemPrompt: string;
|
|
private activeCount = 0;
|
|
|
|
constructor(config: ShadowManagerConfig = {}) {
|
|
this.maxConcurrent = config.maxConcurrent || 5;
|
|
this.defaultSystemPrompt = config.defaultSystemPrompt || `You are a helpful coding assistant. You have access to tools: read, write, edit, bash, grep, ls. Use them to help the user. Be concise and practical.`;
|
|
}
|
|
|
|
async createShadow(worktreePath: string, customPrompt?: string): Promise<Shadow> {
|
|
// Check concurrency limit
|
|
if (this.activeCount >= this.maxConcurrent) {
|
|
throw new Error(`Max concurrent shadows reached (${this.maxConcurrent})`);
|
|
}
|
|
|
|
const id = `shadow-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
|
|
|
|
const shadow = new Shadow({
|
|
id,
|
|
systemPrompt: customPrompt || this.defaultSystemPrompt,
|
|
worktreePath,
|
|
});
|
|
|
|
this.shadows.set(id, shadow);
|
|
this.activeCount++;
|
|
|
|
console.log(`📦 Created shadow ${id} (active: ${this.activeCount}/${this.maxConcurrent})`);
|
|
|
|
return shadow;
|
|
}
|
|
|
|
getShadow(id: string): Shadow | undefined {
|
|
return this.shadows.get(id);
|
|
}
|
|
|
|
listShadows(): ShadowState[] {
|
|
return Array.from(this.shadows.values()).map((s) => s.state);
|
|
}
|
|
|
|
async terminateShadow(id: string): Promise<void> {
|
|
const shadow = this.shadows.get(id);
|
|
if (!shadow) {
|
|
throw new Error(`Shadow ${id} not found`);
|
|
}
|
|
|
|
shadow.abort();
|
|
this.shadows.delete(id);
|
|
this.activeCount--;
|
|
|
|
console.log(`🗑️ Terminated shadow ${id} (active: ${this.activeCount}/${this.maxConcurrent})`);
|
|
}
|
|
|
|
getStats() {
|
|
return {
|
|
active: this.activeCount,
|
|
maxConcurrent: this.maxConcurrent,
|
|
totalShadows: this.shadows.size,
|
|
shadows: this.listShadows(),
|
|
};
|
|
}
|
|
}
|
|
|
|
// ============== MAIN ==============
|
|
async function main() {
|
|
console.log("🚀 Level 2: Shadow + Shadow Manager\n");
|
|
|
|
// Initialize
|
|
registerBuiltInApiProviders();
|
|
|
|
// Create manager
|
|
const manager = new ShadowManager({
|
|
maxConcurrent: 3,
|
|
});
|
|
|
|
// Create a shadow
|
|
console.log("📦 Creating shadow...");
|
|
const shadow = await manager.createShadow("/home/shoko/repositories/shadows");
|
|
|
|
// Subscribe to events
|
|
shadow.onEvent((event) => {
|
|
switch (event.type) {
|
|
case "agent_start":
|
|
console.log("🤖 Agent started");
|
|
break;
|
|
case "turn_start":
|
|
console.log("🔄 Turn started");
|
|
break;
|
|
case "message_update":
|
|
const ev = event as any;
|
|
if (ev.assistantMessageEvent?.type === "text_delta") {
|
|
process.stdout.write(ev.assistantMessageEvent.delta || "");
|
|
}
|
|
break;
|
|
case "tool_execution_start":
|
|
console.log(`\n🔧 Tool: ${event.toolName}`);
|
|
break;
|
|
case "tool_execution_end":
|
|
console.log(` → Done (error: ${(event as any).isError})`);
|
|
break;
|
|
case "turn_end":
|
|
console.log("\n✅ Turn ended");
|
|
break;
|
|
case "agent_end":
|
|
console.log("\n🏁 Agent finished");
|
|
break;
|
|
}
|
|
});
|
|
|
|
// Run a task
|
|
console.log("\n📝 Running task: List files and check current directory\n");
|
|
await shadow.prompt("List the files in the current directory, then run 'pwd' to check the current directory.");
|
|
|
|
// Show stats
|
|
console.log("\n📊 Manager Stats:", manager.getStats());
|
|
|
|
// Cleanup
|
|
await manager.terminateShadow(shadow.id);
|
|
console.log("\n✅ Done!");
|
|
}
|
|
|
|
main().catch(console.error);
|