diff --git a/apps/api/src/agent/llm/loop.test.ts b/apps/api/src/agent/llm/loop.test.ts index 5fa7baf6..416c5203 100644 --- a/apps/api/src/agent/llm/loop.test.ts +++ b/apps/api/src/agent/llm/loop.test.ts @@ -162,3 +162,144 @@ describe("runAgentLoop context compression", () => { expect(result.finalMessage).toBe("done"); }); }); + +function assistantWithNamedTool(toolName: string, callId: string, inputToken = 100): LlmResponse { + return { + message: { + role: "assistant", + content: null, + toolCalls: [{ id: callId, name: toolName, arguments: "{}" }], + }, + finishReason: "tool_calls", + requestId: null, + usage: { inputToken, cachedInputToken: 0, outputToken: 10 }, + }; +} + +function recordingTool(name: string, calls: string[]): AgentTool { + return { + name, + description: name, + parameters: { type: "object", properties: {} }, + execute: async () => { + calls.push(name); + return "ok"; + }, + }; +} + +describe("runAgentLoop ending once the required tools have been called", () => { + it("stops after the required tool runs instead of spending another request", async () => { + const sent: number[] = []; + const calls: string[] = []; + const sender: LlmSender = { + async send() { + sent.push(sent.length); + if (sent.length === 1) { + return assistantWithNamedTool("submit", "tc1"); + } + return assistantDone("this turn should never be requested"); + }, + } as unknown as LlmSender; + + const result = await runAgentLoop(sender, "sys", "prompt", "label", { + requiredToolList: [recordingTool("submit", calls)], + endWhenRequiredToolsCalled: true, + activityId: 1, + }); + + // One request, not two: the confirming turn is what this option removes. + expect(sent).toHaveLength(1); + expect(calls).toEqual(["submit"]); + expect(result.toolCallCount.submit).toBe(1); + expect(result.stepCount).toBe(1); + }); + + it("still asks for another turn when the option is off", async () => { + const sent: number[] = []; + const calls: string[] = []; + const sender: LlmSender = { + async send() { + sent.push(sent.length); + if (sent.length === 1) { + return assistantWithNamedTool("submit", "tc1"); + } + return assistantDone("done"); + }, + } as unknown as LlmSender; + + const result = await runAgentLoop(sender, "sys", "prompt", "label", { + requiredToolList: [recordingTool("submit", calls)], + activityId: 1, + }); + + expect(sent).toHaveLength(2); + expect(result.finalMessage).toBe("done"); + }); + + it("waits for every required tool, not just the first", async () => { + const sent: number[] = []; + const calls: string[] = []; + const sender: LlmSender = { + async send() { + sent.push(sent.length); + if (sent.length === 1) return assistantWithNamedTool("first", "tc1"); + if (sent.length === 2) return assistantWithNamedTool("second", "tc2"); + return assistantDone("this turn should never be requested"); + }, + } as unknown as LlmSender; + + const result = await runAgentLoop(sender, "sys", "prompt", "label", { + requiredToolList: [recordingTool("first", calls), recordingTool("second", calls)], + endWhenRequiredToolsCalled: true, + activityId: 1, + }); + + expect(sent).toHaveLength(2); + expect(calls).toEqual(["first", "second"]); + expect(result.stepCount).toBe(2); + }); + + it("does not end early when only an optional tool was called", async () => { + const sent: number[] = []; + const calls: string[] = []; + const sender: LlmSender = { + async send() { + sent.push(sent.length); + if (sent.length === 1) return assistantWithNamedTool("noop", "tc1"); + if (sent.length === 2) return assistantWithNamedTool("submit", "tc2"); + return assistantDone("this turn should never be requested"); + }, + } as unknown as LlmSender; + + await runAgentLoop(sender, "sys", "prompt", "label", { + toolList: [recordingTool("noop", calls)], + requiredToolList: [recordingTool("submit", calls)], + endWhenRequiredToolsCalled: true, + activityId: 1, + }); + + expect(sent).toHaveLength(2); + expect(calls).toEqual(["noop", "submit"]); + }); + + it("leaves an agent with no required tool unaffected", async () => { + const sent: number[] = []; + const sender: LlmSender = { + async send() { + sent.push(sent.length); + if (sent.length === 1) return assistantWithTool(); + return assistantDone("finished"); + }, + } as unknown as LlmSender; + + const result = await runAgentLoop(sender, "sys", "prompt", "label", { + toolList: [noopTool], + endWhenRequiredToolsCalled: true, + activityId: 1, + }); + + expect(sent).toHaveLength(2); + expect(result.finalMessage).toBe("finished"); + }); +}); diff --git a/apps/api/src/agent/llm/loop.ts b/apps/api/src/agent/llm/loop.ts index a26a53ae..2d49ca03 100644 --- a/apps/api/src/agent/llm/loop.ts +++ b/apps/api/src/agent/llm/loop.ts @@ -65,6 +65,16 @@ export async function runAgentLoop( toolList?: (AgentTool | null)[]; maxSteps?: number; requiredToolList?: (AgentTool | null)[]; + /** + * Ends the loop as soon as every required tool has been called, instead of + * asking the model for one further turn that only confirms it is finished. + * + * An agent whose output is delivered by a required tool has nothing left to + * say once that tool has run, so that extra turn re-sends the whole + * transcript to be told so. Off by default, because an agent that is meant + * to keep working after calling a required tool would be cut short. + */ + endWhenRequiredToolsCalled?: boolean; activityId: number; onUsage?: (usage: ActivityTokenUsage) => Promise; }, @@ -344,6 +354,27 @@ export async function runAgentLoop( toolCallId: r.toolCallId, }); } + + if ( + options.endWhenRequiredToolsCalled && + requiredToolNameList.length > 0 && + requiredToolNameList.every((toolName) => (toolCallCount[toolName] ?? 0) > 0) + ) { + // The agent delivered its output through the required tool, so the + // turn that would follow can only confirm it is done, at the cost of + // re-sending the entire transcript. This is each agent's largest + // request, because the transcript is never longer than it is here. + const result: AgentRunResult = { + finalMessage: response.message.content, + messages, + stepCount, + toolCallCount, + usage, + }; + logAgent(activityId, "every required tool has been called, ending the loop", label); + logAgentResult(activityId, label, result, performance.now() - startedAt, "completed"); + return result; + } } } catch (error: unknown) { if (error instanceof JobCanceledError) { diff --git a/apps/api/src/agent/pull-request/review/sub.service.ts b/apps/api/src/agent/pull-request/review/sub.service.ts index 55b1727e..fd21f48c 100644 --- a/apps/api/src/agent/pull-request/review/sub.service.ts +++ b/apps/api/src/agent/pull-request/review/sub.service.ts @@ -55,6 +55,7 @@ export async function runReviewSubAgent( const result = await runAgentLoop(sender, system, prompt, `[PR #${prIid}] Sub ${index}/${totalIndex}`, { toolList, requiredToolList: [submitReviewHandoffTool(reviewHandoffList, reviewUnit)], + endWhenRequiredToolsCalled: true, activityId, onUsage: (stepUsage) => activityService.addTokenUsage(activityId, stepUsage), }); diff --git a/apps/api/src/agent/pull-request/review/writing.service.ts b/apps/api/src/agent/pull-request/review/writing.service.ts index dc81b796..84999148 100644 --- a/apps/api/src/agent/pull-request/review/writing.service.ts +++ b/apps/api/src/agent/pull-request/review/writing.service.ts @@ -116,6 +116,7 @@ export async function runReviewWritingAgent( postPullRequestCommentTool(provider, prIid, language, activityId), isFollowUpReview ? null : evaluatePullRequestTool(provider, prIid), ], + endWhenRequiredToolsCalled: true, activityId, onUsage: (stepUsage) => activityService.addTokenUsage(activityId, stepUsage), });