diff --git a/packages/coding-agent/src/README.md b/packages/coding-agent/src/README.md index ab7a526a48..63072ebb26 100644 --- a/packages/coding-agent/src/README.md +++ b/packages/coding-agent/src/README.md @@ -108,6 +108,42 @@ Preserve these boundaries when changing context behavior: - A cancelled continuation settles its own waiters. Late results cannot clear a replacement operation or consume its messages. The commit lease is released before awaiting a continued agent turn, and checkpoint waiters are removed when cancellation wins. - Public session methods delegate without additional asynchronous wrappers. Optional waits retain their original positions, and dependency callbacks read current runtime state when invoked. +## Child agent lifecycle + +| File | Responsibility | +| --- | --- | +| `session/children.ts` | Child registry, admission, publication, deletion retries, cancellation, quiescence, retention, and cleanup. | +| `session/child-run.ts` | Detached child execution, publication barriers, terminal state, and event attribution. | +| `session/child-runtime.ts` | Inline child construction and child-specific directory creation. | +| `session/child-state.ts` | Depth and maximum-depth settings, parent replies, and recap state. | +| `session/child-usage.ts` | Child usage attribution, origin batches, flush timers, and retry bookkeeping. | +| `session/child-projection.ts` | Read-only child list and snapshot projections. | +| `session/child-types.ts` | Child contracts and shared child data helpers. | + +The registry owns child identity and lifecycle transitions. Execution and usage components operate on the same child records; they do not create competing copies of run state. Child state exposes read-only properties and named mutations. Runtime hosts, inherited depth, model selection, and event queues are read through live operations supplied by the session. + +- Reserve and publish children in the original order. A late completion cannot replace a newer run or clear another run's cancellation state. +- Retain deletion reservations, retryable cleanup state, and descendant quiescence until their existing completion conditions hold. Parent continuation still waits for the appropriate child work. +- Flush child usage once at the existing parent event boundary. Origin batches and timers have one cleanup owner. +- Complete child cleanup before kernel teardown. The session supplies the following teardown operation so an empty child set does not add a scheduling delay before kernel disposal begins. +- Keep calls that previously passed through public session methods live, including descendant receivers, registration, deletion, and maximum-depth status after settings updates. + +The root kernel directory belongs to `session/kernel-environment.ts`. Child directory construction receives a lazy operation for that directory rather than keeping a second root-directory field. The existing child runtime-options factory retains its public parent-session contract at the facade; child owners receive only the operations they use. + +## Tools, extensions, and kernel resources + +| File | Responsibility | +| --- | --- | +| `session/tools.ts` | Tool definitions and active selection, allowlists, prompt contributions, and ACP tool updates. | +| `session/extensions.ts` | Extension runner bindings, resource reload, and extension lifecycle. | +| `session/kernel.ts` | Kernel construction, snapshot restoration, prewarming, and disposal. | +| `session/kernel-environment.ts` | Kernel provisioning environment and root or ephemeral session directories. | +| `session/kernel-host-handlers.ts` | Typed host-handler composition from live session operations. | + +The session coordinates these owners with children, models, and input admission. A kernel replacement uses the previous kernel's disposal promise as its readiness gate. First-build restoration notices and snapshot-directory ownership stay with the kernel owner. Host handlers read the current runtime when invoked, including after replacement. + +ACP resource cleanup retains its input pause until queued work and cleanup finish, and releases the pause on failure. Extension bindings preserve public session dispatch and callback receivers, including shutdown and partial rebinding. Pure facade delegates do not add asynchronous wrappers around already asynchronous owner operations. + ## Validation Controller tests live in `test/goals/`. Session integration coverage remains in `test/suite/agent-session-goal.test.ts`, `test/suite/agent-session-compaction-continuation.test.ts`, and `test/goal-continuation-quiescence.test.ts`. @@ -118,4 +154,6 @@ Shell-owner tests also live in `test/session/`. Session bash/persistence, prompt Compaction and continuation owner tests in `test/session/` exercise lifecycle and cancellation boundaries. Compaction, refinement, serialized refinement, queue, concurrency, and semantic-edge suites cover their integration with session persistence, goals, and disposal. `test/suite/session-refinement-owner.test.ts` checks the refinement owner boundary using the shared faux-provider harness. +Child usage and recursion suites cover accounting, cancellation, publication, and cleanup. Kernel, environment, and tool tests in `test/session/` cover resource ownership and replacement. The child/runtime facade regressions in `test/suite/regressions/` preserve public dispatch, callback receivers, and teardown ordering. Real Python background-bash cases require the configured runtime environment; record missing-environment skips explicitly. + Run the focused files from the coding-agent package root with the repository's prescribed Vitest command, then run `npm run check` from the repository root. Use faux providers for session tests. diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 0e11c81250..8942dfbf11 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -1,18 +1,17 @@ import { AsyncLocalStorage } from "node:async_hooks"; import { randomUUID } from "node:crypto"; -import { existsSync, mkdirSync, mkdtempSync, readFileSync, writeFileSync } from "node:fs"; -import { tmpdir } from "node:os"; -import { basename, dirname, join, resolve } from "node:path"; -import { +import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs"; +import { dirname, resolve } from "node:path"; +import type { Agent, - type AgentContext, - type AgentEvent, - type AgentMessage, - type AgentState, - type AgentTool, - type GetContinuationMessagesContext, - type ShouldStopAfterTurnContext, - type ThinkingLevel, + AgentContext, + AgentEvent, + AgentMessage, + AgentState, + AgentTool, + GetContinuationMessagesContext, + ShouldStopAfterTurnContext, + ThinkingLevel, } from "@earendil-works/pi-agent-core"; import type { Api, @@ -29,7 +28,6 @@ import { cleanupSessionResources, getSupportedThinkingLevels, modelsAreEqual, - resetApiProviders, supportsFastMode, } from "@earendil-works/pi-ai"; import { parseGoalSlashCommand } from "../goals/commands.js"; @@ -42,6 +40,16 @@ import { SessionBash, type SessionBashEvent, } from "../session/bash.js"; +import { createChildSessionDir, createInlineChildRuntime } from "../session/child-runtime.js"; +import { SessionChildState } from "../session/child-state.js"; +import { + compactRlmText, + type RlmChildAgentSnapshot, + type RlmChildAgentStatus, + rlmChildLabel, +} from "../session/child-types.js"; +import { SessionChildUsage } from "../session/child-usage.js"; +import { SessionChildren } from "../session/children.js"; import { SessionCommitFence, type SessionCommitLease } from "../session/commit-fence.js"; import { SessionCompaction, type SessionCompactionEvent } from "../session/compaction.js"; import { @@ -51,8 +59,12 @@ import { performSessionCompaction, } from "../session/compaction-execution.js"; import { type ContinuationToken, SessionContinuation } from "../session/continuation.js"; +import { type ExtensionBindings, installExtensionToolHooks, SessionExtensions } from "../session/extensions.js"; import { SessionInputDispatcher } from "../session/input-dispatcher.js"; import { SessionInputScheduler } from "../session/input-scheduler.js"; +import { SessionKernel } from "../session/kernel.js"; +import { KernelEnvironment } from "../session/kernel-environment.js"; +import { createSessionKernelHostHandlers } from "../session/kernel-host-handlers.js"; import { buildPromptContent, cloneCustomMessage, @@ -78,25 +90,18 @@ import { } from "../session/prepared-actions.js"; import { type AutoRefineReviewer, SessionRefinement } from "../session/refinement.js"; import { SessionRetry, type SessionRetryEvent } from "../session/retry.js"; +import { SessionTools } from "../session/tools.js"; import { createTurnExecutionPolicy, type TurnExecutionPolicy, TurnPreparer } from "../session/turn-preparation.js"; import { stripFrontmatter } from "../utils/frontmatter.js"; import { waitForPromiseOrAbort } from "../utils/wait-for-abort.js"; import { - AGENT_MESSAGE_CUSTOM_TYPE, AGENT_MESSAGE_SKILL_NAME, - type AgentFamilyCatalogEntry, type AgentSessionMessage, - type AgentSessionMessageAgentSummary, type AgentSessionMessageController, - type AgentSessionMessageListResult, type AgentSessionMessageReceipt, - agentFamilyMemberName, assertAgentMessageQueueCapacity, - assertAgentSessionNameAvailable, assertDirectAgentMessageTarget, - createAgentMessageHostHandlers, DEFAULT_AGENT_MESSAGE_MAX_PENDING_PER_SESSION, - formatAgentSessionNameUnavailable, isAgentSessionMessage, isAgentSessionMessagePrompt, normalizeAgentSessionMessage, @@ -109,7 +114,6 @@ import { type AgentObserveController, type AgentObserveListResult, type AgentObserveRecentMessagesResult, - createAgentObserveHostHandlers, normalizeObserveLimit, normalizeObserveMaxChars, ORCHESTRATION_HEARTBEAT_SKILL_NAME, @@ -159,31 +163,25 @@ import { normalizeHeartbeatDeliveryMode } from "./cron-jobs.js"; import { DEFAULT_THINKING_LEVEL } from "./defaults.js"; import { exportSessionToHtml, type ToolHtmlRenderer } from "./export-html/index.js"; import { createToolHtmlRenderer } from "./export-html/tool-renderer.js"; -import { - type ContextUsage, - type ExtensionCommandContextActions, - type ExtensionErrorListener, +import type { + ContextUsage, ExtensionRunner, - type ExtensionUIContext, - type InputSource, - type MessageEndEvent, - type MessageStartEvent, - type MessageUpdateEvent, - type ReplacedSessionContext, - type SessionBeforeTreeResult, - type SessionStartEvent, - type ShutdownHandler, - type ToolDefinition, - type ToolExecutionEndEvent, - type ToolExecutionStartEvent, - type ToolExecutionUpdateEvent, - type ToolInfo, - type TreePreparation, - type TurnEndEvent, - type TurnStartEvent, - wrapRegisteredTools, + InputSource, + MessageEndEvent, + MessageStartEvent, + MessageUpdateEvent, + ReplacedSessionContext, + SessionBeforeTreeResult, + SessionStartEvent, + ToolDefinition, + ToolExecutionEndEvent, + ToolExecutionStartEvent, + ToolExecutionUpdateEvent, + ToolInfo, + TreePreparation, + TurnEndEvent, + TurnStartEvent, } from "./extensions/index.js"; -import { emitSessionShutdownEvent } from "./extensions/runner.js"; import { createGoalContextMessage, GOAL_CONTEXT_CUSTOM_TYPE, @@ -196,26 +194,22 @@ import { validateGoalObjective, } from "./goals.js"; import type { HostRequestHandlers, KernelSentAgentMessage } from "./kernel/index.js"; -import { type RestoreResult, snapshotPathIn } from "./kernel/state-snapshot.js"; import type { AcpMcpServerConfig } from "./mcp/acp-mcp-types.js"; import type { McpManager } from "./mcp/mcp-manager.js"; +import type { AsyncBashCompletionDetails } from "./messages.js"; import { ASYNC_BASH_COMPLETION_CUSTOM_TYPE, ASYNC_BASH_COMPLETION_PREVIEW_LABEL, - type AsyncBashCompletionDetails, type CustomMessage, createAsyncBashCompletionMessage, createHarnessDigestMessage, createHeartbeatPromptMessage, - createRlmChildFailureMessage, - createRlmChildTerminalNoticeMessage, createSessionSlashCommandMessage, createSessionSlashCommandResultMessage, HARNESS_DIGEST_CUSTOM_TYPE, type HarnessDigestDetails, HEARTBEAT_PROMPT_CUSTOM_TYPE, HEARTBEAT_PROMPT_PREVIEW_LABEL, - IPYTHON_STATE_RESTORED_CUSTOM_TYPE, isSessionSlashCommandMessage, type RefinementSource, RLM_CHILD_FAILURE_CUSTOM_TYPE, @@ -225,35 +219,16 @@ import type { ModelRegistry } from "./model-registry.js"; import { throwIfPromptAdmissionCancelled } from "./prompt-admission.js"; import { expandPromptTemplate, type PromptTemplate, parseCommandArgs } from "./prompt-templates.js"; import { providerRetryPolicy } from "./provider-retry.js"; -import { - formatHarnessStateForPrompt, - getGlobalHarnessStateDir, - getLocalHarnessStateDir, - REFINE_SKILL_NAME, - type RefinementResult, -} from "./refinement/index.js"; -import { resolveConfigValue } from "./resolve-config-value.js"; -import type { ResourceExtensionPaths, ResourceLoader } from "./resource-loader.js"; +import { formatHarnessStateForPrompt, REFINE_SKILL_NAME, type RefinementResult } from "./refinement/index.js"; +import type { ResourceLoader } from "./resource-loader.js"; import { type CreateRlmSubagentRuntimeOptions, - createAsyncBashCompletionHostHandler, - createAsyncBashConsumedHostHandler, - createDefaultRlmSubagentSessionName, - createRlmCreateSessionHostHandler, - createRlmDeleteSubagentHostHandler, - createRlmFindModelsHostHandler, - createRlmListSubagentsHostHandler, - createRlmRunHostHandler, findRlmModelMatches, - normalizeRequestedRlmSubagentModel, - normalizeRequestedRlmSubagentSessionName, - normalizeRequestedRlmSubagentThinkingLevel, type RlmCreateSessionResult, type RlmDeleteSubagentResult, type RlmFindModelsResult, type RlmListSubagentsResult, type RlmSpawnHandle, - type RlmSubagentRegistryEntry, type RlmSubagentRuntime, type SubagentRuntimeHost, } from "./rlm-runtime.js"; @@ -273,18 +248,12 @@ import { type SessionActionSnapshot, transitionSessionAction, } from "./session-action-store.js"; -import type { - BranchSummaryEntry, - ChildUsageAttributionEntry, - SessionContext, - SessionEntry, - SessionMessageEntry, -} from "./session-manager.js"; +import type { BranchSummaryEntry, SessionContext, SessionEntry } from "./session-manager.js"; import { CURRENT_SESSION_VERSION, getLatestCompactionEntry, type SessionHeader, - SessionManager, + type SessionManager, } from "./session-manager.js"; import type { SessionStats } from "./session-stats.js"; import type { SettingsManager } from "./settings-manager.js"; @@ -294,57 +263,19 @@ import { parseSessionSlashCommand, parseSlashCommand, type SessionSlashCommand, - type SlashCommandInfo, } from "./slash-commands.js"; -import { createSyntheticSourceInfo, type SourceInfo } from "./source-info.js"; -import { type BuildSystemPromptOptions, buildSystemPrompt } from "./system-prompt.js"; +import type { BuildSystemPromptOptions } from "./system-prompt.js"; import { THINKING_LEVELS } from "./thinking-levels.js"; -import { acpMcpToolNames, createAcpMcpToolDefinitions } from "./tools/acp-mcp.js"; -import { createAllToolDefinitions } from "./tools/index.js"; -import { IpythonKernelProvisioner } from "./tools/ipython.js"; -import { createToolDefinitionFromAgentTool } from "./tools/tool-definition-wrapper.js"; -import { - addAssistantUsage, - cloneUsage, - emptyUsage, - type SessionUsageSummary, - sessionUsageSummaryFrom, - subtractAssistantUsage, -} from "./usage.js"; -import { SERPER_CREDENTIAL_ID, SERPER_ENV_VAR, WEBSEARCH_SKILL_NAME } from "./websearch-credential.js"; +import type { IpythonKernelProvisioner } from "./tools/ipython.js"; +import { emptyUsage, type SessionUsageSummary, sessionUsageSummaryFrom } from "./usage.js"; +export type { RlmChildAgentActivity, RlmChildAgentSnapshot, RlmChildAgentStatus } from "../session/child-types.js"; +export { compactRlmText, rlmChildLabel } from "../session/child-types.js"; +export type { CompactionReason } from "../session/compaction.js"; export type { GoalState, GoalStatus } from "./goals.js"; export type { SessionStats } from "./session-stats.js"; export { type ParsedSkillBlock, parseSkillBlock } from "./skill-blocks.js"; -export type RlmChildAgentStatus = "queued" | "running" | "done" | "error" | "cancelled"; - -export interface RlmChildAgentActivity { - kind: "waiting" | "writing" | "executing"; - toolName?: string; -} - -export interface RlmChildAgentSnapshot { - id: string; - parentId?: string; - activeSessionId?: string; - sessionName?: string; - model?: string; - label: string; - status: RlmChildAgentStatus; - durationMs?: number; - answerPreview?: string; - toolUseCount?: number; - tokenCount?: number; - recap?: string; - sessionDir: string; - activity?: RlmChildAgentActivity; - repliedSinceTask?: boolean; - error?: string; -} - -export type { CompactionReason } from "../session/compaction.js"; - export type AgentSessionEvent = | AgentEvent | { @@ -451,12 +382,7 @@ export interface AgentSessionConfig { initialGoal?: { objective: string; tokenBudget?: number }; } -export interface ExtensionBindings { - uiContext?: ExtensionUIContext; - commandContextActions?: ExtensionCommandContextActions; - shutdownHandler?: ShutdownHandler; - onError?: ExtensionErrorListener; -} +export type { ExtensionBindings } from "../session/extensions.js"; export type { AutoRefineReviewer, AutoRefineReviewRequest } from "../session/refinement.js"; export interface PromptOptions { @@ -623,107 +549,21 @@ interface ModelSelectOptions { waitForExtensions?: boolean; } -interface ToolDefinitionEntry { - definition: ToolDefinition; - sourceInfo: SourceInfo; -} - type AutonomousSlashCommand = { kind: "status" } | { kind: "on"; config?: AgentAutonomousConfig } | { kind: "off" }; -import type { RlmMaxDepthSource, RlmMaxDepthStatus, SetRlmMaxDepthResult } from "./rlm-max-depth.js"; +import type { RlmMaxDepthStatus, SetRlmMaxDepthResult } from "./rlm-max-depth.js"; export type { RlmMaxDepthSource, RlmMaxDepthStatus, SetRlmMaxDepthResult } from "./rlm-max-depth.js"; -interface PersistedRlmMaxDepthState { - maxDepth: number; -} - type AutonomousRuntimeSnapshot = Pick< AutonomousRuntimeState, "continuationsUsed" | "gateAttempts" | "lastGateFailure" | "lastGateFailureSnapshot" >; -interface RlmChildRun { - id: string; - prompt: string; - sessionName: string; - sessionDir: string; - model: Model; - status: RlmChildAgentStatus; - durationMs?: number; - answerPreview?: string; - toolUseCount: number; - activity?: RlmChildAgentActivity; - error?: string; - abort: () => void; - publication: AgentMessageDeferred; - /** Resolves after terminal result publication and detached-run cleanup finish. */ - settlement: AgentMessageDeferred; - /** Child session, once its runtime exists. Used to cancel nested child runs. */ - session?: AgentSession; - settled: boolean; - /** Do not inject a late terminal notice after the parent session is aborted. */ - suppressTerminalNotice?: boolean; - /** Excluded from future strong barriers after an authoritative cancellation cut. */ - abandonedForQuiescence?: boolean; - /** Selector snapshot for an admitted explicit delete. */ - detachedDeletion?: RlmSubagentRegistryEntry; - /** Shared physical runtime cleanup owned by the explicit-delete path. */ - deletionCleanup?: Promise; - deletionCleanupObserver?: Promise; - /** Resolves when a deletion may release its selector reservation. */ - deletionReservation: AgentMessageDeferred; - deletionCleanupFailed?: boolean; - deletionRunFinished?: boolean; - deletionNotice?: Promise; - deletionFailureNotice?: Promise; - deletionNeedsCompletionNotice?: boolean; - completeDeletion?: () => Promise; - reportDeletionCleanupFailure?: (error: unknown) => Promise; - emitUpdate?: () => void; - lastEmittedUpdate?: string; - unsubscribe?: () => void; -} - -interface RetainedRlmChild { - session: AgentSession; - run?: RlmChildRun; -} - interface RlmSubagentModelSelection { model: Model; } -const KERNEL_STATE_LISTING_TIMEOUT_MS = 5000; -const RLM_MAX_DEPTH_STATE_CUSTOM_TYPE = "rlm_max_depth_state"; - -function noopRlmChildAbort(): void {} -function noopRlmChildEventUnsubscribe(): void {} - -function isNonNegativeInteger(value: unknown): value is number { - return typeof value === "number" && Number.isSafeInteger(value) && value >= 0; -} - -function parseDepth(value: string | undefined, fallback: number, name: string): number { - if (value === undefined || value === "") { - return fallback; - } - if (!/^\d+$/.test(value)) { - throw new Error(`${name} must be a non-negative integer`); - } - const parsed = Number(value); - if (!isNonNegativeInteger(parsed)) { - throw new Error(`${name} must be a non-negative integer`); - } - return parsed; -} - -function isPersistedRlmMaxDepthState(value: unknown): value is PersistedRlmMaxDepthState { - return ( - typeof value === "object" && value !== null && isNonNegativeInteger((value as PersistedRlmMaxDepthState).maxDepth) - ); -} - const AUTONOMOUS_STATUS_NUMBER_FORMAT = new Intl.NumberFormat("en-US"); const AUTONOMOUS_BUDGET_USAGE = @@ -829,60 +669,85 @@ function parseAutonomousBudgetOptions(tokens: string[]): AgentAutonomousConfig { return config; } -export function compactRlmText(text: string, maxLength = 160): string { - const compact = text.replace(/\s+/g, " ").trim(); - if (compact.length <= maxLength) { - return compact; +export class AgentSession { + private readonly _tools: SessionTools; + private readonly _extensions: SessionExtensions; + private readonly _kernel: SessionKernel; + private readonly _kernelEnvironment: KernelEnvironment; + private get _extensionRunner(): ExtensionRunner { + return this._extensions.runner; } - return `${compact.slice(0, Math.max(0, maxLength - 3)).trimEnd()}...`; -} - -// Child-agent label: collapse to one line but keep the full prompt — the TUI -// truncates to the visible width and elides shared prefixes, so capping here -// would only hide the divergence between near-identical sibling prompts. -export function rlmChildLabel(prompt: string): string { - return prompt.replace(/\s+/g, " ").trim() || "child agent"; -} - -function readAssistantText(message: AssistantMessage): string { - return message.content - .filter((block) => block.type === "text") - .map((block) => block.text) - .join(""); -} - -// Bounds how much accumulated child usage a parent process crash can lose. -const RLM_CHILD_USAGE_FLUSH_MAX_PENDING_MS = 60_000; - -/** Label a child completion's usage by the nearest preceding prompt that triggered it. */ -function rlmChildUsageOrigin( - messages: readonly AgentMessage[], - assistant: AssistantMessage, -): ChildUsageAttributionEntry["origin"] { - for (let index = messages.lastIndexOf(assistant) - 1; index >= 0; index--) { - const message = messages[index]; - if (message.role !== "user" && message.role !== "custom") continue; - return message.role === "custom" && isAgentSessionMessage(message) - ? message.details.id.startsWith("spawn:") - ? "spawn_task" - : "agent_message" - : "direct_user"; - } - return "direct_user"; -} + private get _ipythonKernelProvisioner(): IpythonKernelProvisioner | undefined { + return this._kernel.provisioner; + } + private get _rlmSessionDir(): string | undefined { + return this._kernelEnvironment.sessionDir; + } + private get _allowedToolNames(): ReadonlySet | undefined { + return this._tools.allowedToolNames; + } + private get _customTools(): ToolDefinition[] { + return this._tools.customTools; + } + private get _toolRegistry(): ReadonlyMap { + return this._tools.registry; + } + private get _baseSystemPrompt(): string { + return this._tools.baseSystemPrompt; + } + private set _baseSystemPrompt(prompt: string) { + this._tools.baseSystemPrompt = prompt; + } + private get _baseSystemPromptOptions(): BuildSystemPromptOptions { + return this._tools.baseSystemPromptOptions; + } + private readonly _childState: SessionChildState; -function attributeChildUsage(parentUsage: Usage, childUsage: Usage): void { - const parentContextTokens = - parentUsage.totalTokens || - parentUsage.input + parentUsage.output + parentUsage.cacheRead + parentUsage.cacheWrite; - // Recursive children are launched from an assistant tool call, so the parent assistant - // message carries their billable usage for session-level cost totals. - addAssistantUsage(parentUsage, childUsage); - // Child work affects session-level billable totals, not the parent's model-facing context size. - parentUsage.totalTokens = parentContextTokens; -} + private readonly _childUsage = new SessionChildUsage({ + sessionManager: { + getEntries: () => this.sessionManager.getEntries(), + appendChildUsageAttribution: (...args) => this.sessionManager.appendChildUsageAttribution(...args), + }, + invalidateOwnUsage: () => this._invalidateOwnUsage(), + afterParentDrain: (flush) => { + this._agentEventQueue = this._agentEventQueue.then(flush, flush); + this._agentEventQueue.catch(() => {}); + }, + }); + private readonly _children = new SessionChildren({ + isDisposed: () => this._disposed || this._disposing, + isInputSuspended: () => this._inputScheduler.suspended, + isStreaming: () => this.isStreaming, + isSessionActive: () => this.isSessionActive, + getMessageController: () => this._agentMessageController, + getDepth: () => this._childState.depth, + getMaxDepth: () => this._childState.maxDepth, + getParentNodeId: () => this._rlmParentNodeId, + getCwd: () => this._cwd, + getSessionId: () => this.sessionId, + getSessionName: () => this.sessionName, + getSessionFile: () => this.sessionFile, + getThinkingLevel: () => this.thinkingLevel, + getSemanticEdges: () => this._semanticEdges, + getChildOwner: (child) => child._children, + getParentReplyCount: (child) => child._childState.replyCount, + getChildSessionDir: (child) => child._rlmSessionDir, + listRlmSubagents: () => this.listRlmSubagents(), + deleteRlmSubagent: (target) => this.deleteRlmSubagent(target), + registerRlmChildSession: (id, child) => this.registerRlmChildSession(id, child), + resolveModel: (reference, target) => this._resolveRlmSubagentModel(reference, target), + createSessionDir: () => this._createChildRlmSessionDir(), + createRuntimeOptions: (request) => this._createRlmSubagentRuntimeOptions(request), + createRuntime: (options) => this._createRlmSubagentRuntime(options), + createUsageTracker: () => this._childUsage.createTracker(this._findLastAssistantMessage()), + hasDeferredTerminalNotices: () => this._hasDeferredRlmTerminalNotices(), + waitForHeadlessIdle: () => this.waitForHeadlessIdle(), + waitForActivityChange: (signal) => this._waitForSessionActivityChange(signal), + deliverTerminalNotice: (message) => this._deferRlmTerminalNotice(message), + emit: (event) => this._emit(event), + onSettled: () => this._maybeResumeGoalContinuationAfterRlmWork(), + }); -export class AgentSession { private readonly _refinement: SessionRefinement; readonly agent: Agent; readonly sessionManager: SessionManager; @@ -1068,35 +933,21 @@ export class AgentSession { recordBashResult: (command, result, options) => this.recordBashResult(command, result, options), }); - private _extensionRunner!: ExtensionRunner; - private _execEnvProvider?: () => Record | undefined; private _turnIndex = 0; private _modelSelectEmitQueue: Promise = Promise.resolve(); private _modelSelectEmitQueueIdle = true; private _modelSelectEmitContext = new AsyncLocalStorage(); private _resourceLoader: ResourceLoader; - private _customTools: ToolDefinition[]; - private _acpMcpTools: ToolDefinition[] = []; - private _baseToolDefinitions: Map = new Map(); private _cwd: string; private _agentDir?: string; - private _extensionRunnerRef?: { current?: ExtensionRunner }; private _initialActiveToolNames?: string[]; - private _allowedToolNames?: Set; private _includeGoals: boolean; private _includeCompactSkill: boolean; private _rlmHeartbeatController?: AgentRlmHeartbeatController; private _agentMessageController?: AgentSessionMessageController; private _agentObserveController?: AgentObserveController; private _mcpManager?: McpManager; - private _baseToolsOverride?: Record; - private _sessionStartEvent: SessionStartEvent; - private _extensionUIContext?: ExtensionUIContext; - private _extensionCommandContextActions?: ExtensionCommandContextActions; - private _extensionShutdownHandler?: ShutdownHandler; - private _extensionErrorListener?: ExtensionErrorListener; - private _extensionErrorUnsubscriber?: () => void; private _disposed = false; private readonly _disposeCallbacks = new Set<() => void | Promise>(); private _disposeCallbacksPromise?: Promise; @@ -1104,61 +955,12 @@ export class AgentSession { // re-populate the retained map after it's been cleared. private _disposing = false; private _disposeAsyncPromise?: Promise; - private _ipythonKernelProvisioner?: IpythonKernelProvisioner; - /** Artifact dir backing the current provisioner's kernel snapshot, if any. */ - private _ipythonKernelSnapshotDir?: string; - /** True once the runtime has been built once; later builds are in-process rebuilds (/reload). */ - private _ipythonRuntimeBuilt = false; - private readonly _prewarmIpythonKernel: boolean; - private _rlmDepth: number; - private readonly _configuredRlmMaxDepth: number | undefined; - private _rlmMaxDepth: number; - private _rlmMaxDepthSource: RlmMaxDepthSource; - private _rlmSessionDir?: string; private readonly _semanticEdges: SemanticEdgeRecorder; private _rlmParentNodeId?: string; private _rlmParentAgent?: string; - private _repliedToParentSinceTask: boolean | undefined; - private _parentReplyCount = 0; - private _subagentRuntimeHost?: SubagentRuntimeHost; - // Shared by children charged to the same assistant; excludes usage not yet attributed on disk. - private _rlmDurableParentUsage = new WeakMap(); - // Child usage not yet represented by an indexed attribution, including a delayed parent entry. - private _rlmUnindexedChildUsage = new WeakMap(); - private _activeRlmChildRuns = new Map(); - private _unsettledRlmChildRuns = new Set(); - private _abandonedRlmQuiescenceChildIds = new Set(); - private _rlmQuiescenceWaitAborts = new Set(); - private _pendingRlmSubagentSessionNames = new Set(); - // Inline mode keeps finished child sessions so the inspector can still read them; - // the daemon does the same by leaving the child session resident in its registry. - private _rlmChildSessions = new Map(); - private _deletedRlmChildIds = new Set(); - // Failed explicit deletes stay hidden from listings but retain their original - // selector so a later delete can retry cleanup without orphaning the runtime. - private _rlmChildCleanupFailures = new Map(); - private _deletingRlmChildren = new Map< - string, - { - subagent: RlmSubagentRegistryEntry; - promise: Promise; - } - >(); - // Kept alive for retained children so nested updates (e.g. a grandchild cancel) - // still forward to root; torn down when the retained child is disposed. - private _rlmChildUnsubscribes = new Map void>(); - /** Latest recap for this session, written by the daemon summarizer; read by a parent to label its child snapshots. */ - private _currentRecap?: string; private _modelRegistry: ModelRegistry; - private _toolRegistry: Map = new Map(); - private _toolDefinitions: Map = new Map(); - private _toolPromptSnippets: Map = new Map(); - private _toolPromptGuidelines: Map = new Map(); - - private _baseSystemPrompt = ""; - private _baseSystemPromptOptions!: BuildSystemPromptOptions; private readonly _continuation = new SessionContinuation({ waitForAgentIdle: () => this.agent.waitForIdle(), waitForRetry: () => this.waitForRetry(), @@ -1204,7 +1006,7 @@ export class AgentSession { isDisposing: () => this._disposing, isStreaming: () => this.isStreaming, isCompacting: () => this.isCompacting, - getDepth: () => this._rlmDepth, + getDepth: () => this._childState.depth, getRlmSessionDir: () => this._rlmSessionDir, getModel: () => this.model, getThinkingLevel: () => this.thinkingLevel, @@ -1233,37 +1035,138 @@ export class AgentSession { this._serviceTierPreference = config.serviceTierPreference ?? config.agent.state.serviceTier; this._scopedModels = config.scopedModels ?? []; this._resourceLoader = config.resourceLoader; - this._customTools = config.customTools ?? []; this._cwd = config.cwd; this._agentDir = config.agentDir; this._modelRegistry = config.modelRegistry; - this._extensionRunnerRef = config.extensionRunnerRef; this._initialActiveToolNames = config.initialActiveToolNames; - this._allowedToolNames = config.allowedToolNames ? new Set(config.allowedToolNames) : undefined; this._includeGoals = config.includeGoals ?? true; this._includeCompactSkill = config.includeCompactSkill ?? this.settingsManager.getCompactionAgentCallable(); this._rlmHeartbeatController = config.rlmHeartbeatController; this._agentMessageController = config.agentMessageController; this._agentObserveController = config.agentObserveController; this._mcpManager = config.mcpManager; - this._baseToolsOverride = config.baseToolsOverride; - this._sessionStartEvent = config.sessionStartEvent ?? { type: "session_start", reason: "startup" }; - const headerRlmDepth = this.sessionManager.getHeader()?.rlmDepth; - this._rlmDepth = - config.rlmDepth ?? - (isNonNegativeInteger(headerRlmDepth) ? headerRlmDepth : parseDepth(process.env.RLM_DEPTH, 0, "RLM_DEPTH")); - this._configuredRlmMaxDepth = config.rlmMaxDepth; - if (this._configuredRlmMaxDepth !== undefined && !isNonNegativeInteger(this._configuredRlmMaxDepth)) { - throw new Error("rlmMaxDepth must be a non-negative integer"); - } - const resolvedRlmMaxDepth = this._resolveRlmMaxDepth(); - this._rlmMaxDepth = resolvedRlmMaxDepth.maxDepth; - this._rlmMaxDepthSource = resolvedRlmMaxDepth.source; - this._prewarmIpythonKernel = (config.prewarmIpythonKernel ?? false) && this._rlmDepth === 0; - - this._rlmSessionDir = config.rlmSessionDir; + this._childState = new SessionChildState( + { + sessionManager: this.sessionManager, + settingsManager: this.settingsManager, + getRlmMaxDepthStatus: () => this.getRlmMaxDepthStatus(), + refreshPrompt: (preserveExtensionPrompt) => { + const oldBase = this._baseSystemPrompt; + this._baseSystemPrompt = this._rebuildSystemPrompt(this.getActiveToolNames()); + this.agent.state.systemPrompt = preserveExtensionPrompt + ? this._refreshExtensionSystemPrompt(this.agent.state.systemPrompt, oldBase) + : this._baseSystemPrompt; + }, + emitRecap: (recap) => this._emit({ type: "recap_update", recap }), + }, + config, + ); + this._rlmParentNodeId = config.rlmParentNodeId; this._rlmParentAgent = config.rlmParentAgent; + this._kernelEnvironment = new KernelEnvironment( + { + agentDir: this._agentDir, + authStorage: this._modelRegistry.authStorage, + resourceLoader: this._resourceLoader, + getDepth: () => this._childState.depth, + getMaxDepth: () => this._childState.maxDepth, + getArtifactDir: () => this.sessionManager.getSessionArtifactDir(), + getLocalHarnessStateDir: () => this._refinement._localHarnessStateDir(), + }, + config.rlmSessionDir, + ); + this._kernel = new SessionKernel( + { + cwd: this._cwd, + getArtifactDir: () => this.sessionManager.getSessionArtifactDir(), + getSessionId: () => this.sessionId, + getEnv: () => this._rlmKernelEnv(), + getShellCommandPrefix: () => this.settingsManager.getShellCommandPrefix(), + getShellPath: () => this.settingsManager.getShellPath(), + createHostHandlers: () => this._createKernelHostHandlers(), + recordLateSentAgentMessage: (id, message) => this._recordLateIpythonSentAgentMessage(id, message), + getMessages: () => this.agent.state.messages, + appendCustomMessageEntry: (...args) => this.sessionManager.appendCustomMessageEntry(...args), + emit: (event) => this._emit(event), + sendCustomMessage: (message, options) => this.sendCustomMessage(message, options), + }, + (config.prewarmIpythonKernel ?? false) && this._childState.depth === 0, + ); + this._tools = new SessionTools( + { + cwd: this._cwd, + resourceLoader: this._resourceLoader, + getExtensionRunner: () => this._extensionRunner, + getSessionFile: () => this.sessionManager.getSessionFile(), + getModelVisibleSkills: () => this._modelVisibleSkills(), + getDepth: () => this._childState.depth, + getMaxDepth: () => this._childState.maxDepth, + getParentAgent: () => this._rlmParentAgent, + getMcpManager: () => this._mcpManager, + getProvisioner: () => this._kernel.provisioner, + getActiveToolNames: () => this.getActiveToolNames(), + setActiveToolsByName: (names) => this.setActiveToolsByName(names), + getActiveTools: () => this.agent.state.tools, + setActiveTools: (tools) => { + this.agent.state.tools = tools; + }, + setSystemPrompt: (prompt) => { + this.agent.state.systemPrompt = prompt; + }, + isStreaming: () => this.isStreaming, + rebuildRuntime: (options) => this._buildRuntime(options), + acquireInputPause: () => this.acquireSessionInputPause(), + waitForAgentIdle: () => this.agent.waitForIdle(), + getEventQueue: () => this._agentEventQueue, + }, + { + customTools: config.customTools, + allowedToolNames: config.allowedToolNames, + baseToolsOverride: config.baseToolsOverride, + }, + ); + this._extensions = new SessionExtensions( + { + cwd: this._cwd, + sessionManager: this.sessionManager, + resourceLoader: this._resourceLoader, + modelRegistry: this._modelRegistry, + getModelRegistry: () => this.modelRegistry, + getPromptTemplates: () => this.promptTemplates, + bindShutdownHandler: (handler) => handler?.bind(this), + getAgentMessageController: () => this._agentMessageController, + refreshCurrentModel: () => this._refreshCurrentModelFromRegistry(), + sendCustomMessage: (message, options) => this.sendCustomMessage(message, options), + sendUserMessage: (content, options) => this.sendUserMessage(content, options), + setSessionName: (name) => this.setSessionName(name), + getActiveToolNames: () => this.getActiveToolNames(), + getAllTools: () => this.getAllTools(), + setActiveToolsByName: (names) => this.setActiveToolsByName(names), + refreshTools: () => this._refreshToolRegistry(), + setModel: (model) => this.setModel(model), + getThinkingLevel: () => this.thinkingLevel, + setThinkingLevel: (level) => this.setThinkingLevel(level), + getModel: () => this.model, + isStreaming: () => this.isStreaming, + getSignal: () => this.agent.signal, + abort: () => this.abort(), + getQueuedActionCount: () => this.queuedActionCount, + getContextUsage: () => this.getContextUsage(), + compact: (instructions) => this.compact(instructions), + getSystemPrompt: () => this.systemPrompt, + rebuildSystemPrompt: () => { + this._baseSystemPrompt = this._rebuildSystemPrompt(this.getActiveToolNames()); + this.agent.state.systemPrompt = this._baseSystemPrompt; + }, + reloadSettings: () => this.settingsManager.reload(), + getMcpManager: () => this._mcpManager, + rebuildRuntime: (options) => this._buildRuntime(options), + }, + config.sessionStartEvent ?? { type: "session_start", reason: "startup" }, + config.extensionRunnerRef, + ); + this._semanticEdges = new SemanticEdgeRecorder({ ledgerPath: semanticEdgeLedgerPath({ rlmSessionDir: this._rlmSessionDir, @@ -1274,13 +1177,8 @@ export class AgentSession { spawnedByRequestId: config.semanticSpawnedByRequestId, }); this.agent.streamFn = wrapStreamFnWithSemanticEdges(this.agent.streamFn, this._semanticEdges); - // A resumed child may have replied before this process started; false would - // claim knowledge that is not present in the session transcript. - this._repliedToParentSinceTask = - this._rlmDepth > 0 && this.sessionManager.getBranch().some((entry) => entry.type === "message") - ? undefined - : false; - this._subagentRuntimeHost = config.subagentRuntimeHost; + this._childState.initializeParentReply(); + this._children.setRuntimeHost(config.subagentRuntimeHost); this._autonomousState = createAutonomousRuntimeState(config.autonomous, { cwd: this._cwd, }); @@ -1291,7 +1189,7 @@ export class AgentSession { // thinking_level_change, service_tier_change) and no persisted // thread_goal_state. This prevents reseeding after clear/complete/error // or restart/rehydration of a session that already has messages or a goal. - if (this._rlmDepth === 0 && config.initialGoal && goalPersistence.canSeed()) { + if (this._childState.depth === 0 && config.initialGoal && goalPersistence.canSeed()) { this._startGoal(config.initialGoal.objective, config.initialGoal.tokenBudget); // Goal context is the model's only source of goal visibility; action // admission is unavailable mid-construction, so ride the next turn. @@ -1336,89 +1234,11 @@ export class AgentSession { } replaceAcpMcpServers(servers: readonly AcpMcpServerConfig[], ownerId: string): void { - if (this.isStreaming) throw new Error("Cannot replace ACP MCP servers while the agent is running"); - if (!this._mcpManager) { - if (servers.length > 0) throw new Error("MCP is unavailable in this session"); - return; - } - if (servers.length > 0 && !this._ipythonKernelProvisioner) { - throw new Error("ACP MCP servers require the built-in cpython tool"); - } - this._assertAcpMcpToolNamesAvailable(acpMcpToolNames(servers)); - if (!this._mcpManager.replaceAcpServers(servers, ownerId)) return; - this._rebuildRuntimeForAcpMcpServers(); - } - - async releaseAcpMcpServers(ownerId: string, serverNames: readonly string[]): Promise { - if (!this._mcpManager?.canReleaseAcpServers(ownerId)) return; - if (this._mcpManager.replaceAcpServers([], ownerId)) { - const removedToolNames = new Set(this._acpMcpTools.map((tool) => tool.name)); - const activeToolNames = this.getActiveToolNames().filter((name) => !removedToolNames.has(name)); - for (const name of removedToolNames) this._allowedToolNames?.delete(name); - this._acpMcpTools = []; - this._refreshToolRegistry({ activeToolNames, includeAllExtensionTools: true }); - this._baseSystemPrompt = this._rebuildSystemPrompt(this.getActiveToolNames()); - this.agent.state.systemPrompt = this._baseSystemPrompt; - } - const names = [...new Set(serverNames)]; - if (names.length === 0) return; - - const inputPause = this.acquireSessionInputPause(); - try { - // Do not rebuild or kill the notebook. Wait for the current turn, then ask - // the kernel-owned MCP registry to close only these cached transports. - await this.agent.waitForIdle(); - await this._agentEventQueue; - const manager = this._ipythonKernelProvisioner?.manager; - if (!manager?.isRunning) return; - const code = [ - "import importlib as _prime_importlib", - '_prime_mcp = _prime_importlib.import_module("rlm.mcp")', - `_prime_mcp_names = ${JSON.stringify(names)}`, - "_prime_mcp_errors = []", - "for _prime_mcp_name in _prime_mcp_names:", - " try:", - " await _prime_mcp.reload(_prime_mcp_name)", - " except BaseException as _prime_mcp_error:", - " _prime_mcp_errors.append(_prime_mcp_error)", - "if _prime_mcp_errors:", - " raise _prime_mcp_errors[0]", - "del _prime_mcp, _prime_importlib, _prime_mcp_names, _prime_mcp_errors, _prime_mcp_name", - ].join("\n"); - const result = await manager.execute(code); - if (result.status !== "ok") { - throw new Error(`Failed to close ACP MCP kernel transports: ${result.stderr || "kernel error"}`); - } - } finally { - inputPause.release(); - } - } - - private _assertAcpMcpToolNamesAvailable(names: readonly string[]): void { - const occupiedNames = new Set([ - ...this._baseToolDefinitions.keys(), - ...this._customTools.map((tool) => tool.name), - ...this._extensionRunner.getAllRegisteredTools().map((tool) => tool.definition.name), - ]); - for (const name of names) { - if (occupiedNames.has(name)) { - throw new Error(`ACP MCP tool name conflicts with an existing tool: ${name}`); - } - } + this._tools.replaceAcpMcpServers(servers, ownerId); } - private _rebuildRuntimeForAcpMcpServers(): void { - const previousToolNames = new Set(this._acpMcpTools.map((tool) => tool.name)); - const nextToolNames = acpMcpToolNames(this._mcpManager?.getAcpServers() ?? []); - this._assertAcpMcpToolNamesAvailable(nextToolNames); - const activeToolNames = this.getActiveToolNames().filter((name) => !previousToolNames.has(name)); - activeToolNames.push(...nextToolNames); - this._buildRuntime({ - activeToolNames, - includeAllExtensionTools: true, - }); - this._baseSystemPrompt = this._rebuildSystemPrompt(this.getActiveToolNames()); - this.agent.state.systemPrompt = this._baseSystemPrompt; + releaseAcpMcpServers(ownerId: string, serverNames: readonly string[]): Promise { + return this._tools.releaseAcpMcpServers(ownerId, serverNames); } get modelRegistry(): ModelRegistry { @@ -1426,7 +1246,7 @@ export class AgentSession { } setSubagentRuntimeHost(host?: SubagentRuntimeHost): void { - this._subagentRuntimeHost = host; + this._children.setRuntimeHost(host); } private async _getRequiredRequestAuth(model: Model): Promise<{ @@ -1460,55 +1280,11 @@ export class AgentSession { * happens here instead of in wrappers. */ private _installAgentToolHooks(): void { - this.agent.beforeToolCall = async ({ toolCall, args }) => { - const runner = this._extensionRunner; - if (!runner.hasHandlers("tool_call")) { - return undefined; - } - - await this._agentEventQueue; - - try { - return await runner.emitToolCall({ - type: "tool_call", - toolName: toolCall.name, - toolCallId: toolCall.id, - input: args as Record, - }); - } catch (err) { - if (err instanceof Error) { - throw err; - } - throw new Error(`Extension failed, blocking execution: ${String(err)}`); - } - }; - - this.agent.afterToolCall = async ({ toolCall, args, result, isError }) => { - const runner = this._extensionRunner; - if (!runner.hasHandlers("tool_result")) { - return undefined; - } - - const hookResult = await runner.emitToolResult({ - type: "tool_result", - toolName: toolCall.name, - toolCallId: toolCall.id, - input: args as Record, - content: result.content, - details: result.details, - isError, - }); - - if (!hookResult) { - return undefined; - } - - return { - content: hookResult.content, - details: hookResult.details, - isError: hookResult.isError ?? isError, - }; - }; + installExtensionToolHooks( + this.agent, + () => this._extensionRunner, + () => this._agentEventQueue, + ); } private _installAgentContinuationHook(): void { @@ -1591,52 +1367,8 @@ export class AgentSession { this._emit({ type: "goal_update", goal: this.goalState }); } - private _loadPersistedRlmMaxDepthState(): PersistedRlmMaxDepthState | undefined { - const branch = this.sessionManager.getBranch(); - for (let i = branch.length - 1; i >= 0; i--) { - const entry = branch[i]; - if ( - entry.type === "custom" && - entry.customType === RLM_MAX_DEPTH_STATE_CUSTOM_TYPE && - isPersistedRlmMaxDepthState(entry.data) - ) { - return entry.data; - } - } - return undefined; - } - - private _resolveRlmMaxDepth(): { - maxDepth: number; - source: RlmMaxDepthSource; - } { - const persisted = this._loadPersistedRlmMaxDepthState(); - if (persisted) { - return { maxDepth: persisted.maxDepth, source: "chat" }; - } - if (this._configuredRlmMaxDepth !== undefined) { - return { maxDepth: this._configuredRlmMaxDepth, source: "inherited" }; - } - const global = this.settingsManager.getRlmMaxDepth(); - if (global !== undefined && isNonNegativeInteger(global)) { - return { maxDepth: global, source: "global" }; - } - const env = process.env.RLM_MAX_DEPTH; - if (env !== undefined && env !== "") { - return { maxDepth: parseDepth(env, 1, "RLM_MAX_DEPTH"), source: "env" }; - } - return { maxDepth: 2, source: "default" }; - } - private _reloadRlmMaxDepthFromBranch(): void { - const previousMaxDepth = this._rlmMaxDepth; - const resolved = this._resolveRlmMaxDepth(); - this._rlmMaxDepth = resolved.maxDepth; - this._rlmMaxDepthSource = resolved.source; - if (resolved.maxDepth !== previousMaxDepth) { - this._baseSystemPrompt = this._rebuildSystemPrompt(this.getActiveToolNames()); - this.agent.state.systemPrompt = this._baseSystemPrompt; - } + this._childState.reloadFromBranch(); } private _cancelSessionActions( @@ -3070,44 +2802,15 @@ export class AgentSession { return this._disposeAsyncPromise; } - private async _disposeAsyncOnce(kernelSnapshot: boolean): Promise { + private _disposeAsyncOnce(kernelSnapshot: boolean): Promise { // Flush kernels/traces for both still-running and retained children; the sync // dispose() below only tears them down synchronously. - for (const run of [...this._activeRlmChildRuns.values()]) { - const childSession = run.session; - if (!childSession) continue; - if (run.detachedDeletion) { - run.suppressTerminalNotice = true; - if (run.deletionCleanupObserver) { - await run.deletionCleanupObserver.catch(() => false); - } else if (run.deletionCleanup) { - await run.deletionCleanup.catch(() => childSession.disposeAsync().catch(() => undefined)); - } else { - // Cleanup already failed and was exposed for retry before disposal. - await childSession.disposeAsync().catch(() => undefined); - } - if (!run.settled) await this._finishRlmRunDeletion(run); - } else { - await childSession.disposeAsync().catch(() => undefined); - } - } - for (const unsubscribe of this._rlmChildUnsubscribes.values()) { - unsubscribe(); - } - this._rlmChildUnsubscribes.clear(); - for (const { session } of this._rlmChildSessions.values()) { - await session.disposeAsync().catch(() => undefined); - } - this._rlmChildSessions.clear(); - this._rlmChildCleanupFailures.clear(); - this._deletedRlmChildIds.clear(); - try { - await this._ipythonKernelProvisioner?.dispose({ snapshot: kernelSnapshot }); - } catch { - // a failed kernel startup already cleaned up after itself - } - this.dispose(); - await this._disposeCallbacksPromise; + return this._children.disposeAsync(() => + this._kernel.dispose(kernelSnapshot, () => { + this.dispose(); + return this._disposeCallbacksPromise; + }), + ); } private _startDisposeCallbacks(): Promise { @@ -3135,24 +2838,13 @@ export class AgentSession { return; } this._disposed = true; - for (const run of this._unsettledRlmChildRuns) run.suppressTerminalNotice = true; - for (const controller of this._rlmQuiescenceWaitAborts) controller.abort(); + this._children.beginDisposal(); this._commitFence.dispose(); try { // Invalidate scheduled timers and abort any in-flight review so a late // resolution cannot write harness state or re-subscribe handlers. this._refinement.dispose(); - this._cancelActiveRlmChildRuns("Parent session disposed"); - for (const unsubscribe of this._rlmChildUnsubscribes.values()) { - unsubscribe(); - } - this._rlmChildUnsubscribes.clear(); - for (const { session } of this._rlmChildSessions.values()) { - session.dispose(); - } - this._rlmChildSessions.clear(); - this._rlmChildCleanupFailures.clear(); - this._deletedRlmChildIds.clear(); + this._children.dispose(); this._pendingNextTurnMessages = []; const deliveryError = new Error("Session disposed before prompt delivery."); const completionError = new Error("Session disposed before prompt completion."); @@ -3216,41 +2908,19 @@ export class AgentSession { } getActiveToolNames(): string[] { - return this.agent.state.tools.map((t) => t.name); + return this._tools.getActiveToolNames(); } getAllTools(): ToolInfo[] { - return Array.from(this._toolDefinitions.values()).map(({ definition, sourceInfo }) => ({ - name: definition.name, - description: definition.description, - parameters: definition.parameters, - sourceInfo, - })); + return this._tools.getAllTools(); } getToolDefinition(name: string): ToolDefinition | undefined { - return this._toolDefinitions.get(name)?.definition; + return this._tools.getToolDefinition(name); } setActiveToolsByName(toolNames: string[]): void { - const tools: AgentTool[] = []; - const validToolNames: string[] = []; - const seenToolNames = new Set(); - for (const name of toolNames) { - if (seenToolNames.has(name)) { - continue; - } - const tool = this._toolRegistry.get(name); - if (tool) { - seenToolNames.add(name); - tools.push(tool); - validToolNames.push(name); - } - } - this.agent.state.tools = tools; - - this._baseSystemPrompt = this._rebuildSystemPrompt(validToolNames); - this.agent.state.systemPrompt = this._baseSystemPrompt; + this._tools.setActiveToolsByName(toolNames); } get isCompacting(): boolean { @@ -3297,7 +2967,7 @@ export class AgentSession { } get rlmDepth(): number { - return this._rlmDepth; + return this._childState.depth; } get semanticEdges(): SemanticEdgeRecorder { @@ -3305,7 +2975,7 @@ export class AgentSession { } get rlmMaxDepth(): number { - return this._rlmMaxDepth; + return this._childState.maxDepth; } get sessionName(): string | undefined { @@ -3358,79 +3028,12 @@ export class AgentSession { return this._resourceLoader.getPrompts().prompts; } - private _normalizePromptSnippet(text: string | undefined): string | undefined { - if (!text) return undefined; - const oneLine = text - .replace(/[\r\n]+/g, " ") - .replace(/\s+/g, " ") - .trim(); - return oneLine.length > 0 ? oneLine : undefined; - } - - private _normalizePromptGuidelines(guidelines: string[] | undefined): string[] { - if (!guidelines || guidelines.length === 0) { - return []; - } - - const unique = new Set(); - for (const guideline of guidelines) { - const normalized = guideline.trim(); - if (normalized.length > 0) { - unique.add(normalized); - } - } - return Array.from(unique); - } - private _rebuildSystemPrompt(toolNames: string[]): string { - const validToolNames = toolNames.filter((name) => this._toolRegistry.has(name)); - const toolSnippets: Record = {}; - const promptGuidelines: string[] = []; - for (const name of validToolNames) { - const snippet = this._toolPromptSnippets.get(name); - if (snippet) { - toolSnippets[name] = snippet; - } - - const toolGuidelines = this._toolPromptGuidelines.get(name); - if (toolGuidelines) { - promptGuidelines.push(...toolGuidelines); - } - } - - const loaderSystemPrompt = this._resourceLoader.getSystemPrompt(); - const loaderAppendSystemPrompt = this._resourceLoader.getAppendSystemPrompt(); - const appendSystemPrompt = - loaderAppendSystemPrompt.length > 0 ? loaderAppendSystemPrompt.join("\n\n") : undefined; - const loadedSkills = this._modelVisibleSkills(); - const loadedContextFiles = this._resourceLoader.getAgentsFiles().agentsFiles; - - this._baseSystemPromptOptions = { - cwd: this._cwd, - skills: loadedSkills, - contextFiles: loadedContextFiles, - customPrompt: loaderSystemPrompt, - appendSystemPrompt, - messagesPath: this.sessionManager.getSessionFile(), - selectedTools: validToolNames, - toolSnippets, - promptGuidelines, - allowRecursion: this._rlmDepth < this._rlmMaxDepth, - rlmDepth: this._rlmDepth, - rlmParentAgent: this._rlmParentAgent, - genericMcpServers: this._mcpManager?.getEnabledPersistentGenericServers(), - }; - return buildSystemPrompt(this._baseSystemPromptOptions); + return this._tools.rebuildSystemPrompt(toolNames); } private _refreshExtensionSystemPrompt(extensionPrompt: string, baseSnapshot: string): string { - if (this._baseSystemPrompt === baseSnapshot) { - return extensionPrompt; - } - if (!extensionPrompt.includes(baseSnapshot)) { - return extensionPrompt; - } - return extensionPrompt.replace(baseSnapshot, () => this._baseSystemPrompt); + return this._tools.refreshExtensionSystemPrompt(extensionPrompt, baseSnapshot); } private _finishSubmissionNormalization( @@ -3596,7 +3199,7 @@ export class AgentSession { customMessage, admissionCommitted, }); - if (customMessage?.details.fromRelationship === "parent") this._repliedToParentSinceTask = false; + if (customMessage?.details.fromRelationship === "parent") this._childState.resetReply(); } async queueAgentMessagePrompt( @@ -3610,14 +3213,14 @@ export class AgentSession { agentMessageId, message: customMessage, }); - if (customMessage?.details.fromRelationship === "parent") this._repliedToParentSinceTask = false; + if (customMessage?.details.fromRelationship === "parent") this._childState.resetReply(); return true; } const queued = await this._queuePreparedPrompt("followUp", text, undefined, { agentMessageId, message: customMessage, }); - if (queued && customMessage?.details.fromRelationship === "parent") this._repliedToParentSinceTask = false; + if (queued && customMessage?.details.fromRelationship === "parent") this._childState.resetReply(); return queued; } @@ -5592,10 +5195,7 @@ export class AgentSession { } requestAbort(): void { - for (const run of [...this._unsettledRlmChildRuns]) { - if (run.status === "cancelled") this._abandonRlmRunForQuiescence(run); - } - for (const controller of this._rlmQuiescenceWaitAborts) controller.abort(); + this._children.requestAbort(); this._inputScheduler.suspend("abort"); this._demoteRlmTerminalNoticeActions(); this._cancelSessionActions( @@ -5638,7 +5238,7 @@ export class AgentSession { this._inputScheduler.suspend("update-restart"); this._cancelPostCompactionContinue(); this.abortRetry(); - for (const controller of this._rlmQuiescenceWaitAborts) controller.abort(); + this._children.cancelQuiescenceWaits(); this._cancelActiveRlmChildRuns("Parent session aborted for update restart"); this._goalAbortInProgress = this._goals.state.status === "active"; this.agent.abort(); @@ -5937,80 +5537,8 @@ export class AgentSession { return this.model ? (clampThinkingLevel(this.model, level) as ThinkingLevel) : "off"; } - private async _syncKernelStateAfterCompaction(): Promise { - const provisioner = this._ipythonKernelProvisioner; - if (!provisioner?.hasRunningKernel) return; - const pruned = await provisioner.pruneOversizedVariables().catch(() => null); - const abort = new AbortController(); - const timer = setTimeout(() => abort.abort(), KERNEL_STATE_LISTING_TIMEOUT_MS); - if (typeof timer === "object" && "unref" in timer) timer.unref(); - let names: string[] | null; - try { - names = await provisioner.listNamespaceNames(abort.signal).catch(() => null); - } finally { - clearTimeout(timer); - } - if (names === null && !provisioner.hasRunningKernel) return; - const detail = - names === null - ? "" - : names.length > 0 - ? ` These names are still defined: ${names.join(", ")}.` - : " You have not defined any names yet."; - const prunedDetail = - pruned && pruned.length > 0 - ? ` Variables above the per-variable snapshot limit were removed: ${pruned.join(", ")}.` - : ""; - const content = [ - "[python-state]", - "", - `Your Python kernel persisted through compaction; its remaining variables, imports, and helpers are still available.${prunedDetail}${detail}`, - ].join("\n"); - const message = { - role: "custom" as const, - customType: "ipython_state", - content, - display: false, - timestamp: Date.now(), - } satisfies CustomMessage; - const messages = this.agent.state.messages; - const last = messages[messages.length - 1]; - const insertBeforeError = last?.role === "assistant" && (last as AssistantMessage).stopReason === "error"; - if (insertBeforeError) { - messages.splice(messages.length - 1, 0, message); - } else { - messages.push(message); - } - this.sessionManager.appendCustomMessageEntry(message.customType, message.content, message.display, undefined); - this._emit({ type: "message_start", message }); - this._emit({ type: "message_end", message }); - } - - private _onIpythonStateRestored(result: RestoreResult): void { - const lines = ["[python-state-restored]", ""]; - if (result.restored.length > 0) { - lines.push( - `Your Python kernel state was revived from your previous session. These names are available again: ${result.restored.join(", ")}.`, - ); - } else { - lines.push( - "Your previous Python kernel state could not be revived; the kernel is starting fresh, so re-create any variables, imports, or loaded data you need.", - ); - } - if (result.failed.length > 0) { - lines.push( - `These could not be restored and must be recreated if needed: ${result.failed.map((f) => f.name).join(", ")}.`, - ); - } - void this.sendCustomMessage( - { - customType: IPYTHON_STATE_RESTORED_CUSTOM_TYPE, - content: lines.join("\n"), - display: true, - details: { restored: result.restored.length > 0 }, - }, - { deliverAs: "nextTurn" }, - ).catch(() => {}); + private _syncKernelStateAfterCompaction(): Promise { + return this._kernel.syncAfterCompaction(); } setSteeringMode(mode: "all" | "one-at-a-time"): void { @@ -6060,11 +5588,8 @@ export class AgentSession { return performSessionCompaction(this._compactionExecution, options); } - private async _reapDeletedRlmSubagentRuntimesAfterCompaction(): Promise { - const childIds = [...this._rlmChildCleanupFailures.keys()].filter( - (childId) => !this._activeRlmChildRuns.get(childId)?.detachedDeletion, - ); - await Promise.allSettled(childIds.map((childId) => this.deleteRlmSubagent(childId))); + private _reapDeletedRlmSubagentRuntimesAfterCompaction(): Promise { + return this._children.reapAfterCompaction(); } abortCompaction(): void { @@ -6203,320 +5728,29 @@ export class AgentSession { * the daemon) can update the underlying value per attach without rebinding. */ setExecEnvProvider(provider: (() => Record | undefined) | undefined): void { - this._execEnvProvider = provider; - const extensions = this._resourceLoader.getExtensions(); - extensions.runtime.getExecEnv = provider; + this._extensions.setExecEnvProvider(provider); } - async bindExtensions(bindings: ExtensionBindings): Promise { - if (bindings.uiContext !== undefined) { - this._extensionUIContext = bindings.uiContext; - } - if (bindings.commandContextActions !== undefined) { - this._extensionCommandContextActions = bindings.commandContextActions; - } - if (bindings.shutdownHandler !== undefined) { - this._extensionShutdownHandler = bindings.shutdownHandler; - } - if (bindings.onError !== undefined) { - this._extensionErrorListener = bindings.onError; - } - - this._applyExtensionBindings(this._extensionRunner); - await this._extensionRunner.emit(this._sessionStartEvent); - await this.extendResourcesFromExtensions(this._sessionStartEvent.reason === "reload" ? "reload" : "startup"); + bindExtensions(bindings: ExtensionBindings): Promise { + return this._extensions.bindExtensions(bindings); } - private async extendResourcesFromExtensions(reason: "startup" | "reload"): Promise { - if (!this._extensionRunner.hasHandlers("resources_discover")) { + private _refreshCurrentModelFromRegistry(): void { + const currentModel = this.model; + if (!currentModel) { return; } - const { skillPaths, promptPaths, themePaths } = await this._extensionRunner.emitResourcesDiscover( - this._cwd, - reason, - ); - - if (skillPaths.length === 0 && promptPaths.length === 0 && themePaths.length === 0) { + const refreshedModel = this._modelRegistry.find(currentModel.provider, currentModel.id); + if (!refreshedModel || refreshedModel === currentModel) { return; } - const extensionPaths: ResourceExtensionPaths = { - skillPaths: this.buildExtensionResourcePaths(skillPaths), - promptPaths: this.buildExtensionResourcePaths(promptPaths), - themePaths: this.buildExtensionResourcePaths(themePaths), - }; + this.agent.state.model = refreshedModel; + } - this._resourceLoader.extendResources(extensionPaths); - this._baseSystemPrompt = this._rebuildSystemPrompt(this.getActiveToolNames()); - this.agent.state.systemPrompt = this._baseSystemPrompt; - } - - private buildExtensionResourcePaths(entries: Array<{ path: string; extensionPath: string }>): Array<{ - path: string; - metadata: { - source: string; - scope: "temporary"; - origin: "top-level"; - baseDir?: string; - }; - }> { - return entries.map((entry) => { - const source = this.getExtensionSourceLabel(entry.extensionPath); - const baseDir = entry.extensionPath.startsWith("<") ? undefined : dirname(entry.extensionPath); - return { - path: entry.path, - metadata: { - source, - scope: "temporary", - origin: "top-level", - baseDir, - }, - }; - }); - } - - private getExtensionSourceLabel(extensionPath: string): string { - if (extensionPath.startsWith("<")) { - return `extension:${extensionPath.replace(/[<>]/g, "")}`; - } - const base = basename(extensionPath); - const name = base.replace(/\.(ts|js)$/, ""); - return `extension:${name}`; - } - - private _applyExtensionBindings(runner: ExtensionRunner): void { - runner.setUIContext(this._extensionUIContext); - runner.bindCommandContext(this._extensionCommandContextActions); - - this._extensionErrorUnsubscriber?.(); - this._extensionErrorUnsubscriber = this._extensionErrorListener - ? runner.onError(this._extensionErrorListener) - : undefined; - } - - private _refreshCurrentModelFromRegistry(): void { - const currentModel = this.model; - if (!currentModel) { - return; - } - - const refreshedModel = this._modelRegistry.find(currentModel.provider, currentModel.id); - if (!refreshedModel || refreshedModel === currentModel) { - return; - } - - this.agent.state.model = refreshedModel; - } - - private _bindExtensionCore(runner: ExtensionRunner): void { - const getCommands = (): SlashCommandInfo[] => { - const extensionCommands: SlashCommandInfo[] = runner.getRegisteredCommands().map((command) => ({ - name: command.invocationName, - description: command.description, - source: "extension", - sourceInfo: command.sourceInfo, - })); - - const templates: SlashCommandInfo[] = this.promptTemplates.map((template) => ({ - name: template.name, - description: template.description, - source: "prompt", - sourceInfo: template.sourceInfo, - })); - - const skills: SlashCommandInfo[] = this._resourceLoader.getSkills().skills.map((skill) => ({ - name: `skill:${skill.name}`, - description: skill.description, - source: "skill", - sourceInfo: skill.sourceInfo, - })); - - return [...extensionCommands, ...templates, ...skills]; - }; - - runner.bindCore( - { - sendMessage: (message, options) => { - this.sendCustomMessage(message, options).catch((err) => { - runner.emitError({ - extensionPath: "", - event: "send_message", - error: err instanceof Error ? err.message : String(err), - }); - }); - }, - sendUserMessage: (content, options) => { - this.sendUserMessage(content, options).catch((err) => { - runner.emitError({ - extensionPath: "", - event: "send_user_message", - error: err instanceof Error ? err.message : String(err), - }); - }); - }, - appendEntry: (customType, data) => { - this.sessionManager.appendCustomEntry(customType, data); - }, - setSessionName: async (name) => { - if (this._agentMessageController?.setSessionName) { - await this._agentMessageController.setSessionName(name); - return; - } - this.setSessionName(name); - }, - getSessionName: () => { - return this.sessionManager.getSessionName(); - }, - setLabel: (entryId, label) => { - this.sessionManager.appendLabelChange(entryId, label); - }, - getActiveTools: () => this.getActiveToolNames(), - getAllTools: () => this.getAllTools(), - setActiveTools: (toolNames) => this.setActiveToolsByName(toolNames), - refreshTools: () => this._refreshToolRegistry(), - getCommands, - setModel: async (model) => { - if (!this.modelRegistry.hasConfiguredAuth(model)) return false; - await this.setModel(model); - return true; - }, - getThinkingLevel: () => this.thinkingLevel, - setThinkingLevel: (level) => this.setThinkingLevel(level), - }, - { - getModel: () => this.model, - isIdle: () => !this.isStreaming, - getSignal: () => this.agent.signal, - abort: () => this.abort(), - hasPendingMessages: () => this.queuedActionCount > 0, - shutdown: () => { - this._extensionShutdownHandler?.(); - }, - getContextUsage: () => this.getContextUsage(), - compact: (options) => { - void (async () => { - try { - const result = await this.compact(options?.customInstructions); - options?.onComplete?.(result); - } catch (error) { - const err = error instanceof Error ? error : new Error(String(error)); - options?.onError?.(err); - } - })(); - }, - getSystemPrompt: () => this.systemPrompt, - }, - { - registerProvider: (name, config) => { - this._modelRegistry.registerProvider(name, config); - this._refreshCurrentModelFromRegistry(); - }, - unregisterProvider: (name) => { - this._modelRegistry.unregisterProvider(name); - this._refreshCurrentModelFromRegistry(); - }, - }, - ); - } - - private _refreshToolRegistry(options?: { activeToolNames?: string[]; includeAllExtensionTools?: boolean }): void { - const previousRegistryNames = new Set(this._toolRegistry.keys()); - const previousActiveToolNames = this.getActiveToolNames(); - const allowedToolNames = this._allowedToolNames; - const registeredTools = this._extensionRunner.getAllRegisteredTools(); - const sdkToolEntry = (definition: ToolDefinition) => ({ - definition, - sourceInfo: createSyntheticSourceInfo(``, { - source: "sdk" as const, - }), - }); - const allCustomTools = [ - ...registeredTools, - ...this._customTools.map(sdkToolEntry), - ...this._acpMcpTools.map(sdkToolEntry), - ]; - const isAllowedTool = (name: string): boolean => !allowedToolNames || allowedToolNames.has(name); - const allowedCustomTools = allCustomTools.filter((tool) => isAllowedTool(tool.definition.name)); - const definitionRegistry = new Map( - Array.from(this._baseToolDefinitions.entries()) - .filter(([name]) => isAllowedTool(name)) - .map(([name, definition]) => [ - name, - { - definition, - sourceInfo: createSyntheticSourceInfo(``, { - source: "builtin", - }), - }, - ]), - ); - for (const tool of allowedCustomTools) { - definitionRegistry.set(tool.definition.name, { - definition: tool.definition, - sourceInfo: tool.sourceInfo, - }); - } - this._toolDefinitions = definitionRegistry; - this._toolPromptSnippets = new Map( - Array.from(definitionRegistry.values()) - .map(({ definition }) => { - const snippet = this._normalizePromptSnippet(definition.promptSnippet); - return snippet ? ([definition.name, snippet] as const) : undefined; - }) - .filter((entry): entry is readonly [string, string] => entry !== undefined), - ); - this._toolPromptGuidelines = new Map( - Array.from(definitionRegistry.values()) - .map(({ definition }) => { - const guidelines = this._normalizePromptGuidelines(definition.promptGuidelines); - return guidelines.length > 0 ? ([definition.name, guidelines] as const) : undefined; - }) - .filter((entry): entry is readonly [string, string[]] => entry !== undefined), - ); - const runner = this._extensionRunner; - const wrappedExtensionTools = wrapRegisteredTools(allowedCustomTools, runner); - // Resolve the runner at call time so a rebuild/reload rebinds built-in tools to the - // live runner instead of wedging them on the invalidated one's stale-ctx guard. - const wrappedBuiltInTools = wrapRegisteredTools( - Array.from(this._baseToolDefinitions.values()) - .filter((definition) => isAllowedTool(definition.name)) - .map((definition) => ({ - definition, - sourceInfo: createSyntheticSourceInfo(``, { source: "builtin" }), - })), - () => this._extensionRunner, - ); - - const toolRegistry = new Map(wrappedBuiltInTools.map((tool) => [tool.name, tool])); - for (const tool of wrappedExtensionTools as AgentTool[]) { - toolRegistry.set(tool.name, tool); - } - this._toolRegistry = toolRegistry; - - const nextActiveToolNames = ( - options?.activeToolNames ? [...options.activeToolNames] : [...previousActiveToolNames] - ).filter((name) => isAllowedTool(name)); - - if (allowedToolNames) { - for (const toolName of this._toolRegistry.keys()) { - if (allowedToolNames.has(toolName)) { - nextActiveToolNames.push(toolName); - } - } - } else if (options?.includeAllExtensionTools) { - for (const tool of wrappedExtensionTools) { - nextActiveToolNames.push(tool.name); - } - } else if (!options?.activeToolNames) { - for (const toolName of this._toolRegistry.keys()) { - if (!previousRegistryNames.has(toolName)) { - nextActiveToolNames.push(toolName); - } - } - } - - this.setActiveToolsByName([...new Set(nextActiveToolNames)]); + private _refreshToolRegistry(options?: { activeToolNames?: string[]; includeAllExtensionTools?: boolean }): void { + this._tools.refreshToolRegistry(options); } private _buildRuntime(options: { @@ -6525,113 +5759,16 @@ export class AgentSession { includeAllExtensionTools?: boolean; }): void { const pythonSkills = getPythonSkillRuntimeInfo(this._modelVisibleSkills()); - let configuredBaseToolDefinitions: Record; - if (this._baseToolsOverride) { - configuredBaseToolDefinitions = Object.fromEntries( - Object.entries(this._baseToolsOverride).map(([name, tool]) => [ - name, - createToolDefinitionFromAgentTool(tool), - ]), - ); - } else { - // Rebuilding (e.g. /reload) replaces the provisioner; drop the previous - // kernel so the session never holds two live kernels. Gate the new kernel's - // startup on the old one's dispose (which flushes a final snapshot), so a - // reload can't restore from a snapshot the old kernel is still writing. - const previousDispose = this._ipythonKernelProvisioner?.dispose(); - this._ipythonKernelSnapshotDir = this.sessionManager.getSessionArtifactDir(); - // Only surface the "revived from your previous session" notice on the first - // build (a genuine resume). A later rebuild (/reload) restores state silently - // for continuity — the conversation is unchanged, so there's nothing to flag. - const notifyRestore = !this._ipythonRuntimeBuilt; - this._ipythonKernelProvisioner = new IpythonKernelProvisioner(this._cwd, { - env: this._rlmKernelEnv(), - commandPrefix: this.settingsManager.getShellCommandPrefix(), - shellPath: this.settingsManager.getShellPath(), - sessionId: this.sessionId, - hostHandlers: this._createKernelHostHandlers(), - pythonSkills, - snapshotDir: this._ipythonKernelSnapshotDir, - readyGate: previousDispose, - onRestore: notifyRestore ? (result) => this._onIpythonStateRestored(result) : undefined, - }); - configuredBaseToolDefinitions = createAllToolDefinitions(this._cwd, { - ipython: { - provisioner: this._ipythonKernelProvisioner, - commandPrefix: this.settingsManager.getShellCommandPrefix(), - shellPath: this.settingsManager.getShellPath(), - onLateSentAgentMessage: (toolCallId, message) => - this._recordLateIpythonSentAgentMessage(toolCallId, message), - }, - }); - } - - this._baseToolDefinitions = new Map( - Object.entries(configuredBaseToolDefinitions).map(([name, tool]) => [name, tool as ToolDefinition]), - ); - - const extensionsResult = this._resourceLoader.getExtensions(); - if (options.flagValues) { - for (const [name, value] of options.flagValues) { - extensionsResult.runtime.flagValues.set(name, value); - } - } - // Re-apply on (re)build so the provider survives /reload. Guarded: the - // runtime object can be shared across sessions from one ResourceLoader - // (RLM children), so a provider-less session must not wipe the owner's. - if (this._execEnvProvider) { - extensionsResult.runtime.getExecEnv = this._execEnvProvider; - } - - this._extensionRunner = new ExtensionRunner( - extensionsResult.extensions, - extensionsResult.runtime, - this._cwd, - this.sessionManager, - this._modelRegistry, - ); - if (this._extensionRunnerRef) { - this._extensionRunnerRef.current = this._extensionRunner; - } - this._bindExtensionCore(this._extensionRunner); - this._applyExtensionBindings(this._extensionRunner); - - const previousAcpMcpToolNames = new Set(this._acpMcpTools.map((tool) => tool.name)); - const acpServers = this._mcpManager?.getAcpServers() ?? []; - if (acpServers.length > 0 && !this._ipythonKernelProvisioner) { - throw new Error("ACP MCP servers require the built-in cpython tool"); - } - const acpMcpTools = this._ipythonKernelProvisioner - ? createAcpMcpToolDefinitions(acpServers, this._ipythonKernelProvisioner) - : []; - this._assertAcpMcpToolNamesAvailable(acpMcpTools.map((tool) => tool.name)); - for (const name of previousAcpMcpToolNames) this._allowedToolNames?.delete(name); - for (const tool of acpMcpTools) this._allowedToolNames?.add(tool.name); - this._acpMcpTools = acpMcpTools; - - const defaultActiveToolNames = this._baseToolsOverride ? Object.keys(this._baseToolsOverride) : ["ipython"]; - const baseActiveToolNames = [...(options.activeToolNames ?? defaultActiveToolNames)]; - if (this._goals.state.status === "active" && this._includeGoals) { - // An active goal needs ipython so the model can reach the goal skill. - baseActiveToolNames.push("ipython"); - } + this._tools.setBaseDefinitions(this._tools.buildBaseOverrides() ?? this._kernel.build(pythonSkills)); + this._extensions.build(options.flagValues); + this._tools.updateAcpDefinitions(); + const baseActiveToolNames = [...(options.activeToolNames ?? this._tools.defaultActiveToolNames)]; + if (this._goals.state.status === "active" && this._includeGoals) baseActiveToolNames.push("ipython"); this._refreshToolRegistry({ activeToolNames: [...new Set(baseActiveToolNames)], includeAllExtensionTools: options.includeAllExtensionTools, }); - - // Prewarm when configured, or whenever we're resuming a session that already - // has a kernel snapshot — so its state is revived and the model is told what - // came back before the first turn, rather than a turn later when the kernel - // would otherwise lazily start on first use. - const hasSnapshot = - !!this._ipythonKernelSnapshotDir && existsSync(snapshotPathIn(this._ipythonKernelSnapshotDir)); - if ((this._prewarmIpythonKernel || hasSnapshot) && this.getActiveToolNames().includes("ipython")) { - this._ipythonKernelProvisioner?.prewarm(); - } - - // Subsequent builds are in-process rebuilds (/reload), not a fresh resume. - this._ipythonRuntimeBuilt = true; + this._kernel.finishBuild(this.getActiveToolNames()); } /** @@ -6662,260 +5799,81 @@ export class AgentSession { } private _createKernelHostHandlers(): HostRequestHandlers { - const handlers: HostRequestHandlers = { - "rlm.run": createRlmRunHostHandler(async ({ prompt, kwargs, cellSourceCode }) => ({ - ...(await this.runRlmChild(prompt, kwargs, cellSourceCode)), - })), - "rlm.create_session": createRlmCreateSessionHostHandler(async ({ prompt, kwargs }) => ({ - ...(await this.createRlmSession(prompt, kwargs)), - })), - "bash.completed": createAsyncBashCompletionHostHandler(async (details) => { - const message = createAsyncBashCompletionMessage(details); - const disposeSignal = this._commitFence.disposeSignal; - while (true) { - let admissionCommitted = false; - try { - await this._promptInjectedMessage(message.content, message, { - streamingBehavior: "steer", - queueIfBusy: true, - resumeIfIdle: true, - returnAfterAccepted: true, - suppressAutonomousContinuation: true, - admissionCommitted: () => { - admissionCommitted = true; - }, - }); - return; - } catch (error) { - if (admissionCommitted || !(error instanceof SessionInputAdmissionPausedError)) throw error; - while (this._inputScheduler.admissionPaused && !disposeSignal.aborted) { - await this._waitForSessionActivityChange(disposeSignal); - } - } - } - }), - "bash.consumed": createAsyncBashConsumedHostHandler((details) => { - this._withdrawAsyncBashCompletionNotice(details); - }), - "rlm.find_models": createRlmFindModelsHostHandler((query, limit) => this.findRlmModels(query, limit)), - "rlm.list_subagents": createRlmListSubagentsHostHandler(() => this.listRlmSubagents()), - "rlm.delete_subagent": createRlmDeleteSubagentHostHandler((target) => this.deleteRlmSubagent(target)), - "model.info": async () => ({ - id: this.model?.id ?? null, - provider: this.model?.provider ?? null, - input: this.model?.input ?? [], - }), - }; - if (this._includeGoals) { - for (const type of ["goal.get", "goal.create", "goal.complete"]) { - handlers[type] = async (payload) => this.handleGoalHostRequest(type, payload); - } - } - if (this._includeCompactSkill) { - for (const type of ["compact.run", "compact.status"]) { - handlers[type] = async (payload) => this.handleCompactHostRequest(type, payload); - } - } - if (this._refinement._autoRefineAllowedForSession()) { - for (const type of ["refine.run", "refine.status"]) { - handlers[type] = async (payload) => this.handleRefineHostRequest(type, payload); - } - } - if (this._rlmHeartbeatController) { - for (const type of [ - "rlm_heartbeat.list", - "rlm_heartbeat.create", - "rlm_heartbeat.update", - "rlm_heartbeat.delete", - ]) { - handlers[type] = async (payload) => this.handleRlmHeartbeatHostRequest(type, payload); - } - } - const visibleKernelSkillNames = new Set( - this._modelVisibleSkills() - .filter((skill) => !skill.disableModelInvocation) - .map((skill) => skill.name), - ); - const messageController = this._agentMessageController; - if (messageController && visibleKernelSkillNames.has(AGENT_MESSAGE_SKILL_NAME)) { - Object.assign( - handlers, - createAgentMessageHostHandlers({ - family: async () => { - if (!messageController.family) - throw new Error("agent family roster is not available in this session"); - return messageController.family(); - }, - awaitPendingChildPublication: (selector) => this._awaitPendingRlmChildPublication(selector), - sendAgentMessage: async (input) => { - const receipt = (await this.handleAgentMessageHostRequest("agent_message.send", { - target: input.target, - message: input.message, - })) as AgentSessionMessageReceipt; - if (this._rlmDepth > 0) { - let addressedParent = input.receiverRole === "parent"; - if (input.receiverRole === undefined && messageController.family) { - try { - addressedParent = (await messageController.family()).some( - (member) => - member.relationship === "parent" && - (member.entry.id === input.target || - agentFamilyMemberName(member.entry) === input.target), - ); - } catch { - addressedParent = false; - } - } - if (addressedParent) { - this._repliedToParentSinceTask = true; - this._parentReplyCount += 1; - } - } - return receipt; - }, - }), - ); - } - if (this._agentObserveController) { - Object.assign( - handlers, - createAgentObserveHostHandlers({ - listAgents: () => this.handleAgentObserveHostRequest("agent_observe.list") as AgentObserveListResult, - getAgent: (target) => - this.handleAgentObserveHostRequest("agent_observe.get", { - target, - }) as AgentObserveAgentSnapshot, - recentMessages: (input) => - this.handleAgentObserveHostRequest("agent_observe.recent", { - target: input.target, - limit: input.limit, - max_chars: input.maxChars, - }) as AgentObserveRecentMessagesResult, - }), - ); - } - if (this._mcpManager) { - Object.assign(handlers, this._mcpManager.hostHandlers()); - } - return handlers; - } - - async reload(): Promise { - const previousFlagValues = this._extensionRunner.getFlagValues(); - await emitSessionShutdownEvent(this._extensionRunner, { - type: "session_shutdown", - reason: "reload", - }); - await this.settingsManager.reload(); - // Re-read auth.json: a login saved by the client process (daemon mode) must be - // visible here so MCP skill gating sees the new credentials. - this._modelRegistry.authStorage.reload(); - resetApiProviders(); - this._mcpManager?.refresh(); - await this._resourceLoader.reload(); - this._buildRuntime({ - activeToolNames: this.getActiveToolNames(), - flagValues: previousFlagValues, - includeAllExtensionTools: true, + return createSessionKernelHostHandlers({ + runChild: (prompt, kwargs, code) => this.runRlmChild(prompt, kwargs, code), + createSession: (prompt, kwargs) => this.createRlmSession(prompt, kwargs), + findModels: (query, limit) => this.findRlmModels(query, limit), + listSubagents: () => this.listRlmSubagents(), + deleteSubagent: (target) => this.deleteRlmSubagent(target), + handleBashCompletion: (details) => this._handleKernelBashCompletion(details), + withdrawBashCompletion: (details) => this._withdrawAsyncBashCompletionNotice(details), + getModel: () => this.model, + includeGoals: this._includeGoals, + includeCompactSkill: this._includeCompactSkill, + isRefineAllowed: () => this._refinement._autoRefineAllowedForSession(), + hasHeartbeatController: () => !!this._rlmHeartbeatController, + getModelVisibleSkills: () => this._modelVisibleSkills(), + getAgentMessageController: () => this._agentMessageController, + hasObserveController: () => !!this._agentObserveController, + getMcpManager: () => this._mcpManager, + getDepth: () => this._childState.depth, + awaitChildPublication: (selector) => this._awaitPendingRlmChildPublication(selector), + recordParentReply: () => this._childState.recordReply(), + handleGoal: (type, payload) => this.handleGoalHostRequest(type, payload), + handleCompact: (type, payload) => this.handleCompactHostRequest(type, payload), + handleRefine: (type, payload) => this.handleRefineHostRequest(type, payload), + handleHeartbeat: (type, payload) => this.handleRlmHeartbeatHostRequest(type, payload), + handleMessage: (type, payload) => this.handleAgentMessageHostRequest(type, payload), + handleObserve: (type, payload) => this.handleAgentObserveHostRequest(type, payload), }); - - const hasBindings = - this._extensionUIContext || - this._extensionCommandContextActions || - this._extensionShutdownHandler || - this._extensionErrorListener; - if (hasBindings) { - await this._extensionRunner.emit({ - type: "session_start", - reason: "reload", - }); - await this.extendResourcesFromExtensions("reload"); - } } - private _rlmKernelEnv(): Record { - // Kernel env is provisioning-time only: RLM_MAX_DEPTH may be stale in an already-running kernel; - // the TypeScript-side spawn check remains authoritative. - const env: Record = { - RLM_DEPTH: String(this._rlmDepth), - RLM_MAX_DEPTH: String(this._rlmMaxDepth), - RLM_GLOBAL_HARNESS_STATE_DIR: getGlobalHarnessStateDir(), - }; - const rlmSessionDir = this._ensureRlmSessionDir(); - if (rlmSessionDir) { - env.RLM_SESSION_DIR = rlmSessionDir; - // Keep kernel writes and host reads (system prompt, review, /refine) on - // the same local harness path. Subagents prefer their own artifact dir; - // ephemeral sessions fall back to the RLM session dir once it exists. - env.RLM_HARNESS_STATE_DIR = - this._refinement._localHarnessStateDir() ?? getLocalHarnessStateDir(rlmSessionDir)!; + private async _handleKernelBashCompletion(details: AsyncBashCompletionDetails): Promise { + const message = createAsyncBashCompletionMessage(details); + const disposeSignal = this._commitFence.disposeSignal; + while (true) { + let admissionCommitted = false; + try { + await this._promptInjectedMessage(message.content, message, { + streamingBehavior: "steer", + queueIfBusy: true, + resumeIfIdle: true, + returnAfterAccepted: true, + suppressAutonomousContinuation: true, + admissionCommitted: () => { + admissionCommitted = true; + }, + }); + return; + } catch (error) { + if (admissionCommitted || !(error instanceof SessionInputAdmissionPausedError)) throw error; + while (this._inputScheduler.admissionPaused && !disposeSignal.aborted) { + await this._waitForSessionActivityChange(disposeSignal); + } + } } - this._addWebsearchKeyEnv(env); - return env; } - private _addWebsearchKeyEnv(env: Record): void { - if (this._agentDir) { - env.PRIME_AGENT_CODING_AGENT_DIR = this._agentDir; - } - - if (process.env[SERPER_ENV_VAR]?.trim()) { - return; - } - // Inject only when a websearch skill (bundled or custom) is actually loaded, - // so the key isn't exposed to kernels that can't use it. - if (!this._resourceLoader.getSkills().skills.some((skill) => skill.name === WEBSEARCH_SKILL_NAME)) { - return; - } - const cred = this._modelRegistry.authStorage.get(SERPER_CREDENTIAL_ID); - if (cred?.type !== "api_key") { - return; - } - const resolved = resolveConfigValue(cred.key)?.trim(); - if (resolved) { - env[SERPER_ENV_VAR] = resolved; - } + reload(): Promise { + return this._extensions.reload(); } // Undefined when there's no persistent artifact dir (e.g. the viewer client): // don't mkdtemp here, since this runs on every kernel build but a viewer never // does RLM work. The temp dir is created lazily in _createChildRlmSessionDir. - private _ensureRlmSessionDir(): string | undefined { - if (this._rlmSessionDir) { - mkdirSync(this._rlmSessionDir, { recursive: true }); - return this._rlmSessionDir; - } - - const sessionArtifactDir = this.sessionManager.getSessionArtifactDir(); - if (sessionArtifactDir) { - mkdirSync(sessionArtifactDir, { recursive: true }); - this._rlmSessionDir = sessionArtifactDir; - return sessionArtifactDir; - } - - return undefined; - } private _createChildRlmSessionDir(): string { - const parentDir = this._ensureRlmSessionDir() ?? this._createEphemeralRlmSessionDir(); - for (let i = 0; i < 100; i++) { - const childDir = join(parentDir, `sub-${randomUUID().slice(0, 8)}`); - try { - mkdirSync(childDir); - return childDir; - } catch (error) { - if (error instanceof Error && "code" in error && error.code === "EEXIST") { - continue; - } - throw error; - } - } - throw new Error("Unable to create unique RLM child session directory"); + return createChildSessionDir(() => this._ensureRlmSessionDir() ?? this._createEphemeralRlmSessionDir()); } + private _rlmKernelEnv(): Record { + return this._kernelEnvironment.buildEnv(); + } + private _ensureRlmSessionDir(): string | undefined { + return this._kernelEnvironment.ensureSessionDir(); + } private _createEphemeralRlmSessionDir(): string { - this._rlmSessionDir = mkdtempSync(join(tmpdir(), "prime-agent-rlm-")); - return this._rlmSessionDir; + return this._kernelEnvironment.createEphemeralSessionDir(); } _contextTokensForCurrentMessages(): number | undefined { @@ -6924,23 +5882,15 @@ export class AgentSession { } setCurrentRecap(recap: string | undefined): void { - if (this._currentRecap === recap) return; - this._currentRecap = recap; - this._emit({ type: "recap_update", recap }); + this._childState.setCurrentRecap(recap); } get repliedToParentSinceTask(): boolean | undefined { - return this._repliedToParentSinceTask; + return this._childState.repliedSinceTask; } getCurrentRecap(): string | undefined { - return this._currentRecap; - } - - private _findAssistantEntryForMessage(message: AssistantMessage): SessionMessageEntry | undefined { - return this.sessionManager - .getEntries() - .find((entry): entry is SessionMessageEntry => entry.type === "message" && entry.message === message); + return this._childState.getCurrentRecap(); } private _createRlmSubagentRuntimeOptions(options: { @@ -6961,563 +5911,71 @@ export class AgentSession { spawnCode: options.spawnCode, sessionDir: options.sessionDir, model: options.model, - thinkingLevel: - options.thinkingLevel ?? (clampThinkingLevel(options.model, this.thinkingLevel) as ThinkingLevel), - serviceTier: - this.serviceTier === "priority" && !supportsFastMode(options.model) ? "default" : this.serviceTier, - scopedModels: [...this._scopedModels], - activeToolNames: this.getActiveToolNames(), - allowedToolNames: this._allowedToolNames ? [...this._allowedToolNames] : undefined, - customTools: [...this._customTools], - includeGoals: this._includeGoals, - includeCompactSkill: this._includeCompactSkill, - rlmDepth: this._rlmDepth + 1, - rlmMaxDepth: this._rlmMaxDepth, - rlmParentNodeId: options.id, - spawnedByRequestId: options.spawnedByRequestId, - }; - } - - private async _createRlmSubagentRuntime(options: CreateRlmSubagentRuntimeOptions): Promise { - if (this._subagentRuntimeHost) { - return await this._subagentRuntimeHost.createRlmSubagentRuntime(options); - } - - return this._createInlineRlmSubagentRuntime(options); - } - - private _createInlineRlmSubagentRuntime(options: CreateRlmSubagentRuntimeOptions): RlmSubagentRuntime { - const childSessionManager = SessionManager.create(this._cwd, options.sessionDir); - if (options.parentSession.sessionFile) { - childSessionManager.newSession({ - parentSession: options.parentSession.sessionFile, - rlmDepth: options.rlmDepth, - }); - } - childSessionManager.appendModelChange(options.model.provider, options.model.id); - childSessionManager.appendThinkingLevelChange(options.thinkingLevel); - childSessionManager.appendServiceTierChange(options.serviceTier); - - const childAgent = new Agent({ - initialState: { - systemPrompt: "", - model: options.model, - thinkingLevel: options.thinkingLevel, - serviceTier: options.serviceTier, - tools: [], - }, - convertToLlm: this.agent.convertToLlm, - transformContext: this.agent.transformContext, - streamFn: this.agent.streamFn, - getApiKey: this.agent.getApiKey, - onPayload: this.agent.onPayload, - onResponse: this.agent.onResponse, - steeringMode: this.settingsManager.getSteeringMode(), - followUpMode: this.settingsManager.getFollowUpMode(), - sessionId: childSessionManager.getSessionId(), - thinkingBudgets: this.settingsManager.getThinkingBudgets(), - transport: this.settingsManager.getTransport(), - toolExecution: this.agent.toolExecution, - }); - - const child = new AgentSession({ - agent: childAgent, - sessionManager: childSessionManager, - settingsManager: this.settingsManager, - cwd: this._cwd, - agentDir: this._agentDir, - scopedModels: options.scopedModels, - resourceLoader: this._resourceLoader, - customTools: options.customTools, - modelRegistry: this._modelRegistry, - initialActiveToolNames: options.activeToolNames, - allowedToolNames: options.allowedToolNames, - includeGoals: options.includeGoals, - includeCompactSkill: options.includeCompactSkill, - rlmDepth: options.rlmDepth, - rlmMaxDepth: options.rlmMaxDepth, - rlmSessionDir: options.sessionDir, - rlmParentNodeId: options.rlmParentNodeId, - rlmParentAgent: options.parentSession.sessionName ?? options.parentSession.sessionId, - semanticParentSessionId: options.parentSession.sessionId, - semanticSpawnedByRequestId: options.spawnedByRequestId, - sessionStartEvent: { type: "session_start", reason: "startup" }, - }); - if (child.sessionName !== options.sessionName) { - try { - child.setSessionName(options.sessionName); - } catch (error) { - child.dispose(); - throw error; - } - } - options.onSessionPublished?.(child); - - return { session: child }; - } - - private _abandonRlmRunForQuiescence(run: RlmChildRun): void { - run.suppressTerminalNotice = true; - run.abandonedForQuiescence = true; - this._abandonedRlmQuiescenceChildIds.add(run.id); - this._unsettledRlmChildRuns.delete(run); - run.settlement.resolve(); - this._maybeResumeGoalContinuationAfterRlmWork(); - } - - private _cancelActiveRlmChildRuns(reason: string): void { - for (const run of this._activeRlmChildRuns.values()) { - this._cancelRlmChildRun(run, reason); - } - } - - private _cancelRlmChildRun(run: RlmChildRun, reason: string): boolean { - if (run.status !== "running" && run.status !== "queued") { - return false; - } - run.status = "cancelled"; - if (this._inputScheduler.suspended) this._abandonRlmRunForQuiescence(run); - run.error = reason; - run.publication.reject(new Error(reason)); - run.abort(); - // Surface the cancellation immediately; the run's own terminal update is - // delayed indefinitely when the child is stuck mid-stream, which is - // exactly when users reach for the kill. - run.emitUpdate?.(); - return true; - } - - getRlmChildRunStatus(childId: string): RlmChildAgentStatus | undefined { - return this._activeRlmChildRuns.get(childId)?.status; - } - - private async _currentActiveSessionId(): Promise { - try { - return (await this._agentMessageController?.listAgents())?.current?.activeSessionId; - } catch { - return undefined; - } - } - - private async _awaitPendingRlmChildPublication(selector: string): Promise { - const run = [...this._activeRlmChildRuns.values()].find( - (candidate) => - (candidate.status === "queued" || candidate.status === "running" || candidate.status === "done") && - !candidate.detachedDeletion && - (candidate.id === selector || candidate.sessionName === selector), - ); - if (!run) return undefined; - await run.publication.promise; - return run.session?.sessionId; - } - - async listRlmSubagents(): Promise { - return this._buildRlmSubagentList(await this._agentMessageController?.listAgents()); - } - - private _buildRlmSubagentList(listedAgents?: AgentSessionMessageListResult): RlmListSubagentsResult { - const daemonChildren = new Map(); - const parentActiveSessionId = listedAgents?.current?.activeSessionId; - if (parentActiveSessionId) { - for (const agent of listedAgents.agents) { - if ( - agent.runtimeKind === "subagent" && - agent.parentActiveSessionId === parentActiveSessionId && - agent.rlmChildId - ) { - daemonChildren.set(agent.rlmChildId, agent); - } - } - } - - const subagents: RlmListSubagentsResult["subagents"] = []; - const recorded = new Set(); - for (const run of this._activeRlmChildRuns.values()) { - if (this._deletingRlmChildren.has(run.id) || run.detachedDeletion || run.status === "cancelled") { - continue; - } - const daemonChild = daemonChildren.get(run.id); - subagents.push({ - rlm_child_id: run.id, - active_session_id: daemonChild?.activeSessionId ?? null, - session_id: daemonChild?.sessionId ?? run.session?.sessionId ?? null, - session_name: daemonChild?.sessionName ?? run.session?.sessionName ?? run.sessionName, - session_dir: run.sessionDir, - status: run.status === "done" ? "completed" : run.status === "error" ? "error" : "running", - }); - recorded.add(run.id); - } - for (const [childId, { session: childSession }] of this._rlmChildSessions) { - if ( - this._deletingRlmChildren.has(childId) || - recorded.has(childId) || - this._rlmChildCleanupFailures.has(childId) - ) { - continue; - } - const daemonChild = daemonChildren.get(childId); - const sessionDir = childSession._rlmSessionDir; - if (!sessionDir) { - continue; - } - subagents.push({ - rlm_child_id: childId, - active_session_id: daemonChild?.activeSessionId ?? null, - session_id: daemonChild?.sessionId ?? childSession.sessionId, - session_name: - daemonChild?.sessionName ?? childSession.sessionName ?? createDefaultRlmSubagentSessionName("", childId), - session_dir: sessionDir, - status: "completed", - }); - recorded.add(childId); - } - for (const [childId, daemonChild] of daemonChildren) { - if ( - recorded.has(childId) || - this._deletingRlmChildren.has(childId) || - this._deletedRlmChildIds.has(childId) || - this._rlmChildCleanupFailures.has(childId) || - !daemonChild.sessionDir - ) { - continue; - } - subagents.push({ - rlm_child_id: childId, - active_session_id: daemonChild.activeSessionId, - session_id: daemonChild.sessionId, - session_name: daemonChild.sessionName ?? createDefaultRlmSubagentSessionName("", childId), - session_dir: daemonChild.sessionDir, - status: daemonChild.rlmChildRegistryStatus === "completed" ? "completed" : "error", - }); - } - return { subagents }; - } - - private _rlmSubagentMatchesTarget(entry: RlmSubagentRegistryEntry, target: string): boolean { - return ( - entry.rlm_child_id === target || - entry.active_session_id === target || - entry.session_id === target || - entry.session_name === target - ); - } - - private async _resolveDirectRlmSubagent(target: string): Promise { - const candidates = [...(await this.listRlmSubagents()).subagents, ...this._rlmChildCleanupFailures.values()]; - const matches = candidates.filter((entry) => this._rlmSubagentMatchesTarget(entry, target)); - if (matches.length === 0) { - throw new Error(`No direct RLM subagent matches "${target}" in the current parent session`); - } - if (matches.length > 1) { - throw new Error(`RLM subagent selector "${target}" is ambiguous in the current parent session`); - } - return matches[0]!; - } - - async deleteInactiveRlmSubagent( - childId: string, - isExternallyRunning: () => boolean = () => false, - ): Promise<"deleted" | "not_found" | "running"> { - for (const owner of this._rlmSubtreeSessions()) { - const isRunning = (): boolean => { - const status = owner._activeRlmChildRuns.get(childId)?.status; - return status === "queued" || status === "running" || isExternallyRunning(); - }; - if (isRunning()) { - return "running"; - } - const subagent = [ - ...(await owner.listRlmSubagents()).subagents, - ...owner._rlmChildCleanupFailures.values(), - ].find((entry) => entry.rlm_child_id === childId); - if (!subagent) continue; - if (isRunning()) { - return "running"; - } - const result = await owner._trackRlmSubagentDeletion(subagent, () => { - if (isRunning()) { - return Promise.resolve({ subagent, outcome: "skipped_running" }); - } - return owner._deleteResolvedRlmSubagent(subagent); - }); - return result.outcome === "skipped_running" ? "running" : "deleted"; - } - return "not_found"; - } - - async deleteRlmSubagent(target: string): Promise { - const inFlight = [...this._deletingRlmChildren.values()].filter(({ subagent }) => - this._rlmSubagentMatchesTarget(subagent, target), - ); - if (inFlight.length > 1) { - throw new Error(`RLM subagent selector "${target}" is ambiguous in the current parent session`); - } - - // Running and retained children can be reserved synchronously. This keeps - // them hidden immediately while the async daemon listing checks for a - // conflicting passive selector. - const localMatches = [ - ...this._buildRlmSubagentList().subagents, - ...this._rlmChildCleanupFailures.values(), - ].filter((entry) => this._rlmSubagentMatchesTarget(entry, target)); - const matchingChildIds = new Set([ - ...inFlight.map(({ subagent }) => subagent.rlm_child_id), - ...localMatches.map((subagent) => subagent.rlm_child_id), - ]); - if (matchingChildIds.size > 1 || localMatches.length > 1) { - throw new Error(`RLM subagent selector "${target}" is ambiguous in the current parent session`); - } - if (inFlight[0]) { - return inFlight[0].promise; - } - if (localMatches[0]) { - const subagent = localMatches[0]; - return this._trackRlmSubagentDeletion(subagent, async () => { - const listedAgents = await this._agentMessageController?.listAgents(); - const listedSubagents = this._buildRlmSubagentList(listedAgents).subagents; - const passiveMatches = listedSubagents.filter( - (entry) => entry.rlm_child_id !== subagent.rlm_child_id && this._rlmSubagentMatchesTarget(entry, target), - ); - if (passiveMatches.length > 0) { - throw new Error(`RLM subagent selector "${target}" is ambiguous in the current parent session`); - } - const parentActiveSessionId = listedAgents?.current?.activeSessionId; - const daemonChild = listedAgents?.agents.find( - (agent) => - agent.rlmChildId === subagent.rlm_child_id && agent.parentActiveSessionId === parentActiveSessionId, - ); - const resolvedSubagent = daemonChild - ? { - ...subagent, - active_session_id: daemonChild.activeSessionId, - session_id: daemonChild.sessionId, - session_name: daemonChild.sessionName ?? subagent.session_name, - } - : subagent; - return this._deleteResolvedRlmSubagent(resolvedSubagent); - }); - } - - const directMatches = [ - ...(await this.listRlmSubagents()).subagents, - ...this._rlmChildCleanupFailures.values(), - ].filter((entry) => this._rlmSubagentMatchesTarget(entry, target)); - const directChildIds = new Set(directMatches.map((subagent) => subagent.rlm_child_id)); - if (directChildIds.size > 1) { - throw new Error(`RLM subagent selector "${target}" is ambiguous in the current parent session`); - } - const subagent = directMatches[0] ?? (await this._resolveDirectRlmSubagent(target)); - return this._trackRlmSubagentDeletion(subagent, () => this._deleteResolvedRlmSubagent(subagent)); - } - - private async _trackRlmSubagentDeletion( - subagent: RlmSubagentRegistryEntry, - startDeletion: () => Promise, - ): Promise { - const existing = this._deletingRlmChildren.get(subagent.rlm_child_id); - if (existing) return existing.promise; - const deletion = Promise.resolve().then(startDeletion); - this._deletingRlmChildren.set(subagent.rlm_child_id, { - subagent, - promise: deletion, - }); - try { - return await deletion; - } finally { - const clearReservation = () => { - if (this._deletingRlmChildren.get(subagent.rlm_child_id)?.promise === deletion) { - this._deletingRlmChildren.delete(subagent.rlm_child_id); - } - }; - const run = this._activeRlmChildRuns.get(subagent.rlm_child_id); - if (run?.detachedDeletion) { - // Keep every selector reserved until the run settles, or until a failed - // cleanup is exposed for an explicit retry. Repeated deletes before that - // boundary return the same accepted result. - void run.deletionReservation.promise.then(clearReservation, clearReservation); - } else { - clearReservation(); - } - } + thinkingLevel: + options.thinkingLevel ?? (clampThinkingLevel(options.model, this.thinkingLevel) as ThinkingLevel), + serviceTier: + this.serviceTier === "priority" && !supportsFastMode(options.model) ? "default" : this.serviceTier, + scopedModels: [...this._scopedModels], + activeToolNames: this.getActiveToolNames(), + allowedToolNames: this._allowedToolNames ? [...this._allowedToolNames] : undefined, + customTools: [...this._customTools], + includeGoals: this._includeGoals, + includeCompactSkill: this._includeCompactSkill, + rlmDepth: this._childState.depth + 1, + rlmMaxDepth: this._childState.maxDepth, + rlmParentNodeId: options.id, + spawnedByRequestId: options.spawnedByRequestId, + }; } - private _deleteRlmSubagentSession(childId: string, session?: AgentSession): Promise { - if (this._subagentRuntimeHost) { - return this._subagentRuntimeHost.deleteRlmSubagentRuntime(childId, session); + private async _createRlmSubagentRuntime(options: CreateRlmSubagentRuntimeOptions): Promise { + const host = this._children.getRuntimeHost(); + if (host) { + return await host.createRlmSubagentRuntime(options); } - return session?.disposeAsync() ?? Promise.resolve(); - } - - private _ensureRlmRunDeletionCleanup(run: RlmChildRun, session: AgentSession): Promise { - if (run.deletionCleanup) return run.deletionCleanup; - const cleanup = Promise.resolve().then(() => this._deleteRlmSubagentSession(run.id, session)); - run.deletionCleanup = cleanup; - // Deletion admission is intentionally nonblocking. The detached run owner - // joins this exact promise before settlement and records any failure. - void cleanup.catch(() => undefined); - return cleanup; - } - private async _recordRlmRunDeletionCleanupFailure( - run: RlmChildRun, - subagent: RlmSubagentRegistryEntry, - session: AgentSession, - error: unknown, - ): Promise { - if (this._disposed || this._disposing) { - run.suppressTerminalNotice = true; - await session.disposeAsync().catch(() => undefined); - if (!run.settled) await this._finishRlmRunDeletion(run); - return; - } - run.deletionCleanup = undefined; - run.deletionCleanupObserver = undefined; - run.deletionCleanupFailed = true; - run.session = session; - this._rlmChildCleanupFailures.set(run.id, subagent); - // Make retry admission available before waking the parent model with the - // retry-required notice. - run.deletionReservation.resolve(); - await Promise.resolve(); - await run.reportDeletionCleanupFailure?.(error); - } - - private async _finishRlmRunDeletion(run: RlmChildRun): Promise { - await run.completeDeletion?.(); - if (this._activeRlmChildRuns.get(run.id) === run) { - this._removeRlmSubagentTracking(run.id, run); - } - run.settled = true; - run.settlement.resolve(); - run.deletionReservation.resolve(); - this._unsettledRlmChildRuns.delete(run); - this._maybeResumeGoalContinuationAfterRlmWork(); + return this._createInlineRlmSubagentRuntime(options); } - private _observeRlmRunDeletionCleanup( - run: RlmChildRun, - subagent: RlmSubagentRegistryEntry, - session: AgentSession, - cleanup: Promise, - ): Promise { - if (run.deletionCleanupObserver) return run.deletionCleanupObserver; - const observer = cleanup.then( - () => true, - async (error) => { - await this._recordRlmRunDeletionCleanupFailure(run, subagent, session, error); - return false; + private _createInlineRlmSubagentRuntime(options: CreateRlmSubagentRuntimeOptions): RlmSubagentRuntime { + return createInlineChildRuntime( + { + cwd: this._cwd, + agentDir: this._agentDir, + agent: this.agent, + settingsManager: this.settingsManager, + resourceLoader: this._resourceLoader, + modelRegistry: this._modelRegistry, }, + options, ); - run.deletionCleanupObserver = observer; - void observer.catch(() => undefined); - return observer; } - private _continueFinishedRlmRunDeletion( - run: RlmChildRun, - subagent: RlmSubagentRegistryEntry, - session: AgentSession, - ): void { - const cleanup = this._ensureRlmRunDeletionCleanup(run, session); - const observer = this._observeRlmRunDeletionCleanup(run, subagent, session, cleanup); - if (!run.deletionRunFinished) return; - void observer - .then(async (cleanupSucceeded) => { - if (cleanupSucceeded) await this._finishRlmRunDeletion(run); - }) - .catch(() => undefined); + private _cancelActiveRlmChildRuns(reason: string): void { + this._children.cancelActiveRuns(reason); } - private _removeRlmSubagentTracking(childId: string, run?: RlmChildRun): void { - run?.unsubscribe?.(); - this._rlmChildUnsubscribes.get(childId)?.(); - this._rlmChildUnsubscribes.delete(childId); - this._rlmChildSessions.delete(childId); - this._rlmChildCleanupFailures.delete(childId); - this._abandonedRlmQuiescenceChildIds.delete(childId); - if (!run || this._activeRlmChildRuns.get(childId) === run) { - this._activeRlmChildRuns.delete(childId); - } - if (run) { - run.abort = noopRlmChildAbort; - run.unsubscribe = undefined; - run.session = undefined; - } + getRlmChildRunStatus(childId: string): RlmChildAgentStatus | undefined { + return this._children.getRlmChildRunStatus(childId); } - private _emitRlmSubagentRemoval(subagent: RlmSubagentRegistryEntry): void { - this._emit({ - type: "rlm_child_update", - child: { - id: subagent.rlm_child_id, - parentId: this._rlmParentNodeId, - activeSessionId: subagent.active_session_id ?? undefined, - sessionName: subagent.session_name, - label: subagent.session_name, - status: "cancelled", - sessionDir: subagent.session_dir, - error: "Deleted by parent orchestrator", - }, - }); + private _awaitPendingRlmChildPublication(selector: string): Promise { + return this._children.awaitPublication(selector); } - private async _deleteResolvedRlmSubagent(subagent: RlmSubagentRegistryEntry): Promise { - const childId = subagent.rlm_child_id; - const run = this._activeRlmChildRuns.get(childId); - if (run) { - if (run.deletionCleanupFailed) { - // Reset retry coordination only after selector preflight reaches the - // resolved child. A failed preflight must leave the prior retry boundary - // intact so a later call can acquire it. - run.deletionCleanupFailed = false; - run.deletionFailureNotice = undefined; - run.deletionReservation = createAgentMessageDeferred(); - } - // The detached task remains the sole lifecycle owner. Mark deletion before - // cancellation so its catch/finally path cannot race a normal release or - // terminal notice against the physical delete. - run.detachedDeletion = subagent; - if (this._cancelRlmChildRun(run, "Deleted by parent orchestrator")) { - run.deletionNeedsCompletionNotice = true; - } else { - this._emitRlmSubagentRemoval(subagent); - } - const liveSession = run.session; - if (run.status === "error" && !liveSession && run.settled) { - this._deletedRlmChildIds.add(childId); - this._removeRlmSubagentTracking(childId, run); - return { subagent }; - } - if (liveSession && run.settled) { - run.deletionRunFinished = true; - run.settlement = createAgentMessageDeferred(); - run.settled = false; - this._unsettledRlmChildRuns.add(run); - } - if (liveSession) this._continueFinishedRlmRunDeletion(run, subagent, liveSession); + listRlmSubagents(): Promise { + return this._children.listRlmSubagents(); + } - // Return once deletion is accepted. The run stays hidden but unsettled until - // abort-insensitive model/tool work unwinds and the shared cleanup finishes. - this._deletedRlmChildIds.add(childId); - return { subagent }; - } + deleteInactiveRlmSubagent( + childId: string, + isExternallyRunning: () => boolean = () => false, + ): Promise<"deleted" | "not_found" | "running"> { + return this._children.deleteInactiveRlmSubagent(childId, isExternallyRunning); + } - this._emitRlmSubagentRemoval(subagent); - const retained = this._rlmChildSessions.get(childId)?.session; - try { - await this._deleteRlmSubagentSession(childId, retained); - } catch (error) { - if (this._disposed || this._disposing) { - this._removeRlmSubagentTracking(childId); - void retained?.disposeAsync().catch(() => undefined); - } else { - this._rlmChildCleanupFailures.set(childId, subagent); - } - throw error; - } - this._deletedRlmChildIds.add(childId); - this._removeRlmSubagentTracking(childId); - return { subagent }; + deleteRlmSubagent(target: string): Promise { + return this._children.deleteRlmSubagent(target); } /** @@ -7527,174 +5985,25 @@ export class AgentSession { * matching event forwarder too. */ registerRlmChildSession(childId: string, session: AgentSession, unsubscribe?: () => void): boolean { - // A child can finish concurrently while the parent is (or has) torn down; don't - // resurrect the map (it would never be disposed), just drop the child now. - if (this._deletingRlmChildren.has(childId) || this._deletedRlmChildIds.has(childId)) { - return false; - } - if (this._subagentRuntimeHost?.completeRlmSubagentRuntime?.(childId, session) === false) { - return false; - } - if (this._disposed || this._disposing) { - void session.disposeAsync().catch(() => undefined); - return false; - } - this._rlmChildSessions.set(childId, { session, run: this._activeRlmChildRuns.get(childId) }); - if (unsubscribe) { - this._rlmChildUnsubscribes.set(childId, unsubscribe); - } - return true; + return this._children.registerRlmChildSession(childId, session, unsubscribe); } releaseRlmChildSession(childId: string, session: AgentSession): (() => void) | false { - const run = this._activeRlmChildRuns.get(childId); - if (run?.session === session && run.status === "done") { - const unsubscribe = run.unsubscribe ?? noopRlmChildEventUnsubscribe; - return () => { - run.unsubscribe = undefined; - this._activeRlmChildRuns.delete(childId); - unsubscribe(); - }; - } - if (this._rlmChildSessions.get(childId)?.session !== session) return false; - const unsubscribe = this._rlmChildUnsubscribes.get(childId) ?? noopRlmChildEventUnsubscribe; - return () => { - this._rlmChildUnsubscribes.delete(childId); - this._rlmChildSessions.delete(childId); - unsubscribe(); - }; - } - - private _rlmChildSnapshotForRun( - run: RlmChildRun, - child = run.session ?? this._rlmChildSessions.get(run.id)?.session, - ): RlmChildAgentSnapshot { - const model = child?.model ?? run.model; - return { - id: run.id, - parentId: this._rlmParentNodeId, - sessionName: child?.sessionName ?? run.sessionName, - model: `${model.provider}/${model.id}`, - label: rlmChildLabel(run.prompt), - status: run.status, - durationMs: run.durationMs, - answerPreview: run.answerPreview, - toolUseCount: run.toolUseCount > 0 ? run.toolUseCount : undefined, - tokenCount: child?._contextTokensForCurrentMessages(), - recap: child?.getCurrentRecap(), - sessionDir: run.sessionDir, - activity: run.activity, - repliedSinceTask: child?._repliedToParentSinceTask, - error: run.error, - }; - } - - private _rlmChildSnapshotForSession(childId: string, child: AgentSession): RlmChildAgentSnapshot { - let answerPreview: string | undefined; - let toolUseCount = 0; - const messages = - child.state.streamingMessage?.role === "assistant" - ? [...child.messages, child.state.streamingMessage] - : child.messages; - for (const message of messages) { - if (message.role !== "assistant") continue; - const text = compactRlmText(readAssistantText(message)); - if (text) answerPreview = text; - toolUseCount += message.content.filter((block) => block.type === "toolCall").length; - } - return { - id: childId, - parentId: this._rlmParentNodeId, - sessionName: child.sessionName, - model: child.model ? `${child.model.provider}/${child.model.id}` : undefined, - label: child.sessionName ?? "child agent", - status: "done", - answerPreview, - toolUseCount: toolUseCount > 0 ? toolUseCount : undefined, - tokenCount: child._contextTokensForCurrentMessages(), - recap: child.getCurrentRecap(), - sessionDir: child._rlmSessionDir ?? child.sessionManager.getSessionDir(), - // No run exists (e.g. a child rehydrated after daemon recovery), so live - // session state is the only source for in-flight follow-up work. Mirror - // the run projection's convention: status stays "done" (the recorded task - // finished) and current work surfaces through activity. - activity: child.isSessionActive ? { kind: child.isStreaming ? "writing" : "waiting" } : undefined, - repliedSinceTask: child._repliedToParentSinceTask, - }; - } - - private _isUnboundTerminalRlmChildRun(run: RlmChildRun): boolean { - if (run.session !== undefined || this._rlmChildSessions.has(run.id)) return false; - return run.status === "done" || run.status === "error" || run.status === "cancelled"; + return this._children.releaseRlmChildSession(childId, session); } /** Live recursive child roster from lifecycle state, including nested work under retained parents. */ getRlmChildSnapshots(): RlmChildAgentSnapshot[] { - const snapshots: RlmChildAgentSnapshot[] = []; - const recorded = new Set(); - const traversed = new Set(); - for (const run of this._activeRlmChildRuns.values()) { - const hidden = - run.detachedDeletion || - this._deletingRlmChildren.has(run.id) || - this._deletedRlmChildIds.has(run.id) || - this._isUnboundTerminalRlmChildRun(run); - const child = run.session; - if (!hidden) { - snapshots.push(this._rlmChildSnapshotForRun(run)); - recorded.add(run.id); - } - if (child) { - traversed.add(run.id); - snapshots.push(...child.getRlmChildSnapshots()); - } - } - for (const [childId, { session: child, run }] of this._rlmChildSessions) { - if (recorded.has(childId) || traversed.has(childId)) continue; - const hidden = this._deletingRlmChildren.has(childId) || this._deletedRlmChildIds.has(childId); - if (!hidden) { - const snapshot = run - ? this._rlmChildSnapshotForRun(run, child) - : this._rlmChildSnapshotForSession(childId, child); - snapshots.push({ - ...snapshot, - status: this._rlmChildCleanupFailures.has(childId) ? "cancelled" : snapshot.status, - }); - } - snapshots.push(...child.getRlmChildSnapshots()); - } - return snapshots; + return this._children.getRlmChildSnapshots(); } /** True when any direct or nested subagent is still running or queued. */ hasRunningRlmChildren(): boolean { - for (const session of this._rlmSubtreeSessions()) { - for (const run of session._activeRlmChildRuns.values()) { - if (run.status === "running" || run.status === "queued") { - return true; - } - } - } - return false; - } - - private _rlmChildSessionSnapshot(): AgentSession[] { - const sessions = new Set(); - for (const [childId, { session }] of this._rlmChildSessions) { - if (!this._abandonedRlmQuiescenceChildIds.has(childId)) sessions.add(session); - } - for (const run of this._activeRlmChildRuns.values()) { - if (run.session && !run.abandonedForQuiescence) sessions.add(run.session); - } - return [...sessions]; + return this._children.hasRunningRlmChildren(); } private _hasUnsettledRlmQuiescenceWork(): boolean { - if (this._hasDeferredRlmTerminalNotices()) return true; - if ([...this._unsettledRlmChildRuns].some((run) => !run.settled)) return true; - return this._rlmChildSessionSnapshot().some( - (child) => child.isSessionActive || child._hasUnsettledRlmQuiescenceWork(), - ); + return this._children.hasUnsettledWork(); } /** @@ -7702,60 +6011,13 @@ export class AgentSession { * message and for the resulting parent turns to drain. Re-snapshotting after * each drain includes descendants spawned while earlier results were consumed. */ - async waitForRlmQuiescence(externalSignal?: AbortSignal): Promise { - const cancellation = new AbortController(); - const cancelFromParent = () => cancellation.abort(); - if (externalSignal?.aborted) cancellation.abort(); - else externalSignal?.addEventListener("abort", cancelFromParent, { once: true }); - this._rlmQuiescenceWaitAborts.add(cancellation); - let rejectCancelled = (_error: Error) => {}; - const cancelled = new Promise((_resolve, reject) => { - rejectCancelled = reject; - }); - const onCancelled = () => rejectCancelled(new Error("RLM quiescence wait cancelled")); - cancellation.signal.addEventListener("abort", onCancelled, { once: true }); - if (cancellation.signal.aborted) onCancelled(); - const wait = (operation: Promise): Promise => Promise.race([operation, cancelled]); - try { - while (true) { - await wait(this.waitForHeadlessIdle()); - // Strong RLM quiescence also owns work that interactive waitForIdle ignores. - if (this.isSessionActive || this._hasDeferredRlmTerminalNotices()) { - await wait(this._waitForSessionActivityChange(cancellation.signal)); - continue; - } - const unsettledRuns = [...this._unsettledRlmChildRuns].filter((run) => !run.settled); - const childSessions = this._rlmChildSessionSnapshot(); - if (unsettledRuns.length === 0 && !this._hasUnsettledRlmQuiescenceWork()) return; - await wait( - Promise.all([ - ...unsettledRuns.map((run) => run.settlement.promise), - ...childSessions.map((child) => child.waitForRlmQuiescence(cancellation.signal)), - ]), - ); - // Always loop through the self-active/deferred checks again. Work may - // start at the child-settlement boundary. - } - } finally { - // A local descendant error must cancel sibling recursive waits owned by - // this barrier before their propagation listeners are removed. - cancellation.abort(); - externalSignal?.removeEventListener("abort", cancelFromParent); - cancellation.signal.removeEventListener("abort", onCancelled); - this._rlmQuiescenceWaitAborts.delete(cancellation); - } + waitForRlmQuiescence(externalSignal?: AbortSignal): Promise { + return this._children.waitForRlmQuiescence(externalSignal); } // Inline (non-daemon) mode only; daemon clients attach to the child session directly. getRlmChildSession(childId: string): AgentSession | undefined { - for (const session of this._rlmSubtreeSessions()) { - const direct = - session._activeRlmChildRuns.get(childId)?.session ?? session._rlmChildSessions.get(childId)?.session; - if (direct) { - return direct; - } - } - return undefined; + return this._children.getRlmChildSession(childId); } /** @@ -7765,102 +6027,12 @@ export class AgentSession { * was suppressed; false when the id is unknown or the run already settled. */ cancelRlmChildRun(childId: string, reason = "Cancelled by user"): boolean { - for (const session of this._rlmSubtreeSessions()) { - const run = session._activeRlmChildRuns.get(childId); - if (run) { - if (run.status !== "running" && run.status !== "queued" && !run.settled) { - if (session._inputScheduler.suspended) session._abandonRlmRunForQuiescence(run); - else run.suppressTerminalNotice = true; - return true; - } - // The abort cascade never reaches running work retained under a settled descendant. - const cancelled = session._cancelRlmChildRun(run, reason); - const descendantsCancelled = run.session?.cancelRunningRlmDescendants(reason) ?? false; - if (cancelled || descendantsCancelled) { - return true; - } - } - // A fruitless match keeps walking: child ids are only mkdir-unique among - // siblings, so a colliding live run elsewhere must stay reachable. - if (session._rlmChildSessions.get(childId)?.session.cancelRunningRlmDescendants(reason)) { - return true; - } - } - return false; - } - - // A done child sits in BOTH maps until passivation; the visited set keeps that dual membership from doubling the walk. - private *_rlmSubtreeSessions(): Generator { - const visited = new Set([this]); - const stack: AgentSession[] = [this]; - while (stack.length > 0) { - const session = stack.pop()!; - yield session; - for (const run of session._activeRlmChildRuns.values()) { - if (run.session && !visited.has(run.session)) { - visited.add(run.session); - stack.push(run.session); - } - } - for (const { session: retained } of session._rlmChildSessions.values()) { - if (!visited.has(retained)) { - visited.add(retained); - stack.push(retained); - } - } - } + return this._children.cancelRlmChildRun(childId, reason); } /** Cancel every running or queued run in this session's subtree. */ cancelRunningRlmDescendants(reason = "Cancelled by user"): boolean { - let cancelled = false; - for (const session of this._rlmSubtreeSessions()) { - for (const run of session._activeRlmChildRuns.values()) { - if (session._cancelRlmChildRun(run, reason)) cancelled = true; - } - } - return cancelled; - } - - private async _assertRlmSubagentSessionNameAvailable(name: string, ignorePendingReservation = false): Promise { - const depth = this._rlmDepth + 1; - if (!ignorePendingReservation && this._pendingRlmSubagentSessionNames.has(name)) { - throw new Error(formatAgentSessionNameUnavailable(name, depth)); - } - const localConflict = - [...this._activeRlmChildRuns.values()].some( - (run) => run.session?.sessionName === name || (!run.session && run.sessionName === name), - ) || - [...this._rlmChildSessions.values()].some(({ session }) => session.sessionName === name) || - [...this._rlmChildCleanupFailures.values()].some((entry) => entry.session_name === name); - if (localConflict) { - throw new Error(formatAgentSessionNameUnavailable(name, depth)); - } - const controller = this._agentMessageController; - if (!controller) return; - const input = { - name, - depth, - parentSessionId: this.sessionId, - parentSessionPath: this.sessionFile, - }; - if (controller.assertSessionNameAvailable) { - await controller.assertSessionNameAvailable(input); - return; - } - const listed = await controller.listAgents(); - const catalog = listed.agents.map( - (agent): AgentFamilyCatalogEntry => ({ - id: agent.sessionId, - ...(agent.sessionName ? { name: agent.sessionName } : {}), - depth: agent.rlmDepth ?? 0, - status: agent.status ?? "idle", - ...(agent.parentSessionId ? { parentSessionId: agent.parentSessionId } : {}), - ...(agent.parentSessionPath ? { parentSessionPath: agent.parentSessionPath } : {}), - ...(agent.sessionPath ? { sessionPath: agent.sessionPath } : {}), - }), - ); - assertAgentSessionNameAvailable(catalog, input); + return this._children.cancelRunningRlmDescendants(reason); } private async _authenticatedRlmModels(): Promise[]> { @@ -7906,538 +6078,8 @@ export class AgentSession { return { model }; } - private async _startRlmChildRun( - prompt: string, - kwargs: Record = {}, - spawnCode?: string, - ): Promise { - // Snapshot before any await: the spawning request is the turn whose tool call is - // executing now. A spawn arriving outside an active run (a detached kernel task - // firing while the parent is idle) has no such turn; an absent edge beats a wrong one. - const spawnedByRequestId = this.isStreaming ? this._semanticEdges.lastTurnRequestId : undefined; - const { name: rawName, model: rawModel, thinking: rawThinking, ...unsupported } = kwargs; - const unsupportedKwargs = Object.keys(unsupported); - if (unsupportedKwargs.length > 0) { - throw new Error(`Unsupported rlm.spawn kwargs: ${unsupportedKwargs.sort().join(", ")}`); - } - const requestedSessionName = normalizeRequestedRlmSubagentSessionName(rawName); - const requestedModel = normalizeRequestedRlmSubagentModel(rawModel); - const requestedThinkingLevel = normalizeRequestedRlmSubagentThinkingLevel(rawThinking); - if (requestedSessionName) assertDirectAgentMessageTarget(requestedSessionName); - if (this._rlmDepth >= this._rlmMaxDepth) { - throw new Error( - `RLM recursion depth limit reached (RLM_DEPTH=${this._rlmDepth}, RLM_MAX_DEPTH=${this._rlmMaxDepth})`, - ); - } - if (requestedSessionName) { - if (this._pendingRlmSubagentSessionNames.has(requestedSessionName)) { - throw new Error(formatAgentSessionNameUnavailable(requestedSessionName, this._rlmDepth + 1)); - } - this._pendingRlmSubagentSessionNames.add(requestedSessionName); - } - let modelSelection: RlmSubagentModelSelection; - try { - if (requestedSessionName) await this._assertRlmSubagentSessionNameAvailable(requestedSessionName, true); - modelSelection = await this._resolveRlmSubagentModel(requestedModel); - } finally { - if (requestedSessionName) this._pendingRlmSubagentSessionNames.delete(requestedSessionName); - } - if (requestedThinkingLevel !== undefined) { - const supported = getSupportedThinkingLevels(modelSelection.model) as ThinkingLevel[]; - if (!supported.includes(requestedThinkingLevel)) { - throw new Error( - `Requested thinking level "${requestedThinkingLevel}" is not supported by model "${modelSelection.model.provider}/${modelSelection.model.id}"; supported levels: ${supported.join(", ")}`, - ); - } - } - if (this._disposed || this._disposing) throw new Error("Cannot spawn a subagent after its parent was disposed"); - - const childSessionDir = this._createChildRlmSessionDir(); - const childNodeId = basename(childSessionDir); - const sessionName = requestedSessionName ?? createDefaultRlmSubagentSessionName(prompt, childNodeId); - if (!requestedSessionName) await this._assertRlmSubagentSessionNameAvailable(sessionName); - const startedAt = Date.now(); - const parentAssistantForUsage = this._findLastAssistantMessage(); - if (parentAssistantForUsage && !this._rlmDurableParentUsage.has(parentAssistantForUsage)) { - this._rlmDurableParentUsage.set(parentAssistantForUsage, cloneUsage(parentAssistantForUsage.usage)); - } - // Child completions accumulate per origin and flush one durable entry per - // settle boundary (agent_end, settlement); the staleness checkpoints and - // timer bound crash loss to one window of accumulated usage. - const pendingChildUsage = new Map(); - let pendingChildUsageSince = 0; - let pendingChildUsageTimer: ReturnType | undefined; - let parentEntryDrainScheduled = false; - const flushPendingChildUsageAttribution = (afterParentDrain = false) => { - if (pendingChildUsageTimer !== undefined) { - clearTimeout(pendingChildUsageTimer); - pendingChildUsageTimer = undefined; - } - if (pendingChildUsage.size === 0 || !parentAssistantForUsage) return; - const parentEntry = this._findAssistantEntryForMessage(parentAssistantForUsage); - if (!parentEntry) { - if (!afterParentDrain && !parentEntryDrainScheduled) { - parentEntryDrainScheduled = true; - const flushAfterParentDrain = () => { - parentEntryDrainScheduled = false; - flushPendingChildUsageAttribution(true); - }; - // A message_end extension may still be holding the parent assistant before its append. - // The parent drain owns this retry; child settlement never waits for that queue. - this._agentEventQueue = this._agentEventQueue.then(flushAfterParentDrain, flushAfterParentDrain); - this._agentEventQueue.catch(() => {}); - } - return; - } - const batches = [...pendingChildUsage.entries()]; - pendingChildUsage.clear(); - for (const [origin, childUsage] of batches) { - const aggregateUsage = cloneUsage(this._rlmDurableParentUsage.get(parentAssistantForUsage)!); - attributeChildUsage(aggregateUsage, childUsage); - const liveUsage = parentAssistantForUsage.usage; - const entryCount = this.sessionManager.getEntries().length; - try { - this.sessionManager.appendChildUsageAttribution(parentEntry.id, childUsage, aggregateUsage, origin); - this._rlmDurableParentUsage.set(parentAssistantForUsage, aggregateUsage); - } catch { - // Attribution is recoverable bookkeeping; a failed append must not break run settlement. - } finally { - // The manager updates this same message; retain siblings' still-pending live usage. - parentAssistantForUsage.usage = liveUsage; - const indexed = this.sessionManager.getEntries()[entryCount]; - const unindexedUsage = this._rlmUnindexedChildUsage.get(parentAssistantForUsage); - // _persist can throw after indexing. That row already participates in live own-usage subtraction. - if ( - indexed?.type === "child_usage_attributed" && - indexed.targetId === parentEntry.id && - unindexedUsage - ) { - subtractAssistantUsage(unindexedUsage, childUsage); - } - this._ownUsageMemo = undefined; - } - } - }; - const flushPendingChildUsageIfStale = () => { - if ( - pendingChildUsage.size > 0 && - Date.now() - pendingChildUsageSince >= RLM_CHILD_USAGE_FLUSH_MAX_PENDING_MS - ) { - flushPendingChildUsageAttribution(); - } - }; - let runningToolCount = 0; - let childSession: AgentSession | undefined; - const run: RlmChildRun = { - id: childNodeId, - prompt, - sessionName, - sessionDir: childSessionDir, - model: modelSelection.model, - status: "queued", - toolUseCount: 0, - settled: false, - abort: noopRlmChildAbort, - publication: createAgentMessageDeferred(), - settlement: createAgentMessageDeferred(), - deletionReservation: createAgentMessageDeferred(), - }; - const throwIfCancelled = () => { - if (run.status === "cancelled") throw new Error(run.error ?? "RLM child cancelled"); - }; - this._activeRlmChildRuns.set(run.id, run); - this._unsettledRlmChildRuns.add(run); - const emitChildUpdate = () => { - const child = this._rlmChildSnapshotForRun(run); - const serialized = JSON.stringify(child); - if (serialized === run.lastEmittedUpdate) return; - run.lastEmittedUpdate = serialized; - this._emit({ type: "rlm_child_update", child }); - }; - run.emitUpdate = emitChildUpdate; - emitChildUpdate(); - - const publishChildSession = (child: AgentSession) => { - childSession = child; - if (this._activeRlmChildRuns.get(run.id) !== run) return; - run.session = child; - run.abort = () => void child.abort(); - run.publication.resolve(); - // Cancellation may have been admitted while runtime construction was - // blocked and run.abort was still a no-op. - if (run.status === "cancelled") run.abort(); - }; - const subagentOptions: CreateRlmSubagentRuntimeOptions = { - ...this._createRlmSubagentRuntimeOptions({ - id: childNodeId, - prompt, - sessionName, - spawnCode, - sessionDir: childSessionDir, - model: modelSelection.model, - thinkingLevel: requestedThinkingLevel, - spawnedByRequestId, - }), - onSessionPublished: publishChildSession, - }; - - const deliverTerminalMessageToParent = async (message: CustomMessage): Promise => { - // Synthesized lifecycle notices always use the parent's private durable - // path. Explicit child replies continue through agent_message separately. - await this._deferRlmTerminalNotice(message); - }; - - run.completeDeletion = () => { - if (!run.deletionNeedsCompletionNotice || run.suppressTerminalNotice || this._disposed || this._disposing) { - return Promise.resolve(); - } - if (run.deletionNotice) return run.deletionNotice; - const notice = deliverTerminalMessageToParent( - createRlmChildTerminalNoticeMessage({ - kind: "cancelled", - childId: run.id, - sessionName, - reason: run.error ?? "Deleted by parent orchestrator", - }), - ); - run.deletionNotice = notice; - return notice; - }; - - run.reportDeletionCleanupFailure = (error) => { - if (run.suppressTerminalNotice || this._disposed || this._disposing) return Promise.resolve(); - if (run.deletionFailureNotice) return run.deletionFailureNotice; - const cleanupError = error instanceof Error ? error.message : String(error); - const notice = deliverTerminalMessageToParent( - createRlmChildFailureMessage({ - childId: run.id, - sessionName, - error: `Deletion cleanup failed; retry rlm.delete_subagent("${run.id}") before completion: ${cleanupError}`, - }), - ); - run.deletionFailureNotice = notice; - return notice; - }; - - // Runtime startup and the task run are deliberately detached. The public - // spawn resolves at admission, while this task owns live tracking, usage, - // retention, cancellation, and late-startup cleanup. - void (async () => { - let childRuntime: RlmSubagentRuntime | undefined; - try { - childRuntime = await this._createRlmSubagentRuntime(subagentOptions); - const child = childRuntime.session; - if (run.status === "cancelled") throw new Error(run.error ?? "RLM child cancelled"); - if (child.sessionName !== sessionName) child.setSessionName(sessionName); - publishChildSession(child); - throwIfCancelled(); - run.status = "running"; - emitChildUpdate(); - const unsubscribeChildEvents = child.subscribe((event) => { - if (event.type === "rlm_child_update") { - this._emit(event); - return; - } - if (event.type === "agent_start") { - run.activity = { kind: "waiting" }; - emitChildUpdate(); - } else if (event.type === "agent_end") { - flushPendingChildUsageAttribution(); - run.activity = undefined; - emitChildUpdate(); - } else if (event.type === "message_end" && event.message.role === "assistant") { - const assistant = event.message as AssistantMessage; - if (assistant.stopReason !== "error" && assistant.stopReason !== "aborted") { - // Flush before the fold: a persisted aggregate may only include - // completions whose childUsage is durable with or before it. - flushPendingChildUsageIfStale(); - attributeChildUsage(parentAssistantForUsage?.usage ?? emptyUsage(), assistant.usage); - if (parentAssistantForUsage) { - const unindexedUsage = - this._rlmUnindexedChildUsage.get(parentAssistantForUsage) ?? emptyUsage(); - addAssistantUsage(unindexedUsage, assistant.usage); - this._rlmUnindexedChildUsage.set(parentAssistantForUsage, unindexedUsage); - this._ownUsageMemo = undefined; - const origin = rlmChildUsageOrigin(child.messages, assistant); - if (pendingChildUsage.size === 0) { - pendingChildUsageSince = Date.now(); - // Wall-clock backstop for long tool runs without checkpoints. - pendingChildUsageTimer = setTimeout( - flushPendingChildUsageAttribution, - RLM_CHILD_USAGE_FLUSH_MAX_PENDING_MS, - ); - pendingChildUsageTimer.unref?.(); - } - const bucket = pendingChildUsage.get(origin) ?? emptyUsage(); - addAssistantUsage(bucket, assistant.usage); - pendingChildUsage.set(origin, bucket); - } - } - const text = compactRlmText(readAssistantText(assistant)); - if (text) run.answerPreview = text; - emitChildUpdate(); - } else if (event.type === "message_start" || event.type === "message_update") { - if (event.message.role === "assistant") { - const text = compactRlmText(readAssistantText(event.message as AssistantMessage)); - if (text) run.answerPreview = text; - run.activity = { kind: "writing" }; - emitChildUpdate(); - } - } else if (event.type === "tool_execution_start") { - flushPendingChildUsageIfStale(); - run.toolUseCount += 1; - runningToolCount += 1; - run.activity = { kind: "executing", toolName: event.toolName }; - emitChildUpdate(); - } else if (event.type === "tool_execution_end") { - runningToolCount = Math.max(0, runningToolCount - 1); - if (runningToolCount === 0) run.activity = { kind: "waiting" }; - emitChildUpdate(); - } else if (event.type === "session_info_changed" || event.type === "recap_update") { - emitChildUpdate(); - } - }); - run.unsubscribe = unsubscribeChildEvents; - const content = `[task from parent]\n\n${prompt}`; - const spawnMessage: AgentSessionMessage = { - role: "custom", - customType: AGENT_MESSAGE_CUSTOM_TYPE, - content, - display: true, - details: { - id: `spawn:${run.id}`, - message: prompt, - from: { - sessionId: this.sessionId, - sessionName: this.sessionName, - activeSessionId: await this._currentActiveSessionId(), - }, - fromRelationship: "parent", - }, - timestamp: Date.now(), - }; - throwIfCancelled(); - const parentReplyCountBeforeRun = child._parentReplyCount; - await child.promptAndWait(content, { - expandPromptTemplates: false, - source: "extension", - customMessage: spawnMessage, - }); - await child.waitForRlmQuiescence(); - if (run.error) throw new Error(run.error); - run.status = "done"; - // Only successful completions return; the edge lands on the parent's next commit. - const childLastCommitted = child.semanticEdges.lastCommittedRequestId; - if (childLastCommitted !== undefined) { - this._semanticEdges.recordChildReturned(child.sessionId, childLastCommitted); - } - run.durationMs = Date.now() - startedAt; - run.activity = undefined; - emitChildUpdate(); - if ( - !run.detachedDeletion && - !run.suppressTerminalNotice && - child._parentReplyCount === parentReplyCountBeforeRun - ) { - const lastAssistantText = child.getLastAssistantText(); - await deliverTerminalMessageToParent( - createRlmChildTerminalNoticeMessage({ - kind: "completed_without_reply", - childId: run.id, - sessionName, - lastAssistantTextPreview: lastAssistantText ? compactRlmText(lastAssistantText) : undefined, - }), - ); - } - if (!this.registerRlmChildSession(run.id, child) && !run.detachedDeletion) { - if (childRuntime && this._subagentRuntimeHost?.releaseRlmSubagentRuntime) { - await this._subagentRuntimeHost - .releaseRlmSubagentRuntime(childRuntime, subagentOptions, "error") - .catch(() => void child.disposeAsync().catch(() => undefined)); - } else { - await child.disposeAsync().catch(() => undefined); - } - } - } catch (error) { - const runError = error instanceof Error ? error : new Error(String(error)); - run.publication.reject(runError); - if (run.status !== "cancelled") { - run.status = "error"; - run.error = runError.message; - } - // A failed child still returns an error outcome the parent consumes; - // cancelled runs and zero-commit children return nothing. - const failedChild = childSession ?? childRuntime?.session; - const failedLastCommitted = failedChild?.semanticEdges.lastCommittedRequestId; - if (run.status === "error" && failedChild && failedLastCommitted !== undefined) { - this._semanticEdges.recordChildReturned(failedChild.sessionId, failedLastCommitted); - } - run.durationMs = Date.now() - startedAt; - run.activity = undefined; - if (run.status === "error" && childSession === undefined) { - // A pre-bind failure leaves no row: "cancelled" is the wire's removal signal. - this._emit({ - type: "rlm_child_update", - child: { ...this._rlmChildSnapshotForRun(run), status: "cancelled" }, - }); - } else { - emitChildUpdate(); - } - if (!run.detachedDeletion && !run.suppressTerminalNotice) { - if (run.status === "error") { - await deliverTerminalMessageToParent( - createRlmChildFailureMessage({ - childId: run.id, - sessionName, - error: run.error ?? "unknown error", - }), - ); - } else if (run.status === "cancelled") { - await deliverTerminalMessageToParent( - createRlmChildTerminalNoticeMessage({ - kind: "cancelled", - childId: run.id, - sessionName, - reason: run.error, - }), - ); - } - } - if (!run.detachedDeletion && childSession && this._subagentRuntimeHost?.releaseRlmSubagentRuntime) { - try { - await this._subagentRuntimeHost.releaseRlmSubagentRuntime( - childRuntime ?? { session: childSession }, - subagentOptions, - run.status === "cancelled" ? "cancelled" : "error", - ); - if (run.status === "cancelled" && !this._disposed && !this._disposing) { - this._deletedRlmChildIds.add(run.id); - this._removeRlmSubagentTracking(run.id); - } - } catch { - await childSession?.disposeAsync().catch(() => undefined); - } - } else if (!run.detachedDeletion) { - try { - if (childRuntime && this._subagentRuntimeHost) { - await this._subagentRuntimeHost.deleteRlmSubagentRuntime(run.id, childRuntime.session); - } else if (childSession) { - await childSession.disposeAsync(); - } - if (run.status === "cancelled" && !this._disposed && !this._disposing) { - this._deletedRlmChildIds.add(run.id); - this._removeRlmSubagentTracking(run.id); - } - } catch { - // A failed best-effort retry remains available through the retained cleanup maps. - } - } - } finally { - flushPendingChildUsageAttribution(); - if (run.detachedDeletion) { - run.deletionRunFinished = true; - if (!run.settled) { - let cleanupSucceeded = !run.deletionCleanupFailed; - if (childRuntime && cleanupSucceeded) { - const cleanup = - run.deletionCleanup ?? this._ensureRlmRunDeletionCleanup(run, childRuntime.session); - cleanupSucceeded = await this._observeRlmRunDeletionCleanup( - run, - run.detachedDeletion, - childRuntime.session, - cleanup, - ); - } - if (cleanupSucceeded) await this._finishRlmRunDeletion(run); - } - } else { - if (this._activeRlmChildRuns.get(run.id) === run) { - if (this._rlmChildSessions.has(run.id)) { - this._activeRlmChildRuns.delete(run.id); - if (run.unsubscribe) this._rlmChildUnsubscribes.set(run.id, run.unsubscribe); - run.abort = noopRlmChildAbort; - run.unsubscribe = undefined; - run.session = undefined; - } else if (run.status !== "error") { - this._removeRlmSubagentTracking(run.id, run); - } else { - run.unsubscribe?.(); - run.abort = noopRlmChildAbort; - run.unsubscribe = undefined; - } - } - run.settled = true; - run.settlement.resolve(); - this._unsettledRlmChildRuns.delete(run); - this._maybeResumeGoalContinuationAfterRlmWork(); - } - } - })().catch(() => undefined); - - return { - rlm_child_id: childNodeId, - name: sessionName, - session_dir: childSessionDir, - model: `${modelSelection.model.provider}/${modelSelection.model.id}`, - }; - } - - async createRlmSession(prompt: string, kwargs: Record = {}): Promise { - const { name: rawName, model: rawModel, thinking: rawThinking, cwd: rawCwd, ...unsupported } = kwargs; - const unsupportedKeys = Object.keys(unsupported); - if (unsupportedKeys.length > 0) { - throw new Error(`Unsupported rlm.create_session kwargs: ${unsupportedKeys.sort().join(", ")}`); - } - if (!prompt.trim()) { - throw new Error("rlm.create_session prompt must not be empty"); - } - if (this._rlmDepth !== 0) { - throw new Error("rlm.create_session is available only from a depth-0 session"); - } - if (this._disposed || this._disposing) { - throw new Error("Cannot create a top-level session after the current session was disposed"); - } - const host = this._subagentRuntimeHost; - if (!host?.createRlmRootSession) { - throw new Error("rlm.create_session requires a daemon-backed depth-0 session"); - } - - const operation = "rlm.create_session"; - const sessionName = normalizeRequestedRlmSubagentSessionName(rawName, operation); - const requestedModel = normalizeRequestedRlmSubagentModel(rawModel, operation); - const requestedThinkingLevel = normalizeRequestedRlmSubagentThinkingLevel(rawThinking, operation); - if (sessionName) { - assertDirectAgentMessageTarget(sessionName); - const controller = this._agentMessageController; - if (controller?.assertSessionNameAvailable) { - await controller.assertSessionNameAvailable({ name: sessionName, depth: 0 }); - } - } - if (rawCwd !== undefined && (typeof rawCwd !== "string" || !rawCwd.trim())) { - throw new Error("rlm.create_session cwd must be a non-empty string"); - } - const cwd = rawCwd === undefined ? this._cwd : resolve(this._cwd, rawCwd.trim()); - const modelSelection = await this._resolveRlmSubagentModel(requestedModel, "top-level session"); - if (requestedThinkingLevel !== undefined) { - const supported = getSupportedThinkingLevels(modelSelection.model) as ThinkingLevel[]; - if (!supported.includes(requestedThinkingLevel)) { - throw new Error( - `Requested thinking level "${requestedThinkingLevel}" is not supported by model "${modelSelection.model.provider}/${modelSelection.model.id}"; supported levels: ${supported.join(", ")}`, - ); - } - } - const thinkingLevel = - requestedThinkingLevel ?? (clampThinkingLevel(modelSelection.model, this.thinkingLevel) as ThinkingLevel); - if (this._disposed || this._disposing) { - throw new Error("Cannot create a top-level session after the current session was disposed"); - } - return host.createRlmRootSession({ - prompt, - sessionName, - cwd, - model: modelSelection.model, - thinkingLevel, - }); + createRlmSession(prompt: string, kwargs: Record = {}): Promise { + return this._children.createRlmSession(prompt, kwargs); } async runRlmChild( @@ -8445,7 +6087,7 @@ export class AgentSession { kwargs: Record = {}, spawnCode?: string, ): Promise { - return this._startRlmChildRun(prompt, kwargs, spawnCode); + return this._children.run(prompt, kwargs, spawnCode); } abortRetry(): void { @@ -8516,39 +6158,11 @@ export class AgentSession { } getRlmMaxDepthStatus(): RlmMaxDepthStatus { - return { maxDepth: this._rlmMaxDepth, source: this._rlmMaxDepthSource }; + return this._childState.getRlmMaxDepthStatus(); } - async setRlmMaxDepth(maxDepth: number, options: { global?: boolean } = {}): Promise { - if (!isNonNegativeInteger(maxDepth)) { - throw new Error("RLM max depth must be a non-negative integer."); - } - - this.sessionManager.appendCustomEntryWithRollback(RLM_MAX_DEPTH_STATE_CUSTOM_TYPE, { maxDepth }); - this._rlmMaxDepth = maxDepth; - this._rlmMaxDepthSource = "chat"; - const oldBase = this._baseSystemPrompt; - this._baseSystemPrompt = this._rebuildSystemPrompt(this.getActiveToolNames()); - this.agent.state.systemPrompt = this._refreshExtensionSystemPrompt(this.agent.state.systemPrompt, oldBase); - - let globalError: string | undefined; - if (options.global) { - await this.settingsManager.flush(); - const staleErrors = this.settingsManager.drainErrors("global"); - for (const { error } of staleErrors) { - console.warn(`Warning: Earlier global settings write failed: ${error.message}`); - } - this.settingsManager.setRlmMaxDepth(maxDepth); - await this.settingsManager.flush(); - const errors = this.settingsManager.drainErrors("global"); - globalError = errors.map(({ error }) => error.message).join("; ") || undefined; - } - - return { - ...this.getRlmMaxDepthStatus(), - globalSaved: options.global === true && globalError === undefined, - ...(globalError ? { globalError } : {}), - }; + setRlmMaxDepth(maxDepth: number, options: { global?: boolean } = {}): Promise { + return this._childState.setRlmMaxDepth(maxDepth, options); } setSessionName(name: string): void { @@ -8957,13 +6571,12 @@ export class AgentSession { } private _subtractUnindexedChildUsage(ownUsage: Usage, entries: SessionEntry[]): void { - for (const entry of entries) { - if (entry.type !== "message" || entry.message.role !== "assistant") continue; - const unindexedUsage = this._rlmUnindexedChildUsage.get(entry.message); - if (unindexedUsage) subtractAssistantUsage(ownUsage, unindexedUsage); - } + this._childUsage.subtractUnindexed(ownUsage, entries); } + private _invalidateOwnUsage(): void { + this._ownUsageMemo = undefined; + } private _ownUsageMemo?: { count: number; tailId: string | undefined; usage: SessionUsageSummary | undefined }; // Whole-file own spend, identical to the catalog scan so rows never shift at passivation. @@ -8995,7 +6608,7 @@ export class AgentSession { const children: ContextTreeNode[] = []; const liveIds = new Set(); - for (const run of this._activeRlmChildRuns.values()) { + for (const run of this._children.getActiveRuns()) { liveIds.add(run.id); const node = run.session?.getContextTree() ?? loadContextTreeChildFromDisk(run.sessionDir, resolveContextWindow); diff --git a/packages/coding-agent/src/session/child-projection.ts b/packages/coding-agent/src/session/child-projection.ts new file mode 100644 index 0000000000..c59c429ebb --- /dev/null +++ b/packages/coding-agent/src/session/child-projection.ts @@ -0,0 +1,158 @@ +import type { AgentSessionMessageAgentSummary, AgentSessionMessageListResult } from "../core/agent-messages.js"; +import type { AgentSession } from "../core/agent-session.js"; +import { createDefaultRlmSubagentSessionName, type RlmListSubagentsResult } from "../core/rlm-runtime.js"; +import { + compactRlmText, + type RetainedRlmChild, + type RlmChildAgentSnapshot, + type RlmChildRun, + readAssistantText, + rlmChildLabel, +} from "./child-types.js"; + +interface ChildVisibility { + isDeleting(id: string): boolean; + isDeleted(id: string): boolean; + hasCleanupFailure(id: string): boolean; +} +export function buildChildList( + activeRuns: Iterable, + retainedChildren: Iterable<[string, RetainedRlmChild]>, + visibility: ChildVisibility, + getSessionDir: (child: AgentSession) => string | undefined, + listedAgents?: AgentSessionMessageListResult, +): RlmListSubagentsResult { + const daemonChildren = new Map(); + const parentActiveSessionId = listedAgents?.current?.activeSessionId; + if (parentActiveSessionId) { + for (const agent of listedAgents.agents) { + if ( + agent.runtimeKind === "subagent" && + agent.parentActiveSessionId === parentActiveSessionId && + agent.rlmChildId + ) { + daemonChildren.set(agent.rlmChildId, agent); + } + } + } + + const subagents: RlmListSubagentsResult["subagents"] = []; + const recorded = new Set(); + for (const run of activeRuns) { + if (visibility.isDeleting(run.id) || run.detachedDeletion || run.status === "cancelled") { + continue; + } + const daemonChild = daemonChildren.get(run.id); + subagents.push({ + rlm_child_id: run.id, + active_session_id: daemonChild?.activeSessionId ?? null, + session_id: daemonChild?.sessionId ?? run.session?.sessionId ?? null, + session_name: daemonChild?.sessionName ?? run.session?.sessionName ?? run.sessionName, + session_dir: run.sessionDir, + status: run.status === "done" ? "completed" : run.status === "error" ? "error" : "running", + }); + recorded.add(run.id); + } + for (const [childId, { session: childSession }] of retainedChildren) { + if (visibility.isDeleting(childId) || recorded.has(childId) || visibility.hasCleanupFailure(childId)) { + continue; + } + const daemonChild = daemonChildren.get(childId); + const sessionDir = getSessionDir(childSession); + if (!sessionDir) { + continue; + } + subagents.push({ + rlm_child_id: childId, + active_session_id: daemonChild?.activeSessionId ?? null, + session_id: daemonChild?.sessionId ?? childSession.sessionId, + session_name: + daemonChild?.sessionName ?? childSession.sessionName ?? createDefaultRlmSubagentSessionName("", childId), + session_dir: sessionDir, + status: "completed", + }); + recorded.add(childId); + } + for (const [childId, daemonChild] of daemonChildren) { + if ( + recorded.has(childId) || + visibility.isDeleting(childId) || + visibility.isDeleted(childId) || + visibility.hasCleanupFailure(childId) || + !daemonChild.sessionDir + ) { + continue; + } + subagents.push({ + rlm_child_id: childId, + active_session_id: daemonChild.activeSessionId, + session_id: daemonChild.sessionId, + session_name: daemonChild.sessionName ?? createDefaultRlmSubagentSessionName("", childId), + session_dir: daemonChild.sessionDir, + status: daemonChild.rlmChildRegistryStatus === "completed" ? "completed" : "error", + }); + } + return { subagents }; +} +export function snapshotChildRun( + run: RlmChildRun, + child: AgentSession | undefined, + parentId: string | undefined, +): RlmChildAgentSnapshot { + const model = child?.model ?? run.model; + return { + id: run.id, + parentId: parentId, + sessionName: child?.sessionName ?? run.sessionName, + model: `${model.provider}/${model.id}`, + label: rlmChildLabel(run.prompt), + status: run.status, + durationMs: run.durationMs, + answerPreview: run.answerPreview, + toolUseCount: run.toolUseCount > 0 ? run.toolUseCount : undefined, + tokenCount: child?._contextTokensForCurrentMessages(), + recap: child?.getCurrentRecap(), + sessionDir: run.sessionDir, + activity: run.activity, + repliedSinceTask: child?.repliedToParentSinceTask, + error: run.error, + }; +} +export function snapshotRetainedChild( + childId: string, + child: AgentSession, + parentId: string | undefined, + sessionDir: string | undefined, +): RlmChildAgentSnapshot { + let answerPreview: string | undefined; + let toolUseCount = 0; + const messages = + child.state.streamingMessage?.role === "assistant" + ? [...child.messages, child.state.streamingMessage] + : child.messages; + for (const message of messages) { + if (message.role !== "assistant") continue; + const text = compactRlmText(readAssistantText(message)); + if (text) answerPreview = text; + toolUseCount += message.content.filter((block) => block.type === "toolCall").length; + } + return { + id: childId, + parentId: parentId, + sessionName: child.sessionName, + model: child.model ? `${child.model.provider}/${child.model.id}` : undefined, + label: child.sessionName ?? "child agent", + status: "done", + answerPreview, + toolUseCount: toolUseCount > 0 ? toolUseCount : undefined, + tokenCount: child._contextTokensForCurrentMessages(), + recap: child.getCurrentRecap(), + sessionDir: sessionDir ?? child.sessionManager.getSessionDir(), + // No run exists (e.g. a child rehydrated after daemon recovery), so live + // session state is the only source for in-flight follow-up work. Mirror + // the run projection's convention: status stays "done" (the recorded task + // finished) and current work surfaces through activity. + activity: child.isSessionActive ? { kind: child.isStreaming ? "writing" : "waiting" } : undefined, + repliedSinceTask: child.repliedToParentSinceTask, + }; +} diff --git a/packages/coding-agent/src/session/child-run.ts b/packages/coding-agent/src/session/child-run.ts new file mode 100644 index 0000000000..8f1c6c3c13 --- /dev/null +++ b/packages/coding-agent/src/session/child-run.ts @@ -0,0 +1,390 @@ +import type { AssistantMessage } from "@earendil-works/pi-ai"; +import { AGENT_MESSAGE_CUSTOM_TYPE, type AgentSessionMessage } from "../core/agent-messages.js"; +import type { AgentSession } from "../core/agent-session.js"; +import { + type CustomMessage, + createRlmChildFailureMessage, + createRlmChildTerminalNoticeMessage, +} from "../core/messages.js"; +import type { + CreateRlmSubagentRuntimeOptions, + RlmSpawnHandle, + RlmSubagentRegistryEntry, + RlmSubagentRuntime, + SubagentRuntimeHost, +} from "../core/rlm-runtime.js"; +import { + compactRlmText, + createChildDeferred, + noopRlmChildAbort, + type RlmChildAgentSnapshot, + type RlmChildRun, + readAssistantText, +} from "./child-types.js"; +import type { ChildRuntimeRequest, SessionChildrenHost } from "./children.js"; + +interface ChildTaskLifecycle { + admitRun(run: RlmChildRun): void; + isCurrentRun(run: RlmChildRun): boolean; + snapshotForRun(run: RlmChildRun): RlmChildAgentSnapshot; + registerSession(id: string, session: AgentSession): boolean; + currentActiveSessionId(): Promise; + getRuntimeHost(): SubagentRuntimeHost | undefined; + recordDeleted(id: string): void; + removeTracking(id: string, run?: RlmChildRun): void; + ensureDeletionCleanup(run: RlmChildRun, session: AgentSession): Promise; + observeDeletionCleanup( + run: RlmChildRun, + subagent: RlmSubagentRegistryEntry, + session: AgentSession, + cleanup: Promise, + ): Promise; + finishDeletion(run: RlmChildRun): Promise; + finishRun(run: RlmChildRun): void; +} +type ChildTaskHost = Pick< + SessionChildrenHost, + | "createUsageTracker" + | "emit" + | "createRuntimeOptions" + | "deliverTerminalNotice" + | "isDisposed" + | "createRuntime" + | "getSessionId" + | "getSessionName" + | "getParentReplyCount" + | "getSemanticEdges" +>; + +/** Owns the admitted task through publication, result delivery and final cleanup. */ +export function launchChildTask( + host: ChildTaskHost, + lifecycle: ChildTaskLifecycle, + request: ChildRuntimeRequest, +): RlmSpawnHandle { + const { + id: childNodeId, + prompt, + sessionName, + spawnCode, + sessionDir: childSessionDir, + thinkingLevel: requestedThinkingLevel, + spawnedByRequestId, + } = request; + const modelSelection = { model: request.model }; + const startedAt = Date.now(); + const usage = host.createUsageTracker(); + let runningToolCount = 0; + let childSession: AgentSession | undefined; + const run: RlmChildRun = { + id: childNodeId, + prompt, + sessionName, + sessionDir: childSessionDir, + model: modelSelection.model, + status: "queued", + toolUseCount: 0, + settled: false, + abort: noopRlmChildAbort, + publication: createChildDeferred(), + settlement: createChildDeferred(), + deletionReservation: createChildDeferred(), + }; + const throwIfCancelled = () => { + if (run.status === "cancelled") throw new Error(run.error ?? "RLM child cancelled"); + }; + lifecycle.admitRun(run); + const emitChildUpdate = () => { + const child = lifecycle.snapshotForRun(run); + const serialized = JSON.stringify(child); + if (serialized === run.lastEmittedUpdate) return; + run.lastEmittedUpdate = serialized; + host.emit({ type: "rlm_child_update", child }); + }; + run.emitUpdate = emitChildUpdate; + emitChildUpdate(); + + const publishChildSession = (child: AgentSession) => { + childSession = child; + if (!lifecycle.isCurrentRun(run)) return; + run.session = child; + run.abort = () => void child.abort(); + run.publication.resolve(); + // Cancellation may have been admitted while runtime construction was + // blocked and run.abort was still a no-op. + if (run.status === "cancelled") run.abort(); + }; + const subagentOptions: CreateRlmSubagentRuntimeOptions = { + ...host.createRuntimeOptions({ + id: childNodeId, + prompt, + sessionName, + spawnCode, + sessionDir: childSessionDir, + model: modelSelection.model, + thinkingLevel: requestedThinkingLevel, + spawnedByRequestId, + }), + onSessionPublished: publishChildSession, + }; + + const deliverTerminalMessageToParent = async (message: CustomMessage): Promise => { + // Synthesized lifecycle notices always use the parent's private durable + // path. Explicit child replies continue through agent_message separately. + await host.deliverTerminalNotice(message); + }; + + run.completeDeletion = () => { + if (!run.deletionNeedsCompletionNotice || run.suppressTerminalNotice || host.isDisposed()) { + return Promise.resolve(); + } + if (run.deletionNotice) return run.deletionNotice; + const notice = deliverTerminalMessageToParent( + createRlmChildTerminalNoticeMessage({ + kind: "cancelled", + childId: run.id, + sessionName, + reason: run.error ?? "Deleted by parent orchestrator", + }), + ); + run.deletionNotice = notice; + return notice; + }; + + run.reportDeletionCleanupFailure = (error) => { + if (run.suppressTerminalNotice || host.isDisposed()) return Promise.resolve(); + if (run.deletionFailureNotice) return run.deletionFailureNotice; + const cleanupError = error instanceof Error ? error.message : String(error); + const notice = deliverTerminalMessageToParent( + createRlmChildFailureMessage({ + childId: run.id, + sessionName, + error: `Deletion cleanup failed; retry rlm.delete_subagent("${run.id}") before completion: ${cleanupError}`, + }), + ); + run.deletionFailureNotice = notice; + return notice; + }; + + // Runtime startup and the task run are deliberately detached. The public + // spawn resolves at admission, while this task owns live tracking, usage, + // retention, cancellation, and late-startup cleanup. + void (async () => { + let childRuntime: RlmSubagentRuntime | undefined; + try { + childRuntime = await host.createRuntime(subagentOptions); + const child = childRuntime.session; + if (run.status === "cancelled") throw new Error(run.error ?? "RLM child cancelled"); + if (child.sessionName !== sessionName) child.setSessionName(sessionName); + publishChildSession(child); + throwIfCancelled(); + run.status = "running"; + emitChildUpdate(); + const unsubscribeChildEvents = child.subscribe((event) => { + if (event.type === "rlm_child_update") { + host.emit(event); + return; + } + if (event.type === "agent_start") { + run.activity = { kind: "waiting" }; + emitChildUpdate(); + } else if (event.type === "agent_end") { + usage.flush(); + run.activity = undefined; + emitChildUpdate(); + } else if (event.type === "message_end" && event.message.role === "assistant") { + const assistant = event.message as AssistantMessage; + if (assistant.stopReason !== "error" && assistant.stopReason !== "aborted") { + usage.record(child.messages, assistant); + } + const text = compactRlmText(readAssistantText(assistant)); + if (text) run.answerPreview = text; + emitChildUpdate(); + } else if (event.type === "message_start" || event.type === "message_update") { + if (event.message.role === "assistant") { + const text = compactRlmText(readAssistantText(event.message as AssistantMessage)); + if (text) run.answerPreview = text; + run.activity = { kind: "writing" }; + emitChildUpdate(); + } + } else if (event.type === "tool_execution_start") { + usage.flushIfStale(); + run.toolUseCount += 1; + runningToolCount += 1; + run.activity = { kind: "executing", toolName: event.toolName }; + emitChildUpdate(); + } else if (event.type === "tool_execution_end") { + runningToolCount = Math.max(0, runningToolCount - 1); + if (runningToolCount === 0) run.activity = { kind: "waiting" }; + emitChildUpdate(); + } else if (event.type === "session_info_changed" || event.type === "recap_update") { + emitChildUpdate(); + } + }); + run.unsubscribe = unsubscribeChildEvents; + const content = `[task from parent]\n\n${prompt}`; + const spawnMessage: AgentSessionMessage = { + role: "custom", + customType: AGENT_MESSAGE_CUSTOM_TYPE, + content, + display: true, + details: { + id: `spawn:${run.id}`, + message: prompt, + from: { + sessionId: host.getSessionId(), + sessionName: host.getSessionName(), + activeSessionId: await lifecycle.currentActiveSessionId(), + }, + fromRelationship: "parent", + }, + timestamp: Date.now(), + }; + throwIfCancelled(); + const parentReplyCountBeforeRun = host.getParentReplyCount(child); + await child.promptAndWait(content, { + expandPromptTemplates: false, + source: "extension", + customMessage: spawnMessage, + }); + await child.waitForRlmQuiescence(); + if (run.error) throw new Error(run.error); + run.status = "done"; + // Only successful completions return; the edge lands on the parent's next commit. + const childLastCommitted = child.semanticEdges.lastCommittedRequestId; + if (childLastCommitted !== undefined) { + host.getSemanticEdges().recordChildReturned(child.sessionId, childLastCommitted); + } + run.durationMs = Date.now() - startedAt; + run.activity = undefined; + emitChildUpdate(); + if ( + !run.detachedDeletion && + !run.suppressTerminalNotice && + host.getParentReplyCount(child) === parentReplyCountBeforeRun + ) { + const lastAssistantText = child.getLastAssistantText(); + await deliverTerminalMessageToParent( + createRlmChildTerminalNoticeMessage({ + kind: "completed_without_reply", + childId: run.id, + sessionName, + lastAssistantTextPreview: lastAssistantText ? compactRlmText(lastAssistantText) : undefined, + }), + ); + } + if (!lifecycle.registerSession(run.id, child) && !run.detachedDeletion) { + if (childRuntime && lifecycle.getRuntimeHost()?.releaseRlmSubagentRuntime) { + await lifecycle.getRuntimeHost()!.releaseRlmSubagentRuntime!( + childRuntime, + subagentOptions, + "error", + ).catch(() => void child.disposeAsync().catch(() => undefined)); + } else { + await child.disposeAsync().catch(() => undefined); + } + } + } catch (error) { + const runError = error instanceof Error ? error : new Error(String(error)); + run.publication.reject(runError); + if (run.status !== "cancelled") { + run.status = "error"; + run.error = runError.message; + } + // A failed child still returns an error outcome the parent consumes; + // cancelled runs and zero-commit children return nothing. + const failedChild = childSession ?? childRuntime?.session; + const failedLastCommitted = failedChild?.semanticEdges.lastCommittedRequestId; + if (run.status === "error" && failedChild && failedLastCommitted !== undefined) { + host.getSemanticEdges().recordChildReturned(failedChild.sessionId, failedLastCommitted); + } + run.durationMs = Date.now() - startedAt; + run.activity = undefined; + if (run.status === "error" && childSession === undefined) { + // A pre-bind failure leaves no row: "cancelled" is the wire's removal signal. + host.emit({ + type: "rlm_child_update", + child: { ...lifecycle.snapshotForRun(run), status: "cancelled" }, + }); + } else { + emitChildUpdate(); + } + if (!run.detachedDeletion && !run.suppressTerminalNotice) { + if (run.status === "error") { + await deliverTerminalMessageToParent( + createRlmChildFailureMessage({ + childId: run.id, + sessionName, + error: run.error ?? "unknown error", + }), + ); + } else if (run.status === "cancelled") { + await deliverTerminalMessageToParent( + createRlmChildTerminalNoticeMessage({ + kind: "cancelled", + childId: run.id, + sessionName, + reason: run.error, + }), + ); + } + } + if (!run.detachedDeletion && childSession && lifecycle.getRuntimeHost()?.releaseRlmSubagentRuntime) { + try { + await lifecycle.getRuntimeHost()!.releaseRlmSubagentRuntime!( + childRuntime ?? { session: childSession }, + subagentOptions, + run.status === "cancelled" ? "cancelled" : "error", + ); + if (run.status === "cancelled" && !host.isDisposed()) { + lifecycle.recordDeleted(run.id); + lifecycle.removeTracking(run.id); + } + } catch { + await childSession?.disposeAsync().catch(() => undefined); + } + } else if (!run.detachedDeletion) { + try { + if (childRuntime && lifecycle.getRuntimeHost()) { + await lifecycle.getRuntimeHost()!.deleteRlmSubagentRuntime(run.id, childRuntime.session); + } else if (childSession) { + await childSession.disposeAsync(); + } + if (run.status === "cancelled" && !host.isDisposed()) { + lifecycle.recordDeleted(run.id); + lifecycle.removeTracking(run.id); + } + } catch { + // A failed best-effort retry remains available through the retained cleanup maps. + } + } + } finally { + usage.flush(); + if (run.detachedDeletion) { + run.deletionRunFinished = true; + if (!run.settled) { + let cleanupSucceeded = !run.deletionCleanupFailed; + if (childRuntime && cleanupSucceeded) { + const cleanup = run.deletionCleanup ?? lifecycle.ensureDeletionCleanup(run, childRuntime.session); + cleanupSucceeded = await lifecycle.observeDeletionCleanup( + run, + run.detachedDeletion, + childRuntime.session, + cleanup, + ); + } + if (cleanupSucceeded) await lifecycle.finishDeletion(run); + } + } else { + lifecycle.finishRun(run); + } + } + })().catch(() => undefined); + + return { + rlm_child_id: childNodeId, + name: sessionName, + session_dir: childSessionDir, + model: `${modelSelection.model.provider}/${modelSelection.model.id}`, + }; +} diff --git a/packages/coding-agent/src/session/child-runtime.ts b/packages/coding-agent/src/session/child-runtime.ts new file mode 100644 index 0000000000..9f825c0b6c --- /dev/null +++ b/packages/coding-agent/src/session/child-runtime.ts @@ -0,0 +1,111 @@ +import { randomUUID } from "node:crypto"; +import { mkdirSync } from "node:fs"; +import { join } from "node:path"; +import { Agent } from "@earendil-works/pi-agent-core"; +import { AgentSession } from "../core/agent-session.js"; +import type { ModelRegistry } from "../core/model-registry.js"; +import type { ResourceLoader } from "../core/resource-loader.js"; +import type { CreateRlmSubagentRuntimeOptions, RlmSubagentRuntime } from "../core/rlm-runtime.js"; +import { SessionManager } from "../core/session-manager.js"; +import type { SettingsManager } from "../core/settings-manager.js"; + +export interface InlineChildRuntimeHost { + cwd: string; + agentDir?: string; + agent: Pick< + Agent, + "convertToLlm" | "transformContext" | "streamFn" | "getApiKey" | "onPayload" | "onResponse" | "toolExecution" + >; + settingsManager: SettingsManager; + resourceLoader: ResourceLoader; + modelRegistry: ModelRegistry; +} +export function createInlineChildRuntime( + host: InlineChildRuntimeHost, + options: CreateRlmSubagentRuntimeOptions, +): RlmSubagentRuntime { + const childSessionManager = SessionManager.create(host.cwd, options.sessionDir); + if (options.parentSession.sessionFile) { + childSessionManager.newSession({ + parentSession: options.parentSession.sessionFile, + rlmDepth: options.rlmDepth, + }); + } + childSessionManager.appendModelChange(options.model.provider, options.model.id); + childSessionManager.appendThinkingLevelChange(options.thinkingLevel); + childSessionManager.appendServiceTierChange(options.serviceTier); + + const childAgent = new Agent({ + initialState: { + systemPrompt: "", + model: options.model, + thinkingLevel: options.thinkingLevel, + serviceTier: options.serviceTier, + tools: [], + }, + convertToLlm: host.agent.convertToLlm, + transformContext: host.agent.transformContext, + streamFn: host.agent.streamFn, + getApiKey: host.agent.getApiKey, + onPayload: host.agent.onPayload, + onResponse: host.agent.onResponse, + steeringMode: host.settingsManager.getSteeringMode(), + followUpMode: host.settingsManager.getFollowUpMode(), + sessionId: childSessionManager.getSessionId(), + thinkingBudgets: host.settingsManager.getThinkingBudgets(), + transport: host.settingsManager.getTransport(), + toolExecution: host.agent.toolExecution, + }); + + const child = new AgentSession({ + agent: childAgent, + sessionManager: childSessionManager, + settingsManager: host.settingsManager, + cwd: host.cwd, + agentDir: host.agentDir, + scopedModels: options.scopedModels, + resourceLoader: host.resourceLoader, + customTools: options.customTools, + modelRegistry: host.modelRegistry, + initialActiveToolNames: options.activeToolNames, + allowedToolNames: options.allowedToolNames, + includeGoals: options.includeGoals, + includeCompactSkill: options.includeCompactSkill, + rlmDepth: options.rlmDepth, + rlmMaxDepth: options.rlmMaxDepth, + rlmSessionDir: options.sessionDir, + rlmParentNodeId: options.rlmParentNodeId, + rlmParentAgent: options.parentSession.sessionName ?? options.parentSession.sessionId, + semanticParentSessionId: options.parentSession.sessionId, + semanticSpawnedByRequestId: options.spawnedByRequestId, + sessionStartEvent: { type: "session_start", reason: "startup" }, + }); + if (child.sessionName !== options.sessionName) { + try { + child.setSessionName(options.sessionName); + } catch (error) { + child.dispose(); + throw error; + } + } + options.onSessionPublished?.(child); + + return { session: child }; +} + +export function createChildSessionDir(getParentDir: () => string): string { + const parentDir = getParentDir(); + for (let i = 0; i < 100; i++) { + const childDir = join(parentDir, `sub-${randomUUID().slice(0, 8)}`); + try { + mkdirSync(childDir); + return childDir; + } catch (error) { + if (error instanceof Error && "code" in error && error.code === "EEXIST") { + continue; + } + throw error; + } + } + throw new Error("Unable to create unique RLM child session directory"); +} diff --git a/packages/coding-agent/src/session/child-state.ts b/packages/coding-agent/src/session/child-state.ts new file mode 100644 index 0000000000..19b6caf20f --- /dev/null +++ b/packages/coding-agent/src/session/child-state.ts @@ -0,0 +1,170 @@ +import type { RlmMaxDepthSource, RlmMaxDepthStatus, SetRlmMaxDepthResult } from "../core/rlm-max-depth.js"; +import type { SessionManager } from "../core/session-manager.js"; +import type { SettingsManager } from "../core/settings-manager.js"; + +interface PersistedRlmMaxDepthState { + maxDepth: number; +} +const RLM_MAX_DEPTH_STATE_CUSTOM_TYPE = "rlm_max_depth_state"; +function isNonNegativeInteger(value: unknown): value is number { + return typeof value === "number" && Number.isSafeInteger(value) && value >= 0; +} + +function parseDepth(value: string | undefined, fallback: number, name: string): number { + if (value === undefined || value === "") { + return fallback; + } + if (!/^\d+$/.test(value)) { + throw new Error(`${name} must be a non-negative integer`); + } + const parsed = Number(value); + if (!isNonNegativeInteger(parsed)) { + throw new Error(`${name} must be a non-negative integer`); + } + return parsed; +} + +function isPersistedRlmMaxDepthState(value: unknown): value is PersistedRlmMaxDepthState { + return ( + typeof value === "object" && value !== null && isNonNegativeInteger((value as PersistedRlmMaxDepthState).maxDepth) + ); +} + +export interface SessionChildStateHost { + sessionManager: Pick; + settingsManager: Pick; + getRlmMaxDepthStatus(): RlmMaxDepthStatus; + refreshPrompt(preserveExtensionPrompt: boolean): void; + emitRecap(recap: string | undefined): void; +} +export class SessionChildState { + readonly depth: number; + private readonly configuredMaxDepth: number | undefined; + private _maxDepth: number; + private _maxDepthSource: RlmMaxDepthSource; + private _repliedSinceTask: boolean | undefined; + private _replyCount = 0; + private recap: string | undefined; + constructor( + private readonly host: SessionChildStateHost, + config: { rlmDepth?: number; rlmMaxDepth?: number }, + ) { + const headerRlmDepth = this.host.sessionManager.getHeader()?.rlmDepth; + this.depth = + config.rlmDepth ?? + (isNonNegativeInteger(headerRlmDepth) ? headerRlmDepth : parseDepth(process.env.RLM_DEPTH, 0, "RLM_DEPTH")); + this.configuredMaxDepth = config.rlmMaxDepth; + if (this.configuredMaxDepth !== undefined && !isNonNegativeInteger(this.configuredMaxDepth)) { + throw new Error("rlmMaxDepth must be a non-negative integer"); + } + const resolvedRlmMaxDepth = this._resolveRlmMaxDepth(); + this._maxDepth = resolvedRlmMaxDepth.maxDepth; + this._maxDepthSource = resolvedRlmMaxDepth.source; + } + get maxDepth(): number { + return this._maxDepth; + } + get repliedSinceTask(): boolean | undefined { + return this._repliedSinceTask; + } + get replyCount(): number { + return this._replyCount; + } + initializeParentReply(): void { + // Resumed transcripts do not prove whether the child already replied. + this._repliedSinceTask = + this.depth > 0 && this.host.sessionManager.getBranch().some((entry) => entry.type === "message") + ? undefined + : false; + } + recordReply(): void { + this._repliedSinceTask = true; + this._replyCount += 1; + } + resetReply(): void { + this._repliedSinceTask = false; + } + getCurrentRecap(): string | undefined { + return this.recap; + } + setCurrentRecap(recap: string | undefined): void { + if (this.recap === recap) return; + this.recap = recap; + this.host.emitRecap(recap); + } + reloadFromBranch(): void { + const previousMaxDepth = this._maxDepth; + const resolved = this._resolveRlmMaxDepth(); + this._maxDepth = resolved.maxDepth; + this._maxDepthSource = resolved.source; + if (resolved.maxDepth !== previousMaxDepth) this.host.refreshPrompt(false); + } + private _loadPersistedRlmMaxDepthState(): PersistedRlmMaxDepthState | undefined { + const branch = this.host.sessionManager.getBranch(); + for (let i = branch.length - 1; i >= 0; i--) { + const entry = branch[i]; + if ( + entry.type === "custom" && + entry.customType === RLM_MAX_DEPTH_STATE_CUSTOM_TYPE && + isPersistedRlmMaxDepthState(entry.data) + ) { + return entry.data; + } + } + return undefined; + } + private _resolveRlmMaxDepth(): { + maxDepth: number; + source: RlmMaxDepthSource; + } { + const persisted = this._loadPersistedRlmMaxDepthState(); + if (persisted) { + return { maxDepth: persisted.maxDepth, source: "chat" }; + } + if (this.configuredMaxDepth !== undefined) { + return { maxDepth: this.configuredMaxDepth, source: "inherited" }; + } + const global = this.host.settingsManager.getRlmMaxDepth(); + if (global !== undefined && isNonNegativeInteger(global)) { + return { maxDepth: global, source: "global" }; + } + const env = process.env.RLM_MAX_DEPTH; + if (env !== undefined && env !== "") { + return { maxDepth: parseDepth(env, 1, "RLM_MAX_DEPTH"), source: "env" }; + } + return { maxDepth: 2, source: "default" }; + } + + getRlmMaxDepthStatus(): RlmMaxDepthStatus { + return { maxDepth: this._maxDepth, source: this._maxDepthSource }; + } + async setRlmMaxDepth(maxDepth: number, options: { global?: boolean } = {}): Promise { + if (!isNonNegativeInteger(maxDepth)) { + throw new Error("RLM max depth must be a non-negative integer."); + } + + this.host.sessionManager.appendCustomEntryWithRollback(RLM_MAX_DEPTH_STATE_CUSTOM_TYPE, { maxDepth }); + this._maxDepth = maxDepth; + this._maxDepthSource = "chat"; + this.host.refreshPrompt(true); + + let globalError: string | undefined; + if (options.global) { + await this.host.settingsManager.flush(); + const staleErrors = this.host.settingsManager.drainErrors("global"); + for (const { error } of staleErrors) { + console.warn(`Warning: Earlier global settings write failed: ${error.message}`); + } + this.host.settingsManager.setRlmMaxDepth(maxDepth); + await this.host.settingsManager.flush(); + const errors = this.host.settingsManager.drainErrors("global"); + globalError = errors.map(({ error }) => error.message).join("; ") || undefined; + } + + return { + ...this.host.getRlmMaxDepthStatus(), + globalSaved: options.global === true && globalError === undefined, + ...(globalError ? { globalError } : {}), + }; + } +} diff --git a/packages/coding-agent/src/session/child-types.ts b/packages/coding-agent/src/session/child-types.ts new file mode 100644 index 0000000000..69e61e2b39 --- /dev/null +++ b/packages/coding-agent/src/session/child-types.ts @@ -0,0 +1,117 @@ +import type { Api, AssistantMessage, Model } from "@earendil-works/pi-ai"; +import type { AgentSession } from "../core/agent-session.js"; +import type { RlmSubagentRegistryEntry } from "../core/rlm-runtime.js"; + +export type RlmChildAgentStatus = "queued" | "running" | "done" | "error" | "cancelled"; + +export interface RlmChildAgentActivity { + kind: "waiting" | "writing" | "executing"; + toolName?: string; +} + +export interface RlmChildAgentSnapshot { + id: string; + parentId?: string; + activeSessionId?: string; + sessionName?: string; + model?: string; + label: string; + status: RlmChildAgentStatus; + durationMs?: number; + answerPreview?: string; + toolUseCount?: number; + tokenCount?: number; + recap?: string; + sessionDir: string; + activity?: RlmChildAgentActivity; + repliedSinceTask?: boolean; + error?: string; +} + +export interface RlmChildRun { + id: string; + prompt: string; + sessionName: string; + sessionDir: string; + model: Model; + status: RlmChildAgentStatus; + durationMs?: number; + answerPreview?: string; + toolUseCount: number; + activity?: RlmChildAgentActivity; + error?: string; + abort: () => void; + publication: ChildDeferred; + /** Resolves after terminal result publication and detached-run cleanup finish. */ + settlement: ChildDeferred; + /** Child session, once its runtime exists. Used to cancel nested child runs. */ + session?: AgentSession; + settled: boolean; + /** Do not inject a late terminal notice after the parent session is aborted. */ + suppressTerminalNotice?: boolean; + /** Excluded from future strong barriers after an authoritative cancellation cut. */ + abandonedForQuiescence?: boolean; + /** Selector snapshot for an admitted explicit delete. */ + detachedDeletion?: RlmSubagentRegistryEntry; + /** Shared physical runtime cleanup owned by the explicit-delete path. */ + deletionCleanup?: Promise; + deletionCleanupObserver?: Promise; + /** Resolves when a deletion may release its selector reservation. */ + deletionReservation: ChildDeferred; + deletionCleanupFailed?: boolean; + deletionRunFinished?: boolean; + deletionNotice?: Promise; + deletionFailureNotice?: Promise; + deletionNeedsCompletionNotice?: boolean; + completeDeletion?: () => Promise; + reportDeletionCleanupFailure?: (error: unknown) => Promise; + emitUpdate?: () => void; + lastEmittedUpdate?: string; + unsubscribe?: () => void; +} + +export interface RetainedRlmChild { + session: AgentSession; + run?: RlmChildRun; +} + +export interface ChildDeferred { + promise: Promise; + resolve: () => void; + reject: (error: Error) => void; +} + +export function createChildDeferred(): ChildDeferred { + const deferred = {} as ChildDeferred; + deferred.promise = new Promise((resolve, reject) => { + deferred.resolve = resolve; + deferred.reject = reject; + }); + deferred.promise.catch(() => undefined); + return deferred; +} + +export function compactRlmText(text: string, maxLength = 160): string { + const compact = text.replace(/\s+/g, " ").trim(); + if (compact.length <= maxLength) { + return compact; + } + return `${compact.slice(0, Math.max(0, maxLength - 3)).trimEnd()}...`; +} + +// Child-agent label: collapse to one line but keep the full prompt — the TUI +// truncates to the visible width and elides shared prefixes, so capping here +// would only hide the divergence between near-identical sibling prompts. +export function rlmChildLabel(prompt: string): string { + return prompt.replace(/\s+/g, " ").trim() || "child agent"; +} + +export function readAssistantText(message: AssistantMessage): string { + return message.content + .filter((block) => block.type === "text") + .map((block) => block.text) + .join(""); +} + +export function noopRlmChildAbort(): void {} +export function noopRlmChildEventUnsubscribe(): void {} diff --git a/packages/coding-agent/src/session/child-usage.ts b/packages/coding-agent/src/session/child-usage.ts new file mode 100644 index 0000000000..3bf7f29c13 --- /dev/null +++ b/packages/coding-agent/src/session/child-usage.ts @@ -0,0 +1,169 @@ +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import type { AssistantMessage, Usage } from "@earendil-works/pi-ai"; +import { isAgentSessionMessage } from "../core/agent-messages.js"; +import type { + ChildUsageAttributionEntry, + SessionEntry, + SessionManager, + SessionMessageEntry, +} from "../core/session-manager.js"; +import { addAssistantUsage, cloneUsage, emptyUsage, subtractAssistantUsage } from "../core/usage.js"; + +export interface ChildUsageHost { + sessionManager: Pick; + afterParentDrain(flush: () => void): void; + invalidateOwnUsage(): void; +} +export interface ChildUsageTracker { + record(messages: readonly AgentMessage[], assistant: AssistantMessage): void; + flush(): void; + flushIfStale(): void; +} + +// Bounds how much accumulated child usage a parent process crash can lose. +const RLM_CHILD_USAGE_FLUSH_MAX_PENDING_MS = 60_000; + +/** Label a child completion's usage by the nearest preceding prompt that triggered it. */ +function rlmChildUsageOrigin( + messages: readonly AgentMessage[], + assistant: AssistantMessage, +): ChildUsageAttributionEntry["origin"] { + for (let index = messages.lastIndexOf(assistant) - 1; index >= 0; index--) { + const message = messages[index]; + if (message.role !== "user" && message.role !== "custom") continue; + return message.role === "custom" && isAgentSessionMessage(message) + ? message.details.id.startsWith("spawn:") + ? "spawn_task" + : "agent_message" + : "direct_user"; + } + return "direct_user"; +} + +function attributeChildUsage(parentUsage: Usage, childUsage: Usage): void { + const parentContextTokens = + parentUsage.totalTokens || + parentUsage.input + parentUsage.output + parentUsage.cacheRead + parentUsage.cacheWrite; + // Recursive children are launched from an assistant tool call, so the parent assistant + // message carries their billable usage for session-level cost totals. + addAssistantUsage(parentUsage, childUsage); + // Child work affects session-level billable totals, not the parent's model-facing context size. + parentUsage.totalTokens = parentContextTokens; +} + +export class SessionChildUsage { + private _rlmDurableParentUsage = new WeakMap(); + private _rlmUnindexedChildUsage = new WeakMap(); + constructor(private readonly host: ChildUsageHost) {} + createTracker(parent: AssistantMessage | undefined): ChildUsageTracker { + const parentAssistantForUsage = parent; + if (parentAssistantForUsage && !this._rlmDurableParentUsage.has(parentAssistantForUsage)) { + this._rlmDurableParentUsage.set(parentAssistantForUsage, cloneUsage(parentAssistantForUsage.usage)); + } + // Child completions accumulate per origin and flush one durable entry per + // settle boundary (agent_end, settlement); the staleness checkpoints and + // timer bound crash loss to one window of accumulated usage. + const pendingChildUsage = new Map(); + let pendingChildUsageSince = 0; + let pendingChildUsageTimer: ReturnType | undefined; + let parentEntryDrainScheduled = false; + const flushPendingChildUsageAttribution = (afterParentDrain = false) => { + if (pendingChildUsageTimer !== undefined) { + clearTimeout(pendingChildUsageTimer); + pendingChildUsageTimer = undefined; + } + if (pendingChildUsage.size === 0 || !parentAssistantForUsage) return; + const parentEntry = this.host.sessionManager + .getEntries() + .find( + (entry): entry is SessionMessageEntry => + entry.type === "message" && entry.message === parentAssistantForUsage, + ); + if (!parentEntry) { + if (!afterParentDrain && !parentEntryDrainScheduled) { + parentEntryDrainScheduled = true; + const flushAfterParentDrain = () => { + parentEntryDrainScheduled = false; + flushPendingChildUsageAttribution(true); + }; + // A message_end extension may still be holding the parent assistant before its append. + // The parent drain owns this retry; child settlement never waits for that queue. + this.host.afterParentDrain(flushAfterParentDrain); + } + return; + } + const batches = [...pendingChildUsage.entries()]; + pendingChildUsage.clear(); + for (const [origin, childUsage] of batches) { + const aggregateUsage = cloneUsage(this._rlmDurableParentUsage.get(parentAssistantForUsage)!); + attributeChildUsage(aggregateUsage, childUsage); + const liveUsage = parentAssistantForUsage.usage; + const entryCount = this.host.sessionManager.getEntries().length; + try { + this.host.sessionManager.appendChildUsageAttribution(parentEntry.id, childUsage, aggregateUsage, origin); + this._rlmDurableParentUsage.set(parentAssistantForUsage, aggregateUsage); + } catch { + // Attribution is recoverable bookkeeping; a failed append must not break run settlement. + } finally { + // The manager updates this same message; retain siblings' still-pending live usage. + parentAssistantForUsage.usage = liveUsage; + const indexed = this.host.sessionManager.getEntries()[entryCount]; + const unindexedUsage = this._rlmUnindexedChildUsage.get(parentAssistantForUsage); + // _persist can throw after indexing. That row already participates in live own-usage subtraction. + if ( + indexed?.type === "child_usage_attributed" && + indexed.targetId === parentEntry.id && + unindexedUsage + ) { + subtractAssistantUsage(unindexedUsage, childUsage); + } + this.host.invalidateOwnUsage(); + } + } + }; + const flushPendingChildUsageIfStale = () => { + if ( + pendingChildUsage.size > 0 && + Date.now() - pendingChildUsageSince >= RLM_CHILD_USAGE_FLUSH_MAX_PENDING_MS + ) { + flushPendingChildUsageAttribution(); + } + }; + return { + flush: flushPendingChildUsageAttribution, + flushIfStale: flushPendingChildUsageIfStale, + record: (messages, assistant) => { + // Flush before the fold: a persisted aggregate may only include + // completions whose childUsage is durable with or before it. + flushPendingChildUsageIfStale(); + attributeChildUsage(parentAssistantForUsage?.usage ?? emptyUsage(), assistant.usage); + if (parentAssistantForUsage) { + const unindexedUsage = this._rlmUnindexedChildUsage.get(parentAssistantForUsage) ?? emptyUsage(); + addAssistantUsage(unindexedUsage, assistant.usage); + this._rlmUnindexedChildUsage.set(parentAssistantForUsage, unindexedUsage); + this.host.invalidateOwnUsage(); + const origin = rlmChildUsageOrigin(messages, assistant); + if (pendingChildUsage.size === 0) { + pendingChildUsageSince = Date.now(); + // Wall-clock backstop for long tool runs without checkpoints. + pendingChildUsageTimer = setTimeout( + flushPendingChildUsageAttribution, + RLM_CHILD_USAGE_FLUSH_MAX_PENDING_MS, + ); + pendingChildUsageTimer.unref?.(); + } + const bucket = pendingChildUsage.get(origin) ?? emptyUsage(); + addAssistantUsage(bucket, assistant.usage); + pendingChildUsage.set(origin, bucket); + } + }, + }; + } + subtractUnindexed(ownUsage: Usage, entries: SessionEntry[]): void { + for (const entry of entries) { + if (entry.type !== "message" || entry.message.role !== "assistant") continue; + const unindexedUsage = this._rlmUnindexedChildUsage.get(entry.message); + if (unindexedUsage) subtractAssistantUsage(ownUsage, unindexedUsage); + } + } +} diff --git a/packages/coding-agent/src/session/children.ts b/packages/coding-agent/src/session/children.ts new file mode 100644 index 0000000000..1162810e0c --- /dev/null +++ b/packages/coding-agent/src/session/children.ts @@ -0,0 +1,1022 @@ +import { basename, resolve } from "node:path"; +import type { ThinkingLevel } from "@earendil-works/pi-agent-core"; +import { type Api, clampThinkingLevel, getSupportedThinkingLevels, type Model } from "@earendil-works/pi-ai"; +import { + type AgentFamilyCatalogEntry, + type AgentSessionMessageController, + type AgentSessionMessageListResult, + assertAgentSessionNameAvailable, + assertDirectAgentMessageTarget, + formatAgentSessionNameUnavailable, +} from "../core/agent-messages.js"; +import type { AgentSession, AgentSessionEvent } from "../core/agent-session.js"; +import type { CustomMessage } from "../core/messages.js"; +import { + type CreateRlmSubagentRuntimeOptions, + createDefaultRlmSubagentSessionName, + normalizeRequestedRlmSubagentModel, + normalizeRequestedRlmSubagentSessionName, + normalizeRequestedRlmSubagentThinkingLevel, + type RlmCreateSessionResult, + type RlmDeleteSubagentResult, + type RlmListSubagentsResult, + type RlmSpawnHandle, + type RlmSubagentRegistryEntry, + type RlmSubagentRuntime, + type SubagentRuntimeHost, +} from "../core/rlm-runtime.js"; +import type { SemanticEdgeRecorder } from "../core/semantic-edges.js"; +import { buildChildList, snapshotChildRun, snapshotRetainedChild } from "./child-projection.js"; +import { launchChildTask } from "./child-run.js"; +import { + createChildDeferred, + noopRlmChildAbort, + noopRlmChildEventUnsubscribe, + type RetainedRlmChild, + type RlmChildAgentSnapshot, + type RlmChildAgentStatus, + type RlmChildRun, +} from "./child-types.js"; +import type { ChildUsageTracker } from "./child-usage.js"; + +export interface ChildRuntimeRequest { + id: string; + prompt: string; + sessionName: string; + spawnCode?: string; + sessionDir: string; + model: Model; + thinkingLevel?: ThinkingLevel; + spawnedByRequestId?: string; +} +export interface SessionChildrenHost { + isDisposed(): boolean; + isInputSuspended(): boolean; + isStreaming(): boolean; + isSessionActive(): boolean; + getMessageController(): AgentSessionMessageController | undefined; + getDepth(): number; + getMaxDepth(): number; + getParentNodeId(): string | undefined; + getCwd(): string; + getSessionId(): string; + getSessionName(): string | undefined; + getSessionFile(): string | undefined; + getThinkingLevel(): ThinkingLevel; + getSemanticEdges(): SemanticEdgeRecorder; + getChildOwner(child: AgentSession): SessionChildren; + getParentReplyCount(child: AgentSession): number; + getChildSessionDir(child: AgentSession): string | undefined; + listRlmSubagents(): Promise; + deleteRlmSubagent(target: string): Promise; + registerRlmChildSession(childId: string, session: AgentSession): boolean; + resolveModel(reference: string | undefined, target?: string): Promise<{ model: Model }>; + createSessionDir(): string; + createRuntimeOptions(request: ChildRuntimeRequest): CreateRlmSubagentRuntimeOptions; + createRuntime(options: CreateRlmSubagentRuntimeOptions): Promise; + createUsageTracker(): ChildUsageTracker; + hasDeferredTerminalNotices(): boolean; + waitForHeadlessIdle(): Promise; + waitForActivityChange(signal: AbortSignal): Promise; + deliverTerminalNotice(message: CustomMessage): Promise; + emit(event: AgentSessionEvent): void; + onSettled(): void; +} +type RlmSubagentModelSelection = { model: Model }; + +export class SessionChildren { + constructor(private readonly host: SessionChildrenHost) {} + private _subagentRuntimeHost?: SubagentRuntimeHost; + private _activeRlmChildRuns = new Map(); + private _unsettledRlmChildRuns = new Set(); + private _abandonedRlmQuiescenceChildIds = new Set(); + private _rlmQuiescenceWaitAborts = new Set(); + private _pendingRlmSubagentSessionNames = new Set(); + // Inline mode keeps finished child sessions so the inspector can still read them; + // the daemon does the same by leaving the child session resident in its registry. + private _rlmChildSessions = new Map(); + private _deletedRlmChildIds = new Set(); + // Failed explicit deletes stay hidden from listings but retain their original + // selector so a later delete can retry cleanup without orphaning the runtime. + private _rlmChildCleanupFailures = new Map(); + private _deletingRlmChildren = new Map< + string, + { + subagent: RlmSubagentRegistryEntry; + promise: Promise; + } + >(); + // Kept alive for retained children so nested updates (e.g. a grandchild cancel) + // still forward to root; torn down when the retained child is disposed. + private _rlmChildUnsubscribes = new Map void>(); + private _abandonRlmRunForQuiescence(run: RlmChildRun): void { + run.suppressTerminalNotice = true; + run.abandonedForQuiescence = true; + this._abandonedRlmQuiescenceChildIds.add(run.id); + this._unsettledRlmChildRuns.delete(run); + run.settlement.resolve(); + this.host.onSettled(); + } + + cancelActiveRuns(reason: string): void { + for (const run of this._activeRlmChildRuns.values()) { + this._cancelRlmChildRun(run, reason); + } + } + + private _cancelRlmChildRun(run: RlmChildRun, reason: string): boolean { + if (run.status !== "running" && run.status !== "queued") { + return false; + } + run.status = "cancelled"; + if (this.host.isInputSuspended()) this._abandonRlmRunForQuiescence(run); + run.error = reason; + run.publication.reject(new Error(reason)); + run.abort(); + // Surface the cancellation immediately; the run's own terminal update is + // delayed indefinitely when the child is stuck mid-stream, which is + // exactly when users reach for the kill. + run.emitUpdate?.(); + return true; + } + + getRlmChildRunStatus(childId: string): RlmChildAgentStatus | undefined { + return this._activeRlmChildRuns.get(childId)?.status; + } + + private async _currentActiveSessionId(): Promise { + try { + return (await this.host.getMessageController()?.listAgents())?.current?.activeSessionId; + } catch { + return undefined; + } + } + + async awaitPublication(selector: string): Promise { + const run = [...this._activeRlmChildRuns.values()].find( + (candidate) => + (candidate.status === "queued" || candidate.status === "running" || candidate.status === "done") && + !candidate.detachedDeletion && + (candidate.id === selector || candidate.sessionName === selector), + ); + if (!run) return undefined; + await run.publication.promise; + return run.session?.sessionId; + } + + async listRlmSubagents(): Promise { + return this._buildRlmSubagentList(await this.host.getMessageController()?.listAgents()); + } + + private _buildRlmSubagentList(listedAgents?: AgentSessionMessageListResult): RlmListSubagentsResult { + return buildChildList( + this._activeRlmChildRuns.values(), + this._rlmChildSessions, + { + isDeleting: (id) => this._deletingRlmChildren.has(id), + isDeleted: (id) => this._deletedRlmChildIds.has(id), + hasCleanupFailure: (id) => this._rlmChildCleanupFailures.has(id), + }, + (child) => this.host.getChildSessionDir(child), + listedAgents, + ); + } + + private _rlmSubagentMatchesTarget(entry: RlmSubagentRegistryEntry, target: string): boolean { + return ( + entry.rlm_child_id === target || + entry.active_session_id === target || + entry.session_id === target || + entry.session_name === target + ); + } + + private async _resolveDirectRlmSubagent(target: string): Promise { + const candidates = [...(await this.host.listRlmSubagents()).subagents, ...this._rlmChildCleanupFailures.values()]; + const matches = candidates.filter((entry) => this._rlmSubagentMatchesTarget(entry, target)); + if (matches.length === 0) { + throw new Error(`No direct RLM subagent matches "${target}" in the current parent session`); + } + if (matches.length > 1) { + throw new Error(`RLM subagent selector "${target}" is ambiguous in the current parent session`); + } + return matches[0]!; + } + + async deleteInactiveRlmSubagent( + childId: string, + isExternallyRunning: () => boolean = () => false, + ): Promise<"deleted" | "not_found" | "running"> { + for (const owner of this._rlmSubtreeSessions()) { + const isRunning = (): boolean => { + const status = owner._activeRlmChildRuns.get(childId)?.status; + return status === "queued" || status === "running" || isExternallyRunning(); + }; + if (isRunning()) { + return "running"; + } + const subagent = [ + ...(await owner.host.listRlmSubagents()).subagents, + ...owner._rlmChildCleanupFailures.values(), + ].find((entry) => entry.rlm_child_id === childId); + if (!subagent) continue; + if (isRunning()) { + return "running"; + } + const result = await owner._trackRlmSubagentDeletion(subagent, () => { + if (isRunning()) { + return Promise.resolve({ subagent, outcome: "skipped_running" }); + } + return owner._deleteResolvedRlmSubagent(subagent); + }); + return result.outcome === "skipped_running" ? "running" : "deleted"; + } + return "not_found"; + } + + async deleteRlmSubagent(target: string): Promise { + const inFlight = [...this._deletingRlmChildren.values()].filter(({ subagent }) => + this._rlmSubagentMatchesTarget(subagent, target), + ); + if (inFlight.length > 1) { + throw new Error(`RLM subagent selector "${target}" is ambiguous in the current parent session`); + } + + // Running and retained children can be reserved synchronously. This keeps + // them hidden immediately while the async daemon listing checks for a + // conflicting passive selector. + const localMatches = [ + ...this._buildRlmSubagentList().subagents, + ...this._rlmChildCleanupFailures.values(), + ].filter((entry) => this._rlmSubagentMatchesTarget(entry, target)); + const matchingChildIds = new Set([ + ...inFlight.map(({ subagent }) => subagent.rlm_child_id), + ...localMatches.map((subagent) => subagent.rlm_child_id), + ]); + if (matchingChildIds.size > 1 || localMatches.length > 1) { + throw new Error(`RLM subagent selector "${target}" is ambiguous in the current parent session`); + } + if (inFlight[0]) { + return inFlight[0].promise; + } + if (localMatches[0]) { + const subagent = localMatches[0]; + return this._trackRlmSubagentDeletion(subagent, async () => { + const listedAgents = await this.host.getMessageController()?.listAgents(); + const listedSubagents = this._buildRlmSubagentList(listedAgents).subagents; + const passiveMatches = listedSubagents.filter( + (entry) => entry.rlm_child_id !== subagent.rlm_child_id && this._rlmSubagentMatchesTarget(entry, target), + ); + if (passiveMatches.length > 0) { + throw new Error(`RLM subagent selector "${target}" is ambiguous in the current parent session`); + } + const parentActiveSessionId = listedAgents?.current?.activeSessionId; + const daemonChild = listedAgents?.agents.find( + (agent) => + agent.rlmChildId === subagent.rlm_child_id && agent.parentActiveSessionId === parentActiveSessionId, + ); + const resolvedSubagent = daemonChild + ? { + ...subagent, + active_session_id: daemonChild.activeSessionId, + session_id: daemonChild.sessionId, + session_name: daemonChild.sessionName ?? subagent.session_name, + } + : subagent; + return this._deleteResolvedRlmSubagent(resolvedSubagent); + }); + } + + const directMatches = [ + ...(await this.host.listRlmSubagents()).subagents, + ...this._rlmChildCleanupFailures.values(), + ].filter((entry) => this._rlmSubagentMatchesTarget(entry, target)); + const directChildIds = new Set(directMatches.map((subagent) => subagent.rlm_child_id)); + if (directChildIds.size > 1) { + throw new Error(`RLM subagent selector "${target}" is ambiguous in the current parent session`); + } + const subagent = directMatches[0] ?? (await this._resolveDirectRlmSubagent(target)); + return this._trackRlmSubagentDeletion(subagent, () => this._deleteResolvedRlmSubagent(subagent)); + } + + private async _trackRlmSubagentDeletion( + subagent: RlmSubagentRegistryEntry, + startDeletion: () => Promise, + ): Promise { + const existing = this._deletingRlmChildren.get(subagent.rlm_child_id); + if (existing) return existing.promise; + const deletion = Promise.resolve().then(startDeletion); + this._deletingRlmChildren.set(subagent.rlm_child_id, { + subagent, + promise: deletion, + }); + try { + return await deletion; + } finally { + const clearReservation = () => { + if (this._deletingRlmChildren.get(subagent.rlm_child_id)?.promise === deletion) { + this._deletingRlmChildren.delete(subagent.rlm_child_id); + } + }; + const run = this._activeRlmChildRuns.get(subagent.rlm_child_id); + if (run?.detachedDeletion) { + // Keep every selector reserved until the run settles, or until a failed + // cleanup is exposed for an explicit retry. Repeated deletes before that + // boundary return the same accepted result. + void run.deletionReservation.promise.then(clearReservation, clearReservation); + } else { + clearReservation(); + } + } + } + + private _deleteRlmSubagentSession(childId: string, session?: AgentSession): Promise { + if (this._subagentRuntimeHost) { + return this._subagentRuntimeHost.deleteRlmSubagentRuntime(childId, session); + } + return session?.disposeAsync() ?? Promise.resolve(); + } + + private _ensureRlmRunDeletionCleanup(run: RlmChildRun, session: AgentSession): Promise { + if (run.deletionCleanup) return run.deletionCleanup; + const cleanup = Promise.resolve().then(() => this._deleteRlmSubagentSession(run.id, session)); + run.deletionCleanup = cleanup; + // Deletion admission is intentionally nonblocking. The detached run owner + // joins this exact promise before settlement and records any failure. + void cleanup.catch(() => undefined); + return cleanup; + } + + private async _recordRlmRunDeletionCleanupFailure( + run: RlmChildRun, + subagent: RlmSubagentRegistryEntry, + session: AgentSession, + error: unknown, + ): Promise { + if (this.host.isDisposed()) { + run.suppressTerminalNotice = true; + await session.disposeAsync().catch(() => undefined); + if (!run.settled) await this._finishRlmRunDeletion(run); + return; + } + run.deletionCleanup = undefined; + run.deletionCleanupObserver = undefined; + run.deletionCleanupFailed = true; + run.session = session; + this._rlmChildCleanupFailures.set(run.id, subagent); + // Make retry admission available before waking the parent model with the + // retry-required notice. + run.deletionReservation.resolve(); + await Promise.resolve(); + await run.reportDeletionCleanupFailure?.(error); + } + + private async _finishRlmRunDeletion(run: RlmChildRun): Promise { + await run.completeDeletion?.(); + if (this._activeRlmChildRuns.get(run.id) === run) { + this._removeRlmSubagentTracking(run.id, run); + } + run.settled = true; + run.settlement.resolve(); + run.deletionReservation.resolve(); + this._unsettledRlmChildRuns.delete(run); + this.host.onSettled(); + } + + private _observeRlmRunDeletionCleanup( + run: RlmChildRun, + subagent: RlmSubagentRegistryEntry, + session: AgentSession, + cleanup: Promise, + ): Promise { + if (run.deletionCleanupObserver) return run.deletionCleanupObserver; + const observer = cleanup.then( + () => true, + async (error) => { + await this._recordRlmRunDeletionCleanupFailure(run, subagent, session, error); + return false; + }, + ); + run.deletionCleanupObserver = observer; + void observer.catch(() => undefined); + return observer; + } + + private _continueFinishedRlmRunDeletion( + run: RlmChildRun, + subagent: RlmSubagentRegistryEntry, + session: AgentSession, + ): void { + const cleanup = this._ensureRlmRunDeletionCleanup(run, session); + const observer = this._observeRlmRunDeletionCleanup(run, subagent, session, cleanup); + if (!run.deletionRunFinished) return; + void observer + .then(async (cleanupSucceeded) => { + if (cleanupSucceeded) await this._finishRlmRunDeletion(run); + }) + .catch(() => undefined); + } + + private _removeRlmSubagentTracking(childId: string, run?: RlmChildRun): void { + run?.unsubscribe?.(); + this._rlmChildUnsubscribes.get(childId)?.(); + this._rlmChildUnsubscribes.delete(childId); + this._rlmChildSessions.delete(childId); + this._rlmChildCleanupFailures.delete(childId); + this._abandonedRlmQuiescenceChildIds.delete(childId); + if (!run || this._activeRlmChildRuns.get(childId) === run) { + this._activeRlmChildRuns.delete(childId); + } + if (run) { + run.abort = noopRlmChildAbort; + run.unsubscribe = undefined; + run.session = undefined; + } + } + + private _emitRlmSubagentRemoval(subagent: RlmSubagentRegistryEntry): void { + this.host.emit({ + type: "rlm_child_update", + child: { + id: subagent.rlm_child_id, + parentId: this.host.getParentNodeId(), + activeSessionId: subagent.active_session_id ?? undefined, + sessionName: subagent.session_name, + label: subagent.session_name, + status: "cancelled", + sessionDir: subagent.session_dir, + error: "Deleted by parent orchestrator", + }, + }); + } + + private async _deleteResolvedRlmSubagent(subagent: RlmSubagentRegistryEntry): Promise { + const childId = subagent.rlm_child_id; + const run = this._activeRlmChildRuns.get(childId); + if (run) { + if (run.deletionCleanupFailed) { + // Reset retry coordination only after selector preflight reaches the + // resolved child. A failed preflight must leave the prior retry boundary + // intact so a later call can acquire it. + run.deletionCleanupFailed = false; + run.deletionFailureNotice = undefined; + run.deletionReservation = createChildDeferred(); + } + // The detached task remains the sole lifecycle owner. Mark deletion before + // cancellation so its catch/finally path cannot race a normal release or + // terminal notice against the physical delete. + run.detachedDeletion = subagent; + if (this._cancelRlmChildRun(run, "Deleted by parent orchestrator")) { + run.deletionNeedsCompletionNotice = true; + } else { + this._emitRlmSubagentRemoval(subagent); + } + const liveSession = run.session; + if (run.status === "error" && !liveSession && run.settled) { + this._deletedRlmChildIds.add(childId); + this._removeRlmSubagentTracking(childId, run); + return { subagent }; + } + if (liveSession && run.settled) { + run.deletionRunFinished = true; + run.settlement = createChildDeferred(); + run.settled = false; + this._unsettledRlmChildRuns.add(run); + } + if (liveSession) this._continueFinishedRlmRunDeletion(run, subagent, liveSession); + + // Return once deletion is accepted. The run stays hidden but unsettled until + // abort-insensitive model/tool work unwinds and the shared cleanup finishes. + this._deletedRlmChildIds.add(childId); + return { subagent }; + } + + this._emitRlmSubagentRemoval(subagent); + const retained = this._rlmChildSessions.get(childId)?.session; + try { + await this._deleteRlmSubagentSession(childId, retained); + } catch (error) { + if (this.host.isDisposed()) { + this._removeRlmSubagentTracking(childId); + void retained?.disposeAsync().catch(() => undefined); + } else { + this._rlmChildCleanupFailures.set(childId, subagent); + } + throw error; + } + this._deletedRlmChildIds.add(childId); + this._removeRlmSubagentTracking(childId); + return { subagent }; + } + + registerRlmChildSession(childId: string, session: AgentSession, unsubscribe?: () => void): boolean { + // A child can finish concurrently while the parent is (or has) torn down; don't + // resurrect the map (it would never be disposed), just drop the child now. + if (this._deletingRlmChildren.has(childId) || this._deletedRlmChildIds.has(childId)) { + return false; + } + if (this._subagentRuntimeHost?.completeRlmSubagentRuntime?.(childId, session) === false) { + return false; + } + if (this.host.isDisposed()) { + void session.disposeAsync().catch(() => undefined); + return false; + } + this._rlmChildSessions.set(childId, { session, run: this._activeRlmChildRuns.get(childId) }); + if (unsubscribe) { + this._rlmChildUnsubscribes.set(childId, unsubscribe); + } + return true; + } + + releaseRlmChildSession(childId: string, session: AgentSession): (() => void) | false { + const run = this._activeRlmChildRuns.get(childId); + if (run?.session === session && run.status === "done") { + const unsubscribe = run.unsubscribe ?? noopRlmChildEventUnsubscribe; + return () => { + run.unsubscribe = undefined; + this._activeRlmChildRuns.delete(childId); + unsubscribe(); + }; + } + if (this._rlmChildSessions.get(childId)?.session !== session) return false; + const unsubscribe = this._rlmChildUnsubscribes.get(childId) ?? noopRlmChildEventUnsubscribe; + return () => { + this._rlmChildUnsubscribes.delete(childId); + this._rlmChildSessions.delete(childId); + unsubscribe(); + }; + } + + private _rlmChildSnapshotForRun( + run: RlmChildRun, + child = run.session ?? this._rlmChildSessions.get(run.id)?.session, + ): RlmChildAgentSnapshot { + return snapshotChildRun(run, child, this.host.getParentNodeId()); + } + + private _rlmChildSnapshotForSession(childId: string, child: AgentSession): RlmChildAgentSnapshot { + return snapshotRetainedChild(childId, child, this.host.getParentNodeId(), this.host.getChildSessionDir(child)); + } + + private _isUnboundTerminalRlmChildRun(run: RlmChildRun): boolean { + if (run.session !== undefined || this._rlmChildSessions.has(run.id)) return false; + return run.status === "done" || run.status === "error" || run.status === "cancelled"; + } + + getRlmChildSnapshots(): RlmChildAgentSnapshot[] { + const snapshots: RlmChildAgentSnapshot[] = []; + const recorded = new Set(); + const traversed = new Set(); + for (const run of this._activeRlmChildRuns.values()) { + const hidden = + run.detachedDeletion || + this._deletingRlmChildren.has(run.id) || + this._deletedRlmChildIds.has(run.id) || + this._isUnboundTerminalRlmChildRun(run); + const child = run.session; + if (!hidden) { + snapshots.push(this._rlmChildSnapshotForRun(run)); + recorded.add(run.id); + } + if (child) { + traversed.add(run.id); + snapshots.push(...child.getRlmChildSnapshots()); + } + } + for (const [childId, { session: child, run }] of this._rlmChildSessions) { + if (recorded.has(childId) || traversed.has(childId)) continue; + const hidden = this._deletingRlmChildren.has(childId) || this._deletedRlmChildIds.has(childId); + if (!hidden) { + const snapshot = run + ? this._rlmChildSnapshotForRun(run, child) + : this._rlmChildSnapshotForSession(childId, child); + snapshots.push({ + ...snapshot, + status: this._rlmChildCleanupFailures.has(childId) ? "cancelled" : snapshot.status, + }); + } + snapshots.push(...child.getRlmChildSnapshots()); + } + return snapshots; + } + + hasRunningRlmChildren(): boolean { + for (const session of this._rlmSubtreeSessions()) { + for (const run of session._activeRlmChildRuns.values()) { + if (run.status === "running" || run.status === "queued") { + return true; + } + } + } + return false; + } + + private _rlmChildSessionSnapshot(): AgentSession[] { + const sessions = new Set(); + for (const [childId, { session }] of this._rlmChildSessions) { + if (!this._abandonedRlmQuiescenceChildIds.has(childId)) sessions.add(session); + } + for (const run of this._activeRlmChildRuns.values()) { + if (run.session && !run.abandonedForQuiescence) sessions.add(run.session); + } + return [...sessions]; + } + + hasUnsettledWork(): boolean { + if (this.host.hasDeferredTerminalNotices()) return true; + if ([...this._unsettledRlmChildRuns].some((run) => !run.settled)) return true; + return this._rlmChildSessionSnapshot().some( + (child) => child.isSessionActive || this.host.getChildOwner(child).hasUnsettledWork(), + ); + } + + async waitForRlmQuiescence(externalSignal?: AbortSignal): Promise { + const cancellation = new AbortController(); + const cancelFromParent = () => cancellation.abort(); + if (externalSignal?.aborted) cancellation.abort(); + else externalSignal?.addEventListener("abort", cancelFromParent, { once: true }); + this._rlmQuiescenceWaitAborts.add(cancellation); + let rejectCancelled = (_error: Error) => {}; + const cancelled = new Promise((_resolve, reject) => { + rejectCancelled = reject; + }); + const onCancelled = () => rejectCancelled(new Error("RLM quiescence wait cancelled")); + cancellation.signal.addEventListener("abort", onCancelled, { once: true }); + if (cancellation.signal.aborted) onCancelled(); + const wait = (operation: Promise): Promise => Promise.race([operation, cancelled]); + try { + while (true) { + await wait(this.host.waitForHeadlessIdle()); + // Strong RLM quiescence also owns work that interactive waitForIdle ignores. + if (this.host.isSessionActive() || this.host.hasDeferredTerminalNotices()) { + await wait(this.host.waitForActivityChange(cancellation.signal)); + continue; + } + const unsettledRuns = [...this._unsettledRlmChildRuns].filter((run) => !run.settled); + const childSessions = this._rlmChildSessionSnapshot(); + if (unsettledRuns.length === 0 && !this.hasUnsettledWork()) return; + await wait( + Promise.all([ + ...unsettledRuns.map((run) => run.settlement.promise), + ...childSessions.map((child) => child.waitForRlmQuiescence(cancellation.signal)), + ]), + ); + // Always loop through the self-active/deferred checks again. Work may + // start at the child-settlement boundary. + } + } finally { + // A local descendant error must cancel sibling recursive waits owned by + // this barrier before their propagation listeners are removed. + cancellation.abort(); + externalSignal?.removeEventListener("abort", cancelFromParent); + cancellation.signal.removeEventListener("abort", onCancelled); + this._rlmQuiescenceWaitAborts.delete(cancellation); + } + } + + getRlmChildSession(childId: string): AgentSession | undefined { + for (const session of this._rlmSubtreeSessions()) { + const direct = + session._activeRlmChildRuns.get(childId)?.session ?? session._rlmChildSessions.get(childId)?.session; + if (direct) { + return direct; + } + } + return undefined; + } + + cancelRlmChildRun(childId: string, reason = "Cancelled by user"): boolean { + for (const session of this._rlmSubtreeSessions()) { + const run = session._activeRlmChildRuns.get(childId); + if (run) { + if (run.status !== "running" && run.status !== "queued" && !run.settled) { + if (session.host.isInputSuspended()) session._abandonRlmRunForQuiescence(run); + else run.suppressTerminalNotice = true; + return true; + } + // The abort cascade never reaches running work retained under a settled descendant. + const cancelled = session._cancelRlmChildRun(run, reason); + const descendantsCancelled = run.session?.cancelRunningRlmDescendants(reason) ?? false; + if (cancelled || descendantsCancelled) { + return true; + } + } + // A fruitless match keeps walking: child ids are only mkdir-unique among + // siblings, so a colliding live run elsewhere must stay reachable. + if (session._rlmChildSessions.get(childId)?.session.cancelRunningRlmDescendants(reason)) { + return true; + } + } + return false; + } + + // A completed child may belong to both maps; visit each owner once. + private *_rlmSubtreeSessions(): Generator { + const visited = new Set([this]); + const stack: SessionChildren[] = [this]; + while (stack.length > 0) { + const session = stack.pop()!; + yield session; + for (const run of session._activeRlmChildRuns.values()) { + if (run.session && !visited.has(this.host.getChildOwner(run.session))) { + visited.add(this.host.getChildOwner(run.session)); + stack.push(this.host.getChildOwner(run.session)); + } + } + for (const { session: retained } of session._rlmChildSessions.values()) { + if (!visited.has(this.host.getChildOwner(retained))) { + visited.add(this.host.getChildOwner(retained)); + stack.push(this.host.getChildOwner(retained)); + } + } + } + } + + cancelRunningRlmDescendants(reason = "Cancelled by user"): boolean { + let cancelled = false; + for (const session of this._rlmSubtreeSessions()) { + for (const run of session._activeRlmChildRuns.values()) { + if (session._cancelRlmChildRun(run, reason)) cancelled = true; + } + } + return cancelled; + } + + private async _assertRlmSubagentSessionNameAvailable(name: string, ignorePendingReservation = false): Promise { + const depth = this.host.getDepth() + 1; + if (!ignorePendingReservation && this._pendingRlmSubagentSessionNames.has(name)) { + throw new Error(formatAgentSessionNameUnavailable(name, depth)); + } + const localConflict = + [...this._activeRlmChildRuns.values()].some( + (run) => run.session?.sessionName === name || (!run.session && run.sessionName === name), + ) || + [...this._rlmChildSessions.values()].some(({ session }) => session.sessionName === name) || + [...this._rlmChildCleanupFailures.values()].some((entry) => entry.session_name === name); + if (localConflict) { + throw new Error(formatAgentSessionNameUnavailable(name, depth)); + } + const controller = this.host.getMessageController(); + if (!controller) return; + const input = { + name, + depth, + parentSessionId: this.host.getSessionId(), + parentSessionPath: this.host.getSessionFile(), + }; + if (controller.assertSessionNameAvailable) { + await controller.assertSessionNameAvailable(input); + return; + } + const listed = await controller.listAgents(); + const catalog = listed.agents.map( + (agent): AgentFamilyCatalogEntry => ({ + id: agent.sessionId, + ...(agent.sessionName ? { name: agent.sessionName } : {}), + depth: agent.rlmDepth ?? 0, + status: agent.status ?? "idle", + ...(agent.parentSessionId ? { parentSessionId: agent.parentSessionId } : {}), + ...(agent.parentSessionPath ? { parentSessionPath: agent.parentSessionPath } : {}), + ...(agent.sessionPath ? { sessionPath: agent.sessionPath } : {}), + }), + ); + assertAgentSessionNameAvailable(catalog, input); + } + + async reapAfterCompaction(): Promise { + const childIds = [...this._rlmChildCleanupFailures.keys()].filter( + (childId) => !this._activeRlmChildRuns.get(childId)?.detachedDeletion, + ); + await Promise.allSettled(childIds.map((childId) => this.host.deleteRlmSubagent(childId))); + } + + async run(prompt: string, kwargs: Record = {}, spawnCode?: string): Promise { + // Snapshot before any await: the spawning request is the turn whose tool call is + // executing now. A spawn arriving outside an active run (a detached kernel task + // firing while the parent is idle) has no such turn; an absent edge beats a wrong one. + const spawnedByRequestId = this.host.isStreaming() ? this.host.getSemanticEdges().lastTurnRequestId : undefined; + const { name: rawName, model: rawModel, thinking: rawThinking, ...unsupported } = kwargs; + const unsupportedKwargs = Object.keys(unsupported); + if (unsupportedKwargs.length > 0) { + throw new Error(`Unsupported rlm.spawn kwargs: ${unsupportedKwargs.sort().join(", ")}`); + } + const requestedSessionName = normalizeRequestedRlmSubagentSessionName(rawName); + const requestedModel = normalizeRequestedRlmSubagentModel(rawModel); + const requestedThinkingLevel = normalizeRequestedRlmSubagentThinkingLevel(rawThinking); + if (requestedSessionName) assertDirectAgentMessageTarget(requestedSessionName); + if (this.host.getDepth() >= this.host.getMaxDepth()) { + throw new Error( + `RLM recursion depth limit reached (RLM_DEPTH=${this.host.getDepth()}, RLM_MAX_DEPTH=${this.host.getMaxDepth()})`, + ); + } + if (requestedSessionName) { + if (this._pendingRlmSubagentSessionNames.has(requestedSessionName)) { + throw new Error(formatAgentSessionNameUnavailable(requestedSessionName, this.host.getDepth() + 1)); + } + this._pendingRlmSubagentSessionNames.add(requestedSessionName); + } + let modelSelection: RlmSubagentModelSelection; + try { + if (requestedSessionName) await this._assertRlmSubagentSessionNameAvailable(requestedSessionName, true); + modelSelection = await this.host.resolveModel(requestedModel); + } finally { + if (requestedSessionName) this._pendingRlmSubagentSessionNames.delete(requestedSessionName); + } + if (requestedThinkingLevel !== undefined) { + const supported = getSupportedThinkingLevels(modelSelection.model) as ThinkingLevel[]; + if (!supported.includes(requestedThinkingLevel)) { + throw new Error( + `Requested thinking level "${requestedThinkingLevel}" is not supported by model "${modelSelection.model.provider}/${modelSelection.model.id}"; supported levels: ${supported.join(", ")}`, + ); + } + } + if (this.host.isDisposed()) throw new Error("Cannot spawn a subagent after its parent was disposed"); + + const childSessionDir = this.host.createSessionDir(); + const childNodeId = basename(childSessionDir); + const sessionName = requestedSessionName ?? createDefaultRlmSubagentSessionName(prompt, childNodeId); + if (!requestedSessionName) await this._assertRlmSubagentSessionNameAvailable(sessionName); + + return launchChildTask( + this.host, + { + admitRun: (run) => { + this._activeRlmChildRuns.set(run.id, run); + this._unsettledRlmChildRuns.add(run); + }, + isCurrentRun: (run) => this._activeRlmChildRuns.get(run.id) === run, + snapshotForRun: (run) => this._rlmChildSnapshotForRun(run), + registerSession: (id, child) => this.host.registerRlmChildSession(id, child), + currentActiveSessionId: () => this._currentActiveSessionId(), + getRuntimeHost: () => this._subagentRuntimeHost, + recordDeleted: (id) => { + this._deletedRlmChildIds.add(id); + }, + removeTracking: (id, run) => this._removeRlmSubagentTracking(id, run), + ensureDeletionCleanup: (run, child) => this._ensureRlmRunDeletionCleanup(run, child), + observeDeletionCleanup: (run, entry, child, cleanup) => + this._observeRlmRunDeletionCleanup(run, entry, child, cleanup), + finishDeletion: (run) => this._finishRlmRunDeletion(run), + finishRun: (run) => this.finishRun(run), + }, + { + id: childNodeId, + prompt, + sessionName, + spawnCode, + sessionDir: childSessionDir, + model: modelSelection.model, + thinkingLevel: requestedThinkingLevel, + spawnedByRequestId, + }, + ); + } + private finishRun(run: RlmChildRun): void { + if (this._activeRlmChildRuns.get(run.id) === run) { + if (this._rlmChildSessions.has(run.id)) { + this._activeRlmChildRuns.delete(run.id); + if (run.unsubscribe) this._rlmChildUnsubscribes.set(run.id, run.unsubscribe); + run.abort = noopRlmChildAbort; + run.unsubscribe = undefined; + run.session = undefined; + } else if (run.status !== "error") { + this._removeRlmSubagentTracking(run.id, run); + } else { + run.unsubscribe?.(); + run.abort = noopRlmChildAbort; + run.unsubscribe = undefined; + } + } + run.settled = true; + run.settlement.resolve(); + this._unsettledRlmChildRuns.delete(run); + this.host.onSettled(); + } + + async createRlmSession(prompt: string, kwargs: Record = {}): Promise { + const { name: rawName, model: rawModel, thinking: rawThinking, cwd: rawCwd, ...unsupported } = kwargs; + const unsupportedKeys = Object.keys(unsupported); + if (unsupportedKeys.length > 0) { + throw new Error(`Unsupported rlm.create_session kwargs: ${unsupportedKeys.sort().join(", ")}`); + } + if (!prompt.trim()) { + throw new Error("rlm.create_session prompt must not be empty"); + } + if (this.host.getDepth() !== 0) { + throw new Error("rlm.create_session is available only from a depth-0 session"); + } + if (this.host.isDisposed()) { + throw new Error("Cannot create a top-level session after the current session was disposed"); + } + const host = this._subagentRuntimeHost; + if (!host?.createRlmRootSession) { + throw new Error("rlm.create_session requires a daemon-backed depth-0 session"); + } + + const operation = "rlm.create_session"; + const sessionName = normalizeRequestedRlmSubagentSessionName(rawName, operation); + const requestedModel = normalizeRequestedRlmSubagentModel(rawModel, operation); + const requestedThinkingLevel = normalizeRequestedRlmSubagentThinkingLevel(rawThinking, operation); + if (sessionName) { + assertDirectAgentMessageTarget(sessionName); + const controller = this.host.getMessageController(); + if (controller?.assertSessionNameAvailable) { + await controller.assertSessionNameAvailable({ name: sessionName, depth: 0 }); + } + } + if (rawCwd !== undefined && (typeof rawCwd !== "string" || !rawCwd.trim())) { + throw new Error("rlm.create_session cwd must be a non-empty string"); + } + const cwd = rawCwd === undefined ? this.host.getCwd() : resolve(this.host.getCwd(), rawCwd.trim()); + const modelSelection = await this.host.resolveModel(requestedModel, "top-level session"); + if (requestedThinkingLevel !== undefined) { + const supported = getSupportedThinkingLevels(modelSelection.model) as ThinkingLevel[]; + if (!supported.includes(requestedThinkingLevel)) { + throw new Error( + `Requested thinking level "${requestedThinkingLevel}" is not supported by model "${modelSelection.model.provider}/${modelSelection.model.id}"; supported levels: ${supported.join(", ")}`, + ); + } + } + const thinkingLevel = + requestedThinkingLevel ?? + (clampThinkingLevel(modelSelection.model, this.host.getThinkingLevel()) as ThinkingLevel); + if (this.host.isDisposed()) { + throw new Error("Cannot create a top-level session after the current session was disposed"); + } + return host.createRlmRootSession({ + prompt, + sessionName, + cwd, + model: modelSelection.model, + thinkingLevel, + }); + } + async disposeAsync(afterChildren: () => Promise): Promise { + for (const run of [...this._activeRlmChildRuns.values()]) { + const childSession = run.session; + if (!childSession) continue; + if (run.detachedDeletion) { + run.suppressTerminalNotice = true; + if (run.deletionCleanupObserver) { + await run.deletionCleanupObserver.catch(() => false); + } else if (run.deletionCleanup) { + await run.deletionCleanup.catch(() => childSession.disposeAsync().catch(() => undefined)); + } else { + // Cleanup already failed and was exposed for retry before disposal. + await childSession.disposeAsync().catch(() => undefined); + } + if (!run.settled) await this._finishRlmRunDeletion(run); + } else { + await childSession.disposeAsync().catch(() => undefined); + } + } + for (const unsubscribe of this._rlmChildUnsubscribes.values()) { + unsubscribe(); + } + this._rlmChildUnsubscribes.clear(); + for (const { session } of this._rlmChildSessions.values()) { + await session.disposeAsync().catch(() => undefined); + } + this._rlmChildSessions.clear(); + this._rlmChildCleanupFailures.clear(); + this._deletedRlmChildIds.clear(); + return afterChildren(); + } + dispose(): void { + this.cancelActiveRuns("Parent session disposed"); + for (const unsubscribe of this._rlmChildUnsubscribes.values()) { + unsubscribe(); + } + this._rlmChildUnsubscribes.clear(); + for (const { session } of this._rlmChildSessions.values()) { + session.dispose(); + } + this._rlmChildSessions.clear(); + this._rlmChildCleanupFailures.clear(); + this._deletedRlmChildIds.clear(); + } + + beginDisposal(): void { + for (const run of this._unsettledRlmChildRuns) run.suppressTerminalNotice = true; + this.cancelQuiescenceWaits(); + } + cancelQuiescenceWaits(): void { + for (const controller of this._rlmQuiescenceWaitAborts) controller.abort(); + } + requestAbort(): void { + for (const run of [...this._unsettledRlmChildRuns]) { + if (run.status === "cancelled") this._abandonRlmRunForQuiescence(run); + } + this.cancelQuiescenceWaits(); + } + setRuntimeHost(host?: SubagentRuntimeHost): void { + this._subagentRuntimeHost = host; + } + getRuntimeHost(): SubagentRuntimeHost | undefined { + return this._subagentRuntimeHost; + } + getActiveRuns(): Iterable> { + return this._activeRlmChildRuns.values(); + } +} diff --git a/packages/coding-agent/src/session/extensions.ts b/packages/coding-agent/src/session/extensions.ts new file mode 100644 index 0000000000..30fb280176 --- /dev/null +++ b/packages/coding-agent/src/session/extensions.ts @@ -0,0 +1,404 @@ +import { basename, dirname } from "node:path"; +import type { Agent, ThinkingLevel } from "@earendil-works/pi-agent-core"; +import { type Api, type Model, resetApiProviders } from "@earendil-works/pi-ai"; +import type { AgentSessionMessageController } from "../core/agent-messages.js"; +import type { CompactionResult } from "../core/compaction/index.js"; +import { + type ContextUsage, + type ExtensionActions, + type ExtensionCommandContextActions, + type ExtensionErrorListener, + ExtensionRunner, + type ExtensionUIContext, + type SessionStartEvent, + type ShutdownHandler, + type ToolInfo, +} from "../core/extensions/index.js"; +import { emitSessionShutdownEvent } from "../core/extensions/runner.js"; +import type { McpManager } from "../core/mcp/mcp-manager.js"; +import type { ModelRegistry } from "../core/model-registry.js"; +import type { PromptTemplate } from "../core/prompt-templates.js"; +import type { ResourceExtensionPaths, ResourceLoader } from "../core/resource-loader.js"; +import type { SessionManager } from "../core/session-manager.js"; +import type { SlashCommandInfo } from "../core/slash-commands.js"; + +export interface ExtensionBindings { + uiContext?: ExtensionUIContext; + commandContextActions?: ExtensionCommandContextActions; + shutdownHandler?: ShutdownHandler; + onError?: ExtensionErrorListener; +} +export interface SessionExtensionsHost { + cwd: string; + sessionManager: SessionManager; + resourceLoader: ResourceLoader; + modelRegistry: ModelRegistry; + getModelRegistry(): ModelRegistry; + getPromptTemplates(): ReadonlyArray; + bindShutdownHandler(handler: ShutdownHandler | undefined): ShutdownHandler | undefined; + getAgentMessageController(): AgentSessionMessageController | undefined; + refreshCurrentModel(): void; + sendCustomMessage(...args: Parameters): Promise; + sendUserMessage(...args: Parameters): Promise; + setSessionName(name: string): void; + getActiveToolNames(): string[]; + getAllTools(): ToolInfo[]; + setActiveToolsByName(names: string[]): void; + refreshTools(): void; + setModel(model: Model): Promise; + getThinkingLevel(): ThinkingLevel; + setThinkingLevel(level: ThinkingLevel): void; + getModel(): Model | undefined; + isStreaming(): boolean; + getSignal(): AbortSignal | undefined; + abort(): Promise; + getQueuedActionCount(): number; + getContextUsage(): ContextUsage | undefined; + compact(instructions?: string): Promise; + getSystemPrompt(): string; + rebuildSystemPrompt(): void; + reloadSettings(): Promise; + getMcpManager(): McpManager | undefined; + rebuildRuntime(options: { + activeToolNames?: string[]; + flagValues?: Map; + includeAllExtensionTools?: boolean; + }): void; +} +export class SessionExtensions { + private _extensionRunner!: ExtensionRunner; + private _execEnvProvider?: () => Record | undefined; + private _extensionUIContext?: ExtensionUIContext; + private _extensionCommandContextActions?: ExtensionCommandContextActions; + private _extensionShutdownHandler?: ShutdownHandler; + private _extensionErrorListener?: ExtensionErrorListener; + private _extensionErrorUnsubscriber?: () => void; + constructor( + private readonly host: SessionExtensionsHost, + private readonly _sessionStartEvent: SessionStartEvent, + private readonly _extensionRunnerRef?: { current?: ExtensionRunner }, + ) {} + get runner(): ExtensionRunner { + return this._extensionRunner; + } + setExecEnvProvider(provider: (() => Record | undefined) | undefined): void { + this._execEnvProvider = provider; + const extensions = this.host.resourceLoader.getExtensions(); + extensions.runtime.getExecEnv = provider; + } + + async bindExtensions(bindings: ExtensionBindings): Promise { + if (bindings.uiContext !== undefined) { + this._extensionUIContext = bindings.uiContext; + } + if (bindings.commandContextActions !== undefined) { + this._extensionCommandContextActions = bindings.commandContextActions; + } + if (bindings.shutdownHandler !== undefined) { + this._extensionShutdownHandler = this.host.bindShutdownHandler(bindings.shutdownHandler); + } + if (bindings.onError !== undefined) { + this._extensionErrorListener = bindings.onError; + } + + this._applyExtensionBindings(this._extensionRunner); + await this._extensionRunner.emit(this._sessionStartEvent); + await this.extendResourcesFromExtensions(this._sessionStartEvent.reason === "reload" ? "reload" : "startup"); + } + + private async extendResourcesFromExtensions(reason: "startup" | "reload"): Promise { + if (!this._extensionRunner.hasHandlers("resources_discover")) { + return; + } + + const { skillPaths, promptPaths, themePaths } = await this._extensionRunner.emitResourcesDiscover( + this.host.cwd, + reason, + ); + + if (skillPaths.length === 0 && promptPaths.length === 0 && themePaths.length === 0) { + return; + } + + const extensionPaths: ResourceExtensionPaths = { + skillPaths: this.buildExtensionResourcePaths(skillPaths), + promptPaths: this.buildExtensionResourcePaths(promptPaths), + themePaths: this.buildExtensionResourcePaths(themePaths), + }; + + this.host.resourceLoader.extendResources(extensionPaths); + this.host.rebuildSystemPrompt(); + } + + private buildExtensionResourcePaths(entries: Array<{ path: string; extensionPath: string }>): Array<{ + path: string; + metadata: { + source: string; + scope: "temporary"; + origin: "top-level"; + baseDir?: string; + }; + }> { + return entries.map((entry) => { + const source = this.getExtensionSourceLabel(entry.extensionPath); + const baseDir = entry.extensionPath.startsWith("<") ? undefined : dirname(entry.extensionPath); + return { + path: entry.path, + metadata: { + source, + scope: "temporary", + origin: "top-level", + baseDir, + }, + }; + }); + } + + private getExtensionSourceLabel(extensionPath: string): string { + if (extensionPath.startsWith("<")) { + return `extension:${extensionPath.replace(/[<>]/g, "")}`; + } + const base = basename(extensionPath); + const name = base.replace(/\.(ts|js)$/, ""); + return `extension:${name}`; + } + + private _applyExtensionBindings(runner: ExtensionRunner): void { + runner.setUIContext(this._extensionUIContext); + runner.bindCommandContext(this._extensionCommandContextActions); + + this._extensionErrorUnsubscriber?.(); + this._extensionErrorUnsubscriber = this._extensionErrorListener + ? runner.onError(this._extensionErrorListener) + : undefined; + } + + private _bindExtensionCore(runner: ExtensionRunner): void { + const getCommands = (): SlashCommandInfo[] => { + const extensionCommands: SlashCommandInfo[] = runner.getRegisteredCommands().map((command) => ({ + name: command.invocationName, + description: command.description, + source: "extension", + sourceInfo: command.sourceInfo, + })); + + const templates: SlashCommandInfo[] = this.host.getPromptTemplates().map((template) => ({ + name: template.name, + description: template.description, + source: "prompt", + sourceInfo: template.sourceInfo, + })); + + const skills: SlashCommandInfo[] = this.host.resourceLoader.getSkills().skills.map((skill) => ({ + name: `skill:${skill.name}`, + description: skill.description, + source: "skill", + sourceInfo: skill.sourceInfo, + })); + + return [...extensionCommands, ...templates, ...skills]; + }; + + runner.bindCore( + { + sendMessage: (message, options) => { + this.host.sendCustomMessage(message, options).catch((err) => { + runner.emitError({ + extensionPath: "", + event: "send_message", + error: err instanceof Error ? err.message : String(err), + }); + }); + }, + sendUserMessage: (content, options) => { + this.host.sendUserMessage(content, options).catch((err) => { + runner.emitError({ + extensionPath: "", + event: "send_user_message", + error: err instanceof Error ? err.message : String(err), + }); + }); + }, + appendEntry: (customType, data) => { + this.host.sessionManager.appendCustomEntry(customType, data); + }, + setSessionName: async (name) => { + const controller = this.host.getAgentMessageController(); + if (controller?.setSessionName) { + await controller.setSessionName(name); + return; + } + this.host.setSessionName(name); + }, + getSessionName: () => { + return this.host.sessionManager.getSessionName(); + }, + setLabel: (entryId, label) => { + this.host.sessionManager.appendLabelChange(entryId, label); + }, + getActiveTools: () => this.host.getActiveToolNames(), + getAllTools: () => this.host.getAllTools(), + setActiveTools: (toolNames) => this.host.setActiveToolsByName(toolNames), + refreshTools: () => this.host.refreshTools(), + getCommands, + setModel: async (model) => { + if (!this.host.getModelRegistry().hasConfiguredAuth(model)) return false; + await this.host.setModel(model); + return true; + }, + getThinkingLevel: () => this.host.getThinkingLevel(), + setThinkingLevel: (level) => this.host.setThinkingLevel(level), + }, + { + getModel: () => this.host.getModel(), + isIdle: () => !this.host.isStreaming(), + getSignal: () => this.host.getSignal(), + abort: () => this.host.abort(), + hasPendingMessages: () => this.host.getQueuedActionCount() > 0, + shutdown: () => { + this._extensionShutdownHandler?.(); + }, + getContextUsage: () => this.host.getContextUsage(), + compact: (options) => { + void (async () => { + try { + const result = await this.host.compact(options?.customInstructions); + options?.onComplete?.(result); + } catch (error) { + const err = error instanceof Error ? error : new Error(String(error)); + options?.onError?.(err); + } + })(); + }, + getSystemPrompt: () => this.host.getSystemPrompt(), + }, + { + registerProvider: (name, config) => { + this.host.modelRegistry.registerProvider(name, config); + this.host.refreshCurrentModel(); + }, + unregisterProvider: (name) => { + this.host.modelRegistry.unregisterProvider(name); + this.host.refreshCurrentModel(); + }, + }, + ); + } + + async reload(): Promise { + const previousFlagValues = this._extensionRunner.getFlagValues(); + await emitSessionShutdownEvent(this._extensionRunner, { + type: "session_shutdown", + reason: "reload", + }); + await this.host.reloadSettings(); + // Re-read auth.json: a login saved by the client process (daemon mode) must be + // visible here so MCP skill gating sees the new credentials. + this.host.modelRegistry.authStorage.reload(); + resetApiProviders(); + this.host.getMcpManager()?.refresh(); + await this.host.resourceLoader.reload(); + this.host.rebuildRuntime({ + activeToolNames: this.host.getActiveToolNames(), + flagValues: previousFlagValues, + includeAllExtensionTools: true, + }); + + const hasBindings = + this._extensionUIContext || + this._extensionCommandContextActions || + this._extensionShutdownHandler || + this._extensionErrorListener; + if (hasBindings) { + await this._extensionRunner.emit({ + type: "session_start", + reason: "reload", + }); + await this.extendResourcesFromExtensions("reload"); + } + } + + build(flagValues?: Map): void { + const extensionsResult = this.host.resourceLoader.getExtensions(); + if (flagValues) { + for (const [name, value] of flagValues) { + extensionsResult.runtime.flagValues.set(name, value); + } + } + // Re-apply on (re)build so the provider survives /reload. Guarded: the + // runtime object can be shared across sessions from one ResourceLoader + // (RLM children), so a provider-less session must not wipe the owner's. + if (this._execEnvProvider) { + extensionsResult.runtime.getExecEnv = this._execEnvProvider; + } + + this._extensionRunner = new ExtensionRunner( + extensionsResult.extensions, + extensionsResult.runtime, + this.host.cwd, + this.host.sessionManager, + this.host.modelRegistry, + ); + if (this._extensionRunnerRef) { + this._extensionRunnerRef.current = this._extensionRunner; + } + this._bindExtensionCore(this._extensionRunner); + this._applyExtensionBindings(this._extensionRunner); + } +} + +export function installExtensionToolHooks( + agent: Pick, + getRunner: () => ExtensionRunner, + getEventQueue: () => Promise, +): void { + agent.beforeToolCall = async ({ toolCall, args }) => { + const runner = getRunner(); + if (!runner.hasHandlers("tool_call")) { + return undefined; + } + + await getEventQueue(); + + try { + return await runner.emitToolCall({ + type: "tool_call", + toolName: toolCall.name, + toolCallId: toolCall.id, + input: args as Record, + }); + } catch (err) { + if (err instanceof Error) { + throw err; + } + throw new Error(`Extension failed, blocking execution: ${String(err)}`); + } + }; + + agent.afterToolCall = async ({ toolCall, args, result, isError }) => { + const runner = getRunner(); + if (!runner.hasHandlers("tool_result")) { + return undefined; + } + + const hookResult = await runner.emitToolResult({ + type: "tool_result", + toolName: toolCall.name, + toolCallId: toolCall.id, + input: args as Record, + content: result.content, + details: result.details, + isError, + }); + + if (!hookResult) { + return undefined; + } + + return { + content: hookResult.content, + details: hookResult.details, + isError: hookResult.isError ?? isError, + }; + }; +} diff --git a/packages/coding-agent/src/session/kernel-environment.ts b/packages/coding-agent/src/session/kernel-environment.ts new file mode 100644 index 0000000000..dc45cb74e6 --- /dev/null +++ b/packages/coding-agent/src/session/kernel-environment.ts @@ -0,0 +1,91 @@ +import { mkdirSync, mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { AuthStorage } from "../core/auth-storage.js"; +import { getGlobalHarnessStateDir, getLocalHarnessStateDir } from "../core/refinement/index.js"; +import { resolveConfigValue } from "../core/resolve-config-value.js"; +import type { ResourceLoader } from "../core/resource-loader.js"; +import { SERPER_CREDENTIAL_ID, SERPER_ENV_VAR, WEBSEARCH_SKILL_NAME } from "../core/websearch-credential.js"; + +export interface KernelEnvironmentHost { + agentDir?: string; + authStorage: Pick; + resourceLoader: Pick; + getDepth(): number; + getMaxDepth(): number; + getArtifactDir(): string | undefined; + getLocalHarnessStateDir(): string | undefined; +} + +export class KernelEnvironment { + constructor( + private readonly host: KernelEnvironmentHost, + private _rlmSessionDir?: string, + ) {} + get sessionDir(): string | undefined { + return this._rlmSessionDir; + } + buildEnv(): Record { + // Kernel env is provisioning-time only: RLM_MAX_DEPTH may be stale in an already-running kernel; + // the TypeScript-side spawn check remains authoritative. + const env: Record = { + RLM_DEPTH: String(this.host.getDepth()), + RLM_MAX_DEPTH: String(this.host.getMaxDepth()), + RLM_GLOBAL_HARNESS_STATE_DIR: getGlobalHarnessStateDir(), + }; + const rlmSessionDir = this.ensureSessionDir(); + if (rlmSessionDir) { + env.RLM_SESSION_DIR = rlmSessionDir; + // Keep kernel writes and host reads (system prompt, review, /refine) on + // the same local harness path. Subagents prefer their own artifact dir; + // ephemeral sessions fall back to the RLM session dir once it exists. + env.RLM_HARNESS_STATE_DIR = this.host.getLocalHarnessStateDir() ?? getLocalHarnessStateDir(rlmSessionDir)!; + } + this._addWebsearchKeyEnv(env); + return env; + } + + private _addWebsearchKeyEnv(env: Record): void { + if (this.host.agentDir) { + env.PRIME_AGENT_CODING_AGENT_DIR = this.host.agentDir; + } + + if (process.env[SERPER_ENV_VAR]?.trim()) { + return; + } + // Inject only when a websearch skill (bundled or custom) is actually loaded, + // so the key isn't exposed to kernels that can't use it. + if (!this.host.resourceLoader.getSkills().skills.some((skill) => skill.name === WEBSEARCH_SKILL_NAME)) { + return; + } + const cred = this.host.authStorage.get(SERPER_CREDENTIAL_ID); + if (cred?.type !== "api_key") { + return; + } + const resolved = resolveConfigValue(cred.key)?.trim(); + if (resolved) { + env[SERPER_ENV_VAR] = resolved; + } + } + + ensureSessionDir(): string | undefined { + if (this._rlmSessionDir) { + mkdirSync(this._rlmSessionDir, { recursive: true }); + return this._rlmSessionDir; + } + + const sessionArtifactDir = this.host.getArtifactDir(); + if (sessionArtifactDir) { + mkdirSync(sessionArtifactDir, { recursive: true }); + this._rlmSessionDir = sessionArtifactDir; + return sessionArtifactDir; + } + + return undefined; + } + + createEphemeralSessionDir(): string { + this._rlmSessionDir = mkdtempSync(join(tmpdir(), "prime-agent-rlm-")); + return this._rlmSessionDir; + } +} diff --git a/packages/coding-agent/src/session/kernel-host-handlers.ts b/packages/coding-agent/src/session/kernel-host-handlers.ts new file mode 100644 index 0000000000..e3fdffddaf --- /dev/null +++ b/packages/coding-agent/src/session/kernel-host-handlers.ts @@ -0,0 +1,175 @@ +import type { Api, Model } from "@earendil-works/pi-ai"; +import { + AGENT_MESSAGE_SKILL_NAME, + type AgentSessionMessageController, + type AgentSessionMessageReceipt, + agentFamilyMemberName, + createAgentMessageHostHandlers, +} from "../core/agent-messages.js"; +import { + type AgentObserveAgentSnapshot, + type AgentObserveListResult, + type AgentObserveRecentMessagesResult, + createAgentObserveHostHandlers, +} from "../core/agent-observe.js"; +import type { HostRequestHandlers } from "../core/kernel/index.js"; +import type { McpManager } from "../core/mcp/mcp-manager.js"; +import type { AsyncBashCompletionDetails } from "../core/messages.js"; +import { + createAsyncBashCompletionHostHandler, + createAsyncBashConsumedHostHandler, + createRlmCreateSessionHostHandler, + createRlmDeleteSubagentHostHandler, + createRlmFindModelsHostHandler, + createRlmListSubagentsHostHandler, + createRlmRunHostHandler, + type RlmCreateSessionResult, + type RlmDeleteSubagentResult, + type RlmFindModelsResult, + type RlmListSubagentsResult, + type RlmSpawnHandle, +} from "../core/rlm-runtime.js"; +import type { Skill } from "../core/skills.js"; + +type ObserveResult = AgentObserveListResult | AgentObserveAgentSnapshot | AgentObserveRecentMessagesResult; +export interface SessionKernelOperations { + runChild(prompt: string, kwargs: Record, cellSourceCode?: string): Promise; + createSession(prompt: string, kwargs: Record): Promise; + findModels(query: string, limit: number): Promise; + listSubagents(): Promise; + deleteSubagent(target: string): Promise; + handleBashCompletion(details: AsyncBashCompletionDetails): Promise; + withdrawBashCompletion(details: { pid: number; command: string }): void; + getModel(): Model | undefined; + includeGoals: boolean; + includeCompactSkill: boolean; + isRefineAllowed(): boolean; + hasHeartbeatController(): boolean; + getModelVisibleSkills(): Skill[]; + getAgentMessageController(): AgentSessionMessageController | undefined; + hasObserveController(): boolean; + getMcpManager(): McpManager | undefined; + getDepth(): number; + awaitChildPublication(selector: string): Promise; + recordParentReply(): void; + handleGoal(type: string, payload: Record): Record; + handleCompact(type: string, payload: Record): Record; + handleRefine(type: string, payload: Record): Record; + handleHeartbeat(type: string, payload: Record): Record; + handleMessage(type: string, payload?: Record): Promise; + handleObserve(type: string, payload?: Record): ObserveResult | Promise; +} +export function createSessionKernelHostHandlers(host: SessionKernelOperations): HostRequestHandlers { + const handlers: HostRequestHandlers = { + "rlm.run": createRlmRunHostHandler(async ({ prompt, kwargs, cellSourceCode }) => ({ + ...(await host.runChild(prompt, kwargs, cellSourceCode)), + })), + "rlm.create_session": createRlmCreateSessionHostHandler(async ({ prompt, kwargs }) => ({ + ...(await host.createSession(prompt, kwargs)), + })), + "bash.completed": createAsyncBashCompletionHostHandler((details) => host.handleBashCompletion(details)), + "bash.consumed": createAsyncBashConsumedHostHandler((details) => { + host.withdrawBashCompletion(details); + }), + "rlm.find_models": createRlmFindModelsHostHandler((query, limit) => host.findModels(query, limit)), + "rlm.list_subagents": createRlmListSubagentsHostHandler(() => host.listSubagents()), + "rlm.delete_subagent": createRlmDeleteSubagentHostHandler((target) => host.deleteSubagent(target)), + "model.info": async () => ({ + id: host.getModel()?.id ?? null, + provider: host.getModel()?.provider ?? null, + input: host.getModel()?.input ?? [], + }), + }; + if (host.includeGoals) { + for (const type of ["goal.get", "goal.create", "goal.complete"]) { + handlers[type] = async (payload) => host.handleGoal(type, payload); + } + } + if (host.includeCompactSkill) { + for (const type of ["compact.run", "compact.status"]) { + handlers[type] = async (payload) => host.handleCompact(type, payload); + } + } + if (host.isRefineAllowed()) { + for (const type of ["refine.run", "refine.status"]) { + handlers[type] = async (payload) => host.handleRefine(type, payload); + } + } + if (host.hasHeartbeatController()) { + for (const type of [ + "rlm_heartbeat.list", + "rlm_heartbeat.create", + "rlm_heartbeat.update", + "rlm_heartbeat.delete", + ]) { + handlers[type] = async (payload) => host.handleHeartbeat(type, payload); + } + } + const visibleKernelSkillNames = new Set( + host + .getModelVisibleSkills() + .filter((skill) => !skill.disableModelInvocation) + .map((skill) => skill.name), + ); + const messageController = host.getAgentMessageController(); + if (messageController && visibleKernelSkillNames.has(AGENT_MESSAGE_SKILL_NAME)) { + Object.assign( + handlers, + createAgentMessageHostHandlers({ + family: async () => { + if (!messageController.family) throw new Error("agent family roster is not available in this session"); + return messageController.family(); + }, + awaitPendingChildPublication: (selector) => host.awaitChildPublication(selector), + sendAgentMessage: async (input) => { + const receipt = (await host.handleMessage("agent_message.send", { + target: input.target, + message: input.message, + })) as AgentSessionMessageReceipt; + if (host.getDepth() > 0) { + let addressedParent = input.receiverRole === "parent"; + if (input.receiverRole === undefined && messageController.family) { + try { + addressedParent = (await messageController.family()).some( + (member) => + member.relationship === "parent" && + (member.entry.id === input.target || + agentFamilyMemberName(member.entry) === input.target), + ); + } catch { + addressedParent = false; + } + } + if (addressedParent) { + host.recordParentReply(); + } + } + return receipt; + }, + }), + ); + } + if (host.hasObserveController()) { + Object.assign( + handlers, + createAgentObserveHostHandlers({ + listAgents: () => host.handleObserve("agent_observe.list") as AgentObserveListResult, + getAgent: (target) => + host.handleObserve("agent_observe.get", { + target, + }) as AgentObserveAgentSnapshot, + recentMessages: (input) => + host.handleObserve("agent_observe.recent", { + target: input.target, + limit: input.limit, + max_chars: input.maxChars, + }) as AgentObserveRecentMessagesResult, + }), + ); + } + const mcpManager = host.getMcpManager(); + if (mcpManager) { + Object.assign(handlers, mcpManager.hostHandlers()); + } + return handlers; +} diff --git a/packages/coding-agent/src/session/kernel.ts b/packages/coding-agent/src/session/kernel.ts new file mode 100644 index 0000000000..02473260be --- /dev/null +++ b/packages/coding-agent/src/session/kernel.ts @@ -0,0 +1,161 @@ +import { existsSync } from "node:fs"; +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import type { AssistantMessage } from "@earendil-works/pi-ai"; +import type { ToolDefinition } from "../core/extensions/index.js"; +import type { HostRequestHandlers, KernelSentAgentMessage } from "../core/kernel/index.js"; +import { type RestoreResult, snapshotPathIn } from "../core/kernel/state-snapshot.js"; +import { type CustomMessage, IPYTHON_STATE_RESTORED_CUSTOM_TYPE } from "../core/messages.js"; +import type { SessionManager } from "../core/session-manager.js"; +import type { PythonSkillRuntimeInfo } from "../core/skills.js"; +import { createAllToolDefinitions } from "../core/tools/index.js"; +import { IpythonKernelProvisioner } from "../core/tools/ipython.js"; + +const KERNEL_STATE_LISTING_TIMEOUT_MS = 5000; +export interface SessionKernelHost { + cwd: string; + getArtifactDir(): string | undefined; + getSessionId(): string; + getEnv(): Record; + getShellCommandPrefix(): string | undefined; + getShellPath(): string | undefined; + createHostHandlers(): HostRequestHandlers; + recordLateSentAgentMessage(toolCallId: string, message: KernelSentAgentMessage): void; + getMessages(): AgentMessage[]; + appendCustomMessageEntry: SessionManager["appendCustomMessageEntry"]; + emit(event: { type: "message_start" | "message_end"; message: AgentMessage }): void; + sendCustomMessage( + message: Pick, + options: { deliverAs: "nextTurn" }, + ): Promise; +} + +export class SessionKernel { + provisioner?: IpythonKernelProvisioner; + private snapshotDir?: string; + private built = false; + constructor( + private readonly host: SessionKernelHost, + private readonly prewarm: boolean, + ) {} + build(pythonSkills: PythonSkillRuntimeInfo[]): Record { + // Rebuilding (e.g. /reload) replaces the provisioner; drop the previous + // kernel so the session never holds two live kernels. Gate the new kernel's + // startup on the old one's dispose (which flushes a final snapshot), so a + // reload can't restore from a snapshot the old kernel is still writing. + const previousDispose = this.provisioner?.dispose(); + this.snapshotDir = this.host.getArtifactDir(); + // Only surface the "revived from your previous session" notice on the first + // build (a genuine resume). A later rebuild (/reload) restores state silently + // for continuity — the conversation is unchanged, so there's nothing to flag. + const notifyRestore = !this.built; + this.provisioner = new IpythonKernelProvisioner(this.host.cwd, { + env: this.host.getEnv(), + commandPrefix: this.host.getShellCommandPrefix(), + shellPath: this.host.getShellPath(), + sessionId: this.host.getSessionId(), + hostHandlers: this.host.createHostHandlers(), + pythonSkills, + snapshotDir: this.snapshotDir, + readyGate: previousDispose, + onRestore: notifyRestore ? (result) => this.onStateRestored(result) : undefined, + }); + return createAllToolDefinitions(this.host.cwd, { + ipython: { + provisioner: this.provisioner, + commandPrefix: this.host.getShellCommandPrefix(), + shellPath: this.host.getShellPath(), + onLateSentAgentMessage: (toolCallId, message) => this.host.recordLateSentAgentMessage(toolCallId, message), + }, + }); + } + finishBuild(activeToolNames: string[]): void { + const hasSnapshot = !!this.snapshotDir && existsSync(snapshotPathIn(this.snapshotDir)); + if ((this.prewarm || hasSnapshot) && activeToolNames.includes("ipython")) this.provisioner?.prewarm(); + this.built = true; + } + async dispose(snapshot: boolean, afterKernel?: () => Promise | undefined): Promise { + try { + await this.provisioner?.dispose({ snapshot }); + } catch { + /* Failed startup already cleaned up. */ + } + if (afterKernel) await afterKernel(); + } + async syncAfterCompaction(): Promise { + const provisioner = this.provisioner; + if (!provisioner?.hasRunningKernel) return; + const pruned = await provisioner.pruneOversizedVariables().catch(() => null); + const abort = new AbortController(); + const timer = setTimeout(() => abort.abort(), KERNEL_STATE_LISTING_TIMEOUT_MS); + if (typeof timer === "object" && "unref" in timer) timer.unref(); + let names: string[] | null; + try { + names = await provisioner.listNamespaceNames(abort.signal).catch(() => null); + } finally { + clearTimeout(timer); + } + if (names === null && !provisioner.hasRunningKernel) return; + const detail = + names === null + ? "" + : names.length > 0 + ? ` These names are still defined: ${names.join(", ")}.` + : " You have not defined any names yet."; + const prunedDetail = + pruned && pruned.length > 0 + ? ` Variables above the per-variable snapshot limit were removed: ${pruned.join(", ")}.` + : ""; + const content = [ + "[python-state]", + "", + `Your Python kernel persisted through compaction; its remaining variables, imports, and helpers are still available.${prunedDetail}${detail}`, + ].join("\n"); + const message = { + role: "custom" as const, + customType: "ipython_state", + content, + display: false, + timestamp: Date.now(), + } satisfies CustomMessage; + const messages = this.host.getMessages(); + const last = messages[messages.length - 1]; + const insertBeforeError = last?.role === "assistant" && (last as AssistantMessage).stopReason === "error"; + if (insertBeforeError) { + messages.splice(messages.length - 1, 0, message); + } else { + messages.push(message); + } + this.host.appendCustomMessageEntry(message.customType, message.content, message.display, undefined); + this.host.emit({ type: "message_start", message }); + this.host.emit({ type: "message_end", message }); + } + + onStateRestored(result: RestoreResult): void { + const lines = ["[python-state-restored]", ""]; + if (result.restored.length > 0) { + lines.push( + `Your Python kernel state was revived from your previous session. These names are available again: ${result.restored.join(", ")}.`, + ); + } else { + lines.push( + "Your previous Python kernel state could not be revived; the kernel is starting fresh, so re-create any variables, imports, or loaded data you need.", + ); + } + if (result.failed.length > 0) { + lines.push( + `These could not be restored and must be recreated if needed: ${result.failed.map((f) => f.name).join(", ")}.`, + ); + } + void this.host + .sendCustomMessage( + { + customType: IPYTHON_STATE_RESTORED_CUSTOM_TYPE, + content: lines.join("\n"), + display: true, + details: { restored: result.restored.length > 0 }, + }, + { deliverAs: "nextTurn" }, + ) + .catch(() => {}); + } +} diff --git a/packages/coding-agent/src/session/tools.ts b/packages/coding-agent/src/session/tools.ts new file mode 100644 index 0000000000..540868e51d --- /dev/null +++ b/packages/coding-agent/src/session/tools.ts @@ -0,0 +1,406 @@ +import type { AgentTool } from "@earendil-works/pi-agent-core"; +import { + type ExtensionRunner, + type ToolDefinition, + type ToolInfo, + wrapRegisteredTools, +} from "../core/extensions/index.js"; +import type { AcpMcpServerConfig } from "../core/mcp/acp-mcp-types.js"; +import type { McpManager } from "../core/mcp/mcp-manager.js"; +import type { ResourceLoader } from "../core/resource-loader.js"; +import type { Skill } from "../core/skills.js"; +import { createSyntheticSourceInfo, type SourceInfo } from "../core/source-info.js"; +import { type BuildSystemPromptOptions, buildSystemPrompt } from "../core/system-prompt.js"; +import { acpMcpToolNames, createAcpMcpToolDefinitions } from "../core/tools/acp-mcp.js"; +import type { IpythonKernelProvisioner } from "../core/tools/ipython.js"; +import { createToolDefinitionFromAgentTool } from "../core/tools/tool-definition-wrapper.js"; + +interface ToolDefinitionEntry { + definition: ToolDefinition; + sourceInfo: SourceInfo; +} +export interface SessionToolsHost { + cwd: string; + resourceLoader: Pick; + getExtensionRunner(): ExtensionRunner; + getSessionFile(): string | undefined; + getModelVisibleSkills(): Skill[]; + getDepth(): number; + getMaxDepth(): number; + getParentAgent(): string | undefined; + getMcpManager(): + | Pick< + McpManager, + "getAcpServers" | "getEnabledPersistentGenericServers" | "replaceAcpServers" | "canReleaseAcpServers" + > + | undefined; + getProvisioner(): IpythonKernelProvisioner | undefined; + getActiveToolNames(): string[]; + setActiveToolsByName(names: string[]): void; + getActiveTools(): AgentTool[]; + setActiveTools(tools: AgentTool[]): void; + setSystemPrompt(prompt: string): void; + isStreaming(): boolean; + rebuildRuntime(options: { activeToolNames?: string[]; includeAllExtensionTools?: boolean }): void; + acquireInputPause(): { release(): void }; + waitForAgentIdle(): Promise; + getEventQueue(): Promise; +} + +export class SessionTools { + private _toolRegistry = new Map(); + private _toolDefinitions = new Map(); + private _toolPromptSnippets = new Map(); + private _toolPromptGuidelines = new Map(); + private _baseToolDefinitions = new Map(); + private _acpMcpTools: ToolDefinition[] = []; + readonly customTools: ToolDefinition[]; + private readonly baseToolsOverride?: Record; + readonly allowedToolNames?: Set; + baseSystemPrompt = ""; + baseSystemPromptOptions!: BuildSystemPromptOptions; + constructor( + private readonly host: SessionToolsHost, + config: { + customTools?: ToolDefinition[]; + allowedToolNames?: string[]; + baseToolsOverride?: Record; + }, + ) { + this.customTools = config.customTools ?? []; + this.baseToolsOverride = config.baseToolsOverride; + this.allowedToolNames = config.allowedToolNames ? new Set(config.allowedToolNames) : undefined; + } + get registry(): ReadonlyMap { + return this._toolRegistry; + } + get defaultActiveToolNames(): string[] { + return this.baseToolsOverride ? Object.keys(this.baseToolsOverride) : ["ipython"]; + } + buildBaseOverrides(): Record | undefined { + return this.baseToolsOverride + ? Object.fromEntries( + Object.entries(this.baseToolsOverride).map(([name, tool]) => [ + name, + createToolDefinitionFromAgentTool(tool), + ]), + ) + : undefined; + } + setBaseDefinitions(definitions: Record): void { + this._baseToolDefinitions = new Map(Object.entries(definitions)); + } + updateAcpDefinitions(): void { + const previousAcpMcpToolNames = new Set(this._acpMcpTools.map((tool) => tool.name)); + const acpServers = this.host.getMcpManager()?.getAcpServers() ?? []; + const provisioner = this.host.getProvisioner(); + if (acpServers.length > 0 && !provisioner) throw new Error("ACP MCP servers require the built-in cpython tool"); + const acpMcpTools = provisioner ? createAcpMcpToolDefinitions(acpServers, provisioner) : []; + this._assertAcpMcpToolNamesAvailable(acpMcpTools.map((tool) => tool.name)); + for (const name of previousAcpMcpToolNames) this.allowedToolNames?.delete(name); + for (const tool of acpMcpTools) this.allowedToolNames?.add(tool.name); + this._acpMcpTools = acpMcpTools; + } + getActiveToolNames(): string[] { + return this.host.getActiveTools().map((t) => t.name); + } + + getAllTools(): ToolInfo[] { + return Array.from(this._toolDefinitions.values()).map(({ definition, sourceInfo }) => ({ + name: definition.name, + description: definition.description, + parameters: definition.parameters, + sourceInfo, + })); + } + + getToolDefinition(name: string): ToolDefinition | undefined { + return this._toolDefinitions.get(name)?.definition; + } + + setActiveToolsByName(toolNames: string[]): void { + const tools: AgentTool[] = []; + const validToolNames: string[] = []; + const seenToolNames = new Set(); + for (const name of toolNames) { + if (seenToolNames.has(name)) { + continue; + } + const tool = this._toolRegistry.get(name); + if (tool) { + seenToolNames.add(name); + tools.push(tool); + validToolNames.push(name); + } + } + this.host.setActiveTools(tools); + + this.baseSystemPrompt = this.rebuildSystemPrompt(validToolNames); + this.host.setSystemPrompt(this.baseSystemPrompt); + } + + private _normalizePromptSnippet(text: string | undefined): string | undefined { + if (!text) return undefined; + const oneLine = text + .replace(/[\r\n]+/g, " ") + .replace(/\s+/g, " ") + .trim(); + return oneLine.length > 0 ? oneLine : undefined; + } + + private _normalizePromptGuidelines(guidelines: string[] | undefined): string[] { + if (!guidelines || guidelines.length === 0) { + return []; + } + + const unique = new Set(); + for (const guideline of guidelines) { + const normalized = guideline.trim(); + if (normalized.length > 0) { + unique.add(normalized); + } + } + return Array.from(unique); + } + + rebuildSystemPrompt(toolNames: string[]): string { + const validToolNames = toolNames.filter((name) => this._toolRegistry.has(name)); + const toolSnippets: Record = {}; + const promptGuidelines: string[] = []; + for (const name of validToolNames) { + const snippet = this._toolPromptSnippets.get(name); + if (snippet) { + toolSnippets[name] = snippet; + } + + const toolGuidelines = this._toolPromptGuidelines.get(name); + if (toolGuidelines) { + promptGuidelines.push(...toolGuidelines); + } + } + + const loaderSystemPrompt = this.host.resourceLoader.getSystemPrompt(); + const loaderAppendSystemPrompt = this.host.resourceLoader.getAppendSystemPrompt(); + const appendSystemPrompt = + loaderAppendSystemPrompt.length > 0 ? loaderAppendSystemPrompt.join("\n\n") : undefined; + const loadedSkills = this.host.getModelVisibleSkills(); + const loadedContextFiles = this.host.resourceLoader.getAgentsFiles().agentsFiles; + + this.baseSystemPromptOptions = { + cwd: this.host.cwd, + skills: loadedSkills, + contextFiles: loadedContextFiles, + customPrompt: loaderSystemPrompt, + appendSystemPrompt, + messagesPath: this.host.getSessionFile(), + selectedTools: validToolNames, + toolSnippets, + promptGuidelines, + allowRecursion: this.host.getDepth() < this.host.getMaxDepth(), + rlmDepth: this.host.getDepth(), + rlmParentAgent: this.host.getParentAgent(), + genericMcpServers: this.host.getMcpManager()?.getEnabledPersistentGenericServers(), + }; + return buildSystemPrompt(this.baseSystemPromptOptions); + } + + refreshExtensionSystemPrompt(extensionPrompt: string, baseSnapshot: string): string { + if (this.baseSystemPrompt === baseSnapshot) { + return extensionPrompt; + } + if (!extensionPrompt.includes(baseSnapshot)) { + return extensionPrompt; + } + return extensionPrompt.replace(baseSnapshot, () => this.baseSystemPrompt); + } + + refreshToolRegistry(options?: { activeToolNames?: string[]; includeAllExtensionTools?: boolean }): void { + const previousRegistryNames = new Set(this._toolRegistry.keys()); + const previousActiveToolNames = this.host.getActiveToolNames(); + const allowedToolNames = this.allowedToolNames; + const registeredTools = this.host.getExtensionRunner().getAllRegisteredTools(); + const sdkToolEntry = (definition: ToolDefinition) => ({ + definition, + sourceInfo: createSyntheticSourceInfo(``, { + source: "sdk" as const, + }), + }); + const allCustomTools = [ + ...registeredTools, + ...this.customTools.map(sdkToolEntry), + ...this._acpMcpTools.map(sdkToolEntry), + ]; + const isAllowedTool = (name: string): boolean => !allowedToolNames || allowedToolNames.has(name); + const allowedCustomTools = allCustomTools.filter((tool) => isAllowedTool(tool.definition.name)); + const definitionRegistry = new Map( + Array.from(this._baseToolDefinitions.entries()) + .filter(([name]) => isAllowedTool(name)) + .map(([name, definition]) => [ + name, + { + definition, + sourceInfo: createSyntheticSourceInfo(``, { + source: "builtin", + }), + }, + ]), + ); + for (const tool of allowedCustomTools) { + definitionRegistry.set(tool.definition.name, { + definition: tool.definition, + sourceInfo: tool.sourceInfo, + }); + } + this._toolDefinitions = definitionRegistry; + this._toolPromptSnippets = new Map( + Array.from(definitionRegistry.values()) + .map(({ definition }) => { + const snippet = this._normalizePromptSnippet(definition.promptSnippet); + return snippet ? ([definition.name, snippet] as const) : undefined; + }) + .filter((entry): entry is readonly [string, string] => entry !== undefined), + ); + this._toolPromptGuidelines = new Map( + Array.from(definitionRegistry.values()) + .map(({ definition }) => { + const guidelines = this._normalizePromptGuidelines(definition.promptGuidelines); + return guidelines.length > 0 ? ([definition.name, guidelines] as const) : undefined; + }) + .filter((entry): entry is readonly [string, string[]] => entry !== undefined), + ); + const runner = this.host.getExtensionRunner(); + const wrappedExtensionTools = wrapRegisteredTools(allowedCustomTools, runner); + // Resolve the runner at call time so a rebuild/reload rebinds built-in tools to the + // live runner instead of wedging them on the invalidated one's stale-ctx guard. + const wrappedBuiltInTools = wrapRegisteredTools( + Array.from(this._baseToolDefinitions.values()) + .filter((definition) => isAllowedTool(definition.name)) + .map((definition) => ({ + definition, + sourceInfo: createSyntheticSourceInfo(``, { source: "builtin" }), + })), + () => this.host.getExtensionRunner(), + ); + + const toolRegistry = new Map(wrappedBuiltInTools.map((tool) => [tool.name, tool])); + for (const tool of wrappedExtensionTools as AgentTool[]) { + toolRegistry.set(tool.name, tool); + } + this._toolRegistry = toolRegistry; + + const nextActiveToolNames = ( + options?.activeToolNames ? [...options.activeToolNames] : [...previousActiveToolNames] + ).filter((name) => isAllowedTool(name)); + + if (allowedToolNames) { + for (const toolName of this._toolRegistry.keys()) { + if (allowedToolNames.has(toolName)) { + nextActiveToolNames.push(toolName); + } + } + } else if (options?.includeAllExtensionTools) { + for (const tool of wrappedExtensionTools) { + nextActiveToolNames.push(tool.name); + } + } else if (!options?.activeToolNames) { + for (const toolName of this._toolRegistry.keys()) { + if (!previousRegistryNames.has(toolName)) { + nextActiveToolNames.push(toolName); + } + } + } + + this.host.setActiveToolsByName([...new Set(nextActiveToolNames)]); + } + + replaceAcpMcpServers(servers: readonly AcpMcpServerConfig[], ownerId: string): void { + if (this.host.isStreaming()) throw new Error("Cannot replace ACP MCP servers while the agent is running"); + const mcpManager = this.host.getMcpManager(); + if (!mcpManager) { + if (servers.length > 0) throw new Error("MCP is unavailable in this session"); + return; + } + if (servers.length > 0 && !this.host.getProvisioner()) { + throw new Error("ACP MCP servers require the built-in cpython tool"); + } + this._assertAcpMcpToolNamesAvailable(acpMcpToolNames(servers)); + if (!mcpManager.replaceAcpServers(servers, ownerId)) return; + this._rebuildRuntimeForAcpMcpServers(); + } + + async releaseAcpMcpServers(ownerId: string, serverNames: readonly string[]): Promise { + const mcpManager = this.host.getMcpManager(); + if (!mcpManager?.canReleaseAcpServers(ownerId)) return; + if (mcpManager.replaceAcpServers([], ownerId)) { + const removedToolNames = new Set(this._acpMcpTools.map((tool) => tool.name)); + const activeToolNames = this.host.getActiveToolNames().filter((name) => !removedToolNames.has(name)); + for (const name of removedToolNames) this.allowedToolNames?.delete(name); + this._acpMcpTools = []; + this.refreshToolRegistry({ activeToolNames, includeAllExtensionTools: true }); + this.baseSystemPrompt = this.rebuildSystemPrompt(this.host.getActiveToolNames()); + this.host.setSystemPrompt(this.baseSystemPrompt); + } + const names = [...new Set(serverNames)]; + if (names.length === 0) return; + + const inputPause = this.host.acquireInputPause(); + try { + // Do not rebuild or kill the notebook. Wait for the current turn, then ask + // the kernel-owned MCP registry to close only these cached transports. + await this.host.waitForAgentIdle(); + await this.host.getEventQueue(); + const manager = this.host.getProvisioner()?.manager; + if (!manager?.isRunning) return; + const code = [ + "import importlib as _prime_importlib", + '_prime_mcp = _prime_importlib.import_module("rlm.mcp")', + `_prime_mcp_names = ${JSON.stringify(names)}`, + "_prime_mcp_errors = []", + "for _prime_mcp_name in _prime_mcp_names:", + " try:", + " await _prime_mcp.reload(_prime_mcp_name)", + " except BaseException as _prime_mcp_error:", + " _prime_mcp_errors.append(_prime_mcp_error)", + "if _prime_mcp_errors:", + " raise _prime_mcp_errors[0]", + "del _prime_mcp, _prime_importlib, _prime_mcp_names, _prime_mcp_errors, _prime_mcp_name", + ].join("\n"); + const result = await manager.execute(code); + if (result.status !== "ok") { + throw new Error(`Failed to close ACP MCP kernel transports: ${result.stderr || "kernel error"}`); + } + } finally { + inputPause.release(); + } + } + + private _assertAcpMcpToolNamesAvailable(names: readonly string[]): void { + const occupiedNames = new Set([ + ...this._baseToolDefinitions.keys(), + ...this.customTools.map((tool) => tool.name), + ...this.host + .getExtensionRunner() + .getAllRegisteredTools() + .map((tool) => tool.definition.name), + ]); + for (const name of names) { + if (occupiedNames.has(name)) { + throw new Error(`ACP MCP tool name conflicts with an existing tool: ${name}`); + } + } + } + + private _rebuildRuntimeForAcpMcpServers(): void { + const previousToolNames = new Set(this._acpMcpTools.map((tool) => tool.name)); + const nextToolNames = acpMcpToolNames(this.host.getMcpManager()?.getAcpServers() ?? []); + this._assertAcpMcpToolNamesAvailable(nextToolNames); + const activeToolNames = this.host.getActiveToolNames().filter((name) => !previousToolNames.has(name)); + activeToolNames.push(...nextToolNames); + this.host.rebuildRuntime({ + activeToolNames, + includeAllExtensionTools: true, + }); + this.baseSystemPrompt = this.rebuildSystemPrompt(this.host.getActiveToolNames()); + this.host.setSystemPrompt(this.baseSystemPrompt); + } +} diff --git a/packages/coding-agent/test/agent-session-concurrent.test.ts b/packages/coding-agent/test/agent-session-concurrent.test.ts index c4bc2bd5e6..a9791ee829 100644 --- a/packages/coding-agent/test/agent-session-concurrent.test.ts +++ b/packages/coding-agent/test/agent-session-concurrent.test.ts @@ -20,7 +20,7 @@ import type { AgentCronJob } from "../src/core/cron-jobs.js"; import { ModelRegistry } from "../src/core/model-registry.js"; import { SessionManager } from "../src/core/session-manager.js"; import { SettingsManager } from "../src/core/settings-manager.js"; -import type { BuildSystemPromptOptions } from "../src/core/system-prompt.js"; +import { IpythonKernelProvisioner } from "../src/core/tools/ipython.js"; import type { SessionRefinement } from "../src/session/refinement.js"; import { createTestExtensionsResult, createTestResourceLoader } from "./utilities.js"; @@ -67,6 +67,7 @@ describe("AgentSession concurrent prompt guard", () => { }); afterEach(async () => { + vi.restoreAllMocks(); delete (globalThis as typeof globalThis & { testExtensionApi?: unknown }).testExtensionApi; delete (globalThis as typeof globalThis & { testCommandRuns?: unknown }).testCommandRuns; if (session) { @@ -191,11 +192,11 @@ describe("AgentSession concurrent prompt guard", () => { it("forwards kernelSnapshot: false to the kernel provisioner during disposal", async () => { createSession(); - const dispose = vi.fn(async () => {}); - Reflect.set(session, "_ipythonKernelProvisioner", { dispose }); + const dispose = vi.spyOn(IpythonKernelProvisioner.prototype, "dispose").mockResolvedValue(); await session.disposeAsync({ kernelSnapshot: false }); expect(dispose).toHaveBeenCalledWith({ snapshot: false }); + dispose.mockRestore(); }); it("should throw when prompt() called while streaming", async () => { @@ -919,43 +920,16 @@ describe("AgentSession concurrent prompt guard", () => { }); const snapshots: string[][] = []; - const sessionWithRunner = session as unknown as { - _extensionRunner?: { - hasHandlers: (eventType: string) => boolean; - emit: (event: { type: string; message?: { role?: string } }) => Promise; - emitMessageEnd: (event: { type: string; message?: { role?: string } }) => Promise; - emitToolCall: (event: { type: string; toolCallId: string }) => Promise; - emitInput: ( - text: string, - images: unknown, - source: "interactive" | "rpc" | "extension", - ) => Promise<{ action: "continue" }>; - emitBeforeAgentStart: ( - prompt: string, - images: unknown, - systemPrompt: string, - systemPromptOptions: BuildSystemPromptOptions, - ) => Promise; - invalidate: (message?: string) => void; - }; - }; - sessionWithRunner._extensionRunner = { - hasHandlers: (eventType) => eventType === "tool_call", - emit: async () => {}, - emitMessageEnd: async () => undefined, - emitToolCall: async () => { - snapshots.push( - sessionManager - .getEntries() - .filter((entry) => entry.type === "message") - .map((entry) => entry.message.role), - ); - return undefined; - }, - emitInput: async () => ({ action: "continue" }), - emitBeforeAgentStart: async () => undefined, - invalidate: () => {}, - }; + vi.spyOn(session.extensionRunner, "hasHandlers").mockImplementation((eventType) => eventType === "tool_call"); + vi.spyOn(session.extensionRunner, "emitToolCall").mockImplementation(async () => { + snapshots.push( + sessionManager + .getEntries() + .filter((entry) => entry.type === "message") + .map((entry) => entry.message.role), + ); + return undefined; + }); await session.prompt("hi"); await session.agent.waitForIdle(); @@ -1064,38 +1038,10 @@ describe("AgentSession concurrent prompt guard", () => { baseToolsOverride: { dummy: tool }, }); - const sessionWithRunner = session as unknown as { - _extensionRunner?: { - hasHandlers: (eventType: string) => boolean; - emit: (event: { type: string; message?: { role?: string } }) => Promise; - emitMessageEnd: (event: { type: string; message?: { role?: string } }) => Promise; - emitInput: ( - text: string, - images: unknown, - source: "interactive" | "rpc" | "extension", - ) => Promise<{ action: "continue" }>; - emitBeforeAgentStart: ( - prompt: string, - images: unknown, - systemPrompt: string, - systemPromptOptions: BuildSystemPromptOptions, - ) => Promise; - invalidate: (message?: string) => void; - }; - }; - sessionWithRunner._extensionRunner = { - hasHandlers: () => false, - emit: async () => {}, - emitMessageEnd: async (event) => { - if (event.type === "message_end" && event.message?.role === "assistant") { - await new Promise((resolve) => setTimeout(resolve, 40)); - } - return undefined; - }, - emitInput: async () => ({ action: "continue" }), - emitBeforeAgentStart: async () => undefined, - invalidate: () => {}, - }; + vi.spyOn(session.extensionRunner, "emitMessageEnd").mockImplementation(async (event) => { + if (event.message.role === "assistant") await new Promise((resolve) => setTimeout(resolve, 40)); + return undefined; + }); await session.prompt("hi"); await session.agent.waitForIdle(); diff --git a/packages/coding-agent/test/agent-session-dynamic-tools.test.ts b/packages/coding-agent/test/agent-session-dynamic-tools.test.ts index 5185c8cc54..1bddd808a5 100644 --- a/packages/coding-agent/test/agent-session-dynamic-tools.test.ts +++ b/packages/coding-agent/test/agent-session-dynamic-tools.test.ts @@ -3,7 +3,7 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import { getModel } from "@earendil-works/pi-ai"; import { Type } from "typebox"; -import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { DefaultResourceLoader } from "../src/core/resource-loader.js"; import { createAgentSession } from "../src/core/sdk.js"; import { SessionManager } from "../src/core/session-manager.js"; @@ -65,7 +65,13 @@ describe("AgentSession dynamic tool registration", () => { expect(session.getAllTools().map((tool) => tool.name)).not.toContain("dynamic_tool"); + const getActiveTools = vi.spyOn(session, "getActiveToolNames"); + const setActiveTools = vi.spyOn(session, "setActiveToolsByName"); await session.bindExtensions({}); + expect(getActiveTools).toHaveBeenCalled(); + expect(setActiveTools).toHaveBeenCalled(); + expect(getActiveTools.mock.contexts.every((receiver) => receiver === session)).toBe(true); + expect(setActiveTools.mock.contexts.every((receiver) => receiver === session)).toBe(true); const allTools = session.getAllTools(); const dynamicTool = allTools.find((tool) => tool.name === "dynamic_tool"); diff --git a/packages/coding-agent/test/agent-session-recursion.test.ts b/packages/coding-agent/test/agent-session-recursion.test.ts index 039b3ab861..9883dbe158 100644 --- a/packages/coding-agent/test/agent-session-recursion.test.ts +++ b/packages/coding-agent/test/agent-session-recursion.test.ts @@ -112,8 +112,7 @@ interface InspectableRlmRun { session?: AgentSession; } -interface InspectableRlmSession { - _disposing: boolean; +interface InspectableRlmChildren { _activeRlmChildRuns: Map; _unsettledRlmChildRuns: Set; _deletingRlmChildren: Map< @@ -128,6 +127,12 @@ interface InspectableRlmSession { _rlmChildUnsubscribes: Map void>; _deletedRlmChildIds: Set; _rlmQuiescenceWaitAborts: Set; +} + +interface InspectableRlmSession { + _disposing: boolean; + _children: InspectableRlmChildren; + _childState: { recordReply(): void }; _createKernelHostHandlers(): HostRequestHandlers; _reapDeletedRlmSubagentRuntimesAfterCompaction(): Promise; } @@ -689,13 +694,13 @@ describe("AgentSession rlm recursion", () => { await expect(root.deleteRlmSubagent("retained-retry-worker")).rejects.toThrow("retained close failed"); const internals = root as unknown as InspectableRlmSession; - expect(internals._rlmChildCleanupFailures.size).toBe(1); + expect(internals._children._rlmChildCleanupFailures.size).toBe(1); await root.compact(); expect(deleteRuntime).toHaveBeenCalledTimes(2); - expect(internals._rlmChildCleanupFailures.size).toBe(0); - expect(internals._rlmChildSessions.size).toBe(0); + expect(internals._children._rlmChildCleanupFailures.size).toBe(0); + expect(internals._children._rlmChildSessions.size).toBe(0); await expect(root.runRlmChild("replacement", { name: "retained-retry-worker" })).resolves.toMatchObject({ name: "retained-retry-worker", }); @@ -716,7 +721,7 @@ describe("AgentSession rlm recursion", () => { const parentInternals = parent as unknown as InspectableRlmSession; parentInternalsToClear.push(parentInternals); const nestedId = `${id}-live-grandchild`; - parentInternals._activeRlmChildRuns.set(nestedId, { + parentInternals._children._activeRlmChildRuns.set(nestedId, { id: nestedId, prompt: "still working", sessionName: nestedId, @@ -728,7 +733,7 @@ describe("AgentSession rlm recursion", () => { }); if (hiding === "detached") { - rootInternals._activeRlmChildRuns.set(id, { + rootInternals._children._activeRlmChildRuns.set(id, { id, prompt: "hidden parent", sessionName: id, @@ -749,13 +754,13 @@ describe("AgentSession rlm recursion", () => { }); // The same child can be visible in both lifecycle registries while // deletion settles; it must be traversed exactly once and remain hidden. - rootInternals._rlmChildSessions.set(id, { session: parent }); + rootInternals._children._rlmChildSessions.set(id, { session: parent }); } else { - rootInternals._rlmChildSessions.set(id, { session: parent }); + rootInternals._children._rlmChildSessions.set(id, { session: parent }); if (hiding === "deleted") { - rootInternals._deletedRlmChildIds.add(id); + rootInternals._children._deletedRlmChildIds.add(id); } else { - rootInternals._deletingRlmChildren.set(id, { + rootInternals._children._deletingRlmChildren.set(id, { subagent: { rlm_child_id: id, active_session_id: null, @@ -786,9 +791,9 @@ describe("AgentSession rlm recursion", () => { expect(snapshots.map((snapshot) => snapshot.status).sort()).toEqual(["queued", "running", "running"]); // These are deliberately minimal lifecycle records; remove them before // fixture teardown asks real runs to settle. - rootInternals._activeRlmChildRuns.clear(); - rootInternals._rlmChildSessions.clear(); - for (const parentInternals of parentInternalsToClear) parentInternals._activeRlmChildRuns.clear(); + rootInternals._children._activeRlmChildRuns.clear(); + rootInternals._children._rlmChildSessions.clear(); + for (const parentInternals of parentInternalsToClear) parentInternals._children._activeRlmChildRuns.clear(); root.dispose(); }); @@ -1108,7 +1113,7 @@ describe("AgentSession rlm recursion", () => { promptInjectedMessage; await root.runRlmChild("failing task", { name: "shared-child" }); await waitFor(() => - [...(root as unknown as InspectableRlmSession)._activeRlmChildRuns.values()].some( + [...(root as unknown as InspectableRlmSession)._children._activeRlmChildRuns.values()].some( (run) => run.status === "error", ), ); @@ -1160,7 +1165,7 @@ describe("AgentSession rlm recursion", () => { 'No child matches "deleted-child"', ); releaseRuntimeCreation(); - await waitFor(() => (root as unknown as InspectableRlmSession)._activeRlmChildRuns.size === 0); + await waitFor(() => (root as unknown as InspectableRlmSession)._children._activeRlmChildRuns.size === 0); }); it("marks a broadcast delivery to the parent as replied without reloading the roster", async () => { @@ -1201,7 +1206,7 @@ describe("AgentSession rlm recursion", () => { parent.setSessionName("parent"); const child = createSession({ depth: 1 }); child.setSessionName("worker"); - (child as unknown as { _repliedToParentSinceTask: boolean })._repliedToParentSinceTask = true; + (child as unknown as InspectableRlmSession)._childState.recordReply(); const daemon = new AgentDaemon(join(tempDir, "daemon.sock"), { defaultSessionConfig: { agentDir: tempDir, cwd: tempDir }, @@ -1270,7 +1275,7 @@ describe("AgentSession rlm recursion", () => { it("resets replied state when a parent message is accepted", async () => { const child = createSession({ depth: 1 }); - (child as unknown as { _repliedToParentSinceTask: boolean })._repliedToParentSinceTask = true; + (child as unknown as InspectableRlmSession)._childState.recordReply(); const message = createAgentSessionMessage({ id: "agentmsg-parent-task", source: "agent_message", @@ -1286,7 +1291,7 @@ describe("AgentSession rlm recursion", () => { it("resets replied state when a parent follow-up is queued", async () => { const child = createSession({ depth: 1 }); - (child as unknown as { _repliedToParentSinceTask: boolean })._repliedToParentSinceTask = true; + (child as unknown as InspectableRlmSession)._childState.recordReply(); const message = createAgentSessionMessage({ id: "agentmsg-parent-follow-up", source: "agent_message", @@ -1495,7 +1500,9 @@ describe("AgentSession rlm recursion", () => { await waitFor(() => runtimeCreationStarted); await root.deleteRlmSubagent(spawned.rlm_child_id); releaseRuntimeCreation(); - await waitFor(() => !(root as unknown as InspectableRlmSession)._activeRlmChildRuns.has(spawned.rlm_child_id)); + await waitFor( + () => !(root as unknown as InspectableRlmSession)._children._activeRlmChildRuns.has(spawned.rlm_child_id), + ); expect( root.messages.filter( (message) => message.role === "custom" && message.customType === "rlm_child_terminal_notice", @@ -1518,14 +1525,16 @@ describe("AgentSession rlm recursion", () => { }); const spawned = await root.runRlmChild("start failing child", { name: "reusable-worker" }); const internals = root as unknown as InspectableRlmSession; - await vi.waitFor(() => expect(internals._activeRlmChildRuns.get(spawned.rlm_child_id)?.settled).toBe(true)); + await vi.waitFor(() => + expect(internals._children._activeRlmChildRuns.get(spawned.rlm_child_id)?.settled).toBe(true), + ); await expect(root.deleteRlmSubagent(spawned.rlm_child_id)).resolves.toMatchObject({ subagent: { rlm_child_id: spawned.rlm_child_id, session_name: "reusable-worker" }, }); expect(await root.listRlmSubagents()).toEqual({ subagents: [] }); - expect(internals._activeRlmChildRuns.has(spawned.rlm_child_id)).toBe(false); + expect(internals._children._activeRlmChildRuns.has(spawned.rlm_child_id)).toBe(false); await expect(root.runRlmChild("replacement child", { name: "reusable-worker" })).resolves.toMatchObject({ name: "reusable-worker", }); @@ -1670,11 +1679,13 @@ describe("AgentSession rlm recursion", () => { expect(root.registerRlmChildSession("cancelled-bash-active-child", child)).toBe(true); const quiescence = root.waitForRlmQuiescence(); - await vi.waitFor(() => expect((child as unknown as InspectableRlmSession)._rlmQuiescenceWaitAborts.size).toBe(1)); + await vi.waitFor(() => + expect((child as unknown as InspectableRlmSession)._children._rlmQuiescenceWaitAborts.size).toBe(1), + ); root.requestAbort(); await expect(quiescence).rejects.toThrow("RLM quiescence wait cancelled"); expect(child.isBashRunning).toBe(true); - expect((child as unknown as InspectableRlmSession)._rlmQuiescenceWaitAborts.size).toBe(0); + expect((child as unknown as InspectableRlmSession)._children._rlmQuiescenceWaitAborts.size).toBe(0); bashCompletion.resolve(); await bash; @@ -1712,13 +1723,13 @@ describe("AgentSession rlm recursion", () => { const quiescence = root.waitForRlmQuiescence(); await vi.waitFor(() => { - expect((childA as unknown as InspectableRlmSession)._rlmQuiescenceWaitAborts.size).toBe(1); - expect((childB as unknown as InspectableRlmSession)._rlmQuiescenceWaitAborts.size).toBe(1); + expect((childA as unknown as InspectableRlmSession)._children._rlmQuiescenceWaitAborts.size).toBe(1); + expect((childB as unknown as InspectableRlmSession)._children._rlmQuiescenceWaitAborts.size).toBe(1); }); childA.requestAbort(); await expect(quiescence).rejects.toThrow("RLM quiescence wait cancelled"); await vi.waitFor(() => - expect((childB as unknown as InspectableRlmSession)._rlmQuiescenceWaitAborts.size).toBe(0), + expect((childB as unknown as InspectableRlmSession)._children._rlmQuiescenceWaitAborts.size).toBe(0), ); expect(childB.isBashRunning).toBe(true); @@ -1786,7 +1797,7 @@ describe("AgentSession rlm recursion", () => { await vi.waitFor(() => expect(deferredNotices()).toHaveLength(1)); expect(synthesizedAgentMessageSend).not.toHaveBeenCalled(); const restartSnapshot = root.getPendingNextTurnMessageSnapshots(); - await vi.waitFor(() => expect(internals._unsettledRlmChildRuns.size).toBe(0)); + await vi.waitFor(() => expect(internals._children._unsettledRlmChildRuns.size).toBe(0)); expect(root.unfinishedActionCount).toBe(0); const closeIdleBoundary = await Promise.race([ root.waitForIdle().then(() => "idle" as const), @@ -1931,7 +1942,9 @@ describe("AgentSession rlm recursion", () => { const updatePause = root.acquireQueuedWorkPause(); childCompletion.resolve(); const internals = root as unknown as InspectableRlmSession; - await vi.waitFor(() => expect(internals._activeRlmChildRuns.get(spawned.rlm_child_id)?.status).toBe("done")); + await vi.waitFor(() => + expect(internals._children._activeRlmChildRuns.get(spawned.rlm_child_id)?.status).toBe("done"), + ); await root.waitForSessionInputCheckpoint(); const pendingSnapshot = root.getPendingNextTurnMessageSnapshots(); const actionSnapshot = root.getSessionActionRecoverySnapshot(); @@ -1943,8 +1956,8 @@ describe("AgentSession rlm recursion", () => { action.payload.customMessage?.customType === "rlm_child_terminal_notice", ), ).toHaveLength(0); - expect(internals._unsettledRlmChildRuns.size).toBe(1); - expect(internals._activeRlmChildRuns.get(spawned.rlm_child_id)?.settled).toBe(false); + expect(internals._children._unsettledRlmChildRuns.size).toBe(1); + expect(internals._children._activeRlmChildRuns.get(spawned.rlm_child_id)?.settled).toBe(false); updatePause.release(); await root.waitForRlmQuiescence(); @@ -1977,13 +1990,15 @@ describe("AgentSession rlm recursion", () => { const spawned = await root.runRlmChild("fail after startup", { name: "settled-error-worker" }); const internals = root as unknown as InspectableRlmSession; - await vi.waitFor(() => expect(internals._activeRlmChildRuns.get(spawned.rlm_child_id)?.settled).toBe(true)); - const run = internals._activeRlmChildRuns.get(spawned.rlm_child_id); + await vi.waitFor(() => + expect(internals._children._activeRlmChildRuns.get(spawned.rlm_child_id)?.settled).toBe(true), + ); + const run = internals._children._activeRlmChildRuns.get(spawned.rlm_child_id); if (!run) throw new Error("Missing settled error run"); expect(run.status).toBe("error"); expect(run.session).toBe(child); - expect(internals._rlmChildSessions.has(spawned.rlm_child_id)).toBe(false); - expect(internals._unsettledRlmChildRuns.has(run)).toBe(false); + expect(internals._children._rlmChildSessions.has(spawned.rlm_child_id)).toBe(false); + expect(internals._children._unsettledRlmChildRuns.has(run)).toBe(false); await expect(root.deleteRlmSubagent("settled-error-worker")).resolves.toMatchObject({ subagent: { rlm_child_id: spawned.rlm_child_id }, @@ -1993,7 +2008,7 @@ describe("AgentSession rlm recursion", () => { }); expect(deleteRlmSubagentRuntime).toHaveBeenCalledOnce(); expect(run.settled).toBe(false); - expect(internals._unsettledRlmChildRuns.has(run)).toBe(true); + expect(internals._children._unsettledRlmChildRuns.has(run)).toBe(true); let quiesced = false; const quiescence = root.waitForRlmQuiescence().then(() => { quiesced = true; @@ -2002,11 +2017,11 @@ describe("AgentSession rlm recursion", () => { expect(quiesced).toBe(false); firstCleanup.reject(new Error("first cleanup failed")); - await waitFor(() => internals._rlmChildCleanupFailures.has(spawned.rlm_child_id)); + await waitFor(() => internals._children._rlmChildCleanupFailures.has(spawned.rlm_child_id)); await sleep(20); expect(quiesced).toBe(false); - expect(internals._activeRlmChildRuns.has(spawned.rlm_child_id)).toBe(true); - expect(internals._unsettledRlmChildRuns.has(run)).toBe(true); + expect(internals._children._activeRlmChildRuns.has(spawned.rlm_child_id)).toBe(true); + expect(internals._children._unsettledRlmChildRuns.has(run)).toBe(true); await expect(root.runRlmChild("replacement before retry", { name: "settled-error-worker" })).rejects.toThrow( "an agent of that name already exists at depth 1 under this parent", ); @@ -2023,9 +2038,9 @@ describe("AgentSession rlm recursion", () => { retryCleanup.resolve(); await quiescence; - expect(internals._activeRlmChildRuns.has(spawned.rlm_child_id)).toBe(false); - expect(internals._unsettledRlmChildRuns.has(run)).toBe(false); - expect(internals._deletingRlmChildren.has(spawned.rlm_child_id)).toBe(false); + expect(internals._children._activeRlmChildRuns.has(spawned.rlm_child_id)).toBe(false); + expect(internals._children._unsettledRlmChildRuns.has(run)).toBe(false); + expect(internals._children._deletingRlmChildren.has(spawned.rlm_child_id)).toBe(false); expect(deleteRlmSubagentRuntime).toHaveBeenCalledTimes(2); await expect( root.runRlmChild("replacement after cleanup", { name: "settled-error-worker" }), @@ -2160,12 +2175,12 @@ describe("AgentSession rlm recursion", () => { session_name: daemonChildId, status: "completed" as const, }; - inspectable._deletingRlmChildren.set("deleting-child", { + inspectable._children._deletingRlmChildren.set("deleting-child", { subagent: conflictingDeletion, promise: Promise.resolve({ subagent: conflictingDeletion }), }); await expect(root.deleteRlmSubagent(daemonChildId)).rejects.toThrow("is ambiguous"); - inspectable._deletingRlmChildren.delete("deleting-child"); + inspectable._children._deletingRlmChildren.delete("deleting-child"); const handlers = inspectable._createKernelHostHandlers(); const listHandler = handlers["rlm.list_subagents"]; @@ -2325,7 +2340,7 @@ describe("AgentSession rlm recursion", () => { throw new Error("Missing retained child session"); } const rootInternals = root as unknown as InspectableRlmSession; - await waitFor(() => !rootInternals._activeRlmChildRuns.has(childId)); + await waitFor(() => !rootInternals._children._activeRlmChildRuns.has(childId)); const completeRelease = root.releaseRlmChildSession(childId, child); if (!completeRelease) throw new Error("Failed to release retained child"); @@ -2372,7 +2387,7 @@ describe("AgentSession rlm recursion", () => { }, agentMessageController: { listAgents: () => { - const run = [...(root as unknown as InspectableRlmSession)._activeRlmChildRuns.values()][0]; + const run = [...(root as unknown as InspectableRlmSession)._children._activeRlmChildRuns.values()][0]; return { current: { activeSessionId: "parent-active", sessionId: root.sessionId }, agents: run?.session @@ -2407,7 +2422,7 @@ describe("AgentSession rlm recursion", () => { await root.runRlmChild("slow shard"); await waitFor(() => childStarted); - const rootRun = [...(root as unknown as InspectableRlmSession)._activeRlmChildRuns.values()][0]; + const rootRun = [...(root as unknown as InspectableRlmSession)._children._activeRlmChildRuns.values()][0]; if (!rootRun?.session) { throw new Error("Missing child session on root run"); } @@ -2457,7 +2472,7 @@ describe("AgentSession rlm recursion", () => { await root.runRlmChild("slow shard"); await waitFor(() => childStarted); - const run = [...(root as unknown as InspectableRlmSession)._activeRlmChildRuns.values()][0]; + const run = [...(root as unknown as InspectableRlmSession)._children._activeRlmChildRuns.values()][0]; if (!run?.emitUpdate || !run.session) throw new Error("Missing child run emit"); await waitFor(() => run.activity?.kind === "waiting"); const before = updates; @@ -3084,7 +3099,7 @@ describe("AgentSession rlm recursion", () => { await waitFor(() => root.getRlmChildSession(childResult.rlm_child_id) !== undefined); const child = root.getRlmChildSession(childResult.rlm_child_id); if (!child?.sessionFile) throw new Error("Missing persisted child session"); - await waitFor(() => (root as unknown as InspectableRlmSession)._activeRlmChildRuns.size === 0); + await root.waitForRlmQuiescence(); expect(child.getRlmMaxDepthStatus()).toEqual({ maxDepth: 2, source: "inherited" }); await child.setRlmMaxDepth(3); @@ -3104,7 +3119,7 @@ describe("AgentSession rlm recursion", () => { await waitFor(() => root.getRlmChildSession(childResult.rlm_child_id) !== undefined); const child = root.getRlmChildSession(childResult.rlm_child_id); if (!child) throw new Error("Missing retained child session"); - await waitFor(() => (root as unknown as InspectableRlmSession)._activeRlmChildRuns.size === 0); + await root.waitForRlmQuiescence(); expect(child.rlmMaxDepth).toBe(2); await child.setRlmMaxDepth(3); @@ -3282,7 +3297,7 @@ describe("AgentSession rlm recursion", () => { const spawned = await root.runRlmChild("slow shard"); await waitFor(() => childStarted); - const runs = (root as unknown as InspectableRlmSession)._activeRlmChildRuns; + const runs = (root as unknown as InspectableRlmSession)._children._activeRlmChildRuns; expect(runs.size).toBe(1); const run = [...runs.values()][0]; @@ -3316,7 +3331,7 @@ describe("AgentSession rlm recursion", () => { const spawned = await root.runRlmChild("slow shard"); await waitFor(() => childStarted); - const runs = (root as unknown as InspectableRlmSession)._activeRlmChildRuns; + const runs = (root as unknown as InspectableRlmSession)._children._activeRlmChildRuns; expect(runs.size).toBe(1); const run = [...runs.values()][0]; @@ -3361,7 +3376,7 @@ describe("AgentSession rlm recursion", () => { const spawned = await root.runRlmChild("cancel before admission", { name: "cancelled-worker" }); await waitFor(() => agentListStarted); - const runs = (root as unknown as InspectableRlmSession)._activeRlmChildRuns; + const runs = (root as unknown as InspectableRlmSession)._children._activeRlmChildRuns; const run = runs.get(spawned.rlm_child_id); if (!run) throw new Error("Missing running child"); @@ -3387,7 +3402,7 @@ describe("AgentSession rlm recursion", () => { }; const child = createSession({ rlmSessionDir: join(tempDir, "update-restart-parent") }); const childInternals = child as unknown as InspectableRlmSession; - childInternals._activeRlmChildRuns.set("live-grandchild", { + childInternals._children._activeRlmChildRuns.set("live-grandchild", { id: "live-grandchild", prompt: "still working", sessionName: "live-grandchild", @@ -3412,8 +3427,8 @@ describe("AgentSession rlm recursion", () => { settlement: deferred(), session: child, }; - rootInternals._activeRlmChildRuns.set(run.id, run); - rootInternals._unsettledRlmChildRuns.add(run); + rootInternals._children._activeRlmChildRuns.set(run.id, run); + rootInternals._children._unsettledRlmChildRuns.add(run); const quiescence = root.waitForRlmQuiescence(); await Promise.resolve(); @@ -3425,9 +3440,9 @@ describe("AgentSession rlm recursion", () => { expect.arrayContaining([expect.objectContaining({ id: "live-grandchild", status: "running" })]), ); - rootInternals._activeRlmChildRuns.clear(); - rootInternals._unsettledRlmChildRuns.clear(); - childInternals._activeRlmChildRuns.clear(); + rootInternals._children._activeRlmChildRuns.clear(); + rootInternals._children._unsettledRlmChildRuns.clear(); + childInternals._children._activeRlmChildRuns.clear(); root.dispose(); child.dispose(); }); @@ -3463,7 +3478,9 @@ describe("AgentSession rlm recursion", () => { await expect(root.waitForRlmQuiescence()).resolves.toBeUndefined(); releaseStartup(); - await waitFor(() => !(root as unknown as InspectableRlmSession)._activeRlmChildRuns.has(spawned.rlm_child_id)); + await waitFor( + () => !(root as unknown as InspectableRlmSession)._children._activeRlmChildRuns.has(spawned.rlm_child_id), + ); expect(promptAndWait).not.toHaveBeenCalled(); }); @@ -3489,7 +3506,7 @@ describe("AgentSession rlm recursion", () => { await root.runRlmChild("slow shard"); await waitFor(() => childStarted); - const runs = (root as unknown as InspectableRlmSession)._activeRlmChildRuns; + const runs = (root as unknown as InspectableRlmSession)._children._activeRlmChildRuns; expect(runs.size).toBe(1); const run = [...runs.values()][0]; @@ -3529,7 +3546,7 @@ describe("AgentSession rlm recursion", () => { await root.runRlmChild("slow shard"); await waitFor(() => childStarted); - const runs = (root as unknown as InspectableRlmSession)._activeRlmChildRuns; + const runs = (root as unknown as InspectableRlmSession)._children._activeRlmChildRuns; expect(runs.size).toBe(1); const childId = [...runs.keys()][0]; if (!childId) { @@ -3553,14 +3570,12 @@ describe("AgentSession rlm recursion", () => { publication: { reject: vi.fn() }, emitUpdate: vi.fn(), }; - (deepHost as unknown as { _activeRlmChildRuns: Map })._activeRlmChildRuns.set( - "deep-1", - deepRun, - ); - (run?.session as unknown as { _rlmChildSessions: Map })._rlmChildSessions.set( - "deep-host", - { session: deepHost }, - ); + ( + deepHost as unknown as { _children: { _activeRlmChildRuns: Map } } + )._children._activeRlmChildRuns.set("deep-1", deepRun); + ( + run?.session as unknown as { _children: { _rlmChildSessions: Map } } + )._children._rlmChildSessions.set("deep-host", { session: deepHost }); expect(root.cancelRlmChildRun(childId)).toBe(true); expect(run?.status).toBe("cancelled"); @@ -3583,10 +3598,10 @@ describe("AgentSession rlm recursion", () => { const root = createSession({ rlmSessionDir: join(tempDir, "collide-root") }); const finished = createSession({ rlmSessionDir: join(tempDir, "collide-finished") }); const otherParent = createSession({ rlmSessionDir: join(tempDir, "collide-other") }); - const rootMaps = root as unknown as { _rlmChildSessions: Map }; + const rootMaps = root as unknown as { _children: { _rlmChildSessions: Map } }; // Child ids are only mkdir-unique among siblings: "sub-dup" exists twice. - rootMaps._rlmChildSessions.set("sub-dup", { session: finished }); - rootMaps._rlmChildSessions.set("other-parent", { session: otherParent }); + rootMaps._children._rlmChildSessions.set("sub-dup", { session: finished }); + rootMaps._children._rlmChildSessions.set("other-parent", { session: otherParent }); const abort = vi.fn(); const collidingRun = { id: "sub-dup", @@ -3596,10 +3611,9 @@ describe("AgentSession rlm recursion", () => { publication: { reject: vi.fn() }, emitUpdate: vi.fn(), }; - (otherParent as unknown as { _activeRlmChildRuns: Map })._activeRlmChildRuns.set( - "sub-dup", - collidingRun, - ); + ( + otherParent as unknown as { _children: { _activeRlmChildRuns: Map } } + )._children._activeRlmChildRuns.set("sub-dup", collidingRun); expect(root.cancelRlmChildRun("sub-dup")).toBe(true); expect(collidingRun.status).toBe("cancelled"); @@ -3614,12 +3628,16 @@ describe("AgentSession rlm recursion", () => { let cancelPrimitiveCalls = 0; let runMapIterations = 0; for (const [level, session] of sessions.entries()) { - const target = session as unknown as { - _activeRlmChildRuns: Map; - _rlmChildSessions: Map; - _cancelRlmChildRun(run: unknown, reason: string): boolean; - }; - const original = target._cancelRlmChildRun.bind(session); + const target = ( + session as unknown as { + _children: { + _activeRlmChildRuns: Map; + _rlmChildSessions: Map; + _cancelRlmChildRun(run: unknown, reason: string): boolean; + }; + } + )._children; + const original = target._cancelRlmChildRun.bind(target); target._cancelRlmChildRun = (run, reason) => { cancelPrimitiveCalls++; return original(run, reason); @@ -3632,10 +3650,12 @@ describe("AgentSession rlm recursion", () => { if (level === 0) continue; // A finished intermediate lives in BOTH parent maps until passivation. const parent = sessions[level - 1] as unknown as { - _activeRlmChildRuns: Map; - _rlmChildSessions: Map; + _children: { + _activeRlmChildRuns: Map; + _rlmChildSessions: Map; + }; }; - parent._activeRlmChildRuns.set(`chain-${level}`, { + parent._children._activeRlmChildRuns.set(`chain-${level}`, { id: `chain-${level}`, status: "done", settled: true, @@ -3644,7 +3664,7 @@ describe("AgentSession rlm recursion", () => { publication: { reject: vi.fn() }, emitUpdate: vi.fn(), }); - parent._rlmChildSessions.set(`chain-${level}`, { session }); + parent._children._rlmChildSessions.set(`chain-${level}`, { session }); } const leafAbort = vi.fn(); const leafRun = { @@ -3655,10 +3675,9 @@ describe("AgentSession rlm recursion", () => { publication: { reject: vi.fn() }, emitUpdate: vi.fn(), }; - (sessions[levels] as unknown as { _activeRlmChildRuns: Map })._activeRlmChildRuns.set( - "leaf-run", - leafRun, - ); + ( + sessions[levels] as unknown as { _children: { _activeRlmChildRuns: Map } } + )._children._activeRlmChildRuns.set("leaf-run", leafRun); expect(sessions[0]!.hasRunningRlmChildren()).toBe(true); expect(runMapIterations).toBeLessThanOrEqual(3 * (levels + 1)); @@ -3684,13 +3703,12 @@ describe("AgentSession rlm recursion", () => { }, }); - await root.runRlmChild("quick shard"); - const runs = (root as unknown as InspectableRlmSession)._activeRlmChildRuns; - await waitFor(() => runs.size === 0); - const retained = (root as unknown as { _rlmChildSessions: Map }) - ._rlmChildSessions; - expect(retained.size).toBe(1); - const [childId, { session: childSession }] = [...retained.entries()][0]!; + const spawned = await root.runRlmChild("quick shard"); + await root.waitForRlmQuiescence(); + const childId = spawned.rlm_child_id; + expect(root.getRlmChildSnapshots()).toEqual([expect.objectContaining({ id: childId, status: "done" })]); + const childSession = root.getRlmChildSession(childId); + if (!childSession) throw new Error("Missing retained child session"); const abort = vi.fn(); const grandchild = { id: "grandchild-1", @@ -3700,10 +3718,9 @@ describe("AgentSession rlm recursion", () => { publication: { reject: vi.fn() }, emitUpdate: vi.fn(), }; - (childSession as unknown as { _activeRlmChildRuns: Map })._activeRlmChildRuns.set( - "grandchild-1", - grandchild, - ); + ( + childSession as unknown as { _children: { _activeRlmChildRuns: Map } } + )._children._activeRlmChildRuns.set("grandchild-1", grandchild); expect(root.cancelRlmChildRun(childId)).toBe(true); expect(grandchild.status).toBe("cancelled"); @@ -3782,14 +3799,15 @@ describe("AgentSession rlm recursion", () => { const runPromise = root.runRlmChild("slow child", { name: "retained-worker" }); await waitFor(() => childStarted); - const childId = [...(root as unknown as InspectableRlmSession)._activeRlmChildRuns.keys()][0]!; + const childId = [...(root as unknown as InspectableRlmSession)._children._activeRlmChildRuns.keys()][0]!; await expect(root.deleteInactiveRlmSubagent(childId)).resolves.toBe("running"); expect(deleteRuntime).not.toHaveBeenCalled(); releaseChild(); await expect(runPromise).resolves.toMatchObject({ name: "retained-worker" }); await waitFor( - () => (root as unknown as InspectableRlmSession)._activeRlmChildRuns.get(childId)?.status !== "running", + () => + (root as unknown as InspectableRlmSession)._children._activeRlmChildRuns.get(childId)?.status !== "running", ); await expect(root.deleteInactiveRlmSubagent(childId)).resolves.toBe("deleted"); expect(deleteRuntime).toHaveBeenCalledWith(childId, retainedChild); @@ -3818,7 +3836,7 @@ describe("AgentSession rlm recursion", () => { ); }); expect(disposeChild).not.toHaveBeenCalled(); - expect((root as unknown as InspectableRlmSession)._activeRlmChildRuns.size).toBe(0); + expect((root as unknown as InspectableRlmSession)._children._activeRlmChildRuns.size).toBe(0); expect(root.getRlmChildSession(spawned.rlm_child_id)).toBeUndefined(); }); @@ -4026,7 +4044,7 @@ describe("AgentSession rlm recursion", () => { subagent: { rlm_child_id: spawned.rlm_child_id }, }); const internals = root as unknown as InspectableRlmSession; - await waitFor(() => internals._rlmChildCleanupFailures.size === 1); + await waitFor(() => internals._children._rlmChildCleanupFailures.size === 1); const failureContent = await vi.waitFor(() => { const notice = root.messages.find( (message) => message.role === "custom" && message.customType === "rlm_child_failure", @@ -4060,7 +4078,7 @@ describe("AgentSession rlm recursion", () => { (message) => message.role === "custom" && message.customType === "rlm_child_terminal_notice", ), ).toHaveLength(1); - expect(internals._rlmChildCleanupFailures.size).toBe(0); + expect(internals._children._rlmChildCleanupFailures.size).toBe(0); await expect(root.runRlmChild("replacement", { name: "retry-worker" })).resolves.toMatchObject({ name: "retry-worker", }); @@ -4092,7 +4110,7 @@ describe("AgentSession rlm recursion", () => { cleanups[attempt]!.reject(new Error(`cleanup failure ${attempt + 1}`)); await vi.waitFor(() => { expect(failures()).toHaveLength(attempt + 1); - expect(internals._deletingRlmChildren.size).toBe(0); + expect(internals._children._deletingRlmChildren.size).toBe(0); }); expect(terminalNotices()).toHaveLength(0); expect((await root.listRlmSubagents()).subagents).toEqual([]); @@ -4107,7 +4125,7 @@ describe("AgentSession rlm recursion", () => { expect(terminalNotices()).toEqual([ expect.objectContaining({ details: expect.objectContaining({ kind: "cancelled" }) }), ]); - expect(internals._rlmChildCleanupFailures.size).toBe(0); + expect(internals._children._rlmChildCleanupFailures.size).toBe(0); expect(root.getRlmChildSession(spawned.rlm_child_id)).toBeUndefined(); await hostedChild.disposeAsync(); }); @@ -4127,15 +4145,15 @@ describe("AgentSession rlm recursion", () => { await waitFor(hasStarted); await root.deleteRlmSubagent(spawned.rlm_child_id); const internals = root as unknown as InspectableRlmSession; - const run = internals._activeRlmChildRuns.get(spawned.rlm_child_id); + const run = internals._children._activeRlmChildRuns.get(spawned.rlm_child_id); if (!run) throw new Error("Missing deleting run"); - await waitFor(() => internals._rlmChildCleanupFailures.size === 1); + await waitFor(() => internals._children._rlmChildCleanupFailures.size === 1); await root.disposeAsync(); expect(deleteRuntime).toHaveBeenCalledOnce(); expect(disposeHostedChild).toHaveBeenCalled(); - expect(internals._activeRlmChildRuns.has(spawned.rlm_child_id)).toBe(false); - expect(internals._unsettledRlmChildRuns.has(run)).toBe(false); + expect(internals._children._activeRlmChildRuns.has(spawned.rlm_child_id)).toBe(false); + expect(internals._children._unsettledRlmChildRuns.has(run)).toBe(false); expect( root.messages.filter( (message) => message.role === "custom" && message.customType === "rlm_child_terminal_notice", @@ -4159,7 +4177,7 @@ describe("AgentSession rlm recursion", () => { await waitFor(hasStarted); await root.deleteRlmSubagent(spawned.rlm_child_id); const internals = root as unknown as InspectableRlmSession; - const run = internals._activeRlmChildRuns.get(spawned.rlm_child_id); + const run = internals._children._activeRlmChildRuns.get(spawned.rlm_child_id); if (!run) throw new Error("Missing deleting run"); const disposal = root.disposeAsync(); @@ -4167,8 +4185,8 @@ describe("AgentSession rlm recursion", () => { cleanup.resolve(); await disposal; expect(deleteRuntime).toHaveBeenCalledOnce(); - expect(internals._activeRlmChildRuns.has(spawned.rlm_child_id)).toBe(false); - expect(internals._unsettledRlmChildRuns.has(run)).toBe(false); + expect(internals._children._activeRlmChildRuns.has(spawned.rlm_child_id)).toBe(false); + expect(internals._children._unsettledRlmChildRuns.has(run)).toBe(false); expect( root.messages.filter( (message) => message.role === "custom" && message.customType === "rlm_child_terminal_notice", @@ -4193,7 +4211,7 @@ describe("AgentSession rlm recursion", () => { await waitFor(hasStarted); await root.deleteRlmSubagent(spawned.rlm_child_id); const internals = root as unknown as InspectableRlmSession; - const run = internals._activeRlmChildRuns.get(spawned.rlm_child_id); + const run = internals._children._activeRlmChildRuns.get(spawned.rlm_child_id); if (!run) throw new Error("Missing deleting run"); const disposal = root.disposeAsync(); @@ -4203,8 +4221,8 @@ describe("AgentSession rlm recursion", () => { await disposal; expect(deleteRuntime).toHaveBeenCalledOnce(); expect(disposeHostedChild).toHaveBeenCalled(); - expect(internals._activeRlmChildRuns.has(spawned.rlm_child_id)).toBe(false); - expect(internals._unsettledRlmChildRuns.has(run)).toBe(false); + expect(internals._children._activeRlmChildRuns.has(spawned.rlm_child_id)).toBe(false); + expect(internals._children._unsettledRlmChildRuns.has(run)).toBe(false); expect( root.messages.filter( (message) => @@ -4269,7 +4287,7 @@ describe("AgentSession rlm recursion", () => { await waitFor( () => root.getRlmChildSession(spawned.rlm_child_id) !== undefined && - !internals._activeRlmChildRuns.has(spawned.rlm_child_id), + !internals._children._activeRlmChildRuns.has(spawned.rlm_child_id), ); expect(root.unfinishedActionCount).toBe(1); @@ -4285,7 +4303,7 @@ describe("AgentSession rlm recursion", () => { (message) => message.role === "custom" && message.customType === "rlm_child_terminal_notice", ), ).toEqual([expect.objectContaining({ details: expect.objectContaining({ kind: "completed_without_reply" }) })]); - expect(internals._activeRlmChildRuns.has(spawned.rlm_child_id)).toBe(false); + expect(internals._children._activeRlmChildRuns.has(spawned.rlm_child_id)).toBe(false); }); it("keeps failed closure retryable without hanging or late resurrection", async () => { @@ -4304,11 +4322,11 @@ describe("AgentSession rlm recursion", () => { expect(root.registerRlmChildSession("retry-child", child)).toBe(true); await expect(root.deleteRlmSubagent("release-worker")).rejects.toThrow("close failed"); expect(await root.listRlmSubagents()).toEqual({ subagents: [] }); - expect((root as unknown as InspectableRlmSession)._rlmChildCleanupFailures.size).toBe(1); + expect((root as unknown as InspectableRlmSession)._children._rlmChildCleanupFailures.size).toBe(1); await expect(root.deleteRlmSubagent("release-worker")).resolves.toMatchObject({ subagent: { rlm_child_id: "retry-child" }, }); - expect((root as unknown as InspectableRlmSession)._rlmChildCleanupFailures.size).toBe(0); + expect((root as unknown as InspectableRlmSession)._children._rlmChildCleanupFailures.size).toBe(0); }); it("does not restore failed delete retry state after parent teardown", async () => { @@ -4326,10 +4344,10 @@ describe("AgentSession rlm recursion", () => { await expect(root.deleteRlmSubagent("teardown-worker")).rejects.toThrow("close failed during teardown"); root.dispose(); const internals = root as unknown as InspectableRlmSession; - expect(internals._activeRlmChildRuns.size).toBe(0); - expect(internals._rlmChildSessions.size).toBe(0); - expect(internals._rlmChildUnsubscribes.size).toBe(0); - expect(internals._rlmChildCleanupFailures.size).toBe(0); + expect(internals._children._activeRlmChildRuns.size).toBe(0); + expect(internals._children._rlmChildSessions.size).toBe(0); + expect(internals._children._rlmChildUnsubscribes.size).toBe(0); + expect(internals._children._rlmChildCleanupFailures.size).toBe(0); }); it("keeps an errored startup deletable after its failure notice is durably admitted", async () => { @@ -4352,7 +4370,7 @@ describe("AgentSession rlm recursion", () => { const failed = (await root.listRlmSubagents()).subagents[0]; await expect(root.deleteRlmSubagent("failed-worker")).resolves.toEqual({ subagent: failed }); const internals = root as unknown as InspectableRlmSession; - expect(internals._activeRlmChildRuns.size).toBe(0); + expect(internals._children._activeRlmChildRuns.size).toBe(0); expect(await root.listRlmSubagents()).toEqual({ subagents: [] }); await expect(root.runRlmChild("replacement", { name: "failed-worker" })).resolves.toMatchObject({ name: "failed-worker", @@ -4386,7 +4404,7 @@ describe("AgentSession rlm recursion", () => { ); releaseRuntimeCreation(); - await waitFor(() => (root as unknown as InspectableRlmSession)._activeRlmChildRuns.size === 0); + await waitFor(() => (root as unknown as InspectableRlmSession)._children._activeRlmChildRuns.size === 0); expect(setSessionName).not.toHaveBeenCalled(); await expect(root.runRlmChild("replacement", { name: "reserved-worker" })).resolves.toMatchObject({ name: "reserved-worker", @@ -4419,12 +4437,12 @@ describe("AgentSession rlm recursion", () => { expect(queued).toBeDefined(); await expect(root.deleteRlmSubagent("queued-worker")).resolves.toEqual({ subagent: queued }); - expect((root as unknown as InspectableRlmSession)._activeRlmChildRuns.size).toBe(1); + expect((root as unknown as InspectableRlmSession)._children._activeRlmChildRuns.size).toBe(1); expect(await root.listRlmSubagents()).toEqual({ subagents: [] }); releaseRuntimeCreation(); await waitFor(() => deleteRuntime.mock.calls.length === 1); - await waitFor(() => (root as unknown as InspectableRlmSession)._activeRlmChildRuns.size === 0); + await waitFor(() => (root as unknown as InspectableRlmSession)._children._activeRlmChildRuns.size === 0); expect(await root.listRlmSubagents()).toEqual({ subagents: [] }); }); @@ -4470,7 +4488,7 @@ describe("AgentSession rlm recursion", () => { const internals = root as unknown as InspectableRlmSession; await waitFor(() => deleteRuntime.mock.calls.length === 1); - await waitFor(() => internals._rlmChildCleanupFailures.size === 1); + await waitFor(() => internals._children._rlmChildCleanupFailures.size === 1); expect(await root.listRlmSubagents()).toEqual({ subagents: [] }); expect(disposeHostedChild).not.toHaveBeenCalled(); @@ -4478,7 +4496,7 @@ describe("AgentSession rlm recursion", () => { await root.waitForRlmQuiescence(); expect(deleteRuntime).toHaveBeenCalledTimes(2); expect(disposeHostedChild).toHaveBeenCalledOnce(); - expect(internals._rlmChildCleanupFailures.size).toBe(0); + expect(internals._children._rlmChildCleanupFailures.size).toBe(0); }); it("accepts deletion of a running direct child without waiting for task unwind", async () => { @@ -4534,7 +4552,7 @@ describe("AgentSession rlm recursion", () => { const parentPromise = root.runRlmChild("slow parent"); await waitFor(() => parentStarted); - const parentRun = [...(root as unknown as InspectableRlmSession)._activeRlmChildRuns.values()][0]; + const parentRun = [...(root as unknown as InspectableRlmSession)._children._activeRlmChildRuns.values()][0]; if (!parentRun?.session) { throw new Error("Missing parent child session"); } @@ -4550,7 +4568,9 @@ describe("AgentSession rlm recursion", () => { } const disposeNested = vi.spyOn(nestedSession, "disposeAsync"); await waitFor(() => { - const status = (parentSession as unknown as InspectableRlmSession)._activeRlmChildRuns.get(nestedId)?.status; + const status = (parentSession as unknown as InspectableRlmSession)._children._activeRlmChildRuns.get( + nestedId, + )?.status; return status !== "queued" && status !== "running"; }); @@ -4562,7 +4582,9 @@ describe("AgentSession rlm recursion", () => { releaseParent(); await parentPromise; await waitFor(() => { - const status = (root as unknown as InspectableRlmSession)._activeRlmChildRuns.get(parentRun.id)?.status; + const status = (root as unknown as InspectableRlmSession)._children._activeRlmChildRuns.get( + parentRun.id, + )?.status; return status === undefined || status === "done"; }); }); @@ -4602,7 +4624,7 @@ describe("AgentSession rlm recursion", () => { await root.runRlmChild("slow shard"); await waitFor(() => childStarted); - const rootRuns = (root as unknown as InspectableRlmSession)._activeRlmChildRuns; + const rootRuns = (root as unknown as InspectableRlmSession)._children._activeRlmChildRuns; const rootRun = [...rootRuns.values()][0]; if (!rootRun?.session) { throw new Error("Missing child session on root run"); @@ -4611,7 +4633,7 @@ describe("AgentSession rlm recursion", () => { const childSession = rootRun.session; const nestedSpawned = await childSession.runRlmChild("nested shard"); await waitFor(() => nestedStarted); - const nestedRuns = (childSession as unknown as InspectableRlmSession)._activeRlmChildRuns; + const nestedRuns = (childSession as unknown as InspectableRlmSession)._children._activeRlmChildRuns; expect(nestedRuns.size).toBe(1); const nestedId = [...nestedRuns.keys()][0]; if (!nestedId) { diff --git a/packages/coding-agent/test/agent-session-services.test.ts b/packages/coding-agent/test/agent-session-services.test.ts index ca25424d6d..1d9411f51a 100644 --- a/packages/coding-agent/test/agent-session-services.test.ts +++ b/packages/coding-agent/test/agent-session-services.test.ts @@ -7,15 +7,18 @@ import { AGENT_MESSAGE_SKILL_NAME, type AgentSessionMessageController } from ".. import { AGENT_OBSERVE_SKILL_NAME, type AgentObserveController } from "../src/core/agent-observe.js"; import { createAgentSessionFromServices, createAgentSessionServices } from "../src/core/agent-session-services.js"; import { AuthStorage } from "../src/core/auth-storage.js"; +import type { KernelClient } from "../src/core/kernel/index.js"; import { SessionManager } from "../src/core/session-manager.js"; import { SettingsManager } from "../src/core/settings-manager.js"; import { createSyntheticSourceInfo } from "../src/core/source-info.js"; +import { IpythonKernelProvisioner } from "../src/core/tools/ipython.js"; describe("createAgentSessionFromServices", () => { const cleanupPaths: string[] = []; const unregisters: Array<() => void> = []; afterEach(() => { + vi.restoreAllMocks(); vi.unstubAllEnvs(); while (unregisters.length > 0) { unregisters.pop()?.(); @@ -172,8 +175,8 @@ describe("createAgentSessionFromServices", () => { expect(initialPrompt).not.toContain("Enabled generic MCP servers: `linear`"); const rebuildRuntime = vi.spyOn( - session as unknown as { _rebuildRuntimeForAcpMcpServers(): void }, - "_rebuildRuntimeForAcpMcpServers", + session as unknown as { _buildRuntime(options: { activeToolNames?: string[] }): void }, + "_buildRuntime", ); session.replaceAcpMcpServers( [ @@ -194,11 +197,12 @@ describe("createAgentSessionFromServices", () => { const waitForIdle = vi.spyOn(session.agent, "waitForIdle"); await session.releaseAcpMcpServers("unknown-owner", ["task"]); expect(waitForIdle).not.toHaveBeenCalled(); - const originalProvisioner = Reflect.get(session, "_ipythonKernelProvisioner"); const execute = vi.fn(async (_code: string) => ({ status: "ok" })); - Reflect.set(session, "_ipythonKernelProvisioner", { manager: { isRunning: true, execute } }); + const managerGetter = vi + .spyOn(IpythonKernelProvisioner.prototype, "manager", "get") + .mockReturnValue({ isRunning: true, execute } as unknown as KernelClient); await session.releaseAcpMcpServers("owner-a", ["task"]); - Reflect.set(session, "_ipythonKernelProvisioner", originalProvisioner); + managerGetter.mockRestore(); expect(rebuildRuntime).not.toHaveBeenCalled(); expect(execute).toHaveBeenCalledOnce(); expect(execute.mock.calls[0]?.[0]).toContain("await _prime_mcp.reload(_prime_mcp_name)"); diff --git a/packages/coding-agent/test/session-child-usage.test.ts b/packages/coding-agent/test/session-child-usage.test.ts new file mode 100644 index 0000000000..83c7b317fd --- /dev/null +++ b/packages/coding-agent/test/session-child-usage.test.ts @@ -0,0 +1,210 @@ +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import { type AssistantMessage, fauxAssistantMessage, type Usage } from "@earendil-works/pi-ai"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { createAgentSessionMessage } from "../src/core/agent-messages.js"; +import { SessionManager } from "../src/core/session-manager.js"; +import { cloneUsage, emptyUsage } from "../src/core/usage.js"; +import { type ChildUsageHost, type ChildUsageTracker, SessionChildUsage } from "../src/session/child-usage.js"; + +function usage(input: number, output: number): Usage { + return { + input, + output, + cacheRead: 0, + cacheWrite: 0, + totalTokens: input + output, + cost: { input, output, cacheRead: 0, cacheWrite: 0, total: input + output }, + }; +} + +function record(tracker: ChildUsageTracker, childUsage: Usage, prompt?: AgentMessage): void { + const assistant: AssistantMessage = { ...fauxAssistantMessage("child completion"), usage: childUsage }; + tracker.record(prompt ? [prompt, assistant] : [assistant], assistant); +} + +function unindexedUsage(owner: SessionChildUsage, parent: AssistantMessage): Usage | undefined { + // Exact snapshots distinguish retained zero entries from missing entries and avoid subtraction's token clamp. + return ( + owner as unknown as { _rlmUnindexedChildUsage: WeakMap } + )._rlmUnindexedChildUsage.get(parent); +} + +function setup(appendParent = true) { + const manager = SessionManager.inMemory(); + const parent: AssistantMessage = { ...fauxAssistantMessage("spawn children"), usage: usage(2, 1) }; + if (appendParent) manager.appendMessage(parent); + const drains: Array<() => void> = []; + const host = { + sessionManager: manager, + afterParentDrain: vi.fn((flush: () => void) => drains.push(flush)), + invalidateOwnUsage: vi.fn(), + } satisfies ChildUsageHost; + const owner = new SessionChildUsage(host); + const attributions = () => manager.getEntries().filter((entry) => entry.type === "child_usage_attributed"); + return { manager, parent, drains, host, owner, attributions }; +} + +describe("SessionChildUsage boundaries", () => { + beforeEach(() => { + vi.useFakeTimers(); + vi.setSystemTime(1_000); + }); + + afterEach(() => { + vi.clearAllTimers(); + vi.useRealTimers(); + vi.restoreAllMocks(); + }); + + it("shares durable aggregation across siblings while preserving pending live usage and context size", () => { + const { manager, parent, owner, attributions } = setup(); + const first = owner.createTracker(parent); + record(first, usage(7, 3)); + const second = owner.createTracker(parent); + record(second, usage(11, 5)); + + expect(parent.usage).toEqual({ ...usage(20, 9), totalTokens: 3 }); + expect(unindexedUsage(owner, parent)).toEqual(usage(18, 8)); + const own = cloneUsage(parent.usage); + owner.subtractUnindexed(own, manager.getEntries()); + expect(own).toEqual({ ...usage(2, 1), totalTokens: 0 }); + + first.flush(); + expect(attributions().map((entry) => entry.aggregateUsage)).toEqual([{ ...usage(9, 4), totalTokens: 3 }]); + expect(unindexedUsage(owner, parent)).toEqual(usage(11, 5)); + expect(parent.usage).toEqual({ ...usage(20, 9), totalTokens: 3 }); + + second.flush(); + expect(attributions().map((entry) => entry.aggregateUsage)).toEqual([ + { ...usage(9, 4), totalTokens: 3 }, + { ...usage(20, 9), totalTokens: 3 }, + ]); + expect(unindexedUsage(owner, parent)).toEqual(emptyUsage()); + expect(vi.getTimerCount()).toBe(0); + }); + + it("returns from settlement flush before a delayed parent append and drains it exactly once", () => { + const { manager, parent, owner, host, drains, attributions } = setup(false); + const tracker = owner.createTracker(parent); + record(tracker, usage(7, 3)); + + expect(tracker.flush()).toBeUndefined(); + tracker.flush(); + expect(host.afterParentDrain).toHaveBeenCalledOnce(); + expect(attributions()).toEqual([]); + expect(unindexedUsage(owner, parent)).toEqual(usage(7, 3)); + expect(vi.getTimerCount()).toBe(0); + + manager.appendMessage(parent); + drains[0]!(); + tracker.flush(); + expect(attributions()).toHaveLength(1); + expect(attributions()[0]).toMatchObject({ + childUsage: usage(7, 3), + aggregateUsage: { ...usage(9, 4), totalTokens: 3 }, + }); + expect(unindexedUsage(owner, parent)).toEqual(emptyUsage()); + }); + + it("does not subtract indexed usage twice when append fails after indexing", () => { + const { manager, parent, owner, host, attributions } = setup(); + const tracker = owner.createTracker(parent); + const append = manager.appendChildUsageAttribution.bind(manager); + vi.spyOn(manager, "appendChildUsageAttribution").mockImplementationOnce((...args) => { + append(...args); + throw new Error("persist failed after indexing"); + }); + record(tracker, usage(7, 3)); + const liveUsage = parent.usage; + + expect(() => tracker.flush()).not.toThrow(); + expect(attributions()).toHaveLength(1); + expect(parent.usage).toBe(liveUsage); + expect(parent.usage).toEqual({ ...usage(9, 4), totalTokens: 3 }); + expect(unindexedUsage(owner, parent)).toEqual(emptyUsage()); + const ownAfterIndexedSubtraction = usage(2, 1); + owner.subtractUnindexed(ownAfterIndexedSubtraction, manager.getEntries()); + expect(ownAfterIndexedSubtraction).toEqual(usage(2, 1)); + expect(host.invalidateOwnUsage).toHaveBeenCalledTimes(2); + tracker.flush(); + expect(attributions()).toHaveLength(1); + }); + + it("retains unindexed subtraction when append fails before indexing without breaking settlement", () => { + const { manager, parent, owner, attributions } = setup(); + const tracker = owner.createTracker(parent); + vi.spyOn(manager, "appendChildUsageAttribution").mockImplementationOnce(() => { + throw new Error("append rejected"); + }); + record(tracker, usage(7, 3)); + expect(() => tracker.flush()).not.toThrow(); + expect(attributions()).toEqual([]); + expect(unindexedUsage(owner, parent)).toEqual(usage(7, 3)); + const own = cloneUsage(parent.usage); + owner.subtractUnindexed(own, manager.getEntries()); + expect(own.cost.total).toBe(3); + expect(own.input).toBe(2); + expect(own.output).toBe(1); + }); + + it("flushes the wall-clock backstop once even without another completion or checkpoint", async () => { + const { parent, owner, attributions } = setup(); + const tracker = owner.createTracker(parent); + record(tracker, usage(7, 3)); + await vi.advanceTimersByTimeAsync(59_999); + expect(attributions()).toEqual([]); + await vi.advanceTimersByTimeAsync(1); + expect(attributions()).toHaveLength(1); + expect(attributions()[0]?.childUsage).toEqual(usage(7, 3)); + tracker.flush(); + await vi.advanceTimersByTimeAsync(60_000); + expect(attributions()).toHaveLength(1); + }); + + it.each(["completion", "checkpoint"] as const)("flushes stale usage before a later %s", (boundary) => { + const { parent, owner, attributions } = setup(); + const tracker = owner.createTracker(parent); + record(tracker, usage(7, 3)); + vi.setSystemTime(61_000); + if (boundary === "checkpoint") { + tracker.flushIfStale(); + expect(attributions().map((entry) => entry.childUsage.cost.total)).toEqual([10]); + expect(unindexedUsage(owner, parent)).toEqual(emptyUsage()); + } + record(tracker, usage(11, 5)); + expect(attributions().map((entry) => entry.aggregateUsage.cost.total)).toEqual([13]); + expect(unindexedUsage(owner, parent)).toEqual(usage(11, 5)); + tracker.flush(); + expect(attributions().map((entry) => entry.aggregateUsage.cost.total)).toEqual([13, 29]); + expect(attributions().map((entry) => entry.childUsage.cost.total)).toEqual([10, 16]); + }); + + it("tracks future retained-child completions after initial settlement with their own prompt origins", async () => { + const { parent, owner, attributions } = setup(); + const tracker = owner.createTracker(parent); + const prompt = (id: string) => + createAgentSessionMessage({ + id, + source: "agent_message", + message: "child work", + fromRelationship: "parent", + target: { activeSessionId: "child-active", sessionId: "child" }, + }); + record(tracker, usage(7, 3), prompt("spawn:child")); + tracker.flush(); + expect(vi.getTimerCount()).toBe(0); + + record(tracker, usage(11, 5), prompt("follow-up")); + expect(vi.getTimerCount()).toBe(1); + await vi.advanceTimersByTimeAsync(60_000); + record(tracker, usage(2, 2), { role: "user", content: "direct follow-up", timestamp: Date.now() }); + tracker.flush(); + + expect(attributions().map((entry) => entry.origin)).toEqual(["spawn_task", "agent_message", "direct_user"]); + expect(attributions().map((entry) => entry.aggregateUsage.cost.total)).toEqual([13, 29, 33]); + expect(attributions().map((entry) => entry.childUsage.cost.total)).toEqual([10, 16, 4]); + expect(parent.usage.totalTokens).toBe(3); + expect(unindexedUsage(owner, parent)).toEqual(emptyUsage()); + expect(vi.getTimerCount()).toBe(0); + }); +}); diff --git a/packages/coding-agent/test/session/kernel-environment.test.ts b/packages/coding-agent/test/session/kernel-environment.test.ts new file mode 100644 index 0000000000..23ec72d7f8 --- /dev/null +++ b/packages/coding-agent/test/session/kernel-environment.test.ts @@ -0,0 +1,103 @@ +import { existsSync, mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { AuthStorage } from "../../src/core/auth-storage.js"; +import type { Skill } from "../../src/core/skills.js"; +import { createSyntheticSourceInfo } from "../../src/core/source-info.js"; +import { SERPER_CREDENTIAL_ID, SERPER_ENV_VAR, WEBSEARCH_SKILL_NAME } from "../../src/core/websearch-credential.js"; +import { KernelEnvironment, type KernelEnvironmentHost } from "../../src/session/kernel-environment.js"; + +function createEnvironment(overrides: Partial = {}, sessionDir?: string) { + return new KernelEnvironment( + { + authStorage: AuthStorage.inMemory(), + resourceLoader: { getSkills: () => ({ skills: [], diagnostics: [] }) }, + getDepth: () => 0, + getMaxDepth: () => 3, + getArtifactDir: () => undefined, + getLocalHarnessStateDir: () => undefined, + ...overrides, + }, + sessionDir, + ); +} + +describe("KernelEnvironment", () => { + const directories: string[] = []; + afterEach(() => { + vi.unstubAllEnvs(); + for (const directory of directories.splice(0)) rmSync(directory, { recursive: true, force: true }); + }); + + it("reuses the explicit root while reading current depth and harness location", () => { + const root = mkdtempSync(join(tmpdir(), "session-env-")); + directories.push(root); + const sessionDir = join(root, "shared"); + let depth = 0; + let maxDepth = 2; + let harness = join(root, "first"); + const environment = createEnvironment( + { + getDepth: () => depth, + getMaxDepth: () => maxDepth, + getLocalHarnessStateDir: () => harness, + getArtifactDir: () => join(root, "other"), + }, + sessionDir, + ); + expect(existsSync(sessionDir)).toBe(false); + expect(environment.buildEnv()).toMatchObject({ + RLM_SESSION_DIR: sessionDir, + RLM_DEPTH: "0", + RLM_MAX_DEPTH: "2", + RLM_HARNESS_STATE_DIR: harness, + }); + expect(existsSync(sessionDir)).toBe(true); + depth = 1; + maxDepth = 4; + harness = join(root, "second"); + expect(environment.buildEnv()).toMatchObject({ + RLM_SESSION_DIR: sessionDir, + RLM_DEPTH: "1", + RLM_MAX_DEPTH: "4", + RLM_HARNESS_STATE_DIR: harness, + }); + }); + + it("keeps ephemeral allocation lazy and uses its directory on later provisioning", () => { + const environment = createEnvironment(); + expect(environment.ensureSessionDir()).toBeUndefined(); + expect(environment.buildEnv()).not.toHaveProperty("RLM_SESSION_DIR"); + const directory = environment.createEphemeralSessionDir(); + directories.push(directory); + expect(environment.ensureSessionDir()).toBe(directory); + expect(environment.buildEnv().RLM_SESSION_DIR).toBe(directory); + }); + + it("reads loaded skills and credentials again without overriding an inherited websearch key", () => { + vi.stubEnv(SERPER_ENV_VAR, ""); + const authStorage = AuthStorage.inMemory(); + authStorage.set(SERPER_CREDENTIAL_ID, { type: "api_key", key: "first" }); + const skills: Skill[] = []; + const environment = createEnvironment({ + authStorage, + resourceLoader: { getSkills: () => ({ skills, diagnostics: [] }) }, + }); + expect(environment.buildEnv()).not.toHaveProperty(SERPER_ENV_VAR); + skills.push({ + kind: "markdown", + name: WEBSEARCH_SKILL_NAME, + description: "search", + filePath: "/skills/search/SKILL.md", + baseDir: "/skills/search", + disableModelInvocation: false, + sourceInfo: createSyntheticSourceInfo("", { source: "test" }), + }); + expect(environment.buildEnv()[SERPER_ENV_VAR]).toBe("first"); + authStorage.set(SERPER_CREDENTIAL_ID, { type: "api_key", key: "second" }); + expect(environment.buildEnv()[SERPER_ENV_VAR]).toBe("second"); + vi.stubEnv(SERPER_ENV_VAR, "inherited"); + expect(environment.buildEnv()).not.toHaveProperty(SERPER_ENV_VAR); + }); +}); diff --git a/packages/coding-agent/test/session/kernel.test.ts b/packages/coding-agent/test/session/kernel.test.ts new file mode 100644 index 0000000000..0cd4f26c89 --- /dev/null +++ b/packages/coding-agent/test/session/kernel.test.ts @@ -0,0 +1,122 @@ +import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { snapshotPathIn } from "../../src/core/kernel/state-snapshot.js"; +import type { IpythonToolOptions } from "../../src/core/tools/ipython.js"; +import { SessionKernel, type SessionKernelHost } from "../../src/session/kernel.js"; + +const mocks = vi.hoisted(() => ({ + instances: [] as Array<{ + options: IpythonToolOptions; + dispose: ReturnType Promise>>; + prewarm: ReturnType void>>; + }>, +})); + +vi.mock("../../src/core/tools/ipython.js", () => ({ + IpythonKernelProvisioner: class { + dispose = vi.fn(async () => {}); + prewarm = vi.fn(); + constructor( + _cwd: string, + readonly options: IpythonToolOptions, + ) { + mocks.instances.push(this); + } + }, +})); +vi.mock("../../src/core/tools/index.js", () => ({ createAllToolDefinitions: () => ({}) })); + +function createKernel(overrides: Partial = {}, prewarm = false) { + const sendCustomMessage = vi.fn(async () => {}); + const host: SessionKernelHost = { + cwd: "/workspace", + getArtifactDir: () => undefined, + getSessionId: () => "session", + getEnv: () => ({}), + getShellCommandPrefix: () => undefined, + getShellPath: () => undefined, + createHostHandlers: () => ({}), + recordLateSentAgentMessage: () => {}, + getMessages: () => [], + appendCustomMessageEntry: () => "entry", + emit: () => {}, + sendCustomMessage, + ...overrides, + }; + return { kernel: new SessionKernel(host, prewarm), sendCustomMessage }; +} + +describe("SessionKernel lifecycle", () => { + const directories: string[] = []; + beforeEach(() => { + mocks.instances.length = 0; + }); + afterEach(() => { + for (const directory of directories.splice(0)) rmSync(directory, { recursive: true, force: true }); + }); + + it("gates a replacement on the old snapshot flush and only announces a first-build restore", async () => { + let finishDispose = () => {}; + const disposing = new Promise((resolve) => { + finishDispose = resolve; + }); + let sessionId = "first"; + let depth = "0"; + const { kernel, sendCustomMessage } = createKernel({ + getSessionId: () => sessionId, + getEnv: () => ({ RLM_DEPTH: depth }), + }); + kernel.build([]); + kernel.finishBuild(["ipython"]); + const first = mocks.instances[0]!; + first.dispose.mockReturnValue(disposing); + first.options.onRestore?.({ restored: ["value"], failed: [], path: "/snapshot" }); + expect(sendCustomMessage).toHaveBeenCalledWith( + expect.objectContaining({ display: true, details: { restored: true } }), + { deliverAs: "nextTurn" }, + ); + sessionId = "second"; + depth = "1"; + kernel.build([]); + const second = mocks.instances[1]!; + expect(first.dispose).toHaveBeenCalledOnce(); + expect(second.options.readyGate).toBe(disposing); + expect(second.options.onRestore).toBeUndefined(); + expect(second.options).toMatchObject({ sessionId: "second", env: { RLM_DEPTH: "1" } }); + let ready = false; + void second.options.readyGate?.then(() => { + ready = true; + }); + await Promise.resolve(); + expect(ready).toBe(false); + finishDispose(); + await second.options.readyGate; + expect(ready).toBe(true); + }); + + it("prewarms resumed state only when ipython is active", () => { + const directory = mkdtempSync(join(tmpdir(), "session-kernel-")); + directories.push(directory); + writeFileSync(snapshotPathIn(directory), "snapshot"); + const { kernel } = createKernel({ getArtifactDir: () => directory }); + kernel.build([]); + kernel.finishBuild([]); + expect(mocks.instances[0]!.prewarm).not.toHaveBeenCalled(); + kernel.build([]); + kernel.finishBuild(["ipython"]); + expect(mocks.instances[1]!.prewarm).toHaveBeenCalledOnce(); + }); + + it("passes the teardown snapshot policy and tolerates failed startup cleanup", async () => { + const { kernel } = createKernel({}, true); + kernel.build([]); + kernel.finishBuild(["ipython"]); + const current = mocks.instances[0]!; + expect(current.prewarm).toHaveBeenCalledOnce(); + current.dispose.mockRejectedValue(new Error("startup failed")); + await expect(kernel.dispose(false)).resolves.toBeUndefined(); + expect(current.dispose).toHaveBeenCalledWith({ snapshot: false }); + }); +}); diff --git a/packages/coding-agent/test/session/tools.test.ts b/packages/coding-agent/test/session/tools.test.ts new file mode 100644 index 0000000000..56e5a1b950 --- /dev/null +++ b/packages/coding-agent/test/session/tools.test.ts @@ -0,0 +1,89 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { createDeferred, type KernelClient } from "../../src/core/kernel/index.js"; +import { IpythonKernelProvisioner } from "../../src/core/tools/ipython.js"; +import { SessionTools, type SessionToolsHost } from "../../src/session/tools.js"; + +describe("SessionTools ACP release", () => { + afterEach(() => vi.restoreAllMocks()); + + it.each(["ok", "error"])( + "reads the replacement kernel after queued work and releases input on %s", + async (status) => { + const idle = createDeferred(); + const events = createDeferred(); + const sequence: string[] = []; + const previous = new IpythonKernelProvisioner("/workspace"); + const current = new IpythonKernelProvisioner("/workspace"); + let provisioner = previous; + const oldManager = vi.spyOn(previous, "manager", "get"); + const execute = vi.fn(async (_code: string) => { + sequence.push("execute"); + return { status, stderr: "close failed" }; + }); + vi.spyOn(current, "manager", "get").mockReturnValue({ isRunning: true, execute } as unknown as KernelClient); + const host: SessionToolsHost = { + cwd: "/workspace", + resourceLoader: { + getSystemPrompt: () => undefined, + getAppendSystemPrompt: () => [], + getAgentsFiles: () => ({ agentsFiles: [] }), + }, + getExtensionRunner: () => { + throw new Error("unchanged registry should not rebind extensions"); + }, + getSessionFile: () => undefined, + getModelVisibleSkills: () => [], + getDepth: () => 0, + getMaxDepth: () => 3, + getParentAgent: () => undefined, + getMcpManager: () => ({ + getAcpServers: () => [], + getEnabledPersistentGenericServers: () => [], + canReleaseAcpServers: () => true, + replaceAcpServers: () => false, + }), + getProvisioner: () => provisioner, + getActiveToolNames: () => [], + setActiveToolsByName: () => {}, + getActiveTools: () => [], + setActiveTools: () => {}, + setSystemPrompt: () => {}, + isStreaming: () => false, + rebuildRuntime: () => { + throw new Error("release must preserve the live notebook"); + }, + acquireInputPause: () => { + sequence.push("pause"); + return { + release: () => { + sequence.push("release"); + }, + }; + }, + waitForAgentIdle: () => { + sequence.push("idle"); + return idle.promise; + }, + getEventQueue: () => { + sequence.push("events"); + return events.promise; + }, + }; + const tools = new SessionTools(host, {}); + const release = tools.releaseAcpMcpServers("owner", ["server", "server"]); + expect(sequence).toEqual(["pause", "idle"]); + provisioner = current; + idle.resolve(); + await Promise.resolve(); + expect(sequence).toEqual(["pause", "idle", "events"]); + expect(execute).not.toHaveBeenCalled(); + events.resolve(); + if (status === "error") await expect(release).rejects.toThrow("close failed"); + else await release; + expect(oldManager).not.toHaveBeenCalled(); + expect(sequence).toEqual(["pause", "idle", "events", "execute", "release"]); + expect(execute).toHaveBeenCalledOnce(); + expect(execute.mock.calls[0]![0]).toContain('_prime_mcp_names = ["server"]'); + }, + ); +}); diff --git a/packages/coding-agent/test/suite/agent-session-compaction.test.ts b/packages/coding-agent/test/suite/agent-session-compaction.test.ts index 3d6d6cef83..aa6bf9dd8b 100644 --- a/packages/coding-agent/test/suite/agent-session-compaction.test.ts +++ b/packages/coding-agent/test/suite/agent-session-compaction.test.ts @@ -11,6 +11,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { convertToLlm } from "../../src/core/messages.js"; import { getLocalHarnessStateDir, loadHarnessState, saveHarnessState } from "../../src/core/refinement/index.js"; import { SessionManager } from "../../src/core/session-manager.js"; +import { IpythonKernelProvisioner } from "../../src/core/tools/ipython.js"; import type { SessionCompaction } from "../../src/session/compaction.js"; import type { SessionContinuation } from "../../src/session/continuation.js"; import { createHarness, getMessageText, type Harness } from "./harness.js"; @@ -101,21 +102,17 @@ describe("AgentSession compaction characterization", () => { await harness.session.prompt("one"); await harness.session.prompt("two"); - const pruneOversizedVariables = vi.fn(async () => ["large_text"]); - const listNamespaceNames = vi.fn(async () => ["small_value"]); - const internals = harness.session as unknown as { _ipythonKernelProvisioner?: unknown }; - const previousProvisioner = internals._ipythonKernelProvisioner; - internals._ipythonKernelProvisioner = { - hasRunningKernel: true, - pruneOversizedVariables, - listNamespaceNames, - }; - let result!: Awaited>; - try { - result = await harness.session.compact(); - } finally { - internals._ipythonKernelProvisioner = previousProvisioner; - } + const hasRunningKernel = vi + .spyOn(IpythonKernelProvisioner.prototype, "hasRunningKernel", "get") + .mockReturnValue(true); + const pruneOversizedVariables = vi + .spyOn(IpythonKernelProvisioner.prototype, "pruneOversizedVariables") + .mockResolvedValue(["large_text"]); + const listNamespaceNames = vi + .spyOn(IpythonKernelProvisioner.prototype, "listNamespaceNames") + .mockResolvedValue(["small_value"]); + const result = await harness.session.compact(); + hasRunningKernel.mockRestore(); const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction"); expect(pruneOversizedVariables).toHaveBeenCalledOnce(); diff --git a/packages/coding-agent/test/suite/regressions/2002-acp-mcp-native-tools.test.ts b/packages/coding-agent/test/suite/regressions/2002-acp-mcp-native-tools.test.ts index 862231173d..dc1b5ec2c7 100644 --- a/packages/coding-agent/test/suite/regressions/2002-acp-mcp-native-tools.test.ts +++ b/packages/coding-agent/test/suite/regressions/2002-acp-mcp-native-tools.test.ts @@ -106,11 +106,9 @@ describe("PR 2002 ACP MCP native tools", () => { expect(harness.session.getActiveToolNames()).toContain("mcp_call_task"); expect(harness.session.systemPrompt).not.toContain('await mcp.list_tools("task")'); - const beforeReload = Reflect.get(harness.session, "_toolDefinitions") as Map; - const originalCallTool = beforeReload.get("mcp_call_task"); + const originalCallTool = harness.session.getToolDefinition("mcp_call_task"); await harness.session.reload(); - const afterReload = Reflect.get(harness.session, "_toolDefinitions") as Map; - expect(afterReload.get("mcp_call_task")).not.toBe(originalCallTool); + expect(harness.session.getToolDefinition("mcp_call_task")).not.toBe(originalCallTool); await harness.session.releaseAcpMcpServers("owner-a", ["task"]); expect(harness.session.getAllTools().map((tool) => tool.name)).not.toContain("mcp_call_task"); diff --git a/packages/coding-agent/test/suite/regressions/2053-background-bash-passivation.test.ts b/packages/coding-agent/test/suite/regressions/2053-background-bash-passivation.test.ts index 4da5a7745d..8377a52fb5 100644 --- a/packages/coding-agent/test/suite/regressions/2053-background-bash-passivation.test.ts +++ b/packages/coding-agent/test/suite/regressions/2053-background-bash-passivation.test.ts @@ -12,6 +12,7 @@ import { import { ASYNC_BASH_COMPLETION_CUSTOM_TYPE } from "../../../src/core/messages.js"; import { canEvictWorker, canPassivateSession } from "../../../src/core/session-action-store.js"; import { IpythonKernelProvisioner } from "../../../src/core/tools/ipython.js"; +import type { SessionKernel } from "../../../src/session/kernel.js"; import { createHarness, type Harness } from "../harness.js"; const runtimeDir = resolve(__dirname, "../../../../../prime-agent-runtime"); @@ -19,7 +20,7 @@ const python = resolve(runtimeDir, ".venv/bin/python"); const describeRuntime = existsSync(python) ? describe : describe.skip; interface KernelSession { - _ipythonKernelProvisioner?: IpythonKernelProvisioner; + _kernel: SessionKernel; _createKernelHostHandlers(): HostRequestHandlers; } @@ -76,7 +77,7 @@ describeRuntime("#2053 background kernel bash residency", () => { }); const provisioner = new IpythonKernelProvisioner(harness.tempDir); vi.spyOn(provisioner, "manager", "get").mockReturnValue(manager); - internals._ipythonKernelProvisioner = provisioner; + internals._kernel.provisioner = provisioner; return { session, kernel: manager }; } diff --git a/packages/coding-agent/test/suite/regressions/4530-ipython-state-restore-message.test.ts b/packages/coding-agent/test/suite/regressions/4530-ipython-state-restore-message.test.ts index 9a790cc2f1..7968c8ccd6 100644 --- a/packages/coding-agent/test/suite/regressions/4530-ipython-state-restore-message.test.ts +++ b/packages/coding-agent/test/suite/regressions/4530-ipython-state-restore-message.test.ts @@ -2,7 +2,6 @@ import type { AgentTool } from "@earendil-works/pi-agent-core"; import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai"; import { Type } from "typebox"; import { afterEach, beforeAll, describe, expect, it, vi } from "vitest"; -import type { RestoreResult } from "../../../src/core/kernel/state-snapshot.js"; import { type CustomMessage, IPYTHON_STATE_RESTORED_CUSTOM_TYPE, @@ -13,10 +12,11 @@ import { isInjectedPromptMessage, } from "../../../src/modes/interactive/components/injected-prompt-message.js"; import { initTheme } from "../../../src/modes/interactive/theme/theme.js"; +import type { SessionKernel } from "../../../src/session/kernel.js"; import { conversationMessages, createHarness, getMessageText, getUserTexts, type Harness } from "../harness.js"; type StateRestoreHost = { - _onIpythonStateRestored(result: RestoreResult): void; + _kernel: SessionKernel; }; function stripAnsi(text: string): string { @@ -79,7 +79,7 @@ describe("ENG-4530 IPython state restore message", () => { const firstPrompt = harness.session.prompt("start"); await toolStarted; - (harness.session as unknown as StateRestoreHost)._onIpythonStateRestored({ + (harness.session as unknown as StateRestoreHost)._kernel.onStateRestored({ restored: ["alpha", "beta"], failed: [], path: "/tmp/kernel-state.dill", diff --git a/packages/coding-agent/test/suite/regressions/5939-child-owner-boundaries.test.ts b/packages/coding-agent/test/suite/regressions/5939-child-owner-boundaries.test.ts new file mode 100644 index 0000000000..06809abd57 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/5939-child-owner-boundaries.test.ts @@ -0,0 +1,257 @@ +import { fauxAssistantMessage } from "@earendil-works/pi-ai"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import type { AgentSession } from "../../../src/core/agent-session.js"; +import type { CreateRlmSubagentRuntimeOptions, RlmSubagentRegistryEntry } from "../../../src/core/rlm-runtime.js"; +import { IpythonKernelProvisioner } from "../../../src/core/tools/ipython.js"; +import { createHarness, type Harness } from "../harness.js"; + +function deferred() { + let resolve = () => {}; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +describe("ENG-5939 child owner boundaries", () => { + const harnesses: Harness[] = []; + afterEach(async () => { + for (const harness of harnesses.reverse()) { + await harness.session.disposeAsync(); + harness.cleanup(); + } + harnesses.length = 0; + vi.restoreAllMocks(); + }); + + it("dispatches both missing-selector listings through the live facade", async () => { + const parent = await createHarness(); + harnesses.push(parent); + const list = vi.spyOn(parent.session, "listRlmSubagents"); + await expect(parent.session.deleteRlmSubagent("missing-child")).rejects.toThrow("No direct RLM subagent"); + expect(list).toHaveBeenCalledTimes(2); + expect(list.mock.contexts).toEqual([parent.session, parent.session]); + }); + + it("uses each descendant facade when resolving an inactive child", async () => { + const parent = await createHarness(); + const child = await createHarness(); + harnesses.push(child, parent); + parent.session.registerRlmChildSession("retained-child", child.session); + const subagent: RlmSubagentRegistryEntry = { + rlm_child_id: "passive-grandchild", + active_session_id: null, + session_id: null, + session_name: "passive-grandchild", + session_dir: child.tempDir, + status: "completed", + }; + const list = vi.spyOn(child.session, "listRlmSubagents").mockResolvedValue({ subagents: [subagent] }); + const deleteRuntime = vi.fn(async () => {}); + child.session.setSubagentRuntimeHost({ + createRlmSubagentRuntime: async () => { + throw new Error("unexpected runtime creation"); + }, + deleteRlmSubagentRuntime: deleteRuntime, + }); + await expect(parent.session.deleteInactiveRlmSubagent(subagent.rlm_child_id)).resolves.toBe("deleted"); + expect(list.mock.contexts).toEqual([child.session]); + expect(deleteRuntime).toHaveBeenCalledWith(subagent.rlm_child_id, undefined); + }); + + it("dispatches compaction cleanup retries through the current delete facade", async () => { + const parent = await createHarness(); + const child = await createHarness(); + harnesses.push(child, parent); + parent.session.registerRlmChildSession("retry-child", child.session); + vi.spyOn(parent.session, "listRlmSubagents").mockResolvedValueOnce({ + subagents: [ + { + rlm_child_id: "retry-child", + active_session_id: null, + session_id: child.session.sessionId, + session_name: "retry-child", + session_dir: child.tempDir, + status: "completed", + }, + ], + }); + const deleteRuntime = vi.fn(async () => {}); + deleteRuntime.mockRejectedValueOnce(new Error("cleanup failed")); + parent.session.setSubagentRuntimeHost({ + createRlmSubagentRuntime: async () => { + throw new Error("unexpected runtime creation"); + }, + deleteRlmSubagentRuntime: deleteRuntime, + }); + await expect(parent.session.deleteRlmSubagent("retry-child")).rejects.toThrow("cleanup failed"); + const remove = vi.spyOn(parent.session, "deleteRlmSubagent"); + const lifecycle = parent.session as unknown as { + _reapDeletedRlmSubagentRuntimesAfterCompaction(): Promise; + }; + await lifecycle._reapDeletedRlmSubagentRuntimesAfterCompaction(); + expect(remove).toHaveBeenCalledWith("retry-child"); + expect(remove.mock.contexts).toEqual([parent.session]); + expect(deleteRuntime).toHaveBeenCalledTimes(2); + }); + + it("propagates a facade listing rejection before deletion or fallback resolution", async () => { + const parent = await createHarness(); + harnesses.push(parent); + const failure = new Error("listing rejected"); + const list = vi.spyOn(parent.session, "listRlmSubagents").mockRejectedValue(failure); + const remove = vi.fn(async () => {}); + parent.session.setSubagentRuntimeHost({ + createRlmSubagentRuntime: async () => { + throw new Error("unexpected runtime creation"); + }, + deleteRlmSubagentRuntime: remove, + }); + await expect(parent.session.deleteRlmSubagent("missing-child")).rejects.toBe(failure); + expect(list).toHaveBeenCalledOnce(); + expect(remove).not.toHaveBeenCalled(); + }); + + it("reads the live status facade after global max-depth settings finish", async () => { + const parent = await createHarness(); + harnesses.push(parent); + const pendingFlush = deferred(); + vi.spyOn(parent.settingsManager, "flush").mockImplementationOnce(() => pendingFlush.promise); + const setting = parent.session.setRlmMaxDepth(3, { global: true }); + const status = vi + .spyOn(parent.session, "getRlmMaxDepthStatus") + .mockReturnValue({ maxDepth: 7, source: "global" }); + expect(status).not.toHaveBeenCalled(); + pendingFlush.resolve(); + await expect(setting).resolves.toEqual({ maxDepth: 7, source: "global", globalSaved: true }); + expect(status.mock.contexts).toEqual([parent.session]); + expect(parent.session.rlmMaxDepth).toBe(3); + }); + + it("starts parent kernel disposal synchronously when there are no children", async () => { + const parent = await createHarness(); + harnesses.push(parent); + const kernel = deferred(); + const order: string[] = []; + const lifecycle = parent.session as unknown as { + _disposeAsyncOnce(kernelSnapshot: boolean): Promise; + }; + const dispose = vi.spyOn(IpythonKernelProvisioner.prototype, "dispose").mockImplementation(async (options) => { + const snapshot = options?.snapshot; + expect(snapshot).toBe(false); + order.push("kernel"); + await kernel.promise; + }); + parent.session.registerDisposeCallback(() => { + order.push("callback"); + }); + const disposing = lifecycle._disposeAsyncOnce(false); + order.push("caller"); + expect(order).toEqual(["kernel", "caller"]); + kernel.resolve(); + await disposing; + expect(order).toEqual(["kernel", "caller", "callback"]); + dispose.mockRestore(); + }); + + it("reads the runtime host and inherited settings after asynchronous name preflight", async () => { + const preflight = deferred(); + const child = await createHarness(); + harnesses.push(child); + child.setResponses([fauxAssistantMessage("child result")]); + const initialCreate = vi.fn(async () => ({ session: child.session })); + const checkName = vi.fn(() => preflight.promise); + const parent = await createHarness({ + persistSession: true, + rlmMaxDepth: 2, + agentMessageController: { + listAgents: () => ({ agents: [] }), + sendAgentMessage: async () => { + throw new Error("unexpected explicit reply"); + }, + assertSessionNameAvailable: checkName, + }, + subagentRuntimeHost: { createRlmSubagentRuntime: initialCreate, deleteRlmSubagentRuntime: async () => {} }, + }); + harnesses.push(parent); + parent.setResponses([fauxAssistantMessage("parent consumed result")]); + const spawning = parent.session.runRlmChild("child task", { name: "live-settings-worker" }); + expect(checkName).toHaveBeenCalledOnce(); + const create = vi.fn(async (_options: CreateRlmSubagentRuntimeOptions) => ({ session: child.session })); + parent.session.setSubagentRuntimeHost({ + createRlmSubagentRuntime: create, + deleteRlmSubagentRuntime: async () => {}, + }); + await parent.session.setRlmMaxDepth(3); + preflight.resolve(); + const spawned = await spawning; + await parent.session.waitForRlmQuiescence(); + expect(initialCreate).not.toHaveBeenCalled(); + expect(create).toHaveBeenCalledOnce(); + expect(create.mock.calls[0]?.[0]).toMatchObject({ + rlmDepth: 1, + rlmMaxDepth: 3, + sessionName: "live-settings-worker", + }); + expect(parent.session.getRlmChildSession(spawned.rlm_child_id)).toBe(child.session); + }); + + it("uses the current host for completion and release after delayed runtime publication", async () => { + const publication = deferred(); + const child = await createHarness(); + harnesses.push(child); + child.setResponses([fauxAssistantMessage("child result")]); + const originalComplete = vi.fn(() => true); + const create = vi.fn(async () => { + await publication.promise; + return { session: child.session }; + }); + const parent = await createHarness({ + persistSession: true, + subagentRuntimeHost: { + createRlmSubagentRuntime: create, + completeRlmSubagentRuntime: originalComplete, + deleteRlmSubagentRuntime: async () => {}, + }, + }); + harnesses.push(parent); + parent.setResponses([fauxAssistantMessage("parent consumed result")]); + const spawned = await parent.session.runRlmChild("delayed child"); + expect(create).toHaveBeenCalledOnce(); + const originalRegister = parent.session.registerRlmChildSession; + const register = vi.spyOn(parent.session, "registerRlmChildSession").mockImplementation(function ( + this: AgentSession, + id, + session, + unsubscribe, + ) { + expect(this).toBe(parent.session); + return originalRegister.call(this, id, session, unsubscribe); + }); + const complete = vi.fn(() => false); + const release = vi.fn(async () => { + await child.session.disposeAsync(); + }); + parent.session.setSubagentRuntimeHost({ + createRlmSubagentRuntime: async () => { + throw new Error("unexpected second create"); + }, + completeRlmSubagentRuntime: complete, + releaseRlmSubagentRuntime: release, + deleteRlmSubagentRuntime: async () => {}, + }); + publication.resolve(); + await parent.session.waitForRlmQuiescence(); + expect(originalComplete).not.toHaveBeenCalled(); + expect(register).toHaveBeenCalledWith(spawned.rlm_child_id, child.session); + expect(complete).toHaveBeenCalledWith(spawned.rlm_child_id, child.session); + expect(release).toHaveBeenCalledOnce(); + expect(release).toHaveBeenCalledWith( + { session: child.session }, + expect.objectContaining({ id: spawned.rlm_child_id }), + "error", + ); + expect(parent.session.getRlmChildSession(spawned.rlm_child_id)).toBeUndefined(); + expect(parent.session.getRlmChildSnapshots()).toEqual([]); + }); +}); diff --git a/packages/coding-agent/test/suite/regressions/5939-runtime-facade-boundaries.test.ts b/packages/coding-agent/test/suite/regressions/5939-runtime-facade-boundaries.test.ts new file mode 100644 index 0000000000..d998a8b0cb --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/5939-runtime-facade-boundaries.test.ts @@ -0,0 +1,207 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import type { AgentSession } from "../../../src/core/agent-session.js"; +import { AuthStorage } from "../../../src/core/auth-storage.js"; +import type { ExtensionAPI } from "../../../src/core/extensions/index.js"; +import { createDeferred, type HostRequestHandlers } from "../../../src/core/kernel/index.js"; +import { ModelRegistry } from "../../../src/core/model-registry.js"; +import { createSyntheticSourceInfo } from "../../../src/core/source-info.js"; +import { IpythonKernelProvisioner } from "../../../src/core/tools/ipython.js"; +import { createHarness, type Harness } from "../harness.js"; + +describe("ENG-5939 runtime facade boundaries", () => { + const harnesses: Harness[] = []; + afterEach(async () => { + vi.restoreAllMocks(); + for (const harness of harnesses.splice(0)) { + await harness.session.disposeAsync(); + harness.cleanup(); + } + }); + + it("retains the session receiver for extension shutdown after partial binding and reload", async () => { + let shutdown = () => {}; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("session_start", (_event, ctx) => { + shutdown = () => ctx.shutdown(); + }); + }, + ], + }); + harnesses.push(harness); + const receivers: AgentSession[] = []; + await harness.session.bindExtensions({ + shutdownHandler: function (this: AgentSession) { + receivers.push(this); + this.setActiveToolsByName([]); + }, + }); + shutdown(); + await harness.session.bindExtensions({ onError: () => {} }); + shutdown(); + await harness.session.reload(); + shutdown(); + expect(receivers).toEqual([harness.session, harness.session, harness.session]); + }); + + it("reads current public template and registry getters from extension actions", async () => { + let api!: ExtensionAPI; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + api = pi; + }, + ], + }); + harnesses.push(harness); + const templates = vi.spyOn(harness.session, "promptTemplates", "get").mockReturnValue([ + { + name: "intercepted-template", + description: "live public template", + content: "template", + filePath: "/templates/intercepted.md", + sourceInfo: createSyntheticSourceInfo("", { source: "test" }), + }, + ]); + expect(api.getCommands()).toContainEqual( + expect.objectContaining({ name: "intercepted-template", source: "prompt" }), + ); + expect(templates.mock.contexts).toEqual([harness.session]); + const registry = ModelRegistry.inMemory(AuthStorage.inMemory()); + const auth = vi.spyOn(registry, "hasConfiguredAuth").mockReturnValue(false); + vi.spyOn(harness.session, "modelRegistry", "get").mockReturnValue(registry); + await expect(api.setModel(harness.models[0])).resolves.toBe(false); + expect(auth.mock.contexts).toEqual([registry]); + }); + + it("rejects accessor failures asynchronously without reading later bindings", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const failure = new Error("binding getter failed"); + const onError = vi.fn(); + let binding: Promise | undefined; + expect(() => { + binding = harness.session.bindExtensions({ + get shutdownHandler(): never { + throw failure; + }, + get onError() { + onError(); + return undefined; + }, + }); + }).not.toThrow(); + await expect(binding).rejects.toBe(failure); + expect(onError).not.toHaveBeenCalled(); + }); + + it("keeps original accessor order and earlier binding updates when a later getter fails", async () => { + let shutdown = () => {}; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("session_start", (_event, ctx) => { + shutdown = () => ctx.shutdown(); + }); + }, + ], + }); + harnesses.push(harness); + const previous = vi.fn(); + await harness.session.bindExtensions({ shutdownHandler: previous }); + const reads: string[] = []; + const current = vi.fn(function (this: AgentSession) { + expect(this).toBe(harness.session); + }); + const failure = new Error("later binding getter failed"); + await expect( + harness.session.bindExtensions({ + get shutdownHandler() { + reads.push("shutdown"); + return current; + }, + get onError(): never { + reads.push("onError"); + throw failure; + }, + }), + ).rejects.toBe(failure); + expect(reads).toEqual(["shutdown", "shutdown", "onError"]); + shutdown(); + expect(previous).not.toHaveBeenCalled(); + expect(current).toHaveBeenCalledOnce(); + }); + + it("settles a no-op ACP release before work queued by its caller", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const order: string[] = []; + const released = harness.session.releaseAcpMcpServers("missing-owner", []); + void released.then(() => { + order.push("released"); + }); + queueMicrotask(() => { + order.push("caller"); + }); + await released; + expect(order).toEqual(["released", "caller"]); + }); + + it("runs disposal callbacks at the kernel completion boundary", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const kernel = createDeferred(); + const order: string[] = []; + vi.spyOn(IpythonKernelProvisioner.prototype, "dispose").mockImplementation(() => { + order.push("kernel"); + return kernel.promise; + }); + harness.session.registerDisposeCallback(() => { + order.push("callback"); + }); + const lifecycle = harness.session as unknown as { _disposeAsyncOnce(snapshot: boolean): Promise }; + const disposed = lifecycle._disposeAsyncOnce(false); + order.push("caller"); + expect(order).toEqual(["kernel", "caller"]); + kernel.resolve(); + queueMicrotask(() => { + order.push("after-kernel"); + }); + await disposed; + expect(order).toEqual(["kernel", "caller", "callback", "after-kernel"]); + }); + + it("acknowledges bash completion without delaying work after message admission", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const admitted = createDeferred(); + const lifecycle = harness.session as unknown as { + _promptInjectedMessage(): Promise; + _createKernelHostHandlers(): HostRequestHandlers; + }; + vi.spyOn(lifecycle, "_promptInjectedMessage").mockReturnValue(admitted.promise); + const order: string[] = []; + const completed = lifecycle._createKernelHostHandlers()["bash.completed"]!({ + pid: 123, + command: "echo ready", + exitCode: 0, + }); + void completed.then(() => { + order.push("acknowledged"); + }); + admitted.resolve(); + queueMicrotask(() => { + order.push("first"); + queueMicrotask(() => { + order.push("second"); + queueMicrotask(() => { + order.push("third"); + }); + }); + }); + await completed; + await Promise.resolve(); + expect(order).toEqual(["first", "second", "acknowledged", "third"]); + }); +});