Skip to content

Commit 5ce61ff

Browse files
committed
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 <ckelsey@nvidia.com>
1 parent 6abb9ec commit 5ce61ff

45 files changed

Lines changed: 3752 additions & 172 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎desktop/docs/services-api.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,7 @@
9797
| `errors:clear` | notification (we consume) | ✅ yes |
9898
| `errors:report` | notification (we consume) | ✅ yes |
9999
| `engine:action` | request (we call) | ✅ yes |
100+
| `engine:cancel-pull` | request (we call) | ✅ yes |
100101
| `engine:configure-launch` | request (we call) | ⚠️ not called |
101102
| `engine:configured-ports` | request (we call) | ⚠️ not called |
102103
| `engine:describe` | request (we call) | ⚠️ not called |
@@ -109,6 +110,7 @@
109110
| `engine:prepare-shutdown` | request (we call) | ✅ yes |
110111
| `engine:preview-launch` | request (we call) | ⚠️ not called |
111112
| `engine:remote-apply-settings` | request (we call) | ⚠️ not called |
113+
| `engine:remote-cancel-pull` | request (we call) | ✅ yes |
112114
| `engine:remote-delete-model` | request (we call) | ✅ yes |
113115
| `engine:remote-get-installed` | request (we call) | ✅ yes |
114116
| `engine:remote-get-settings` | request (we call) | ⚠️ not called |

‎desktop/docs/services-parity.md‎

Lines changed: 107 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -267,15 +267,116 @@ the engine simply shows as stopped and then running.
267267

