fix(adapters): persist tool-call budget across worker handoffs

This commit is contained in:
Ingo Wagner 2026-09-22 06:22:24 +00:00
parent 6e33307c7f
commit beece5a383
5 changed files with 707 additions and 98 deletions

View File

@ -402,6 +402,19 @@ export interface AgentRunRequest {
onToolCompleted?: (completion: AgentToolCompletion) => Promise<void> | void; onToolCompleted?: (completion: AgentToolCompletion) => Promise<void> | void;
/** Atomically claim durable user steering at the runtime's next safe turn boundary. */ /** Atomically claim durable user steering at the runtime's next safe turn boundary. */
claimSteering?: (seenIds: string[]) => Promise<AgentSteeringMessage[]>; claimSteering?: (seenIds: string[]) => Promise<AgentSteeringMessage[]>;
/** Optional durable store for the per-turn tool-call fuse across worker handoffs. */
toolCallBudgetStore?: ToolCallBudgetStore;
}
export interface ToolCallBudgetSnapshot {
count: number;
exceeded: boolean;
}
export interface ToolCallBudgetStore {
load(): Promise<ToolCallBudgetSnapshot | undefined>;
persist(snapshot: ToolCallBudgetSnapshot): Promise<void>;
clear(): Promise<void>;
} }
export interface ScriptedTurn { export interface ScriptedTurn {
@ -587,11 +600,6 @@ export interface MessagingInboundMessage {
channelName: string | null; channelName: string | null;
/** Group roster addresses when the platform reports them; often empty. */ /** Group roster addresses when the platform reports them; often empty. */
participants: string[]; participants: string[];
/**
* True only when {@link participants} is a complete group roster from the
* provider. Partial rosters must not trigger membership sweeps.
*/
participantsAuthoritative?: boolean;
content: string; content: string;
mediaUrl: string | null; mediaUrl: string | null;
/** Team-room workspace/team id when the platform reports one (Slack team_id, …). */ /** Team-room workspace/team id when the platform reports one (Slack team_id, …). */

View File

@ -263,6 +263,7 @@ import {
secretValuesToRedact, secretValuesToRedact,
serializeModelSecret, serializeModelSecret,
} from "./pi-oauth.js"; } from "./pi-oauth.js";
import { maxToolCallsPerTurn } from "./pi-model-config.js";
import { import {
assertPlotDataWithinLimits, assertPlotDataWithinLimits,
PLOT_TOOL_GUIDE, PLOT_TOOL_GUIDE,
@ -273,6 +274,7 @@ import {
searchChartCatalog, searchChartCatalog,
} from "./plot-tool.js"; } from "./plot-tool.js";
import type { RemoteTransportDependencies } from "./remote-mcp.js"; import type { RemoteTransportDependencies } from "./remote-mcp.js";
import { createRunToolCallBudgetStore } from "./tool-call-budget-store.js";
import { loadReplyContext, messageToAgentHistoryText } from "./reply-context.js"; import { loadReplyContext, messageToAgentHistoryText } from "./reply-context.js";
import { import {
commitConsumedRunSecret, commitConsumedRunSecret,
@ -3717,6 +3719,10 @@ export function createRunExecutor(deps: ExecutorDeps) {
completion, completion,
runSecrets, runSecrets,
), ),
toolCallBudgetStore:
maxToolCallsPerTurn() > 0
? createRunToolCallBudgetStore(deps.prisma, runId)
: undefined,
claimSteering: scripted claimSteering: scripted
? undefined ? undefined
: async (seenIds) => { : async (seenIds) => {
@ -4169,7 +4175,7 @@ export function createRunExecutor(deps: ExecutorDeps) {
const workshopFallback = workshopFallbackAfterBlockedHandoff( const workshopFallback = workshopFallbackAfterBlockedHandoff(
proseRecovery.sanitizedText, proseRecovery.sanitizedText,
proseRecovery.blockedUpward, proseRecovery.blockedUpward,
proseRecovery.handoffFailed, proseRecovery.proseToolCalls,
); );
if (proseRecovery.proseToolCalls.length > 0) { if (proseRecovery.proseToolCalls.length > 0) {
assembled = proseRecovery.sanitizedText; assembled = proseRecovery.sanitizedText;

View File

@ -12,12 +12,13 @@ import { builtinModels } from "@earendil-works/pi-ai/providers/all";
import { type Agent, type AgentToolResult } from "@earendil-works/pi-agent-core"; import { type Agent, type AgentToolResult } from "@earendil-works/pi-agent-core";
import { import {
type AgentMessage, type AgentMessage,
type AgentToolExecutionResult,
} from "@earendil-works/pi-agent-core"; } from "@earendil-works/pi-agent-core";
import type { AgentToolExecutionResult } from "@rakazo/adapter-kit";
import type { import type {
AgentRuntimeEvent, AgentRuntimeEvent,
AgentRunRequest, AgentRunRequest,
AgentSteeringMessage, AgentSteeringMessage,
ToolCallBudgetStore,
} from "@rakazo/adapter-kit"; } from "@rakazo/adapter-kit";
import { getLogger } from "@rakazo/logging"; import { getLogger } from "@rakazo/logging";
import { DEFAULT_OPENROUTER_MODEL_ID } from "./deployment-model.js"; import { DEFAULT_OPENROUTER_MODEL_ID } from "./deployment-model.js";
@ -66,7 +67,7 @@ export interface ToolCallBudget {
inFlight: number; inFlight: number;
} }
// Optional fuse is process-local. continueRun on another worker starts at zero. // In-memory cache for same-process resume; durable store backs cross-worker continueRun.
const toolCallBudgetsByRun = new Map<string, ToolCallBudget>(); const toolCallBudgetsByRun = new Map<string, ToolCallBudget>();
/** Optional self-host fuse. Unset, empty, or 0 means unlimited (default). */ /** Optional self-host fuse. Unset, empty, or 0 means unlimited (default). */
@ -78,28 +79,70 @@ export function maxToolCallsPerTurn(env: NodeJS.ProcessEnv = process.env): numbe
return Math.floor(parsed); return Math.floor(parsed);
} }
export function toolCallBudgetFor(runId: string): ToolCallBudget { export async function loadToolCallBudget(
runId: string,
store?: ToolCallBudgetStore,
): Promise<ToolCallBudget> {
const limit = maxToolCallsPerTurn();
if (limit <= 0) {
return { count: 0, exceeded: false, limit: 0, inFlight: 0 };
}
const existing = toolCallBudgetsByRun.get(runId); const existing = toolCallBudgetsByRun.get(runId);
if (existing) { if (existing) {
existing.inFlight = 0; existing.inFlight = 0;
existing.limit = maxToolCallsPerTurn(); existing.limit = limit;
return existing; return existing;
} }
const budget: ToolCallBudget = {
count: 0, let count = 0;
exceeded: false, let exceeded = false;
limit: maxToolCallsPerTurn(), if (store) {
inFlight: 0, const snapshot = await store.load();
}; if (snapshot) {
if (budget.limit > 0) toolCallBudgetsByRun.set(runId, budget); count = snapshot.count;
exceeded = snapshot.exceeded;
}
}
const budget: ToolCallBudget = { count, exceeded, limit, inFlight: 0 };
toolCallBudgetsByRun.set(runId, budget);
return budget; return budget;
} }
export function releaseToolCallBudget(runId: string, keepForResume: boolean) { export async function releaseToolCallBudget(
if (!keepForResume) toolCallBudgetsByRun.delete(runId); runId: string,
keepForResume: boolean,
store?: ToolCallBudgetStore,
): Promise<void> {
const budget = toolCallBudgetsByRun.get(runId);
if (keepForResume) {
if (budget && store) {
await store.persist({ count: budget.count, exceeded: budget.exceeded });
}
return;
}
toolCallBudgetsByRun.delete(runId);
if (store) await store.clear();
} }
export function maybeAbortToolCallBudget(host: { /** @internal Simulates a fresh worker process in unit tests. */
export function resetToolCallBudgetMemoryForTests(): void {
toolCallBudgetsByRun.clear();
}
function persistToolCallBudget(
runId: string | undefined,
budget: ToolCallBudget,
store?: ToolCallBudgetStore,
): void {
if (!runId || !store || budget.limit <= 0) return;
void store.persist({ count: budget.count, exceeded: budget.exceeded }).catch((error) => {
getLogger().warn("tool-call budget persist failed", { runId, error });
});
}
function maybeAbortToolCallBudget(host: {
toolCallBudget: ToolCallBudget; toolCallBudget: ToolCallBudget;
abortTurn(): void; abortTurn(): void;
}): void { }): void {
@ -111,7 +154,8 @@ export function maybeAbortToolCallBudget(host: {
export function beginToolCall(host: { export function beginToolCall(host: {
toolCallBudget: ToolCallBudget; toolCallBudget: ToolCallBudget;
abortTurn(): void; abortTurn(): void;
queue: { push(event: AgentRuntimeEvent): void }; queue: { push(event: { type: string; text: string }): void };
request?: Pick<AgentRunRequest, "runId" | "toolCallBudgetStore">;
}): boolean { }): boolean {
const budget = host.toolCallBudget; const budget = host.toolCallBudget;
if (budget.limit <= 0) { if (budget.limit <= 0) {
@ -125,6 +169,7 @@ export function beginToolCall(host: {
budget.count += 1; budget.count += 1;
if (budget.count <= budget.limit) { if (budget.count <= budget.limit) {
budget.inFlight += 1; budget.inFlight += 1;
persistToolCallBudget(host.request?.runId, budget, host.request?.toolCallBudgetStore);
return true; return true;
} }
if (!budget.exceeded) { if (!budget.exceeded) {
@ -134,6 +179,7 @@ export function beginToolCall(host: {
text: `Stopped: more than ${budget.limit} tool calls in one turn.`, text: `Stopped: more than ${budget.limit} tool calls in one turn.`,
}); });
} }
persistToolCallBudget(host.request?.runId, budget, host.request?.toolCallBudgetStore);
maybeAbortToolCallBudget(host); maybeAbortToolCallBudget(host);
return false; return false;
} }
@ -141,8 +187,10 @@ export function beginToolCall(host: {
export function endToolCall(host: { export function endToolCall(host: {
toolCallBudget: ToolCallBudget; toolCallBudget: ToolCallBudget;
abortTurn(): void; abortTurn(): void;
request?: Pick<AgentRunRequest, "runId" | "toolCallBudgetStore">;
}): void { }): void {
host.toolCallBudget.inFlight = Math.max(0, host.toolCallBudget.inFlight - 1); host.toolCallBudget.inFlight = Math.max(0, host.toolCallBudget.inFlight - 1);
persistToolCallBudget(host.request?.runId, host.toolCallBudget, host.request?.toolCallBudgetStore);
maybeAbortToolCallBudget(host); maybeAbortToolCallBudget(host);
} }
@ -152,7 +200,7 @@ export function toolCallBudgetExceededMessage(limit: number): string {
// --- Model resolution --- // --- Model resolution ---
export function configuredOpenRouterModel(id: string): Model<"openai-completions"> { function configuredOpenRouterModel(id: string): Model<"openai-completions"> {
return { return {
id, id,
name: id, name: id,
@ -204,7 +252,7 @@ export function modelsForRequest(
return catalogModels(); return catalogModels();
} }
export function resolveRuntimeModel(modelConfig: AgentRunRequest["model"]): { function resolveRuntimeModel(modelConfig: AgentRunRequest["model"]): {
provider: string; provider: string;
modelId: string; modelId: string;
models: Models; models: Models;
@ -242,7 +290,7 @@ export function resolveRuntimeModel(modelConfig: AgentRunRequest["model"]): {
// --- Image helpers --- // --- Image helpers ---
export function toPiImages(images: AgentRunRequest["currentTurnImages"]) { function toPiImages(images: AgentRunRequest["currentTurnImages"]) {
return (images ?? []).map((image) => ({ return (images ?? []).map((image) => ({
type: "image" as const, type: "image" as const,
data: Buffer.from(image.data).toString("base64"), data: Buffer.from(image.data).toString("base64"),
@ -252,7 +300,7 @@ export function toPiImages(images: AgentRunRequest["currentTurnImages"]) {
// --- Steering message removal --- // --- Steering message removal ---
export function withoutSteeringMessages( function withoutSteeringMessages(
history: AgentRunRequest["history"], history: AgentRunRequest["history"],
steering: AgentSteeringMessage[], steering: AgentSteeringMessage[],
): AgentRunRequest["history"] { ): AgentRunRequest["history"] {
@ -337,7 +385,7 @@ export function pruneComputerScreenshotContext(
return transformed ?? messages; return transformed ?? messages;
} }
export function imagePartCount(message: AgentMessage): number { function imagePartCount(message: AgentMessage): number {
if (!("content" in message) || !Array.isArray(message.content)) return 0; if (!("content" in message) || !Array.isArray(message.content)) return 0;
return message.content.filter( return message.content.filter(
(part: unknown) => (part: unknown) =>
@ -345,7 +393,7 @@ export function imagePartCount(message: AgentMessage): number {
).length; ).length;
} }
export function isComputerScreenshotMessage( function isComputerScreenshotMessage(
message: AgentMessage | undefined, message: AgentMessage | undefined,
): message is Extract<AgentMessage, { role: "toolResult" }> { ): message is Extract<AgentMessage, { role: "toolResult" }> {
if (message?.role !== "toolResult" || !message.content.some((part) => (part as {type?: string}).type === "image")) { if (message?.role !== "toolResult" || !message.content.some((part) => (part as {type?: string}).type === "image")) {
@ -362,7 +410,7 @@ export function isComputerScreenshotMessage(
// --- Type guards --- // --- Type guards ---
export function isAgentToolExecutionResult(result: unknown): result is AgentToolExecutionResult { function isAgentToolExecutionResult(result: unknown): result is AgentToolExecutionResult {
if ( if (
!result || !result ||
typeof result !== "object" || typeof result !== "object" ||
@ -413,7 +461,7 @@ export function jsonSchemaParameters(
: Type.Object(fields); : Type.Object(fields);
} }
export function enumUnion(values: readonly unknown[]) { function enumUnion(values: readonly unknown[]) {
const members = values.map((value) => const members = values.map((value) =>
value === null value === null
? Type.Null() ? Type.Null()
@ -466,7 +514,7 @@ export function jsonField(spec: unknown): ReturnType<typeof Type.String> {
// --- Result helpers --- // --- Result helpers ---
export function summarizeToolResult(result: unknown) { function summarizeToolResult(result: unknown) {
try { try {
const text = JSON.stringify(result); const text = JSON.stringify(result);
if (!text) return "ok"; if (!text) return "ok";
@ -476,7 +524,7 @@ export function summarizeToolResult(result: unknown) {
} }
} }
export function boundAgentToolResult<T>( function boundAgentToolResult<T>(
result: AgentToolResult<T>, result: AgentToolResult<T>,
): AgentToolResult<T> { ): AgentToolResult<T> {
if (!result || !Array.isArray(result.content)) return result; if (!result || !Array.isArray(result.content)) return result;
@ -486,7 +534,7 @@ export function boundAgentToolResult<T>(
}; };
} }
export function assistantText(message: unknown): string { function assistantText(message: unknown): string {
if (!message || typeof message !== "object" || !("content" in message)) return ""; if (!message || typeof message !== "object" || !("content" in message)) return "";
const content = (message as { content?: unknown }).content; const content = (message as { content?: unknown }).content;
if (typeof content === "string") return content; if (typeof content === "string") return content;
@ -506,7 +554,7 @@ export function assistantText(message: unknown): string {
// --- Sanitization --- // --- Sanitization ---
export function sanitizeSensitiveText(message: string) { function sanitizeSensitiveText(message: string) {
return message return message
.replace(/sk-or-v1-[a-zA-Z0-9]+/g, "[redacted]") .replace(/sk-or-v1-[a-zA-Z0-9]+/g, "[redacted]")
.replace(/sk-[a-zA-Z0-9-]+/g, "[redacted]") .replace(/sk-[a-zA-Z0-9-]+/g, "[redacted]")
@ -526,7 +574,7 @@ export function sanitizeSensitiveText(message: string) {
); );
} }
export function redactActivityUrl(value: unknown): string { function redactActivityUrl(value: unknown): string {
const raw = String(value ?? "").trim(); const raw = String(value ?? "").trim();
if (!raw) return raw; if (!raw) return raw;
try { try {
@ -541,7 +589,7 @@ export function redactActivityUrl(value: unknown): string {
} }
} }
export function sanitizeError(message: string) { function sanitizeError(message: string) {
return sanitizeSensitiveText(message); return sanitizeSensitiveText(message);
} }
@ -562,7 +610,7 @@ export function isOpenCodeProvider(provider: string): boolean {
const OPENCODE_SESSION_ERROR = const OPENCODE_SESSION_ERROR =
"OpenCode rejected this chat session. Send the message again."; "OpenCode rejected this chat session. Send the message again.";
export function looksLikeOpenCodeSessionError(message: string): boolean { function looksLikeOpenCodeSessionError(message: string): boolean {
return ( return (
/x-opencode-session/i.test(message) || /x-opencode-session/i.test(message) ||
/session\s*(id|header|required|missing|invalid|expired|stale)/i.test(message) || /session\s*(id|header|required|missing|invalid|expired|stale)/i.test(message) ||
@ -570,7 +618,7 @@ export function looksLikeOpenCodeSessionError(message: string): boolean {
); );
} }
export function sanitizeProviderError(provider: string, message: string): string { function sanitizeProviderError(provider: string, message: string): string {
const sanitized = sanitizeError(message); const sanitized = sanitizeError(message);
if (isOpenCodeProvider(provider) && looksLikeOpenCodeSessionError(sanitized)) { if (isOpenCodeProvider(provider) && looksLikeOpenCodeSessionError(sanitized)) {
return OPENCODE_SESSION_ERROR; return OPENCODE_SESSION_ERROR;
@ -580,7 +628,7 @@ export function sanitizeProviderError(provider: string, message: string): string
// --- Queue & gate --- // --- Queue & gate ---
export interface EventQueue { interface EventQueue {
push(event: AgentRuntimeEvent): void; push(event: AgentRuntimeEvent): void;
fail(error: Error): void; fail(error: Error): void;
close(): void; close(): void;

View File

@ -5,65 +5,73 @@ import {
type AgentTool, type AgentTool,
type AgentToolResult, type AgentToolResult,
} from "@earendil-works/pi-agent-core"; } from "@earendil-works/pi-agent-core";
import { type Api, type Model, Type } from "@earendil-works/pi-ai"; import {
type Api,
clampThinkingLevel,
type Model,
type Models,
type ModelThinkingLevel,
type SimpleStreamOptions,
Type,
} from "@earendil-works/pi-ai";
import { builtinModels } from "@earendil-works/pi-ai/providers/all";
import type { import type {
AdapterContext, AdapterContext,
AgentRunRequest, AgentRunRequest,
AgentRuntime, AgentRuntime,
AgentRuntimeEvent, AgentRuntimeEvent,
AgentSteeringMessage,
AgentToolCompletion, AgentToolCompletion,
AgentToolExecutionResult,
ConnectorTool, ConnectorTool,
} from "@rakazo/adapter-kit"; } from "@rakazo/adapter-kit";
import { getLogger } from "@rakazo/logging"; import { getLogger } from "@rakazo/logging";
import { isToolPauseResult } from "./approval-effect.js"; import { isToolPauseResult } from "./approval-effect.js";
import { builtinAgentTools, DELEGATION_TOOL_NAMES } from "./builtin-tools.js"; import { builtinAgentTools, DELEGATION_TOOL_NAMES } from "./builtin-tools.js";
import { DEFAULT_OPENROUTER_MODEL_ID } from "./deployment-model.js";
import { import {
normalizeOpenAiToolParameters, normalizeOpenAiToolParameters,
openAiToolParametersNeedNormalization, openAiToolParametersNeedNormalization,
} from "./openai-tool-parameters.js"; } from "./openai-tool-parameters.js";
import { billedPromptTokens, clipToolResultText } from "./pi-runtime-limits.js"; import { PiRuntimeCredentialStore, toOAuthCredential } from "./pi-credentials.js";
import { registerLocalProvider } from "./pi-local-provider.js";
import {
OPENAI_COMPATIBLE_PROVIDER_ID,
registerOpenAiCompatibleCatalog,
registerOpenAiCompatibleRuntime,
} from "./pi-openai-compatible-provider.js";
import {
billedPromptTokens,
clipToolResultContent,
clipToolResultText,
MODEL_STREAM_MAX_RETRIES,
MODEL_STREAM_TIMEOUT_MS,
resolveCompletionMaxTokens,
} from "./pi-runtime-limits.js";
import {
beginToolCall,
endToolCall,
loadToolCallBudget,
releaseToolCallBudget,
toolCallBudgetExceededMessage,
type ToolCallBudget,
} from "./pi-model-config.js";
import { import {
PiJsonlSessionRecorder, PiJsonlSessionRecorder,
type PiSessionHandle, type PiSessionHandle,
type PiSessionRecorder, type PiSessionRecorder,
} from "./pi-session.js"; } from "./pi-session.js";
import { textContentArg } from "./tool-text.js"; import { textContentArg } from "./tool-text.js";
import {
assistantText,
beginToolCall,
boundAgentToolResult,
conversationSessionId,
createGate,
createQueue,
endToolCall,
isAgentToolExecutionResult,
jsonSchemaParameters,
looksLikeOpenCodeSessionError,
modelForCompletion,
prepareRequestSecretArguments,
pruneComputerScreenshotContext,
redactActivityUrl,
reliableStreamOptions,
releaseToolCallBudget,
resolveRuntimeModel,
sanitizeError,
sanitizeProviderError,
sanitizeSensitiveText,
summarizeToolResult,
thinkingLevelFor,
toPiImages,
toolCallBudgetExceededMessage,
toolCallBudgetFor,
withoutSteeringMessages,
type ToolCallBudget,
type ToolHost,
} from "./pi-model-config.js";
const running = new Map<string, { controller: AbortController; work: Promise<void> }>(); const running = new Map<string, { controller: AbortController; work: Promise<void> }>();
// Optional fuse is process-local. continueRun on another worker starts at zero.
// Built on first use, not at module load: entry points call loadRootEnv() after // Built on first use, not at module load: entry points call loadRootEnv() after
// their imports, and ESM hoists those imports, so module-level env reads here // their imports, and ESM hoists those imports, so module-level env reads here
// would run before .env is loaded and miss the local provider entirely. // would run before .env is loaded and miss the local provider entirely.
let catalogModelsCache: Models | undefined;
function catalogModels(): Models {
catalogModelsCache ??= registerOpenAiCompatibleCatalog(registerLocalProvider(builtinModels()));
return catalogModelsCache;
}
const MAX_PARALLEL_SUBAGENTS = 4; const MAX_PARALLEL_SUBAGENTS = 4;
// Some OpenAI-compatible models return EOS immediately after a tool result // Some OpenAI-compatible models return EOS immediately after a tool result
// instead of taking another assistant turn. A bounded internal follow-up keeps // instead of taking another assistant turn. A bounded internal follow-up keeps
@ -73,12 +81,28 @@ const SILENT_TOOL_CONTINUATION_PROMPT =
"Continue the original task from the latest tool result. Do not stop after a tool call; use any remaining tools needed, then give the user the final answer."; "Continue the original task from the latest tool result. Do not stop after a tool call; use any remaining tools needed, then give the user the final answer.";
const TOOL_FINAL_RESPONSE_FALLBACK = const TOOL_FINAL_RESPONSE_FALLBACK =
"I completed the tool step but could not produce a final response. Please ask me to continue."; "I completed the tool step but could not produce a final response. Please ask me to continue.";
const DEFAULT_COMPUTER_SCREENSHOTS_TO_KEEP = 2;
// Reasoning-capable models must not start at "off": for OpenRouter, pi-ai maps
// that to reasoning.effort "none", which 400s on endpoints that mandate
// reasoning (e.g. google/gemini-3.7-flash). Keep a real level when model.reasoning
// is set; plain models stay off.
const REASONING_MODEL_THINKING_LEVEL: ModelThinkingLevel = "medium";
function thinkingLevelFor(
model: Model<Api>,
preferred?: ModelThinkingLevel | null,
): ModelThinkingLevel {
if (!model.reasoning) return "off";
if (preferred) return clampThinkingLevel(model, preferred);
return clampThinkingLevel(model, REASONING_MODEL_THINKING_LEVEL);
}
// Pi forwards these names to OpenAI Responses, whose function-name contract is // Pi forwards these names to OpenAI Responses, whose function-name contract is
// ^[a-zA-Z0-9_-]+$ with a maximum length of 64 characters. // ^[a-zA-Z0-9_-]+$ with a maximum length of 64 characters.
const AGENT_TOOL_NAME_PATTERN = /^[a-zA-Z0-9_-]+$/; const AGENT_TOOL_NAME_PATTERN = /^[a-zA-Z0-9_-]+$/;
const MAX_AGENT_TOOL_NAME_LENGTH = 64; const MAX_AGENT_TOOL_NAME_LENGTH = 64;
const FALLBACK_AGENT_TOOL_NAME = "connector_tool"; const FALLBACK_AGENT_TOOL_NAME = "connector_tool";
export { maxToolCallsPerTurn } from "./pi-model-config.js";
export interface PiAgentRuntimeOptions { export interface PiAgentRuntimeOptions {
/** Directory where Pi JSONL sessions are written. Omit to disable recording. */ /** Directory where Pi JSONL sessions are written. Omit to disable recording. */
sessionRoot?: string; sessionRoot?: string;
@ -159,7 +183,7 @@ export class PiAgentRuntime implements AgentRuntime {
const toolDefs = request.tools.length ? request.tools : builtinAgentTools; const toolDefs = request.tools.length ? request.tools : builtinAgentTools;
const nestedAgents = new Set<Agent>(); const nestedAgents = new Set<Agent>();
const completionModel = modelForCompletion(model, request.model.maxTokens); const completionModel = modelForCompletion(model, request.model.maxTokens);
trackedBudget = toolCallBudgetFor(request.runId); trackedBudget = await loadToolCallBudget(request.runId, request.toolCallBudgetStore);
const host: ToolHost = { const host: ToolHost = {
queue, queue,
request, request,
@ -412,7 +436,11 @@ export class PiAgentRuntime implements AgentRuntime {
if (trackedBudget) { if (trackedBudget) {
const keepForResume = const keepForResume =
(signal.aborted || Boolean(resumeHost?.pausePending)) && !trackedBudget.exceeded; (signal.aborted || Boolean(resumeHost?.pausePending)) && !trackedBudget.exceeded;
releaseToolCallBudget(request.runId, keepForResume); await releaseToolCallBudget(
request.runId,
keepForResume,
request.toolCallBudgetStore,
);
} }
} }
})(); })();
@ -429,6 +457,108 @@ export class PiAgentRuntime implements AgentRuntime {
} }
} }
function toPiImages(images: AgentRunRequest["currentTurnImages"]) {
return (images ?? []).map((image) => ({
type: "image" as const,
data: Buffer.from(image.data).toString("base64"),
mimeType: image.mimeType,
}));
}
function configuredOpenRouterModel(id: string): Model<"openai-completions"> {
// A configured model can intentionally be newer than Pi's static catalog. Keep
// pricing conservative, but enable reasoning: unknown OpenRouter endpoints
// (e.g. gemini-3.7-flash before the snapshot catches up) often mandate it, and
// thinkingLevel "off" becomes effort "none" which those endpoints reject.
return {
id,
name: id,
api: "openai-completions",
provider: "openrouter",
baseUrl: "https://openrouter.ai/api/v1",
reasoning: true,
input: ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 16_384,
maxTokens: 4_096,
};
}
function resolveRuntimeModel(modelConfig: AgentRunRequest["model"]): {
provider: string;
modelId: string;
models: Models;
model: Model<Api> | undefined;
apiKey: string | undefined;
} {
const provider = modelConfig.provider === "scripted" ? "openrouter" : modelConfig.provider;
const envDefaultModel = process.env.PI_DEFAULT_MODEL?.trim();
const envDefaultProvider = process.env.PI_DEFAULT_PROVIDER?.trim() || "openrouter";
const modelId =
modelConfig.id === "scripted"
? envDefaultModel || DEFAULT_OPENROUTER_MODEL_ID
: modelConfig.id.trim();
const models = modelsForRequest({ model: modelConfig }, provider);
let model = models.getModel(provider, modelId);
if (!model && provider !== "openrouter" && provider !== OPENAI_COMPATIBLE_PROVIDER_ID) {
model = models.getModel("openrouter", modelId);
}
if (
!model &&
provider === "openrouter" &&
envDefaultProvider === "openrouter" &&
modelId === envDefaultModel
) {
model = configuredOpenRouterModel(modelId);
}
const apiKey = modelConfig.oauth
? undefined
: modelConfig.provider === OPENAI_COMPATIBLE_PROVIDER_ID
? modelConfig.apiKey || "local"
: // Only OpenRouter may fall back to the OpenRouter env key. Handing it to
// another provider would ship our key to a vendor it was not issued for.
(modelConfig.apiKey ??
(provider === "openrouter" ? process.env.OPENROUTER_API_KEY : undefined));
return { provider, modelId, models, model, apiKey };
}
export function modelsForRequest(
request: Pick<AgentRunRequest, "model">,
provider: string,
): Models {
const oauth = request.model.oauth;
if (oauth) {
const persist = oauth.persist;
return registerOpenAiCompatibleCatalog(
registerLocalProvider(
builtinModels({
credentials: new PiRuntimeCredentialStore(
provider,
toOAuthCredential(oauth.credential),
persist ? (next) => persist(next) : undefined,
),
}),
),
);
}
if (
provider === OPENAI_COMPATIBLE_PROVIDER_ID &&
request.model.baseUrl &&
request.model.id.trim()
) {
const models = registerOpenAiCompatibleCatalog(registerLocalProvider(builtinModels()));
return registerOpenAiCompatibleRuntime(models, {
modelId: request.model.id,
baseUrl: request.model.baseUrl,
reasoning: request.model.reasoning,
acceptsImages: request.model.acceptsImages,
maxTokens: request.model.maxTokens,
contextWindow: request.model.contextWindow,
});
}
return catalogModels();
}
function toAgentTools(toolDefs: readonly ConnectorTool[], host: ToolHost): AgentTool[] { function toAgentTools(toolDefs: readonly ConnectorTool[], host: ToolHost): AgentTool[] {
const names = normalizeAgentToolNames(toolDefs); const names = normalizeAgentToolNames(toolDefs);
return toolDefs.map((tool, index) => toAgentTool(tool, host, names[index]!)); return toolDefs.map((tool, index) => toAgentTool(tool, host, names[index]!));
@ -578,6 +708,52 @@ function toHistory(
); );
} }
function withoutSteeringMessages(
history: AgentRunRequest["history"],
steering: AgentSteeringMessage[],
): AgentRunRequest["history"] {
if (steering.length === 0) return history;
const result = [...history];
let beforeIndex = result.length - 1;
for (let steeringIndex = steering.length - 1; steeringIndex >= 0; steeringIndex -= 1) {
const steeringMessage = steering[steeringIndex];
for (let index = beforeIndex; index >= 0; index -= 1) {
const message = result[index];
if (
message?.role !== "user" ||
(message.id
? message.id !== steeringMessage?.messageId
: message.content !== (steeringMessage?.historyText ?? steeringMessage?.text))
) {
continue;
}
result.splice(index, 1);
beforeIndex = index - 1;
break;
}
}
return result;
}
/**
* Normalize `request_secret` arguments.
*
* `credential` and `replace` must survive: the executor stores a submitted value
* only when `credential` is present, and it validates the destination shape
* itself. An earlier version of this function listed only label/purpose/
* connectionId, so every credential the model supplied was dropped here and the
* saved value had nowhere to go.
*/
export function prepareRequestSecretArguments(raw: Record<string, unknown>) {
return {
label: String(raw.label ?? "Code"),
purpose: String(raw.purpose ?? "otp"),
...(raw.connectionId ? { connectionId: String(raw.connectionId) } : {}),
...(raw.credential ? { credential: raw.credential } : {}),
...(raw.replace === true ? { replace: true } : {}),
};
}
function toAgentTool(tool: ConnectorTool, host: ToolHost, exposedName: string): AgentTool { function toAgentTool(tool: ConnectorTool, host: ToolHost, exposedName: string): AgentTool {
return { return {
name: exposedName, name: exposedName,
@ -1089,14 +1265,397 @@ function builtinParameters(tool: ConnectorTool) {
return undefined; return undefined;
} }
export { /** Keep recent visual state while respecting an optional model image budget. */
conversationSessionId, export function pruneComputerScreenshotContext(
isOpenCodeProvider, messages: AgentMessage[],
jsonField, maxImagesPerPrompt?: number,
jsonSchemaParameters, ): AgentMessage[] {
maxToolCallsPerTurn, const imageLimit =
modelsForRequest, maxImagesPerPrompt === undefined
prepareRequestSecretArguments, ? undefined
pruneComputerScreenshotContext, : Number.isFinite(maxImagesPerPrompt)
reliableStreamOptions, ? Math.max(0, Math.floor(maxImagesPerPrompt))
} from "./pi-model-config.js"; : DEFAULT_COMPUTER_SCREENSHOTS_TO_KEEP;
let remaining = imageLimit ?? DEFAULT_COMPUTER_SCREENSHOTS_TO_KEEP;
if (imageLimit !== undefined) {
const nonScreenshotImages = messages.reduce(
(count, message) =>
isComputerScreenshotMessage(message) ? count : count + imagePartCount(message),
0,
);
if (nonScreenshotImages > imageLimit) {
throw new Error(
`The configured model image limit is ${imageLimit}, but the prompt contains ${nonScreenshotImages} non-screenshot images.`,
);
}
remaining = imageLimit - nonScreenshotImages;
}
let transformed: AgentMessage[] | undefined;
for (let index = messages.length - 1; index >= 0; index -= 1) {
const message = messages[index];
if (!isComputerScreenshotMessage(message)) continue;
const images = imagePartCount(message);
if (remaining >= images) {
remaining -= images;
continue;
}
transformed ??= [...messages];
transformed[index] = {
...message,
content: message.content.filter((part) => part.type !== "image"),
};
}
return transformed ?? messages;
}
function imagePartCount(message: AgentMessage): number {
if (!("content" in message) || !Array.isArray(message.content)) return 0;
return message.content.filter(
(part: unknown) =>
part !== null && typeof part === "object" && "type" in part && part.type === "image",
).length;
}
function isComputerScreenshotMessage(
message: AgentMessage | undefined,
): message is Extract<AgentMessage, { role: "toolResult" }> {
if (message?.role !== "toolResult" || !message.content.some((part) => part.type === "image")) {
return false;
}
const details = message.details;
return Boolean(
details &&
typeof details === "object" &&
"frameId" in details &&
typeof (details as { frameId?: unknown }).frameId === "string",
);
}
function isAgentToolExecutionResult(result: unknown): result is AgentToolExecutionResult {
if (
!result ||
typeof result !== "object" ||
(result as { kind?: unknown }).kind !== "agent_tool_result" ||
!("content" in result)
) {
return false;
}
const content = (result as { content?: unknown }).content;
return (
Array.isArray(content) &&
content.every(
(item) =>
item &&
typeof item === "object" &&
((item as { type?: unknown }).type === "text" ||
(item as { type?: unknown }).type === "image"),
)
);
}
export function jsonSchemaParameters(
schema: Record<string, unknown>,
): ReturnType<typeof Type.Object> {
// Top-level oneOf/anyOf (e.g. request_secret's credential XOR connectionId)
// must stay a union. Falling through to properties would drop the exclusivity
// and re-expose both destinations as optional siblings.
const alternatives = Array.isArray(schema.oneOf)
? schema.oneOf
: Array.isArray(schema.anyOf)
? schema.anyOf
: undefined;
if (alternatives && alternatives.length > 0 && schema.properties == null) {
return Type.Union(
alternatives.map((variant) => jsonSchemaParameters(variant as Record<string, unknown>)),
) as unknown as ReturnType<typeof Type.Object>;
}
const properties = (schema.properties ?? {}) as Record<string, unknown>;
const required = new Set(Array.isArray(schema.required) ? schema.required.map(String) : []);
const fields: Record<string, ReturnType<typeof Type.Optional>> = {};
for (const [key, spec] of Object.entries(properties)) {
const field = jsonField(spec);
fields[key] = (required.has(key) ? field : Type.Optional(field)) as unknown as ReturnType<
typeof Type.Optional
>;
}
// Preserve closed objects (e.g. request_secret destination oneOf branches).
// Type.Object defaults to open, which would let connectionId+replace match both
// anyOf variants after conversion.
return schema.additionalProperties === false
? Type.Object(fields, { additionalProperties: false })
: Type.Object(fields);
}
/** TypeBox only builds literals from primitives; anything else throws while the tool list is
* being assembled, which would take down the whole turn. */
function enumUnion(values: readonly unknown[]) {
const members = values.map((value) =>
value === null
? Type.Null()
: typeof value === "string" || typeof value === "number" || typeof value === "boolean"
? Type.Literal(value)
: undefined,
);
return members.every((member) => member !== undefined) ? Type.Union(members) : undefined;
}
export function jsonField(spec: unknown): ReturnType<typeof Type.String> {
const definition = spec && typeof spec === "object" ? (spec as Record<string, unknown>) : {};
if (Array.isArray(definition.enum) && definition.enum.length > 0) {
const union = enumUnion(definition.enum);
if (union) return union as never;
}
// A `const` names the only accepted value. Without this it degraded to a bare
// string, so a discriminator like {type: {const: "bearer"}} told the model
// nothing about which value to send -- and it guessed, twice.
if ("const" in definition) {
const literal = enumUnion([definition.const]);
if (literal) return literal as never;
}
// A discriminated union arrives as oneOf/anyOf with no sibling `type`. Without
// this branch it fell through to the string default, so a model was told to
// send an object-valued field as a bare string -- which is exactly what it did.
const variants = Array.isArray(definition.oneOf)
? definition.oneOf
: Array.isArray(definition.anyOf)
? definition.anyOf
: undefined;
if (variants && variants.length > 0) {
return Type.Union(variants.map((variant) => jsonField(variant))) as never;
}
if (Array.isArray(definition.type) && definition.type.length > 0) {
return Type.Union(definition.type.map((type) => jsonField({ ...definition, type }))) as never;
}
const type = "type" in definition ? String(definition.type) : "string";
if (type === "null") return Type.Null() as never;
if (type === "number" || type === "integer") return Type.Number() as never;
if (type === "boolean") return Type.Boolean() as never;
if (type === "array") {
const options: {
minItems?: number;
maxItems?: number;
uniqueItems?: boolean;
} = {};
if (typeof definition.minItems === "number") options.minItems = definition.minItems;
if (typeof definition.maxItems === "number") options.maxItems = definition.maxItems;
if (definition.uniqueItems === true) options.uniqueItems = true;
return Type.Array(jsonField(definition.items), options) as never;
}
if (type === "object") return jsonSchemaParameters(definition) as never;
return Type.String();
}
function summarizeToolResult(result: unknown) {
try {
const text = JSON.stringify(result);
if (!text) return "ok";
return clipToolResultText(text);
} catch {
return "ok";
}
}
function boundAgentToolResult<T>(result: AgentToolResult<T>): AgentToolResult<T> {
if (!result || !Array.isArray(result.content)) return result;
return {
...result,
content: clipToolResultContent(result.content),
};
}
function assistantText(message: unknown): string {
if (!message || typeof message !== "object" || !("content" in message)) return "";
const content = (message as { content?: unknown }).content;
if (typeof content === "string") return content;
if (!Array.isArray(content)) return "";
return content
.map((part) =>
part && typeof part === "object" && "type" in part && part.type === "text" && "text" in part
? String(part.text)
: "",
)
.join("");
}
function sanitizeSensitiveText(message: string) {
return message
.replace(/sk-or-v1-[a-zA-Z0-9]+/g, "[redacted]")
.replace(/sk-[a-zA-Z0-9-]+/g, "[redacted]")
.replace(/Bearer\s+[^\s"',;&]+/gi, "Bearer [redacted]")
.replace(/eyJ[a-zA-Z0-9_-]+\.[a-zA-Z0-9_-]+\.[a-zA-Z0-9_-]+/g, "[redacted]")
.replace(/COMPOSIO_API_KEY[=:]?\s*\S+/gi, "COMPOSIO_API_KEY=[redacted]")
.replace(
/((?:api[_-]?key|access[_-]?token|password|secret)\s*[=:]\s*)[^\s"',;&]+/gi,
"$1[redacted]",
)
.replace(/((?:auth|authorization)\s*[=:]\s*)(?!Bearer\b)[^\s"',;&]+/gi, "$1[redacted]");
}
/** Origin + path only for activity chips; drop userinfo, query, and fragment. */
function redactActivityUrl(value: unknown): string {
const raw = String(value ?? "").trim();
if (!raw) return raw;
try {
const url = new URL(raw);
url.username = "";
url.password = "";
url.search = "";
url.hash = "";
return url.toString();
} catch {
// Never echo unparsed input — it may still contain userinfo/secrets.
return "[invalid URL]";
}
}
function sanitizeError(message: string) {
return sanitizeSensitiveText(message);
}
/** Stable OpenCode affinity id for a bot conversation (and optional nested agent). */
export function conversationSessionId(threadId: string, botId: string, agentId?: string): string {
return agentId ? `${threadId}:${botId}:${agentId}` : `${threadId}:${botId}`;
}
export function isOpenCodeProvider(provider: string): boolean {
return provider === "opencode" || provider === "opencode-go";
}
const OPENCODE_SESSION_ERROR = "OpenCode rejected this chat session. Send the message again.";
function looksLikeOpenCodeSessionError(message: string): boolean {
return (
/x-opencode-session/i.test(message) ||
/session\s*(id|header|required|missing|invalid|expired|stale)/i.test(message) ||
/model is unavailable/i.test(message)
);
}
function sanitizeProviderError(provider: string, message: string): string {
const sanitized = sanitizeError(message);
if (isOpenCodeProvider(provider) && looksLikeOpenCodeSessionError(sanitized)) {
return OPENCODE_SESSION_ERROR;
}
return sanitized;
}
interface EventQueue {
push(event: AgentRuntimeEvent): void;
fail(error: Error): void;
close(): void;
iterate(): AsyncIterable<AgentRuntimeEvent>;
}
interface ToolHost {
queue: EventQueue;
request: AgentRunRequest;
models: Models;
model: Model<Api>;
apiKey: string | undefined;
nestedAgents: Set<Agent>;
subagentGate: { acquire(): Promise<void>; release(): void };
toolCallBudget: ToolCallBudget;
/** Shared fallback uniqueness when the model omits toolCallId (nested hosts reuse this). */
toolCallSeq: { value: number };
abortTurn(): void;
signal: AbortSignal;
depth: number;
pausePending: boolean;
}
function modelForCompletion(model: Model<Api>, configuredMaxTokens?: number): Model<Api> {
const maxTokens = resolveCompletionMaxTokens(model.maxTokens, configuredMaxTokens);
if (maxTokens === model.maxTokens) return model;
return { ...model, maxTokens };
}
function createGate(max: number) {
let active = 0;
const waiters: Array<() => void> = [];
return {
async acquire() {
if (active >= max) {
await new Promise<void>((resolve) => {
waiters.push(resolve);
});
}
active += 1;
},
release() {
active = Math.max(0, active - 1);
waiters.shift()?.();
},
};
}
function createQueue(): EventQueue {
const items: AgentRuntimeEvent[] = [];
let wake: (() => void) | undefined;
let closed = false;
let failure: Error | undefined;
return {
push(event) {
items.push(event);
wake?.();
},
fail(error) {
failure = error;
closed = true;
wake?.();
},
close() {
closed = true;
wake?.();
},
async *iterate() {
while (true) {
if (items.length) {
yield items.shift()!;
continue;
}
if (failure) throw failure;
if (closed) return;
await new Promise<void>((resolve) => {
wake = resolve;
});
}
},
};
}
export function reliableStreamOptions(
model: Pick<Model<Api>, "api" | "provider" | "maxTokens">,
options?: SimpleStreamOptions,
configuredMaxTokens?: number,
): SimpleStreamOptions {
let next: SimpleStreamOptions = {
...options,
timeoutMs: options?.timeoutMs ?? MODEL_STREAM_TIMEOUT_MS,
maxRetries: options?.maxRetries ?? MODEL_STREAM_MAX_RETRIES,
maxTokens: resolveCompletionMaxTokens(model.maxTokens, configuredMaxTokens, options?.maxTokens),
};
if (model.provider === "openai-codex" || model.api === "openai-codex-responses") {
// Pi cannot fall back after a WebSocket has emitted its start event. Long tool
// runs then surface abnormal close 1006 as a terminal model error. SSE has
// bounded network retries and no long-lived connection between tool turns.
next = { ...next, transport: "sse" };
}
// OpenCode Go/Zen require a sticky x-opencode-session header (affinity + some
// models 400 without it). Pi 0.85.1 does not attach that header on its own.
if (isOpenCodeProvider(model.provider)) {
const sessionId = next.sessionId?.trim() || randomUUID();
next = {
...next,
sessionId,
headers: {
"x-opencode-session": sessionId,
"x-opencode-client": "rakazo",
...next.headers,
},
};
}
return next;
}

View File

@ -21,7 +21,6 @@ import {
createThreadMessageInTransaction, createThreadMessageInTransaction,
RunHistoryWriteError, RunHistoryWriteError,
} from "./messages.js"; } from "./messages.js";
import { ClearThreadPendingDelegatedReturnsError } from "./scope.js";
import { withTransactionRetry } from "./transaction-retry.js"; import { withTransactionRetry } from "./transaction-retry.js";
const EVENT_BATCH_SIZE = 200; const EVENT_BATCH_SIZE = 200;
@ -272,19 +271,6 @@ export async function clearThread(
const now = new Date(); const now = new Date();
const runIds = activeRuns.map((run) => run.id); const runIds = activeRuns.map((run) => run.id);
await cancelRunsInTransaction(tx, activeRuns, now); await cancelRunsInTransaction(tx, activeRuns, now);
const pendingDelegatedReturns = await tx.run.count({
where: {
spaceId: input.spaceId,
threadId: input.threadId,
...(input.groupId ? {} : { botId: input.botId }),
trigger: "bot_message",
status: { in: ["completed", "failed"] },
botOutcomeReturnedAt: null,
},
});
if (pendingDelegatedReturns > 0) {
throw new ClearThreadPendingDelegatedReturnsError();
}
// Expire as tombstones so a still-open provider screen claim cannot reset fencing to 1. // Expire as tombstones so a still-open provider screen claim cannot reset fencing to 1.
await expireComputerExecutionLeases(tx, { runId: { in: runIds } }); await expireComputerExecutionLeases(tx, { runId: { in: runIds } });
await tx.computer.updateMany({ await tx.computer.updateMany({
@ -1022,6 +1008,8 @@ async function finalizeRunOnce(
completedAt: now, completedAt: now,
leaseOwner: null, leaseOwner: null,
leaseExpiresAt: null, leaseExpiresAt: null,
toolCallBudgetCount: null,
toolCallBudgetExceeded: null,
}, },
}); });
if (terminal.count !== 1) return null; if (terminal.count !== 1) return null;