diff --git a/README.md b/README.md index d6d98df..29c961a 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,7 @@ This keeps Downshift simple: premium handles orientation, economy handles contin - Switches to a configured economy model after a token or percent threshold - Supports token thresholds, percent thresholds, or both - Remembers whether the session is premium, economy, paused, or mid-handoff +- Restores the premium/economy phase for the active session-tree branch after navigation and existing-session restoration - Pauses automatically after manual model changes - Optionally switches back to premium after successful compaction, even while immediate post-compaction usage telemetry is still unknown - Shows a compact status indicator in the UI @@ -60,6 +61,12 @@ It does not try to be clever. It is a deterministic context-cost governor. +## Session trees and forks 🌳 + +Downshift associates the premium or economy phase with the active session-tree branch. After `/tree` navigation, it restores that branch's phase and reconciles both the model and thinking level without evaluating the context threshold or creating a handoff. Reloading, resuming, forking, and cloning an existing Downshift session use the same restoration behavior. + +Session-level enablement and safety pauses remain active across branches. Navigating to a branch does not clear `/downshift off`, a manual-change pause, or a captured premium target. + ## Install 📦 ```bash diff --git a/src/downshift-core.test.ts b/src/downshift-core.test.ts index c73dced..fdf351f 100644 --- a/src/downshift-core.test.ts +++ b/src/downshift-core.test.ts @@ -9,6 +9,8 @@ import { handleManualThinkingLevelSelect, maybeDownshift, maybeUpshiftAfterCompaction, + reconcilePositionTarget, + restoreBranchPhaseFromEntries, restoreStateFromEntries, statusText, thresholdReached, @@ -55,6 +57,7 @@ type TestDeps = { reason: string, ) => Promise >; + getActiveTarget: Mock<() => ModelTarget | undefined>; updateStatus: Mock<() => void>; notify: Mock<(message: string, level?: string) => void>; }; @@ -65,6 +68,7 @@ function createDeps(config = baseConfig): TestDeps { saveState: vi.fn(), sendUserMessage: vi.fn(async () => undefined), switchToTarget: vi.fn(async () => true), + getActiveTarget: vi.fn(() => undefined), updateStatus: vi.fn(), notify: vi.fn(), }; @@ -148,6 +152,125 @@ describe("downshift core", () => { return { state: createState({ position: "economy" }) }; } + describe("restoreBranchPhaseFromEntries", () => { + it("defaults to premium and idle without a state entry", () => { + expect(restoreBranchPhaseFromEntries([])).toEqual({ + position: "premium", + handoff: "idle", + interrupted: false, + }); + }); + + it("ignores non-Downshift and invalid entries", () => { + expect( + restoreBranchPhaseFromEntries([ + { type: "message", data: {} }, + { type: "custom", customType: "downshift-state", data: "invalid" }, + ]), + ).toEqual({ + position: "premium", + handoff: "idle", + interrupted: false, + }); + }); + + it("uses the newest valid state on the branch", () => { + expect( + restoreBranchPhaseFromEntries([ + { + type: "custom", + customType: "downshift-state", + data: { position: "economy", handoff: "done" }, + }, + { + type: "custom", + customType: "downshift-state", + data: { position: "premium", handoff: "done" }, + }, + ]), + ).toEqual({ + position: "premium", + handoff: "idle", + interrupted: false, + }); + }); + + it("restores premium as idle", () => { + expect( + restoreBranchPhaseFromEntries([ + { + type: "custom", + customType: "downshift-state", + data: { position: "premium", handoff: "done" }, + }, + ]), + ).toEqual({ + position: "premium", + handoff: "idle", + interrupted: false, + }); + }); + + it("restores economy as done", () => { + expect( + restoreBranchPhaseFromEntries([ + { + type: "custom", + customType: "downshift-state", + data: { position: "economy", handoff: "idle" }, + }, + ]), + ).toEqual({ + position: "economy", + handoff: "done", + interrupted: false, + }); + }); + + it.each(["requested", "active"] as const)( + "marks a %s handoff as interrupted and normalizes it", + (handoff) => { + expect( + restoreBranchPhaseFromEntries([ + { + type: "custom", + customType: "downshift-state", + data: { position: "premium", handoff, continueAfterHandoff: true }, + }, + ]), + ).toEqual({ + position: "premium", + handoff: "idle", + interrupted: true, + }); + }, + ); + + it("does not restore continuation or session-wide fields", () => { + expect( + restoreBranchPhaseFromEntries([ + { + type: "custom", + customType: "downshift-state", + data: { + position: "economy", + handoff: "done", + continueAfterHandoff: true, + sessionMode: "off", + paused: true, + capturedPremium: premium, + lastError: "manual model change", + }, + }, + ]), + ).toEqual({ + position: "economy", + handoff: "done", + interrupted: false, + }); + }); + }); + it("formats premium status with remaining token and percent budget", () => { expect( statusText( @@ -562,6 +685,141 @@ describe("downshift core", () => { }); }); + describe("reconcilePositionTarget", () => { + it("selects the configured economy target for economy position", async () => { + const deps = createDeps(); + deps.getActiveTarget.mockReturnValue(premium); + const runtime = { state: createState({ position: "economy" }) }; + + await reconcilePositionTarget(deps, runtime, "restored branch"); + + expect(deps.switchToTarget).toHaveBeenCalledWith( + economy, + "economy", + "restored branch", + ); + }); + + it("selects an explicit premium target for premium position", async () => { + const config = { ...baseConfig, premiumSource: "explicit" as const, premium }; + const deps = createDeps(config); + deps.getActiveTarget.mockReturnValue(economy); + const runtime = { state: createState() }; + + await reconcilePositionTarget(deps, runtime, "restored session"); + + expect(deps.switchToTarget).toHaveBeenCalledWith( + premium, + "premium", + "restored session", + ); + }); + + it("selects the captured premium target for current premium source", async () => { + const deps = createDeps(); + deps.getActiveTarget.mockReturnValue(economy); + const runtime = { state: createState({ capturedPremium: premium }) }; + + await reconcilePositionTarget(deps, runtime, "restored session"); + + expect(deps.switchToTarget).toHaveBeenCalledWith( + premium, + "premium", + "restored session", + ); + }); + + it("does not switch or persist when the complete target matches", async () => { + const config = { + ...baseConfig, + premiumSource: "explicit" as const, + premium, + }; + const deps = createDeps(config); + deps.getActiveTarget.mockReturnValue(premium); + const runtime = { state: createState() }; + + await reconcilePositionTarget(deps, runtime, "restored branch"); + + expect(deps.switchToTarget).not.toHaveBeenCalled(); + expect(deps.saveState).not.toHaveBeenCalled(); + expect(deps.notify).not.toHaveBeenCalled(); + expect(deps.updateStatus).toHaveBeenCalledWith(config); + }); + + it.each([ + ["thinking level", premium, { ...premium, thinkingLevel: "high" }], + ["model", premium, { ...premium, model: "premium-new" }], + ] as const)("switches for a %s-only mismatch", async (_kind, expected, active) => { + const config = { + ...baseConfig, + premiumSource: "explicit" as const, + premium: expected, + }; + const deps = createDeps(config); + deps.getActiveTarget.mockReturnValue(active); + const runtime = { state: createState() }; + + await reconcilePositionTarget(deps, runtime, "restored branch"); + + expect(deps.switchToTarget).toHaveBeenCalledWith( + expected, + "premium", + "restored branch", + ); + expect(deps.sendUserMessage).not.toHaveBeenCalled(); + }); + + it.each([ + ["disabled config", { ...baseConfig, enabled: false }, createState()], + ["off session", baseConfig, createState({ sessionMode: "off" })], + ["paused runtime", baseConfig, createState({ paused: true })], + ["pending handoff", baseConfig, createState({ handoff: "requested" })], + ] as const)("does not switch for %s", async (_name, config, state) => { + const deps = createDeps(config); + deps.getActiveTarget.mockReturnValue(economy); + const runtime = { state }; + + await reconcilePositionTarget(deps, runtime, "restored branch"); + + expect(deps.switchToTarget).not.toHaveBeenCalled(); + expect(deps.sendUserMessage).not.toHaveBeenCalled(); + }); + + it("pauses when the premium target is missing", async () => { + const config = { + ...baseConfig, + premiumSource: "explicit" as const, + premium: undefined, + }; + const deps = createDeps(config); + const runtime = { state: createState() }; + + await reconcilePositionTarget(deps, runtime, "restored session"); + + expect(runtime.state.paused).toBe(true); + expect(runtime.state.lastError).toBe("premium target is unset"); + expect(deps.switchToTarget).not.toHaveBeenCalled(); + expect(deps.sendUserMessage).not.toHaveBeenCalled(); + }); + + it("returns the current runtime state after a successful switch", async () => { + const config = { + ...baseConfig, + premiumSource: "explicit" as const, + premium, + }; + const deps = createDeps(config); + deps.getActiveTarget.mockReturnValue(economy); + const runtime = { state: createState() }; + + const result = await reconcilePositionTarget(deps, runtime, "restored session"); + + expect(result).toBe(runtime.state); + expect(deps.updateStatus).toHaveBeenCalledWith(config); + }); + }); + it("restores interrupted handoff as paused instead of stranded", () => { for (const handoff of ["requested", "active"] as const) { const restored = restoreStateFromEntries( @@ -603,6 +861,21 @@ describe("downshift core", () => { expect(restored?.paused).toBe(true); }); + it("uses the current session ID instead of a persisted parent ID", () => { + const restored = restoreStateFromEntries( + [ + { + type: "custom", + customType: "downshift-state", + data: { ...createState(), sessionId: "parent-session" }, + }, + ], + "current-session", + ); + + expect(restored?.sessionId).toBe("current-session"); + }); + it("restores legacy session flags", () => { const restoredOverride = restoreStateFromEntries( [ diff --git a/src/downshift-core.ts b/src/downshift-core.ts index da86a5e..873a243 100644 --- a/src/downshift-core.ts +++ b/src/downshift-core.ts @@ -52,6 +52,12 @@ type PersistedStateEntry = { data?: unknown; }; +export type RestoredBranchPhase = { + position: Position; + handoff: "idle" | "done"; + interrupted: boolean; +}; + export type Runtime = { state: DownshiftState }; export type HandoffDelivery = "immediate" | "steer"; @@ -69,6 +75,7 @@ export type CoreDeps = { position: Position, reason: string, ) => Promise; + getActiveTarget: () => ModelTarget | undefined; updateStatus: (config?: DownshiftConfig) => void; notify: (message: string, level?: string) => void; }; @@ -129,7 +136,7 @@ function parseStateEntry(entry: PersistedStateEntry): StateEntry | undefined { function restoredState(data: StateEntry, sessionId: string): DownshiftState { const interrupted = data.handoff === "requested" || data.handoff === "active"; return { - sessionId: typeof data.sessionId === "string" ? data.sessionId : sessionId, + sessionId, sessionMode: restoredSessionMode(data), paused: interrupted || data.paused === true, position: data.position === "economy" ? "economy" : "premium", @@ -140,6 +147,24 @@ function restoredState(data: StateEntry, sessionId: string): DownshiftState { }; } +export function restoreBranchPhaseFromEntries( + entries: PersistedStateEntry[], +): RestoredBranchPhase { + for (let index = entries.length - 1; index >= 0; index -= 1) { + const data = parseStateEntry(entries[index]); + if (!data) continue; + + const position: Position = data.position === "economy" ? "economy" : "premium"; + const interrupted = data.handoff === "requested" || data.handoff === "active"; + return { + position, + handoff: position === "economy" ? "done" : "idle", + interrupted, + }; + } + return { position: "premium", handoff: "idle", interrupted: false }; +} + function restoredSessionMode(data: StateEntry): SessionMode { if (data.sessionMode === "on" || data.sessionMode === "off") return data.sessionMode; @@ -465,6 +490,50 @@ function resolvePremium( : state.capturedPremium; } +function targetsEqual( + left: ModelTarget | undefined, + right: ModelTarget, +): boolean { + return ( + left?.provider === right.provider && + left.model === right.model && + left.thinkingLevel === right.thinkingLevel + ); +} + +export async function reconcilePositionTarget( + deps: CoreDeps, + runtime: Runtime, + reason: string, +): Promise { + const config = await deps.readConfig(); + + if ( + !isDownshiftEnabled(config, runtime.state) || + runtime.state.paused || + hasPendingHandoff(runtime.state) + ) { + deps.updateStatus(config); + return runtime.state; + } + + const expectedTarget = + runtime.state.position === "economy" + ? config.economy + : resolvePremium(config, runtime.state); + if (!expectedTarget) { + pauseForMissingPremium(deps, runtime); + deps.updateStatus(config); + return runtime.state; + } + + if (!targetsEqual(deps.getActiveTarget(), expectedTarget)) { + await deps.switchToTarget(expectedTarget, runtime.state.position, reason); + } + deps.updateStatus(config); + return runtime.state; +} + function pauseForMissingPremium(deps: CoreDeps, runtime: Runtime): void { setState(deps, runtime, { paused: true, diff --git a/src/downshift.test.ts b/src/downshift.test.ts index 1e315b7..218339e 100644 --- a/src/downshift.test.ts +++ b/src/downshift.test.ts @@ -111,6 +111,8 @@ function createContext( model: models[0], thinkingLevel: "off", }, + sessionEntries: unknown[] = [], + branchEntries: unknown[] = [], ) { const status = vi.fn(); const select = vi.fn(); @@ -130,7 +132,8 @@ function createContext( getAvailable: () => models, }, sessionManager: { - getEntries: () => [], + getEntries: () => sessionEntries, + getBranch: () => branchEntries, getSessionId: () => "test-session", }, ui: { @@ -209,6 +212,535 @@ describe("downshift lifecycle adapter", () => { expect(handlers.has("thinking_level_select")).toBe(true); }); + it("registers a session tree handler", () => { + const fixture = createContext({ current: { tokens: 100, percent: 10 } }); + const { handlers } = createExtension(fixture.active); + + expect(handlers.has("session_tree")).toBe(true); + }); + + it("reconciles the active branch phase after tree navigation", async () => { + setConfig({ + premiumSource: "explicit", + premium: target("premium-new", "high"), + handoffBeforeDownshift: false, + }); + const sessionEntries = [ + { + type: "custom", + customType: "downshift-state", + data: { + sessionId: "parent-session", + sessionMode: "inherit", + paused: false, + position: "economy", + handoff: "done", + }, + }, + ]; + const branchEntries = [ + { + type: "custom", + customType: "downshift-state", + data: { position: "economy", handoff: "done" }, + }, + ]; + const fixture = createContext( + { current: { tokens: 100, percent: 10 } }, + { model: models[2], thinkingLevel: "off" }, + sessionEntries, + branchEntries, + ); + const extension = createExtension(fixture.active); + + await extension.handlers + .get("session_start") + ?.({ reason: "resume" }, fixture.context); + expect(extension.setModel).not.toHaveBeenCalled(); + + branchEntries.splice(0, branchEntries.length, { + type: "custom", + customType: "downshift-state", + data: { position: "premium", handoff: "done" }, + }); + await extension.handlers.get("session_tree")?.({}, fixture.context); + + expect(extension.setModel).toHaveBeenCalledOnce(); + expect(extension.setModel).toHaveBeenCalledWith(models[1]); + expect(extension.setThinkingLevel).toHaveBeenCalledWith("high"); + expect(extension.pi.sendUserMessage).not.toHaveBeenCalled(); + expect(fixture.notify).toHaveBeenCalledWith( + "downshift: restored branch to test/premium-new:high", + "info", + ); + }); + + it("reconciles from a premium branch to an economy branch", async () => { + setConfig({ + premiumSource: "explicit", + premium: target("premium-new", "high"), + handoffBeforeDownshift: false, + }); + const premiumEntry = { + type: "custom", + customType: "downshift-state", + data: { position: "premium", handoff: "idle" }, + }; + const economyEntry = { + type: "custom", + customType: "downshift-state", + data: { position: "economy", handoff: "done" }, + }; + const fixture = createContext( + { current: { tokens: 100, percent: 10 } }, + { model: models[1], thinkingLevel: "high" }, + [premiumEntry], + [premiumEntry], + ); + const extension = createExtension(fixture.active); + + await extension.handlers + .get("session_start") + ?.({ reason: "resume" }, fixture.context); + fixture.context.sessionManager.getBranch = () => [economyEntry] as any; + + await extension.handlers.get("session_tree")?.({}, fixture.context); + + expect(extension.setModel).toHaveBeenCalledWith(models[2]); + expect(extension.setThinkingLevel).toHaveBeenCalledWith("off"); + expect(extension.pi.sendUserMessage).not.toHaveBeenCalled(); + expect(latestState(extension.pi)).toMatchObject({ + position: "economy", + handoff: "done", + continueAfterHandoff: false, + }); + }); + + it("pauses when a source handoff is pending during tree navigation", async () => { + setConfig({ handoffBeforeDownshift: true }); + const premiumEntry = { + type: "custom", + customType: "downshift-state", + data: { position: "premium", handoff: "idle" }, + }; + const economyEntry = { + type: "custom", + customType: "downshift-state", + data: { position: "economy", handoff: "done" }, + }; + const fixture = createContext( + { current: { tokens: 100, percent: 10 } }, + { model: models[0], thinkingLevel: "off" }, + [premiumEntry], + [premiumEntry], + ); + const extension = createExtension(fixture.active); + + await extension.handlers + .get("session_start") + ?.({ reason: "resume" }, fixture.context); + await extension.commands.get("downshift")?.("now", fixture.commandContext); + extension.setModel.mockClear(); + extension.setThinkingLevel.mockClear(); + vi.mocked(extension.pi.sendUserMessage).mockClear(); + fixture.context.sessionManager.getBranch = () => [economyEntry] as any; + + await extension.handlers.get("session_tree")?.({}, fixture.context); + + expect(latestState(extension.pi)).toMatchObject({ + paused: true, + handoff: "done", + continueAfterHandoff: false, + lastError: "handoff interrupted by tree navigation", + }); + expect(extension.setModel).not.toHaveBeenCalled(); + expect(extension.setThinkingLevel).not.toHaveBeenCalled(); + expect(extension.pi.sendUserMessage).not.toHaveBeenCalled(); + }); + + it.each(["requested", "active"] as const)( + "pauses when the destination has an interrupted %s handoff", + async (handoff) => { + setConfig({ handoffBeforeDownshift: true }); + const premiumEntry = { + type: "custom", + customType: "downshift-state", + data: { position: "premium", handoff: "idle" }, + }; + const interruptedEntry = { + type: "custom", + customType: "downshift-state", + data: { position: "premium", handoff }, + }; + const fixture = createContext( + { current: { tokens: 100, percent: 10 } }, + { model: models[0], thinkingLevel: "off" }, + [premiumEntry], + [premiumEntry], + ); + const extension = createExtension(fixture.active); + + await extension.handlers + .get("session_start") + ?.({ reason: "resume" }, fixture.context); + extension.setModel.mockClear(); + extension.setThinkingLevel.mockClear(); + vi.mocked(extension.pi.sendUserMessage).mockClear(); + fixture.context.sessionManager.getBranch = () => [interruptedEntry] as any; + + await extension.handlers.get("session_tree")?.({}, fixture.context); + + expect(latestState(extension.pi)).toMatchObject({ + paused: true, + handoff: "idle", + continueAfterHandoff: false, + lastError: "handoff interrupted by tree navigation", + }); + expect(extension.setModel).not.toHaveBeenCalled(); + expect(extension.setThinkingLevel).not.toHaveBeenCalled(); + expect(extension.pi.sendUserMessage).not.toHaveBeenCalled(); + }, + ); + + it("defaults an untracked branch to premium after tree navigation", async () => { + setConfig({ + premiumSource: "explicit", + premium: target("premium-new", "high"), + handoffBeforeDownshift: false, + }); + const sessionEntries = [ + { + type: "custom", + customType: "downshift-state", + data: { position: "economy", handoff: "done" }, + }, + ]; + const fixture = createContext( + { current: { tokens: 100, percent: 10 } }, + { model: models[2], thinkingLevel: "off" }, + sessionEntries, + [ + { + type: "custom", + customType: "downshift-state", + data: { position: "economy", handoff: "done" }, + }, + ], + ); + const extension = createExtension(fixture.active); + + await extension.handlers + .get("session_start") + ?.({ reason: "resume" }, fixture.context); + fixture.active.model = models[2]; + fixture.active.thinkingLevel = "off"; + extension.setModel.mockClear(); + extension.setThinkingLevel.mockClear(); + + const branch = fixture.context.sessionManager.getBranch() as unknown[]; + branch.splice(0, branch.length); + + await extension.handlers.get("session_tree")?.({}, fixture.context); + + expect(extension.setModel).toHaveBeenCalledWith(models[1]); + expect(extension.setThinkingLevel).toHaveBeenCalledWith("high"); + expect(extension.pi.sendUserMessage).not.toHaveBeenCalled(); + }); + + it.each(["reload", "resume", "fork", "clone"])( + "reconciles an existing branch during %s session start", + async (reason) => { + setConfig({ + premiumSource: "explicit", + premium: target("premium-new", "high"), + handoffBeforeDownshift: false, + }); + const stateEntry = { + type: "custom", + customType: "downshift-state", + data: { sessionId: "parent-session", position: "premium", handoff: "idle" }, + }; + const fixture = createContext( + { current: { tokens: 100, percent: 90 } }, + { model: models[2], thinkingLevel: "off" }, + [stateEntry], + [stateEntry], + ); + const extension = createExtension(fixture.active); + + await extension.handlers.get("session_start")?.( + { reason } as any, + fixture.context, + ); + + expect(extension.setModel).toHaveBeenCalledWith(models[1]); + expect(extension.setThinkingLevel).toHaveBeenCalledWith("high"); + expect(extension.pi.sendUserMessage).not.toHaveBeenCalled(); + expect(fixture.notify).not.toHaveBeenCalledWith( + expect.stringContaining("preparing handoff"), + expect.anything(), + ); + }, + ); + + it("keeps fresh session startup behavior when no state exists", async () => { + setConfig({ + premiumSource: "explicit", + premium: target("premium-new", "high"), + startOnPremium: true, + }); + const fixture = createContext( + { current: { tokens: 100, percent: 10 } }, + { model: models[0], thinkingLevel: "off" }, + ); + const extension = createExtension(fixture.active); + + await extension.handlers + .get("session_start") + ?.({ reason: "new" }, fixture.context); + + expect(extension.setModel).toHaveBeenCalledWith(models[1]); + expect(extension.setThinkingLevel).toHaveBeenCalledWith("high"); + }); + + it.each([ + ["disabled", { enabled: false }, { sessionMode: "inherit" }], + ["off", {}, { sessionMode: "off" }], + ["paused", {}, { paused: true }], + ] as const)("does not switch a %s destination branch", async (_name, configPatch, statePatch) => { + setConfig({ + ...configPatch, + premiumSource: "explicit", + premium: target("premium-new", "high"), + handoffBeforeDownshift: false, + }); + const stateEntry = { + type: "custom", + customType: "downshift-state", + data: { position: "economy", handoff: "done", ...statePatch }, + }; + const destinationEntry = { + type: "custom", + customType: "downshift-state", + data: { position: "premium", handoff: "idle" }, + }; + const fixture = createContext( + { current: { tokens: 100, percent: 10 } }, + { model: models[2], thinkingLevel: "off" }, + [stateEntry], + [destinationEntry], + ); + const extension = createExtension(fixture.active); + + await extension.handlers + .get("session_start") + ?.({ reason: "resume" }, fixture.context); + extension.setModel.mockClear(); + extension.setThinkingLevel.mockClear(); + vi.mocked(extension.pi.sendUserMessage).mockClear(); + + await extension.handlers.get("session_tree")?.({}, fixture.context); + + expect(extension.setModel).not.toHaveBeenCalled(); + expect(extension.setThinkingLevel).not.toHaveBeenCalled(); + expect(extension.pi.sendUserMessage).not.toHaveBeenCalled(); + if (_name === "paused") { + expect(fixture.status).toHaveBeenLastCalledWith("downshift", "⇣ paused"); + } + }); + + it("does not generate a handoff while restoring a branch", async () => { + setConfig({ + premiumSource: "explicit", + premium: target("premium-new", "high"), + handoffBeforeDownshift: true, + }); + const economyEntry = { + type: "custom", + customType: "downshift-state", + data: { position: "economy", handoff: "done" }, + }; + const premiumEntry = { + type: "custom", + customType: "downshift-state", + data: { position: "premium", handoff: "idle" }, + }; + const fixture = createContext( + { current: { tokens: 100, percent: 90 } }, + { model: models[2], thinkingLevel: "off" }, + [economyEntry], + [economyEntry], + ); + const extension = createExtension(fixture.active); + + await extension.handlers + .get("session_start") + ?.({ reason: "resume" }, fixture.context); + fixture.context.sessionManager.getBranch = () => [premiumEntry] as any; + + await extension.handlers.get("session_tree")?.({}, fixture.context); + + expect(extension.setModel).toHaveBeenCalledWith(models[1]); + expect(extension.pi.sendUserMessage).not.toHaveBeenCalled(); + }); + + it("preserves internal target changes during branch reconciliation", async () => { + setConfig({ + premiumSource: "explicit", + premium: target("premium-new", "high"), + handoffBeforeDownshift: false, + }); + const economyEntry = { + type: "custom", + customType: "downshift-state", + data: { position: "economy", handoff: "done" }, + }; + const premiumEntry = { + type: "custom", + customType: "downshift-state", + data: { position: "premium", handoff: "idle" }, + }; + const fixture = createContext( + { current: { tokens: 100, percent: 90 } }, + { model: models[2], thinkingLevel: "off" }, + [economyEntry], + [economyEntry], + ); + const extension = createExtension(fixture.active); + extension.setModel.mockImplementationOnce(async (model: Model) => { + const previousModel = fixture.active.model; + fixture.active.model = model; + await extension.handlers.get("model_select")?.( + { model, previousModel, source: "set" }, + fixture.context, + ); + return true; + }); + extension.setThinkingLevel.mockImplementationOnce((level: ThinkingLevel) => { + const previousLevel = fixture.active.thinkingLevel; + fixture.active.thinkingLevel = level; + void extension.handlers.get("thinking_level_select")?.( + { level, previousLevel }, + fixture.context, + ); + }); + + await extension.handlers + .get("session_start") + ?.({ reason: "resume" }, fixture.context); + fixture.context.sessionManager.getBranch = () => [premiumEntry] as any; + + await extension.handlers.get("session_tree")?.({}, fixture.context); + + expect(latestState(extension.pi)).toMatchObject({ + paused: false, + position: "premium", + }); + expect(fixture.notify).not.toHaveBeenCalledWith( + "downshift paused: manual model change", + "error", + ); + }); + + it("continues to downshift after restoring a premium branch", async () => { + setConfig({ + premiumSource: "explicit", + premium: target("premium-new", "high"), + handoffBeforeDownshift: false, + }); + const premiumEntry = { + type: "custom", + customType: "downshift-state", + data: { position: "premium", handoff: "idle" }, + }; + const fixture = createContext( + { current: { tokens: 100, percent: 90 } }, + { model: models[1], thinkingLevel: "high" }, + [premiumEntry], + [premiumEntry], + ); + const extension = createExtension(fixture.active); + + await extension.handlers + .get("session_start") + ?.({ reason: "resume" }, fixture.context); + await extension.handlers.get("context")?.({}, fixture.context); + + expect(extension.setModel).toHaveBeenCalledWith(models[2]); + expect(latestState(extension.pi)).toMatchObject({ + position: "economy", + handoff: "done", + }); + }); + + it("continues to upshift after restoring an economy branch and compacting", async () => { + setConfig({ + premiumSource: "explicit", + premium: target("premium-new", "high"), + upshiftAfterCompaction: true, + }); + const economyEntry = { + type: "custom", + customType: "downshift-state", + data: { position: "economy", handoff: "done" }, + }; + const fixture = createContext( + { current: { tokens: 100, percent: 10 } }, + { model: models[2], thinkingLevel: "off" }, + [economyEntry], + [economyEntry], + ); + const extension = createExtension(fixture.active); + + await extension.handlers + .get("session_start") + ?.({ reason: "resume" }, fixture.context); + await extension.handlers.get("session_compact")?.( + { compactionEntry: {} }, + fixture.context, + ); + + expect(extension.setModel).toHaveBeenCalledWith(models[1]); + expect(extension.setThinkingLevel).toHaveBeenCalledWith("high"); + expect(latestState(extension.pi)).toMatchObject({ + position: "premium", + handoff: "idle", + }); + }); + + it("does nothing when branch reconciliation already matches the target", async () => { + setConfig({ + premiumSource: "explicit", + premium: target("premium-new", "high"), + handoffBeforeDownshift: false, + }); + const premiumEntry = { + type: "custom", + customType: "downshift-state", + data: { position: "premium", handoff: "idle" }, + }; + const fixture = createContext( + { current: { tokens: 100, percent: 10 } }, + { model: models[1], thinkingLevel: "high" }, + [premiumEntry], + [premiumEntry], + ); + const extension = createExtension(fixture.active); + + await extension.handlers + .get("session_start") + ?.({ reason: "resume" }, fixture.context); + vi.mocked(extension.pi.appendEntry).mockClear(); + fixture.notify.mockClear(); + + await extension.handlers.get("session_tree")?.({}, fixture.context); + + expect(extension.setModel).not.toHaveBeenCalled(); + expect(extension.setThinkingLevel).not.toHaveBeenCalled(); + expect(extension.pi.appendEntry).not.toHaveBeenCalled(); + expect(fixture.notify).not.toHaveBeenCalled(); + }); + it("opens configuration for the bare command", async () => { const usage = { current: { tokens: 100, percent: 10 } }; const fixture = createContext(usage); diff --git a/src/downshift.ts b/src/downshift.ts index 946cd9e..0656ab3 100644 --- a/src/downshift.ts +++ b/src/downshift.ts @@ -22,6 +22,8 @@ import { maybeUpshiftAfterCompaction, parseTarget, formatCompactNumber, + reconcilePositionTarget, + restoreBranchPhaseFromEntries, restoreStateFromEntries, statusText, thresholdReached, @@ -29,6 +31,7 @@ import { type DownshiftState, type ModelTarget, type Position, + type RestoredBranchPhase, type Threshold, } from "./downshift-core"; @@ -204,6 +207,17 @@ function restoreState(ctx: ExtensionContext): boolean { return true; } +function restoreActiveBranchPhase(ctx: ExtensionContext): RestoredBranchPhase { + const phase = restoreBranchPhaseFromEntries(ctx.sessionManager.getBranch()); + runtime.state = { + ...runtime.state, + position: phase.position, + handoff: phase.handoff, + continueAfterHandoff: false, + }; + return phase; +} + function pause( pi: ExtensionAPI, ctx: ExtensionContext, @@ -270,6 +284,7 @@ function coreDeps(pi: ExtensionAPI, ctx: ExtensionContext) { ) => pi.sendUserMessage(prompt, options), switchToTarget: (target: ModelTarget, position: Position, reason: string) => switchToTarget(pi, ctx, target, position, reason), + getActiveTarget: () => getActiveTarget(pi, ctx), updateStatus: (config?: DownshiftConfig) => updateStatus(ctx, config), notify: (message: string, level?: string) => ctx.ui.notify(message, level as any), @@ -902,10 +917,41 @@ async function handleSessionStart( ): Promise { const hadState = restoreState(ctx); const config = await readConfig(); - if (!hadState) await initializeSessionState(pi, event, ctx, config); + if (!hadState) { + await initializeSessionState(pi, event, ctx, config); + } else { + const phase = restoreActiveBranchPhase(ctx); + if (phase.interrupted && !runtime.state.paused) { + pause(pi, ctx, "handoff interrupted by session restore", { + position: phase.position, + handoff: phase.handoff, + continueAfterHandoff: false, + }); + } else { + await reconcilePositionTarget(coreDeps(pi, ctx), runtime, "restored session"); + } + } updateStatus(ctx, config); } +async function handleSessionTree( + pi: ExtensionAPI, + ctx: ExtensionContext, +): Promise { + const sourceHandoffPending = + runtime.state.handoff === "requested" || runtime.state.handoff === "active"; + const phase = restoreActiveBranchPhase(ctx); + if (sourceHandoffPending || phase.interrupted) { + pause(pi, ctx, "handoff interrupted by tree navigation", { + position: phase.position, + handoff: phase.handoff, + continueAfterHandoff: false, + }); + return; + } + await reconcilePositionTarget(coreDeps(pi, ctx), runtime, "restored branch"); +} + async function initializeSessionState( pi: ExtensionAPI, event: { reason?: string }, @@ -958,6 +1004,10 @@ export default function downshift(pi: ExtensionAPI): void { await handleSessionStart(pi, event, ctx); }); + pi.on("session_tree", async (_event, ctx) => { + await handleSessionTree(pi, ctx); + }); + pi.on("context", async (_event, ctx) => { await maybeDownshift( coreDeps(pi, ctx),