268268
A local `pull_model` streams live download progress: the engine-manager routes
269269
`engine:action{pull_model}` through its streaming pull path and emits
270-
`engine:pull-progress` (`{ engine, op, stage, percent, message }`) — the local
270+
`engine:pull-progress` (`{ engine, model, op, stage, percent, message }`) — the local
271271
counterpart of `engine:remote-progress`. Personal AI Router consumes it in
272272
`applyLocalEngineProgress` (`modular-supervisor.ts` → `modular-state.ts`),
273-
backfilling the dispatched model (the frame carries none) and advancing the
274-
optimistic pull entry's percent in place, so a local pull shows "Pulling · N%"
273+
keying the frame to its optimistic pull entry by model and advancing that
274+
entry's percent in place, so a local pull shows "Pulling · N%"
275275
to completion just like a remote pull. The awaited action response owns
276-
completion (clearing the entry and refreshing the model list); a CLI-driven pull
277-
(LM Studio) emits a single `pulling` marker and degrades to the indeterminate
278-
spinner.
276+
completion (clearing the entry and refreshing the model list). LM Studio's CLI
277+
progress updates the same percentage display. A pull for an engine that already
278+
has one running is reported as `stage:"queued"` until its turn.
279+
280+
The pull request itself stays pending for the whole download, so it is the
281+
request's own response that reports the outcome — a completion or a
282+
cancellation — while progress arrives out of band on the notification:
283+
284+
```mermaid
285+
sequenceDiagram
286+
participant UI as Model Manager
287+
participant Bridge as Electron bridge
288+
participant Manager as Engine manager
289+
participant Engine as Ollama or LM Studio
290+
291+
UI->>Bridge: Download model
292+
Bridge->>UI: Show optimistic "Pulling…" row
293+
Bridge->>Manager: Start pull (request remains pending)
294+
Manager->>Manager: Record the partial files already on disk
295+
Manager->>Engine: Begin download
296+
297+
loop While downloading
298+
Engine-->>Manager: Download status
299+
Manager-->>Bridge: Progress {engine, model, stage, percent}
300+
Bridge-->>UI: Update the matching model row
301+
end
302+
303+
alt Download completes
304+
Engine-->>Manager: Success
305+
Manager-->>Bridge: Original pull request completes
306+
else User cancels
307+
UI->>Bridge: Cancel download
308+
Bridge->>UI: Show "Canceling…"
309+
Bridge->>Manager: Cancel {engine, model}
310+
Manager->>Engine: Stop transfer
311+
Manager->>Manager: Remove the partial files this pull created
312+
Manager-->>Bridge: Cancellation completes
313+
Manager-->>Bridge: Original pull settles as canceled
314+
end
315+
316+
Bridge->>Manager: Refresh model list
317+
Manager-->>Bridge: Authoritative models
318+
Bridge->>UI: Replace temporary download state
319+
```
320+
321+
`engine:cancel-pull` and `engine:remote-cancel-pull` cancel the selected model
322+
download. The UI keeps "Canceling…" visible until the pull settles, including
323+
when the cancel RPC's own budget elapses first — the backend is still working,
324+
so the row is not dropped back to "Downloading". Ollama closes the pull request
325+
and removes partial blobs named by its progress digests; LM Studio receives
326+
Ctrl+C and a negative answer to its background-download prompt, and after
327+
acknowledgement the partial files carrying the requested quantization are
328+
removed. Completed model files, other quantizations, and unrelated downloads are
329+
retained, and a cancellation the CLI never confirmed deletes nothing. Only a
330+
cancellation someone requested removes files: a pull interrupted by app
331+
shutdown, a dropped remote connection, or the action timeout leaves its partial
332+
data resumable.
333+
334+
Both engines funnel every client on the machine into one cache, so naming a file
335+
is not the same as owning it — Ollama blobs are content-addressed, and the LM
336+
Studio app downloads into the same repository directory. Each pull therefore
337+
records the partial files already present before it starts, and cancelling it
338+
considers only the ones that appeared afterwards. A file that is still growing
339+
once the transfer has stopped is shared with another client even by that
340+
measure, so it survives as well:
341+
342+
```mermaid
343+
flowchart TD
344+
Request["Model pull requested"] --> Tracked["Track by engine + model"]
345+
Tracked --> Busy{"Another PAIR pull active<br/>for this engine?"}
346+
347+
Busy -- No --> Snapshot["Record the partial files<br/>already on disk"]
348+
Busy -- Yes --> Queued["Show as queued"]
349+
350+
Queued --> QueueCancel{"Canceled while queued?"}
351+
QueueCancel -- Yes --> NeverStarted["Remove operation<br/>without starting download"]
352+
QueueCancel -- No --> Wait["Wait for active pull to settle"]
353+
Wait --> Snapshot
354+
355+
Snapshot --> Active["Start download"]
356+
Active --> Cancel{"Cancellation requested?"}
357+
Cancel -- No --> Complete["Download completes normally"]
358+
Cancel -- Yes --> Engine{"Which engine?"}
359+
360+
Engine -- Ollama --> OllamaStop["Close HTTP pull request"]
361+
OllamaStop --> OllamaCandidates["Select partial blobs named<br/>by this pull's digests"]
362+
363+
Engine -- LM Studio --> LMSStop["Send Ctrl+C"]
364+
LMSStop --> LMSPrompt["Answer no to background download"]
365+
LMSPrompt --> LMSCandidates[".part files carrying the<br/>requested quantization"]
366+
367+
OllamaCandidates --> Owned{"Absent from<br/>the snapshot?"}
368+
LMSCandidates --> Owned
369+
Owned -- No --> Preserve["Preserve completed,<br/>shared, and unrelated files"]
370+
Owned -- Yes --> Quiet{"Held still across<br/>repeated checks?"}
371+
Quiet -- No --> Preserve
372+
Quiet -- Yes --> Delete["Remove the partial file"]
373+
374+
Delete --> Settle["Settle pull and refresh models"]
375+
Preserve --> Settle
376+
Complete --> Settle
377+
NeverStarted --> Settle
378+
Settle --> Next["Allow next queued pull to start"]
379+
```
279380

280381
`engine:models` (and the `em` `GET /v1/models` surface) returns the flat model
281382
union, the per-engine breakdown (`modelsByEngine`), and the per-engine set of

