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

feat(agent): merge main into agent-harness-tools

Browse Source
This commit is contained in:
Mario Zechner
2026-07-22 11:56:35 +02:00
parent 37eb243d26 bc41f612da
commit 2d0b2294cd
93 changed files with 2398 additions and 241 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/agent/src/agent-loop.ts
+9 -8
View File
@@ -10,6 +10,7 @@ import {
type ToolResultMessage,
validateToolArguments,
} from "@earendil-works/pi-ai";
import { getDefaultStreamFn } from "./stream-fn.ts";
import type {
AgentContext,
AgentEvent,
@@ -32,7 +33,7 @@ export function agentLoop(
context: AgentContext,
config: AgentLoopConfig,
signal: AbortSignal | undefined,
streamFunction: StreamFn,
streamFn: StreamFn,
): EventStream<AgentEvent, AgentMessage[]> {
const stream = createAgentStream();
@@ -44,7 +45,7 @@ export function agentLoop(
stream.push(event);
},
signal,
streamFunction,
streamFn,
).then((messages) => {
stream.end(messages);
});
@@ -64,7 +65,7 @@ export function agentLoopContinue(
context: AgentContext,
config: AgentLoopConfig,
signal: AbortSignal | undefined,
streamFunction: StreamFn,
streamFn: StreamFn,
): EventStream<AgentEvent, AgentMessage[]> {
if (context.messages.length === 0) {
throw new Error("Cannot continue: no messages in context");
@@ -83,7 +84,7 @@ export function agentLoopContinue(
stream.push(event);
},
signal,
streamFunction,
streamFn,
).then((messages) => {
stream.end(messages);
});
@@ -97,7 +98,7 @@ export async function runAgentLoop(
config: AgentLoopConfig,
emit: AgentEventSink,
signal: AbortSignal | undefined,
streamFunction: StreamFn,
streamFn: StreamFn,
): Promise<AgentMessage[]> {
const newMessages: AgentMessage[] = [...prompts];
const currentContext: AgentContext = {
@@ -112,7 +113,7 @@ export async function runAgentLoop(
await emit({ type: "message_end", message: prompt });
}
await runLoop(currentContext, newMessages, config, signal, emit, streamFunction);
await runLoop(currentContext, newMessages, config, signal, emit, streamFn ?? getDefaultStreamFn());
return newMessages;
}
@@ -121,7 +122,7 @@ export async function runAgentLoopContinue(
config: AgentLoopConfig,
emit: AgentEventSink,
signal: AbortSignal | undefined,
streamFunction: StreamFn,
streamFn: StreamFn,
): Promise<AgentMessage[]> {
if (context.messages.length === 0) {
throw new Error("Cannot continue: no messages in context");
@@ -137,7 +138,7 @@ export async function runAgentLoopContinue(
await emit({ type: "agent_start" });
await emit({ type: "turn_start" });
await runLoop(currentContext, newMessages, config, signal, emit, streamFunction);
await runLoop(currentContext, newMessages, config, signal, emit, streamFn ?? getDefaultStreamFn());
return newMessages;
}
packages/agent/src/agent.ts
+22 -19
View File
@@ -8,6 +8,7 @@ import type {
Transport,
} from "@earendil-works/pi-ai";
import { runAgentLoop, runAgentLoopContinue } from "./agent-loop.ts";
import { getDefaultStreamFn } from "./stream-fn.ts";
import type {
AfterToolCallContext,
AfterToolCallResult,
@@ -97,7 +98,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[]>;
streamFunction: StreamFn;
streamFn: StreamFn;
getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
onPayload?: SimpleStreamOptions["onPayload"];
onResponse?: SimpleStreamOptions["onResponse"];
@@ -207,24 +208,26 @@ export class Agent {
public toolExecution: ToolExecutionMode;
constructor(options: AgentOptions) {
this._state = createMutableAgentState(options.initialState);
this.convertToLlm = options.convertToLlm ?? defaultConvertToLlm;
this.transformContext = options.transformContext;
this.streamFunction = options.streamFunction;
this.getApiKey = options.getApiKey;
this.onPayload = options.onPayload;
this.onResponse = options.onResponse;
this.beforeToolCall = options.beforeToolCall;
this.afterToolCall = options.afterToolCall;
this.prepareNextTurn = options.prepareNextTurn;
this.prepareNextTurnWithContext = options.prepareNextTurnWithContext;
this.steeringQueue = new PendingMessageQueue(options.steeringMode ?? "one-at-a-time");
this.followUpQueue = new PendingMessageQueue(options.followUpMode ?? "one-at-a-time");
this.sessionId = options.sessionId;
this.thinkingBudgets = options.thinkingBudgets;
this.transport = options.transport ?? "auto";
this.maxRetryDelayMs = options.maxRetryDelayMs;
this.toolExecution = options.toolExecution ?? "parallel";
// Older compiled consumers may omit options or streamFn even though the current API requires them.
const runtimeOptions: Partial<AgentOptions> = options ?? {};
this._state = createMutableAgentState(runtimeOptions.initialState);
this.convertToLlm = runtimeOptions.convertToLlm ?? defaultConvertToLlm;
this.transformContext = runtimeOptions.transformContext;
this.streamFunction = runtimeOptions.streamFn ?? getDefaultStreamFn();
this.getApiKey = runtimeOptions.getApiKey;
this.onPayload = runtimeOptions.onPayload;
this.onResponse = runtimeOptions.onResponse;
this.beforeToolCall = runtimeOptions.beforeToolCall;
this.afterToolCall = runtimeOptions.afterToolCall;
this.prepareNextTurn = runtimeOptions.prepareNextTurn;
this.prepareNextTurnWithContext = runtimeOptions.prepareNextTurnWithContext;
this.steeringQueue = new PendingMessageQueue(runtimeOptions.steeringMode ?? "one-at-a-time");
this.followUpQueue = new PendingMessageQueue(runtimeOptions.followUpMode ?? "one-at-a-time");
this.sessionId = runtimeOptions.sessionId;
this.thinkingBudgets = runtimeOptions.thinkingBudgets;
this.transport = runtimeOptions.transport ?? "auto";
this.maxRetryDelayMs = runtimeOptions.maxRetryDelayMs;
this.toolExecution = runtimeOptions.toolExecution ?? "parallel";
}
/**
packages/agent/src/harness/agent-harness.ts
+25 -1
View File
@@ -4,6 +4,8 @@ import {
type ImageContent,
type Model,
type Models,
type RetryCallbacks,
type RetryPolicy,
type UserMessage,
} from "@earendil-works/pi-ai";
import { runAgentLoop } from "../agent-loop.ts";
@@ -182,6 +184,7 @@ export class AgentHarness<
private systemPrompt: AgentHarnessOptions<TContext, TSkill, TPromptTemplate, TTool>["systemPrompt"];
private toolContext: AgentHarnessToolContextSource<TContext> | undefined;
private streamOptions: AgentHarnessStreamOptions;
private retry: RetryPolicy | undefined;
private resources: AgentHarnessResources<TSkill, TPromptTemplate>;
private tools = new Map<string, TTool>();
private activeToolNames: string[];
@@ -197,6 +200,7 @@ export class AgentHarness<
this.models = options.models;
this.resources = options.resources ?? {};
this.streamOptions = cloneStreamOptions(options.streamOptions);
this.retry = options.retry;
this.systemPrompt = options.systemPrompt;
this.toolContext = options.toolContext;
this.validateUniqueNames(
@@ -260,6 +264,15 @@ export class AgentHarness<
return lastResult;
}
private retryCallbacks(operation: "compaction" | "branch_summary"): RetryCallbacks {
return {
onRetryScheduled: (attempt, maxAttempts, delayMs, errorMessage) =>
this.emitOwn({ type: "retry_scheduled", operation, attempt, maxAttempts, delayMs, errorMessage }),
onRetryAttemptStart: () => this.emitOwn({ type: "retry_attempt_start", operation }),
onRetryFinished: () => this.emitOwn({ type: "retry_finished", operation }),
};
}
private async emitBeforeProviderRequest(
model: Model<any>,
sessionId: string,
@@ -741,7 +754,16 @@ export class AgentHarness<
const provided = hookResult?.compaction;
const compactResult = provided
? { ok: true as const, value: provided }
: await compact(preparation, this.models, model, customInstructions, undefined, this.thinkingLevel);
: await compact(
preparation,
this.models,
model,
customInstructions,
undefined,
this.thinkingLevel,
this.retry,
this.retryCallbacks("compaction"),
);
if (!compactResult.ok) throw compactResult.error;
const result = compactResult.value;
const entryId = await this.session.appendCompaction(
@@ -803,6 +825,8 @@ export class AgentHarness<
signal: new AbortController().signal,
customInstructions: hookResult?.customInstructions ?? options?.customInstructions,
replaceInstructions: hookResult?.replaceInstructions ?? options?.replaceInstructions,
retry: this.retry,
callbacks: this.retryCallbacks("branch_summary"),
});
if (!branchSummary.ok) {
if (branchSummary.error.code === "aborted") return { cancelled: true };
packages/agent/src/harness/compaction/branch-summarization.ts
+20 -4
View File
@@ -1,4 +1,4 @@
import { contentText, type Model, type Models } from "@earendil-works/pi-ai";
import { contentText, type Model, type Models, type RetryCallbacks, type RetryPolicy } from "@earendil-works/pi-ai";
import type { AgentMessage } from "../../types.ts";
import {
@@ -9,7 +9,7 @@ import {
} from "../messages.ts";
import type { BranchSummaryResult, Session, SessionTreeEntry } from "../types.ts";
import { BranchSummaryError, err, ok, type Result, SessionError } from "../types.ts";
import { estimateTokens, SUMMARIZATION_SYSTEM_PROMPT } from "./compaction.ts";
import { completeSimpleWithRetries, estimateTokens, SUMMARIZATION_SYSTEM_PROMPT } from "./compaction.ts";
import {
computeFileLists,
createFileOps,
@@ -61,6 +61,10 @@ export interface GenerateBranchSummaryOptions {
replaceInstructions?: boolean;
/** Tokens reserved for prompt and model output. Defaults to 16384. */
reserveTokens?: number;
/** Optional retry policy for transient summarization errors. */
retry?: RetryPolicy;
/** Optional callbacks for retry reporting. */
callbacks?: RetryCallbacks;
}
/** Collect entries that should be summarized before navigating to a different session tree entry. */
@@ -200,7 +204,16 @@ export async function generateBranchSummary(
entries: SessionTreeEntry[],
options: GenerateBranchSummaryOptions,
): Promise<Result<BranchSummaryResult, BranchSummaryError>> {
const { models, model, signal, customInstructions, replaceInstructions, reserveTokens = 16384 } = options;
const {
models,
model,
signal,
customInstructions,
replaceInstructions,
reserveTokens = 16384,
retry,
callbacks,
} = options;
const contextWindow = model.contextWindow || 128000;
const tokenBudget = contextWindow - reserveTokens;
@@ -228,10 +241,13 @@ export async function generateBranchSummary(
timestamp: Date.now(),
},
];
const response = await models.completeSimple(
const response = await completeSimpleWithRetries(
models,
model,
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
{ signal, maxTokens: 2048 },
retry,
callbacks,
);
if (response.stopReason === "aborted") {
return err(new BranchSummaryError("aborted", response.errorMessage || "Branch summary aborted"));
packages/agent/src/harness/compaction/compaction.ts
+45 -5
View File
@@ -1,9 +1,14 @@
import {
type AssistantMessage,
type Context,
contentText,
type ImageContent,
type Model,
type Models,
type RetryCallbacks,
type RetryPolicy,
retryAssistantCall,
type SimpleStreamOptions,
type TextContent,
type Usage,
} from "@earendil-works/pi-ai";
@@ -109,6 +114,17 @@ export interface CompactionResult<T = unknown> {
details?: T;
}
export async function completeSimpleWithRetries(
models: Models,
model: Model<any>,
context: Context,
options: SimpleStreamOptions,
retry?: RetryPolicy,
callbacks?: RetryCallbacks,
): Promise<AssistantMessage> {
return retryAssistantCall(() => models.completeSimple(model, context, options), retry, options.signal, callbacks);
}
function combineUsage(first: Usage, second: Usage): Usage {
return {
input: first.input + second.input,
@@ -501,6 +517,8 @@ export async function generateSummary(
customInstructions?: string,
previousSummary?: string,
thinkingLevel?: ThinkingLevel,
retry?: RetryPolicy,
callbacks?: RetryCallbacks,
): Promise<Result<string, CompactionError>> {
const result = await generateSummaryWithUsage(
currentMessages,
@@ -511,6 +529,8 @@ export async function generateSummary(
customInstructions,
previousSummary,
thinkingLevel,
retry,
callbacks,
);
return result.ok ? ok(result.value.text) : err(result.error);
}
@@ -525,6 +545,8 @@ export async function generateSummaryWithUsage(
customInstructions?: string,
previousSummary?: string,
thinkingLevel?: ThinkingLevel,
retry?: RetryPolicy,
callbacks?: RetryCallbacks,
): Promise<Result<{ text: string; usage: Usage }, CompactionError>> {
const maxTokens = Math.min(
Math.floor(0.8 * reserveTokens),
@@ -555,10 +577,13 @@ export async function generateSummaryWithUsage(
? { maxTokens, signal, reasoning: thinkingLevel }
: { maxTokens, signal };
const response = await models.completeSimple(
const response = await completeSimpleWithRetries(
models,
model,
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
completionOptions,
retry,
callbacks,
);
if (response.stopReason === "aborted") {
return err(new CompactionError("aborted", response.errorMessage || "Summarization aborted"));
@@ -700,6 +725,8 @@ export async function compact(
customInstructions?: string,
signal?: AbortSignal,
thinkingLevel?: ThinkingLevel,
retry?: RetryPolicy,
callbacks?: RetryCallbacks,
): Promise<Result<CompactionResult, CompactionError>> {
const {
firstKeptEntryId,
@@ -733,6 +760,8 @@ export async function compact(
customInstructions,
previousSummary,
thinkingLevel,
retry,
callbacks,
);
if (!historyResult.ok) return err(historyResult.error);
historyText = historyResult.value.text;
@@ -745,6 +774,8 @@ export async function compact(
settings.reserveTokens,
signal,
thinkingLevel,
retry,
callbacks,
);
if (!turnPrefixResult.ok) return err(turnPrefixResult.error);
summary = `${historyText}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.value.text}`;
@@ -761,6 +792,8 @@ export async function compact(
customInstructions,
previousSummary,
thinkingLevel,
retry,
callbacks,
);
if (!summaryResult.ok) return err(summaryResult.error);
summary = summaryResult.value.text;
@@ -786,6 +819,8 @@ async function generateTurnPrefixSummary(
reserveTokens: number,
signal?: AbortSignal,
thinkingLevel?: ThinkingLevel,
retry?: RetryPolicy,
callbacks?: RetryCallbacks,
): Promise<Result<{ text: string; usage: Usage }, CompactionError>> {
const maxTokens = Math.min(
Math.floor(0.5 * reserveTokens),
@@ -802,12 +837,17 @@ async function generateTurnPrefixSummary(
},
];
const response = await models.completeSimple(
model,
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
const completionOptions =
model.reasoning && thinkingLevel && thinkingLevel !== "off"
? { maxTokens, signal, reasoning: thinkingLevel }
: { maxTokens, signal },
: { maxTokens, signal };
const response = await completeSimpleWithRetries(
models,
model,
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
completionOptions,
retry,
callbacks,
);
if (response.stopReason === "aborted") {
return err(new CompactionError("aborted", response.errorMessage || "Turn prefix summarization aborted"));
packages/agent/src/harness/types.ts
+28
View File
@@ -2,6 +2,7 @@ import type {
ImageContent,
Model,
Models,
RetryPolicy,
SimpleStreamOptions,
TextContent,
Transport,
@@ -659,6 +660,25 @@ export interface SessionTreeEvent {
fromHook?: boolean;
}
export interface RetryScheduledEvent {
type: "retry_scheduled";
operation: "compaction" | "branch_summary";
attempt: number;
maxAttempts: number;
delayMs: number;
errorMessage: string;
}
export interface RetryAttemptStartEvent {
type: "retry_attempt_start";
operation: "compaction" | "branch_summary";
}
export interface RetryFinishedEvent {
type: "retry_finished";
operation: "compaction" | "branch_summary";
}
export interface ModelUpdateEvent {
type: "model_update";
model: Model<any>;
@@ -709,6 +729,9 @@ export type AgentHarnessOwnEvent<
| SessionCompactEvent
| SessionBeforeTreeEvent
| SessionTreeEvent
| RetryScheduledEvent
| RetryAttemptStartEvent
| RetryFinishedEvent
| ModelUpdateEvent
| ThinkingLevelUpdateEvent
| ResourcesUpdateEvent<TSkill, TPromptTemplate>
@@ -778,6 +801,9 @@ export type AgentHarnessEventResultMap = {
session_compact: undefined;
session_before_tree: SessionBeforeTreeResult | undefined;
session_tree: undefined;
retry_scheduled: undefined;
retry_attempt_start: undefined;
retry_finished: undefined;
model_update: undefined;
thinking_level_update: undefined;
resources_update: undefined;
@@ -897,6 +923,8 @@ export interface AgentHarnessOptions<
}) => string | Promise<string>);
/** Curated stream/provider request options. Snapshotted at turn start. */
streamOptions?: AgentHarnessStreamOptions;
/** Optional retry policy for generated compaction and branch-summary requests. */
retry?: RetryPolicy;
model: Model<any>;
thinkingLevel?: ThinkingLevel;
activeToolNames?: string[];
packages/agent/src/index.ts
+2
View File
@@ -44,5 +44,7 @@ export * from "./harness/utils/shell-output.ts";
export * from "./harness/utils/truncate.ts";
// Proxy utilities
export * from "./proxy.ts";
// Stream defaults
export { setDefaultStreamFn } from "./stream-fn.ts";
// Types
export * from "./types.ts";
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 `streamFunction` option when creating an Agent that needs to go through a proxy.
* Use this as the `streamFn` option when creating an Agent that needs to go through a proxy.
*
* @example
* ```typescript
* const agent = new Agent({
* streamFunction: (model, context, options) =>
* streamFn: (model, context, options) =>
* streamProxy(model, context, {
* ...options,
* authToken: await getAuthToken(),
packages/agent/src/stream-fn.ts
+20
View File
@@ -0,0 +1,20 @@
import type { StreamFn } from "./types.ts";
let defaultStreamFn: StreamFn | undefined;
/**
* Configure the fallback used by Agent and low-level loops when callers omit streamFn.
*
* Hosts that provide a default model runtime can install its stream function here
* without making pi-agent-core depend on a provider catalog or compatibility layer.
*/
export function setDefaultStreamFn(streamFn: StreamFn | undefined): void {
defaultStreamFn = streamFn;
}
export function getDefaultStreamFn(): StreamFn {
if (!defaultStreamFn) {
throw new Error("No default stream function configured. Pass streamFn explicitly or call setDefaultStreamFn().");
}
return defaultStreamFn;
}