From 5ce61ff3d551cf76ae6e28bf392758320c14e22f Mon Sep 17 00:00:00 2001 From: Chris Kelsey Date: Tue, 22 Sep 2026 11:36:12 -0700 Subject: [PATCH 01/12] Add model download cancellation and progress Model downloads now report live progress and can be canceled, for both Ollama and LM Studio, on local and remote nodes. The row holds "Canceling" until the backend confirms the transfer has stopped and its files are settled. Cancellation removes only the partial files that download created. Both engines funnel every client on a machine into one cache, so naming a file is not the same as owning it: each pull records the partial files already present before it starts and considers only those that appear afterwards, and a file still growing once the transfer stops is treated as shared and preserved. Completed models, other quantizations, and unrelated downloads are never touched, and a cancellation the LM Studio CLI never confirmed deletes nothing. Only a cancellation someone requested removes files; a pull interrupted by shutdown, a dropped remote connection, or the action timeout leaves its partial data resumable. Downloads are serialized per engine so cleanup cannot race another download, and a queued download can be canceled or retried before it starts. Signed-off-by: Chris Kelsey --- desktop/docs/services-api.md | 2 + desktop/docs/services-parity.md | 113 ++- .../electron/service-bridge/empty-handlers.ts | 10 + .../service-bridge/json-rpc-subprocess.ts | 9 +- .../electron/service-bridge/modular-state.ts | 97 +- .../service-bridge/modular-supervisor.ts | 63 ++ .../src/shared/constants/modular-runtime.ts | 10 + desktop/src/shared/types/engine-api.ts | 1 + desktop/src/ui/api/engine-api.ts | 3 + .../ModelManager/IncomingSyncPullRow.tsx | 78 +- .../components/ModelManager/ModelManager.tsx | 28 +- .../ModelManager/TransientModelStatusRow.tsx | 76 +- .../modular/cancel-model-download.test.ts | 73 ++ .../tests/modular/cancel-pull-timeout.test.ts | 217 +++++ .../modular/model-download-controls.test.ts | 70 ++ services/nvpair-engine-manager/README.md | 57 +- services/nvpair-engine-manager/actions.go | 3 + .../nvpair-engine-manager/controlmodels.go | 13 +- .../nvpair-engine-manager/controlserver.go | 1 + .../nvpair-engine-manager/controlstream.go | 24 +- .../controlstream_test.go | 35 +- .../nvpair-engine-manager/download_unix.go | 34 + .../nvpair-engine-manager/download_windows.go | 143 +++ .../download_windows_test.go | 50 + services/nvpair-engine-manager/executor.go | 15 + services/nvpair-engine-manager/lmspull.go | 337 +++++++ services/nvpair-engine-manager/main.go | 3 + services/nvpair-engine-manager/main_test.go | 3 + services/nvpair-engine-manager/manager.go | 50 +- services/nvpair-engine-manager/ollamapull.go | 177 ++++ .../nvpair-engine-manager/ollamapull_test.go | 109 +++ services/nvpair-engine-manager/partials.go | 169 ++++ .../nvpair-engine-manager/partials_test.go | 92 ++ services/nvpair-engine-manager/progress.go | 1 + services/nvpair-engine-manager/pull.go | 78 +- services/nvpair-engine-manager/pull_test.go | 16 +- services/nvpair-engine-manager/pullcancel.go | 217 +++++ .../nvpair-engine-manager/pullcancel_test.go | 876 ++++++++++++++++++ services/nvpair-engine-manager/remote.go | 150 ++- services/nvpair-engine-manager/remote_test.go | 133 +++ .../nvpair-engine-manager/remoteclient.go | 16 +- services/nvpair-engine-manager/spec.md | 12 +- .../testdata/fakeengine/main.go | 61 ++ services/nvpair-tui/ui/engines.go | 120 ++- services/nvpair-tui/ui/engines_test.go | 79 +- 45 files changed, 3752 insertions(+), 172 deletions(-) create mode 100644 desktop/tests/modular/cancel-model-download.test.ts create mode 100644 desktop/tests/modular/cancel-pull-timeout.test.ts create mode 100644 desktop/tests/modular/model-download-controls.test.ts create mode 100644 services/nvpair-engine-manager/download_unix.go create mode 100644 services/nvpair-engine-manager/download_windows.go create mode 100644 services/nvpair-engine-manager/download_windows_test.go create mode 100644 services/nvpair-engine-manager/lmspull.go create mode 100644 services/nvpair-engine-manager/ollamapull.go create mode 100644 services/nvpair-engine-manager/ollamapull_test.go create mode 100644 services/nvpair-engine-manager/partials.go create mode 100644 services/nvpair-engine-manager/partials_test.go create mode 100644 services/nvpair-engine-manager/pullcancel.go create mode 100644 services/nvpair-engine-manager/pullcancel_test.go diff --git a/desktop/docs/services-api.md b/desktop/docs/services-api.md index 7aea40f9..51638ba3 100644 --- a/desktop/docs/services-api.md +++ b/desktop/docs/services-api.md @@ -97,6 +97,7 @@ | `errors:clear` | notification (we consume) | ✅ yes | | `errors:report` | notification (we consume) | ✅ yes | | `engine:action` | request (we call) | ✅ yes | +| `engine:cancel-pull` | request (we call) | ✅ yes | | `engine:configure-launch` | request (we call) | ⚠️ not called | | `engine:configured-ports` | request (we call) | ⚠️ not called | | `engine:describe` | request (we call) | ⚠️ not called | @@ -109,6 +110,7 @@ | `engine:prepare-shutdown` | request (we call) | ✅ yes | | `engine:preview-launch` | request (we call) | ⚠️ not called | | `engine:remote-apply-settings` | request (we call) | ⚠️ not called | +| `engine:remote-cancel-pull` | request (we call) | ✅ yes | | `engine:remote-delete-model` | request (we call) | ✅ yes | | `engine:remote-get-installed` | request (we call) | ✅ yes | | `engine:remote-get-settings` | request (we call) | ⚠️ not called | diff --git a/desktop/docs/services-parity.md b/desktop/docs/services-parity.md index c0b8cbc5..ffd775e7 100644 --- a/desktop/docs/services-parity.md +++ b/desktop/docs/services-parity.md @@ -267,15 +267,116 @@ the engine simply shows as stopped and then running. A local `pull_model` streams live download progress: the engine-manager routes `engine:action{pull_model}` through its streaming pull path and emits -`engine:pull-progress` (`{ engine, op, stage, percent, message }`) — the local +`engine:pull-progress` (`{ engine, model, op, stage, percent, message }`) — the local counterpart of `engine:remote-progress`. Personal AI Router consumes it in `applyLocalEngineProgress` (`modular-supervisor.ts` → `modular-state.ts`), -backfilling the dispatched model (the frame carries none) and advancing the -optimistic pull entry's percent in place, so a local pull shows "Pulling · N%" +keying the frame to its optimistic pull entry by model and advancing that +entry's percent in place, so a local pull shows "Pulling · N%" to completion just like a remote pull. The awaited action response owns -completion (clearing the entry and refreshing the model list); a CLI-driven pull -(LM Studio) emits a single `pulling` marker and degrades to the indeterminate -spinner. +completion (clearing the entry and refreshing the model list). LM Studio's CLI +progress updates the same percentage display. A pull for an engine that already +has one running is reported as `stage:"queued"` until its turn. + +The pull request itself stays pending for the whole download, so it is the +request's own response that reports the outcome — a completion or a +cancellation — while progress arrives out of band on the notification: + +```mermaid +sequenceDiagram + participant UI as Model Manager + participant Bridge as Electron bridge + participant Manager as Engine manager + participant Engine as Ollama or LM Studio + + UI->>Bridge: Download model + Bridge->>UI: Show optimistic "Pulling…" row + Bridge->>Manager: Start pull (request remains pending) + Manager->>Manager: Record the partial files already on disk + Manager->>Engine: Begin download + + loop While downloading + Engine-->>Manager: Download status + Manager-->>Bridge: Progress {engine, model, stage, percent} + Bridge-->>UI: Update the matching model row + end + + alt Download completes + Engine-->>Manager: Success + Manager-->>Bridge: Original pull request completes + else User cancels + UI->>Bridge: Cancel download + Bridge->>UI: Show "Canceling…" + Bridge->>Manager: Cancel {engine, model} + Manager->>Engine: Stop transfer + Manager->>Manager: Remove the partial files this pull created + Manager-->>Bridge: Cancellation completes + Manager-->>Bridge: Original pull settles as canceled + end + + Bridge->>Manager: Refresh model list + Manager-->>Bridge: Authoritative models + Bridge->>UI: Replace temporary download state +``` + +`engine:cancel-pull` and `engine:remote-cancel-pull` cancel the selected model +download. The UI keeps "Canceling…" visible until the pull settles, including +when the cancel RPC's own budget elapses first — the backend is still working, +so the row is not dropped back to "Downloading". Ollama closes the pull request +and removes partial blobs named by its progress digests; LM Studio receives +Ctrl+C and a negative answer to its background-download prompt, and after +acknowledgement the partial files carrying the requested quantization are +removed. Completed model files, other quantizations, and unrelated downloads are +retained, and a cancellation the CLI never confirmed deletes nothing. Only a +cancellation someone requested removes files: a pull interrupted by app +shutdown, a dropped remote connection, or the action timeout leaves its partial +data resumable. + +Both engines funnel every client on the machine into one cache, so naming a file +is not the same as owning it — Ollama blobs are content-addressed, and the LM +Studio app downloads into the same repository directory. Each pull therefore +records the partial files already present before it starts, and cancelling it +considers only the ones that appeared afterwards. A file that is still growing +once the transfer has stopped is shared with another client even by that +measure, so it survives as well: + +```mermaid +flowchart TD + Request["Model pull requested"] --> Tracked["Track by engine + model"] + Tracked --> Busy{"Another PAIR pull active
for this engine?"} + + Busy -- No --> Snapshot["Record the partial files
already on disk"] + Busy -- Yes --> Queued["Show as queued"] + + Queued --> QueueCancel{"Canceled while queued?"} + QueueCancel -- Yes --> NeverStarted["Remove operation
without starting download"] + QueueCancel -- No --> Wait["Wait for active pull to settle"] + Wait --> Snapshot + + Snapshot --> Active["Start download"] + Active --> Cancel{"Cancellation requested?"} + Cancel -- No --> Complete["Download completes normally"] + Cancel -- Yes --> Engine{"Which engine?"} + + Engine -- Ollama --> OllamaStop["Close HTTP pull request"] + OllamaStop --> OllamaCandidates["Select partial blobs named
by this pull's digests"] + + Engine -- LM Studio --> LMSStop["Send Ctrl+C"] + LMSStop --> LMSPrompt["Answer no to background download"] + LMSPrompt --> LMSCandidates[".part files carrying the
requested quantization"] + + OllamaCandidates --> Owned{"Absent from
the snapshot?"} + LMSCandidates --> Owned + Owned -- No --> Preserve["Preserve completed,
shared, and unrelated files"] + Owned -- Yes --> Quiet{"Held still across
repeated checks?"} + Quiet -- No --> Preserve + Quiet -- Yes --> Delete["Remove the partial file"] + + Delete --> Settle["Settle pull and refresh models"] + Preserve --> Settle + Complete --> Settle + NeverStarted --> Settle + Settle --> Next["Allow next queued pull to start"] +``` `engine:models` (and the `em` `GET /v1/models` surface) returns the flat model union, the per-engine breakdown (`modelsByEngine`), and the per-engine set of diff --git a/desktop/src/electron/service-bridge/empty-handlers.ts b/desktop/src/electron/service-bridge/empty-handlers.ts index 14fd3d89..0570cca0 100644 --- a/desktop/src/electron/service-bridge/empty-handlers.ts +++ b/desktop/src/electron/service-bridge/empty-handlers.ts @@ -250,6 +250,11 @@ function routeEngineManagerCommand(payload: WsInvokeRequest<'engine:command'>): void supervisor.pullModel(engine, payload.engineType, payload.model) } break + case 'cancelModelPull': + if (payload.model) { + void supervisor.cancelModelPull(engine, payload.engineType, payload.model) + } + break case 'deleteModel': if (payload.model) { void supervisor.deleteModel(engine, payload.engineType, payload.model) @@ -365,6 +370,11 @@ function routeRemoteEngineCommand(payload: WsInvokeRequest<'engine:command'>): v void supervisor.pullModelRemote(nodeId, engine, payload.engineType, payload.model) } break + case 'cancelModelPull': + if (payload.model) { + void supervisor.cancelModelPull(engine, payload.engineType, payload.model, nodeId) + } + break case 'uninstall': case 'update': refuseRemote( diff --git a/desktop/src/electron/service-bridge/json-rpc-subprocess.ts b/desktop/src/electron/service-bridge/json-rpc-subprocess.ts index 87497015..92b32952 100644 --- a/desktop/src/electron/service-bridge/json-rpc-subprocess.ts +++ b/desktop/src/electron/service-bridge/json-rpc-subprocess.ts @@ -22,6 +22,13 @@ interface JsonRpcError { export class JsonRpcResponseError extends Error {} +/** + * The request's own budget elapsed. The backend has not answered and has not + * failed either — it is still working — so a caller that showed optimistic + * state must decide whether to keep it rather than assume the operation lost. + */ +export class JsonRpcTimeoutError extends Error {} + type JsonRpcId = number | string interface JsonRpcMessage { @@ -148,7 +155,7 @@ export class JsonRpcSubprocess extends EventEmitter { ? undefined : setTimeout(() => { this.pending.delete(id) - reject(new Error(`${this.name} ${method} timed out`)) + reject(new JsonRpcTimeoutError(`${this.name} ${method} timed out`)) }, timeoutMs) this.pending.set(id, { resolve, reject, timeout }) this.write(message).catch(err => { diff --git a/desktop/src/electron/service-bridge/modular-state.ts b/desktop/src/electron/service-bridge/modular-state.ts index 151f1745..02de0825 100644 --- a/desktop/src/electron/service-bridge/modular-state.ts +++ b/desktop/src/electron/service-bridge/modular-state.ts @@ -948,6 +948,7 @@ class ModularBridgeState { * replay it after a UI refresh mid-download. */ private activePulls = new Map() + private cancelingPulls = new Map() /** * Authoritative remote-engine facts keyed by {@link remoteOpKey} * (`${nodeId}:${engineType}`). Populated by the supervisor from each @@ -960,27 +961,6 @@ class ModularBridgeState { * evicted. */ private remoteEngineFacts = new Map() - /** - * The model of an in-flight remote pull keyed by {@link remoteOpKey}. The - * backend's `engine:remote-progress` frames carry **no** `model` field - * (`nvpair-engine-manager/remote.go` `remoteProgress`), so we stamp the model - * we dispatched with here and backfill it onto every incoming frame — that - * keeps the optimistic entry, the streamed percent updates, and the terminal - * clear all on one {@link engineProgressKey}. Cleared by - * {@link finishRemoteModelPull}. - */ - private remotePullModels = new Map() - /** - * The model of an in-flight local pull keyed by `engineType`. The backend's - * `engine:pull-progress` frames (the local counterpart of - * `engine:remote-progress`) carry **no** `model` field - * (`nvpair-engine-manager/executor.go` `emitPullProgress`), so we stamp the - * model we dispatched here and backfill it onto every incoming frame — one - * active local pull per engine, matching the backend's engine-scoped - * progress hub. Set by {@link beginModelPull}, cleared by - * {@link finishModelPull}. - */ - private localPullModels = new Map() /** * The authoritative set of live inbound invites awaiting the local user's PIN * entry, keyed by `inviteId`. The cluster-manager pushes `cluster:invite-received` @@ -1492,9 +1472,6 @@ class ModularBridgeState { model, status: 'pulling' } - // Remember the model so the model-less `engine:remote-progress` frames can - // be re-associated with this exact entry (see remotePullModels). - this.remotePullModels.set(this.remoteOpKey(nodeId, engineType), model) this.activePulls.set(engineProgressKey(progress), progress) emitBridgePush('engines:progress-changed', progress) } @@ -1515,8 +1492,8 @@ class ModularBridgeState { * is the only signal that removes the spinner. */ finishRemoteModelPull(nodeId: string, engineType: EngineType, model: string): void { - this.remotePullModels.delete(this.remoteOpKey(nodeId, engineType)) const key = engineProgressKey({ nodeId, engineType, operation: 'pull', model }) + this.cancelingPulls.delete(key) if (!this.activePulls.delete(key)) return emitBridgePush('engines:progress-cleared', { key }) } @@ -1531,14 +1508,11 @@ class ModularBridgeState { const op = stringValue(obj.op) const operation = op === 'pull' || op === 'pull_model' ? 'pull' : 'install' const stage = stringValue(obj.stage) || 'working' - // `engine:remote-progress` carries no `model`, so for a pull we backfill - // the model captured at dispatch — otherwise the frame would emit under a - // different key than the optimistic entry and never update/clear it. - const model = - stringValue(obj.model) || - (operation === 'pull' - ? (this.remotePullModels.get(this.remoteOpKey(nodeId, engineType)) ?? '') - : '') + // The peer stamps the model on every pull frame it relays + // (`nvpair-engine-manager/remote.go` `remoteProgressFn`), which is what + // keys a frame to its optimistic entry. An install carries none, and + // needs none: its progress is tracked per engine. + const model = stringValue(obj.model) if (stage === 'done' || stage === 'already-installed' || stage === 'failed') { const progressKey = @@ -1546,13 +1520,10 @@ class ModularBridgeState { ? engineProgressKey({ nodeId, engineType, operation: 'pull', model }) : `${nodeId}:${engineType}:${operation}` emitBridgePush('engines:progress-cleared', { key: progressKey }) - if (operation === 'pull') { - this.remotePullModels.delete(this.remoteOpKey(nodeId, engineType)) - if (model) { - this.activePulls.delete( - engineProgressKey({ nodeId, engineType, operation: 'pull', model }) - ) - } + if (operation === 'pull' && model) { + this.activePulls.delete( + engineProgressKey({ nodeId, engineType, operation: 'pull', model }) + ) } if (stage === 'already-installed' || stage === 'failed') { this.clearRemoteEngineOp(nodeId, engineType) @@ -1590,7 +1561,7 @@ class ModularBridgeState { nodeId, nodeName: node?.name ?? nodeId, operation, - status: stage, + status: this.cancelingPulls.has(pullKey) ? 'canceling' : stage, percent, model: model || undefined } @@ -1820,7 +1791,6 @@ class ModularBridgeState { private dropRemoteEngineFacts(nodeId: string): void { for (const engine of PROXY_ENGINES) { this.remoteEngineFacts.delete(this.remoteOpKey(nodeId, engine)) - this.remotePullModels.delete(this.remoteOpKey(nodeId, engine)) } } @@ -2180,9 +2150,6 @@ class ModularBridgeState { model, status: 'pulling' } - // Remember the model so the model-less `engine:pull-progress` frames can - // be re-associated with this exact entry (see localPullModels). - this.localPullModels.set(engineType, model) this.activePulls.set(engineProgressKey(progress), progress) emitBridgePush('engines:progress-changed', progress) } @@ -2191,12 +2158,13 @@ class ModularBridgeState { * Project a `nvpair-engine-manager` `engine:pull-progress` for the local node * — the local counterpart of `engine:remote-progress`: live * download progress for a model pull driven via `engine:action` - * `pull_model`. The frame carries `{ engine, op, stage, percent, message }` - * and **no** `model`, so we backfill the model captured at dispatch and - * refresh the optimistic entry in place. Clearing stays with the awaited pull - * ({@link finishModelPull}), so this only advances percent/stage: it ignores - * the terminal error sentinel (`percent: -1`) and the non-download frames' - * `percent: 0` so the rendered bar never jumps backwards. + * `pull_model`. The frame carries `{ engine, model, op, stage, percent, + * message }` and its `model` identifies which optimistic entry to refresh — + * the backend queues concurrent pulls per engine, so the engine alone does + * not. Clearing stays with the awaited pull ({@link finishModelPull}), so + * this only advances percent/stage: it ignores the terminal error sentinel + * (`percent: -1`) and the non-download frames' `percent: 0` so the rendered + * bar never jumps backwards. */ applyLocalEngineProgress(params: JsonValue | undefined): void { const nodeId = this.selfId @@ -2206,7 +2174,7 @@ class ModularBridgeState { const engineType = engineTypeFromManagerName(stringValue(obj.engine)) if (!engineType) return - const model = this.localPullModels.get(engineType) + const model = stringValue(obj.model) if (!model) return const key = engineProgressKey({ nodeId, engineType, operation: 'pull', model }) const existing = this.activePulls.get(key) @@ -2215,7 +2183,9 @@ class ModularBridgeState { const percent = numberValue(obj.percent) const progress: EngineProgress = { ...existing, - status: stringValue(obj.stage) || existing.status, + status: this.cancelingPulls.has(key) + ? 'canceling' + : stringValue(obj.stage) || existing.status, percent: mergePullProgressPercent(percent, existing.percent) } this.activePulls.set(key, progress) @@ -2235,12 +2205,31 @@ class ModularBridgeState { finishModelPull(engineType: EngineType, model: string): void { const nodeId = this.selfId if (!nodeId) return - this.localPullModels.delete(engineType) const key = engineProgressKey({ nodeId, engineType, operation: 'pull', model }) + this.cancelingPulls.delete(key) if (!this.activePulls.delete(key)) return emitBridgePush('engines:progress-cleared', { key }) } + setModelPullCanceling( + engineType: EngineType, + model: string, + canceling: boolean, + nodeId = this.selfId + ): boolean { + if (!nodeId) return false + const key = engineProgressKey({ nodeId, engineType, operation: 'pull', model }) + const existing = this.activePulls.get(key) + if (!existing || (canceling && this.cancelingPulls.has(key))) return false + const status = canceling ? 'canceling' : (this.cancelingPulls.get(key) ?? existing.status) + if (canceling) this.cancelingPulls.set(key, existing.status) + else this.cancelingPulls.delete(key) + const progress = { ...existing, status } + this.activePulls.set(key, progress) + emitBridgePush('engines:progress-changed', progress) + return true + } + getLogs(): LogEntry[] { return [...this.logs] } diff --git a/desktop/src/electron/service-bridge/modular-supervisor.ts b/desktop/src/electron/service-bridge/modular-supervisor.ts index 8cd1012b..0b372555 100644 --- a/desktop/src/electron/service-bridge/modular-supervisor.ts +++ b/desktop/src/electron/service-bridge/modular-supervisor.ts @@ -7,6 +7,7 @@ import { app } from 'electron' import { JsonRpcResponseError, JsonRpcSubprocess, + JsonRpcTimeoutError, type JsonObject, type JsonRpcInboundRequest, type JsonRpcNotification, @@ -34,6 +35,7 @@ import { isFirstRun } from '@/electron/config/ui-config' import { parseClusterNodes, parseInvite, parseNodeIdentity } from './cluster-json' import { startNodeInfoPoller, stopNodeInfoPoller } from './node-info-poller' import { + MODULAR_CANCEL_PULL_TIMEOUT_MS, MODULAR_DEFAULT_LOG_LEVEL, MODULAR_INVITE_STATUS_POLL_INTERVAL_MS, MODULAR_MODEL_ACTION_TIMEOUT_MS, @@ -1724,6 +1726,67 @@ class ModularSupervisor { } } + /** + * Stop a model download. The backend acknowledges only after the transfer + * has stopped and its partial files are cleaned up, so the row stays in + * "Canceling" for the whole of that — see + * {@link MODULAR_CANCEL_PULL_TIMEOUT_MS} for what bounds it. + * + * A timeout is not a failure: the backend is still cancelling, and dropping + * the row back to "downloading" seconds before it finishes would be a lie + * the user acts on. Only a real rejection restores the previous status, and + * either way the pull's own settling event clears the entry. + */ + async cancelModelPull( + engine: string, + engineType: EngineType, + model: string, + nodeId?: string + ): Promise { + const state = getModularBridgeState() + const active = nodeId + ? state.isRemoteModelPullActive(nodeId, engineType, model) + : state.isModelPullActive(engineType, model) + if (!active) return + if (!state.setModelPullCanceling(engineType, model, true, nodeId)) return + try { + if (nodeId) { + await this.callProcess( + 'broker', + 'engine:remote-cancel-pull', + { + node: nodeId, + engine, + model + }, + MODULAR_CANCEL_PULL_TIMEOUT_MS + ) + } else { + await this.callProcess( + 'broker', + 'engine:cancel-pull', + { engine, model }, + MODULAR_CANCEL_PULL_TIMEOUT_MS + ) + } + } catch (err) { + if (err instanceof JsonRpcTimeoutError) { + log.warn({ + sublevel: 'broker', + message: `cancel of ${engine} download is still running after ${MODULAR_CANCEL_PULL_TIMEOUT_MS}ms` + }) + return + } + state.setModelPullCanceling(engineType, model, false, nodeId) + this.reportError( + `Failed to cancel download of ${model}: ${getErrorString(err)}`, + 'error', + `engine-cancel-pull:${nodeId ?? 'local'}:${engine}:${model}`, + { engineType, nodeId, operation: 'pull', modelName: model } + ) + } + } + /** * Delete a downloaded model locally. Awaits the backend action, refreshes * `list_models`, and surfaces RPC failures so the optimistic spinner clears diff --git a/desktop/src/shared/constants/modular-runtime.ts b/desktop/src/shared/constants/modular-runtime.ts index 412827c6..6d46b63c 100644 --- a/desktop/src/shared/constants/modular-runtime.ts +++ b/desktop/src/shared/constants/modular-runtime.ts @@ -79,6 +79,16 @@ export const MODULAR_ENGINE_LIFECYCLE_CALL_TIMEOUT_MS = 14 * 60_000 // docs/services-parity.md. export const MODULAR_MODEL_ACTION_TIMEOUT_MS = 120_000 +// Upper bound for `engine:cancel-pull` / `engine:remote-cancel-pull`. The +// backend answers only once the transfer has stopped and its partial files are +// settled: LM Studio gets two interrupt windows before its CLI is reclaimed, +// and Ollama retries blob removal while its writers release. The remote variant +// adds the peer's own `remoteResponseHeaderTimeout` +// (services/nvpair-engine-manager/remoteclient.go), which starts before the +// peer writes a header — so this has to sit above that budget, not equal it, or +// the desktop reports a failure while the cancel is still running. +export const MODULAR_CANCEL_PULL_TIMEOUT_MS = 90_000 + // Poll interval for `cluster:invite-status` while a pairing handshake is open. // Also drives the pending-invite reconciliation sweep, the backstop for a missed // receiver-side `cluster:invite-canceled` / `cluster:invite-expired` push. diff --git a/desktop/src/shared/types/engine-api.ts b/desktop/src/shared/types/engine-api.ts index 9187e658..4ac88385 100644 --- a/desktop/src/shared/types/engine-api.ts +++ b/desktop/src/shared/types/engine-api.ts @@ -45,6 +45,7 @@ export type EngineCommandType = | 'uninstall' | 'update' | 'pullModel' + | 'cancelModelPull' | 'loadModel' | 'unloadModel' | 'deleteModel' diff --git a/desktop/src/ui/api/engine-api.ts b/desktop/src/ui/api/engine-api.ts index 71d8f163..d3ad48dc 100644 --- a/desktop/src/ui/api/engine-api.ts +++ b/desktop/src/ui/api/engine-api.ts @@ -54,6 +54,7 @@ export interface IEngineApi { uninstall(engineType: EngineType, nodeId: string): void /** Pull (download) a model on a node. */ pullModel(engineType: EngineType, nodeId: string, model: string): void + cancelModelPull(engineType: EngineType, nodeId: string, model: string): void /** Load a model into memory on a node. */ loadModel(engineType: EngineType, nodeId: string, model: string): void /** Unload a model from memory on a node. */ @@ -100,6 +101,8 @@ export function createEngineApi(transport: ServiceTransport): IEngineApi { fireCommand(transport, { command: 'uninstall', engineType, nodeId }), pullModel: (engineType, nodeId, model) => fireCommand(transport, { command: 'pullModel', engineType, nodeId, model }), + cancelModelPull: (engineType, nodeId, model) => + fireCommand(transport, { command: 'cancelModelPull', engineType, nodeId, model }), loadModel: (engineType, nodeId, model) => fireCommand(transport, { command: 'loadModel', engineType, nodeId, model }), unloadModel: (engineType, nodeId, model) => diff --git a/desktop/src/ui/components/ModelManager/IncomingSyncPullRow.tsx b/desktop/src/ui/components/ModelManager/IncomingSyncPullRow.tsx index 428f40ec..52cea3eb 100644 --- a/desktop/src/ui/components/ModelManager/IncomingSyncPullRow.tsx +++ b/desktop/src/ui/components/ModelManager/IncomingSyncPullRow.tsx @@ -1,38 +1,62 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -import { Flex, Text } from '@nvidia/foundations-react-core' +import { Button, Flex, ProgressBar, Stack, Text } from '@nvidia/foundations-react-core' import { formatPullProgressLabel } from '@/ui/utils/formatters' import type { IncomingSyncRow } from '@/ui/types/model-manager' -export function IncomingSyncPullRow({ row }: { row: IncomingSyncRow }) { +export function IncomingSyncPullRow({ + row, + onCancel +}: { + row: IncomingSyncRow + onCancel: () => void +}) { + const canceling = row.status === 'canceling' return ( - - - - - - {row.label} - - - {formatPullProgressLabel(row)} - + + + + + + + {row.label} + + + {formatPullProgressLabel(row)} + + + - + + ) } diff --git a/desktop/src/ui/components/ModelManager/ModelManager.tsx b/desktop/src/ui/components/ModelManager/ModelManager.tsx index e63dacf2..cfd93867 100644 --- a/desktop/src/ui/components/ModelManager/ModelManager.tsx +++ b/desktop/src/ui/components/ModelManager/ModelManager.tsx @@ -70,7 +70,7 @@ export function ModelManager({ backend, nodeId }: { backend: BackendInfo; nodeId const isBusy = models.some( m => m.status === 'loading' || m.status === 'ejecting' || m.status === 'pulling' ) - const transientModel = models.find( + const transientModels = models.filter( m => m.status === 'loading' || m.status === 'ejecting' || m.status === 'pulling' ) @@ -151,10 +151,6 @@ export function ModelManager({ backend, nodeId }: { backend: BackendInfo; nodeId [backendType, nodeId] ) - const transientPullProgress = transientModel - ? getProgress(nodeId, backend.type, 'pull', transientModel.name) - : undefined - const hasModelSearchOnlyWhenRunning = useMemo( () => caps?.hasModelSearchOnlyWhenRunning ?? false, [caps?.hasModelSearchOnlyWhenRunning] @@ -186,16 +182,30 @@ export function ModelManager({ backend, nodeId }: { backend: BackendInfo; nodeId return ( - {transientModel && ( + {transientModels.map(transientModel => ( + window.pairApi.engines.cancelModelPull( + backendType, + nodeId, + transientModel.name + ) + } /> - )} + ))} {incomingSyncs.map(p => ( - + + window.pairApi.engines.cancelModelPull(backendType, nodeId, p.rawModel) + } + /> ))} {!isBusy && ( diff --git a/desktop/src/ui/components/ModelManager/TransientModelStatusRow.tsx b/desktop/src/ui/components/ModelManager/TransientModelStatusRow.tsx index 6f224f2c..b363f04f 100644 --- a/desktop/src/ui/components/ModelManager/TransientModelStatusRow.tsx +++ b/desktop/src/ui/components/ModelManager/TransientModelStatusRow.tsx @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -import { Flex, Text } from '@nvidia/foundations-react-core' +import { Button, Flex, ProgressBar, Stack, Text } from '@nvidia/foundations-react-core' import type { ModelItem } from '@/ui/types/engine-info' import { formatPullProgressLabel } from '@/ui/utils/formatters' import type { EngineProgress } from '@/shared/types/engines' @@ -9,34 +9,64 @@ import type { EngineProgress } from '@/shared/types/engines' export function TransientModelStatusRow({ transientModel, displayName, - pullProgress + pullProgress, + onCancel }: { transientModel: ModelItem displayName: (name: string) => string pullProgress: EngineProgress | undefined + onCancel: () => void }) { + const pulling = transientModel.status === 'pulling' + const canceling = pullProgress?.status === 'canceling' + const percent = pullProgress?.percent return ( - - - - - {transientModel.status === 'loading' && - `Loading ${displayName(transientModel.name)}...`} - {transientModel.status === 'ejecting' && - `Ejecting ${displayName(transientModel.name)}...`} - {transientModel.status === 'pulling' && - `Pulling ${displayName(transientModel.name)}${ - pullProgress && pullProgress.status !== 'idle' - ? ` · ${formatPullProgressLabel(pullProgress)}` - : '...' - }`} - + + + + + + {transientModel.status === 'loading' && + `Loading ${displayName(transientModel.name)}...`} + {transientModel.status === 'ejecting' && + `Ejecting ${displayName(transientModel.name)}...`} + {transientModel.status === 'pulling' && + `Pulling ${displayName(transientModel.name)}${ + pullProgress && pullProgress.status !== 'idle' + ? ` · ${formatPullProgressLabel(pullProgress)}` + : '...' + }`} + + + {pulling && ( + + )} - + {pulling && ( + + )} + ) } diff --git a/desktop/tests/modular/cancel-model-download.test.ts b/desktop/tests/modular/cancel-model-download.test.ts new file mode 100644 index 00000000..7bbae609 --- /dev/null +++ b/desktop/tests/modular/cancel-model-download.test.ts @@ -0,0 +1,73 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { afterEach, describe, expect, it, vi } from 'vitest' + +vi.mock('electron', () => ({ + BrowserWindow: { getAllWindows: () => [] } +})) +vi.mock('@/electron/window', () => ({ createOverviewWindow: vi.fn() })) + +import { getModularBridgeState } from '@/electron/service-bridge/modular-state' +import { subscribePush } from '@/electron/service-bridge/push-bus' +import type { EngineProgress } from '@/shared/types/engines' + +let unsubscribe: (() => void) | undefined +afterEach(() => { + unsubscribe?.() + const state = getModularBridgeState() + state.finishRemoteModelPull('peer', 'ollama', 'demo') + state.finishRemoteModelPull('peer', 'lm-studio', 'owner/model') +}) + +describe('model download cancellation', () => { + it('keeps Canceling visible across late progress and rejects duplicate cancellation', () => { + const state = getModularBridgeState() + const events: EngineProgress[] = [] + unsubscribe = subscribePush(event => { + if (event.channel === 'engines:progress-changed') events.push(event.payload) + }) + state.beginRemoteModelPull('peer', 'lm-studio', 'owner/model') + expect(state.setModelPullCanceling('lm-studio', 'owner/model', true, 'peer')).toBe(true) + expect(state.setModelPullCanceling('lm-studio', 'owner/model', true, 'peer')).toBe(false) + state.applyRemoteEngineProgress({ + node: 'peer', + engine: 'lmstudio', + model: 'owner/model', + op: 'pull', + stage: 'downloading', + percent: 35 + }) + expect(events.at(-1)).toMatchObject({ status: 'canceling', percent: 35 }) + }) + + it('restores the previous status when cancellation fails', () => { + const state = getModularBridgeState() + const events: EngineProgress[] = [] + unsubscribe = subscribePush(event => { + if (event.channel === 'engines:progress-changed') events.push(event.payload) + }) + state.beginRemoteModelPull('peer', 'ollama', 'demo') + state.applyRemoteEngineProgress({ + node: 'peer', + engine: 'ollama', + model: 'demo', + op: 'pull', + stage: 'downloading', + percent: 25 + }) + state.setModelPullCanceling('ollama', 'demo', true, 'peer') + state.setModelPullCanceling('ollama', 'demo', false, 'peer') + expect(events.at(-1)).toMatchObject({ status: 'downloading', percent: 25 }) + }) + + it('clears cancellation state when the pull finishes and allows retry', () => { + const state = getModularBridgeState() + state.beginRemoteModelPull('peer', 'ollama', 'demo') + state.setModelPullCanceling('ollama', 'demo', true, 'peer') + state.finishRemoteModelPull('peer', 'ollama', 'demo') + expect(state.isRemoteModelPullActive('peer', 'ollama', 'demo')).toBe(false) + state.beginRemoteModelPull('peer', 'ollama', 'demo') + expect(state.setModelPullCanceling('ollama', 'demo', true, 'peer')).toBe(true) + }) +}) diff --git a/desktop/tests/modular/cancel-pull-timeout.test.ts b/desktop/tests/modular/cancel-pull-timeout.test.ts new file mode 100644 index 00000000..e803c025 --- /dev/null +++ b/desktop/tests/modular/cancel-pull-timeout.test.ts @@ -0,0 +1,217 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { JsonRpcTimeoutError } from '@/electron/service-bridge/json-rpc-subprocess' +import { MODULAR_CANCEL_PULL_TIMEOUT_MS } from '@/shared/constants/modular-runtime' + +const mocks = vi.hoisted(() => ({ + bridgeState: { + handleNotification: vi.fn(), + getSelfId: vi.fn(() => null), + getProxyPort: vi.fn(() => null), + isModelPullActive: vi.fn(() => true), + isRemoteModelPullActive: vi.fn(() => true), + setModelPullCanceling: vi.fn(() => true) + }, + emitBridgePush: vi.fn() +})) + +vi.mock('electron', () => ({ + app: { + isPackaged: false, + getAppPath: () => process.cwd() + } +})) + +vi.mock('@/shared/utils/log', () => ({ + createStructuredLogger: () => ({ + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + verbose: vi.fn() + }) +})) + +vi.mock('@/electron/config/ui-config', () => ({ + isFirstRun: () => false +})) + +vi.mock('@/electron/service-bridge/broadcaster', () => ({ + emitBridgePush: mocks.emitBridgePush +})) + +vi.mock('@/electron/service-bridge/manual-nodes-store', () => ({ + listManualNodeEntries: () => [] +})) + +vi.mock('@/electron/service-bridge/node-info-poller', () => ({ + startNodeInfoPoller: vi.fn(), + stopNodeInfoPoller: vi.fn() +})) + +vi.mock('@/electron/service-bridge/modular-state', () => ({ + getModularBridgeState: () => mocks.bridgeState, + isUpstreamUnreachableError: () => false, + parseServiceErrors: () => [], + parseWorkloadsInitial: () => [], + PROXY_ENGINES: ['ollama', 'lm-studio'], + PROXY_NODE_SOURCES: ['ollama-proxy', 'lmstudio-proxy'] +})) + +import { getModularSupervisor } from '@/electron/service-bridge/modular-supervisor' + +const supervisor = getModularSupervisor() + +describe('cancelling a model download', () => { + beforeEach(() => { + vi.restoreAllMocks() + mocks.bridgeState.setModelPullCanceling.mockClear().mockReturnValue(true) + mocks.bridgeState.isModelPullActive.mockReturnValue(true) + mocks.bridgeState.isRemoteModelPullActive.mockReturnValue(true) + }) + + it('gives the backend a budget above the peer response-header timeout', async () => { + const call = vi.spyOn(supervisor, 'callProcess').mockResolvedValue(null) + + await supervisor.cancelModelPull('ollama', 'ollama', 'demo') + + expect(call).toHaveBeenCalledWith( + 'broker', + 'engine:cancel-pull', + { engine: 'ollama', model: 'demo' }, + MODULAR_CANCEL_PULL_TIMEOUT_MS + ) + // The peer's own header budget is 30s and starts before it writes one, + // so an outer budget that merely matched it would expire first. + expect(MODULAR_CANCEL_PULL_TIMEOUT_MS).toBeGreaterThan(30_000) + }) + + // The backend is still stopping the transfer and cleaning up after it. Its + // own settling event clears the row; dropping it back to "downloading" here + // would tell the user the cancel failed seconds before it succeeds. + it('leaves the row in Canceling when the request outlives its own budget', async () => { + const reportError = vi.spyOn(supervisor, 'reportError').mockImplementation(() => {}) + vi.spyOn(supervisor, 'callProcess').mockRejectedValue( + new JsonRpcTimeoutError('broker engine:cancel-pull timed out') + ) + + await supervisor.cancelModelPull('ollama', 'ollama', 'demo') + + expect(mocks.bridgeState.setModelPullCanceling).toHaveBeenCalledTimes(1) + expect(mocks.bridgeState.setModelPullCanceling).toHaveBeenCalledWith( + 'ollama', + 'demo', + true, + undefined + ) + expect(reportError).not.toHaveBeenCalled() + }) + + it('restores the previous status when the backend rejects the cancel', async () => { + const reportError = vi.spyOn(supervisor, 'reportError').mockImplementation(() => {}) + vi.spyOn(supervisor, 'callProcess').mockRejectedValue( + new Error('engine ollama is not running') + ) + + await supervisor.cancelModelPull('ollama', 'ollama', 'demo') + + expect(mocks.bridgeState.setModelPullCanceling).toHaveBeenCalledWith( + 'ollama', + 'demo', + false, + undefined + ) + expect(reportError).toHaveBeenCalledWith( + expect.stringContaining('Failed to cancel download of demo'), + 'error', + 'engine-cancel-pull:local:ollama:demo', + expect.objectContaining({ engineType: 'ollama', operation: 'pull', modelName: 'demo' }) + ) + }) + + // A download on a peer is cancelled through the initiating node's + // engine-manager, which relays it over the cluster's `ec` surface. The row + // and its budget are the local ones either way, so only the method, the + // node it names, and the error's own key change. + it('routes a peer node cancellation through the remote method', async () => { + const call = vi.spyOn(supervisor, 'callProcess').mockResolvedValue(null) + + await supervisor.cancelModelPull('ollama', 'ollama', 'demo', 'uuid-b') + + expect(call).toHaveBeenCalledWith( + 'broker', + 'engine:remote-cancel-pull', + { node: 'uuid-b', engine: 'ollama', model: 'demo' }, + MODULAR_CANCEL_PULL_TIMEOUT_MS + ) + expect(mocks.bridgeState.setModelPullCanceling).toHaveBeenCalledWith( + 'ollama', + 'demo', + true, + 'uuid-b' + ) + }) + + it('leaves a peer node row in Canceling when the request outlives its budget', async () => { + const reportError = vi.spyOn(supervisor, 'reportError').mockImplementation(() => {}) + vi.spyOn(supervisor, 'callProcess').mockRejectedValue( + new JsonRpcTimeoutError('broker engine:remote-cancel-pull timed out') + ) + + await supervisor.cancelModelPull('ollama', 'ollama', 'demo', 'uuid-b') + + expect(mocks.bridgeState.setModelPullCanceling).toHaveBeenCalledTimes(1) + expect(mocks.bridgeState.setModelPullCanceling).toHaveBeenCalledWith( + 'ollama', + 'demo', + true, + 'uuid-b' + ) + expect(reportError).not.toHaveBeenCalled() + }) + + // The error key carries the node, so two peers downloading the same model + // report separately instead of overwriting one another's row. + it('restores the peer node row and keys its error by node when the cancel is rejected', async () => { + const reportError = vi.spyOn(supervisor, 'reportError').mockImplementation(() => {}) + vi.spyOn(supervisor, 'callProcess').mockRejectedValue( + new Error('node uuid-b is not a discovered ec peer') + ) + + await supervisor.cancelModelPull('ollama', 'ollama', 'demo', 'uuid-b') + + expect(mocks.bridgeState.setModelPullCanceling).toHaveBeenCalledWith( + 'ollama', + 'demo', + false, + 'uuid-b' + ) + expect(reportError).toHaveBeenCalledWith( + expect.stringContaining('Failed to cancel download of demo'), + 'error', + 'engine-cancel-pull:uuid-b:ollama:demo', + expect.objectContaining({ + engineType: 'ollama', + nodeId: 'uuid-b', + operation: 'pull', + modelName: 'demo' + }) + ) + }) + + // Each path checks its own registry before sending anything: a stale click + // on a row whose download already settled must not reach the backend, and + // must not put the row back into "Canceling". + it('sends nothing when the download it names is no longer active', async () => { + const call = vi.spyOn(supervisor, 'callProcess').mockResolvedValue(null) + mocks.bridgeState.isModelPullActive.mockReturnValue(false) + mocks.bridgeState.isRemoteModelPullActive.mockReturnValue(false) + + await supervisor.cancelModelPull('ollama', 'ollama', 'demo') + await supervisor.cancelModelPull('ollama', 'ollama', 'demo', 'uuid-b') + + expect(call).not.toHaveBeenCalled() + expect(mocks.bridgeState.setModelPullCanceling).not.toHaveBeenCalled() + }) +}) diff --git a/desktop/tests/modular/model-download-controls.test.ts b/desktop/tests/modular/model-download-controls.test.ts new file mode 100644 index 00000000..67cd9b20 --- /dev/null +++ b/desktop/tests/modular/model-download-controls.test.ts @@ -0,0 +1,70 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { createElement } from 'react' +import { renderToStaticMarkup } from 'react-dom/server' +import { describe, expect, it } from 'vitest' +import { IncomingSyncPullRow } from '@/ui/components/ModelManager/IncomingSyncPullRow' +import { TransientModelStatusRow } from '@/ui/components/ModelManager/TransientModelStatusRow' +import type { ModelItem } from '@/ui/types/engine-info' + +const model: ModelItem = { + name: 'demo', + size: 0, + downloaded: false, + status: 'pulling', + parameterSize: '', + quantization: '', + family: '', + digest: '', + sizeVram: null, + expiresAt: null, + expiry: '10m', + capabilities: [] +} + +const rows = [ + { + name: 'incoming download', + render: (status: string) => + createElement(IncomingSyncPullRow, { + row: { rawModel: 'demo', label: 'Demo', status }, + onCancel: () => {} + }) + }, + { + name: 'transient model', + render: (status: string) => + createElement(TransientModelStatusRow, { + transientModel: model, + displayName: name => name, + pullProgress: { + engineType: 'ollama', + nodeId: 'local', + nodeName: 'Local', + operation: 'pull', + status + }, + onCancel: () => {} + }) + } +] + +for (const row of rows) { + describe(row.name, () => { + it.each([ + { status: 'pulling', disabled: false }, + { status: 'queued', disabled: false }, + { status: 'downloading', disabled: false }, + { status: 'canceling', disabled: true }, + { status: 'idle', disabled: true }, + { status: 'error', disabled: true } + ])('cancellation availability during $status', ({ status, disabled }) => { + const html = renderToStaticMarkup(row.render(status)) + const button = html.match(/]*>[\s\S]*?<\/button>/)?.[0] + expect(button).toBeDefined() + expect(button).toContain(status === 'canceling' ? 'Canceling…' : 'Cancel download') + expect(/]*\bdisabled(?:\s|=|>)/.test(button!)).toBe(disabled) + }) + }) +} diff --git a/services/nvpair-engine-manager/README.md b/services/nvpair-engine-manager/README.md index b401c883..e200a4fd 100644 --- a/services/nvpair-engine-manager/README.md +++ b/services/nvpair-engine-manager/README.md @@ -40,12 +40,14 @@ Requests (caller → service): | `engine:restart` | `{ engine }` | `EngineStatus` | | `engine:set-port` | `{ engine, port }` | `EngineStatus` (after rebind) | | `engine:action` | `{ engine, action, params }` | the engine's raw response. `action:"pull_model"` is streamed: it emits live `engine:pull-progress` notifications and returns the pull's terminal result (see below). An action whose manifest declares `restart_after` (LM Studio's `delete_model`) restarts a running engine before replying, so the response also means the engine is back and healthy | +| `engine:cancel-pull` | `{ engine, model }` | `null`, once the transfer has stopped **and** its partial files are settled — so a UI can hold "Canceling" until the backend is genuinely done. See "Cancelling a pull" below | | `engine:logs` | `{ engine }` | `{ lines: [LogLine] }` | | `engine:errors` | — | `{ errors: [ServiceError] }` | | `engine:models` | — | `{ models: [string], modelsByEngine: { : [string] }, loadedByEngine: { : [string] } }` — the flat de-duplicated union of every running engine's models, the per-engine breakdown keyed by engine name, and the per-engine set of models currently **loaded in memory** (all normalized from each engine's `list_models` / `loaded_models` action `result` spec). `modelsByEngine` carries a key for every running engine whose inventory was successfully queried, including an empty list = "running, no models available"; a missing key means not running / not queryable / invalid response. `loadedByEngine` uses the same known-empty distinction for residency and also omits engines with no loaded endpoint. The `/v1/models` HTTP surface returns the same shape. | | `engine:remote-get-installed` | `{ node }` | `{ engines: [EngineStatus] }` fetched from the remote node over `ec` mTLS | | `engine:remote-install` | `{ node, engine, start? }` | `{ opId, status: EngineStatus }` after the remote install (live progress via `engine:remote-progress`) | | `engine:remote-pull-model` | `{ node, engine, model?, params? }` | `{ opId, result }` after the remote pull (live progress via `engine:remote-progress`) | +| `engine:remote-cancel-pull` | `{ node, engine, model }` | the remote cancellation result. Held until the peer has accepted the matching `engine:remote-pull-model`, so a cancel dispatched right behind a pull cannot reach the peer first | | `engine:remote-start` | `{ node, engine, port? }` | `EngineStatus` from the remote node (always the manifest's `runtime.bind`; no per-call bind override on the remote path) | | `engine:remote-stop` | `{ node, engine }` | `EngineStatus` from the remote node | | `shutdown` | — | `null` | @@ -60,14 +62,19 @@ loaded (in-memory) models changes (explicit load/unload, JIT auto-load, or TTL/idle eviction); `models` is the full `engine:models` shape (incl. `loadedByEngine`) so a consumer swaps its whole snapshot, `engine:install-progress{engine, stage, percent}`, -`engine:pull-progress{engine, op, stage, percent, message}` (live progress for a -local model pull driven via `engine:action{action:"pull_model"}` — the local -counterpart of `engine:remote-progress`; frames are coalesced to changes in +`engine:pull-progress{engine, model, op, stage, percent, message}` (live progress +for a local model pull driven via `engine:action{action:"pull_model"}` — the +local counterpart of `engine:remote-progress`; frames are coalesced to changes in stage/percent, the engine's terminal success surfaces as `stage:"success"`, and a failed pull emits a terminal `stage:"error", percent:-1, message` frame so a -UI converges even if its synchronous call already timed out), -`engine:remote-progress{opId, node, engine, op, stage, percent, message}` -(relayed live progress for a remote install/pull), and — for the error +UI converges even if its synchronous call already timed out. `model` names the +download the frame belongs to, so a consumer can tell two concurrent pulls on +one engine apart; a pull waiting on another of the same engine's downloads +reports `stage:"queued"` until its turn), +`engine:remote-progress{opId, node, engine, model, op, stage, percent, message}` +(relayed live progress for a remote install/pull; `model` is carried through +from the peer's frame for the same reason it appears on the local +notification), and — for the error pipeline — `errors:report` / `errors:clear` (consumed by `nvpair-errors` via the broker; see below). @@ -182,6 +189,43 @@ external app first and re-starting on its usual port works too. Auto-assigned ports (manifest `runtime.port: 0`) never adopt — there's no fixed address to probe — so they always spawn an owned process. +### Cancelling a pull + +`engine:cancel-pull` answers only once the transfer has stopped and its partial +files are settled, because a UI reads that answer as "it is safe to download +this again". The pull's own `engine:action{pull_model}` request settles +separately, as `{"status":"cancelled"}`. + +**Only a cancellation someone asked for deletes anything.** A pull's context +also dies on shutdown, on a dropped remote connection, and on the action +timeout, and none of those mean the user gave up on the bytes already on disk — +those leave the partial data for the next attempt to resume. + +**Naming a file is not owning it.** Both vendors funnel every client on the +machine into one cache: Ollama coalesces a pull of the same layer — from its +CLI, its desktop app, another engine-manager, any client on the daemon — onto +the very same content-addressed partial, and the LM Studio app downloads into +the same repository directory. So each pull records the partial files already +present before it starts, and cancelling it considers only the ones that +appeared afterwards. A file that is still growing once the transfer has stopped +is shared by that measure too, and also survives; cleanup retries within a +bounded budget and leaves anything still busy at the deadline, which is not a +failure. Completed model files are never candidates — they may be layers of an +installed model. + +LM Studio additionally requires the CLI to confirm: declining its +background-download prompt is what aborts the daemon's task, so a CLI that +exited without answering it may still be fetching, and nothing is deleted. + +A cancel can also arrive **before** the pull it names has registered, since the +read loop dispatches each onto its own goroutine. That leaves a short-lived +tombstone the pull consumes instead of starting, rather than a success report +for a download that then runs on under a UI stuck on "Canceling". The tombstone +is scoped to the accepted request, so it cannot be inherited by a later retry. +On the remote path the same ordering is enforced by holding +`engine:remote-cancel-pull` until the peer has accepted the matching +`engine:remote-pull-model`. + ## Remote engine management When the parent passes `--control-port`, engine-manager serves the **`ec` @@ -211,6 +255,7 @@ Endpoints (all under `/v1`): | `GET /v1/engines` | JSON | remote `engine:get-installed` | | `POST /v1/engines/install` | NDJSON stream | remote install (+ optional start) with live progress | | `POST /v1/models/pull` | NDJSON stream | remote model pull with live progress | +| `POST /v1/models/cancel-pull` | JSON | remote `engine:cancel-pull` | | `POST /v1/engines/start` | JSON | remote start → `EngineStatus` | | `POST /v1/engines/stop` | JSON | remote stop → `EngineStatus` | diff --git a/services/nvpair-engine-manager/actions.go b/services/nvpair-engine-manager/actions.go index 372ee104..d63e95d7 100644 --- a/services/nvpair-engine-manager/actions.go +++ b/services/nvpair-engine-manager/actions.go @@ -294,6 +294,9 @@ func (e *Executor) runCommandOutput(ctx context.Context, argv []string) (string, if len(argv) == 0 { return "", nil } + if output, ok := ctx.Value(downloadOutputKey{}).(*lmsDownloadOutput); ok { + return runLMSDownloadCommand(ctx, argv, output) + } cmd := exec.CommandContext(ctx, argv[0], argv[1:]...) configureSysProcAttr(cmd) out, err := cmd.Output() diff --git a/services/nvpair-engine-manager/controlmodels.go b/services/nvpair-engine-manager/controlmodels.go index 8bdce1ab..082cf062 100644 --- a/services/nvpair-engine-manager/controlmodels.go +++ b/services/nvpair-engine-manager/controlmodels.go @@ -14,15 +14,20 @@ import ( ) const ( - controlLoadPath = "/v1/models/load" - controlUnloadPath = "/v1/models/unload" - controlDeletePath = "/v1/models/delete" + controlCancelPullPath = "/v1/models/cancel-pull" + controlLoadPath = "/v1/models/load" + controlUnloadPath = "/v1/models/unload" + controlDeletePath = "/v1/models/delete" ) func (s *controlServer) handleLoad(w http.ResponseWriter, r *http.Request) { s.handleModelAction(w, r, "load") } +func (s *controlServer) handleCancelPull(w http.ResponseWriter, r *http.Request) { + s.handleModelAction(w, r, "cancel-pull") +} + func (s *controlServer) handleUnload(w http.ResponseWriter, r *http.Request) { s.handleModelAction(w, r, "unload") } @@ -55,6 +60,8 @@ func (s *controlServer) handleModelAction(w http.ResponseWriter, r *http.Request err error ) switch op { + case "cancel-pull": + err = s.exec.CancelModelPull(r.Context(), req.Engine, req.Model) case "load": res, err = s.exec.ModelLoad(r.Context(), req.Engine, req.Model) case "unload": diff --git a/services/nvpair-engine-manager/controlserver.go b/services/nvpair-engine-manager/controlserver.go index 8a23a5f5..eb132fe8 100644 --- a/services/nvpair-engine-manager/controlserver.go +++ b/services/nvpair-engine-manager/controlserver.go @@ -63,6 +63,7 @@ func (s *controlServer) mux() *http.ServeMux { mux.HandleFunc(controlEnginesPath, s.requirePin(s.handleEngines)) mux.HandleFunc(controlInstallPath, s.requirePin(s.handleInstall)) mux.HandleFunc(controlPullPath, s.requirePin(s.handlePull)) + mux.HandleFunc(controlCancelPullPath, s.requirePin(s.handleCancelPull)) mux.HandleFunc(controlLoadPath, s.requirePin(s.handleLoad)) mux.HandleFunc(controlUnloadPath, s.requirePin(s.handleUnload)) mux.HandleFunc(controlDeletePath, s.requirePin(s.handleDelete)) diff --git a/services/nvpair-engine-manager/controlstream.go b/services/nvpair-engine-manager/controlstream.go index 13bd584d..a7e56dc4 100644 --- a/services/nvpair-engine-manager/controlstream.go +++ b/services/nvpair-engine-manager/controlstream.go @@ -32,6 +32,7 @@ const maxControlBody = 1 << 20 // 1 MiB // streamFrame is one NDJSON frame on the ec streaming endpoints. Type is // "progress", "result", or "error". Percent is omitted when indeterminate (0). type streamFrame struct { + Model string `json:"model,omitempty"` Type string `json:"type"` OpID string `json:"opId,omitempty"` Engine string `json:"engine,omitempty"` @@ -77,7 +78,7 @@ func (s *controlServer) handleInstall(w http.ResponseWriter, r *http.Request) { http.Error(w, `"engine" is required`, http.StatusBadRequest) return } - s.streamOp(w, r, req.OpID, req.Engine, "install", func(ctx context.Context) (streamFrame, error) { + s.streamOp(w, r, req.OpID, req.Engine, "install", "", func(ctx context.Context) (streamFrame, error) { if err := s.exec.Install(ctx, req.Engine); err != nil { return streamFrame{}, err } @@ -115,7 +116,14 @@ func (s *controlServer) handlePull(w http.ResponseWriter, r *http.Request) { http.Error(w, `"model" or "params" is required`, http.StatusBadRequest) return } - s.streamOp(w, r, req.OpID, req.Engine, "pull", func(ctx context.Context) (streamFrame, error) { + if model := modelFromParams(req.Params); model != "" { + req.Model = model + } + // The initiator's cancel arrives as its own request, so claim the pull + // before starting it: a cancel that beats the download to the executor is + // then held for it instead of being read as a cancel for nothing. + defer s.exec.claimPull(req.Engine, req.Model, true)() + s.streamOp(w, r, req.OpID, req.Engine, "pull", req.Model, func(ctx context.Context) (streamFrame, error) { res, err := s.exec.PullModelStream(ctx, req.Engine, req.Model, req.Params) if err != nil { model := req.Model @@ -134,7 +142,7 @@ func (s *controlServer) handlePull(w http.ResponseWriter, r *http.Request) { // forwards each progress event as an NDJSON frame, and writes run's terminal // frame (or an error frame) last. run executes on the request context, so a // disconnected initiator cancels the operation. -func (s *controlServer) streamOp(w http.ResponseWriter, r *http.Request, opID, engine, op string, run func(ctx context.Context) (streamFrame, error)) { +func (s *controlServer) streamOp(w http.ResponseWriter, r *http.Request, opID, engine, op, modelFilter string, run func(ctx context.Context) (streamFrame, error)) { flusher, ok := w.(http.Flusher) if !ok { http.Error(w, "streaming unsupported", http.StatusInternalServerError) @@ -148,9 +156,17 @@ func (s *controlServer) streamOp(w http.ResponseWriter, r *http.Request, opID, e defer cancelSub() enc := json.NewEncoder(w) + // The subscription is per engine, so a model-scoped stream also sees the + // engine's install steps and any other model's pull. An install carries no + // model, so matching on the model alone would let it through — and the + // initiator would receive install progress stamped with its pull's opID. writeProgress := func(ev ProgressEvent) { + if modelFilter != "" && (ev.Op != "pull" || ev.Model != modelFilter) { + return + } f := streamFrame{ - Type: "progress", OpID: opID, Engine: ev.Engine, Op: ev.Op, + Model: ev.Model, + Type: "progress", OpID: opID, Engine: ev.Engine, Op: ev.Op, Stage: ev.Stage, Message: ev.Message, } if wirePercentIncluded(ev.Percent) { diff --git a/services/nvpair-engine-manager/controlstream_test.go b/services/nvpair-engine-manager/controlstream_test.go index deb37132..48f62206 100644 --- a/services/nvpair-engine-manager/controlstream_test.go +++ b/services/nvpair-engine-manager/controlstream_test.go @@ -22,7 +22,7 @@ func TestStreamOpEmitsProgressThenResult(t *testing.T) { req := httptest.NewRequest("POST", controlInstallPath, nil) st := EngineStatus{Engine: "ollama", Installed: true, Running: true} - s.streamOp(rec, req, "op1", "ollama", "install", func(ctx context.Context) (streamFrame, error) { + s.streamOp(rec, req, "op1", "ollama", "install", "", func(ctx context.Context) (streamFrame, error) { exec.progress.publish(ProgressEvent{Engine: "ollama", Op: "install", Stage: "downloading", Percent: 42}) return streamFrame{Type: "result", OpID: "op1", Engine: "ollama", Op: "install", Status: &st}, nil }) @@ -51,7 +51,7 @@ func TestStreamOpEmitsErrorFrame(t *testing.T) { rec := httptest.NewRecorder() req := httptest.NewRequest("POST", controlInstallPath, nil) - s.streamOp(rec, req, "op2", "ollama", "install", func(ctx context.Context) (streamFrame, error) { + s.streamOp(rec, req, "op2", "ollama", "install", "", func(ctx context.Context) (streamFrame, error) { return streamFrame{}, context.DeadlineExceeded }) @@ -61,6 +61,37 @@ func TestStreamOpEmitsErrorFrame(t *testing.T) { } } +// A subscription is per engine, so a pull stream also sees the engine's +// install steps and any other model's pull. Only the requested model's +// download belongs to the initiator that asked for it. +func TestStreamOpScopedToAModelDropsEveryOtherOperation(t *testing.T) { + exec := &Executor{progress: newProgressHub()} + s := &controlServer{exec: exec} + + rec := httptest.NewRecorder() + req := httptest.NewRequest("POST", controlPullPath, nil) + + s.streamOp(rec, req, "op3", "ollama", "pull", "demo", func(ctx context.Context) (streamFrame, error) { + // An install carries no model, so matching on the model alone let it + // through and stamped it with this pull's opID. + exec.progress.publish(ProgressEvent{Engine: "ollama", Op: "install", Stage: "downloading", Percent: 10}) + exec.progress.publish(ProgressEvent{Engine: "ollama", Model: "other", Op: "pull", Stage: "downloading", Percent: 20}) + exec.progress.publish(ProgressEvent{Engine: "ollama", Model: "demo", Op: "pull", Stage: "downloading", Percent: 30}) + return streamFrame{Type: "result", OpID: "op3", Engine: "ollama", Op: "pull"}, nil + }) + + frames := decodeFrames(t, rec.Body.String()) + if len(frames) != 2 { + t.Fatalf("expected the requested model's progress and the result, got %d frames: %s", len(frames), rec.Body.String()) + } + if frames[0].Type != "progress" || frames[0].Model != "demo" || frames[0].Percent != 30 { + t.Fatalf("bad progress frame: %+v", frames[0]) + } + if frames[1].Type != "result" { + t.Fatalf("bad result frame: %+v", frames[1]) + } +} + // TestHandleInstallRejectsMissingEngine verifies request-body validation. func TestHandleInstallRejectsMissingEngine(t *testing.T) { s := &controlServer{exec: &Executor{progress: newProgressHub()}} diff --git a/services/nvpair-engine-manager/download_unix.go b/services/nvpair-engine-manager/download_unix.go new file mode 100644 index 00000000..ea7b735f --- /dev/null +++ b/services/nvpair-engine-manager/download_unix.go @@ -0,0 +1,34 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !windows + +package main + +import ( + "os/exec" + "syscall" +) + +func configureDownloadProcess(cmd *exec.Cmd) error { + configureSysProcAttr(cmd) + return nil +} + +func killDownload(cmd *exec.Cmd) error { return signalDownload(cmd, syscall.SIGKILL) } +func interruptDownload(cmd *exec.Cmd) error { return signalDownload(cmd, syscall.SIGINT) } +func handleDownloadProcess() bool { return false } + +// signalDownload signals the whole download process group. configureSysProcAttr +// puts the CLI in its own group precisely so a signal reaches the helpers it +// forks; addressing the leader alone leaves the process doing the downloading +// running. +func signalDownload(cmd *exec.Cmd, sig syscall.Signal) error { + if cmd == nil || cmd.Process == nil { + return nil + } + if pgid, err := syscall.Getpgid(cmd.Process.Pid); err == nil { + return syscall.Kill(-pgid, sig) + } + return cmd.Process.Signal(sig) +} diff --git a/services/nvpair-engine-manager/download_windows.go b/services/nvpair-engine-manager/download_windows.go new file mode 100644 index 00000000..10992a23 --- /dev/null +++ b/services/nvpair-engine-manager/download_windows.go @@ -0,0 +1,143 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//go:build windows + +package main + +import ( + "context" + "fmt" + "os" + "os/exec" + "strconv" + "strings" + "syscall" + "time" + + "golang.org/x/sys/windows" +) + +func configureDownloadProcess(cmd *exec.Cmd) error { + executable, err := os.Executable() + if err != nil { + return err + } + // The broker disables Ctrl+C for its workers; that setting survives even + // CREATE_NEW_CONSOLE. A launcher clears it before the CLI inherits it. + cmd.Args = append([]string{executable, "--run-download", cmd.Path}, cmd.Args[1:]...) + cmd.Path = executable + cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true, CreationFlags: windows.CREATE_NEW_CONSOLE} + return nil +} + +func killDownload(cmd *exec.Cmd) error { return taskkill(cmd, true) } + +// runDownloadProcess owns the CLI in a private hidden console. Reset only this +// launcher's inherited Ctrl+C setting, leaving the broker and worker untouched. +func runDownloadProcess(argv []string) int { + ignore := windows.NewLazySystemDLL("kernel32.dll").NewProc("SetConsoleCtrlHandler") + if ok, _, err := ignore.Call(0, 0); ok == 0 { + fmt.Fprintln(os.Stderr, err) + return 1 + } + cmd := exec.Command(argv[0], argv[1:]...) + cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true} + cmd.Stdin, cmd.Stdout, cmd.Stderr = os.Stdin, os.Stdout, os.Stderr + if err := cmd.Start(); err != nil { + fmt.Fprintln(os.Stderr, err) + return 1 + } + // The child has inherited enabled Ctrl+C. Ignore it in the launcher only, + // so the launcher stays alive to relay the CLI's exit and close its pipes. + if ok, _, err := ignore.Call(0, 1); ok == 0 { + _ = cmd.Process.Kill() + _ = cmd.Wait() + fmt.Fprintln(os.Stderr, err) + return 1 + } + if err := cmd.Wait(); err != nil { + if exitErr, ok := err.(*exec.ExitError); ok { + return exitErr.ExitCode() + } + fmt.Fprintln(os.Stderr, err) + return 1 + } + return 0 +} + +// isDownloadLauncher reports whether pid still names a launcher this worker +// started, by checking that its image is this same executable. Windows reissues +// a PID as soon as the last handle to the exited process closes, so a cancel +// racing the CLI's exit could otherwise attach to a stranger's console — and +// GenerateConsoleCtrlEvent(CTRL_C_EVENT, 0) reaches everything sharing it. +func isDownloadLauncher(pid uint32) bool { + self, err := os.Executable() + if err != nil { + return false + } + handle, err := windows.OpenProcess(windows.PROCESS_QUERY_LIMITED_INFORMATION, false, pid) + if err != nil { + return false + } + defer windows.CloseHandle(handle) + image := make([]uint16, 32768) + size := uint32(len(image)) + if err := windows.QueryFullProcessImageName(handle, 0, &image[0], &size); err != nil { + return false + } + return strings.EqualFold(windows.UTF16ToString(image[:size]), self) +} + +func interruptDownload(cmd *exec.Cmd) error { + executable, err := os.Executable() + if err != nil { + return err + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + helper := exec.CommandContext(ctx, executable, "--interrupt-download", strconv.Itoa(cmd.Process.Pid)) + configureSysProcAttr(helper) + return helper.Run() +} + +// Download launch and console attachment run in helpers so engine-manager's +// stdio and signal handlers are not affected. +func handleDownloadProcess() bool { + if len(os.Args) >= 3 && os.Args[1] == "--run-download" { + os.Exit(runDownloadProcess(os.Args[2:])) + } + if len(os.Args) != 3 || os.Args[1] != "--interrupt-download" { + return false + } + pid, err := strconv.ParseUint(os.Args[2], 10, 32) + if err != nil || pid == 0 { + os.Exit(1) + } + if !isDownloadLauncher(uint32(pid)) { + fmt.Fprintln(os.Stderr, "download launcher is gone") + os.Exit(1) + } + kernel := windows.NewLazySystemDLL("kernel32.dll") + attach := kernel.NewProc("AttachConsole") + free := kernel.NewProc("FreeConsole") + ignore := kernel.NewProc("SetConsoleCtrlHandler") + generate := kernel.NewProc("GenerateConsoleCtrlEvent") + free.Call() + if ok, _, err := attach.Call(uintptr(pid)); ok == 0 { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + if ok, _, err := ignore.Call(0, 1); ok == 0 { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + if ok, _, err := generate.Call(0, 0); ok == 0 { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + // Control-event delivery is asynchronous. + time.Sleep(100 * time.Millisecond) + free.Call() + return true +} diff --git a/services/nvpair-engine-manager/download_windows_test.go b/services/nvpair-engine-manager/download_windows_test.go new file mode 100644 index 00000000..1c831e30 --- /dev/null +++ b/services/nvpair-engine-manager/download_windows_test.go @@ -0,0 +1,50 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//go:build windows + +package main + +import ( + "context" + "os" + "os/exec" + "testing" + "time" + + "golang.org/x/sys/windows" +) + +// A cancel racing the CLI's exit can reach interruptDownload after the PID has +// been reissued, and the helper's GenerateConsoleCtrlEvent(CTRL_C_EVENT, 0) +// reaches every process sharing that console. The image check is what keeps the +// signal inside the app. +func TestIsDownloadLauncherRejectsForeignProcesses(t *testing.T) { + if !isDownloadLauncher(uint32(os.Getpid())) { + t.Fatal("this test binary was not recognized as its own image") + } + // PID 4 is the Windows System process; a PID no process holds is also a + // reissue candidate. + for _, pid := range []uint32{4, 0xFFFFFFF0} { + if isDownloadLauncher(pid) { + t.Errorf("pid %d was accepted as a download launcher", pid) + } + } +} + +// The broker launches engine-manager in a new process group, which disables +// Ctrl+C and passes that setting to descendants, including a download CLI. +func TestLMSDownloadInBrokerProcessGroup(t *testing.T) { + executable, err := os.Executable() + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second) + defer cancel() + cmd := exec.CommandContext(ctx, executable, "-test.run=^TestLMSDownloadInterruptsChild$", "-test.v") + configureSysProcAttr(cmd) + cmd.SysProcAttr.CreationFlags |= windows.CREATE_NEW_PROCESS_GROUP + if output, err := cmd.CombinedOutput(); err != nil { + t.Fatalf("download cancellation under broker process flags: %v\n%s", err, output) + } +} diff --git a/services/nvpair-engine-manager/executor.go b/services/nvpair-engine-manager/executor.go index de6bcfdd..70bad976 100644 --- a/services/nvpair-engine-manager/executor.go +++ b/services/nvpair-engine-manager/executor.go @@ -71,6 +71,18 @@ type engineState struct { // layer runs the long ones (install, start) in goroutines so the read // loop stays responsive. type Executor struct { + pullMu sync.Mutex + pulls map[string]*activePull + // pendingCancels holds cancels that arrived before their pull registered, + // so the download aborts instead of running on under a UI that already + // shows "Canceling". Entries are consumed by the pull they name and expire + // after pendingCancelWindow. + pendingCancels map[string]time.Time + // pullClaims counts the accepted pull requests per engine+model that have + // not finished, so a cancel can tell a pull that has not registered yet + // from one that already completed. Claimed where the request is accepted; + // see claimPull. + pullClaims map[string]int settingsHub settings.Hub settingsParent func(context.Context, string, settings.Request, string) (json.RawMessage, error) reg *Registry @@ -236,6 +248,9 @@ func (e *Executor) emitPullProgress(ev ProgressEvent) { params := map[string]any{ "engine": ev.Engine, "op": ev.Op, "stage": ev.Stage, "message": ev.Message, } + if ev.Model != "" { + params["model"] = ev.Model + } if wirePercentIncluded(ev.Percent) { params["percent"] = ev.Percent } diff --git a/services/nvpair-engine-manager/lmspull.go b/services/nvpair-engine-manager/lmspull.go new file mode 100644 index 00000000..71d1e5d6 --- /dev/null +++ b/services/nvpair-engine-manager/lmspull.go @@ -0,0 +1,337 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "encoding/json" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "regexp" + "strconv" + "strings" + "time" +) + +type downloadOutputKey struct{} + +// LMS redraws with carriage returns and ANSI controls, not newlines. +// A bounded tail handles split writes without keeping the entire download log. +type lmsDownloadOutput struct { + text string + lastPercent int + cancelled bool + completed bool + answered bool + stdin io.Writer + progress func(int) +} + +var lmsPercentPattern = regexp.MustCompile(`\]\s+(\d+(?:\.\d+)?)%`) +var ansiPattern = regexp.MustCompile("\x1b\\[[0-9;?]*[a-zA-Z]") + +func (w *lmsDownloadOutput) Write(data []byte) (int, error) { + w.text += string(data) + clean := ansiPattern.ReplaceAllString(w.text, "") + if strings.Contains(clean, "Continue to download in the background?") && !w.answered { + if _, err := io.WriteString(w.stdin, "n\n"); err != nil { + return 0, err + } + w.answered = true + } + w.cancelled = w.cancelled || strings.Contains(clean, "Download canceled.") + w.completed = w.completed || strings.Contains(clean, "Download completed.") + matches := lmsPercentPattern.FindAllStringSubmatch(clean, -1) + if len(matches) > 0 { + value, err := strconv.ParseFloat(matches[len(matches)-1][1], 64) + if err == nil { + percent := max(0, min(100, int(value))) + if percent != w.lastPercent { + w.lastPercent = percent + w.progress(percent) + } + } + } + if len(w.text) > 8192 { + w.text = w.text[len(w.text)-8192:] + } + return len(data), nil +} + +func (e *Executor) pullLMSModel(ctx context.Context, st *engineState, engine, model string) (json.RawMessage, error) { + before, err := lmsPartialFiles(lmstudioModelsDir(), model) + if err != nil { + return nil, fmt.Errorf("inspect partial downloads: %w", err) + } + st.mu.Lock() + port := st.port + st.mu.Unlock() + params, err := json.Marshal(map[string]string{"model": model}) + if err != nil { + return nil, err + } + output := &lmsDownloadOutput{lastPercent: -1, progress: func(percent int) { + if ctx.Err() == nil { + e.emitPullProgress(ProgressEvent{Engine: engine, Model: model, Op: "pull", Stage: "downloading", Percent: percent, Message: model}) + } + }} + e.emitPullProgress(ProgressEvent{Engine: engine, Model: model, Op: "pull", Stage: "pulling", Message: model}) + result, err := e.runCmdAction(context.WithValue(ctx, downloadOutputKey{}, output), st, st.manifest.Actions[pullModelAction], port, params) + // Only a cancellation someone asked for may delete partial files. The same + // context also dies when PAIR quits mid-download, when a remote initiator + // disconnects, and when the action timeout elapses — and `lms get` is built + // to resume every one of those on the next attempt. + if err != nil && cancelRequested(ctx) && output.cancelled { + // Cleanup outlives the cancellation that triggered it, so it runs on a + // context that keeps this one's values without its deadline. + if cleanupErr := cleanupLMSPartials(context.WithoutCancel(ctx), lmstudioModelsDir(), model, before); cleanupErr != nil { + return nil, fmt.Errorf("download stopped, but partial-file cleanup failed: %w", cleanupErr) + } + return nil, context.Canceled + } + return result, err +} + +// downloadWaitDelay bounds how long Wait blocks on the stdio pipes after the +// CLI itself has exited. `lms` commonly ships as a shim that execs Node, so the +// process actually writing bytes is a grandchild; one that outlives the signal +// keeps the inherited pipe open and would otherwise park Wait forever, wedging +// every queued pull for that engine. It is set here rather than per platform +// because that parked Wait is not platform-specific. +const downloadWaitDelay = 5 * time.Second + +func runLMSDownloadCommand(ctx context.Context, argv []string, output *lmsDownloadOutput) (string, error) { + if err := ctx.Err(); err != nil { + return "", err + } + cmd := exec.Command(argv[0], argv[1:]...) + cmd.WaitDelay = downloadWaitDelay + if err := configureDownloadProcess(cmd); err != nil { + return "", err + } + stdin, err := cmd.StdinPipe() + if err != nil { + return "", err + } + defer stdin.Close() + output.stdin = stdin + output.text = "" + output.cancelled = false + output.completed = false + output.answered = false + cmd.Stdout = output + cmd.Stderr = output + if err := cmd.Start(); err != nil { + return "", err + } + finished := make(chan error, 1) + go func() { finished <- cmd.Wait() }() + var runErr error + select { + case runErr = <-finished: + case <-ctx.Done(): + // Both cases can be ready at once — the user cancels exactly as the CLI + // exits — and select picks between them at random. Prefer the finished + // run: cmd.Wait has already reaped the process by then, so the OS is + // free to reissue its PID and an interrupt would reach a stranger. + select { + case runErr = <-finished: + default: + return stopLMSDownload(lmsCancelGracesFrom(ctx), cmd, finished, output) + } + } + if runErr != nil { + return "", fmt.Errorf("%w: %s", runErr, lmsErrorDetail(output.text)) + } + return output.text, nil +} + +// lmsCancelGraces bounds each stage of stopping the CLI. No wait here may be +// unbounded: cmd.Wait does not return until the stdio pipes close, and `lms` +// commonly ships as a shim that execs Node, so a grandchild holding an +// inherited pipe can outlive every signal sent to the group. downloadWaitDelay +// covers that for a process that exited, but a process the kill cannot remove +// at all would park the reap forever — and with it the cancel's JSON-RPC +// response, leaving the row on "Canceling" for good and wedging every queued +// pull for the engine. Giving up on the reap leaks one goroutine holding a +// pipe, which is the cheaper failure. +type lmsCancelGraces struct { + // interrupt is how long `lms get` has to acknowledge the first interrupt. + // It answers with the "Continue to download in the background?" prompt, + // which the output writer declines, then prints "Download canceled." + interrupt time.Duration + // retry bounds the second interrupt. The first can land while the CLI is + // mid-redraw and never reach the handler that prints the prompt. + retry time.Duration + // reap bounds the wait for a killed process. Reaching it means the kill did + // not remove the process, so the cancel is reported as failed. + reap time.Duration +} + +// lmsCancelGracesKey addresses the shortened graces a test installs on a context. +type lmsCancelGracesKey struct{} + +func lmsCancelGracesFrom(ctx context.Context) lmsCancelGraces { + if graces, ok := ctx.Value(lmsCancelGracesKey{}).(lmsCancelGraces); ok { + return graces + } + return lmsCancelGraces{interrupt: 15 * time.Second, retry: 5 * time.Second, reap: 10 * time.Second} +} + +// stopLMSDownload interrupts the CLI and reports whether it confirmed the +// cancellation. Confirmation matters beyond the exit code: declining the +// background-download prompt is what aborts the daemon's task, so a client +// killed before it answers can leave the transfer running. The interrupt is +// repeated before the process is reclaimed, and a stop the CLI never +// acknowledged is reported as a failed cancel that deletes nothing. +func stopLMSDownload( + graces lmsCancelGraces, + cmd *exec.Cmd, + finished <-chan error, + output *lmsDownloadOutput, +) (string, error) { + if err := interruptDownload(cmd); err != nil { + _ = killDownload(cmd) + if reaped, runErr := waitLMSExit(finished, graces.reap); reaped && runErr == nil && output.completed { + return output.text, nil + } + return "", fmt.Errorf("could not confirm LM Studio cancellation: %w", err) + } + exited, _ := waitLMSExit(finished, graces.interrupt) + if !exited && interruptDownload(cmd) == nil { + exited, _ = waitLMSExit(finished, graces.retry) + } + if exited { + if output.completed { + return output.text, nil + } + if !output.cancelled { + return "", fmt.Errorf("LM Studio exited without confirming download cancellation") + } + return "", context.Canceled + } + // Reclaim the process — it holds this worker's stdio pipes — but report the + // cancel as failed so no partial file is deleted on the strength of a + // cancellation the CLI never acknowledged. A process that outlasts even the + // kill is abandoned rather than waited on; see lmsCancelGraces. + _ = killDownload(cmd) + waitLMSExit(finished, graces.reap) + output.cancelled = false + return "", fmt.Errorf("LM Studio did not confirm download cancellation; the download may still be running in LM Studio") +} + +// waitLMSExit reports whether the CLI was reaped within grace, along with what +// cmd.Wait returned for it. +func waitLMSExit(finished <-chan error, grace time.Duration) (exited bool, runErr error) { + timer := time.NewTimer(grace) + defer timer.Stop() + select { + case runErr = <-finished: + return true, runErr + case <-timer.C: + return false, nil + } +} + +// lmsErrorDetailMax bounds the CLI output quoted in a returned error. The full +// transcript is kilobytes of redraw frames and local filesystem paths, and it +// travels into the persisted errors pipeline and back to remote initiators. +const lmsErrorDetailMax = 200 + +func lmsErrorDetail(text string) string { + clean := strings.TrimSpace(ansiPattern.ReplaceAllString(text, "")) + runes := []rune(clean) + if len(runes) <= lmsErrorDetailMax { + return clean + } + return "…" + string(runes[len(runes)-lmsErrorDetailMax:]) +} + +// Remove only the partial files this `lms get` produced, preserving completed +// shards. The LM Studio app and any other client download into the same +// repository directory, so a file is deleted only when it carries the requested +// quantization, differs from the snapshot taken before this download started, +// and has stopped moving. Never delete a model directory or follow a symlink +// outside the cache. +func cleanupLMSPartials(ctx context.Context, root, model string, before map[string]os.FileInfo) error { + files, err := lmsPartialFiles(root, model) + if err != nil { + return err + } + touched := make([]string, 0, len(files)) + for path, info := range files { + if old, ok := before[path]; ok && old.Size() == info.Size() && old.ModTime().Equal(info.ModTime()) { + continue + } + touched = append(touched, path) + } + stable, _ := quiescentPaths(ctx, touched) + for _, candidate := range stable { + if _, err := removeIfUnchanged(candidate); err != nil { + return err + } + } + return nil +} + +// lmsPartialFiles lists the in-progress downloads in a model's repository +// directory that could belong to the requested quantization. +func lmsPartialFiles(root, model string) (map[string]os.FileInfo, error) { + owner, repo, ok := lmsOwnerName(model) + if !ok { + return nil, nil + } + repo, quant, _ := strings.Cut(repo, "@") + for _, part := range []string{owner, repo} { + if part == "" || part == "." || part == ".." || strings.ContainsAny(part, `\:`) { + return nil, fmt.Errorf("invalid model repository") + } + } + target := filepath.Join(root, owner, repo) + resolvedRoot, err := filepath.EvalSymlinks(root) + if os.IsNotExist(err) { + return nil, nil + } + if err != nil { + return nil, err + } + resolvedTarget, err := filepath.EvalSymlinks(target) + if os.IsNotExist(err) { + return nil, nil + } + if err != nil { + return nil, err + } + if !pathWithinRoot(resolvedRoot, resolvedTarget) { + return nil, fmt.Errorf("partial download path escapes model directory") + } + entries, err := os.ReadDir(resolvedTarget) + if err != nil { + return nil, err + } + files := make(map[string]os.FileInfo) + for _, entry := range entries { + if entry.IsDir() || entry.Type()&os.ModeSymlink != 0 || !strings.HasPrefix(entry.Name(), "downloading_") || !strings.HasSuffix(entry.Name(), ".part") { + continue + } + // LM Studio names a GGUF after its quantization, so when the request + // pinned one, only files carrying it can be this download's — pulling + // @Q4_K_M must not touch a @Q8_0 the LM Studio app is fetching beside + // it. Without a pinned quantization the CLI picks its own and every + // partial here stays a candidate, left to the snapshot and settle + // checks in cleanupLMSPartials. + if quant != "" && !strings.Contains(strings.ToUpper(entry.Name()), strings.ToUpper(quant)) { + continue + } + info, err := entry.Info() + if err != nil { + return nil, err + } + files[filepath.Join(resolvedTarget, entry.Name())] = info + } + return files, nil +} diff --git a/services/nvpair-engine-manager/main.go b/services/nvpair-engine-manager/main.go index 4a9b5765..2caa89af 100644 --- a/services/nvpair-engine-manager/main.go +++ b/services/nvpair-engine-manager/main.go @@ -31,6 +31,9 @@ import ( var bundledManifests embed.FS func main() { + if handleDownloadProcess() { + return + } ipcPath := flag.String("ipc", "", "IPC endpoint: Unix domain socket path or Windows named pipe (default: stdin/stdout)") httpPort := flag.Int("http-port", 0, "if >0, serve the LAN HTTP surface (/v1/models) on this port so peers can enrich this node's model list; 0 disables it") controlPort := flag.Int("control-port", 0, "if >0 and this node is clustered, serve the cluster-scoped mTLS remote-control surface (ec: /v1/engines + remote install/pull/start/stop) on this port") diff --git a/services/nvpair-engine-manager/main_test.go b/services/nvpair-engine-manager/main_test.go index 0368a589..2f769ef6 100644 --- a/services/nvpair-engine-manager/main_test.go +++ b/services/nvpair-engine-manager/main_test.go @@ -20,6 +20,9 @@ var ( ) func TestMain(m *testing.M) { + if handleDownloadProcess() { + return + } tmp, err := os.MkdirTemp("", "nvpair-em-test-*") if err != nil { panic(err) diff --git a/services/nvpair-engine-manager/manager.go b/services/nvpair-engine-manager/manager.go index e12f62dc..f9d38397 100644 --- a/services/nvpair-engine-manager/manager.go +++ b/services/nvpair-engine-manager/manager.go @@ -48,6 +48,17 @@ type actionParam struct { Params json.RawMessage `json:"params,omitempty"` } +// pullClaimFrom reports the engine and model of an engine:action that starts a +// download, for Executor.claimPull. Malformed params are left to the handler, +// which owns the error response. +func pullClaimFrom(params json.RawMessage) (engine, model string, isPull bool) { + var p actionParam + if err := json.Unmarshal(params, &p); err != nil || p.Action != pullModelAction { + return "", "", false + } + return p.Engine, modelFromParams(p.Params), true +} + // setPortParam is the engine:set-port input: the engine whose server port to // change and the new port. The chosen port is persisted as a manifest // override and applied (the engine is bounced onto it if running). @@ -72,7 +83,10 @@ type Manager struct { // uses the longer header budget for start/delete (see waitsForEngineReadiness). remoteHTTP *clustertrust.PeerClientPool readyHTTP *clustertrust.PeerClientPool - cancel context.CancelFunc + // remotePulls orders an engine:remote-cancel-pull behind the remote pull it + // targets, the way claimPull orders a local one. + remotePulls remotePullGate + cancel context.CancelFunc } func NewManager(codec *Codec, exec *Executor, mesh *clustertrust.Mesh) *Manager { @@ -270,12 +284,38 @@ func (m *Manager) handleMessage(ctx context.Context, msg *Message) { m.codec.Respond(msg.ID, map[string]int{"port": p.Port}) case "engine:action": - go m.runAction(ctx, msg) + // Claim a pull here rather than in the goroutine. This case and + // engine:cancel-pull below each dispatch their own, so a cancel can + // reach the executor first; claiming on the read loop preserves the + // order the client sent them in. + release := m.exec.claimPull(pullClaimFrom(msg.Params)) + go func() { + defer release() + m.runAction(ctx, msg) + }() + + case "engine:cancel-pull": + var p remoteParam + if !m.parse(msg, &p) { + return + } + if p.Engine == "" || p.Model == "" { + m.codec.RespondError(msg.ID, -32602, "engine and model are required") + return + } + go func() { m.respondOrErr(msg, nil, m.exec.CancelModelPull(ctx, p.Engine, p.Model)) }() case "engine:remote-get-installed", "engine:remote-install", "engine:remote-pull-model", - "engine:remote-load-model", "engine:remote-unload-model", "engine:remote-delete-model", + "engine:remote-load-model", "engine:remote-unload-model", "engine:remote-delete-model", "engine:remote-cancel-pull", "engine:remote-start", "engine:remote-stop": - go m.runRemote(ctx, msg) + // Register a remote pull here for the same reason engine:action claims a + // local one: engine:remote-cancel-pull dispatches its own goroutine, so + // a cancel can otherwise reach the peer first. See remotePullGate. + release := m.remotePulls.register(remotePullClaimFrom(msg.Method, msg.Params)) + go func() { + defer release() + m.runRemote(ctx, msg) + }() default: m.codec.RespondError(msg.ID, -32601, fmt.Sprintf("method not found: %s", msg.Method)) @@ -379,7 +419,7 @@ func (m *Manager) runAction(ctx context.Context, msg *Message) { // error frame so a local subscriber converges off "pulling" even in // that case, mirroring install's failed progress step. userMsg := m.exec.reportPullFailed(p.Engine, model, err) - m.exec.emitPullProgress(ProgressEvent{Engine: p.Engine, Op: "pull", Stage: "error", Percent: -1, Message: userMsg}) + m.exec.emitPullProgress(ProgressEvent{Engine: p.Engine, Model: model, Op: "pull", Stage: "error", Percent: -1, Message: userMsg}) m.codec.RespondError(msg.ID, -32000, userMsg) return } diff --git a/services/nvpair-engine-manager/ollamapull.go b/services/nvpair-engine-manager/ollamapull.go new file mode 100644 index 00000000..ebfd7f1c --- /dev/null +++ b/services/nvpair-engine-manager/ollamapull.go @@ -0,0 +1,177 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +// ollamapull.go is Ollama's half of partial-file cleanup: naming the candidate +// files and deciding which of them this pull created. The safety test they are +// then put through is engine-neutral and lives in partials.go; lmspull.go is +// the LM Studio counterpart of this file. + +import ( + "context" + "fmt" + "os" + "path/filepath" + "regexp" + "strings" + "time" +) + +var ollamaDigestPattern = regexp.MustCompile(`^sha256:[a-f0-9]{64}$`) + +// ollamaPartialsBefore records the partial blobs already on disk when a +// download started. A nil map means no snapshot was taken, so nothing is +// attributable to that download and cancelling it deletes nothing. +type ollamaPartialsBefore map[string]bool + +// ollamaPartialSnapshot lists the partial blobs present before a download +// starts, which is what decides ownership at cancellation. +// +// A blob is content-addressed, so a digest names shared content rather than +// ownership: Ollama coalesces a pull of the same layer — from its CLI, its +// desktop app, the TUI's own engine-manager, or any other client on the daemon +// — onto the very same partial file. A partial that was already there is +// therefore not this download's to remove. It is either another client's +// transfer, which may sit still for far longer than any settle window while it +// waits on a slow server, or one PAIR deliberately kept when an earlier +// attempt's context died without a cancel request. Only the files this pull +// creates are unambiguously its own. +func ollamaPartialSnapshot(root string) (ollamaPartialsBefore, error) { + entries, err := os.ReadDir(root) + if os.IsNotExist(err) { + return ollamaPartialsBefore{}, nil + } + if err != nil { + return nil, err + } + before := make(ollamaPartialsBefore, len(entries)) + for _, entry := range entries { + if !entry.IsDir() && strings.Contains(entry.Name(), "-partial") { + before[filepath.Join(root, entry.Name())] = true + } + } + return before, nil +} + +// cleanupOllamaPartials removes the partial blobs this pull both reported and +// created. Completed blobs may be shared by installed models, and a partial +// that predates this download belongs to someone else; both are deliberately +// retained. A file this pull created can still be shared — another client can +// coalesce onto it afterwards — so one that is still moving is left alone and +// reported through busy, letting the caller look again. +func cleanupOllamaPartials( + ctx context.Context, + root string, + digests map[string]bool, + before ollamaPartialsBefore, +) (busy bool, err error) { + if before == nil { + return false, nil + } + candidates, err := ollamaPartialPaths(root, digests) + if err != nil || len(candidates) == 0 { + return false, err + } + created := make([]string, 0, len(candidates)) + for _, path := range candidates { + if !before[path] { + created = append(created, path) + } + } + stable, busy := quiescentPaths(ctx, created) + for _, candidate := range stable { + removed, err := removeIfUnchanged(candidate) + if err != nil { + return true, err + } + if !removed { + // Claimed between the last observation and the unlink. Leave it + // and let the retry decide. + busy = true + } + } + return busy, nil +} + +// ollamaPartialPaths lists the blob files matching any digest this pull +// reported. The directory holds one entry per layer of every installed model, +// so it is read once rather than per digest. +func ollamaPartialPaths(root string, digests map[string]bool) ([]string, error) { + prefixes := make([]string, 0, len(digests)) + for digest := range digests { + if ollamaDigestPattern.MatchString(digest) { + prefixes = append(prefixes, strings.ReplaceAll(digest, ":", "-")+"-partial") + } + } + if len(prefixes) == 0 { + return nil, nil + } + entries, err := os.ReadDir(root) + if os.IsNotExist(err) { + return nil, nil + } + if err != nil { + return nil, err + } + var paths []string + for _, entry := range entries { + if entry.IsDir() || entry.Type()&os.ModeSymlink != 0 { + continue + } + for _, prefix := range prefixes { + if entry.Name() == prefix || strings.HasPrefix(entry.Name(), prefix+"-") { + paths = append(paths, filepath.Join(root, entry.Name())) + break + } + } + } + return paths, nil +} + +func ollamaBlobsDir(st *engineState) string { + root := st.plat.Runtime.Env["OLLAMA_MODELS"] + if root == "" { + root = os.Getenv("OLLAMA_MODELS") + } + if root == "" { + root = "~/.ollama/models" + } + return filepath.Join(expandPath(root), "blobs") +} + +// partialCleanupRetryInterval spaces the cleanup passes after a cancellation +// and partialCleanupBudget bounds them. Ollama releases its background blob +// writers asynchronously once the pull request closes, so a Windows sharing +// violation and a file that has not stopped moving yet both tend to resolve +// within a few passes. The budget holds several full settle windows, and a file +// still busy when it runs out belongs to someone else. +const ( + partialCleanupRetryInterval = 100 * time.Millisecond + partialCleanupBudget = 5 * time.Second +) + +// cleanupOllamaAfterCancel retries until every partial this pull created is +// removed or confirmed to be another client's. One still busy at the deadline +// is left in place; that is the only safe outcome, and is not a cleanup failure. +func cleanupOllamaAfterCancel( + ctx context.Context, + root string, + digests map[string]bool, + before ollamaPartialsBefore, +) error { + deadline := time.Now().Add(partialCleanupBudget) + for { + time.Sleep(partialCleanupRetryInterval) + busy, err := cleanupOllamaPartials(ctx, root, digests, before) + if err == nil && !busy { + return nil + } + if time.Now().After(deadline) { + if err == nil { + return nil + } + return fmt.Errorf("download stopped, but partial-file cleanup failed: %w", err) + } + } +} diff --git a/services/nvpair-engine-manager/ollamapull_test.go b/services/nvpair-engine-manager/ollamapull_test.go new file mode 100644 index 00000000..27663a2b --- /dev/null +++ b/services/nvpair-engine-manager/ollamapull_test.go @@ -0,0 +1,109 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "testing" +) + +// Cancelling a download removes the partial blobs that download created and +// nothing else. A completed blob can be a layer of an installed model, and a +// partial for a digest this pull never reported belongs to another transfer. +func TestOllamaPartialCleanupRemovesOnlyThisPullsPartials(t *testing.T) { + root := t.TempDir() + expectKeep := []string{ + writeBlob(t, root, blobA), + writeBlob(t, root, blobB+"-partial"), + } + before, err := ollamaPartialSnapshot(root) + if err != nil { + t.Fatalf("snapshot partials: %v", err) + } + expectDelete := []string{ + writeBlob(t, root, blobA+"-partial"), + writeBlob(t, root, blobA+"-partial-0"), + } + + // "../escape" stands for a digest the engine never could have reported; it + // must not be turned into a path. + digests := map[string]bool{digestA: true, "../escape": true} + busy, err := cleanupOllamaPartials(settledContext(), root, digests, before) + if err != nil { + t.Fatalf("cleanup: %v", err) + } + if busy { + t.Error("settled partials were reported busy") + } + for _, path := range expectDelete { + assertRemoved(t, path) + } + for _, path := range expectKeep { + assertPresent(t, path) + } +} + +// A partial already on disk when the download started is not that download's to +// delete. Blobs are content-addressed, so it is either another client's +// transfer — which can sit still for far longer than any settle window while it +// waits on a slow server — or one PAIR deliberately kept when an earlier +// attempt's context died without a cancel request. Only the files a pull +// creates are unambiguously its own. +func TestOllamaPartialCleanupPreservesPartialsOlderThanThePull(t *testing.T) { + root := t.TempDir() + stalled := writeBlob(t, root, blobA+"-partial") + before, err := ollamaPartialSnapshot(root) + if err != nil { + t.Fatalf("snapshot partials: %v", err) + } + + busy, err := cleanupOllamaPartials(settledContext(), root, map[string]bool{digestA: true}, before) + if err != nil { + t.Fatalf("cleanup: %v", err) + } + if busy { + t.Error("a partial the snapshot excluded was reported busy") + } + assertPresent(t, stalled) +} + +// Without a snapshot nothing can be attributed to the pull, so cancelling it +// deletes nothing rather than guess at ownership. +func TestOllamaPartialCleanupWithoutASnapshotDeletesNothing(t *testing.T) { + root := t.TempDir() + partial := writeBlob(t, root, blobA+"-partial") + + busy, err := cleanupOllamaPartials(settledContext(), root, map[string]bool{digestA: true}, nil) + if err != nil { + t.Fatalf("cleanup: %v", err) + } + if busy { + t.Error("cleanup that ran on no snapshot reported busy") + } + assertPresent(t, partial) +} + +// Ollama coalesces a concurrent pull of the same layer — from its CLI, its +// desktop app, the TUI's own engine-manager, or any other client on the daemon +// — onto the very same partial file, including one this pull created. A file +// still growing after this transfer stopped is therefore shared, so it survives +// and is reported busy for the caller to look at again. +func TestOllamaPartialCleanupPreservesAPartialAnotherClientIsWriting(t *testing.T) { + root := t.TempDir() + before, err := ollamaPartialSnapshot(root) + if err != nil { + t.Fatalf("snapshot partials: %v", err) + } + shared := writeBlob(t, root, blobC+"-partial") + ctx := withPartialSettle(context.Background(), func() { growFile(t, shared) }) + + busy, err := cleanupOllamaPartials(ctx, root, map[string]bool{digestC: true}, before) + if err != nil { + t.Fatalf("cleanup: %v", err) + } + if !busy { + t.Error("a partial another client was still writing was not reported busy") + } + assertPresent(t, shared) +} diff --git a/services/nvpair-engine-manager/partials.go b/services/nvpair-engine-manager/partials.go new file mode 100644 index 00000000..e3f17583 --- /dev/null +++ b/services/nvpair-engine-manager/partials.go @@ -0,0 +1,169 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +// partials.go holds the engine-neutral half of partial-file cleanup: deciding +// whether a file is safe to delete. Which files are candidates in the first +// place is engine-specific and lives with that engine — ollamapull.go for +// Ollama's content-addressed blobs, lmspull.go for LM Studio's repository +// `.part` files. +// +// An engine's cleanup is expected to answer two questions in this order, and +// a new engine's should too: +// +// 1. Attribution — is this file ours? Each pull snapshots the partials +// present before it starts, and only files that appeared afterwards are +// candidates. This is what the engine supplies. +// 2. Quiescence — is anyone still writing it? A file this pull created can +// still be shared, because both vendors coalesce a concurrent download of +// the same content onto the same file. That test is the same for every +// engine, so it lives here. +// +// Deletion requires both. Every ambiguous case resolves toward keeping the +// file: a partial left behind costs disk space the vendor's own cleanup +// reclaims, while one deleted out from under a live download costs its bytes. + +import ( + "context" + "errors" + "io/fs" + "os" + "time" +) + +// A partial file is deleted only after it has matched its first observation at +// every one of partialSettleSamples checks, partialSettleInterval apart. +// +// No single observation can tell a finished writer from a stalled one: a +// transfer waiting on a slow server, a large flush, or a machine that just woke +// looks exactly like one that has stopped. Requiring several in a row makes one +// write anywhere in the window enough to spare the file, which is the error to +// prefer. Attribution does the real work; this is the backstop for the one case +// attribution cannot settle, a file two clients are writing at once. +const ( + partialSettleInterval = 250 * time.Millisecond + partialSettleSamples = 3 +) + +// partialSettleKey addresses the settle override a test installs on a context. +type partialSettleKey struct{} + +// settlePartials waits out one observation window. Cleanup runs after the pull's +// own context is already cancelled, so this deliberately does not observe +// ctx.Done: callers hand it a context.WithoutCancel of the pull's, and the +// bounded retry in cleanupOllamaAfterCancel is what limits the total wait. +func settlePartials(ctx context.Context) { + if settle, ok := ctx.Value(partialSettleKey{}).(func()); ok { + settle() + return + } + time.Sleep(partialSettleInterval) +} + +// observation is one stat of a candidate path. A path that could not be read +// is held distinct from one that is absent: "gone" means there is nothing left +// to do, while "unreadable" means we do not know, and the two must not lead to +// the same conclusion. +type observation struct { + info os.FileInfo + err error +} + +// gone reports whether the path is known to be absent, as opposed to present +// or merely unreadable. +func (o observation) gone() bool { return errors.Is(o.err, fs.ErrNotExist) } + +// moved reports whether a file differs from when it was first observed. An +// observation that could not be read counts as moved: an unreadable file is +// not a still one, and treating it as stable would delete on no evidence. +func (o observation) moved(first os.FileInfo) bool { + return o.err != nil || + first.Size() != o.info.Size() || + !first.ModTime().Equal(o.info.ModTime()) +} + +// statPaths observes every path, recording why each could not be read rather +// than dropping it. An entry is returned for every requested path. +func statPaths(paths []string) map[string]observation { + observed := make(map[string]observation, len(paths)) + for _, path := range paths { + info, err := os.Stat(path) + observed[path] = observation{info: info, err: err} + } + return observed +} + +// stablePath is a candidate that held still, carried with the observation it +// settled on so the removal can confirm nothing has changed since. +type stablePath struct { + path string + info os.FileInfo +} + +// quiescentPaths splits paths into those whose size and modification time held +// still across every observation, and reports whether any of the rest is still +// unfinished business. A path that moved, or appeared mid-window, is another +// client's and the caller can look again; one that is gone is already handled; +// one that could not be read at all is unknown, which is reported as busy +// rather than quietly treated as success. +func quiescentPaths(ctx context.Context, paths []string) (stable []stablePath, busy bool) { + watching := make(map[string]os.FileInfo, len(paths)) + for path, first := range statPaths(paths) { + if first.err == nil { + watching[path] = first.info + } + } + for sample := 0; sample < partialSettleSamples && len(watching) > 0; sample++ { + settlePartials(ctx) + now := statPaths(paths) + for path, first := range watching { + if now[path].moved(first) { + delete(watching, path) + } + } + } + final := statPaths(paths) + for _, path := range paths { + if info, quiet := watching[path]; quiet { + stable = append(stable, stablePath{path: path, info: info}) + continue + } + if !final[path].gone() { + busy = true + } + } + return stable, busy +} + +// removeIfUnchanged deletes path only if it still matches the observation +// quiescentPaths settled on, reporting whether the file is now gone. +// +// This narrows a race it cannot close. Between the final observation and the +// unlink, another client can open the file and resume writing it — both +// vendors coalesce a new download of the same content onto whatever partial is +// already there. Re-checking immediately before the unlink cuts the exposure +// from the settle window to the gap between two adjacent syscalls. Closing it +// outright would need a lock protocol neither vendor participates in. +// +// The platforms fail differently, and only one fails safely. On Windows an +// open handle without FILE_SHARE_DELETE makes Remove return a sharing +// violation, so the file survives on its own. On Unix the unlink succeeds and +// a writer holding a descriptor goes on filling an inode with no name, losing +// its download silently — so on that side this check is the only guard there +// is. +func removeIfUnchanged(candidate stablePath) (removed bool, err error) { + now := statPaths([]string{candidate.path})[candidate.path] + switch { + case now.gone(): + return true, nil + case now.moved(candidate.info): + // Unreadable, or claimed since we looked. Either way it is not ours to + // delete on this pass. + return false, nil + } + if err := os.Remove(candidate.path); err != nil && !errors.Is(err, fs.ErrNotExist) { + return false, err + } + return true, nil +} diff --git a/services/nvpair-engine-manager/partials_test.go b/services/nvpair-engine-manager/partials_test.go new file mode 100644 index 00000000..51081ba9 --- /dev/null +++ b/services/nvpair-engine-manager/partials_test.go @@ -0,0 +1,92 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "os" + "path/filepath" + "testing" +) + +// unreadablePath is a path os.Stat rejects with something other than "does not +// exist". The errors this stands in for — a permission denial on the vendor's +// cache, an I/O error on a failing disk — cannot be provoked portably, and the +// distinction under test is only between "absent" and "could not be read", so +// any non-NotExist error exercises it. Go rejects an interior NUL on every +// platform. +const unreadablePath = "partial\x00name" + +// A file that cannot be read is not a file that is gone. Treating the two +// alike let cleanup report success over a partial it never managed to look at, +// which is the one outcome worse than leaving the file: the caller stops +// retrying and the cancel is acknowledged as complete. +func TestQuiescentPathsSeparatesUnreadableFromAbsent(t *testing.T) { + root := t.TempDir() + settled := writeBlob(t, root, blobA+"-partial") + absent := filepath.Join(root, blobB+"-partial") + + stable, busy := quiescentPaths(settledContext(), []string{settled, absent, unreadablePath}) + + if !busy { + t.Error("a path that could not be read was not reported busy, so cleanup would stop retrying") + } + if len(stable) != 1 || stable[0].path != settled { + t.Fatalf("stable = %v, want only %s", stable, filepath.Base(settled)) + } +} + +// An absent path on its own settles: there is nothing left to remove, so the +// caller is done rather than retrying to the deadline. +func TestQuiescentPathsIgnoresAnAbsentPath(t *testing.T) { + root := t.TempDir() + absent := filepath.Join(root, blobA+"-partial") + + stable, busy := quiescentPaths(settledContext(), []string{absent}) + + if busy { + t.Error("an absent path was reported busy") + } + if len(stable) != 0 { + t.Fatalf("stable = %v, want nothing", stable) + } +} + +// quiescentPaths can only report what was true when it last looked. Both +// vendors coalesce a new download onto whatever partial is already there, so a +// file can be claimed in the gap before the unlink — and on Unix the unlink +// would succeed, leaving the new writer filling an inode with no name. The +// observation is therefore rechecked immediately before the removal. +func TestRemoveIfUnchangedSkipsAFileClaimedSinceItWasObserved(t *testing.T) { + test := func(name string, claim func(t *testing.T, path string), wantRemoved, wantPresent bool) { + t.Run(name, func(t *testing.T) { + root := t.TempDir() + path := writeBlob(t, root, blobA+"-partial") + info, err := os.Stat(path) + if err != nil { + t.Fatalf("observe %s: %v", path, err) + } + claim(t, path) + + removed, err := removeIfUnchanged(stablePath{path: path, info: info}) + if err != nil { + t.Fatalf("remove: %v", err) + } + if removed != wantRemoved { + t.Errorf("removed = %v, want %v", removed, wantRemoved) + } + if _, err := os.Stat(path); (err == nil) != wantPresent { + t.Errorf("file present = %v, want %v", err == nil, wantPresent) + } + }) + } + test("still matches the observation", func(*testing.T, string) {}, true, false) + test("another client resumed writing it", growFile, false, true) + // Already gone is a completed removal, not a failure: something else + // finished the job and the caller has nothing left to retry. + test("removed by someone else first", func(t *testing.T, path string) { + if err := os.Remove(path); err != nil { + t.Fatalf("remove %s: %v", path, err) + } + }, true, false) +} diff --git a/services/nvpair-engine-manager/progress.go b/services/nvpair-engine-manager/progress.go index e8c4dd53..10fb071c 100644 --- a/services/nvpair-engine-manager/progress.go +++ b/services/nvpair-engine-manager/progress.go @@ -10,6 +10,7 @@ import "sync" // engine:install-progress notification path (this node's own UI) and the ec // streaming handlers that relay live progress to a remote initiator. type ProgressEvent struct { + Model string `json:"model,omitempty"` Engine string `json:"engine"` Op string `json:"op"` // "install" | "pull" Stage string `json:"stage,omitempty"` diff --git a/services/nvpair-engine-manager/pull.go b/services/nvpair-engine-manager/pull.go index 6ffd8cca..b2344d7b 100644 --- a/services/nvpair-engine-manager/pull.go +++ b/services/nvpair-engine-manager/pull.go @@ -14,10 +14,8 @@ package main // Ollama's /api/pull streams newline-delimited JSON status objects // ({"status":...,"total":N,"completed":M}); each line maps to a progress event, // coalesced so only changes in stage/percent are emitted (a single layer streams -// many byte-progress lines at the same rendered percent). CLI-driven pulls (LM -// Studio's `lms get`) don't expose structured line progress here, so they emit a -// single "pulling" marker and return the final result — the security/trust -// boundary and result contract are identical. +// many byte-progress lines at the same rendered percent). LM Studio's lms get +// output is streamed through lmspull.go, including its cancellation prompt. import ( "bufio" @@ -58,6 +56,15 @@ func modelFromParams(params json.RawMessage) string { // params is empty it defaults to {"name","model"} (covering Ollama's `name` body // key and the {model} CLI placeholder), so a caller can pass just a model name. func (e *Executor) PullModelStream(ctx context.Context, engine, model string, params json.RawMessage) (json.RawMessage, error) { + if fromParams := modelFromParams(params); fromParams != "" { + model = fromParams + } + return e.trackedPull(ctx, engine, model, func(ctx context.Context) (json.RawMessage, error) { + return e.pullModelStream(ctx, engine, model, params) + }) +} + +func (e *Executor) pullModelStream(ctx context.Context, engine, model string, params json.RawMessage) (result json.RawMessage, resultErr error) { st, err := e.state(engine) if err != nil { return nil, err @@ -73,17 +80,58 @@ func (e *Executor) PullModelStream(ctx context.Context, engine, model string, pa ctx, cancel := context.WithTimeout(ctx, e.actionTimeout) defer cancel() - // CLI action (e.g. lms get): no structured line progress; emit a start + if act.ModelResolution == modelResolutionLMSGet { + return e.pullLMSModel(ctx, st, engine, model) + } + + // Other CLI actions without a progress adapter emit a start // marker and return the final result via the existing runner. if len(act.Cmd) > 0 { st.mu.Lock() port := st.port st.mu.Unlock() - e.emitPullProgress(ProgressEvent{Engine: engine, Op: "pull", Stage: "pulling", Message: model}) + e.emitPullProgress(ProgressEvent{Engine: engine, Model: model, Op: "pull", Stage: "pulling", Message: model}) return e.runCmdAction(ctx, st, act, port, params) } // HTTP action (e.g. Ollama /api/pull): stream NDJSON progress. + digests := make(map[string]bool) + completed := false + // Record the partial blobs on disk before a byte moves. Only the files this + // pull goes on to create are unambiguously its own, so this snapshot is + // what keeps a cancellation off another client's transfer and off a partial + // an earlier attempt deliberately left resumable. A snapshot that could not + // be taken stays nil, and cleanup then deletes nothing rather than guess. + var before ollamaPartialsBefore + if engine == "ollama" { + if snapshot, snapshotErr := ollamaPartialSnapshot(ollamaBlobsDir(st)); snapshotErr == nil { + before = snapshot + } + } + defer func() { + if ctx.Err() != context.Canceled { + return + } + if completed { + result = json.RawMessage(`{"status":"success"}`) + resultErr = nil + return + } + // Only a cancellation someone asked for may delete partial blobs. The + // same context also dies when PAIR quits mid-download and when a remote + // initiator's connection drops, and a resumable transfer must survive + // both rather than restart from zero on the next attempt. + if engine == "ollama" && cancelRequested(ctx) { + // Cleanup outlives the cancellation that triggered it, so it runs + // on a context that keeps this one's values without its deadline. + cleanupCtx := context.WithoutCancel(ctx) + if err := cleanupOllamaAfterCancel(cleanupCtx, ollamaBlobsDir(st), digests, before); err != nil { + resultErr = err + return + } + } + resultErr = context.Canceled + }() st.mu.Lock() running := st.running port := st.port @@ -127,7 +175,25 @@ func (e *Executor) PullModelStream(ctx context.Context, engine, model string, pa continue } last = append(json.RawMessage(nil), line...) + var frame struct { + Digest string `json:"digest"` + Error string `json:"error"` + Status string `json:"status"` + } + if err := json.Unmarshal(line, &frame); err != nil { + return nil, err + } + if frame.Error != "" { + return nil, fmt.Errorf("%s", frame.Error) + } + if frame.Status == "success" { + completed = true + } + if frame.Digest != "" { + digests[frame.Digest] = true + } ev := pullProgressFromLine(engine, line) + ev.Model = model if ev.Stage == lastStage && ev.Percent == lastPct { continue } diff --git a/services/nvpair-engine-manager/pull_test.go b/services/nvpair-engine-manager/pull_test.go index 76f48197..ec3ce3aa 100644 --- a/services/nvpair-engine-manager/pull_test.go +++ b/services/nvpair-engine-manager/pull_test.go @@ -117,7 +117,7 @@ func TestActionPullModelStreamsProgress(t *testing.T) { } } -// TestActionPullModelCmdMarkerAndResult covers the CLI (LM Studio `lms get`) +// TestActionPullModelCmdMarkerAndResult covers the CLI fallback // pull path through the same engine:action routing: a Cmd-based pull_model can't // expose structured line progress, so it emits a single "pulling" marker and // returns the command's terminal result — the counterpart to the HTTP streaming @@ -140,6 +140,9 @@ func TestActionPullModelCmdMarkerAndResult(t *testing.T) { mu.Unlock() }, t.TempDir()) + progress, unsubscribe := ex.progress.subscribe(m.Engine) + defer unsubscribe() + // A Cmd action just runs a binary; the engine need not be started. var out bytes.Buffer mgr := NewManager(NewCodec(&out), ex, nil) @@ -159,6 +162,17 @@ func TestActionPullModelCmdMarkerAndResult(t *testing.T) { if pulls[0]["op"] != "pull" || pulls[0]["stage"] != "pulling" || pulls[0]["message"] != "demo:1b" { t.Fatalf("unexpected CLI pull marker: %+v", pulls[0]) } + if pulls[0]["model"] != "demo:1b" { + t.Fatalf("CLI pull marker model = %v, want demo:1b", pulls[0]["model"]) + } + select { + case event := <-progress: + if event.Model != "demo:1b" { + t.Fatalf("CLI pull hub model = %q, want demo:1b", event.Model) + } + default: + t.Fatal("CLI pull did not publish progress to the hub") + } if _, hasPercent := pulls[0]["percent"]; hasPercent { t.Fatalf("CLI pull marker must omit indeterminate percent, got %+v", pulls[0]) } diff --git a/services/nvpair-engine-manager/pullcancel.go b/services/nvpair-engine-manager/pullcancel.go new file mode 100644 index 00000000..20c08aba --- /dev/null +++ b/services/nvpair-engine-manager/pullcancel.go @@ -0,0 +1,217 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +// pullcancel.go owns the engine-neutral cancellation lifecycle: tracking the +// pull in flight for an engine and model, stopping it, and serializing an +// engine's pulls. Removing the files a cancelled download left behind is +// separate — see partials.go for the shared safety test and ollamapull.go / +// lmspull.go for each engine's own. + +import ( + "context" + "encoding/json" + "strings" + "sync" + "sync/atomic" + "time" +) + +type activePull struct { + cancel context.CancelFunc + done chan struct{} + result json.RawMessage + err error + // requested records that someone asked for this download to stop. A + // context dies for three other reasons — the manager's run context ending + // when PAIR quits, a remote initiator's connection dropping, the action + // timeout elapsing — and none of them mean the user gave up on the bytes + // already on disk. Only a requested cancellation may delete partial files. + requested atomic.Bool +} + +type activePullKey struct{} + +// cancelRequested reports whether the pull running on this context was stopped +// on someone's behalf rather than because the context died underneath it. +func cancelRequested(ctx context.Context) bool { + p, ok := ctx.Value(activePullKey{}).(*activePull) + return ok && p.requested.Load() +} + +func pullKey(engine, model string) string { return engine + "\x00" + model } + +func cancelledPull() json.RawMessage { return json.RawMessage(`{"status":"cancelled"}`) } + +// pendingCancelWindow bounds how long a cancel that arrived before its pull +// registered stays armed. Both handlers are dispatched onto their own goroutine +// from the same read loop, so this covers the scheduler running them out of +// order — not a user cancelling a download that never starts. +const pendingCancelWindow = 10 * time.Second + +// claimPull marks a pull request accepted and not yet finished, and returns +// the release to run when it settles. The claim is what lets CancelModelPull +// tell a pull that has not registered yet from one that is already over, so it +// has to be taken where the request is accepted rather than where the download +// starts. Claiming something that is not a pull is a no-op. +func (e *Executor) claimPull(engine, model string, isPull bool) func() { + if !isPull || engine == "" || model == "" { + return func() {} + } + key := pullKey(engine, model) + e.pullMu.Lock() + if e.pullClaims == nil { + e.pullClaims = make(map[string]int) + } + e.pullClaims[key]++ + e.pullMu.Unlock() + var once sync.Once + return func() { + once.Do(func() { + e.pullMu.Lock() + defer e.pullMu.Unlock() + e.pullClaims[key]-- + if e.pullClaims[key] > 0 { + return + } + delete(e.pullClaims, key) + // A cancel held for a pull that ended without consuming it has + // nothing left to stop. Drop it here rather than letting it sit + // out its window, where a retry would inherit the cancellation. + delete(e.pendingCancels, key) + }) + } +} + +// CancelModelPull waits for the transfer and its cleanup before acknowledging. +// Cancellation must never delete a model. +// +// A cancel can arrive before its pull has registered, because the pull and the +// cancel are dispatched on separate goroutines. Acknowledging that as "already +// complete" would leave the download running under a UI stuck on "Canceling", +// so it leaves a tombstone the pull consumes instead of starting. That is only +// right while the pull is still owed to someone: an unclaimed engine and model +// has no download to stop, and a tombstone left there would cancel whatever +// pull came next. +func (e *Executor) CancelModelPull(ctx context.Context, engine, model string) error { + key := pullKey(engine, model) + e.pullMu.Lock() + p := e.pulls[key] + if p != nil { + p.requested.Store(true) + p.cancel() + } else if e.pullClaims[key] > 0 { + e.armPendingCancel(key) + } + e.pullMu.Unlock() + if p == nil { + return nil + } + select { + case <-p.done: + return p.err + case <-ctx.Done(): + return ctx.Err() + } +} + +// armPendingCancel records a cancel for a pull that has not registered yet, and +// drops the entries that have aged out. Callers hold pullMu. +func (e *Executor) armPendingCancel(key string) { + if e.pendingCancels == nil { + e.pendingCancels = make(map[string]time.Time) + } + now := time.Now() + for pending, armed := range e.pendingCancels { + if now.Sub(armed) > pendingCancelWindow { + delete(e.pendingCancels, pending) + } + } + e.pendingCancels[key] = now +} + +// takePendingCancel consumes a tombstone armed for this pull. Callers hold pullMu. +func (e *Executor) takePendingCancel(key string) bool { + armed, ok := e.pendingCancels[key] + if !ok { + return false + } + delete(e.pendingCancels, key) + return time.Since(armed) <= pendingCancelWindow +} + +func (e *Executor) trackedPull(ctx context.Context, engine, model string, run func(context.Context) (json.RawMessage, error)) (json.RawMessage, error) { + ctx, cancel := context.WithCancel(ctx) + defer cancel() + p := &activePull{cancel: cancel, done: make(chan struct{})} + ctx = context.WithValue(ctx, activePullKey{}, p) + key := pullKey(engine, model) + e.pullMu.Lock() + if e.pulls == nil { + e.pulls = make(map[string]*activePull) + } + if e.takePendingCancel(key) { + e.pullMu.Unlock() + return cancelledPull(), nil + } + // The vendor caches can share partial files across models. Keep pulls for + // an engine serialized so cleanup cannot remove another PAIR pull's data. + var predecessors []<-chan struct{} + for activeKey, active := range e.pulls { + if !strings.HasPrefix(activeKey, engine+"\x00") { + continue + } + if activeKey == key { + // A second request for a download already in flight joins it. An + // error here would be reported against the live pull's own model, + // and a terminal error frame disables that row's Cancel button. + e.pullMu.Unlock() + return joinPull(ctx, active) + } + predecessors = append(predecessors, active.done) + } + e.pulls[key] = p + e.pullMu.Unlock() + var result json.RawMessage + var err error + if len(predecessors) > 0 { + e.emitPullProgress(ProgressEvent{Engine: engine, Model: model, Op: "pull", Stage: "queued"}) + } +waitForPredecessors: + for _, done := range predecessors { + select { + case <-done: + case <-ctx.Done(): + break waitForPredecessors + } + } + if ctx.Err() != nil { + err = ctx.Err() + } else { + result, err = run(ctx) + } + e.pullMu.Lock() + delete(e.pulls, key) + // "Cancelled" is a successful outcome only for a cancellation someone asked + // for. A context that died on its own leaves a transfer the next attempt + // resumes, which is a failure to report rather than a request fulfilled. + if p.requested.Load() && ctx.Err() == context.Canceled && err == context.Canceled { + result, err = cancelledPull(), nil + } + p.result, p.err = result, err + close(p.done) + e.pullMu.Unlock() + return result, err +} + +// joinPull reports the outcome of a download already in flight, so a duplicate +// request is idempotent instead of a failure. +func joinPull(ctx context.Context, p *activePull) (json.RawMessage, error) { + select { + case <-p.done: + return p.result, p.err + case <-ctx.Done(): + return nil, ctx.Err() + } +} diff --git a/services/nvpair-engine-manager/pullcancel_test.go b/services/nvpair-engine-manager/pullcancel_test.go new file mode 100644 index 00000000..2cc85662 --- /dev/null +++ b/services/nvpair-engine-manager/pullcancel_test.go @@ -0,0 +1,876 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "path/filepath" + "strconv" + "strings" + "testing" + "time" +) + +// Digests and blob names are written out in full rather than built with +// strings.Repeat, so a failure message can be grepped straight back to the case +// that produced it. +const ( + digestA = "sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + blobA = "sha256-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + blobB = "sha256-bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + digestC = "sha256:cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" + blobC = "sha256-cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" +) + +// fakeDownloadPercent is the progress the fake CLI reports before it waits to +// be interrupted. The fake engine is a package main under testdata/, so the +// test binary cannot import it; passing the value in on the command line keeps +// it declared once, here, instead of matching a literal across two files. +const fakeDownloadPercent = 25 + +// testLMSCancelGraces shrinks the twenty seconds a real `lms get` is given to +// answer an interrupt, so a case that models a CLI ignoring one still runs +// quickly. reap stays generous: it covers a real kill and the pipe closing +// behind it. +var testLMSCancelGraces = lmsCancelGraces{ + interrupt: 500 * time.Millisecond, + retry: 500 * time.Millisecond, + reap: 30 * time.Second, +} + +// outcome is a pull's return pair, carried off the goroutine that ran it. +type outcome struct { + result json.RawMessage + err error +} + +// withPartialSettle replaces the real wait between the stat passes in +// quiescentPaths. A case where nothing is writing installs a no-op so the +// window costs nothing; one that models another client still writing appends +// from the hook, which lands inside the window on every run where a background +// goroutine racing a real sleep only usually does. +func withPartialSettle(ctx context.Context, settle func()) context.Context { + return context.WithValue(ctx, partialSettleKey{}, settle) +} + +// settledContext is the context for cleanup that has nothing to wait for. +func settledContext() context.Context { + return withPartialSettle(context.Background(), func() {}) +} + +func withLMSCancelGraces(ctx context.Context, graces lmsCancelGraces) context.Context { + return context.WithValue(ctx, lmsCancelGracesKey{}, graces) +} + +// fakeDownloadArgv builds the argv for one of the fake engine's download +// subcommands. Each takes the percentage to report before it waits for its +// interrupt. +func fakeDownloadArgv(subcommand string) []string { + return []string{fakeEngineBin, subcommand, strconv.Itoa(fakeDownloadPercent)} +} + +func writeBlob(t *testing.T, root, name string) string { + t.Helper() + path := filepath.Join(root, name) + if err := os.WriteFile(path, []byte("data"), 0600); err != nil { + t.Fatalf("seed blob %s: %v", name, err) + } + return path +} + +func writePartial(t *testing.T, path, data string) { + t.Helper() + if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil { + t.Fatalf("create %s: %v", filepath.Dir(path), err) + } + if err := os.WriteFile(path, []byte(data), 0600); err != nil { + t.Fatalf("seed partial %s: %v", path, err) + } +} + +// growFile appends to path, standing in for another client still writing it. +func growFile(t *testing.T, path string) { + t.Helper() + file, err := os.OpenFile(path, os.O_WRONLY|os.O_APPEND, 0600) + if err != nil { + t.Fatalf("open %s for append: %v", path, err) + } + if _, err := file.WriteString("more"); err != nil { + t.Errorf("append to %s: %v", path, err) + } + if err := file.Close(); err != nil { + t.Errorf("close %s: %v", path, err) + } +} + +func assertRemoved(t *testing.T, path string) { + t.Helper() + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Errorf("%s was not removed: %v", filepath.Base(path), err) + } +} + +func assertPresent(t *testing.T, path string) { + t.Helper() + if _, err := os.Stat(path); err != nil { + t.Errorf("%s was removed: %v", filepath.Base(path), err) + } +} + +func assertPullStatus(t *testing.T, result json.RawMessage, want string) { + t.Helper() + var terminal struct { + Status string `json:"status"` + } + if err := json.Unmarshal(result, &terminal); err != nil { + t.Fatalf("decode pull result %s: %v", result, err) + } + if terminal.Status != want { + t.Errorf("status = %q, want %q", terminal.Status, want) + } +} + +func succeedingPull(context.Context) (json.RawMessage, error) { + return json.RawMessage(`{"status":"success"}`), nil +} + +func TestLMSDownloadProgressChunks(t *testing.T) { + test := func(name string, chunks []string, want int) { + t.Run(name, func(t *testing.T) { + got := -1 + output := &lmsDownloadOutput{lastPercent: -1, progress: func(percent int) { got = percent }} + for _, chunk := range chunks { + if _, err := output.Write([]byte(chunk)); err != nil { + t.Fatal(err) + } + } + if got != want { + t.Fatalf("progress = %d, want %d", got, want) + } + }) + } + test("split percentage", []string{"\r[==== ] 2", "5.80", "%"}, 25) + test("ANSI redraw", []string{"\x1b[?25l\x1b[s\r[== ] 30.25%\x1b[u", "\r[=== ] 40.90%"}, 40) + test("indeterminate output", []string{"Resolving model..."}, -1) + test("clamps percentage", []string{"\r[=== ] 120.00%"}, 100) +} + +func TestLMSDownloadAnswersCancellationPrompt(t *testing.T) { + var stdin bytes.Buffer + output := &lmsDownloadOutput{stdin: &stdin, lastPercent: -1, progress: func(int) {}} + for _, chunk := range []string{"Continue to download ", "in the background? (Y/N): ", "Download canceled."} { + if _, err := output.Write([]byte(chunk)); err != nil { + t.Fatal(err) + } + } + if stdin.String() != "n\n" { + t.Fatalf("answer = %q, want n followed by newline", stdin.String()) + } + if !output.cancelled { + t.Fatal("cancellation acknowledgement was not recorded") + } +} + +func TestLMSDownloadInterruptsChild(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + progress := make(chan int, 1) + output := &lmsDownloadOutput{lastPercent: -1, progress: func(percent int) { progress <- percent }} + done := make(chan error, 1) + go func() { + _, err := runLMSDownloadCommand(ctx, fakeDownloadArgv("canceldownload"), output) + done <- err + }() + select { + case percent := <-progress: + if percent != fakeDownloadPercent { + t.Fatalf("progress = %d, want %d", percent, fakeDownloadPercent) + } + case <-ctx.Done(): + t.Fatal("no progress from child") + } + cancel() + select { + case err := <-done: + if !errors.Is(err, context.Canceled) { + t.Fatalf("cancel: %v", err) + } + case <-time.After(20 * time.Second): + t.Fatal("download did not stop") + } +} + +// What the CLI does with its interrupt decides what cancellation may do next, +// so each way stopLMSDownload can end gets a case. Only an acknowledged +// cancellation may delete a partial file: declining the background-download +// prompt is what aborts the daemon's task, so a CLI that exited without +// answering it can have left the transfer running. Every other outcome reports +// a failed cancel and leaves the files alone. +func TestLMSDownloadCancellationOutcomes(t *testing.T) { + test := func(name, subcommand, wantErrText string, wantCancelled bool) { + t.Run(name, func(t *testing.T) { + ctx, cancel := context.WithCancel(withLMSCancelGraces(context.Background(), testLMSCancelGraces)) + defer cancel() + progress := make(chan int, 1) + output := &lmsDownloadOutput{lastPercent: -1, progress: func(percent int) { + select { + case progress <- percent: + default: + } + }} + done := make(chan error, 1) + go func() { + _, err := runLMSDownloadCommand(ctx, fakeDownloadArgv(subcommand), output) + done <- err + }() + select { + case <-progress: + case <-time.After(30 * time.Second): + t.Fatal("no progress from the download CLI") + } + cancel() + + var err error + select { + case err = <-done: + case <-time.After(60 * time.Second): + t.Fatal("cancellation never settled") + } + switch { + case wantErrText == "" && err != nil: + t.Fatalf("cancellation = %v, want the completed run", err) + case wantErrText == "": + case err == nil: + t.Fatalf("cancellation succeeded, want an error containing %q", wantErrText) + case !strings.Contains(err.Error(), wantErrText): + t.Fatalf("cancellation = %v, want an error containing %q", err, wantErrText) + } + if output.cancelled != wantCancelled { + t.Errorf("acknowledged = %v, want %v (this is what gates deleting partial files)", output.cancelled, wantCancelled) + } + }) + } + // The CLI answers the prompt and confirms: context.Canceled is the + // cancellation the caller asked for, and cleanup may run. + test("acknowledged cancellation", "canceldownload", context.Canceled.Error(), true) + // The download finished as the interrupt landed. There is nothing to cancel + // and nothing to clean up, so the completed run is reported instead. + test("download completed first", "downloadcompletes", "", false) + // Exited without the prompt, so the daemon may still be fetching. + test("exit without confirmation", "downloadsilent", "exited without confirming download cancellation", false) + // Never exited at all, so the process is killed and the cancel fails. + test("interrupt ignored", "downloadignoresinterrupt", "did not confirm download cancellation", false) +} + +// An interrupt that cannot be delivered leaves the CLI's state unknown, so the +// only safe report is a failed cancel — unless the run had in fact already +// finished, which its reaped result still proves. A process that has been +// reaped is the reachable way to make delivery fail on every platform. +func TestStopLMSDownloadWhenTheInterruptCannotBeDelivered(t *testing.T) { + test := func(name string, completed bool, wantText, wantErrText string) { + t.Run(name, func(t *testing.T) { + cmd := exec.Command(fakeEngineBin, "echo", "done") + cmd.WaitDelay = downloadWaitDelay + if err := configureDownloadProcess(cmd); err != nil { + t.Fatalf("configure download process: %v", err) + } + if err := cmd.Start(); err != nil { + t.Fatalf("start fake CLI: %v", err) + } + finished := make(chan error, 1) + finished <- cmd.Wait() + + output := &lmsDownloadOutput{ + lastPercent: -1, + progress: func(int) {}, + completed: completed, + text: "transcript", + } + text, err := stopLMSDownload(testLMSCancelGraces, cmd, finished, output) + if text != wantText { + t.Errorf("text = %q, want %q", text, wantText) + } + if wantErrText == "" { + if err != nil { + t.Fatalf("stop = %v, want the completed run", err) + } + return + } + if err == nil || !strings.Contains(err.Error(), wantErrText) { + t.Fatalf("stop = %v, want an error containing %q", err, wantErrText) + } + }) + } + test("the run had already completed", true, "transcript", "") + test("the interrupt went nowhere", false, "", "could not confirm LM Studio cancellation") +} + +func TestLMSPartialCleanupPreservesCompletedShards(t *testing.T) { + root := t.TempDir() + target := filepath.Join(root, "owner", "repo") + partial := filepath.Join(target, "downloading_model-Q4_K_M.gguf.part") + complete := filepath.Join(target, "model-00001-of-00002.gguf") + writePartial(t, partial, "data") + writePartial(t, complete, "data") + + if err := cleanupLMSPartials(settledContext(), root, "https://huggingface.co/owner/repo@Q4_K_M", nil); err != nil { + t.Fatalf("cleanup: %v", err) + } + assertRemoved(t, partial) + assertPresent(t, complete) +} + +func TestLMSPartialCleanupPreservesUntouchedDownloads(t *testing.T) { + root := t.TempDir() + target := filepath.Join(root, "owner", "repo") + oldPartial := filepath.Join(target, "downloading_other.gguf.part") + writePartial(t, oldPartial, "other download") + + before, err := lmsPartialFiles(root, "owner/repo") + if err != nil { + t.Fatalf("snapshot partials: %v", err) + } + newPartial := filepath.Join(target, "downloading_selected.gguf.part") + writePartial(t, newPartial, "new download") + + if err := cleanupLMSPartials(settledContext(), root, "owner/repo", before); err != nil { + t.Fatalf("cleanup: %v", err) + } + assertPresent(t, oldPartial) + assertRemoved(t, newPartial) +} + +// The LM Studio app downloads into the same repository directory PAIR does, so +// cancelling a @Q4_K_M pull must not remove the @Q8_0 the app is fetching — +// neither the copy that grew since the snapshot nor one created after it. +func TestLMSPartialCleanupPreservesOtherQuantizations(t *testing.T) { + root := t.TempDir() + target := filepath.Join(root, "lmstudio-community", "Qwen3-8B-GGUF") + growing := filepath.Join(target, "downloading_Qwen3-8B-Q8_0.gguf.part") + writePartial(t, growing, "app download") + + const model = "lmstudio-community/Qwen3-8B-GGUF@Q4_K_M" + before, err := lmsPartialFiles(root, model) + if err != nil { + t.Fatalf("snapshot partials: %v", err) + } + if len(before) != 0 { + t.Fatalf("another quantization was captured as this pull's: %v", before) + } + appeared := filepath.Join(target, "downloading_Qwen3-8B-Q6_K.gguf.part") + mine := filepath.Join(target, "downloading_Qwen3-8B-Q4_K_M.gguf.part") + writePartial(t, growing, "app download, now longer") + writePartial(t, appeared, "app started this one after our snapshot") + writePartial(t, mine, "our download") + + if err := cleanupLMSPartials(settledContext(), root, model, before); err != nil { + t.Fatalf("cleanup: %v", err) + } + assertPresent(t, growing) + assertPresent(t, appeared) + assertRemoved(t, mine) +} + +// The app can also be fetching the very quantization PAIR asked for, into the +// same file. Carrying the requested quantization in its name is what makes a +// partial a candidate, not what makes it this download's, so one still growing +// after this transfer stopped has to survive too. +func TestLMSPartialCleanupPreservesTheRequestedQuantizationWhileItGrows(t *testing.T) { + root := t.TempDir() + target := filepath.Join(root, "lmstudio-community", "Qwen3-8B-GGUF") + const model = "lmstudio-community/Qwen3-8B-GGUF@Q4_K_M" + if err := os.MkdirAll(target, 0700); err != nil { + t.Fatalf("create %s: %v", target, err) + } + before, err := lmsPartialFiles(root, model) + if err != nil { + t.Fatalf("snapshot partials: %v", err) + } + shared := filepath.Join(target, "downloading_Qwen3-8B-Q4_K_M.gguf.part") + writePartial(t, shared, "shared download") + ctx := withPartialSettle(context.Background(), func() { growFile(t, shared) }) + + if err := cleanupLMSPartials(ctx, root, model, before); err != nil { + t.Fatalf("cleanup: %v", err) + } + assertPresent(t, shared) +} + +// The cancel may only be acknowledged once the partial files are gone, because +// the UI reads that acknowledgement as "it is safe to download this again". +// +// The case drives that order explicitly: +// +// 1. the download registers and blocks, so there is something to cancel; +// 2. the cancel arrives and closes the download's context; +// 3. the download observes the cancellation and enters its cleanup, where it +// blocks. Reaching this point proves the cancel is already past +// CancelModelPull's own call to p.cancel, so it must still be waiting; +// 4. cleanup finishes, and only then does the cancel return. +func TestCancelModelPullWaitsForCleanup(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + ex := NewExecutor(NewRegistry(), NewReporter(nil), nil, t.TempDir()) + downloading := make(chan struct{}) + cleaningUp := make(chan struct{}) + finishCleanup := make(chan struct{}) + pull := make(chan outcome, 1) + go func() { + var got outcome + got.result, got.err = ex.trackedPull(ctx, "ollama", "demo", func(runCtx context.Context) (json.RawMessage, error) { + close(downloading) + <-runCtx.Done() + close(cleaningUp) + <-finishCleanup + return nil, runCtx.Err() + }) + pull <- got + }() + <-downloading + + cancelled := make(chan error, 1) + go func() { cancelled <- ex.CancelModelPull(ctx, "ollama", "demo") }() + <-cleaningUp + select { + case err := <-cancelled: + t.Fatalf("cancel acknowledged before cleanup finished: %v", err) + default: + } + + close(finishCleanup) + if err := <-cancelled; err != nil { + t.Fatalf("cancel: %v", err) + } + got := <-pull + if got.err != nil { + t.Fatalf("pull: %v", got.err) + } + assertPullStatus(t, got.result, "cancelled") +} + +// Because the cancel waits for the transfer and its cleanup, its own caller +// giving up — a remote initiator disconnecting, the RPC's budget elapsing — has +// to end that wait. The download still stops; only the acknowledgement is +// abandoned. +func TestCancelModelPullStopsWaitingWhenItsCallerGivesUp(t *testing.T) { + ex := NewExecutor(NewRegistry(), NewReporter(nil), nil, t.TempDir()) + downloading := make(chan struct{}) + finishCleanup := make(chan struct{}) + pull := make(chan outcome, 1) + go func() { + var got outcome + got.result, got.err = ex.trackedPull(context.Background(), "ollama", "demo", func(runCtx context.Context) (json.RawMessage, error) { + close(downloading) + <-runCtx.Done() + <-finishCleanup + return nil, runCtx.Err() + }) + pull <- got + }() + <-downloading + // Let the download settle after the assertion, so the executor is not left + // with a pull in flight when the case ends. + defer func() { + close(finishCleanup) + <-pull + }() + + abandoned, abandon := context.WithCancel(context.Background()) + waited := make(chan error, 1) + go func() { waited <- ex.CancelModelPull(abandoned, "ollama", "demo") }() + abandon() + + select { + case err := <-waited: + if !errors.Is(err, context.Canceled) { + t.Fatalf("cancel = %v, want context.Canceled", err) + } + case <-time.After(30 * time.Second): + t.Fatal("cancel kept waiting after its own caller gave up") + } +} + +// A cancel dispatched before its pull registers must stop the download rather +// than report success while it runs on under a UI stuck on "Canceling". The +// pull and the cancel each run on their own goroutine, so the read loop's +// ordering does not settle which reaches the registry first. +func TestCancelBeforePullRegistersStopsIt(t *testing.T) { + ex := NewExecutor(NewRegistry(), NewReporter(nil), nil, t.TempDir()) + // The read loop claims a pull before dispatching it, so this is the window + // in which the cancel's goroutine can reach the executor first. + release := ex.claimPull("ollama", "demo", true) + defer release() + if err := ex.CancelModelPull(context.Background(), "ollama", "demo"); err != nil { + t.Fatalf("cancel: %v", err) + } + result, err := ex.trackedPull(context.Background(), "ollama", "demo", func(context.Context) (json.RawMessage, error) { + t.Error("cancelled pull started its download") + return nil, nil + }) + if err != nil { + t.Fatalf("pull: %v", err) + } + assertPullStatus(t, result, "cancelled") + + // The tombstone is consumed, so the same click cannot cancel a retry. + ran := false + if _, err := ex.trackedPull(context.Background(), "ollama", "demo", func(context.Context) (json.RawMessage, error) { + ran = true + return succeedingPull(context.Background()) + }); err != nil { + t.Fatalf("retry: %v", err) + } + if !ran { + t.Fatal("retry after a consumed cancellation did not start") + } +} + +// A cancel can also land just after the pull it targets finished, and that one +// has nothing to stop. Holding it for the next pull of the same model made a +// retry report "cancelled" without downloading anything. +func TestCancelAfterPullCompletesLeavesTheRetryAlone(t *testing.T) { + ex := NewExecutor(NewRegistry(), NewReporter(nil), nil, t.TempDir()) + release := ex.claimPull("ollama", "demo", true) + if _, err := ex.trackedPull(context.Background(), "ollama", "demo", succeedingPull); err != nil { + t.Fatalf("first pull: %v", err) + } + release() + + if err := ex.CancelModelPull(context.Background(), "ollama", "demo"); err != nil { + t.Fatalf("cancel: %v", err) + } + + ran := false + result, err := ex.trackedPull(context.Background(), "ollama", "demo", func(context.Context) (json.RawMessage, error) { + ran = true + return succeedingPull(context.Background()) + }) + if err != nil { + t.Fatalf("retry: %v", err) + } + if !ran { + t.Fatal("retry inherited the cancellation of a download that had already finished") + } + assertPullStatus(t, result, "success") +} + +// A pull can also end without ever consuming the cancel held for it — it +// failed before registering, or never registered at all. That tombstone has to +// go with the attempt rather than sit out its window. +func TestPendingCancelDoesNotOutliveTheAttemptItWasHeldFor(t *testing.T) { + ex := NewExecutor(NewRegistry(), NewReporter(nil), nil, t.TempDir()) + release := ex.claimPull("ollama", "demo", true) + if err := ex.CancelModelPull(context.Background(), "ollama", "demo"); err != nil { + t.Fatalf("cancel: %v", err) + } + release() + + ran := false + result, err := ex.trackedPull(context.Background(), "ollama", "demo", func(context.Context) (json.RawMessage, error) { + ran = true + return succeedingPull(context.Background()) + }) + if err != nil { + t.Fatalf("pull: %v", err) + } + if !ran { + t.Fatal("a cancellation held for an abandoned attempt was inherited by the next pull") + } + assertPullStatus(t, result, "success") +} + +// An unclaimed engine and model has no download to stop, so the cancel is a +// no-op rather than a tombstone. +func TestCancelWithNoDownloadRequestedIsANoOp(t *testing.T) { + ex := NewExecutor(NewRegistry(), NewReporter(nil), nil, t.TempDir()) + if err := ex.CancelModelPull(context.Background(), "ollama", "demo"); err != nil { + t.Fatalf("cancel: %v", err) + } + ex.pullMu.Lock() + pending := len(ex.pendingCancels) + ex.pullMu.Unlock() + if pending != 0 { + t.Fatalf("armed %d cancellation(s) for a download nobody requested", pending) + } +} + +// A duplicate request joins the download already in flight and reports its +// outcome. Failing it instead produced a terminal error frame attributed to the +// live pull's own model, which disables that row's Cancel button — and a +// cancellation the user asked for has to settle the duplicate's row too, not +// just the row the Cancel was clicked on. +func TestDuplicatePullJoinsTheActiveDownload(t *testing.T) { + test := func(name string, activeResult json.RawMessage, wantStatus string) { + t.Run(name, func(t *testing.T) { + ex := NewExecutor(NewRegistry(), NewReporter(nil), nil, t.TempDir()) + active := &activePull{cancel: func() {}, done: make(chan struct{})} + ex.pulls = map[string]*activePull{pullKey("ollama", "demo"): active} + + joined := make(chan outcome, 1) + go func() { + var got outcome + got.result, got.err = ex.trackedPull(context.Background(), "ollama", "demo", func(context.Context) (json.RawMessage, error) { + return nil, fmt.Errorf("duplicate request started a second download") + }) + joined <- got + }() + // Settling the live pull is what releases the joiner, so the result + // is published before the close that hands it over. The order holds + // however the two goroutines are scheduled. + active.result = activeResult + close(active.done) + + got := <-joined + if got.err != nil { + t.Fatalf("joined request: %v", got.err) + } + assertPullStatus(t, got.result, wantStatus) + }) + } + test("completed download", json.RawMessage(`{"status":"success"}`), "success") + test("cancelled download", cancelledPull(), "cancelled") +} + +func TestCancelledModelPullAllowsRetry(t *testing.T) { + ex := NewExecutor(NewRegistry(), NewReporter(nil), nil, t.TempDir()) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := ex.trackedPull(ctx, "ollama", "demo", func(context.Context) (json.RawMessage, error) { + t.Error("cancelled pull ran") + return nil, nil + }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("pull on a dead context = %v, want context.Canceled", err) + } + result, err := ex.trackedPull(context.Background(), "ollama", "demo", succeedingPull) + if err != nil { + t.Fatalf("retry: %v", err) + } + assertPullStatus(t, result, "success") +} + +// newOllamaPullExecutor wires an executor whose ollama pull_model action points +// at a fake /api/pull server. +func newOllamaPullExecutor(t *testing.T, serverURL string, emit func(method string, params any)) *Executor { + t.Helper() + _, portText, err := net.SplitHostPort(strings.TrimPrefix(serverURL, "http://")) + if err != nil { + t.Fatal(err) + } + port, err := strconv.Atoi(portText) + if err != nil { + t.Fatal(err) + } + manifest := testEngineManifest(fakeEngineBin) + manifest.Engine = "ollama" + manifest.Actions[pullModelAction] = Action{HTTP: &ActionHTTP{Method: http.MethodPost, Path: "/api/pull"}} + reg := NewRegistry() + reg.engines["ollama"] = manifest + ex := NewExecutor(reg, NewReporter(nil), emit, t.TempDir()) + st, err := ex.state("ollama") + if err != nil { + t.Fatal(err) + } + st.running = true + st.port = port + return ex +} + +// newOllamaBlobsDir points the engine at an empty blobs directory. The partial +// arrives once the download starts, the way a real one does, so the snapshot +// the pull takes beforehand can attribute it. +func newOllamaBlobsDir(t *testing.T) string { + t.Helper() + root := t.TempDir() + t.Setenv("OLLAMA_MODELS", root) + blobs := filepath.Join(root, "blobs") + if err := os.MkdirAll(blobs, 0700); err != nil { + t.Fatalf("create blobs dir: %v", err) + } + return blobs +} + +// The pull's context also dies when PAIR quits mid-download, when a remote +// initiator disconnects, and when the action timeout elapses. None of those +// mean the user gave up on the bytes already on disk, so a resumable transfer +// has to survive them instead of restarting from zero on the next attempt. +func TestUnrequestedCancellationKeepsPartials(t *testing.T) { + blobs := newOllamaBlobsDir(t) + partial := filepath.Join(blobs, blobA+"-partial") + streaming := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + writePartial(t, partial, "partial") + w.Header().Set("Content-Type", "application/x-ndjson") + if _, err := fmt.Fprintf(w, `{"status":"pulling layer","digest":%q,"total":100,"completed":25}`+"\n", digestA); err != nil { + t.Error(err) + return + } + w.(http.Flusher).Flush() + close(streaming) + <-r.Context().Done() + })) + defer server.Close() + ex := newOllamaPullExecutor(t, server.URL, nil) + + ctx, cancel := context.WithCancel(settledContext()) + defer cancel() + go func() { + <-streaming + cancel() + }() + if _, err := ex.PullModelStream(ctx, "ollama", "demo", nil); !errors.Is(err, context.Canceled) { + t.Fatalf("pull error = %v, want context.Canceled", err) + } + if data, err := os.ReadFile(partial); err != nil || string(data) != "partial" { + t.Fatalf("a resumable partial was deleted without a cancel request: data=%q err=%v", data, err) + } +} + +func TestOllamaCancellationHandlesCompletion(t *testing.T) { + test := func(name string, statuses []string, wantStatus string) { + t.Run(name, func(t *testing.T) { + blobs := newOllamaBlobsDir(t) + partial := filepath.Join(blobs, blobA+"-partial") + started := make(chan struct{}) + disconnected := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + writePartial(t, partial, "partial") + w.Header().Set("Content-Type", "application/x-ndjson") + for _, status := range statuses { + if _, err := fmt.Fprintf(w, `{"status":%q,"digest":%q,"total":100,"completed":25}`+"\n", status, digestA); err != nil { + t.Error(err) + return + } + } + w.(http.Flusher).Flush() + <-r.Context().Done() + close(disconnected) + })) + defer server.Close() + progressFrames := 0 + ex := newOllamaPullExecutor(t, server.URL, func(method string, _ any) { + if method == "engine:pull-progress" { + progressFrames++ + if progressFrames == len(statuses) { + close(started) + } + } + }) + ctx, cancel := context.WithTimeout(settledContext(), 30*time.Second) + defer cancel() + done := make(chan outcome, 1) + go func() { + var got outcome + got.result, got.err = ex.PullModelStream(ctx, "ollama", "demo", nil) + done <- got + }() + select { + case <-started: + case <-ctx.Done(): + t.Fatal("pull did not start") + } + if err := ex.CancelModelPull(ctx, "ollama", "demo"); err != nil { + t.Fatalf("cancel: %v", err) + } + got := <-done + if got.err != nil { + t.Fatalf("pull: %v", got.err) + } + assertPullStatus(t, got.result, wantStatus) + select { + case <-disconnected: + case <-ctx.Done(): + t.Fatal("HTTP transfer did not disconnect") + } + if wantStatus == "success" { + data, err := os.ReadFile(partial) + if err != nil || string(data) != "partial" { + t.Fatalf("late cancellation changed partial file after completion: data=%q, err=%v", data, err) + } + return + } + assertRemoved(t, partial) + }) + } + test("active transfer is canceled and cleaned", []string{"pulling layer"}, "cancelled") + test("success survives late cancellation", []string{"pulling layer", "success"}, "success") + test("success survives later frame and cancellation", []string{"pulling layer", "success", ""}, "success") +} + +// A predecessor here is another of this engine's downloads that registered +// first. trackedPull serializes pulls per engine — the vendor caches share +// partial files across models, so one pull's cleanup must not run while another +// PAIR pull is writing — which leaves a later request queued behind the ones +// already in flight. +// +// Cancelling a queued download has to settle on its own, without waiting for a +// turn it will never take: the row already says "Canceling", and a predecessor +// can be a multi-gigabyte transfer. The channels below stay open for the whole +// case to hold the predecessors in flight and prove that. +func TestQueuedModelPullCancelsBeforePredecessorsFinish(t *testing.T) { + test := func(name string, predecessorCount int) { + t.Run(name, func(t *testing.T) { + queued := make(chan struct{}, 1) + ex := NewExecutor(NewRegistry(), NewReporter(nil), func(method string, _ any) { + if method == "engine:pull-progress" { + queued <- struct{}{} + } + }, t.TempDir()) + ex.pulls = make(map[string]*activePull) + for i := 0; i < predecessorCount; i++ { + ex.pulls[pullKey("ollama", fmt.Sprintf("predecessor-%d", i))] = &activePull{done: make(chan struct{})} + } + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + finished := make(chan outcome, 1) + ran := false + go func() { + var got outcome + got.result, got.err = ex.trackedPull(ctx, "ollama", "queued", func(context.Context) (json.RawMessage, error) { + ran = true + return nil, nil + }) + finished <- got + }() + select { + case <-queued: + case <-ctx.Done(): + t.Fatal("pull was not queued") + } + if err := ex.CancelModelPull(ctx, "ollama", "queued"); err != nil { + t.Fatalf("cancel queued pull: %v", err) + } + select { + case got := <-finished: + if got.err != nil { + t.Fatalf("queued pull: %v", got.err) + } + if ran { + t.Error("canceled queued transfer started") + } + assertPullStatus(t, got.result, "cancelled") + case <-ctx.Done(): + t.Fatal("queued pull did not finish after cancellation") + } + ex.pullMu.Lock() + remaining := len(ex.pulls) + _, stillQueued := ex.pulls[pullKey("ollama", "queued")] + ex.pullMu.Unlock() + if stillQueued || remaining != predecessorCount { + t.Fatalf("registry after cancellation: queued=%v, remaining=%d", stillQueued, remaining) + } + }) + } + test("one active predecessor", 1) + test("multiple blocked predecessors", 3) +} diff --git a/services/nvpair-engine-manager/remote.go b/services/nvpair-engine-manager/remote.go index 5d296441..bc9b44cd 100644 --- a/services/nvpair-engine-manager/remote.go +++ b/services/nvpair-engine-manager/remote.go @@ -15,6 +15,8 @@ import ( "crypto/rand" "encoding/hex" "encoding/json" + "sync" + "time" ) // remoteParam is the shared input for the engine:remote-* methods. Fields not @@ -31,6 +33,7 @@ type remoteParam struct { // remoteProgress is the engine:remote-progress notification payload the broker // forwards to a subscribed UI. type remoteProgress struct { + Model string `json:"model,omitempty"` OpID string `json:"opId"` Node string `json:"node"` Engine string `json:"engine,omitempty"` @@ -47,6 +50,134 @@ func newOpID() string { return hex.EncodeToString(b[:]) } +// remotePullGate holds an engine:remote-cancel-pull until the +// engine:remote-pull-model it targets has reached the peer. +// +// Both methods dispatch their own goroutine from the same read loop, so nothing +// downstream preserves the order the UI sent them in. A cancel that wins that +// race arrives at the peer with no download registered and nothing claimed, is +// answered as a cancel for nothing, and leaves the transfer running under a row +// stuck on "Canceling". Executor.claimPull is the local cure for exactly this, +// and it cannot help here: the claim that matters belongs to the peer. +// +// handlePull claims the pull before streamOp writes the stream's response +// header, so that header arriving is proof the peer will recognize the cancel. +// Registering the pull on the read loop, where the client's order still holds, +// is what gives the cancel something to wait for. +type remotePullGate struct { + mu sync.Mutex + pulls map[string]*remotePullAttempt +} + +type remotePullAttempt struct { + // accepted closes when the peer has taken the pull, or when the attempt + // ended without ever getting that far, so a waiting cancel proceeds either + // way rather than outliving the download it was chasing. + accepted chan struct{} + settle sync.Once + // refs counts the requests sharing this attempt. The peer joins a duplicate + // pull onto the download already in flight, so two can be outstanding. + refs int +} + +// remotePullGateWindow bounds how long a cancel waits for its pull to reach the +// peer. Dialing a peer's ec surface can be slow, but a cancel that waits +// indefinitely would hold its own JSON-RPC response open for as long as the +// pull runs. Sending it late is better than never: the peer either has the +// download by then or has nothing to stop. +const remotePullGateWindow = 10 * time.Second + +func remotePullKey(node, engine, model string) string { + return node + "\x00" + engine + "\x00" + model +} + +// remotePullClaimFrom reports the gate key of an engine:remote-pull-model +// request. Malformed or incomplete params are left to runRemote, which owns the +// error response. +func remotePullClaimFrom(method string, params json.RawMessage) (string, bool) { + if method != "engine:remote-pull-model" { + return "", false + } + var p remoteParam + if err := json.Unmarshal(params, &p); err != nil { + return "", false + } + model := p.Model + if fromParams := modelFromParams(p.Params); fromParams != "" { + model = fromParams + } + if p.Node == "" || p.Engine == "" || model == "" { + return "", false + } + return remotePullKey(p.Node, p.Engine, model), true +} + +// register records a remote pull for the cancel that may be chasing it, and +// returns the release to run when the attempt ends. Registering anything that +// is not a remote pull is a no-op. +func (g *remotePullGate) register(key string, isPull bool) func() { + if !isPull { + return func() {} + } + g.mu.Lock() + if g.pulls == nil { + g.pulls = make(map[string]*remotePullAttempt) + } + attempt := g.pulls[key] + if attempt == nil { + attempt = &remotePullAttempt{accepted: make(chan struct{})} + g.pulls[key] = attempt + } + attempt.refs++ + g.mu.Unlock() + var once sync.Once + return func() { + once.Do(func() { + g.mu.Lock() + attempt.refs-- + if attempt.refs == 0 { + delete(g.pulls, key) + } + g.mu.Unlock() + // An attempt that ended without the peer accepting it has nothing + // left for a cancel to chase. + attempt.settle.Do(func() { close(attempt.accepted) }) + }) + } +} + +// accepted marks the peer as holding this pull, releasing any cancel waiting on +// it. Returns a callback so the caller can hand it straight to remoteClient.stream. +func (g *remotePullGate) accepted(key string) func() { + return func() { + g.mu.Lock() + attempt := g.pulls[key] + g.mu.Unlock() + if attempt != nil { + attempt.settle.Do(func() { close(attempt.accepted) }) + } + } +} + +// awaitAccepted blocks until the pull this cancel targets has reached the peer, +// the attempt ended, or the window elapses. A cancel for a download nobody +// requested finds no attempt and proceeds immediately. +func (g *remotePullGate) awaitAccepted(ctx context.Context, key string) { + g.mu.Lock() + attempt := g.pulls[key] + g.mu.Unlock() + if attempt == nil { + return + } + timer := time.NewTimer(remotePullGateWindow) + defer timer.Stop() + select { + case <-attempt.accepted: + case <-timer.C: + case <-ctx.Done(): + } +} + // runRemote dispatches an engine:remote-* request. It runs on its own goroutine // (like the other long ops) so the read loop stays responsive during a // multi-minute remote install/pull. @@ -82,7 +213,7 @@ func (m *Manager) runRemote(ctx context.Context, msg *Message) { } opID := newOpID() body := installRequest{OpID: opID, Engine: p.Engine, Start: p.Start} - terminal, err := client.stream(ctx, controlInstallPath, body, m.remoteProgressFn(opID, peer.nodeID)) + terminal, err := client.stream(ctx, controlInstallPath, body, m.remoteProgressFn(opID, peer.nodeID), nil) if err != nil { m.codec.RespondError(msg.ID, -32000, err.Error()) return @@ -100,14 +231,19 @@ func (m *Manager) runRemote(ctx context.Context, msg *Message) { } opID := newOpID() body := pullRequest{OpID: opID, Engine: p.Engine, Model: p.Model, Params: p.Params} - terminal, err := client.stream(ctx, controlPullPath, body, m.remoteProgressFn(opID, peer.nodeID)) + key, isPull := remotePullClaimFrom(msg.Method, msg.Params) + var onAccepted func() + if isPull { + onAccepted = m.remotePulls.accepted(key) + } + terminal, err := client.stream(ctx, controlPullPath, body, m.remoteProgressFn(opID, peer.nodeID), onAccepted) if err != nil { m.codec.RespondError(msg.ID, -32000, err.Error()) return } m.codec.Respond(msg.ID, map[string]any{"opId": opID, "result": terminal.Result}) - case "engine:remote-load-model", "engine:remote-unload-model", "engine:remote-delete-model": + case "engine:remote-load-model", "engine:remote-unload-model", "engine:remote-delete-model", "engine:remote-cancel-pull": if p.Engine == "" { m.codec.RespondError(msg.ID, -32602, "engine is required") return @@ -118,6 +254,11 @@ func (m *Manager) runRemote(ctx context.Context, msg *Message) { } path := controlLoadPath switch msg.Method { + case "engine:remote-cancel-pull": + path = controlCancelPullPath + // Let the download this cancel names reach the peer first; see + // remotePullGate. + m.remotePulls.awaitAccepted(ctx, remotePullKey(p.Node, p.Engine, p.Model)) case "engine:remote-unload-model": path = controlUnloadPath case "engine:remote-delete-model": @@ -153,7 +294,8 @@ func (m *Manager) runRemote(ctx context.Context, msg *Message) { func (m *Manager) remoteProgressFn(opID, node string) func(streamFrame) { return func(f streamFrame) { p := remoteProgress{ - OpID: opID, Node: node, Engine: f.Engine, Op: f.Op, + Model: f.Model, + OpID: opID, Node: node, Engine: f.Engine, Op: f.Op, Stage: f.Stage, Message: f.Message, } if wirePercentIncluded(f.Percent) { diff --git a/services/nvpair-engine-manager/remote_test.go b/services/nvpair-engine-manager/remote_test.go index 490f48dc..fb3d4b44 100644 --- a/services/nvpair-engine-manager/remote_test.go +++ b/services/nvpair-engine-manager/remote_test.go @@ -8,6 +8,7 @@ import ( "context" "encoding/json" "testing" + "time" "nvpair-shared/noderec" ) @@ -46,3 +47,135 @@ func TestRunRemoteNotClustered(t *testing.T) { Method: "engine:remote-install", Params: json.RawMessage(`{"node":"uuid-b","engine":"ollama"}`)}) mustContain(t, out.String(), "not clustered") } + +// Only a remote pull is registered, and its key has to match the one a cancel +// builds from its own params — including when the model arrived inside params +// rather than as a top-level field. +func TestRemotePullClaimFrom(t *testing.T) { + test := func(name, method, params, wantKey string) { + t.Run(name, func(t *testing.T) { + key, isPull := remotePullClaimFrom(method, json.RawMessage(params)) + if isPull != (wantKey != "") { + t.Fatalf("isPull = %v, want %v", isPull, wantKey != "") + } + if key != wantKey { + t.Errorf("key = %q, want %q", key, wantKey) + } + }) + } + test("pull with a top-level model", "engine:remote-pull-model", + `{"node":"uuid-b","engine":"ollama","model":"demo"}`, remotePullKey("uuid-b", "ollama", "demo")) + test("pull naming its model in params", "engine:remote-pull-model", + `{"node":"uuid-b","engine":"ollama","params":{"name":"demo"}}`, remotePullKey("uuid-b", "ollama", "demo")) + test("cancel is not registered", "engine:remote-cancel-pull", + `{"node":"uuid-b","engine":"ollama","model":"demo"}`, "") + test("install is not registered", "engine:remote-install", + `{"node":"uuid-b","engine":"ollama"}`, "") + test("pull without a model", "engine:remote-pull-model", + `{"node":"uuid-b","engine":"ollama"}`, "") + test("malformed params", "engine:remote-pull-model", `{`, "") +} + +// A cancel that reached the peer before its pull would find no download +// registered and nothing claimed, be answered as a cancel for nothing, and +// leave the transfer running under a row stuck on "Canceling". +func TestRemotePullGateHoldsACancelUntilThePeerHasThePull(t *testing.T) { + var gate remotePullGate + key := remotePullKey("uuid-b", "ollama", "demo") + release := gate.register(key, true) + defer release() + + waited := make(chan struct{}) + go func() { + defer close(waited) + gate.awaitAccepted(context.Background(), key) + }() + select { + case <-waited: + t.Fatal("cancel was sent before the peer had the pull") + case <-time.After(50 * time.Millisecond): + } + + gate.accepted(key)() + select { + case <-waited: + case <-time.After(10 * time.Second): + t.Fatal("cancel kept waiting after the peer accepted the pull") + } +} + +// The pull can also never reach the peer — an unreachable node, a rejected +// request. Its release has to let the cancel go rather than hold it for the +// whole window chasing a download that no longer exists. +func TestRemotePullGateReleasesACancelWhenThePullNeverLands(t *testing.T) { + var gate remotePullGate + key := remotePullKey("uuid-b", "ollama", "demo") + release := gate.register(key, true) + + waited := make(chan struct{}) + go func() { + defer close(waited) + gate.awaitAccepted(context.Background(), key) + }() + select { + case <-waited: + t.Fatal("cancel was sent while the pull attempt was still in flight") + case <-time.After(50 * time.Millisecond): + } + + release() + select { + case <-waited: + case <-time.After(10 * time.Second): + t.Fatal("cancel outlived the pull attempt it was chasing") + } + // The attempt is gone, so a later cancel for the same model has nothing to + // wait for and must not be held by the entry the first one used. + gate.awaitAccepted(context.Background(), key) +} + +// Two requests for the same remote download share one attempt, because the peer +// joins the duplicate onto the transfer already in flight. The cancel waits for +// the peer to accept, not for both requests to end. +func TestRemotePullGateSharesOneAttemptAcrossDuplicateRequests(t *testing.T) { + var gate remotePullGate + key := remotePullKey("uuid-b", "ollama", "demo") + first := gate.register(key, true) + second := gate.register(key, true) + + first() + waited := make(chan struct{}) + go func() { + defer close(waited) + gate.awaitAccepted(context.Background(), key) + }() + select { + case <-waited: + case <-time.After(10 * time.Second): + t.Fatal("cancel kept waiting after the first of two requests ended") + } + second() + + gate.mu.Lock() + remaining := len(gate.pulls) + gate.mu.Unlock() + if remaining != 0 { + t.Fatalf("gate retained %d attempt(s) after both requests ended", remaining) + } +} + +// A cancel for a download nobody requested has nothing to chase, so it goes +// straight through rather than spending the whole window. +func TestRemotePullGatePassesACancelWithNoPullRegistered(t *testing.T) { + var gate remotePullGate + done := make(chan struct{}) + go func() { + defer close(done) + gate.awaitAccepted(context.Background(), remotePullKey("uuid-b", "ollama", "demo")) + }() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("cancel waited on a pull that was never registered") + } +} diff --git a/services/nvpair-engine-manager/remoteclient.go b/services/nvpair-engine-manager/remoteclient.go index 3a19bdf0..9bc7ca91 100644 --- a/services/nvpair-engine-manager/remoteclient.go +++ b/services/nvpair-engine-manager/remoteclient.go @@ -157,7 +157,18 @@ func (c *remoteClient) postJSON(ctx context.Context, path, engine string, body a // stream POSTs body to a streaming ec endpoint and consumes its NDJSON frames, // calling onProgress for each progress frame and returning the terminal result // frame. An error frame (or a stream that ends without a result) is an error. -func (c *remoteClient) stream(ctx context.Context, path string, body any, onProgress func(streamFrame)) (streamFrame, error) { +// +// onAccepted, when set, is called once the peer has answered with a success +// status. The peer's handler has taken the operation by the time it writes that +// header, so this is the point at which a second request about the same +// operation can expect the peer to recognize it — see remotePullGate. +func (c *remoteClient) stream( + ctx context.Context, + path string, + body any, + onProgress func(streamFrame), + onAccepted func(), +) (streamFrame, error) { b, err := json.Marshal(body) if err != nil { return streamFrame{}, err @@ -177,6 +188,9 @@ func (c *remoteClient) stream(ctx context.Context, path string, body any, onProg data, _ := io.ReadAll(io.LimitReader(resp.Body, 64<<10)) return streamFrame{}, fmt.Errorf("remote %s: HTTP %d: %s", path, resp.StatusCode, strings.TrimSpace(string(data))) } + if onAccepted != nil { + onAccepted() + } dec := json.NewDecoder(resp.Body) var result streamFrame diff --git a/services/nvpair-engine-manager/spec.md b/services/nvpair-engine-manager/spec.md index ba2e6324..53f8fd3e 100644 --- a/services/nvpair-engine-manager/spec.md +++ b/services/nvpair-engine-manager/spec.md @@ -118,6 +118,7 @@ Requests (caller → service): | `engine:stop` | `{ engine }` | `EngineStatus` | | `engine:restart` | `{ engine }` | `EngineStatus` | | `engine:action` | `{ engine, action, params }` | the engine's raw response | +| `engine:cancel-pull` | `{ engine, model }` | `null` once the transfer has stopped and its partial files are cleaned up | | `engine:logs` | `{ engine }` | `{ lines: [LogLine] }` | | `engine:errors` | — | `{ errors: [ServiceError] }` | | `engine:remote-get-installed` | `{ node }` | `{ engines: [EngineStatus] }` from the remote node | @@ -126,12 +127,13 @@ Requests (caller → service): | `engine:remote-load-model` | `{ node, engine, model }` | the remote action result | | `engine:remote-unload-model` | `{ node, engine, model }` | the remote action result | | `engine:remote-delete-model` | `{ node, engine, model }` | the remote action result | +| `engine:remote-cancel-pull` | `{ node, engine, model }` | the remote cancellation result | | `engine:remote-start` | `{ node, engine, port? }` | `EngineStatus` from the remote node (manifest `runtime.bind`; no per-call bind on the remote path) | | `engine:remote-stop` | `{ node, engine }` | `EngineStatus` from the remote node | | `shutdown` | — | `null` | | `log/set-level` | `{ level }` | `{ level }` | -Notifications (service → caller): `ready{version}`, `engine:state-changed{EngineStatus}`, `engine:models-changed{engine, models}` (pushed when an engine's loaded-in-memory model set changes; `models` is the full `engine:models` shape incl. `loadedByEngine`), `engine:install-progress{engine, stage, percent}`, `engine:pull-progress{engine, op, stage, percent, message}` (live progress for a local model pull driven via `engine:action{action:"pull_model"}` — the local counterpart of `engine:remote-progress`), `engine:remote-progress{opId, node, engine, op, stage, percent, message}` (live progress relayed from a remote install/pull), and `errors:report` / `errors:clear` (consumed by `nvpair-errors` via the Broker). `install` / `start` / `stop` / `restart` / `action` / `remote-*` each run in their own goroutine so the read loop never blocks; their responses arrive when the op completes. +Notifications (service → caller): `ready{version}`, `engine:state-changed{EngineStatus}`, `engine:models-changed{engine, models}` (pushed when an engine's loaded-in-memory model set changes; `models` is the full `engine:models` shape incl. `loadedByEngine`), `engine:install-progress{engine, stage, percent}`, `engine:pull-progress{engine, model, op, stage, percent, message}` (live progress for a local model pull driven via `engine:action{action:"pull_model"}` — the local counterpart of `engine:remote-progress`), `engine:remote-progress{opId, node, engine, model, op, stage, percent, message}` (live progress relayed from a remote install/pull; `model` is set for a pull and absent for an install), and `errors:report` / `errors:clear` (consumed by `nvpair-errors` via the Broker). `install` / `start` / `stop` / `restart` / `action` / `remote-*` each run in their own goroutine so the read loop never blocks; their responses arrive when the op completes. Example `engine:install-progress` (stdout): ```json @@ -140,7 +142,7 @@ Example `engine:install-progress` (stdout): Example `engine:pull-progress` (stdout): ```json -{"jsonrpc":"2.0","method":"engine:pull-progress","params":{"engine":"ollama","op":"pull","stage":"pulling","percent":62,"message":"pulling"}} +{"jsonrpc":"2.0","method":"engine:pull-progress","params":{"engine":"ollama","model":"llama3.2","op":"pull","stage":"pulling","percent":62,"message":"pulling"}} ``` ### 7.1 Versioning @@ -148,7 +150,7 @@ Example `engine:pull-progress` (stdout): - `manifest_version` gates manifest-schema evolution: unknown optional fields are ignored (backward-compatible growth); a version higher than supported is rejected. ### 7.2 Remote engine management (the `ec` surface) -With `--control-port` and a clustered `--cluster-dir`, engine-manager serves a cluster-scoped remote-control surface over pin-based mTLS (`nvpair-shared/clustertrust`): it presents this node's cluster leaf, requires a client cert, and `403`s any caller that isn't a byte-for-byte pinned cluster peer. The listener is bound whenever `--control-port` is set and admits callers by live membership — its leaf is resolved per handshake, so an unclustered node presents none and every handshake is refused — so no restart is needed on `cluster:identity-changed`; the broker registers `ec` whenever a cluster dir is configured. Routes under `/v1`: `GET /engines`, streaming `POST /engines/install` and `/models/pull` (chunked NDJSON — zero+ `{"type":"progress"}` frames then one terminal `{"type":"result"}`/`{"type":"error"}` frame), non-streaming `POST /models/{load,unload,delete}`, and non-streaming `POST /engines/{start,stop}` returning `EngineStatus`. +With `--control-port` and a clustered `--cluster-dir`, engine-manager serves a cluster-scoped remote-control surface over pin-based mTLS (`nvpair-shared/clustertrust`): it presents this node's cluster leaf, requires a client cert, and `403`s any caller that isn't a byte-for-byte pinned cluster peer. The listener is bound whenever `--control-port` is set and admits callers by live membership — its leaf is resolved per handshake, so an unclustered node presents none and every handshake is refused — so no restart is needed on `cluster:identity-changed`; the broker registers `ec` whenever a cluster dir is configured. Routes under `/v1`: `GET /engines`, streaming `POST /engines/install` and `/models/pull` (chunked NDJSON — zero+ `{"type":"progress"}` frames then one terminal `{"type":"result"}`/`{"type":"error"}` frame), non-streaming `POST /models/{load,unload,delete,cancel-pull}`, and non-streaming `POST /engines/{start,stop}` returning `EngineStatus`. The `engine:remote-*` methods are the client half: engine-manager resolves the target `node` in an `ec` peer directory (fed by its own `discovery:subscribe{services:[ec]}` to the broker relay), dials the peer's `ec` surface with the same pinned identity, relays each streamed progress frame up as `engine:remote-progress` (keyed by a minted `opId`), and settles the request on the terminal frame. Remote install/pull run for the operation's full duration (no broker-imposed timeout); remote stop and model unload/delete are fast request/response, remote Ollama model load uses the readiness-sized response budget, and remote start waits for the engine's bounded readiness result. They error if this node isn't clustered or the target isn't a pinned peer. @@ -194,7 +196,9 @@ The operator starts it: `engine:start {engine:"ollama"}` resolves the manifest r The operator stops it: `engine:stop {engine:"ollama"}` signals a process the service owns. For an **adopted** engine (no owned process), it resolves the PID bound to the port and terminates it only when that process is running the binary we manage — reclaiming an orphan a prior run left on our own managed port; a genuinely foreign listener (a different image on a different port) is declined with an error naming its PID and image. A user-initiated `stop` records the OFF intent regardless (even when the RPC returns an error), so the health loop and restore-on-restart don't flip the engine back on — clients must not treat a stop error as proof the OFF choice was discarded. The cluster `ec` stop endpoint shares this semantics and may return HTTP 500 while OFF is persisted. -The operator pulls a model: `engine:action {engine:"ollama", action:"pull_model", params:{name:"llama3.2"}}` issues the manifest-declared `POST 127.0.0.1:{port}/api/pull`. Because the action is `pull_model`, the request is routed through the streaming pull path (not the buffered `engine:action` reader): each `/api/pull` status line is emitted as an `engine:pull-progress` notification — so a local pull shows live download progress just like a remote pull's `engine:remote-progress` — and the request settles with the pull's terminal result line. Frames are coalesced (only a change in `stage` or `percent` is emitted) so a chatty engine that streams many byte-progress lines per layer doesn't flood subscribers. The engine's terminal `{"status":"success"}` surfaces as a `stage:"success"` frame; a **failed** pull emits a terminal `stage:"error", percent:-1, message:` frame in addition to the JSON-RPC error, so a UI whose synchronous call already timed out on a long download still converges off "pulling". A CLI-driven pull (LM Studio's `lms get`) has no line-level progress, so it emits one `stage:"pulling"` marker and returns the command's result. On `shutdown` (or stdin EOF) the service stops every running engine first, so none are orphaned. +The operator pulls a model: `engine:action {engine:"ollama", action:"pull_model", params:{name:"llama3.2"}}` issues the manifest-declared `POST 127.0.0.1:{port}/api/pull`. Because the action is `pull_model`, the request is routed through the streaming pull path (not the buffered `engine:action` reader): each `/api/pull` status line is emitted as an `engine:pull-progress` notification — so a local pull shows live download progress just like a remote pull's `engine:remote-progress` — and the request settles with the pull's terminal result line. Frames are coalesced (only a change in `stage` or `percent` is emitted) so a chatty engine that streams many byte-progress lines per layer doesn't flood subscribers. The engine's terminal `{"status":"success"}` surfaces as a `stage:"success"` frame; a **failed** pull emits a terminal `stage:"error", percent:-1, message:` frame in addition to the JSON-RPC error, so a UI whose synchronous call already timed out on a long download still converges off "pulling". A CLI-driven pull (LM Studio's `lms get`) is streamed the same way: its carriage-return redraws are parsed into `stage:"downloading"` frames carrying a percent, after the one `stage:"pulling"` marker. Every frame carries the `model` it belongs to, because pulls for one engine are queued rather than rejected — a request for an engine that is already downloading emits `stage:"queued"` and starts when its predecessor finishes, and a second request for a download already in flight joins it instead of failing. + +The operator cancels a download: `engine:cancel-pull {engine:"ollama", model:"llama3.2"}` stops the transfer and answers only once its partial files are settled, so a UI can hold "Canceling" until the backend is genuinely done. Cancellation never deletes a completed model, and it only deletes partial data when someone asked for it — a pull cut short by shutdown, a dropped remote connection, or the action timeout leaves its partials for the next attempt to resume. Ollama's partial blobs are content-addressed and shared, so one another client on the same daemon is still writing is left alone; LM Studio's cleanup is scoped to the requested quantization's files, and is skipped entirely when the CLI did not confirm the cancellation. A cancel that arrives before its pull has registered is held for that pull, which then stops instead of starting; a cancel for a model with no download outstanding stops nothing and cannot be inherited by a later attempt. On `shutdown` (or stdin EOF) the service stops every running engine first, so none are orphaned. ## 15. Current integration / wiring diff --git a/services/nvpair-engine-manager/testdata/fakeengine/main.go b/services/nvpair-engine-manager/testdata/fakeengine/main.go index 9f12ba82..92d71ef7 100644 --- a/services/nvpair-engine-manager/testdata/fakeengine/main.go +++ b/services/nvpair-engine-manager/testdata/fakeengine/main.go @@ -13,12 +13,14 @@ package main import ( + "bufio" "encoding/json" "fmt" "log" "net" "net/http" "os" + "os/signal" "strconv" "strings" "sync" @@ -70,6 +72,28 @@ func bstr(m map[string]any, k string) string { return "" } +// reportDownloadProgress renders one `lms get` progress redraw at the +// percentage named in os.Args[2], in the CLI's own two-decimal shape. +func reportDownloadProgress() { + if len(os.Args) < 3 { + os.Exit(2) + } + percent, err := strconv.ParseFloat(os.Args[2], 64) + if err != nil { + os.Exit(2) + } + fmt.Printf("\r[==== ] %.2f%%", percent) +} + +// awaitInterrupt blocks until this process is interrupted, the way `lms get` +// waits while it downloads. +func awaitInterrupt() { + signals := make(chan os.Signal, 1) + signal.Notify(signals, os.Interrupt) + defer signal.Stop(signals) + <-signals +} + func main() { if path := os.Getenv("PAIR_TEST_ENV_FILE"); path != "" { values := map[string]string{} @@ -99,6 +123,43 @@ func main() { // exit (no server), standing in for a daemon's control CLI. if len(os.Args) > 1 { switch os.Args[1] { + // The four download subcommands stand in for `lms get` and cover the + // ways it can answer an interrupt. Each takes the percentage to report + // so the caller owns that number and can assert on the same constant it + // passed in. + case "canceldownload": // answers the prompt and confirms the cancellation + reportDownloadProgress() + awaitInterrupt() + fmt.Fprint(os.Stderr, "Continue to download in the background? (Y/N): ") + answer, err := bufio.NewReader(os.Stdin).ReadString('\n') + if err != nil || strings.TrimSpace(answer) != "n" { + os.Exit(2) + } + fmt.Fprintln(os.Stderr, "Download canceled.") + os.Exit(1) + case "downloadcompletes": // finishes as the interrupt lands + reportDownloadProgress() + awaitInterrupt() + fmt.Fprintln(os.Stderr, "Download completed.") + return + case "downloadsilent": // exits without answering the prompt + reportDownloadProgress() + awaitInterrupt() + os.Exit(1) + case "downloadignoresinterrupt": // never acknowledges, so it has to be killed + // Take delivery of the interrupt and do nothing with it. Calling + // signal.Ignore instead is worse than a no-op on Windows: the + // runtime's console handler exits the process for an event no + // receiver wants, which is the opposite of what this stands for. + swallowed := make(chan os.Signal, 1) + signal.Notify(swallowed, os.Interrupt) + defer signal.Stop(swallowed) + reportDownloadProgress() + // Sleeping rather than blocking forever: a bare receive is the + // runtime's deadlock condition, and that panic would exit the very + // process this case needs to outlive its interrupt. + time.Sleep(time.Hour) + os.Exit(2) case "captureargs": // record exact argv for launch-text round-trip tests if len(os.Args) < 3 { os.Exit(2) diff --git a/services/nvpair-tui/ui/engines.go b/services/nvpair-tui/ui/engines.go index 809febfb..9f0ffcda 100644 --- a/services/nvpair-tui/ui/engines.go +++ b/services/nvpair-tui/ui/engines.go @@ -29,11 +29,20 @@ type engineStatus struct { Port int `json:"port"` } +// enginePull identifies one in-flight model download. Several can be +// registered per engine — the engine-manager queues rather than rejects them — +// so the cancel key needs the model, not just the engine. +type enginePull struct { + engine string + model string +} + // enginesView manages local inference engines via the engine-manager // control plane: an installed/running/healthy table plus lifecycle // actions, kept live from engine:state-changed and engine:install-progress. -// It can also pull a model (engine:action{action:"pull_model"}), rendering -// the live engine:pull-progress feed the way remote pulls already show. +// It can also pull a model (engine:action{action:"pull_model"}) and cancel one +// (engine:cancel-pull), rendering the live engine:pull-progress feed the way +// remote pulls already show. type enginesView struct { client *rpc.Client table table.Model @@ -43,6 +52,9 @@ type enginesView struct { input textinput.Model pulling bool pullEngine string + // active is every download this view started, oldest first, so the cancel + // key can name the newest one on the selected engine. + active []enginePull width, height int } @@ -58,13 +70,22 @@ type engineOpMsg struct { err error } +// enginePullDoneMsg retires a download from active once its request settles, +// whether it succeeded or failed. A failed entry left behind stays selectable, +// and the next cancel would target a download that is not running. +type enginePullDoneMsg struct { + pull enginePull + err error +} + var ( - engStartKey = key.NewBinding(key.WithKeys("s"), key.WithHelp("s", "start")) - engStopKey = key.NewBinding(key.WithKeys("x"), key.WithHelp("x", "stop")) - engRestartKey = key.NewBinding(key.WithKeys("r"), key.WithHelp("r", "restart")) - engInstallKey = key.NewBinding(key.WithKeys("i"), key.WithHelp("i", "install")) - engUninstallKey = key.NewBinding(key.WithKeys("u"), key.WithHelp("u", "uninstall")) - engPullKey = key.NewBinding(key.WithKeys("p"), key.WithHelp("p", "pull model")) + engStartKey = key.NewBinding(key.WithKeys("s"), key.WithHelp("s", "start")) + engStopKey = key.NewBinding(key.WithKeys("x"), key.WithHelp("x", "stop")) + engRestartKey = key.NewBinding(key.WithKeys("r"), key.WithHelp("r", "restart")) + engInstallKey = key.NewBinding(key.WithKeys("i"), key.WithHelp("i", "install")) + engUninstallKey = key.NewBinding(key.WithKeys("u"), key.WithHelp("u", "uninstall")) + engPullKey = key.NewBinding(key.WithKeys("p"), key.WithHelp("p", "pull model")) + engCancelPullKey = key.NewBinding(key.WithKeys("c"), key.WithHelp("c", "cancel pull")) ) func newEnginesView(client *rpc.Client) *enginesView { @@ -151,30 +172,49 @@ func (v *enginesView) Update(msg tea.Msg) tea.Cmd { case "engine:pull-progress": var p struct { Engine string `json:"engine"` + Model string `json:"model"` Stage string `json:"stage"` Percent int `json:"percent"` Message string `json:"message"` } _ = decodeParams(msg.Msg.Params, &p) + // The engine-manager queues concurrent pulls per engine, so every + // frame is attributed by model — without it one download's percent + // would be rendered against another's name. + what := p.Engine + if p.Model != "" { + what = p.Engine + " " + p.Model + } // Terminal stages carry no meaningful percent (success is implicitly // 100%; error uses -1), so render them as outcomes rather than a // misleading "success (0%)". A late failure that arrives after the // synchronous call timed out still surfaces here. switch p.Stage { case "success": - v.status = fmt.Sprintf("pull %s: done", p.Engine) + v.status = fmt.Sprintf("pull %s: done", what) case "error": detail := p.Message if detail == "" { detail = "failed" } - v.status = fmt.Sprintf("pull %s failed: %s", p.Engine, detail) + v.status = fmt.Sprintf("pull %s failed: %s", what, detail) + case "queued": + // The engine already has a download running; this one starts + // when that finishes. There is no percent to report yet. + v.status = fmt.Sprintf("pull %s: queued", what) default: - v.status = fmt.Sprintf("pull %s: %s (%d%%)", p.Engine, p.Stage, p.Percent) + v.status = fmt.Sprintf("pull %s: %s (%d%%)", what, p.Stage, p.Percent) } } return nil + case enginePullDoneMsg: + v.retire(msg.pull) + if msg.err != nil { + v.status = fmt.Sprintf("pull %s %s failed: %s", msg.pull.engine, msg.pull.model, msg.err.Error()) + } + return nil + case tea.KeyMsg: return v.handleKey(msg) } @@ -208,6 +248,9 @@ func (v *enginesView) handleKey(msg tea.KeyMsg) tea.Cmd { v.input.Focus() return textinput.Blink } + if key.Matches(msg, engCancelPullKey) { + return v.cancelPull() + } if cmd, handled := v.handleAction(msg); handled { return cmd } @@ -240,16 +283,63 @@ func (v *enginesView) submitPull() tea.Cmd { v.status = "model name required" return nil } - v.status = fmt.Sprintf("pull %s: %s...", engine, model) + pull := enginePull{engine: engine, model: model} + v.active = append(v.active, pull) + v.status = fmt.Sprintf("pull %s %s...", engine, model) params := pullParams(engine, model) return call(v.client, "engine:action", params, func(_ *rpc.Message, err error) tea.Msg { - if err != nil && !errors.Is(err, context.DeadlineExceeded) { - return engineOpMsg{what: "pull " + model, engine: engine, err: err} + // A deadline here is the call timeout, not the download's outcome: the + // pull runs on and the progress feed still reports it, so the entry has + // to stay cancellable. + if errors.Is(err, context.DeadlineExceeded) { + return nil } + return enginePullDoneMsg{pull: pull, err: err} + }) +} + +// cancelPull stops the newest download this view started on the selected +// engine. The engine-manager acknowledges only after the transfer has stopped +// and its partial files are cleaned up, so the reply can be slow enough to hit +// the call timeout — which the progress feed then settles. +func (v *enginesView) cancelPull() tea.Cmd { + engine := v.selectedEngine() + if engine == "" { + return nil + } + pull, ok := v.newestPull(engine) + if !ok { + v.status = "no download in progress on " + engine return nil + } + v.status = fmt.Sprintf("canceling %s %s...", pull.engine, pull.model) + params := map[string]string{"engine": pull.engine, "model": pull.model} + return call(v.client, "engine:cancel-pull", params, func(_ *rpc.Message, err error) tea.Msg { + if err != nil && !errors.Is(err, context.DeadlineExceeded) { + return engineOpMsg{what: "cancel " + pull.model, engine: pull.engine, err: err} + } + return enginePullDoneMsg{pull: pull} }) } +func (v *enginesView) newestPull(engine string) (enginePull, bool) { + for i := len(v.active) - 1; i >= 0; i-- { + if v.active[i].engine == engine { + return v.active[i], true + } + } + return enginePull{}, false +} + +func (v *enginesView) retire(pull enginePull) { + for i, active := range v.active { + if active == pull { + v.active = append(v.active[:i], v.active[i+1:]...) + return + } + } +} + func (v *enginesView) handleAction(msg tea.KeyMsg) (tea.Cmd, bool) { var method, what string switch { @@ -333,7 +423,7 @@ func (v *enginesView) View() string { } func (v *enginesView) Help() []key.Binding { - return []key.Binding{engStartKey, engStopKey, engRestartKey, engInstallKey, engUninstallKey, engPullKey} + return []key.Binding{engStartKey, engStopKey, engRestartKey, engInstallKey, engUninstallKey, engPullKey, engCancelPullKey} } func yesNo(b bool) string { diff --git a/services/nvpair-tui/ui/engines_test.go b/services/nvpair-tui/ui/engines_test.go index e04a2c77..4edf4b94 100644 --- a/services/nvpair-tui/ui/engines_test.go +++ b/services/nvpair-tui/ui/engines_test.go @@ -3,7 +3,14 @@ package ui -import "testing" +import ( + "encoding/json" + "errors" + "strings" + "testing" + + "nvpair-tui/rpc" +) // TestPullParamsSendsBothKeys guards the LM Studio pull fix: the pull params // must carry the model under BOTH "name" (Ollama's /api/pull body key) and @@ -29,3 +36,73 @@ func TestPullParamsSendsBothKeys(t *testing.T) { t.Fatalf(`params["model"] = %q, want "owner/model" (LM Studio reads this key)`, inner["model"]) } } + +// The engine-manager queues concurrent pulls per engine and attributes every +// engine:pull-progress frame by model, so the status line must name the model +// rather than credit one download's percent to another. +func TestPullProgressRendersModelAndQueuedStage(t *testing.T) { + test := func(name, params, want string) { + t.Run(name, func(t *testing.T) { + v := newEnginesView(nil) + v.Update(NotificationMsg{Msg: &rpc.Message{ + Method: "engine:pull-progress", + Params: json.RawMessage(params), + }}) + if v.status != want { + t.Fatalf("status = %q, want %q", v.status, want) + } + }) + } + test("downloading", + `{"engine":"ollama","model":"llama3.2","op":"pull","stage":"downloading","percent":42}`, + "pull ollama llama3.2: downloading (42%)") + test("queued behind another download", + `{"engine":"ollama","model":"qwen3:8b","op":"pull","stage":"queued"}`, + "pull ollama qwen3:8b: queued") + test("success", + `{"engine":"lmstudio","model":"owner/model","op":"pull","stage":"success"}`, + "pull lmstudio owner/model: done") +} + +// The cancel key needs the model, not just the engine: several downloads can be +// registered against one engine at a time. +func TestCancelTargetsNewestPullOnEngine(t *testing.T) { + v := newEnginesView(nil) + v.active = []enginePull{ + {engine: "ollama", model: "llama3.2"}, + {engine: "lmstudio", model: "owner/model"}, + {engine: "ollama", model: "qwen3:8b"}, + } + pull, ok := v.newestPull("ollama") + if !ok || pull.model != "qwen3:8b" { + t.Fatalf("newestPull = %+v (ok=%v), want qwen3:8b", pull, ok) + } + v.retire(pull) + if pull, ok = v.newestPull("ollama"); !ok || pull.model != "llama3.2" { + t.Fatalf("after retiring, newestPull = %+v (ok=%v), want llama3.2", pull, ok) + } + v.retire(pull) + if _, ok = v.newestPull("ollama"); ok { + t.Fatal("ollama still reports a download after both were retired") + } + if pull, ok = v.newestPull("lmstudio"); !ok || pull.model != "owner/model" { + t.Fatalf("another engine's download was retired: %+v (ok=%v)", pull, ok) + } +} + +// A pull whose request failed is over. Leaving it tracked kept it selectable, +// so the next cancel key would target a download that was never running. +func TestFailedPullIsRetiredAndNotCancelable(t *testing.T) { + v := newEnginesView(nil) + failed := enginePull{engine: "ollama", model: "llama3.2"} + v.active = []enginePull{failed} + + v.Update(enginePullDoneMsg{pull: failed, err: errors.New("engine ollama is not running")}) + + if _, ok := v.newestPull("ollama"); ok { + t.Fatal("a failed download is still offered as the cancel target") + } + if !strings.Contains(v.status, "engine ollama is not running") { + t.Fatalf("status = %q, want the failure reason", v.status) + } +} From 946fe0109add0c166d081c1e020aaaea4244a1fd Mon Sep 17 00:00:00 2001 From: Chris Kelsey Date: Tue, 22 Sep 2026 15:05:29 -0700 Subject: [PATCH 02/12] fix(engine-manager): stop LM Studio cleanup deleting partials it does not own cleanupLMSPartials treated a partial as this pull's whenever it differed from the pre-download snapshot. Bytes gained during the pull are not evidence of ownership: another client writing into the same repository directory produces exactly the same difference, and the model hub pulls unpinned ids, so quant is empty and the quantization filter excludes nothing. Every partial in the repository was reachable. Match Ollama's rule instead - presence in the snapshot disqualifies a path outright - and treat a nil snapshot as a failed inspection that deletes nothing rather than as an empty directory. A cancelled resume now keeps its own partial, since nothing on disk tells a resume target from another client's file. That costs disk the next attempt reuses, where the alternative costs a live download its bytes. Signed-off-by: Chris Kelsey --- services/nvpair-engine-manager/lmspull.go | 41 ++++++--- .../nvpair-engine-manager/pullcancel_test.go | 85 ++++++++++++++++++- 2 files changed, 111 insertions(+), 15 deletions(-) diff --git a/services/nvpair-engine-manager/lmspull.go b/services/nvpair-engine-manager/lmspull.go index 71d1e5d6..c3f2c548 100644 --- a/services/nvpair-engine-manager/lmspull.go +++ b/services/nvpair-engine-manager/lmspull.go @@ -251,25 +251,37 @@ func lmsErrorDetail(text string) string { return "…" + string(runes[len(runes)-lmsErrorDetailMax:]) } -// Remove only the partial files this `lms get` produced, preserving completed +// Remove only the partial files this `lms get` created, preserving completed // shards. The LM Studio app and any other client download into the same // repository directory, so a file is deleted only when it carries the requested -// quantization, differs from the snapshot taken before this download started, -// and has stopped moving. Never delete a model directory or follow a symlink -// outside the cache. +// quantization, was absent from the snapshot taken before this download +// started, and has stopped moving. Never delete a model directory or follow a +// symlink outside the cache. +// +// Presence in the snapshot is disqualifying on its own; what the file has done +// since is not consulted. Bytes gained during this pull look identical whether +// another client is writing them or `lms get` is resuming the file in place, +// and a writer that has merely paused cannot be told from one that finished. A +// resumed download therefore keeps its partial through a cancellation, which +// costs disk the vendor reuses on the next attempt; the alternative costs +// another client its transfer. func cleanupLMSPartials(ctx context.Context, root, model string, before map[string]os.FileInfo) error { + // No snapshot means the inspection failed, not that the directory was + // empty, and nothing is attributable without one. + if before == nil { + return nil + } files, err := lmsPartialFiles(root, model) if err != nil { return err } - touched := make([]string, 0, len(files)) - for path, info := range files { - if old, ok := before[path]; ok && old.Size() == info.Size() && old.ModTime().Equal(info.ModTime()) { - continue + created := make([]string, 0, len(files)) + for path := range files { + if _, existed := before[path]; !existed { + created = append(created, path) } - touched = append(touched, path) } - stable, _ := quiescentPaths(ctx, touched) + stable, _ := quiescentPaths(ctx, created) for _, candidate := range stable { if _, err := removeIfUnchanged(candidate); err != nil { return err @@ -280,6 +292,11 @@ func cleanupLMSPartials(ctx context.Context, root, model string, before map[stri // lmsPartialFiles lists the in-progress downloads in a model's repository // directory that could belong to the requested quantization. +// +// A nil map means the listing could not be made — the reference names no +// repository — and is distinct from an empty one, which means the repository +// holds no partials. cleanupLMSPartials deletes nothing on the former, so the +// two must not be conflated. func lmsPartialFiles(root, model string) (map[string]os.FileInfo, error) { owner, repo, ok := lmsOwnerName(model) if !ok { @@ -294,14 +311,14 @@ func lmsPartialFiles(root, model string) (map[string]os.FileInfo, error) { target := filepath.Join(root, owner, repo) resolvedRoot, err := filepath.EvalSymlinks(root) if os.IsNotExist(err) { - return nil, nil + return map[string]os.FileInfo{}, nil } if err != nil { return nil, err } resolvedTarget, err := filepath.EvalSymlinks(target) if os.IsNotExist(err) { - return nil, nil + return map[string]os.FileInfo{}, nil } if err != nil { return nil, err diff --git a/services/nvpair-engine-manager/pullcancel_test.go b/services/nvpair-engine-manager/pullcancel_test.go index 2cc85662..d3d8fe59 100644 --- a/services/nvpair-engine-manager/pullcancel_test.go +++ b/services/nvpair-engine-manager/pullcancel_test.go @@ -318,12 +318,18 @@ func TestStopLMSDownloadWhenTheInterruptCannotBeDelivered(t *testing.T) { func TestLMSPartialCleanupPreservesCompletedShards(t *testing.T) { root := t.TempDir() target := filepath.Join(root, "owner", "repo") - partial := filepath.Join(target, "downloading_model-Q4_K_M.gguf.part") complete := filepath.Join(target, "model-00001-of-00002.gguf") - writePartial(t, partial, "data") writePartial(t, complete, "data") - if err := cleanupLMSPartials(settledContext(), root, "https://huggingface.co/owner/repo@Q4_K_M", nil); err != nil { + const model = "https://huggingface.co/owner/repo@Q4_K_M" + before, err := lmsPartialFiles(root, model) + if err != nil { + t.Fatalf("snapshot partials: %v", err) + } + partial := filepath.Join(target, "downloading_model-Q4_K_M.gguf.part") + writePartial(t, partial, "data") + + if err := cleanupLMSPartials(settledContext(), root, model, before); err != nil { t.Fatalf("cleanup: %v", err) } assertRemoved(t, partial) @@ -406,6 +412,79 @@ func TestLMSPartialCleanupPreservesTheRequestedQuantizationWhileItGrows(t *testi assertPresent(t, shared) } +// Growing since the snapshot is not evidence of ownership. A writer that paused +// looks exactly like one that finished, so a partial already present when this +// pull started stays another client's however many bytes it gained meanwhile. +// +// The model hub pulls unpinned ids, so quant is empty and the quantization +// filter excludes nothing — every partial in the repository reaches this rule, +// which is why it has to be the strict one. +func TestLMSPartialCleanupPreservesAPreexistingPartialThatGrew(t *testing.T) { + root := t.TempDir() + target := filepath.Join(root, "lmstudio-community", "Qwen3-8B-GGUF") + const model = "lmstudio-community/Qwen3-8B-GGUF" + theirs := filepath.Join(target, "downloading_Qwen3-8B-Q4_K_M.gguf.part") + writePartial(t, theirs, "app download") + + before, err := lmsPartialFiles(root, model) + if err != nil { + t.Fatalf("snapshot partials: %v", err) + } + if len(before) != 1 { + t.Fatalf("snapshot missed the pre-existing partial: %v", before) + } + // The app writes more while our download runs, then stalls, so the file is + // perfectly still by the time cleanup observes it. + growFile(t, theirs) + + if err := cleanupLMSPartials(settledContext(), root, model, before); err != nil { + t.Fatalf("cleanup: %v", err) + } + assertPresent(t, theirs) +} + +// `lms get` resumes in place, so the same rule means a cancelled resume leaves +// its own partial behind. That is deliberate: nothing on disk tells a resume +// target apart from another client's file, the bytes stay useful to the next +// attempt, and the alternative is deleting a live download whenever the guess +// goes the other way. Files this pull actually created are still removed. +func TestLMSPartialCleanupLeavesAResumedDownloadsPartial(t *testing.T) { + root := t.TempDir() + target := filepath.Join(root, "lmstudio-community", "Qwen3-8B-GGUF") + const model = "lmstudio-community/Qwen3-8B-GGUF" + resuming := filepath.Join(target, "downloading_Qwen3-8B-Q4_K_M.gguf.part") + writePartial(t, resuming, "bytes the attempt being resumed left behind") + + before, err := lmsPartialFiles(root, model) + if err != nil { + t.Fatalf("snapshot partials: %v", err) + } + growFile(t, resuming) + fresh := filepath.Join(target, "downloading_Qwen3-8B-Q6_K.gguf.part") + writePartial(t, fresh, "a shard this pull started from nothing") + + if err := cleanupLMSPartials(settledContext(), root, model, before); err != nil { + t.Fatalf("cleanup: %v", err) + } + assertPresent(t, resuming) + assertRemoved(t, fresh) +} + +// A nil snapshot means the inspection failed, not that the directory was empty. +// Attribution is impossible without one, so cleanup deletes nothing — the rule +// cleanupOllamaPartials already applies. +func TestLMSPartialCleanupDeletesNothingWithoutASnapshot(t *testing.T) { + root := t.TempDir() + target := filepath.Join(root, "owner", "repo") + partial := filepath.Join(target, "downloading_model-Q4_K_M.gguf.part") + writePartial(t, partial, "data") + + if err := cleanupLMSPartials(settledContext(), root, "owner/repo@Q4_K_M", nil); err != nil { + t.Fatalf("cleanup: %v", err) + } + assertPresent(t, partial) +} + // The cancel may only be acknowledged once the partial files are gone, because // the UI reads that acknowledgement as "it is safe to download this again". // From 8f29e683dc6d0fdda3a530de164ba3367946aea4 Mon Sep 17 00:00:00 2001 From: Chris Kelsey Date: Tue, 22 Sep 2026 15:25:07 -0700 Subject: [PATCH 03/12] fix(engine-manager): give a remote cancel the readiness response budget A peer answering cancel_pull interrupts the CLI, waits for it to acknowledge, and removes partial files before it writes a header. The ordinary 30s response-header budget can expire while that is still in progress, so the initiator saw a transport timeout from a cancel the peer was carrying out. Route controlCancelPullPath through the readyHTTP pool, alongside the other endpoints whose peer withholds its header while doing slow work in its own handler. Being cut off never undid the cancellation - that is latched before the peer starts waiting - so the damage was confined to reporting: a cancel that was succeeding came back as failed, and the row rolled back to "Downloading" under a transfer that was stopping. Signed-off-by: Chris Kelsey --- services/nvpair-engine-manager/remoteclient.go | 15 ++++++++++++--- .../nvpair-engine-manager/remoteclient_test.go | 7 +++++++ 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/services/nvpair-engine-manager/remoteclient.go b/services/nvpair-engine-manager/remoteclient.go index 9bc7ca91..98c977ac 100644 --- a/services/nvpair-engine-manager/remoteclient.go +++ b/services/nvpair-engine-manager/remoteclient.go @@ -37,8 +37,9 @@ type remoteClient struct { forget func() } -// waitsForEngineReadiness reports whether a peer may only answer after an -// engine is healthy or an Ollama model is loaded, which the ordinary 30s +// waitsForEngineReadiness reports whether a peer may only answer after slow +// work inside its own handler — an engine becoming healthy, an Ollama model +// loading, or a download being stopped — which the ordinary 30s // response-header budget cannot cover. // // Cutting such a call off is worse than slow. The initiator's cancellation @@ -51,10 +52,18 @@ func waitsForEngineReadiness(path, engine string) bool { } // controlDeletePath: LM Studio's delete_model declares restart_after, so the // peer replies only after the post-delete restart is ready. + // + // controlCancelPullPath: the peer interrupts the CLI, waits for it to + // acknowledge, and removes partial files before writing a header. The + // ordinary budget can expire while that is still in progress. Here the + // cancellation is already latched before the peer starts waiting, so being + // cut off does not undo it — the damage is that the initiator reports a + // cancel that is succeeding as failed, and the row rolls back to + // "Downloading" under a transfer that is stopping. if path == controlLoadPath { return engine == "ollama" } - return path == controlStartPath || path == controlDeletePath + return path == controlStartPath || path == controlDeletePath || path == controlCancelPullPath } func newRemoteHTTPClient(base *http.Transport, responseHeaderTimeout time.Duration) *http.Client { diff --git a/services/nvpair-engine-manager/remoteclient_test.go b/services/nvpair-engine-manager/remoteclient_test.go index c4d16e67..7400b2a5 100644 --- a/services/nvpair-engine-manager/remoteclient_test.go +++ b/services/nvpair-engine-manager/remoteclient_test.go @@ -84,6 +84,13 @@ func TestRemoteReadinessBudgetCoversEngineStartupAllowance(t *testing.T) { {controlDeletePath, "lmstudio", true}, {controlLoadPath, "ollama", true}, {controlLoadPath, "lmstudio", false}, + // Cancelling is the peer stopping a live transfer, not answering a + // question: it interrupts the CLI, waits for the acknowledgement, and + // cleans up before writing a header. Both engines can outlast the + // ordinary budget, and being cut off there reports a cancel that + // succeeded as failed. + {controlCancelPullPath, "lmstudio", true}, + {controlCancelPullPath, "ollama", true}, {controlStopPath, "ollama", false}, {controlUnloadPath, "ollama", false}, {controlEnginesPath, "", false}, From 4f3e24de3bedf4a807ea0718617dfd0e32c410ea Mon Sep 17 00:00:00 2001 From: Chris Kelsey Date: Tue, 22 Sep 2026 15:25:33 -0700 Subject: [PATCH 04/12] fix(engine-manager): hold a remote cancel until every duplicate has ended The peer joins a duplicate pull onto the download already in flight, so two requests can share one attempt. Their release closed attempt.accepted on the first one to end, which freed a waiting cancel while the other was still on its way to the peer. The cancel then arrived at a peer with nothing registered, was answered as a cancel for nothing, and left the transfer running under a row stuck on "Canceling". Close accepted only once refs reaches zero. Acceptance keeps its own closer, so the peer taking the pull still frees the cancel immediately, however many requests are outstanding. Size remotePullGateWindow to the pull's pre-header budget while here. The old 10s could elapse with the pull legitimately still in flight, producing the same cancel-for-nothing. The window is a backstop, not the mechanism: dial, handshake and header chain only on success, so a pull that never lands releases the gate itself well inside the new value. Signed-off-by: Chris Kelsey --- services/nvpair-engine-manager/remote.go | 41 +++++++++++++---- services/nvpair-engine-manager/remote_test.go | 46 +++++++++++++++++-- 2 files changed, 75 insertions(+), 12 deletions(-) diff --git a/services/nvpair-engine-manager/remote.go b/services/nvpair-engine-manager/remote.go index bc9b44cd..b9cce2bd 100644 --- a/services/nvpair-engine-manager/remote.go +++ b/services/nvpair-engine-manager/remote.go @@ -80,12 +80,27 @@ type remotePullAttempt struct { refs int } -// remotePullGateWindow bounds how long a cancel waits for its pull to reach the -// peer. Dialing a peer's ec surface can be slow, but a cancel that waits -// indefinitely would hold its own JSON-RPC response open for as long as the -// pull runs. Sending it late is better than never: the peer either has the -// download by then or has nothing to stop. -const remotePullGateWindow = 10 * time.Second +// remotePullGateWindow is the backstop under a cancel waiting for its pull to +// reach the peer, not the mechanism. The two signals that actually free it +// both close attempt.accepted: the peer taking the pull, or every request for +// it ending. One of them always arrives, because the pull's own transport +// gives up by itself — 30s to dial and 30s for the handshake, clustertrust's +// fallback when no PeerClientOptions.Timeout is set, then +// remoteResponseHeaderTimeout for the header, which is the point stream +// reports acceptance. +// +// Those three phases chain only on success, so the sum is a ceiling the pull +// cannot actually reach: getting as far as the header wait means the dial and +// the handshake each finished inside their own 30s, putting the failure +// strictly under this. The cancel's timer then starts later still, when the +// cancel arrives rather than when the pull did. +// +// A timer that fires first is the bug rather than the safety net: it frees the +// cancel while the pull is still legitimately in flight, and a cancel arriving +// at a peer that has nothing registered is answered as a cancel for nothing, +// leaving the transfer running under a row stuck on "Canceling". Waiting is +// the cheaper error, because a pull that never lands releases this itself. +const remotePullGateWindow = 90 * time.Second func remotePullKey(node, engine, model string) string { return node + "\x00" + engine + "\x00" + model @@ -135,13 +150,21 @@ func (g *remotePullGate) register(key string, isPull bool) func() { once.Do(func() { g.mu.Lock() attempt.refs-- - if attempt.refs == 0 { + last := attempt.refs == 0 + if last { delete(g.pulls, key) } g.mu.Unlock() // An attempt that ended without the peer accepting it has nothing - // left for a cancel to chase. - attempt.settle.Do(func() { close(attempt.accepted) }) + // left for a cancel to chase — but only once every request sharing + // it has ended. One duplicate failing while another is still on + // its way to the peer is not the attempt ending, and freeing the + // cancel there sends it ahead of a pull that may still be + // accepted. Acceptance has its own closer and does not wait on + // this count. + if last { + attempt.settle.Do(func() { close(attempt.accepted) }) + } }) } } diff --git a/services/nvpair-engine-manager/remote_test.go b/services/nvpair-engine-manager/remote_test.go index fb3d4b44..c496c1e2 100644 --- a/services/nvpair-engine-manager/remote_test.go +++ b/services/nvpair-engine-manager/remote_test.go @@ -135,8 +135,12 @@ func TestRemotePullGateReleasesACancelWhenThePullNeverLands(t *testing.T) { } // Two requests for the same remote download share one attempt, because the peer -// joins the duplicate onto the transfer already in flight. The cancel waits for -// the peer to accept, not for both requests to end. +// joins the duplicate onto the transfer already in flight. What releases the +// cancel is the peer accepting, which happens however many requests are +// outstanding — not the first of them ending. While one is still on its way to +// the peer there is still a pull to chase, and letting the cancel go early +// sends it to a peer with nothing registered, which is answered as a cancel for +// nothing and leaves the row stuck on "Canceling". func TestRemotePullGateSharesOneAttemptAcrossDuplicateRequests(t *testing.T) { var gate remotePullGate key := remotePullKey("uuid-b", "ollama", "demo") @@ -150,9 +154,17 @@ func TestRemotePullGateSharesOneAttemptAcrossDuplicateRequests(t *testing.T) { gate.awaitAccepted(context.Background(), key) }() select { + case <-waited: + t.Fatal("cancel was sent while the second request was still in flight") + case <-time.After(50 * time.Millisecond): + } + + // Acceptance frees it without that request having ended. + gate.accepted(key)() + select { case <-waited: case <-time.After(10 * time.Second): - t.Fatal("cancel kept waiting after the first of two requests ended") + t.Fatal("cancel kept waiting after the peer accepted the pull") } second() @@ -164,6 +176,34 @@ func TestRemotePullGateSharesOneAttemptAcrossDuplicateRequests(t *testing.T) { } } +// When no request reaches the peer, the last one ending is what frees the +// cancel — the shared attempt has to be exhausted, not merely reduced. +func TestRemotePullGateHoldsACancelUntilEveryDuplicateHasEnded(t *testing.T) { + var gate remotePullGate + key := remotePullKey("uuid-b", "ollama", "demo") + first := gate.register(key, true) + second := gate.register(key, true) + + first() + waited := make(chan struct{}) + go func() { + defer close(waited) + gate.awaitAccepted(context.Background(), key) + }() + select { + case <-waited: + t.Fatal("cancel was sent while the second request was still in flight") + case <-time.After(50 * time.Millisecond): + } + + second() + select { + case <-waited: + case <-time.After(10 * time.Second): + t.Fatal("cancel outlived every request it was chasing") + } +} + // A cancel for a download nobody requested has nothing to chase, so it goes // straight through rather than spending the whole window. func TestRemotePullGatePassesACancelWithNoPullRegistered(t *testing.T) { From 2fc0e3052aefffc58255b0cec0ddb7b70e4ec2a6 Mon Sep 17 00:00:00 2001 From: Chris Kelsey Date: Tue, 22 Sep 2026 15:25:48 -0700 Subject: [PATCH 05/12] fix(engine-manager): require a run of still passes before deleting a partial A partial this pull created can still be shared, because another client can coalesce onto it afterwards. Cleanup left a file alone while it was moving, but one quiescence check per pass was not enough once the retry loop ran several of them: each pass observes its own window, so every extra attempt was another independent chance for a writer that merely paused to look finished. A transfer does pause, which made the loop - there to catch up with Ollama's asynchronous blob release - likeliest to delete a shared file exactly when it tried hardest. Carry a per-path count of consecutive still passes across one cancellation and delete only at partialCleanupSteadyPasses. Any movement, or the path dropping out of the candidate set, resets the run, so a file now has to hold still across the whole run to qualify. This pull's own released partial pays one extra pass for it. Signed-off-by: Chris Kelsey --- services/nvpair-engine-manager/ollamapull.go | 44 ++++++++++- .../nvpair-engine-manager/ollamapull_test.go | 74 ++++++++++++++----- 2 files changed, 98 insertions(+), 20 deletions(-) diff --git a/services/nvpair-engine-manager/ollamapull.go b/services/nvpair-engine-manager/ollamapull.go index ebfd7f1c..2e844d8c 100644 --- a/services/nvpair-engine-manager/ollamapull.go +++ b/services/nvpair-engine-manager/ollamapull.go @@ -54,17 +54,40 @@ func ollamaPartialSnapshot(root string) (ollamaPartialsBefore, error) { return before, nil } +// ollamaSteadyPasses counts, per path, how many cleanup passes in a row have +// found the file holding still. One cancellation's retry loop keeps one of +// these; it must not be nil, since a pass records into it. +type ollamaSteadyPasses map[string]int + +// partialCleanupSteadyPasses is how many consecutive passes must find a +// candidate unchanged before it is deleted. +// +// A single pass is not enough once there are several of them. Each observes +// its own window, so every extra pass is another independent chance for a +// writer that merely paused to look finished, and a transfer does pause. That +// makes the retry loop — there to catch up with Ollama's asynchronous blob +// release — likeliest to delete a shared file exactly when it tries hardest, +// which is backwards. Requiring a run instead means any movement resets the +// count, so a file has to hold still across the whole run to qualify, while +// this pull's own released partial pays one extra pass. +const partialCleanupSteadyPasses = 2 + // cleanupOllamaPartials removes the partial blobs this pull both reported and // created. Completed blobs may be shared by installed models, and a partial // that predates this download belongs to someone else; both are deliberately // retained. A file this pull created can still be shared — another client can // coalesce onto it afterwards — so one that is still moving is left alone and // reported through busy, letting the caller look again. +// +// steady carries the consecutive-stillness count between passes of one +// cancellation, and a candidate is removed only once it has reached +// partialCleanupSteadyPasses. func cleanupOllamaPartials( ctx context.Context, root string, digests map[string]bool, before ollamaPartialsBefore, + steady ollamaSteadyPasses, ) (busy bool, err error) { if before == nil { return false, nil @@ -80,7 +103,24 @@ func cleanupOllamaPartials( } } stable, busy := quiescentPaths(ctx, created) + held := make(map[string]bool, len(stable)) for _, candidate := range stable { + held[candidate.path] = true + } + // Anything that moved, vanished, or stopped being a candidate this pass + // starts its run over. + for path := range steady { + if !held[path] { + delete(steady, path) + } + } + for _, candidate := range stable { + steady[candidate.path]++ + if steady[candidate.path] < partialCleanupSteadyPasses { + // Still enough for now, but not for long enough to act on. + busy = true + continue + } removed, err := removeIfUnchanged(candidate) if err != nil { return true, err @@ -88,6 +128,7 @@ func cleanupOllamaPartials( if !removed { // Claimed between the last observation and the unlink. Leave it // and let the retry decide. + delete(steady, candidate.path) busy = true } } @@ -160,10 +201,11 @@ func cleanupOllamaAfterCancel( digests map[string]bool, before ollamaPartialsBefore, ) error { + steady := make(ollamaSteadyPasses) deadline := time.Now().Add(partialCleanupBudget) for { time.Sleep(partialCleanupRetryInterval) - busy, err := cleanupOllamaPartials(ctx, root, digests, before) + busy, err := cleanupOllamaPartials(ctx, root, digests, before, steady) if err == nil && !busy { return nil } diff --git a/services/nvpair-engine-manager/ollamapull_test.go b/services/nvpair-engine-manager/ollamapull_test.go index 27663a2b..f19ef6bb 100644 --- a/services/nvpair-engine-manager/ollamapull_test.go +++ b/services/nvpair-engine-manager/ollamapull_test.go @@ -8,6 +8,29 @@ import ( "testing" ) +// cleanupRun drives the consecutive passes a candidate must hold still through +// before cleanup will remove it, which is what cleanupOllamaAfterCancel does +// inside its budget, and reports the last pass's busy result. +func cleanupRun( + t *testing.T, + ctx context.Context, + root string, + digests map[string]bool, + before ollamaPartialsBefore, +) bool { + t.Helper() + steady := make(ollamaSteadyPasses) + var busy bool + for pass := range partialCleanupSteadyPasses { + var err error + busy, err = cleanupOllamaPartials(ctx, root, digests, before, steady) + if err != nil { + t.Fatalf("cleanup pass %d: %v", pass, err) + } + } + return busy +} + // Cancelling a download removes the partial blobs that download created and // nothing else. A completed blob can be a layer of an installed model, and a // partial for a digest this pull never reported belongs to another transfer. @@ -29,11 +52,7 @@ func TestOllamaPartialCleanupRemovesOnlyThisPullsPartials(t *testing.T) { // "../escape" stands for a digest the engine never could have reported; it // must not be turned into a path. digests := map[string]bool{digestA: true, "../escape": true} - busy, err := cleanupOllamaPartials(settledContext(), root, digests, before) - if err != nil { - t.Fatalf("cleanup: %v", err) - } - if busy { + if cleanupRun(t, settledContext(), root, digests, before) { t.Error("settled partials were reported busy") } for _, path := range expectDelete { @@ -58,11 +77,7 @@ func TestOllamaPartialCleanupPreservesPartialsOlderThanThePull(t *testing.T) { t.Fatalf("snapshot partials: %v", err) } - busy, err := cleanupOllamaPartials(settledContext(), root, map[string]bool{digestA: true}, before) - if err != nil { - t.Fatalf("cleanup: %v", err) - } - if busy { + if cleanupRun(t, settledContext(), root, map[string]bool{digestA: true}, before) { t.Error("a partial the snapshot excluded was reported busy") } assertPresent(t, stalled) @@ -74,11 +89,7 @@ func TestOllamaPartialCleanupWithoutASnapshotDeletesNothing(t *testing.T) { root := t.TempDir() partial := writeBlob(t, root, blobA+"-partial") - busy, err := cleanupOllamaPartials(settledContext(), root, map[string]bool{digestA: true}, nil) - if err != nil { - t.Fatalf("cleanup: %v", err) - } - if busy { + if cleanupRun(t, settledContext(), root, map[string]bool{digestA: true}, nil) { t.Error("cleanup that ran on no snapshot reported busy") } assertPresent(t, partial) @@ -98,12 +109,37 @@ func TestOllamaPartialCleanupPreservesAPartialAnotherClientIsWriting(t *testing. shared := writeBlob(t, root, blobC+"-partial") ctx := withPartialSettle(context.Background(), func() { growFile(t, shared) }) - busy, err := cleanupOllamaPartials(ctx, root, map[string]bool{digestC: true}, before) + if !cleanupRun(t, ctx, root, map[string]bool{digestC: true}, before) { + t.Error("a partial another client was still writing was not reported busy") + } + assertPresent(t, shared) +} + +// A transfer does not write continuously. The client sharing this partial goes +// quiet for a whole observation window and then resumes, and cleanup runs +// several passes inside its retry budget — so deciding afresh on each pass +// hands a paused writer a new chance to look finished every time, and the file +// is likeliest to be deleted precisely when cleanup tries hardest. +func TestOllamaPartialCleanupKeepsAPartialSharedByAnIntermittentWriter(t *testing.T) { + root := t.TempDir() + before, err := ollamaPartialSnapshot(root) if err != nil { - t.Fatalf("cleanup: %v", err) + t.Fatalf("snapshot partials: %v", err) } - if !busy { - t.Error("a partial another client was still writing was not reported busy") + shared := writeBlob(t, root, blobC+"-partial") + + writing := true + ctx := withPartialSettle(context.Background(), func() { + if writing { + growFile(t, shared) + } + }) + steady := make(ollamaSteadyPasses) + for pass := range 4 { + writing = pass%2 == 0 + if _, err := cleanupOllamaPartials(ctx, root, map[string]bool{digestC: true}, before, steady); err != nil { + t.Fatalf("cleanup pass %d: %v", pass, err) + } } assertPresent(t, shared) } From 425fff94457b5d3b2f596cb104438a8c30edf0fe Mon Sep 17 00:00:00 2001 From: Chris Kelsey Date: Tue, 22 Sep 2026 16:06:06 -0700 Subject: [PATCH 06/12] fix(engine-manager): fix what a cancellation may report and may delete Two defects in the same path, both about the evidence a cancellation acts on. The download output was shared without synchronisation. The exec copy goroutine writes lmsDownloadOutput while the goroutine handling the cancellation reads it to decide whether any partial file may be deleted, and a process the worker gives up on keeps its writer alive past the call that started it - so the reader in pullLMSModel and stopLMSDownload writing cancelled back to false genuinely overlap. Guard every mutable field with a mutex and read them through state(), which returns one consistent snapshot so a caller cannot pair a stale completed with a fresh cancelled. The prompt answer and the progress callback run outside the lock, since each reaches outside the writer and the stdin write can block on the child. Clearing cancelled did not hold anyway. The process is still running, so its next redraw could set the flag again and re-arm deletion of files nothing confirmed had stopped. disown() withdraws the confirmation for good, and the abandonment path now branches on the reap result it discarded, so a process that outlasts the kill is reported as its own outcome. Separately, both engines folded a cleanup failure into the pull's result, which reported a failed cancellation for a transfer that had genuinely stopped and failed the pull's own request with it. Leftover bytes are the vendor's to resume, so this denied the one outcome that did happen and left the row on "Canceling" over a download that was already gone. Warn and keep the cancellation. The race needs -race to demonstrate, which cannot run here (CGO_ENABLED=0, no C compiler), so the concurrency test asserts nothing by itself and is there for CI. The disown and cleanup-failure behaviours are asserted directly. Signed-off-by: Chris Kelsey --- services/nvpair-engine-manager/lmspull.go | 129 ++++++++++++---- services/nvpair-engine-manager/pull.go | 9 +- .../nvpair-engine-manager/pullcancel_test.go | 140 +++++++++++++++++- 3 files changed, 245 insertions(+), 33 deletions(-) diff --git a/services/nvpair-engine-manager/lmspull.go b/services/nvpair-engine-manager/lmspull.go index c3f2c548..3f1ec9af 100644 --- a/services/nvpair-engine-manager/lmspull.go +++ b/services/nvpair-engine-manager/lmspull.go @@ -8,12 +8,14 @@ import ( "encoding/json" "fmt" "io" + "log/slog" "os" "os/exec" "path/filepath" "regexp" "strconv" "strings" + "sync" "time" ) @@ -21,47 +23,104 @@ type downloadOutputKey struct{} // LMS redraws with carriage returns and ANSI controls, not newlines. // A bounded tail handles split writes without keeping the entire download log. +// +// mu guards every mutable field. The exec copy goroutine writes them while the +// goroutine that asked for the cancellation reads them to decide whether any +// partial file may be deleted, and a process this worker gives up on keeps its +// writer alive past the call that started it — so the reader and the writer +// genuinely overlap rather than merely appearing to. type lmsDownloadOutput struct { + mu sync.Mutex text string lastPercent int cancelled bool completed bool answered bool - stdin io.Writer - progress func(int) + // disowned records that this worker stopped waiting on the process behind + // this writer. Its later output cannot be this cancellation's + // confirmation, so state reports no cancellation once it is set, and a + // "Download canceled." arriving afterwards cannot re-arm file deletion. + disowned bool + stdin io.Writer + progress func(int) } var lmsPercentPattern = regexp.MustCompile(`\]\s+(\d+(?:\.\d+)?)%`) var ansiPattern = regexp.MustCompile("\x1b\\[[0-9;?]*[a-zA-Z]") func (w *lmsDownloadOutput) Write(data []byte) (int, error) { + w.mu.Lock() w.text += string(data) clean := ansiPattern.ReplaceAllString(w.text, "") - if strings.Contains(clean, "Continue to download in the background?") && !w.answered { - if _, err := io.WriteString(w.stdin, "n\n"); err != nil { - return 0, err - } - w.answered = true - } + answer := strings.Contains(clean, "Continue to download in the background?") && !w.answered w.cancelled = w.cancelled || strings.Contains(clean, "Download canceled.") w.completed = w.completed || strings.Contains(clean, "Download completed.") + percent := -1 matches := lmsPercentPattern.FindAllStringSubmatch(clean, -1) if len(matches) > 0 { value, err := strconv.ParseFloat(matches[len(matches)-1][1], 64) if err == nil { - percent := max(0, min(100, int(value))) - if percent != w.lastPercent { - w.lastPercent = percent - w.progress(percent) + latest := max(0, min(100, int(value))) + if latest != w.lastPercent { + w.lastPercent = latest + percent = latest } } } if len(w.text) > 8192 { w.text = w.text[len(w.text)-8192:] } + stdin := w.stdin + w.mu.Unlock() + + // Declining the prompt and publishing progress both reach outside this + // writer, and the stdin write can block on the child. Neither runs under + // the lock the cancelling goroutine needs to read the state above. + if answer { + if _, err := io.WriteString(stdin, "n\n"); err != nil { + return 0, err + } + w.mu.Lock() + w.answered = true + w.mu.Unlock() + } + if percent >= 0 { + w.progress(percent) + } return len(data), nil } +// state reports the output as one consistent snapshot, so a caller cannot pair +// a stale completed with a fresh cancelled. +func (w *lmsDownloadOutput) state() (cancelled, completed bool, text string) { + w.mu.Lock() + defer w.mu.Unlock() + return w.cancelled && !w.disowned, w.completed, w.text +} + +// reset clears the previous run's output. It is called before the process +// starts, where nothing is writing yet, but takes the lock anyway: an earlier +// disowned process may still hold this writer. +func (w *lmsDownloadOutput) reset(stdin io.Writer) { + w.mu.Lock() + defer w.mu.Unlock() + w.stdin = stdin + w.text = "" + w.cancelled = false + w.completed = false + w.answered = false +} + +// disown withdraws this writer's cancellation confirmation for good, for a +// process the worker could not reclaim. Clearing the flag alone would not +// hold: the process is still running, so its next redraw could set it again +// and delete files on the strength of a stop nothing confirmed. +func (w *lmsDownloadOutput) disown() { + w.mu.Lock() + defer w.mu.Unlock() + w.disowned = true +} + func (e *Executor) pullLMSModel(ctx context.Context, st *engineState, engine, model string) (json.RawMessage, error) { before, err := lmsPartialFiles(lmstudioModelsDir(), model) if err != nil { @@ -85,11 +144,17 @@ func (e *Executor) pullLMSModel(ctx context.Context, st *engineState, engine, mo // context also dies when PAIR quits mid-download, when a remote initiator // disconnects, and when the action timeout elapses — and `lms get` is built // to resume every one of those on the next attempt. - if err != nil && cancelRequested(ctx) && output.cancelled { + cancelled, _, _ := output.state() + if err != nil && cancelRequested(ctx) && cancelled { // Cleanup outlives the cancellation that triggered it, so it runs on a // context that keeps this one's values without its deadline. if cleanupErr := cleanupLMSPartials(context.WithoutCancel(ctx), lmstudioModelsDir(), model, before); cleanupErr != nil { - return nil, fmt.Errorf("download stopped, but partial-file cleanup failed: %w", cleanupErr) + // The download did stop, which is what was asked for. Leftover + // bytes are the vendor's to resume, so reporting a failed cancel + // here would deny the one outcome that did happen and leave the + // row on "Canceling" over a transfer that is gone. + slog.Warn("partial-file cleanup failed after cancellation", + "engine", engine, "model", model, "err", cleanupErr) } return nil, context.Canceled } @@ -118,11 +183,7 @@ func runLMSDownloadCommand(ctx context.Context, argv []string, output *lmsDownlo return "", err } defer stdin.Close() - output.stdin = stdin - output.text = "" - output.cancelled = false - output.completed = false - output.answered = false + output.reset(stdin) cmd.Stdout = output cmd.Stderr = output if err := cmd.Start(); err != nil { @@ -144,10 +205,11 @@ func runLMSDownloadCommand(ctx context.Context, argv []string, output *lmsDownlo return stopLMSDownload(lmsCancelGracesFrom(ctx), cmd, finished, output) } } + _, _, text := output.state() if runErr != nil { - return "", fmt.Errorf("%w: %s", runErr, lmsErrorDetail(output.text)) + return "", fmt.Errorf("%w: %s", runErr, lmsErrorDetail(text)) } - return output.text, nil + return text, nil } // lmsCancelGraces bounds each stage of stopping the CLI. No wait here may be @@ -196,8 +258,13 @@ func stopLMSDownload( ) (string, error) { if err := interruptDownload(cmd); err != nil { _ = killDownload(cmd) - if reaped, runErr := waitLMSExit(finished, graces.reap); reaped && runErr == nil && output.completed { - return output.text, nil + reaped, runErr := waitLMSExit(finished, graces.reap) + if !reaped { + output.disown() + return "", fmt.Errorf("could not confirm LM Studio cancellation and could not reclaim its process: %w", err) + } + if _, completed, text := output.state(); runErr == nil && completed { + return text, nil } return "", fmt.Errorf("could not confirm LM Studio cancellation: %w", err) } @@ -206,10 +273,11 @@ func stopLMSDownload( exited, _ = waitLMSExit(finished, graces.retry) } if exited { - if output.completed { - return output.text, nil + cancelled, completed, text := output.state() + if completed { + return text, nil } - if !output.cancelled { + if !cancelled { return "", fmt.Errorf("LM Studio exited without confirming download cancellation") } return "", context.Canceled @@ -219,8 +287,13 @@ func stopLMSDownload( // cancellation the CLI never acknowledged. A process that outlasts even the // kill is abandoned rather than waited on; see lmsCancelGraces. _ = killDownload(cmd) - waitLMSExit(finished, graces.reap) - output.cancelled = false + reaped, _ := waitLMSExit(finished, graces.reap) + output.disown() + if !reaped { + // The writer is still attached to a live process, so its output from + // here belongs to nobody this worker is waiting on. + return "", fmt.Errorf("LM Studio did not confirm download cancellation and its process could not be reclaimed; the download may still be running in LM Studio") + } return "", fmt.Errorf("LM Studio did not confirm download cancellation; the download may still be running in LM Studio") } diff --git a/services/nvpair-engine-manager/pull.go b/services/nvpair-engine-manager/pull.go index b2344d7b..78e89df7 100644 --- a/services/nvpair-engine-manager/pull.go +++ b/services/nvpair-engine-manager/pull.go @@ -24,6 +24,7 @@ import ( "encoding/json" "fmt" "io" + "log/slog" "net/http" "strconv" "strings" @@ -126,8 +127,12 @@ func (e *Executor) pullModelStream(ctx context.Context, engine, model string, pa // on a context that keeps this one's values without its deadline. cleanupCtx := context.WithoutCancel(ctx) if err := cleanupOllamaAfterCancel(cleanupCtx, ollamaBlobsDir(st), digests, before); err != nil { - resultErr = err - return + // The download did stop, which is what was asked for. Leftover + // blobs are Ollama's to resume, so reporting a failed cancel + // here would deny the one outcome that did happen and leave + // the row on "Canceling" over a transfer that is gone. + slog.Warn("partial-blob cleanup failed after cancellation", + "engine", engine, "model", model, "err", err) } } resultErr = context.Canceled diff --git a/services/nvpair-engine-manager/pullcancel_test.go b/services/nvpair-engine-manager/pullcancel_test.go index d3d8fe59..93f0022a 100644 --- a/services/nvpair-engine-manager/pullcancel_test.go +++ b/services/nvpair-engine-manager/pullcancel_test.go @@ -15,6 +15,7 @@ import ( "os" "os/exec" "path/filepath" + "runtime" "strconv" "strings" "testing" @@ -176,11 +177,61 @@ func TestLMSDownloadAnswersCancellationPrompt(t *testing.T) { if stdin.String() != "n\n" { t.Fatalf("answer = %q, want n followed by newline", stdin.String()) } - if !output.cancelled { + if cancelled, _, _ := output.state(); !cancelled { t.Fatal("cancellation acknowledgement was not recorded") } } +// A process the worker could not reclaim keeps writing into this writer, so +// withdrawing its confirmation has to stick. Clearing the flag alone would not: +// the next redraw carrying "Download canceled." would set it again and re-arm +// deletion of files nothing confirmed had stopped. +func TestLMSDownloadDisownedOutputCannotConfirmACancellation(t *testing.T) { + output := &lmsDownloadOutput{lastPercent: -1, progress: func(int) {}} + if _, err := output.Write([]byte("Download canceled.")); err != nil { + t.Fatal(err) + } + if cancelled, _, _ := output.state(); !cancelled { + t.Fatal("acknowledgement was not recorded before the process was disowned") + } + + output.disown() + if cancelled, _, _ := output.state(); cancelled { + t.Fatal("a disowned process still confirmed the cancellation") + } + // The abandoned CLI carries on writing; none of it may re-arm deletion. + if _, err := output.Write([]byte("\rDownload canceled.")); err != nil { + t.Fatal(err) + } + if cancelled, _, _ := output.state(); cancelled { + t.Fatal("output written after disowning re-armed partial-file deletion") + } +} + +// The exec copy goroutine writes this struct while the goroutine handling the +// cancellation reads it to decide whether a partial file may be deleted, and a +// disowned process makes them overlap for real rather than in principle. The +// assertion is the race detector's: this fails under -race if the fields are +// touched without the lock. +func TestLMSDownloadOutputToleratesConcurrentAccess(t *testing.T) { + output := &lmsDownloadOutput{lastPercent: -1, progress: func(int) {}} + writing := make(chan struct{}) + go func() { + defer close(writing) + for i := range 200 { + if _, err := output.Write([]byte(fmt.Sprintf("\r[=== ] %d.00%%", i%100))); err != nil { + t.Errorf("write: %v", err) + return + } + } + }() + for range 200 { + output.state() + } + output.disown() + <-writing +} + func TestLMSDownloadInterruptsChild(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() @@ -255,8 +306,8 @@ func TestLMSDownloadCancellationOutcomes(t *testing.T) { case !strings.Contains(err.Error(), wantErrText): t.Fatalf("cancellation = %v, want an error containing %q", err, wantErrText) } - if output.cancelled != wantCancelled { - t.Errorf("acknowledged = %v, want %v (this is what gates deleting partial files)", output.cancelled, wantCancelled) + if cancelled, _, _ := output.state(); cancelled != wantCancelled { + t.Errorf("acknowledged = %v, want %v (this is what gates deleting partial files)", cancelled, wantCancelled) } }) } @@ -886,6 +937,89 @@ func TestOllamaCancellationHandlesCompletion(t *testing.T) { test("success survives later frame and cancellation", []string{"pulling layer", "success", ""}, "success") } +// blockPartialRemoval makes the unlink of path fail, standing in for the +// sharing violations and permission errors cleanup meets in the field. The +// platforms need different levers: Windows refuses to delete a file that has +// an open handle which did not opt into FILE_SHARE_DELETE, which is the handle +// Go's Open returns, while Unix refuses the unlink when the parent directory +// is not writable. +func blockPartialRemoval(t *testing.T, path string) { + t.Helper() + if runtime.GOOS == "windows" { + file, err := os.Open(path) + if err != nil { + t.Fatalf("hold %s open: %v", path, err) + } + t.Cleanup(func() { _ = file.Close() }) + return + } + if os.Geteuid() == 0 { + t.Skip("root ignores directory permissions, so the unlink cannot be made to fail") + } + dir := filepath.Dir(path) + info, err := os.Stat(dir) + if err != nil { + t.Fatalf("stat %s: %v", dir, err) + } + if err := os.Chmod(dir, 0500); err != nil { + t.Fatalf("make %s read-only: %v", dir, err) + } + t.Cleanup(func() { _ = os.Chmod(dir, info.Mode().Perm()) }) +} + +// Cleanup that cannot finish is not a failed cancellation. The transfer has +// already stopped, which is what was asked for, and the bytes left behind are +// the vendor's to resume. Reporting the cleanup error as the pull's outcome +// denied the one thing that did happen and left the row on "Canceling" over a +// download that was already gone. +func TestOllamaCancellationSucceedsWhenCleanupCannotFinish(t *testing.T) { + blobs := newOllamaBlobsDir(t) + partial := filepath.Join(blobs, blobA+"-partial") + started := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + writePartial(t, partial, "partial") + w.Header().Set("Content-Type", "application/x-ndjson") + if _, err := fmt.Fprintf(w, `{"status":"pulling layer","digest":%q,"total":100,"completed":25}`+"\n", digestA); err != nil { + t.Error(err) + return + } + w.(http.Flusher).Flush() + close(started) + <-r.Context().Done() + })) + defer server.Close() + ex := newOllamaPullExecutor(t, server.URL, nil) + + ctx, cancel := context.WithTimeout(settledContext(), 60*time.Second) + defer cancel() + done := make(chan outcome, 1) + go func() { + var got outcome + got.result, got.err = ex.PullModelStream(ctx, "ollama", "demo", nil) + done <- got + }() + select { + case <-started: + case <-ctx.Done(): + t.Fatal("pull did not start") + } + // The partial exists and is attributable to this pull by now, so this is + // the point its removal can be made to fail. + blockPartialRemoval(t, partial) + + if err := ex.CancelModelPull(ctx, "ollama", "demo"); err != nil { + t.Fatalf("cancel: %v", err) + } + got := <-done + if got.err != nil { + t.Fatalf("pull = %v, want the cancellation to be reported as a cancellation", got.err) + } + assertPullStatus(t, got.result, "cancelled") + // Cleanup could not remove it, and saying so is the warning's job, not the + // cancel's. The file stays for the next attempt to resume. + assertPresent(t, partial) +} + // A predecessor here is another of this engine's downloads that registered // first. trackedPull serializes pulls per engine — the vendor caches share // partial files across models, so one pull's cleanup must not run while another From 6340c6c26db956abc11f0119c3caf1928ca85709 Mon Sep 17 00:00:00 2001 From: Chris Kelsey Date: Tue, 22 Sep 2026 16:18:04 -0700 Subject: [PATCH 07/12] fix(engine-manager): reject a remote pull that names no model The model is what scopes a pull: it keys the claim a cancel is matched against, and it is streamOp's progress filter. Params that named none satisfied the old check purely by being non-empty, then left req.Model empty - which made claimPull a no-op and emptied the filter. An empty filter disables filtering, so the initiator received every other model's pull progress on that engine stamped with its own opID, the one thing the comment above the guard says must not happen. Resolve the model from params first, then require one. That also retires the params fallback in the failure path, which can no longer be reached. Signed-off-by: Chris Kelsey --- .../nvpair-engine-manager/controlstream.go | 20 +++++++------- services/nvpair-engine-manager/pull_test.go | 26 ++++++++++++++----- 2 files changed, 31 insertions(+), 15 deletions(-) diff --git a/services/nvpair-engine-manager/controlstream.go b/services/nvpair-engine-manager/controlstream.go index a7e56dc4..9843a777 100644 --- a/services/nvpair-engine-manager/controlstream.go +++ b/services/nvpair-engine-manager/controlstream.go @@ -112,13 +112,19 @@ func (s *controlServer) handlePull(w http.ResponseWriter, r *http.Request) { http.Error(w, `"engine" is required`, http.StatusBadRequest) return } - if req.Model == "" && len(req.Params) == 0 { - http.Error(w, `"model" or "params" is required`, http.StatusBadRequest) - return - } if model := modelFromParams(req.Params); model != "" { req.Model = model } + // Resolve the model before validating, because the model is what scopes + // everything below. Params that name none used to pass this check and + // leave it empty, which made claimPull a no-op and emptied streamOp's + // modelFilter — and an empty filter disables filtering, so the initiator + // received every other model's pull progress on this engine stamped with + // its own opID. + if req.Model == "" { + http.Error(w, `"model", or "params" naming one, is required`, http.StatusBadRequest) + return + } // The initiator's cancel arrives as its own request, so claim the pull // before starting it: a cancel that beats the download to the executor is // then held for it instead of being read as a cancel for nothing. @@ -126,11 +132,7 @@ func (s *controlServer) handlePull(w http.ResponseWriter, r *http.Request) { s.streamOp(w, r, req.OpID, req.Engine, "pull", req.Model, func(ctx context.Context) (streamFrame, error) { res, err := s.exec.PullModelStream(ctx, req.Engine, req.Model, req.Params) if err != nil { - model := req.Model - if model == "" { - model = modelFromParams(req.Params) - } - msg := s.exec.reportPullFailed(req.Engine, model, err) + msg := s.exec.reportPullFailed(req.Engine, req.Model, err) return streamFrame{}, fmt.Errorf("%s", msg) } return streamFrame{Type: "result", OpID: req.OpID, Engine: req.Engine, Op: "pull", Result: res}, nil diff --git a/services/nvpair-engine-manager/pull_test.go b/services/nvpair-engine-manager/pull_test.go index ec3ce3aa..0a5d4586 100644 --- a/services/nvpair-engine-manager/pull_test.go +++ b/services/nvpair-engine-manager/pull_test.go @@ -32,14 +32,28 @@ func TestPullProgressFromLine(t *testing.T) { } } +// The model is what scopes a pull: it keys the claim a cancel is matched +// against, and it is streamOp's progress filter. A request that does not name +// one has to be refused rather than run unscoped — an empty filter disables +// filtering, so the initiator would receive every other model's pull progress +// on that engine stamped with its own opID. func TestHandlePullRejectsMissingTarget(t *testing.T) { - s := &controlServer{exec: &Executor{progress: newProgressHub()}} - rec := httptest.NewRecorder() - req := httptest.NewRequest("POST", controlPullPath, strings.NewReader(`{"opId":"x","engine":"ollama"}`)) - s.handlePull(rec, req) - if rec.Code != 400 { - t.Fatalf("expected 400 when neither model nor params set, got %d", rec.Code) + test := func(name, body string) { + t.Run(name, func(t *testing.T) { + s := &controlServer{exec: &Executor{progress: newProgressHub()}} + rec := httptest.NewRecorder() + req := httptest.NewRequest("POST", controlPullPath, strings.NewReader(body)) + s.handlePull(rec, req) + if rec.Code != 400 { + t.Fatalf("code = %d, want 400 for a pull that names no model", rec.Code) + } + }) } + test("neither model nor params", `{"opId":"x","engine":"ollama"}`) + // Params carrying no recognisable model name used to satisfy the check + // purely by being non-empty. + test("params naming no model", `{"opId":"x","engine":"ollama","params":{"insecure":true}}`) + test("params with an empty model", `{"opId":"x","engine":"ollama","params":{"name":""}}`) } func TestModelFromParams(t *testing.T) { From 71f780a9ac6c68806e9cc0c9c59ee974b5d4831a Mon Sep 17 00:00:00 2001 From: Chris Kelsey Date: Tue, 22 Sep 2026 17:24:38 -0700 Subject: [PATCH 08/12] fix(engine-manager): stop a download signal reaching a recycled PID Neither cancellation signal addresses the CLI alone. On Unix it goes to the whole download process group; on Windows it becomes a console control event for everything sharing the launcher's console. So a signal sent after the process is reaped does not merely miss - the OS is free to have reissued that PID, and the signal reaches a stranger. The Windows guard did not close the window it was written for. isDownloadLauncher verified the image and closed its process handle on return, and only afterwards did the helper call AttachConsole(pid). Holding a handle is what keeps Windows from reusing a PID, so the check proved only what had been true a moment earlier. Return the handle from openDownloadLauncher instead and close it once the attach has happened. Unix had no guard at all. Both signal sites in stopLMSDownload now go through helpers that decline once the process is reaped, which covers Windows too. The flag is an atomic set before Wait's result is published: cmd.ProcessState cannot serve, because Wait writes it on one goroutine while the cancelling goroutine would read it - trading this race for another. Signed-off-by: Chris Kelsey --- .../nvpair-engine-manager/download_windows.go | 45 ++++++++++++------ .../download_windows_test.go | 23 +++++++-- services/nvpair-engine-manager/lmspull.go | 47 +++++++++++++++---- .../nvpair-engine-manager/pullcancel_test.go | 8 +++- 4 files changed, 95 insertions(+), 28 deletions(-) diff --git a/services/nvpair-engine-manager/download_windows.go b/services/nvpair-engine-manager/download_windows.go index 10992a23..3abd95ae 100644 --- a/services/nvpair-engine-manager/download_windows.go +++ b/services/nvpair-engine-manager/download_windows.go @@ -66,27 +66,38 @@ func runDownloadProcess(argv []string) int { return 0 } -// isDownloadLauncher reports whether pid still names a launcher this worker -// started, by checking that its image is this same executable. Windows reissues -// a PID as soon as the last handle to the exited process closes, so a cancel -// racing the CLI's exit could otherwise attach to a stranger's console — and -// GenerateConsoleCtrlEvent(CTRL_C_EVENT, 0) reaches everything sharing it. -func isDownloadLauncher(pid uint32) bool { +// openDownloadLauncher returns an open handle to pid when it still names a +// launcher this worker started, established by its image being this same +// executable. A cancel racing the CLI's exit could otherwise attach to a +// stranger's console — and GenerateConsoleCtrlEvent(CTRL_C_EVENT, 0) reaches +// everything sharing it. +// +// The handle is the guard, not the image check. Windows reissues a PID as soon +// as the last handle to the exited process closes, so verifying the image and +// then closing the handle proves only what was true a moment ago: the PID can +// be recycled before the caller attaches. Holding this handle keeps the +// process object, and therefore the PID, from being reused — so the caller +// must not close it until it has attached to the console. +func openDownloadLauncher(pid uint32) (windows.Handle, bool) { self, err := os.Executable() if err != nil { - return false + return 0, false } handle, err := windows.OpenProcess(windows.PROCESS_QUERY_LIMITED_INFORMATION, false, pid) if err != nil { - return false + return 0, false } - defer windows.CloseHandle(handle) image := make([]uint16, 32768) size := uint32(len(image)) if err := windows.QueryFullProcessImageName(handle, 0, &image[0], &size); err != nil { - return false + _ = windows.CloseHandle(handle) + return 0, false } - return strings.EqualFold(windows.UTF16ToString(image[:size]), self) + if !strings.EqualFold(windows.UTF16ToString(image[:size]), self) { + _ = windows.CloseHandle(handle) + return 0, false + } + return handle, true } func interruptDownload(cmd *exec.Cmd) error { @@ -114,7 +125,8 @@ func handleDownloadProcess() bool { if err != nil || pid == 0 { os.Exit(1) } - if !isDownloadLauncher(uint32(pid)) { + launcher, ok := openDownloadLauncher(uint32(pid)) + if !ok { fmt.Fprintln(os.Stderr, "download launcher is gone") os.Exit(1) } @@ -124,8 +136,13 @@ func handleDownloadProcess() bool { ignore := kernel.NewProc("SetConsoleCtrlHandler") generate := kernel.NewProc("GenerateConsoleCtrlEvent") free.Call() - if ok, _, err := attach.Call(uintptr(pid)); ok == 0 { - fmt.Fprintln(os.Stderr, err) + attached, _, attachErr := attach.Call(uintptr(pid)) + // The handle has held the PID against reuse up to here, which is the point + // the console stops being addressed by number, so it has done its job. + // os.Exit below would skip a defer, so close it explicitly. + _ = windows.CloseHandle(launcher) + if attached == 0 { + fmt.Fprintln(os.Stderr, attachErr) os.Exit(1) } if ok, _, err := ignore.Call(0, 1); ok == 0 { diff --git a/services/nvpair-engine-manager/download_windows_test.go b/services/nvpair-engine-manager/download_windows_test.go index 1c831e30..ea8e4ec7 100644 --- a/services/nvpair-engine-manager/download_windows_test.go +++ b/services/nvpair-engine-manager/download_windows_test.go @@ -18,16 +18,31 @@ import ( // A cancel racing the CLI's exit can reach interruptDownload after the PID has // been reissued, and the helper's GenerateConsoleCtrlEvent(CTRL_C_EVENT, 0) // reaches every process sharing that console. The image check is what keeps the -// signal inside the app. -func TestIsDownloadLauncherRejectsForeignProcesses(t *testing.T) { - if !isDownloadLauncher(uint32(os.Getpid())) { +// signal inside the app, and the returned handle is what keeps the PID from +// being recycled between that check and the attach. +func TestOpenDownloadLauncherRejectsForeignProcesses(t *testing.T) { + handle, ok := openDownloadLauncher(uint32(os.Getpid())) + if !ok { t.Fatal("this test binary was not recognized as its own image") } + if handle == 0 { + t.Error("accepted launcher came back without a handle, so nothing holds its PID") + } + if err := windows.CloseHandle(handle); err != nil { + t.Errorf("close launcher handle: %v", err) + } // PID 4 is the Windows System process; a PID no process holds is also a // reissue candidate. for _, pid := range []uint32{4, 0xFFFFFFF0} { - if isDownloadLauncher(pid) { + handle, ok := openDownloadLauncher(pid) + if ok { t.Errorf("pid %d was accepted as a download launcher", pid) + _ = windows.CloseHandle(handle) + continue + } + // A rejection must not leak the handle it opened to look. + if handle != 0 { + t.Errorf("pid %d was rejected but returned handle %v", pid, handle) } } } diff --git a/services/nvpair-engine-manager/lmspull.go b/services/nvpair-engine-manager/lmspull.go index 3f1ec9af..94b33410 100644 --- a/services/nvpair-engine-manager/lmspull.go +++ b/services/nvpair-engine-manager/lmspull.go @@ -16,6 +16,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "time" ) @@ -190,7 +191,16 @@ func runLMSDownloadCommand(ctx context.Context, argv []string, output *lmsDownlo return "", err } finished := make(chan error, 1) - go func() { finished <- cmd.Wait() }() + // reaped is set before the result is published, so a reader that sees it + // knows the PID is already free for the OS to reissue. cmd.ProcessState + // cannot serve here: Wait writes it on this goroutine while the + // cancelling one would be reading it. + var reaped atomic.Bool + go func() { + err := cmd.Wait() + reaped.Store(true) + finished <- err + }() var runErr error select { case runErr = <-finished: @@ -202,7 +212,7 @@ func runLMSDownloadCommand(ctx context.Context, argv []string, output *lmsDownlo select { case runErr = <-finished: default: - return stopLMSDownload(lmsCancelGracesFrom(ctx), cmd, finished, output) + return stopLMSDownload(lmsCancelGracesFrom(ctx), cmd, finished, &reaped, output) } } _, _, text := output.state() @@ -254,12 +264,31 @@ func stopLMSDownload( graces lmsCancelGraces, cmd *exec.Cmd, finished <-chan error, + reaped *atomic.Bool, output *lmsDownloadOutput, ) (string, error) { - if err := interruptDownload(cmd); err != nil { + // Neither signal addresses the CLI alone: on Unix it goes to the whole + // download process group, and on Windows it becomes a control event for + // everything sharing the launcher's console. A reaped PID is free for the + // OS to reissue, so one sent afterwards does not merely miss — it can + // reach a stranger. runLMSDownloadCommand's select prefers a finished run + // for this reason, but Wait can also complete while this is running. + interrupt := func() error { + if reaped.Load() { + return os.ErrProcessDone + } + return interruptDownload(cmd) + } + kill := func() { + if reaped.Load() { + return + } _ = killDownload(cmd) - reaped, runErr := waitLMSExit(finished, graces.reap) - if !reaped { + } + if err := interrupt(); err != nil { + kill() + wasReaped, runErr := waitLMSExit(finished, graces.reap) + if !wasReaped { output.disown() return "", fmt.Errorf("could not confirm LM Studio cancellation and could not reclaim its process: %w", err) } @@ -269,7 +298,7 @@ func stopLMSDownload( return "", fmt.Errorf("could not confirm LM Studio cancellation: %w", err) } exited, _ := waitLMSExit(finished, graces.interrupt) - if !exited && interruptDownload(cmd) == nil { + if !exited && interrupt() == nil { exited, _ = waitLMSExit(finished, graces.retry) } if exited { @@ -286,10 +315,10 @@ func stopLMSDownload( // cancel as failed so no partial file is deleted on the strength of a // cancellation the CLI never acknowledged. A process that outlasts even the // kill is abandoned rather than waited on; see lmsCancelGraces. - _ = killDownload(cmd) - reaped, _ := waitLMSExit(finished, graces.reap) + kill() + wasReaped, _ := waitLMSExit(finished, graces.reap) output.disown() - if !reaped { + if !wasReaped { // The writer is still attached to a live process, so its output from // here belongs to nobody this worker is waiting on. return "", fmt.Errorf("LM Studio did not confirm download cancellation and its process could not be reclaimed; the download may still be running in LM Studio") diff --git a/services/nvpair-engine-manager/pullcancel_test.go b/services/nvpair-engine-manager/pullcancel_test.go index 93f0022a..a1ba024f 100644 --- a/services/nvpair-engine-manager/pullcancel_test.go +++ b/services/nvpair-engine-manager/pullcancel_test.go @@ -18,6 +18,7 @@ import ( "runtime" "strconv" "strings" + "sync/atomic" "testing" "time" ) @@ -340,6 +341,11 @@ func TestStopLMSDownloadWhenTheInterruptCannotBeDelivered(t *testing.T) { } finished := make(chan error, 1) finished <- cmd.Wait() + // The process is reaped, which is exactly the state that must not + // be signalled: its PID is free for the OS to reissue, and both + // signals reach more than the CLI. + var reaped atomic.Bool + reaped.Store(true) output := &lmsDownloadOutput{ lastPercent: -1, @@ -347,7 +353,7 @@ func TestStopLMSDownloadWhenTheInterruptCannotBeDelivered(t *testing.T) { completed: completed, text: "transcript", } - text, err := stopLMSDownload(testLMSCancelGraces, cmd, finished, output) + text, err := stopLMSDownload(testLMSCancelGraces, cmd, finished, &reaped, output) if text != wantText { t.Errorf("text = %q, want %q", text, wantText) } From 78378cd1daa8b33bfd529059c4e1df2209a6816a Mon Sep 17 00:00:00 2001 From: Chris Kelsey Date: Tue, 22 Sep 2026 17:33:36 -0700 Subject: [PATCH 09/12] fix(tui): reconcile the active-download list with the progress feed The list was wrong in both directions. A cancel that hit callTimeout was read as a finished cancel. The engine-manager acknowledges one only after the transfer has stopped and its partial files are cleaned up, so the reply can be slower than the call, and retiring the entry there dropped the only cancel target for a download that was still running - a second press reported no active download on the engine while the transfer carried on. Treat the deadline as the call giving up rather than an outcome, the way the pull path already does, and leave the entry cancellable. Nothing then retired an entry whose own request had outlived that same ignored deadline, so a finished download stayed selectable for the rest of the session. Retire on the terminal progress frame, which the engine-manager emits on failure precisely so a client that stopped waiting still converges. The two halves interlock: ignoring the deadline keeps a download cancellable while it runs, and the terminal frame is what ends it. A successful pull that outlasts callTimeout on an engine whose stream emits no terminal frame still leaks its entry. The residue is a stale cancel target, which the engine-manager already answers harmlessly, and closing it means emitting a terminal success frame - a backend contract change rather than a TUI one. Signed-off-by: Chris Kelsey --- services/nvpair-tui/ui/engines.go | 28 ++++++- services/nvpair-tui/ui/engines_test.go | 101 +++++++++++++++++++++++++ 2 files changed, 126 insertions(+), 3 deletions(-) diff --git a/services/nvpair-tui/ui/engines.go b/services/nvpair-tui/ui/engines.go index 9f0ffcda..0dcc19ad 100644 --- a/services/nvpair-tui/ui/engines.go +++ b/services/nvpair-tui/ui/engines.go @@ -189,15 +189,23 @@ func (v *enginesView) Update(msg tea.Msg) tea.Cmd { // 100%; error uses -1), so render them as outcomes rather than a // misleading "success (0%)". A late failure that arrives after the // synchronous call timed out still surfaces here. + // + // They also retire the entry, which is the only thing that does so + // for a download whose own request outlived callTimeout. That + // deadline is ignored deliberately, so without this the entry + // would stay in active for the rest of the session and keep + // offering a finished download as the cancel target. switch p.Stage { case "success": v.status = fmt.Sprintf("pull %s: done", what) + v.retire(enginePull{engine: p.Engine, model: p.Model}) case "error": detail := p.Message if detail == "" { detail = "failed" } v.status = fmt.Sprintf("pull %s failed: %s", what, detail) + v.retire(enginePull{engine: p.Engine, model: p.Model}) case "queued": // The engine already has a download running; this one starts // when that finishes. There is no percent to report yet. @@ -314,12 +322,26 @@ func (v *enginesView) cancelPull() tea.Cmd { } v.status = fmt.Sprintf("canceling %s %s...", pull.engine, pull.model) params := map[string]string{"engine": pull.engine, "model": pull.model} - return call(v.client, "engine:cancel-pull", params, func(_ *rpc.Message, err error) tea.Msg { - if err != nil && !errors.Is(err, context.DeadlineExceeded) { + return call(v.client, "engine:cancel-pull", params, v.decodeCancel(pull)) +} + +// decodeCancel maps a cancel request's outcome onto the update loop. +func (v *enginesView) decodeCancel(pull enginePull) func(*rpc.Message, error) tea.Msg { + return func(_ *rpc.Message, err error) tea.Msg { + // A deadline here is the call timeout, not the cancel's outcome, the + // same way it is for submitPull. Retiring the entry on it would drop + // the only cancel target for a download that is still running, so a + // second press would report nothing active on the engine while the + // transfer carried on. Leave it and let the pull's own terminal + // progress frame retire it. + if errors.Is(err, context.DeadlineExceeded) { + return nil + } + if err != nil { return engineOpMsg{what: "cancel " + pull.model, engine: pull.engine, err: err} } return enginePullDoneMsg{pull: pull} - }) + } } func (v *enginesView) newestPull(engine string) (enginePull, bool) { diff --git a/services/nvpair-tui/ui/engines_test.go b/services/nvpair-tui/ui/engines_test.go index 4edf4b94..7aea2c47 100644 --- a/services/nvpair-tui/ui/engines_test.go +++ b/services/nvpair-tui/ui/engines_test.go @@ -4,6 +4,7 @@ package ui import ( + "context" "encoding/json" "errors" "strings" @@ -64,6 +65,106 @@ func TestPullProgressRendersModelAndQueuedStage(t *testing.T) { "pull lmstudio owner/model: done") } +// A pull's own request routinely outlives callTimeout, and that deadline is +// ignored on purpose because the download carries on. The terminal progress +// frame is then the only thing that can retire the entry — without it a +// finished download stays in active for the rest of the session and keeps +// being offered as the cancel target. +func TestTerminalPullProgressRetiresTheDownload(t *testing.T) { + test := func(name, params string) { + t.Run(name, func(t *testing.T) { + v := newEnginesView(nil) + v.active = []enginePull{ + {engine: "ollama", model: "llama3.2"}, + {engine: "ollama", model: "qwen3:8b"}, + } + v.Update(NotificationMsg{Msg: &rpc.Message{ + Method: "engine:pull-progress", + Params: json.RawMessage(params), + }}) + pull, ok := v.newestPull("ollama") + if !ok { + t.Fatal("the other download was retired too") + } + if pull.model != "llama3.2" { + t.Fatalf("newestPull = %q, want llama3.2 once qwen3:8b has finished", pull.model) + } + }) + } + test("success", `{"engine":"ollama","model":"qwen3:8b","op":"pull","stage":"success"}`) + test("error", `{"engine":"ollama","model":"qwen3:8b","op":"pull","stage":"error","percent":-1,"message":"no space left on device"}`) +} + +// The engine-manager acknowledges a cancel only after the transfer has stopped +// and its partial files are cleaned up, so the reply can outlast callTimeout. +// Retiring the entry on that deadline dropped the only cancel target for a +// download that was still running: a second press reported no active download +// on the engine while the transfer carried on. +func TestCancelKeepsTheDownloadWhenTheCallTimesOut(t *testing.T) { + v := newEnginesView(nil) + pull := enginePull{engine: "ollama", model: "llama3.2"} + v.active = []enginePull{pull} + + if msg := v.decodeCancel(pull)(nil, context.DeadlineExceeded); msg != nil { + t.Fatalf("a timed-out cancel produced %#v, want no message", msg) + } + if _, ok := v.newestPull("ollama"); !ok { + t.Fatal("the download is no longer cancelable after its cancel timed out") + } +} + +// A cancel that actually failed is a different matter: it reports, and leaves +// the entry so it can be tried again. +func TestCancelReportsARealFailure(t *testing.T) { + v := newEnginesView(nil) + pull := enginePull{engine: "ollama", model: "llama3.2"} + v.active = []enginePull{pull} + + msg := v.decodeCancel(pull)(nil, errors.New("engine ollama is not running")) + op, ok := msg.(engineOpMsg) + if !ok { + t.Fatalf("msg = %#v, want engineOpMsg", msg) + } + if !strings.Contains(op.what, pull.model) { + t.Errorf("what = %q, want it to name %q", op.what, pull.model) + } + if _, ok := v.newestPull("ollama"); !ok { + t.Error("a failed cancel left the download unselectable, so it cannot be retried") + } +} + +// A cancel the engine-manager confirmed is the one case that retires the entry +// on the reply itself. +func TestConfirmedCancelRetiresTheDownload(t *testing.T) { + v := newEnginesView(nil) + pull := enginePull{engine: "ollama", model: "llama3.2"} + v.active = []enginePull{pull} + + msg := v.decodeCancel(pull)(nil, nil) + if _, ok := msg.(enginePullDoneMsg); !ok { + t.Fatalf("msg = %#v, want enginePullDoneMsg", msg) + } + v.Update(msg) + if _, ok := v.newestPull("ollama"); ok { + t.Error("a confirmed cancel left the download offered as a cancel target") + } +} + +// A non-terminal frame says the download is still going, so it must leave the +// entry alone. +func TestProgressFrameDoesNotRetireARunningDownload(t *testing.T) { + v := newEnginesView(nil) + running := enginePull{engine: "ollama", model: "llama3.2"} + v.active = []enginePull{running} + v.Update(NotificationMsg{Msg: &rpc.Message{ + Method: "engine:pull-progress", + Params: json.RawMessage(`{"engine":"ollama","model":"llama3.2","op":"pull","stage":"downloading","percent":42}`), + }}) + if _, ok := v.newestPull("ollama"); !ok { + t.Fatal("a running download was retired on a progress frame") + } +} + // The cancel key needs the model, not just the engine: several downloads can be // registered against one engine at a time. func TestCancelTargetsNewestPullOnEngine(t *testing.T) { From 654dce681a0efde3ff3b61ad4a0f0195f5a5b4e0 Mon Sep 17 00:00:00 2001 From: Chris Kelsey Date: Tue, 22 Sep 2026 21:31:09 -0700 Subject: [PATCH 10/12] fix(desktop): let a cancel that outlived its budget be asked again MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cancelling a download was one-shot. The first attempt latched an entry in `cancelingPulls`, and `setModelPullCanceling` refused any later cancel whose key was already there, so once a request stopped being awaited the row sat on "Canceling" with nothing behind it. That entry is reachable. `MODULAR_CANCEL_PULL_TIMEOUT_MS` is 90s, but a peer answering `engine:remote-cancel-pull` is served by the readiness pool, whose header budget is eleven minutes — deliberately, because cutting a peer off mid-cancel is worse than waiting. A peer that takes longer than 90s to stop the transfer and settle its partial files is therefore ordinary, not exceptional. The desktop stops awaiting, keeps the row in "Canceling" (correct: the cancel really is still running), and the only other thing that clears the row is the pull's own RPC settling, which is bounded at six hours. Until then the download keeps running and the user has no way to ask again. The fix separates the two jobs that one map was doing. `cancelingPulls` keeps only the job it is named for: remembering the status the download had before the first cancel, so a rejection can restore it. It now records that on the first cancel alone. Recording it again on a retry would capture "canceling" as the status to restore, and a later rejection would put the row back into the very state it was being told the backend refused to enter. The guard against concurrent cancels moves to the supervisor, which is the layer that knows when a request is actually outstanding. It is held across the await and released in a `finally`, so a mashed button still sends one request while a timed-out one leaves nothing latched behind. Both cancel buttons drop `canceling` from `disabled`. Without that the retry is unreachable and the rest of this change is invisible: the row is the user's only handle on the download. The label still reads "Canceling…" because that remains true, and repeat clicks are absorbed by the guard above rather than by greying out the control. The comment on `MODULAR_CANCEL_PULL_TIMEOUT_MS` claimed the 90s sits above the peer's header budget. That stopped being true when the remote cancel moved onto the readiness pool earlier in this branch; it now describes why the desktop deliberately gives up first and relies on retry instead of matching a budget measured in minutes. Each part was verified by reverting it alone: restoring the old refusal fails the re-issue case, making the status memory unconditional restores a rejected row to "canceling", and moving the release out of `finally` leaks the guard into every later cancel. Signed-off-by: Chris Kelsey --- .../electron/service-bridge/modular-state.ts | 26 +++++++++++-- .../service-bridge/modular-supervisor.ts | 18 +++++++++ .../src/shared/constants/modular-runtime.ts | 17 +++++--- .../ModelManager/IncomingSyncPullRow.tsx | 3 +- .../ModelManager/TransientModelStatusRow.tsx | 11 +++--- .../modular/cancel-model-download.test.ts | 28 ++++++++++++- .../tests/modular/cancel-pull-timeout.test.ts | 39 ++++++++++++++++++- .../modular/model-download-controls.test.ts | 6 ++- 8 files changed, 128 insertions(+), 20 deletions(-) diff --git a/desktop/src/electron/service-bridge/modular-state.ts b/desktop/src/electron/service-bridge/modular-state.ts index 02de0825..3fce0b03 100644 --- a/desktop/src/electron/service-bridge/modular-state.ts +++ b/desktop/src/electron/service-bridge/modular-state.ts @@ -2211,6 +2211,16 @@ class ModularBridgeState { emitBridgePush('engines:progress-cleared', { key }) } + /** + * Move a tracked download into or out of "Canceling", remembering the + * status it had so a rejected cancel can restore it. + * + * Returns false only when there is no such download to mark. It is not the + * guard against concurrent cancels — the supervisor holds that for as long + * as it awaits a reply, because a cancel that outlives its budget is still + * running and the row must stay in "Canceling" while the user remains able + * to ask again. + */ setModelPullCanceling( engineType: EngineType, model: string, @@ -2220,10 +2230,18 @@ class ModularBridgeState { if (!nodeId) return false const key = engineProgressKey({ nodeId, engineType, operation: 'pull', model }) const existing = this.activePulls.get(key) - if (!existing || (canceling && this.cancelingPulls.has(key))) return false - const status = canceling ? 'canceling' : (this.cancelingPulls.get(key) ?? existing.status) - if (canceling) this.cancelingPulls.set(key, existing.status) - else this.cancelingPulls.delete(key) + if (!existing) return false + const remembered = this.cancelingPulls.get(key) + if (canceling) { + // Record the pre-cancel status on the first cancel only. A retry + // would otherwise record "canceling" as the status to restore, and + // a later rejection would put the row back into the very state it + // is trying to leave. + if (remembered === undefined) this.cancelingPulls.set(key, existing.status) + } else { + this.cancelingPulls.delete(key) + } + const status = canceling ? 'canceling' : (remembered ?? existing.status) const progress = { ...existing, status } this.activePulls.set(key, progress) emitBridgePush('engines:progress-changed', progress) diff --git a/desktop/src/electron/service-bridge/modular-supervisor.ts b/desktop/src/electron/service-bridge/modular-supervisor.ts index 0b372555..1a79dde0 100644 --- a/desktop/src/electron/service-bridge/modular-supervisor.ts +++ b/desktop/src/electron/service-bridge/modular-supervisor.ts @@ -344,6 +344,12 @@ class ModularSupervisor { private modelRefreshGenerations = new Map() private discoveryModelRetryTimers = new Map>() private discoveryModelRetryAttempts = new Map() + + // Download cancellations this is currently awaiting a broker reply for, + // keyed by node, engine and model. Held only for the duration of the await + // so a cancel that outlives its budget can be re-issued; see + // {@link ModularSupervisor.cancelModelPull}. + private cancelsInFlight = new Set() // Authoritative value is the persisted ui-config `modularLogLevel`; the // connector seeds it via setLogLevel() before start() so spawn args use it. // This default only applies if start() runs before the connector seeds. @@ -1736,6 +1742,13 @@ class ModularSupervisor { * the row back to "downloading" seconds before it finishes would be a lie * the user acts on. Only a real rejection restores the previous status, and * either way the pull's own settling event clears the entry. + * + * It does not make the cancel one-shot, though. A request that outlives + * its budget has stopped being awaited, and the row it left in "Canceling" + * is the user's only handle on a download that may still be running, so + * asking again has to reach the backend. Only a cancel this is currently + * waiting on is refused, which is what keeps a mashed button from fanning + * out duplicate requests. */ async cancelModelPull( engine: string, @@ -1748,7 +1761,10 @@ class ModularSupervisor { ? state.isRemoteModelPullActive(nodeId, engineType, model) : state.isModelPullActive(engineType, model) if (!active) return + const inFlight = `${nodeId ?? 'local'}:${engineType}:${model}` + if (this.cancelsInFlight.has(inFlight)) return if (!state.setModelPullCanceling(engineType, model, true, nodeId)) return + this.cancelsInFlight.add(inFlight) try { if (nodeId) { await this.callProcess( @@ -1784,6 +1800,8 @@ class ModularSupervisor { `engine-cancel-pull:${nodeId ?? 'local'}:${engine}:${model}`, { engineType, nodeId, operation: 'pull', modelName: model } ) + } finally { + this.cancelsInFlight.delete(inFlight) } } diff --git a/desktop/src/shared/constants/modular-runtime.ts b/desktop/src/shared/constants/modular-runtime.ts index 6d46b63c..a9627b56 100644 --- a/desktop/src/shared/constants/modular-runtime.ts +++ b/desktop/src/shared/constants/modular-runtime.ts @@ -82,11 +82,18 @@ export const MODULAR_MODEL_ACTION_TIMEOUT_MS = 120_000 // Upper bound for `engine:cancel-pull` / `engine:remote-cancel-pull`. The // backend answers only once the transfer has stopped and its partial files are // settled: LM Studio gets two interrupt windows before its CLI is reclaimed, -// and Ollama retries blob removal while its writers release. The remote variant -// adds the peer's own `remoteResponseHeaderTimeout` -// (services/nvpair-engine-manager/remoteclient.go), which starts before the -// peer writes a header — so this has to sit above that budget, not equal it, or -// the desktop reports a failure while the cancel is still running. +// and Ollama retries blob removal while its writers release. This covers that +// locally with room to spare. +// +// It deliberately does not cover the remote variant. A peer answering a cancel +// is one of the endpoints served by the readiness pool, whose header budget is +// far longer (`remoteReadyResponseHeaderTimeout` in +// services/nvpair-engine-manager/remoteclient.go), because cutting the peer off +// mid-cancel is worse than waiting. Matching that here would leave the row +// disabled for minutes before the desktop said anything. Instead a timeout is +// treated as "still running": the row stays in Canceling, the pull's own +// settling clears it, and the cancel stays retryable in the meantime. See +// ModularSupervisor.cancelModelPull. export const MODULAR_CANCEL_PULL_TIMEOUT_MS = 90_000 // Poll interval for `cluster:invite-status` while a pairing handshake is open. diff --git a/desktop/src/ui/components/ModelManager/IncomingSyncPullRow.tsx b/desktop/src/ui/components/ModelManager/IncomingSyncPullRow.tsx index 52cea3eb..59f9cb95 100644 --- a/desktop/src/ui/components/ModelManager/IncomingSyncPullRow.tsx +++ b/desktop/src/ui/components/ModelManager/IncomingSyncPullRow.tsx @@ -41,10 +41,11 @@ export function IncomingSyncPullRow({ + {/* Still clickable while canceling — see TransientModelStatusRow. */}