‎desktop/src/electron/service-bridge/empty-handlers.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -250,6 +250,11 @@ function routeEngineManagerCommand(payload: WsInvokeRequest<'engine:command'>):
250250
void supervisor.pullModel(engine, payload.engineType, payload.model)
251251
}
252252
break
253+
case 'cancelModelPull':
254+
if (payload.model) {
255+
void supervisor.cancelModelPull(engine, payload.engineType, payload.model)
256+
}
257+
break
253258
case 'deleteModel':
254259
if (payload.model) {
255260
void supervisor.deleteModel(engine, payload.engineType, payload.model)
@@ -365,6 +370,11 @@ function routeRemoteEngineCommand(payload: WsInvokeRequest<'engine:command'>): v
365370
void supervisor.pullModelRemote(nodeId, engine, payload.engineType, payload.model)
366371
}
367372
break
373+
case 'cancelModelPull':
374+
if (payload.model) {
375+
void supervisor.cancelModelPull(engine, payload.engineType, payload.model, nodeId)
376+
}
377+
break
368378
case 'uninstall':
369379
case 'update':
370380
refuseRemote(

‎desktop/src/electron/service-bridge/json-rpc-subprocess.ts‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,13 @@ interface JsonRpcError {
2222

2323
export class JsonRpcResponseError extends Error {}
2424

25+
/**
26+
* The request's own budget elapsed. The backend has not answered and has not
27+
* failed either — it is still working — so a caller that showed optimistic
28+
* state must decide whether to keep it rather than assume the operation lost.
29+
*/
30+
export class JsonRpcTimeoutError extends Error {}
31+
2532
type JsonRpcId = number | string
2633

2734
interface JsonRpcMessage {
@@ -148,7 +155,7 @@ export class JsonRpcSubprocess extends EventEmitter<JsonRpcSubprocessEvents> {
148155
? undefined
149156
: setTimeout(() => {
150157
this.pending.delete(id)
151-
reject(new Error(`${this.name} ${method} timed out`))
158+
reject(new JsonRpcTimeoutError(`${this.name} ${method} timed out`))
152159
}, timeoutMs)
153160
this.pending.set(id, { resolve, reject, timeout })
154161
this.write(message).catch(err => {

‎desktop/src/electron/service-bridge/modular-state.ts‎

Lines changed: 43 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -948,6 +948,7 @@ class ModularBridgeState {
948948
* replay it after a UI refresh mid-download.
949949
*/
950950
private activePulls = new Map<string, EngineProgress>()
951+
private cancelingPulls = new Map<string, string>()
951952
/**
952953
* Authoritative remote-engine facts keyed by {@link remoteOpKey}
953954
* (`${nodeId}:${engineType}`). Populated by the supervisor from each
@@ -960,27 +961,6 @@ class ModularBridgeState {
960961
* evicted.
961962
*/
962963
private remoteEngineFacts = new Map<string, RemoteEngineFacts>()
963-
/**
964-
* The model of an in-flight remote pull keyed by {@link remoteOpKey}. The
965-
* backend's `engine:remote-progress` frames carry **no** `model` field
966-
* (`nvpair-engine-manager/remote.go` `remoteProgress`), so we stamp the model
967-
* we dispatched with here and backfill it onto every incoming frame — that
968-
* keeps the optimistic entry, the streamed percent updates, and the terminal
969-
* clear all on one {@link engineProgressKey}. Cleared by
970-
* {@link finishRemoteModelPull}.
971-
*/
972-
private remotePullModels = new Map<string, string>()
973-
/**
974-
* The model of an in-flight local pull keyed by `engineType`. The backend's
975-
* `engine:pull-progress` frames (the local counterpart of
976-
* `engine:remote-progress`) carry **no** `model` field
977-
* (`nvpair-engine-manager/executor.go` `emitPullProgress`), so we stamp the
978-
* model we dispatched here and backfill it onto every incoming frame — one
979-
* active local pull per engine, matching the backend's engine-scoped
980-
* progress hub. Set by {@link beginModelPull}, cleared by
981-
* {@link finishModelPull}.
982-
*/
983-
private localPullModels = new Map<EngineType, string>()
984964
/**
985965
* The authoritative set of live inbound invites awaiting the local user's PIN
986966
* entry, keyed by `inviteId`. The cluster-manager pushes `cluster:invite-received`
@@ -1492,9 +1472,6 @@ class ModularBridgeState {
14921472
model,
14931473
status: 'pulling'
14941474
}
1495-
// Remember the model so the model-less `engine:remote-progress` frames can
1496-
// be re-associated with this exact entry (see remotePullModels).
1497-
this.remotePullModels.set(this.remoteOpKey(nodeId, engineType), model)
14981475
this.activePulls.set(engineProgressKey(progress), progress)
14991476
emitBridgePush('engines:progress-changed', progress)
15001477
}
@@ -1515,8 +1492,8 @@ class ModularBridgeState {
15151492
* is the only signal that removes the spinner.
15161493
*/
15171494
finishRemoteModelPull(nodeId: string, engineType: EngineType, model: string): void {
1518-
this.remotePullModels.delete(this.remoteOpKey(nodeId, engineType))
15191495
const key = engineProgressKey({ nodeId, engineType, operation: 'pull', model })
1496+
this.cancelingPulls.delete(key)
15201497
if (!this.activePulls.delete(key)) return
15211498
emitBridgePush('engines:progress-cleared', { key })
15221499
}
@@ -1531,28 +1508,22 @@ class ModularBridgeState {
15311508
const op = stringValue(obj.op)
15321509
const operation = op === 'pull' || op === 'pull_model' ? 'pull' : 'install'
15331510
const stage = stringValue(obj.stage) || 'working'
1534-
// `engine:remote-progress` carries no `model`, so for a pull we backfill
1535-
// the model captured at dispatch — otherwise the frame would emit under a
1536-
// different key than the optimistic entry and never update/clear it.
1537-
const model =
1538-
stringValue(obj.model) ||
1539-
(operation === 'pull'
1540-
? (this.remotePullModels.get(this.remoteOpKey(nodeId, engineType)) ?? '')
1541-
: '')
1511+
// The peer stamps the model on every pull frame it relays
1512+
// (`nvpair-engine-manager/remote.go` `remoteProgressFn`), which is what
1513+
// keys a frame to its optimistic entry. An install carries none, and
1514+
// needs none: its progress is tracked per engine.
1515+
const model = stringValue(obj.model)
15421516

15431517
if (stage === 'done' || stage === 'already-installed' || stage === 'failed') {
15441518
const progressKey =
15451519
operation === 'pull' && model
15461520
? engineProgressKey({ nodeId, engineType, operation: 'pull', model })
15471521
: `${nodeId}:${engineType}:${operation}`
15481522
emitBridgePush('engines:progress-cleared', { key: progressKey })
1549-
if (operation === 'pull') {
1550-
this.remotePullModels.delete(this.remoteOpKey(nodeId, engineType))
1551-
if (model) {
1552-
this.activePulls.delete(
1553-
engineProgressKey({ nodeId, engineType, operation: 'pull', model })
1554-
)
1555-
}
1523+
if (operation === 'pull' && model) {
1524+
this.activePulls.delete(
1525+
engineProgressKey({ nodeId, engineType, operation: 'pull', model })
1526+
)
15561527
}
15571528
if (stage === 'already-installed' || stage === 'failed') {
15581529
this.clearRemoteEngineOp(nodeId, engineType)
@@ -1590,7 +1561,7 @@ class ModularBridgeState {
15901561
nodeId,
15911562
nodeName: node?.name ?? nodeId,
15921563
operation,
1593-
status: stage,
1564+
status: this.cancelingPulls.has(pullKey) ? 'canceling' : stage,
15941565
percent,
15951566
model: model || undefined
15961567
}
@@ -1820,7 +1791,6 @@ class ModularBridgeState {
18201791
private dropRemoteEngineFacts(nodeId: string): void {
18211792
for (const engine of PROXY_ENGINES) {
18221793
this.remoteEngineFacts.delete(this.remoteOpKey(nodeId, engine))
1823-
this.remotePullModels.delete(this.remoteOpKey(nodeId, engine))
18241794
}
18251795
}
18261796

@@ -2180,9 +2150,6 @@ class ModularBridgeState {
21802150
model,
21812151
status: 'pulling'
21822152
}
2183-
// Remember the model so the model-less `engine:pull-progress` frames can
2184-
// be re-associated with this exact entry (see localPullModels).
2185-
this.localPullModels.set(engineType, model)
21862153
this.activePulls.set(engineProgressKey(progress), progress)
21872154
emitBridgePush('engines:progress-changed', progress)
21882155
}
@@ -2191,12 +2158,13 @@ class ModularBridgeState {
21912158
* Project a `nvpair-engine-manager` `engine:pull-progress` for the local node
21922159
* — the local counterpart of `engine:remote-progress`: live
21932160
* download progress for a model pull driven via `engine:action`
2194-
* `pull_model`. The frame carries `{ engine, op, stage, percent, message }`
2195-
* and **no** `model`, so we backfill the model captured at dispatch and
2196-
* refresh the optimistic entry in place. Clearing stays with the awaited pull
2197-
* ({@link finishModelPull}), so this only advances percent/stage: it ignores
2198-
* the terminal error sentinel (`percent: -1`) and the non-download frames'
2199-
* `percent: 0` so the rendered bar never jumps backwards.
2161+
* `pull_model`. The frame carries `{ engine, model, op, stage, percent,
2162+
* message }` and its `model` identifies which optimistic entry to refresh —
2163+
* the backend queues concurrent pulls per engine, so the engine alone does
2164+
* not. Clearing stays with the awaited pull ({@link finishModelPull}), so
2165+
* this only advances percent/stage: it ignores the terminal error sentinel
2166+
* (`percent: -1`) and the non-download frames' `percent: 0` so the rendered
2167+
* bar never jumps backwards.
22002168
*/
22012169
applyLocalEngineProgress(params: JsonValue | undefined): void {
22022170
const nodeId = this.selfId
@@ -2206,7 +2174,7 @@ class ModularBridgeState {
22062174
const engineType = engineTypeFromManagerName(stringValue(obj.engine))
22072175
if (!engineType) return
22082176

2209-
const model = this.localPullModels.get(engineType)
2177+
const model = stringValue(obj.model)
22102178
if (!model) return
22112179
const key = engineProgressKey({ nodeId, engineType, operation: 'pull', model })
22122180
const existing = this.activePulls.get(key)
@@ -2215,7 +2183,9 @@ class ModularBridgeState {
22152183
const percent = numberValue(obj.percent)
22162184
const progress: EngineProgress = {
22172185
...existing,
2218-
status: stringValue(obj.stage) || existing.status,
2186+
status: this.cancelingPulls.has(key)
2187+
? 'canceling'
2188+
: stringValue(obj.stage) || existing.status,
22192189
percent: mergePullProgressPercent(percent, existing.percent)
22202190
}
22212191
this.activePulls.set(key, progress)
@@ -2235,12 +2205,31 @@ class ModularBridgeState {
22352205
finishModelPull(engineType: EngineType, model: string): void {
22362206
const nodeId = this.selfId
22372207
if (!nodeId) return
2238-
this.localPullModels.delete(engineType)
22392208
const key = engineProgressKey({ nodeId, engineType, operation: 'pull', model })
2209+
this.cancelingPulls.delete(key)
22402210
if (!this.activePulls.delete(key)) return
22412211
emitBridgePush('engines:progress-cleared', { key })
22422212
}
22432213

2214+
setModelPullCanceling(
2215+
engineType: EngineType,
2216+
model: string,
2217+
canceling: boolean,
2218+
nodeId = this.selfId
2219+
): boolean {
2220+
if (!nodeId) return false
2221+
const key = engineProgressKey({ nodeId, engineType, operation: 'pull', model })
2222+
const existing = this.activePulls.get(key)
2223+
if (!existing || (canceling && this.cancelingPulls.has(key))) return false
2224+
const status = canceling ? 'canceling' : (this.cancelingPulls.get(key) ?? existing.status)
2225+
if (canceling) this.cancelingPulls.set(key, existing.status)
2226+
else this.cancelingPulls.delete(key)
2227+
const progress = { ...existing, status }
2228+
this.activePulls.set(key, progress)
2229+
emitBridgePush('engines:progress-changed', progress)
2230+
return true
2231+
}
2232+
22442233
getLogs(): LogEntry[] {
22452234
return [...this.logs]
22462235
}

0 commit comments

Comments
 (0)