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.
13 kB · 344 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345import { join } from "node:path";import type { ThinkingLevel } from "@earendil-works/pi-agent-core";import type { Model, ServiceTier } from "@earendil-works/pi-ai";import { getAgentDir } from "../config.js";import type { AgentSessionMessageController } from "./agent-messages.js";import type { AgentObserveController } from "./agent-observe.js";import type { AgentExecutionMode } from "./agent-session-config.js";import type { AgentSessionController } from "./agent-sessions.js";import { installAgentTraceUpload } from "./agent-traces.js";import { AuthStorage } from "./auth-storage.js";import type { AgentAutonomousConfig } from "./autonomous.js";import type { AgentRlmHeartbeatController } from "./cron-jobs.js";import { createHerdrAgentStateExtension } from "./extensions/builtin/herdr-agent-state.js";import type { SessionStartEvent, ToolDefinition } from "./extensions/index.js";import { McpManager } from "./mcp/mcp-manager.js";import { ModelRegistry } from "./model-registry.js";import { DefaultResourceLoader, type DefaultResourceLoaderOptions, type ResourceLoader } from "./resource-loader.js";import type { SubagentRuntimeHost } from "./rlm-runtime.js";import { type CreateAgentSessionResult, createAgentSession } from "./sdk.js";import { semanticEdgeLedgerPath } from "./semantic-edges.js";import type { SessionManager } from "./session-manager.js";import { SettingsManager } from "./settings-manager.js";import { installAgentTelemetry, isTelemetryEnabled } from "./telemetry.js";import { time } from "./timings.js";
export interface AgentSessionRuntimeDiagnostic { type: "info" | "warning" | "error"; message: string;}
export interface CreateAgentSessionServicesOptions { cwd: string; agentDir?: string; authStorage?: AuthStorage; settingsManager?: SettingsManager; modelRegistry?: ModelRegistry; /** Pre-built MCP manager (tests inject stub probes and stores). */ mcpManager?: McpManager; extensionFlagValues?: Map<string, boolean | string>; resourceLoaderOptions?: Omit<DefaultResourceLoaderOptions, "cwd" | "agentDir" | "settingsManager">; /** * Skip the built-in Herdr reporter for these services. Set for RLM subagent * runtimes: they inherit the parent's HERDR_* pane identity, so their own * reporter would race the parent's on the same pane and a subagent quit * would release the pane while the parent is still running. */ noBuiltinHerdrReporter?: boolean; telemetryDisabled?: true; /** * Hold the telemetry disclosure back on a first interactive launch, where it * would land on the onboarding screen. Onboarding marks itself shown, so the * notice appears on the next launch; sessions that never onboard disclose now. */ deferTelemetryNoticeForOnboarding?: boolean; /** * Start the resource reload without waiting for it. Daemon clients attach to a * session the daemon already owns, so the extensions, skills and themes this * client loads are not needed to show the transcript: waiting for them would * put ~130ms of TypeScript extension loading in front of the first frame. Only * the settings, auth and model registry read before the reload are guaranteed * for a deferred caller. */ deferResourceReload?: boolean;}
export interface AgentSessionCreationOptions { model?: Model<any>; thinkingLevel?: ThinkingLevel; serviceTier?: ServiceTier; scopedModels?: Array<{ model: Model<any>; thinkingLevel?: ThinkingLevel }>; tools?: string[]; noTools?: "all" | "builtin"; customTools?: ToolDefinition[]; initialActiveToolNames?: string[]; allowedToolNames?: string[]; includeGoals?: boolean; includeCompactSkill?: boolean; agentMessageController?: AgentSessionMessageController; agentObserveController?: AgentObserveController; agentSessionController?: AgentSessionController; rlmDepth?: number; rlmMaxDepth?: number; /** Let this session's skills publish without a further authorization (spawn allow_push). */ rlmAllowPush?: boolean; rlmSessionDir?: string; rlmParentNodeId?: string; rlmParentAgent?: string; semanticParentSessionId?: string; semanticSpawnedByRequestId?: string; subagentRuntimeHost?: SubagentRuntimeHost; rlmHeartbeatController?: AgentRlmHeartbeatController; prewarmIpythonKernel?: boolean; autonomous?: AgentAutonomousConfig; serializedRefine?: boolean; executionMode?: AgentExecutionMode; telemetryDisabled?: true; initialGoal?: { objective: string; tokenBudget?: number }; remoteHost?: string;}
export interface CreateAgentSessionFromServicesOptions extends AgentSessionCreationOptions { services: AgentSessionServices; sessionManager: SessionManager; sessionStartEvent?: SessionStartEvent;}
export interface AgentSessionServices { cwd: string; agentDir: string; authStorage: AuthStorage; settingsManager: SettingsManager; modelRegistry: ModelRegistry; resourceLoader: ResourceLoader; mcpManager: McpManager; diagnostics: AgentSessionRuntimeDiagnostic[]; /** * Present only when the extension reload was deferred. Resolves once the * loader has reloaded and provider registrations, flag values and the * resources the loader exposes are all in place. */ resourceReload?: Promise<void>;}
function applyExtensionFlagValues( resourceLoader: ResourceLoader, extensionFlagValues: Map<string, boolean | string> | undefined,): AgentSessionRuntimeDiagnostic[] { if (!extensionFlagValues) { return []; }
const diagnostics: AgentSessionRuntimeDiagnostic[] = []; const extensionsResult = resourceLoader.getExtensions(); const registeredFlags = new Map<string, { type: "boolean" | "string" }>(); for (const extension of extensionsResult.extensions) { for (const [name, flag] of extension.flags) { registeredFlags.set(name, { type: flag.type }); } }
const unknownFlags: string[] = []; for (const [name, value] of extensionFlagValues) { const flag = registeredFlags.get(name); if (!flag) { unknownFlags.push(name); continue; } if (flag.type === "boolean") { extensionsResult.runtime.flagValues.set(name, true); continue; } if (typeof value === "string") { extensionsResult.runtime.flagValues.set(name, value); continue; } diagnostics.push({ type: "error", message: `Extension flag "--${name}" requires a value`, }); }
if (unknownFlags.length > 0) { diagnostics.push({ type: "error", message: `Unknown option${unknownFlags.length === 1 ? "" : "s"}: ${unknownFlags.map((name) => `--${name}`).join(", ")}`, }); }
return diagnostics;}
export async function createAgentSessionServices( options: CreateAgentSessionServicesOptions,): Promise<AgentSessionServices> { const cwd = options.cwd; const agentDir = options.agentDir ?? getAgentDir(); const authStorage = options.authStorage ?? AuthStorage.create(options.agentDir === undefined ? undefined : join(agentDir, "auth.json")); const settingsManager = options.settingsManager ?? SettingsManager.create(cwd, agentDir); time("services.settings"); const modelRegistry = options.modelRegistry ?? ModelRegistry.create(authStorage, join(agentDir, "models.json")); time("services.modelRegistry");
// MCP integrations: registers OAuth providers and gates the built-in // integration skills by whether the user is logged in (enable-by-login). const mcpManager = options.mcpManager ?? new McpManager({ authStorage, getUserServers: () => settingsManager.getGlobalMcpServers(), getCatalogSources: () => settingsManager.getMcpCatalogSources(), }); // refresh() resets the OAuth registry to built-ins; re-add user MCP providers too. modelRegistry.setOnOAuthProvidersReset(() => mcpManager.registerAllProviders());
const userExtensionFactories = options.resourceLoaderOptions?.extensionFactories ?? []; // The built-in Herdr reporter defers to Herdr's own file-based integration // when the loader actually loaded it; two reporters would race on the same // pane. Deferral is late-bound to the loader's loaded paths (inline // factories run after file extensions load), so a file that exists but is // disabled or never discovered does not silence the built-in. // noExtensions is a full opt-out: it disables the built-in reporter too, // not just discovered extension files. const skipHerdrReporter = options.noBuiltinHerdrReporter || options.resourceLoaderOptions?.noExtensions; const builtinExtensionFactories = skipHerdrReporter ? [] : [createHerdrAgentStateExtension(() => resourceLoader.getLoadedExtensionPaths())]; const resourceLoader: DefaultResourceLoader = new DefaultResourceLoader({ ...(options.resourceLoaderOptions ?? {}), extensionFactories: [...builtinExtensionFactories, ...userExtensionFactories], cwd, agentDir, settingsManager, extraBuiltinSkillOverrides: () => mcpManager.getDisabledBuiltinSkillOverrides(), }); const reloadStarted = resourceLoader.reload(); if (!options.deferResourceReload) { await reloadStarted; } time("services.resourceLoader");
const diagnostics: AgentSessionRuntimeDiagnostic[] = []; if ( !options.telemetryDisabled && isTelemetryEnabled(settingsManager) && // A first interactive launch belongs to onboarding, where the notice would // land on the welcome screen; it surfaces on the next launch once // onboarding marks itself shown. Sessions that never onboard disclose now. (settingsManager.getOnboardingShown() || !options.deferTelemetryNoticeForOnboarding) && !settingsManager.getTelemetryNoticeShown() ) { diagnostics.push({ type: "info", message: "Prime Agent sends pseudonymous usage and performance metrics without prompts, responses, tool content, file paths, or repository data. Disable this with telemetry.enabled=false, PRIME_AGENT_TELEMETRY=0, DO_NOT_TRACK=1, or offline mode.", }); settingsManager.setTelemetryNoticeShown(true); } /** Everything that depends on the reloaded loader: providers, flags, errors. */ const applyLoadedResources = async (): Promise<void> => { const extensionsResult = resourceLoader.getExtensions(); for (const { name, config, extensionPath } of extensionsResult.runtime.pendingProviderRegistrations) { try { modelRegistry.registerProvider(name, config); } catch (error) { const message = error instanceof Error ? error.message : String(error); diagnostics.push({ type: "error", message: `Extension "${extensionPath}" error: ${message}`, }); } } extensionsResult.runtime.pendingProviderRegistrations = []; diagnostics.push(...applyExtensionFlagValues(resourceLoader, options.extensionFlagValues)); };
if (options.deferResourceReload) { return { cwd, agentDir, authStorage, settingsManager, modelRegistry, resourceLoader, mcpManager, diagnostics, resourceReload: reloadStarted.then(applyLoadedResources), }; }
await applyLoadedResources();
return { cwd, agentDir, authStorage, settingsManager, modelRegistry, resourceLoader, mcpManager, diagnostics, };}
export async function createAgentSessionFromServices( options: CreateAgentSessionFromServicesOptions,): Promise<CreateAgentSessionResult> { installAgentTraceUpload(options.sessionManager, { authStorage: options.services.authStorage, settingsManager: options.services.settingsManager, semanticEdgesLedgerPath: semanticEdgeLedgerPath({ rlmSessionDir: options.rlmSessionDir, sessionArtifactDir: options.sessionManager.getSessionArtifactDir(), }), }); const result = await createAgentSession({ cwd: options.services.cwd, agentDir: options.services.agentDir, authStorage: options.services.authStorage, settingsManager: options.services.settingsManager, modelRegistry: options.services.modelRegistry, resourceLoader: options.services.resourceLoader, mcpManager: options.services.mcpManager, sessionManager: options.sessionManager, model: options.model, thinkingLevel: options.thinkingLevel, serviceTier: options.serviceTier, scopedModels: options.scopedModels, tools: options.tools, noTools: options.noTools, customTools: options.customTools, initialActiveToolNames: options.initialActiveToolNames, allowedToolNames: options.allowedToolNames, includeGoals: options.includeGoals, includeCompactSkill: options.includeCompactSkill, agentMessageController: options.agentMessageController, agentObserveController: options.agentObserveController, agentSessionController: options.agentSessionController, rlmDepth: options.rlmDepth, rlmMaxDepth: options.rlmMaxDepth, rlmSessionDir: options.rlmSessionDir, rlmParentNodeId: options.rlmParentNodeId, rlmParentAgent: options.rlmParentAgent, semanticParentSessionId: options.semanticParentSessionId, semanticSpawnedByRequestId: options.semanticSpawnedByRequestId, subagentRuntimeHost: options.subagentRuntimeHost, rlmHeartbeatController: options.rlmHeartbeatController, sessionStartEvent: options.sessionStartEvent, prewarmIpythonKernel: options.prewarmIpythonKernel, autonomous: options.autonomous, serializedRefine: options.serializedRefine, initialGoal: options.initialGoal, remoteHost: options.remoteHost, }); if (result.session.rlmDepth === 0 && !options.telemetryDisabled) { installAgentTelemetry(result.session, { agentDir: options.services.agentDir, settingsManager: options.services.settingsManager, executionMode: options.executionMode, }); } return result;}