mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-17 10:02:23 +00:00
Simplify compaction: remove proactive abort, use Agent.continue() for retry
- Add agentLoopContinue() to pi-ai for resuming from existing context - Add Agent.continue() method and transport.continue() interface - Simplify AgentSession compaction to two cases: overflow (auto-retry) and threshold (no retry) - Remove proactive mid-turn compaction abort - Merge turn prefix summary into main summary - Add isCompacting property to AgentSession and RPC state - Block input during compaction in interactive mode - Show compaction count on session resume - Rename RPC.md to rpc.md for consistency Related to #128
This commit is contained in:
Mario Zechner
parent
d67c69c6e9
commit
5a9d844f9a
27 changed files with 1261 additions and 1011 deletions
|
|
@ -2,6 +2,10 @@
|
|||
|
||||
## [Unreleased]
|
||||
|
||||
### Added
|
||||
|
||||
- **`agentLoopContinue` function**: Continue an agent loop from existing context without adding a new user message. Validates that the last message is `user` or `toolResult`. Useful for retry after context overflow or resuming from manually-added tool results.
|
||||
|
||||
### Breaking Changes
|
||||
|
||||
- Removed provider-level tool argument validation. Validation now happens in `agentLoop` via `executeToolCalls`, allowing models to retry on validation errors. For manual tool execution, use `validateToolCall(tools, toolCall)` or `validateToolArguments(tool, toolCall)`.
|
||||
|
|
|
|||
|
|
@ -898,6 +898,34 @@ const messages = await stream.result();
|
|||
context.messages.push(...messages);
|
||||
```
|
||||
|
||||
### Continuing from Existing Context
|
||||
|
||||
Use `agentLoopContinue` to resume an agent loop without adding a new user message. This is useful for:
|
||||
- Retrying after context overflow (after compaction reduces context size)
|
||||
- Resuming from tool results that were added manually to the context
|
||||
|
||||
```typescript
|
||||
import { agentLoopContinue, AgentContext } from '@mariozechner/pi-ai';
|
||||
|
||||
// Context already has messages - last must be 'user' or 'toolResult'
|
||||
const context: AgentContext = {
|
||||
systemPrompt: 'You are helpful.',
|
||||
messages: [userMessage, assistantMessage, toolResult],
|
||||
tools: [myTool]
|
||||
};
|
||||
|
||||
// Continue processing from the tool result
|
||||
const stream = agentLoopContinue(context, { model });
|
||||
|
||||
for await (const event of stream) {
|
||||
// Same events as agentLoop, but no user message events emitted
|
||||
}
|
||||
|
||||
const newMessages = await stream.result();
|
||||
```
|
||||
|
||||
**Validation**: Throws if context has no messages or if the last message is an assistant message.
|
||||
|
||||
### Defining Tools with TypeBox
|
||||
|
||||
Tools use TypeBox schemas for runtime validation and type inference:
|
||||
|
|
|
|||
|
|
@ -4,7 +4,10 @@ import { EventStream } from "../utils/event-stream.js";
|
|||
import { validateToolArguments } from "../utils/validation.js";
|
||||
import type { AgentContext, AgentEvent, AgentLoopConfig, AgentTool, AgentToolResult, QueuedMessage } from "./types.js";
|
||||
|
||||
// Main prompt function - returns a stream of events
|
||||
/**
|
||||
* Start an agent loop with a new user message.
|
||||
* The prompt is added to the context and events are emitted for it.
|
||||
*/
|
||||
export function agentLoop(
|
||||
prompt: UserMessage,
|
||||
context: AgentContext,
|
||||
|
|
@ -12,92 +15,137 @@ export function agentLoop(
|
|||
signal?: AbortSignal,
|
||||
streamFn?: typeof streamSimple,
|
||||
): EventStream<AgentEvent, AgentContext["messages"]> {
|
||||
const stream = new EventStream<AgentEvent, AgentContext["messages"]>(
|
||||
(event: AgentEvent) => event.type === "agent_end",
|
||||
(event: AgentEvent) => (event.type === "agent_end" ? event.messages : []),
|
||||
);
|
||||
const stream = createAgentStream();
|
||||
|
||||
// Run the prompt async
|
||||
(async () => {
|
||||
// Track new messages generated during this prompt
|
||||
const newMessages: AgentContext["messages"] = [];
|
||||
// Create user message for the prompt
|
||||
const messages = [...context.messages, prompt];
|
||||
newMessages.push(prompt);
|
||||
const newMessages: AgentContext["messages"] = [prompt];
|
||||
const currentContext: AgentContext = {
|
||||
...context,
|
||||
messages: [...context.messages, prompt],
|
||||
};
|
||||
|
||||
stream.push({ type: "agent_start" });
|
||||
stream.push({ type: "turn_start" });
|
||||
stream.push({ type: "message_start", message: prompt });
|
||||
stream.push({ type: "message_end", message: prompt });
|
||||
|
||||
// Update context with new messages
|
||||
const currentContext: AgentContext = {
|
||||
...context,
|
||||
messages,
|
||||
};
|
||||
|
||||
// Keep looping while we have tool calls or queued messages
|
||||
let hasMoreToolCalls = true;
|
||||
let firstTurn = true;
|
||||
let queuedMessages: QueuedMessage<any>[] = (await config.getQueuedMessages?.()) || [];
|
||||
|
||||
while (hasMoreToolCalls || queuedMessages.length > 0) {
|
||||
if (!firstTurn) {
|
||||
stream.push({ type: "turn_start" });
|
||||
} else {
|
||||
firstTurn = false;
|
||||
}
|
||||
|
||||
// Process queued messages first (inject before next assistant response)
|
||||
if (queuedMessages.length > 0) {
|
||||
for (const { original, llm } of queuedMessages) {
|
||||
stream.push({ type: "message_start", message: original });
|
||||
stream.push({ type: "message_end", message: original });
|
||||
if (llm) {
|
||||
currentContext.messages.push(llm);
|
||||
newMessages.push(llm);
|
||||
}
|
||||
}
|
||||
queuedMessages = [];
|
||||
}
|
||||
|
||||
// console.log("agent-loop: ", [...currentContext.messages]);
|
||||
|
||||
// Stream assistant response
|
||||
const message = await streamAssistantResponse(currentContext, config, signal, stream, streamFn);
|
||||
newMessages.push(message);
|
||||
|
||||
if (message.stopReason === "error" || message.stopReason === "aborted") {
|
||||
// Stop the loop on error or abort
|
||||
stream.push({ type: "turn_end", message, toolResults: [] });
|
||||
stream.push({ type: "agent_end", messages: newMessages });
|
||||
stream.end(newMessages);
|
||||
return;
|
||||
}
|
||||
|
||||
// Check for tool calls
|
||||
const toolCalls = message.content.filter((c) => c.type === "toolCall");
|
||||
hasMoreToolCalls = toolCalls.length > 0;
|
||||
|
||||
const toolResults: ToolResultMessage[] = [];
|
||||
if (hasMoreToolCalls) {
|
||||
// Execute tool calls
|
||||
toolResults.push(...(await executeToolCalls(currentContext.tools, message, signal, stream)));
|
||||
currentContext.messages.push(...toolResults);
|
||||
newMessages.push(...toolResults);
|
||||
}
|
||||
stream.push({ type: "turn_end", message, toolResults: toolResults });
|
||||
|
||||
// Get queued messages after turn completes
|
||||
queuedMessages = (await config.getQueuedMessages?.()) || [];
|
||||
}
|
||||
stream.push({ type: "agent_end", messages: newMessages });
|
||||
stream.end(newMessages);
|
||||
await runLoop(currentContext, newMessages, config, signal, stream, streamFn);
|
||||
})();
|
||||
|
||||
return stream;
|
||||
}
|
||||
|
||||
/**
|
||||
* Continue an agent loop from the current context without adding a new message.
|
||||
* Used for retry after overflow - context already has user message or tool results.
|
||||
* Throws if the last message is not a user message or tool result.
|
||||
*/
|
||||
export function agentLoopContinue(
|
||||
context: AgentContext,
|
||||
config: AgentLoopConfig,
|
||||
signal?: AbortSignal,
|
||||
streamFn?: typeof streamSimple,
|
||||
): EventStream<AgentEvent, AgentContext["messages"]> {
|
||||
// Validate that we can continue from this context
|
||||
const lastMessage = context.messages[context.messages.length - 1];
|
||||
if (!lastMessage) {
|
||||
throw new Error("Cannot continue: no messages in context");
|
||||
}
|
||||
if (lastMessage.role !== "user" && lastMessage.role !== "toolResult") {
|
||||
throw new Error(`Cannot continue from message role: ${lastMessage.role}. Expected 'user' or 'toolResult'.`);
|
||||
}
|
||||
|
||||
const stream = createAgentStream();
|
||||
|
||||
(async () => {
|
||||
const newMessages: AgentContext["messages"] = [];
|
||||
const currentContext: AgentContext = { ...context };
|
||||
|
||||
stream.push({ type: "agent_start" });
|
||||
stream.push({ type: "turn_start" });
|
||||
// No user message events - we're continuing from existing context
|
||||
|
||||
await runLoop(currentContext, newMessages, config, signal, stream, streamFn);
|
||||
})();
|
||||
|
||||
return stream;
|
||||
}
|
||||
|
||||
function createAgentStream(): EventStream<AgentEvent, AgentContext["messages"]> {
|
||||
return new EventStream<AgentEvent, AgentContext["messages"]>(
|
||||
(event: AgentEvent) => event.type === "agent_end",
|
||||
(event: AgentEvent) => (event.type === "agent_end" ? event.messages : []),
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Shared loop logic for both agentLoop and agentLoopContinue.
|
||||
*/
|
||||
async function runLoop(
|
||||
currentContext: AgentContext,
|
||||
newMessages: AgentContext["messages"],
|
||||
config: AgentLoopConfig,
|
||||
signal: AbortSignal | undefined,
|
||||
stream: EventStream<AgentEvent, AgentContext["messages"]>,
|
||||
streamFn?: typeof streamSimple,
|
||||
): Promise<void> {
|
||||
let hasMoreToolCalls = true;
|
||||
let firstTurn = true;
|
||||
let queuedMessages: QueuedMessage<any>[] = (await config.getQueuedMessages?.()) || [];
|
||||
|
||||
while (hasMoreToolCalls || queuedMessages.length > 0) {
|
||||
if (!firstTurn) {
|
||||
stream.push({ type: "turn_start" });
|
||||
} else {
|
||||
firstTurn = false;
|
||||
}
|
||||
|
||||
// Process queued messages first (inject before next assistant response)
|
||||
if (queuedMessages.length > 0) {
|
||||
for (const { original, llm } of queuedMessages) {
|
||||
stream.push({ type: "message_start", message: original });
|
||||
stream.push({ type: "message_end", message: original });
|
||||
if (llm) {
|
||||
currentContext.messages.push(llm);
|
||||
newMessages.push(llm);
|
||||
}
|
||||
}
|
||||
queuedMessages = [];
|
||||
}
|
||||
|
||||
// Stream assistant response
|
||||
const message = await streamAssistantResponse(currentContext, config, signal, stream, streamFn);
|
||||
newMessages.push(message);
|
||||
|
||||
if (message.stopReason === "error" || message.stopReason === "aborted") {
|
||||
// Stop the loop on error or abort
|
||||
stream.push({ type: "turn_end", message, toolResults: [] });
|
||||
stream.push({ type: "agent_end", messages: newMessages });
|
||||
stream.end(newMessages);
|
||||
return;
|
||||
}
|
||||
|
||||
// Check for tool calls
|
||||
const toolCalls = message.content.filter((c) => c.type === "toolCall");
|
||||
hasMoreToolCalls = toolCalls.length > 0;
|
||||
|
||||
const toolResults: ToolResultMessage[] = [];
|
||||
if (hasMoreToolCalls) {
|
||||
// Execute tool calls
|
||||
toolResults.push(...(await executeToolCalls(currentContext.tools, message, signal, stream)));
|
||||
currentContext.messages.push(...toolResults);
|
||||
newMessages.push(...toolResults);
|
||||
}
|
||||
stream.push({ type: "turn_end", message, toolResults: toolResults });
|
||||
|
||||
// Get queued messages after turn completes
|
||||
queuedMessages = (await config.getQueuedMessages?.()) || [];
|
||||
}
|
||||
|
||||
stream.push({ type: "agent_end", messages: newMessages });
|
||||
stream.end(newMessages);
|
||||
}
|
||||
|
||||
// Helper functions
|
||||
async function streamAssistantResponse(
|
||||
context: AgentContext,
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
export { agentLoop } from "./agent-loop.js";
|
||||
export { agentLoop, agentLoopContinue } from "./agent-loop.js";
|
||||
export * from "./tools/index.js";
|
||||
export type { AgentContext, AgentEvent, AgentLoopConfig, AgentTool, QueuedMessage } from "./types.js";
|
||||
|
|
|
|||
|
|
@ -3275,13 +3275,13 @@ export const MODELS = {
|
|||
reasoning: true,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 0.19999999999999998,
|
||||
output: 0.7999999999999999,
|
||||
input: 0.15,
|
||||
output: 0.75,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 163840,
|
||||
maxTokens: 163840,
|
||||
contextWindow: 8192,
|
||||
maxTokens: 7168,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"openai/gpt-4o-audio-preview": {
|
||||
id: "openai/gpt-4o-audio-preview",
|
||||
|
|
@ -4516,8 +4516,8 @@ export const MODELS = {
|
|||
reasoning: false,
|
||||
input: ["text", "image"],
|
||||
cost: {
|
||||
input: 0.049999999999999996,
|
||||
output: 0.22,
|
||||
input: 0.04,
|
||||
output: 0.15,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
|
|
|
|||
|
|
@ -1,9 +1,17 @@
|
|||
import { describe, expect, it } from "vitest";
|
||||
import { agentLoop } from "../src/agent/agent-loop.js";
|
||||
import { agentLoop, agentLoopContinue } from "../src/agent/agent-loop.js";
|
||||
import { calculateTool } from "../src/agent/tools/calculate.js";
|
||||
import type { AgentContext, AgentEvent, AgentLoopConfig } from "../src/agent/types.js";
|
||||
import { getModel } from "../src/models.js";
|
||||
import type { Api, Message, Model, OptionsForApi, UserMessage } from "../src/types.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
Message,
|
||||
Model,
|
||||
OptionsForApi,
|
||||
ToolResultMessage,
|
||||
UserMessage,
|
||||
} from "../src/types.js";
|
||||
|
||||
async function calculateTest<TApi extends Api>(model: Model<TApi>, options: OptionsForApi<TApi> = {}) {
|
||||
// Create the agent context with the calculator tool
|
||||
|
|
@ -282,7 +290,7 @@ describe("Agent Calculator Tests", () => {
|
|||
});
|
||||
|
||||
describe.skipIf(!process.env.ANTHROPIC_API_KEY)("Anthropic Provider Agent", () => {
|
||||
const model = getModel("anthropic", "claude-3-5-haiku-20241022");
|
||||
const model = getModel("anthropic", "claude-haiku-4-5");
|
||||
|
||||
it("should calculate multiple expressions and sum the results", async () => {
|
||||
const result = await calculateTest(model);
|
||||
|
|
@ -351,3 +359,175 @@ describe("Agent Calculator Tests", () => {
|
|||
}, 30000);
|
||||
});
|
||||
});
|
||||
|
||||
describe("agentLoopContinue", () => {
|
||||
describe("validation", () => {
|
||||
const model = getModel("anthropic", "claude-haiku-4-5");
|
||||
const baseContext: AgentContext = {
|
||||
systemPrompt: "You are a helpful assistant.",
|
||||
messages: [],
|
||||
tools: [],
|
||||
};
|
||||
const config: AgentLoopConfig = { model };
|
||||
|
||||
it("should throw when context has no messages", () => {
|
||||
expect(() => agentLoopContinue(baseContext, config)).toThrow("Cannot continue: no messages in context");
|
||||
});
|
||||
|
||||
it("should throw when last message is an assistant message", () => {
|
||||
const assistantMessage: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text: "Hello" }],
|
||||
api: "anthropic-messages",
|
||||
provider: "anthropic",
|
||||
model: "claude-haiku-4-5",
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
const context: AgentContext = {
|
||||
...baseContext,
|
||||
messages: [assistantMessage],
|
||||
};
|
||||
expect(() => agentLoopContinue(context, config)).toThrow(
|
||||
"Cannot continue from message role: assistant. Expected 'user' or 'toolResult'.",
|
||||
);
|
||||
});
|
||||
|
||||
// Note: "should not throw" tests for valid inputs are covered by the E2E tests below
|
||||
// which actually consume the stream and verify the output
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.ANTHROPIC_API_KEY)("continue from user message", () => {
|
||||
const model = getModel("anthropic", "claude-haiku-4-5");
|
||||
|
||||
it("should continue and get assistant response when last message is user", async () => {
|
||||
const userMessage: UserMessage = {
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "Say exactly: HELLO WORLD" }],
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
const context: AgentContext = {
|
||||
systemPrompt: "You are a helpful assistant. Follow instructions exactly.",
|
||||
messages: [userMessage],
|
||||
tools: [],
|
||||
};
|
||||
|
||||
const config: AgentLoopConfig = { model };
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
const stream = agentLoopContinue(context, config);
|
||||
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
}
|
||||
|
||||
const messages = await stream.result();
|
||||
|
||||
// Should have gotten an assistant response
|
||||
expect(messages.length).toBe(1);
|
||||
expect(messages[0].role).toBe("assistant");
|
||||
|
||||
// Verify event sequence - no user message events since we're continuing
|
||||
const eventTypes = events.map((e) => e.type);
|
||||
expect(eventTypes).toContain("agent_start");
|
||||
expect(eventTypes).toContain("turn_start");
|
||||
expect(eventTypes).toContain("message_start");
|
||||
expect(eventTypes).toContain("message_end");
|
||||
expect(eventTypes).toContain("turn_end");
|
||||
expect(eventTypes).toContain("agent_end");
|
||||
|
||||
// Should NOT have user message events (that's the difference from agentLoop)
|
||||
const messageEndEvents = events.filter((e) => e.type === "message_end");
|
||||
expect(messageEndEvents.length).toBe(1); // Only assistant message
|
||||
expect((messageEndEvents[0] as any).message.role).toBe("assistant");
|
||||
}, 30000);
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.ANTHROPIC_API_KEY)("continue from tool result", () => {
|
||||
const model = getModel("anthropic", "claude-haiku-4-5");
|
||||
|
||||
it("should continue processing after tool results", async () => {
|
||||
// Simulate a conversation where:
|
||||
// 1. User asked to calculate something
|
||||
// 2. Assistant made a tool call
|
||||
// 3. Tool result is ready
|
||||
// 4. We continue from here
|
||||
|
||||
const userMessage: UserMessage = {
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "What is 5 + 3? Use the calculator." }],
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
const assistantMessage: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [
|
||||
{ type: "text", text: "Let me calculate that for you." },
|
||||
{ type: "toolCall", id: "calc-1", name: "calculate", arguments: { expression: "5 + 3" } },
|
||||
],
|
||||
api: "anthropic-messages",
|
||||
provider: "anthropic",
|
||||
model: "claude-haiku-4-5",
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "toolUse",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
const toolResult: ToolResultMessage = {
|
||||
role: "toolResult",
|
||||
toolCallId: "calc-1",
|
||||
toolName: "calculate",
|
||||
content: [{ type: "text", text: "5 + 3 = 8" }],
|
||||
isError: false,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
const context: AgentContext = {
|
||||
systemPrompt: "You are a helpful assistant. After getting a calculation result, state the answer clearly.",
|
||||
messages: [userMessage, assistantMessage, toolResult],
|
||||
tools: [calculateTool],
|
||||
};
|
||||
|
||||
const config: AgentLoopConfig = { model };
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
const stream = agentLoopContinue(context, config);
|
||||
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
}
|
||||
|
||||
const messages = await stream.result();
|
||||
|
||||
// Should have gotten an assistant response
|
||||
expect(messages.length).toBeGreaterThanOrEqual(1);
|
||||
const lastMessage = messages[messages.length - 1];
|
||||
expect(lastMessage.role).toBe("assistant");
|
||||
|
||||
// The assistant should mention the result (8)
|
||||
if (lastMessage.role === "assistant") {
|
||||
const textContent = lastMessage.content
|
||||
.filter((c) => c.type === "text")
|
||||
.map((c) => (c as any).text)
|
||||
.join(" ");
|
||||
expect(textContent).toMatch(/8/);
|
||||
}
|
||||
}, 30000);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue