chore(vendor): import pi 0.85.1
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: d981de1229ef899957bbe968bc8dcda02a21f477
This commit is contained in:
803
third_party/pi/packages/agent/src/agent-loop.ts
vendored
Normal file
803
third_party/pi/packages/agent/src/agent-loop.ts
vendored
Normal file
@@ -0,0 +1,803 @@
|
||||
/**
|
||||
* Agent loop that works with AgentMessage throughout.
|
||||
* Transforms to Message[] only at the LLM call boundary.
|
||||
*/
|
||||
|
||||
import {
|
||||
type AssistantMessage,
|
||||
type Context,
|
||||
EventStream,
|
||||
type ToolResultMessage,
|
||||
validateToolArguments,
|
||||
} from "@earendil-works/pi-ai";
|
||||
import { getDefaultStreamFn } from "./stream-fn.ts";
|
||||
import type {
|
||||
AgentContext,
|
||||
AgentEvent,
|
||||
AgentLoopConfig,
|
||||
AgentMessage,
|
||||
AgentTool,
|
||||
AgentToolCall,
|
||||
AgentToolResult,
|
||||
PrepareNextTurnContext,
|
||||
StreamFn,
|
||||
} from "./types.ts";
|
||||
|
||||
export type AgentEventSink = (event: AgentEvent) => Promise<void> | void;
|
||||
|
||||
/**
|
||||
* Start an agent loop with a new prompt message.
|
||||
* The prompt is added to the context and events are emitted for it.
|
||||
*/
|
||||
export function agentLoop(
|
||||
prompts: AgentMessage[],
|
||||
context: AgentContext,
|
||||
config: AgentLoopConfig,
|
||||
signal: AbortSignal | undefined,
|
||||
streamFn: StreamFn,
|
||||
): EventStream<AgentEvent, AgentMessage[]> {
|
||||
const stream = createAgentStream();
|
||||
|
||||
void runAgentLoop(
|
||||
prompts,
|
||||
context,
|
||||
config,
|
||||
async (event) => {
|
||||
stream.push(event);
|
||||
},
|
||||
signal,
|
||||
streamFn,
|
||||
).then((messages) => {
|
||||
stream.end(messages);
|
||||
});
|
||||
|
||||
return stream;
|
||||
}
|
||||
|
||||
/**
|
||||
* Continue an agent loop from the current context without adding a new message.
|
||||
* Used for retries - context already has user message or tool results.
|
||||
*
|
||||
* **Important:** The last message in context must convert to a `user` or `toolResult` message
|
||||
* via `convertToLlm`. If it doesn't, the LLM provider will reject the request.
|
||||
* This cannot be validated here since `convertToLlm` is only called once per turn.
|
||||
*/
|
||||
export function agentLoopContinue(
|
||||
context: AgentContext,
|
||||
config: AgentLoopConfig,
|
||||
signal: AbortSignal | undefined,
|
||||
streamFn: StreamFn,
|
||||
): EventStream<AgentEvent, AgentMessage[]> {
|
||||
if (context.messages.length === 0) {
|
||||
throw new Error("Cannot continue: no messages in context");
|
||||
}
|
||||
|
||||
if (context.messages[context.messages.length - 1].role === "assistant") {
|
||||
throw new Error("Cannot continue from message role: assistant");
|
||||
}
|
||||
|
||||
const stream = createAgentStream();
|
||||
|
||||
void runAgentLoopContinue(
|
||||
context,
|
||||
config,
|
||||
async (event) => {
|
||||
stream.push(event);
|
||||
},
|
||||
signal,
|
||||
streamFn,
|
||||
).then((messages) => {
|
||||
stream.end(messages);
|
||||
});
|
||||
|
||||
return stream;
|
||||
}
|
||||
|
||||
export async function runAgentLoop(
|
||||
prompts: AgentMessage[],
|
||||
context: AgentContext,
|
||||
config: AgentLoopConfig,
|
||||
emit: AgentEventSink,
|
||||
signal: AbortSignal | undefined,
|
||||
streamFn: StreamFn,
|
||||
): Promise<AgentMessage[]> {
|
||||
const newMessages: AgentMessage[] = [...prompts];
|
||||
const currentContext: AgentContext = {
|
||||
...context,
|
||||
messages: [...context.messages, ...prompts],
|
||||
};
|
||||
|
||||
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 });
|
||||
}
|
||||
|
||||
await runLoop(currentContext, newMessages, config, signal, emit, streamFn ?? getDefaultStreamFn());
|
||||
return newMessages;
|
||||
}
|
||||
|
||||
export async function runAgentLoopContinue(
|
||||
context: AgentContext,
|
||||
config: AgentLoopConfig,
|
||||
emit: AgentEventSink,
|
||||
signal: AbortSignal | undefined,
|
||||
streamFn: StreamFn,
|
||||
): Promise<AgentMessage[]> {
|
||||
if (context.messages.length === 0) {
|
||||
throw new Error("Cannot continue: no messages in context");
|
||||
}
|
||||
|
||||
if (context.messages[context.messages.length - 1].role === "assistant") {
|
||||
throw new Error("Cannot continue from message role: assistant");
|
||||
}
|
||||
|
||||
const newMessages: AgentMessage[] = [];
|
||||
const currentContext: AgentContext = { ...context };
|
||||
|
||||
await emit({ type: "agent_start" });
|
||||
await emit({ type: "turn_start" });
|
||||
|
||||
await runLoop(currentContext, newMessages, config, signal, emit, streamFn ?? getDefaultStreamFn());
|
||||
return newMessages;
|
||||
}
|
||||
|
||||
function createAgentStream(): EventStream<AgentEvent, AgentMessage[]> {
|
||||
return new EventStream<AgentEvent, AgentMessage[]>(
|
||||
(event: AgentEvent) => event.type === "agent_end",
|
||||
(event: AgentEvent) => (event.type === "agent_end" ? event.messages : []),
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Main loop logic shared by agentLoop and agentLoopContinue.
|
||||
*/
|
||||
async function runLoop(
|
||||
initialContext: AgentContext,
|
||||
newMessages: AgentMessage[],
|
||||
initialConfig: AgentLoopConfig,
|
||||
signal: AbortSignal | undefined,
|
||||
emit: AgentEventSink,
|
||||
streamFunction: StreamFn,
|
||||
): Promise<void> {
|
||||
let currentContext = initialContext;
|
||||
let config = initialConfig;
|
||||
let lastCompletedTurn: PrepareNextTurnContext | undefined;
|
||||
// Check for steering messages at start (user may have typed while waiting)
|
||||
let pendingMessages: AgentMessage[] = (await config.getSteeringMessages?.()) || [];
|
||||
|
||||
// Outer loop: continues when queued follow-up messages arrive after agent would stop
|
||||
while (true) {
|
||||
let hasMoreToolCalls = true;
|
||||
|
||||
// Inner loop: process tool calls and steering messages
|
||||
while (hasMoreToolCalls || pendingMessages.length > 0) {
|
||||
if (lastCompletedTurn) {
|
||||
const nextTurnSnapshot = await config.prepareNextTurn?.(lastCompletedTurn);
|
||||
if (nextTurnSnapshot) {
|
||||
currentContext = nextTurnSnapshot.context ?? currentContext;
|
||||
config = {
|
||||
...config,
|
||||
model: nextTurnSnapshot.model ?? config.model,
|
||||
reasoning:
|
||||
nextTurnSnapshot.thinkingLevel === undefined
|
||||
? config.reasoning
|
||||
: nextTurnSnapshot.thinkingLevel === "off"
|
||||
? undefined
|
||||
: nextTurnSnapshot.thinkingLevel,
|
||||
};
|
||||
}
|
||||
// Preparation can be long-running (for example, compaction). Pick up steering
|
||||
// queued while it ran. Only poll again if the earlier poll returned nothing;
|
||||
// otherwise one-at-a-time mode would deliver two messages in this turn.
|
||||
if (pendingMessages.length === 0) {
|
||||
pendingMessages = (await config.getSteeringMessages?.()) || [];
|
||||
}
|
||||
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 = [];
|
||||
}
|
||||
|
||||
// Stream assistant response
|
||||
const message = await streamAssistantResponse(currentContext, config, signal, emit, streamFunction);
|
||||
newMessages.push(message);
|
||||
|
||||
if (message.stopReason === "error" || message.stopReason === "aborted") {
|
||||
await emit({ type: "turn_end", message, toolResults: [] });
|
||||
await emit({ type: "agent_end", messages: newMessages });
|
||||
return;
|
||||
}
|
||||
|
||||
// Check for tool calls
|
||||
const toolCalls = message.content.filter((c) => c.type === "toolCall");
|
||||
|
||||
const toolResults: ToolResultMessage[] = [];
|
||||
hasMoreToolCalls = false;
|
||||
if (toolCalls.length > 0) {
|
||||
// A "length" stop means the output was cut off by the token limit, so
|
||||
// every tool call in the message may carry truncated arguments. Fail
|
||||
// them all instead of executing potentially borked calls.
|
||||
const executedToolBatch =
|
||||
message.stopReason === "length"
|
||||
? await failToolCallsFromTruncatedMessage(toolCalls, emit)
|
||||
: await executeToolCalls(currentContext, message, config, signal, emit);
|
||||
toolResults.push(...executedToolBatch.messages);
|
||||
hasMoreToolCalls = !executedToolBatch.terminate;
|
||||
|
||||
for (const result of toolResults) {
|
||||
currentContext.messages.push(result);
|
||||
newMessages.push(result);
|
||||
}
|
||||
}
|
||||
|
||||
await emit({ type: "turn_end", message, toolResults });
|
||||
|
||||
lastCompletedTurn = {
|
||||
message,
|
||||
toolResults,
|
||||
context: currentContext,
|
||||
newMessages,
|
||||
};
|
||||
|
||||
if (await config.shouldStopAfterTurn?.(lastCompletedTurn)) {
|
||||
await emit({ type: "agent_end", messages: newMessages });
|
||||
return;
|
||||
}
|
||||
|
||||
pendingMessages = (await config.getSteeringMessages?.()) || [];
|
||||
}
|
||||
|
||||
// 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
|
||||
pendingMessages = followUpMessages;
|
||||
continue;
|
||||
}
|
||||
|
||||
// No more messages, exit
|
||||
break;
|
||||
}
|
||||
|
||||
await emit({ type: "agent_end", messages: newMessages });
|
||||
}
|
||||
|
||||
/**
|
||||
* Stream an assistant response from the LLM.
|
||||
* This is where AgentMessage[] gets transformed to Message[] for the LLM.
|
||||
*/
|
||||
async function streamAssistantResponse(
|
||||
context: AgentContext,
|
||||
config: AgentLoopConfig,
|
||||
signal: AbortSignal | undefined,
|
||||
emit: AgentEventSink,
|
||||
streamFunction: StreamFn,
|
||||
): Promise<AssistantMessage> {
|
||||
// Apply context transform if configured (AgentMessage[] → AgentMessage[])
|
||||
let messages = context.messages;
|
||||
if (config.transformContext) {
|
||||
messages = await config.transformContext(messages, signal);
|
||||
}
|
||||
|
||||
// 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,
|
||||
};
|
||||
|
||||
// Resolve API key (important for expiring tokens)
|
||||
const resolvedApiKey =
|
||||
(config.getApiKey ? await config.getApiKey(config.model.provider) : undefined) || config.apiKey;
|
||||
|
||||
const response = await streamFunction(config.model, llmContext, {
|
||||
...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;
|
||||
await emit({ 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;
|
||||
await emit({
|
||||
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) {
|
||||
await emit({ type: "message_start", message: { ...finalMessage } });
|
||||
}
|
||||
await emit({ type: "message_end", message: finalMessage });
|
||||
return finalMessage;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const finalMessage = await response.result();
|
||||
if (addedPartial) {
|
||||
context.messages[context.messages.length - 1] = finalMessage;
|
||||
} else {
|
||||
context.messages.push(finalMessage);
|
||||
await emit({ type: "message_start", message: { ...finalMessage } });
|
||||
}
|
||||
await emit({ type: "message_end", message: finalMessage });
|
||||
return finalMessage;
|
||||
}
|
||||
|
||||
/**
|
||||
* Fail all tool calls from an assistant message that was truncated by the
|
||||
* output token limit. Streamed tool-call arguments are finalized with a
|
||||
* best-effort JSON salvage parser, so a truncated message can yield tool calls
|
||||
* whose arguments parse and validate but are silently incomplete. None of them
|
||||
* are safe to execute; report each as an error so the model can re-issue them.
|
||||
*/
|
||||
async function failToolCallsFromTruncatedMessage(
|
||||
toolCalls: AgentToolCall[],
|
||||
emit: AgentEventSink,
|
||||
): Promise<ExecutedToolCallBatch> {
|
||||
const messages: ToolResultMessage[] = [];
|
||||
for (const toolCall of toolCalls) {
|
||||
await emit({
|
||||
type: "tool_execution_start",
|
||||
toolCallId: toolCall.id,
|
||||
toolName: toolCall.name,
|
||||
args: toolCall.arguments,
|
||||
});
|
||||
const finalized: FinalizedToolCallOutcome = {
|
||||
toolCall,
|
||||
result: createErrorToolResult(
|
||||
`Tool call "${toolCall.name}" was not executed: the response hit the output token limit, so its arguments may be truncated. Re-issue the tool call with complete arguments.`,
|
||||
),
|
||||
isError: true,
|
||||
};
|
||||
await emitToolExecutionEnd(finalized, emit);
|
||||
const toolResultMessage = createToolResultMessage(finalized);
|
||||
await emitToolResultMessage(toolResultMessage, emit);
|
||||
messages.push(toolResultMessage);
|
||||
}
|
||||
return { messages, terminate: false };
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute tool calls from an assistant message.
|
||||
*/
|
||||
async function executeToolCalls(
|
||||
currentContext: AgentContext,
|
||||
assistantMessage: AssistantMessage,
|
||||
config: AgentLoopConfig,
|
||||
signal: AbortSignal | undefined,
|
||||
emit: AgentEventSink,
|
||||
): Promise<ExecutedToolCallBatch> {
|
||||
const toolCalls = assistantMessage.content.filter((c) => c.type === "toolCall");
|
||||
const hasSequentialToolCall = toolCalls.some(
|
||||
(tc) => currentContext.tools?.find((t) => t.name === tc.name)?.executionMode === "sequential",
|
||||
);
|
||||
if (config.toolExecution === "sequential" || hasSequentialToolCall) {
|
||||
return executeToolCallsSequential(currentContext, assistantMessage, toolCalls, config, signal, emit);
|
||||
}
|
||||
return executeToolCallsParallel(currentContext, assistantMessage, toolCalls, config, signal, emit);
|
||||
}
|
||||
|
||||
type ExecutedToolCallBatch = {
|
||||
messages: ToolResultMessage[];
|
||||
terminate: boolean;
|
||||
};
|
||||
|
||||
async function executeToolCallsSequential(
|
||||
currentContext: AgentContext,
|
||||
assistantMessage: AssistantMessage,
|
||||
toolCalls: AgentToolCall[],
|
||||
config: AgentLoopConfig,
|
||||
signal: AbortSignal | undefined,
|
||||
emit: AgentEventSink,
|
||||
): Promise<ExecutedToolCallBatch> {
|
||||
const finalizedCalls: FinalizedToolCallOutcome[] = [];
|
||||
const messages: ToolResultMessage[] = [];
|
||||
|
||||
for (const toolCall of toolCalls) {
|
||||
await emit({
|
||||
type: "tool_execution_start",
|
||||
toolCallId: toolCall.id,
|
||||
toolName: toolCall.name,
|
||||
args: toolCall.arguments,
|
||||
});
|
||||
|
||||
const preparation = await prepareToolCall(currentContext, assistantMessage, toolCall, config, signal);
|
||||
let finalized: FinalizedToolCallOutcome;
|
||||
if (preparation.kind === "immediate") {
|
||||
finalized = {
|
||||
toolCall,
|
||||
result: preparation.result,
|
||||
isError: preparation.isError,
|
||||
};
|
||||
} else {
|
||||
const executed = await executePreparedToolCall(preparation, signal, emit);
|
||||
finalized = await finalizeExecutedToolCall(
|
||||
currentContext,
|
||||
assistantMessage,
|
||||
preparation,
|
||||
executed,
|
||||
config,
|
||||
signal,
|
||||
);
|
||||
}
|
||||
|
||||
await emitToolExecutionEnd(finalized, emit);
|
||||
const toolResultMessage = createToolResultMessage(finalized);
|
||||
await emitToolResultMessage(toolResultMessage, emit);
|
||||
finalizedCalls.push(finalized);
|
||||
messages.push(toolResultMessage);
|
||||
|
||||
if (signal?.aborted) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
messages,
|
||||
terminate: shouldTerminateToolBatch(finalizedCalls),
|
||||
};
|
||||
}
|
||||
|
||||
async function executeToolCallsParallel(
|
||||
currentContext: AgentContext,
|
||||
assistantMessage: AssistantMessage,
|
||||
toolCalls: AgentToolCall[],
|
||||
config: AgentLoopConfig,
|
||||
signal: AbortSignal | undefined,
|
||||
emit: AgentEventSink,
|
||||
): Promise<ExecutedToolCallBatch> {
|
||||
const finalizedCalls: FinalizedToolCallEntry[] = [];
|
||||
|
||||
for (const toolCall of toolCalls) {
|
||||
await emit({
|
||||
type: "tool_execution_start",
|
||||
toolCallId: toolCall.id,
|
||||
toolName: toolCall.name,
|
||||
args: toolCall.arguments,
|
||||
});
|
||||
|
||||
const preparation = await prepareToolCall(currentContext, assistantMessage, toolCall, config, signal);
|
||||
if (preparation.kind === "immediate") {
|
||||
const finalized = {
|
||||
toolCall,
|
||||
result: preparation.result,
|
||||
isError: preparation.isError,
|
||||
} satisfies FinalizedToolCallOutcome;
|
||||
await emitToolExecutionEnd(finalized, emit);
|
||||
finalizedCalls.push(finalized);
|
||||
if (signal?.aborted) {
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
finalizedCalls.push(async () => {
|
||||
if (signal?.aborted) {
|
||||
const finalized = {
|
||||
toolCall,
|
||||
result: createErrorToolResult("Operation aborted"),
|
||||
isError: true,
|
||||
} satisfies FinalizedToolCallOutcome;
|
||||
await emitToolExecutionEnd(finalized, emit);
|
||||
return finalized;
|
||||
}
|
||||
const executed = await executePreparedToolCall(preparation, signal, emit);
|
||||
const finalized = await finalizeExecutedToolCall(
|
||||
currentContext,
|
||||
assistantMessage,
|
||||
preparation,
|
||||
executed,
|
||||
config,
|
||||
signal,
|
||||
);
|
||||
await emitToolExecutionEnd(finalized, emit);
|
||||
return finalized;
|
||||
});
|
||||
if (signal?.aborted) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
const orderedFinalizedCalls = await Promise.all(
|
||||
finalizedCalls.map((entry) => (typeof entry === "function" ? entry() : Promise.resolve(entry))),
|
||||
);
|
||||
const messages: ToolResultMessage[] = [];
|
||||
for (const finalized of orderedFinalizedCalls) {
|
||||
const toolResultMessage = createToolResultMessage(finalized);
|
||||
await emitToolResultMessage(toolResultMessage, emit);
|
||||
messages.push(toolResultMessage);
|
||||
}
|
||||
|
||||
return {
|
||||
messages,
|
||||
terminate: shouldTerminateToolBatch(orderedFinalizedCalls),
|
||||
};
|
||||
}
|
||||
|
||||
type PreparedToolCall = {
|
||||
kind: "prepared";
|
||||
toolCall: AgentToolCall;
|
||||
tool: AgentTool<any>;
|
||||
args: unknown;
|
||||
};
|
||||
|
||||
type ImmediateToolCallOutcome = {
|
||||
kind: "immediate";
|
||||
result: AgentToolResult<any>;
|
||||
isError: boolean;
|
||||
};
|
||||
|
||||
type ExecutedToolCallOutcome = {
|
||||
result: AgentToolResult<any>;
|
||||
isError: boolean;
|
||||
};
|
||||
|
||||
type FinalizedToolCallOutcome = {
|
||||
toolCall: AgentToolCall;
|
||||
result: AgentToolResult<any>;
|
||||
isError: boolean;
|
||||
};
|
||||
|
||||
type FinalizedToolCallEntry = FinalizedToolCallOutcome | (() => Promise<FinalizedToolCallOutcome>);
|
||||
|
||||
function shouldTerminateToolBatch(finalizedCalls: FinalizedToolCallOutcome[]): boolean {
|
||||
return finalizedCalls.length > 0 && finalizedCalls.every((finalized) => finalized.result.terminate === true);
|
||||
}
|
||||
|
||||
function prepareToolCallArguments(tool: AgentTool<any>, toolCall: AgentToolCall): AgentToolCall {
|
||||
if (!tool.prepareArguments) {
|
||||
return toolCall;
|
||||
}
|
||||
const preparedArguments = tool.prepareArguments(toolCall.arguments);
|
||||
if (preparedArguments === toolCall.arguments) {
|
||||
return toolCall;
|
||||
}
|
||||
return {
|
||||
...toolCall,
|
||||
arguments: preparedArguments as Record<string, any>,
|
||||
};
|
||||
}
|
||||
|
||||
async function prepareToolCall(
|
||||
currentContext: AgentContext,
|
||||
assistantMessage: AssistantMessage,
|
||||
toolCall: AgentToolCall,
|
||||
config: AgentLoopConfig,
|
||||
signal: AbortSignal | undefined,
|
||||
): Promise<PreparedToolCall | ImmediateToolCallOutcome> {
|
||||
const tool = currentContext.tools?.find((t) => t.name === toolCall.name);
|
||||
if (!tool) {
|
||||
return {
|
||||
kind: "immediate",
|
||||
result: createErrorToolResult(`Tool ${toolCall.name} not found`),
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
|
||||
try {
|
||||
const preparedToolCall = prepareToolCallArguments(tool, toolCall);
|
||||
const validatedArgs = validateToolArguments(tool, preparedToolCall);
|
||||
if (config.beforeToolCall) {
|
||||
const beforeResult = await config.beforeToolCall(
|
||||
{
|
||||
assistantMessage,
|
||||
toolCall,
|
||||
args: validatedArgs,
|
||||
context: currentContext,
|
||||
},
|
||||
signal,
|
||||
);
|
||||
if (signal?.aborted) {
|
||||
return {
|
||||
kind: "immediate",
|
||||
result: createErrorToolResult("Operation aborted"),
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
if (beforeResult?.block) {
|
||||
const result = createErrorToolResult(beforeResult.reason || "Tool execution was blocked");
|
||||
if (beforeResult.terminate === true) {
|
||||
result.terminate = true;
|
||||
}
|
||||
return {
|
||||
kind: "immediate",
|
||||
result,
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
}
|
||||
if (signal?.aborted) {
|
||||
return {
|
||||
kind: "immediate",
|
||||
result: createErrorToolResult("Operation aborted"),
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
return {
|
||||
kind: "prepared",
|
||||
toolCall,
|
||||
tool,
|
||||
args: validatedArgs,
|
||||
};
|
||||
} catch (error) {
|
||||
return {
|
||||
kind: "immediate",
|
||||
result: createErrorToolResult(error instanceof Error ? error.message : String(error)),
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
async function executePreparedToolCall(
|
||||
prepared: PreparedToolCall,
|
||||
signal: AbortSignal | undefined,
|
||||
emit: AgentEventSink,
|
||||
): Promise<ExecutedToolCallOutcome> {
|
||||
const updateEvents: Promise<void>[] = [];
|
||||
let acceptingUpdates = true;
|
||||
|
||||
try {
|
||||
const result = await prepared.tool.execute(
|
||||
prepared.toolCall.id,
|
||||
prepared.args as never,
|
||||
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,
|
||||
}),
|
||||
),
|
||||
);
|
||||
},
|
||||
);
|
||||
acceptingUpdates = false;
|
||||
await Promise.all(updateEvents);
|
||||
return { result, isError: false };
|
||||
} catch (error) {
|
||||
acceptingUpdates = false;
|
||||
await Promise.all(updateEvents);
|
||||
return {
|
||||
result: createErrorToolResult(error instanceof Error ? error.message : String(error)),
|
||||
isError: true,
|
||||
};
|
||||
} finally {
|
||||
acceptingUpdates = false;
|
||||
}
|
||||
}
|
||||
|
||||
async function finalizeExecutedToolCall(
|
||||
currentContext: AgentContext,
|
||||
assistantMessage: AssistantMessage,
|
||||
prepared: PreparedToolCall,
|
||||
executed: ExecutedToolCallOutcome,
|
||||
config: AgentLoopConfig,
|
||||
signal: AbortSignal | undefined,
|
||||
): Promise<FinalizedToolCallOutcome> {
|
||||
let result = executed.result;
|
||||
let isError = executed.isError;
|
||||
|
||||
if (config.afterToolCall) {
|
||||
try {
|
||||
const afterResult = await config.afterToolCall(
|
||||
{
|
||||
assistantMessage,
|
||||
toolCall: prepared.toolCall,
|
||||
args: prepared.args,
|
||||
result,
|
||||
isError,
|
||||
context: currentContext,
|
||||
},
|
||||
signal,
|
||||
);
|
||||
if (afterResult) {
|
||||
result = {
|
||||
...result,
|
||||
content: afterResult.content ?? result.content,
|
||||
details: afterResult.details ?? result.details,
|
||||
usage: afterResult.usage ?? result.usage,
|
||||
terminate: afterResult.terminate ?? result.terminate,
|
||||
};
|
||||
isError = afterResult.isError ?? isError;
|
||||
}
|
||||
} catch (error) {
|
||||
result = createErrorToolResult(error instanceof Error ? error.message : String(error));
|
||||
isError = true;
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
toolCall: prepared.toolCall,
|
||||
result,
|
||||
isError,
|
||||
};
|
||||
}
|
||||
|
||||
function createErrorToolResult(message: string): AgentToolResult<any> {
|
||||
return {
|
||||
content: [{ type: "text", text: message }],
|
||||
details: {},
|
||||
};
|
||||
}
|
||||
|
||||
async function emitToolExecutionEnd(finalized: FinalizedToolCallOutcome, emit: AgentEventSink): Promise<void> {
|
||||
await emit({
|
||||
type: "tool_execution_end",
|
||||
toolCallId: finalized.toolCall.id,
|
||||
toolName: finalized.toolCall.name,
|
||||
result: finalized.result,
|
||||
isError: finalized.isError,
|
||||
});
|
||||
}
|
||||
|
||||
function createToolResultMessage(finalized: FinalizedToolCallOutcome): ToolResultMessage {
|
||||
return {
|
||||
role: "toolResult",
|
||||
toolCallId: finalized.toolCall.id,
|
||||
toolName: finalized.toolCall.name,
|
||||
// Untyped tools (JS extensions) can return results without content; normalize
|
||||
// so the null never enters session history or provider payloads.
|
||||
content: finalized.result.content ?? [],
|
||||
details: finalized.result.details,
|
||||
usage: finalized.result.usage,
|
||||
...(finalized.result.addedToolNames?.length ? { addedToolNames: finalized.result.addedToolNames } : {}),
|
||||
isError: finalized.isError,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
async function emitToolResultMessage(toolResultMessage: ToolResultMessage, emit: AgentEventSink): Promise<void> {
|
||||
await emit({ type: "message_start", message: toolResultMessage });
|
||||
await emit({ type: "message_end", message: toolResultMessage });
|
||||
}
|
||||
592
third_party/pi/packages/agent/src/agent.ts
vendored
Normal file
592
third_party/pi/packages/agent/src/agent.ts
vendored
Normal file
@@ -0,0 +1,592 @@
|
||||
import type {
|
||||
ImageContent,
|
||||
Message,
|
||||
Model,
|
||||
SimpleStreamOptions,
|
||||
TextContent,
|
||||
ThinkingBudgets,
|
||||
Transport,
|
||||
} from "@earendil-works/pi-ai";
|
||||
import { runAgentLoop, runAgentLoopContinue } from "./agent-loop.ts";
|
||||
import { getDefaultStreamFn } from "./stream-fn.ts";
|
||||
import type {
|
||||
AfterToolCallContext,
|
||||
AfterToolCallResult,
|
||||
AgentContext,
|
||||
AgentEvent,
|
||||
AgentLoopConfig,
|
||||
AgentLoopTurnUpdate,
|
||||
AgentMessage,
|
||||
AgentState,
|
||||
AgentTool,
|
||||
BeforeToolCallContext,
|
||||
BeforeToolCallResult,
|
||||
PrepareNextTurnContext,
|
||||
QueueMode,
|
||||
ShouldStopAfterTurnContext,
|
||||
StreamFn,
|
||||
ToolExecutionMode,
|
||||
} from "./types.ts";
|
||||
|
||||
export type { QueueMode } from "./types.ts";
|
||||
|
||||
function defaultConvertToLlm(messages: AgentMessage[]): Message[] {
|
||||
return messages.filter(
|
||||
(message) => message.role === "user" || message.role === "assistant" || message.role === "toolResult",
|
||||
);
|
||||
}
|
||||
|
||||
const EMPTY_USAGE = {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
|
||||
const DEFAULT_MODEL = {
|
||||
id: "unknown",
|
||||
name: "unknown",
|
||||
api: "unknown",
|
||||
provider: "unknown",
|
||||
baseUrl: "",
|
||||
reasoning: false,
|
||||
input: [],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 0,
|
||||
maxTokens: 0,
|
||||
} satisfies Model<any>;
|
||||
|
||||
type MutableAgentState = Omit<AgentState, "isStreaming" | "streamingMessage" | "pendingToolCalls" | "errorMessage"> & {
|
||||
isStreaming: boolean;
|
||||
streamingMessage?: AgentMessage;
|
||||
pendingToolCalls: Set<string>;
|
||||
errorMessage?: string;
|
||||
};
|
||||
|
||||
function createMutableAgentState(
|
||||
initialState?: Partial<Omit<AgentState, "pendingToolCalls" | "isStreaming" | "streamingMessage" | "errorMessage">>,
|
||||
): MutableAgentState {
|
||||
let tools = initialState?.tools?.slice() ?? [];
|
||||
let messages = initialState?.messages?.slice() ?? [];
|
||||
|
||||
return {
|
||||
systemPrompt: initialState?.systemPrompt ?? "",
|
||||
model: initialState?.model ?? DEFAULT_MODEL,
|
||||
thinkingLevel: initialState?.thinkingLevel ?? "off",
|
||||
get tools() {
|
||||
return tools;
|
||||
},
|
||||
set tools(nextTools: AgentTool<any>[]) {
|
||||
tools = nextTools.slice();
|
||||
},
|
||||
get messages() {
|
||||
return messages;
|
||||
},
|
||||
set messages(nextMessages: AgentMessage[]) {
|
||||
messages = nextMessages.slice();
|
||||
},
|
||||
isStreaming: false,
|
||||
streamingMessage: undefined,
|
||||
pendingToolCalls: new Set<string>(),
|
||||
errorMessage: undefined,
|
||||
};
|
||||
}
|
||||
|
||||
/** Options for constructing an {@link Agent}. */
|
||||
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;
|
||||
getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
|
||||
onPayload?: SimpleStreamOptions["onPayload"];
|
||||
onResponse?: SimpleStreamOptions["onResponse"];
|
||||
beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>;
|
||||
afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>;
|
||||
shouldStopAfterTurn?: (context: ShouldStopAfterTurnContext, signal?: AbortSignal) => boolean | Promise<boolean>;
|
||||
prepareNextTurn?: (
|
||||
signal?: AbortSignal,
|
||||
) => Promise<AgentLoopTurnUpdate | undefined> | AgentLoopTurnUpdate | undefined;
|
||||
prepareNextTurnWithContext?: (
|
||||
context: PrepareNextTurnContext,
|
||||
signal?: AbortSignal,
|
||||
) => Promise<AgentLoopTurnUpdate | undefined> | AgentLoopTurnUpdate | undefined;
|
||||
steeringMode?: QueueMode;
|
||||
followUpMode?: QueueMode;
|
||||
sessionId?: string;
|
||||
thinkingBudgets?: ThinkingBudgets;
|
||||
transport?: Transport;
|
||||
maxRetryDelayMs?: number;
|
||||
toolExecution?: ToolExecutionMode;
|
||||
}
|
||||
|
||||
class PendingMessageQueue {
|
||||
private messages: AgentMessage[] = [];
|
||||
public mode: QueueMode;
|
||||
|
||||
constructor(mode: QueueMode) {
|
||||
this.mode = mode;
|
||||
}
|
||||
|
||||
enqueue(message: AgentMessage): void {
|
||||
this.messages.push(message);
|
||||
}
|
||||
|
||||
hasItems(): boolean {
|
||||
return this.messages.length > 0;
|
||||
}
|
||||
|
||||
drain(): AgentMessage[] {
|
||||
if (this.mode === "all") {
|
||||
const drained = this.messages.slice();
|
||||
this.messages = [];
|
||||
return drained;
|
||||
}
|
||||
|
||||
const first = this.messages[0];
|
||||
if (!first) {
|
||||
return [];
|
||||
}
|
||||
this.messages = this.messages.slice(1);
|
||||
return [first];
|
||||
}
|
||||
|
||||
clear(): void {
|
||||
this.messages = [];
|
||||
}
|
||||
}
|
||||
|
||||
type ActiveRun = {
|
||||
promise: Promise<void>;
|
||||
resolve: () => void;
|
||||
abortController: AbortController;
|
||||
};
|
||||
|
||||
/**
|
||||
* Stateful wrapper around the low-level agent loop.
|
||||
*
|
||||
* `Agent` owns the current transcript, emits lifecycle events, executes tools,
|
||||
* and exposes queueing APIs for steering and follow-up messages.
|
||||
*/
|
||||
export class Agent {
|
||||
private _state: MutableAgentState;
|
||||
private readonly listeners = new Set<(event: AgentEvent, signal: AbortSignal) => Promise<void> | void>();
|
||||
private readonly steeringQueue: PendingMessageQueue;
|
||||
private readonly followUpQueue: PendingMessageQueue;
|
||||
|
||||
public convertToLlm: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
|
||||
public transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
|
||||
public streamFunction: StreamFn;
|
||||
public getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
|
||||
public onPayload?: SimpleStreamOptions["onPayload"];
|
||||
public onResponse?: SimpleStreamOptions["onResponse"];
|
||||
public beforeToolCall?: (
|
||||
context: BeforeToolCallContext,
|
||||
signal?: AbortSignal,
|
||||
) => Promise<BeforeToolCallResult | undefined>;
|
||||
public afterToolCall?: (
|
||||
context: AfterToolCallContext,
|
||||
signal?: AbortSignal,
|
||||
) => Promise<AfterToolCallResult | undefined>;
|
||||
public shouldStopAfterTurn?: (
|
||||
context: ShouldStopAfterTurnContext,
|
||||
signal?: AbortSignal,
|
||||
) => boolean | Promise<boolean>;
|
||||
public prepareNextTurn?: (
|
||||
signal?: AbortSignal,
|
||||
) => Promise<AgentLoopTurnUpdate | undefined> | AgentLoopTurnUpdate | undefined;
|
||||
public prepareNextTurnWithContext?: (
|
||||
context: PrepareNextTurnContext,
|
||||
signal?: AbortSignal,
|
||||
) => Promise<AgentLoopTurnUpdate | undefined> | AgentLoopTurnUpdate | undefined;
|
||||
private activeRun?: ActiveRun;
|
||||
/** Session identifier forwarded to providers for cache-aware backends. */
|
||||
public sessionId?: string;
|
||||
/** Optional per-level thinking token budgets forwarded to the stream function. */
|
||||
public thinkingBudgets?: ThinkingBudgets;
|
||||
/** Preferred transport forwarded to the stream function. */
|
||||
public transport: Transport;
|
||||
/** Optional cap for provider-requested retry delays. */
|
||||
public maxRetryDelayMs?: number;
|
||||
/** Tool execution strategy for assistant messages that contain multiple tool calls. */
|
||||
public toolExecution: ToolExecutionMode;
|
||||
|
||||
constructor(options: AgentOptions) {
|
||||
// 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.shouldStopAfterTurn = runtimeOptions.shouldStopAfterTurn;
|
||||
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";
|
||||
}
|
||||
|
||||
/**
|
||||
* Subscribe to agent lifecycle events.
|
||||
*
|
||||
* Listener promises are awaited in subscription order and are included in
|
||||
* the current run's settlement. Listeners also receive the active abort
|
||||
* signal for the current run.
|
||||
*
|
||||
* `agent_end` is the final emitted event for a run, but the agent does not
|
||||
* become idle until all awaited listeners for that event have settled.
|
||||
*/
|
||||
subscribe(listener: (event: AgentEvent, signal: AbortSignal) => Promise<void> | void): () => void {
|
||||
this.listeners.add(listener);
|
||||
return () => this.listeners.delete(listener);
|
||||
}
|
||||
|
||||
/**
|
||||
* Current agent state.
|
||||
*
|
||||
* Assigning `state.tools` or `state.messages` copies the provided top-level array.
|
||||
*/
|
||||
get state(): AgentState {
|
||||
return this._state;
|
||||
}
|
||||
|
||||
/** Controls how queued steering messages are drained. */
|
||||
set steeringMode(mode: QueueMode) {
|
||||
this.steeringQueue.mode = mode;
|
||||
}
|
||||
|
||||
get steeringMode(): QueueMode {
|
||||
return this.steeringQueue.mode;
|
||||
}
|
||||
|
||||
/** Controls how queued follow-up messages are drained. */
|
||||
set followUpMode(mode: QueueMode) {
|
||||
this.followUpQueue.mode = mode;
|
||||
}
|
||||
|
||||
get followUpMode(): QueueMode {
|
||||
return this.followUpQueue.mode;
|
||||
}
|
||||
|
||||
/** Queue a message to be injected after the current assistant turn finishes. */
|
||||
steer(message: AgentMessage): void {
|
||||
this.steeringQueue.enqueue(message);
|
||||
}
|
||||
|
||||
/** Queue a message to run only after the agent would otherwise stop. */
|
||||
followUp(message: AgentMessage): void {
|
||||
this.followUpQueue.enqueue(message);
|
||||
}
|
||||
|
||||
/** Remove all queued steering messages. */
|
||||
clearSteeringQueue(): void {
|
||||
this.steeringQueue.clear();
|
||||
}
|
||||
|
||||
/** Remove all queued follow-up messages. */
|
||||
clearFollowUpQueue(): void {
|
||||
this.followUpQueue.clear();
|
||||
}
|
||||
|
||||
/** Remove all queued steering and follow-up messages. */
|
||||
clearAllQueues(): void {
|
||||
this.clearSteeringQueue();
|
||||
this.clearFollowUpQueue();
|
||||
}
|
||||
|
||||
/** Returns true when either queue still contains pending messages. */
|
||||
hasQueuedMessages(): boolean {
|
||||
return this.steeringQueue.hasItems() || this.followUpQueue.hasItems();
|
||||
}
|
||||
|
||||
/** Active abort signal for the current run, if any. */
|
||||
get signal(): AbortSignal | undefined {
|
||||
return this.activeRun?.abortController.signal;
|
||||
}
|
||||
|
||||
/** Abort the current run, if one is active. */
|
||||
abort(): void {
|
||||
this.activeRun?.abortController.abort();
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve when the current run and all awaited event listeners have finished.
|
||||
*
|
||||
* This resolves after `agent_end` listeners settle.
|
||||
*/
|
||||
waitForIdle(): Promise<void> {
|
||||
return this.activeRun?.promise ?? Promise.resolve();
|
||||
}
|
||||
|
||||
/** Clear transcript state, runtime state, and queued messages. */
|
||||
reset(): void {
|
||||
if (this.activeRun) {
|
||||
throw new Error("Agent is already processing. Wait for completion before resetting.");
|
||||
}
|
||||
|
||||
this._state.messages = [];
|
||||
this._state.isStreaming = false;
|
||||
this._state.streamingMessage = undefined;
|
||||
this._state.pendingToolCalls = new Set<string>();
|
||||
this._state.errorMessage = undefined;
|
||||
this.clearFollowUpQueue();
|
||||
this.clearSteeringQueue();
|
||||
}
|
||||
|
||||
/** Start a new prompt from text, a single message, or a batch of messages. */
|
||||
async prompt(message: AgentMessage | AgentMessage[]): Promise<void>;
|
||||
async prompt(input: string, images?: ImageContent[]): Promise<void>;
|
||||
async prompt(input: string | AgentMessage | AgentMessage[], images?: ImageContent[]): Promise<void> {
|
||||
if (this.activeRun) {
|
||||
throw new Error(
|
||||
"Agent is already processing a prompt. Use steer() or followUp() to queue messages, or wait for completion.",
|
||||
);
|
||||
}
|
||||
const messages = this.normalizePromptInput(input, images);
|
||||
await this.runPromptMessages(messages);
|
||||
}
|
||||
|
||||
/** Continue from the current transcript. The last message must be a user or tool-result message. */
|
||||
async continue(): Promise<void> {
|
||||
if (this.activeRun) {
|
||||
throw new Error("Agent is already processing. Wait for completion before continuing.");
|
||||
}
|
||||
|
||||
const lastMessage = this._state.messages[this._state.messages.length - 1];
|
||||
if (!lastMessage) {
|
||||
throw new Error("No messages to continue from");
|
||||
}
|
||||
|
||||
if (lastMessage.role === "assistant") {
|
||||
const queuedSteering = this.steeringQueue.drain();
|
||||
if (queuedSteering.length > 0) {
|
||||
await this.runPromptMessages(queuedSteering, { skipInitialSteeringPoll: true });
|
||||
return;
|
||||
}
|
||||
|
||||
const queuedFollowUps = this.followUpQueue.drain();
|
||||
if (queuedFollowUps.length > 0) {
|
||||
await this.runPromptMessages(queuedFollowUps);
|
||||
return;
|
||||
}
|
||||
|
||||
throw new Error("Cannot continue from message role: assistant");
|
||||
}
|
||||
|
||||
await this.runContinuation();
|
||||
}
|
||||
|
||||
private normalizePromptInput(
|
||||
input: string | AgentMessage | AgentMessage[],
|
||||
images?: ImageContent[],
|
||||
): AgentMessage[] {
|
||||
if (Array.isArray(input)) {
|
||||
return input;
|
||||
}
|
||||
|
||||
if (typeof input !== "string") {
|
||||
return [input];
|
||||
}
|
||||
|
||||
const content: Array<TextContent | ImageContent> = [{ type: "text", text: input }];
|
||||
if (images && images.length > 0) {
|
||||
content.push(...images);
|
||||
}
|
||||
return [{ role: "user", content, timestamp: Date.now() }];
|
||||
}
|
||||
|
||||
private async runPromptMessages(
|
||||
messages: AgentMessage[],
|
||||
options: { skipInitialSteeringPoll?: boolean } = {},
|
||||
): Promise<void> {
|
||||
await this.runWithLifecycle(async (signal) => {
|
||||
await runAgentLoop(
|
||||
messages,
|
||||
this.createContextSnapshot(),
|
||||
this.createLoopConfig(options),
|
||||
(event) => this.processEvents(event),
|
||||
signal,
|
||||
this.streamFunction,
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
private async runContinuation(): Promise<void> {
|
||||
await this.runWithLifecycle(async (signal) => {
|
||||
await runAgentLoopContinue(
|
||||
this.createContextSnapshot(),
|
||||
this.createLoopConfig(),
|
||||
(event) => this.processEvents(event),
|
||||
signal,
|
||||
this.streamFunction,
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
private createContextSnapshot(): AgentContext {
|
||||
return {
|
||||
systemPrompt: this._state.systemPrompt,
|
||||
messages: this._state.messages.slice(),
|
||||
tools: this._state.tools.slice(),
|
||||
};
|
||||
}
|
||||
|
||||
private createLoopConfig(options: { skipInitialSteeringPoll?: boolean } = {}): AgentLoopConfig {
|
||||
let skipInitialSteeringPoll = options.skipInitialSteeringPoll === true;
|
||||
const shouldStopAfterTurn = this.shouldStopAfterTurn;
|
||||
return {
|
||||
model: this._state.model,
|
||||
reasoning: this._state.thinkingLevel === "off" ? undefined : this._state.thinkingLevel,
|
||||
sessionId: this.sessionId,
|
||||
onPayload: this.onPayload,
|
||||
onResponse: this.onResponse,
|
||||
transport: this.transport,
|
||||
thinkingBudgets: this.thinkingBudgets,
|
||||
maxRetryDelayMs: this.maxRetryDelayMs,
|
||||
toolExecution: this.toolExecution,
|
||||
beforeToolCall: this.beforeToolCall,
|
||||
afterToolCall: this.afterToolCall,
|
||||
shouldStopAfterTurn: shouldStopAfterTurn
|
||||
? async (context) => await shouldStopAfterTurn(context, this.signal)
|
||||
: undefined,
|
||||
prepareNextTurn:
|
||||
this.prepareNextTurnWithContext || this.prepareNextTurn
|
||||
? async (context) => {
|
||||
if (this.prepareNextTurnWithContext) {
|
||||
return await this.prepareNextTurnWithContext(context, this.signal);
|
||||
}
|
||||
return await this.prepareNextTurn?.(this.signal);
|
||||
}
|
||||
: undefined,
|
||||
convertToLlm: this.convertToLlm,
|
||||
transformContext: this.transformContext,
|
||||
getApiKey: this.getApiKey,
|
||||
getSteeringMessages: async () => {
|
||||
if (skipInitialSteeringPoll) {
|
||||
skipInitialSteeringPoll = false;
|
||||
return [];
|
||||
}
|
||||
return this.steeringQueue.drain();
|
||||
},
|
||||
getFollowUpMessages: async () => this.followUpQueue.drain(),
|
||||
};
|
||||
}
|
||||
|
||||
private async runWithLifecycle(executor: (signal: AbortSignal) => Promise<void>): Promise<void> {
|
||||
if (this.activeRun) {
|
||||
throw new Error("Agent is already processing.");
|
||||
}
|
||||
|
||||
const abortController = new AbortController();
|
||||
let resolvePromise = () => {};
|
||||
const promise = new Promise<void>((resolve) => {
|
||||
resolvePromise = resolve;
|
||||
});
|
||||
this.activeRun = { promise, resolve: resolvePromise, abortController };
|
||||
|
||||
this._state.isStreaming = true;
|
||||
this._state.streamingMessage = undefined;
|
||||
this._state.errorMessage = undefined;
|
||||
|
||||
try {
|
||||
await executor(abortController.signal);
|
||||
} catch (error) {
|
||||
await this.handleRunFailure(error, abortController.signal.aborted);
|
||||
} finally {
|
||||
this.finishRun();
|
||||
}
|
||||
}
|
||||
|
||||
private async handleRunFailure(error: unknown, aborted: boolean): Promise<void> {
|
||||
const failureMessage = {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text: "" }],
|
||||
api: this._state.model.api,
|
||||
provider: this._state.model.provider,
|
||||
model: this._state.model.id,
|
||||
usage: EMPTY_USAGE,
|
||||
stopReason: aborted ? "aborted" : "error",
|
||||
errorMessage: error instanceof Error ? error.message : String(error),
|
||||
timestamp: Date.now(),
|
||||
} satisfies AgentMessage;
|
||||
await this.processEvents({ type: "message_start", message: failureMessage });
|
||||
await this.processEvents({ type: "message_end", message: failureMessage });
|
||||
await this.processEvents({ type: "turn_end", message: failureMessage, toolResults: [] });
|
||||
await this.processEvents({ type: "agent_end", messages: [failureMessage] });
|
||||
}
|
||||
|
||||
private finishRun(): void {
|
||||
this._state.isStreaming = false;
|
||||
this._state.streamingMessage = undefined;
|
||||
this._state.pendingToolCalls = new Set<string>();
|
||||
this.activeRun?.resolve();
|
||||
this.activeRun = undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Reduce internal state for a loop event, then await listeners.
|
||||
*
|
||||
* `agent_end` only means no further loop events will be emitted. The run is
|
||||
* considered idle later, after all awaited listeners for `agent_end` finish
|
||||
* and `finishRun()` clears runtime-owned state.
|
||||
*/
|
||||
private async processEvents(event: AgentEvent): Promise<void> {
|
||||
switch (event.type) {
|
||||
case "message_start":
|
||||
this._state.streamingMessage = event.message;
|
||||
break;
|
||||
|
||||
case "message_update":
|
||||
this._state.streamingMessage = event.message;
|
||||
break;
|
||||
|
||||
case "message_end":
|
||||
this._state.streamingMessage = undefined;
|
||||
this._state.messages.push(event.message);
|
||||
break;
|
||||
|
||||
case "tool_execution_start": {
|
||||
const pendingToolCalls = new Set(this._state.pendingToolCalls);
|
||||
pendingToolCalls.add(event.toolCallId);
|
||||
this._state.pendingToolCalls = pendingToolCalls;
|
||||
break;
|
||||
}
|
||||
|
||||
case "tool_execution_end": {
|
||||
const pendingToolCalls = new Set(this._state.pendingToolCalls);
|
||||
pendingToolCalls.delete(event.toolCallId);
|
||||
this._state.pendingToolCalls = pendingToolCalls;
|
||||
break;
|
||||
}
|
||||
|
||||
case "turn_end":
|
||||
if (event.message.role === "assistant" && event.message.errorMessage) {
|
||||
this._state.errorMessage = event.message.errorMessage;
|
||||
}
|
||||
break;
|
||||
|
||||
case "agent_end":
|
||||
this._state.streamingMessage = undefined;
|
||||
break;
|
||||
}
|
||||
|
||||
const signal = this.activeRun?.abortController.signal;
|
||||
if (!signal) {
|
||||
throw new Error("Agent listener invoked outside active run");
|
||||
}
|
||||
for (const listener of this.listeners) {
|
||||
await listener(event, signal);
|
||||
}
|
||||
}
|
||||
}
|
||||
622
third_party/pi/packages/agent/src/harness/agent-harness.ts
vendored
Normal file
622
third_party/pi/packages/agent/src/harness/agent-harness.ts
vendored
Normal file
@@ -0,0 +1,622 @@
|
||||
import type { JsonRepresentation } from "@earendil-works/chord";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
AssistantMessageEvent,
|
||||
AssistantMessageFrame,
|
||||
DeferredHandle,
|
||||
ImageContent,
|
||||
Message,
|
||||
Model,
|
||||
Models,
|
||||
RetryPolicy,
|
||||
ToolResultMessage,
|
||||
Usage,
|
||||
} from "@earendil-works/pi-ai";
|
||||
import type { AgentMessage, AgentToolResult, QueueMode, ThinkingLevel } from "../types.ts";
|
||||
import type { BranchPreparation, BranchSummaryResult } from "./compaction/branch-summarization.ts";
|
||||
import type { CompactionPreparation, CompactionSettings, CompactResult } from "./compaction/compaction.ts";
|
||||
import type { Context } from "./context.ts";
|
||||
import type {
|
||||
Closed,
|
||||
InvalidMessage,
|
||||
InvalidNavigation,
|
||||
LaneBusy,
|
||||
NoActiveOperation,
|
||||
NothingToCompact,
|
||||
NothingToResume,
|
||||
OperationMismatch,
|
||||
Result,
|
||||
UnknownSkill,
|
||||
UnknownTarget,
|
||||
UnknownTemplate,
|
||||
} from "./result.ts";
|
||||
|
||||
export {
|
||||
Closed,
|
||||
HarnessClosed,
|
||||
HarnessFault,
|
||||
InvalidLane,
|
||||
InvalidMessage,
|
||||
InvalidNavigation,
|
||||
LaneBusy,
|
||||
NoActiveOperation,
|
||||
NoActiveRun,
|
||||
NothingToCompact,
|
||||
NothingToResume,
|
||||
OperationMismatch,
|
||||
UnknownSkill,
|
||||
UnknownTarget,
|
||||
UnknownTemplate,
|
||||
} from "./result.ts";
|
||||
export { SliceNotImplemented } from "./runtime/types.ts";
|
||||
|
||||
import { createAgentHarness } from "./runtime/harness.ts";
|
||||
import type {
|
||||
BranchScan,
|
||||
Entry,
|
||||
EntryProjector,
|
||||
JsonValue,
|
||||
LaneConfiguration,
|
||||
OperationError,
|
||||
OperationResultRecord,
|
||||
Session,
|
||||
SessionStats,
|
||||
SettledAssistantMessage,
|
||||
UsageRow,
|
||||
} from "./session/types.ts";
|
||||
import type {
|
||||
AgentHarnessResources,
|
||||
AgentHarnessStreamOptions,
|
||||
AgentHarnessStreamOptionsPatch,
|
||||
AgentHarnessTool,
|
||||
PromptTemplate,
|
||||
Skill,
|
||||
} from "./types.ts";
|
||||
|
||||
/** Convenience-only suspended run observation, constructed when M8 exposes public drive. */
|
||||
export interface SuspendedRun {
|
||||
operationId: string;
|
||||
status: "suspended";
|
||||
deferred: DeferredHandle;
|
||||
}
|
||||
|
||||
export type RunResult = Result<
|
||||
OperationResultRecord | SuspendedRun,
|
||||
LaneBusy | InvalidMessage | UnknownSkill | UnknownTemplate | Closed
|
||||
>;
|
||||
export type CompactionResult = Result<
|
||||
{ compaction: OperationResultRecord; run?: OperationResultRecord | SuspendedRun },
|
||||
LaneBusy | NothingToCompact | Closed
|
||||
>;
|
||||
export type NavigationResult = Result<
|
||||
{ navigation: OperationResultRecord; run?: OperationResultRecord | SuspendedRun },
|
||||
LaneBusy | InvalidNavigation | UnknownTarget | Closed
|
||||
>;
|
||||
export type ResumeResult = Result<OperationResultRecord | SuspendedRun, NothingToResume | Closed>;
|
||||
export type QueueResult = Result<{ entryId: string }, InvalidMessage | Closed>;
|
||||
export type CancelQueuedResult = Result<{ kind: "cancelled" | "already_consumed" | "not_found" }, Closed>;
|
||||
export type AbortResult = Result<
|
||||
{ operationId: string; steer: AgentMessage[]; followUp: AgentMessage[] },
|
||||
NoActiveOperation | Closed
|
||||
>;
|
||||
export type RecordUsageResult = Result<{ usageId: string }, Closed>;
|
||||
|
||||
export interface NavigateOptions {
|
||||
summarize?: boolean;
|
||||
label?: string;
|
||||
customInstructions?: string;
|
||||
}
|
||||
|
||||
export type OperationRequest =
|
||||
| { kind: "prompt"; operationId?: string; prompt: string; images?: ImageContent[] }
|
||||
| { kind: "prompt"; operationId?: string; prompt: AgentMessage | AgentMessage[]; images?: never }
|
||||
| { kind: "skill"; operationId?: string; name: string; additionalInstructions?: string }
|
||||
| { kind: "prompt_template"; operationId?: string; name: string; args?: string[] }
|
||||
| { kind: "compaction"; operationId?: string; customInstructions?: string }
|
||||
| { kind: "navigation"; operationId?: string; targetId: string | null; options?: NavigateOptions };
|
||||
|
||||
export interface OperationAdmission {
|
||||
operationId: string;
|
||||
kind: "run" | "compaction" | "navigation";
|
||||
startedAt: number;
|
||||
}
|
||||
|
||||
export type OperationAdmissionError =
|
||||
| LaneBusy
|
||||
| InvalidMessage
|
||||
| UnknownSkill
|
||||
| UnknownTemplate
|
||||
| NothingToCompact
|
||||
| InvalidNavigation
|
||||
| UnknownTarget
|
||||
| Closed;
|
||||
export type OperationAdmissionResult = Result<OperationAdmission, OperationAdmissionError>;
|
||||
|
||||
export interface DriveOptions {
|
||||
operationId: string;
|
||||
waitForRetry?: boolean;
|
||||
pollDeferred?: boolean;
|
||||
}
|
||||
|
||||
export interface ModelIdentity {
|
||||
provider: string;
|
||||
modelId: string;
|
||||
}
|
||||
|
||||
export type OperationStatus = "running" | "open" | "aborting";
|
||||
|
||||
export interface CurrentOperationInfo {
|
||||
id: string;
|
||||
kind: "run" | "compaction" | "navigation";
|
||||
startedAt: number;
|
||||
status: OperationStatus;
|
||||
capturedModel?: ModelIdentity;
|
||||
}
|
||||
|
||||
export interface LaneExecutionInfo {
|
||||
lane: string;
|
||||
tipId: string | null;
|
||||
configuredModel: ModelIdentity;
|
||||
current: CurrentOperationInfo | null;
|
||||
lastOperationId: string | null;
|
||||
}
|
||||
|
||||
export type DriveOutcome =
|
||||
| { kind: "settled"; outcome: OperationResultRecord }
|
||||
| { kind: "waiting"; operationId: string; reason: "retry"; notBefore: number }
|
||||
| { kind: "waiting"; operationId: string; reason: "deferred"; deferred: DeferredHandle };
|
||||
export type DriveResult = Result<DriveOutcome, OperationMismatch | Closed>;
|
||||
|
||||
export type AbortRequestResult = Result<
|
||||
{
|
||||
operationId: string;
|
||||
newlyRequested: boolean;
|
||||
steer: AgentMessage[];
|
||||
followUp: AgentMessage[];
|
||||
},
|
||||
OperationMismatch | Closed
|
||||
>;
|
||||
|
||||
export interface WatchHandle<T> {
|
||||
snapshot: T;
|
||||
start(listener: EventListener): void;
|
||||
resnapshot(context: Context): Promise<T>;
|
||||
unsubscribe(): void;
|
||||
}
|
||||
|
||||
export interface LaneInfo {
|
||||
name: string;
|
||||
tipId: string | null;
|
||||
operation: CurrentOperationInfo | null;
|
||||
}
|
||||
|
||||
export type LaneSnapshotTool =
|
||||
| {
|
||||
status: "running";
|
||||
toolCallId: string;
|
||||
toolName: string;
|
||||
args: unknown;
|
||||
result?: AgentToolResult<unknown>;
|
||||
}
|
||||
| {
|
||||
status: "settled";
|
||||
toolCallId: string;
|
||||
toolName: string;
|
||||
args: unknown;
|
||||
result: AgentToolResult<unknown>;
|
||||
isError: boolean;
|
||||
};
|
||||
|
||||
export interface OpenOperation {
|
||||
lane: string;
|
||||
operationId: string;
|
||||
kind: "run" | "compaction" | "navigation";
|
||||
startedAt: number;
|
||||
aborting?: true;
|
||||
}
|
||||
|
||||
export type LaneQueuedItem =
|
||||
| {
|
||||
entryId: string;
|
||||
kind: "steer" | "followUp" | "nextRun" | "write";
|
||||
type: "message";
|
||||
message: AgentMessage;
|
||||
}
|
||||
| { entryId: string; kind: "write"; type: "custom"; customType: string; data?: JsonValue };
|
||||
|
||||
export interface LaneSnapshot {
|
||||
lane: string;
|
||||
transcript: Entry[];
|
||||
tipId: string | null;
|
||||
lastResult?: OperationResultRecord;
|
||||
configuration: LaneConfiguration;
|
||||
stats: SessionStats;
|
||||
operation: null | {
|
||||
id: string;
|
||||
kind: "run" | "compaction" | "navigation";
|
||||
startedAt: number;
|
||||
fromTipId: string | null;
|
||||
status: OperationStatus;
|
||||
retry?: { attempt: number; maxAttempts: number; nextAttemptAt: number };
|
||||
deferred?: { handle: DeferredHandle; poll: number };
|
||||
streamingMessage?: AssistantMessage;
|
||||
runningTools: LaneSnapshotTool[];
|
||||
};
|
||||
queues: LaneQueuedItem[];
|
||||
faulted: boolean;
|
||||
}
|
||||
|
||||
export interface SessionSnapshot {
|
||||
lanes: LaneInfo[];
|
||||
faulted: boolean;
|
||||
}
|
||||
|
||||
export type HarnessEventPayload =
|
||||
| { type: "run_start"; runId: string; startedAt: number }
|
||||
| { type: "run_resume"; runId: string }
|
||||
| { type: "run_suspend"; runId: string; reason: "deferred"; deferred: DeferredHandle; poll: number }
|
||||
| { type: "operation_abort"; operationId: string; steer: AgentMessage[]; followUp: AgentMessage[] }
|
||||
| ({ type: "run_end"; runId: string; fromTipId: string | null; tipId: string | null; endedAt: number } & (
|
||||
| { status: "completed" | "aborted"; error?: never }
|
||||
| { status: "failed"; error: OperationError }
|
||||
))
|
||||
| { type: "fault"; code: string; message: string }
|
||||
| ({ type: "handler_error"; error: string; stack?: string } & (
|
||||
| { kind: "hook"; hook: string }
|
||||
| { kind: "event"; event: string }
|
||||
))
|
||||
| { type: "turn_start"; runId: string; turnId: string }
|
||||
| {
|
||||
type: "turn_end";
|
||||
runId: string;
|
||||
turnId: string;
|
||||
message: AssistantMessage;
|
||||
toolResults: ToolResultMessage[];
|
||||
}
|
||||
| {
|
||||
type: "retry_scheduled";
|
||||
runId: string;
|
||||
step: string;
|
||||
attempt: number;
|
||||
maxAttempts: number;
|
||||
delayMs: number;
|
||||
notBefore: number;
|
||||
errorMessage: string;
|
||||
}
|
||||
| { type: "retry_start"; runId: string; step: string; attempt: number }
|
||||
| {
|
||||
type: "retry_end";
|
||||
runId: string;
|
||||
step: string;
|
||||
attempt: number;
|
||||
success: boolean;
|
||||
finalError?: string;
|
||||
}
|
||||
| { type: "message_start"; runId?: string; message: AgentMessage }
|
||||
| {
|
||||
type: "message_update";
|
||||
runId: string;
|
||||
message: AgentMessage;
|
||||
event: AssistantMessageEvent;
|
||||
frame?: AssistantMessageFrame;
|
||||
}
|
||||
| { type: "message_end"; runId?: string; message: AgentMessage; entryId?: string }
|
||||
| {
|
||||
type: "tool_start";
|
||||
runId: string;
|
||||
turnId: string;
|
||||
toolCallId: string;
|
||||
toolName: string;
|
||||
args: unknown;
|
||||
}
|
||||
| {
|
||||
type: "tool_update";
|
||||
runId: string;
|
||||
turnId: string;
|
||||
toolCallId: string;
|
||||
toolName: string;
|
||||
partialResult: AgentToolResult<unknown>;
|
||||
}
|
||||
| {
|
||||
type: "tool_end";
|
||||
runId: string;
|
||||
turnId: string;
|
||||
toolCallId: string;
|
||||
toolName: string;
|
||||
result: AgentToolResult<unknown>;
|
||||
isError: boolean;
|
||||
terminate: boolean;
|
||||
}
|
||||
| { type: "entry_added"; entry: Entry }
|
||||
| { type: "queue_update"; queues: LaneQueuedItem[] }
|
||||
| ({ type: "value_update" } & (
|
||||
| { value: "session_name"; name: string | undefined }
|
||||
| { value: "entry_label"; targetId: string; label: string | undefined }
|
||||
))
|
||||
| ({ type: "config_update" } & (
|
||||
| {
|
||||
property: "model";
|
||||
value: { provider: string; modelId: string };
|
||||
previous: unknown;
|
||||
}
|
||||
| { property: "thinkingLevel"; value: ThinkingLevel; previous: ThinkingLevel }
|
||||
| { property: "activeTools"; value: string[]; previous: string[] }
|
||||
| { property: "tools" | "resources" }
|
||||
| {
|
||||
property: "streamOptions";
|
||||
value: AgentHarnessStreamOptions;
|
||||
previous: AgentHarnessStreamOptions;
|
||||
}
|
||||
| { property: "retryPolicy"; value: RetryPolicy; previous: RetryPolicy }
|
||||
| { property: "compactionSettings"; value: CompactionSettings; previous: CompactionSettings }
|
||||
| { property: "steeringMode"; value: QueueMode; previous: QueueMode }
|
||||
| { property: "followUpMode"; value: QueueMode; previous: QueueMode }
|
||||
))
|
||||
| {
|
||||
type: "compaction_start";
|
||||
runId: string;
|
||||
reason: "manual" | "threshold" | "overflow";
|
||||
startedAt: number;
|
||||
}
|
||||
| ({ type: "compaction_end"; runId: string; reason: "manual" | "threshold" | "overflow"; endedAt: number } & (
|
||||
| { status: "completed"; entryId: string; error?: never }
|
||||
| { status: "declined" | "aborted"; entryId?: never; error?: never }
|
||||
| { status: "failed"; entryId?: never; error: OperationError }
|
||||
))
|
||||
| { type: "navigation_start"; runId: string; targetId: string | null; startedAt: number }
|
||||
| ({ type: "navigation_end"; runId: string; fromTipId: string | null; tipId: string | null; endedAt: number } & (
|
||||
| { status: "completed" | "declined" | "aborted"; error?: never }
|
||||
| { status: "failed"; error: OperationError }
|
||||
))
|
||||
| { type: "lane_created"; at: string | null }
|
||||
| { type: "usage"; lane: string; row: UsageRow; totals: Usage };
|
||||
|
||||
export type SpecialEventPayload = Extract<
|
||||
HarnessEventPayload,
|
||||
{ type: "fault" | "value_update" | "usage" | "config_update" | "handler_error" }
|
||||
>;
|
||||
export type LaneEventPayload = Exclude<HarnessEventPayload, SpecialEventPayload>;
|
||||
export type ConfigEventPayload = Extract<HarnessEventPayload, { type: "config_update" }>;
|
||||
export type LaneConfigEventPayload = Extract<
|
||||
ConfigEventPayload,
|
||||
{ property: "model" | "thinkingLevel" | "activeTools" }
|
||||
>;
|
||||
export type GlobalConfigEventPayload = Exclude<ConfigEventPayload, LaneConfigEventPayload>;
|
||||
export type HandlerErrorPayload = Extract<HarnessEventPayload, { type: "handler_error" }>;
|
||||
|
||||
export type HarnessEvent =
|
||||
| (LaneEventPayload & { lane: string; recovery?: true })
|
||||
| (LaneConfigEventPayload & { lane: string; recovery?: true })
|
||||
| (Extract<HarnessEventPayload, { type: "fault" | "value_update" }> & {
|
||||
lane?: never;
|
||||
recovery?: never;
|
||||
})
|
||||
| (Extract<HarnessEventPayload, { type: "usage" }> & { recovery?: never })
|
||||
| (GlobalConfigEventPayload & { lane?: never; recovery?: never })
|
||||
| (HandlerErrorPayload & ({ lane: string; recovery?: true } | { lane?: never; recovery?: never }));
|
||||
|
||||
type LaneWatchSourceEvent =
|
||||
| Exclude<
|
||||
HarnessEvent,
|
||||
| { type: "handler_error" | "turn_start" | "turn_end" | "value_update" | "lane_created" | "message_update" }
|
||||
| ({ type: "config_update" } & { property: string })
|
||||
>
|
||||
| Extract<HarnessEvent, { type: "config_update"; property: "model" | "thinkingLevel" | "activeTools" }>
|
||||
| Omit<Extract<HarnessEvent, { type: "message_update" }>, "event">;
|
||||
|
||||
/** Strict-JSON snapshot representation published to remote transcript consumers. */
|
||||
export type LaneTranscriptSnapshot = JsonRepresentation<LaneSnapshot>;
|
||||
/** Reducer-relevant strict-JSON Harness events published to remote transcript consumers. */
|
||||
export type LaneWatchEvent = JsonRepresentation<LaneWatchSourceEvent>;
|
||||
|
||||
export type HarnessEventType = HarnessEvent["type"];
|
||||
export type EventListener<TEvent extends HarnessEvent = HarnessEvent> = (
|
||||
event: TEvent,
|
||||
context: Context,
|
||||
) => void | Promise<void>;
|
||||
|
||||
export interface Events {
|
||||
on<TType extends HarnessEventType>(
|
||||
type: TType,
|
||||
listener: EventListener<Extract<HarnessEvent, { type: TType }>>,
|
||||
): () => void;
|
||||
}
|
||||
|
||||
export type Resources = AgentHarnessResources<Skill, PromptTemplate>;
|
||||
|
||||
type VoidHookResult = ReturnType<() => void>;
|
||||
|
||||
export interface HookMap {
|
||||
before_run: {
|
||||
event: { prompt: AgentMessage[]; resources: Resources };
|
||||
result: { messages?: AgentMessage[] } | undefined;
|
||||
};
|
||||
before_drive: {
|
||||
event: { operation: "run" | "compaction" | "navigation" };
|
||||
result: VoidHookResult;
|
||||
};
|
||||
before_run_end: {
|
||||
event: { runId: string; messages: AgentMessage[] };
|
||||
result: { followUp?: string } | undefined;
|
||||
};
|
||||
transform_context: {
|
||||
event: { messages: AgentMessage[]; systemPrompt: string };
|
||||
result: { messages?: AgentMessage[]; systemPrompt?: string } | undefined;
|
||||
};
|
||||
before_request: {
|
||||
event: {
|
||||
model: Model<Api>;
|
||||
step: "assistant" | "deferred" | "compaction" | "branch_summary";
|
||||
attempt: number;
|
||||
streamOptions: AgentHarnessStreamOptions;
|
||||
};
|
||||
result: { streamOptions?: AgentHarnessStreamOptionsPatch } | undefined;
|
||||
};
|
||||
before_payload: {
|
||||
event: { model: Model<Api>; payload: unknown };
|
||||
result: { payload: unknown } | undefined;
|
||||
};
|
||||
after_response: {
|
||||
event: { status?: number; headers?: Record<string, string>; message: SettledAssistantMessage };
|
||||
result: { message?: SettledAssistantMessage } | undefined;
|
||||
};
|
||||
before_tool: {
|
||||
event: { toolCallId: string; toolName: string; args: Record<string, JsonValue> };
|
||||
result: { args?: Record<string, JsonValue>; block?: { reason: string; terminate?: boolean } } | undefined;
|
||||
};
|
||||
after_tool: {
|
||||
event: {
|
||||
toolCallId: string;
|
||||
toolName: string;
|
||||
args: Record<string, JsonValue>;
|
||||
content: AgentToolResult<unknown>["content"];
|
||||
details?: JsonValue;
|
||||
isError: boolean;
|
||||
usage?: Usage;
|
||||
};
|
||||
result:
|
||||
| {
|
||||
content?: AgentToolResult<unknown>["content"];
|
||||
details?: JsonValue;
|
||||
isError?: boolean;
|
||||
usage?: Usage;
|
||||
terminate?: boolean;
|
||||
}
|
||||
| undefined;
|
||||
};
|
||||
before_compaction: {
|
||||
event: {
|
||||
reason: "manual" | "threshold" | "overflow";
|
||||
preparation: CompactionPreparation;
|
||||
customInstructions?: string;
|
||||
};
|
||||
result: { decline?: boolean; compaction?: CompactResult } | undefined;
|
||||
};
|
||||
before_navigation: {
|
||||
event: { targetId: string; preparation: BranchPreparation; customInstructions?: string };
|
||||
result: { decline?: boolean; summary?: BranchSummaryResult } | undefined;
|
||||
};
|
||||
}
|
||||
|
||||
export type HookName = keyof HookMap;
|
||||
export type HookInvocation<TName extends HookName> = HookMap[TName]["event"] & {
|
||||
lane: string;
|
||||
runId: string;
|
||||
};
|
||||
export type HookHandler<TName extends HookName> = (
|
||||
event: HookInvocation<TName>,
|
||||
context: Context,
|
||||
) => Promise<HookMap[TName]["result"]> | HookMap[TName]["result"];
|
||||
|
||||
export interface Hooks {
|
||||
on<TName extends HookName>(name: TName, handler: HookHandler<TName>, options?: { id?: string }): () => void;
|
||||
}
|
||||
|
||||
export type { EntryProjector } from "./session/types.ts";
|
||||
|
||||
export interface AgentHarnessOptions<TContext extends object | undefined = object | undefined> {
|
||||
session: Session;
|
||||
models: Models;
|
||||
model: Model<Api>;
|
||||
thinkingLevel?: ThinkingLevel;
|
||||
activeToolNames?: string[];
|
||||
tools?: AgentHarnessTool<TContext>[];
|
||||
toolContext?: TContext | ((context: Context) => TContext | Promise<TContext>);
|
||||
systemPrompt?: string | ((toolContext: TContext, context: Context) => string | Promise<string>);
|
||||
resources?: Resources;
|
||||
streamOptions?: AgentHarnessStreamOptions;
|
||||
retry?: RetryPolicy;
|
||||
compaction?: CompactionSettings;
|
||||
steeringMode?: QueueMode;
|
||||
followUpMode?: QueueMode;
|
||||
toolExecution?: "sequential" | "parallel";
|
||||
toProviderMessages?: (messages: AgentMessage[], context: Context) => Message[] | Promise<Message[]>;
|
||||
entryProjectors?: Record<string, EntryProjector>;
|
||||
}
|
||||
|
||||
export interface AgentLane {
|
||||
readonly name: string;
|
||||
getTipId(context: Context): Promise<string | null>;
|
||||
findEntries(query: BranchScan | undefined, context: Context): Promise<Entry[]>;
|
||||
findEntry(query: BranchScan | undefined, context: Context): Promise<Entry | undefined>;
|
||||
appendMessage(message: AgentMessage, context: Context): Promise<string>;
|
||||
appendCustomEntry(customType: string, data: JsonValue | undefined, context: Context): Promise<string>;
|
||||
getResult(operationId: string, context: Context): Promise<OperationResultRecord | undefined>;
|
||||
accept(request: OperationRequest, context: Context): Promise<OperationAdmissionResult>;
|
||||
drive(options: DriveOptions, context: Context): Promise<DriveResult>;
|
||||
requestAbort(operationId: string, context: Context): Promise<AbortRequestResult>;
|
||||
inspectExecution(context: Context): Promise<LaneExecutionInfo>;
|
||||
prompt(text: string, images: ImageContent[] | undefined, context: Context): Promise<RunResult>;
|
||||
prompt(message: AgentMessage | AgentMessage[], context: Context): Promise<RunResult>;
|
||||
skill(name: string, additionalInstructions: string | undefined, context: Context): Promise<RunResult>;
|
||||
promptFromTemplate(name: string, args: string[] | undefined, context: Context): Promise<RunResult>;
|
||||
compact(options: { customInstructions?: string } | undefined, context: Context): Promise<CompactionResult>;
|
||||
navigateTree(
|
||||
targetId: string | null,
|
||||
options: NavigateOptions | undefined,
|
||||
context: Context,
|
||||
): Promise<NavigationResult>;
|
||||
resume(context: Context): Promise<ResumeResult>;
|
||||
abort(context: Context): Promise<AbortResult>;
|
||||
steer(message: string | AgentMessage, images: ImageContent[] | undefined, context: Context): Promise<QueueResult>;
|
||||
followUp(message: string | AgentMessage, images: ImageContent[] | undefined, context: Context): Promise<QueueResult>;
|
||||
nextRun(message: string | AgentMessage, images: ImageContent[] | undefined, context: Context): Promise<QueueResult>;
|
||||
cancelQueued(entryId: string, context: Context): Promise<CancelQueuedResult>;
|
||||
recordUsage(
|
||||
usage: Usage,
|
||||
options: { entryId?: string; details?: JsonValue } | undefined,
|
||||
context: Context,
|
||||
): Promise<RecordUsageResult>;
|
||||
waitForIdle(context: Context): Promise<void>;
|
||||
runWhenIdle(callback: (context: Context) => void | Promise<void>, context: Context): Promise<void>;
|
||||
getModel(context: Context): Promise<Model<Api> | undefined>;
|
||||
setModel(model: ModelIdentity, context: Context): Promise<void>;
|
||||
getThinkingLevel(context: Context): Promise<ThinkingLevel>;
|
||||
setThinkingLevel(level: ThinkingLevel, context: Context): Promise<void>;
|
||||
getActiveTools(context: Context): Promise<string[]>;
|
||||
setActiveTools(names: string[], context: Context): Promise<void>;
|
||||
watch(context: Context): Promise<WatchHandle<LaneSnapshot>>;
|
||||
}
|
||||
|
||||
export interface AcquireLaneOptions {
|
||||
createAt?: string | null;
|
||||
}
|
||||
|
||||
export interface AgentHarness<TContext extends object | undefined = object | undefined> {
|
||||
lane(name: string, context: Context): Promise<AgentLane>;
|
||||
lane(name: string, options: AcquireLaneOptions, context: Context): Promise<AgentLane>;
|
||||
lanes(context: Context): Promise<LaneInfo[]>;
|
||||
getName(context: Context): Promise<string | undefined>;
|
||||
setName(name: string | undefined, context: Context): Promise<void>;
|
||||
getLabel(targetId: string, context: Context): Promise<string | undefined>;
|
||||
setLabel(targetId: string, label: string | undefined, context: Context): Promise<void>;
|
||||
getTools(context: Context): Promise<AgentHarnessTool<TContext>[]>;
|
||||
setTools(tools: AgentHarnessTool<TContext>[], context: Context): Promise<void>;
|
||||
getResources(context: Context): Promise<Resources>;
|
||||
setResources(resources: Resources, context: Context): Promise<void>;
|
||||
getStreamOptions(context: Context): Promise<AgentHarnessStreamOptions>;
|
||||
setStreamOptions(options: AgentHarnessStreamOptions, context: Context): Promise<void>;
|
||||
getRetryPolicy(context: Context): Promise<RetryPolicy>;
|
||||
setRetryPolicy(policy: RetryPolicy, context: Context): Promise<void>;
|
||||
getCompactionSettings(context: Context): Promise<CompactionSettings>;
|
||||
setCompactionSettings(settings: CompactionSettings, context: Context): Promise<void>;
|
||||
getSteeringMode(context: Context): Promise<QueueMode>;
|
||||
setSteeringMode(mode: QueueMode, context: Context): Promise<void>;
|
||||
getFollowUpMode(context: Context): Promise<QueueMode>;
|
||||
setFollowUpMode(mode: QueueMode, context: Context): Promise<void>;
|
||||
watchSession(context: Context): Promise<WatchHandle<SessionSnapshot>>;
|
||||
readonly hooks: Hooks;
|
||||
readonly events: Events;
|
||||
close(context: Context): Promise<void>;
|
||||
}
|
||||
|
||||
export interface AgentHarnessConstructor {
|
||||
create<TContext extends object | undefined = object | undefined>(
|
||||
options: AgentHarnessOptions<TContext>,
|
||||
context: Context,
|
||||
): Promise<{ harness: AgentHarness<TContext>; open: OpenOperation[] }>;
|
||||
}
|
||||
|
||||
/** Runtime constructor for attaching the durable harness to one open session. */
|
||||
export const AgentHarness = { create: createAgentHarness } satisfies AgentHarnessConstructor;
|
||||
300
third_party/pi/packages/agent/src/harness/compaction/branch-summarization.ts
vendored
Normal file
300
third_party/pi/packages/agent/src/harness/compaction/branch-summarization.ts
vendored
Normal file
@@ -0,0 +1,300 @@
|
||||
import {
|
||||
type Api,
|
||||
contentText,
|
||||
type Model,
|
||||
type Models,
|
||||
type RetryCallbacks,
|
||||
type RetryPolicy,
|
||||
type Usage,
|
||||
} from "@earendil-works/pi-ai";
|
||||
|
||||
import type { AgentMessage } from "../../types.ts";
|
||||
import type { Context } from "../context.ts";
|
||||
import { convertToLlm, createBranchSummaryMessage, createCompactionSummaryMessage } from "../messages.ts";
|
||||
import type { Branch, Entry, Session } from "../session/index.ts";
|
||||
import { BranchSummaryError, err, ok, type Result } from "../types.ts";
|
||||
import {
|
||||
completeSimpleWithRetries,
|
||||
createSummaryRequestOptions,
|
||||
estimateTokens,
|
||||
SUMMARIZATION_SYSTEM_PROMPT,
|
||||
type SummaryRequest,
|
||||
} from "./compaction.ts";
|
||||
import {
|
||||
computeFileLists,
|
||||
createFileOps,
|
||||
extractFileOpsFromMessage,
|
||||
type FileOperations,
|
||||
formatFileOperations,
|
||||
serializeConversation,
|
||||
} from "./utils.ts";
|
||||
|
||||
/** Generated branch summary data ready to be persisted as a branch-summary entry. */
|
||||
export interface BranchSummaryResult {
|
||||
summary: string;
|
||||
usage?: Usage;
|
||||
readFiles: string[];
|
||||
modifiedFiles: string[];
|
||||
}
|
||||
|
||||
/** File-operation details stored on generated branch summary entries. */
|
||||
export interface BranchSummaryDetails {
|
||||
/** Files read while exploring the summarized branch. */
|
||||
readFiles: string[];
|
||||
/** Files modified while exploring the summarized branch. */
|
||||
modifiedFiles: string[];
|
||||
}
|
||||
|
||||
export type { FileOperations } from "./utils.ts";
|
||||
|
||||
/** Prepared branch content for summarization. */
|
||||
export interface BranchPreparation {
|
||||
/** Messages selected for the branch summary. */
|
||||
messages: AgentMessage[];
|
||||
/** File operations extracted from the branch. */
|
||||
fileOps: FileOperations;
|
||||
/** Estimated token count for selected messages. */
|
||||
totalTokens: number;
|
||||
}
|
||||
|
||||
/** Entries selected for branch summarization. */
|
||||
export interface CollectEntriesResult {
|
||||
/** Entries to summarize in chronological order. */
|
||||
entries: Entry[];
|
||||
/** Deepest common ancestor between the previous tip and target entry. */
|
||||
commonAncestorId: string | null;
|
||||
}
|
||||
|
||||
/** Options for generating a branch summary. */
|
||||
export interface GenerateBranchSummaryOptions {
|
||||
/** Provider collection the summarization request goes through; owns auth resolution. */
|
||||
models: Models;
|
||||
/** Model used for summarization. */
|
||||
model: Model<Api>;
|
||||
/** Optional instructions appended to or replacing the default prompt. */
|
||||
customInstructions?: string;
|
||||
/** Replace the default prompt with custom instructions instead of appending them. */
|
||||
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. */
|
||||
export async function collectEntriesForBranchSummary(
|
||||
branch: Pick<Branch, "findEntries">,
|
||||
session: Pick<Session, "getEntry">,
|
||||
oldTipId: string | null,
|
||||
targetId: string,
|
||||
context: Context,
|
||||
): Promise<CollectEntriesResult> {
|
||||
if (!oldTipId) {
|
||||
return { entries: [], commonAncestorId: null };
|
||||
}
|
||||
const oldPath = new Set((await branch.findEntries({ start: oldTipId }, context)).map((entry) => entry.id));
|
||||
const targetPath = await branch.findEntries({ start: targetId }, context);
|
||||
let commonAncestorId: string | null = null;
|
||||
for (const entry of targetPath) {
|
||||
if (oldPath.has(entry.id)) {
|
||||
commonAncestorId = entry.id;
|
||||
break;
|
||||
}
|
||||
}
|
||||
const entries: Entry[] = [];
|
||||
let current: string | null = oldTipId;
|
||||
|
||||
while (current && current !== commonAncestorId) {
|
||||
const entry = await session.getEntry(current, context);
|
||||
if (!entry) throw new Error(`Corrupt session: entry ${current} not found`);
|
||||
entries.push(entry);
|
||||
current = entry.parentId;
|
||||
}
|
||||
entries.reverse();
|
||||
|
||||
return { entries, commonAncestorId };
|
||||
}
|
||||
function getMessageFromEntry(entry: Entry): AgentMessage | undefined {
|
||||
switch (entry.type) {
|
||||
case "message":
|
||||
if (entry.message.role === "toolResult") return undefined;
|
||||
return entry.message;
|
||||
|
||||
case "branch_summary":
|
||||
return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp);
|
||||
|
||||
case "compaction":
|
||||
return createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp);
|
||||
case "custom":
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
|
||||
/** Prepare branch entries for summarization within an optional token budget. */
|
||||
export function prepareBranchEntries(entries: Entry[], tokenBudget: number = 0): BranchPreparation {
|
||||
const messages: AgentMessage[] = [];
|
||||
const fileOps = createFileOps();
|
||||
let totalTokens = 0;
|
||||
for (const entry of entries) {
|
||||
if (
|
||||
entry.type !== "branch_summary" ||
|
||||
typeof entry.details !== "object" ||
|
||||
entry.details === null ||
|
||||
Array.isArray(entry.details)
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
if (Array.isArray(entry.details.readFiles)) {
|
||||
for (const path of entry.details.readFiles) {
|
||||
if (typeof path === "string") fileOps.read.add(path);
|
||||
}
|
||||
}
|
||||
if (Array.isArray(entry.details.modifiedFiles)) {
|
||||
for (const path of entry.details.modifiedFiles) {
|
||||
if (typeof path === "string") fileOps.edited.add(path);
|
||||
}
|
||||
}
|
||||
}
|
||||
for (let i = entries.length - 1; i >= 0; i--) {
|
||||
const entry = entries[i];
|
||||
const message = getMessageFromEntry(entry);
|
||||
if (!message) continue;
|
||||
extractFileOpsFromMessage(message, fileOps);
|
||||
|
||||
const tokens = estimateTokens(message);
|
||||
if (tokenBudget > 0 && totalTokens + tokens > tokenBudget) {
|
||||
if (entry.type === "compaction" || entry.type === "branch_summary") {
|
||||
if (totalTokens < tokenBudget * 0.9) {
|
||||
messages.unshift(message);
|
||||
totalTokens += tokens;
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
messages.unshift(message);
|
||||
totalTokens += tokens;
|
||||
}
|
||||
|
||||
return { messages, fileOps, totalTokens };
|
||||
}
|
||||
|
||||
const BRANCH_SUMMARY_PREAMBLE = `The user explored a different conversation branch before returning here.
|
||||
Summary of that exploration:
|
||||
|
||||
`;
|
||||
|
||||
const BRANCH_SUMMARY_PROMPT = `Create a structured summary of this conversation branch for context when returning later.
|
||||
|
||||
Use this EXACT format:
|
||||
|
||||
## Goal
|
||||
[What was the user trying to accomplish in this branch?]
|
||||
|
||||
## Constraints & Preferences
|
||||
- [Any constraints, preferences, or requirements mentioned]
|
||||
- [Or "(none)" if none were mentioned]
|
||||
|
||||
## Progress
|
||||
### Done
|
||||
- [x] [Completed tasks/changes]
|
||||
|
||||
### In Progress
|
||||
- [ ] [Work that was started but not finished]
|
||||
|
||||
### Blocked
|
||||
- [Issues preventing progress, if any]
|
||||
|
||||
## Key Decisions
|
||||
- **[Decision]**: [Brief rationale]
|
||||
|
||||
## Next Steps
|
||||
1. [What should happen next to continue this work]
|
||||
|
||||
Keep each section concise. Preserve exact file paths, function names, and error messages.`;
|
||||
|
||||
/** Generate a summary for abandoned branch entries. */
|
||||
export function generateBranchSummary(
|
||||
entries: Entry[],
|
||||
options: GenerateBranchSummaryOptions,
|
||||
context: Context,
|
||||
): Promise<Result<BranchSummaryResult, BranchSummaryError>> {
|
||||
const { models, model, customInstructions, replaceInstructions, reserveTokens = 16384, retry, callbacks } = options;
|
||||
const contextWindow = model.contextWindow || 128000;
|
||||
const preparation = prepareBranchEntries(entries, contextWindow - reserveTokens);
|
||||
return generateBranchSummaryWithRequest(
|
||||
preparation,
|
||||
{ customInstructions, replaceInstructions },
|
||||
(aiContext, requestOptions, requestContext) =>
|
||||
completeSimpleWithRetries(models, model, aiContext, requestOptions, retry, callbacks, requestContext),
|
||||
context,
|
||||
);
|
||||
}
|
||||
|
||||
export interface PreparedBranchSummaryOptions {
|
||||
customInstructions?: string;
|
||||
replaceInstructions?: boolean;
|
||||
}
|
||||
|
||||
/** Generate a prepared branch summary through a caller-owned one-request boundary. */
|
||||
export async function generateBranchSummaryWithRequest(
|
||||
preparation: BranchPreparation,
|
||||
options: PreparedBranchSummaryOptions,
|
||||
request: SummaryRequest,
|
||||
context: Context,
|
||||
): Promise<Result<BranchSummaryResult, BranchSummaryError>> {
|
||||
const { customInstructions, replaceInstructions } = options;
|
||||
const { messages, fileOps } = preparation;
|
||||
if (messages.length === 0) {
|
||||
return ok({ summary: "No content to summarize", readFiles: [], modifiedFiles: [] });
|
||||
}
|
||||
const llmMessages = convertToLlm(messages);
|
||||
const conversationText = serializeConversation(llmMessages);
|
||||
let instructions: string;
|
||||
if (replaceInstructions && customInstructions) {
|
||||
instructions = customInstructions;
|
||||
} else if (customInstructions) {
|
||||
instructions = `${BRANCH_SUMMARY_PROMPT}\n\nAdditional focus: ${customInstructions}`;
|
||||
} else {
|
||||
instructions = BRANCH_SUMMARY_PROMPT;
|
||||
}
|
||||
const promptText = `<conversation>\n${conversationText}\n</conversation>\n\n${instructions}`;
|
||||
|
||||
const summarizationMessages = [
|
||||
{
|
||||
role: "user" as const,
|
||||
content: [{ type: "text" as const, text: promptText }],
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
];
|
||||
const response = await request(
|
||||
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
|
||||
createSummaryRequestOptions({ maxTokens: 2048 }, context),
|
||||
context,
|
||||
);
|
||||
if (response.stopReason === "aborted") {
|
||||
return err(new BranchSummaryError("aborted", response.errorMessage || "Branch summary aborted"));
|
||||
}
|
||||
if (response.stopReason === "error") {
|
||||
return err(
|
||||
new BranchSummaryError(
|
||||
"summarization_failed",
|
||||
`Branch summary failed: ${response.errorMessage || "Unknown error"}`,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
let summary = contentText(response.content);
|
||||
summary = BRANCH_SUMMARY_PREAMBLE + summary;
|
||||
const { readFiles, modifiedFiles } = computeFileLists(fileOps);
|
||||
summary += formatFileOperations(readFiles, modifiedFiles);
|
||||
|
||||
return ok({
|
||||
summary: summary || "No summary generated",
|
||||
usage: response.usage,
|
||||
readFiles,
|
||||
modifiedFiles,
|
||||
});
|
||||
}
|
||||
865
third_party/pi/packages/agent/src/harness/compaction/compaction.ts
vendored
Normal file
865
third_party/pi/packages/agent/src/harness/compaction/compaction.ts
vendored
Normal file
@@ -0,0 +1,865 @@
|
||||
import {
|
||||
type Context as AiContext,
|
||||
type Api,
|
||||
type AssistantMessage,
|
||||
contentText,
|
||||
type Model,
|
||||
type Models,
|
||||
type RetryCallbacks,
|
||||
type RetryPolicy,
|
||||
retryAssistantCall,
|
||||
type SimpleStreamOptions,
|
||||
type Usage,
|
||||
uuidv7,
|
||||
} from "@earendil-works/pi-ai";
|
||||
import type { AgentMessage, ThinkingLevel } from "../../types.ts";
|
||||
import { type Context, getTelemetryContext } from "../context.ts";
|
||||
import { convertToLlm, createBranchSummaryMessage, createCompactionSummaryMessage } from "../messages.ts";
|
||||
import { buildContextEntries, sessionEntryToContextMessages } from "../session/context.ts";
|
||||
import type { CompactionEntry, Entry, JsonValue } from "../session/types.ts";
|
||||
import { CompactionError, err, ok, type Result } from "../types.ts";
|
||||
import { addUsage } from "../utils/usage.ts";
|
||||
import {
|
||||
computeFileLists,
|
||||
createFileOps,
|
||||
extractFileOpsFromMessage,
|
||||
type FileOperations,
|
||||
formatFileOperations,
|
||||
serializeConversation,
|
||||
} from "./utils.ts";
|
||||
|
||||
/** File-operation details stored on generated compaction entries. */
|
||||
export interface CompactionDetails extends Record<string, JsonValue> {
|
||||
/** Files read in the compacted history. */
|
||||
readFiles: string[];
|
||||
/** Files modified in the compacted history. */
|
||||
modifiedFiles: string[];
|
||||
}
|
||||
function safeJsonStringify(value: unknown): string {
|
||||
try {
|
||||
return JSON.stringify(value) ?? "undefined";
|
||||
} catch {
|
||||
return "[unserializable]";
|
||||
}
|
||||
}
|
||||
|
||||
function extractFileOperations(
|
||||
messages: AgentMessage[],
|
||||
entries: Entry[],
|
||||
prevCompactionIndex: number,
|
||||
): FileOperations {
|
||||
const fileOps = createFileOps();
|
||||
if (prevCompactionIndex >= 0) {
|
||||
const prevCompaction = entries[prevCompactionIndex] as CompactionEntry;
|
||||
if (
|
||||
typeof prevCompaction.details === "object" &&
|
||||
prevCompaction.details !== null &&
|
||||
!Array.isArray(prevCompaction.details)
|
||||
) {
|
||||
if (Array.isArray(prevCompaction.details.readFiles)) {
|
||||
for (const path of prevCompaction.details.readFiles) {
|
||||
if (typeof path === "string") fileOps.read.add(path);
|
||||
}
|
||||
}
|
||||
if (Array.isArray(prevCompaction.details.modifiedFiles)) {
|
||||
for (const path of prevCompaction.details.modifiedFiles) {
|
||||
if (typeof path === "string") fileOps.edited.add(path);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
for (const msg of messages) {
|
||||
extractFileOpsFromMessage(msg, fileOps);
|
||||
}
|
||||
|
||||
return fileOps;
|
||||
}
|
||||
function getMessageFromEntry(entry: Entry): AgentMessage | undefined {
|
||||
if (entry.type === "message") {
|
||||
return entry.message as AgentMessage;
|
||||
}
|
||||
if (entry.type === "branch_summary") {
|
||||
return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp);
|
||||
}
|
||||
if (entry.type === "compaction") {
|
||||
return createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp);
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function getMessageFromEntryForCompaction(entry: Entry): AgentMessage | undefined {
|
||||
if (entry.type === "compaction") {
|
||||
return undefined;
|
||||
}
|
||||
return getMessageFromEntry(entry);
|
||||
}
|
||||
|
||||
/** Generated compaction data ready to be persisted as a compaction entry. */
|
||||
export interface CompactResult<T = JsonValue> {
|
||||
/** Summary text that replaces compacted history in future context. */
|
||||
summary: string;
|
||||
/** Estimated context tokens before compaction. */
|
||||
tokensBefore: number;
|
||||
/** Usage from the LLM call(s) that generated this summary, if available. */
|
||||
usage?: Usage;
|
||||
/** Retained recent messages stored directly on the compaction entry. */
|
||||
retainedTail: AgentMessage[];
|
||||
/** Optional implementation-specific details stored with the compaction entry. */
|
||||
details?: T;
|
||||
}
|
||||
|
||||
export type SummaryRequest = (
|
||||
aiContext: AiContext,
|
||||
options: SimpleStreamOptions,
|
||||
context: Context,
|
||||
) => Promise<AssistantMessage>;
|
||||
|
||||
export function createSummaryRequestOptions(options: SimpleStreamOptions, context: Context): SimpleStreamOptions {
|
||||
return {
|
||||
...options,
|
||||
signal: context.abortSignal,
|
||||
telemetryContext: getTelemetryContext(context),
|
||||
cacheRetention: "none",
|
||||
sessionId: options.sessionId ?? uuidv7(),
|
||||
};
|
||||
}
|
||||
|
||||
export async function completeSimpleWithRetries(
|
||||
models: Models,
|
||||
model: Model<Api>,
|
||||
aiContext: AiContext,
|
||||
options: SimpleStreamOptions,
|
||||
retry: RetryPolicy | undefined,
|
||||
callbacks: RetryCallbacks | undefined,
|
||||
context: Context,
|
||||
): Promise<AssistantMessage> {
|
||||
// Summaries are standalone requests, so isolate routing and avoid cache writes that cannot be reused.
|
||||
const requestOptions = createSummaryRequestOptions(options, context);
|
||||
return retryAssistantCall(
|
||||
() => models.completeSimple(model, aiContext, requestOptions),
|
||||
retry,
|
||||
requestOptions.signal,
|
||||
callbacks,
|
||||
);
|
||||
}
|
||||
|
||||
/** Compaction thresholds and retention settings. */
|
||||
export interface CompactionSettings {
|
||||
/** Enable automatic compaction decisions. */
|
||||
enabled: boolean;
|
||||
/** Tokens reserved for summary prompt and output. */
|
||||
reserveTokens: number;
|
||||
/** Approximate recent-context tokens to keep after compaction. */
|
||||
keepRecentTokens: number;
|
||||
}
|
||||
|
||||
/** Default compaction settings used by the harness. */
|
||||
export const DEFAULT_COMPACTION_SETTINGS: CompactionSettings = {
|
||||
enabled: true,
|
||||
reserveTokens: 16384,
|
||||
keepRecentTokens: 20000,
|
||||
};
|
||||
|
||||
/** Calculate total context tokens from provider usage. */
|
||||
export function calculateContextTokens(usage: Usage): number {
|
||||
return usage.totalTokens || usage.input + usage.output + usage.cacheRead + usage.cacheWrite;
|
||||
}
|
||||
function getAssistantUsage(msg: AgentMessage): Usage | undefined {
|
||||
if (msg.role === "assistant" && "usage" in msg) {
|
||||
const assistantMsg = msg as AssistantMessage;
|
||||
if (
|
||||
assistantMsg.stopReason !== "aborted" &&
|
||||
assistantMsg.stopReason !== "error" &&
|
||||
assistantMsg.usage &&
|
||||
calculateContextTokens(assistantMsg.usage) > 0
|
||||
) {
|
||||
return assistantMsg.usage;
|
||||
}
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/** Return usage from the last valid assistant message in session entries. */
|
||||
export function getLastAssistantUsage(entries: Entry[]): Usage | undefined {
|
||||
for (let i = entries.length - 1; i >= 0; i--) {
|
||||
const entry = entries[i];
|
||||
if (entry.type === "message") {
|
||||
const usage = getAssistantUsage(entry.message as AgentMessage);
|
||||
if (usage) return usage;
|
||||
}
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/** Estimated context-token usage for a message list. */
|
||||
export interface ContextUsageEstimate {
|
||||
/** Estimated total context tokens. */
|
||||
tokens: number;
|
||||
/** Tokens reported by the most recent assistant usage block. */
|
||||
usageTokens: number;
|
||||
/** Estimated tokens after the most recent assistant usage block. */
|
||||
trailingTokens: number;
|
||||
/** Index of the message that provided usage, or null when none exists. */
|
||||
lastUsageIndex: number | null;
|
||||
}
|
||||
|
||||
function getLastAssistantUsageInfo(messages: AgentMessage[]): { usage: Usage; index: number } | undefined {
|
||||
for (let i = messages.length - 1; i >= 0; i--) {
|
||||
const usage = getAssistantUsage(messages[i]);
|
||||
if (usage) return { usage, index: i };
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/** Estimate context tokens for messages using provider usage when available. */
|
||||
export function estimateContextTokens(messages: AgentMessage[]): ContextUsageEstimate {
|
||||
const usageInfo = getLastAssistantUsageInfo(messages);
|
||||
|
||||
if (!usageInfo) {
|
||||
let estimated = 0;
|
||||
for (const message of messages) {
|
||||
estimated += estimateTokens(message);
|
||||
}
|
||||
return {
|
||||
tokens: estimated,
|
||||
usageTokens: 0,
|
||||
trailingTokens: estimated,
|
||||
lastUsageIndex: null,
|
||||
};
|
||||
}
|
||||
|
||||
const usageTokens = calculateContextTokens(usageInfo.usage);
|
||||
let trailingTokens = 0;
|
||||
for (let i = usageInfo.index + 1; i < messages.length; i++) {
|
||||
trailingTokens += estimateTokens(messages[i]);
|
||||
}
|
||||
|
||||
return {
|
||||
tokens: usageTokens + trailingTokens,
|
||||
usageTokens,
|
||||
trailingTokens,
|
||||
lastUsageIndex: usageInfo.index,
|
||||
};
|
||||
}
|
||||
|
||||
/** Return whether context usage exceeds the configured compaction threshold. */
|
||||
export function shouldCompact(contextTokens: number, contextWindow: number, settings: CompactionSettings): boolean {
|
||||
if (!settings.enabled) return false;
|
||||
return contextTokens > contextWindow - settings.reserveTokens;
|
||||
}
|
||||
|
||||
const ESTIMATED_IMAGE_CHARS = 4800;
|
||||
|
||||
function estimateTextAndImageContentChars(content: string | Array<{ type: string; text?: string }>): number {
|
||||
if (typeof content === "string") {
|
||||
return content.length;
|
||||
}
|
||||
|
||||
let chars = 0;
|
||||
for (const block of content) {
|
||||
if (block.type === "text" && block.text) {
|
||||
chars += block.text.length;
|
||||
} else if (block.type === "image") {
|
||||
chars += ESTIMATED_IMAGE_CHARS;
|
||||
}
|
||||
}
|
||||
return chars;
|
||||
}
|
||||
|
||||
/** Estimate token count for one message using a conservative character heuristic. */
|
||||
export function estimateTokens(message: AgentMessage): number {
|
||||
let chars = 0;
|
||||
|
||||
switch (message.role) {
|
||||
case "user": {
|
||||
chars = estimateTextAndImageContentChars(
|
||||
(message as { content: string | Array<{ type: string; text?: string }> }).content,
|
||||
);
|
||||
return Math.ceil(chars / 4);
|
||||
}
|
||||
case "assistant": {
|
||||
const assistant = message as AssistantMessage;
|
||||
for (const block of assistant.content) {
|
||||
if (block.type === "text") {
|
||||
chars += block.text.length;
|
||||
} else if (block.type === "thinking") {
|
||||
chars += block.thinking.length;
|
||||
} else if (block.type === "toolCall") {
|
||||
chars += block.name.length + safeJsonStringify(block.arguments).length;
|
||||
}
|
||||
}
|
||||
return Math.ceil(chars / 4);
|
||||
}
|
||||
case "custom":
|
||||
case "toolResult": {
|
||||
chars = estimateTextAndImageContentChars(message.content);
|
||||
return Math.ceil(chars / 4);
|
||||
}
|
||||
case "bashExecution": {
|
||||
chars = message.command.length + message.output.length;
|
||||
return Math.ceil(chars / 4);
|
||||
}
|
||||
case "branchSummary":
|
||||
case "compactionSummary": {
|
||||
chars = message.summary.length;
|
||||
return Math.ceil(chars / 4);
|
||||
}
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
function findValidCutPoints(entries: Entry[], startIndex: number, endIndex: number): number[] {
|
||||
const cutPoints: number[] = [];
|
||||
for (let i = startIndex; i < endIndex; i++) {
|
||||
const entry = entries[i];
|
||||
switch (entry.type) {
|
||||
case "message": {
|
||||
const role = entry.message.role;
|
||||
switch (role) {
|
||||
case "bashExecution":
|
||||
case "custom":
|
||||
case "branchSummary":
|
||||
case "compactionSummary":
|
||||
case "user":
|
||||
case "assistant":
|
||||
cutPoints.push(i);
|
||||
break;
|
||||
case "toolResult":
|
||||
break;
|
||||
}
|
||||
break;
|
||||
}
|
||||
case "compaction":
|
||||
case "branch_summary":
|
||||
case "custom":
|
||||
break;
|
||||
}
|
||||
if (entry.type === "branch_summary") cutPoints.push(i);
|
||||
}
|
||||
return cutPoints;
|
||||
}
|
||||
|
||||
/** Find the user-visible message that starts the turn containing an entry. */
|
||||
export function findTurnStartIndex(entries: Entry[], entryIndex: number, startIndex: number): number {
|
||||
for (let i = entryIndex; i >= startIndex; i--) {
|
||||
const entry = entries[i];
|
||||
if (entry.type === "branch_summary") {
|
||||
return i;
|
||||
}
|
||||
if (entry.type === "message") {
|
||||
const role = entry.message.role;
|
||||
if (role === "user" || role === "bashExecution") {
|
||||
return i;
|
||||
}
|
||||
}
|
||||
}
|
||||
return -1;
|
||||
}
|
||||
|
||||
/** Cut point selected for compaction. */
|
||||
export interface CutPointResult {
|
||||
/** Index of the first entry retained after compaction. */
|
||||
firstKeptEntryIndex: number;
|
||||
/** Index of the turn-start entry when the cut splits a turn, otherwise -1. */
|
||||
turnStartIndex: number;
|
||||
/** Whether the selected cut point splits an in-progress turn. */
|
||||
isSplitTurn: boolean;
|
||||
}
|
||||
|
||||
/** Find the compaction cut point that keeps approximately the requested recent-token budget. */
|
||||
export function findCutPoint(
|
||||
entries: Entry[],
|
||||
startIndex: number,
|
||||
endIndex: number,
|
||||
keepRecentTokens: number,
|
||||
): CutPointResult {
|
||||
const cutPoints = findValidCutPoints(entries, startIndex, endIndex);
|
||||
|
||||
if (cutPoints.length === 0) {
|
||||
return { firstKeptEntryIndex: startIndex, turnStartIndex: -1, isSplitTurn: false };
|
||||
}
|
||||
let accumulatedTokens = 0;
|
||||
let cutIndex = cutPoints[0];
|
||||
|
||||
for (let i = endIndex - 1; i >= startIndex; i--) {
|
||||
const entry = entries[i];
|
||||
if (entry.type !== "message") continue;
|
||||
const messageTokens = estimateTokens(entry.message as AgentMessage);
|
||||
accumulatedTokens += messageTokens;
|
||||
if (accumulatedTokens >= keepRecentTokens) {
|
||||
for (let c = 0; c < cutPoints.length; c++) {
|
||||
if (cutPoints[c] >= i) {
|
||||
cutIndex = cutPoints[c];
|
||||
break;
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
while (cutIndex > startIndex) {
|
||||
const prevEntry = entries[cutIndex - 1];
|
||||
if (prevEntry.type === "compaction") {
|
||||
break;
|
||||
}
|
||||
if (prevEntry.type === "message") {
|
||||
break;
|
||||
}
|
||||
cutIndex--;
|
||||
}
|
||||
const cutEntry = entries[cutIndex];
|
||||
const isUserMessage = cutEntry.type === "message" && cutEntry.message.role === "user";
|
||||
const turnStartIndex = isUserMessage ? -1 : findTurnStartIndex(entries, cutIndex, startIndex);
|
||||
|
||||
return {
|
||||
firstKeptEntryIndex: cutIndex,
|
||||
turnStartIndex,
|
||||
isSplitTurn: !isUserMessage && turnStartIndex !== -1,
|
||||
};
|
||||
}
|
||||
|
||||
export const SUMMARIZATION_SYSTEM_PROMPT = `You are a context summarization assistant. Your task is to read a conversation between a user and an AI assistant, then produce a structured summary following the exact format specified.
|
||||
|
||||
Do NOT continue the conversation. Do NOT respond to any questions in the conversation. ONLY output the structured summary.`;
|
||||
|
||||
const SUMMARIZATION_PROMPT = `The messages above are a conversation to summarize. Create a structured context checkpoint summary that another LLM will use to continue the work.
|
||||
|
||||
Use this EXACT format:
|
||||
|
||||
## Goal
|
||||
[What is the user trying to accomplish? Can be multiple items if the session covers different tasks.]
|
||||
|
||||
## Constraints & Preferences
|
||||
- [Any constraints, preferences, or requirements mentioned by user]
|
||||
- [Or "(none)" if none were mentioned]
|
||||
|
||||
## Progress
|
||||
### Done
|
||||
- [x] [Completed tasks/changes]
|
||||
|
||||
### In Progress
|
||||
- [ ] [Current work]
|
||||
|
||||
### Blocked
|
||||
- [Issues preventing progress, if any]
|
||||
|
||||
## Key Decisions
|
||||
- **[Decision]**: [Brief rationale]
|
||||
|
||||
## Next Steps
|
||||
1. [Ordered list of what should happen next]
|
||||
|
||||
## Critical Context
|
||||
- [Any data, examples, or references needed to continue]
|
||||
- [Or "(none)" if not applicable]
|
||||
|
||||
Keep each section concise. Preserve exact file paths, function names, and error messages.`;
|
||||
|
||||
const UPDATE_SUMMARIZATION_PROMPT = `The messages above are NEW conversation messages to incorporate into the existing summary provided in <previous-summary> tags.
|
||||
|
||||
Update the existing structured summary with new information. RULES:
|
||||
- PRESERVE all existing information from the previous summary
|
||||
- ADD new progress, decisions, and context from the new messages
|
||||
- UPDATE the Progress section: move items from "In Progress" to "Done" when completed
|
||||
- UPDATE "Next Steps" based on what was accomplished
|
||||
- PRESERVE exact file paths, function names, and error messages
|
||||
- If something is no longer relevant, you may remove it
|
||||
|
||||
Use this EXACT format:
|
||||
|
||||
## Goal
|
||||
[Preserve existing goals, add new ones if the task expanded]
|
||||
|
||||
## Constraints & Preferences
|
||||
- [Preserve existing, add new ones discovered]
|
||||
|
||||
## Progress
|
||||
### Done
|
||||
- [x] [Include previously done items AND newly completed items]
|
||||
|
||||
### In Progress
|
||||
- [ ] [Current work - update based on progress]
|
||||
|
||||
### Blocked
|
||||
- [Current blockers - remove if resolved]
|
||||
|
||||
## Key Decisions
|
||||
- **[Decision]**: [Brief rationale] (preserve all previous, add new)
|
||||
|
||||
## Next Steps
|
||||
1. [Update based on current state]
|
||||
|
||||
## Critical Context
|
||||
- [Preserve important context, add new if needed]
|
||||
|
||||
Keep each section concise. Preserve exact file paths, function names, and error messages.`;
|
||||
|
||||
/** Generate or update a conversation summary for compaction. */
|
||||
export async function generateSummary(
|
||||
currentMessages: AgentMessage[],
|
||||
models: Models,
|
||||
model: Model<Api>,
|
||||
reserveTokens: number,
|
||||
customInstructions: string | undefined,
|
||||
previousSummary: string | undefined,
|
||||
thinkingLevel: ThinkingLevel | undefined,
|
||||
retry: RetryPolicy | undefined,
|
||||
callbacks: RetryCallbacks | undefined,
|
||||
context: Context,
|
||||
): Promise<Result<string, CompactionError>> {
|
||||
const result = await generateSummaryWithUsage(
|
||||
currentMessages,
|
||||
models,
|
||||
model,
|
||||
reserveTokens,
|
||||
customInstructions,
|
||||
previousSummary,
|
||||
thinkingLevel,
|
||||
retry,
|
||||
callbacks,
|
||||
context,
|
||||
);
|
||||
return result.ok ? ok(result.value.text) : err(result.error);
|
||||
}
|
||||
|
||||
/** Generate or update a conversation summary and return its provider usage. */
|
||||
export function generateSummaryWithUsage(
|
||||
currentMessages: AgentMessage[],
|
||||
models: Models,
|
||||
model: Model<Api>,
|
||||
reserveTokens: number,
|
||||
customInstructions: string | undefined,
|
||||
previousSummary: string | undefined,
|
||||
thinkingLevel: ThinkingLevel | undefined,
|
||||
retry: RetryPolicy | undefined,
|
||||
callbacks: RetryCallbacks | undefined,
|
||||
context: Context,
|
||||
): Promise<Result<{ text: string; usage: Usage }, CompactionError>> {
|
||||
return generateSummaryWithRequest(
|
||||
currentMessages,
|
||||
{ model, reserveTokens, customInstructions, previousSummary, thinkingLevel },
|
||||
(aiContext, options, requestContext) =>
|
||||
completeSimpleWithRetries(models, model, aiContext, options, retry, callbacks, requestContext),
|
||||
context,
|
||||
);
|
||||
}
|
||||
|
||||
export interface SummaryGenerationOptions {
|
||||
model: Model<Api>;
|
||||
reserveTokens: number;
|
||||
customInstructions?: string;
|
||||
previousSummary?: string;
|
||||
thinkingLevel?: ThinkingLevel;
|
||||
}
|
||||
|
||||
/** Generate one summary through a caller-owned one-request boundary. */
|
||||
export async function generateSummaryWithRequest(
|
||||
currentMessages: AgentMessage[],
|
||||
options: SummaryGenerationOptions,
|
||||
request: SummaryRequest,
|
||||
context: Context,
|
||||
): Promise<Result<{ text: string; usage: Usage }, CompactionError>> {
|
||||
const { model, reserveTokens, customInstructions, previousSummary, thinkingLevel } = options;
|
||||
const maxTokens = Math.min(
|
||||
Math.floor(0.8 * reserveTokens),
|
||||
model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY,
|
||||
);
|
||||
let basePrompt = previousSummary ? UPDATE_SUMMARIZATION_PROMPT : SUMMARIZATION_PROMPT;
|
||||
if (customInstructions) {
|
||||
basePrompt = `${basePrompt}\n\nAdditional focus: ${customInstructions}`;
|
||||
}
|
||||
const llmMessages = convertToLlm(currentMessages);
|
||||
const conversationText = serializeConversation(llmMessages);
|
||||
let promptText = `<conversation>\n${conversationText}\n</conversation>\n\n`;
|
||||
if (previousSummary) {
|
||||
promptText += `<previous-summary>\n${previousSummary}\n</previous-summary>\n\n`;
|
||||
}
|
||||
promptText += basePrompt;
|
||||
|
||||
const summarizationMessages = [
|
||||
{
|
||||
role: "user" as const,
|
||||
content: [{ type: "text" as const, text: promptText }],
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
];
|
||||
|
||||
const completionOptions =
|
||||
model.reasoning && thinkingLevel && thinkingLevel !== "off"
|
||||
? { maxTokens, reasoning: thinkingLevel }
|
||||
: { maxTokens };
|
||||
|
||||
const response = await request(
|
||||
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
|
||||
createSummaryRequestOptions(completionOptions, context),
|
||||
context,
|
||||
);
|
||||
if (response.stopReason === "aborted") {
|
||||
return err(new CompactionError("aborted", response.errorMessage || "Summarization aborted"));
|
||||
}
|
||||
if (response.stopReason === "error") {
|
||||
return err(
|
||||
new CompactionError(
|
||||
"summarization_failed",
|
||||
`Summarization failed: ${response.errorMessage || "Unknown error"}`,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
const textContent = contentText(response.content);
|
||||
|
||||
return ok({ text: textContent, usage: response.usage });
|
||||
}
|
||||
|
||||
/** Prepared inputs for a compaction run. */
|
||||
export interface CompactionPreparation {
|
||||
/** Messages summarized into the history summary. */
|
||||
messagesToSummarize: AgentMessage[];
|
||||
/** Prefix messages summarized separately when compaction splits a turn. */
|
||||
turnPrefixMessages: AgentMessage[];
|
||||
/** Recent messages retained after compaction and stored on the compaction entry. */
|
||||
retainedTail: AgentMessage[];
|
||||
/** Whether compaction splits a turn. */
|
||||
isSplitTurn: boolean;
|
||||
/** Estimated context tokens before compaction. */
|
||||
tokensBefore: number;
|
||||
/** Previous compaction summary used for iterative updates. */
|
||||
previousSummary?: string;
|
||||
/** File operations extracted from summarized history. */
|
||||
fileOps: FileOperations;
|
||||
/** Settings used to prepare compaction. */
|
||||
settings: CompactionSettings;
|
||||
}
|
||||
|
||||
/** Prepare session entries for compaction, or return undefined when compaction is not applicable. */
|
||||
export function prepareCompaction(
|
||||
pathEntries: Entry[],
|
||||
settings: CompactionSettings,
|
||||
): Result<CompactionPreparation | undefined, CompactionError> {
|
||||
if (pathEntries.length === 0 || pathEntries[pathEntries.length - 1].type === "compaction") {
|
||||
return ok(undefined);
|
||||
}
|
||||
|
||||
let prevCompactionIndex = -1;
|
||||
for (let i = pathEntries.length - 1; i >= 0; i--) {
|
||||
if (pathEntries[i].type === "compaction") {
|
||||
prevCompactionIndex = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let previousSummary: string | undefined;
|
||||
let compactableEntries = pathEntries;
|
||||
if (prevCompactionIndex >= 0) {
|
||||
const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry;
|
||||
previousSummary = prevCompaction.summary;
|
||||
const virtualRetainedEntries: Entry[] = prevCompaction.retainedTail.map((message, index) => ({
|
||||
type: "message",
|
||||
id: `${prevCompaction.id}:retained:${index}`,
|
||||
parentId: index === 0 ? prevCompaction.id : `${prevCompaction.id}:retained:${index - 1}`,
|
||||
seq: prevCompaction.seq,
|
||||
timestamp: message.timestamp,
|
||||
message,
|
||||
}));
|
||||
compactableEntries = [...virtualRetainedEntries, ...pathEntries.slice(prevCompactionIndex + 1)];
|
||||
}
|
||||
const boundaryEnd = compactableEntries.length;
|
||||
|
||||
const tokensBefore = estimateContextTokens(
|
||||
buildContextEntries(pathEntries).flatMap(sessionEntryToContextMessages),
|
||||
).tokens;
|
||||
|
||||
const cutPoint = findCutPoint(compactableEntries, 0, boundaryEnd, settings.keepRecentTokens);
|
||||
const historyEnd = cutPoint.isSplitTurn ? cutPoint.turnStartIndex : cutPoint.firstKeptEntryIndex;
|
||||
const messagesToSummarize: AgentMessage[] = [];
|
||||
for (let i = 0; i < historyEnd; i++) {
|
||||
const msg = getMessageFromEntryForCompaction(compactableEntries[i]);
|
||||
if (msg) messagesToSummarize.push(msg);
|
||||
}
|
||||
const turnPrefixMessages: AgentMessage[] = [];
|
||||
if (cutPoint.isSplitTurn) {
|
||||
for (let i = cutPoint.turnStartIndex; i < cutPoint.firstKeptEntryIndex; i++) {
|
||||
const msg = getMessageFromEntryForCompaction(compactableEntries[i]);
|
||||
if (msg) turnPrefixMessages.push(msg);
|
||||
}
|
||||
}
|
||||
const retainedTail: AgentMessage[] = [];
|
||||
for (let i = cutPoint.firstKeptEntryIndex; i < boundaryEnd; i++) {
|
||||
const msg = getMessageFromEntryForCompaction(compactableEntries[i]);
|
||||
if (msg) retainedTail.push(msg);
|
||||
}
|
||||
const fileOps = extractFileOperations(messagesToSummarize, pathEntries, prevCompactionIndex);
|
||||
if (cutPoint.isSplitTurn) {
|
||||
for (const msg of turnPrefixMessages) {
|
||||
extractFileOpsFromMessage(msg, fileOps);
|
||||
}
|
||||
}
|
||||
|
||||
return ok({
|
||||
messagesToSummarize,
|
||||
turnPrefixMessages,
|
||||
retainedTail,
|
||||
isSplitTurn: cutPoint.isSplitTurn,
|
||||
tokensBefore,
|
||||
previousSummary,
|
||||
fileOps,
|
||||
settings,
|
||||
});
|
||||
}
|
||||
|
||||
const TURN_PREFIX_SUMMARIZATION_PROMPT = `This is the PREFIX of a turn that was too large to keep. The SUFFIX (recent work) is retained.
|
||||
|
||||
Summarize the prefix to provide context for the retained suffix:
|
||||
|
||||
## Original Request
|
||||
[What did the user ask for in this turn?]
|
||||
|
||||
## Early Progress
|
||||
- [Key decisions and work done in the prefix]
|
||||
|
||||
## Context for Suffix
|
||||
- [Information needed to understand the retained recent work]
|
||||
|
||||
Be concise. Focus on what's needed to understand the kept suffix.`;
|
||||
|
||||
export { serializeConversation } from "./utils.ts";
|
||||
|
||||
/** Generate compaction summary data from prepared session history. */
|
||||
export function compact(
|
||||
preparation: CompactionPreparation,
|
||||
models: Models,
|
||||
model: Model<Api>,
|
||||
customInstructions: string | undefined,
|
||||
thinkingLevel: ThinkingLevel | undefined,
|
||||
retry: RetryPolicy | undefined,
|
||||
callbacks: RetryCallbacks | undefined,
|
||||
context: Context,
|
||||
): Promise<Result<CompactResult, CompactionError>> {
|
||||
return compactWithRequest(
|
||||
preparation,
|
||||
{ model, customInstructions, thinkingLevel },
|
||||
(aiContext, options, requestContext) =>
|
||||
completeSimpleWithRetries(models, model, aiContext, options, retry, callbacks, requestContext),
|
||||
context,
|
||||
);
|
||||
}
|
||||
|
||||
export interface CompactGenerationOptions {
|
||||
model: Model<Api>;
|
||||
customInstructions?: string;
|
||||
thinkingLevel?: ThinkingLevel;
|
||||
}
|
||||
|
||||
/** Generate compaction data through a caller-owned boundary for each provider request. */
|
||||
export async function compactWithRequest(
|
||||
preparation: CompactionPreparation,
|
||||
options: CompactGenerationOptions,
|
||||
request: SummaryRequest,
|
||||
context: Context,
|
||||
): Promise<Result<CompactResult, CompactionError>> {
|
||||
const { model, customInstructions, thinkingLevel } = options;
|
||||
const {
|
||||
messagesToSummarize,
|
||||
turnPrefixMessages,
|
||||
retainedTail,
|
||||
isSplitTurn,
|
||||
tokensBefore,
|
||||
previousSummary,
|
||||
fileOps,
|
||||
settings,
|
||||
} = preparation;
|
||||
|
||||
let summary: string;
|
||||
let summaryUsage: Usage;
|
||||
|
||||
if (isSplitTurn && turnPrefixMessages.length > 0) {
|
||||
let historyText = "No prior history.";
|
||||
let historyUsage: Usage | undefined;
|
||||
if (messagesToSummarize.length > 0) {
|
||||
const historyResult = await generateSummaryWithRequest(
|
||||
messagesToSummarize,
|
||||
{ model, reserveTokens: settings.reserveTokens, customInstructions, previousSummary, thinkingLevel },
|
||||
request,
|
||||
context,
|
||||
);
|
||||
if (!historyResult.ok) return err(historyResult.error);
|
||||
historyText = historyResult.value.text;
|
||||
historyUsage = historyResult.value.usage;
|
||||
}
|
||||
const turnPrefixResult = await generateTurnPrefixSummary(
|
||||
turnPrefixMessages,
|
||||
model,
|
||||
settings.reserveTokens,
|
||||
thinkingLevel,
|
||||
request,
|
||||
context,
|
||||
);
|
||||
if (!turnPrefixResult.ok) return err(turnPrefixResult.error);
|
||||
summary = `${historyText}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.value.text}`;
|
||||
summaryUsage = historyUsage ? addUsage(historyUsage, turnPrefixResult.value.usage) : turnPrefixResult.value.usage;
|
||||
} else {
|
||||
const summaryResult = await generateSummaryWithRequest(
|
||||
messagesToSummarize,
|
||||
{ model, reserveTokens: settings.reserveTokens, customInstructions, previousSummary, thinkingLevel },
|
||||
request,
|
||||
context,
|
||||
);
|
||||
if (!summaryResult.ok) return err(summaryResult.error);
|
||||
summary = summaryResult.value.text;
|
||||
summaryUsage = summaryResult.value.usage;
|
||||
}
|
||||
|
||||
const { readFiles, modifiedFiles } = computeFileLists(fileOps);
|
||||
summary += formatFileOperations(readFiles, modifiedFiles);
|
||||
const details: CompactionDetails = { readFiles, modifiedFiles };
|
||||
|
||||
return ok({ summary, tokensBefore, usage: summaryUsage, retainedTail, details });
|
||||
}
|
||||
async function generateTurnPrefixSummary(
|
||||
messages: AgentMessage[],
|
||||
model: Model<Api>,
|
||||
reserveTokens: number,
|
||||
thinkingLevel: ThinkingLevel | undefined,
|
||||
request: SummaryRequest,
|
||||
context: Context,
|
||||
): Promise<Result<{ text: string; usage: Usage }, CompactionError>> {
|
||||
const maxTokens = Math.min(
|
||||
Math.floor(0.5 * reserveTokens),
|
||||
model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY,
|
||||
);
|
||||
const llmMessages = convertToLlm(messages);
|
||||
const conversationText = serializeConversation(llmMessages);
|
||||
const promptText = `<conversation>\n${conversationText}\n</conversation>\n\n${TURN_PREFIX_SUMMARIZATION_PROMPT}`;
|
||||
const summarizationMessages = [
|
||||
{
|
||||
role: "user" as const,
|
||||
content: [{ type: "text" as const, text: promptText }],
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
];
|
||||
|
||||
const completionOptions =
|
||||
model.reasoning && thinkingLevel && thinkingLevel !== "off"
|
||||
? { maxTokens, reasoning: thinkingLevel }
|
||||
: { maxTokens };
|
||||
const response = await request(
|
||||
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
|
||||
createSummaryRequestOptions(completionOptions, context),
|
||||
context,
|
||||
);
|
||||
if (response.stopReason === "aborted") {
|
||||
return err(new CompactionError("aborted", response.errorMessage || "Turn prefix summarization aborted"));
|
||||
}
|
||||
if (response.stopReason === "error") {
|
||||
return err(
|
||||
new CompactionError(
|
||||
"summarization_failed",
|
||||
`Turn prefix summarization failed: ${response.errorMessage || "Unknown error"}`,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
return ok({
|
||||
text: contentText(response.content),
|
||||
usage: response.usage,
|
||||
});
|
||||
}
|
||||
132
third_party/pi/packages/agent/src/harness/compaction/utils.ts
vendored
Normal file
132
third_party/pi/packages/agent/src/harness/compaction/utils.ts
vendored
Normal file
@@ -0,0 +1,132 @@
|
||||
import { contentText, type Message } from "@earendil-works/pi-ai";
|
||||
import type { AgentMessage } from "../../types.ts";
|
||||
|
||||
/** File paths touched by a session branch or compaction range. */
|
||||
export interface FileOperations {
|
||||
/** Files read but not necessarily modified. */
|
||||
read: Set<string>;
|
||||
/** Files written by full-file write operations. */
|
||||
written: Set<string>;
|
||||
/** Files modified by edit operations. */
|
||||
edited: Set<string>;
|
||||
}
|
||||
|
||||
/** Create an empty file-operation accumulator. */
|
||||
export function createFileOps(): FileOperations {
|
||||
return {
|
||||
read: new Set(),
|
||||
written: new Set(),
|
||||
edited: new Set(),
|
||||
};
|
||||
}
|
||||
|
||||
/** Add file operations from assistant tool calls to an accumulator. */
|
||||
export function extractFileOpsFromMessage(message: AgentMessage, fileOps: FileOperations): void {
|
||||
if (message.role !== "assistant") return;
|
||||
if (!("content" in message) || !Array.isArray(message.content)) return;
|
||||
|
||||
for (const block of message.content) {
|
||||
if (typeof block !== "object" || block === null) continue;
|
||||
if (!("type" in block) || block.type !== "toolCall") continue;
|
||||
if (!("arguments" in block) || !("name" in block)) continue;
|
||||
|
||||
const args = block.arguments as Record<string, unknown> | undefined;
|
||||
if (!args) continue;
|
||||
|
||||
const path = typeof args.path === "string" ? args.path : undefined;
|
||||
if (!path) continue;
|
||||
|
||||
switch (block.name) {
|
||||
case "read":
|
||||
fileOps.read.add(path);
|
||||
break;
|
||||
case "write":
|
||||
fileOps.written.add(path);
|
||||
break;
|
||||
case "edit":
|
||||
fileOps.edited.add(path);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/** Compute sorted read-only and modified file lists from accumulated operations. */
|
||||
export function computeFileLists(fileOps: FileOperations): { readFiles: string[]; modifiedFiles: string[] } {
|
||||
const modified = new Set([...fileOps.edited, ...fileOps.written]);
|
||||
const readOnly = [...fileOps.read].filter((f) => !modified.has(f)).sort();
|
||||
const modifiedFiles = [...modified].sort();
|
||||
return { readFiles: readOnly, modifiedFiles };
|
||||
}
|
||||
|
||||
/** Format file lists as summary metadata tags. */
|
||||
export function formatFileOperations(readFiles: string[], modifiedFiles: string[]): string {
|
||||
const sections: string[] = [];
|
||||
if (readFiles.length > 0) {
|
||||
sections.push(`<read-files>\n${readFiles.join("\n")}\n</read-files>`);
|
||||
}
|
||||
if (modifiedFiles.length > 0) {
|
||||
sections.push(`<modified-files>\n${modifiedFiles.join("\n")}\n</modified-files>`);
|
||||
}
|
||||
if (sections.length === 0) return "";
|
||||
return `\n\n${sections.join("\n\n")}`;
|
||||
}
|
||||
|
||||
const TOOL_RESULT_MAX_CHARS = 2000;
|
||||
|
||||
function safeJsonStringify(value: unknown): string {
|
||||
try {
|
||||
return JSON.stringify(value) ?? "undefined";
|
||||
} catch {
|
||||
return "[unserializable]";
|
||||
}
|
||||
}
|
||||
|
||||
function truncateForSummary(text: string, maxChars: number): string {
|
||||
if (text.length <= maxChars) return text;
|
||||
const truncatedChars = text.length - maxChars;
|
||||
return `${text.slice(0, maxChars)}\n\n[... ${truncatedChars} more characters truncated]`;
|
||||
}
|
||||
|
||||
/** Serialize LLM messages to plain text for summarization prompts. */
|
||||
export function serializeConversation(messages: Message[]): string {
|
||||
const parts: string[] = [];
|
||||
|
||||
for (const msg of messages) {
|
||||
if (msg.role === "user") {
|
||||
const content = contentText(msg.content, "");
|
||||
if (content) parts.push(`[User]: ${content}`);
|
||||
} else if (msg.role === "assistant") {
|
||||
const thinkingParts: string[] = [];
|
||||
const toolCalls: string[] = [];
|
||||
|
||||
for (const block of msg.content) {
|
||||
if (block.type === "thinking") {
|
||||
thinkingParts.push(block.thinking);
|
||||
} else if (block.type === "toolCall") {
|
||||
const args = block.arguments as Record<string, unknown>;
|
||||
const argsStr = Object.entries(args)
|
||||
.map(([k, v]) => `${k}=${safeJsonStringify(v)}`)
|
||||
.join(", ");
|
||||
toolCalls.push(`${block.name}(${argsStr})`);
|
||||
}
|
||||
}
|
||||
|
||||
if (thinkingParts.length > 0) {
|
||||
parts.push(`[Assistant thinking]: ${thinkingParts.join("\n")}`);
|
||||
}
|
||||
if (msg.content.some((block) => block.type === "text")) {
|
||||
parts.push(`[Assistant]: ${contentText(msg.content)}`);
|
||||
}
|
||||
if (toolCalls.length > 0) {
|
||||
parts.push(`[Assistant tool calls]: ${toolCalls.join("; ")}`);
|
||||
}
|
||||
} else if (msg.role === "toolResult") {
|
||||
const content = contentText(msg.content, "");
|
||||
if (content) {
|
||||
parts.push(`[Tool result]: ${truncateForSummary(content, TOOL_RESULT_MAX_CHARS)}`);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return parts.join("\n\n");
|
||||
}
|
||||
35
third_party/pi/packages/agent/src/harness/config.ts
vendored
Normal file
35
third_party/pi/packages/agent/src/harness/config.ts
vendored
Normal file
@@ -0,0 +1,35 @@
|
||||
import type { RetryPolicy } from "@earendil-works/pi-ai";
|
||||
import type { CompactionSettings } from "./compaction/compaction.ts";
|
||||
|
||||
export const DEFAULT_RETRY_POLICY: RetryPolicy = { enabled: true, maxRetries: 3, baseDelayMs: 1_000 };
|
||||
|
||||
export function validateToolNames(tools: readonly { name: string }[]): void {
|
||||
const names = new Set<string>();
|
||||
for (const tool of tools) {
|
||||
if (names.has(tool.name)) throw new TypeError(`Duplicate tool name: ${JSON.stringify(tool.name)}`);
|
||||
names.add(tool.name);
|
||||
}
|
||||
}
|
||||
|
||||
export function validateRetryPolicy(policy: RetryPolicy): void {
|
||||
if (
|
||||
!Number.isSafeInteger(policy.maxRetries) ||
|
||||
policy.maxRetries < 0 ||
|
||||
policy.maxRetries === Number.MAX_SAFE_INTEGER ||
|
||||
!Number.isSafeInteger(policy.baseDelayMs) ||
|
||||
policy.baseDelayMs < 0
|
||||
) {
|
||||
throw new RangeError("Retry policy values must be finite non-negative safe integers");
|
||||
}
|
||||
}
|
||||
|
||||
export function validateCompactionSettings(settings: CompactionSettings): void {
|
||||
if (
|
||||
!Number.isSafeInteger(settings.reserveTokens) ||
|
||||
settings.reserveTokens < 0 ||
|
||||
!Number.isSafeInteger(settings.keepRecentTokens) ||
|
||||
settings.keepRecentTokens < 0
|
||||
) {
|
||||
throw new RangeError("Compaction token counts must be finite non-negative safe integers");
|
||||
}
|
||||
}
|
||||
37
third_party/pi/packages/agent/src/harness/context.ts
vendored
Normal file
37
third_party/pi/packages/agent/src/harness/context.ts
vendored
Normal file
@@ -0,0 +1,37 @@
|
||||
import type { Context, ContextKey } from "@earendil-works/chord";
|
||||
import {
|
||||
awaitWithContext,
|
||||
BACKGROUND_CONTEXT,
|
||||
createContextKey,
|
||||
TODO_CONTEXT,
|
||||
withAbortSignal,
|
||||
withCancel,
|
||||
withContextValue,
|
||||
withoutAbortSignal,
|
||||
} from "@earendil-works/chord/context";
|
||||
import { NOOP_TELEMETRY_CONTEXT, type TelemetryContext } from "@earendil-works/pi-telemetry";
|
||||
|
||||
export {
|
||||
awaitWithContext,
|
||||
BACKGROUND_CONTEXT,
|
||||
type Context,
|
||||
type ContextKey,
|
||||
createContextKey,
|
||||
TODO_CONTEXT,
|
||||
withAbortSignal,
|
||||
withCancel,
|
||||
withContextValue,
|
||||
withoutAbortSignal,
|
||||
};
|
||||
|
||||
const TELEMETRY_CONTEXT_KEY = createContextKey<TelemetryContext>("pi.telemetryContext");
|
||||
|
||||
/** Return the telemetry parent attached to a context, or the shared no-op parent. */
|
||||
export function getTelemetryContext(context: Context): TelemetryContext {
|
||||
return context.value(TELEMETRY_CONTEXT_KEY) ?? NOOP_TELEMETRY_CONTEXT;
|
||||
}
|
||||
|
||||
/** Derive a context whose telemetry children use the supplied parent or active span. */
|
||||
export function withTelemetryContext(telemetryContext: TelemetryContext, context: Context): Context {
|
||||
return withContextValue(TELEMETRY_CONTEXT_KEY, telemetryContext, context);
|
||||
}
|
||||
851
third_party/pi/packages/agent/src/harness/env/nodejs.ts
vendored
Normal file
851
third_party/pi/packages/agent/src/harness/env/nodejs.ts
vendored
Normal file
@@ -0,0 +1,851 @@
|
||||
import { type ChildProcess, spawn } from "node:child_process";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { constants, createReadStream, createWriteStream, type WriteStream } from "node:fs";
|
||||
import {
|
||||
access,
|
||||
appendFile,
|
||||
lstat,
|
||||
mkdir,
|
||||
mkdtemp,
|
||||
readdir,
|
||||
readFile,
|
||||
realpath,
|
||||
rename,
|
||||
rm,
|
||||
writeFile,
|
||||
} from "node:fs/promises";
|
||||
import { homedir, constants as osConstants, tmpdir } from "node:os";
|
||||
import { basename, isAbsolute, join, resolve } from "node:path";
|
||||
import { createInterface } from "node:readline";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import type { Context } from "../context.ts";
|
||||
import {
|
||||
type ExecutionEnv,
|
||||
ExecutionError,
|
||||
err,
|
||||
FileError,
|
||||
type FileInfo,
|
||||
type FileKind,
|
||||
ok,
|
||||
type Result,
|
||||
type ShellExecOptions,
|
||||
type ShellExecResult,
|
||||
toError,
|
||||
} from "../types.ts";
|
||||
import { OutputCapture } from "../utils/output-capture.ts";
|
||||
|
||||
const MAX_TIMEOUT_MS = 2_147_483_647;
|
||||
const MAX_TIMEOUT_SECONDS = MAX_TIMEOUT_MS / 1000;
|
||||
const EXIT_STDIO_GRACE_MS = 100;
|
||||
const SPILL_HIGH_WATER_MARK = 8 * 1024 * 1024;
|
||||
|
||||
type SpillChunk = string | Uint8Array;
|
||||
|
||||
function resolveTimeoutMs(timeout: number | undefined): Result<number | undefined, ExecutionError> {
|
||||
if (timeout === undefined) return ok(undefined);
|
||||
if (!Number.isFinite(timeout) || timeout <= 0) {
|
||||
return err(new ExecutionError("timeout", "Invalid timeout: must be a finite number of seconds"));
|
||||
}
|
||||
|
||||
const timeoutMs = timeout * 1000;
|
||||
if (timeoutMs > MAX_TIMEOUT_MS) {
|
||||
return err(new ExecutionError("timeout", `Invalid timeout: maximum is ${MAX_TIMEOUT_SECONDS} seconds`));
|
||||
}
|
||||
return ok(timeoutMs);
|
||||
}
|
||||
|
||||
function resolvePath(cwd: string, path: string): string {
|
||||
let normalized = path;
|
||||
if (normalized === "~") {
|
||||
normalized = homedir();
|
||||
} else if (normalized.startsWith("~/") || (process.platform === "win32" && normalized.startsWith("~\\"))) {
|
||||
normalized = join(homedir(), normalized.slice(2));
|
||||
} else if (normalized.startsWith("file://")) {
|
||||
try {
|
||||
normalized = fileURLToPath(normalized);
|
||||
} catch {
|
||||
// Keep malformed URLs as ordinary paths so filesystem methods preserve their non-throwing contract.
|
||||
}
|
||||
}
|
||||
return isAbsolute(normalized) ? resolve(normalized) : resolve(cwd, normalized);
|
||||
}
|
||||
|
||||
function fileKindFromStats(stats: {
|
||||
isFile(): boolean;
|
||||
isDirectory(): boolean;
|
||||
isSymbolicLink(): boolean;
|
||||
}): FileKind | undefined {
|
||||
if (stats.isFile()) return "file";
|
||||
if (stats.isDirectory()) return "directory";
|
||||
if (stats.isSymbolicLink()) return "symlink";
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function fileInfoFromStats(
|
||||
path: string,
|
||||
stats: { isFile(): boolean; isDirectory(): boolean; isSymbolicLink(): boolean; size: number; mtimeMs: number },
|
||||
): Result<FileInfo, FileError> {
|
||||
const kind = fileKindFromStats(stats);
|
||||
if (!kind) return err(new FileError("invalid", "Unsupported file type", path));
|
||||
return ok({
|
||||
name: basename(path),
|
||||
path,
|
||||
kind,
|
||||
size: stats.size,
|
||||
mtimeMs: stats.mtimeMs,
|
||||
});
|
||||
}
|
||||
|
||||
function isNodeError(error: unknown): error is NodeJS.ErrnoException {
|
||||
return error instanceof Error && "code" in error;
|
||||
}
|
||||
|
||||
function toFileError(error: unknown, fallbackPath?: string): FileError {
|
||||
if (error instanceof FileError) return error;
|
||||
const cause = toError(error);
|
||||
const nodeError = isNodeError(error) ? error : undefined;
|
||||
const path = typeof nodeError?.path === "string" ? nodeError.path : fallbackPath;
|
||||
if (nodeError) {
|
||||
const message = nodeError.message;
|
||||
switch (nodeError.code) {
|
||||
case "ABORT_ERR":
|
||||
return new FileError("aborted", message, path, cause);
|
||||
case "ENOENT":
|
||||
return new FileError("not_found", message, path, cause);
|
||||
case "EACCES":
|
||||
case "EPERM":
|
||||
return new FileError("permission_denied", message, path, cause);
|
||||
case "ENOTDIR":
|
||||
return new FileError("not_directory", message, path, cause);
|
||||
case "EISDIR":
|
||||
return new FileError("is_directory", message, path, cause);
|
||||
case "EINVAL":
|
||||
return new FileError("invalid", message, path, cause);
|
||||
}
|
||||
}
|
||||
return new FileError("unknown", cause.message, path, cause);
|
||||
}
|
||||
|
||||
function abortResult<TValue>(signal: AbortSignal | undefined, path?: string): Result<TValue, FileError> | undefined {
|
||||
return signal?.aborted ? err(new FileError("aborted", "aborted", path)) : undefined;
|
||||
}
|
||||
|
||||
async function pathExists(path: string): Promise<boolean> {
|
||||
try {
|
||||
await access(path, constants.F_OK);
|
||||
return true;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
async function runCommand(
|
||||
command: string,
|
||||
args: string[],
|
||||
timeoutMs: number,
|
||||
): Promise<{ stdout: string; status: number | null }> {
|
||||
return await new Promise((resolve) => {
|
||||
let stdout = "";
|
||||
let child: ReturnType<typeof spawn>;
|
||||
try {
|
||||
child = spawn(command, args, {
|
||||
stdio: ["ignore", "pipe", "ignore"],
|
||||
windowsHide: true,
|
||||
});
|
||||
} catch {
|
||||
resolve({ stdout: "", status: null });
|
||||
return;
|
||||
}
|
||||
const timeout = setTimeout(() => {
|
||||
if (child.pid) killProcessTree(child.pid);
|
||||
}, timeoutMs);
|
||||
child.stdout?.setEncoding("utf8");
|
||||
child.stdout?.on("data", (chunk: string) => {
|
||||
stdout += chunk;
|
||||
});
|
||||
child.on("error", () => {
|
||||
clearTimeout(timeout);
|
||||
resolve({ stdout: "", status: null });
|
||||
});
|
||||
child.on("close", (status) => {
|
||||
clearTimeout(timeout);
|
||||
resolve({ stdout, status });
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
async function findBashOnPath(): Promise<string | null> {
|
||||
const result =
|
||||
process.platform === "win32"
|
||||
? await runCommand("where", ["bash.exe"], 5000)
|
||||
: await runCommand("which", ["bash"], 5000);
|
||||
if (result.status !== 0 || !result.stdout) return null;
|
||||
const firstMatch = result.stdout.trim().split(/\r?\n/)[0];
|
||||
return firstMatch && (await pathExists(firstMatch)) ? firstMatch : null;
|
||||
}
|
||||
|
||||
interface ShellConfig {
|
||||
shell: string;
|
||||
args: string[];
|
||||
commandTransport?: "argv" | "stdin";
|
||||
}
|
||||
|
||||
function isLegacyWslBashPath(path: string): boolean {
|
||||
const normalized = path.replace(/\//g, "\\").toLowerCase();
|
||||
return /^[a-z]:\\windows\\(?:system32|sysnative)\\bash\.exe$/.test(normalized);
|
||||
}
|
||||
|
||||
function getBashShellConfig(shell: string): ShellConfig {
|
||||
return isLegacyWslBashPath(shell) ? { shell, args: ["-s"], commandTransport: "stdin" } : { shell, args: ["-c"] };
|
||||
}
|
||||
|
||||
async function getShellConfig(customShellPath?: string): Promise<Result<ShellConfig, ExecutionError>> {
|
||||
if (customShellPath) {
|
||||
if (await pathExists(customShellPath)) {
|
||||
return ok(getBashShellConfig(customShellPath));
|
||||
}
|
||||
return err(new ExecutionError("shell_unavailable", `Custom shell path not found: ${customShellPath}`));
|
||||
}
|
||||
if (process.platform === "win32") {
|
||||
const candidates: string[] = [];
|
||||
const programFiles = process.env.ProgramFiles;
|
||||
if (programFiles) candidates.push(`${programFiles}\\Git\\bin\\bash.exe`);
|
||||
const programFilesX86 = process.env["ProgramFiles(x86)"];
|
||||
if (programFilesX86) candidates.push(`${programFilesX86}\\Git\\bin\\bash.exe`);
|
||||
for (const candidate of candidates) {
|
||||
if (await pathExists(candidate)) {
|
||||
return ok(getBashShellConfig(candidate));
|
||||
}
|
||||
}
|
||||
const bashOnPath = await findBashOnPath();
|
||||
if (bashOnPath) {
|
||||
return ok(getBashShellConfig(bashOnPath));
|
||||
}
|
||||
return err(
|
||||
new ExecutionError(
|
||||
"shell_unavailable",
|
||||
`No bash shell found. Options:\n` +
|
||||
` 1. Install Git for Windows: https://git-scm.com/download/win\n` +
|
||||
` 2. Add your bash to PATH (Cygwin, MSYS2, etc.)\n` +
|
||||
" 3. Configure an explicit shellPath\n\n" +
|
||||
`Searched Git Bash in:\n${candidates.map((path) => ` ${path}`).join("\n")}`,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
if (await pathExists("/bin/bash")) {
|
||||
return ok(getBashShellConfig("/bin/bash"));
|
||||
}
|
||||
const bashOnPath = await findBashOnPath();
|
||||
if (bashOnPath) {
|
||||
return ok(getBashShellConfig(bashOnPath));
|
||||
}
|
||||
return ok({ shell: "sh", args: ["-c"] });
|
||||
}
|
||||
|
||||
function getShellEnv(
|
||||
baseEnv?: NodeJS.ProcessEnv,
|
||||
extraEnv?: Record<string, string>,
|
||||
inheritEnv = true,
|
||||
): NodeJS.ProcessEnv {
|
||||
if (!inheritEnv) return { ...extraEnv };
|
||||
return {
|
||||
...process.env,
|
||||
...baseEnv,
|
||||
...extraEnv,
|
||||
};
|
||||
}
|
||||
|
||||
function killProcessTree(pid: number): void {
|
||||
if (process.platform === "win32") {
|
||||
try {
|
||||
const child = spawn(
|
||||
join(process.env.SystemRoot ?? "C:\\Windows", "System32", "taskkill.exe"),
|
||||
["/F", "/T", "/PID", String(pid)],
|
||||
{
|
||||
stdio: "ignore",
|
||||
detached: true,
|
||||
windowsHide: true,
|
||||
},
|
||||
);
|
||||
// A failed spawn emits "error" asynchronously; consume it to avoid crashing Node.
|
||||
child.once("error", () => {});
|
||||
} catch {
|
||||
// Ignore errors.
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
process.kill(-pid, "SIGKILL");
|
||||
} catch {
|
||||
try {
|
||||
process.kill(pid, "SIGKILL");
|
||||
} catch {
|
||||
// Process already dead.
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function waitForChildProcess(
|
||||
child: ChildProcess,
|
||||
spillIsDraining: () => boolean,
|
||||
): Promise<{ code: number | null; signal: NodeJS.Signals | null }> {
|
||||
return new Promise((resolvePromise, reject) => {
|
||||
let settled = false;
|
||||
let exited = false;
|
||||
let exitCode: number | null = null;
|
||||
let exitSignal: NodeJS.Signals | null = null;
|
||||
let postExitTimer: ReturnType<typeof setTimeout> | undefined;
|
||||
let stdoutEnded = child.stdout === null;
|
||||
let stderrEnded = child.stderr === null;
|
||||
|
||||
const cleanup = (): void => {
|
||||
if (postExitTimer) clearTimeout(postExitTimer);
|
||||
child.removeListener("error", onError);
|
||||
child.removeListener("exit", onExit);
|
||||
child.removeListener("close", onClose);
|
||||
child.stdout?.removeListener("end", onStdoutEnd);
|
||||
child.stderr?.removeListener("end", onStderrEnd);
|
||||
child.stdout?.removeListener("data", onData);
|
||||
child.stderr?.removeListener("data", onData);
|
||||
};
|
||||
const finalize = (): void => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
cleanup();
|
||||
child.stdout?.destroy();
|
||||
child.stderr?.destroy();
|
||||
resolvePromise({ code: exitCode, signal: exitSignal });
|
||||
};
|
||||
const maybeFinalizeAfterExit = (): void => {
|
||||
if (exited && stdoutEnded && stderrEnded) finalize();
|
||||
};
|
||||
const armIdleTimer = (): void => {
|
||||
if (postExitTimer) clearTimeout(postExitTimer);
|
||||
postExitTimer = setTimeout(() => {
|
||||
if (spillIsDraining()) armIdleTimer();
|
||||
else finalize();
|
||||
}, EXIT_STDIO_GRACE_MS);
|
||||
};
|
||||
const onData = (): void => {
|
||||
if (exited && !settled) armIdleTimer();
|
||||
};
|
||||
const onStdoutEnd = (): void => {
|
||||
stdoutEnded = true;
|
||||
maybeFinalizeAfterExit();
|
||||
};
|
||||
const onStderrEnd = (): void => {
|
||||
stderrEnded = true;
|
||||
maybeFinalizeAfterExit();
|
||||
};
|
||||
const onError = (error: Error): void => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
cleanup();
|
||||
reject(error);
|
||||
};
|
||||
const onExit = (code: number | null, signal: NodeJS.Signals | null): void => {
|
||||
exited = true;
|
||||
exitCode = code;
|
||||
exitSignal = signal;
|
||||
maybeFinalizeAfterExit();
|
||||
if (!settled) armIdleTimer();
|
||||
};
|
||||
const onClose = (code: number | null, signal: NodeJS.Signals | null): void => {
|
||||
exitCode = code;
|
||||
exitSignal = signal;
|
||||
finalize();
|
||||
};
|
||||
|
||||
child.stdout?.once("end", onStdoutEnd);
|
||||
child.stderr?.once("end", onStderrEnd);
|
||||
child.stdout?.on("data", onData);
|
||||
child.stderr?.on("data", onData);
|
||||
child.once("error", onError);
|
||||
child.once("exit", onExit);
|
||||
child.once("close", onClose);
|
||||
});
|
||||
}
|
||||
|
||||
export class NodeExecutionEnv implements ExecutionEnv {
|
||||
cwd: string;
|
||||
private shellPath?: string;
|
||||
private shellEnv?: NodeJS.ProcessEnv;
|
||||
private activeChildPids = new Set<number>();
|
||||
|
||||
constructor(options: { cwd: string; shellPath?: string; shellEnv?: NodeJS.ProcessEnv }) {
|
||||
this.cwd = options.cwd;
|
||||
this.shellPath = options.shellPath;
|
||||
this.shellEnv = options.shellEnv;
|
||||
}
|
||||
|
||||
async absolutePath(path: string, _context: Context): Promise<Result<string, FileError>> {
|
||||
return ok(resolvePath(this.cwd, path));
|
||||
}
|
||||
|
||||
async joinPath(parts: string[], _context: Context): Promise<Result<string, FileError>> {
|
||||
return ok(join(...parts));
|
||||
}
|
||||
|
||||
async exec(
|
||||
command: string,
|
||||
options: ShellExecOptions | undefined,
|
||||
context: Context,
|
||||
): Promise<Result<ShellExecResult, ExecutionError>> {
|
||||
const signal = context.abortSignal;
|
||||
if (signal?.aborted) return err(new ExecutionError("aborted", "aborted"));
|
||||
const timeoutMsResult = resolveTimeoutMs(options?.timeout);
|
||||
if (!timeoutMsResult.ok) return err(timeoutMsResult.error);
|
||||
const timeoutMs = timeoutMsResult.value;
|
||||
|
||||
const cwd = options?.cwd ? resolvePath(this.cwd, options.cwd) : this.cwd;
|
||||
const shellConfig = await getShellConfig(this.shellPath);
|
||||
if (!shellConfig.ok) return shellConfig;
|
||||
try {
|
||||
await access(cwd, constants.F_OK);
|
||||
} catch (error) {
|
||||
const cause = toError(error);
|
||||
return err(
|
||||
new ExecutionError(
|
||||
"spawn_error",
|
||||
`Working directory does not exist: ${cwd}\nCannot execute bash commands.`,
|
||||
cause,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
return await new Promise((resolvePromise) => {
|
||||
let settled = false;
|
||||
let timedOut = false;
|
||||
let callbackError: ExecutionError | undefined;
|
||||
let spillError: ExecutionError | undefined;
|
||||
let child: ReturnType<typeof spawn> | undefined;
|
||||
let timeoutId: ReturnType<typeof setTimeout> | undefined;
|
||||
const spillPrefix: SpillChunk[] = [];
|
||||
let spillPath: string | undefined;
|
||||
const spillQueue: SpillChunk[] = [];
|
||||
let spillStart: Promise<void> | undefined;
|
||||
let spillStream: WriteStream | undefined;
|
||||
let spillBackpressured = false;
|
||||
|
||||
const onAbort = () => {
|
||||
if (child?.pid) killProcessTree(child.pid);
|
||||
};
|
||||
const failCallback = (error: unknown) => {
|
||||
if (callbackError !== undefined) return;
|
||||
const cause = toError(error);
|
||||
callbackError = new ExecutionError("callback_error", cause.message, cause);
|
||||
onAbort();
|
||||
};
|
||||
let capture: OutputCapture;
|
||||
try {
|
||||
capture = new OutputCapture(options?.capture, context, {
|
||||
onUpdate: options?.onUpdate,
|
||||
onError: failCallback,
|
||||
});
|
||||
} catch (error) {
|
||||
const cause = toError(error);
|
||||
resolvePromise(err(new ExecutionError("unknown", cause.message, cause)));
|
||||
return;
|
||||
}
|
||||
|
||||
const settle = (result: Result<ShellExecResult, ExecutionError>) => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
if (timeoutId) clearTimeout(timeoutId);
|
||||
if (signal) signal.removeEventListener("abort", onAbort);
|
||||
if (child?.pid) this.activeChildPids.delete(child.pid);
|
||||
capture.dispose();
|
||||
resolvePromise(result);
|
||||
};
|
||||
const pauseOutput = () => {
|
||||
child?.stdout?.pause();
|
||||
child?.stderr?.pause();
|
||||
};
|
||||
const resumeOutput = () => {
|
||||
if (callbackError || spillError || timedOut || signal?.aborted || spillBackpressured) return;
|
||||
child?.stdout?.resume();
|
||||
child?.stderr?.resume();
|
||||
};
|
||||
const failSpill = (error: unknown) => {
|
||||
if (spillError !== undefined) return;
|
||||
const cause = toError(error);
|
||||
spillError = new ExecutionError(
|
||||
"unknown",
|
||||
`Failed to preserve complete shell output: ${cause.message}`,
|
||||
cause,
|
||||
);
|
||||
spillBackpressured = false;
|
||||
onAbort();
|
||||
};
|
||||
const writeSpill = (chunk: SpillChunk): void => {
|
||||
if (spillStream === undefined || chunk.length === 0) return;
|
||||
if (spillStream.write(chunk) || spillBackpressured) return;
|
||||
spillBackpressured = true;
|
||||
pauseOutput();
|
||||
spillStream.once("drain", () => {
|
||||
spillBackpressured = false;
|
||||
resumeOutput();
|
||||
});
|
||||
};
|
||||
const startSpill = (chunk: SpillChunk): void => {
|
||||
if (spillStream !== undefined) {
|
||||
writeSpill(chunk);
|
||||
return;
|
||||
}
|
||||
spillQueue.push(chunk);
|
||||
if (spillStart !== undefined) return;
|
||||
pauseOutput();
|
||||
spillStart = (async () => {
|
||||
const created = await this.createTempFile({ prefix: "pi-output-", suffix: ".log" }, context);
|
||||
if (!created.ok) throw created.error;
|
||||
spillPath = created.value;
|
||||
capture.setSpillPath(spillPath);
|
||||
spillStream = createWriteStream(spillPath, { flags: "a", highWaterMark: SPILL_HIGH_WATER_MARK });
|
||||
spillStream.on("error", failSpill);
|
||||
for (const queued of spillQueue) writeSpill(queued);
|
||||
spillQueue.length = 0;
|
||||
})()
|
||||
.catch(failSpill)
|
||||
.finally(resumeOutput);
|
||||
};
|
||||
const finishSpill = async (): Promise<void> => {
|
||||
await spillStart;
|
||||
const stream = spillStream;
|
||||
if (stream === undefined || spillError !== undefined || stream.destroyed) return;
|
||||
await new Promise<void>((resolveFinish) => {
|
||||
stream.once("error", () => resolveFinish());
|
||||
stream.once("finish", resolveFinish);
|
||||
stream.end();
|
||||
});
|
||||
};
|
||||
|
||||
try {
|
||||
const commandFromStdin = shellConfig.value.commandTransport === "stdin";
|
||||
child = spawn(
|
||||
shellConfig.value.shell,
|
||||
commandFromStdin ? shellConfig.value.args : [...shellConfig.value.args, command],
|
||||
{
|
||||
cwd,
|
||||
detached: process.platform !== "win32",
|
||||
env: getShellEnv(this.shellEnv, options?.env, options?.inheritEnv),
|
||||
stdio: [commandFromStdin ? "pipe" : "ignore", "pipe", "pipe"],
|
||||
windowsHide: true,
|
||||
},
|
||||
);
|
||||
if (child.pid) this.activeChildPids.add(child.pid);
|
||||
if (commandFromStdin) {
|
||||
child.stdin?.on("error", () => {});
|
||||
child.stdin?.end(command);
|
||||
}
|
||||
} catch (error) {
|
||||
const cause = toError(error);
|
||||
settle(err(new ExecutionError("spawn_error", cause.message, cause)));
|
||||
return;
|
||||
}
|
||||
|
||||
timeoutId =
|
||||
timeoutMs === undefined
|
||||
? undefined
|
||||
: setTimeout(() => {
|
||||
timedOut = true;
|
||||
onAbort();
|
||||
}, timeoutMs);
|
||||
|
||||
if (signal) {
|
||||
if (signal.aborted) onAbort();
|
||||
else signal.addEventListener("abort", onAbort, { once: true });
|
||||
}
|
||||
|
||||
const feed = (chunk: Uint8Array) => {
|
||||
try {
|
||||
const wasTruncated = capture.truncated;
|
||||
capture.push(chunk);
|
||||
if (!options?.capture?.spill || chunk.length === 0) return;
|
||||
if (spillPath !== undefined || wasTruncated) {
|
||||
startSpill(chunk);
|
||||
} else if (capture.truncated) {
|
||||
for (const prefix of spillPrefix) startSpill(prefix);
|
||||
spillPrefix.length = 0;
|
||||
startSpill(chunk);
|
||||
} else {
|
||||
spillPrefix.push(chunk);
|
||||
}
|
||||
} catch (error) {
|
||||
failCallback(error);
|
||||
}
|
||||
};
|
||||
child.stdout?.on("data", feed);
|
||||
child.stderr?.on("data", feed);
|
||||
|
||||
void waitForChildProcess(
|
||||
child,
|
||||
() =>
|
||||
spillError === undefined &&
|
||||
spillStart !== undefined &&
|
||||
(spillStream === undefined || spillBackpressured),
|
||||
).then(
|
||||
async ({ code, signal: exitSignal }) => {
|
||||
await finishSpill();
|
||||
try {
|
||||
capture.finish();
|
||||
capture.flush();
|
||||
} catch (error) {
|
||||
failCallback(error);
|
||||
}
|
||||
if (callbackError) {
|
||||
settle(err(callbackError));
|
||||
return;
|
||||
}
|
||||
if (timedOut) {
|
||||
settle(err(new ExecutionError("timeout", `timeout:${options?.timeout}`)));
|
||||
return;
|
||||
}
|
||||
if (signal?.aborted) {
|
||||
settle(err(new ExecutionError("aborted", "aborted")));
|
||||
return;
|
||||
}
|
||||
if (spillError) {
|
||||
settle(err(spillError));
|
||||
return;
|
||||
}
|
||||
const output = capture.snapshot();
|
||||
// A process killed by a signal (e.g. OOM killer) has no exit code; map it
|
||||
// to the conventional 128 + signal number so callers do not mistake it
|
||||
// for a successful exit.
|
||||
const exitCode = code ?? (exitSignal ? 128 + (osConstants.signals[exitSignal] ?? 0) : 1);
|
||||
settle(
|
||||
ok({
|
||||
exitCode,
|
||||
truncation: output.truncation,
|
||||
...(output.spillPath === undefined ? {} : { spillPath: output.spillPath }),
|
||||
...(output.lastLineBytes === undefined ? {} : { lastLineBytes: output.lastLineBytes }),
|
||||
}),
|
||||
);
|
||||
},
|
||||
(error: Error) => settle(err(new ExecutionError("spawn_error", error.message, error))),
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
async readTextFile(path: string, context: Context): Promise<Result<string, FileError>> {
|
||||
const resolved = resolvePath(this.cwd, path);
|
||||
const signal = context.abortSignal;
|
||||
const aborted = abortResult<string>(signal, resolved);
|
||||
if (aborted) return aborted;
|
||||
try {
|
||||
return ok(await readFile(resolved, { encoding: "utf8", signal }));
|
||||
} catch (error) {
|
||||
return err(toFileError(error, resolved));
|
||||
}
|
||||
}
|
||||
|
||||
async readTextLines(
|
||||
path: string,
|
||||
options: { maxLines?: number } | undefined,
|
||||
context: Context,
|
||||
): Promise<Result<string[], FileError>> {
|
||||
const resolved = resolvePath(this.cwd, path);
|
||||
const signal = context.abortSignal;
|
||||
const aborted = abortResult<string[]>(signal, resolved);
|
||||
if (aborted) return aborted;
|
||||
if (options?.maxLines !== undefined && options.maxLines <= 0) return ok([]);
|
||||
let stream: ReturnType<typeof createReadStream> | undefined;
|
||||
let lineReader: ReturnType<typeof createInterface> | undefined;
|
||||
try {
|
||||
stream = createReadStream(resolved, { encoding: "utf8", signal });
|
||||
lineReader = createInterface({ input: stream, crlfDelay: Infinity });
|
||||
const lines: string[] = [];
|
||||
for await (const line of lineReader) {
|
||||
const loopAbort = abortResult<string[]>(signal, resolved);
|
||||
if (loopAbort) return loopAbort;
|
||||
lines.push(line);
|
||||
if (options?.maxLines !== undefined && lines.length >= options.maxLines) break;
|
||||
}
|
||||
const afterReadAbort = abortResult<string[]>(signal, resolved);
|
||||
if (afterReadAbort) return afterReadAbort;
|
||||
return ok(lines);
|
||||
} catch (error) {
|
||||
return err(toFileError(error, resolved));
|
||||
} finally {
|
||||
lineReader?.close();
|
||||
stream?.destroy();
|
||||
}
|
||||
}
|
||||
|
||||
async readBinaryFile(path: string, context: Context): Promise<Result<Uint8Array, FileError>> {
|
||||
const resolved = resolvePath(this.cwd, path);
|
||||
const signal = context.abortSignal;
|
||||
const aborted = abortResult<Uint8Array>(signal, resolved);
|
||||
if (aborted) return aborted;
|
||||
try {
|
||||
return ok(await readFile(resolved, { signal }));
|
||||
} catch (error) {
|
||||
return err(toFileError(error, resolved));
|
||||
}
|
||||
}
|
||||
|
||||
async writeFile(path: string, content: string | Uint8Array, context: Context): Promise<Result<void, FileError>> {
|
||||
const resolved = resolvePath(this.cwd, path);
|
||||
const signal = context.abortSignal;
|
||||
const aborted = abortResult<void>(signal, resolved);
|
||||
if (aborted) return aborted;
|
||||
try {
|
||||
await mkdir(resolve(resolved, ".."), { recursive: true });
|
||||
const afterMkdirAbort = abortResult<void>(signal, resolved);
|
||||
if (afterMkdirAbort) return afterMkdirAbort;
|
||||
await writeFile(resolved, content, { signal });
|
||||
return ok(undefined);
|
||||
} catch (error) {
|
||||
return err(toFileError(error, resolved));
|
||||
}
|
||||
}
|
||||
|
||||
async appendFile(path: string, content: string | Uint8Array, context: Context): Promise<Result<void, FileError>> {
|
||||
const resolved = resolvePath(this.cwd, path);
|
||||
const signal = context.abortSignal;
|
||||
const aborted = abortResult<void>(signal, resolved);
|
||||
if (aborted) return aborted;
|
||||
try {
|
||||
await mkdir(resolve(resolved, ".."), { recursive: true });
|
||||
const afterMkdirAbort = abortResult<void>(signal, resolved);
|
||||
if (afterMkdirAbort) return afterMkdirAbort;
|
||||
await appendFile(resolved, content);
|
||||
const afterAppendAbort = abortResult<void>(signal, resolved);
|
||||
return afterAppendAbort ?? ok(undefined);
|
||||
} catch (error) {
|
||||
return err(toFileError(error, resolved));
|
||||
}
|
||||
}
|
||||
|
||||
async renameFile(sourcePath: string, destinationPath: string, context: Context): Promise<Result<void, FileError>> {
|
||||
const source = resolvePath(this.cwd, sourcePath);
|
||||
const destination = resolvePath(this.cwd, destinationPath);
|
||||
const aborted = abortResult<void>(context.abortSignal, destination);
|
||||
if (aborted) return aborted;
|
||||
try {
|
||||
await rename(source, destination);
|
||||
return ok(undefined);
|
||||
} catch (error) {
|
||||
return err(toFileError(error, source));
|
||||
}
|
||||
}
|
||||
|
||||
async fileInfo(path: string, context: Context): Promise<Result<FileInfo, FileError>> {
|
||||
const resolved = resolvePath(this.cwd, path);
|
||||
const aborted = abortResult<FileInfo>(context.abortSignal, resolved);
|
||||
if (aborted) return aborted;
|
||||
try {
|
||||
return fileInfoFromStats(resolved, await lstat(resolved));
|
||||
} catch (error) {
|
||||
return err(toFileError(error, resolved));
|
||||
}
|
||||
}
|
||||
|
||||
async listDir(path: string, context: Context): Promise<Result<FileInfo[], FileError>> {
|
||||
const resolved = resolvePath(this.cwd, path);
|
||||
const signal = context.abortSignal;
|
||||
const aborted = abortResult<FileInfo[]>(signal, resolved);
|
||||
if (aborted) return aborted;
|
||||
try {
|
||||
const entries = await readdir(resolved, { withFileTypes: true });
|
||||
const infos: FileInfo[] = [];
|
||||
for (const entry of entries) {
|
||||
const loopAbort = abortResult<FileInfo[]>(signal, resolved);
|
||||
if (loopAbort) return loopAbort;
|
||||
const entryPath = resolve(resolved, entry.name);
|
||||
try {
|
||||
const info = fileInfoFromStats(entryPath, await lstat(entryPath));
|
||||
if (info.ok) infos.push(info.value);
|
||||
} catch (error) {
|
||||
return err(toFileError(error, entryPath));
|
||||
}
|
||||
}
|
||||
return ok(infos);
|
||||
} catch (error) {
|
||||
return err(toFileError(error, resolved));
|
||||
}
|
||||
}
|
||||
|
||||
async canonicalPath(path: string, context: Context): Promise<Result<string, FileError>> {
|
||||
const resolved = resolvePath(this.cwd, path);
|
||||
const aborted = abortResult<string>(context.abortSignal, resolved);
|
||||
if (aborted) return aborted;
|
||||
try {
|
||||
return ok(await realpath(resolved));
|
||||
} catch (error) {
|
||||
return err(toFileError(error, resolved));
|
||||
}
|
||||
}
|
||||
|
||||
async exists(path: string, context: Context): Promise<Result<boolean, FileError>> {
|
||||
const result = await this.fileInfo(path, context);
|
||||
if (result.ok) return ok(true);
|
||||
if (result.error.code === "not_found") return ok(false);
|
||||
return err(result.error);
|
||||
}
|
||||
|
||||
async createDir(
|
||||
path: string,
|
||||
options: { recursive?: boolean } | undefined,
|
||||
context: Context,
|
||||
): Promise<Result<void, FileError>> {
|
||||
const resolved = resolvePath(this.cwd, path);
|
||||
const aborted = abortResult<void>(context.abortSignal, resolved);
|
||||
if (aborted) return aborted;
|
||||
try {
|
||||
await mkdir(resolved, { recursive: options?.recursive ?? true });
|
||||
return ok(undefined);
|
||||
} catch (error) {
|
||||
return err(toFileError(error, resolved));
|
||||
}
|
||||
}
|
||||
|
||||
async remove(
|
||||
path: string,
|
||||
options: { recursive?: boolean; force?: boolean } | undefined,
|
||||
context: Context,
|
||||
): Promise<Result<void, FileError>> {
|
||||
const resolved = resolvePath(this.cwd, path);
|
||||
const aborted = abortResult<void>(context.abortSignal, resolved);
|
||||
if (aborted) return aborted;
|
||||
try {
|
||||
await rm(resolved, { recursive: options?.recursive ?? false, force: options?.force ?? false });
|
||||
return ok(undefined);
|
||||
} catch (error) {
|
||||
return err(toFileError(error, resolved));
|
||||
}
|
||||
}
|
||||
|
||||
async createTempDir(prefix: string | undefined, context: Context): Promise<Result<string, FileError>> {
|
||||
const aborted = abortResult<string>(context.abortSignal);
|
||||
if (aborted) return aborted;
|
||||
try {
|
||||
prefix ??= "tmp-";
|
||||
return ok(await mkdtemp(join(tmpdir(), prefix)));
|
||||
} catch (error) {
|
||||
return err(toFileError(error));
|
||||
}
|
||||
}
|
||||
|
||||
async createTempFile(
|
||||
options: { prefix?: string; suffix?: string } | undefined,
|
||||
context: Context,
|
||||
): Promise<Result<string, FileError>> {
|
||||
const dir = await this.createTempDir("tmp-", context);
|
||||
if (!dir.ok) return dir;
|
||||
const filePath = join(dir.value, `${options?.prefix ?? ""}${randomUUID()}${options?.suffix ?? ""}`);
|
||||
try {
|
||||
await writeFile(filePath, "");
|
||||
return ok(filePath);
|
||||
} catch (error) {
|
||||
return err(toFileError(error, filePath));
|
||||
}
|
||||
}
|
||||
|
||||
async cleanup(_context: Context): Promise<void> {
|
||||
for (const pid of this.activeChildPids) killProcessTree(pid);
|
||||
this.activeChildPids.clear();
|
||||
}
|
||||
}
|
||||
286
third_party/pi/packages/agent/src/harness/events.ts
vendored
Normal file
286
third_party/pi/packages/agent/src/harness/events.ts
vendored
Normal file
@@ -0,0 +1,286 @@
|
||||
import type { EventListener, Events, HarnessEvent, HarnessEventType, WatchHandle } from "./agent-harness.ts";
|
||||
import type { Context } from "./context.ts";
|
||||
|
||||
type UntypedEventListener = (event: HarnessEvent, context: Context) => void | Promise<void>;
|
||||
type ResnapshotCapture<T> = (context: Context, markBoundary: () => void) => Promise<T>;
|
||||
|
||||
/** Passive harness event bus with isolated handler failures. */
|
||||
export class HarnessEventBus implements Events {
|
||||
private readonly listeners = new Map<HarnessEventType, Set<UntypedEventListener>>();
|
||||
private readonly watchListeners = new Set<UntypedEventListener>();
|
||||
private deliveryTail: Promise<void> = Promise.resolve();
|
||||
private closedError: Error | undefined;
|
||||
|
||||
on<TType extends HarnessEventType>(
|
||||
type: TType,
|
||||
listener: EventListener<Extract<HarnessEvent, { type: TType }>>,
|
||||
): () => void {
|
||||
if (this.closedError !== undefined) throw this.closedError;
|
||||
const wrapped: UntypedEventListener = (event, context) =>
|
||||
listener(event as Extract<HarnessEvent, { type: TType }>, context);
|
||||
let listeners = this.listeners.get(type);
|
||||
if (listeners === undefined) {
|
||||
listeners = new Set();
|
||||
this.listeners.set(type, listeners);
|
||||
}
|
||||
listeners.add(wrapped);
|
||||
return () => listeners?.delete(wrapped);
|
||||
}
|
||||
|
||||
emit(event: HarnessEvent, context: Context): Promise<void> {
|
||||
return this.emitBatch([event], context);
|
||||
}
|
||||
|
||||
/** Bind current recipients and append one contiguous batch to the global delivery tail. */
|
||||
emitBatch(events: readonly HarnessEvent[], context: Context): Promise<void> {
|
||||
if (this.closedError !== undefined || events.length === 0) return Promise.resolve();
|
||||
const bound = events.map((event) => {
|
||||
const payload = structuredClone(event);
|
||||
return { payload, recipients: this.snapshotRecipients(payload) };
|
||||
});
|
||||
const delivery = this.deliveryTail.then(async () => {
|
||||
for (const { payload, recipients } of bound) await this.deliver(payload, recipients, true, context);
|
||||
});
|
||||
this.deliveryTail = delivery.catch(() => {});
|
||||
return delivery;
|
||||
}
|
||||
|
||||
watch<T>(
|
||||
snapshot: T,
|
||||
filter: (event: HarnessEvent) => boolean,
|
||||
_context: Context,
|
||||
resnapshot?: ResnapshotCapture<T>,
|
||||
): WatchHandle<T> {
|
||||
if (this.closedError !== undefined) throw this.closedError;
|
||||
return this.installWatcher(snapshot, filter, resnapshot);
|
||||
}
|
||||
|
||||
async watchFromSnapshot<T>(
|
||||
capture: (context: Context) => Promise<T>,
|
||||
filter: (event: HarnessEvent) => boolean,
|
||||
context: Context,
|
||||
): Promise<WatchHandle<T>> {
|
||||
if (this.closedError !== undefined) throw this.closedError;
|
||||
const watcher = this.installWatcher<T>(undefined, filter, async (captureContext, markBoundary) => {
|
||||
const snapshot = await capture(captureContext);
|
||||
markBoundary();
|
||||
return snapshot;
|
||||
});
|
||||
try {
|
||||
watcher.setSnapshot(await capture(context));
|
||||
return watcher;
|
||||
} catch (error) {
|
||||
watcher.unsubscribe();
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
close(error: Error): void {
|
||||
this.closedError ??= error;
|
||||
void this.deliveryTail.finally(() => {
|
||||
this.listeners.clear();
|
||||
this.watchListeners.clear();
|
||||
});
|
||||
}
|
||||
|
||||
private installWatcher<T>(
|
||||
snapshot: T | undefined,
|
||||
filter: (event: HarnessEvent) => boolean,
|
||||
resnapshot: ResnapshotCapture<T> | undefined,
|
||||
): BufferedEventWatcher<T> {
|
||||
let watcher!: BufferedEventWatcher<T>;
|
||||
const capture =
|
||||
resnapshot === undefined
|
||||
? undefined
|
||||
: async (context: Context) => {
|
||||
let marked = false;
|
||||
const next = await resnapshot(context, () => {
|
||||
if (marked) throw new Error("Resnapshot boundary was already marked");
|
||||
marked = true;
|
||||
this.enqueueBarrier(() => watcher.markResnapshotBoundary());
|
||||
});
|
||||
if (!marked) throw new Error("Resnapshot capture did not mark its boundary");
|
||||
return next;
|
||||
};
|
||||
watcher = new BufferedEventWatcher(snapshot, capture, async (error, event, context) => {
|
||||
if (event.type === "handler_error") return;
|
||||
const normalized = error instanceof Error ? error : new Error(String(error));
|
||||
const lane = "lane" in event && typeof event.lane === "string" ? event.lane : undefined;
|
||||
await this.emit(
|
||||
{
|
||||
type: "handler_error",
|
||||
kind: "event",
|
||||
event: event.type,
|
||||
error: normalized.message,
|
||||
...(normalized.stack === undefined ? {} : { stack: normalized.stack }),
|
||||
...(lane === undefined ? {} : { lane }),
|
||||
},
|
||||
context,
|
||||
);
|
||||
});
|
||||
const watchListener: UntypedEventListener = (event, context) => {
|
||||
if (filter(event)) watcher.push(event, context);
|
||||
};
|
||||
this.watchListeners.add(watchListener);
|
||||
watcher.setUnsubscribe(() => this.watchListeners.delete(watchListener));
|
||||
return watcher;
|
||||
}
|
||||
|
||||
private enqueueBarrier(barrier: () => void): void {
|
||||
const delivery = this.deliveryTail.then(barrier);
|
||||
this.deliveryTail = delivery.catch(() => {});
|
||||
}
|
||||
|
||||
private snapshotRecipients(event: HarnessEvent): UntypedEventListener[] {
|
||||
return [...(this.listeners.get(event.type) ?? []), ...this.watchListeners];
|
||||
}
|
||||
|
||||
private async deliver(
|
||||
event: HarnessEvent,
|
||||
recipients: readonly UntypedEventListener[],
|
||||
reportErrors: boolean,
|
||||
context: Context,
|
||||
): Promise<void> {
|
||||
for (const listener of recipients) {
|
||||
try {
|
||||
await listener(structuredClone(event), context);
|
||||
} catch (error) {
|
||||
if (!reportErrors || event.type === "handler_error") continue;
|
||||
const normalized = error instanceof Error ? error : new Error(String(error));
|
||||
const lane = "lane" in event && typeof event.lane === "string" ? event.lane : undefined;
|
||||
const handlerError: HarnessEvent = {
|
||||
type: "handler_error",
|
||||
kind: "event",
|
||||
event: event.type,
|
||||
error: normalized.message,
|
||||
...(normalized.stack === undefined ? {} : { stack: normalized.stack }),
|
||||
...(lane === undefined ? {} : { lane }),
|
||||
};
|
||||
await this.deliver(handlerError, this.snapshotRecipients(handlerError), false, context);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
class BufferedEventWatcher<T> implements WatchHandle<T> {
|
||||
snapshot: T;
|
||||
private readonly resnapshotCallback: ((context: Context) => Promise<T>) | undefined;
|
||||
private readonly onError: (error: unknown, event: HarnessEvent, context: Context) => void | Promise<void>;
|
||||
private buffer: Array<{ event: HarnessEvent; context: Context; epoch: number }> = [];
|
||||
private listener: EventListener | undefined;
|
||||
private unsubscribeCallback: (() => void) | undefined;
|
||||
private deliveryTail: Promise<void> = Promise.resolve();
|
||||
private epoch = 0;
|
||||
private resnapshotState:
|
||||
| {
|
||||
phase: "dropping" | "holding";
|
||||
held: Array<{ event: HarnessEvent; context: Context }>;
|
||||
reached: Promise<void>;
|
||||
resolveReached: () => void;
|
||||
}
|
||||
| undefined;
|
||||
private state: "buffering" | "started" | "unsubscribed" = "buffering";
|
||||
|
||||
constructor(
|
||||
snapshot: T | undefined,
|
||||
resnapshot: ((context: Context) => Promise<T>) | undefined,
|
||||
onError: (error: unknown, event: HarnessEvent, context: Context) => void | Promise<void>,
|
||||
) {
|
||||
this.snapshot = snapshot as T;
|
||||
this.resnapshotCallback = resnapshot;
|
||||
this.onError = onError;
|
||||
}
|
||||
|
||||
setSnapshot(snapshot: T): void {
|
||||
this.snapshot = snapshot;
|
||||
}
|
||||
|
||||
start(listener: EventListener): void {
|
||||
if (this.state !== "buffering") throw new Error("WatchHandle.start() may be called only once");
|
||||
this.state = "started";
|
||||
this.listener = listener;
|
||||
const buffered = this.buffer;
|
||||
this.buffer = [];
|
||||
for (const bufferedEvent of buffered) {
|
||||
this.enqueue(bufferedEvent.event, bufferedEvent.context, bufferedEvent.epoch);
|
||||
}
|
||||
}
|
||||
|
||||
async resnapshot(context: Context): Promise<T> {
|
||||
if (this.state === "unsubscribed") throw new Error("WatchHandle is unsubscribed");
|
||||
if (this.resnapshotCallback === undefined) throw new Error("WatchHandle does not support resnapshot");
|
||||
if (this.resnapshotState !== undefined) throw new Error("WatchHandle resnapshot is already in progress");
|
||||
let resolveReached!: () => void;
|
||||
const reached = new Promise<void>((resolve) => {
|
||||
resolveReached = resolve;
|
||||
});
|
||||
const resnapshotState = {
|
||||
phase: "dropping" as const,
|
||||
held: [] as Array<{ event: HarnessEvent; context: Context }>,
|
||||
reached,
|
||||
resolveReached,
|
||||
};
|
||||
this.epoch++;
|
||||
this.resnapshotState = resnapshotState;
|
||||
try {
|
||||
const snapshot = await this.resnapshotCallback(context);
|
||||
await reached;
|
||||
this.snapshot = snapshot;
|
||||
this.resnapshotState = undefined;
|
||||
for (const held of resnapshotState.held) this.push(held.event, held.context);
|
||||
return snapshot;
|
||||
} catch (error) {
|
||||
this.resnapshotState = undefined;
|
||||
for (const held of resnapshotState.held) this.push(held.event, held.context);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
markResnapshotBoundary(): void {
|
||||
const resnapshot = this.resnapshotState;
|
||||
if (resnapshot === undefined || resnapshot.phase !== "dropping") return;
|
||||
resnapshot.phase = "holding";
|
||||
resnapshot.resolveReached();
|
||||
}
|
||||
|
||||
unsubscribe(): void {
|
||||
if (this.state === "unsubscribed") return;
|
||||
this.state = "unsubscribed";
|
||||
this.buffer = [];
|
||||
this.listener = undefined;
|
||||
this.unsubscribeCallback?.();
|
||||
this.unsubscribeCallback = undefined;
|
||||
}
|
||||
|
||||
push(event: HarnessEvent, context: Context): void {
|
||||
if (this.state === "unsubscribed") return;
|
||||
if (this.resnapshotState?.phase === "dropping") return;
|
||||
if (this.resnapshotState?.phase === "holding") {
|
||||
this.resnapshotState.held.push({ event, context });
|
||||
return;
|
||||
}
|
||||
if (this.state === "buffering") {
|
||||
this.buffer.push({ event, context, epoch: this.epoch });
|
||||
return;
|
||||
}
|
||||
this.enqueue(event, context, this.epoch);
|
||||
}
|
||||
|
||||
setUnsubscribe(callback: () => void): void {
|
||||
this.unsubscribeCallback = callback;
|
||||
}
|
||||
|
||||
private enqueue(event: HarnessEvent, context: Context, epoch: number): void {
|
||||
const listener = this.listener;
|
||||
if (listener === undefined) return;
|
||||
this.deliveryTail = this.deliveryTail
|
||||
.then(async () => {
|
||||
if (this.state === "started" && epoch === this.epoch) await listener(event, context);
|
||||
})
|
||||
.catch(async (error) => {
|
||||
try {
|
||||
await this.onError(error, event, context);
|
||||
} catch {}
|
||||
});
|
||||
}
|
||||
}
|
||||
175
third_party/pi/packages/agent/src/harness/execution/assistant.ts
vendored
Normal file
175
third_party/pi/packages/agent/src/harness/execution/assistant.ts
vendored
Normal file
@@ -0,0 +1,175 @@
|
||||
import type {
|
||||
Context as AiContext,
|
||||
Api,
|
||||
AssistantMessage,
|
||||
AssistantMessageEvent,
|
||||
AssistantMessageEventStream,
|
||||
Message,
|
||||
Model,
|
||||
SimpleStreamOptions,
|
||||
Tool,
|
||||
} from "@earendil-works/pi-ai";
|
||||
import type { AgentMessage, ThinkingLevel } from "../../types.ts";
|
||||
import { type Context, getTelemetryContext } from "../context.ts";
|
||||
import type { SettledAssistantMessage } from "../session/types.ts";
|
||||
import type { AgentHarnessStreamOptions } from "../types.ts";
|
||||
import { AbortRequested } from "./effect-gate.ts";
|
||||
|
||||
/** HTTP response metadata captured before the provider response body is consumed. */
|
||||
export interface AssistantResponseMetadata {
|
||||
status?: number;
|
||||
headers?: Record<string, string>;
|
||||
}
|
||||
|
||||
/** Process-local lifecycle observer for one assistant stream. */
|
||||
export interface AssistantStreamObserver {
|
||||
start(
|
||||
message: AssistantMessage,
|
||||
event: Extract<AssistantMessageEvent, { type: "start" }>,
|
||||
context: Context,
|
||||
): void | Promise<void>;
|
||||
update(message: AssistantMessage, event: AssistantMessageEvent, context: Context): void | Promise<void>;
|
||||
end(message: SettledAssistantMessage, context: Context): void | Promise<void>;
|
||||
}
|
||||
|
||||
/** Executable inputs for one already-approved assistant provider request. */
|
||||
export interface HarnessAssistantStreamConfig {
|
||||
model: Model<Api>;
|
||||
systemPrompt: string;
|
||||
tools?: Tool[];
|
||||
thinkingLevel: ThinkingLevel;
|
||||
streamOptions: AgentHarnessStreamOptions;
|
||||
transformContext?: (
|
||||
requestContext: { messages: AgentMessage[]; systemPrompt: string },
|
||||
context: Context,
|
||||
) => Promise<{ messages: AgentMessage[]; systemPrompt: string }>;
|
||||
toProviderMessages: (messages: AgentMessage[], context: Context) => Message[] | Promise<Message[]>;
|
||||
beforePayload?: (
|
||||
payload: unknown,
|
||||
model: Model<Api>,
|
||||
context: Context,
|
||||
) => unknown | undefined | Promise<unknown | undefined>;
|
||||
afterResponse?: (
|
||||
message: SettledAssistantMessage,
|
||||
metadata: AssistantResponseMetadata,
|
||||
context: Context,
|
||||
) => Promise<SettledAssistantMessage>;
|
||||
request(
|
||||
aiContext: AiContext,
|
||||
options: SimpleStreamOptions,
|
||||
context: Context,
|
||||
): AssistantMessageEventStream | Promise<AssistantMessageEventStream>;
|
||||
observer: AssistantStreamObserver;
|
||||
}
|
||||
|
||||
function createRequestOptions(
|
||||
config: HarnessAssistantStreamConfig,
|
||||
captureMetadata: (metadata: AssistantResponseMetadata) => void,
|
||||
context: Context,
|
||||
): SimpleStreamOptions {
|
||||
const options = config.streamOptions;
|
||||
return {
|
||||
transport: options.transport,
|
||||
timeoutMs: options.timeoutMs,
|
||||
maxRetries: options.maxRetries,
|
||||
maxRetryDelayMs: options.maxRetryDelayMs,
|
||||
headers: options.headers,
|
||||
metadata: options.metadata,
|
||||
cacheRetention: options.cacheRetention,
|
||||
deferred: options.deferred,
|
||||
...(config.thinkingLevel === "off" ? {} : { reasoning: config.thinkingLevel }),
|
||||
signal: context.abortSignal,
|
||||
telemetryContext: getTelemetryContext(context),
|
||||
onPayload:
|
||||
config.beforePayload === undefined
|
||||
? undefined
|
||||
: (payload, model) => config.beforePayload?.(payload, model, context),
|
||||
onResponse: (response) => {
|
||||
captureMetadata({ status: response.status, headers: response.headers });
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function isUpdateEvent(
|
||||
event: AssistantMessageEvent,
|
||||
): event is Exclude<AssistantMessageEvent, { type: "start" | "done" | "error" }> {
|
||||
return event.type !== "start" && event.type !== "done" && event.type !== "error";
|
||||
}
|
||||
|
||||
export async function consumeAssistantStream(
|
||||
stream: AssistantMessageEventStream,
|
||||
observer: AssistantStreamObserver,
|
||||
afterResponse:
|
||||
| ((message: SettledAssistantMessage, context: Context) => Promise<SettledAssistantMessage>)
|
||||
| undefined,
|
||||
context: Context,
|
||||
): Promise<SettledAssistantMessage> {
|
||||
let started = false;
|
||||
for await (const event of stream) {
|
||||
if (event.type === "start") {
|
||||
if (started) throw new Error("Assistant message stream emitted more than one start event");
|
||||
started = true;
|
||||
await observer.start({ ...event.partial }, event, context);
|
||||
} else if (isUpdateEvent(event)) {
|
||||
if (!started) throw new Error(`Assistant message stream emitted ${event.type} before start`);
|
||||
await observer.update({ ...event.partial }, event, context);
|
||||
} else if (event.type === "done" && !started) {
|
||||
throw new Error("Assistant message stream emitted done before start");
|
||||
}
|
||||
}
|
||||
|
||||
const settled = (await stream.result()) as SettledAssistantMessage;
|
||||
let finalMessage = settled;
|
||||
if (afterResponse !== undefined) {
|
||||
try {
|
||||
finalMessage = await afterResponse(settled, context);
|
||||
} catch (error) {
|
||||
if (!(error instanceof AbortRequested)) throw error;
|
||||
await error.cancellation;
|
||||
}
|
||||
}
|
||||
await observer.end(finalMessage, context);
|
||||
return finalMessage;
|
||||
}
|
||||
|
||||
/** Stream one assistant response without mutating the caller's message list. */
|
||||
export async function streamHarnessAssistant(
|
||||
messages: AgentMessage[],
|
||||
config: HarnessAssistantStreamConfig,
|
||||
context: Context,
|
||||
): Promise<SettledAssistantMessage> {
|
||||
let requestContext = { messages: messages.slice(), systemPrompt: config.systemPrompt };
|
||||
if (config.transformContext) {
|
||||
requestContext = await config.transformContext(requestContext, context);
|
||||
}
|
||||
|
||||
const providerMessages = await config.toProviderMessages(requestContext.messages, context);
|
||||
const aiContext: AiContext = {
|
||||
systemPrompt: requestContext.systemPrompt,
|
||||
messages: providerMessages,
|
||||
tools: config.tools,
|
||||
};
|
||||
|
||||
let metadata: AssistantResponseMetadata = {};
|
||||
const stream = await config.request(
|
||||
aiContext,
|
||||
createRequestOptions(
|
||||
config,
|
||||
(nextMetadata) => {
|
||||
metadata = nextMetadata;
|
||||
},
|
||||
context,
|
||||
),
|
||||
context,
|
||||
);
|
||||
|
||||
const afterResponse = config.afterResponse;
|
||||
return consumeAssistantStream(
|
||||
stream,
|
||||
config.observer,
|
||||
afterResponse === undefined
|
||||
? undefined
|
||||
: (message, afterContext) => afterResponse(message, metadata, afterContext),
|
||||
context,
|
||||
);
|
||||
}
|
||||
64
third_party/pi/packages/agent/src/harness/execution/effect-gate.ts
vendored
Normal file
64
third_party/pi/packages/agent/src/harness/execution/effect-gate.ts
vendored
Normal file
@@ -0,0 +1,64 @@
|
||||
/** Expected internal control flow when cancellation wins effect admission. */
|
||||
export class AbortRequested extends Error {
|
||||
readonly cancellation: Promise<void>;
|
||||
|
||||
constructor(cancellation: Promise<void>) {
|
||||
super("Abort requested");
|
||||
this.name = "AbortRequested";
|
||||
this.cancellation = cancellation;
|
||||
}
|
||||
}
|
||||
|
||||
/** Procedure-facing synchronous admission capability for one drive pass. */
|
||||
export interface Gate {
|
||||
readonly signal: AbortSignal;
|
||||
admit<T>(invoke: () => T): T;
|
||||
}
|
||||
|
||||
/** Owner-facing lifecycle controls for one drive pass. */
|
||||
export interface GateControl {
|
||||
beginAbort(cancellation: Promise<void>): void;
|
||||
signalAbort(): void;
|
||||
close(error: Error): void;
|
||||
}
|
||||
|
||||
type GateState =
|
||||
| { status: "open" }
|
||||
| { status: "aborting"; cancellation: Promise<void> }
|
||||
| { status: "closed"; error: Error };
|
||||
|
||||
/** Create separate procedure-facing and owner-facing views of one effect gate. */
|
||||
export function createGate(): { gate: Gate; control: GateControl } {
|
||||
let state: GateState = { status: "open" };
|
||||
const controller = new AbortController();
|
||||
|
||||
const check = (): void => {
|
||||
if (state.status === "aborting") throw new AbortRequested(state.cancellation);
|
||||
if (state.status === "closed") throw state.error;
|
||||
};
|
||||
|
||||
return {
|
||||
gate: {
|
||||
admit<T>(invoke: () => T): T {
|
||||
check();
|
||||
return invoke();
|
||||
},
|
||||
signal: controller.signal,
|
||||
},
|
||||
control: {
|
||||
beginAbort(cancellation) {
|
||||
if (state.status !== "open") return;
|
||||
state = { status: "aborting", cancellation };
|
||||
},
|
||||
signalAbort() {
|
||||
if (state.status !== "aborting" || controller.signal.aborted) return;
|
||||
controller.abort(new AbortRequested(state.cancellation));
|
||||
},
|
||||
close(error) {
|
||||
if (state.status === "closed") return;
|
||||
state = { status: "closed", error };
|
||||
if (!controller.signal.aborted) controller.abort(error);
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
210
third_party/pi/packages/agent/src/harness/execution/tools.ts
vendored
Normal file
210
third_party/pi/packages/agent/src/harness/execution/tools.ts
vendored
Normal file
@@ -0,0 +1,210 @@
|
||||
import { type ToolResultMessage, validateToolArguments } from "@earendil-works/pi-ai";
|
||||
import type { AgentToolCall, AgentToolResult } from "../../types.ts";
|
||||
import { type Context, withAbortSignal } from "../context.ts";
|
||||
import type { JsonValue } from "../session/types.ts";
|
||||
import type { AgentHarnessTool, AgentHarnessToolInvocation, AgentHarnessToolUpdateCallback } from "../types.ts";
|
||||
import type { Gate } from "./effect-gate.ts";
|
||||
|
||||
/** A tool call whose tool exists and whose prepared arguments passed validation. */
|
||||
export interface PreparedToolCall<TContext extends object | undefined> {
|
||||
toolCall: AgentToolCall;
|
||||
tool: AgentHarnessTool<TContext>;
|
||||
args: Record<string, JsonValue>;
|
||||
}
|
||||
|
||||
/** Synthetic result produced without crossing the external tool-effect boundary. */
|
||||
export interface ImmediateToolOutcome {
|
||||
kind: "immediate";
|
||||
toolCall: AgentToolCall;
|
||||
result: AgentToolResult<unknown>;
|
||||
isError: true;
|
||||
terminate: boolean;
|
||||
}
|
||||
|
||||
/** Aggregated decision from the before-tool hook pipeline. */
|
||||
export interface BeforeToolDecision {
|
||||
args?: Record<string, JsonValue>;
|
||||
block?: { reason: string; terminate?: boolean };
|
||||
}
|
||||
|
||||
/** A prepared call cleared for durable intent publication and execution. */
|
||||
export interface ClearedToolCall<TContext extends object | undefined> {
|
||||
toolCall: AgentToolCall;
|
||||
tool: AgentHarnessTool<TContext>;
|
||||
args: Record<string, JsonValue>;
|
||||
}
|
||||
|
||||
/** Raw phase-two tool output before after-tool patching. */
|
||||
export interface ExecutedToolCall {
|
||||
result: AgentToolResult<unknown>;
|
||||
isError: boolean;
|
||||
}
|
||||
|
||||
/** Aggregated patch from the after-tool hook pipeline. */
|
||||
export interface AfterToolPatch {
|
||||
content?: AgentToolResult<unknown>["content"];
|
||||
details?: JsonValue;
|
||||
isError?: boolean;
|
||||
usage?: AgentToolResult<unknown>["usage"];
|
||||
terminate?: boolean;
|
||||
}
|
||||
|
||||
/** Final tool output ready to become a durable tool-result message. */
|
||||
export interface FinalizedToolCall {
|
||||
toolCall: AgentToolCall;
|
||||
result: AgentToolResult<unknown>;
|
||||
isError: boolean;
|
||||
terminate: boolean;
|
||||
}
|
||||
|
||||
function createErrorToolResult(message: string): AgentToolResult<unknown> {
|
||||
return {
|
||||
content: [{ type: "text", text: message }],
|
||||
details: undefined,
|
||||
};
|
||||
}
|
||||
|
||||
function immediateError(toolCall: AgentToolCall, message: string, terminate = false): ImmediateToolOutcome {
|
||||
return {
|
||||
kind: "immediate",
|
||||
toolCall,
|
||||
result: createErrorToolResult(message),
|
||||
isError: true,
|
||||
terminate,
|
||||
};
|
||||
}
|
||||
|
||||
/** Resolve a tool, apply its deterministic argument preparation, and validate the result. */
|
||||
export function prepareToolCall<TContext extends object | undefined>(
|
||||
call: AgentToolCall,
|
||||
tools: AgentHarnessTool<TContext>[],
|
||||
): PreparedToolCall<TContext> | ImmediateToolOutcome {
|
||||
const tool = tools.find((candidate) => candidate.name === call.name);
|
||||
if (!tool) {
|
||||
return immediateError(call, `Tool ${JSON.stringify(call.name)} is unavailable`);
|
||||
}
|
||||
|
||||
try {
|
||||
const preparedArguments = tool.prepareArguments ? tool.prepareArguments(call.arguments) : call.arguments;
|
||||
const preparedCall: AgentToolCall =
|
||||
preparedArguments === call.arguments
|
||||
? call
|
||||
: { ...call, arguments: preparedArguments as Record<string, JsonValue> };
|
||||
const args = validateToolArguments(tool, preparedCall) as Record<string, JsonValue>;
|
||||
return { toolCall: call, tool, args };
|
||||
} catch (error) {
|
||||
return immediateError(call, error instanceof Error ? error.message : String(error));
|
||||
}
|
||||
}
|
||||
|
||||
/** Apply an explicit hook decision and revalidate replacement arguments. */
|
||||
export function applyBeforeToolDecision<TContext extends object | undefined>(
|
||||
prepared: PreparedToolCall<TContext>,
|
||||
decision: BeforeToolDecision | undefined,
|
||||
): ClearedToolCall<TContext> | ImmediateToolOutcome {
|
||||
if (decision?.block) {
|
||||
return immediateError(prepared.toolCall, decision.block.reason, decision.block.terminate === true);
|
||||
}
|
||||
|
||||
if (!decision?.args) {
|
||||
return { toolCall: prepared.toolCall, tool: prepared.tool, args: prepared.args };
|
||||
}
|
||||
|
||||
try {
|
||||
const validatedArgs = validateToolArguments(prepared.tool, {
|
||||
...prepared.toolCall,
|
||||
arguments: decision.args,
|
||||
}) as Record<string, JsonValue>;
|
||||
return { toolCall: prepared.toolCall, tool: prepared.tool, args: validatedArgs };
|
||||
} catch (error) {
|
||||
return immediateError(prepared.toolCall, error instanceof Error ? error.message : String(error));
|
||||
}
|
||||
}
|
||||
|
||||
/** Execute one cleared external tool effect, converting expected tool throws to error output. */
|
||||
export function executeToolCall<TContext extends object | undefined>(
|
||||
call: ClearedToolCall<TContext>,
|
||||
gate: Gate,
|
||||
onUpdate: AgentHarnessToolUpdateCallback<unknown>,
|
||||
toolContext: TContext,
|
||||
invocation: AgentHarnessToolInvocation,
|
||||
context: Context,
|
||||
): Promise<ExecutedToolCall> {
|
||||
let acceptingUpdates = true;
|
||||
return gate.admit(async () => {
|
||||
const admittedContext = withAbortSignal(gate.signal, context);
|
||||
admittedContext.abortSignal?.throwIfAborted();
|
||||
try {
|
||||
const result = await call.tool.execute(
|
||||
call.toolCall.id,
|
||||
call.args,
|
||||
(partial, options) => {
|
||||
if (acceptingUpdates) onUpdate(partial, options);
|
||||
},
|
||||
toolContext,
|
||||
invocation,
|
||||
admittedContext,
|
||||
);
|
||||
return { result, isError: false };
|
||||
} catch (error) {
|
||||
return {
|
||||
result: createErrorToolResult(error instanceof Error ? error.message : String(error)),
|
||||
isError: true,
|
||||
};
|
||||
} finally {
|
||||
acceptingUpdates = false;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/** Apply an after-tool patch field by field. */
|
||||
export function finalizeToolCall<TContext extends object | undefined>(
|
||||
call: ClearedToolCall<TContext>,
|
||||
executed: ExecutedToolCall,
|
||||
patch: AfterToolPatch | undefined,
|
||||
): FinalizedToolCall {
|
||||
const result: AgentToolResult<unknown> = patch
|
||||
? {
|
||||
...executed.result,
|
||||
content: patch.content === undefined ? executed.result.content : patch.content,
|
||||
details: patch.details === undefined ? executed.result.details : patch.details,
|
||||
usage: patch.usage === undefined ? executed.result.usage : patch.usage,
|
||||
terminate: patch.terminate === undefined ? executed.result.terminate : patch.terminate,
|
||||
}
|
||||
: executed.result;
|
||||
return {
|
||||
toolCall: call.toolCall,
|
||||
result,
|
||||
isError: patch?.isError ?? executed.isError,
|
||||
terminate: result.terminate === true,
|
||||
};
|
||||
}
|
||||
|
||||
/** Reconstruct the canonical tool result represented by a staged transcript message. */
|
||||
export function toolResultFromMessage(
|
||||
message: ToolResultMessage<unknown>,
|
||||
terminate: boolean,
|
||||
): AgentToolResult<unknown> {
|
||||
return {
|
||||
content: message.content,
|
||||
details: message.details,
|
||||
...(message.usage === undefined ? {} : { usage: message.usage }),
|
||||
...(message.addedToolNames === undefined ? {} : { addedToolNames: message.addedToolNames }),
|
||||
...(terminate ? { terminate: true } : {}),
|
||||
};
|
||||
}
|
||||
|
||||
/** Convert finalized tool output to the provider-facing transcript message. */
|
||||
export function createToolResultMessage(call: FinalizedToolCall): ToolResultMessage {
|
||||
return {
|
||||
role: "toolResult",
|
||||
toolCallId: call.toolCall.id,
|
||||
toolName: call.toolCall.name,
|
||||
content: call.result.content ?? [],
|
||||
...(call.result.details === undefined ? {} : { details: call.result.details }),
|
||||
...(call.result.usage === undefined ? {} : { usage: call.result.usage }),
|
||||
...(call.result.addedToolNames?.length ? { addedToolNames: call.result.addedToolNames } : {}),
|
||||
isError: call.isError,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
533
third_party/pi/packages/agent/src/harness/hooks.ts
vendored
Normal file
533
third_party/pi/packages/agent/src/harness/hooks.ts
vendored
Normal file
@@ -0,0 +1,533 @@
|
||||
import type { HookHandler, HookInvocation, HookMap, HookName, Hooks } from "./agent-harness.ts";
|
||||
import { type Context, withAbortSignal } from "./context.ts";
|
||||
import type { Gate } from "./execution/effect-gate.ts";
|
||||
import { startHarnessSpan } from "./telemetry.ts";
|
||||
import type { AgentHarnessStreamOptions, AgentHarnessStreamOptionsPatch } from "./types.ts";
|
||||
|
||||
interface HookRegistration {
|
||||
id?: string;
|
||||
handler: (event: unknown, context: Context) => unknown | Promise<unknown>;
|
||||
}
|
||||
|
||||
type HookErrorReporter = (error: Error, hook: HookName, lane: string, context: Context) => void | Promise<void>;
|
||||
|
||||
/** Ordered harness hook registry and aggregate runner. */
|
||||
export class HookRegistry implements Hooks {
|
||||
private readonly registrations = new Map<HookName, HookRegistration[]>();
|
||||
private readonly reportError: HookErrorReporter;
|
||||
private closedError: Error | undefined;
|
||||
|
||||
constructor(reportError: HookErrorReporter) {
|
||||
this.reportError = reportError;
|
||||
}
|
||||
|
||||
on<TName extends HookName>(name: TName, handler: HookHandler<TName>, options: { id?: string } = {}): () => void {
|
||||
if (this.closedError !== undefined) throw this.closedError;
|
||||
const registrations = this.registrations.get(name) ?? [];
|
||||
const registration: HookRegistration = {
|
||||
...(options.id === undefined ? {} : { id: options.id }),
|
||||
handler: (event, context) => handler(event as HookInvocation<TName>, context),
|
||||
};
|
||||
registrations.push(registration);
|
||||
this.registrations.set(name, registrations);
|
||||
return () => {
|
||||
const index = registrations.indexOf(registration);
|
||||
if (index !== -1) registrations.splice(index, 1);
|
||||
};
|
||||
}
|
||||
|
||||
has(name: HookName): boolean {
|
||||
return (this.registrations.get(name)?.length ?? 0) !== 0;
|
||||
}
|
||||
|
||||
/** Invoke one accepted-operation aggregate after synchronously passing its effect gate. */
|
||||
runWithGate<TName extends HookName>(
|
||||
name: TName,
|
||||
event: HookInvocation<TName>,
|
||||
gate: Gate,
|
||||
context: Context,
|
||||
): Promise<HookMap[TName]["result"]> {
|
||||
return gate.admit(() => {
|
||||
const admittedContext = withAbortSignal(gate.signal, context);
|
||||
admittedContext.abortSignal?.throwIfAborted();
|
||||
return this.runAdmitted(name, event, admittedContext);
|
||||
});
|
||||
}
|
||||
|
||||
/** Invoke a tool-hook aggregate with one telemetry span per registered handler. */
|
||||
runToolWithGate<TName extends "before_tool" | "after_tool">(
|
||||
name: TName,
|
||||
event: HookInvocation<TName>,
|
||||
gate: Gate,
|
||||
context: Context,
|
||||
): Promise<HookMap[TName]["result"]> {
|
||||
return gate.admit(() => {
|
||||
const admittedContext = withAbortSignal(gate.signal, context);
|
||||
admittedContext.abortSignal?.throwIfAborted();
|
||||
return (
|
||||
name === "before_tool"
|
||||
? this.beforeTool(event as HookInvocation<"before_tool">, admittedContext)
|
||||
: this.afterTool(event as HookInvocation<"after_tool">, admittedContext)
|
||||
) as Promise<HookMap[TName]["result"]>;
|
||||
});
|
||||
}
|
||||
|
||||
close(error: Error): void {
|
||||
this.closedError ??= error;
|
||||
}
|
||||
|
||||
private async runAdmitted<TName extends HookName>(
|
||||
name: TName,
|
||||
event: HookInvocation<TName>,
|
||||
context: Context,
|
||||
): Promise<HookMap[TName]["result"]> {
|
||||
if (this.closedError !== undefined) throw this.closedError;
|
||||
const result = await this.aggregate(name, event, context);
|
||||
return result as HookMap[TName]["result"];
|
||||
}
|
||||
|
||||
private async aggregate(name: HookName, event: HookInvocation<HookName>, context: Context): Promise<unknown> {
|
||||
switch (name) {
|
||||
case "before_run":
|
||||
return this.beforeRun(event as HookInvocation<"before_run">, context);
|
||||
case "before_drive":
|
||||
await this.invokeAllFailClosed(name, event as HookInvocation<"before_drive">, context);
|
||||
return undefined;
|
||||
case "before_run_end": {
|
||||
let followUp: string | undefined;
|
||||
await this.invokeAll(
|
||||
name,
|
||||
event,
|
||||
(value) => {
|
||||
const result = value as HookMap["before_run_end"]["result"];
|
||||
if (result?.followUp !== undefined) followUp = result.followUp;
|
||||
},
|
||||
context,
|
||||
);
|
||||
return followUp === undefined ? undefined : { followUp };
|
||||
}
|
||||
case "transform_context":
|
||||
return this.transformContext(event as HookInvocation<"transform_context">, context);
|
||||
case "before_request":
|
||||
return this.beforeRequest(event as HookInvocation<"before_request">, context);
|
||||
case "before_payload":
|
||||
return this.beforePayload(event as HookInvocation<"before_payload">, context);
|
||||
case "after_response":
|
||||
return this.afterResponse(event as HookInvocation<"after_response">, context);
|
||||
case "before_tool":
|
||||
return this.beforeTool(event as HookInvocation<"before_tool">, context);
|
||||
case "after_tool":
|
||||
return this.afterTool(event as HookInvocation<"after_tool">, context);
|
||||
case "before_compaction":
|
||||
return this.firstStructural(name, event as HookInvocation<"before_compaction">, "compaction", context);
|
||||
case "before_navigation":
|
||||
return this.firstStructural(name, event as HookInvocation<"before_navigation">, "summary", context);
|
||||
}
|
||||
}
|
||||
|
||||
private async beforeRun(
|
||||
event: HookInvocation<"before_run">,
|
||||
context: Context,
|
||||
): Promise<HookMap["before_run"]["result"]> {
|
||||
let prompt = event.prompt;
|
||||
let injected: HookMap["before_run"]["event"]["prompt"] = [];
|
||||
for (const registration of this.registrationsFor("before_run")) {
|
||||
try {
|
||||
const result = (await registration.handler(
|
||||
{ ...event, prompt },
|
||||
context,
|
||||
)) as HookMap["before_run"]["result"];
|
||||
if (result?.messages !== undefined) {
|
||||
injected = [...injected, ...result.messages];
|
||||
prompt = [...prompt, ...result.messages];
|
||||
}
|
||||
} catch (error) {
|
||||
await this.reportError(
|
||||
error instanceof Error ? error : new Error(String(error)),
|
||||
"before_run",
|
||||
event.lane,
|
||||
context,
|
||||
);
|
||||
}
|
||||
}
|
||||
return injected.length === 0 ? undefined : { messages: injected };
|
||||
}
|
||||
|
||||
private async beforeTool(
|
||||
event: HookInvocation<"before_tool">,
|
||||
context: Context,
|
||||
): Promise<HookMap["before_tool"]["result"]> {
|
||||
let args = event.args;
|
||||
let block: { reason: string; terminate?: boolean } | undefined;
|
||||
for (const registration of this.registrationsFor("before_tool")) {
|
||||
try {
|
||||
const result = (await this.invokeToolRegistration(
|
||||
"before_tool",
|
||||
registration,
|
||||
{ ...event, args },
|
||||
context,
|
||||
)) as HookMap["before_tool"]["result"];
|
||||
if (result?.args !== undefined) args = result.args;
|
||||
if (result?.block !== undefined) {
|
||||
block = result.block;
|
||||
break;
|
||||
}
|
||||
} catch (error) {
|
||||
const normalized = error instanceof Error ? error : new Error(String(error));
|
||||
await this.reportError(normalized, "before_tool", event.lane, context);
|
||||
block = { reason: normalized.message };
|
||||
break;
|
||||
}
|
||||
}
|
||||
return {
|
||||
...(args === event.args ? {} : { args }),
|
||||
...(block === undefined ? {} : { block }),
|
||||
};
|
||||
}
|
||||
|
||||
private async transformContext(
|
||||
event: HookInvocation<"transform_context">,
|
||||
context: Context,
|
||||
): Promise<HookMap["transform_context"]["result"]> {
|
||||
let messages = event.messages;
|
||||
let systemPrompt = event.systemPrompt;
|
||||
for (const registration of this.registrationsFor("transform_context")) {
|
||||
try {
|
||||
const result = (await registration.handler(
|
||||
{
|
||||
...event,
|
||||
messages,
|
||||
systemPrompt,
|
||||
},
|
||||
context,
|
||||
)) as HookMap["transform_context"]["result"];
|
||||
if (result?.messages !== undefined) messages = result.messages;
|
||||
if (result?.systemPrompt !== undefined) systemPrompt = result.systemPrompt;
|
||||
} catch (error) {
|
||||
await this.reportError(
|
||||
error instanceof Error ? error : new Error(String(error)),
|
||||
"transform_context",
|
||||
event.lane,
|
||||
context,
|
||||
);
|
||||
}
|
||||
}
|
||||
return { messages, systemPrompt };
|
||||
}
|
||||
|
||||
private async beforeRequest(
|
||||
event: HookInvocation<"before_request">,
|
||||
context: Context,
|
||||
): Promise<HookMap["before_request"]["result"]> {
|
||||
let streamOptions = event.streamOptions;
|
||||
let changed = false;
|
||||
for (const registration of this.registrationsFor("before_request")) {
|
||||
try {
|
||||
const result = (await registration.handler(
|
||||
{ ...event, streamOptions },
|
||||
context,
|
||||
)) as HookMap["before_request"]["result"];
|
||||
if (result?.streamOptions !== undefined) {
|
||||
streamOptions = applyStreamOptionsPatch(streamOptions, result.streamOptions);
|
||||
changed = true;
|
||||
}
|
||||
} catch (error) {
|
||||
await this.reportError(
|
||||
error instanceof Error ? error : new Error(String(error)),
|
||||
"before_request",
|
||||
event.lane,
|
||||
context,
|
||||
);
|
||||
}
|
||||
}
|
||||
return changed ? { streamOptions: createStreamOptionsPatch(event.streamOptions, streamOptions) } : undefined;
|
||||
}
|
||||
|
||||
private async beforePayload(
|
||||
event: HookInvocation<"before_payload">,
|
||||
context: Context,
|
||||
): Promise<HookMap["before_payload"]["result"]> {
|
||||
let payload = event.payload;
|
||||
for (const registration of this.registrationsFor("before_payload")) {
|
||||
try {
|
||||
const result = (await registration.handler(
|
||||
{ ...event, payload },
|
||||
context,
|
||||
)) as HookMap["before_payload"]["result"];
|
||||
if (result?.payload !== undefined) payload = result.payload;
|
||||
} catch (error) {
|
||||
await this.reportError(
|
||||
error instanceof Error ? error : new Error(String(error)),
|
||||
"before_payload",
|
||||
event.lane,
|
||||
context,
|
||||
);
|
||||
}
|
||||
}
|
||||
return { payload };
|
||||
}
|
||||
|
||||
private async afterResponse(
|
||||
event: HookInvocation<"after_response">,
|
||||
context: Context,
|
||||
): Promise<HookMap["after_response"]["result"]> {
|
||||
let message = event.message;
|
||||
for (const registration of this.registrationsFor("after_response")) {
|
||||
try {
|
||||
const result = (await registration.handler(
|
||||
{ ...event, message },
|
||||
context,
|
||||
)) as HookMap["after_response"]["result"];
|
||||
if (result?.message !== undefined) message = result.message;
|
||||
} catch (error) {
|
||||
await this.reportError(
|
||||
error instanceof Error ? error : new Error(String(error)),
|
||||
"after_response",
|
||||
event.lane,
|
||||
context,
|
||||
);
|
||||
}
|
||||
}
|
||||
return { message };
|
||||
}
|
||||
|
||||
private async afterTool(
|
||||
event: HookInvocation<"after_tool">,
|
||||
context: Context,
|
||||
): Promise<HookMap["after_tool"]["result"]> {
|
||||
let current = {
|
||||
content: event.content,
|
||||
details: event.details,
|
||||
isError: event.isError,
|
||||
usage: event.usage,
|
||||
};
|
||||
const aggregate: NonNullable<HookMap["after_tool"]["result"]> = {};
|
||||
for (const registration of this.registrationsFor("after_tool")) {
|
||||
try {
|
||||
const result = (await this.invokeToolRegistration(
|
||||
"after_tool",
|
||||
registration,
|
||||
{ ...event, ...current },
|
||||
context,
|
||||
)) as HookMap["after_tool"]["result"];
|
||||
if (result === undefined) continue;
|
||||
if (result.content !== undefined) aggregate.content = result.content;
|
||||
if (result.details !== undefined) aggregate.details = result.details;
|
||||
if (result.isError !== undefined) aggregate.isError = result.isError;
|
||||
if (result.usage !== undefined) aggregate.usage = result.usage;
|
||||
if (result.terminate !== undefined) aggregate.terminate = result.terminate;
|
||||
current = {
|
||||
content: result.content === undefined ? current.content : result.content,
|
||||
details: result.details === undefined ? current.details : result.details,
|
||||
isError: result.isError === undefined ? current.isError : result.isError,
|
||||
usage: result.usage === undefined ? current.usage : result.usage,
|
||||
};
|
||||
} catch (error) {
|
||||
await this.reportError(
|
||||
error instanceof Error ? error : new Error(String(error)),
|
||||
"after_tool",
|
||||
event.lane,
|
||||
context,
|
||||
);
|
||||
}
|
||||
}
|
||||
return Object.keys(aggregate).length === 0 ? undefined : aggregate;
|
||||
}
|
||||
|
||||
private async firstStructural(
|
||||
name: "before_compaction" | "before_navigation",
|
||||
event: HookInvocation<"before_compaction"> | HookInvocation<"before_navigation">,
|
||||
resultField: "compaction" | "summary",
|
||||
context: Context,
|
||||
): Promise<unknown> {
|
||||
for (const registration of this.registrationsFor(name)) {
|
||||
try {
|
||||
const value = await registration.handler(event, context);
|
||||
if (value === undefined || value === null || typeof value !== "object") continue;
|
||||
const result = value as Record<string, unknown>;
|
||||
if (result.decline === true && result[resultField] !== undefined) {
|
||||
await this.reportError(
|
||||
new Error(`${name} hook cannot return both decline and ${resultField}`),
|
||||
name,
|
||||
event.lane,
|
||||
context,
|
||||
);
|
||||
continue;
|
||||
}
|
||||
if (result.decline === true || result[resultField] !== undefined) return value;
|
||||
} catch (error) {
|
||||
await this.reportError(
|
||||
error instanceof Error ? error : new Error(String(error)),
|
||||
name,
|
||||
event.lane,
|
||||
context,
|
||||
);
|
||||
}
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
private invokeToolRegistration(
|
||||
name: "before_tool" | "after_tool",
|
||||
registration: HookRegistration,
|
||||
event: HookInvocation<"before_tool"> | HookInvocation<"after_tool">,
|
||||
context: Context,
|
||||
): Promise<unknown> {
|
||||
return startHarnessSpan(
|
||||
"pi.harness.hook",
|
||||
{
|
||||
"pi.lane.name": event.lane,
|
||||
"pi.operation.id": event.runId,
|
||||
"pi.hook.name": name,
|
||||
...(registration.id === undefined ? {} : { "pi.hook.registration_id": registration.id }),
|
||||
},
|
||||
async (span, spanContext) => {
|
||||
try {
|
||||
const result = await registration.handler(event, spanContext);
|
||||
const blocked =
|
||||
name === "before_tool" &&
|
||||
result !== null &&
|
||||
typeof result === "object" &&
|
||||
"block" in result &&
|
||||
result.block !== undefined;
|
||||
span.setAttributes({ "pi.hook.outcome": blocked ? "blocked" : "completed" });
|
||||
return result;
|
||||
} catch (error) {
|
||||
span.setAttributes({ "pi.hook.outcome": "failed" });
|
||||
span.setStatus({ status: "error" });
|
||||
throw error;
|
||||
}
|
||||
},
|
||||
context,
|
||||
);
|
||||
}
|
||||
|
||||
private registrationsFor(name: HookName): HookRegistration[] {
|
||||
return [...(this.registrations.get(name) ?? [])];
|
||||
}
|
||||
|
||||
private async invokeAllFailClosed(
|
||||
name: "before_drive",
|
||||
event: HookInvocation<"before_drive">,
|
||||
context: Context,
|
||||
): Promise<void> {
|
||||
for (const registration of this.registrationsFor(name)) {
|
||||
try {
|
||||
await registration.handler(event, context);
|
||||
} catch (error) {
|
||||
const normalized = error instanceof Error ? error : new Error(String(error));
|
||||
await this.reportError(normalized, name, event.lane, context);
|
||||
throw normalized;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private async invokeAll(
|
||||
name: HookName,
|
||||
event: HookInvocation<HookName>,
|
||||
apply: (value: unknown) => void,
|
||||
context: Context,
|
||||
): Promise<void> {
|
||||
for (const registration of this.registrationsFor(name)) {
|
||||
try {
|
||||
apply(await registration.handler(event, context));
|
||||
} catch (error) {
|
||||
await this.reportError(
|
||||
error instanceof Error ? error : new Error(String(error)),
|
||||
name,
|
||||
event.lane,
|
||||
context,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export function applyStreamOptionsPatch(
|
||||
base: AgentHarnessStreamOptions,
|
||||
patch: AgentHarnessStreamOptionsPatch,
|
||||
): AgentHarnessStreamOptions {
|
||||
const next: AgentHarnessStreamOptions = { ...base };
|
||||
for (const key of [
|
||||
"transport",
|
||||
"timeoutMs",
|
||||
"maxRetries",
|
||||
"maxRetryDelayMs",
|
||||
"cacheRetention",
|
||||
"deferred",
|
||||
] as const) {
|
||||
if (!(key in patch)) continue;
|
||||
const value = patch[key];
|
||||
if (value === undefined) delete next[key];
|
||||
else Object.assign(next, { [key]: value });
|
||||
}
|
||||
if ("headers" in patch) {
|
||||
if (patch.headers === undefined) delete next.headers;
|
||||
else {
|
||||
const headers = { ...next.headers };
|
||||
for (const [key, value] of Object.entries(patch.headers)) {
|
||||
if (value === undefined) delete headers[key];
|
||||
else headers[key] = value;
|
||||
}
|
||||
next.headers = headers;
|
||||
}
|
||||
}
|
||||
if ("metadata" in patch) {
|
||||
if (patch.metadata === undefined) delete next.metadata;
|
||||
else {
|
||||
const metadata = { ...next.metadata };
|
||||
for (const [key, value] of Object.entries(patch.metadata)) {
|
||||
if (value === undefined) delete metadata[key];
|
||||
else metadata[key] = value;
|
||||
}
|
||||
next.metadata = metadata;
|
||||
}
|
||||
}
|
||||
return next;
|
||||
}
|
||||
|
||||
function createStreamOptionsPatch(
|
||||
base: AgentHarnessStreamOptions,
|
||||
value: AgentHarnessStreamOptions,
|
||||
): AgentHarnessStreamOptionsPatch {
|
||||
const patch: AgentHarnessStreamOptionsPatch = {};
|
||||
for (const key of [
|
||||
"transport",
|
||||
"timeoutMs",
|
||||
"maxRetries",
|
||||
"maxRetryDelayMs",
|
||||
"cacheRetention",
|
||||
"deferred",
|
||||
] as const) {
|
||||
if (base[key] !== value[key]) Object.assign(patch, { [key]: value[key] });
|
||||
}
|
||||
if (base.headers !== value.headers) {
|
||||
if (value.headers === undefined) patch.headers = undefined;
|
||||
else {
|
||||
const headers: Record<string, string | undefined> = {};
|
||||
for (const key of Object.keys(base.headers ?? {})) {
|
||||
if (!(key in value.headers)) headers[key] = undefined;
|
||||
}
|
||||
for (const [key, header] of Object.entries(value.headers)) {
|
||||
if (base.headers?.[key] !== header) headers[key] = header;
|
||||
}
|
||||
if (base.headers === undefined && Object.keys(headers).length === 0) patch.headers = {};
|
||||
else if (Object.keys(headers).length !== 0) patch.headers = headers;
|
||||
}
|
||||
}
|
||||
if (base.metadata !== value.metadata) {
|
||||
if (value.metadata === undefined) patch.metadata = undefined;
|
||||
else {
|
||||
const metadata: Record<string, unknown | undefined> = {};
|
||||
for (const key of Object.keys(base.metadata ?? {})) {
|
||||
if (!(key in value.metadata)) metadata[key] = undefined;
|
||||
}
|
||||
for (const [key, metadataValue] of Object.entries(value.metadata)) {
|
||||
if (base.metadata?.[key] !== metadataValue) metadata[key] = metadataValue;
|
||||
}
|
||||
if (base.metadata === undefined && Object.keys(metadata).length === 0) patch.metadata = {};
|
||||
else if (Object.keys(metadata).length !== 0) patch.metadata = metadata;
|
||||
}
|
||||
}
|
||||
return patch;
|
||||
}
|
||||
168
third_party/pi/packages/agent/src/harness/messages.ts
vendored
Normal file
168
third_party/pi/packages/agent/src/harness/messages.ts
vendored
Normal file
@@ -0,0 +1,168 @@
|
||||
import type { ImageContent, Message, TextContent } from "@earendil-works/pi-ai";
|
||||
import type { AgentMessage } from "../types.ts";
|
||||
|
||||
export const COMPACTION_SUMMARY_PREFIX = `The conversation history before this point was compacted into the following summary:
|
||||
|
||||
<summary>
|
||||
`;
|
||||
|
||||
export const COMPACTION_SUMMARY_SUFFIX = `
|
||||
</summary>`;
|
||||
|
||||
export const BRANCH_SUMMARY_PREFIX = `The following is a summary of a branch that this conversation came back from:
|
||||
|
||||
<summary>
|
||||
`;
|
||||
|
||||
export const BRANCH_SUMMARY_SUFFIX = `</summary>`;
|
||||
|
||||
export interface BashExecutionMessage {
|
||||
role: "bashExecution";
|
||||
command: string;
|
||||
output: string;
|
||||
exitCode: number | undefined;
|
||||
cancelled: boolean;
|
||||
truncated: boolean;
|
||||
fullOutputPath?: string;
|
||||
timestamp: number;
|
||||
excludeFromContext?: boolean;
|
||||
}
|
||||
|
||||
export interface CustomMessage<T = unknown> {
|
||||
role: "custom";
|
||||
customType: string;
|
||||
content: string | (TextContent | ImageContent)[];
|
||||
display: boolean;
|
||||
details?: T;
|
||||
timestamp: number;
|
||||
}
|
||||
|
||||
export interface BranchSummaryMessage {
|
||||
role: "branchSummary";
|
||||
summary: string;
|
||||
fromId: string | null;
|
||||
timestamp: number;
|
||||
}
|
||||
|
||||
export interface CompactionSummaryMessage {
|
||||
role: "compactionSummary";
|
||||
summary: string;
|
||||
tokensBefore: number;
|
||||
timestamp: number;
|
||||
}
|
||||
|
||||
declare module "../types.ts" {
|
||||
interface CustomAgentMessages {
|
||||
bashExecution: BashExecutionMessage;
|
||||
custom: CustomMessage;
|
||||
branchSummary: BranchSummaryMessage;
|
||||
compactionSummary: CompactionSummaryMessage;
|
||||
}
|
||||
}
|
||||
|
||||
export function bashExecutionToText(msg: BashExecutionMessage): string {
|
||||
let text = `Ran \`${msg.command}\`\n`;
|
||||
if (msg.output) {
|
||||
text += `\`\`\`\n${msg.output}\n\`\`\``;
|
||||
} else {
|
||||
text += "(no output)";
|
||||
}
|
||||
if (msg.cancelled) {
|
||||
text += "\n\n(command cancelled)";
|
||||
} else if (msg.exitCode !== null && msg.exitCode !== undefined && msg.exitCode !== 0) {
|
||||
text += `\n\nCommand exited with code ${msg.exitCode}`;
|
||||
}
|
||||
if (msg.truncated && msg.fullOutputPath) {
|
||||
text += `\n\n[Output truncated. Full output: ${msg.fullOutputPath}]`;
|
||||
}
|
||||
return text;
|
||||
}
|
||||
|
||||
export function createBranchSummaryMessage(
|
||||
summary: string,
|
||||
fromId: string | null,
|
||||
timestamp: string | number,
|
||||
): BranchSummaryMessage {
|
||||
return {
|
||||
role: "branchSummary",
|
||||
summary,
|
||||
fromId,
|
||||
timestamp: typeof timestamp === "number" ? timestamp : new Date(timestamp).getTime(),
|
||||
};
|
||||
}
|
||||
|
||||
export function createCompactionSummaryMessage(
|
||||
summary: string,
|
||||
tokensBefore: number,
|
||||
timestamp: string | number,
|
||||
): CompactionSummaryMessage {
|
||||
return {
|
||||
role: "compactionSummary",
|
||||
summary,
|
||||
tokensBefore,
|
||||
timestamp: typeof timestamp === "number" ? timestamp : new Date(timestamp).getTime(),
|
||||
};
|
||||
}
|
||||
|
||||
export function createCustomMessage(
|
||||
customType: string,
|
||||
content: string | (TextContent | ImageContent)[],
|
||||
display: boolean,
|
||||
details: unknown | undefined,
|
||||
timestamp: string | number,
|
||||
): CustomMessage {
|
||||
return {
|
||||
role: "custom",
|
||||
customType,
|
||||
content,
|
||||
display,
|
||||
details,
|
||||
timestamp: typeof timestamp === "number" ? timestamp : new Date(timestamp).getTime(),
|
||||
};
|
||||
}
|
||||
|
||||
export function convertToLlm(messages: AgentMessage[]): Message[] {
|
||||
return messages
|
||||
.map((m): Message | undefined => {
|
||||
switch (m.role) {
|
||||
case "bashExecution":
|
||||
if (m.excludeFromContext) {
|
||||
return undefined;
|
||||
}
|
||||
return {
|
||||
role: "user",
|
||||
content: [{ type: "text", text: bashExecutionToText(m) }],
|
||||
timestamp: m.timestamp,
|
||||
};
|
||||
case "custom": {
|
||||
const content = typeof m.content === "string" ? [{ type: "text" as const, text: m.content }] : m.content;
|
||||
return {
|
||||
role: "user",
|
||||
content,
|
||||
timestamp: m.timestamp,
|
||||
};
|
||||
}
|
||||
case "branchSummary":
|
||||
return {
|
||||
role: "user",
|
||||
content: [{ type: "text" as const, text: BRANCH_SUMMARY_PREFIX + m.summary + BRANCH_SUMMARY_SUFFIX }],
|
||||
timestamp: m.timestamp,
|
||||
};
|
||||
case "compactionSummary":
|
||||
return {
|
||||
role: "user",
|
||||
content: [
|
||||
{ type: "text" as const, text: COMPACTION_SUMMARY_PREFIX + m.summary + COMPACTION_SUMMARY_SUFFIX },
|
||||
],
|
||||
timestamp: m.timestamp,
|
||||
};
|
||||
case "user":
|
||||
case "assistant":
|
||||
case "toolResult":
|
||||
return m;
|
||||
default:
|
||||
return undefined;
|
||||
}
|
||||
})
|
||||
.filter((m): m is Message => m !== undefined);
|
||||
}
|
||||
270
third_party/pi/packages/agent/src/harness/prompt-templates.ts
vendored
Normal file
270
third_party/pi/packages/agent/src/harness/prompt-templates.ts
vendored
Normal file
@@ -0,0 +1,270 @@
|
||||
import { parse } from "yaml";
|
||||
import type { Context } from "./context.ts";
|
||||
import { type ExecutionEnv, type FileInfo, type PromptTemplate, type Result, toError } from "./types.ts";
|
||||
|
||||
export type PromptTemplateDiagnosticCode = "file_info_failed" | "list_failed" | "read_failed" | "parse_failed";
|
||||
|
||||
/** Warning produced while loading prompt templates. */
|
||||
export interface PromptTemplateDiagnostic {
|
||||
/** Diagnostic severity. Currently only warnings are emitted. */
|
||||
type: "warning";
|
||||
/** Stable diagnostic code. */
|
||||
code: PromptTemplateDiagnosticCode;
|
||||
/** Human-readable diagnostic message. */
|
||||
message: string;
|
||||
/** Path associated with the diagnostic. */
|
||||
path: string;
|
||||
}
|
||||
|
||||
interface PromptTemplateFrontmatter {
|
||||
description?: string;
|
||||
"argument-hint"?: string;
|
||||
[key: string]: unknown;
|
||||
}
|
||||
|
||||
/**
|
||||
* Load prompt templates from one or more paths.
|
||||
*
|
||||
* Directory inputs load direct `.md` children non-recursively. File inputs load explicit `.md` files. Missing paths and
|
||||
* non-markdown files are skipped. Read and parse failures are returned as diagnostics.
|
||||
*/
|
||||
export async function loadPromptTemplates(
|
||||
env: ExecutionEnv,
|
||||
paths: string | string[],
|
||||
context: Context,
|
||||
): Promise<{ promptTemplates: PromptTemplate[]; diagnostics: PromptTemplateDiagnostic[] }> {
|
||||
const promptTemplates: PromptTemplate[] = [];
|
||||
const diagnostics: PromptTemplateDiagnostic[] = [];
|
||||
for (const path of Array.isArray(paths) ? paths : [paths]) {
|
||||
const infoResult = await env.fileInfo(path, context);
|
||||
if (!infoResult.ok) {
|
||||
if (infoResult.error.code !== "not_found") {
|
||||
diagnostics.push({
|
||||
type: "warning",
|
||||
code: "file_info_failed",
|
||||
message: infoResult.error.message,
|
||||
path,
|
||||
});
|
||||
}
|
||||
continue;
|
||||
}
|
||||
const info = infoResult.value;
|
||||
const kind = await resolveKind(env, info, diagnostics, context);
|
||||
if (kind === "directory") {
|
||||
const result = await loadTemplatesFromDir(env, info.path, context);
|
||||
promptTemplates.push(...result.promptTemplates);
|
||||
diagnostics.push(...result.diagnostics);
|
||||
} else if (kind === "file" && info.name.endsWith(".md")) {
|
||||
const result = await loadTemplateFromFile(env, info.path, info.name, context);
|
||||
if (result.promptTemplate) promptTemplates.push(result.promptTemplate);
|
||||
diagnostics.push(...result.diagnostics);
|
||||
}
|
||||
}
|
||||
return { promptTemplates, diagnostics };
|
||||
}
|
||||
|
||||
/**
|
||||
* Load prompt templates from source-tagged paths.
|
||||
*
|
||||
* Source values are preserved exactly and attached to every loaded prompt template and diagnostic. The agent package does
|
||||
* not interpret source values; applications define their own provenance shape.
|
||||
*/
|
||||
export async function loadSourcedPromptTemplates<TSource, TPromptTemplate extends PromptTemplate = PromptTemplate>(
|
||||
env: ExecutionEnv,
|
||||
inputs: Array<{ path: string; source: TSource }>,
|
||||
mapPromptTemplate:
|
||||
| ((promptTemplate: PromptTemplate, source: TSource, context: Context) => TPromptTemplate)
|
||||
| undefined,
|
||||
context: Context,
|
||||
): Promise<{
|
||||
promptTemplates: Array<{ promptTemplate: TPromptTemplate; source: TSource }>;
|
||||
diagnostics: Array<PromptTemplateDiagnostic & { source: TSource }>;
|
||||
}> {
|
||||
const promptTemplates: Array<{ promptTemplate: TPromptTemplate; source: TSource }> = [];
|
||||
const diagnostics: Array<PromptTemplateDiagnostic & { source: TSource }> = [];
|
||||
for (const input of inputs) {
|
||||
const result = await loadPromptTemplates(env, input.path, context);
|
||||
for (const promptTemplate of result.promptTemplates) {
|
||||
promptTemplates.push({
|
||||
promptTemplate: mapPromptTemplate
|
||||
? mapPromptTemplate(promptTemplate, input.source, context)
|
||||
: (promptTemplate as TPromptTemplate),
|
||||
source: input.source,
|
||||
});
|
||||
}
|
||||
for (const diagnostic of result.diagnostics) diagnostics.push({ ...diagnostic, source: input.source });
|
||||
}
|
||||
return { promptTemplates, diagnostics };
|
||||
}
|
||||
|
||||
async function loadTemplatesFromDir(
|
||||
env: ExecutionEnv,
|
||||
dir: string,
|
||||
context: Context,
|
||||
): Promise<{ promptTemplates: PromptTemplate[]; diagnostics: PromptTemplateDiagnostic[] }> {
|
||||
const promptTemplates: PromptTemplate[] = [];
|
||||
const diagnostics: PromptTemplateDiagnostic[] = [];
|
||||
const entriesResult = await env.listDir(dir, context);
|
||||
if (!entriesResult.ok) {
|
||||
diagnostics.push({
|
||||
type: "warning",
|
||||
code: "list_failed",
|
||||
message: entriesResult.error.message,
|
||||
path: dir,
|
||||
});
|
||||
return { promptTemplates, diagnostics };
|
||||
}
|
||||
const entries = entriesResult.value;
|
||||
|
||||
for (const entry of entries.sort((a, b) => a.name.localeCompare(b.name))) {
|
||||
const kind = await resolveKind(env, entry, diagnostics, context);
|
||||
if (kind !== "file" || !entry.name.endsWith(".md")) continue;
|
||||
const result = await loadTemplateFromFile(env, entry.path, entry.name, context);
|
||||
if (result.promptTemplate) promptTemplates.push(result.promptTemplate);
|
||||
diagnostics.push(...result.diagnostics);
|
||||
}
|
||||
return { promptTemplates, diagnostics };
|
||||
}
|
||||
|
||||
async function loadTemplateFromFile(
|
||||
env: ExecutionEnv,
|
||||
filePath: string,
|
||||
fileName: string,
|
||||
context: Context,
|
||||
): Promise<{ promptTemplate: PromptTemplate | null; diagnostics: PromptTemplateDiagnostic[] }> {
|
||||
const diagnostics: PromptTemplateDiagnostic[] = [];
|
||||
const rawContent = await env.readTextFile(filePath, context);
|
||||
if (!rawContent.ok) {
|
||||
diagnostics.push({
|
||||
type: "warning",
|
||||
code: "read_failed",
|
||||
message: rawContent.error.message,
|
||||
path: filePath,
|
||||
});
|
||||
return { promptTemplate: null, diagnostics };
|
||||
}
|
||||
|
||||
const parsed = parseFrontmatter<PromptTemplateFrontmatter>(rawContent.value);
|
||||
if (!parsed.ok) {
|
||||
diagnostics.push({
|
||||
type: "warning",
|
||||
code: "parse_failed",
|
||||
message: parsed.error.message,
|
||||
path: filePath,
|
||||
});
|
||||
return { promptTemplate: null, diagnostics };
|
||||
}
|
||||
|
||||
const { frontmatter, body } = parsed.value;
|
||||
const firstLine = body.split("\n").find((line) => line.trim());
|
||||
let description = typeof frontmatter.description === "string" ? frontmatter.description : "";
|
||||
if (!description && firstLine) {
|
||||
description = firstLine.slice(0, 60);
|
||||
if (firstLine.length > 60) description += "...";
|
||||
}
|
||||
return {
|
||||
promptTemplate: {
|
||||
name: fileName.replace(/\.md$/i, ""),
|
||||
description,
|
||||
content: body,
|
||||
},
|
||||
diagnostics,
|
||||
};
|
||||
}
|
||||
|
||||
async function resolveKind(
|
||||
env: ExecutionEnv,
|
||||
info: FileInfo,
|
||||
diagnostics: PromptTemplateDiagnostic[],
|
||||
context: Context,
|
||||
): Promise<"file" | "directory" | undefined> {
|
||||
if (info.kind === "file" || info.kind === "directory") return info.kind;
|
||||
const canonicalPath = await env.canonicalPath(info.path, context);
|
||||
if (!canonicalPath.ok) {
|
||||
if (canonicalPath.error.code !== "not_found") {
|
||||
diagnostics.push({
|
||||
type: "warning",
|
||||
code: "file_info_failed",
|
||||
message: canonicalPath.error.message,
|
||||
path: info.path,
|
||||
});
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
const target = await env.fileInfo(canonicalPath.value, context);
|
||||
if (!target.ok) {
|
||||
if (target.error.code !== "not_found") {
|
||||
diagnostics.push({
|
||||
type: "warning",
|
||||
code: "file_info_failed",
|
||||
message: target.error.message,
|
||||
path: info.path,
|
||||
});
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
return target.value.kind === "file" || target.value.kind === "directory" ? target.value.kind : undefined;
|
||||
}
|
||||
|
||||
function parseFrontmatter<T extends Record<string, unknown>>(
|
||||
content: string,
|
||||
): Result<{ frontmatter: T; body: string }, Error> {
|
||||
try {
|
||||
const normalized = content.replace(/\r\n/g, "\n").replace(/\r/g, "\n");
|
||||
if (!normalized.startsWith("---")) return { ok: true, value: { frontmatter: {} as T, body: normalized } };
|
||||
const endIndex = normalized.indexOf("\n---", 3);
|
||||
if (endIndex === -1) return { ok: true, value: { frontmatter: {} as T, body: normalized } };
|
||||
const yamlString = normalized.slice(4, endIndex);
|
||||
const body = normalized.slice(endIndex + 4).trim();
|
||||
return { ok: true, value: { frontmatter: (parse(yamlString) ?? {}) as T, body } };
|
||||
} catch (error) {
|
||||
return { ok: false, error: toError(error) };
|
||||
}
|
||||
}
|
||||
|
||||
/** Parse an argument string using simple shell-style single and double quotes. */
|
||||
export function parseCommandArgs(argsString: string): string[] {
|
||||
const args: string[] = [];
|
||||
let current = "";
|
||||
let inQuote: string | null = null;
|
||||
|
||||
for (let i = 0; i < argsString.length; i++) {
|
||||
const char = argsString[i]!;
|
||||
if (inQuote) {
|
||||
if (char === inQuote) inQuote = null;
|
||||
else current += char;
|
||||
} else if (char === '"' || char === "'") {
|
||||
inQuote = char;
|
||||
} else if (char === " " || char === "\t") {
|
||||
if (current) {
|
||||
args.push(current);
|
||||
current = "";
|
||||
}
|
||||
} else {
|
||||
current += char;
|
||||
}
|
||||
}
|
||||
if (current) args.push(current);
|
||||
return args;
|
||||
}
|
||||
|
||||
/** Substitute prompt template placeholders (`$1`, `$@`, `$ARGUMENTS`, `${@:N}`, `${@:N:L}`) with command arguments. */
|
||||
export function substituteArgs(content: string, args: string[]): string {
|
||||
let result = content;
|
||||
result = result.replace(/\$(\d+)/g, (_, num: string) => args[parseInt(num, 10) - 1] ?? "");
|
||||
result = result.replace(/\$\{@:(\d+)(?::(\d+))?\}/g, (_, startStr: string, lengthStr?: string) => {
|
||||
let start = parseInt(startStr, 10) - 1;
|
||||
if (start < 0) start = 0;
|
||||
if (lengthStr) return args.slice(start, start + parseInt(lengthStr, 10)).join(" ");
|
||||
return args.slice(start).join(" ");
|
||||
});
|
||||
const allArgs = args.join(" ");
|
||||
result = result.replace(/\$ARGUMENTS/g, allArgs);
|
||||
result = result.replace(/\$@/g, allArgs);
|
||||
return result;
|
||||
}
|
||||
|
||||
/** Format a prompt template invocation with positional arguments. */
|
||||
export function formatPromptTemplateInvocation(template: PromptTemplate, args: string[] = []): string {
|
||||
return substituteArgs(template.content, args);
|
||||
}
|
||||
117
third_party/pi/packages/agent/src/harness/result.ts
vendored
Normal file
117
third_party/pi/packages/agent/src/harness/result.ts
vendored
Normal file
@@ -0,0 +1,117 @@
|
||||
export type Result<TValue, TError> = { ok: true; value: TValue } | { ok: false; error: TError };
|
||||
|
||||
export const Result = {
|
||||
ok<TValue>(value: TValue): Result<TValue, never> {
|
||||
return { ok: true, value };
|
||||
},
|
||||
err<TError>(error: TError): Result<never, TError> {
|
||||
return { ok: false, error };
|
||||
},
|
||||
isOk<TValue, TError>(result: Result<TValue, TError>): result is { ok: true; value: TValue } {
|
||||
return result.ok;
|
||||
},
|
||||
isErr<TValue, TError>(result: Result<TValue, TError>): result is { ok: false; error: TError } {
|
||||
return !result.ok;
|
||||
},
|
||||
};
|
||||
|
||||
export interface TaggedErrorValue<Tag extends string> extends Error {
|
||||
readonly _tag: Tag;
|
||||
toJSON(): { _tag: Tag; message: string } & Record<string, unknown>;
|
||||
}
|
||||
|
||||
export interface TaggedErrorFactory<Tag extends string> {
|
||||
new <Props extends { message: string }>(props: Props): TaggedErrorValue<Tag> & Readonly<Props>;
|
||||
is(value: unknown): value is TaggedErrorValue<Tag>;
|
||||
}
|
||||
|
||||
export function TaggedError<Tag extends string>(tag: Tag): TaggedErrorFactory<Tag> {
|
||||
class TaggedErrorClass extends Error {
|
||||
readonly _tag = tag;
|
||||
|
||||
constructor(props: { message: string } & Record<string, unknown>) {
|
||||
super(props.message);
|
||||
this.name = tag;
|
||||
Object.assign(this, props);
|
||||
}
|
||||
|
||||
toJSON(): { _tag: Tag; message: string } & Record<string, unknown> {
|
||||
const payload: Record<string, unknown> = {};
|
||||
for (const key of Object.keys(this)) {
|
||||
if (key !== "_tag") payload[key] = (this as unknown as Record<string, unknown>)[key];
|
||||
}
|
||||
return { _tag: tag, message: this.message, ...payload };
|
||||
}
|
||||
|
||||
static is(value: unknown): value is TaggedErrorValue<Tag> {
|
||||
return value instanceof TaggedErrorClass;
|
||||
}
|
||||
}
|
||||
return TaggedErrorClass as unknown as TaggedErrorFactory<Tag>;
|
||||
}
|
||||
|
||||
export class LaneBusy extends TaggedError("LaneBusy")<{
|
||||
lane: string;
|
||||
operationId: string;
|
||||
operationKind: "run" | "compaction" | "navigation";
|
||||
message: string;
|
||||
}> {}
|
||||
export class OperationMismatch extends TaggedError("OperationMismatch")<{
|
||||
lane: string;
|
||||
expectedOperationId: string;
|
||||
currentOperationId?: string;
|
||||
lastOperationId?: string;
|
||||
message: string;
|
||||
}> {}
|
||||
export class NoActiveRun extends TaggedError("NoActiveRun")<{ lane: string; message: string }> {}
|
||||
export class NoActiveOperation extends TaggedError("NoActiveOperation")<{ lane: string; message: string }> {}
|
||||
export class NothingToResume extends TaggedError("NothingToResume")<{ lane: string; message: string }> {}
|
||||
export class NothingToCompact extends TaggedError("NothingToCompact")<{ lane: string; message: string }> {}
|
||||
export class InvalidMessage extends TaggedError("InvalidMessage")<{
|
||||
lane: string;
|
||||
reason: string;
|
||||
message: string;
|
||||
}> {}
|
||||
export class InvalidNavigation extends TaggedError("InvalidNavigation")<{
|
||||
lane: string;
|
||||
reason: string;
|
||||
message: string;
|
||||
}> {}
|
||||
export class UnknownSkill extends TaggedError("UnknownSkill")<{ name: string; message: string }> {}
|
||||
export class UnknownTemplate extends TaggedError("UnknownTemplate")<{ name: string; message: string }> {}
|
||||
export class UnknownTarget extends TaggedError("UnknownTarget")<{ targetId: string; message: string }> {}
|
||||
export class InvalidLane extends TaggedError("InvalidLane")<{
|
||||
lane: string;
|
||||
reason: string;
|
||||
message: string;
|
||||
}> {}
|
||||
export class Closed extends TaggedError("Closed")<{ message: string }> {}
|
||||
|
||||
export class HarnessFault extends Error {
|
||||
readonly cause: unknown;
|
||||
|
||||
constructor(message: string, cause: unknown) {
|
||||
super(message);
|
||||
this.name = "HarnessFault";
|
||||
this.cause = cause;
|
||||
}
|
||||
}
|
||||
|
||||
export class HarnessClosed extends Error {
|
||||
constructor() {
|
||||
super("AgentHarness was closed while the operation was active");
|
||||
this.name = "HarnessClosed";
|
||||
}
|
||||
}
|
||||
|
||||
export type ErrorMatchers<TError extends TaggedErrorValue<string>, TValue> = {
|
||||
[Tag in TError["_tag"]]: (error: Extract<TError, { _tag: Tag }>) => TValue;
|
||||
};
|
||||
|
||||
export function matchError<TError extends TaggedErrorValue<string>, TValue>(
|
||||
error: TError,
|
||||
matchers: ErrorMatchers<TError, TValue>,
|
||||
): TValue {
|
||||
const matcher = (matchers as unknown as Record<string, (value: TError) => TValue>)[error._tag];
|
||||
return matcher(error);
|
||||
}
|
||||
106
third_party/pi/packages/agent/src/harness/runtime/drive.ts
vendored
Normal file
106
third_party/pi/packages/agent/src/harness/runtime/drive.ts
vendored
Normal file
@@ -0,0 +1,106 @@
|
||||
import type { DriveOutcome } from "../agent-harness.ts";
|
||||
import { AbortRequested } from "../execution/effect-gate.ts";
|
||||
import { SessionInvariantError } from "../session/session.ts";
|
||||
import { runCheckpoint, startRun } from "./drive/checkpoint.ts";
|
||||
import { runDeferred } from "./drive/deferred.ts";
|
||||
import { runGeneration } from "./drive/generation.ts";
|
||||
import { reconcileOperation } from "./drive/reconcile.ts";
|
||||
import { recoverAssistantGeneration } from "./drive/recovery.ts";
|
||||
import {
|
||||
commitNavigation,
|
||||
recoverStructuralGeneration,
|
||||
runStructuralDecision,
|
||||
runStructuralGeneration,
|
||||
runStructuralRetryWait,
|
||||
} from "./drive/structural.ts";
|
||||
import { runTools } from "./drive/tools.ts";
|
||||
import type { Lane } from "./lane.ts";
|
||||
import type { Drive, ProcedureResult } from "./types.ts";
|
||||
|
||||
function currentOperation<TContext extends object | undefined>(lane: Lane<TContext>, drive: Drive) {
|
||||
const operation = lane.state.operation;
|
||||
if (operation === null || operation.meta.operationId !== drive.operationId) {
|
||||
throw new SessionInvariantError(`Drive ${drive.operationId} has no matching current operation`);
|
||||
}
|
||||
return operation;
|
||||
}
|
||||
|
||||
/** Drive one installed pass through direct durable procedures until settlement or a durable wait. */
|
||||
export async function driveOperation<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
): Promise<DriveOutcome> {
|
||||
let operation = currentOperation(lane, drive);
|
||||
if (operation.state.control.status === "running") {
|
||||
try {
|
||||
await lane.hooks.runWithGate(
|
||||
"before_drive",
|
||||
{ lane: lane.name, runId: drive.operationId, operation: operation.meta.intent.kind },
|
||||
drive.gate,
|
||||
drive.context,
|
||||
);
|
||||
} catch (error) {
|
||||
if (!(error instanceof AbortRequested)) throw error;
|
||||
await error.cancellation;
|
||||
}
|
||||
}
|
||||
|
||||
for (;;) {
|
||||
operation = currentOperation(lane, drive);
|
||||
const state = operation.state;
|
||||
let result: ProcedureResult;
|
||||
try {
|
||||
if (state.control.status === "cancel_requested") {
|
||||
result = await reconcileOperation(lane, drive);
|
||||
} else
|
||||
switch (state.at) {
|
||||
case "starting":
|
||||
result = await startRun(lane, drive, state);
|
||||
break;
|
||||
case "checkpoint":
|
||||
result = await runCheckpoint(lane, drive, state);
|
||||
break;
|
||||
case "assistant.ready":
|
||||
case "assistant.retry_wait":
|
||||
result = await runGeneration(lane, drive, state);
|
||||
break;
|
||||
case "assistant.effect_pending":
|
||||
result = await recoverAssistantGeneration(lane, drive, state);
|
||||
break;
|
||||
case "tools":
|
||||
result = await runTools(lane, drive, state);
|
||||
break;
|
||||
case "deferred.suspended":
|
||||
case "deferred.effect_pending":
|
||||
result = await runDeferred(lane, drive, state);
|
||||
break;
|
||||
case "summary.deciding":
|
||||
result = await runStructuralDecision(lane, drive, state);
|
||||
break;
|
||||
case "summary.ready":
|
||||
result = await runStructuralGeneration(lane, drive, state);
|
||||
break;
|
||||
case "summary.effect_pending":
|
||||
result = await recoverStructuralGeneration(lane, drive, state);
|
||||
break;
|
||||
case "summary.retry_wait":
|
||||
result = await runStructuralRetryWait(lane, drive, state);
|
||||
break;
|
||||
case "navigation.ready_to_commit":
|
||||
result = await commitNavigation(lane, drive, state);
|
||||
break;
|
||||
}
|
||||
} catch (error) {
|
||||
if (!(error instanceof AbortRequested)) throw error;
|
||||
await error.cancellation;
|
||||
result = { kind: "continue" };
|
||||
}
|
||||
|
||||
if (result.kind === "settled") return { kind: "settled", outcome: result.outcome };
|
||||
if (result.kind === "waiting") return result.outcome;
|
||||
const next = currentOperation(lane, drive).state;
|
||||
if (next === state && next.control.status !== "cancel_requested") {
|
||||
throw new SessionInvariantError(`Drive procedure made no progress from ${state.at}`);
|
||||
}
|
||||
}
|
||||
}
|
||||
259
third_party/pi/packages/agent/src/harness/runtime/drive/boundary.ts
vendored
Normal file
259
third_party/pi/packages/agent/src/harness/runtime/drive/boundary.ts
vendored
Normal file
@@ -0,0 +1,259 @@
|
||||
import type { HarnessEvent, LaneQueuedItem } from "../../agent-harness.ts";
|
||||
import { insertEntry } from "../../session/commit.ts";
|
||||
import { SessionInvariantError } from "../../session/session.ts";
|
||||
import type {
|
||||
AssistantReadyOperation,
|
||||
CheckpointOperation,
|
||||
CommitResult,
|
||||
NewEntry,
|
||||
NormalizedRetryPolicy,
|
||||
OperationScope,
|
||||
SessionReader,
|
||||
SummaryDecidingOperation,
|
||||
SummaryEffectPendingOperation,
|
||||
SummaryReadyOperation,
|
||||
Write,
|
||||
} from "../../session/types.ts";
|
||||
import { branchTip, deleteValue, pendingEntry, setValue } from "../../session/values.ts";
|
||||
import type { Lane } from "../lane.ts";
|
||||
import { committedEntryEvents, entryLifecycleEvents, readBoundedContext, readLaneQueues } from "../transcript.ts";
|
||||
import type { Drive, LaneState, ProcedureResult } from "../types.ts";
|
||||
import { operationCleanupWrites, operationResultRecord } from "./terminal.ts";
|
||||
|
||||
export interface BoundaryFinishPending {
|
||||
kind: "finish_pending";
|
||||
entryIds: string[];
|
||||
}
|
||||
|
||||
export interface BoundaryPlacement {
|
||||
entries: NewEntry[];
|
||||
writes: Write[];
|
||||
tipId: string | null;
|
||||
inbox: LaneState["inbox"];
|
||||
triggerEntryId?: string;
|
||||
queues?: LaneQueuedItem[];
|
||||
}
|
||||
|
||||
type FinishBoundaryOperation =
|
||||
| CheckpointOperation
|
||||
| SummaryDecidingOperation
|
||||
| SummaryReadyOperation
|
||||
| SummaryEffectPendingOperation;
|
||||
|
||||
export function normalizedRetryPolicy<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
): NormalizedRetryPolicy {
|
||||
const retry = lane.readConfig().retryPolicy;
|
||||
return retry.enabled
|
||||
? { maxAttempts: retry.maxRetries + 1, baseDelayMs: retry.baseDelayMs }
|
||||
: { maxAttempts: 1, baseDelayMs: retry.baseDelayMs };
|
||||
}
|
||||
|
||||
export function assistantReadyAtBoundary<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
state: LaneState,
|
||||
scope: OperationScope,
|
||||
triggerEntryId: string,
|
||||
overflowRecoveryUsed: boolean,
|
||||
): AssistantReadyOperation {
|
||||
const config = lane.readConfig();
|
||||
return {
|
||||
...scope,
|
||||
at: "assistant.ready",
|
||||
generationContext: {
|
||||
stepId: lane.session.idGenerator.next(),
|
||||
triggerEntryId,
|
||||
configuration: state.configuration,
|
||||
streamOptions: config.streamOptions,
|
||||
retryPolicy: normalizedRetryPolicy(lane),
|
||||
overflowRecoveryUsed,
|
||||
},
|
||||
nextAttempt: 1,
|
||||
};
|
||||
}
|
||||
|
||||
/** Select and materialize one boundary's lane-owned input without committing it. */
|
||||
export async function planBoundaryInbox<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
state: LaneState,
|
||||
scope: OperationScope,
|
||||
reader: SessionReader,
|
||||
tipId: string | null,
|
||||
followUpWhenNoTrigger: boolean,
|
||||
): Promise<BoundaryPlacement> {
|
||||
const steer = state.inbox.filter((item) => item.kind === "steer");
|
||||
const selectedSteer = scope.settings.steeringMode === "all" ? steer : steer.slice(0, 1);
|
||||
let selected = state.inbox.filter(
|
||||
(item) => item.kind === "write" || selectedSteer.some((candidate) => candidate.entryId === item.entryId),
|
||||
);
|
||||
const load = (items: typeof selected) =>
|
||||
Promise.all(
|
||||
items.map(async (item) => {
|
||||
const stored = await reader.getValue(pendingEntry(item.entryId), drive.context);
|
||||
if (stored === undefined) {
|
||||
throw new SessionInvariantError(`Pending ${item.kind} entry ${item.entryId} is missing its payload`);
|
||||
}
|
||||
if (item.kind !== "write" && stored.value.type !== "message") {
|
||||
throw new SessionInvariantError(`Queued ${item.kind} entry ${item.entryId} is not a message`);
|
||||
}
|
||||
return { item, pending: stored.value };
|
||||
}),
|
||||
);
|
||||
let pending = await load(selected);
|
||||
const projects = (value: (typeof pending)[number]["pending"]): boolean =>
|
||||
value.type === "message" || lane.readConfig().entryProjectors[value.customType] !== undefined;
|
||||
if (followUpWhenNoTrigger && !pending.some(({ pending: value }) => projects(value))) {
|
||||
const followUp = state.inbox.filter((item) => item.kind === "followUp");
|
||||
const selectedFollowUp = scope.settings.followUpMode === "all" ? followUp : followUp.slice(0, 1);
|
||||
selected = [...selected, ...selectedFollowUp].sort((a, b) => state.inbox.indexOf(a) - state.inbox.indexOf(b));
|
||||
pending = await load(selected);
|
||||
}
|
||||
|
||||
let parentId = tipId;
|
||||
let triggerEntryId: string | undefined;
|
||||
const entries: NewEntry[] = pending.map(({ item, pending: value }) => {
|
||||
const entry: NewEntry =
|
||||
value.type === "message"
|
||||
? { id: item.entryId, parentId, type: "message", message: value.payload }
|
||||
: {
|
||||
id: item.entryId,
|
||||
parentId,
|
||||
type: "custom",
|
||||
customType: value.customType,
|
||||
...(value.payload === undefined ? {} : { data: value.payload }),
|
||||
};
|
||||
parentId = item.entryId;
|
||||
if (projects(value)) triggerEntryId = item.entryId;
|
||||
return entry;
|
||||
});
|
||||
const selectedIds = new Set(selected.map((item) => item.entryId));
|
||||
const inbox = state.inbox.filter((item) => !selectedIds.has(item.entryId));
|
||||
const queues = selected.length === 0 ? undefined : await readLaneQueues(reader, inbox, drive.context);
|
||||
return {
|
||||
entries,
|
||||
writes: [
|
||||
...entries.map((entry) => insertEntry(entry)),
|
||||
...selected.map((item) => deleteValue(pendingEntry(item.entryId))),
|
||||
...(entries.length === 0 ? [] : [setValue(branchTip(lane.name), parentId)]),
|
||||
],
|
||||
tipId: parentId,
|
||||
inbox,
|
||||
...(triggerEntryId === undefined ? {} : { triggerEntryId }),
|
||||
...(queues === undefined ? {} : { queues }),
|
||||
};
|
||||
}
|
||||
|
||||
export function boundaryPlacementEvents(
|
||||
placement: BoundaryPlacement,
|
||||
commit: CommitResult,
|
||||
firstWriteIndex: number,
|
||||
lane: string,
|
||||
runId: string,
|
||||
): HarnessEvent[] {
|
||||
return [
|
||||
...committedEntryEvents(placement.entries, commit, lane, runId, firstWriteIndex),
|
||||
...(placement.queues === undefined ? [] : [{ type: "queue_update" as const, lane, queues: placement.queues }]),
|
||||
];
|
||||
}
|
||||
|
||||
/** Replan after before_run_end and commit either renewed work or the terminal run result. */
|
||||
export async function finishRunBoundary<TContext extends object | undefined, TState extends FinishBoundaryOperation>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
capability: TState,
|
||||
continuation: Extract<CheckpointOperation["continuation"], { kind: "may_finish" }>,
|
||||
plannedEntryIds: readonly string[],
|
||||
pendingEvents: HarnessEvent[] = [],
|
||||
): Promise<ProcedureResult> {
|
||||
const context = await readBoundedContext(lane, drive, capability);
|
||||
if (context.kind === "cancel_requested") return { kind: "continue" };
|
||||
const hook = await lane.hooks.runWithGate(
|
||||
"before_run_end",
|
||||
{ lane: lane.name, runId: drive.operationId, messages: context.value },
|
||||
drive.gate,
|
||||
drive.context,
|
||||
);
|
||||
const followUp =
|
||||
hook?.followUp === undefined
|
||||
? undefined
|
||||
: {
|
||||
id: lane.session.idGenerator.next(),
|
||||
message: { role: "user" as const, content: hook.followUp, timestamp: Date.now() },
|
||||
};
|
||||
const result = await lane.continueOperation<TState, ProcedureResult>(
|
||||
capability,
|
||||
async (state, current, meta, reader) => {
|
||||
const placement = await planBoundaryInbox(lane, drive, state, current, reader, state.tipId, true);
|
||||
if (placement.triggerEntryId !== undefined) {
|
||||
return {
|
||||
kind: "commit",
|
||||
writes: placement.writes,
|
||||
operationState: assistantReadyAtBoundary(lane, state, current, placement.triggerEntryId, false),
|
||||
lane: { tipId: placement.tipId, inbox: placement.inbox },
|
||||
materialize: () => ({ kind: "continue" }) as const,
|
||||
events: (commit) => [
|
||||
...pendingEvents,
|
||||
...boundaryPlacementEvents(placement, commit, 0, lane.name, drive.operationId),
|
||||
],
|
||||
};
|
||||
}
|
||||
const hookPlanIsCurrent =
|
||||
placement.entries.length === plannedEntryIds.length &&
|
||||
placement.entries.every((entry, index) => entry.id === plannedEntryIds[index]);
|
||||
if (hookPlanIsCurrent && followUp !== undefined) {
|
||||
const entry: NewEntry = {
|
||||
id: followUp.id,
|
||||
parentId: placement.tipId,
|
||||
type: "message",
|
||||
message: followUp.message,
|
||||
};
|
||||
const entryWriteIndex = placement.writes.length;
|
||||
return {
|
||||
kind: "commit",
|
||||
writes: [...placement.writes, insertEntry(entry), setValue(branchTip(lane.name), followUp.id)],
|
||||
operationState: assistantReadyAtBoundary(lane, state, current, followUp.id, false),
|
||||
lane: { tipId: followUp.id, inbox: placement.inbox },
|
||||
materialize: () => ({ kind: "continue" }) as const,
|
||||
events: (commit) => [
|
||||
...pendingEvents,
|
||||
...boundaryPlacementEvents(placement, commit, 0, lane.name, drive.operationId),
|
||||
...entryLifecycleEvents(
|
||||
{ ...entry, seq: commit.seqs[entryWriteIndex]!, timestamp: commit.timestamp },
|
||||
lane.name,
|
||||
drive.operationId,
|
||||
),
|
||||
],
|
||||
};
|
||||
}
|
||||
if (placement.tipId === null) throw new SessionInvariantError("Completed run has no tip");
|
||||
if (continuation.includeFinalAssistant && current.latestAssistantEntryId === null) {
|
||||
throw new SessionInvariantError("Completed run is missing its final assistant");
|
||||
}
|
||||
const record = operationResultRecord(meta, "completed", placement.tipId);
|
||||
const cleanup = await operationCleanupWrites(reader, drive.operationId, current, drive.context);
|
||||
return {
|
||||
kind: "finish",
|
||||
writes: [...placement.writes, ...cleanup],
|
||||
record,
|
||||
lane: { tipId: placement.tipId, inbox: placement.inbox },
|
||||
materialize: () => ({ kind: "settled", outcome: record }) as const,
|
||||
events: (commit) => [
|
||||
...pendingEvents,
|
||||
...boundaryPlacementEvents(placement, commit, 0, lane.name, drive.operationId),
|
||||
{
|
||||
type: "run_end",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
status: "completed",
|
||||
fromTipId: meta.sourceTipId,
|
||||
tipId: placement.tipId,
|
||||
endedAt: record.endedAt,
|
||||
},
|
||||
],
|
||||
};
|
||||
},
|
||||
drive.context,
|
||||
);
|
||||
return result.kind === "cancel_requested" ? { kind: "continue" } : result.value;
|
||||
}
|
||||
190
third_party/pi/packages/agent/src/harness/runtime/drive/checkpoint.ts
vendored
Normal file
190
third_party/pi/packages/agent/src/harness/runtime/drive/checkpoint.ts
vendored
Normal file
@@ -0,0 +1,190 @@
|
||||
import { insertEntry } from "../../session/commit.ts";
|
||||
import { SessionInvariantError } from "../../session/session.ts";
|
||||
import {
|
||||
type CheckpointOperation,
|
||||
type MessageEntry,
|
||||
type NewEntry,
|
||||
operationScopeOf,
|
||||
type StartingOperation,
|
||||
type SummaryDecidingOperation,
|
||||
} from "../../session/types.ts";
|
||||
import { branchTip, operationPreparation, setValue } from "../../session/values.ts";
|
||||
import type { Lane } from "../lane.ts";
|
||||
import { chainEntries, committedEntryEvents } from "../transcript.ts";
|
||||
import type { Drive, ProcedureResult } from "../types.ts";
|
||||
import {
|
||||
assistantReadyAtBoundary,
|
||||
type BoundaryFinishPending,
|
||||
boundaryPlacementEvents,
|
||||
finishRunBoundary,
|
||||
planBoundaryInbox,
|
||||
} from "./boundary.ts";
|
||||
import { prepareCompactionThreshold } from "./structural.ts";
|
||||
|
||||
/** Consume before_run and commit the initial checkpoint. */
|
||||
export async function startRun<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
run: StartingOperation,
|
||||
): Promise<ProcedureResult> {
|
||||
const prompt = await lane.continueOperation(
|
||||
run,
|
||||
async (_state, _current, meta, reader) => {
|
||||
if (meta.intent.kind !== "run") throw new SessionInvariantError("Run operation has non-run intent");
|
||||
const entries = await reader.getEntries(meta.intent.promptEntryIds, drive.context);
|
||||
const messages = meta.intent.promptEntryIds.map((id) => {
|
||||
const entry = entries.get(id);
|
||||
if (entry?.type !== "message") {
|
||||
throw new SessionInvariantError(`Run prompt entry ${id} is missing its message`);
|
||||
}
|
||||
return entry.message;
|
||||
});
|
||||
return { kind: "return", result: messages };
|
||||
},
|
||||
drive.context,
|
||||
);
|
||||
if (prompt.kind === "cancel_requested") return { kind: "continue" };
|
||||
|
||||
const hook = await lane.hooks.runWithGate(
|
||||
"before_run",
|
||||
{ lane: lane.name, runId: drive.operationId, prompt: prompt.value, resources: lane.readConfig().resources },
|
||||
drive.gate,
|
||||
drive.context,
|
||||
);
|
||||
const injected = hook?.messages ?? [];
|
||||
for (const message of injected) {
|
||||
if (message.role === "assistant" && message.stopReason === "pending") {
|
||||
throw new SessionInvariantError("before_run returned a pending assistant message");
|
||||
}
|
||||
}
|
||||
const reserved = injected.map((message) => ({ id: lane.session.idGenerator.next(), message }));
|
||||
|
||||
const result = await lane.continueOperation(
|
||||
run,
|
||||
(state, current) => {
|
||||
const entries: NewEntry<MessageEntry>[] = chainEntries(
|
||||
state.tipId,
|
||||
reserved.map(({ id, message }) => ({ id, type: "message" as const, message })),
|
||||
);
|
||||
const triggerEntryId = entries.at(-1)?.id ?? state.tipId;
|
||||
if (triggerEntryId === null) throw new SessionInvariantError("Run start has no trigger entry");
|
||||
const nextState: CheckpointOperation = {
|
||||
...operationScopeOf(current),
|
||||
at: "checkpoint",
|
||||
continuation: { kind: "need_assistant", overflowRecoveryUsed: false },
|
||||
triggerEntryId,
|
||||
};
|
||||
return {
|
||||
kind: "commit",
|
||||
writes: [
|
||||
...entries.map((entry) => insertEntry(entry)),
|
||||
...(entries.length === 0 ? [] : [setValue(branchTip(lane.name), triggerEntryId)]),
|
||||
],
|
||||
operationState: nextState,
|
||||
lane: { tipId: triggerEntryId },
|
||||
materialize: () => ({ kind: "continue" }) as const,
|
||||
events: (commit) => committedEntryEvents(entries, commit, lane.name, drive.operationId),
|
||||
};
|
||||
},
|
||||
drive.context,
|
||||
);
|
||||
return result.kind === "cancel_requested" ? { kind: "continue" } : result.value;
|
||||
}
|
||||
|
||||
/** Advance one durable run boundary with at most one commit. */
|
||||
export async function runCheckpoint<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
run: CheckpointOperation,
|
||||
): Promise<ProcedureResult> {
|
||||
const threshold = await prepareCompactionThreshold(lane, drive, run);
|
||||
if (threshold.kind === "cancel_requested") return { kind: "continue" };
|
||||
const planned = await lane.continueOperation<CheckpointOperation, ProcedureResult | BoundaryFinishPending>(
|
||||
run,
|
||||
async (state, current, _meta, reader) => {
|
||||
const placement = await planBoundaryInbox(
|
||||
lane,
|
||||
drive,
|
||||
state,
|
||||
current,
|
||||
reader,
|
||||
state.tipId,
|
||||
threshold.value === undefined && current.continuation.kind === "may_finish",
|
||||
);
|
||||
if (placement.triggerEntryId !== undefined) {
|
||||
return {
|
||||
kind: "commit",
|
||||
writes: placement.writes,
|
||||
operationState: assistantReadyAtBoundary(lane, state, current, placement.triggerEntryId, false),
|
||||
lane: { tipId: placement.tipId, inbox: placement.inbox },
|
||||
materialize: () => ({ kind: "continue" }) as const,
|
||||
events: (commit) => boundaryPlacementEvents(placement, commit, 0, lane.name, drive.operationId),
|
||||
};
|
||||
}
|
||||
if (threshold.value !== undefined) {
|
||||
const structural: SummaryDecidingOperation = {
|
||||
...operationScopeOf(current),
|
||||
at: "summary.deciding",
|
||||
task: {
|
||||
taskId: threshold.value.taskId,
|
||||
reason: "threshold",
|
||||
boundary: {
|
||||
kind: "resume_checkpoint",
|
||||
resumeAfter: { continuation: current.continuation, triggerEntryId: current.triggerEntryId },
|
||||
},
|
||||
},
|
||||
};
|
||||
return {
|
||||
kind: "commit",
|
||||
writes: [
|
||||
...placement.writes,
|
||||
setValue(
|
||||
operationPreparation(drive.operationId, threshold.value.taskId),
|
||||
threshold.value.preparation,
|
||||
),
|
||||
],
|
||||
operationState: structural,
|
||||
lane: { tipId: placement.tipId, inbox: placement.inbox },
|
||||
materialize: () => ({ kind: "continue" }) as const,
|
||||
events: (commit) => [
|
||||
...boundaryPlacementEvents(placement, commit, 0, lane.name, drive.operationId),
|
||||
{
|
||||
type: "compaction_start",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
reason: "threshold",
|
||||
startedAt: commit.timestamp,
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
if (current.continuation.kind === "need_assistant") {
|
||||
return {
|
||||
kind: "commit",
|
||||
writes: placement.writes,
|
||||
operationState: assistantReadyAtBoundary(
|
||||
lane,
|
||||
state,
|
||||
current,
|
||||
current.triggerEntryId,
|
||||
current.continuation.overflowRecoveryUsed,
|
||||
),
|
||||
lane: { tipId: placement.tipId, inbox: placement.inbox },
|
||||
materialize: () => ({ kind: "continue" }) as const,
|
||||
events: (commit) => boundaryPlacementEvents(placement, commit, 0, lane.name, drive.operationId),
|
||||
};
|
||||
}
|
||||
return {
|
||||
kind: "return",
|
||||
result: { kind: "finish_pending", entryIds: placement.entries.map((entry) => entry.id) } as const,
|
||||
};
|
||||
},
|
||||
drive.context,
|
||||
);
|
||||
if (planned.kind === "cancel_requested") return { kind: "continue" };
|
||||
if (planned.value.kind !== "finish_pending") return planned.value;
|
||||
if (run.continuation.kind !== "may_finish") {
|
||||
throw new SessionInvariantError("Checkpoint finish mediation requires a finish continuation");
|
||||
}
|
||||
return finishRunBoundary(lane, drive, run, run.continuation, planned.value.entryIds);
|
||||
}
|
||||
287
third_party/pi/packages/agent/src/harness/runtime/drive/deferred.ts
vendored
Normal file
287
third_party/pi/packages/agent/src/harness/runtime/drive/deferred.ts
vendored
Normal file
@@ -0,0 +1,287 @@
|
||||
import type { Api, DeferredHandle, Model } from "@earendil-works/pi-ai";
|
||||
import { type Context, getTelemetryContext, withAbortSignal } from "../../context.ts";
|
||||
import { consumeAssistantStream } from "../../execution/assistant.ts";
|
||||
import { applyStreamOptionsPatch } from "../../hooks.ts";
|
||||
import { SessionInvariantError } from "../../session/session.ts";
|
||||
import {
|
||||
type DeferredEffectPendingOperation,
|
||||
type DeferredSuspendedOperation,
|
||||
type LaneConfiguration,
|
||||
type OperationError,
|
||||
type OperationScope,
|
||||
operationScopeOf,
|
||||
type SessionReader,
|
||||
type SettledAssistantMessage,
|
||||
} from "../../session/types.ts";
|
||||
import { deleteList, pendingAssistantFrames } from "../../session/values.ts";
|
||||
import type { AgentHarnessStreamOptions } from "../../types.ts";
|
||||
import type { Lane } from "../lane.ts";
|
||||
import type { ContinueOperationResult, Drive, ProcedureResult } from "../types.ts";
|
||||
import { openAssistantResponse, publishConfigurationFailure, publishResponse } from "./response.ts";
|
||||
|
||||
type DeferredLeaf = DeferredSuspendedOperation | DeferredEffectPendingOperation;
|
||||
/** Deferred effect payload without the run-wide scope fields. */
|
||||
type EffectPendingFields = Omit<DeferredEffectPendingOperation, keyof OperationScope>;
|
||||
|
||||
type PreparedDeferredPoll = {
|
||||
kind: "ready";
|
||||
source: DeferredHandle;
|
||||
model: Model<Api>;
|
||||
poll: number;
|
||||
streamOptions: AgentHarnessStreamOptions;
|
||||
};
|
||||
|
||||
type DeferredPreparation =
|
||||
| PreparedDeferredPoll
|
||||
| { kind: "cancel_requested" }
|
||||
| { kind: "waiting"; source: DeferredHandle }
|
||||
| { kind: "configuration_failure" };
|
||||
|
||||
function configurationError(identity: LaneConfiguration["model"]): OperationError {
|
||||
return {
|
||||
code: "model_unavailable",
|
||||
message: "The configured model is unavailable in this process",
|
||||
details: identity,
|
||||
};
|
||||
}
|
||||
|
||||
export async function readDeferredSourceHandle(
|
||||
reader: SessionReader,
|
||||
deferred: DeferredLeaf,
|
||||
context: Context,
|
||||
): Promise<DeferredHandle> {
|
||||
const source = (await reader.getEntries([deferred.sourceEntryId], context)).get(deferred.sourceEntryId);
|
||||
if (
|
||||
source?.type !== "message" ||
|
||||
source.message.role !== "assistant" ||
|
||||
source.message.stopReason !== "deferred" ||
|
||||
source.message.deferred === undefined
|
||||
) {
|
||||
throw new SessionInvariantError(`Deferred source ${deferred.sourceEntryId} is missing its assistant handle`);
|
||||
}
|
||||
const handle = source.message.deferred;
|
||||
const identity = deferred.configuration.model;
|
||||
if (
|
||||
handle.id.length === 0 ||
|
||||
handle.provider !== identity.provider ||
|
||||
handle.modelId !== identity.modelId ||
|
||||
handle.api !== source.message.api
|
||||
) {
|
||||
throw new SessionInvariantError(`Deferred source ${deferred.sourceEntryId} has an invalid handle`);
|
||||
}
|
||||
return handle;
|
||||
}
|
||||
|
||||
async function readSourceHandle<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
deferred: DeferredLeaf,
|
||||
): Promise<ContinueOperationResult<DeferredHandle>> {
|
||||
return lane.continueOperation(
|
||||
deferred,
|
||||
async (_state, _current, _meta, reader) => ({
|
||||
kind: "return",
|
||||
result: await readDeferredSourceHandle(reader, deferred, drive.context),
|
||||
}),
|
||||
drive.context,
|
||||
);
|
||||
}
|
||||
|
||||
async function prepareDeferredPoll<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
expected: DeferredLeaf,
|
||||
): Promise<DeferredPreparation> {
|
||||
const source = await readSourceHandle(lane, drive, expected);
|
||||
if (source.kind === "cancel_requested") return source;
|
||||
if (drive.deferredPermits === 0) return { kind: "waiting", source: source.value };
|
||||
|
||||
const identity = expected.configuration.model;
|
||||
const model = lane.models.getModel(identity.provider, identity.modelId);
|
||||
if (model === undefined) return { kind: "configuration_failure" };
|
||||
const baseOptions: AgentHarnessStreamOptions = { ...expected.streamOptions, deferred: false };
|
||||
const poll = expected.at === "deferred.suspended" ? expected.poll + 1 : expected.poll;
|
||||
const beforeRequest = await lane.hooks.runWithGate(
|
||||
"before_request",
|
||||
{
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
model,
|
||||
step: "deferred",
|
||||
attempt: poll,
|
||||
streamOptions: baseOptions,
|
||||
},
|
||||
drive.gate,
|
||||
drive.context,
|
||||
);
|
||||
return {
|
||||
kind: "ready",
|
||||
source: source.value,
|
||||
model,
|
||||
poll,
|
||||
streamOptions: {
|
||||
...(beforeRequest?.streamOptions === undefined
|
||||
? baseOptions
|
||||
: applyStreamOptionsPatch(baseOptions, beforeRequest.streamOptions)),
|
||||
deferred: false,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
async function publishPollIntent<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
deferred: DeferredLeaf,
|
||||
prepared: PreparedDeferredPoll,
|
||||
recovery: boolean,
|
||||
): Promise<ContinueOperationResult<DeferredEffectPendingOperation>> {
|
||||
const at = Date.now();
|
||||
const next: EffectPendingFields = {
|
||||
at: "deferred.effect_pending",
|
||||
stepId: deferred.stepId,
|
||||
sourceEntryId: deferred.sourceEntryId,
|
||||
poll: prepared.poll,
|
||||
responseEntryId: lane.session.idGenerator.next(at),
|
||||
usageId: lane.session.idGenerator.next(at),
|
||||
configuration: deferred.configuration,
|
||||
streamOptions: deferred.streamOptions,
|
||||
};
|
||||
return lane.continueOperation(
|
||||
deferred,
|
||||
(_state, current) => {
|
||||
const nextState: DeferredEffectPendingOperation = { ...operationScopeOf(current), ...next };
|
||||
return {
|
||||
kind: "commit",
|
||||
writes:
|
||||
deferred.at === "deferred.effect_pending"
|
||||
? [deleteList(pendingAssistantFrames(drive.operationId, deferred.responseEntryId))]
|
||||
: [],
|
||||
operationState: nextState,
|
||||
materialize: () => {
|
||||
drive.deferredPermits--;
|
||||
return nextState;
|
||||
},
|
||||
events: () => [
|
||||
{
|
||||
type: "run_resume",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
...(recovery ? { recovery: true as const } : {}),
|
||||
},
|
||||
{
|
||||
type: "turn_start",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
turnId: `${next.stepId}:poll:${next.poll}`,
|
||||
...(recovery ? { recovery: true as const } : {}),
|
||||
},
|
||||
],
|
||||
};
|
||||
},
|
||||
drive.context,
|
||||
);
|
||||
}
|
||||
|
||||
async function performDeferredPoll<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
prepared: PreparedDeferredPoll,
|
||||
intent: DeferredEffectPendingOperation,
|
||||
recovery: boolean,
|
||||
): Promise<SettledAssistantMessage> {
|
||||
const response = openAssistantResponse(lane, drive, intent.responseEntryId, recovery);
|
||||
let metadata: { status?: number; headers?: Record<string, string> } = {};
|
||||
const admitted = withAbortSignal(drive.gate.signal, drive.context);
|
||||
const stream = drive.gate.admit(() =>
|
||||
lane.models.streamDeferred(prepared.model, prepared.source, {
|
||||
wait: 0,
|
||||
signal: admitted.abortSignal,
|
||||
telemetryContext: getTelemetryContext(admitted),
|
||||
timeoutMs: prepared.streamOptions.timeoutMs,
|
||||
maxRetries: prepared.streamOptions.maxRetries,
|
||||
maxRetryDelayMs: prepared.streamOptions.maxRetryDelayMs,
|
||||
headers: prepared.streamOptions.headers,
|
||||
onPayload: async (payload, requestModel) => {
|
||||
const result = await lane.hooks.runWithGate(
|
||||
"before_payload",
|
||||
{ lane: lane.name, runId: drive.operationId, model: requestModel, payload },
|
||||
drive.gate,
|
||||
drive.context,
|
||||
);
|
||||
return result?.payload;
|
||||
},
|
||||
onResponse: (response) => {
|
||||
metadata = { status: response.status, headers: response.headers };
|
||||
},
|
||||
}),
|
||||
);
|
||||
|
||||
try {
|
||||
return await consumeAssistantStream(
|
||||
stream,
|
||||
response.observer,
|
||||
(message, context) => response.afterResponse(message, metadata, context),
|
||||
drive.context,
|
||||
);
|
||||
} finally {
|
||||
await response.close();
|
||||
}
|
||||
}
|
||||
|
||||
async function pollDeferred<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
expected: DeferredLeaf,
|
||||
recovery: boolean,
|
||||
): Promise<ProcedureResult> {
|
||||
const prepared = await prepareDeferredPoll(lane, drive, expected);
|
||||
if (prepared.kind === "cancel_requested") return { kind: "continue" };
|
||||
if (prepared.kind === "waiting") {
|
||||
return {
|
||||
kind: "waiting",
|
||||
outcome: {
|
||||
kind: "waiting",
|
||||
operationId: drive.operationId,
|
||||
reason: "deferred",
|
||||
deferred: prepared.source,
|
||||
},
|
||||
};
|
||||
}
|
||||
if (prepared.kind === "configuration_failure") {
|
||||
return publishConfigurationFailure(lane, drive, expected, configurationError(expected.configuration.model));
|
||||
}
|
||||
|
||||
const intent = await publishPollIntent(lane, drive, expected, prepared, recovery);
|
||||
if (intent.kind === "cancel_requested") return { kind: "continue" };
|
||||
const response = await performDeferredPoll(lane, drive, prepared, intent.value, recovery);
|
||||
return publishResponse(lane, drive, intent.value, response, recovery ? { recovery: true } : {});
|
||||
}
|
||||
|
||||
/** Poll one durably suspended deferred response when this pass carries a permit. */
|
||||
export function runDeferredSuspended<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
deferred: DeferredSuspendedOperation,
|
||||
): Promise<ProcedureResult> {
|
||||
return pollDeferred(lane, drive, deferred, false);
|
||||
}
|
||||
|
||||
/** Replace one orphaned unknown-outcome poll under fresh ids when this pass carries a permit. */
|
||||
export function recoverDeferredPoll<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
deferred: DeferredEffectPendingOperation,
|
||||
): Promise<ProcedureResult> {
|
||||
return pollDeferred(lane, drive, deferred, true);
|
||||
}
|
||||
|
||||
/** Advance or report the wait for one deferred run phase. */
|
||||
export function runDeferred<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
deferred: DeferredSuspendedOperation | DeferredEffectPendingOperation,
|
||||
): Promise<ProcedureResult> {
|
||||
return deferred.at === "deferred.suspended"
|
||||
? runDeferredSuspended(lane, drive, deferred)
|
||||
: recoverDeferredPoll(lane, drive, deferred);
|
||||
}
|
||||
302
third_party/pi/packages/agent/src/harness/runtime/drive/generation.ts
vendored
Normal file
302
third_party/pi/packages/agent/src/harness/runtime/drive/generation.ts
vendored
Normal file
@@ -0,0 +1,302 @@
|
||||
import type { Api, Model, Tool } from "@earendil-works/pi-ai";
|
||||
import type { AgentMessage } from "../../../types.ts";
|
||||
import { type Context, getTelemetryContext, withAbortSignal } from "../../context.ts";
|
||||
import { type HarnessAssistantStreamConfig, streamHarnessAssistant } from "../../execution/assistant.ts";
|
||||
import { applyStreamOptionsPatch } from "../../hooks.ts";
|
||||
import { SessionInvariantError } from "../../session/session.ts";
|
||||
import {
|
||||
type AssistantEffectPendingOperation,
|
||||
type AssistantReadyOperation,
|
||||
type AssistantRetryWaitOperation,
|
||||
type JsonValue,
|
||||
type OperationError,
|
||||
type OperationScope,
|
||||
operationScopeOf,
|
||||
type SettledAssistantMessage,
|
||||
} from "../../session/types.ts";
|
||||
import type { AgentHarnessStreamOptions } from "../../types.ts";
|
||||
import type { Lane } from "../lane.ts";
|
||||
import { readBoundedContext } from "../transcript.ts";
|
||||
import type { ContinueOperationResult, Drive, ProcedureResult } from "../types.ts";
|
||||
import { openAssistantResponse, publishConfigurationFailure, publishResponse } from "./response.ts";
|
||||
import { waitUntil } from "./retry.ts";
|
||||
|
||||
/** Assistant effect-pending payload without the run-wide scope fields. */
|
||||
type AssistantEffectPending = Omit<AssistantEffectPendingOperation, keyof OperationScope>;
|
||||
|
||||
type PreparedGeneration = {
|
||||
kind: "ready";
|
||||
model: Model<Api>;
|
||||
tools: Tool[];
|
||||
messages: AgentMessage[];
|
||||
systemPrompt: string;
|
||||
streamOptions: AgentHarnessStreamOptions;
|
||||
toProviderMessages: HarnessAssistantStreamConfig["toProviderMessages"];
|
||||
};
|
||||
|
||||
type GenerationPreparation =
|
||||
| PreparedGeneration
|
||||
| { kind: "configuration_failure"; error: OperationError }
|
||||
| { kind: "cancel_requested" };
|
||||
|
||||
function configurationError(
|
||||
code: "model_unavailable" | "configured_tools_unavailable",
|
||||
details: JsonValue,
|
||||
): OperationError {
|
||||
return {
|
||||
code,
|
||||
message:
|
||||
code === "model_unavailable"
|
||||
? "The configured model is unavailable in this process"
|
||||
: "One or more configured tools are unavailable in this process",
|
||||
details,
|
||||
};
|
||||
}
|
||||
|
||||
async function resolveSystemPrompt<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
context: Context,
|
||||
): Promise<string> {
|
||||
const config = lane.readConfig();
|
||||
if (config.systemPrompt === undefined) return "";
|
||||
if (typeof config.systemPrompt === "string") return config.systemPrompt;
|
||||
const source = config.toolContext;
|
||||
const toolContext = typeof source === "function" ? await source(context) : source;
|
||||
return config.systemPrompt(toolContext as TContext, context);
|
||||
}
|
||||
|
||||
async function prepareGeneration<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
generation: AssistantReadyOperation,
|
||||
): Promise<GenerationPreparation> {
|
||||
const identity = generation.generationContext.configuration.model;
|
||||
const model = lane.models.getModel(identity.provider, identity.modelId);
|
||||
if (model === undefined) {
|
||||
return { kind: "configuration_failure", error: configurationError("model_unavailable", identity) };
|
||||
}
|
||||
|
||||
const config = lane.readConfig();
|
||||
const toolsByName = new Map(config.tools.map((tool) => [tool.name, tool]));
|
||||
const missingTools = generation.generationContext.configuration.activeToolNames.filter(
|
||||
(name) => !toolsByName.has(name),
|
||||
);
|
||||
if (missingTools.length !== 0) {
|
||||
return {
|
||||
kind: "configuration_failure",
|
||||
error: configurationError("configured_tools_unavailable", { tools: missingTools }),
|
||||
};
|
||||
}
|
||||
const tools: Tool[] = generation.generationContext.configuration.activeToolNames.map((name) => {
|
||||
const tool = toolsByName.get(name);
|
||||
if (tool === undefined) throw new SessionInvariantError(`Configured tool ${name} disappeared during resolution`);
|
||||
return {
|
||||
name: tool.name,
|
||||
description: tool.description,
|
||||
parameters: tool.parameters,
|
||||
...(tool.constrainedSampling === undefined ? {} : { constrainedSampling: tool.constrainedSampling }),
|
||||
};
|
||||
});
|
||||
|
||||
const messages = await readBoundedContext(lane, drive, generation);
|
||||
if (messages.kind === "cancel_requested") return messages;
|
||||
const systemPrompt = await resolveSystemPrompt(lane, drive.context);
|
||||
const beforeRequest = await lane.hooks.runWithGate(
|
||||
"before_request",
|
||||
{
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
model,
|
||||
step: "assistant",
|
||||
attempt: generation.nextAttempt,
|
||||
streamOptions: generation.generationContext.streamOptions,
|
||||
},
|
||||
drive.gate,
|
||||
drive.context,
|
||||
);
|
||||
const streamOptions =
|
||||
beforeRequest?.streamOptions === undefined
|
||||
? generation.generationContext.streamOptions
|
||||
: applyStreamOptionsPatch(generation.generationContext.streamOptions, beforeRequest.streamOptions);
|
||||
return {
|
||||
kind: "ready",
|
||||
model,
|
||||
tools,
|
||||
messages: messages.value,
|
||||
systemPrompt,
|
||||
streamOptions,
|
||||
toProviderMessages: config.toProviderMessages,
|
||||
};
|
||||
}
|
||||
|
||||
async function publishGenerationIntent<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
ready: AssistantReadyOperation,
|
||||
prepared: PreparedGeneration,
|
||||
): Promise<ContinueOperationResult<AssistantEffectPendingOperation>> {
|
||||
const at = Date.now();
|
||||
const pending: AssistantEffectPending = {
|
||||
at: "assistant.effect_pending",
|
||||
generationContext: ready.generationContext,
|
||||
attempt: ready.nextAttempt,
|
||||
responseEntryId: lane.session.idGenerator.next(at),
|
||||
usageId: lane.session.idGenerator.next(at),
|
||||
intendedOutputLimit: prepared.model.maxTokens,
|
||||
contextWindow: prepared.model.contextWindow,
|
||||
};
|
||||
return lane.continueOperation(
|
||||
ready,
|
||||
(_state, current) => {
|
||||
const nextState: AssistantEffectPendingOperation = { ...operationScopeOf(current), ...pending };
|
||||
return {
|
||||
kind: "commit",
|
||||
writes: [],
|
||||
operationState: nextState,
|
||||
materialize: () => nextState,
|
||||
events: () =>
|
||||
ready.nextAttempt === 1
|
||||
? [
|
||||
{
|
||||
type: "turn_start",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
turnId: ready.generationContext.stepId,
|
||||
},
|
||||
]
|
||||
: [],
|
||||
};
|
||||
},
|
||||
drive.context,
|
||||
);
|
||||
}
|
||||
|
||||
async function performGeneration<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
intent: AssistantEffectPendingOperation,
|
||||
prepared: PreparedGeneration,
|
||||
): Promise<SettledAssistantMessage> {
|
||||
const response = openAssistantResponse(lane, drive, intent.responseEntryId);
|
||||
try {
|
||||
return await streamHarnessAssistant(
|
||||
prepared.messages,
|
||||
{
|
||||
model: prepared.model,
|
||||
systemPrompt: prepared.systemPrompt,
|
||||
tools: prepared.tools,
|
||||
thinkingLevel: intent.generationContext.configuration.thinkingLevel,
|
||||
streamOptions: prepared.streamOptions,
|
||||
transformContext: async (requestContext, context) => {
|
||||
const result = await lane.hooks.runWithGate(
|
||||
"transform_context",
|
||||
{ lane: lane.name, runId: drive.operationId, ...requestContext },
|
||||
drive.gate,
|
||||
context,
|
||||
);
|
||||
return {
|
||||
messages: result?.messages ?? requestContext.messages,
|
||||
systemPrompt: result?.systemPrompt ?? requestContext.systemPrompt,
|
||||
};
|
||||
},
|
||||
toProviderMessages: prepared.toProviderMessages,
|
||||
beforePayload: async (payload, requestModel, context) => {
|
||||
const result = await lane.hooks.runWithGate(
|
||||
"before_payload",
|
||||
{ lane: lane.name, runId: drive.operationId, model: requestModel, payload },
|
||||
drive.gate,
|
||||
context,
|
||||
);
|
||||
return result?.payload;
|
||||
},
|
||||
afterResponse: response.afterResponse,
|
||||
request: (aiContext, options, context) => {
|
||||
const admitted = withAbortSignal(drive.gate.signal, context);
|
||||
return drive.gate.admit(() =>
|
||||
lane.models.streamSimple(prepared.model, aiContext, {
|
||||
...options,
|
||||
sessionId: `${lane.session.metadata.id}:${lane.name}`,
|
||||
signal: admitted.abortSignal,
|
||||
telemetryContext: getTelemetryContext(admitted),
|
||||
}),
|
||||
);
|
||||
},
|
||||
observer: response.observer,
|
||||
},
|
||||
drive.context,
|
||||
);
|
||||
} finally {
|
||||
await response.close();
|
||||
}
|
||||
}
|
||||
|
||||
/** Advance one durable assistant retry wait according to this pass's local wait policy. */
|
||||
export async function runRetryWait<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
generation: AssistantRetryWaitOperation,
|
||||
): Promise<ProcedureResult> {
|
||||
if (Date.now() < generation.notBefore) {
|
||||
if (!drive.waitForRetry) {
|
||||
return {
|
||||
kind: "waiting",
|
||||
outcome: {
|
||||
kind: "waiting",
|
||||
operationId: drive.operationId,
|
||||
reason: "retry",
|
||||
notBefore: generation.notBefore,
|
||||
},
|
||||
};
|
||||
}
|
||||
await drive.gate.admit(() => waitUntil(generation.notBefore, drive.gate.signal));
|
||||
}
|
||||
|
||||
const result = await lane.continueOperation(
|
||||
generation,
|
||||
(_state, current) => {
|
||||
const nextState: AssistantReadyOperation = {
|
||||
...operationScopeOf(current),
|
||||
at: "assistant.ready",
|
||||
generationContext: generation.generationContext,
|
||||
nextAttempt: generation.nextAttempt,
|
||||
};
|
||||
return {
|
||||
kind: "commit",
|
||||
writes: [],
|
||||
operationState: nextState,
|
||||
materialize: () => ({ kind: "continue" }) as const,
|
||||
events: () => [
|
||||
{
|
||||
type: "retry_start",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
step: generation.generationContext.stepId,
|
||||
attempt: generation.nextAttempt,
|
||||
},
|
||||
],
|
||||
};
|
||||
},
|
||||
drive.context,
|
||||
);
|
||||
return result.kind === "cancel_requested" ? { kind: "continue" } : result.value;
|
||||
}
|
||||
|
||||
/** Execute one ready assistant generation or advance its durable retry wait. */
|
||||
export async function runGeneration<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
generation: AssistantReadyOperation | AssistantRetryWaitOperation,
|
||||
): Promise<ProcedureResult> {
|
||||
if (generation.at === "assistant.retry_wait") return runRetryWait(lane, drive, generation);
|
||||
|
||||
const prepared = await prepareGeneration(lane, drive, generation);
|
||||
if (prepared.kind === "configuration_failure") {
|
||||
return publishConfigurationFailure(lane, drive, generation, prepared.error);
|
||||
}
|
||||
if (prepared.kind === "cancel_requested") return { kind: "continue" };
|
||||
|
||||
const intent = await publishGenerationIntent(lane, drive, generation, prepared);
|
||||
if (intent.kind === "cancel_requested") return { kind: "continue" };
|
||||
const response = await performGeneration(lane, drive, intent.value, prepared);
|
||||
return publishResponse(lane, drive, intent.value, response);
|
||||
}
|
||||
173
third_party/pi/packages/agent/src/harness/runtime/drive/reconcile.ts
vendored
Normal file
173
third_party/pi/packages/agent/src/harness/runtime/drive/reconcile.ts
vendored
Normal file
@@ -0,0 +1,173 @@
|
||||
import type { DeferredHandle } from "@earendil-works/pi-ai";
|
||||
import type { HarnessEvent } from "../../agent-harness.ts";
|
||||
import { getTelemetryContext } from "../../context.ts";
|
||||
import { SessionInvariantError } from "../../session/session.ts";
|
||||
import type {
|
||||
DeferredEffectPendingOperation,
|
||||
DeferredSuspendedOperation,
|
||||
OperationState,
|
||||
} from "../../session/types.ts";
|
||||
import type { Lane } from "../lane.ts";
|
||||
import type { Drive, ProcedureResult } from "../types.ts";
|
||||
import { readDeferredSourceHandle } from "./deferred.ts";
|
||||
import { recoverCancelledAssistantEffect } from "./recovery.ts";
|
||||
import { operationCleanupWrites, operationResultRecord } from "./terminal.ts";
|
||||
import { runTools } from "./tools.ts";
|
||||
|
||||
async function cancelDeferredBestEffort<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
deferred: DeferredSuspendedOperation | DeferredEffectPendingOperation,
|
||||
handle: DeferredHandle,
|
||||
): Promise<void> {
|
||||
const identity = deferred.configuration.model;
|
||||
const model = lane.models.getModel(identity.provider, identity.modelId);
|
||||
if (model === undefined) return;
|
||||
try {
|
||||
await lane.models.cancelDeferred(model, handle, {
|
||||
signal: drive.closeSignal,
|
||||
telemetryContext: getTelemetryContext(drive.context),
|
||||
timeoutMs: deferred.streamOptions.timeoutMs,
|
||||
maxRetries: deferred.streamOptions.maxRetries,
|
||||
maxRetryDelayMs: deferred.streamOptions.maxRetryDelayMs,
|
||||
headers: deferred.streamOptions.headers,
|
||||
});
|
||||
} catch {
|
||||
// Remote cancellation is best-effort; durable local reconciliation must continue.
|
||||
}
|
||||
}
|
||||
|
||||
async function readDeferredHandle<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
deferred: DeferredSuspendedOperation | DeferredEffectPendingOperation,
|
||||
): Promise<DeferredHandle> {
|
||||
return lane.settleOperation(
|
||||
deferred,
|
||||
async (_state, _current, _meta, reader) => ({
|
||||
kind: "return",
|
||||
result: await readDeferredSourceHandle(reader, deferred, drive.context),
|
||||
}),
|
||||
drive.context,
|
||||
);
|
||||
}
|
||||
|
||||
async function publishAbortedTerminal<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
capability: OperationState,
|
||||
): Promise<ProcedureResult> {
|
||||
return lane.settleOperation(
|
||||
capability,
|
||||
async (state, current, meta, reader) => {
|
||||
if (current.control.status !== "cancel_requested") {
|
||||
throw new SessionInvariantError("Cancellation reconciliation requires cancelled durable control");
|
||||
}
|
||||
const record = operationResultRecord(meta, "aborted", state.tipId);
|
||||
const cleanup = await operationCleanupWrites(reader, drive.operationId, current, drive.context);
|
||||
const events: HarnessEvent[] = [];
|
||||
if (meta.intent.kind === "run") {
|
||||
switch (current.at) {
|
||||
case "summary.deciding":
|
||||
case "summary.ready":
|
||||
case "summary.effect_pending":
|
||||
case "summary.retry_wait":
|
||||
if (current.task.boundary.kind !== "resume_checkpoint" || current.task.reason === undefined) {
|
||||
throw new SessionInvariantError("Cancelled run summary has an invalid result boundary");
|
||||
}
|
||||
events.push({
|
||||
type: "compaction_end",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
reason: current.task.reason,
|
||||
status: "aborted",
|
||||
endedAt: record.endedAt,
|
||||
});
|
||||
break;
|
||||
default:
|
||||
break;
|
||||
}
|
||||
events.push({
|
||||
type: "run_end",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
status: "aborted",
|
||||
fromTipId: meta.sourceTipId,
|
||||
tipId: state.tipId,
|
||||
endedAt: record.endedAt,
|
||||
});
|
||||
} else if (meta.intent.kind === "compaction") {
|
||||
events.push({
|
||||
type: "compaction_end",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
reason: "manual",
|
||||
status: "aborted",
|
||||
endedAt: record.endedAt,
|
||||
});
|
||||
} else {
|
||||
events.push({
|
||||
type: "navigation_end",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
status: "aborted",
|
||||
fromTipId: meta.sourceTipId,
|
||||
tipId: state.tipId,
|
||||
endedAt: record.endedAt,
|
||||
});
|
||||
}
|
||||
return {
|
||||
kind: "finish",
|
||||
writes: cleanup,
|
||||
record,
|
||||
materialize: () => ({ kind: "settled", outcome: record }) as const,
|
||||
events: () => events,
|
||||
};
|
||||
},
|
||||
drive.context,
|
||||
);
|
||||
}
|
||||
|
||||
/** Advance one cancelled durable leaf without starting new ordinary work. */
|
||||
export async function reconcileOperation<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
): Promise<ProcedureResult> {
|
||||
const operation = lane.state.operation;
|
||||
if (operation === null || operation.meta.operationId !== drive.operationId) {
|
||||
throw new SessionInvariantError(`Drive ${drive.operationId} has no matching operation to reconcile`);
|
||||
}
|
||||
if (operation.state.control.status !== "cancel_requested") {
|
||||
throw new SessionInvariantError(`Operation ${drive.operationId} is not cancelled`);
|
||||
}
|
||||
drive.beginAbort(Promise.resolve());
|
||||
drive.signalAbort();
|
||||
|
||||
const state = operation.state;
|
||||
switch (state.at) {
|
||||
case "assistant.effect_pending":
|
||||
return recoverCancelledAssistantEffect(lane, drive, state);
|
||||
case "tools":
|
||||
return runTools(lane, drive, state);
|
||||
case "deferred.suspended": {
|
||||
const handle = await readDeferredHandle(lane, drive, state);
|
||||
await cancelDeferredBestEffort(lane, drive, state, handle);
|
||||
return publishAbortedTerminal(lane, drive, state);
|
||||
}
|
||||
case "deferred.effect_pending": {
|
||||
const handle = await readDeferredHandle(lane, drive, state);
|
||||
await cancelDeferredBestEffort(lane, drive, state, handle);
|
||||
return recoverCancelledAssistantEffect(lane, drive, state);
|
||||
}
|
||||
case "starting":
|
||||
case "checkpoint":
|
||||
case "assistant.ready":
|
||||
case "assistant.retry_wait":
|
||||
case "summary.deciding":
|
||||
case "summary.ready":
|
||||
case "summary.effect_pending":
|
||||
case "summary.retry_wait":
|
||||
case "navigation.ready_to_commit":
|
||||
return publishAbortedTerminal(lane, drive, state);
|
||||
}
|
||||
}
|
||||
126
third_party/pi/packages/agent/src/harness/runtime/drive/recovery.ts
vendored
Normal file
126
third_party/pi/packages/agent/src/harness/runtime/drive/recovery.ts
vendored
Normal file
@@ -0,0 +1,126 @@
|
||||
import { type AssistantMessage, reduceAssistantMessageFrames } from "@earendil-works/pi-ai";
|
||||
import type {
|
||||
AssistantEffectPendingOperation,
|
||||
DeferredEffectPendingOperation,
|
||||
LaneConfiguration,
|
||||
SettledAssistantMessage,
|
||||
} from "../../session/types.ts";
|
||||
import type { Lane } from "../lane.ts";
|
||||
import { readAssistantFrames } from "../progress.ts";
|
||||
import type { Drive, ProcedureResult } from "../types.ts";
|
||||
import { publishResponse } from "./response.ts";
|
||||
|
||||
const ZERO_USAGE = {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
|
||||
function interruptedAssistantMessage(
|
||||
identity: LaneConfiguration["model"],
|
||||
partial: AssistantMessage | undefined,
|
||||
): SettledAssistantMessage {
|
||||
const warning =
|
||||
"Assistant request was interrupted. The preceding content is the latest committed partial; newer live output may be missing and the external outcome is unknown.";
|
||||
return partial === undefined
|
||||
? {
|
||||
role: "assistant",
|
||||
content: [],
|
||||
api: "unknown",
|
||||
provider: identity.provider,
|
||||
model: identity.modelId,
|
||||
usage: ZERO_USAGE,
|
||||
stopReason: "error",
|
||||
errorMessage: warning,
|
||||
timestamp: Date.now(),
|
||||
}
|
||||
: { ...partial, usage: ZERO_USAGE, stopReason: "error", errorMessage: warning };
|
||||
}
|
||||
|
||||
/** Settle an orphaned assistant request from its bounded committed frame prefix without another provider call. */
|
||||
export async function recoverAssistantGeneration<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
generation: AssistantEffectPendingOperation,
|
||||
): Promise<ProcedureResult> {
|
||||
const frames = await lane.continueOperation(
|
||||
generation,
|
||||
async (_state, _current, _meta, reader) => ({
|
||||
kind: "return",
|
||||
result: await readAssistantFrames(reader, drive.operationId, generation.responseEntryId, drive.context),
|
||||
}),
|
||||
drive.context,
|
||||
);
|
||||
if (frames.kind === "cancel_requested") return { kind: "continue" };
|
||||
|
||||
const message = interruptedAssistantMessage(
|
||||
generation.generationContext.configuration.model,
|
||||
reduceAssistantMessageFrames(frames.value),
|
||||
);
|
||||
await lane.emitBatch(
|
||||
[
|
||||
{
|
||||
type: "message_start",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
message,
|
||||
recovery: true,
|
||||
},
|
||||
{
|
||||
type: "message_end",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
message,
|
||||
entryId: generation.responseEntryId,
|
||||
recovery: true,
|
||||
},
|
||||
],
|
||||
drive.context,
|
||||
);
|
||||
return publishResponse(lane, drive, generation, message, { recovery: true });
|
||||
}
|
||||
|
||||
/** Synthetically settle one cancelled orphaned assistant or deferred effect under its reserved ids. */
|
||||
export async function recoverCancelledAssistantEffect<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
effect: AssistantEffectPendingOperation | DeferredEffectPendingOperation,
|
||||
): Promise<ProcedureResult> {
|
||||
const frames = await lane.settleOperation(
|
||||
effect,
|
||||
async (_state, _current, _meta, reader) => ({
|
||||
kind: "return",
|
||||
result: await readAssistantFrames(reader, drive.operationId, effect.responseEntryId, drive.context),
|
||||
}),
|
||||
drive.context,
|
||||
);
|
||||
const identity =
|
||||
effect.at === "assistant.effect_pending"
|
||||
? effect.generationContext.configuration.model
|
||||
: effect.configuration.model;
|
||||
const message = interruptedAssistantMessage(identity, reduceAssistantMessageFrames(frames));
|
||||
await lane.emitBatch(
|
||||
[
|
||||
{
|
||||
type: "message_start",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
message,
|
||||
recovery: true,
|
||||
},
|
||||
{
|
||||
type: "message_end",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
message,
|
||||
entryId: effect.responseEntryId,
|
||||
recovery: true,
|
||||
},
|
||||
],
|
||||
drive.context,
|
||||
);
|
||||
return publishResponse(lane, drive, effect, message, { recovery: true });
|
||||
}
|
||||
484
third_party/pi/packages/agent/src/harness/runtime/drive/response.ts
vendored
Normal file
484
third_party/pi/packages/agent/src/harness/runtime/drive/response.ts
vendored
Normal file
@@ -0,0 +1,484 @@
|
||||
import {
|
||||
type AssistantMessageEvent,
|
||||
AssistantMessageFrameEncoder,
|
||||
isContextOverflow,
|
||||
isRecoverableLength,
|
||||
isRetryableAssistantError,
|
||||
} from "@earendil-works/pi-ai";
|
||||
import type { HarnessEvent } from "../../agent-harness.ts";
|
||||
import type { Context } from "../../context.ts";
|
||||
import type { AssistantResponseMetadata, AssistantStreamObserver } from "../../execution/assistant.ts";
|
||||
import { insertEntry, insertUsage } from "../../session/commit.ts";
|
||||
import { SessionInvariantError } from "../../session/session.ts";
|
||||
import {
|
||||
type AssistantEffectPendingOperation,
|
||||
type CommitResult,
|
||||
type DeferredEffectPendingOperation,
|
||||
type MessageEntry,
|
||||
type NewEntry,
|
||||
type OperationError,
|
||||
type OperationState,
|
||||
operationScopeOf,
|
||||
type SettledAssistantMessage,
|
||||
type SummaryDecidingOperation,
|
||||
type ToolCall,
|
||||
type UsageRow,
|
||||
} from "../../session/types.ts";
|
||||
import { branchTip, deleteList, operationPreparation, pendingAssistantFrames, setValue } from "../../session/values.ts";
|
||||
import type { Lane } from "../lane.ts";
|
||||
import { openFrameProgress } from "../progress.ts";
|
||||
import type { Drive, ProcedureResult } from "../types.ts";
|
||||
import { retryDelay, retryNotBefore } from "./retry.ts";
|
||||
import { prepareOverflowCompaction } from "./structural.ts";
|
||||
import { operationCleanupWrites, operationResultRecord } from "./terminal.ts";
|
||||
|
||||
export type AssistantResponseLifecycle = {
|
||||
observer: AssistantStreamObserver;
|
||||
afterResponse(
|
||||
message: SettledAssistantMessage,
|
||||
metadata: AssistantResponseMetadata,
|
||||
context: Context,
|
||||
): Promise<SettledAssistantMessage>;
|
||||
close(): Promise<void>;
|
||||
};
|
||||
|
||||
export function openAssistantResponse<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
responseEntryId: string,
|
||||
recovery = false,
|
||||
): AssistantResponseLifecycle {
|
||||
const progress = openFrameProgress(lane, drive, responseEntryId);
|
||||
const frameEncoder = new AssistantMessageFrameEncoder();
|
||||
const eventContext = {
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
...(recovery ? { recovery: true as const } : {}),
|
||||
};
|
||||
const close = async () => {
|
||||
progress.seal();
|
||||
await progress.drain();
|
||||
};
|
||||
return {
|
||||
observer: {
|
||||
start(message, event, context) {
|
||||
const frame = frameEncoder.encode(event);
|
||||
if (frame !== undefined) progress.write(frame);
|
||||
return lane.emitBatch([{ type: "message_start", ...eventContext, message }], context);
|
||||
},
|
||||
update(message, event: AssistantMessageEvent, context) {
|
||||
const frame = frameEncoder.encode(event);
|
||||
if (frame !== undefined) progress.write(frame);
|
||||
return lane.emitBatch(
|
||||
[{ type: "message_update", ...eventContext, message, event, ...(frame === undefined ? {} : { frame }) }],
|
||||
context,
|
||||
);
|
||||
},
|
||||
end(message, context) {
|
||||
return lane.emitBatch(
|
||||
[{ type: "message_end", ...eventContext, message, entryId: responseEntryId }],
|
||||
context,
|
||||
);
|
||||
},
|
||||
},
|
||||
async afterResponse(message, metadata, context) {
|
||||
await close();
|
||||
const result = await lane.hooks.runWithGate(
|
||||
"after_response",
|
||||
{ lane: lane.name, runId: drive.operationId, ...metadata, message },
|
||||
drive.gate,
|
||||
context,
|
||||
);
|
||||
return result?.message ?? message;
|
||||
},
|
||||
close,
|
||||
};
|
||||
}
|
||||
|
||||
type ResponseIntent = AssistantEffectPendingOperation | DeferredEffectPendingOperation;
|
||||
type ConfigurationFailureState = Extract<
|
||||
OperationState,
|
||||
{ at: "assistant.ready" | "assistant.retry_wait" | "deferred.suspended" | "deferred.effect_pending" }
|
||||
>;
|
||||
|
||||
/** Publish a non-retryable request-configuration failure before reserving response ids. */
|
||||
export async function publishConfigurationFailure<
|
||||
TContext extends object | undefined,
|
||||
TState extends ConfigurationFailureState,
|
||||
>(lane: Lane<TContext>, drive: Drive, capability: TState, error: OperationError): Promise<ProcedureResult> {
|
||||
const result = await lane.continueOperation(
|
||||
capability,
|
||||
async (state, current, meta, reader) => {
|
||||
if (state.tipId === null) throw new SessionInvariantError("Failed run has no Branch tip");
|
||||
const record = operationResultRecord(meta, "failed", state.tipId, error);
|
||||
const cleanup = await operationCleanupWrites(reader, drive.operationId, current, drive.context);
|
||||
return {
|
||||
kind: "finish",
|
||||
writes: cleanup,
|
||||
record,
|
||||
materialize: () => ({ kind: "settled", outcome: record }) as const,
|
||||
events: () => [
|
||||
{
|
||||
type: "run_end",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
status: "failed",
|
||||
error,
|
||||
fromTipId: meta.sourceTipId,
|
||||
tipId: state.tipId,
|
||||
endedAt: record.endedAt,
|
||||
},
|
||||
],
|
||||
};
|
||||
},
|
||||
drive.context,
|
||||
);
|
||||
return result.kind === "cancel_requested" ? { kind: "continue" } : result.value;
|
||||
}
|
||||
|
||||
function uuidV7Timestamp(id: string): number {
|
||||
const timestamp = Number.parseInt(id.slice(0, 8) + id.slice(9, 13), 16);
|
||||
if (!Number.isSafeInteger(timestamp)) throw new SessionInvariantError(`Invalid reserved UUIDv7 ${id}`);
|
||||
return timestamp;
|
||||
}
|
||||
|
||||
function providerError(source: "assistant" | "deferred", message: SettledAssistantMessage): OperationError {
|
||||
return {
|
||||
code: "assistant_error",
|
||||
message:
|
||||
message.errorMessage ??
|
||||
`${source === "assistant" ? "Assistant" : "Deferred"} request ended with ${message.stopReason}`,
|
||||
};
|
||||
}
|
||||
|
||||
function normalizeError(message: SettledAssistantMessage, errorMessage: string): SettledAssistantMessage {
|
||||
return { ...message, stopReason: "error", errorMessage };
|
||||
}
|
||||
|
||||
function normalizeAborted(source: "assistant" | "deferred", message: SettledAssistantMessage): SettledAssistantMessage {
|
||||
return {
|
||||
...message,
|
||||
stopReason: "aborted",
|
||||
errorMessage:
|
||||
message.errorMessage ?? `${source === "assistant" ? "Assistant" : "Deferred"} request was cancelled`,
|
||||
};
|
||||
}
|
||||
|
||||
function deferredHandleIsValid(message: SettledAssistantMessage, generation: AssistantEffectPendingOperation): boolean {
|
||||
const handle = message.deferred;
|
||||
const identity = generation.generationContext.configuration.model;
|
||||
return (
|
||||
message.stopReason === "deferred" &&
|
||||
handle !== undefined &&
|
||||
handle.id.length !== 0 &&
|
||||
handle.provider === identity.provider &&
|
||||
handle.modelId === identity.modelId &&
|
||||
handle.api === message.api
|
||||
);
|
||||
}
|
||||
|
||||
/** Classify and atomically settle one assistant-generation or deferred-poll response. */
|
||||
export async function publishResponse<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
intent: ResponseIntent,
|
||||
response: SettledAssistantMessage,
|
||||
options: { recovery?: true } = {},
|
||||
): Promise<ProcedureResult> {
|
||||
const overflow =
|
||||
intent.at === "assistant.effect_pending" &&
|
||||
(isContextOverflow(response, intent.contextWindow) || isRecoverableLength(response, intent.intendedOutputLimit));
|
||||
const overflowPreparation =
|
||||
overflow && !intent.generationContext.overflowRecoveryUsed
|
||||
? await prepareOverflowCompaction(lane, drive, intent)
|
||||
: undefined;
|
||||
return lane.settleOperation<ResponseIntent, ProcedureResult>(
|
||||
intent,
|
||||
async (state, current, meta, reader) => {
|
||||
const source = current.at === "assistant.effect_pending" ? "assistant" : "deferred";
|
||||
const responseEntryId = intent.responseEntryId;
|
||||
const configuration =
|
||||
current.at === "assistant.effect_pending" ? current.generationContext.configuration : current.configuration;
|
||||
const turnId =
|
||||
current.at === "assistant.effect_pending"
|
||||
? current.generationContext.stepId
|
||||
: `${current.stepId}:poll:${current.poll}`;
|
||||
const scope = { ...operationScopeOf(current), latestAssistantEntryId: responseEntryId };
|
||||
let committed = response;
|
||||
let settled: OperationState | undefined;
|
||||
let failure: OperationError | undefined;
|
||||
|
||||
if (current.control.status === "cancel_requested") {
|
||||
committed = normalizeAborted(source, response);
|
||||
settled = {
|
||||
...scope,
|
||||
at: "checkpoint",
|
||||
continuation: { kind: "may_finish", includeFinalAssistant: true },
|
||||
triggerEntryId: responseEntryId,
|
||||
};
|
||||
} else if (response.stopReason === "aborted") {
|
||||
throw new SessionInvariantError(
|
||||
`${source === "assistant" ? "Assistant" : "Deferred"} response is aborted while durable control is running`,
|
||||
);
|
||||
} else if (current.at === "assistant.effect_pending" && overflow) {
|
||||
committed = normalizeError(
|
||||
response,
|
||||
response.errorMessage ?? "Assistant request exceeded the context window",
|
||||
);
|
||||
if (current.generationContext.overflowRecoveryUsed || overflowPreparation === undefined) {
|
||||
failure = providerError(source, committed);
|
||||
} else {
|
||||
const structural: SummaryDecidingOperation = {
|
||||
...scope,
|
||||
at: "summary.deciding",
|
||||
task: {
|
||||
taskId: overflowPreparation.taskId,
|
||||
reason: "overflow",
|
||||
boundary: {
|
||||
kind: "resume_checkpoint",
|
||||
resumeAfter: {
|
||||
continuation: { kind: "need_assistant", overflowRecoveryUsed: true },
|
||||
triggerEntryId: current.generationContext.triggerEntryId,
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
settled = structural;
|
||||
}
|
||||
} else if (response.stopReason === "deferred") {
|
||||
if (current.at === "assistant.effect_pending") {
|
||||
if (deferredHandleIsValid(response, current)) {
|
||||
settled = {
|
||||
...scope,
|
||||
at: "deferred.suspended",
|
||||
stepId: current.generationContext.stepId,
|
||||
sourceEntryId: responseEntryId,
|
||||
poll: 0,
|
||||
configuration,
|
||||
streamOptions: current.generationContext.streamOptions,
|
||||
};
|
||||
} else {
|
||||
committed = normalizeError(response, "Provider returned an invalid deferred handle");
|
||||
failure = providerError(source, committed);
|
||||
}
|
||||
} else {
|
||||
settled = {
|
||||
...scope,
|
||||
at: "deferred.suspended",
|
||||
stepId: current.stepId,
|
||||
sourceEntryId: responseEntryId,
|
||||
poll: current.poll,
|
||||
configuration,
|
||||
streamOptions: current.streamOptions,
|
||||
};
|
||||
}
|
||||
} else if (response.stopReason === "error") {
|
||||
if (
|
||||
current.at === "assistant.effect_pending" &&
|
||||
(options.recovery === true || isRetryableAssistantError(response)) &&
|
||||
current.attempt < current.generationContext.retryPolicy.maxAttempts
|
||||
) {
|
||||
settled = {
|
||||
...scope,
|
||||
at: "assistant.retry_wait",
|
||||
generationContext: current.generationContext,
|
||||
nextAttempt: current.attempt + 1,
|
||||
notBefore: retryNotBefore(current.generationContext.retryPolicy.baseDelayMs, current.attempt),
|
||||
errorMessage: response.errorMessage ?? "Assistant request failed",
|
||||
};
|
||||
} else {
|
||||
failure = providerError(source, response);
|
||||
}
|
||||
} else {
|
||||
const calls = response.content.flatMap((content, sourceIndex) =>
|
||||
content.type === "toolCall" ? [{ sourceIndex }] : [],
|
||||
);
|
||||
if (calls.length !== 0) {
|
||||
const timestamp = uuidV7Timestamp(responseEntryId);
|
||||
const planned: ToolCall[] = calls.map(({ sourceIndex }) => ({
|
||||
status: "planned",
|
||||
sourceIndex,
|
||||
resultEntryId: lane.session.idGenerator.next(timestamp),
|
||||
}));
|
||||
settled = {
|
||||
...scope,
|
||||
at: "tools",
|
||||
batch: { assistantEntryId: responseEntryId, configuration, turnId, calls: planned },
|
||||
};
|
||||
} else if (response.stopReason === "toolUse") {
|
||||
committed = normalizeError(response, "Provider reported tool use without any tool calls");
|
||||
failure = providerError(source, committed);
|
||||
} else {
|
||||
settled = {
|
||||
...scope,
|
||||
at: "checkpoint",
|
||||
continuation: { kind: "may_finish", includeFinalAssistant: true },
|
||||
triggerEntryId: responseEntryId,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
const responseEntry: NewEntry<MessageEntry> = {
|
||||
id: responseEntryId,
|
||||
parentId: state.tipId,
|
||||
type: "message",
|
||||
message: committed,
|
||||
};
|
||||
const usageRow: Omit<UsageRow, "seq"> = {
|
||||
id: intent.usageId,
|
||||
usage: committed.usage,
|
||||
entryId: responseEntryId,
|
||||
adjustment: false,
|
||||
};
|
||||
if (settled === undefined && failure === undefined) {
|
||||
throw new SessionInvariantError("Response settlement has no durable disposition");
|
||||
}
|
||||
const record =
|
||||
failure === undefined ? undefined : operationResultRecord(meta, "failed", responseEntryId, failure);
|
||||
const cleanup =
|
||||
record === undefined ? [] : await operationCleanupWrites(reader, drive.operationId, current, drive.context);
|
||||
const writes = [
|
||||
insertEntry(responseEntry),
|
||||
insertUsage(usageRow),
|
||||
setValue(branchTip(lane.name), responseEntryId),
|
||||
...(record === undefined
|
||||
? [deleteList(pendingAssistantFrames(drive.operationId, responseEntryId))]
|
||||
: cleanup),
|
||||
...(settled?.at === "summary.deciding" && overflowPreparation !== undefined
|
||||
? [
|
||||
setValue(
|
||||
operationPreparation(drive.operationId, overflowPreparation.taskId),
|
||||
overflowPreparation.preparation,
|
||||
),
|
||||
]
|
||||
: []),
|
||||
];
|
||||
const materializeEntry = (commit: CommitResult): MessageEntry => ({
|
||||
...responseEntry,
|
||||
seq: commit.seqs[0]!,
|
||||
timestamp: commit.timestamp,
|
||||
});
|
||||
const events = (commit: CommitResult): HarnessEvent[] => {
|
||||
const entry = materializeEntry(commit);
|
||||
const batch: HarnessEvent[] = [
|
||||
{ type: "entry_added", lane: lane.name, entry, ...options },
|
||||
{
|
||||
type: "usage",
|
||||
lane: lane.name,
|
||||
row: { ...usageRow, seq: commit.seqs[1]! },
|
||||
totals: commit.stats.usage,
|
||||
},
|
||||
];
|
||||
if (current.at === "assistant.effect_pending") {
|
||||
if (options.recovery !== true && current.attempt > 1 && settled?.at !== "assistant.retry_wait") {
|
||||
const success = committed.stopReason !== "error" && committed.stopReason !== "aborted";
|
||||
batch.push({
|
||||
type: "retry_end",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
step: turnId,
|
||||
attempt: current.attempt,
|
||||
success,
|
||||
...(success
|
||||
? {}
|
||||
: {
|
||||
finalError:
|
||||
committed.errorMessage ?? `Assistant request ended with ${committed.stopReason}`,
|
||||
}),
|
||||
});
|
||||
}
|
||||
if (options.recovery !== true && settled?.at === "assistant.retry_wait") {
|
||||
batch.push({
|
||||
type: "retry_scheduled",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
step: turnId,
|
||||
attempt: settled.nextAttempt,
|
||||
maxAttempts: settled.generationContext.retryPolicy.maxAttempts,
|
||||
delayMs: retryDelay(current.generationContext.retryPolicy.baseDelayMs, current.attempt),
|
||||
notBefore: settled.notBefore,
|
||||
errorMessage: settled.errorMessage,
|
||||
});
|
||||
}
|
||||
if (options.recovery !== true && settled?.at !== "tools" && settled?.at !== "assistant.retry_wait") {
|
||||
batch.push({
|
||||
type: "turn_end",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
turnId,
|
||||
message: committed,
|
||||
toolResults: [],
|
||||
});
|
||||
}
|
||||
if (settled?.at === "summary.deciding") {
|
||||
batch.push({
|
||||
type: "compaction_start",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
reason: "overflow",
|
||||
startedAt: commit.timestamp,
|
||||
});
|
||||
}
|
||||
} else if (settled?.at !== "tools") {
|
||||
batch.push({
|
||||
type: "turn_end",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
turnId,
|
||||
message: committed,
|
||||
toolResults: [],
|
||||
...options,
|
||||
});
|
||||
}
|
||||
if (
|
||||
(current.at === "deferred.effect_pending" || options.recovery !== true) &&
|
||||
settled?.at === "deferred.suspended" &&
|
||||
committed.deferred !== undefined
|
||||
) {
|
||||
batch.push({
|
||||
type: "run_suspend",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
reason: "deferred",
|
||||
deferred: committed.deferred,
|
||||
poll: settled.poll,
|
||||
...(current.at === "deferred.effect_pending" ? options : {}),
|
||||
});
|
||||
}
|
||||
if (record !== undefined && failure !== undefined) {
|
||||
batch.push({
|
||||
type: "run_end",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
status: "failed",
|
||||
error: failure,
|
||||
fromTipId: meta.sourceTipId,
|
||||
tipId: responseEntryId,
|
||||
endedAt: record.endedAt,
|
||||
});
|
||||
}
|
||||
return batch;
|
||||
};
|
||||
if (record !== undefined) {
|
||||
return {
|
||||
kind: "finish",
|
||||
writes,
|
||||
record,
|
||||
lane: { tipId: responseEntryId },
|
||||
materialize: () => ({ kind: "settled", outcome: record }) as const,
|
||||
events,
|
||||
};
|
||||
}
|
||||
if (settled === undefined) throw new SessionInvariantError("Response settlement is missing its next state");
|
||||
return {
|
||||
kind: "commit",
|
||||
writes,
|
||||
operationState: settled,
|
||||
lane: { tipId: responseEntryId },
|
||||
materialize: () => ({ kind: "continue" }) as const,
|
||||
events,
|
||||
};
|
||||
},
|
||||
drive.context,
|
||||
);
|
||||
}
|
||||
35
third_party/pi/packages/agent/src/harness/runtime/drive/retry.ts
vendored
Normal file
35
third_party/pi/packages/agent/src/harness/runtime/drive/retry.ts
vendored
Normal file
@@ -0,0 +1,35 @@
|
||||
export function retryDelay(baseDelayMs: number, attempt: number): number {
|
||||
const delay = baseDelayMs * 2 ** Math.max(0, attempt - 1);
|
||||
return Number.isSafeInteger(delay) ? delay : Number.MAX_SAFE_INTEGER;
|
||||
}
|
||||
|
||||
export function retryNotBefore(baseDelayMs: number, attempt: number, now = Date.now()): number {
|
||||
const sum = now + retryDelay(baseDelayMs, attempt);
|
||||
return Number.isSafeInteger(sum) ? sum : Number.MAX_SAFE_INTEGER;
|
||||
}
|
||||
|
||||
export function waitUntil(notBefore: number, signal: AbortSignal): Promise<void> {
|
||||
return new Promise<void>((resolve, reject) => {
|
||||
let timer: ReturnType<typeof setTimeout> | undefined;
|
||||
const cleanup = () => {
|
||||
if (timer !== undefined) clearTimeout(timer);
|
||||
signal.removeEventListener("abort", onAbort);
|
||||
};
|
||||
const onAbort = () => {
|
||||
cleanup();
|
||||
reject(signal.reason);
|
||||
};
|
||||
const check = () => {
|
||||
const remaining = notBefore - Date.now();
|
||||
if (remaining <= 0) {
|
||||
cleanup();
|
||||
resolve();
|
||||
return;
|
||||
}
|
||||
timer = setTimeout(check, Math.min(remaining, 2_147_483_647));
|
||||
};
|
||||
signal.addEventListener("abort", onAbort, { once: true });
|
||||
if (signal.aborted) onAbort();
|
||||
else check();
|
||||
});
|
||||
}
|
||||
1221
third_party/pi/packages/agent/src/harness/runtime/drive/structural.ts
vendored
Normal file
1221
third_party/pi/packages/agent/src/harness/runtime/drive/structural.ts
vendored
Normal file
File diff suppressed because it is too large
Load Diff
82
third_party/pi/packages/agent/src/harness/runtime/drive/terminal.ts
vendored
Normal file
82
third_party/pi/packages/agent/src/harness/runtime/drive/terminal.ts
vendored
Normal file
@@ -0,0 +1,82 @@
|
||||
import type { Context } from "../../context.ts";
|
||||
import { SessionInvariantError } from "../../session/session.ts";
|
||||
import type {
|
||||
OperationError,
|
||||
OperationMeta,
|
||||
OperationResultRecord,
|
||||
OperationState,
|
||||
SessionReader,
|
||||
TerminalStatus,
|
||||
Write,
|
||||
} from "../../session/types.ts";
|
||||
import {
|
||||
deleteList,
|
||||
deleteValue,
|
||||
operationMeta,
|
||||
operationPreparationPrefix,
|
||||
operationState as operationStateValue,
|
||||
operationToolArgsPrefix,
|
||||
operationToolMemoPrefix,
|
||||
pendingAssistantFrames,
|
||||
pendingEntry,
|
||||
pendingToolOutputPrefix,
|
||||
} from "../../session/values.ts";
|
||||
|
||||
/** Build the mechanical operation-owned suffix used by an owning procedure's terminal transaction. */
|
||||
export async function operationCleanupWrites(
|
||||
reader: SessionReader,
|
||||
operationId: string,
|
||||
state: OperationState,
|
||||
context: Context,
|
||||
): Promise<Write[]> {
|
||||
const [toolArguments, toolMemos, preparations, toolOutputs] = await Promise.all([
|
||||
reader.scanValues(operationToolArgsPrefix(operationId), context),
|
||||
reader.scanValues(operationToolMemoPrefix(operationId), context),
|
||||
reader.scanValues(operationPreparationPrefix(operationId), context),
|
||||
reader.scanValues(pendingToolOutputPrefix(operationId), context),
|
||||
]);
|
||||
|
||||
const pendingIds = new Set<string>();
|
||||
if (state.at === "tools") {
|
||||
for (const call of state.batch.calls) {
|
||||
if (call.status === "outcome_ready") pendingIds.add(call.resultEntryId);
|
||||
}
|
||||
}
|
||||
let frameDelete: Write | undefined;
|
||||
if (state.at === "assistant.effect_pending" || state.at === "deferred.effect_pending") {
|
||||
frameDelete = deleteList(pendingAssistantFrames(operationId, state.responseEntryId));
|
||||
}
|
||||
|
||||
return [
|
||||
deleteValue(operationMeta(operationId)),
|
||||
deleteValue(operationStateValue(operationId)),
|
||||
...toolArguments.map(({ address }) => deleteValue(address)),
|
||||
...toolMemos.map(({ address }) => deleteValue(address)),
|
||||
...preparations.map(({ address }) => deleteValue(address)),
|
||||
...toolOutputs.map(({ address }) => deleteValue(address)),
|
||||
...(frameDelete === undefined ? [] : [frameDelete]),
|
||||
...[...pendingIds].map((id) => deleteValue(pendingEntry(id))),
|
||||
];
|
||||
}
|
||||
|
||||
/** Construct the immutable observation record for one terminal decision. */
|
||||
export function operationResultRecord(
|
||||
meta: OperationMeta,
|
||||
status: TerminalStatus,
|
||||
tipId: string | null,
|
||||
error?: OperationError,
|
||||
): OperationResultRecord {
|
||||
if ((status === "failed") !== (error !== undefined)) {
|
||||
throw new SessionInvariantError("Only a failed operation result may carry an error");
|
||||
}
|
||||
return {
|
||||
operationId: meta.operationId,
|
||||
kind: meta.intent.kind,
|
||||
status,
|
||||
...(error === undefined ? {} : { error }),
|
||||
fromTipId: meta.sourceTipId,
|
||||
tipId,
|
||||
startedAt: meta.startedAt,
|
||||
endedAt: Date.now(),
|
||||
};
|
||||
}
|
||||
326
third_party/pi/packages/agent/src/harness/runtime/drive/tool-placement.ts
vendored
Normal file
326
third_party/pi/packages/agent/src/harness/runtime/drive/tool-placement.ts
vendored
Normal file
@@ -0,0 +1,326 @@
|
||||
import type { AssistantMessage, ToolResultMessage } from "@earendil-works/pi-ai";
|
||||
import type { AgentToolCall } from "../../../types.ts";
|
||||
import type { HarnessEvent } from "../../agent-harness.ts";
|
||||
import { insertEntry, insertUsage } from "../../session/commit.ts";
|
||||
import { SessionInvariantError } from "../../session/session.ts";
|
||||
import {
|
||||
type CheckpointOperation,
|
||||
type MessageEntry,
|
||||
type NewEntry,
|
||||
type OperationState,
|
||||
operationScopeOf,
|
||||
type ToolBatch,
|
||||
type ToolCall,
|
||||
type ToolsOperation,
|
||||
type UsageRow,
|
||||
type Write,
|
||||
} from "../../session/types.ts";
|
||||
import {
|
||||
branchTip,
|
||||
deleteValue,
|
||||
laneConfig,
|
||||
operationToolArgsPrefix,
|
||||
pendingEntry,
|
||||
setValue,
|
||||
} from "../../session/values.ts";
|
||||
import type { Lane } from "../lane.ts";
|
||||
import type { Drive } from "../types.ts";
|
||||
|
||||
export type ToolBatchSource = {
|
||||
assistant: AssistantMessage;
|
||||
calls: Map<number, AgentToolCall>;
|
||||
};
|
||||
|
||||
export async function readToolBatchSource<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
batch: ToolBatch,
|
||||
): Promise<ToolBatchSource> {
|
||||
return lane.command<ToolBatchSource>(async (_state, reader) => {
|
||||
const entry = (await reader.getEntries([batch.assistantEntryId], drive.context)).get(batch.assistantEntryId);
|
||||
if (entry?.type !== "message" || entry.message.role !== "assistant") {
|
||||
throw new SessionInvariantError("Tool batch assistant entry is invalid");
|
||||
}
|
||||
const calls = new Map<number, AgentToolCall>();
|
||||
for (const call of batch.calls) {
|
||||
const block = entry.message.content[call.sourceIndex];
|
||||
if (block?.type !== "toolCall") {
|
||||
throw new SessionInvariantError(
|
||||
`Tool call source index ${call.sourceIndex} does not name a tool-call block`,
|
||||
);
|
||||
}
|
||||
calls.set(call.sourceIndex, block);
|
||||
}
|
||||
return { kind: "return", result: { assistant: entry.message, calls } };
|
||||
}, drive.context);
|
||||
}
|
||||
|
||||
type PlacementItem = {
|
||||
call: Extract<ToolCall, { status: "outcome_ready" }>;
|
||||
message: ToolResultMessage<unknown>;
|
||||
};
|
||||
|
||||
type PlacementRead = {
|
||||
items: PlacementItem[];
|
||||
turnResults?: ToolResultMessage<unknown>[];
|
||||
};
|
||||
|
||||
function isToolResultMessage(value: unknown): value is ToolResultMessage<unknown> {
|
||||
return typeof value === "object" && value !== null && "role" in value && value.role === "toolResult";
|
||||
}
|
||||
|
||||
export function toolCallFor(sources: ToolBatchSource, call: ToolCall): AgentToolCall {
|
||||
const source = sources.calls.get(call.sourceIndex);
|
||||
if (source === undefined) {
|
||||
throw new SessionInvariantError(`Tool call source index ${call.sourceIndex} is invalid`);
|
||||
}
|
||||
return source;
|
||||
}
|
||||
|
||||
export function withToolBatch(run: ToolsOperation, batch: ToolBatch): ToolsOperation {
|
||||
return { ...operationScopeOf(run), at: "tools", batch };
|
||||
}
|
||||
|
||||
async function readPlacement<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
sources: ToolBatchSource,
|
||||
): Promise<PlacementRead | undefined> {
|
||||
return lane.command<PlacementRead | undefined>(async (state, reader) => {
|
||||
const operation = state.operation;
|
||||
if (operation?.state.at !== "tools") return { kind: "return", result: undefined };
|
||||
const current = operation.state.batch;
|
||||
let first = current.calls.findIndex((call) => call.status !== "completed");
|
||||
if (first === -1) return { kind: "return", result: undefined };
|
||||
const ready: Array<Extract<ToolCall, { status: "outcome_ready" }>> = [];
|
||||
while (first < current.calls.length) {
|
||||
const call = current.calls[first]!;
|
||||
if (call.status !== "outcome_ready") break;
|
||||
ready.push(call);
|
||||
first += 1;
|
||||
}
|
||||
if (ready.length === 0) return { kind: "return", result: undefined };
|
||||
|
||||
const items: PlacementItem[] = [];
|
||||
for (const call of ready) {
|
||||
const stored = await reader.getValue(pendingEntry(call.resultEntryId), drive.context);
|
||||
if (stored?.value.type !== "message" || !isToolResultMessage(stored.value.payload)) {
|
||||
throw new SessionInvariantError(`Tool call ${call.resultEntryId} is missing its staged result`);
|
||||
}
|
||||
const source = toolCallFor(sources, call);
|
||||
if (stored.value.payload.toolCallId !== source.id || stored.value.payload.toolName !== source.name) {
|
||||
throw new SessionInvariantError(`Tool call ${call.resultEntryId} has a mismatched staged result`);
|
||||
}
|
||||
items.push({ call, message: stored.value.payload });
|
||||
}
|
||||
|
||||
let turnResults: ToolResultMessage<unknown>[] | undefined;
|
||||
if (first === current.calls.length) {
|
||||
const placedIds = current.calls
|
||||
.filter((call) => call.status === "completed")
|
||||
.map((call) => call.resultEntryId);
|
||||
const placed = await reader.getEntries(placedIds, drive.context);
|
||||
const staged = new Map(items.map((item) => [item.call.resultEntryId, item.message]));
|
||||
turnResults = current.calls.map((call) => {
|
||||
const message =
|
||||
staged.get(call.resultEntryId) ??
|
||||
(() => {
|
||||
const entry = placed.get(call.resultEntryId);
|
||||
return entry?.type === "message" && isToolResultMessage(entry.message) ? entry.message : undefined;
|
||||
})();
|
||||
if (message === undefined) {
|
||||
throw new SessionInvariantError(`Completed tool call ${call.resultEntryId} is missing its result entry`);
|
||||
}
|
||||
return message;
|
||||
});
|
||||
}
|
||||
return {
|
||||
kind: "return",
|
||||
result: { items, ...(turnResults === undefined ? {} : { turnResults }) },
|
||||
};
|
||||
}, drive.context);
|
||||
}
|
||||
|
||||
async function commitPlacement<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
capability: ToolsOperation,
|
||||
read: PlacementRead,
|
||||
): Promise<boolean> {
|
||||
const usageIds = read.items.map((item) =>
|
||||
item.message.usage === undefined ? undefined : lane.session.idGenerator.next(),
|
||||
);
|
||||
return lane.settleOperation<ToolsOperation, boolean>(
|
||||
capability,
|
||||
async (state, run, _meta, reader) => {
|
||||
const current = run.batch;
|
||||
const writes: Write[] = [];
|
||||
const eventEntries: Array<{ entry: NewEntry<MessageEntry>; seqIndex: number }> = [];
|
||||
const eventUsage: Array<{ row: Omit<UsageRow, "seq">; seqIndex: number }> = [];
|
||||
let parentId = state.tipId;
|
||||
for (const [index, item] of read.items.entries()) {
|
||||
const entry: NewEntry<MessageEntry> = {
|
||||
id: item.call.resultEntryId,
|
||||
parentId,
|
||||
type: "message",
|
||||
message: item.message,
|
||||
...(item.call.terminate ? { terminate: true } : {}),
|
||||
};
|
||||
eventEntries.push({ entry, seqIndex: writes.length });
|
||||
writes.push(insertEntry(entry), deleteValue(pendingEntry(item.call.resultEntryId)));
|
||||
const usageId = usageIds[index];
|
||||
if (usageId !== undefined && item.message.usage !== undefined) {
|
||||
const row: Omit<UsageRow, "seq"> = {
|
||||
id: usageId,
|
||||
usage: item.message.usage,
|
||||
entryId: item.call.resultEntryId,
|
||||
adjustment: false,
|
||||
};
|
||||
eventUsage.push({ row, seqIndex: writes.length });
|
||||
writes.push(insertUsage(row));
|
||||
}
|
||||
parentId = item.call.resultEntryId;
|
||||
}
|
||||
|
||||
const completedCalls = current.calls.map((call) => {
|
||||
const item = read.items.find(
|
||||
(candidate) =>
|
||||
candidate.call.sourceIndex === call.sourceIndex &&
|
||||
candidate.call.resultEntryId === call.resultEntryId,
|
||||
);
|
||||
return item === undefined
|
||||
? call
|
||||
: {
|
||||
status: "completed" as const,
|
||||
sourceIndex: call.sourceIndex,
|
||||
resultEntryId: call.resultEntryId,
|
||||
terminate: item.call.terminate,
|
||||
};
|
||||
});
|
||||
const complete = completedCalls.every((call) => call.status === "completed");
|
||||
let nextConfiguration = state.configuration;
|
||||
const addedNames: string[] = [];
|
||||
for (const item of read.items) {
|
||||
for (const name of item.message.addedToolNames ?? []) {
|
||||
if (!nextConfiguration.activeToolNames.includes(name) && !addedNames.includes(name)) {
|
||||
addedNames.push(name);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (addedNames.length !== 0) {
|
||||
nextConfiguration = {
|
||||
...nextConfiguration,
|
||||
activeToolNames: [...nextConfiguration.activeToolNames, ...addedNames],
|
||||
};
|
||||
writes.push(setValue(laneConfig(lane.name), nextConfiguration));
|
||||
}
|
||||
writes.push(setValue(branchTip(lane.name), parentId));
|
||||
|
||||
let nextRun: OperationState;
|
||||
if (complete) {
|
||||
const allTerminate = completedCalls.every((call) => call.status === "completed" && call.terminate);
|
||||
const checkpoint: CheckpointOperation = {
|
||||
...operationScopeOf(run),
|
||||
at: "checkpoint",
|
||||
continuation: allTerminate
|
||||
? { kind: "may_finish", includeFinalAssistant: false }
|
||||
: { kind: "need_assistant", overflowRecoveryUsed: false },
|
||||
triggerEntryId: parentId!,
|
||||
};
|
||||
nextRun = checkpoint;
|
||||
const args = await reader.scanValues(
|
||||
operationToolArgsPrefix(drive.operationId, capability.batch.turnId),
|
||||
drive.context,
|
||||
);
|
||||
writes.push(...args.map(({ address }) => deleteValue(address)));
|
||||
} else {
|
||||
nextRun = withToolBatch(run, { ...current, calls: completedCalls });
|
||||
}
|
||||
return {
|
||||
kind: "commit",
|
||||
writes,
|
||||
operationState: nextRun,
|
||||
lane: { tipId: parentId, configuration: nextConfiguration },
|
||||
materialize: () => complete,
|
||||
events: (commit) => {
|
||||
const events: HarnessEvent[] = [];
|
||||
for (const { entry, seqIndex } of eventEntries) {
|
||||
events.push({
|
||||
type: "entry_added",
|
||||
lane: lane.name,
|
||||
entry: { ...entry, seq: commit.seqs[seqIndex]!, timestamp: commit.timestamp },
|
||||
});
|
||||
const usage = eventUsage.find((candidate) => candidate.row.entryId === entry.id);
|
||||
if (usage !== undefined) {
|
||||
events.push({
|
||||
type: "usage",
|
||||
lane: lane.name,
|
||||
row: { ...usage.row, seq: commit.seqs[usage.seqIndex]! },
|
||||
totals: commit.stats.usage,
|
||||
});
|
||||
}
|
||||
}
|
||||
if (addedNames.length !== 0) {
|
||||
events.push({
|
||||
type: "config_update",
|
||||
lane: lane.name,
|
||||
property: "activeTools",
|
||||
previous: state.configuration.activeToolNames,
|
||||
value: nextConfiguration.activeToolNames,
|
||||
});
|
||||
}
|
||||
return events;
|
||||
},
|
||||
};
|
||||
},
|
||||
drive.context,
|
||||
);
|
||||
}
|
||||
|
||||
export async function materializeReady<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
capability: ToolsOperation,
|
||||
sources: ToolBatchSource,
|
||||
recovery: boolean,
|
||||
): Promise<void> {
|
||||
const read = await readPlacement(lane, drive, sources);
|
||||
if (read === undefined) return;
|
||||
await lane.emitBatch(
|
||||
read.items.flatMap(({ call, message }) => [
|
||||
{
|
||||
type: "message_start" as const,
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
message,
|
||||
...(recovery ? { recovery: true as const } : {}),
|
||||
},
|
||||
{
|
||||
type: "message_end" as const,
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
message,
|
||||
entryId: call.resultEntryId,
|
||||
...(recovery ? { recovery: true as const } : {}),
|
||||
},
|
||||
]),
|
||||
drive.context,
|
||||
);
|
||||
const complete = await commitPlacement(lane, drive, capability, read);
|
||||
if (complete && read.turnResults !== undefined) {
|
||||
await lane.emitBatch(
|
||||
[
|
||||
{
|
||||
type: "turn_end",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
turnId: capability.batch.turnId,
|
||||
message: sources.assistant,
|
||||
toolResults: read.turnResults,
|
||||
...(recovery ? { recovery: true as const } : {}),
|
||||
},
|
||||
],
|
||||
drive.context,
|
||||
);
|
||||
}
|
||||
}
|
||||
692
third_party/pi/packages/agent/src/harness/runtime/drive/tools.ts
vendored
Normal file
692
third_party/pi/packages/agent/src/harness/runtime/drive/tools.ts
vendored
Normal file
@@ -0,0 +1,692 @@
|
||||
import type { ToolResultMessage } from "@earendil-works/pi-ai";
|
||||
import type { AgentToolCall, AgentToolResult } from "../../../types.ts";
|
||||
import { AbortRequested } from "../../execution/effect-gate.ts";
|
||||
import {
|
||||
applyBeforeToolDecision,
|
||||
type ClearedToolCall,
|
||||
createToolResultMessage,
|
||||
type ExecutedToolCall,
|
||||
executeToolCall,
|
||||
type FinalizedToolCall,
|
||||
finalizeToolCall,
|
||||
prepareToolCall,
|
||||
toolResultFromMessage,
|
||||
} from "../../execution/tools.ts";
|
||||
import { SessionInvariantError } from "../../session/session.ts";
|
||||
import type { JsonValue, ToolBatch, ToolCall, ToolsOperation } from "../../session/types.ts";
|
||||
import {
|
||||
deleteValue,
|
||||
operationToolArgs,
|
||||
operationToolMemo,
|
||||
operationToolMemoPrefix,
|
||||
pendingEntry,
|
||||
pendingToolOutput,
|
||||
setValue,
|
||||
} from "../../session/values.ts";
|
||||
import type { AgentHarnessTool, AgentHarnessToolInvocation } from "../../types.ts";
|
||||
import type { Lane } from "../lane.ts";
|
||||
import { openToolProgress } from "../progress.ts";
|
||||
import type { ContinueOperationResult, Drive, LaneState, ProcedureResult } from "../types.ts";
|
||||
import {
|
||||
materializeReady,
|
||||
readToolBatchSource,
|
||||
type ToolBatchSource,
|
||||
toolCallFor,
|
||||
withToolBatch,
|
||||
} from "./tool-placement.ts";
|
||||
|
||||
type ToolCallTask = { completion: Promise<void> };
|
||||
type ToolOutcome = { toolCall: AgentToolCall; message: ToolResultMessage<unknown>; terminate: boolean };
|
||||
type PreparedToolInvocation<TContext extends object | undefined> =
|
||||
| { kind: "ready"; cleared: ClearedToolCall<TContext> }
|
||||
| { kind: "outcome"; outcome: ToolOutcome };
|
||||
|
||||
const INTERRUPTION_MARKER =
|
||||
"[Tool execution was interrupted. The preceding output is the latest durable progress snapshot; newer live output may be missing, and the external outcome is unknown.]";
|
||||
|
||||
class ToolInvocationEnded extends Error {
|
||||
constructor() {
|
||||
super("Tool invocation no longer owns its durable effect");
|
||||
this.name = "ToolInvocationEnded";
|
||||
}
|
||||
}
|
||||
|
||||
function currentBatch<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
): { run: ToolsOperation; batch: ToolBatch } | undefined {
|
||||
const operation = lane.state.operation;
|
||||
if (operation?.state.at !== "tools") return undefined;
|
||||
return { run: operation.state, batch: operation.state.batch };
|
||||
}
|
||||
|
||||
function findCall(batch: ToolBatch, sourceIndex: number, resultEntryId: string): ToolCall | undefined {
|
||||
return batch.calls.find((call) => call.sourceIndex === sourceIndex && call.resultEntryId === resultEntryId);
|
||||
}
|
||||
|
||||
function replaceCall(batch: ToolBatch, replacement: ToolCall): ToolBatch {
|
||||
return {
|
||||
...batch,
|
||||
calls: batch.calls.map((call) =>
|
||||
call.sourceIndex === replacement.sourceIndex && call.resultEntryId === replacement.resultEntryId
|
||||
? replacement
|
||||
: call,
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
function validateMemoName(name: string): void {
|
||||
if (name.length === 0) throw new TypeError("Tool invocation memo name must not be empty");
|
||||
if (name.includes(":")) throw new TypeError("Tool invocation memo name must not contain ':'");
|
||||
}
|
||||
|
||||
function invocationCapability<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
batch: ToolBatch,
|
||||
call: Extract<ToolCall, { status: "effect_pending" }>,
|
||||
): { invocation: AgentHarnessToolInvocation; expire(): void } {
|
||||
let active = true;
|
||||
const ownsEffect = (state: LaneState): boolean => {
|
||||
const operation = state.operation;
|
||||
if (operation?.state.at !== "tools") return false;
|
||||
return findCall(operation.state.batch, call.sourceIndex, call.resultEntryId)?.status === "effect_pending";
|
||||
};
|
||||
const ended = (): ToolInvocationEnded => new ToolInvocationEnded();
|
||||
return {
|
||||
invocation: {
|
||||
invocationId: call.resultEntryId,
|
||||
operationId: drive.operationId,
|
||||
turnId: batch.turnId,
|
||||
getMemo(name) {
|
||||
validateMemoName(name);
|
||||
if (!active) return Promise.reject(ended());
|
||||
return lane.command<JsonValue | undefined>(async (state, reader) => {
|
||||
if (!ownsEffect(state)) return { kind: "reject", error: ended() };
|
||||
const stored = await reader.getValue(
|
||||
operationToolMemo(drive.operationId, call.resultEntryId, name),
|
||||
drive.context,
|
||||
);
|
||||
return { kind: "return", result: stored?.value };
|
||||
}, drive.context);
|
||||
},
|
||||
setMemo(name, value) {
|
||||
validateMemoName(name);
|
||||
if (!active) return Promise.reject(ended());
|
||||
return lane.command<void>((state) => {
|
||||
if (!ownsEffect(state)) return { kind: "reject", error: ended() };
|
||||
const address = operationToolMemo(drive.operationId, call.resultEntryId, name);
|
||||
return {
|
||||
kind: "commit",
|
||||
writes: [value === undefined ? deleteValue(address) : setValue(address, value)],
|
||||
next: state,
|
||||
materialize: () => undefined,
|
||||
};
|
||||
}, drive.context);
|
||||
},
|
||||
},
|
||||
expire() {
|
||||
active = false;
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function syntheticMessage(
|
||||
toolCall: AgentToolCall,
|
||||
content: AgentToolResult<unknown>["content"],
|
||||
options: { details?: unknown; usage?: AgentToolResult<unknown>["usage"] } = {},
|
||||
): ToolResultMessage<unknown> {
|
||||
return {
|
||||
role: "toolResult",
|
||||
toolCallId: toolCall.id,
|
||||
toolName: toolCall.name,
|
||||
content,
|
||||
...(options.details === undefined ? {} : { details: options.details }),
|
||||
...(options.usage === undefined ? {} : { usage: options.usage }),
|
||||
isError: true,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
function abortedOutcome(toolCall: AgentToolCall): ToolOutcome {
|
||||
return {
|
||||
toolCall,
|
||||
message: syntheticMessage(toolCall, [{ type: "text", text: "Tool execution was cancelled before completion." }]),
|
||||
terminate: false,
|
||||
};
|
||||
}
|
||||
|
||||
function interruptedOutcome(toolCall: AgentToolCall, checkpoint: AgentToolResult<unknown> | undefined): ToolOutcome {
|
||||
return {
|
||||
toolCall,
|
||||
message: syntheticMessage(
|
||||
toolCall,
|
||||
[...(checkpoint?.content ?? []), { type: "text", text: INTERRUPTION_MARKER }],
|
||||
checkpoint === undefined ? {} : { details: checkpoint.details, usage: checkpoint.usage },
|
||||
),
|
||||
terminate: false,
|
||||
};
|
||||
}
|
||||
|
||||
function truncatedOutcome(toolCall: AgentToolCall): ToolOutcome {
|
||||
return {
|
||||
toolCall,
|
||||
message: syntheticMessage(toolCall, [
|
||||
{
|
||||
type: "text",
|
||||
text: `Tool call ${JSON.stringify(toolCall.name)} was not executed because the assistant response hit the output token limit, so its arguments may be truncated. Re-issue the tool call with complete arguments.`,
|
||||
},
|
||||
]),
|
||||
terminate: false,
|
||||
};
|
||||
}
|
||||
|
||||
function outcomeFromFinalizedCall(finalized: FinalizedToolCall): ToolOutcome {
|
||||
return { toolCall: finalized.toolCall, message: createToolResultMessage(finalized), terminate: finalized.terminate };
|
||||
}
|
||||
|
||||
async function publishToolIntent<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
run: ToolsOperation,
|
||||
planned: Extract<ToolCall, { status: "planned" }>,
|
||||
toolCall: AgentToolCall,
|
||||
args: Record<string, JsonValue>,
|
||||
replay: "never" | "safe",
|
||||
recovery: boolean,
|
||||
): Promise<ContinueOperationResult<Extract<ToolCall, { status: "effect_pending" }>>> {
|
||||
return lane.continueOperation(
|
||||
run,
|
||||
(_state, run) => {
|
||||
const effectPending: Extract<ToolCall, { status: "effect_pending" }> = {
|
||||
status: "effect_pending",
|
||||
sourceIndex: planned.sourceIndex,
|
||||
resultEntryId: planned.resultEntryId,
|
||||
replay,
|
||||
};
|
||||
return {
|
||||
kind: "commit",
|
||||
writes: [setValue(operationToolArgs(drive.operationId, run.batch.turnId, planned.sourceIndex), args)],
|
||||
operationState: withToolBatch(run, replaceCall(run.batch, effectPending)),
|
||||
materialize: () => effectPending,
|
||||
events: () => [
|
||||
{
|
||||
type: "tool_start",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
turnId: run.batch.turnId,
|
||||
toolCallId: toolCall.id,
|
||||
toolName: toolCall.name,
|
||||
args,
|
||||
...(recovery ? { recovery: true as const } : {}),
|
||||
},
|
||||
],
|
||||
};
|
||||
},
|
||||
drive.context,
|
||||
);
|
||||
}
|
||||
|
||||
async function publishToolOutcome<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
capability: ToolsOperation,
|
||||
call: Extract<ToolCall, { status: "planned" | "effect_pending" }>,
|
||||
finalized: ToolOutcome,
|
||||
recovery: boolean,
|
||||
): Promise<void> {
|
||||
await lane.settleOperation(
|
||||
capability,
|
||||
async (_state, run, _meta, reader) => {
|
||||
const { toolCall } = finalized;
|
||||
const memos = await reader.scanValues(
|
||||
operationToolMemoPrefix(drive.operationId, call.resultEntryId),
|
||||
drive.context,
|
||||
);
|
||||
const durableTerminate = run.control.status === "running" && finalized.terminate;
|
||||
const outcome: Extract<ToolCall, { status: "outcome_ready" }> = {
|
||||
status: "outcome_ready",
|
||||
sourceIndex: call.sourceIndex,
|
||||
resultEntryId: call.resultEntryId,
|
||||
terminate: durableTerminate,
|
||||
};
|
||||
return {
|
||||
kind: "commit",
|
||||
writes: [
|
||||
setValue(pendingEntry(call.resultEntryId), { type: "message", payload: finalized.message }),
|
||||
deleteValue(pendingToolOutput(drive.operationId, call.resultEntryId)),
|
||||
...memos.map(({ address }) => deleteValue(address)),
|
||||
],
|
||||
operationState: withToolBatch(run, replaceCall(run.batch, outcome)),
|
||||
materialize: () => undefined,
|
||||
events: () => [
|
||||
...(call.status === "planned"
|
||||
? [
|
||||
{
|
||||
type: "tool_start" as const,
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
turnId: run.batch.turnId,
|
||||
toolCallId: toolCall.id,
|
||||
toolName: toolCall.name,
|
||||
args: toolCall.arguments,
|
||||
...(recovery ? { recovery: true as const } : {}),
|
||||
},
|
||||
]
|
||||
: []),
|
||||
{
|
||||
type: "tool_end",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
turnId: run.batch.turnId,
|
||||
toolCallId: toolCall.id,
|
||||
toolName: toolCall.name,
|
||||
result: toolResultFromMessage(finalized.message, durableTerminate),
|
||||
isError: finalized.message.isError,
|
||||
terminate: durableTerminate,
|
||||
...(recovery ? { recovery: true as const } : {}),
|
||||
},
|
||||
],
|
||||
};
|
||||
},
|
||||
drive.context,
|
||||
);
|
||||
}
|
||||
|
||||
async function clearReplayCheckpoint<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
batch: ToolBatch,
|
||||
call: Extract<ToolCall, { status: "effect_pending" }>,
|
||||
toolCall: AgentToolCall,
|
||||
): Promise<Record<string, JsonValue>> {
|
||||
return lane.command(async (state, reader) => {
|
||||
const stored = await reader.getValue(
|
||||
operationToolArgs(drive.operationId, batch.turnId, call.sourceIndex),
|
||||
drive.context,
|
||||
);
|
||||
if (stored === undefined) {
|
||||
throw new SessionInvariantError(`Tool call ${call.resultEntryId} is missing persisted arguments`);
|
||||
}
|
||||
return {
|
||||
kind: "commit",
|
||||
writes: [deleteValue(pendingToolOutput(drive.operationId, call.resultEntryId))],
|
||||
next: state,
|
||||
materialize: () => stored.value,
|
||||
events: () => [
|
||||
{
|
||||
type: "tool_start",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
turnId: batch.turnId,
|
||||
toolCallId: toolCall.id,
|
||||
toolName: toolCall.name,
|
||||
args: stored.value,
|
||||
recovery: true,
|
||||
},
|
||||
],
|
||||
};
|
||||
}, drive.context);
|
||||
}
|
||||
|
||||
function readCheckpoint<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
call: Extract<ToolCall, { status: "effect_pending" }>,
|
||||
): Promise<AgentToolResult<unknown> | undefined> {
|
||||
return lane.command(async (_state, reader) => {
|
||||
const stored = await reader.getValue(pendingToolOutput(drive.operationId, call.resultEntryId), drive.context);
|
||||
return { kind: "return", result: stored?.value };
|
||||
}, drive.context);
|
||||
}
|
||||
|
||||
async function resolveToolContext<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
): Promise<TContext> {
|
||||
const source = lane.readConfig().toolContext;
|
||||
return (typeof source === "function" ? await source(drive.context) : source) as TContext;
|
||||
}
|
||||
|
||||
async function performToolInvocation<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
batch: ToolBatch,
|
||||
call: Extract<ToolCall, { status: "effect_pending" }>,
|
||||
cleared: ClearedToolCall<TContext>,
|
||||
toolContext: TContext,
|
||||
recovery: boolean,
|
||||
): Promise<ToolOutcome> {
|
||||
const capability = invocationCapability(lane, drive, batch, call);
|
||||
const progress = openToolProgress(lane, drive, batch.turnId, call.sourceIndex, call.resultEntryId);
|
||||
let latestUpdateDelivery: Promise<void> = Promise.resolve();
|
||||
const publishUpdate = (partial: AgentToolResult<unknown>): void => {
|
||||
latestUpdateDelivery = lane.emitBatch(
|
||||
[
|
||||
{
|
||||
type: "tool_update",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
turnId: batch.turnId,
|
||||
toolCallId: cleared.toolCall.id,
|
||||
toolName: cleared.toolCall.name,
|
||||
partialResult: partial,
|
||||
...(recovery ? { recovery: true as const } : {}),
|
||||
},
|
||||
],
|
||||
drive.context,
|
||||
);
|
||||
void latestUpdateDelivery.catch(() => {});
|
||||
};
|
||||
|
||||
let execution: Promise<ExecutedToolCall>;
|
||||
try {
|
||||
execution = executeToolCall(
|
||||
cleared,
|
||||
drive.gate,
|
||||
(partial, options) => {
|
||||
publishUpdate(partial);
|
||||
if (options?.checkpoint === true) progress.write(partial);
|
||||
},
|
||||
toolContext,
|
||||
capability.invocation,
|
||||
drive.context,
|
||||
);
|
||||
} catch (error) {
|
||||
capability.expire();
|
||||
progress.seal();
|
||||
await progress.drain();
|
||||
if (!(error instanceof AbortRequested)) throw error;
|
||||
await error.cancellation;
|
||||
return recovery ? interruptedOutcome(cleared.toolCall, undefined) : abortedOutcome(cleared.toolCall);
|
||||
}
|
||||
|
||||
const executed = await execution.finally(() => {
|
||||
capability.expire();
|
||||
progress.seal();
|
||||
});
|
||||
await latestUpdateDelivery;
|
||||
await progress.drain();
|
||||
|
||||
let patch: Awaited<ReturnType<typeof lane.hooks.runToolWithGate<"after_tool">>>;
|
||||
try {
|
||||
patch = await lane.hooks.runToolWithGate(
|
||||
"after_tool",
|
||||
{
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
toolCallId: cleared.toolCall.id,
|
||||
toolName: cleared.toolCall.name,
|
||||
args: cleared.args,
|
||||
content: executed.result.content,
|
||||
...(executed.result.details === undefined ? {} : { details: executed.result.details as JsonValue }),
|
||||
isError: executed.isError,
|
||||
...(executed.result.usage === undefined ? {} : { usage: executed.result.usage }),
|
||||
},
|
||||
drive.gate,
|
||||
drive.context,
|
||||
);
|
||||
} catch (error) {
|
||||
if (!(error instanceof AbortRequested)) throw error;
|
||||
patch = undefined;
|
||||
}
|
||||
const finalized = finalizeToolCall(cleared, executed, patch);
|
||||
return { toolCall: finalized.toolCall, message: createToolResultMessage(finalized), terminate: finalized.terminate };
|
||||
}
|
||||
|
||||
async function prepareToolInvocation<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
sources: ToolBatchSource,
|
||||
call: Extract<ToolCall, { status: "planned" }>,
|
||||
tools: AgentHarnessTool<TContext>[],
|
||||
): Promise<PreparedToolInvocation<TContext>> {
|
||||
const toolCall = toolCallFor(sources, call);
|
||||
if (sources.assistant.stopReason === "length") {
|
||||
return { kind: "outcome", outcome: truncatedOutcome(toolCall) };
|
||||
}
|
||||
const prepared = prepareToolCall(toolCall, tools);
|
||||
if ("kind" in prepared) return { kind: "outcome", outcome: outcomeFromFinalizedCall(prepared) };
|
||||
|
||||
let decision: Awaited<ReturnType<typeof lane.hooks.runToolWithGate<"before_tool">>>;
|
||||
try {
|
||||
decision = await lane.hooks.runToolWithGate(
|
||||
"before_tool",
|
||||
{
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
toolCallId: toolCall.id,
|
||||
toolName: toolCall.name,
|
||||
args: prepared.args,
|
||||
},
|
||||
drive.gate,
|
||||
drive.context,
|
||||
);
|
||||
} catch (error) {
|
||||
if (!(error instanceof AbortRequested)) throw error;
|
||||
await error.cancellation;
|
||||
return { kind: "outcome", outcome: abortedOutcome(toolCall) };
|
||||
}
|
||||
const cleared = applyBeforeToolDecision(prepared, decision);
|
||||
return "kind" in cleared
|
||||
? { kind: "outcome", outcome: outcomeFromFinalizedCall(cleared) }
|
||||
: { kind: "ready", cleared };
|
||||
}
|
||||
|
||||
async function startToolInvocation<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
run: ToolsOperation,
|
||||
sources: ToolBatchSource,
|
||||
call: Extract<ToolCall, { status: "planned" }>,
|
||||
tools: AgentHarnessTool<TContext>[],
|
||||
toolContext: TContext,
|
||||
recovery: boolean,
|
||||
): Promise<ToolCallTask> {
|
||||
const prepared = await prepareToolInvocation(lane, drive, sources, call, tools);
|
||||
if (prepared.kind === "outcome") {
|
||||
return { completion: publishToolOutcome(lane, drive, run, call, prepared.outcome, recovery) };
|
||||
}
|
||||
const effectPending = await publishToolIntent(
|
||||
lane,
|
||||
drive,
|
||||
run,
|
||||
call,
|
||||
prepared.cleared.toolCall,
|
||||
prepared.cleared.args,
|
||||
prepared.cleared.tool.replay ?? "never",
|
||||
recovery,
|
||||
);
|
||||
return {
|
||||
completion:
|
||||
effectPending.kind === "cancel_requested"
|
||||
? publishToolOutcome(lane, drive, run, call, abortedOutcome(prepared.cleared.toolCall), recovery)
|
||||
: performToolInvocation(
|
||||
lane,
|
||||
drive,
|
||||
run.batch,
|
||||
effectPending.value,
|
||||
prepared.cleared,
|
||||
toolContext,
|
||||
recovery,
|
||||
).then((outcome) => publishToolOutcome(lane, drive, run, effectPending.value, outcome, recovery)),
|
||||
};
|
||||
}
|
||||
|
||||
async function recoverToolInvocation<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
run: ToolsOperation,
|
||||
sources: ToolBatchSource,
|
||||
call: Extract<ToolCall, { status: "effect_pending" }>,
|
||||
toolsByName: Map<string, AgentHarnessTool<TContext>>,
|
||||
toolContext: TContext,
|
||||
cancelled: boolean,
|
||||
): Promise<ToolCallTask> {
|
||||
const toolCall = toolCallFor(sources, call);
|
||||
const tool = toolsByName.get(toolCall.name);
|
||||
if (!cancelled && call.replay === "safe" && tool?.replay === "safe") {
|
||||
const args = await clearReplayCheckpoint(lane, drive, run.batch, call, toolCall);
|
||||
const cleared: ClearedToolCall<TContext> = { toolCall, tool, args };
|
||||
return {
|
||||
completion: performToolInvocation(lane, drive, run.batch, call, cleared, toolContext, true).then((outcome) =>
|
||||
publishToolOutcome(lane, drive, run, call, outcome, true),
|
||||
),
|
||||
};
|
||||
}
|
||||
const checkpoint = await readCheckpoint(lane, drive, call);
|
||||
return {
|
||||
completion: publishToolOutcome(lane, drive, run, call, interruptedOutcome(toolCall, checkpoint), true),
|
||||
};
|
||||
}
|
||||
|
||||
async function runSequential<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
run: ToolsOperation,
|
||||
sources: ToolBatchSource,
|
||||
execution:
|
||||
| {
|
||||
tools: AgentHarnessTool<TContext>[];
|
||||
toolsByName: Map<string, AgentHarnessTool<TContext>>;
|
||||
toolContext: TContext;
|
||||
}
|
||||
| undefined,
|
||||
recovery: boolean,
|
||||
): Promise<ProcedureResult> {
|
||||
const { batch } = run;
|
||||
for (let transition = 0; transition <= batch.calls.length * 2 + 1; transition += 1) {
|
||||
await materializeReady(lane, drive, run, sources, recovery);
|
||||
const current = currentBatch(lane);
|
||||
if (current === undefined) return { kind: "continue" };
|
||||
const call = current.batch.calls.find((candidate) => candidate.status !== "completed");
|
||||
if (call === undefined) throw new SessionInvariantError("Tool batch remained open after every call completed");
|
||||
if (call.status === "outcome_ready") throw new SessionInvariantError("Ready tool outcome was not materialized");
|
||||
|
||||
if (current.run.control.status === "cancel_requested") {
|
||||
const toolCall = toolCallFor(sources, call);
|
||||
if (call.status === "planned") {
|
||||
await publishToolOutcome(lane, drive, current.run, call, abortedOutcome(toolCall), recovery);
|
||||
} else {
|
||||
const checkpoint = await readCheckpoint(lane, drive, call);
|
||||
await publishToolOutcome(
|
||||
lane,
|
||||
drive,
|
||||
current.run,
|
||||
call,
|
||||
interruptedOutcome(toolCall, checkpoint),
|
||||
recovery,
|
||||
);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if (execution === undefined) throw new SessionInvariantError("Running tool batch is missing execution context");
|
||||
const started =
|
||||
call.status === "planned"
|
||||
? await startToolInvocation(
|
||||
lane,
|
||||
drive,
|
||||
current.run,
|
||||
sources,
|
||||
call,
|
||||
execution.tools,
|
||||
execution.toolContext,
|
||||
recovery,
|
||||
)
|
||||
: await recoverToolInvocation(
|
||||
lane,
|
||||
drive,
|
||||
current.run,
|
||||
sources,
|
||||
call,
|
||||
execution.toolsByName,
|
||||
execution.toolContext,
|
||||
false,
|
||||
);
|
||||
await started.completion;
|
||||
}
|
||||
throw new SessionInvariantError("Sequential tool batch exceeded its bounded transition count");
|
||||
}
|
||||
|
||||
async function runParallel<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
run: ToolsOperation,
|
||||
sources: ToolBatchSource,
|
||||
tools: AgentHarnessTool<TContext>[],
|
||||
toolsByName: Map<string, AgentHarnessTool<TContext>>,
|
||||
toolContext: TContext,
|
||||
recovery: boolean,
|
||||
): Promise<ProcedureResult> {
|
||||
const { batch } = run;
|
||||
let materialization = Promise.resolve();
|
||||
const scheduleMaterialization = (): Promise<void> => {
|
||||
const scheduled = materialization.then(() => materializeReady(lane, drive, run, sources, recovery));
|
||||
materialization = scheduled.catch(() => {});
|
||||
return scheduled;
|
||||
};
|
||||
const jobs: Promise<void>[] = [];
|
||||
for (const call of batch.calls) {
|
||||
if (call.status === "completed" || call.status === "outcome_ready") continue;
|
||||
const started =
|
||||
call.status === "planned"
|
||||
? await startToolInvocation(lane, drive, run, sources, call, tools, toolContext, recovery)
|
||||
: await recoverToolInvocation(
|
||||
lane,
|
||||
drive,
|
||||
run,
|
||||
sources,
|
||||
call,
|
||||
toolsByName,
|
||||
toolContext,
|
||||
lane.state.operation!.state.control.status === "cancel_requested",
|
||||
);
|
||||
const job = started.completion.then(async () => {
|
||||
await scheduleMaterialization();
|
||||
});
|
||||
void job.catch(() => {});
|
||||
jobs.push(job);
|
||||
}
|
||||
await Promise.all(jobs);
|
||||
await scheduleMaterialization();
|
||||
return { kind: "continue" };
|
||||
}
|
||||
|
||||
/** Execute, recover, stage, and source-order one complete durable tool batch. */
|
||||
export async function runTools<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
run: ToolsOperation,
|
||||
): Promise<ProcedureResult> {
|
||||
const batch = run.batch;
|
||||
const recovery = batch.calls.some((call) => call.status === "effect_pending" || call.status === "outcome_ready");
|
||||
if (recovery) {
|
||||
await lane.emitBatch(
|
||||
[
|
||||
{
|
||||
type: "turn_start",
|
||||
lane: lane.name,
|
||||
runId: drive.operationId,
|
||||
turnId: batch.turnId,
|
||||
recovery: true,
|
||||
},
|
||||
],
|
||||
drive.context,
|
||||
);
|
||||
}
|
||||
const sources = await readToolBatchSource(lane, drive, batch);
|
||||
await materializeReady(lane, drive, run, sources, recovery);
|
||||
const current = currentBatch(lane);
|
||||
if (current === undefined) return { kind: "continue" };
|
||||
if (current.run.control.status === "cancel_requested") {
|
||||
return runSequential(lane, drive, current.run, sources, undefined, recovery);
|
||||
}
|
||||
const config = lane.readConfig();
|
||||
const active = new Set(batch.configuration.activeToolNames);
|
||||
const tools = config.tools.filter((tool) => active.has(tool.name));
|
||||
const toolsByName = new Map(tools.map((tool) => [tool.name, tool]));
|
||||
const toolContext = await resolveToolContext(lane, drive);
|
||||
return run.settings.toolExecution === "sequential"
|
||||
? runSequential(lane, drive, current.run, sources, { tools, toolsByName, toolContext }, recovery)
|
||||
: runParallel(lane, drive, current.run, sources, tools, toolsByName, toolContext, recovery);
|
||||
}
|
||||
408
third_party/pi/packages/agent/src/harness/runtime/harness.ts
vendored
Normal file
408
third_party/pi/packages/agent/src/harness/runtime/harness.ts
vendored
Normal file
@@ -0,0 +1,408 @@
|
||||
import type { Models, RetryPolicy } from "@earendil-works/pi-ai";
|
||||
import type { QueueMode } from "../../types.ts";
|
||||
import type {
|
||||
AcquireLaneOptions,
|
||||
AgentHarness,
|
||||
AgentHarnessOptions,
|
||||
AgentLane,
|
||||
GlobalConfigEventPayload,
|
||||
LaneInfo,
|
||||
OpenOperation,
|
||||
Resources,
|
||||
} from "../agent-harness.ts";
|
||||
import { type CompactionSettings, DEFAULT_COMPACTION_SETTINGS } from "../compaction/compaction.ts";
|
||||
import { DEFAULT_RETRY_POLICY, validateCompactionSettings, validateRetryPolicy, validateToolNames } from "../config.ts";
|
||||
import type { Context } from "../context.ts";
|
||||
import { HarnessEventBus } from "../events.ts";
|
||||
import { HookRegistry } from "../hooks.ts";
|
||||
import { convertToLlm } from "../messages.ts";
|
||||
import { HarnessClosed, HarnessFault, InvalidLane, UnknownTarget } from "../result.ts";
|
||||
import { SessionInvariantError } from "../session/session.ts";
|
||||
import type { LaneConfiguration, Session } from "../session/types.ts";
|
||||
import { branchTip, deleteValue, entryLabel, laneConfig, laneState, sessionName, setValue } from "../session/values.ts";
|
||||
import type { AgentHarnessStreamOptions, AgentHarnessTool } from "../types.ts";
|
||||
import { Lane } from "./lane.ts";
|
||||
import { readLaneStorage, restoreLaneState, restoreSession } from "./restore.ts";
|
||||
import { type Config, type LaneState, SliceNotImplemented } from "./types.ts";
|
||||
|
||||
/** Runtime implementation of AgentHarness. The harness manages lanes but is not itself a lane. */
|
||||
export class Harness<TContext extends object | undefined> implements AgentHarness<TContext> {
|
||||
readonly session: Session;
|
||||
readonly models: Models;
|
||||
readonly hooks: HookRegistry;
|
||||
readonly events: HarnessEventBus;
|
||||
readonly lanesByName = new Map<string, Lane<TContext>>();
|
||||
private readonly seed: LaneConfiguration;
|
||||
private readonly configStore: { value: Config<TContext> };
|
||||
private closePromise: Promise<void> | undefined;
|
||||
private closedError: Error | undefined;
|
||||
private faultError: HarnessFault | undefined;
|
||||
|
||||
constructor(options: AgentHarnessOptions<TContext>, seed: LaneConfiguration, restored: Map<string, LaneState>) {
|
||||
this.session = options.session;
|
||||
this.models = options.models;
|
||||
this.seed = seed;
|
||||
this.events = new HarnessEventBus();
|
||||
this.hooks = new HookRegistry((error, hook, lane, context) =>
|
||||
this.events.emit(
|
||||
{
|
||||
type: "handler_error",
|
||||
kind: "hook",
|
||||
hook,
|
||||
error: error.message,
|
||||
...(error.stack === undefined ? {} : { stack: error.stack }),
|
||||
lane,
|
||||
},
|
||||
context,
|
||||
),
|
||||
);
|
||||
this.configStore = {
|
||||
value: {
|
||||
tools: options.tools ?? [],
|
||||
resources: options.resources ?? {},
|
||||
streamOptions: options.streamOptions ?? {},
|
||||
retryPolicy: options.retry ?? DEFAULT_RETRY_POLICY,
|
||||
compaction: options.compaction ?? DEFAULT_COMPACTION_SETTINGS,
|
||||
steeringMode: options.steeringMode ?? "all",
|
||||
followUpMode: options.followUpMode ?? "all",
|
||||
toolExecution: options.toolExecution ?? "parallel",
|
||||
toolContext: options.toolContext,
|
||||
systemPrompt: options.systemPrompt,
|
||||
toProviderMessages: options.toProviderMessages ?? ((messages) => convertToLlm(messages)),
|
||||
entryProjectors: options.entryProjectors ?? {},
|
||||
},
|
||||
};
|
||||
for (const [name, state] of restored) this.lanesByName.set(name, this.buildLane(name, state));
|
||||
}
|
||||
|
||||
lane(name: string, context: Context): Promise<AgentLane>;
|
||||
lane(name: string, options: AcquireLaneOptions, context: Context): Promise<AgentLane>;
|
||||
async lane(
|
||||
name: string,
|
||||
optionsOrContext: AcquireLaneOptions | Context,
|
||||
maybeContext?: Context,
|
||||
): Promise<AgentLane> {
|
||||
this.assertOpen();
|
||||
if (name.length === 0 || name.includes("\u0000")) {
|
||||
const reason = name.length === 0 ? "lane name must not be empty" : "lane name must not contain \\u0000";
|
||||
throw new InvalidLane({
|
||||
lane: name,
|
||||
reason,
|
||||
message: `Invalid lane ${JSON.stringify(name)}: ${reason}`,
|
||||
});
|
||||
}
|
||||
const options = maybeContext === undefined ? {} : (optionsOrContext as AcquireLaneOptions);
|
||||
const context = maybeContext ?? (optionsOrContext as Context);
|
||||
let lane: Lane<TContext> | undefined;
|
||||
let delivery: Promise<void> | undefined;
|
||||
try {
|
||||
await this.session.mutate(async (mutator) => {
|
||||
this.assertOpen();
|
||||
lane = this.lanesByName.get(name);
|
||||
if (lane !== undefined) return;
|
||||
|
||||
const stored = await readLaneStorage(mutator, name, context);
|
||||
if (stored.kind === "lane") {
|
||||
const restored = await restoreLaneState(mutator, name, stored, context);
|
||||
lane = this.buildLane(name, restored);
|
||||
this.lanesByName.set(name, lane);
|
||||
return;
|
||||
}
|
||||
|
||||
const tipId = stored.kind === "branch" ? stored.tip.value : (options.createAt ?? null);
|
||||
if (
|
||||
stored.kind === "absent" &&
|
||||
tipId !== null &&
|
||||
!(await mutator.getEntries([tipId], context)).has(tipId)
|
||||
) {
|
||||
throw new UnknownTarget({
|
||||
targetId: tipId,
|
||||
message: `Unknown target: ${tipId}`,
|
||||
});
|
||||
}
|
||||
const attachedConfiguration: LaneConfiguration = {
|
||||
model: { ...this.seed.model },
|
||||
thinkingLevel: this.seed.thinkingLevel,
|
||||
activeToolNames: [...this.seed.activeToolNames],
|
||||
};
|
||||
const state: LaneState = {
|
||||
tipId,
|
||||
configuration: attachedConfiguration,
|
||||
inbox: [],
|
||||
lastOperationId: null,
|
||||
operation: null,
|
||||
};
|
||||
const writes = [
|
||||
...(stored.kind === "absent" ? [setValue(branchTip(name), tipId)] : []),
|
||||
setValue(laneConfig(name), attachedConfiguration),
|
||||
setValue(laneState(name), {
|
||||
currentOperationId: null,
|
||||
lastOperationId: null,
|
||||
inbox: [],
|
||||
}),
|
||||
];
|
||||
await mutator.commit(writes, context);
|
||||
lane = this.buildLane(name, state);
|
||||
this.lanesByName.set(name, lane);
|
||||
delivery = this.events.emitBatch([{ type: "lane_created", lane: name, at: tipId }], context);
|
||||
}, context);
|
||||
} catch (error) {
|
||||
if (this.closedError !== undefined) throw this.closedError;
|
||||
if (error instanceof InvalidLane || error instanceof UnknownTarget) throw error;
|
||||
throw this.fault(error, context);
|
||||
}
|
||||
await delivery;
|
||||
if (lane === undefined)
|
||||
throw this.fault(new SessionInvariantError(`Lane ${JSON.stringify(name)} was not published`), context);
|
||||
return lane;
|
||||
}
|
||||
|
||||
async lanes(context: Context): Promise<LaneInfo[]> {
|
||||
this.assertOpen();
|
||||
const executions = await Promise.all(
|
||||
[...this.lanesByName.values()].map((lane) => lane.inspectExecution(context)),
|
||||
);
|
||||
return executions.map((execution) => ({
|
||||
name: execution.lane,
|
||||
tipId: execution.tipId,
|
||||
operation: execution.current,
|
||||
}));
|
||||
}
|
||||
|
||||
getName(context: Context): Promise<string | undefined> {
|
||||
this.assertOpen();
|
||||
return this.session.getName(context);
|
||||
}
|
||||
|
||||
async setName(name: string | undefined, context: Context): Promise<void> {
|
||||
this.assertOpen();
|
||||
let delivery: Promise<void> | undefined;
|
||||
try {
|
||||
await this.session.mutate(async (mutator) => {
|
||||
this.assertOpen();
|
||||
await mutator.commit(
|
||||
[name === undefined ? deleteValue(sessionName) : setValue(sessionName, name)],
|
||||
context,
|
||||
);
|
||||
delivery = this.events.emitBatch([{ type: "value_update", value: "session_name", name }], context);
|
||||
}, context);
|
||||
} catch (error) {
|
||||
if (this.closedError !== undefined) throw this.closedError;
|
||||
throw this.fault(error, context);
|
||||
}
|
||||
await delivery;
|
||||
}
|
||||
|
||||
getLabel(targetId: string, context: Context): Promise<string | undefined> {
|
||||
this.assertOpen();
|
||||
return this.session.getLabel(targetId, context);
|
||||
}
|
||||
|
||||
async setLabel(targetId: string, label: string | undefined, context: Context): Promise<void> {
|
||||
this.assertOpen();
|
||||
let delivery: Promise<void> | undefined;
|
||||
try {
|
||||
await this.session.mutate(async (mutator) => {
|
||||
this.assertOpen();
|
||||
const address = entryLabel(targetId);
|
||||
await mutator.commit([label === undefined ? deleteValue(address) : setValue(address, label)], context);
|
||||
delivery = this.events.emitBatch(
|
||||
[{ type: "value_update", value: "entry_label", targetId, label }],
|
||||
context,
|
||||
);
|
||||
}, context);
|
||||
} catch (error) {
|
||||
if (this.closedError !== undefined) throw this.closedError;
|
||||
throw this.fault(error, context);
|
||||
}
|
||||
await delivery;
|
||||
}
|
||||
|
||||
getTools(context: Context): Promise<AgentHarnessTool<TContext>[]> {
|
||||
return this.getConfig("tools", context);
|
||||
}
|
||||
|
||||
setTools(tools: AgentHarnessTool<TContext>[], context: Context): Promise<void> {
|
||||
validateToolNames(tools);
|
||||
return this.setConfig("tools", tools, () => ({ type: "config_update", property: "tools" }), context);
|
||||
}
|
||||
|
||||
getResources(context: Context): Promise<Resources> {
|
||||
return this.getConfig("resources", context);
|
||||
}
|
||||
|
||||
setResources(resources: Resources, context: Context): Promise<void> {
|
||||
return this.setConfig("resources", resources, () => ({ type: "config_update", property: "resources" }), context);
|
||||
}
|
||||
|
||||
getStreamOptions(context: Context): Promise<AgentHarnessStreamOptions> {
|
||||
return this.getConfig("streamOptions", context);
|
||||
}
|
||||
|
||||
setStreamOptions(options: AgentHarnessStreamOptions, context: Context): Promise<void> {
|
||||
return this.setConfig(
|
||||
"streamOptions",
|
||||
options,
|
||||
(previous, value) => ({ type: "config_update", property: "streamOptions", previous, value }),
|
||||
context,
|
||||
);
|
||||
}
|
||||
|
||||
getRetryPolicy(context: Context): Promise<RetryPolicy> {
|
||||
return this.getConfig("retryPolicy", context);
|
||||
}
|
||||
|
||||
setRetryPolicy(policy: RetryPolicy, context: Context): Promise<void> {
|
||||
validateRetryPolicy(policy);
|
||||
return this.setConfig(
|
||||
"retryPolicy",
|
||||
policy,
|
||||
(previous, value) => ({ type: "config_update", property: "retryPolicy", previous, value }),
|
||||
context,
|
||||
);
|
||||
}
|
||||
|
||||
getCompactionSettings(context: Context): Promise<CompactionSettings> {
|
||||
return this.getConfig("compaction", context);
|
||||
}
|
||||
|
||||
setCompactionSettings(compaction: CompactionSettings, context: Context): Promise<void> {
|
||||
validateCompactionSettings(compaction);
|
||||
return this.setConfig(
|
||||
"compaction",
|
||||
compaction,
|
||||
(previous, value) => ({ type: "config_update", property: "compactionSettings", previous, value }),
|
||||
context,
|
||||
);
|
||||
}
|
||||
|
||||
getSteeringMode(context: Context): Promise<QueueMode> {
|
||||
return this.getConfig("steeringMode", context);
|
||||
}
|
||||
|
||||
setSteeringMode(steeringMode: QueueMode, context: Context): Promise<void> {
|
||||
return this.setConfig(
|
||||
"steeringMode",
|
||||
steeringMode,
|
||||
(previous, value) => ({ type: "config_update", property: "steeringMode", previous, value }),
|
||||
context,
|
||||
);
|
||||
}
|
||||
|
||||
getFollowUpMode(context: Context): Promise<QueueMode> {
|
||||
return this.getConfig("followUpMode", context);
|
||||
}
|
||||
|
||||
setFollowUpMode(followUpMode: QueueMode, context: Context): Promise<void> {
|
||||
return this.setConfig(
|
||||
"followUpMode",
|
||||
followUpMode,
|
||||
(previous, value) => ({ type: "config_update", property: "followUpMode", previous, value }),
|
||||
context,
|
||||
);
|
||||
}
|
||||
|
||||
async watchSession(_context: Context): Promise<never> {
|
||||
throw new SliceNotImplemented("watchSession");
|
||||
}
|
||||
|
||||
fault(cause: unknown, context: Context): Error {
|
||||
if (this.faultError !== undefined) return this.faultError;
|
||||
if (this.closedError !== undefined) return this.closedError;
|
||||
const normalized = cause instanceof Error ? cause : new Error(String(cause));
|
||||
const fault = new HarnessFault("AgentHarness storage or invariant fault", normalized);
|
||||
this.faultError = fault;
|
||||
for (const lane of this.lanesByName.values()) lane.seal(fault);
|
||||
this.hooks.close(fault);
|
||||
void this.events.emit({ type: "fault", code: "harness_fault", message: fault.message }, context);
|
||||
this.events.close(fault);
|
||||
return fault;
|
||||
}
|
||||
|
||||
close(context: Context): Promise<void> {
|
||||
if (this.closePromise !== undefined) return this.closePromise;
|
||||
const error = new HarnessClosed();
|
||||
this.closedError = error;
|
||||
const idleCallbacks = [...this.lanesByName.values()].map((lane) => lane.seal(error));
|
||||
this.hooks.close(error);
|
||||
this.events.close(error);
|
||||
const sessionClose = this.session.close(context);
|
||||
this.closePromise = Promise.all([sessionClose, ...idleCallbacks]).then(() => undefined);
|
||||
return this.closePromise;
|
||||
}
|
||||
|
||||
private buildLane(name: string, state: LaneState): Lane<TContext> {
|
||||
return new Lane<TContext>(
|
||||
name,
|
||||
this.session,
|
||||
this.models,
|
||||
this.hooks,
|
||||
state,
|
||||
(cause, context) => this.fault(cause, context),
|
||||
(events, context) => this.events.emitBatch(events, context),
|
||||
(snapshot, filter, context, resnapshot) => this.events.watch(snapshot, filter, context, resnapshot),
|
||||
() => this.configStore.value,
|
||||
);
|
||||
}
|
||||
|
||||
private async getConfig<TKey extends keyof Config<TContext>>(
|
||||
key: TKey,
|
||||
_context: Context,
|
||||
): Promise<Config<TContext>[TKey]> {
|
||||
this.assertOpen();
|
||||
return this.configStore.value[key];
|
||||
}
|
||||
|
||||
private async setConfig<TKey extends keyof Config<TContext>>(
|
||||
key: TKey,
|
||||
value: Config<TContext>[TKey],
|
||||
event: (previous: Config<TContext>[TKey], value: Config<TContext>[TKey]) => GlobalConfigEventPayload,
|
||||
context: Context,
|
||||
): Promise<void> {
|
||||
this.assertOpen();
|
||||
const previous = this.configStore.value[key];
|
||||
this.configStore.value = { ...this.configStore.value, [key]: value };
|
||||
await this.events.emit(event(previous, value), context);
|
||||
}
|
||||
|
||||
private assertOpen(): void {
|
||||
if (this.faultError !== undefined) throw this.faultError;
|
||||
if (this.closedError !== undefined) throw this.closedError;
|
||||
}
|
||||
}
|
||||
|
||||
/** Attach runtime without starting provider, tool, hook, or timer effects. */
|
||||
export async function createAgentHarness<TContext extends object | undefined = object | undefined>(
|
||||
options: AgentHarnessOptions<TContext>,
|
||||
context: Context,
|
||||
): Promise<{ harness: AgentHarness<TContext>; open: OpenOperation[] }> {
|
||||
const tools = options.tools ?? [];
|
||||
validateToolNames(tools);
|
||||
validateRetryPolicy(options.retry ?? DEFAULT_RETRY_POLICY);
|
||||
validateCompactionSettings(options.compaction ?? DEFAULT_COMPACTION_SETTINGS);
|
||||
const seed: LaneConfiguration = {
|
||||
model: { provider: options.model.provider, modelId: options.model.id },
|
||||
thinkingLevel: options.thinkingLevel ?? "off",
|
||||
activeToolNames: [...(options.activeToolNames ?? tools.map((tool) => tool.name))],
|
||||
};
|
||||
try {
|
||||
const restored = await restoreSession(options.session, context);
|
||||
const open = [...restored].flatMap(([lane, state]): OpenOperation[] => {
|
||||
const operation = state.operation;
|
||||
return operation === null
|
||||
? []
|
||||
: [
|
||||
{
|
||||
lane,
|
||||
operationId: operation.meta.operationId,
|
||||
kind: operation.meta.intent.kind,
|
||||
startedAt: operation.meta.startedAt,
|
||||
...(operation.state.control.status === "cancel_requested" ? { aborting: true as const } : {}),
|
||||
},
|
||||
];
|
||||
});
|
||||
return { harness: new Harness(options, seed, restored), open };
|
||||
} catch (error) {
|
||||
throw new HarnessFault("AgentHarness storage or invariant fault", error);
|
||||
}
|
||||
}
|
||||
1
third_party/pi/packages/agent/src/harness/runtime/index.ts
vendored
Normal file
1
third_party/pi/packages/agent/src/harness/runtime/index.ts
vendored
Normal file
@@ -0,0 +1 @@
|
||||
export { createAgentHarness } from "./harness.ts";
|
||||
2012
third_party/pi/packages/agent/src/harness/runtime/lane.ts
vendored
Normal file
2012
third_party/pi/packages/agent/src/harness/runtime/lane.ts
vendored
Normal file
File diff suppressed because it is too large
Load Diff
117
third_party/pi/packages/agent/src/harness/runtime/progress.ts
vendored
Normal file
117
third_party/pi/packages/agent/src/harness/runtime/progress.ts
vendored
Normal file
@@ -0,0 +1,117 @@
|
||||
import type { AssistantMessageFrame } from "@earendil-works/pi-ai";
|
||||
import type { AgentToolResult } from "../../types.ts";
|
||||
import type { Context } from "../context.ts";
|
||||
import type { SessionReader, Write } from "../session/types.ts";
|
||||
import { appendList, pendingAssistantFrames, pendingToolOutput, setValue } from "../session/values.ts";
|
||||
import type { Lane } from "./lane.ts";
|
||||
import type { Drive, LaneState } from "./types.ts";
|
||||
|
||||
export interface ProgressChannel<T> {
|
||||
write(item: T): void;
|
||||
seal(): void;
|
||||
drain(): Promise<void>;
|
||||
}
|
||||
|
||||
export async function readAssistantFrames(
|
||||
reader: SessionReader,
|
||||
operationId: string,
|
||||
responseEntryId: string,
|
||||
context: Context,
|
||||
): Promise<AssistantMessageFrame[]> {
|
||||
const frames: AssistantMessageFrame[] = [];
|
||||
let cursor: { seq: number } | undefined;
|
||||
for (;;) {
|
||||
const page = await reader.readList(
|
||||
pendingAssistantFrames(operationId, responseEntryId),
|
||||
{ order: "asc", limit: 1_000, ...(cursor === undefined ? {} : { cursor }) },
|
||||
context,
|
||||
);
|
||||
frames.push(...page.map(({ value }) => value));
|
||||
if (page.length < 1_000) return frames;
|
||||
cursor = { seq: page[page.length - 1]!.seq };
|
||||
}
|
||||
}
|
||||
|
||||
function openProgress<TContext extends object | undefined, T>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
commitWrite: (item: T) => Write,
|
||||
stillOwns: (state: LaneState) => boolean,
|
||||
): ProgressChannel<T> {
|
||||
let sealed = false;
|
||||
let latest: Promise<void> = Promise.resolve();
|
||||
return {
|
||||
write(item) {
|
||||
if (sealed) return;
|
||||
const write = lane
|
||||
.command((projection) => {
|
||||
if (!stillOwns(projection)) return { kind: "return", result: undefined };
|
||||
return {
|
||||
kind: "commit",
|
||||
writes: [commitWrite(item)],
|
||||
next: projection,
|
||||
materialize: () => undefined,
|
||||
};
|
||||
}, drive.context)
|
||||
.then(() => undefined);
|
||||
latest = write;
|
||||
void write.catch(() => {});
|
||||
},
|
||||
seal() {
|
||||
sealed = true;
|
||||
},
|
||||
async drain() {
|
||||
await latest;
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
export function openFrameProgress<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
responseEntryId: string,
|
||||
): ProgressChannel<AssistantMessageFrame> {
|
||||
const address = pendingAssistantFrames(drive.operationId, responseEntryId);
|
||||
return openProgress(
|
||||
lane,
|
||||
drive,
|
||||
(frame) => appendList(address, frame),
|
||||
(state) => {
|
||||
const run = state.operation?.state;
|
||||
if (run === undefined) return false;
|
||||
return (
|
||||
(run.at === "assistant.effect_pending" || run.at === "deferred.effect_pending") &&
|
||||
run.responseEntryId === responseEntryId
|
||||
);
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
export function openToolProgress<TContext extends object | undefined>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
turnId: string,
|
||||
sourceIndex: number,
|
||||
invocationId: string,
|
||||
): ProgressChannel<AgentToolResult<unknown>> {
|
||||
const address = pendingToolOutput(drive.operationId, invocationId);
|
||||
return openProgress(
|
||||
lane,
|
||||
drive,
|
||||
(snapshot) => setValue(address, snapshot),
|
||||
(state) => {
|
||||
const operation = state.operation;
|
||||
if (operation?.state.at !== "tools") return false;
|
||||
const batch = operation.state.batch;
|
||||
return (
|
||||
batch.turnId === turnId &&
|
||||
batch.calls.some(
|
||||
(call) =>
|
||||
call.sourceIndex === sourceIndex &&
|
||||
call.resultEntryId === invocationId &&
|
||||
call.status === "effect_pending",
|
||||
)
|
||||
);
|
||||
},
|
||||
);
|
||||
}
|
||||
232
third_party/pi/packages/agent/src/harness/runtime/reducer.ts
vendored
Normal file
232
third_party/pi/packages/agent/src/harness/runtime/reducer.ts
vendored
Normal file
@@ -0,0 +1,232 @@
|
||||
import type { HarnessEvent, LaneSnapshot, LaneWatchEvent } from "../agent-harness.ts";
|
||||
import type { OperationResultRecord } from "../session/types.ts";
|
||||
|
||||
export type LaneSnapshotReduction = "rebase" | undefined;
|
||||
|
||||
type LaneOperationSnapshot = NonNullable<LaneSnapshot["operation"]>;
|
||||
|
||||
function upsertTool(operation: LaneOperationSnapshot, tool: LaneOperationSnapshot["runningTools"][number]): void {
|
||||
const index = operation.runningTools.findIndex((candidate) => candidate.toolCallId === tool.toolCallId);
|
||||
if (index === -1) operation.runningTools.push(tool);
|
||||
else operation.runningTools[index] = tool;
|
||||
}
|
||||
|
||||
function matchingOperation(
|
||||
snapshot: LaneSnapshot,
|
||||
operationId: string,
|
||||
): NonNullable<LaneSnapshot["operation"]> | undefined {
|
||||
return snapshot.operation?.id === operationId ? snapshot.operation : undefined;
|
||||
}
|
||||
|
||||
/** Apply one harness event to a mutable lane snapshot. Navigation completion requires a fresh snapshot. */
|
||||
export function reduceLaneSnapshot(
|
||||
snapshot: LaneSnapshot,
|
||||
event: HarnessEvent | LaneWatchEvent,
|
||||
): LaneSnapshotReduction {
|
||||
if ("lane" in event && event.lane !== undefined && event.lane !== snapshot.lane && event.type !== "usage") return;
|
||||
switch (event.type) {
|
||||
case "run_start":
|
||||
snapshot.operation = {
|
||||
id: event.runId,
|
||||
kind: "run",
|
||||
startedAt: event.startedAt,
|
||||
fromTipId: snapshot.tipId,
|
||||
status: "open",
|
||||
runningTools: [],
|
||||
};
|
||||
return;
|
||||
case "compaction_start":
|
||||
if (snapshot.operation !== null) return;
|
||||
snapshot.operation = {
|
||||
id: event.runId,
|
||||
kind: "compaction",
|
||||
startedAt: event.startedAt,
|
||||
fromTipId: snapshot.tipId,
|
||||
status: "open",
|
||||
runningTools: [],
|
||||
};
|
||||
return;
|
||||
case "navigation_start":
|
||||
snapshot.operation = {
|
||||
id: event.runId,
|
||||
kind: "navigation",
|
||||
startedAt: event.startedAt,
|
||||
fromTipId: snapshot.tipId,
|
||||
status: "open",
|
||||
runningTools: [],
|
||||
};
|
||||
return;
|
||||
case "operation_abort": {
|
||||
const operation = matchingOperation(snapshot, event.operationId);
|
||||
if (operation !== undefined) operation.status = "aborting";
|
||||
return;
|
||||
}
|
||||
case "run_resume": {
|
||||
const operation = matchingOperation(snapshot, event.runId);
|
||||
if (operation !== undefined) delete operation.deferred;
|
||||
return;
|
||||
}
|
||||
case "run_suspend": {
|
||||
const operation = matchingOperation(snapshot, event.runId);
|
||||
if (operation === undefined) return;
|
||||
delete operation.streamingMessage;
|
||||
operation.deferred = { handle: event.deferred, poll: event.poll };
|
||||
return;
|
||||
}
|
||||
case "retry_scheduled": {
|
||||
const operation = matchingOperation(snapshot, event.runId);
|
||||
if (operation !== undefined) {
|
||||
operation.retry = {
|
||||
attempt: event.attempt,
|
||||
maxAttempts: event.maxAttempts,
|
||||
nextAttemptAt: event.notBefore,
|
||||
};
|
||||
}
|
||||
return;
|
||||
}
|
||||
case "retry_start":
|
||||
case "retry_end": {
|
||||
const operation = matchingOperation(snapshot, event.runId);
|
||||
if (operation !== undefined) delete operation.retry;
|
||||
return;
|
||||
}
|
||||
case "message_start": {
|
||||
if (
|
||||
event.runId === undefined ||
|
||||
event.message.role !== "assistant" ||
|
||||
event.message.stopReason !== "pending"
|
||||
) {
|
||||
return;
|
||||
}
|
||||
const operation = matchingOperation(snapshot, event.runId);
|
||||
if (operation !== undefined) operation.streamingMessage = event.message;
|
||||
return;
|
||||
}
|
||||
case "message_update": {
|
||||
if (event.message.role !== "assistant") return;
|
||||
const operation = matchingOperation(snapshot, event.runId);
|
||||
if (operation !== undefined) operation.streamingMessage = event.message;
|
||||
return;
|
||||
}
|
||||
case "message_end": {
|
||||
if (event.runId === undefined) return;
|
||||
const operation = matchingOperation(snapshot, event.runId);
|
||||
if (operation !== undefined) delete operation.streamingMessage;
|
||||
return;
|
||||
}
|
||||
case "tool_start": {
|
||||
const operation = matchingOperation(snapshot, event.runId);
|
||||
if (operation === undefined) return;
|
||||
upsertTool(operation, {
|
||||
status: "running",
|
||||
toolCallId: event.toolCallId,
|
||||
toolName: event.toolName,
|
||||
args: event.args,
|
||||
});
|
||||
return;
|
||||
}
|
||||
case "tool_update": {
|
||||
const operation = matchingOperation(snapshot, event.runId);
|
||||
const tool = operation?.runningTools.find((candidate) => candidate.toolCallId === event.toolCallId);
|
||||
if (tool?.status === "running") tool.result = event.partialResult;
|
||||
return;
|
||||
}
|
||||
case "tool_end": {
|
||||
const operation = matchingOperation(snapshot, event.runId);
|
||||
if (operation === undefined) return;
|
||||
const index = operation.runningTools.findIndex((candidate) => candidate.toolCallId === event.toolCallId);
|
||||
const current = operation.runningTools[index];
|
||||
if (current === undefined) return;
|
||||
operation.runningTools[index] = {
|
||||
status: "settled",
|
||||
toolCallId: event.toolCallId,
|
||||
toolName: event.toolName,
|
||||
args: current.args,
|
||||
result: event.result,
|
||||
isError: event.isError,
|
||||
};
|
||||
return;
|
||||
}
|
||||
case "entry_added": {
|
||||
if (event.entry.type === "message" && event.entry.message.role === "toolResult") {
|
||||
const operation = snapshot.operation;
|
||||
if (operation !== null) {
|
||||
const toolCallId = event.entry.message.toolCallId;
|
||||
const index = operation.runningTools.findIndex((candidate) => candidate.toolCallId === toolCallId);
|
||||
if (index !== -1) operation.runningTools.splice(index, 1);
|
||||
}
|
||||
}
|
||||
if (event.entry.type === "compaction") snapshot.transcript.splice(0, snapshot.transcript.length, event.entry);
|
||||
else snapshot.transcript.push(event.entry);
|
||||
snapshot.tipId = event.entry.id;
|
||||
if (event.entry.type === "message") snapshot.stats.messageCount += 1;
|
||||
return;
|
||||
}
|
||||
case "queue_update":
|
||||
snapshot.queues = event.queues;
|
||||
return;
|
||||
case "usage":
|
||||
snapshot.stats.usage = event.totals;
|
||||
return;
|
||||
case "config_update":
|
||||
if (!("lane" in event) || event.lane !== snapshot.lane) return;
|
||||
switch (event.property) {
|
||||
case "model":
|
||||
snapshot.configuration.model = event.value;
|
||||
return;
|
||||
case "thinkingLevel":
|
||||
snapshot.configuration.thinkingLevel = event.value;
|
||||
return;
|
||||
case "activeTools":
|
||||
snapshot.configuration.activeToolNames = event.value;
|
||||
return;
|
||||
}
|
||||
return;
|
||||
case "run_end": {
|
||||
const operation = matchingOperation(snapshot, event.runId);
|
||||
if (operation?.kind !== "run") return;
|
||||
const record: OperationResultRecord = {
|
||||
operationId: event.runId,
|
||||
kind: "run",
|
||||
status: event.status,
|
||||
...(event.status === "failed" ? { error: event.error } : {}),
|
||||
fromTipId: event.fromTipId,
|
||||
tipId: event.tipId,
|
||||
startedAt: operation.startedAt,
|
||||
endedAt: event.endedAt,
|
||||
};
|
||||
snapshot.lastResult = record;
|
||||
snapshot.operation = null;
|
||||
snapshot.tipId = event.tipId;
|
||||
return;
|
||||
}
|
||||
case "compaction_end": {
|
||||
const operation = matchingOperation(snapshot, event.runId);
|
||||
if (operation?.kind !== "compaction") return;
|
||||
const record: OperationResultRecord = {
|
||||
operationId: event.runId,
|
||||
kind: "compaction",
|
||||
status: event.status,
|
||||
...(event.status === "failed" ? { error: event.error } : {}),
|
||||
fromTipId: operation.fromTipId,
|
||||
tipId: snapshot.tipId,
|
||||
startedAt: operation.startedAt,
|
||||
endedAt: event.endedAt,
|
||||
};
|
||||
snapshot.lastResult = record;
|
||||
snapshot.operation = null;
|
||||
return;
|
||||
}
|
||||
case "navigation_end":
|
||||
return "rebase";
|
||||
case "fault":
|
||||
snapshot.faulted = true;
|
||||
return;
|
||||
case "handler_error":
|
||||
case "turn_start":
|
||||
case "turn_end":
|
||||
case "value_update":
|
||||
case "lane_created":
|
||||
return;
|
||||
}
|
||||
}
|
||||
171
third_party/pi/packages/agent/src/harness/runtime/restore.ts
vendored
Normal file
171
third_party/pi/packages/agent/src/harness/runtime/restore.ts
vendored
Normal file
@@ -0,0 +1,171 @@
|
||||
import type { Context } from "../context.ts";
|
||||
import { SessionInvariantError } from "../session/session.ts";
|
||||
import type {
|
||||
LaneState as DurableLaneState,
|
||||
LaneConfiguration,
|
||||
OperationMeta,
|
||||
OperationState,
|
||||
Session,
|
||||
SessionReader,
|
||||
} from "../session/types.ts";
|
||||
import {
|
||||
branchTip,
|
||||
branchTipInventoryPrefix,
|
||||
laneConfig,
|
||||
laneState as laneStateValue,
|
||||
operationMeta,
|
||||
operationState,
|
||||
type StoredValue,
|
||||
} from "../session/values.ts";
|
||||
import type { LaneState } from "./types.ts";
|
||||
|
||||
function isSummaryState(state: OperationState): state is Extract<OperationState, { at: `summary.${string}` }> {
|
||||
return state.at.startsWith("summary.");
|
||||
}
|
||||
|
||||
function stateMatchesIntent(intent: OperationMeta["intent"], state: OperationState): boolean {
|
||||
if (intent.kind === "compaction") return isSummaryState(state) && state.task.boundary.kind === "finish";
|
||||
if (intent.kind === "navigation") {
|
||||
if (state.at === "navigation.ready_to_commit") {
|
||||
return !intent.summarize && state.targetId === intent.targetId && state.label === intent.label;
|
||||
}
|
||||
return (
|
||||
intent.summarize &&
|
||||
isSummaryState(state) &&
|
||||
state.task.boundary.kind === "commit_navigation" &&
|
||||
state.task.boundary.targetId === intent.targetId &&
|
||||
state.task.boundary.label === intent.label &&
|
||||
state.task.customInstructions === intent.customInstructions
|
||||
);
|
||||
}
|
||||
return (
|
||||
state.at !== "navigation.ready_to_commit" &&
|
||||
(!isSummaryState(state) || state.task.boundary.kind === "resume_checkpoint")
|
||||
);
|
||||
}
|
||||
|
||||
type LaneValues = {
|
||||
tip: StoredValue<string | null> | undefined;
|
||||
configuration: StoredValue<LaneConfiguration> | undefined;
|
||||
laneState: StoredValue<DurableLaneState> | undefined;
|
||||
};
|
||||
|
||||
export type ClassifiedLaneStorage =
|
||||
| { kind: "absent" }
|
||||
| { kind: "branch"; tip: StoredValue<string | null> }
|
||||
| {
|
||||
kind: "lane";
|
||||
tip: StoredValue<string | null>;
|
||||
configuration: StoredValue<LaneConfiguration>;
|
||||
laneState: StoredValue<DurableLaneState>;
|
||||
};
|
||||
|
||||
function classifyLaneStorage(lane: string, values: LaneValues): ClassifiedLaneStorage {
|
||||
const { tip, configuration, laneState } = values;
|
||||
if (tip === undefined && configuration === undefined && laneState === undefined) {
|
||||
return { kind: "absent" };
|
||||
}
|
||||
if (tip !== undefined && configuration === undefined && laneState === undefined) {
|
||||
return { kind: "branch", tip };
|
||||
}
|
||||
if (tip === undefined) throw new SessionInvariantError(`Lane ${JSON.stringify(lane)} is missing branch.tip`);
|
||||
if (configuration === undefined)
|
||||
throw new SessionInvariantError(`Lane ${JSON.stringify(lane)} is missing lane.config`);
|
||||
if (laneState === undefined) throw new SessionInvariantError(`Lane ${JSON.stringify(lane)} is missing lane.state`);
|
||||
return { kind: "lane", tip, configuration, laneState };
|
||||
}
|
||||
|
||||
export async function readLaneStorage(
|
||||
reader: SessionReader,
|
||||
lane: string,
|
||||
context: Context,
|
||||
): Promise<ClassifiedLaneStorage> {
|
||||
const [tip, configuration, laneState] = await Promise.all([
|
||||
reader.getValue(branchTip(lane), context),
|
||||
reader.getValue(laneConfig(lane), context),
|
||||
reader.getValue(laneStateValue(lane), context),
|
||||
]);
|
||||
return classifyLaneStorage(lane, { tip, configuration, laneState });
|
||||
}
|
||||
|
||||
/** Restore every complete configured AgentLane in one coherent Session read. */
|
||||
export function restoreSession(session: Session, context: Context): Promise<Map<string, LaneState>> {
|
||||
return session.mutate(async (reader) => {
|
||||
const [tips, configurations, states] = await Promise.all([
|
||||
reader.scanValues(branchTipInventoryPrefix(), context),
|
||||
reader.scanValues(laneConfig(""), context),
|
||||
reader.scanValues(laneStateValue(""), context),
|
||||
]);
|
||||
const tipByLane = new Map(tips.map((value) => [value.address.key, value]));
|
||||
const configurationByLane = new Map(configurations.map((value) => [value.address.key, value]));
|
||||
const stateByLane = new Map(states.map((value) => [value.address.key, value]));
|
||||
const names = new Set([...tipByLane.keys(), ...configurationByLane.keys(), ...stateByLane.keys()]);
|
||||
const restored = new Map<string, LaneState>();
|
||||
for (const lane of names) {
|
||||
const stored = classifyLaneStorage(lane, {
|
||||
tip: tipByLane.get(lane),
|
||||
configuration: configurationByLane.get(lane),
|
||||
laneState: stateByLane.get(lane),
|
||||
});
|
||||
if (stored.kind !== "lane") continue;
|
||||
restored.set(lane, await restoreLaneState(reader, lane, stored, context));
|
||||
}
|
||||
return restored;
|
||||
}, context);
|
||||
}
|
||||
|
||||
/** Restore one configured lane without starting work or interpreting its state. */
|
||||
export function restoreLane(session: Session, lane: string, context: Context): Promise<LaneState> {
|
||||
return session.mutate(async (reader) => {
|
||||
const stored = await readLaneStorage(reader, lane, context);
|
||||
if (stored.kind === "absent") {
|
||||
throw new SessionInvariantError(`Lane ${JSON.stringify(lane)} is missing branch.tip`);
|
||||
}
|
||||
if (stored.kind === "branch") {
|
||||
throw new SessionInvariantError(`Lane ${JSON.stringify(lane)} is missing lane.config`);
|
||||
}
|
||||
return restoreLaneState(reader, lane, stored, context);
|
||||
}, context);
|
||||
}
|
||||
|
||||
export async function restoreLaneState(
|
||||
reader: SessionReader,
|
||||
lane: string,
|
||||
stored: Extract<ClassifiedLaneStorage, { kind: "lane" }>,
|
||||
context: Context,
|
||||
): Promise<LaneState> {
|
||||
const operationId = stored.laneState.value.currentOperationId;
|
||||
let operation: LaneState["operation"] = null;
|
||||
if (operationId !== null) {
|
||||
const [meta, state] = await Promise.all([
|
||||
reader.getValue(operationMeta(operationId), context),
|
||||
reader.getValue(operationState(operationId), context),
|
||||
]);
|
||||
if (meta === undefined) throw new SessionInvariantError(`Operation ${operationId} is missing op.meta`);
|
||||
if (state === undefined) throw new SessionInvariantError(`Operation ${operationId} is missing op.state`);
|
||||
if (meta.value.operationId !== operationId) {
|
||||
throw new SessionInvariantError(
|
||||
`Operation ${operationId} metadata names operation ${JSON.stringify(meta.value.operationId)}`,
|
||||
);
|
||||
}
|
||||
if (meta.value.lane !== lane) {
|
||||
throw new SessionInvariantError(
|
||||
`Operation ${operationId} belongs to lane ${JSON.stringify(meta.value.lane)}, not ${JSON.stringify(lane)}`,
|
||||
);
|
||||
}
|
||||
if (!stateMatchesIntent(meta.value.intent, state.value)) {
|
||||
throw new SessionInvariantError(
|
||||
`Operation ${operationId} intent ${meta.value.intent.kind} does not match state ${state.value.at}`,
|
||||
);
|
||||
}
|
||||
operation = { meta: meta.value, state: state.value };
|
||||
}
|
||||
|
||||
return {
|
||||
tipId: stored.tip.value,
|
||||
configuration: stored.configuration.value,
|
||||
inbox: stored.laneState.value.inbox,
|
||||
lastOperationId: stored.laneState.value.lastOperationId,
|
||||
operation,
|
||||
};
|
||||
}
|
||||
129
third_party/pi/packages/agent/src/harness/runtime/transcript.ts
vendored
Normal file
129
third_party/pi/packages/agent/src/harness/runtime/transcript.ts
vendored
Normal file
@@ -0,0 +1,129 @@
|
||||
import type { AgentMessage } from "../../types.ts";
|
||||
import type { HarnessEvent, LaneQueuedItem } from "../agent-harness.ts";
|
||||
import type { Context } from "../context.ts";
|
||||
import { materializeCommittedEntry } from "../session/commit.ts";
|
||||
import { buildSessionContext } from "../session/context.ts";
|
||||
import { SessionInvariantError } from "../session/session.ts";
|
||||
import type { CommitResult, Entry, InboxItem, NewEntry, OperationState, SessionReader } from "../session/types.ts";
|
||||
import { pendingEntry } from "../session/values.ts";
|
||||
import type { Lane } from "./lane.ts";
|
||||
import type { ContinueOperationResult, Drive } from "./types.ts";
|
||||
|
||||
export function chainEntries<T extends { id: string }>(
|
||||
parentId: string | null,
|
||||
items: readonly T[],
|
||||
): Array<T & { parentId: string | null }> {
|
||||
return items.map((item) => {
|
||||
const entry = { ...item, parentId };
|
||||
parentId = item.id;
|
||||
return entry;
|
||||
});
|
||||
}
|
||||
|
||||
export function entryLifecycleEvents(entry: Entry, lane: string, runId?: string): HarnessEvent[] {
|
||||
const operation = runId === undefined ? {} : { runId };
|
||||
return entry.type === "message"
|
||||
? [
|
||||
{ type: "message_start", lane, ...operation, message: entry.message },
|
||||
{ type: "message_end", lane, ...operation, message: entry.message, entryId: entry.id },
|
||||
{ type: "entry_added", lane, entry },
|
||||
]
|
||||
: [{ type: "entry_added", lane, entry }];
|
||||
}
|
||||
|
||||
export function committedEntryEvents(
|
||||
entries: readonly NewEntry[],
|
||||
commit: CommitResult,
|
||||
lane: string,
|
||||
runId?: string,
|
||||
firstWriteIndex = 0,
|
||||
): HarnessEvent[] {
|
||||
return entries.flatMap((entry, index) =>
|
||||
entryLifecycleEvents(
|
||||
materializeCommittedEntry(entry, commit.seqs[firstWriteIndex + index]!, commit.timestamp),
|
||||
lane,
|
||||
runId,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
export function readBoundedEntries<TContext extends object | undefined, TState extends OperationState>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
capability: TState,
|
||||
): Promise<ContinueOperationResult<Entry[]>> {
|
||||
return lane.continueOperation(
|
||||
capability,
|
||||
async (state, _current, _meta, reader) => {
|
||||
if (state.tipId === null) throw new SessionInvariantError("Run operation has no Branch tip");
|
||||
const entries = await reader.scanBranch(
|
||||
{ start: state.tipId, stopAtType: "compaction", order: "newestFirst" },
|
||||
drive.context,
|
||||
);
|
||||
return { kind: "return", result: entries.reverse() };
|
||||
},
|
||||
drive.context,
|
||||
);
|
||||
}
|
||||
|
||||
export async function readBoundedContext<TContext extends object | undefined, TState extends OperationState>(
|
||||
lane: Lane<TContext>,
|
||||
drive: Drive,
|
||||
capability: TState,
|
||||
): Promise<ContinueOperationResult<AgentMessage[]>> {
|
||||
const entries = await readBoundedEntries(lane, drive, capability);
|
||||
if (entries.kind === "cancel_requested") return entries;
|
||||
return {
|
||||
kind: "result",
|
||||
value: await buildSessionContext(
|
||||
entries.value,
|
||||
{ entryProjectors: lane.readConfig().entryProjectors },
|
||||
drive.context,
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
export function readLaneQueues(
|
||||
reader: SessionReader,
|
||||
inbox: readonly InboxItem[],
|
||||
context: Context,
|
||||
): Promise<LaneQueuedItem[]> {
|
||||
return Promise.all(
|
||||
inbox.map(async (item): Promise<LaneQueuedItem> => {
|
||||
const stored = await reader.getValue(pendingEntry(item.entryId), context);
|
||||
if (stored === undefined) {
|
||||
throw new SessionInvariantError(`Pending ${item.kind} entry ${item.entryId} is missing its payload`);
|
||||
}
|
||||
if (stored.value.type === "message") {
|
||||
return { entryId: item.entryId, kind: item.kind, type: "message", message: stored.value.payload };
|
||||
}
|
||||
if (item.kind !== "write") {
|
||||
throw new SessionInvariantError(`Pending ${item.kind} entry ${item.entryId} is not a message`);
|
||||
}
|
||||
return {
|
||||
entryId: item.entryId,
|
||||
kind: "write",
|
||||
type: "custom",
|
||||
customType: stored.value.customType,
|
||||
...(stored.value.payload === undefined ? {} : { data: stored.value.payload }),
|
||||
};
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
export function readPendingMessages(
|
||||
reader: SessionReader,
|
||||
ids: readonly string[],
|
||||
description: string,
|
||||
context: Context,
|
||||
): Promise<Array<{ entryId: string; message: AgentMessage }>> {
|
||||
return Promise.all(
|
||||
ids.map(async (entryId) => {
|
||||
const value = await reader.getValue(pendingEntry(entryId), context);
|
||||
if (value?.value.type !== "message") {
|
||||
throw new SessionInvariantError(`${description} ${entryId} is missing its message payload`);
|
||||
}
|
||||
return { entryId, message: value.value.payload };
|
||||
}),
|
||||
);
|
||||
}
|
||||
146
third_party/pi/packages/agent/src/harness/runtime/types.ts
vendored
Normal file
146
third_party/pi/packages/agent/src/harness/runtime/types.ts
vendored
Normal file
@@ -0,0 +1,146 @@
|
||||
import type { RetryPolicy } from "@earendil-works/pi-ai";
|
||||
import type { QueueMode } from "../../types.ts";
|
||||
import type { AgentHarnessOptions, DriveOptions, DriveOutcome, HarnessEvent, Resources } from "../agent-harness.ts";
|
||||
import type { CompactionSettings } from "../compaction/compaction.ts";
|
||||
import { type Context, withoutAbortSignal } from "../context.ts";
|
||||
import { createGate, type Gate, type GateControl } from "../execution/effect-gate.ts";
|
||||
import type {
|
||||
CommitResult,
|
||||
InboxItem,
|
||||
LaneConfiguration,
|
||||
Operation,
|
||||
OperationResultRecord,
|
||||
OperationState,
|
||||
Write,
|
||||
} from "../session/types.ts";
|
||||
import type { AgentHarnessStreamOptions, AgentHarnessTool } from "../types.ts";
|
||||
|
||||
export class SliceNotImplemented extends Error {
|
||||
constructor(operation: string) {
|
||||
super(`${operation} is not implemented until its later AgentHarness slice`);
|
||||
this.name = "SliceNotImplemented";
|
||||
}
|
||||
}
|
||||
|
||||
/** Current process-local harness configuration. */
|
||||
export interface Config<TContext extends object | undefined> {
|
||||
readonly tools: AgentHarnessTool<TContext>[];
|
||||
readonly resources: Resources;
|
||||
readonly streamOptions: AgentHarnessStreamOptions;
|
||||
readonly retryPolicy: RetryPolicy;
|
||||
readonly compaction: CompactionSettings;
|
||||
readonly steeringMode: QueueMode;
|
||||
readonly followUpMode: QueueMode;
|
||||
readonly toolExecution: "sequential" | "parallel";
|
||||
readonly toolContext: AgentHarnessOptions<TContext>["toolContext"];
|
||||
readonly systemPrompt: AgentHarnessOptions<TContext>["systemPrompt"];
|
||||
readonly toProviderMessages: NonNullable<AgentHarnessOptions<TContext>["toProviderMessages"]>;
|
||||
readonly entryProjectors: Readonly<NonNullable<AgentHarnessOptions<TContext>["entryProjectors"]>>;
|
||||
}
|
||||
|
||||
/** The current durable state owned by one lane. */
|
||||
export interface LaneState {
|
||||
readonly tipId: string | null;
|
||||
readonly configuration: LaneConfiguration;
|
||||
readonly inbox: InboxItem[];
|
||||
readonly lastOperationId: string | null;
|
||||
readonly operation: Operation | null;
|
||||
}
|
||||
|
||||
type Synchronous<TResult> = TResult extends PromiseLike<unknown> ? never : TResult;
|
||||
|
||||
interface CommitDecision<TResult> {
|
||||
kind: "commit";
|
||||
writes: Write[];
|
||||
materialize(commit: CommitResult): Synchronous<TResult>;
|
||||
events?(commit: CommitResult): HarnessEvent[];
|
||||
}
|
||||
|
||||
/** One effect-free decision made on a lane's serialized mutation line. */
|
||||
export type LaneCommand<TResult> =
|
||||
| (CommitDecision<TResult> & { next: LaneState })
|
||||
| { kind: "return"; result: TResult }
|
||||
| { kind: "reject"; error: Error };
|
||||
|
||||
export type ContinueOperationResult<TResult> = { kind: "cancel_requested" } | { kind: "result"; value: TResult };
|
||||
|
||||
/** A durable operation transition. The Lane pairs the state write with projection publication. */
|
||||
type LanePatch = Partial<Pick<LaneState, "tipId" | "configuration" | "inbox">>;
|
||||
|
||||
interface FinishDecision<TResult> {
|
||||
kind: "finish";
|
||||
writes: Write[];
|
||||
record: OperationResultRecord;
|
||||
lane?: LanePatch;
|
||||
materialize(commit: CommitResult): Synchronous<TResult>;
|
||||
events?(commit: CommitResult): HarnessEvent[];
|
||||
}
|
||||
|
||||
export type OperationCommand<TResult> =
|
||||
| (CommitDecision<TResult> & { operationState: OperationState; lane?: LanePatch })
|
||||
| FinishDecision<TResult>
|
||||
| { kind: "return"; result: TResult };
|
||||
|
||||
/** One installed process-local drive pass. */
|
||||
export class Drive {
|
||||
readonly operationId: string;
|
||||
readonly completion: Promise<DriveOutcome>;
|
||||
readonly gate: Gate;
|
||||
readonly context: Context;
|
||||
readonly waitForRetry: boolean;
|
||||
readonly closeSignal: AbortSignal;
|
||||
deferredPermits: number;
|
||||
|
||||
private readonly control: GateControl;
|
||||
private readonly closeController: AbortController;
|
||||
private readonly resolveCompletion: (outcome: DriveOutcome) => void;
|
||||
private readonly rejectCompletion: (error: unknown) => void;
|
||||
|
||||
constructor(options: DriveOptions, context: Context) {
|
||||
this.operationId = options.operationId;
|
||||
this.context = withoutAbortSignal(context);
|
||||
this.waitForRetry = options.waitForRetry ?? false;
|
||||
this.deferredPermits = options.pollDeferred === true ? 1 : 0;
|
||||
let resolveCompletion!: (outcome: DriveOutcome) => void;
|
||||
let rejectCompletion!: (error: unknown) => void;
|
||||
this.completion = new Promise<DriveOutcome>((resolve, reject) => {
|
||||
resolveCompletion = resolve;
|
||||
rejectCompletion = reject;
|
||||
});
|
||||
this.resolveCompletion = resolveCompletion;
|
||||
this.rejectCompletion = rejectCompletion;
|
||||
void this.completion.catch(() => {});
|
||||
const { gate, control } = createGate();
|
||||
this.gate = gate;
|
||||
this.control = control;
|
||||
this.closeController = new AbortController();
|
||||
this.closeSignal = this.closeController.signal;
|
||||
}
|
||||
|
||||
settle(outcome: DriveOutcome): void {
|
||||
this.resolveCompletion(outcome);
|
||||
}
|
||||
|
||||
fail(error: unknown): void {
|
||||
this.rejectCompletion(error);
|
||||
}
|
||||
|
||||
beginAbort(cancellation: Promise<void>): void {
|
||||
this.control.beginAbort(cancellation);
|
||||
}
|
||||
|
||||
signalAbort(): void {
|
||||
this.control.signalAbort();
|
||||
}
|
||||
|
||||
closeGate(error: Error): void {
|
||||
this.control.close(error);
|
||||
if (!this.closeSignal.aborted) this.closeController.abort(error);
|
||||
this.rejectCompletion(error);
|
||||
}
|
||||
}
|
||||
|
||||
export type ProcedureResult =
|
||||
| { kind: "continue" }
|
||||
| { kind: "waiting"; outcome: DriveOutcome }
|
||||
| { kind: "settled"; outcome: OperationResultRecord };
|
||||
116
third_party/pi/packages/agent/src/harness/session/commit.ts
vendored
Normal file
116
third_party/pi/packages/agent/src/harness/session/commit.ts
vendored
Normal file
@@ -0,0 +1,116 @@
|
||||
import type { CommitResult, Entry, EntryWrite, NewEntry, UsageRow, UsageWrite, Write } from "./types.ts";
|
||||
|
||||
export type CommittedEntryWrite = Entry & { kind: "entry" };
|
||||
export type CommittedUsageWrite = UsageRow & { kind: "usage" };
|
||||
export interface CommittedValueSetWrite {
|
||||
kind: "value";
|
||||
op: "set";
|
||||
seq: number;
|
||||
namespace: string;
|
||||
key: string;
|
||||
value: unknown;
|
||||
}
|
||||
export interface CommittedValueDeleteWrite {
|
||||
kind: "value";
|
||||
op: "delete";
|
||||
seq: number;
|
||||
namespace: string;
|
||||
key: string;
|
||||
}
|
||||
export interface CommittedListAppendWrite {
|
||||
kind: "list";
|
||||
op: "append";
|
||||
seq: number;
|
||||
namespace: string;
|
||||
key: string;
|
||||
value: unknown;
|
||||
}
|
||||
export interface CommittedListDeleteWrite {
|
||||
kind: "list";
|
||||
op: "delete";
|
||||
seq: number;
|
||||
namespace: string;
|
||||
key: string;
|
||||
}
|
||||
export type CommittedWrite =
|
||||
| CommittedEntryWrite
|
||||
| CommittedUsageWrite
|
||||
| CommittedValueSetWrite
|
||||
| CommittedValueDeleteWrite
|
||||
| CommittedListAppendWrite
|
||||
| CommittedListDeleteWrite;
|
||||
|
||||
export interface PreparedCommit {
|
||||
writes: CommittedWrite[];
|
||||
result: Omit<CommitResult, "stats">;
|
||||
}
|
||||
|
||||
export interface CommitValidationState {
|
||||
hasEntryOrUsageId(id: string): boolean;
|
||||
hasEntryId(id: string): boolean;
|
||||
}
|
||||
|
||||
export function insertEntry(entry: NewEntry): EntryWrite {
|
||||
return { kind: "entry", entry };
|
||||
}
|
||||
|
||||
export function insertUsage(row: Omit<UsageRow, "seq">): UsageWrite {
|
||||
return { kind: "usage", row };
|
||||
}
|
||||
|
||||
export function commitWrite(write: Write, seq: number, timestamp: number): CommittedWrite {
|
||||
switch (write.kind) {
|
||||
case "entry":
|
||||
return { kind: "entry", ...write.entry, seq, timestamp };
|
||||
case "usage":
|
||||
return { kind: "usage", ...write.row, seq };
|
||||
case "value":
|
||||
return write.op === "set"
|
||||
? { kind: "value", op: "set", seq, namespace: write.namespace, key: write.key, value: write.value }
|
||||
: { kind: "value", op: "delete", seq, namespace: write.namespace, key: write.key };
|
||||
case "list":
|
||||
return write.op === "append"
|
||||
? { kind: "list", op: "append", seq, namespace: write.namespace, key: write.key, value: write.value }
|
||||
: { kind: "list", op: "delete", seq, namespace: write.namespace, key: write.key };
|
||||
}
|
||||
}
|
||||
|
||||
export function materializeCommittedEntry(entry: NewEntry, seq: number, timestamp: number): Entry {
|
||||
return { ...entry, seq, timestamp };
|
||||
}
|
||||
|
||||
export function prepareStorageCommit(writes: Write[], firstSeq: number, timestamp: number): PreparedCommit {
|
||||
const committedWrites = writes.map((write, index) => commitWrite(write, firstSeq + index, timestamp));
|
||||
return {
|
||||
writes: committedWrites,
|
||||
result: { firstSeq, seqs: committedWrites.map((write) => write.seq), timestamp },
|
||||
};
|
||||
}
|
||||
|
||||
export function validateCommittedWrites(
|
||||
writes: readonly CommittedWrite[],
|
||||
firstSeq: number,
|
||||
state: CommitValidationState,
|
||||
): void {
|
||||
let previousSeq = firstSeq - 1;
|
||||
const transactionIds = new Set<string>();
|
||||
const transactionEntryIds = new Set<string>();
|
||||
for (const write of writes) {
|
||||
if (write.seq <= previousSeq) throw new Error(`Non-monotonic storage sequence: ${write.seq}`);
|
||||
previousSeq = write.seq;
|
||||
if (write.kind !== "entry" && write.kind !== "usage") continue;
|
||||
if (state.hasEntryOrUsageId(write.id) || transactionIds.has(write.id)) {
|
||||
throw new Error(`Duplicate entry or usage id: ${write.id}`);
|
||||
}
|
||||
if (
|
||||
write.kind === "entry" &&
|
||||
write.parentId !== null &&
|
||||
!state.hasEntryId(write.parentId) &&
|
||||
!transactionEntryIds.has(write.parentId)
|
||||
) {
|
||||
throw new Error(`Missing parent entry: ${write.parentId}`);
|
||||
}
|
||||
transactionIds.add(write.id);
|
||||
if (write.kind === "entry") transactionEntryIds.add(write.id);
|
||||
}
|
||||
}
|
||||
64
third_party/pi/packages/agent/src/harness/session/context.ts
vendored
Normal file
64
third_party/pi/packages/agent/src/harness/session/context.ts
vendored
Normal file
@@ -0,0 +1,64 @@
|
||||
import type { AgentMessage } from "../../types.ts";
|
||||
import type { Context } from "../context.ts";
|
||||
import { createBranchSummaryMessage, createCompactionSummaryMessage } from "../messages.ts";
|
||||
import type { CompactionEntry, Entry, EntryProjector } from "./types.ts";
|
||||
|
||||
export interface SessionContextBuildOptions {
|
||||
entryProjectors?: Readonly<Record<string, EntryProjector>>;
|
||||
}
|
||||
|
||||
export function buildContextEntries(pathEntries: readonly Entry[]): Entry[] {
|
||||
let compaction: CompactionEntry | undefined;
|
||||
let compactionIndex = -1;
|
||||
for (let index = pathEntries.length - 1; index >= 0; index--) {
|
||||
const entry = pathEntries[index];
|
||||
if (entry?.type === "compaction") {
|
||||
compaction = entry;
|
||||
compactionIndex = index;
|
||||
break;
|
||||
}
|
||||
}
|
||||
return compaction === undefined ? [...pathEntries] : [compaction, ...pathEntries.slice(compactionIndex + 1)];
|
||||
}
|
||||
|
||||
function isContextMessage(message: AgentMessage): boolean {
|
||||
return (
|
||||
message.role !== "assistant" ||
|
||||
(message.stopReason !== "error" && message.stopReason !== "aborted" && message.stopReason !== "deferred")
|
||||
);
|
||||
}
|
||||
|
||||
export function sessionEntryToContextMessages(entry: Entry): AgentMessage[] {
|
||||
switch (entry.type) {
|
||||
case "message":
|
||||
return isContextMessage(entry.message) ? [entry.message] : [];
|
||||
case "compaction":
|
||||
return [
|
||||
createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp),
|
||||
...entry.retainedTail.filter(isContextMessage),
|
||||
];
|
||||
case "branch_summary":
|
||||
return entry.summary ? [createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp)] : [];
|
||||
case "custom":
|
||||
return [];
|
||||
}
|
||||
}
|
||||
|
||||
export async function buildSessionContext(
|
||||
pathEntries: readonly Entry[],
|
||||
options: SessionContextBuildOptions | undefined,
|
||||
context: Context,
|
||||
): Promise<AgentMessage[]> {
|
||||
options ??= {};
|
||||
const entries = buildContextEntries(pathEntries);
|
||||
const messages: AgentMessage[] = [];
|
||||
for (const entry of entries) {
|
||||
if (entry.type !== "custom") {
|
||||
messages.push(...sessionEntryToContextMessages(entry));
|
||||
continue;
|
||||
}
|
||||
const projector = options.entryProjectors?.[entry.customType];
|
||||
if (projector !== undefined) messages.push(...((await projector(entry, context)) ?? []));
|
||||
}
|
||||
return messages;
|
||||
}
|
||||
32
third_party/pi/packages/agent/src/harness/session/fork-policy.ts
vendored
Normal file
32
third_party/pi/packages/agent/src/harness/session/fork-policy.ts
vendored
Normal file
@@ -0,0 +1,32 @@
|
||||
/**
|
||||
* Final action for current state at one address:
|
||||
* - copy: emit the current value in the destination;
|
||||
* - exclude: omit it from the destination;
|
||||
* - reconstruct: do not copy the row because lane handling emits a coherent replacement.
|
||||
*/
|
||||
export type ForkDisposition = "copy" | "exclude" | "reconstruct";
|
||||
|
||||
/** Decide the final fork action for one current scalar or list address. */
|
||||
export function classifyForkAddress(
|
||||
address: { readonly namespace: string; readonly key: string },
|
||||
scope: "branch" | "tree",
|
||||
isEntryCopied: (entryId: string) => boolean,
|
||||
): ForkDisposition {
|
||||
switch (address.namespace) {
|
||||
case "pi.session.name":
|
||||
return "copy";
|
||||
case "pi.entry.label":
|
||||
return isEntryCopied(address.key) ? "copy" : "exclude";
|
||||
case "pi.branch.tip":
|
||||
case "pi.lane.config":
|
||||
case "pi.lane.state":
|
||||
return "reconstruct";
|
||||
case "pi.result":
|
||||
return "exclude";
|
||||
}
|
||||
if (address.namespace.startsWith("pi.op.") || address.namespace.startsWith("pi.pending.")) return "exclude";
|
||||
if (address.namespace === "pi" || address.namespace.startsWith("pi.")) {
|
||||
throw new Error(`Unknown reserved fork namespace: ${address.namespace}`);
|
||||
}
|
||||
return scope === "tree" ? "copy" : "exclude";
|
||||
}
|
||||
161
third_party/pi/packages/agent/src/harness/session/fork.ts
vendored
Normal file
161
third_party/pi/packages/agent/src/harness/session/fork.ts
vendored
Normal file
@@ -0,0 +1,161 @@
|
||||
import type { CommittedWrite } from "./commit.ts";
|
||||
import { classifyForkAddress } from "./fork-policy.ts";
|
||||
import type { Entry, ForkOptions } from "./types.ts";
|
||||
import { branchTip, laneConfig, laneState, type StoredValue, type Value, value } from "./values.ts";
|
||||
|
||||
export interface ForkSourceSnapshot {
|
||||
entries: Entry[];
|
||||
scalarValues: StoredValue<unknown>[];
|
||||
/** False when a backend supplied only the requested branch rather than the full tree. */
|
||||
entriesComplete?: boolean;
|
||||
}
|
||||
|
||||
export interface ForkDestinationSnapshot {
|
||||
entries: Map<string, Entry>;
|
||||
scalarValues: StoredValue<unknown>[];
|
||||
nextSeq: number;
|
||||
}
|
||||
|
||||
function storedValuesInNamespace<T>(values: readonly StoredValue<unknown>[], address: Value<T>): StoredValue<T>[] {
|
||||
return values.filter((stored) => stored.address.namespace === address.namespace) as StoredValue<T>[];
|
||||
}
|
||||
|
||||
function findStoredValue<T>(values: readonly StoredValue<unknown>[], address: Value<T>): StoredValue<T> | undefined {
|
||||
return values.find(
|
||||
(stored) => stored.address.namespace === address.namespace && stored.address.key === address.key,
|
||||
) as StoredValue<T> | undefined;
|
||||
}
|
||||
|
||||
/** Build the complete logical state for a forked destination session. */
|
||||
export function createForkSnapshot(source: ForkSourceSnapshot, options: ForkOptions): ForkDestinationSnapshot {
|
||||
const sourceEntries = new Map(source.entries.map((entry) => [entry.id, entry]));
|
||||
const sourceTips = storedValuesInNamespace(source.scalarValues, branchTip(""));
|
||||
validateForkSourceSnapshot(source, sourceEntries, sourceTips, options);
|
||||
|
||||
const { entryIds, destinationTips } = selectForkContents(sourceEntries, sourceTips, options);
|
||||
const entries = new Map<string, Entry>();
|
||||
for (const id of entryIds) entries.set(id, sourceEntries.get(id)!);
|
||||
|
||||
const scalarValues: StoredValue<unknown>[] = [];
|
||||
let nextSeq = Math.max(0, ...[...entries.values()].map((entry) => entry.seq)) + 1;
|
||||
const store = <T>(address: Value<T>, storedValue: T): void => {
|
||||
scalarValues.push({
|
||||
address: value<unknown>(address.namespace, address.key),
|
||||
value: storedValue,
|
||||
seq: nextSeq++,
|
||||
});
|
||||
};
|
||||
for (const [branch, tipId] of destinationTips) {
|
||||
const configuration = findStoredValue(source.scalarValues, laneConfig(branch));
|
||||
store(branchTip(branch), tipId);
|
||||
if (configuration !== undefined) {
|
||||
store(laneConfig(branch), configuration.value);
|
||||
store(laneState(branch), { currentOperationId: null, lastOperationId: null, inbox: [] });
|
||||
}
|
||||
}
|
||||
for (const stored of source.scalarValues) {
|
||||
switch (classifyForkAddress(stored.address, options.scope, (entryId) => entryIds.has(entryId))) {
|
||||
case "copy":
|
||||
store(stored.address, stored.value);
|
||||
break;
|
||||
case "exclude":
|
||||
case "reconstruct":
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
return { entries, scalarValues, nextSeq };
|
||||
}
|
||||
|
||||
export function forkSnapshotWrites(snapshot: ForkDestinationSnapshot): CommittedWrite[] {
|
||||
const writes: CommittedWrite[] = [];
|
||||
for (const entry of snapshot.entries.values()) writes.push({ kind: "entry", ...entry });
|
||||
for (const stored of snapshot.scalarValues) {
|
||||
writes.push({
|
||||
kind: "value",
|
||||
op: "set",
|
||||
seq: stored.seq,
|
||||
namespace: stored.address.namespace,
|
||||
key: stored.address.key,
|
||||
value: stored.value,
|
||||
});
|
||||
}
|
||||
return writes.sort((left, right) => left.seq - right.seq);
|
||||
}
|
||||
|
||||
function selectForkContents(
|
||||
sourceEntries: Map<string, Entry>,
|
||||
sourceTips: StoredValue<string | null>[],
|
||||
options: ForkOptions,
|
||||
): {
|
||||
entryIds: Set<string>;
|
||||
destinationTips: Map<string, string | null>;
|
||||
} {
|
||||
const entryIds = new Set<string>();
|
||||
const destinationTips = new Map<string, string | null>();
|
||||
if (options.scope === "tree") {
|
||||
for (const id of sourceEntries.keys()) entryIds.add(id);
|
||||
for (const stored of sourceTips) destinationTips.set(stored.address.key, stored.value);
|
||||
} else {
|
||||
const sourceTip = sourceTips.find((stored) => stored.address.key === options.branch);
|
||||
if (sourceTip === undefined) throw new Error(`Unknown source branch: ${options.branch}`);
|
||||
|
||||
const requested = options.entryId ?? sourceTip.value;
|
||||
let found = requested === null;
|
||||
let tipId: string | null = null;
|
||||
let entryId = sourceTip.value;
|
||||
while (entryId !== null) {
|
||||
const entry = sourceEntries.get(entryId);
|
||||
if (entry === undefined) throw new Error(`Corrupt source branch: missing parent ${entryId}`);
|
||||
if (entry.id === requested) {
|
||||
found = true;
|
||||
tipId = options.position === "before" ? entry.parentId : entry.id;
|
||||
if (options.position !== "before") entryIds.add(entry.id);
|
||||
} else if (found) {
|
||||
entryIds.add(entry.id);
|
||||
}
|
||||
entryId = entry.parentId;
|
||||
}
|
||||
if (!found) {
|
||||
throw new Error(`Fork entry ${requested} is not on source branch ${JSON.stringify(options.branch)}`);
|
||||
}
|
||||
destinationTips.set(options.branch, tipId);
|
||||
}
|
||||
return { entryIds, destinationTips };
|
||||
}
|
||||
|
||||
function validateForkSourceSnapshot(
|
||||
source: ForkSourceSnapshot,
|
||||
sourceEntries: Map<string, Entry>,
|
||||
sourceTips: StoredValue<string | null>[],
|
||||
options: ForkOptions,
|
||||
): void {
|
||||
const sourceTipKeys = new Set(sourceTips.map((stored) => stored.address.key));
|
||||
|
||||
for (const stored of source.scalarValues) {
|
||||
if (
|
||||
(stored.address.namespace === laneConfig("").namespace ||
|
||||
stored.address.namespace === laneState("").namespace) &&
|
||||
!sourceTipKeys.has(stored.address.key)
|
||||
) {
|
||||
throw new Error(`Source session branch ${JSON.stringify(stored.address.key)} is missing branch.tip`);
|
||||
}
|
||||
}
|
||||
for (const tip of sourceTips) {
|
||||
const configuration = findStoredValue(source.scalarValues, laneConfig(tip.address.key));
|
||||
const state = findStoredValue(source.scalarValues, laneState(tip.address.key));
|
||||
if ((configuration === undefined) !== (state === undefined)) {
|
||||
throw new Error(`Source session branch ${JSON.stringify(tip.address.key)} has incomplete lane state`);
|
||||
}
|
||||
if (options.scope === "branch" && tip.address.key === options.branch && configuration === undefined) {
|
||||
throw new Error(`Source branch ${JSON.stringify(options.branch)} is not a configured AgentLane`);
|
||||
}
|
||||
if (
|
||||
(source.entriesComplete !== false || options.scope === "tree") &&
|
||||
tip.value !== null &&
|
||||
!sourceEntries.has(tip.value)
|
||||
) {
|
||||
throw new Error(`Source session branch ${JSON.stringify(tip.address.key)} has an unknown tip`);
|
||||
}
|
||||
}
|
||||
}
|
||||
264
third_party/pi/packages/agent/src/harness/session/in-memory-storage-state.ts
vendored
Normal file
264
third_party/pi/packages/agent/src/harness/session/in-memory-storage-state.ts
vendored
Normal file
@@ -0,0 +1,264 @@
|
||||
import { addUsage, emptyUsage } from "../utils/usage.ts";
|
||||
import { type CommittedWrite, type PreparedCommit, prepareStorageCommit, validateCommittedWrites } from "./commit.ts";
|
||||
|
||||
export type {
|
||||
CommittedEntryWrite,
|
||||
CommittedListAppendWrite,
|
||||
CommittedListDeleteWrite,
|
||||
CommittedUsageWrite,
|
||||
CommittedValueDeleteWrite,
|
||||
CommittedValueSetWrite,
|
||||
CommittedWrite,
|
||||
PreparedCommit,
|
||||
} from "./commit.ts";
|
||||
|
||||
import type {
|
||||
Entry,
|
||||
EntryScan,
|
||||
EntryStructure,
|
||||
SessionStats,
|
||||
StorageBranchScan,
|
||||
UsageRow,
|
||||
UsageScan,
|
||||
Write,
|
||||
} from "./types.ts";
|
||||
import {
|
||||
type ListElement,
|
||||
type ListReadOptions,
|
||||
list,
|
||||
resolveListReadOptions,
|
||||
type StoredValue,
|
||||
type Value,
|
||||
type ValueList,
|
||||
value,
|
||||
} from "./values.ts";
|
||||
|
||||
interface StoredListSnapshot {
|
||||
address: ValueList<unknown>;
|
||||
elements: ListElement<unknown>[];
|
||||
}
|
||||
|
||||
function physicalKey(namespace: string, key: string): string {
|
||||
return `${namespace}\u0000${key}`;
|
||||
}
|
||||
|
||||
function compareKeys(left: string, right: string): number {
|
||||
const leftCodePoints = Array.from(left, (character) => character.codePointAt(0)!);
|
||||
const rightCodePoints = Array.from(right, (character) => character.codePointAt(0)!);
|
||||
const length = Math.min(leftCodePoints.length, rightCodePoints.length);
|
||||
for (let index = 0; index < length; index++) {
|
||||
const difference = leftCodePoints[index]! - rightCodePoints[index]!;
|
||||
if (difference !== 0) return difference;
|
||||
}
|
||||
return leftCodePoints.length - rightCodePoints.length;
|
||||
}
|
||||
|
||||
/**
|
||||
* Complete materialized session state for MemoryStorage and JsonlStorage.
|
||||
*
|
||||
* This is intentionally unsuitable for database backends and long-running sessions that may not fit in memory.
|
||||
* Those backends should query indexed durable state and update durable aggregates within each commit transaction.
|
||||
*/
|
||||
export class InMemoryStorageState {
|
||||
private readonly entries: Map<string, Entry>;
|
||||
private readonly entriesBySeq: Entry[];
|
||||
private readonly scalarValues: Map<string, StoredValue<unknown>>;
|
||||
private readonly listValues: Map<string, StoredListSnapshot>;
|
||||
private readonly usage: Map<string, UsageRow>;
|
||||
private stats: SessionStats;
|
||||
private nextSeq: number;
|
||||
|
||||
constructor() {
|
||||
this.entries = new Map();
|
||||
this.entriesBySeq = [];
|
||||
this.scalarValues = new Map();
|
||||
this.listValues = new Map();
|
||||
this.usage = new Map();
|
||||
this.stats = { messageCount: 0, usage: emptyUsage() };
|
||||
this.nextSeq = 1;
|
||||
}
|
||||
|
||||
prepareCommit(writes: Write[], timestamp: number): PreparedCommit {
|
||||
const prepared = prepareStorageCommit(writes, this.nextSeq, timestamp);
|
||||
this.validateCommitted(prepared.writes);
|
||||
return prepared;
|
||||
}
|
||||
|
||||
validateCommitted(writes: readonly CommittedWrite[]): void {
|
||||
validateCommittedWrites(writes, this.nextSeq, {
|
||||
hasEntryOrUsageId: (id) => this.entries.has(id) || this.usage.has(id),
|
||||
hasEntryId: (id) => this.entries.has(id),
|
||||
});
|
||||
}
|
||||
|
||||
/** Apply writes already accepted by validateCommitted() and return the post-apply totals. */
|
||||
applyValidated(writes: readonly CommittedWrite[]): SessionStats {
|
||||
for (const write of writes) {
|
||||
switch (write.kind) {
|
||||
case "entry": {
|
||||
const { kind: _kind, ...entry } = write;
|
||||
this.entries.set(entry.id, entry);
|
||||
this.entriesBySeq.push(entry);
|
||||
if (entry.type === "message") this.stats = { ...this.stats, messageCount: this.stats.messageCount + 1 };
|
||||
break;
|
||||
}
|
||||
case "usage": {
|
||||
const { kind: _kind, ...row } = write;
|
||||
this.usage.set(row.id, row);
|
||||
this.stats = { ...this.stats, usage: addUsage(this.stats.usage, row.usage) };
|
||||
break;
|
||||
}
|
||||
case "value": {
|
||||
const key = physicalKey(write.namespace, write.key);
|
||||
if (write.op === "delete") {
|
||||
this.scalarValues.delete(key);
|
||||
} else {
|
||||
this.scalarValues.set(key, {
|
||||
address: value<unknown>(write.namespace, write.key),
|
||||
value: write.value,
|
||||
seq: write.seq,
|
||||
});
|
||||
}
|
||||
break;
|
||||
}
|
||||
case "list": {
|
||||
const key = physicalKey(write.namespace, write.key);
|
||||
if (write.op === "delete") {
|
||||
this.listValues.delete(key);
|
||||
} else {
|
||||
const stored = this.listValues.get(key);
|
||||
const element = { seq: write.seq, value: write.value };
|
||||
if (stored === undefined) {
|
||||
this.listValues.set(key, {
|
||||
address: list<unknown>(write.namespace, write.key),
|
||||
elements: [element],
|
||||
});
|
||||
} else {
|
||||
stored.elements.push(element);
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
this.nextSeq = write.seq + 1;
|
||||
}
|
||||
return this.stats;
|
||||
}
|
||||
|
||||
advanceNextSeq(nextSeq: number): void {
|
||||
if (!Number.isSafeInteger(nextSeq) || nextSeq < 1) {
|
||||
throw new Error(`Invalid storage sequence high-water mark: ${nextSeq}`);
|
||||
}
|
||||
this.nextSeq = Math.max(this.nextSeq, nextSeq);
|
||||
}
|
||||
|
||||
getEntries(ids: readonly string[]): Map<string, Entry> {
|
||||
const found = new Map<string, Entry>();
|
||||
for (const id of ids) {
|
||||
const entry = this.entries.get(id);
|
||||
if (entry !== undefined) found.set(id, entry);
|
||||
}
|
||||
return found;
|
||||
}
|
||||
|
||||
getValue<T>(address: Value<T>): StoredValue<T> | undefined {
|
||||
return this.scalarValues.get(physicalKey(address.namespace, address.key)) as StoredValue<T> | undefined;
|
||||
}
|
||||
|
||||
scanValues<T>(prefix: Value<T>): StoredValue<T>[] {
|
||||
return [...this.scalarValues.values()]
|
||||
.filter((stored) => stored.address.namespace === prefix.namespace && stored.address.key.startsWith(prefix.key))
|
||||
.sort((left, right) => compareKeys(left.address.key, right.address.key)) as StoredValue<T>[];
|
||||
}
|
||||
|
||||
readList<T>(address: ValueList<T>, options?: ListReadOptions): ListElement<T>[] {
|
||||
const resolved = resolveListReadOptions(options);
|
||||
const elements = this.listValues.get(physicalKey(address.namespace, address.key))?.elements ?? [];
|
||||
const filtered = elements.filter((element) => {
|
||||
if (resolved.cursor === undefined) return true;
|
||||
return resolved.order === "asc" ? element.seq > resolved.cursor.seq : element.seq < resolved.cursor.seq;
|
||||
});
|
||||
const ordered = resolved.order === "asc" ? filtered : [...filtered].reverse();
|
||||
return ordered.slice(0, resolved.limit) as ListElement<T>[];
|
||||
}
|
||||
|
||||
scanBranch(query: StorageBranchScan): Entry[] {
|
||||
const start = this.entries.get(query.start);
|
||||
if (start === undefined) throw new Error(`Unknown branch start: ${query.start}`);
|
||||
|
||||
const path: Entry[] = [];
|
||||
let entry: Entry | undefined = start;
|
||||
while (entry !== undefined) {
|
||||
path.push(entry);
|
||||
if (entry.parentId === null) break;
|
||||
entry = this.entries.get(entry.parentId);
|
||||
if (entry === undefined) throw new Error("Corrupt branch: missing parent");
|
||||
}
|
||||
if (query.order === "oldestFirst") path.reverse();
|
||||
|
||||
const stopped: Entry[] = [];
|
||||
for (const candidate of path) {
|
||||
stopped.push(candidate);
|
||||
if (candidate.id === query.stopAtId || candidate.type === query.stopAtType) break;
|
||||
}
|
||||
const filtered = stopped
|
||||
.filter((candidate) => query.type === undefined || candidate.type === query.type)
|
||||
.filter((candidate) => query.customType === undefined || candidate.customType === query.customType)
|
||||
.filter(
|
||||
(candidate) =>
|
||||
query.cursor === undefined ||
|
||||
(query.order === "oldestFirst" ? candidate.seq > query.cursor.seq : candidate.seq < query.cursor.seq),
|
||||
);
|
||||
return query.limit === undefined ? filtered : filtered.slice(0, Math.max(0, query.limit));
|
||||
}
|
||||
|
||||
scanBranchStructure(query: StorageBranchScan): EntryStructure[] {
|
||||
return this.scanBranch(query).map((entry) => ({
|
||||
id: entry.id,
|
||||
parentId: entry.parentId,
|
||||
seq: entry.seq,
|
||||
timestamp: entry.timestamp,
|
||||
type: entry.type,
|
||||
...(entry.customType === undefined ? {} : { customType: entry.customType }),
|
||||
}));
|
||||
}
|
||||
|
||||
scanEntries(query: EntryScan): Entry[] {
|
||||
const limit = query.limit === undefined ? Number.POSITIVE_INFINITY : Math.max(0, Math.trunc(query.limit));
|
||||
const entries: Entry[] = [];
|
||||
const descending = query.order === "desc";
|
||||
let index = descending ? this.entriesBySeq.length - 1 : 0;
|
||||
while (index >= 0 && index < this.entriesBySeq.length && entries.length < limit) {
|
||||
const entry = this.entriesBySeq[index]!;
|
||||
if (
|
||||
(query.type === undefined || entry.type === query.type) &&
|
||||
(query.customType === undefined || entry.customType === query.customType) &&
|
||||
(query.fromSeq === undefined || entry.seq >= query.fromSeq) &&
|
||||
(query.toSeq === undefined || entry.seq <= query.toSeq)
|
||||
) {
|
||||
entries.push(entry);
|
||||
}
|
||||
index += descending ? -1 : 1;
|
||||
}
|
||||
return entries;
|
||||
}
|
||||
|
||||
scanUsage(query: UsageScan): UsageRow[] {
|
||||
const rows = [...this.usage.values()]
|
||||
.filter((row) => query.fromSeq === undefined || row.seq >= query.fromSeq)
|
||||
.filter((row) => query.toSeq === undefined || row.seq <= query.toSeq)
|
||||
.sort((left, right) => (query.order === "desc" ? right.seq - left.seq : left.seq - right.seq));
|
||||
return query.limit === undefined ? rows : rows.slice(0, Math.max(0, query.limit));
|
||||
}
|
||||
|
||||
getStats(): SessionStats {
|
||||
return this.stats;
|
||||
}
|
||||
|
||||
snapshotEntriesAndValues(): { entries: Entry[]; scalarValues: StoredValue<unknown>[] } {
|
||||
return {
|
||||
entries: [...this.entries.values()].sort((left, right) => left.seq - right.seq),
|
||||
scalarValues: [...this.scalarValues.values()],
|
||||
};
|
||||
}
|
||||
}
|
||||
35
third_party/pi/packages/agent/src/harness/session/index.ts
vendored
Normal file
35
third_party/pi/packages/agent/src/harness/session/index.ts
vendored
Normal file
@@ -0,0 +1,35 @@
|
||||
export type {
|
||||
CommittedEntryWrite,
|
||||
CommittedListAppendWrite,
|
||||
CommittedListDeleteWrite,
|
||||
CommittedUsageWrite,
|
||||
CommittedValueDeleteWrite,
|
||||
CommittedValueSetWrite,
|
||||
CommittedWrite,
|
||||
CommitValidationState,
|
||||
PreparedCommit,
|
||||
} from "./commit.ts";
|
||||
export { commitWrite, insertEntry, insertUsage, prepareStorageCommit, validateCommittedWrites } from "./commit.ts";
|
||||
export { createForkSnapshot, type ForkSourceSnapshot } from "./fork.ts";
|
||||
export { classifyForkAddress, type ForkDisposition } from "./fork-policy.ts";
|
||||
export {
|
||||
JSONL_STORAGE_VERSION,
|
||||
type JsonlSessionCreateOptions,
|
||||
type JsonlSessionListOptions,
|
||||
type JsonlSessionMetadata,
|
||||
JsonlSessionRepo,
|
||||
type JsonlSessionRepoOptions,
|
||||
} from "./jsonl/index.ts";
|
||||
export type { MemorySessionRepoOptions } from "./memory.ts";
|
||||
export { MemorySessionRepo } from "./memory.ts";
|
||||
export {
|
||||
SessionBranchExistsError,
|
||||
SessionInvalidBranchError,
|
||||
SessionInvariantError,
|
||||
SessionPendingAssistantMessageError,
|
||||
SessionUnknownTargetError,
|
||||
StorageBackedSession,
|
||||
type StorageBackedSessionOptions,
|
||||
} from "./session.ts";
|
||||
export * from "./types.ts";
|
||||
export * from "./values.ts";
|
||||
63
third_party/pi/packages/agent/src/harness/session/jsonl/codec.ts
vendored
Normal file
63
third_party/pi/packages/agent/src/harness/session/jsonl/codec.ts
vendored
Normal file
@@ -0,0 +1,63 @@
|
||||
import { err, ok, type Result } from "../../types.ts";
|
||||
import { JSONL_FORMAT_VERSION, type JsonlStorageHeader } from "./types.ts";
|
||||
|
||||
export interface LegacyV3SessionHeader {
|
||||
type: "session";
|
||||
version: 3;
|
||||
id: string;
|
||||
timestamp: string;
|
||||
cwd: string;
|
||||
parentSession?: string;
|
||||
}
|
||||
|
||||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||
return typeof value === "object" && value !== null && !Array.isArray(value);
|
||||
}
|
||||
|
||||
function isSafeIntegerAtLeast(value: unknown, minimum: number): value is number {
|
||||
return Number.isSafeInteger(value) && (value as number) >= minimum;
|
||||
}
|
||||
|
||||
export function isLegacyV3SessionHeader(value: unknown): value is LegacyV3SessionHeader {
|
||||
return (
|
||||
isRecord(value) &&
|
||||
value.type === "session" &&
|
||||
value.version === 3 &&
|
||||
typeof value.id === "string" &&
|
||||
typeof value.cwd === "string" &&
|
||||
typeof value.timestamp === "string" &&
|
||||
Number.isFinite(Date.parse(value.timestamp)) &&
|
||||
(value.parentSession === undefined || typeof value.parentSession === "string")
|
||||
);
|
||||
}
|
||||
|
||||
export function isJsonlStorageHeader(value: unknown): value is JsonlStorageHeader {
|
||||
return (
|
||||
isRecord(value) &&
|
||||
value.kind === "header" &&
|
||||
value.v === JSONL_FORMAT_VERSION &&
|
||||
typeof value.id === "string" &&
|
||||
typeof value.cwd === "string" &&
|
||||
isSafeIntegerAtLeast(value.storageVersion, 1) &&
|
||||
isSafeIntegerAtLeast(value.createdAt, 0) &&
|
||||
(value.nextSeq === undefined || isSafeIntegerAtLeast(value.nextSeq, 1)) &&
|
||||
(value.parentSessionId === undefined || typeof value.parentSessionId === "string") &&
|
||||
(value.legacyParentSessionPath === undefined || typeof value.legacyParentSessionPath === "string")
|
||||
);
|
||||
}
|
||||
|
||||
export type JsonlParsedSessionHeader =
|
||||
| { format: "v4"; header: JsonlStorageHeader }
|
||||
| { format: "v3-legacy"; header: LegacyV3SessionHeader };
|
||||
|
||||
export function parseJsonlSessionHeader(line: string): Result<JsonlParsedSessionHeader, Error> {
|
||||
let value: unknown;
|
||||
try {
|
||||
value = JSON.parse(line);
|
||||
} catch (error) {
|
||||
return err(new Error("Invalid JSONL session header: not valid JSON", { cause: error }));
|
||||
}
|
||||
if (isJsonlStorageHeader(value)) return ok({ format: "v4", header: value });
|
||||
if (isLegacyV3SessionHeader(value)) return ok({ format: "v3-legacy", header: value });
|
||||
return err(new Error("Unsupported JSONL session header"));
|
||||
}
|
||||
3
third_party/pi/packages/agent/src/harness/session/jsonl/index.ts
vendored
Normal file
3
third_party/pi/packages/agent/src/harness/session/jsonl/index.ts
vendored
Normal file
@@ -0,0 +1,3 @@
|
||||
export * from "./repo.ts";
|
||||
export * from "./storage.ts";
|
||||
export * from "./types.ts";
|
||||
539
third_party/pi/packages/agent/src/harness/session/jsonl/legacy-v3.ts
vendored
Normal file
539
third_party/pi/packages/agent/src/harness/session/jsonl/legacy-v3.ts
vendored
Normal file
@@ -0,0 +1,539 @@
|
||||
import type { ImageContent, TextContent, Usage } from "@earendil-works/pi-ai";
|
||||
import { uuidv7 } from "@earendil-works/pi-ai/utils/uuid";
|
||||
import type { AgentMessage, ThinkingLevel } from "../../../types.ts";
|
||||
import type { Context } from "../../context.ts";
|
||||
import { createBranchSummaryMessage, createCompactionSummaryMessage } from "../../messages.ts";
|
||||
import type { FileSystem } from "../../types.ts";
|
||||
import { addUsage, emptyUsage } from "../../utils/usage.ts";
|
||||
import type { CommittedEntryWrite, CommittedValueSetWrite, CommittedWrite } from "../commit.ts";
|
||||
import type { JsonValue, LaneConfiguration } from "../types.ts";
|
||||
import { branchTip, entryLabel, laneConfig, laneState, sessionName } from "../values.ts";
|
||||
import { type LegacyV3SessionHeader, parseJsonlSessionHeader } from "./codec.ts";
|
||||
import {
|
||||
JSONL_FORMAT_VERSION,
|
||||
JSONL_STORAGE_VERSION,
|
||||
type JsonlSessionMetadata,
|
||||
type JsonlStorageHeader,
|
||||
} from "./types.ts";
|
||||
|
||||
interface LegacyV3EntryBase {
|
||||
type: string;
|
||||
id: string;
|
||||
parentId: string | null;
|
||||
timestamp: string;
|
||||
}
|
||||
|
||||
interface LegacyV3MessageEntry extends LegacyV3EntryBase {
|
||||
type: "message";
|
||||
message: AgentMessage;
|
||||
}
|
||||
|
||||
interface LegacyV3CustomEntry extends LegacyV3EntryBase {
|
||||
type: "custom";
|
||||
customType: string;
|
||||
data?: JsonValue;
|
||||
}
|
||||
|
||||
interface LegacyV3CustomMessageEntry extends LegacyV3EntryBase {
|
||||
type: "custom_message";
|
||||
customType: string;
|
||||
content: string | (TextContent | ImageContent)[];
|
||||
details?: unknown;
|
||||
display: boolean;
|
||||
}
|
||||
|
||||
interface LegacyV3BranchSummaryEntry extends LegacyV3EntryBase {
|
||||
type: "branch_summary";
|
||||
fromId: string;
|
||||
summary: string;
|
||||
details?: JsonValue;
|
||||
usage?: Usage;
|
||||
fromHook?: boolean;
|
||||
}
|
||||
|
||||
interface LegacyV3CompactionEntry extends LegacyV3EntryBase {
|
||||
type: "compaction";
|
||||
summary: string;
|
||||
firstKeptEntryId: string;
|
||||
tokensBefore: number;
|
||||
details?: JsonValue;
|
||||
usage?: Usage;
|
||||
fromHook?: boolean;
|
||||
}
|
||||
|
||||
interface LegacyV3ModelChangeEntry extends LegacyV3EntryBase {
|
||||
type: "model_change";
|
||||
provider: string;
|
||||
modelId: string;
|
||||
}
|
||||
|
||||
interface LegacyV3ThinkingLevelChangeEntry extends LegacyV3EntryBase {
|
||||
type: "thinking_level_change";
|
||||
thinkingLevel: string;
|
||||
}
|
||||
|
||||
interface LegacyV3ActiveToolsChangeEntry extends LegacyV3EntryBase {
|
||||
type: "active_tools_change";
|
||||
activeToolNames: string[];
|
||||
}
|
||||
|
||||
interface LegacyV3SessionInfoEntry extends LegacyV3EntryBase {
|
||||
type: "session_info";
|
||||
name?: string;
|
||||
}
|
||||
|
||||
interface LegacyV3LabelEntry extends LegacyV3EntryBase {
|
||||
type: "label";
|
||||
targetId: string;
|
||||
label?: string;
|
||||
}
|
||||
|
||||
interface ImportedCustomMessage {
|
||||
role: "custom";
|
||||
customType: string;
|
||||
content: string | (TextContent | ImageContent)[];
|
||||
details?: unknown;
|
||||
display: boolean;
|
||||
timestamp: number;
|
||||
}
|
||||
|
||||
type RetainedLegacyV3Entry =
|
||||
| LegacyV3MessageEntry
|
||||
| LegacyV3CustomEntry
|
||||
| LegacyV3CustomMessageEntry
|
||||
| LegacyV3BranchSummaryEntry
|
||||
| LegacyV3CompactionEntry;
|
||||
|
||||
type DiscardedLegacyV3Entry =
|
||||
| LegacyV3ModelChangeEntry
|
||||
| LegacyV3ThinkingLevelChangeEntry
|
||||
| LegacyV3ActiveToolsChangeEntry
|
||||
| LegacyV3SessionInfoEntry
|
||||
| LegacyV3LabelEntry;
|
||||
|
||||
type LegacyV3Entry = RetainedLegacyV3Entry | DiscardedLegacyV3Entry;
|
||||
|
||||
export interface NormalizedLegacyV3Records {
|
||||
writes: CommittedWrite[];
|
||||
importedUsage: Usage;
|
||||
nextSeq: number;
|
||||
}
|
||||
|
||||
type JsonlSessionMetadataBase = Omit<JsonlSessionMetadata, "path" | "modifiedAt">;
|
||||
|
||||
async function resolveLegacyV3ParentSessionId(
|
||||
fileSystem: FileSystem,
|
||||
parentSessionPath: string,
|
||||
context: Context,
|
||||
): Promise<string | undefined> {
|
||||
const lines = await fileSystem.readTextLines(parentSessionPath, { maxLines: 1 }, context);
|
||||
if (!lines.ok || lines.value[0] === undefined) return undefined;
|
||||
const parsed = parseJsonlSessionHeader(lines.value[0]);
|
||||
return parsed.ok ? parsed.value.header.id : undefined;
|
||||
}
|
||||
|
||||
export async function metadataFromLegacyV3Header(
|
||||
fileSystem: FileSystem,
|
||||
header: LegacyV3SessionHeader,
|
||||
context: Context,
|
||||
): Promise<JsonlSessionMetadataBase> {
|
||||
const metadata: JsonlSessionMetadataBase = {
|
||||
id: header.id,
|
||||
createdAt: Date.parse(header.timestamp),
|
||||
storageVersion: JSONL_STORAGE_VERSION,
|
||||
cwd: header.cwd,
|
||||
};
|
||||
if (header.parentSession !== undefined) {
|
||||
const parentSessionId = await resolveLegacyV3ParentSessionId(fileSystem, header.parentSession, context);
|
||||
if (parentSessionId !== undefined) metadata.parentSessionId = parentSessionId;
|
||||
else metadata.legacyParentSessionPath = header.parentSession;
|
||||
}
|
||||
return metadata;
|
||||
}
|
||||
|
||||
export async function normalizeLegacyV3Header(
|
||||
fileSystem: FileSystem,
|
||||
header: LegacyV3SessionHeader,
|
||||
context: Context,
|
||||
): Promise<JsonlStorageHeader> {
|
||||
return {
|
||||
v: JSONL_FORMAT_VERSION,
|
||||
kind: "header",
|
||||
...(await metadataFromLegacyV3Header(fileSystem, header, context)),
|
||||
};
|
||||
}
|
||||
|
||||
function parseLegacyV3Entry(line: string, lineNumber: number): LegacyV3Entry {
|
||||
let entry: LegacyV3Entry;
|
||||
try {
|
||||
entry = JSON.parse(line) as LegacyV3Entry;
|
||||
} catch (error) {
|
||||
throw new Error(`Invalid legacy v3 JSONL record at line ${lineNumber}: not valid JSON`, { cause: error });
|
||||
}
|
||||
const recordType: unknown = entry.type;
|
||||
if (
|
||||
recordType !== "message" &&
|
||||
recordType !== "custom" &&
|
||||
recordType !== "custom_message" &&
|
||||
recordType !== "branch_summary" &&
|
||||
recordType !== "compaction" &&
|
||||
recordType !== "model_change" &&
|
||||
recordType !== "thinking_level_change" &&
|
||||
recordType !== "active_tools_change" &&
|
||||
recordType !== "session_info" &&
|
||||
recordType !== "label"
|
||||
) {
|
||||
throw new Error(`Unsupported legacy v3 record type at line ${lineNumber}: ${String(recordType)}`);
|
||||
}
|
||||
return entry;
|
||||
}
|
||||
|
||||
function importedCustomMessage(entry: LegacyV3CustomMessageEntry): AgentMessage {
|
||||
const message: ImportedCustomMessage = {
|
||||
role: "custom",
|
||||
customType: entry.customType,
|
||||
content: entry.content,
|
||||
details: entry.details,
|
||||
display: entry.display,
|
||||
timestamp: Date.parse(entry.timestamp),
|
||||
};
|
||||
// The coding-agent CustomAgentMessages declaration merge is not visible in this package.
|
||||
return message as unknown as AgentMessage;
|
||||
}
|
||||
|
||||
function isRetainedEntry(entry: LegacyV3Entry): entry is RetainedLegacyV3Entry {
|
||||
return (
|
||||
entry.type !== "model_change" &&
|
||||
entry.type !== "thinking_level_change" &&
|
||||
entry.type !== "active_tools_change" &&
|
||||
entry.type !== "session_info" &&
|
||||
entry.type !== "label"
|
||||
);
|
||||
}
|
||||
|
||||
class RetainedIdResolver {
|
||||
/** Caches the reminted ID of each discarded legacy node's nearest retained ancestor, or null. */
|
||||
private resolvedIds = new Map<string, string | null>();
|
||||
private entriesById: Map<string, LegacyV3Entry>;
|
||||
private remintedIds: Map<string, string>;
|
||||
|
||||
/**
|
||||
* @param entriesById Complete inventory of legacy physical nodes, including discarded nodes.
|
||||
* @param remintedIds Legacy-to-current ID map containing retained nodes only.
|
||||
*/
|
||||
constructor(entriesById: Map<string, LegacyV3Entry>, remintedIds: Map<string, string>) {
|
||||
this.entriesById = entriesById;
|
||||
this.remintedIds = remintedIds;
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve a legacy node to its reminted ID, or to its nearest retained ancestor when discarded.
|
||||
* Returns null when the reference is null or no retained ancestor exists.
|
||||
*/
|
||||
resolve(legacyId: string | null): string | null {
|
||||
const traversedIds: string[] = [];
|
||||
const visitedIds = new Set<string>();
|
||||
let currentId = legacyId;
|
||||
let resolvedId: string | null = null;
|
||||
|
||||
while (currentId !== null) {
|
||||
const remintedId = this.remintedIds.get(currentId);
|
||||
if (remintedId !== undefined) {
|
||||
resolvedId = remintedId;
|
||||
break;
|
||||
}
|
||||
if (this.resolvedIds.has(currentId)) {
|
||||
resolvedId = this.resolvedIds.get(currentId) ?? null;
|
||||
break;
|
||||
}
|
||||
if (visitedIds.has(currentId)) throw new Error(`Cycle in legacy v3 parent chain at entry: ${currentId}`);
|
||||
visitedIds.add(currentId);
|
||||
const entry = this.entriesById.get(currentId);
|
||||
if (entry === undefined) throw new Error(`Missing legacy v3 entry reference: ${currentId}`);
|
||||
traversedIds.push(currentId);
|
||||
currentId = entry.parentId;
|
||||
}
|
||||
|
||||
for (const traversedId of traversedIds) this.resolvedIds.set(traversedId, resolvedId);
|
||||
return resolvedId;
|
||||
}
|
||||
}
|
||||
|
||||
function requireRetainedId(resolver: RetainedIdResolver, legacyId: string): string {
|
||||
const importedId = resolver.resolve(legacyId);
|
||||
if (importedId === null) throw new Error(`Legacy v3 entry reference has no retained ancestor: ${legacyId}`);
|
||||
return importedId;
|
||||
}
|
||||
|
||||
function resolveBranchSummaryFromId(resolver: RetainedIdResolver, legacyFromId: string): string | null {
|
||||
// Legacy branchWithSummary() encoded a root source as the "root" sentinel instead of null.
|
||||
return legacyFromId === "root" ? null : resolver.resolve(legacyFromId);
|
||||
}
|
||||
|
||||
function projectContextMessages(entry: LegacyV3Entry, resolver: RetainedIdResolver): AgentMessage[] {
|
||||
switch (entry.type) {
|
||||
case "message":
|
||||
return [entry.message];
|
||||
case "custom_message":
|
||||
return [importedCustomMessage(entry)];
|
||||
case "branch_summary":
|
||||
return entry.summary
|
||||
? [
|
||||
createBranchSummaryMessage(
|
||||
entry.summary,
|
||||
resolveBranchSummaryFromId(resolver, entry.fromId),
|
||||
entry.timestamp,
|
||||
),
|
||||
]
|
||||
: [];
|
||||
case "compaction":
|
||||
return [createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp)];
|
||||
case "custom":
|
||||
case "model_change":
|
||||
case "thinking_level_change":
|
||||
case "active_tools_change":
|
||||
case "session_info":
|
||||
case "label":
|
||||
return [];
|
||||
}
|
||||
}
|
||||
|
||||
function materializeRetainedTail(
|
||||
compaction: LegacyV3CompactionEntry,
|
||||
entriesById: ReadonlyMap<string, LegacyV3Entry>,
|
||||
resolver: RetainedIdResolver,
|
||||
): AgentMessage[] {
|
||||
const reversedTail: LegacyV3Entry[] = [];
|
||||
const visited = new Set<string>();
|
||||
let currentId = compaction.parentId;
|
||||
while (currentId !== null) {
|
||||
if (visited.has(currentId)) throw new Error(`Cycle in legacy v3 parent chain at entry: ${currentId}`);
|
||||
visited.add(currentId);
|
||||
const entry = entriesById.get(currentId);
|
||||
if (entry === undefined) throw new Error(`Missing legacy v3 parent entry: ${currentId}`);
|
||||
reversedTail.push(entry);
|
||||
if (currentId === compaction.firstKeptEntryId) {
|
||||
return reversedTail.reverse().flatMap((tailEntry) => projectContextMessages(tailEntry, resolver));
|
||||
}
|
||||
currentId = entry.parentId;
|
||||
}
|
||||
throw new Error(
|
||||
`Legacy v3 compaction ${compaction.id} firstKeptEntryId is not on its parent branch: ${compaction.firstKeptEntryId}`,
|
||||
);
|
||||
}
|
||||
|
||||
function normalizeRetainedEntry(
|
||||
entry: RetainedLegacyV3Entry,
|
||||
seq: number,
|
||||
entriesById: ReadonlyMap<string, LegacyV3Entry>,
|
||||
resolver: RetainedIdResolver,
|
||||
): CommittedEntryWrite {
|
||||
const committedBase = {
|
||||
kind: "entry" as const,
|
||||
id: requireRetainedId(resolver, entry.id),
|
||||
parentId: resolver.resolve(entry.parentId),
|
||||
seq,
|
||||
timestamp: Date.parse(entry.timestamp),
|
||||
};
|
||||
if (entry.type === "message") {
|
||||
return { ...committedBase, type: "message", message: entry.message };
|
||||
}
|
||||
if (entry.type === "custom_message") {
|
||||
return {
|
||||
...committedBase,
|
||||
type: "message",
|
||||
message: importedCustomMessage(entry),
|
||||
};
|
||||
}
|
||||
if (entry.type === "branch_summary") {
|
||||
return {
|
||||
...committedBase,
|
||||
type: "branch_summary",
|
||||
fromId: resolveBranchSummaryFromId(resolver, entry.fromId),
|
||||
summary: entry.summary,
|
||||
details: entry.details,
|
||||
usage: entry.usage,
|
||||
fromHook: entry.fromHook ?? false,
|
||||
};
|
||||
}
|
||||
if (entry.type === "compaction") {
|
||||
return {
|
||||
...committedBase,
|
||||
type: "compaction",
|
||||
summary: entry.summary,
|
||||
retainedTail: materializeRetainedTail(entry, entriesById, resolver),
|
||||
tokensBefore: entry.tokensBefore,
|
||||
details: entry.details,
|
||||
usage: entry.usage,
|
||||
fromHook: entry.fromHook ?? false,
|
||||
};
|
||||
}
|
||||
return {
|
||||
...committedBase,
|
||||
type: "custom",
|
||||
customType: entry.customType,
|
||||
data: entry.data,
|
||||
};
|
||||
}
|
||||
|
||||
const THINKING_LEVELS = new Set<ThinkingLevel>(["off", "minimal", "low", "medium", "high", "xhigh", "max"]);
|
||||
|
||||
function selectedConfiguration(
|
||||
entriesById: ReadonlyMap<string, LegacyV3Entry>,
|
||||
selectedId: string | null,
|
||||
): LaneConfiguration | undefined {
|
||||
let model: LaneConfiguration["model"] | undefined;
|
||||
let thinkingLevel: ThinkingLevel | undefined;
|
||||
let activeToolNames: string[] | undefined;
|
||||
let sawModel = false;
|
||||
let sawThinkingLevel = false;
|
||||
let sawActiveToolNames = false;
|
||||
const visited = new Set<string>();
|
||||
let currentId = selectedId;
|
||||
while (currentId !== null && (!sawModel || !sawThinkingLevel || !sawActiveToolNames)) {
|
||||
if (visited.has(currentId)) throw new Error(`Cycle in legacy v3 parent chain at entry: ${currentId}`);
|
||||
visited.add(currentId);
|
||||
const entry = entriesById.get(currentId);
|
||||
if (entry === undefined) throw new Error(`Missing legacy v3 entry reference: ${currentId}`);
|
||||
if (entry.type === "model_change" && !sawModel) {
|
||||
sawModel = true;
|
||||
if (
|
||||
typeof entry.provider === "string" &&
|
||||
entry.provider.length !== 0 &&
|
||||
typeof entry.modelId === "string" &&
|
||||
entry.modelId.length !== 0
|
||||
) {
|
||||
model = { provider: entry.provider, modelId: entry.modelId };
|
||||
}
|
||||
} else if (entry.type === "thinking_level_change" && !sawThinkingLevel) {
|
||||
sawThinkingLevel = true;
|
||||
if (typeof entry.thinkingLevel === "string" && THINKING_LEVELS.has(entry.thinkingLevel as ThinkingLevel)) {
|
||||
thinkingLevel = entry.thinkingLevel as ThinkingLevel;
|
||||
}
|
||||
} else if (entry.type === "active_tools_change" && !sawActiveToolNames) {
|
||||
sawActiveToolNames = true;
|
||||
if (Array.isArray(entry.activeToolNames) && entry.activeToolNames.every((name) => typeof name === "string")) {
|
||||
activeToolNames = [...entry.activeToolNames];
|
||||
}
|
||||
}
|
||||
currentId = entry.parentId;
|
||||
}
|
||||
if (model === undefined || thinkingLevel === undefined) return undefined;
|
||||
return { model, thinkingLevel, activeToolNames: activeToolNames ?? [] };
|
||||
}
|
||||
|
||||
function normalizeLegacyV3Values(
|
||||
entries: readonly LegacyV3Entry[],
|
||||
entriesById: ReadonlyMap<string, LegacyV3Entry>,
|
||||
resolver: RetainedIdResolver,
|
||||
firstSeq: number,
|
||||
): CommittedValueSetWrite[] {
|
||||
const writes: CommittedValueSetWrite[] = [];
|
||||
let latestSessionInfo: LegacyV3SessionInfoEntry | undefined;
|
||||
for (const entry of entries) {
|
||||
if (entry.type === "session_info") latestSessionInfo = entry;
|
||||
}
|
||||
if (latestSessionInfo?.name) {
|
||||
writes.push({
|
||||
kind: "value",
|
||||
op: "set",
|
||||
seq: firstSeq + writes.length,
|
||||
namespace: sessionName.namespace,
|
||||
key: sessionName.key,
|
||||
value: latestSessionInfo.name,
|
||||
});
|
||||
}
|
||||
|
||||
const labels = new Map<string, string>();
|
||||
for (const entry of entries) {
|
||||
if (entry.type !== "label") continue;
|
||||
const targetId = resolver.resolve(entry.targetId);
|
||||
// Labels have no current address when their target has no retained ancestor.
|
||||
if (targetId === null) continue;
|
||||
// Legacy v3 treated both undefined and the empty string as clearing a label.
|
||||
if (entry.label) labels.set(targetId, entry.label);
|
||||
else labels.delete(targetId);
|
||||
}
|
||||
for (const [targetId, label] of labels) {
|
||||
const address = entryLabel(targetId);
|
||||
writes.push({
|
||||
kind: "value",
|
||||
op: "set",
|
||||
seq: firstSeq + writes.length,
|
||||
namespace: address.namespace,
|
||||
key: address.key,
|
||||
value: label,
|
||||
});
|
||||
}
|
||||
|
||||
const finalEntry = entries.at(-1);
|
||||
const tipAddress = branchTip("main");
|
||||
writes.push({
|
||||
kind: "value",
|
||||
op: "set",
|
||||
seq: firstSeq + writes.length,
|
||||
namespace: tipAddress.namespace,
|
||||
key: tipAddress.key,
|
||||
value: finalEntry === undefined ? null : resolver.resolve(finalEntry.id),
|
||||
});
|
||||
const configuration = selectedConfiguration(entriesById, finalEntry?.id ?? null);
|
||||
if (configuration !== undefined) {
|
||||
const configAddress = laneConfig("main");
|
||||
writes.push({
|
||||
kind: "value",
|
||||
op: "set",
|
||||
seq: firstSeq + writes.length,
|
||||
namespace: configAddress.namespace,
|
||||
key: configAddress.key,
|
||||
value: configuration,
|
||||
});
|
||||
const stateAddress = laneState("main");
|
||||
writes.push({
|
||||
kind: "value",
|
||||
op: "set",
|
||||
seq: firstSeq + writes.length,
|
||||
namespace: stateAddress.namespace,
|
||||
key: stateAddress.key,
|
||||
value: { currentOperationId: null, lastOperationId: null, inbox: [] },
|
||||
});
|
||||
}
|
||||
return writes;
|
||||
}
|
||||
|
||||
function aggregateImportedUsage(entries: readonly LegacyV3Entry[]): Usage {
|
||||
let aggregate = emptyUsage();
|
||||
for (const entry of entries) {
|
||||
let usage: Usage | undefined;
|
||||
if (entry.type === "message") {
|
||||
if (entry.message.role === "assistant") usage = entry.message.usage;
|
||||
else if (entry.message.role === "toolResult") usage = entry.message.usage;
|
||||
} else if (entry.type === "compaction" || entry.type === "branch_summary") {
|
||||
usage = entry.usage;
|
||||
}
|
||||
if (usage !== undefined) aggregate = addUsage(aggregate, usage);
|
||||
}
|
||||
return aggregate;
|
||||
}
|
||||
|
||||
/** Normalize the currently supported v3 records without touching their source file. */
|
||||
export function normalizeLegacyV3Records(recordLines: readonly string[]): NormalizedLegacyV3Records {
|
||||
const entries = recordLines.map((line, index) => parseLegacyV3Entry(line, index + 2));
|
||||
const entriesById = new Map<string, LegacyV3Entry>();
|
||||
for (const entry of entries) {
|
||||
if (entriesById.has(entry.id)) throw new Error(`Duplicate legacy v3 entry id: ${entry.id}`);
|
||||
entriesById.set(entry.id, entry);
|
||||
}
|
||||
const retainedEntries = entries.filter(isRetainedEntry);
|
||||
const remintedIds = new Map(retainedEntries.map((entry) => [entry.id, uuidv7(Date.parse(entry.timestamp))]));
|
||||
const resolver = new RetainedIdResolver(entriesById, remintedIds);
|
||||
const entryWrites = retainedEntries.map((entry, index) =>
|
||||
normalizeRetainedEntry(entry, index + 1, entriesById, resolver),
|
||||
);
|
||||
const valueWrites = normalizeLegacyV3Values(entries, entriesById, resolver, entryWrites.length + 1);
|
||||
const writes: CommittedWrite[] = [...entryWrites, ...valueWrites];
|
||||
return {
|
||||
writes,
|
||||
importedUsage: aggregateImportedUsage(entries),
|
||||
nextSeq: writes.length + 1,
|
||||
};
|
||||
}
|
||||
368
third_party/pi/packages/agent/src/harness/session/jsonl/repo.ts
vendored
Normal file
368
third_party/pi/packages/agent/src/harness/session/jsonl/repo.ts
vendored
Normal file
@@ -0,0 +1,368 @@
|
||||
import { uuidv7 } from "@earendil-works/pi-ai/utils/uuid";
|
||||
import type { Context } from "../../context.ts";
|
||||
import type { FileError, FileInfo, FileSystem, Result } from "../../types.ts";
|
||||
import { createForkSnapshot, type ForkSourceSnapshot } from "../fork.ts";
|
||||
import { StorageBackedSession } from "../session.ts";
|
||||
import type { ForkOptions, Session, SessionRepo } from "../types.ts";
|
||||
import { parseJsonlSessionHeader } from "./codec.ts";
|
||||
import { metadataFromLegacyV3Header } from "./legacy-v3.ts";
|
||||
import { JsonlStorage } from "./storage.ts";
|
||||
import {
|
||||
JSONL_FORMAT_VERSION,
|
||||
JSONL_STORAGE_VERSION,
|
||||
type JsonlSessionCreateOptions,
|
||||
type JsonlSessionListOptions,
|
||||
type JsonlSessionMetadata,
|
||||
type JsonlSessionRepoOptions,
|
||||
type JsonlStorageHeader,
|
||||
} from "./types.ts";
|
||||
|
||||
function fileValue<T>(result: Result<T, FileError>, action: string): T {
|
||||
if (!result.ok) throw new Error(`${action}: ${result.error.message}`, { cause: result.error });
|
||||
return result.value;
|
||||
}
|
||||
|
||||
function metadataFromHeader(header: JsonlStorageHeader, path: string, modifiedAt: number): JsonlSessionMetadata {
|
||||
return {
|
||||
id: header.id,
|
||||
createdAt: header.createdAt,
|
||||
storageVersion: header.storageVersion,
|
||||
cwd: header.cwd,
|
||||
path,
|
||||
modifiedAt,
|
||||
...(header.parentSessionId === undefined ? {} : { parentSessionId: header.parentSessionId }),
|
||||
...(header.legacyParentSessionPath === undefined
|
||||
? {}
|
||||
: { legacyParentSessionPath: header.legacyParentSessionPath }),
|
||||
};
|
||||
}
|
||||
|
||||
function sessionDirectoryName(cwd: string): string {
|
||||
return `--${cwd.replace(/^[/\\]/, "").replace(/[/\\:]/g, "-")}--`;
|
||||
}
|
||||
|
||||
function sessionFileName(createdAt: number, id: string): string {
|
||||
const timestamp = new Date(createdAt).toISOString().replace(/[:.]/g, "-");
|
||||
return `${timestamp}_${encodeURIComponent(id)}.jsonl`;
|
||||
}
|
||||
|
||||
/** File-backed format-4 session repository lifecycle. */
|
||||
export class JsonlSessionRepo
|
||||
implements SessionRepo<JsonlSessionMetadata, JsonlSessionCreateOptions, JsonlSessionListOptions>
|
||||
{
|
||||
private readonly fileSystem: FileSystem;
|
||||
private readonly sessionsRootInput: string;
|
||||
private readonly now: () => number;
|
||||
private readonly openSessions = new Map<string, JsonlStorage>();
|
||||
private readonly pendingCreates = new Set<string>();
|
||||
private closed = false;
|
||||
private closePromise: Promise<void> | undefined;
|
||||
|
||||
constructor(options: JsonlSessionRepoOptions) {
|
||||
this.fileSystem = options.fileSystem;
|
||||
this.sessionsRootInput = options.sessionsRoot;
|
||||
this.now = options.now ?? Date.now;
|
||||
}
|
||||
|
||||
async create(options: JsonlSessionCreateOptions, context: Context): Promise<Session<JsonlSessionMetadata>> {
|
||||
this.assertOpen();
|
||||
const createdAt = this.now();
|
||||
const { cwd, id } = await this.resolveCreateDestination(options.cwd, options.id, createdAt, context);
|
||||
const key = this.sessionKey(cwd, id);
|
||||
if (this.openSessions.has(key) || this.pendingCreates.has(key)) throw new Error(`Session already exists: ${id}`);
|
||||
this.pendingCreates.add(key);
|
||||
let path: string | undefined;
|
||||
let storage: JsonlStorage | undefined;
|
||||
try {
|
||||
path = await this.resolveNewSessionPath(cwd, createdAt, id, context);
|
||||
const header: JsonlStorageHeader = {
|
||||
v: JSONL_FORMAT_VERSION,
|
||||
kind: "header",
|
||||
id,
|
||||
storageVersion: JSONL_STORAGE_VERSION,
|
||||
createdAt,
|
||||
cwd,
|
||||
...(options.parentSessionId === undefined ? {} : { parentSessionId: options.parentSessionId }),
|
||||
};
|
||||
storage = await JsonlStorage.create({ fileSystem: this.fileSystem, path, now: this.now }, header, [], context);
|
||||
const info = fileValue(await this.fileSystem.fileInfo(path, context), `Failed to read session ${path}`);
|
||||
return this.publishOpenSession(metadataFromHeader(header, path, info.mtimeMs), storage, key);
|
||||
} catch (error) {
|
||||
await storage?.close(context).catch(() => undefined);
|
||||
if (path !== undefined) await this.fileSystem.remove(path, { force: true }, context);
|
||||
throw error;
|
||||
} finally {
|
||||
this.pendingCreates.delete(key);
|
||||
}
|
||||
}
|
||||
|
||||
async open(metadata: JsonlSessionMetadata, context: Context): Promise<Session<JsonlSessionMetadata>> {
|
||||
this.assertOpen();
|
||||
const key = this.sessionKey(metadata.cwd, metadata.id);
|
||||
if (this.openSessions.has(key)) throw new Error(`Session is already open: ${metadata.id}`);
|
||||
let storage: JsonlStorage | undefined;
|
||||
try {
|
||||
storage = await this.loadStorage(metadata, context);
|
||||
return this.publishOpenSession(metadata, storage, key);
|
||||
} catch (error) {
|
||||
await storage?.close(context);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
async list(options: JsonlSessionListOptions | undefined, context: Context): Promise<JsonlSessionMetadata[]> {
|
||||
options ??= {};
|
||||
this.assertOpen();
|
||||
const cwd =
|
||||
options.cwd === undefined
|
||||
? undefined
|
||||
: fileValue(
|
||||
await this.fileSystem.absolutePath(options.cwd, context),
|
||||
`Failed to resolve session cwd ${options.cwd}`,
|
||||
);
|
||||
const root = await this.root(context);
|
||||
if (!fileValue(await this.fileSystem.exists(root, context), `Failed to check sessions root ${root}`)) return [];
|
||||
const directories =
|
||||
cwd === undefined ? await this.sessionDirectories(root, context) : [await this.sessionDirectory(cwd, context)];
|
||||
const metadata: JsonlSessionMetadata[] = [];
|
||||
for (const directory of directories) metadata.push(...(await this.listDirectory(directory, cwd, context)));
|
||||
return metadata.sort(
|
||||
(left, right) =>
|
||||
right.createdAt - left.createdAt || left.id.localeCompare(right.id) || left.cwd.localeCompare(right.cwd),
|
||||
);
|
||||
}
|
||||
|
||||
async delete(metadata: JsonlSessionMetadata, context: Context): Promise<void> {
|
||||
this.assertOpen();
|
||||
const key = this.sessionKey(metadata.cwd, metadata.id);
|
||||
if (this.openSessions.has(key)) throw new Error(`Session is open: ${metadata.id}`);
|
||||
if (
|
||||
!fileValue(await this.fileSystem.exists(metadata.path, context), `Failed to check session ${metadata.path}`)
|
||||
) {
|
||||
throw new Error(`Session file does not exist: ${metadata.path}`);
|
||||
}
|
||||
fileValue(
|
||||
await this.fileSystem.remove(metadata.path, undefined, context),
|
||||
`Failed to delete session ${metadata.path}`,
|
||||
);
|
||||
}
|
||||
|
||||
async fork(
|
||||
source: JsonlSessionMetadata,
|
||||
options: ForkOptions,
|
||||
context: Context,
|
||||
): Promise<Session<JsonlSessionMetadata>> {
|
||||
this.assertOpen();
|
||||
const createdAt = this.now();
|
||||
const sourceStorage = this.openSessions.get(this.sessionKey(source.cwd, source.id));
|
||||
const sourceSnapshot = await (sourceStorage === undefined
|
||||
? this.loadClosedForkSourceSnapshot(source, context)
|
||||
: sourceStorage.captureForkSource(context));
|
||||
const { cwd, id } = await this.resolveCreateDestination(source.cwd, options.id, createdAt, context);
|
||||
const destinationKey = this.sessionKey(cwd, id);
|
||||
if (this.openSessions.has(destinationKey) || this.pendingCreates.has(destinationKey)) {
|
||||
throw new Error(`Session already exists: ${id}`);
|
||||
}
|
||||
this.pendingCreates.add(destinationKey);
|
||||
|
||||
let path: string | undefined;
|
||||
let storage: JsonlStorage | undefined;
|
||||
try {
|
||||
path = await this.resolveNewSessionPath(cwd, createdAt, id, context);
|
||||
const snapshot = createForkSnapshot(sourceSnapshot, options);
|
||||
const header: JsonlStorageHeader = {
|
||||
v: JSONL_FORMAT_VERSION,
|
||||
kind: "header",
|
||||
id,
|
||||
storageVersion: JSONL_STORAGE_VERSION,
|
||||
createdAt,
|
||||
cwd,
|
||||
parentSessionId: source.id,
|
||||
};
|
||||
storage = await JsonlStorage.createFromForkSnapshot(
|
||||
{ fileSystem: this.fileSystem, path, now: this.now },
|
||||
header,
|
||||
snapshot,
|
||||
context,
|
||||
);
|
||||
const info = fileValue(await this.fileSystem.fileInfo(path, context), `Failed to read session ${path}`);
|
||||
return this.publishOpenSession(metadataFromHeader(header, path, info.mtimeMs), storage, destinationKey);
|
||||
} catch (error) {
|
||||
await storage?.close(context).catch(() => undefined);
|
||||
if (path !== undefined) await this.fileSystem.remove(path, { force: true }, context);
|
||||
throw error;
|
||||
} finally {
|
||||
this.pendingCreates.delete(destinationKey);
|
||||
}
|
||||
}
|
||||
|
||||
close(_context: Context): Promise<void> {
|
||||
if (this.closePromise !== undefined) return this.closePromise;
|
||||
this.closed = true;
|
||||
// TODO: Define ownership semantics before deciding whether repository close should close session handles.
|
||||
this.closePromise = Promise.resolve();
|
||||
return this.closePromise;
|
||||
}
|
||||
|
||||
private async resolveCreateDestination(
|
||||
cwdInput: string,
|
||||
id: string | undefined,
|
||||
createdAt: number,
|
||||
context: Context,
|
||||
): Promise<{ cwd: string; id: string }> {
|
||||
const destinationId = id ?? uuidv7(createdAt);
|
||||
const cwd = fileValue(
|
||||
await this.fileSystem.absolutePath(cwdInput, context),
|
||||
`Failed to resolve session cwd ${cwdInput}`,
|
||||
);
|
||||
return { cwd, id: destinationId };
|
||||
}
|
||||
|
||||
private async listDirectory(
|
||||
directory: string,
|
||||
cwd: string | undefined,
|
||||
context: Context,
|
||||
): Promise<JsonlSessionMetadata[]> {
|
||||
if (
|
||||
!fileValue(await this.fileSystem.exists(directory, context), `Failed to check sessions directory ${directory}`)
|
||||
) {
|
||||
return [];
|
||||
}
|
||||
const files = fileValue(
|
||||
await this.fileSystem.listDir(directory, context),
|
||||
`Failed to list sessions directory ${directory}`,
|
||||
).filter((file) => file.kind !== "directory" && file.name.endsWith(".jsonl"));
|
||||
const metadata: JsonlSessionMetadata[] = [];
|
||||
for (const file of files) {
|
||||
const discovered = await this.readSessionMetadata(file, context);
|
||||
if (discovered === undefined) continue;
|
||||
// Directory encoding is lossy: /a/b and /a-b both map to --a-b--.
|
||||
if (cwd === undefined || discovered.cwd === cwd) metadata.push(discovered);
|
||||
}
|
||||
return metadata;
|
||||
}
|
||||
|
||||
private async readSessionMetadata(file: FileInfo, context: Context): Promise<JsonlSessionMetadata | undefined> {
|
||||
const lines = fileValue(
|
||||
await this.fileSystem.readTextLines(file.path, { maxLines: 1 }, context),
|
||||
`Failed to read session header ${file.path}`,
|
||||
);
|
||||
if (lines[0] === undefined) return undefined;
|
||||
const parsedHeader = parseJsonlSessionHeader(lines[0]);
|
||||
if (!parsedHeader.ok) return undefined;
|
||||
if (parsedHeader.value.format === "v3-legacy") {
|
||||
const metadata = await metadataFromLegacyV3Header(this.fileSystem, parsedHeader.value.header, context);
|
||||
return { ...metadata, path: file.path, modifiedAt: file.mtimeMs };
|
||||
}
|
||||
return metadataFromHeader(parsedHeader.value.header, file.path, file.mtimeMs);
|
||||
}
|
||||
|
||||
private async sessionDirectories(root: string, context: Context): Promise<string[]> {
|
||||
return fileValue(await this.fileSystem.listDir(root, context), `Failed to list sessions root ${root}`)
|
||||
.filter((entry) => entry.kind === "directory")
|
||||
.map((entry) => entry.path);
|
||||
}
|
||||
|
||||
private async sessionDirectory(cwd: string, context: Context): Promise<string> {
|
||||
return fileValue(
|
||||
await this.fileSystem.joinPath([await this.root(context), sessionDirectoryName(cwd)], context),
|
||||
`Failed to resolve sessions directory for ${cwd}`,
|
||||
);
|
||||
}
|
||||
|
||||
private async resolveNewSessionPath(cwd: string, createdAt: number, id: string, context: Context): Promise<string> {
|
||||
const directory = await this.sessionDirectory(cwd, context);
|
||||
await this.assertSessionIdAvailable(directory, id, context);
|
||||
fileValue(
|
||||
await this.fileSystem.createDir(directory, undefined, context),
|
||||
`Failed to create sessions directory ${directory}`,
|
||||
);
|
||||
return fileValue(
|
||||
await this.fileSystem.joinPath([directory, sessionFileName(createdAt, id)], context),
|
||||
`Failed to resolve path for session ${id}`,
|
||||
);
|
||||
}
|
||||
|
||||
private async assertSessionIdAvailable(directory: string, id: string, context: Context): Promise<void> {
|
||||
if (
|
||||
!fileValue(await this.fileSystem.exists(directory, context), `Failed to check sessions directory ${directory}`)
|
||||
)
|
||||
return;
|
||||
|
||||
const suffix = `_${encodeURIComponent(id)}.jsonl`;
|
||||
const idExists = fileValue(
|
||||
await this.fileSystem.listDir(directory, context),
|
||||
`Failed to list sessions directory ${directory}`,
|
||||
).some((entry) => entry.kind !== "directory" && entry.name.endsWith(suffix));
|
||||
if (idExists) throw new Error(`Session already exists: ${id}`);
|
||||
}
|
||||
|
||||
private async loadClosedForkSourceSnapshot(
|
||||
source: JsonlSessionMetadata,
|
||||
context: Context,
|
||||
): Promise<ForkSourceSnapshot> {
|
||||
const storage = await this.loadStorage(source, context);
|
||||
try {
|
||||
return await storage.captureForkSource(context);
|
||||
} finally {
|
||||
await storage.close(context);
|
||||
}
|
||||
}
|
||||
|
||||
private publishOpenSession(
|
||||
metadata: JsonlSessionMetadata,
|
||||
storage: JsonlStorage,
|
||||
key: string,
|
||||
): StorageBackedSession<JsonlSessionMetadata> {
|
||||
if (this.openSessions.has(key)) throw new Error(`Session is already open: ${metadata.id}`);
|
||||
const session = new StorageBackedSession(metadata, storage, {
|
||||
onClose: () => {
|
||||
if (this.openSessions.get(key) === storage) this.openSessions.delete(key);
|
||||
},
|
||||
});
|
||||
this.openSessions.set(key, storage);
|
||||
return session;
|
||||
}
|
||||
|
||||
private sessionKey(cwd: string, id: string): string {
|
||||
return `${cwd}\0${id}`;
|
||||
}
|
||||
|
||||
private async loadStorage(metadata: JsonlSessionMetadata, context: Context): Promise<JsonlStorage> {
|
||||
if (
|
||||
!fileValue(await this.fileSystem.exists(metadata.path, context), `Failed to check session ${metadata.path}`)
|
||||
) {
|
||||
throw new Error(`Session file does not exist: ${metadata.path}`);
|
||||
}
|
||||
const storage = await JsonlStorage.open(
|
||||
{
|
||||
fileSystem: this.fileSystem,
|
||||
path: metadata.path,
|
||||
now: this.now,
|
||||
},
|
||||
context,
|
||||
);
|
||||
try {
|
||||
if (storage.header.id !== metadata.id || storage.header.cwd !== metadata.cwd) {
|
||||
throw new Error(`Session identity does not match header: ${metadata.id}`);
|
||||
}
|
||||
if (storage.header.storageVersion !== JSONL_STORAGE_VERSION) {
|
||||
throw new Error(`Session ${metadata.id} uses unsupported storage version ${storage.header.storageVersion}`);
|
||||
}
|
||||
return storage;
|
||||
} catch (error) {
|
||||
await storage.close(context);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
private root(context: Context): Promise<string> {
|
||||
return this.fileSystem
|
||||
.absolutePath(this.sessionsRootInput, context)
|
||||
.then((result) => fileValue(result, `Failed to resolve sessions root ${this.sessionsRootInput}`));
|
||||
}
|
||||
|
||||
private assertOpen(): void {
|
||||
if (this.closed) throw new Error("JsonlSessionRepo is closed");
|
||||
}
|
||||
}
|
||||
376
third_party/pi/packages/agent/src/harness/session/jsonl/storage.ts
vendored
Normal file
376
third_party/pi/packages/agent/src/harness/session/jsonl/storage.ts
vendored
Normal file
@@ -0,0 +1,376 @@
|
||||
import type { Usage } from "@earendil-works/pi-ai";
|
||||
import { uuidv7 } from "@earendil-works/pi-ai/utils/uuid";
|
||||
import type { Context } from "../../context.ts";
|
||||
import type { FileError, FileSystem, Result } from "../../types.ts";
|
||||
import { insertUsage } from "../commit.ts";
|
||||
import { type ForkDestinationSnapshot, type ForkSourceSnapshot, forkSnapshotWrites } from "../fork.ts";
|
||||
import {
|
||||
type CommittedEntryWrite,
|
||||
type CommittedListAppendWrite,
|
||||
type CommittedListDeleteWrite,
|
||||
type CommittedUsageWrite,
|
||||
type CommittedValueDeleteWrite,
|
||||
type CommittedValueSetWrite,
|
||||
type CommittedWrite,
|
||||
InMemoryStorageState,
|
||||
} from "../in-memory-storage-state.ts";
|
||||
import type {
|
||||
CommitResult,
|
||||
Entry,
|
||||
EntryScan,
|
||||
EntryStructure,
|
||||
SessionStats,
|
||||
Storage,
|
||||
StorageBranchScan,
|
||||
UsageRow,
|
||||
UsageScan,
|
||||
Write,
|
||||
} from "../types.ts";
|
||||
import type { ListElement, ListReadOptions, StoredValue, Value, ValueList } from "../values.ts";
|
||||
import { type LegacyV3SessionHeader, parseJsonlSessionHeader } from "./codec.ts";
|
||||
import { normalizeLegacyV3Header, normalizeLegacyV3Records } from "./legacy-v3.ts";
|
||||
import { JSONL_STORAGE_VERSION, type JsonlStorageHeader, type JsonlStorageOptions } from "./types.ts";
|
||||
|
||||
function fileValue<T>(result: Result<T, FileError>, action: string): T {
|
||||
if (!result.ok) throw new Error(`${action}: ${result.error.message}`, { cause: result.error });
|
||||
return result.value;
|
||||
}
|
||||
|
||||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||
return typeof value === "object" && value !== null && !Array.isArray(value);
|
||||
}
|
||||
|
||||
function requireSafeInteger(value: unknown, field: string, minimum: number): void {
|
||||
if (!Number.isSafeInteger(value) || (value as number) < minimum) throw new Error(`Invalid JSONL ${field}`);
|
||||
}
|
||||
|
||||
function parseCommittedWrite(value: unknown): CommittedWrite {
|
||||
if (!isRecord(value)) throw new Error("Invalid JSONL transaction write");
|
||||
requireSafeInteger(value.seq, "write seq", 1);
|
||||
switch (value.kind) {
|
||||
case "entry":
|
||||
requireSafeInteger(value.timestamp, "entry timestamp", 0);
|
||||
return value as unknown as CommittedEntryWrite;
|
||||
case "usage":
|
||||
return value as unknown as CommittedUsageWrite;
|
||||
case "value":
|
||||
if (value.op === "set") return value as unknown as CommittedValueSetWrite;
|
||||
if (value.op === "delete") return value as unknown as CommittedValueDeleteWrite;
|
||||
throw new Error(`Invalid JSONL value operation: ${String(value.op)}`);
|
||||
case "list":
|
||||
if (value.op === "append") return value as unknown as CommittedListAppendWrite;
|
||||
if (value.op === "delete") return value as unknown as CommittedListDeleteWrite;
|
||||
throw new Error(`Invalid JSONL list operation: ${String(value.op)}`);
|
||||
default:
|
||||
throw new Error(`Invalid JSONL write kind: ${String(value.kind)}`);
|
||||
}
|
||||
}
|
||||
|
||||
function parseTransaction(line: string): CommittedWrite[] {
|
||||
let value: unknown;
|
||||
try {
|
||||
value = JSON.parse(line);
|
||||
} catch (error) {
|
||||
throw new Error("Invalid JSONL transaction: not valid JSON", { cause: error });
|
||||
}
|
||||
return (Array.isArray(value) ? value : [value]).map(parseCommittedWrite);
|
||||
}
|
||||
|
||||
function serializeTransaction(writes: CommittedWrite[]): string {
|
||||
return JSON.stringify(writes.length === 1 ? writes[0] : writes);
|
||||
}
|
||||
|
||||
function serializeStorage(header: JsonlStorageHeader, transactions: CommittedWrite[][]): string {
|
||||
return `${[JSON.stringify(header), ...transactions.map(serializeTransaction)].join("\n")}\n`;
|
||||
}
|
||||
|
||||
function splitCompleteLines(content: string): { lines: string[]; torn: boolean } {
|
||||
if (content.endsWith("\n")) return { lines: content.slice(0, -1).split("\n"), torn: false };
|
||||
const lastNewline = content.lastIndexOf("\n");
|
||||
if (lastNewline === -1) return { lines: [], torn: true };
|
||||
return { lines: content.slice(0, lastNewline).split("\n"), torn: true };
|
||||
}
|
||||
|
||||
async function publishFileAtomically(
|
||||
fileSystem: FileSystem,
|
||||
destinationPath: string,
|
||||
content: string,
|
||||
context: Context,
|
||||
): Promise<void> {
|
||||
const tempPath = `${destinationPath}.tmp`;
|
||||
try {
|
||||
fileValue(
|
||||
await fileSystem.writeFile(tempPath, content, context),
|
||||
`Failed to stage JSONL storage ${destinationPath}`,
|
||||
);
|
||||
fileValue(
|
||||
await fileSystem.renameFile(tempPath, destinationPath, context),
|
||||
`Failed to publish JSONL storage ${destinationPath}`,
|
||||
);
|
||||
} catch (error) {
|
||||
await fileSystem.remove(tempPath, { force: true }, context);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
type LegacyV3Backing = {
|
||||
kind: "v3";
|
||||
importedUsage: Usage;
|
||||
baselineWrites: readonly CommittedWrite[];
|
||||
};
|
||||
|
||||
type JsonlBacking = { kind: "v4" } | LegacyV3Backing;
|
||||
|
||||
/** JSONL storage backed by an injected filesystem capability. */
|
||||
export class JsonlStorage implements Storage {
|
||||
private readonly fileSystem: FileSystem;
|
||||
private readonly path: string;
|
||||
private readonly now: () => number;
|
||||
readonly header: JsonlStorageHeader;
|
||||
private backing: JsonlBacking;
|
||||
private readonly storageState = new InMemoryStorageState();
|
||||
private commitQueue: Promise<void> = Promise.resolve();
|
||||
private state: "open" | "closing" | "closed" = "open";
|
||||
private closePromise: Promise<void> | undefined;
|
||||
|
||||
private constructor(options: JsonlStorageOptions, header: JsonlStorageHeader, backing: JsonlBacking) {
|
||||
this.fileSystem = options.fileSystem;
|
||||
this.path = options.path;
|
||||
this.now = options.now ?? Date.now;
|
||||
this.header = header;
|
||||
this.backing = backing;
|
||||
}
|
||||
|
||||
static async create(
|
||||
options: JsonlStorageOptions,
|
||||
header: JsonlStorageHeader,
|
||||
initialWrites: Write[],
|
||||
context: Context,
|
||||
): Promise<JsonlStorage> {
|
||||
const storage = new JsonlStorage(options, header, { kind: "v4" });
|
||||
const prepared = storage.storageState.prepareCommit(initialWrites, storage.now());
|
||||
const transactions = prepared.writes.length === 0 ? [] : [prepared.writes];
|
||||
await publishFileAtomically(options.fileSystem, options.path, serializeStorage(header, transactions), context);
|
||||
storage.storageState.applyValidated(prepared.writes);
|
||||
return storage;
|
||||
}
|
||||
|
||||
/** Atomically create storage from a complete prepared snapshot. */
|
||||
static async createFromForkSnapshot(
|
||||
options: JsonlStorageOptions,
|
||||
header: JsonlStorageHeader,
|
||||
snapshot: ForkDestinationSnapshot,
|
||||
context: Context,
|
||||
): Promise<JsonlStorage> {
|
||||
const writes = forkSnapshotWrites(snapshot);
|
||||
const snapshotHeader = { ...header, nextSeq: snapshot.nextSeq };
|
||||
await publishFileAtomically(
|
||||
options.fileSystem,
|
||||
options.path,
|
||||
serializeStorage(
|
||||
snapshotHeader,
|
||||
writes.map((write) => [write]),
|
||||
),
|
||||
context,
|
||||
);
|
||||
return JsonlStorage.open(options, context);
|
||||
}
|
||||
|
||||
static async open(options: JsonlStorageOptions, context: Context): Promise<JsonlStorage> {
|
||||
const content = fileValue(
|
||||
await options.fileSystem.readTextFile(options.path, context),
|
||||
`Failed to read JSONL storage ${options.path}`,
|
||||
);
|
||||
const { lines, torn } = splitCompleteLines(content);
|
||||
if (lines[0] === undefined || lines[0] === "") {
|
||||
throw new Error(`Invalid JSONL storage ${options.path}: missing header`);
|
||||
}
|
||||
const parsedHeader = parseJsonlSessionHeader(lines[0]);
|
||||
if (!parsedHeader.ok) {
|
||||
throw new Error(`Invalid JSONL storage ${options.path}: invalid header`, { cause: parsedHeader.error });
|
||||
}
|
||||
if (parsedHeader.value.format === "v3-legacy") {
|
||||
return JsonlStorage.openLegacyV3(options, parsedHeader.value.header, lines.slice(1), context);
|
||||
}
|
||||
|
||||
const header = parsedHeader.value.header;
|
||||
if (header.storageVersion !== JSONL_STORAGE_VERSION) {
|
||||
throw new Error(`Session ${header.id} uses unsupported storage version ${header.storageVersion}`);
|
||||
}
|
||||
const storage = new JsonlStorage(options, header, { kind: "v4" });
|
||||
for (let index = 1; index < lines.length; index++) {
|
||||
const line = lines[index]!;
|
||||
try {
|
||||
storage.replayCommitted(parseTransaction(line));
|
||||
} catch (error) {
|
||||
throw new Error(`Invalid JSONL storage ${options.path}: line ${index + 1}`, { cause: error });
|
||||
}
|
||||
}
|
||||
if (header.nextSeq !== undefined) storage.storageState.advanceNextSeq(header.nextSeq);
|
||||
if (torn) await publishFileAtomically(options.fileSystem, options.path, `${lines.join("\n")}\n`, context);
|
||||
return storage;
|
||||
}
|
||||
|
||||
private static async openLegacyV3(
|
||||
options: JsonlStorageOptions,
|
||||
header: LegacyV3SessionHeader,
|
||||
recordLines: readonly string[],
|
||||
context: Context,
|
||||
): Promise<JsonlStorage> {
|
||||
const { writes, importedUsage, nextSeq } = normalizeLegacyV3Records(recordLines);
|
||||
const targetHeader = {
|
||||
...(await normalizeLegacyV3Header(options.fileSystem, header, context)),
|
||||
nextSeq,
|
||||
};
|
||||
const storage = new JsonlStorage(options, targetHeader, {
|
||||
kind: "v3",
|
||||
importedUsage,
|
||||
baselineWrites: writes,
|
||||
});
|
||||
storage.replayCommitted(writes);
|
||||
return storage;
|
||||
}
|
||||
|
||||
private replayCommitted(writes: readonly CommittedWrite[]): void {
|
||||
this.storageState.validateCommitted(writes);
|
||||
this.storageState.applyValidated(writes);
|
||||
}
|
||||
|
||||
async commit(writes: Write[], context: Context): Promise<CommitResult> {
|
||||
if (this.state !== "open") throw new Error("JsonlStorage is closed");
|
||||
const result = this.commitQueue.then(() => this.applyCommit(writes, context));
|
||||
this.commitQueue = result.then(
|
||||
() => undefined,
|
||||
() => undefined,
|
||||
);
|
||||
return result;
|
||||
}
|
||||
|
||||
private async applyCommit(writes: Write[], context: Context): Promise<CommitResult> {
|
||||
if (this.backing.kind === "v3" && writes.length !== 0) {
|
||||
return this.upgradeLegacyV3ToV4(this.backing, writes, context);
|
||||
}
|
||||
const prepared = this.storageState.prepareCommit(writes, this.now());
|
||||
if (prepared.writes.length !== 0) {
|
||||
fileValue(
|
||||
await this.fileSystem.appendFile(this.path, `${serializeTransaction(prepared.writes)}\n`, context),
|
||||
`Failed to append JSONL storage ${this.path}`,
|
||||
);
|
||||
}
|
||||
const stats = this.storageState.applyValidated(prepared.writes);
|
||||
return { ...prepared.result, stats: this.withImportedUsage(stats) };
|
||||
}
|
||||
|
||||
/** Atomically upgrade legacy v3 backing and preserve the first caller write as a v4 transaction. */
|
||||
private async upgradeLegacyV3ToV4(
|
||||
backing: LegacyV3Backing,
|
||||
callerWrites: Write[],
|
||||
context: Context,
|
||||
): Promise<CommitResult> {
|
||||
const timestamp = this.now();
|
||||
const prepared = this.storageState.prepareCommit(
|
||||
[
|
||||
insertUsage({
|
||||
id: uuidv7(timestamp),
|
||||
usage: backing.importedUsage,
|
||||
adjustment: true,
|
||||
details: { source: "v3-import" },
|
||||
}),
|
||||
...callerWrites,
|
||||
],
|
||||
timestamp,
|
||||
);
|
||||
|
||||
const nextSeq = prepared.result.firstSeq + prepared.writes.length;
|
||||
const upgradedHeader = { ...this.header, nextSeq };
|
||||
await publishFileAtomically(
|
||||
this.fileSystem,
|
||||
this.path,
|
||||
serializeStorage(upgradedHeader, [...backing.baselineWrites.map((write) => [write]), prepared.writes]),
|
||||
context,
|
||||
);
|
||||
|
||||
const stats = this.storageState.applyValidated(prepared.writes);
|
||||
this.backing = { kind: "v4" };
|
||||
// The first sequence belongs to the internal usage adjustment; return only caller-write sequences.
|
||||
return {
|
||||
...prepared.result,
|
||||
firstSeq: prepared.result.firstSeq + 1,
|
||||
seqs: prepared.result.seqs.slice(1),
|
||||
stats,
|
||||
};
|
||||
}
|
||||
|
||||
getEntries(ids: string[], _context: Context): Promise<Map<string, Entry>> {
|
||||
if (this.state !== "open") return Promise.reject(new Error("JsonlStorage is closed"));
|
||||
return Promise.resolve(this.storageState.getEntries(ids));
|
||||
}
|
||||
|
||||
getValue<T>(address: Value<T>, _context: Context): Promise<StoredValue<T> | undefined> {
|
||||
if (this.state !== "open") return Promise.reject(new Error("JsonlStorage is closed"));
|
||||
return Promise.resolve(this.storageState.getValue(address));
|
||||
}
|
||||
|
||||
scanValues<T>(prefix: Value<T>, _context: Context): Promise<StoredValue<T>[]> {
|
||||
if (this.state !== "open") return Promise.reject(new Error("JsonlStorage is closed"));
|
||||
return Promise.resolve(this.storageState.scanValues(prefix));
|
||||
}
|
||||
|
||||
async readList<T>(
|
||||
address: ValueList<T>,
|
||||
options: ListReadOptions | undefined,
|
||||
_context: Context,
|
||||
): Promise<ListElement<T>[]> {
|
||||
if (this.state !== "open") throw new Error("JsonlStorage is closed");
|
||||
return this.storageState.readList(address, options);
|
||||
}
|
||||
|
||||
async scanBranch(query: StorageBranchScan, _context: Context): Promise<Entry[]> {
|
||||
if (this.state !== "open") throw new Error("JsonlStorage is closed");
|
||||
return this.storageState.scanBranch(query);
|
||||
}
|
||||
|
||||
async scanBranchStructure(query: StorageBranchScan, _context: Context): Promise<EntryStructure[]> {
|
||||
if (this.state !== "open") throw new Error("JsonlStorage is closed");
|
||||
return this.storageState.scanBranchStructure(query);
|
||||
}
|
||||
|
||||
scanEntries(query: EntryScan, _context: Context): Promise<Entry[]> {
|
||||
if (this.state !== "open") return Promise.reject(new Error("JsonlStorage is closed"));
|
||||
return Promise.resolve(this.storageState.scanEntries(query));
|
||||
}
|
||||
|
||||
scanUsage(query: UsageScan, _context: Context): Promise<UsageRow[]> {
|
||||
if (this.state !== "open") return Promise.reject(new Error("JsonlStorage is closed"));
|
||||
return Promise.resolve(this.storageState.scanUsage(query));
|
||||
}
|
||||
|
||||
getStats(_context: Context): Promise<SessionStats> {
|
||||
if (this.state !== "open") return Promise.reject(new Error("JsonlStorage is closed"));
|
||||
return Promise.resolve(this.withImportedUsage(this.storageState.getStats()));
|
||||
}
|
||||
|
||||
private withImportedUsage(stats: SessionStats): SessionStats {
|
||||
return this.backing.kind === "v4" ? stats : { ...stats, usage: this.backing.importedUsage };
|
||||
}
|
||||
|
||||
/** Capture the state needed to fork at one serialized boundary between commits. */
|
||||
captureForkSource(_context: Context): Promise<ForkSourceSnapshot> {
|
||||
if (this.state !== "open") return Promise.reject(new Error("JsonlStorage is closed"));
|
||||
const result = this.commitQueue.then(() => this.storageState.snapshotEntriesAndValues());
|
||||
this.commitQueue = result.then(
|
||||
() => undefined,
|
||||
() => undefined,
|
||||
);
|
||||
return result;
|
||||
}
|
||||
|
||||
close(_context: Context): Promise<void> {
|
||||
if (this.closePromise !== undefined) return this.closePromise;
|
||||
this.state = "closing";
|
||||
this.closePromise = this.commitQueue.then(() => {
|
||||
this.state = "closed";
|
||||
});
|
||||
return this.closePromise;
|
||||
}
|
||||
}
|
||||
45
third_party/pi/packages/agent/src/harness/session/jsonl/types.ts
vendored
Normal file
45
third_party/pi/packages/agent/src/harness/session/jsonl/types.ts
vendored
Normal file
@@ -0,0 +1,45 @@
|
||||
import type { FileSystem } from "../../types.ts";
|
||||
import type { SessionCreateOptions, SessionMetadata } from "../types.ts";
|
||||
|
||||
export const JSONL_FORMAT_VERSION = 4;
|
||||
export const JSONL_STORAGE_VERSION = 1;
|
||||
|
||||
export interface JsonlStorageHeader {
|
||||
v: typeof JSONL_FORMAT_VERSION;
|
||||
kind: "header";
|
||||
id: string;
|
||||
storageVersion: number;
|
||||
createdAt: number;
|
||||
cwd: string;
|
||||
parentSessionId?: string;
|
||||
legacyParentSessionPath?: string;
|
||||
/** Sequence high-water mark written by snapshot rewrites. */
|
||||
nextSeq?: number;
|
||||
}
|
||||
|
||||
export interface JsonlStorageOptions {
|
||||
fileSystem: FileSystem;
|
||||
path: string;
|
||||
now?: () => number;
|
||||
}
|
||||
|
||||
export interface JsonlSessionMetadata extends SessionMetadata {
|
||||
cwd: string;
|
||||
path: string;
|
||||
/** Filesystem modification time as milliseconds since Unix epoch. */
|
||||
modifiedAt: number;
|
||||
}
|
||||
|
||||
export interface JsonlSessionCreateOptions extends SessionCreateOptions {
|
||||
cwd: string;
|
||||
}
|
||||
|
||||
export interface JsonlSessionListOptions {
|
||||
cwd?: string;
|
||||
}
|
||||
|
||||
export interface JsonlSessionRepoOptions {
|
||||
fileSystem: FileSystem;
|
||||
sessionsRoot: string;
|
||||
now?: () => number;
|
||||
}
|
||||
464
third_party/pi/packages/agent/src/harness/session/memory.ts
vendored
Normal file
464
third_party/pi/packages/agent/src/harness/session/memory.ts
vendored
Normal file
@@ -0,0 +1,464 @@
|
||||
import { uuidv7 } from "@earendil-works/pi-ai/utils/uuid";
|
||||
import type { Context } from "../context.ts";
|
||||
import {
|
||||
createForkSnapshot,
|
||||
type ForkDestinationSnapshot,
|
||||
type ForkSourceSnapshot,
|
||||
forkSnapshotWrites,
|
||||
} from "./fork.ts";
|
||||
import { InMemoryStorageState } from "./in-memory-storage-state.ts";
|
||||
import { StorageBackedSession } from "./session.ts";
|
||||
import type {
|
||||
Branch,
|
||||
CommitResult,
|
||||
Entry,
|
||||
EntryQuery,
|
||||
EntryScan,
|
||||
EntryStructure,
|
||||
ForkOptions,
|
||||
IdGenerator,
|
||||
Session,
|
||||
SessionCreateOptions,
|
||||
SessionMetadata,
|
||||
SessionMutation,
|
||||
SessionMutationCallback,
|
||||
SessionRepo,
|
||||
SessionStats,
|
||||
Storage,
|
||||
StorageBranchScan,
|
||||
UsageRow,
|
||||
UsageScan,
|
||||
Write,
|
||||
} from "./types.ts";
|
||||
import type { ListElement, ListReadOptions, StoredValue, Value, ValueList } from "./values.ts";
|
||||
|
||||
export interface MemoryStorageOptions {
|
||||
now?: () => number;
|
||||
}
|
||||
|
||||
export interface MemorySessionRepoOptions {
|
||||
now?: () => number;
|
||||
}
|
||||
|
||||
export class MemoryStorage implements Storage {
|
||||
private readonly now: () => number;
|
||||
private storageState = new InMemoryStorageState();
|
||||
private commitQueue: Promise<void> = Promise.resolve();
|
||||
private state: "open" | "closing" | "closed" = "open";
|
||||
private closePromise: Promise<void> | undefined;
|
||||
|
||||
constructor(options: MemoryStorageOptions = {}) {
|
||||
this.now = options.now ?? Date.now;
|
||||
}
|
||||
|
||||
async commit(writes: Write[], _context: Context): Promise<CommitResult> {
|
||||
if (this.state !== "open") throw new Error("MemoryStorage is closed");
|
||||
const result = this.commitQueue.then(() => {
|
||||
const prepared = this.storageState.prepareCommit(writes, this.now());
|
||||
const stats = this.storageState.applyValidated(prepared.writes);
|
||||
return { ...prepared.result, stats };
|
||||
});
|
||||
this.commitQueue = result.then(
|
||||
() => undefined,
|
||||
() => undefined,
|
||||
);
|
||||
return result;
|
||||
}
|
||||
|
||||
getEntries(ids: string[], _context: Context): Promise<Map<string, Entry>> {
|
||||
if (this.state !== "open") return Promise.reject(new Error("MemoryStorage is closed"));
|
||||
return Promise.resolve(this.storageState.getEntries(ids));
|
||||
}
|
||||
|
||||
getValue<T>(address: Value<T>, _context: Context): Promise<StoredValue<T> | undefined> {
|
||||
if (this.state !== "open") return Promise.reject(new Error("MemoryStorage is closed"));
|
||||
return Promise.resolve(this.storageState.getValue(address));
|
||||
}
|
||||
|
||||
scanValues<T>(prefix: Value<T>, _context: Context): Promise<StoredValue<T>[]> {
|
||||
if (this.state !== "open") return Promise.reject(new Error("MemoryStorage is closed"));
|
||||
return Promise.resolve(this.storageState.scanValues(prefix));
|
||||
}
|
||||
|
||||
async readList<T>(
|
||||
address: ValueList<T>,
|
||||
options: ListReadOptions | undefined,
|
||||
_context: Context,
|
||||
): Promise<ListElement<T>[]> {
|
||||
if (this.state !== "open") throw new Error("MemoryStorage is closed");
|
||||
return this.storageState.readList(address, options);
|
||||
}
|
||||
|
||||
async scanBranch(query: StorageBranchScan, _context: Context): Promise<Entry[]> {
|
||||
if (this.state !== "open") throw new Error("MemoryStorage is closed");
|
||||
return this.storageState.scanBranch(query);
|
||||
}
|
||||
|
||||
async scanBranchStructure(query: StorageBranchScan, _context: Context): Promise<EntryStructure[]> {
|
||||
if (this.state !== "open") throw new Error("MemoryStorage is closed");
|
||||
return this.storageState.scanBranchStructure(query);
|
||||
}
|
||||
|
||||
scanEntries(query: EntryScan, _context: Context): Promise<Entry[]> {
|
||||
if (this.state !== "open") return Promise.reject(new Error("MemoryStorage is closed"));
|
||||
return Promise.resolve(this.storageState.scanEntries(query));
|
||||
}
|
||||
|
||||
scanUsage(query: UsageScan, _context: Context): Promise<UsageRow[]> {
|
||||
if (this.state !== "open") return Promise.reject(new Error("MemoryStorage is closed"));
|
||||
return Promise.resolve(this.storageState.scanUsage(query));
|
||||
}
|
||||
|
||||
getStats(_context: Context): Promise<SessionStats> {
|
||||
if (this.state !== "open") return Promise.reject(new Error("MemoryStorage is closed"));
|
||||
return Promise.resolve(this.storageState.getStats());
|
||||
}
|
||||
|
||||
/** Capture the state needed to fork at one serialized boundary between commits. */
|
||||
captureForkSource(_context: Context): Promise<ForkSourceSnapshot> {
|
||||
if (this.state !== "open") return Promise.reject(new Error("MemoryStorage is closed"));
|
||||
const result = this.commitQueue.then(() => this.storageState.snapshotEntriesAndValues());
|
||||
this.commitQueue = result.then(
|
||||
() => undefined,
|
||||
() => undefined,
|
||||
);
|
||||
return result;
|
||||
}
|
||||
|
||||
close(_context: Context): Promise<void> {
|
||||
if (this.closePromise !== undefined) return this.closePromise;
|
||||
this.state = "closing";
|
||||
this.closePromise = this.commitQueue.then(() => {
|
||||
this.state = "closed";
|
||||
});
|
||||
return this.closePromise;
|
||||
}
|
||||
|
||||
static fromSnapshot(options: MemoryStorageOptions, snapshot: ForkDestinationSnapshot): MemoryStorage {
|
||||
const storage = new MemoryStorage(options);
|
||||
const writes = forkSnapshotWrites(snapshot);
|
||||
storage.storageState.validateCommitted(writes);
|
||||
storage.storageState.applyValidated(writes);
|
||||
return storage;
|
||||
}
|
||||
}
|
||||
|
||||
const MEMORY_STORAGE_VERSION = 1;
|
||||
|
||||
interface MemorySessionRecord {
|
||||
metadata: SessionMetadata;
|
||||
storage: MemoryStorage;
|
||||
session: StorageBackedSession;
|
||||
open: boolean;
|
||||
}
|
||||
|
||||
class MemorySessionFacade implements Session {
|
||||
readonly metadata: SessionMetadata;
|
||||
readonly idGenerator: IdGenerator;
|
||||
private readonly session: StorageBackedSession;
|
||||
private readonly onClose: () => void;
|
||||
private readonly admitted = new Set<Promise<unknown>>();
|
||||
private readonly closedError = new Error("Session is closed");
|
||||
private state: "open" | "closing" | "closed" = "open";
|
||||
private closePromise: Promise<void> | undefined;
|
||||
|
||||
constructor(session: StorageBackedSession, onClose: () => void) {
|
||||
this.session = session;
|
||||
this.metadata = session.metadata;
|
||||
this.idGenerator = session.idGenerator;
|
||||
this.onClose = onClose;
|
||||
}
|
||||
|
||||
async beginMutation(context: Context): Promise<SessionMutation> {
|
||||
let resolveFinished!: () => void;
|
||||
const finished = new Promise<void>((resolve) => {
|
||||
resolveFinished = resolve;
|
||||
});
|
||||
this.admitted.add(finished);
|
||||
let source: SessionMutation;
|
||||
try {
|
||||
source = await this.admit(() => this.session.beginMutation(context));
|
||||
} catch (error) {
|
||||
this.admitted.delete(finished);
|
||||
resolveFinished();
|
||||
throw error;
|
||||
}
|
||||
if (this.state !== "open") {
|
||||
await source.end(context);
|
||||
this.admitted.delete(finished);
|
||||
resolveFinished();
|
||||
throw this.closedError;
|
||||
}
|
||||
let ended = false;
|
||||
return {
|
||||
commit: (writes, commitContext) => source.commit(writes, commitContext),
|
||||
end: async (endContext) => {
|
||||
try {
|
||||
await source.end(endContext);
|
||||
} finally {
|
||||
if (!ended) {
|
||||
ended = true;
|
||||
this.admitted.delete(finished);
|
||||
resolveFinished();
|
||||
}
|
||||
}
|
||||
},
|
||||
getEntries: (ids, readContext) => source.getEntries(ids, readContext),
|
||||
getStats: (readContext) => source.getStats(readContext),
|
||||
getValue: (address, readContext) => source.getValue(address, readContext),
|
||||
scanValues: (prefix, readContext) => source.scanValues(prefix, readContext),
|
||||
readList: (address, options, readContext) => source.readList(address, options, readContext),
|
||||
scanBranch: (query, readContext) => source.scanBranch(query, readContext),
|
||||
};
|
||||
}
|
||||
|
||||
mutate<T>(mutation: SessionMutationCallback<T>, context: Context): Promise<T> {
|
||||
return this.admit(() =>
|
||||
this.session.mutate((mutator, mutationContext) => {
|
||||
if (this.state !== "open") throw this.closedError;
|
||||
return mutation(mutator, mutationContext);
|
||||
}, context),
|
||||
);
|
||||
}
|
||||
|
||||
getEntries(ids: string[], context: Context): Promise<Map<string, Entry>> {
|
||||
return this.admit(() => this.session.getEntries(ids, context));
|
||||
}
|
||||
|
||||
getEntry(id: string, context: Context): Promise<Entry | undefined> {
|
||||
return this.admit(() => this.session.getEntry(id, context));
|
||||
}
|
||||
|
||||
getValue<T>(address: Value<T>, context: Context): Promise<StoredValue<T> | undefined> {
|
||||
return this.admit(() => this.session.getValue(address, context));
|
||||
}
|
||||
|
||||
scanValues<T>(prefix: Value<T>, context: Context): Promise<StoredValue<T>[]> {
|
||||
return this.admit(() => this.session.scanValues(prefix, context));
|
||||
}
|
||||
|
||||
readList<T>(
|
||||
address: ValueList<T>,
|
||||
options: ListReadOptions | undefined,
|
||||
context: Context,
|
||||
): Promise<ListElement<T>[]> {
|
||||
return this.admit(() => this.session.readList(address, options, context));
|
||||
}
|
||||
|
||||
scanBranch(query: StorageBranchScan, context: Context): Promise<Entry[]> {
|
||||
return this.admit(() => this.session.scanBranch(query, context));
|
||||
}
|
||||
|
||||
getStats(context: Context): Promise<SessionStats> {
|
||||
return this.admit(() => this.session.getStats(context));
|
||||
}
|
||||
|
||||
getName(context: Context): Promise<string | undefined> {
|
||||
return this.admit(() => this.session.getName(context));
|
||||
}
|
||||
|
||||
getLabel(targetId: string, context: Context): Promise<string | undefined> {
|
||||
return this.admit(() => this.session.getLabel(targetId, context));
|
||||
}
|
||||
|
||||
findEntries(query: EntryQuery | undefined, context: Context): Promise<Entry[]> {
|
||||
return this.admit(() => this.session.findEntries(query, context));
|
||||
}
|
||||
|
||||
findEntry(query: EntryQuery | undefined, context: Context): Promise<Entry | undefined> {
|
||||
return this.admit(() => this.session.findEntry(query, context));
|
||||
}
|
||||
|
||||
async branch(name: string, context: Context): Promise<Branch | undefined> {
|
||||
const branch = await this.admit(() => this.session.branch(name, context));
|
||||
return branch === undefined ? undefined : this.wrapBranch(branch);
|
||||
}
|
||||
|
||||
async createBranch(name: string, at: string | null, context: Context): Promise<Branch> {
|
||||
return this.wrapBranch(await this.admit(() => this.session.createBranch(name, at, context)));
|
||||
}
|
||||
|
||||
setValue<T>(address: Value<T>, next: NoInfer<T>, context: Context): Promise<void> {
|
||||
return this.admit(() => this.session.setValue(address, next, context));
|
||||
}
|
||||
|
||||
deleteValue<T>(address: Value<T>, context: Context): Promise<void> {
|
||||
return this.admit(() => this.session.deleteValue(address, context));
|
||||
}
|
||||
|
||||
appendList<T>(address: ValueList<T>, element: NoInfer<T>, context: Context): Promise<void> {
|
||||
return this.admit(() => this.session.appendList(address, element, context));
|
||||
}
|
||||
|
||||
deleteList<T>(address: ValueList<T>, context: Context): Promise<void> {
|
||||
return this.admit(() => this.session.deleteList(address, context));
|
||||
}
|
||||
|
||||
setName(name: string | undefined, context: Context): Promise<void> {
|
||||
return this.admit(() => this.session.setName(name, context));
|
||||
}
|
||||
|
||||
setLabel(targetId: string, label: string | undefined, context: Context): Promise<void> {
|
||||
return this.admit(() => this.session.setLabel(targetId, label, context));
|
||||
}
|
||||
|
||||
close(_context: Context): Promise<void> {
|
||||
if (this.closePromise !== undefined) return this.closePromise;
|
||||
this.state = "closing";
|
||||
this.closePromise = Promise.allSettled([...this.admitted]).then(() => {
|
||||
this.state = "closed";
|
||||
this.onClose();
|
||||
});
|
||||
return this.closePromise;
|
||||
}
|
||||
|
||||
private wrapBranch(branch: Branch): Branch {
|
||||
return {
|
||||
name: branch.name,
|
||||
getTipId: (context) => this.admit(() => branch.getTipId(context)),
|
||||
findEntries: (query, context) => this.admit(() => branch.findEntries(query, context)),
|
||||
findEntry: (query, context) => this.admit(() => branch.findEntry(query, context)),
|
||||
appendMessage: (message, context) => this.admit(() => branch.appendMessage(message, context)),
|
||||
appendCustomEntry: (customType, data, context) =>
|
||||
this.admit(() => branch.appendCustomEntry(customType, data, context)),
|
||||
};
|
||||
}
|
||||
|
||||
private admit<T>(operation: () => Promise<T>): Promise<T> {
|
||||
if (this.state !== "open") return Promise.reject(this.closedError);
|
||||
let result: Promise<T>;
|
||||
try {
|
||||
result = operation();
|
||||
} catch (error) {
|
||||
result = Promise.reject(error);
|
||||
}
|
||||
this.admitted.add(result);
|
||||
void result.then(
|
||||
() => this.admitted.delete(result),
|
||||
() => this.admitted.delete(result),
|
||||
);
|
||||
return result;
|
||||
}
|
||||
}
|
||||
|
||||
export class MemorySessionRepo implements SessionRepo {
|
||||
private readonly now: () => number;
|
||||
private readonly sessions = new Map<string, MemorySessionRecord>();
|
||||
private readonly pendingIds = new Set<string>();
|
||||
private closed = false;
|
||||
private closePromise: Promise<void> | undefined;
|
||||
|
||||
constructor(options: MemorySessionRepoOptions = {}) {
|
||||
this.now = options.now ?? Date.now;
|
||||
}
|
||||
|
||||
async create(options: SessionCreateOptions, context: Context): Promise<Session> {
|
||||
this.assertOpen();
|
||||
const createdAt = this.now();
|
||||
const id = options.id ?? uuidv7(createdAt);
|
||||
this.reserveId(id);
|
||||
const metadata: SessionMetadata = {
|
||||
id,
|
||||
createdAt,
|
||||
storageVersion: MEMORY_STORAGE_VERSION,
|
||||
...(options.parentSessionId === undefined ? {} : { parentSessionId: options.parentSessionId }),
|
||||
};
|
||||
const storage = new MemoryStorage({ now: this.now });
|
||||
const session = new StorageBackedSession(metadata, storage);
|
||||
try {
|
||||
const record: MemorySessionRecord = {
|
||||
metadata,
|
||||
storage,
|
||||
session,
|
||||
open: true,
|
||||
};
|
||||
this.sessions.set(id, record);
|
||||
return this.openRecord(record);
|
||||
} catch (error) {
|
||||
await session.close(context);
|
||||
throw error;
|
||||
} finally {
|
||||
this.pendingIds.delete(id);
|
||||
}
|
||||
}
|
||||
|
||||
open(metadata: SessionMetadata, _context: Context): Promise<Session> {
|
||||
// Memory sessions are always created at the current storage version, so
|
||||
// persistent-backend version gating does not apply here.
|
||||
this.assertOpen();
|
||||
const record = this.sessions.get(metadata.id);
|
||||
if (record === undefined) return Promise.reject(new Error(`Unknown session: ${metadata.id}`));
|
||||
if (record.open) return Promise.reject(new Error(`Session is already open: ${metadata.id}`));
|
||||
record.open = true;
|
||||
return Promise.resolve(this.openRecord(record));
|
||||
}
|
||||
|
||||
list(_options: undefined, _context: Context): Promise<SessionMetadata[]> {
|
||||
this.assertOpen();
|
||||
return Promise.resolve([...this.sessions.values()].map(({ metadata }) => metadata));
|
||||
}
|
||||
|
||||
async delete(metadata: SessionMetadata, context: Context): Promise<void> {
|
||||
this.assertOpen();
|
||||
const record = this.sessions.get(metadata.id);
|
||||
if (record === undefined) throw new Error(`Unknown session: ${metadata.id}`);
|
||||
if (record.open) throw new Error(`Session is open: ${metadata.id}`);
|
||||
await record.session.close(context);
|
||||
this.sessions.delete(metadata.id);
|
||||
}
|
||||
|
||||
async fork(source: SessionMetadata, options: ForkOptions, context: Context): Promise<Session> {
|
||||
this.assertOpen();
|
||||
const sourceRecord = this.sessions.get(source.id);
|
||||
if (sourceRecord === undefined) throw new Error(`Unknown session: ${source.id}`);
|
||||
const createdAt = this.now();
|
||||
const id = options.id ?? uuidv7(createdAt);
|
||||
this.reserveId(id);
|
||||
|
||||
try {
|
||||
const snapshot = createForkSnapshot(await sourceRecord.storage.captureForkSource(context), options);
|
||||
const storage = MemoryStorage.fromSnapshot({ now: this.now }, snapshot);
|
||||
const metadata: SessionMetadata = {
|
||||
id,
|
||||
createdAt,
|
||||
storageVersion: MEMORY_STORAGE_VERSION,
|
||||
parentSessionId: sourceRecord.metadata.id,
|
||||
};
|
||||
const session = new StorageBackedSession(metadata, storage);
|
||||
const record: MemorySessionRecord = {
|
||||
metadata,
|
||||
storage,
|
||||
session,
|
||||
open: true,
|
||||
};
|
||||
this.sessions.set(id, record);
|
||||
return this.openRecord(record);
|
||||
} finally {
|
||||
this.pendingIds.delete(id);
|
||||
}
|
||||
}
|
||||
|
||||
close(context: Context): Promise<void> {
|
||||
if (this.closePromise !== undefined) return this.closePromise;
|
||||
this.closed = true;
|
||||
this.closePromise = Promise.all([...this.sessions.values()].map(({ session }) => session.close(context))).then(
|
||||
() => undefined,
|
||||
);
|
||||
return this.closePromise;
|
||||
}
|
||||
|
||||
private openRecord(record: MemorySessionRecord): Session {
|
||||
return new MemorySessionFacade(record.session, () => {
|
||||
record.open = false;
|
||||
});
|
||||
}
|
||||
|
||||
private reserveId(id: string): void {
|
||||
if (this.sessions.has(id) || this.pendingIds.has(id)) throw new Error(`Session already exists: ${id}`);
|
||||
this.pendingIds.add(id);
|
||||
}
|
||||
|
||||
private assertOpen(): void {
|
||||
if (this.closed) throw new Error("MemorySessionRepo is closed");
|
||||
}
|
||||
}
|
||||
23
third_party/pi/packages/agent/src/harness/session/mutation-line.ts
vendored
Normal file
23
third_party/pi/packages/agent/src/harness/session/mutation-line.ts
vendored
Normal file
@@ -0,0 +1,23 @@
|
||||
/** Serializes complete read-modify-write jobs for one Session. */
|
||||
export class MutationLine {
|
||||
private tail: Promise<void> = Promise.resolve();
|
||||
private sealedError: Error | undefined;
|
||||
|
||||
run<T>(operation: () => T | Promise<T>): Promise<T> {
|
||||
if (this.sealedError !== undefined) return Promise.reject(this.sealedError);
|
||||
const result = this.tail.then(() => {
|
||||
if (this.sealedError !== undefined) throw this.sealedError;
|
||||
return operation();
|
||||
});
|
||||
this.tail = result.then(
|
||||
() => undefined,
|
||||
() => undefined,
|
||||
);
|
||||
return result;
|
||||
}
|
||||
|
||||
seal(error: Error): Promise<void> {
|
||||
this.sealedError ??= error;
|
||||
return this.tail;
|
||||
}
|
||||
}
|
||||
475
third_party/pi/packages/agent/src/harness/session/session.ts
vendored
Normal file
475
third_party/pi/packages/agent/src/harness/session/session.ts
vendored
Normal file
@@ -0,0 +1,475 @@
|
||||
import { uuidv7 } from "@earendil-works/pi-ai/utils/uuid";
|
||||
import type { AgentMessage } from "../../types.ts";
|
||||
import type { Context } from "../context.ts";
|
||||
import { insertEntry } from "./commit.ts";
|
||||
import { MutationLine } from "./mutation-line.ts";
|
||||
import type {
|
||||
Branch,
|
||||
BranchScan,
|
||||
CommitResult,
|
||||
Entry,
|
||||
EntryQuery,
|
||||
IdGenerator,
|
||||
JsonValue,
|
||||
Session,
|
||||
SessionMetadata,
|
||||
SessionMutation,
|
||||
SessionMutationCallback,
|
||||
SessionStats,
|
||||
Storage,
|
||||
StorageBranchScan,
|
||||
Write,
|
||||
} from "./types.ts";
|
||||
import {
|
||||
appendList as appendListWrite,
|
||||
branchTip,
|
||||
deleteList as deleteListWrite,
|
||||
deleteValue as deleteValueWrite,
|
||||
entryLabel,
|
||||
type ListElement,
|
||||
type ListReadOptions,
|
||||
type StoredValue,
|
||||
sessionName,
|
||||
setValue as setValueWrite,
|
||||
type Value,
|
||||
type ValueList,
|
||||
} from "./values.ts";
|
||||
|
||||
export interface StorageBackedSessionOptions {
|
||||
mutationLine?: MutationLine;
|
||||
idGenerator?: IdGenerator;
|
||||
onClose?: () => void;
|
||||
}
|
||||
|
||||
/** Durable session state is internally inconsistent and cannot be safely advanced. */
|
||||
export class SessionInvariantError extends Error {
|
||||
constructor(message: string) {
|
||||
super(message);
|
||||
this.name = "SessionInvariantError";
|
||||
}
|
||||
}
|
||||
|
||||
/** A requested Branch name is invalid. */
|
||||
export class SessionInvalidBranchError extends Error {
|
||||
readonly branch: string;
|
||||
readonly reason: string;
|
||||
|
||||
constructor(branch: string, reason: string) {
|
||||
super(`Invalid branch ${JSON.stringify(branch)}: ${reason}`);
|
||||
this.name = "SessionInvalidBranchError";
|
||||
this.branch = branch;
|
||||
this.reason = reason;
|
||||
}
|
||||
}
|
||||
|
||||
/** A requested branch already exists. */
|
||||
export class SessionBranchExistsError extends Error {
|
||||
readonly branch: string;
|
||||
|
||||
constructor(branch: string) {
|
||||
super(`Branch already exists: ${branch}`);
|
||||
this.name = "SessionBranchExistsError";
|
||||
this.branch = branch;
|
||||
}
|
||||
}
|
||||
|
||||
/** A pending assistant message cannot be persisted as a session entry. */
|
||||
export class SessionPendingAssistantMessageError extends Error {
|
||||
constructor() {
|
||||
super("Cannot persist a pending assistant message");
|
||||
this.name = "SessionPendingAssistantMessageError";
|
||||
}
|
||||
}
|
||||
|
||||
/** A requested session entry target does not exist. */
|
||||
export class SessionUnknownTargetError extends Error {
|
||||
readonly targetId: string;
|
||||
|
||||
constructor(targetId: string) {
|
||||
super(`Unknown target: ${targetId}`);
|
||||
this.name = "SessionUnknownTargetError";
|
||||
this.targetId = targetId;
|
||||
}
|
||||
}
|
||||
|
||||
class StorageBackedSessionMutation implements SessionMutation {
|
||||
private readonly storage: Storage;
|
||||
private readonly release: () => void;
|
||||
private active = true;
|
||||
private commitResult: Promise<CommitResult> | undefined;
|
||||
private endPromise: Promise<void> | undefined;
|
||||
|
||||
constructor(storage: Storage, release: () => void) {
|
||||
this.storage = storage;
|
||||
this.release = release;
|
||||
}
|
||||
|
||||
commit(writes: Write[], context: Context): Promise<CommitResult> {
|
||||
this.assertActive();
|
||||
if (this.commitResult !== undefined) return Promise.reject(new Error("SessionMutator commit already attempted"));
|
||||
try {
|
||||
for (const write of writes) {
|
||||
if (
|
||||
write.kind === "entry" &&
|
||||
write.entry.type === "message" &&
|
||||
write.entry.message.role === "assistant" &&
|
||||
write.entry.message.stopReason === "pending"
|
||||
) {
|
||||
throw new SessionPendingAssistantMessageError();
|
||||
}
|
||||
}
|
||||
this.commitResult = this.storage.commit(writes, context);
|
||||
} catch (error) {
|
||||
this.commitResult = Promise.reject(error);
|
||||
}
|
||||
return this.commitResult;
|
||||
}
|
||||
|
||||
end(_context: Context): Promise<void> {
|
||||
if (this.endPromise !== undefined) return this.endPromise;
|
||||
this.active = false;
|
||||
this.endPromise = this.settle().finally(this.release);
|
||||
return this.endPromise;
|
||||
}
|
||||
|
||||
getEntries(ids: string[], context: Context): Promise<Map<string, Entry>> {
|
||||
this.assertActive();
|
||||
return this.storage.getEntries(ids, context);
|
||||
}
|
||||
|
||||
getStats(context: Context): Promise<SessionStats> {
|
||||
this.assertActive();
|
||||
return this.storage.getStats(context);
|
||||
}
|
||||
|
||||
getValue<T>(address: Value<T>, context: Context): Promise<StoredValue<T> | undefined> {
|
||||
this.assertActive();
|
||||
return this.storage.getValue(address, context);
|
||||
}
|
||||
|
||||
scanValues<T>(prefix: Value<T>, context: Context): Promise<StoredValue<T>[]> {
|
||||
this.assertActive();
|
||||
return this.storage.scanValues(prefix, context);
|
||||
}
|
||||
|
||||
readList<T>(
|
||||
address: ValueList<T>,
|
||||
options: ListReadOptions | undefined,
|
||||
context: Context,
|
||||
): Promise<ListElement<T>[]> {
|
||||
this.assertActive();
|
||||
return this.storage.readList(address, options, context);
|
||||
}
|
||||
|
||||
scanBranch(query: StorageBranchScan, context: Context): Promise<Entry[]> {
|
||||
this.assertActive();
|
||||
return this.storage.scanBranch(query, context);
|
||||
}
|
||||
|
||||
private settle(): Promise<void> {
|
||||
return (
|
||||
this.commitResult?.then(
|
||||
() => undefined,
|
||||
() => undefined,
|
||||
) ?? Promise.resolve()
|
||||
);
|
||||
}
|
||||
|
||||
private assertActive(): void {
|
||||
if (!this.active) throw new Error("SessionMutator cannot be used outside its mutation callback");
|
||||
}
|
||||
}
|
||||
|
||||
class StorageBackedBranch implements Branch {
|
||||
readonly name: string;
|
||||
private readonly session: StorageBackedSession;
|
||||
|
||||
constructor(name: string, session: StorageBackedSession) {
|
||||
this.name = name;
|
||||
this.session = session;
|
||||
}
|
||||
|
||||
getTipId(context: Context): Promise<string | null> {
|
||||
return this.session.getBranchTip(this.name, context);
|
||||
}
|
||||
|
||||
async findEntries(query: BranchScan | undefined, context: Context): Promise<Entry[]> {
|
||||
query ??= {};
|
||||
const start = query.start ?? (await this.getTipId(context));
|
||||
if (start === null) return [];
|
||||
return this.session.scanBranch({ ...query, start, order: query.order ?? "newestFirst" }, context);
|
||||
}
|
||||
|
||||
async findEntry(query: BranchScan | undefined, context: Context): Promise<Entry | undefined> {
|
||||
query ??= {};
|
||||
return (
|
||||
await this.findEntries({ ...query, limit: query.limit === undefined ? 1 : Math.min(query.limit, 1) }, context)
|
||||
)[0];
|
||||
}
|
||||
|
||||
appendMessage(message: AgentMessage, context: Context): Promise<string> {
|
||||
return this.session.appendToBranch(this.name, { type: "message", message }, context);
|
||||
}
|
||||
|
||||
appendCustomEntry(customType: string, data: JsonValue | undefined, context: Context): Promise<string> {
|
||||
return this.session.appendToBranch(
|
||||
this.name,
|
||||
{ type: "custom", customType, ...(data === undefined ? {} : { data }) },
|
||||
context,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/** Package-internal typed boundary shared by concrete session repositories. */
|
||||
export class StorageBackedSession<TMetadata extends SessionMetadata = SessionMetadata> implements Session<TMetadata> {
|
||||
readonly metadata: TMetadata;
|
||||
readonly idGenerator: IdGenerator;
|
||||
private readonly storage: Storage;
|
||||
private readonly mutationLine: MutationLine;
|
||||
private readonly onClose: (() => void) | undefined;
|
||||
private readonly branches = new Map<string, StorageBackedBranch>();
|
||||
private readonly closedError = new Error("Session is closed");
|
||||
private state: "open" | "closing" | "closed" = "open";
|
||||
private closePromise: Promise<void> | undefined;
|
||||
|
||||
constructor(metadata: TMetadata, storage: Storage, options: StorageBackedSessionOptions = {}) {
|
||||
this.metadata = metadata;
|
||||
this.idGenerator = options.idGenerator ?? { next: uuidv7 };
|
||||
this.storage = storage;
|
||||
this.mutationLine = options.mutationLine ?? new MutationLine();
|
||||
this.onClose = options.onClose;
|
||||
}
|
||||
|
||||
async beginMutation(_context: Context): Promise<SessionMutation> {
|
||||
this.assertOpen();
|
||||
let grant!: (mutation: SessionMutation) => void;
|
||||
let rejectGrant!: (error: unknown) => void;
|
||||
const granted = new Promise<SessionMutation>((resolve, reject) => {
|
||||
grant = resolve;
|
||||
rejectGrant = reject;
|
||||
});
|
||||
let release!: () => void;
|
||||
const finished = new Promise<void>((resolve) => {
|
||||
release = resolve;
|
||||
});
|
||||
const line = this.mutationLine.run(async () => {
|
||||
grant(new StorageBackedSessionMutation(this.storage, release));
|
||||
await finished;
|
||||
});
|
||||
void line.catch(rejectGrant);
|
||||
return granted;
|
||||
}
|
||||
|
||||
async mutate<T>(mutation: SessionMutationCallback<T>, context: Context): Promise<T> {
|
||||
const mutator = await this.beginMutation(context);
|
||||
try {
|
||||
return await mutation(mutator, context);
|
||||
} finally {
|
||||
await mutator.end(context);
|
||||
}
|
||||
}
|
||||
|
||||
async getEntries(ids: string[], context: Context): Promise<Map<string, Entry>> {
|
||||
this.assertOpen();
|
||||
return this.storage.getEntries(ids, context);
|
||||
}
|
||||
|
||||
async getEntry(id: string, context: Context): Promise<Entry | undefined> {
|
||||
return (await this.getEntries([id], context)).get(id);
|
||||
}
|
||||
|
||||
async getValue<T>(address: Value<T>, context: Context): Promise<StoredValue<T> | undefined> {
|
||||
this.assertOpen();
|
||||
return this.storage.getValue(address, context);
|
||||
}
|
||||
|
||||
async scanValues<T>(prefix: Value<T>, context: Context): Promise<StoredValue<T>[]> {
|
||||
this.assertOpen();
|
||||
return this.storage.scanValues(prefix, context);
|
||||
}
|
||||
|
||||
async readList<T>(
|
||||
address: ValueList<T>,
|
||||
options: ListReadOptions | undefined,
|
||||
context: Context,
|
||||
): Promise<ListElement<T>[]> {
|
||||
this.assertOpen();
|
||||
return this.storage.readList(address, options, context);
|
||||
}
|
||||
|
||||
async scanBranch(query: StorageBranchScan, context: Context): Promise<Entry[]> {
|
||||
this.assertOpen();
|
||||
return this.storage.scanBranch(query, context);
|
||||
}
|
||||
|
||||
async getStats(context: Context): Promise<SessionStats> {
|
||||
this.assertOpen();
|
||||
return this.storage.getStats(context);
|
||||
}
|
||||
|
||||
async getName(context: Context): Promise<string | undefined> {
|
||||
return (await this.getValue(sessionName, context))?.value;
|
||||
}
|
||||
|
||||
async getLabel(targetId: string, context: Context): Promise<string | undefined> {
|
||||
return (await this.getValue(entryLabel(targetId), context))?.value;
|
||||
}
|
||||
|
||||
async findEntries(query: EntryQuery | undefined, context: Context): Promise<Entry[]> {
|
||||
query ??= {};
|
||||
this.assertOpen();
|
||||
const order = query.order ?? "desc";
|
||||
if (query.cursor !== undefined) {
|
||||
if (order === "asc" && query.cursor.seq === Number.MAX_SAFE_INTEGER) return [];
|
||||
if (order === "desc" && query.cursor.seq <= 1) return [];
|
||||
}
|
||||
return this.storage.scanEntries(
|
||||
{
|
||||
type: query.type,
|
||||
customType: query.customType,
|
||||
order,
|
||||
limit: query.limit,
|
||||
...(query.cursor === undefined
|
||||
? {}
|
||||
: order === "asc"
|
||||
? { fromSeq: query.cursor.seq + 1 }
|
||||
: { toSeq: query.cursor.seq - 1 }),
|
||||
},
|
||||
context,
|
||||
);
|
||||
}
|
||||
|
||||
async findEntry(query: EntryQuery | undefined, context: Context): Promise<Entry | undefined> {
|
||||
query ??= {};
|
||||
return (
|
||||
await this.findEntries({ ...query, limit: query.limit === undefined ? 1 : Math.min(query.limit, 1) }, context)
|
||||
)[0];
|
||||
}
|
||||
|
||||
async branch(name: string, context: Context): Promise<Branch | undefined> {
|
||||
this.assertValidBranchName(name);
|
||||
if ((await this.getValue(branchTip(name), context)) === undefined) return undefined;
|
||||
return this.getOrCreateBranchObject(name);
|
||||
}
|
||||
|
||||
async createBranch(name: string, at: string | null, context: Context): Promise<Branch> {
|
||||
this.assertOpen();
|
||||
this.assertValidBranchName(name);
|
||||
await this.mutate(async (mutator) => {
|
||||
if ((await mutator.getValue(branchTip(name), context)) !== undefined) {
|
||||
throw new SessionBranchExistsError(name);
|
||||
}
|
||||
if (at !== null && !(await mutator.getEntries([at], context)).has(at)) {
|
||||
throw new SessionUnknownTargetError(at);
|
||||
}
|
||||
await mutator.commit([setValueWrite(branchTip(name), at)], context);
|
||||
}, context);
|
||||
return this.getOrCreateBranchObject(name);
|
||||
}
|
||||
|
||||
setValue<T>(address: Value<T>, next: NoInfer<T>, context: Context): Promise<void> {
|
||||
return this.mutate(async (mutator) => {
|
||||
await mutator.commit([setValueWrite(address, next)], context);
|
||||
}, context);
|
||||
}
|
||||
|
||||
deleteValue<T>(address: Value<T>, context: Context): Promise<void> {
|
||||
return this.mutate(async (mutator) => {
|
||||
await mutator.commit([deleteValueWrite(address)], context);
|
||||
}, context);
|
||||
}
|
||||
|
||||
appendList<T>(address: ValueList<T>, element: NoInfer<T>, context: Context): Promise<void> {
|
||||
return this.mutate(async (mutator) => {
|
||||
await mutator.commit([appendListWrite(address, element)], context);
|
||||
}, context);
|
||||
}
|
||||
|
||||
deleteList<T>(address: ValueList<T>, context: Context): Promise<void> {
|
||||
return this.mutate(async (mutator) => {
|
||||
await mutator.commit([deleteListWrite(address)], context);
|
||||
}, context);
|
||||
}
|
||||
|
||||
setName(name: string | undefined, context: Context): Promise<void> {
|
||||
return name === undefined ? this.deleteValue(sessionName, context) : this.setValue(sessionName, name, context);
|
||||
}
|
||||
|
||||
setLabel(targetId: string, label: string | undefined, context: Context): Promise<void> {
|
||||
const address = entryLabel(targetId);
|
||||
return label === undefined ? this.deleteValue(address, context) : this.setValue(address, label, context);
|
||||
}
|
||||
|
||||
close(context: Context): Promise<void> {
|
||||
if (this.closePromise !== undefined) return this.closePromise;
|
||||
this.state = "closing";
|
||||
this.closePromise = this.mutationLine
|
||||
.seal(this.closedError)
|
||||
.then(() => this.storage.close(context))
|
||||
.finally(() => {
|
||||
this.state = "closed";
|
||||
this.onClose?.();
|
||||
});
|
||||
return this.closePromise;
|
||||
}
|
||||
|
||||
async getBranchTip(name: string, context: Context): Promise<string | null> {
|
||||
const stored = await this.getValue(branchTip(name), context);
|
||||
if (stored === undefined) throw new SessionInvariantError(`Unknown branch: ${name}`);
|
||||
return stored.value;
|
||||
}
|
||||
|
||||
async appendToBranch(
|
||||
name: string,
|
||||
entry: { type: "message"; message: AgentMessage } | { type: "custom"; customType: string; data?: JsonValue },
|
||||
context: Context,
|
||||
): Promise<string> {
|
||||
this.assertOpen();
|
||||
if (entry.type === "message" && entry.message.role === "assistant" && entry.message.stopReason === "pending") {
|
||||
throw new SessionPendingAssistantMessageError();
|
||||
}
|
||||
const id = this.idGenerator.next();
|
||||
await this.mutate(async (mutator) => {
|
||||
const tip = await mutator.getValue(branchTip(name), context);
|
||||
if (tip === undefined) throw new SessionInvariantError(`Unknown branch: ${name}`);
|
||||
await mutator.commit(
|
||||
[
|
||||
insertEntry(
|
||||
entry.type === "message"
|
||||
? { id, parentId: tip.value, type: "message", message: entry.message }
|
||||
: {
|
||||
id,
|
||||
parentId: tip.value,
|
||||
type: "custom",
|
||||
customType: entry.customType,
|
||||
...(entry.data === undefined ? {} : { data: entry.data }),
|
||||
},
|
||||
),
|
||||
setValueWrite(branchTip(name), id),
|
||||
],
|
||||
context,
|
||||
);
|
||||
}, context);
|
||||
return id;
|
||||
}
|
||||
|
||||
private getOrCreateBranchObject(name: string): Branch {
|
||||
let branch = this.branches.get(name);
|
||||
if (branch === undefined) {
|
||||
branch = new StorageBackedBranch(name, this);
|
||||
this.branches.set(name, branch);
|
||||
}
|
||||
return branch;
|
||||
}
|
||||
|
||||
private assertValidBranchName(name: string): void {
|
||||
if (name.length === 0) throw new SessionInvalidBranchError(name, "branch name must not be empty");
|
||||
if (name.includes("\u0000")) {
|
||||
throw new SessionInvalidBranchError(name, "branch name must not contain \\u0000");
|
||||
}
|
||||
}
|
||||
|
||||
private assertOpen(): void {
|
||||
if (this.state !== "open") throw this.closedError;
|
||||
}
|
||||
}
|
||||
33
third_party/pi/packages/agent/src/harness/session/testing/benchmark/datasets.ts
vendored
Normal file
33
third_party/pi/packages/agent/src/harness/session/testing/benchmark/datasets.ts
vendored
Normal file
@@ -0,0 +1,33 @@
|
||||
export interface StorageBenchmarkDataset {
|
||||
readonly name: string;
|
||||
readonly entryCount: number;
|
||||
readonly payloadBytes: number;
|
||||
readonly lookupIds: readonly string[];
|
||||
readonly tipId: string;
|
||||
}
|
||||
|
||||
/** Package-internal deterministic id shared by dataset and transaction generation. */
|
||||
export function storageBenchmarkEntryId(index: number): string {
|
||||
return `benchmark-entry-${index.toString().padStart(8, "0")}`;
|
||||
}
|
||||
|
||||
function createDataset(scale: string, entryCount: number): StorageBenchmarkDataset {
|
||||
const lookupCount = Math.min(100, entryCount);
|
||||
return {
|
||||
name: `synthetic linear branch: ${scale}, 256-byte payloads`,
|
||||
entryCount,
|
||||
payloadBytes: 256,
|
||||
lookupIds: Array.from({ length: lookupCount }, (_, index) => {
|
||||
const entryIndex = Math.floor((index * (entryCount - 1)) / Math.max(1, lookupCount - 1));
|
||||
return storageBenchmarkEntryId(entryIndex);
|
||||
}),
|
||||
tipId: storageBenchmarkEntryId(entryCount - 1),
|
||||
};
|
||||
}
|
||||
|
||||
/** Deterministic synthetic linear branches shared by all storage measurements. */
|
||||
export const STORAGE_BENCHMARK_DATASETS: readonly StorageBenchmarkDataset[] = [
|
||||
createDataset("1k entries", 1_000),
|
||||
createDataset("10k entries", 10_000),
|
||||
createDataset("100k entries", 100_000),
|
||||
];
|
||||
181
third_party/pi/packages/agent/src/harness/session/testing/benchmark/session-repo.ts
vendored
Normal file
181
third_party/pi/packages/agent/src/harness/session/testing/benchmark/session-repo.ts
vendored
Normal file
@@ -0,0 +1,181 @@
|
||||
import { BACKGROUND_CONTEXT } from "../../../context.ts";
|
||||
import type { SessionMetadata, SessionRepo } from "../../types.ts";
|
||||
import { branchTip, laneConfig, laneState, setValue } from "../../values.ts";
|
||||
import { STORAGE_BENCHMARK_DATASETS, type StorageBenchmarkDataset } from "./datasets.ts";
|
||||
import { generateStorageBenchmarkSeedTransactions } from "./storage.ts";
|
||||
|
||||
export interface SessionRepoCatalogBenchmarkDataset {
|
||||
readonly name: string;
|
||||
readonly sessionCount: number;
|
||||
}
|
||||
|
||||
interface SessionRepoCatalogReadBenchmarkScenario {
|
||||
readonly name: string;
|
||||
expectedResult(dataset: SessionRepoCatalogBenchmarkDataset): number;
|
||||
run(repo: SessionRepo): Promise<number>;
|
||||
}
|
||||
|
||||
interface SessionRepoCatalogWriteBenchmarkOperation {
|
||||
run(): Promise<number>;
|
||||
}
|
||||
|
||||
interface SessionRepoCatalogWriteBenchmarkScenario {
|
||||
readonly name: string;
|
||||
readonly expectedResult: number;
|
||||
prepare(repo: SessionRepo): Promise<SessionRepoCatalogWriteBenchmarkOperation>;
|
||||
}
|
||||
|
||||
interface SessionRepoForkWriteBenchmarkScenario {
|
||||
readonly name: string;
|
||||
expectedResult(dataset: StorageBenchmarkDataset): number;
|
||||
run(repo: SessionRepo, source: SessionMetadata, dataset: StorageBenchmarkDataset): Promise<number>;
|
||||
}
|
||||
|
||||
/** Package-internal deterministic session id shared by repository benchmark workloads. */
|
||||
export function sessionRepoBenchmarkSessionId(index: number): string {
|
||||
return `benchmark-session-${index.toString().padStart(8, "0")}`;
|
||||
}
|
||||
|
||||
const BENCHMARK_SESSION_ID = sessionRepoBenchmarkSessionId(0);
|
||||
const FORK_DESTINATION_SESSION_ID = sessionRepoBenchmarkSessionId(1);
|
||||
|
||||
function createCatalogDataset(scale: string, sessionCount: number): SessionRepoCatalogBenchmarkDataset {
|
||||
return {
|
||||
name: `synthetic catalog: ${scale} closed sessions`,
|
||||
sessionCount,
|
||||
};
|
||||
}
|
||||
|
||||
/** Deterministic closed-session catalogs shared by repository measurements. */
|
||||
export const SESSION_REPO_CATALOG_BENCHMARK_DATASETS: readonly SessionRepoCatalogBenchmarkDataset[] = [
|
||||
createCatalogDataset("100", 100),
|
||||
createCatalogDataset("1k", 1_000),
|
||||
createCatalogDataset("10k", 10_000),
|
||||
];
|
||||
|
||||
/** Seeds one deterministic catalog and returns its durable metadata in creation order. */
|
||||
export async function seedSessionRepoCatalogBenchmark(
|
||||
repo: SessionRepo,
|
||||
dataset: SessionRepoCatalogBenchmarkDataset,
|
||||
): Promise<SessionMetadata[]> {
|
||||
const metadata: SessionMetadata[] = [];
|
||||
for (let index = 0; index < dataset.sessionCount; index++) {
|
||||
const session = await repo.create({ id: sessionRepoBenchmarkSessionId(index) }, BACKGROUND_CONTEXT);
|
||||
metadata.push(session.metadata);
|
||||
await session.close(BACKGROUND_CONTEXT);
|
||||
}
|
||||
return metadata;
|
||||
}
|
||||
|
||||
/** Shared catalog reads. Returning a number ensures each result is consumed. */
|
||||
export const SESSION_REPO_CATALOG_READ_BENCHMARK_SCENARIOS: readonly SessionRepoCatalogReadBenchmarkScenario[] = [
|
||||
{
|
||||
name: "list sessions",
|
||||
expectedResult(dataset) {
|
||||
return dataset.sessionCount;
|
||||
},
|
||||
async run(repo) {
|
||||
return (await repo.list(undefined, BACKGROUND_CONTEXT)).length;
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
/** Shared catalog writes. Every invocation receives an equivalent independently prepared repository. */
|
||||
export const SESSION_REPO_CATALOG_WRITE_BENCHMARK_SCENARIOS: readonly SessionRepoCatalogWriteBenchmarkScenario[] = [
|
||||
{
|
||||
name: "create empty session",
|
||||
expectedResult: 1,
|
||||
prepare(repo) {
|
||||
return Promise.resolve({
|
||||
async run() {
|
||||
const session = await repo.create({ id: BENCHMARK_SESSION_ID }, BACKGROUND_CONTEXT);
|
||||
return session.metadata.id === BENCHMARK_SESSION_ID ? 1 : 0;
|
||||
},
|
||||
});
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "open closed empty session",
|
||||
expectedResult: 1,
|
||||
async prepare(repo) {
|
||||
const session = await repo.create({ id: BENCHMARK_SESSION_ID }, BACKGROUND_CONTEXT);
|
||||
const { metadata } = session;
|
||||
await session.close(BACKGROUND_CONTEXT);
|
||||
return {
|
||||
async run() {
|
||||
const reopened = await repo.open(metadata, BACKGROUND_CONTEXT);
|
||||
return reopened.metadata.id === metadata.id ? 1 : 0;
|
||||
},
|
||||
};
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "delete closed empty session",
|
||||
expectedResult: 1,
|
||||
async prepare(repo) {
|
||||
const session = await repo.create({ id: BENCHMARK_SESSION_ID }, BACKGROUND_CONTEXT);
|
||||
const { metadata } = session;
|
||||
await session.close(BACKGROUND_CONTEXT);
|
||||
return {
|
||||
async run() {
|
||||
await repo.delete(metadata, BACKGROUND_CONTEXT);
|
||||
return 1;
|
||||
},
|
||||
};
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
/** Initial fork timing uses a bounded source because each iteration owns an equivalent seeded repository. */
|
||||
export const SESSION_REPO_FORK_BENCHMARK_DATASETS: readonly StorageBenchmarkDataset[] = [
|
||||
STORAGE_BENCHMARK_DATASETS[0]!,
|
||||
STORAGE_BENCHMARK_DATASETS[1]!,
|
||||
];
|
||||
|
||||
/** Seeds one deterministic open source session with a linear main branch. */
|
||||
export async function seedSessionRepoForkBenchmark(
|
||||
repo: SessionRepo,
|
||||
dataset: StorageBenchmarkDataset,
|
||||
): Promise<SessionMetadata> {
|
||||
const session = await repo.create({ id: BENCHMARK_SESSION_ID }, BACKGROUND_CONTEXT);
|
||||
for (const transaction of generateStorageBenchmarkSeedTransactions(dataset)) {
|
||||
await session.mutate((mutator) => mutator.commit(transaction, BACKGROUND_CONTEXT), BACKGROUND_CONTEXT);
|
||||
}
|
||||
await session.mutate(
|
||||
(mutator) =>
|
||||
mutator.commit(
|
||||
[
|
||||
setValue(branchTip("main"), dataset.tipId),
|
||||
setValue(laneConfig("main"), {
|
||||
model: { provider: "benchmark", modelId: "benchmark" },
|
||||
thinkingLevel: "off",
|
||||
activeToolNames: [],
|
||||
}),
|
||||
setValue(laneState("main"), { currentOperationId: null, lastOperationId: null, inbox: [] }),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
return session.metadata;
|
||||
}
|
||||
|
||||
/** Shared fork writes. Every invocation receives an equivalent source repository. */
|
||||
export const SESSION_REPO_FORK_WRITE_BENCHMARK_SCENARIOS: readonly SessionRepoForkWriteBenchmarkScenario[] = [
|
||||
{
|
||||
name: "fork open current branch",
|
||||
expectedResult(dataset) {
|
||||
return dataset.entryCount;
|
||||
},
|
||||
async run(repo, source, dataset) {
|
||||
const fork = await repo.fork(
|
||||
source,
|
||||
{ id: FORK_DESTINATION_SESSION_ID, scope: "branch", branch: "main" },
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
return fork.metadata.id === FORK_DESTINATION_SESSION_ID && fork.metadata.parentSessionId === source.id
|
||||
? dataset.entryCount
|
||||
: 0;
|
||||
},
|
||||
},
|
||||
];
|
||||
143
third_party/pi/packages/agent/src/harness/session/testing/benchmark/storage.ts
vendored
Normal file
143
third_party/pi/packages/agent/src/harness/session/testing/benchmark/storage.ts
vendored
Normal file
@@ -0,0 +1,143 @@
|
||||
import { BACKGROUND_CONTEXT } from "../../../context.ts";
|
||||
import { insertEntry, insertUsage } from "../../commit.ts";
|
||||
import type { MessageEntry, NewEntry, Storage, Write } from "../../types.ts";
|
||||
import { branchTip, setValue } from "../../values.ts";
|
||||
import { STORAGE_BENCHMARK_DATASETS, type StorageBenchmarkDataset, storageBenchmarkEntryId } from "./datasets.ts";
|
||||
|
||||
const MESSAGE_TIMESTAMP = 1_650_000_000_000;
|
||||
const SEED_BATCH_SIZE = 250;
|
||||
const WRITE_BASELINE_DATASET = STORAGE_BENCHMARK_DATASETS[0]!;
|
||||
|
||||
function createEntry(index: number, payloadBytes: number): NewEntry<MessageEntry> {
|
||||
const id = storageBenchmarkEntryId(index);
|
||||
const prefix = `${id}:`;
|
||||
return {
|
||||
id,
|
||||
parentId: index === 0 ? null : storageBenchmarkEntryId(index - 1),
|
||||
type: "message",
|
||||
message: {
|
||||
role: "user",
|
||||
content: [{ type: "text", text: prefix + "x".repeat(Math.max(0, payloadBytes - prefix.length)) }],
|
||||
timestamp: MESSAGE_TIMESTAMP,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function createStorageBenchmarkTransaction(startIndex: number, entryCount: number, payloadBytes: number): Write[] {
|
||||
return Array.from({ length: entryCount }, (_, offset) =>
|
||||
insertEntry(createEntry(startIndex + offset, payloadBytes)),
|
||||
);
|
||||
}
|
||||
|
||||
/** Generates deterministic seed transactions without claiming a production data distribution. */
|
||||
export function* generateStorageBenchmarkSeedTransactions(dataset: StorageBenchmarkDataset): Generator<Write[]> {
|
||||
for (let startIndex = 0; startIndex < dataset.entryCount; startIndex += SEED_BATCH_SIZE) {
|
||||
yield createStorageBenchmarkTransaction(
|
||||
startIndex,
|
||||
Math.min(SEED_BATCH_SIZE, dataset.entryCount - startIndex),
|
||||
dataset.payloadBytes,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/** Seeds one deterministic synthetic linear branch. */
|
||||
export async function seedStorageBenchmark(storage: Storage, dataset: StorageBenchmarkDataset): Promise<void> {
|
||||
for (const transaction of generateStorageBenchmarkSeedTransactions(dataset))
|
||||
await storage.commit(transaction, BACKGROUND_CONTEXT);
|
||||
}
|
||||
|
||||
/** A steady-state read operation run against a pre-seeded fixture. */
|
||||
interface StorageReadBenchmarkScenario {
|
||||
readonly name: string;
|
||||
expectedResult(dataset: StorageBenchmarkDataset): number;
|
||||
run(storage: Storage, dataset: StorageBenchmarkDataset): Promise<number>;
|
||||
}
|
||||
|
||||
/** A write operation run once against each independently prepared fixture. */
|
||||
interface StorageWriteBenchmarkScenario {
|
||||
readonly name: string;
|
||||
readonly writeCount: number;
|
||||
prepare?(storage: Storage): Promise<void>;
|
||||
run(storage: Storage): Promise<number>;
|
||||
}
|
||||
|
||||
const singleEntryTransaction = createStorageBenchmarkTransaction(0, 1, 256);
|
||||
const hundredEntryTransaction = createStorageBenchmarkTransaction(0, 100, 256);
|
||||
const appendedEntryId = storageBenchmarkEntryId(WRITE_BASELINE_DATASET.entryCount);
|
||||
const mixedAppendTransaction: Write[] = [
|
||||
...createStorageBenchmarkTransaction(WRITE_BASELINE_DATASET.entryCount, 1, WRITE_BASELINE_DATASET.payloadBytes),
|
||||
setValue(branchTip("main"), appendedEntryId),
|
||||
insertUsage({
|
||||
id: "benchmark-usage",
|
||||
entryId: appendedEntryId,
|
||||
adjustment: false,
|
||||
usage: {
|
||||
input: 1_000,
|
||||
output: 250,
|
||||
cacheRead: 500,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 1_250,
|
||||
cost: { input: 0.001, output: 0.001, cacheRead: 0.0001, cacheWrite: 0, total: 0.0021 },
|
||||
},
|
||||
}),
|
||||
];
|
||||
|
||||
/** Shared read scenarios. Returning a number ensures each result is consumed. */
|
||||
export const STORAGE_READ_BENCHMARK_SCENARIOS: readonly StorageReadBenchmarkScenario[] = [
|
||||
{
|
||||
name: "get 100 distributed entries",
|
||||
expectedResult(dataset) {
|
||||
return dataset.lookupIds.length;
|
||||
},
|
||||
async run(storage, dataset) {
|
||||
return (await storage.getEntries([...dataset.lookupIds], BACKGROUND_CONTEXT)).size;
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "scan latest 50 entries",
|
||||
expectedResult(dataset) {
|
||||
return Math.min(50, dataset.entryCount);
|
||||
},
|
||||
async run(storage) {
|
||||
return (await storage.scanEntries({ order: "desc", limit: 50 }, BACKGROUND_CONTEXT)).length;
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "scan full branch structure",
|
||||
expectedResult(dataset) {
|
||||
return dataset.entryCount;
|
||||
},
|
||||
async run(storage, dataset) {
|
||||
return (await storage.scanBranchStructure({ start: dataset.tipId, order: "newestFirst" }, BACKGROUND_CONTEXT))
|
||||
.length;
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
/** Shared writes. Every invocation receives equivalent pre-benchmark state. */
|
||||
export const STORAGE_WRITE_BENCHMARK_SCENARIOS: readonly StorageWriteBenchmarkScenario[] = [
|
||||
{
|
||||
name: "commit one message entry",
|
||||
writeCount: 1,
|
||||
async run(storage) {
|
||||
return (await storage.commit(singleEntryTransaction, BACKGROUND_CONTEXT)).seqs.length;
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "commit 100 message entries",
|
||||
writeCount: 100,
|
||||
async run(storage) {
|
||||
return (await storage.commit(hundredEntryTransaction, BACKGROUND_CONTEXT)).seqs.length;
|
||||
},
|
||||
},
|
||||
{
|
||||
name: `commit mixed append (${WRITE_BASELINE_DATASET.name})`,
|
||||
writeCount: 3,
|
||||
prepare(storage) {
|
||||
return seedStorageBenchmark(storage, WRITE_BASELINE_DATASET);
|
||||
},
|
||||
async run(storage) {
|
||||
return (await storage.commit(mixedAppendTransaction, BACKGROUND_CONTEXT)).seqs.length;
|
||||
},
|
||||
},
|
||||
];
|
||||
846
third_party/pi/packages/agent/src/harness/session/testing/conformance/session-repo.ts
vendored
Normal file
846
third_party/pi/packages/agent/src/harness/session/testing/conformance/session-repo.ts
vendored
Normal file
@@ -0,0 +1,846 @@
|
||||
import { deepStrictEqual, rejects, strictEqual } from "node:assert/strict";
|
||||
import type { AssistantMessage, StopReason } from "@earendil-works/pi-ai";
|
||||
import { BACKGROUND_CONTEXT } from "../../../context.ts";
|
||||
import { insertEntry, insertUsage } from "../../commit.ts";
|
||||
import type {
|
||||
JsonValue,
|
||||
LaneConfiguration,
|
||||
LaneState,
|
||||
Session,
|
||||
SessionMetadata,
|
||||
SessionRepo,
|
||||
UsageRow,
|
||||
} from "../../types.ts";
|
||||
import {
|
||||
appendList,
|
||||
branchTip,
|
||||
entryLabel,
|
||||
laneConfig,
|
||||
laneState,
|
||||
list,
|
||||
operationMeta,
|
||||
operationPreparation,
|
||||
operationResult,
|
||||
operationState,
|
||||
operationToolArgs,
|
||||
pendingEntry,
|
||||
sessionName,
|
||||
setValue,
|
||||
value,
|
||||
} from "../../values.ts";
|
||||
import type { ConformanceCase } from "../types.ts";
|
||||
|
||||
const ROOT_ID = "00000000-0000-7000-8000-000000000001";
|
||||
const CHILD_ID = "00000000-0000-7000-8000-000000000002";
|
||||
const SIBLING_ID = "00000000-0000-7000-8000-000000000003";
|
||||
const USAGE_ID = "00000000-0000-7000-8000-000000000004";
|
||||
const OPERATION_ID = "00000000-0000-7000-8000-000000000005";
|
||||
const PENDING_ID = "00000000-0000-7000-8000-000000000006";
|
||||
const UNKNOWN_ID = "00000000-0000-7000-8000-000000000007";
|
||||
const idleLaneState = {
|
||||
currentOperationId: null,
|
||||
lastOperationId: null,
|
||||
inbox: [],
|
||||
} satisfies LaneState;
|
||||
const applicationValue = value<JsonValue>("test.application.value");
|
||||
const applicationList = list<JsonValue>("test.application.list");
|
||||
const configuration = {
|
||||
model: { provider: "provider", modelId: "model" },
|
||||
thinkingLevel: "off",
|
||||
activeToolNames: ["read"],
|
||||
} satisfies LaneConfiguration;
|
||||
|
||||
async function getBranchTip(session: Session, name = "main"): Promise<string | null | undefined> {
|
||||
return (await session.branch(name, BACKGROUND_CONTEXT))?.getTipId(BACKGROUND_CONTEXT);
|
||||
}
|
||||
|
||||
function assistantMessage(stopReason: StopReason): AssistantMessage {
|
||||
return {
|
||||
role: "assistant",
|
||||
content:
|
||||
stopReason === "toolUse"
|
||||
? [{ type: "toolCall", id: "call", name: "read", arguments: {} }]
|
||||
: [{ type: "text", text: stopReason }],
|
||||
api: "anthropic-messages",
|
||||
provider: "anthropic",
|
||||
model: "claude-sonnet-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,
|
||||
...(stopReason === "deferred"
|
||||
? {
|
||||
deferred: {
|
||||
provider: "anthropic",
|
||||
modelId: "claude-sonnet-4-5",
|
||||
api: "anthropic-messages",
|
||||
id: "job",
|
||||
},
|
||||
}
|
||||
: {}),
|
||||
timestamp: 1,
|
||||
};
|
||||
}
|
||||
|
||||
function usageRow(): Omit<UsageRow, "seq"> {
|
||||
return {
|
||||
id: USAGE_ID,
|
||||
adjustment: true,
|
||||
usage: {
|
||||
input: 1,
|
||||
output: 2,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 3,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
interface RepoCaseContext<TRepo> {
|
||||
repo: TRepo;
|
||||
close?: () => void | Promise<void>;
|
||||
}
|
||||
|
||||
function prepareRepoCaseFactory<TRepo>(
|
||||
factory: () => Promise<TRepo>,
|
||||
onClose?: () => void | Promise<void>,
|
||||
): () => Promise<RepoCaseContext<TRepo>> {
|
||||
return async () => ({
|
||||
repo: await factory(),
|
||||
...(onClose === undefined ? {} : { close: onClose }),
|
||||
});
|
||||
}
|
||||
|
||||
function createCase<TRepo>(
|
||||
factory: () => Promise<RepoCaseContext<TRepo>>,
|
||||
group: string,
|
||||
name: string,
|
||||
test: (context: { repo: TRepo }) => Promise<void>,
|
||||
): ConformanceCase {
|
||||
return {
|
||||
group,
|
||||
name,
|
||||
async run() {
|
||||
const context = await factory();
|
||||
try {
|
||||
await test(context);
|
||||
} finally {
|
||||
await context.close?.();
|
||||
}
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
/** Creates lifecycle cases for repositories that support creation, discovery, open, and deletion. */
|
||||
export function createSessionRepoLifecycleConformance<TMetadata extends SessionMetadata>(
|
||||
// Require only the methods this group exercises so partial backends can run it independently.
|
||||
backendFactory: () => Promise<Pick<SessionRepo<TMetadata>, "create" | "open" | "list" | "delete">>,
|
||||
onClose?: () => void | Promise<void>,
|
||||
): readonly ConformanceCase[] {
|
||||
const factory = prepareRepoCaseFactory(backendFactory, onClose);
|
||||
return [
|
||||
createCase(
|
||||
factory,
|
||||
"lifecycle",
|
||||
"creates a session with no implicit branch and rejects duplicate ids",
|
||||
async ({ repo }) => {
|
||||
const session = await repo.create({ id: "session" }, BACKGROUND_CONTEXT);
|
||||
|
||||
strictEqual(session.metadata.id, "session");
|
||||
strictEqual(Number.isSafeInteger(session.metadata.createdAt), true);
|
||||
strictEqual(session.metadata.storageVersion, 1);
|
||||
strictEqual(await session.branch("main", BACKGROUND_CONTEXT), undefined);
|
||||
strictEqual(await session.getValue(laneState("main"), BACKGROUND_CONTEXT), undefined);
|
||||
strictEqual(await session.getValue(laneConfig("main"), BACKGROUND_CONTEXT), undefined);
|
||||
await rejects(repo.create({ id: "session" }, BACKGROUND_CONTEXT));
|
||||
await session.close(BACKGROUND_CONTEXT);
|
||||
},
|
||||
),
|
||||
createCase(
|
||||
factory,
|
||||
"lifecycle",
|
||||
"close drains an acquired scope and rejects a queued mutation callback",
|
||||
async ({ repo }) => {
|
||||
const session = await repo.create({ id: "session" }, BACKGROUND_CONTEXT);
|
||||
const active = await session.beginMutation(BACKGROUND_CONTEXT);
|
||||
let queuedStarted = false;
|
||||
const queued = session.mutate(() => {
|
||||
queuedStarted = true;
|
||||
}, BACKGROUND_CONTEXT);
|
||||
void queued.catch(() => {});
|
||||
const closing = session.close(BACKGROUND_CONTEXT);
|
||||
|
||||
await active.end(BACKGROUND_CONTEXT);
|
||||
await rejects(queued);
|
||||
await closing;
|
||||
strictEqual(queuedStarted, false);
|
||||
},
|
||||
),
|
||||
createCase(
|
||||
factory,
|
||||
"lifecycle",
|
||||
"lists metadata and preserves state across close and reopen",
|
||||
async ({ repo }) => {
|
||||
const first = await repo.create({ id: "first" }, BACKGROUND_CONTEXT);
|
||||
await first.setName("preserved", BACKGROUND_CONTEXT);
|
||||
const second = await repo.create({ id: "second", parentSessionId: "parent" }, BACKGROUND_CONTEXT);
|
||||
await second.close(BACKGROUND_CONTEXT);
|
||||
|
||||
const listed = await repo.list(undefined, BACKGROUND_CONTEXT);
|
||||
deepStrictEqual(
|
||||
listed
|
||||
.map(({ id, parentSessionId }) => ({ id, parentSessionId }))
|
||||
.sort((left, right) => left.id.localeCompare(right.id)),
|
||||
[
|
||||
{ id: "first", parentSessionId: undefined },
|
||||
{ id: "second", parentSessionId: "parent" },
|
||||
],
|
||||
);
|
||||
await first.close(BACKGROUND_CONTEXT);
|
||||
await rejects(first.getName(BACKGROUND_CONTEXT));
|
||||
const reopened = await repo.open(first.metadata, BACKGROUND_CONTEXT);
|
||||
strictEqual(reopened === first, false);
|
||||
strictEqual(await reopened.getName(BACKGROUND_CONTEXT), "preserved");
|
||||
await reopened.close(BACKGROUND_CONTEXT);
|
||||
},
|
||||
),
|
||||
createCase(factory, "lifecycle", "deletes closed sessions without affecting other sessions", async ({ repo }) => {
|
||||
const removed = await repo.create({ id: "removed" }, BACKGROUND_CONTEXT);
|
||||
const retained = await repo.create({ id: "retained" }, BACKGROUND_CONTEXT);
|
||||
await Promise.all([removed.close(BACKGROUND_CONTEXT), retained.close(BACKGROUND_CONTEXT)]);
|
||||
|
||||
await repo.delete(removed.metadata, BACKGROUND_CONTEXT);
|
||||
deepStrictEqual(
|
||||
(await repo.list(undefined, BACKGROUND_CONTEXT)).map(({ id }) => id),
|
||||
["retained"],
|
||||
);
|
||||
await rejects(repo.open(removed.metadata, BACKGROUND_CONTEXT));
|
||||
await rejects(repo.delete(removed.metadata, BACKGROUND_CONTEXT));
|
||||
}),
|
||||
];
|
||||
}
|
||||
|
||||
/** Creates exclusive-open cases for repositories that own active session handles. */
|
||||
export function createSessionRepoOwnershipConformance<TMetadata extends SessionMetadata>(
|
||||
backendFactory: () => Promise<Pick<SessionRepo<TMetadata>, "create" | "open">>,
|
||||
onClose?: () => void | Promise<void>,
|
||||
): readonly ConformanceCase[] {
|
||||
const factory = prepareRepoCaseFactory(backendFactory, onClose);
|
||||
return [
|
||||
createCase(factory, "ownership", "rejects opening an already-open session", async ({ repo }) => {
|
||||
const session = await repo.create({ id: "session" }, BACKGROUND_CONTEXT);
|
||||
await rejects(repo.open(session.metadata, BACKGROUND_CONTEXT));
|
||||
await session.close(BACKGROUND_CONTEXT);
|
||||
|
||||
const reopened = await repo.open(session.metadata, BACKGROUND_CONTEXT);
|
||||
await rejects(repo.open(session.metadata, BACKGROUND_CONTEXT));
|
||||
await reopened.close(BACKGROUND_CONTEXT);
|
||||
}),
|
||||
];
|
||||
}
|
||||
|
||||
/** Creates message cases for repositories that support session creation. */
|
||||
export function createSessionRepoMessageConformance<TMetadata extends SessionMetadata>(
|
||||
backendFactory: () => Promise<Pick<SessionRepo<TMetadata>, "create">>,
|
||||
onClose?: () => void | Promise<void>,
|
||||
): readonly ConformanceCase[] {
|
||||
const factory = prepareRepoCaseFactory(backendFactory, onClose);
|
||||
return [
|
||||
createCase(
|
||||
factory,
|
||||
"messages",
|
||||
"rejects pending assistant messages without changing the tree",
|
||||
async ({ repo }) => {
|
||||
const session = await repo.create({ id: "session" }, BACKGROUND_CONTEXT);
|
||||
const branch = await session.createBranch("main", null, BACKGROUND_CONTEXT);
|
||||
|
||||
await rejects(branch.appendMessage(assistantMessage("pending"), BACKGROUND_CONTEXT));
|
||||
|
||||
strictEqual(await getBranchTip(session), null);
|
||||
deepStrictEqual(await session.findEntries(undefined, BACKGROUND_CONTEXT), []);
|
||||
await session.close(BACKGROUND_CONTEXT);
|
||||
},
|
||||
),
|
||||
createCase(factory, "messages", "preserves every settled assistant stop reason", async ({ repo }) => {
|
||||
const session = await repo.create({ id: "session" }, BACKGROUND_CONTEXT);
|
||||
const messagesByStopReason = {
|
||||
stop: assistantMessage("stop"),
|
||||
length: assistantMessage("length"),
|
||||
toolUse: assistantMessage("toolUse"),
|
||||
error: assistantMessage("error"),
|
||||
aborted: assistantMessage("aborted"),
|
||||
deferred: assistantMessage("deferred"),
|
||||
} satisfies Record<Exclude<StopReason, "pending">, AssistantMessage>;
|
||||
const branch = await session.createBranch("main", null, BACKGROUND_CONTEXT);
|
||||
const messages = Object.values(messagesByStopReason);
|
||||
const ids: string[] = [];
|
||||
|
||||
for (const message of messages) ids.push(await branch.appendMessage(message, BACKGROUND_CONTEXT));
|
||||
|
||||
const entries = await session.findEntries({ order: "asc", type: "message" }, BACKGROUND_CONTEXT);
|
||||
deepStrictEqual(
|
||||
entries.map((entry) => entry.id),
|
||||
ids,
|
||||
);
|
||||
for (const [index, entry] of entries.entries()) {
|
||||
if (entry.type !== "message") throw new Error("Expected message entry");
|
||||
deepStrictEqual(entry.message, messages[index]);
|
||||
}
|
||||
strictEqual(await getBranchTip(session), ids.at(-1));
|
||||
await session.close(BACKGROUND_CONTEXT);
|
||||
}),
|
||||
];
|
||||
}
|
||||
|
||||
/** Creates fork-content cases that do not require concurrent repository coordination. */
|
||||
export function createSessionRepoForkBehaviorConformance<TMetadata extends SessionMetadata>(
|
||||
backendFactory: () => Promise<Pick<SessionRepo<TMetadata>, "create" | "list" | "fork">>,
|
||||
onClose?: () => void | Promise<void>,
|
||||
): readonly ConformanceCase[] {
|
||||
const factory = prepareRepoCaseFactory(backendFactory, onClose);
|
||||
return [
|
||||
createCase(factory, "forks", "tree-forks a fresh session before first attachment", async ({ repo }) => {
|
||||
const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT);
|
||||
const fork = await repo.fork(source.metadata, { id: "fork", scope: "tree" }, BACKGROUND_CONTEXT);
|
||||
|
||||
strictEqual(fork.metadata.id, "fork");
|
||||
strictEqual(fork.metadata.parentSessionId, "source");
|
||||
strictEqual(await fork.branch("main", BACKGROUND_CONTEXT), undefined);
|
||||
strictEqual(await fork.getValue(laneConfig("main"), BACKGROUND_CONTEXT), undefined);
|
||||
strictEqual(await fork.getValue(laneState("main"), BACKGROUND_CONTEXT), undefined);
|
||||
deepStrictEqual(await fork.findEntries(undefined, BACKGROUND_CONTEXT), []);
|
||||
deepStrictEqual(await fork.getStats(BACKGROUND_CONTEXT), {
|
||||
messageCount: 0,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
total: 0,
|
||||
},
|
||||
},
|
||||
});
|
||||
await Promise.all([source.close(BACKGROUND_CONTEXT), fork.close(BACKGROUND_CONTEXT)]);
|
||||
}),
|
||||
createCase(factory, "forks", "rejects a data-only branch and releases its destination id", async ({ repo }) => {
|
||||
const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT);
|
||||
await source.createBranch("data", null, BACKGROUND_CONTEXT);
|
||||
|
||||
await rejects(
|
||||
repo.fork(source.metadata, { id: "destination", scope: "branch", branch: "data" }, BACKGROUND_CONTEXT),
|
||||
);
|
||||
deepStrictEqual(
|
||||
(await repo.list(undefined, BACKGROUND_CONTEXT)).map(({ id }) => id),
|
||||
["source"],
|
||||
);
|
||||
|
||||
const destination = await repo.create({ id: "destination" }, BACKGROUND_CONTEXT);
|
||||
await Promise.all([source.close(BACKGROUND_CONTEXT), destination.close(BACKGROUND_CONTEXT)]);
|
||||
}),
|
||||
createCase(
|
||||
factory,
|
||||
"forks",
|
||||
"forks one named configured branch with scoped values and a zero ledger",
|
||||
async ({ repo }) => {
|
||||
const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT);
|
||||
await source.mutate(
|
||||
(mutator) =>
|
||||
mutator.commit(
|
||||
[
|
||||
insertEntry({
|
||||
id: ROOT_ID,
|
||||
parentId: null,
|
||||
type: "custom",
|
||||
customType: "root",
|
||||
}),
|
||||
insertEntry({
|
||||
id: CHILD_ID,
|
||||
parentId: ROOT_ID,
|
||||
type: "message",
|
||||
message: { role: "user", content: "child", timestamp: 1 },
|
||||
}),
|
||||
insertEntry({
|
||||
id: SIBLING_ID,
|
||||
parentId: ROOT_ID,
|
||||
type: "custom",
|
||||
customType: "sibling",
|
||||
}),
|
||||
setValue(branchTip("main"), SIBLING_ID),
|
||||
setValue(branchTip("review"), CHILD_ID),
|
||||
setValue(laneConfig("review"), configuration),
|
||||
setValue(laneState("review"), {
|
||||
currentOperationId: OPERATION_ID,
|
||||
lastOperationId: "previous",
|
||||
inbox: [{ entryId: PENDING_ID, kind: "write" }],
|
||||
}),
|
||||
setValue(operationResult("previous"), {
|
||||
operationId: "previous",
|
||||
kind: "navigation",
|
||||
status: "completed",
|
||||
fromTipId: ROOT_ID,
|
||||
tipId: CHILD_ID,
|
||||
startedAt: 1,
|
||||
endedAt: 2,
|
||||
}),
|
||||
setValue(sessionName, "source name"),
|
||||
setValue(applicationValue, { copied: false }),
|
||||
appendList(applicationList, { copied: false }),
|
||||
setValue(entryLabel(ROOT_ID), "root label"),
|
||||
setValue(entryLabel(SIBLING_ID), "sibling label"),
|
||||
setValue(pendingEntry(PENDING_ID), {
|
||||
type: "custom",
|
||||
customType: "pending",
|
||||
}),
|
||||
setValue(operationMeta(OPERATION_ID), {
|
||||
operationId: OPERATION_ID,
|
||||
lane: "review",
|
||||
sourceTipId: CHILD_ID,
|
||||
startedAt: 1,
|
||||
intent: { kind: "compaction" },
|
||||
}),
|
||||
setValue(operationState(OPERATION_ID), {
|
||||
at: "summary.deciding",
|
||||
control: { status: "running" },
|
||||
settings: {
|
||||
compaction: { enabled: true, reserveTokens: 1, keepRecentTokens: 1 },
|
||||
steeringMode: "all",
|
||||
followUpMode: "all",
|
||||
toolExecution: "sequential",
|
||||
},
|
||||
latestAssistantEntryId: null,
|
||||
task: { taskId: OPERATION_ID, reason: "manual", boundary: { kind: "finish" } },
|
||||
}),
|
||||
setValue(operationToolArgs(OPERATION_ID, ROOT_ID, 0), {
|
||||
argument: true,
|
||||
}),
|
||||
setValue(operationPreparation(OPERATION_ID, OPERATION_ID), {
|
||||
kind: "compaction",
|
||||
messagesToSummarize: [],
|
||||
turnPrefixMessages: [],
|
||||
retainedTail: [],
|
||||
isSplitTurn: false,
|
||||
tokensBefore: 0,
|
||||
fileOps: { read: [], written: [], edited: [] },
|
||||
settings: {
|
||||
enabled: true,
|
||||
reserveTokens: 1,
|
||||
keepRecentTokens: 1,
|
||||
},
|
||||
}),
|
||||
insertUsage(usageRow()),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
const fork = await repo.fork(
|
||||
source.metadata,
|
||||
{ id: "fork", scope: "branch", branch: "review", entryId: CHILD_ID, position: "at" },
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
deepStrictEqual(
|
||||
(await fork.findEntries({ order: "asc" }, BACKGROUND_CONTEXT)).map(({ id }) => id),
|
||||
[ROOT_ID, CHILD_ID],
|
||||
);
|
||||
strictEqual(await fork.branch("main", BACKGROUND_CONTEXT), undefined);
|
||||
strictEqual(await getBranchTip(fork, "review"), CHILD_ID);
|
||||
deepStrictEqual((await fork.getValue(laneConfig("review"), BACKGROUND_CONTEXT))?.value, configuration);
|
||||
deepStrictEqual((await fork.getValue(laneState("review"), BACKGROUND_CONTEXT))?.value, idleLaneState);
|
||||
strictEqual(await fork.getValue(laneConfig("main"), BACKGROUND_CONTEXT), undefined);
|
||||
strictEqual(await fork.getValue(laneState("main"), BACKGROUND_CONTEXT), undefined);
|
||||
strictEqual(await fork.getName(BACKGROUND_CONTEXT), "source name");
|
||||
strictEqual(await fork.getValue(applicationValue, BACKGROUND_CONTEXT), undefined);
|
||||
deepStrictEqual(await fork.readList(applicationList, undefined, BACKGROUND_CONTEXT), []);
|
||||
strictEqual(await fork.getLabel(ROOT_ID, BACKGROUND_CONTEXT), "root label");
|
||||
strictEqual(await fork.getLabel(SIBLING_ID, BACKGROUND_CONTEXT), undefined);
|
||||
strictEqual(await fork.getValue(operationResult("previous"), BACKGROUND_CONTEXT), undefined);
|
||||
strictEqual(await fork.getValue(pendingEntry(PENDING_ID), BACKGROUND_CONTEXT), undefined);
|
||||
strictEqual(await fork.getValue(operationMeta(OPERATION_ID), BACKGROUND_CONTEXT), undefined);
|
||||
strictEqual(await fork.getValue(operationState(OPERATION_ID), BACKGROUND_CONTEXT), undefined);
|
||||
strictEqual(
|
||||
await fork.getValue(operationToolArgs(OPERATION_ID, ROOT_ID, 0), BACKGROUND_CONTEXT),
|
||||
undefined,
|
||||
);
|
||||
strictEqual(
|
||||
await fork.getValue(operationPreparation(OPERATION_ID, OPERATION_ID), BACKGROUND_CONTEXT),
|
||||
undefined,
|
||||
);
|
||||
const stats = await fork.getStats(BACKGROUND_CONTEXT);
|
||||
strictEqual(stats.messageCount, 1);
|
||||
deepStrictEqual(stats.usage, {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
});
|
||||
await Promise.all([source.close(BACKGROUND_CONTEXT), fork.close(BACKGROUND_CONTEXT)]);
|
||||
},
|
||||
),
|
||||
createCase(factory, "forks", "enforces branch ancestry for at and before placement", async ({ repo }) => {
|
||||
const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT);
|
||||
await source.mutate(
|
||||
(mutator) =>
|
||||
mutator.commit(
|
||||
[
|
||||
insertEntry({ id: ROOT_ID, parentId: null, type: "custom", customType: "root" }),
|
||||
insertEntry({ id: CHILD_ID, parentId: ROOT_ID, type: "custom", customType: "child" }),
|
||||
insertEntry({ id: SIBLING_ID, parentId: ROOT_ID, type: "custom", customType: "sibling" }),
|
||||
setValue(branchTip("main"), CHILD_ID),
|
||||
setValue(laneConfig("main"), configuration),
|
||||
setValue(laneState("main"), idleLaneState),
|
||||
setValue(branchTip("empty"), null),
|
||||
setValue(laneConfig("empty"), configuration),
|
||||
setValue(laneState("empty"), idleLaneState),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
const before = await repo.fork(
|
||||
source.metadata,
|
||||
{ id: "before", scope: "branch", branch: "main", entryId: CHILD_ID, position: "before" },
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
strictEqual(await getBranchTip(before), ROOT_ID);
|
||||
deepStrictEqual(
|
||||
(await before.findEntries({ order: "asc" }, BACKGROUND_CONTEXT)).map(({ id }) => id),
|
||||
[ROOT_ID],
|
||||
);
|
||||
|
||||
const mid = await repo.fork(
|
||||
source.metadata,
|
||||
{ id: "mid", scope: "branch", branch: "main", entryId: ROOT_ID, position: "at" },
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
strictEqual(await getBranchTip(mid), ROOT_ID);
|
||||
deepStrictEqual(
|
||||
(await mid.findEntries({ order: "asc" }, BACKGROUND_CONTEXT)).map(({ id }) => id),
|
||||
[ROOT_ID],
|
||||
);
|
||||
|
||||
const beforeRoot = await repo.fork(
|
||||
source.metadata,
|
||||
{ id: "before-root", scope: "branch", branch: "main", entryId: ROOT_ID, position: "before" },
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
strictEqual(await getBranchTip(beforeRoot), null);
|
||||
deepStrictEqual(await beforeRoot.findEntries({ order: "asc" }, BACKGROUND_CONTEXT), []);
|
||||
|
||||
const empty = await repo.fork(
|
||||
source.metadata,
|
||||
{ id: "empty", scope: "branch", branch: "empty" },
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
strictEqual(await getBranchTip(empty, "empty"), null);
|
||||
deepStrictEqual(await empty.findEntries({ order: "asc" }, BACKGROUND_CONTEXT), []);
|
||||
|
||||
for (const [id, branch, entryId] of [
|
||||
["off-branch", "main", SIBLING_ID],
|
||||
["unknown", "main", UNKNOWN_ID],
|
||||
["null-tip", "empty", ROOT_ID],
|
||||
] as const) {
|
||||
await rejects(repo.fork(source.metadata, { id, scope: "branch", branch, entryId }, BACKGROUND_CONTEXT));
|
||||
}
|
||||
deepStrictEqual((await repo.list(undefined, BACKGROUND_CONTEXT)).map(({ id }) => id).sort(), [
|
||||
"before",
|
||||
"before-root",
|
||||
"empty",
|
||||
"mid",
|
||||
"source",
|
||||
]);
|
||||
await Promise.all([
|
||||
source.close(BACKGROUND_CONTEXT),
|
||||
before.close(BACKGROUND_CONTEXT),
|
||||
beforeRoot.close(BACKGROUND_CONTEXT),
|
||||
empty.close(BACKGROUND_CONTEXT),
|
||||
mid.close(BACKGROUND_CONTEXT),
|
||||
]);
|
||||
}),
|
||||
createCase(factory, "forks", "forks a closed source session", async ({ repo }) => {
|
||||
const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT);
|
||||
await source.mutate(
|
||||
(mutator) =>
|
||||
mutator.commit(
|
||||
[
|
||||
insertEntry({
|
||||
id: ROOT_ID,
|
||||
parentId: null,
|
||||
type: "custom",
|
||||
customType: "root",
|
||||
}),
|
||||
setValue(branchTip("main"), ROOT_ID),
|
||||
setValue(laneConfig("main"), configuration),
|
||||
setValue(laneState("main"), idleLaneState),
|
||||
setValue(applicationValue, "excluded"),
|
||||
appendList(applicationList, "excluded"),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
await source.close(BACKGROUND_CONTEXT);
|
||||
|
||||
const fork = await repo.fork(
|
||||
source.metadata,
|
||||
{ id: "fork", scope: "branch", branch: "main" },
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
strictEqual(await getBranchTip(fork), ROOT_ID);
|
||||
deepStrictEqual((await fork.getValue(laneConfig("main"), BACKGROUND_CONTEXT))?.value, configuration);
|
||||
deepStrictEqual((await fork.getValue(laneState("main"), BACKGROUND_CONTEXT))?.value, idleLaneState);
|
||||
strictEqual(await fork.getValue(applicationValue, BACKGROUND_CONTEXT), undefined);
|
||||
deepStrictEqual(await fork.readList(applicationList, undefined, BACKGROUND_CONTEXT), []);
|
||||
await fork.close(BACKGROUND_CONTEXT);
|
||||
}),
|
||||
|
||||
createCase(factory, "forks", "forks the whole configured tree with fresh lane state", async ({ repo }) => {
|
||||
const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT);
|
||||
await source.mutate(
|
||||
(mutator) =>
|
||||
mutator.commit(
|
||||
[
|
||||
insertEntry({
|
||||
id: ROOT_ID,
|
||||
parentId: null,
|
||||
type: "custom",
|
||||
customType: "root",
|
||||
}),
|
||||
insertEntry({
|
||||
id: CHILD_ID,
|
||||
parentId: ROOT_ID,
|
||||
type: "custom",
|
||||
customType: "child",
|
||||
}),
|
||||
insertEntry({
|
||||
id: SIBLING_ID,
|
||||
parentId: ROOT_ID,
|
||||
type: "custom",
|
||||
customType: "sibling",
|
||||
}),
|
||||
setValue(branchTip("main"), CHILD_ID),
|
||||
setValue(laneConfig("main"), configuration),
|
||||
setValue(laneState("main"), idleLaneState),
|
||||
setValue(branchTip("review"), SIBLING_ID),
|
||||
setValue(laneConfig("review"), configuration),
|
||||
setValue(laneState("review"), idleLaneState),
|
||||
setValue(branchTip("notes"), ROOT_ID),
|
||||
setValue(applicationValue, { copied: true }),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
const fork = await repo.fork(source.metadata, { id: "fork", scope: "tree" }, BACKGROUND_CONTEXT);
|
||||
|
||||
deepStrictEqual(
|
||||
(await fork.findEntries({ order: "asc" }, BACKGROUND_CONTEXT)).map(({ id }) => id),
|
||||
[ROOT_ID, CHILD_ID, SIBLING_ID],
|
||||
);
|
||||
strictEqual(await getBranchTip(fork), CHILD_ID);
|
||||
strictEqual(await getBranchTip(fork, "review"), SIBLING_ID);
|
||||
strictEqual(await getBranchTip(fork, "notes"), ROOT_ID);
|
||||
strictEqual(await fork.getValue(laneConfig("notes"), BACKGROUND_CONTEXT), undefined);
|
||||
strictEqual(await fork.getValue(laneState("notes"), BACKGROUND_CONTEXT), undefined);
|
||||
deepStrictEqual((await fork.getValue(laneConfig("review"), BACKGROUND_CONTEXT))?.value, configuration);
|
||||
deepStrictEqual((await fork.getValue(laneState("review"), BACKGROUND_CONTEXT))?.value, idleLaneState);
|
||||
deepStrictEqual((await fork.getValue(applicationValue, BACKGROUND_CONTEXT))?.value, { copied: true });
|
||||
await Promise.all([source.close(BACKGROUND_CONTEXT), fork.close(BACKGROUND_CONTEXT)]);
|
||||
}),
|
||||
createCase(factory, "forks", "rejects only surviving unknown reserved scalar state", async ({ repo }) => {
|
||||
const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT);
|
||||
await source.createBranch("main", null, BACKGROUND_CONTEXT);
|
||||
await source.mutate(
|
||||
(mutator) =>
|
||||
mutator.commit(
|
||||
[setValue(laneConfig("main"), configuration), setValue(laneState("main"), idleLaneState)],
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
for (const namespace of ["pi", "pi.unknown"] as const) {
|
||||
const address = value<JsonValue>(namespace);
|
||||
await source.setValue(address, true, BACKGROUND_CONTEXT);
|
||||
await rejects(repo.fork(source.metadata, { id: "tree", scope: "tree" }, BACKGROUND_CONTEXT));
|
||||
await rejects(
|
||||
repo.fork(source.metadata, { id: "branch", scope: "branch", branch: "main" }, BACKGROUND_CONTEXT),
|
||||
);
|
||||
await source.deleteValue(address, BACKGROUND_CONTEXT);
|
||||
}
|
||||
await source.close(BACKGROUND_CONTEXT);
|
||||
|
||||
const tree = await repo.fork(source.metadata, { id: "tree", scope: "tree" }, BACKGROUND_CONTEXT);
|
||||
const branch = await repo.fork(
|
||||
source.metadata,
|
||||
{ id: "branch", scope: "branch", branch: "main" },
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
await Promise.all([tree.close(BACKGROUND_CONTEXT), branch.close(BACKGROUND_CONTEXT)]);
|
||||
}),
|
||||
];
|
||||
}
|
||||
|
||||
/** Creates fork cases that require destination reservation across create and fork. */
|
||||
export function createSessionRepoForkDestinationReservationConformance<TMetadata extends SessionMetadata>(
|
||||
backendFactory: () => Promise<Pick<SessionRepo<TMetadata>, "create" | "fork">>,
|
||||
onClose?: () => void | Promise<void>,
|
||||
): readonly ConformanceCase[] {
|
||||
const factory = prepareRepoCaseFactory(backendFactory, onClose);
|
||||
return [
|
||||
createCase(
|
||||
factory,
|
||||
"fork coordination",
|
||||
"publishes create when it reserves a shared destination id first",
|
||||
async ({ repo }) => {
|
||||
const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT);
|
||||
const results = await Promise.allSettled([
|
||||
repo.create({ id: "destination" }, BACKGROUND_CONTEXT),
|
||||
repo.fork(source.metadata, { id: "destination", scope: "tree" }, BACKGROUND_CONTEXT),
|
||||
]);
|
||||
strictEqual(results[0].status, "fulfilled");
|
||||
strictEqual(results[1].status, "rejected");
|
||||
if (results[0].status === "fulfilled") await results[0].value.close(BACKGROUND_CONTEXT);
|
||||
await source.close(BACKGROUND_CONTEXT);
|
||||
},
|
||||
),
|
||||
createCase(
|
||||
factory,
|
||||
"fork coordination",
|
||||
"publishes fork when it reserves a shared destination id first",
|
||||
async ({ repo }) => {
|
||||
const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT);
|
||||
const results = await Promise.allSettled([
|
||||
repo.fork(source.metadata, { id: "destination", scope: "tree" }, BACKGROUND_CONTEXT),
|
||||
repo.create({ id: "destination" }, BACKGROUND_CONTEXT),
|
||||
]);
|
||||
strictEqual(results[0].status, "fulfilled");
|
||||
strictEqual(results[1].status, "rejected");
|
||||
if (results[0].status === "fulfilled") await results[0].value.close(BACKGROUND_CONTEXT);
|
||||
await source.close(BACKGROUND_CONTEXT);
|
||||
},
|
||||
),
|
||||
];
|
||||
}
|
||||
|
||||
/** Creates fork cases that require a snapshot boundary on an active source storage queue. */
|
||||
export function createSessionRepoForkSourceSnapshotConformance<TMetadata extends SessionMetadata>(
|
||||
backendFactory: () => Promise<Pick<SessionRepo<TMetadata>, "create" | "fork">>,
|
||||
onClose?: () => void | Promise<void>,
|
||||
): readonly ConformanceCase[] {
|
||||
const factory = prepareRepoCaseFactory(backendFactory, onClose);
|
||||
return [
|
||||
createCase(
|
||||
factory,
|
||||
"fork coordination",
|
||||
"captures one coherent boundary between source commits",
|
||||
async ({ repo }) => {
|
||||
const source = await repo.create({ id: "source" }, BACKGROUND_CONTEXT);
|
||||
const firstMutation = await source.beginMutation(BACKGROUND_CONTEXT);
|
||||
const firstCommit = firstMutation.commit(
|
||||
[
|
||||
insertEntry({
|
||||
id: ROOT_ID,
|
||||
parentId: null,
|
||||
type: "custom",
|
||||
customType: "first",
|
||||
}),
|
||||
setValue(branchTip("main"), ROOT_ID),
|
||||
setValue(laneConfig("main"), configuration),
|
||||
setValue(laneState("main"), idleLaneState),
|
||||
setValue(sessionName, "first name"),
|
||||
setValue(entryLabel(ROOT_ID), "first label"),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
const fork = repo.fork(
|
||||
source.metadata,
|
||||
{ id: "fork", scope: "branch", branch: "main" },
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
const secondCommit = source.mutate(
|
||||
(mutator) =>
|
||||
mutator.commit(
|
||||
[
|
||||
insertEntry({
|
||||
id: CHILD_ID,
|
||||
parentId: ROOT_ID,
|
||||
type: "custom",
|
||||
customType: "second",
|
||||
}),
|
||||
setValue(branchTip("main"), CHILD_ID),
|
||||
setValue(sessionName, "second name"),
|
||||
setValue(entryLabel(ROOT_ID), "second label"),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
const [, forked] = await Promise.all([firstCommit, fork]);
|
||||
await firstMutation.end(BACKGROUND_CONTEXT);
|
||||
await secondCommit;
|
||||
strictEqual(await getBranchTip(forked), ROOT_ID);
|
||||
deepStrictEqual(
|
||||
(await forked.findEntries({ order: "asc" }, BACKGROUND_CONTEXT)).map(({ id }) => id),
|
||||
[ROOT_ID],
|
||||
);
|
||||
strictEqual(await forked.getName(BACKGROUND_CONTEXT), "first name");
|
||||
strictEqual(await forked.getLabel(ROOT_ID, BACKGROUND_CONTEXT), "first label");
|
||||
await Promise.all([source.close(BACKGROUND_CONTEXT), forked.close(BACKGROUND_CONTEXT)]);
|
||||
},
|
||||
),
|
||||
];
|
||||
}
|
||||
|
||||
/** Creates every fork coordination case. */
|
||||
export function createSessionRepoForkCoordinationConformance<TMetadata extends SessionMetadata>(
|
||||
backendFactory: () => Promise<Pick<SessionRepo<TMetadata>, "create" | "fork">>,
|
||||
onClose?: () => void | Promise<void>,
|
||||
): readonly ConformanceCase[] {
|
||||
return [
|
||||
...createSessionRepoForkDestinationReservationConformance(backendFactory, onClose),
|
||||
...createSessionRepoForkSourceSnapshotConformance(backendFactory, onClose),
|
||||
];
|
||||
}
|
||||
|
||||
/** Creates every fork conformance case. */
|
||||
export function createSessionRepoForkConformance<TMetadata extends SessionMetadata>(
|
||||
backendFactory: () => Promise<Pick<SessionRepo<TMetadata>, "create" | "list" | "fork">>,
|
||||
onClose?: () => void | Promise<void>,
|
||||
): readonly ConformanceCase[] {
|
||||
return [
|
||||
...createSessionRepoForkBehaviorConformance(backendFactory, onClose),
|
||||
...createSessionRepoForkCoordinationConformance(backendFactory, onClose),
|
||||
];
|
||||
}
|
||||
|
||||
/** Creates every SessionRepo conformance case. */
|
||||
export function createSessionRepoConformance<TMetadata extends SessionMetadata>(
|
||||
factory: () => Promise<SessionRepo<TMetadata>>,
|
||||
onClose?: () => void | Promise<void>,
|
||||
): readonly ConformanceCase[] {
|
||||
return [
|
||||
...createSessionRepoLifecycleConformance(factory, onClose),
|
||||
...createSessionRepoOwnershipConformance(factory, onClose),
|
||||
...createSessionRepoMessageConformance(factory, onClose),
|
||||
...createSessionRepoForkConformance(factory, onClose),
|
||||
];
|
||||
}
|
||||
920
third_party/pi/packages/agent/src/harness/session/testing/conformance/storage.ts
vendored
Normal file
920
third_party/pi/packages/agent/src/harness/session/testing/conformance/storage.ts
vendored
Normal file
@@ -0,0 +1,920 @@
|
||||
import { deepStrictEqual, ok, rejects, strictEqual } from "node:assert/strict";
|
||||
import type { Usage } from "@earendil-works/pi-ai";
|
||||
import { BACKGROUND_CONTEXT } from "../../../context.ts";
|
||||
import { insertEntry, insertUsage } from "../../commit.ts";
|
||||
import type {
|
||||
CommitResult,
|
||||
CompactionEntry,
|
||||
CustomEntry,
|
||||
Entry,
|
||||
EntryStructure,
|
||||
JsonValue,
|
||||
MessageEntry,
|
||||
NewEntry,
|
||||
Storage,
|
||||
Write,
|
||||
} from "../../types.ts";
|
||||
import {
|
||||
appendList,
|
||||
branchTip,
|
||||
deleteList,
|
||||
deleteValue,
|
||||
entryLabel,
|
||||
list,
|
||||
pendingEntry,
|
||||
setValue,
|
||||
value,
|
||||
} from "../../values.ts";
|
||||
import type { ConformanceCase, StorageFixture } from "../types.ts";
|
||||
|
||||
const MESSAGE_TIMESTAMP = 1_650_000_000_000;
|
||||
const testName = value<string>("test.session.name");
|
||||
const testValue = (key: string) => value<JsonValue>("test.value", key);
|
||||
const testValuePrefix = (prefix = "") => value<JsonValue>("test.value", prefix);
|
||||
const testList = (key: string) => list<JsonValue>("test.list", key);
|
||||
|
||||
type ConformanceTest = (fixture: StorageFixture) => Promise<void>;
|
||||
|
||||
function createCase(
|
||||
factory: () => Promise<StorageFixture>,
|
||||
group: string,
|
||||
name: string,
|
||||
test: ConformanceTest,
|
||||
): ConformanceCase {
|
||||
return {
|
||||
group,
|
||||
name,
|
||||
async run() {
|
||||
await using fixture = await factory();
|
||||
await test(fixture);
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function usage(input: number, output: number, options: { cacheWrite1h?: number; reasoning?: number } = {}): Usage {
|
||||
return {
|
||||
input,
|
||||
output,
|
||||
cacheRead: input + 1,
|
||||
cacheWrite: output + 1,
|
||||
...(options.cacheWrite1h === undefined ? {} : { cacheWrite1h: options.cacheWrite1h }),
|
||||
...(options.reasoning === undefined ? {} : { reasoning: options.reasoning }),
|
||||
totalTokens: input + output,
|
||||
cost: {
|
||||
input: input / 100,
|
||||
output: output / 100,
|
||||
cacheRead: (input + 1) / 100,
|
||||
cacheWrite: (output + 1) / 100,
|
||||
total: (input + output + 2) / 100,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function zeroUsage(): Usage {
|
||||
return {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
}
|
||||
|
||||
function userEntry(id: string, parentId: string | null = null, text = id): NewEntry<MessageEntry> {
|
||||
return {
|
||||
id,
|
||||
parentId,
|
||||
type: "message",
|
||||
message: {
|
||||
role: "user",
|
||||
content: [{ type: "text", text }],
|
||||
timestamp: MESSAGE_TIMESTAMP,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function customEntry(
|
||||
id: string,
|
||||
parentId: string | null,
|
||||
customType = "note",
|
||||
data: CustomEntry["data"] = { id },
|
||||
): NewEntry<CustomEntry> {
|
||||
return { id, parentId, type: "custom", customType, data };
|
||||
}
|
||||
|
||||
function compactionEntry(id: string, parentId: string | null): NewEntry<CompactionEntry> {
|
||||
return {
|
||||
id,
|
||||
parentId,
|
||||
type: "compaction",
|
||||
summary: `summary:${id}`,
|
||||
retainedTail: [],
|
||||
tokensBefore: 10,
|
||||
fromHook: false,
|
||||
};
|
||||
}
|
||||
|
||||
// EntryStructure is included for scanBranchStructure conformance, especially SQLite's payload-free branch scans.
|
||||
function ids(entries: readonly (Entry | EntryStructure)[]): string[] {
|
||||
return entries.map((entry) => entry.id);
|
||||
}
|
||||
|
||||
function assertStrictlyIncreasing(values: readonly number[]): void {
|
||||
for (let index = 1; index < values.length; index++) {
|
||||
ok(values[index - 1] < values[index], `Expected ${values.join(", ")} to be strictly increasing`);
|
||||
}
|
||||
}
|
||||
|
||||
async function assertCommitStats(storage: Storage, result: CommitResult): Promise<void> {
|
||||
deepStrictEqual(result.stats, await storage.getStats(BACKGROUND_CONTEXT));
|
||||
}
|
||||
|
||||
/** Creates fresh, runner-independent cases for the durable Storage contract. */
|
||||
export function createStorageConformance(factory: () => Promise<StorageFixture>): readonly ConformanceCase[] {
|
||||
return [
|
||||
createCase(factory, "transactions", "commits mixed writes atomically in write order", async ({ storage }) => {
|
||||
const result = await storage.commit(
|
||||
[
|
||||
insertEntry(userEntry("entry")),
|
||||
setValue(testName, "session"),
|
||||
insertUsage({ id: "usage", usage: usage(2, 3), adjustment: false, entryId: "entry" }),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
strictEqual(result.seqs.length, 3);
|
||||
strictEqual(result.firstSeq, result.seqs[0]);
|
||||
await assertCommitStats(storage, result);
|
||||
assertStrictlyIncreasing(result.seqs);
|
||||
ok(Number.isSafeInteger(result.timestamp) && result.timestamp >= 0);
|
||||
deepStrictEqual(
|
||||
await storage.getEntries(["entry"], BACKGROUND_CONTEXT),
|
||||
new Map([["entry", { ...userEntry("entry"), seq: result.seqs[0], timestamp: result.timestamp }]]),
|
||||
);
|
||||
deepStrictEqual(await storage.getValue(testName, BACKGROUND_CONTEXT), {
|
||||
address: testName,
|
||||
value: "session",
|
||||
seq: result.seqs[1],
|
||||
});
|
||||
deepStrictEqual(await storage.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT), [
|
||||
{ id: "usage", seq: result.seqs[2], usage: usage(2, 3), adjustment: false, entryId: "entry" },
|
||||
]);
|
||||
}),
|
||||
|
||||
createCase(
|
||||
factory,
|
||||
"transactions",
|
||||
"rolls back every store when a mixed transaction fails",
|
||||
async ({ storage }) => {
|
||||
await storage.commit(
|
||||
[insertEntry(userEntry("root")), insertUsage({ id: "taken", usage: usage(1, 1), adjustment: false })],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
const entriesBefore = await storage.scanEntries({ order: "asc" }, BACKGROUND_CONTEXT);
|
||||
const usageBefore = await storage.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT);
|
||||
const statsBefore = await storage.getStats(BACKGROUND_CONTEXT);
|
||||
|
||||
await rejects(
|
||||
storage.commit(
|
||||
[
|
||||
setValue(testName, "transient"),
|
||||
insertEntry(customEntry("transient-entry", "root")),
|
||||
insertUsage({ id: "transient-usage", usage: usage(5, 8), adjustment: true }),
|
||||
insertEntry(customEntry("taken", "root")),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
);
|
||||
|
||||
deepStrictEqual(await storage.scanEntries({ order: "asc" }, BACKGROUND_CONTEXT), entriesBefore);
|
||||
deepStrictEqual(await storage.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT), usageBefore);
|
||||
deepStrictEqual(await storage.getStats(BACKGROUND_CONTEXT), statsBefore);
|
||||
strictEqual(await storage.getValue(testName, BACKGROUND_CONTEXT), undefined);
|
||||
},
|
||||
),
|
||||
|
||||
createCase(
|
||||
factory,
|
||||
"transactions",
|
||||
"preserves overwritten and deleted values when a transaction fails",
|
||||
async ({ storage }) => {
|
||||
await storage.commit(
|
||||
[
|
||||
setValue(testValue("overwritten"), "original"),
|
||||
setValue(testValue("deleted"), { kept: true }),
|
||||
insertEntry(userEntry("taken")),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
const overwrittenBefore = await storage.getValue(testValue("overwritten"), BACKGROUND_CONTEXT);
|
||||
const deletedBefore = await storage.getValue(testValue("deleted"), BACKGROUND_CONTEXT);
|
||||
|
||||
await rejects(
|
||||
storage.commit(
|
||||
[
|
||||
setValue(testValue("overwritten"), "transient"),
|
||||
deleteValue(testValue("deleted")),
|
||||
insertEntry(customEntry("transient", "taken")),
|
||||
insertEntry(customEntry("taken", null)),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
);
|
||||
|
||||
deepStrictEqual(await storage.getValue(testValue("overwritten"), BACKGROUND_CONTEXT), overwrittenBefore);
|
||||
deepStrictEqual(await storage.getValue(testValue("deleted"), BACKGROUND_CONTEXT), deletedBefore);
|
||||
strictEqual((await storage.getEntries(["transient"], BACKGROUND_CONTEXT)).has("transient"), false);
|
||||
},
|
||||
),
|
||||
|
||||
createCase(factory, "transactions", "enforces one shared entry and usage id namespace", async ({ storage }) => {
|
||||
await storage.commit(
|
||||
[
|
||||
insertEntry(userEntry("existing-entry")),
|
||||
insertUsage({ id: "existing-usage", usage: usage(1, 1), adjustment: false }),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
await rejects(
|
||||
storage.commit(
|
||||
[insertUsage({ id: "existing-entry", usage: usage(2, 2), adjustment: false })],
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
);
|
||||
await rejects(storage.commit([insertEntry(customEntry("existing-usage", null))], BACKGROUND_CONTEXT));
|
||||
|
||||
for (const [id, writes] of [
|
||||
[
|
||||
"entry-then-usage",
|
||||
[
|
||||
insertEntry(customEntry("entry-then-usage", null)),
|
||||
insertUsage({ id: "entry-then-usage", usage: usage(3, 3), adjustment: false }),
|
||||
],
|
||||
],
|
||||
[
|
||||
"usage-then-entry",
|
||||
[
|
||||
insertUsage({ id: "usage-then-entry", usage: usage(4, 4), adjustment: false }),
|
||||
insertEntry(customEntry("usage-then-entry", null)),
|
||||
],
|
||||
],
|
||||
] satisfies Array<[string, Write[]]>) {
|
||||
await rejects(storage.commit(writes, BACKGROUND_CONTEXT), `Expected duplicate id ${id} to reject`);
|
||||
}
|
||||
|
||||
deepStrictEqual(ids(await storage.scanEntries({ order: "asc" }, BACKGROUND_CONTEXT)), ["existing-entry"]);
|
||||
deepStrictEqual(
|
||||
(await storage.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT)).map((row) => row.id),
|
||||
["existing-usage"],
|
||||
);
|
||||
}),
|
||||
|
||||
createCase(
|
||||
factory,
|
||||
"transactions",
|
||||
"resolves parents only from prior entries and earlier writes",
|
||||
async ({ storage }) => {
|
||||
await storage.commit([insertEntry(userEntry("root"))], BACKGROUND_CONTEXT);
|
||||
await storage.commit(
|
||||
[insertEntry(customEntry("child", "root")), insertEntry(customEntry("grandchild", "child"))],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
deepStrictEqual(
|
||||
ids(await storage.scanBranch({ start: "grandchild", order: "oldestFirst" }, BACKGROUND_CONTEXT)),
|
||||
["root", "child", "grandchild"],
|
||||
);
|
||||
|
||||
await rejects(
|
||||
storage.commit(
|
||||
[
|
||||
insertEntry(customEntry("before-parent", "later-parent")),
|
||||
insertEntry(customEntry("later-parent", "root")),
|
||||
setValue(entryLabel("before-parent"), "transient"),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
);
|
||||
await rejects(storage.commit([insertEntry(customEntry("orphan", "missing"))], BACKGROUND_CONTEXT));
|
||||
await storage.commit(
|
||||
[insertUsage({ id: "usage-is-not-parent", usage: usage(1, 1), adjustment: false })],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
await rejects(
|
||||
storage.commit([insertEntry(customEntry("usage-child", "usage-is-not-parent"))], BACKGROUND_CONTEXT),
|
||||
);
|
||||
|
||||
deepStrictEqual(
|
||||
await storage.getEntries(["before-parent", "later-parent", "orphan", "usage-child"], BACKGROUND_CONTEXT),
|
||||
new Map(),
|
||||
);
|
||||
strictEqual(await storage.getValue(entryLabel("before-parent"), BACKGROUND_CONTEXT), undefined);
|
||||
},
|
||||
),
|
||||
|
||||
createCase(factory, "transactions", "places pending content under its reserved entry id", async ({ storage }) => {
|
||||
const entry = userEntry("reserved", null, "queued");
|
||||
await storage.commit(
|
||||
[
|
||||
setValue(pendingEntry(entry.id), { type: "message", payload: entry.message }),
|
||||
setValue(branchTip("main"), null),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
strictEqual((await storage.getEntries([entry.id], BACKGROUND_CONTEXT)).has(entry.id), false);
|
||||
deepStrictEqual((await storage.getValue(pendingEntry(entry.id), BACKGROUND_CONTEXT))?.value, {
|
||||
type: "message",
|
||||
payload: entry.message,
|
||||
});
|
||||
strictEqual((await storage.getValue(branchTip("main"), BACKGROUND_CONTEXT))?.value, null);
|
||||
|
||||
const placement = await storage.commit(
|
||||
[insertEntry(entry), deleteValue(pendingEntry(entry.id)), setValue(branchTip("main"), entry.id)],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
deepStrictEqual(
|
||||
await storage.getEntries([entry.id], BACKGROUND_CONTEXT),
|
||||
new Map([[entry.id, { ...entry, seq: placement.seqs[0], timestamp: placement.timestamp }]]),
|
||||
);
|
||||
strictEqual(await storage.getValue(pendingEntry(entry.id), BACKGROUND_CONTEXT), undefined);
|
||||
deepStrictEqual(await storage.getValue(branchTip("main"), BACKGROUND_CONTEXT), {
|
||||
address: branchTip("main"),
|
||||
value: entry.id,
|
||||
seq: placement.seqs[2],
|
||||
});
|
||||
}),
|
||||
|
||||
createCase(
|
||||
factory,
|
||||
"values",
|
||||
"sets, replaces, deletes, and recreates values without tombstones",
|
||||
async ({ storage }) => {
|
||||
const first = await storage.commit(
|
||||
[
|
||||
setValue(testValue("prefix/b"), 1),
|
||||
setValue(testValue("prefix/a"), 2),
|
||||
setValue(testValue("other"), 3),
|
||||
setValue(testValue("prefix/\ue000"), 4),
|
||||
setValue(testValue("prefix/\u{10000}"), 5),
|
||||
setValue(testValue("prefix/a"), null),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
deepStrictEqual(await storage.getValue(testValue("prefix/a"), BACKGROUND_CONTEXT), {
|
||||
address: testValue("prefix/a"),
|
||||
value: null,
|
||||
seq: first.seqs[5],
|
||||
});
|
||||
|
||||
const second = await storage.commit(
|
||||
[
|
||||
deleteValue(testValue("prefix/a")),
|
||||
deleteValue(testValue("absent")),
|
||||
setValue(testValue("prefix/a"), "recreated"),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
deepStrictEqual(await storage.scanValues(testValuePrefix("prefix/"), BACKGROUND_CONTEXT), [
|
||||
{ address: testValue("prefix/a"), value: "recreated", seq: second.seqs[2] },
|
||||
{ address: testValue("prefix/b"), value: 1, seq: first.seqs[0] },
|
||||
{ address: testValue("prefix/\ue000"), value: 4, seq: first.seqs[3] },
|
||||
{ address: testValue("prefix/\u{10000}"), value: 5, seq: first.seqs[4] },
|
||||
]);
|
||||
strictEqual(await storage.getValue(testValue("absent"), BACKGROUND_CONTEXT), undefined);
|
||||
},
|
||||
),
|
||||
|
||||
createCase(
|
||||
factory,
|
||||
"values",
|
||||
"applies same-transaction value and list operations in write order",
|
||||
async ({ storage }) => {
|
||||
const keptValue = testValue("write-order/kept");
|
||||
const deletedValue = testValue("write-order/deleted");
|
||||
const keptList = testList("write-order/kept");
|
||||
const deletedList = testList("write-order/deleted");
|
||||
const result = await storage.commit(
|
||||
[
|
||||
setValue(deletedValue, "transient"),
|
||||
deleteValue(deletedValue),
|
||||
setValue(keptValue, "transient"),
|
||||
setValue(keptValue, "kept"),
|
||||
appendList(keptList, "transient"),
|
||||
deleteList(keptList),
|
||||
appendList(keptList, "kept"),
|
||||
appendList(deletedList, "transient"),
|
||||
deleteList(deletedList),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
strictEqual(await storage.getValue(deletedValue, BACKGROUND_CONTEXT), undefined);
|
||||
deepStrictEqual(await storage.getValue(keptValue, BACKGROUND_CONTEXT), {
|
||||
address: keptValue,
|
||||
value: "kept",
|
||||
seq: result.seqs[3],
|
||||
});
|
||||
deepStrictEqual(await storage.readList(keptList, undefined, BACKGROUND_CONTEXT), [
|
||||
{ seq: result.seqs[6], value: "kept" },
|
||||
]);
|
||||
deepStrictEqual(await storage.readList(deletedList, undefined, BACKGROUND_CONTEXT), []);
|
||||
},
|
||||
),
|
||||
|
||||
createCase(
|
||||
factory,
|
||||
"values",
|
||||
"does not change historical stores during value-only commits",
|
||||
async ({ storage }) => {
|
||||
await storage.commit(
|
||||
[
|
||||
insertEntry(userEntry("root")),
|
||||
insertUsage({ id: "historical-usage", usage: usage(2, 3), adjustment: false }),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
const entriesBefore = await storage.scanEntries({ order: "asc" }, BACKGROUND_CONTEXT);
|
||||
const usageBefore = await storage.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT);
|
||||
const statsBefore = await storage.getStats(BACKGROUND_CONTEXT);
|
||||
|
||||
const result = await storage.commit(
|
||||
[setValue(testName, "first"), setValue(testName, "second")],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
deepStrictEqual(await storage.scanEntries({ order: "asc" }, BACKGROUND_CONTEXT), entriesBefore);
|
||||
deepStrictEqual(await storage.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT), usageBefore);
|
||||
deepStrictEqual(await storage.getStats(BACKGROUND_CONTEXT), statsBefore);
|
||||
deepStrictEqual(await storage.getValue(testName, BACKGROUND_CONTEXT), {
|
||||
address: testName,
|
||||
value: "second",
|
||||
seq: result.seqs[1],
|
||||
});
|
||||
},
|
||||
),
|
||||
|
||||
createCase(factory, "lists", "pages appends by global sequence and deletes whole lists", async ({ storage }) => {
|
||||
const address = testList("events");
|
||||
strictEqual((await storage.readList(address, undefined, BACKGROUND_CONTEXT)).length, 0);
|
||||
const result = await storage.commit(
|
||||
[appendList(address, "a"), setValue(testName, "gap"), appendList(address, "b"), appendList(address, "c")],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
deepStrictEqual(await storage.readList(address, undefined, BACKGROUND_CONTEXT), [
|
||||
{ seq: result.seqs[0], value: "a" },
|
||||
{ seq: result.seqs[2], value: "b" },
|
||||
{ seq: result.seqs[3], value: "c" },
|
||||
]);
|
||||
deepStrictEqual(await storage.readList(address, { limit: 2 }, BACKGROUND_CONTEXT), [
|
||||
{ seq: result.seqs[0], value: "a" },
|
||||
{ seq: result.seqs[2], value: "b" },
|
||||
]);
|
||||
deepStrictEqual(
|
||||
await storage.readList(address, { cursor: { seq: result.seqs[0]! }, limit: 2 }, BACKGROUND_CONTEXT),
|
||||
[
|
||||
{ seq: result.seqs[2], value: "b" },
|
||||
{ seq: result.seqs[3], value: "c" },
|
||||
],
|
||||
);
|
||||
deepStrictEqual(await storage.readList(address, { order: "desc", limit: 2 }, BACKGROUND_CONTEXT), [
|
||||
{ seq: result.seqs[3], value: "c" },
|
||||
{ seq: result.seqs[2], value: "b" },
|
||||
]);
|
||||
deepStrictEqual(
|
||||
await storage.readList(
|
||||
address,
|
||||
{ order: "desc", cursor: { seq: result.seqs[3]! }, limit: 2 },
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
[
|
||||
{ seq: result.seqs[2], value: "b" },
|
||||
{ seq: result.seqs[0], value: "a" },
|
||||
],
|
||||
);
|
||||
await rejects(storage.readList(address, { limit: 0 }, BACKGROUND_CONTEXT));
|
||||
await rejects(storage.readList(address, { limit: Number.MAX_VALUE }, BACKGROUND_CONTEXT));
|
||||
|
||||
await storage.commit(
|
||||
[deleteList(address), deleteList(testList("absent")), appendList(address, "new")],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
deepStrictEqual(
|
||||
(await storage.readList(address, undefined, BACKGROUND_CONTEXT)).map(({ value }) => value),
|
||||
["new"],
|
||||
);
|
||||
}),
|
||||
|
||||
createCase(factory, "lists", "clamps one read page without limiting list growth", async ({ storage }) => {
|
||||
const address = testList("large");
|
||||
await storage.commit(
|
||||
Array.from({ length: 10_001 }, (_, index) => appendList(address, index)),
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
const firstPage = await storage.readList(address, undefined, BACKGROUND_CONTEXT);
|
||||
strictEqual(firstPage.length, 1_000);
|
||||
strictEqual((await storage.readList(address, { limit: 20_000 }, BACKGROUND_CONTEXT)).length, 10_000);
|
||||
strictEqual(
|
||||
(
|
||||
await storage.readList(
|
||||
address,
|
||||
{
|
||||
cursor: { seq: firstPage.at(-1)!.seq },
|
||||
},
|
||||
BACKGROUND_CONTEXT,
|
||||
)
|
||||
).length,
|
||||
1_000,
|
||||
);
|
||||
}),
|
||||
|
||||
createCase(
|
||||
factory,
|
||||
"lists",
|
||||
"commits mixed list writes atomically and rolls them back with siblings",
|
||||
async ({ storage }) => {
|
||||
const address = testList("atomic");
|
||||
const committed = await storage.commit(
|
||||
[
|
||||
insertEntry(userEntry("mixed")),
|
||||
appendList(address, "kept"),
|
||||
setValue(testName, "kept"),
|
||||
insertUsage({ id: "mixed-usage", usage: usage(1, 2), adjustment: false }),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
deepStrictEqual(await storage.readList(address, undefined, BACKGROUND_CONTEXT), [
|
||||
{ seq: committed.seqs[1], value: "kept" },
|
||||
]);
|
||||
|
||||
await rejects(
|
||||
storage.commit(
|
||||
[appendList(address, "transient"), deleteValue(testName), insertEntry(userEntry("mixed"))],
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
);
|
||||
deepStrictEqual(await storage.readList(address, undefined, BACKGROUND_CONTEXT), [
|
||||
{ seq: committed.seqs[1], value: "kept" },
|
||||
]);
|
||||
strictEqual((await storage.getValue(testName, BACKGROUND_CONTEXT))?.value, "kept");
|
||||
},
|
||||
),
|
||||
|
||||
createCase(factory, "entry queries", "stores custom entries with and without data", async ({ storage }) => {
|
||||
const result = await storage.commit(
|
||||
[
|
||||
insertEntry({ id: "without-data", parentId: null, type: "custom", customType: "marker" }),
|
||||
insertEntry(customEntry("with-data", "without-data", "note", { nested: [1, 2] })),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
deepStrictEqual(
|
||||
await storage.getEntries(["without-data", "with-data"], BACKGROUND_CONTEXT),
|
||||
new Map([
|
||||
[
|
||||
"without-data",
|
||||
{
|
||||
id: "without-data",
|
||||
parentId: null,
|
||||
type: "custom",
|
||||
customType: "marker",
|
||||
seq: result.seqs[0],
|
||||
timestamp: result.timestamp,
|
||||
},
|
||||
],
|
||||
[
|
||||
"with-data",
|
||||
{
|
||||
...customEntry("with-data", "without-data", "note", { nested: [1, 2] }),
|
||||
seq: result.seqs[1],
|
||||
timestamp: result.timestamp,
|
||||
},
|
||||
],
|
||||
]),
|
||||
);
|
||||
}),
|
||||
|
||||
createCase(
|
||||
factory,
|
||||
"entry queries",
|
||||
"scans global entries with explicit ranges, filters, orders, and limits",
|
||||
async ({ storage }) => {
|
||||
const result = await storage.commit(
|
||||
[
|
||||
insertEntry(userEntry("root")),
|
||||
insertEntry(customEntry("note-1", "root", "note")),
|
||||
insertEntry(customEntry("other", "note-1", "other")),
|
||||
insertEntry(customEntry("note-2", "other", "note")),
|
||||
insertEntry(userEntry("tail", "note-2")),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
deepStrictEqual(
|
||||
ids(
|
||||
await storage.scanEntries(
|
||||
{
|
||||
type: "custom",
|
||||
customType: "note",
|
||||
fromSeq: result.seqs[1],
|
||||
toSeq: result.seqs[3],
|
||||
order: "desc",
|
||||
},
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
),
|
||||
["note-2", "note-1"],
|
||||
);
|
||||
deepStrictEqual(ids(await storage.scanEntries({ order: "asc", limit: 2 }, BACKGROUND_CONTEXT)), [
|
||||
"root",
|
||||
"note-1",
|
||||
]);
|
||||
deepStrictEqual(ids(await storage.scanEntries({ order: "desc", limit: 2 }, BACKGROUND_CONTEXT)), [
|
||||
"tail",
|
||||
"note-2",
|
||||
]);
|
||||
},
|
||||
),
|
||||
|
||||
createCase(
|
||||
factory,
|
||||
"branch queries",
|
||||
"applies stops before filters and cursors before limits",
|
||||
async ({ storage }) => {
|
||||
const result = await storage.commit(
|
||||
[
|
||||
insertEntry(userEntry("root")),
|
||||
insertEntry(customEntry("marker", "root", "marker")),
|
||||
insertEntry(userEntry("middle", "marker")),
|
||||
insertEntry(compactionEntry("compact", "middle")),
|
||||
insertEntry(customEntry("note", "compact", "note")),
|
||||
insertEntry(userEntry("leaf", "note")),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
deepStrictEqual(
|
||||
ids(
|
||||
await storage.scanBranch(
|
||||
{ start: "leaf", stopAtType: "compaction", type: "message" },
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
),
|
||||
["leaf"],
|
||||
);
|
||||
deepStrictEqual(
|
||||
ids(
|
||||
await storage.scanBranch(
|
||||
{
|
||||
start: "leaf",
|
||||
order: "oldestFirst",
|
||||
stopAtId: "middle",
|
||||
type: "custom",
|
||||
},
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
),
|
||||
["marker"],
|
||||
);
|
||||
deepStrictEqual(
|
||||
ids(
|
||||
await storage.scanBranch(
|
||||
{
|
||||
start: "leaf",
|
||||
order: "newestFirst",
|
||||
cursor: { seq: result.seqs[4] },
|
||||
limit: 2,
|
||||
},
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
),
|
||||
["compact", "middle"],
|
||||
);
|
||||
deepStrictEqual(
|
||||
ids(
|
||||
await storage.scanBranch(
|
||||
{
|
||||
start: "leaf",
|
||||
order: "oldestFirst",
|
||||
cursor: { seq: result.seqs[1] },
|
||||
limit: 2,
|
||||
},
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
),
|
||||
["middle", "compact"],
|
||||
);
|
||||
deepStrictEqual(
|
||||
ids(await storage.scanBranch({ start: "leaf", stopAtId: "leaf", type: "custom" }, BACKGROUND_CONTEXT)),
|
||||
[],
|
||||
);
|
||||
deepStrictEqual(ids(await storage.scanBranch({ start: "leaf", customType: "note" }, BACKGROUND_CONTEXT)), [
|
||||
"note",
|
||||
]);
|
||||
await rejects(storage.scanBranch({ start: "missing" }, BACKGROUND_CONTEXT));
|
||||
},
|
||||
),
|
||||
|
||||
createCase(factory, "branch queries", "returns branch structure without payload fields", async ({ storage }) => {
|
||||
const result = await storage.commit(
|
||||
[insertEntry(userEntry("root")), insertEntry(customEntry("child", "root", "note"))],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
deepStrictEqual(
|
||||
await storage.scanBranchStructure({ start: "child", order: "oldestFirst" }, BACKGROUND_CONTEXT),
|
||||
[
|
||||
{
|
||||
id: "root",
|
||||
parentId: null,
|
||||
seq: result.seqs[0],
|
||||
timestamp: result.timestamp,
|
||||
type: "message",
|
||||
},
|
||||
{
|
||||
id: "child",
|
||||
parentId: "root",
|
||||
seq: result.seqs[1],
|
||||
timestamp: result.timestamp,
|
||||
type: "custom",
|
||||
customType: "note",
|
||||
},
|
||||
],
|
||||
);
|
||||
}),
|
||||
|
||||
createCase(
|
||||
factory,
|
||||
"branch queries",
|
||||
"applies branch query semantics to structure scans",
|
||||
async ({ storage }) => {
|
||||
const result = await storage.commit(
|
||||
[
|
||||
insertEntry(userEntry("root")),
|
||||
insertEntry(customEntry("marker", "root", "marker")),
|
||||
insertEntry(userEntry("middle", "marker")),
|
||||
insertEntry(compactionEntry("compact", "middle")),
|
||||
insertEntry(customEntry("note", "compact", "note")),
|
||||
insertEntry(userEntry("leaf", "note")),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
deepStrictEqual(
|
||||
ids(
|
||||
await storage.scanBranchStructure(
|
||||
{ start: "leaf", stopAtType: "compaction", type: "message" },
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
),
|
||||
["leaf"],
|
||||
);
|
||||
deepStrictEqual(
|
||||
ids(
|
||||
await storage.scanBranchStructure(
|
||||
{ start: "leaf", order: "oldestFirst", cursor: { seq: result.seqs[1] }, limit: 2 },
|
||||
BACKGROUND_CONTEXT,
|
||||
),
|
||||
),
|
||||
["middle", "compact"],
|
||||
);
|
||||
await rejects(storage.scanBranchStructure({ start: "missing" }, BACKGROUND_CONTEXT));
|
||||
},
|
||||
),
|
||||
|
||||
createCase(
|
||||
factory,
|
||||
"usage and stats",
|
||||
"scans the usage ledger with explicit ranges, orders, and limits",
|
||||
async ({ storage }) => {
|
||||
const result = await storage.commit(
|
||||
[
|
||||
insertUsage({ id: "usage-1", usage: usage(1, 1), adjustment: false }),
|
||||
setValue(testName, "sequence gap"),
|
||||
insertUsage({ id: "usage-2", usage: usage(2, 2), adjustment: false }),
|
||||
insertUsage({ id: "usage-3", usage: usage(3, 3), adjustment: true }),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
|
||||
deepStrictEqual(
|
||||
(
|
||||
await storage.scanUsage(
|
||||
{ fromSeq: result.seqs[1], toSeq: result.seqs[2], order: "asc" },
|
||||
BACKGROUND_CONTEXT,
|
||||
)
|
||||
).map((row) => row.id),
|
||||
["usage-2"],
|
||||
);
|
||||
deepStrictEqual(
|
||||
(await storage.scanUsage({ order: "desc", limit: 2 }, BACKGROUND_CONTEXT)).map((row) => row.id),
|
||||
["usage-3", "usage-2"],
|
||||
);
|
||||
deepStrictEqual(
|
||||
(await storage.scanUsage({ order: "asc", limit: 2 }, BACKGROUND_CONTEXT)).map((row) => row.id),
|
||||
["usage-1", "usage-2"],
|
||||
);
|
||||
},
|
||||
),
|
||||
|
||||
createCase(
|
||||
factory,
|
||||
"usage and stats",
|
||||
"keeps stats equal to message count and ledger totals",
|
||||
async ({ storage }) => {
|
||||
deepStrictEqual(await storage.getStats(BACKGROUND_CONTEXT), { messageCount: 0, usage: zeroUsage() });
|
||||
|
||||
const firstUsage = usage(2, 3, { cacheWrite1h: 4, reasoning: 1 });
|
||||
const first = await storage.commit(
|
||||
[
|
||||
insertEntry(userEntry("message")),
|
||||
insertUsage({ id: "usage-1", usage: firstUsage, adjustment: false }),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
await assertCommitStats(storage, first);
|
||||
deepStrictEqual(first.stats, { messageCount: 1, usage: firstUsage });
|
||||
|
||||
const secondUsage = usage(5, 7, { cacheWrite1h: 6, reasoning: 2 });
|
||||
const second = await storage.commit(
|
||||
[
|
||||
insertEntry(customEntry("custom", "message")),
|
||||
insertEntry(compactionEntry("compaction", "custom")),
|
||||
insertUsage({ id: "usage-2", usage: secondUsage, adjustment: true }),
|
||||
],
|
||||
BACKGROUND_CONTEXT,
|
||||
);
|
||||
await assertCommitStats(storage, second);
|
||||
deepStrictEqual(second.stats, {
|
||||
messageCount: 1,
|
||||
usage: {
|
||||
input: 7,
|
||||
output: 10,
|
||||
cacheRead: 9,
|
||||
cacheWrite: 12,
|
||||
cacheWrite1h: 10,
|
||||
reasoning: 3,
|
||||
totalTokens: 17,
|
||||
cost: {
|
||||
input: firstUsage.cost.input + secondUsage.cost.input,
|
||||
output: firstUsage.cost.output + secondUsage.cost.output,
|
||||
cacheRead: firstUsage.cost.cacheRead + secondUsage.cost.cacheRead,
|
||||
cacheWrite: firstUsage.cost.cacheWrite + secondUsage.cost.cacheWrite,
|
||||
total: firstUsage.cost.total + secondUsage.cost.total,
|
||||
},
|
||||
},
|
||||
});
|
||||
},
|
||||
),
|
||||
|
||||
createCase(
|
||||
factory,
|
||||
"serialization",
|
||||
"serializes back-to-back commits in admission order",
|
||||
async ({ storage }) => {
|
||||
const first = storage.commit([insertEntry(userEntry("first"))], BACKGROUND_CONTEXT);
|
||||
const second = storage.commit([insertEntry(userEntry("second", "first"))], BACKGROUND_CONTEXT);
|
||||
const [firstResult, secondResult] = await Promise.all([first, second]);
|
||||
|
||||
ok(firstResult.seqs[0]! < secondResult.seqs[0]!);
|
||||
deepStrictEqual(firstResult.stats, { messageCount: 1, usage: zeroUsage() });
|
||||
deepStrictEqual(secondResult.stats, { messageCount: 2, usage: zeroUsage() });
|
||||
await assertCommitStats(storage, secondResult);
|
||||
deepStrictEqual(ids(await storage.scanEntries({ order: "asc" }, BACKGROUND_CONTEXT)), ["first", "second"]);
|
||||
},
|
||||
),
|
||||
|
||||
createCase(
|
||||
factory,
|
||||
"lifecycle",
|
||||
"seals admission, drains admitted commits, and closes idempotently",
|
||||
async ({ storage }) => {
|
||||
const admitted = storage.commit([insertEntry(userEntry("admitted"))], BACKGROUND_CONTEXT);
|
||||
const firstClose = storage.close(BACKGROUND_CONTEXT);
|
||||
const secondClose = storage.close(BACKGROUND_CONTEXT);
|
||||
|
||||
await rejects(storage.getStats(BACKGROUND_CONTEXT));
|
||||
await rejects(storage.commit([], BACKGROUND_CONTEXT));
|
||||
strictEqual((await admitted).seqs.length, 1);
|
||||
await Promise.all([firstClose, secondClose]);
|
||||
|
||||
const rejectedReads = [
|
||||
storage.getEntries([], BACKGROUND_CONTEXT),
|
||||
storage.getValue(testName, BACKGROUND_CONTEXT),
|
||||
storage.scanValues(testName, BACKGROUND_CONTEXT),
|
||||
storage.readList(testList("events"), undefined, BACKGROUND_CONTEXT),
|
||||
storage.scanBranch({ start: "admitted" }, BACKGROUND_CONTEXT),
|
||||
storage.scanBranchStructure({ start: "admitted" }, BACKGROUND_CONTEXT),
|
||||
storage.scanEntries({ order: "asc" }, BACKGROUND_CONTEXT),
|
||||
storage.scanUsage({ order: "asc" }, BACKGROUND_CONTEXT),
|
||||
storage.getStats(BACKGROUND_CONTEXT),
|
||||
];
|
||||
for (const read of rejectedReads) await rejects(read);
|
||||
},
|
||||
),
|
||||
];
|
||||
}
|
||||
114
third_party/pi/packages/agent/src/harness/session/testing/gating-storage.ts
vendored
Normal file
114
third_party/pi/packages/agent/src/harness/session/testing/gating-storage.ts
vendored
Normal file
@@ -0,0 +1,114 @@
|
||||
import type { Context } from "../../context.ts";
|
||||
import type { CommitResult, Write } from "../types.ts";
|
||||
import { StorageDecorator } from "./storage-decorator.ts";
|
||||
|
||||
/** Thrown for every commit rejected after simulated storage loss. */
|
||||
export class CommitDiscarded extends Error {
|
||||
constructor(message: string) {
|
||||
super(message);
|
||||
this.name = "CommitDiscarded";
|
||||
}
|
||||
}
|
||||
|
||||
interface ParkedCommit {
|
||||
release(): void;
|
||||
drop(error: Error): void;
|
||||
readonly landing: Promise<void>;
|
||||
}
|
||||
|
||||
interface PendingWaiter {
|
||||
count: number;
|
||||
resolve(): void;
|
||||
reject(error: Error): void;
|
||||
}
|
||||
|
||||
/** Test-only storage decorator that deterministically parks admitted commits. */
|
||||
export class GatingStorage extends StorageDecorator {
|
||||
private armed = false;
|
||||
private discarded = false;
|
||||
private readonly queue: ParkedCommit[] = [];
|
||||
private readonly waiters: PendingWaiter[] = [];
|
||||
|
||||
/** Fixture setup bypasses gating until explicitly armed. */
|
||||
arm(): void {
|
||||
this.armed = true;
|
||||
}
|
||||
|
||||
pending(): number {
|
||||
return this.queue.length;
|
||||
}
|
||||
|
||||
/** Wait until at least `count` commits are parked. */
|
||||
waitPending(count = 1): Promise<void> {
|
||||
if (!Number.isSafeInteger(count) || count < 1) {
|
||||
return Promise.reject(new RangeError("Pending commit count must be a positive safe integer"));
|
||||
}
|
||||
if (this.discarded) return Promise.reject(new CommitDiscarded("storage discarded"));
|
||||
if (this.queue.length >= count) return Promise.resolve();
|
||||
return new Promise<void>((resolve, reject) => {
|
||||
this.waiters.push({ count, resolve, reject });
|
||||
});
|
||||
}
|
||||
|
||||
override async commit(writes: Write[], context: Context): Promise<CommitResult> {
|
||||
if (this.discarded) throw new CommitDiscarded("commit rejected: storage discarded");
|
||||
if (!this.armed) return super.commit(writes, context);
|
||||
|
||||
let landed!: () => void;
|
||||
let lost!: (error: Error) => void;
|
||||
const landing = new Promise<void>((resolve, reject) => {
|
||||
landed = resolve;
|
||||
lost = reject;
|
||||
});
|
||||
void landing.catch(() => {});
|
||||
|
||||
const released = new Promise<void>((resolve, reject) => {
|
||||
this.queue.push({ release: resolve, drop: reject, landing });
|
||||
});
|
||||
this.notifyWaiters();
|
||||
|
||||
try {
|
||||
await released;
|
||||
if (this.discarded) throw new CommitDiscarded("commit rejected: storage discarded");
|
||||
const result = await super.commit(writes, context);
|
||||
landed();
|
||||
return result;
|
||||
} catch (error) {
|
||||
const normalized = error instanceof Error ? error : new Error(String(error));
|
||||
lost(normalized);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
/** Release `count` commits in FIFO order and wait until each write lands. */
|
||||
async next(count = 1): Promise<void> {
|
||||
if (!Number.isSafeInteger(count) || count < 1) {
|
||||
throw new RangeError("Released commit count must be a positive safe integer");
|
||||
}
|
||||
for (let index = 0; index < count; index++) {
|
||||
await this.waitPending();
|
||||
const parked = this.queue.shift();
|
||||
if (parked === undefined) throw new Error("No parked commit");
|
||||
parked.release();
|
||||
await parked.landing;
|
||||
}
|
||||
}
|
||||
|
||||
/** Drop parked commits and permanently reject every later commit. */
|
||||
discard(): void {
|
||||
if (this.discarded) return;
|
||||
this.discarded = true;
|
||||
const error = new CommitDiscarded("commit discarded");
|
||||
for (const parked of this.queue.splice(0)) parked.drop(error);
|
||||
for (const waiter of this.waiters.splice(0)) waiter.reject(error);
|
||||
}
|
||||
|
||||
private notifyWaiters(): void {
|
||||
for (let index = this.waiters.length - 1; index >= 0; index--) {
|
||||
const waiter = this.waiters[index]!;
|
||||
if (this.queue.length < waiter.count) continue;
|
||||
this.waiters.splice(index, 1);
|
||||
waiter.resolve();
|
||||
}
|
||||
}
|
||||
}
|
||||
34
third_party/pi/packages/agent/src/harness/session/testing/index.ts
vendored
Normal file
34
third_party/pi/packages/agent/src/harness/session/testing/index.ts
vendored
Normal file
@@ -0,0 +1,34 @@
|
||||
export { STORAGE_BENCHMARK_DATASETS } from "./benchmark/datasets.ts";
|
||||
export {
|
||||
SESSION_REPO_CATALOG_BENCHMARK_DATASETS,
|
||||
SESSION_REPO_CATALOG_READ_BENCHMARK_SCENARIOS,
|
||||
SESSION_REPO_CATALOG_WRITE_BENCHMARK_SCENARIOS,
|
||||
SESSION_REPO_FORK_BENCHMARK_DATASETS,
|
||||
SESSION_REPO_FORK_WRITE_BENCHMARK_SCENARIOS,
|
||||
type SessionRepoCatalogBenchmarkDataset,
|
||||
seedSessionRepoCatalogBenchmark,
|
||||
seedSessionRepoForkBenchmark,
|
||||
sessionRepoBenchmarkSessionId,
|
||||
} from "./benchmark/session-repo.ts";
|
||||
export {
|
||||
generateStorageBenchmarkSeedTransactions,
|
||||
STORAGE_READ_BENCHMARK_SCENARIOS,
|
||||
STORAGE_WRITE_BENCHMARK_SCENARIOS,
|
||||
seedStorageBenchmark,
|
||||
} from "./benchmark/storage.ts";
|
||||
export {
|
||||
createSessionRepoConformance,
|
||||
createSessionRepoForkBehaviorConformance,
|
||||
createSessionRepoForkConformance,
|
||||
createSessionRepoForkCoordinationConformance,
|
||||
createSessionRepoForkDestinationReservationConformance,
|
||||
createSessionRepoForkSourceSnapshotConformance,
|
||||
createSessionRepoLifecycleConformance,
|
||||
createSessionRepoMessageConformance,
|
||||
createSessionRepoOwnershipConformance,
|
||||
} from "./conformance/session-repo.ts";
|
||||
export { createStorageConformance } from "./conformance/storage.ts";
|
||||
export { CommitDiscarded, GatingStorage } from "./gating-storage.ts";
|
||||
export { InstrumentedStorage } from "./instrumented-storage.ts";
|
||||
export { StorageDecorator } from "./storage-decorator.ts";
|
||||
export type { ConformanceCase, StorageFixture } from "./types.ts";
|
||||
21
third_party/pi/packages/agent/src/harness/session/testing/instrumented-storage.ts
vendored
Normal file
21
third_party/pi/packages/agent/src/harness/session/testing/instrumented-storage.ts
vendored
Normal file
@@ -0,0 +1,21 @@
|
||||
import type { Context } from "../../context.ts";
|
||||
import type { CommitResult, Write } from "../types.ts";
|
||||
import { StorageDecorator } from "./storage-decorator.ts";
|
||||
|
||||
/** Test-only transparent Storage decorator that records commit admission. */
|
||||
export class InstrumentedStorage extends StorageDecorator {
|
||||
private readonly commitAttempts: Write[][] = [];
|
||||
|
||||
getCommitAttempts(): readonly Write[][] {
|
||||
return this.commitAttempts.slice();
|
||||
}
|
||||
|
||||
clearCommitAttempts(): void {
|
||||
this.commitAttempts.length = 0;
|
||||
}
|
||||
|
||||
override commit(writes: Write[], context: Context): Promise<CommitResult> {
|
||||
this.commitAttempts.push(writes);
|
||||
return this.delegate.commit(writes, context);
|
||||
}
|
||||
}
|
||||
71
third_party/pi/packages/agent/src/harness/session/testing/storage-decorator.ts
vendored
Normal file
71
third_party/pi/packages/agent/src/harness/session/testing/storage-decorator.ts
vendored
Normal file
@@ -0,0 +1,71 @@
|
||||
import type { Context } from "../../context.ts";
|
||||
import type {
|
||||
CommitResult,
|
||||
Entry,
|
||||
EntryScan,
|
||||
EntryStructure,
|
||||
SessionStats,
|
||||
Storage,
|
||||
StorageBranchScan,
|
||||
UsageRow,
|
||||
UsageScan,
|
||||
Write,
|
||||
} from "../types.ts";
|
||||
import type { ListElement, ListReadOptions, StoredValue, Value, ValueList } from "../values.ts";
|
||||
|
||||
/** Test-only forwarding base for decorators that alter one part of Storage behavior. */
|
||||
export class StorageDecorator implements Storage {
|
||||
protected readonly delegate: Storage;
|
||||
|
||||
constructor(delegate: Storage) {
|
||||
this.delegate = delegate;
|
||||
}
|
||||
|
||||
commit(writes: Write[], context: Context): Promise<CommitResult> {
|
||||
return this.delegate.commit(writes, context);
|
||||
}
|
||||
|
||||
getEntries(ids: string[], context: Context): Promise<Map<string, Entry>> {
|
||||
return this.delegate.getEntries(ids, context);
|
||||
}
|
||||
|
||||
getValue<T>(address: Value<T>, context: Context): Promise<StoredValue<T> | undefined> {
|
||||
return this.delegate.getValue(address, context);
|
||||
}
|
||||
|
||||
scanValues<T>(prefix: Value<T>, context: Context): Promise<StoredValue<T>[]> {
|
||||
return this.delegate.scanValues(prefix, context);
|
||||
}
|
||||
|
||||
readList<T>(
|
||||
address: ValueList<T>,
|
||||
options: ListReadOptions | undefined,
|
||||
context: Context,
|
||||
): Promise<ListElement<T>[]> {
|
||||
return this.delegate.readList(address, options, context);
|
||||
}
|
||||
|
||||
scanBranch(query: StorageBranchScan, context: Context): Promise<Entry[]> {
|
||||
return this.delegate.scanBranch(query, context);
|
||||
}
|
||||
|
||||
scanBranchStructure(query: StorageBranchScan, context: Context): Promise<EntryStructure[]> {
|
||||
return this.delegate.scanBranchStructure(query, context);
|
||||
}
|
||||
|
||||
scanEntries(query: EntryScan, context: Context): Promise<Entry[]> {
|
||||
return this.delegate.scanEntries(query, context);
|
||||
}
|
||||
|
||||
scanUsage(query: UsageScan, context: Context): Promise<UsageRow[]> {
|
||||
return this.delegate.scanUsage(query, context);
|
||||
}
|
||||
|
||||
getStats(context: Context): Promise<SessionStats> {
|
||||
return this.delegate.getStats(context);
|
||||
}
|
||||
|
||||
close(context: Context): Promise<void> {
|
||||
return this.delegate.close(context);
|
||||
}
|
||||
}
|
||||
13
third_party/pi/packages/agent/src/harness/session/testing/types.ts
vendored
Normal file
13
third_party/pi/packages/agent/src/harness/session/testing/types.ts
vendored
Normal file
@@ -0,0 +1,13 @@
|
||||
import type { Storage } from "../types.ts";
|
||||
|
||||
/** A fresh backend storage instance owned by one test or benchmark case. */
|
||||
export interface StorageFixture extends AsyncDisposable {
|
||||
readonly storage: Storage;
|
||||
}
|
||||
|
||||
/** A runner-independent conformance case that can be registered with any test framework. */
|
||||
export interface ConformanceCase {
|
||||
readonly group: string;
|
||||
readonly name: string;
|
||||
run(): Promise<void>;
|
||||
}
|
||||
601
third_party/pi/packages/agent/src/harness/session/types.ts
vendored
Normal file
601
third_party/pi/packages/agent/src/harness/session/types.ts
vendored
Normal file
@@ -0,0 +1,601 @@
|
||||
import type { JsonValue } from "@earendil-works/chord";
|
||||
import type { AssistantMessage, StopReason, Usage } from "@earendil-works/pi-ai";
|
||||
import type { AgentMessage, QueueMode, ThinkingLevel } from "../../types.ts";
|
||||
import type { BranchPreparation } from "../compaction/branch-summarization.ts";
|
||||
import type { CompactionPreparation, CompactionSettings } from "../compaction/compaction.ts";
|
||||
import type { Context } from "../context.ts";
|
||||
import type { AgentHarnessStreamOptions } from "../types.ts";
|
||||
import type { ListElement, ListReadOptions, ListWrite, StoredValue, Value, ValueList, ValueWrite } from "./values.ts";
|
||||
|
||||
export type { JsonValue } from "@earendil-works/chord";
|
||||
|
||||
export type SettledAssistantMessage = AssistantMessage & {
|
||||
stopReason: Exclude<StopReason, "pending">;
|
||||
};
|
||||
|
||||
export type EntryType = "message" | "compaction" | "branch_summary" | "custom";
|
||||
|
||||
export interface EntryBase {
|
||||
id: string;
|
||||
parentId: string | null;
|
||||
seq: number;
|
||||
timestamp: number;
|
||||
type: EntryType;
|
||||
customType?: string;
|
||||
}
|
||||
|
||||
export interface MessageEntry extends EntryBase {
|
||||
type: "message";
|
||||
message: AgentMessage;
|
||||
terminate?: true;
|
||||
}
|
||||
|
||||
export interface CompactionEntry extends EntryBase {
|
||||
type: "compaction";
|
||||
summary: string;
|
||||
retainedTail: AgentMessage[];
|
||||
tokensBefore: number;
|
||||
details?: JsonValue;
|
||||
usage?: Usage;
|
||||
fromHook: boolean;
|
||||
}
|
||||
|
||||
export interface BranchSummaryEntry extends EntryBase {
|
||||
type: "branch_summary";
|
||||
fromId: string | null;
|
||||
summary: string;
|
||||
details?: JsonValue;
|
||||
usage?: Usage;
|
||||
fromHook: boolean;
|
||||
}
|
||||
|
||||
export interface CustomEntry extends EntryBase {
|
||||
type: "custom";
|
||||
customType: string;
|
||||
data?: JsonValue;
|
||||
}
|
||||
|
||||
/** Convert an application-defined custom entry into model context. */
|
||||
export type EntryProjector = (
|
||||
entry: CustomEntry,
|
||||
context: Context,
|
||||
) => AgentMessage[] | undefined | Promise<AgentMessage[] | undefined>;
|
||||
|
||||
export type Entry = MessageEntry | CompactionEntry | BranchSummaryEntry | CustomEntry;
|
||||
|
||||
/** Entry supplied to a transaction before storage assigns sequence and timestamp. */
|
||||
export type NewEntry<TEntry extends Entry = Entry> = TEntry extends Entry ? Omit<TEntry, "seq" | "timestamp"> : never;
|
||||
|
||||
export interface LaneConfiguration {
|
||||
model: { provider: string; modelId: string };
|
||||
thinkingLevel: ThinkingLevel;
|
||||
activeToolNames: string[];
|
||||
}
|
||||
|
||||
export interface OperationMeta {
|
||||
operationId: string;
|
||||
lane: string;
|
||||
sourceTipId: string | null;
|
||||
startedAt: number;
|
||||
intent:
|
||||
| { kind: "run"; promptEntryIds: string[] }
|
||||
| { kind: "compaction"; customInstructions?: string }
|
||||
| {
|
||||
kind: "navigation";
|
||||
targetId: string | null;
|
||||
summarize: boolean;
|
||||
label?: string;
|
||||
customInstructions?: string;
|
||||
};
|
||||
}
|
||||
|
||||
export type Control =
|
||||
| { status: "running" }
|
||||
| {
|
||||
status: "cancel_requested";
|
||||
requestedAt: number;
|
||||
};
|
||||
|
||||
export interface OperationError {
|
||||
code: string;
|
||||
message: string;
|
||||
details?: JsonValue;
|
||||
}
|
||||
|
||||
export type TerminalStatus = "completed" | "declined" | "aborted" | "failed";
|
||||
|
||||
/** Immutable lane-lived observation record written by one terminal transaction. */
|
||||
export interface OperationResultRecord {
|
||||
operationId: string;
|
||||
kind: OperationMeta["intent"]["kind"];
|
||||
status: TerminalStatus;
|
||||
error?: OperationError;
|
||||
fromTipId: string | null;
|
||||
tipId: string | null;
|
||||
startedAt: number;
|
||||
endedAt: number;
|
||||
}
|
||||
|
||||
export type Continuation =
|
||||
| { kind: "need_assistant"; overflowRecoveryUsed: boolean }
|
||||
| { kind: "may_finish"; includeFinalAssistant: boolean };
|
||||
|
||||
/** Checkpoint payload; the flat leaf literal replaces the old nested phase tag. */
|
||||
export interface CheckpointData {
|
||||
continuation: Continuation;
|
||||
triggerEntryId: string;
|
||||
}
|
||||
|
||||
export type InboxItemKind = "steer" | "followUp" | "nextRun" | "write";
|
||||
|
||||
export interface InboxItem {
|
||||
entryId: string;
|
||||
kind: InboxItemKind;
|
||||
}
|
||||
|
||||
export interface NormalizedRetryPolicy {
|
||||
maxAttempts: number;
|
||||
baseDelayMs: number;
|
||||
}
|
||||
|
||||
export interface GenerationContext {
|
||||
stepId: string;
|
||||
triggerEntryId: string;
|
||||
configuration: LaneConfiguration;
|
||||
streamOptions: AgentHarnessStreamOptions;
|
||||
retryPolicy: NormalizedRetryPolicy;
|
||||
overflowRecoveryUsed: boolean;
|
||||
}
|
||||
|
||||
interface ToolCallSource {
|
||||
/** Zero-based index in the assistant message's complete content array, not a filtered tool-call ordinal. */
|
||||
sourceIndex: number;
|
||||
resultEntryId: string;
|
||||
}
|
||||
|
||||
export type ToolCall = ToolCallSource &
|
||||
(
|
||||
| { status: "planned" }
|
||||
| { status: "effect_pending"; replay: "never" | "safe" }
|
||||
| { status: "outcome_ready"; terminate: boolean }
|
||||
| { status: "completed"; terminate: boolean }
|
||||
);
|
||||
|
||||
export interface ToolBatch {
|
||||
assistantEntryId: string;
|
||||
configuration: LaneConfiguration;
|
||||
turnId: string;
|
||||
calls: ToolCall[];
|
||||
}
|
||||
|
||||
export interface SummaryContext {
|
||||
resultEntryId: string;
|
||||
configuration: LaneConfiguration;
|
||||
streamOptions: AgentHarnessStreamOptions;
|
||||
retryPolicy: NormalizedRetryPolicy;
|
||||
}
|
||||
|
||||
/*
|
||||
* Durable operation state is one flat union with one family-neutral discriminator
|
||||
* per dispatcher leaf. ToolBatch/ToolCall remain the nested child collection state
|
||||
* machine, and cancellation stays orthogonal via Control.
|
||||
*/
|
||||
|
||||
export interface Cancellable {
|
||||
control: Control;
|
||||
}
|
||||
|
||||
export interface RunSettings {
|
||||
compaction: CompactionSettings;
|
||||
steeringMode: QueueMode;
|
||||
followUpMode: QueueMode;
|
||||
toolExecution: "sequential" | "parallel";
|
||||
}
|
||||
|
||||
/** Uniform scope carried by every operation leaf. */
|
||||
export interface OperationScope extends Cancellable {
|
||||
settings: RunSettings;
|
||||
latestAssistantEntryId: string | null;
|
||||
}
|
||||
|
||||
/** Shared backoff data for every retry-wait leaf. */
|
||||
export interface RetryWait {
|
||||
nextAttempt: number;
|
||||
notBefore: number;
|
||||
errorMessage: string;
|
||||
}
|
||||
|
||||
export interface AssistantGenerationScope {
|
||||
generationContext: GenerationContext;
|
||||
}
|
||||
|
||||
export type ResultBoundary =
|
||||
| { kind: "resume_checkpoint"; resumeAfter: CheckpointData }
|
||||
| { kind: "finish" }
|
||||
| { kind: "commit_navigation"; targetId: string; label?: string };
|
||||
|
||||
export interface SummaryTask {
|
||||
taskId: string;
|
||||
reason?: "manual" | "threshold" | "overflow";
|
||||
customInstructions?: string;
|
||||
boundary: ResultBoundary;
|
||||
}
|
||||
|
||||
export interface SummaryGenerationScope {
|
||||
task: SummaryTask;
|
||||
summaryContext: SummaryContext;
|
||||
}
|
||||
|
||||
export interface SummaryGenerationReady extends SummaryGenerationScope {
|
||||
nextAttempt: number;
|
||||
}
|
||||
|
||||
export interface SummaryGenerationEffectPending extends SummaryGenerationScope {
|
||||
attempt: number;
|
||||
request?: { index: number; usageId: string };
|
||||
usageIds: string[];
|
||||
}
|
||||
|
||||
export interface SummaryGenerationRetryWait extends SummaryGenerationScope, RetryWait {}
|
||||
|
||||
export interface DeferredScope extends OperationScope {
|
||||
stepId: string;
|
||||
sourceEntryId: string;
|
||||
poll: number;
|
||||
configuration: LaneConfiguration;
|
||||
streamOptions: AgentHarnessStreamOptions;
|
||||
}
|
||||
|
||||
export interface StartingOperation extends OperationScope {
|
||||
at: "starting";
|
||||
}
|
||||
|
||||
export interface CheckpointOperation extends OperationScope, CheckpointData {
|
||||
at: "checkpoint";
|
||||
}
|
||||
|
||||
export interface AssistantReadyOperation extends OperationScope, AssistantGenerationScope {
|
||||
at: "assistant.ready";
|
||||
nextAttempt: number;
|
||||
}
|
||||
|
||||
export interface AssistantEffectPendingOperation extends OperationScope, AssistantGenerationScope {
|
||||
at: "assistant.effect_pending";
|
||||
attempt: number;
|
||||
responseEntryId: string;
|
||||
usageId: string;
|
||||
intendedOutputLimit: number;
|
||||
contextWindow: number;
|
||||
}
|
||||
|
||||
export interface AssistantRetryWaitOperation extends OperationScope, AssistantGenerationScope, RetryWait {
|
||||
at: "assistant.retry_wait";
|
||||
}
|
||||
|
||||
export interface ToolsOperation extends OperationScope {
|
||||
at: "tools";
|
||||
batch: ToolBatch;
|
||||
}
|
||||
|
||||
export interface DeferredSuspendedOperation extends DeferredScope {
|
||||
at: "deferred.suspended";
|
||||
}
|
||||
|
||||
export interface DeferredEffectPendingOperation extends DeferredScope {
|
||||
at: "deferred.effect_pending";
|
||||
responseEntryId: string;
|
||||
usageId: string;
|
||||
}
|
||||
|
||||
export interface SummaryDecidingOperation extends OperationScope {
|
||||
at: "summary.deciding";
|
||||
task: SummaryTask;
|
||||
}
|
||||
|
||||
export interface SummaryReadyOperation extends OperationScope, SummaryGenerationReady {
|
||||
at: "summary.ready";
|
||||
}
|
||||
|
||||
export interface SummaryEffectPendingOperation extends OperationScope, SummaryGenerationEffectPending {
|
||||
at: "summary.effect_pending";
|
||||
}
|
||||
|
||||
export interface SummaryRetryWaitOperation extends OperationScope, SummaryGenerationRetryWait {
|
||||
at: "summary.retry_wait";
|
||||
}
|
||||
|
||||
export interface NavigationReadyToCommitOperation extends OperationScope {
|
||||
at: "navigation.ready_to_commit";
|
||||
/** Unsummarized navigation may target the branch root (null). */
|
||||
targetId: string | null;
|
||||
label?: string;
|
||||
}
|
||||
|
||||
/** Flat durable operation state: exactly 13 family-neutral dispatcher leaves. */
|
||||
export type OperationState =
|
||||
| StartingOperation
|
||||
| CheckpointOperation
|
||||
| AssistantReadyOperation
|
||||
| AssistantEffectPendingOperation
|
||||
| AssistantRetryWaitOperation
|
||||
| ToolsOperation
|
||||
| DeferredSuspendedOperation
|
||||
| DeferredEffectPendingOperation
|
||||
| SummaryDecidingOperation
|
||||
| SummaryReadyOperation
|
||||
| SummaryEffectPendingOperation
|
||||
| SummaryRetryWaitOperation
|
||||
| NavigationReadyToCommitOperation;
|
||||
|
||||
export type OperationAt = OperationState["at"];
|
||||
|
||||
/** Copy only the uniform operation scope when constructing a successor leaf. */
|
||||
export function operationScopeOf(state: OperationState): OperationScope {
|
||||
return {
|
||||
control: state.control,
|
||||
settings: state.settings,
|
||||
latestAssistantEntryId: state.latestAssistantEntryId,
|
||||
};
|
||||
}
|
||||
|
||||
export type Operation = { meta: OperationMeta; state: OperationState };
|
||||
|
||||
export interface LaneState {
|
||||
currentOperationId: string | null;
|
||||
lastOperationId: string | null;
|
||||
inbox: InboxItem[];
|
||||
}
|
||||
|
||||
export type PendingEntry =
|
||||
| { type: "message"; payload: AgentMessage }
|
||||
| { type: "custom"; customType: string; payload?: JsonValue };
|
||||
|
||||
export interface DurableFileOperations {
|
||||
read: string[];
|
||||
written: string[];
|
||||
edited: string[];
|
||||
}
|
||||
|
||||
export type DurableStructuralPreparation =
|
||||
| {
|
||||
kind: "compaction";
|
||||
messagesToSummarize: CompactionPreparation["messagesToSummarize"];
|
||||
turnPrefixMessages: CompactionPreparation["turnPrefixMessages"];
|
||||
retainedTail: CompactionPreparation["retainedTail"];
|
||||
isSplitTurn: boolean;
|
||||
tokensBefore: number;
|
||||
previousSummary?: string;
|
||||
fileOps: DurableFileOperations;
|
||||
settings: CompactionSettings;
|
||||
}
|
||||
| {
|
||||
kind: "branch_summary";
|
||||
messages: BranchPreparation["messages"];
|
||||
fileOps: DurableFileOperations;
|
||||
totalTokens: number;
|
||||
};
|
||||
|
||||
export interface UsageRow {
|
||||
id: string;
|
||||
seq: number;
|
||||
usage: Usage;
|
||||
entryId?: string;
|
||||
adjustment: boolean;
|
||||
details?: JsonValue;
|
||||
}
|
||||
|
||||
export interface EntryWrite {
|
||||
kind: "entry";
|
||||
entry: NewEntry;
|
||||
}
|
||||
|
||||
export interface UsageWrite {
|
||||
kind: "usage";
|
||||
row: Omit<UsageRow, "seq">;
|
||||
}
|
||||
|
||||
export type Write = EntryWrite | UsageWrite | ValueWrite | ListWrite;
|
||||
|
||||
export interface CommitResult {
|
||||
firstSeq: number;
|
||||
seqs: number[];
|
||||
timestamp: number;
|
||||
/** Session totals immediately after this commit was applied. */
|
||||
stats: SessionStats;
|
||||
}
|
||||
|
||||
export interface EntryStructure {
|
||||
id: string;
|
||||
parentId: string | null;
|
||||
seq: number;
|
||||
timestamp: number;
|
||||
type: EntryType;
|
||||
customType?: string;
|
||||
}
|
||||
|
||||
export interface EntryCursor {
|
||||
seq: number;
|
||||
}
|
||||
|
||||
export interface BranchScan {
|
||||
start?: string;
|
||||
stopAtType?: EntryType;
|
||||
stopAtId?: string;
|
||||
type?: EntryType;
|
||||
customType?: string;
|
||||
order?: "newestFirst" | "oldestFirst";
|
||||
limit?: number;
|
||||
cursor?: EntryCursor;
|
||||
}
|
||||
|
||||
export type StorageBranchScan = BranchScan & { start: string };
|
||||
|
||||
export interface EntryScan {
|
||||
type?: EntryType;
|
||||
customType?: string;
|
||||
fromSeq?: number;
|
||||
toSeq?: number;
|
||||
order?: "asc" | "desc";
|
||||
limit?: number;
|
||||
}
|
||||
|
||||
export interface UsageScan {
|
||||
fromSeq?: number;
|
||||
toSeq?: number;
|
||||
order?: "asc" | "desc";
|
||||
limit?: number;
|
||||
}
|
||||
|
||||
export interface SessionStats {
|
||||
messageCount: number;
|
||||
usage: Usage;
|
||||
}
|
||||
|
||||
export interface Storage {
|
||||
commit(writes: Write[], context: Context): Promise<CommitResult>;
|
||||
getEntries(ids: string[], context: Context): Promise<Map<string, Entry>>;
|
||||
getValue<T>(address: Value<T>, context: Context): Promise<StoredValue<T> | undefined>;
|
||||
scanValues<T>(prefix: Value<T>, context: Context): Promise<StoredValue<T>[]>;
|
||||
readList<T>(
|
||||
address: ValueList<T>,
|
||||
options: ListReadOptions | undefined,
|
||||
context: Context,
|
||||
): Promise<ListElement<T>[]>;
|
||||
scanBranch(query: StorageBranchScan, context: Context): Promise<Entry[]>;
|
||||
scanBranchStructure(query: StorageBranchScan, context: Context): Promise<EntryStructure[]>;
|
||||
scanEntries(query: EntryScan, context: Context): Promise<Entry[]>;
|
||||
scanUsage(query: UsageScan, context: Context): Promise<UsageRow[]>;
|
||||
getStats(context: Context): Promise<SessionStats>;
|
||||
close(context: Context): Promise<void>;
|
||||
}
|
||||
|
||||
export interface SessionMetadata {
|
||||
id: string;
|
||||
createdAt: number;
|
||||
storageVersion: number;
|
||||
cwd?: string;
|
||||
parentSessionId?: string;
|
||||
legacyParentSessionPath?: string;
|
||||
}
|
||||
|
||||
export interface IdGenerator {
|
||||
next(timestampMs?: number): string;
|
||||
}
|
||||
|
||||
export interface EntryQuery {
|
||||
type?: EntryType;
|
||||
customType?: string;
|
||||
order?: "asc" | "desc";
|
||||
limit?: number;
|
||||
cursor?: EntryCursor;
|
||||
}
|
||||
|
||||
export interface SessionReader {
|
||||
getEntries(ids: string[], context: Context): Promise<Map<string, Entry>>;
|
||||
getStats(context: Context): Promise<SessionStats>;
|
||||
getValue<T>(address: Value<T>, context: Context): Promise<StoredValue<T> | undefined>;
|
||||
scanValues<T>(prefix: Value<T>, context: Context): Promise<StoredValue<T>[]>;
|
||||
readList<T>(
|
||||
address: ValueList<T>,
|
||||
options: ListReadOptions | undefined,
|
||||
context: Context,
|
||||
): Promise<ListElement<T>[]>;
|
||||
/** Scan a branch from an explicit entry while this reader capability remains valid. */
|
||||
scanBranch(query: StorageBranchScan, context: Context): Promise<Entry[]>;
|
||||
}
|
||||
|
||||
/** Exclusive keyless mutation barrier for one Session. */
|
||||
export interface SessionMutation extends SessionReader {
|
||||
/** Exactly zero or one commit attempt. A second attempt rejects. */
|
||||
commit(writes: Write[], context: Context): Promise<CommitResult>;
|
||||
/** Wait for any commit attempt, invalidate the capability, and release the barrier. */
|
||||
end(context: Context): Promise<void>;
|
||||
}
|
||||
|
||||
/** Callback-scoped mutation capability without authority to release its Session barrier. */
|
||||
export type SessionMutator = Omit<SessionMutation, "end">;
|
||||
|
||||
export type SessionMutationCallback<T> = (mutator: SessionMutator, context: Context) => T | Promise<T>;
|
||||
|
||||
export interface Branch {
|
||||
readonly name: string;
|
||||
getTipId(context: Context): Promise<string | null>;
|
||||
findEntries(query: BranchScan | undefined, context: Context): Promise<Entry[]>;
|
||||
findEntry(query: BranchScan | undefined, context: Context): Promise<Entry | undefined>;
|
||||
appendMessage(message: AgentMessage, context: Context): Promise<string>;
|
||||
appendCustomEntry(customType: string, data: JsonValue | undefined, context: Context): Promise<string>;
|
||||
}
|
||||
|
||||
export interface Session<TMetadata extends SessionMetadata = SessionMetadata> extends SessionReader {
|
||||
readonly metadata: TMetadata;
|
||||
readonly idGenerator: IdGenerator;
|
||||
getEntry(id: string, context: Context): Promise<Entry | undefined>;
|
||||
getStats(context: Context): Promise<SessionStats>;
|
||||
getName(context: Context): Promise<string | undefined>;
|
||||
getLabel(targetId: string, context: Context): Promise<string | undefined>;
|
||||
findEntries(query: EntryQuery | undefined, context: Context): Promise<Entry[]>;
|
||||
findEntry(query: EntryQuery | undefined, context: Context): Promise<Entry | undefined>;
|
||||
branch(name: string, context: Context): Promise<Branch | undefined>;
|
||||
createBranch(name: string, at: string | null, context: Context): Promise<Branch>;
|
||||
beginMutation(context: Context): Promise<SessionMutation>;
|
||||
/**
|
||||
* Trusted exclusive callback over the Session mutation line. Calling a public Session writer from
|
||||
* this callback queues it behind the callback; awaiting that nested writer therefore deadlocks.
|
||||
* Use the supplied mutator for the callback's sole commit.
|
||||
*/
|
||||
mutate<T>(mutation: SessionMutationCallback<T>, context: Context): Promise<T>;
|
||||
setValue<T>(address: Value<T>, next: NoInfer<T>, context: Context): Promise<void>;
|
||||
deleteValue<T>(address: Value<T>, context: Context): Promise<void>;
|
||||
appendList<T>(address: ValueList<T>, element: NoInfer<T>, context: Context): Promise<void>;
|
||||
deleteList<T>(address: ValueList<T>, context: Context): Promise<void>;
|
||||
setName(name: string | undefined, context: Context): Promise<void>;
|
||||
setLabel(targetId: string, label: string | undefined, context: Context): Promise<void>;
|
||||
close(context: Context): Promise<void>;
|
||||
}
|
||||
|
||||
export interface SessionCreateOptions {
|
||||
id?: string;
|
||||
parentSessionId?: string;
|
||||
}
|
||||
|
||||
export type ForkOptions =
|
||||
| {
|
||||
/**
|
||||
* Copy one path from a complete configured source AgentLane under the same
|
||||
* Branch name, with copied configuration and fresh idle lane state.
|
||||
*/
|
||||
scope: "branch";
|
||||
/** Source Branch to copy. */
|
||||
branch: string;
|
||||
/** Entry on the source Branch's current tip ancestry. Defaults to the current tip. */
|
||||
entryId?: string;
|
||||
/**
|
||||
* Whether the fork includes the selected entry or stops at its parent.
|
||||
* Defaults to including the selected entry.
|
||||
*/
|
||||
position?: "before" | "at";
|
||||
/** Optional destination session id. */
|
||||
id?: string;
|
||||
}
|
||||
| {
|
||||
/**
|
||||
* Copy the whole conversation tree and every Branch tip. Each configured
|
||||
* AgentLane copies configuration plus fresh idle state; data-only Branches
|
||||
* remain data-only. Operation/pending/result/usage state is excluded.
|
||||
*/
|
||||
scope: "tree";
|
||||
/** Optional destination session id. */
|
||||
id?: string;
|
||||
};
|
||||
|
||||
export interface SessionRepo<
|
||||
TMetadata extends SessionMetadata = SessionMetadata,
|
||||
TCreateOptions extends { id?: string; parentSessionId?: string } = SessionCreateOptions,
|
||||
TListOptions = void,
|
||||
> {
|
||||
create(options: TCreateOptions, context: Context): Promise<Session<TMetadata>>;
|
||||
open(metadata: TMetadata, context: Context): Promise<Session<TMetadata>>;
|
||||
list(options: TListOptions | undefined, context: Context): Promise<TMetadata[]>;
|
||||
delete(metadata: TMetadata, context: Context): Promise<void>;
|
||||
fork(source: TMetadata, options: ForkOptions, context: Context): Promise<Session<TMetadata>>;
|
||||
}
|
||||
195
third_party/pi/packages/agent/src/harness/session/values.ts
vendored
Normal file
195
third_party/pi/packages/agent/src/harness/session/values.ts
vendored
Normal file
@@ -0,0 +1,195 @@
|
||||
import type { AssistantMessageFrame } from "@earendil-works/pi-ai";
|
||||
import type { AgentToolResult } from "../../types.ts";
|
||||
import type {
|
||||
DurableStructuralPreparation,
|
||||
JsonValue,
|
||||
LaneConfiguration,
|
||||
LaneState,
|
||||
OperationMeta,
|
||||
OperationResultRecord,
|
||||
OperationState,
|
||||
PendingEntry,
|
||||
} from "./types.ts";
|
||||
|
||||
declare const storedValueType: unique symbol;
|
||||
|
||||
export interface StoredAddressBase {
|
||||
readonly namespace: string;
|
||||
readonly key: string;
|
||||
readonly kind: "value" | "list";
|
||||
}
|
||||
|
||||
export interface Value<T> extends StoredAddressBase {
|
||||
readonly kind: "value";
|
||||
readonly [storedValueType]?: (value: T) => T;
|
||||
}
|
||||
|
||||
export interface ValueList<T> extends StoredAddressBase {
|
||||
readonly kind: "list";
|
||||
readonly [storedValueType]?: (value: T) => T;
|
||||
}
|
||||
|
||||
export interface StoredValue<T> {
|
||||
address: Value<T>;
|
||||
value: T;
|
||||
seq: number;
|
||||
}
|
||||
|
||||
export interface ListElement<T> {
|
||||
seq: number;
|
||||
value: T;
|
||||
}
|
||||
|
||||
export interface ListCursor {
|
||||
seq: number;
|
||||
}
|
||||
|
||||
export interface ListReadOptions {
|
||||
cursor?: ListCursor;
|
||||
order?: "asc" | "desc";
|
||||
limit?: number;
|
||||
}
|
||||
|
||||
export interface ResolvedListReadOptions {
|
||||
cursor?: ListCursor;
|
||||
order: "asc" | "desc";
|
||||
limit: number;
|
||||
}
|
||||
|
||||
export interface ValueSetWrite {
|
||||
kind: "value";
|
||||
op: "set";
|
||||
namespace: string;
|
||||
key: string;
|
||||
value: unknown;
|
||||
}
|
||||
|
||||
export interface ValueDeleteWrite {
|
||||
kind: "value";
|
||||
op: "delete";
|
||||
namespace: string;
|
||||
key: string;
|
||||
}
|
||||
|
||||
export interface ListAppendWrite {
|
||||
kind: "list";
|
||||
op: "append";
|
||||
namespace: string;
|
||||
key: string;
|
||||
value: unknown;
|
||||
}
|
||||
|
||||
export interface ListDeleteWrite {
|
||||
kind: "list";
|
||||
op: "delete";
|
||||
namespace: string;
|
||||
key: string;
|
||||
}
|
||||
|
||||
export type ValueWrite = ValueSetWrite | ValueDeleteWrite;
|
||||
export type ListWrite = ListAppendWrite | ListDeleteWrite;
|
||||
|
||||
function validateAddress(namespace: string, key: string): void {
|
||||
if (namespace.length === 0) throw new TypeError("Value namespace must not be empty");
|
||||
if (namespace.includes("\u0000")) throw new TypeError("Value namespace must not contain \\u0000");
|
||||
if (key.includes("\u0000")) throw new TypeError("Value key must not contain \\u0000");
|
||||
}
|
||||
|
||||
export function value<T>(namespace: string, key = ""): Value<T> {
|
||||
validateAddress(namespace, key);
|
||||
return Object.freeze({ namespace, key, kind: "value" }) as Value<T>;
|
||||
}
|
||||
|
||||
export function list<T>(namespace: string, key = ""): ValueList<T> {
|
||||
validateAddress(namespace, key);
|
||||
return Object.freeze({ namespace, key, kind: "list" }) as ValueList<T>;
|
||||
}
|
||||
|
||||
export function setValue<T>(address: Value<T>, next: NoInfer<T>): ValueSetWrite {
|
||||
return {
|
||||
kind: "value",
|
||||
op: "set",
|
||||
namespace: address.namespace,
|
||||
key: address.key,
|
||||
value: next,
|
||||
};
|
||||
}
|
||||
|
||||
export function deleteValue<T>(address: Value<T>): ValueDeleteWrite {
|
||||
return {
|
||||
kind: "value",
|
||||
op: "delete",
|
||||
namespace: address.namespace,
|
||||
key: address.key,
|
||||
};
|
||||
}
|
||||
|
||||
export function appendList<T>(address: ValueList<T>, element: NoInfer<T>): ListAppendWrite {
|
||||
return {
|
||||
kind: "list",
|
||||
op: "append",
|
||||
namespace: address.namespace,
|
||||
key: address.key,
|
||||
value: element,
|
||||
};
|
||||
}
|
||||
|
||||
export function deleteList<T>(address: ValueList<T>): ListDeleteWrite {
|
||||
return {
|
||||
kind: "list",
|
||||
op: "delete",
|
||||
namespace: address.namespace,
|
||||
key: address.key,
|
||||
};
|
||||
}
|
||||
|
||||
export function resolveListReadOptions(options: ListReadOptions = {}): ResolvedListReadOptions {
|
||||
const requestedLimit = options.limit ?? 1_000;
|
||||
if (!Number.isSafeInteger(requestedLimit) || requestedLimit <= 0) {
|
||||
throw new TypeError("List read limit must be a positive safe integer");
|
||||
}
|
||||
return {
|
||||
...(options.cursor === undefined ? {} : { cursor: options.cursor }),
|
||||
order: options.order ?? "asc",
|
||||
limit: Math.min(requestedLimit, 10_000),
|
||||
};
|
||||
}
|
||||
|
||||
export const branchTip = (branch: string) => value<string | null>("pi.branch.tip", branch);
|
||||
export const branchTipInventoryPrefix = () => value<string | null>("pi.branch.tip");
|
||||
export const laneConfig = (lane: string) => value<LaneConfiguration>("pi.lane.config", lane);
|
||||
export const laneState = (lane: string) => value<LaneState>("pi.lane.state", lane);
|
||||
export const operationResult = (operationId: string) => value<OperationResultRecord>("pi.result", operationId);
|
||||
|
||||
export const operationMeta = (operationId: string) => value<OperationMeta>("pi.op.meta", operationId);
|
||||
export const operationState = (operationId: string) => value<OperationState>("pi.op.state", operationId);
|
||||
export const operationToolArgs = (operationId: string, stepId: string, sourceIndex: number) =>
|
||||
value<Record<string, JsonValue>>("pi.op.tool_args", `${operationId}:${stepId}:${sourceIndex}`);
|
||||
export const operationToolMemo = (operationId: string, invocationId: string, name: string) =>
|
||||
value<JsonValue>("pi.op.tool_memo", `${operationId}:${invocationId}:${name}`);
|
||||
export const operationPreparation = (operationId: string, taskId: string) =>
|
||||
value<DurableStructuralPreparation>("pi.op.preparation", `${operationId}:${taskId}`);
|
||||
|
||||
export const operationToolArgsPrefix = (operationId: string, stepId?: string) =>
|
||||
value<Record<string, JsonValue>>(
|
||||
"pi.op.tool_args",
|
||||
stepId === undefined ? `${operationId}:` : `${operationId}:${stepId}:`,
|
||||
);
|
||||
export const operationToolMemoPrefix = (operationId: string, invocationId?: string) =>
|
||||
value<JsonValue>(
|
||||
"pi.op.tool_memo",
|
||||
invocationId === undefined ? `${operationId}:` : `${operationId}:${invocationId}:`,
|
||||
);
|
||||
export const operationPreparationPrefix = (operationId: string) =>
|
||||
value<DurableStructuralPreparation>("pi.op.preparation", `${operationId}:`);
|
||||
|
||||
export const pendingEntry = (entryId: string) => value<PendingEntry>("pi.pending.entry", entryId);
|
||||
export const pendingToolOutput = (operationId: string, invocationId: string) =>
|
||||
value<AgentToolResult<unknown>>("pi.pending.tool_output", `${operationId}:${invocationId}`);
|
||||
export const pendingAssistantFrames = (operationId: string, responseEntryId: string) =>
|
||||
list<AssistantMessageFrame>("pi.pending.assistant_frame", `${operationId}:${responseEntryId}`);
|
||||
export const pendingToolOutputPrefix = (operationId: string) =>
|
||||
value<AgentToolResult<unknown>>("pi.pending.tool_output", `${operationId}:`);
|
||||
|
||||
export const sessionName = value<string>("pi.session.name");
|
||||
export const entryLabel = (entryId: string) => value<string>("pi.entry.label", entryId);
|
||||
396
third_party/pi/packages/agent/src/harness/skills.ts
vendored
Normal file
396
third_party/pi/packages/agent/src/harness/skills.ts
vendored
Normal file
@@ -0,0 +1,396 @@
|
||||
import ignore from "ignore";
|
||||
import { parse } from "yaml";
|
||||
import type { Context } from "./context.ts";
|
||||
import { type ExecutionEnv, type FileInfo, type Result, type Skill, toError } from "./types.ts";
|
||||
|
||||
const MAX_NAME_LENGTH = 64;
|
||||
const MAX_DESCRIPTION_LENGTH = 1024;
|
||||
const IGNORE_FILE_NAMES = [".gitignore", ".ignore", ".fdignore"];
|
||||
|
||||
type IgnoreMatcher = ReturnType<typeof ignore>;
|
||||
|
||||
export type SkillDiagnosticCode =
|
||||
| "file_info_failed"
|
||||
| "list_failed"
|
||||
| "read_failed"
|
||||
| "parse_failed"
|
||||
| "invalid_metadata";
|
||||
|
||||
/** Warning produced while loading skills. */
|
||||
export interface SkillDiagnostic {
|
||||
/** Diagnostic severity. Currently only warnings are emitted. */
|
||||
type: "warning";
|
||||
/** Stable diagnostic code. */
|
||||
code: SkillDiagnosticCode;
|
||||
/** Human-readable diagnostic message. */
|
||||
message: string;
|
||||
/** Path associated with the diagnostic. */
|
||||
path: string;
|
||||
}
|
||||
|
||||
interface SkillFrontmatter {
|
||||
name?: string;
|
||||
description?: string;
|
||||
"disable-model-invocation"?: boolean;
|
||||
[key: string]: unknown;
|
||||
}
|
||||
|
||||
/** Format a skill invocation prompt, optionally appending additional user instructions. */
|
||||
export function formatSkillInvocation(skill: Skill, additionalInstructions?: string): string {
|
||||
const skillBlock = `<skill name="${skill.name}" location="${skill.filePath}">\nReferences are relative to ${dirnameEnvPath(skill.filePath)}.\n\n${skill.content}\n</skill>`;
|
||||
return additionalInstructions ? `${skillBlock}\n\n${additionalInstructions}` : skillBlock;
|
||||
}
|
||||
|
||||
/**
|
||||
* Load skills from one or more directories.
|
||||
*
|
||||
* Traverses directories recursively, loads `SKILL.md` files, loads direct root `.md` files with skill
|
||||
* frontmatter, honors ignore files, and returns diagnostics for invalid declared skill files. Missing input
|
||||
* directories are skipped.
|
||||
*/
|
||||
export async function loadSkills(
|
||||
env: ExecutionEnv,
|
||||
dirs: string | string[],
|
||||
context: Context,
|
||||
): Promise<{ skills: Skill[]; diagnostics: SkillDiagnostic[] }> {
|
||||
const skills: Skill[] = [];
|
||||
const diagnostics: SkillDiagnostic[] = [];
|
||||
for (const dir of Array.isArray(dirs) ? dirs : [dirs]) {
|
||||
const rootInfoResult = await env.fileInfo(dir, context);
|
||||
if (!rootInfoResult.ok) {
|
||||
if (rootInfoResult.error.code !== "not_found") {
|
||||
diagnostics.push({
|
||||
type: "warning",
|
||||
code: "file_info_failed",
|
||||
message: rootInfoResult.error.message,
|
||||
path: dir,
|
||||
});
|
||||
}
|
||||
continue;
|
||||
}
|
||||
const rootInfo = rootInfoResult.value;
|
||||
if ((await resolveKind(env, rootInfo, diagnostics, context)) !== "directory") continue;
|
||||
const result = await loadSkillsFromDirInternal(env, rootInfo.path, true, ignore(), rootInfo.path, context);
|
||||
skills.push(...result.skills);
|
||||
diagnostics.push(...result.diagnostics);
|
||||
}
|
||||
return { skills, diagnostics };
|
||||
}
|
||||
|
||||
/**
|
||||
* Load skills from source-tagged directories.
|
||||
*
|
||||
* Source values are preserved exactly and attached to every loaded skill and diagnostic. The agent package does not
|
||||
* interpret source values; applications define their own provenance shape.
|
||||
*/
|
||||
export async function loadSourcedSkills<TSource, TSkill extends Skill = Skill>(
|
||||
env: ExecutionEnv,
|
||||
inputs: Array<{ path: string; source: TSource }>,
|
||||
mapSkill: ((skill: Skill, source: TSource, context: Context) => TSkill) | undefined,
|
||||
context: Context,
|
||||
): Promise<{
|
||||
skills: Array<{ skill: TSkill; source: TSource }>;
|
||||
diagnostics: Array<SkillDiagnostic & { source: TSource }>;
|
||||
}> {
|
||||
const skills: Array<{ skill: TSkill; source: TSource }> = [];
|
||||
const diagnostics: Array<SkillDiagnostic & { source: TSource }> = [];
|
||||
for (const input of inputs) {
|
||||
const result = await loadSkills(env, input.path, context);
|
||||
for (const skill of result.skills) {
|
||||
skills.push({
|
||||
skill: mapSkill ? mapSkill(skill, input.source, context) : (skill as TSkill),
|
||||
source: input.source,
|
||||
});
|
||||
}
|
||||
for (const diagnostic of result.diagnostics) diagnostics.push({ ...diagnostic, source: input.source });
|
||||
}
|
||||
return { skills, diagnostics };
|
||||
}
|
||||
|
||||
async function loadSkillsFromDirInternal(
|
||||
env: ExecutionEnv,
|
||||
dir: string,
|
||||
includeRootFiles: boolean,
|
||||
ignoreMatcher: IgnoreMatcher,
|
||||
rootDir: string,
|
||||
context: Context,
|
||||
): Promise<{ skills: Skill[]; diagnostics: SkillDiagnostic[] }> {
|
||||
const skills: Skill[] = [];
|
||||
const diagnostics: SkillDiagnostic[] = [];
|
||||
|
||||
const dirInfoResult = await env.fileInfo(dir, context);
|
||||
if (!dirInfoResult.ok) {
|
||||
if (dirInfoResult.error.code !== "not_found") {
|
||||
diagnostics.push({
|
||||
type: "warning",
|
||||
code: "file_info_failed",
|
||||
message: dirInfoResult.error.message,
|
||||
path: dir,
|
||||
});
|
||||
}
|
||||
return { skills, diagnostics };
|
||||
}
|
||||
const dirInfo = dirInfoResult.value;
|
||||
if ((await resolveKind(env, dirInfo, diagnostics, context)) !== "directory") return { skills, diagnostics };
|
||||
|
||||
await addIgnoreRules(env, ignoreMatcher, dir, rootDir, diagnostics, context);
|
||||
|
||||
const entriesResult = await env.listDir(dir, context);
|
||||
if (!entriesResult.ok) {
|
||||
diagnostics.push({ type: "warning", code: "list_failed", message: entriesResult.error.message, path: dir });
|
||||
return { skills, diagnostics };
|
||||
}
|
||||
const entries = entriesResult.value;
|
||||
|
||||
for (const entry of entries) {
|
||||
if (entry.name !== "SKILL.md") continue;
|
||||
const fullPath = entry.path;
|
||||
const kind = await resolveKind(env, entry, diagnostics, context);
|
||||
if (kind !== "file") continue;
|
||||
const relPath = relativeEnvPath(rootDir, fullPath);
|
||||
if (ignoreMatcher.ignores(relPath)) continue;
|
||||
|
||||
const result = await loadSkillFromFile(env, fullPath, dirInfo.name, context);
|
||||
if (result.skill) skills.push(result.skill);
|
||||
diagnostics.push(...result.diagnostics);
|
||||
return { skills, diagnostics };
|
||||
}
|
||||
|
||||
for (const entry of entries.sort((a, b) => a.name.localeCompare(b.name))) {
|
||||
if (entry.name.startsWith(".") || entry.name === "node_modules") continue;
|
||||
const fullPath = entry.path;
|
||||
const kind = await resolveKind(env, entry, diagnostics, context);
|
||||
if (!kind) continue;
|
||||
|
||||
const relPath = relativeEnvPath(rootDir, fullPath);
|
||||
const ignorePath = kind === "directory" ? `${relPath}/` : relPath;
|
||||
if (ignoreMatcher.ignores(ignorePath)) continue;
|
||||
|
||||
if (kind === "directory") {
|
||||
const result = await loadSkillsFromDirInternal(env, fullPath, false, ignoreMatcher, rootDir, context);
|
||||
skills.push(...result.skills);
|
||||
diagnostics.push(...result.diagnostics);
|
||||
continue;
|
||||
}
|
||||
|
||||
if (kind !== "file" || !includeRootFiles || !entry.name.endsWith(".md")) continue;
|
||||
const result = await loadSkillFromFile(env, fullPath, dirInfo.name, context);
|
||||
if (result.skill) skills.push(result.skill);
|
||||
diagnostics.push(...result.diagnostics);
|
||||
}
|
||||
|
||||
return { skills, diagnostics };
|
||||
}
|
||||
|
||||
async function addIgnoreRules(
|
||||
env: ExecutionEnv,
|
||||
ig: IgnoreMatcher,
|
||||
dir: string,
|
||||
rootDir: string,
|
||||
diagnostics: SkillDiagnostic[],
|
||||
context: Context,
|
||||
): Promise<void> {
|
||||
const relativeDir = relativeEnvPath(rootDir, dir);
|
||||
const prefix = relativeDir ? `${relativeDir}/` : "";
|
||||
|
||||
for (const filename of IGNORE_FILE_NAMES) {
|
||||
const ignorePathResult = await env.joinPath([dir, filename], context);
|
||||
if (!ignorePathResult.ok) {
|
||||
diagnostics.push({
|
||||
type: "warning",
|
||||
code: "file_info_failed",
|
||||
message: ignorePathResult.error.message,
|
||||
path: dir,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
const ignorePath = ignorePathResult.value;
|
||||
const info = await env.fileInfo(ignorePath, context);
|
||||
if (!info.ok) {
|
||||
if (info.error.code !== "not_found") {
|
||||
diagnostics.push({
|
||||
type: "warning",
|
||||
code: "file_info_failed",
|
||||
message: info.error.message,
|
||||
path: ignorePath,
|
||||
});
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if (info.value.kind !== "file") continue;
|
||||
const content = await env.readTextFile(ignorePath, context);
|
||||
if (!content.ok) {
|
||||
diagnostics.push({ type: "warning", code: "read_failed", message: content.error.message, path: ignorePath });
|
||||
continue;
|
||||
}
|
||||
const patterns = content.value
|
||||
.split(/\r?\n/)
|
||||
.map((line) => prefixIgnorePattern(line, prefix))
|
||||
.filter((line): line is string => Boolean(line));
|
||||
if (patterns.length > 0) ig.add(patterns);
|
||||
}
|
||||
}
|
||||
|
||||
function prefixIgnorePattern(line: string, prefix: string): string | null {
|
||||
const trimmed = line.trim();
|
||||
if (!trimmed) return null;
|
||||
if (trimmed.startsWith("#") && !trimmed.startsWith("\\#")) return null;
|
||||
|
||||
let pattern = line;
|
||||
let negated = false;
|
||||
if (pattern.startsWith("!")) {
|
||||
negated = true;
|
||||
pattern = pattern.slice(1);
|
||||
} else if (pattern.startsWith("\\!")) {
|
||||
pattern = pattern.slice(1);
|
||||
}
|
||||
if (pattern.startsWith("/")) pattern = pattern.slice(1);
|
||||
const prefixed = prefix ? `${prefix}${pattern}` : pattern;
|
||||
return negated ? `!${prefixed}` : prefixed;
|
||||
}
|
||||
|
||||
async function loadSkillFromFile(
|
||||
env: ExecutionEnv,
|
||||
filePath: string,
|
||||
parentDirName: string,
|
||||
context: Context,
|
||||
): Promise<{ skill: Skill | null; diagnostics: SkillDiagnostic[] }> {
|
||||
const diagnostics: SkillDiagnostic[] = [];
|
||||
const isDeclaredSkill =
|
||||
filePath
|
||||
.replace(/[\\/]+$/, "")
|
||||
.split(/[\\/]/)
|
||||
.pop() === "SKILL.md";
|
||||
const rawContent = await env.readTextFile(filePath, context);
|
||||
if (!rawContent.ok) {
|
||||
diagnostics.push({ type: "warning", code: "read_failed", message: rawContent.error.message, path: filePath });
|
||||
return { skill: null, diagnostics };
|
||||
}
|
||||
|
||||
const parsed = parseFrontmatter<SkillFrontmatter>(rawContent.value);
|
||||
if (!parsed.ok) {
|
||||
if (isDeclaredSkill) {
|
||||
diagnostics.push({ type: "warning", code: "parse_failed", message: parsed.error.message, path: filePath });
|
||||
}
|
||||
return { skill: null, diagnostics };
|
||||
}
|
||||
|
||||
const { frontmatter, body } = parsed.value;
|
||||
const description = typeof frontmatter.description === "string" ? frontmatter.description : undefined;
|
||||
if (!isDeclaredSkill && (!description || description.trim() === "")) {
|
||||
return { skill: null, diagnostics };
|
||||
}
|
||||
|
||||
for (const error of validateDescription(description)) {
|
||||
diagnostics.push({ type: "warning", code: "invalid_metadata", message: error, path: filePath });
|
||||
}
|
||||
|
||||
const frontmatterName = typeof frontmatter.name === "string" ? frontmatter.name : undefined;
|
||||
const name = frontmatterName || parentDirName;
|
||||
for (const error of validateName(name, parentDirName)) {
|
||||
diagnostics.push({ type: "warning", code: "invalid_metadata", message: error, path: filePath });
|
||||
}
|
||||
|
||||
if (!description || description.trim() === "") {
|
||||
return { skill: null, diagnostics };
|
||||
}
|
||||
|
||||
return {
|
||||
skill: {
|
||||
name,
|
||||
description,
|
||||
content: body,
|
||||
filePath,
|
||||
disableModelInvocation: frontmatter["disable-model-invocation"] === true,
|
||||
},
|
||||
diagnostics,
|
||||
};
|
||||
}
|
||||
|
||||
function validateName(name: string, parentDirName: string): string[] {
|
||||
const errors: string[] = [];
|
||||
if (name !== parentDirName) errors.push(`name "${name}" does not match parent directory "${parentDirName}"`);
|
||||
if (name.length > MAX_NAME_LENGTH) errors.push(`name exceeds ${MAX_NAME_LENGTH} characters (${name.length})`);
|
||||
if (!/^[a-z0-9-]+$/.test(name)) {
|
||||
errors.push("name contains invalid characters (must be lowercase a-z, 0-9, hyphens only)");
|
||||
}
|
||||
if (name.startsWith("-") || name.endsWith("-")) errors.push("name must not start or end with a hyphen");
|
||||
if (name.includes("--")) errors.push("name must not contain consecutive hyphens");
|
||||
return errors;
|
||||
}
|
||||
|
||||
function validateDescription(description: string | undefined): string[] {
|
||||
const errors: string[] = [];
|
||||
if (!description || description.trim() === "") {
|
||||
errors.push("description is required");
|
||||
} else if (description.length > MAX_DESCRIPTION_LENGTH) {
|
||||
errors.push(`description exceeds ${MAX_DESCRIPTION_LENGTH} characters (${description.length})`);
|
||||
}
|
||||
return errors;
|
||||
}
|
||||
|
||||
function parseFrontmatter<T extends Record<string, unknown>>(
|
||||
content: string,
|
||||
): Result<{ frontmatter: T; body: string }, Error> {
|
||||
try {
|
||||
const normalized = content.replace(/\r\n/g, "\n").replace(/\r/g, "\n");
|
||||
if (!normalized.startsWith("---")) return { ok: true, value: { frontmatter: {} as T, body: normalized } };
|
||||
const endIndex = normalized.indexOf("\n---", 3);
|
||||
if (endIndex === -1) return { ok: true, value: { frontmatter: {} as T, body: normalized } };
|
||||
const yamlString = normalized.slice(4, endIndex);
|
||||
const body = normalized.slice(endIndex + 4).trim();
|
||||
return { ok: true, value: { frontmatter: (parse(yamlString) ?? {}) as T, body } };
|
||||
} catch (error) {
|
||||
return { ok: false, error: toError(error) };
|
||||
}
|
||||
}
|
||||
|
||||
async function resolveKind(
|
||||
env: ExecutionEnv,
|
||||
info: FileInfo,
|
||||
diagnostics: SkillDiagnostic[],
|
||||
context: Context,
|
||||
): Promise<"file" | "directory" | undefined> {
|
||||
if (info.kind === "file" || info.kind === "directory") return info.kind;
|
||||
const canonicalPath = await env.canonicalPath(info.path, context);
|
||||
if (!canonicalPath.ok) {
|
||||
if (canonicalPath.error.code !== "not_found") {
|
||||
diagnostics.push({
|
||||
type: "warning",
|
||||
code: "file_info_failed",
|
||||
message: canonicalPath.error.message,
|
||||
path: info.path,
|
||||
});
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
const target = await env.fileInfo(canonicalPath.value, context);
|
||||
if (!target.ok) {
|
||||
if (target.error.code !== "not_found") {
|
||||
diagnostics.push({
|
||||
type: "warning",
|
||||
code: "file_info_failed",
|
||||
message: target.error.message,
|
||||
path: info.path,
|
||||
});
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
return target.value.kind === "file" || target.value.kind === "directory" ? target.value.kind : undefined;
|
||||
}
|
||||
|
||||
function dirnameEnvPath(path: string): string {
|
||||
const normalized = path.replace(/[\\/]+$/, "");
|
||||
const separatorIndex = Math.max(normalized.lastIndexOf("/"), normalized.lastIndexOf("\\"));
|
||||
if (separatorIndex === 2 && normalized[1] === ":") return normalized.slice(0, 3);
|
||||
return separatorIndex <= 0 ? "/" : normalized.slice(0, separatorIndex);
|
||||
}
|
||||
|
||||
function relativeEnvPath(root: string, path: string): string {
|
||||
const normalizedRoot = root.replace(/\\/g, "/").replace(/\/+$/, "");
|
||||
const normalizedPath = path.replace(/\\/g, "/").replace(/\/+$/, "");
|
||||
if (normalizedPath === normalizedRoot) return "";
|
||||
return normalizedPath.startsWith(`${normalizedRoot}/`)
|
||||
? normalizedPath.slice(normalizedRoot.length + 1)
|
||||
: normalizedPath.replace(/^\/+/, "");
|
||||
}
|
||||
34
third_party/pi/packages/agent/src/harness/system-prompt.ts
vendored
Normal file
34
third_party/pi/packages/agent/src/harness/system-prompt.ts
vendored
Normal file
@@ -0,0 +1,34 @@
|
||||
import type { Skill } from "./types.ts";
|
||||
|
||||
export function formatSkillsForSystemPrompt(skills: Skill[]): string {
|
||||
const visibleSkills = skills.filter((skill) => !skill.disableModelInvocation);
|
||||
if (visibleSkills.length === 0) return "";
|
||||
|
||||
const lines = [
|
||||
"The following skills provide specialized instructions for specific tasks.",
|
||||
"Read the full skill file when the task matches its description.",
|
||||
"When a skill file references a relative path, resolve it against the skill directory (parent of SKILL.md / dirname of the path) and use that absolute path in tool commands.",
|
||||
"",
|
||||
"<available_skills>",
|
||||
];
|
||||
|
||||
for (const skill of visibleSkills) {
|
||||
lines.push(" <skill>");
|
||||
lines.push(` <name>${escapeXml(skill.name)}</name>`);
|
||||
lines.push(` <description>${escapeXml(skill.description)}</description>`);
|
||||
lines.push(` <location>${escapeXml(skill.filePath)}</location>`);
|
||||
lines.push(" </skill>");
|
||||
}
|
||||
|
||||
lines.push("</available_skills>");
|
||||
return lines.join("\n");
|
||||
}
|
||||
|
||||
function escapeXml(value: string): string {
|
||||
return value
|
||||
.replace(/&/g, "&")
|
||||
.replace(/</g, "<")
|
||||
.replace(/>/g, ">")
|
||||
.replace(/"/g, """)
|
||||
.replace(/'/g, "'");
|
||||
}
|
||||
635
third_party/pi/packages/agent/src/harness/telemetry.ts
vendored
Normal file
635
third_party/pi/packages/agent/src/harness/telemetry.ts
vendored
Normal file
@@ -0,0 +1,635 @@
|
||||
import type {
|
||||
ExactTelemetryAttributes,
|
||||
SchemaTelemetrySpan,
|
||||
TelemetrySchemaDefinition,
|
||||
TelemetrySchemaSpanEndAttributes,
|
||||
TelemetrySchemaSpanEventAttributes,
|
||||
TelemetrySchemaSpanEventName,
|
||||
TelemetrySchemaSpanName,
|
||||
TelemetrySchemaSpanStartAttributes,
|
||||
TelemetrySchemaSpanUnion,
|
||||
TelemetrySpan,
|
||||
} from "@earendil-works/pi-telemetry";
|
||||
import { type Context, getTelemetryContext, withTelemetryContext } from "./context.ts";
|
||||
|
||||
export type {
|
||||
AttributeValue,
|
||||
ExactTelemetryAttributes,
|
||||
SchemaTelemetrySpan,
|
||||
SpanAttributes,
|
||||
SpanOptions,
|
||||
SpanStatus,
|
||||
TelemetryAttributeDefinition,
|
||||
TelemetryAttributeMetadata,
|
||||
TelemetryAttributeType,
|
||||
TelemetryContext,
|
||||
TelemetryEventAttributeDefinition,
|
||||
TelemetryEventDefinition,
|
||||
TelemetryParentDefinition,
|
||||
TelemetrySchemaDefinition,
|
||||
TelemetrySchemaSpanEndAttributes,
|
||||
TelemetrySchemaSpanEventAttributes,
|
||||
TelemetrySchemaSpanEventName,
|
||||
TelemetrySchemaSpanName,
|
||||
TelemetrySchemaSpanStartAttributes,
|
||||
TelemetrySchemaSpanUnion,
|
||||
TelemetrySpan,
|
||||
TelemetrySpanDefinition,
|
||||
TelemetryStartAttributeDefinition,
|
||||
TypedSpanStarter,
|
||||
} from "@earendil-works/pi-telemetry";
|
||||
|
||||
export const AI_TELEMETRY_SCHEMA = {
|
||||
version: 1,
|
||||
spans: {
|
||||
"pi.ai.request": {
|
||||
description: "One logical request to an AI provider",
|
||||
parents: { kind: "any" },
|
||||
startAttributes: {
|
||||
"pi.ai.operation": {
|
||||
type: "string",
|
||||
required: true,
|
||||
values: ["stream", "fetch_deferred", "cancel_deferred", "generate_images"],
|
||||
description: "Logical provider operation",
|
||||
},
|
||||
"pi.ai.provider": {
|
||||
type: "string",
|
||||
required: true,
|
||||
description: "Selected provider id",
|
||||
},
|
||||
"pi.ai.model": {
|
||||
type: "string",
|
||||
required: true,
|
||||
description: "Requested model id",
|
||||
},
|
||||
"pi.ai.api": {
|
||||
type: "string",
|
||||
required: true,
|
||||
description: "Provider API id",
|
||||
},
|
||||
"pi.ai.streaming": {
|
||||
type: "boolean",
|
||||
required: true,
|
||||
description: "Whether this operation returns a stream",
|
||||
},
|
||||
"pi.ai.deferred": {
|
||||
type: "boolean",
|
||||
required: false,
|
||||
description: "Whether the operation requests or participates in deferred execution",
|
||||
},
|
||||
},
|
||||
endAttributes: {
|
||||
"pi.ai.response.model": { type: "string", description: "Concrete response model" },
|
||||
"pi.ai.response.id": {
|
||||
type: "string",
|
||||
cardinality: "high",
|
||||
description: "Provider response id",
|
||||
},
|
||||
"pi.ai.response.stop_reason": {
|
||||
type: "string",
|
||||
values: ["stop", "length", "tool_use", "error", "aborted", "deferred"],
|
||||
description: "Normalized terminal response reason",
|
||||
},
|
||||
"pi.ai.http.status_code": { type: "number", description: "Final HTTP status" },
|
||||
"pi.ai.usage.input_tokens": { type: "number", description: "Reported input tokens" },
|
||||
"pi.ai.usage.output_tokens": { type: "number", description: "Reported output tokens" },
|
||||
"pi.ai.usage.cache_read_tokens": { type: "number", description: "Reported cache-read tokens" },
|
||||
"pi.ai.usage.cache_write_tokens": {
|
||||
type: "number",
|
||||
description: "Reported cache-write tokens",
|
||||
},
|
||||
"pi.ai.usage.reasoning_tokens": { type: "number", description: "Reported reasoning tokens" },
|
||||
"pi.ai.usage.total_tokens": { type: "number", description: "Reported total tokens" },
|
||||
"pi.ai.usage.cost": { type: "number", description: "Reported total cost" },
|
||||
"pi.ai.stream.chunk_count": { type: "number", description: "Streamed update chunk count" },
|
||||
"pi.ai.stream.time_to_first_chunk_ms": {
|
||||
type: "number",
|
||||
description: "Elapsed milliseconds to first update chunk",
|
||||
},
|
||||
"pi.ai.error.type": {
|
||||
type: "string",
|
||||
cardinality: "low",
|
||||
description: "Provider or transport error class",
|
||||
},
|
||||
},
|
||||
status: { default: "ok", errorWhen: "The operation throws or returns an error result" },
|
||||
},
|
||||
},
|
||||
} as const satisfies TelemetrySchemaDefinition;
|
||||
|
||||
export type AiSpanName = TelemetrySchemaSpanName<typeof AI_TELEMETRY_SCHEMA>;
|
||||
export type AiSpanStartAttributes<Name extends AiSpanName> = TelemetrySchemaSpanStartAttributes<
|
||||
typeof AI_TELEMETRY_SCHEMA,
|
||||
Name
|
||||
>;
|
||||
export type AiSpanEndAttributes<Name extends AiSpanName> = TelemetrySchemaSpanEndAttributes<
|
||||
typeof AI_TELEMETRY_SCHEMA,
|
||||
Name
|
||||
>;
|
||||
export type AiSpanAttributes<Name extends AiSpanName> = AiSpanStartAttributes<Name> & AiSpanEndAttributes<Name>;
|
||||
export type AiSpanEventName<Name extends AiSpanName> = TelemetrySchemaSpanEventName<typeof AI_TELEMETRY_SCHEMA, Name>;
|
||||
export type AiSpanEventAttributes<
|
||||
Name extends AiSpanName,
|
||||
EventName extends AiSpanEventName<Name>,
|
||||
> = TelemetrySchemaSpanEventAttributes<typeof AI_TELEMETRY_SCHEMA, Name, EventName>;
|
||||
export type AiTelemetrySpan<Name extends AiSpanName> = SchemaTelemetrySpan<typeof AI_TELEMETRY_SCHEMA, Name>;
|
||||
export type AiSpan = TelemetrySchemaSpanUnion<typeof AI_TELEMETRY_SCHEMA>;
|
||||
|
||||
export function startAiSpan<Name extends AiSpanName, const Attributes extends AiSpanStartAttributes<Name>, Result>(
|
||||
name: Name,
|
||||
attributes: ExactTelemetryAttributes<AiSpanStartAttributes<Name>, Attributes>,
|
||||
callback: (span: AiTelemetrySpan<Name>, context: Context) => Result | Promise<Result>,
|
||||
context: Context,
|
||||
): Promise<Result> {
|
||||
return getTelemetryContext(context).startSpan({ name, attributes }, (span) =>
|
||||
callback(span as AiTelemetrySpan<Name>, withTelemetryContext(span, context)),
|
||||
);
|
||||
}
|
||||
|
||||
const HOOK_NAMES = [
|
||||
"before_run",
|
||||
"before_drive",
|
||||
"before_run_end",
|
||||
"transform_context",
|
||||
"before_request",
|
||||
"before_payload",
|
||||
"after_response",
|
||||
"before_tool",
|
||||
"after_tool",
|
||||
"before_compaction",
|
||||
"before_navigation",
|
||||
] as const;
|
||||
|
||||
const EVENT_TYPES = [
|
||||
"run_start",
|
||||
"run_resume",
|
||||
"run_suspend",
|
||||
"operation_abort",
|
||||
"run_end",
|
||||
"fault",
|
||||
"handler_error",
|
||||
"turn_start",
|
||||
"turn_end",
|
||||
"retry_scheduled",
|
||||
"retry_start",
|
||||
"retry_end",
|
||||
"message_start",
|
||||
"message_update",
|
||||
"message_end",
|
||||
"tool_start",
|
||||
"tool_update",
|
||||
"tool_end",
|
||||
"entry_added",
|
||||
"queue_update",
|
||||
"value_update",
|
||||
"config_update",
|
||||
"compaction_start",
|
||||
"compaction_end",
|
||||
"navigation_start",
|
||||
"navigation_end",
|
||||
"lane_created",
|
||||
"usage",
|
||||
] as const;
|
||||
|
||||
const operationStartAttributes = {
|
||||
"pi.session.id": {
|
||||
type: "string",
|
||||
required: true,
|
||||
cardinality: "high",
|
||||
description: "Session id",
|
||||
},
|
||||
"pi.lane.name": {
|
||||
type: "string",
|
||||
required: true,
|
||||
cardinality: "high",
|
||||
description: "Lane name",
|
||||
},
|
||||
"pi.operation.id": {
|
||||
type: "string",
|
||||
required: true,
|
||||
cardinality: "high",
|
||||
description: "Durable operation id",
|
||||
},
|
||||
"pi.operation.recovery": {
|
||||
type: "boolean",
|
||||
required: true,
|
||||
description: "Whether this invocation resumes durable work",
|
||||
},
|
||||
} as const;
|
||||
|
||||
const operationErrorAttributes = {
|
||||
"pi.error.code": {
|
||||
type: "string",
|
||||
cardinality: "low",
|
||||
description: "Stable operation error code",
|
||||
},
|
||||
"pi.error.type": {
|
||||
type: "string",
|
||||
cardinality: "low",
|
||||
description: "Low-cardinality operation error class",
|
||||
},
|
||||
} as const;
|
||||
|
||||
export const HARNESS_TELEMETRY_SCHEMA = {
|
||||
version: 1,
|
||||
spans: {
|
||||
"pi.harness.run": {
|
||||
description: "One admitted in-process run invocation",
|
||||
parents: { kind: "root_or_external" },
|
||||
startAttributes: {
|
||||
...operationStartAttributes,
|
||||
"pi.operation.kind": {
|
||||
type: "string",
|
||||
required: true,
|
||||
values: ["run"],
|
||||
description: "Run operation kind",
|
||||
},
|
||||
},
|
||||
endAttributes: {
|
||||
"pi.operation.outcome": {
|
||||
type: "string",
|
||||
values: ["completed", "aborted", "failed", "suspended"],
|
||||
description: "Run invocation outcome",
|
||||
},
|
||||
...operationErrorAttributes,
|
||||
},
|
||||
status: { default: "ok", errorWhen: "The run fails or throws" },
|
||||
},
|
||||
"pi.harness.compaction": {
|
||||
description: "One admitted in-process manual compaction invocation",
|
||||
parents: { kind: "root_or_external" },
|
||||
startAttributes: {
|
||||
...operationStartAttributes,
|
||||
"pi.operation.kind": {
|
||||
type: "string",
|
||||
required: true,
|
||||
values: ["compaction"],
|
||||
description: "Compaction operation kind",
|
||||
},
|
||||
},
|
||||
endAttributes: {
|
||||
"pi.operation.outcome": {
|
||||
type: "string",
|
||||
values: ["completed", "declined", "aborted", "failed"],
|
||||
description: "Compaction invocation outcome",
|
||||
},
|
||||
...operationErrorAttributes,
|
||||
},
|
||||
status: { default: "ok", errorWhen: "The compaction fails or throws" },
|
||||
},
|
||||
"pi.harness.navigation": {
|
||||
description: "One admitted in-process navigation invocation",
|
||||
parents: { kind: "root_or_external" },
|
||||
startAttributes: {
|
||||
...operationStartAttributes,
|
||||
"pi.operation.kind": {
|
||||
type: "string",
|
||||
required: true,
|
||||
values: ["navigation"],
|
||||
description: "Navigation operation kind",
|
||||
},
|
||||
},
|
||||
endAttributes: {
|
||||
"pi.operation.outcome": {
|
||||
type: "string",
|
||||
values: ["completed", "declined", "aborted", "failed"],
|
||||
description: "Navigation invocation outcome",
|
||||
},
|
||||
...operationErrorAttributes,
|
||||
},
|
||||
status: { default: "ok", errorWhen: "The navigation fails or throws" },
|
||||
},
|
||||
"pi.harness.checkpoint": {
|
||||
description: "One run checkpoint",
|
||||
parents: { kind: "spans", spans: ["pi.harness.run"] },
|
||||
startAttributes: {
|
||||
"pi.lane.name": {
|
||||
type: "string",
|
||||
required: true,
|
||||
cardinality: "high",
|
||||
description: "Lane name",
|
||||
},
|
||||
"pi.operation.id": {
|
||||
type: "string",
|
||||
required: true,
|
||||
cardinality: "high",
|
||||
description: "Durable operation id",
|
||||
},
|
||||
"pi.checkpoint.kind": {
|
||||
type: "string",
|
||||
required: true,
|
||||
values: ["normal", "abort_reconcile"],
|
||||
description: "Checkpoint purpose",
|
||||
},
|
||||
},
|
||||
endAttributes: {},
|
||||
status: { default: "ok", errorWhen: "Checkpoint work throws" },
|
||||
},
|
||||
"pi.harness.turn": {
|
||||
description: "One assistant response and its tool batch",
|
||||
parents: { kind: "spans", spans: ["pi.harness.run"] },
|
||||
startAttributes: {
|
||||
"pi.lane.name": {
|
||||
type: "string",
|
||||
required: true,
|
||||
cardinality: "high",
|
||||
description: "Lane name",
|
||||
},
|
||||
"pi.operation.id": {
|
||||
type: "string",
|
||||
required: true,
|
||||
cardinality: "high",
|
||||
description: "Durable operation id",
|
||||
},
|
||||
"pi.turn.id": {
|
||||
type: "string",
|
||||
required: true,
|
||||
cardinality: "high",
|
||||
description: "Invocation-local turn id",
|
||||
},
|
||||
},
|
||||
endAttributes: {},
|
||||
status: { default: "ok", errorWhen: "Turn work throws" },
|
||||
},
|
||||
"pi.harness.step": {
|
||||
description: "One durable retry attempt",
|
||||
parents: {
|
||||
kind: "spans",
|
||||
spans: ["pi.harness.turn", "pi.harness.checkpoint", "pi.harness.compaction", "pi.harness.navigation"],
|
||||
},
|
||||
startAttributes: {
|
||||
"pi.lane.name": {
|
||||
type: "string",
|
||||
required: true,
|
||||
cardinality: "high",
|
||||
description: "Lane name",
|
||||
},
|
||||
"pi.operation.id": {
|
||||
type: "string",
|
||||
required: true,
|
||||
cardinality: "high",
|
||||
description: "Durable operation id",
|
||||
},
|
||||
"pi.step.kind": {
|
||||
type: "string",
|
||||
required: true,
|
||||
values: ["assistant", "compaction", "branch_summary"],
|
||||
description: "Retryable step kind",
|
||||
},
|
||||
"pi.step.attempt": {
|
||||
type: "number",
|
||||
required: true,
|
||||
description: "One-based durable attempt number",
|
||||
},
|
||||
"pi.compaction.reason": {
|
||||
type: "string",
|
||||
required: false,
|
||||
values: ["manual", "threshold", "overflow"],
|
||||
description: "Compaction trigger",
|
||||
},
|
||||
},
|
||||
endAttributes: {
|
||||
"pi.step.outcome": {
|
||||
type: "string",
|
||||
values: ["succeeded", "retry", "failed", "aborted", "deferred", "overflow"],
|
||||
description: "Attempt outcome",
|
||||
},
|
||||
},
|
||||
status: { default: "ok", errorWhen: "The attempt retries, fails, or throws" },
|
||||
},
|
||||
"pi.harness.tool": {
|
||||
description: "One raw phase-2 tool execution",
|
||||
parents: { kind: "spans", spans: ["pi.harness.turn", "pi.harness.run"] },
|
||||
startAttributes: {
|
||||
"pi.lane.name": {
|
||||
type: "string",
|
||||
required: true,
|
||||
cardinality: "high",
|
||||
description: "Lane name",
|
||||
},
|
||||
"pi.operation.id": {
|
||||
type: "string",
|
||||
required: true,
|
||||
cardinality: "high",
|
||||
description: "Durable operation id",
|
||||
},
|
||||
"pi.turn.id": {
|
||||
type: "string",
|
||||
required: false,
|
||||
cardinality: "high",
|
||||
description: "Invocation-local live turn id",
|
||||
},
|
||||
"pi.tool.name": {
|
||||
type: "string",
|
||||
required: true,
|
||||
description: "Tool name",
|
||||
},
|
||||
"pi.tool.call_id": {
|
||||
type: "string",
|
||||
required: true,
|
||||
cardinality: "high",
|
||||
description: "Tool call id",
|
||||
},
|
||||
"pi.tool.replay": {
|
||||
type: "string",
|
||||
required: true,
|
||||
values: ["never", "safe"],
|
||||
description: "Declared replay policy",
|
||||
},
|
||||
"pi.tool.recovery": {
|
||||
type: "boolean",
|
||||
required: true,
|
||||
description: "Whether this is recovery execution",
|
||||
},
|
||||
},
|
||||
endAttributes: {
|
||||
"pi.tool.is_error": {
|
||||
type: "boolean",
|
||||
description: "Whether raw phase-2 execution returned an error",
|
||||
},
|
||||
},
|
||||
status: { default: "ok", errorWhen: "Raw phase-2 execution returns an error" },
|
||||
},
|
||||
"pi.harness.hook": {
|
||||
description: "One registered hook handler invocation",
|
||||
parents: { kind: "any" },
|
||||
startAttributes: {
|
||||
"pi.lane.name": {
|
||||
type: "string",
|
||||
required: true,
|
||||
cardinality: "high",
|
||||
description: "Lane name",
|
||||
},
|
||||
"pi.operation.id": {
|
||||
type: "string",
|
||||
required: false,
|
||||
cardinality: "high",
|
||||
description: "Durable operation id when accepted",
|
||||
},
|
||||
"pi.hook.name": {
|
||||
type: "string",
|
||||
required: true,
|
||||
values: HOOK_NAMES,
|
||||
description: "Hook name",
|
||||
},
|
||||
"pi.hook.registration_id": {
|
||||
type: "string",
|
||||
required: false,
|
||||
description: "Optional hook registration metadata",
|
||||
},
|
||||
},
|
||||
endAttributes: {
|
||||
"pi.hook.outcome": {
|
||||
type: "string",
|
||||
values: ["completed", "skipped", "blocked", "failed"],
|
||||
description: "Handler outcome",
|
||||
},
|
||||
},
|
||||
status: { default: "ok", errorWhen: "The handler throws" },
|
||||
},
|
||||
"pi.harness.sleep": {
|
||||
description: "One retry delay",
|
||||
parents: {
|
||||
kind: "spans",
|
||||
spans: [
|
||||
"pi.harness.run",
|
||||
"pi.harness.compaction",
|
||||
"pi.harness.navigation",
|
||||
"pi.harness.turn",
|
||||
"pi.harness.checkpoint",
|
||||
],
|
||||
},
|
||||
startAttributes: {
|
||||
"pi.operation.id": {
|
||||
type: "string",
|
||||
required: true,
|
||||
cardinality: "high",
|
||||
description: "Durable operation id",
|
||||
},
|
||||
"pi.sleep.delay_ms": {
|
||||
type: "number",
|
||||
required: true,
|
||||
description: "Requested delay in milliseconds",
|
||||
},
|
||||
},
|
||||
endAttributes: {
|
||||
"pi.sleep.outcome": {
|
||||
type: "string",
|
||||
values: ["elapsed", "aborted"],
|
||||
description: "Delay outcome",
|
||||
},
|
||||
},
|
||||
status: { default: "ok", errorWhen: "Sleep work throws" },
|
||||
},
|
||||
"pi.harness.event_handler": {
|
||||
description: "One passive event listener invocation",
|
||||
parents: { kind: "any" },
|
||||
startAttributes: {
|
||||
"pi.event.type": {
|
||||
type: "string",
|
||||
required: true,
|
||||
cardinality: "low",
|
||||
values: EVENT_TYPES,
|
||||
description: "Delivered harness event type",
|
||||
},
|
||||
"pi.lane.name": {
|
||||
type: "string",
|
||||
required: false,
|
||||
cardinality: "high",
|
||||
description: "Lane name for lane-scoped events",
|
||||
},
|
||||
},
|
||||
endAttributes: {},
|
||||
status: { default: "ok", errorWhen: "The listener throws" },
|
||||
},
|
||||
"pi.session.write": {
|
||||
description: "One committed session transaction",
|
||||
parents: { kind: "any" },
|
||||
startAttributes: {
|
||||
"pi.session.id": {
|
||||
type: "string",
|
||||
required: true,
|
||||
cardinality: "high",
|
||||
description: "Session id",
|
||||
},
|
||||
"pi.lane.name": {
|
||||
type: "string",
|
||||
required: false,
|
||||
cardinality: "high",
|
||||
description: "Lane name when supplied by the caller",
|
||||
},
|
||||
"pi.operation.id": {
|
||||
type: "string",
|
||||
required: false,
|
||||
cardinality: "high",
|
||||
description: "Durable operation id when supplied by the caller",
|
||||
},
|
||||
"pi.session.item_count": {
|
||||
type: "number",
|
||||
required: true,
|
||||
description: "Number of writes in the transaction",
|
||||
},
|
||||
"pi.session.item_kinds": {
|
||||
type: "string[]",
|
||||
required: true,
|
||||
elementValues: ["entry", "usage", "value", "list"],
|
||||
description: "Distinct write kinds in the transaction",
|
||||
},
|
||||
},
|
||||
endAttributes: {
|
||||
"pi.session.first_seq": {
|
||||
type: "number",
|
||||
description: "First committed sequence in the transaction",
|
||||
},
|
||||
"pi.session.last_seq": {
|
||||
type: "number",
|
||||
description: "Last committed sequence in the transaction",
|
||||
},
|
||||
},
|
||||
status: { default: "ok", errorWhen: "Storage rejects the transaction" },
|
||||
},
|
||||
},
|
||||
} as const satisfies TelemetrySchemaDefinition;
|
||||
|
||||
/** Combined typed span vocabulary for agent-owned AI-request and harness telemetry. */
|
||||
export const AGENT_TELEMETRY_SCHEMAS = [AI_TELEMETRY_SCHEMA, HARNESS_TELEMETRY_SCHEMA] as const;
|
||||
|
||||
export type HarnessSpanName = TelemetrySchemaSpanName<typeof HARNESS_TELEMETRY_SCHEMA>;
|
||||
export type HarnessSpanStartAttributes<Name extends HarnessSpanName> = TelemetrySchemaSpanStartAttributes<
|
||||
typeof HARNESS_TELEMETRY_SCHEMA,
|
||||
Name
|
||||
>;
|
||||
export type HarnessSpanEndAttributes<Name extends HarnessSpanName> = TelemetrySchemaSpanEndAttributes<
|
||||
typeof HARNESS_TELEMETRY_SCHEMA,
|
||||
Name
|
||||
>;
|
||||
export type HarnessSpanAttributes<Name extends HarnessSpanName> = HarnessSpanStartAttributes<Name> &
|
||||
HarnessSpanEndAttributes<Name>;
|
||||
export type HarnessSpanEventName<Name extends HarnessSpanName> = TelemetrySchemaSpanEventName<
|
||||
typeof HARNESS_TELEMETRY_SCHEMA,
|
||||
Name
|
||||
>;
|
||||
export type HarnessSpanEventAttributes<
|
||||
Name extends HarnessSpanName,
|
||||
EventName extends HarnessSpanEventName<Name>,
|
||||
> = TelemetrySchemaSpanEventAttributes<typeof HARNESS_TELEMETRY_SCHEMA, Name, EventName>;
|
||||
export type HarnessTelemetrySpan<Name extends HarnessSpanName> = SchemaTelemetrySpan<
|
||||
typeof HARNESS_TELEMETRY_SCHEMA,
|
||||
Name
|
||||
>;
|
||||
export type HarnessSpan = TelemetrySchemaSpanUnion<typeof HARNESS_TELEMETRY_SCHEMA>;
|
||||
|
||||
export function startHarnessSpan<
|
||||
Name extends HarnessSpanName,
|
||||
const Attributes extends HarnessSpanStartAttributes<Name>,
|
||||
Result,
|
||||
>(
|
||||
name: Name,
|
||||
attributes: ExactTelemetryAttributes<HarnessSpanStartAttributes<Name>, Attributes>,
|
||||
callback: (span: HarnessTelemetrySpan<Name>, context: Context) => Result | Promise<Result>,
|
||||
context: Context,
|
||||
): Promise<Result> {
|
||||
return getTelemetryContext(context).startSpan({ name, attributes }, (span: TelemetrySpan) =>
|
||||
callback(span as HarnessTelemetrySpan<Name>, withTelemetryContext(span, context)),
|
||||
);
|
||||
}
|
||||
145
third_party/pi/packages/agent/src/harness/tools/bash.ts
vendored
Normal file
145
third_party/pi/packages/agent/src/harness/tools/bash.ts
vendored
Normal file
@@ -0,0 +1,145 @@
|
||||
import { type Static, Type } from "typebox";
|
||||
import type { Context } from "../context.ts";
|
||||
import type { AgentHarnessTool, ShellOutputTruncation, ShellOutputView } from "../types.ts";
|
||||
import { applyShellOutputUpdate } from "../utils/output-capture.ts";
|
||||
import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, formatSize } from "../utils/truncate.ts";
|
||||
import type { ExecutionToolContext } from "./tool-context.ts";
|
||||
|
||||
const MAX_TIMEOUT_SECONDS = 2_147_483_647 / 1000;
|
||||
const BASH_CHECKPOINT_INTERVAL_MS = 2_000;
|
||||
|
||||
const bashSchema = Type.Object({
|
||||
command: Type.String({ description: "Bash command to execute" }),
|
||||
timeout: Type.Optional(Type.Number({ description: "Timeout in seconds (optional, no default timeout)" })),
|
||||
});
|
||||
|
||||
export type BashToolInput = Static<typeof bashSchema>;
|
||||
|
||||
export interface BashToolDetails {
|
||||
truncation?: ShellOutputTruncation;
|
||||
fullOutputPath?: string;
|
||||
}
|
||||
|
||||
export interface BashExecution {
|
||||
command: string;
|
||||
cwd: string;
|
||||
env: Record<string, string>;
|
||||
inheritEnv: boolean;
|
||||
}
|
||||
|
||||
export type BashPrepare<TContext extends ExecutionToolContext = ExecutionToolContext> = (
|
||||
execution: BashExecution,
|
||||
toolContext: TContext,
|
||||
context: Context,
|
||||
) => void | Promise<void>;
|
||||
|
||||
export interface BashToolOptions<TContext extends ExecutionToolContext = ExecutionToolContext> {
|
||||
commandPrefix?: string;
|
||||
prepare?: BashPrepare<TContext>;
|
||||
}
|
||||
|
||||
function validateTimeout(timeout: number | undefined): void {
|
||||
if (timeout === undefined) return;
|
||||
if (!Number.isFinite(timeout) || timeout <= 0) {
|
||||
throw new Error("Invalid timeout: must be a finite number of seconds");
|
||||
}
|
||||
if (timeout > MAX_TIMEOUT_SECONDS) {
|
||||
throw new Error(`Invalid timeout: maximum is ${MAX_TIMEOUT_SECONDS} seconds`);
|
||||
}
|
||||
}
|
||||
|
||||
export function createBashTool<TContext extends ExecutionToolContext = ExecutionToolContext>(
|
||||
options?: BashToolOptions<TContext>,
|
||||
): AgentHarnessTool<TContext, typeof bashSchema, BashToolDetails | undefined> {
|
||||
return {
|
||||
name: "bash",
|
||||
label: "bash",
|
||||
description: `Execute a bash command in the current working directory. Returns combined stdout and stderr. Output is truncated to last ${DEFAULT_MAX_LINES} lines or ${DEFAULT_MAX_BYTES / 1024}KB (whichever is hit first). If truncated, full output is saved to a temp file. Optionally provide a timeout in seconds.`,
|
||||
parameters: bashSchema,
|
||||
async execute(_toolCallId, { command, timeout }, onUpdate, toolContext, _invocation, context) {
|
||||
validateTimeout(timeout);
|
||||
const { env } = toolContext;
|
||||
const execution: BashExecution = {
|
||||
command: options?.commandPrefix ? `${options.commandPrefix}\n${command}` : command,
|
||||
cwd: env.cwd,
|
||||
env: {},
|
||||
inheritEnv: true,
|
||||
};
|
||||
await options?.prepare?.(execution, toolContext, context);
|
||||
let view: ShellOutputView | undefined;
|
||||
let lastCheckpointAt = Date.now();
|
||||
let lastCheckpoint: string | undefined;
|
||||
let acceptingUpdates = true;
|
||||
|
||||
onUpdate({ content: [], details: undefined });
|
||||
const result = await env.exec(
|
||||
execution.command,
|
||||
{
|
||||
cwd: execution.cwd,
|
||||
env: execution.env,
|
||||
inheritEnv: execution.inheritEnv,
|
||||
timeout,
|
||||
capture: {
|
||||
limits: { maxBytes: DEFAULT_MAX_BYTES, maxLines: DEFAULT_MAX_LINES, retain: "tail" },
|
||||
spill: true,
|
||||
},
|
||||
onUpdate: (update) => {
|
||||
if (!acceptingUpdates) return;
|
||||
view = applyShellOutputUpdate(view, update);
|
||||
const snapshot = {
|
||||
content: [{ type: "text" as const, text: view.text }],
|
||||
details: {
|
||||
truncation: view.truncation.truncated ? view.truncation : undefined,
|
||||
fullOutputPath: view.spillPath,
|
||||
},
|
||||
};
|
||||
const now = Date.now();
|
||||
const encoded = JSON.stringify(snapshot);
|
||||
const checkpoint =
|
||||
now - lastCheckpointAt >= BASH_CHECKPOINT_INTERVAL_MS && encoded !== lastCheckpoint;
|
||||
onUpdate(snapshot, checkpoint ? { checkpoint: true } : undefined);
|
||||
if (checkpoint) {
|
||||
lastCheckpointAt = now;
|
||||
lastCheckpoint = encoded;
|
||||
}
|
||||
},
|
||||
},
|
||||
context,
|
||||
);
|
||||
acceptingUpdates = false;
|
||||
|
||||
let outputText = view?.text ?? "";
|
||||
const capture = result.ok ? { text: outputText, ...result.value } : view;
|
||||
let details: BashToolDetails | undefined;
|
||||
if (capture?.truncation.truncated) {
|
||||
details = { truncation: capture.truncation, fullOutputPath: capture.spillPath };
|
||||
const startLine = capture.truncation.totalLines - capture.truncation.outputLines + 1;
|
||||
const endLine = capture.truncation.totalLines;
|
||||
if (capture.truncation.lastLinePartial) {
|
||||
const lastLineSize = formatSize(capture.lastLineBytes ?? capture.truncation.outputBytes);
|
||||
outputText += `\n\n[Showing last ${formatSize(capture.truncation.outputBytes)} of line ${endLine} (line is ${lastLineSize}). Full output: ${capture.spillPath}]`;
|
||||
} else if (capture.truncation.truncatedBy === "lines") {
|
||||
outputText += `\n\n[Showing lines ${startLine}-${endLine} of ${capture.truncation.totalLines}. Full output: ${capture.spillPath}]`;
|
||||
} else {
|
||||
outputText += `\n\n[Showing lines ${startLine}-${endLine} of ${capture.truncation.totalLines} (${formatSize(DEFAULT_MAX_BYTES)} limit). Full output: ${capture.spillPath}]`;
|
||||
}
|
||||
}
|
||||
|
||||
if (!result.ok) {
|
||||
const status =
|
||||
result.error.code === "timeout"
|
||||
? `Command timed out after ${timeout} seconds`
|
||||
: result.error.code === "aborted"
|
||||
? "Command aborted"
|
||||
: result.error.message;
|
||||
throw new Error(outputText ? `${outputText}\n\n${status}` : status, { cause: result.error });
|
||||
}
|
||||
if (result.value.exitCode !== 0) {
|
||||
throw new Error(
|
||||
`${outputText ? `${outputText}\n\n` : ""}Command exited with code ${result.value.exitCode}`,
|
||||
);
|
||||
}
|
||||
return { content: [{ type: "text", text: outputText || "(no output)" }], details };
|
||||
},
|
||||
};
|
||||
}
|
||||
500
third_party/pi/packages/agent/src/harness/tools/edit-diff.ts
vendored
Normal file
500
third_party/pi/packages/agent/src/harness/tools/edit-diff.ts
vendored
Normal file
@@ -0,0 +1,500 @@
|
||||
/**
|
||||
* Shared diff computation utilities for the edit and similar tools.
|
||||
*/
|
||||
|
||||
import * as Diff from "diff";
|
||||
|
||||
export function detectLineEnding(content: string): "\r\n" | "\n" {
|
||||
const crlfIdx = content.indexOf("\r\n");
|
||||
const lfIdx = content.indexOf("\n");
|
||||
if (lfIdx === -1) return "\n";
|
||||
if (crlfIdx === -1) return "\n";
|
||||
return crlfIdx < lfIdx ? "\r\n" : "\n";
|
||||
}
|
||||
|
||||
export function normalizeToLF(text: string): string {
|
||||
return text.replace(/\r\n/g, "\n").replace(/\r/g, "\n");
|
||||
}
|
||||
|
||||
export function restoreLineEndings(text: string, ending: "\r\n" | "\n"): string {
|
||||
return ending === "\r\n" ? text.replace(/\n/g, "\r\n") : text;
|
||||
}
|
||||
|
||||
/**
|
||||
* Normalize text for fuzzy matching. Applies progressive transformations:
|
||||
* - Strip trailing whitespace from each line
|
||||
* - Normalize smart quotes to ASCII equivalents
|
||||
* - Normalize Unicode dashes/hyphens to ASCII hyphen
|
||||
* - Normalize special Unicode spaces to regular space
|
||||
*/
|
||||
export function normalizeForFuzzyMatch(text: string): string {
|
||||
return (
|
||||
text
|
||||
.normalize("NFKC")
|
||||
// Strip trailing whitespace per line
|
||||
.split("\n")
|
||||
.map((line) => line.trimEnd())
|
||||
.join("\n")
|
||||
// Smart single quotes → '
|
||||
.replace(/[\u2018\u2019\u201A\u201B]/g, "'")
|
||||
// Smart double quotes → "
|
||||
.replace(/[\u201C\u201D\u201E\u201F]/g, '"')
|
||||
// Various dashes/hyphens → -
|
||||
// U+2010 hyphen, U+2011 non-breaking hyphen, U+2012 figure dash,
|
||||
// U+2013 en-dash, U+2014 em-dash, U+2015 horizontal bar, U+2212 minus
|
||||
.replace(/[\u2010\u2011\u2012\u2013\u2014\u2015\u2212]/g, "-")
|
||||
// Special spaces → regular space
|
||||
// U+00A0 NBSP, U+2002-U+200A various spaces, U+202F narrow NBSP,
|
||||
// U+205F medium math space, U+3000 ideographic space
|
||||
.replace(/[\u00A0\u2002-\u200A\u202F\u205F\u3000]/g, " ")
|
||||
);
|
||||
}
|
||||
|
||||
function splitLinesWithEndings(content: string): string[] {
|
||||
return content.match(/[^\n]*\n|[^\n]+/g) ?? [];
|
||||
}
|
||||
|
||||
interface LineSpan {
|
||||
start: number;
|
||||
end: number;
|
||||
}
|
||||
|
||||
interface MatchedEdit {
|
||||
editIndex: number;
|
||||
matchIndex: number;
|
||||
matchLength: number;
|
||||
newText: string;
|
||||
}
|
||||
|
||||
type TextReplacement = Pick<MatchedEdit, "matchIndex" | "matchLength" | "newText">;
|
||||
|
||||
function getLineSpans(content: string): LineSpan[] {
|
||||
let offset = 0;
|
||||
return splitLinesWithEndings(content).map((line) => {
|
||||
const span = { start: offset, end: offset + line.length };
|
||||
offset = span.end;
|
||||
return span;
|
||||
});
|
||||
}
|
||||
|
||||
function getReplacementLineRange(lines: LineSpan[], replacement: TextReplacement) {
|
||||
const replacementStart = replacement.matchIndex;
|
||||
const replacementEnd = replacement.matchIndex + replacement.matchLength;
|
||||
|
||||
let startLine = -1;
|
||||
for (let i = 0; i < lines.length; i++) {
|
||||
const line = lines[i];
|
||||
if (replacementStart >= line.start && replacementStart < line.end) {
|
||||
startLine = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (startLine === -1) {
|
||||
throw new Error("Replacement range is outside the base content.");
|
||||
}
|
||||
|
||||
let endLine = startLine;
|
||||
while (endLine < lines.length && lines[endLine].end < replacementEnd) {
|
||||
endLine++;
|
||||
}
|
||||
if (endLine >= lines.length) {
|
||||
throw new Error("Replacement range is outside the base content.");
|
||||
}
|
||||
|
||||
return { startLine, endLine: endLine + 1 };
|
||||
}
|
||||
|
||||
function applyReplacements(content: string, replacements: TextReplacement[], offset = 0): string {
|
||||
let result = content;
|
||||
for (let i = replacements.length - 1; i >= 0; i--) {
|
||||
const replacement = replacements[i];
|
||||
const matchIndex = replacement.matchIndex - offset;
|
||||
result =
|
||||
result.substring(0, matchIndex) + replacement.newText + result.substring(matchIndex + replacement.matchLength);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
* Apply replacements matched against `baseContent` to `originalContent` while
|
||||
* preserving unchanged line blocks from the original.
|
||||
*
|
||||
* This is useful when `baseContent` is a normalized view of the original. Each
|
||||
* replacement is widened to the lines it actually touches, those touched lines
|
||||
* are rewritten from the normalized base, and all other lines are copied back
|
||||
* from `originalContent`. The actual replacement ranges drive preservation so
|
||||
* duplicate normalized lines cannot be aligned to the wrong occurrence.
|
||||
*/
|
||||
export function applyReplacementsPreservingUnchangedLines(
|
||||
originalContent: string,
|
||||
baseContent: string,
|
||||
replacements: TextReplacement[],
|
||||
): string {
|
||||
const originalLines = splitLinesWithEndings(originalContent);
|
||||
const baseLines = getLineSpans(baseContent);
|
||||
if (originalLines.length !== baseLines.length) {
|
||||
throw new Error("Cannot preserve unchanged lines because the base content has a different line count.");
|
||||
}
|
||||
|
||||
const groups: Array<{ startLine: number; endLine: number; replacements: TextReplacement[] }> = [];
|
||||
const sortedReplacements = [...replacements].sort((a, b) => a.matchIndex - b.matchIndex);
|
||||
for (const replacement of sortedReplacements) {
|
||||
const range = getReplacementLineRange(baseLines, replacement);
|
||||
const current = groups[groups.length - 1];
|
||||
if (current && range.startLine < current.endLine) {
|
||||
current.endLine = Math.max(current.endLine, range.endLine);
|
||||
current.replacements.push(replacement);
|
||||
continue;
|
||||
}
|
||||
groups.push({ ...range, replacements: [replacement] });
|
||||
}
|
||||
|
||||
let originalLineIndex = 0;
|
||||
let result = "";
|
||||
for (const group of groups) {
|
||||
result += originalLines.slice(originalLineIndex, group.startLine).join("");
|
||||
|
||||
const groupStartOffset = baseLines[group.startLine].start;
|
||||
const groupEndOffset = baseLines[group.endLine - 1].end;
|
||||
result += applyReplacements(
|
||||
baseContent.slice(groupStartOffset, groupEndOffset),
|
||||
group.replacements,
|
||||
groupStartOffset,
|
||||
);
|
||||
originalLineIndex = group.endLine;
|
||||
}
|
||||
result += originalLines.slice(originalLineIndex).join("");
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
export interface FuzzyMatchResult {
|
||||
/** Whether a match was found */
|
||||
found: boolean;
|
||||
/** The index where the match starts (in the content that should be used for replacement) */
|
||||
index: number;
|
||||
/** Length of the matched text */
|
||||
matchLength: number;
|
||||
/** Whether fuzzy matching was used (false = exact match) */
|
||||
usedFuzzyMatch: boolean;
|
||||
/**
|
||||
* The content to use for replacement operations.
|
||||
* When exact match: original content. When fuzzy match: normalized content.
|
||||
*/
|
||||
contentForReplacement: string;
|
||||
}
|
||||
|
||||
export interface Edit {
|
||||
oldText: string;
|
||||
newText: string;
|
||||
}
|
||||
|
||||
export interface AppliedEditsResult {
|
||||
baseContent: string;
|
||||
newContent: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Find oldText in content, trying exact match first, then fuzzy match.
|
||||
* When fuzzy matching is used, the returned contentForReplacement is the
|
||||
* fuzzy-normalized version of the content (trailing whitespace stripped,
|
||||
* Unicode quotes/dashes normalized to ASCII).
|
||||
*/
|
||||
export function fuzzyFindText(content: string, oldText: string): FuzzyMatchResult {
|
||||
// Try exact match first
|
||||
const exactIndex = content.indexOf(oldText);
|
||||
if (exactIndex !== -1) {
|
||||
return {
|
||||
found: true,
|
||||
index: exactIndex,
|
||||
matchLength: oldText.length,
|
||||
usedFuzzyMatch: false,
|
||||
contentForReplacement: content,
|
||||
};
|
||||
}
|
||||
|
||||
// Try fuzzy match - work entirely in normalized space
|
||||
const fuzzyContent = normalizeForFuzzyMatch(content);
|
||||
const fuzzyOldText = normalizeForFuzzyMatch(oldText);
|
||||
const fuzzyIndex = fuzzyContent.indexOf(fuzzyOldText);
|
||||
|
||||
if (fuzzyIndex === -1) {
|
||||
return {
|
||||
found: false,
|
||||
index: -1,
|
||||
matchLength: 0,
|
||||
usedFuzzyMatch: false,
|
||||
contentForReplacement: content,
|
||||
};
|
||||
}
|
||||
|
||||
// When fuzzy matching, return offsets in normalized space. Callers can use
|
||||
// the normalized content to compute replacements, then decide how much of
|
||||
// that normalized output should be written back.
|
||||
return {
|
||||
found: true,
|
||||
index: fuzzyIndex,
|
||||
matchLength: fuzzyOldText.length,
|
||||
usedFuzzyMatch: true,
|
||||
contentForReplacement: fuzzyContent,
|
||||
};
|
||||
}
|
||||
|
||||
/** Strip UTF-8 BOM if present, return both the BOM (if any) and the text without it */
|
||||
export function stripBom(content: string): { bom: string; text: string } {
|
||||
return content.startsWith("\uFEFF") ? { bom: "\uFEFF", text: content.slice(1) } : { bom: "", text: content };
|
||||
}
|
||||
|
||||
function countOccurrences(content: string, oldText: string): number {
|
||||
const fuzzyContent = normalizeForFuzzyMatch(content);
|
||||
const fuzzyOldText = normalizeForFuzzyMatch(oldText);
|
||||
return fuzzyContent.split(fuzzyOldText).length - 1;
|
||||
}
|
||||
|
||||
function getNotFoundError(path: string, editIndex: number, totalEdits: number): Error {
|
||||
if (totalEdits === 1) {
|
||||
return new Error(
|
||||
`Could not find the exact text in ${path}. The old text must match exactly including all whitespace and newlines.`,
|
||||
);
|
||||
}
|
||||
return new Error(
|
||||
`Could not find edits[${editIndex}] in ${path}. The oldText must match exactly including all whitespace and newlines.`,
|
||||
);
|
||||
}
|
||||
|
||||
function getDuplicateError(path: string, editIndex: number, totalEdits: number, occurrences: number): Error {
|
||||
if (totalEdits === 1) {
|
||||
return new Error(
|
||||
`Found ${occurrences} occurrences of the text in ${path}. The text must be unique. Please provide more context to make it unique.`,
|
||||
);
|
||||
}
|
||||
return new Error(
|
||||
`Found ${occurrences} occurrences of edits[${editIndex}] in ${path}. Each oldText must be unique. Please provide more context to make it unique.`,
|
||||
);
|
||||
}
|
||||
|
||||
function getEmptyOldTextError(path: string, editIndex: number, totalEdits: number): Error {
|
||||
if (totalEdits === 1) {
|
||||
return new Error(`oldText must not be empty in ${path}.`);
|
||||
}
|
||||
return new Error(`edits[${editIndex}].oldText must not be empty in ${path}.`);
|
||||
}
|
||||
|
||||
function getNoChangeError(path: string, totalEdits: number): Error {
|
||||
if (totalEdits === 1) {
|
||||
return new Error(
|
||||
`No changes made to ${path}. The replacement produced identical content. This might indicate an issue with special characters or the text not existing as expected.`,
|
||||
);
|
||||
}
|
||||
return new Error(`No changes made to ${path}. The replacements produced identical content.`);
|
||||
}
|
||||
|
||||
/**
|
||||
* Apply one or more exact-text replacements to LF-normalized content.
|
||||
*
|
||||
* All edits are matched against the same original content. Replacements are
|
||||
* then applied in reverse order so offsets remain stable. If any edit needs
|
||||
* fuzzy matching, the operation runs in fuzzy-normalized content space and then
|
||||
* overlays those line-level changes onto the original content so unchanged line
|
||||
* blocks keep their original bytes.
|
||||
*/
|
||||
export function applyEditsToNormalizedContent(
|
||||
normalizedContent: string,
|
||||
edits: Edit[],
|
||||
path: string,
|
||||
): AppliedEditsResult {
|
||||
const normalizedEdits = edits.map((edit) => ({
|
||||
oldText: normalizeToLF(edit.oldText),
|
||||
newText: normalizeToLF(edit.newText),
|
||||
}));
|
||||
|
||||
for (let i = 0; i < normalizedEdits.length; i++) {
|
||||
if (normalizedEdits[i].oldText.length === 0) {
|
||||
throw getEmptyOldTextError(path, i, normalizedEdits.length);
|
||||
}
|
||||
}
|
||||
|
||||
const initialMatches = normalizedEdits.map((edit) => fuzzyFindText(normalizedContent, edit.oldText));
|
||||
const usedFuzzyMatch = initialMatches.some((match) => match.usedFuzzyMatch);
|
||||
const replacementBaseContent = usedFuzzyMatch ? normalizeForFuzzyMatch(normalizedContent) : normalizedContent;
|
||||
|
||||
const matchedEdits: MatchedEdit[] = [];
|
||||
for (let i = 0; i < normalizedEdits.length; i++) {
|
||||
const edit = normalizedEdits[i];
|
||||
const matchResult = fuzzyFindText(replacementBaseContent, edit.oldText);
|
||||
if (!matchResult.found) {
|
||||
throw getNotFoundError(path, i, normalizedEdits.length);
|
||||
}
|
||||
|
||||
const occurrences = countOccurrences(replacementBaseContent, edit.oldText);
|
||||
if (occurrences > 1) {
|
||||
throw getDuplicateError(path, i, normalizedEdits.length, occurrences);
|
||||
}
|
||||
|
||||
matchedEdits.push({
|
||||
editIndex: i,
|
||||
matchIndex: matchResult.index,
|
||||
matchLength: matchResult.matchLength,
|
||||
newText: edit.newText,
|
||||
});
|
||||
}
|
||||
|
||||
matchedEdits.sort((a, b) => a.matchIndex - b.matchIndex);
|
||||
for (let i = 1; i < matchedEdits.length; i++) {
|
||||
const previous = matchedEdits[i - 1];
|
||||
const current = matchedEdits[i];
|
||||
if (previous.matchIndex + previous.matchLength > current.matchIndex) {
|
||||
throw new Error(
|
||||
`edits[${previous.editIndex}] and edits[${current.editIndex}] overlap in ${path}. Merge them into one edit or target disjoint regions.`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
const baseContent = normalizedContent;
|
||||
const newContent = usedFuzzyMatch
|
||||
? applyReplacementsPreservingUnchangedLines(normalizedContent, replacementBaseContent, matchedEdits)
|
||||
: applyReplacements(replacementBaseContent, matchedEdits);
|
||||
|
||||
if (baseContent === newContent) {
|
||||
throw getNoChangeError(path, normalizedEdits.length);
|
||||
}
|
||||
|
||||
return { baseContent, newContent };
|
||||
}
|
||||
|
||||
/** Generate a standard unified patch. */
|
||||
export function generateUnifiedPatch(path: string, oldContent: string, newContent: string, contextLines = 4): string {
|
||||
return Diff.createTwoFilesPatch(path, path, oldContent, newContent, undefined, undefined, {
|
||||
context: contextLines,
|
||||
headerOptions: Diff.FILE_HEADERS_ONLY,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate a display-oriented diff string with line numbers and context.
|
||||
* Returns both the diff string and the first changed line number (in the new file).
|
||||
*/
|
||||
export function generateDiffString(
|
||||
oldContent: string,
|
||||
newContent: string,
|
||||
contextLines = 4,
|
||||
): { diff: string; firstChangedLine: number | undefined } {
|
||||
const parts = Diff.diffLines(oldContent, newContent);
|
||||
const output: string[] = [];
|
||||
|
||||
const oldLines = oldContent.split("\n");
|
||||
const newLines = newContent.split("\n");
|
||||
const maxLineNum = Math.max(oldLines.length, newLines.length);
|
||||
const lineNumWidth = String(maxLineNum).length;
|
||||
|
||||
let oldLineNum = 1;
|
||||
let newLineNum = 1;
|
||||
let lastWasChange = false;
|
||||
let firstChangedLine: number | undefined;
|
||||
|
||||
for (let i = 0; i < parts.length; i++) {
|
||||
const part = parts[i];
|
||||
const raw = part.value.split("\n");
|
||||
if (raw[raw.length - 1] === "") {
|
||||
raw.pop();
|
||||
}
|
||||
|
||||
if (part.added || part.removed) {
|
||||
// Capture the first changed line (in the new file)
|
||||
if (firstChangedLine === undefined) {
|
||||
firstChangedLine = newLineNum;
|
||||
}
|
||||
|
||||
// Show the change
|
||||
for (const line of raw) {
|
||||
if (part.added) {
|
||||
const lineNum = String(newLineNum).padStart(lineNumWidth, " ");
|
||||
output.push(`+${lineNum} ${line}`);
|
||||
newLineNum++;
|
||||
} else {
|
||||
// removed
|
||||
const lineNum = String(oldLineNum).padStart(lineNumWidth, " ");
|
||||
output.push(`-${lineNum} ${line}`);
|
||||
oldLineNum++;
|
||||
}
|
||||
}
|
||||
lastWasChange = true;
|
||||
} else {
|
||||
// Context lines - only show a few before/after changes
|
||||
const nextPartIsChange = i < parts.length - 1 && (parts[i + 1].added || parts[i + 1].removed);
|
||||
const hasLeadingChange = lastWasChange;
|
||||
const hasTrailingChange = nextPartIsChange;
|
||||
|
||||
if (hasLeadingChange && hasTrailingChange) {
|
||||
if (raw.length <= contextLines * 2) {
|
||||
for (const line of raw) {
|
||||
const lineNum = String(oldLineNum).padStart(lineNumWidth, " ");
|
||||
output.push(` ${lineNum} ${line}`);
|
||||
oldLineNum++;
|
||||
newLineNum++;
|
||||
}
|
||||
} else {
|
||||
const leadingLines = raw.slice(0, contextLines);
|
||||
const trailingLines = raw.slice(raw.length - contextLines);
|
||||
const skippedLines = raw.length - leadingLines.length - trailingLines.length;
|
||||
|
||||
for (const line of leadingLines) {
|
||||
const lineNum = String(oldLineNum).padStart(lineNumWidth, " ");
|
||||
output.push(` ${lineNum} ${line}`);
|
||||
oldLineNum++;
|
||||
newLineNum++;
|
||||
}
|
||||
|
||||
output.push(` ${"".padStart(lineNumWidth, " ")} ...`);
|
||||
oldLineNum += skippedLines;
|
||||
newLineNum += skippedLines;
|
||||
|
||||
for (const line of trailingLines) {
|
||||
const lineNum = String(oldLineNum).padStart(lineNumWidth, " ");
|
||||
output.push(` ${lineNum} ${line}`);
|
||||
oldLineNum++;
|
||||
newLineNum++;
|
||||
}
|
||||
}
|
||||
} else if (hasLeadingChange) {
|
||||
const shownLines = raw.slice(0, contextLines);
|
||||
const skippedLines = raw.length - shownLines.length;
|
||||
|
||||
for (const line of shownLines) {
|
||||
const lineNum = String(oldLineNum).padStart(lineNumWidth, " ");
|
||||
output.push(` ${lineNum} ${line}`);
|
||||
oldLineNum++;
|
||||
newLineNum++;
|
||||
}
|
||||
|
||||
if (skippedLines > 0) {
|
||||
output.push(` ${"".padStart(lineNumWidth, " ")} ...`);
|
||||
oldLineNum += skippedLines;
|
||||
newLineNum += skippedLines;
|
||||
}
|
||||
} else if (hasTrailingChange) {
|
||||
const skippedLines = Math.max(0, raw.length - contextLines);
|
||||
if (skippedLines > 0) {
|
||||
output.push(` ${"".padStart(lineNumWidth, " ")} ...`);
|
||||
oldLineNum += skippedLines;
|
||||
newLineNum += skippedLines;
|
||||
}
|
||||
|
||||
for (const line of raw.slice(skippedLines)) {
|
||||
const lineNum = String(oldLineNum).padStart(lineNumWidth, " ");
|
||||
output.push(` ${lineNum} ${line}`);
|
||||
oldLineNum++;
|
||||
newLineNum++;
|
||||
}
|
||||
} else {
|
||||
// Skip these context lines entirely
|
||||
oldLineNum += raw.length;
|
||||
newLineNum += raw.length;
|
||||
}
|
||||
|
||||
lastWasChange = false;
|
||||
}
|
||||
}
|
||||
|
||||
return { diff: output.join("\n"), firstChangedLine };
|
||||
}
|
||||
145
third_party/pi/packages/agent/src/harness/tools/edit.ts
vendored
Normal file
145
third_party/pi/packages/agent/src/harness/tools/edit.ts
vendored
Normal file
@@ -0,0 +1,145 @@
|
||||
import { type Static, Type } from "typebox";
|
||||
import type { AgentHarnessTool, FileError } from "../types.ts";
|
||||
import {
|
||||
applyEditsToNormalizedContent,
|
||||
detectLineEnding,
|
||||
type Edit,
|
||||
generateDiffString,
|
||||
generateUnifiedPatch,
|
||||
normalizeToLF,
|
||||
restoreLineEndings,
|
||||
stripBom,
|
||||
} from "./edit-diff.ts";
|
||||
import { withFileMutationQueue } from "./file-mutation-queue.ts";
|
||||
import { resolveToolPath } from "./path-utils.ts";
|
||||
import type { ExecutionToolContext } from "./tool-context.ts";
|
||||
|
||||
const replaceEditSchema = Type.Object(
|
||||
{
|
||||
oldText: Type.String({
|
||||
description:
|
||||
"Exact text for one targeted replacement. It must be unique in the original file and must not overlap with any other edits[].oldText in the same call.",
|
||||
}),
|
||||
newText: Type.String({ description: "Replacement text for this targeted edit." }),
|
||||
},
|
||||
{},
|
||||
);
|
||||
|
||||
const editSchema = Type.Object(
|
||||
{
|
||||
path: Type.String({ description: "Path to the file to edit (relative or absolute)" }),
|
||||
edits: Type.Array(replaceEditSchema, {
|
||||
description:
|
||||
"One or more targeted replacements. Each edit is matched against the original file, not incrementally. Do not include overlapping or nested edits. If two changes touch the same block or nearby lines, merge them into one edit instead.",
|
||||
}),
|
||||
},
|
||||
{},
|
||||
);
|
||||
|
||||
export type EditToolInput = Static<typeof editSchema>;
|
||||
type LegacyEditToolInput = EditToolInput & { oldText?: unknown; newText?: unknown };
|
||||
type SingleEditInput = { oldText: string; newText: string };
|
||||
|
||||
function isSingleEditInput(value: unknown): value is SingleEditInput {
|
||||
if (!value || typeof value !== "object" || Array.isArray(value)) return false;
|
||||
const edit = value as Record<string, unknown>;
|
||||
return typeof edit.oldText === "string" && typeof edit.newText === "string";
|
||||
}
|
||||
|
||||
export interface EditToolDetails {
|
||||
diff: string;
|
||||
patch: string;
|
||||
firstChangedLine?: number;
|
||||
}
|
||||
|
||||
function prepareEditArguments(input: unknown): EditToolInput {
|
||||
if (!input || typeof input !== "object") return input as EditToolInput;
|
||||
const args = input as Record<string, unknown>;
|
||||
if (typeof args.edits === "string") {
|
||||
try {
|
||||
const parsed: unknown = JSON.parse(args.edits);
|
||||
if (Array.isArray(parsed)) {
|
||||
args.edits = parsed;
|
||||
} else if (isSingleEditInput(parsed)) {
|
||||
args.edits = [parsed];
|
||||
}
|
||||
} catch {}
|
||||
} else if (isSingleEditInput(args.edits)) {
|
||||
args.edits = [args.edits];
|
||||
}
|
||||
|
||||
const legacy = args as LegacyEditToolInput;
|
||||
if (typeof legacy.oldText !== "string" || typeof legacy.newText !== "string") return args as EditToolInput;
|
||||
const edits = Array.isArray(legacy.edits) ? [...legacy.edits] : [];
|
||||
edits.push({ oldText: legacy.oldText, newText: legacy.newText });
|
||||
const { oldText: _oldText, newText: _newText, ...rest } = legacy;
|
||||
return { ...rest, edits } as EditToolInput;
|
||||
}
|
||||
|
||||
function validateEditInput(input: EditToolInput): { path: string; edits: Edit[] } {
|
||||
if (!Array.isArray(input.edits) || input.edits.length === 0) {
|
||||
throw new Error("Edit tool input is invalid. edits must contain at least one replacement.");
|
||||
}
|
||||
return { path: input.path, edits: input.edits };
|
||||
}
|
||||
|
||||
function editAccessError(path: string, error: FileError): Error {
|
||||
return new Error(`Could not edit file: ${path}. Error code: ${error.code}.`, { cause: error });
|
||||
}
|
||||
|
||||
export function createEditTool<TContext extends ExecutionToolContext = ExecutionToolContext>(): AgentHarnessTool<
|
||||
TContext,
|
||||
typeof editSchema,
|
||||
EditToolDetails | undefined
|
||||
> {
|
||||
return {
|
||||
name: "edit",
|
||||
label: "edit",
|
||||
description:
|
||||
"Edit a single file using exact text replacement. Every edits[].oldText must match a unique, non-overlapping region of the original file. If two changes affect the same block or nearby lines, merge them into one edit instead of emitting overlapping edits. Do not include large unchanged regions just to connect distant changes.",
|
||||
parameters: editSchema,
|
||||
prepareArguments: prepareEditArguments,
|
||||
async execute(_toolCallId, input, _onUpdate, { env }, _invocation, context) {
|
||||
const { path, edits } = validateEditInput(input);
|
||||
const absolutePath = await resolveToolPath(env, path, context);
|
||||
return withFileMutationQueue(
|
||||
env,
|
||||
absolutePath,
|
||||
async () => {
|
||||
if (context.abortSignal?.aborted) throw new Error("Operation aborted");
|
||||
const info = await env.fileInfo(absolutePath, context);
|
||||
if (!info.ok) throw editAccessError(path, info.error);
|
||||
if (info.value.kind !== "file" && info.value.kind !== "symlink") {
|
||||
throw new Error(`Could not edit file: ${path}. Path is not a file.`);
|
||||
}
|
||||
|
||||
const readResult = await env.readTextFile(absolutePath, context);
|
||||
if (!readResult.ok) throw editAccessError(path, readResult.error);
|
||||
if (context.abortSignal?.aborted) throw new Error("Operation aborted");
|
||||
|
||||
const { bom, text: content } = stripBom(readResult.value);
|
||||
const originalEnding = detectLineEnding(content);
|
||||
const normalizedContent = normalizeToLF(content);
|
||||
const { baseContent, newContent } = applyEditsToNormalizedContent(normalizedContent, edits, path);
|
||||
if (context.abortSignal?.aborted) throw new Error("Operation aborted");
|
||||
|
||||
const finalContent = bom + restoreLineEndings(newContent, originalEnding);
|
||||
const writeResult = await env.writeFile(absolutePath, finalContent, context);
|
||||
if (!writeResult.ok) throw editAccessError(path, writeResult.error);
|
||||
if (context.abortSignal?.aborted) throw new Error("Operation aborted");
|
||||
|
||||
const diffResult = generateDiffString(baseContent, newContent);
|
||||
return {
|
||||
content: [{ type: "text", text: `Successfully replaced ${edits.length} block(s) in ${path}.` }],
|
||||
details: {
|
||||
diff: diffResult.diff,
|
||||
patch: generateUnifiedPatch(path, baseContent, newContent),
|
||||
firstChangedLine: diffResult.firstChangedLine,
|
||||
},
|
||||
};
|
||||
},
|
||||
context,
|
||||
);
|
||||
},
|
||||
};
|
||||
}
|
||||
62
third_party/pi/packages/agent/src/harness/tools/file-mutation-queue.ts
vendored
Normal file
62
third_party/pi/packages/agent/src/harness/tools/file-mutation-queue.ts
vendored
Normal file
@@ -0,0 +1,62 @@
|
||||
import type { Context } from "../context.ts";
|
||||
import type { ExecutionEnv } from "../types.ts";
|
||||
import { getOrThrow } from "../types.ts";
|
||||
|
||||
type MutationQueueState = {
|
||||
queues: Map<string, Promise<void>>;
|
||||
registration: Promise<void>;
|
||||
};
|
||||
|
||||
const states = new WeakMap<ExecutionEnv, MutationQueueState>();
|
||||
|
||||
function getState(env: ExecutionEnv): MutationQueueState {
|
||||
let state = states.get(env);
|
||||
if (!state) {
|
||||
state = { queues: new Map(), registration: Promise.resolve() };
|
||||
states.set(env, state);
|
||||
}
|
||||
return state;
|
||||
}
|
||||
|
||||
async function getMutationQueueKey(env: ExecutionEnv, path: string, context: Context): Promise<string> {
|
||||
const absolutePath = getOrThrow(await env.absolutePath(path, context));
|
||||
const canonicalPath = await env.canonicalPath(absolutePath, context);
|
||||
if (canonicalPath.ok) return canonicalPath.value;
|
||||
if (canonicalPath.error.code === "not_found" || canonicalPath.error.code === "not_supported") return absolutePath;
|
||||
throw canonicalPath.error;
|
||||
}
|
||||
|
||||
/** Serialize file mutations targeting the same environment and canonical path. */
|
||||
export async function withFileMutationQueue<T>(
|
||||
env: ExecutionEnv,
|
||||
path: string,
|
||||
fn: () => Promise<T>,
|
||||
context: Context,
|
||||
): Promise<T> {
|
||||
const state = getState(env);
|
||||
const registration = state.registration.then(async () => {
|
||||
const key = await getMutationQueueKey(env, path, context);
|
||||
const currentQueue = state.queues.get(key) ?? Promise.resolve();
|
||||
|
||||
let releaseNext = () => {};
|
||||
const nextQueue = new Promise<void>((resolve) => {
|
||||
releaseNext = resolve;
|
||||
});
|
||||
const chainedQueue = currentQueue.then(() => nextQueue);
|
||||
state.queues.set(key, chainedQueue);
|
||||
return { key, currentQueue, chainedQueue, releaseNext };
|
||||
});
|
||||
state.registration = registration.then(
|
||||
() => undefined,
|
||||
() => undefined,
|
||||
);
|
||||
|
||||
const { key, currentQueue, chainedQueue, releaseNext } = await registration;
|
||||
await currentQueue;
|
||||
try {
|
||||
return await fn();
|
||||
} finally {
|
||||
releaseNext();
|
||||
if (state.queues.get(key) === chainedQueue) state.queues.delete(key);
|
||||
}
|
||||
}
|
||||
104
third_party/pi/packages/agent/src/harness/tools/image.ts
vendored
Normal file
104
third_party/pi/packages/agent/src/harness/tools/image.ts
vendored
Normal file
@@ -0,0 +1,104 @@
|
||||
const PNG_SIGNATURE = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a];
|
||||
|
||||
export function detectSupportedImageMimeType(buffer: Uint8Array): string | undefined {
|
||||
if (startsWith(buffer, [0xff, 0xd8, 0xff])) return buffer[3] === 0xf7 ? undefined : "image/jpeg";
|
||||
if (startsWith(buffer, PNG_SIGNATURE)) return isPng(buffer) && !isAnimatedPng(buffer) ? "image/png" : undefined;
|
||||
if (startsWithAscii(buffer, 0, "GIF")) return "image/gif";
|
||||
if (startsWithAscii(buffer, 0, "RIFF") && startsWithAscii(buffer, 8, "WEBP")) return "image/webp";
|
||||
if (startsWithAscii(buffer, 0, "BM") && isBmp(buffer)) return "image/bmp";
|
||||
return undefined;
|
||||
}
|
||||
|
||||
export function encodeBase64(bytes: Uint8Array): string {
|
||||
const alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
|
||||
let output = "";
|
||||
for (let index = 0; index < bytes.length; index += 3) {
|
||||
const first = bytes[index] ?? 0;
|
||||
const second = bytes[index + 1];
|
||||
const third = bytes[index + 2];
|
||||
output += alphabet[first >> 2];
|
||||
output += alphabet[((first & 0x03) << 4) | ((second ?? 0) >> 4)];
|
||||
output += second === undefined ? "=" : alphabet[((second & 0x0f) << 2) | ((third ?? 0) >> 6)];
|
||||
output += third === undefined ? "=" : alphabet[third & 0x3f];
|
||||
}
|
||||
return output;
|
||||
}
|
||||
|
||||
function isPng(buffer: Uint8Array): boolean {
|
||||
return (
|
||||
buffer.length >= 16 && readUint32BE(buffer, PNG_SIGNATURE.length) === 13 && startsWithAscii(buffer, 12, "IHDR")
|
||||
);
|
||||
}
|
||||
|
||||
function isAnimatedPng(buffer: Uint8Array): boolean {
|
||||
let offset = PNG_SIGNATURE.length;
|
||||
while (offset + 8 <= buffer.length) {
|
||||
const chunkLength = readUint32BE(buffer, offset);
|
||||
const chunkTypeOffset = offset + 4;
|
||||
if (startsWithAscii(buffer, chunkTypeOffset, "acTL")) return true;
|
||||
if (startsWithAscii(buffer, chunkTypeOffset, "IDAT")) return false;
|
||||
const nextOffset = offset + 8 + chunkLength + 4;
|
||||
if (nextOffset <= offset || nextOffset > buffer.length) return false;
|
||||
offset = nextOffset;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
function isBmp(buffer: Uint8Array): boolean {
|
||||
if (buffer.length < 26) return false;
|
||||
const declaredFileSize = readUint32LE(buffer, 2);
|
||||
const pixelDataOffset = readUint32LE(buffer, 10);
|
||||
const dibHeaderSize = readUint32LE(buffer, 14);
|
||||
if (declaredFileSize !== 0 && declaredFileSize < 26) return false;
|
||||
if (pixelDataOffset < 14 + dibHeaderSize) return false;
|
||||
if (declaredFileSize !== 0 && pixelDataOffset >= declaredFileSize) return false;
|
||||
|
||||
let colorPlanes: number;
|
||||
let bitsPerPixel: number;
|
||||
if (dibHeaderSize === 12) {
|
||||
colorPlanes = readUint16LE(buffer, 22);
|
||||
bitsPerPixel = readUint16LE(buffer, 24);
|
||||
} else if (dibHeaderSize >= 40 && dibHeaderSize <= 124) {
|
||||
if (buffer.length < 30) return false;
|
||||
colorPlanes = readUint16LE(buffer, 26);
|
||||
bitsPerPixel = readUint16LE(buffer, 28);
|
||||
} else {
|
||||
return false;
|
||||
}
|
||||
return colorPlanes === 1 && [1, 4, 8, 16, 24, 32].includes(bitsPerPixel);
|
||||
}
|
||||
|
||||
function readUint16LE(buffer: Uint8Array, offset: number): number {
|
||||
return (buffer[offset] ?? 0) + ((buffer[offset + 1] ?? 0) << 8);
|
||||
}
|
||||
|
||||
function readUint32BE(buffer: Uint8Array, offset: number): number {
|
||||
return (
|
||||
(buffer[offset] ?? 0) * 0x1000000 +
|
||||
((buffer[offset + 1] ?? 0) << 16) +
|
||||
((buffer[offset + 2] ?? 0) << 8) +
|
||||
(buffer[offset + 3] ?? 0)
|
||||
);
|
||||
}
|
||||
|
||||
function readUint32LE(buffer: Uint8Array, offset: number): number {
|
||||
return (
|
||||
(buffer[offset] ?? 0) +
|
||||
((buffer[offset + 1] ?? 0) << 8) +
|
||||
((buffer[offset + 2] ?? 0) << 16) +
|
||||
(buffer[offset + 3] ?? 0) * 0x1000000
|
||||
);
|
||||
}
|
||||
|
||||
function startsWith(buffer: Uint8Array, bytes: number[]): boolean {
|
||||
if (buffer.length < bytes.length) return false;
|
||||
return bytes.every((byte, index) => buffer[index] === byte);
|
||||
}
|
||||
|
||||
function startsWithAscii(buffer: Uint8Array, offset: number, text: string): boolean {
|
||||
if (buffer.length < offset + text.length) return false;
|
||||
for (let index = 0; index < text.length; index++) {
|
||||
if (buffer[offset + index] !== text.charCodeAt(index)) return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
23
third_party/pi/packages/agent/src/harness/tools/index.ts
vendored
Normal file
23
third_party/pi/packages/agent/src/harness/tools/index.ts
vendored
Normal file
@@ -0,0 +1,23 @@
|
||||
export {
|
||||
type BashExecution,
|
||||
type BashPrepare,
|
||||
type BashToolDetails,
|
||||
type BashToolInput,
|
||||
type BashToolOptions,
|
||||
createBashTool,
|
||||
} from "./bash.ts";
|
||||
export {
|
||||
createEditTool,
|
||||
type EditToolDetails,
|
||||
type EditToolInput,
|
||||
} from "./edit.ts";
|
||||
export {
|
||||
createReadTool,
|
||||
type ReadImageProcessor,
|
||||
type ReadImageProcessorResult,
|
||||
type ReadToolDetails,
|
||||
type ReadToolInput,
|
||||
type ReadToolOptions,
|
||||
} from "./read.ts";
|
||||
export type { ExecutionToolContext } from "./tool-context.ts";
|
||||
export { createWriteTool, type WriteToolInput } from "./write.ts";
|
||||
31
third_party/pi/packages/agent/src/harness/tools/path-utils.ts
vendored
Normal file
31
third_party/pi/packages/agent/src/harness/tools/path-utils.ts
vendored
Normal file
@@ -0,0 +1,31 @@
|
||||
import type { Context } from "../context.ts";
|
||||
import type { ExecutionEnv } from "../types.ts";
|
||||
import { getOrThrow } from "../types.ts";
|
||||
|
||||
const UNICODE_SPACES = /[\u00A0\u2000-\u200A\u202F\u205F\u3000]/g;
|
||||
const NARROW_NO_BREAK_SPACE = "\u202F";
|
||||
|
||||
function normalizeToolPath(path: string): string {
|
||||
const normalized = path.replace(UNICODE_SPACES, " ");
|
||||
return normalized.startsWith("@") ? normalized.slice(1) : normalized;
|
||||
}
|
||||
|
||||
export async function resolveToolPath(env: ExecutionEnv, path: string, context: Context): Promise<string> {
|
||||
return getOrThrow(await env.absolutePath(normalizeToolPath(path), context));
|
||||
}
|
||||
|
||||
export async function resolveReadToolPath(env: ExecutionEnv, path: string, context: Context): Promise<string> {
|
||||
const resolved = await resolveToolPath(env, path, context);
|
||||
const variants = [
|
||||
resolved,
|
||||
resolved.replace(/ (AM|PM)\./gi, `${NARROW_NO_BREAK_SPACE}$1.`),
|
||||
resolved.normalize("NFD"),
|
||||
resolved.replace(/'/g, "\u2019"),
|
||||
resolved.normalize("NFD").replace(/'/g, "\u2019"),
|
||||
];
|
||||
|
||||
for (const variant of new Set(variants)) {
|
||||
if (getOrThrow(await env.exists(variant, context))) return variant;
|
||||
}
|
||||
return resolved;
|
||||
}
|
||||
149
third_party/pi/packages/agent/src/harness/tools/read.ts
vendored
Normal file
149
third_party/pi/packages/agent/src/harness/tools/read.ts
vendored
Normal file
@@ -0,0 +1,149 @@
|
||||
import type { ImageContent, TextContent } from "@earendil-works/pi-ai";
|
||||
import { type Static, Type } from "typebox";
|
||||
import type { Context } from "../context.ts";
|
||||
import type { AgentHarnessTool } from "../types.ts";
|
||||
import { getOrThrow } from "../types.ts";
|
||||
import {
|
||||
DEFAULT_MAX_BYTES,
|
||||
DEFAULT_MAX_LINES,
|
||||
formatSize,
|
||||
type TruncationResult,
|
||||
truncateHead,
|
||||
} from "../utils/truncate.ts";
|
||||
import { detectSupportedImageMimeType, encodeBase64 } from "./image.ts";
|
||||
import { resolveReadToolPath } from "./path-utils.ts";
|
||||
import type { ExecutionToolContext } from "./tool-context.ts";
|
||||
|
||||
const readSchema = Type.Object({
|
||||
path: Type.String({ description: "Path to the file to read (relative or absolute)" }),
|
||||
offset: Type.Optional(Type.Number({ description: "Line number to start reading from (1-indexed)" })),
|
||||
limit: Type.Optional(Type.Number({ description: "Maximum number of lines to read" })),
|
||||
});
|
||||
|
||||
export type ReadToolInput = Static<typeof readSchema>;
|
||||
|
||||
export interface ReadToolDetails {
|
||||
truncation?: TruncationResult;
|
||||
}
|
||||
|
||||
export type ReadImageProcessorResult =
|
||||
| { ok: true; data: string; mimeType: string; hints: string[] }
|
||||
| { ok: false; message: string };
|
||||
|
||||
export type ReadImageProcessor = (
|
||||
bytes: Uint8Array,
|
||||
mimeType: string,
|
||||
options: { autoResizeImages: boolean },
|
||||
context: Context,
|
||||
) => Promise<ReadImageProcessorResult>;
|
||||
|
||||
export interface ReadToolOptions {
|
||||
/** Whether an injected image processor should resize images. Default: true. */
|
||||
autoResizeImages?: boolean;
|
||||
/** Optional image conversion/resizing implementation. */
|
||||
imageProcessor?: ReadImageProcessor;
|
||||
}
|
||||
|
||||
export function createReadTool<TContext extends ExecutionToolContext = ExecutionToolContext>(
|
||||
options?: ReadToolOptions,
|
||||
): AgentHarnessTool<TContext, typeof readSchema, ReadToolDetails | undefined> {
|
||||
return {
|
||||
name: "read",
|
||||
label: "read",
|
||||
description: `Read the contents of a file. Supports text files and images (jpg, png, gif, webp, bmp). Images are sent as attachments. For text files, output is truncated to ${DEFAULT_MAX_LINES} lines or ${DEFAULT_MAX_BYTES / 1024}KB (whichever is hit first). Use offset/limit for large files. When you need the full file, continue with offset until complete.`,
|
||||
parameters: readSchema,
|
||||
async execute(_toolCallId, { path, offset, limit }, _onUpdate, { env }, _invocation, context) {
|
||||
const absolutePath = await resolveReadToolPath(env, path, context);
|
||||
const bytes = getOrThrow(await env.readBinaryFile(absolutePath, context));
|
||||
const mimeType = detectSupportedImageMimeType(bytes);
|
||||
if (mimeType) {
|
||||
if (options?.imageProcessor) {
|
||||
const processed = await options.imageProcessor(
|
||||
bytes,
|
||||
mimeType,
|
||||
{ autoResizeImages: options.autoResizeImages ?? true },
|
||||
context,
|
||||
);
|
||||
if (!processed.ok) {
|
||||
return {
|
||||
content: [{ type: "text", text: `Read image file [${mimeType}]\n${processed.message}` }],
|
||||
details: undefined,
|
||||
};
|
||||
}
|
||||
const hints = processed.hints.length > 0 ? `\n${processed.hints.join("\n")}` : "";
|
||||
return {
|
||||
content: [
|
||||
{ type: "text", text: `Read image file [${processed.mimeType}]${hints}` },
|
||||
{ type: "image", data: processed.data, mimeType: processed.mimeType },
|
||||
] satisfies Array<TextContent | ImageContent>,
|
||||
details: undefined,
|
||||
};
|
||||
}
|
||||
if (mimeType === "image/bmp") {
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: "Read image file [image/bmp]\n[Image omitted: configure an imageProcessor to convert BMP images.]",
|
||||
},
|
||||
],
|
||||
details: undefined,
|
||||
};
|
||||
}
|
||||
return {
|
||||
content: [
|
||||
{ type: "text", text: `Read image file [${mimeType}]` },
|
||||
{ type: "image", data: encodeBase64(bytes), mimeType },
|
||||
] satisfies Array<TextContent | ImageContent>,
|
||||
details: undefined,
|
||||
};
|
||||
}
|
||||
|
||||
const textContent = new TextDecoder().decode(bytes);
|
||||
const allLines = textContent.split("\n");
|
||||
const totalFileLines = allLines.length;
|
||||
const startLine = offset ? Math.max(0, offset - 1) : 0;
|
||||
const startLineDisplay = startLine + 1;
|
||||
if (startLine >= allLines.length) {
|
||||
throw new Error(`Offset ${offset} is beyond end of file (${allLines.length} lines total)`);
|
||||
}
|
||||
|
||||
let selectedContent: string;
|
||||
let userLimitedLines: number | undefined;
|
||||
if (limit !== undefined) {
|
||||
const endLine = Math.min(startLine + limit, allLines.length);
|
||||
selectedContent = allLines.slice(startLine, endLine).join("\n");
|
||||
userLimitedLines = endLine - startLine;
|
||||
} else {
|
||||
selectedContent = allLines.slice(startLine).join("\n");
|
||||
}
|
||||
|
||||
const truncation = truncateHead(selectedContent);
|
||||
let outputText: string;
|
||||
let details: ReadToolDetails | undefined;
|
||||
if (truncation.firstLineExceedsLimit) {
|
||||
const firstLineSize = formatSize(new TextEncoder().encode(allLines[startLine]).byteLength);
|
||||
outputText = `[Line ${startLineDisplay} is ${firstLineSize}, exceeds ${formatSize(DEFAULT_MAX_BYTES)} limit. Use bash: sed -n '${startLineDisplay}p' ${path} | head -c ${DEFAULT_MAX_BYTES}]`;
|
||||
details = { truncation };
|
||||
} else if (truncation.truncated) {
|
||||
const endLineDisplay = startLineDisplay + truncation.outputLines - 1;
|
||||
const nextOffset = endLineDisplay + 1;
|
||||
outputText = truncation.content;
|
||||
if (truncation.truncatedBy === "lines") {
|
||||
outputText += `\n\n[Showing lines ${startLineDisplay}-${endLineDisplay} of ${totalFileLines}. Use offset=${nextOffset} to continue.]`;
|
||||
} else {
|
||||
outputText += `\n\n[Showing lines ${startLineDisplay}-${endLineDisplay} of ${totalFileLines} (${formatSize(DEFAULT_MAX_BYTES)} limit). Use offset=${nextOffset} to continue.]`;
|
||||
}
|
||||
details = { truncation };
|
||||
} else if (userLimitedLines !== undefined && startLine + userLimitedLines < allLines.length) {
|
||||
const remaining = allLines.length - (startLine + userLimitedLines);
|
||||
const nextOffset = startLine + userLimitedLines + 1;
|
||||
outputText = `${truncation.content}\n\n[${remaining} more lines in file. Use offset=${nextOffset} to continue.]`;
|
||||
} else {
|
||||
outputText = truncation.content;
|
||||
}
|
||||
|
||||
return { content: [{ type: "text", text: outputText }], details };
|
||||
},
|
||||
};
|
||||
}
|
||||
6
third_party/pi/packages/agent/src/harness/tools/tool-context.ts
vendored
Normal file
6
third_party/pi/packages/agent/src/harness/tools/tool-context.ts
vendored
Normal file
@@ -0,0 +1,6 @@
|
||||
import type { ExecutionEnv } from "../types.ts";
|
||||
|
||||
/** Filesystem and shell context required by the built-in execution tools. */
|
||||
export interface ExecutionToolContext {
|
||||
env: ExecutionEnv;
|
||||
}
|
||||
44
third_party/pi/packages/agent/src/harness/tools/write.ts
vendored
Normal file
44
third_party/pi/packages/agent/src/harness/tools/write.ts
vendored
Normal file
@@ -0,0 +1,44 @@
|
||||
import { type Static, Type } from "typebox";
|
||||
import type { AgentHarnessTool } from "../types.ts";
|
||||
import { getOrThrow } from "../types.ts";
|
||||
import { withFileMutationQueue } from "./file-mutation-queue.ts";
|
||||
import { resolveToolPath } from "./path-utils.ts";
|
||||
import type { ExecutionToolContext } from "./tool-context.ts";
|
||||
|
||||
const writeSchema = Type.Object({
|
||||
path: Type.String({ description: "Path to the file to write (relative or absolute)" }),
|
||||
content: Type.String({ description: "Content to write to the file" }),
|
||||
});
|
||||
|
||||
export type WriteToolInput = Static<typeof writeSchema>;
|
||||
|
||||
export function createWriteTool<TContext extends ExecutionToolContext = ExecutionToolContext>(): AgentHarnessTool<
|
||||
TContext,
|
||||
typeof writeSchema,
|
||||
undefined
|
||||
> {
|
||||
return {
|
||||
name: "write",
|
||||
label: "write",
|
||||
description:
|
||||
"Write content to a file. Creates the file if it doesn't exist, overwrites if it does. Automatically creates parent directories.",
|
||||
parameters: writeSchema,
|
||||
async execute(_toolCallId, { path, content }, _onUpdate, { env }, _invocation, context) {
|
||||
const absolutePath = await resolveToolPath(env, path, context);
|
||||
return withFileMutationQueue(
|
||||
env,
|
||||
absolutePath,
|
||||
async () => {
|
||||
if (context.abortSignal?.aborted) throw new Error("Operation aborted");
|
||||
getOrThrow(await env.writeFile(absolutePath, content, context));
|
||||
if (context.abortSignal?.aborted) throw new Error("Operation aborted");
|
||||
return {
|
||||
content: [{ type: "text", text: `Successfully wrote to ${path}` }],
|
||||
details: undefined,
|
||||
};
|
||||
},
|
||||
context,
|
||||
);
|
||||
},
|
||||
};
|
||||
}
|
||||
391
third_party/pi/packages/agent/src/harness/types.ts
vendored
Normal file
391
third_party/pi/packages/agent/src/harness/types.ts
vendored
Normal file
@@ -0,0 +1,391 @@
|
||||
import type { SimpleStreamOptions, Transport } from "@earendil-works/pi-ai";
|
||||
import type { Static, TSchema } from "typebox";
|
||||
import type { AgentTool, AgentToolResult } from "../types.ts";
|
||||
import type { Context } from "./context.ts";
|
||||
import type { JsonValue } from "./session/types.ts";
|
||||
import type { TruncationResult } from "./utils/truncate.ts";
|
||||
|
||||
/** Result of a fallible operation. Expected failures are returned as `ok: false` instead of thrown. */
|
||||
export type Result<TValue, TError> = { ok: true; value: TValue } | { ok: false; error: TError };
|
||||
|
||||
/** Create a successful {@link Result}. */
|
||||
export function ok<TValue, TError>(value: TValue): Result<TValue, TError> {
|
||||
return { ok: true, value };
|
||||
}
|
||||
|
||||
/** Create a failed {@link Result}. */
|
||||
export function err<TValue, TError>(error: TError): Result<TValue, TError> {
|
||||
return { ok: false, error };
|
||||
}
|
||||
|
||||
/** Return the success value or throw the failure error. Intended for tests and explicit adapter boundaries. */
|
||||
export function getOrThrow<TValue, TError>(result: Result<TValue, TError>): TValue {
|
||||
if (!result.ok) throw result.error;
|
||||
return result.value;
|
||||
}
|
||||
|
||||
/** Return the success value or `undefined`. Only object values are allowed to avoid truthiness bugs with primitives. */
|
||||
export function getOrUndefined<TValue extends object, TError>(result: Result<TValue, TError>): TValue | undefined {
|
||||
return result.ok ? result.value : undefined;
|
||||
}
|
||||
|
||||
/** Normalize unknown thrown values into Error instances before using them as typed error causes. */
|
||||
export function toError(error: unknown): Error {
|
||||
if (error instanceof Error) return error;
|
||||
if (typeof error === "string") return new Error(error);
|
||||
try {
|
||||
return new Error(JSON.stringify(error));
|
||||
} catch {
|
||||
return new Error(String(error));
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Skill loaded from a `SKILL.md` file or provided by an application.
|
||||
*
|
||||
* `name`, `description`, and `filePath` are inserted into the system prompt in an XML-formatted block as suggested by agentskills.io.
|
||||
* Use {@link formatSkillsForSystemPrompt} to generate the spec-compatible system prompt block.
|
||||
*/
|
||||
export interface Skill {
|
||||
/** Stable skill name used for lookup and model-visible listings. */
|
||||
name: string;
|
||||
/** Short model-visible description of when to use the skill. */
|
||||
description: string;
|
||||
/** Full skill instructions. */
|
||||
content: string;
|
||||
/** Absolute path to the skill file. Used for model-visible location and resolving relative references. */
|
||||
filePath: string;
|
||||
/** Exclude this skill from model-visible skill lists while still allowing explicit application invocation. */
|
||||
disableModelInvocation?: boolean;
|
||||
}
|
||||
|
||||
/** Prompt template that can be formatted into a prompt for explicit invocation. */
|
||||
export interface PromptTemplate {
|
||||
/** Stable template name used for lookup or application command routing. */
|
||||
name: string;
|
||||
/** Optional description for command lists or autocomplete. */
|
||||
description?: string;
|
||||
/** Template content. Argument placeholders are formatted by `formatPromptTemplateInvocation`. */
|
||||
content: string;
|
||||
}
|
||||
|
||||
/** Resources made available to explicit invocation methods and system-prompt callbacks. */
|
||||
export interface AgentHarnessResources<
|
||||
TSkill extends Skill = Skill,
|
||||
TPromptTemplate extends PromptTemplate = PromptTemplate,
|
||||
> {
|
||||
/** Prompt templates available for explicit invocation. */
|
||||
promptTemplates?: TPromptTemplate[];
|
||||
/** Skills available to the model and explicit skill invocation. */
|
||||
skills?: TSkill[];
|
||||
}
|
||||
|
||||
/** Options for one live harness tool progress update. */
|
||||
export interface AgentHarnessToolUpdateOptions {
|
||||
/** Request replacement of this invocation's durable recovery checkpoint. */
|
||||
checkpoint?: true;
|
||||
}
|
||||
|
||||
/** Synchronous full-snapshot progress callback supplied to harness-native tools. */
|
||||
export type AgentHarnessToolUpdateCallback<TDetails> = (
|
||||
partialResult: AgentToolResult<TDetails>,
|
||||
options?: AgentHarnessToolUpdateOptions,
|
||||
) => void;
|
||||
|
||||
/** Stable harness identity for one logical tool call, unchanged during safe replay. */
|
||||
export interface AgentHarnessToolInvocation {
|
||||
/** Opaque session-unique id equal to the call's reserved result-entry id. */
|
||||
readonly invocationId: string;
|
||||
readonly operationId: string;
|
||||
readonly turnId: string;
|
||||
/** Read one invocation-scoped durable replay memo. */
|
||||
getMemo(name: string): Promise<JsonValue | undefined>;
|
||||
/** Set or delete one invocation-scoped durable replay memo. */
|
||||
setMemo(name: string, value: JsonValue | undefined): Promise<void>;
|
||||
}
|
||||
|
||||
/** Tool definition executed by an {@link AgentHarness} with an application-defined context. */
|
||||
export type AgentHarnessTool<
|
||||
TContext extends object | undefined,
|
||||
TParameters extends TSchema = TSchema,
|
||||
TDetails = unknown,
|
||||
> = Omit<AgentTool<TParameters, TDetails>, "execute"> & {
|
||||
/** Execute the tool call with the context resolved for the current turn snapshot. */
|
||||
execute(
|
||||
toolCallId: string,
|
||||
params: Static<TParameters>,
|
||||
onUpdate: AgentHarnessToolUpdateCallback<TDetails>,
|
||||
toolContext: TContext,
|
||||
invocation: AgentHarnessToolInvocation,
|
||||
context: Context,
|
||||
): Promise<AgentToolResult<TDetails>>;
|
||||
};
|
||||
|
||||
/** Static tool context or provider resolved for each turn snapshot. */
|
||||
export type AgentHarnessToolContextSource<TContext extends object | undefined> =
|
||||
| TContext
|
||||
| ((context: Context) => TContext | Promise<TContext>);
|
||||
|
||||
/** Curated provider request options owned by the harness and snapshotted per turn. */
|
||||
export interface AgentHarnessStreamOptions {
|
||||
/** Preferred transport forwarded to the stream function. */
|
||||
transport?: Transport;
|
||||
/** Provider request timeout in milliseconds. */
|
||||
timeoutMs?: number;
|
||||
/** Maximum provider retry attempts. */
|
||||
maxRetries?: number;
|
||||
/** Optional cap for provider-requested retry delays. */
|
||||
maxRetryDelayMs?: number;
|
||||
/** Additional request headers merged with auth and lifecycle headers. */
|
||||
headers?: Record<string, string>;
|
||||
/** Provider metadata forwarded with requests. */
|
||||
metadata?: SimpleStreamOptions["metadata"];
|
||||
/** Provider cache retention hint. */
|
||||
cacheRetention?: SimpleStreamOptions["cacheRetention"];
|
||||
/** Ask a capable provider to continue generation asynchronously. */
|
||||
deferred?: boolean | { window?: "15m" | "1h" | "24h" };
|
||||
}
|
||||
|
||||
/** Per-request stream option patch returned by provider hooks. */
|
||||
export interface AgentHarnessStreamOptionsPatch
|
||||
extends Omit<Partial<AgentHarnessStreamOptions>, "headers" | "metadata"> {
|
||||
/** Header patch. `undefined` values delete keys; explicit `headers: undefined` clears all headers. */
|
||||
headers?: Record<string, string | undefined>;
|
||||
/** Metadata patch. `undefined` values delete keys; explicit `metadata: undefined` clears all metadata. */
|
||||
metadata?: Record<string, unknown | undefined>;
|
||||
}
|
||||
|
||||
/** Kind of filesystem object as addressed by a {@link FileSystem}. Symlinks are not followed automatically. */
|
||||
export type FileKind = "file" | "directory" | "symlink";
|
||||
|
||||
/** Stable, backend-independent file error codes returned by {@link FileSystem} file operations. */
|
||||
export type FileErrorCode =
|
||||
| "aborted"
|
||||
| "not_found"
|
||||
| "permission_denied"
|
||||
| "not_directory"
|
||||
| "is_directory"
|
||||
| "invalid"
|
||||
| "not_supported"
|
||||
| "unknown";
|
||||
|
||||
/** Error returned by {@link FileSystem} file operations. */
|
||||
export class FileError extends Error {
|
||||
/** Backend-independent error code. */
|
||||
public code: FileErrorCode;
|
||||
/** Absolute addressed path associated with the failure, when available. */
|
||||
public path?: string;
|
||||
|
||||
constructor(code: FileErrorCode, message: string, path?: string, cause?: Error) {
|
||||
super(message, cause === undefined ? undefined : { cause });
|
||||
this.name = "FileError";
|
||||
this.code = code;
|
||||
this.path = path;
|
||||
}
|
||||
}
|
||||
|
||||
/** Stable, backend-independent execution error codes returned by {@link ExecutionEnv.exec}. */
|
||||
export type ExecutionErrorCode =
|
||||
| "aborted"
|
||||
| "timeout"
|
||||
| "shell_unavailable"
|
||||
| "spawn_error"
|
||||
| "callback_error"
|
||||
| "unknown";
|
||||
|
||||
/** Error returned by {@link ExecutionEnv.exec}. */
|
||||
export class ExecutionError extends Error {
|
||||
/** Backend-independent error code. */
|
||||
public code: ExecutionErrorCode;
|
||||
|
||||
constructor(code: ExecutionErrorCode, message: string, cause?: Error) {
|
||||
super(message, cause === undefined ? undefined : { cause });
|
||||
this.name = "ExecutionError";
|
||||
this.code = code;
|
||||
}
|
||||
}
|
||||
|
||||
/** Stable compaction error codes returned by compaction helpers. */
|
||||
export type CompactionErrorCode = "aborted" | "summarization_failed";
|
||||
|
||||
/** Error returned by compaction helpers. */
|
||||
export class CompactionError extends Error {
|
||||
/** Backend-independent error code. */
|
||||
public code: CompactionErrorCode;
|
||||
|
||||
constructor(code: CompactionErrorCode, message: string, cause?: Error) {
|
||||
super(message, cause === undefined ? undefined : { cause });
|
||||
this.name = "CompactionError";
|
||||
this.code = code;
|
||||
}
|
||||
}
|
||||
|
||||
/** Stable branch-summary error codes returned by branch summarization helpers. */
|
||||
export type BranchSummaryErrorCode = "aborted" | "summarization_failed";
|
||||
|
||||
/** Error returned by branch summarization helpers. */
|
||||
export class BranchSummaryError extends Error {
|
||||
/** Backend-independent error code. */
|
||||
public code: BranchSummaryErrorCode;
|
||||
|
||||
constructor(code: BranchSummaryErrorCode, message: string, cause?: Error) {
|
||||
super(message, cause === undefined ? undefined : { cause });
|
||||
this.name = "BranchSummaryError";
|
||||
this.code = code;
|
||||
}
|
||||
}
|
||||
|
||||
/** Metadata for one filesystem object in a {@link FileSystem}. */
|
||||
export interface FileInfo {
|
||||
/** Basename of {@link path}. */
|
||||
name: string;
|
||||
/** Absolute, syntactically normalized addressed path in the execution environment. Symlinks are not followed. */
|
||||
path: string;
|
||||
/** Object kind. Symlink targets are not followed; use {@link FileSystem.canonicalPath} explicitly. */
|
||||
kind: FileKind;
|
||||
/** Size in bytes for the addressed filesystem object. */
|
||||
size: number;
|
||||
/** Modification time as milliseconds since Unix epoch. */
|
||||
mtimeMs: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Filesystem capability used by the harness.
|
||||
*
|
||||
* Paths passed to methods may be absolute or relative to {@link cwd}. Paths returned by file operations are addressed paths
|
||||
* in the filesystem namespace, but are not canonicalized through symlinks unless returned by {@link canonicalPath}.
|
||||
*
|
||||
* Operation methods must never throw or reject. All filesystem failures, including unexpected backend failures, must be
|
||||
* encoded in the returned {@link Result}. Implementations must preserve this invariant.
|
||||
*/
|
||||
export interface FileSystem {
|
||||
/** Current working directory for relative paths. */
|
||||
cwd: string;
|
||||
|
||||
/** Return an absolute addressed path without requiring it to exist and without resolving symlinks. */
|
||||
absolutePath(path: string, context: Context): Promise<Result<string, FileError>>;
|
||||
/** Join path segments in the filesystem namespace without requiring the result to exist. */
|
||||
joinPath(parts: string[], context: Context): Promise<Result<string, FileError>>;
|
||||
/** Read a UTF-8 text file. */
|
||||
readTextFile(path: string, context: Context): Promise<Result<string, FileError>>;
|
||||
/** Read UTF-8 text lines. Implementations should stop once `maxLines` lines have been read. */
|
||||
readTextLines(
|
||||
path: string,
|
||||
options: { maxLines?: number } | undefined,
|
||||
context: Context,
|
||||
): Promise<Result<string[], FileError>>;
|
||||
/** Read a binary file. */
|
||||
readBinaryFile(path: string, context: Context): Promise<Result<Uint8Array, FileError>>;
|
||||
/** Create or overwrite a file, creating parent directories when supported. */
|
||||
writeFile(path: string, content: string | Uint8Array, context: Context): Promise<Result<void, FileError>>;
|
||||
/** Create or append to a file, creating parent directories when supported. */
|
||||
appendFile(path: string, content: string | Uint8Array, context: Context): Promise<Result<void, FileError>>;
|
||||
/** Atomically rename a file, replacing the destination when it exists. Does not copy across filesystems. */
|
||||
renameFile(sourcePath: string, destinationPath: string, context: Context): Promise<Result<void, FileError>>;
|
||||
/** Return metadata for the addressed path without following symlinks. */
|
||||
fileInfo(path: string, context: Context): Promise<Result<FileInfo, FileError>>;
|
||||
/** List direct children of a directory without following symlinks. */
|
||||
listDir(path: string, context: Context): Promise<Result<FileInfo[], FileError>>;
|
||||
/** Return the canonical path for an existing path, resolving symlinks where supported. */
|
||||
canonicalPath(path: string, context: Context): Promise<Result<string, FileError>>;
|
||||
/** Return false for missing paths. Other errors, such as permission failures, return a {@link FileError}. */
|
||||
exists(path: string, context: Context): Promise<Result<boolean, FileError>>;
|
||||
/** Create a directory. Defaults to `recursive: true`. */
|
||||
createDir(
|
||||
path: string,
|
||||
options: { recursive?: boolean } | undefined,
|
||||
context: Context,
|
||||
): Promise<Result<void, FileError>>;
|
||||
/** Remove a file or directory. Defaults to `recursive: false` and `force: false`. */
|
||||
remove(
|
||||
path: string,
|
||||
options: { recursive?: boolean; force?: boolean } | undefined,
|
||||
context: Context,
|
||||
): Promise<Result<void, FileError>>;
|
||||
/** Create a temporary directory and return its absolute path. Defaults to `prefix: "tmp-"`. */
|
||||
createTempDir(prefix: string | undefined, context: Context): Promise<Result<string, FileError>>;
|
||||
/** Create a temporary file and return its absolute path. Defaults to `prefix: ""` and `suffix: ""`. */
|
||||
createTempFile(
|
||||
options: { prefix?: string; suffix?: string } | undefined,
|
||||
context: Context,
|
||||
): Promise<Result<string, FileError>>;
|
||||
|
||||
/** Release filesystem resources. Must be best-effort and must not throw or reject. */
|
||||
cleanup(context: Context): Promise<void>;
|
||||
}
|
||||
|
||||
/** Which portion of bounded output survives after the limit is crossed. */
|
||||
export type ShellOutputRetention = "head" | "tail";
|
||||
|
||||
/** Source-side limits for one combined shell output view. */
|
||||
export interface ShellOutputLimits {
|
||||
maxBytes: number;
|
||||
maxLines: number;
|
||||
/** Defaults to `"tail"`. */
|
||||
retain?: ShellOutputRetention;
|
||||
}
|
||||
|
||||
/** Bounded shell capture requested by the caller. */
|
||||
export interface ShellOutputCaptureOptions {
|
||||
limits: ShellOutputLimits;
|
||||
/** Preserve complete output in an execution-environment-local file after the limits are crossed. */
|
||||
spill?: boolean;
|
||||
}
|
||||
|
||||
/** Truncation metadata without a duplicate copy of the retained text. */
|
||||
export type ShellOutputTruncation = Omit<TruncationResult, "content">;
|
||||
|
||||
/** Metadata accompanying a bounded shell output view. */
|
||||
export interface ShellOutputMetadata {
|
||||
truncation: ShellOutputTruncation;
|
||||
spillPath?: string;
|
||||
lastLineBytes?: number;
|
||||
}
|
||||
|
||||
/** Complete bounded shell output view. */
|
||||
export interface ShellOutputView extends ShellOutputMetadata {
|
||||
text: string;
|
||||
}
|
||||
|
||||
/** Incremental source-side change to one bounded shell output view. */
|
||||
export type ShellOutputUpdate =
|
||||
| { kind: "replace"; output: ShellOutputView }
|
||||
| { kind: "append"; text: string; metadata: ShellOutputMetadata }
|
||||
| { kind: "slide"; drop: number; text: string; metadata: ShellOutputMetadata }
|
||||
| { kind: "metadata"; metadata: ShellOutputMetadata };
|
||||
|
||||
/** Bounded shell completion. Output text is delivered through {@link ShellExecOptions.onUpdate}. */
|
||||
export interface ShellExecResult extends ShellOutputMetadata {
|
||||
exitCode: number;
|
||||
}
|
||||
|
||||
/** Options for {@link Shell.exec}. */
|
||||
export interface ShellExecOptions {
|
||||
/** Working directory for the command. Relative paths are resolved against {@link ExecutionEnv.cwd}. Defaults to {@link ExecutionEnv.cwd}. */
|
||||
cwd?: string;
|
||||
/** Environment variables for the command. Values override inherited defaults when `inheritEnv` is true. */
|
||||
env?: Record<string, string>;
|
||||
/** Whether to inherit the execution environment's default variables. Defaults to true. */
|
||||
inheritEnv?: boolean;
|
||||
/** Timeout in seconds. Implementations should return a timeout error when the command exceeds this duration. Defaults to no timeout. */
|
||||
timeout?: number;
|
||||
/** Source-side bounded capture. Output is discarded when this and `onUpdate` are both absent. */
|
||||
capture?: ShellOutputCaptureOptions;
|
||||
/** Called with bounded output changes. */
|
||||
onUpdate?: (update: ShellOutputUpdate, context: Context) => void;
|
||||
}
|
||||
|
||||
/** Shell execution capability used by the harness. */
|
||||
export interface Shell {
|
||||
/** Execute a shell command in {@link FileSystem.cwd} unless `options.cwd` is provided. */
|
||||
exec(
|
||||
command: string,
|
||||
options: ShellExecOptions | undefined,
|
||||
context: Context,
|
||||
): Promise<Result<ShellExecResult, ExecutionError>>;
|
||||
/** Release shell resources. Must be best-effort and must not throw or reject. */
|
||||
cleanup(context: Context): Promise<void>;
|
||||
}
|
||||
|
||||
/** Filesystem and process execution environment used by the harness. */
|
||||
export interface ExecutionEnv extends FileSystem, Shell {}
|
||||
87
third_party/pi/packages/agent/src/harness/utils/adaptive-publisher.ts
vendored
Normal file
87
third_party/pi/packages/agent/src/harness/utils/adaptive-publisher.ts
vendored
Normal file
@@ -0,0 +1,87 @@
|
||||
export interface AdaptivePublisherOptions<TValue, TUpdate> {
|
||||
snapshot(): TValue;
|
||||
update(previous: TValue | undefined, current: TValue): TUpdate | undefined;
|
||||
measure(update: TUpdate): number;
|
||||
publish(update: TUpdate): void;
|
||||
onError(error: unknown): void;
|
||||
minIntervalMs?: number;
|
||||
targetBytesPerSecond?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Publishes the latest state without queuing intermediate mutations.
|
||||
*
|
||||
* The first dirty state after idle is immediate. Each publication then buys a
|
||||
* delay proportional to its encoded size, with a minimum interval that also
|
||||
* bounds event count. A single trailing timer guarantees eventual publication.
|
||||
*/
|
||||
export class AdaptivePublisher<TValue, TUpdate> {
|
||||
readonly #options: AdaptivePublisherOptions<TValue, TUpdate>;
|
||||
readonly #minIntervalMs: number;
|
||||
readonly #targetBytesPerSecond: number;
|
||||
#published: TValue | undefined;
|
||||
#dirty = false;
|
||||
#nextEmitAt = 0;
|
||||
#timer: ReturnType<typeof setTimeout> | undefined;
|
||||
#disposed = false;
|
||||
|
||||
constructor(options: AdaptivePublisherOptions<TValue, TUpdate>) {
|
||||
this.#options = options;
|
||||
this.#minIntervalMs = options.minIntervalMs ?? 100;
|
||||
this.#targetBytesPerSecond = options.targetBytesPerSecond ?? 100 * 1024;
|
||||
}
|
||||
|
||||
markDirty(): void {
|
||||
if (this.#disposed) return;
|
||||
this.#dirty = true;
|
||||
const wait = this.#nextEmitAt - Date.now();
|
||||
if (wait <= 0) {
|
||||
this.flush();
|
||||
return;
|
||||
}
|
||||
this.#armTimer(wait);
|
||||
}
|
||||
|
||||
flush(force = false): void {
|
||||
if (this.#disposed || !this.#dirty) return;
|
||||
const now = Date.now();
|
||||
if (!force && now < this.#nextEmitAt) {
|
||||
this.#armTimer(this.#nextEmitAt - now);
|
||||
return;
|
||||
}
|
||||
if (this.#timer !== undefined) clearTimeout(this.#timer);
|
||||
this.#timer = undefined;
|
||||
const current = this.#options.snapshot();
|
||||
const update = this.#options.update(this.#published, current);
|
||||
if (update === undefined) {
|
||||
this.#published = current;
|
||||
this.#dirty = false;
|
||||
return;
|
||||
}
|
||||
const encodedBytes = this.#options.measure(update);
|
||||
this.#published = current;
|
||||
this.#dirty = false;
|
||||
this.#nextEmitAt = now + Math.max(this.#minIntervalMs, (encodedBytes * 1000) / this.#targetBytesPerSecond);
|
||||
// Commit before delivery. A consumer may apply the update and then throw or
|
||||
// reenter the producer; retaining the old baseline would duplicate that delta.
|
||||
this.#options.publish(update);
|
||||
}
|
||||
|
||||
dispose(): void {
|
||||
if (this.#timer !== undefined) clearTimeout(this.#timer);
|
||||
this.#timer = undefined;
|
||||
this.#disposed = true;
|
||||
}
|
||||
|
||||
#armTimer(wait: number): void {
|
||||
if (this.#timer !== undefined) return;
|
||||
this.#timer = setTimeout(() => {
|
||||
this.#timer = undefined;
|
||||
try {
|
||||
this.flush();
|
||||
} catch (error) {
|
||||
this.#options.onError(error);
|
||||
}
|
||||
}, wait);
|
||||
}
|
||||
}
|
||||
238
third_party/pi/packages/agent/src/harness/utils/output-capture.ts
vendored
Normal file
238
third_party/pi/packages/agent/src/harness/utils/output-capture.ts
vendored
Normal file
@@ -0,0 +1,238 @@
|
||||
import type { Context } from "../context.ts";
|
||||
import type { ShellOutputCaptureOptions, ShellOutputMetadata, ShellOutputUpdate, ShellOutputView } from "../types.ts";
|
||||
import { AdaptivePublisher } from "./adaptive-publisher.ts";
|
||||
import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, truncateHead, truncateTail, utf8ByteLength } from "./truncate.ts";
|
||||
|
||||
export const OUTPUT_MIN_EMIT_INTERVAL_MS = 100;
|
||||
export const OUTPUT_TARGET_BYTES_PER_SECOND = 100 * 1024;
|
||||
|
||||
const INVALID_SHELL_OUTPUT = /[\x00-\x08\x0b-\x1f\ufff9-\ufffb]/g;
|
||||
const textEncoder = new TextEncoder();
|
||||
const textDecoder = new TextDecoder();
|
||||
|
||||
interface OutputCaptureHandlers {
|
||||
onUpdate?: (update: ShellOutputUpdate, context: Context) => void;
|
||||
onError(error: unknown): void;
|
||||
}
|
||||
|
||||
/**
|
||||
* Maintains and publishes one bounded shell-output view.
|
||||
*
|
||||
* Writes received while publication is rate-limited collapse into the latest
|
||||
* view. Small changes remain responsive; complete window turnovers purchase a
|
||||
* proportionally longer delay. The first update after idle and an explicit
|
||||
* final flush are immediate.
|
||||
*/
|
||||
export class OutputCapture {
|
||||
readonly #maxBytes: number;
|
||||
readonly #maxLines: number;
|
||||
readonly #retain: "head" | "tail";
|
||||
readonly #context: Context;
|
||||
readonly #onUpdate: OutputCaptureHandlers["onUpdate"];
|
||||
|
||||
readonly #decoder = new TextDecoder();
|
||||
#buffer = "";
|
||||
#bufferBytes = 0;
|
||||
#totalBytes = 0;
|
||||
#newlines = 0;
|
||||
#endsWithNewline = true;
|
||||
#currentLineBytes = 0;
|
||||
#spillPath: string | undefined;
|
||||
#disposed = false;
|
||||
readonly #publisher: AdaptivePublisher<ShellOutputView, ShellOutputUpdate>;
|
||||
|
||||
constructor(options: ShellOutputCaptureOptions | undefined, context: Context, handlers: OutputCaptureHandlers) {
|
||||
this.#maxBytes = options?.limits.maxBytes ?? DEFAULT_MAX_BYTES;
|
||||
this.#maxLines = options?.limits.maxLines ?? DEFAULT_MAX_LINES;
|
||||
this.#retain = options?.limits.retain ?? "tail";
|
||||
this.#context = context;
|
||||
this.#onUpdate = handlers.onUpdate;
|
||||
if (!Number.isFinite(this.#maxBytes) || this.#maxBytes <= 0) {
|
||||
throw new TypeError("Output maxBytes must be a positive finite number");
|
||||
}
|
||||
if (!Number.isInteger(this.#maxLines) || this.#maxLines <= 0) {
|
||||
throw new TypeError("Output maxLines must be a positive integer");
|
||||
}
|
||||
this.#publisher = new AdaptivePublisher({
|
||||
snapshot: () => this.snapshot(),
|
||||
update: updateFrom,
|
||||
measure: (update) => utf8ByteLength(JSON.stringify(update)),
|
||||
publish: (update) => this.#onUpdate?.(update, this.#context),
|
||||
onError: handlers.onError,
|
||||
minIntervalMs: OUTPUT_MIN_EMIT_INTERVAL_MS,
|
||||
targetBytesPerSecond: OUTPUT_TARGET_BYTES_PER_SECOND,
|
||||
});
|
||||
}
|
||||
|
||||
get truncated(): boolean {
|
||||
return this.#totalBytes > this.#maxBytes || this.#totalLines() > this.#maxLines;
|
||||
}
|
||||
|
||||
push(chunk: string | Uint8Array): void {
|
||||
if (this.#disposed) return;
|
||||
if (typeof chunk === "string") {
|
||||
this.#appendText(this.#decoder.decode());
|
||||
this.#appendText(chunk);
|
||||
return;
|
||||
}
|
||||
this.#appendText(this.#decoder.decode(chunk, { stream: true }));
|
||||
}
|
||||
|
||||
finish(): void {
|
||||
if (this.#disposed) return;
|
||||
this.#appendText(this.#decoder.decode());
|
||||
}
|
||||
|
||||
setSpillPath(path: string): void {
|
||||
if (this.#disposed || this.#spillPath === path) return;
|
||||
this.#spillPath = path;
|
||||
this.#publisher.markDirty();
|
||||
this.flush();
|
||||
}
|
||||
|
||||
snapshot(): ShellOutputView {
|
||||
const retained =
|
||||
this.#retain === "head"
|
||||
? truncateHead(this.#buffer, { maxBytes: this.#maxBytes, maxLines: this.#maxLines })
|
||||
: truncateTail(this.#buffer, { maxBytes: this.#maxBytes, maxLines: this.#maxLines });
|
||||
const totalLines = this.#totalLines();
|
||||
const truncated = this.truncated;
|
||||
const { content, ...truncation } = retained;
|
||||
return {
|
||||
text: sanitizeShellOutput(content),
|
||||
truncation: {
|
||||
...truncation,
|
||||
truncated,
|
||||
truncatedBy: truncated ? (totalLines > this.#maxLines ? "lines" : "bytes") : null,
|
||||
totalBytes: this.#totalBytes,
|
||||
totalLines,
|
||||
},
|
||||
...(this.#spillPath === undefined ? {} : { spillPath: this.#spillPath }),
|
||||
...(retained.lastLinePartial ? { lastLineBytes: this.#currentLineBytes } : {}),
|
||||
};
|
||||
}
|
||||
|
||||
flush(): void {
|
||||
if (this.#disposed) return;
|
||||
this.#publisher.flush(true);
|
||||
}
|
||||
|
||||
dispose(): void {
|
||||
this.#publisher.dispose();
|
||||
this.#disposed = true;
|
||||
}
|
||||
|
||||
#appendText(text: string): void {
|
||||
if (text === "") return;
|
||||
const textBytes = utf8ByteLength(text);
|
||||
this.#totalBytes += textBytes;
|
||||
this.#newlines += countNewlines(text);
|
||||
this.#endsWithNewline = text.endsWith("\n");
|
||||
const lastNewline = text.lastIndexOf("\n");
|
||||
this.#currentLineBytes =
|
||||
lastNewline === -1 ? this.#currentLineBytes + textBytes : utf8ByteLength(text.slice(lastNewline + 1));
|
||||
this.#buffer += text;
|
||||
this.#bufferBytes += textBytes;
|
||||
|
||||
const guard = this.#maxBytes * 2;
|
||||
if (this.#bufferBytes > guard * 2) {
|
||||
this.#buffer =
|
||||
this.#retain === "tail"
|
||||
? trimToLastUtf8Bytes(this.#buffer, guard)
|
||||
: trimToFirstUtf8Bytes(this.#buffer, guard);
|
||||
this.#bufferBytes = utf8ByteLength(this.#buffer);
|
||||
}
|
||||
this.#publisher.markDirty();
|
||||
}
|
||||
|
||||
#totalLines(): number {
|
||||
return this.#newlines + (this.#endsWithNewline || this.#totalBytes === 0 ? 0 : 1);
|
||||
}
|
||||
}
|
||||
|
||||
export function applyShellOutputUpdate(
|
||||
current: ShellOutputView | undefined,
|
||||
update: ShellOutputUpdate,
|
||||
): ShellOutputView {
|
||||
switch (update.kind) {
|
||||
case "replace":
|
||||
return update.output;
|
||||
case "append":
|
||||
return { text: `${current?.text ?? ""}${update.text}`, ...update.metadata };
|
||||
case "slide":
|
||||
return { text: `${current?.text.slice(update.drop) ?? ""}${update.text}`, ...update.metadata };
|
||||
case "metadata":
|
||||
return { text: current?.text ?? "", ...update.metadata };
|
||||
}
|
||||
}
|
||||
|
||||
function updateFrom(previous: ShellOutputView | undefined, current: ShellOutputView): ShellOutputUpdate {
|
||||
if (previous === undefined) return { kind: "replace", output: current };
|
||||
const metadata: ShellOutputMetadata = {
|
||||
truncation: current.truncation,
|
||||
...(current.spillPath === undefined ? {} : { spillPath: current.spillPath }),
|
||||
...(current.lastLineBytes === undefined ? {} : { lastLineBytes: current.lastLineBytes }),
|
||||
};
|
||||
if (current.text === previous.text) return { kind: "metadata", metadata };
|
||||
if (current.text.length > previous.text.length && current.text.slice(0, previous.text.length) === previous.text) {
|
||||
return { kind: "append", text: current.text.slice(previous.text.length), metadata };
|
||||
}
|
||||
const shared = suffixPrefixOverlap(
|
||||
previous.text,
|
||||
current.text,
|
||||
Math.min(previous.text.length, current.text.length, current.truncation.maxBytes * 2),
|
||||
);
|
||||
if (shared > 0) {
|
||||
return {
|
||||
kind: "slide",
|
||||
drop: previous.text.length - shared,
|
||||
text: current.text.slice(shared),
|
||||
metadata,
|
||||
};
|
||||
}
|
||||
return { kind: "replace", output: current };
|
||||
}
|
||||
|
||||
function suffixPrefixOverlap(before: string, after: string, scan: number): number {
|
||||
if (before.length === 0 || after.length === 0 || scan === 0) return 0;
|
||||
const tail = before.length > scan ? before.slice(before.length - scan) : before;
|
||||
for (const probeLength of [Math.min(64, after.length), 1]) {
|
||||
const probe = after.slice(0, probeLength);
|
||||
let candidates = 0;
|
||||
for (let index = tail.indexOf(probe); index !== -1; index = tail.indexOf(probe, index + 1)) {
|
||||
if (++candidates > 8) break;
|
||||
const overlapLength = tail.length - index;
|
||||
if (overlapLength <= after.length && tail.slice(index) === after.slice(0, overlapLength)) {
|
||||
return overlapLength;
|
||||
}
|
||||
}
|
||||
if (probeLength === 1) break;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
export function sanitizeShellOutput(text: string): string {
|
||||
return text.replace(INVALID_SHELL_OUTPUT, "");
|
||||
}
|
||||
|
||||
function countNewlines(text: string): number {
|
||||
let count = 0;
|
||||
for (let index = text.indexOf("\n"); index !== -1; index = text.indexOf("\n", index + 1)) count++;
|
||||
return count;
|
||||
}
|
||||
|
||||
function trimToLastUtf8Bytes(text: string, maxBytes: number): string {
|
||||
const bytes = textEncoder.encode(text);
|
||||
if (bytes.length <= maxBytes) return text;
|
||||
let start = bytes.length - maxBytes;
|
||||
while (start < bytes.length && ((bytes[start] ?? 0) & 0xc0) === 0x80) start++;
|
||||
return textDecoder.decode(bytes.subarray(start));
|
||||
}
|
||||
|
||||
function trimToFirstUtf8Bytes(text: string, maxBytes: number): string {
|
||||
const bytes = textEncoder.encode(text);
|
||||
if (bytes.length <= maxBytes) return text;
|
||||
let end = maxBytes;
|
||||
while (end > 0 && ((bytes[end] ?? 0) & 0xc0) === 0x80) end--;
|
||||
return textDecoder.decode(bytes.subarray(0, end));
|
||||
}
|
||||
112
third_party/pi/packages/agent/src/harness/utils/shell-output.ts
vendored
Normal file
112
third_party/pi/packages/agent/src/harness/utils/shell-output.ts
vendored
Normal file
@@ -0,0 +1,112 @@
|
||||
import type { Context } from "../context.ts";
|
||||
import {
|
||||
type ExecutionEnv,
|
||||
type ExecutionError,
|
||||
err,
|
||||
ok,
|
||||
type Result,
|
||||
type ShellExecOptions,
|
||||
type ShellOutputView,
|
||||
} from "../types.ts";
|
||||
import { applyShellOutputUpdate, sanitizeShellOutput } from "./output-capture.ts";
|
||||
import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, type TruncationResult, truncateTail } from "./truncate.ts";
|
||||
|
||||
export interface ShellCaptureProgress {
|
||||
output: string;
|
||||
truncation: TruncationResult;
|
||||
fullOutputPath?: string;
|
||||
lastLineBytes: number;
|
||||
}
|
||||
|
||||
export interface ShellCaptureOptions extends Omit<ShellExecOptions, "capture" | "onUpdate"> {
|
||||
onChunk?: (chunk: string, getProgress: () => ShellCaptureProgress, context: Context) => void;
|
||||
/** Return shell execution failures with captured output instead of as a failed Result. */
|
||||
returnExecutionErrors?: boolean;
|
||||
}
|
||||
|
||||
export interface ShellCaptureResult extends ShellCaptureProgress {
|
||||
exitCode: number | undefined;
|
||||
cancelled: boolean;
|
||||
truncated: boolean;
|
||||
executionError?: ExecutionError;
|
||||
}
|
||||
|
||||
function progressFrom(output: ShellOutputView): ShellCaptureProgress {
|
||||
return {
|
||||
output: output.text,
|
||||
truncation: { content: output.text, ...output.truncation },
|
||||
...(output.spillPath === undefined ? {} : { fullOutputPath: output.spillPath }),
|
||||
lastLineBytes: output.lastLineBytes ?? 0,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Compatibility collector for callers that need one bounded final view.
|
||||
* Source-side capture, adaptive publication, and spilling remain owned by the
|
||||
* execution environment.
|
||||
*/
|
||||
export async function executeShellWithCapture(
|
||||
env: ExecutionEnv,
|
||||
command: string,
|
||||
options: ShellCaptureOptions | undefined,
|
||||
context: Context,
|
||||
): Promise<Result<ShellCaptureResult, ExecutionError>> {
|
||||
let output: ShellOutputView | undefined;
|
||||
const result = await env.exec(
|
||||
command,
|
||||
{
|
||||
cwd: options?.cwd,
|
||||
env: options?.env,
|
||||
inheritEnv: options?.inheritEnv,
|
||||
timeout: options?.timeout,
|
||||
capture: {
|
||||
limits: { maxBytes: DEFAULT_MAX_BYTES, maxLines: DEFAULT_MAX_LINES, retain: "tail" },
|
||||
spill: true,
|
||||
},
|
||||
onUpdate: (update, updateContext) => {
|
||||
const previous = output;
|
||||
output = applyShellOutputUpdate(output, update);
|
||||
const chunk =
|
||||
update.kind === "append" || update.kind === "slide"
|
||||
? update.text
|
||||
: update.kind === "replace" && previous === undefined
|
||||
? output.text
|
||||
: undefined;
|
||||
// A metadata-only update and a post-cap replacement contain no new
|
||||
// incremental chunk. Reporting their complete view would duplicate bytes
|
||||
// for callers that accumulate this compatibility callback.
|
||||
if (chunk) options?.onChunk?.(chunk, () => progressFrom(output!), updateContext);
|
||||
},
|
||||
},
|
||||
context,
|
||||
);
|
||||
|
||||
if (output === undefined) {
|
||||
const { content, ...truncation } = truncateTail("");
|
||||
output = { text: content, truncation };
|
||||
}
|
||||
const progress = progressFrom(output);
|
||||
if (!result.ok) {
|
||||
if (result.error.code === "aborted" || context.abortSignal?.aborted) {
|
||||
return ok({ ...progress, exitCode: undefined, cancelled: true, truncated: progress.truncation.truncated });
|
||||
}
|
||||
if (options?.returnExecutionErrors) {
|
||||
return ok({
|
||||
...progress,
|
||||
exitCode: undefined,
|
||||
cancelled: false,
|
||||
truncated: progress.truncation.truncated,
|
||||
executionError: result.error,
|
||||
});
|
||||
}
|
||||
return err(result.error);
|
||||
}
|
||||
return ok({
|
||||
...progress,
|
||||
exitCode: result.value.exitCode,
|
||||
cancelled: false,
|
||||
truncated: result.value.truncation.truncated,
|
||||
});
|
||||
}
|
||||
|
||||
export { sanitizeShellOutput as sanitizeBinaryOutput };
|
||||
350
third_party/pi/packages/agent/src/harness/utils/truncate.ts
vendored
Normal file
350
third_party/pi/packages/agent/src/harness/utils/truncate.ts
vendored
Normal file
@@ -0,0 +1,350 @@
|
||||
/**
|
||||
* Shared truncation utilities for tool outputs.
|
||||
*
|
||||
* Truncation is based on two independent limits - whichever is hit first wins:
|
||||
* - Line limit (default: 2000 lines)
|
||||
* - Byte limit (default: 50KB)
|
||||
*
|
||||
* Never returns partial lines (except bash tail truncation edge case).
|
||||
*/
|
||||
|
||||
export const DEFAULT_MAX_LINES = 2000;
|
||||
export const DEFAULT_MAX_BYTES = 50 * 1024; // 50KB
|
||||
export const GREP_MAX_LINE_LENGTH = 500; // Max chars per grep match line
|
||||
|
||||
export interface TruncationResult {
|
||||
/** The truncated content */
|
||||
content: string;
|
||||
/** Whether truncation occurred */
|
||||
truncated: boolean;
|
||||
/** Which limit was hit: "lines", "bytes", or null if not truncated */
|
||||
truncatedBy: "lines" | "bytes" | null;
|
||||
/** Total number of lines in the original content */
|
||||
totalLines: number;
|
||||
/** Total number of bytes in the original content */
|
||||
totalBytes: number;
|
||||
/** Number of complete lines in the truncated output */
|
||||
outputLines: number;
|
||||
/** Number of bytes in the truncated output */
|
||||
outputBytes: number;
|
||||
/** Whether the last line was partially truncated (only for tail truncation edge case) */
|
||||
lastLinePartial: boolean;
|
||||
/** Whether the first line exceeded the byte limit (for head truncation) */
|
||||
firstLineExceedsLimit: boolean;
|
||||
/** The max lines limit that was applied */
|
||||
maxLines: number;
|
||||
/** The max bytes limit that was applied */
|
||||
maxBytes: number;
|
||||
}
|
||||
|
||||
export interface TruncationOptions {
|
||||
/** Maximum number of lines (default: 2000) */
|
||||
maxLines?: number;
|
||||
/** Maximum number of bytes (default: 50KB) */
|
||||
maxBytes?: number;
|
||||
}
|
||||
|
||||
interface RuntimeBuffer {
|
||||
byteLength(content: string, encoding: "utf8"): number;
|
||||
}
|
||||
|
||||
const runtimeBuffer = (globalThis as { Buffer?: RuntimeBuffer }).Buffer;
|
||||
const nonAsciiPattern = /[^\x00-\x7f]/;
|
||||
|
||||
export function utf8ByteLength(content: string): number {
|
||||
if (runtimeBuffer) return runtimeBuffer.byteLength(content, "utf8");
|
||||
|
||||
const firstNonAscii = content.search(nonAsciiPattern);
|
||||
if (firstNonAscii === -1) return content.length;
|
||||
|
||||
let bytes = firstNonAscii;
|
||||
for (let i = firstNonAscii; i < content.length; i++) {
|
||||
const code = content.charCodeAt(i);
|
||||
if (code <= 0x7f) {
|
||||
bytes += 1;
|
||||
} else if (code <= 0x7ff) {
|
||||
bytes += 2;
|
||||
} else if (code >= 0xd800 && code <= 0xdbff && i + 1 < content.length) {
|
||||
const next = content.charCodeAt(i + 1);
|
||||
if (next >= 0xdc00 && next <= 0xdfff) {
|
||||
bytes += 4;
|
||||
i++;
|
||||
} else {
|
||||
bytes += 3;
|
||||
}
|
||||
} else {
|
||||
bytes += 3;
|
||||
}
|
||||
}
|
||||
return bytes;
|
||||
}
|
||||
|
||||
function splitLinesForCounting(content: string): string[] {
|
||||
if (content.length === 0) return [];
|
||||
const lines = content.split("\n");
|
||||
if (content.endsWith("\n")) lines.pop();
|
||||
return lines;
|
||||
}
|
||||
|
||||
function replaceUnpairedSurrogates(content: string): string {
|
||||
let output = "";
|
||||
for (let i = 0; i < content.length; i++) {
|
||||
const code = content.charCodeAt(i);
|
||||
if (code >= 0xd800 && code <= 0xdbff) {
|
||||
if (i + 1 < content.length) {
|
||||
const next = content.charCodeAt(i + 1);
|
||||
if (next >= 0xdc00 && next <= 0xdfff) {
|
||||
output += content[i] + content[i + 1];
|
||||
i++;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
output += "<22>";
|
||||
} else if (code >= 0xdc00 && code <= 0xdfff) {
|
||||
output += "<22>";
|
||||
} else {
|
||||
output += content[i];
|
||||
}
|
||||
}
|
||||
return output;
|
||||
}
|
||||
|
||||
/**
|
||||
* Format bytes as human-readable size.
|
||||
*/
|
||||
export function formatSize(bytes: number): string {
|
||||
if (bytes < 1024) {
|
||||
return `${bytes}B`;
|
||||
} else if (bytes < 1024 * 1024) {
|
||||
return `${(bytes / 1024).toFixed(1)}KB`;
|
||||
} else {
|
||||
return `${(bytes / (1024 * 1024)).toFixed(1)}MB`;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Truncate content from the head (keep first N lines/bytes).
|
||||
* Suitable for file reads where you want to see the beginning.
|
||||
*
|
||||
* Never returns partial lines. If first line exceeds byte limit,
|
||||
* returns empty content with firstLineExceedsLimit=true.
|
||||
*/
|
||||
export function truncateHead(content: string, options: TruncationOptions = {}): TruncationResult {
|
||||
const maxLines = options.maxLines ?? DEFAULT_MAX_LINES;
|
||||
const maxBytes = options.maxBytes ?? DEFAULT_MAX_BYTES;
|
||||
|
||||
const totalBytes = utf8ByteLength(content);
|
||||
const lines = splitLinesForCounting(content);
|
||||
const totalLines = lines.length;
|
||||
|
||||
// Check if no truncation needed
|
||||
if (totalLines <= maxLines && totalBytes <= maxBytes) {
|
||||
return {
|
||||
content,
|
||||
truncated: false,
|
||||
truncatedBy: null,
|
||||
totalLines,
|
||||
totalBytes,
|
||||
outputLines: totalLines,
|
||||
outputBytes: totalBytes,
|
||||
lastLinePartial: false,
|
||||
firstLineExceedsLimit: false,
|
||||
maxLines,
|
||||
maxBytes,
|
||||
};
|
||||
}
|
||||
|
||||
// Check if first line alone exceeds byte limit
|
||||
const firstLineBytes = utf8ByteLength(lines[0]);
|
||||
if (firstLineBytes > maxBytes) {
|
||||
return {
|
||||
content: "",
|
||||
truncated: true,
|
||||
truncatedBy: "bytes",
|
||||
totalLines,
|
||||
totalBytes,
|
||||
outputLines: 0,
|
||||
outputBytes: 0,
|
||||
lastLinePartial: false,
|
||||
firstLineExceedsLimit: true,
|
||||
maxLines,
|
||||
maxBytes,
|
||||
};
|
||||
}
|
||||
|
||||
// Collect complete lines that fit
|
||||
const outputLinesArr: string[] = [];
|
||||
let outputBytesCount = 0;
|
||||
let truncatedBy: "lines" | "bytes" = "lines";
|
||||
|
||||
for (let i = 0; i < lines.length && i < maxLines; i++) {
|
||||
const line = lines[i];
|
||||
const lineBytes = utf8ByteLength(line) + (i > 0 ? 1 : 0); // +1 for newline
|
||||
|
||||
if (outputBytesCount + lineBytes > maxBytes) {
|
||||
truncatedBy = "bytes";
|
||||
break;
|
||||
}
|
||||
|
||||
outputLinesArr.push(line);
|
||||
outputBytesCount += lineBytes;
|
||||
}
|
||||
|
||||
// If we exited due to line limit
|
||||
if (outputLinesArr.length >= maxLines && outputBytesCount <= maxBytes) {
|
||||
truncatedBy = "lines";
|
||||
}
|
||||
|
||||
const outputContent = outputLinesArr.join("\n");
|
||||
const finalOutputBytes = utf8ByteLength(outputContent);
|
||||
|
||||
return {
|
||||
content: outputContent,
|
||||
truncated: true,
|
||||
truncatedBy,
|
||||
totalLines,
|
||||
totalBytes,
|
||||
outputLines: outputLinesArr.length,
|
||||
outputBytes: finalOutputBytes,
|
||||
lastLinePartial: false,
|
||||
firstLineExceedsLimit: false,
|
||||
maxLines,
|
||||
maxBytes,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Truncate content from the tail (keep last N lines/bytes).
|
||||
* Suitable for bash output where you want to see the end (errors, final results).
|
||||
*
|
||||
* May return partial first line if the last line of original content exceeds byte limit.
|
||||
*/
|
||||
export function truncateTail(content: string, options: TruncationOptions = {}): TruncationResult {
|
||||
const maxLines = options.maxLines ?? DEFAULT_MAX_LINES;
|
||||
const maxBytes = options.maxBytes ?? DEFAULT_MAX_BYTES;
|
||||
|
||||
const totalBytes = utf8ByteLength(content);
|
||||
const lines = splitLinesForCounting(content);
|
||||
const totalLines = lines.length;
|
||||
|
||||
// Check if no truncation needed
|
||||
if (totalLines <= maxLines && totalBytes <= maxBytes) {
|
||||
return {
|
||||
content,
|
||||
truncated: false,
|
||||
truncatedBy: null,
|
||||
totalLines,
|
||||
totalBytes,
|
||||
outputLines: totalLines,
|
||||
outputBytes: totalBytes,
|
||||
lastLinePartial: false,
|
||||
firstLineExceedsLimit: false,
|
||||
maxLines,
|
||||
maxBytes,
|
||||
};
|
||||
}
|
||||
|
||||
// Work backwards from the end
|
||||
const outputLinesArr: string[] = [];
|
||||
let outputBytesCount = 0;
|
||||
let truncatedBy: "lines" | "bytes" = "lines";
|
||||
let lastLinePartial = false;
|
||||
|
||||
for (let i = lines.length - 1; i >= 0 && outputLinesArr.length < maxLines; i--) {
|
||||
const line = lines[i];
|
||||
const lineBytes = utf8ByteLength(line) + (outputLinesArr.length > 0 ? 1 : 0); // +1 for newline
|
||||
|
||||
if (outputBytesCount + lineBytes > maxBytes) {
|
||||
truncatedBy = "bytes";
|
||||
// Edge case: if we haven't added ANY lines yet and this line exceeds maxBytes,
|
||||
// take the end of the line (partial)
|
||||
if (outputLinesArr.length === 0) {
|
||||
const truncatedLine = truncateStringToBytesFromEnd(line, maxBytes);
|
||||
outputLinesArr.unshift(truncatedLine);
|
||||
outputBytesCount = utf8ByteLength(truncatedLine);
|
||||
lastLinePartial = true;
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
outputLinesArr.unshift(line);
|
||||
outputBytesCount += lineBytes;
|
||||
}
|
||||
|
||||
// If we exited due to line limit
|
||||
if (outputLinesArr.length >= maxLines && outputBytesCount <= maxBytes) {
|
||||
truncatedBy = "lines";
|
||||
}
|
||||
|
||||
const outputContent = outputLinesArr.join("\n");
|
||||
const finalOutputBytes = utf8ByteLength(outputContent);
|
||||
|
||||
return {
|
||||
content: outputContent,
|
||||
truncated: true,
|
||||
truncatedBy,
|
||||
totalLines,
|
||||
totalBytes,
|
||||
outputLines: outputLinesArr.length,
|
||||
outputBytes: finalOutputBytes,
|
||||
lastLinePartial,
|
||||
firstLineExceedsLimit: false,
|
||||
maxLines,
|
||||
maxBytes,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Truncate a string to fit within a byte limit (from the end).
|
||||
* Handles multi-byte UTF-8 characters correctly.
|
||||
*/
|
||||
function truncateStringToBytesFromEnd(str: string, maxBytes: number): string {
|
||||
if (maxBytes <= 0) return "";
|
||||
|
||||
let outputBytes = 0;
|
||||
let start = str.length;
|
||||
let needsReplacement = false;
|
||||
for (let i = str.length; i > 0; ) {
|
||||
let characterStart = i - 1;
|
||||
const code = str.charCodeAt(characterStart);
|
||||
let characterBytes: number;
|
||||
let unpairedSurrogate = false;
|
||||
if (code >= 0xdc00 && code <= 0xdfff && characterStart > 0) {
|
||||
const previous = str.charCodeAt(characterStart - 1);
|
||||
if (previous >= 0xd800 && previous <= 0xdbff) {
|
||||
characterStart--;
|
||||
characterBytes = 4;
|
||||
} else {
|
||||
characterBytes = 3;
|
||||
unpairedSurrogate = true;
|
||||
}
|
||||
} else if (code >= 0xd800 && code <= 0xdfff) {
|
||||
characterBytes = 3;
|
||||
unpairedSurrogate = true;
|
||||
} else {
|
||||
characterBytes = code <= 0x7f ? 1 : code <= 0x7ff ? 2 : 3;
|
||||
}
|
||||
if (outputBytes + characterBytes > maxBytes) break;
|
||||
outputBytes += characterBytes;
|
||||
start = characterStart;
|
||||
needsReplacement ||= unpairedSurrogate;
|
||||
i = characterStart;
|
||||
}
|
||||
|
||||
const output = str.slice(start);
|
||||
return needsReplacement ? replaceUnpairedSurrogates(output) : output;
|
||||
}
|
||||
|
||||
/**
|
||||
* Truncate a single line to max characters, adding [truncated] suffix.
|
||||
* Used for grep match lines.
|
||||
*/
|
||||
export function truncateLine(
|
||||
line: string,
|
||||
maxChars: number = GREP_MAX_LINE_LENGTH,
|
||||
): { text: string; wasTruncated: boolean } {
|
||||
if (line.length <= maxChars) {
|
||||
return { text: line, wasTruncated: false };
|
||||
}
|
||||
return { text: `${line.slice(0, maxChars)}... [truncated]`, wasTruncated: true };
|
||||
}
|
||||
35
third_party/pi/packages/agent/src/harness/utils/usage.ts
vendored
Normal file
35
third_party/pi/packages/agent/src/harness/utils/usage.ts
vendored
Normal file
@@ -0,0 +1,35 @@
|
||||
import type { Usage } from "@earendil-works/pi-ai";
|
||||
|
||||
export function emptyUsage(): Usage {
|
||||
return {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
}
|
||||
|
||||
export function addUsage(left: Usage, right: Usage): Usage {
|
||||
return {
|
||||
input: left.input + right.input,
|
||||
output: left.output + right.output,
|
||||
cacheRead: left.cacheRead + right.cacheRead,
|
||||
cacheWrite: left.cacheWrite + right.cacheWrite,
|
||||
...(left.cacheWrite1h === undefined && right.cacheWrite1h === undefined
|
||||
? {}
|
||||
: { cacheWrite1h: (left.cacheWrite1h ?? 0) + (right.cacheWrite1h ?? 0) }),
|
||||
...(left.reasoning === undefined && right.reasoning === undefined
|
||||
? {}
|
||||
: { reasoning: (left.reasoning ?? 0) + (right.reasoning ?? 0) }),
|
||||
totalTokens: left.totalTokens + right.totalTokens,
|
||||
cost: {
|
||||
input: left.cost.input + right.cost.input,
|
||||
output: left.cost.output + right.cost.output,
|
||||
cacheRead: left.cost.cacheRead + right.cost.cacheRead,
|
||||
cacheWrite: left.cost.cacheWrite + right.cost.cacheWrite,
|
||||
total: left.cost.total + right.cost.total,
|
||||
},
|
||||
};
|
||||
}
|
||||
152
third_party/pi/packages/agent/src/index.ts
vendored
Normal file
152
third_party/pi/packages/agent/src/index.ts
vendored
Normal file
@@ -0,0 +1,152 @@
|
||||
export { uuidv7 } from "@earendil-works/pi-ai";
|
||||
export type {
|
||||
AttributeValue,
|
||||
ExactTelemetryAttributes,
|
||||
InferEventAttributes,
|
||||
InferOptionalAttributes,
|
||||
InferRequiredAndOptionalAttributes,
|
||||
InferStartAttributes,
|
||||
RecordedTelemetryEvent,
|
||||
RecordedTelemetrySpan,
|
||||
SchemaTelemetrySpan,
|
||||
SpanAttributes,
|
||||
SpanAttributes as TelemetrySpanAttributes,
|
||||
SpanOptions,
|
||||
SpanStatus,
|
||||
TelemetryAttributeDefinition,
|
||||
TelemetryAttributeMetadata,
|
||||
TelemetryAttributeType,
|
||||
TelemetryContext,
|
||||
TelemetryEventAttributeDefinition,
|
||||
TelemetryEventDefinition,
|
||||
TelemetryParentDefinition,
|
||||
TelemetrySchemaDefinition,
|
||||
TelemetrySchemaSpanEndAttributes,
|
||||
TelemetrySchemaSpanEventAttributes,
|
||||
TelemetrySchemaSpanEventName,
|
||||
TelemetrySchemaSpanName,
|
||||
TelemetrySchemaSpanStartAttributes,
|
||||
TelemetrySchemaSpanUnion,
|
||||
TelemetrySpan,
|
||||
TelemetrySpanDefinition,
|
||||
TelemetryStartAttributeDefinition,
|
||||
TypedSpanStarter,
|
||||
} from "@earendil-works/pi-telemetry";
|
||||
export {
|
||||
createTypedSpanStarter,
|
||||
defineTelemetrySchema,
|
||||
InMemoryTelemetryContext,
|
||||
NOOP_TELEMETRY_CONTEXT,
|
||||
} from "@earendil-works/pi-telemetry";
|
||||
export * from "./agent.ts";
|
||||
export * from "./agent-loop.ts";
|
||||
export * from "./harness/agent-harness.ts";
|
||||
export {
|
||||
type BranchPreparation,
|
||||
type BranchSummaryDetails,
|
||||
type BranchSummaryResult,
|
||||
type CollectEntriesResult,
|
||||
collectEntriesForBranchSummary,
|
||||
type FileOperations,
|
||||
type GenerateBranchSummaryOptions,
|
||||
generateBranchSummary,
|
||||
prepareBranchEntries,
|
||||
} from "./harness/compaction/branch-summarization.ts";
|
||||
export {
|
||||
type CompactionPreparation,
|
||||
type CompactionSettings,
|
||||
type CompactResult,
|
||||
calculateContextTokens,
|
||||
compact,
|
||||
DEFAULT_COMPACTION_SETTINGS,
|
||||
estimateContextTokens,
|
||||
estimateTokens,
|
||||
findCutPoint,
|
||||
findTurnStartIndex,
|
||||
generateSummary,
|
||||
generateSummaryWithUsage,
|
||||
getLastAssistantUsage,
|
||||
prepareCompaction,
|
||||
serializeConversation,
|
||||
shouldCompact,
|
||||
} from "./harness/compaction/compaction.ts";
|
||||
export * from "./harness/context.ts";
|
||||
export * from "./harness/messages.ts";
|
||||
export * from "./harness/prompt-templates.ts";
|
||||
export * from "./harness/result.ts";
|
||||
export { type LaneSnapshotReduction, reduceLaneSnapshot } from "./harness/runtime/reducer.ts";
|
||||
export * from "./harness/session/index.ts";
|
||||
export * from "./harness/skills.ts";
|
||||
export * from "./harness/system-prompt.ts";
|
||||
export type {
|
||||
AiSpan,
|
||||
AiSpanAttributes,
|
||||
AiSpanEndAttributes,
|
||||
AiSpanEventAttributes,
|
||||
AiSpanEventName,
|
||||
AiSpanName,
|
||||
AiSpanStartAttributes,
|
||||
AiTelemetrySpan,
|
||||
HarnessSpan,
|
||||
HarnessSpanAttributes,
|
||||
HarnessSpanEndAttributes,
|
||||
HarnessSpanEventAttributes,
|
||||
HarnessSpanEventName,
|
||||
HarnessSpanName,
|
||||
HarnessSpanStartAttributes,
|
||||
HarnessTelemetrySpan,
|
||||
} from "./harness/telemetry.ts";
|
||||
export {
|
||||
AGENT_TELEMETRY_SCHEMAS,
|
||||
AI_TELEMETRY_SCHEMA,
|
||||
HARNESS_TELEMETRY_SCHEMA,
|
||||
startAiSpan,
|
||||
startHarnessSpan,
|
||||
} from "./harness/telemetry.ts";
|
||||
export * from "./harness/tools/index.ts";
|
||||
export {
|
||||
type AgentHarnessResources,
|
||||
type AgentHarnessStreamOptions,
|
||||
type AgentHarnessStreamOptionsPatch,
|
||||
type AgentHarnessTool,
|
||||
type AgentHarnessToolContextSource,
|
||||
type AgentHarnessToolInvocation,
|
||||
type AgentHarnessToolUpdateCallback,
|
||||
type AgentHarnessToolUpdateOptions,
|
||||
BranchSummaryError,
|
||||
type BranchSummaryErrorCode,
|
||||
CompactionError,
|
||||
type CompactionErrorCode,
|
||||
type ExecutionEnv,
|
||||
ExecutionError,
|
||||
type ExecutionErrorCode,
|
||||
err,
|
||||
FileError,
|
||||
type FileErrorCode,
|
||||
type FileInfo,
|
||||
type FileKind,
|
||||
type FileSystem,
|
||||
getOrThrow,
|
||||
getOrUndefined,
|
||||
ok,
|
||||
type PromptTemplate,
|
||||
type Shell,
|
||||
type ShellExecOptions,
|
||||
type ShellExecResult,
|
||||
type ShellOutputCaptureOptions,
|
||||
type ShellOutputLimits,
|
||||
type ShellOutputMetadata,
|
||||
type ShellOutputRetention,
|
||||
type ShellOutputTruncation,
|
||||
type ShellOutputUpdate,
|
||||
type ShellOutputView,
|
||||
type Skill,
|
||||
toError,
|
||||
} from "./harness/types.ts";
|
||||
export { applyShellOutputUpdate } from "./harness/utils/output-capture.ts";
|
||||
export * from "./harness/utils/shell-output.ts";
|
||||
export * from "./harness/utils/truncate.ts";
|
||||
export * from "./proxy.ts";
|
||||
export * from "./search/index.ts";
|
||||
export { setDefaultStreamFn } from "./stream-fn.ts";
|
||||
export * from "./types.ts";
|
||||
2
third_party/pi/packages/agent/src/node.ts
vendored
Normal file
2
third_party/pi/packages/agent/src/node.ts
vendored
Normal file
@@ -0,0 +1,2 @@
|
||||
export { NodeExecutionEnv } from "./harness/env/nodejs.ts";
|
||||
export * from "./index.ts";
|
||||
402
third_party/pi/packages/agent/src/proxy.ts
vendored
Normal file
402
third_party/pi/packages/agent/src/proxy.ts
vendored
Normal file
@@ -0,0 +1,402 @@
|
||||
/**
|
||||
* Proxy stream function for apps that route LLM calls through a server.
|
||||
* The server manages auth and proxies requests to LLM providers.
|
||||
*/
|
||||
|
||||
// Internal import for JSON parsing utility
|
||||
import {
|
||||
type AssistantMessage,
|
||||
type AssistantMessageEvent,
|
||||
type Context,
|
||||
EventStream,
|
||||
type Model,
|
||||
parseStreamingJson,
|
||||
type SimpleStreamOptions,
|
||||
type StopReason,
|
||||
type ToolCall,
|
||||
} from "@earendil-works/pi-ai";
|
||||
|
||||
// Create stream class matching ProxyMessageEventStream
|
||||
class ProxyMessageEventStream extends EventStream<AssistantMessageEvent, AssistantMessage> {
|
||||
constructor() {
|
||||
super(
|
||||
(event) => event.type === "done" || event.type === "error",
|
||||
(event) => {
|
||||
if (event.type === "done") return event.message;
|
||||
if (event.type === "error") return event.error;
|
||||
throw new Error("Unexpected event type");
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Proxy event types - server sends these with partial field stripped to reduce bandwidth.
|
||||
*/
|
||||
export type ProxyAssistantMessageEvent =
|
||||
| { type: "start" }
|
||||
| { type: "text_start"; contentIndex: number }
|
||||
| { type: "text_delta"; contentIndex: number; delta: string }
|
||||
| { type: "text_end"; contentIndex: number; contentSignature?: string }
|
||||
| { type: "thinking_start"; contentIndex: number }
|
||||
| { type: "thinking_delta"; contentIndex: number; delta: string }
|
||||
| { type: "thinking_end"; contentIndex: number; contentSignature?: string }
|
||||
| { type: "toolcall_start"; contentIndex: number; id: string; toolName: string }
|
||||
| { type: "toolcall_delta"; contentIndex: number; delta: string }
|
||||
| { type: "toolcall_end"; contentIndex: number; toolCall: ToolCall }
|
||||
| {
|
||||
type: "done";
|
||||
reason: Extract<StopReason, "stop" | "length" | "toolUse">;
|
||||
usage: AssistantMessage["usage"];
|
||||
providerThinkingLevel?: string;
|
||||
}
|
||||
| {
|
||||
type: "error";
|
||||
reason: Extract<StopReason, "aborted" | "error">;
|
||||
errorMessage?: string;
|
||||
usage: AssistantMessage["usage"];
|
||||
providerThinkingLevel?: string;
|
||||
};
|
||||
|
||||
type ProxySerializableStreamOptions = Pick<
|
||||
SimpleStreamOptions,
|
||||
| "temperature"
|
||||
| "samplingParams"
|
||||
| "maxTokens"
|
||||
| "reasoning"
|
||||
| "cacheRetention"
|
||||
| "sessionId"
|
||||
| "headers"
|
||||
| "metadata"
|
||||
| "transport"
|
||||
| "thinkingBudgets"
|
||||
| "maxRetryDelayMs"
|
||||
>;
|
||||
|
||||
export interface ProxyStreamOptions extends ProxySerializableStreamOptions {
|
||||
/** Local abort signal for the proxy request */
|
||||
signal?: AbortSignal;
|
||||
/** Auth token for the proxy server */
|
||||
authToken: string;
|
||||
/** Proxy server URL (e.g., "https://genai.example.com") */
|
||||
proxyUrl: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Stream function that proxies through a server instead of calling LLM providers directly.
|
||||
* 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.
|
||||
*
|
||||
* @example
|
||||
* ```typescript
|
||||
* const agent = new Agent({
|
||||
* streamFn: (model, context, options) =>
|
||||
* streamProxy(model, context, {
|
||||
* ...options,
|
||||
* authToken: await getAuthToken(),
|
||||
* proxyUrl: "https://genai.example.com",
|
||||
* }),
|
||||
* });
|
||||
* ```
|
||||
*/
|
||||
function buildProxyRequestOptions(options: ProxyStreamOptions): ProxySerializableStreamOptions {
|
||||
return {
|
||||
temperature: options.temperature,
|
||||
samplingParams: options.samplingParams,
|
||||
maxTokens: options.maxTokens,
|
||||
reasoning: options.reasoning,
|
||||
cacheRetention: options.cacheRetention,
|
||||
sessionId: options.sessionId,
|
||||
headers: options.headers,
|
||||
metadata: options.metadata,
|
||||
transport: options.transport,
|
||||
thinkingBudgets: options.thinkingBudgets,
|
||||
maxRetryDelayMs: options.maxRetryDelayMs,
|
||||
};
|
||||
}
|
||||
|
||||
export function streamProxy(model: Model<any>, context: Context, options: ProxyStreamOptions): ProxyMessageEventStream {
|
||||
const stream = new ProxyMessageEventStream();
|
||||
|
||||
(async () => {
|
||||
// Initialize the partial message that we'll build up from events
|
||||
const partial: AssistantMessage = {
|
||||
role: "assistant",
|
||||
stopReason: "pending",
|
||||
content: [],
|
||||
api: model.api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
let reader: ReadableStreamDefaultReader<Uint8Array> | undefined;
|
||||
|
||||
const abortHandler = () => {
|
||||
if (reader) {
|
||||
reader.cancel("Request aborted by user").catch(() => {});
|
||||
}
|
||||
};
|
||||
|
||||
if (options.signal) {
|
||||
options.signal.addEventListener("abort", abortHandler);
|
||||
}
|
||||
|
||||
try {
|
||||
const response = await fetch(`${options.proxyUrl}/api/stream`, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
Authorization: `Bearer ${options.authToken}`,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
body: JSON.stringify({
|
||||
model,
|
||||
context,
|
||||
options: buildProxyRequestOptions(options),
|
||||
}),
|
||||
signal: options.signal,
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
let errorMessage = `Proxy error: ${response.status} ${response.statusText}`;
|
||||
try {
|
||||
const errorData = (await response.json()) as { error?: string };
|
||||
if (errorData.error) {
|
||||
errorMessage = `Proxy error: ${errorData.error}`;
|
||||
}
|
||||
} catch {
|
||||
// Couldn't parse error response
|
||||
}
|
||||
throw new Error(errorMessage);
|
||||
}
|
||||
|
||||
reader = response.body!.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = "";
|
||||
let sawTerminalEvent = false;
|
||||
|
||||
const processLine = (line: string): void => {
|
||||
if (!line.startsWith("data: ")) return;
|
||||
const data = line.slice(6).trim();
|
||||
if (!data) return;
|
||||
const proxyEvent = JSON.parse(data) as ProxyAssistantMessageEvent;
|
||||
const event = processProxyEvent(proxyEvent, partial);
|
||||
if (event) {
|
||||
if (event.type === "done" || event.type === "error") sawTerminalEvent = true;
|
||||
stream.push(event);
|
||||
}
|
||||
};
|
||||
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
|
||||
if (options.signal?.aborted) {
|
||||
throw new Error("Request aborted by user");
|
||||
}
|
||||
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
const lines = buffer.split("\n");
|
||||
buffer = lines.pop() || "";
|
||||
|
||||
for (const line of lines) {
|
||||
processLine(line);
|
||||
}
|
||||
}
|
||||
|
||||
if (options.signal?.aborted) {
|
||||
throw new Error("Request aborted by user");
|
||||
}
|
||||
|
||||
// The final event may not be newline-terminated; flush the decoder and
|
||||
// process whatever is left in the buffer.
|
||||
buffer += decoder.decode();
|
||||
if (buffer) {
|
||||
processLine(buffer);
|
||||
}
|
||||
|
||||
if (!sawTerminalEvent) {
|
||||
// A clean EOF without a done/error event means the server dropped the
|
||||
// response mid-stream. Surface it as an error instead of leaving
|
||||
// consumers waiting on a result that never arrives.
|
||||
partial.stopReason = "error";
|
||||
partial.errorMessage = "Connection closed by proxy server before the response completed";
|
||||
stream.push({
|
||||
type: "error",
|
||||
reason: "error",
|
||||
error: partial,
|
||||
});
|
||||
}
|
||||
|
||||
stream.end();
|
||||
} catch (error) {
|
||||
const errorMessage = error instanceof Error ? error.message : String(error);
|
||||
const reason = options.signal?.aborted ? "aborted" : "error";
|
||||
partial.stopReason = reason;
|
||||
partial.errorMessage = errorMessage;
|
||||
stream.push({
|
||||
type: "error",
|
||||
reason,
|
||||
error: partial,
|
||||
});
|
||||
stream.end();
|
||||
} finally {
|
||||
if (options.signal) {
|
||||
options.signal.removeEventListener("abort", abortHandler);
|
||||
}
|
||||
}
|
||||
})();
|
||||
|
||||
return stream;
|
||||
}
|
||||
|
||||
/**
|
||||
* Process a proxy event and update the partial message.
|
||||
*/
|
||||
function processProxyEvent(
|
||||
proxyEvent: ProxyAssistantMessageEvent,
|
||||
partial: AssistantMessage,
|
||||
): AssistantMessageEvent | undefined {
|
||||
switch (proxyEvent.type) {
|
||||
case "start":
|
||||
return { type: "start", partial };
|
||||
|
||||
case "text_start":
|
||||
partial.content[proxyEvent.contentIndex] = { type: "text", text: "" };
|
||||
return { type: "text_start", contentIndex: proxyEvent.contentIndex, partial };
|
||||
|
||||
case "text_delta": {
|
||||
const content = partial.content[proxyEvent.contentIndex];
|
||||
if (content?.type === "text") {
|
||||
content.text += proxyEvent.delta;
|
||||
return {
|
||||
type: "text_delta",
|
||||
contentIndex: proxyEvent.contentIndex,
|
||||
delta: proxyEvent.delta,
|
||||
partial,
|
||||
};
|
||||
}
|
||||
throw new Error("Received text_delta for non-text content");
|
||||
}
|
||||
|
||||
case "text_end": {
|
||||
const content = partial.content[proxyEvent.contentIndex];
|
||||
if (content?.type === "text") {
|
||||
content.textSignature = proxyEvent.contentSignature;
|
||||
return {
|
||||
type: "text_end",
|
||||
contentIndex: proxyEvent.contentIndex,
|
||||
content: content.text,
|
||||
partial,
|
||||
};
|
||||
}
|
||||
throw new Error("Received text_end for non-text content");
|
||||
}
|
||||
|
||||
case "thinking_start":
|
||||
partial.content[proxyEvent.contentIndex] = { type: "thinking", thinking: "" };
|
||||
return { type: "thinking_start", contentIndex: proxyEvent.contentIndex, partial };
|
||||
|
||||
case "thinking_delta": {
|
||||
const content = partial.content[proxyEvent.contentIndex];
|
||||
if (content?.type === "thinking") {
|
||||
content.thinking += proxyEvent.delta;
|
||||
return {
|
||||
type: "thinking_delta",
|
||||
contentIndex: proxyEvent.contentIndex,
|
||||
delta: proxyEvent.delta,
|
||||
partial,
|
||||
};
|
||||
}
|
||||
throw new Error("Received thinking_delta for non-thinking content");
|
||||
}
|
||||
|
||||
case "thinking_end": {
|
||||
const content = partial.content[proxyEvent.contentIndex];
|
||||
if (content?.type === "thinking") {
|
||||
content.thinkingSignature = proxyEvent.contentSignature;
|
||||
return {
|
||||
type: "thinking_end",
|
||||
contentIndex: proxyEvent.contentIndex,
|
||||
content: content.thinking,
|
||||
partial,
|
||||
};
|
||||
}
|
||||
throw new Error("Received thinking_end for non-thinking content");
|
||||
}
|
||||
|
||||
case "toolcall_start":
|
||||
partial.content[proxyEvent.contentIndex] = {
|
||||
type: "toolCall",
|
||||
id: proxyEvent.id,
|
||||
name: proxyEvent.toolName,
|
||||
arguments: {},
|
||||
partialJson: "",
|
||||
} satisfies ToolCall & { partialJson: string } as ToolCall;
|
||||
return { type: "toolcall_start", contentIndex: proxyEvent.contentIndex, partial };
|
||||
|
||||
case "toolcall_delta": {
|
||||
const content = partial.content[proxyEvent.contentIndex];
|
||||
if (content?.type === "toolCall") {
|
||||
(content as any).partialJson += proxyEvent.delta;
|
||||
content.arguments = parseStreamingJson((content as any).partialJson) || {};
|
||||
partial.content[proxyEvent.contentIndex] = { ...content }; // Trigger reactivity
|
||||
return {
|
||||
type: "toolcall_delta",
|
||||
contentIndex: proxyEvent.contentIndex,
|
||||
delta: proxyEvent.delta,
|
||||
partial,
|
||||
};
|
||||
}
|
||||
throw new Error("Received toolcall_delta for non-toolCall content");
|
||||
}
|
||||
|
||||
case "toolcall_end": {
|
||||
const content = partial.content[proxyEvent.contentIndex];
|
||||
if (content?.type === "toolCall") {
|
||||
Object.assign(content, proxyEvent.toolCall);
|
||||
delete (content as any).partialJson;
|
||||
return {
|
||||
type: "toolcall_end",
|
||||
contentIndex: proxyEvent.contentIndex,
|
||||
toolCall: content,
|
||||
partial,
|
||||
};
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
case "done":
|
||||
partial.stopReason = proxyEvent.reason;
|
||||
partial.usage = proxyEvent.usage;
|
||||
if (proxyEvent.providerThinkingLevel !== undefined) {
|
||||
partial.providerThinkingLevel = proxyEvent.providerThinkingLevel;
|
||||
}
|
||||
return { type: "done", reason: proxyEvent.reason, message: partial };
|
||||
|
||||
case "error":
|
||||
partial.stopReason = proxyEvent.reason;
|
||||
partial.errorMessage = proxyEvent.errorMessage;
|
||||
partial.usage = proxyEvent.usage;
|
||||
if (proxyEvent.providerThinkingLevel !== undefined) {
|
||||
partial.providerThinkingLevel = proxyEvent.providerThinkingLevel;
|
||||
}
|
||||
return { type: "error", reason: proxyEvent.reason, error: partial };
|
||||
|
||||
default: {
|
||||
const _exhaustiveCheck: never = proxyEvent;
|
||||
console.warn(`Unhandled proxy event type: ${(proxyEvent as any).type}`);
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
}
|
||||
27
third_party/pi/packages/agent/src/search/index.ts
vendored
Normal file
27
third_party/pi/packages/agent/src/search/index.ts
vendored
Normal file
@@ -0,0 +1,27 @@
|
||||
export interface SearchQuery {
|
||||
text: string;
|
||||
limit?: number;
|
||||
}
|
||||
|
||||
export interface SessionSearchHit {
|
||||
sessionId: string;
|
||||
score?: number;
|
||||
top?: { entryId: string; snippet?: string; timestamp: number };
|
||||
}
|
||||
|
||||
export interface EntrySearchHit {
|
||||
sessionId: string;
|
||||
entryId: string;
|
||||
timestamp: number;
|
||||
snippet?: string;
|
||||
score?: number;
|
||||
}
|
||||
|
||||
export interface SessionSearchService {
|
||||
searchSessions(query: SearchQuery): Promise<SessionSearchHit[]>;
|
||||
searchEntries?(query: SearchQuery): Promise<EntrySearchHit[]>;
|
||||
sync(): Promise<void>;
|
||||
notify(sessionId: string): void;
|
||||
remove(sessionId: string): Promise<void>;
|
||||
close(): Promise<void>;
|
||||
}
|
||||
20
third_party/pi/packages/agent/src/stream-fn.ts
vendored
Normal file
20
third_party/pi/packages/agent/src/stream-fn.ts
vendored
Normal 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;
|
||||
}
|
||||
446
third_party/pi/packages/agent/src/types.ts
vendored
Normal file
446
third_party/pi/packages/agent/src/types.ts
vendored
Normal file
@@ -0,0 +1,446 @@
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
AssistantMessageEvent,
|
||||
AssistantMessageEventStream,
|
||||
Context,
|
||||
ImageContent,
|
||||
Message,
|
||||
Model,
|
||||
SimpleStreamOptions,
|
||||
TextContent,
|
||||
Tool,
|
||||
ToolResultMessage,
|
||||
Usage,
|
||||
} from "@earendil-works/pi-ai";
|
||||
import type { Static, TSchema } from "typebox";
|
||||
|
||||
/**
|
||||
* Stream function used by the agent loop. `Models.streamSimple` satisfies
|
||||
* this shape.
|
||||
*
|
||||
* Contract:
|
||||
* - Must not throw or return a rejected promise for request/model/runtime failures.
|
||||
* - Must return an AssistantMessageEventStream.
|
||||
* - Failures must be encoded in the returned stream via protocol events and a
|
||||
* final AssistantMessage with stopReason "error" or "aborted" and errorMessage.
|
||||
*/
|
||||
export type StreamFn = (
|
||||
model: Model<Api>,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
) => AssistantMessageEventStream | Promise<AssistantMessageEventStream>;
|
||||
|
||||
/**
|
||||
* Configuration for how tool calls from a single assistant message are executed.
|
||||
*
|
||||
* - "sequential": each tool call is prepared, executed, and finalized before the next one starts.
|
||||
* - "parallel": tool calls are prepared sequentially, then allowed tools execute concurrently.
|
||||
* `tool_execution_end` is emitted in tool completion order after each tool is finalized,
|
||||
* while tool-result message artifacts are emitted later in assistant source order.
|
||||
*/
|
||||
export type ToolExecutionMode = "sequential" | "parallel";
|
||||
|
||||
/**
|
||||
* Controls how many queued user messages are injected when the agent loop reaches a queue drain point.
|
||||
*
|
||||
* - "all": drain and inject every queued message at that point.
|
||||
* - "one-at-a-time": drain and inject only the oldest queued message, leaving the rest queued for later drain points.
|
||||
*/
|
||||
export type QueueMode = "all" | "one-at-a-time";
|
||||
|
||||
/** A single tool call content block emitted by an assistant message. */
|
||||
export type AgentToolCall = Extract<AssistantMessage["content"][number], { type: "toolCall" }>;
|
||||
|
||||
/**
|
||||
* Result returned from `beforeToolCall`.
|
||||
*
|
||||
* Returning `{ block: true }` prevents the tool from executing. The loop emits an error tool result instead.
|
||||
* `reason` becomes the text shown in that error result. If omitted, a default blocked message is used.
|
||||
*/
|
||||
export interface BeforeToolCallResult {
|
||||
block?: boolean;
|
||||
reason?: string;
|
||||
/**
|
||||
* Hint that the agent should stop after the current tool batch when this call is blocked.
|
||||
* Early termination only happens when every finalized tool result in the batch sets this to true.
|
||||
*/
|
||||
terminate?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Partial override returned from `afterToolCall`.
|
||||
*
|
||||
* Merge semantics are field-by-field:
|
||||
* - `content`: if provided, replaces the tool result content array in full
|
||||
* - `details`: if provided, replaces the tool result details value in full
|
||||
* - `isError`: if provided, replaces the tool result error flag
|
||||
* - `usage`: if provided, replaces the tool result usage
|
||||
* - `terminate`: if provided, replaces the early-termination hint
|
||||
*
|
||||
* Omitted fields keep the original executed tool result values.
|
||||
* There is no deep merge for `content`, `details`, or `usage`.
|
||||
*/
|
||||
export interface AfterToolCallResult {
|
||||
content?: (TextContent | ImageContent)[];
|
||||
details?: unknown;
|
||||
isError?: boolean;
|
||||
/** Usage from the final tool execution itself, if available. Not used for main LLM context accounting. */
|
||||
usage?: Usage;
|
||||
/**
|
||||
* Hint that the agent should stop after the current tool batch.
|
||||
* Early termination only happens when every finalized tool result in the batch sets this to true.
|
||||
*/
|
||||
terminate?: boolean;
|
||||
}
|
||||
|
||||
/** Context passed to `beforeToolCall`. */
|
||||
export interface BeforeToolCallContext {
|
||||
/** The assistant message that requested the tool call. */
|
||||
assistantMessage: AssistantMessage;
|
||||
/** The raw tool call block from `assistantMessage.content`. */
|
||||
toolCall: AgentToolCall;
|
||||
/** Validated tool arguments for the target tool schema. */
|
||||
args: unknown;
|
||||
/** Current agent context at the time the tool call is prepared. */
|
||||
context: AgentContext;
|
||||
}
|
||||
|
||||
/** Context passed to `afterToolCall`. */
|
||||
export interface AfterToolCallContext {
|
||||
/** The assistant message that requested the tool call. */
|
||||
assistantMessage: AssistantMessage;
|
||||
/** The raw tool call block from `assistantMessage.content`. */
|
||||
toolCall: AgentToolCall;
|
||||
/** Validated tool arguments for the target tool schema. */
|
||||
args: unknown;
|
||||
/** The executed tool result before any `afterToolCall` overrides are applied. */
|
||||
result: AgentToolResult<any>;
|
||||
/** Whether the executed tool result is currently treated as an error. */
|
||||
isError: boolean;
|
||||
/** Current agent context at the time the tool call is finalized. */
|
||||
context: AgentContext;
|
||||
}
|
||||
|
||||
/** Context passed to `shouldStopAfterTurn`. */
|
||||
export interface ShouldStopAfterTurnContext {
|
||||
/** The assistant message that completed the turn. */
|
||||
message: AssistantMessage;
|
||||
/** Tool result messages passed to the preceding `turn_end` event. */
|
||||
toolResults: ToolResultMessage[];
|
||||
/** Current agent context after the turn's assistant message and tool results have been appended. */
|
||||
context: AgentContext;
|
||||
/** Messages that this loop invocation will return if it exits at this point. Prompt runs include the initial prompt messages; continuation runs do not include pre-existing context messages. */
|
||||
newMessages: AgentMessage[];
|
||||
}
|
||||
|
||||
/** Replacement runtime state used by the agent loop before starting another provider request. */
|
||||
export interface AgentLoopTurnUpdate {
|
||||
/** Context for the next provider request. */
|
||||
context?: AgentContext;
|
||||
/** Model for the next provider request. */
|
||||
model?: Model<any>;
|
||||
/** Thinking level for the next provider request. */
|
||||
thinkingLevel?: ThinkingLevel;
|
||||
}
|
||||
|
||||
export interface PrepareNextTurnContext extends ShouldStopAfterTurnContext {}
|
||||
|
||||
export interface AgentLoopConfig extends SimpleStreamOptions {
|
||||
model: Model<any>;
|
||||
|
||||
/**
|
||||
* Converts AgentMessage[] to LLM-compatible Message[] before each LLM call.
|
||||
*
|
||||
* Each AgentMessage must be converted to a UserMessage, AssistantMessage, or ToolResultMessage
|
||||
* that the LLM can understand. AgentMessages that cannot be converted (e.g., UI-only notifications,
|
||||
* status messages) should be filtered out.
|
||||
*
|
||||
* Contract: must not throw or reject. Return a safe fallback value instead.
|
||||
* Throwing interrupts the low-level agent loop without producing a normal event sequence.
|
||||
*
|
||||
* @example
|
||||
* ```typescript
|
||||
* convertToLlm: (messages) => messages.flatMap(m => {
|
||||
* if (m.role === "custom") {
|
||||
* // Convert custom message to user message
|
||||
* return [{ role: "user", content: m.content, timestamp: m.timestamp }];
|
||||
* }
|
||||
* if (m.role === "notification") {
|
||||
* // Filter out UI-only messages
|
||||
* return [];
|
||||
* }
|
||||
* // Pass through standard LLM messages
|
||||
* return [m];
|
||||
* })
|
||||
* ```
|
||||
*/
|
||||
convertToLlm: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
|
||||
|
||||
/**
|
||||
* Optional transform applied to the context before `convertToLlm`.
|
||||
*
|
||||
* Use this for operations that work at the AgentMessage level:
|
||||
* - Context window management (pruning old messages)
|
||||
* - Injecting context from external sources
|
||||
*
|
||||
* Contract: must not throw or reject. Return the original messages or another
|
||||
* safe fallback value instead.
|
||||
*
|
||||
* @example
|
||||
* ```typescript
|
||||
* transformContext: async (messages) => {
|
||||
* if (estimateTokens(messages) > MAX_TOKENS) {
|
||||
* return pruneOldMessages(messages);
|
||||
* }
|
||||
* return messages;
|
||||
* }
|
||||
* ```
|
||||
*/
|
||||
transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
|
||||
|
||||
/**
|
||||
* Resolves an API key dynamically for each LLM call.
|
||||
*
|
||||
* Useful for short-lived OAuth tokens (e.g., GitHub Copilot) that may expire
|
||||
* during long-running tool execution phases.
|
||||
*
|
||||
* Contract: must not throw or reject. Return undefined when no key is available.
|
||||
*/
|
||||
getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
|
||||
|
||||
/**
|
||||
* Called after each turn fully completes and `turn_end` has been emitted.
|
||||
*
|
||||
* If it returns true, the loop emits `agent_end` and exits before polling steering or follow-up queues,
|
||||
* without starting another LLM call. The current assistant response and any tool executions finish normally.
|
||||
* This callback sees the completed-turn context and runs before `prepareNextTurn`.
|
||||
*
|
||||
* Use this to request a graceful stop after the current turn, e.g. before context gets too full.
|
||||
*
|
||||
* Contract: must not throw or reject. Throwing interrupts the low-level agent loop without producing a normal event sequence.
|
||||
*/
|
||||
shouldStopAfterTurn?: (context: ShouldStopAfterTurnContext) => boolean | Promise<boolean>;
|
||||
|
||||
/**
|
||||
* Called after `turn_end` when the loop will continue, immediately before the next turn starts.
|
||||
* Return replacement context/model/thinking state to affect that turn.
|
||||
* Return undefined to keep using the current context/config.
|
||||
*/
|
||||
prepareNextTurn?: (
|
||||
context: PrepareNextTurnContext,
|
||||
) => AgentLoopTurnUpdate | undefined | Promise<AgentLoopTurnUpdate | undefined>;
|
||||
|
||||
/**
|
||||
* Returns steering messages to inject into the conversation mid-run.
|
||||
*
|
||||
* Called after the current assistant turn finishes executing its tool calls, unless `shouldStopAfterTurn` exits first.
|
||||
* If messages are returned, they are added to the context before the next LLM call.
|
||||
* Tool calls from the current assistant message are not skipped.
|
||||
*
|
||||
* Use this for "steering" the agent while it's working.
|
||||
*
|
||||
* Contract: must not throw or reject. Return [] when no steering messages are available.
|
||||
*/
|
||||
getSteeringMessages?: () => Promise<AgentMessage[]>;
|
||||
|
||||
/**
|
||||
* Returns follow-up messages to process after the agent would otherwise stop.
|
||||
*
|
||||
* Called when the agent has no more tool calls and no steering messages.
|
||||
* If messages are returned, they're added to the context and the agent
|
||||
* continues with another turn.
|
||||
*
|
||||
* Use this for follow-up messages that should wait until the agent finishes.
|
||||
*
|
||||
* Contract: must not throw or reject. Return [] when no follow-up messages are available.
|
||||
*/
|
||||
getFollowUpMessages?: () => Promise<AgentMessage[]>;
|
||||
|
||||
/**
|
||||
* Tool execution mode.
|
||||
* - "sequential": execute tool calls one by one
|
||||
* - "parallel": preflight tool calls sequentially, then execute allowed tools concurrently;
|
||||
* emit `tool_execution_end` in tool completion order after each tool is finalized,
|
||||
* then emit tool-result message artifacts later in assistant source order
|
||||
*
|
||||
* Default: "parallel"
|
||||
*/
|
||||
toolExecution?: ToolExecutionMode;
|
||||
|
||||
/**
|
||||
* Called before a tool is executed, after arguments have been validated.
|
||||
*
|
||||
* Return `{ block: true }` to prevent execution. The loop emits an error tool result instead.
|
||||
* A blocked result can also set `terminate: true` to participate in the batch early-termination rule.
|
||||
* The hook receives the agent abort signal and is responsible for honoring it.
|
||||
*/
|
||||
beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>;
|
||||
|
||||
/**
|
||||
* Called after a tool finishes executing, before `tool_execution_end` and tool-result message events are emitted.
|
||||
*
|
||||
* Return an `AfterToolCallResult` to override parts of the executed tool result:
|
||||
* - `content` replaces the full content array
|
||||
* - `details` replaces the full details payload
|
||||
* - `isError` replaces the error flag
|
||||
* - `usage` replaces the tool result usage
|
||||
* - `terminate` replaces the early-termination hint
|
||||
*
|
||||
* Any omitted fields keep their original values. No deep merge is performed.
|
||||
* The hook receives the agent abort signal and is responsible for honoring it.
|
||||
*/
|
||||
afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Thinking/reasoning level for models that support it.
|
||||
* Note: "xhigh" and "max" are only supported by selected model families. Use model
|
||||
* thinking-level metadata from @earendil-works/pi-ai to detect support for a concrete model.
|
||||
*/
|
||||
export type ThinkingLevel = "off" | "minimal" | "low" | "medium" | "high" | "xhigh" | "max";
|
||||
|
||||
/**
|
||||
* Extensible interface for custom app messages.
|
||||
* Apps can extend via declaration merging:
|
||||
*
|
||||
* @example
|
||||
* ```typescript
|
||||
* declare module "@mariozechner/agent" {
|
||||
* interface CustomAgentMessages {
|
||||
* artifact: ArtifactMessage;
|
||||
* notification: NotificationMessage;
|
||||
* }
|
||||
* }
|
||||
* ```
|
||||
*/
|
||||
export interface CustomAgentMessages {
|
||||
// Empty by default - apps extend via declaration merging
|
||||
}
|
||||
|
||||
/**
|
||||
* AgentMessage: Union of LLM messages + custom messages.
|
||||
* This abstraction allows apps to add custom message types while maintaining
|
||||
* type safety and compatibility with the base LLM messages.
|
||||
*/
|
||||
export type AgentMessage = Message | CustomAgentMessages[keyof CustomAgentMessages];
|
||||
|
||||
/**
|
||||
* Public agent state.
|
||||
*
|
||||
* `tools` and `messages` use accessor properties so implementations can copy
|
||||
* assigned arrays before storing them.
|
||||
*/
|
||||
export interface AgentState {
|
||||
/** System prompt sent with each model request. */
|
||||
systemPrompt: string;
|
||||
/** Active model used for future turns. */
|
||||
model: Model<any>;
|
||||
/** Requested reasoning level for future turns. */
|
||||
thinkingLevel: ThinkingLevel;
|
||||
/** Available tools. Assigning a new array copies the top-level array. */
|
||||
set tools(tools: AgentTool<any>[]);
|
||||
get tools(): AgentTool<any>[];
|
||||
/** Conversation transcript. Assigning a new array copies the top-level array. */
|
||||
set messages(messages: AgentMessage[]);
|
||||
get messages(): AgentMessage[];
|
||||
/**
|
||||
* True while the agent is processing a prompt or continuation.
|
||||
*
|
||||
* This remains true until awaited `agent_end` listeners settle.
|
||||
*/
|
||||
readonly isStreaming: boolean;
|
||||
/** Partial assistant message for the current streamed response, if any. */
|
||||
readonly streamingMessage?: AgentMessage;
|
||||
/** Tool call ids currently executing. */
|
||||
readonly pendingToolCalls: ReadonlySet<string>;
|
||||
/** Error message from the most recent failed or aborted assistant turn, if any. */
|
||||
readonly errorMessage?: string;
|
||||
}
|
||||
|
||||
/** Final or partial result produced by a tool. */
|
||||
export interface AgentToolResult<T> {
|
||||
/** Text or image content returned to the model. */
|
||||
content: (TextContent | ImageContent)[];
|
||||
/** Arbitrary structured details for logs or UI rendering. */
|
||||
details: T;
|
||||
/** Usage from the final tool execution itself, if available. Not used for main LLM context accounting. */
|
||||
usage?: Usage;
|
||||
/** Names of tools introduced by this result and available from this transcript point onward. */
|
||||
addedToolNames?: string[];
|
||||
/**
|
||||
* Hint that the agent should stop after the current tool batch.
|
||||
* Early termination only happens when every finalized tool result in the batch sets this to true.
|
||||
*/
|
||||
terminate?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Callback used by tools to stream partial execution updates.
|
||||
*
|
||||
* The callback is scoped to the current `execute()` invocation. Calls made after
|
||||
* the tool promise settles are ignored.
|
||||
*/
|
||||
export type AgentToolUpdateCallback<T = any> = (partialResult: AgentToolResult<T>) => void;
|
||||
|
||||
/** Tool definition used by the agent runtime. */
|
||||
export interface AgentTool<TParameters extends TSchema = TSchema, TDetails = any> extends Tool<TParameters> {
|
||||
/** Human-readable label for UI display. */
|
||||
label: string;
|
||||
/**
|
||||
* Optional compatibility shim for raw tool-call arguments before schema validation.
|
||||
* Must return an object that matches `TParameters`.
|
||||
*/
|
||||
prepareArguments?: (args: unknown) => Static<TParameters>;
|
||||
/** Execute the tool call. Throw on failure instead of encoding errors in `content`. */
|
||||
execute: (
|
||||
toolCallId: string,
|
||||
params: Static<TParameters>,
|
||||
signal?: AbortSignal,
|
||||
onUpdate?: AgentToolUpdateCallback<TDetails>,
|
||||
) => Promise<AgentToolResult<TDetails>>;
|
||||
/** Recovery policy for an effect whose durable intent exists but whose outcome is unknown. */
|
||||
replay?: "never" | "safe";
|
||||
/**
|
||||
* Per-tool execution mode override.
|
||||
* - "sequential": this tool must execute one at a time with other tool calls.
|
||||
* - "parallel": this tool can execute concurrently with other tool calls.
|
||||
*
|
||||
* If omitted, the default execution mode applies.
|
||||
*/
|
||||
executionMode?: ToolExecutionMode;
|
||||
}
|
||||
|
||||
/** Context snapshot passed into the low-level agent loop. */
|
||||
export interface AgentContext {
|
||||
/** System prompt included with the request. */
|
||||
systemPrompt: string;
|
||||
/** Transcript visible to the model. */
|
||||
messages: AgentMessage[];
|
||||
/** Tools available for this run. */
|
||||
tools?: AgentTool<any>[];
|
||||
}
|
||||
|
||||
/**
|
||||
* Events emitted by the Agent for UI updates.
|
||||
*
|
||||
* `agent_end` is the last event emitted for a run, but awaited `Agent.subscribe()`
|
||||
* listeners for that event are still part of run settlement. The agent becomes
|
||||
* idle only after those listeners finish.
|
||||
*/
|
||||
export type AgentEvent =
|
||||
// Agent lifecycle
|
||||
| { type: "agent_start" }
|
||||
| { type: "agent_end"; messages: AgentMessage[] }
|
||||
// Turn lifecycle - a turn is one assistant response + any tool calls/results
|
||||
| { type: "turn_start" }
|
||||
| { type: "turn_end"; message: AgentMessage; toolResults: ToolResultMessage[] }
|
||||
// Message lifecycle - emitted for user, assistant, and toolResult messages
|
||||
| { type: "message_start"; message: AgentMessage }
|
||||
// Only emitted for assistant messages during streaming
|
||||
| { type: "message_update"; message: AgentMessage; assistantMessageEvent: AssistantMessageEvent }
|
||||
| { type: "message_end"; message: AgentMessage }
|
||||
// Tool execution lifecycle
|
||||
| { type: "tool_execution_start"; toolCallId: string; toolName: string; args: any }
|
||||
| { type: "tool_execution_update"; toolCallId: string; toolName: string; args: any; partialResult: any }
|
||||
| { type: "tool_execution_end"; toolCallId: string; toolName: string; result: any; isError: boolean };
|
||||
Reference in New Issue
Block a user