WIP: Refactor agent package - not compiling

- Renamed AppMessage to AgentMessage throughout
- New agent-loop.ts with AgentLoopContext, AgentLoopConfig
- Removed transport abstraction, Agent now takes streamFn directly
- Extracted streamProxy to proxy.ts utility
- Removed agent-loop from pi-ai (now in agent package)
- Updated consumers (coding-agent, mom) for AgentMessage rename
- Tests updated but some consumers still need migration

Known issues:
- AgentTool, AgentToolResult not exported from pi-ai
- Attachment not exported from pi-agent-core
- ProviderTransport removed but still referenced
- messageTransformer -> convertToLlm migration incomplete
- CustomMessages declaration merging not working properly
This commit is contained in:
Mario Zechner 2025-12-28 09:23:38 +01:00
parent f7ef44dc38
commit a055fd4481
32 changed files with 1312 additions and 2009 deletions

View file

@ -1,376 +0,0 @@
import { streamSimple } from "../stream.js";
import type { AssistantMessage, Context, Message, ToolResultMessage, UserMessage } from "../types.js";
import { EventStream } from "../utils/event-stream.js";
import { validateToolArguments } from "../utils/validation.js";
import type { AgentContext, AgentEvent, AgentLoopConfig, AgentTool, AgentToolResult, QueuedMessage } from "./types.js";
/**
* 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,
config: AgentLoopConfig,
signal?: AbortSignal,
streamFn?: typeof streamSimple,
): EventStream<AgentEvent, AgentContext["messages"]> {
const stream = createAgentStream();
(async () => {
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 });
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.
*/
/**
* 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?.()) || [];
let queuedAfterTools: QueuedMessage<any>[] | null = null;
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
const toolExecution = await executeToolCalls(
currentContext.tools,
message,
signal,
stream,
config.getQueuedMessages,
);
toolResults.push(...toolExecution.toolResults);
queuedAfterTools = toolExecution.queuedMessages ?? null;
currentContext.messages.push(...toolResults);
newMessages.push(...toolResults);
}
stream.push({ type: "turn_end", message, toolResults: toolResults });
// Get queued messages after turn completes
if (queuedAfterTools && queuedAfterTools.length > 0) {
queuedMessages = queuedAfterTools;
queuedAfterTools = null;
} else {
queuedMessages = (await config.getQueuedMessages?.()) || [];
}
}
stream.push({ type: "agent_end", messages: newMessages });
stream.end(newMessages);
}
// Helper functions
async function streamAssistantResponse(
context: AgentContext,
config: AgentLoopConfig,
signal: AbortSignal | undefined,
stream: EventStream<AgentEvent, AgentContext["messages"]>,
streamFn?: typeof streamSimple,
): Promise<AssistantMessage> {
// Convert AgentContext to Context for streamSimple
// Use a copy of messages to avoid mutating the original context
const processedMessages = config.preprocessor
? await config.preprocessor(context.messages, signal)
: [...context.messages];
const processedContext: Context = {
systemPrompt: context.systemPrompt,
messages: [...processedMessages].map((m) => {
if (m.role === "toolResult") {
// biome-ignore lint/correctness/noUnusedVariables: fine here
const { details, ...rest } = m;
return rest;
} else {
return m;
}
}),
tools: context.tools, // AgentTool extends Tool, so this works
};
// Use custom stream function if provided, otherwise use default streamSimple
const streamFunction = streamFn || streamSimple;
// Resolve API key for every assistant response (important for expiring tokens)
const resolvedApiKey =
(config.getApiKey ? await config.getApiKey(config.model.provider) : undefined) || config.apiKey;
const response = await streamFunction(config.model, processedContext, { ...config, apiKey: resolvedApiKey, signal });
let partialMessage: AssistantMessage | null = null;
let addedPartial = false;
for await (const event of response) {
switch (event.type) {
case "start":
partialMessage = event.partial;
context.messages.push(partialMessage);
addedPartial = true;
stream.push({ type: "message_start", message: { ...partialMessage } });
break;
case "text_start":
case "text_delta":
case "text_end":
case "thinking_start":
case "thinking_delta":
case "thinking_end":
case "toolcall_start":
case "toolcall_delta":
case "toolcall_end":
if (partialMessage) {
partialMessage = event.partial;
context.messages[context.messages.length - 1] = partialMessage;
stream.push({ type: "message_update", assistantMessageEvent: event, message: { ...partialMessage } });
}
break;
case "done":
case "error": {
const finalMessage = await response.result();
if (addedPartial) {
context.messages[context.messages.length - 1] = finalMessage;
} else {
context.messages.push(finalMessage);
}
if (!addedPartial) {
stream.push({ type: "message_start", message: { ...finalMessage } });
}
stream.push({ type: "message_end", message: finalMessage });
return finalMessage;
}
}
}
return await response.result();
}
async function executeToolCalls<T>(
tools: AgentTool<any, T>[] | undefined,
assistantMessage: AssistantMessage,
signal: AbortSignal | undefined,
stream: EventStream<AgentEvent, Message[]>,
getQueuedMessages?: AgentLoopConfig["getQueuedMessages"],
): Promise<{ toolResults: ToolResultMessage<T>[]; queuedMessages?: QueuedMessage<any>[] }> {
const toolCalls = assistantMessage.content.filter((c) => c.type === "toolCall");
const results: ToolResultMessage<any>[] = [];
let queuedMessages: QueuedMessage<any>[] | undefined;
for (let index = 0; index < toolCalls.length; index++) {
const toolCall = toolCalls[index];
const tool = tools?.find((t) => t.name === toolCall.name);
stream.push({
type: "tool_execution_start",
toolCallId: toolCall.id,
toolName: toolCall.name,
args: toolCall.arguments,
});
let result: AgentToolResult<T>;
let isError = false;
try {
if (!tool) throw new Error(`Tool ${toolCall.name} not found`);
// Validate arguments using shared validation function
const validatedArgs = validateToolArguments(tool, toolCall);
// Execute with validated, typed arguments, passing update callback
result = await tool.execute(toolCall.id, validatedArgs, signal, (partialResult) => {
stream.push({
type: "tool_execution_update",
toolCallId: toolCall.id,
toolName: toolCall.name,
args: toolCall.arguments,
partialResult,
});
});
} catch (e) {
result = {
content: [{ type: "text", text: e instanceof Error ? e.message : String(e) }],
details: {} as T,
};
isError = true;
}
stream.push({
type: "tool_execution_end",
toolCallId: toolCall.id,
toolName: toolCall.name,
result,
isError,
});
const toolResultMessage: ToolResultMessage<T> = {
role: "toolResult",
toolCallId: toolCall.id,
toolName: toolCall.name,
content: result.content,
details: result.details,
isError,
timestamp: Date.now(),
};
results.push(toolResultMessage);
stream.push({ type: "message_start", message: toolResultMessage });
stream.push({ type: "message_end", message: toolResultMessage });
if (getQueuedMessages) {
const queued = await getQueuedMessages();
if (queued.length > 0) {
queuedMessages = queued;
const remainingCalls = toolCalls.slice(index + 1);
for (const skipped of remainingCalls) {
results.push(skipToolCall(skipped, stream));
}
break;
}
}
}
return { toolResults: results, queuedMessages };
}
function skipToolCall<T>(
toolCall: Extract<AssistantMessage["content"][number], { type: "toolCall" }>,
stream: EventStream<AgentEvent, Message[]>,
): ToolResultMessage<T> {
const result: AgentToolResult<T> = {
content: [{ type: "text", text: "Skipped due to queued user message." }],
details: {} as T,
};
stream.push({
type: "tool_execution_start",
toolCallId: toolCall.id,
toolName: toolCall.name,
args: toolCall.arguments,
});
stream.push({
type: "tool_execution_end",
toolCallId: toolCall.id,
toolName: toolCall.name,
result,
isError: true,
});
const toolResultMessage: ToolResultMessage<T> = {
role: "toolResult",
toolCallId: toolCall.id,
toolName: toolCall.name,
content: result.content,
details: result.details,
isError: true,
timestamp: Date.now(),
};
stream.push({ type: "message_start", message: toolResultMessage });
stream.push({ type: "message_end", message: toolResultMessage });
return toolResultMessage;
}

View file

@ -1,11 +0,0 @@
export { agentLoop, agentLoopContinue } from "./agent-loop.js";
export * from "./tools/index.js";
export type {
AgentContext,
AgentEvent,
AgentLoopConfig,
AgentTool,
AgentToolResult,
AgentToolUpdateCallback,
QueuedMessage,
} from "./types.js";

View file

@ -1,32 +0,0 @@
import { type Static, Type } from "@sinclair/typebox";
import type { AgentTool, AgentToolResult } from "../../agent/types.js";
export interface CalculateResult extends AgentToolResult<undefined> {
content: Array<{ type: "text"; text: string }>;
details: undefined;
}
export function calculate(expression: string): CalculateResult {
try {
const result = new Function(`return ${expression}`)();
return { content: [{ type: "text", text: `${expression} = ${result}` }], details: undefined };
} catch (e: any) {
throw new Error(e.message || String(e));
}
}
const calculateSchema = Type.Object({
expression: Type.String({ description: "The mathematical expression to evaluate" }),
});
type CalculateParams = Static<typeof calculateSchema>;
export const calculateTool: AgentTool<typeof calculateSchema, undefined> = {
label: "Calculator",
name: "calculate",
description: "Evaluate mathematical expressions",
parameters: calculateSchema,
execute: async (_toolCallId: string, args: CalculateParams) => {
return calculate(args.expression);
},
};

View file

@ -1,47 +0,0 @@
import { type Static, Type } from "@sinclair/typebox";
import type { AgentTool } from "../../agent/index.js";
import type { AgentToolResult } from "../types.js";
export interface GetCurrentTimeResult extends AgentToolResult<{ utcTimestamp: number }> {}
export async function getCurrentTime(timezone?: string): Promise<GetCurrentTimeResult> {
const date = new Date();
if (timezone) {
try {
const timeStr = date.toLocaleString("en-US", {
timeZone: timezone,
dateStyle: "full",
timeStyle: "long",
});
return {
content: [{ type: "text", text: timeStr }],
details: { utcTimestamp: date.getTime() },
};
} catch (_e) {
throw new Error(`Invalid timezone: ${timezone}. Current UTC time: ${date.toISOString()}`);
}
}
const timeStr = date.toLocaleString("en-US", { dateStyle: "full", timeStyle: "long" });
return {
content: [{ type: "text", text: timeStr }],
details: { utcTimestamp: date.getTime() },
};
}
const getCurrentTimeSchema = Type.Object({
timezone: Type.Optional(
Type.String({ description: "Optional timezone (e.g., 'America/New_York', 'Europe/London')" }),
),
});
type GetCurrentTimeParams = Static<typeof getCurrentTimeSchema>;
export const getCurrentTimeTool: AgentTool<typeof getCurrentTimeSchema, { utcTimestamp: number }> = {
label: "Current Time",
name: "get_current_time",
description: "Get the current date and time",
parameters: getCurrentTimeSchema,
execute: async (_toolCallId: string, args: GetCurrentTimeParams) => {
return getCurrentTime(args.timezone);
},
};

View file

@ -1,2 +0,0 @@
export { calculate, calculateTool } from "./calculate.js";
export { getCurrentTime, getCurrentTimeTool } from "./get-current-time.js";

View file

@ -1,105 +0,0 @@
import type { Static, TSchema } from "@sinclair/typebox";
import type {
AssistantMessage,
AssistantMessageEvent,
ImageContent,
Message,
Model,
SimpleStreamOptions,
TextContent,
Tool,
ToolResultMessage,
} from "../types.js";
export interface AgentToolResult<T> {
// Content blocks supporting text and images
content: (TextContent | ImageContent)[];
// Details to be displayed in a UI or logged
details: T;
}
// Callback for streaming tool execution updates
export type AgentToolUpdateCallback<T = any> = (partialResult: AgentToolResult<T>) => void;
// AgentTool extends Tool but adds the execute function
export interface AgentTool<TParameters extends TSchema = TSchema, TDetails = any> extends Tool<TParameters> {
// A human-readable label for the tool to be displayed in UI
label: string;
execute: (
toolCallId: string,
params: Static<TParameters>,
signal?: AbortSignal,
onUpdate?: AgentToolUpdateCallback<TDetails>,
) => Promise<AgentToolResult<TDetails>>;
}
// AgentContext is like Context but uses AgentTool
export interface AgentContext {
systemPrompt: string;
messages: Message[];
tools?: AgentTool<any>[];
}
// Event types
export type AgentEvent =
// Emitted when the agent starts. An agent can emit multiple turns
| { type: "agent_start" }
// Emitted when a turn starts. A turn can emit an optional user message (initial prompt), an assistant message (response) and multiple tool result messages
| { type: "turn_start" }
// Emitted when a user, assistant or tool result message starts
| { type: "message_start"; message: Message }
// Emitted when an asssitant messages is updated due to streaming
| { type: "message_update"; assistantMessageEvent: AssistantMessageEvent; message: AssistantMessage }
// Emitted when a user, assistant or tool result message is complete
| { type: "message_end"; message: Message }
// Emitted when a tool execution starts
| { type: "tool_execution_start"; toolCallId: string; toolName: string; args: any }
// Emitted when a tool execution produces output (streaming)
| {
type: "tool_execution_update";
toolCallId: string;
toolName: string;
args: any;
partialResult: AgentToolResult<any>;
}
// Emitted when a tool execution completes
| {
type: "tool_execution_end";
toolCallId: string;
toolName: string;
result: AgentToolResult<any>;
isError: boolean;
}
// Emitted when a full turn completes
| { type: "turn_end"; message: AssistantMessage; toolResults: ToolResultMessage[] }
// Emitted when the agent has completed all its turns. All messages from every turn are
// contained in messages, which can be appended to the context
| { type: "agent_end"; messages: AgentContext["messages"] };
// Queued message with optional LLM representation
export interface QueuedMessage<TApp = Message> {
original: TApp; // Original message for UI events
llm?: Message; // Optional transformed message for loop context (undefined if filtered)
}
// Configuration for agent loop execution
export interface AgentLoopConfig extends SimpleStreamOptions {
model: Model<any>;
/**
* Optional hook to resolve an API key dynamically for each LLM call.
*
* This is useful for short-lived OAuth tokens (e.g. GitHub Copilot) that may
* expire during long-running tool execution phases.
*
* The agent loop will call this before each assistant response and pass the
* returned value as `apiKey` to `streamSimple()` (or a custom `streamFn`).
*
* If it returns `undefined`, the loop falls back to `config.apiKey`, and then
* to `streamSimple()`'s own provider key lookup (setApiKey/env vars).
*/
getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
preprocessor?: (messages: AgentContext["messages"], abortSignal?: AbortSignal) => Promise<AgentContext["messages"]>;
getQueuedMessages?: <T>() => Promise<QueuedMessage<T>[]>;
}

View file

@ -1,4 +1,3 @@
export * from "./agent/index.js";
export * from "./models.js";
export * from "./providers/anthropic.js";
export * from "./providers/google.js";
@ -7,6 +6,7 @@ export * from "./providers/openai-completions.js";
export * from "./providers/openai-responses.js";
export * from "./stream.js";
export * from "./types.js";
export * from "./utils/event-stream.js";
export * from "./utils/oauth/index.js";
export * from "./utils/overflow.js";
export * from "./utils/typebox-helpers.js";