diff --git a/src/agent/routing/compute-step-complexity.test.ts b/src/agent/routing/compute-step-complexity.test.ts new file mode 100644 index 00000000..104fa6cf --- /dev/null +++ b/src/agent/routing/compute-step-complexity.test.ts @@ -0,0 +1,91 @@ +import { describe, expect, it } from "vitest"; + +import { + computeStepComplexity, + type StepComplexitySignals, +} from "./compute-step-complexity.js"; + +const base: StepComplexitySignals = { + promptTokens: 0, + stablePrefixTokens: 0, + stepIndex: 0, + maxSteps: 25, + conversationMaxTokens: 32_000, + hasTransientNotice: false, +}; + +const at = (over: Partial): number => + computeStepComplexity({ ...base, ...over }); + +describe("computeStepComplexity", () => { + it("scores a fresh, empty step at zero", () => { + expect(at({})).toBe(0); + }); + + it("saturates at 100 when every signal is maxed", () => { + expect( + at({ + promptTokens: 64_000, + stablePrefixTokens: 0, + stepIndex: 25, + hasTransientNotice: true, + }), + ).toBe(100); + }); + + it("always returns an integer inside 0-100", () => { + const samples = [ + at({ promptTokens: 7_777, stablePrefixTokens: 1_234, stepIndex: 3 }), + at({ promptTokens: 31_999, stablePrefixTokens: 12_001, stepIndex: 7 }), + at({ promptTokens: 1, stablePrefixTokens: 0, stepIndex: 1 }), + ]; + for (const score of samples) { + expect(Number.isInteger(score)).toBe(true); + expect(score).toBeGreaterThanOrEqual(0); + expect(score).toBeLessThanOrEqual(100); + } + }); + + it("is monotonic in context pressure", () => { + const low = at({ promptTokens: 4_000, stablePrefixTokens: 4_000 }); + const high = at({ promptTokens: 16_000, stablePrefixTokens: 16_000 }); + expect(high).toBeGreaterThan(low); + }); + + it("is monotonic in turn depth", () => { + expect(at({ stepIndex: 12 })).toBeGreaterThan(at({ stepIndex: 2 })); + }); + + it("is monotonic in tail growth at a fixed prompt size", () => { + const mostlyStable = at({ promptTokens: 20_000, stablePrefixTokens: 19_000 }); + const mostlyTail = at({ promptTokens: 20_000, stablePrefixTokens: 1_000 }); + expect(mostlyTail).toBeGreaterThan(mostlyStable); + }); + + it("adds exactly the transient-notice weight", () => { + const quiet = at({ promptTokens: 8_000, stablePrefixTokens: 6_000 }); + const noisy = at({ + promptTokens: 8_000, + stablePrefixTokens: 6_000, + hasTransientNotice: true, + }); + expect(noisy - quiet).toBe(20); + }); + + it("treats a tail larger than the prompt as zero, never negative", () => { + expect(at({ promptTokens: 100, stablePrefixTokens: 5_000 })).toBe(0); + }); + + it("survives zero and non-finite budgets without producing NaN", () => { + for (const bad of [0, -1, Number.NaN, Number.POSITIVE_INFINITY]) { + const score = at({ + promptTokens: 10_000, + conversationMaxTokens: bad, + maxSteps: bad, + stepIndex: 5, + }); + expect(Number.isInteger(score)).toBe(true); + expect(score).toBeGreaterThanOrEqual(0); + } + }); +}); diff --git a/src/agent/routing/compute-step-complexity.ts b/src/agent/routing/compute-step-complexity.ts new file mode 100644 index 00000000..c88b0072 --- /dev/null +++ b/src/agent/routing/compute-step-complexity.ts @@ -0,0 +1,94 @@ +/** + * Signals available at routing time — i.e. after `buildPrompt` but + * BEFORE `slotManager.acquire`, because the slot depends on which + * provider we route to. + * + * That ordering is why `cacheReused` is deliberately absent: it is + * produced by `slotManager.acquire`, so feeding it back into the + * routing decision would be circular. Do not add it. + */ +export interface StepComplexitySignals { + /** `prompt.tokens.total` for the step about to run. */ + promptTokens: number; + /** `prompt.tokens.stablePrefix` — the KV-stable head of the prompt. */ + stablePrefixTokens: number; + /** 0-based index of this step inside the current turn. */ + stepIndex: number; + /** `config.agent.maxSteps` — the turn's step budget. */ + maxSteps: number; + /** `config.agent.conversationMaxTokens` — the conversation budget. */ + conversationMaxTokens: number; + /** + * Whether a one-shot notice is being rendered into this step's prompt + * (loop detector fired, or a tool batch was trimmed). The model just + * did something wrong, so the step deserves the stronger model. + */ + hasTransientNotice: boolean; +} + +/** + * Weights sum to 100 so the score is directly comparable to the + * operator's `cloudShare` dial without any rescaling. + */ +const WEIGHT_CONTEXT_PRESSURE = 40; +const WEIGHT_TURN_DEPTH = 25; +const WEIGHT_TRANSIENT_NOTICE = 20; +const WEIGHT_TAIL_GROWTH = 15; + +/** + * The tail is judged against half the conversation budget: a turn whose + * accumulated tool output has eaten that much is already synthesis-shaped, + * and waiting for the full budget would only escalate on the very last + * step or two. + */ +const TAIL_BUDGET_FRACTION = 2; + +function clamp01(value: number): number { + if (!Number.isFinite(value) || value <= 0) return 0; + return value >= 1 ? 1 : value; +} + +function ratio(numerator: number, denominator: number): number { + if (!Number.isFinite(denominator) || denominator <= 0) return 0; + return clamp01(numerator / denominator); +} + +/** + * Score one step's difficulty on a bounded 0-100 scale. + * + * Deliberately a *heuristic over cheap signals*, not a model call: it + * runs before every inference in fusion mode, so it has to be free and + * deterministic. The four terms, in weight order: + * + * 1. **Context pressure** (40) — how full the context is. This is the + * dominant term on purpose. It is also how a final synthesis step + * ends up on the cloud without the loop being able to know a step is + * final: by the time the model is ready to answer, it is carrying the + * whole turn's context. + * 2. **Turn depth** (25) — later steps in a long turn are the ones that + * have to hold more state together. + * 3. **Transient notice** (20) — a binary "the model just misbehaved" + * signal from the loop detector / batch trimmer. + * 4. **Tail growth** (15) — how much of the prompt is accumulated tool + * output rather than the stable prefix, i.e. how much raw material + * this step has to reconcile. + */ +export function computeStepComplexity( + signals: StepComplexitySignals, +): number { + const tailTokens = Math.max( + 0, + signals.promptTokens - signals.stablePrefixTokens, + ); + const score = + WEIGHT_CONTEXT_PRESSURE * + ratio(signals.promptTokens, signals.conversationMaxTokens) + + WEIGHT_TURN_DEPTH * ratio(signals.stepIndex, signals.maxSteps) + + WEIGHT_TRANSIENT_NOTICE * (signals.hasTransientNotice ? 1 : 0) + + WEIGHT_TAIL_GROWTH * + ratio( + tailTokens, + signals.conversationMaxTokens / TAIL_BUDGET_FRACTION, + ); + return Math.round(score); +} diff --git a/src/agent/routing/decide-routing-role.test.ts b/src/agent/routing/decide-routing-role.test.ts new file mode 100644 index 00000000..a170ff50 --- /dev/null +++ b/src/agent/routing/decide-routing-role.test.ts @@ -0,0 +1,86 @@ +import { describe, expect, it } from "vitest"; + +import { + decideRoutingRole, + ROUTING_HYSTERESIS, +} from "./decide-routing-role.js"; + +describe("decideRoutingRole", () => { + it("keeps everything local at cloudShare 0, even a maximal score", () => { + expect( + decideRoutingRole({ score: 100, cloudShare: 0, stepIndex: 0 }), + ).toBe("executor"); + }); + + it("sends everything to the cloud at cloudShare 100, even a zero score", () => { + expect( + decideRoutingRole({ score: 0, cloudShare: 100, stepIndex: 9 }), + ).toBe("orchestrator"); + }); + + it("always orchestrates step 0 when the cloud leg is in play", () => { + expect( + decideRoutingRole({ score: 0, cloudShare: 1, stepIndex: 0 }), + ).toBe("orchestrator"); + }); + + it("routes on the cutoff at 100 - cloudShare with no prior role", () => { + // cloudShare 40 ⇒ cutoff 60. + expect( + decideRoutingRole({ score: 60, cloudShare: 40, stepIndex: 1 }), + ).toBe("orchestrator"); + expect( + decideRoutingRole({ score: 59, cloudShare: 40, stepIndex: 1 }), + ).toBe("executor"); + }); + + it("makes it harder to leave the local leg", () => { + // cutoff 60, previously executor ⇒ effective bar 70. + const args = { cloudShare: 40, stepIndex: 1, previousRole: "executor" } as const; + expect(decideRoutingRole({ ...args, score: 69 })).toBe("executor"); + expect(decideRoutingRole({ ...args, score: 70 })).toBe("orchestrator"); + }); + + it("makes it harder to leave the cloud leg", () => { + // cutoff 60, previously orchestrator ⇒ effective bar 50. + const args = { + cloudShare: 40, + stepIndex: 1, + previousRole: "orchestrator", + } as const; + expect(decideRoutingRole({ ...args, score: 50 })).toBe("orchestrator"); + expect(decideRoutingRole({ ...args, score: 49 })).toBe("executor"); + }); + + it("applies the hysteresis symmetrically", () => { + expect(ROUTING_HYSTERESIS).toBe(10); + const score = 55; + expect( + decideRoutingRole({ + score, + cloudShare: 40, + stepIndex: 1, + previousRole: "executor", + }), + ).toBe("executor"); + expect( + decideRoutingRole({ + score, + cloudShare: 40, + stepIndex: 1, + previousRole: "orchestrator", + }), + ).toBe("orchestrator"); + }); + + it("treats a null previous role like no prior state", () => { + expect( + decideRoutingRole({ + score: 60, + cloudShare: 40, + stepIndex: 1, + previousRole: null, + }), + ).toBe("orchestrator"); + }); +}); diff --git a/src/agent/routing/decide-routing-role.ts b/src/agent/routing/decide-routing-role.ts new file mode 100644 index 00000000..74e5fce7 --- /dev/null +++ b/src/agent/routing/decide-routing-role.ts @@ -0,0 +1,62 @@ +/** + * Which leg of a fusion pair serves one inference. + * + * `orchestrator` is the cloud provider (plans, reconciles, synthesises); + * `executor` is the local provider (mechanical continuation steps). + */ +export type RoutingRole = "orchestrator" | "executor"; + +/** + * Score margin applied against the direction of travel so a step near + * the cutoff does not flip the provider back and forth. + * + * This is load-bearing, not cosmetic. llama-server reuses its KV cache + * by longest common prefix, so every return to the local leg after a + * cloud step has to reprocess the tail that grew in between. Hysteresis + * produces RUNS of consecutive local steps, which is what makes the + * local cache pay for itself. + */ +export const ROUTING_HYSTERESIS = 10; + +export interface RoutingDecisionArgs { + /** 0-100 from `computeStepComplexity`. */ + score: number; + /** 0-100 operator dial from `llm.runMode.fusion.cloudShare`. */ + cloudShare: number; + /** 0-based step index inside the turn. */ + stepIndex: number; + /** Role the previous step of this session resolved to, if any. */ + previousRole?: RoutingRole | null; +} + +/** + * Map a complexity score onto a fusion leg. + * + * The dial sets a cutoff at `100 - cloudShare`: a bigger share means a + * lower bar for reaching the cloud. It is a DIAL, NOT A QUOTA — it does + * not promise that N% of steps go to the cloud, and it must not be + * turned into a running-counter scheduler, which would necessarily send + * some trivial steps to the cloud and keep some hard ones local. + * + * Two rules override the score: + * - `cloudShare` 0 / 100 short-circuit to pure local / pure cloud, so + * the extremes are exact rather than merely very likely. + * - Step 0 always orchestrates (when the cloud leg is in play at all): + * it forms the turn's plan and picks the first tool batch, which + * determines everything downstream. It is exactly one call per turn, + * so the cost is bounded and predictable. + */ +export function decideRoutingRole(args: RoutingDecisionArgs): RoutingRole { + if (args.cloudShare <= 0) return "executor"; + if (args.cloudShare >= 100) return "orchestrator"; + if (args.stepIndex === 0) return "orchestrator"; + + const cutoff = 100 - args.cloudShare; + const margin = + args.previousRole === "executor" + ? ROUTING_HYSTERESIS + : args.previousRole === "orchestrator" + ? -ROUTING_HYSTERESIS + : 0; + return args.score >= cutoff + margin ? "orchestrator" : "executor"; +} diff --git a/src/agent/routing/index.ts b/src/agent/routing/index.ts new file mode 100644 index 00000000..b5775b8e --- /dev/null +++ b/src/agent/routing/index.ts @@ -0,0 +1,11 @@ +export { computeStepComplexity } from "./compute-step-complexity.js"; +export type { StepComplexitySignals } from "./compute-step-complexity.js"; +export { decideRoutingRole, ROUTING_HYSTERESIS } from "./decide-routing-role.js"; +export type { RoutingDecisionArgs, RoutingRole } from "./decide-routing-role.js"; +export { StepRouter } from "./step-router.js"; +export type { + FusionRoutingSnapshot, + RouteStepArgs, + StepRouterDeps, + StepRouting, +} from "./step-router.js"; diff --git a/src/agent/routing/step-router.test.ts b/src/agent/routing/step-router.test.ts new file mode 100644 index 00000000..8bda5e10 --- /dev/null +++ b/src/agent/routing/step-router.test.ts @@ -0,0 +1,143 @@ +import { describe, expect, it } from "vitest"; + +import { StepRouter, type FusionRoutingSnapshot } from "./step-router.js"; + +const FUSION: FusionRoutingSnapshot = { + cloudProviderId: "openrouter", + localProviderId: "local-llama", + cloudShare: 40, + subRunners: "local", + maxSteps: 25, + conversationMaxTokens: 32_000, +}; + +function router(snapshot: FusionRoutingSnapshot | null = FUSION): StepRouter { + return new StepRouter({ resolveFusion: () => snapshot }); +} + +/** + * Scores 55 with the FUSION snapshot: above the cutoff a prior cloud + * step lowers it to (50), below the bare cutoff (60). That band is + * exactly where hysteresis is observable. + */ +const MEDIUM = { promptTokens: 24_000, stablePrefixTokens: 4_000, stepIndex: 10 }; +/** Scores 68: above the bare cutoff, below the "came from local" bar (70). */ +const HEAVY = { promptTokens: 30_000, stablePrefixTokens: 4_000, stepIndex: 15 }; + +const step = (over: Partial[0]> = {}) => ({ + sessionId: "s1", + stepIndex: 1, + promptTokens: 1_000, + stablePrefixTokens: 900, + hasTransientNotice: false, + ...over, +}); + +describe("StepRouter", () => { + it("returns null when fusion is not the effective mode", () => { + expect(router(null).routeStep(step())).toBeNull(); + }); + + it("routes step 0 to the cloud orchestrator", () => { + const routing = router().routeStep(step({ stepIndex: 0 })); + expect(routing).toMatchObject({ + role: "orchestrator", + providerId: "openrouter", + cloudShare: 40, + }); + }); + + it("routes a cheap continuation step to the local executor", () => { + const routing = router().routeStep(step()); + expect(routing).toMatchObject({ + role: "executor", + providerId: "local-llama", + }); + }); + + it("escalates a heavy continuation step to the cloud", () => { + const routing = router().routeStep( + step({ + stepIndex: 20, + promptTokens: 30_000, + stablePrefixTokens: 4_000, + hasTransientNotice: true, + }), + ); + expect(routing?.role).toBe("orchestrator"); + expect(routing?.complexity).toBeGreaterThanOrEqual(60); + }); + + it("re-reads the live snapshot on every step", () => { + let snapshot: FusionRoutingSnapshot = { ...FUSION, cloudShare: 0 }; + const r = new StepRouter({ resolveFusion: () => snapshot }); + expect(r.routeStep(step({ stepIndex: 0 }))?.role).toBe("executor"); + snapshot = { ...FUSION, cloudShare: 100 }; + expect(r.routeStep(step({ stepIndex: 1 }))?.role).toBe("orchestrator"); + }); + + it("keeps hysteresis state per session", () => { + const r = router(); + // Drive session A to the cloud (step 0 always orchestrates) and + // session B to local (a cheap continuation step). + expect(r.routeStep(step({ sessionId: "a", stepIndex: 0 }))?.role).toBe( + "orchestrator", + ); + expect(r.routeStep(step({ sessionId: "b", stepIndex: 1 }))?.role).toBe( + "executor", + ); + // Identical score, opposite prior roles ⇒ opposite decisions. + const a = r.routeStep(step({ sessionId: "a", ...HEAVY })); + const b = r.routeStep(step({ sessionId: "b", ...HEAVY })); + expect(a?.complexity).toBe(b?.complexity); + expect(a?.role).toBe("orchestrator"); + expect(b?.role).toBe("executor"); + }); + + it("forgets a session on request", () => { + const r = router(); + // Same MEDIUM step decided twice: sticky to the cloud while the + // prior role survives, back to the bare cutoff once it is dropped. + r.routeStep(step({ sessionId: "a", stepIndex: 0 })); + expect(r.routeStep(step({ sessionId: "a", ...MEDIUM }))?.role).toBe( + "orchestrator", + ); + r.forgetSession("a"); + expect(r.routeStep(step({ sessionId: "a", ...MEDIUM }))?.role).toBe( + "executor", + ); + }); + + it("drops hysteresis state when fusion is switched off", () => { + let snapshot: FusionRoutingSnapshot | null = FUSION; + const r = new StepRouter({ resolveFusion: () => snapshot }); + r.routeStep(step({ stepIndex: 0 })); + snapshot = null; + expect(r.routeStep(step())).toBeNull(); + snapshot = FUSION; + // The prior cloud role is gone, so the bare cutoff applies again. + expect(r.routeStep(step({ ...MEDIUM }))?.role).toBe("executor"); + }); + + it("sends sub-runners to the local leg by default", () => { + expect(router().subRunnerProviderId("s1")).toBe("local-llama"); + }); + + it("sends sub-runners to the cloud when configured", () => { + expect( + router({ ...FUSION, subRunners: "cloud" }).subRunnerProviderId("s1"), + ).toBe("openrouter"); + }); + + it("follows the session's last main-loop leg when configured", () => { + const r = router({ ...FUSION, subRunners: "follow" }); + r.routeStep(step({ sessionId: "s1", stepIndex: 0 })); + expect(r.subRunnerProviderId("s1")).toBe("openrouter"); + r.forgetSession("s1"); + expect(r.subRunnerProviderId("s1")).toBe("local-llama"); + }); + + it("has no sub-runner opinion outside fusion", () => { + expect(router(null).subRunnerProviderId("s1")).toBeNull(); + }); +}); diff --git a/src/agent/routing/step-router.ts b/src/agent/routing/step-router.ts new file mode 100644 index 00000000..a2fb4d77 --- /dev/null +++ b/src/agent/routing/step-router.ts @@ -0,0 +1,140 @@ +import type { RunModeSubRunners } from "../../config/llm-run-mode-config.js"; +import { computeStepComplexity } from "./compute-step-complexity.js"; +import { decideRoutingRole, type RoutingRole } from "./decide-routing-role.js"; + +/** + * Live fusion parameters. Re-read before every step so a mode switch or + * a dial change made in the TUI takes effect on the next inference + * rather than at the next process start — the same late-binding + * discipline `bootstrap` uses for `toolTransport` and slot affinity. + */ +export interface FusionRoutingSnapshot { + cloudProviderId: string; + localProviderId: string; + cloudShare: number; + subRunners: RunModeSubRunners; + maxSteps: number; + conversationMaxTokens: number; +} + +export interface StepRouterDeps { + /** Returns `null` whenever the effective run mode is not fusion. */ + resolveFusion: () => FusionRoutingSnapshot | null; +} + +export interface RouteStepArgs { + sessionId: string; + stepIndex: number; + promptTokens: number; + stablePrefixTokens: number; + hasTransientNotice: boolean; +} + +export interface StepRouting { + role: RoutingRole; + providerId: string; + complexity: number; + cloudShare: number; +} + +/** + * Bound on the per-session role memory. Only exists so a long-lived + * `serve` process cannot accumulate one entry per session forever; + * eviction is oldest-first, and losing an entry costs nothing but one + * step of hysteresis. + */ +const MAX_TRACKED_SESSIONS = 256; + +/** + * Chooses the fusion leg for each inference. + * + * Deliberately thin: scoring and the cutoff rule live in the two pure + * modules beside it, so the only thing here is the per-session memory + * that hysteresis needs. + * + * Note what is NOT here: repair-retry stickiness. The parse-repair call + * in `step-executor` spreads the original `LlmStreamParams`, so it + * inherits `preferredProviderId` from the attempt it is repairing for + * free — which is exactly the required behaviour, since a repair must + * be judged by the model that made the mistake and against the same + * transport. + */ +export class StepRouter { + private readonly resolveFusion: StepRouterDeps["resolveFusion"]; + private readonly lastRole = new Map(); + + constructor(deps: StepRouterDeps) { + this.resolveFusion = deps.resolveFusion; + } + + /** `null` ⇒ not in fusion; the caller leaves provider selection alone. */ + routeStep(args: RouteStepArgs): StepRouting | null { + const fusion = this.resolveFusion(); + if (!fusion) { + this.lastRole.delete(args.sessionId); + return null; + } + const complexity = computeStepComplexity({ + promptTokens: args.promptTokens, + stablePrefixTokens: args.stablePrefixTokens, + stepIndex: args.stepIndex, + hasTransientNotice: args.hasTransientNotice, + maxSteps: fusion.maxSteps, + conversationMaxTokens: fusion.conversationMaxTokens, + }); + const role = decideRoutingRole({ + score: complexity, + cloudShare: fusion.cloudShare, + stepIndex: args.stepIndex, + previousRole: this.lastRole.get(args.sessionId) ?? null, + }); + this.remember(args.sessionId, role); + return { + role, + providerId: + role === "orchestrator" + ? fusion.cloudProviderId + : fusion.localProviderId, + complexity, + cloudShare: fusion.cloudShare, + }; + } + + /** + * Provider for a memory sub-runner (reflection, link generation, + * curation votes, query rewriting, distillation). + * + * `local` by default: these are cold-path structured-JSON jobs that + * ride the reserved reflection slot and are already KV-warm on the + * local server, so routing them to the cloud multiplies per-turn cost + * with no user-visible latency win. `follow` reuses the leg the last + * main-loop step of that session used. + */ + subRunnerProviderId(sessionId?: string): string | null { + const fusion = this.resolveFusion(); + if (!fusion) return null; + const target: RunModeSubRunners = fusion.subRunners; + if (target === "cloud") return fusion.cloudProviderId; + if (target === "local") return fusion.localProviderId; + const last = sessionId ? this.lastRole.get(sessionId) : undefined; + return last === "orchestrator" + ? fusion.cloudProviderId + : fusion.localProviderId; + } + + /** Drop a finished session's hysteresis memory. */ + forgetSession(sessionId: string): void { + this.lastRole.delete(sessionId); + } + + private remember(sessionId: string, role: RoutingRole): void { + // Re-insert so the Map's insertion order doubles as recency. + this.lastRole.delete(sessionId); + this.lastRole.set(sessionId, role); + while (this.lastRole.size > MAX_TRACKED_SESSIONS) { + const oldest = this.lastRole.keys().next(); + if (oldest.done === true) break; + this.lastRole.delete(oldest.value); + } + } +} diff --git a/src/agent/step-events.ts b/src/agent/step-events.ts index 9cbbcad0..98f43880 100644 --- a/src/agent/step-events.ts +++ b/src/agent/step-events.ts @@ -32,6 +32,21 @@ export interface PromptCapturedTokens { */ export type StepEvent = | { type: "prompt_built"; prompt: BuiltPrompt; slotId: number } + /** + * Fusion routing picked a leg for this step. Emitted before the slot + * is acquired and only while fusion is the effective run mode, so its + * absence is the normal single-provider case rather than a gap. + */ + | { + type: "step_routed"; + stepIndex: number; + role: "orchestrator" | "executor"; + providerId: string; + /** 0-100 score from `computeStepComplexity`. */ + complexity: number; + /** The operator dial in force for this decision. */ + cloudShare: number; + } /** * Trace-oriented sibling of `prompt_built`: carries the salted hash of * the stable prefix (so per-step records stay small) together with the diff --git a/src/agent/step-executor-routing.test.ts b/src/agent/step-executor-routing.test.ts new file mode 100644 index 00000000..f888331b --- /dev/null +++ b/src/agent/step-executor-routing.test.ts @@ -0,0 +1,198 @@ +import { describe, it, expect } from "vitest"; +import { join } from "node:path"; + +import { executeStep } from "./step-executor.js"; +import type { LlmStreamParams, StepDependencies } from "./step-executor.js"; +import type { StepEvent } from "./step-events.js"; +import { StepRouter, type FusionRoutingSnapshot } from "./routing/index.js"; +import { ToolRegistry } from "../tools/tool-registry.js"; +import { compressToolResult } from "../compressor/result-compressor.js"; +import { SlotManager } from "../llm/slot-manager.js"; +import { PLAIN_INSTRUCT_PROFILE } from "../llm/model-profile.js"; +import { buildGrammar } from "../llm/grammar/build-grammar.js"; +import { createEmptySessionState } from "../session/session-state.js"; +import { DEFAULT_TOOL_DESCRIPTORS } from "../prompt/tool-descriptors.js"; +import type { + CapabilitiesSummary, + SkillCatalogEntry, +} from "../prompt/stable-prefix.js"; + +const CAPS: CapabilitiesSummary = { + platform: "darwin", + arch: "arm64", + browserChannel: "chrome", + workingDir: "/work", + hasClipboard: true, + hasWmctrl: false, + hasNotifications: true, +}; +const SKILLS: SkillCatalogEntry[] = []; + +const FUSION: FusionRoutingSnapshot = { + cloudProviderId: "cloud", + localProviderId: "local", + cloudShare: 40, + subRunners: "local", + maxSteps: 25, + conversationMaxTokens: 32_000, +}; + +function replyRegistry(): ToolRegistry { + const registry = new ToolRegistry(); + registry.register({ + name: "reply", + description: "reply", + readonly: true, + async run(args: Record) { + return compressToolResult({ + tool: "reply", + status: "ok", + output: String(args.text ?? ""), + }); + }, + }); + return registry; +} + +function completion(content: string) { + return { + content, + reasoningContent: "", + stop: true, + truncated: false, + timing: { promptMs: 1, predictedMs: 1, promptTokens: 20, predictedTokens: 5 }, + cacheHitTokens: 0, + slotId: 0, + modelId: "mock", + }; +} + +const REPLY_BODY = JSON.stringify({ tool: "reply", args: { text: "done" } }); + +/** + * Run one step and report what the LLM seam actually received, plus the + * events the step emitted. + */ +async function runStep(opts: { + router?: StepRouter; + resolveSlotAffinity?: (providerId: string) => boolean; + supportsSlotAffinity?: boolean; + bodies?: string[]; +}): Promise<{ seen: LlmStreamParams[]; events: StepEvent[] }> { + const grammar = await buildGrammar( + PLAIN_INSTRUCT_PROFILE, + join(process.cwd(), "grammars"), + ); + const seen: LlmStreamParams[] = []; + const events: StepEvent[] = []; + const bodies = opts.bodies ?? [REPLY_BODY]; + let call = 0; + + const deps = { + registry: replyRegistry(), + slotManager: new SlotManager(2), + llmComplete: async (params: LlmStreamParams) => { + seen.push(params); + const body = bodies[Math.min(call, bodies.length - 1)]!; + call += 1; + return completion(body); + }, + grammar, + profile: PLAIN_INSTRUCT_PROFILE, + supportsSlotAffinity: opts.supportsSlotAffinity ?? false, + onEvent: (event: StepEvent) => events.push(event), + ...(opts.router ? { stepRouter: opts.router } : {}), + ...(opts.resolveSlotAffinity + ? { resolveSlotAffinity: opts.resolveSlotAffinity } + : {}), + } as unknown as StepDependencies; + + await executeStep( + { + session: createEmptySessionState({ id: "s-route", workingDir: "/w" }), + toolDescriptors: DEFAULT_TOOL_DESCRIPTORS, + capabilities: CAPS, + skillCatalog: SKILLS, + stepIndex: 0, + signal: new AbortController().signal, + userMessage: "x", + }, + deps, + ); + return { seen, events }; +} + +const routerWith = (over: Partial = {}): StepRouter => + new StepRouter({ resolveFusion: () => ({ ...FUSION, ...over }) }); + +describe("executeStep fusion routing", () => { + it("sets no preferredProviderId when no router is wired", async () => { + const { seen, events } = await runStep({}); + expect(seen).toHaveLength(1); + expect(seen[0]).not.toHaveProperty("preferredProviderId"); + expect(events.some((e) => e.type === "step_routed")).toBe(false); + }); + + it("sets no preferredProviderId when the router declines (not fusion)", async () => { + const router = new StepRouter({ resolveFusion: () => null }); + const { seen, events } = await runStep({ router }); + expect(seen[0]).not.toHaveProperty("preferredProviderId"); + expect(events.some((e) => e.type === "step_routed")).toBe(false); + }); + + it("forwards the routed provider to the LLM seam", async () => { + // Step 0 always orchestrates ⇒ the cloud leg. + const { seen } = await runStep({ router: routerWith() }); + expect(seen[0]?.preferredProviderId).toBe("cloud"); + }); + + it("forwards the local leg when the dial is fully local", async () => { + const { seen } = await runStep({ router: routerWith({ cloudShare: 0 }) }); + expect(seen[0]?.preferredProviderId).toBe("local"); + }); + + it("emits step_routed describing the decision", async () => { + const { events } = await runStep({ router: routerWith() }); + const routed = events.find((e) => e.type === "step_routed"); + expect(routed).toMatchObject({ + type: "step_routed", + stepIndex: 0, + role: "orchestrator", + providerId: "cloud", + cloudShare: 40, + }); + }); + + it("acquires a real slot when the ROUTED provider has slot affinity", async () => { + // The active provider reports no affinity (cloud), but the step is + // routed to the local leg, which does. Without `resolveSlotAffinity` + // this would run at slotId -1 and reprocess the whole prompt. + const { seen } = await runStep({ + router: routerWith({ cloudShare: 0 }), + supportsSlotAffinity: false, + resolveSlotAffinity: (id) => id === "local", + }); + expect(seen[0]?.slotId).toBeGreaterThanOrEqual(0); + }); + + it("drops to slotId -1 when the routed provider has no slot affinity", async () => { + const { seen } = await runStep({ + router: routerWith({ cloudShare: 100 }), + supportsSlotAffinity: true, + resolveSlotAffinity: (id) => id === "local", + }); + expect(seen[0]?.preferredProviderId).toBe("cloud"); + expect(seen[0]?.slotId).toBe(-1); + }); + + it("repairs on the same leg that produced the malformed call", async () => { + const { seen } = await runStep({ + router: routerWith({ cloudShare: 0 }), + bodies: ["not json at all", REPLY_BODY], + }); + expect(seen.length).toBeGreaterThanOrEqual(2); + // A repair judged by the OTHER leg would parse a different model's + // mistake against a different transport. + expect(seen[1]?.preferredProviderId).toBe("local"); + }); +}); diff --git a/src/agent/step-executor.ts b/src/agent/step-executor.ts index 321f6066..0854d4b1 100644 --- a/src/agent/step-executor.ts +++ b/src/agent/step-executor.ts @@ -84,6 +84,7 @@ import type { ProfileFact } from "../memory/profile-store.js"; import type { AgentMetrics } from "../tracing/agent-metrics.js"; import type { StructuredLogger } from "../tracing/structured-logger.js"; import type { StepEvent } from "./step-events.js"; +import type { StepRouter } from "./routing/index.js"; export type { PromptCapturedTokens, StepEvent } from "./step-events.js"; export interface LlmStreamParams { @@ -119,6 +120,17 @@ export interface LlmStreamParams { * instead of waiting for the current step to finish on its own. */ signal?: AbortSignal; + /** + * Fusion routing: the provider id this call should START at. Only the + * starting link changes — the fallback chain still owns health, so a + * preferred provider in cooldown is ignored and a failure still + * advances through the chain. + * + * Deliberately a provider id rather than a role: the seam stays + * ignorant of run modes, and the policy that picked the leg lives in + * `src/agent/routing/`. Absent ⇒ today's behaviour, unchanged. + */ + preferredProviderId?: string; } export type LlmCompleteStream = ( @@ -145,6 +157,23 @@ export interface StepDependencies { toolCallAdapter: ToolCallAdapter | null; /** When false, completions use slotId -1 (cloud providers). */ supportsSlotAffinity: boolean; + /** + * Fusion step router. When present and fusion is the effective run + * mode, it picks the leg for each step; absent (or returning null) ⇒ + * provider selection is left entirely to the fallback chain, i.e. + * today's behaviour. + */ + stepRouter?: StepRouter; + /** + * Slot affinity for a SPECIFIC provider, used when `stepRouter` routes + * a step away from the active provider. + * + * Without this, fusion would read `supportsSlotAffinity` off the + * active (cloud) provider and run every locally-routed step with + * `slotId: -1` and `cachePrompt: false` — forcing llama-server to + * reprocess the whole prompt on each one. + */ + resolveSlotAffinity?: (providerId: string) => boolean; /** * Invoked after every LLM completion (initial call and one-shot parse * retry alike). Used by the agent loop to feed the served `modelId` @@ -302,7 +331,33 @@ async function executeStepInner( ? { userMessage: ctx.userMessage } : {}), }); - const slot = deps.supportsSlotAffinity + // Route BEFORE acquiring a slot: which provider serves this step + // decides whether a slot is worth acquiring at all. That ordering is + // also why the complexity score cannot use `cacheReused` — it does + // not exist yet, and making it an input would be circular. + const routing = + deps.stepRouter?.routeStep({ + sessionId: ctx.session.id, + stepIndex: ctx.stepIndex, + promptTokens: prompt.tokens.total, + stablePrefixTokens: prompt.tokens.stablePrefix, + hasTransientNotice: ctx.transientNotice !== undefined, + }) ?? null; + if (routing) { + deps.onEvent?.({ + type: "step_routed", + stepIndex: ctx.stepIndex, + role: routing.role, + providerId: routing.providerId, + complexity: routing.complexity, + cloudShare: routing.cloudShare, + }); + } + const slotAffinity = routing + ? (deps.resolveSlotAffinity?.(routing.providerId) ?? + deps.supportsSlotAffinity) + : deps.supportsSlotAffinity; + const slot = slotAffinity ? deps.slotManager.acquire(ctx.session.id, prompt.stablePrefix) : { slotId: -1, @@ -348,6 +403,7 @@ async function executeStepInner( sessionId: ctx.session.id, toolDescriptors: ctx.toolDescriptors, signal: ctx.signal, + ...(routing ? { preferredProviderId: routing.providerId } : {}), }); const firstAttempt = await runInitialCompletion({ @@ -1118,6 +1174,7 @@ function buildLlmStreamParams(args: { sessionId: string; toolDescriptors: readonly ToolDescriptor[]; signal?: AbortSignal; + preferredProviderId?: string; }): LlmStreamParams { const base: LlmStreamParams = { prompt: args.promptText, @@ -1125,6 +1182,9 @@ function buildLlmStreamParams(args: { slotId: args.slotId, sessionId: args.sessionId, ...(args.signal ? { signal: args.signal } : {}), + ...(args.preferredProviderId + ? { preferredProviderId: args.preferredProviderId } + : {}), }; if (args.deps.toolTransport !== "native_tools") { return base; diff --git a/src/llm/fallback/index.ts b/src/llm/fallback/index.ts index a4dc1f11..cfd7ec33 100644 --- a/src/llm/fallback/index.ts +++ b/src/llm/fallback/index.ts @@ -11,7 +11,10 @@ export { type ResolvedFallbackChain, } from "./fallback-config.js"; export { shouldAdvance, type AdvanceDecision } from "./should-advance.js"; -export { runWithFallback } from "./run-with-fallback.js"; +export { + runWithFallback, + type RunWithFallbackOptions, +} from "./run-with-fallback.js"; export { primeStream, replayPrimedStream, diff --git a/src/llm/fallback/provider-fallback-chain.test.ts b/src/llm/fallback/provider-fallback-chain.test.ts index d93b9581..59fe5a50 100644 --- a/src/llm/fallback/provider-fallback-chain.test.ts +++ b/src/llm/fallback/provider-fallback-chain.test.ts @@ -409,3 +409,101 @@ describe("ProviderFallbackChain", () => { }); }); }); + +describe("ProviderFallbackChain — fusion preferred start", () => { + it("starts at the preferred leg instead of the chain primary", () => { + const chain = new ProviderFallbackChain({ + resolve: () => chainOf(["cloud", "local"]), + }); + expect(chain.pickProvider("s1", "local")).toEqual({ + providerId: "local", + isProbe: false, + }); + }); + + it("never marks a preferred pick as a probe", () => { + const chain = new ProviderFallbackChain({ + resolve: () => chainOf(["cloud", "local"]), + }); + expect(chain.pickProvider("s1", "local").isProbe).toBe(false); + }); + + it("never sets the sticky override", () => { + const chain = new ProviderFallbackChain({ + resolve: () => chainOf(["cloud", "local"]), + }); + chain.pickProvider("s1", "local"); + expect(chain.activeOverrideFor("s1")).toBeNull(); + }); + + it("never clears an override that a real fallover established", () => { + const chain = new ProviderFallbackChain({ + resolve: () => chainOf(["cloud", "local"]), + }); + chain.advanceFrom("cloud", http(429), "s1"); + expect(chain.activeOverrideFor("s1")).toBe("local"); + chain.pickProvider("s1", "local"); + expect(chain.activeOverrideFor("s1")).toBe("local"); + }); + + it("accepts a preferred id that is not a chain member", () => { + const chain = new ProviderFallbackChain({ + resolve: () => chainOf(["cloud"]), + }); + expect(chain.pickProvider("s1", "local").providerId).toBe("local"); + }); + + it("ignores the preference while THAT leg is in cooldown", () => { + const clock = makeClock(); + const chain = new ProviderFallbackChain({ + resolve: () => chainOf(["cloud", "local"]), + now: clock.now, + }); + // A 429 on the local leg trips its breaker immediately. + chain.advanceFrom("local", http(429), "s1"); + expect(chain.pickProvider("s1", "local").providerId).not.toBe("local"); + // Once the cooldown elapses the preference is honoured again. + clock.advance(DEFAULT_FALLBACK_TIMING.cooldownMs[0]! + 1); + expect(chain.pickProvider("s1", "local").providerId).toBe("local"); + }); + + it("honours the preference while a DIFFERENT leg is in cooldown", () => { + const chain = new ProviderFallbackChain({ + resolve: () => chainOf(["cloud", "local"]), + }); + chain.advanceFrom("cloud", http(429), "s1"); + expect(chain.pickProvider("s1", "local").providerId).toBe("local"); + }); + + it("resumes from the chain head when a preferred TAIL start fails", () => { + const chain = new ProviderFallbackChain({ + resolve: () => chainOf(["cloud", "local"]), + }); + // Without `restartFromHead` this returns null (nothing after the + // tail) and the turn dies even though the cloud leg is healthy. + expect( + chain.advanceFrom("local", http(503), "s1", { restartFromHead: true }), + ).toBe("cloud"); + }); + + it("keeps the default advance behaviour when the flag is absent", () => { + const chain = new ProviderFallbackChain({ + resolve: () => chainOf(["cloud", "local"]), + }); + expect(chain.advanceFrom("local", http(503), "s1")).toBeNull(); + }); + + it("leaves an unpreferred pick byte-identical to today", () => { + const chain = new ProviderFallbackChain({ + resolve: () => chainOf(["cloud", "local"]), + }); + expect(chain.pickProvider("s1")).toEqual({ + providerId: "cloud", + isProbe: false, + }); + expect(chain.pickProvider("s1", undefined)).toEqual({ + providerId: "cloud", + isProbe: false, + }); + }); +}); diff --git a/src/llm/fallback/provider-fallback-chain.ts b/src/llm/fallback/provider-fallback-chain.ts index 05658ab0..43dbe6ca 100644 --- a/src/llm/fallback/provider-fallback-chain.ts +++ b/src/llm/fallback/provider-fallback-chain.ts @@ -36,6 +36,17 @@ export interface ProviderSwitchNotice { reason: string; } +/** Extra switching policy for one `advanceFrom` call. */ +export interface AdvanceOptions { + /** + * Resume the scan from the chain head rather than from just after + * `fromId`. Set when the failed attempt was a fusion-preferred start, + * whose position in the chain carries no "already tried everything + * above me" meaning. + */ + restartFromHead?: boolean; +} + /** What `pickProvider` decided for the turn about to run. */ export interface ProviderPick { /** Provider id to route this turn through. */ @@ -117,8 +128,20 @@ export class ProviderFallbackChain { * chain, drops a stale override that no longer names a chain member, * and — when the primary's cooldown has elapsed and the probe throttle * allows — routes this one turn back to the primary as a probe. + * + * `preferredId` is the fusion router's chosen leg for THIS call. It + * changes only the starting link, and health still wins: the + * preference is ignored while that specific provider is in cooldown, + * and a failure still advances through the chain as usual. It never + * sets or clears `overrideId` and is never reported as a probe — + * both of those are primary-recovery concepts, and a fusion pick is + * not a fallover. The id need not be a chain member; `advanceFrom` + * already restarts from the chain head for a non-member. */ - pickProvider(partitionKey: string = DEFAULT_PARTITION): ProviderPick { + pickProvider( + partitionKey: string = DEFAULT_PARTITION, + preferredId?: string, + ): ProviderPick { const { chain } = this.resolve(); const primary = chain[0]; if (!primary) { @@ -135,6 +158,16 @@ export class ProviderFallbackChain { this.clearOverride(p); } + // Fusion routing preference, checked before the override/probe + // logic so a healthy preferred leg is honoured — but only while + // that leg itself is healthy, so a tripped breaker still wins. + if (preferredId !== undefined && preferredId.length > 0) { + const preferred = this.breaker(p, preferredId); + if (this.now() >= preferred.cooldownUntil) { + return { providerId: preferredId, isProbe: false }; + } + } + if (!p.overrideId) { return { providerId: primary, isProbe: false }; } @@ -161,6 +194,7 @@ export class ProviderFallbackChain { fromId: string, err: unknown, partitionKey: string = DEFAULT_PARTITION, + options?: AdvanceOptions, ): string | null { const decision = shouldAdvance(err); if (!decision.advance) return null; @@ -171,8 +205,14 @@ export class ProviderFallbackChain { const idx = chain.indexOf(fromId); // Next healthy link after `fromId`. When `fromId` is not in the chain - // (raced config edit) start from the top. - const startFrom = idx < 0 ? 0 : idx + 1; + // (raced config edit) start from the top — and likewise when the + // failure came from a fusion-preferred start, which can sit anywhere + // in the chain and is commonly its TAIL. Advancing "after" the tail + // would strand a recoverable turn with the rest of the chain untried. + // The `candidate === fromId` guard below keeps the failed link out + // of the scan either way. + const startFrom = + idx < 0 || options?.restartFromHead === true ? 0 : idx + 1; for (let i = startFrom; i < chain.length; i += 1) { const candidate = chain[i]!; if (candidate === fromId) continue; diff --git a/src/llm/fallback/run-with-fallback.ts b/src/llm/fallback/run-with-fallback.ts index 0ae01894..a4cca257 100644 --- a/src/llm/fallback/run-with-fallback.ts +++ b/src/llm/fallback/run-with-fallback.ts @@ -17,14 +17,31 @@ import type { ProviderFallbackChain } from "./provider-fallback-chain.js"; * chunks is never restarted (mirrors the openai-http "stream is live" * contract), so failures after the first chunk propagate as-is. */ +export interface RunWithFallbackOptions { + /** + * Provider to START at for this unit of work (fusion routing). Only + * the starting link changes: the switching policy below is untouched, + * so a failure still advances through the chain and a preferred + * provider in cooldown is ignored. + */ + preferredProviderId?: string; +} + export async function runWithFallback( chain: ProviderFallbackChain, attempt: (providerId: string) => Promise, partitionKey?: string, + options?: RunWithFallbackOptions, ): Promise { - const pick = chain.pickProvider(partitionKey); + const pick = chain.pickProvider(partitionKey, options?.preferredProviderId); let currentId = pick.providerId; let wasProbe = pick.isProbe; + // True only while we are still sitting on the fusion-preferred start. + // A failure there resumes the chain from its head, because a preferred + // leg's position in the chain says nothing about what has been tried. + let onPreferredStart = + options?.preferredProviderId !== undefined && + currentId === options.preferredProviderId; // Guard against a pathological empty chain: no provider to try. if (!currentId) { @@ -37,12 +54,15 @@ export async function runWithFallback( chain.recordSuccess(currentId, wasProbe, partitionKey); return result; } catch (err) { - const nextId = chain.advanceFrom(currentId, err, partitionKey); + const nextId = chain.advanceFrom(currentId, err, partitionKey, { + restartFromHead: onPreferredStart, + }); if (nextId === null) throw err; currentId = nextId; // Only the very first pick can be a probe; every advance is a real // fallover on the working path. wasProbe = false; + onPreferredStart = false; } } } diff --git a/src/llm/provider/completion-types.ts b/src/llm/provider/completion-types.ts index d6cd5644..41109097 100644 --- a/src/llm/provider/completion-types.ts +++ b/src/llm/provider/completion-types.ts @@ -98,6 +98,15 @@ export interface CompletionResult { * authoritative. */ servedTransport?: ToolCallTransport; + /** + * Id of the provider that actually served this completion. Stamped by + * the same wrapper, and for the same reason as `servedTransport`: + * under fusion routing (and after any fallover) the link that answered + * is not necessarily `llm.activeTextProvider`, so anything that + * attributes the result — cost/pricing lookup above all — has to ask + * who served it rather than who was active. + */ + servedProviderId?: string; } export interface OpenAiToolCall { diff --git a/src/llm/provider/registry/provider-registry.test.ts b/src/llm/provider/registry/provider-registry.test.ts index 2fb5eef5..33bb6136 100644 --- a/src/llm/provider/registry/provider-registry.test.ts +++ b/src/llm/provider/registry/provider-registry.test.ts @@ -87,3 +87,68 @@ describe("ProviderRegistry", () => { ).rejects.toThrow(/unknown llm provider kind/); }); }); + +describe("ProviderRegistry pinned providers", () => { + /** Minimal stand-in that records whether it was torn down. */ + function fake(id: string) { + const state = { closed: false }; + const provider = { + id, + name: id, + capabilities: {}, + toolCallAdapter: null, + streamConsumer: null, + async complete() { + throw new Error("unused"); + }, + async *completeStream() { + throw new Error("unused"); + }, + async describeImage() { + throw new Error("unused"); + }, + async health() { + return { reachable: true, status: 200, error: null, latencyMs: 1 }; + }, + async close() { + state.closed = true; + }, + }; + return { provider, state }; + } + + function registryOf(ids: string[]) { + const fakes = ids.map((id) => fake(id)); + const map = new Map( + fakes.map((f) => [f.provider.id, f.provider as never] as const), + ); + // `new ProviderRegistry(...)` is private to the module's factory, so + // reach it the same way `fromConfig` does. + const registry = Reflect.construct(ProviderRegistry, [ids[0], map]) as + ProviderRegistry; + return { registry, fakes }; + } + + it("closes the previous provider on a plain swap", async () => { + const { registry, fakes } = registryOf(["cloud", "local"]); + await registry.swapActive("local"); + expect(fakes[0]!.state.closed).toBe(true); + }); + + it("keeps a pinned provider open across a swap", async () => { + // Fusion keeps both legs live; closing the one it is about to route + // to would break the executor leg on the very next step. + const { registry, fakes } = registryOf(["cloud", "local"]); + registry.setPinnedProviderIds(() => new Set(["cloud", "local"])); + await registry.swapActive("local"); + expect(fakes[0]!.state.closed).toBe(false); + expect(registry.activeText.id).toBe("local"); + }); + + it("resumes closing once nothing is pinned", async () => { + const { registry, fakes } = registryOf(["cloud", "local"]); + registry.setPinnedProviderIds(() => new Set()); + await registry.swapActive("local"); + expect(fakes[0]!.state.closed).toBe(true); + }); +}); diff --git a/src/llm/provider/registry/provider-registry.ts b/src/llm/provider/registry/provider-registry.ts index a8ac4243..cb3dda62 100644 --- a/src/llm/provider/registry/provider-registry.ts +++ b/src/llm/provider/registry/provider-registry.ts @@ -28,6 +28,8 @@ export class ProviderRegistry { this.providers = providers; } + private pinnedProviderIds?: () => ReadonlySet; + static async fromConfig( config: AtomicAgentConfig, ctx: Omit & { @@ -76,6 +78,21 @@ export class ProviderRegistry { return [...this.providers.keys()]; } + /** + * Providers that must stay open even when they stop being active. + * + * Fusion keeps two legs live at once and only one of them can be the + * active provider, so switching INTO fusion would otherwise close the + * very provider it is about to route to. `close()` is a no-op on both + * shipped provider kinds today, which is why this is not currently a + * visible crash — but the interface promises teardown, and the first + * provider kind that honours it (pooled sockets, a WS transport) + * would break fusion silently without this. + */ + setPinnedProviderIds(pinned: () => ReadonlySet): void { + this.pinnedProviderIds = pinned; + } + async swapActive(id: string): Promise { const next = this.providers.get(id); if (!next) { @@ -83,7 +100,7 @@ export class ProviderRegistry { } const prev = this.providers.get(this.activeTextId); this.activeTextId = id; - if (prev && prev.id !== id) { + if (prev && prev.id !== id && !this.pinnedProviderIds?.().has(prev.id)) { await prev.close().catch(() => undefined); } return next; diff --git a/src/runtime/bootstrap.ts b/src/runtime/bootstrap.ts index 68000664..143f4a2c 100644 --- a/src/runtime/bootstrap.ts +++ b/src/runtime/bootstrap.ts @@ -10,6 +10,8 @@ import { } from "../config/index.js"; import type { LlmStreamParams } from "../agent/step-executor.js"; +import { StepRouter } from "../agent/routing/index.js"; +import { resolveRunMode } from "../llm/run-mode/index.js"; import { TurnController } from "./turn-controller.js"; import type { TurnEventHook, TurnOrigin } from "./turn-controller.js"; import type { ChannelStatus } from "./channel-status.js"; @@ -1188,11 +1190,16 @@ export async function createAgentRuntime( */ const resolveModelPricing = ( modelId: string | null, + servedProviderId?: string, ): ResolvedModel | undefined => { if (!modelId) return undefined; const resolved = resolveLlmConfig(getConfig()); + // Price against the provider that actually SERVED the completion, + // not the active one. They differ after any fallover, and routinely + // under fusion — pricing a local completion against the cloud + // provider's catalog reports a cost that was never incurred. const entry = resolved.providers.find( - (p) => p.id === resolved.activeTextProvider, + (p) => p.id === (servedProviderId ?? resolved.activeTextProvider), ); if (!entry) return undefined; return resolveModel(entry, modelId, catalogForProvider(entry)); @@ -1343,9 +1350,10 @@ export async function createAgentRuntime( const recordUnaryUsage = ( params: LlmStreamParams, result: CompletionResult, + servedProviderId?: string, ): void => { if (!result.usage) return; - const model = resolveModelPricing(result.modelId); + const model = resolveModelPricing(result.modelId, servedProviderId); if (costAccumulator) { costAccumulator.recordTurn({ modelId: result.modelId, @@ -1367,7 +1375,7 @@ export async function createAgentRuntime( result: CompletionResult, ): void => { if (!result.usage || !sessionId) return; - const model = resolveModelPricing(result.modelId); + const model = resolveModelPricing(result.modelId, result.servedProviderId); turnUsageMeter.record({ sessionId, usage: result.usage, @@ -1375,6 +1383,40 @@ export async function createAgentRuntime( }); }; + // Keep both fusion legs open across an active-provider swap. Without + // this, switching into fusion closes the provider it routes to. + providerRegistry.setPinnedProviderIds(() => { + const runMode = resolveRunMode(resolveLlmConfig(getConfig())); + if (runMode.effective !== "fusion") return new Set(); + return new Set( + [runMode.cloudProviderId, runMode.localProviderId].filter( + (id): id is string => id !== null, + ), + ); + }); + + /** + * Fusion step router. Always constructed; it resolves the live config + * on every step and returns `null` unless fusion is the effective run + * mode, so a non-fusion install pays one config read and nothing else. + */ + const stepRouter = new StepRouter({ + resolveFusion: () => { + const live = getConfig(); + const runMode = resolveRunMode(resolveLlmConfig(live)); + if (runMode.effective !== "fusion") return null; + if (!runMode.cloudProviderId || !runMode.localProviderId) return null; + return { + cloudProviderId: runMode.cloudProviderId, + localProviderId: runMode.localProviderId, + cloudShare: runMode.fusion.cloudShare, + subRunners: runMode.fusion.subRunners, + maxSteps: live.agent.maxSteps, + conversationMaxTokens: live.agent.conversationMaxTokens, + }; + }, + }); + const fallbackSeamDeps: FallbackSeamDeps = { fallbackChain, resolveSlice: (providerId) => { @@ -1396,6 +1438,26 @@ export async function createAgentRuntime( ? undefined : createFallbackStreamer(fallbackSeamDeps)); + /** + * Completion seam for the memory sub-runners (reflection, link + * generation, curation votes, query rewriting, distillation). + * + * Under fusion these default to the LOCAL leg: they are cold-path, + * fire-and-forget structured-JSON jobs that ride the reserved + * reflection slot and are already KV-warm on the local server, so + * sending them to the cloud multiplies per-turn cost with no + * user-visible latency win. `llm.runMode.fusion.subRunners` overrides + * it. Outside fusion this is `llmComplete` with no added behaviour. + */ + const subRunnerLlmComplete = ( + params: LlmStreamParams, + ): Promise => { + const providerId = stepRouter.subRunnerProviderId(params.sessionId); + return llmComplete( + providerId ? { ...params, preferredProviderId: providerId } : params, + ); + }; + const taskStore = new TaskStore({ dbFile: config.paths.tasksDbFile }); const webhookSessionStore = new WebhookSessionStore( resolve(config.paths.stateDir, "webhook-sessions.json"), @@ -1426,7 +1488,7 @@ export async function createAgentRuntime( const baseReflectionRunner = buildReflectionRunner({ config, slotManager, - llmComplete, + llmComplete: subRunnerLlmComplete, toolTransport: bootstrapLlmSlice.transport, profileStore, notesStore, @@ -1483,7 +1545,7 @@ export async function createAgentRuntime( { once: true }, ); }); - const completionPromise = llmComplete({ + const completionPromise = subRunnerLlmComplete({ prompt: params.prompt, grammar: params.grammar, slotId: params.slotId, @@ -1552,7 +1614,7 @@ export async function createAgentRuntime( { once: true }, ); }); - const completionPromise = llmComplete({ + const completionPromise = subRunnerLlmComplete({ prompt: params.prompt, grammar: params.grammar, slotId: params.slotId, @@ -1688,7 +1750,7 @@ export async function createAgentRuntime( { once: true }, ); }); - const completionPromise = llmComplete({ + const completionPromise = subRunnerLlmComplete({ prompt: params.prompt, grammar: params.grammar, slotId: params.slotId, @@ -1752,6 +1814,7 @@ export async function createAgentRuntime( registry: toolRegistry, slotManager, grammar, + stepRouter, llmComplete, ...(llmCompleteStream ? { llmCompleteStream } : {}), toolDescriptors: effectiveToolDescriptors, @@ -1833,6 +1896,15 @@ export async function createAgentRuntime( enumerable: true, get: () => resolveActiveLlmSlice().slotAffinity, }); + // Slot affinity for a specific routed provider. Without this, fusion + // would read affinity off the active (cloud) provider and run every + // locally-routed step with slotId -1 — no prompt cache at all on the + // local leg. + Object.defineProperty(loopDeps, "resolveSlotAffinity", { + enumerable: true, + get: () => (providerId: string) => + resolveActiveLlmSlice(providerId).slotAffinity, + }); const loop = new AgentLoop( loopDeps as typeof loopDeps & { skillCatalog: readonly SkillCatalogEntry[]; @@ -2224,7 +2296,7 @@ export async function createAgentRuntime( { once: true }, ); }); - const completionPromise = llmComplete({ + const completionPromise = subRunnerLlmComplete({ prompt: params.prompt, grammar: params.grammar, slotId: params.slotId, diff --git a/src/runtime/llm-fallback-seam.test.ts b/src/runtime/llm-fallback-seam.test.ts index a2a8197a..32ac7719 100644 --- a/src/runtime/llm-fallback-seam.test.ts +++ b/src/runtime/llm-fallback-seam.test.ts @@ -184,3 +184,73 @@ describe("createFallbackStreamer (real bootstrap seam)", () => { expect(result.servedTransport).toBe("native_tools"); }); }); + +describe("fusion routing through the real seam", () => { + const providers = () => + new Map([ + ["cloud", fakeProvider("cloud", "native_tools", async () => answer("cloud"))], + ["local", fakeProvider("local", "grammar", async () => answer("local"))], + ]); + + it("stamps servedProviderId with the link that answered", async () => { + const complete = createFallbackCompleter(seamDeps(providers())); + const result = await complete(baseParams); + expect(result.servedProviderId).toBe("cloud"); + }); + + it("starts at preferredProviderId instead of the chain primary", async () => { + const complete = createFallbackCompleter(seamDeps(providers())); + const result = await complete({ + ...baseParams, + preferredProviderId: "local", + }); + // Load-bearing: "local" is the chain TAIL, so without the + // preference plumbing this would answer from "cloud". + expect(result.servedProviderId).toBe("local"); + expect(result.modelId).toBe("local-model"); + // And the transport stamp must follow the routed leg, not the primary. + expect(result.servedTransport).toBe("grammar"); + }); + + it("still falls over on health when the preferred leg fails", async () => { + const map = new Map([ + ["cloud", fakeProvider("cloud", "native_tools", async () => answer("cloud"))], + [ + "local", + fakeProvider("local", "grammar", async () => { + throw new OpenAiHttpError("boom", 503, "http://local", false, null, "local"); + }), + ], + ]); + const complete = createFallbackCompleter(seamDeps(map)); + const result = await complete({ + ...baseParams, + preferredProviderId: "local", + }); + expect(result.servedProviderId).toBe("cloud"); + }); + + it("prices against the served leg, not the active one", async () => { + // Guards the fusion cost-attribution fix in bootstrap: the recorder + // is handed the id of the link that answered. + const seen: string[] = []; + const deps = seamDeps(providers()); + const complete = createFallbackCompleter({ + ...deps, + recordUnaryUsage: (_params, _result, servedProviderId) => { + seen.push(servedProviderId); + }, + }); + await complete({ ...baseParams, preferredProviderId: "local" }); + expect(seen).toEqual(["local"]); + }); + + it("routes the stream seam by preference and stamps the served id", async () => { + const stream = createFallbackStreamer(seamDeps(providers())); + const gen = stream({ ...baseParams, preferredProviderId: "local" }); + let next = await gen.next(); + while (next.done !== true) next = await gen.next(); + expect(next.value.servedProviderId).toBe("local"); + expect(next.value.servedTransport).toBe("grammar"); + }); +}); diff --git a/src/runtime/llm-fallback-seam.ts b/src/runtime/llm-fallback-seam.ts index d77f7e87..ed54348d 100644 --- a/src/runtime/llm-fallback-seam.ts +++ b/src/runtime/llm-fallback-seam.ts @@ -34,8 +34,17 @@ export interface FallbackSeamDeps { fallbackChain: ProviderFallbackChain; /** Resolve the served link's provider + transport for `providerId`. */ resolveSlice: (providerId: string) => ResolvedLinkSlice; - /** Fold a unary completion's usage into cost + meter (no-op when absent). */ - recordUnaryUsage: (params: LlmStreamParams, result: CompletionResult) => void; + /** + * Fold a unary completion's usage into cost + meter (no-op when + * absent). `servedProviderId` names the link that answered, which is + * what pricing must be looked up against — under fusion it routinely + * differs from `llm.activeTextProvider`. + */ + recordUnaryUsage: ( + params: LlmStreamParams, + result: CompletionResult, + servedProviderId: string, + ) => void; /** Fold a streamed completion's usage into the meter. */ recordStreamUsage: ( sessionId: string | undefined, @@ -92,10 +101,13 @@ export function createFallbackCompleter( slotId: params.slotId, cachePrompt: params.slotId >= 0, }); - deps.recordUnaryUsage(params, result); - return { ...result, servedTransport: transport }; + deps.recordUnaryUsage(params, result, providerId); + return { ...result, servedTransport: transport, servedProviderId: providerId }; }, params.sessionId, + { ...(params.preferredProviderId + ? { preferredProviderId: params.preferredProviderId } + : {}) }, ); } @@ -117,6 +129,7 @@ export function createFallbackStreamer( ): Promise<{ primed: PrimedStream; transport: ToolCallTransport; + providerId: string; }> => { const { provider, transport } = deps.resolveSlice(providerId); const base = { @@ -142,18 +155,21 @@ export function createFallbackStreamer( slotId: params.slotId, cachePrompt: params.slotId >= 0, }); - return { primed: await primeStream(stream), transport }; + return { primed: await primeStream(stream), transport, providerId }; }; return (params) => { async function* run(): AsyncGenerator { - const { primed, transport } = await runWithFallback( + const { primed, transport, providerId } = await runWithFallback( deps.fallbackChain, (id) => openStreamPrimed(id, params), params.sessionId, + { ...(params.preferredProviderId + ? { preferredProviderId: params.preferredProviderId } + : {}) }, ); const result = yield* replayPrimedStream(primed); - return { ...result, servedTransport: transport }; + return { ...result, servedTransport: transport, servedProviderId: providerId }; } return meterStream(deps, params.sessionId, run()); }; diff --git a/src/tui/agent-event-reducer.ts b/src/tui/agent-event-reducer.ts index b72d64b9..7aabe437 100644 --- a/src/tui/agent-event-reducer.ts +++ b/src/tui/agent-event-reducer.ts @@ -375,6 +375,21 @@ function reduceStepEvent( }, }; } + case "step_routed": { + return appendFeed(state, { + kind: "runtime_info", + stepIndex: event.stepIndex, + line: formatFeedLine({ + type: "step_routed", + stepIndex: event.stepIndex, + role: event.role, + providerId: event.providerId, + complexity: event.complexity, + cloudShare: event.cloudShare, + }), + color: "blue", + }); + } case "parse_retry": { const withFeed = appendFeed(state, { kind: "runtime_info", diff --git a/src/tui/format-event.ts b/src/tui/format-event.ts index 696eb816..8e33dd80 100644 --- a/src/tui/format-event.ts +++ b/src/tui/format-event.ts @@ -40,6 +40,14 @@ export type FeedLineInput = | { type: "rare_tool_autoloaded"; tool: string; + } + | { + type: "step_routed"; + stepIndex: number; + role: "orchestrator" | "executor"; + providerId: string; + complexity: number; + cloudShare: number; }; const ARGS_PREVIEW_LIMIT = 160; @@ -66,6 +74,15 @@ export function formatFeedLine(input: FeedLineInput): string { case "rare_tool_autoloaded": { return ` ↻ loaded schema for ${input.tool} after tool error`; } + case "step_routed": { + // Show the cutoff alongside the score so the line explains the + // decision rather than just announcing it. + const cutoff = 100 - input.cloudShare; + const leg = + input.role === "orchestrator" ? "cloud orchestrator" : "local executor"; + const comparison = input.role === "orchestrator" ? "≥" : "<"; + return `[step ${input.stepIndex}] → ${leg} ${input.providerId} (complexity ${input.complexity} ${comparison} ${cutoff})`; + } default: return ""; }