diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 9c5eee224..609ac0679 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed aborted usage-limit recovery waiting on a local usage fetch or marking the resolved credential blocked after the caller had already moved to another session ([#6883](https://github.com/can1357/oh-my-pi/pull/6883) by [@paolomazzitti](https://github.com/paolomazzitti)). + ## [17.1.8] - 2026-07-28 ### Fixed diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index ba7e2ffa6..194e444e3 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -4205,7 +4205,10 @@ export class AuthStorage { let blockedUntil = now + (options?.retryAfterMs ?? AuthStorage.#defaultBackoffMs); if (credentialType === "oauth" && target.credential.type === "oauth" && strategy) { - const report = await this.#getUsageReport(provider, target.credential, options); + const report = await raceUsageWithSignal( + this.#getUsageReport(provider, target.credential, options), + options?.signal, + ); if (report) { const scopedLimits = this.#getScopedUsageLimits(strategy, report, rankingContext); if (this.#isUsageLimitReached(scopedLimits)) { @@ -4216,6 +4219,7 @@ export class AuthStorage { } } } + options?.signal?.throwIfAborted(); // Usage lookup may refresh, disable, or remove a row. Re-resolve its // durable id before applying positional in-memory and persisted blocks. diff --git a/packages/ai/test/auth-storage-claude-fable-fallback.test.ts b/packages/ai/test/auth-storage-claude-fable-fallback.test.ts index a94395807..c5d337acb 100644 --- a/packages/ai/test/auth-storage-claude-fable-fallback.test.ts +++ b/packages/ai/test/auth-storage-claude-fable-fallback.test.ts @@ -45,7 +45,7 @@ function makeStore(rows: StoredAuthCredential[]): ObservableStore { }; } -function oauthRow(id: number, email: string): StoredAuthCredential { +function oauthRow(id: number, email: string, provider = "anthropic"): StoredAuthCredential { const credential: AuthCredential = { type: "oauth", access: `oat-${id}`, @@ -54,7 +54,7 @@ function oauthRow(id: number, email: string): StoredAuthCredential { accountId: `account-${id}`, email, }; - return { id, provider: "anthropic", credential, disabledCause: null }; + return { id, provider, credential, disabledCause: null }; } function baseReport(email: string): UsageReport { @@ -300,6 +300,63 @@ describe("AuthStorage Claude Fable tier fallback", () => { expect(["oat-2", "oat-3"]).toContain(retryKey as string); }); + it("aborts a local usage lookup before marking the credential blocked", async () => { + let blockUsage = false; + const usageStarted = Promise.withResolvers(); + const releaseUsage = Promise.withResolvers(); + vi.spyOn(claudeUsage.claudeUsageProvider, "fetchUsage").mockImplementation(async () => { + if (!blockUsage) return baseReport("a@example.com"); + usageStarted.resolve(); + return releaseUsage.promise; + }); + + const firstKey = await storage.getApiKey("anthropic", "session-3", { modelId: "claude-fable-5" }); + store.cache.clear(); + blockUsage = true; + const controller = new AbortController(); + const marking = storage.markUsageLimitReached("anthropic", "session-3", { + modelId: "claude-fable-5", + signal: controller.signal, + }); + await usageStarted.promise; + controller.abort(); + let rejection: unknown; + const rejectedQuickly = await Promise.race([ + marking.then( + () => false, + error => { + rejection = error; + return true; + }, + ), + Bun.sleep(50).then(() => false), + ]); + releaseUsage.resolve(baseReport("a@example.com")); + await marking.catch(() => {}); + + expect(rejectedQuickly).toBe(true); + expect(String(rejection)).toContain("usage fetch aborted"); + expect(await storage.getApiKey("anthropic", "session-3", { modelId: "claude-fable-5" })).toBe(firstKey); + }); + + it("does not mark a credential when aborted during target resolution", async () => { + storage.close(); + const provider = "no-usage-provider"; + store = makeStore([oauthRow(1, "a@example.com", provider)]); + storage = new AuthStorage(store); + await storage.reload(); + const firstKey = await storage.getApiKey(provider, "session-3"); + const controller = new AbortController(); + + const marking = storage.markUsageLimitReached(provider, "session-3", { + signal: controller.signal, + }); + controller.abort(); + + await expect(marking).rejects.toThrow(); + expect(await storage.getApiKey(provider, "session-3")).toBe(firstKey); + }); + it("extends a live Fable rate-limit block to the confirmed Fable reset", async () => { const startNow = Date.now(); let now = startNow; diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index d97a1f633..a2a07b3a0 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed the Advisor cost in the status line following the active session across conversation boundaries: a handoff no longer credits the replacement session with the conversation it replaced, a fork carries only the same conversation's finalized spend, and resuming or switching to an existing session restores that session's spend from its persisted Advisor transcripts instead of restarting `(adv)` at zero. Session transitions now also cancel in-flight Advisor maintenance, credential recovery, and retry backoff before swapping conversations, preventing stale Advisor work from blocking `/new`, switch, fork, branch, or handoff ([#6883](https://github.com/can1357/oh-my-pi/pull/6883) by [@paolomazzitti](https://github.com/paolomazzitti)). + ## [17.1.8] - 2026-07-28 ### Breaking Changes diff --git a/packages/coding-agent/src/advisor/runtime.ts b/packages/coding-agent/src/advisor/runtime.ts index b237c1c11..230e75dea 100644 --- a/packages/coding-agent/src/advisor/runtime.ts +++ b/packages/coding-agent/src/advisor/runtime.ts @@ -1,7 +1,9 @@ +import { scheduler } from "node:timers/promises"; import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; import { estimateTokens } from "@oh-my-pi/pi-agent-core/compaction"; import type { AssistantMessage, ImageContent, TextContent } from "@oh-my-pi/pi-ai"; import * as AIError from "@oh-my-pi/pi-ai/error"; +import { raceWithSignal } from "@oh-my-pi/pi-ai/utils/abort"; import { type CursorExecResolvedCarrier, kCursorExecResolved } from "@oh-my-pi/pi-ai/utils/block-symbols"; import { logger } from "@oh-my-pi/pi-utils"; import { obfuscateToolArguments, type SecretObfuscator } from "../secrets/obfuscator"; @@ -47,7 +49,7 @@ export interface AdvisorRuntimeHost { * recovery path must never replay the full primary transcript. * Optional: hosts that omit it get no proactive maintenance. */ - maintainContext?(incomingTokens: number): Promise; + maintainContext?(incomingTokens: number, signal: AbortSignal): Promise; /** * Called immediately before each `agent.prompt(batch)` cycle. Lets the host * clear per-update advisor state — currently the one-advise-per-update gate @@ -67,6 +69,7 @@ export interface AdvisorRuntimeHost { onTurnError?( error: unknown, failedMessages: readonly AgentMessage[], + signal: AbortSignal, ): Promise | boolean | undefined; /** Called after a successful advisor turn so the host can finish fallback lifecycle reporting. */ onTurnSuccess?(): Promise | void; @@ -285,6 +288,9 @@ export class AdvisorRuntime { #advisorRegexSecretValues = new Set(); #pending: PendingDelta[] = []; #busy = false; + #sessionTransitionPaused = false; + #promptInFlight: Promise | undefined; + #iterationAbort: AbortController | undefined; #backlog = 0; #consecutiveFailures = 0; #failureNotified = false; @@ -428,6 +434,7 @@ export class AdvisorRuntime { } dispose(): void { + this.#iterationAbort?.abort("advisor disposed"); this.disposed = true; this.#epoch++; this.#pending = []; @@ -490,6 +497,31 @@ export class AdvisorRuntime { }); } + /** Stop new advisor work and wait only for the active prompt's recorder-visible events. */ + pauseForSessionTransition(): Promise { + if (!this.#sessionTransitionPaused) { + this.#sessionTransitionPaused = true; + this.#wakeAllWaiters(); + this.#iterationAbort?.abort("advisor session transition"); + try { + this.agent.abort("advisor session transition"); + } catch {} + } + return ( + this.#promptInFlight?.then( + () => {}, + () => {}, + ) ?? Promise.resolve() + ); + } + + /** Continue queued work after a session transition rolls back or preserves the conversation. */ + resumeAfterSessionTransition(): void { + if (!this.#sessionTransitionPaused) return; + this.#sessionTransitionPaused = false; + if (!this.#quotaExhausted && !this.#halted) void this.#drain(); + } + /** * Re-prime the advisor after a history rewrite (compaction, session * switch/resume, branch). Clears the advisor's own (non-persisted) context @@ -498,7 +530,9 @@ export class AdvisorRuntime { * leaving it blind to everything before the rewrite. */ reset(): void { + this.#iterationAbort?.abort("advisor reset"); this.#epoch++; + this.#sessionTransitionPaused = false; this.#quotaExhausted = false; this.#halted = false; this.#failing = false; @@ -696,6 +730,7 @@ export class AdvisorRuntime { epoch: number, initial: PendingDelta[], recoveringOverflow: boolean, + signal: AbortSignal, ): Promise<{ batch: string | null; rawMessages: AgentMessage[]; @@ -712,11 +747,12 @@ export class AdvisorRuntime { let wip = initial.at(-1)?.wip ?? false; for (let round = 0; round < MAX_COALESCE_ROUNDS; round++) { + if (this.#sessionTransitionPaused) break; if (this.host.maintainContext) { const incomingTokens = estimateTokens({ role: "user", content: batchText, timestamp: Date.now() }); let shouldResetContext = false; try { - shouldResetContext = await this.host.maintainContext(incomingTokens); + shouldResetContext = await this.host.maintainContext(incomingTokens, signal); } catch (err) { logger.debug("advisor context maintenance failed", { err: String(err) }); } @@ -733,6 +769,7 @@ export class AdvisorRuntime { // remain queued and ship as their own subsequent batch. if (round > 0) { const lateItems = this.#pending.splice(0); + initial.push(...lateItems); turns += lateItems.reduce((sum, b) => sum + b.turns, 0); if (lateItems.length > 0) { wip = lateItems.at(-1)!.wip; @@ -769,6 +806,7 @@ export class AdvisorRuntime { // update WIP state, and re-check the maintenance budget. const late = this.#pending.splice(0); if (late.length === 0) break; + initial.push(...late); batchText = [batchText, ...late.map(b => b.text)].join("\n\n"); rawMessages = rawMessages.concat(late.flatMap(b => b.rawMessages)); turns += late.reduce((sum, b) => sum + b.turns, 0); @@ -802,10 +840,10 @@ export class AdvisorRuntime { } async #drain(): Promise { - if (this.#busy) return; + if (this.#busy || this.#sessionTransitionPaused) return; this.#busy = true; try { - while (!this.disposed && this.#pending.length) { + while (!this.disposed && !this.#sessionTransitionPaused && this.#pending.length) { let popped: PendingDelta[]; if (this.#pending[0]?.overflowRecovery) { const recovery = this.#pending.shift(); @@ -814,6 +852,8 @@ export class AdvisorRuntime { } else { popped = this.#pending.splice(0); } + const iterationAbort = new AbortController(); + this.#iterationAbort = iterationAbort; const epoch = this.#epoch; for (const delta of popped) { if (delta.renderRevision === this.#renderRevision) continue; @@ -822,10 +862,19 @@ export class AdvisorRuntime { delta.renderRevision = this.#renderRevision; } const recoveringOverflow = popped.some(delta => delta.overflowRecovery === true); - const result = await this.#collectAndMaintainBatch(epoch, popped, recoveringOverflow); + const result = await this.#collectAndMaintainBatch( + epoch, + popped, + recoveringOverflow, + iterationAbort.signal, + ); // Epoch was invalidated during batch collection; restart the loop. if (result === null) continue; + if (this.#sessionTransitionPaused) { + this.#pending.unshift(...popped); + continue; + } const { batch, rawMessages, finalTurns, wip, resetContext } = result; @@ -847,7 +896,13 @@ export class AdvisorRuntime { // Reset the host's per-update advisor state (one-advise-per-update // gate) before each model cycle so the new batch starts fresh. this.host.beginAdvisorUpdate?.(); - await this.agent.prompt(batch); + const prompt = this.agent.prompt(batch); + this.#promptInFlight = prompt; + try { + await prompt; + } finally { + if (this.#promptInFlight === prompt) this.#promptInFlight = undefined; + } // Agent.#runLoop catches provider/stream failures internally and // resolves prompt() cleanly with stopReason: "error". Treat that // as a failed turn so endpoint rejections trip the retry path. @@ -868,12 +923,17 @@ export class AdvisorRuntime { this.#consecutiveQuarantines = 0; if (this.host.onTurnSuccess) { try { - await this.host.onTurnSuccess(); + await raceWithSignal(Promise.resolve(this.host.onTurnSuccess()), iterationAbort.signal); } catch (hookErr) { logger.debug("advisor onTurnSuccess hook failed", { err: String(hookErr) }); } } } catch (err) { + if (this.#sessionTransitionPaused) { + this.#rollbackFailedTurn(messageSnapshot); + this.#pending.unshift(...popped); + continue; + } // reset()/dispose() aborts the in-flight prompt; treat it as a // reset, not a transient failure — drop the stale batch. if (this.#epoch !== epoch) continue; @@ -902,10 +962,18 @@ export class AdvisorRuntime { logger.debug("advisor turn failed", { err: String(err) }); let recovered = false; try { - recovered = (await this.host.onTurnError?.(err, failedMessages)) === true; + recovered = + (await raceWithSignal( + Promise.resolve(this.host.onTurnError?.(err, failedMessages, iterationAbort.signal)), + iterationAbort.signal, + )) === true; } catch (hookErr) { logger.debug("advisor onTurnError hook failed", { err: String(hookErr) }); } + if (this.#sessionTransitionPaused) { + this.#pending.unshift(...popped); + continue; + } if (err instanceof AdvisorOutputQuarantinedError) { // A quarantine discards the advisor's whole turn before dispatch, so // its advice never reaches the primary. One re-prime is allowed to @@ -1020,7 +1088,15 @@ export class AdvisorRuntime { wip, overflowRecovery: recoveringOverflow || undefined, }); - await Bun.sleep(this.retryDelayMs); + if (this.retryDelayMs <= 0) { + await Bun.sleep(0); + } else { + try { + await scheduler.wait(this.retryDelayMs, { signal: iterationAbort.signal }); + } catch (sleepError) { + if (!iterationAbort.signal.aborted) throw sleepError; + } + } } } } @@ -1031,6 +1107,7 @@ export class AdvisorRuntime { } } } finally { + this.#iterationAbort = undefined; this.#busy = false; } } diff --git a/packages/coding-agent/src/advisor/transcript-recorder.ts b/packages/coding-agent/src/advisor/transcript-recorder.ts index 5c8602f2c..b893780f1 100644 --- a/packages/coding-agent/src/advisor/transcript-recorder.ts +++ b/packages/coding-agent/src/advisor/transcript-recorder.ts @@ -1,7 +1,9 @@ +import * as fs from "node:fs/promises"; import * as path from "node:path"; import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; import type { Message, UserMessage } from "@oh-my-pi/pi-ai"; import { logger } from "@oh-my-pi/pi-utils"; +import { visitEntriesFromFileStream } from "../session/session-loader"; import { SessionManager } from "../session/session-manager"; /** @@ -31,6 +33,60 @@ export function isAdvisorTranscriptName(name: string): boolean { ); } +/** + * Sum the advisor spend already persisted next to a primary session transcript, + * keyed by advisor slug. + * + * The ledger a session keeps in memory only covers the current process, so a + * resumed session would report zero until the next advisor turn. The recorded + * transcripts are the durable copy of exactly the same finalized messages, so + * they are read back through the shared loader - no lock, no writer, and no + * second parser to keep in step with the session format. + * + * Only the session's own advisors count: subagent advisors write to + * `//__advisor.jsonl`, and their spend belongs to the subagent, + * not to this roster. Hence the scan stays at the top level of the directory. + */ +export async function loadAdvisorTranscriptCosts(sessionFile: string | undefined): Promise> { + const costs = new Map(); + if (!sessionFile?.endsWith(JSONL_SUFFIX)) return costs; + const directory = sessionFile.slice(0, -JSONL_SUFFIX.length); + const dirents = await fs.readdir(directory, { withFileTypes: true }).catch(() => []); + for (const dirent of dirents) { + if (!dirent.isFile() || !isAdvisorTranscriptName(dirent.name)) continue; + const slug = + dirent.name === ADVISOR_TRANSCRIPT_FILENAME + ? "" + : dirent.name.slice(`${ADVISOR_TRANSCRIPT_STEM}.`.length, -JSONL_SUFFIX.length); + let total = 0; + let validHeader: boolean | undefined; + try { + await visitEntriesFromFileStream(path.join(directory, dirent.name), entry => { + const isObject = typeof entry === "object" && entry !== null; + if (validHeader === undefined) { + validHeader = isObject && entry.type === "session" && typeof entry.id === "string"; + return; + } + // A syntactically valid but non-object entry (e.g. a bare `null` + // line) must cost only itself, not crash entry.type access and + // discard everything accumulated for this transcript. + if (!validHeader || !isObject || entry.type !== "message") return; + const message = entry.message; + if (!message || typeof message !== "object" || message.role !== "assistant") return; + // One malformed usage block must cost that entry only, not the + // whole transcript's total. + const total_ = message.usage?.cost?.total; + if (typeof total_ === "number" && Number.isFinite(total_)) total += total_; + }); + } catch (err) { + logger.debug("advisor transcript cost read failed", { file: dirent.name, err: String(err) }); + continue; + } + if (total > 0) costs.set(slug, total); + } + return costs; +} + /** * Append-only persister for an advisor agent's transcript. * diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index afe318051..9778466bd 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -33,6 +33,7 @@ import { discoverWatchdogFiles, formatActiveRepoWatchdogPrompt, formatAdvisorContextPrompt, + loadAdvisorTranscriptCosts, } from "./advisor"; import { AsyncJobManager } from "./async"; import { AutoLearnController, buildAutoLearnInstructions } from "./autolearn/controller"; @@ -3126,6 +3127,9 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} // Owned only when this session created the manager; subagents receive a // parent's manager via `options.mcpManager` and MUST NOT disconnect it. const ownedMcpManager = options.mcpManager ? undefined : mcpManager; + // A resumed session already has advisor turns on disk; without this the status + // line would restart its `(adv)` total at zero for the rest of the session. + const initialAdvisorCosts = await loadAdvisorTranscriptCosts(sessionManager.getSessionFile()); session = new AgentSession({ advisorWatchdogPrompt, advisorContextPrompt, @@ -3140,6 +3144,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} planYolo: options.planYolo, serviceTierByFamily: initialServiceTierByFamily, sessionManager, + initialAdvisorCosts, settings, autoApprove: options.autoApprove, evalKernelOwnerId, diff --git a/packages/coding-agent/src/session/agent-session-types.ts b/packages/coding-agent/src/session/agent-session-types.ts index a2177cb11..75828db4f 100644 --- a/packages/coding-agent/src/session/agent-session-types.ts +++ b/packages/coding-agent/src/session/agent-session-types.ts @@ -161,6 +161,8 @@ export interface AgentSessionConfig { sideStreamFn?: StreamFn; /** Stream wrapper for advisor requests. */ advisorStreamFn?: StreamFn; + /** Advisor spend already recorded for the session being opened, restored on resume. */ + initialAdvisorCosts?: ReadonlyMap; /** Prefer websocket transport for OpenAI Codex requests when supported. */ preferWebsockets?: boolean; /** Provider payload hook used by the active session request path. */ diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index d2eb9e249..e7ebffbe6 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -95,7 +95,7 @@ import { stringProperty, withTimeout, } from "@oh-my-pi/pi-utils"; -import type { AdvisorConfig, AdvisorRuntimeStatus } from "../advisor"; +import { type AdvisorConfig, type AdvisorRuntimeStatus, loadAdvisorTranscriptCosts } from "../advisor"; import { type AsyncJob, AsyncJobManager } from "../async"; import { shouldEnableAppendOnlyContext } from "../config/append-only-context-mode"; import type { ModelRegistry } from "../config/model-registry"; @@ -1274,8 +1274,8 @@ export class AgentSession { hasPendingNextTurnMessages: () => this.#pendingNextTurnMessages.length > 0, convertToLlmForSideRequest: messages => this.#convertToLlmForSideRequest(messages), effectiveServiceTier: model => this.#models.effectiveServiceTier(model), - resolveContextPromotionTarget: (model, contextWindow) => - this.#maintenance.resolveContextPromotionTarget(model, contextWindow), + resolveContextPromotionTarget: (model, contextWindow, signal) => + this.#maintenance.resolveContextPromotionTarget(model, contextWindow, signal), resolveCompactionModelCandidates: (model, availableModels) => this.#maintenance.resolveCompactionModelCandidates(model, availableModels), resolveRetryFallbackRole: (selector, model) => this.#recovery.resolveRetryFallbackRole(selector, model), @@ -1296,6 +1296,7 @@ export class AgentSession { configs: config.advisorConfigs, streamFn: config.advisorStreamFn, transformProviderContext: config.transformProviderContext, + initialCosts: config.initialAdvisorCosts, }); const maintenanceHost: SessionMaintenanceHost = { @@ -1401,7 +1402,10 @@ export class AgentSession { }, resetTodoCycle: () => this.#todo.resetCycle(), buildDisplaySessionContext: () => this.buildDisplaySessionContext(), - resetAdvisorRuntimes: () => this.#advisors.resetAllRuntimes(), + resetAdvisorSessionState: () => this.#advisors.resetSessionState(), + drainAndDetachAdvisorRecorders: () => this.#advisors.drainAndDetachRecorders(), + reattachAdvisorRecorderFeeds: () => this.#advisors.reattachRecorderFeeds(), + clearAdvisorCost: () => this.#advisors.clearCost(), syncTodoPhasesFromBranch: () => this.#todo.syncFromBranch(), }; this.#handoff = new SessionHandoff(handoffHost); @@ -5988,6 +5992,7 @@ export class AgentSession { } this.#disconnectFromAgent(); + let advisorRecordersDetached = false; await this.abort(); this.#cancelOwnAsyncJobs(); this.#closeAllProviderSessions("new session"); @@ -5995,69 +6000,73 @@ export class AgentSession { const bashTransition = this.#bash.beginSessionTransition({ persistDetached: options?.drop !== true }); let sessionTransitioned = false; try { - this.agent.reset(); - if (options?.drop && previousSessionFile) { - // Detach the advisor recorder feed and drain its writer BEFORE deleting the - // old artifacts dir: `await this.abort()` only stops the primary, so a still- - // running advisor turn could otherwise finish, emit `message_end`, and recreate - // `/__advisor.jsonl`. #resetAdvisorSessionState (after newSession) re-primes - // the advisor and re-attaches the feed at the new session's path. - await this.#advisors.detachAndCloseRecorders(); - try { - await this.sessionManager.dropSession(previousSessionFile); - } catch (err) { - logger.error("Failed to delete session during /drop", { err }); + advisorRecordersDetached = true; + await this.#advisors.drainAndDetachRecorders(); + try { + this.agent.reset(); + if (options?.drop && previousSessionFile) { + try { + await this.sessionManager.dropSession(previousSessionFile); + } catch (err) { + logger.error("Failed to delete session during /drop", { err }); + } + } else { + await this.sessionManager.flush(); } - } else { - await this.sessionManager.flush(); + await this.sessionManager.newSession({ + ...options, + additionalDirectories: this.settings.get("workspace.additionalDirectories"), + }); + this.#bash.markSessionTransition(bashTransition); + // The new session owns the transcript from here, so the previous + // conversation's advisor spend is retired with it. Clearing at the commit + // point keeps the status line honest even if a later step below throws. + this.#advisors.clearCost(); + sessionTransitioned = true; + } finally { + this.#bash.finishSessionTransition(bashTransition, sessionTransitioned); } - await this.sessionManager.newSession({ - ...options, - additionalDirectories: this.settings.get("workspace.additionalDirectories"), - }); - this.#bash.markSessionTransition(bashTransition); - // The new session owns the transcript from here, so the previous - // conversation's advisor spend is retired with it. Clearing at the commit - // point keeps the status line honest even if a later step below throws. - this.#advisors.clearCost(); - sessionTransitioned = true; + + this.#clearSessionScopedToolState(); + this.#clearCheckpointRuntimeState(); + this.setTodoPhases([]); + this.#freshProviderSessionId = undefined; + this.#clearInheritedProviderPromptCacheKey(); + this.#syncAgentSessionId(); + this.#memory.rekeyForCurrentSessionId(); + await this.#memory.resetContextForNewTranscript(); + this.#pendingNextTurnMessages = []; + this.#scheduledHiddenNextTurnGeneration = undefined; + + this.sessionManager.appendThinkingLevelChange(this.thinkingLevel, this.configuredThinkingLevel()); + this.sessionManager.appendServiceTierChange(this.#models.serviceTierEntry()); + + this.#todo.resetCycle(); + this.#planReferenceSent = false; + this.#planReferencePath = "local://PLAN.md"; + this.#advisors.resetSessionState(); + advisorRecordersDetached = false; + this.#reconnectToAgent(); + // The workspace-roots block must reflect the new session's directory set, + // not the previous session's — refresh before the next turn goes out. + await this.refreshBaseSystemPrompt(); + + // Emit session_switch event with reason "new" to hooks + if (this.#extensionRunner) { + await this.#extensionRunner.emit({ + type: "session_switch", + reason: "new", + previousSessionFile, + }); + } + + return true; } finally { - this.#bash.finishSessionTransition(bashTransition, sessionTransitioned); + if (advisorRecordersDetached) { + if (sessionTransitioned) this.#advisors.resetSessionState(); + else this.#advisors.reattachRecorderFeeds(); + } } - - this.#clearSessionScopedToolState(); - this.#clearCheckpointRuntimeState(); - this.setTodoPhases([]); - this.#freshProviderSessionId = undefined; - this.#clearInheritedProviderPromptCacheKey(); - this.#syncAgentSessionId(); - this.#memory.rekeyForCurrentSessionId(); - await this.#memory.resetContextForNewTranscript(); - this.#pendingNextTurnMessages = []; - this.#scheduledHiddenNextTurnGeneration = undefined; - - this.sessionManager.appendThinkingLevelChange(this.thinkingLevel, this.configuredThinkingLevel()); - this.sessionManager.appendServiceTierChange(this.#models.serviceTierEntry()); - - this.#todo.resetCycle(); - this.#planReferenceSent = false; - this.#planReferencePath = "local://PLAN.md"; - this.#advisors.resetSessionState(); - this.#reconnectToAgent(); - // The workspace-roots block must reflect the new session's directory set, - // not the previous session's — refresh before the next turn goes out. - await this.refreshBaseSystemPrompt(); - - // Emit session_switch event with reason "new" to hooks - if (this.#extensionRunner) { - await this.#extensionRunner.emit({ - type: "session_switch", - reason: "new", - previousSessionFile, - }); - } - - return true; } /** @@ -6093,59 +6102,70 @@ export class AgentSession { await this.#bash.flushPending(); // Flush current session to ensure all entries are written await this.sessionManager.flush(); - const bashTransition = this.#bash.beginSessionTransition(); - - // Fork the session (creates new session file with same entries) - let forkResult: { oldSessionFile: string; newSessionFile: string } | undefined; + let advisorRecordersDetached = false; try { - forkResult = await this.sessionManager.fork(); - } catch (error) { - this.#bash.finishSessionTransition(bashTransition, false); - throw error; - } - if (!forkResult) { - this.#bash.finishSessionTransition(bashTransition, false); - return false; - } - this.#bash.markSessionTransition(bashTransition); - this.#bash.finishSessionTransition(bashTransition, true); + advisorRecordersDetached = true; + // Fork keeps the conversation, but still needs a quiet artifact boundary: + // stop and settle in-flight advisors before muting their feeds. + await this.#advisors.drainAndDetachRecorders(); + const bashTransition = this.#bash.beginSessionTransition(); - // Copy artifacts directory if it exists - const oldArtifactDir = forkResult.oldSessionFile.slice(0, -6); - const newArtifactDir = forkResult.newSessionFile.slice(0, -6); - - try { - const oldDirStat = await fs.promises.stat(oldArtifactDir); - if (oldDirStat.isDirectory()) { - await fs.promises.cp(oldArtifactDir, newArtifactDir, { recursive: true }); + // Fork the session (creates new session file with same entries) + let forkResult: { oldSessionFile: string; newSessionFile: string } | undefined; + try { + forkResult = await this.sessionManager.fork(); + } catch (error) { + this.#bash.finishSessionTransition(bashTransition, false); + throw error; } - } catch (err) { - if (!isEnoent(err)) { - logger.warn("Failed to copy artifacts during fork", { - oldArtifactDir, - newArtifactDir, - error: err instanceof Error ? err.message : String(err), + if (!forkResult) { + this.#bash.finishSessionTransition(bashTransition, false); + return false; + } + this.#bash.markSessionTransition(bashTransition); + this.#bash.finishSessionTransition(bashTransition, true); + + // Copy artifacts directory if it exists + const oldArtifactDir = forkResult.oldSessionFile.slice(0, -6); + const newArtifactDir = forkResult.newSessionFile.slice(0, -6); + + try { + const oldDirStat = await fs.promises.stat(oldArtifactDir); + if (oldDirStat.isDirectory()) { + await fs.promises.cp(oldArtifactDir, newArtifactDir, { recursive: true }); + } + } catch (err) { + if (!isEnoent(err)) { + logger.warn("Failed to copy artifacts during fork", { + oldArtifactDir, + newArtifactDir, + error: err instanceof Error ? err.message : String(err), + }); + } + } + + // Update agent session ID + this.#freshProviderSessionId = undefined; + this.#adoptInheritedProviderPromptCacheKey(); + this.#syncAgentSessionId(); + this.#memory.rekeyForCurrentSessionId(); + this.#advisors.reattachRecorderFeeds(); + advisorRecordersDetached = false; + await this.#memory.resetContextForNewTranscript(); + + // Emit session_switch event with reason "fork" to hooks + if (this.#extensionRunner) { + await this.#extensionRunner.emit({ + type: "session_switch", + reason: "fork", + previousSessionFile, }); } + + return true; + } finally { + if (advisorRecordersDetached) this.#advisors.reattachRecorderFeeds(); } - - // Update agent session ID - this.#freshProviderSessionId = undefined; - this.#adoptInheritedProviderPromptCacheKey(); - this.#syncAgentSessionId(); - this.#memory.rekeyForCurrentSessionId(); - await this.#memory.resetContextForNewTranscript(); - - // Emit session_switch event with reason "fork" to hooks - if (this.#extensionRunner) { - await this.#extensionRunner.emit({ - type: "session_switch", - reason: "fork", - previousSessionFile, - }); - } - - return true; } /** Move the active session and artifacts after enforcing mode transition invariants. */ @@ -7044,6 +7064,11 @@ export class AgentSession { this.#scheduledHiddenNextTurnGeneration = undefined; try { + if (switchingToDifferentSession) { + // Stop and settle in-flight advisors while the old-session feeds can + // still observe message_end, then mute before swapping files. + await this.#advisors.drainAndDetachRecorders(); + } await this.sessionManager.setSessionFile(sessionPath); this.#bash.markSessionTransition(bashTransition); if (switchingToDifferentSession) { @@ -7179,9 +7204,14 @@ export class AgentSession { error: String(refreshErr), }); } - // Only a committed switch retires the previous conversation's advisor spend: - // an earlier clear would be lost work if any step above rolled the switch back. - if (switchingToDifferentSession) this.#advisors.clearCost(); + // Hand the ledger over to the session that just took over, and only once the + // switch has committed: an earlier swap would be lost work if any step above + // rolled it back. The target's own advisor transcripts are the record of what + // it already spent, so a session with history resumes with its total instead + // of restarting at zero. + if (switchingToDifferentSession) { + this.#advisors.restoreCost(await loadAdvisorTranscriptCosts(this.sessionFile)); + } this.#bash.finishSessionTransition(bashTransition, true); return true; } catch (error) { @@ -7209,6 +7239,7 @@ export class AgentSession { this.#models.restoreServiceTiers(previousServiceTierByFamily); this.#todo.syncFromBranch(); this.#advisors.resetAllRuntimes(); + this.#advisors.reattachRecorderFeeds(); this.#reconnectToAgent(); try { await this.#sessionSwitchReconciler?.(); @@ -7276,45 +7307,57 @@ export class AgentSession { await this.#drainAutolearnCapture(); let sessionTransitioned = false; + let advisorRecordersDetached = false; try { - if (!selectedEntry.parentId) { - await this.sessionManager.newSession({ parentSession: previousSessionFile }); - } else { - this.sessionManager.createBranchedSession(selectedEntry.parentId); + advisorRecordersDetached = true; + await this.#advisors.drainAndDetachRecorders(); + try { + if (!selectedEntry.parentId) { + await this.sessionManager.newSession({ parentSession: previousSessionFile }); + } else { + this.sessionManager.createBranchedSession(selectedEntry.parentId); + } + this.#bash.markSessionTransition(bashTransition); + this.#advisors.clearCost(); + sessionTransitioned = true; + } finally { + this.#bash.finishSessionTransition(bashTransition, sessionTransitioned); } - this.#bash.markSessionTransition(bashTransition); - this.#advisors.clearCost(); - sessionTransitioned = true; + this.#clearSessionScopedToolState(); + this.#rehydrateCheckpointRewindState(); + this.#todo.syncFromBranch(); + this.#freshProviderSessionId = undefined; + this.#clearInheritedProviderPromptCacheKey(); + this.#syncAgentSessionId(); + this.#memory.rekeyForCurrentSessionId(); + await this.#memory.resetContextForNewTranscript(); + + // Reload messages from entries (works for both file and in-memory mode) + const sessionContext = this.buildDisplaySessionContext(); + + // Emit session_branch event to hooks (after branch completes) + if (this.#extensionRunner) { + await this.#extensionRunner.emit({ + type: "session_branch", + previousSessionFile, + }); + } + + if (!skipConversationRestore) { + this.agent.replaceMessages(sessionContext.messages); + this.#advisors.resetSessionState(); + this.#closeCodexProviderSessionsForHistoryRewrite(); + } + + this.#advisors.reattachRecorderFeeds(); + advisorRecordersDetached = false; + return { selectedText, selectedImages, cancelled: false }; } finally { - this.#bash.finishSessionTransition(bashTransition, sessionTransitioned); + if (advisorRecordersDetached) { + if (sessionTransitioned) this.#advisors.resetSessionState(); + else this.#advisors.reattachRecorderFeeds(); + } } - this.#clearSessionScopedToolState(); - this.#rehydrateCheckpointRewindState(); - this.#todo.syncFromBranch(); - this.#freshProviderSessionId = undefined; - this.#clearInheritedProviderPromptCacheKey(); - this.#syncAgentSessionId(); - this.#memory.rekeyForCurrentSessionId(); - await this.#memory.resetContextForNewTranscript(); - - // Reload messages from entries (works for both file and in-memory mode) - const sessionContext = this.buildDisplaySessionContext(); - - // Emit session_branch event to hooks (after branch completes) - if (this.#extensionRunner) { - await this.#extensionRunner.emit({ - type: "session_branch", - previousSessionFile, - }); - } - - if (!skipConversationRestore) { - this.agent.replaceMessages(sessionContext.messages); - this.#advisors.resetSessionState(); - this.#closeCodexProviderSessionsForHistoryRewrite(); - } - - return { selectedText, selectedImages, cancelled: false }; } async branchFromBtw( @@ -7378,44 +7421,55 @@ export class AgentSession { await this.#drainAutolearnCapture(); let sessionTransitioned = false; + let advisorRecordersDetached = false; try { - this.sessionManager.createBranchedSession(leafId); - this.#bash.markSessionTransition(bashTransition); - this.#advisors.clearCost(); - sessionTransitioned = true; - } finally { - this.#bash.finishSessionTransition(bashTransition, sessionTransitioned); - } + advisorRecordersDetached = true; + await this.#advisors.drainAndDetachRecorders(); + try { + this.sessionManager.createBranchedSession(leafId); + this.#bash.markSessionTransition(bashTransition); + this.#advisors.clearCost(); + sessionTransitioned = true; + } finally { + this.#bash.finishSessionTransition(bashTransition, sessionTransitioned); + } - this.#clearSessionScopedToolState(); + this.#clearSessionScopedToolState(); - this.#rehydrateCheckpointRewindState(); - this.sessionManager.appendMessage({ - role: "user", - content: [{ type: "text", text: question }], - timestamp: Date.now(), - }); - this.sessionManager.appendMessage(sanitizeAssistantForReparentedHistory(assistantMessage)); - this.#todo.syncFromBranch(); - this.#freshProviderSessionId = undefined; - this.#syncAgentSessionId(); - this.#memory.rekeyForCurrentSessionId(); - await this.#memory.resetContextForNewTranscript(); - - const sessionContext = this.buildDisplaySessionContext(); - - if (this.#extensionRunner) { - await this.#extensionRunner.emit({ - type: "session_branch", - previousSessionFile, + this.#rehydrateCheckpointRewindState(); + this.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: question }], + timestamp: Date.now(), }); + this.sessionManager.appendMessage(sanitizeAssistantForReparentedHistory(assistantMessage)); + this.#todo.syncFromBranch(); + this.#freshProviderSessionId = undefined; + this.#syncAgentSessionId(); + this.#memory.rekeyForCurrentSessionId(); + await this.#memory.resetContextForNewTranscript(); + + const sessionContext = this.buildDisplaySessionContext(); + + if (this.#extensionRunner) { + await this.#extensionRunner.emit({ + type: "session_branch", + previousSessionFile, + }); + } + + this.agent.replaceMessages(sessionContext.messages); + this.#advisors.resetSessionState(); + this.#closeCodexProviderSessionsForHistoryRewrite(); + advisorRecordersDetached = false; + + return { cancelled: false, sessionFile: this.sessionFile }; + } finally { + if (advisorRecordersDetached) { + if (sessionTransitioned) this.#advisors.resetSessionState(); + else this.#advisors.reattachRecorderFeeds(); + } } - - this.agent.replaceMessages(sessionContext.messages); - this.#advisors.resetSessionState(); - this.#closeCodexProviderSessionsForHistoryRewrite(); - - return { cancelled: false, sessionFile: this.sessionFile }; } // ========================================================================= diff --git a/packages/coding-agent/src/session/session-advisors.ts b/packages/coding-agent/src/session/session-advisors.ts index c0ece6a39..b18d7db08 100644 --- a/packages/coding-agent/src/session/session-advisors.ts +++ b/packages/coding-agent/src/session/session-advisors.ts @@ -183,6 +183,8 @@ export interface SessionAdvisorsOptions { configs?: AdvisorConfig[]; streamFn?: StreamFn; transformProviderContext?: (context: Context, model: Model) => Context | Promise; + /** Advisor spend already persisted for this session, restored on resume. */ + initialCosts?: ReadonlyMap; } /** Options accepted when an advisor injects a primary-session message. */ @@ -221,7 +223,11 @@ export interface SessionAdvisorsHost { hasPendingNextTurnMessages(): boolean; convertToLlmForSideRequest(messages: AgentMessage[]): Message[]; effectiveServiceTier(model: Model): ServiceTier | undefined; - resolveContextPromotionTarget(currentModel: Model, contextWindow: number): Promise; + resolveContextPromotionTarget( + currentModel: Model, + contextWindow: number, + signal: AbortSignal, + ): Promise; resolveCompactionModelCandidates(preferredModel: Model | null | undefined, availableModels: Model[]): Model[]; resolveRetryFallbackRole(currentSelector: string, currentModel?: Model | null): string | undefined; findRetryFallbackCandidates( @@ -272,6 +278,7 @@ export class SessionAdvisors { this.#advisorConfigs = options.configs; this.#advisorStreamFn = options.streamFn; this.#transformProviderContext = options.transformProviderContext; + if (options.initialCosts) this.#advisorCosts = new Map(options.initialCosts); if (this.#advisorEnabled) this.#buildAdvisorRuntime(); } @@ -313,6 +320,15 @@ export class SessionAdvisors { this.#stopAdvisorRuntime(); } + /** + * Pause advisor work while old-session recorder feeds remain attached, then + * detach only after any active prompt has settled. + */ + async drainAndDetachRecorders(): Promise { + await Promise.all(this.#advisors.map(advisor => advisor.runtime.pauseForSessionTransition())); + await this.detachAndCloseRecorders(); + } + /** Detaches and drains recorder feeds before transcript artifacts are removed. */ async detachAndCloseRecorders(): Promise { const closes: Promise[] = []; @@ -325,6 +341,14 @@ export class SessionAdvisors { await Promise.all(closes); } + /** Reattach recorder feeds and resume work after a rolled-back or preserving transition. */ + reattachRecorderFeeds(): void { + for (const advisor of this.#advisors) { + if (!advisor.agentUnsubscribe) this.#attachAdvisorRecorderFeed(advisor); + advisor.runtime.resumeAfterSessionTransition(); + } + } + /** Re-primes advisor transcript views across a conversation boundary. */ resetSessionState(options: { preserveCost?: boolean } = {}): void { this.#resetAdvisorSessionState(options.preserveCost === true); @@ -335,6 +359,11 @@ export class SessionAdvisors { this.#advisorCosts.clear(); } + /** Replace the ledger with the spend recorded for the session becoming active. */ + restoreCost(costs: ReadonlyMap): void { + this.#advisorCosts = new Map(costs); + } + /** * Rebind every live advisor to the active primary conversation's provider * identity (session id, prompt-cache key, credential + metadata resolvers, @@ -784,10 +813,12 @@ export class SessionAdvisors { const runtime = new AdvisorRuntime(advisorAgentFacade, { snapshotMessages: () => this.#host.agent.state.messages, enqueueAdvice: (note, severity) => this.#routeAdvice(advisorRef, note, severity), - maintainContext: incomingTokens => this.#maintainAdvisorContext(advisorRef, incomingTokens), + maintainContext: (incomingTokens, signal) => + this.#maintainAdvisorContext(advisorRef, incomingTokens, signal), obfuscator: this.#host.obfuscator, beginAdvisorUpdate: () => advisorRef.emissionGuard.beginUpdate(), - onTurnError: (error, failedMessages) => this.#recoverAdvisorTurn(advisorRef, error, failedMessages), + onTurnError: (error, failedMessages, signal) => + this.#recoverAdvisorTurn(advisorRef, error, failedMessages, signal), onTurnSuccess: async () => { const fallback = advisorRef.retryFallback; if (!advisorRef.retryFallbackPendingSuccess || !fallback) return; @@ -1015,7 +1046,7 @@ export class SessionAdvisors { } /** Restore an advisor's configured primary once its fallback cooldown expires. */ - async #maybeRestoreAdvisorRetryFallbackPrimary(advisor: ActiveAdvisor): Promise { + async #maybeRestoreAdvisorRetryFallbackPrimary(advisor: ActiveAdvisor, signal: AbortSignal): Promise { const fallback = advisor.retryFallback; if (!fallback || getRetryFallbackRevertPolicy(this.#host.settings) !== "cooldown-expiry") return; @@ -1043,8 +1074,9 @@ export class SessionAdvisors { const primaryModel = resolvedPrimary.model ?? this.#host.modelRegistry.find(originalSelector.provider, originalSelector.id); if (!primaryModel) return; - const apiKey = await this.#host.modelRegistry.getApiKey(primaryModel, advisor.providerSessionId); + const apiKey = await this.#host.modelRegistry.getApiKey(primaryModel, advisor.providerSessionId, { signal }); if (!apiKey) return; + signal.throwIfAborted(); const thinkingToApply = advisor.thinkingLevel === fallback.lastAppliedThinkingLevel @@ -1064,6 +1096,7 @@ export class SessionAdvisors { advisor: ActiveAdvisor, error: unknown, failedMessages: readonly AgentMessage[], + signal: AbortSignal, ): Promise { if (error instanceof AdvisorOutputQuarantinedError) return false; @@ -1086,6 +1119,7 @@ export class SessionAdvisors { retryAfterMs: extractRetryHint(undefined, message), baseUrl: currentModel.baseUrl, modelId: currentModel.id, + signal, }, ); return outcome.switched; @@ -1117,6 +1151,7 @@ export class SessionAdvisors { retryAfterMs, baseUrl: currentModel.baseUrl, modelId: currentModel.id, + signal, }, ); if (outcome.switched) return true; @@ -1134,8 +1169,9 @@ export class SessionAdvisors { const resolved = resolveModelOverride([selector.raw], this.#host.modelRegistry, this.#host.settings); const candidate = resolved.model ?? this.#host.modelRegistry.find(selector.provider, selector.id); if (!candidate || modelsAreEqual(candidate, currentModel)) continue; - const apiKey = await this.#host.modelRegistry.getApiKey(candidate, advisor.providerSessionId); + const apiKey = await this.#host.modelRegistry.getApiKey(candidate, advisor.providerSessionId, { signal }); if (!apiKey) continue; + signal.throwIfAborted(); const originalThinkingLevel = advisor.thinkingLevel; const requestedThinkingLevel = selector.thinkingLevel ?? originalThinkingLevel; @@ -1163,13 +1199,18 @@ export class SessionAdvisors { return false; } - async #promoteAdvisorContextModel(advisor: ActiveAdvisor, currentModel: Model): Promise { + async #promoteAdvisorContextModel( + advisor: ActiveAdvisor, + currentModel: Model, + signal: AbortSignal, + ): Promise { const promotionSettings = this.#host.settings.getGroup("contextPromotion"); if (!promotionSettings.enabled) return false; const contextWindow = currentModel.contextWindow ?? 0; if (contextWindow <= 0) return false; - const targetModel = await this.#host.resolveContextPromotionTarget(currentModel, contextWindow); + const targetModel = await this.#host.resolveContextPromotionTarget(currentModel, contextWindow, signal); if (!targetModel) return false; + signal.throwIfAborted(); // Preserve this advisor's own thinking level (a configured `model:...:high` // keeps its suffix across a promotion); only the model changes. @@ -1193,8 +1234,12 @@ export class SessionAdvisors { } } - async #maintainAdvisorContext(advisor: ActiveAdvisor, incomingTokens: number): Promise { - await this.#maybeRestoreAdvisorRetryFallbackPrimary(advisor); + async #maintainAdvisorContext( + advisor: ActiveAdvisor, + incomingTokens: number, + signal: AbortSignal, + ): Promise { + await this.#maybeRestoreAdvisorRetryFallbackPrimary(advisor, signal); const agent = advisor.agent; const compactionSettings = this.#host.settings.getGroup("compaction"); @@ -1229,7 +1274,7 @@ export class SessionAdvisors { } // 1. Try promotion first - if (await this.#promoteAdvisorContextModel(advisor, advisorModel)) { + if (await this.#promoteAdvisorContextModel(advisor, advisorModel, signal)) { // Promotion succeeded, check if new model has enough space const newModel = agent.state.model; const newWindow = newModel.contextWindow ?? 0; @@ -1307,7 +1352,7 @@ export class SessionAdvisors { }); for (const candidate of candidates) { - const apiKey = await this.#host.modelRegistry.getApiKey(candidate, advisorProviderSessionId); + const apiKey = await this.#host.modelRegistry.getApiKey(candidate, advisorProviderSessionId, { signal }); if (!apiKey) continue; // The advisor overflow-compaction one-shot bypasses the advisor `Agent`, // so its installed metadata resolver never runs. Emit the same @@ -1324,7 +1369,7 @@ export class SessionAdvisors { candidate, this.#host.modelRegistry.resolver(candidate, advisorProviderSessionId), undefined, - undefined, + signal, { thinkingLevel: advisorCompactionThinkingLevel, convertToLlm: messages => this.#host.convertToLlmForSideRequest(messages), @@ -1339,6 +1384,7 @@ export class SessionAdvisors { ); break; } catch (error) { + if (signal.aborted) throw error; lastError = error; } } diff --git a/packages/coding-agent/src/session/session-handoff.ts b/packages/coding-agent/src/session/session-handoff.ts index 33dda9b5f..7e268e5f4 100644 --- a/packages/coding-agent/src/session/session-handoff.ts +++ b/packages/coding-agent/src/session/session-handoff.ts @@ -64,7 +64,10 @@ export interface SessionHandoffHost { clearPendingNextTurnMessages(): void; resetTodoCycle(): void; buildDisplaySessionContext(): SessionContext; - resetAdvisorRuntimes(): void; + resetAdvisorSessionState(): void; + drainAndDetachAdvisorRecorders(): Promise; + reattachAdvisorRecorderFeeds(): void; + clearAdvisorCost(): void; syncTodoPhasesFromBranch(): void; } @@ -124,6 +127,8 @@ export class SessionHandoff { } } + let advisorRecordersDetached = false; + let sessionTransitioned = false; try { if (handoffSignal.aborted) { throw new Error("Handoff cancelled"); @@ -224,14 +229,21 @@ export class SessionHandoff { } await this.#host.flushPendingBash(); await this.#host.sessionManager.flush(); + advisorRecordersDetached = true; + // Stop and settle in-flight advisors while the old-session feeds can still + // observe message_end, then mute before opening the replacement session. + await this.#host.drainAndDetachAdvisorRecorders(); const bashTransition = this.#host.beginBashSessionTransition(); this.#host.cancelOwnAsyncJobs(); - let sessionTransitioned = false; try { await this.#host.sessionManager.newSession( previousSessionFile ? { parentSession: previousSessionFile } : undefined, ); this.#host.markBashSessionTransition(bashTransition); + // The handoff opens a fresh conversation, so the spend of the one it + // summarizes stays with it. Clearing here, at the commit point, keeps the + // status line honest even if a later step throws. + this.#host.clearAdvisorCost(); sessionTransitioned = true; } finally { this.#host.finishBashSessionTransition(bashTransition, sessionTransitioned); @@ -284,7 +296,8 @@ export class SessionHandoff { // Rebuild agent messages from session const sessionContext = this.#host.buildDisplaySessionContext(); this.#host.agent.replaceMessages(sessionContext.messages); - this.#host.resetAdvisorRuntimes(); + this.#host.resetAdvisorSessionState(); + advisorRecordersDetached = false; this.#host.syncTodoPhasesFromBranch(); if (this.#host.extensionRunner) { await this.#host.extensionRunner.emit({ @@ -301,6 +314,10 @@ export class SessionHandoff { } throw error; } finally { + if (advisorRecordersDetached) { + if (sessionTransitioned) this.#host.resetAdvisorSessionState(); + else this.#host.reattachAdvisorRecorderFeeds(); + } sourceSignal?.removeEventListener("abort", onSourceAbort); this.#handoffAbortController = undefined; } diff --git a/packages/coding-agent/src/session/session-loader.ts b/packages/coding-agent/src/session/session-loader.ts index 3ddce2814..f493e974e 100644 --- a/packages/coding-agent/src/session/session-loader.ts +++ b/packages/coding-agent/src/session/session-loader.ts @@ -56,14 +56,14 @@ export function parseSessionContent(content: string): { return { entries: foldTitleSlot(entries, slot), titleSlot: slot }; } -/** Exported for testing — the ≥8MiB streaming path (works on any file size). */ -export async function loadEntriesFromFileStream(filePath: string): Promise<{ - entries: FileEntry[]; - titleSlot: SessionTitleUpdate | undefined; -}> { - const entries: FileEntry[] = []; +/** Parse session JSONL and visit each entry without retaining prior entries. */ +export async function visitEntriesFromFileStream( + filePath: string, + visit: (entry: FileEntry) => void, +): Promise { let titleSlot: SessionTitleUpdate | undefined; let sawFirstLine = false; + let visitorThrew = false; // Byte buffer (NOT a decoded string): multibyte UTF-8 sequences that straddle // a stream-chunk boundary stay intact, and Bun.JSONL.parseChunk accepts typed // arrays directly. Only the unconsumed remainder is held (≤ one record + a @@ -75,8 +75,13 @@ export async function loadEntriesFromFileStream(filePath: string): Promise<{ const drain = () => { while (buffer.length > 0) { const { values, error, read, done } = Bun.JSONL.parseChunk(buffer); - if (values.length > 0) { - for (const value of values) entries.push(value as FileEntry); + for (const value of values) { + try { + visit(value as FileEntry); + } catch (err) { + visitorThrew = true; + throw err; + } } if (error) { // Malformed record: skip past the next newline and continue. @@ -125,10 +130,21 @@ export async function loadEntriesFromFileStream(filePath: string): Promise<{ } drain(); } catch (err) { - if (isEnoent(err)) return { entries: [], titleSlot: undefined }; + if (visitorThrew) throw err; + if (isEnoent(err)) return undefined; throw err; } + return titleSlot; +} + +/** Exported for testing — the ≥8MiB streaming path (works on any file size). */ +export async function loadEntriesFromFileStream(filePath: string): Promise<{ + entries: FileEntry[]; + titleSlot: SessionTitleUpdate | undefined; +}> { + const entries: FileEntry[] = []; + const titleSlot = await visitEntriesFromFileStream(filePath, entry => entries.push(entry)); return { entries: foldTitleSlot(entries, titleSlot), titleSlot }; } diff --git a/packages/coding-agent/src/session/session-maintenance.ts b/packages/coding-agent/src/session/session-maintenance.ts index d282b4b68..df65f9a32 100644 --- a/packages/coding-agent/src/session/session-maintenance.ts +++ b/packages/coding-agent/src/session/session-maintenance.ts @@ -1373,7 +1373,11 @@ export class SessionMaintenance { } } - async resolveContextPromotionTarget(currentModel: Model, contextWindow: number): Promise { + async resolveContextPromotionTarget( + currentModel: Model, + contextWindow: number, + signal?: AbortSignal, + ): Promise { const availableModels = this.#host.modelRegistry.getAvailable(); if (availableModels.length === 0) return undefined; @@ -1381,7 +1385,7 @@ export class SessionMaintenance { if (!candidate) return undefined; if (modelsAreEqual(candidate, currentModel)) return undefined; if (candidate.contextWindow == null || candidate.contextWindow <= contextWindow) return undefined; - const apiKey = await this.#host.modelRegistry.getApiKey(candidate, this.#host.sessionId()); + const apiKey = await this.#host.modelRegistry.getApiKey(candidate, this.#host.sessionId(), { signal }); if (!apiKey) return undefined; return candidate; } diff --git a/packages/coding-agent/test/advisor-context-maintenance.test.ts b/packages/coding-agent/test/advisor-context-maintenance.test.ts index 831588a80..423565e7e 100644 --- a/packages/coding-agent/test/advisor-context-maintenance.test.ts +++ b/packages/coding-agent/test/advisor-context-maintenance.test.ts @@ -19,6 +19,7 @@ const OUTPUT_TOKENS = 150; interface MaintenanceHarness { advisor: Agent; advisorMock: MockModel; + modelRegistry: ModelRegistry; settings: Settings; } @@ -44,7 +45,7 @@ describe("AgentSession advisor context maintenance", () => { await tempDir.remove(); }); - function createHarness(): MaintenanceHarness { + function createHarness(contextPromotionTarget?: string, contextPromotionEnabled = false): MaintenanceHarness { const primaryMock = createMockModel({ provider: "anthropic", responses: [{ content: ["primary complete"] }], @@ -54,12 +55,13 @@ describe("AgentSession advisor context maintenance", () => { contextWindow: CONTEXT_WINDOW, responses: [{ content: ["advisor reviewed current update"] }], }); + Object.assign(advisorMock, { contextPromotionTarget }); const modelRegistry = new ModelRegistry(authStorage, tempDir.join("models.yml")); const settings = Settings.isolated({ "advisor.syncBacklog": "1", "compaction.enabled": true, "compaction.strategy": "context-full", - "contextPromotion.enabled": false, + "contextPromotion.enabled": contextPromotionEnabled, }); const agent = new Agent({ getApiKey: () => "test-key", @@ -85,7 +87,7 @@ describe("AgentSession advisor context maintenance", () => { vi.spyOn(modelRegistry, "getApiKey").mockImplementation(async model => model === primaryMock ? "test-key" : undefined, ); - return { advisor, advisorMock, settings }; + return { advisor, advisorMock, modelRegistry, settings }; } function usageAnchor(advisorMock: MockModel, timestamp: number, cost = 0): AssistantMessage { @@ -144,6 +146,44 @@ describe("AgentSession advisor context maintenance", () => { expect(session.getAdvisorCost()).toBeCloseTo(0.5, 8); }); + it("ignores late context-promotion credentials after a session transition", async () => { + const promotion = createMockModel({ + id: "advisor-promotion-target", + provider: "anthropic", + contextWindow: CONTEXT_WINDOW + 1, + }); + const { advisor, advisorMock, modelRegistry } = createHarness(`${promotion.provider}/${promotion.id}`, true); + vi.spyOn(modelRegistry, "getAvailable").mockReturnValue([advisor.state.model, promotion]); + const credentialStarted = Promise.withResolvers(); + const releaseCredential = Promise.withResolvers(); + const credentialReturned = Promise.withResolvers(); + let credentialSignal: AbortSignal | undefined; + vi.spyOn(modelRegistry, "getApiKey").mockImplementation(async (model, _sessionId, options) => { + if (model === promotion) { + credentialSignal = options?.signal; + credentialStarted.resolve(); + await releaseCredential.promise; + credentialReturned.resolve(); + } + return "test-key"; + }); + advisor.emitExternalEvent({ + type: "message_end", + message: usageAnchor(advisorMock, Date.now() - 1_000), + }); + + const prompt = session.prompt("trigger advisor context promotion"); + await credentialStarted.promise; + await session.newSession(); + releaseCredential.resolve(); + await credentialReturned.promise; + await prompt; + await Bun.sleep(0); + + expect(credentialSignal?.aborted).toBe(true); + expect(session.getAdvisorAgent()?.state.model).toBe(advisorMock); + }); + it("includes advisor system prompt and tool schemas in the local maintenance floor", async () => { const { advisor, advisorMock, settings } = createHarness(); const seed: AgentMessage = { role: "user", content: "small stored advisor message", timestamp: 1 }; @@ -199,7 +239,7 @@ describe("AgentSession advisor context maintenance", () => { expect(sentContext).not.toContain("fresh post-compaction output"); }); - it("forwards the advisor session metadata to overflow compaction requests", async () => { + it("forwards compaction metadata and aborts transitions without fallback or re-prime", async () => { // Regression for #6625 review: advisor overflow compaction issues a direct // `compact(...)` request that bypasses the advisor `Agent`, so the metadata // resolver installed on the agent never runs for it. The direct call must @@ -208,6 +248,9 @@ describe("AgentSession advisor context maintenance", () => { // API lets the compaction one-shot's `completeSimple` route to it so the // summarization request actually reaches the mock (and its recorded calls). registerMockApi(); + const compactionStarted = Promise.withResolvers(); + const releaseCompaction = Promise.withResolvers(); + let fallbackCalls = 0; const primaryMock = createMockModel({ provider: "anthropic", responses: [{ content: ["primary complete"] }], @@ -215,7 +258,28 @@ describe("AgentSession advisor context maintenance", () => { const advisorMock = createMockModel({ provider: "anthropic", contextWindow: CONTEXT_WINDOW, - handler: () => ({ content: ["bounded advisor summary"] }), + handler: async (context, options) => { + if (!JSON.stringify(context.messages).includes("")) { + return { content: ["advisor reviewed current update"] }; + } + compactionStarted.resolve(); + const signal = options?.signal; + if (!signal) throw new Error("Expected compaction abort signal"); + const compactionAborted = Promise.withResolvers(); + signal.addEventListener("abort", () => compactionAborted.resolve(), { once: true }); + await Promise.race([releaseCompaction.promise, compactionAborted.promise]); + signal.throwIfAborted(); + return { content: ["bounded advisor summary"] }; + }, + }); + const fallbackMock = createMockModel({ + id: "advisor-compaction-fallback", + provider: "anthropic", + contextWindow: CONTEXT_WINDOW, + handler: () => { + fallbackCalls++; + return { content: ["unexpected fallback"] }; + }, }); const modelRegistry = new ModelRegistry(authStorage, tempDir.join("models.yml")); const settings = Settings.isolated({ @@ -242,9 +306,10 @@ describe("AgentSession advisor context maintenance", () => { const advisor = session.getAdvisorAgent(); if (!advisor?.sessionId) throw new Error("Expected advisor agent with a provider session id"); advisor.setModel(advisorMock); + vi.spyOn(modelRegistry, "getAvailable").mockReturnValue([advisorMock, fallbackMock]); // Unlike the recovery-branch harness, the advisor holds usable credentials // so maintenance runs the LLM summarization compaction path. - vi.spyOn(modelRegistry, "getApiKey").mockResolvedValue("test-key"); + const getApiKey = vi.spyOn(modelRegistry, "getApiKey").mockResolvedValue("test-key"); // Two accumulated turns so compaction has older history to summarize while // retaining the most recent one (a single message would be fully retained, @@ -253,8 +318,24 @@ describe("AgentSession advisor context maintenance", () => { usageAnchor(advisorMock, Date.now() - 2_000), usageAnchor(advisorMock, Date.now() - 1_000), ); + const previousAdvisorMessages = [...advisor.state.messages]; - await session.prompt("small current update"); + const prompt = session.prompt("small current update"); + await compactionStarted.promise; + const failure = new Error("new session failed"); + vi.spyOn(session.sessionManager, "newSession").mockRejectedValue(failure); + const transition = session.newSession(); + try { + await expect(transition).rejects.toThrow(failure); + expect(fallbackCalls).toBe(0); + expect(advisor.state.messages).toEqual(previousAdvisorMessages); + expect(getApiKey).toHaveBeenCalledWith(advisorMock, advisor.sessionId, { + signal: expect.any(AbortSignal), + }); + } finally { + releaseCompaction.resolve(); + await prompt; + } // A summarization compaction one-shot actually ran (its prompt wraps the // conversation in tags). @@ -262,6 +343,7 @@ describe("AgentSession advisor context maintenance", () => { JSON.stringify(call.context.messages).includes(""), ); expect(compactionCalls.length).toBeGreaterThan(0); + expect(compactionCalls.every(call => call.options?.signal instanceof AbortSignal)).toBe(true); // Every advisor request — the compaction one-shot and the advisor turn — // carries the advisor's own provider session id via metadata.user_id. diff --git a/packages/coding-agent/test/advisor-toggle.test.ts b/packages/coding-agent/test/advisor-toggle.test.ts index eb8698c37..41c35746b 100644 --- a/packages/coding-agent/test/advisor-toggle.test.ts +++ b/packages/coding-agent/test/advisor-toggle.test.ts @@ -1,14 +1,17 @@ import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "bun:test"; -import * as fs from "node:fs"; +import * as fs from "node:fs/promises"; import * as path from "node:path"; import { Agent, type AgentMessage } from "@oh-my-pi/pi-agent-core"; -import type { Model } from "@oh-my-pi/pi-ai"; +import * as compactionModule from "@oh-my-pi/pi-agent-core/compaction"; +import type { AssistantMessage, Model } from "@oh-my-pi/pi-ai"; import * as AIError from "@oh-my-pi/pi-ai/error"; import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock"; import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; +import { loadAdvisorTranscriptCosts } from "@oh-my-pi/pi-coding-agent/advisor/transcript-recorder"; import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import type { ExtensionRunner } from "@oh-my-pi/pi-coding-agent/extensibility/extensions"; +import { createAgentSession } from "@oh-my-pi/pi-coding-agent/sdk"; import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; import { AgentStorage } from "@oh-my-pi/pi-coding-agent/session/agent-storage"; import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; @@ -76,8 +79,8 @@ describe("AgentSession advisor toggle", () => { } catch {} }); - function appendAdvisorCost(advisor: Agent, cost: number, timestamp: number): void { - const message: AgentMessage = { + function advisorMessage(cost: number, timestamp: number): AssistantMessage { + return { role: "assistant", content: [{ type: "text", text: "reviewed" }], api: "anthropic-messages", @@ -94,7 +97,59 @@ describe("AgentSession advisor toggle", () => { stopReason: "stop", timestamp, }; - advisor.emitExternalEvent({ type: "message_end", message }); + } + + function appendAdvisorCost(advisor: Agent, cost: number, timestamp: number): void { + advisor.emitExternalEvent({ type: "message_end", message: advisorMessage(cost, timestamp) }); + } + + function enableAdvisor(target: AgentSession = session): Agent { + target.settings.setModelRole("advisor", `${model.provider}/${model.id}`); + target.toggleAdvisorEnabled(); + const advisor = target.getAdvisorAgent(); + if (!advisor) throw new Error("Expected advisor agent to exist"); + return advisor; + } + + /** + * Persist advisor turns beside a session file the same way the recorder does, + * so the fixture stays valid if the transcript format ever moves. + */ + async function writeAdvisorTranscript(sessionFile: string, filename: string, costs: number[]): Promise { + const dir = sessionFile.slice(0, -".jsonl".length); + await fs.mkdir(dir, { recursive: true }); + const manager = await SessionManager.open(path.join(dir, filename), undefined, undefined, { + initialCwd: dir, + suppressBreadcrumb: true, + }); + try { + for (const [index, cost] of costs.entries()) manager.appendMessage(advisorMessage(cost, index + 1)); + } finally { + await manager.close(); + } + } + + function prepareHandoffConversation(advisor: Agent): void { + sessionManager.appendMessage({ role: "user", content: "work to hand off", timestamp: 1 }); + sessionManager.appendMessage({ + role: "assistant", + content: [{ type: "text", text: "done" }], + api: "anthropic-messages", + provider: model.provider, + model: model.id, + usage: { + input: 1, + output: 1, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 2, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: 2, + }); + session.agent.replaceMessages(sessionManager.buildSessionContext().messages); + appendAdvisorCost(advisor, 0.5, 1); } it("starts with advisor disabled", () => { @@ -161,9 +216,7 @@ describe("AgentSession advisor toggle", () => { const projectA = path.join(tempDir.path(), "project-a"); const projectB = path.join(tempDir.path(), "project-b"); const agentDir = path.join(tempDir.path(), "agent"); - fs.mkdirSync(getProjectAgentDir(projectA), { recursive: true }); - fs.mkdirSync(getProjectAgentDir(projectB), { recursive: true }); - fs.mkdirSync(agentDir, { recursive: true }); + await fs.mkdir(agentDir, { recursive: true }); await Bun.write( path.join(getProjectAgentDir(projectA), "settings.json"), JSON.stringify({ modelRoles: { advisor: `${model.provider}/${model.id}` } }), @@ -321,10 +374,7 @@ describe("AgentSession advisor toggle", () => { expect(sid).not.toContain("-advisor"); }); it("retains cumulative advisor cost after the advisor is disabled", () => { - session.settings.setModelRole("advisor", `${model.provider}/${model.id}`); - session.toggleAdvisorEnabled(); - const advisor = session.getAdvisorAgent(); - if (!advisor) throw new Error("Expected advisor agent to exist"); + const advisor = enableAdvisor(); appendAdvisorCost(advisor, 0.41, 1); appendAdvisorCost(advisor, 0.09, 2); @@ -334,10 +384,7 @@ describe("AgentSession advisor toggle", () => { expect(session.getAdvisorCost()).toBeCloseTo(0.5, 8); }); it("retains total advisor cost after the live roster changes", () => { - session.settings.setModelRole("advisor", `${model.provider}/${model.id}`); - session.toggleAdvisorEnabled(); - const advisor = session.getAdvisorAgent(); - if (!advisor) throw new Error("Expected advisor agent to exist"); + const advisor = enableAdvisor(); appendAdvisorCost(advisor, 0.5, 1); expect(session.applyAdvisorConfigs([{ name: "Security" }], undefined)).toBe(1); @@ -345,10 +392,7 @@ describe("AgentSession advisor toggle", () => { expect(session.formatAdvisorStatus()).toContain("$0.5000"); }); it("retains cumulative advisor cost after an in-session history rewrite", async () => { - session.settings.setModelRole("advisor", `${model.provider}/${model.id}`); - session.toggleAdvisorEnabled(); - const advisor = session.getAdvisorAgent(); - if (!advisor) throw new Error("Expected advisor agent to exist"); + const advisor = enableAdvisor(); appendAdvisorCost(advisor, 0.5, 1); sessionManager.appendMessage({ role: "user", @@ -367,26 +411,22 @@ describe("AgentSession advisor toggle", () => { expect(session.formatAdvisorStatus()).toContain("$0.5000"); }); it("retains cumulative advisor cost when reloading the same session", async () => { - session.settings.setModelRole("advisor", `${model.provider}/${model.id}`); - session.toggleAdvisorEnabled(); - const advisor = session.getAdvisorAgent(); - if (!advisor) throw new Error("Expected advisor agent to exist"); + const advisor = enableAdvisor(); appendAdvisorCost(advisor, 0.5, 1); await session.reload(); expect(session.getAdvisorCost()).toBeCloseTo(0.5, 8); }); - it("keeps advisor cost when switching sessions fails after the reset", async () => { - session.settings.setModelRole("advisor", `${model.provider}/${model.id}`); - session.toggleAdvisorEnabled(); - const advisor = session.getAdvisorAgent(); - if (!advisor) throw new Error("Expected advisor agent to exist"); + it("restores advisor recording when a session switch fails before reset", async () => { + const advisor = enableAdvisor(); appendAdvisorCost(advisor, 0.5, 1); const previousSessionFile = sessionManager.getSessionFile(); const targetSessionFile = SessionManager.createEmptySessionFile(tempDir.path()); - const failure = new Error("switch failed after advisor reset"); - vi.spyOn(sessionManager, "getLastModelChangeRole").mockImplementation(() => { + const failure = new Error("switch failed before advisor reset"); + const setSessionFile = sessionManager.setSessionFile.bind(sessionManager); + vi.spyOn(sessionManager, "setSessionFile").mockImplementation(async file => { + await setSessionFile(file); throw failure; }); @@ -394,29 +434,250 @@ describe("AgentSession advisor toggle", () => { expect(sessionManager.getSessionFile()).toBe(previousSessionFile); expect(session.getAdvisorCost()).toBeCloseTo(0.5, 8); + if (!previousSessionFile) throw new Error("Expected the previous session to be persisted"); + appendAdvisorCost(advisor, 0.25, 2); + expect(session.getAdvisorCost()).toBeCloseTo(0.75, 8); + await session.dispose(); + expect((await loadAdvisorTranscriptCosts(previousSessionFile)).get("")).toBeCloseTo(0.75, 8); }); - it("clears advisor cost once a switch to a different session commits", async () => { - session.settings.setModelRole("advisor", `${model.provider}/${model.id}`); - session.toggleAdvisorEnabled(); - const advisor = session.getAdvisorAgent(); - if (!advisor) throw new Error("Expected advisor agent to exist"); + it("adopts only the target session's recorded advisor cost after a switch", async () => { + const advisor = enableAdvisor(); appendAdvisorCost(advisor, 0.5, 1); const targetSessionFile = SessionManager.createEmptySessionFile(tempDir.path()); + await writeAdvisorTranscript(targetSessionFile, "__advisor.jsonl", [0.25]); + const setSessionFile = sessionManager.setSessionFile.bind(sessionManager); + vi.spyOn(sessionManager, "setSessionFile").mockImplementation(async file => { + await setSessionFile(file); + // Reproduce an old advisor finishing after the target file became active. + appendAdvisorCost(advisor, 9, 2); + }); expect(await session.switchSession(targetSessionFile)).toBe(true); - - expect(session.getAdvisorCost()).toBe(0); + expect(session.getAdvisorCost()).toBeCloseTo(0.25, 8); + await session.dispose(); + expect((await loadAdvisorTranscriptCosts(targetSessionFile)).get("")).toBeCloseTo(0.25, 8); }); - it("clears cumulative advisor cost for a new session", async () => { - session.settings.setModelRole("advisor", `${model.provider}/${model.id}`); - session.toggleAdvisorEnabled(); - const advisor = session.getAdvisorAgent(); - if (!advisor) throw new Error("Expected advisor agent to exist"); + it("hydrates persisted advisor cost during SDK session startup", async () => { + const sessionFile = SessionManager.createEmptySessionFile(tempDir.path()); + await writeAdvisorTranscript(sessionFile, "__advisor.jsonl", [0.5]); + // A subagent advisor writes one directory deeper; its spend belongs to that + // subagent and must not inflate the resumed primary conversation. + await writeAdvisorTranscript( + path.join(sessionFile.slice(0, -".jsonl".length), "SubAgent.jsonl"), + "__advisor.jsonl", + [9], + ); + const settings = Settings.isolated({ + "async.enabled": false, + "advisor.enabled": true, + "compaction.enabled": false, + }); + settings.setModelRole("advisor", `${model.provider}/${model.id}`); + const result = await createAgentSession({ + cwd: tempDir.path(), + agentDir: tempDir.path(), + sessionManager: await SessionManager.open(sessionFile), + authStorage, + modelRegistry, + settings, + model, + disableExtensionDiscovery: true, + skills: [], + contextFiles: [], + workspaceTree: { + rootPath: tempDir.path(), + rendered: "", + truncated: false, + totalLines: 0, + agentsMdFiles: [], + }, + promptTemplates: [], + slashCommands: [], + enableMCP: false, + enableLsp: false, + }); + try { + expect(result.session.getAdvisorCost()).toBeCloseTo(0.5, 8); + } finally { + await result.session.dispose(); + } + }); + it("starts a new session with only post-transition advisor cost", async () => { + const advisor = enableAdvisor(); appendAdvisorCost(advisor, 0.5, 1); + const newSession = sessionManager.newSession.bind(sessionManager); + vi.spyOn(sessionManager, "newSession").mockImplementation(async options => { + const result = await newSession(options); + appendAdvisorCost(advisor, 9, 2); + return result; + }); await session.newSession(); + const replacementSessionFile = session.sessionFile; + if (!replacementSessionFile) throw new Error("Expected the replacement session to be persisted"); + appendAdvisorCost(advisor, 0.25, 3); + expect(session.getAdvisorCost()).toBeCloseTo(0.25, 8); + await session.dispose(); + expect((await loadAdvisorTranscriptCosts(replacementSessionFile)).get("")).toBeCloseTo(0.25, 8); + }); + it("records an advisor completion that races abort onto the previous session", async () => { + const advisor = enableAdvisor(); + const previousSessionFile = session.sessionFile; + if (!previousSessionFile) throw new Error("Expected the previous session to be persisted"); + appendAdvisorCost(advisor, 0.5, 1); + let injected = false; + const abort = advisor.abort.bind(advisor); + vi.spyOn(advisor, "abort").mockImplementation((reason?: unknown) => { + if (!injected) { + injected = true; + // Provider completes with billed usage while we stop the advisor for /new. + appendAdvisorCost(advisor, 0.41, 2); + } + return abort(reason); + }); + + await session.newSession(); + const replacementSessionFile = session.sessionFile; + if (!replacementSessionFile) throw new Error("Expected the replacement session to be persisted"); expect(session.getAdvisorCost()).toBe(0); + await session.dispose(); + + expect((await loadAdvisorTranscriptCosts(previousSessionFile)).get("")).toBeCloseTo(0.91, 8); + expect((await loadAdvisorTranscriptCosts(replacementSessionFile)).get("")).toBeUndefined(); + }); + it("restores advisor recording when a new session fails before commit", async () => { + const advisor = enableAdvisor(); + const previousSessionFile = session.sessionFile; + if (!previousSessionFile) throw new Error("Expected the previous session to be persisted"); + appendAdvisorCost(advisor, 0.5, 1); + const failure = new Error("new session failed"); + vi.spyOn(sessionManager, "newSession").mockRejectedValue(failure); + + await expect(session.newSession()).rejects.toThrow(failure); + + expect(advisor.state.messages).toHaveLength(1); + appendAdvisorCost(advisor, 0.25, 4); + expect(session.getAdvisorCost()).toBeCloseTo(0.75, 8); + await session.dispose(); + + expect((await loadAdvisorTranscriptCosts(previousSessionFile)).get("")).toBeCloseTo(0.75, 8); + }); + it("does not record a late advisor turn into a branched session", async () => { + const advisor = enableAdvisor(); + sessionManager.appendMessage({ role: "user", content: "ancestor", timestamp: 1 }); + sessionManager.appendMessage({ role: "user", content: "branch point", timestamp: 2 }); + const entryId = sessionManager.getLeafId(); + if (!entryId) throw new Error("Expected a branchable entry"); + const createBranchedSession = sessionManager.createBranchedSession.bind(sessionManager); + vi.spyOn(sessionManager, "createBranchedSession").mockImplementation(parentId => { + const result = createBranchedSession(parentId); + queueMicrotask(() => appendAdvisorCost(advisor, 9, 3)); + return result; + }); + + await expect(session.branch(entryId)).resolves.toMatchObject({ cancelled: false }); + const replacementSessionFile = session.sessionFile; + if (!replacementSessionFile) throw new Error("Expected the replacement session to be persisted"); + appendAdvisorCost(advisor, 0.25, 4); + expect(session.getAdvisorCost()).toBeCloseTo(0.25, 8); + await session.dispose(); + + expect((await loadAdvisorTranscriptCosts(replacementSessionFile)).get("")).toBeCloseTo(0.25, 8); + }); + it("keeps advisor cost across a fork of the same conversation", async () => { + const advisor = enableAdvisor(); + sessionManager.appendMessage({ role: "user", content: "keep me", timestamp: 1 }); + appendAdvisorCost(advisor, 0.5, 1); + const previousSessionFile = sessionManager.getSessionFile(); + const fork = sessionManager.fork.bind(sessionManager); + vi.spyOn(sessionManager, "fork").mockImplementation(async () => { + const result = await fork(); + // Reproduce the outgoing advisor finalizing after the fork selected its file. + appendAdvisorCost(advisor, 9, 2); + return result; + }); + + expect(await session.fork()).toBe(true); + + // A fork copies the entries and artifacts and keeps the messages, so the + // conversation continues under a new file and its spend continues with it. + expect(sessionManager.getSessionFile()).not.toBe(previousSessionFile); + expect(session.getAdvisorCost()).toBeCloseTo(0.5, 8); + expect(advisor.state.messages).toContainEqual(advisorMessage(0.5, 1)); + const forkedSessionFile = sessionManager.getSessionFile(); + if (!forkedSessionFile) throw new Error("Expected the forked session to be persisted"); + await session.dispose(); + expect((await loadAdvisorTranscriptCosts(forkedSessionFile)).get("")).toBeCloseTo(0.5, 8); + }); + it("restores advisor recording when a fork fails", async () => { + const advisor = enableAdvisor(); + appendAdvisorCost(advisor, 0.5, 1); + const previousSessionFile = sessionManager.getSessionFile(); + if (!previousSessionFile) throw new Error("Expected the previous session to be persisted"); + const failure = new Error("fork failed"); + vi.spyOn(sessionManager, "fork").mockRejectedValue(failure); + + await expect(session.fork()).rejects.toThrow(failure); + + expect(sessionManager.getSessionFile()).toBe(previousSessionFile); + expect(session.getAdvisorCost()).toBeCloseTo(0.5, 8); + appendAdvisorCost(advisor, 0.25, 2); + expect(session.getAdvisorCost()).toBeCloseTo(0.75, 8); + await session.dispose(); + expect((await loadAdvisorTranscriptCosts(previousSessionFile)).get("")).toBeCloseTo(0.75, 8); + }); + it("clears advisor cost when a handoff opens the replacement session", async () => { + vi.spyOn(compactionModule, "generateHandoffFromContext").mockResolvedValue("## Goal\nContinue from here"); + try { + const advisor = enableAdvisor(); + prepareHandoffConversation(advisor); + const previousSessionFile = session.sessionFile; + const newSession = sessionManager.newSession.bind(sessionManager); + vi.spyOn(sessionManager, "newSession").mockImplementation(async options => { + const result = await newSession(options); + // The outgoing advisor finalizes after the replacement file is selected. + appendAdvisorCost(advisor, 9, 3); + return result; + }); + + await session.handoff(); + + // The handoff hands the work over to a fresh conversation, so the spend of + // the one it summarizes must not follow it. + expect(session.sessionFile).not.toBe(previousSessionFile); + expect(session.getAdvisorCost()).toBe(0); + const replacementSessionFile = session.sessionFile; + if (!replacementSessionFile) throw new Error("Expected the replacement session to be persisted"); + appendAdvisorCost(advisor, 0.25, 4); + expect(session.getAdvisorCost()).toBeCloseTo(0.25, 8); + await session.dispose(); + expect((await loadAdvisorTranscriptCosts(replacementSessionFile)).get("")).toBeCloseTo(0.25, 8); + } finally { + vi.restoreAllMocks(); + } + }); + it("restores advisor recording when a handoff fails before replacing the session", async () => { + vi.spyOn(compactionModule, "generateHandoffFromContext").mockResolvedValue("## Goal\nContinue from here"); + try { + const advisor = enableAdvisor(); + prepareHandoffConversation(advisor); + const previousSessionFile = session.sessionFile; + if (!previousSessionFile) throw new Error("Expected the previous session to be persisted"); + const failure = new Error("replacement session failed"); + vi.spyOn(sessionManager, "newSession").mockRejectedValue(failure); + + await expect(session.handoff()).rejects.toThrow(failure); + + expect(session.sessionFile).toBe(previousSessionFile); + expect(session.getAdvisorCost()).toBeCloseTo(0.5, 8); + appendAdvisorCost(advisor, 0.25, 3); + expect(session.getAdvisorCost()).toBeCloseTo(0.75, 8); + await session.dispose(); + expect((await loadAdvisorTranscriptCosts(previousSessionFile)).get("")).toBeCloseTo(0.75, 8); + } finally { + vi.restoreAllMocks(); + } }); it("clears advisor cost when a branch skips conversation restore", async () => { const extensionRunner = { @@ -436,10 +697,7 @@ describe("AgentSession advisor toggle", () => { extensionRunner, }); try { - branchSession.settings.setModelRole("advisor", `${model.provider}/${model.id}`); - branchSession.toggleAdvisorEnabled(); - const advisor = branchSession.getAdvisorAgent(); - if (!advisor) throw new Error("Expected advisor agent to exist"); + const advisor = enableAdvisor(branchSession); const branchPoint = { role: "user" as const, content: "branch point", timestamp: 1 }; branchManager.appendMessage(branchPoint); const entryId = branchManager.getLeafId(); @@ -482,10 +740,7 @@ describe("AgentSession advisor toggle", () => { extensionRunner, }); try { - branchSession.settings.setModelRole("advisor", `${model.provider}/${model.id}`); - branchSession.toggleAdvisorEnabled(); - const advisor = branchSession.getAdvisorAgent(); - if (!advisor) throw new Error("Expected advisor agent to exist"); + const advisor = enableAdvisor(branchSession); branchManager.appendMessage({ role: "user", content: "branch point", timestamp: 1 }); const entryId = branchManager.getLeafId(); if (!entryId) throw new Error("Expected a branchable entry"); @@ -499,6 +754,8 @@ describe("AgentSession advisor toggle", () => { // so the abandoned conversation's spend must not be billed to the new one. expect(branchManager.getSessionFile()).not.toBe(previousSessionFile); expect(branchSession.getAdvisorCost()).toBe(0); + appendAdvisorCost(advisor, 0.25, 2); + expect(branchSession.getAdvisorCost()).toBeCloseTo(0.25, 8); } finally { await branchSession.dispose(); await branchDir.remove().catch(() => {}); diff --git a/packages/coding-agent/test/advisor/advisor.test.ts b/packages/coding-agent/test/advisor/advisor.test.ts index 7523e21e2..9ebf05c48 100644 --- a/packages/coding-agent/test/advisor/advisor.test.ts +++ b/packages/coding-agent/test/advisor/advisor.test.ts @@ -4000,6 +4000,128 @@ describe("advisor", () => { expect(promptInputs[1]).toContain("new-conversation"); expect(promptInputs[1]).not.toContain("old-conversation"); }); + + it("retries the interrupted batch after a session transition rolls back", async () => { + const promptInputs: string[] = []; + const firstPromptStarted = Promise.withResolvers(); + let rejectInFlight: ((reason?: unknown) => void) | undefined; + const agent: AdvisorAgent = { + prompt: input => { + promptInputs.push(input); + if (promptInputs.length > 1) return Promise.resolve(); + const gate = Promise.withResolvers(); + rejectInFlight = gate.reject; + firstPromptStarted.resolve(); + return gate.promise; + }, + abort: () => rejectInFlight?.(new Error("session transition")), + reset: () => {}, + state: { messages: [] }, + }; + const messages: AgentMessage[] = [{ role: "user", content: "keep me", timestamp: 1 } as AgentMessage]; + const runtime = new AdvisorRuntime(agent, { + snapshotMessages: () => messages, + enqueueAdvice: () => {}, + }); + + runtime.onTurnEnd(messages); + await firstPromptStarted.promise; + await runtime.pauseForSessionTransition(); + expect(promptInputs).toHaveLength(1); + + runtime.resumeAfterSessionTransition(); + await settleUntil(() => runtime.backlog === 0); + expect(promptInputs).toHaveLength(2); + expect(promptInputs[1]).toContain("keep me"); + }); + + it.each(["success", "error"] as const)( + "releases blocked %s hooks so reset can run replacement work", + async hookKind => { + const hookStarted = Promise.withResolvers(); + const releaseHook = Promise.withResolvers(); + const replacementPromptStarted = Promise.withResolvers(); + let promptCalls = 0; + let hookCalls = 0; + const blockHook = async () => { + if (++hookCalls !== 1) return; + hookStarted.resolve(); + await releaseHook.promise; + }; + const agent: AdvisorAgent = { + prompt: async () => { + promptCalls++; + if (promptCalls === 1 && hookKind === "error") throw new Error("provider failure"); + if (promptCalls === 2) replacementPromptStarted.resolve(); + }, + abort: () => {}, + reset: () => {}, + state: { messages: [] }, + }; + const runtime = new AdvisorRuntime(agent, { + snapshotMessages: () => [], + enqueueAdvice: () => {}, + ...(hookKind === "success" + ? { onTurnSuccess: blockHook } + : { + onTurnError: async () => { + await blockHook(); + return false; + }, + }), + }); + + runtime.onTurnEnd([{ role: "user", content: "old session", timestamp: 1 } as AgentMessage]); + await hookStarted.promise; + const pause = runtime.pauseForSessionTransition(); + const pausedQuickly = await Promise.race([pause.then(() => true), Bun.sleep(50).then(() => false)]); + runtime.reset(); + runtime.onTurnEnd([{ role: "user", content: "replacement session", timestamp: 2 } as AgentMessage]); + const replacementRan = await Promise.race([ + replacementPromptStarted.promise.then(() => true), + Bun.sleep(50).then(() => false), + ]); + releaseHook.resolve(); + await pause; + runtime.dispose(); + + expect(pausedQuickly).toBe(true); + expect(replacementRan).toBe(true); + }, + ); + it("aborts retry backoff before pausing for a session transition", async () => { + const recoveryStarted = Promise.withResolvers(); + const agent: AdvisorAgent = { + prompt: async () => { + throw new Error("provider failure"); + }, + abort: () => {}, + reset: () => {}, + state: { messages: [] }, + }; + const runtime = new AdvisorRuntime( + agent, + { + snapshotMessages: () => [], + enqueueAdvice: () => {}, + onTurnError: () => { + recoveryStarted.resolve(); + return false; + }, + }, + 250, + ); + + runtime.onTurnEnd([{ role: "user", content: "retry me", timestamp: 1 } as AgentMessage]); + await recoveryStarted.promise; + await Bun.sleep(0); + const pause = runtime.pauseForSessionTransition(); + const pausedQuickly = await Promise.race([pause.then(() => true), Bun.sleep(50).then(() => false)]); + if (!pausedQuickly) await pause; + runtime.dispose(); + + expect(pausedQuickly).toBe(true); + }); }); describe("AdvisorRuntime quota classification", () => { @@ -4106,6 +4228,11 @@ describe("advisor", () => { expect(promptInputs).toHaveLength(1); expect(promptInputs[0]).toContain("quota-turn"); + await runtime.pauseForSessionTransition(); + runtime.resumeAfterSessionTransition(); + await Promise.resolve(); + expect(promptInputs).toHaveLength(1); + // After reset() clears the quota pause, the next onTurnEnd drains the // retained batch — proving it was never lost. shouldFail = false; @@ -4268,15 +4395,22 @@ describe("advisor", () => { }; let quotaNotified = 0; let hookInvocations = 0; + const maintenanceSignals: AbortSignal[] = []; const { promise: hookEntered, resolve: allowHook } = Promise.withResolvers(); - const { promise: hookProceed, resolve: proceedHook } = Promise.withResolvers(); const host: AdvisorRuntimeHost = { snapshotMessages: () => [], enqueueAdvice: () => {}, - onTurnError: async () => { + maintainContext: async (_incomingTokens, signal) => { + maintenanceSignals.push(signal); + return false; + }, + onTurnError: async (_error, _failedMessages, signal) => { hookInvocations++; allowHook(); - await hookProceed; + const hookAborted = Promise.withResolvers(); + signal.addEventListener("abort", () => hookAborted.resolve(), { once: true }); + await hookAborted.promise; + signal.throwIfAborted(); return false; }, notifyQuotaExhausted: () => { @@ -4289,17 +4423,49 @@ describe("advisor", () => { await hookEntered; runtime.reset(); runtime.onTurnEnd([{ role: "user", content: "fresh-turn", timestamp: 2 } as AgentMessage]); - proceedHook(); await runtime.waitForCatchup(1000, 1); expect(hookInvocations).toBe(1); expect(promptInputs).toHaveLength(2); expect(promptInputs[0]).toContain("stale-turn"); expect(promptInputs[1]).toContain("fresh-turn"); + expect(maintenanceSignals).toHaveLength(2); + expect(maintenanceSignals[0]?.aborted).toBe(true); + expect(maintenanceSignals[1]?.aborted).toBe(false); expect(runtime.quotaExhausted).toBe(false); expect(runtime.backlog).toBe(0); expect(quotaNotified).toBe(0); }); + it("aborts the active recovery hook when disposed", async () => { + const hookEntered = Promise.withResolvers(); + let recoverySignal!: AbortSignal; + const agent: AdvisorAgent = { + prompt: async () => { + throw new Error("provider failure"); + }, + abort: () => {}, + reset: () => {}, + state: { messages: [] }, + }; + const runtime = new AdvisorRuntime(agent, { + snapshotMessages: () => [], + enqueueAdvice: () => {}, + onTurnError: async (_error, _failedMessages, signal) => { + recoverySignal = signal; + hookEntered.resolve(); + const hookAborted = Promise.withResolvers(); + signal.addEventListener("abort", () => hookAborted.resolve(), { once: true }); + await hookAborted.promise; + signal.throwIfAborted(); + }, + }); + + runtime.onTurnEnd([{ role: "user", content: "stale-turn", timestamp: 1 } as AgentMessage]); + await hookEntered.promise; + runtime.dispose(); + + expect(recoverySignal.aborted).toBe(true); + }); it("uses generic failure path when switched retry hits a non-quota error", async () => { const promptInputs: string[] = []; let callCount = 0; diff --git a/packages/coding-agent/test/advisor/transcript-recorder.test.ts b/packages/coding-agent/test/advisor/transcript-recorder.test.ts index 7beed7168..62eb69778 100644 --- a/packages/coding-agent/test/advisor/transcript-recorder.test.ts +++ b/packages/coding-agent/test/advisor/transcript-recorder.test.ts @@ -19,6 +19,8 @@ import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; import { ADVISOR_TRANSCRIPT_FILENAME, AdvisorTranscriptRecorder, + advisorTranscriptFilename, + loadAdvisorTranscriptCosts, } from "@oh-my-pi/pi-coding-agent/advisor/transcript-recorder"; import { removeWithRetries } from "@oh-my-pi/pi-utils"; @@ -54,7 +56,7 @@ async function readMessageEntries(file: string): Promise { return entries.filter(entry => entry.type === "message"); } -function assistantMessage(text: string, inputTokens: number): AgentMessage { +function assistantMessage(text: string, inputTokens: number, cost = 0): AgentMessage { const message = { role: "assistant" as const, content: [{ type: "text" as const, text }], @@ -67,7 +69,7 @@ function assistantMessage(text: string, inputTokens: number): AgentMessage { cacheRead: 0, cacheWrite: 0, totalTokens: inputTokens + 3, - cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + cost: { input: 0, output: cost, cacheRead: 0, cacheWrite: 0, total: cost }, }, stopReason: "stop" as const, timestamp: 1, @@ -161,4 +163,53 @@ describe("AdvisorTranscriptRecorder", () => { expect(second[0].message?.usage?.input).toBe(2); }); }); + + it("loads cumulative costs by advisor slug", async () => { + await withTempDir(async dir => { + const sessionFile = path.join(dir, "sess.jsonl"); + const primary = new AdvisorTranscriptRecorder( + () => sessionFile, + () => dir, + ); + const security = new AdvisorTranscriptRecorder( + () => sessionFile, + () => dir, + advisorTranscriptFilename("security"), + ); + primary.record(assistantMessage("primary", 1, 0.25)); + security.record(assistantMessage("first", 1, 0.25)); + security.record(assistantMessage("second", 1, 0.5)); + await Promise.all([primary.close(), security.close()]); + + expect(Object.fromEntries(await loadAdvisorTranscriptCosts(sessionFile))).toEqual({ + "": 0.25, + security: 0.75, + }); + }); + }); + + it("keeps valid costs when persisted entries are malformed", async () => { + await withTempDir(async dir => { + const sessionFile = path.join(dir, "sess.jsonl"); + const recorder = new AdvisorTranscriptRecorder( + () => sessionFile, + () => dir, + ); + recorder.record(assistantMessage("valid", 1, 0.25)); + await recorder.close(); + const transcript = path.join(dir, "sess", ADVISOR_TRANSCRIPT_FILENAME); + const lines = (await fs.readFile(transcript, "utf8")).trimEnd().split("\n"); + lines.splice( + -1, + 0, + JSON.stringify({ type: "message", message: { role: "assistant" } }), + "{ this is not valid json", + JSON.stringify({ type: "message" }), + "null", + ); + await fs.writeFile(transcript, `${lines.join("\n")}\n`); + + expect((await loadAdvisorTranscriptCosts(sessionFile)).get("")).toBe(0.25); + }); + }); }); diff --git a/packages/coding-agent/test/agent-session-advisor-suppression.test.ts b/packages/coding-agent/test/agent-session-advisor-suppression.test.ts index ad6540722..1de6e489f 100644 --- a/packages/coding-agent/test/agent-session-advisor-suppression.test.ts +++ b/packages/coding-agent/test/agent-session-advisor-suppression.test.ts @@ -192,6 +192,7 @@ describe("AgentSession advisor auto-resume suppression", () => { }, ], }, + { content: [], stopReason: "stop" }, ], }); const agent = new Agent({ diff --git a/packages/coding-agent/test/agent-session-btw-branch.test.ts b/packages/coding-agent/test/agent-session-btw-branch.test.ts index 0d6708591..052584b44 100644 --- a/packages/coding-agent/test/agent-session-btw-branch.test.ts +++ b/packages/coding-agent/test/agent-session-btw-branch.test.ts @@ -6,6 +6,7 @@ import { Agent } from "@oh-my-pi/pi-agent-core"; import type { AssistantMessage } from "@oh-my-pi/pi-ai"; import { createMockModel, type MockHandler } from "@oh-my-pi/pi-ai/providers/mock"; import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; +import { loadAdvisorTranscriptCosts } from "@oh-my-pi/pi-coding-agent/advisor/transcript-recorder"; import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import type { ExtensionRunner } from "@oh-my-pi/pi-coding-agent/extensibility/extensions"; @@ -135,6 +136,34 @@ describe("AgentSession.branchFromBtw", () => { if (promoted?.role !== "assistant") throw new Error("Expected promoted assistant message"); expectSanitizedBtwAssistant(promoted); }); + it("does not record a late advisor turn into a /btw branch", async () => { + const activeSession = await createSession(); + activeSession.settings.setModelRole("advisor", "anthropic/claude-sonnet-4-5"); + activeSession.toggleAdvisorEnabled(); + const advisor = activeSession.getAdvisorAgent(); + if (!advisor) throw new Error("Expected advisor agent to exist"); + activeSession.sessionManager.appendMessage({ role: "user", content: "seed", timestamp: Date.now() }); + await activeSession.sessionManager.flush(); + const createBranchedSession = activeSession.sessionManager.createBranchedSession.bind( + activeSession.sessionManager, + ); + vi.spyOn(activeSession.sessionManager, "createBranchedSession").mockImplementation(parentId => { + const result = createBranchedSession(parentId); + const lateMessage = createBtwAssistant(); + lateMessage.usage.cost.total = 9; + advisor.emitExternalEvent({ type: "message_end", message: lateMessage }); + return result; + }); + + const result = await activeSession.branchFromBtw("question", createBtwAssistant()); + expect(result.cancelled).toBe(false); + const replacementSessionFile = activeSession.sessionFile; + if (!replacementSessionFile) throw new Error("Expected the replacement session to be persisted"); + await activeSession.dispose(); + session = undefined; + + expect((await loadAdvisorTranscriptCosts(replacementSessionFile)).get("")).toBeUndefined(); + }); it("honors session_before_branch cancellation without creating a branch", async () => { const emit = vi.fn(async () => ({ cancel: true })); diff --git a/packages/coding-agent/test/agent-session-retry-fallback.test.ts b/packages/coding-agent/test/agent-session-retry-fallback.test.ts index 6c538623b..5ee906014 100644 --- a/packages/coding-agent/test/agent-session-retry-fallback.test.ts +++ b/packages/coding-agent/test/agent-session-retry-fallback.test.ts @@ -669,10 +669,16 @@ describe("AgentSession retry fallback", () => { ]); expect(advisorFailures).toEqual([]); + const getApiKey = vi.spyOn(modelRegistry, "getApiKey"); const afterCooldown = Date.now() + 2_000; vi.spyOn(Date, "now").mockReturnValue(afterCooldown); await session.prompt("Complete another primary turn after the advisor cooldown"); await session.waitForIdle(); + expect(getApiKey).toHaveBeenCalledWith( + expect.objectContaining({ provider: advisorPrimary.provider, id: advisorPrimary.id }), + expect.any(String), + { signal: expect.any(AbortSignal) }, + ); expect(requestedAdvisorModels).toEqual([advisorPrimarySelector, advisorFallbackSelector, advisorPrimarySelector]); expect(session.getAdvisorAgent()?.state.model).toMatchObject({ @@ -681,6 +687,77 @@ describe("AgentSession retry fallback", () => { }); }); + it("ignores late advisor fallback credentials after a session transition", async () => { + const mainModel = getBundledModel("openai", "gpt-4o-mini"); + const advisorPrimary = getBundledModel("anthropic", "claude-sonnet-4-5"); + const advisorFallback = getBundledModel("openai", "gpt-4o"); + if (!mainModel || !advisorPrimary || !advisorFallback) { + throw new Error("Expected bundled advisor fallback models to exist"); + } + + const mainMock = createMockModel({ responses: [{ content: ["Primary complete"] }] }); + const advisorMock = createMockModel({ + responses: [{ throw: "service unavailable: 503 overloaded" }], + }); + const agent = new Agent({ + getApiKey: model => `${model.provider}-test-key`, + initialState: { + model: mainModel, + systemPrompt: ["Test"], + tools: [], + messages: [], + }, + streamFn: mainMock.stream, + }); + const advisorPrimarySelector = `${advisorPrimary.provider}/${advisorPrimary.id}`; + const advisorFallbackSelector = `${advisorFallback.provider}/${advisorFallback.id}`; + const settings = Settings.isolated({ + "advisor.syncBacklog": "1", + "compaction.enabled": false, + "retry.baseDelayMs": 5, + "retry.fallbackChains": { + [advisorPrimarySelector]: [advisorFallbackSelector], + }, + }); + settings.setModelRole("advisor", advisorPrimarySelector); + session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(), + settings, + modelRegistry, + advisorTools: [], + advisorStreamFn: advisorMock.stream, + }); + expect(session.setAdvisorEnabled(true)).toBe(true); + + const credentialStarted = Promise.withResolvers(); + const releaseCredential = Promise.withResolvers(); + const credentialReturned = Promise.withResolvers(); + let credentialSignal: AbortSignal | undefined; + vi.spyOn(modelRegistry, "getApiKey").mockImplementation(async (model, _sessionId, options) => { + if (model.provider === advisorFallback.provider && model.id === advisorFallback.id) { + credentialSignal = options?.signal; + credentialStarted.resolve(); + await releaseCredential.promise; + credentialReturned.resolve(); + } + return `${model.provider}-test-key`; + }); + + await session.prompt("Trigger advisor fallback"); + await credentialStarted.promise; + await session.newSession(); + releaseCredential.resolve(); + await credentialReturned.promise; + await Bun.sleep(0); + + expect(credentialSignal?.aborted).toBe(true); + expect(session.getAdvisorAgent()?.state.model).toMatchObject({ + provider: advisorPrimary.provider, + id: advisorPrimary.id, + }); + }); + it("activates a model-keyed fallback chain without any role assignment", async () => { const primaryModel = getBundledModel("anthropic", "claude-sonnet-4-5"); const fallbackModel = getBundledModel("openai", "gpt-4o-mini"); diff --git a/packages/coding-agent/test/session-loader-stream.test.ts b/packages/coding-agent/test/session-loader-stream.test.ts index 6d31703cd..61484c5d8 100644 --- a/packages/coding-agent/test/session-loader-stream.test.ts +++ b/packages/coding-agent/test/session-loader-stream.test.ts @@ -3,7 +3,7 @@ import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; import type { FileEntry } from "@oh-my-pi/pi-coding-agent/session/session-entries"; -import { loadEntriesFromFileStream, parseSessionContent } from "@oh-my-pi/pi-coding-agent/session/session-loader"; +import * as sessionLoader from "@oh-my-pi/pi-coding-agent/session/session-loader"; import { serializeTitleSlot } from "@oh-my-pi/pi-coding-agent/session/session-title-slot"; // Parity contract for the ≥8MiB streaming loader (now Bun.JSONL-based): it must @@ -75,6 +75,48 @@ function messageTexts(entries: FileEntry[]): string[] { } describe("loadEntriesFromFileStream (Bun.JSONL parity)", () => { + it("visits entries incrementally while skipping malformed lines", async () => { + const slotLine = serializeTitleSlot({ title: "Visitor", source: "user", updatedAt: ISO }); + const content = [ + slotLine, + JSON.stringify(HEADER), + JSON.stringify(msg("m1", "s1", "first")), + "{ this is not valid json", + JSON.stringify(msg("m2", "m1", "second")), + ].join("\n"); + const file = await writeTemp(content); + const visited: FileEntry[] = []; + const titleSlot = await sessionLoader.visitEntriesFromFileStream(file, entry => visited.push(entry)); + + expect(titleSlot?.title).toBe("Visitor"); + expect(entryIds(visited)).toEqual(["s1", "m1", "m2"]); + }); + it("does not revisit entries before a malformed line spanning stream chunks", async () => { + const content = [ + JSON.stringify(HEADER), + JSON.stringify(msg("m1", "s1", "first")), + `{ this is not valid json ${"x".repeat(256 * 1024)}`, + JSON.stringify(msg("m2", "m1", "second")), + ].join("\n"); + const file = await writeTemp(content); + const visited: FileEntry[] = []; + + await sessionLoader.visitEntriesFromFileStream(file, entry => visited.push(entry)); + + expect(entryIds(visited)).toEqual(["s1", "m1", "m2"]); + }); + + it("propagates ENOENT errors thrown by the visitor", async () => { + const file = await writeTemp(`${JSON.stringify(HEADER)}\n`); + const failure = Object.assign(new Error("visitor failed"), { code: "ENOENT" }); + + await expect( + sessionLoader.visitEntriesFromFileStream(file, () => { + throw failure; + }), + ).rejects.toBe(failure); + }); + it("matches parseSessionContent on title slot + valid + malformed + blank lines", async () => { const slotLine = serializeTitleSlot({ title: "Hello world", source: "user", updatedAt: ISO }); // title slot | header | valid | blank | malformed | valid | malformed-no-newline-at-EOF @@ -89,8 +131,8 @@ describe("loadEntriesFromFileStream (Bun.JSONL parity)", () => { const content = lines.join("\n"); // no trailing newline on the last line const file = await writeTemp(content); - const stream = await loadEntriesFromFileStream(file); - const reference = parseSessionContent(content); + const stream = await sessionLoader.loadEntriesFromFileStream(file); + const reference = sessionLoader.parseSessionContent(content); // Parity: the stream path must agree with the common path exactly. expect(stream).toEqual(reference); @@ -111,8 +153,8 @@ describe("loadEntriesFromFileStream (Bun.JSONL parity)", () => { const content = lines.join("\n"); const file = await writeTemp(content); - const stream = await loadEntriesFromFileStream(file); - const reference = parseSessionContent(content); + const stream = await sessionLoader.loadEntriesFromFileStream(file); + const reference = sessionLoader.parseSessionContent(content); expect(stream).toEqual(reference); expect(stream.titleSlot).toBeUndefined(); @@ -132,8 +174,8 @@ describe("loadEntriesFromFileStream (Bun.JSONL parity)", () => { const content = lines.join("\n"); const file = await writeTemp(content); - const stream = await loadEntriesFromFileStream(file); - const reference = parseSessionContent(content); + const stream = await sessionLoader.loadEntriesFromFileStream(file); + const reference = sessionLoader.parseSessionContent(content); // Parity (a corrupted multibyte sequence would diverge here) ... expect(stream).toEqual(reference); @@ -146,7 +188,7 @@ describe("loadEntriesFromFileStream (Bun.JSONL parity)", () => { it("returns empty for a missing file (ENOENT)", async () => { const missing = path.join(os.tmpdir(), `does-not-exist-${Date.now()}.jsonl`); - const stream = await loadEntriesFromFileStream(missing); + const stream = await sessionLoader.loadEntriesFromFileStream(missing); expect(stream.entries).toEqual([]); expect(stream.titleSlot).toBeUndefined(); });