import type { AgentToolResult, AgentToolUpdateCallback, ExtensionAPI, ExtensionContext, ToolDefinition, } from "@earendil-works/pi-coding-agent"; import { StringEnum } from "@earendil-works/pi-ai"; import { Type, type Static, type TSchema } from "typebox"; import type { ContextManagementMode } from "../adapter/activation/config.ts"; import type { ContextRouter, SharedContextRequest, RemoteNoteReference } from "../context-sharing.ts"; import { contextAccountScope, contextAgentIdentity, contextTargetAgent } from "./agent-identity.ts"; import { assertRemoteBackendScope, bindRemoteBackendScope, isRemoteNestedContext, remoteBackendScope, remoteContextScope, resolveRemoteContextProvider, REMOTE_ACCOUNT_MISMATCH, REMOTE_CONTEXT_UNAVAILABLE } from "./remote-scope.ts"; import { historyNotesRenderers } from "./rendering.ts"; import { projectSharedNotesResult } from "./notes-response-projection.ts"; import { codexToolProviderHeaders } from "../adapter/codex-tool-provider.ts"; import { getPiSessionHistoryRecoveryHint, readPiSessionHistory, } from "./local-history.ts"; import { renderPiSessionNotesThreadHint, usePiSessionNotes, } from "./local-notes.ts"; import { HISTORY_ACTIONS, HISTORY_DESCRIPTION, HISTORY_GUIDANCE, type HistoryAction, NOTES_ACTIONS, NOTES_DESCRIPTION, NOTES_GUIDANCE, type NotesAction, } from "./tool-contract.ts"; const BACKEND_TIMEOUT_MS = 35_000; const THREAD_HINT_MAX_BYTES = 4_000; const TOOL_OUTPUT_TOKEN_LIMIT = 10_000; const HISTORY_ENDPOINTS = { list_windows: "alpha/history/v2/list_windows", list_items: "alpha/history/v2/list_items", read_item: "alpha/history/v2/read_item", search_contents: "alpha/history/v2/search_contents", } as const satisfies Record; const NOTES_ENDPOINTS = { list_files_by_prefix: "alpha/notes/v2/list_files_by_prefix", read_file: "alpha/notes/v2/read_file", search_contents: "alpha/notes/v2/search_contents", append_to_file: "alpha/notes/v2/append_to_file", write_file: "alpha/notes/v2/write_file", } as const satisfies Record; const ENCRYPTED_ARGUMENT_ENDPOINTS = new Set([ HISTORY_ENDPOINTS.search_contents, NOTES_ENDPOINTS.search_contents, NOTES_ENDPOINTS.append_to_file, NOTES_ENDPOINTS.write_file, ]); const HISTORY_ACTION_SET = new Set(HISTORY_ACTIONS); const NOTES_ACTION_SET = new Set(NOTES_ACTIONS); const HISTORY_ACTION_FIELDS = { list_windows: ["agent_name", "limit", "recent_first"], list_items: [ "agent_name", "limit", "max_chars_per_item", "recent_first", "role", "tool_name", "tool_namespace", "window_id", ], read_item: [ "agent_name", "item_id", "limit_chars", "offset_chars", "window_id", ], search_contents: [ "agent_name", "limit", "query", "recent_first", "role", "tool_name", "tool_namespace", "window_id", ], } satisfies Record; const NOTES_ACTION_FIELDS = { list_files_by_prefix: ["file_order", "file_order_by", "max_results", "prefix"], read_file: ["path", "start_line", "stop_line"], search_contents: [ "max_files", "max_matches_per_file", "path_prefix", "query", "recent_file_first", ], append_to_file: ["path", "text"], write_file: ["path", "text"], } satisfies Record; const HISTORY_PARAMETERS = Type.Object( { action: Type.Optional(StringEnum(["help", ...HISTORY_ACTIONS])), agent_name: Type.Optional(Type.Union([Type.String(), Type.Null()])), item_id: Type.Optional(Type.String()), limit: Type.Optional(Type.Integer({ minimum: 1 })), limit_chars: Type.Optional(Type.Integer({ minimum: 1 })), max_chars_per_item: Type.Optional(Type.Integer({ minimum: 1 })), offset_chars: Type.Optional(Type.Integer({ minimum: 0 })), query: Type.Optional(Type.String()), recent_first: Type.Optional(Type.Boolean()), role: Type.Optional( Type.Union([ StringEnum([ "user", "assistant", "tool", "system", "developer", ] as const), Type.Null(), ]), ), tool_name: Type.Optional(Type.Union([Type.String(), Type.Null()])), tool_namespace: Type.Optional( Type.Union([Type.String(), Type.Null()]), ), window_id: Type.Optional(Type.Union([Type.String(), Type.Null()])), }, { additionalProperties: false }, ); const NOTES_PARAMETERS = Type.Object( { action: Type.Optional(StringEnum(["help", ...NOTES_ACTIONS])), file_order: Type.Optional( StringEnum(["ascending", "descending"] as const), ), file_order_by: Type.Optional( StringEnum(["name", "created_at", "updated_at"] as const), ), max_files: Type.Optional(Type.Integer({ minimum: 1 })), max_matches_per_file: Type.Optional(Type.Integer({ minimum: 1 })), max_results: Type.Optional(Type.Integer({ minimum: 1 })), path: Type.Optional(Type.String()), path_prefix: Type.Optional(Type.Union([Type.String(), Type.Null()])), prefix: Type.Optional(Type.Union([Type.String(), Type.Null()])), query: Type.Optional(Type.String()), recent_file_first: Type.Optional(Type.Boolean()), start_line: Type.Optional(Type.Union([Type.Integer(), Type.Null()])), stop_line: Type.Optional(Type.Union([Type.Integer(), Type.Null()])), text: Type.Optional(Type.String()), }, { additionalProperties: false }, ); export interface CodexHistoryNotesDetails { codexHistoryNotes: Record; } type HistoryNotesTool = Omit, "execute"> & { execute( toolCallId: string, params: Static, signal: AbortSignal | undefined, onUpdate: AgentToolUpdateCallback | undefined, ctx: ExtensionContext, ): Promise>; }; export function createHistoryNotesTools( pi?: Pick, resolveMode: (ctx: ExtensionContext) => ContextManagementMode = () => "local", prepareNoteWrite?: (action: NotesAction, path: unknown, ctx: ExtensionContext) => () => boolean, route?: ContextRouter, ): [HistoryNotesTool, HistoryNotesTool] { return [ { name: "history", label: "history", description: HISTORY_DESCRIPTION, parameters: HISTORY_PARAMETERS, ...historyNotesRenderers("history"), prepareArguments(args) { // Ignore extra search fields without advertising them or forwarding them to storage. if (args && typeof args === "object" && !Array.isArray(args) && "action" in args && args.action === "search_contents") { args = Object.fromEntries(Object.entries(args).filter(([field]) => field === "action" || HISTORY_ACTION_FIELDS.search_contents.includes(field), )); } // Pi validates the prepared value against parameters before execution. return args as Static; }, async execute(_id, params, signal, _update, ctx: ExtensionContext) { if (params.action === undefined || params.action === "help") return historyNotesHelp("history"); const action = historyAction(params.action); validateHistoryArguments(action, params); return callHistoryNotesTool( "history", action, HISTORY_ENDPOINTS[action], params, ctx, signal, resolveMode(ctx), pi, route, _id, ); }, }, { name: "notes", label: "notes", description: NOTES_DESCRIPTION, parameters: NOTES_PARAMETERS, ...historyNotesRenderers("notes"), executionMode: "sequential", prepareArguments(args) { if (args && typeof args === "object" && !Array.isArray(args) && "action" in args && (args.action === "write_file" || args.action === "append_to_file") && "content" in args) { const { content, ...rest } = args; args = { ...rest, text: "text" in rest ? rest.text : content }; } if (args && typeof args === "object" && !Array.isArray(args) && "action" in args && args.action === "list_files_by_prefix" && "max_files" in args) { const { max_files, ...rest } = args; args = { ...rest, max_results: "max_results" in rest ? rest.max_results : max_files }; } return args as Static; }, async execute(_id, params, signal, _update, ctx: ExtensionContext) { if (params.action === undefined || params.action === "help") return historyNotesHelp("notes"); const action = notesAction(params.action); validateNotesArguments(action, params); const finishNoteWrite = action === "write_file" || action === "append_to_file" ? prepareNoteWrite?.(action, params.path, ctx) : undefined; const result = await callHistoryNotesTool( "notes", action, NOTES_ENDPOINTS[action], params, ctx, signal, resolveMode(ctx), pi, route, _id, ); return finishNoteWrite?.() ? { ...result, terminate: true } : result; }, }, ]; } function historyNotesHelp(namespace: "history" | "notes"): AgentToolResult { return historyNotesResult(namespace, { usage: `${namespace}({ action, ...args }); no arguments or action:help returns usage`, actions: namespace === "history" ? HISTORY_ACTION_FIELDS : NOTES_ACTION_FIELDS, required: namespace === "history" ? { read_item: ["item_id", "window_id"], search_contents: ["query"] } : { read_file: ["path"], search_contents: ["query"], append_to_file: ["path", "text"], write_file: ["path", "text"] }, schema: namespace === "history" ? HISTORY_PARAMETERS : NOTES_PARAMETERS, guidance: namespace === "history" ? HISTORY_GUIDANCE : NOTES_GUIDANCE, }); } export async function loadHistoryNotesThreadHint( ctx: ExtensionContext, mode: ContextManagementMode, signal?: AbortSignal, ): Promise { if (mode === "local" || mode === "tree") return piSessionThreadHint(ctx, mode); if (!usesRemoteHistoryNotes(ctx, mode)) return undefined; try { const result = await callHistoryNotesBackend( "alpha/notes/v2/thread_hint", {}, ctx, signal, { mode: "bytes", limit: THREAD_HINT_MAX_BYTES }, ); const text = typeof result["text"] === "string" ? result["text"] : ""; return text && Buffer.byteLength(text, "utf8") <= THREAD_HINT_MAX_BYTES ? text : undefined; } catch (error) { if (signal?.aborted) throw error; return undefined; } } function piSessionThreadHint( ctx: ExtensionContext, mode: ContextManagementMode, ): string | undefined { const recovery = getPiSessionHistoryRecoveryHint(ctx, mode); const recoveryHint = recovery ? `Previous window history IDs: ${JSON.stringify(recovery)}` : undefined; const recoveryBytes = recoveryHint ? Buffer.byteLength(`\n${recoveryHint}`, "utf8") : 0; const notesHint = renderPiSessionNotesThreadHint( ctx.sessionManager.getBranch(), Math.max(0, THREAD_HINT_MAX_BYTES - recoveryBytes), ); const hint = [notesHint, recoveryHint].filter( (value): value is string => Boolean(value), ).join("\n"); return hint && Buffer.byteLength(hint, "utf8") <= THREAD_HINT_MAX_BYTES ? hint : undefined; } async function callHistoryNotesTool( namespace: "history" | "notes", action: HistoryAction | NotesAction, endpoint: string, params: Record, ctx: ExtensionContext, signal: AbortSignal | undefined, mode: ContextManagementMode, pi: Pick | undefined, route: ContextRouter | undefined, callId: string, ): Promise> { const identity = contextAgentIdentity(ctx); if (identity.storage && identity.storage !== (mode === "remote" ? "remote" : "session")) throw new Error("Shared notes are unavailable with this storage. Use messages to exchange the context you need."); const target = contextTargetAgent(namespace, params, identity.agentName); if (target !== identity.agentName && route) { const routed = await route(ctx, { sessionId: identity.sessionId, agentName: target, namespace, params: namespace === "history" ? { ...params, agent_name: target } : params, encryptedArguments: mode === "remote" && ENCRYPTED_ARGUMENT_ENDPOINTS.has(endpoint) && !isRemoteNestedContext(ctx), }, signal); if (routed) return projectSharedNotesResult(routed, callId, ctx, signal); } let result: Record; if (mode === "remote") { if (!usesRemoteHistoryNotes(ctx, mode)) throw new Error(REMOTE_CONTEXT_UNAVAILABLE); result = await callHistoryNotesBackend( endpoint, stripAction(params), ctx, signal, { mode: "tokens", limit: TOOL_OUTPUT_TOKEN_LIMIT }, ); } else { if (target !== identity.agentName) { throw new Error("Cross-agent context router is unavailable"); } result = callLocalHistoryNotes(namespace, action, params, ctx, pi, mode); } return historyNotesResult(namespace, result, action); } function historyNotesResult(namespace: "history" | "notes", result: Record, action?: HistoryAction | NotesAction): AgentToolResult { const modelResult = { ...result }; const backendScope = remoteBackendScope(result); if (backendScope) bindRemoteBackendScope(modelResult, backendScope); delete modelResult["images"]; // Local freshness consumes source in details, not in the model's write receipt. const visibleResult = namespace === "notes" && modelResult["source"] === "pi-session" && (action === "write_file" || action === "append_to_file") ? { file: modelResult["file"] } : modelResult; const content: AgentToolResult["content"] = [ { type: "text", text: typeof modelResult["encrypted_output"] === "string" ? `${namespace} operation completed` : JSON.stringify(visibleResult), }, ]; for (const image of parseBackendImages(result["images"])) content.push(image); return { content, details: { codexHistoryNotes: modelResult }, }; } async function callHistoryNotesBackend( endpoint: string, arguments_: Record, ctx: ExtensionContext, signal: AbortSignal | undefined, truncationPolicy: { mode: "bytes" | "tokens"; limit: number }, target?: RemoteNoteReference, ordinaryArguments = false, capture?: (bytes: Uint8Array, encoding: string, scope: string) => void, ): Promise> { const before = JSON.stringify(contextAgentIdentity(ctx)); const model = JSON.stringify([ctx.model?.api, ctx.model?.provider, ctx.model?.id, ctx.model?.baseUrl]); const provider = await resolveRemoteContextProvider(ctx); const identity = contextAgentIdentity(ctx); if (before !== JSON.stringify(identity) || model !== JSON.stringify([ctx.model?.api, ctx.model?.provider, ctx.model?.id, ctx.model?.baseUrl])) throw new Error("Context changed during authentication; retry in the original session"); if (identity.accountScope && contextAccountScope(provider.accountId) !== identity.accountScope) throw new Error(REMOTE_ACCOUNT_MISMATCH); if (identity.backendUrl && identity.backendUrl !== provider.baseUrl) throw new Error(REMOTE_CONTEXT_UNAVAILABLE); if (target && target.identity.accountScope !== contextAccountScope(provider.accountId)) throw new Error(REMOTE_ACCOUNT_MISMATCH); if (target && target.baseUrl !== provider.baseUrl) throw new Error(REMOTE_CONTEXT_UNAVAILABLE); const scope = remoteContextScope(identity, provider.accountId, provider.baseUrl); assertRemoteBackendScope(ctx, scope); const headers = codexToolProviderHeaders(provider); headers.set( "x-openai-tool-output-truncation-policy", JSON.stringify(truncationPolicy), ); if (ENCRYPTED_ARGUMENT_ENDPOINTS.has(endpoint) && !isRemoteNestedContext(ctx) && !ordinaryArguments) headers.set("x-openai-encrypted-tool-arguments", "true"); const timeoutSignal = AbortSignal.timeout(BACKEND_TIMEOUT_MS); const response = await fetch( `${provider.baseUrl.replace(/\/+$/, "")}/${endpoint}`, { method: "POST", headers, signal: signal ? AbortSignal.any([signal, timeoutSignal]) : timeoutSignal, body: JSON.stringify({ ...arguments_, context: { session_id: target?.identity.sessionId ?? identity.sessionId, current_agent_name: target?.identity.agentName ?? identity.agentName, }, }), }, ); if (!response.ok) throw new Error(`History and notes backend failed (${response.status})`); const bytes = new Uint8Array(await response.arrayBuffer()); const result: unknown = JSON.parse(new TextDecoder().decode(bytes)); if (!result || typeof result !== "object" || Array.isArray(result)) throw new Error("History and notes backend returned invalid data"); if (before !== JSON.stringify(contextAgentIdentity(ctx)) || model !== JSON.stringify([ctx.model?.api, ctx.model?.provider, ctx.model?.id, ctx.model?.baseUrl])) throw new Error("Context changed after Remote execution; verify note state before repeating a write"); bindRemoteBackendScope(result, scope); capture?.(bytes, response.headers.get("content-type") ?? "application/json", scope); return result as Record; } /** Read-only bridge access captures the wire body before envelope decoding. */ export async function lookupRemoteHistoryNotes( namespace: "history" | "notes", params: Record, ctx: ExtensionContext, signal?: AbortSignal, ): Promise<{ bytes: Uint8Array; encoding: string; source: string }> { const action = namespace === "history" ? historyAction(params["action"]) : notesAction(params["action"]); if (action === "write_file" || action === "append_to_file") throw new Error("Remote lookup is read-only"); if (namespace === "history") validateHistoryArguments(action as HistoryAction, params); else validateNotesArguments(action as NotesAction, params); const endpoint = namespace === "history" ? HISTORY_ENDPOINTS[action as HistoryAction] : NOTES_ENDPOINTS[action as NotesAction]; let captured: { bytes: Uint8Array; encoding: string; source: string } | undefined; await callHistoryNotesBackend(endpoint, stripAction(params), ctx, signal, { mode: "tokens", limit: TOOL_OUTPUT_TOKEN_LIMIT }, undefined, true, (bytes, encoding, source) => { captured = { bytes, encoding, source }; }); if (!captured) throw new Error("Remote lookup returned no response"); return captured; } /** The caller authenticates and owns protected delivery; the reference owns backend addressing. */ export async function executeRemoteAttachment( ctx: ExtensionContext, reference: RemoteNoteReference, request: SharedContextRequest, signal?: AbortSignal, ): Promise> { signal?.throwIfAborted(); if ((request.namespace !== "history" && request.namespace !== "notes") || !request.params || typeof request.params !== "object" || Array.isArray(request.params)) throw new Error("Invalid Remote attachment request"); const action = request.namespace === "history" ? historyAction(request.params["action"]) : notesAction(request.params["action"]); if (request.sessionId !== reference.identity.sessionId || (request.agentName !== reference.identity.agentName && !request.agentName.startsWith(`${reference.identity.agentName}/`)) || contextTargetAgent(request.namespace, request.params, reference.identity.agentName) !== request.agentName) throw new Error("Remote attachment request does not belong to its checkpoint owner"); if (request.namespace === "history") validateHistoryArguments(action as HistoryAction, request.params); else validateNotesArguments(action as NotesAction, request.params); const endpoint = request.namespace === "history" ? HISTORY_ENDPOINTS[action as HistoryAction] : NOTES_ENDPOINTS[action as NotesAction]; if (ENCRYPTED_ARGUMENT_ENDPOINTS.has(endpoint) && !isRemoteNestedContext(ctx) && !request.encryptedArguments) throw new Error("Use Code or Notebook history/notes with ordinary query/text for this attached context"); return historyNotesResult(request.namespace, await callHistoryNotesBackend(endpoint, stripAction(request.params), ctx, signal, { mode: "tokens", limit: TOOL_OUTPUT_TOKEN_LIMIT }, reference)); } export function usesRemoteHistoryNotes( ctx: Pick, mode: ContextManagementMode, ): boolean { return mode === "remote" && (ctx.model?.api ?? "").trim().toLowerCase() === "openai-codex-responses"; } function callLocalHistoryNotes( namespace: "history" | "notes", action: HistoryAction | NotesAction, params: Record, ctx: ExtensionContext, pi: Pick | undefined, mode: ContextManagementMode, ): Record { if (namespace === "history") return readPiSessionHistory(action as HistoryAction, params, ctx, mode); if (!pi) throw new Error("Local notes require an active Pi session"); return usePiSessionNotes(pi, action as NotesAction, params, ctx); } function stripAction(params: Record): Record { const result = { ...params }; delete result["action"]; delete result["context"]; return result; } function historyAction(value: unknown): HistoryAction { if (typeof value === "string" && HISTORY_ACTION_SET.has(value)) return value as HistoryAction; throw new Error("history requires a supported action"); } function notesAction(value: unknown): NotesAction { if (typeof value === "string" && NOTES_ACTION_SET.has(value)) return value as NotesAction; throw new Error("notes requires a supported action"); } function validateActionFields( namespace: "history" | "notes", action: HistoryAction | NotesAction, params: Record, allowed: readonly string[], ): void { const unexpected = Object.keys(params).find( (field) => field !== "action" && !allowed.includes(field), ); if (unexpected) throw new Error(`${namespace} ${action} does not accept ${unexpected}`); } function validateHistoryArguments( action: HistoryAction, params: Record, ): void { validateActionFields("history", action, params, HISTORY_ACTION_FIELDS[action]); if (action === "read_item") { if (typeof params["item_id"] !== "string" || !params["item_id"]) throw new Error("history read_item requires item_id"); if (typeof params["window_id"] !== "string" || !params["window_id"]) throw new Error("history read_item requires window_id"); } if ( action === "search_contents" && (typeof params["query"] !== "string" || !params["query"]) ) throw new Error("history search_contents requires query"); } function validateNotesArguments( action: NotesAction, params: Record, ): void { validateActionFields("notes", action, params, NOTES_ACTION_FIELDS[action]); if ( (action === "read_file" || action === "append_to_file" || action === "write_file") && (typeof params["path"] !== "string" || !params["path"]) ) throw new Error(`notes ${action} requires path`); if ( (action === "search_contents" && (typeof params["query"] !== "string" || !params["query"])) || ((action === "append_to_file" || action === "write_file") && typeof params["text"] !== "string") ) throw new Error( `notes ${action} requires ${action === "search_contents" ? "query" : "text"}`, ); } function parseBackendImages( value: unknown, ): Array<{ type: "image"; data: string; mimeType: string; detail?: "auto" | "high" | "original" | undefined; }> { if (value === undefined) return []; if (!Array.isArray(value)) throw new Error("History backend returned invalid image content"); return value.map((item) => { if (!item || typeof item !== "object") throw new Error("History backend returned invalid image content"); const image = item as Record; if ( typeof image["data"] !== "string" || typeof image["mime_type"] !== "string" ) throw new Error("History backend returned invalid image content"); const detail = image["detail"]; if ( detail !== undefined && detail !== null && detail !== "auto" && detail !== "high" && detail !== "original" ) throw new Error("History backend returned invalid image detail"); return { type: "image" as const, data: image["data"], mimeType: image["mime_type"], ...(detail ? { detail } : {}), }; }); }