ton/pi_harness
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(agent): decouple agent streams from compat

Browse Source 
closes #6851
This commit is contained in:
Mario Zechner
2026-07-20 17:54:17 +02:00
parent 3a40794ea1
commit 1235c0ec64
29 changed files with 219 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/agent/src/agent-loop.ts
+16 -19
View File
@@ -7,10 +7,9 @@ import {
type AssistantMessage,
type Context,
EventStream,
streamSimple,
type ToolResultMessage,
validateToolArguments,
} from "@earendil-works/pi-ai/compat";
} from "@earendil-works/pi-ai";
import type {
AgentContext,
AgentEvent,
@@ -32,8 +31,8 @@ export function agentLoop(
prompts: AgentMessage[],
context: AgentContext,
config: AgentLoopConfig,
signal?: AbortSignal,
streamFn?: StreamFn,
signal: AbortSignal | undefined,
streamFunction: StreamFn,
): EventStream<AgentEvent, AgentMessage[]> {
const stream = createAgentStream();
@@ -45,7 +44,7 @@ export function agentLoop(
stream.push(event);
},
signal,
streamFn,
streamFunction,
).then((messages) => {
stream.end(messages);
});
@@ -64,8 +63,8 @@ export function agentLoop(
export function agentLoopContinue(
context: AgentContext,
config: AgentLoopConfig,
signal?: AbortSignal,
streamFn?: StreamFn,
signal: AbortSignal | undefined,
streamFunction: StreamFn,
): EventStream<AgentEvent, AgentMessage[]> {
if (context.messages.length === 0) {
throw new Error("Cannot continue: no messages in context");
@@ -84,7 +83,7 @@ export function agentLoopContinue(
stream.push(event);
},
signal,
streamFn,
streamFunction,
).then((messages) => {
stream.end(messages);
});
@@ -97,8 +96,8 @@ export async function runAgentLoop(
context: AgentContext,
config: AgentLoopConfig,
emit: AgentEventSink,
signal?: AbortSignal,
streamFn?: StreamFn,
signal: AbortSignal | undefined,
streamFunction: StreamFn,
): Promise<AgentMessage[]> {
const newMessages: AgentMessage[] = [...prompts];
const currentContext: AgentContext = {
@@ -113,7 +112,7 @@ export async function runAgentLoop(
await emit({ type: "message_end", message: prompt });
}
await runLoop(currentContext, newMessages, config, signal, emit, streamFn);
await runLoop(currentContext, newMessages, config, signal, emit, streamFunction);
return newMessages;
}
@@ -121,8 +120,8 @@ export async function runAgentLoopContinue(
context: AgentContext,
config: AgentLoopConfig,
emit: AgentEventSink,
signal?: AbortSignal,
streamFn?: StreamFn,
signal: AbortSignal | undefined,
streamFunction: StreamFn,
): Promise<AgentMessage[]> {
if (context.messages.length === 0) {
throw new Error("Cannot continue: no messages in context");
@@ -138,7 +137,7 @@ export async function runAgentLoopContinue(
await emit({ type: "agent_start" });
await emit({ type: "turn_start" });
await runLoop(currentContext, newMessages, config, signal, emit, streamFn);
await runLoop(currentContext, newMessages, config, signal, emit, streamFunction);
return newMessages;
}
@@ -158,7 +157,7 @@ async function runLoop(
initialConfig: AgentLoopConfig,
signal: AbortSignal | undefined,
emit: AgentEventSink,
streamFn?: StreamFn,
streamFunction: StreamFn,
): Promise<void> {
let currentContext = initialContext;
let config = initialConfig;
@@ -190,7 +189,7 @@ async function runLoop(
}
// Stream assistant response
const message = await streamAssistantResponse(currentContext, config, signal, emit, streamFn);
const message = await streamAssistantResponse(currentContext, config, signal, emit, streamFunction);
newMessages.push(message);
if (message.stopReason === "error" || message.stopReason === "aborted") {
@@ -283,7 +282,7 @@ async function streamAssistantResponse(
config: AgentLoopConfig,
signal: AbortSignal | undefined,
emit: AgentEventSink,
streamFn?: StreamFn,
streamFunction: StreamFn,
): Promise<AssistantMessage> {
// Apply context transform if configured (AgentMessage[] → AgentMessage[])
let messages = context.messages;
@@ -301,8 +300,6 @@ async function streamAssistantResponse(
tools: context.tools,
};
const streamFunction = streamFn || streamSimple;
// Resolve API key (important for expiring tokens)
const resolvedApiKey =
(config.getApiKey ? await config.getApiKey(config.model.provider) : undefined) || config.apiKey;
packages/agent/src/agent.ts
+15 -16
View File
@@ -1,13 +1,12 @@
import {
type ImageContent,
type Message,
type Model,
type SimpleStreamOptions,
streamSimple,
type TextContent,
type ThinkingBudgets,
type Transport,
} from "@earendil-works/pi-ai/compat";
import type {
ImageContent,
Message,
Model,
SimpleStreamOptions,
TextContent,
ThinkingBudgets,
Transport,
} from "@earendil-works/pi-ai";
import { runAgentLoop, runAgentLoopContinue } from "./agent-loop.ts";
import type {
AfterToolCallContext,
@@ -98,7 +97,7 @@ export interface AgentOptions {
initialState?: Partial<Omit<AgentState, "pendingToolCalls" | "isStreaming" | "streamingMessage" | "errorMessage">>;
convertToLlm?: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
streamFn?: StreamFn;
streamFunction: StreamFn;
getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
onPayload?: SimpleStreamOptions["onPayload"];
onResponse?: SimpleStreamOptions["onResponse"];
@@ -176,7 +175,7 @@ export class Agent {
public convertToLlm: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
public transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
public streamFn: StreamFn;
public streamFunction: StreamFn;
public getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
public onPayload?: SimpleStreamOptions["onPayload"];
public onResponse?: SimpleStreamOptions["onResponse"];
@@ -207,11 +206,11 @@ export class Agent {
/** Tool execution strategy for assistant messages that contain multiple tool calls. */
public toolExecution: ToolExecutionMode;
constructor(options: AgentOptions = {}) {
constructor(options: AgentOptions) {
this._state = createMutableAgentState(options.initialState);
this.convertToLlm = options.convertToLlm ?? defaultConvertToLlm;
this.transformContext = options.transformContext;
this.streamFn = options.streamFn ?? streamSimple;
this.streamFunction = options.streamFunction;
this.getApiKey = options.getApiKey;
this.onPayload = options.onPayload;
this.onResponse = options.onResponse;
@@ -404,7 +403,7 @@ export class Agent {
this.createLoopConfig(options),
(event) => this.processEvents(event),
signal,
this.streamFn,
this.streamFunction,
);
});
}
@@ -416,7 +415,7 @@ export class Agent {
this.createLoopConfig(),
(event) => this.processEvents(event),
signal,
this.streamFn,
this.streamFunction,
);
});
}
packages/agent/src/proxy.ts
+2 -2
View File
@@ -84,12 +84,12 @@ export interface ProxyStreamOptions extends ProxySerializableStreamOptions {
* The server strips the partial field from delta events to reduce bandwidth.
* We reconstruct the partial message client-side.
*
* Use this as the `streamFn` option when creating an Agent that needs to go through a proxy.
* Use this as the `streamFunction` option when creating an Agent that needs to go through a proxy.
*
* @example
* ```typescript
* const agent = new Agent({
* streamFn: (model, context, options) =>
* streamFunction: (model, context, options) =>
* streamProxy(model, context, {
* ...options,
* authToken: await getAuthToken(),