Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
66 kB · 2230 lines
TypeScript
at main
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231import { type AssistantMessage, type AssistantMessageEvent, EventStream, type Message, type Model, type UserMessage,} from "@earendil-works/pi-ai";import { Type } from "typebox";import { describe, expect, it, vi } from "vitest";import { agentLoop, agentLoopContinue, runAgentLoop } from "../src/agent-loop.js";import type { AgentContext, AgentEvent, AgentLoopConfig, AgentMessage, AgentTool } from "../src/types.js";
class MockAssistantStream extends EventStream<AssistantMessageEvent, AssistantMessage> { constructor() { super( (event) => event.type === "done" || event.type === "error", (event) => { if (event.type === "done") return event.message; if (event.type === "error") return event.error; throw new Error("Unexpected event type"); }, ); }}
class DelayedResultStream extends MockAssistantStream { constructor(private readonly getDelayedResult: () => Promise<AssistantMessage>) { super(); }
override result(): Promise<AssistantMessage> { return this.getDelayedResult(); }}
class ThrowingResultStream extends MockAssistantStream { constructor( private readonly onResult: () => void, private readonly error: Error, ) { super(); }
override [Symbol.asyncIterator](): AsyncIterator<AssistantMessageEvent> { return { next: async () => ({ done: true, value: undefined as never }), }; }
override result(): Promise<AssistantMessage> { this.onResult(); throw this.error; }}
function createUsage() { return { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, };}
function createModel(): Model<"openai-responses"> { return { id: "mock", name: "mock", api: "openai-responses", provider: "openai", baseUrl: "https://example.invalid", reasoning: false, input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 8192, maxTokens: 2048, };}
function createAssistantMessage( content: AssistantMessage["content"], stopReason: AssistantMessage["stopReason"] = "stop",): AssistantMessage { return { role: "assistant", content, api: "openai-responses", provider: "openai", model: "mock", usage: createUsage(), stopReason, timestamp: Date.now(), };}
function createUserMessage(text: string): UserMessage { return { role: "user", content: text, timestamp: Date.now(), };}
function identityConverter(messages: AgentMessage[]): Message[] { return messages.filter((m) => m.role === "user" || m.role === "assistant" || m.role === "toolResult") as Message[];}
describe("agentLoop with AgentMessage", () => { it("should preserve a terminal response when abort fires after done", async () => { const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [], }; const controller = new AbortController(); const finalMessage = createAssistantMessage([{ type: "text", text: "complete" }]); const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, }; const streamFn = () => { const stream = new DelayedResultStream(() => { controller.abort(); return Promise.resolve(finalMessage); }); queueMicrotask(() => { stream.push({ type: "done", reason: "stop", message: finalMessage }); }); return stream; };
const stream = agentLoop([createUserMessage("Hello")], context, config, controller.signal, streamFn); for await (const _event of stream) { }
const messages = await stream.result(); const assistant = messages.find((message) => message.role === "assistant"); expect(assistant?.role).toBe("assistant"); if (assistant?.role === "assistant") { expect(assistant.stopReason).toBe("stop"); expect(assistant.content).toEqual([{ type: "text", text: "complete" }]); } });
it("should not wait for a pending terminal result after abort", async () => { const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [], }; const controller = new AbortController(); const finalMessage = createAssistantMessage([{ type: "text", text: "complete" }]); let resolveResult: ((message: AssistantMessage) => void) | undefined; const pendingResult = new Promise<AssistantMessage>((resolve) => { resolveResult = resolve; }); const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, }; const streamFn = () => { const stream = new DelayedResultStream(() => { controller.abort(); return pendingResult; }); queueMicrotask(() => { stream.push({ type: "done", reason: "stop", message: finalMessage }); }); return stream; };
const stream = agentLoop([createUserMessage("Hello")], context, config, controller.signal, streamFn); const events: AgentEvent[] = []; const consume = (async () => { for await (const event of stream) { events.push(event); } return stream.result(); })(); const result = await Promise.race([ consume, new Promise<"timeout">((resolve) => { setTimeout(() => resolve("timeout"), 50); }), ]); if (result === "timeout") { resolveResult?.(finalMessage); await consume; throw new Error("agent loop waited for a pending terminal result after abort"); }
const assistant = result.find((message) => message.role === "assistant"); expect(assistant?.role).toBe("assistant"); if (assistant?.role === "assistant") { expect(assistant.stopReason).toBe("stop"); expect(assistant.content).toEqual([{ type: "text", text: "complete" }]); } expect(events.some((event) => event.type === "agent_end")).toBe(true); });
it("should not mask stream result errors when abort is already signaled", async () => { const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [], }; const controller = new AbortController(); const providerError = new Error("provider parse failed"); const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, }; const events: AgentEvent[] = [];
await expect( runAgentLoop( [createUserMessage("Hello")], context, config, (event) => { events.push(event); }, controller.signal, () => new ThrowingResultStream(() => { controller.abort(); }, providerError), ), ).rejects.toThrow("provider parse failed"); expect( events.some( (event) => event.type === "message_end" && event.message.role === "assistant" && event.message.stopReason === "aborted", ), ).toBe(false); });
it("should return an aborted assistant when abort fires before the stream starts", async () => { const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [], }; const controller = new AbortController(); const streamFn = vi.fn(() => new MockAssistantStream()); const config: AgentLoopConfig = { model: createModel(), convertToLlm: (messages) => { controller.abort(); return identityConverter(messages); }, };
const stream = agentLoop([createUserMessage("Hello")], context, config, controller.signal, streamFn); for await (const _event of stream) { }
const messages = await stream.result(); const assistant = messages.find((message) => message.role === "assistant"); expect(streamFn).not.toHaveBeenCalled(); expect(assistant?.role).toBe("assistant"); if (assistant?.role === "assistant") { expect(assistant.stopReason).toBe("aborted"); expect(assistant.errorMessage).toBe("Request was aborted"); } });
it("should end without adding an aborted assistant when abort fires after turn completion", async () => { const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [], }; const controller = new AbortController(); const finalMessage = createAssistantMessage([{ type: "text", text: "complete" }]); const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, shouldStopAfterTurn: async () => { controller.abort(); await new Promise((resolve) => setTimeout(resolve, 0)); return false; }, }; const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "stop", message: finalMessage }); }); return stream; };
const events: AgentEvent[] = []; const stream = agentLoop([createUserMessage("Hello")], context, config, controller.signal, streamFn); for await (const event of stream) { events.push(event); }
const messages = await stream.result(); const assistantMessages = messages.filter((message) => message.role === "assistant");
expect(assistantMessages).toHaveLength(1); expect(assistantMessages[0]?.role).toBe("assistant"); if (assistantMessages[0]?.role === "assistant") { expect(assistantMessages[0].stopReason).toBe("stop"); expect(assistantMessages[0].content).toEqual([{ type: "text", text: "complete" }]); } expect(events.filter((event) => event.type === "agent_end")).toHaveLength(1); });
it("should not invoke tools when the signal aborts before tool execution", async () => { const controller = new AbortController(); const toolExecute: AgentTool["execute"] = vi.fn(async () => ({ content: [{ type: "text" as const, text: "should not run" }], details: {}, })); const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [ { name: "wait", label: "Wait", description: "Wait", parameters: Type.Object({}), execute: toolExecute, }, ], }; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, beforeToolCall: async () => { controller.abort(); return undefined; }, }; const assistantMessage = createAssistantMessage( [{ type: "toolCall", id: "tool_1", name: "wait", arguments: {} }], "toolUse", ); const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "toolUse", message: assistantMessage }); }); return stream; };
const stream = agentLoop([createUserMessage("Hello")], context, config, controller.signal, streamFn); for await (const _event of stream) { } await stream.result();
expect(toolExecute).not.toHaveBeenCalled(); });
it("should stop a sequential tool batch after aborting a tool call", async () => { const controller = new AbortController(); const toolSchema = Type.Object({ value: Type.String() }); const executed: string[] = []; const tool: AgentTool<typeof toolSchema, { value: string }> = { name: "work", label: "Work", description: "Work", parameters: toolSchema, execute: async (_toolCallId, params) => { executed.push(params.value); if (params.value === "first") { controller.abort(); } return { content: [{ type: "text", text: `done:${params.value}` }], details: { value: params.value }, }; }, }; const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [tool], }; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, toolExecution: "sequential", }; const assistantMessage = createAssistantMessage( [ { type: "toolCall", id: "tool_1", name: "work", arguments: { value: "first" } }, { type: "toolCall", id: "tool_2", name: "work", arguments: { value: "second" } }, ], "toolUse", ); const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "toolUse", message: assistantMessage }); }); return stream; }; const events: AgentEvent[] = [];
await runAgentLoop( [createUserMessage("Hello")], context, config, (event) => { events.push(event); }, controller.signal, streamFn, );
const toolStartIds = events.flatMap((event) => (event.type === "tool_execution_start" ? [event.toolCallId] : [])); const toolResultIds = events.flatMap((event) => event.type === "message_end" && event.message.role === "toolResult" ? [event.message.toolCallId] : [], );
expect(executed).toEqual(["first"]); expect(toolStartIds).toEqual(["tool_1"]); expect(toolResultIds).toEqual(["tool_1"]); expect(events.some((event) => event.type === "agent_end")).toBe(true); });
it("should coalesce synchronous tool update bursts and flush the newest partial", async () => { const toolSchema = Type.Object({}); const tool: AgentTool<typeof toolSchema, { index: number }> = { name: "work", label: "Work", description: "Work", parameters: toolSchema, execute: async (_toolCallId, _params, _signal, onUpdate) => { for (let index = 0; index < 1000; index++) { onUpdate?.({ content: [{ type: "text", text: `progress:${index}` }], details: { index } }); } return { content: [{ type: "text", text: "done" }], details: { index: 1000 } }; }, }; const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [tool], }; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, }; const assistantMessage = createAssistantMessage( [{ type: "toolCall", id: "tool_1", name: "work", arguments: {} }], "toolUse", ); let streamCall = 0; const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { if (streamCall++ === 0) { stream.push({ type: "done", reason: "toolUse", message: assistantMessage }); } else { const done = createAssistantMessage([{ type: "text", text: "done" }]); stream.push({ type: "done", reason: "stop", message: done }); } }); return stream; }; const events: AgentEvent[] = [];
await runAgentLoop( [createUserMessage("Hello")], context, config, (event) => { events.push(event); }, undefined, streamFn, );
const updates = events.flatMap((event) => (event.type === "tool_execution_update" ? [event] : [])); const toolEndIndex = events.findIndex((event) => event.type === "tool_execution_end"); expect(updates).toHaveLength(2); expect(updates.at(-1)?.partialResult).toEqual({ content: [{ type: "text", text: "progress:999" }], details: { index: 999 }, }); expect(events.indexOf(updates.at(-1)!)).toBeLessThan(toolEndIndex); });
it("should preserve a successful tool result when abort fires during update flush", async () => { const controller = new AbortController(); const toolSchema = Type.Object({}); const tool: AgentTool<typeof toolSchema, Record<string, never>> = { name: "work", label: "Work", description: "Work", parameters: toolSchema, execute: async (_toolCallId, _params, _signal, onUpdate) => { onUpdate?.({ content: [{ type: "text", text: "progress" }], details: {} }); return { content: [{ type: "text", text: "done" }], details: {}, }; }, }; const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [tool], }; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, toolExecution: "sequential", }; const assistantMessage = createAssistantMessage( [{ type: "toolCall", id: "tool_1", name: "work", arguments: {} }], "toolUse", ); const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "toolUse", message: assistantMessage }); }); return stream; }; const events: AgentEvent[] = [];
const messages = await runAgentLoop( [createUserMessage("Hello")], context, config, (event) => { events.push(event); if (event.type === "tool_execution_update") { return new Promise<void>(() => { setTimeout(() => controller.abort(), 0); }); } }, controller.signal, streamFn, );
const toolResult = messages.find((message) => message.role === "toolResult"); expect(toolResult?.role).toBe("toolResult"); if (toolResult?.role === "toolResult") { expect(toolResult.isError).toBe(false); expect(toolResult.content).toEqual([{ type: "text", text: "done" }]); } expect(events.some((event) => event.type === "agent_end")).toBe(true); });
it("should freeze synthetic aborted messages against later partial mutation", async () => { const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [], }; const controller = new AbortController(); const partialMessage = createAssistantMessage([{ type: "text", text: "partial" }]); partialMessage.usage.cost.total = 1; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, }; const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "start", partial: partialMessage }); }); return stream; };
const stream = agentLoop([createUserMessage("Hello")], context, config, controller.signal, streamFn); for await (const event of stream) { if (event.type === "message_start" && event.message.role === "assistant") { controller.abort(); } }
const messages = await stream.result(); const assistant = messages.find((message) => message.role === "assistant"); const text = partialMessage.content[0]; if (text?.type === "text") { text.text = "mutated"; } partialMessage.usage.cost.total = 99;
expect(assistant?.role).toBe("assistant"); if (assistant?.role === "assistant") { expect(assistant.stopReason).toBe("aborted"); expect(assistant.content).toEqual([{ type: "text", text: "partial" }]); expect(assistant.usage.cost.total).toBe(1); } });
it("should handle custom message types via convertToLlm", async () => { interface CustomNotification { role: "notification"; text: string; timestamp: number; }
const notification: CustomNotification = { role: "notification", text: "This is a notification", timestamp: Date.now(), };
const context: AgentContext = { systemPrompt: "You are helpful.", messages: [notification as unknown as AgentMessage], // Custom message in context tools: [], };
const userPrompt: AgentMessage = createUserMessage("Hello");
let convertedMessages: Message[] = []; const config: AgentLoopConfig = { model: createModel(), convertToLlm: (messages) => { convertedMessages = messages .filter((m) => (m as { role: string }).role !== "notification") .filter((m) => m.role === "user" || m.role === "assistant" || m.role === "toolResult") as Message[]; return convertedMessages; }, };
const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { const message = createAssistantMessage([{ type: "text", text: "Response" }]); stream.push({ type: "done", reason: "stop", message }); }); return stream; };
const events: AgentEvent[] = []; const stream = agentLoop([userPrompt], context, config, undefined, streamFn);
for await (const event of stream) { events.push(event); }
expect(convertedMessages.length).toBe(1); // Only user message expect(convertedMessages[0].role).toBe("user"); });
it("should resolve the system prompt after asynchronous API key lookup", async () => { let systemPrompt = "before lookup"; let resolveApiKey: ((apiKey: string) => void) | undefined; const apiKey = new Promise<string>((resolve) => { resolveApiKey = resolve; }); let markLookupStarted: (() => void) | undefined; const lookupStarted = new Promise<void>((resolve) => { markLookupStarted = resolve; }); let providerSystemPrompt: string | undefined; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, getApiKey: async () => { markLookupStarted?.(); return apiKey; }, getSystemPrompt: () => systemPrompt, }; const stream = agentLoop( [createUserMessage("Hello")], { systemPrompt: "fallback", messages: [], tools: [] }, config, undefined, (_model, context, options) => { providerSystemPrompt = context.systemPrompt; expect(options?.apiKey).toBe("resolved key"); const mockStream = new MockAssistantStream(); queueMicrotask(() => { mockStream.push({ type: "done", reason: "stop", message: createAssistantMessage([{ type: "text", text: "Response" }]), }); }); return mockStream; }, ); const consume = (async () => { for await (const _event of stream) { } })();
await lookupStarted; systemPrompt = "after lookup"; resolveApiKey?.("resolved key"); await consume;
expect(providerSystemPrompt).toBe("after lookup"); });
it("should apply transformContext before convertToLlm", async () => { const context: AgentContext = { systemPrompt: "You are helpful.", messages: [ createUserMessage("old message 1"), createAssistantMessage([{ type: "text", text: "old response 1" }]), createUserMessage("old message 2"), createAssistantMessage([{ type: "text", text: "old response 2" }]), ], tools: [], };
const userPrompt: AgentMessage = createUserMessage("new message");
let transformedMessages: AgentMessage[] = []; let convertedMessages: Message[] = [];
const config: AgentLoopConfig = { model: createModel(), transformContext: async (messages) => { transformedMessages = messages.slice(-2); return transformedMessages; }, convertToLlm: (messages) => { convertedMessages = messages.filter( (m) => m.role === "user" || m.role === "assistant" || m.role === "toolResult", ) as Message[]; return convertedMessages; }, };
const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { const message = createAssistantMessage([{ type: "text", text: "Response" }]); stream.push({ type: "done", reason: "stop", message }); }); return stream; };
const stream = agentLoop([userPrompt], context, config, undefined, streamFn);
for await (const _ of stream) { }
expect(transformedMessages.length).toBe(2); expect(convertedMessages.length).toBe(2); });
it("should execute mutated beforeToolCall args without revalidation", async () => { const toolSchema = Type.Object({ value: Type.String() }); const executed: Array<string | number> = []; const tool: AgentTool<typeof toolSchema, { value: string | number }> = { name: "echo", label: "Echo", description: "Echo tool", parameters: toolSchema, async execute(_toolCallId, params) { executed.push(params.value as string | number); return { content: [{ type: "text", text: `echoed: ${String(params.value)}` }], details: { value: params.value as string | number }, }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool], };
const userPrompt: AgentMessage = createUserMessage("echo something");
const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, beforeToolCall: async ({ args }) => { const mutableArgs = args as { value: string | number }; mutableArgs.value = 123; return undefined; }, };
let callIndex = 0; const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { if (callIndex === 0) { const message = createAssistantMessage( [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }], "toolUse", ); stream.push({ type: "done", reason: "toolUse", message }); } else { const message = createAssistantMessage([{ type: "text", text: "done" }]); stream.push({ type: "done", reason: "stop", message }); } callIndex++; }); return stream; };
const stream = agentLoop([userPrompt], context, config, undefined, streamFn); for await (const _event of stream) { }
expect(executed).toEqual([123]); });
it("should prepare tool arguments for validation", async () => { const replaceSchema = Type.Object({ oldText: Type.String(), newText: Type.String() }); const toolSchema = Type.Object({ edits: Type.Array(replaceSchema) }); const executed: Array<Array<{ oldText: string; newText: string }>> = []; const tool: AgentTool<typeof toolSchema, { count: number }> = { name: "edit", label: "Edit", description: "Edit tool", parameters: toolSchema, prepareArguments(args) { if (!args || typeof args !== "object") { return args as { edits: { oldText: string; newText: string }[] }; } const input = args as { edits?: Array<{ oldText: string; newText: string }>; oldText?: string; newText?: string; }; if (typeof input.oldText !== "string" || typeof input.newText !== "string") { return args as { edits: { oldText: string; newText: string }[] }; } return { edits: [...(input.edits ?? []), { oldText: input.oldText, newText: input.newText }], }; }, async execute(_toolCallId, params) { executed.push(params.edits); return { content: [{ type: "text", text: `edited ${params.edits.length}` }], details: { count: params.edits.length }, }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool], };
const userPrompt: AgentMessage = createUserMessage("edit something"); const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, };
let callIndex = 0; const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { if (callIndex === 0) { const message = createAssistantMessage( [ { type: "toolCall", id: "tool-1", name: "edit", arguments: { oldText: "before", newText: "after" }, }, ], "toolUse", ); stream.push({ type: "done", reason: "toolUse", message }); } else { const message = createAssistantMessage([{ type: "text", text: "done" }]); stream.push({ type: "done", reason: "stop", message }); } callIndex++; }); return stream; };
const stream = agentLoop([userPrompt], context, config, undefined, streamFn); for await (const _event of stream) { }
expect(executed).toEqual([[{ oldText: "before", newText: "after" }]]); });
it("should emit tool_execution_end in completion order but persist tool results in source order", async () => { const toolSchema = Type.Object({ value: Type.String() }); let firstResolved = false; let parallelObserved = false; let releaseFirst: (() => void) | undefined; const firstDone = new Promise<void>((resolve) => { releaseFirst = resolve; });
const tool: AgentTool<typeof toolSchema, { value: string }> = { name: "echo", label: "Echo", description: "Echo tool", parameters: toolSchema, async execute(_toolCallId, params) { if (params.value === "first") { await firstDone; firstResolved = true; } if (params.value === "second" && !firstResolved) { parallelObserved = true; } return { content: [{ type: "text", text: `echoed: ${params.value}` }], details: { value: params.value }, }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool], };
const userPrompt: AgentMessage = createUserMessage("echo both"); const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, toolExecution: "parallel", };
let callIndex = 0; const stream = agentLoop([userPrompt], context, config, undefined, () => { const mockStream = new MockAssistantStream(); queueMicrotask(() => { if (callIndex === 0) { const message = createAssistantMessage( [ { type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "first" } }, { type: "toolCall", id: "tool-2", name: "echo", arguments: { value: "second" } }, ], "toolUse", ); mockStream.push({ type: "done", reason: "toolUse", message }); setTimeout(() => releaseFirst?.(), 20); } else { const message = createAssistantMessage([{ type: "text", text: "done" }]); mockStream.push({ type: "done", reason: "stop", message }); } callIndex++; }); return mockStream; });
const events: AgentEvent[] = []; for await (const event of stream) { events.push(event); }
const toolExecutionEndIds = events.flatMap((event) => { if (event.type !== "tool_execution_end") { return []; } return [event.toolCallId]; }); const toolResultIds = events.flatMap((event) => { if (event.type !== "message_end" || event.message.role !== "toolResult") { return []; } return [event.message.toolCallId]; }); const turnToolResultIds = events.flatMap((event) => { if (event.type !== "turn_end") { return []; } return event.toolResults.map((toolResult) => toolResult.toolCallId); });
expect(parallelObserved).toBe(true); expect(toolExecutionEndIds).toEqual(["tool-2", "tool-1"]); expect(toolResultIds).toEqual(["tool-1", "tool-2"]); expect(turnToolResultIds).toEqual(["tool-1", "tool-2"]); });
it("should measure and persist tool call duration", async () => { const toolSchema = Type.Object({ value: Type.String() }); const tool: AgentTool<typeof toolSchema> = { name: "echo", label: "Echo", description: "Echo tool", parameters: toolSchema, async execute() { await new Promise((resolve) => setTimeout(resolve, 20)); return { content: [{ type: "text", text: "ok" }], details: {}, }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool], };
const userPrompt: AgentMessage = createUserMessage("measure"); const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, };
let callIndex = 0; const stream = agentLoop([userPrompt], context, config, undefined, () => { const mockStream = new MockAssistantStream(); queueMicrotask(() => { if (callIndex === 0) { const message = createAssistantMessage( [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "x" } }], "toolUse", ); mockStream.push({ type: "done", reason: "toolUse", message }); } else { const message = createAssistantMessage([{ type: "text", text: "done" }]); mockStream.push({ type: "done", reason: "stop", message }); } callIndex++; }); return mockStream; });
const events: AgentEvent[] = []; for await (const event of stream) { events.push(event); }
const endEvents = events.filter( (e): e is Extract<AgentEvent, { type: "tool_execution_end" }> => e.type === "tool_execution_end", ); expect(endEvents.length).toBe(1); expect(endEvents[0].durationMs).toBeGreaterThanOrEqual(15);
const toolResults = events.flatMap((event) => { if (event.type !== "message_end" || event.message.role !== "toolResult") { return []; } return [event.message]; }); expect(toolResults.length).toBe(1); expect(toolResults[0].durationMs).toBeGreaterThanOrEqual(15); });
it("should inject queued messages after all tool calls complete", async () => { const toolSchema = Type.Object({ value: Type.String() }); const executed: string[] = []; const tool: AgentTool<typeof toolSchema, { value: string }> = { name: "echo", label: "Echo", description: "Echo tool", parameters: toolSchema, async execute(_toolCallId, params) { executed.push(params.value); return { content: [{ type: "text", text: `ok:${params.value}` }], details: { value: params.value }, }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool], };
const userPrompt: AgentMessage = createUserMessage("start"); const queuedUserMessage: AgentMessage = createUserMessage("interrupt");
let queuedDelivered = false; let callIndex = 0; let sawInterruptInContext = false;
const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, toolExecution: "sequential", getSteeringMessages: async () => { if (executed.length >= 1 && !queuedDelivered) { queuedDelivered = true; return [queuedUserMessage]; } return []; }, };
const events: AgentEvent[] = []; const stream = agentLoop([userPrompt], context, config, undefined, (_model, ctx, _options) => { if (callIndex === 1) { sawInterruptInContext = ctx.messages.some( (m) => m.role === "user" && typeof m.content === "string" && m.content === "interrupt", ); }
const mockStream = new MockAssistantStream(); queueMicrotask(() => { if (callIndex === 0) { const message = createAssistantMessage( [ { type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "first" } }, { type: "toolCall", id: "tool-2", name: "echo", arguments: { value: "second" } }, ], "toolUse", ); mockStream.push({ type: "done", reason: "toolUse", message }); } else { const message = createAssistantMessage([{ type: "text", text: "done" }]); mockStream.push({ type: "done", reason: "stop", message }); } callIndex++; }); return mockStream; });
for await (const event of stream) { events.push(event); }
expect(executed).toEqual(["first", "second"]);
const toolEnds = events.filter( (e): e is Extract<AgentEvent, { type: "tool_execution_end" }> => e.type === "tool_execution_end", ); expect(toolEnds.length).toBe(2); expect(toolEnds[0].isError).toBe(false); expect(toolEnds[1].isError).toBe(false);
const eventSequence = events.flatMap((event) => { if (event.type !== "message_start") return []; if (event.message.role === "toolResult") return [`tool:${event.message.toolCallId}`]; if (event.message.role === "user" && typeof event.message.content === "string") { return [event.message.content]; } return []; }); expect(eventSequence).toContain("interrupt"); expect(eventSequence.indexOf("tool:tool-1")).toBeLessThan(eventSequence.indexOf("interrupt")); expect(eventSequence.indexOf("tool:tool-2")).toBeLessThan(eventSequence.indexOf("interrupt"));
expect(sawInterruptInContext).toBe(true); });
it("should inject continuation messages when the agent would otherwise stop", async () => { const context: AgentContext = { systemPrompt: "", messages: [], tools: [], };
let continuationPolls = 0; let sawContinuationInContext = false; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, getContinuationMessages: async ({ message }) => { continuationPolls++; expect(message.role).toBe("assistant"); if (continuationPolls === 1) { return [createUserMessage("continue")]; } return []; }, };
let callIndex = 0; const stream = agentLoop([createUserMessage("start")], context, config, undefined, (_model, ctx) => { if (callIndex === 1) { sawContinuationInContext = ctx.messages.some( (message) => message.role === "user" && message.content === "continue", ); } const mockStream = new MockAssistantStream(); queueMicrotask(() => { const message = createAssistantMessage([{ type: "text", text: callIndex === 0 ? "paused" : "done" }]); mockStream.push({ type: "done", reason: "stop", message }); callIndex++; }); return mockStream; });
for await (const _event of stream) { }
const messages = await stream.result(); expect(callIndex).toBe(2); expect(continuationPolls).toBe(2); expect(sawContinuationInContext).toBe(true); expect(messages.map((message) => message.role)).toEqual(["user", "assistant", "user", "assistant"]); });
it("should prefer explicit follow-up messages before continuation messages", async () => { const context: AgentContext = { systemPrompt: "", messages: [], tools: [], };
let followUpDelivered = false; let continuationPolls = 0; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, getFollowUpMessages: async () => { if (followUpDelivered) { return []; } followUpDelivered = true; return [createUserMessage("follow up")]; }, getContinuationMessages: async () => { continuationPolls++; return []; }, };
let callIndex = 0; let sawFollowUpInContext = false; const stream = agentLoop([createUserMessage("start")], context, config, undefined, (_model, ctx) => { if (callIndex === 1) { sawFollowUpInContext = ctx.messages.some( (message) => message.role === "user" && message.content === "follow up", ); } const mockStream = new MockAssistantStream(); queueMicrotask(() => { const message = createAssistantMessage([{ type: "text", text: callIndex === 0 ? "paused" : "done" }]); mockStream.push({ type: "done", reason: "stop", message }); callIndex++; }); return mockStream; });
for await (const _event of stream) { }
expect(callIndex).toBe(2); expect(continuationPolls).toBe(1); expect(sawFollowUpInContext).toBe(true); });
it("should force sequential execution when a tool has executionMode=sequential even with default parallel config", async () => { const toolSchema = Type.Object({ value: Type.String() }); let firstResolved = false; let parallelObserved = false; let releaseFirst: (() => void) | undefined; const firstDone = new Promise<void>((resolve) => { releaseFirst = resolve; });
const slowTool: AgentTool<typeof toolSchema, { value: string }> = { name: "slow", label: "Slow", description: "Slow tool", parameters: toolSchema, executionMode: "sequential", async execute(_toolCallId, params) { if (params.value === "first") { await firstDone; firstResolved = true; } if (params.value === "second" && !firstResolved) { parallelObserved = true; } return { content: [{ type: "text", text: `slow: ${params.value}` }], details: { value: params.value }, }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [slowTool], };
const userPrompt: AgentMessage = createUserMessage("run both"); const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, };
let callIndex = 0; const stream = agentLoop([userPrompt], context, config, undefined, () => { const mockStream = new MockAssistantStream(); queueMicrotask(() => { if (callIndex === 0) { const message = createAssistantMessage( [ { type: "toolCall", id: "tool-1", name: "slow", arguments: { value: "first" } }, { type: "toolCall", id: "tool-2", name: "slow", arguments: { value: "second" } }, ], "toolUse", ); mockStream.push({ type: "done", reason: "toolUse", message }); setTimeout(() => releaseFirst?.(), 20); } else { const message = createAssistantMessage([{ type: "text", text: "done" }]); mockStream.push({ type: "done", reason: "stop", message }); } callIndex++; }); return mockStream; });
const events: AgentEvent[] = []; for await (const event of stream) { events.push(event); }
expect(parallelObserved).toBe(false);
const toolResultIds = events.flatMap((event) => { if (event.type !== "message_end" || event.message.role !== "toolResult") { return []; } return [event.message.toolCallId]; }); expect(toolResultIds).toEqual(["tool-1", "tool-2"]); });
it("should force sequential execution when one of multiple tools has executionMode=sequential", async () => { const toolSchema = Type.Object({ value: Type.String() }); const executionOrder: string[] = []; let releaseSlow: (() => void) | undefined; const slowDone = new Promise<void>((resolve) => { releaseSlow = resolve; });
const slowTool: AgentTool<typeof toolSchema, { value: string }> = { name: "slow", label: "Slow", description: "Slow tool", parameters: toolSchema, executionMode: "sequential", async execute(_toolCallId, params) { executionOrder.push(`slow:${params.value}`); if (params.value === "a") { await slowDone; } return { content: [{ type: "text", text: `slow: ${params.value}` }], details: { value: params.value }, }; }, };
const fastTool: AgentTool<typeof toolSchema, { value: string }> = { name: "fast", label: "Fast", description: "Fast tool", parameters: toolSchema, async execute(_toolCallId, params) { executionOrder.push(`fast:${params.value}`); return { content: [{ type: "text", text: `fast: ${params.value}` }], details: { value: params.value }, }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [slowTool, fastTool], };
const userPrompt: AgentMessage = createUserMessage("run both"); const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, };
let callIndex = 0; const stream = agentLoop([userPrompt], context, config, undefined, () => { const mockStream = new MockAssistantStream(); queueMicrotask(() => { if (callIndex === 0) { const message = createAssistantMessage( [ { type: "toolCall", id: "tool-1", name: "slow", arguments: { value: "a" } }, { type: "toolCall", id: "tool-2", name: "fast", arguments: { value: "b" } }, ], "toolUse", ); mockStream.push({ type: "done", reason: "toolUse", message }); setTimeout(() => releaseSlow?.(), 20); } else { const message = createAssistantMessage([{ type: "text", text: "done" }]); mockStream.push({ type: "done", reason: "stop", message }); } callIndex++; }); return mockStream; });
const events: AgentEvent[] = []; for await (const event of stream) { events.push(event); }
expect(executionOrder[0]).toBe("slow:a"); expect(executionOrder).toContain("fast:b"); });
it("should allow parallel execution when all tools have executionMode=parallel", async () => { const toolSchema = Type.Object({ value: Type.String() }); let firstResolved = false; let parallelObserved = false; let releaseFirst: (() => void) | undefined; const firstDone = new Promise<void>((resolve) => { releaseFirst = resolve; });
const tool: AgentTool<typeof toolSchema, { value: string }> = { name: "echo", label: "Echo", description: "Echo tool", parameters: toolSchema, executionMode: "parallel", async execute(_toolCallId, params) { if (params.value === "first") { await firstDone; firstResolved = true; } if (params.value === "second" && !firstResolved) { parallelObserved = true; } return { content: [{ type: "text", text: `echoed: ${params.value}` }], details: { value: params.value }, }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool], };
const userPrompt: AgentMessage = createUserMessage("echo both"); const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, };
let callIndex = 0; const stream = agentLoop([userPrompt], context, config, undefined, () => { const mockStream = new MockAssistantStream(); queueMicrotask(() => { if (callIndex === 0) { const message = createAssistantMessage( [ { type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "first" } }, { type: "toolCall", id: "tool-2", name: "echo", arguments: { value: "second" } }, ], "toolUse", ); mockStream.push({ type: "done", reason: "toolUse", message }); setTimeout(() => releaseFirst?.(), 20); } else { const message = createAssistantMessage([{ type: "text", text: "done" }]); mockStream.push({ type: "done", reason: "stop", message }); } callIndex++; }); return mockStream; });
const events: AgentEvent[] = []; for await (const event of stream) { events.push(event); }
expect(parallelObserved).toBe(true); });
const naturalCompletionCases: Array<{ name: string; stopBefore: "always" | "afterTurn" | "never"; steerOnce?: boolean; naturalWork?: boolean; expected: { llmCalls: number; stopChecks?: number; steeringPolls?: number; naturalPolls?: number }; }> = [ { name: "stops only after completing the initial turn", stopBefore: "always", expected: { llmCalls: 1, stopChecks: 1 }, }, { name: "observes a stop requested after the turn before polling steering", stopBefore: "afterTurn", expected: { llmCalls: 1, steeringPolls: 1 }, }, { name: "stops before polling natural-completion work", stopBefore: "always", naturalWork: true, expected: { llmCalls: 1, naturalPolls: 0 }, }, { name: "continues with steering when no stop is requested", stopBefore: "never", steerOnce: true, expected: { llmCalls: 2 }, }, ];
it.each(naturalCompletionCases)("$name", async ({ stopBefore, steerOnce, naturalWork, expected }) => { let llmCalls = 0; let stopRequested = false; let steeringDelivered = false; const shouldStopBeforeTurn = vi.fn(() => stopBefore === "always" || stopRequested); const getSteeringMessages = vi.fn(async () => { if (steerOnce && llmCalls === 1 && !steeringDelivered) { steeringDelivered = true; return [createUserMessage("steer")]; } return []; }); const getFollowUpMessages = vi.fn(async () => (naturalWork ? [createUserMessage("follow up")] : [])); const getContinuationMessages = vi.fn(async () => (naturalWork ? [createUserMessage("continue")] : [])); const stream = agentLoop( [createUserMessage("start")], { systemPrompt: "", messages: [], tools: [] }, { model: createModel(), convertToLlm: identityConverter, shouldStopBeforeTurn, getSteeringMessages, getFollowUpMessages, getContinuationMessages, ...(stopBefore === "afterTurn" && { shouldStopAfterTurn: async () => { await Promise.resolve(); stopRequested = true; return false; }, }), }, undefined, () => { llmCalls++; const mockStream = new MockAssistantStream(); queueMicrotask(() => mockStream.push({ type: "done", reason: "stop", message: createAssistantMessage([{ type: "text", text: "done" }]), }), ); return mockStream; }, ); for await (const _event of stream) { }
expect(llmCalls).toBe(expected.llmCalls); if (expected.stopChecks !== undefined) expect(shouldStopBeforeTurn).toHaveBeenCalledTimes(expected.stopChecks); if (expected.steeringPolls !== undefined) expect(getSteeringMessages).toHaveBeenCalledTimes(expected.steeringPolls); if (expected.naturalPolls !== undefined) { expect(getFollowUpMessages).toHaveBeenCalledTimes(expected.naturalPolls); expect(getContinuationMessages).toHaveBeenCalledTimes(expected.naturalPolls); } });
it("rechecks a tool-cycle stop after polling steering", async () => { const execute = vi.fn(async () => ({ content: [{ type: "text" as const, text: "done" }], details: {} })); let stopRequested = false; let llmCalls = 0; const stream = agentLoop( [createUserMessage("start")], { systemPrompt: "", messages: [], tools: [{ name: "work", label: "Work", description: "Work", parameters: Type.Object({}), execute }], }, { model: createModel(), convertToLlm: identityConverter, shouldStopBeforeTurn: () => stopRequested, getSteeringMessages: async () => { if (llmCalls > 0) stopRequested = true; return []; }, }, undefined, () => { llmCalls++; const mockStream = new MockAssistantStream(); queueMicrotask(() => mockStream.push({ type: "done", reason: "toolUse", message: createAssistantMessage( [{ type: "toolCall", id: "tool-1", name: "work", arguments: {} }], "toolUse", ), }), ); return mockStream; }, ); for await (const _event of stream) { } expect(llmCalls).toBe(1); });
it("injects steering drained by the poll even when a stop flips during it", async () => { const execute = vi.fn(async () => ({ content: [{ type: "text" as const, text: "done" }], details: {} })); let stopRequested = false; const injected: string[] = []; let llmCalls = 0; const stream = agentLoop( [createUserMessage("start")], { systemPrompt: "", messages: [], tools: [{ name: "work", label: "Work", description: "Work", parameters: Type.Object({}), execute }], }, { model: createModel(), convertToLlm: identityConverter, shouldStopBeforeTurn: () => stopRequested, getSteeringMessages: async () => { if (llmCalls !== 1 || stopRequested) return []; stopRequested = true; return [createUserMessage("late steer")]; }, }, undefined, () => { llmCalls++; const mockStream = new MockAssistantStream(); const calls = llmCalls; queueMicrotask(() => mockStream.push({ type: "done", reason: calls === 1 ? "toolUse" : "stop", message: calls === 1 ? createAssistantMessage( [{ type: "toolCall", id: "tool-1", name: "work", arguments: {} }], "toolUse", ) : createAssistantMessage([{ type: "text", text: "steered reply" }], "stop"), }), ); return mockStream; }, ); for await (const event of stream) { if (event.type === "message_end" && event.message.role === "user") { const text = event.message.content; if (typeof text !== "string") continue; injected.push(text); } } expect(injected).toContain("late steer"); expect(llmCalls).toBe(2); });
it("checks shouldStopBeforeTurn after a tool batch and before another model call", async () => { const execute = vi.fn(async () => ({ content: [{ type: "text" as const, text: "done" }], details: {} })); const context: AgentContext = { systemPrompt: "", messages: [], tools: [ { name: "work", label: "Work", description: "Work", parameters: Type.Object({}), execute, }, ], }; let llmCalls = 0; const stream = agentLoop( [createUserMessage("start")], context, { model: createModel(), convertToLlm: identityConverter, shouldStopBeforeTurn: () => true }, undefined, () => { llmCalls++; const mockStream = new MockAssistantStream(); queueMicrotask(() => mockStream.push({ type: "done", reason: "toolUse", message: createAssistantMessage( [{ type: "toolCall", id: "tool-1", name: "work", arguments: {} }], "toolUse", ), }), ); return mockStream; }, ); for await (const _event of stream) { }
expect(execute).toHaveBeenCalledOnce(); expect(llmCalls).toBe(1); });
it("should stop after the current turn when shouldStopAfterTurn returns true", async () => { const toolSchema = Type.Object({ value: Type.String() }); const executed: string[] = []; const tool: AgentTool<typeof toolSchema, { value: string }> = { name: "echo", label: "Echo", description: "Echo tool", parameters: toolSchema, async execute(_toolCallId, params) { executed.push(params.value); return { content: [{ type: "text", text: `echoed: ${params.value}` }], details: { value: params.value }, }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool], };
let steeringPolls = 0; let followUpPolls = 0; let callbackToolResultIds: string[] = []; let callbackContextRoles: string[] = []; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, getSteeringMessages: async () => { steeringPolls++; return []; }, getFollowUpMessages: async () => { followUpPolls++; return [createUserMessage("follow up should stay queued")]; }, shouldStopAfterTurn: async ({ message, toolResults, context }) => { expect(message.role).toBe("assistant"); callbackToolResultIds = toolResults.map((toolResult) => toolResult.toolCallId); callbackContextRoles = context.messages.map((contextMessage) => contextMessage.role); return true; }, };
let llmCalls = 0; const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, () => { llmCalls++; const mockStream = new MockAssistantStream(); queueMicrotask(() => { if (llmCalls === 1) { const message = createAssistantMessage( [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }], "toolUse", ); mockStream.push({ type: "done", reason: "toolUse", message }); } else { mockStream.push({ type: "done", reason: "stop", message: createAssistantMessage([{ type: "text", text: "should not run" }]), }); } }); return mockStream; });
const events: AgentEvent[] = []; for await (const event of stream) { events.push(event); }
const messages = await stream.result(); expect(llmCalls).toBe(1); expect(executed).toEqual(["hello"]); expect(steeringPolls).toBe(1); expect(followUpPolls).toBe(0); expect(callbackToolResultIds).toEqual(["tool-1"]); expect(callbackContextRoles).toEqual(["user", "assistant", "toolResult"]); expect(messages.map((message) => message.role)).toEqual(["user", "assistant", "toolResult"]); expect(events.map((event) => event.type)).toEqual([ "agent_start", "turn_start", "message_start", "message_end", "message_start", "message_end", "tool_execution_start", "tool_execution_end", "message_start", "message_end", "turn_end", "agent_end", ]); });
it("should stop after a tool batch when every tool result sets terminate=true", async () => { const toolSchema = Type.Object({ value: Type.String() }); const tool: AgentTool<typeof toolSchema, { value: string }> = { name: "echo", label: "Echo", description: "Echo tool", parameters: toolSchema, async execute(_toolCallId, params) { return { content: [{ type: "text", text: `echoed: ${params.value}` }], details: { value: params.value }, terminate: true, }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool], };
const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, };
let llmCalls = 0; const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, () => { llmCalls++; const mockStream = new MockAssistantStream(); queueMicrotask(() => { const message = createAssistantMessage( [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }], "toolUse", ); mockStream.push({ type: "done", reason: "toolUse", message }); }); return mockStream; });
const events: AgentEvent[] = []; for await (const event of stream) { events.push(event); }
const messages = await stream.result(); expect(llmCalls).toBe(1); expect(messages.map((message) => message.role)).toEqual(["user", "assistant", "toolResult"]); expect(events.filter((event) => event.type === "turn_end")).toHaveLength(1); });
it("should continue after parallel tool calls when not all tool results terminate", async () => { const toolSchema = Type.Object({ value: Type.String() }); const tool: AgentTool<typeof toolSchema, { value: string }> = { name: "echo", label: "Echo", description: "Echo tool", parameters: toolSchema, async execute(_toolCallId, params) { return { content: [{ type: "text", text: `echoed: ${params.value}` }], details: { value: params.value }, terminate: params.value === "first", }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool], };
const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, toolExecution: "parallel", };
let callIndex = 0; const stream = agentLoop([createUserMessage("echo both")], context, config, undefined, () => { const mockStream = new MockAssistantStream(); queueMicrotask(() => { if (callIndex === 0) { const message = createAssistantMessage( [ { type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "first" } }, { type: "toolCall", id: "tool-2", name: "echo", arguments: { value: "second" } }, ], "toolUse", ); mockStream.push({ type: "done", reason: "toolUse", message }); } else { const message = createAssistantMessage([{ type: "text", text: "done" }]); mockStream.push({ type: "done", reason: "stop", message }); } callIndex++; }); return mockStream; });
for await (const _event of stream) { }
const messages = await stream.result(); expect(callIndex).toBe(2); expect(messages.map((message) => message.role)).toEqual([ "user", "assistant", "toolResult", "toolResult", "assistant", ]); });
it("should allow afterToolCall to mark a tool batch as terminating", async () => { const toolSchema = Type.Object({ value: Type.String() }); const tool: AgentTool<typeof toolSchema, { value: string }> = { name: "echo", label: "Echo", description: "Echo tool", parameters: toolSchema, async execute(_toolCallId, params) { return { content: [{ type: "text", text: `echoed: ${params.value}` }], details: { value: params.value }, }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool], };
const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, afterToolCall: async () => ({ terminate: true }), };
let llmCalls = 0; const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, () => { llmCalls++; const mockStream = new MockAssistantStream(); queueMicrotask(() => { const message = createAssistantMessage( [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }], "toolUse", ); mockStream.push({ type: "done", reason: "toolUse", message }); }); return mockStream; });
for await (const _event of stream) { }
expect(llmCalls).toBe(1); });});
describe("agentLoopContinue with AgentMessage", () => { it("should throw when context has no messages", () => { const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [], };
const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, };
expect(() => agentLoopContinue(context, config)).toThrow("Cannot continue: no messages in context"); });
it("should continue from existing context without emitting user message events", async () => { const userMessage: AgentMessage = createUserMessage("Hello");
const context: AgentContext = { systemPrompt: "You are helpful.", messages: [userMessage], tools: [], };
const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, };
const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { const message = createAssistantMessage([{ type: "text", text: "Response" }]); stream.push({ type: "done", reason: "stop", message }); }); return stream; };
const events: AgentEvent[] = []; const stream = agentLoopContinue(context, config, undefined, streamFn);
for await (const event of stream) { events.push(event); }
const messages = await stream.result();
expect(messages.length).toBe(1); expect(messages[0].role).toBe("assistant");
const messageEndEvents = events.filter((e) => e.type === "message_end"); expect(messageEndEvents.length).toBe(1); expect((messageEndEvents[0] as any).message.role).toBe("assistant"); });
it("should allow custom message types as last message (caller responsibility)", async () => { interface CustomMessage { role: "custom"; text: string; timestamp: number; }
const customMessage: CustomMessage = { role: "custom", text: "Hook content", timestamp: Date.now(), };
const context: AgentContext = { systemPrompt: "You are helpful.", messages: [customMessage as unknown as AgentMessage], tools: [], };
const config: AgentLoopConfig = { model: createModel(), convertToLlm: (messages) => { return messages .map((m) => { if ((m as any).role === "custom") { return { role: "user" as const, content: (m as any).text, timestamp: m.timestamp, }; } return m; }) .filter((m) => m.role === "user" || m.role === "assistant" || m.role === "toolResult") as Message[]; }, };
const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { const message = createAssistantMessage([{ type: "text", text: "Response to custom message" }]); stream.push({ type: "done", reason: "stop", message }); }); return stream; };
const stream = agentLoopContinue(context, config, undefined, streamFn);
const events: AgentEvent[] = []; for await (const event of stream) { events.push(event); }
const messages = await stream.result(); expect(messages.length).toBe(1); expect(messages[0].role).toBe("assistant"); });
it("should report a tool-declared failure on the result message and the end event", async () => { const toolSchema = Type.Object({ value: Type.String() }); const tool: AgentTool<typeof toolSchema, { value: string }> = { name: "flaky", label: "Flaky", description: "Returns a failure result instead of throwing", parameters: toolSchema, async execute(_toolCallId, params) { return { content: [{ type: "text", text: `failed:${params.value}` }], details: { value: params.value }, isError: true, }; }, }; const okTool: AgentTool<typeof toolSchema, { value: string }> = { name: "steady", label: "Steady", description: "Returns a success result", parameters: toolSchema, async execute(_toolCallId, params) { return { content: [{ type: "text", text: `ok:${params.value}` }], details: { value: params.value }, }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool, okTool], };
let callIndex = 0; const events: AgentEvent[] = []; const stream = agentLoop( [createUserMessage("start")], context, { model: createModel(), convertToLlm: identityConverter }, undefined, () => { const mockStream = new MockAssistantStream(); queueMicrotask(() => { if (callIndex === 0) { mockStream.push({ type: "done", reason: "toolUse", message: createAssistantMessage( [ { type: "toolCall", id: "tool-1", name: "flaky", arguments: { value: "a" } }, { type: "toolCall", id: "tool-2", name: "steady", arguments: { value: "b" } }, ], "toolUse", ), }); } else { mockStream.push({ type: "done", reason: "stop", message: createAssistantMessage([{ type: "text", text: "done" }]), }); } callIndex++; }); return mockStream; }, );
for await (const event of stream) { events.push(event); }
const toolResults = events.flatMap((event) => event.type === "message_end" && event.message.role === "toolResult" ? [event.message] : [], ); const toolEnds = events.filter( (event): event is Extract<AgentEvent, { type: "tool_execution_end" }> => event.type === "tool_execution_end", ); expect(toolResults.map((message) => [message.toolName, message.isError])).toEqual([ ["flaky", true], ["steady", false], ]); expect(toolEnds.map((event) => [event.toolName, event.isError])).toEqual([ ["flaky", true], ["steady", false], ]); });});
describe("AgentLoopConfig.getModel", () => { it("retargets the model between turns of one loop", async () => { const first = createModel(); const second: Model<"openai-responses"> = { ...createModel(), id: "mock-2", provider: "other" }; let selected = first; const models: string[] = []; const tool: AgentTool = { name: "wait", label: "Wait", description: "Wait", parameters: Type.Object({}), execute: async () => ({ content: [{ type: "text" as const, text: "waited" }], details: {} }), }; const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [tool], }; const config: AgentLoopConfig = { model: first, getModel: () => selected, convertToLlm: identityConverter, }; let turn = 0; const streamFn = (model: Model<any>) => { models.push(`${model.provider}/${model.id}`); const stream = new MockAssistantStream(); const message = turn++ === 0 ? createAssistantMessage([{ type: "toolCall", id: "tool_1", name: "wait", arguments: {} }], "toolUse") : createAssistantMessage([{ type: "text", text: "done" }]); if (turn === 1) selected = second; queueMicrotask(() => { stream.push({ type: "done", reason: message.stopReason === "toolUse" ? "toolUse" : "stop", message }); }); return stream; };
await runAgentLoop([createUserMessage("Hello")], context, config, () => {}, undefined, streamFn as never);
expect(models).toEqual(["openai/mock", "other/mock-2"]); });});