diff --git a/packages/coding-agent/src/README.md b/packages/coding-agent/src/README.md new file mode 100644 index 0000000000..c8e290f5b5 --- /dev/null +++ b/packages/coding-agent/src/README.md @@ -0,0 +1,38 @@ +# 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. + +`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. + +## 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. | + +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. + +`AgentSession` retains responsibilities that cross features: authentication and tool readiness, queue admission, cancellation, compaction, and waiting for child agents. It validates requests and tells the goal controller when to transition or account for a message. + +Three ordering rules matter during future extractions: + +- Account for assistant usage before executing its tools, so a completing turn is included. Repeated delivery of the same assistant message must not count twice. +- 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. + +## 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. + +The next design review should cover input admission and turn lifecycle. The existing `SessionActionStore` already owns action transitions and tickets; build around that ownership when extracting queue dispatch, pause/cancel, and continuation decisions. Do not migrate all callers in the same change as the goal extraction. + +## 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`. + +Run the focused files from the coding-agent package root with the repository's prescribed Vitest command, then run `npm run check` from the repository root. Use faux providers for session tests. diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 431c108378..674fe40064 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -34,6 +34,9 @@ import { resetApiProviders, 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 { stripFrontmatter } from "../utils/frontmatter.js"; import { sleep } from "../utils/sleep.js"; @@ -148,18 +151,12 @@ import { import { emitSessionShutdownEvent } from "./extensions/runner.js"; import { createGoalContextMessage, - emptyGoalState, GOAL_CONTEXT_CUSTOM_TYPE, GOAL_CONTEXT_PREVIEW_LABEL, GOAL_SKILL_NAME, - GOAL_STATE_CUSTOM_TYPE, type GoalHostResponse, type GoalState, - type GoalStatus, goalHostResponse, - goalTokenDeltaForUsage, - isPersistedGoalState, - normalizeGoalState, validateGoalBudget, validateGoalObjective, } from "./goals.js"; @@ -915,13 +912,6 @@ interface ToolDefinitionEntry { sourceInfo: SourceInfo; } -type GoalSlashCommand = - | { kind: "status" } - | { kind: "clear" } - | { kind: "pause" } - | { kind: "resume" } - | { kind: "start"; objective: string; tokenBudget?: number }; - type AutonomousSlashCommand = { kind: "status" } | { kind: "on"; config?: AgentAutonomousConfig } | { kind: "off" }; import type { RlmMaxDepthSource, RlmMaxDepthStatus, SetRlmMaxDepthResult } from "./rlm-max-depth.js"; @@ -1026,17 +1016,6 @@ function isPersistedRlmMaxDepthState(value: unknown): value is PersistedRlmMaxDe ); } -function parseGoalBudgetValue(value: string): number { - if (!/^[1-9]\d*$/.test(value)) { - throw new Error("Goal token budget must be a positive integer."); - } - const budget = validateGoalBudget(Number(value)); - if (budget === undefined) { - throw new Error("Goal token budget must be a positive integer."); - } - return budget; -} - const AUTONOMOUS_STATUS_NUMBER_FORMAT = new Intl.NumberFormat("en-US"); const AUTONOMOUS_BUDGET_USAGE = @@ -1267,10 +1246,8 @@ export class AgentSession { private readonly _sessionInputCheckpointWaiters = new Set<() => void>(); private _pendingNextTurnMessages: CustomMessage[] = []; - private _goalState: GoalState = emptyGoalState(); - private _goalAccountingStartedAt: number | undefined = undefined; + private readonly _goals: GoalController; private _goalContinuationAwaitsRlmWork = false; - private _goalAccountedAssistantMessages = new WeakSet(); private _goalAbortInProgress = false; private _autonomousState: AutonomousRuntimeState; private _autonomousContinuationSuppressionDepth = 0; @@ -1488,22 +1465,21 @@ export class AgentSession { this._autonomousState = createAutonomousRuntimeState(config.autonomous, { cwd: this._cwd, }); - this._goalState = this._loadPersistedGoalState(); + const goalPersistence = createGoalPersistence(this.sessionManager); + this._goals = new GoalController(goalPersistence, (goal) => this._emit({ type: "goal_update", goal })); // 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 // thread_goal_state. This prevents reseeding after clear/complete/error // or restart/rehydration of a session that already has messages or a goal. - if (this._rlmDepth === 0 && config.initialGoal && this._isBranchSeedable()) { - this._goalState = this._startGoal(config.initialGoal.objective, config.initialGoal.tokenBudget); + if (this._rlmDepth === 0 && config.initialGoal && goalPersistence.canSeed()) { + this._startGoal(config.initialGoal.objective, config.initialGoal.tokenBudget); // Goal context is the model's only source of goal visibility; action // admission is unavailable mid-construction, so ride the next turn. - this._pendingNextTurnMessages.push(createGoalContextMessage(this._goalState, "continuation")); + this._pendingNextTurnMessages.push(createGoalContextMessage(this._goals.state, "continuation")); } this._restoreLateIpythonSentAgentMessages(); - if (this._goalState.status === "active") { - this._goalAccountingStartedAt = Date.now(); - } + this._goals.restartAccounting(); this._unsubscribeAgent = this.agent.subscribe(this._handleAgentEvent); this._installAgentToolHooks(); @@ -1833,55 +1809,6 @@ export class AgentSession { return { maxDepth: 2, source: "default" }; } - private _loadPersistedGoalState(): GoalState { - const branch = this.sessionManager.getBranch(); - for (let i = branch.length - 1; i >= 0; i--) { - const entry = branch[i]; - if ( - entry.type === "custom" && - entry.customType === GOAL_STATE_CUSTOM_TYPE && - isPersistedGoalState(entry.data) - ) { - return normalizeGoalState(entry.data); - } - } - return emptyGoalState(); - } - - /** - * Whether the session branch is seedable for an initial goal. Returns true - * only when the branch contains exclusively bootstrap entry types - * (model_change, thinking_level_change, service_tier_change) and no - * thread_goal_state custom entry. Any message, custom entry, or persisted - * goal (including cleared/complete/error) means the session has been used - * and should not be reseeded. - */ - private _isBranchSeedable(): boolean { - const branch = this.sessionManager.getBranch(); - for (const entry of branch) { - switch (entry.type) { - case "model_change": - case "thinking_level_change": - case "service_tier_change": - continue; - case "custom": - if (entry.customType === GOAL_STATE_CUSTOM_TYPE) { - return false; - } - return false; - default: - return false; - } - } - return true; - } - - private _reloadGoalStateFromBranch(): void { - this._goalState = this._loadPersistedGoalState(); - this._goalAccountingStartedAt = this._goalState.status === "active" ? Date.now() : undefined; - this._emitGoalUpdate(); - } - private _reloadRlmMaxDepthFromBranch(): void { const previousMaxDepth = this._rlmMaxDepth; const resolved = this._resolveRlmMaxDepth(); @@ -1893,54 +1820,6 @@ export class AgentSession { } } - private _persistGoalState(goal: GoalState): void { - this.sessionManager.appendCustomEntry(GOAL_STATE_CUSTOM_TYPE, goal); - // Force flush so the goal state is durable on disk immediately, - // even before the first assistant response. This ensures idempotent - // restart/rehydration can detect the persisted goal. - this.sessionManager.flushNow(); - } - - private _setGoalState(next: GoalState, options: { persist?: boolean } = {}): void { - const normalized = normalizeGoalState({ - ...next, - updatedAt: Date.now(), - }); - this._goalState = normalized; - if (normalized.status === "active") { - this._goalAccountingStartedAt ??= Date.now(); - } else { - this._goalAccountingStartedAt = undefined; - } - if (options.persist !== false) { - this._persistGoalState(normalized); - } - this._emitGoalUpdate(); - } - - private _goalWithCurrentWallClock(now = Date.now()): GoalState { - if (this._goalState.status !== "active" || !this._goalAccountingStartedAt) { - return this._goalState; - } - const elapsedSeconds = Math.floor((now - this._goalAccountingStartedAt) / 1000); - if (elapsedSeconds <= 0) { - return this._goalState; - } - return { - ...this._goalState, - timeUsedSeconds: this._goalState.timeUsedSeconds + elapsedSeconds, - }; - } - - private _goalWithAccountedWallClock(): GoalState { - const now = Date.now(); - const goal = this._goalWithCurrentWallClock(now); - if (goal !== this._goalState) { - this._goalAccountingStartedAt = now; - } - return goal; - } - private _cancelSessionActions( predicate: (action: QueuedSessionAction) => boolean, error: Error, @@ -2034,86 +1913,28 @@ export class AgentSession { private _startGoal(objectiveText: string, tokenBudget: number | undefined): GoalState { const objective = validateGoalObjective(objectiveText); const budget = validateGoalBudget(tokenBudget); - const now = Date.now(); - const goal: GoalState = { - active: true, - status: "active", - goalId: randomUUID(), - objective, - tokenBudget: budget, - tokensUsed: 0, - timeUsedSeconds: 0, - continuationsUsed: 0, - createdAt: now, - updatedAt: now, - }; - this._goalAccountingStartedAt = now; this._goalContinuationAwaitsRlmWork = false; - this._setGoalState(goal); - return this._goalState; + return this._goals.start(objective, budget); } private _clearGoal(): void { this._clearQueuedGoalContexts(); - this._setGoalState(emptyGoalState()); + this._goals.clear(); } - private _pauseGoal(reason = "Paused by user"): void { + private _pauseGoal(): void { this._clearQueuedGoalContexts(); - if (this._goalState.status !== "active") { - this._emitGoalUpdate(); - return; - } - const goal = this._goalWithAccountedWallClock(); - this._setGoalState({ - ...goal, - active: false, - status: "paused", - lastReason: reason, - lastError: undefined, - }); + this._goals.pause(); } private async _resumeGoal(): Promise { - if (!this._goalState.objective) { - this._emitGoalUpdate(); - return; - } - if (this._goalState.status !== "paused" && this._goalState.status !== "budget_limited") { - this._emitGoalUpdate(); - return; - } - const exhausted = - this._goalState.tokenBudget !== undefined && this._goalState.tokensUsed >= this._goalState.tokenBudget; - const nextStatus: GoalStatus = exhausted ? "budget_limited" : "active"; - this._setGoalState({ - ...this._goalState, - active: nextStatus === "active", - status: nextStatus, - lastReason: exhausted ? "Goal token budget already reached" : undefined, - lastError: undefined, - }); - if (nextStatus === "active") { + if (this._goals.resume()) { await this._runOrQueueGoalContext("continuation"); } } - private _finishGoalWithError(errorMessage: string): void { - if (!this._goalState.objective || this._goalState.status !== "active") { - return; - } - const goal = this._goalWithAccountedWallClock(); - this._setGoalState({ - ...goal, - active: false, - status: "error", - lastReason: errorMessage, - lastError: errorMessage, - }); - } - private _finishGoalForTerminalAssistantMessage(message: AssistantMessage): void { - if (this._goalState.status !== "active") { + if (this._goals.state.status !== "active") { return; } @@ -2127,7 +1948,7 @@ export class AgentSession { this._goalAbortInProgress = false; return; } - this._finishGoalWithError(message.errorMessage || "Assistant response failed"); + this._goals.fail(message.errorMessage || "Assistant response failed"); } } @@ -2143,58 +1964,6 @@ export class AgentSession { return true; } - private _parseGoalSlashCommand(text: string): GoalSlashCommand | undefined { - const command = parseSessionSlashCommand(text); - if (command?.name !== "goal") return undefined; - - const rest = command.args; - const normalized = rest.toLowerCase(); - if (!rest || normalized === "status") { - return { kind: "status" }; - } - if (normalized === "clear" || normalized === "stop") { - return { kind: "clear" }; - } - if (normalized === "pause") { - return { kind: "pause" }; - } - if (normalized === "resume") { - return { kind: "resume" }; - } - - let tokenBudget: number | undefined; - let objective = rest; - const firstToken = rest.split(/\s+/, 1)[0] ?? ""; - if ( - firstToken === "--budget" || - firstToken === "--token-budget" || - firstToken.startsWith("--budget=") || - firstToken.startsWith("--token-budget=") - ) { - let valueText: string; - if (firstToken === "--budget" || firstToken === "--token-budget") { - const withoutFlag = rest.slice(firstToken.length).trimStart(); - const nextSpace = withoutFlag.search(/\s/); - if (nextSpace < 0) { - throw new Error("Usage: /goal [--budget ] "); - } - valueText = withoutFlag.slice(0, nextSpace); - objective = withoutFlag.slice(nextSpace + 1).trim(); - } else { - const separator = firstToken.indexOf("="); - valueText = firstToken.slice(separator + 1); - objective = rest.slice(firstToken.length).trim(); - } - tokenBudget = parseGoalBudgetValue(valueText); - } - - return { - kind: "start", - objective: validateGoalObjective(objective), - tokenBudget, - }; - } - private _parseAutonomousSlashCommand(text: string): AutonomousSlashCommand | undefined { const command = parseSessionSlashCommand(text); if (command?.name !== "autonomous") return undefined; @@ -2328,23 +2097,18 @@ export class AgentSession { private _maybeResumeGoalContinuationAfterRlmWork(): void { if (!this._goalContinuationAwaitsRlmWork) return; if (this._disposed || this._disposing || this._hasUnsettledRlmQuiescenceWork()) return; - if (this._goalState.status !== "active" || !this._goalState.objective) { + 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._sessionInputAdmissionPauses.size > 0 || this._sessionInputPumpSuspended) return; - const goalBeforeResume = this._goalState; + const goalBeforeResume = this._goals.checkpoint(); try { this._ensureGoalRuntimeActive(); - this._setGoalState({ - ...this._goalState, - continuationsUsed: this._goalState.continuationsUsed + 1, - lastReason: undefined, - lastError: undefined, - }); - const message = createGoalContextMessage(this._goalState, "continuation"); + 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( @@ -2356,14 +2120,14 @@ export class AgentSession { this._goalContinuationAwaitsRlmWork = false; } catch { // Admission can race a new pause; roll back so the retry re-counts. - this._setGoalState(goalBeforeResume); + this._goals.restore(goalBeforeResume, { restoreClock: false }); } } private _runOrQueueGoalContext(kind: "continuation" | "objective_updated", images?: ImageContent[]): void { - if (!this._goalState.objective) return; + if (!this._goals.state.objective) return; this._ensureGoalRuntimeActive(); - const message = createGoalContextMessage(this._goalState, kind, images); + const message = createGoalContextMessage(this._goals.state, kind, images); const normalized = normalizeMessageContent(message.content); const action = this._createPreparedTurnAction("followUp", normalized.text, normalized.images, { message, @@ -2373,7 +2137,7 @@ export class AgentSession { } private async _handleGoalSlashCommand(text: string, images: ImageContent[] | undefined): Promise { - const command = this._parseGoalSlashCommand(text); + const command = parseGoalSlashCommand(text); if (!command) { return false; } @@ -2398,7 +2162,7 @@ export class AgentSession { return true; } - const previousWasActive = this._goalState.status === "active"; + const previousWasActive = this._goals.state.status === "active"; if (!this.isStreaming) { await this._validateCanStartAgentRun(); } @@ -2409,46 +2173,6 @@ export class AgentSession { return true; } - private _accountGoalUsageForAssistantMessage(message: AssistantMessage): boolean { - if (!this._goalState.objective) { - return false; - } - if (message.stopReason === "error" || message.stopReason === "aborted") { - return false; - } - if (this._goalAccountedAssistantMessages.has(message)) { - return false; - } - // Usage is attributed at the assistant message's message_end, which fires - // before that turn's ipython cell runs. goal.complete() only arrives later - // over the kernel host bridge, so the completing turn is always accounted - // while the goal is still active. Only count turns spent pursuing the goal; - // post-completion turns (e.g. a closing summary) must not be attributed. - if (this._goalState.status !== "active") { - return false; - } - this._goalAccountedAssistantMessages.add(message); - const tokenDelta = goalTokenDeltaForUsage(message.usage); - const goal = this._goalWithAccountedWallClock(); - const nextGoal: GoalState = { - ...goal, - tokensUsed: goal.tokensUsed + tokenDelta, - }; - const budgetReached = nextGoal.tokenBudget !== undefined && nextGoal.tokensUsed >= nextGoal.tokenBudget; - if (!budgetReached) { - this._setGoalState(nextGoal); - return false; - } - this._setGoalState({ - ...nextGoal, - active: false, - status: "budget_limited", - lastReason: `Reached ${nextGoal.tokenBudget} token goal budget`, - lastError: undefined, - }); - return true; - } - private get _steeringStopPending(): boolean { return ( this._actionStore.queuedActions("next_turn_boundary").length > 0 || @@ -2471,8 +2195,8 @@ export class AgentSession { return true; } try { - if (this._accountGoalUsageForAssistantMessage(context.message)) { - const message = createGoalContextMessage(this._goalState, "budget_limit"); + 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, @@ -3081,7 +2805,7 @@ export class AgentSession { if (message.stopReason === "error" || message.stopReason === "aborted") { return false; } - if (this._goalState.status !== "active" || !this._goalState.objective) { + if (this._goals.state.status !== "active" || !this._goals.state.objective) { return false; } const alreadyQueued = this._queuedGoalThresholdContinuation; @@ -3102,13 +2826,8 @@ export class AgentSession { } try { this._ensureGoalRuntimeActive(); - this._setGoalState({ - ...this._goalState, - continuationsUsed: this._goalState.continuationsUsed + 1, - lastReason: undefined, - lastError: undefined, - }); - const goalMessage = createGoalContextMessage(this._goalState, "continuation"); + this._goals.recordContinuation(); + const goalMessage = createGoalContextMessage(this._goals.state, "continuation"); const normalized = normalizeMessageContent(goalMessage.content); this._admitSessionInput( this._createPreparedTurnAction("followUp", normalized.text, normalized.images, { @@ -3136,7 +2855,7 @@ export class AgentSession { // 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._setGoalState({ ...this._goalState, continuationsUsed: this._goalState.continuationsUsed - 1 }); + this._goals.cancelContinuation(); this._emitQueueUpdate(); } @@ -3499,7 +3218,7 @@ export class AgentSession { } private _createGoalFromHost(objective: string, tokenBudget: number | undefined): GoalState { - switch (this._goalState.status) { + 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", @@ -3519,22 +3238,9 @@ export class AgentSession { } private _completeGoalFromHost(): GoalState { - if (!this._goalState.objective || this._goalState.status === "idle") { - throw new Error("cannot complete goal because this thread has no goal"); - } - const goal = this._goalWithAccountedWallClock(); - // A turn can cross the budget and complete the goal at once: accounting - // runs at message_end, before the completing ipython cell executes, so a - // budget-limit context may already be steered. It is stale now — drop it. - this._clearQueuedGoalContexts(); - this._setGoalState({ - ...goal, - active: false, - status: "complete", - lastReason: "Goal achieved", - lastError: undefined, - }); - return this._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( @@ -3544,7 +3250,7 @@ export class AgentSession { if (this._stopGoalContinuationForTerminalMessage(context.message)) { return []; } - if (signal?.aborted || this._goalState.status !== "active" || !this._goalState.objective) { + if (signal?.aborted || this._goals.state.status !== "active" || !this._goals.state.objective) { return []; } // Delegating and ending the turn is correct behavior; hold the continuation @@ -3556,18 +3262,12 @@ export class AgentSession { this._goalContinuationAwaitsRlmWork = false; try { this._ensureGoalRuntimeActive(context.context); - const nextGoal = { - ...this._goalState, - continuationsUsed: this._goalState.continuationsUsed + 1, - lastReason: undefined, - lastError: undefined, - }; - this._setGoalState(nextGoal); - return [createGoalContextMessage(this._goalState, "continuation")]; + this._goals.recordContinuation(); + return [createGoalContextMessage(this._goals.state, "continuation")]; } catch (error) { const message = error instanceof Error ? error.message : String(error); try { - this._finishGoalWithError(message); + this._goals.fail(message); } catch { // The continuation hook must not reject; listener failures should not crash the agent loop. } @@ -3583,13 +3283,11 @@ export class AgentSession { return []; } const arrivalEpoch = this._sessionInputArrivalEpoch; - const goalSnapshot = this._goalState; - const goalAccountingStartedAt = this._goalAccountingStartedAt; + 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._setGoalState(goalSnapshot); - this._goalAccountingStartedAt = goalAccountingStartedAt; + this._goals.restore(goalSnapshot); return []; } return goalMessages; @@ -3906,8 +3604,8 @@ export class AgentSession { this._retryAttempt = 0; this._retryAuthFailureSources = []; } - if (this._accountGoalUsageForAssistantMessage(assistantMsg)) { - const message = createGoalContextMessage(this._goalState, "budget_limit"); + 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, @@ -4542,7 +4240,7 @@ export class AgentSession { } get goalState(): GoalState { - return { ...this._goalWithCurrentWallClock() }; + return this._goals.current; } getAutonomousStatus(): AgentAutonomousStatus { @@ -6397,8 +6095,8 @@ export class AgentSession { } case "goal": await this._handleGoalSlashCommand(input.text, input.images); - resultText = this._goalState.objective - ? `Goal ${this._goalState.status}: ${this._goalState.objective}` + resultText = this._goals.state.objective + ? `Goal ${this._goals.state.status}: ${this._goals.state.objective}` : "No active goal."; break; case "autonomous": @@ -7284,7 +6982,7 @@ export class AgentSession { const branchSummaryOperation = this._branchSummaryOperation; this.requestAbort(); this._cancelActiveRlmChildRuns("Parent session aborted"); - this._goalAbortInProgress = this._goalState.status === "active"; + this._goalAbortInProgress = this._goals.state.status === "active"; try { await Promise.allSettled([ this.agent.waitForIdle(), @@ -7308,7 +7006,7 @@ export class AgentSession { this.abortRetry(); for (const controller of this._rlmQuiescenceWaitAborts) controller.abort(); this._cancelActiveRlmChildRuns("Parent session aborted for update restart"); - this._goalAbortInProgress = this._goalState.status === "active"; + this._goalAbortInProgress = this._goals.state.status === "active"; this.agent.abort(); if (this._goalAbortInProgress) { void this.agent @@ -7766,7 +7464,7 @@ export class AgentSession { this._scheduleSessionInputPump(); if (didCompact) { this._discardPendingAutoRefine({ cancelPostCompactionContinue: true }); - if (this._goalState.status === "active" && !compactionAbort.signal.aborted) { + if (this._goals.state.status === "active" && !compactionAbort.signal.aborted) { this._goalContinuationAwaitsRlmWork ||= !this.agent.hasQueuedMessages(); this.resumeQueuedWork(); if (this.agent.hasQueuedMessages()) this._schedulePostCompactionContinue(); @@ -9626,7 +9324,7 @@ export class AgentSession { const defaultActiveToolNames = this._baseToolsOverride ? Object.keys(this._baseToolsOverride) : ["ipython"]; const baseActiveToolNames = [...(options.activeToolNames ?? defaultActiveToolNames)]; - if (this._goalState.status === "active" && this._includeGoals) { + if (this._goals.state.status === "active" && this._includeGoals) { // An active goal needs ipython so the model can reach the goal skill. baseActiveToolNames.push("ipython"); } @@ -12286,7 +11984,7 @@ export class AgentSession { this._restoreLateIpythonSentAgentMessages(); // Context rebuild = cold boundary: refresh the digest like resume. this._ensureHarnessDigestContext(); - this._reloadGoalStateFromBranch(); + this._goals.reload(); this._reloadRlmMaxDepthFromBranch(); this._invalidateQueuedPromptPreparation(); diff --git a/packages/coding-agent/src/goals/commands.ts b/packages/coding-agent/src/goals/commands.ts new file mode 100644 index 0000000000..cca2fa43b8 --- /dev/null +++ b/packages/coding-agent/src/goals/commands.ts @@ -0,0 +1,72 @@ +import { validateGoalBudget, validateGoalObjective } from "../core/goals.js"; +import { parseSessionSlashCommand } from "../core/slash-commands.js"; + +type GoalSlashCommand = + | { kind: "status" } + | { kind: "clear" } + | { kind: "pause" } + | { kind: "resume" } + | { kind: "start"; objective: string; tokenBudget?: number }; + +export function parseGoalSlashCommand(text: string): GoalSlashCommand | undefined { + const command = parseSessionSlashCommand(text); + if (command?.name !== "goal") return undefined; + + const rest = command.args; + const normalized = rest.toLowerCase(); + if (!rest || normalized === "status") { + return { kind: "status" }; + } + if (normalized === "clear" || normalized === "stop") { + return { kind: "clear" }; + } + if (normalized === "pause") { + return { kind: "pause" }; + } + if (normalized === "resume") { + return { kind: "resume" }; + } + + let tokenBudget: number | undefined; + let objective = rest; + const firstToken = rest.split(/\s+/, 1)[0] ?? ""; + if ( + firstToken === "--budget" || + firstToken === "--token-budget" || + firstToken.startsWith("--budget=") || + firstToken.startsWith("--token-budget=") + ) { + let valueText: string; + if (firstToken === "--budget" || firstToken === "--token-budget") { + const withoutFlag = rest.slice(firstToken.length).trimStart(); + const nextSpace = withoutFlag.search(/\s/); + if (nextSpace < 0) { + throw new Error("Usage: /goal [--budget ] "); + } + valueText = withoutFlag.slice(0, nextSpace); + objective = withoutFlag.slice(nextSpace + 1).trim(); + } else { + const separator = firstToken.indexOf("="); + valueText = firstToken.slice(separator + 1); + objective = rest.slice(firstToken.length).trim(); + } + tokenBudget = parseGoalBudgetValue(valueText); + } + + return { + kind: "start", + objective: validateGoalObjective(objective), + tokenBudget, + }; +} + +function parseGoalBudgetValue(value: string): number { + if (!/^[1-9]\d*$/.test(value)) { + throw new Error("Goal token budget must be a positive integer."); + } + const budget = validateGoalBudget(Number(value)); + if (budget === undefined) { + throw new Error("Goal token budget must be a positive integer."); + } + return budget; +} diff --git a/packages/coding-agent/src/goals/controller.ts b/packages/coding-agent/src/goals/controller.ts new file mode 100644 index 0000000000..fb33987535 --- /dev/null +++ b/packages/coding-agent/src/goals/controller.ts @@ -0,0 +1,194 @@ +import { randomUUID } from "node:crypto"; +import type { AssistantMessage } from "@earendil-works/pi-ai"; +import { emptyGoalState, type GoalState, goalTokenDeltaForUsage, normalizeGoalState } from "../core/goals.js"; +import type { GoalPersistence } from "./persistence.js"; + +export interface GoalCheckpoint { + readonly state: Readonly; + readonly accountingStartedAt: number | undefined; +} + +/** Owns goal state and accounting; the session owns scheduling and runtime readiness. */ +export class GoalController { + private _state: GoalState; + private _accountingStartedAt: number | undefined; + private readonly _accountedAssistantMessages = new WeakSet(); + + constructor( + private readonly _persistence: GoalPersistence, + private readonly _onUpdate: (goal: GoalState) => void, + private readonly _now: () => number = () => Date.now(), + ) { + this._state = _persistence.load(); + } + + get state(): Readonly { + return this._state; + } + + get current(): GoalState { + return { ...this._withCurrentWallClock() }; + } + + restartAccounting(): void { + this._accountingStartedAt = this._state.status === "active" ? this._now() : undefined; + } + + reload(): void { + this._state = this._persistence.load(); + this.restartAccounting(); + this._onUpdate(this.current); + } + + // The session request boundary validates the objective and budget before starting. + start(objective: string, tokenBudget: number | undefined): GoalState { + const now = this._now(); + this._accountingStartedAt = now; + this._setState({ + active: true, + status: "active", + goalId: randomUUID(), + objective, + tokenBudget, + tokensUsed: 0, + timeUsedSeconds: 0, + continuationsUsed: 0, + createdAt: now, + updatedAt: now, + }); + return this._state; + } + + clear(): void { + this._setState(emptyGoalState()); + } + + pause(reason = "Paused by user"): void { + if (this._state.status !== "active") { + this._onUpdate(this.current); + return; + } + this._setState({ + ...this._withAccountedWallClock(), + active: false, + status: "paused", + lastReason: reason, + lastError: undefined, + }); + } + + /** Returns whether the session should schedule a continuation. */ + resume(): boolean { + if (!this._state.objective || (this._state.status !== "paused" && this._state.status !== "budget_limited")) { + this._onUpdate(this.current); + return false; + } + const exhausted = this._state.tokenBudget !== undefined && this._state.tokensUsed >= this._state.tokenBudget; + this._setState({ + ...this._state, + active: !exhausted, + status: exhausted ? "budget_limited" : "active", + lastReason: exhausted ? "Goal token budget already reached" : undefined, + lastError: undefined, + }); + return !exhausted; + } + + fail(errorMessage: string): void { + if (!this._state.objective || this._state.status !== "active") return; + this._setState({ + ...this._withAccountedWallClock(), + active: false, + status: "error", + lastReason: errorMessage, + lastError: errorMessage, + }); + } + + complete(clearQueuedContexts: () => void): GoalState { + if (!this._state.objective || this._state.status === "idle") { + throw new Error("cannot complete goal because this thread has no goal"); + } + const goal = this._withAccountedWallClock(); + // Capture usage before clearing stale budget messages, then publish completion. + clearQueuedContexts(); + this._setState({ + ...goal, + active: false, + status: "complete", + lastReason: "Goal achieved", + lastError: undefined, + }); + return this._state; + } + + recordContinuation(): void { + this._setState({ + ...this._state, + continuationsUsed: this._state.continuationsUsed + 1, + lastReason: undefined, + lastError: undefined, + }); + } + + cancelContinuation(): void { + this._setState({ ...this._state, continuationsUsed: this._state.continuationsUsed - 1 }); + } + + checkpoint(): GoalCheckpoint { + return { state: this._state, accountingStartedAt: this._accountingStartedAt }; + } + + restore(checkpoint: GoalCheckpoint, options: { restoreClock?: boolean } = {}): void { + this._setState(checkpoint.state); + if (options.restoreClock !== false) this._accountingStartedAt = checkpoint.accountingStartedAt; + } + + /** Returns true only when this message newly reaches the goal budget. */ + accountAssistantMessage(message: AssistantMessage): boolean { + if (!this._state.objective || message.stopReason === "error" || message.stopReason === "aborted") return false; + if (this._accountedAssistantMessages.has(message) || this._state.status !== "active") return false; + this._accountedAssistantMessages.add(message); + const goal = this._withAccountedWallClock(); + const nextGoal = { ...goal, tokensUsed: goal.tokensUsed + goalTokenDeltaForUsage(message.usage) }; + const budgetReached = nextGoal.tokenBudget !== undefined && nextGoal.tokensUsed >= nextGoal.tokenBudget; + if (!budgetReached) { + this._setState(nextGoal); + return false; + } + this._setState({ + ...nextGoal, + active: false, + status: "budget_limited", + lastReason: `Reached ${nextGoal.tokenBudget} token goal budget`, + lastError: undefined, + }); + return true; + } + + private _setState(next: Readonly): void { + const normalized = normalizeGoalState({ ...next, updatedAt: this._now() }); + this._state = normalized; + if (normalized.status === "active") { + this._accountingStartedAt ??= this._now(); + } else { + this._accountingStartedAt = undefined; + } + this._persistence.save(normalized); + this._onUpdate(this.current); + } + + private _withCurrentWallClock(now = this._now()): GoalState { + if (this._state.status !== "active" || !this._accountingStartedAt) return this._state; + const elapsedSeconds = Math.floor((now - this._accountingStartedAt) / 1000); + if (elapsedSeconds <= 0) return this._state; + return { ...this._state, timeUsedSeconds: this._state.timeUsedSeconds + elapsedSeconds }; + } + + private _withAccountedWallClock(): GoalState { + const now = this._now(); + const goal = this._withCurrentWallClock(now); + if (goal !== this._state) this._accountingStartedAt = now; + return goal; + } +} diff --git a/packages/coding-agent/src/goals/persistence.ts b/packages/coding-agent/src/goals/persistence.ts new file mode 100644 index 0000000000..25e077791c --- /dev/null +++ b/packages/coding-agent/src/goals/persistence.ts @@ -0,0 +1,50 @@ +import { + emptyGoalState, + GOAL_STATE_CUSTOM_TYPE, + type GoalState, + isPersistedGoalState, + normalizeGoalState, +} from "../core/goals.js"; +import type { SessionManager } from "../core/session-manager.js"; + +export interface GoalPersistence { + load(): GoalState; + save(goal: GoalState): void; +} + +export function createGoalPersistence( + session: Pick, +): GoalPersistence & { canSeed(): boolean } { + return { + load() { + const branch = session.getBranch(); + for (let i = branch.length - 1; i >= 0; i--) { + const entry = branch[i]; + if ( + entry.type === "custom" && + entry.customType === GOAL_STATE_CUSTOM_TYPE && + isPersistedGoalState(entry.data) + ) { + return normalizeGoalState(entry.data); + } + } + return emptyGoalState(); + }, + save(goal) { + session.appendCustomEntry(GOAL_STATE_CUSTOM_TYPE, goal); + // Restart must see the goal even before the first assistant response. + session.flushNow(); + }, + canSeed() { + // Any message or custom entry, including a cleared goal, means this branch was used. + return session + .getBranch() + .every( + (entry) => + entry.type === "model_change" || + entry.type === "thinking_level_change" || + entry.type === "service_tier_change", + ); + }, + }; +} diff --git a/packages/coding-agent/src/modes/daemon/daemon-mode.ts b/packages/coding-agent/src/modes/daemon/daemon-mode.ts index 7ba94b8366..2f4b15e08b 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-mode.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-mode.ts @@ -3404,9 +3404,7 @@ export class AgentDaemon { ...(entry.repliedSinceTask !== undefined ? { repliedSinceTask: entry.repliedSinceTask } : {}), ...(entry.parentSessionId ? { parentSessionId: entry.parentSessionId } : {}), ...(entry.rlmChildId ? { rlmChildId: entry.rlmChildId } : {}), - ...(entry.firstMessage - ? { firstMessage: entry.firstMessage.slice(0, AGENT_OBSERVE_PREVIEW_MAX_CHARS) } - : {}), + ...(entry.firstMessage ? { firstMessage: entry.firstMessage.slice(0, AGENT_OBSERVE_PREVIEW_MAX_CHARS) } : {}), }; } @@ -3482,7 +3480,11 @@ export class AgentDaemon { ...(summary.firstMessage ? { firstMessage: summary.firstMessage } : {}), ...(latest ? { - latestMessage: createAgentObserveMessagePreview(latest, messages.length - 1, AGENT_OBSERVE_PREVIEW_MAX_CHARS), + latestMessage: createAgentObserveMessagePreview( + latest, + messages.length - 1, + AGENT_OBSERVE_PREVIEW_MAX_CHARS, + ), } : {}), }; diff --git a/packages/coding-agent/test/goal-continuation-quiescence.test.ts b/packages/coding-agent/test/goal-continuation-quiescence.test.ts index 53e639f5db..1ab9434a8c 100644 --- a/packages/coding-agent/test/goal-continuation-quiescence.test.ts +++ b/packages/coding-agent/test/goal-continuation-quiescence.test.ts @@ -1,8 +1,10 @@ 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"; type Harness = { - _goalState: { status: string; objective?: string; continuationsUsed: number }; + _goals: GoalController; _goalContinuationAwaitsRlmWork: boolean; _disposed: boolean; _disposing: boolean; @@ -11,7 +13,6 @@ type Harness = { _hasUnsettledRlmQuiescenceWork: () => boolean; _stopGoalContinuationForTerminalMessage: () => boolean; _ensureGoalRuntimeActive: () => void; - _setGoalState: (goal: unknown) => void; _createPreparedTurnAction: ReturnType; _admitSessionInput: ReturnType; }; @@ -25,8 +26,10 @@ const maybeResume = Reflect.get(AgentSession.prototype, "_maybeResumeGoalContinu ) => void; function harness(overrides: Partial = {}): Harness { + const goals = new GoalController({ load: emptyGoalState, save: () => {} }, () => {}); + goals.start("ship it", undefined); return { - _goalState: { status: "active", objective: "ship it", continuationsUsed: 0 }, + _goals: goals, _goalContinuationAwaitsRlmWork: false, _disposed: false, _disposing: false, @@ -35,9 +38,6 @@ function harness(overrides: Partial = {}): Harness { _hasUnsettledRlmQuiescenceWork: () => false, _stopGoalContinuationForTerminalMessage: () => false, _ensureGoalRuntimeActive: () => {}, - _setGoalState: function (this: Harness, goal: unknown) { - this._goalState = goal as Harness["_goalState"]; - }, _createPreparedTurnAction: vi.fn((schedule: string, _text: string, _images: unknown, options: unknown) => ({ schedule, options, @@ -54,7 +54,7 @@ describe("goal continuation vs unsettled subagent work", () => { const mode = harness({ _hasUnsettledRlmQuiescenceWork: () => true }); await expect(getGoalContinuation.call(mode, context)).resolves.toEqual([]); expect(mode._goalContinuationAwaitsRlmWork).toBe(true); - expect(mode._goalState.continuationsUsed).toBe(0); + expect(mode._goals.state.continuationsUsed).toBe(0); }); it("continues normally when no descendant work is pending", async () => { @@ -62,7 +62,7 @@ describe("goal continuation vs unsettled subagent work", () => { const messages = await getGoalContinuation.call(mode, context); expect(messages).toHaveLength(1); expect(mode._goalContinuationAwaitsRlmWork).toBe(false); - expect(mode._goalState.continuationsUsed).toBe(1); + expect(mode._goals.state.continuationsUsed).toBe(1); }); it("resumes a deferred continuation exactly once, unqueued, idle-waking, and counted", () => { @@ -73,7 +73,7 @@ describe("goal continuation vs unsettled subagent work", () => { const [action, options] = mode._admitSessionInput.mock.calls[0]!; expect((action as { options: { resumeIfIdle: boolean } }).options.resumeIfIdle).toBe(true); expect(options).toBeUndefined(); - expect(mode._goalState.continuationsUsed).toBe(1); + expect(mode._goals.state.continuationsUsed).toBe(1); }); it("keeps the deferral while admission is paused and retries after release", () => { @@ -107,7 +107,7 @@ describe("goal continuation vs unsettled subagent work", () => { }); maybeResume.call(mode); expect(mode._goalContinuationAwaitsRlmWork).toBe(true); - expect(mode._goalState.continuationsUsed).toBe(0); + expect(mode._goals.state.continuationsUsed).toBe(0); }); it("stays deferred while work remains and drops the deferral for inactive goals", () => { @@ -117,7 +117,7 @@ describe("goal continuation vs unsettled subagent work", () => { expect(busy._goalContinuationAwaitsRlmWork).toBe(true); const inactive = harness({ _goalContinuationAwaitsRlmWork: true }); - inactive._goalState = { status: "paused", objective: "ship it", continuationsUsed: 0 }; + inactive._goals.pause(); maybeResume.call(inactive); expect(inactive._admitSessionInput).not.toHaveBeenCalled(); expect(inactive._goalContinuationAwaitsRlmWork).toBe(false); diff --git a/packages/coding-agent/test/goals/controller.test.ts b/packages/coding-agent/test/goals/controller.test.ts new file mode 100644 index 0000000000..99c0423617 --- /dev/null +++ b/packages/coding-agent/test/goals/controller.test.ts @@ -0,0 +1,123 @@ +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"; + +function createController(initial = emptyGoalState()) { + let persisted = initial; + const save = vi.fn((goal: GoalState) => { + persisted = goal; + }); + const onUpdate = vi.fn(); + let now = 1_000; + const goals = new GoalController({ load: () => persisted, save }, onUpdate, () => now); + goals.restartAccounting(); + return { + goals, + save, + onUpdate, + persisted: () => persisted, + setTime: (time: number) => { + now = time; + }, + }; +} + +function assistant(input: number, output: number, stopReason: AssistantMessage["stopReason"] = "stop") { + const message = fauxAssistantMessage("step", { stopReason }); + return { ...message, usage: { ...message.usage, input, output, cacheRead: 100, cacheWrite: 100 } }; +} + +describe("GoalController", () => { + it("counts active time across pause and resume without persisting status reads", () => { + const { goals, save, setTime } = createController(); + goals.start("finish the task", undefined); + setTime(6_500); + expect(goals.current.timeUsedSeconds).toBe(5); + expect(save).toHaveBeenCalledTimes(1); + goals.pause(); + setTime(60_000); + expect(goals.current.timeUsedSeconds).toBe(5); + expect(goals.resume()).toBe(true); + setTime(63_000); + expect(goals.current.timeUsedSeconds).toBe(8); + const visible = goals.current; + visible.tokensUsed = 999; + expect(goals.current.tokensUsed).toBe(0); + }); + + it("counts each assistant message once and excludes cached tokens", () => { + const { goals } = createController(); + goals.start("finish the task", 10); + const first = assistant(4, 2); + expect(goals.accountAssistantMessage(first)).toBe(false); + expect(goals.accountAssistantMessage(first)).toBe(false); + expect(goals.current.tokensUsed).toBe(6); + expect(goals.accountAssistantMessage(assistant(3, 2))).toBe(true); + expect(goals.current).toMatchObject({ status: "budget_limited", tokensUsed: 11 }); + expect(goals.resume()).toBe(false); + }); + + it("keeps completion usage and clears queued context before persisting completion", () => { + const { goals, persisted } = createController(); + goals.start("finish the task", 10); + goals.accountAssistantMessage(assistant(6, 5, "toolUse")); + const clearQueuedContexts = vi.fn(() => { + expect(persisted().status).toBe("budget_limited"); + }); + goals.complete(clearQueuedContexts); + expect(clearQueuedContexts).toHaveBeenCalledOnce(); + goals.accountAssistantMessage(assistant(20, 10)); + expect(persisted()).toMatchObject({ status: "complete", tokensUsed: 11 }); + expect(goals.resume()).toBe(false); + }); + + it("excludes aborted and errored messages, and accounts time on terminal failure", () => { + const { goals, setTime } = createController(); + goals.start("finish the task", undefined); + expect(goals.accountAssistantMessage(assistant(3, 2, "error"))).toBe(false); + expect(goals.accountAssistantMessage(assistant(3, 2, "aborted"))).toBe(false); + setTime(4_000); + goals.fail("provider failed"); + setTime(9_000); + expect(goals.current).toMatchObject({ status: "error", tokensUsed: 0, timeUsedSeconds: 3 }); + expect(goals.resume()).toBe(false); + }); + + it("restores the continuation count and accounting clock when admission loses a race", () => { + const { goals, persisted, setTime } = createController(); + goals.start("finish the task", undefined); + const checkpoint = goals.checkpoint(); + setTime(6_000); + goals.recordContinuation(); + goals.pause(); + goals.restore(checkpoint); + setTime(8_000); + expect(goals.current).toMatchObject({ status: "active", continuationsUsed: 0, timeUsedSeconds: 7 }); + expect(persisted().continuationsUsed).toBe(0); + }); + + it("reloads the selected branch without counting time spent away", () => { + const { goals, save, setTime } = createController(); + goals.start("first branch", undefined); + goals.accountAssistantMessage(assistant(2, 3)); + const firstBranch = goals.current; + goals.start("second branch", undefined); + setTime(50_000); + save(firstBranch); + goals.reload(); + setTime(52_000); + expect(goals.current).toMatchObject({ objective: "first branch", tokensUsed: 5, timeUsedSeconds: 2 }); + }); + + it("does not publish a successful update when persistence fails", () => { + const { goals, save, onUpdate } = createController(); + goals.start("finish the task", undefined); + onUpdate.mockClear(); + save.mockImplementationOnce(() => { + throw new Error("disk full"); + }); + expect(() => goals.pause()).toThrow("disk full"); + expect(onUpdate).not.toHaveBeenCalled(); + }); +});