evOo-local/packages/adapters/src/pi-model-config.ts

import { randomUUID } from "node:crypto";
import type { Agent, AgentMessage, AgentToolResult } from "@earendil-works/pi-agent-core";
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 {
  AgentRunRequest,
  AgentRuntimeEvent,
  AgentSteeringMessage,
  AgentToolExecutionResult,
} from "@rakazo/adapter-kit";
import { getLogger } from "@rakazo/logging";
import { DEFAULT_OPENROUTER_MODEL_ID } from "./deployment-model.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";

// --- Model catalog ---

const REASONING_MODEL_THINKING_LEVEL: ModelThinkingLevel = "medium";

let catalogModelsCache: Models | undefined;

export function catalogModels(): Models {
  catalogModelsCache ??= registerOpenAiCompatibleCatalog(registerLocalProvider(builtinModels()));
  return catalogModelsCache;
}

export 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);
}

// --- Tool call budget ---

export interface ToolCallBudget {
  count: number;
  exceeded: boolean;
  limit: number;
  inFlight: number;
}

// Optional fuse is process-local. continueRun on another worker starts at zero.
const toolCallBudgetsByRun = new Map<string, ToolCallBudget>();

/** Optional self-host fuse. Unset, empty, or 0 means unlimited (default). */
export function maxToolCallsPerTurn(env: NodeJS.ProcessEnv = process.env): number {
  const raw = env.MAX_TOOL_CALLS_PER_TURN?.trim();
  if (!raw) return 0;
  const parsed = Number(raw);
  if (!Number.isFinite(parsed) || parsed <= 0) return 0;
  return Math.floor(parsed);
}

export function toolCallBudgetFor(runId: string): ToolCallBudget {
  const existing = toolCallBudgetsByRun.get(runId);
  if (existing) {
    existing.inFlight = 0;
    existing.limit = maxToolCallsPerTurn();
    return existing;
  }
  const budget: ToolCallBudget = {
    count: 0,
    exceeded: false,
    limit: maxToolCallsPerTurn(),
    inFlight: 0,
  };
  if (budget.limit > 0) toolCallBudgetsByRun.set(runId, budget);
  return budget;
}

export function releaseToolCallBudget(runId: string, keepForResume: boolean) {
  if (!keepForResume) toolCallBudgetsByRun.delete(runId);
}

export function maybeAbortToolCallBudget(host: {
  toolCallBudget: ToolCallBudget;
  abortTurn(): void;
}): void {
  if (host.toolCallBudget.exceeded && host.toolCallBudget.inFlight === 0) {
    host.abortTurn();
  }
}

export function beginToolCall(host: {
  toolCallBudget: ToolCallBudget;
  abortTurn(): void;
  queue: { push(event: AgentRuntimeEvent): void };
}): boolean {
  const budget = host.toolCallBudget;
  if (budget.limit <= 0) {
    budget.count += 1;
    return true;
  }
  if (budget.exceeded) {
    maybeAbortToolCallBudget(host);
    return false;
  }
  budget.count += 1;
  if (budget.count <= budget.limit) {
    budget.inFlight += 1;
    return true;
  }
  if (!budget.exceeded) {
    budget.exceeded = true;
    host.queue.push({
      type: "progress",
      text: `Stopped: more than ${budget.limit} tool calls in one turn.`,
    });
  }
  maybeAbortToolCallBudget(host);
  return false;
}

export function endToolCall(host: { toolCallBudget: ToolCallBudget; abortTurn(): void }): void {
  host.toolCallBudget.inFlight = Math.max(0, host.toolCallBudget.inFlight - 1);
  maybeAbortToolCallBudget(host);
}

export function toolCallBudgetExceededMessage(limit: number): string {
  return `I stopped after reaching the limit of ${limit} tool calls in this turn. Send another message to continue.`;
}

// --- Model resolution ---

export function configuredOpenRouterModel(id: string): Model<"openai-completions"> {
  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,
  };
}

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();
}

export 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"
      : (modelConfig.apiKey ??
        (provider === "openrouter" ? process.env.OPENROUTER_API_KEY : undefined));
  return { provider, modelId, models, model, apiKey };
}

// --- Image helpers ---

export function toPiImages(images: AgentRunRequest["currentTurnImages"]) {
  return (images ?? []).map((image) => ({
    type: "image" as const,
    data: Buffer.from(image.data).toString("base64"),
    mimeType: image.mimeType,
  }));
}

// --- Steering message removal ---

export 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;
}

// --- Request secret arguments ---

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 } : {}),
  };
}

// --- Screenshot context pruning ---

const DEFAULT_COMPUTER_SCREENSHOTS_TO_KEEP = 2;

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;
}

export function imagePartCount(message: AgentMessage): number {
  if (!("content" in message) || !Array.isArray(message.content)) return 0;
  return message.content.filter(
    (part: unknown) =>
      typeof part === "object" &&
      part !== null &&
      "type" in part &&
      (part as { type?: string }).type === "image",
  ).length;
}

export 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")
  ) {
    return false;
  }
  const details = message.details;
  return Boolean(
    details &&
      typeof details === "object" &&
      "frameId" in details &&
      typeof (details as { frameId?: unknown }).frameId === "string",
  );
}

// --- Type guards ---

export 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"),
    )
  );
}

// --- JSON Schema helpers ---

export function jsonSchemaParameters(
  schema: Record<string, unknown>,
): ReturnType<typeof Type.Object> {
  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
    >;
  }
  return schema.additionalProperties === false
    ? Type.Object(fields, { additionalProperties: false })
    : Type.Object(fields);
}

export 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;
  }
  if ("const" in definition) {
    const literal = enumUnion([definition.const]);
    if (literal) return literal as never;
  }
  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();
}

// --- Result helpers ---

export function summarizeToolResult(result: unknown) {
  try {
    const text = JSON.stringify(result);
    if (!text) return "ok";
    return clipToolResultText(text);
  } catch {
    return "ok";
  }
}

export function boundAgentToolResult<T>(result: AgentToolResult<T>): AgentToolResult<T> {
  if (!result || !Array.isArray(result.content)) return result;
  return {
    ...result,
    content: clipToolResultContent(result.content),
  };
}

export 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("");
}

// --- Sanitization ---

export 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]");
}

export 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 {
    return "[invalid URL]";
  }
}

export function sanitizeError(message: string) {
  return sanitizeSensitiveText(message);
}

// --- Session ID & provider helpers ---

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.";

export 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)
  );
}

export function sanitizeProviderError(provider: string, message: string): string {
  const sanitized = sanitizeError(message);
  if (isOpenCodeProvider(provider) && looksLikeOpenCodeSessionError(sanitized)) {
    return OPENCODE_SESSION_ERROR;
  }
  return sanitized;
}

// --- Queue & gate ---

export interface EventQueue {
  push(event: AgentRuntimeEvent): void;
  fail(error: Error): void;
  close(): void;
  iterate(): AsyncIterable<AgentRuntimeEvent>;
}

export 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 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()?.();
    },
  };
}

// --- Model completion ---

export function modelForCompletion(model: Model<Api>, configuredMaxTokens?: number): Model<Api> {
  const maxTokens = resolveCompletionMaxTokens(model.maxTokens, configuredMaxTokens);
  if (maxTokens === model.maxTokens) return model;
  return { ...model, maxTokens };
}

// --- Reliable streaming ---

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") {
    next = { ...next, transport: "sse" };
  }

  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;
}

// --- Tool host interface ---

export 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;
  toolCallSeq: { value: number };
  abortTurn(): void;
  signal: AbortSignal;
  depth: number;
  pausePending: boolean;
}