diff --git a/AGENTS.md b/AGENTS.md index 85e64ff351..6d1666bc6a 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -20,6 +20,12 @@ - Never hardcode key checks with, eg. `matchesKey(keyData, "ctrl+x")`. All keybindings must be configurable. Add default to matching object (`DEFAULT_EDITOR_KEYBINDINGS` or `DEFAULT_APP_KEYBINDINGS`) - NEVER modify `packages/ai/src/models.generated.ts` directly. Update `packages/ai/scripts/generate-models.ts` instead. +## Source Organization + +- Follow the [source ownership rules](packages/coding-agent/docs/architecture.md#source-ownership-and-module-boundaries) when adding or extracting modules. The [source map](packages/coding-agent/src/README.md) documents current owners and ordering invariants. +- Before a structural change, identify the feature owner, its state and lifecycle, its public API, and its allowed dependencies. Update the source map when ownership changes. +- Group complete feature responsibilities into reviewable PRs. File size alone does not justify an abstraction or a separate PR. + ## Commands - After code changes (not documentation changes): `npm run check` (get full output, no tail). Fix all errors, warnings, and infos before committing. diff --git a/packages/coding-agent/docs/architecture.md b/packages/coding-agent/docs/architecture.md index 9b6a235d80..394c98acf0 100644 --- a/packages/coding-agent/docs/architecture.md +++ b/packages/coding-agent/docs/architecture.md @@ -85,6 +85,47 @@ sequenceDiagram From the session queue onward, the same execution and persistence path is used when a prompt comes from a heartbeat, cron schedule, goal continuation, autonomous mode, or another agent instead of an attached user. +## Source Ownership and Module Boundaries + +Place code inside the smallest feature that owns its behavior. Promote it outside that feature when it provides an independently useful capability with a clear API and actual consumers across subsystems. This rule applies to new code and incremental extractions; existing paths are not precedents for new exceptions. + +### Choosing a directory + +| Responsibility | Home | +| --- | --- | +| A feature of one session, including its state, policies, persistence adapters, and cleanup | `src/session//` | +| Coordination across session features, such as choosing when a turn compacts, refines, or continues | Session composition or `src/session/turns/` | +| An independent capability used across subsystems, with its own API and dependency boundary | Its own named feature directory outside `session/` | +| Presentation or process coordination specific to a mode | The owning directory under `src/modes/` | + +Use feature names and keep the structure shallow: files directly inside a feature directory are the default. Add another directory only for a coherent subfeature. Application source stays under `src/`; tests, scripts, docs, and build output stay at the package root. `core/` is a transitional location for existing code, not the default home for new shared code. + +Independent testability, a pure function, or a small dependency interface does not require a top-level directory. Multiple importers are evidence to inspect, not a promotion rule: a UI reading a session goal's status does not make goal execution independent of sessions. Do not promote code for hypothetical future reuse. + +### Keeping a feature together + +- Keep a responsibility's state, transitions, cancellation, recovery, and cleanup under one feature owner. Parsing, persistence, and execution may use separate files within that feature when their dependencies differ. +- Keep one authoritative copy of mutable state. Derived views may read it; extracting a module must not introduce a second queue, registry, transcript, or competing lifecycle owner. +- Put feature contracts beside their owner in lightweight modules. Clients may import those contracts without importing controllers or session orchestration. Contract modules must not depend on their feature's runtime implementation. +- Pass only the named operations and state views a module needs. Avoid passing the entire session, exposing mutable maps, or introducing a universal context object. Keep imports acyclic; compose dependencies at the owner that coordinates the features. +- Keep cross-feature ordering visible in the composition layer. A feature may request scheduling or persistence through a narrow interface without owning the scheduler or storage implementation. +- Keep feature tests with the corresponding feature in `test/` where practical. Integration tests continue to cover behavior through the public session or mode API. + +For example, goal state, accounting, persistence, command parsing, and goal-specific continuation live in `session/goals/`. General turn selection and coordination stay in `session/turns/`. Goal contracts stay with goals even when the UI and protocol types consume them. Kernel transport and process APIs are separate from the session-specific code that creates, replaces, and disposes its kernel. The [source map](../src/README.md) records the current implementation and its ordering invariants. + +### Reviewing a structural change + +Before choosing files, identify four things in the change description: + +1. The feature that owns the behavior and why this is its directory. +2. The state and lifecycle operations that must move together. +3. The public operations and contracts, including their consumers. +4. The allowed dependencies and the layer responsible for cross-feature ordering. + +Group complete feature responsibilities into a few substantial, reviewable changes. Avoid one file per method, arbitrary line-count targets, or new frameworks introduced only to make files smaller. Preserve behavior during extraction, including event order, cancellation, persistence, and public contracts; functional changes should be identified and reviewed explicitly. Keep protocol compatibility requirements in force whenever a wire shape changes. + +Validate affected behavior with focused tests and the repository's required checks. Record missing-environment skips. File moves alone do not establish lower memory use, faster execution, or smaller bundles; performance claims require measurements of the affected workload. + ## Detailed Architecture - [Agent Connection Architecture](agent-connection.md) explains the client/runtime boundary, snapshots, replay, and reconnect behavior. diff --git a/packages/coding-agent/src/README.md b/packages/coding-agent/src/README.md index 63072ebb26..92b7d3f337 100644 --- a/packages/coding-agent/src/README.md +++ b/packages/coding-agent/src/README.md @@ -1,17 +1,18 @@ # Source organization -`src/` contains application source; tests, scripts, docs, examples, and build output stay at the package root. Put feature folders directly under `src/`, such as `goals/`, `session/`, and `kernel/`. Add another level only when a feature has distinct subparts that benefit from being grouped. +`src/` contains application source; tests, scripts, docs, examples, and build output stay at the package root. Follow the [source ownership rules](../docs/architecture.md#source-ownership-and-module-boundaries) when choosing module boundaries. Session features, including goals, live below `session/`; each feature groups its state, dependencies, and lifecycle. -`core/` currently contains most execution logic. Migrate its responsibilities into sibling feature folders as they are extracted. `AgentSession` remains the public entry point and coordinates work across features. Each feature owner keeps its state and transitions together and receives only the dependencies it uses. +`core/` currently contains most execution logic. Migrate responsibilities into their owning feature folders as their boundaries are established. `AgentSession` remains the public entry point and coordinates work across features. Each feature owner keeps its state and transitions together and receives only the dependencies it uses. ## Goals | File | Responsibility | | --- | --- | -| `goals/controller.ts` | Goal transitions, token and time accounting, continuation counts, and rollback checkpoints. | -| `goals/persistence.ts` | Reading the selected branch, flushing goal records, and deciding whether a branch can receive an initial goal. | -| `goals/commands.ts` | Parsing `/goal` arguments into typed commands. | -| `core/goals.ts` | Shared goal types, validation, serialization, and context-message formatting used by session clients. | +| `session/goals/controller.ts` | Goal transitions, token and time accounting, continuation counts, and rollback checkpoints. | +| `session/goals/persistence.ts` | Reading the selected branch, flushing goal records, and deciding whether a branch can receive an initial goal. | +| `session/goals/commands.ts` | Parsing `/goal` arguments into typed commands. | +| `session/goals/continuation.ts` | Goal continuation admission, budget notices, child-wait coordination, and rollback. | +| `session/goals/contracts.ts` | Lightweight goal types, validation, serialization, and context-message formatting used by session clients. | The controller depends on a load/save interface, an update callback, and a clock. It does not receive `AgentSession`, the agent loop, a kernel, or a UI object. Its state is read-only to callers; mutations go through named operations. @@ -23,17 +24,35 @@ Three ordering rules matter during future extractions: - Capture completion usage, clear stale queued goal context, and then persist and publish completion. The explicit completion callback preserves this order. - A continuation rejected by new input restores both goal state and the accounting clock. Deferred child-work admission preserves the existing clock while restoring goal state. -The shared goal types and message formatting remain in `core/goals.ts` during this extraction; their existing consumers can migrate together in a later change. The public goal payload and persisted `thread_goal_state` format remain shared contracts. Internal organization does not require a new daemon command or schema. +Goal state belongs to a session branch, so the controller, persistence adapter, and continuation owner share `session/goals/`. UI and protocol consumers import the lightweight contracts without loading the controller. The public goal payload and persisted `thread_goal_state` format remain unchanged; the internal move does not require a new daemon command or schema. ## Extending this structure -Use the same ownership rule for the next extraction: move a responsibility's fields, transitions, and cleanup together. Keep request parsing and storage adapters separate when they have independent dependencies. Avoid generic helper folders, modules that receive the entire session, and duplicate copies of feature state. +Apply the architecture guide's placement and dependency rules to each extraction. This document records the resulting owners and ordering invariants; update it when those boundaries change. + +## Session feature folders + +| Folder | Ownership | +| --- | --- | +| `session/input/` | Submission normalization, admission, scheduling, commit fencing, action queues, delivery, and recovery. | +| `session/goals/` | Goal state and contracts, accounting, persistence, command parsing, and goal-specific continuation. | +| `session/turns/` | Turn preparation and execution, session commands, ordered events, retry, and autonomous/shared continuation. | +| `session/compaction/` | Session compaction lifecycle and execution. | +| `session/refinement/` | Refinement planning/application lifecycle, automatic review, and execution. | +| `session/context/` | Pending context, harness context, transcript views, branch navigation, and export. | +| `session/children/` | Child records, runtime creation, execution, projections, and usage accounting. | +| `session/kernel/` | Kernel lifecycle, environment, host-handler composition, and stateless host-request adapters. | +| `session/models/` | Model selection, thinking preferences, and authenticated availability. | +| `session/tools/` | Tool selection and shell-command execution. | +| `session/extensions/` | Extension bindings, resource reload, and tool hooks. | + +`session/prepared-actions.ts` remains the shared action and recovery contract used across input, turns, and context. Shared continuation belongs to `session/turns/`, including continuation after compaction. Child usage belongs to `session/children/` because its accounting and cleanup follow child records. Core compaction/refinement algorithms and their public exports remain in `core/` for their broader callers. ## Session input scheduling -`session/input-scheduler.ts` owns the serialized pump, its preparation epoch, pause leases, and abort/restart suspension. It receives two callbacks: whether the session has work eligible for scheduling, and the operation that runs that work. The scheduler exposes read-only state and named operations; callers cannot change its pause sets or scheduling flags. +`session/input/input-scheduler.ts` owns the serialized pump, its preparation epoch, pause leases, and abort/restart suspension. It receives two callbacks: whether the session has work eligible for scheduling, and the operation that runs that work. The scheduler exposes read-only state and named operations; callers cannot change its pause sets or scheduling flags. -The existing `ActionStore` in `core/session-action-store.ts` owns queued actions, their transitions, and delivery/completion tickets. `session/input-dispatcher.ts` selects and batches those actions, reconciles durable delivery after dispatch, rolls undelivered work back, and settles completion or failure. `AgentSession` supplies turn execution and session-command operations and coordinates goals, child agents, and compaction. The dispatcher shares the existing `ActionStore`; it does not create a second queue or copy the transcript. +The existing `ActionStore` in `core/session-action-store.ts` owns queued actions, their transitions, and delivery/completion tickets. `session/input/input-dispatcher.ts` selects and batches those actions, reconciles durable delivery after dispatch, rolls undelivered work back, and settles completion or failure. `AgentSession` supplies turn execution and session-command operations and coordinates goals, child agents, and compaction. The dispatcher shares the existing `ActionStore`; it does not create a second queue or copy the transcript. Preserve these distinctions when extending the scheduler: @@ -45,7 +64,7 @@ Preserve these distinctions when extending the scheduler: ## Session commit coordination -`session/commit-fence.ts` owns the FIFO commit queue, its current owner and waiters, asynchronous reentrancy context, and disposal signal. Prompt dispatch, session commands, and branch navigation acquire a lease and run their critical section within its owner context. The session still decides when admission is allowed and when to release the lease. +`session/input/commit-fence.ts` owns the FIFO commit queue, its current owner and waiters, asynchronous reentrancy context, and disposal signal. Prompt dispatch, session commands, and branch navigation acquire a lease and run their critical section within its owner context. The session still decides when admission is allowed and when to release the lease. - Reentrant work shares the current owner's lease; releasing that nested lease does not release the outer operation. An asynchronous callback from an old owner must queue behind the current owner. - Cancelling a waiter rejects it promptly but retains its place in the promise chain until its predecessor releases. Later operations cannot overtake that predecessor. @@ -58,7 +77,7 @@ The input dispatcher preserves selection and settlement ordering. Batches includ ## Session shell commands -`session/bash.ts` owns shell-command execution, abort controllers, the user-command slot, abort requests during extension dispatch, and deferred transcript output. Its host supplies current shell settings and working directory, extension interception, event delivery, transcript append, and session scheduling notifications. These callbacks read the current runtime so rebuilding extensions or changing settings does not retain stale dependencies. +`session/tools/bash.ts` owns shell-command execution, abort controllers, the user-command slot, abort requests during extension dispatch, and deferred transcript output. Its host supplies current shell settings and working directory, extension interception, event delivery, transcript append, and session scheduling notifications. These callbacks read the current runtime so rebuilding extensions or changing settings does not retain stale dependencies. `AgentSession` keeps its public shell methods and the cross-feature decision about when to flush deferred output. It also appends messages to agent state before persistence and schedules queued input after the agent becomes idle. The shell owner does not receive the session, kernel, agent loop, or storage manager. @@ -73,7 +92,7 @@ Execution and recording callbacks preserve dispatch through the public session m ## Session retry handling -`session/retry.ts` owns retry attempts, backoff cancellation, retry completion, and authentication-failure tracking. The session reports assistant and agent completion at their existing points in event processing. The retry owner receives current settings, model authentication operations, context inspection, and named operations for continuing or ending a turn. +`session/turns/retry.ts` owns retry attempts, backoff cancellation, retry completion, and authentication-failure tracking. The session reports assistant and agent completion at their existing points in event processing. The retry owner receives current settings, model authentication operations, context inspection, and named operations for continuing or ending a turn. - Reserve retry completion synchronously when receiving `agent_end`, before asynchronous event processing. Callers waiting for retry must observe the same pending work. - Resolve completion before notifying waiters and scheduling queued input. Generation checks keep a rejected continuation from terminating a later retry. @@ -82,22 +101,48 @@ Execution and recording callbacks preserve dispatch through the public session m ## Turn preparation and action records -`session/turn-preparation.ts` contains the execution policies for direct, queued, injected, and custom-triggered turns and the ordered preparation pipeline. `TurnPreparer` receives six operations for validation, pending shell output, model selection, compaction, and refinement. The session supplies their implementations and retains transcript dispatch and context rollback. +`session/turns/turn-preparation.ts` contains the execution policies for direct, queued, injected, and custom-triggered turns and the ordered preparation pipeline. `TurnPreparer` receives six operations for validation, pending shell output, model selection, compaction, and refinement. The session supplies their implementations and retains transcript dispatch and context rollback. Preserve the policy differences: direct prompts flush shell output before validation and compact after model selection; queued turns validate before flushing and compact before model selection. Conditional refinement barriers are checked when reached, so a refinement started during preparation is still awaited. Withdrawing prepared work skips the final barrier and commit. The exported `TurnExecutionPolicy` shape remains available from the session facade. `session/prepared-actions.ts` contains prepared action types, delivery records, recovery contracts, input copying, action factories, and queue projections. It has no session dependency. Primary messages retain their identity for durable-delivery checks; separately stored input blocks and prefix messages retain their existing copy behavior. Recovery format version 1 and the public exports from `AgentSession` stay unchanged. +## Session input and turns + +| File | Responsibility | +| --- | --- | +| `session/input/submission-normalization.ts` | Copying and validating submission content, options, and provenance. | +| `session/input/prompt-submission.ts` | Prompt preparation and admission, steering, follow-ups, custom/user messages, and background shell completion messages. | +| `session/input/message-delivery.ts` | Accepted agent-message receipts, completion settlement, and restoration of late Python messages. | +| `session/input/input-admission.ts` | Readiness checks and admission predicates. | +| `session/input/input-checkpoints.ts` | Checkpoint waiters, notification, cancellation, input-dispatch barriers, and headless waiting. | +| `session/input/action-queue.ts` | Queue operations and projections over the existing ActionStore. | +| `session/input/action-recovery.ts` | Capturing and restoring pending input when the session runtime changes. | +| `session/turns/turn-execution.ts` | Executing direct, queued, injected, and custom-triggered turns. | +| `session/turns/command-execution.ts` | Executing queued session commands and recording their outcomes. | +| `session/turns/events.ts` | Ordered agent-event processing, listener delivery, and transcript/accounting coordination. | +| `session/context/pending-context.ts` | Pending messages, notices, and retention for the next turn. | +| `session/goals/continuation.ts` | Goal continuation admission, budget notices, child-wait coordination, and rollback. | +| `session/turns/autonomous-continuation.ts` | Autonomous continuation messages, snapshots, and rollback. | +| `session/turns/turn-policy.ts` | Turn stopping, threshold compaction, and continuation decisions from current session state. | + +Input has one durable ActionStore, one scheduling pump, one commit fence, and one ordered event queue. Queue operations, admission, preparation, execution, and recovery use these same owners; they do not maintain parallel queues or transcripts. Dependencies are named operations and small state views. Callbacks read the current model, controllers, and runtime where the original operation did. + +- Pending-context state changes through named synchronous operations. Raw rollback restores the same messages without waking input; recovery restores copied envelopes and flushes deferred work. Message identity, shared details, and the existing copy boundaries matter for delivery and rollback. +- Goal and autonomous continuation state stays with its owner. Consuming a threshold message or harness digest and rearming it are explicit transitions. Preserve the existing Boolean abort/child-wait timing and message-keyed snapshot identity. +- Public session entry points remain live dispatch boundaries for extensions and callers. Delegates retain callback receivers and their existing synchronous or asynchronous return behavior. +- Session startup, abort, disposal, pause release, and work after compaction remain coordinated in `AgentSession` because they span several owners. Keep those operation sequences visible instead of introducing a general lifecycle framework. + ## Session context | File | Responsibility | | --- | --- | -| `session/compaction.ts` | Manual and automatic compaction lifecycle, pending requests, cancellation, thresholds, and overflow recovery. | -| `session/compaction-execution.ts` | Summary generation, extension interception, request accounting, persistence, and context rebuild ordering. | -| `session/refinement.ts` | Refinement admission, planning and application barriers, serialized plan ownership, and disposal drains. | -| `session/auto-refinement.ts` | Review triggers, cooldowns, pending reviews, timers, and automatic operation cleanup. | -| `session/refinement-execution.ts` | Planning against current dependencies, applying harness edits, and persisting outcomes and notices. | -| `session/continuation.ts` | Resuming work after compaction, settlement, cancellation, and ownership of continuation messages. | +| `session/compaction/compaction.ts` | Manual and automatic compaction lifecycle, pending requests, cancellation, thresholds, and overflow recovery. | +| `session/compaction/compaction-execution.ts` | Summary generation, extension interception, request accounting, persistence, and context rebuild ordering. | +| `session/refinement/refinement.ts` | Refinement admission, planning and application barriers, serialized plan ownership, and disposal drains. | +| `session/refinement/auto-refinement.ts` | Review triggers, cooldowns, pending reviews, timers, and automatic operation cleanup. | +| `session/refinement/refinement-execution.ts` | Planning against current dependencies, applying harness edits, and persisting outcomes and notices. | +| `session/turns/continuation.ts` | Resuming work after compaction, settlement, cancellation, and ownership of continuation messages. | Each owner keeps its mutable state and cleanup together. Typed host operations connect the owners to current model, authentication, extensions, storage, and scheduling. `AgentSession` composes them and retains public methods, events, and decisions that cross features, including goal and autonomous continuation admission. The summary algorithms and harness storage remain in their existing `core/` feature modules. @@ -112,13 +157,13 @@ Preserve these boundaries when changing context behavior: | 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. | +| `session/children/children.ts` | Child registry, admission, publication, deletion retries, cancellation, quiescence, retention, and cleanup. | +| `session/children/child-run.ts` | Detached child execution, publication barriers, terminal state, and event attribution. | +| `session/children/child-runtime.ts` | Inline child construction and child-specific directory creation. | +| `session/children/child-state.ts` | Depth and maximum-depth settings, parent replies, and recap state. | +| `session/children/child-usage.ts` | Child usage attribution, origin batches, flush timers, and retry bookkeeping. | +| `session/children/child-projection.ts` | Read-only child list and snapshot projections. | +| `session/children/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. @@ -128,25 +173,48 @@ The registry owns child identity and lifecycle transitions. Execution and usage - 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. +The root kernel directory belongs to `session/kernel/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. | +| `session/tools/tools.ts` | Tool definitions and active selection, allowlists, prompt contributions, and ACP tool updates. | +| `session/extensions/extensions.ts` | Extension runner bindings, resource reload, and extension lifecycle. | +| `session/kernel/kernel.ts` | Kernel construction, snapshot restoration, prewarming, and disposal. | +| `session/kernel/kernel-environment.ts` | Kernel provisioning environment and root or ephemeral session directories. | +| `session/kernel/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. +## Models, history, and host requests + +| File | Responsibility | +| --- | --- | +| `session/models/model-selection.ts` | Model and thinking preferences, authenticated availability, preflight checks, cycling, and child-model selection. | +| `session/context/history-navigation.ts` | Switching sessions, forking, and navigating branches with their existing barriers. | +| `session/context/context-view.ts` | Context usage, session statistics, and tree views over live transcript and child usage. | +| `session/context/harness-context.ts` | Harness changes, digest consumption, and context for subsequent turns. | +| `session/context/export.ts` | Session export using current model and extension rendering dependencies. | +| `session/kernel/heartbeat-host-requests.ts` | Heartbeat request validation and controller operations. | +| `session/kernel/message-host-requests.ts` | Message request validation and controller operations. | +| `session/kernel/observe-host-requests.ts` | Observation request validation, controller operations, and result encoding. | + +The three host-request modules are stateless adapters. Kernel handler composition still calls the public session methods. Heartbeat and observation requests capture their controller on entry; messaging reads its controller at each operation. Compaction request interpretation belongs to the compaction owner. + +Model selection captures the parent model before awaiting authenticated availability. Preserve the unauthenticated parent-model fast path, expired-credential filtering, original validation errors, and public model/thinking getter dispatch. Context and export read existing records; neither owns a second transcript. + +Context views own aggregation and derived usage memos. Child usage owns the adjustment for child spend not yet indexed in the transcript. The view calls that existing operation with the same usage object and entries. + +Feature contracts live with their owner, including prompt options and session events. The facade re-exports the existing public types and retains session construction configuration. Constructor-only references are readonly; replaceable runtime bindings stay mutable. A host interface should expose only the state and operations its consumer needs. + +Integration regressions in `test/suite/regressions/` cover model/history boundaries, prompt and message dispatch, host-controller capture, callback receivers, and admission behavior. `test/session/pending-context.test.ts` covers pending-context identities and recovery. Existing queue, action, goal, autonomous, branch, extension, and kernel suites cover composition through the public session API. + ## 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`. +Controller tests live in `test/session/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`. Scheduler and commit-fence tests live in `test/session/`. Existing queue, action-contract, action-race, and compaction suites cover the integration with `AgentSession`, including pause, cancellation, restart, branch navigation, and goal continuation. diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 72dd9811b0..e46f24bfcc 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -1,112 +1,81 @@ -import { AsyncLocalStorage } from "node:async_hooks"; -import { randomUUID } from "node:crypto"; -import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs"; -import { dirname, resolve } from "node:path"; import type { Agent, AgentContext, - AgentEvent, AgentMessage, AgentState, AgentTool, - GetContinuationMessagesContext, - ShouldStopAfterTurnContext, ThinkingLevel, } from "@earendil-works/pi-agent-core"; -import type { - Api, - AssistantMessage, - ImageContent, - Model, - ServiceTier, - TextContent, - Usage, - UserMessage, -} from "@earendil-works/pi-ai"; -import { - clampThinkingLevel, - cleanupSessionResources, - getSupportedThinkingLevels, - modelsAreEqual, - supportsFastMode, -} from "@earendil-works/pi-ai"; -import { parseGoalSlashCommand } from "../goals/commands.js"; -import { GoalController } from "../goals/controller.js"; -import { createGoalPersistence } from "../goals/persistence.js"; -import { theme } from "../modes/interactive/theme/theme.js"; -import { - type ExecuteBashOptions, - type RunUserBashOptions, - SessionBash, - type SessionBashEvent, -} from "../session/bash.js"; -import { createChildSessionDir, createInlineChildRuntime } from "../session/child-runtime.js"; -import { SessionChildState } from "../session/child-state.js"; +import type { AssistantMessage, Model, ServiceTier } from "@earendil-works/pi-ai"; +import { clampThinkingLevel, cleanupSessionResources, supportsFastMode } from "@earendil-works/pi-ai"; +import { createChildSessionDir, createInlineChildRuntime } from "../session/children/child-runtime.js"; +import { SessionChildState } from "../session/children/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"; +} from "../session/children/child-types.js"; +import { SessionChildUsage } from "../session/children/child-usage.js"; +import { SessionChildren } from "../session/children/children.js"; +import { SessionCompaction } from "../session/compaction/compaction.js"; import { type CompactionExecutionHost, type CompactionExecutionOptions, - CompactionSkippedError, 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"; +} from "../session/compaction/compaction-execution.js"; +import { type ContextViewChild, SessionContextView } from "../session/context/context-view.js"; +import { SessionExport } from "../session/context/export.js"; +import { SessionHarnessContext } from "../session/context/harness-context.js"; +import { SessionHistoryNavigation } from "../session/context/history-navigation.js"; +import { SessionPendingContext } from "../session/context/pending-context.js"; import { - buildPromptContent, - cloneCustomMessage, - cloneQueuedAgentMessage, - createDeliveryRecord, - createPreparedTurnAction, - createSessionCommandAction, - DeferredSessionInputError, - normalizeMessageContent, - type PreparedCommandPayload, - type PreparedPromptPreparation, - type PreparedTurnPayload, - primaryDeliveryRecord, - type QueuedAgentMessage, - type QueuedSessionAction, - queuedAgentMessagePreview, - type RestoredPromptInput, - SESSION_ACTION_RECOVERY_FORMAT_VERSION, - type SessionActionRecoverySnapshot, - SessionInputAdmissionPausedError, - type SessionInputSchedule, - visibleSessionActionProjection, -} 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"; + type ExtensionBindings, + installExtensionToolHooks, + SessionExtensions, +} from "../session/extensions/extensions.js"; +import { SessionGoalContinuation } from "../session/goals/continuation.js"; +import { + createGoalContextMessage, + GOAL_CONTEXT_CUSTOM_TYPE, + GOAL_SKILL_NAME, + type GoalState, +} from "../session/goals/contracts.js"; +import { GoalController } from "../session/goals/controller.js"; +import { createGoalPersistence } from "../session/goals/persistence.js"; +import { SessionActionQueue } from "../session/input/action-queue.js"; +import { SessionActionRecovery } from "../session/input/action-recovery.js"; +import { SessionCommitFence, type SessionCommitLease } from "../session/input/commit-fence.js"; +import { SessionInputAdmission } from "../session/input/input-admission.js"; +import { SessionInputCheckpoints } from "../session/input/input-checkpoints.js"; +import { SessionInputDispatcher } from "../session/input/input-dispatcher.js"; +import { SessionInputScheduler } from "../session/input/input-scheduler.js"; +import { SessionMessageDelivery } from "../session/input/message-delivery.js"; +import { type PromptOptions, SessionPromptSubmission } from "../session/input/prompt-submission.js"; +import { SubmissionNormalizer } from "../session/input/submission-normalization.js"; +import { handleRlmHeartbeatHostRequest } from "../session/kernel/heartbeat-host-requests.js"; +import { SessionKernel } from "../session/kernel/kernel.js"; +import { KernelEnvironment } from "../session/kernel/kernel-environment.js"; +import { createSessionKernelHostHandlers } from "../session/kernel/kernel-host-handlers.js"; +import { handleAgentMessageHostRequest } from "../session/kernel/message-host-requests.js"; +import { handleAgentObserveHostRequest } from "../session/kernel/observe-host-requests.js"; +import { SessionModelSelection } from "../session/models/model-selection.js"; +import { type QueuedSessionAction, visibleSessionActionProjection } from "../session/prepared-actions.js"; +import { type AutoRefineReviewer, SessionRefinement } from "../session/refinement/refinement.js"; +import { type ExecuteBashOptions, type RunUserBashOptions, SessionBash } from "../session/tools/bash.js"; +import { SessionTools } from "../session/tools/tools.js"; +import { SessionAutonomousContinuation } from "../session/turns/autonomous-continuation.js"; +import { SessionCommandExecution } from "../session/turns/command-execution.js"; +import { SessionContinuation } from "../session/turns/continuation.js"; +import { SessionEvents } from "../session/turns/events.js"; +import { SessionRetry } from "../session/turns/retry.js"; +import { SessionTurnExecution } from "../session/turns/turn-execution.js"; +import { SessionTurnPolicy } from "../session/turns/turn-policy.js"; +import { TurnPreparer } from "../session/turns/turn-preparation.js"; import { AGENT_MESSAGE_SKILL_NAME, - type AgentSessionMessage, type AgentSessionMessageController, type AgentSessionMessageReceipt, - assertAgentMessageQueueCapacity, - assertDirectAgentMessageTarget, - DEFAULT_AGENT_MESSAGE_MAX_PENDING_PER_SESSION, - isAgentSessionMessage, - isAgentSessionMessagePrompt, - normalizeAgentSessionMessage, - parseAgentSessionMessagePromptId, - startsAgentRun, } from "./agent-messages.js"; import { AGENT_OBSERVE_SKILL_NAME, @@ -114,188 +83,54 @@ import { type AgentObserveController, type AgentObserveListResult, type AgentObserveRecentMessagesResult, - normalizeObserveLimit, - normalizeObserveMaxChars, ORCHESTRATION_HEARTBEAT_SKILL_NAME, } from "./agent-observe.js"; -import { - addLoginGuidanceToAuthError, - formatAuthenticationFailedMessage, - formatNoApiKeyFoundMessage, - formatNoModelSelectedMessage, - isLikelyAuthenticationError, -} from "./auth-guidance.js"; -import { - type AgentAutonomousConfig, - type AgentAutonomousStatus, - type AutonomousRuntimeState, - addAutonomousContinuation, - addAutonomousUsage, - autonomousStatus, - createAutonomousRuntimeState, - isUnlimitedAutonomousLimit, - nextAutonomousContinuation, - refreshAutonomousQualityGates, - setAutonomousEnabled, - setAutonomousLimits, - UNLIMITED_AUTONOMOUS_LIMIT, -} from "./autonomous.js"; +import type { AgentAutonomousConfig } from "./autonomous.js"; import type { BashResult } from "./bash-executor.js"; -import { - COMPACT_SKILL_NAME, - type CompactionResult, - calculateContextTokens, - collectEntriesForBranchSummary, - estimateContextTokens, - generateBranchSummary, - prepareCompaction, - shouldCompact, -} from "./compaction/index.js"; -import { - type ContextTreeNode, - type ContextWindowResolver, - computeOwnAndTotalUsage, - loadContextTreeChildFromDisk, - loadContextTreeChildrenFromDisk, -} from "./context-tree.js"; -import type { AgentCronJob, AgentRlmHeartbeatController, AgentRlmHeartbeatStatusUpdate } from "./cron-jobs.js"; -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 { COMPACT_SKILL_NAME, type CompactionResult, calculateContextTokens } from "./compaction/index.js"; +import type { AgentCronJob, AgentRlmHeartbeatController } from "./cron-jobs.js"; import type { - ContextUsage, ExtensionRunner, - InputSource, - MessageEndEvent, - MessageStartEvent, - MessageUpdateEvent, ReplacedSessionContext, - SessionBeforeTreeResult, SessionStartEvent, ToolDefinition, - ToolExecutionEndEvent, - ToolExecutionStartEvent, - ToolExecutionUpdateEvent, ToolInfo, - TreePreparation, - TurnEndEvent, - TurnStartEvent, } from "./extensions/index.js"; -import { - createGoalContextMessage, - GOAL_CONTEXT_CUSTOM_TYPE, - GOAL_CONTEXT_PREVIEW_LABEL, - GOAL_SKILL_NAME, - type GoalHostResponse, - type GoalState, - goalHostResponse, - validateGoalBudget, - validateGoalObjective, -} from "./goals.js"; -import type { HostRequestHandlers, KernelSentAgentMessage } from "./kernel/index.js"; +import type { HostRequestHandlers } from "./kernel/index.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 CustomMessage, - createAsyncBashCompletionMessage, - createHarnessDigestMessage, - createHeartbeatPromptMessage, - createSessionSlashCommandMessage, - createSessionSlashCommandResultMessage, - HARNESS_DIGEST_CUSTOM_TYPE, - type HarnessDigestDetails, - HEARTBEAT_PROMPT_CUSTOM_TYPE, - HEARTBEAT_PROMPT_PREVIEW_LABEL, - isSessionSlashCommandMessage, - type RefinementSource, - RLM_CHILD_FAILURE_CUSTOM_TYPE, - RLM_CHILD_TERMINAL_NOTICE_CUSTOM_TYPE, -} from "./messages.js"; +import { type CustomMessage, createHeartbeatPromptMessage, type RefinementSource } from "./messages.js"; import type { ModelRegistry } from "./model-registry.js"; -import { throwIfPromptAdmissionCancelled } from "./prompt-admission.js"; -import { expandPromptTemplate, type PromptTemplate, parseCommandArgs } from "./prompt-templates.js"; +import type { PromptTemplate } from "./prompt-templates.js"; import { providerRetryPolicy } from "./provider-retry.js"; -import { formatHarnessStateForPrompt, REFINE_SKILL_NAME, type RefinementResult } from "./refinement/index.js"; +import { REFINE_SKILL_NAME, type RefinementResult } from "./refinement/index.js"; import type { ResourceLoader } from "./resource-loader.js"; -import { - type CreateRlmSubagentRuntimeOptions, - findRlmModelMatches, - type RlmCreateSessionResult, - type RlmDeleteSubagentResult, - type RlmFindModelsResult, - type RlmListSubagentsResult, - type RlmSpawnHandle, - type RlmSubagentRuntime, - type SubagentRuntimeHost, +import type { + CreateRlmSubagentRuntimeOptions, + RlmCreateSessionResult, + RlmDeleteSubagentResult, + RlmListSubagentsResult, + RlmSpawnHandle, + RlmSubagentRuntime, + SubagentRuntimeHost, } from "./rlm-runtime.js"; import { SemanticEdgeRecorder, semanticEdgeLedgerPath, wrapStreamFnWithSemanticEdges } from "./semantic-edges.js"; -import { - ActionStore, - type ActionTicket, - canSelectSessionAction, - type DeliveryPolicy, - type DeliveryRecord, - type QueuedMessageLane, - type QueuedMessageMutation, - type QueuedMessageMutationStatus, - queuedMessageLaneDeliveryPolicy, - type RuntimeActivity, - type SessionAction, - type SessionActionSnapshot, - transitionSessionAction, -} from "./session-action-store.js"; -import type { BranchSummaryEntry, SessionContext, SessionEntry } from "./session-manager.js"; -import { - CURRENT_SESSION_VERSION, - getLatestCompactionEntry, - type SessionHeader, - type SessionManager, -} from "./session-manager.js"; -import type { SessionStats } from "./session-stats.js"; +import { ActionStore, type RuntimeActivity } from "./session-action-store.js"; +import type { SessionManager } from "./session-manager.js"; import type { SettingsManager } from "./settings-manager.js"; import { getPythonSkillRuntimeInfo, type Skill } from "./skills.js"; -import { - parseRefineCommandOptions, - parseSessionSlashCommand, - parseSlashCommand, - type SessionSlashCommand, -} from "./slash-commands.js"; import type { BuildSystemPromptOptions } from "./system-prompt.js"; -import { THINKING_LEVELS } from "./thinking-levels.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 AgentSessionEvent = - | AgentEvent - | { - type: "ipython_sent_agent_message"; - toolCallId: string; - message: KernelSentAgentMessage; - } - | { type: "session_action_update"; actions: SessionActionSnapshot } - | SessionCompactionEvent - | { type: "session_info_changed"; name: string | undefined } - | { type: "thinking_level_changed"; level: ThinkingLevel } - | { type: "service_tier_changed"; serviceTier: ServiceTier } - | SessionRetryEvent - | { type: "rlm_child_update"; child: RlmChildAgentSnapshot } - | { type: "recap_update"; recap: string | undefined } - | { type: "goal_update"; goal: GoalState } - | SessionBashEvent - | { type: "refine_complete"; result: RefinementResult } - | { type: "refine_failed"; error: string }; +export type { + RlmChildAgentActivity, + RlmChildAgentSnapshot, + RlmChildAgentStatus, +} from "../session/children/child-types.js"; +export { compactRlmText, rlmChildLabel } from "../session/children/child-types.js"; +export type { CompactionReason } from "../session/compaction/compaction.js"; +export { CompactionSkippedError } from "../session/compaction/compaction-execution.js"; +export type { GoalState, GoalStatus } from "../session/goals/contracts.js"; export { SESSION_ACTION_RECOVERY_FORMAT_VERSION, type SessionActionRecoveryAction, @@ -303,14 +138,11 @@ export { type SessionActionRecoveryRecord, type SessionActionRecoverySnapshot, } from "../session/prepared-actions.js"; - -export type { TurnExecutionPolicy } from "../session/turn-preparation.js"; - -export type AgentSessionEventListener = (event: AgentSessionEvent) => void; - -export { CompactionSkippedError } from "../session/compaction-execution.js"; - -export { RefineSkippedError } from "../session/refinement.js"; +export { RefineSkippedError } from "../session/refinement/refinement.js"; +export type { AgentSessionEvent, AgentSessionEventListener } from "../session/turns/events.js"; +export type { TurnExecutionPolicy } from "../session/turns/turn-preparation.js"; +export type { SessionStats } from "./session-stats.js"; +export { type ParsedSkillBlock, parseSkillBlock } from "./skill-blocks.js"; export interface AgentSessionConfig { agent: Agent; @@ -382,293 +214,15 @@ export interface AgentSessionConfig { initialGoal?: { objective: string; tokenBudget?: number }; } -export type { ExtensionBindings } from "../session/extensions.js"; - -export type { AutoRefineReviewer, AutoRefineReviewRequest } from "../session/refinement.js"; -export interface PromptOptions { - expandPromptTemplates?: boolean; - images?: ImageContent[]; - streamingBehavior?: "steer" | "followUp"; - followUpQueueKey?: string; - source?: InputSource; - preflightResult?: (success: boolean, queued?: boolean) => void; - queueIfBusy?: boolean; - resumeIfIdle?: boolean; - internalPrompt?: boolean; - suppressAutonomousContinuation?: boolean; - skipInputHandlers?: boolean; - signal?: AbortSignal; - admissionCommitted?: () => void; - agentMessageId?: string; - content?: (TextContent | ImageContent)[]; - customMessage?: CustomMessage; -} - -interface InternalPromptOptions extends PromptOptions { - skipPrePromptWork?: boolean; - returnAfterAccepted?: boolean; - agentMessageId?: string; -} - -type SubmissionExtensionCommandPolicy = "execute" | "reject" | "ignore"; - -interface SubmissionNormalizationPolicy { - parseSessionCommands: boolean; - extensionCommands: SubmissionExtensionCommandPolicy; - inputSource?: InputSource; - expandSkills: boolean; - expandPromptTemplates: boolean; -} - -type NormalizedSubmission = - | { kind: "prompt"; text: string; images?: ImageContent[] } - | { - kind: "sessionCommand"; - text: string; - images?: ImageContent[]; - command: SessionSlashCommand; - } - | { kind: "extensionCommand"; completion: Promise } - | { kind: "handled" }; - -function oncePreflight( - preflightResult: ((success: boolean, queued?: boolean) => void) | undefined, -): (success: boolean, queued?: boolean) => void { - let settled = false; - return (success, queued = false) => { - if (!settled) { - settled = true; - preflightResult?.(success, queued); - } - }; -} - -const IPYTHON_SENT_AGENT_MESSAGE_CUSTOM_ENTRY = "ipython_sent_agent_message"; - -interface PersistedIpythonSentAgentMessage { - toolCallId: string; - message: KernelSentAgentMessage; -} - -function isObjectRecord(value: unknown): value is Record { - return typeof value === "object" && value !== null && !Array.isArray(value); -} - -function parsePersistedIpythonSentAgentMessage(value: unknown): PersistedIpythonSentAgentMessage | undefined { - if (!isObjectRecord(value) || typeof value.toolCallId !== "string" || !isObjectRecord(value.message)) { - return undefined; - } - const { id, message, deliveryStatus, target } = value.message; - if ( - typeof id !== "string" || - typeof message !== "string" || - (deliveryStatus !== "delivered" && deliveryStatus !== "queued") || - !isObjectRecord(target) || - typeof target.activeSessionId !== "string" || - typeof target.sessionId !== "string" - ) { - return undefined; - } - return { - toolCallId: value.toolCallId, - message: { - id, - message, - deliveryStatus, - target: { - activeSessionId: target.activeSessionId, - sessionId: target.sessionId, - ...(typeof target.sessionName === "string" ? { sessionName: target.sessionName } : {}), - }, - }, - }; -} - -function appendSentAgentMessageToToolResult( - message: AgentMessage, - toolCallId: string, - sentMessage: KernelSentAgentMessage, -): boolean { - if (message.role !== "toolResult" || message.toolName !== "ipython" || message.toolCallId !== toolCallId) { - return false; - } - const details = isObjectRecord(message.details) ? message.details : {}; - const current = Array.isArray(details.sentAgentMessages) ? details.sentAgentMessages : []; - if (current.some((entry) => isObjectRecord(entry) && entry.id === sentMessage.id)) { - return true; - } - message.details = { - ...details, - sentAgentMessages: [...current, sentMessage], - }; - return true; -} - -function injectedMessagePreviewLabel(message: CustomMessage): string | undefined { - switch (message.customType) { - case HEARTBEAT_PROMPT_CUSTOM_TYPE: - return HEARTBEAT_PROMPT_PREVIEW_LABEL; - case ASYNC_BASH_COMPLETION_CUSTOM_TYPE: - return ASYNC_BASH_COMPLETION_PREVIEW_LABEL; - case GOAL_CONTEXT_CUSTOM_TYPE: - return GOAL_CONTEXT_PREVIEW_LABEL; - default: - return undefined; - } -} - -interface AgentMessageDeferred { - promise: Promise; - resolve: () => void; - reject: (error: Error) => void; -} - -interface AgentMessageOutcome { - delivery?: AgentMessageDeferred; - completion?: AgentMessageDeferred; -} - -function createAgentMessageDeferred(): AgentMessageDeferred { - const deferred = {} as AgentMessageDeferred; - deferred.promise = new Promise((resolve, reject) => { - deferred.resolve = resolve; - deferred.reject = reject; - }); - deferred.promise.catch(() => undefined); - return deferred; -} - -export interface ModelCycleResult { - model: Model; - thinkingLevel: ThinkingLevel; - serviceTier: ServiceTier; - isScoped: boolean; -} - -interface ModelSelectOptions { - waitForExtensions?: boolean; -} - -type AutonomousSlashCommand = { kind: "status" } | { kind: "on"; config?: AgentAutonomousConfig } | { kind: "off" }; +export type { ExtensionBindings } from "../session/extensions/extensions.js"; +export type { PromptOptions } from "../session/input/prompt-submission.js"; +export type { ModelCycleResult } from "../session/models/model-selection.js"; +export type { AutoRefineReviewer, AutoRefineReviewRequest } from "../session/refinement/refinement.js"; import type { RlmMaxDepthStatus, SetRlmMaxDepthResult } from "./rlm-max-depth.js"; export type { RlmMaxDepthSource, RlmMaxDepthStatus, SetRlmMaxDepthResult } from "./rlm-max-depth.js"; -type AutonomousRuntimeSnapshot = Pick< - AutonomousRuntimeState, - "continuationsUsed" | "gateAttempts" | "lastGateFailure" | "lastGateFailureSnapshot" ->; - -interface RlmSubagentModelSelection { - model: Model; -} - -const AUTONOMOUS_STATUS_NUMBER_FORMAT = new Intl.NumberFormat("en-US"); - -const AUTONOMOUS_BUDGET_USAGE = - "Usage: /autonomous [status|off] or /autonomous on [--max-continuations ] [--max-turns ] [--max-tokens ] [--timeout-ms ] [--gate ] [--gate-retries ] [--gate-timeout-ms ]"; - -// `/autonomous` budget flags mirror the `--autonomous-*` CLI options. The CLI -// spelling (`--autonomous-max-continuations`) is accepted as an alias so the -// exact CLI budget flags also work from the slash command. -const AUTONOMOUS_BUDGET_FLAGS: ReadonlySet = new Set([ - "max-continuations", - "max-turns", - "max-tokens", - "timeout-ms", - "gate", - "gate-retries", - "gate-timeout-ms", -]); - -function parseAutonomousBudgetInt(flag: string, value: string, allowUnlimited = false): number { - if (allowUnlimited && value.toLowerCase() === "unlimited") { - return UNLIMITED_AUTONOMOUS_LIMIT; - } - // Commas and underscores are accepted as digit separators (100,000,000). - const digits = value.replace(/[,_]/g, ""); - if (!/^[1-9]\d*$/.test(digits)) { - throw new Error( - `--${flag} must be a positive integer${allowUnlimited ? ' or "unlimited"' : ""}. ${AUTONOMOUS_BUDGET_USAGE}`, - ); - } - return Number(digits); -} - -function parseAutonomousBudgetOptions(tokens: string[]): AgentAutonomousConfig { - const config: AgentAutonomousConfig = {}; - const gateCommands: string[] = []; - for (let i = 0; i < tokens.length; i++) { - const token = tokens[i]!; - if (!token.startsWith("--")) { - throw new Error(`Unexpected autonomous argument: ${token}. ${AUTONOMOUS_BUDGET_USAGE}`); - } - const equalsIndex = token.indexOf("="); - const rawFlag = equalsIndex === -1 ? token : token.slice(0, equalsIndex); - const inlineValue = equalsIndex === -1 ? undefined : token.slice(equalsIndex + 1); - const flag = rawFlag.startsWith("--autonomous-") ? rawFlag.slice("--autonomous-".length) : rawFlag.slice(2); - if (!AUTONOMOUS_BUDGET_FLAGS.has(flag)) { - throw new Error(`Unknown autonomous budget flag: ${rawFlag}. ${AUTONOMOUS_BUDGET_USAGE}`); - } - let value = inlineValue; - if (value === undefined) { - const next = tokens[i + 1]; - if (next === undefined || next.startsWith("--")) { - throw new Error(`Missing value for ${rawFlag}. ${AUTONOMOUS_BUDGET_USAGE}`); - } - value = next; - i++; - } - if (value === "") { - throw new Error(`Missing value for ${rawFlag}. ${AUTONOMOUS_BUDGET_USAGE}`); - } - switch (flag) { - case "gate": - gateCommands.push(value); - break; - case "gate-retries": - config.gates = config.gates ?? {}; - config.gates.maxRetries = parseAutonomousBudgetInt(flag, value); - break; - case "gate-timeout-ms": - config.gates = config.gates ?? {}; - config.gates.timeoutMs = parseAutonomousBudgetInt(flag, value); - break; - case "max-continuations": - config.maxContinuations = parseAutonomousBudgetInt(flag, value, true); - break; - case "max-turns": - config.maxTurns = parseAutonomousBudgetInt(flag, value, true); - break; - case "max-tokens": - config.maxTokens = parseAutonomousBudgetInt(flag, value, true); - break; - case "timeout-ms": - config.timeoutMs = parseAutonomousBudgetInt(flag, value, true); - break; - } - } - if (gateCommands.length > 0) { - config.gates = { ...config.gates, commands: gateCommands }; - } - // Named budget flags define the whole budget: any limit the user did not - // name stops cutting the run short. With no budget flags at all, the - // configured or default limits still apply. - if ( - config.maxContinuations !== undefined || - config.maxTurns !== undefined || - config.maxTokens !== undefined || - config.timeoutMs !== undefined - ) { - config.maxContinuations ??= UNLIMITED_AUTONOMOUS_LIMIT; - config.maxTurns ??= UNLIMITED_AUTONOMOUS_LIMIT; - config.maxTokens ??= UNLIMITED_AUTONOMOUS_LIMIT; - config.timeoutMs ??= UNLIMITED_AUTONOMOUS_LIMIT; - } - return config; -} - export class AgentSession { private readonly _tools: SessionTools; private readonly _extensions: SessionExtensions; @@ -710,8 +264,7 @@ export class AgentSession { }, invalidateOwnUsage: () => this._invalidateOwnUsage(), afterParentDrain: (flush) => { - this._agentEventQueue = this._agentEventQueue.then(flush, flush); - this._agentEventQueue.catch(() => {}); + this._events.enqueue(flush); }, }); private readonly _children = new SessionChildren({ @@ -748,28 +301,209 @@ export class AgentSession { onSettled: () => this._maybeResumeGoalContinuationAfterRlmWork(), }); + private readonly _turnPolicy = new SessionTurnPolicy({ + steeringStopPending: () => this._steeringStopPending, + stopGoalForTerminalMessage: (message) => this._stopGoalContinuationForTerminalMessage(message), + getGoals: () => this._goals, + accountAssistantBudget: (message) => this._goalContinuation.accountAssistantBudget(message), + getRefinement: () => this._refinement, + getEventQueue: () => this._events.queue, + getCompaction: () => this._compaction, + getMessages: () => this.agent.state.messages, + getSettings: () => this.settingsManager, + getModel: () => this.model, + getStore: () => this.sessionManager, + queueThresholdGoal: (message) => this._queueGoalContinuationForThresholdCompaction(message), + queueThresholdAutonomous: (message) => this._queueAutonomousContinuationForThresholdCompaction(message), + getQueuedCount: () => this.queuedActionCount, + getArrivalEpoch: () => this._inputAdmission.arrivalEpoch, + getGoalMessages: (context, signal) => this._getGoalContinuationMessages(context, signal), + getAutonomous: () => this._autonomousContinuation, + snapshotAutonomous: () => this._snapshotAutonomousRuntimeState(), + restoreAutonomous: (snapshot) => this._restoreAutonomousRuntimeSnapshot(snapshot), + }); + private readonly _turnExecution = new SessionTurnExecution({ + getPreparer: () => this._turnPreparer, + getFence: () => this._commitFence, + acquireFence: () => this._acquireSessionActionCommitFence(), + isDeferred: (epoch) => this._isSessionInputHandoffDeferred(epoch), + isStreaming: () => this.isStreaming, + getBasePrompt: () => this._baseSystemPrompt, + getBasePromptOptions: () => this._baseSystemPromptOptions, + refreshExtensionSystemPrompt: (prompt, snapshot) => this._refreshExtensionSystemPrompt(prompt, snapshot), + getExtensions: () => this._extensionRunner, + getAgent: () => this.agent, + takeNextTurnMessages: () => this._takePendingNextTurnMessages(), + restoreNextTurnMessages: (messages) => this._pendingContext.prependMessages(messages), + consumePendingDigest: () => this._harnessContext.consumePendingDigest(), + rearmDigest: () => this._harnessContext.rearmDigest(), + getDigest: () => this._harnessDigest(), + getLatestDigest: () => this._latestContextHarnessDigest(), + suppressForMessage: (message) => this._markAutonomousContinuationSuppressed(message), + runSuppressed: (run) => this._runWithAutonomousContinuationSuppressed(run), + notifyCheckpoints: () => this._notifySessionInputCheckpointChange(), + emitQueueUpdate: () => this._emitQueueUpdate(), + hasCancelledCapture: () => this._hasCancelledDispatchCapture(), + getEventQueue: () => this._events.queue, + waitForRetry: () => this.waitForRetry(), + forgetContinuations: (messages) => this._forgetConsumedPostCompactionContinuations(messages), + }); + /** Session-owned actions. Items are never fed into Agent.steer/followUp. */ + private readonly _actionStore = new ActionStore(); + private readonly _inputCheckpoints = new SessionInputCheckpoints(this._actionStore, { + getFence: () => this._commitFence, + getScheduler: () => this._inputScheduler, + getEventQueue: () => this._events.queue, + acquireFence: (signal) => this._acquireSessionActionCommitFence(signal), + getStore: () => this.sessionManager, + assertAdmissionAvailable: () => this._assertSessionActionAdmissionAvailable(), + getContinuation: () => this._continuation, + scheduleInput: () => this._scheduleSessionInputPump(), + getAgent: () => this.agent, + getUnfinishedCount: () => this.unfinishedActionCount, + waitForIdle: () => this.waitForIdle(), + }); + private readonly _promptSubmission = new SessionPromptSubmission(this._actionStore, { + promptInjectedMessage: (text, message, options) => this._promptInjectedMessage(text, message, options), + waitForActivityChange: (signal) => this._waitForSessionActivityChange(signal), + queueAgentMessagePrompt: (text, streamingBehavior, customMessage) => + this.queueAgentMessagePrompt(text, streamingBehavior, customMessage), + getScheduler: () => this._inputScheduler, + getFence: () => this._commitFence, + isStreaming: () => this.isStreaming, + isCompacting: () => this.isCompacting, + isRetrying: () => this.isRetrying, + isBashRunning: () => this.isBashRunning, + resumeAdmission: () => this._resumeSessionInputAdmission(), + assertAdmissionAvailable: () => this._assertSessionActionAdmissionAvailable(), + acquireAdmissionFence: (signal) => this._acquireDirectTurnAdmissionFence(signal), + normalize: (text, images, policy) => this._normalizeSubmission(text, images, policy), + settleAgentMessage: (id, leg, error) => this._settleAgentMessage(id, leg, error), + canStartImmediately: () => this._canStartSessionActionImmediately(), + admit: (action, options) => this._admitSessionInput(action, options), + waitForInputIdle: () => this.waitForSessionInputIdle(), + isBusy: (point) => this._isBusyForSessionInput(point), + takeNextTurnMessages: () => this._takePendingNextTurnMessages(), + restoreNextTurnMessages: (messages) => this._pendingContext.prependMessages(messages), + appendNextTurnMessage: (message) => this._pendingContext.appendMessages(message), + getActivity: () => this._runtimeActivity(), + suppressForMessage: (message) => this._markAutonomousContinuationSuppressed(message), + observeDeferral: (action) => this._observeSessionActionDeferral(action), + rejectAgentMessage: (id, error) => this._rejectAgentMessage(id, error), + cancelActions: (predicate, error) => this._cancelSessionActions(predicate, error), + emitQueueUpdate: () => this._emitQueueUpdate(), + getClearEpoch: () => this._actionQueue.clearEpoch, + queuePrompt: (schedule, text, images, options) => this._queuePreparedPrompt(schedule, text, images, options), + resetParentReply: () => { + this._childState.resetReply(); + }, + getAgent: () => this.agent, + getStore: () => this.sessionManager, + emit: (event) => this._emit(event), + }); + private readonly _inputAdmission = new SessionInputAdmission(this._actionStore, { + getScheduler: () => this._inputScheduler, + isDisposed: () => this._disposed, + isDisposing: () => this._disposing, + isStreaming: () => this.isStreaming, + rejectAgentMessage: (id, error) => this._rejectAgentMessage(id, error), + emitQueueUpdate: () => this._emitQueueUpdate(), + resumeAdmission: () => this._resumeSessionInputAdmission(), + scheduleInput: () => this._scheduleSessionInputPump(), + suppressForMessage: (message) => this._markAutonomousContinuationSuppressed(message), + }); + private readonly _pendingContext = new SessionPendingContext(this._actionStore, { + getScheduler: () => this._inputScheduler, + getFence: () => this._commitFence, + isDisposed: () => this._disposed, + isDisposing: () => this._disposing, + admit: (action, options) => this._admitSessionInput(action, options), + scheduleInput: () => this._scheduleSessionInputPump(), + addCheckpointWaiter: (waiter) => this._inputCheckpoints.add(waiter), + removeCheckpointWaiter: (waiter) => this._inputCheckpoints.remove(waiter), + acquireFence: (signal) => this._acquireSessionActionCommitFence(signal), + cancelActions: (predicate, error, candidates) => this._cancelSessionActions(predicate, error, candidates), + }); + private readonly _actionQueue = new SessionActionQueue(this._actionStore, { + formatLabel: (text) => compactRlmText(text), + getScheduler: () => this._inputScheduler, + getAgent: () => this.agent, + rearmDigest: () => this._harnessContext.rearmDigest(), + restoreNextTurnMessages: (messages) => this._pendingContext.prependMessages(messages), + notifyCheckpoints: () => this._notifySessionInputCheckpointChange(), + emitQueueUpdate: () => this._emitQueueUpdate(), + settleAgentMessage: (id, leg, error) => this._settleAgentMessage(id, leg, error), + rejectAgentMessage: (id, error) => this._rejectAgentMessage(id, error), + admit: (action, options) => this._admitSessionInput(action, options), + queuePrompt: (schedule, text, images, options) => this._queuePreparedPrompt(schedule, text, images, options), + resumeQueuedWork: () => this.resumeQueuedWork(), + }); + private readonly _actionRecovery = new SessionActionRecovery(this._actionStore, { + isTerminalNoticeAction: (action) => this._isRlmTerminalNoticeAction(action), + retainTerminalNotice: (id) => this._pendingContext.retainTerminalNotice(id), + releaseTerminalNotice: (id) => this._pendingContext.releaseTerminalNotice(id), + admit: (action, options) => this._admitSessionInput(action, options), + }); + private readonly _commandExecution = new SessionCommandExecution(this._actionStore, { + getFence: () => this._commitFence, + acquireFence: () => this._acquireSessionActionCommitFence(), + getRefinement: () => this._refinement, + isDeferred: (epoch) => this._isSessionInputHandoffDeferred(epoch), + getActivity: () => this._runtimeActivity(), + notifyCheckpoints: () => this._notifySessionInputCheckpointChange(), + emitQueueUpdate: () => this._emitQueueUpdate(), + settleAgentMessage: (id, leg, error) => this._settleAgentMessage(id, leg, error), + rejectAgentMessage: (id, error) => this._rejectAgentMessage(id, error), + compact: (instructions, options) => this.compact(instructions, options), + refine: (options, internal) => this.refine(options, internal), + handleGoalCommand: (text, images) => this._handleGoalSlashCommand(text, images), + handleAutonomousCommand: (text) => this._handleAutonomousSlashCommand(text), + getGoalState: () => this._goals.state, + getStore: () => this.sessionManager, + getAgent: () => this.agent, + emit: (event) => this._emit(event), + }); + private readonly _events = new SessionEvents(this._actionStore, { + getAgent: () => this.agent, + getStore: () => this.sessionManager, + getExtensions: () => this._extensionRunner, + getRetry: () => this._retry, + getCompaction: () => this._compaction, + getRefinement: () => this._refinement, + addAutonomousUsage: (usage) => this._autonomousContinuation.recordUsage(usage), + applyLateMessages: (message) => this._applyLateIpythonSentAgentMessages(message), + notifyCheckpoints: () => this._notifySessionInputCheckpointChange(), + settleAgentMessage: (id, leg, error) => this._settleAgentMessage(id, leg, error), + getSnapshot: () => this.getSessionActionSnapshot(), + accountAssistantBudget: (message) => this._goalContinuation.accountAssistantBudget(message), + finishGoal: (message) => this._finishGoalForTerminalAssistantMessage(message), + checkCompaction: (message) => this._checkCompaction(message), + }); + private readonly _messageDelivery = new SessionMessageDelivery(this._actionStore, { + isDisposed: () => this._disposed, + getMessages: () => this.agent.state.messages, + getStore: () => this.sessionManager, + enqueue: (work) => this._events.enqueue(work), + emit: (event) => this._emit(event), + promptUntilAccepted: (text, options) => this.promptUntilAccepted(text, options), + cancelActions: (predicate, error) => this._cancelSessionActions(predicate, error), + }); + private readonly _submissionNormalizer = new SubmissionNormalizer({ + getExtensions: () => this._extensionRunner, + getPrompts: () => this.promptTemplates, + getSkills: () => this.resourceLoader.getSkills().skills, + }); private readonly _refinement: SessionRefinement; readonly agent: Agent; readonly sessionManager: SessionManager; readonly settingsManager: SettingsManager; - private _serviceTierPreference: ServiceTier; - - private _scopedModels: Array<{ - model: Model; - thinkingLevel?: ThinkingLevel; - }>; - - private _unsubscribeAgent?: () => void; - private _eventListeners: AgentSessionEventListener[] = []; - private _lastSessionActionSnapshot: SessionActionSnapshot = { - queuedCount: 0, - steering: [], - followUps: [], - }; - private _agentEventQueue: Promise = Promise.resolve(); + private readonly _export: SessionExport; + private readonly _contextView: SessionContextView; + private readonly _modelSelection: SessionModelSelection; + private get _scopedModels() { + return this._modelSelection.scopedModels; + } - /** Session-owned actions. Items are never fed into Agent.steer/followUp. */ - private readonly _actionStore = new ActionStore(); private readonly _inputScheduler = new SessionInputScheduler({ canSchedule: () => !this._disposed && !this._disposing && this._hasSelectableSessionInput(), run: (epoch) => this._inputDispatcher.run(epoch), @@ -783,22 +517,20 @@ export class AgentSession { getDeliveryMode: (delivery) => (delivery === "next_turn_boundary" ? this.steeringMode : this.followUpMode), waitForAgentIdle: () => this.agent.waitForIdle(), hasCancelledDispatchCapture: () => this._hasCancelledDispatchCapture(), - getEventQueue: () => this._agentEventQueue, + getEventQueue: () => this._events.queue, waitForRefinement: () => this._refinement._waitForRefineIdle(), getTranscript: () => this.agent.state.messages, startTurns: (actions, epoch) => this._startPreparedTurnActions(actions, epoch), executeCommand: (action, epoch) => this._executeSelectedSessionCommand(action, epoch), settleAgentMessage: (id, leg, error) => this._settleAgentMessage(id, leg, error), releaseTurn: (id) => { - this._durableRlmTerminalNoticeActionIds.delete(id); + this._pendingContext.releaseTerminalNotice(id); }, notifyCheckpoints: () => this._notifySessionInputCheckpointChange(), emitQueueUpdate: () => this._emitQueueUpdate(), surfaceError: (error) => this._surfaceSessionInputError(error), schedule: () => this._scheduleSessionInputPump(), }); - private _sessionInputArrivalEpoch = 0; - private readonly _durableRlmTerminalNoticeActionIds = new Set(); private readonly _commitFence = new SessionCommitFence(); private readonly _turnPreparer = new TurnPreparer({ hasRefinement: () => this._refinement.isApplying, @@ -809,17 +541,16 @@ export class AgentSession { pendingModelSelection: () => this._pendingModelSelectEmit(), }); // Checkpoint, handoff, and activity waiters share lifecycle-edge notifications to avoid polling. - private readonly _sessionInputCheckpointWaiters = new Set<() => void>(); - private _pendingNextTurnMessages: CustomMessage[] = []; - private readonly _goals: GoalController; - private _goalContinuationAwaitsRlmWork = false; - private _goalAbortInProgress = false; - private _autonomousState: AutonomousRuntimeState; - private _autonomousContinuationSuppressionDepth = 0; - private _autonomousContinuationSuppressedMessages = new WeakSet(); + private readonly _autonomousContinuation: SessionAutonomousContinuation; + private readonly _goalContinuation: SessionGoalContinuation; + private get _goals(): GoalController { + return this._goalContinuation.controller; + } private readonly _compaction = new SessionCompaction({ + includesCompactSkill: () => this._includeCompactSkill, + getContextUsage: () => this.getContextUsage(), getModel: () => this.model, isStreaming: () => this.isStreaming, getSettings: () => this.settingsManager.getCompactionSettings(), @@ -847,14 +578,14 @@ export class AgentSession { hasPendingSessionWork: () => this.hasPendingSessionWork, scheduleContinuation: (continueAfterInput) => this._schedulePostCompactionContinue(continueAfterInput), scheduleRefinement: (willContinue) => this._refinement._scheduleAutoRefineAfterCompaction(willContinue), - takeThresholdAutonomousMessages: () => this._pendingThresholdCompactionAutonomousMessages.splice(0), - getThresholdGoalContinuation: () => this._queuedGoalThresholdContinuation, + takeThresholdAutonomousMessages: () => this._autonomousContinuation.takePendingThresholdMessages(), + getThresholdGoalContinuation: () => this._goalContinuation.thresholdContinuation, clearAutonomousContinuations: (shouldContinue, messages) => this._clearQueuedAutonomousContinuationsAfterSkippedThresholdCompaction(shouldContinue, messages), clearGoalContinuation: (message) => this._clearQueuedGoalContinuationAfterCancelledThresholdCompaction(message), getSessionStore: () => this.sessionManager, retainUnpersistedOutcome: (message) => { - this._unpersistedOutcomes.push(message); + this._harnessContext.retainOutcome(message); }, emit: (event) => this._emit(event), notifyCheckpoints: () => this._notifySessionInputCheckpointChange(), @@ -878,8 +609,10 @@ export class AgentSession { reapDeletedChildren: () => this._reapDeletedRlmSubagentRuntimesAfterCompaction(), }; - private _branchSummaryAbortController: AbortController | undefined = undefined; - private _branchSummaryOperation: Promise | undefined = undefined; + private get _branchSummaryOperation(): Promise | undefined { + return this._history.operation; + } + private readonly _history: SessionHistoryNavigation; private readonly _retry = new SessionRetry({ getRetrySettings: () => this.settingsManager.getRetrySettings(), @@ -909,13 +642,8 @@ export class AgentSession { this._scheduleSessionInputPump(); }, }); - private _agentMessageClearEpoch = 0; - private _agentMessageOutcomes = new Map(); - private _lateIpythonSentAgentMessages = new Map(); - /** Outcome disclosures whose session-file append failed; retained for context rebuilds. */ - private readonly _unpersistedOutcomes: CustomMessage[] = []; /** Fresh/empty contexts defer digest injection to the first committed turn so untouched sessions stay empty. */ - private _harnessDigestPending = false; + private readonly _harnessContext: SessionHarnessContext; private readonly _bash = new SessionBash({ getCwd: () => this.sessionManager.getCwd(), @@ -934,21 +662,16 @@ export class AgentSession { recordBashResult: (command, result, options) => this.recordBashResult(command, result, options), }); - private _turnIndex = 0; - private _modelSelectEmitQueue: Promise = Promise.resolve(); - private _modelSelectEmitQueueIdle = true; - private _modelSelectEmitContext = new AsyncLocalStorage(); - - private _resourceLoader: ResourceLoader; - private _cwd: string; - private _agentDir?: string; - private _initialActiveToolNames?: string[]; - private _includeGoals: boolean; - private _includeCompactSkill: boolean; + private readonly _resourceLoader: ResourceLoader; + private readonly _cwd: string; + private readonly _agentDir?: string; + private readonly _initialActiveToolNames?: string[]; + private readonly _includeGoals: boolean; + private readonly _includeCompactSkill: boolean; private _rlmHeartbeatController?: AgentRlmHeartbeatController; - private _agentMessageController?: AgentSessionMessageController; - private _agentObserveController?: AgentObserveController; - private _mcpManager?: McpManager; + private readonly _agentMessageController?: AgentSessionMessageController; + private readonly _agentObserveController?: AgentObserveController; + private readonly _mcpManager?: McpManager; private _disposed = false; private readonly _disposeCallbacks = new Set<() => void | Promise>(); private _disposeCallbacksPromise?: Promise; @@ -957,10 +680,10 @@ export class AgentSession { private _disposing = false; private _disposeAsyncPromise?: Promise; private readonly _semanticEdges: SemanticEdgeRecorder; - private _rlmParentNodeId?: string; - private _rlmParentAgent?: string; + private readonly _rlmParentNodeId?: string; + private readonly _rlmParentAgent?: string; - private _modelRegistry: ModelRegistry; + private readonly _modelRegistry: ModelRegistry; private readonly _continuation = new SessionContinuation({ waitForAgentIdle: () => this.agent.waitForIdle(), @@ -968,10 +691,10 @@ export class AgentSession { waitForRefinement: () => this._refinement._waitForRefineIdle(), queuedWorkPauseCount: () => this._inputScheduler.queuedWorkPauseCount, addCheckpointWaiter: (waiter) => { - this._sessionInputCheckpointWaiters.add(waiter); + this._inputCheckpoints.add(waiter); }, removeCheckpointWaiter: (waiter) => { - this._sessionInputCheckpointWaiters.delete(waiter); + this._inputCheckpoints.remove(waiter); }, notifyCheckpoints: () => this._notifySessionInputCheckpointChange(), compactionOperation: () => this._compaction.operation, @@ -986,18 +709,65 @@ export class AgentSession { removeQueuedMessages: (predicate) => this.agent.removeQueuedMessages(predicate), followUp: (message) => this.agent.followUp(message), onMessageConsumed: (message) => { - this._queuedAutonomousContinuationSnapshots.delete(message); + this._autonomousContinuation.forgetSnapshot(message); }, }); - private _queuedAutonomousThresholdContinuations = new WeakMap(); - private _queuedAutonomousContinuationSnapshots = new WeakMap(); - private _pendingThresholdCompactionAutonomousMessages: AgentMessage[] = []; - private _queuedGoalThresholdContinuation: AgentMessage | undefined; constructor(config: AgentSessionConfig) { this.agent = config.agent; this.sessionManager = config.sessionManager; this.settingsManager = config.settingsManager; + this._harnessContext = new SessionHarnessContext({ + sessionManager: this.sessionManager, + getMessages: () => this.messages, + getActiveToolNames: () => this.getActiveToolNames(), + getVisibleSkills: () => this._modelVisibleSkills(), + loadHarnessState: () => this._refinement._loadMergedHarnessState(), + applyLateSentMessages: (message) => this._applyLateIpythonSentAgentMessages(message), + }); + this._export = new SessionExport({ + sessionManager: this.sessionManager, + getState: () => this.state, + getTheme: () => this.settingsManager.getTheme(), + getToolDefinition: (name) => this.getToolDefinition(name), + }); + this._contextView = new SessionContextView({ + getContextUsage: () => this.getContextUsage(), + sessionManager: this.sessionManager, + getMessages: () => this.messages, + getModel: () => this.model, + findModel: (provider, modelId) => this._modelRegistry.find(provider, modelId), + subtractUnindexedChildUsage: (ownUsage, entries) => this._childUsage.subtractUnindexed(ownUsage, entries), + getRlmSessionDir: () => this._rlmSessionDirForReading(), + getLiveChildren: () => this._contextViewChildren(), + }); + + this._history = new SessionHistoryNavigation({ + getSessionId: () => this.sessionId, + getRetryPolicy: () => providerRetryPolicy(this.settingsManager), + sessionManager: this.sessionManager, + settingsManager: this.settingsManager, + getModel: () => this.model, + getExtensions: () => this._extensionRunner, + getRequiredAuth: (model) => this._getRequiredRequestAuth(model), + acquireQueuedWorkPause: () => this.acquireQueuedWorkPause(), + acquireCommitFence: () => this._acquireSessionActionCommitFence(), + runWithCommitFence: (lease, run) => this._commitFence.run(lease, run), + waitForAgentIdle: () => this.agent.waitForIdle(), + getEventQueue: () => this._events.queue, + invalidateRefinement: () => this._refinement._invalidatePendingAutoRefineForBranchChange(), + rebuildBranchContext: () => { + this.agent.state.messages = this.sessionManager.buildSessionContext().messages; + this._mergeUnpersistedOutcomes(this.agent.state.messages); + this._restoreLateIpythonSentAgentMessages(); + this._ensureHarnessDigestContext(); + this._goals.reload(); + this._reloadRlmMaxDepthFromBranch(); + this._invalidateQueuedPromptPreparation(); + }, + notifyCheckpoints: () => this._notifySessionInputCheckpointChange(), + }); + this._refinement = new SessionRefinement( { sessionManager: this.sessionManager, @@ -1015,7 +785,7 @@ export class AgentSession { getMessages: () => this.agent.state.messages, getRequiredRequestAuth: (model) => this._getRequiredRequestAuth(model), getExtensionRunner: () => this._extensionRunner, - getEventQueue: () => this._agentEventQueue, + getEventQueue: () => this._events.queue, getCompactionOperation: () => this._compaction.operation, getBranchSummaryOperation: () => this._branchSummaryOperation, waitForAgentIdle: () => this.agent.waitForIdle(), @@ -1023,7 +793,7 @@ export class AgentSession { disconnect: () => this._disconnectFromAgent(), reconnect: () => this._reconnectToAgent(), emit: (event) => this._emit(event), - retainUnpersistedOutcome: (message) => this._unpersistedOutcomes.push(message), + retainUnpersistedOutcome: (message) => this._harnessContext.retainOutcome(message), notifyCheckpoints: () => this._notifySessionInputCheckpointChange(), scheduleInputPump: () => this._scheduleSessionInputPump(), isContinuationScheduled: () => this._continuation.isScheduled, @@ -1034,8 +804,22 @@ export class AgentSession { autoRefineReviewer: config.autoRefineReviewer?.bind(this), }, ); - this._serviceTierPreference = config.serviceTierPreference ?? config.agent.state.serviceTier; - this._scopedModels = config.scopedModels ?? []; + this._modelSelection = new SessionModelSelection( + { + getModel: () => this.model, + getState: () => this.agent.state, + setThinkingLevel: (level) => this.setThinkingLevel(level), + getAvailableThinkingLevels: () => this.getAvailableThinkingLevels(), + supportsThinking: () => this.supportsThinking(), + getRegistry: () => this._modelRegistry, + getExtensions: () => this._extensionRunner, + sessionManager: this.sessionManager, + settingsManager: this.settingsManager, + emit: (event) => this._emit(event), + }, + config.serviceTierPreference ?? config.agent.state.serviceTier, + config.scopedModels ?? [], + ); this._resourceLoader = config.resourceLoader; this._cwd = config.cwd; this._agentDir = config.agentDir; @@ -1120,7 +904,7 @@ export class AgentSession { rebuildRuntime: (options) => this._buildRuntime(options), acquireInputPause: () => this.acquireSessionInputPause(), waitForAgentIdle: () => this.agent.waitForIdle(), - getEventQueue: () => this._agentEventQueue, + getEventQueue: () => this._events.queue, }, { customTools: config.customTools, @@ -1181,11 +965,49 @@ export class AgentSession { this.agent.streamFn = wrapStreamFnWithSemanticEdges(this.agent.streamFn, this._semanticEdges); this._childState.initializeParentReply(); this._children.setRuntimeHost(config.subagentRuntimeHost); - this._autonomousState = createAutonomousRuntimeState(config.autonomous, { - cwd: this._cwd, + this._autonomousContinuation = new SessionAutonomousContinuation(config.autonomous, { + getStatus: () => this.getAutonomousStatus(), + getCwd: () => this._cwd, + getAgent: () => this.agent, + getStore: () => this.sessionManager, + emit: (event) => this._emit(event), + getContinuation: () => this._continuation, + getArrivalEpoch: () => this._inputAdmission.arrivalEpoch, + admit: (action, options) => this._admitSessionInput(action, options), + cancelActions: (predicate, error) => this._cancelSessionActions(predicate, error), + emitQueueUpdate: () => this._emitQueueUpdate(), + getCompaction: () => this._compaction, + getUnfinishedActionCount: () => this.unfinishedActionCount, + cancelContinuation: () => this._cancelPostCompactionContinue(), }); const goalPersistence = createGoalPersistence(this.sessionManager); - this._goals = new GoalController(goalPersistence, (goal) => this._emit({ type: "goal_update", goal })); + this._goalContinuation = new SessionGoalContinuation( + new GoalController(goalPersistence, (goal) => this._emit({ type: "goal_update", goal })), + this._actionStore, + { + getGoalState: () => this.goalState, + queuePrompt: (schedule, text, images, options) => + this._queuePreparedPrompt(schedule, text, images, options), + getScheduler: () => this._inputScheduler, + isDisposed: () => this._disposed, + isDisposing: () => this._disposing, + hasUnsettledChildWork: () => this._hasUnsettledRlmQuiescenceWork(), + ensureRuntimeActive: (context) => this._ensureGoalRuntimeActive(context), + admit: (action, options) => this._admitSessionInput(action, options), + cancelActions: (predicate, error) => this._cancelSessionActions(predicate, error), + clearPendingGoalContexts: () => { + this._pendingContext.removeMessagesMatching( + (message) => message.customType === GOAL_CONTEXT_CUSTOM_TYPE, + ); + }, + emitQueueUpdate: () => this._emitQueueUpdate(), + emitGoalUpdate: () => this._emitGoalUpdate(), + validate: () => this._validateCanStartAgentRun(), + isStreaming: () => this.isStreaming, + includesGoals: () => this._includeGoals, + getAgent: () => this.agent, + }, + ); // Seed initial goal from CLI --goal flag, but only for top-level sessions // and only when the branch contains only bootstrap entry types (model_change, // thinking_level_change, service_tier_change) and no persisted @@ -1195,12 +1017,12 @@ export class AgentSession { 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. - this._pendingNextTurnMessages.push(createGoalContextMessage(this._goals.state, "continuation")); + this._pendingContext.appendMessages(createGoalContextMessage(this._goals.state, "continuation")); } this._restoreLateIpythonSentAgentMessages(); this._goals.restartAccounting(); - this._unsubscribeAgent = this.agent.subscribe(this._handleAgentEvent); + this._events.reconnectToAgent(); this._installAgentToolHooks(); this._installAgentTurnHook(); this._installAgentContinuationHook(); @@ -1251,27 +1073,10 @@ export class AgentSession { this._children.setRuntimeHost(host); } - private async _getRequiredRequestAuth(model: Model): Promise<{ - apiKey: string; - headers?: Record; - requestModel: Model; - }> { - const result = await this._modelRegistry.getApiKeyAndHeaders(model); - if (!result.ok) { - if (result.error.startsWith("No API key found")) { - throw new Error(formatNoApiKeyFoundMessage(model.provider)); - } - throw new Error(result.error); - } - if (result.apiKey) { - return { apiKey: result.apiKey, headers: result.headers, requestModel: result.requestModel ?? model }; - } - - const isOAuth = this._modelRegistry.isUsingOAuth(model); - if (isOAuth) { - throw new Error(formatAuthenticationFailedMessage(model.provider)); - } - throw new Error(formatNoApiKeyFoundMessage(model.provider)); + private _getRequiredRequestAuth( + ...args: Parameters + ): ReturnType { + return this._modelSelection.getRequiredRequestAuth(...args); } /** @@ -1286,7 +1091,7 @@ export class AgentSession { installExtensionToolHooks( this.agent, () => this._extensionRunner, - () => this._agentEventQueue, + () => this._events.queue, ); } @@ -1299,71 +1104,32 @@ export class AgentSession { this.agent.shouldStopAfterTurn = (context) => this._shouldStopAfterTurn(context); } - private _emit(event: AgentSessionEvent): void { - for (const l of this._eventListeners) { - try { - l(event); - } catch { - // A failing observer must not prevent other subscribers from - // receiving lifecycle and persistence events. - } - } - } - - private _emitQueueUpdate(): void { - const actions = this.getSessionActionSnapshot(); - if (JSON.stringify(actions) === JSON.stringify(this._lastSessionActionSnapshot)) return; - this._lastSessionActionSnapshot = actions; - this._emit({ type: "session_action_update", actions }); + private _emit(...args: Parameters): ReturnType { + return this._events.emit(...args); } - private _restoreLateIpythonSentAgentMessages(): void { - this._lateIpythonSentAgentMessages.clear(); - for (const entry of this.sessionManager.getBranch()) { - if (entry.type !== "custom" || entry.customType !== IPYTHON_SENT_AGENT_MESSAGE_CUSTOM_ENTRY) { - continue; - } - const persisted = parsePersistedIpythonSentAgentMessage(entry.data); - if (persisted) { - this._rememberLateIpythonSentAgentMessage(persisted.toolCallId, persisted.message); - } - } + private _emitQueueUpdate( + ...args: Parameters + ): ReturnType { + return this._events.emitQueueUpdate(...args); } - private _rememberLateIpythonSentAgentMessage(toolCallId: string, message: KernelSentAgentMessage): boolean { - const messages = this._lateIpythonSentAgentMessages.get(toolCallId) ?? []; - const isNew = !messages.some((entry) => entry.id === message.id); - if (isNew) { - messages.push(message); - this._lateIpythonSentAgentMessages.set(toolCallId, messages); - } - for (let index = this.agent.state.messages.length - 1; index >= 0; index -= 1) { - if (appendSentAgentMessageToToolResult(this.agent.state.messages[index], toolCallId, message)) { - break; - } - } - return isNew; + private _restoreLateIpythonSentAgentMessages( + ...args: Parameters + ): ReturnType { + return this._messageDelivery.restoreLateIpythonSentAgentMessages(...args); } - private _applyLateIpythonSentAgentMessages(message: AgentMessage): void { - if (message.role !== "toolResult" || message.toolName !== "ipython") { - return; - } - for (const sentMessage of this._lateIpythonSentAgentMessages.get(message.toolCallId) ?? []) { - appendSentAgentMessageToToolResult(message, message.toolCallId, sentMessage); - } + private _applyLateIpythonSentAgentMessages( + ...args: Parameters + ): ReturnType { + return this._messageDelivery.applyLateIpythonSentAgentMessages(...args); } - private _recordLateIpythonSentAgentMessage(toolCallId: string, message: KernelSentAgentMessage): void { - const record = () => { - if (this._disposed || !this._rememberLateIpythonSentAgentMessage(toolCallId, message)) { - return; - } - this.sessionManager.appendCustomEntry(IPYTHON_SENT_AGENT_MESSAGE_CUSTOM_ENTRY, { toolCallId, message }); - this._emit({ type: "ipython_sent_agent_message", toolCallId, message }); - }; - this._agentEventQueue = this._agentEventQueue.then(record, record); - this._agentEventQueue.catch(() => {}); + private _recordLateIpythonSentAgentMessage( + ...args: Parameters + ): ReturnType { + return this._messageDelivery.recordLateIpythonSentAgentMessage(...args); } private _emitGoalUpdate(): void { @@ -1375,250 +1141,39 @@ export class AgentSession { } private _cancelSessionActions( - predicate: (action: QueuedSessionAction) => boolean, - error: Error, - candidates = this._actionStore.clearableActions(), - ): QueuedSessionAction[] { - const matching = candidates.filter(predicate); - const previousStates = new Map(matching.map((action) => [action.id, action.lifecycle.state])); - const preparing = this._actionStore - .activeActions() - .filter( - (action): action is SessionAction => - action.payload.kind === "turn" && action.lifecycle.state === "preparing", - ); - const previousAnchor = preparing.at(-1); - const actions = this._actionStore.remove(predicate, candidates); - const restorableMessages: CustomMessage[] = []; - const removed = new Set(actions); - if (previousAnchor && removed.has(previousAnchor)) { - for (const action of preparing) { - if (!removed.has(action)) action.payload.prepared = undefined; - } - } - for (const action of actions) { - const ticket = this._actionStore.ticketFor(action); - if ( - action.payload.kind === "turn" && - (action.payload.acceptedAgentMessage || - !action.payload.queueVisible || - previousStates.get(action.id) !== "queued") - ) { - ticket.rejectDelivered(error); - } else { - ticket.settleDelivered({ status: "not_applicable" }); - } - ticket.settleCompleted(error); - const dispatched = previousStates.get(action.id) === "committing" && action.payload.kind === "turn"; - if (action.payload.kind === "turn") { - const payload = action.payload; - const restorable = payload.records - .filter( - (record): record is DeliveryRecord & { message: CustomMessage } => - (record.role === "next_turn" || (payload.acceptedAgentMessage && record.role === "prefix")) && - record.message.role === "custom" && - record.message.customType !== HARNESS_DIGEST_CUSTOM_TYPE && - !record.durable, - ) - .map((record) => cloneCustomMessage(record.message)); - restorableMessages.push(...restorable); - // Lazy injection owns digest delivery: a cancelled turn re-arms it - // instead of restoring a possibly stale digest message. - if ( - payload.records.some( - (record) => - record.message.role === "custom" && record.message.customType === HARNESS_DIGEST_CUSTOM_TYPE, - ) - ) { - this._harnessDigestPending = true; - } - if (dispatched) { - payload.captureRunMessages = new Set(payload.records.map((record) => record.message)); - this.agent.state.messages = this.agent.state.messages.filter( - (message) => !payload.captureRunMessages?.has(message), - ); - } - } - if (!dispatched) { - this._actionStore.releaseTerminal(action); - } - } - this._pendingNextTurnMessages.unshift(...restorableMessages); - if (actions.length > 0) this._notifySessionInputCheckpointChange(); - return actions; - } - - private _clearQueuedGoalContexts(): void { - this._goalContinuationAwaitsRlmWork = false; - this._pendingNextTurnMessages = this._pendingNextTurnMessages.filter( - (message) => message.customType !== GOAL_CONTEXT_CUSTOM_TYPE, - ); - this.agent.removeQueuedMessages( - (message) => message.role === "custom" && message.customType === GOAL_CONTEXT_CUSTOM_TYPE, - ); - this._cancelSessionActions( - (action) => - action.payload.kind === "turn" && action.payload.customMessage?.customType === GOAL_CONTEXT_CUSTOM_TYPE, - new Error("Queued goal context was cleared before delivery."), - ); - this._emitQueueUpdate(); - } - - private _startGoal(objectiveText: string, tokenBudget: number | undefined): GoalState { - const objective = validateGoalObjective(objectiveText); - const budget = validateGoalBudget(tokenBudget); - this._goalContinuationAwaitsRlmWork = false; - return this._goals.start(objective, budget); - } - - private _clearGoal(): void { - this._clearQueuedGoalContexts(); - this._goals.clear(); - } - - private _pauseGoal(): void { - this._clearQueuedGoalContexts(); - this._goals.pause(); - } - - private async _resumeGoal(): Promise { - if (this._goals.resume()) { - await this._runOrQueueGoalContext("continuation"); - } - } - - private _finishGoalForTerminalAssistantMessage(message: AssistantMessage): void { - if (this._goals.state.status !== "active") { - return; - } - - if (message.stopReason === "aborted") { - this._goalAbortInProgress = false; - return; - } - - if (message.stopReason === "error") { - if (this._goalAbortInProgress) { - this._goalAbortInProgress = false; - return; - } - this._goals.fail(message.errorMessage || "Assistant response failed"); - } + ...args: Parameters + ): ReturnType { + return this._actionQueue.cancelSessionActions(...args); } - private _stopGoalContinuationForTerminalMessage(message: AssistantMessage): boolean { - if (message.stopReason !== "error" && message.stopReason !== "aborted") { - return false; - } - try { - this._finishGoalForTerminalAssistantMessage(message); - } catch { - // Goal hooks must not reject; listener failures should not crash the agent loop. - } - return true; + private _startGoal( + ...args: Parameters + ): ReturnType { + return this._goalContinuation.startGoal(...args); } - private _parseAutonomousSlashCommand(text: string): AutonomousSlashCommand | undefined { - const command = parseSessionSlashCommand(text); - if (command?.name !== "autonomous") return undefined; - const tokens = parseCommandArgs(command.args); - if (tokens.length === 0 || tokens[0]!.toLowerCase() === "status") { - if (tokens.length > 1) { - throw new Error(`Unexpected autonomous argument: ${tokens[1]}. ${AUTONOMOUS_BUDGET_USAGE}`); - } - return { kind: "status" }; - } - const subcommand = tokens[0]!.toLowerCase(); - if (subcommand === "on" || subcommand === "enable" || subcommand === "enabled") { - return { kind: "on", config: parseAutonomousBudgetOptions(tokens.slice(1)) }; - } - if (subcommand === "off" || subcommand === "disable" || subcommand === "disabled") { - if (tokens.length > 1) { - throw new Error(`Unexpected autonomous argument: ${tokens[1]}. ${AUTONOMOUS_BUDGET_USAGE}`); - } - return { kind: "off" }; - } - throw new Error(AUTONOMOUS_BUDGET_USAGE); - } - - private _formatAutonomousStatus(): string { - const status = this.getAutonomousStatus(); - const state = status.enabled ? "on" : "off"; - const elapsedSeconds = status.startedAt ? Math.round((Date.now() - status.startedAt) / 1000) : 0; - const gateSummary = - status.gates.commands.length > 0 ? status.gates.commands.map((command) => `"${command}"`).join(", ") : "none"; - const formatCount = (value: number): string => - isUnlimitedAutonomousLimit(value) ? "unlimited" : AUTONOMOUS_STATUS_NUMBER_FORMAT.format(value); - const timeBudget = isUnlimitedAutonomousLimit(status.limits.timeoutMs) - ? "unlimited" - : `${AUTONOMOUS_STATUS_NUMBER_FORMAT.format(Math.round(status.limits.timeoutMs / 1000))}s`; - return `[autonomous-status: ${state}]\n\nContinuations: ${formatCount(status.continuationsUsed)}/${formatCount(status.limits.maxContinuations)}. Turns: ${formatCount(status.turnsUsed)}/${formatCount(status.limits.maxTurns)}. Tokens: ${formatCount(status.tokensUsed)}/${formatCount(status.limits.maxTokens)}. Time: ${elapsedSeconds}s/${timeBudget}. Gates: ${gateSummary}.`; - } - - private _emitAutonomousStatus(): void { - const message = { - role: "custom" as const, - customType: "autonomous_status", - content: this._formatAutonomousStatus(), - display: true, - details: this.getAutonomousStatus(), - timestamp: Date.now(), - } satisfies CustomMessage; - this.agent.state.messages.push(message); - this.sessionManager.appendCustomMessageEntry( - message.customType, - message.content, - message.display, - message.details, - ); - this._emit({ type: "message_start", message }); - this._emit({ type: "message_end", message }); + private _finishGoalForTerminalAssistantMessage( + ...args: Parameters + ): ReturnType { + return this._goalContinuation.finishGoalForTerminalAssistantMessage(...args); } - private async _handleAutonomousSlashCommand(text: string): Promise { - const command = this._parseAutonomousSlashCommand(text); - if (!command) { - return false; - } - if (command.kind === "on") { - setAutonomousEnabled(this._autonomousState, true, { cwd: this._cwd }); - setAutonomousLimits(this._autonomousState, command.config); - } else if (command.kind === "off") { - setAutonomousEnabled(this._autonomousState, false); - this._clearQueuedAutonomousContinuations(); - } - this._emitAutonomousStatus(); - return true; + private _stopGoalContinuationForTerminalMessage( + ...args: Parameters + ): ReturnType { + return this._goalContinuation.stopGoalContinuationForTerminalMessage(...args); } - private _appendBeforeAgentStartMessages( - messages: AgentMessage[], - result: Awaited>, - ): void { - if (!result?.messages) return; - for (const message of result.messages) { - messages.push({ - role: "custom", - customType: message.customType, - content: message.content, - display: message.display, - details: message.details, - timestamp: Date.now(), - }); - } + private _handleAutonomousSlashCommand( + ...args: Parameters + ): ReturnType { + return this._autonomousContinuation.handleAutonomousSlashCommand(...args); } - private async _validateCanStartAgentRun(): Promise { - if (!this.model) { - throw new Error(formatNoModelSelectedMessage()); - } - if (!this._modelRegistry.hasConfiguredAuth(this.model)) { - const isOAuth = this._modelRegistry.isUsingOAuth(this.model); - if (isOAuth) { - throw new Error(formatAuthenticationFailedMessage(this.model.provider)); - } - throw new Error(formatNoApiKeyFoundMessage(this.model.provider)); - } + private _validateCanStartAgentRun( + ...args: Parameters + ): ReturnType { + return this._modelSelection.validateCanStartAgentRun(...args); } /** @@ -1648,343 +1203,74 @@ export class AgentSession { } } - private _maybeResumeGoalContinuationAfterRlmWork(): void { - if (!this._goalContinuationAwaitsRlmWork) return; - if (this._disposed || this._disposing || this._hasUnsettledRlmQuiescenceWork()) return; - if (this._goals.state.status !== "active" || !this._goals.state.objective) { - this._goalContinuationAwaitsRlmWork = false; - return; - } - // Keep the deferral while admission is paused or the pump is suspended - // (post-abort); the pause release and resumeQueuedWork retry. - if (this._inputScheduler.admissionPaused || this._inputScheduler.suspended) return; - const goalBeforeResume = this._goals.checkpoint(); - try { - this._ensureGoalRuntimeActive(); - this._goals.recordContinuation(); - const message = createGoalContextMessage(this._goals.state, "continuation"); - const normalized = normalizeMessageContent(message.content); - // No front: a settling child's terminal notice must be read first. - this._admitSessionInput( - this._createPreparedTurnAction("followUp", normalized.text, normalized.images, { - message, - resumeIfIdle: true, - }), - ); - this._goalContinuationAwaitsRlmWork = false; - } catch { - // Admission can race a new pause; roll back so the retry re-counts. - this._goals.restore(goalBeforeResume, { restoreClock: false }); - } + private _maybeResumeGoalContinuationAfterRlmWork( + ...args: Parameters + ): ReturnType { + return this._goalContinuation.maybeResumeGoalContinuationAfterRlmWork(...args); } - private _runOrQueueGoalContext(kind: "continuation" | "objective_updated", images?: ImageContent[]): void { - if (!this._goals.state.objective) return; - this._ensureGoalRuntimeActive(); - const message = createGoalContextMessage(this._goals.state, kind, images); - const normalized = normalizeMessageContent(message.content); - const action = this._createPreparedTurnAction("followUp", normalized.text, normalized.images, { - message, - resumeIfIdle: true, - }); - this._admitSessionInput(action, { front: true, wake: false }); - } - - private async _handleGoalSlashCommand(text: string, images: ImageContent[] | undefined): Promise { - const command = parseGoalSlashCommand(text); - if (!command) { - return false; - } - - if (command.kind === "status") { - this._emitGoalUpdate(); - return true; - } - - if (command.kind === "clear") { - this._clearGoal(); - return true; - } - - if (command.kind === "pause") { - this._pauseGoal(); - return true; - } - - if (command.kind === "resume") { - await this._resumeGoal(); - return true; - } - - const previousWasActive = this._goals.state.status === "active"; - if (!this.isStreaming) { - await this._validateCanStartAgentRun(); - } - this._ensureGoalRuntimeActive(); - this._clearQueuedGoalContexts(); - this._startGoal(command.objective, command.tokenBudget); - await this._runOrQueueGoalContext(previousWasActive ? "objective_updated" : "continuation", images); - return true; + private _handleGoalSlashCommand( + ...args: Parameters + ): ReturnType { + return this._goalContinuation.handleGoalSlashCommand(...args); } private get _steeringStopPending(): boolean { - return ( - this._actionStore.queuedActions("next_turn_boundary").length > 0 || - this._actionStore - .activeActions("next_turn_boundary") - .some( - (action) => - action.payload.kind === "turn" && - (action.lifecycle.state === "selected" || action.lifecycle.state === "preparing"), - ) - ); - } - - private _shouldStopBeforeTurn(): boolean { - return this._steeringStopPending; - } - - private async _shouldStopAfterTurn(context: ShouldStopAfterTurnContext): Promise { - if (this._stopGoalContinuationForTerminalMessage(context.message)) { - return true; - } - try { - if (this._goals.accountAssistantMessage(context.message)) { - const message = createGoalContextMessage(this._goals.state, "budget_limit"); - const normalized = normalizeMessageContent(message.content); - await this._queuePreparedPrompt("steer", normalized.text, normalized.images, { - message, - resumeIfIdle: true, - }); - } - } catch { - // Goal accounting must not interrupt the core agent loop. - } - // Serialized refine checkpoint: in print/headless mode, run refinement - // planning+apply synchronously here — the quiescent boundary between - // turns — so it never overlaps the primary model request. - // This MUST run BEFORE threshold compaction to prevent the - // compaction model call from overlapping an in-flight refine - // plan/apply that was started at message_end. - if (this._refinement.serialized) { - // Ensure the preceding message_end processing (counter increment, - // background plan kickoff) has completed before the checkpoint. - await this._agentEventQueue; - await this._refinement._runSerializedRefineCheckpoint(); - } - if (await this._shouldStopForThresholdCompaction(context)) { - return true; - } - // Steering stops continuation only after mandatory serialized checkpoints. - // Returning true here still prevents the agent loop from starting another turn. - return this._steeringStopPending; + return this._actionQueue.steeringStopPending; } - private async _shouldStopForThresholdCompaction(context: ShouldStopAfterTurnContext): Promise { - this._compaction.resetContinuation(); - if (!this._compaction.hasPendingRequest && !(await this._thresholdCompactionNeeded(context))) { - return false; - } - - const lastMessage = this.agent.state.messages[this.agent.state.messages.length - 1]; - // A queued continuation disproves the assistant-last "task finished" heuristic, so preserve a true set above. - if (lastMessage !== undefined && lastMessage.role !== "assistant") this._compaction.requestContinuation(); - return true; + private _shouldStopBeforeTurn( + ...args: Parameters + ): ReturnType { + return this._turnPolicy.shouldStopBeforeTurn(...args); } - private async _thresholdCompactionNeeded(context: ShouldStopAfterTurnContext): Promise { - const settings = this.settingsManager.getCompactionSettings(); - if (!settings.enabled) return false; - - const contextWindow = this.model?.contextWindow ?? 0; - const compactionEntry = getLatestCompactionEntry(this.sessionManager.getBranch()); - const compactionTimestamp = compactionEntry ? new Date(compactionEntry.timestamp).getTime() : undefined; - if (compactionTimestamp !== undefined && context.message.timestamp <= compactionTimestamp) { - return false; - } - - const contextTokens = this._compaction.getThresholdContextTokens(context.message, compactionTimestamp); - if (contextTokens === undefined || !shouldCompact(contextTokens, contextWindow, settings)) { - return false; - } - - // Goal continuation takes exclusive priority over autonomous continuation, matching _getContinuationMessages. - if (this._queueGoalContinuationForThresholdCompaction(context.message)) { - this._compaction.requestContinuation(); - } else if (await this._queueAutonomousContinuationForThresholdCompaction(context.message)) { - this._compaction.requestContinuation(); - } - return true; + private _shouldStopAfterTurn( + ...args: Parameters + ): ReturnType { + return this._turnPolicy.shouldStopAfterTurn(...args); } - private _snapshotAutonomousRuntimeState(): AutonomousRuntimeSnapshot { - return { - continuationsUsed: this._autonomousState.continuationsUsed, - gateAttempts: { ...this._autonomousState.gateAttempts }, - lastGateFailure: this._autonomousState.lastGateFailure - ? { ...this._autonomousState.lastGateFailure } - : undefined, - lastGateFailureSnapshot: this._autonomousState.lastGateFailureSnapshot - ? { ...this._autonomousState.lastGateFailureSnapshot } - : undefined, - }; + private _snapshotAutonomousRuntimeState( + ...args: Parameters + ): ReturnType { + return this._autonomousContinuation.snapshotAutonomousRuntimeState(...args); } - private _restoreAutonomousRuntimeSnapshot(snapshot: AutonomousRuntimeSnapshot): void { - this._autonomousState.continuationsUsed = snapshot.continuationsUsed; - this._autonomousState.gateAttempts = { ...snapshot.gateAttempts }; - this._autonomousState.lastGateFailure = snapshot.lastGateFailure ? { ...snapshot.lastGateFailure } : undefined; - this._autonomousState.lastGateFailureSnapshot = snapshot.lastGateFailureSnapshot - ? { ...snapshot.lastGateFailureSnapshot } - : undefined; + private _restoreAutonomousRuntimeSnapshot( + ...args: Parameters + ): ReturnType { + return this._autonomousContinuation.restoreAutonomousRuntimeSnapshot(...args); } - private async _queueAutonomousContinuationForThresholdCompaction( - message: AssistantMessage, - ): Promise { - const queuedMessage = this._queuedAutonomousThresholdContinuations.get(message); - if (queuedMessage && this._continuation.messages.includes(queuedMessage)) { - return queuedMessage; - } - const snapshot = this._snapshotAutonomousRuntimeState(); - const arrivalEpoch = this._sessionInputArrivalEpoch; - const autonomousMessage = await nextAutonomousContinuation(this._autonomousState, message, { - cwd: this._cwd, - signal: this.agent.signal, - }); - if (!autonomousMessage) { - return undefined; - } - if (this._sessionInputArrivalEpoch !== arrivalEpoch) { - this._restoreAutonomousRuntimeSnapshot(snapshot); - return undefined; - } - this._queuedAutonomousThresholdContinuations.set(message, autonomousMessage); - this._queuedAutonomousContinuationSnapshots.set(autonomousMessage, snapshot); - this._continuation.track(autonomousMessage); - this._pendingThresholdCompactionAutonomousMessages.push(autonomousMessage); - const text = - typeof autonomousMessage.content === "string" - ? autonomousMessage.content - : autonomousMessage.content.map((block) => (block.type === "text" ? block.text : "")).join("\n"); - this._admitSessionInput( - this._createPreparedTurnAction("followUp", text, undefined, { - message: autonomousMessage, - }), - ); - return autonomousMessage; + private _queueAutonomousContinuationForThresholdCompaction( + ...args: Parameters + ): ReturnType { + return this._autonomousContinuation.queueAutonomousContinuationForThresholdCompaction(...args); } // The role heuristic reads an assistant-last threshold stop as "task finished" and // agent.continue() cannot resume from it, so the goal continuation is queued as a session input. - private _queueGoalContinuationForThresholdCompaction(message: AssistantMessage): boolean { - if (message.stopReason === "error" || message.stopReason === "aborted") { - return false; - } - if (this._goals.state.status !== "active" || !this._goals.state.objective) { - return false; - } - const alreadyQueued = this._queuedGoalThresholdContinuation; - if ( - alreadyQueued !== undefined && - this._actionStore.unfinishedActions().some((action) => { - if (action.payload.kind !== "turn" || primaryDeliveryRecord(action).message !== alreadyQueued) return false; - // A running continuation may already need a successor; only undelivered actions deduplicate. - return ( - action.lifecycle.state === "queued" || - action.lifecycle.state === "selected" || - action.lifecycle.state === "preparing" || - action.lifecycle.state === "committing" - ); - }) - ) { - return true; - } - try { - this._ensureGoalRuntimeActive(); - this._goals.recordContinuation(); - const goalMessage = createGoalContextMessage(this._goals.state, "continuation"); - const normalized = normalizeMessageContent(goalMessage.content); - this._admitSessionInput( - this._createPreparedTurnAction("followUp", normalized.text, normalized.images, { - message: goalMessage, - }), - ); - this._queuedGoalThresholdContinuation = goalMessage; - return true; - } catch { - return false; - } + private _queueGoalContinuationForThresholdCompaction( + ...args: Parameters + ): ReturnType { + return this._goalContinuation.queueGoalContinuationForThresholdCompaction(...args); } // Withdraws a goal continuation queued for a threshold compaction the user cancelled, // rolling back the continuationsUsed increment so the next natural stop re-queues it. private _clearQueuedGoalContinuationAfterCancelledThresholdCompaction( - queuedGoalContinuation: AgentMessage | undefined, - ): void { - if (queuedGoalContinuation === undefined) return; - const cancelled = this._cancelSessionActions( - (action) => action.payload.kind === "turn" && primaryDeliveryRecord(action).message === queuedGoalContinuation, - new Error("Queued goal continuation was cleared before delivery."), - ); - this._queuedGoalThresholdContinuation = undefined; - // A stale marker (continuation already consumed) matches no action; only an - // actual cancellation may roll back its queue-time continuationsUsed increment. - if (cancelled.length === 0) return; - this._goals.cancelContinuation(); - this._emitQueueUpdate(); - } - - private _clearQueuedAutonomousContinuations( - options: { restoreAutonomousState?: boolean; messages?: AgentMessage[] } = {}, - ): void { - const requestedMessages = options.messages ?? [...this._continuation.messages]; - const requestedMessageSet = new Set(requestedMessages); - const queuedMessages = this._continuation.messages.filter((message) => requestedMessageSet.has(message)); - if (queuedMessages.length === 0) { - return; - } - const queuedMessageSet = new Set(queuedMessages); - this._continuation.remove(queuedMessageSet); - this.agent.removeQueuedMessages((message) => queuedMessageSet.has(message)); - this._cancelSessionActions( - (action) => action.payload.kind === "turn" && queuedMessageSet.has(primaryDeliveryRecord(action).message), - new Error("Queued autonomous continuation was cleared before delivery."), - ); - this._emitQueueUpdate(); - if (options.restoreAutonomousState) { - for (const queuedMessage of queuedMessages) { - const snapshot = this._queuedAutonomousContinuationSnapshots.get(queuedMessage); - if (snapshot) { - this._restoreAutonomousRuntimeSnapshot(snapshot); - break; - } - } - } - for (const queuedMessage of queuedMessages) { - this._queuedAutonomousContinuationSnapshots.delete(queuedMessage); - } - this._pendingThresholdCompactionAutonomousMessages = this._pendingThresholdCompactionAutonomousMessages.filter( - (message) => !queuedMessageSet.has(message), - ); - if (options.messages === undefined) { - this._compaction.resetContinuation(); - } - if (!this.agent.hasQueuedMessages() && this.unfinishedActionCount === 0) { - this._cancelPostCompactionContinue(); - } + ...args: Parameters + ): ReturnType { + return this._goalContinuation.clearQueuedGoalContinuationAfterCancelledThresholdCompaction(...args); } private _clearQueuedAutonomousContinuationsAfterSkippedThresholdCompaction( - shouldContinueAfterThreshold: boolean, - queuedMessages: AgentMessage[], - ): void { - if (shouldContinueAfterThreshold) { - this._clearQueuedAutonomousContinuations({ - restoreAutonomousState: true, - messages: queuedMessages, - }); - } + ...args: Parameters< + SessionAutonomousContinuation["clearQueuedAutonomousContinuationsAfterSkippedThresholdCompaction"] + > + ): ReturnType { + return this._autonomousContinuation.clearQueuedAutonomousContinuationsAfterSkippedThresholdCompaction(...args); } /** @@ -1992,27 +1278,10 @@ export class AgentSession { * goal skill). All goal state stays host-side; the kernel only sees the * serialized snake_case response. */ - handleGoalHostRequest(type: string, payload: Record = {}): GoalHostResponse { - if (!this._includeGoals) { - throw new Error("goals are disabled in this session"); - } - switch (type) { - case "goal.get": - return goalHostResponse(this.goalState, false); - case "goal.create": { - if (typeof payload.objective !== "string") { - throw new Error("goal.create objective must be a string"); - } - if (payload.token_budget !== undefined && typeof payload.token_budget !== "number") { - throw new Error("goal.create token_budget must be an integer when provided"); - } - return goalHostResponse(this._createGoalFromHost(payload.objective, payload.token_budget), false); - } - case "goal.complete": - return goalHostResponse(this._completeGoalFromHost(), true); - default: - throw new Error(`unknown goal request type "${type}"`); - } + handleGoalHostRequest( + ...args: Parameters + ): ReturnType { + return this._goalContinuation.handleGoalHostRequest(...args); } /** @@ -2020,51 +1289,10 @@ export class AgentSession { * abort the run executing the requesting cell, so compact.run only schedules * it; _checkCompaction consumes the request at the turn boundary. */ - handleCompactHostRequest(type: string, payload: Record = {}): Record { - if (!this._includeCompactSkill) { - throw new Error("the compact skill is disabled in this session"); - } - switch (type) { - case "compact.status": { - const usage = this.getContextUsage(); - return { - tokens: usage?.tokens ?? null, - context_window: usage?.contextWindow ?? null, - percent: usage?.percent ?? null, - scheduled: this._compaction.hasPendingRequest, - }; - } - case "compact.run": { - const instructions = payload.instructions; - if (instructions !== undefined && typeof instructions !== "string") { - throw new Error("compact.run instructions must be a string when provided"); - } - if (!this.isStreaming) { - return { - scheduled: false, - reason: "no active turn; compaction can only be requested while a turn is running", - }; - } - const preparation = prepareCompaction( - this.sessionManager.getBranch(), - this.settingsManager.getCompactionSettings(), - ); - if (!preparation) { - const lastEntry = this.sessionManager.getBranch().at(-1); - return { - scheduled: false, - reason: lastEntry?.type === "compaction" ? "already compacted" : "session is too short to compact", - }; - } - this._compaction.request(instructions); - return { - scheduled: true, - note: "Compaction runs when the current turn ends; you resume automatically afterwards. Continue working normally.", - }; - } - default: - throw new Error(`unknown compact request type "${type}"`); - } + handleCompactHostRequest( + ...args: Parameters + ): ReturnType { + return this._compaction.handleCompactHostRequest(...args); } /** @@ -2084,117 +1312,14 @@ export class AgentSession { * mutate the user-level /heartbeat. */ handleRlmHeartbeatHostRequest(type: string, payload: Record = {}): Record { - const controller = this._rlmHeartbeatController; - if (!controller) { - throw new Error("RLM heartbeat skill is not available in this session"); - } - switch (type) { - case "rlm_heartbeat.list": { - const includeInactive = payload.include_inactive === true || payload.includeInactive === true; - return { - heartbeats: controller - .listRlmHeartbeats({ includeInactive }) - .map((heartbeat) => rlmHeartbeatHostResponse(heartbeat)), - }; - } - case "rlm_heartbeat.create": { - if (typeof payload.instruction !== "string") { - throw new Error("rlm_heartbeat.create instruction must be a string"); - } - if (payload.interval !== undefined && typeof payload.interval !== "string") { - throw new Error("rlm_heartbeat.create interval must be a string when provided"); - } - if (payload.label !== undefined && typeof payload.label !== "string") { - throw new Error("rlm_heartbeat.create label must be a string when provided"); - } - const deliveryMode = normalizeHeartbeatDeliveryMode(payload.delivery_mode ?? payload.deliveryMode); - return { - heartbeat: rlmHeartbeatHostResponse( - controller.createRlmHeartbeat({ - instruction: payload.instruction, - interval: payload.interval, - label: payload.label, - deliveryMode, - }), - ), - }; - } - case "rlm_heartbeat.update": { - if (typeof payload.id !== "string") { - throw new Error("rlm_heartbeat.update id must be a string"); - } - if (payload.instruction !== undefined && typeof payload.instruction !== "string") { - throw new Error("rlm_heartbeat.update instruction must be a string when provided"); - } - if (payload.interval !== undefined && typeof payload.interval !== "string") { - throw new Error("rlm_heartbeat.update interval must be a string when provided"); - } - if (payload.label !== undefined && typeof payload.label !== "string") { - throw new Error("rlm_heartbeat.update label must be a string when provided"); - } - if (payload.status !== undefined && !isRlmHeartbeatStatusUpdate(payload.status)) { - throw new Error('rlm_heartbeat.update status must be "pause" or "resume" when provided'); - } - const rawDeliveryMode = payload.delivery_mode ?? payload.deliveryMode; - const deliveryMode = normalizeHeartbeatDeliveryMode(rawDeliveryMode); - if ( - payload.instruction === undefined && - payload.interval === undefined && - payload.label === undefined && - payload.status === undefined && - rawDeliveryMode === undefined - ) { - throw new Error("rlm_heartbeat.update requires at least one field to update"); - } - const heartbeat = controller.updateRlmHeartbeat({ - id: payload.id, - instruction: payload.instruction, - interval: payload.interval, - label: payload.label, - status: payload.status, - deliveryMode, - }); - return { - heartbeat: heartbeat ? rlmHeartbeatHostResponse(heartbeat) : null, - }; - } - case "rlm_heartbeat.delete": { - if (typeof payload.id !== "string") { - throw new Error("rlm_heartbeat.delete id must be a string"); - } - const heartbeat = controller.deleteRlmHeartbeat(payload.id); - return { - heartbeat: heartbeat ? rlmHeartbeatHostResponse(heartbeat) : null, - }; - } - default: - throw new Error(`unknown RLM heartbeat request type "${type}"`); - } + return handleRlmHeartbeatHostRequest(this._rlmHeartbeatController, type, payload); } handleAgentMessageHostRequest( type: string, payload: Record = {}, ): Promise { - if (!this._agentMessageController) { - throw new Error("agent messaging is not available in this session"); - } - switch (type) { - case "agent_message.send": { - if (typeof payload.target !== "string") { - throw new Error("agent_message.send target must be a string"); - } - if (typeof payload.message !== "string") { - throw new Error("agent_message.send message must be a string"); - } - return this._agentMessageController.sendAgentMessage({ - target: assertDirectAgentMessageTarget(payload.target), - message: normalizeAgentSessionMessage(payload.message), - }); - } - default: - throw new Error(`unknown agent message request type "${type}"`); - } + return handleAgentMessageHostRequest(() => this._agentMessageController, type, payload); } handleAgentObserveHostRequest( @@ -2205,534 +1330,51 @@ export class AgentSession { | AgentObserveAgentSnapshot | AgentObserveRecentMessagesResult | Promise { - const controller = this._agentObserveController; - if (!controller) { - throw new Error("agent observation is not available in this session"); - } - switch (type) { - case "agent_observe.list": - return controller.listAgents(); - case "agent_observe.get": { - if (typeof payload.target !== "string") { - throw new Error("agent_observe.get target must be a string"); - } - return controller.getAgent(payload.target); - } - case "agent_observe.recent": { - if (typeof payload.target !== "string") { - throw new Error("agent_observe.recent target must be a string"); - } - return controller.recentMessages({ - target: payload.target, - limit: normalizeObserveLimit(payload.limit as number | undefined), - maxChars: normalizeObserveMaxChars((payload.max_chars ?? payload.maxChars) as number | undefined), - }); - } - default: - throw new Error(`unknown agent observe request type "${type}"`); - } - } - - private _createGoalFromHost(objective: string, tokenBudget: number | undefined): GoalState { - switch (this._goals.state.status) { - case "active": - throw new Error( - "cannot create a new goal because this thread already has an active goal; run `await goal.complete()` when it is achieved, or ask the user to clear it with /goal clear", - ); - case "paused": - throw new Error( - "cannot create a new goal because a paused goal exists; ask the user to resume it with /goal resume or clear it with /goal clear", - ); - case "budget_limited": - throw new Error( - "cannot create a new goal because a budget-limited goal exists; ask the user to resume it with /goal resume or clear it with /goal clear", - ); - default: - // idle, or a terminal record (complete / error): nothing pending, start fresh. - return this._startGoal(objective, tokenBudget); - } - } - - private _completeGoalFromHost(): GoalState { - // Accounting precedes the completing ipython cell, so its budget-limit - // context may already be queued and must be withdrawn before completion. - return this._goals.complete(() => this._clearQueuedGoalContexts()); - } - - private async _getGoalContinuationMessages( - context: GetContinuationMessagesContext, - signal?: AbortSignal, - ): Promise { - if (this._stopGoalContinuationForTerminalMessage(context.message)) { - return []; - } - if (signal?.aborted || this._goals.state.status !== "active" || !this._goals.state.objective) { - return []; - } - // Delegating and ending the turn is correct behavior; hold the continuation - // until descendants settle instead of re-prompting a waiting parent. - if (this._hasUnsettledRlmQuiescenceWork()) { - this._goalContinuationAwaitsRlmWork = true; - return []; - } - this._goalContinuationAwaitsRlmWork = false; - try { - this._ensureGoalRuntimeActive(context.context); - this._goals.recordContinuation(); - return [createGoalContextMessage(this._goals.state, "continuation")]; - } catch (error) { - const message = error instanceof Error ? error.message : String(error); - try { - this._goals.fail(message); - } catch { - // The continuation hook must not reject; listener failures should not crash the agent loop. - } - return []; - } + return handleAgentObserveHostRequest(this._agentObserveController, type, payload); } - private async _getContinuationMessages( - context: GetContinuationMessagesContext, - signal?: AbortSignal, - ): Promise { - if (this.queuedActionCount > 0) { - return []; - } - const arrivalEpoch = this._sessionInputArrivalEpoch; - const goalSnapshot = this._goals.checkpoint(); - const goalMessages = await this._getGoalContinuationMessages(context, signal); - if (goalMessages.length > 0 || signal?.aborted) { - if (goalMessages.length > 0 && this._sessionInputArrivalEpoch !== arrivalEpoch) { - this._goals.restore(goalSnapshot); - return []; - } - return goalMessages; - } - if ( - this._autonomousContinuationSuppressionDepth > 0 || - context.newMessages.some((message) => this._autonomousContinuationSuppressedMessages.has(message)) - ) { - return []; - } - const autonomousSnapshot = this._snapshotAutonomousRuntimeState(); - const autonomousMessage = await nextAutonomousContinuation(this._autonomousState, context.message, { - cwd: this._cwd, - signal, - }); - if (autonomousMessage && this._sessionInputArrivalEpoch !== arrivalEpoch) { - this._restoreAutonomousRuntimeSnapshot(autonomousSnapshot); - return []; - } - return autonomousMessage ? [autonomousMessage] : []; + private _getGoalContinuationMessages( + ...args: Parameters + ): ReturnType { + return this._goalContinuation.getGoalContinuationMessages(...args); } - private _lastAssistantMessage: AssistantMessage | undefined = undefined; - - private _agentMessageOutcome(agentMessageId: string): AgentMessageOutcome { - let outcome = this._agentMessageOutcomes.get(agentMessageId); - if (!outcome) { - outcome = {}; - this._agentMessageOutcomes.set(agentMessageId, outcome); - } - return outcome; + private _getContinuationMessages( + ...args: Parameters + ): ReturnType { + return this._turnPolicy.getContinuationMessages(...args); } /** * Register a delivery waiter before submitting the prompt. Delivery outcomes are not retained * for late lookup, so callers that register after admission may wait for a future use of the id. */ - waitForAgentMessagePromptDelivery(agentMessageId: string): Promise { - const outcome = this._agentMessageOutcome(agentMessageId); - outcome.delivery ??= createAgentMessageDeferred(); - return outcome.delivery.promise; + waitForAgentMessagePromptDelivery( + ...args: Parameters + ): ReturnType { + return this._messageDelivery.waitForAgentMessagePromptDelivery(...args); } private _settleAgentMessage( - agentMessageId: string | undefined, - leg: "delivery" | "completion", - error?: Error, - ): void { - if (agentMessageId === undefined) return; - const outcome = this._agentMessageOutcomes.get(agentMessageId); - if (!outcome) return; - const deferred = outcome[leg]; - if (!deferred) return; - outcome[leg] = undefined; - if (!outcome.delivery && !outcome.completion) { - this._agentMessageOutcomes.delete(agentMessageId); - } - if (error) deferred.reject(error); - else deferred.resolve(); - } - - private _rejectAgentMessage(agentMessageId: string | undefined, error: Error): void { - if (agentMessageId === undefined) return; - this._settleAgentMessage(agentMessageId, "delivery", error); - this._settleAgentMessage(agentMessageId, "completion", error); - } - - private _rejectQueuedAgentMessageDeliveries(deliveryError: Error, completionError = deliveryError): void { - for (const action of this._actionStore.unfinishedActions()) { - this._settleAgentMessage(action.agentMessageId, "delivery", deliveryError); - this._settleAgentMessage(action.agentMessageId, "completion", completionError); - } - } - - private _capturingCancelledAction(message: AgentMessage): QueuedSessionAction | undefined { - return this._actionStore - .ownedActions() - .find( - (action) => - action.lifecycle.state === "cancelled" && - action.payload.kind === "turn" && - action.payload.captureRunMessages?.has(message) === true, - ); - } - - private _hasCancelledDispatchCapture(): boolean { - return this._actionStore - .ownedActions() - .some( - (action) => - action.lifecycle.state === "cancelled" && - action.payload.kind === "turn" && - action.payload.captureRunMessages !== undefined, - ); - } - - private _handleAgentEvent = (event: AgentEvent): void => { - this._retry.observeAgentEnd(event); - if (event.type === "message_start" || event.type === "message_end") { - for (const action of this._actionStore.ownedActions()) { - if ( - action.payload.kind !== "turn" || - !action.payload.captureRunMessages || - action.payload.cancelledDispatchEnded - ) { - continue; - } - const primary = primaryDeliveryRecord(action); - if (event.message === primary.message || primary.started) { - action.payload.captureRunMessages.add(event.message); - } - } - } else if (event.type === "agent_end") { - const captured = new Set(); - for (const action of this._actionStore.ownedActions()) { - if (action.payload.kind === "turn" && action.payload.captureRunMessages) { - for (const message of action.payload.captureRunMessages) captured.add(message); - action.payload.cancelledDispatchEnded = true; - } - } - if (captured.size > 0) { - this.agent.state.messages = this.agent.state.messages.filter((message) => !captured.has(message)); - } - } - if (event.type === "message_start" && (event.message.role === "user" || event.message.role === "custom")) { - for (const action of this._actionStore.actionsForMessage(event.message)) { - const record = - action.payload.kind === "turn" - ? action.payload.records.find((candidate) => candidate.message === event.message) - : undefined; - if (record) record.started = true; - if (record?.role === "primary") { - this._actionStore.ticketFor(action).settleDelivered({ status: "delivered" }); - this._settleAgentMessage(action.agentMessageId, "delivery"); - } - } - } else if (event.type === "message_end" && (event.message.role === "user" || event.message.role === "custom")) { - for (const action of this._actionStore.actionsForMessage(event.message)) { - const record = - action.payload.kind === "turn" - ? action.payload.records.find((candidate) => candidate.message === event.message) - : undefined; - if (record) record.durable = true; - if (record?.role === "primary" && action.lifecycle.state === "committing") { - transitionSessionAction(action, { - state: "running", - execution: "agent_turn", - }); - this._notifySessionInputCheckpointChange(); - this._emitQueueUpdate(); - } - } - } - this._agentEventQueue = this._agentEventQueue.then( - () => this._processAgentEvent(event), - () => this._processAgentEvent(event), - ); - this._agentEventQueue.catch(() => {}); - }; - - private _findLastAssistantInMessages(messages: AgentMessage[]): AssistantMessage | undefined { - for (let i = messages.length - 1; i >= 0; i--) { - const message = messages[i]; - if (message.role === "assistant") { - return message as AssistantMessage; - } - } - return undefined; + ...args: Parameters + ): ReturnType { + return this._messageDelivery.settleAgentMessage(...args); } - private _addLoginGuidanceToAuthError(event: AgentEvent): void { - const message = - event.type === "message_end" && event.message.role === "assistant" - ? (event.message as AssistantMessage) - : event.type === "agent_end" - ? this._findLastAssistantInMessages(event.messages) - : undefined; - if (!message || message.stopReason !== "error" || !message.errorMessage) { - return; - } - if (!isLikelyAuthenticationError(message.errorMessage)) { - return; - } - message.errorMessage = addLoginGuidanceToAuthError(message.errorMessage); + private _rejectAgentMessage( + ...args: Parameters + ): ReturnType { + return this._messageDelivery.rejectAgentMessage(...args); } - private async _processAgentEvent(event: AgentEvent): Promise { - let clearedDispatchEnded = false; - if ((event.type === "message_start" || event.type === "message_end") && event.message.role === "toolResult") { - this._applyLateIpythonSentAgentMessages(event.message); - } - if (event.type === "message_start" || event.type === "message_end") { - const cleared = this._capturingCancelledAction(event.message); - if (cleared?.payload.kind === "turn" && cleared.payload.captureRunMessages) { - const captured = cleared.payload.captureRunMessages; - this.agent.state.messages = this.agent.state.messages.filter((message) => !captured.has(message)); - return; - } - } - if (event.type === "agent_end") { - const cleared = this._actionStore - .ownedActions() - .filter( - (action) => - action.lifecycle.state === "cancelled" && - action.payload.kind === "turn" && - action.payload.captureRunMessages !== undefined, - ); - if (cleared.length > 0) { - clearedDispatchEnded = true; - const removed = new Set( - cleared.flatMap((action) => - action.payload.kind === "turn" ? [...(action.payload.captureRunMessages ?? [])] : [], - ), - ); - this.agent.state.messages = this.agent.state.messages.filter((message) => !removed.has(message)); - (this.agent.state as { errorMessage?: string }).errorMessage = undefined; - this._lastAssistantMessage = undefined; - for (const action of cleared) this._actionStore.releaseTerminal(action); - this._notifySessionInputCheckpointChange(); - this._retry.resolve(); - } - } - - if (event.type === "message_start" && startsAgentRun(event.message)) { - this._compaction.resetOverflowRecovery(); - } - - await this._emitExtensionEvent(event); - if (event.type === "message_start" || event.type === "message_end") { - const cleared = this._capturingCancelledAction(event.message); - if (cleared?.payload.kind === "turn" && cleared.payload.captureRunMessages) { - const captured = cleared.payload.captureRunMessages; - this.agent.state.messages = this.agent.state.messages.filter((message) => !captured.has(message)); - return; - } - } - - this._addLoginGuidanceToAuthError(event); - - this._emit(event); - - if (event.type === "message_end") { - if (event.message.role === "custom") { - this.sessionManager.appendCustomMessageEntry( - event.message.customType, - event.message.content, - event.message.display, - event.message.details, - ); - } else if ( - event.message.role === "user" || - event.message.role === "assistant" || - event.message.role === "toolResult" - ) { - this.sessionManager.appendMessage(event.message); - } - - if (event.message.role === "assistant") { - this._lastAssistantMessage = event.message; - - const assistantMsg = event.message as AssistantMessage; - if (assistantMsg.stopReason !== "error") { - addAutonomousUsage(this._autonomousState, assistantMsg.usage); - } - if (assistantMsg.stopReason !== "error" && assistantMsg.stopReason !== "aborted") { - this._refinement.observeAssistantEnd(); - // In serialized mode, kick off background refinement planning - // immediately after the primary stream finishes, while tools - // are still executing. The plan is awaited at shouldStopAfterTurn - // before applying, so planning overlaps tools only — never another - // model request. - } - if (assistantMsg.stopReason !== "error") { - this._compaction.resetOverflowRecovery(); - } - this._retry.observeAssistantEnd(assistantMsg); - if (this._goals.accountAssistantMessage(assistantMsg)) { - const message = createGoalContextMessage(this._goals.state, "budget_limit"); - const normalized = normalizeMessageContent(message.content); - await this._queuePreparedPrompt("steer", normalized.text, normalized.images, { - message, - resumeIfIdle: true, - }); - } - } - } - - if (clearedDispatchEnded) { - return; - } - - if (event.type === "agent_end") { - const msg = - this._lastAssistantMessage ?? - (this._retry.isRetrying ? this._findLastAssistantInMessages(event.messages) : undefined); - this._lastAssistantMessage = undefined; - if (!msg) { - this._retry.resolve(); - return; - } - - const retry = this._retry.retryError(msg); - if (retry && (await retry)) return; - - const compactionWillRetry = await this._checkCompaction(msg); - if (compactionWillRetry && this._retry.attempt > 0) { - return; - } - this._retry.finishActiveRetryWithFailure(msg); - this._retry.resolve(); - if (!compactionWillRetry) { - this._finishGoalForTerminalAssistantMessage(msg); - // In serialized mode, agent-callable refine.run is serviced - // at the shouldStopAfterTurn boundary, not here at agent_end. - if (!this._refinement.serialized) { - const consumedRequestedRefine = this._refinement._consumePendingRequestedRefine(); - if (!consumedRequestedRefine) { - this._refinement._scheduleAutoRefineAfterAgentEnd(); - } - } - } - } + private _hasCancelledDispatchCapture( + ...args: Parameters + ): ReturnType { + return this._events.hasCancelledDispatchCapture(...args); } private _findLastAssistantMessage(): AssistantMessage | undefined { - const messages = this.agent.state.messages; - for (let i = messages.length - 1; i >= 0; i--) { - const msg = messages[i]; - if (msg.role === "assistant") { - return msg as AssistantMessage; - } - } - return undefined; - } - - private _replaceMessageInPlace(target: AgentMessage, replacement: AgentMessage): void { - // Agent-core stores the finalized message object in its state before emitting message_end. - // SessionManager persistence happens later in _processAgentEvent() with event.message. - // Mutating this object in place keeps agent state, later turn/agent events, listeners, - // and the eventual SessionManager.appendMessage(event.message) persistence in sync. - if (target === replacement) { - return; - } - - const targetRecord = target as unknown as Record; - for (const key of Object.keys(targetRecord)) { - delete targetRecord[key]; - } - Object.assign(targetRecord, replacement); - } - - private async _emitExtensionEvent(event: AgentEvent): Promise { - if (event.type === "agent_start") { - this._turnIndex = 0; - this.sessionManager.recordGitStateIfChanged(); - await this._extensionRunner.emit({ type: "agent_start" }); - } else if (event.type === "agent_end") { - // Also capture at end of turn so commits made during the run (e.g. via a bash tool) land. - this.sessionManager.recordGitStateIfChanged(); - await this._extensionRunner.emit({ - type: "agent_end", - messages: event.messages, - }); - } else if (event.type === "turn_start") { - const extensionEvent: TurnStartEvent = { - type: "turn_start", - turnIndex: this._turnIndex, - timestamp: Date.now(), - }; - await this._extensionRunner.emit(extensionEvent); - } else if (event.type === "turn_end") { - const extensionEvent: TurnEndEvent = { - type: "turn_end", - turnIndex: this._turnIndex, - message: event.message, - toolResults: event.toolResults, - }; - await this._extensionRunner.emit(extensionEvent); - this._turnIndex++; - } else if (event.type === "message_start") { - const extensionEvent: MessageStartEvent = { - type: "message_start", - message: event.message, - }; - await this._extensionRunner.emit(extensionEvent); - } else if (event.type === "message_update") { - const extensionEvent: MessageUpdateEvent = { - type: "message_update", - message: event.message, - assistantMessageEvent: event.assistantMessageEvent, - }; - await this._extensionRunner.emit(extensionEvent); - } else if (event.type === "message_end") { - const extensionEvent: MessageEndEvent = { - type: "message_end", - message: event.message, - }; - const replacement = await this._extensionRunner.emitMessageEnd(extensionEvent); - if (replacement) { - this._replaceMessageInPlace(event.message, replacement); - } - } else if (event.type === "tool_execution_start") { - const extensionEvent: ToolExecutionStartEvent = { - type: "tool_execution_start", - toolCallId: event.toolCallId, - toolName: event.toolName, - args: event.args, - }; - await this._extensionRunner.emit(extensionEvent); - } else if (event.type === "tool_execution_update") { - const extensionEvent: ToolExecutionUpdateEvent = { - type: "tool_execution_update", - toolCallId: event.toolCallId, - toolName: event.toolName, - args: event.args, - partialResult: event.partialResult, - }; - await this._extensionRunner.emit(extensionEvent); - } else if (event.type === "tool_execution_end") { - const extensionEvent: ToolExecutionEndEvent = { - type: "tool_execution_end", - toolCallId: event.toolCallId, - toolName: event.toolName, - result: event.result, - isError: event.isError, - }; - await this._extensionRunner.emit(extensionEvent); - } + return this._events.findLastAssistantInMessages(this.agent.state.messages); } /** @@ -2740,15 +1382,8 @@ export class AgentSession { * Session persistence is handled internally (saves messages on message_end). * Multiple listeners can be added. Returns unsubscribe function for this listener. */ - subscribe(listener: AgentSessionEventListener): () => void { - this._eventListeners.push(listener); - - return () => { - const index = this._eventListeners.indexOf(listener); - if (index !== -1) { - this._eventListeners.splice(index, 1); - } - }; + subscribe(...args: Parameters): ReturnType { + return this._events.subscribe(...args); } /** @@ -2756,20 +1391,20 @@ export class AgentSession { * User listeners are preserved and will receive events again after resubscribe(). * Used internally during operations that need to pause event processing. */ - private _disconnectFromAgent(): void { - if (this._unsubscribeAgent) { - this._unsubscribeAgent(); - this._unsubscribeAgent = undefined; - } + private _disconnectFromAgent( + ...args: Parameters + ): ReturnType { + return this._events.disconnectFromAgent(...args); } /** * Reconnect to agent events after _disconnectFromAgent(). * Preserves all existing listeners. */ - private _reconnectToAgent(): void { - if (this._unsubscribeAgent) return; // Already connected - this._unsubscribeAgent = this.agent.subscribe(this._handleAgentEvent); + private _reconnectToAgent( + ...args: Parameters + ): ReturnType { + return this._events.reconnectToAgent(...args); } /** @@ -2848,21 +1483,17 @@ export class AgentSession { // resolution cannot write harness state or re-subscribe handlers. this._refinement.dispose(); this._children.dispose(); - this._pendingNextTurnMessages = []; + this._pendingContext.dispose(); const deliveryError = new Error("Session disposed before prompt delivery."); const completionError = new Error("Session disposed before prompt completion."); - this._rejectQueuedAgentMessageDeliveries(deliveryError, completionError); - for (const [agentMessageId, outcome] of this._agentMessageOutcomes) { - if (outcome.delivery) this._settleAgentMessage(agentMessageId, "delivery", deliveryError); - if (outcome.completion) this._settleAgentMessage(agentMessageId, "completion", completionError); - } + this._messageDelivery.dispose(deliveryError, completionError); this._cancelSessionActions(() => true, deliveryError); this.agent.clearAllQueues(); this._extensionRunner.invalidate( "This extension ctx is stale after session replacement or reload. Do not use a captured pi or command ctx after ctx.newSession(), ctx.fork(), ctx.switchSession(), or ctx.reload(). For newSession, fork, and switchSession, move post-replacement work into withSession and use the ctx passed to withSession. For reload, do not use the old ctx after await ctx.reload().", ); this._disconnectFromAgent(); - this._eventListeners = []; + this._events.dispose(); cleanupSessionResources(this.sessionId); } finally { void this._startDisposeCallbacks(); @@ -2927,30 +1558,23 @@ export class AgentSession { } get isCompacting(): boolean { - return this._compaction.isRunning || this._branchSummaryAbortController !== undefined; + return this._compaction.isRunning || this._history.isSummarizing; } get messages(): AgentMessage[] { return this.agent.state.messages; } - buildSessionContext(): SessionContext { - const context = this.sessionManager.buildSessionContext(); - for (const message of context.messages) { - this._applyLateIpythonSentAgentMessages(message); - } - this._mergeUnpersistedOutcomes(context.messages); - return context; + buildSessionContext( + ...args: Parameters + ): ReturnType { + return this._harnessContext.buildSessionContext(...args); } - private _mergeUnpersistedOutcomes(messages: AgentMessage[]): void { - for (const outcome of this._unpersistedOutcomes) { - let insertAt = messages.length; - while (insertAt > 0 && messages[insertAt - 1]!.timestamp > outcome.timestamp) { - insertAt -= 1; - } - messages.splice(insertAt, 0, outcome); - } + private _mergeUnpersistedOutcomes( + ...args: Parameters + ): ReturnType { + return this._harnessContext.mergeUnpersistedOutcomes(...args); } get steeringMode(): "all" | "one-at-a-time" { @@ -2989,31 +1613,32 @@ export class AgentSession { return this._goals.current; } - getAutonomousStatus(): AgentAutonomousStatus { - return autonomousStatus(this._autonomousState); + getAutonomousStatus( + ...args: Parameters + ): ReturnType { + return this._autonomousContinuation.getAutonomousStatus(...args); } - recordHostAutonomousContinuation(): void { - addAutonomousContinuation(this._autonomousState); + recordHostAutonomousContinuation( + ...args: Parameters + ): ReturnType { + return this._autonomousContinuation.recordHostAutonomousContinuation(...args); } - async refreshAutonomousGates(): Promise { - await refreshAutonomousQualityGates(this._autonomousState, { - cwd: this._cwd, - }); + refreshAutonomousGates( + ...args: Parameters + ): ReturnType { + return this._autonomousContinuation.refreshAutonomousGates(...args); } - private async _runWithAutonomousContinuationSuppressed(fn: () => Promise): Promise { - this._autonomousContinuationSuppressionDepth++; - try { - return await fn(); - } finally { - this._autonomousContinuationSuppressionDepth--; - } + private _runWithAutonomousContinuationSuppressed(fn: () => Promise): Promise { + return this._autonomousContinuation.runWithAutonomousContinuationSuppressed(fn); } - private _markAutonomousContinuationSuppressed(message: AgentMessage): void { - this._autonomousContinuationSuppressedMessages.add(message); + private _markAutonomousContinuationSuppressed( + ...args: Parameters + ): ReturnType { + return this._autonomousContinuation.markAutonomousContinuationSuppressed(...args); } get scopedModels(): ReadonlyArray<{ @@ -3024,7 +1649,7 @@ export class AgentSession { } setScopedModels(scopedModels: Array<{ model: Model; thinkingLevel?: ThinkingLevel }>): void { - this._scopedModels = scopedModels; + this._modelSelection.setScopedModels(scopedModels); } get promptTemplates(): ReadonlyArray { @@ -3035,53 +1660,16 @@ export class AgentSession { return this._tools.rebuildSystemPrompt(toolNames); } - private _refreshExtensionSystemPrompt(extensionPrompt: string, baseSnapshot: string): string { - return this._tools.refreshExtensionSystemPrompt(extensionPrompt, baseSnapshot); - } - - private _finishSubmissionNormalization( - text: string, - images: ImageContent[] | undefined, - policy: SubmissionNormalizationPolicy, - ): NormalizedSubmission { - let expandedText = text; - if (policy.expandSkills) expandedText = this._expandSkillCommand(expandedText); - if (policy.expandPromptTemplates) { - expandedText = expandPromptTemplate(expandedText, [...this.promptTemplates]); - } - return { kind: "prompt", text: expandedText, images }; + private _refreshExtensionSystemPrompt( + ...args: Parameters + ): ReturnType { + return this._tools.refreshExtensionSystemPrompt(...args); } private _normalizeSubmission( - text: string, - images: ImageContent[] | undefined, - policy: SubmissionNormalizationPolicy, - ): NormalizedSubmission | Promise { - if (policy.parseSessionCommands) { - const command = parseSessionSlashCommand(text); - if (command) return { kind: "sessionCommand", text, images, command }; - } - - if (text.startsWith("/")) { - if (policy.extensionCommands === "execute") { - const completion = this._executeExtensionCommand(text); - if (completion) return { kind: "extensionCommand", completion }; - } else if (policy.extensionCommands === "reject") { - this._throwIfExtensionCommand(text); - } - } - - if (policy.inputSource !== undefined && this._extensionRunner.hasHandlers("input")) { - return this._extensionRunner.emitInput(text, images, policy.inputSource).then((result) => { - if (result.action === "handled") return { kind: "handled" }; - if (result.action === "transform") { - return this._finishSubmissionNormalization(result.text, result.images ?? images, policy); - } - return this._finishSubmissionNormalization(text, images, policy); - }); - } - - return this._finishSubmissionNormalization(text, images, policy); + ...args: Parameters + ): ReturnType { + return this._submissionNormalizer.normalizeSubmission(...args); } private async _runPreTurnCompaction(): Promise { @@ -3089,20 +1677,6 @@ export class AgentSession { if (lastAssistant) await this._checkCompaction(lastAssistant, false, false); } - private _applyPreparedSystemPrompt( - preparation: PreparedPromptPreparation | undefined, - preserveEmptyExtensionPrompt: boolean, - ): void { - const extensionPrompt = preparation?.result?.systemPrompt; - const hasExtensionPrompt = preserveEmptyExtensionPrompt - ? extensionPrompt !== undefined - : Boolean(extensionPrompt); - this.agent.state.systemPrompt = - hasExtensionPrompt && extensionPrompt !== undefined && preparation !== undefined - ? this._refreshExtensionSystemPrompt(extensionPrompt, preparation.basePromptSnapshot) - : this._baseSystemPrompt; - } - private _canStartSessionActionImmediately(): boolean { return ( !this.isStreaming && @@ -3133,98 +1707,22 @@ export class AgentSession { return this._prompt(text, { ...options, returnAfterAccepted: true }); } - async promptAndWait(text: string, options?: PromptOptions): Promise { - const agentMessageId = options?.agentMessageId ?? `prompt-wait:${randomUUID()}`; - if (this._agentMessageOutcomes.get(agentMessageId)?.completion) { - throw new Error(`Prompt completion id is already in use: ${agentMessageId}`); - } - const outcome = this._agentMessageOutcome(agentMessageId); - outcome.completion = createAgentMessageDeferred(); - const completion = outcome.completion.promise; - const signal = options?.signal; - let cancelQueuedPrompt: (() => void) | undefined; - try { - await this.promptUntilAccepted(text, { ...options, agentMessageId }); - if (signal) { - cancelQueuedPrompt = () => { - const error = new Error("Prompt was cancelled before it started."); - const cancelled = this._cancelSessionActions( - (action) => action.agentMessageId === agentMessageId && action.payload.kind === "turn", - error, - ); - if (cancelled.length > 0) { - this._settleAgentMessage(agentMessageId, "completion", error); - } - }; - signal.addEventListener("abort", cancelQueuedPrompt, { once: true }); - if (signal.aborted) cancelQueuedPrompt(); - } - await completion; - } catch (error) { - this._settleAgentMessage(agentMessageId, "completion", this._asError(error)); - throw error; - } finally { - if (signal && cancelQueuedPrompt) { - signal.removeEventListener("abort", cancelQueuedPrompt); - } - } + promptAndWait( + ...args: Parameters + ): ReturnType { + return this._messageDelivery.promptAndWait(...args); } - async acceptAgentMessagePrompt(text: string, options?: PromptOptions): Promise { - const customMessage = - options?.customMessage && isAgentSessionMessage(options.customMessage) ? options.customMessage : undefined; - const clearEpoch = this._agentMessageClearEpoch; - const admissionCommitted = () => { - options?.admissionCommitted?.(); - if (clearEpoch !== this._agentMessageClearEpoch) { - throw new Error("Agent message was cleared before admission"); - } - }; - if ( - this._inputScheduler.suspended && - this._isBusyForSessionInput("preflight") && - options?.queueIfBusy === true && - options.streamingBehavior - ) { - admissionCommitted(); - const queued = await this.queueAgentMessagePrompt(text, options.streamingBehavior, customMessage); - options.preflightResult?.(queued, queued); - return; - } - await this._prompt(text, { - ...options, - resumeIfIdle: false, - expandPromptTemplates: false, - skipInputHandlers: true, - skipPrePromptWork: true, - returnAfterAccepted: true, - agentMessageId: options?.agentMessageId ?? customMessage?.details.id ?? parseAgentSessionMessagePromptId(text), - customMessage, - admissionCommitted, - }); - if (customMessage?.details.fromRelationship === "parent") this._childState.resetReply(); + acceptAgentMessagePrompt( + ...args: Parameters + ): ReturnType { + return this._promptSubmission.acceptAgentMessagePrompt(...args); } - async queueAgentMessagePrompt( - text: string, - streamingBehavior: "steer" | "followUp", - customMessage?: AgentSessionMessage, - ): Promise { - const agentMessageId = customMessage?.details.id ?? parseAgentSessionMessagePromptId(text); - if (streamingBehavior === "steer") { - await this._queuePreparedPrompt("steer", text, undefined, { - agentMessageId, - message: customMessage, - }); - 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._childState.resetReply(); - return queued; + queueAgentMessagePrompt( + ...args: Parameters + ): ReturnType { + return this._promptSubmission.queueAgentMessagePrompt(...args); } async promptHeartbeat(job: AgentCronJob, options?: PromptOptions): Promise { @@ -3236,493 +1734,46 @@ export class AgentSession { }); } - private _isRlmTerminalNotice(message: CustomMessage): boolean { - return ( - message.customType === RLM_CHILD_TERMINAL_NOTICE_CUSTOM_TYPE || - message.customType === RLM_CHILD_FAILURE_CUSTOM_TYPE - ); - } - - private _assertRlmTerminalNotice(message: CustomMessage): void { - if (!this._isRlmTerminalNotice(message)) { - throw new Error("Deferred terminal admission only accepts RLM child terminal notices."); - } - } - - private _isRlmTerminalNoticeAction(action: QueuedSessionAction): boolean { - if (action.payload.kind !== "turn") return false; - const message = primaryDeliveryRecord(action).message; - return message.role === "custom" && this._isRlmTerminalNotice(message); - } - - private _hasDeferredRlmTerminalNotices(): boolean { - return this._pendingNextTurnMessages.some((message) => this._isRlmTerminalNotice(message)); - } - - private _enqueueRlmTerminalNoticeAction(message: CustomMessage): void { - this._assertRlmTerminalNotice(message); - const action = this._createPreparedTurnAction("followUp", message.content as string, undefined, { - message, - suppressAutonomousContinuation: true, - resumeIfIdle: false, - source: "internal", - executionPolicy: createTurnExecutionPolicy("injected"), - queueVisible: false, - }); - this._durableRlmTerminalNoticeActionIds.add(action.id); - try { - const result = this._admitSessionInput(action, { wake: false }); - if (!result.accepted) throw new Error("RLM child terminal notice was not admitted."); - } catch (error) { - this._durableRlmTerminalNoticeActionIds.delete(action.id); - throw error; - } - } - - private _flushDeferredRlmTerminalNotices(): void { - if ( - this._inputScheduler.admissionPaused || - this._inputScheduler.suspended || - this._inputScheduler.queuedWorkPauseCount > 0 || - this._disposed || - this._disposing - ) { - return; - } - while (true) { - const index = this._pendingNextTurnMessages.findIndex((message) => this._isRlmTerminalNotice(message)); - if (index < 0) break; - const message = this._pendingNextTurnMessages[index]; - try { - this._enqueueRlmTerminalNoticeAction(message); - } catch { - return; - } - this._pendingNextTurnMessages.splice(index, 1); - } - this._scheduleSessionInputPump(); - } - - private async _acquireRlmTerminalNoticeRetentionFence(): Promise<{ owner: symbol; release(): void } | undefined> { - const disposeSignal = this._commitFence.disposeSignal; - while (!this._disposed && !this._disposing && !disposeSignal.aborted) { - if (this._inputScheduler.queuedWorkPauseCount > 0) { - let wake = () => {}; - const pauseReleased = new Promise((resolve) => { - wake = resolve; - this._sessionInputCheckpointWaiters.add(resolve); - }); - try { - await waitForPromiseOrAbort(pauseReleased, disposeSignal, "Terminal notice retention cancelled"); - } catch { - return undefined; - } finally { - this._sessionInputCheckpointWaiters.delete(wake); - } - continue; - } - let fence: { owner: symbol; release(): void }; - try { - fence = await this._acquireSessionActionCommitFence(disposeSignal); - } catch { - return undefined; - } - if (this._inputScheduler.queuedWorkPauseCount === 0 && !this._disposed && !this._disposing) return fence; - fence.release(); - } - return undefined; + private _isRlmTerminalNoticeAction( + ...args: Parameters + ): ReturnType { + return this._pendingContext.isRlmTerminalNoticeAction(...args); } - private async _deferRlmTerminalNotice(message: CustomMessage): Promise { - this._assertRlmTerminalNotice(message); - const fence = await this._acquireRlmTerminalNoticeRetentionFence(); - if (!fence) return; - try { - if (this._disposed || this._disposing) return; - this._pendingNextTurnMessages.push(cloneCustomMessage(message)); - this._flushDeferredRlmTerminalNotices(); - } finally { - fence.release(); - } - } - - private _demoteRlmTerminalNoticeActions(): void { - const actions = this._actionStore - .clearableActions() - .filter((action) => this._durableRlmTerminalNoticeActionIds.has(action.id)); - if (actions.length === 0) return; - for (const action of actions) { - if (!this._isRlmTerminalNoticeAction(action)) continue; - const message = primaryDeliveryRecord(action).message; - if (message.role === "custom") this._pendingNextTurnMessages.push(cloneCustomMessage(message)); - } - const ids = new Set(actions.map((action) => action.id)); - this._cancelSessionActions( - (action) => ids.has(action.id), - new Error("RLM child terminal notice deferred across session input suspension."), - actions, - ); - for (const id of ids) this._durableRlmTerminalNoticeActionIds.delete(id); + private _hasDeferredRlmTerminalNotices( + ...args: Parameters + ): ReturnType { + return this._pendingContext.hasDeferredRlmTerminalNotices(...args); } - /** - * The kernel read the command's result before the notice reached the model, so - * the notice has nothing left to report: drop it while it is still queued. - * Delivered notices are no longer clearable, which makes this a no-op. - */ - private _withdrawAsyncBashCompletionNotice(details: { pid: number; command: string }): void { - // One read withdraws one notice: pid reuse can queue an identical key twice, - // and the read belongs to the older handle, which is the earlier notice. - const notice = this._actionStore - .clearableActions() - .find((action) => this._isAsyncBashCompletionActionFor(action, details)); - if (!notice) return; - this._cancelSessionActions( - (action) => action === notice, - new Error("Background command completion notice withdrawn: the kernel read the result first."), - ); - this._emitQueueUpdate(); + private _flushDeferredRlmTerminalNotices( + ...args: Parameters + ): ReturnType { + return this._pendingContext.flushDeferredRlmTerminalNotices(...args); } - private _isAsyncBashCompletionActionFor( - action: QueuedSessionAction, - details: { pid: number; command: string }, - ): boolean { - if (action.payload.kind !== "turn") return false; - const message = primaryDeliveryRecord(action).message; - if (message.role !== "custom" || message.customType !== ASYNC_BASH_COMPLETION_CUSTOM_TYPE) return false; - // pids are reused across handles, so the command has to match too. - const completion = message.details as AsyncBashCompletionDetails | undefined; - return completion?.pid === details.pid && completion.command === details.command; + private _deferRlmTerminalNotice( + ...args: Parameters + ): ReturnType { + return this._pendingContext.deferRlmTerminalNotice(...args); } - private async _promptInjectedMessage( - text: string, - message: CustomMessage, - options?: InternalPromptOptions & { executionPolicy?: TurnExecutionPolicy }, - ): Promise { - if (!this.isStreaming && options?.resumeIfIdle) this._resumeSessionInputAdmission(); - const admissionEpoch = this._inputScheduler.epoch; - const admissionFence = await this._acquireDirectTurnAdmissionFence(options?.signal).catch((error: unknown) => { - throwIfPromptAdmissionCancelled(options?.signal); - throw error; - }); - const reportPreflight = oncePreflight(options?.preflightResult); - try { - throwIfPromptAdmissionCancelled(options?.signal); - if (admissionEpoch !== this._inputScheduler.epoch) { - throw new Error("Injected session input was invalidated before admission"); - } - options?.admissionCommitted?.(); - const queueForStreaming = this.isStreaming; - const queueForBusy = options?.queueIfBusy === true && this._isBusyForSessionInput("preflight"); - const visibleQueued = queueForStreaming || queueForBusy; - if (visibleQueued && !options?.streamingBehavior) { - const stateDescription = queueForStreaming ? "Agent is already processing" : "Agent has queued work"; - throw new Error( - `${stateDescription}. Specify streamingBehavior ('steer' or 'followUp') to queue the message.`, - ); - } - const schedule = options?.streamingBehavior ?? "followUp"; - const prefixMessages = visibleQueued ? this._takePendingNextTurnMessages() : undefined; - const action = this._createPreparedTurnAction(schedule, text, undefined, { - message, - prefixMessages, - queueKey: options?.followUpQueueKey, - previewLabel: injectedMessagePreviewLabel(message), - suppressAutonomousContinuation: options?.suppressAutonomousContinuation, - resumeIfIdle: - !visibleQueued || - options?.resumeIfIdle || - (options?.queueIfBusy === true && canSelectSessionAction(this._runtimeActivity())), - source: options?.source ?? "internal", - executionPolicy: - options?.executionPolicy ?? - (visibleQueued ? createTurnExecutionPolicy("queued") : createTurnExecutionPolicy("injected")), - queueVisible: visibleQueued, - }); - const result = this._admitSessionInput(action, { - immediatelyEligible: !visibleQueued, - }); - admissionFence.release(); - if (!result.accepted || !result.ticket) { - if (prefixMessages) this._pendingNextTurnMessages.unshift(...prefixMessages); - reportPreflight(false, false); - return; - } - if (result.disposition === "queued") { - reportPreflight(true, true); - } else { - void result.ticket.delivered.then( - () => reportPreflight(true), - () => reportPreflight(false), - ); - } - if (options?.returnAfterAccepted) { - if (result.disposition === "starts_when_admitted") await result.ticket.delivered; - return; - } - if (visibleQueued) return; - await result.ticket.completed; - } catch (error) { - reportPreflight(false); - throw error; - } finally { - admissionFence.release(); - } + private _demoteRlmTerminalNoticeActions( + ...args: Parameters + ): ReturnType { + return this._pendingContext.demoteRlmTerminalNoticeActions(...args); } - private async _prompt(text: string, options?: InternalPromptOptions): Promise { - const resumeSuspendedInput = options?.resumeIfIdle !== false; - if (!this.isStreaming) { - if (resumeSuspendedInput) this._resumeSessionInputAdmission(); - this._assertSessionActionAdmissionAvailable(); - } - const admissionEpoch = this._inputScheduler.epoch; - const commitFence = this.isStreaming - ? undefined - : await this._acquireDirectTurnAdmissionFence(options?.signal).catch((error: unknown) => { - throwIfPromptAdmissionCancelled(options?.signal); - throw error; - }); - const reportPreflight = oncePreflight(options?.preflightResult); - const run = async () => { - try { - throwIfPromptAdmissionCancelled(options?.signal); - if (!resumeSuspendedInput && admissionEpoch !== this._inputScheduler.epoch) { - throw new Error("Session input was invalidated before admission"); - } - options?.admissionCommitted?.(); - const isInternalPrompt = options?.internalPrompt === true; - const expandPromptTemplates = isInternalPrompt ? false : (options?.expandPromptTemplates ?? true); - const normalizationResult = this._normalizeSubmission(text, options?.images, { - parseSessionCommands: !isInternalPrompt && !options?.skipPrePromptWork, - extensionCommands: expandPromptTemplates ? "execute" : "ignore", - inputSource: - !isInternalPrompt && !options?.skipInputHandlers ? (options?.source ?? "interactive") : undefined, - expandSkills: expandPromptTemplates, - expandPromptTemplates, - }); - const normalized = normalizationResult instanceof Promise ? await normalizationResult : normalizationResult; - // Async input handlers ran between the admission check above and - // admission itself; re-check so content invalidated during that - // await (e.g. a cron job cancelled or updated) is not admitted. - if (normalizationResult instanceof Promise) options?.admissionCommitted?.(); - if (normalized.kind === "extensionCommand") { - commitFence?.release(); - reportPreflight(true); - void normalized.completion.then( - () => this._settleAgentMessage(options?.agentMessageId, "completion"), - (error) => this._settleAgentMessage(options?.agentMessageId, "completion", error), - ); - void normalized.completion.catch(() => undefined); - if (!options?.returnAfterAccepted) await normalized.completion.catch(() => undefined); - return; - } - if (normalized.kind === "handled") { - commitFence?.release(); - reportPreflight(true); - this._settleAgentMessage(options?.agentMessageId, "completion"); - return; - } - - const pendingOwnedWork = this._actionStore.unfinishedActions().length > 0; - const wasRuntimeBusy = this.isStreaming || this.isCompacting || this.isRetrying || this.isBashRunning; - const wasBusy = wasRuntimeBusy || pendingOwnedWork; - if (normalized.kind === "sessionCommand") { - const schedule = options?.streamingBehavior ?? (this.isStreaming ? "steer" : "followUp"); - const action = createSessionCommandAction( - normalized.text, - normalized.command, - normalized.images, - schedule, - { - agentMessageId: options?.agentMessageId, - source: isInternalPrompt ? "internal" : (options?.source ?? "interactive"), - }, - ); - const result = this._admitSessionInput(action, { - immediatelyEligible: !wasBusy && this._canStartSessionActionImmediately(), - }); - commitFence?.release(); - reportPreflight(result.accepted, result.disposition === "queued"); - if (!result.accepted || !result.ticket) return; - if (options?.returnAfterAccepted) { - if (result.disposition === "starts_when_admitted") await result.ticket.delivered; - return; - } - if (result.disposition === "queued") return; - await this.waitForSessionInputIdle(); - return; - } - - const queueForStreaming = this.isStreaming; - const queueForBusy = options?.queueIfBusy === true && this._isBusyForSessionInput("preflight"); - const visibleQueued = queueForStreaming || queueForBusy; - if (visibleQueued && !options?.streamingBehavior) { - const stateDescription = queueForStreaming ? "Agent is already processing" : "Agent has queued work"; - throw new Error( - `${stateDescription}. Specify streamingBehavior ('steer' or 'followUp') to queue the message.`, - ); - } - const schedule = options?.streamingBehavior ?? "followUp"; - const prefixMessages = visibleQueued ? this._takePendingNextTurnMessages() : undefined; - const content = options?.content - ? options.content.map((block) => ({ ...block })) - : buildPromptContent(normalized.text, normalized.images); - const suppliedMessage = options?.customMessage; - const primaryMessage = suppliedMessage - ? visibleQueued - ? suppliedMessage - : cloneCustomMessage(suppliedMessage) - : ({ - role: "user", - content: content.map((block) => ({ ...block })), - timestamp: Date.now(), - } satisfies UserMessage); - const acceptedAgentMessage = options?.skipPrePromptWork === true && options.returnAfterAccepted === true; - const action = this._createPreparedTurnAction(schedule, normalized.text, normalized.images, { - agentMessageId: options?.agentMessageId, - queueKey: options?.followUpQueueKey, - content, - message: primaryMessage, - prefixMessages, - suppressAutonomousContinuation: options?.suppressAutonomousContinuation, - resumeIfIdle: - !visibleQueued || - options?.resumeIfIdle || - (options?.queueIfBusy === true && canSelectSessionAction(this._runtimeActivity())), - source: isInternalPrompt ? "internal" : (options?.source ?? "interactive"), - executionPolicy: visibleQueued - ? createTurnExecutionPolicy("queued") - : createTurnExecutionPolicy("directPrompt", { - returnAfterAccepted: options?.returnAfterAccepted, - skipPrePromptWork: options?.skipPrePromptWork, - }), - queueVisible: visibleQueued, - acceptedAgentMessage, - acceptedBeforeCompletion: options?.returnAfterAccepted === true, - }); - if (action.suppressAutonomousContinuation) { - this._markAutonomousContinuationSuppressed(primaryDeliveryRecord(action).message); - } - const result = this._admitSessionInput(action, { - immediatelyEligible: !visibleQueued && this._canStartSessionActionImmediately(), - }); - commitFence?.release(); - if (!result.accepted || !result.ticket) { - if (prefixMessages) this._pendingNextTurnMessages.unshift(...prefixMessages); - reportPreflight(false, false); - return; - } - if (result.disposition === "queued") { - reportPreflight(true, true); - } else { - void result.ticket.delivered.then( - () => reportPreflight(true), - () => reportPreflight(false), - ); - } - const deferralObserver = - acceptedAgentMessage && - options?.queueIfBusy === true && - !options.streamingBehavior && - result.disposition === "starts_when_admitted" - ? this._observeSessionActionDeferral(action) - : undefined; - if (acceptedAgentMessage && !queueForStreaming && !queueForBusy && !options?.streamingBehavior) { - try { - const outcome = deferralObserver - ? await Promise.race([ - result.ticket.delivered.then(() => "delivered" as const), - deferralObserver.deferred.then(() => "deferred" as const), - ]) - : await result.ticket.delivered.then(() => "delivered" as const); - if (outcome === "deferred" && !options?.streamingBehavior) { - const error = new Error( - "Agent became busy before prompt delivery. Specify streamingBehavior ('steer' or 'followUp') to queue the message.", - ); - this._rejectAgentMessage(action.agentMessageId, error); - this._cancelSessionActions((candidate) => candidate === action, error); - this._emitQueueUpdate(); - throw error; - } - return; - } finally { - deferralObserver?.stop(); - } - } - if (options?.returnAfterAccepted) { - if (result.disposition === "starts_when_admitted" || (acceptedAgentMessage && !visibleQueued)) { - await result.ticket.delivered; - } - return; - } - if (visibleQueued) return; - await result.ticket.completed; - await this.waitForSessionInputIdle(); - } catch (error) { - reportPreflight(false); - throw error; - } finally { - commitFence?.release(); - } - }; - return commitFence ? this._commitFence.run(commitFence, run) : run(); - } - - private _executeExtensionCommand(text: string): Promise | undefined { - const parsed = parseSlashCommand(text); - if (!parsed) return undefined; - const commandName = parsed.name; - const args = parsed.args; - - const command = this._extensionRunner.getCommand(commandName); - if (!command) return undefined; - const context = this._extensionRunner.createCommandContext(); - return Promise.resolve() - .then(() => command.handler(args, context)) - - .catch((error: unknown) => { - const commandError = error instanceof Error ? error : new Error(String(error)); - this._extensionRunner.emitError({ - extensionPath: `command:${commandName}`, - event: "command", - error: commandError.message, - }); - throw commandError; - }); + private _promptInjectedMessage( + ...args: Parameters + ): ReturnType { + return this._promptSubmission.promptInjectedMessage(...args); } - /** - * Expand skill commands (/skill:name args) to their full content. - * Returns the expanded text, or the original text if not a skill command or skill not found. - * Emits errors via extension runner if file read fails. - */ - private _expandSkillCommand(text: string): string { - if (!text.startsWith("/skill:")) return text; - - const parsed = parseSlashCommand(text); - if (!parsed?.name.startsWith("skill:")) return text; - const skillName = parsed.name.slice("skill:".length); - const args = parsed.args; - - const skill = this.resourceLoader.getSkills().skills.find((s) => s.name === skillName); - if (!skill) return text; // Unknown skill, pass through - - try { - const content = readFileSync(skill.filePath, "utf-8"); - const body = stripFrontmatter(content).trim(); - const skillBlock = `\nReferences are relative to ${skill.baseDir}.\n\n${body}\n`; - return args ? `${skillBlock}\n\n${args}` : skillBlock; - } catch (err) { - this._extensionRunner.emitError({ - extensionPath: skill.filePath, - event: "skill_expansion", - error: err instanceof Error ? err.message : String(err), - }); - return text; // Return original on error - } + private _prompt( + ...args: Parameters + ): ReturnType { + return this._promptSubmission.prompt(...args); } /** @@ -3733,30 +1784,8 @@ export class AgentSession { * @param images Optional image attachments to include with the message * @throws Error if text is an extension command */ - async steer( - text: string, - images?: ImageContent[], - options: { - queueKey?: string; - agentMessageId?: string; - resumeIfIdle?: boolean; - } = {}, - ): Promise { - const normalized = this._normalizeSubmission(text, images, { - parseSessionCommands: false, - extensionCommands: "reject", - expandSkills: true, - expandPromptTemplates: true, - }); - if (normalized instanceof Promise || normalized.kind !== "prompt") { - throw new Error("Queued prompt normalization did not produce a prompt"); - } - - await this._queuePreparedPrompt("steer", normalized.text, normalized.images, { - queueKey: options.queueKey, - agentMessageId: options.agentMessageId, - resumeIfIdle: options.resumeIfIdle, - }); + steer(...args: Parameters): ReturnType { + return this._promptSubmission.steer(...args); } /** @@ -3766,344 +1795,50 @@ export class AgentSession { * @param images Optional image attachments to include with the message * @throws Error if text is an extension command */ - async followUp( - text: string, - images?: ImageContent[], - options: { - queueKey?: string; - agentMessageId?: string; - resumeIfIdle?: boolean; - } = {}, - ): Promise { - const normalized = this._normalizeSubmission(text, images, { - parseSessionCommands: false, - extensionCommands: "reject", - expandSkills: true, - expandPromptTemplates: true, - }); - if (normalized instanceof Promise || normalized.kind !== "prompt") { - throw new Error("Queued prompt normalization did not produce a prompt"); - } - - return this._queuePreparedPrompt("followUp", normalized.text, normalized.images, { - queueKey: options.queueKey, - agentMessageId: options.agentMessageId, - resumeIfIdle: options.resumeIfIdle, - }); + followUp(...args: Parameters): ReturnType { + return this._promptSubmission.followUp(...args); } - async restoreSessionActions(snapshot: SessionActionRecoverySnapshot): Promise { - if (snapshot.formatVersion !== SESSION_ACTION_RECOVERY_FORMAT_VERSION) { - throw new Error(`Unsupported session action recovery format version: ${snapshot.formatVersion}`); - } - const actionIds = new Set(this._actionStore.ownedActions().map((action) => action.id)); - const actions = snapshot.actions.map((recovered): QueuedSessionAction => { - if (actionIds.has(recovered.id)) throw new Error(`Duplicate session action id: ${recovered.id}`); - actionIds.add(recovered.id); - if ( - recovered.payload.kind === "turn" && - recovered.payload.records.some((record) => record.ownerActionId !== recovered.id) - ) { - throw new Error(`Session action ${recovered.id} has invalid delivery correlation`); - } - const payload: PreparedTurnPayload | PreparedCommandPayload = - recovered.payload.kind === "turn" - ? { - kind: "turn", - text: recovered.payload.text, - ...(recovered.payload.preview ? { preview: recovered.payload.preview } : {}), - records: recovered.payload.records.map((record) => ({ - id: record.id, - role: record.role, - message: cloneQueuedAgentMessage(record.message), - started: false, - durable: false, - ownerActionId: record.ownerActionId, - })), - ...(recovered.payload.images - ? { - images: recovered.payload.images.map((image) => ({ - ...image, - })), - } - : {}), - ...(recovered.payload.content - ? { - content: recovered.payload.content.map((block) => ({ - ...block, - })), - } - : {}), - ...(recovered.payload.customMessage - ? { - customMessage: cloneCustomMessage(recovered.payload.customMessage), - } - : {}), - executionPolicy: { - ...recovered.payload.executionPolicy, - preparation: { - ...recovered.payload.executionPolicy.preparation, - }, - }, - queueVisible: recovered.payload.queueVisible, - acceptedAgentMessage: recovered.payload.acceptedAgentMessage, - acceptedBeforeCompletion: recovered.payload.acceptedBeforeCompletion, - } - : { - kind: "session_command", - text: recovered.payload.text, - command: { ...recovered.payload.command }, - ...(recovered.payload.images - ? { - images: recovered.payload.images.map((image) => ({ - ...image, - })), - } - : {}), - }; - return { - id: recovered.id, - source: recovered.source, - delivery: recovered.delivery, - wake: recovered.wake, - payload, - lifecycle: { state: "queued" }, - ...(recovered.queueKey ? { queueKey: recovered.queueKey } : {}), - ...(recovered.agentMessageId ? { agentMessageId: recovered.agentMessageId } : {}), - ...(recovered.suppressAutonomousContinuation ? { suppressAutonomousContinuation: true } : {}), - }; - }); - for (const action of actions) { - const durableTerminalNotice = this._isRlmTerminalNoticeAction(action); - if (durableTerminalNotice) this._durableRlmTerminalNoticeActionIds.add(action.id); - try { - this._admitSessionInput(action, { restore: true }); - } catch (error) { - if (durableTerminalNotice) this._durableRlmTerminalNoticeActionIds.delete(action.id); - throw error; - } - } - return actions.length; - } - - private _restoreSessionCommand( - text: string, - customMessage: CustomMessage | undefined, - images: ImageContent[] | undefined, - schedule: SessionInputSchedule, - agentMessageId: string | undefined, - ): boolean | undefined { - if (!isSessionSlashCommandMessage(customMessage) || text !== customMessage.details.command.text) { - return undefined; - } - return this._admitSessionInput( - createSessionCommandAction(text, customMessage.details.command, images, schedule, { - agentMessageId, - source: "internal", - }), - { restore: true }, - ).accepted; - } - - private _restorePromptInput(schedule: SessionInputSchedule, snapshot: RestoredPromptInput): Promise { - return this._queuePreparedPrompt(schedule, snapshot.text, snapshot.images, { - queueKey: snapshot.queueKey, - agentMessageId: snapshot.agentMessageId, - content: snapshot.content, - message: snapshot.customMessage, - prefixMessages: snapshot.prefixMessages, - source: "internal", - }); + restoreSessionActions( + ...args: Parameters + ): ReturnType { + return this._actionRecovery.restoreSessionActions(...args); } - async restoreSteeringMessage( - text: string, - images?: ImageContent[], - options: { - queueKey?: string; - agentMessageId?: string; - content?: (TextContent | ImageContent)[]; - customMessage?: CustomMessage; - prefixMessages?: CustomMessage[]; - } = {}, - ): Promise { - if ( - this._restoreSessionCommand(text, options.customMessage, images, "steer", options.agentMessageId) !== undefined - ) - return; - - await this._restorePromptInput("steer", { - text, - images, - queueKey: options.queueKey, - agentMessageId: options.agentMessageId, - content: options.content, - customMessage: options.customMessage, - prefixMessages: options.prefixMessages, - }); + restoreSteeringMessage( + ...args: Parameters + ): ReturnType { + return this._actionQueue.restoreSteeringMessage(...args); } - async restoreFollowUpMessage( - text: string, - images?: ImageContent[], - options: { - queueKey?: string; - agentMessageId?: string; - content?: (TextContent | ImageContent)[]; - customMessage?: CustomMessage; - prefixMessages?: CustomMessage[]; - } = {}, - ): Promise { - const restoredCommand = this._restoreSessionCommand( - text, - options.customMessage, - images, - "followUp", - options.agentMessageId, - ); - if (restoredCommand !== undefined) return restoredCommand; - - return this._restorePromptInput("followUp", { - text, - images, - queueKey: options.queueKey, - agentMessageId: options.agentMessageId, - content: options.content, - customMessage: options.customMessage, - prefixMessages: options.prefixMessages, - }); + restoreFollowUpMessage( + ...args: Parameters + ): ReturnType { + return this._actionQueue.restoreFollowUpMessage(...args); } - private _takePendingNextTurnMessages(): CustomMessage[] { - const messages = this._pendingNextTurnMessages; - this._pendingNextTurnMessages = []; - return messages; + private _takePendingNextTurnMessages( + ...args: Parameters + ): ReturnType { + return this._pendingContext.takePendingNextTurnMessages(...args); } - private _createPreparedTurnAction(...args: Parameters): QueuedSessionAction { - return createPreparedTurnAction(...args); + private _assertSessionActionAdmissionAvailable( + ...args: Parameters + ): ReturnType { + return this._inputAdmission.assertSessionActionAdmissionAvailable(...args); } - private _coalescedFollowUpOwner(action: QueuedSessionAction): QueuedSessionAction | undefined { - if (action.delivery !== "when_run_idle" || action.payload.kind !== "turn" || !action.queueKey) return undefined; - return this._actionStore - .unfinishedActions() - .find( - (candidate) => - candidate.queueKey === action.queueKey && - (candidate.lifecycle.state === "queued" || - candidate.lifecycle.state === "selected" || - candidate.lifecycle.state === "preparing"), - ); + private _admitSessionInput( + ...args: Parameters + ): ReturnType { + return this._inputAdmission.admitSessionInput(...args); } - private _assertSessionActionAdmissionAvailable(): void { - if (this._disposed || this._disposing) { - throw new Error("Cannot admit a session action because the session is disposing or disposed."); - } - if (this._inputScheduler.admissionPaused) { - throw new SessionInputAdmissionPausedError( - "Cannot admit a session action while session input admission is paused.", - ); - } - if (this._inputScheduler.suspended) { - throw new Error("Cannot admit a session action while queued session input is suspended."); - } - } - - private _admitSessionInput( - action: QueuedSessionAction, - options: { - restore?: boolean; - front?: boolean; - wake?: boolean; - immediatelyEligible?: boolean; - } = {}, - ): { - accepted: boolean; - disposition: "starts_when_admitted" | "queued"; - ticket?: ActionTicket; - } { - if (this._disposed || this._disposing) { - throw new Error("Cannot admit a session action because the session is disposing or disposed."); - } - if (this._inputScheduler.admissionPaused) { - throw new SessionInputAdmissionPausedError( - "Cannot admit a session action while session input admission is paused.", - ); - } - if ( - options.restore !== true && - action.payload.kind === "turn" && - isAgentSessionMessage(primaryDeliveryRecord(action).message) - ) { - assertAgentMessageQueueCapacity( - this._actionStore.unfinishedActions().length, - DEFAULT_AGENT_MESSAGE_MAX_PENDING_PER_SESSION, - ); - } - const coalescedOwner = options.restore ? undefined : this._coalescedFollowUpOwner(action); - if (coalescedOwner) { - if (action.agentMessageId !== coalescedOwner.agentMessageId) { - this._rejectAgentMessage( - action.agentMessageId, - new Error("Prompt was not queued because an equivalent follow-up is already pending."), - ); - } - return { accepted: false, disposition: "queued" }; - } - const canStartImmediately = - options.immediatelyEligible === true && - (this._actionStore.unfinishedActions().length === 0 || options.front === true); - if (options.front) this._actionStore.enqueueFront(action); - else this._actionStore.enqueue(action); - let disposition: "starts_when_admitted" | "queued" = "queued"; - if (canStartImmediately && this._actionStore.selectFirst() === action) disposition = "starts_when_admitted"; - const controller = this._actionStore.ticketFor(action); - controller.settleAccepted({ - status: "accepted", - actionId: action.id, - disposition, - }); - this._sessionInputArrivalEpoch++; - this._emitQueueUpdate(); - if ( - !options.restore && - options.wake !== false && - (disposition === "starts_when_admitted" || - (action.delivery === "next_turn_boundary" && this.isStreaming) || - action.payload.kind === "session_command" || - action.wake === "immediate") - ) { - if (action.payload.kind === "turn" && action.wake === "immediate") { - this._resumeSessionInputAdmission(); - } - this._scheduleSessionInputPump(); - } - return { accepted: true, disposition, ticket: controller.ticket }; - } - - private async _queuePreparedPrompt( - schedule: SessionInputSchedule, - text: string, - images?: ImageContent[], - options: { - agentMessageId?: string; - queueKey?: string; - content?: (TextContent | ImageContent)[]; - message?: QueuedAgentMessage; - prefixMessages?: CustomMessage[]; - previewLabel?: string; - suppressAutonomousContinuation?: boolean; - resumeIfIdle?: boolean; - source?: InputSource | "internal"; - } = {}, - ): Promise { - const action = this._createPreparedTurnAction(schedule, text, images, options); - if (action.suppressAutonomousContinuation) { - this._markAutonomousContinuationSuppressed(primaryDeliveryRecord(action).message); - } - return this._admitSessionInput(action).accepted; + private _queuePreparedPrompt( + ...args: Parameters + ): ReturnType { + return this._inputAdmission.queuePreparedPrompt(...args); } private _runtimeActivity(): RuntimeActivity { @@ -4124,74 +1859,21 @@ export class AgentSession { } get hasPendingSessionWork(): boolean { - return this._actionStore.unfinishedActions().some((action) => { - const state = action.lifecycle.state; - return ( - state === "queued" || - state === "selected" || - state === "preparing" || - (state === "committing" && action.payload.kind === "turn" && !primaryDeliveryRecord(action).durable) - ); - }); + return this._actionQueue.hasPendingSessionWork; } get hasPendingAdmissionWaiters(): boolean { - return this._commitFence.hasPendingWork || this._sessionInputCheckpointWaiters.size > 0; + return this._commitFence.hasPendingWork || this._inputCheckpoints.hasWaiters; } private _scheduleSessionInputPump(): void { this._inputScheduler.schedule(); } - private async _executeSelectedSessionCommand(action: QueuedSessionAction, epoch: number): Promise { - if (action.payload.kind !== "session_command") throw new Error("Expected a selected session command"); - const input = action.payload; - const commitFence = await this._acquireSessionActionCommitFence(); - try { - await this._commitFence.run(commitFence, async () => { - const isCancelled = () => action.lifecycle.state === "cancelled"; - if (isCancelled()) return; - await this._refinement._waitForRefineIdle(); - if (isCancelled()) return; - if (this._isSessionInputHandoffDeferred(epoch) || !canSelectSessionAction(this._runtimeActivity())) { - this._actionStore.rollback(action); - this._notifySessionInputCheckpointChange(); - this._emitQueueUpdate(); - return; - } - transitionSessionAction(action, { - state: "running", - execution: "session_command", - }); - this._notifySessionInputCheckpointChange(); - this._emitQueueUpdate(); - try { - this._appendDurableSessionCommandMessage(input.text, input.command, false); - this._actionStore.ticketFor(action).settleDelivered({ status: "not_applicable" }); - this._settleAgentMessage(action.agentMessageId, "delivery"); - await this._executeQueuedSessionCommand(action); - transitionSessionAction(action, { state: "completed" }); - this._actionStore.ticketFor(action).settleCompleted(); - this._settleAgentMessage(action.agentMessageId, "completion"); - } catch (error) { - const commandError = this._asError(error); - transitionSessionAction(action, { - state: "failed", - error: commandError, - }); - const ticket = this._actionStore.ticketFor(action); - ticket.rejectDelivered(commandError); - ticket.settleCompleted(commandError); - this._rejectAgentMessage(action.agentMessageId, commandError); - } finally { - this._actionStore.releaseTerminal(action); - this._notifySessionInputCheckpointChange(); - this._emitQueueUpdate(); - } - }); - } finally { - commitFence.release(); - } + private _executeSelectedSessionCommand( + ...args: Parameters + ): ReturnType { + return this._commandExecution.executeSelectedSessionCommand(...args); } private _isBusyForSessionInput(point: "preflight" | "pump"): boolean { @@ -4231,260 +1913,10 @@ export class AgentSession { } } - private async _startPreparedTurnActions(actions: QueuedSessionAction[], epoch: number): Promise { - let nextTurnMessages: CustomMessage[] = []; - const activeTurns = () => - actions.filter( - (action): action is SessionAction => - action.payload.kind === "turn" && action.lifecycle.state === "preparing", - ); - const firstTurn = activeTurns()[0]; - if (!firstTurn) return; - const executionPolicy = firstTurn.payload.executionPolicy; - // The digest is never parked as pending context; lazy injection re-arms instead. - const parkNextTurnMessages = (messages: CustomMessage[]) => { - const parked = messages.filter((message) => message.customType !== HARNESS_DIGEST_CUSTOM_TYPE); - if (parked.length !== messages.length) this._harnessDigestPending = true; - this._pendingNextTurnMessages.unshift(...parked); - }; - const restoreNextTurnContext = () => { - parkNextTurnMessages(nextTurnMessages); - nextTurnMessages = []; - }; - try { - const preparedTurn = await this._turnPreparer.prepare(executionPolicy.preparation, { - afterValidation: () => { - if (this._isSessionInputHandoffDeferred(epoch)) { - throw new DeferredSessionInputError("Session input paused before preflight"); - } - }, - prepare: async () => { - if (executionPolicy.nextTurnContextTiming === "preparation") { - nextTurnMessages = this._takePendingNextTurnMessages(); - } - if (!executionPolicy.runBeforeAgentStart) return undefined; - while (activeTurns().some((action) => action.payload.prepared === undefined)) { - if (this._isSessionInputHandoffDeferred(epoch)) { - throw new DeferredSessionInputError("Session input paused before preparation"); - } - const preparationAction = activeTurns().at(-1); - if (!preparationAction) return undefined; - const basePromptSnapshot = this._baseSystemPrompt; - const result = await this._extensionRunner.emitBeforeAgentStart( - preparationAction.payload.text, - preparationAction.payload.images, - basePromptSnapshot, - this._baseSystemPromptOptions, - ); - if (activeTurns().at(-1) !== preparationAction) continue; - const prepared = { result, basePromptSnapshot }; - for (const action of activeTurns()) action.payload.prepared = prepared; - } - if (this._isSessionInputHandoffDeferred(epoch)) { - throw new DeferredSessionInputError("Session input paused before handoff"); - } - return activeTurns()[0]?.payload.prepared; - }, - shouldCommit: () => activeTurns().length > 0, - commit: (prepared) => { - if (this._isSessionInputHandoffDeferred(epoch)) { - throw new DeferredSessionInputError("Session input paused before handoff"); - } - const turns = activeTurns(); - if (turns.length === 0) return undefined; - return { prepared, turns }; - }, - }); - if (!preparedTurn) { - restoreNextTurnContext(); - return; - } - const { prepared, turns } = preparedTurn; - const commitFence = await this._acquireSessionActionCommitFence(); - let promptPromise: Promise; - try { - promptPromise = this._commitFence.run(commitFence, () => { - if ( - this._isSessionInputHandoffDeferred(epoch) || - this.isStreaming || - turns.some((action) => action.lifecycle.state !== "preparing") - ) { - throw new DeferredSessionInputError("Agent became active before session input handoff"); - } - if (executionPolicy.nextTurnContextTiming === "commit") { - nextTurnMessages = this._takePendingNextTurnMessages(); - } - if (this._harnessDigestPending) { - // The first-turn digest rides the turn's delivery records so a - // cancelled first turn strips it with the rest of the turn. - this._harnessDigestPending = false; - const digest = this._harnessDigest(); - if (this._latestContextHarnessDigest() !== digest) { - nextTurnMessages = [createHarnessDigestMessage(digest), ...nextTurnMessages]; - } - } - const contextRecords = nextTurnMessages.map((message) => - createDeliveryRecord(turns[0].id, "next_turn", message), - ); - const firstPrimaryIndex = turns[0].payload.records.indexOf(primaryDeliveryRecord(turns[0])); - turns[0].payload.records.splice(firstPrimaryIndex, 0, ...contextRecords); - const preparedMessages: AgentMessage[] = turns.flatMap((action) => - action.payload.records.map((record) => record.message), - ); - for (const action of turns) { - if (action.suppressAutonomousContinuation) { - this._markAutonomousContinuationSuppressed(primaryDeliveryRecord(action).message); - } - } - if (executionPolicy.runBeforeAgentStart) { - this._appendBeforeAgentStartMessages(preparedMessages, prepared?.result); - this._applyPreparedSystemPrompt(prepared, executionPolicy.preserveEmptyExtensionPrompt); - } else if (executionPolicy.nextTurnContextTiming !== "skip") { - this.agent.state.systemPrompt = this._baseSystemPrompt; - } - for (const action of turns) transitionSessionAction(action, { state: "committing" }); - this._notifySessionInputCheckpointChange(); - this._emitQueueUpdate(); - return turns.some((action) => action.suppressAutonomousContinuation) - ? this._runWithAutonomousContinuationSuppressed(() => this.agent.prompt(preparedMessages)) - : this.agent.prompt(preparedMessages); - }); - } finally { - commitFence.release(); - } - await promptPromise; - if (executionPolicy.completionIncludesRetryChain) await this.waitForRetry(); - if (!this._hasCancelledDispatchCapture()) await this._agentEventQueue; - if ( - turns.some( - (action) => - action.lifecycle.state !== "cancelled" && - !primaryDeliveryRecord(action).durable && - !this.agent.state.messages.includes(primaryDeliveryRecord(action).message), - ) - ) { - throw new Error("Session input dispatch settled without durable delivery"); - } - this._forgetConsumedPostCompactionContinuations(turns.map((action) => primaryDeliveryRecord(action).message)); - } catch (error) { - const delivered = new Set(this.agent.state.messages); - parkNextTurnMessages(nextTurnMessages.filter((message) => !delivered.has(message))); - for (const action of actions) { - if (action.payload.kind === "turn") { - action.payload.records = action.payload.records.filter((record) => record.role !== "next_turn"); - } - } - throw error; - } - } - - private async _executeQueuedSessionCommand(action: QueuedSessionAction): Promise { - if (action.payload.kind !== "session_command") throw new Error("Expected a session command action"); - const input = action.payload; - try { - let resultText: string | undefined; - let displayResult = true; - switch (input.command.name) { - case "compact": - await this.compact(input.command.args || undefined, { - skipAbort: true, - }); - break; - case "refine": { - let result: RefinementResult; - try { - const options = parseRefineCommandOptions(input.command.args); - result = await this.refine(options, { skipAbort: true }); - } catch (error) { - // Only a failure of the refinement itself is a refine failure; a later - // result-row persist error must not report a completed refinement as failed. - this._refinement._emitRefineFailed(this._asError(error)); - throw error; - } - const applied = result.appliedEdits.filter((edit) => edit.applied).length; - resultText = `Refined continual harness state: ${applied} edit${applied === 1 ? "" : "s"} applied.`; - displayResult = false; - break; - } - case "goal": - await this._handleGoalSlashCommand(input.text, input.images); - resultText = this._goals.state.objective - ? `Goal ${this._goals.state.status}: ${this._goals.state.objective}` - : "No active goal."; - break; - case "autonomous": - await this._handleAutonomousSlashCommand(input.text); - break; - } - if (resultText) { - this._appendDurableSessionCommandMessage(resultText, input.command, true, false, displayResult); - } - } catch (error) { - if (error instanceof CompactionSkippedError) return; - const commandError = error instanceof Error ? error : new Error(String(error)); - try { - this._appendDurableSessionCommandMessage( - `Command failed: ${commandError.message}`, - input.command, - true, - true, - ); - } catch { - // The result row is also the command-correlated UI settle edge. - const message = createSessionSlashCommandResultMessage(`Command failed: ${commandError.message}`, { - command: input.command, - success: false, - severity: "error", - error: commandError.message, - }); - this._emit({ type: "message_start", message }); - this._emit({ type: "message_end", message }); - } - throw commandError; - } - } - - private _appendDurableSessionCommandMessage( - content: string, - command: SessionSlashCommand, - isResult: boolean, - isError = false, - display = true, - ): void { - const message: CustomMessage = isResult - ? createSessionSlashCommandResultMessage( - content, - { - command, - success: !isError, - severity: isError ? "error" : "info", - ...(isError ? { error: content.replace(/^Command failed:\s*/, "") } : {}), - }, - display, - ) - : createSessionSlashCommandMessage(command); - // Persist before touching live state so a failed write cannot leave an - // unsaved leaf that the next entry would silently parent onto. - this.sessionManager.appendCustomMessageEntryWithRollback( - message.customType, - message.content, - message.display, - message.details, - ); - this.agent.state.messages.push(message); - this._emit({ type: "message_start", message }); - this._emit({ type: "message_end", message }); - } - - private _throwIfExtensionCommand(text: string): void { - const commandName = parseSlashCommand(text)?.name ?? ""; - const command = this._extensionRunner.getCommand(commandName); - - if (command) { - throw new Error( - `Extension command "/${commandName}" cannot be queued. Use prompt() or execute the command when not streaming.`, - ); - } + private _startPreparedTurnActions( + ...args: Parameters + ): ReturnType { + return this._turnExecution.startPreparedTurnActions(...args); } /** @@ -4499,66 +1931,14 @@ export class AgentSession { * @param options.triggerTurn If true and not streaming, triggers a new LLM turn * @param options.deliverAs Delivery mode: "steer", "followUp", or "nextTurn" */ - async sendCustomMessage( + sendCustomMessage( message: Pick, "customType" | "content" | "display" | "details">, options?: { triggerTurn?: boolean; deliverAs?: "steer" | "followUp" | "nextTurn"; }, ): Promise { - const appMessage = { - role: "custom" as const, - customType: message.customType, - content: message.content, - display: message.display, - details: message.details, - timestamp: Date.now(), - } satisfies CustomMessage; - if (options?.deliverAs === "nextTurn") { - this._pendingNextTurnMessages.push(appMessage); - } else if (this.isStreaming) { - const normalized = normalizeMessageContent(message.content); - if (options?.deliverAs === "followUp") { - await this._queuePreparedPrompt("followUp", normalized.text, normalized.images, { - message: appMessage, - resumeIfIdle: true, - }); - } else { - await this._queuePreparedPrompt("steer", normalized.text, normalized.images, { - message: appMessage, - resumeIfIdle: true, - }); - } - } else if (options?.triggerTurn) { - if (!this._inputScheduler.suspendedForUpdateRestart) this._resumeSessionInputAdmission(); - const admissionFence = await this._acquireDirectTurnAdmissionFence(); - try { - const normalized = normalizeMessageContent(message.content); - const immediatelyEligible = this._canStartSessionActionImmediately(); - const action = this._createPreparedTurnAction("followUp", normalized.text, normalized.images, { - message: appMessage, - resumeIfIdle: true, - executionPolicy: createTurnExecutionPolicy("customTrigger"), - queueVisible: false, - }); - const result = this._admitSessionInput(action, { immediatelyEligible }); - admissionFence.release(); - if (!result.ticket) return; - await result.ticket.completed; - } finally { - admissionFence.release(); - } - } else { - this.agent.state.messages.push(appMessage); - this.sessionManager.appendCustomMessageEntry( - message.customType, - message.content, - message.display, - message.details, - ); - this._emit({ type: "message_start", message: appMessage }); - this._emit({ type: "message_end", message: appMessage }); - } + return this._promptSubmission.sendCustomMessage(message, options); } /** @@ -4568,146 +1948,32 @@ export class AgentSession { * @param content User message content (string or content array) * @param options.deliverAs Delivery mode when streaming: "steer" or "followUp" */ - async sendUserMessage( - content: string | (TextContent | ImageContent)[], - options?: { deliverAs?: "steer" | "followUp" }, - ): Promise { - let text: string; - let images: ImageContent[] | undefined; - - if (typeof content === "string") { - text = content; - } else { - const textParts: string[] = []; - images = []; - for (const part of content) { - if (part.type === "text") { - textParts.push(part.text); - } else { - images.push(part); - } - } - text = textParts.join("\n"); - if (images.length === 0) images = undefined; - } - - await this._prompt(text, { - expandPromptTemplates: false, - streamingBehavior: options?.deliverAs, - images, - source: "extension", - resumeIfIdle: true, - }); + sendUserMessage( + ...args: Parameters + ): ReturnType { + return this._promptSubmission.sendUserMessage(...args); } - clearQueue(): { steering: string[]; followUp: string[] } { - const clearable = this._actionStore - .clearableActions() - .filter((action) => action.payload.kind === "session_command" || action.payload.queueVisible); - if (clearable.some((action) => action.payload.kind === "turn" && action.lifecycle.state === "preparing")) { - this._inputScheduler.invalidatePreparation(); - } - const steering = clearable - .filter((action) => action.delivery === "next_turn_boundary") - .map((action) => action.payload.text); - const followUp = clearable - .filter((action) => action.delivery === "when_run_idle") - .map((action) => action.payload.text); - const promptError = new Error("Queued prompt was cleared before delivery."); - const agentMessageError = new Error("Queued agent message was cleared before delivery."); - for (const action of clearable) { - const error = - action.payload.kind === "turn" && action.lifecycle.state === "preparing" ? promptError : agentMessageError; - this._settleAgentMessage(action.agentMessageId, "delivery", error); - this._settleAgentMessage(action.agentMessageId, "completion", error); - } - const clearableIds = new Set(clearable.map((action) => action.id)); - this._cancelSessionActions((action) => clearableIds.has(action.id), agentMessageError); - this.agent.clearAllQueues(); - this._emitQueueUpdate(); - return { steering, followUp }; + clearQueue(...args: Parameters): ReturnType { + return this._actionQueue.clearQueue(...args); } - private _invalidateQueuedPromptPreparation(): void { - for (const action of this._actionStore.clearableActions()) { - if (action.payload.kind === "turn") action.payload.prepared = undefined; - } + private _invalidateQueuedPromptPreparation( + ...args: Parameters + ): ReturnType { + return this._actionQueue.invalidateQueuedPromptPreparation(...args); } - clearQueuedAgentMessages(): { steering: string[]; followUp: string[] } { - this._agentMessageClearEpoch++; - // customType identifies agent messages; the text parser covers persisted pre-grammar prompts. - return this._clearQueuedTurnActionsMatching( - (action) => - isAgentSessionMessage(primaryDeliveryRecord(action).message) || - isAgentSessionMessagePrompt(action.payload.text), - ); + clearQueuedAgentMessages( + ...args: Parameters + ): ReturnType { + return this._actionQueue.clearQueuedAgentMessages(...args); } - clearQueuedUserMessagesMatching(predicate: (text: string) => boolean): { steering: string[]; followUp: string[] } { - return this._clearQueuedTurnActionsMatching((action) => predicate(action.payload.text)); - } - - private _clearQueuedTurnActionsMatching(matches: (action: QueuedSessionAction) => boolean): { - steering: string[]; - followUp: string[]; - } { - const ownedActions = this._actionStore.ownedActions(); - const dispatchedTurnCount = ownedActions.filter( - (action) => - action.payload.kind === "turn" && - (action.lifecycle.state === "committing" || action.lifecycle.state === "running"), - ).length; - const matching = ownedActions.filter( - (action) => - action.payload.kind === "turn" && - action.agentMessageId !== undefined && - matches(action) && - (action.lifecycle.state === "queued" || - action.lifecycle.state === "selected" || - action.lifecycle.state === "preparing" || - (action.lifecycle.state === "committing" && - dispatchedTurnCount === 1 && - !primaryDeliveryRecord(action).started)), - ); - if (matching.length === 0) return { steering: [], followUp: [] }; - const removedTexts = (delivery: DeliveryPolicy) => - [ - ...matching.filter((action) => action.delivery === delivery && action.lifecycle.state === "queued"), - ...matching.filter((action) => action.delivery === delivery && action.lifecycle.state !== "queued"), - ].map((action) => action.payload.text); - const removedSteering = removedTexts("next_turn_boundary"); - const removedFollowUp = removedTexts("when_run_idle"); - const acceptedError = new Error("Accepted agent message was cleared before delivery."); - const queuedError = new Error("Queued agent message was cleared before delivery."); - for (const action of matching) { - const error = - action.payload.kind === "turn" && action.payload.acceptedAgentMessage ? acceptedError : queuedError; - this._rejectAgentMessage(action.agentMessageId, error); - } - for (const [accepted, error] of [ - [true, acceptedError], - [false, queuedError], - ] as const) { - const ids = new Set( - matching - .filter((action) => action.payload.kind === "turn" && action.payload.acceptedAgentMessage === accepted) - .map((action) => action.id), - ); - if (ids.size > 0) this._cancelSessionActions((action) => ids.has(action.id), error, matching); - } - if ( - matching.some( - (action) => - action.lifecycle.state === "cancelled" && - action.payload.kind === "turn" && - action.payload.captureRunMessages, - ) - ) { - this.agent.abort(); - } - this._emitQueueUpdate(); - return { steering: removedSteering, followUp: removedFollowUp }; + clearQueuedUserMessagesMatching( + ...args: Parameters + ): ReturnType { + return this._actionQueue.clearQueuedUserMessagesMatching(...args); } /** @@ -4716,70 +1982,9 @@ export class AgentSession { * item's current preview so clients never edit a shifted queue by accident. */ mutateQueuedMessage( - lane: QueuedMessageLane, - index: number, - expectedText: string, - mutation: QueuedMessageMutation, - ): QueuedMessageMutationStatus { - const policy = queuedMessageLaneDeliveryPolicy(lane); - const projection = visibleSessionActionProjection(this._actionStore.queuedActions(policy)); - const item = projection[index]; - if (!item || queuedAgentMessagePreview(item) !== expectedText) return "rejected"; - if (mutation.type === "delete") { - const error = new Error("Queued prompt was deleted before delivery."); - this._rejectAgentMessage(item.agentMessageId, error); - this._cancelSessionActions((candidate) => candidate === item, error); - this._emitQueueUpdate(); - this.resumeQueuedWork(); - return "applied"; - } - if (mutation.type === "move") { - const neighbor = projection[index + mutation.direction]; - if (!neighbor) return "rejected"; - this._actionStore.swapQueued(item, neighbor); - this._emitQueueUpdate(); - return "applied"; - } - if ( - item.payload.kind === "turn" && - (item.payload.acceptedAgentMessage || - item.payload.records.some((record) => record.role === "primary" && record.message.role !== "user")) - ) { - return "rejected"; - } - const images = mutation.images?.map((image) => ({ ...image })); - if (item.payload.kind === "session_command") { - const command = parseSessionSlashCommand(mutation.text); - if (!command) return "invalid"; - item.payload.text = mutation.text; - item.payload.command = command; - if (mutation.images !== undefined) item.payload.images = images?.length ? images : undefined; - } else { - item.payload.text = mutation.text; - const text = { type: "text" as const, text: mutation.text }; - if (mutation.images !== undefined) { - item.payload.images = images?.length ? images : undefined; - item.payload.content = [text, ...(images?.map((image) => ({ ...image })) ?? [])]; - } else if (item.payload.content) { - item.payload.content = [text, ...item.payload.content.filter((block) => block.type !== "text")]; - } - item.payload.preview = undefined; - item.payload.prepared = undefined; - for (const record of item.payload.records) { - if (record.role === "primary" && record.message.role === "user") { - record.message.content = item.payload.content?.map((block) => ({ ...block })) ?? mutation.text; - } - } - } - const targetPolicy = queuedMessageLaneDeliveryPolicy(mutation.lane); - if (targetPolicy !== policy) { - item.queueKey = undefined; - item.wake = mutation.lane === "steering" ? "on_lower_boundary" : "external_resume"; - this._actionStore.moveQueued(item, targetPolicy, this._actionStore.queuedActions(targetPolicy).length); - } - this.resumeQueuedWork(); - this._emitQueueUpdate(); - return "applied"; + ...args: Parameters + ): ReturnType { + return this._actionQueue.mutateQueuedMessage(...args); } get queuedActionCount(): number { @@ -4808,212 +2013,64 @@ export class AgentSession { ); } - getSessionActionSnapshot(): SessionActionSnapshot { - const steering = visibleSessionActionProjection(this._actionStore.queuedActions("next_turn_boundary")).map( - queuedAgentMessagePreview, - ); - const followUps = visibleSessionActionProjection(this._actionStore.queuedActions("when_run_idle")).map( - queuedAgentMessagePreview, - ); - const active = visibleSessionActionProjection(this._actionStore.activeActions())[0]; - const activeState = active?.lifecycle.state; - const phase = - activeState === "selected" - ? "preparing" - : activeState === "preparing" || activeState === "committing" || activeState === "running" - ? activeState - : undefined; - return { - queuedCount: steering.length + followUps.length, - steering, - followUps, - ...(active && phase - ? { - active: { - kind: active.payload.kind, - phase, - label: compactRlmText(active.payload.text), - }, - } - : {}), - }; - } - - getSteeringMessages(): readonly string[] { - return visibleSessionActionProjection(this._actionStore.queuedActions("next_turn_boundary")).map( - (action) => action.payload.text, - ); + getSessionActionSnapshot( + ...args: Parameters + ): ReturnType { + return this._actionQueue.getSessionActionSnapshot(...args); } - getSteeringMessagePreviews(): readonly string[] { - return visibleSessionActionProjection(this._actionStore.queuedActions("next_turn_boundary")).map( - queuedAgentMessagePreview, - ); + getSteeringMessages( + ...args: Parameters + ): ReturnType { + return this._actionQueue.getSteeringMessages(...args); } - getFollowUpMessages(): readonly string[] { - return visibleSessionActionProjection(this._actionStore.queuedActions("when_run_idle")).map( - (action) => action.payload.text, - ); + getSteeringMessagePreviews( + ...args: Parameters + ): ReturnType { + return this._actionQueue.getSteeringMessagePreviews(...args); } - getFollowUpMessagePreviews(): readonly string[] { - return visibleSessionActionProjection(this._actionStore.queuedActions("when_run_idle")).map( - queuedAgentMessagePreview, - ); + getFollowUpMessages( + ...args: Parameters + ): ReturnType { + return this._actionQueue.getFollowUpMessages(...args); } - getSessionActionRecoverySnapshot(): SessionActionRecoverySnapshot { - return { - formatVersion: SESSION_ACTION_RECOVERY_FORMAT_VERSION, - actions: this._actionStore.snapshotActions().map((action) => ({ - id: action.id, - source: action.source, - delivery: action.delivery, - wake: action.wake, - ...(action.queueKey ? { queueKey: action.queueKey } : {}), - ...(action.agentMessageId ? { agentMessageId: action.agentMessageId } : {}), - ...(action.suppressAutonomousContinuation ? { suppressAutonomousContinuation: true } : {}), - payload: - action.payload.kind === "turn" - ? { - kind: "turn", - text: action.payload.text, - ...(action.payload.preview ? { preview: action.payload.preview } : {}), - records: action.payload.records.map((record) => ({ - id: record.id, - role: record.role, - message: cloneQueuedAgentMessage(record.message), - ownerActionId: record.ownerActionId, - })), - ...(action.payload.images - ? { - images: action.payload.images.map((image) => ({ - ...image, - })), - } - : {}), - ...(action.payload.content - ? { - content: action.payload.content.map((block) => ({ - ...block, - })), - } - : {}), - ...(action.payload.customMessage - ? { - customMessage: cloneCustomMessage(action.payload.customMessage), - } - : {}), - executionPolicy: { - ...action.payload.executionPolicy, - preparation: { - ...action.payload.executionPolicy.preparation, - }, - }, - queueVisible: action.payload.queueVisible, - acceptedAgentMessage: action.payload.acceptedAgentMessage, - acceptedBeforeCompletion: action.payload.acceptedBeforeCompletion, - } - : { - kind: "session_command", - text: action.payload.text, - command: { ...action.payload.command }, - ...(action.payload.images - ? { - images: action.payload.images.map((image) => ({ - ...image, - })), - } - : {}), - }, - })), - }; + getFollowUpMessagePreviews( + ...args: Parameters + ): ReturnType { + return this._actionQueue.getFollowUpMessagePreviews(...args); } - private _notifySessionInputCheckpointChange(): void { - const waiters = [...this._sessionInputCheckpointWaiters]; - this._sessionInputCheckpointWaiters.clear(); - for (const resolve of waiters) resolve(); + getSessionActionRecoverySnapshot( + ...args: Parameters + ): ReturnType { + return this._actionRecovery.getSessionActionRecoverySnapshot(...args); } - private _waitForSessionActivityChange(signal: AbortSignal): Promise { - return new Promise((resolve) => { - const finish = () => { - this._sessionInputCheckpointWaiters.delete(finish); - signal.removeEventListener("abort", finish); - resolve(); - }; - this._sessionInputCheckpointWaiters.add(finish); - signal.addEventListener("abort", finish, { once: true }); - if (signal.aborted) finish(); - }); + private _notifySessionInputCheckpointChange( + ...args: Parameters + ): ReturnType { + return this._inputCheckpoints.notifySessionInputCheckpointChange(...args); } - - private _observeSessionActionDeferral(action: QueuedSessionAction): { - deferred: Promise; - stop(): void; - } { - let resolveDeferral = () => {}; - const deferred = new Promise((resolve) => { - resolveDeferral = resolve; - }); - const check = () => { - if (action.lifecycle.state === "queued") resolveDeferral(); - else this._sessionInputCheckpointWaiters.add(check); - }; - this._sessionInputCheckpointWaiters.add(check); - return { - deferred, - stop: () => this._sessionInputCheckpointWaiters.delete(check), - }; + + private _waitForSessionActivityChange( + ...args: Parameters + ): ReturnType { + return this._inputCheckpoints.waitForSessionActivityChange(...args); } - async waitForSessionInputCheckpoint(signal?: AbortSignal): Promise { - const blocksCheckpoint = () => - this._actionStore.activeActions().some((action) => { - if (action.payload.kind === "session_command") { - return action.lifecycle.state === "selected" || action.lifecycle.state === "running"; - } - return ( - action.lifecycle.state === "selected" || - action.lifecycle.state === "preparing" || - (action.lifecycle.state === "committing" && !primaryDeliveryRecord(action).durable) - ); - }); - while (true) { - while (blocksCheckpoint()) { - if (signal?.aborted) throw new Error("Update restart preparation cancelled"); - await new Promise((resolve, reject) => { - const onChange = () => { - cleanup(); - resolve(); - }; - const onAbort = () => { - cleanup(); - reject(new Error("Update restart preparation cancelled")); - }; - const cleanup = () => { - this._sessionInputCheckpointWaiters.delete(onChange); - signal?.removeEventListener("abort", onAbort); - }; - this._sessionInputCheckpointWaiters.add(onChange); - signal?.addEventListener("abort", onAbort, { once: true }); - if (signal?.aborted) onAbort(); - }); - } - const commitFence = await this._acquireSessionActionCommitFence(signal); - try { - if (blocksCheckpoint()) continue; - if (signal?.aborted) throw new Error("Update restart preparation cancelled"); - await waitForPromiseOrAbort(this._agentEventQueue, signal, "Update restart preparation cancelled"); - if (signal?.aborted) throw new Error("Update restart preparation cancelled"); - this.sessionManager.flushNow(); - return; - } finally { - commitFence.release(); - } - } + private _observeSessionActionDeferral( + ...args: Parameters + ): ReturnType { + return this._inputCheckpoints.observeSessionActionDeferral(...args); + } + + waitForSessionInputCheckpoint( + ...args: Parameters + ): ReturnType { + return this._inputCheckpoints.waitForSessionInputCheckpoint(...args); } acquireSessionInputPause(): { release(): void } { @@ -5033,45 +2090,10 @@ export class AgentSession { }); } - private async _acquireDirectTurnAdmissionFence(signal?: AbortSignal): Promise<{ owner: symbol; release(): void }> { - if (this._commitFence.isHeldByCurrentContext) { - this._assertSessionActionAdmissionAvailable(); - return this._acquireSessionActionCommitFence(signal); - } - const disposeSignal = this._commitFence.disposeSignal; - const waitSignal = signal ? AbortSignal.any([signal, disposeSignal]) : disposeSignal; - while (true) { - this._assertSessionActionAdmissionAvailable(); - if (this._inputScheduler.queuedWorkPauseCount > 0) { - let wake = () => {}; - const pauseReleased = new Promise((resolve) => { - wake = resolve; - this._sessionInputCheckpointWaiters.add(resolve); - }); - try { - await waitForPromiseOrAbort(pauseReleased, waitSignal, "Update restart preparation cancelled"); - } catch (error) { - if (disposeSignal.aborted) { - throw new Error("Cannot admit a session action because the session is disposing or disposed."); - } - throw error; - } finally { - this._sessionInputCheckpointWaiters.delete(wake); - } - continue; - } - const fence = await this._acquireSessionActionCommitFence(signal); - try { - if (this._inputScheduler.queuedWorkPauseCount === 0) { - this._assertSessionActionAdmissionAvailable(); - return fence; - } - } catch (error) { - fence.release(); - throw error; - } - fence.release(); - } + private _acquireDirectTurnAdmissionFence( + ...args: Parameters + ): ReturnType { + return this._inputCheckpoints.acquireDirectTurnAdmissionFence(...args); } private _acquireSessionActionCommitFence(signal?: AbortSignal): Promise { @@ -5106,91 +2128,33 @@ export class AgentSession { * waiter registered (a leaked waiter holds hasPendingAdmissionWaiters true and * blocks daemon passivation). */ - private async _waitForIdleOrSettlement(settlement?: ContinuationToken): Promise { - while (settlement === undefined || this._continuation.current === settlement) { - if (this._actionStore.queuedActions().length > 0) { - if (this._inputScheduler.suspended || this._inputScheduler.queuedWorkPauseCount > 0) { - let wake = () => {}; - const changed = new Promise((resolve) => { - wake = resolve; - this._sessionInputCheckpointWaiters.add(resolve); - }); - try { - await (settlement ? Promise.race([changed, settlement.promise]) : changed); - } finally { - this._sessionInputCheckpointWaiters.delete(wake); - } - continue; - } - this._scheduleSessionInputPump(); - } - const pump = this._inputScheduler.pendingPump; - await pump; - await this.agent.waitForIdle(); - const agentEventQueue = this._agentEventQueue; - await agentEventQueue; - if ( - pump === this._inputScheduler.pendingPump && - agentEventQueue === this._agentEventQueue && - !this._inputScheduler.requested && - !this.agent.state.isStreaming && - this.unfinishedActionCount === 0 - ) { - return; - } - } + private _waitForIdleOrSettlement( + ...args: Parameters + ): ReturnType { + return this._inputCheckpoints.waitForIdleOrSettlement(...args); } /** Waits out any owned post-compaction continuation and rejects when one cannot start; {@link waitForIdle} never rejects. */ - async waitForHeadlessIdle(): Promise { - while (true) { - await this.waitForIdle(); - const postCompactionContinuation = this._continuation.current?.promise; - if (!postCompactionContinuation) return; - await postCompactionContinuation; - } + waitForHeadlessIdle(): Promise { + return this._inputCheckpoints.waitForHeadlessIdle(); } - getPendingNextTurnMessageSnapshots(): readonly CustomMessage[] { - const messages = this._pendingNextTurnMessages.map((message) => cloneCustomMessage(message)); - for (const action of this._actionStore.unfinishedActions()) { - if ( - action.payload.kind !== "turn" || - !action.payload.acceptedAgentMessage || - !primaryDeliveryRecord(action).started - ) { - continue; - } - messages.push( - ...action.payload.records - .filter( - (record): record is DeliveryRecord & { message: CustomMessage } => - (record.role === "next_turn" || record.role === "prefix") && - record.message.role === "custom" && - !record.durable, - ) - .map((record) => cloneCustomMessage(record.message)), - ); - } - return messages; + getPendingNextTurnMessageSnapshots( + ...args: Parameters + ): ReturnType { + return this._pendingContext.getPendingNextTurnMessageSnapshots(...args); } - restorePendingNextTurnMessages(messages: readonly CustomMessage[]): void { - this._pendingNextTurnMessages.push(...messages.map((message) => cloneCustomMessage(message))); - this._flushDeferredRlmTerminalNotices(); + restorePendingNextTurnMessages( + ...args: Parameters + ): ReturnType { + return this._pendingContext.restorePendingNextTurnMessages(...args); } - removeQueuedFollowUp(queueKey: string): boolean { - const matching = this._actionStore - .clearableActions() - .filter((action) => action.payload.kind === "turn" && action.queueKey === queueKey); - if (matching.length === 0) return false; - const error = new Error("Queued agent message was cleared before delivery."); - for (const action of matching) this._rejectAgentMessage(action.agentMessageId, error); - const ids = new Set(matching.map((action) => action.id)); - this._cancelSessionActions((action) => ids.has(action.id), error); - this._emitQueueUpdate(); - return true; + removeQueuedFollowUp( + ...args: Parameters + ): ReturnType { + return this._actionQueue.removeQueuedFollowUp(...args); } get resourceLoader(): ResourceLoader { @@ -5205,7 +2169,7 @@ export class AgentSession { (action) => action.payload.kind === "turn" && !action.payload.queueVisible && - !this._durableRlmTerminalNoticeActionIds.has(action.id), + !this._pendingContext.isRetainedTerminalNotice(action.id), new Error("Prompt aborted before delivery."), ); this._cancelPostCompactionContinue(); @@ -5222,16 +2186,16 @@ export class AgentSession { const branchSummaryOperation = this._branchSummaryOperation; this.requestAbort(); this._cancelActiveRlmChildRuns("Parent session aborted"); - this._goalAbortInProgress = this._goals.state.status === "active"; + this._goalContinuation.beginAbort(); try { await Promise.allSettled([ this.agent.waitForIdle(), - this._agentEventQueue, + this._events.queue, ...(compactionOperation ? [compactionOperation] : []), ...(branchSummaryOperation ? [branchSummaryOperation] : []), ]); } finally { - this._goalAbortInProgress = false; + this._goalContinuation.finishAbort(); } } @@ -5243,301 +2207,63 @@ export class AgentSession { this.abortRetry(); this._children.cancelQuiescenceWaits(); this._cancelActiveRlmChildRuns("Parent session aborted for update restart"); - this._goalAbortInProgress = this._goals.state.status === "active"; + this._goalContinuation.beginAbort(); this.agent.abort(); - if (this._goalAbortInProgress) { + if (this._goalContinuation.abortInProgress) { void this.agent .waitForIdle() - .then(() => this._agentEventQueue) + .then(() => this._events.queue) .catch(() => undefined) .finally(() => { - this._goalAbortInProgress = false; + this._goalContinuation.finishAbort(); }); } } - private async _emitModelSelect( - nextModel: Model, - previousModel: Model | undefined, - source: "set" | "cycle" | "restore", - ): Promise { - if (modelsAreEqual(previousModel, nextModel)) return; - await this._extensionRunner.emit({ - type: "model_select", - model: nextModel, - previousModel, - source, - }); - } - - private _queueModelSelectEmit( - nextModel: Model, - previousModel: Model | undefined, - source: "set" | "cycle" | "restore", - ): Promise { - const emit = () => - this._modelSelectEmitContext.run(true, () => this._emitModelSelect(nextModel, previousModel, source)); - this._modelSelectEmitQueueIdle = false; - const promise = this._modelSelectEmitQueue.then(emit, emit); - const queued = promise.catch(() => {}); - this._modelSelectEmitQueue = queued; - void queued.finally(() => { - if (this._modelSelectEmitQueue === queued) { - this._modelSelectEmitQueueIdle = true; - } - }); - return promise; - } - - async setModel(model: Model, options: ModelSelectOptions = {}): Promise { - // Explicit selection recovers from a stale-auth lockout, but only a fully - // validated switch commits the clear (single owner): failed selections never unlock. - const staleOnly = - !this._modelRegistry.hasConfiguredAuth(model) && - this._modelRegistry.getProviderAuthStatus(model.provider).source === "stale"; - if (!staleOnly && !this._modelRegistry.hasConfiguredAuth(model)) { - throw new Error(`No API key for ${model.provider}/${model.id}`); - } - if (!(await this._modelRegistry.canUseModel(model, { assumeAuthConfigured: staleOnly }))) { - throw new Error(`Model "${model.provider}/${model.id}" is not available for the current Prime team.`); - } - if (staleOnly) { - this._modelRegistry.clearProviderAuthStale(model.provider); - if (!this._modelRegistry.hasConfiguredAuth(model)) { - throw new Error(`No API key for ${model.provider}/${model.id}`); - } - } - - const previousModel = this.model; - const thinkingLevel = this._getThinkingLevelForModelSwitch(); - const serviceTier = this._getServiceTierForModelSwitch(); - this.agent.state.model = model; - this.sessionManager.appendModelChange(model.provider, model.id); - this.settingsManager.setDefaultModelAndProvider(model.provider, model.id); - - this.setThinkingLevel(thinkingLevel); - this._clampServiceTierForModel(serviceTier); - - const emitPromise = this._queueModelSelectEmit(model, previousModel, "set"); - if (this._shouldWaitForModelSelectEmit(options)) { - await emitPromise; - } else { - this._trackModelSelectEmitError(emitPromise); - } - } - - private _trackModelSelectEmitError(emitPromise: Promise): void { - void emitPromise.catch((error) => { - this._extensionRunner.emitError({ - extensionPath: "", - event: "model_select", - error: error instanceof Error ? error.message : String(error), - stack: error instanceof Error ? error.stack : undefined, - }); - }); - } - - private _shouldWaitForModelSelectEmit(options: ModelSelectOptions): boolean { - return options.waitForExtensions !== false && !this._modelSelectEmitContext.getStore(); - } - - private _pendingModelSelectEmit(): Promise | undefined { - if (!this._modelSelectEmitContext.getStore() && !this._modelSelectEmitQueueIdle) { - return this._modelSelectEmitQueue; - } - return undefined; - } - - async cycleModel( - direction: "forward" | "backward" = "forward", - options: ModelSelectOptions = {}, - ): Promise { - if (this._scopedModels.length > 0) { - return this._cycleScopedModel(direction, options); - } - return this._cycleAvailableModel(direction, options); - } - - private async _cycleScopedModel( - direction: "forward" | "backward", - options: ModelSelectOptions, - ): Promise { - const availableModels = await this._modelRegistry.refreshAvailableModels(); - const scopedModels = this._scopedModels.filter((scoped) => - availableModels.some((model) => modelsAreEqual(model, scoped.model)), - ); - if (scopedModels.length <= 1) return undefined; - - const currentModel = this.model; - let currentIndex = scopedModels.findIndex((sm) => modelsAreEqual(sm.model, currentModel)); - - if (currentIndex === -1) currentIndex = 0; - const len = scopedModels.length; - const nextIndex = direction === "forward" ? (currentIndex + 1) % len : (currentIndex - 1 + len) % len; - const next = scopedModels[nextIndex]; - const thinkingLevel = this._getThinkingLevelForModelSwitch(next.thinkingLevel); - const serviceTier = this._getServiceTierForModelSwitch(); - - this.agent.state.model = next.model; - this.sessionManager.appendModelChange(next.model.provider, next.model.id); - this.settingsManager.setDefaultModelAndProvider(next.model.provider, next.model.id); - - this.setThinkingLevel(thinkingLevel); - this._clampServiceTierForModel(serviceTier); - - const emitPromise = this._queueModelSelectEmit(next.model, currentModel, "cycle"); - if (this._shouldWaitForModelSelectEmit(options)) { - await emitPromise; - } else { - this._trackModelSelectEmitError(emitPromise); - } - - return { - model: next.model, - thinkingLevel: this.thinkingLevel, - serviceTier: this.serviceTier, - isScoped: true, - }; - } - - private async _cycleAvailableModel( - direction: "forward" | "backward", - options: ModelSelectOptions, - ): Promise { - const availableModels = await this._modelRegistry.refreshAvailableModels(); - if (availableModels.length <= 1) return undefined; - - const currentModel = this.model; - let currentIndex = availableModels.findIndex((m) => modelsAreEqual(m, currentModel)); - - if (currentIndex === -1) currentIndex = 0; - const len = availableModels.length; - const nextIndex = direction === "forward" ? (currentIndex + 1) % len : (currentIndex - 1 + len) % len; - const nextModel = availableModels[nextIndex]; - - const thinkingLevel = this._getThinkingLevelForModelSwitch(); - const serviceTier = this._getServiceTierForModelSwitch(); - this.agent.state.model = nextModel; - this.sessionManager.appendModelChange(nextModel.provider, nextModel.id); - this.settingsManager.setDefaultModelAndProvider(nextModel.provider, nextModel.id); - - this.setThinkingLevel(thinkingLevel); - this._clampServiceTierForModel(serviceTier); - - const emitPromise = this._queueModelSelectEmit(nextModel, currentModel, "cycle"); - if (this._shouldWaitForModelSelectEmit(options)) { - await emitPromise; - } else { - this._trackModelSelectEmitError(emitPromise); - } - - return { - model: nextModel, - thinkingLevel: this.thinkingLevel, - serviceTier: this.serviceTier, - isScoped: false, - }; - } - - setThinkingLevel(level: ThinkingLevel): void { - const availableLevels = this.getAvailableThinkingLevels(); - const effectiveLevel = availableLevels.includes(level) ? level : this._clampThinkingLevel(level, availableLevels); - - const previousLevel = this.agent.state.thinkingLevel; - const isChanging = effectiveLevel !== previousLevel; - - this.agent.state.thinkingLevel = effectiveLevel; - - if (isChanging) { - this.sessionManager.appendThinkingLevelChange(effectiveLevel); - if (this.supportsThinking() || effectiveLevel !== "off") { - this.settingsManager.setDefaultThinkingLevel(effectiveLevel); - } - this._emit({ type: "thinking_level_changed", level: effectiveLevel }); - void this._extensionRunner.emit({ - type: "thinking_level_select", - level: effectiveLevel, - previousLevel, - }); - } - } - - setServiceTier(serviceTier: ServiceTier): void { - const effectiveServiceTier = this._getEffectiveServiceTier(serviceTier); - const preferenceChanged = effectiveServiceTier !== this._serviceTierPreference; - const effectiveTierChanged = effectiveServiceTier !== this.agent.state.serviceTier; - if (!preferenceChanged && !effectiveTierChanged) { - return; - } - this._serviceTierPreference = effectiveServiceTier; - if (preferenceChanged) { - this.sessionManager.appendServiceTierChange(effectiveServiceTier); - if (this.model && supportsFastMode(this.model)) { - this.settingsManager.setDefaultServiceTier(effectiveServiceTier); - } - } - if (effectiveTierChanged) { - this.agent.state.serviceTier = effectiveServiceTier; - this._emit({ - type: "service_tier_changed", - serviceTier: effectiveServiceTier, - }); - } - } - - private _getEffectiveServiceTier(serviceTier: ServiceTier): ServiceTier { - return serviceTier === "priority" && (!this.model || !supportsFastMode(this.model)) ? "default" : serviceTier; + setModel(...args: Parameters): ReturnType { + return this._modelSelection.setModel(...args); } - private _getServiceTierForModelSwitch(): ServiceTier { - return this._serviceTierPreference; + private _pendingModelSelectEmit( + ...args: Parameters + ): ReturnType { + return this._modelSelection.pendingModelSelectEmit(...args); } - private _clampServiceTierForModel(serviceTier: ServiceTier = this.serviceTier): void { - const effectiveServiceTier = this._getEffectiveServiceTier(serviceTier); - if (effectiveServiceTier === this.agent.state.serviceTier) { - return; - } - this.agent.state.serviceTier = effectiveServiceTier; - this._emit({ - type: "service_tier_changed", - serviceTier: effectiveServiceTier, - }); + cycleModel( + ...args: Parameters + ): ReturnType { + return this._modelSelection.cycleModel(...args); } - cycleThinkingLevel(): ThinkingLevel | undefined { - if (!this.supportsThinking()) return undefined; - - const levels = this.getAvailableThinkingLevels(); - const currentIndex = levels.indexOf(this.thinkingLevel); - const nextIndex = (currentIndex + 1) % levels.length; - const nextLevel = levels[nextIndex]; - - this.setThinkingLevel(nextLevel); - return nextLevel; + setThinkingLevel( + ...args: Parameters + ): ReturnType { + return this._modelSelection.setThinkingLevel(...args); } - getAvailableThinkingLevels(): ThinkingLevel[] { - if (!this.model) return THINKING_LEVELS; - return getSupportedThinkingLevels(this.model) as ThinkingLevel[]; + setServiceTier( + ...args: Parameters + ): ReturnType { + return this._modelSelection.setServiceTier(...args); } - supportsThinking(): boolean { - return !!this.model?.reasoning; + cycleThinkingLevel( + ...args: Parameters + ): ReturnType { + return this._modelSelection.cycleThinkingLevel(...args); } - private _getThinkingLevelForModelSwitch(explicitLevel?: ThinkingLevel): ThinkingLevel { - if (explicitLevel !== undefined) { - return explicitLevel; - } - if (!this.supportsThinking()) { - return this.settingsManager.getDefaultThinkingLevel() ?? DEFAULT_THINKING_LEVEL; - } - return this.thinkingLevel; + getAvailableThinkingLevels( + ...args: Parameters + ): ReturnType { + return this._modelSelection.getAvailableThinkingLevels(...args); } - private _clampThinkingLevel(level: ThinkingLevel, _availableLevels: ThinkingLevel[]): ThinkingLevel { - return this.model ? (clampThinkingLevel(this.model, level) as ThinkingLevel) : "off"; + supportsThinking( + ...args: Parameters + ): ReturnType { + return this._modelSelection.supportsThinking(...args); } private _syncKernelStateAfterCompaction(): Promise { @@ -5565,7 +2291,9 @@ export class AgentSession { ): void { this._refinement._discardPendingAutoRefine({ cancelPostCompactionContinue: true }); if (this._goals.state.status === "active" && !signal.aborted) { - this._goalContinuationAwaitsRlmWork ||= !this.agent.hasQueuedMessages(); + if (!this._goalContinuation.awaitsChildWork && !this.agent.hasQueuedMessages()) { + this._goalContinuation.deferUntilChildSettlement(); + } this.resumeQueuedWork(); if (this.agent.hasQueuedMessages()) this._schedulePostCompactionContinue(); } @@ -5575,7 +2303,7 @@ export class AgentSession { // Queued agent or session-owned inputs resume the loop; defer refine // behind them instead of interleaving it before their turns. this._refinement._scheduleAutoRefineAfterCompaction( - this._goalContinuationAwaitsRlmWork || + this._goalContinuation.awaitsChildWork || hadPostCompactionContinue || this.agent.hasQueuedMessages() || this.unfinishedActionCount > 0, @@ -5612,62 +2340,23 @@ export class AgentSession { } /** The compact harness digest delivered at cold context boundaries (session start, resume, compaction head). */ - private _harnessDigest(): string { - const tools = this.getActiveToolNames(); - const hasIpython = tools.includes("ipython"); - const visibleSkills = this._modelVisibleSkills().filter((skill) => !skill.disableModelInvocation); - const hasRefineSkill = visibleSkills.some((skill) => skill.name === REFINE_SKILL_NAME); - return formatHarnessStateForPrompt(this._refinement._loadMergedHarnessState(), { - includeIpythonExamples: hasIpython, - includeShellExamples: tools.includes("bash"), - includeRefineExamples: hasIpython && hasRefineSkill, - }); + private _harnessDigest( + ...args: Parameters + ): ReturnType { + return this._harnessContext.harnessDigest(...args); } /** Cold-boundary digest delivery: empty contexts defer to the first committed turn (untouched sessions must stay empty); non-empty contexts append only when the newest in-context digest mismatches disk. */ - private _ensureHarnessDigestContext(): void { - if (this.agent.state.messages.length === 0) { - this._harnessDigestPending = true; - return; - } - this._harnessDigestPending = false; - this._appendHarnessDigestIfStale(); + private _ensureHarnessDigestContext( + ...args: Parameters + ): ReturnType { + return this._harnessContext.ensureHarnessDigestContext(...args); } - private _appendHarnessDigestIfStale(): void { - const digest = this._harnessDigest(); - if (this._latestContextHarnessDigest() === digest) return; - const message = createHarnessDigestMessage(digest); - try { - this.sessionManager.appendCustomMessageEntryWithRollback( - message.customType, - message.content, - message.display, - message.details, - ); - } catch { - // Unpersisted session: context-only injection. - } - this.agent.state.messages.push(message); - } - - private _latestContextHarnessDigest(): string | undefined { - // Retained pre-compaction messages follow the compaction head, so recency is by timestamp, not position. - let latest: { timestamp: number; digest: string } | undefined; - for (const message of this.agent.state.messages) { - let digest: string | undefined; - if (message.role === "custom" && message.customType === HARNESS_DIGEST_CUSTOM_TYPE) { - digest = (message.details as HarnessDigestDetails | undefined)?.digest; - } else if (message.role === "compactionSummary") { - digest = message.harnessDigest; - } else { - continue; - } - if (digest !== undefined && (!latest || message.timestamp >= latest.timestamp)) { - latest = { timestamp: message.timestamp, digest }; - } - } - return latest?.digest; + private _latestContextHarnessDigest( + ...args: Parameters + ): ReturnType { + return this._harnessContext.latestContextHarnessDigest(...args); } /** @@ -5685,8 +2374,10 @@ export class AgentSession { return this._refinement.refine(options, internal); } - abortBranchSummary(): void { - this._branchSummaryAbortController?.abort(); + abortBranchSummary( + ...args: Parameters + ): ReturnType { + return this._history.abortBranchSummary(...args); } /** @@ -5739,30 +2430,13 @@ export class AgentSession { } refreshModelMetadata(): void { - if (this.model?.provider === "xai") { - this.agent.state.model = this._modelRegistry.getModelForCurrentAuth(this.model); - this.setThinkingLevel(this.thinkingLevel); - this._clampServiceTierForModel(); - } - this._scopedModels = this._scopedModels.map((scoped) => - scoped.model.provider === "xai" - ? { ...scoped, model: this._modelRegistry.getModelForCurrentAuth(scoped.model) } - : scoped, - ); + this._modelSelection.refreshModelMetadata(); } - 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 _refreshCurrentModelFromRegistry( + ...args: Parameters + ): ReturnType { + return this._modelSelection.refreshCurrentModelFromRegistry(...args); } private _refreshToolRegistry(options?: { activeToolNames?: string[]; includeAllExtensionTools?: boolean }): void { @@ -5822,7 +2496,7 @@ export class AgentSession { listSubagents: () => this.listRlmSubagents(), deleteSubagent: (target) => this.deleteRlmSubagent(target), handleBashCompletion: (details) => this._handleKernelBashCompletion(details), - withdrawBashCompletion: (details) => this._withdrawAsyncBashCompletionNotice(details), + withdrawBashCompletion: (details) => this._actionQueue.withdrawAsyncBashCompletionNotice(details), getModel: () => this.model, includeGoals: this._includeGoals, includeCompactSkill: this._includeCompactSkill, @@ -5844,30 +2518,10 @@ export class AgentSession { }); } - 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); - } - } - } + private _handleKernelBashCompletion( + ...args: Parameters + ): ReturnType { + return this._promptSubmission.handleKernelBashCompletion(...args); } reload(): Promise { @@ -6051,47 +2705,16 @@ export class AgentSession { return this._children.cancelRunningRlmDescendants(reason); } - private async _authenticatedRlmModels(): Promise[]> { - return (await this._modelRegistry.getExecutableModels()).filter((model) => { - const status = this._modelRegistry.getProviderAuthStatus(model.provider); - return status.source !== "stale" && status.label !== "expired"; - }); - } - - async findRlmModels(query: string, limit: number): Promise { - return { - models: findRlmModelMatches(query, await this._authenticatedRlmModels(), limit), - }; + findRlmModels( + ...args: Parameters + ): ReturnType { + return this._modelSelection.findRlmModels(...args); } - private async _resolveRlmSubagentModel( - reference: string | undefined, - target = "subagent", - ): Promise { - const parentModel = this.model; - if (!parentModel) { - throw new Error(formatNoModelSelectedMessage()); - } - if (!reference) { - return { model: parentModel }; - } - - const normalizedReference = reference.toLowerCase(); - if (`${parentModel.provider}/${parentModel.id}`.toLowerCase() === normalizedReference) { - return { model: parentModel }; - } - const model = (await this._authenticatedRlmModels()).find( - (candidate) => `${candidate.provider}/${candidate.id}`.toLowerCase() === normalizedReference, - ); - if (!model) { - throw new Error(`Requested ${target} model "${reference}" is unavailable, unauthenticated, or expired`); - } - - const auth = await this._modelRegistry.getApiKeyAndHeaders(model); - if (!auth.ok) { - throw new Error(`Requested ${target} model "${reference}" failed authentication preflight`); - } - return { model }; + private _resolveRlmSubagentModel( + ...args: Parameters + ): ReturnType { + return this._modelSelection.resolveRlmSubagentModel(...args); } createRlmSession(prompt: string, kwargs: Record = {}): Promise { @@ -6119,14 +2742,7 @@ export class AgentSession { } get hasAcceptedPromptInFlight(): boolean { - return this._actionStore - .unfinishedActions() - .some( - (action) => - action.payload.kind === "turn" && - !action.payload.queueVisible && - action.payload.acceptedBeforeCompletion, - ); + return this._actionQueue.hasAcceptedPromptInFlight; } get autoRetryEnabled(): boolean { @@ -6200,414 +2816,60 @@ export class AgentSession { * @param options.label Label to attach to the branch summary entry * @returns Result with editorText (if user message) and cancelled status */ - private _branchNavigationQueue: Promise = Promise.resolve(); - - async navigateTree( - targetId: string, - options: { - summarize?: boolean; - customInstructions?: string; - replaceInstructions?: boolean; - label?: string; - } = {}, - ): Promise<{ - editorText?: string; - cancelled: boolean; - aborted?: boolean; - summaryEntry?: BranchSummaryEntry; - }> { - const previous = this._branchNavigationQueue; - let release = () => {}; - this._branchNavigationQueue = new Promise((resolve) => { - release = resolve; - }); - await previous; - try { - return await this._navigateTree(targetId, options); - } finally { - release(); - } - } - - private async _navigateTree( - targetId: string, - options: { - summarize?: boolean; - customInstructions?: string; - replaceInstructions?: boolean; - label?: string; - } = {}, - ): Promise<{ - editorText?: string; - cancelled: boolean; - aborted?: boolean; - summaryEntry?: BranchSummaryEntry; - }> { - if (options.summarize && !this.model) { - throw new Error("No model available for summarization"); - } - - const targetEntry = this.sessionManager.getEntry(targetId); - if (!targetEntry) { - throw new Error(`Entry ${targetId} not found`); - } - - const queuedWorkPause = this.acquireQueuedWorkPause(); - let commitFence: { owner: symbol; release(): void } | undefined; - try { - // Branch navigation and turn dispatch mutate the same transcript leaf. - commitFence = await this._acquireSessionActionCommitFence(); - return await this._commitFence.run(commitFence, async () => { - await this.agent.waitForIdle(); - await this._agentEventQueue; - return this._navigateTreeUnderPause(targetId, targetEntry, options); - }); - } finally { - queuedWorkPause.release(); - commitFence?.release(); - } - } - - private async _navigateTreeUnderPause( - targetId: string, - targetEntry: NonNullable>, - options: { - summarize?: boolean; - customInstructions?: string; - replaceInstructions?: boolean; - label?: string; - }, - ): Promise<{ - editorText?: string; - cancelled: boolean; - aborted?: boolean; - summaryEntry?: BranchSummaryEntry; - }> { - const oldLeafId = this.sessionManager.getLeafId(); - - // No-op if already at target after admitted work has settled. - if (targetId === oldLeafId) { - return { cancelled: false }; - } - - // Do not switch branches while /refine has detached event handling and is - // about to persist harness/session entries for the current branch. - await this._refinement._invalidatePendingAutoRefineForBranchChange(); - - const { entries: entriesToSummarize, commonAncestorId } = collectEntriesForBranchSummary( - this.sessionManager, - oldLeafId, - targetId, - ); - - let customInstructions = options.customInstructions; - let replaceInstructions = options.replaceInstructions; - let label = options.label; - - const preparation: TreePreparation = { - targetId, - oldLeafId, - commonAncestorId, - entriesToSummarize, - userWantsSummary: options.summarize ?? false, - customInstructions, - replaceInstructions, - label, - }; - - this._branchSummaryAbortController = new AbortController(); - let resolveBranchSummaryOperation: () => void = () => {}; - const branchSummaryOperation = new Promise((resolve) => { - resolveBranchSummaryOperation = resolve; - }); - this._branchSummaryOperation = branchSummaryOperation; - - try { - let extensionSummary: { summary: string; details?: unknown } | undefined; - let fromExtension = false; - - if (this._extensionRunner.hasHandlers("session_before_tree")) { - const result = (await this._extensionRunner.emit({ - type: "session_before_tree", - preparation, - signal: this._branchSummaryAbortController.signal, - })) as SessionBeforeTreeResult | undefined; - - if (result?.cancel) { - return { cancelled: true }; - } - - if (result?.summary && options.summarize) { - extensionSummary = result.summary; - fromExtension = true; - } - - if (result?.customInstructions !== undefined) { - customInstructions = result.customInstructions; - } - if (result?.replaceInstructions !== undefined) { - replaceInstructions = result.replaceInstructions; - } - if (result?.label !== undefined) { - label = result.label; - } - } - - let summaryText: string | undefined; - let summaryDetails: unknown; - let summaryUsage: Usage | undefined; - if (options.summarize && entriesToSummarize.length > 0 && !extensionSummary) { - const { apiKey, headers, requestModel: model } = await this._getRequiredRequestAuth(this.model!); - const branchSummarySettings = this.settingsManager.getBranchSummarySettings(); - const result = await generateBranchSummary(entriesToSummarize, { - model, - apiKey, - headers, - signal: this._branchSummaryAbortController.signal, - sessionId: this.sessionId, - customInstructions, - replaceInstructions, - reserveTokens: branchSummarySettings.reserveTokens, - retry: providerRetryPolicy(this.settingsManager), - }); - if (result.aborted) { - return { cancelled: true, aborted: true }; - } - if (result.error) { - throw new Error(result.error); - } - summaryText = result.summary; - summaryUsage = result.usage; - summaryDetails = { - readFiles: result.readFiles || [], - modifiedFiles: result.modifiedFiles || [], - }; - } else if (extensionSummary) { - summaryText = extensionSummary.summary; - summaryDetails = extensionSummary.details; - } - - let newLeafId: string | null; - let editorText: string | undefined; - - if (targetEntry.type === "message" && targetEntry.message.role === "user") { - newLeafId = targetEntry.parentId; - editorText = this._extractUserMessageText(targetEntry.message.content); - } else if (targetEntry.type === "custom_message") { - newLeafId = targetEntry.parentId; - editorText = - typeof targetEntry.content === "string" - ? targetEntry.content - : targetEntry.content - .filter((c): c is { type: "text"; text: string } => c.type === "text") - .map((c) => c.text) - .join(""); - } else { - newLeafId = targetId; - } - - let summaryEntry: BranchSummaryEntry | undefined; - if (summaryText) { - const summaryId = this.sessionManager.branchWithSummary( - newLeafId, - summaryText, - summaryDetails, - fromExtension, - summaryUsage, - ); - summaryEntry = this.sessionManager.getEntry(summaryId) as BranchSummaryEntry; - - if (label) { - this.sessionManager.appendLabelChange(summaryId, label); - } - } else if (newLeafId === null) { - this.sessionManager.resetLeaf(); - } else { - this.sessionManager.branch(newLeafId); - } - - if (label && !summaryText) { - this.sessionManager.appendLabelChange(targetId, label); - } - - const sessionContext = this.sessionManager.buildSessionContext(); - this.agent.state.messages = sessionContext.messages; - this._mergeUnpersistedOutcomes(this.agent.state.messages); - this._restoreLateIpythonSentAgentMessages(); - // Context rebuild = cold boundary: refresh the digest like resume. - this._ensureHarnessDigestContext(); - this._goals.reload(); - this._reloadRlmMaxDepthFromBranch(); - this._invalidateQueuedPromptPreparation(); - - await this._extensionRunner.emit({ - type: "session_tree", - newLeafId: this.sessionManager.getLeafId(), - oldLeafId, - summaryEntry, - fromExtension: summaryText ? fromExtension : undefined, - }); - return { editorText, cancelled: false, summaryEntry }; - } finally { - this._branchSummaryAbortController = undefined; - if (this._branchSummaryOperation === branchSummaryOperation) { - this._branchSummaryOperation = undefined; - } - resolveBranchSummaryOperation(); - this._notifySessionInputCheckpointChange(); - } + navigateTree( + ...args: Parameters + ): ReturnType { + return this._history.navigateTree(...args); } - getUserMessagesForForking(): Array<{ entryId: string; text: string }> { - const entries = this.sessionManager.getEntries(); - const result: Array<{ entryId: string; text: string }> = []; - - for (const entry of entries) { - if (entry.type !== "message") continue; - if (entry.message.role !== "user") continue; - - const text = this._extractUserMessageText(entry.message.content); - if (text) { - result.push({ entryId: entry.id, text }); - } - } - - return result; + getUserMessagesForForking( + ...args: Parameters + ): ReturnType { + return this._history.getUserMessagesForForking(...args); } - private _extractUserMessageText(content: string | Array<{ type: string; text?: string }>): string { - if (typeof content === "string") return content; - if (Array.isArray(content)) { - return content - .filter((c): c is { type: "text"; text: string } => c.type === "text") - .map((c) => c.text) - .join(""); - } - return ""; - } - - getSessionStats(): SessionStats { - const state = this.state; - const userMessages = state.messages.filter((m) => m.role === "user").length; - const assistantMessages = state.messages.filter((m) => m.role === "assistant").length; - const toolResults = state.messages.filter((m) => m.role === "toolResult").length; - - let toolCalls = 0; - let totalInput = 0; - let totalOutput = 0; - let totalCacheRead = 0; - let totalCacheWrite = 0; - let totalCost = 0; - - for (const message of state.messages) { - if (message.role === "assistant") { - const assistantMsg = message as AssistantMessage; - toolCalls += assistantMsg.content.filter((c) => c.type === "toolCall").length; - totalInput += assistantMsg.usage.input; - totalOutput += assistantMsg.usage.output; - totalCacheRead += assistantMsg.usage.cacheRead; - totalCacheWrite += assistantMsg.usage.cacheWrite; - totalCost += assistantMsg.usage.cost.total; - } - } - - return { - sessionFile: this.sessionFile, - sessionId: this.sessionId, - userMessages, - assistantMessages, - toolCalls, - toolResults, - totalMessages: state.messages.length, - tokens: { - input: totalInput, - output: totalOutput, - cacheRead: totalCacheRead, - cacheWrite: totalCacheWrite, - total: totalInput + totalOutput + totalCacheRead + totalCacheWrite, - }, - cost: totalCost, - contextUsage: this.getContextUsage(), - }; + getSessionStats( + ...args: Parameters + ): ReturnType { + return this._contextView.getSessionStats(...args); } - getContextUsage(): ContextUsage | undefined { - const model = this.model; - if (!model) return undefined; - - const contextWindow = model.contextWindow ?? 0; - if (contextWindow <= 0) return undefined; - - // After compaction, the last assistant usage reflects pre-compaction context size. - // We can only trust usage from an assistant that responded after the latest compaction. - // If no such assistant exists, context token count is unknown until the next LLM response. - const branchEntries = this.sessionManager.getBranch(); - const latestCompaction = getLatestCompactionEntry(branchEntries); - - if (latestCompaction) { - // Check if there's a valid assistant usage after the compaction boundary - const compactionIndex = branchEntries.lastIndexOf(latestCompaction); - let hasPostCompactionUsage = false; - for (let i = branchEntries.length - 1; i > compactionIndex; i--) { - const entry = branchEntries[i]; - if (entry.type === "message" && entry.message.role === "assistant") { - const assistant = entry.message; - if (assistant.stopReason !== "aborted" && assistant.stopReason !== "error") { - const contextTokens = calculateContextTokens(assistant.usage); - if (contextTokens > 0) { - hasPostCompactionUsage = true; - } - break; - } - } - } - - if (!hasPostCompactionUsage) { - return { tokens: null, contextWindow, percent: null }; - } - } - - const estimate = estimateContextTokens(this.messages); - const percent = (estimate.tokens / contextWindow) * 100; - - return { - tokens: estimate.tokens, - contextWindow, - percent, - }; + getContextUsage( + ...args: Parameters + ): ReturnType { + return this._contextView.getContextUsage(...args); } private _rlmSessionDirForReading(): string | undefined { return this._rlmSessionDir ?? this.sessionManager.getSessionArtifactDir(); } - private _contextWindowResolver(): ContextWindowResolver { - return (provider, modelId) => this._modelRegistry.find(provider, modelId)?.contextWindow; - } - - private _subtractUnindexedChildUsage(ownUsage: Usage, entries: SessionEntry[]): void { - this._childUsage.subtractUnindexed(ownUsage, entries); + private *_contextViewChildren(): Generator { + for (const run of this._children.getActiveRuns()) { + yield { + id: run.id, + get label() { + return rlmChildLabel(run.prompt); + }, + get status() { + return run.status; + }, + sessionDir: run.sessionDir, + getContextTree: run.session ? () => run.session!.getContextTree() : undefined, + }; + } } private _invalidateOwnUsage(): void { - this._ownUsageMemo = undefined; + this._contextView.invalidateOwnUsage(); } - 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. - getOwnUsageSummary(): SessionUsageSummary | undefined { - const entries = this.sessionManager.getEntries(); - const tailId = entries.at(-1)?.id; - const memo = this._ownUsageMemo; - if (memo && memo.count === entries.length && memo.tailId === tailId) { - return memo.usage; - } - const { ownUsage } = computeOwnAndTotalUsage(entries, entries); - this._subtractUnindexedChildUsage(ownUsage, entries); - const usage = sessionUsageSummaryFrom(ownUsage); - this._ownUsageMemo = { count: entries.length, tailId, usage }; - return usage; + getOwnUsageSummary( + ...args: Parameters + ): ReturnType { + return this._contextView.getOwnUsageSummary(...args); } /** @@ -6616,42 +2878,10 @@ export class AgentSession { * from their live sessions; completed children from their persisted session * dirs, so the tree survives child disposal and session resume. */ - getContextTree(): ContextTreeNode { - const resolveContextWindow = this._contextWindowResolver(); - const branch = this.sessionManager.getBranch(); - const { ownUsage, totalUsage } = computeOwnAndTotalUsage(branch, this.sessionManager.getEntries()); - this._subtractUnindexedChildUsage(ownUsage, branch); - - const children: ContextTreeNode[] = []; - const liveIds = new Set(); - for (const run of this._children.getActiveRuns()) { - liveIds.add(run.id); - const node = - run.session?.getContextTree() ?? loadContextTreeChildFromDisk(run.sessionDir, resolveContextWindow); - children.push({ - ...(node ?? { - ownUsage: emptyUsage(), - totalUsage: emptyUsage(), - children: [], - }), - id: run.id, - label: rlmChildLabel(run.prompt), - status: run.status, - }); - } - children.push(...loadContextTreeChildrenFromDisk(this._rlmSessionDirForReading(), resolveContextWindow, liveIds)); - - const model = this.model; - return { - id: "root", - label: this.sessionName ?? "main agent", - status: "active", - model: model ? { provider: model.provider, id: model.id } : undefined, - ownUsage, - totalUsage, - contextUsage: this.getContextUsage(), - children, - }; + getContextTree( + ...args: Parameters + ): ReturnType { + return this._contextView.getContextTree(...args); } /** @@ -6659,20 +2889,8 @@ export class AgentSession { * @param outputPath Optional output path (defaults to session directory) * @returns Path to exported file */ - async exportToHtml(outputPath?: string): Promise { - const themeName = this.settingsManager.getTheme(); - - const toolRenderer: ToolHtmlRenderer = createToolHtmlRenderer({ - getToolDefinition: (name) => this.getToolDefinition(name), - theme, - cwd: this.sessionManager.getCwd(), - }); - - return await exportSessionToHtml(this.sessionManager, this.state, { - outputPath, - themeName, - toolRenderer, - }); + exportToHtml(...args: Parameters): ReturnType { + return this._export.exportToHtml(...args); } /** @@ -6681,34 +2899,8 @@ export class AgentSession { * @param outputPath Target file path. If omitted, generates a timestamped file in cwd. * @returns The resolved output file path. */ - exportToJsonl(outputPath?: string): string { - const filePath = resolve(outputPath ?? `session-${new Date().toISOString().replace(/[:.]/g, "-")}.jsonl`); - const dir = dirname(filePath); - if (!existsSync(dir)) { - mkdirSync(dir, { recursive: true }); - } - - const header: SessionHeader = { - type: "session", - version: CURRENT_SESSION_VERSION, - id: this.sessionManager.getSessionId(), - timestamp: new Date().toISOString(), - cwd: this.sessionManager.getCwd(), - }; - - const branchEntries = this.sessionManager.getBranch(); - const lines = [JSON.stringify(header)]; - - // Re-chain parentIds to form a linear sequence - let prevId: string | null = null; - for (const entry of branchEntries) { - const linear = { ...entry, parentId: prevId }; - lines.push(JSON.stringify(linear)); - prevId = entry.id; - } - - writeFileSync(filePath, `${lines.join("\n")}\n`); - return filePath; + exportToJsonl(...args: Parameters): ReturnType { + return this._export.exportToJsonl(...args); } /** @@ -6716,28 +2908,10 @@ export class AgentSession { * Useful for /copy command. * @returns Text content, or undefined if no assistant message exists */ - getLastAssistantText(): string | undefined { - const lastAssistant = this.messages - .slice() - .reverse() - .find((m) => { - if (m.role !== "assistant") return false; - const msg = m as AssistantMessage; - // Skip aborted messages with no content - if (msg.stopReason === "aborted" && msg.content.length === 0) return false; - return true; - }); - - if (!lastAssistant) return undefined; - - let text = ""; - for (const content of (lastAssistant as AssistantMessage).content) { - if (content.type === "text") { - text += content.text; - } - } - - return text.trim() || undefined; + getLastAssistantText( + ...args: Parameters + ): ReturnType { + return this._contextView.getLastAssistantText(...args); } // ================================================================== // Extension System @@ -6760,24 +2934,3 @@ export class AgentSession { return this._extensionRunner; } } - -function isRlmHeartbeatStatusUpdate(value: unknown): value is AgentRlmHeartbeatStatusUpdate { - return value === "pause" || value === "resume"; -} - -function rlmHeartbeatHostResponse(job: AgentCronJob): Record { - return { - id: job.id, - status: job.status, - label: job.label ?? null, - delivery_mode: job.deliveryMode ?? "steer", - instruction: job.prompt, - schedule: job.schedule, - created_at: job.createdAt, - updated_at: job.updatedAt, - next_run_at: job.nextRunAt ?? null, - last_run_at: job.lastRunAt ?? null, - last_error: job.lastError ?? null, - run_count: job.runCount, - }; -} diff --git a/packages/coding-agent/src/modes/agent-connection/types.ts b/packages/coding-agent/src/modes/agent-connection/types.ts index 0ae4094b37..f075c28419 100644 --- a/packages/coding-agent/src/modes/agent-connection/types.ts +++ b/packages/coding-agent/src/modes/agent-connection/types.ts @@ -14,7 +14,6 @@ import type { } from "../../core/cron-jobs.js"; import type { ReplayBuiltInToolName } from "../../core/extensions/index.js"; import type { InputSource } from "../../core/extensions/types.js"; -import type { GoalState } from "../../core/goals.js"; import type { KernelSentAgentMessage } from "../../core/kernel/index.js"; import type { AcpMcpServerConfig } from "../../core/mcp/acp-mcp-types.js"; import type { RefinementResult } from "../../core/refinement/index.js"; @@ -28,6 +27,7 @@ import type { import type { DeleteSessionFileResult } from "../../core/session-file-actions.js"; import type { SessionStats } from "../../core/session-stats.js"; import type { SessionUsageSummary } from "../../core/usage.js"; +import type { GoalState } from "../../session/goals/contracts.js"; import type { SessionSummary } from "../daemon/daemon-session-list.js"; /** diff --git a/packages/coding-agent/src/modes/interactive/components/injected-prompt-message.ts b/packages/coding-agent/src/modes/interactive/components/injected-prompt-message.ts index a639f930b3..f1f96d1ae6 100644 --- a/packages/coding-agent/src/modes/interactive/components/injected-prompt-message.ts +++ b/packages/coding-agent/src/modes/interactive/components/injected-prompt-message.ts @@ -8,7 +8,6 @@ import { truncateToWidth, visibleWidth, } from "@earendil-works/pi-tui"; -import { GOAL_CONTEXT_CUSTOM_TYPE, type GoalContextDetails } from "../../../core/goals.js"; import { ASYNC_BASH_COMPLETION_CUSTOM_TYPE, type AsyncBashCompletionDetails, @@ -22,6 +21,7 @@ import { type RlmChildFailureDetails, type RlmChildTerminalNoticeDetails, } from "../../../core/messages.js"; +import { GOAL_CONTEXT_CUSTOM_TYPE, type GoalContextDetails } from "../../../session/goals/contracts.js"; import { getMarkdownTheme, theme } from "../theme/theme.js"; import { expandCollapseHint } from "./keybinding-hints.js"; import { ShellCompletionComponent } from "./shell-completion.js"; diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 69a6a22366..86652ea322 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -99,7 +99,6 @@ import type { ExtensionWidgetOptions, } from "../../core/extensions/index.js"; import { FooterDataProvider, type ReadonlyFooterDataProvider } from "../../core/footer-data-provider.js"; -import { emptyGoalState, formatGoalUsage, GOAL_CONTEXT_PREVIEW_LABEL, type GoalState } from "../../core/goals.js"; import type { KernelSentAgentMessage } from "../../core/kernel/index.js"; import { type AppKeybinding, KeybindingsManager } from "../../core/keybindings.js"; import { runMcpManagementCommand } from "../../core/mcp/mcp-command.js"; @@ -140,6 +139,12 @@ import { type TelemetryOnboardingOutcome, } from "../../core/telemetry.js"; import { type TruncationResult, truncateTail } from "../../core/tools/truncate.js"; +import { + emptyGoalState, + formatGoalUsage, + GOAL_CONTEXT_PREVIEW_LABEL, + type GoalState, +} from "../../session/goals/contracts.js"; import { PRIME_COMPACT_BUTTERFLY_LOGO } from "../../themes/prime-logo.js"; import { getChangelogPath, parseChangelog } from "../../utils/changelog.js"; import { spawnHidden, spawnSyncHidden } from "../../utils/child-process.js"; diff --git a/packages/coding-agent/src/modes/rpc/rpc-types.ts b/packages/coding-agent/src/modes/rpc/rpc-types.ts index 7cba4f8a9b..f7119fc508 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-types.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-types.ts @@ -16,10 +16,10 @@ import type { AgentHeartbeatManagementAction, AgentHeartbeatUpdateAction, } from "../../core/cron-jobs.js"; -import type { GoalState } from "../../core/goals.js"; import type { RefinementResult } from "../../core/refinement/index.js"; import type { SessionActionSnapshot } from "../../core/session-action-store.js"; import type { SessionStats } from "../../core/session-stats.js"; +import type { GoalState } from "../../session/goals/contracts.js"; import type { AgentConnectionHeartbeat, AgentConnectionSourceInfo } from "../agent-connection/types.js"; // ============================================================================ diff --git a/packages/coding-agent/src/session/child-projection.ts b/packages/coding-agent/src/session/children/child-projection.ts similarity index 97% rename from packages/coding-agent/src/session/child-projection.ts rename to packages/coding-agent/src/session/children/child-projection.ts index c59c429ebb..067ebaf99a 100644 --- a/packages/coding-agent/src/session/child-projection.ts +++ b/packages/coding-agent/src/session/children/child-projection.ts @@ -1,6 +1,6 @@ -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 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, diff --git a/packages/coding-agent/src/session/child-run.ts b/packages/coding-agent/src/session/children/child-run.ts similarity index 98% rename from packages/coding-agent/src/session/child-run.ts rename to packages/coding-agent/src/session/children/child-run.ts index 8f1c6c3c13..4bfdd30dfb 100644 --- a/packages/coding-agent/src/session/child-run.ts +++ b/packages/coding-agent/src/session/children/child-run.ts @@ -1,18 +1,18 @@ 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 { 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"; +} from "../../core/messages.js"; import type { CreateRlmSubagentRuntimeOptions, RlmSpawnHandle, RlmSubagentRegistryEntry, RlmSubagentRuntime, SubagentRuntimeHost, -} from "../core/rlm-runtime.js"; +} from "../../core/rlm-runtime.js"; import { compactRlmText, createChildDeferred, diff --git a/packages/coding-agent/src/session/child-runtime.ts b/packages/coding-agent/src/session/children/child-runtime.ts similarity index 90% rename from packages/coding-agent/src/session/child-runtime.ts rename to packages/coding-agent/src/session/children/child-runtime.ts index 9f825c0b6c..e9048b5cd3 100644 --- a/packages/coding-agent/src/session/child-runtime.ts +++ b/packages/coding-agent/src/session/children/child-runtime.ts @@ -2,12 +2,12 @@ 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"; +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; diff --git a/packages/coding-agent/src/session/child-state.ts b/packages/coding-agent/src/session/children/child-state.ts similarity index 96% rename from packages/coding-agent/src/session/child-state.ts rename to packages/coding-agent/src/session/children/child-state.ts index 19b6caf20f..034eb55bc1 100644 --- a/packages/coding-agent/src/session/child-state.ts +++ b/packages/coding-agent/src/session/children/child-state.ts @@ -1,6 +1,6 @@ -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"; +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; diff --git a/packages/coding-agent/src/session/child-types.ts b/packages/coding-agent/src/session/children/child-types.ts similarity index 96% rename from packages/coding-agent/src/session/child-types.ts rename to packages/coding-agent/src/session/children/child-types.ts index 69e61e2b39..0a645bbef5 100644 --- a/packages/coding-agent/src/session/child-types.ts +++ b/packages/coding-agent/src/session/children/child-types.ts @@ -1,6 +1,6 @@ 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"; +import type { AgentSession } from "../../core/agent-session.js"; +import type { RlmSubagentRegistryEntry } from "../../core/rlm-runtime.js"; export type RlmChildAgentStatus = "queued" | "running" | "done" | "error" | "cancelled"; diff --git a/packages/coding-agent/src/session/child-usage.ts b/packages/coding-agent/src/session/children/child-usage.ts similarity index 97% rename from packages/coding-agent/src/session/child-usage.ts rename to packages/coding-agent/src/session/children/child-usage.ts index 3bf7f29c13..f33f120bf2 100644 --- a/packages/coding-agent/src/session/child-usage.ts +++ b/packages/coding-agent/src/session/children/child-usage.ts @@ -1,13 +1,13 @@ 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 { 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"; +} from "../../core/session-manager.js"; +import { addAssistantUsage, cloneUsage, emptyUsage, subtractAssistantUsage } from "../../core/usage.js"; export interface ChildUsageHost { sessionManager: Pick; diff --git a/packages/coding-agent/src/session/children.ts b/packages/coding-agent/src/session/children/children.ts similarity index 98% rename from packages/coding-agent/src/session/children.ts rename to packages/coding-agent/src/session/children/children.ts index 1162810e0c..a73ae6be96 100644 --- a/packages/coding-agent/src/session/children.ts +++ b/packages/coding-agent/src/session/children/children.ts @@ -8,9 +8,9 @@ import { assertAgentSessionNameAvailable, assertDirectAgentMessageTarget, formatAgentSessionNameUnavailable, -} from "../core/agent-messages.js"; -import type { AgentSession, AgentSessionEvent } from "../core/agent-session.js"; -import type { CustomMessage } from "../core/messages.js"; +} 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, @@ -24,8 +24,8 @@ import { type RlmSubagentRegistryEntry, type RlmSubagentRuntime, type SubagentRuntimeHost, -} from "../core/rlm-runtime.js"; -import type { SemanticEdgeRecorder } from "../core/semantic-edges.js"; +} 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 { @@ -1016,7 +1016,11 @@ export class SessionChildren { getRuntimeHost(): SubagentRuntimeHost | undefined { return this._subagentRuntimeHost; } - getActiveRuns(): Iterable> { + getActiveRuns(): Iterable< + Readonly> & { + readonly session?: Pick; + } + > { return this._activeRlmChildRuns.values(); } } diff --git a/packages/coding-agent/src/session/compaction-execution.ts b/packages/coding-agent/src/session/compaction/compaction-execution.ts similarity index 95% rename from packages/coding-agent/src/session/compaction-execution.ts rename to packages/coding-agent/src/session/compaction/compaction-execution.ts index 3a8aeb4df8..818825d303 100644 --- a/packages/coding-agent/src/session/compaction-execution.ts +++ b/packages/coding-agent/src/session/compaction/compaction-execution.ts @@ -5,11 +5,11 @@ import { type CompactionSettings, compact, prepareCompaction, -} from "../core/compaction/index.js"; -import type { ExtensionRunner, SessionBeforeCompactResult } from "../core/extensions/index.js"; -import type { ProviderRetryPolicy } from "../core/provider-retry.js"; -import { modelRequestHeaders, type SemanticEdgeRecorder } from "../core/semantic-edges.js"; -import type { CompactionEntry, SessionManager } from "../core/session-manager.js"; +} from "../../core/compaction/index.js"; +import type { ExtensionRunner, SessionBeforeCompactResult } from "../../core/extensions/index.js"; +import type { ProviderRetryPolicy } from "../../core/provider-retry.js"; +import { modelRequestHeaders, type SemanticEdgeRecorder } from "../../core/semantic-edges.js"; +import type { CompactionEntry, SessionManager } from "../../core/session-manager.js"; export class CompactionSkippedError extends Error {} diff --git a/packages/coding-agent/src/session/compaction.ts b/packages/coding-agent/src/session/compaction/compaction.ts similarity index 90% rename from packages/coding-agent/src/session/compaction.ts rename to packages/coding-agent/src/session/compaction/compaction.ts index 3f048df354..cb4dfc916b 100644 --- a/packages/coding-agent/src/session/compaction.ts +++ b/packages/coding-agent/src/session/compaction/compaction.ts @@ -1,21 +1,23 @@ import type { AgentEvent, AgentMessage } from "@earendil-works/pi-agent-core"; import { type Api, type AssistantMessage, isContextOverflow, type Model } from "@earendil-works/pi-ai"; -import { formatNoModelSelectedMessage } from "../core/auth-guidance.js"; +import { formatNoModelSelectedMessage } from "../../core/auth-guidance.js"; import { type CompactionResult, type CompactionSettings, calculateContextTokens, estimateContextTokens, + prepareCompaction, shouldCompact, -} from "../core/compaction/index.js"; +} from "../../core/compaction/index.js"; +import type { ContextUsage } from "../../core/extensions/index.js"; import { type CompactionOutcome, type CompactionOutcomeReason, type CustomMessage, createCompactionOutcomeMessage, -} from "../core/messages.js"; -import type { ModelRegistry } from "../core/model-registry.js"; -import { getLatestCompactionEntry, type SessionManager } from "../core/session-manager.js"; +} from "../../core/messages.js"; +import type { ModelRegistry } from "../../core/model-registry.js"; +import { getLatestCompactionEntry, type SessionManager } from "../../core/session-manager.js"; import { type CompactionExecutionOptions, CompactionSkippedError } from "./compaction-execution.js"; export type CompactionReason = "manual" | "threshold" | "overflow" | "requested"; @@ -33,6 +35,8 @@ export type SessionCompactionEvent = }; export interface SessionCompactionHost { + includesCompactSkill(): boolean; + getContextUsage(): ContextUsage | undefined; getSettings(): CompactionSettings; runAutomatic(reason: "overflow" | "threshold" | "requested", willRetry: boolean): Promise; queueGoalContinuation(message: AssistantMessage): boolean; @@ -77,6 +81,50 @@ export class SessionCompaction { constructor(private readonly host: SessionCompactionHost) {} + handleCompactHostRequest(type: string, payload: Record = {}): Record { + if (!this.host.includesCompactSkill()) { + throw new Error("the compact skill is disabled in this session"); + } + switch (type) { + case "compact.status": { + const usage = this.host.getContextUsage(); + return { + tokens: usage?.tokens ?? null, + context_window: usage?.contextWindow ?? null, + percent: usage?.percent ?? null, + scheduled: this.hasPendingRequest, + }; + } + case "compact.run": { + const instructions = payload.instructions; + if (instructions !== undefined && typeof instructions !== "string") { + throw new Error("compact.run instructions must be a string when provided"); + } + if (!this.host.isStreaming()) { + return { + scheduled: false, + reason: "no active turn; compaction can only be requested while a turn is running", + }; + } + const preparation = prepareCompaction(this.host.getSessionStore().getBranch(), this.host.getSettings()); + if (!preparation) { + const lastEntry = this.host.getSessionStore().getBranch().at(-1); + return { + scheduled: false, + reason: lastEntry?.type === "compaction" ? "already compacted" : "session is too short to compact", + }; + } + this.request(instructions); + return { + scheduled: true, + note: "Compaction runs when the current turn ends; you resume automatically afterwards. Continue working normally.", + }; + } + default: + throw new Error(`unknown compact request type "${type}"`); + } + } + private get model(): Model | undefined { return this.host.getModel(); } diff --git a/packages/coding-agent/src/session/context/context-view.ts b/packages/coding-agent/src/session/context/context-view.ts new file mode 100644 index 0000000000..e30311417d --- /dev/null +++ b/packages/coding-agent/src/session/context/context-view.ts @@ -0,0 +1,212 @@ +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import type { Api, AssistantMessage, Model, Usage } from "@earendil-works/pi-ai"; +import { calculateContextTokens, estimateContextTokens } from "../../core/compaction/index.js"; +import { + type ContextTreeNode, + type ContextWindowResolver, + computeOwnAndTotalUsage, + loadContextTreeChildFromDisk, + loadContextTreeChildrenFromDisk, +} from "../../core/context-tree.js"; +import type { ContextUsage } from "../../core/extensions/index.js"; +import { getLatestCompactionEntry, type SessionEntry, type SessionManager } from "../../core/session-manager.js"; +import type { SessionStats } from "../../core/session-stats.js"; +import { emptyUsage, type SessionUsageSummary, sessionUsageSummaryFrom } from "../../core/usage.js"; + +export interface ContextViewChild { + id: string; + label: string; + status: ContextTreeNode["status"]; + sessionDir: string; + getContextTree?: () => ContextTreeNode; +} +export interface ContextViewHost { + sessionManager: Pick< + SessionManager, + "getEntries" | "getBranch" | "getSessionId" | "getSessionFile" | "getSessionName" + >; + getMessages(): AgentMessage[]; + getContextUsage(): ContextUsage | undefined; + getModel(): Model | undefined; + findModel(provider: string, modelId: string): Model | undefined; + subtractUnindexedChildUsage(ownUsage: Usage, entries: SessionEntry[]): void; + getLiveChildren(): Iterable; + getRlmSessionDir(): string | undefined; +} +export class SessionContextView { + private _ownUsageMemo?: { count: number; tailId: string | undefined; usage: SessionUsageSummary | undefined }; + constructor(private readonly host: ContextViewHost) {} + invalidateOwnUsage(): void { + this._ownUsageMemo = undefined; + } + getSessionStats(): SessionStats { + const state = { messages: this.host.getMessages() }; + const userMessages = state.messages.filter((m) => m.role === "user").length; + const assistantMessages = state.messages.filter((m) => m.role === "assistant").length; + const toolResults = state.messages.filter((m) => m.role === "toolResult").length; + + let toolCalls = 0; + let totalInput = 0; + let totalOutput = 0; + let totalCacheRead = 0; + let totalCacheWrite = 0; + let totalCost = 0; + + for (const message of state.messages) { + if (message.role === "assistant") { + const assistantMsg = message as AssistantMessage; + toolCalls += assistantMsg.content.filter((c) => c.type === "toolCall").length; + totalInput += assistantMsg.usage.input; + totalOutput += assistantMsg.usage.output; + totalCacheRead += assistantMsg.usage.cacheRead; + totalCacheWrite += assistantMsg.usage.cacheWrite; + totalCost += assistantMsg.usage.cost.total; + } + } + + return { + sessionFile: this.host.sessionManager.getSessionFile(), + sessionId: this.host.sessionManager.getSessionId(), + userMessages, + assistantMessages, + toolCalls, + toolResults, + totalMessages: state.messages.length, + tokens: { + input: totalInput, + output: totalOutput, + cacheRead: totalCacheRead, + cacheWrite: totalCacheWrite, + total: totalInput + totalOutput + totalCacheRead + totalCacheWrite, + }, + cost: totalCost, + contextUsage: this.host.getContextUsage(), + }; + } + + getContextUsage(): ContextUsage | undefined { + const model = this.host.getModel(); + if (!model) return undefined; + + const contextWindow = model.contextWindow ?? 0; + if (contextWindow <= 0) return undefined; + + // After compaction, the last assistant usage reflects pre-compaction context size. + // We can only trust usage from an assistant that responded after the latest compaction. + // If no such assistant exists, context token count is unknown until the next LLM response. + const branchEntries = this.host.sessionManager.getBranch(); + const latestCompaction = getLatestCompactionEntry(branchEntries); + + if (latestCompaction) { + // Check if there's a valid assistant usage after the compaction boundary + const compactionIndex = branchEntries.lastIndexOf(latestCompaction); + let hasPostCompactionUsage = false; + for (let i = branchEntries.length - 1; i > compactionIndex; i--) { + const entry = branchEntries[i]; + if (entry.type === "message" && entry.message.role === "assistant") { + const assistant = entry.message; + if (assistant.stopReason !== "aborted" && assistant.stopReason !== "error") { + const contextTokens = calculateContextTokens(assistant.usage); + if (contextTokens > 0) { + hasPostCompactionUsage = true; + } + break; + } + } + } + + if (!hasPostCompactionUsage) { + return { tokens: null, contextWindow, percent: null }; + } + } + + const estimate = estimateContextTokens(this.host.getMessages()); + const percent = (estimate.tokens / contextWindow) * 100; + + return { + tokens: estimate.tokens, + contextWindow, + percent, + }; + } + + private _contextWindowResolver(): ContextWindowResolver { + return (provider, modelId) => this.host.findModel(provider, modelId)?.contextWindow; + } + + getOwnUsageSummary(): SessionUsageSummary | undefined { + const entries = this.host.sessionManager.getEntries(); + const tailId = entries.at(-1)?.id; + const memo = this._ownUsageMemo; + if (memo && memo.count === entries.length && memo.tailId === tailId) { + return memo.usage; + } + const { ownUsage } = computeOwnAndTotalUsage(entries, entries); + this.host.subtractUnindexedChildUsage(ownUsage, entries); + const usage = sessionUsageSummaryFrom(ownUsage); + this._ownUsageMemo = { count: entries.length, tailId, usage }; + return usage; + } + + getContextTree(): ContextTreeNode { + const resolveContextWindow = this._contextWindowResolver(); + const branch = this.host.sessionManager.getBranch(); + const { ownUsage, totalUsage } = computeOwnAndTotalUsage(branch, this.host.sessionManager.getEntries()); + this.host.subtractUnindexedChildUsage(ownUsage, branch); + + const children: ContextTreeNode[] = []; + const liveIds = new Set(); + for (const run of this.host.getLiveChildren()) { + liveIds.add(run.id); + const node = run.getContextTree?.() ?? loadContextTreeChildFromDisk(run.sessionDir, resolveContextWindow); + children.push({ + ...(node ?? { + ownUsage: emptyUsage(), + totalUsage: emptyUsage(), + children: [], + }), + id: run.id, + label: run.label, + status: run.status, + }); + } + children.push(...loadContextTreeChildrenFromDisk(this.host.getRlmSessionDir(), resolveContextWindow, liveIds)); + + const model = this.host.getModel(); + return { + id: "root", + label: this.host.sessionManager.getSessionName() ?? "main agent", + status: "active", + model: model ? { provider: model.provider, id: model.id } : undefined, + ownUsage, + totalUsage, + contextUsage: this.host.getContextUsage(), + children, + }; + } + + getLastAssistantText(): string | undefined { + const lastAssistant = this.host + .getMessages() + .slice() + .reverse() + .find((m) => { + if (m.role !== "assistant") return false; + const msg = m as AssistantMessage; + // Skip aborted messages with no content + if (msg.stopReason === "aborted" && msg.content.length === 0) return false; + return true; + }); + + if (!lastAssistant) return undefined; + + let text = ""; + for (const content of (lastAssistant as AssistantMessage).content) { + if (content.type === "text") { + text += content.text; + } + } + + return text.trim() || undefined; + } +} diff --git a/packages/coding-agent/src/session/context/export.ts b/packages/coding-agent/src/session/context/export.ts new file mode 100644 index 0000000000..252d198967 --- /dev/null +++ b/packages/coding-agent/src/session/context/export.ts @@ -0,0 +1,62 @@ +import { existsSync, mkdirSync, writeFileSync } from "node:fs"; +import { dirname, resolve } from "node:path"; +import type { AgentState } from "@earendil-works/pi-agent-core"; +import { exportSessionToHtml, type ToolHtmlRenderer } from "../../core/export-html/index.js"; +import { createToolHtmlRenderer } from "../../core/export-html/tool-renderer.js"; +import type { ToolDefinition } from "../../core/extensions/index.js"; +import { CURRENT_SESSION_VERSION, type SessionHeader, type SessionManager } from "../../core/session-manager.js"; +import { theme } from "../../modes/interactive/theme/theme.js"; +export interface SessionExportHost { + sessionManager: SessionManager; + getState(): AgentState; + getTheme(): string | undefined; + getToolDefinition(name: string): ToolDefinition | undefined; +} +export class SessionExport { + constructor(private readonly host: SessionExportHost) {} + async exportToHtml(outputPath?: string): Promise { + const themeName = this.host.getTheme(); + + const toolRenderer: ToolHtmlRenderer = createToolHtmlRenderer({ + getToolDefinition: (name) => this.host.getToolDefinition(name), + theme, + cwd: this.host.sessionManager.getCwd(), + }); + + return await exportSessionToHtml(this.host.sessionManager, this.host.getState(), { + outputPath, + themeName, + toolRenderer, + }); + } + + exportToJsonl(outputPath?: string): string { + const filePath = resolve(outputPath ?? `session-${new Date().toISOString().replace(/[:.]/g, "-")}.jsonl`); + const dir = dirname(filePath); + if (!existsSync(dir)) { + mkdirSync(dir, { recursive: true }); + } + + const header: SessionHeader = { + type: "session", + version: CURRENT_SESSION_VERSION, + id: this.host.sessionManager.getSessionId(), + timestamp: new Date().toISOString(), + cwd: this.host.sessionManager.getCwd(), + }; + + const branchEntries = this.host.sessionManager.getBranch(); + const lines = [JSON.stringify(header)]; + + // Re-chain parentIds to form a linear sequence + let prevId: string | null = null; + for (const entry of branchEntries) { + const linear = { ...entry, parentId: prevId }; + lines.push(JSON.stringify(linear)); + prevId = entry.id; + } + + writeFileSync(filePath, `${lines.join("\n")}\n`); + return filePath; + } +} diff --git a/packages/coding-agent/src/session/context/harness-context.ts b/packages/coding-agent/src/session/context/harness-context.ts new file mode 100644 index 0000000000..5a9995436b --- /dev/null +++ b/packages/coding-agent/src/session/context/harness-context.ts @@ -0,0 +1,114 @@ +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import { + type CustomMessage, + createHarnessDigestMessage, + HARNESS_DIGEST_CUSTOM_TYPE, + type HarnessDigestDetails, +} from "../../core/messages.js"; +import { formatHarnessStateForPrompt, type HarnessState, REFINE_SKILL_NAME } from "../../core/refinement/index.js"; +import type { SessionContext, SessionManager } from "../../core/session-manager.js"; +import type { Skill } from "../../core/skills.js"; +export interface HarnessContextHost { + sessionManager: Pick; + getMessages(): AgentMessage[]; + getActiveToolNames(): string[]; + getVisibleSkills(): Pick[]; + loadHarnessState(): HarnessState; + applyLateSentMessages(message: AgentMessage): void; +} +export class SessionHarnessContext { + private _harnessDigestPending = false; + /** Disclosures retained when their session append failed. */ + private readonly _unpersistedOutcomes: CustomMessage[] = []; + constructor(private readonly host: HarnessContextHost) {} + get digestPending(): boolean { + return this._harnessDigestPending; + } + consumePendingDigest(): boolean { + const pending = this._harnessDigestPending; + this._harnessDigestPending = false; + return pending; + } + + rearmDigest(): void { + this._harnessDigestPending = true; + } + retainOutcome(message: CustomMessage): void { + this._unpersistedOutcomes.push(message); + } + harnessDigest(): string { + const tools = this.host.getActiveToolNames(); + const hasIpython = tools.includes("ipython"); + const visibleSkills = this.host.getVisibleSkills().filter((skill) => !skill.disableModelInvocation); + const hasRefineSkill = visibleSkills.some((skill) => skill.name === REFINE_SKILL_NAME); + return formatHarnessStateForPrompt(this.host.loadHarnessState(), { + includeIpythonExamples: hasIpython, + includeShellExamples: tools.includes("bash"), + includeRefineExamples: hasIpython && hasRefineSkill, + }); + } + + ensureHarnessDigestContext(): void { + if (this.host.getMessages().length === 0) { + this._harnessDigestPending = true; + return; + } + this._harnessDigestPending = false; + this.appendHarnessDigestIfStale(); + } + + private appendHarnessDigestIfStale(): void { + const digest = this.harnessDigest(); + if (this.latestContextHarnessDigest() === digest) return; + const message = createHarnessDigestMessage(digest); + try { + this.host.sessionManager.appendCustomMessageEntryWithRollback( + message.customType, + message.content, + message.display, + message.details, + ); + } catch { + // Unpersisted session: context-only injection. + } + this.host.getMessages().push(message); + } + + latestContextHarnessDigest(): string | undefined { + // Retained pre-compaction messages follow the compaction head, so recency is by timestamp, not position. + let latest: { timestamp: number; digest: string } | undefined; + for (const message of this.host.getMessages()) { + let digest: string | undefined; + if (message.role === "custom" && message.customType === HARNESS_DIGEST_CUSTOM_TYPE) { + digest = (message.details as HarnessDigestDetails | undefined)?.digest; + } else if (message.role === "compactionSummary") { + digest = message.harnessDigest; + } else { + continue; + } + if (digest !== undefined && (!latest || message.timestamp >= latest.timestamp)) { + latest = { timestamp: message.timestamp, digest }; + } + } + return latest?.digest; + } + + buildSessionContext(): SessionContext { + const context = this.host.sessionManager.buildSessionContext(); + for (const message of context.messages) { + this.host.applyLateSentMessages(message); + } + this.mergeUnpersistedOutcomes(context.messages); + return context; + } + + mergeUnpersistedOutcomes(messages: AgentMessage[]): void { + for (const outcome of this._unpersistedOutcomes) { + let insertAt = messages.length; + while (insertAt > 0 && messages[insertAt - 1]!.timestamp > outcome.timestamp) { + insertAt -= 1; + } + messages.splice(insertAt, 0, outcome); + } + } +} diff --git a/packages/coding-agent/src/session/context/history-navigation.ts b/packages/coding-agent/src/session/context/history-navigation.ts new file mode 100644 index 0000000000..ee75bcf701 --- /dev/null +++ b/packages/coding-agent/src/session/context/history-navigation.ts @@ -0,0 +1,319 @@ +import type { Api, Model, Usage } from "@earendil-works/pi-ai"; +import { collectEntriesForBranchSummary, generateBranchSummary } from "../../core/compaction/index.js"; +import type { ExtensionRunner, SessionBeforeTreeResult, TreePreparation } from "../../core/extensions/index.js"; +import type { ProviderRetryPolicy } from "../../core/provider-retry.js"; +import type { BranchSummaryEntry, SessionManager } from "../../core/session-manager.js"; +import type { SettingsManager } from "../../core/settings-manager.js"; +import type { SessionCommitLease } from "../input/commit-fence.js"; + +export interface HistoryNavigationHost { + getSessionId?(): string; + sessionManager: SessionManager; + settingsManager: Pick; + getRetryPolicy(): ProviderRetryPolicy; + getModel(): Model | undefined; + getExtensions(): Pick; + getRequiredAuth( + model: Model, + ): Promise<{ apiKey: string; headers?: Record; requestModel?: Model }>; + acquireQueuedWorkPause(): { release(): void }; + acquireCommitFence(): Promise; + runWithCommitFence(lease: SessionCommitLease, run: () => T): T; + waitForAgentIdle(): Promise; + getEventQueue(): Promise; + invalidateRefinement(): Promise; + rebuildBranchContext(): void; + notifyCheckpoints(): void; +} +export class SessionHistoryNavigation { + private _branchNavigationQueue: Promise = Promise.resolve(); + private _branchSummaryAbortController?: AbortController; + private _branchSummaryOperation?: Promise; + constructor(private readonly host: HistoryNavigationHost) {} + get operation(): Promise | undefined { + return this._branchSummaryOperation; + } + get isSummarizing(): boolean { + return this._branchSummaryAbortController !== undefined; + } + abortBranchSummary(): void { + this._branchSummaryAbortController?.abort(); + } + async navigateTree( + targetId: string, + options: { + summarize?: boolean; + customInstructions?: string; + replaceInstructions?: boolean; + label?: string; + } = {}, + ): Promise<{ + editorText?: string; + cancelled: boolean; + aborted?: boolean; + summaryEntry?: BranchSummaryEntry; + }> { + const previous = this._branchNavigationQueue; + let release = () => {}; + this._branchNavigationQueue = new Promise((resolve) => { + release = resolve; + }); + await previous; + try { + return await this._navigateTree(targetId, options); + } finally { + release(); + } + } + + private async _navigateTree( + targetId: string, + options: { + summarize?: boolean; + customInstructions?: string; + replaceInstructions?: boolean; + label?: string; + } = {}, + ): Promise<{ + editorText?: string; + cancelled: boolean; + aborted?: boolean; + summaryEntry?: BranchSummaryEntry; + }> { + if (options.summarize && !this.host.getModel()) { + throw new Error("No model available for summarization"); + } + + const targetEntry = this.host.sessionManager.getEntry(targetId); + if (!targetEntry) { + throw new Error(`Entry ${targetId} not found`); + } + + const queuedWorkPause = this.host.acquireQueuedWorkPause(); + let commitFence: { owner: symbol; release(): void } | undefined; + try { + // Branch navigation and turn dispatch mutate the same transcript leaf. + commitFence = await this.host.acquireCommitFence(); + return await this.host.runWithCommitFence(commitFence, async () => { + await this.host.waitForAgentIdle(); + await this.host.getEventQueue(); + return this._navigateTreeUnderPause(targetId, targetEntry, options); + }); + } finally { + queuedWorkPause.release(); + commitFence?.release(); + } + } + + private async _navigateTreeUnderPause( + targetId: string, + targetEntry: NonNullable>, + options: { + summarize?: boolean; + customInstructions?: string; + replaceInstructions?: boolean; + label?: string; + }, + ): Promise<{ + editorText?: string; + cancelled: boolean; + aborted?: boolean; + summaryEntry?: BranchSummaryEntry; + }> { + const oldLeafId = this.host.sessionManager.getLeafId(); + + // No-op if already at target after admitted work has settled. + if (targetId === oldLeafId) { + return { cancelled: false }; + } + + // Do not switch branches while /refine has detached event handling and is + // about to persist harness/session entries for the current branch. + await this.host.invalidateRefinement(); + + const { entries: entriesToSummarize, commonAncestorId } = collectEntriesForBranchSummary( + this.host.sessionManager, + oldLeafId, + targetId, + ); + + let customInstructions = options.customInstructions; + let replaceInstructions = options.replaceInstructions; + let label = options.label; + + const preparation: TreePreparation = { + targetId, + oldLeafId, + commonAncestorId, + entriesToSummarize, + userWantsSummary: options.summarize ?? false, + customInstructions, + replaceInstructions, + label, + }; + + this._branchSummaryAbortController = new AbortController(); + let resolveBranchSummaryOperation: () => void = () => {}; + const branchSummaryOperation = new Promise((resolve) => { + resolveBranchSummaryOperation = resolve; + }); + this._branchSummaryOperation = branchSummaryOperation; + + try { + let extensionSummary: { summary: string; details?: unknown } | undefined; + let fromExtension = false; + + if (this.host.getExtensions().hasHandlers("session_before_tree")) { + const result = (await this.host.getExtensions().emit({ + type: "session_before_tree", + preparation, + signal: this._branchSummaryAbortController.signal, + })) as SessionBeforeTreeResult | undefined; + + if (result?.cancel) { + return { cancelled: true }; + } + + if (result?.summary && options.summarize) { + extensionSummary = result.summary; + fromExtension = true; + } + + if (result?.customInstructions !== undefined) { + customInstructions = result.customInstructions; + } + if (result?.replaceInstructions !== undefined) { + replaceInstructions = result.replaceInstructions; + } + if (result?.label !== undefined) { + label = result.label; + } + } + + let summaryText: string | undefined; + let summaryDetails: unknown; + let summaryUsage: Usage | undefined; + if (options.summarize && entriesToSummarize.length > 0 && !extensionSummary) { + const model = this.host.getModel()!; + const { apiKey, headers, requestModel } = await this.host.getRequiredAuth(model); + const branchSummarySettings = this.host.settingsManager.getBranchSummarySettings(); + const result = await generateBranchSummary(entriesToSummarize, { + model: requestModel ?? model, + sessionId: this.host.getSessionId?.(), + apiKey, + headers, + signal: this._branchSummaryAbortController.signal, + customInstructions, + replaceInstructions, + reserveTokens: branchSummarySettings.reserveTokens, + retry: this.host.getRetryPolicy(), + }); + if (result.aborted) { + return { cancelled: true, aborted: true }; + } + if (result.error) { + throw new Error(result.error); + } + summaryText = result.summary; + summaryUsage = result.usage; + summaryDetails = { + readFiles: result.readFiles || [], + modifiedFiles: result.modifiedFiles || [], + }; + } else if (extensionSummary) { + summaryText = extensionSummary.summary; + summaryDetails = extensionSummary.details; + } + + let newLeafId: string | null; + let editorText: string | undefined; + + if (targetEntry.type === "message" && targetEntry.message.role === "user") { + newLeafId = targetEntry.parentId; + editorText = this._extractUserMessageText(targetEntry.message.content); + } else if (targetEntry.type === "custom_message") { + newLeafId = targetEntry.parentId; + editorText = + typeof targetEntry.content === "string" + ? targetEntry.content + : targetEntry.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + } else { + newLeafId = targetId; + } + + let summaryEntry: BranchSummaryEntry | undefined; + if (summaryText) { + const summaryId = this.host.sessionManager.branchWithSummary( + newLeafId, + summaryText, + summaryDetails, + fromExtension, + summaryUsage, + ); + summaryEntry = this.host.sessionManager.getEntry(summaryId) as BranchSummaryEntry; + + if (label) { + this.host.sessionManager.appendLabelChange(summaryId, label); + } + } else if (newLeafId === null) { + this.host.sessionManager.resetLeaf(); + } else { + this.host.sessionManager.branch(newLeafId); + } + + if (label && !summaryText) { + this.host.sessionManager.appendLabelChange(targetId, label); + } + + this.host.rebuildBranchContext(); + + await this.host.getExtensions().emit({ + type: "session_tree", + newLeafId: this.host.sessionManager.getLeafId(), + oldLeafId, + summaryEntry, + fromExtension: summaryText ? fromExtension : undefined, + }); + + return { editorText, cancelled: false, summaryEntry }; + } finally { + this._branchSummaryAbortController = undefined; + if (this._branchSummaryOperation === branchSummaryOperation) { + this._branchSummaryOperation = undefined; + } + resolveBranchSummaryOperation(); + this.host.notifyCheckpoints(); + } + } + + getUserMessagesForForking(): Array<{ entryId: string; text: string }> { + const entries = this.host.sessionManager.getEntries(); + const result: Array<{ entryId: string; text: string }> = []; + + for (const entry of entries) { + if (entry.type !== "message") continue; + if (entry.message.role !== "user") continue; + + const text = this._extractUserMessageText(entry.message.content); + if (text) { + result.push({ entryId: entry.id, text }); + } + } + + return result; + } + + private _extractUserMessageText(content: string | Array<{ type: string; text?: string }>): string { + if (typeof content === "string") return content; + if (Array.isArray(content)) { + return content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join(""); + } + return ""; + } +} diff --git a/packages/coding-agent/src/session/context/pending-context.ts b/packages/coding-agent/src/session/context/pending-context.ts new file mode 100644 index 0000000000..2676642a01 --- /dev/null +++ b/packages/coding-agent/src/session/context/pending-context.ts @@ -0,0 +1,231 @@ +import { + type CustomMessage, + RLM_CHILD_FAILURE_CUSTOM_TYPE, + RLM_CHILD_TERMINAL_NOTICE_CUSTOM_TYPE, +} from "../../core/messages.js"; +import type { ActionStore, DeliveryRecord } from "../../core/session-action-store.js"; +import { waitForPromiseOrAbort } from "../../utils/wait-for-abort.js"; +import type { SessionCommitFence, SessionCommitLease } from "../input/commit-fence.js"; +import type { SessionInputScheduler } from "../input/input-scheduler.js"; +import { + cloneCustomMessage, + createPreparedTurnAction, + primaryDeliveryRecord, + type QueuedSessionAction, +} from "../prepared-actions.js"; +import { createTurnExecutionPolicy } from "../turns/turn-preparation.js"; + +export interface SessionPendingContextHost { + getScheduler(): Pick; + getFence(): Pick; + isDisposed(): boolean; + isDisposing(): boolean; + admit(action: QueuedSessionAction, options: { wake: false }): { accepted: boolean }; + scheduleInput(): void; + addCheckpointWaiter(waiter: () => void): void; + removeCheckpointWaiter(waiter: () => void): void; + acquireFence(signal: AbortSignal): Promise; + cancelActions( + predicate: (action: QueuedSessionAction) => boolean, + error: Error, + candidates: QueuedSessionAction[], + ): QueuedSessionAction[]; +} +export class SessionPendingContext { + private messages: CustomMessage[] = []; + private readonly terminalNoticeActionIds = new Set(); + constructor( + private readonly actions: Pick, "clearableActions" | "unfinishedActions">, + private readonly host: SessionPendingContextHost, + ) {} + appendMessages(...messages: CustomMessage[]): void { + this.messages.push(...messages); + } + + prependMessages(messages: readonly CustomMessage[]): void { + this.messages.unshift(...messages); + } + + removeMessagesMatching(predicate: (message: CustomMessage) => boolean): void { + this.messages = this.messages.filter((message) => !predicate(message)); + } + + retainTerminalNotice(id: string): void { + this.terminalNoticeActionIds.add(id); + } + + releaseTerminalNotice(id: string): void { + this.terminalNoticeActionIds.delete(id); + } + + isRetainedTerminalNotice(id: string): boolean { + return this.terminalNoticeActionIds.has(id); + } + + dispose(): void { + this.messages = []; + } + + isRlmTerminalNotice(message: CustomMessage): boolean { + return ( + message.customType === RLM_CHILD_TERMINAL_NOTICE_CUSTOM_TYPE || + message.customType === RLM_CHILD_FAILURE_CUSTOM_TYPE + ); + } + + assertRlmTerminalNotice(message: CustomMessage): void { + if (!this.isRlmTerminalNotice(message)) { + throw new Error("Deferred terminal admission only accepts RLM child terminal notices."); + } + } + + isRlmTerminalNoticeAction(action: QueuedSessionAction): boolean { + if (action.payload.kind !== "turn") return false; + const message = primaryDeliveryRecord(action).message; + return message.role === "custom" && this.isRlmTerminalNotice(message); + } + + hasDeferredRlmTerminalNotices(): boolean { + return this.messages.some((message) => this.isRlmTerminalNotice(message)); + } + + enqueueRlmTerminalNoticeAction(message: CustomMessage): void { + this.assertRlmTerminalNotice(message); + const action = createPreparedTurnAction("followUp", message.content as string, undefined, { + message, + suppressAutonomousContinuation: true, + resumeIfIdle: false, + source: "internal", + executionPolicy: createTurnExecutionPolicy("injected"), + queueVisible: false, + }); + this.terminalNoticeActionIds.add(action.id); + try { + const result = this.host.admit(action, { wake: false }); + if (!result.accepted) throw new Error("RLM child terminal notice was not admitted."); + } catch (error) { + this.terminalNoticeActionIds.delete(action.id); + throw error; + } + } + + flushDeferredRlmTerminalNotices(): void { + if ( + this.host.getScheduler().admissionPaused || + this.host.getScheduler().suspended || + this.host.getScheduler().queuedWorkPauseCount > 0 || + this.host.isDisposed() || + this.host.isDisposing() + ) { + return; + } + while (true) { + const index = this.messages.findIndex((message) => this.isRlmTerminalNotice(message)); + if (index < 0) break; + const message = this.messages[index]; + try { + this.enqueueRlmTerminalNoticeAction(message); + } catch { + return; + } + this.messages.splice(index, 1); + } + this.host.scheduleInput(); + } + + async acquireRlmTerminalNoticeRetentionFence(): Promise<{ owner: symbol; release(): void } | undefined> { + const disposeSignal = this.host.getFence().disposeSignal; + while (!this.host.isDisposed() && !this.host.isDisposing() && !disposeSignal.aborted) { + if (this.host.getScheduler().queuedWorkPauseCount > 0) { + let wake = () => {}; + const pauseReleased = new Promise((resolve) => { + wake = resolve; + this.host.addCheckpointWaiter(resolve); + }); + try { + await waitForPromiseOrAbort(pauseReleased, disposeSignal, "Terminal notice retention cancelled"); + } catch { + return undefined; + } finally { + this.host.removeCheckpointWaiter(wake); + } + continue; + } + let fence: { owner: symbol; release(): void }; + try { + fence = await this.host.acquireFence(disposeSignal); + } catch { + return undefined; + } + if (this.host.getScheduler().queuedWorkPauseCount === 0 && !this.host.isDisposed() && !this.host.isDisposing()) + return fence; + fence.release(); + } + return undefined; + } + + async deferRlmTerminalNotice(message: CustomMessage): Promise { + this.assertRlmTerminalNotice(message); + const fence = await this.acquireRlmTerminalNoticeRetentionFence(); + if (!fence) return; + try { + if (this.host.isDisposed() || this.host.isDisposing()) return; + this.messages.push(cloneCustomMessage(message)); + this.flushDeferredRlmTerminalNotices(); + } finally { + fence.release(); + } + } + + demoteRlmTerminalNoticeActions(): void { + const actions = this.actions.clearableActions().filter((action) => this.terminalNoticeActionIds.has(action.id)); + if (actions.length === 0) return; + for (const action of actions) { + if (!this.isRlmTerminalNoticeAction(action)) continue; + const message = primaryDeliveryRecord(action).message; + if (message.role === "custom") this.messages.push(cloneCustomMessage(message)); + } + const ids = new Set(actions.map((action) => action.id)); + this.host.cancelActions( + (action) => ids.has(action.id), + new Error("RLM child terminal notice deferred across session input suspension."), + actions, + ); + for (const id of ids) this.terminalNoticeActionIds.delete(id); + } + + takePendingNextTurnMessages(): CustomMessage[] { + const messages = this.messages; + this.messages = []; + return messages; + } + + getPendingNextTurnMessageSnapshots(): readonly CustomMessage[] { + const messages = this.messages.map((message) => cloneCustomMessage(message)); + for (const action of this.actions.unfinishedActions()) { + if ( + action.payload.kind !== "turn" || + !action.payload.acceptedAgentMessage || + !primaryDeliveryRecord(action).started + ) { + continue; + } + messages.push( + ...action.payload.records + .filter( + (record): record is DeliveryRecord & { message: CustomMessage } => + (record.role === "next_turn" || record.role === "prefix") && + record.message.role === "custom" && + !record.durable, + ) + .map((record) => cloneCustomMessage(record.message)), + ); + } + return messages; + } + + restorePendingNextTurnMessages(messages: readonly CustomMessage[]): void { + this.messages.push(...messages.map((message) => cloneCustomMessage(message))); + this.flushDeferredRlmTerminalNotices(); + } +} diff --git a/packages/coding-agent/src/session/extensions.ts b/packages/coding-agent/src/session/extensions/extensions.ts similarity index 95% rename from packages/coding-agent/src/session/extensions.ts rename to packages/coding-agent/src/session/extensions/extensions.ts index 30fb280176..eee3aa3593 100644 --- a/packages/coding-agent/src/session/extensions.ts +++ b/packages/coding-agent/src/session/extensions/extensions.ts @@ -1,8 +1,8 @@ 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 { AgentSessionMessageController } from "../../core/agent-messages.js"; +import type { CompactionResult } from "../../core/compaction/index.js"; import { type ContextUsage, type ExtensionActions, @@ -13,14 +13,14 @@ import { 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"; +} 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; diff --git a/packages/coding-agent/src/goals/commands.ts b/packages/coding-agent/src/session/goals/commands.ts similarity index 93% rename from packages/coding-agent/src/goals/commands.ts rename to packages/coding-agent/src/session/goals/commands.ts index cca2fa43b8..d5bb17f06e 100644 --- a/packages/coding-agent/src/goals/commands.ts +++ b/packages/coding-agent/src/session/goals/commands.ts @@ -1,5 +1,5 @@ -import { validateGoalBudget, validateGoalObjective } from "../core/goals.js"; -import { parseSessionSlashCommand } from "../core/slash-commands.js"; +import { parseSessionSlashCommand } from "../../core/slash-commands.js"; +import { validateGoalBudget, validateGoalObjective } from "./contracts.js"; type GoalSlashCommand = | { kind: "status" } diff --git a/packages/coding-agent/src/session/goals/continuation.ts b/packages/coding-agent/src/session/goals/continuation.ts new file mode 100644 index 0000000000..1e9440d609 --- /dev/null +++ b/packages/coding-agent/src/session/goals/continuation.ts @@ -0,0 +1,370 @@ +import type { Agent, AgentContext, AgentMessage, GetContinuationMessagesContext } from "@earendil-works/pi-agent-core"; +import type { AssistantMessage, ImageContent } from "@earendil-works/pi-ai"; +import type { ActionStore } from "../../core/session-action-store.js"; +import type { SessionInputAdmission } from "../input/input-admission.js"; +import type { SessionInputScheduler } from "../input/input-scheduler.js"; +import { + createPreparedTurnAction, + normalizeMessageContent, + primaryDeliveryRecord, + type QueuedSessionAction, +} from "../prepared-actions.js"; +import { parseGoalSlashCommand } from "./commands.js"; +import { + createGoalContextMessage, + GOAL_CONTEXT_CUSTOM_TYPE, + type GoalHostResponse, + type GoalState, + goalHostResponse, + validateGoalBudget, + validateGoalObjective, +} from "./contracts.js"; +import type { GoalController } from "./controller.js"; + +export interface SessionGoalContinuationHost { + getGoalState(): GoalState; + queuePrompt: SessionInputAdmission["queuePreparedPrompt"]; + getScheduler(): Pick; + isDisposed(): boolean; + isDisposing(): boolean; + hasUnsettledChildWork(): boolean; + ensureRuntimeActive(context?: AgentContext): void; + admit: SessionInputAdmission["admitSessionInput"]; + cancelActions(predicate: (action: QueuedSessionAction) => boolean, error: Error): QueuedSessionAction[]; + clearPendingGoalContexts(): void; + emitQueueUpdate(): void; + emitGoalUpdate(): void; + validate(): Promise; + isStreaming(): boolean; + includesGoals(): boolean; + getAgent(): Pick; +} +export class SessionGoalContinuation { + private _awaitsChildWork = false; + private _abortInProgress = false; + private _thresholdContinuation: AgentMessage | undefined; + constructor( + readonly controller: GoalController, + private readonly actions: Pick, "unfinishedActions">, + private readonly host: SessionGoalContinuationHost, + ) {} + get awaitsChildWork(): boolean { + return this._awaitsChildWork; + } + + get abortInProgress(): boolean { + return this._abortInProgress; + } + + get thresholdContinuation(): AgentMessage | undefined { + return this._thresholdContinuation; + } + + beginAbort(): void { + this._abortInProgress = this.controller.state.status === "active"; + } + + finishAbort(): void { + this._abortInProgress = false; + } + + deferUntilChildSettlement(): void { + this._awaitsChildWork = true; + } + + accountAssistantBudget(message: AssistantMessage): Promise | undefined { + if (!this.controller.accountAssistantMessage(message)) return undefined; + const notice = createGoalContextMessage(this.controller.state, "budget_limit"); + const normalized = normalizeMessageContent(notice.content); + return this.host.queuePrompt("steer", normalized.text, normalized.images, { + message: notice, + resumeIfIdle: true, + }); + } + + clearQueuedGoalContexts(): void { + this._awaitsChildWork = false; + this.host.clearPendingGoalContexts(); + this.host + .getAgent() + .removeQueuedMessages( + (message) => message.role === "custom" && message.customType === GOAL_CONTEXT_CUSTOM_TYPE, + ); + this.host.cancelActions( + (action) => + action.payload.kind === "turn" && action.payload.customMessage?.customType === GOAL_CONTEXT_CUSTOM_TYPE, + new Error("Queued goal context was cleared before delivery."), + ); + this.host.emitQueueUpdate(); + } + + startGoal(objectiveText: string, tokenBudget: number | undefined): GoalState { + const objective = validateGoalObjective(objectiveText); + const budget = validateGoalBudget(tokenBudget); + this._awaitsChildWork = false; + return this.controller.start(objective, budget); + } + + clearGoal(): void { + this.clearQueuedGoalContexts(); + this.controller.clear(); + } + + pauseGoal(): void { + this.clearQueuedGoalContexts(); + this.controller.pause(); + } + + async resumeGoal(): Promise { + if (this.controller.resume()) { + await this.runOrQueueGoalContext("continuation"); + } + } + + finishGoalForTerminalAssistantMessage(message: AssistantMessage): void { + if (this.controller.state.status !== "active") { + return; + } + + if (message.stopReason === "aborted") { + this._abortInProgress = false; + return; + } + + if (message.stopReason === "error") { + if (this._abortInProgress) { + this._abortInProgress = false; + return; + } + this.controller.fail(message.errorMessage || "Assistant response failed"); + } + } + + stopGoalContinuationForTerminalMessage(message: AssistantMessage): boolean { + if (message.stopReason !== "error" && message.stopReason !== "aborted") { + return false; + } + try { + this.finishGoalForTerminalAssistantMessage(message); + } catch { + // Goal hooks must not reject; listener failures should not crash the agent loop. + } + return true; + } + + maybeResumeGoalContinuationAfterRlmWork(): void { + if (!this._awaitsChildWork) return; + if (this.host.isDisposed() || this.host.isDisposing() || this.host.hasUnsettledChildWork()) return; + if (this.controller.state.status !== "active" || !this.controller.state.objective) { + this._awaitsChildWork = false; + return; + } + // Keep the deferral while admission is paused or the pump is suspended + // (post-abort); the pause release and resumeQueuedWork retry. + if (this.host.getScheduler().admissionPaused || this.host.getScheduler().suspended) return; + const goalBeforeResume = this.controller.checkpoint(); + try { + this.host.ensureRuntimeActive(); + this.controller.recordContinuation(); + const message = createGoalContextMessage(this.controller.state, "continuation"); + const normalized = normalizeMessageContent(message.content); + // No front: a settling child's terminal notice must be read first. + this.host.admit( + createPreparedTurnAction("followUp", normalized.text, normalized.images, { + message, + resumeIfIdle: true, + }), + ); + this._awaitsChildWork = false; + } catch { + // Admission can race a new pause; roll back so the retry re-counts. + this.controller.restore(goalBeforeResume, { restoreClock: false }); + } + } + + runOrQueueGoalContext(kind: "continuation" | "objective_updated", images?: ImageContent[]): void { + if (!this.controller.state.objective) return; + this.host.ensureRuntimeActive(); + const message = createGoalContextMessage(this.controller.state, kind, images); + const normalized = normalizeMessageContent(message.content); + const action = createPreparedTurnAction("followUp", normalized.text, normalized.images, { + message, + resumeIfIdle: true, + }); + this.host.admit(action, { front: true, wake: false }); + } + + async handleGoalSlashCommand(text: string, images: ImageContent[] | undefined): Promise { + const command = parseGoalSlashCommand(text); + if (!command) { + return false; + } + + if (command.kind === "status") { + this.host.emitGoalUpdate(); + return true; + } + + if (command.kind === "clear") { + this.clearGoal(); + return true; + } + + if (command.kind === "pause") { + this.pauseGoal(); + return true; + } + + if (command.kind === "resume") { + await this.resumeGoal(); + return true; + } + + const previousWasActive = this.controller.state.status === "active"; + if (!this.host.isStreaming()) { + await this.host.validate(); + } + this.host.ensureRuntimeActive(); + this.clearQueuedGoalContexts(); + this.startGoal(command.objective, command.tokenBudget); + await this.runOrQueueGoalContext(previousWasActive ? "objective_updated" : "continuation", images); + return true; + } + + queueGoalContinuationForThresholdCompaction(message: AssistantMessage): boolean { + if (message.stopReason === "error" || message.stopReason === "aborted") { + return false; + } + if (this.controller.state.status !== "active" || !this.controller.state.objective) { + return false; + } + const alreadyQueued = this._thresholdContinuation; + if ( + alreadyQueued !== undefined && + this.actions.unfinishedActions().some((action) => { + if (action.payload.kind !== "turn" || primaryDeliveryRecord(action).message !== alreadyQueued) return false; + // A running continuation may already need a successor; only undelivered actions deduplicate. + return ( + action.lifecycle.state === "queued" || + action.lifecycle.state === "selected" || + action.lifecycle.state === "preparing" || + action.lifecycle.state === "committing" + ); + }) + ) { + return true; + } + try { + this.host.ensureRuntimeActive(); + this.controller.recordContinuation(); + const goalMessage = createGoalContextMessage(this.controller.state, "continuation"); + const normalized = normalizeMessageContent(goalMessage.content); + this.host.admit( + createPreparedTurnAction("followUp", normalized.text, normalized.images, { + message: goalMessage, + }), + ); + this._thresholdContinuation = goalMessage; + return true; + } catch { + return false; + } + } + + clearQueuedGoalContinuationAfterCancelledThresholdCompaction( + queuedGoalContinuation: AgentMessage | undefined, + ): void { + if (queuedGoalContinuation === undefined) return; + const cancelled = this.host.cancelActions( + (action) => action.payload.kind === "turn" && primaryDeliveryRecord(action).message === queuedGoalContinuation, + new Error("Queued goal continuation was cleared before delivery."), + ); + this._thresholdContinuation = undefined; + // A stale marker (continuation already consumed) matches no action; only an + // actual cancellation may roll back its queue-time continuationsUsed increment. + if (cancelled.length === 0) return; + this.controller.cancelContinuation(); + this.host.emitQueueUpdate(); + } + + handleGoalHostRequest(type: string, payload: Record = {}): GoalHostResponse { + if (!this.host.includesGoals()) { + throw new Error("goals are disabled in this session"); + } + switch (type) { + case "goal.get": + return goalHostResponse(this.host.getGoalState(), false); + case "goal.create": { + if (typeof payload.objective !== "string") { + throw new Error("goal.create objective must be a string"); + } + if (payload.token_budget !== undefined && typeof payload.token_budget !== "number") { + throw new Error("goal.create token_budget must be an integer when provided"); + } + return goalHostResponse(this.createGoalFromHost(payload.objective, payload.token_budget), false); + } + case "goal.complete": + return goalHostResponse(this.completeGoalFromHost(), true); + default: + throw new Error(`unknown goal request type "${type}"`); + } + } + + createGoalFromHost(objective: string, tokenBudget: number | undefined): GoalState { + switch (this.controller.state.status) { + case "active": + throw new Error( + "cannot create a new goal because this thread already has an active goal; run `await goal.complete()` when it is achieved, or ask the user to clear it with /goal clear", + ); + case "paused": + throw new Error( + "cannot create a new goal because a paused goal exists; ask the user to resume it with /goal resume or clear it with /goal clear", + ); + case "budget_limited": + throw new Error( + "cannot create a new goal because a budget-limited goal exists; ask the user to resume it with /goal resume or clear it with /goal clear", + ); + default: + // idle, or a terminal record (complete / error): nothing pending, start fresh. + return this.startGoal(objective, tokenBudget); + } + } + + completeGoalFromHost(): GoalState { + // Accounting precedes the completing ipython cell, so its budget-limit + // context may already be queued and must be withdrawn before completion. + return this.controller.complete(() => this.clearQueuedGoalContexts()); + } + + async getGoalContinuationMessages( + context: GetContinuationMessagesContext, + signal?: AbortSignal, + ): Promise { + if (this.stopGoalContinuationForTerminalMessage(context.message)) { + return []; + } + if (signal?.aborted || this.controller.state.status !== "active" || !this.controller.state.objective) { + return []; + } + // Delegating and ending the turn is correct behavior; hold the continuation + // until descendants settle instead of re-prompting a waiting parent. + if (this.host.hasUnsettledChildWork()) { + this._awaitsChildWork = true; + return []; + } + this._awaitsChildWork = false; + try { + this.host.ensureRuntimeActive(context.context); + this.controller.recordContinuation(); + return [createGoalContextMessage(this.controller.state, "continuation")]; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + try { + this.controller.fail(message); + } catch { + // The continuation hook must not reject; listener failures should not crash the agent loop. + } + return []; + } + } +} diff --git a/packages/coding-agent/src/core/goals.ts b/packages/coding-agent/src/session/goals/contracts.ts similarity index 99% rename from packages/coding-agent/src/core/goals.ts rename to packages/coding-agent/src/session/goals/contracts.ts index 4bff368176..1be3f30eda 100644 --- a/packages/coding-agent/src/core/goals.ts +++ b/packages/coding-agent/src/session/goals/contracts.ts @@ -1,5 +1,5 @@ import type { ImageContent, TextContent } from "@earendil-works/pi-ai"; -import type { CustomMessage } from "./messages.js"; +import type { CustomMessage } from "../../core/messages.js"; export const GOAL_STATE_CUSTOM_TYPE = "thread_goal_state"; export const GOAL_CONTEXT_CUSTOM_TYPE = "goal_context"; diff --git a/packages/coding-agent/src/goals/controller.ts b/packages/coding-agent/src/session/goals/controller.ts similarity index 99% rename from packages/coding-agent/src/goals/controller.ts rename to packages/coding-agent/src/session/goals/controller.ts index 4115015642..c099e9b48a 100644 --- a/packages/coding-agent/src/goals/controller.ts +++ b/packages/coding-agent/src/session/goals/controller.ts @@ -1,6 +1,6 @@ import { randomUUID } from "node:crypto"; import type { AssistantMessage } from "@earendil-works/pi-ai"; -import { emptyGoalState, type GoalState, goalTokenDeltaForUsage, normalizeGoalState } from "../core/goals.js"; +import { emptyGoalState, type GoalState, goalTokenDeltaForUsage, normalizeGoalState } from "./contracts.js"; import type { GoalPersistence } from "./persistence.js"; export interface GoalCheckpoint { diff --git a/packages/coding-agent/src/goals/persistence.ts b/packages/coding-agent/src/session/goals/persistence.ts similarity index 92% rename from packages/coding-agent/src/goals/persistence.ts rename to packages/coding-agent/src/session/goals/persistence.ts index 25e077791c..c7b74b2b54 100644 --- a/packages/coding-agent/src/goals/persistence.ts +++ b/packages/coding-agent/src/session/goals/persistence.ts @@ -1,11 +1,11 @@ +import type { SessionManager } from "../../core/session-manager.js"; import { emptyGoalState, GOAL_STATE_CUSTOM_TYPE, type GoalState, isPersistedGoalState, normalizeGoalState, -} from "../core/goals.js"; -import type { SessionManager } from "../core/session-manager.js"; +} from "./contracts.js"; export interface GoalPersistence { load(): GoalState; diff --git a/packages/coding-agent/src/session/input/action-queue.ts b/packages/coding-agent/src/session/input/action-queue.ts new file mode 100644 index 0000000000..b0d7c41ab7 --- /dev/null +++ b/packages/coding-agent/src/session/input/action-queue.ts @@ -0,0 +1,536 @@ +import type { Agent } from "@earendil-works/pi-agent-core"; +import type { ImageContent, TextContent } from "@earendil-works/pi-ai"; +import { isAgentSessionMessage, isAgentSessionMessagePrompt } from "../../core/agent-messages.js"; +import { + ASYNC_BASH_COMPLETION_CUSTOM_TYPE, + type AsyncBashCompletionDetails, + type CustomMessage, + HARNESS_DIGEST_CUSTOM_TYPE, + isSessionSlashCommandMessage, +} from "../../core/messages.js"; +import { + type ActionStore, + type DeliveryPolicy, + type DeliveryRecord, + type QueuedMessageLane, + type QueuedMessageMutation, + type QueuedMessageMutationStatus, + queuedMessageLaneDeliveryPolicy, + type SessionAction, + type SessionActionSnapshot, +} from "../../core/session-action-store.js"; +import { parseSessionSlashCommand } from "../../core/slash-commands.js"; +import { + cloneCustomMessage, + type createPreparedTurnAction, + createSessionCommandAction, + type PreparedTurnPayload, + primaryDeliveryRecord, + type QueuedSessionAction, + queuedAgentMessagePreview, + type RestoredPromptInput, + type SessionInputSchedule, + visibleSessionActionProjection, +} from "../prepared-actions.js"; +import type { SessionInputScheduler } from "./input-scheduler.js"; + +export interface SessionActionQueueHost { + formatLabel(text: string): string; + getScheduler(): Pick; + getAgent(): Pick; + rearmDigest(): void; + restoreNextTurnMessages(messages: CustomMessage[]): void; + notifyCheckpoints(): void; + emitQueueUpdate(): void; + settleAgentMessage(id: string | undefined, leg: "delivery" | "completion", error?: Error): void; + rejectAgentMessage(id: string | undefined, error: Error): void; + admit(action: QueuedSessionAction, options: { restore: true }): { accepted: boolean }; + queuePrompt( + schedule: SessionInputSchedule, + text: string, + images: ImageContent[] | undefined, + options: Parameters[3], + ): Promise; + resumeQueuedWork(): boolean; +} +export class SessionActionQueue { + private _clearEpoch: number = 0; + get clearEpoch(): number { + return this._clearEpoch; + } + constructor( + private readonly actions: ActionStore, + private readonly host: SessionActionQueueHost, + ) {} + get steeringStopPending(): boolean { + return ( + this.actions.queuedActions("next_turn_boundary").length > 0 || + this.actions + .activeActions("next_turn_boundary") + .some( + (action) => + action.payload.kind === "turn" && + (action.lifecycle.state === "selected" || action.lifecycle.state === "preparing"), + ) + ); + } + + get hasPendingSessionWork(): boolean { + return this.actions.unfinishedActions().some((action) => { + const state = action.lifecycle.state; + return ( + state === "queued" || + state === "selected" || + state === "preparing" || + (state === "committing" && action.payload.kind === "turn" && !primaryDeliveryRecord(action).durable) + ); + }); + } + + get hasAcceptedPromptInFlight(): boolean { + return this.actions + .unfinishedActions() + .some( + (action) => + action.payload.kind === "turn" && + !action.payload.queueVisible && + action.payload.acceptedBeforeCompletion, + ); + } + + cancelSessionActions( + predicate: (action: QueuedSessionAction) => boolean, + error: Error, + candidates = this.actions.clearableActions(), + ): QueuedSessionAction[] { + const matching = candidates.filter(predicate); + const previousStates = new Map(matching.map((action) => [action.id, action.lifecycle.state])); + const preparing = this.actions + .activeActions() + .filter( + (action): action is SessionAction => + action.payload.kind === "turn" && action.lifecycle.state === "preparing", + ); + const previousAnchor = preparing.at(-1); + const actions = this.actions.remove(predicate, candidates); + const restorableMessages: CustomMessage[] = []; + const removed = new Set(actions); + if (previousAnchor && removed.has(previousAnchor)) { + for (const action of preparing) { + if (!removed.has(action)) action.payload.prepared = undefined; + } + } + for (const action of actions) { + const ticket = this.actions.ticketFor(action); + if ( + action.payload.kind === "turn" && + (action.payload.acceptedAgentMessage || + !action.payload.queueVisible || + previousStates.get(action.id) !== "queued") + ) { + ticket.rejectDelivered(error); + } else { + ticket.settleDelivered({ status: "not_applicable" }); + } + ticket.settleCompleted(error); + const dispatched = previousStates.get(action.id) === "committing" && action.payload.kind === "turn"; + if (action.payload.kind === "turn") { + const payload = action.payload; + const restorable = payload.records + .filter( + (record): record is DeliveryRecord & { message: CustomMessage } => + (record.role === "next_turn" || (payload.acceptedAgentMessage && record.role === "prefix")) && + record.message.role === "custom" && + record.message.customType !== HARNESS_DIGEST_CUSTOM_TYPE && + !record.durable, + ) + .map((record) => cloneCustomMessage(record.message)); + restorableMessages.push(...restorable); + // Lazy injection owns digest delivery: a cancelled turn re-arms it + // instead of restoring a possibly stale digest message. + if ( + payload.records.some( + (record) => + record.message.role === "custom" && record.message.customType === HARNESS_DIGEST_CUSTOM_TYPE, + ) + ) { + this.host.rearmDigest(); + } + if (dispatched) { + payload.captureRunMessages = new Set(payload.records.map((record) => record.message)); + this.host.getAgent().state.messages = this.host + .getAgent() + .state.messages.filter((message) => !payload.captureRunMessages?.has(message)); + } + } + if (!dispatched) { + this.actions.releaseTerminal(action); + } + } + this.host.restoreNextTurnMessages(restorableMessages); + if (actions.length > 0) this.host.notifyCheckpoints(); + return actions; + } + + /** + * The kernel read the command's result before the notice reached the model, so + * the notice has nothing left to report: drop it while it is still queued. + * Delivered notices are no longer clearable, which makes this a no-op. + */ + withdrawAsyncBashCompletionNotice(details: { pid: number; command: string }): void { + // One read withdraws one notice: pid reuse can queue an identical key twice, + // and the read belongs to the older handle, which is the earlier notice. + const notice = this.actions + .clearableActions() + .find((action) => this.isAsyncBashCompletionActionFor(action, details)); + if (!notice) return; + this.cancelSessionActions( + (action) => action === notice, + new Error("Background command completion notice withdrawn: the kernel read the result first."), + ); + this.host.emitQueueUpdate(); + } + + private isAsyncBashCompletionActionFor( + action: QueuedSessionAction, + details: { pid: number; command: string }, + ): boolean { + if (action.payload.kind !== "turn") return false; + const message = primaryDeliveryRecord(action).message; + if (message.role !== "custom" || message.customType !== ASYNC_BASH_COMPLETION_CUSTOM_TYPE) return false; + // pids are reused across handles, so the command has to match too. + const completion = message.details as AsyncBashCompletionDetails | undefined; + return completion?.pid === details.pid && completion.command === details.command; + } + + restoreSessionCommand( + text: string, + customMessage: CustomMessage | undefined, + images: ImageContent[] | undefined, + schedule: SessionInputSchedule, + agentMessageId: string | undefined, + ): boolean | undefined { + if (!isSessionSlashCommandMessage(customMessage) || text !== customMessage.details.command.text) { + return undefined; + } + return this.host.admit( + createSessionCommandAction(text, customMessage.details.command, images, schedule, { + agentMessageId, + source: "internal", + }), + { restore: true }, + ).accepted; + } + + restorePromptInput(schedule: SessionInputSchedule, snapshot: RestoredPromptInput): Promise { + return this.host.queuePrompt(schedule, snapshot.text, snapshot.images, { + queueKey: snapshot.queueKey, + agentMessageId: snapshot.agentMessageId, + content: snapshot.content, + message: snapshot.customMessage, + prefixMessages: snapshot.prefixMessages, + source: "internal", + }); + } + + async restoreSteeringMessage( + text: string, + images?: ImageContent[], + options: { + queueKey?: string; + agentMessageId?: string; + content?: (TextContent | ImageContent)[]; + customMessage?: CustomMessage; + prefixMessages?: CustomMessage[]; + } = {}, + ): Promise { + if ( + this.restoreSessionCommand(text, options.customMessage, images, "steer", options.agentMessageId) !== undefined + ) + return; + + await this.restorePromptInput("steer", { + text, + images, + queueKey: options.queueKey, + agentMessageId: options.agentMessageId, + content: options.content, + customMessage: options.customMessage, + prefixMessages: options.prefixMessages, + }); + } + + async restoreFollowUpMessage( + text: string, + images?: ImageContent[], + options: { + queueKey?: string; + agentMessageId?: string; + content?: (TextContent | ImageContent)[]; + customMessage?: CustomMessage; + prefixMessages?: CustomMessage[]; + } = {}, + ): Promise { + const restoredCommand = this.restoreSessionCommand( + text, + options.customMessage, + images, + "followUp", + options.agentMessageId, + ); + if (restoredCommand !== undefined) return restoredCommand; + + return this.restorePromptInput("followUp", { + text, + images, + queueKey: options.queueKey, + agentMessageId: options.agentMessageId, + content: options.content, + customMessage: options.customMessage, + prefixMessages: options.prefixMessages, + }); + } + + clearQueue(): { steering: string[]; followUp: string[] } { + const clearable = this.actions + .clearableActions() + .filter((action) => action.payload.kind === "session_command" || action.payload.queueVisible); + if (clearable.some((action) => action.payload.kind === "turn" && action.lifecycle.state === "preparing")) { + this.host.getScheduler().invalidatePreparation(); + } + const steering = clearable + .filter((action) => action.delivery === "next_turn_boundary") + .map((action) => action.payload.text); + const followUp = clearable + .filter((action) => action.delivery === "when_run_idle") + .map((action) => action.payload.text); + const promptError = new Error("Queued prompt was cleared before delivery."); + const agentMessageError = new Error("Queued agent message was cleared before delivery."); + for (const action of clearable) { + const error = + action.payload.kind === "turn" && action.lifecycle.state === "preparing" ? promptError : agentMessageError; + this.host.settleAgentMessage(action.agentMessageId, "delivery", error); + this.host.settleAgentMessage(action.agentMessageId, "completion", error); + } + const clearableIds = new Set(clearable.map((action) => action.id)); + this.cancelSessionActions((action) => clearableIds.has(action.id), agentMessageError); + this.host.getAgent().clearAllQueues(); + this.host.emitQueueUpdate(); + return { steering, followUp }; + } + + invalidateQueuedPromptPreparation(): void { + for (const action of this.actions.clearableActions()) { + if (action.payload.kind === "turn") action.payload.prepared = undefined; + } + } + + clearQueuedAgentMessages(): { steering: string[]; followUp: string[] } { + this._clearEpoch++; + // customType identifies agent messages; the text parser covers persisted pre-grammar prompts. + return this.clearQueuedTurnActionsMatching( + (action) => + isAgentSessionMessage(primaryDeliveryRecord(action).message) || + isAgentSessionMessagePrompt(action.payload.text), + ); + } + + clearQueuedUserMessagesMatching(predicate: (text: string) => boolean): { steering: string[]; followUp: string[] } { + return this.clearQueuedTurnActionsMatching((action) => predicate(action.payload.text)); + } + + private clearQueuedTurnActionsMatching(matches: (action: QueuedSessionAction) => boolean): { + steering: string[]; + followUp: string[]; + } { + const ownedActions = this.actions.ownedActions(); + const dispatchedTurnCount = ownedActions.filter( + (action) => + action.payload.kind === "turn" && + (action.lifecycle.state === "committing" || action.lifecycle.state === "running"), + ).length; + const matching = ownedActions.filter( + (action) => + action.payload.kind === "turn" && + action.agentMessageId !== undefined && + matches(action) && + (action.lifecycle.state === "queued" || + action.lifecycle.state === "selected" || + action.lifecycle.state === "preparing" || + (action.lifecycle.state === "committing" && + dispatchedTurnCount === 1 && + !primaryDeliveryRecord(action).started)), + ); + if (matching.length === 0) return { steering: [], followUp: [] }; + const removedTexts = (delivery: DeliveryPolicy) => + [ + ...matching.filter((action) => action.delivery === delivery && action.lifecycle.state === "queued"), + ...matching.filter((action) => action.delivery === delivery && action.lifecycle.state !== "queued"), + ].map((action) => action.payload.text); + const removedSteering = removedTexts("next_turn_boundary"); + const removedFollowUp = removedTexts("when_run_idle"); + const acceptedError = new Error("Accepted agent message was cleared before delivery."); + const queuedError = new Error("Queued agent message was cleared before delivery."); + for (const action of matching) { + const error = + action.payload.kind === "turn" && action.payload.acceptedAgentMessage ? acceptedError : queuedError; + this.host.rejectAgentMessage(action.agentMessageId, error); + } + for (const [accepted, error] of [ + [true, acceptedError], + [false, queuedError], + ] as const) { + const ids = new Set( + matching + .filter((action) => action.payload.kind === "turn" && action.payload.acceptedAgentMessage === accepted) + .map((action) => action.id), + ); + if (ids.size > 0) this.cancelSessionActions((action) => ids.has(action.id), error, matching); + } + if ( + matching.some( + (action) => + action.lifecycle.state === "cancelled" && + action.payload.kind === "turn" && + action.payload.captureRunMessages, + ) + ) { + this.host.getAgent().abort(); + } + this.host.emitQueueUpdate(); + return { steering: removedSteering, followUp: removedFollowUp }; + } + + mutateQueuedMessage( + lane: QueuedMessageLane, + index: number, + expectedText: string, + mutation: QueuedMessageMutation, + ): QueuedMessageMutationStatus { + const policy = queuedMessageLaneDeliveryPolicy(lane); + const projection = visibleSessionActionProjection(this.actions.queuedActions(policy)); + const item = projection[index]; + if (!item || queuedAgentMessagePreview(item) !== expectedText) return "rejected"; + if (mutation.type === "delete") { + const error = new Error("Queued prompt was deleted before delivery."); + this.host.rejectAgentMessage(item.agentMessageId, error); + this.cancelSessionActions((candidate) => candidate === item, error); + this.host.emitQueueUpdate(); + this.host.resumeQueuedWork(); + return "applied"; + } + if (mutation.type === "move") { + const neighbor = projection[index + mutation.direction]; + if (!neighbor) return "rejected"; + this.actions.swapQueued(item, neighbor); + this.host.emitQueueUpdate(); + return "applied"; + } + if ( + item.payload.kind === "turn" && + (item.payload.acceptedAgentMessage || + item.payload.records.some((record) => record.role === "primary" && record.message.role !== "user")) + ) { + return "rejected"; + } + const images = mutation.images?.map((image) => ({ ...image })); + if (item.payload.kind === "session_command") { + const command = parseSessionSlashCommand(mutation.text); + if (!command) return "invalid"; + item.payload.text = mutation.text; + item.payload.command = command; + if (mutation.images !== undefined) item.payload.images = images?.length ? images : undefined; + } else { + item.payload.text = mutation.text; + const text = { type: "text" as const, text: mutation.text }; + if (mutation.images !== undefined) { + item.payload.images = images?.length ? images : undefined; + item.payload.content = [text, ...(images?.map((image) => ({ ...image })) ?? [])]; + } else if (item.payload.content) { + item.payload.content = [text, ...item.payload.content.filter((block) => block.type !== "text")]; + } + item.payload.preview = undefined; + item.payload.prepared = undefined; + for (const record of item.payload.records) { + if (record.role === "primary" && record.message.role === "user") { + record.message.content = item.payload.content?.map((block) => ({ ...block })) ?? mutation.text; + } + } + } + const targetPolicy = queuedMessageLaneDeliveryPolicy(mutation.lane); + if (targetPolicy !== policy) { + item.queueKey = undefined; + item.wake = mutation.lane === "steering" ? "on_lower_boundary" : "external_resume"; + this.actions.moveQueued(item, targetPolicy, this.actions.queuedActions(targetPolicy).length); + } + this.host.resumeQueuedWork(); + this.host.emitQueueUpdate(); + return "applied"; + } + + getSessionActionSnapshot(): SessionActionSnapshot { + const steering = visibleSessionActionProjection(this.actions.queuedActions("next_turn_boundary")).map( + queuedAgentMessagePreview, + ); + const followUps = visibleSessionActionProjection(this.actions.queuedActions("when_run_idle")).map( + queuedAgentMessagePreview, + ); + const active = visibleSessionActionProjection(this.actions.activeActions())[0]; + const activeState = active?.lifecycle.state; + const phase = + activeState === "selected" + ? "preparing" + : activeState === "preparing" || activeState === "committing" || activeState === "running" + ? activeState + : undefined; + return { + queuedCount: steering.length + followUps.length, + steering, + followUps, + ...(active && phase + ? { + active: { + kind: active.payload.kind, + phase, + label: this.host.formatLabel(active.payload.text), + }, + } + : {}), + }; + } + + getSteeringMessages(): readonly string[] { + return visibleSessionActionProjection(this.actions.queuedActions("next_turn_boundary")).map( + (action) => action.payload.text, + ); + } + + getSteeringMessagePreviews(): readonly string[] { + return visibleSessionActionProjection(this.actions.queuedActions("next_turn_boundary")).map( + queuedAgentMessagePreview, + ); + } + + getFollowUpMessages(): readonly string[] { + return visibleSessionActionProjection(this.actions.queuedActions("when_run_idle")).map( + (action) => action.payload.text, + ); + } + + getFollowUpMessagePreviews(): readonly string[] { + return visibleSessionActionProjection(this.actions.queuedActions("when_run_idle")).map(queuedAgentMessagePreview); + } + + removeQueuedFollowUp(queueKey: string): boolean { + const matching = this.actions + .clearableActions() + .filter((action) => action.payload.kind === "turn" && action.queueKey === queueKey); + if (matching.length === 0) return false; + const error = new Error("Queued agent message was cleared before delivery."); + for (const action of matching) this.host.rejectAgentMessage(action.agentMessageId, error); + const ids = new Set(matching.map((action) => action.id)); + this.cancelSessionActions((action) => ids.has(action.id), error); + this.host.emitQueueUpdate(); + return true; + } +} diff --git a/packages/coding-agent/src/session/input/action-recovery.ts b/packages/coding-agent/src/session/input/action-recovery.ts new file mode 100644 index 0000000000..0793420376 --- /dev/null +++ b/packages/coding-agent/src/session/input/action-recovery.ts @@ -0,0 +1,184 @@ +import type { ActionStore } from "../../core/session-action-store.js"; +import { + cloneCustomMessage, + cloneQueuedAgentMessage, + type PreparedCommandPayload, + type PreparedTurnPayload, + type QueuedSessionAction, + SESSION_ACTION_RECOVERY_FORMAT_VERSION, + type SessionActionRecoverySnapshot, +} from "../prepared-actions.js"; + +export interface SessionActionRecoveryHost { + isTerminalNoticeAction(action: QueuedSessionAction): boolean; + retainTerminalNotice(id: string): void; + releaseTerminalNotice(id: string): void; + admit(action: QueuedSessionAction, options: { restore: true }): unknown; +} +export class SessionActionRecovery { + constructor( + private readonly actions: Pick, "ownedActions" | "snapshotActions">, + private readonly host: SessionActionRecoveryHost, + ) {} + async restoreSessionActions(snapshot: SessionActionRecoverySnapshot): Promise { + if (snapshot.formatVersion !== SESSION_ACTION_RECOVERY_FORMAT_VERSION) { + throw new Error(`Unsupported session action recovery format version: ${snapshot.formatVersion}`); + } + const actionIds = new Set(this.actions.ownedActions().map((action) => action.id)); + const actions = snapshot.actions.map((recovered): QueuedSessionAction => { + if (actionIds.has(recovered.id)) throw new Error(`Duplicate session action id: ${recovered.id}`); + actionIds.add(recovered.id); + if ( + recovered.payload.kind === "turn" && + recovered.payload.records.some((record) => record.ownerActionId !== recovered.id) + ) { + throw new Error(`Session action ${recovered.id} has invalid delivery correlation`); + } + const payload: PreparedTurnPayload | PreparedCommandPayload = + recovered.payload.kind === "turn" + ? { + kind: "turn", + text: recovered.payload.text, + ...(recovered.payload.preview ? { preview: recovered.payload.preview } : {}), + records: recovered.payload.records.map((record) => ({ + id: record.id, + role: record.role, + message: cloneQueuedAgentMessage(record.message), + started: false, + durable: false, + ownerActionId: record.ownerActionId, + })), + ...(recovered.payload.images + ? { + images: recovered.payload.images.map((image) => ({ + ...image, + })), + } + : {}), + ...(recovered.payload.content + ? { + content: recovered.payload.content.map((block) => ({ + ...block, + })), + } + : {}), + ...(recovered.payload.customMessage + ? { + customMessage: cloneCustomMessage(recovered.payload.customMessage), + } + : {}), + executionPolicy: { + ...recovered.payload.executionPolicy, + preparation: { + ...recovered.payload.executionPolicy.preparation, + }, + }, + queueVisible: recovered.payload.queueVisible, + acceptedAgentMessage: recovered.payload.acceptedAgentMessage, + acceptedBeforeCompletion: recovered.payload.acceptedBeforeCompletion, + } + : { + kind: "session_command", + text: recovered.payload.text, + command: { ...recovered.payload.command }, + ...(recovered.payload.images + ? { + images: recovered.payload.images.map((image) => ({ + ...image, + })), + } + : {}), + }; + return { + id: recovered.id, + source: recovered.source, + delivery: recovered.delivery, + wake: recovered.wake, + payload, + lifecycle: { state: "queued" }, + ...(recovered.queueKey ? { queueKey: recovered.queueKey } : {}), + ...(recovered.agentMessageId ? { agentMessageId: recovered.agentMessageId } : {}), + ...(recovered.suppressAutonomousContinuation ? { suppressAutonomousContinuation: true } : {}), + }; + }); + for (const action of actions) { + const durableTerminalNotice = this.host.isTerminalNoticeAction(action); + if (durableTerminalNotice) this.host.retainTerminalNotice(action.id); + try { + this.host.admit(action, { restore: true }); + } catch (error) { + if (durableTerminalNotice) this.host.releaseTerminalNotice(action.id); + throw error; + } + } + return actions.length; + } + + getSessionActionRecoverySnapshot(): SessionActionRecoverySnapshot { + return { + formatVersion: SESSION_ACTION_RECOVERY_FORMAT_VERSION, + actions: this.actions.snapshotActions().map((action) => ({ + id: action.id, + source: action.source, + delivery: action.delivery, + wake: action.wake, + ...(action.queueKey ? { queueKey: action.queueKey } : {}), + ...(action.agentMessageId ? { agentMessageId: action.agentMessageId } : {}), + ...(action.suppressAutonomousContinuation ? { suppressAutonomousContinuation: true } : {}), + payload: + action.payload.kind === "turn" + ? { + kind: "turn", + text: action.payload.text, + ...(action.payload.preview ? { preview: action.payload.preview } : {}), + records: action.payload.records.map((record) => ({ + id: record.id, + role: record.role, + message: cloneQueuedAgentMessage(record.message), + ownerActionId: record.ownerActionId, + })), + ...(action.payload.images + ? { + images: action.payload.images.map((image) => ({ + ...image, + })), + } + : {}), + ...(action.payload.content + ? { + content: action.payload.content.map((block) => ({ + ...block, + })), + } + : {}), + ...(action.payload.customMessage + ? { + customMessage: cloneCustomMessage(action.payload.customMessage), + } + : {}), + executionPolicy: { + ...action.payload.executionPolicy, + preparation: { + ...action.payload.executionPolicy.preparation, + }, + }, + queueVisible: action.payload.queueVisible, + acceptedAgentMessage: action.payload.acceptedAgentMessage, + acceptedBeforeCompletion: action.payload.acceptedBeforeCompletion, + } + : { + kind: "session_command", + text: action.payload.text, + command: { ...action.payload.command }, + ...(action.payload.images + ? { + images: action.payload.images.map((image) => ({ + ...image, + })), + } + : {}), + }, + })), + }; + } +} diff --git a/packages/coding-agent/src/session/commit-fence.ts b/packages/coding-agent/src/session/input/commit-fence.ts similarity index 96% rename from packages/coding-agent/src/session/commit-fence.ts rename to packages/coding-agent/src/session/input/commit-fence.ts index a28f5c49d6..7cd9e68931 100644 --- a/packages/coding-agent/src/session/commit-fence.ts +++ b/packages/coding-agent/src/session/input/commit-fence.ts @@ -1,5 +1,5 @@ import { AsyncLocalStorage } from "node:async_hooks"; -import { waitForPromiseOrAbort } from "../utils/wait-for-abort.js"; +import { waitForPromiseOrAbort } from "../../utils/wait-for-abort.js"; export interface SessionCommitLease { readonly owner: symbol; diff --git a/packages/coding-agent/src/session/input/input-admission.ts b/packages/coding-agent/src/session/input/input-admission.ts new file mode 100644 index 0000000000..c2f89f936a --- /dev/null +++ b/packages/coding-agent/src/session/input/input-admission.ts @@ -0,0 +1,162 @@ +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import type { ImageContent, TextContent } from "@earendil-works/pi-ai"; +import { + assertAgentMessageQueueCapacity, + DEFAULT_AGENT_MESSAGE_MAX_PENDING_PER_SESSION, + isAgentSessionMessage, +} from "../../core/agent-messages.js"; +import type { InputSource } from "../../core/extensions/index.js"; +import type { CustomMessage } from "../../core/messages.js"; +import type { ActionStore, ActionTicket } from "../../core/session-action-store.js"; +import { + createPreparedTurnAction, + primaryDeliveryRecord, + type QueuedAgentMessage, + type QueuedSessionAction, + SessionInputAdmissionPausedError, + type SessionInputSchedule, +} from "../prepared-actions.js"; +import type { SessionInputScheduler } from "./input-scheduler.js"; + +export interface SessionInputAdmissionHost { + getScheduler(): Pick; + isDisposed(): boolean; + isDisposing(): boolean; + isStreaming(): boolean; + rejectAgentMessage(id: string | undefined, error: Error): void; + emitQueueUpdate(): void; + resumeAdmission(): void; + scheduleInput(): void; + suppressForMessage(message: AgentMessage): void; +} +export class SessionInputAdmission { + private _arrivalEpoch: number = 0; + get arrivalEpoch(): number { + return this._arrivalEpoch; + } + constructor( + private readonly actions: ActionStore, + private readonly host: SessionInputAdmissionHost, + ) {} + coalescedFollowUpOwner(action: QueuedSessionAction): QueuedSessionAction | undefined { + if (action.delivery !== "when_run_idle" || action.payload.kind !== "turn" || !action.queueKey) return undefined; + return this.actions + .unfinishedActions() + .find( + (candidate) => + candidate.queueKey === action.queueKey && + (candidate.lifecycle.state === "queued" || + candidate.lifecycle.state === "selected" || + candidate.lifecycle.state === "preparing"), + ); + } + + assertSessionActionAdmissionAvailable(): void { + if (this.host.isDisposed() || this.host.isDisposing()) { + throw new Error("Cannot admit a session action because the session is disposing or disposed."); + } + if (this.host.getScheduler().admissionPaused) { + throw new SessionInputAdmissionPausedError( + "Cannot admit a session action while session input admission is paused.", + ); + } + if (this.host.getScheduler().suspended) { + throw new Error("Cannot admit a session action while queued session input is suspended."); + } + } + + admitSessionInput( + action: QueuedSessionAction, + options: { + restore?: boolean; + front?: boolean; + wake?: boolean; + immediatelyEligible?: boolean; + } = {}, + ): { + accepted: boolean; + disposition: "starts_when_admitted" | "queued"; + ticket?: ActionTicket; + } { + if (this.host.isDisposed() || this.host.isDisposing()) { + throw new Error("Cannot admit a session action because the session is disposing or disposed."); + } + if (this.host.getScheduler().admissionPaused) { + throw new SessionInputAdmissionPausedError( + "Cannot admit a session action while session input admission is paused.", + ); + } + if ( + options.restore !== true && + action.payload.kind === "turn" && + isAgentSessionMessage(primaryDeliveryRecord(action).message) + ) { + assertAgentMessageQueueCapacity( + this.actions.unfinishedActions().length, + DEFAULT_AGENT_MESSAGE_MAX_PENDING_PER_SESSION, + ); + } + const coalescedOwner = options.restore ? undefined : this.coalescedFollowUpOwner(action); + if (coalescedOwner) { + if (action.agentMessageId !== coalescedOwner.agentMessageId) { + this.host.rejectAgentMessage( + action.agentMessageId, + new Error("Prompt was not queued because an equivalent follow-up is already pending."), + ); + } + return { accepted: false, disposition: "queued" }; + } + const canStartImmediately = + options.immediatelyEligible === true && + (this.actions.unfinishedActions().length === 0 || options.front === true); + if (options.front) this.actions.enqueueFront(action); + else this.actions.enqueue(action); + let disposition: "starts_when_admitted" | "queued" = "queued"; + if (canStartImmediately && this.actions.selectFirst() === action) disposition = "starts_when_admitted"; + const controller = this.actions.ticketFor(action); + controller.settleAccepted({ + status: "accepted", + actionId: action.id, + disposition, + }); + this._arrivalEpoch++; + this.host.emitQueueUpdate(); + if ( + !options.restore && + options.wake !== false && + (disposition === "starts_when_admitted" || + (action.delivery === "next_turn_boundary" && this.host.isStreaming()) || + action.payload.kind === "session_command" || + action.wake === "immediate") + ) { + if (action.payload.kind === "turn" && action.wake === "immediate") { + this.host.resumeAdmission(); + } + this.host.scheduleInput(); + } + return { accepted: true, disposition, ticket: controller.ticket }; + } + + async queuePreparedPrompt( + schedule: SessionInputSchedule, + text: string, + images?: ImageContent[], + options: { + agentMessageId?: string; + queueKey?: string; + content?: (TextContent | ImageContent)[]; + message?: QueuedAgentMessage; + prefixMessages?: CustomMessage[]; + previewLabel?: string; + suppressAutonomousContinuation?: boolean; + resumeIfIdle?: boolean; + source?: InputSource | "internal"; + } = {}, + ): Promise { + const action = createPreparedTurnAction(schedule, text, images, options); + if (action.suppressAutonomousContinuation) { + this.host.suppressForMessage(primaryDeliveryRecord(action).message); + } + return this.admitSessionInput(action).accepted; + } +} diff --git a/packages/coding-agent/src/session/input/input-checkpoints.ts b/packages/coding-agent/src/session/input/input-checkpoints.ts new file mode 100644 index 0000000000..66987c89ae --- /dev/null +++ b/packages/coding-agent/src/session/input/input-checkpoints.ts @@ -0,0 +1,208 @@ +import type { Agent } from "@earendil-works/pi-agent-core"; +import type { ActionStore } from "../../core/session-action-store.js"; +import type { SessionManager } from "../../core/session-manager.js"; +import { waitForPromiseOrAbort } from "../../utils/wait-for-abort.js"; +import { primaryDeliveryRecord, type QueuedSessionAction } from "../prepared-actions.js"; +import type { ContinuationToken, SessionContinuation } from "../turns/continuation.js"; +import type { SessionCommitFence, SessionCommitLease } from "./commit-fence.js"; +import type { SessionInputScheduler } from "./input-scheduler.js"; + +export interface SessionInputCheckpointsHost { + getFence(): Pick; + getScheduler(): Pick; + getEventQueue(): Promise; + acquireFence(signal?: AbortSignal): Promise; + getStore(): Pick; + assertAdmissionAvailable(): void; + getContinuation(): Pick; + scheduleInput(): void; + getAgent(): Pick; + getUnfinishedCount(): number; + waitForIdle(): Promise; +} +export class SessionInputCheckpoints { + private readonly waiters = new Set<() => void>(); + constructor( + private readonly actions: Pick, "activeActions" | "queuedActions">, + private readonly host: SessionInputCheckpointsHost, + ) {} + get hasWaiters(): boolean { + return this.waiters.size > 0; + } + add(waiter: () => void): void { + this.waiters.add(waiter); + } + remove(waiter: () => void): void { + this.waiters.delete(waiter); + } + + notifySessionInputCheckpointChange(): void { + const waiters = [...this.waiters]; + this.waiters.clear(); + for (const resolve of waiters) resolve(); + } + + waitForSessionActivityChange(signal: AbortSignal): Promise { + return new Promise((resolve) => { + const finish = () => { + this.waiters.delete(finish); + signal.removeEventListener("abort", finish); + resolve(); + }; + this.waiters.add(finish); + signal.addEventListener("abort", finish, { once: true }); + if (signal.aborted) finish(); + }); + } + + observeSessionActionDeferral(action: QueuedSessionAction): { + deferred: Promise; + stop(): void; + } { + let resolveDeferral = () => {}; + const deferred = new Promise((resolve) => { + resolveDeferral = resolve; + }); + const check = () => { + if (action.lifecycle.state === "queued") resolveDeferral(); + else this.waiters.add(check); + }; + this.waiters.add(check); + return { + deferred, + stop: () => this.waiters.delete(check), + }; + } + + async waitForSessionInputCheckpoint(signal?: AbortSignal): Promise { + const blocksCheckpoint = () => + this.actions.activeActions().some((action) => { + if (action.payload.kind === "session_command") { + return action.lifecycle.state === "selected" || action.lifecycle.state === "running"; + } + return ( + action.lifecycle.state === "selected" || + action.lifecycle.state === "preparing" || + (action.lifecycle.state === "committing" && !primaryDeliveryRecord(action).durable) + ); + }); + while (true) { + while (blocksCheckpoint()) { + if (signal?.aborted) throw new Error("Update restart preparation cancelled"); + await new Promise((resolve, reject) => { + const onChange = () => { + cleanup(); + resolve(); + }; + const onAbort = () => { + cleanup(); + reject(new Error("Update restart preparation cancelled")); + }; + const cleanup = () => { + this.waiters.delete(onChange); + signal?.removeEventListener("abort", onAbort); + }; + this.waiters.add(onChange); + signal?.addEventListener("abort", onAbort, { once: true }); + if (signal?.aborted) onAbort(); + }); + } + const commitFence = await this.host.acquireFence(signal); + try { + if (blocksCheckpoint()) continue; + if (signal?.aborted) throw new Error("Update restart preparation cancelled"); + await waitForPromiseOrAbort(this.host.getEventQueue(), signal, "Update restart preparation cancelled"); + if (signal?.aborted) throw new Error("Update restart preparation cancelled"); + this.host.getStore().flushNow(); + return; + } finally { + commitFence.release(); + } + } + } + + async acquireDirectTurnAdmissionFence(signal?: AbortSignal): Promise<{ owner: symbol; release(): void }> { + if (this.host.getFence().isHeldByCurrentContext) { + this.host.assertAdmissionAvailable(); + return this.host.acquireFence(signal); + } + const disposeSignal = this.host.getFence().disposeSignal; + const waitSignal = signal ? AbortSignal.any([signal, disposeSignal]) : disposeSignal; + while (true) { + this.host.assertAdmissionAvailable(); + if (this.host.getScheduler().queuedWorkPauseCount > 0) { + let wake = () => {}; + const pauseReleased = new Promise((resolve) => { + wake = resolve; + this.waiters.add(resolve); + }); + try { + await waitForPromiseOrAbort(pauseReleased, waitSignal, "Update restart preparation cancelled"); + } catch (error) { + if (disposeSignal.aborted) { + throw new Error("Cannot admit a session action because the session is disposing or disposed."); + } + throw error; + } finally { + this.waiters.delete(wake); + } + continue; + } + const fence = await this.host.acquireFence(signal); + try { + if (this.host.getScheduler().queuedWorkPauseCount === 0) { + this.host.assertAdmissionAvailable(); + return fence; + } + } catch (error) { + fence.release(); + throw error; + } + fence.release(); + } + } + + async waitForHeadlessIdle(): Promise { + while (true) { + await this.host.waitForIdle(); + const postCompactionContinuation = this.host.getContinuation().current?.promise; + if (!postCompactionContinuation) return; + await postCompactionContinuation; + } + } + + async waitForIdleOrSettlement(settlement?: ContinuationToken): Promise { + while (settlement === undefined || this.host.getContinuation().current === settlement) { + if (this.actions.queuedActions().length > 0) { + if (this.host.getScheduler().suspended || this.host.getScheduler().queuedWorkPauseCount > 0) { + let wake = () => {}; + const changed = new Promise((resolve) => { + wake = resolve; + this.waiters.add(resolve); + }); + try { + await (settlement ? Promise.race([changed, settlement.promise]) : changed); + } finally { + this.waiters.delete(wake); + } + continue; + } + this.host.scheduleInput(); + } + const pump = this.host.getScheduler().pendingPump; + await pump; + await this.host.getAgent().waitForIdle(); + const agentEventQueue = this.host.getEventQueue(); + await agentEventQueue; + if ( + pump === this.host.getScheduler().pendingPump && + agentEventQueue === this.host.getEventQueue() && + !this.host.getScheduler().requested && + !this.host.getAgent().state.isStreaming && + this.host.getUnfinishedCount() === 0 + ) { + return; + } + } + } +} diff --git a/packages/coding-agent/src/session/input-dispatcher.ts b/packages/coding-agent/src/session/input/input-dispatcher.ts similarity index 97% rename from packages/coding-agent/src/session/input-dispatcher.ts rename to packages/coding-agent/src/session/input/input-dispatcher.ts index baef151b13..965e1e72ab 100644 --- a/packages/coding-agent/src/session/input-dispatcher.ts +++ b/packages/coding-agent/src/session/input/input-dispatcher.ts @@ -5,9 +5,9 @@ import { type DeliveryPolicy, type RuntimeActivity, transitionSessionAction, -} from "../core/session-action-store.js"; -import { DeferredSessionInputError, primaryDeliveryRecord, type QueuedSessionAction } from "./prepared-actions.js"; -import { turnExecutionPoliciesEqual } from "./turn-preparation.js"; +} from "../../core/session-action-store.js"; +import { DeferredSessionInputError, primaryDeliveryRecord, type QueuedSessionAction } from "../prepared-actions.js"; +import { turnExecutionPoliciesEqual } from "../turns/turn-preparation.js"; export interface SessionInputDispatcherHost { isDisposed(): boolean; diff --git a/packages/coding-agent/src/session/input-scheduler.ts b/packages/coding-agent/src/session/input/input-scheduler.ts similarity index 100% rename from packages/coding-agent/src/session/input-scheduler.ts rename to packages/coding-agent/src/session/input/input-scheduler.ts diff --git a/packages/coding-agent/src/session/input/message-delivery.ts b/packages/coding-agent/src/session/input/message-delivery.ts new file mode 100644 index 0000000000..401f7cf267 --- /dev/null +++ b/packages/coding-agent/src/session/input/message-delivery.ts @@ -0,0 +1,243 @@ +import { randomUUID } from "node:crypto"; +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import type { AgentSessionEvent, PromptOptions } from "../../core/agent-session.js"; +import type { KernelSentAgentMessage } from "../../core/kernel/index.js"; +import type { ActionStore } from "../../core/session-action-store.js"; +import type { SessionManager } from "../../core/session-manager.js"; +import type { QueuedSessionAction } from "../prepared-actions.js"; + +interface AgentMessageDeferred { + promise: Promise; + resolve: () => void; + reject: (error: Error) => void; +} + +interface AgentMessageOutcome { + delivery?: AgentMessageDeferred; + completion?: AgentMessageDeferred; +} + +function createAgentMessageDeferred(): AgentMessageDeferred { + const deferred = {} as AgentMessageDeferred; + deferred.promise = new Promise((resolve, reject) => { + deferred.resolve = resolve; + deferred.reject = reject; + }); + deferred.promise.catch(() => undefined); + return deferred; +} + +const IPYTHON_SENT_AGENT_MESSAGE_CUSTOM_ENTRY = "ipython_sent_agent_message"; + +interface PersistedIpythonSentAgentMessage { + toolCallId: string; + message: KernelSentAgentMessage; +} + +function isObjectRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function parsePersistedIpythonSentAgentMessage(value: unknown): PersistedIpythonSentAgentMessage | undefined { + if (!isObjectRecord(value) || typeof value.toolCallId !== "string" || !isObjectRecord(value.message)) { + return undefined; + } + const { id, message, deliveryStatus, target } = value.message; + if ( + typeof id !== "string" || + typeof message !== "string" || + (deliveryStatus !== "delivered" && deliveryStatus !== "queued") || + !isObjectRecord(target) || + typeof target.activeSessionId !== "string" || + typeof target.sessionId !== "string" + ) { + return undefined; + } + return { + toolCallId: value.toolCallId, + message: { + id, + message, + deliveryStatus, + target: { + activeSessionId: target.activeSessionId, + sessionId: target.sessionId, + ...(typeof target.sessionName === "string" ? { sessionName: target.sessionName } : {}), + }, + }, + }; +} + +function appendSentAgentMessageToToolResult( + message: AgentMessage, + toolCallId: string, + sentMessage: KernelSentAgentMessage, +): boolean { + if (message.role !== "toolResult" || message.toolName !== "ipython" || message.toolCallId !== toolCallId) { + return false; + } + const details = isObjectRecord(message.details) ? message.details : {}; + const current = Array.isArray(details.sentAgentMessages) ? details.sentAgentMessages : []; + if (current.some((entry) => isObjectRecord(entry) && entry.id === sentMessage.id)) { + return true; + } + message.details = { + ...details, + sentAgentMessages: [...current, sentMessage], + }; + return true; +} + +function asError(error: unknown): Error { + return error instanceof Error ? error : new Error(String(error)); +} +export interface SessionMessageDeliveryHost { + isDisposed(): boolean; + getMessages(): AgentMessage[]; + getStore(): Pick; + enqueue(work: () => void): void; + emit(event: AgentSessionEvent): void; + promptUntilAccepted(text: string, options?: PromptOptions): Promise; + cancelActions(predicate: (action: QueuedSessionAction) => boolean, error: Error): QueuedSessionAction[]; +} +export class SessionMessageDelivery { + private readonly outcomes = new Map(); + private readonly lateMessages = new Map(); + constructor( + private readonly actions: Pick, "unfinishedActions">, + private readonly host: SessionMessageDeliveryHost, + ) {} + dispose(deliveryError: Error, completionError: Error): void { + this.rejectQueuedAgentMessageDeliveries(deliveryError, completionError); + for (const [id, outcome] of this.outcomes) { + if (outcome.delivery) this.settleAgentMessage(id, "delivery", deliveryError); + if (outcome.completion) this.settleAgentMessage(id, "completion", completionError); + } + } + agentMessageOutcome(agentMessageId: string): AgentMessageOutcome { + let outcome = this.outcomes.get(agentMessageId); + if (!outcome) { + outcome = {}; + this.outcomes.set(agentMessageId, outcome); + } + return outcome; + } + + waitForAgentMessagePromptDelivery(agentMessageId: string): Promise { + const outcome = this.agentMessageOutcome(agentMessageId); + outcome.delivery ??= createAgentMessageDeferred(); + return outcome.delivery.promise; + } + + settleAgentMessage(agentMessageId: string | undefined, leg: "delivery" | "completion", error?: Error): void { + if (agentMessageId === undefined) return; + const outcome = this.outcomes.get(agentMessageId); + if (!outcome) return; + const deferred = outcome[leg]; + if (!deferred) return; + outcome[leg] = undefined; + if (!outcome.delivery && !outcome.completion) { + this.outcomes.delete(agentMessageId); + } + if (error) deferred.reject(error); + else deferred.resolve(); + } + + rejectAgentMessage(agentMessageId: string | undefined, error: Error): void { + if (agentMessageId === undefined) return; + this.settleAgentMessage(agentMessageId, "delivery", error); + this.settleAgentMessage(agentMessageId, "completion", error); + } + + rejectQueuedAgentMessageDeliveries(deliveryError: Error, completionError = deliveryError): void { + for (const action of this.actions.unfinishedActions()) { + this.settleAgentMessage(action.agentMessageId, "delivery", deliveryError); + this.settleAgentMessage(action.agentMessageId, "completion", completionError); + } + } + + async promptAndWait(text: string, options?: PromptOptions): Promise { + const agentMessageId = options?.agentMessageId ?? `prompt-wait:${randomUUID()}`; + if (this.outcomes.get(agentMessageId)?.completion) { + throw new Error(`Prompt completion id is already in use: ${agentMessageId}`); + } + const outcome = this.agentMessageOutcome(agentMessageId); + outcome.completion = createAgentMessageDeferred(); + const completion = outcome.completion.promise; + const signal = options?.signal; + let cancelQueuedPrompt: (() => void) | undefined; + try { + await this.host.promptUntilAccepted(text, { ...options, agentMessageId }); + if (signal) { + cancelQueuedPrompt = () => { + const error = new Error("Prompt was cancelled before it started."); + const cancelled = this.host.cancelActions( + (action) => action.agentMessageId === agentMessageId && action.payload.kind === "turn", + error, + ); + if (cancelled.length > 0) { + this.settleAgentMessage(agentMessageId, "completion", error); + } + }; + signal.addEventListener("abort", cancelQueuedPrompt, { once: true }); + if (signal.aborted) cancelQueuedPrompt(); + } + await completion; + } catch (error) { + this.settleAgentMessage(agentMessageId, "completion", asError(error)); + throw error; + } finally { + if (signal && cancelQueuedPrompt) { + signal.removeEventListener("abort", cancelQueuedPrompt); + } + } + } + + restoreLateIpythonSentAgentMessages(): void { + this.lateMessages.clear(); + for (const entry of this.host.getStore().getBranch()) { + if (entry.type !== "custom" || entry.customType !== IPYTHON_SENT_AGENT_MESSAGE_CUSTOM_ENTRY) { + continue; + } + const persisted = parsePersistedIpythonSentAgentMessage(entry.data); + if (persisted) { + this.rememberLateIpythonSentAgentMessage(persisted.toolCallId, persisted.message); + } + } + } + + rememberLateIpythonSentAgentMessage(toolCallId: string, message: KernelSentAgentMessage): boolean { + const messages = this.lateMessages.get(toolCallId) ?? []; + const isNew = !messages.some((entry) => entry.id === message.id); + if (isNew) { + messages.push(message); + this.lateMessages.set(toolCallId, messages); + } + for (let index = this.host.getMessages().length - 1; index >= 0; index -= 1) { + if (appendSentAgentMessageToToolResult(this.host.getMessages()[index], toolCallId, message)) { + break; + } + } + return isNew; + } + + applyLateIpythonSentAgentMessages(message: AgentMessage): void { + if (message.role !== "toolResult" || message.toolName !== "ipython") { + return; + } + for (const sentMessage of this.lateMessages.get(message.toolCallId) ?? []) { + appendSentAgentMessageToToolResult(message, message.toolCallId, sentMessage); + } + } + + recordLateIpythonSentAgentMessage(toolCallId: string, message: KernelSentAgentMessage): void { + const record = () => { + if (this.host.isDisposed() || !this.rememberLateIpythonSentAgentMessage(toolCallId, message)) { + return; + } + this.host.getStore().appendCustomEntry(IPYTHON_SENT_AGENT_MESSAGE_CUSTOM_ENTRY, { toolCallId, message }); + this.host.emit({ type: "ipython_sent_agent_message", toolCallId, message }); + }; + this.host.enqueue(record); + } +} diff --git a/packages/coding-agent/src/session/input/prompt-submission.ts b/packages/coding-agent/src/session/input/prompt-submission.ts new file mode 100644 index 0000000000..d78ebe1559 --- /dev/null +++ b/packages/coding-agent/src/session/input/prompt-submission.ts @@ -0,0 +1,633 @@ +import type { Agent, AgentMessage } from "@earendil-works/pi-agent-core"; +import type { ImageContent, TextContent, UserMessage } from "@earendil-works/pi-ai"; +import { + type AgentSessionMessage, + isAgentSessionMessage, + parseAgentSessionMessagePromptId, +} from "../../core/agent-messages.js"; +import type { InputSource } from "../../core/extensions/index.js"; +import { + ASYNC_BASH_COMPLETION_CUSTOM_TYPE, + ASYNC_BASH_COMPLETION_PREVIEW_LABEL, + type AsyncBashCompletionDetails, + type CustomMessage, + createAsyncBashCompletionMessage, + HEARTBEAT_PROMPT_CUSTOM_TYPE, + HEARTBEAT_PROMPT_PREVIEW_LABEL, +} from "../../core/messages.js"; +import { throwIfPromptAdmissionCancelled } from "../../core/prompt-admission.js"; +import { type ActionStore, canSelectSessionAction, type RuntimeActivity } from "../../core/session-action-store.js"; +import type { SessionManager } from "../../core/session-manager.js"; +import { GOAL_CONTEXT_CUSTOM_TYPE, GOAL_CONTEXT_PREVIEW_LABEL } from "../goals/contracts.js"; +import { + buildPromptContent, + cloneCustomMessage, + createPreparedTurnAction, + createSessionCommandAction, + normalizeMessageContent, + primaryDeliveryRecord, + type QueuedSessionAction, + SessionInputAdmissionPausedError, +} from "../prepared-actions.js"; +import type { AgentSessionEvent } from "../turns/events.js"; +import { createTurnExecutionPolicy, type TurnExecutionPolicy } from "../turns/turn-preparation.js"; +import type { SessionCommitFence, SessionCommitLease } from "./commit-fence.js"; +import type { SessionInputAdmission } from "./input-admission.js"; +import type { SessionInputScheduler } from "./input-scheduler.js"; +import type { SubmissionNormalizer } from "./submission-normalization.js"; +export interface PromptOptions { + expandPromptTemplates?: boolean; + images?: ImageContent[]; + streamingBehavior?: "steer" | "followUp"; + followUpQueueKey?: string; + source?: InputSource; + preflightResult?: (success: boolean, queued?: boolean) => void; + queueIfBusy?: boolean; + resumeIfIdle?: boolean; + internalPrompt?: boolean; + suppressAutonomousContinuation?: boolean; + skipInputHandlers?: boolean; + signal?: AbortSignal; + admissionCommitted?: () => void; + agentMessageId?: string; + content?: (TextContent | ImageContent)[]; + customMessage?: CustomMessage; +} + +export interface InternalPromptOptions extends PromptOptions { + skipPrePromptWork?: boolean; + returnAfterAccepted?: boolean; + agentMessageId?: string; +} +function oncePreflight( + preflightResult: ((success: boolean, queued?: boolean) => void) | undefined, +): (success: boolean, queued?: boolean) => void { + let settled = false; + return (success, queued = false) => { + if (!settled) { + settled = true; + preflightResult?.(success, queued); + } + }; +} +function injectedMessagePreviewLabel(message: CustomMessage): string | undefined { + switch (message.customType) { + case HEARTBEAT_PROMPT_CUSTOM_TYPE: + return HEARTBEAT_PROMPT_PREVIEW_LABEL; + case ASYNC_BASH_COMPLETION_CUSTOM_TYPE: + return ASYNC_BASH_COMPLETION_PREVIEW_LABEL; + case GOAL_CONTEXT_CUSTOM_TYPE: + return GOAL_CONTEXT_PREVIEW_LABEL; + default: + return undefined; + } +} +export interface SessionPromptSubmissionHost { + promptInjectedMessage: SessionPromptSubmission["promptInjectedMessage"]; + queueAgentMessagePrompt( + text: string, + streamingBehavior: "steer" | "followUp", + customMessage?: AgentSessionMessage, + ): Promise; + getScheduler(): Pick; + getFence(): Pick; + waitForActivityChange(signal: AbortSignal): Promise; + isStreaming(): boolean; + isCompacting(): boolean; + isRetrying(): boolean; + isBashRunning(): boolean; + resumeAdmission(): void; + assertAdmissionAvailable(): void; + acquireAdmissionFence(signal?: AbortSignal): Promise; + normalize: SubmissionNormalizer["normalizeSubmission"]; + settleAgentMessage(id: string | undefined, leg: "delivery" | "completion", error?: Error): void; + canStartImmediately(): boolean; + admit: SessionInputAdmission["admitSessionInput"]; + waitForInputIdle(): Promise; + isBusy(point: "preflight" | "pump"): boolean; + takeNextTurnMessages(): CustomMessage[]; + restoreNextTurnMessages(messages: CustomMessage[]): void; + appendNextTurnMessage(message: CustomMessage): void; + getActivity(): RuntimeActivity; + suppressForMessage(message: AgentMessage): void; + observeDeferral(action: QueuedSessionAction): { deferred: Promise; stop(): void }; + rejectAgentMessage(id: string | undefined, error: Error): void; + cancelActions(predicate: (action: QueuedSessionAction) => boolean, error: Error): QueuedSessionAction[]; + emitQueueUpdate(): void; + getClearEpoch(): number; + queuePrompt: SessionInputAdmission["queuePreparedPrompt"]; + resetParentReply(): void; + getAgent(): Pick; + getStore(): Pick; + emit(event: AgentSessionEvent): void; +} +export class SessionPromptSubmission { + constructor( + private readonly actions: Pick, "unfinishedActions">, + private readonly host: SessionPromptSubmissionHost, + ) {} + async handleKernelBashCompletion(details: AsyncBashCompletionDetails): Promise { + const message = createAsyncBashCompletionMessage(details); + const disposeSignal = this.host.getFence().disposeSignal; + while (true) { + let admissionCommitted = false; + try { + await this.host.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.host.getScheduler().admissionPaused && !disposeSignal.aborted) { + await this.host.waitForActivityChange(disposeSignal); + } + } + } + } + + async promptInjectedMessage( + text: string, + message: CustomMessage, + options?: InternalPromptOptions & { executionPolicy?: TurnExecutionPolicy }, + ): Promise { + if (!this.host.isStreaming() && options?.resumeIfIdle) this.host.resumeAdmission(); + const admissionEpoch = this.host.getScheduler().epoch; + const admissionFence = await this.host.acquireAdmissionFence(options?.signal).catch((error: unknown) => { + throwIfPromptAdmissionCancelled(options?.signal); + throw error; + }); + const reportPreflight = oncePreflight(options?.preflightResult); + try { + throwIfPromptAdmissionCancelled(options?.signal); + if (admissionEpoch !== this.host.getScheduler().epoch) { + throw new Error("Injected session input was invalidated before admission"); + } + options?.admissionCommitted?.(); + const queueForStreaming = this.host.isStreaming(); + const queueForBusy = options?.queueIfBusy === true && this.host.isBusy("preflight"); + const visibleQueued = queueForStreaming || queueForBusy; + if (visibleQueued && !options?.streamingBehavior) { + const stateDescription = queueForStreaming ? "Agent is already processing" : "Agent has queued work"; + throw new Error( + `${stateDescription}. Specify streamingBehavior ('steer' or 'followUp') to queue the message.`, + ); + } + const schedule = options?.streamingBehavior ?? "followUp"; + const prefixMessages = visibleQueued ? this.host.takeNextTurnMessages() : undefined; + const action = createPreparedTurnAction(schedule, text, undefined, { + message, + prefixMessages, + queueKey: options?.followUpQueueKey, + previewLabel: injectedMessagePreviewLabel(message), + suppressAutonomousContinuation: options?.suppressAutonomousContinuation, + resumeIfIdle: + !visibleQueued || + options?.resumeIfIdle || + (options?.queueIfBusy === true && canSelectSessionAction(this.host.getActivity())), + source: options?.source ?? "internal", + executionPolicy: + options?.executionPolicy ?? + (visibleQueued ? createTurnExecutionPolicy("queued") : createTurnExecutionPolicy("injected")), + queueVisible: visibleQueued, + }); + const result = this.host.admit(action, { + immediatelyEligible: !visibleQueued, + }); + admissionFence.release(); + if (!result.accepted || !result.ticket) { + if (prefixMessages) this.host.restoreNextTurnMessages(prefixMessages); + reportPreflight(false, false); + return; + } + if (result.disposition === "queued") { + reportPreflight(true, true); + } else { + void result.ticket.delivered.then( + () => reportPreflight(true), + () => reportPreflight(false), + ); + } + if (options?.returnAfterAccepted) { + if (result.disposition === "starts_when_admitted") await result.ticket.delivered; + return; + } + if (visibleQueued) return; + await result.ticket.completed; + } catch (error) { + reportPreflight(false); + throw error; + } finally { + admissionFence.release(); + } + } + + async prompt(text: string, options?: InternalPromptOptions): Promise { + const resumeSuspendedInput = options?.resumeIfIdle !== false; + if (!this.host.isStreaming()) { + if (resumeSuspendedInput) this.host.resumeAdmission(); + this.host.assertAdmissionAvailable(); + } + const admissionEpoch = this.host.getScheduler().epoch; + const commitFence = this.host.isStreaming() + ? undefined + : await this.host.acquireAdmissionFence(options?.signal).catch((error: unknown) => { + throwIfPromptAdmissionCancelled(options?.signal); + throw error; + }); + const reportPreflight = oncePreflight(options?.preflightResult); + const run = async () => { + try { + throwIfPromptAdmissionCancelled(options?.signal); + if (!resumeSuspendedInput && admissionEpoch !== this.host.getScheduler().epoch) { + throw new Error("Session input was invalidated before admission"); + } + options?.admissionCommitted?.(); + const isInternalPrompt = options?.internalPrompt === true; + const expandPromptTemplates = isInternalPrompt ? false : (options?.expandPromptTemplates ?? true); + const normalizationResult = this.host.normalize(text, options?.images, { + parseSessionCommands: !isInternalPrompt && !options?.skipPrePromptWork, + extensionCommands: expandPromptTemplates ? "execute" : "ignore", + inputSource: + !isInternalPrompt && !options?.skipInputHandlers ? (options?.source ?? "interactive") : undefined, + expandSkills: expandPromptTemplates, + expandPromptTemplates, + }); + const normalized = normalizationResult instanceof Promise ? await normalizationResult : normalizationResult; + // Async input handlers ran between the admission check above and + // admission itself; re-check so content invalidated during that + // await (e.g. a cron job cancelled or updated) is not admitted. + if (normalizationResult instanceof Promise) options?.admissionCommitted?.(); + if (normalized.kind === "extensionCommand") { + commitFence?.release(); + reportPreflight(true); + void normalized.completion.then( + () => this.host.settleAgentMessage(options?.agentMessageId, "completion"), + (error) => this.host.settleAgentMessage(options?.agentMessageId, "completion", error), + ); + void normalized.completion.catch(() => undefined); + if (!options?.returnAfterAccepted) await normalized.completion.catch(() => undefined); + return; + } + if (normalized.kind === "handled") { + commitFence?.release(); + reportPreflight(true); + this.host.settleAgentMessage(options?.agentMessageId, "completion"); + return; + } + + const pendingOwnedWork = this.actions.unfinishedActions().length > 0; + const wasRuntimeBusy = + this.host.isStreaming() || + this.host.isCompacting() || + this.host.isRetrying() || + this.host.isBashRunning(); + const wasBusy = wasRuntimeBusy || pendingOwnedWork; + if (normalized.kind === "sessionCommand") { + const schedule = options?.streamingBehavior ?? (this.host.isStreaming() ? "steer" : "followUp"); + const action = createSessionCommandAction( + normalized.text, + normalized.command, + normalized.images, + schedule, + { + agentMessageId: options?.agentMessageId, + source: isInternalPrompt ? "internal" : (options?.source ?? "interactive"), + }, + ); + const result = this.host.admit(action, { + immediatelyEligible: !wasBusy && this.host.canStartImmediately(), + }); + commitFence?.release(); + reportPreflight(result.accepted, result.disposition === "queued"); + if (!result.accepted || !result.ticket) return; + if (options?.returnAfterAccepted) { + if (result.disposition === "starts_when_admitted") await result.ticket.delivered; + return; + } + if (result.disposition === "queued") return; + await this.host.waitForInputIdle(); + return; + } + + const queueForStreaming = this.host.isStreaming(); + const queueForBusy = options?.queueIfBusy === true && this.host.isBusy("preflight"); + const visibleQueued = queueForStreaming || queueForBusy; + if (visibleQueued && !options?.streamingBehavior) { + const stateDescription = queueForStreaming ? "Agent is already processing" : "Agent has queued work"; + throw new Error( + `${stateDescription}. Specify streamingBehavior ('steer' or 'followUp') to queue the message.`, + ); + } + const schedule = options?.streamingBehavior ?? "followUp"; + const prefixMessages = visibleQueued ? this.host.takeNextTurnMessages() : undefined; + const content = options?.content + ? options.content.map((block) => ({ ...block })) + : buildPromptContent(normalized.text, normalized.images); + const suppliedMessage = options?.customMessage; + const primaryMessage = suppliedMessage + ? visibleQueued + ? suppliedMessage + : cloneCustomMessage(suppliedMessage) + : ({ + role: "user", + content: content.map((block) => ({ ...block })), + timestamp: Date.now(), + } satisfies UserMessage); + const acceptedAgentMessage = options?.skipPrePromptWork === true && options.returnAfterAccepted === true; + const action = createPreparedTurnAction(schedule, normalized.text, normalized.images, { + agentMessageId: options?.agentMessageId, + queueKey: options?.followUpQueueKey, + content, + message: primaryMessage, + prefixMessages, + suppressAutonomousContinuation: options?.suppressAutonomousContinuation, + resumeIfIdle: + !visibleQueued || + options?.resumeIfIdle || + (options?.queueIfBusy === true && canSelectSessionAction(this.host.getActivity())), + source: isInternalPrompt ? "internal" : (options?.source ?? "interactive"), + executionPolicy: visibleQueued + ? createTurnExecutionPolicy("queued") + : createTurnExecutionPolicy("directPrompt", { + returnAfterAccepted: options?.returnAfterAccepted, + skipPrePromptWork: options?.skipPrePromptWork, + }), + queueVisible: visibleQueued, + acceptedAgentMessage, + acceptedBeforeCompletion: options?.returnAfterAccepted === true, + }); + if (action.suppressAutonomousContinuation) { + this.host.suppressForMessage(primaryDeliveryRecord(action).message); + } + const result = this.host.admit(action, { + immediatelyEligible: !visibleQueued && this.host.canStartImmediately(), + }); + commitFence?.release(); + if (!result.accepted || !result.ticket) { + if (prefixMessages) this.host.restoreNextTurnMessages(prefixMessages); + reportPreflight(false, false); + return; + } + if (result.disposition === "queued") { + reportPreflight(true, true); + } else { + void result.ticket.delivered.then( + () => reportPreflight(true), + () => reportPreflight(false), + ); + } + const deferralObserver = + acceptedAgentMessage && + options?.queueIfBusy === true && + !options.streamingBehavior && + result.disposition === "starts_when_admitted" + ? this.host.observeDeferral(action) + : undefined; + if (acceptedAgentMessage && !queueForStreaming && !queueForBusy && !options?.streamingBehavior) { + try { + const outcome = deferralObserver + ? await Promise.race([ + result.ticket.delivered.then(() => "delivered" as const), + deferralObserver.deferred.then(() => "deferred" as const), + ]) + : await result.ticket.delivered.then(() => "delivered" as const); + if (outcome === "deferred" && !options?.streamingBehavior) { + const error = new Error( + "Agent became busy before prompt delivery. Specify streamingBehavior ('steer' or 'followUp') to queue the message.", + ); + this.host.rejectAgentMessage(action.agentMessageId, error); + this.host.cancelActions((candidate) => candidate === action, error); + this.host.emitQueueUpdate(); + throw error; + } + return; + } finally { + deferralObserver?.stop(); + } + } + if (options?.returnAfterAccepted) { + if (result.disposition === "starts_when_admitted" || (acceptedAgentMessage && !visibleQueued)) { + await result.ticket.delivered; + } + return; + } + if (visibleQueued) return; + await result.ticket.completed; + await this.host.waitForInputIdle(); + } catch (error) { + reportPreflight(false); + throw error; + } finally { + commitFence?.release(); + } + }; + return commitFence ? this.host.getFence().run(commitFence, run) : run(); + } + + async acceptAgentMessagePrompt(text: string, options?: PromptOptions): Promise { + const customMessage = + options?.customMessage && isAgentSessionMessage(options.customMessage) ? options.customMessage : undefined; + const clearEpoch = this.host.getClearEpoch(); + const admissionCommitted = () => { + options?.admissionCommitted?.(); + if (clearEpoch !== this.host.getClearEpoch()) { + throw new Error("Agent message was cleared before admission"); + } + }; + if ( + this.host.getScheduler().suspended && + this.host.isBusy("preflight") && + options?.queueIfBusy === true && + options.streamingBehavior + ) { + admissionCommitted(); + const queued = await this.host.queueAgentMessagePrompt(text, options.streamingBehavior, customMessage); + options.preflightResult?.(queued, queued); + return; + } + await this.prompt(text, { + ...options, + resumeIfIdle: false, + expandPromptTemplates: false, + skipInputHandlers: true, + skipPrePromptWork: true, + returnAfterAccepted: true, + agentMessageId: options?.agentMessageId ?? customMessage?.details.id ?? parseAgentSessionMessagePromptId(text), + customMessage, + admissionCommitted, + }); + if (customMessage?.details.fromRelationship === "parent") this.host.resetParentReply(); + } + + async queueAgentMessagePrompt( + text: string, + streamingBehavior: "steer" | "followUp", + customMessage?: AgentSessionMessage, + ): Promise { + const agentMessageId = customMessage?.details.id ?? parseAgentSessionMessagePromptId(text); + if (streamingBehavior === "steer") { + await this.host.queuePrompt("steer", text, undefined, { + agentMessageId, + message: customMessage, + }); + if (customMessage?.details.fromRelationship === "parent") this.host.resetParentReply(); + return true; + } + const queued = await this.host.queuePrompt("followUp", text, undefined, { + agentMessageId, + message: customMessage, + }); + if (queued && customMessage?.details.fromRelationship === "parent") this.host.resetParentReply(); + return queued; + } + + async steer( + text: string, + images?: ImageContent[], + options: { + queueKey?: string; + agentMessageId?: string; + resumeIfIdle?: boolean; + } = {}, + ): Promise { + const normalized = this.host.normalize(text, images, { + parseSessionCommands: false, + extensionCommands: "reject", + expandSkills: true, + expandPromptTemplates: true, + }); + if (normalized instanceof Promise || normalized.kind !== "prompt") { + throw new Error("Queued prompt normalization did not produce a prompt"); + } + + await this.host.queuePrompt("steer", normalized.text, normalized.images, { + queueKey: options.queueKey, + agentMessageId: options.agentMessageId, + resumeIfIdle: options.resumeIfIdle, + }); + } + + async followUp( + text: string, + images?: ImageContent[], + options: { + queueKey?: string; + agentMessageId?: string; + resumeIfIdle?: boolean; + } = {}, + ): Promise { + const normalized = this.host.normalize(text, images, { + parseSessionCommands: false, + extensionCommands: "reject", + expandSkills: true, + expandPromptTemplates: true, + }); + if (normalized instanceof Promise || normalized.kind !== "prompt") { + throw new Error("Queued prompt normalization did not produce a prompt"); + } + + return this.host.queuePrompt("followUp", normalized.text, normalized.images, { + queueKey: options.queueKey, + agentMessageId: options.agentMessageId, + resumeIfIdle: options.resumeIfIdle, + }); + } + + async sendCustomMessage( + message: Pick, "customType" | "content" | "display" | "details">, + options?: { + triggerTurn?: boolean; + deliverAs?: "steer" | "followUp" | "nextTurn"; + }, + ): Promise { + const appMessage = { + role: "custom" as const, + customType: message.customType, + content: message.content, + display: message.display, + details: message.details, + timestamp: Date.now(), + } satisfies CustomMessage; + if (options?.deliverAs === "nextTurn") { + this.host.appendNextTurnMessage(appMessage); + } else if (this.host.isStreaming()) { + const normalized = normalizeMessageContent(message.content); + if (options?.deliverAs === "followUp") { + await this.host.queuePrompt("followUp", normalized.text, normalized.images, { + message: appMessage, + resumeIfIdle: true, + }); + } else { + await this.host.queuePrompt("steer", normalized.text, normalized.images, { + message: appMessage, + resumeIfIdle: true, + }); + } + } else if (options?.triggerTurn) { + if (!this.host.getScheduler().suspendedForUpdateRestart) this.host.resumeAdmission(); + const admissionFence = await this.host.acquireAdmissionFence(); + try { + const normalized = normalizeMessageContent(message.content); + const immediatelyEligible = this.host.canStartImmediately(); + const action = createPreparedTurnAction("followUp", normalized.text, normalized.images, { + message: appMessage, + resumeIfIdle: true, + executionPolicy: createTurnExecutionPolicy("customTrigger"), + queueVisible: false, + }); + const result = this.host.admit(action, { immediatelyEligible }); + admissionFence.release(); + if (!result.ticket) return; + await result.ticket.completed; + } finally { + admissionFence.release(); + } + } else { + this.host.getAgent().state.messages.push(appMessage); + this.host + .getStore() + .appendCustomMessageEntry(message.customType, message.content, message.display, message.details); + this.host.emit({ type: "message_start", message: appMessage }); + this.host.emit({ type: "message_end", message: appMessage }); + } + } + + async sendUserMessage( + content: string | (TextContent | ImageContent)[], + options?: { deliverAs?: "steer" | "followUp" }, + ): Promise { + let text: string; + let images: ImageContent[] | undefined; + + if (typeof content === "string") { + text = content; + } else { + const textParts: string[] = []; + images = []; + for (const part of content) { + if (part.type === "text") { + textParts.push(part.text); + } else { + images.push(part); + } + } + text = textParts.join("\n"); + if (images.length === 0) images = undefined; + } + + await this.prompt(text, { + expandPromptTemplates: false, + streamingBehavior: options?.deliverAs, + images, + source: "extension", + resumeIfIdle: true, + }); + } +} diff --git a/packages/coding-agent/src/session/input/submission-normalization.ts b/packages/coding-agent/src/session/input/submission-normalization.ts new file mode 100644 index 0000000000..e0daef016c --- /dev/null +++ b/packages/coding-agent/src/session/input/submission-normalization.ts @@ -0,0 +1,147 @@ +import { readFileSync } from "node:fs"; +import type { ImageContent } from "@earendil-works/pi-ai"; +import type { ExtensionRunner, InputSource } from "../../core/extensions/index.js"; +import { expandPromptTemplate, type PromptTemplate } from "../../core/prompt-templates.js"; +import type { Skill } from "../../core/skills.js"; +import { parseSessionSlashCommand, parseSlashCommand, type SessionSlashCommand } from "../../core/slash-commands.js"; +import { stripFrontmatter } from "../../utils/frontmatter.js"; + +type SubmissionExtensionCommandPolicy = "execute" | "reject" | "ignore"; + +export interface SubmissionNormalizationPolicy { + parseSessionCommands: boolean; + extensionCommands: SubmissionExtensionCommandPolicy; + inputSource?: InputSource; + expandSkills: boolean; + expandPromptTemplates: boolean; +} + +export type NormalizedSubmission = + | { kind: "prompt"; text: string; images?: ImageContent[] } + | { + kind: "sessionCommand"; + text: string; + images?: ImageContent[]; + command: SessionSlashCommand; + } + | { kind: "extensionCommand"; completion: Promise } + | { kind: "handled" }; + +export interface SubmissionNormalizationHost { + getExtensions(): Pick< + ExtensionRunner, + "hasHandlers" | "emitInput" | "getCommand" | "createCommandContext" | "emitError" + >; + getPrompts(): readonly PromptTemplate[]; + getSkills(): readonly Skill[]; +} +export class SubmissionNormalizer { + constructor(private readonly host: SubmissionNormalizationHost) {} + finishSubmissionNormalization( + text: string, + images: ImageContent[] | undefined, + policy: SubmissionNormalizationPolicy, + ): NormalizedSubmission { + let expandedText = text; + if (policy.expandSkills) expandedText = this.expandSkillCommand(expandedText); + if (policy.expandPromptTemplates) { + expandedText = expandPromptTemplate(expandedText, [...this.host.getPrompts()]); + } + return { kind: "prompt", text: expandedText, images }; + } + + normalizeSubmission( + text: string, + images: ImageContent[] | undefined, + policy: SubmissionNormalizationPolicy, + ): NormalizedSubmission | Promise { + if (policy.parseSessionCommands) { + const command = parseSessionSlashCommand(text); + if (command) return { kind: "sessionCommand", text, images, command }; + } + + if (text.startsWith("/")) { + if (policy.extensionCommands === "execute") { + const completion = this.executeExtensionCommand(text); + if (completion) return { kind: "extensionCommand", completion }; + } else if (policy.extensionCommands === "reject") { + this.throwIfExtensionCommand(text); + } + } + + if (policy.inputSource !== undefined && this.host.getExtensions().hasHandlers("input")) { + return this.host + .getExtensions() + .emitInput(text, images, policy.inputSource) + .then((result) => { + if (result.action === "handled") return { kind: "handled" }; + if (result.action === "transform") { + return this.finishSubmissionNormalization(result.text, result.images ?? images, policy); + } + return this.finishSubmissionNormalization(text, images, policy); + }); + } + + return this.finishSubmissionNormalization(text, images, policy); + } + + executeExtensionCommand(text: string): Promise | undefined { + const parsed = parseSlashCommand(text); + if (!parsed) return undefined; + const commandName = parsed.name; + const args = parsed.args; + + const command = this.host.getExtensions().getCommand(commandName); + if (!command) return undefined; + const context = this.host.getExtensions().createCommandContext(); + return Promise.resolve() + .then(() => command.handler(args, context)) + + .catch((error: unknown) => { + const commandError = error instanceof Error ? error : new Error(String(error)); + this.host.getExtensions().emitError({ + extensionPath: `command:${commandName}`, + event: "command", + error: commandError.message, + }); + throw commandError; + }); + } + + expandSkillCommand(text: string): string { + if (!text.startsWith("/skill:")) return text; + + const parsed = parseSlashCommand(text); + if (!parsed?.name.startsWith("skill:")) return text; + const skillName = parsed.name.slice("skill:".length); + const args = parsed.args; + + const skill = this.host.getSkills().find((s) => s.name === skillName); + if (!skill) return text; // Unknown skill, pass through + + try { + const content = readFileSync(skill.filePath, "utf-8"); + const body = stripFrontmatter(content).trim(); + const skillBlock = `\nReferences are relative to ${skill.baseDir}.\n\n${body}\n`; + return args ? `${skillBlock}\n\n${args}` : skillBlock; + } catch (err) { + this.host.getExtensions().emitError({ + extensionPath: skill.filePath, + event: "skill_expansion", + error: err instanceof Error ? err.message : String(err), + }); + return text; // Return original on error + } + } + + throwIfExtensionCommand(text: string): void { + const commandName = parseSlashCommand(text)?.name ?? ""; + const command = this.host.getExtensions().getCommand(commandName); + + if (command) { + throw new Error( + `Extension command "/${commandName}" cannot be queued. Use prompt() or execute the command when not streaming.`, + ); + } + } +} diff --git a/packages/coding-agent/src/session/kernel/heartbeat-host-requests.ts b/packages/coding-agent/src/session/kernel/heartbeat-host-requests.ts new file mode 100644 index 0000000000..42fe6f6f49 --- /dev/null +++ b/packages/coding-agent/src/session/kernel/heartbeat-host-requests.ts @@ -0,0 +1,115 @@ +import type { AgentCronJob, AgentRlmHeartbeatController, AgentRlmHeartbeatStatusUpdate } from "../../core/cron-jobs.js"; +import { normalizeHeartbeatDeliveryMode } from "../../core/cron-jobs.js"; + +export function handleRlmHeartbeatHostRequest( + controller: AgentRlmHeartbeatController | undefined, + type: string, + payload: Record = {}, +): Record { + if (!controller) { + throw new Error("RLM heartbeat skill is not available in this session"); + } + switch (type) { + case "rlm_heartbeat.list": { + const includeInactive = payload.include_inactive === true || payload.includeInactive === true; + return { + heartbeats: controller + .listRlmHeartbeats({ includeInactive }) + .map((heartbeat) => rlmHeartbeatHostResponse(heartbeat)), + }; + } + case "rlm_heartbeat.create": { + if (typeof payload.instruction !== "string") { + throw new Error("rlm_heartbeat.create instruction must be a string"); + } + if (payload.interval !== undefined && typeof payload.interval !== "string") { + throw new Error("rlm_heartbeat.create interval must be a string when provided"); + } + if (payload.label !== undefined && typeof payload.label !== "string") { + throw new Error("rlm_heartbeat.create label must be a string when provided"); + } + const deliveryMode = normalizeHeartbeatDeliveryMode(payload.delivery_mode ?? payload.deliveryMode); + return { + heartbeat: rlmHeartbeatHostResponse( + controller.createRlmHeartbeat({ + instruction: payload.instruction, + interval: payload.interval, + label: payload.label, + deliveryMode, + }), + ), + }; + } + case "rlm_heartbeat.update": { + if (typeof payload.id !== "string") { + throw new Error("rlm_heartbeat.update id must be a string"); + } + if (payload.instruction !== undefined && typeof payload.instruction !== "string") { + throw new Error("rlm_heartbeat.update instruction must be a string when provided"); + } + if (payload.interval !== undefined && typeof payload.interval !== "string") { + throw new Error("rlm_heartbeat.update interval must be a string when provided"); + } + if (payload.label !== undefined && typeof payload.label !== "string") { + throw new Error("rlm_heartbeat.update label must be a string when provided"); + } + if (payload.status !== undefined && !isRlmHeartbeatStatusUpdate(payload.status)) { + throw new Error('rlm_heartbeat.update status must be "pause" or "resume" when provided'); + } + const rawDeliveryMode = payload.delivery_mode ?? payload.deliveryMode; + const deliveryMode = normalizeHeartbeatDeliveryMode(rawDeliveryMode); + if ( + payload.instruction === undefined && + payload.interval === undefined && + payload.label === undefined && + payload.status === undefined && + rawDeliveryMode === undefined + ) { + throw new Error("rlm_heartbeat.update requires at least one field to update"); + } + const heartbeat = controller.updateRlmHeartbeat({ + id: payload.id, + instruction: payload.instruction, + interval: payload.interval, + label: payload.label, + status: payload.status, + deliveryMode, + }); + return { + heartbeat: heartbeat ? rlmHeartbeatHostResponse(heartbeat) : null, + }; + } + case "rlm_heartbeat.delete": { + if (typeof payload.id !== "string") { + throw new Error("rlm_heartbeat.delete id must be a string"); + } + const heartbeat = controller.deleteRlmHeartbeat(payload.id); + return { + heartbeat: heartbeat ? rlmHeartbeatHostResponse(heartbeat) : null, + }; + } + default: + throw new Error(`unknown RLM heartbeat request type "${type}"`); + } +} + +function isRlmHeartbeatStatusUpdate(value: unknown): value is AgentRlmHeartbeatStatusUpdate { + return value === "pause" || value === "resume"; +} + +function rlmHeartbeatHostResponse(job: AgentCronJob): Record { + return { + id: job.id, + status: job.status, + label: job.label ?? null, + delivery_mode: job.deliveryMode ?? "steer", + instruction: job.prompt, + schedule: job.schedule, + created_at: job.createdAt, + updated_at: job.updatedAt, + next_run_at: job.nextRunAt ?? null, + last_run_at: job.lastRunAt ?? null, + last_error: job.lastError ?? null, + run_count: job.runCount, + }; +} diff --git a/packages/coding-agent/src/session/kernel-environment.ts b/packages/coding-agent/src/session/kernel/kernel-environment.ts similarity index 90% rename from packages/coding-agent/src/session/kernel-environment.ts rename to packages/coding-agent/src/session/kernel/kernel-environment.ts index dc45cb74e6..a9363392a7 100644 --- a/packages/coding-agent/src/session/kernel-environment.ts +++ b/packages/coding-agent/src/session/kernel/kernel-environment.ts @@ -1,11 +1,11 @@ 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"; +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; diff --git a/packages/coding-agent/src/session/kernel-host-handlers.ts b/packages/coding-agent/src/session/kernel/kernel-host-handlers.ts similarity index 94% rename from packages/coding-agent/src/session/kernel-host-handlers.ts rename to packages/coding-agent/src/session/kernel/kernel-host-handlers.ts index e3fdffddaf..7803096efe 100644 --- a/packages/coding-agent/src/session/kernel-host-handlers.ts +++ b/packages/coding-agent/src/session/kernel/kernel-host-handlers.ts @@ -5,16 +5,16 @@ import { type AgentSessionMessageReceipt, agentFamilyMemberName, createAgentMessageHostHandlers, -} from "../core/agent-messages.js"; +} 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"; +} 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, @@ -28,8 +28,8 @@ import { type RlmFindModelsResult, type RlmListSubagentsResult, type RlmSpawnHandle, -} from "../core/rlm-runtime.js"; -import type { Skill } from "../core/skills.js"; +} from "../../core/rlm-runtime.js"; +import type { Skill } from "../../core/skills.js"; type ObserveResult = AgentObserveListResult | AgentObserveAgentSnapshot | AgentObserveRecentMessagesResult; export interface SessionKernelOperations { diff --git a/packages/coding-agent/src/session/kernel.ts b/packages/coding-agent/src/session/kernel/kernel.ts similarity index 92% rename from packages/coding-agent/src/session/kernel.ts rename to packages/coding-agent/src/session/kernel/kernel.ts index 02473260be..77ec6db5da 100644 --- a/packages/coding-agent/src/session/kernel.ts +++ b/packages/coding-agent/src/session/kernel/kernel.ts @@ -1,14 +1,14 @@ 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"; +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 { diff --git a/packages/coding-agent/src/session/kernel/message-host-requests.ts b/packages/coding-agent/src/session/kernel/message-host-requests.ts new file mode 100644 index 0000000000..f918c91f21 --- /dev/null +++ b/packages/coding-agent/src/session/kernel/message-host-requests.ts @@ -0,0 +1,33 @@ +import { + type AgentSessionMessageController, + type AgentSessionMessageReceipt, + assertDirectAgentMessageTarget, + normalizeAgentSessionMessage, +} from "../../core/agent-messages.js"; + +export function handleAgentMessageHostRequest( + getController: () => AgentSessionMessageController | undefined, + + type: string, + payload: Record = {}, +): Promise { + if (!getController()) { + throw new Error("agent messaging is not available in this session"); + } + switch (type) { + case "agent_message.send": { + if (typeof payload.target !== "string") { + throw new Error("agent_message.send target must be a string"); + } + if (typeof payload.message !== "string") { + throw new Error("agent_message.send message must be a string"); + } + return getController()!.sendAgentMessage({ + target: assertDirectAgentMessageTarget(payload.target), + message: normalizeAgentSessionMessage(payload.message), + }); + } + default: + throw new Error(`unknown agent message request type "${type}"`); + } +} diff --git a/packages/coding-agent/src/session/kernel/observe-host-requests.ts b/packages/coding-agent/src/session/kernel/observe-host-requests.ts new file mode 100644 index 0000000000..992718750b --- /dev/null +++ b/packages/coding-agent/src/session/kernel/observe-host-requests.ts @@ -0,0 +1,45 @@ +import { + type AgentObserveAgentSnapshot, + type AgentObserveController, + type AgentObserveListResult, + type AgentObserveRecentMessagesResult, + normalizeObserveLimit, + normalizeObserveMaxChars, +} from "../../core/agent-observe.js"; + +export function handleAgentObserveHostRequest( + controller: AgentObserveController | undefined, + + type: string, + payload: Record = {}, +): + | AgentObserveListResult + | AgentObserveAgentSnapshot + | AgentObserveRecentMessagesResult + | Promise { + if (!controller) { + throw new Error("agent observation is not available in this session"); + } + switch (type) { + case "agent_observe.list": + return controller.listAgents(); + case "agent_observe.get": { + if (typeof payload.target !== "string") { + throw new Error("agent_observe.get target must be a string"); + } + return controller.getAgent(payload.target); + } + case "agent_observe.recent": { + if (typeof payload.target !== "string") { + throw new Error("agent_observe.recent target must be a string"); + } + return controller.recentMessages({ + target: payload.target, + limit: normalizeObserveLimit(payload.limit as number | undefined), + maxChars: normalizeObserveMaxChars((payload.max_chars ?? payload.maxChars) as number | undefined), + }); + } + default: + throw new Error(`unknown agent observe request type "${type}"`); + } +} diff --git a/packages/coding-agent/src/session/models/model-selection.ts b/packages/coding-agent/src/session/models/model-selection.ts new file mode 100644 index 0000000000..3fed0bdbdc --- /dev/null +++ b/packages/coding-agent/src/session/models/model-selection.ts @@ -0,0 +1,480 @@ +import { AsyncLocalStorage } from "node:async_hooks"; +import type { AgentState, ThinkingLevel } from "@earendil-works/pi-agent-core"; +import { + type Api, + clampThinkingLevel, + getSupportedThinkingLevels, + type Model, + modelsAreEqual, + type ServiceTier, + supportsFastMode, +} from "@earendil-works/pi-ai"; +import { + formatAuthenticationFailedMessage, + formatNoApiKeyFoundMessage, + formatNoModelSelectedMessage, +} from "../../core/auth-guidance.js"; +import { DEFAULT_THINKING_LEVEL } from "../../core/defaults.js"; +import type { ExtensionRunner } from "../../core/extensions/index.js"; +import type { ModelRegistry } from "../../core/model-registry.js"; +import { findRlmModelMatches, type RlmFindModelsResult } from "../../core/rlm-runtime.js"; +import type { SessionManager } from "../../core/session-manager.js"; +import type { SettingsManager } from "../../core/settings-manager.js"; +import { THINKING_LEVELS } from "../../core/thinking-levels.js"; + +export interface ModelCycleResult { + model: Model; + thinkingLevel: ThinkingLevel; + serviceTier: ServiceTier; + isScoped: boolean; +} +export interface ModelSelectOptions { + waitForExtensions?: boolean; +} +export interface ScopedModel { + model: Model; + thinkingLevel?: ThinkingLevel; +} +export interface ModelSelectionHost { + getModel(): Model | undefined; + getState(): Pick; + setThinkingLevel(level: ThinkingLevel): void; + getAvailableThinkingLevels(): ThinkingLevel[]; + supportsThinking(): boolean; + getRegistry(): Pick< + ModelRegistry, + | "getModelForCurrentAuth" + | "getExecutableModels" + | "getApiKeyAndHeaders" + | "isUsingOAuth" + | "hasConfiguredAuth" + | "getProviderAuthStatus" + | "canUseModel" + | "clearProviderAuthStale" + | "refreshAvailableModels" + | "find" + >; + getExtensions(): Pick; + sessionManager: Pick; + settingsManager: Pick< + SettingsManager, + "setDefaultModelAndProvider" | "setDefaultThinkingLevel" | "getDefaultThinkingLevel" | "setDefaultServiceTier" + >; + emit( + event: + | { type: "thinking_level_changed"; level: ThinkingLevel } + | { type: "service_tier_changed"; serviceTier: ServiceTier }, + ): void; +} + +export class SessionModelSelection { + private _modelSelectEmitQueue: Promise = Promise.resolve(); + private _modelSelectEmitQueueIdle = true; + private _modelSelectEmitContext = new AsyncLocalStorage(); + constructor( + private readonly host: ModelSelectionHost, + private _serviceTierPreference: ServiceTier, + private _scopedModels: ScopedModel[], + ) {} + async validateCanStartAgentRun(): Promise { + if (!this.model) { + throw new Error(formatNoModelSelectedMessage()); + } + if (!this.host.getRegistry().hasConfiguredAuth(this.model)) { + const isOAuth = this.host.getRegistry().isUsingOAuth(this.model); + if (isOAuth) { + throw new Error(formatAuthenticationFailedMessage(this.model.provider)); + } + throw new Error(formatNoApiKeyFoundMessage(this.model.provider)); + } + } + + async authenticatedRlmModels(): Promise[]> { + return (await this.host.getRegistry().getExecutableModels()).filter((model) => { + const status = this.host.getRegistry().getProviderAuthStatus(model.provider); + return status.source !== "stale" && status.label !== "expired"; + }); + } + + async findRlmModels(query: string, limit: number): Promise { + return { + models: findRlmModelMatches(query, await this.authenticatedRlmModels(), limit), + }; + } + + async resolveRlmSubagentModel(reference: string | undefined, target = "subagent"): Promise<{ model: Model }> { + const parentModel = this.model; + if (!parentModel) { + throw new Error(formatNoModelSelectedMessage()); + } + if (!reference) { + return { model: parentModel }; + } + + const normalizedReference = reference.toLowerCase(); + if (`${parentModel.provider}/${parentModel.id}`.toLowerCase() === normalizedReference) { + return { model: parentModel }; + } + const model = (await this.authenticatedRlmModels()).find( + (candidate) => `${candidate.provider}/${candidate.id}`.toLowerCase() === normalizedReference, + ); + if (!model) { + throw new Error(`Requested ${target} model "${reference}" is unavailable, unauthenticated, or expired`); + } + + const auth = await this.host.getRegistry().getApiKeyAndHeaders(model); + if (!auth.ok) { + throw new Error(`Requested ${target} model "${reference}" failed authentication preflight`); + } + return { model }; + } + + get model(): Model | undefined { + return this.host.getModel(); + } + get thinkingLevel(): ThinkingLevel { + return this.host.getState().thinkingLevel; + } + get serviceTier(): ServiceTier { + return this.host.getState().serviceTier; + } + get scopedModels(): readonly ScopedModel[] { + return this._scopedModels; + } + setScopedModels(models: ScopedModel[]): void { + this._scopedModels = models; + } + async getRequiredRequestAuth(model: Model): Promise<{ + apiKey: string; + headers?: Record; + requestModel: Model; + }> { + const result = await this.host.getRegistry().getApiKeyAndHeaders(model); + if (!result.ok) { + if (result.error.startsWith("No API key found")) { + throw new Error(formatNoApiKeyFoundMessage(model.provider)); + } + throw new Error(result.error); + } + if (result.apiKey) { + return { apiKey: result.apiKey, headers: result.headers, requestModel: result.requestModel ?? model }; + } + + const isOAuth = this.host.getRegistry().isUsingOAuth(model); + if (isOAuth) { + throw new Error(formatAuthenticationFailedMessage(model.provider)); + } + throw new Error(formatNoApiKeyFoundMessage(model.provider)); + } + + private async _emitModelSelect( + nextModel: Model, + previousModel: Model | undefined, + source: "set" | "cycle" | "restore", + ): Promise { + if (modelsAreEqual(previousModel, nextModel)) return; + await this.host.getExtensions().emit({ + type: "model_select", + model: nextModel, + previousModel, + source, + }); + } + + private _queueModelSelectEmit( + nextModel: Model, + previousModel: Model | undefined, + source: "set" | "cycle" | "restore", + ): Promise { + const emit = () => + this._modelSelectEmitContext.run(true, () => this._emitModelSelect(nextModel, previousModel, source)); + this._modelSelectEmitQueueIdle = false; + const promise = this._modelSelectEmitQueue.then(emit, emit); + const queued = promise.catch(() => {}); + this._modelSelectEmitQueue = queued; + void queued.finally(() => { + if (this._modelSelectEmitQueue === queued) { + this._modelSelectEmitQueueIdle = true; + } + }); + return promise; + } + + async setModel(model: Model, options: ModelSelectOptions = {}): Promise { + // Explicit selection recovers from a stale-auth lockout, but only a fully + // validated switch commits the clear (single owner): failed selections never unlock. + const staleOnly = + !this.host.getRegistry().hasConfiguredAuth(model) && + this.host.getRegistry().getProviderAuthStatus(model.provider).source === "stale"; + if (!staleOnly && !this.host.getRegistry().hasConfiguredAuth(model)) { + throw new Error(`No API key for ${model.provider}/${model.id}`); + } + if (!(await this.host.getRegistry().canUseModel(model, { assumeAuthConfigured: staleOnly }))) { + throw new Error(`Model "${model.provider}/${model.id}" is not available for the current Prime team.`); + } + if (staleOnly) { + this.host.getRegistry().clearProviderAuthStale(model.provider); + if (!this.host.getRegistry().hasConfiguredAuth(model)) { + throw new Error(`No API key for ${model.provider}/${model.id}`); + } + } + + const previousModel = this.model; + const thinkingLevel = this._getThinkingLevelForModelSwitch(); + const serviceTier = this._getServiceTierForModelSwitch(); + this.host.getState().model = model; + this.host.sessionManager.appendModelChange(model.provider, model.id); + this.host.settingsManager.setDefaultModelAndProvider(model.provider, model.id); + + this.host.setThinkingLevel(thinkingLevel); + this._clampServiceTierForModel(serviceTier); + + const emitPromise = this._queueModelSelectEmit(model, previousModel, "set"); + if (this._shouldWaitForModelSelectEmit(options)) { + await emitPromise; + } else { + this._trackModelSelectEmitError(emitPromise); + } + } + + private _trackModelSelectEmitError(emitPromise: Promise): void { + void emitPromise.catch((error) => { + this.host.getExtensions().emitError({ + extensionPath: "", + event: "model_select", + error: error instanceof Error ? error.message : String(error), + stack: error instanceof Error ? error.stack : undefined, + }); + }); + } + + private _shouldWaitForModelSelectEmit(options: ModelSelectOptions): boolean { + return options.waitForExtensions !== false && !this._modelSelectEmitContext.getStore(); + } + + pendingModelSelectEmit(): Promise | undefined { + if (!this._modelSelectEmitContext.getStore() && !this._modelSelectEmitQueueIdle) { + return this._modelSelectEmitQueue; + } + return undefined; + } + + async cycleModel( + direction: "forward" | "backward" = "forward", + options: ModelSelectOptions = {}, + ): Promise { + if (this._scopedModels.length > 0) { + return this._cycleScopedModel(direction, options); + } + return this._cycleAvailableModel(direction, options); + } + + private async _cycleScopedModel( + direction: "forward" | "backward", + options: ModelSelectOptions, + ): Promise { + const availableModels = await this.host.getRegistry().refreshAvailableModels(); + const scopedModels = this._scopedModels.filter((scoped) => + availableModels.some((model) => modelsAreEqual(model, scoped.model)), + ); + if (scopedModels.length <= 1) return undefined; + + const currentModel = this.model; + let currentIndex = scopedModels.findIndex((sm) => modelsAreEqual(sm.model, currentModel)); + + if (currentIndex === -1) currentIndex = 0; + const len = scopedModels.length; + const nextIndex = direction === "forward" ? (currentIndex + 1) % len : (currentIndex - 1 + len) % len; + const next = scopedModels[nextIndex]; + const thinkingLevel = this._getThinkingLevelForModelSwitch(next.thinkingLevel); + const serviceTier = this._getServiceTierForModelSwitch(); + + this.host.getState().model = next.model; + this.host.sessionManager.appendModelChange(next.model.provider, next.model.id); + this.host.settingsManager.setDefaultModelAndProvider(next.model.provider, next.model.id); + + this.host.setThinkingLevel(thinkingLevel); + this._clampServiceTierForModel(serviceTier); + + const emitPromise = this._queueModelSelectEmit(next.model, currentModel, "cycle"); + if (this._shouldWaitForModelSelectEmit(options)) { + await emitPromise; + } else { + this._trackModelSelectEmitError(emitPromise); + } + + return { + model: next.model, + thinkingLevel: this.thinkingLevel, + serviceTier: this.serviceTier, + isScoped: true, + }; + } + + private async _cycleAvailableModel( + direction: "forward" | "backward", + options: ModelSelectOptions, + ): Promise { + const availableModels = await this.host.getRegistry().refreshAvailableModels(); + if (availableModels.length <= 1) return undefined; + + const currentModel = this.model; + let currentIndex = availableModels.findIndex((m) => modelsAreEqual(m, currentModel)); + + if (currentIndex === -1) currentIndex = 0; + const len = availableModels.length; + const nextIndex = direction === "forward" ? (currentIndex + 1) % len : (currentIndex - 1 + len) % len; + const nextModel = availableModels[nextIndex]; + + const thinkingLevel = this._getThinkingLevelForModelSwitch(); + const serviceTier = this._getServiceTierForModelSwitch(); + this.host.getState().model = nextModel; + this.host.sessionManager.appendModelChange(nextModel.provider, nextModel.id); + this.host.settingsManager.setDefaultModelAndProvider(nextModel.provider, nextModel.id); + + this.host.setThinkingLevel(thinkingLevel); + this._clampServiceTierForModel(serviceTier); + + const emitPromise = this._queueModelSelectEmit(nextModel, currentModel, "cycle"); + if (this._shouldWaitForModelSelectEmit(options)) { + await emitPromise; + } else { + this._trackModelSelectEmitError(emitPromise); + } + + return { + model: nextModel, + thinkingLevel: this.thinkingLevel, + serviceTier: this.serviceTier, + isScoped: false, + }; + } + + setThinkingLevel(level: ThinkingLevel): void { + const availableLevels = this.host.getAvailableThinkingLevels(); + const effectiveLevel = availableLevels.includes(level) ? level : this._clampThinkingLevel(level, availableLevels); + + const previousLevel = this.host.getState().thinkingLevel; + const isChanging = effectiveLevel !== previousLevel; + + this.host.getState().thinkingLevel = effectiveLevel; + + if (isChanging) { + this.host.sessionManager.appendThinkingLevelChange(effectiveLevel); + if (this.host.supportsThinking() || effectiveLevel !== "off") { + this.host.settingsManager.setDefaultThinkingLevel(effectiveLevel); + } + this.host.emit({ type: "thinking_level_changed", level: effectiveLevel }); + void this.host.getExtensions().emit({ + type: "thinking_level_select", + level: effectiveLevel, + previousLevel, + }); + } + } + + setServiceTier(serviceTier: ServiceTier): void { + const effectiveServiceTier = this._getEffectiveServiceTier(serviceTier); + const preferenceChanged = effectiveServiceTier !== this._serviceTierPreference; + const effectiveTierChanged = effectiveServiceTier !== this.host.getState().serviceTier; + if (!preferenceChanged && !effectiveTierChanged) { + return; + } + this._serviceTierPreference = effectiveServiceTier; + if (preferenceChanged) { + this.host.sessionManager.appendServiceTierChange(effectiveServiceTier); + if (this.model && supportsFastMode(this.model)) { + this.host.settingsManager.setDefaultServiceTier(effectiveServiceTier); + } + } + if (effectiveTierChanged) { + this.host.getState().serviceTier = effectiveServiceTier; + this.host.emit({ + type: "service_tier_changed", + serviceTier: effectiveServiceTier, + }); + } + } + + private _getEffectiveServiceTier(serviceTier: ServiceTier): ServiceTier { + return serviceTier === "priority" && (!this.model || !supportsFastMode(this.model)) ? "default" : serviceTier; + } + + private _getServiceTierForModelSwitch(): ServiceTier { + return this._serviceTierPreference; + } + + private _clampServiceTierForModel(serviceTier: ServiceTier = this.serviceTier): void { + const effectiveServiceTier = this._getEffectiveServiceTier(serviceTier); + if (effectiveServiceTier === this.host.getState().serviceTier) { + return; + } + this.host.getState().serviceTier = effectiveServiceTier; + this.host.emit({ + type: "service_tier_changed", + serviceTier: effectiveServiceTier, + }); + } + + cycleThinkingLevel(): ThinkingLevel | undefined { + if (!this.host.supportsThinking()) return undefined; + + const levels = this.host.getAvailableThinkingLevels(); + const currentIndex = levels.indexOf(this.thinkingLevel); + const nextIndex = (currentIndex + 1) % levels.length; + const nextLevel = levels[nextIndex]; + + this.host.setThinkingLevel(nextLevel); + return nextLevel; + } + + getAvailableThinkingLevels(): ThinkingLevel[] { + if (!this.model) return THINKING_LEVELS; + return getSupportedThinkingLevels(this.model) as ThinkingLevel[]; + } + + supportsThinking(): boolean { + return !!this.model?.reasoning; + } + + private _getThinkingLevelForModelSwitch(explicitLevel?: ThinkingLevel): ThinkingLevel { + if (explicitLevel !== undefined) { + return explicitLevel; + } + if (!this.host.supportsThinking()) { + return this.host.settingsManager.getDefaultThinkingLevel() ?? DEFAULT_THINKING_LEVEL; + } + return this.thinkingLevel; + } + + private _clampThinkingLevel(level: ThinkingLevel, _availableLevels: ThinkingLevel[]): ThinkingLevel { + return this.model ? (clampThinkingLevel(this.model, level) as ThinkingLevel) : "off"; + } + + refreshModelMetadata(): void { + if (this.model?.provider === "xai") { + this.host.getState().model = this.host.getRegistry().getModelForCurrentAuth(this.model); + this.host.setThinkingLevel(this.thinkingLevel); + this._clampServiceTierForModel(); + } + this._scopedModels = this._scopedModels.map((scoped) => + scoped.model.provider === "xai" + ? { ...scoped, model: this.host.getRegistry().getModelForCurrentAuth(scoped.model) } + : scoped, + ); + } + + refreshCurrentModelFromRegistry(): void { + const currentModel = this.model; + if (!currentModel) { + return; + } + + const refreshedModel = this.host.getRegistry().find(currentModel.provider, currentModel.id); + if (!refreshedModel || refreshedModel === currentModel) { + return; + } + + this.host.getState().model = refreshedModel; + } +} diff --git a/packages/coding-agent/src/session/prepared-actions.ts b/packages/coding-agent/src/session/prepared-actions.ts index 67a0088b56..fb19ed1e13 100644 --- a/packages/coding-agent/src/session/prepared-actions.ts +++ b/packages/coding-agent/src/session/prepared-actions.ts @@ -18,7 +18,7 @@ import type { WakePolicy, } from "../core/session-action-store.js"; import type { SessionSlashCommand } from "../core/slash-commands.js"; -import { createTurnExecutionPolicy, type TurnExecutionPolicy } from "./turn-preparation.js"; +import { createTurnExecutionPolicy, type TurnExecutionPolicy } from "./turns/turn-preparation.js"; export type QueuedAgentMessage = UserMessage | CustomMessage; export type SessionInputSchedule = "steer" | "followUp"; diff --git a/packages/coding-agent/src/session/auto-refinement.ts b/packages/coding-agent/src/session/refinement/auto-refinement.ts similarity index 99% rename from packages/coding-agent/src/session/auto-refinement.ts rename to packages/coding-agent/src/session/refinement/auto-refinement.ts index bebc5b5860..4f3c7bdbb7 100644 --- a/packages/coding-agent/src/session/auto-refinement.ts +++ b/packages/coding-agent/src/session/refinement/auto-refinement.ts @@ -1,6 +1,6 @@ import type { Api, Model } from "@earendil-works/pi-ai"; -import type { AutoRefineReason, AutoRefineReview, RefinementResult } from "../core/refinement/index.js"; -import type { SettingsManager } from "../core/settings-manager.js"; +import type { AutoRefineReason, AutoRefineReview, RefinementResult } from "../../core/refinement/index.js"; +import type { SettingsManager } from "../../core/settings-manager.js"; import { RefineSkippedError } from "./refinement-execution.js"; export interface AutoRefineReviewRequest { reason: AutoRefineReason; diff --git a/packages/coding-agent/src/session/refinement-execution.ts b/packages/coding-agent/src/session/refinement/refinement-execution.ts similarity index 96% rename from packages/coding-agent/src/session/refinement-execution.ts rename to packages/coding-agent/src/session/refinement/refinement-execution.ts index 11b75376d6..479f59a1ae 100644 --- a/packages/coding-agent/src/session/refinement-execution.ts +++ b/packages/coding-agent/src/session/refinement/refinement-execution.ts @@ -2,17 +2,17 @@ import { existsSync } from "node:fs"; import { dirname, resolve } from "node:path"; import type { AgentMessage, ThinkingLevel } from "@earendil-works/pi-agent-core"; import type { Api, Model } from "@earendil-works/pi-ai"; -import { formatNoModelSelectedMessage } from "../core/auth-guidance.js"; -import { serializeConversation } from "../core/compaction/index.js"; -import type { ExtensionRunner, SessionBeforeRefineResult } from "../core/extensions/index.js"; +import { formatNoModelSelectedMessage } from "../../core/auth-guidance.js"; +import { serializeConversation } from "../../core/compaction/index.js"; +import type { ExtensionRunner, SessionBeforeRefineResult } from "../../core/extensions/index.js"; import { type CustomMessage, convertToLlm, createRefinementNoticeMessage, createRefinementOutcomeMessage, type RefinementSource, -} from "../core/messages.js"; -import type { ProviderRetryPolicy } from "../core/provider-retry.js"; +} from "../../core/messages.js"; +import type { ProviderRetryPolicy } from "../../core/provider-retry.js"; import { type AutoRefineReview, appendGlobalRefinement, @@ -33,8 +33,8 @@ import { type RefinementResult, reviewAutoRefine, saveHarnessState, -} from "../core/refinement/index.js"; -import type { SessionManager } from "../core/session-manager.js"; +} from "../../core/refinement/index.js"; +import type { SessionManager } from "../../core/session-manager.js"; import type { AutoRefineReviewRequest } from "./auto-refinement.js"; export type SessionRefinementEvent = | { type: "refine_complete"; result: RefinementResult } diff --git a/packages/coding-agent/src/session/refinement.ts b/packages/coding-agent/src/session/refinement/refinement.ts similarity index 98% rename from packages/coding-agent/src/session/refinement.ts rename to packages/coding-agent/src/session/refinement/refinement.ts index 57e316773c..77f75094e0 100644 --- a/packages/coding-agent/src/session/refinement.ts +++ b/packages/coding-agent/src/session/refinement/refinement.ts @@ -1,11 +1,11 @@ import type { AgentMessage, ThinkingLevel } from "@earendil-works/pi-agent-core"; import type { Api, Model } from "@earendil-works/pi-ai"; -import type { ExtensionRunner } from "../core/extensions/index.js"; -import type { CustomMessage, RefinementSource } from "../core/messages.js"; -import type { ProviderRetryPolicy } from "../core/provider-retry.js"; -import type { HarnessState, RefinementPlan, RefinementResult } from "../core/refinement/index.js"; -import type { SessionManager } from "../core/session-manager.js"; -import type { SettingsManager } from "../core/settings-manager.js"; +import type { ExtensionRunner } from "../../core/extensions/index.js"; +import type { CustomMessage, RefinementSource } from "../../core/messages.js"; +import type { ProviderRetryPolicy } from "../../core/provider-retry.js"; +import type { HarnessState, RefinementPlan, RefinementResult } from "../../core/refinement/index.js"; +import type { SessionManager } from "../../core/session-manager.js"; +import type { SettingsManager } from "../../core/settings-manager.js"; import { AutoRefinement, type AutoRefineReviewer, autoRefineInstructions } from "./auto-refinement.js"; import { RefinementExecution, RefineSkippedError, type SessionRefinementEvent } from "./refinement-execution.js"; diff --git a/packages/coding-agent/src/session/bash.ts b/packages/coding-agent/src/session/tools/bash.ts similarity index 97% rename from packages/coding-agent/src/session/bash.ts rename to packages/coding-agent/src/session/tools/bash.ts index ee998cbc60..0252ab093b 100644 --- a/packages/coding-agent/src/session/bash.ts +++ b/packages/coding-agent/src/session/tools/bash.ts @@ -1,7 +1,7 @@ -import { type BashResult, executeBashWithOperations } from "../core/bash-executor.js"; -import type { UserBashEvent, UserBashEventResult } from "../core/extensions/types.js"; -import type { BashExecutionMessage } from "../core/messages.js"; -import { type BashOperations, createLocalBashOperations } from "../core/tools/bash.js"; +import { type BashResult, executeBashWithOperations } from "../../core/bash-executor.js"; +import type { UserBashEvent, UserBashEventResult } from "../../core/extensions/types.js"; +import type { BashExecutionMessage } from "../../core/messages.js"; +import { type BashOperations, createLocalBashOperations } from "../../core/tools/bash.js"; export interface ExecuteBashOptions { excludeFromContext?: boolean; diff --git a/packages/coding-agent/src/session/tools.ts b/packages/coding-agent/src/session/tools/tools.ts similarity index 96% rename from packages/coding-agent/src/session/tools.ts rename to packages/coding-agent/src/session/tools/tools.ts index 540868e51d..96fc146d1e 100644 --- a/packages/coding-agent/src/session/tools.ts +++ b/packages/coding-agent/src/session/tools/tools.ts @@ -4,16 +4,16 @@ import { 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"; +} 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; diff --git a/packages/coding-agent/src/session/turns/autonomous-continuation.ts b/packages/coding-agent/src/session/turns/autonomous-continuation.ts new file mode 100644 index 0000000000..6915ceec8d --- /dev/null +++ b/packages/coding-agent/src/session/turns/autonomous-continuation.ts @@ -0,0 +1,387 @@ +import type { Agent, AgentMessage } from "@earendil-works/pi-agent-core"; +import type { AssistantMessage, Usage } from "@earendil-works/pi-ai"; +import type { AgentSessionEvent } from "../../core/agent-session.js"; +import { + type AgentAutonomousConfig, + type AgentAutonomousStatus, + type AutonomousRuntimeState, + addAutonomousContinuation, + addAutonomousUsage, + autonomousStatus, + createAutonomousRuntimeState, + isUnlimitedAutonomousLimit, + nextAutonomousContinuation, + refreshAutonomousQualityGates, + setAutonomousEnabled, + setAutonomousLimits, + UNLIMITED_AUTONOMOUS_LIMIT, +} from "../../core/autonomous.js"; +import type { CustomMessage } from "../../core/messages.js"; +import { parseCommandArgs } from "../../core/prompt-templates.js"; +import type { SessionManager } from "../../core/session-manager.js"; +import { parseSessionSlashCommand } from "../../core/slash-commands.js"; +import type { SessionCompaction } from "../compaction/compaction.js"; +import type { SessionInputAdmission } from "../input/input-admission.js"; +import { createPreparedTurnAction, primaryDeliveryRecord, type QueuedSessionAction } from "../prepared-actions.js"; +import type { SessionContinuation } from "./continuation.js"; + +type AutonomousSlashCommand = { kind: "status" } | { kind: "on"; config?: AgentAutonomousConfig } | { kind: "off" }; +type AutonomousRuntimeSnapshot = Pick< + AutonomousRuntimeState, + "continuationsUsed" | "gateAttempts" | "lastGateFailure" | "lastGateFailureSnapshot" +>; + +const AUTONOMOUS_STATUS_NUMBER_FORMAT = new Intl.NumberFormat("en-US"); + +const AUTONOMOUS_BUDGET_USAGE = + "Usage: /autonomous [status|off] or /autonomous on [--max-continuations ] [--max-turns ] [--max-tokens ] [--timeout-ms ] [--gate ] [--gate-retries ] [--gate-timeout-ms ]"; + +// `/autonomous` budget flags mirror the `--autonomous-*` CLI options. The CLI +// spelling (`--autonomous-max-continuations`) is accepted as an alias so the +// exact CLI budget flags also work from the slash command. +const AUTONOMOUS_BUDGET_FLAGS: ReadonlySet = new Set([ + "max-continuations", + "max-turns", + "max-tokens", + "timeout-ms", + "gate", + "gate-retries", + "gate-timeout-ms", +]); + +function parseAutonomousBudgetInt(flag: string, value: string, allowUnlimited = false): number { + if (allowUnlimited && value.toLowerCase() === "unlimited") { + return UNLIMITED_AUTONOMOUS_LIMIT; + } + // Commas and underscores are accepted as digit separators (100,000,000). + const digits = value.replace(/[,_]/g, ""); + if (!/^[1-9]\d*$/.test(digits)) { + throw new Error( + `--${flag} must be a positive integer${allowUnlimited ? ' or "unlimited"' : ""}. ${AUTONOMOUS_BUDGET_USAGE}`, + ); + } + return Number(digits); +} + +function parseAutonomousBudgetOptions(tokens: string[]): AgentAutonomousConfig { + const config: AgentAutonomousConfig = {}; + const gateCommands: string[] = []; + for (let i = 0; i < tokens.length; i++) { + const token = tokens[i]!; + if (!token.startsWith("--")) { + throw new Error(`Unexpected autonomous argument: ${token}. ${AUTONOMOUS_BUDGET_USAGE}`); + } + const equalsIndex = token.indexOf("="); + const rawFlag = equalsIndex === -1 ? token : token.slice(0, equalsIndex); + const inlineValue = equalsIndex === -1 ? undefined : token.slice(equalsIndex + 1); + const flag = rawFlag.startsWith("--autonomous-") ? rawFlag.slice("--autonomous-".length) : rawFlag.slice(2); + if (!AUTONOMOUS_BUDGET_FLAGS.has(flag)) { + throw new Error(`Unknown autonomous budget flag: ${rawFlag}. ${AUTONOMOUS_BUDGET_USAGE}`); + } + let value = inlineValue; + if (value === undefined) { + const next = tokens[i + 1]; + if (next === undefined || next.startsWith("--")) { + throw new Error(`Missing value for ${rawFlag}. ${AUTONOMOUS_BUDGET_USAGE}`); + } + value = next; + i++; + } + if (value === "") { + throw new Error(`Missing value for ${rawFlag}. ${AUTONOMOUS_BUDGET_USAGE}`); + } + switch (flag) { + case "gate": + gateCommands.push(value); + break; + case "gate-retries": + config.gates = config.gates ?? {}; + config.gates.maxRetries = parseAutonomousBudgetInt(flag, value); + break; + case "gate-timeout-ms": + config.gates = config.gates ?? {}; + config.gates.timeoutMs = parseAutonomousBudgetInt(flag, value); + break; + case "max-continuations": + config.maxContinuations = parseAutonomousBudgetInt(flag, value, true); + break; + case "max-turns": + config.maxTurns = parseAutonomousBudgetInt(flag, value, true); + break; + case "max-tokens": + config.maxTokens = parseAutonomousBudgetInt(flag, value, true); + break; + case "timeout-ms": + config.timeoutMs = parseAutonomousBudgetInt(flag, value, true); + break; + } + } + if (gateCommands.length > 0) { + config.gates = { ...config.gates, commands: gateCommands }; + } + // Named budget flags define the whole budget: any limit the user did not + // name stops cutting the run short. With no budget flags at all, the + // configured or default limits still apply. + if ( + config.maxContinuations !== undefined || + config.maxTurns !== undefined || + config.maxTokens !== undefined || + config.timeoutMs !== undefined + ) { + config.maxContinuations ??= UNLIMITED_AUTONOMOUS_LIMIT; + config.maxTurns ??= UNLIMITED_AUTONOMOUS_LIMIT; + config.maxTokens ??= UNLIMITED_AUTONOMOUS_LIMIT; + config.timeoutMs ??= UNLIMITED_AUTONOMOUS_LIMIT; + } + return config; +} + +export interface SessionAutonomousContinuationHost { + getStatus(): AgentAutonomousStatus; + getCwd(): string; + getAgent(): Pick; + getStore(): Pick; + emit(event: AgentSessionEvent): void; + getContinuation(): Pick; + getArrivalEpoch(): number; + admit: SessionInputAdmission["admitSessionInput"]; + cancelActions(predicate: (action: QueuedSessionAction) => boolean, error: Error): QueuedSessionAction[]; + emitQueueUpdate(): void; + getCompaction(): Pick; + getUnfinishedActionCount(): number; + cancelContinuation(): void; +} +export class SessionAutonomousContinuation { + private readonly state: AutonomousRuntimeState; + private suppressionDepth = 0; + private readonly suppressedMessages = new WeakSet(); + private readonly thresholdContinuations = new WeakMap(); + private readonly snapshots = new WeakMap(); + private pendingThresholdMessages: AgentMessage[] = []; + constructor( + config: AgentAutonomousConfig | undefined, + private readonly host: SessionAutonomousContinuationHost, + ) { + this.state = createAutonomousRuntimeState(config, { cwd: host.getCwd() }); + } + forgetSnapshot(message: AgentMessage): void { + this.snapshots.delete(message); + } + + takePendingThresholdMessages(): AgentMessage[] { + return this.pendingThresholdMessages.splice(0); + } + + recordUsage(usage: Usage): void { + addAutonomousUsage(this.state, usage); + } + isSuppressed(messages: AgentMessage[]): boolean { + return this.suppressionDepth > 0 || messages.some((message) => this.suppressedMessages.has(message)); + } + next(message: AssistantMessage, signal?: AbortSignal): Promise { + return nextAutonomousContinuation(this.state, message, { cwd: this.host.getCwd(), signal }); + } + + parseAutonomousSlashCommand(text: string): AutonomousSlashCommand | undefined { + const command = parseSessionSlashCommand(text); + if (command?.name !== "autonomous") return undefined; + const tokens = parseCommandArgs(command.args); + if (tokens.length === 0 || tokens[0]!.toLowerCase() === "status") { + if (tokens.length > 1) { + throw new Error(`Unexpected autonomous argument: ${tokens[1]}. ${AUTONOMOUS_BUDGET_USAGE}`); + } + return { kind: "status" }; + } + const subcommand = tokens[0]!.toLowerCase(); + if (subcommand === "on" || subcommand === "enable" || subcommand === "enabled") { + return { kind: "on", config: parseAutonomousBudgetOptions(tokens.slice(1)) }; + } + if (subcommand === "off" || subcommand === "disable" || subcommand === "disabled") { + if (tokens.length > 1) { + throw new Error(`Unexpected autonomous argument: ${tokens[1]}. ${AUTONOMOUS_BUDGET_USAGE}`); + } + return { kind: "off" }; + } + throw new Error(AUTONOMOUS_BUDGET_USAGE); + } + + formatAutonomousStatus(): string { + const status = this.host.getStatus(); + const state = status.enabled ? "on" : "off"; + const elapsedSeconds = status.startedAt ? Math.round((Date.now() - status.startedAt) / 1000) : 0; + const gateSummary = + status.gates.commands.length > 0 ? status.gates.commands.map((command) => `"${command}"`).join(", ") : "none"; + const formatCount = (value: number): string => + isUnlimitedAutonomousLimit(value) ? "unlimited" : AUTONOMOUS_STATUS_NUMBER_FORMAT.format(value); + const timeBudget = isUnlimitedAutonomousLimit(status.limits.timeoutMs) + ? "unlimited" + : `${AUTONOMOUS_STATUS_NUMBER_FORMAT.format(Math.round(status.limits.timeoutMs / 1000))}s`; + return `[autonomous-status: ${state}]\n\nContinuations: ${formatCount(status.continuationsUsed)}/${formatCount(status.limits.maxContinuations)}. Turns: ${formatCount(status.turnsUsed)}/${formatCount(status.limits.maxTurns)}. Tokens: ${formatCount(status.tokensUsed)}/${formatCount(status.limits.maxTokens)}. Time: ${elapsedSeconds}s/${timeBudget}. Gates: ${gateSummary}.`; + } + + emitAutonomousStatus(): void { + const message = { + role: "custom" as const, + customType: "autonomous_status", + content: this.formatAutonomousStatus(), + display: true, + details: this.host.getStatus(), + timestamp: Date.now(), + } satisfies CustomMessage; + this.host.getAgent().state.messages.push(message); + this.host + .getStore() + .appendCustomMessageEntry(message.customType, message.content, message.display, message.details); + this.host.emit({ type: "message_start", message }); + this.host.emit({ type: "message_end", message }); + } + + async handleAutonomousSlashCommand(text: string): Promise { + const command = this.parseAutonomousSlashCommand(text); + if (!command) { + return false; + } + if (command.kind === "on") { + setAutonomousEnabled(this.state, true, { cwd: this.host.getCwd() }); + setAutonomousLimits(this.state, command.config); + } else if (command.kind === "off") { + setAutonomousEnabled(this.state, false); + this.clearQueuedAutonomousContinuations(); + } + this.emitAutonomousStatus(); + return true; + } + + snapshotAutonomousRuntimeState(): AutonomousRuntimeSnapshot { + return { + continuationsUsed: this.state.continuationsUsed, + gateAttempts: { ...this.state.gateAttempts }, + lastGateFailure: this.state.lastGateFailure ? { ...this.state.lastGateFailure } : undefined, + lastGateFailureSnapshot: this.state.lastGateFailureSnapshot + ? { ...this.state.lastGateFailureSnapshot } + : undefined, + }; + } + + restoreAutonomousRuntimeSnapshot(snapshot: AutonomousRuntimeSnapshot): void { + this.state.continuationsUsed = snapshot.continuationsUsed; + this.state.gateAttempts = { ...snapshot.gateAttempts }; + this.state.lastGateFailure = snapshot.lastGateFailure ? { ...snapshot.lastGateFailure } : undefined; + this.state.lastGateFailureSnapshot = snapshot.lastGateFailureSnapshot + ? { ...snapshot.lastGateFailureSnapshot } + : undefined; + } + + async queueAutonomousContinuationForThresholdCompaction( + message: AssistantMessage, + ): Promise { + const queuedMessage = this.thresholdContinuations.get(message); + if (queuedMessage && this.host.getContinuation().messages.includes(queuedMessage)) { + return queuedMessage; + } + const snapshot = this.snapshotAutonomousRuntimeState(); + const arrivalEpoch = this.host.getArrivalEpoch(); + const autonomousMessage = await nextAutonomousContinuation(this.state, message, { + cwd: this.host.getCwd(), + signal: this.host.getAgent().signal, + }); + if (!autonomousMessage) { + return undefined; + } + if (this.host.getArrivalEpoch() !== arrivalEpoch) { + this.restoreAutonomousRuntimeSnapshot(snapshot); + return undefined; + } + this.thresholdContinuations.set(message, autonomousMessage); + this.snapshots.set(autonomousMessage, snapshot); + this.host.getContinuation().track(autonomousMessage); + this.pendingThresholdMessages.push(autonomousMessage); + const text = + typeof autonomousMessage.content === "string" + ? autonomousMessage.content + : autonomousMessage.content.map((block) => (block.type === "text" ? block.text : "")).join("\n"); + this.host.admit( + createPreparedTurnAction("followUp", text, undefined, { + message: autonomousMessage, + }), + ); + return autonomousMessage; + } + + clearQueuedAutonomousContinuations( + options: { restoreAutonomousState?: boolean; messages?: AgentMessage[] } = {}, + ): void { + const requestedMessages = options.messages ?? [...this.host.getContinuation().messages]; + const requestedMessageSet = new Set(requestedMessages); + const queuedMessages = this.host.getContinuation().messages.filter((message) => requestedMessageSet.has(message)); + if (queuedMessages.length === 0) { + return; + } + const queuedMessageSet = new Set(queuedMessages); + this.host.getContinuation().remove(queuedMessageSet); + this.host.getAgent().removeQueuedMessages((message) => queuedMessageSet.has(message)); + this.host.cancelActions( + (action) => action.payload.kind === "turn" && queuedMessageSet.has(primaryDeliveryRecord(action).message), + new Error("Queued autonomous continuation was cleared before delivery."), + ); + this.host.emitQueueUpdate(); + if (options.restoreAutonomousState) { + for (const queuedMessage of queuedMessages) { + const snapshot = this.snapshots.get(queuedMessage); + if (snapshot) { + this.restoreAutonomousRuntimeSnapshot(snapshot); + break; + } + } + } + for (const queuedMessage of queuedMessages) { + this.snapshots.delete(queuedMessage); + } + this.pendingThresholdMessages = this.pendingThresholdMessages.filter((message) => !queuedMessageSet.has(message)); + if (options.messages === undefined) { + this.host.getCompaction().resetContinuation(); + } + if (!this.host.getAgent().hasQueuedMessages() && this.host.getUnfinishedActionCount() === 0) { + this.host.cancelContinuation(); + } + } + + clearQueuedAutonomousContinuationsAfterSkippedThresholdCompaction( + shouldContinueAfterThreshold: boolean, + queuedMessages: AgentMessage[], + ): void { + if (shouldContinueAfterThreshold) { + this.clearQueuedAutonomousContinuations({ + restoreAutonomousState: true, + messages: queuedMessages, + }); + } + } + + getAutonomousStatus(): AgentAutonomousStatus { + return autonomousStatus(this.state); + } + + recordHostAutonomousContinuation(): void { + addAutonomousContinuation(this.state); + } + + async refreshAutonomousGates(): Promise { + await refreshAutonomousQualityGates(this.state, { + cwd: this.host.getCwd(), + }); + } + + async runWithAutonomousContinuationSuppressed(fn: () => Promise): Promise { + this.suppressionDepth++; + try { + return await fn(); + } finally { + this.suppressionDepth--; + } + } + + markAutonomousContinuationSuppressed(message: AgentMessage): void { + this.suppressedMessages.add(message); + } +} diff --git a/packages/coding-agent/src/session/turns/command-execution.ts b/packages/coding-agent/src/session/turns/command-execution.ts new file mode 100644 index 0000000000..399507811c --- /dev/null +++ b/packages/coding-agent/src/session/turns/command-execution.ts @@ -0,0 +1,200 @@ +import type { Agent } from "@earendil-works/pi-agent-core"; +import type { ImageContent } from "@earendil-works/pi-ai"; +import type { AgentSessionEvent } from "../../core/agent-session.js"; +import type { CompactionResult } from "../../core/compaction/index.js"; +import { + type CustomMessage, + createSessionSlashCommandMessage, + createSessionSlashCommandResultMessage, +} from "../../core/messages.js"; +import type { RefinementResult } from "../../core/refinement/index.js"; +import { + type ActionStore, + canSelectSessionAction, + type RuntimeActivity, + transitionSessionAction, +} from "../../core/session-action-store.js"; +import type { SessionManager } from "../../core/session-manager.js"; +import { parseRefineCommandOptions, type SessionSlashCommand } from "../../core/slash-commands.js"; +import { CompactionSkippedError } from "../compaction/compaction-execution.js"; +import type { GoalState } from "../goals/contracts.js"; +import type { SessionCommitFence, SessionCommitLease } from "../input/commit-fence.js"; +import type { QueuedSessionAction } from "../prepared-actions.js"; +import type { SessionRefinement } from "../refinement/refinement.js"; + +function asError(error: unknown): Error { + return error instanceof Error ? error : new Error(String(error)); +} +export interface SessionCommandExecutionHost { + getFence(): Pick; + acquireFence(): Promise; + getRefinement(): Pick; + isDeferred(epoch: number): boolean; + getActivity(): RuntimeActivity; + notifyCheckpoints(): void; + emitQueueUpdate(): void; + settleAgentMessage(id: string | undefined, leg: "delivery" | "completion", error?: Error): void; + rejectAgentMessage(id: string | undefined, error: Error): void; + compact(instructions?: string, options?: { skipAbort?: boolean }): Promise; + refine( + options: { instructions?: string; rollbackId?: string; global?: boolean }, + internal: { skipAbort?: boolean }, + ): Promise; + handleGoalCommand(text: string, images: ImageContent[] | undefined): Promise; + handleAutonomousCommand(text: string): Promise; + getGoalState(): GoalState; + getStore(): Pick; + getAgent(): Pick; + emit(event: AgentSessionEvent): void; +} +export class SessionCommandExecution { + constructor( + private readonly actions: ActionStore, + private readonly host: SessionCommandExecutionHost, + ) {} + async executeSelectedSessionCommand(action: QueuedSessionAction, epoch: number): Promise { + if (action.payload.kind !== "session_command") throw new Error("Expected a selected session command"); + const input = action.payload; + const commitFence = await this.host.acquireFence(); + try { + await this.host.getFence().run(commitFence, async () => { + const isCancelled = () => action.lifecycle.state === "cancelled"; + if (isCancelled()) return; + await this.host.getRefinement()._waitForRefineIdle(); + if (isCancelled()) return; + if (this.host.isDeferred(epoch) || !canSelectSessionAction(this.host.getActivity())) { + this.actions.rollback(action); + this.host.notifyCheckpoints(); + this.host.emitQueueUpdate(); + return; + } + transitionSessionAction(action, { + state: "running", + execution: "session_command", + }); + this.host.notifyCheckpoints(); + this.host.emitQueueUpdate(); + try { + this.appendDurableSessionCommandMessage(input.text, input.command, false); + this.actions.ticketFor(action).settleDelivered({ status: "not_applicable" }); + this.host.settleAgentMessage(action.agentMessageId, "delivery"); + await this.executeQueuedSessionCommand(action); + transitionSessionAction(action, { state: "completed" }); + this.actions.ticketFor(action).settleCompleted(); + this.host.settleAgentMessage(action.agentMessageId, "completion"); + } catch (error) { + const commandError = asError(error); + transitionSessionAction(action, { + state: "failed", + error: commandError, + }); + const ticket = this.actions.ticketFor(action); + ticket.rejectDelivered(commandError); + ticket.settleCompleted(commandError); + this.host.rejectAgentMessage(action.agentMessageId, commandError); + } finally { + this.actions.releaseTerminal(action); + this.host.notifyCheckpoints(); + this.host.emitQueueUpdate(); + } + }); + } finally { + commitFence.release(); + } + } + + async executeQueuedSessionCommand(action: QueuedSessionAction): Promise { + if (action.payload.kind !== "session_command") throw new Error("Expected a session command action"); + const input = action.payload; + try { + let resultText: string | undefined; + let displayResult = true; + switch (input.command.name) { + case "compact": + await this.host.compact(input.command.args || undefined, { + skipAbort: true, + }); + break; + case "refine": { + let result: RefinementResult; + try { + const options = parseRefineCommandOptions(input.command.args); + result = await this.host.refine(options, { skipAbort: true }); + } catch (error) { + // Only a failure of the refinement itself is a refine failure; a later + // result-row persist error must not report a completed refinement as failed. + this.host.getRefinement()._emitRefineFailed(asError(error)); + throw error; + } + const applied = result.appliedEdits.filter((edit) => edit.applied).length; + resultText = `Refined continual harness state: ${applied} edit${applied === 1 ? "" : "s"} applied.`; + displayResult = false; + break; + } + case "goal": + await this.host.handleGoalCommand(input.text, input.images); + resultText = this.host.getGoalState().objective + ? `Goal ${this.host.getGoalState().status}: ${this.host.getGoalState().objective}` + : "No active goal."; + break; + case "autonomous": + await this.host.handleAutonomousCommand(input.text); + break; + } + if (resultText) { + this.appendDurableSessionCommandMessage(resultText, input.command, true, false, displayResult); + } + } catch (error) { + if (error instanceof CompactionSkippedError) return; + const commandError = error instanceof Error ? error : new Error(String(error)); + try { + this.appendDurableSessionCommandMessage( + `Command failed: ${commandError.message}`, + input.command, + true, + true, + ); + } catch { + // The result row is also the command-correlated UI settle edge. + const message = createSessionSlashCommandResultMessage(`Command failed: ${commandError.message}`, { + command: input.command, + success: false, + severity: "error", + error: commandError.message, + }); + this.host.emit({ type: "message_start", message }); + this.host.emit({ type: "message_end", message }); + } + throw commandError; + } + } + + appendDurableSessionCommandMessage( + content: string, + command: SessionSlashCommand, + isResult: boolean, + isError = false, + display = true, + ): void { + const message: CustomMessage = isResult + ? createSessionSlashCommandResultMessage( + content, + { + command, + success: !isError, + severity: isError ? "error" : "info", + ...(isError ? { error: content.replace(/^Command failed:\s*/, "") } : {}), + }, + display, + ) + : createSessionSlashCommandMessage(command); + // Persist before touching live state so a failed write cannot leave an + // unsaved leaf that the next entry would silently parent onto. + this.host + .getStore() + .appendCustomMessageEntryWithRollback(message.customType, message.content, message.display, message.details); + this.host.getAgent().state.messages.push(message); + this.host.emit({ type: "message_start", message }); + this.host.emit({ type: "message_end", message }); + } +} diff --git a/packages/coding-agent/src/session/continuation.ts b/packages/coding-agent/src/session/turns/continuation.ts similarity index 99% rename from packages/coding-agent/src/session/continuation.ts rename to packages/coding-agent/src/session/turns/continuation.ts index a68e2e0d6a..27c3b2187a 100644 --- a/packages/coding-agent/src/session/continuation.ts +++ b/packages/coding-agent/src/session/turns/continuation.ts @@ -1,5 +1,5 @@ import { AgentContinueError, type AgentMessage } from "@earendil-works/pi-agent-core"; -import type { SessionCommitLease } from "./commit-fence.js"; +import type { SessionCommitLease } from "../input/commit-fence.js"; export interface ContinuationToken { readonly promise: Promise; diff --git a/packages/coding-agent/src/session/turns/events.ts b/packages/coding-agent/src/session/turns/events.ts new file mode 100644 index 0000000000..89dd609eca --- /dev/null +++ b/packages/coding-agent/src/session/turns/events.ts @@ -0,0 +1,495 @@ +import type { Agent, AgentEvent, AgentMessage, ThinkingLevel } from "@earendil-works/pi-agent-core"; +import type { AssistantMessage, ServiceTier, Usage } from "@earendil-works/pi-ai"; +import { startsAgentRun } from "../../core/agent-messages.js"; +import { addLoginGuidanceToAuthError, isLikelyAuthenticationError } from "../../core/auth-guidance.js"; +import type { + ExtensionRunner, + MessageEndEvent, + MessageStartEvent, + MessageUpdateEvent, + ToolExecutionEndEvent, + ToolExecutionStartEvent, + ToolExecutionUpdateEvent, + TurnEndEvent, + TurnStartEvent, +} from "../../core/extensions/index.js"; +import type { KernelSentAgentMessage } from "../../core/kernel/index.js"; +import type { RefinementResult } from "../../core/refinement/index.js"; +import { + type ActionStore, + type SessionActionSnapshot, + transitionSessionAction, +} from "../../core/session-action-store.js"; +import type { SessionManager } from "../../core/session-manager.js"; +import type { RlmChildAgentSnapshot } from "../children/child-types.js"; +import type { SessionCompaction, SessionCompactionEvent } from "../compaction/compaction.js"; +import type { GoalState } from "../goals/contracts.js"; +import { primaryDeliveryRecord, type QueuedSessionAction } from "../prepared-actions.js"; +import type { SessionRefinement } from "../refinement/refinement.js"; +import type { SessionBashEvent } from "../tools/bash.js"; +import type { SessionRetry, SessionRetryEvent } from "./retry.js"; + +export type AgentSessionEvent = + | AgentEvent + | { + type: "ipython_sent_agent_message"; + toolCallId: string; + message: KernelSentAgentMessage; + } + | { type: "session_action_update"; actions: SessionActionSnapshot } + | SessionCompactionEvent + | { type: "session_info_changed"; name: string | undefined } + | { type: "thinking_level_changed"; level: ThinkingLevel } + | { type: "service_tier_changed"; serviceTier: ServiceTier } + | SessionRetryEvent + | { type: "rlm_child_update"; child: RlmChildAgentSnapshot } + | { type: "recap_update"; recap: string | undefined } + | { type: "goal_update"; goal: GoalState } + | SessionBashEvent + | { type: "refine_complete"; result: RefinementResult } + | { type: "refine_failed"; error: string }; + +export type AgentSessionEventListener = (event: AgentSessionEvent) => void; + +export interface SessionEventsHost { + getAgent(): Pick & { state: { messages: AgentMessage[]; errorMessage?: string } }; + getStore(): Pick; + getExtensions(): Pick; + getRetry(): Pick< + SessionRetry, + | "observeAgentEnd" + | "resolve" + | "observeAssistantEnd" + | "isRetrying" + | "retryError" + | "attempt" + | "finishActiveRetryWithFailure" + >; + getCompaction(): Pick; + getRefinement(): Pick< + SessionRefinement, + "observeAssistantEnd" | "serialized" | "_consumePendingRequestedRefine" | "_scheduleAutoRefineAfterAgentEnd" + >; + addAutonomousUsage(usage: Usage): void; + applyLateMessages(message: AgentMessage): void; + notifyCheckpoints(): void; + settleAgentMessage(id: string | undefined, leg: "delivery" | "completion", error?: Error): void; + getSnapshot(): SessionActionSnapshot; + accountAssistantBudget(message: AssistantMessage): Promise | undefined; + finishGoal(message: AssistantMessage): void; + checkCompaction(message: AssistantMessage): Promise; +} +export class SessionEvents { + private listeners: AgentSessionEventListener[] = []; + private lastSnapshot: SessionActionSnapshot = { queuedCount: 0, steering: [], followUps: [] }; + private _queue: Promise = Promise.resolve(); + get queue(): Promise { + return this._queue; + } + private lastAssistant: AssistantMessage | undefined; + private turnIndex = 0; + private unsubscribeAgent?: () => void; + constructor( + private readonly actions: ActionStore, + private readonly host: SessionEventsHost, + ) {} + enqueue(work: () => void): void { + this._queue = this._queue.then(work, work); + this._queue.catch(() => {}); + } + dispose(): void { + this.disconnectFromAgent(); + this.listeners = []; + } + + emit(event: AgentSessionEvent): void { + for (const l of this.listeners) { + try { + l(event); + } catch { + // A failing observer must not prevent other subscribers from + // receiving lifecycle and persistence events. + } + } + } + + emitQueueUpdate(): void { + const actions = this.host.getSnapshot(); + if (JSON.stringify(actions) === JSON.stringify(this.lastSnapshot)) return; + this.lastSnapshot = actions; + this.emit({ type: "session_action_update", actions }); + } + + capturingCancelledAction(message: AgentMessage): QueuedSessionAction | undefined { + return this.actions + .ownedActions() + .find( + (action) => + action.lifecycle.state === "cancelled" && + action.payload.kind === "turn" && + action.payload.captureRunMessages?.has(message) === true, + ); + } + + hasCancelledDispatchCapture(): boolean { + return this.actions + .ownedActions() + .some( + (action) => + action.lifecycle.state === "cancelled" && + action.payload.kind === "turn" && + action.payload.captureRunMessages !== undefined, + ); + } + + handleAgentEvent = (event: AgentEvent): void => { + this.host.getRetry().observeAgentEnd(event); + if (event.type === "message_start" || event.type === "message_end") { + for (const action of this.actions.ownedActions()) { + if ( + action.payload.kind !== "turn" || + !action.payload.captureRunMessages || + action.payload.cancelledDispatchEnded + ) { + continue; + } + const primary = primaryDeliveryRecord(action); + if (event.message === primary.message || primary.started) { + action.payload.captureRunMessages.add(event.message); + } + } + } else if (event.type === "agent_end") { + const captured = new Set(); + for (const action of this.actions.ownedActions()) { + if (action.payload.kind === "turn" && action.payload.captureRunMessages) { + for (const message of action.payload.captureRunMessages) captured.add(message); + action.payload.cancelledDispatchEnded = true; + } + } + if (captured.size > 0) { + this.host.getAgent().state.messages = this.host + .getAgent() + .state.messages.filter((message) => !captured.has(message)); + } + } + if (event.type === "message_start" && (event.message.role === "user" || event.message.role === "custom")) { + for (const action of this.actions.actionsForMessage(event.message)) { + const record = + action.payload.kind === "turn" + ? action.payload.records.find((candidate) => candidate.message === event.message) + : undefined; + if (record) record.started = true; + if (record?.role === "primary") { + this.actions.ticketFor(action).settleDelivered({ status: "delivered" }); + this.host.settleAgentMessage(action.agentMessageId, "delivery"); + } + } + } else if (event.type === "message_end" && (event.message.role === "user" || event.message.role === "custom")) { + for (const action of this.actions.actionsForMessage(event.message)) { + const record = + action.payload.kind === "turn" + ? action.payload.records.find((candidate) => candidate.message === event.message) + : undefined; + if (record) record.durable = true; + if (record?.role === "primary" && action.lifecycle.state === "committing") { + transitionSessionAction(action, { + state: "running", + execution: "agent_turn", + }); + this.host.notifyCheckpoints(); + this.emitQueueUpdate(); + } + } + } + this._queue = this._queue.then( + () => this.processAgentEvent(event), + () => this.processAgentEvent(event), + ); + this._queue.catch(() => {}); + }; + + findLastAssistantInMessages(messages: AgentMessage[]): AssistantMessage | undefined { + for (let i = messages.length - 1; i >= 0; i--) { + const message = messages[i]; + if (message.role === "assistant") { + return message as AssistantMessage; + } + } + return undefined; + } + + addLoginGuidanceToAuthError(event: AgentEvent): void { + const message = + event.type === "message_end" && event.message.role === "assistant" + ? (event.message as AssistantMessage) + : event.type === "agent_end" + ? this.findLastAssistantInMessages(event.messages) + : undefined; + if (!message || message.stopReason !== "error" || !message.errorMessage) { + return; + } + if (!isLikelyAuthenticationError(message.errorMessage)) { + return; + } + message.errorMessage = addLoginGuidanceToAuthError(message.errorMessage); + } + + async processAgentEvent(event: AgentEvent): Promise { + let clearedDispatchEnded = false; + if ((event.type === "message_start" || event.type === "message_end") && event.message.role === "toolResult") { + this.host.applyLateMessages(event.message); + } + if (event.type === "message_start" || event.type === "message_end") { + const cleared = this.capturingCancelledAction(event.message); + if (cleared?.payload.kind === "turn" && cleared.payload.captureRunMessages) { + const captured = cleared.payload.captureRunMessages; + this.host.getAgent().state.messages = this.host + .getAgent() + .state.messages.filter((message) => !captured.has(message)); + return; + } + } + if (event.type === "agent_end") { + const cleared = this.actions + .ownedActions() + .filter( + (action) => + action.lifecycle.state === "cancelled" && + action.payload.kind === "turn" && + action.payload.captureRunMessages !== undefined, + ); + if (cleared.length > 0) { + clearedDispatchEnded = true; + const removed = new Set( + cleared.flatMap((action) => + action.payload.kind === "turn" ? [...(action.payload.captureRunMessages ?? [])] : [], + ), + ); + this.host.getAgent().state.messages = this.host + .getAgent() + .state.messages.filter((message) => !removed.has(message)); + this.host.getAgent().state.errorMessage = undefined; + this.lastAssistant = undefined; + for (const action of cleared) this.actions.releaseTerminal(action); + this.host.notifyCheckpoints(); + this.host.getRetry().resolve(); + } + } + + if (event.type === "message_start" && startsAgentRun(event.message)) { + this.host.getCompaction().resetOverflowRecovery(); + } + + await this.emitExtensionEvent(event); + if (event.type === "message_start" || event.type === "message_end") { + const cleared = this.capturingCancelledAction(event.message); + if (cleared?.payload.kind === "turn" && cleared.payload.captureRunMessages) { + const captured = cleared.payload.captureRunMessages; + this.host.getAgent().state.messages = this.host + .getAgent() + .state.messages.filter((message) => !captured.has(message)); + return; + } + } + + this.addLoginGuidanceToAuthError(event); + + this.emit(event); + + if (event.type === "message_end") { + if (event.message.role === "custom") { + this.host + .getStore() + .appendCustomMessageEntry( + event.message.customType, + event.message.content, + event.message.display, + event.message.details, + ); + } else if ( + event.message.role === "user" || + event.message.role === "assistant" || + event.message.role === "toolResult" + ) { + this.host.getStore().appendMessage(event.message); + } + + if (event.message.role === "assistant") { + this.lastAssistant = event.message; + + const assistantMsg = event.message as AssistantMessage; + if (assistantMsg.stopReason !== "error") { + this.host.addAutonomousUsage(assistantMsg.usage); + } + if (assistantMsg.stopReason !== "error" && assistantMsg.stopReason !== "aborted") { + this.host.getRefinement().observeAssistantEnd(); + // In serialized mode, kick off background refinement planning + // immediately after the primary stream finishes, while tools + // are still executing. The plan is awaited at shouldStopAfterTurn + // before applying, so planning overlaps tools only — never another + // model request. + } + if (assistantMsg.stopReason !== "error") { + this.host.getCompaction().resetOverflowRecovery(); + } + this.host.getRetry().observeAssistantEnd(assistantMsg); + const budgetNotice = this.host.accountAssistantBudget(assistantMsg); + if (budgetNotice) await budgetNotice; + } + } + + if (clearedDispatchEnded) { + return; + } + + if (event.type === "agent_end") { + const msg = + this.lastAssistant ?? + (this.host.getRetry().isRetrying ? this.findLastAssistantInMessages(event.messages) : undefined); + this.lastAssistant = undefined; + if (!msg) { + this.host.getRetry().resolve(); + return; + } + + const retry = this.host.getRetry().retryError(msg); + if (retry && (await retry)) return; + + const compactionWillRetry = await this.host.checkCompaction(msg); + if (compactionWillRetry && this.host.getRetry().attempt > 0) { + return; + } + this.host.getRetry().finishActiveRetryWithFailure(msg); + this.host.getRetry().resolve(); + if (!compactionWillRetry) { + this.host.finishGoal(msg); + // In serialized mode, agent-callable refine.run is serviced + // at the shouldStopAfterTurn boundary, not here at agent_end. + if (!this.host.getRefinement().serialized) { + const consumedRequestedRefine = this.host.getRefinement()._consumePendingRequestedRefine(); + if (!consumedRequestedRefine) { + this.host.getRefinement()._scheduleAutoRefineAfterAgentEnd(); + } + } + } + } + } + + replaceMessageInPlace(target: AgentMessage, replacement: AgentMessage): void { + // Agent-core stores the finalized message object in its state before emitting message_end. + // SessionManager persistence happens later in _processAgentEvent() with event.message. + // Mutating this object in place keeps agent state, later turn/agent events, listeners, + // and the eventual SessionManager.appendMessage(event.message) persistence in sync. + if (target === replacement) { + return; + } + + const targetRecord = target as unknown as Record; + for (const key of Object.keys(targetRecord)) { + delete targetRecord[key]; + } + Object.assign(targetRecord, replacement); + } + + async emitExtensionEvent(event: AgentEvent): Promise { + if (event.type === "agent_start") { + this.turnIndex = 0; + this.host.getStore().recordGitStateIfChanged(); + await this.host.getExtensions().emit({ type: "agent_start" }); + } else if (event.type === "agent_end") { + // Also capture at end of turn so commits made during the run (e.g. via a bash tool) land. + this.host.getStore().recordGitStateIfChanged(); + await this.host.getExtensions().emit({ + type: "agent_end", + messages: event.messages, + }); + } else if (event.type === "turn_start") { + const extensionEvent: TurnStartEvent = { + type: "turn_start", + turnIndex: this.turnIndex, + timestamp: Date.now(), + }; + await this.host.getExtensions().emit(extensionEvent); + } else if (event.type === "turn_end") { + const extensionEvent: TurnEndEvent = { + type: "turn_end", + turnIndex: this.turnIndex, + message: event.message, + toolResults: event.toolResults, + }; + await this.host.getExtensions().emit(extensionEvent); + this.turnIndex++; + } else if (event.type === "message_start") { + const extensionEvent: MessageStartEvent = { + type: "message_start", + message: event.message, + }; + await this.host.getExtensions().emit(extensionEvent); + } else if (event.type === "message_update") { + const extensionEvent: MessageUpdateEvent = { + type: "message_update", + message: event.message, + assistantMessageEvent: event.assistantMessageEvent, + }; + await this.host.getExtensions().emit(extensionEvent); + } else if (event.type === "message_end") { + const extensionEvent: MessageEndEvent = { + type: "message_end", + message: event.message, + }; + const replacement = await this.host.getExtensions().emitMessageEnd(extensionEvent); + if (replacement) { + this.replaceMessageInPlace(event.message, replacement); + } + } else if (event.type === "tool_execution_start") { + const extensionEvent: ToolExecutionStartEvent = { + type: "tool_execution_start", + toolCallId: event.toolCallId, + toolName: event.toolName, + args: event.args, + }; + await this.host.getExtensions().emit(extensionEvent); + } else if (event.type === "tool_execution_update") { + const extensionEvent: ToolExecutionUpdateEvent = { + type: "tool_execution_update", + toolCallId: event.toolCallId, + toolName: event.toolName, + args: event.args, + partialResult: event.partialResult, + }; + await this.host.getExtensions().emit(extensionEvent); + } else if (event.type === "tool_execution_end") { + const extensionEvent: ToolExecutionEndEvent = { + type: "tool_execution_end", + toolCallId: event.toolCallId, + toolName: event.toolName, + result: event.result, + isError: event.isError, + }; + await this.host.getExtensions().emit(extensionEvent); + } + } + + subscribe(listener: AgentSessionEventListener): () => void { + this.listeners.push(listener); + + return () => { + const index = this.listeners.indexOf(listener); + if (index !== -1) { + this.listeners.splice(index, 1); + } + }; + } + + disconnectFromAgent(): void { + if (this.unsubscribeAgent) { + this.unsubscribeAgent(); + this.unsubscribeAgent = undefined; + } + } + + reconnectToAgent(): void { + if (this.unsubscribeAgent) return; // Already connected + this.unsubscribeAgent = this.host.getAgent().subscribe(this.handleAgentEvent); + } +} diff --git a/packages/coding-agent/src/session/retry.ts b/packages/coding-agent/src/session/turns/retry.ts similarity index 97% rename from packages/coding-agent/src/session/retry.ts rename to packages/coding-agent/src/session/turns/retry.ts index cfa92ae06a..9ed2fcadc9 100644 --- a/packages/coding-agent/src/session/retry.ts +++ b/packages/coding-agent/src/session/turns/retry.ts @@ -1,7 +1,7 @@ import type { AgentEvent } from "@earendil-works/pi-agent-core"; import { type AssistantMessage, isContextOverflow } from "@earendil-works/pi-ai"; -import { addLoginGuidanceToAuthError } from "../core/auth-guidance.js"; -import type { AuthSourceToken } from "../core/auth-storage.js"; +import { addLoginGuidanceToAuthError } from "../../core/auth-guidance.js"; +import type { AuthSourceToken } from "../../core/auth-storage.js"; import { isAgentLifecycleFailure, isFauxProviderQueueExhausted, @@ -9,9 +9,9 @@ import { providerRetryDelay, providerStreamFailureKind, providerStreamFailureRetryAfterMs, -} from "../core/provider-retry.js"; -import type { SettingsManager } from "../core/settings-manager.js"; -import { sleep } from "../utils/sleep.js"; +} from "../../core/provider-retry.js"; +import type { SettingsManager } from "../../core/settings-manager.js"; +import { sleep } from "../../utils/sleep.js"; export type SessionRetryEvent = | { diff --git a/packages/coding-agent/src/session/turns/turn-execution.ts b/packages/coding-agent/src/session/turns/turn-execution.ts new file mode 100644 index 0000000000..d5e0bcaf82 --- /dev/null +++ b/packages/coding-agent/src/session/turns/turn-execution.ts @@ -0,0 +1,224 @@ +import type { Agent, AgentMessage } from "@earendil-works/pi-agent-core"; +import type { ExtensionRunner } from "../../core/extensions/index.js"; +import { type CustomMessage, createHarnessDigestMessage, HARNESS_DIGEST_CUSTOM_TYPE } from "../../core/messages.js"; +import { type SessionAction, transitionSessionAction } from "../../core/session-action-store.js"; +import type { BuildSystemPromptOptions } from "../../core/system-prompt.js"; +import type { SessionCommitFence, SessionCommitLease } from "../input/commit-fence.js"; +import { + createDeliveryRecord, + DeferredSessionInputError, + type PreparedPromptPreparation, + type PreparedTurnPayload, + primaryDeliveryRecord, + type QueuedSessionAction, +} from "../prepared-actions.js"; +import type { TurnPreparer } from "./turn-preparation.js"; + +export interface SessionTurnExecutionHost { + getPreparer(): TurnPreparer; + getFence(): Pick; + acquireFence(): Promise; + isDeferred(epoch: number): boolean; + isStreaming(): boolean; + getBasePrompt(): string; + refreshExtensionSystemPrompt(extensionPrompt: string, baseSnapshot: string): string; + getBasePromptOptions(): BuildSystemPromptOptions; + getExtensions(): Pick; + getAgent(): Pick; + takeNextTurnMessages(): CustomMessage[]; + restoreNextTurnMessages(messages: CustomMessage[]): void; + consumePendingDigest(): boolean; + rearmDigest(): void; + getDigest(): string; + getLatestDigest(): string | undefined; + suppressForMessage(message: AgentMessage): void; + runSuppressed(run: () => Promise): Promise; + notifyCheckpoints(): void; + emitQueueUpdate(): void; + hasCancelledCapture(): boolean; + getEventQueue(): Promise; + waitForRetry(): Promise; + forgetContinuations(messages: AgentMessage[]): void; +} +export class SessionTurnExecution { + constructor(private readonly host: SessionTurnExecutionHost) {} + appendBeforeAgentStartMessages( + messages: AgentMessage[], + result: Awaited>, + ): void { + if (!result?.messages) return; + for (const message of result.messages) { + messages.push({ + role: "custom", + customType: message.customType, + content: message.content, + display: message.display, + details: message.details, + timestamp: Date.now(), + }); + } + } + + applyPreparedSystemPrompt( + preparation: PreparedPromptPreparation | undefined, + preserveEmptyExtensionPrompt: boolean, + ): void { + const extensionPrompt = preparation?.result?.systemPrompt; + const hasExtensionPrompt = preserveEmptyExtensionPrompt + ? extensionPrompt !== undefined + : Boolean(extensionPrompt); + this.host.getAgent().state.systemPrompt = + hasExtensionPrompt && extensionPrompt !== undefined && preparation !== undefined + ? this.host.refreshExtensionSystemPrompt(extensionPrompt, preparation.basePromptSnapshot) + : this.host.getBasePrompt(); + } + + async startPreparedTurnActions(actions: QueuedSessionAction[], epoch: number): Promise { + let nextTurnMessages: CustomMessage[] = []; + const activeTurns = () => + actions.filter( + (action): action is SessionAction => + action.payload.kind === "turn" && action.lifecycle.state === "preparing", + ); + const firstTurn = activeTurns()[0]; + if (!firstTurn) return; + const executionPolicy = firstTurn.payload.executionPolicy; + // The digest is never parked as pending context; lazy injection re-arms instead. + const parkNextTurnMessages = (messages: CustomMessage[]) => { + const parked = messages.filter((message) => message.customType !== HARNESS_DIGEST_CUSTOM_TYPE); + if (parked.length !== messages.length) this.host.rearmDigest(); + this.host.restoreNextTurnMessages(parked); + }; + const restoreNextTurnContext = () => { + parkNextTurnMessages(nextTurnMessages); + nextTurnMessages = []; + }; + try { + const preparedTurn = await this.host.getPreparer().prepare(executionPolicy.preparation, { + afterValidation: () => { + if (this.host.isDeferred(epoch)) { + throw new DeferredSessionInputError("Session input paused before preflight"); + } + }, + prepare: async () => { + if (executionPolicy.nextTurnContextTiming === "preparation") { + nextTurnMessages = this.host.takeNextTurnMessages(); + } + if (!executionPolicy.runBeforeAgentStart) return undefined; + while (activeTurns().some((action) => action.payload.prepared === undefined)) { + if (this.host.isDeferred(epoch)) { + throw new DeferredSessionInputError("Session input paused before preparation"); + } + const preparationAction = activeTurns().at(-1); + if (!preparationAction) return undefined; + const basePromptSnapshot = this.host.getBasePrompt(); + const result = await this.host + .getExtensions() + .emitBeforeAgentStart( + preparationAction.payload.text, + preparationAction.payload.images, + basePromptSnapshot, + this.host.getBasePromptOptions(), + ); + if (activeTurns().at(-1) !== preparationAction) continue; + const prepared = { result, basePromptSnapshot }; + for (const action of activeTurns()) action.payload.prepared = prepared; + } + if (this.host.isDeferred(epoch)) { + throw new DeferredSessionInputError("Session input paused before handoff"); + } + return activeTurns()[0]?.payload.prepared; + }, + shouldCommit: () => activeTurns().length > 0, + commit: (prepared) => { + if (this.host.isDeferred(epoch)) { + throw new DeferredSessionInputError("Session input paused before handoff"); + } + const turns = activeTurns(); + if (turns.length === 0) return undefined; + return { prepared, turns }; + }, + }); + if (!preparedTurn) { + restoreNextTurnContext(); + return; + } + const { prepared, turns } = preparedTurn; + const commitFence = await this.host.acquireFence(); + let promptPromise: Promise; + try { + promptPromise = this.host.getFence().run(commitFence, () => { + if ( + this.host.isDeferred(epoch) || + this.host.isStreaming() || + turns.some((action) => action.lifecycle.state !== "preparing") + ) { + throw new DeferredSessionInputError("Agent became active before session input handoff"); + } + if (executionPolicy.nextTurnContextTiming === "commit") { + nextTurnMessages = this.host.takeNextTurnMessages(); + } + if (this.host.consumePendingDigest()) { + // The first-turn digest rides the turn's delivery records so a + // cancelled first turn strips it with the rest of the turn. + + const digest = this.host.getDigest(); + if (this.host.getLatestDigest() !== digest) { + nextTurnMessages = [createHarnessDigestMessage(digest), ...nextTurnMessages]; + } + } + const contextRecords = nextTurnMessages.map((message) => + createDeliveryRecord(turns[0].id, "next_turn", message), + ); + const firstPrimaryIndex = turns[0].payload.records.indexOf(primaryDeliveryRecord(turns[0])); + turns[0].payload.records.splice(firstPrimaryIndex, 0, ...contextRecords); + const preparedMessages: AgentMessage[] = turns.flatMap((action) => + action.payload.records.map((record) => record.message), + ); + for (const action of turns) { + if (action.suppressAutonomousContinuation) { + this.host.suppressForMessage(primaryDeliveryRecord(action).message); + } + } + if (executionPolicy.runBeforeAgentStart) { + this.appendBeforeAgentStartMessages(preparedMessages, prepared?.result); + this.applyPreparedSystemPrompt(prepared, executionPolicy.preserveEmptyExtensionPrompt); + } else if (executionPolicy.nextTurnContextTiming !== "skip") { + this.host.getAgent().state.systemPrompt = this.host.getBasePrompt(); + } + for (const action of turns) transitionSessionAction(action, { state: "committing" }); + this.host.notifyCheckpoints(); + this.host.emitQueueUpdate(); + return turns.some((action) => action.suppressAutonomousContinuation) + ? this.host.runSuppressed(() => this.host.getAgent().prompt(preparedMessages)) + : this.host.getAgent().prompt(preparedMessages); + }); + } finally { + commitFence.release(); + } + await promptPromise; + if (executionPolicy.completionIncludesRetryChain) await this.host.waitForRetry(); + if (!this.host.hasCancelledCapture()) await this.host.getEventQueue(); + if ( + turns.some( + (action) => + action.lifecycle.state !== "cancelled" && + !primaryDeliveryRecord(action).durable && + !this.host.getAgent().state.messages.includes(primaryDeliveryRecord(action).message), + ) + ) { + throw new Error("Session input dispatch settled without durable delivery"); + } + this.host.forgetContinuations(turns.map((action) => primaryDeliveryRecord(action).message)); + } catch (error) { + const delivered = new Set(this.host.getAgent().state.messages); + parkNextTurnMessages(nextTurnMessages.filter((message) => !delivered.has(message))); + for (const action of actions) { + if (action.payload.kind === "turn") { + action.payload.records = action.payload.records.filter((record) => record.role !== "next_turn"); + } + } + throw error; + } + } +} diff --git a/packages/coding-agent/src/session/turns/turn-policy.ts b/packages/coding-agent/src/session/turns/turn-policy.ts new file mode 100644 index 0000000000..0f75c2a756 --- /dev/null +++ b/packages/coding-agent/src/session/turns/turn-policy.ts @@ -0,0 +1,141 @@ +import type { + AgentMessage, + GetContinuationMessagesContext, + ShouldStopAfterTurnContext, +} from "@earendil-works/pi-agent-core"; +import type { Api, AssistantMessage, Model } from "@earendil-works/pi-ai"; +import { shouldCompact } from "../../core/compaction/index.js"; +import { getLatestCompactionEntry, type SessionManager } from "../../core/session-manager.js"; +import type { SettingsManager } from "../../core/settings-manager.js"; +import type { SessionCompaction } from "../compaction/compaction.js"; +import type { GoalController } from "../goals/controller.js"; +import type { SessionRefinement } from "../refinement/refinement.js"; +import type { SessionAutonomousContinuation } from "./autonomous-continuation.js"; + +export interface SessionTurnPolicyHost { + steeringStopPending(): boolean; + stopGoalForTerminalMessage(message: AssistantMessage): boolean; + getGoals(): Pick; + accountAssistantBudget(message: AssistantMessage): Promise | undefined; + getRefinement(): Pick; + getEventQueue(): Promise; + getCompaction(): Pick< + SessionCompaction, + "resetContinuation" | "hasPendingRequest" | "requestContinuation" | "getThresholdContextTokens" + >; + getMessages(): AgentMessage[]; + getSettings(): Pick; + getModel(): Model | undefined; + getStore(): Pick; + queueThresholdGoal(message: AssistantMessage): boolean; + queueThresholdAutonomous(message: AssistantMessage): Promise; + getQueuedCount(): number; + getArrivalEpoch(): number; + getGoalMessages(context: GetContinuationMessagesContext, signal?: AbortSignal): Promise; + getAutonomous(): Pick; + snapshotAutonomous: SessionAutonomousContinuation["snapshotAutonomousRuntimeState"]; + restoreAutonomous: SessionAutonomousContinuation["restoreAutonomousRuntimeSnapshot"]; +} +export class SessionTurnPolicy { + constructor(private readonly host: SessionTurnPolicyHost) {} + shouldStopBeforeTurn(): boolean { + return this.host.steeringStopPending(); + } + + async shouldStopAfterTurn(context: ShouldStopAfterTurnContext): Promise { + if (this.host.stopGoalForTerminalMessage(context.message)) { + return true; + } + try { + const budgetNotice = this.host.accountAssistantBudget(context.message); + if (budgetNotice) await budgetNotice; + } catch { + // Goal accounting must not interrupt the core agent loop. + } + // Serialized refine checkpoint: in print/headless mode, run refinement + // planning+apply synchronously here — the quiescent boundary between + // turns — so it never overlaps the primary model request. + // This MUST run BEFORE threshold compaction to prevent the + // compaction model call from overlapping an in-flight refine + // plan/apply that was started at message_end. + if (this.host.getRefinement().serialized) { + // Ensure the preceding message_end processing (counter increment, + // background plan kickoff) has completed before the checkpoint. + await this.host.getEventQueue(); + await this.host.getRefinement()._runSerializedRefineCheckpoint(); + } + if (await this.shouldStopForThresholdCompaction(context)) { + return true; + } + // Steering stops continuation only after mandatory serialized checkpoints. + // Returning true here still prevents the agent loop from starting another turn. + return this.host.steeringStopPending(); + } + + async shouldStopForThresholdCompaction(context: ShouldStopAfterTurnContext): Promise { + this.host.getCompaction().resetContinuation(); + if (!this.host.getCompaction().hasPendingRequest && !(await this.thresholdCompactionNeeded(context))) { + return false; + } + + const lastMessage = this.host.getMessages()[this.host.getMessages().length - 1]; + // A queued continuation disproves the assistant-last "task finished" heuristic, so preserve a true set above. + if (lastMessage !== undefined && lastMessage.role !== "assistant") + this.host.getCompaction().requestContinuation(); + return true; + } + + async thresholdCompactionNeeded(context: ShouldStopAfterTurnContext): Promise { + const settings = this.host.getSettings().getCompactionSettings(); + if (!settings.enabled) return false; + + const contextWindow = this.host.getModel()?.contextWindow ?? 0; + const compactionEntry = getLatestCompactionEntry(this.host.getStore().getBranch()); + const compactionTimestamp = compactionEntry ? new Date(compactionEntry.timestamp).getTime() : undefined; + if (compactionTimestamp !== undefined && context.message.timestamp <= compactionTimestamp) { + return false; + } + + const contextTokens = this.host.getCompaction().getThresholdContextTokens(context.message, compactionTimestamp); + if (contextTokens === undefined || !shouldCompact(contextTokens, contextWindow, settings)) { + return false; + } + + // Goal continuation takes exclusive priority over autonomous continuation, matching _getContinuationMessages. + if (this.host.queueThresholdGoal(context.message)) { + this.host.getCompaction().requestContinuation(); + } else if (await this.host.queueThresholdAutonomous(context.message)) { + this.host.getCompaction().requestContinuation(); + } + return true; + } + + async getContinuationMessages( + context: GetContinuationMessagesContext, + signal?: AbortSignal, + ): Promise { + if (this.host.getQueuedCount() > 0) { + return []; + } + const arrivalEpoch = this.host.getArrivalEpoch(); + const goalSnapshot = this.host.getGoals().checkpoint(); + const goalMessages = await this.host.getGoalMessages(context, signal); + if (goalMessages.length > 0 || signal?.aborted) { + if (goalMessages.length > 0 && this.host.getArrivalEpoch() !== arrivalEpoch) { + this.host.getGoals().restore(goalSnapshot); + return []; + } + return goalMessages; + } + if (this.host.getAutonomous().isSuppressed(context.newMessages)) { + return []; + } + const autonomousSnapshot = this.host.snapshotAutonomous(); + const autonomousMessage = await this.host.getAutonomous().next(context.message, signal); + if (autonomousMessage && this.host.getArrivalEpoch() !== arrivalEpoch) { + this.host.restoreAutonomous(autonomousSnapshot); + return []; + } + return autonomousMessage ? [autonomousMessage] : []; + } +} diff --git a/packages/coding-agent/src/session/turn-preparation.ts b/packages/coding-agent/src/session/turns/turn-preparation.ts similarity index 100% rename from packages/coding-agent/src/session/turn-preparation.ts rename to packages/coding-agent/src/session/turns/turn-preparation.ts diff --git a/packages/coding-agent/test/agent-connection-in-process.test.ts b/packages/coding-agent/test/agent-connection-in-process.test.ts index 0c027777ec..b59bc3c6f0 100644 --- a/packages/coding-agent/test/agent-connection-in-process.test.ts +++ b/packages/coding-agent/test/agent-connection-in-process.test.ts @@ -3,9 +3,9 @@ import { getModel } from "@earendil-works/pi-ai"; import { describe, expect, it, vi } from "vitest"; import type { AgentSessionEvent, AgentSessionEventListener, PromptOptions } from "../src/core/agent-session.js"; import type { AgentSessionRuntime } from "../src/core/agent-session-runtime.js"; -import { emptyGoalState } from "../src/core/goals.js"; import { InProcessAgentConnection } from "../src/modes/agent-connection/in-process-agent-connection.js"; import type { AgentConnectionEvent, AgentConnectionState } from "../src/modes/agent-connection/types.js"; +import { emptyGoalState } from "../src/session/goals/contracts.js"; type RuntimeSession = AgentSessionRuntime["session"]; type RuntimeRebindCallback = Parameters[0]; diff --git a/packages/coding-agent/test/agent-session-concurrent.test.ts b/packages/coding-agent/test/agent-session-concurrent.test.ts index a9791ee829..fe8992936b 100644 --- a/packages/coding-agent/test/agent-session-concurrent.test.ts +++ b/packages/coding-agent/test/agent-session-concurrent.test.ts @@ -21,7 +21,7 @@ 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 { IpythonKernelProvisioner } from "../src/core/tools/ipython.js"; -import type { SessionRefinement } from "../src/session/refinement.js"; +import type { SessionRefinement } from "../src/session/refinement/refinement.js"; import { createTestExtensionsResult, createTestResourceLoader } from "./utilities.js"; class MockAssistantStream extends EventStream { diff --git a/packages/coding-agent/test/goal-continuation-quiescence.test.ts b/packages/coding-agent/test/goal-continuation-quiescence.test.ts index 091f5684a3..8f05db2f39 100644 --- a/packages/coding-agent/test/goal-continuation-quiescence.test.ts +++ b/packages/coding-agent/test/goal-continuation-quiescence.test.ts @@ -1,125 +1,127 @@ +import type { GetContinuationMessagesContext } from "@earendil-works/pi-agent-core"; +import { fauxAssistantMessage } from "@earendil-works/pi-ai"; import { describe, expect, it, vi } from "vitest"; -import { AgentSession } from "../src/core/agent-session.js"; -import { emptyGoalState } from "../src/core/goals.js"; -import { GoalController } from "../src/goals/controller.js"; -import { SessionInputScheduler } from "../src/session/input-scheduler.js"; +import { ActionStore } from "../src/core/session-action-store.js"; +import { SessionGoalContinuation, type SessionGoalContinuationHost } from "../src/session/goals/continuation.js"; +import { emptyGoalState } from "../src/session/goals/contracts.js"; +import { GoalController } from "../src/session/goals/controller.js"; +import { SessionInputScheduler } from "../src/session/input/input-scheduler.js"; +import type { QueuedSessionAction } from "../src/session/prepared-actions.js"; -type Harness = { - _goals: GoalController; - _goalContinuationAwaitsRlmWork: boolean; - _disposed: boolean; - _disposing: boolean; - _inputScheduler: SessionInputScheduler; - _hasUnsettledRlmQuiescenceWork: () => boolean; - _stopGoalContinuationForTerminalMessage: () => boolean; - _ensureGoalRuntimeActive: () => void; - _createPreparedTurnAction: ReturnType; - _admitSessionInput: ReturnType; -}; - -const getGoalContinuation = Reflect.get(AgentSession.prototype, "_getGoalContinuationMessages") as ( - this: Harness, - context: { message: unknown; context: unknown }, -) => Promise; -const maybeResume = Reflect.get(AgentSession.prototype, "_maybeResumeGoalContinuationAfterRlmWork") as ( - this: Harness, -) => void; - -function harness(overrides: Partial = {}): Harness { +function harness(overrides: { awaitsChildWork?: boolean } & Partial = {}) { const goals = new GoalController({ load: emptyGoalState, save: () => {} }, () => {}); goals.start("ship it", undefined); - return { - _goals: goals, - _goalContinuationAwaitsRlmWork: false, - _disposed: false, - _disposing: false, - _inputScheduler: new SessionInputScheduler({ canSchedule: () => false, run: async () => {} }), - _hasUnsettledRlmQuiescenceWork: () => false, - _stopGoalContinuationForTerminalMessage: () => false, - _ensureGoalRuntimeActive: () => {}, - _createPreparedTurnAction: vi.fn((schedule: string, _text: string, _images: unknown, options: unknown) => ({ - schedule, - options, - })), - _admitSessionInput: vi.fn(), + const actions = new ActionStore(); + const scheduler = new SessionInputScheduler({ canSchedule: () => false, run: async () => {} }); + const admit = vi.fn( + overrides.admit ?? + ((action) => { + actions.enqueue(action); + return { accepted: true, disposition: "queued" }; + }), + ); + const owner = new SessionGoalContinuation(goals, actions, { + getGoalState: () => goals.current, + queuePrompt: async () => false, + getScheduler: () => scheduler, + isDisposed: () => false, + isDisposing: () => false, + hasUnsettledChildWork: () => false, + ensureRuntimeActive: () => {}, + cancelActions: (predicate) => actions.remove(predicate), + clearPendingGoalContexts: () => {}, + emitQueueUpdate: () => {}, + emitGoalUpdate: () => {}, + validate: async () => {}, + isStreaming: () => false, + includesGoals: () => true, + getAgent: () => ({ removeQueuedMessages: () => [] }), ...overrides, - }; + admit, + }); + if (overrides.awaitsChildWork) owner.deferUntilChildSettlement(); + return { owner, goals, scheduler, actions, admit }; } -const context = { message: { role: "assistant", stopReason: "stop" }, context: {} }; +const context: GetContinuationMessagesContext = { + message: fauxAssistantMessage("done"), + toolResults: [], + context: { systemPrompt: "", messages: [], tools: [] }, + newMessages: [], +}; describe("goal continuation vs unsettled subagent work", () => { it("defers the continuation while descendant work is unsettled", async () => { - const mode = harness({ _hasUnsettledRlmQuiescenceWork: () => true }); - await expect(getGoalContinuation.call(mode, context)).resolves.toEqual([]); - expect(mode._goalContinuationAwaitsRlmWork).toBe(true); - expect(mode._goals.state.continuationsUsed).toBe(0); + const mode = harness({ hasUnsettledChildWork: () => true }); + await expect(mode.owner.getGoalContinuationMessages(context)).resolves.toEqual([]); + expect(mode.owner.awaitsChildWork).toBe(true); + expect(mode.goals.state.continuationsUsed).toBe(0); }); it("continues normally when no descendant work is pending", async () => { const mode = harness(); - const messages = await getGoalContinuation.call(mode, context); + const messages = await mode.owner.getGoalContinuationMessages(context); expect(messages).toHaveLength(1); - expect(mode._goalContinuationAwaitsRlmWork).toBe(false); - expect(mode._goals.state.continuationsUsed).toBe(1); + expect(mode.owner.awaitsChildWork).toBe(false); + expect(mode.goals.state.continuationsUsed).toBe(1); }); it("resumes a deferred continuation exactly once, unqueued, idle-waking, and counted", () => { - const mode = harness({ _goalContinuationAwaitsRlmWork: true }); - maybeResume.call(mode); - maybeResume.call(mode); - expect(mode._admitSessionInput).toHaveBeenCalledTimes(1); - const [action, options] = mode._admitSessionInput.mock.calls[0]!; - expect((action as { options: { resumeIfIdle: boolean } }).options.resumeIfIdle).toBe(true); + const mode = harness({ awaitsChildWork: true }); + mode.owner.maybeResumeGoalContinuationAfterRlmWork(); + mode.owner.maybeResumeGoalContinuationAfterRlmWork(); + expect(mode.admit).toHaveBeenCalledTimes(1); + const [action, options] = mode.admit.mock.calls[0]!; + expect(action?.wake).toBe("immediate"); expect(options).toBeUndefined(); - expect(mode._goals.state.continuationsUsed).toBe(1); + expect(mode.goals.state.continuationsUsed).toBe(1); }); it("keeps the deferral while admission is paused and retries after release", () => { const paused = harness({ - _goalContinuationAwaitsRlmWork: true, + awaitsChildWork: true, }); - const pause = paused._inputScheduler.acquireAdmissionPause(() => {}); - maybeResume.call(paused); - expect(paused._admitSessionInput).not.toHaveBeenCalled(); - expect(paused._goalContinuationAwaitsRlmWork).toBe(true); + const pause = paused.scheduler.acquireAdmissionPause(() => {}); + paused.owner.maybeResumeGoalContinuationAfterRlmWork(); + expect(paused.admit).not.toHaveBeenCalled(); + expect(paused.owner.awaitsChildWork).toBe(true); pause.release(); - maybeResume.call(paused); - expect(paused._admitSessionInput).toHaveBeenCalledTimes(1); - expect(paused._goalContinuationAwaitsRlmWork).toBe(false); + paused.owner.maybeResumeGoalContinuationAfterRlmWork(); + expect(paused.admit).toHaveBeenCalledTimes(1); + expect(paused.owner.awaitsChildWork).toBe(false); }); it("keeps the deferral while the pump is suspended after an abort", () => { - const mode = harness({ _goalContinuationAwaitsRlmWork: true }); - mode._inputScheduler.suspend("abort"); - maybeResume.call(mode); - expect(mode._admitSessionInput).not.toHaveBeenCalled(); - expect(mode._goalContinuationAwaitsRlmWork).toBe(true); + const mode = harness({ awaitsChildWork: true }); + mode.scheduler.suspend("abort"); + mode.owner.maybeResumeGoalContinuationAfterRlmWork(); + expect(mode.admit).not.toHaveBeenCalled(); + expect(mode.owner.awaitsChildWork).toBe(true); }); it("keeps the deferral and rolls back the count when admission throws", () => { const mode = harness({ - _goalContinuationAwaitsRlmWork: true, - _admitSessionInput: vi.fn(() => { + awaitsChildWork: true, + admit: vi.fn(() => { throw new Error("admission race"); }), }); - maybeResume.call(mode); - expect(mode._goalContinuationAwaitsRlmWork).toBe(true); - expect(mode._goals.state.continuationsUsed).toBe(0); + mode.owner.maybeResumeGoalContinuationAfterRlmWork(); + expect(mode.owner.awaitsChildWork).toBe(true); + expect(mode.goals.state.continuationsUsed).toBe(0); }); it("stays deferred while work remains and drops the deferral for inactive goals", () => { - const busy = harness({ _goalContinuationAwaitsRlmWork: true, _hasUnsettledRlmQuiescenceWork: () => true }); - maybeResume.call(busy); - expect(busy._admitSessionInput).not.toHaveBeenCalled(); - expect(busy._goalContinuationAwaitsRlmWork).toBe(true); + const busy = harness({ awaitsChildWork: true, hasUnsettledChildWork: () => true }); + busy.owner.maybeResumeGoalContinuationAfterRlmWork(); + expect(busy.admit).not.toHaveBeenCalled(); + expect(busy.owner.awaitsChildWork).toBe(true); - const inactive = harness({ _goalContinuationAwaitsRlmWork: true }); - inactive._goals.pause(); - maybeResume.call(inactive); - expect(inactive._admitSessionInput).not.toHaveBeenCalled(); - expect(inactive._goalContinuationAwaitsRlmWork).toBe(false); + const inactive = harness({ awaitsChildWork: true }); + inactive.goals.pause(); + inactive.owner.maybeResumeGoalContinuationAfterRlmWork(); + expect(inactive.admit).not.toHaveBeenCalled(); + expect(inactive.owner.awaitsChildWork).toBe(false); }); }); diff --git a/packages/coding-agent/test/interactive-mode-status.test.ts b/packages/coding-agent/test/interactive-mode-status.test.ts index 34756d4e93..dc1adef3ad 100644 --- a/packages/coding-agent/test/interactive-mode-status.test.ts +++ b/packages/coding-agent/test/interactive-mode-status.test.ts @@ -21,7 +21,6 @@ import { formatNoModelsAvailableMessage } from "../src/core/auth-guidance.js"; import { type AuthStatus, AuthStorage } from "../src/core/auth-storage.js"; import type { AgentCronJob } from "../src/core/cron-jobs.js"; import type { AutocompleteProviderFactory } from "../src/core/extensions/types.js"; -import { emptyGoalState, type GoalState } from "../src/core/goals.js"; import { KeybindingsManager } from "../src/core/keybindings.js"; import { createSessionSlashCommandMessage, createSessionSlashCommandResultMessage } from "../src/core/messages.js"; import type { ModelRegistry } from "../src/core/model-registry.js"; @@ -55,6 +54,7 @@ import { formatSplashCwd, InteractiveMode, truncatePathMiddle } from "../src/mod import { ClientPromptStashStore, type PromptStashState } from "../src/modes/interactive/prompt-stash-state.js"; import { QueueSelection } from "../src/modes/interactive/queue-selection.js"; import { initTheme, theme } from "../src/modes/interactive/theme/theme.js"; +import { emptyGoalState, type GoalState } from "../src/session/goals/contracts.js"; function renderLastLine(container: Container, width = 120): string { const last = container.children[container.children.length - 1]; diff --git a/packages/coding-agent/test/session-child-usage.test.ts b/packages/coding-agent/test/session-child-usage.test.ts index 83c7b317fd..442f53882e 100644 --- a/packages/coding-agent/test/session-child-usage.test.ts +++ b/packages/coding-agent/test/session-child-usage.test.ts @@ -4,7 +4,7 @@ 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"; +import { type ChildUsageHost, type ChildUsageTracker, SessionChildUsage } from "../src/session/children/child-usage.js"; function usage(input: number, output: number): Usage { return { diff --git a/packages/coding-agent/test/session-command-messages.test.ts b/packages/coding-agent/test/session-command-messages.test.ts index 5db7082fd5..30182301fb 100644 --- a/packages/coding-agent/test/session-command-messages.test.ts +++ b/packages/coding-agent/test/session-command-messages.test.ts @@ -3,7 +3,6 @@ import type { TUI } from "@earendil-works/pi-tui"; import stripAnsi from "strip-ansi"; import { beforeAll, describe, expect, test, vi } from "vitest"; import { AGENT_MESSAGE_SOURCE, createAgentSessionMessage } from "../src/core/agent-messages.js"; -import { createGoalContextMessage, type GoalState } from "../src/core/goals.js"; import { COMPACTION_OUTCOME_CUSTOM_TYPE, type CustomMessage, @@ -30,6 +29,7 @@ import { buildConversationComponents } from "../src/modes/interactive/components import { SlashCommandMessageComponent } from "../src/modes/interactive/components/slash-command-message.js"; import { SlashCommandResultMessageComponent } from "../src/modes/interactive/components/slash-command-result-message.js"; import { initTheme } from "../src/modes/interactive/theme/theme.js"; +import { createGoalContextMessage, type GoalState } from "../src/session/goals/contracts.js"; const componentOptions = { ui: { requestRender: vi.fn() } as unknown as TUI, diff --git a/packages/coding-agent/test/session/bash.test.ts b/packages/coding-agent/test/session/bash.test.ts index 236d921c29..23414e77d4 100644 --- a/packages/coding-agent/test/session/bash.test.ts +++ b/packages/coding-agent/test/session/bash.test.ts @@ -2,7 +2,7 @@ import { Buffer } from "node:buffer"; import { describe, expect, it, vi } from "vitest"; import type { BashResult } from "../../src/core/bash-executor.js"; import type { BashExecutionMessage } from "../../src/core/messages.js"; -import { SessionBash, type SessionBashEvent, type SessionBashHost } from "../../src/session/bash.js"; +import { SessionBash, type SessionBashEvent, type SessionBashHost } from "../../src/session/tools/bash.js"; function deferred() { let resolve = () => {}; diff --git a/packages/coding-agent/test/session/commit-fence.test.ts b/packages/coding-agent/test/session/commit-fence.test.ts index 731c29086a..94a3a97a03 100644 --- a/packages/coding-agent/test/session/commit-fence.test.ts +++ b/packages/coding-agent/test/session/commit-fence.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { SessionCommitFence } from "../../src/session/commit-fence.js"; +import { SessionCommitFence } from "../../src/session/input/commit-fence.js"; function deferred() { let resolve = () => {}; diff --git a/packages/coding-agent/test/session/compaction.test.ts b/packages/coding-agent/test/session/compaction.test.ts index 5d58a5404e..e9c4ea5a15 100644 --- a/packages/coding-agent/test/session/compaction.test.ts +++ b/packages/coding-agent/test/session/compaction.test.ts @@ -2,8 +2,8 @@ import type { AgentMessage } from "@earendil-works/pi-agent-core"; import { fauxAssistantMessage, registerFauxProvider } from "@earendil-works/pi-ai"; import { afterEach, describe, expect, it, vi } from "vitest"; import { SessionManager } from "../../src/core/session-manager.js"; -import { SessionCompaction, type SessionCompactionHost } from "../../src/session/compaction.js"; -import { CompactionSkippedError } from "../../src/session/compaction-execution.js"; +import { SessionCompaction, type SessionCompactionHost } from "../../src/session/compaction/compaction.js"; +import { CompactionSkippedError } from "../../src/session/compaction/compaction-execution.js"; function deferred() { let resolve!: (value: T) => void; @@ -25,6 +25,8 @@ function setup() { const order: string[] = []; const result = { summary: "summary", firstKeptEntryId: "kept", tokensBefore: 1000 }; const host = { + includesCompactSkill: () => true, + getContextUsage: () => undefined, getSettings: vi.fn(() => ({ enabled: false, reserveTokens: 100, keepRecentTokens: 10 })), runAutomatic: vi.fn(async () => false), queueGoalContinuation: vi.fn(() => false), diff --git a/packages/coding-agent/test/session/continuation.test.ts b/packages/coding-agent/test/session/continuation.test.ts index f698bd90e6..196e440942 100644 --- a/packages/coding-agent/test/session/continuation.test.ts +++ b/packages/coding-agent/test/session/continuation.test.ts @@ -1,8 +1,8 @@ import { AgentContinueError, type AgentMessage } from "@earendil-works/pi-agent-core"; import { fauxAssistantMessage } from "@earendil-works/pi-ai"; import { describe, expect, it, vi } from "vitest"; -import { SessionCommitFence } from "../../src/session/commit-fence.js"; -import { SessionContinuation, type SessionContinuationHost } from "../../src/session/continuation.js"; +import { SessionCommitFence } from "../../src/session/input/commit-fence.js"; +import { SessionContinuation, type SessionContinuationHost } from "../../src/session/turns/continuation.js"; function deferred() { let resolve!: () => void; diff --git a/packages/coding-agent/test/goals/controller.test.ts b/packages/coding-agent/test/session/goals/controller.test.ts similarity index 96% rename from packages/coding-agent/test/goals/controller.test.ts rename to packages/coding-agent/test/session/goals/controller.test.ts index 99c0423617..0ba7cc13b0 100644 --- a/packages/coding-agent/test/goals/controller.test.ts +++ b/packages/coding-agent/test/session/goals/controller.test.ts @@ -1,7 +1,7 @@ import { type AssistantMessage, fauxAssistantMessage } from "@earendil-works/pi-ai"; import { describe, expect, it, vi } from "vitest"; -import { emptyGoalState, type GoalState } from "../../src/core/goals.js"; -import { GoalController } from "../../src/goals/controller.js"; +import { emptyGoalState, type GoalState } from "../../../src/session/goals/contracts.js"; +import { GoalController } from "../../../src/session/goals/controller.js"; function createController(initial = emptyGoalState()) { let persisted = initial; diff --git a/packages/coding-agent/test/session/input-dispatcher.test.ts b/packages/coding-agent/test/session/input-dispatcher.test.ts index 19e1679477..1bc3bd2667 100644 --- a/packages/coding-agent/test/session/input-dispatcher.test.ts +++ b/packages/coding-agent/test/session/input-dispatcher.test.ts @@ -1,7 +1,7 @@ import type { AgentMessage } from "@earendil-works/pi-agent-core"; import { describe, expect, it, vi } from "vitest"; import { ActionStore, transitionSessionAction } from "../../src/core/session-action-store.js"; -import { SessionInputDispatcher, type SessionInputDispatcherHost } from "../../src/session/input-dispatcher.js"; +import { SessionInputDispatcher, type SessionInputDispatcherHost } from "../../src/session/input/input-dispatcher.js"; import { createDeliveryRecord, createPreparedTurnAction, @@ -9,7 +9,7 @@ import { primaryDeliveryRecord, type QueuedSessionAction, } from "../../src/session/prepared-actions.js"; -import { createTurnExecutionPolicy } from "../../src/session/turn-preparation.js"; +import { createTurnExecutionPolicy } from "../../src/session/turns/turn-preparation.js"; function createFixture() { const actions = new ActionStore(); diff --git a/packages/coding-agent/test/session/input-scheduler.test.ts b/packages/coding-agent/test/session/input-scheduler.test.ts index 05829c3f6a..7047beb79a 100644 --- a/packages/coding-agent/test/session/input-scheduler.test.ts +++ b/packages/coding-agent/test/session/input-scheduler.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it, vi } from "vitest"; -import { SessionInputScheduler } from "../../src/session/input-scheduler.js"; +import { SessionInputScheduler } from "../../src/session/input/input-scheduler.js"; function deferred() { let resolve = () => {}; diff --git a/packages/coding-agent/test/session/kernel-environment.test.ts b/packages/coding-agent/test/session/kernel-environment.test.ts index 23ec72d7f8..a0e4f89604 100644 --- a/packages/coding-agent/test/session/kernel-environment.test.ts +++ b/packages/coding-agent/test/session/kernel-environment.test.ts @@ -6,7 +6,7 @@ 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"; +import { KernelEnvironment, type KernelEnvironmentHost } from "../../src/session/kernel/kernel-environment.js"; function createEnvironment(overrides: Partial = {}, sessionDir?: string) { return new KernelEnvironment( diff --git a/packages/coding-agent/test/session/kernel.test.ts b/packages/coding-agent/test/session/kernel.test.ts index 0cd4f26c89..65489d33af 100644 --- a/packages/coding-agent/test/session/kernel.test.ts +++ b/packages/coding-agent/test/session/kernel.test.ts @@ -4,7 +4,7 @@ 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"; +import { SessionKernel, type SessionKernelHost } from "../../src/session/kernel/kernel.js"; const mocks = vi.hoisted(() => ({ instances: [] as Array<{ diff --git a/packages/coding-agent/test/session/pending-context.test.ts b/packages/coding-agent/test/session/pending-context.test.ts new file mode 100644 index 0000000000..537c74a090 --- /dev/null +++ b/packages/coding-agent/test/session/pending-context.test.ts @@ -0,0 +1,78 @@ +import { describe, expect, it, vi } from "vitest"; +import type { CustomMessage } from "../../src/core/messages.js"; +import { ActionStore } from "../../src/core/session-action-store.js"; +import { SessionPendingContext } from "../../src/session/context/pending-context.js"; +import { SessionCommitFence } from "../../src/session/input/commit-fence.js"; +import { SessionInputScheduler } from "../../src/session/input/input-scheduler.js"; +import type { QueuedSessionAction } from "../../src/session/prepared-actions.js"; + +function createOwner() { + const actions = new ActionStore(); + const scheduler = new SessionInputScheduler({ canSchedule: () => false, run: async () => {} }); + const fence = new SessionCommitFence(); + const schedule = vi.fn(); + const owner = new SessionPendingContext(actions, { + getScheduler: () => scheduler, + getFence: () => fence, + isDisposed: () => false, + isDisposing: () => false, + admit: (action) => { + actions.enqueue(action); + return { accepted: true }; + }, + scheduleInput: schedule, + addCheckpointWaiter: () => {}, + removeCheckpointWaiter: () => {}, + acquireFence: (signal) => fence.acquire(signal), + cancelActions: (predicate) => actions.remove(predicate), + }); + return { owner, schedule }; +} + +function message(content: string): CustomMessage { + return { role: "custom", customType: "context", content, display: false, timestamp: 1, details: { content } }; +} + +describe("pending context ownership", () => { + it("keeps rollback message identities and order without cloning or waking input", () => { + const { owner, schedule } = createOwner(); + const first = message("first"); + const last = message("last"); + owner.appendMessages(last); + owner.prependMessages([first]); + const taken = owner.takePendingNextTurnMessages(); + expect(taken).toEqual([first, last]); + expect(taken[0]).toBe(first); + expect(taken[1]).toBe(last); + owner.appendMessages(message("next")); + expect(taken).toEqual([first, last]); + expect(owner.takePendingNextTurnMessages()).toHaveLength(1); + expect(schedule).not.toHaveBeenCalled(); + }); + + it("removes matching context without cloning survivors, while recovery clones and wakes", () => { + const { owner, schedule } = createOwner(); + const keep = message("keep"); + owner.appendMessages(message("drop"), keep); + owner.removeMessagesMatching((item) => item.content === "drop"); + expect(owner.takePendingNextTurnMessages()[0]).toBe(keep); + expect(schedule).not.toHaveBeenCalled(); + owner.restorePendingNextTurnMessages([keep]); + const restored = owner.takePendingNextTurnMessages()[0]; + expect(restored).toEqual(keep); + expect(restored).not.toBe(keep); + expect(restored.details).toBe(keep.details); + expect(schedule).toHaveBeenCalledOnce(); + }); + + it("tracks terminal-notice retention independently from pending-message disposal", () => { + const { owner } = createOwner(); + owner.retainTerminalNotice("notice"); + owner.appendMessages(message("pending")); + owner.dispose(); + expect(owner.takePendingNextTurnMessages()).toEqual([]); + expect(owner.isRetainedTerminalNotice("notice")).toBe(true); + owner.releaseTerminalNotice("notice"); + expect(owner.isRetainedTerminalNotice("notice")).toBe(false); + }); +}); diff --git a/packages/coding-agent/test/session/retry.test.ts b/packages/coding-agent/test/session/retry.test.ts index 9d958c52a6..58d414d2e4 100644 --- a/packages/coding-agent/test/session/retry.test.ts +++ b/packages/coding-agent/test/session/retry.test.ts @@ -1,7 +1,7 @@ import { type AssistantMessage, fauxAssistantMessage } from "@earendil-works/pi-ai"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { AuthSourceToken } from "../../src/core/auth-storage.js"; -import { SessionRetry, type SessionRetryEvent, type SessionRetryHost } from "../../src/session/retry.js"; +import { SessionRetry, type SessionRetryEvent, type SessionRetryHost } from "../../src/session/turns/retry.js"; function failure(kind?: string): AssistantMessage { return { diff --git a/packages/coding-agent/test/session/submission-normalization.test.ts b/packages/coding-agent/test/session/submission-normalization.test.ts new file mode 100644 index 0000000000..0c59afd5bf --- /dev/null +++ b/packages/coding-agent/test/session/submission-normalization.test.ts @@ -0,0 +1,96 @@ +import type { ImageContent } from "@earendil-works/pi-ai"; +import { describe, expect, it, vi } from "vitest"; +import type { PromptTemplate } from "../../src/core/prompt-templates.js"; +import { createSyntheticSourceInfo } from "../../src/core/source-info.js"; +import { + type SubmissionNormalizationHost, + type SubmissionNormalizationPolicy, + SubmissionNormalizer, +} from "../../src/session/input/submission-normalization.js"; +import { createDeferred } from "../suite/scheduling.js"; + +const policy: SubmissionNormalizationPolicy = { + parseSessionCommands: true, + extensionCommands: "ignore", + inputSource: "interactive", + expandSkills: true, + expandPromptTemplates: true, +}; + +function createNormalizer() { + const extensions: ReturnType = { + hasHandlers: vi.fn(() => false), + emitInput: vi.fn(async () => ({ action: "continue" as const })), + getCommand: vi.fn(() => undefined), + createCommandContext: () => { + throw new Error("Unexpected extension command"); + }, + emitError: vi.fn(), + }; + const resources: { prompts: PromptTemplate[] } = { prompts: [] }; + const normalizer = new SubmissionNormalizer({ + getExtensions: () => extensions, + getPrompts: () => resources.prompts, + getSkills: () => [], + }); + return { normalizer, extensions, resources }; +} + +describe("submission normalization boundaries", () => { + it("keeps plain submissions synchronous when there are no input hooks", () => { + const { normalizer, extensions } = createNormalizer(); + const result = normalizer.normalizeSubmission("hello", undefined, policy); + expect(result).not.toBeInstanceOf(Promise); + expect(result).toEqual({ kind: "prompt", text: "hello", images: undefined }); + expect(extensions.emitInput).not.toHaveBeenCalled(); + }); + + it("recognizes session commands before extension or input processing", () => { + const { normalizer, extensions } = createNormalizer(); + vi.mocked(extensions.hasHandlers).mockReturnValue(true); + const result = normalizer.normalizeSubmission("/compact focus", undefined, { + ...policy, + extensionCommands: "execute", + }); + expect(result).toMatchObject({ kind: "sessionCommand", text: "/compact focus" }); + expect(result).not.toBeInstanceOf(Promise); + expect(extensions.getCommand).not.toHaveBeenCalled(); + expect(extensions.emitInput).not.toHaveBeenCalled(); + }); + + it("expands the current templates after an asynchronous input transform", async () => { + const { normalizer, extensions, resources } = createNormalizer(); + const transformed = createDeferred>>(); + vi.mocked(extensions.hasHandlers).mockReturnValue(true); + vi.mocked(extensions.emitInput).mockReturnValue(transformed.promise); + const result = normalizer.normalizeSubmission("original", undefined, policy); + resources.prompts = [ + { + name: "latest", + description: "", + content: "current resource $1", + filePath: "/tmp/latest.md", + sourceInfo: createSyntheticSourceInfo("/tmp/latest.md", { source: "test" }), + }, + ]; + transformed.resolve({ action: "transform", text: "/latest argument" }); + await expect(result).resolves.toEqual({ kind: "prompt", text: "current resource argument", images: undefined }); + }); + + it.each([undefined, []] satisfies Array)( + "preserves transform image semantics for %s", + async (replacement) => { + const { normalizer, extensions } = createNormalizer(); + const images: ImageContent[] = [{ type: "image", data: "image", mimeType: "image/png" }]; + vi.mocked(extensions.hasHandlers).mockReturnValue(true); + vi.mocked(extensions.emitInput).mockResolvedValue({ + action: "transform", + text: "changed", + images: replacement, + }); + const result = await normalizer.normalizeSubmission("original", images, policy); + expect(result).toEqual({ kind: "prompt", text: "changed", images: replacement ?? images }); + if (result.kind === "prompt") expect(result.images).toBe(replacement ?? images); + }, + ); +}); diff --git a/packages/coding-agent/test/session/tools.test.ts b/packages/coding-agent/test/session/tools.test.ts index 56e5a1b950..5e7af855df 100644 --- a/packages/coding-agent/test/session/tools.test.ts +++ b/packages/coding-agent/test/session/tools.test.ts @@ -1,7 +1,7 @@ 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"; +import { SessionTools, type SessionToolsHost } from "../../src/session/tools/tools.js"; describe("SessionTools ACP release", () => { afterEach(() => vi.restoreAllMocks()); diff --git a/packages/coding-agent/test/session/turn-preparation.test.ts b/packages/coding-agent/test/session/turn-preparation.test.ts index 011b5f08ac..625f41a711 100644 --- a/packages/coding-agent/test/session/turn-preparation.test.ts +++ b/packages/coding-agent/test/session/turn-preparation.test.ts @@ -3,7 +3,7 @@ import { createTurnExecutionPolicy, type TurnPreparationHost, TurnPreparer, -} from "../../src/session/turn-preparation.js"; +} from "../../src/session/turns/turn-preparation.js"; function createPreparation(overrides: Partial = {}) { const order: string[] = []; diff --git a/packages/coding-agent/test/suite/agent-session-compact-skill.test.ts b/packages/coding-agent/test/suite/agent-session-compact-skill.test.ts index 92af80fc73..d56c84462d 100644 --- a/packages/coding-agent/test/suite/agent-session-compact-skill.test.ts +++ b/packages/coding-agent/test/suite/agent-session-compact-skill.test.ts @@ -1,7 +1,7 @@ import type { ShouldStopAfterTurnContext } from "@earendil-works/pi-agent-core"; import { type AssistantMessage, fauxAssistantMessage } from "@earendil-works/pi-ai"; import { afterEach, describe, expect, it, vi } from "vitest"; -import type { SessionCompaction } from "../../src/session/compaction.js"; +import type { SessionCompaction } from "../../src/session/compaction/compaction.js"; import { createHarness, type Harness } from "./harness.js"; type SessionInternals = { diff --git a/packages/coding-agent/test/suite/agent-session-compaction-continuation.test.ts b/packages/coding-agent/test/suite/agent-session-compaction-continuation.test.ts index a65438017c..440f03a954 100644 --- a/packages/coding-agent/test/suite/agent-session-compaction-continuation.test.ts +++ b/packages/coding-agent/test/suite/agent-session-compaction-continuation.test.ts @@ -15,7 +15,7 @@ import { import { Type } from "typebox"; import { afterEach, describe, expect, it, vi } from "vitest"; import type { AgentSession } from "../../src/core/agent-session.js"; -import type { SessionCompaction } from "../../src/session/compaction.js"; +import type { SessionCompaction } from "../../src/session/compaction/compaction.js"; import { createHarness, type Harness } from "./harness.js"; type SessionInternals = { 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 aa6bf9dd8b..9b8f81059f 100644 --- a/packages/coding-agent/test/suite/agent-session-compaction.test.ts +++ b/packages/coding-agent/test/suite/agent-session-compaction.test.ts @@ -12,8 +12,9 @@ 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 type { SessionCompaction } from "../../src/session/compaction/compaction.js"; +import { createPreparedTurnAction } from "../../src/session/prepared-actions.js"; +import type { SessionContinuation } from "../../src/session/turns/continuation.js"; import { createHarness, getMessageText, type Harness } from "./harness.js"; import { createDeferred } from "./scheduling.js"; @@ -506,7 +507,7 @@ describe("AgentSession compaction characterization", () => { const internals = session as unknown as { _schedulePostCompactionContinue(): void; _cancelPostCompactionContinue(): void; - _sessionInputCheckpointWaiters: Set<() => void>; + _inputCheckpoints: { hasWaiters: boolean }; }; // A queued follow-up held back by a pause, then a pump suspension (the // requestAbort teardown state): the queue stays populated but undispatchable. @@ -522,7 +523,7 @@ describe("AgentSession compaction characterization", () => { // parks in the session idle wait on a checkpoint waiter. internals._schedulePostCompactionContinue(); await vi.waitFor(() => { - expect(internals._sessionInputCheckpointWaiters.size).toBeGreaterThan(0); + expect(internals._inputCheckpoints.hasWaiters).toBe(true); }); expect(session.hasPendingAdmissionWaiters).toBe(true); @@ -533,7 +534,7 @@ describe("AgentSession compaction characterization", () => { internals._cancelPostCompactionContinue(); await new Promise((resolve) => setTimeout(resolve, 100)); try { - expect(internals._sessionInputCheckpointWaiters.size).toBe(0); + expect(internals._inputCheckpoints.hasWaiters).toBe(false); expect(session.hasPendingAdmissionWaiters).toBe(false); } finally { session.clearQueue(); @@ -1125,12 +1126,6 @@ describe("AgentSession compaction characterization", () => { const sessionInternals = harness.session as unknown as { _schedulePostCompactionContinue(continueAfterSessionInput?: boolean): void; _continuation: SessionContinuation; - _createPreparedTurnAction( - schedule: "followUp", - text: string, - images: undefined, - options: { message?: AgentMessage; resumeIfIdle: boolean }, - ): unknown; _admitSessionInput(action: unknown, options?: { wake?: boolean }): { accepted: boolean }; }; const continuation = { @@ -1141,7 +1136,7 @@ describe("AgentSession compaction characterization", () => { if (tracked) sessionInternals._continuation.track(continuation); harness.setResponses([fauxAssistantMessage(response)]); sessionInternals._admitSessionInput( - sessionInternals._createPreparedTurnAction("followUp", text, undefined, { + createPreparedTurnAction("followUp", text, undefined, { ...(tracked && { message: continuation }), resumeIfIdle: tracked, }), diff --git a/packages/coding-agent/test/suite/agent-session-goal.test.ts b/packages/coding-agent/test/suite/agent-session-goal.test.ts index 969cf111dc..8343d3dc0e 100644 --- a/packages/coding-agent/test/suite/agent-session-goal.test.ts +++ b/packages/coding-agent/test/suite/agent-session-goal.test.ts @@ -7,10 +7,10 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { AgentSession } from "../../src/core/agent-session.js"; import { AuthStorage } from "../../src/core/auth-storage.js"; import type { ExtensionFactory } from "../../src/core/extensions/types.js"; -import type { GoalHostResponse } from "../../src/core/goals.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 { GOAL_STATE_CUSTOM_TYPE, type GoalHostResponse } from "../../src/session/goals/contracts.js"; import { createTestResourceLoader } from "../utilities.js"; import { conversationMessages, createHarness, getAssistantTexts, getMessageText, type Harness } from "./harness.js"; @@ -919,7 +919,6 @@ describe("initial goal seeding from config", () => { }); // Goal is persisted before first prompt - const { GOAL_STATE_CUSTOM_TYPE } = await import("../../src/core/goals.js"); const branch = harness.sessionManager.getBranch(); const goalEntry = branch.find((e) => e.type === "custom" && e.customType === GOAL_STATE_CUSTOM_TYPE); expect(goalEntry).toBeDefined(); diff --git a/packages/coding-agent/test/suite/agent-session-prompt.test.ts b/packages/coding-agent/test/suite/agent-session-prompt.test.ts index bfb656c3fd..33b7b376c4 100644 --- a/packages/coding-agent/test/suite/agent-session-prompt.test.ts +++ b/packages/coding-agent/test/suite/agent-session-prompt.test.ts @@ -1059,7 +1059,7 @@ stale post-hook extension instructions`, const agentPrompt = "Agent-to-agent message received.\nSource: agent_message\nTo: Target, active target, session session-target\nMessage id: agentmsg_handoff_busy\n\nqueue at handoff"; const sessionInternals = harness.session as unknown as { - _sessionInputCheckpointWaiters: Set<() => void>; + _inputCheckpoints: { hasWaiters: boolean }; }; const pause = harness.session.acquireQueuedWorkPause(); let acceptedSettled = false; @@ -1076,7 +1076,7 @@ stale post-hook extension instructions`, expect(harness.session.getFollowUpMessages()).toEqual([]); expect(acceptedSettled).toBe(false); - expect(sessionInternals._sessionInputCheckpointWaiters.size).toBe(1); + expect(sessionInternals._inputCheckpoints.hasWaiters).toBe(true); harness.setResponses([fauxAssistantMessage("delivered")]); pause.release(); @@ -1084,7 +1084,7 @@ stale post-hook extension instructions`, await harness.session.waitForIdle(); expect(getUserTexts(harness)).toEqual([agentPrompt]); - expect(sessionInternals._sessionInputCheckpointWaiters.size).toBe(0); + expect(sessionInternals._inputCheckpoints.hasWaiters).toBe(false); }); it("restores nextTurn context when handoff busy rejection cannot queue", async () => { @@ -1453,7 +1453,7 @@ stale post-hook extension instructions`, }, fauxAssistantMessage("second done"), ]); - const internals = harness.session as unknown as { _agentMessageOutcomes: Map }; + const internals = harness.session as unknown as { _messageDelivery: { outcomes: Map } }; const first = harness.session.prompt("first"); await vi.waitFor(() => expect(harness.session.isStreaming).toBe(true)); @@ -1464,7 +1464,7 @@ stale post-hook extension instructions`, releaseFirst?.(); await Promise.all([first, queued]); - const keys = [...internals._agentMessageOutcomes.keys()]; + const keys = [...internals._messageDelivery.outcomes.keys()]; expect(keys.filter((key) => key.startsWith("prompt-wait:"))).toEqual([]); }); @@ -1934,7 +1934,7 @@ stale post-hook extension instructions`, harness.setResponses([fauxAssistantMessage("done")]); const prompt = harness.session.prompt("pending extension event"); await extensionReached.promise; - const eventQueue = (harness.session as unknown as { _agentEventQueue: Promise })._agentEventQueue; + const eventQueue = (harness.session as unknown as { _events: { queue: Promise } })._events.queue; let queueDrained = false; void eventQueue.then(() => { queueDrained = true; @@ -1958,9 +1958,10 @@ stale post-hook extension instructions`, it("propagates a snapshotted event queue rejection without flushing", async () => { const harness = await createHarness(); harnesses.push(harness); - (harness.session as unknown as { _agentEventQueue: Promise })._agentEventQueue = Promise.reject( - new Error("event queue failed"), - ); + const internals = harness.session as unknown as { _events: { enqueue(work: () => void): void } }; + internals._events.enqueue(() => { + throw new Error("event queue failed"); + }); const flushNow = vi.spyOn(harness.sessionManager, "flushNow"); await expect(harness.session.waitForSessionInputCheckpoint()).rejects.toThrow("event queue failed"); @@ -2113,7 +2114,7 @@ stale post-hook extension instructions`, await expect(accepted).rejects.toThrow("cleared before delivery"); await expect(delivery).rejects.toThrow("cleared before delivery"); await harness.session.agent.waitForIdle(); - await (harness.session as unknown as { _agentEventQueue: Promise })._agentEventQueue; + await (harness.session as unknown as { _events: { queue: Promise } })._events.queue; const persistedAfter = harness.sessionManager.getEntries().filter((entry) => entry.type === "message").length; expect(persistedAfter).toBe(persistedBefore); expect(getUserTexts(harness)).not.toContain(clearedAgentPrompt); diff --git a/packages/coding-agent/test/suite/agent-session-queue.test.ts b/packages/coding-agent/test/suite/agent-session-queue.test.ts index b50305a380..242b5ac078 100644 --- a/packages/coding-agent/test/suite/agent-session-queue.test.ts +++ b/packages/coding-agent/test/suite/agent-session-queue.test.ts @@ -29,7 +29,7 @@ import { saveHarnessState, } from "../../src/core/refinement/index.js"; import { parseSessionSlashCommand } from "../../src/core/slash-commands.js"; -import type { SessionContinuation } from "../../src/session/continuation.js"; +import type { SessionContinuation } from "../../src/session/turns/continuation.js"; import { conversationMessages, createHarness, @@ -69,7 +69,7 @@ type AutoRefineInternals = { type SteeringStopInternals = { _steeringStopPending: boolean; - _clearQueuedGoalContexts(): void; + _goalContinuation: { clearQueuedGoalContexts(): void }; }; function emptyRefinementResult(): RefinementResult { @@ -1502,7 +1502,7 @@ describe("AgentSession queue characterization", () => { pause.release(); await hook.reached; - (harness.session as unknown as SteeringStopInternals)._clearQueuedGoalContexts(); + (harness.session as unknown as SteeringStopInternals)._goalContinuation.clearQueuedGoalContexts(); hook.release(); await harness.session.waitForIdle(); @@ -3014,12 +3014,15 @@ describe("AgentSession queue characterization", () => { harnesses.push(harness); const initialEvent = createDeferred(); const chainedOperation = createDeferred(); - const internals = harness.session as unknown as { _agentEventQueue: Promise }; - let eventQueue: Promise; - eventQueue = initialEvent.promise.then(() => { - internals._agentEventQueue = eventQueue.then(() => chainedOperation.promise); - }); - internals._agentEventQueue = eventQueue; + const internals = harness.session as unknown as { + _events: { readonly queue: Promise; enqueue(work: () => void): void }; + }; + internals._events.enqueue(() => + initialEvent.promise.then(() => { + internals._events.enqueue(() => chainedOperation.promise); + }), + ); + const eventQueue = internals._events.queue; let idle = false; const waiting = harness.session.waitForIdle().then(() => { idle = true; diff --git a/packages/coding-agent/test/suite/agent-session-retry-events.test.ts b/packages/coding-agent/test/suite/agent-session-retry-events.test.ts index ea88d3ce77..15b251bce8 100644 --- a/packages/coding-agent/test/suite/agent-session-retry-events.test.ts +++ b/packages/coding-agent/test/suite/agent-session-retry-events.test.ts @@ -3,10 +3,10 @@ import { type AssistantMessage, fauxAssistantMessage, fauxThinking, fauxToolCall import { Type } from "typebox"; import { afterEach, describe, expect, it, vi } from "vitest"; import type { CompactionResult } from "../../src/core/compaction/index.js"; -import type { SessionCompaction } from "../../src/session/compaction.js"; -import type { CompactionExecutionOptions } from "../../src/session/compaction-execution.js"; -import type { SessionContinuation } from "../../src/session/continuation.js"; -import type { SessionRetry } from "../../src/session/retry.js"; +import type { SessionCompaction } from "../../src/session/compaction/compaction.js"; +import type { CompactionExecutionOptions } from "../../src/session/compaction/compaction-execution.js"; +import type { SessionContinuation } from "../../src/session/turns/continuation.js"; +import type { SessionRetry } from "../../src/session/turns/retry.js"; import { createHarness, type Harness } from "./harness.js"; function normalizeEventOrder(events: Harness["events"]): string[] { @@ -60,7 +60,7 @@ type SessionRetryCompactionInternals = { _compaction: SessionCompaction; _performCompaction(options: CompactionExecutionOptions): Promise; _continuation: SessionContinuation; - _processAgentEvent: (event: AgentEvent) => Promise; + _events: { processAgentEvent(event: AgentEvent): Promise }; _checkCompaction: (message: AssistantMessage) => Promise; _schedulePostCompactionContinue: () => void; _cancelPostCompactionContinue: () => void; @@ -417,7 +417,7 @@ describe("AgentSession retry and event characterization", () => { internals._checkCompaction = async () => true; try { - await internals._processAgentEvent({ type: "agent_end", messages: [overflowMessage] } as AgentEvent); + await internals._events.processAgentEvent({ type: "agent_end", messages: [overflowMessage] } as AgentEvent); expect(harness.session.retryAttempt).toBe(1); expect(harness.session.isRetrying).toBe(true); diff --git a/packages/coding-agent/test/suite/agent-session-serialized-refine.test.ts b/packages/coding-agent/test/suite/agent-session-serialized-refine.test.ts index 7ba5cedba7..726ce5f216 100644 --- a/packages/coding-agent/test/suite/agent-session-serialized-refine.test.ts +++ b/packages/coding-agent/test/suite/agent-session-serialized-refine.test.ts @@ -3,6 +3,7 @@ import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai"; import { Type } from "typebox"; import { afterEach, describe, expect, it, vi } from "vitest"; import { getLocalHarnessStateDir, loadHarnessState, saveHarnessState } from "../../src/core/refinement/index.js"; +import { createPreparedTurnAction } from "../../src/session/prepared-actions.js"; import { createHarness, getMessageText, type Harness } from "./harness.js"; type SerializedInternals = { @@ -58,23 +59,20 @@ type SerializedInternals = { newMessages: unknown[]; }): Promise; - _createPreparedTurnAction( - schedule: "steer", - text: string, - images: undefined, - options: Record, - ): unknown; _admitSessionInput(action: unknown, options?: { wake?: boolean }): { accepted: boolean }; _disposing: boolean; _disposed: boolean; _checkCompaction(message: unknown): Promise; - _lastAssistantMessage: unknown; - _handleAgentEvent(event: { type: string; messages?: unknown[] }): void; - _agentEventQueue: Promise; + _events: { + lastAssistant: unknown; + handleAgentEvent(event: { type: string; messages?: unknown[] }): void; + readonly queue: Promise; + enqueue(work: () => void): void; + }; + _turnPolicy: { shouldStopForThresholdCompaction(ctx: unknown): Promise }; - _branchSummaryOperation?: Promise | undefined; requestAbort(): void; abortCompaction(): void; abortBranchSummary(): void; @@ -335,12 +333,9 @@ describe("Serialized agent-callable refine", () => { const internals = harness.session as unknown as SerializedInternals; const { applyRefine } = mockSerializedRefine(harness); internals._refinement._pendingRequestedRefine = { instructions: "capture a lesson" }; - internals._admitSessionInput(internals._createPreparedTurnAction("steer", "steer", undefined, {})); + internals._admitSessionInput(createPreparedTurnAction("steer", "steer", undefined, {})); const compactionSpy = vi - .spyOn( - internals as unknown as { _shouldStopForThresholdCompaction: (ctx: unknown) => Promise }, - "_shouldStopForThresholdCompaction", - ) + .spyOn(internals._turnPolicy, "shouldStopForThresholdCompaction") .mockResolvedValue(false); const shouldStop = await internals._shouldStopAfterTurn(makeCtx("turn")); @@ -821,10 +816,10 @@ describe("Serialized refine review-fix regressions", () => { const refine = vi.spyOn(internals._refinement, "refine").mockResolvedValue(emptyRefinementResult()); const schedule = vi.spyOn(internals._refinement._auto, "_scheduleAutoRefineAfterAgentEnd"); const assistant = fauxAssistantMessage("done"); - internals._lastAssistantMessage = assistant; + internals._events.lastAssistant = assistant; - internals._handleAgentEvent({ type: "agent_end", messages: [assistant] }); - await internals._agentEventQueue; + internals._events.handleAgentEvent({ type: "agent_end", messages: [assistant] }); + await internals._events.queue; await vi.waitFor(() => expect(refine).toHaveBeenCalledOnce()); expect(schedule).not.toHaveBeenCalled(); @@ -867,7 +862,7 @@ describe("Serialized refine review-fix regressions", () => { // Now simulate the actual agent_end event path. // Set _lastAssistantMessage so agent_end has a non-error assistant to process. const fauxAssistant = fauxAssistantMessage("done"); - internals._lastAssistantMessage = fauxAssistant; + internals._events.lastAssistant = fauxAssistant; // Spy on _scheduleAutoRefineAfterAgentEnd — it should NOT be called in serialized mode. const scheduleSpy = vi.spyOn(internals._refinement._auto, "_scheduleAutoRefineAfterAgentEnd"); @@ -875,8 +870,8 @@ describe("Serialized refine review-fix regressions", () => { vi.spyOn(internals, "_checkCompaction").mockResolvedValue(false); // Drive the real agent_end event through _handleAgentEvent → _processAgentEvent. - internals._handleAgentEvent({ type: "agent_end", messages: [fauxAssistant] }); - await internals._agentEventQueue; + internals._events.handleAgentEvent({ type: "agent_end", messages: [fauxAssistant] }); + await internals._events.queue; // Flush any setTimeout(0) that _scheduleAutoRefine would have used. await new Promise((resolve) => setTimeout(resolve, 20)); @@ -1245,14 +1240,33 @@ describe("Serialized refine review-fix regressions", () => { }); it("public refine waits for active branch summary without aborting it", async () => { - const harness = await createHarness({ persistSession: true }); - harnesses.push(harness); - const internals = harness.session as unknown as SerializedInternals; let releaseBranchSummary: () => void = () => {}; - const branchSummaryOperation = new Promise((resolve) => { + const branchSummaryGate = new Promise((resolve) => { releaseBranchSummary = resolve; }); - internals._branchSummaryOperation = branchSummaryOperation; + let summaryStarted: () => void = () => {}; + const summaryReady = new Promise((resolve) => { + summaryStarted = resolve; + }); + const harness = await createHarness({ + persistSession: true, + extensionFactories: [ + (pi) => { + pi.on("session_before_tree", async () => { + summaryStarted(); + await branchSummaryGate; + return { cancel: true }; + }); + }, + ], + }); + harnesses.push(harness); + const internals = harness.session as unknown as SerializedInternals; + const target = harness.sessionManager.appendMessage({ role: "user", content: "branch target", timestamp: 1 }); + harness.sessionManager.appendMessage(fauxAssistantMessage("answer")); + const navigation = harness.session.navigateTree(target); + await summaryReady; + expect(harness.session.isCompacting).toBe(true); vi.spyOn(internals._refinement._execution, "_planRefine").mockResolvedValue({ id: "public-plan", proposal: { edits: [] }, @@ -1260,19 +1274,17 @@ describe("Serialized refine review-fix regressions", () => { const applyRefine = vi .spyOn(internals._refinement._execution, "_applyRefine") .mockResolvedValue(emptyRefinementResult()); - const abortBranchSummary = vi.spyOn(internals, "abortBranchSummary"); - + const abortBranchSummary = vi.spyOn(harness.session, "abortBranchSummary"); const refinePromise = harness.session.refine({ instructions: "public" }); await vi.waitFor(() => expect(internals._refinement._refineInFlight).toBeDefined()); - expect(abortBranchSummary).not.toHaveBeenCalled(); expect(applyRefine).not.toHaveBeenCalled(); expect(internals._refinement._refineAbortController?.signal.aborted).toBe(false); releaseBranchSummary(); + await expect(navigation).resolves.toEqual({ cancelled: true }); await refinePromise; - expect(applyRefine).toHaveBeenCalledOnce(); - internals._branchSummaryOperation = undefined; + expect(harness.session.isCompacting).toBe(false); internals._refinement._refineAbortController = undefined; }); @@ -1297,9 +1309,10 @@ describe("Serialized refine review-fix regressions", () => { await vi.waitFor(() => expect(waitForIdle).toHaveBeenCalledOnce()); let releaseEventQueue: () => void = () => {}; - internals._agentEventQueue = new Promise((resolve) => { + const eventQueue = new Promise((resolve) => { releaseEventQueue = resolve; }); + internals._events.enqueue(() => eventQueue); let releaseCompaction: () => void = () => {}; const compactionOperation = new Promise((resolve) => { releaseCompaction = resolve; @@ -1665,7 +1678,7 @@ describe("Serialized refine event-ordering integration", () => { it("real message_end/turn_end ordering: threshold planning starts and applies before next model turn", async () => { // This test uses real agent.prompt() with faux responses to verify - // that _shouldStopAfterTurn's await this._agentEventQueue ensures + // that _shouldStopAfterTurn's await this._events.queue ensures // the message_end counter increment and background plan kickoff // happen BEFORE the serialized checkpoint runs. const reviewer = vi.fn(async () => ({ @@ -1749,7 +1762,7 @@ describe("P0 concurrency regressions", () => { }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals; - internals._admitSessionInput(internals._createPreparedTurnAction("steer", "steer", undefined, {})); + internals._admitSessionInput(createPreparedTurnAction("steer", "steer", undefined, {})); let planResolved = false; let applyFinished = false; @@ -1769,10 +1782,7 @@ describe("P0 concurrency regressions", () => { // Spy on _shouldStopForThresholdCompaction: assert apply finished // when compaction check runs, return true to simulate compaction firing. const compactionSpy = vi - .spyOn( - internals as unknown as { _shouldStopForThresholdCompaction: (ctx: unknown) => Promise }, - "_shouldStopForThresholdCompaction", - ) + .spyOn(internals._turnPolicy, "shouldStopForThresholdCompaction") .mockImplementation(async () => { expect(applyFinished).toBe(true); return true; @@ -2112,9 +2122,9 @@ describe("P0 concurrency regressions", () => { // Simulate an aborted assistant message arriving at agent_end. // Do NOT mock _checkCompaction — the real aborted path must clear the pending refine. const abortedAssistant = fauxAssistantMessage("aborted turn", { stopReason: "aborted" }); - internals._lastAssistantMessage = abortedAssistant; - internals._handleAgentEvent({ type: "agent_end", messages: [abortedAssistant] }); - await internals._agentEventQueue; + internals._events.lastAssistant = abortedAssistant; + internals._events.handleAgentEvent({ type: "agent_end", messages: [abortedAssistant] }); + await internals._events.queue; await new Promise((resolve) => setTimeout(resolve, 20)); // _checkCompaction's aborted block cleared _pendingRequestedRefine. @@ -2148,9 +2158,9 @@ describe("P0 concurrency regressions", () => { const toolUseAssistant = fauxAssistantMessage([fauxToolCall("ipython", { code: "await refine.run()" })], { stopReason: "toolUse", }); - internals._lastAssistantMessage = toolUseAssistant; - internals._handleAgentEvent({ type: "agent_end", messages: [toolUseAssistant] }); - await internals._agentEventQueue; + internals._events.lastAssistant = toolUseAssistant; + internals._events.handleAgentEvent({ type: "agent_end", messages: [toolUseAssistant] }); + await internals._events.queue; await new Promise((resolve) => setTimeout(resolve, 20)); expect(refineSpy).not.toHaveBeenCalled(); @@ -2421,9 +2431,9 @@ describe("P0 concurrency regressions", () => { // Simulate an aborted assistant message arriving at agent_end. // Do NOT mock _checkCompaction — the real aborted path must clear the pending refine. const abortedAssistant = fauxAssistantMessage("aborted turn", { stopReason: "aborted" }); - internals._lastAssistantMessage = abortedAssistant; - internals._handleAgentEvent({ type: "agent_end", messages: [abortedAssistant] }); - await internals._agentEventQueue; + internals._events.lastAssistant = abortedAssistant; + internals._events.handleAgentEvent({ type: "agent_end", messages: [abortedAssistant] }); + await internals._events.queue; await new Promise((resolve) => setTimeout(resolve, 20)); // _checkCompaction's aborted block cleared _pendingRequestedRefine. @@ -2674,7 +2684,7 @@ describe("P0 concurrency regressions", () => { const harness = await createHarness({ persistSession: true }); harnesses.push(harness); const internals = harness.session as unknown as SerializedInternals & { - _goalAbortInProgress: boolean; + _goalContinuation: { abortInProgress: boolean }; _cancelActiveRlmChildRuns: () => void; }; @@ -2690,6 +2700,6 @@ describe("P0 concurrency regressions", () => { expect(requestAbort).not.toHaveBeenCalled(); expect(cancelChildRuns).not.toHaveBeenCalled(); - expect(internals._goalAbortInProgress).toBe(false); + expect(internals._goalContinuation.abortInProgress).toBe(false); }); }); 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 8377a52fb5..32a5565d5b 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,7 +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 type { SessionKernel } from "../../../src/session/kernel/kernel.js"; import { createHarness, type Harness } from "../harness.js"; const runtimeDir = resolve(__dirname, "../../../../../prime-agent-runtime"); diff --git a/packages/coding-agent/test/suite/regressions/2190-refinement-dispatch.test.ts b/packages/coding-agent/test/suite/regressions/2190-refinement-dispatch.test.ts index c23425119f..6ef31a45bd 100644 --- a/packages/coding-agent/test/suite/regressions/2190-refinement-dispatch.test.ts +++ b/packages/coding-agent/test/suite/regressions/2190-refinement-dispatch.test.ts @@ -1,6 +1,6 @@ import { afterEach, describe, expect, it, vi } from "vitest"; -import type { AutoRefinement } from "../../../src/session/auto-refinement.js"; -import type { SessionRefinement, SessionRefinementHost } from "../../../src/session/refinement.js"; +import type { AutoRefinement } from "../../../src/session/refinement/auto-refinement.js"; +import type { SessionRefinement, SessionRefinementHost } from "../../../src/session/refinement/refinement.js"; import { createHarness, type Harness } from "../harness.js"; import { withStreaming } from "../scheduling.js"; diff --git a/packages/coding-agent/test/suite/regressions/4257-update-restart-resume.test.ts b/packages/coding-agent/test/suite/regressions/4257-update-restart-resume.test.ts index ddfd58b603..6d60309aeb 100644 --- a/packages/coding-agent/test/suite/regressions/4257-update-restart-resume.test.ts +++ b/packages/coding-agent/test/suite/regressions/4257-update-restart-resume.test.ts @@ -649,17 +649,17 @@ describe("issue #4257 update restart resume", () => { }); const waitForIdleSpy = vi.spyOn(harness.session.agent, "waitForIdle").mockReturnValue(idlePromise); const agentAbortSpy = vi.spyOn(harness.session.agent, "abort"); - const internals = harness.session as unknown as { _goalAbortInProgress: boolean }; + const internals = harness.session as unknown as { _goalContinuation: { abortInProgress: boolean } }; harness.session.abortForUpdateRestart(); expect(agentAbortSpy).toHaveBeenCalledOnce(); - expect(internals._goalAbortInProgress).toBe(true); + expect(internals._goalContinuation.abortInProgress).toBe(true); releaseIdle?.(); - await waitForCondition(() => !internals._goalAbortInProgress); + await waitForCondition(() => !internals._goalContinuation.abortInProgress); - expect(internals._goalAbortInProgress).toBe(false); + expect(internals._goalContinuation.abortInProgress).toBe(false); waitForIdleSpy.mockRestore(); agentAbortSpy.mockRestore(); }); diff --git a/packages/coding-agent/test/suite/regressions/4435-auth-error-login-guidance.test.ts b/packages/coding-agent/test/suite/regressions/4435-auth-error-login-guidance.test.ts index b56fb1f84e..bf056255bc 100644 --- a/packages/coding-agent/test/suite/regressions/4435-auth-error-login-guidance.test.ts +++ b/packages/coding-agent/test/suite/regressions/4435-auth-error-login-guidance.test.ts @@ -40,10 +40,10 @@ describe("issue #4435 auth error login guidance", () => { }); const event = { type: "agent_end", messages: [message] } as AgentEvent; const session = harness.session as unknown as { - _addLoginGuidanceToAuthError(event: AgentEvent): void; + _events: { addLoginGuidanceToAuthError(event: AgentEvent): void }; }; - session._addLoginGuidanceToAuthError(event); + session._events.addLoginGuidanceToAuthError(event); expect(message.errorMessage).toContain("Run /login to update credentials."); }); diff --git a/packages/coding-agent/test/suite/regressions/4482-heartbeat-injected-prompt.test.ts b/packages/coding-agent/test/suite/regressions/4482-heartbeat-injected-prompt.test.ts index a9bc0e9e3e..6dc72fefe0 100644 --- a/packages/coding-agent/test/suite/regressions/4482-heartbeat-injected-prompt.test.ts +++ b/packages/coding-agent/test/suite/regressions/4482-heartbeat-injected-prompt.test.ts @@ -4,7 +4,6 @@ import { Container, type MarkdownTheme } from "@earendil-works/pi-tui"; import { Type } from "typebox"; import { afterEach, beforeAll, describe, expect, it, vi } from "vitest"; import { type AgentCronJob, shouldDeferHeartbeatCronJob } from "../../../src/core/cron-jobs.js"; -import { createGoalContextMessage, type GoalState } from "../../../src/core/goals.js"; import { type CustomMessage, createHeartbeatPromptMessage, @@ -17,7 +16,8 @@ import { } from "../../../src/modes/interactive/components/injected-prompt-message.js"; import { InteractiveMode } from "../../../src/modes/interactive/interactive-mode.js"; import { getMarkdownTheme, initTheme } from "../../../src/modes/interactive/theme/theme.js"; -import type { SessionCompaction } from "../../../src/session/compaction.js"; +import type { SessionCompaction } from "../../../src/session/compaction/compaction.js"; +import { createGoalContextMessage, type GoalState } from "../../../src/session/goals/contracts.js"; import { conversationMessages, createHarness, getMessageText, getUserTexts, type Harness } from "../harness.js"; type AddMessageToChatHost = { diff --git a/packages/coding-agent/test/suite/regressions/4491-provider-stale-after-401.test.ts b/packages/coding-agent/test/suite/regressions/4491-provider-stale-after-401.test.ts index 55d014667f..36090ec47d 100644 --- a/packages/coding-agent/test/suite/regressions/4491-provider-stale-after-401.test.ts +++ b/packages/coding-agent/test/suite/regressions/4491-provider-stale-after-401.test.ts @@ -3,7 +3,7 @@ import { type AssistantMessage, fauxAssistantMessage } from "@earendil-works/pi- import { afterEach, describe, expect, it, vi } from "vitest"; import type { AgentSessionRuntime } from "../../../src/core/agent-session-runtime.js"; import { InProcessAgentConnection } from "../../../src/modes/agent-connection/in-process-agent-connection.js"; -import type { SessionRetry } from "../../../src/session/retry.js"; +import type { SessionRetry } from "../../../src/session/turns/retry.js"; import { createHarness, type Harness } from "../harness.js"; function structuredFailureMessage(kind: string, status: number, errorMessage: string): AssistantMessage { @@ -335,11 +335,11 @@ describe("issue #4491 provider stale after repeated 401", () => { const event = { type: "agent_end", messages: [message] } as AgentEvent; const session = harness.session as unknown as { _retry: SessionRetry; - _processAgentEvent(event: AgentEvent): Promise; + _events: { processAgentEvent(event: AgentEvent): Promise }; }; session._retry.observeAgentEnd(event); - await session._processAgentEvent(event); + await session._events.processAgentEvent(event); expect(harness.session.isRetrying).toBe(false); expect(harness.eventsOfType("auto_retry_end").map((retryEvent) => retryEvent.success)).toEqual([false]); 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 7968c8ccd6..c86286429c 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 @@ -12,7 +12,7 @@ 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 type { SessionKernel } from "../../../src/session/kernel/kernel.js"; import { conversationMessages, createHarness, getMessageText, getUserTexts, type Harness } from "../harness.js"; type StateRestoreHost = { diff --git a/packages/coding-agent/test/suite/regressions/4531-agent-message-ui.test.ts b/packages/coding-agent/test/suite/regressions/4531-agent-message-ui.test.ts index e98fe87c0b..2db7746488 100644 --- a/packages/coding-agent/test/suite/regressions/4531-agent-message-ui.test.ts +++ b/packages/coding-agent/test/suite/regressions/4531-agent-message-ui.test.ts @@ -54,8 +54,8 @@ function render(component: AgentMessageComponent): string { type LateSentAgentMessageHost = { _recordLateIpythonSentAgentMessage: (toolCallId: string, message: KernelSentAgentMessage) => void; - _agentEventQueue: Promise; - _lateIpythonSentAgentMessages: Map; + _events: { queue: Promise }; + _messageDelivery: { lateMessages: Map }; _restoreLateIpythonSentAgentMessages: () => void; }; @@ -202,7 +202,7 @@ describe("ENG-4531 agent message UI", () => { const host = harness.session as unknown as LateSentAgentMessageHost; host._recordLateIpythonSentAgentMessage(toolResult.toolCallId, lateMessage); - await host._agentEventQueue; + await host._events.queue; unsubscribe(); expect(toolResult.details).toMatchObject({ sentAgentMessages: [lateMessage] }); @@ -226,11 +226,11 @@ describe("ENG-4531 agent message UI", () => { expect(toolResult.details).toMatchObject({ sentAgentMessages: [lateMessage] }); toolResult.details = { status: "ok" }; - host._lateIpythonSentAgentMessages = new Map(); + host._messageDelivery.lateMessages = new Map(); host._restoreLateIpythonSentAgentMessages(); expect(toolResult.details).toMatchObject({ sentAgentMessages: [lateMessage] }); - host._lateIpythonSentAgentMessages.set("ipython_other_branch", [ + host._messageDelivery.lateMessages.set("ipython_other_branch", [ { id: "agentmsg_other_branch", message: "Stale branch receipt.", @@ -239,7 +239,7 @@ describe("ENG-4531 agent message UI", () => { }, ]); host._restoreLateIpythonSentAgentMessages(); - expect(host._lateIpythonSentAgentMessages.has("ipython_other_branch")).toBe(false); + expect(host._messageDelivery.lateMessages.has("ipython_other_branch")).toBe(false); }); it("preserves the custom message when direct delivery races with active work", async () => { diff --git a/packages/coding-agent/test/suite/regressions/4685-daemon-client-modes.test.ts b/packages/coding-agent/test/suite/regressions/4685-daemon-client-modes.test.ts index 1dee7be89c..12c6b5d055 100644 --- a/packages/coding-agent/test/suite/regressions/4685-daemon-client-modes.test.ts +++ b/packages/coding-agent/test/suite/regressions/4685-daemon-client-modes.test.ts @@ -209,9 +209,9 @@ describe("ENG-4685 daemon-backed client modes", () => { harnesses.push(harness); const state = ( harness.session as unknown as { - _autonomousState: AutonomousRuntimeState; + _autonomousContinuation: { state: AutonomousRuntimeState }; } - )._autonomousState; + )._autonomousContinuation.state; state.gateAttempts[gate] = 1; state.lastGateFailure = { command: gate, diff --git a/packages/coding-agent/test/suite/regressions/5938-session-finishing-boundaries.test.ts b/packages/coding-agent/test/suite/regressions/5938-session-finishing-boundaries.test.ts new file mode 100644 index 0000000000..be7f5a3bdd --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/5938-session-finishing-boundaries.test.ts @@ -0,0 +1,180 @@ +import type { Api, Model } from "@earendil-works/pi-ai"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import type { AgentSessionMessageController, AgentSessionMessageReceipt } from "../../../src/core/agent-messages.js"; +import type { AgentObserveAgentSnapshot, AgentObserveController } from "../../../src/core/agent-observe.js"; +import type { AgentRlmHeartbeatController } from "../../../src/core/cron-jobs.js"; +import type { SessionTurnExecution } from "../../../src/session/turns/turn-execution.js"; +import { createHarness, type Harness } from "../harness.js"; + +const harnesses: Harness[] = []; +afterEach(() => { + vi.restoreAllMocks(); + while (harnesses.length) harnesses.pop()?.cleanup(); +}); + +describe("session finishing boundaries", () => { + it("keeps compact status on live public usage and rejects invalid instructions synchronously", async () => { + const harness = await createHarness({ tools: [] }); + harnesses.push(harness); + vi.spyOn(harness.session, "getContextUsage").mockReturnValue({ tokens: 7, contextWindow: 70, percent: 10 }); + expect(harness.session.handleCompactHostRequest("compact.status")).toEqual({ + tokens: 7, + context_window: 70, + percent: 10, + scheduled: false, + }); + expect(() => harness.session.handleCompactHostRequest("compact.run", { instructions: 9 })).toThrow( + "compact.run instructions must be a string when provided", + ); + expect(harness.session.handleCompactHostRequest("compact.run")).toMatchObject({ scheduled: false }); + }); + + it("keeps the heartbeat controller captured before payload getters replace the binding", async () => { + const harness = await createHarness({ tools: [] }); + harnesses.push(harness); + const unavailable = () => { + throw new Error("unused heartbeat operation"); + }; + const first: AgentRlmHeartbeatController = { + listRlmHeartbeats(options) { + expect(this).toBe(first); + expect(options).toEqual({ includeInactive: true }); + return []; + }, + createRlmHeartbeat: unavailable, + updateRlmHeartbeat: unavailable, + deleteRlmHeartbeat: unavailable, + }; + const second: AgentRlmHeartbeatController = { ...first, listRlmHeartbeats: vi.fn(() => []) }; + harness.session.setRlmHeartbeatController(first); + const result = harness.session.handleRlmHeartbeatHostRequest("rlm_heartbeat.list", { + get include_inactive() { + harness.session.setRlmHeartbeatController(second); + return true; + }, + }); + expect(result).toEqual({ heartbeats: [] }); + expect(second.listRlmHeartbeats).not.toHaveBeenCalled(); + }); + + it("preserves observation controller capture, receiver and the returned promise", async () => { + const failure = new Error("observation failed"); + const pending = Promise.reject(failure); + void pending.catch(() => {}); + const unused = () => { + throw new Error("unused observation operation"); + }; + const controller: AgentObserveController = { + listAgents: unused, + recentMessages: unused, + getAgent(target) { + expect(this).toBe(controller); + expect(target).toBe("child"); + return pending; + }, + }; + const harness = await createHarness({ tools: [], agentObserveController: controller }); + harnesses.push(harness); + const result = harness.session.handleAgentObserveHostRequest("agent_observe.get", { + get target() { + Object.defineProperty(harness.session, "_agentObserveController", { value: undefined }); + return "child"; + }, + }); + expect(result).toBe(pending); + await expect(result).rejects.toBe(failure); + }); + + it("keeps each messaging-controller read live and preserves receiver and promise identity", async () => { + const failure = new Error("send failed"); + const pending = Promise.reject(failure); + void pending.catch(() => {}); + const unused = () => { + throw new Error("unused messaging operation"); + }; + const first: AgentSessionMessageController = { listAgents: unused, sendAgentMessage: unused }; + const second: AgentSessionMessageController = { + ...first, + sendAgentMessage(input) { + expect(this).toBe(second); + expect(input).toEqual({ target: "child", message: "hello" }); + return pending; + }, + }; + const harness = await createHarness({ tools: [], agentMessageController: first }); + harnesses.push(harness); + let reads = 0; + Object.defineProperty(harness.session, "_agentMessageController", { + get: () => (++reads === 1 ? first : second), + }); + const result = harness.session.handleAgentMessageHostRequest("agent_message.send", { + target: "child", + message: "hello", + }); + expect(reads).toBe(2); + expect(result).toBe(pending); + await expect(result).rejects.toBe(failure); + }); + + it("uses the live public parent model for child defaults without an auth request", async () => { + const harness = await createHarness({ tools: [], models: [{ id: "one" }, { id: "two" }] }); + harnesses.push(harness); + const parent = harness.getModel("two")!; + vi.spyOn(harness.session, "model", "get").mockReturnValue(parent); + const auth = vi.spyOn(harness.session.modelRegistry, "getApiKeyAndHeaders"); + const catalog = vi.spyOn(harness.session.modelRegistry, "getExecutableModels"); + const session = harness.session as unknown as { + _resolveRlmSubagentModel(reference?: string): Promise<{ model: Model }>; + }; + expect((await session._resolveRlmSubagentModel()).model).toBe(parent); + expect((await session._resolveRlmSubagentModel(`${parent.provider}/${parent.id}`)).model).toBe(parent); + expect(auth).not.toHaveBeenCalled(); + expect(catalog).not.toHaveBeenCalled(); + }); + it.each(["facade", "prepared turn"])( + "refreshes %s prompts with the final live base and literal replacement text", + async (path) => { + const harness = await createHarness({ tools: [] }); + harnesses.push(harness); + const session = harness.session as unknown as { + _tools: { baseSystemPrompt: string }; + _turnExecution: Pick; + _refreshExtensionSystemPrompt(prompt: string, snapshot: string): string; + }; + let reads = 0; + Object.defineProperty(session._tools, "baseSystemPrompt", { + configurable: true, + get: () => (++reads === 1 ? "intermediate" : "$& final base"), + }); + if (path === "facade") { + expect(session._refreshExtensionSystemPrompt("prefix old suffix old", "old")).toBe( + "prefix $& final base suffix old", + ); + } else { + session._turnExecution.applyPreparedSystemPrompt( + { basePromptSnapshot: "old", result: { systemPrompt: "prefix old suffix old" } }, + true, + ); + expect(harness.session.agent.state.systemPrompt).toBe("prefix $& final base suffix old"); + } + expect(reads).toBe(2); + }, + ); + + it("headless waiting re-reads the public idle method after continuation settlement", async () => { + const harness = await createHarness({ tools: [] }); + harnesses.push(harness); + const session = harness.session as unknown as { + _continuation: { readonly current?: { promise: Promise } }; + }; + const failure = new Error("replacement idle failed"); + const replacement = vi.fn().mockRejectedValue(failure); + const initial = vi.spyOn(harness.session, "waitForIdle").mockImplementationOnce(async () => { + harness.session.waitForIdle = replacement; + }); + vi.spyOn(session._continuation, "current", "get").mockReturnValue({ promise: Promise.resolve() }); + await expect(harness.session.waitForHeadlessIdle()).rejects.toBe(failure); + expect(initial).toHaveBeenCalledTimes(1); + expect(replacement).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/coding-agent/test/suite/regressions/5938-session-input-turns.test.ts b/packages/coding-agent/test/suite/regressions/5938-session-input-turns.test.ts new file mode 100644 index 0000000000..8aa5f548eb --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/5938-session-input-turns.test.ts @@ -0,0 +1,220 @@ +import { fauxAssistantMessage } from "@earendil-works/pi-ai"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { AGENT_MESSAGE_SOURCE, createAgentSessionMessage } from "../../../src/core/agent-messages.js"; +import type { SessionAutonomousContinuation } from "../../../src/session/turns/autonomous-continuation.js"; +import { createHarness, getUserTexts, type Harness } from "../harness.js"; +import { createDeferred } from "../scheduling.js"; + +describe("session input and turn ownership", () => { + const harnesses: Harness[] = []; + afterEach(() => { + vi.restoreAllMocks(); + while (harnesses.length) harnesses.pop()?.cleanup(); + }); + + it("resolves public admission wrappers installed after construction and waits for completion", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const started = createDeferred(); + const response = createDeferred(); + harness.setResponses([ + async () => { + started.resolve(); + await response.promise; + return fauxAssistantMessage("complete"); + }, + ]); + const admit = vi.spyOn(harness.session, "promptUntilAccepted"); + let completed = false; + const prompt = harness.session.promptAndWait("work", { agentMessageId: "correlation" }).then(() => { + completed = true; + }); + await started.promise; + expect(admit).toHaveBeenCalledWith("work", { agentMessageId: "correlation" }); + expect(completed).toBe(false); + response.resolve(); + await prompt; + expect(completed).toBe(true); + }); + + it("keeps correlation waiters isolated across sessions and rejects only disposed work", async () => { + const first = await createHarness(); + const second = await createHarness(); + harnesses.push(first, second); + const firstDelivery = first.session.waitForAgentMessagePromptDelivery("same-id"); + const secondDelivery = second.session.waitForAgentMessagePromptDelivery("same-id"); + const firstRejected = expect(firstDelivery).rejects.toThrow("disposed"); + first.session.dispose(); + await firstRejected; + second.setResponses([fauxAssistantMessage("received")]); + await second.session.promptAndWait("second session", { agentMessageId: "same-id" }); + await expect(secondDelivery).resolves.toBeUndefined(); + expect(getUserTexts(first)).toEqual([]); + expect(getUserTexts(second)).toEqual(["second session"]); + }); + + it("uses late autonomous status wrappers for both rendered and persisted status", async () => { + const harness = await createHarness({ persistSession: true }); + harnesses.push(harness); + const session = harness.session; + const original = session.getAutonomousStatus; + let calls = 0; + const status = vi.spyOn(session, "getAutonomousStatus").mockImplementation(function (this: typeof session) { + expect(this).toBe(session); + return { ...original.call(this), continuationsUsed: 40 + ++calls }; + }); + + await session.prompt("/autonomous status"); + await session.waitForIdle(); + + expect(status).toHaveBeenCalledTimes(2); + const message = session.messages.find( + (entry) => entry.role === "custom" && entry.customType === "autonomous_status", + ); + expect(message).toMatchObject({ + content: expect.stringContaining("Continuations: 41/"), + details: { continuationsUsed: 42 }, + }); + expect(harness.sessionManager.getEntries()).toContainEqual( + expect.objectContaining({ + type: "custom_message", + customType: "autonomous_status", + content: expect.stringContaining("Continuations: 41/"), + details: expect.objectContaining({ continuationsUsed: 42 }), + }), + ); + }); + + it("surfaces a status wrapper failure without emitting or persisting autonomous status", async () => { + const harness = await createHarness({ persistSession: true }); + harnesses.push(harness); + vi.spyOn(harness.session, "getAutonomousStatus").mockImplementation(() => { + throw new Error("status unavailable"); + }); + + await harness.session.prompt("/autonomous status"); + await harness.session.waitForIdle(); + expect(harness.session.messages).toContainEqual( + expect.objectContaining({ content: "Command failed: status unavailable" }), + ); + + expect(harness.session.messages).not.toContainEqual(expect.objectContaining({ customType: "autonomous_status" })); + expect(harness.sessionManager.getEntries()).not.toContainEqual( + expect.objectContaining({ customType: "autonomous_status" }), + ); + }); + + it("reads goal.get through a getter installed after construction", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const session = harness.session; + session.handleGoalHostRequest("goal.create", { objective: "stored goal", token_budget: 100 }); + const stored = session.goalState; + const getter = vi.spyOn(session, "goalState", "get").mockImplementation(function (this: typeof session) { + expect(this).toBe(session); + return { ...stored, objective: "visible goal", tokensUsed: 30 }; + }); + + expect(session.handleGoalHostRequest("goal.get")).toMatchObject({ + goal: { objective: "visible goal", tokens_used: 30 }, + remaining_tokens: 70, + }); + expect(getter).toHaveBeenCalledOnce(); + getter.mockImplementation(() => { + throw new Error("goal unavailable"); + }); + expect(() => session.handleGoalHostRequest("goal.get")).toThrow("goal unavailable"); + }); + + it.each(["coalesce", "reject"] as const)( + "preserves late queue wrappers for suspended incoming messages: %s", + async (outcome) => { + const harness = await createHarness(); + harnesses.push(harness); + const session = harness.session; + await session.followUp("existing queued work"); + session.requestAbort(); + const queue = vi.spyOn(session, "queueAgentMessagePrompt"); + if (outcome === "coalesce") queue.mockResolvedValue(false); + else queue.mockRejectedValue(new Error("queue unavailable")); + const preflightResult = vi.fn(); + const accepted = session.acceptAgentMessagePrompt("incoming", { + queueIfBusy: true, + streamingBehavior: "followUp", + preflightResult, + }); + + if (outcome === "coalesce") { + await accepted; + expect(preflightResult).toHaveBeenCalledWith(false, false); + } else { + await expect(accepted).rejects.toThrow("queue unavailable"); + expect(preflightResult).not.toHaveBeenCalled(); + } + expect(queue).toHaveBeenCalledWith("incoming", "followUp", undefined); + expect(queue.mock.contexts).toEqual([session]); + expect(session.getFollowUpMessages()).toEqual(["existing queued work"]); + }, + ); + + it("clears agent messages by metadata and advances cancellation identity without text-wrapper dispatch", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const session = harness.session; + await session.followUp("protected queued message"); + const message = createAgentSessionMessage({ + id: "agentmsg-clear-boundary", + source: AGENT_MESSAGE_SOURCE, + message: "clear this agent message", + fromRelationship: "parent", + from: { sessionName: "root" }, + target: { activeSessionId: "active", sessionId: "session" }, + }); + await session.queueAgentMessagePrompt(message.content, "followUp", message); + const owner = (session as unknown as { _actionQueue: { clearEpoch: number } })._actionQueue; + const initialEpoch = owner.clearEpoch; + const clear = vi.spyOn(session, "clearQueuedUserMessagesMatching").mockImplementation(() => { + throw new Error("text wrapper must not handle metadata-based clearing"); + }); + + expect(session.clearQueuedAgentMessages()).toEqual({ steering: [], followUp: [message.content] }); + expect(owner.clearEpoch).toBe(initialEpoch + 1); + expect(clear).not.toHaveBeenCalled(); + expect(session.getFollowUpMessages()).toEqual(["protected queued message"]); + expect(session.clearQueuedAgentMessages()).toEqual({ steering: [], followUp: [] }); + expect(owner.clearEpoch).toBe(initialEpoch + 2); + expect(clear).not.toHaveBeenCalled(); + expect(session.getFollowUpMessages()).toEqual(["protected queued message"]); + }); + + it.each(["nextTurn", undefined] as const)( + "keeps settled custom-message timing for %s delivery", + async (deliverAs) => { + const harness = await createHarness(); + harnesses.push(harness); + const order: string[] = []; + const sent = harness.session + .sendCustomMessage({ customType: "timing", content: "context", display: false }, { deliverAs }) + .then(() => order.push("sent")); + queueMicrotask(() => order.push("tick")); + await sent; + expect(order).toEqual(["sent", "tick"]); + }, + ); + + it("does not add a promise adoption step before autonomous continuation decisions", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const owner = ( + harness.session as unknown as { _autonomousContinuation: Pick } + )._autonomousContinuation; + const order: string[] = []; + const next = owner.next(fauxAssistantMessage("done")).then((message) => { + expect(message).toBeUndefined(); + order.push("continuation"); + }); + queueMicrotask(() => order.push("tick")); + await next; + expect(order).toEqual(["continuation", "tick"]); + }); +}); diff --git a/packages/coding-agent/test/suite/regressions/5938-session-model-history.test.ts b/packages/coding-agent/test/suite/regressions/5938-session-model-history.test.ts new file mode 100644 index 0000000000..f229f7c896 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/5938-session-model-history.test.ts @@ -0,0 +1,409 @@ +import { readFileSync } from "node:fs"; +import { join } from "node:path"; +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import { fauxAssistantMessage, type Usage } from "@earendil-works/pi-ai"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { createHarnessDigestMessage } from "../../../src/core/messages.js"; +import { loadHarnessState } from "../../../src/core/refinement/index.js"; +import type { FileEntry, SessionEntry } from "../../../src/core/session-manager.js"; +import { emptyUsage, subtractAssistantUsage } from "../../../src/core/usage.js"; +import { SessionContextView } from "../../../src/session/context/context-view.js"; +import { SessionHarnessContext } from "../../../src/session/context/harness-context.js"; +import { SessionModelSelection } from "../../../src/session/models/model-selection.js"; +import { createHarness, type Harness } from "../harness.js"; + +function deferred() { + let resolve = () => {}; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +function usage(input: number): Usage { + return { ...emptyUsage(), input, totalTokens: input }; +} + +describe("session model and history ownership boundaries", () => { + const harnesses: Harness[] = []; + afterEach(() => { + vi.restoreAllMocks(); + while (harnesses.length) harnesses.pop()?.cleanup(); + }); + + it.each(["set", "available", "scoped"] as const)( + "honors a late thinking wrapper before tier changes and model hooks during %s selection", + async (selection) => { + const harness = await createHarness({ models: [{ id: "one" }, { id: "two" }] }); + harnesses.push(harness); + const { session, sessionManager, settingsManager } = harness; + const nextModel = harness.getModel("two")!; + if (selection === "scoped") session.setScopedModels(harness.models.map((model) => ({ model }))); + const events: string[] = []; + const original = session.setThinkingLevel.bind(session); + const wrapper = vi.spyOn(session, "setThinkingLevel").mockImplementation((level) => { + expect(session.model).toMatchObject({ provider: nextModel.provider, id: nextModel.id }); + expect(sessionManager.getEntries().at(-1)).toMatchObject({ type: "model_change", modelId: "two" }); + expect(settingsManager.getDefaultModel()).toBe("two"); + expect(session.serviceTier).toBe("priority"); + events.push("thinking"); + original(level); + }); + session.subscribe((event) => { + if (event.type === "service_tier_changed") events.push("tier"); + }); + vi.spyOn(session.extensionRunner, "emit").mockImplementation(async (event) => { + if (event.type === "model_select") events.push("model"); + return undefined; + }); + session.state.serviceTier = "priority"; + if (selection === "set") await session.setModel(nextModel); + else await session.cycleModel(); + expect(wrapper).toHaveBeenCalledOnce(); + expect(events).toEqual(["thinking", "tier", "model"]); + expect(session.serviceTier).toBe("default"); + }, + ); + + it.each(["set", "available", "scoped"] as const)( + "preserves partial updates and rejection when a late thinking wrapper throws during %s selection", + async (selection) => { + const harness = await createHarness({ models: [{ id: "one" }, { id: "two" }] }); + harnesses.push(harness); + const { session, sessionManager, settingsManager } = harness; + if (selection === "scoped") session.setScopedModels(harness.models.map((model) => ({ model }))); + const failure = new Error("thinking wrapper failed"); + vi.spyOn(session, "setThinkingLevel").mockImplementation(() => { + throw failure; + }); + const emit = vi.spyOn(session.extensionRunner, "emit"); + const events = vi.fn(); + session.subscribe(events); + session.state.serviceTier = "priority"; + const selectionPromise = + selection === "set" ? session.setModel(harness.getModel("two")!) : session.cycleModel(); + await expect(selectionPromise).rejects.toBe(failure); + expect(session.model?.id).toBe("two"); + expect(sessionManager.getEntries().at(-1)).toMatchObject({ type: "model_change", modelId: "two" }); + expect(settingsManager.getDefaultModel()).toBe("two"); + expect(session.serviceTier).toBe("priority"); + expect(events).not.toHaveBeenCalled(); + expect(emit).not.toHaveBeenCalled(); + }, + ); + + it("uses late capability and thinking-level overrides for cycling and preference persistence", async () => { + const harness = await createHarness({ models: [{ id: "one", reasoning: false }] }); + harnesses.push(harness); + const { session, settingsManager } = harness; + const supports = vi.spyOn(session, "supportsThinking").mockReturnValue(true); + vi.spyOn(session, "getAvailableThinkingLevels").mockReturnValue(["off", "high"]); + const original = session.setThinkingLevel.bind(session); + const wrapper = vi.spyOn(session, "setThinkingLevel").mockImplementation((level) => original(level)); + session.state.thinkingLevel = "off"; + expect(session.cycleThinkingLevel()).toBe("high"); + expect(session.thinkingLevel).toBe("high"); + expect(wrapper).toHaveBeenLastCalledWith("high"); + expect(session.cycleThinkingLevel()).toBe("off"); + expect(settingsManager.getDefaultThinkingLevel()).toBe("off"); + supports.mockReturnValue(false); + wrapper.mockClear(); + expect(session.cycleThinkingLevel()).toBeUndefined(); + expect(wrapper).not.toHaveBeenCalled(); + }); + + it("uses a late capability override to restore the saved thinking preference during model selection", async () => { + const harness = await createHarness({ + models: [ + { id: "one", reasoning: true }, + { id: "two", reasoning: true }, + ], + }); + harnesses.push(harness); + const { session, settingsManager } = harness; + session.state.thinkingLevel = "low"; + settingsManager.setDefaultThinkingLevel("high"); + vi.spyOn(session, "supportsThinking").mockReturnValue(false); + vi.spyOn(session, "getAvailableThinkingLevels").mockReturnValue(["low", "high"]); + await session.setModel(harness.getModel("two")!); + expect(session.thinkingLevel).toBe("high"); + }); + + it("uses a context-usage wrapper installed after construction for stats and tree aggregates", async () => { + const harness = await createHarness({ tools: [] }); + harnesses.push(harness); + const { session } = harness; + const original = session.getContextUsage.bind(session); + const replacement = { tokens: 42, contextWindow: 100, percent: 42 }; + const wrapper = vi.spyOn(session, "getContextUsage").mockImplementation(() => { + expect(original()).toBeDefined(); + return replacement; + }); + expect(session.getSessionStats().contextUsage).toBe(replacement); + expect(session.getContextTree().contextUsage).toBe(replacement); + expect(wrapper).toHaveBeenCalledTimes(2); + }); + + it("resolves request auth freshly and reads the current registry after replacement", async () => { + const first = await createHarness(); + const second = await createHarness(); + harnesses.push(first, second); + let registry = first.session.modelRegistry; + const owner = new SessionModelSelection( + { + getState: () => first.session.state, + getModel: () => first.session.model, + setThinkingLevel: (level) => first.session.setThinkingLevel(level), + getAvailableThinkingLevels: () => first.session.getAvailableThinkingLevels(), + supportsThinking: () => first.session.supportsThinking(), + getRegistry: () => registry, + getExtensions: () => first.session.extensionRunner, + sessionManager: first.sessionManager, + settingsManager: first.settingsManager, + emit: () => {}, + }, + "default", + [], + ); + const firstAuth = vi + .spyOn(registry, "getApiKeyAndHeaders") + .mockResolvedValue({ ok: true, apiKey: "first", headers: { team: "one" } }); + const secondAuth = vi.spyOn(second.session.modelRegistry, "getApiKeyAndHeaders").mockResolvedValue({ + ok: true, + apiKey: "second", + headers: { team: "two" }, + requestModel: second.getModel(), + }); + expect(await owner.getRequiredRequestAuth(first.getModel())).toEqual({ + apiKey: "first", + headers: { team: "one" }, + requestModel: first.getModel(), + }); + registry = second.session.modelRegistry; + expect(await owner.getRequiredRequestAuth(first.getModel())).toEqual({ + apiKey: "second", + headers: { team: "two" }, + requestModel: second.getModel(), + }); + secondAuth.mockResolvedValue({ ok: false, error: "credential refresh failed" }); + await expect(owner.getRequiredRequestAuth(first.getModel())).rejects.toThrow("credential refresh failed"); + expect(firstAuth).toHaveBeenCalledOnce(); + expect(secondAuth).toHaveBeenCalledTimes(2); + }); + + it("dispatches queued model hooks through the live extension runner", async () => { + const entered = deferred(); + const release = deferred(); + const events: string[] = []; + const first = await createHarness({ + models: [{ id: "one" }, { id: "two" }, { id: "three" }], + extensionFactories: [ + (pi) => + pi.on("model_select", async (event) => { + events.push(`old:${event.model.id}`); + entered.resolve(); + await release.promise; + }), + ], + }); + const second = await createHarness({ + extensionFactories: [ + (pi) => + pi.on("model_select", (event) => { + events.push(`new:${event.model.id}`); + }), + ], + }); + harnesses.push(first, second); + let extensions = first.session.extensionRunner; + const owner = new SessionModelSelection( + { + getState: () => first.session.state, + getModel: () => first.session.model, + setThinkingLevel: (level) => first.session.setThinkingLevel(level), + getAvailableThinkingLevels: () => first.session.getAvailableThinkingLevels(), + supportsThinking: () => first.session.supportsThinking(), + getRegistry: () => first.session.modelRegistry, + getExtensions: () => extensions, + sessionManager: first.sessionManager, + settingsManager: first.settingsManager, + emit: () => {}, + }, + "default", + [], + ); + await owner.setModel(first.getModel("two")!, { waitForExtensions: false }); + await entered.promise; + await owner.setModel(first.getModel("three")!, { waitForExtensions: false }); + const pending = owner.pendingModelSelectEmit(); + expect(pending).toBeDefined(); + extensions = second.session.extensionRunner; + release.resolve(); + await pending; + expect(events).toEqual(["old:two", "new:three"]); + expect(owner.pendingModelSelectEmit()).toBeUndefined(); + expect(first.session.model?.id).toBe("three"); + }); + + it("releases a cancelled navigation before the next queued branch mutation", async () => { + const entered = deferred(); + const calls: string[] = []; + let firstSignal: AbortSignal | undefined; + const harness = await createHarness({ + tools: [], + extensionFactories: [ + (pi) => + pi.on("session_before_tree", async (event) => { + calls.push(event.preparation.targetId); + if (calls.length !== 1) return; + firstSignal = event.signal; + entered.resolve(); + await new Promise((resolve) => + event.signal.addEventListener("abort", () => resolve(), { once: true }), + ); + return { cancel: true }; + }), + ], + }); + harnesses.push(harness); + const firstUser = harness.sessionManager.appendMessage({ role: "user", content: "first", timestamp: 1 }); + harness.sessionManager.appendMessage(fauxAssistantMessage("first answer")); + const secondUser = harness.sessionManager.appendMessage({ role: "user", content: "second", timestamp: 3 }); + harness.sessionManager.appendMessage(fauxAssistantMessage("second answer")); + const originalLeaf = harness.sessionManager.getLeafId(); + const first = harness.session.navigateTree(firstUser); + await entered.promise; + const second = harness.session.navigateTree(secondUser); + expect(calls).toEqual([firstUser]); + expect(harness.session.isCompacting).toBe(true); + expect(harness.sessionManager.getLeafId()).toBe(originalLeaf); + harness.session.abortBranchSummary(); + expect(firstSignal?.aborted).toBe(true); + await expect(first).resolves.toEqual({ cancelled: true }); + await expect(second).resolves.toMatchObject({ cancelled: false, editorText: "second" }); + expect(calls).toEqual([firstUser, secondUser]); + expect(harness.session.isCompacting).toBe(false); + await expect(harness.session.navigateTree("missing")).rejects.toThrow("Entry missing not found"); + await expect(harness.session.navigateTree(harness.sessionManager.getLeafId()!)).resolves.toEqual({ + cancelled: false, + }); + }); + + it("invalidates own spend for live attribution and equal-length entry replacement", async () => { + const harness = await createHarness({ tools: [] }); + harnesses.push(harness); + const message = { ...fauxAssistantMessage("answer"), usage: usage(100) }; + harness.sessionManager.appendMessage(message); + let entries = harness.sessionManager.getEntries(); + let unindexed: Usage | undefined; + const owner = new SessionContextView({ + sessionManager: { + getEntries: () => entries, + getBranch: () => entries, + getSessionId: () => harness.session.sessionId, + getSessionFile: () => undefined, + getSessionName: () => undefined, + }, + getMessages: () => [message], + getContextUsage: () => harness.session.getContextUsage(), + getModel: () => harness.session.model, + findModel: (provider, id) => harness.session.modelRegistry.find(provider, id), + subtractUnindexedChildUsage: (ownUsage, entries) => { + for (const entry of entries) { + if (entry.type === "message" && entry.message === message && unindexed) { + subtractAssistantUsage(ownUsage, unindexed); + } + } + }, + getLiveChildren: () => [], + getRlmSessionDir: () => undefined, + }); + const first = owner.getOwnUsageSummary(); + expect(first?.inputTokens).toBe(100); + expect(owner.getOwnUsageSummary()).toBe(first); + unindexed = usage(40); + owner.invalidateOwnUsage(); + expect(owner.getOwnUsageSummary()?.inputTokens).toBe(60); + expect(owner.getContextTree().totalUsage.input).toBe(100); + expect(owner.getContextTree().ownUsage.input).toBe(60); + entries = [ + { + ...entries[0]!, + id: "replacement", + message: { ...fauxAssistantMessage("replacement"), usage: usage(25) }, + } as SessionEntry, + ]; + expect(owner.getOwnUsageSummary()?.inputTokens).toBe(25); + }); + + it("keeps cold digests lazy and uses timestamp recency after a context replacement", async () => { + const harness = await createHarness({ tools: [] }); + harnesses.push(harness); + let messages: AgentMessage[] = []; + const append = vi.spyOn(harness.sessionManager, "appendCustomMessageEntryWithRollback").mockImplementation(() => { + throw new Error("disk unavailable"); + }); + const owner = new SessionHarnessContext({ + sessionManager: harness.sessionManager, + getMessages: () => messages, + getActiveToolNames: () => [], + getVisibleSkills: () => [], + loadHarnessState: () => loadHarnessState(join(harness.tempDir, "harness"), "local"), + applyLateSentMessages: () => {}, + }); + owner.ensureHarnessDigestContext(); + expect(owner.digestPending).toBe(true); + expect(messages).toEqual([]); + expect(append).not.toHaveBeenCalled(); + messages = [{ role: "user", content: "resume", timestamp: 1 }]; + owner.ensureHarnessDigestContext(); + expect(owner.digestPending).toBe(false); + expect(messages).toHaveLength(2); + expect(append).toHaveBeenCalledOnce(); + owner.ensureHarnessDigestContext(); + expect(messages).toHaveLength(2); + const newer = { ...createHarnessDigestMessage("new"), timestamp: 30 }; + const older = { ...createHarnessDigestMessage("old"), timestamp: 20 }; + messages = [newer, older]; + expect(owner.latestContextHarnessDigest()).toBe("new"); + const outcome = { ...createHarnessDigestMessage("retained"), timestamp: 25 }; + owner.retainOutcome(outcome); + const rebuilt: AgentMessage[] = [older, newer]; + owner.mergeUnpersistedOutcomes(rebuilt); + expect(rebuilt).toEqual([older, outcome, newer]); + }); + + it("exports only the selected branch without mutating persisted parent links", async () => { + const harness = await createHarness({ tools: [], persistSession: true }); + harnesses.push(harness); + const root = harness.sessionManager.appendMessage({ role: "user", content: "root", timestamp: 1 }); + harness.sessionManager.appendMessage(fauxAssistantMessage("abandoned answer")); + harness.sessionManager.branch(root); + const answer = harness.sessionManager.appendMessage(fauxAssistantMessage("selected answer")); + harness.session.state.messages = harness.session.buildSessionContext().messages; + const before = structuredClone(harness.sessionManager.getEntries()); + const path = harness.session.exportToJsonl(join(harness.tempDir, "export", "branch.jsonl")); + const exported = readFileSync(path, "utf8") + .trim() + .split("\n") + .map((line) => JSON.parse(line) as FileEntry); + expect(exported[0]).toMatchObject({ type: "session", id: harness.session.sessionId }); + expect(exported.slice(1)).toMatchObject([ + { id: root, parentId: null }, + { id: answer, parentId: root }, + ]); + expect(exported).toHaveLength(3); + expect(harness.sessionManager.getEntries()).toEqual(before); + const html = await harness.session.exportToHtml(join(harness.tempDir, "session.html")); + const encoded = readFileSync(html, "utf8").match( + /