fix(adapters): persist tool-call budget across worker handoffs
This commit is contained in:
parent
6e33307c7f
commit
beece5a383
@ -402,6 +402,19 @@ export interface AgentRunRequest {
|
||||
onToolCompleted?: (completion: AgentToolCompletion) => Promise<void> | void;
|
||||
/** Atomically claim durable user steering at the runtime's next safe turn boundary. */
|
||||
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 {
|
||||
@ -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, …). */
|
||||
|
||||
@ -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;
|
||||
|
||||
@ -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<string, ToolCallBudget>();
|
||||
|
||||
/** 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<ToolCallBudget> {
|
||||
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<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;
|
||||
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<AgentRunRequest, "runId" | "toolCallBudgetStore">;
|
||||
}): 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<AgentRunRequest, "runId" | "toolCallBudgetStore">;
|
||||
}): 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<AgentMessage, { role: "toolResult" }> {
|
||||
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<typeof Type.String> {
|
||||
|
||||
// --- 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<T>(
|
||||
function boundAgentToolResult<T>(
|
||||
result: AgentToolResult<T>,
|
||||
): AgentToolResult<T> {
|
||||
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 "";
|
||||
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;
|
||||
|
||||
@ -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<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
|
||||
// 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<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
|
||||
// ^[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<Agent>();
|
||||
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<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[] {
|
||||
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<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 {
|
||||
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<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;
|
||||
}
|
||||
|
||||
@ -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;
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user