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