Sub-Agents

import { z } from "zod";import { chat, toolDefinition } from "@tanstack/ai";import { openaiText } from "@tanstack/ai-openai";// Custom fetch that injects ALS-bound inbound x-* headers (e.g.// x-aimock-context) onto every outbound OpenAI call. Required so aimock// can match fixtures by integration context. See ../header-forwarding.ts// for the full rationale; mirrors the Mastra precedent.import { forwardingFetch } from "../header-forwarding";// Each role becomes its own nested chat() with a dedicated system prompt.// They don't share memory or tools with the supervisor — the supervisor// only sees the role's return value via the delegate tool below.//// Tool names match the LangGraph Python reference agent (`subagents.py`)://   research_agent, writing_agent, critique_agent// This alignment is load-bearing: the D5 fixtures are recorded against// the LGP agent's tool names, and aimock matches on tool name.const subagentRoles = [  {    id: "research_agent",    systemPrompt:      "You are a research sub-agent. Given a topic, produce a concise " +      "bulleted list of 3-5 key facts. No preamble, no closing.",  },  {    id: "writing_agent",    systemPrompt:      "You are a writing sub-agent. Given a brief and optional source " +      "facts, produce a polished 1-paragraph draft. Be clear and " +      "concrete. No preamble.",  },  {    id: "critique_agent",    systemPrompt:      "You are an editorial critique sub-agent. Given a draft, give " +      "2-3 crisp, actionable critiques. No preamble.",  },] as const;

import { z } from "zod";import { chat, toolDefinition } from "@tanstack/ai";import { openaiText } from "@tanstack/ai-openai";// Custom fetch that injects ALS-bound inbound x-* headers (e.g.// x-aimock-context) onto every outbound OpenAI call. Required so aimock// can match fixtures by integration context. See ../header-forwarding.ts// for the full rationale; mirrors the Mastra precedent.import { forwardingFetch } from "../header-forwarding";// Each role becomes its own nested chat() with a dedicated system prompt.// They don't share memory or tools with the supervisor — the supervisor// only sees the role's return value via the delegate tool below.//// Tool names match the LangGraph Python reference agent (`subagents.py`)://   research_agent, writing_agent, critique_agent// This alignment is load-bearing: the D5 fixtures are recorded against// the LGP agent's tool names, and aimock matches on tool name.const subagentRoles = [  {    id: "research_agent",    systemPrompt:      "You are a research sub-agent. Given a topic, produce a concise " +      "bulleted list of 3-5 key facts. No preamble, no closing.",  },  {    id: "writing_agent",    systemPrompt:      "You are a writing sub-agent. Given a brief and optional source " +      "facts, produce a polished 1-paragraph draft. Be clear and " +      "concrete. No preamble.",  },  {    id: "critique_agent",    systemPrompt:      "You are an editorial critique sub-agent. Given a draft, give " +      "2-3 crisp, actionable critiques. No preamble.",  },] as const;// Builder takes the parent run's AbortController so subagent `chat()` calls// abort with the parent. Constructing tools at module-import time leaves them// with their own fresh AbortController, which means a user cancel never reaches// the in-flight subagent call — orphan async work, billed tokens, hung// promises. Each parent run threads its controller through here.// Each `<role>_agent` tool wraps a nested chat() call with the// role's system prompt. The supervisor LLM "calls" these tools to// delegate work; each invocation runs the matching subagent and returns// its output for the supervisor's next step.// Cap the supervisor → critique sub-agent loop at a single iteration, matching// the reference's `_MAX_CRITIQUE_ITERATIONS`. Without it the supervisor LLM// occasionally re-calls `critique_agent` on the same draft, which shows up as// stacking 🧐 cards in the chat and duplicate rows in the delegation log. The// critic only adds value once per draft.const MAX_CRITIQUE_ITERATIONS = 1;export function buildSubagentTools(parentAbortController: AbortController) {  // Per-run counter. `buildSubagentTools` is already called once per run (so  // each run's nested chat() aborts with its parent), so a closure here is  // naturally run-scoped — module scope would leak the cap across requests.  let critiqueCalls = 0;  return subagentRoles.map((role) =>    toolDefinition({      name: role.id,      description: `Delegate a task to the ${role.id.replace(/_/g, " ")}.`,      inputSchema: z.object({        task: z          .string()          .describe(`Task description for the ${role.id.replace(/_/g, " ")}`),      }),    }).server(async ({ task }) => {      if (role.id === "critique_agent") {        critiqueCalls += 1;        if (critiqueCalls > MAX_CRITIQUE_ITERATIONS) {          // Return a no-op result rather than throwing: a throw would surface          // as a failed tool call and derail the supervisor's final summary.          return {            role: role.id,            text: "Critique already provided for this draft; skipping further review.",          };        }      }      const text = await chat({        adapter: openaiText("gpt-5.4", { fetch: forwardingFetch }),        messages: [{ role: "user", content: task }],        systemPrompts: [role.systemPrompt],        abortController: parentAbortController,        stream: false,      });      return { role: role.id, text };    }),  );}

  private readonly agentStore = injectAgentStore(this.agentId);  protected readonly delegations = computed(() =>    readDelegations(this.agentStore().state()),  );  constructor() {    if (this.feature === "subagents") {      this.registerSubAgent("research_agent");      this.registerSubAgent("writing_agent");      this.registerSubAgent("critique_agent");    }  }  private registerSubAgent(name: SubAgentName): void {    registerRenderToolCall(subAgentRendererConfig(name));  }