Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
15 kB · 402 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403import { join } from "node:path";import { Agent, type AgentMessage, type ThinkingLevel } from "@earendil-works/pi-agent-core";import { clampThinkingLevel, type Message, type Model, streamSimple, supportsFastMode } from "@earendil-works/pi-ai";import { getAgentDir } from "../config.js";import { AgentSession } from "./agent-session.js";import type { AgentSessionCreationOptions } from "./agent-session-services.js";import { formatNoModelsAvailableMessage } from "./auth-guidance.js";import { AuthStorage } from "./auth-storage.js";import type { AgentAutonomousConfig } from "./autonomous.js";import { DEFAULT_THINKING_LEVEL } from "./defaults.js";import type { ExtensionRunner, LoadExtensionsResult, SessionStartEvent, ToolDefinition } from "./extensions/index.js";import { McpManager } from "./mcp/mcp-manager.js";import { convertToLlm } from "./messages.js";import { ModelRegistry } from "./model-registry.js";import { findInitialModel, findSessionModelWithReadinessWait } from "./model-resolver.js";import type { ResourceLoader } from "./resource-loader.js";import { DefaultResourceLoader } from "./resource-loader.js";import { getDefaultSessionDir, SessionManager } from "./session-manager.js";import { SettingsManager } from "./settings-manager.js";import { time } from "./timings.js";import { createBashTool, createEditTool, createIpythonTool, withFileMutationQueue } from "./tools/index.js";
export interface CreateAgentSessionOptions extends AgentSessionCreationOptions { /** Working directory for project-local discovery. Default: process.cwd() */ cwd?: string; /** Global config directory. Default: ~/.pi/agent */ agentDir?: string;
/** Auth storage for credentials. Default: AuthStorage.create(agentDir/auth.json) */ authStorage?: AuthStorage; /** Model registry. Default: ModelRegistry.create(authStorage, agentDir/models.json) */ modelRegistry?: ModelRegistry;
/** Model to use. Default: from settings, else first available */ model?: Model<any>; /** Thinking level. Default: from settings, else 'medium' (clamped to model capabilities) */ thinkingLevel?: ThinkingLevel; /** Models available for cycling (Ctrl+P in interactive mode) */ scopedModels?: Array<{ model: Model<any>; thinkingLevel?: ThinkingLevel }>;
/** * Optional default tool suppression mode when no explicit allowlist is provided. * * - "all": start with no tools enabled * - "builtin": disable the default built-in tool (ipython) * but keep extension/custom tools enabled */ noTools?: "all" | "builtin"; /** * Optional allowlist of tool names. * * When omitted, pi enables the default built-in tool (ipython) * and leaves extension/custom tools enabled unless `noTools` changes that default. * When provided, only the listed tool names are enabled. */ tools?: string[]; /** Custom tools to register (in addition to built-in tools). */ customTools?: ToolDefinition[];
/** Resource loader. When omitted, DefaultResourceLoader is used. */ resourceLoader?: ResourceLoader;
/** MCP integration manager. When omitted, MCP host handlers are not wired. */ mcpManager?: McpManager;
/** Session manager. Default: SessionManager.create(cwd) */ sessionManager?: SessionManager;
/** Settings manager. Default: SettingsManager.create(cwd, agentDir) */ settingsManager?: SettingsManager; /** Session start event metadata for extension runtime startup. */ sessionStartEvent?: SessionStartEvent; /** Host-side autonomous continuation policy. */ autonomous?: AgentAutonomousConfig;}
/** Result from createAgentSession */export interface CreateAgentSessionResult { /** The created session */ session: AgentSession; /** Extensions result (for UI context setup in interactive mode) */ extensionsResult: LoadExtensionsResult; /** Warning if session was restored with a different model than saved */ modelFallbackMessage?: string;}
export type { AgentSessionRuntimeConfig } from "./agent-session-config.js";export * from "./agent-session-runtime.js";export type { AgentSessionCreationOptions } from "./agent-session-services.js";export type { ExtensionAPI, ExtensionCommandContext, ExtensionContext, ExtensionFactory, SlashCommandInfo, SlashCommandSource, ToolDefinition,} from "./extensions/index.js";export type { PromptTemplate } from "./prompt-templates.js";export type { CreateRlmSubagentRuntimeOptions, RlmSubagentRuntime, SubagentRuntimeHost } from "./rlm-runtime.js";export type { Skill } from "./skills.js";export type { Tool } from "./tools/index.js";
export { createBashTool, createEditTool, createIpythonTool, withFileMutationQueue };
function getDefaultAgentDir(): string { return getAgentDir();}
/** * Create an AgentSession with the specified options. * * @example * ```typescript * // Minimal - uses defaults * const { session } = await createAgentSession(); * * // With explicit model * import { getModel } from '@earendil-works/pi-ai'; * const { session } = await createAgentSession({ * model: getModel('anthropic', 'claude-opus-4-5'), * thinkingLevel: 'high', * }); * * // Continue previous session * const { session, modelFallbackMessage } = await createAgentSession({ * continueSession: true, * }); * * // Full control * const loader = new DefaultResourceLoader({ * cwd: process.cwd(), * agentDir: getAgentDir(), * settingsManager: SettingsManager.create(), * }); * await loader.reload(); * const { session } = await createAgentSession({ * model: myModel, * tools: ["ipython"], * resourceLoader: loader, * sessionManager: SessionManager.inMemory(), * }); * ``` */export async function createAgentSession(options: CreateAgentSessionOptions = {}): Promise<CreateAgentSessionResult> { const cwd = options.cwd ?? options.sessionManager?.getCwd() ?? process.cwd(); const agentDir = options.agentDir ?? getDefaultAgentDir(); let resourceLoader = options.resourceLoader;
const authPath = options.agentDir ? join(agentDir, "auth.json") : undefined; const modelsPath = options.agentDir ? join(agentDir, "models.json") : undefined; const authStorage = options.authStorage ?? AuthStorage.create(authPath); const modelRegistry = options.modelRegistry ?? ModelRegistry.create(authStorage, modelsPath);
const settingsManager = options.settingsManager ?? SettingsManager.create(cwd, agentDir); const sessionManager = options.sessionManager ?? SessionManager.create(cwd, getDefaultSessionDir(cwd, agentDir));
// Ensure MCP providers are registered and built-in MCP skills are gated by // auth even on the bare SDK path (not just the CLI's createAgentSessionServices). const mcpManager = options.mcpManager ?? new McpManager({ authStorage, getUserServers: () => settingsManager.getGlobalMcpServers() }); modelRegistry.setOnOAuthProvidersReset(() => mcpManager.registerAllProviders());
if (!resourceLoader) { resourceLoader = new DefaultResourceLoader({ cwd, agentDir, settingsManager, extraBuiltinSkillOverrides: () => mcpManager.getDisabledBuiltinSkillOverrides(), }); await resourceLoader.reload(); time("resourceLoader.reload"); }
const existingSession = sessionManager.buildSessionContext(); const hasExistingSession = existingSession.messages.length > 0; const hasThinkingEntry = sessionManager.branchHasEntryType("thinking_level_change"); const hasServiceTierEntry = sessionManager.branchHasEntryType("service_tier_change");
let model = options.model; let modelFallbackMessage: string | undefined;
if (!model && hasExistingSession && existingSession.model) { // The saved model may be missing only because the Prime Inference // catalog/auth refresh has not settled yet (e.g. right after a daemon // restart). Give the pending refreshes a bounded window, then retry the // lookup once before falling back to another model. const restoredModel = await findSessionModelWithReadinessWait( modelRegistry, existingSession.model.provider, existingSession.model.modelId, ); if (restoredModel) { model = restoredModel; } if (!model) { modelFallbackMessage = `Could not restore model ${existingSession.model.provider}/${existingSession.model.modelId}`; } }
if (!model) { const result = await findInitialModel({ scopedModels: [], isContinuing: hasExistingSession, defaultProvider: settingsManager.getDefaultProvider(), defaultModelId: settingsManager.getDefaultModel(), defaultThinkingLevel: settingsManager.getDefaultThinkingLevel(), modelRegistry, }); model = result.model; if (!model) { modelFallbackMessage = formatNoModelsAvailableMessage(); } else if (modelFallbackMessage) { modelFallbackMessage += `. Using ${model.provider}/${model.id}`; } }
let thinkingLevel = options.thinkingLevel;
if (thinkingLevel === undefined && hasExistingSession) { thinkingLevel = hasThinkingEntry ? (existingSession.thinkingLevel as ThinkingLevel) : (settingsManager.getDefaultThinkingLevel() ?? DEFAULT_THINKING_LEVEL); }
if (thinkingLevel === undefined) { thinkingLevel = settingsManager.getDefaultThinkingLevel() ?? DEFAULT_THINKING_LEVEL; }
if (!model) { thinkingLevel = "off"; } else { thinkingLevel = clampThinkingLevel(model, thinkingLevel) as ThinkingLevel; }
const serviceTierPreference = options.serviceTier ?? (hasServiceTierEntry ? existingSession.serviceTier : settingsManager.getDefaultServiceTier()); const serviceTier = serviceTierPreference === "priority" && (!model || !supportsFastMode(model)) ? "default" : serviceTierPreference;
const allowedToolNames = options.allowedToolNames ?? options.tools ?? (options.noTools === "all" ? [] : undefined); const includeGoals = options.includeGoals ?? (options.tools !== undefined || options.noTools !== "all"); const initialActiveToolNames: string[] = options.initialActiveToolNames ?? (options.tools ? [...options.tools] : options.noTools ? [] : ["ipython"]);
let agent: Agent;
const convertToLlmWithBlockImages = (messages: AgentMessage[]): Message[] => { const converted = convertToLlm(messages); if (!settingsManager.getBlockImages()) { return converted; } return converted.map((msg) => { if (msg.role === "user" || msg.role === "toolResult") { const content = msg.content; if (Array.isArray(content)) { const hasImages = content.some((c) => c.type === "image"); if (hasImages) { const filteredContent = content .map((c) => c.type === "image" ? { type: "text" as const, text: "Image reading is disabled." } : c, ) .filter( (c, i, arr) => !( c.type === "text" && c.text === "Image reading is disabled." && i > 0 && arr[i - 1].type === "text" && (arr[i - 1] as { type: "text"; text: string }).text === "Image reading is disabled." ), ); return { ...msg, content: filteredContent }; } } } return msg; }); };
const extensionRunnerRef: { current?: ExtensionRunner } = {};
agent = new Agent({ initialState: { systemPrompt: "", model, thinkingLevel, serviceTier, tools: [], }, convertToLlm: convertToLlmWithBlockImages, streamFn: async (model, context, options) => { // R15 measures the payload ceiling and publishes it (see `payload_ceiling` in the // usage lookups and the pinned model metadata), but does not refuse here: compaction // runs at the END of a turn, so a turn refused for being over the ceiling could never // reach the compaction that would shrink it. An over-window payload stays the // provider's own error, reported where it happens. const auth = await modelRegistry.getApiKeyAndHeaders(model, options?.headers, sessionManager.getSessionId()); if (!auth.ok) { throw new Error(auth.error); } const providerRetrySettings = settingsManager.getProviderRetrySettings(); const requestModel = auth.requestModel ?? model; return streamSimple(requestModel, context, { ...options, apiKey: auth.apiKey, timeoutMs: options?.timeoutMs ?? providerRetrySettings.timeoutMs, headers: auth.headers, }); }, onPayload: async (payload, _model) => { const runner = extensionRunnerRef.current; if (!runner?.hasHandlers("before_provider_request")) { return payload; } return runner.emitBeforeProviderRequest(payload); }, onResponse: async (response, _model) => { const runner = extensionRunnerRef.current; if (!runner?.hasHandlers("after_provider_response")) { return; } await runner.emit({ type: "after_provider_response", status: response.status, headers: response.headers, }); }, sessionId: sessionManager.getSessionId(), transformContext: async (messages) => { const runner = extensionRunnerRef.current; if (!runner) return messages; return runner.emitContext(messages); }, steeringMode: settingsManager.getSteeringMode(), followUpMode: settingsManager.getFollowUpMode(), transport: settingsManager.getTransport(), thinkingBudgets: settingsManager.getThinkingBudgets(), });
if (hasExistingSession) { agent.state.messages = existingSession.messages; if (!hasThinkingEntry) { sessionManager.appendThinkingLevelChange(thinkingLevel); } } else { if (model) { sessionManager.appendModelChange(model.provider, model.id); } sessionManager.appendThinkingLevelChange(thinkingLevel); } if (!hasServiceTierEntry) { sessionManager.appendServiceTierChange(serviceTierPreference); }
const session = new AgentSession({ agent, sessionManager, settingsManager, serviceTierPreference, cwd, // Only the explicit dir — the default may not match injected custom storage. agentDir: options.agentDir, scopedModels: options.scopedModels, resourceLoader, customTools: options.customTools, modelRegistry, mcpManager, initialActiveToolNames, allowedToolNames, includeGoals, includeCompactSkill: options.includeCompactSkill, rlmHeartbeatController: options.rlmHeartbeatController, agentMessageController: options.agentMessageController, agentObserveController: options.agentObserveController, agentSessionController: options.agentSessionController, extensionRunnerRef, rlmDepth: options.rlmDepth, rlmMaxDepth: options.rlmMaxDepth, rlmSessionDir: options.rlmSessionDir, rlmParentNodeId: options.rlmParentNodeId, rlmParentAgent: options.rlmParentAgent, semanticParentSessionId: options.semanticParentSessionId, semanticSpawnedByRequestId: options.semanticSpawnedByRequestId, subagentRuntimeHost: options.subagentRuntimeHost, sessionStartEvent: options.sessionStartEvent, prewarmIpythonKernel: options.prewarmIpythonKernel, autonomous: options.autonomous, serializedRefine: options.serializedRefine, initialGoal: options.initialGoal, remoteHost: options.remoteHost, }); const extensionsResult = resourceLoader.getExtensions();
return { session, extensionsResult, modelFallbackMessage, };}