initial commit
This commit is contained in:
336
src/tool.ts
Normal file
336
src/tool.ts
Normal file
@@ -0,0 +1,336 @@
|
||||
import { Text } from "@mariozechner/pi-tui";
|
||||
import { discoverAgents } from "./agents.ts";
|
||||
import {
|
||||
listAvailableModelReferences,
|
||||
normalizeAvailableModelReference,
|
||||
resolveChildModel,
|
||||
} from "./models.ts";
|
||||
import {
|
||||
SubagentParamsSchema,
|
||||
type AgentScope,
|
||||
type SubagentRunResult,
|
||||
type SubagentToolDetails,
|
||||
} from "./schema.ts";
|
||||
|
||||
const MAX_PARALLEL_TASKS = 8;
|
||||
const MAX_CONCURRENCY = 4;
|
||||
|
||||
async function mapWithConcurrencyLimit<TIn, TOut>(
|
||||
items: TIn[],
|
||||
concurrency: number,
|
||||
fn: (item: TIn, index: number) => Promise<TOut>,
|
||||
): Promise<TOut[]> {
|
||||
const limit = Math.max(1, Math.min(concurrency, items.length || 1));
|
||||
const results = new Array<TOut>(items.length);
|
||||
let nextIndex = 0;
|
||||
|
||||
await Promise.all(
|
||||
Array.from({ length: limit }, async () => {
|
||||
while (nextIndex < items.length) {
|
||||
const index = nextIndex++;
|
||||
results[index] = await fn(items[index], index);
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
return results;
|
||||
}
|
||||
|
||||
function isFailure(result: Pick<SubagentRunResult, "exitCode" | "stopReason">) {
|
||||
return result.exitCode !== 0 || result.stopReason === "error" || result.stopReason === "aborted";
|
||||
}
|
||||
|
||||
function makeDetails(
|
||||
mode: "single" | "parallel" | "chain",
|
||||
agentScope: AgentScope,
|
||||
projectAgentsDir: string | null,
|
||||
results: SubagentRunResult[],
|
||||
): SubagentToolDetails {
|
||||
return { mode, agentScope, projectAgentsDir, results };
|
||||
}
|
||||
|
||||
function makeErrorResult(
|
||||
text: string,
|
||||
mode: "single" | "parallel" | "chain",
|
||||
agentScope: AgentScope,
|
||||
projectAgentsDir: string | null,
|
||||
) {
|
||||
return {
|
||||
content: [{ type: "text" as const, text }],
|
||||
details: makeDetails(mode, agentScope, projectAgentsDir, []),
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
|
||||
export function createSubagentTool(deps: {
|
||||
discoverAgents?: typeof discoverAgents;
|
||||
listAvailableModelReferences?: typeof listAvailableModelReferences;
|
||||
normalizeAvailableModelReference?: typeof normalizeAvailableModelReference;
|
||||
parameters?: typeof SubagentParamsSchema;
|
||||
resolveChildModel?: typeof resolveChildModel;
|
||||
runSingleTask?: (input: {
|
||||
cwd: string;
|
||||
meta: Record<string, unknown>;
|
||||
onEvent?: (event: any) => void;
|
||||
}) => Promise<SubagentRunResult>;
|
||||
} = {}) {
|
||||
return {
|
||||
name: "subagent",
|
||||
label: "Subagent",
|
||||
description: "Delegate tasks to specialized agents running in tmux panes.",
|
||||
parameters: deps.parameters ?? SubagentParamsSchema,
|
||||
async execute(_toolCallId: string, params: any, _signal: AbortSignal | undefined, onUpdate: any, ctx: any) {
|
||||
const hasSingle = Boolean(params.agent && params.task);
|
||||
const hasParallel = Boolean(params.tasks?.length);
|
||||
const hasChain = Boolean(params.chain?.length);
|
||||
const modeCount = Number(hasSingle) + Number(hasParallel) + Number(hasChain);
|
||||
const mode = hasParallel ? "parallel" : hasChain ? "chain" : "single";
|
||||
const agentScope = (params.agentScope ?? "user") as AgentScope;
|
||||
|
||||
if (modeCount !== 1) {
|
||||
return makeErrorResult("Provide exactly one mode: single, parallel, or chain.", "single", agentScope, null);
|
||||
}
|
||||
|
||||
const discovery = (deps.discoverAgents ?? discoverAgents)(ctx.cwd, { scope: agentScope });
|
||||
const availableModelReferences = (deps.listAvailableModelReferences ?? listAvailableModelReferences)(ctx.modelRegistry);
|
||||
const availableModelsText = availableModelReferences.join(", ") || "(none)";
|
||||
const normalizeModelReference = (requestedModel?: string) =>
|
||||
(deps.normalizeAvailableModelReference ?? normalizeAvailableModelReference)(requestedModel, availableModelReferences);
|
||||
|
||||
if (availableModelReferences.length === 0) {
|
||||
return makeErrorResult(
|
||||
"No available models are configured. Configure at least one model before using subagent.",
|
||||
mode,
|
||||
agentScope,
|
||||
discovery.projectAgentsDir,
|
||||
);
|
||||
}
|
||||
|
||||
const topLevelModel = normalizeModelReference(params.model);
|
||||
if (!topLevelModel) {
|
||||
const message =
|
||||
typeof params.model !== "string" || params.model.trim().length === 0
|
||||
? `Subagent requires a top-level model chosen from the available models: ${availableModelsText}`
|
||||
: `Invalid top-level model "${params.model}". Choose one of the available models: ${availableModelsText}`;
|
||||
return makeErrorResult(message, mode, agentScope, discovery.projectAgentsDir);
|
||||
}
|
||||
params.model = topLevelModel;
|
||||
|
||||
for (const [index, task] of (params.tasks ?? []).entries()) {
|
||||
if (task.model === undefined) continue;
|
||||
|
||||
const normalizedTaskModel = normalizeModelReference(task.model);
|
||||
if (!normalizedTaskModel) {
|
||||
return makeErrorResult(
|
||||
`Invalid model for parallel task ${index + 1} (${task.agent}): "${task.model}". Choose one of the available models: ${availableModelsText}`,
|
||||
mode,
|
||||
agentScope,
|
||||
discovery.projectAgentsDir,
|
||||
);
|
||||
}
|
||||
task.model = normalizedTaskModel;
|
||||
}
|
||||
|
||||
for (const [index, step] of (params.chain ?? []).entries()) {
|
||||
if (step.model === undefined) continue;
|
||||
|
||||
const normalizedStepModel = normalizeModelReference(step.model);
|
||||
if (!normalizedStepModel) {
|
||||
return makeErrorResult(
|
||||
`Invalid model for chain step ${index + 1} (${step.agent}): "${step.model}". Choose one of the available models: ${availableModelsText}`,
|
||||
mode,
|
||||
agentScope,
|
||||
discovery.projectAgentsDir,
|
||||
);
|
||||
}
|
||||
step.model = normalizedStepModel;
|
||||
}
|
||||
|
||||
const requestedAgentNames = [
|
||||
...(hasSingle ? [params.agent] : []),
|
||||
...((params.tasks ?? []).map((task: any) => task.agent)),
|
||||
...((params.chain ?? []).map((step: any) => step.agent)),
|
||||
];
|
||||
const projectAgents = requestedAgentNames
|
||||
.map((name) => discovery.agents.find((candidate) => candidate.name === name))
|
||||
.filter((agent): agent is NonNullable<typeof agent> => Boolean(agent && agent.source === "project"));
|
||||
|
||||
if (projectAgents.length > 0 && (params.confirmProjectAgents ?? true) && ctx.hasUI) {
|
||||
const ok = await ctx.ui.confirm(
|
||||
"Run project-local agents?",
|
||||
`Agents: ${projectAgents.map((agent) => agent.name).join(", ")}\nSource: ${
|
||||
discovery.projectAgentsDir ?? "(unknown)"
|
||||
}`,
|
||||
);
|
||||
if (!ok) {
|
||||
return makeErrorResult(
|
||||
"Canceled: project-local agents not approved.",
|
||||
mode,
|
||||
agentScope,
|
||||
discovery.projectAgentsDir,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
const resolveAgent = (name: string) => {
|
||||
const agent = discovery.agents.find((candidate) => candidate.name === name);
|
||||
if (!agent) throw new Error(`Unknown agent: ${name}`);
|
||||
return agent;
|
||||
};
|
||||
|
||||
const runTask = async (input: {
|
||||
agentName: string;
|
||||
task: string;
|
||||
cwd?: string;
|
||||
taskModel?: string;
|
||||
taskIndex?: number;
|
||||
step?: number;
|
||||
mode: "single" | "parallel" | "chain";
|
||||
}) => {
|
||||
const agent = resolveAgent(input.agentName);
|
||||
const model = (deps.resolveChildModel ?? resolveChildModel)({
|
||||
taskModel: input.taskModel,
|
||||
topLevelModel: params.model,
|
||||
});
|
||||
|
||||
return deps.runSingleTask?.({
|
||||
cwd: input.cwd ?? ctx.cwd,
|
||||
onEvent(event) {
|
||||
onUpdate?.({
|
||||
content: [{ type: "text", text: `Running ${input.agentName}: ${event.type}` }],
|
||||
details: makeDetails(input.mode, agentScope, discovery.projectAgentsDir, []),
|
||||
});
|
||||
},
|
||||
meta: {
|
||||
mode: input.mode,
|
||||
taskIndex: input.taskIndex,
|
||||
step: input.step,
|
||||
agent: agent.name,
|
||||
agentSource: agent.source,
|
||||
task: input.task,
|
||||
cwd: input.cwd ?? ctx.cwd,
|
||||
requestedModel: model.requestedModel,
|
||||
resolvedModel: model.resolvedModel,
|
||||
systemPrompt: agent.systemPrompt,
|
||||
tools: agent.tools,
|
||||
},
|
||||
}) as Promise<SubagentRunResult>;
|
||||
};
|
||||
|
||||
if (hasSingle) {
|
||||
try {
|
||||
const result = await runTask({
|
||||
agentName: params.agent,
|
||||
task: params.task,
|
||||
cwd: params.cwd,
|
||||
mode: "single",
|
||||
});
|
||||
|
||||
return {
|
||||
content: [{ type: "text" as const, text: result.finalText }],
|
||||
details: makeDetails("single", agentScope, discovery.projectAgentsDir, [result]),
|
||||
isError: isFailure(result),
|
||||
};
|
||||
} catch (error) {
|
||||
return {
|
||||
content: [{ type: "text" as const, text: (error as Error).message }],
|
||||
details: makeDetails("single", agentScope, discovery.projectAgentsDir, []),
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
if (hasParallel) {
|
||||
if (params.tasks.length > MAX_PARALLEL_TASKS) {
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: "text" as const,
|
||||
text: `Too many parallel tasks (${params.tasks.length}). Max is ${MAX_PARALLEL_TASKS}.`,
|
||||
},
|
||||
],
|
||||
details: makeDetails("parallel", agentScope, discovery.projectAgentsDir, []),
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
|
||||
const liveResults: SubagentRunResult[] = [];
|
||||
const results = await mapWithConcurrencyLimit(params.tasks, MAX_CONCURRENCY, async (task: any, index) => {
|
||||
const result = await runTask({
|
||||
agentName: task.agent,
|
||||
task: task.task,
|
||||
cwd: task.cwd,
|
||||
taskModel: task.model,
|
||||
taskIndex: index,
|
||||
mode: "parallel",
|
||||
});
|
||||
liveResults[index] = result;
|
||||
onUpdate?.({
|
||||
content: [{ type: "text", text: `Parallel: ${liveResults.filter(Boolean).length}/${params.tasks.length} finished` }],
|
||||
details: makeDetails("parallel", agentScope, discovery.projectAgentsDir, liveResults.filter(Boolean)),
|
||||
});
|
||||
return result;
|
||||
});
|
||||
|
||||
const successCount = results.filter((result) => !isFailure(result)).length;
|
||||
const summary = results
|
||||
.map((result) => `[${result.agent}] ${isFailure(result) ? "failed" : "completed"}: ${result.finalText || "(no output)"}`)
|
||||
.join("\n\n");
|
||||
|
||||
return {
|
||||
content: [{ type: "text" as const, text: `Parallel: ${successCount}/${results.length} succeeded\n\n${summary}` }],
|
||||
details: makeDetails("parallel", agentScope, discovery.projectAgentsDir, results),
|
||||
isError: successCount !== results.length,
|
||||
};
|
||||
}
|
||||
|
||||
const results: SubagentRunResult[] = [];
|
||||
let previous = "";
|
||||
for (let index = 0; index < params.chain.length; index += 1) {
|
||||
const item = params.chain[index];
|
||||
const task = item.task.replaceAll("{previous}", previous);
|
||||
const result = await runTask({
|
||||
agentName: item.agent,
|
||||
task,
|
||||
cwd: item.cwd,
|
||||
taskModel: item.model,
|
||||
step: index + 1,
|
||||
mode: "chain",
|
||||
});
|
||||
onUpdate?.({
|
||||
content: [{ type: "text", text: `Chain: completed step ${index + 1}/${params.chain.length}` }],
|
||||
details: makeDetails("chain", agentScope, discovery.projectAgentsDir, [...results, result]),
|
||||
});
|
||||
results.push(result);
|
||||
if (isFailure(result)) {
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: "text" as const,
|
||||
text: `Chain stopped at step ${index + 1} (${item.agent}): ${result.finalText || result.stopReason || "failed"}`,
|
||||
},
|
||||
],
|
||||
details: makeDetails("chain", agentScope, discovery.projectAgentsDir, results),
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
previous = result.finalText;
|
||||
}
|
||||
|
||||
const finalResult = results[results.length - 1];
|
||||
return {
|
||||
content: [{ type: "text" as const, text: finalResult?.finalText ?? "" }],
|
||||
details: makeDetails("chain", agentScope, discovery.projectAgentsDir, results),
|
||||
};
|
||||
},
|
||||
renderCall(args: any) {
|
||||
if (args.tasks?.length) return new Text(`subagent parallel (${args.tasks.length} tasks)`, 0, 0);
|
||||
if (args.chain?.length) return new Text(`subagent chain (${args.chain.length} steps)`, 0, 0);
|
||||
return new Text(`subagent ${args.agent ?? ""}`.trim(), 0, 0);
|
||||
},
|
||||
renderResult(result: { content: Array<{ type: string; text?: string }> }) {
|
||||
const first = result.content[0];
|
||||
return new Text(first?.type === "text" ? first.text ?? "" : "", 0, 0);
|
||||
},
|
||||
};
|
||||
}
|
||||
Reference in New Issue
Block a user