311 lines
13 KiB
TypeScript
311 lines
13 KiB
TypeScript
import type { ExtensionAPI } from "@jaeswift/jae-coding-agent";
|
|
import { spawn, type ChildProcess } from "node:child_process";
|
|
import { Type } from "@sinclair/typebox";
|
|
|
|
interface AgentInstance {
|
|
id: string;
|
|
role: string;
|
|
systemPrompt: string;
|
|
status: "running" | "idle" | "stopped" | "error";
|
|
process: ChildProcess | null;
|
|
output: string;
|
|
createdAt: number;
|
|
lastMessage: string;
|
|
}
|
|
|
|
const ROLE_PROMPTS: Record<string, string> = {
|
|
frontend: "You are a frontend development specialist. Focus on UI/UX, React, CSS, accessibility, responsive design, and browser APIs. Write clean component code with proper state management.",
|
|
backend: "You are a backend development specialist. Focus on API design, databases, server architecture, authentication, and performance. Write robust server-side code with proper error handling.",
|
|
tester: "You are a QA and testing specialist. Focus on writing comprehensive tests: unit tests, integration tests, E2E tests. Identify edge cases, write test plans, and ensure code reliability.",
|
|
reviewer: "You are a senior code reviewer. Analyze code for bugs, security issues, performance problems, and style violations. Provide constructive feedback with specific suggestions.",
|
|
devops: "You are a DevOps and infrastructure specialist. Focus on CI/CD pipelines, Docker, Kubernetes, monitoring, and deployment automation. Write infrastructure as code.",
|
|
security: "You are a security specialist. Focus on vulnerability assessment, secure coding practices, authentication/authorization, encryption, and threat modeling.",
|
|
architect: "You are a software architect. Focus on system design, scalability patterns, technology selection, API contracts, and architectural decision records.",
|
|
docs: "You are a technical documentation specialist. Focus on writing clear README files, API docs, tutorials, inline comments, and architectural documentation.",
|
|
};
|
|
|
|
export default function (pi: ExtensionAPI) {
|
|
const agents = new Map<string, AgentInstance>();
|
|
let idCounter = 0;
|
|
|
|
function generateId(role: string): string {
|
|
idCounter++;
|
|
return `${role}-${idCounter}`;
|
|
}
|
|
|
|
function getSystemPrompt(role: string): string {
|
|
return ROLE_PROMPTS[role] || `You are a ${role} specialist. Focus on tasks related to ${role} and provide expert-level assistance.`;
|
|
}
|
|
|
|
async function spawnAgent(role: string, cwd: string): Promise<AgentInstance> {
|
|
const id = generateId(role);
|
|
const systemPrompt = getSystemPrompt(role);
|
|
|
|
const agent: AgentInstance = {
|
|
id,
|
|
role,
|
|
systemPrompt,
|
|
status: "idle",
|
|
process: null,
|
|
output: "",
|
|
createdAt: Date.now(),
|
|
lastMessage: "",
|
|
};
|
|
|
|
agents.set(id, agent);
|
|
return agent;
|
|
}
|
|
|
|
async function sendMessage(
|
|
agent: AgentInstance,
|
|
message: string,
|
|
cwd: string,
|
|
): Promise<string> {
|
|
return new Promise((resolve) => {
|
|
const timeout = 180_000; // 3 minutes
|
|
let output = "";
|
|
let resolved = false;
|
|
|
|
agent.status = "running";
|
|
agent.lastMessage = message;
|
|
|
|
const prompt = `[System: ${agent.systemPrompt}]\n\n${message}`;
|
|
|
|
try {
|
|
const child = spawn(
|
|
"jae",
|
|
["-p", prompt],
|
|
{
|
|
cwd,
|
|
stdio: ["ignore", "pipe", "pipe"],
|
|
timeout,
|
|
env: { ...process.env },
|
|
},
|
|
);
|
|
|
|
agent.process = child;
|
|
|
|
child.stdout?.on("data", (data: Buffer) => {
|
|
output += data.toString();
|
|
});
|
|
|
|
child.stderr?.on("data", (data: Buffer) => {
|
|
output += data.toString();
|
|
});
|
|
|
|
child.on("close", (code) => {
|
|
agent.process = null;
|
|
agent.status = code === 0 ? "idle" : "error";
|
|
agent.output = output.trim();
|
|
if (!resolved) {
|
|
resolved = true;
|
|
resolve(output.trim() || "(no output)");
|
|
}
|
|
});
|
|
|
|
child.on("error", (err) => {
|
|
agent.process = null;
|
|
agent.status = "error";
|
|
if (!resolved) {
|
|
resolved = true;
|
|
resolve(`Agent error: ${err.message}`);
|
|
}
|
|
});
|
|
|
|
setTimeout(() => {
|
|
if (!resolved) {
|
|
resolved = true;
|
|
try { child.kill("SIGTERM"); } catch {}
|
|
agent.status = "error";
|
|
resolve(output.trim() + "\n(timed out after 3 minutes)");
|
|
}
|
|
}, timeout);
|
|
} catch (err: any) {
|
|
agent.status = "error";
|
|
resolve(`Spawn error: ${err.message}`);
|
|
}
|
|
});
|
|
}
|
|
|
|
function killAgent(id: string): boolean {
|
|
const agent = agents.get(id);
|
|
if (!agent) return false;
|
|
if (agent.process) {
|
|
try { agent.process.kill("SIGTERM"); } catch {}
|
|
}
|
|
agent.status = "stopped";
|
|
agent.process = null;
|
|
agents.delete(id);
|
|
return true;
|
|
}
|
|
|
|
function formatAgentList(): string {
|
|
if (agents.size === 0) return "No active agents. Use /a2a spawn <role> to create one.";
|
|
const lines: string[] = [];
|
|
for (const [id, agent] of agents) {
|
|
const uptime = Math.floor((Date.now() - agent.createdAt) / 1000);
|
|
const uptimeStr = uptime > 60 ? `${Math.floor(uptime / 60)}m${uptime % 60}s` : `${uptime}s`;
|
|
const statusIcon = agent.status === "running" ? "\u{1F7E2}" : agent.status === "idle" ? "\u{1F535}" : agent.status === "error" ? "\u{1F534}" : "\u26AB";
|
|
lines.push(` ${statusIcon} ${id} [${agent.role}] — ${agent.status} — uptime: ${uptimeStr}`);
|
|
}
|
|
return lines.join("\n");
|
|
}
|
|
|
|
pi.registerTool({
|
|
name: "a2a_communicate",
|
|
label: "Agent-to-Agent Communication",
|
|
description: "Spawn specialized AI agents and communicate with them. Agents can be given roles like 'frontend', 'backend', 'tester', 'reviewer', 'devops', 'security', 'architect', or 'docs'.",
|
|
parameters: Type.Object({
|
|
action: Type.String({ description: "Action: 'spawn', 'send', 'broadcast', 'list', 'kill'" }),
|
|
role: Type.Optional(Type.String({ description: "Agent role for spawn (e.g., 'frontend', 'backend', 'tester')" })),
|
|
agent_id: Type.Optional(Type.String({ description: "Target agent ID for send/kill" })),
|
|
message: Type.Optional(Type.String({ description: "Message to send to agent(s)" })),
|
|
}),
|
|
execute: async (args, ctx) => {
|
|
const action = args.action.toLowerCase();
|
|
|
|
switch (action) {
|
|
case "spawn": {
|
|
const role = args.role || "general";
|
|
const agent = await spawnAgent(role, ctx.cwd);
|
|
return { output: `\u2705 Agent spawned: ${agent.id} [${role}]\nSystem prompt: ${agent.systemPrompt.substring(0, 100)}...` };
|
|
}
|
|
case "send": {
|
|
if (!args.agent_id) return { error: "Specify agent_id to send message to." };
|
|
if (!args.message) return { error: "Specify message to send." };
|
|
const agent = agents.get(args.agent_id);
|
|
if (!agent) return { error: `Agent '${args.agent_id}' not found. Use action 'list' to see agents.` };
|
|
if (agent.status === "running") return { error: `Agent '${args.agent_id}' is busy processing. Wait for completion.` };
|
|
const response = await sendMessage(agent, args.message, ctx.cwd);
|
|
return { output: `\u{1F4AC} Response from ${args.agent_id}:\n\n${response}` };
|
|
}
|
|
case "broadcast": {
|
|
if (!args.message) return { error: "Specify message to broadcast." };
|
|
if (agents.size === 0) return { error: "No active agents. Spawn some first." };
|
|
const idle = Array.from(agents.values()).filter((a) => a.status === "idle");
|
|
if (idle.length === 0) return { error: "All agents are busy. Wait for them to finish." };
|
|
const promises = idle.map(async (agent) => {
|
|
const response = await sendMessage(agent, args.message!, ctx.cwd);
|
|
return `### ${agent.id} [${agent.role}]:\n${response.substring(0, 500)}`;
|
|
});
|
|
const results = await Promise.all(promises);
|
|
return { output: `\u{1F4E2} Broadcast to ${idle.length} agents:\n\n${results.join("\n\n")}` };
|
|
}
|
|
case "list": {
|
|
return { output: `\u{1F916} Active Agents:\n${formatAgentList()}` };
|
|
}
|
|
case "kill": {
|
|
if (!args.agent_id) return { error: "Specify agent_id to kill." };
|
|
const success = killAgent(args.agent_id);
|
|
return { output: success ? `\u2705 Agent '${args.agent_id}' terminated.` : `\u274C Agent '${args.agent_id}' not found.` };
|
|
}
|
|
default:
|
|
return { error: `Unknown action '${action}'. Use: spawn, send, broadcast, list, kill` };
|
|
}
|
|
},
|
|
});
|
|
|
|
pi.registerCommand("a2a", {
|
|
description: "Agent-to-Agent: /a2a spawn <role> | list | send <id> <msg> | broadcast <msg> | kill <id>",
|
|
handler: async (args, ctx) => {
|
|
const parts = args.trim().split(/\s+/);
|
|
const subcommand = (parts[0] || "list").toLowerCase();
|
|
|
|
switch (subcommand) {
|
|
case "spawn": {
|
|
const role = parts[1] || "general";
|
|
const validRoles = [...Object.keys(ROLE_PROMPTS), "general"];
|
|
const agent = await spawnAgent(role, ctx.cwd);
|
|
ctx.ui.notify(
|
|
`\u{1F916} Agent Spawned!\n${'\u2500'.repeat(30)}\n\n` +
|
|
` ID: ${agent.id}\n` +
|
|
` Role: ${role}\n` +
|
|
` Prompt: ${agent.systemPrompt.substring(0, 80)}...\n\n` +
|
|
`Available roles: ${validRoles.join(", ")}\n` +
|
|
`Send messages: /a2a send ${agent.id} <message>`,
|
|
"info",
|
|
);
|
|
break;
|
|
}
|
|
case "list":
|
|
case "ls": {
|
|
ctx.ui.notify(`\u{1F916} Agent Registry\n${'\u2500'.repeat(30)}\n\n${formatAgentList()}`, "info");
|
|
break;
|
|
}
|
|
case "send": {
|
|
const agentId = parts[1];
|
|
const message = parts.slice(2).join(" ");
|
|
if (!agentId || !message) {
|
|
ctx.ui.notify("Usage: /a2a send <agent-id> <message>", "warning");
|
|
return;
|
|
}
|
|
const agent = agents.get(agentId);
|
|
if (!agent) {
|
|
ctx.ui.notify(`\u274C Agent '${agentId}' not found. Use /a2a list.`, "warning");
|
|
return;
|
|
}
|
|
if (agent.status === "running") {
|
|
ctx.ui.notify(`\u23F3 Agent '${agentId}' is busy. Wait for completion.`, "warning");
|
|
return;
|
|
}
|
|
ctx.ui.notify(`\u{1F4E4} Sending to ${agentId}...`, "info");
|
|
const response = await sendMessage(agent, message, ctx.cwd);
|
|
ctx.ui.notify(`\u{1F4AC} ${agentId} responded:\n\n${response.substring(0, 2000)}`, "info");
|
|
break;
|
|
}
|
|
case "broadcast": {
|
|
const message = parts.slice(1).join(" ");
|
|
if (!message) {
|
|
ctx.ui.notify("Usage: /a2a broadcast <message>", "warning");
|
|
return;
|
|
}
|
|
const idle = Array.from(agents.values()).filter((a) => a.status === "idle");
|
|
if (idle.length === 0) {
|
|
ctx.ui.notify("\u274C No idle agents. Spawn some first with /a2a spawn <role>.", "warning");
|
|
return;
|
|
}
|
|
ctx.ui.notify(`\u{1F4E2} Broadcasting to ${idle.length} agents...`, "info");
|
|
const promises = idle.map(async (agent) => {
|
|
const response = await sendMessage(agent, message, ctx.cwd);
|
|
return `\u{1F916} ${agent.id}:\n${response.substring(0, 500)}`;
|
|
});
|
|
const results = await Promise.all(promises);
|
|
ctx.ui.notify(`\u{1F4E2} Broadcast Results:\n${'\u2500'.repeat(30)}\n\n${results.join("\n\n")}`, "info");
|
|
break;
|
|
}
|
|
case "kill":
|
|
case "stop": {
|
|
const agentId = parts[1];
|
|
if (!agentId) {
|
|
ctx.ui.notify("Usage: /a2a kill <agent-id>", "warning");
|
|
return;
|
|
}
|
|
if (agentId === "all") {
|
|
const count = agents.size;
|
|
for (const id of Array.from(agents.keys())) killAgent(id);
|
|
ctx.ui.notify(`\u2705 Terminated ${count} agents.`, "info");
|
|
return;
|
|
}
|
|
const success = killAgent(agentId);
|
|
ctx.ui.notify(
|
|
success ? `\u2705 Agent '${agentId}' terminated.` : `\u274C Agent '${agentId}' not found.`,
|
|
success ? "info" : "warning",
|
|
);
|
|
break;
|
|
}
|
|
default:
|
|
ctx.ui.notify(
|
|
`\u{1F916} Agent-to-Agent Communication\n${'\u2500'.repeat(30)}\n\n` +
|
|
"Commands:\n" +
|
|
" /a2a spawn <role> - Spawn specialized agent\n" +
|
|
" /a2a list - List active agents\n" +
|
|
" /a2a send <id> <message> - Send message to agent\n" +
|
|
" /a2a broadcast <message> - Send to all idle agents\n" +
|
|
" /a2a kill <id|all> - Terminate agent(s)\n\n" +
|
|
`Roles: ${Object.keys(ROLE_PROMPTS).join(", ")}`,
|
|
"info",
|
|
);
|
|
}
|
|
},
|
|
});
|
|
}
|