From c32dd2b55a78d81474cf91b35c8f6ae617f0f69d Mon Sep 17 00:00:00 2001 From: Nadia Rodionova <273119990+sosidudku1@users.noreply.github.com> Date: Wed, 11 Nov 2026 02:16:55 +0300 Subject: [PATCH] fix(fusion): keep local server health per server With mixLocalModels two local servers can serve one fan-out, but the health mark set after a cut-off worker was shared, and the only probe asked the main server. A wedge on one server could fail a worker sent to the other, and an answer from the main server cleared the other's mark. Marks are now filed under the leg's provider id, and a leg is asked only through its own server's probe. The runtime has a probe for the main local server only, so a second server's leg is not checked rather than checked against the wrong one. (ATO-1318) --- .../composition/runtime-tool-catalog.ts | 14 +++++ src/tools/fusion/local-leg-health.test.ts | 55 ++++++++++++++++++- src/tools/fusion/local-leg-health.ts | 46 ++++++++++------ src/tools/fusion/worker-runner.ts | 17 +++++- 4 files changed, 112 insertions(+), 20 deletions(-) diff --git a/src/runtime/composition/runtime-tool-catalog.ts b/src/runtime/composition/runtime-tool-catalog.ts index 677bda242..0b0e98a01 100644 --- a/src/runtime/composition/runtime-tool-catalog.ts +++ b/src/runtime/composition/runtime-tool-catalog.ts @@ -692,6 +692,20 @@ export function registerRuntimeFusionAndReadScope(input: { // Asked before a local worker is sent after one was cut off // (ATO-917): a server that stopped answering fails the task now. probeLocalHealth: () => llama.probeHealth(), + // Per local server (ATO-1318): this client only ever asks the main + // local server (`localModels.url`), so it is the probe for a leg on + // that server alone. A second local server has no probe here and + // its leg is not checked, rather than checked against the wrong one. + probeLocalHealthFor: (providerId) => { + const entry = resolveLlmConfig(getConfig()).providers.find((p) => p.id === providerId); + if (entry === undefined || entry.kind !== "llama-server") return undefined; + const main = getConfig().localModels.url; + const same = (a: string, b: string): boolean => + a.trim().replace(/\/+$/, "") === b.trim().replace(/\/+$/, ""); + return entry.url === undefined || same(entry.url, main) + ? () => llama.probeHealth() + : undefined; + }, // The turn's $ / time caps (ATO-947), reset at each fusion turn's // start by the loop; read only when a cap is configured. ...(turnBudget === undefined ? {} : { turnBudget }), diff --git a/src/tools/fusion/local-leg-health.test.ts b/src/tools/fusion/local-leg-health.test.ts index 95afd881a..c9a084ed4 100644 --- a/src/tools/fusion/local-leg-health.test.ts +++ b/src/tools/fusion/local-leg-health.test.ts @@ -132,7 +132,7 @@ describe("dispatch after a cut-off local worker (ATO-917)", () => { it("dispatches as before once the server answers, and clears the mark", async () => { const health = new LocalLegHealth(); - health.noteCut("task t9 was cancelled"); + health.noteCut("task t9 was cancelled", "local-llama"); const deps = harness(okTurn, { localLegHealth: health, probeLocalHealth: async () => "answered", @@ -167,4 +167,57 @@ describe("dispatch after a cut-off local worker (ATO-917)", () => { expect(probes).toBe(0); expect(deps.turns).toBe(1); }); + + it("keeps each local server's mark apart and asks only the leg's own server (ATO-1318)", async () => { + const health = new LocalLegHealth(); + health.noteCut("task t9 was timed out", "second-llama"); + const asked: string[] = []; + const deps = harness(okTurn, { + localLegHealth: health, + probeLocalHealthFor: (providerId) => + providerId === "local-llama" + ? async () => { + asked.push(providerId); + return "unanswered"; + } + : undefined, + }); + // The main server was never cut: no question, the worker goes. + const [main] = await runWorkerTasks(deps, { + ...BASE, + tasks: TASK, + signal: new AbortController().signal, + }); + expect(main!.status).not.toBe("failed"); + expect(asked).toEqual([]); + // The cut server has no probe of its own: not checked against the + // main one, and its mark stays. + const [second] = await runWorkerTasks(deps, { + ...BASE, + providerId: "second-llama", + tasks: TASK, + signal: new AbortController().signal, + }); + expect(second!.status).not.toBe("failed"); + expect(asked).toEqual([]); + expect(health.isSuspect("second-llama")).toBe(true); + expect(health.isSuspect("local-llama")).toBe(false); + }); + + it("files a cut under the worker's own server", async () => { + vi.useFakeTimers(); + const health = new LocalLegHealth(); + const deps = harness(servedThenHung, { localLegHealth: health }); + const run = runWorkerTasks(deps, { + ...BASE, + providerId: "second-llama", + tasks: TASK, + signal: new AbortController().signal, + }); + await vi.advanceTimersByTimeAsync(46 * MIN); + await run; + expect(health.isSuspect("second-llama")).toBe(true); + expect(health.isSuspect("local-llama")).toBe(false); + expect(health.cause("second-llama")).toMatch(/task t0 was timed out/); + }); }); diff --git a/src/tools/fusion/local-leg-health.ts b/src/tools/fusion/local-leg-health.ts index b07b4fa0c..1167757dc 100644 --- a/src/tools/fusion/local-leg-health.ts +++ b/src/tools/fusion/local-leg-health.ts @@ -31,36 +31,49 @@ export const LOCAL_LEG_UNRESPONSIVE_ERROR = export const LOCAL_LEG_UNRESPONSIVE_HINT = "the local server is not answering /health since a worker on it was cancelled or timed out mid-generation; restart the local server (Manage → Local models, or restart the app) before delegating again, or do this part yourself"; +/** The key a mark is filed under when the caller names no server. */ +export const DEFAULT_LOCAL_LEG = "local"; + +/** + * Kept per local server (ATO-1318). With `mixLocalModels` two local + * servers can serve one fan-out; one shared mark let a wedge on one + * fail a worker sent to the other, or a healthy answer from one clear + * the other's mark. Each leg is filed under its provider id, and only + * a leg with a probe of its own is ever asked. + */ export class LocalLegHealth { - private suspect = false; - private suspectSince: string | null = null; + private readonly suspects = new Map(); - /** A worker on the local leg ended by an abort; check before the next dispatch. */ - noteCut(reason: string): void { - if (!this.suspect) this.suspectSince = reason; - this.suspect = true; + /** A worker on this local leg ended by an abort; check before the next dispatch there. */ + noteCut(reason: string, leg: string = DEFAULT_LOCAL_LEG): void { + if (!this.suspects.has(leg)) this.suspects.set(leg, reason); } - isSuspect(): boolean { - return this.suspect; + /** Whether `leg` is marked; with no leg named, whether any is. */ + isSuspect(leg?: string): boolean { + return leg === undefined ? this.suspects.size > 0 : this.suspects.has(leg); } - /** Why the leg was marked, for the row's note. */ - cause(): string | null { - return this.suspectSince; + /** Why the leg was marked, for the row's note; with no leg named, the first mark. */ + cause(leg?: string): string | null { + if (leg !== undefined) return this.suspects.get(leg) ?? null; + const first = this.suspects.values().next(); + return first.done === true ? null : first.value; } /** - * Ask the server before dispatching. Resolves `true` when the worker may - * be sent: nothing was cut, there is no probe to ask with, or the server - * answered (which clears the mark). Never throws. + * Ask the leg's server before dispatching. Resolves `true` when the + * worker may be sent: nothing on this leg was cut, there is no probe + * for this leg to ask with, or the server answered (which clears this + * leg's mark only). Never throws. */ async readyForDispatch( probe: (() => Promise) | undefined, sleep: (ms: number) => Promise = defaultSleep, signal?: AbortSignal, + leg: string = DEFAULT_LOCAL_LEG, ): Promise { - if (!this.suspect || probe === undefined) return true; + if (!this.suspects.has(leg) || probe === undefined) return true; for (let attempt = 0; attempt < LOCAL_LEG_HEALTH_ATTEMPTS; attempt += 1) { if (signal?.aborted) return true; if (attempt > 0) await sleep(LOCAL_LEG_HEALTH_RETRY_MS); @@ -71,8 +84,7 @@ export class LocalLegHealth { answer = "unanswered"; } if (answer === "answered") { - this.suspect = false; - this.suspectSince = null; + this.suspects.delete(leg); return true; } } diff --git a/src/tools/fusion/worker-runner.ts b/src/tools/fusion/worker-runner.ts index c0d72619f..3fbd489f6 100644 --- a/src/tools/fusion/worker-runner.ts +++ b/src/tools/fusion/worker-runner.ts @@ -483,6 +483,15 @@ export interface WorkerRunnerDeps { localLegHealth?: LocalLegHealth; /** `LlamaServerClient.probeHealth`, the question the check above asks. */ probeLocalHealth?: () => Promise; + /** + * The probe for one local leg's own server (ATO-1318), or `undefined` + * when nothing can ask that server; then the check is skipped for it. + * When set it replaces `probeLocalHealth`, which only ever asks the + * main local server. + */ + probeLocalHealthFor?: ( + providerId: string, + ) => (() => Promise) | undefined; /** * One completion on the escalation provider (ATO-942). Absent, nothing * is ever escalated, manual or automatic. @@ -719,12 +728,15 @@ async function runOneTask( options.localLeg === true && deps.localLegHealth !== undefined && !(await deps.localLegHealth.readyForDispatch( - deps.probeLocalHealth, + deps.probeLocalHealthFor !== undefined + ? deps.probeLocalHealthFor(options.providerId) + : deps.probeLocalHealth, undefined, options.signal, + options.providerId, )) ) { - const cause = deps.localLegHealth.cause(); + const cause = deps.localLegHealth.cause(options.providerId); const refused: WorkerTaskResult = { ...collector.finish({ id: task.id, @@ -1614,6 +1626,7 @@ async function runOneTask( ) { deps.localLegHealth?.noteCut( `task ${task.id} was ${timedOutOutcome ? "timed out" : queuedOutcome ? "never served" : stoppedByHandBack ? "handed back" : "cancelled"}`, + options.providerId, ); }