alexion/dev
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

chore(vendor): import pi 1.3.0

Browse Source 
Imported as a linear vendor commit so forge rebases do not replay the upstream root tree.

git-subtree-dir: third_party/pi

git-subtree-split: 71023aec1546d6d677a65669c5665c0927ce0dd0
This commit is contained in:
alexion
2026-10-15 17:58:48 -04:00
parent 719370c8f7
commit 6081e3e325
1743 changed files with 205460 additions and 131705 deletions
Download Patch File Download Diff File Expand all files Collapse all files

246
third_party/pi/packages/agent/src/agent-loop.ts vendored
View File

@@ -5,9 +5,14 @@
import {
type AssistantMessage,
type Context,
EventStream,
getCurrentTools,
getToolStateChanges,
normalizeContext,
type SystemMessage,
type ToolResultMessage,
type ToolStateChanges,
toToolDeclaration,
validateToolArguments,
} from "@earendil-works/pi-ai";
import { getDefaultStreamFn } from "./stream-fn.ts";
@@ -18,6 +23,7 @@ import type {
AgentMessage,
AgentTool,
AgentToolCall,
AgentToolCallOutcome,
AgentToolResult,
PrepareNextTurnContext,
StreamFn,
@@ -101,17 +107,18 @@ export async function runAgentLoop(
signal: AbortSignal | undefined,
streamFn: StreamFn,
): Promise<AgentMessage[]> {
const newMessages: AgentMessage[] = [...prompts];
const initialMessages = declareToolChanges(context, prompts);
const newMessages: AgentMessage[] = [...initialMessages];
const currentContext: AgentContext = {
...context,
messages: [...context.messages, ...prompts],
messages: [...context.messages, ...initialMessages],
};
await emit({ type: "agent_start" });
await emit({ type: "turn_start" });
for (const prompt of prompts) {
await emit({ type: "message_start", message: prompt });
await emit({ type: "message_end", message: prompt });
for (const message of initialMessages) {
await emit({ type: "message_start", message });
await emit({ type: "message_end", message });
}
await runLoop(currentContext, newMessages, config, signal, emit, streamFn ?? getDefaultStreamFn());
@@ -164,6 +171,7 @@ async function runLoop(
let currentContext = initialContext;
let config = initialConfig;
let lastCompletedTurn: PrepareNextTurnContext | undefined;
let explicitContinuation = false;
// Check for steering messages at start (user may have typed while waiting)
let pendingMessages: AgentMessage[] = (await config.getSteeringMessages?.()) || [];
@@ -173,10 +181,12 @@ async function runLoop(
// Inner loop: process tool calls and steering messages
while (hasMoreToolCalls || pendingMessages.length > 0) {
let preparedMessages: AgentMessage[] = [];
if (lastCompletedTurn) {
const nextTurnSnapshot = await config.prepareNextTurn?.(lastCompletedTurn);
if (nextTurnSnapshot) {
currentContext = nextTurnSnapshot.context ?? currentContext;
preparedMessages = nextTurnSnapshot.messages ?? [];
config = {
...config,
model: nextTurnSnapshot.model ?? config.model,
@@ -197,15 +207,35 @@ async function runLoop(
await emit({ type: "turn_start" });
}
// Process pending messages (inject before next assistant response)
if (pendingMessages.length > 0) {
for (const message of pendingMessages) {
await emit({ type: "message_start", message });
await emit({ type: "message_end", message });
currentContext.messages.push(message);
newMessages.push(message);
}
pendingMessages = [];
// Process prepared and queued messages before the next assistant response.
for (const message of declareToolChanges(currentContext, [...preparedMessages, ...pendingMessages])) {
await emit({ type: "message_start", message });
await emit({ type: "message_end", message });
currentContext.messages.push(message);
newMessages.push(message);
}
pendingMessages = [];
const requestUpdate = await config.prepareRequest?.(
{
context: currentContext,
model: config.model,
thinkingLevel: config.reasoning ?? "off",
},
signal,
);
if (requestUpdate) {
currentContext = requestUpdate.context ?? currentContext;
config = {
...config,
model: requestUpdate.model ?? config.model,
reasoning:
requestUpdate.thinkingLevel === undefined
? config.reasoning
: requestUpdate.thinkingLevel === "off"
? undefined
: requestUpdate.thinkingLevel,
};
}
// Stream assistant response
@@ -213,6 +243,13 @@ async function runLoop(
newMessages.push(message);
if (message.stopReason === "error" || message.stopReason === "aborted") {
lastCompletedTurn = {
message,
toolResults: [],
context: currentContext,
newMessages,
};
await config.finishTurn?.(lastCompletedTurn, signal);
await emit({ type: "turn_end", message, toolResults: [] });
await emit({ type: "agent_end", messages: newMessages });
return;
@@ -240,31 +277,42 @@ async function runLoop(
}
}
await emit({ type: "turn_end", message, toolResults });
lastCompletedTurn = {
message,
toolResults,
context: currentContext,
newMessages,
};
const decision = await config.finishTurn?.(lastCompletedTurn, signal);
await emit({ type: "turn_end", message, toolResults });
if (await config.shouldStopAfterTurn?.(lastCompletedTurn)) {
if (decision?.action === "end") {
await emit({ type: "agent_end", messages: newMessages });
return;
}
explicitContinuation = decision?.action === "continue";
pendingMessages = (await config.getSteeringMessages?.()) || [];
if (hasMoreToolCalls || pendingMessages.length > 0) {
explicitContinuation = false;
}
}
// Agent would stop here. Check for follow-up messages.
const followUpMessages = (await config.getFollowUpMessages?.()) || [];
if (followUpMessages.length > 0) {
// Set as pending so inner loop processes them
explicitContinuation = false;
pendingMessages = followUpMessages;
continue;
}
// No natural request was selected, so fulfill the continuation decision with one context-only turn.
if (explicitContinuation) {
explicitContinuation = false;
continue;
}
// No more messages, exit
break;
}
@@ -272,6 +320,60 @@ async function runLoop(
await emit({ type: "agent_end", messages: newMessages });
}
/**
* Declare tool loadout changes to the model.
*
* `context.tools` is what the runtime can execute; the transcript's system messages declare
* what the model may call. Before each request the difference becomes `toolsAdded` and
* `toolsRemoved` on a system message. When a pending system message exists, its tool fields
* are treated as intent and replaced with the delta between the committed transcript and
* the executable set, so replay always yields exactly `context.tools`. Otherwise a new
* system message is inserted before the first non-system pending message.
*/
function declareToolChanges(context: AgentContext, pendingMessages: AgentMessage[]): AgentMessage[] {
let systemIndex = -1;
for (let i = pendingMessages.length - 1; i >= 0; i--) {
if (pendingMessages[i].role === "system") {
systemIndex = i;
break;
}
}
const pending = pendingMessages[systemIndex] as SystemMessage | undefined;
const baseline = pending
? pendingMessages.map((message, index) =>
index === systemIndex ? withToolChanges(pending, NO_CHANGES) : message,
)
: pendingMessages;
const changes = getToolStateChanges(
getCurrentTools([...context.messages, ...baseline]),
(context.tools ?? []).map(toToolDeclaration),
);
const unchanged = changes.toolsAdded.length === 0 && changes.toolsRemoved.length === 0;
if (pending) {
// Keep the caller's message object when it already declares no tool changes.
if (unchanged && !pending.toolsAdded?.length && !pending.toolsRemoved?.length) return pendingMessages;
return baseline.map((message, index) => (index === systemIndex ? withToolChanges(pending, changes) : message));
}
if (unchanged) return pendingMessages;
const update = withToolChanges({ role: "system", content: "", timestamp: Date.now() }, changes);
const insertIndex = pendingMessages.findIndex((message) => message.role !== "system");
const index = insertIndex === -1 ? pendingMessages.length : insertIndex;
return [...pendingMessages.slice(0, index), update, ...pendingMessages.slice(index)];
}
const NO_CHANGES: ToolStateChanges = { toolsAdded: [], toolsRemoved: [] };
/** Copy a system message with its tool fields replaced by `changes`; empty lists omit the field. */
function withToolChanges(message: SystemMessage, { toolsAdded, toolsRemoved }: ToolStateChanges): SystemMessage {
const { toolsAdded: _added, toolsRemoved: _removed, ...rest } = message;
return {
...rest,
...(toolsAdded.length > 0 ? { toolsAdded } : {}),
...(toolsRemoved.length > 0 ? { toolsRemoved } : {}),
};
}
/**
* Stream an assistant response from the LLM.
* This is where AgentMessage[] gets transformed to Message[] for the LLM.
@@ -292,12 +394,7 @@ async function streamAssistantResponse(
// Convert to LLM-compatible messages (AgentMessage[] → Message[])
const llmMessages = await config.convertToLlm(messages);
// Build LLM context
const llmContext: Context = {
systemPrompt: context.systemPrompt,
messages: llmMessages,
tools: context.tools,
};
const llmContext = normalizeContext({ messages: llmMessages });
// Resolve API key (important for expiring tokens)
const resolvedApiKey =
@@ -308,6 +405,8 @@ async function streamAssistantResponse(
apiKey: resolvedApiKey,
signal,
});
// Record the requested level, whichever stream function answered.
const result = async () => Object.assign(await response.result(), { thinkingLevel: config.reasoning ?? "off" });
let partialMessage: AssistantMessage | null = null;
let addedPartial = false;
@@ -343,7 +442,7 @@ async function streamAssistantResponse(
case "done":
case "error": {
const finalMessage = await response.result();
const finalMessage = await result();
if (addedPartial) {
context.messages[context.messages.length - 1] = finalMessage;
} else {
@@ -358,7 +457,7 @@ async function streamAssistantResponse(
}
}
const finalMessage = await response.result();
const finalMessage = await result();
if (addedPartial) {
context.messages[context.messages.length - 1] = finalMessage;
} else {
@@ -456,7 +555,7 @@ async function executeToolCallsSequential(
isError: preparation.isError,
};
} else {
const executed = await executePreparedToolCall(preparation, signal, emit);
const executed = await executePreparedToolCall(preparation, signal, emitToolExecutionUpdate(toolCall, emit));
finalized = await finalizeExecutedToolCall(
currentContext,
assistantMessage,
@@ -527,7 +626,7 @@ async function executeToolCallsParallel(
await emitToolExecutionEnd(finalized, emit);
return finalized;
}
const executed = await executePreparedToolCall(preparation, signal, emit);
const executed = await executePreparedToolCall(preparation, signal, emitToolExecutionUpdate(toolCall, emit));
const finalized = await finalizeExecutedToolCall(
currentContext,
assistantMessage,
@@ -576,13 +675,15 @@ type ImmediateToolCallOutcome = {
type ExecutedToolCallOutcome = {
result: AgentToolResult<any>;
isError: boolean;
durationMs: number;
};
type FinalizedToolCallOutcome = {
toolCall: AgentToolCall;
result: AgentToolResult<any>;
isError: boolean;
};
type FinalizedToolCallOutcome = AgentToolCallOutcome;
/** The `beforeToolCall` and `afterToolCall` hooks of {@link AgentLoopConfig}. */
export type ToolCallHooks = Pick<AgentLoopConfig, "beforeToolCall" | "afterToolCall">;
type ToolUpdateSink = (partialResult: AgentToolResult<any>) => Promise<void> | void;
type FinalizedToolCallEntry = FinalizedToolCallOutcome | (() => Promise<FinalizedToolCallOutcome>);
@@ -608,10 +709,11 @@ async function prepareToolCall(
currentContext: AgentContext,
assistantMessage: AssistantMessage,
toolCall: AgentToolCall,
config: AgentLoopConfig,
config: ToolCallHooks,
signal: AbortSignal | undefined,
tools: readonly AgentTool<any>[] = currentContext.tools ?? [],
): Promise<PreparedToolCall | ImmediateToolCallOutcome> {
const tool = currentContext.tools?.find((t) => t.name === toolCall.name);
const tool = tools.find((t) => t.name === toolCall.name);
if (!tool) {
return {
kind: "immediate",
@@ -674,13 +776,57 @@ async function prepareToolCall(
}
}
function emitToolExecutionUpdate(toolCall: AgentToolCall, emit: AgentEventSink): ToolUpdateSink {
return (partialResult) =>
emit({
type: "tool_execution_update",
toolCallId: toolCall.id,
toolName: toolCall.name,
args: toolCall.arguments,
partialResult,
});
}
/** Options for {@link runToolCall}. */
export interface RunToolCallOptions extends ToolCallHooks {
/** Tools the call resolves against. */
tools: readonly AgentTool<any>[];
/** Passed to the hooks as the message that issued the call. */
assistantMessage: AssistantMessage;
/** Passed to the hooks as the current agent context. */
context: AgentContext;
signal?: AbortSignal;
onUpdate?: ToolUpdateSink;
}
/**
* Run one tool call through the same steps as a model-issued call: argument preparation, schema
* validation, `beforeToolCall`, execution, and `afterToolCall`. Emits no events and adds no
* messages. Tools that call other tools use this so the hooks (for example permission checks)
* apply to those calls too.
*
* Never rejects for tool failures: unknown tools, validation errors, blocked calls, and thrown
* errors come back as `isError: true`.
*/
export async function runToolCall(toolCall: AgentToolCall, options: RunToolCallOptions): Promise<AgentToolCallOutcome> {
const { assistantMessage, context, signal } = options;
const preparation = await prepareToolCall(context, assistantMessage, toolCall, options, signal, options.tools);
if (preparation.kind === "immediate") {
return { toolCall, result: preparation.result, isError: preparation.isError };
}
const executed = await executePreparedToolCall(preparation, signal, options.onUpdate ?? (() => {}));
return finalizeExecutedToolCall(context, assistantMessage, preparation, executed, options, signal);
}
async function executePreparedToolCall(
prepared: PreparedToolCall,
signal: AbortSignal | undefined,
emit: AgentEventSink,
onUpdate: ToolUpdateSink,
): Promise<ExecutedToolCallOutcome> {
const updateEvents: Promise<void>[] = [];
let acceptingUpdates = true;
const startedAt = performance.now();
const elapsed = () => Math.round(performance.now() - startedAt);
try {
const result = await prepared.tool.execute(
@@ -689,28 +835,21 @@ async function executePreparedToolCall(
signal,
(partialResult) => {
if (!acceptingUpdates) return;
updateEvents.push(
Promise.resolve(
emit({
type: "tool_execution_update",
toolCallId: prepared.toolCall.id,
toolName: prepared.toolCall.name,
args: prepared.toolCall.arguments,
partialResult,
}),
),
);
updateEvents.push(Promise.resolve(onUpdate(partialResult)));
},
);
const durationMs = elapsed();
acceptingUpdates = false;
await Promise.all(updateEvents);
return { result, isError: false };
return { result, isError: result.isError === true, durationMs };
} catch (error) {
const durationMs = elapsed();
acceptingUpdates = false;
await Promise.all(updateEvents);
return {
result: createErrorToolResult(error instanceof Error ? error.message : String(error)),
isError: true,
durationMs,
};
} finally {
acceptingUpdates = false;
@@ -722,7 +861,7 @@ async function finalizeExecutedToolCall(
assistantMessage: AssistantMessage,
prepared: PreparedToolCall,
executed: ExecutedToolCallOutcome,
config: AgentLoopConfig,
config: ToolCallHooks,
signal: AbortSignal | undefined,
): Promise<FinalizedToolCallOutcome> {
let result = executed.result;
@@ -742,6 +881,9 @@ async function finalizeExecutedToolCall(
signal,
);
if (afterResult) {
// Structured content not replaced along with the content may no longer match it.
const structuredContent =
afterResult.structuredContent ?? (afterResult.content ? undefined : result.structuredContent);
result = {
...result,
content: afterResult.content ?? result.content,
@@ -749,6 +891,8 @@ async function finalizeExecutedToolCall(
usage: afterResult.usage ?? result.usage,
terminate: afterResult.terminate ?? result.terminate,
};
if (structuredContent === undefined) delete result.structuredContent;
else result.structuredContent = structuredContent;
isError = afterResult.isError ?? isError;
}
} catch (error) {
@@ -761,6 +905,7 @@ async function finalizeExecutedToolCall(
toolCall: prepared.toolCall,
result,
isError,
durationMs: executed.durationMs,
};
}
@@ -778,6 +923,7 @@ async function emitToolExecutionEnd(finalized: FinalizedToolCallOutcome, emit: A
toolName: finalized.toolCall.name,
result: finalized.result,
isError: finalized.isError,
...(finalized.durationMs === undefined ? {} : { durationMs: finalized.durationMs }),
});
}
@@ -791,8 +937,8 @@ function createToolResultMessage(finalized: FinalizedToolCallOutcome): ToolResul
content: finalized.result.content ?? [],
details: finalized.result.details,
usage: finalized.result.usage,
...(finalized.result.addedToolNames?.length ? { addedToolNames: finalized.result.addedToolNames } : {}),
isError: finalized.isError,
...(finalized.durationMs === undefined ? {} : { durationMs: finalized.durationMs }),
timestamp: Date.now(),
};
}