diff --git a/packages/adapter-kit/src/types.ts b/packages/adapter-kit/src/types.ts index 4e060be..afd90d4 100644 --- a/packages/adapter-kit/src/types.ts +++ b/packages/adapter-kit/src/types.ts @@ -402,6 +402,19 @@ export interface AgentRunRequest { onToolCompleted?: (completion: AgentToolCompletion) => Promise | void; /** Atomically claim durable user steering at the runtime's next safe turn boundary. */ claimSteering?: (seenIds: string[]) => Promise; + /** 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; + persist(snapshot: ToolCallBudgetSnapshot): Promise; + clear(): Promise; } export interface ScriptedTurn { @@ -587,11 +600,6 @@ export interface MessagingInboundMessage { channelName: string | null; /** Group roster addresses when the platform reports them; often empty. */ 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; mediaUrl: string | null; /** Team-room workspace/team id when the platform reports one (Slack team_id, …). */ diff --git a/packages/adapters/src/executor.ts b/packages/adapters/src/executor.ts index 99ac3a1..68761fa 100644 --- a/packages/adapters/src/executor.ts +++ b/packages/adapters/src/executor.ts @@ -263,6 +263,7 @@ import { secretValuesToRedact, serializeModelSecret, } from "./pi-oauth.js"; +import { maxToolCallsPerTurn } from "./pi-model-config.js"; import { assertPlotDataWithinLimits, PLOT_TOOL_GUIDE, @@ -273,6 +274,7 @@ import { searchChartCatalog, } from "./plot-tool.js"; import type { RemoteTransportDependencies } from "./remote-mcp.js"; +import { createRunToolCallBudgetStore } from "./tool-call-budget-store.js"; import { loadReplyContext, messageToAgentHistoryText } from "./reply-context.js"; import { commitConsumedRunSecret, @@ -3717,6 +3719,10 @@ export function createRunExecutor(deps: ExecutorDeps) { completion, runSecrets, ), + toolCallBudgetStore: + maxToolCallsPerTurn() > 0 + ? createRunToolCallBudgetStore(deps.prisma, runId) + : undefined, claimSteering: scripted ? undefined : async (seenIds) => { @@ -4169,7 +4175,7 @@ export function createRunExecutor(deps: ExecutorDeps) { const workshopFallback = workshopFallbackAfterBlockedHandoff( proseRecovery.sanitizedText, proseRecovery.blockedUpward, - proseRecovery.handoffFailed, + proseRecovery.proseToolCalls, ); if (proseRecovery.proseToolCalls.length > 0) { assembled = proseRecovery.sanitizedText; diff --git a/packages/adapters/src/pi-model-config.ts b/packages/adapters/src/pi-model-config.ts index d38702a..fdc69b3 100644 --- a/packages/adapters/src/pi-model-config.ts +++ b/packages/adapters/src/pi-model-config.ts @@ -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 AgentMessage, + type AgentToolExecutionResult, } from "@earendil-works/pi-agent-core"; -import type { AgentToolExecutionResult } from "@rakazo/adapter-kit"; import type { AgentRuntimeEvent, AgentRunRequest, AgentSteeringMessage, + ToolCallBudgetStore, } from "@rakazo/adapter-kit"; import { getLogger } from "@rakazo/logging"; import { DEFAULT_OPENROUTER_MODEL_ID } from "./deployment-model.js"; @@ -66,7 +67,7 @@ export interface ToolCallBudget { 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(); /** 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); } -export function toolCallBudgetFor(runId: string): ToolCallBudget { +export async function loadToolCallBudget( + runId: string, + store?: ToolCallBudgetStore, +): Promise { + const limit = maxToolCallsPerTurn(); + if (limit <= 0) { + return { count: 0, exceeded: false, limit: 0, inFlight: 0 }; + } + const existing = toolCallBudgetsByRun.get(runId); if (existing) { existing.inFlight = 0; - existing.limit = maxToolCallsPerTurn(); + existing.limit = limit; return existing; } - const budget: ToolCallBudget = { - count: 0, - exceeded: false, - limit: maxToolCallsPerTurn(), - inFlight: 0, - }; - if (budget.limit > 0) toolCallBudgetsByRun.set(runId, budget); + + let count = 0; + let exceeded = false; + if (store) { + const snapshot = await store.load(); + if (snapshot) { + count = snapshot.count; + exceeded = snapshot.exceeded; + } + } + + const budget: ToolCallBudget = { count, exceeded, limit, inFlight: 0 }; + toolCallBudgetsByRun.set(runId, budget); return budget; } -export function releaseToolCallBudget(runId: string, keepForResume: boolean) { - if (!keepForResume) toolCallBudgetsByRun.delete(runId); +export async function releaseToolCallBudget( + runId: string, + keepForResume: boolean, + store?: ToolCallBudgetStore, +): Promise { + 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; abortTurn(): void; }): void { @@ -111,7 +154,8 @@ export function maybeAbortToolCallBudget(host: { export function beginToolCall(host: { toolCallBudget: ToolCallBudget; abortTurn(): void; - queue: { push(event: AgentRuntimeEvent): void }; + queue: { push(event: { type: string; text: string }): void }; + request?: Pick; }): boolean { const budget = host.toolCallBudget; if (budget.limit <= 0) { @@ -125,6 +169,7 @@ export function beginToolCall(host: { budget.count += 1; if (budget.count <= budget.limit) { budget.inFlight += 1; + persistToolCallBudget(host.request?.runId, budget, host.request?.toolCallBudgetStore); return true; } if (!budget.exceeded) { @@ -134,6 +179,7 @@ export function beginToolCall(host: { text: `Stopped: more than ${budget.limit} tool calls in one turn.`, }); } + persistToolCallBudget(host.request?.runId, budget, host.request?.toolCallBudgetStore); maybeAbortToolCallBudget(host); return false; } @@ -141,8 +187,10 @@ export function beginToolCall(host: { export function endToolCall(host: { toolCallBudget: ToolCallBudget; abortTurn(): void; + request?: Pick; }): void { host.toolCallBudget.inFlight = Math.max(0, host.toolCallBudget.inFlight - 1); + persistToolCallBudget(host.request?.runId, host.toolCallBudget, host.request?.toolCallBudgetStore); maybeAbortToolCallBudget(host); } @@ -152,7 +200,7 @@ export function toolCallBudgetExceededMessage(limit: number): string { // --- Model resolution --- -export function configuredOpenRouterModel(id: string): Model<"openai-completions"> { +function configuredOpenRouterModel(id: string): Model<"openai-completions"> { return { id, name: id, @@ -204,7 +252,7 @@ export function modelsForRequest( return catalogModels(); } -export function resolveRuntimeModel(modelConfig: AgentRunRequest["model"]): { +function resolveRuntimeModel(modelConfig: AgentRunRequest["model"]): { provider: string; modelId: string; models: Models; @@ -242,7 +290,7 @@ export function resolveRuntimeModel(modelConfig: AgentRunRequest["model"]): { // --- Image helpers --- -export function toPiImages(images: AgentRunRequest["currentTurnImages"]) { +function toPiImages(images: AgentRunRequest["currentTurnImages"]) { return (images ?? []).map((image) => ({ type: "image" as const, data: Buffer.from(image.data).toString("base64"), @@ -252,7 +300,7 @@ export function toPiImages(images: AgentRunRequest["currentTurnImages"]) { // --- Steering message removal --- -export function withoutSteeringMessages( +function withoutSteeringMessages( history: AgentRunRequest["history"], steering: AgentSteeringMessage[], ): AgentRunRequest["history"] { @@ -337,7 +385,7 @@ export function pruneComputerScreenshotContext( return transformed ?? messages; } -export function imagePartCount(message: AgentMessage): number { +function imagePartCount(message: AgentMessage): number { if (!("content" in message) || !Array.isArray(message.content)) return 0; return message.content.filter( (part: unknown) => @@ -345,7 +393,7 @@ export function imagePartCount(message: AgentMessage): number { ).length; } -export function isComputerScreenshotMessage( +function isComputerScreenshotMessage( message: AgentMessage | undefined, ): message is Extract { if (message?.role !== "toolResult" || !message.content.some((part) => (part as {type?: string}).type === "image")) { @@ -362,7 +410,7 @@ export function isComputerScreenshotMessage( // --- Type guards --- -export function isAgentToolExecutionResult(result: unknown): result is AgentToolExecutionResult { +function isAgentToolExecutionResult(result: unknown): result is AgentToolExecutionResult { if ( !result || typeof result !== "object" || @@ -413,7 +461,7 @@ export function jsonSchemaParameters( : Type.Object(fields); } -export function enumUnion(values: readonly unknown[]) { +function enumUnion(values: readonly unknown[]) { const members = values.map((value) => value === null ? Type.Null() @@ -466,7 +514,7 @@ export function jsonField(spec: unknown): ReturnType { // --- Result helpers --- -export function summarizeToolResult(result: unknown) { +function summarizeToolResult(result: unknown) { try { const text = JSON.stringify(result); if (!text) return "ok"; @@ -476,7 +524,7 @@ export function summarizeToolResult(result: unknown) { } } -export function boundAgentToolResult( +function boundAgentToolResult( result: AgentToolResult, ): AgentToolResult { if (!result || !Array.isArray(result.content)) return result; @@ -486,7 +534,7 @@ export function boundAgentToolResult( }; } -export function assistantText(message: unknown): string { +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; @@ -506,7 +554,7 @@ export function assistantText(message: unknown): string { // --- Sanitization --- -export function sanitizeSensitiveText(message: string) { +function sanitizeSensitiveText(message: string) { return message .replace(/sk-or-v1-[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(); if (!raw) return raw; try { @@ -541,7 +589,7 @@ export function redactActivityUrl(value: unknown): string { } } -export function sanitizeError(message: string) { +function sanitizeError(message: string) { return sanitizeSensitiveText(message); } @@ -562,7 +610,7 @@ export function isOpenCodeProvider(provider: string): boolean { const OPENCODE_SESSION_ERROR = "OpenCode rejected this chat session. Send the message again."; -export function looksLikeOpenCodeSessionError(message: string): boolean { +function looksLikeOpenCodeSessionError(message: string): boolean { return ( /x-opencode-session/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); if (isOpenCodeProvider(provider) && looksLikeOpenCodeSessionError(sanitized)) { return OPENCODE_SESSION_ERROR; @@ -580,7 +628,7 @@ export function sanitizeProviderError(provider: string, message: string): string // --- Queue & gate --- -export interface EventQueue { +interface EventQueue { push(event: AgentRuntimeEvent): void; fail(error: Error): void; close(): void; diff --git a/packages/adapters/src/pi-runtime.ts b/packages/adapters/src/pi-runtime.ts index f537483..a53e1c9 100644 --- a/packages/adapters/src/pi-runtime.ts +++ b/packages/adapters/src/pi-runtime.ts @@ -5,65 +5,73 @@ import { type AgentTool, type AgentToolResult, } 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 { AdapterContext, AgentRunRequest, AgentRuntime, AgentRuntimeEvent, + AgentSteeringMessage, AgentToolCompletion, + AgentToolExecutionResult, ConnectorTool, } from "@rakazo/adapter-kit"; import { getLogger } from "@rakazo/logging"; import { isToolPauseResult } from "./approval-effect.js"; import { builtinAgentTools, DELEGATION_TOOL_NAMES } from "./builtin-tools.js"; +import { DEFAULT_OPENROUTER_MODEL_ID } from "./deployment-model.js"; import { normalizeOpenAiToolParameters, openAiToolParametersNeedNormalization, } 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 { PiJsonlSessionRecorder, type PiSessionHandle, type PiSessionRecorder, } from "./pi-session.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 }>(); -// 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 // 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. +let catalogModelsCache: Models | undefined; +function catalogModels(): Models { + catalogModelsCache ??= registerOpenAiCompatibleCatalog(registerLocalProvider(builtinModels())); + return catalogModelsCache; +} const MAX_PARALLEL_SUBAGENTS = 4; // Some OpenAI-compatible models return EOS immediately after a tool result // 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."; const TOOL_FINAL_RESPONSE_FALLBACK = "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, + 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 // ^[a-zA-Z0-9_-]+$ with a maximum length of 64 characters. const AGENT_TOOL_NAME_PATTERN = /^[a-zA-Z0-9_-]+$/; const MAX_AGENT_TOOL_NAME_LENGTH = 64; const FALLBACK_AGENT_TOOL_NAME = "connector_tool"; +export { maxToolCallsPerTurn } from "./pi-model-config.js"; + export interface PiAgentRuntimeOptions { /** Directory where Pi JSONL sessions are written. Omit to disable recording. */ sessionRoot?: string; @@ -159,7 +183,7 @@ export class PiAgentRuntime implements AgentRuntime { const toolDefs = request.tools.length ? request.tools : builtinAgentTools; const nestedAgents = new Set(); const completionModel = modelForCompletion(model, request.model.maxTokens); - trackedBudget = toolCallBudgetFor(request.runId); + trackedBudget = await loadToolCallBudget(request.runId, request.toolCallBudgetStore); const host: ToolHost = { queue, request, @@ -412,7 +436,11 @@ export class PiAgentRuntime implements AgentRuntime { if (trackedBudget) { const keepForResume = (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 | 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, + 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[] { const names = normalizeAgentToolNames(toolDefs); 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) { + 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 { return { name: exposedName, @@ -1089,14 +1265,397 @@ function builtinParameters(tool: ConnectorTool) { return undefined; } -export { - conversationSessionId, - isOpenCodeProvider, - jsonField, - jsonSchemaParameters, - maxToolCallsPerTurn, - modelsForRequest, - prepareRequestSecretArguments, - pruneComputerScreenshotContext, - reliableStreamOptions, -} from "./pi-model-config.js"; +/** Keep recent visual state while respecting an optional model image budget. */ +export function pruneComputerScreenshotContext( + messages: AgentMessage[], + maxImagesPerPrompt?: number, +): AgentMessage[] { + const imageLimit = + maxImagesPerPrompt === undefined + ? undefined + : Number.isFinite(maxImagesPerPrompt) + ? Math.max(0, Math.floor(maxImagesPerPrompt)) + : 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 { + 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, +): ReturnType { + // 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)), + ) as unknown as ReturnType; + } + const properties = (schema.properties ?? {}) as Record; + const required = new Set(Array.isArray(schema.required) ? schema.required.map(String) : []); + const fields: Record> = {}; + 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 { + const definition = spec && typeof spec === "object" ? (spec as Record) : {}; + 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(result: AgentToolResult): AgentToolResult { + 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; +} + +interface ToolHost { + queue: EventQueue; + request: AgentRunRequest; + models: Models; + model: Model; + apiKey: string | undefined; + nestedAgents: Set; + subagentGate: { acquire(): Promise; 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, configuredMaxTokens?: number): Model { + 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((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((resolve) => { + wake = resolve; + }); + } + }, + }; +} + +export function reliableStreamOptions( + model: Pick, "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; +} diff --git a/packages/db/src/events.ts b/packages/db/src/events.ts index 120219d..5cb4bcc 100644 --- a/packages/db/src/events.ts +++ b/packages/db/src/events.ts @@ -21,7 +21,6 @@ import { createThreadMessageInTransaction, RunHistoryWriteError, } from "./messages.js"; -import { ClearThreadPendingDelegatedReturnsError } from "./scope.js"; import { withTransactionRetry } from "./transaction-retry.js"; const EVENT_BATCH_SIZE = 200; @@ -272,19 +271,6 @@ export async function clearThread( const now = new Date(); const runIds = activeRuns.map((run) => run.id); 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. await expireComputerExecutionLeases(tx, { runId: { in: runIds } }); await tx.computer.updateMany({ @@ -1022,6 +1008,8 @@ async function finalizeRunOnce( completedAt: now, leaseOwner: null, leaseExpiresAt: null, + toolCallBudgetCount: null, + toolCallBudgetExceeded: null, }, }); if (terminal.count !== 1) return null;