feat(studio): run chats in parallel in the Chat tab (#7455)
* feat(studio): run chats in parallel in the Chat tab
New Chat used to cancel whatever the current conversation was generating.
It now leaves it running, like switching to the Train or Export tab: the
sidebar shows which chats are still going, and Stop is per conversation.
Plain `unsloth studio` launched llama-server with one decode slot, so the
admission queue serialised every chat regardless of what the UI did. Both
entry points now default to the same slot count as `unsloth studio run`.
A model swap still ends every running chat, since they all decode on one
llama-server. /load and /unload now refuse with 409 and name those chats
unless the caller passes force_cancel_active, and the UI asks first.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix(studio): scope the composer tool badge to its own conversation
The green "Running Python: ..." badge above the composer read a single
global store value, so one chat's tool call showed above every other
chat's composer, including a brand-new empty one. Its elapsed counter
also restarted at 0 on every thread switch, and a run ending anywhere
cleared the badge everywhere.
Key the status by thread and store the moment it started, so each
conversation shows only its own tool call and the counter resumes rather
than restarts. Also adds a test that every conversation gets its own
tool sandbox directory, which parallel tool calls depend on.
* Fix stalled tool calls while awaiting approval for PR #7455
Three problems, all from the approval prompt behaving as though only one
chat could ever run.
Arguments were not streamed for a gated call, so the chat stayed blank for
as long as the model took to write the payload, which for a large file is
minutes. Nothing runs before the decision either way, and the code is what
is being approved, so python and terminal now stream their card while
gated. render_html stays suppressed: its card renders the payload.
The status read "Running ..." with a climbing timer while the call had not
started. It now reports that it is waiting for approval, then switches to
running once allowed.
The admission lease was held across the wait, so four unanswered prompts
held all four decode slots and no other chat could start while llama-server
sat idle. A parked run keeps its lease but no longer counts against
capacity.
Measured with four prompts left open: every gated call streamed its code,
none reported running, and a fresh chat answered in 0.4s where it
previously waited 290s and never did.
* Fix duplicated and truncated tool cards for PR #7455
A gated tool call rendered two cards: the provisional one that streams the
arguments, plus a second one keyed by the approval id. Only the second ever
got its tool_end, so the first spun "Running" for the rest of the chat.
Reuse the open part when the approval prompt arrives.
The terminal card also showed nothing but a 60-char trigger label, so a long
heredoc read as no progress at all. It now renders the command the same way
the Python card renders its script, and neither is capped at 10k chars.
Both cells moved inside the collapsible, so one chevron hides the code with
the output and Copy / Download exist only while the card is open. A card
parked on the prompt says so instead of counting up "Running".
* Fix review findings on the parallel-chat gate for PR #7455
Backend:
- /unload rechecks active generations under the lifecycle gate, like /load,
and lets its 409 through the catch-all instead of rewriting it as a 500.
- /load gates only once _load_model_impl has decided this is a real reload,
so an Apply on the already-loaded model no longer refuses, and the retry it
asks for no longer cancels every chat before returning already_loaded.
- The direct /v1/responses stream registers in the cancel registry, so a
non-forced unload can no longer tear llama-server down under it.
- run_server defaults to the same slot count as the CLI. colab.py calls it
without the argument, so Colab was still serialising every chat.
Frontend:
- Cancelling a backgrounded chat aborts its own request rather than only
posting a cancel id, which is the only thing that ends an external-provider
or audio run.
- The model-swap dialog counts local runs only, and falls back to the backend
when this tab's map is empty, so a reload or a second tab still gets asked.
- Context usage and the diffusion canvas are scoped to the chat that produced
them; a compare row reads activity from its member threads.
Tests:
- The extracted-source cancel harnesses supply the active-generations module,
which the tracked-cancel class now depends on.
* Fix the swap confirmation scope and cancel timing for PR #7455
A forced load cancelled every chat before the model identifier, GPU selection,
training coexistence and download checks had run, so a load that then failed
those checks stopped the chats and replaced nothing. The refusal still happens
early, but the destructive cancel now sits immediately before the teardown it
is paying for, and rechecks under the gate like /unload does.
The swap dialog only reconciled with the backend when this tab looked idle, so
one local chat was enough to hide a second tab's runs. Confirming then sent
force_cancel_active, which cancels every backend run, including the ones the
dialog never mentioned. The backend snapshot is now merged in every time, so
the dialog names what will actually stop. External-provider runs are never
registered there, so the union stays local-only.
Also drops the active-generations docstring claim about restoring sidebar
spinners, which nothing consumes.
* Defer destructive cancels and track every local stream for PR #7455
/unload cancelled the running chats before it had resolved that it unloads
anything. A stale model_path, which a second tab produces routinely, killed
every chat and then no-opped, leaving the resident model up. It now refuses
early and cancels only at each teardown, matching /load.
The swap dialog also stopped every chat locally the moment the user confirmed,
which threw away the two-phase backend behaviour: a load that then failed
identifier resolution, GPU validation or the training guard had already
truncated the replies. The backend now owns the cancel.
Three local streams decoded on llama-server without registering, so a
non-forced unload counted zero generations and tore the server down mid
response: /v1/completions streaming, and the plain and server-tool Anthropic
streams, the first of which is the default /v1/messages path. Note this makes
a non-forced load return 409 during those runs rather than draining quietly,
the same trade the /v1/responses fix made.
The safetensors tool loop still announced a gated call as running while it
waited on a human; only the GGUF loop had been fixed. A source-level parity
test now pins both.
Also drops stopAllChatThreads, which has no callers left.
* Studio: close three load/unload gate races found in review
Re-check the in-flight load guard after the stop-running-chats confirm.
The confirm always GETs active-generations before its zero-running
early-out, so the guard no longer sits atomically ahead of the
reservation and two picks in that window both reached performLoad over
the same refs. ejectModel had the same shape and gets the same re-check.
Reject a sidecar swap immediately before the forced cancel in both load
branches. The previous check was back at the top of preflight, so an
install reserving during identifier resolution, the tier probe, the
training guard or the download check made the post-drain recheck 409 a
load whose chats had already been stopped.
Enter the Anthropic passthrough's cancel tracker inside its body
generator. It was entered eagerly and returned through
_sse_streaming_response, which sets no unstarted_cleanup, so a response
whose body never started left the run registered forever and 409'd every
later non-forced load and unload.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Trim comments across the files this PR touches
Tightens the comments and doc blocks in the backend, CLI, tests and frontend
files changed by this PR: collapses multi-line explanations to a single line
where they still read clearly, and drops the ones the code already says.
No code changes, verified by an AST comparison against the previous commit.
* Studio: defer the destructive cancel and close two gate gaps
Move the forced cancel behind every check that can still reject a swap.
The drain now runs first with the runs it is about to cancel discounted,
so it waits only for inference the cancel cannot end, then the sidecar
check decides, then the cancel fires, then a second drain lets those runs
unwind before teardown. A sidecar install reserving during the drain no
longer 409s a load whose chats have already been stopped.
Track the non-streaming /v1/completions proxy. It was the last local
decode path missing from active_generations, so an unload, which runs no
drain, tore llama-server down under it and force_cancel_active could not
signal it. It now uses the same tracked cancel event and dedicated client
as the OpenAI pass-through.
Skip the client's preliminary unload while chats are generating and let
/load evict at its own post-preflight point instead. Forwarding
force_cancel_active there truncated replies before identifier
resolution, the GPU and training guards and the download check had run.
Keep per-thread context usage so returning to a chat whose background run
finished restores its bar instead of leaving it blank until the next turn.
Make the running-flag clear run-specific. Every run without a resolved
thread id shares the "__default" key, so concurrent compare panes could
clear each other's flag and strand a live stop handle.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: register the embeddings proxy with the swap gate
/v1/embeddings proxied straight through the pooled client with no tracked
cancel event, so it never appeared in active_generations. /unload runs no
idle drain, so a concurrent non-forced unload counted zero generations and
killed llama-server mid-request, and force_cancel_active had no event to
signal. Mirrors the completions proxy: tracked event, dedicated unpooled
client closed by a cancel/disconnect watcher, unregister in a nested
finally so a close failure cannot leave a phantom generation behind.
* Trim comments on the newest changes in this PR
Comments only, no code changes: shorten the ones added by the load-gate
ordering, embeddings and per-thread usage work down to the same density as
the rest of the diff.
* Studio: register the legacy generate stream with the swap gate
/generate/stream built a cancel event but never entered the tracker, so it
was invisible to active_generations. Being in the keep-warm middleware's
inference suffixes only covers /load, which drains; /unload does not, so a
non-forced unload passed the 409 gate and then blocked on the standard
backend's generation lock, and a forced swap had no event to signal.
Registered inside the body generator under a nested finally so a teardown
failure cannot skip the unregister.
The AST contract test asserted the cleanup finally by overwriting its flag
per Try node, so a nested try made the last one win. Accumulate instead,
which is what the existence claim meant.
* Studio: three more swap-gate gaps found in review
Register /audio/generate with the gate. TTS holds the model for the whole
request and /unload runs no drain, so unregistered a non-forced swap counted
zero generations and tore the model down mid-generation; the orchestrator
path only waits 15s for the generation lock, which real TTS exceeds. No
cancel keys: no backend takes a cancel_event for audio, so the event has no
observer and a forced swap still cannot interrupt audio already in flight.
Thread the tracked cancel event into the /v1/responses admission wait. It
was the only admission caller passing None, so a queued run could not be
reached by cancel_all() and a plain /inference/cancel could not stop it at
all. Same omission fixed at the upstream send there and on /v1/completions.
Let an unforced unload of a stale model path reach the no-op check. Before
this PR that request returned 200 and did nothing; the new gate refused it
with 409 for a request that reaches no teardown branch. Gate both refusal
passes on the disjunction of the route's own teardown conditions, including
not is_loaded, so a mid-load GGUF still refuses.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: register the remaining non-streaming decode paths
stream defaults to false on all three of these, so they are the ordinary
shape of their routes, and each holds a local backend for the whole
request. /unload runs no idle drain, so with no registry entry a non-forced
swap counted zero generations and tore the backend down mid-request instead
of returning 409, and a forced one had no event to signal.
Non-streaming /v1/messages: all three helpers ran with an empty registry,
since only the streaming siblings were tracked. Registered at the call site
because the pass-through takes no cancel_event of its own, and with no
cancel keys, matching those siblings.
Non-streaming standard chat and audio-input chat: the trackers in this route
sit inside their `if payload.stream:` arms, so neither else branch was
covered. The GGUF sibling already registers its own non-streaming branch.
Each exit is in a finally on the branch's existing try, so the except arms
are covered too: a leaked entry 409s every later swap until restart.
* Studio: tighten the swap-gate comments
Comment-only pass over the newest swap-gate registrations: collapse the multi-line rationales in /unload, the legacy generate stream, audio generation and the non-streaming chat branches, and the matching test preambles, to the shortest form that still carries the reason. No code changes.
* Studio: stop the reselect dialog promising a stop that never happens
Picking an external provider leaves the local model resident and stops the
status poll mirroring it, so reselecting that model showed the stop-chats
dialog, and /load then answered already_loaded ahead of its cancel hook.
Confirmed with the live backend: the same pick with force_cancel_active set
still returned already_loaded and the chat kept streaming. Not stopping
those chats is right, since the load never interrupts them, so remove the
prompt rather than honour it. Blanket-skipping is unsafe, because the same
id and variant with one sampling setting changed is a real reload and 409s,
so the branch only fires when a status fetch confirms the resident
checkpoint and variant match, and then adopts it without calling /load.
Redact native model paths from the active-generations response. Registering
/generate/stream recorded backend.active_model_name verbatim, which is an
absolute path for a native local model, and this route is the only place
that serialises it. Redacting at the response covers every tracker rather
than the one that surfaced it.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: keep hydrated context usage in the per-thread map
The history loader restores a saved conversation's usage through
setContextUsage only, and it runs once per mount, so switching away and
back left the bar blank for a hydrated chat even after the per-thread map
landed. setContextUsage now writes the value through to the visible
thread's own entry and clears that entry when passed null, which covers
both hydration call sites and any future writer.
* Studio: unblock load cancellation and share unresolved thread keys
Run the two stop-loading fast paths ahead of the unload route's pre-gate
refusal. _unload_may_evict returns True for exactly the model being
cancelled, so the refusal was blocking the branch that cancels a load which
has replaced nothing and can interrupt no chat. The client made that
unrecoverable: cancelLoading sends the unload without force, drops the
result, and its abort never reaches /load, which takes no signal, so the
load ran on and could later cancel those chats and swap the model. Nothing
else is exempted; an unload that would tear down a serving model matches
neither fast path and still 409s. The comment claiming the client lets that
409 surface is corrected, since it discards it.
Hold every owner behind a shared thread key. Runs with no resolved thread id
share "__default" (concurrent compare panes, since startCompare clears
activeThreadId), so a single owner slot let a second run replace the first's
token and then delete the shared entry while it was still generating, and
the server-cancel map lost the older handle the same way. Both now hold a
list, the running and local flags survive until the last owner clears, and
stopChatThread stops every handle under the key.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: carry a confirmed swap into the sidecar install, key restored usage by thread
Picking a model that needs a newer transformers while chats generate raised the
"stop N chats" prompt, but the answer never reached the install that runs before
the load: /install-latest-transformers refused on those same chats and took no
force flag, so Retry hit the same 409 and nothing in the flow stopped them.
Carry force_cancel_active through the consent dialog into the installer. Only
the pre-gate fast path is skipped: the recheck under the lifecycle gate still
has to pass, so an unconfirmed caller is refused as before. The cancel runs last
inside the gate, after every check that can still reject the install, and the
drain behind it is bounded since it holds the gate and the sidecar reservation.
Also key restored context usage by the thread the loader read. history.load()
captures remoteId before two awaited round trips, so a switch inside that window
filed one thread's usage under another and setActiveThreadId kept re-applying it.
Preserve sibling owners when a run key is cleared without an owner: the image
rejection gate now uses its own token, and the reducer leaves owned runs alone.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: bound the post-cancel drains, and make cancellation reach the paths that ignored it
A forced swap cancels the chats it interrupts, then waits for them to unwind.
That wait had no deadline while holding the lifecycle gate, and TTS on the
subprocess backend observes no cancel event at all, so one audio generation
could pin every load, unload and new request for its whole duration. Bound both
post-cancel drains. Pre-cancel drains stay unbounded: the swap can still be
refused there, so shortening them would weaken what they protect.
/unload had the opposite problem and no drain at all, cancelling and tearing
down on the next line, which turned a clean stream end into a dropped
connection. Give it the same bounded wait, gated on the cancel having cancelled
something so an idle Eject pays nothing.
Make the cancel actually land where it can. GGUF TTS now takes a cancel_event
and a watcher closes its client to break the blocking POST. The Anthropic
non-streaming pass-through did the same thing the completions and embeddings
paths used to: register with the gate, then run both POSTs on the pooled client
that cannot be closed. It now uses a per-request client like they do.
Also: park and unpark the admission queue the reservation actually holds, since
queues are keyed by base_url and a reload mints a new port; key tool output by
remoteId on both sides, so the first turn of a New Chat stops writing under one
key and reading another; and give tool status a run owner, so a finishing run
cannot blank the badge a concurrent one is still showing.
Clamp --parallel to 1 on a llama-server without --kv-unified. The new default of
4 would otherwise split -c four ways on such a build, quartering the context
window for a feature it cannot serve.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: scope a chat's Stop to its own generation, and clear the way before a confirmed install
Safetensors generation is serialized on _gen_lock and the worker has a single
cancel event, so a chat still queued on that lock owns no generation. Its Stop
handler called reset_generation_state() anyway, which set the shared event and
ended whichever conversation was actually running. Parallel chats is what makes
that reachable.
_generate_inner now records its cancel_event as the current holder once it takes
the lock, and reset_generation_state drops a reset from anyone else. Every route
call site passes its own request event. A reset with no event stays global, so
unload and model switch cannot leave a generation alive, and a reset while
nothing runs still resets, so an error path before generation is not a no-op.
The other two backends take the argument too, or the standard one raises
TypeError on every cancel.
The sidecar install had the mirror of the /load ordering problem: it cancelled
the chats first and drained second, so an unrelated counted request the cancel
cannot reach (a count_tokens, say) was still there for the recheck, which then
refused an install that had already stopped every chat for nothing. Drain the
unreachable remainder first, discounting the registered chats, then cancel.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: close the windows the previous round's fixes left open
Three follow-ups, two of them holes in the fixes just before them.
The worker claim went in after _send_cmd, so the command was already running
unclaimed and a queued chat's Stop in that window still reset it. Claim first,
with the send inside the same try, so a failed send releases it too.
Tool status kept one entry per key with an owner. That stops a foreign clear but
not an overwrite: under the shared unresolved-thread key the second run replaced
the first's entry, and its own clear then removed the only one while the first
tool was still running. Keep per-run entries and render the newest.
/unload gated its drain on having cancelled something, so a request that passed
the keep-warm middleware but had not reached its tracker yet was invisible to it
and the teardown landed on an already-admitted request. Drain on the middleware
count instead, which covers that window as well as the cancelled runs, then
re-cancel whatever registered while waiting. Bounded, not a refusal: an unload is
deliberate, and on expiry it proceeds exactly as before.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Trim the parallel-chats comments to their reasons
Compress the multi-line rationales added by this branch into shorter forms and drop
restatements of the code below them. The reasons behind the drain bounds, the deferred
cancel, the per-request generation ownership and the thread-scoped tool and usage keys
are kept, just said in fewer lines.
* Studio: own the worker per generation, and make a resumed chat requeue for its slot
Ownership was a single lock holder, so dispatched runs (compare mode bypasses
_gen_lock by design) never claimed it and the guard fell straight through to the
global reset: a Stop on one of them ended its siblings. Track the generations
actually running instead, claimed before the send and released in the same
finally on both paths. A reset still proceeds when nothing is running, so an
error path ahead of generation is not swallowed.
park() hands the freed slot to a waiter, so a chat resuming from a tool approval
could take it back while that waiter was still decoding, putting two holders on
a one-slot server and sending the resumed tool loop past the admission limit.
unpark_async waits for room; the plain unpark stays for a holder tearing down,
which will not decode again.
Audio only observed its cancel event on a forced swap. An explicit Stop just
aborts the fetch, and this route has no cancel id, so llama-server ran on to the
request timeout after the chat reported it stopped. Watch the disconnect.
Also read tool status by remoteId, matching the key the adapter writes and the
fix already made for tool output, and stop an unresolved run from writing its
usage into whichever conversation the user moved to.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: let only the generation the worker is running speak for it, and hold a slot for a resuming chat
The ownership list recorded admission, but the subprocess runs generations one
at a time, so a dispatched request queued behind another counted as an owner and
its Stop signalled the shared cancel event, ending the request that was actually
running. Keep admission for release bookkeeping and gate ownership on execution
instead, promoted when the worker first answers that request. Nothing executing
still permits a reset, so an error path ahead of generation is not swallowed.
The worker has one cancel event and no per-request cancellation, so this decides
who may pull the lever rather than making the lever per-request.
A resuming chat also polled for a slot it could never see: release() grants to
the next waiter under the same lock, so later arrivals overtook an approved chat
indefinitely. A pending unpark now reserves the next slot and they queue behind
it.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: cover the prefill window, and keep a first turn's tool output readable
Gating worker ownership on execution left the interval between the send and the
first response uncovered: nothing is executing then, and the empty case admitted
anyone, so a queued chat's Stop still ended the one in prefill. Split the empty
case. Nothing claimed at all still permits a reset, so an error path ahead of
generation is not swallowed; claimed but unanswered resolves to the oldest
claim, which is what a FIFO command queue is working on.
Putting both sides of the tool-output scope on remoteId left the first turn of a
New Chat writing under the unresolved scope for its whole life while the readers
recomputed the moment the autosave assigned an id, so the card blanked mid-run.
The readers now fall back to the unresolved scope, which only an unpersisted
first turn can occupy.
* Studio: order the parked approvals, and tie a worker claim to its enqueue
The reservation added for admission fairness was a bare count, so every approved
holder counted against every other: park two chats, approve both, and once the
last decoder released, nothing could ever satisfy the check again. That is a
deadlock where the problem it fixed was only unfairness. Make it a FIFO ticket
so a pending unpark blocks the ones behind it and no others.
_owns_worker reads claim order to decide which request the worker is prefilling,
which only holds if claiming and enqueuing cannot interleave. Hold one lock
across both on the dispatched and the locked path.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: follow a first turn's run when its thread gets an id, and key the denoising canvas per chat
A run started before its thread existed filed every handle under "__default". Nothing
moved them once autosave assigned the real id, so the sidebar row showed no spinner and
Stop could not reach the generation, which kept holding a slot.
adoptDefaultThreadRun re-keys the run maps onto the real id from the thread adapter's
initialize(), where the id first exists; anything already filed under that id wins, since
that is a later run. The adapter captures its key once at run start, so it now resolves
the live key per use through runKeyForOwner, looking its own serverCancel up in the owner
map. Without that the migrated entries are stranded and the spinner never clears.
The denoising canvas was one global slot, so two diffusion chats overwrote each other and
the ownership tag then hid the visible preview until that thread emitted again. It is now
activeDiffusionCanvasByThreadId, written and cleared per thread, and the frame no longer
carries a threadId of its own. The bubble reads threadListItem.remoteId, dropping the dead
threadListItem.id arm: the writer tags unstable_threadId, which is exactly remoteId.
Two existing backend tests needed the same treatment. _bare_orchestrator skips __init__, so
it now sets the claim bookkeeping the worker ownership check reads. The Anthropic
passthrough gate test anchored on comment prose that a rewrap had broken; it anchors on the
code instead.
* Studio: hand the worker over cleanly between generations, and stop unresolved runs sharing each other's state
Worker ownership moved off the consumer and onto the dispatcher. Consumers read their
mailbox whenever they get around to it, so a request whose gen_done had been routed still
owned the worker while the next one ran, and a late Stop for it cancelled that one. The
dispatcher is the only place responses arrive in the order the worker produced them: it
now retires a request at its terminal response and promotes the next one, and answering a
request makes it the sole executor, since the subprocess runs one generation at a time.
reserve()'s immediate path ignored the unpark tickets that _grant_waiters_locked already
honours, so a request arriving between a slot freeing and an approved chat's next poll
took it, repeatedly. It applies the same reservation now.
Three places let concurrent first turns share state through the "__default" key. Nothing
links a run filed there to the id its thread later receives, so rather than guess, each
now declines when the key is ambiguous: adoption only re-keys a lone run, the composer
badge only claims a lone status, and the tool-output fallback only applies to a thread
that is still running. That leaves two concurrent first turns where they were before
adoption existed instead of handing one thread the other's handles.
A first turn's usage was never filed, because its key stayed null for the whole run while
autosave moved activeThreadId to the real id, so the context bar went blank after the
first reply. It resolves the adopted key like the cleanup handles do.
Cancelling a forced load left the UI with no model: the previous one stays resident until
/load's teardown, and the cancel path cleared the checkpoint without rolling back. It now
resyncs from the backend, which is right whether or not the load got that far.
The sidecar install drain is weighted 1:4 rather than halved, total unchanged. Only the
second half benefits from patience, and cutting it short refused installs whose chats had
already been stopped for nothing.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: give a first turn its real thread id before the run starts
A first turn filed every run handle under a shared unresolved key because
assistant-ui binds unstable_threadId before the thread is persisted. Two of them
overlapping there is unresolvable afterwards, and the last round's migration could
only decline rather than guess, which left neither sidebar row showing its run.
The id is available earlier than I claimed. append() already tracks
threadListItem.initialize() by the user message id, and createPersistedRunAdapter
already awaits that promise before invoking the adapter, so the thread is persisted
by the time the run begins. It was only being discarded: the tracked promise resolved
to void. It now resolves to the assigned id, and the wrapper hands it to the adapter
when assistant-ui had none. An id that is already set is never replaced, since that
would move a running chat's handles out from under the row watching them. The
existing unresolved-key guards stay as a safety net but should no longer carry weight.
The sidebar counted running thread ids rather than rows, so one compare conversation
read as two chats. It folds ids into rows through the same threadIds the row spinner
uses, and still counts a running id that matches no row.
_TrackedCancel always registered kind="chat", so an embeddings or raw completions
request appeared in the model-swap prompt as an unnamed conversation and confirming
cancelled it while calling it a chat. The non-conversation routes now pass their own
kind, and the prompt says "requests" whenever the snapshot is not all chats.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: withhold the shared worker cancel from a request the worker has left
Moving ownership to the dispatcher fixed reset_generation_state, but the token loop
signals the shared worker event directly and did not carry the same rule. A dispatched
consumer runs with mark_started off and can still be draining tokens buffered before
its gen_done was routed, so stopping it there ended whichever request the worker had
started next.
It now signals only when _owns_worker agrees, the same predicate reset_generation_state
uses. The local drain and return are unconditional, since those touch nothing but this
stream. The remaining _cancel_generation callers are deliberately global: subprocess
shutdown, the pre-load kill and unload_model.
* Studio: add the AGPL-3.0 header to the first-turn identity test
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: stop the dispatcher and a _gen_lock stream fighting over the response queue
Nothing stopped the dispatcher starting under a _gen_lock generation, so once compare
was opened while an ordinary chat was still streaming, both consumed _resp_queue and
whichever response the dispatcher took without a mailbox was dropped, gen_done included.
That chat truncated or hung. This PR is what makes it reachable, since navigating into
compare no longer ends the chat behind it.
Delaying the dispatcher would serialise compare behind whatever chat happens to be
streaming, so the direct readers get a mailbox instead. _direct_reader returns a reader,
a cancel drain and a release, and files the mailbox under _direct_mailboxes rather than
_mailboxes, which means "compare requests are in flight" to the unload and distributed
paths and must not count an ordinary chat.
Both directions close. The dispatcher finds the direct reader's mailbox instead of
dropping. And this reader can already be blocked on the queue when a compare request's
dispatcher starts, so a response that is not ours goes to its own mailbox rather than
being consumed, which would have corrupted the chat and hung the pane. All three
_gen_lock readers use it, and the cancel drain goes through it too.
The sidebar's return target still picked a raw pane id while the count grouped by row,
and /chat addresses compare with `compare`, not `thread`. It resolves through the same
items now, so a running compare row returns to its pair.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: keep worker ownership honest across audio, API traffic and a replaced worker
The audio-input send got a mailbox last round but stayed unclaimed, so a compare request
queued behind it looked like the oldest owner and stopping that queued request signalled
the shared event into the audio chat. It claims under the send lock and releases in the
finally, like _generate_inner.
Ownership is keyed on cancel-event identity with nothing tying it to a worker generation,
so a consumer still blocked on its mailbox when the process was replaced stayed recorded
as the executor, and a generation on the fresh worker could not be stopped.
_shutdown_subprocess clears that state once the process is confirmed dead, mailboxes
included: nothing routes to them again, and a stale one reads as compare activity to the
unload path. Not on the survived-SIGKILL path, which keeps its handle on purpose.
The four public /v1/messages trackers were registering as chats. The distinction is a
Studio thread, not the protocol, and those branches already say "No thread_id: public API
surface" while the Studio path passes payload.thread_id separately. They carry their own
kind now, so the swap prompt stops calling an external request a chat.
The swap confirmation still counted raw pane ids, so a compare conversation asked to stop
two chats and listed its title twice. It folds panes onto pairId and lowers the count by
what it collapsed, leaving a first turn the backend can count but not name.
Deep Research set runningByThreadId but registered no server-cancel handle, and that map
is how Stop, archive and delete reach a thread that is no longer active. Leaving the
outgoing thread running is this PR's doing, so the run was left unreachable while its
supervisor kept working against a conversation the user could delete.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: tighten the parallel-chats comments
* Studio: replay a Deep Research stop that arrived before the run existed
The handle is registered before createResearchRun resolves because the thread can be
stopped while that request is in flight, but it had no id to act on and dropped the stop.
The supervisor then followed a run the user had already stopped, archived or deleted.
It latches instead: a stop with no id yet sets a flag, and the adapter replays it against
the id the moment creation returns rather than starting to follow.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: fix worker ownership on a raced reroute, and the stop-chats prompt
Four review findings on the parallel-chats work, all reproduced first.
- _direct_reader hands a foreign response to its own mailbox, but skipped the
ownership move the dispatcher makes. A _gen_lock reader already blocked on
resp_queue can beat the compare dispatcher to that request's first response,
and the compare consumer opts out of marking, so nothing promoted it: the
direct request stayed the recorded executor, its late reset cancelled the
compare generation, and the compare chat's own Stop was ignored.
- A chat stopped while queued on _gen_lock was still claimed and sent once the
lock freed. Cancellation is only checked on a token, so a long prefill, or a
generation reaching gen_done without one, occupied the worker after Stop.
Same hole in the audio-input path, which shares the lock.
- The stop-chats prompt counted generation handles, not conversations. One chat
holds several while a tool continuation registers its next leg before the
previous unwinds, so it offered to stop two chats and listed one title.
- Ejecting a model confirms through that dialog, which told the user
"Unloading the model reloads the model" and offered "Stop and reload".
Confirming calls /unload and leaves nothing loaded.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: name the TTS run's thread so the stop prompt counts it once
The audio branch registers its run locally under the thread key but sent no
thread_id, so the backend tracker filed the same generation under no thread.
The stop-chats prompt then had a named local run and an unnamed backend one and,
since e8e7594 started adding unnamed entries to the named ones, counted a single
TTS chat as two requests. The backend already reads payload.thread_id, so
sending it lines both registries up on the same run.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
99e1f402c7
commit
c608649552
73 changed files with 6996 additions and 935 deletions
|
|
@ -2281,8 +2281,13 @@ class InferenceBackend:
|
|||
except Exception as e:
|
||||
logger.warning(f"Could not fully reset model state for {model_name}: {e}")
|
||||
|
||||
def reset_generation_state(self):
|
||||
"""Reset any cached generation state to prevent hanging after errors"""
|
||||
def reset_generation_state(self, caller_cancel_event = None):
|
||||
"""Reset any cached generation state to prevent hanging after errors
|
||||
|
||||
``caller_cancel_event`` is accepted for signature parity with the
|
||||
orchestrator, which uses it to drop a reset from a request that never
|
||||
started. Nothing here cancels a live generation, so it is unused.
|
||||
"""
|
||||
try:
|
||||
# Clear cached state for ALL loaded models
|
||||
for model_name in self.models.keys():
|
||||
|
|
|
|||
|
|
@ -214,7 +214,7 @@ class _Waiter:
|
|||
|
||||
|
||||
class LlamaAdmissionLease:
|
||||
__slots__ = ("_queue", "_slot", "_released", "_release_lock")
|
||||
__slots__ = ("_queue", "_slot", "_released", "_release_lock", "_parked")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -225,20 +225,88 @@ class LlamaAdmissionLease:
|
|||
self._slot = slot
|
||||
self._released = False
|
||||
self._release_lock = threading.Lock()
|
||||
self._parked = False
|
||||
|
||||
@property
|
||||
def slot(self) -> Optional[int]:
|
||||
"""Pool slot this lease holds, or None when admission is disabled."""
|
||||
return self._slot
|
||||
|
||||
def park(self) -> None:
|
||||
"""Hand the slot back while this holder waits on something off the GPU.
|
||||
|
||||
A run stopped on a tool approval prompt is not decoding, so holding its
|
||||
slot would let unanswered prompts fill the pool while llama-server idles.
|
||||
The lease itself stays valid: releasing it after a park is still correct.
|
||||
"""
|
||||
queue = self._queue
|
||||
slot = None
|
||||
with self._release_lock:
|
||||
if queue is None or self._released or self._parked:
|
||||
return
|
||||
self._parked = True
|
||||
slot, self._slot = self._slot, None
|
||||
queue.park(slot)
|
||||
|
||||
def unpark(self) -> None:
|
||||
"""Drop the parked state without reclaiming a slot.
|
||||
|
||||
For a holder that is tearing down: it will not decode again. Resuming
|
||||
holders must use ``unpark_async``, which waits for a slot instead of
|
||||
going back to llama-server past the admission limit.
|
||||
"""
|
||||
with self._release_lock:
|
||||
if not self._parked:
|
||||
return
|
||||
self._parked = False
|
||||
if self._queue is not None:
|
||||
self._queue.unpark()
|
||||
|
||||
async def unpark_async(
|
||||
self,
|
||||
*,
|
||||
cancel_event = None,
|
||||
poll_s: float = 0.02,
|
||||
) -> None:
|
||||
"""Take a slot back, waiting until the pool has room.
|
||||
|
||||
``park`` gave the slot to a waiter, so by the time the user answers the
|
||||
prompt someone else may be decoding in it. Resuming regardless put two
|
||||
holders on a one-slot server. Gives up if the caller is cancelled, since
|
||||
the holder is then leaving anyway and must not be stuck here.
|
||||
"""
|
||||
queue = self._queue
|
||||
if queue is None or not self._parked:
|
||||
return
|
||||
slot = await queue.acquire_parked_slot(cancel_event = cancel_event, poll_s = poll_s)
|
||||
stranded = None
|
||||
with self._release_lock:
|
||||
# release() may have run during the wait; it clears the flag and does
|
||||
# the unpark itself, so only the caller that clears it here repeats one.
|
||||
parked, self._parked = self._parked, False
|
||||
if self._released:
|
||||
# Released while waiting: this lease will never hand the slot
|
||||
# back, so return it here rather than strand it for good.
|
||||
stranded = slot
|
||||
else:
|
||||
self._slot = slot
|
||||
if parked:
|
||||
queue.unpark()
|
||||
if stranded is not None:
|
||||
queue.release(stranded)
|
||||
|
||||
def release(self) -> None:
|
||||
queue = None
|
||||
parked = False
|
||||
with self._release_lock:
|
||||
if self._released:
|
||||
return
|
||||
self._released = True
|
||||
queue = self._queue
|
||||
parked, self._parked = self._parked, False
|
||||
if queue is not None:
|
||||
if parked:
|
||||
queue.unpark()
|
||||
queue.release(self._slot)
|
||||
|
||||
async def __aenter__(self) -> "LlamaAdmissionLease":
|
||||
|
|
@ -338,7 +406,18 @@ class LlamaAdmissionQueue:
|
|||
set to 0. See ``LlamaAdmissionConfig.queue_limit``.
|
||||
"""
|
||||
|
||||
__slots__ = ("key", "_lock", "_capacity", "_free", "_in_use", "_held", "_waiters")
|
||||
__slots__ = (
|
||||
"key",
|
||||
"_lock",
|
||||
"_capacity",
|
||||
"_free",
|
||||
"_in_use",
|
||||
"_held",
|
||||
"_waiters",
|
||||
"_parked",
|
||||
"_unpark_tickets",
|
||||
"_unpark_seq",
|
||||
)
|
||||
|
||||
def __init__(self, key: str):
|
||||
self.key = key
|
||||
|
|
@ -351,6 +430,13 @@ class LlamaAdmissionQueue:
|
|||
self._in_use = 0
|
||||
self._held = 0
|
||||
self._waiters: Deque[_Waiter] = deque()
|
||||
# Holders parked on a tool approval prompt. They hold no slot, so this only
|
||||
# keeps the queue off the idle-eviction list while they are away.
|
||||
self._parked = 0
|
||||
# FIFO tickets for holders resuming from a park (see acquire_parked_slot). A
|
||||
# bare count deadlocked: every approved holder blocked every other one.
|
||||
self._unpark_tickets: Deque[int] = deque()
|
||||
self._unpark_seq = 0
|
||||
|
||||
def _resize_pool_locked(self, capacity: int) -> None:
|
||||
# Slots past a shrunk capacity retire when their holder releases them.
|
||||
|
|
@ -359,13 +445,15 @@ class LlamaAdmissionQueue:
|
|||
self._capacity = capacity
|
||||
self._free = [slot for slot in range(capacity) if not self._in_use >> slot & 1]
|
||||
|
||||
def _can_admit_locked(self) -> bool:
|
||||
def _can_admit_locked(self, reserved: int) -> bool:
|
||||
# Slots still held above a shrunk capacity keep occupying the backend, so
|
||||
# count every held slot against the ceiling, not just the ids below it.
|
||||
return bool(self._free) and self._held < self._capacity
|
||||
# ``reserved`` holds slots back for approved holders waiting to resume;
|
||||
# without it a stream of new arrivals took the next slot, forever.
|
||||
return bool(self._free) and (self._held + reserved) < self._capacity
|
||||
|
||||
def _take_slot_locked(self) -> Optional[int]:
|
||||
if not self._can_admit_locked():
|
||||
def _take_slot_locked(self, reserved: int) -> Optional[int]:
|
||||
if not self._can_admit_locked(reserved):
|
||||
return None
|
||||
slot = self._free.pop()
|
||||
self._in_use |= 1 << slot
|
||||
|
|
@ -386,7 +474,7 @@ class LlamaAdmissionQueue:
|
|||
self._resize_pool_locked(capacity)
|
||||
self._grant_waiters_locked()
|
||||
if not self._waiters:
|
||||
slot = self._take_slot_locked()
|
||||
slot = self._take_slot_locked(len(self._unpark_tickets))
|
||||
if slot is not None:
|
||||
# No snapshot here: callers read it through snapshot_now(),
|
||||
# which re-reads the queue, so building one per admitted
|
||||
|
|
@ -425,6 +513,58 @@ class LlamaAdmissionQueue:
|
|||
self._release_slot_locked(slot)
|
||||
self._grant_waiters_locked()
|
||||
|
||||
def park(self, slot: Optional[int]) -> None:
|
||||
"""Return a parked holder's slot to the pool. See ``LlamaAdmissionLease.park``."""
|
||||
with self._lock:
|
||||
self._parked += 1
|
||||
self._release_slot_locked(slot)
|
||||
self._grant_waiters_locked()
|
||||
|
||||
def unpark(self) -> None:
|
||||
with self._lock:
|
||||
if self._parked > 0:
|
||||
self._parked -= 1
|
||||
|
||||
async def acquire_parked_slot(
|
||||
self,
|
||||
*,
|
||||
cancel_event = None,
|
||||
poll_s: float = 0.02,
|
||||
) -> Optional[int]:
|
||||
"""Wait for a slot for a holder resuming from a park, None if cancelled.
|
||||
|
||||
Ordered by ticket rather than counted, so approvals resume in the order
|
||||
they came back: counting them made every approved holder block every
|
||||
other one, and with nothing decoding that never resolved.
|
||||
"""
|
||||
with self._lock:
|
||||
self._unpark_seq += 1
|
||||
ticket = self._unpark_seq
|
||||
self._unpark_tickets.append(ticket)
|
||||
try:
|
||||
while True:
|
||||
with self._lock:
|
||||
ahead = 0
|
||||
for queued in self._unpark_tickets:
|
||||
if queued == ticket:
|
||||
break
|
||||
ahead += 1
|
||||
# Only the approvals ahead of this one hold slots back from it.
|
||||
slot = self._take_slot_locked(ahead)
|
||||
if slot is not None:
|
||||
return slot
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
return None
|
||||
await asyncio.sleep(poll_s)
|
||||
finally:
|
||||
with self._lock:
|
||||
try:
|
||||
self._unpark_tickets.remove(ticket)
|
||||
except ValueError:
|
||||
pass
|
||||
# This ticket was holding a slot back from the wait line.
|
||||
self._grant_waiters_locked()
|
||||
|
||||
def cancel(self, waiter: _Waiter) -> None:
|
||||
lease_to_release = None
|
||||
with self._lock:
|
||||
|
|
@ -455,15 +595,17 @@ class LlamaAdmissionQueue:
|
|||
def is_idle(self) -> bool:
|
||||
with self._lock:
|
||||
self._prune_waiters_locked()
|
||||
return self._in_use == 0 and not self._waiters
|
||||
# A parked holder owns no slot but is coming back to this queue, so
|
||||
# evicting it here would resume it against a fresh 1-slot pool.
|
||||
return self._in_use == 0 and not self._waiters and not self._parked
|
||||
|
||||
def _grant_waiters_locked(self) -> None:
|
||||
# Dead waiters are skipped as they are popped, so no prune is needed here.
|
||||
while self._waiters and self._can_admit_locked():
|
||||
while self._waiters and self._can_admit_locked(len(self._unpark_tickets)):
|
||||
waiter = self._waiters.popleft()
|
||||
if waiter.cancelled or waiter.future.done():
|
||||
continue
|
||||
slot = self._take_slot_locked()
|
||||
slot = self._take_slot_locked(len(self._unpark_tickets))
|
||||
lease = LlamaAdmissionLease(self, slot)
|
||||
waiter.granted_lease = lease
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -98,6 +98,7 @@ from core.inference.tool_call_parser import (
|
|||
from core.inference.tool_loop_controller import (
|
||||
ToolLoopController,
|
||||
append_deferred_nudges,
|
||||
awaiting_approval_status,
|
||||
tool_event_provenance,
|
||||
)
|
||||
from state.tool_approvals import (
|
||||
|
|
@ -6551,6 +6552,25 @@ class LlamaCppBackend:
|
|||
binary = self._find_llama_server_binary()
|
||||
is_vulkan_backend = self._is_vulkan_backend(binary)
|
||||
|
||||
# Without --kv-unified an explicit --parallel N splits -c into windows of -c/N, so on a
|
||||
# build lacking the flag the default of 4 would quarter every context window for a
|
||||
# feature it cannot serve: fall back to one slot. Ahead of the KV estimates so the
|
||||
# fit matches what launches.
|
||||
if (
|
||||
n_parallel > 1
|
||||
and binary
|
||||
and not self.probe_server_capabilities(binary).get("supports_kv_unified")
|
||||
):
|
||||
logger.warning(
|
||||
"llama-server at %s has no --kv-unified, so %d parallel slots would "
|
||||
"split the context window %d ways. Using 1 slot instead; update "
|
||||
"llama.cpp to run chats in parallel.",
|
||||
binary,
|
||||
n_parallel,
|
||||
n_parallel,
|
||||
)
|
||||
n_parallel = 1
|
||||
|
||||
# ── Vulkan-ordinal preflight (BEFORE the Phase 1 kill) ────────
|
||||
# An explicit Vulkan pin the ggml probe never enumerated cannot be honored.
|
||||
# Validate it ABOVE the kill so an invalid selection leaves the live model
|
||||
|
|
@ -11101,6 +11121,7 @@ class LlamaCppBackend:
|
|||
from core.inference.tools import (
|
||||
build_rag_autoinject,
|
||||
execute_tool,
|
||||
has_text_only_provisional_card,
|
||||
is_always_safe_tool,
|
||||
is_high_risk_tool_call,
|
||||
)
|
||||
|
|
@ -11527,6 +11548,9 @@ class LlamaCppBackend:
|
|||
permission_mode == "auto"
|
||||
and is_always_safe_tool(current_name)
|
||||
)
|
||||
# A text-preview card still streams while gated;
|
||||
# hiding it blanks the chat.
|
||||
and not has_text_only_provisional_card(current_name)
|
||||
)
|
||||
# Keep small-argument tools on the normal path.
|
||||
_args_len = len(
|
||||
|
|
@ -11628,20 +11652,27 @@ class LlamaCppBackend:
|
|||
# TEXT call to a provisional card. Gated on an enabled-name
|
||||
# sniff + size floor so prose/small calls spawn no pane; id
|
||||
# matches the first call so the final tool_start reconciles.
|
||||
if (
|
||||
not has_structured_tc
|
||||
and not _confirm_gated_iteration
|
||||
and _text_args_call_start >= 0
|
||||
):
|
||||
if not has_structured_tc and _text_args_call_start >= 0:
|
||||
if not _text_args_id:
|
||||
_call_text = content_accum[_text_args_call_start:]
|
||||
_sniffed = _sniff_text_tool_name(
|
||||
_call_text, _enabled_tool_names
|
||||
)
|
||||
if _sniffed and (
|
||||
_sniffed == "render_html"
|
||||
or len(_call_text)
|
||||
>= _PROVISIONAL_ARGS_MIN_CHARS
|
||||
# Structured-path rule: gated calls
|
||||
# stream only from a text-preview card.
|
||||
if (
|
||||
_sniffed
|
||||
and not (
|
||||
_confirm_gated_iteration
|
||||
and not has_text_only_provisional_card(
|
||||
_sniffed
|
||||
)
|
||||
)
|
||||
and (
|
||||
_sniffed == "render_html"
|
||||
or len(_call_text)
|
||||
>= _PROVISIONAL_ARGS_MIN_CHARS
|
||||
)
|
||||
):
|
||||
_text_args_id = "call_0"
|
||||
_text_args_name = _sniffed
|
||||
|
|
@ -12230,18 +12261,31 @@ class LlamaCppBackend:
|
|||
start_event["awaiting_confirmation"] = needs_confirm
|
||||
|
||||
try:
|
||||
yield {"type": "status", "text": decision.status_text}
|
||||
# Gated calls are not running yet; a "Running ..." badge
|
||||
# counting up while it waits on a human reads as a hang.
|
||||
yield {
|
||||
"type": "status",
|
||||
"text": (
|
||||
awaiting_approval_status(decision.tool_name)
|
||||
if needs_confirm
|
||||
else decision.status_text
|
||||
),
|
||||
}
|
||||
yield start_event
|
||||
|
||||
if (
|
||||
decision_slot is not None
|
||||
and wait_tool_decision(
|
||||
_decision = (
|
||||
wait_tool_decision(
|
||||
decision_slot,
|
||||
approval_id,
|
||||
cancel_event = cancel_event,
|
||||
)
|
||||
== "deny"
|
||||
):
|
||||
if decision_slot is not None
|
||||
else None
|
||||
)
|
||||
if _decision is not None and _decision != "deny":
|
||||
# Approved: now it really is running.
|
||||
yield {"type": "status", "text": decision.status_text}
|
||||
if _decision == "deny":
|
||||
decision_slot = None
|
||||
resolved_provisional_tool_call_ids.add(decision.tool_call_id)
|
||||
yield {
|
||||
|
|
@ -12809,10 +12853,15 @@ class LlamaCppBackend:
|
|||
min_p: float = 0.0,
|
||||
max_new_tokens: int = 2048,
|
||||
repetition_penalty: float = 1.1,
|
||||
cancel_event: Optional[threading.Event] = None,
|
||||
) -> tuple:
|
||||
"""
|
||||
Generate TTS audio via llama-server /completion + codec decode.
|
||||
Returns (wav_bytes, sample_rate).
|
||||
|
||||
``cancel_event`` lets a Stop or a forced model swap end the request: the
|
||||
decode is one blocking POST, so a watcher closes the client out from under
|
||||
it rather than polling. Raises RuntimeError once cancelled.
|
||||
"""
|
||||
if audio_type not in self._TTS_PROMPTS:
|
||||
raise RuntimeError(f"GGUF TTS does not support '{audio_type}' codec.")
|
||||
|
|
@ -12834,15 +12883,47 @@ class LlamaCppBackend:
|
|||
if need_ids:
|
||||
payload["n_probs"] = 1
|
||||
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
raise RuntimeError("Audio generation cancelled")
|
||||
|
||||
with httpx.Client(
|
||||
timeout = httpx.Timeout(300, connect = 10),
|
||||
headers = self._auth_headers,
|
||||
trust_env = False,
|
||||
) as client:
|
||||
resp = client.post(f"{self.base_url}/completion", json = payload)
|
||||
finished = threading.Event()
|
||||
watcher: Optional[threading.Thread] = None
|
||||
if cancel_event is not None:
|
||||
|
||||
def _close_when_cancelled() -> None:
|
||||
while not finished.wait(0.05):
|
||||
if cancel_event.is_set():
|
||||
# Closing mid-request makes the blocking post raise
|
||||
# httpx.RequestError, the only way out of it.
|
||||
with contextlib.suppress(Exception):
|
||||
client.close()
|
||||
return
|
||||
|
||||
watcher = threading.Thread(target = _close_when_cancelled, daemon = True)
|
||||
watcher.start()
|
||||
try:
|
||||
resp = client.post(f"{self.base_url}/completion", json = payload)
|
||||
except httpx.RequestError:
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
raise RuntimeError("Audio generation cancelled") from None
|
||||
raise
|
||||
finally:
|
||||
finished.set()
|
||||
if watcher is not None:
|
||||
watcher.join(timeout = 0.5)
|
||||
if resp.status_code != 200:
|
||||
raise RuntimeError(f"llama-server returned {resp.status_code}: {resp.text}")
|
||||
|
||||
# The codec decode below is GPU work with no interruption point, so check here:
|
||||
# cancelling after this only wastes the decode it cannot stop.
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
raise RuntimeError("Audio generation cancelled")
|
||||
|
||||
data = resp.json()
|
||||
token_ids = (
|
||||
[p["id"] for p in data.get("completion_probabilities", []) if "id" in p]
|
||||
|
|
|
|||
|
|
@ -1189,7 +1189,8 @@ class MLXInferenceBackend:
|
|||
**gen_kwargs,
|
||||
)
|
||||
|
||||
def reset_generation_state(self):
|
||||
def reset_generation_state(self, caller_cancel_event = None):
|
||||
# caller_cancel_event: signature parity with the orchestrator; unused here.
|
||||
import mlx.core as mx
|
||||
import gc
|
||||
|
||||
|
|
|
|||
|
|
@ -104,6 +104,14 @@ class InferenceOrchestrator:
|
|||
# so a generate queued behind the cancelled one is skipped, not run.
|
||||
self._drain_event: Any = None
|
||||
self._gen_lock = threading.Lock() # Serializes generation
|
||||
# Cancel event of the request holding _gen_lock: lets a Stop tell whether it owns the
|
||||
# running generation or is queued behind it (the worker's event is shared).
|
||||
self._active_cancel_events: list = []
|
||||
self._executing_cancel_events: list = []
|
||||
self._active_cancel_lock = threading.Lock()
|
||||
# Held across claim + _send_cmd so claim order matches the subprocess dequeue order,
|
||||
# which _owns_worker relies on.
|
||||
self._send_order_lock = threading.Lock()
|
||||
# Set during a switch so a generation winning the _gen_lock handoff bails
|
||||
# instead of starting on the outgoing model.
|
||||
self._unload_pending = False
|
||||
|
|
@ -112,6 +120,13 @@ class InferenceOrchestrator:
|
|||
# bypass _gen_lock, send commands directly, read from per-request
|
||||
# mailboxes routed by a dispatcher thread on request_id.
|
||||
self._mailboxes: dict[str, queue.Queue] = {}
|
||||
# request_id -> cancel event, so the dispatcher can move worker ownership as it routes.
|
||||
# Consumers read their mailbox whenever they get to it, so only the dispatcher sees
|
||||
# responses in the order the worker produced them.
|
||||
self._request_cancel_events: dict[str, object] = {}
|
||||
# Mailboxes for the _gen_lock generations. Kept apart from _mailboxes because that map
|
||||
# means "compare requests are in flight" to the unload and distributed paths.
|
||||
self._direct_mailboxes: dict[str, queue.Queue] = {}
|
||||
self._mailbox_lock = threading.Lock()
|
||||
self._dispatcher_thread: Optional[threading.Thread] = None
|
||||
self._dispatcher_stop = threading.Event()
|
||||
|
|
@ -321,9 +336,27 @@ class InferenceOrchestrator:
|
|||
self._resp_queue = None
|
||||
self._cancel_event = None
|
||||
self._drain_event = None
|
||||
self._reset_worker_scoped_state()
|
||||
logger.info("Inference subprocess shut down")
|
||||
return True
|
||||
|
||||
def _reset_worker_scoped_state(self) -> None:
|
||||
"""Drop bookkeeping that only means anything for the worker that just died.
|
||||
|
||||
Ownership is scoped by cancel-event identity alone, so a consumer still blocked
|
||||
on its mailbox when the process was replaced stayed recorded as the executor. A
|
||||
generation on the fresh worker then failed _owns_worker and could not be stopped.
|
||||
Mailboxes go too: nothing will ever route to them, and a stale one reads as
|
||||
compare activity to the unload path.
|
||||
"""
|
||||
with self._active_cancel_lock:
|
||||
self._active_cancel_events.clear()
|
||||
self._executing_cancel_events.clear()
|
||||
with self._mailbox_lock:
|
||||
self._mailboxes.clear()
|
||||
self._direct_mailboxes.clear()
|
||||
self._request_cancel_events.clear()
|
||||
|
||||
def _cleanup(self):
|
||||
"""atexit handler."""
|
||||
self._shutdown_subprocess(timeout = 5.0)
|
||||
|
|
@ -463,6 +496,74 @@ class InferenceOrchestrator:
|
|||
except (EOFError, OSError, ValueError):
|
||||
return events
|
||||
|
||||
def _direct_reader(self, request_id: str):
|
||||
"""Response reader for a _gen_lock generation, safe once compare exists.
|
||||
|
||||
The dispatcher and this reader would otherwise both consume _resp_queue. A
|
||||
dispatcher started mid-stream took our responses and dropped them as
|
||||
unaddressed (truncating or hanging the chat), and this reader, already blocked
|
||||
on the queue, could take a compare request's response before that dispatcher
|
||||
saw it. Registering a mailbox fixes the first; handing foreign responses to
|
||||
their own mailbox fixes the second.
|
||||
|
||||
Returns (read_one, drain, release).
|
||||
"""
|
||||
mailbox: queue.Queue = queue.Queue()
|
||||
with self._mailbox_lock:
|
||||
self._direct_mailboxes[request_id] = mailbox
|
||||
|
||||
def read_one(timeout: float = 1.0):
|
||||
try:
|
||||
return mailbox.get_nowait()
|
||||
except queue.Empty:
|
||||
pass
|
||||
thread = self._dispatcher_thread
|
||||
if thread is not None and thread.is_alive():
|
||||
# It owns the queue now, and it routes to us.
|
||||
try:
|
||||
return mailbox.get(timeout = timeout)
|
||||
except queue.Empty:
|
||||
return None
|
||||
resp = self._read_resp(timeout = timeout)
|
||||
if resp is None:
|
||||
return None
|
||||
rid = resp.get("request_id")
|
||||
if rid and rid != request_id:
|
||||
with self._mailbox_lock:
|
||||
other = self._mailboxes.get(rid) or self._direct_mailboxes.get(rid)
|
||||
owner = self._request_cancel_events.get(rid)
|
||||
if other is not None:
|
||||
# We beat the dispatcher to this response, so make its ownership move here
|
||||
# too. The compare consumer opts out of marking, so nothing else promotes
|
||||
# or retires that request: skipping it left this one recorded as the
|
||||
# executor, ignoring its Stop and letting a late reset cancel it.
|
||||
if owner is not None:
|
||||
if resp.get("type", "") in ("gen_done", "gen_error"):
|
||||
self._release_worker(owner)
|
||||
else:
|
||||
self._mark_worker_started(owner)
|
||||
other.put(resp)
|
||||
return None
|
||||
return resp
|
||||
|
||||
def drain(timeout: float = 5.0) -> None:
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
resp = read_one(timeout = min(0.5, deadline - time.monotonic()))
|
||||
if resp is None:
|
||||
if not self._ensure_subprocess_alive():
|
||||
return
|
||||
continue
|
||||
if resp.get("type", "") in ("gen_done", "gen_error"):
|
||||
return
|
||||
logger.warning("Timed out waiting for gen_done after cancel")
|
||||
|
||||
def release() -> None:
|
||||
with self._mailbox_lock:
|
||||
self._direct_mailboxes.pop(request_id, None)
|
||||
|
||||
return read_one, drain, release
|
||||
|
||||
def _drain_until_gen_done(self, timeout: float = 5.0) -> None:
|
||||
"""Consume resp_queue events until gen_done/gen_error, discarding them.
|
||||
|
||||
|
|
@ -542,6 +643,7 @@ class InferenceOrchestrator:
|
|||
cancel_event = None,
|
||||
stats_holder: Optional[dict] = None,
|
||||
read_timeout: float = 30.0,
|
||||
mark_started: bool = True,
|
||||
) -> Generator[str, None, None]:
|
||||
"""Yield tokens from a response stream until gen_done/gen_error.
|
||||
|
||||
|
|
@ -578,6 +680,11 @@ class InferenceOrchestrator:
|
|||
rtype = resp.get("type", "")
|
||||
if rtype == "status":
|
||||
continue
|
||||
# The worker is answering THIS request, so it is the one executing: only now may its
|
||||
# cancel event speak for the shared worker one. The dispatched path opts out: its
|
||||
# dispatcher already did this in worker order, which a mailbox read can lag behind.
|
||||
if mark_started:
|
||||
self._mark_worker_started(cancel_event)
|
||||
# Subprocess-level error (no request_id); request-scoped failures
|
||||
# arrive as gen_error below.
|
||||
if rtype == "error" and not resp.get("request_id"):
|
||||
|
|
@ -587,7 +694,13 @@ class InferenceOrchestrator:
|
|||
if rtype == "token":
|
||||
# Cancel from route (e.g. SSE connection closed).
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
self._cancel_generation()
|
||||
# Same rule as reset_generation_state: the shared worker event may only be set by
|
||||
# the generation the worker is running. A dispatched request can still be draining
|
||||
# stale mailbox tokens after the dispatcher retired it, and signalling from here
|
||||
# would end the next one instead. Tearing this stream down is always safe, so the
|
||||
# local drain happens either way.
|
||||
if self._owns_worker(cancel_event):
|
||||
self._cancel_generation()
|
||||
drain_on_cancel()
|
||||
return
|
||||
yield resp.get("text", "")
|
||||
|
|
@ -681,8 +794,17 @@ class InferenceOrchestrator:
|
|||
# Route to mailbox if a matching request_id exists
|
||||
if rid:
|
||||
with self._mailbox_lock:
|
||||
mbox = self._mailboxes.get(rid)
|
||||
mbox = self._mailboxes.get(rid) or self._direct_mailboxes.get(rid)
|
||||
owner = self._request_cancel_events.get(rid)
|
||||
if mbox is not None:
|
||||
# Worker order, not consumer order: retire a request the moment its last response
|
||||
# is routed. Waiting for the consumer's finally left it owning the worker after
|
||||
# the worker moved on, so a late Stop for it cancelled whichever request started next.
|
||||
if owner is not None:
|
||||
if rtype in ("gen_done", "gen_error"):
|
||||
self._release_worker(owner)
|
||||
else:
|
||||
self._mark_worker_started(owner)
|
||||
mbox.put(resp)
|
||||
continue
|
||||
|
||||
|
|
@ -798,6 +920,8 @@ class InferenceOrchestrator:
|
|||
)
|
||||
if not unloading:
|
||||
self._mailboxes[request_id] = mailbox
|
||||
if cancel_event is not None:
|
||||
self._request_cancel_events[request_id] = cancel_event
|
||||
# When bailing without a mailbox, note whether any OTHER compare request still
|
||||
# routes through the dispatcher; if none and this call started it, stop it below.
|
||||
orphaned_dispatcher = unloading and not dispatcher_preexisting and not self._mailboxes
|
||||
|
|
@ -813,11 +937,19 @@ class InferenceOrchestrator:
|
|||
yield GenStreamError("Error: model is being unloaded", public = True)
|
||||
return
|
||||
|
||||
# Claim before sending, like the locked path: dispatched runs are concurrent by design,
|
||||
# so without this a Stop on one saw no owner and reset the worker, ending its siblings.
|
||||
# Claim and enqueue under one lock, or two dispatcher threads interleave and claim order
|
||||
# stops matching the subprocess's command order, which _owns_worker reads.
|
||||
try:
|
||||
self._send_cmd(cmd)
|
||||
with self._send_order_lock:
|
||||
self._claim_worker(cancel_event)
|
||||
self._send_cmd(cmd)
|
||||
except RuntimeError as exc:
|
||||
self._release_worker(cancel_event)
|
||||
with self._mailbox_lock:
|
||||
self._mailboxes.pop(request_id, None)
|
||||
self._request_cancel_events.pop(request_id, None)
|
||||
yield GenStreamError(f"Error: {exc}")
|
||||
return
|
||||
|
||||
|
|
@ -836,10 +968,15 @@ class InferenceOrchestrator:
|
|||
cancel_event = cancel_event,
|
||||
stats_holder = stats_holder,
|
||||
read_timeout = _DISPATCH_READ_TIMEOUT,
|
||||
mark_started = False,
|
||||
)
|
||||
finally:
|
||||
# Normally already retired by the dispatcher at gen_done; this covers streams that
|
||||
# end without one (cancel, disconnect, a dead subprocess).
|
||||
self._release_worker(cancel_event)
|
||||
with self._mailbox_lock:
|
||||
self._mailboxes.pop(request_id, None)
|
||||
self._request_cancel_events.pop(request_id, None)
|
||||
|
||||
def _drain_mailbox(
|
||||
self,
|
||||
|
|
@ -1578,6 +1715,11 @@ class InferenceOrchestrator:
|
|||
# Won the lock handoff during a switch; don't start on the outgoing model.
|
||||
yield GenStreamError("Error: model is being unloaded", public = True)
|
||||
return
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
# Stopped while queued on the lock. Sending anyway occupied the worker with a
|
||||
# run the user ended: the cancel is only seen on a token, so a long prefill
|
||||
# (or a generation that reaches gen_done without one) held up its siblings.
|
||||
return
|
||||
request_id = str(uuid.uuid4())
|
||||
image_b64 = self._pil_to_base64(image) if image is not None else None
|
||||
cmd = self._build_generate_cmd(
|
||||
|
|
@ -1599,22 +1741,95 @@ class InferenceOrchestrator:
|
|||
preserve_thinking = preserve_thinking,
|
||||
)
|
||||
|
||||
# Claim the worker BEFORE sending, so a Stop on some OTHER chat -- still queued on the
|
||||
# lock above, having generated nothing -- cannot reset the generation this is starting.
|
||||
# Claiming after the send left the command running unclaimed. Released in the finally.
|
||||
# Own mailbox: a compare request can start the dispatcher while this is streaming,
|
||||
# and it would otherwise consume our responses and drop them.
|
||||
read_one, drain, release_mailbox = self._direct_reader(request_id)
|
||||
try:
|
||||
self._send_cmd(cmd)
|
||||
except RuntimeError as exc:
|
||||
yield GenStreamError(f"Error: {exc}")
|
||||
return
|
||||
try:
|
||||
with self._send_order_lock:
|
||||
self._claim_worker(cancel_event)
|
||||
self._send_cmd(cmd)
|
||||
except RuntimeError as exc:
|
||||
yield GenStreamError(f"Error: {exc}")
|
||||
return
|
||||
|
||||
yield from self._consume_token_stream(
|
||||
self._read_resp,
|
||||
lambda: self._drain_until_gen_done(timeout = 5.0),
|
||||
crash_context = "generation",
|
||||
cancel_event = cancel_event,
|
||||
stats_holder = stats_holder,
|
||||
)
|
||||
yield from self._consume_token_stream(
|
||||
read_one,
|
||||
lambda: drain(timeout = 5.0),
|
||||
crash_context = "generation",
|
||||
cancel_event = cancel_event,
|
||||
stats_holder = stats_holder,
|
||||
)
|
||||
finally:
|
||||
self._release_worker(cancel_event)
|
||||
release_mailbox()
|
||||
|
||||
def reset_generation_state(self):
|
||||
"""Cancel any ongoing generation and reset state."""
|
||||
def _claim_worker(self, cancel_event) -> None:
|
||||
"""Record this request as one the worker will run.
|
||||
|
||||
Admission only. The subprocess executes generations one at a time, so a
|
||||
dispatched request sitting behind another in the command queue is claimed
|
||||
but not executing, and must not be able to signal the shared cancel event
|
||||
(that would end whichever request IS executing). _mark_worker_started
|
||||
promotes it once the worker answers it.
|
||||
"""
|
||||
with self._active_cancel_lock:
|
||||
self._active_cancel_events.append(cancel_event)
|
||||
|
||||
def _mark_worker_started(self, cancel_event) -> None:
|
||||
"""Promote a claimed request to executing, on its first worker response.
|
||||
|
||||
Sole executor: the subprocess runs one generation at a time, so answering
|
||||
this one means it has left the previous one behind.
|
||||
"""
|
||||
if cancel_event is None:
|
||||
return
|
||||
with self._active_cancel_lock:
|
||||
if self._executing_cancel_events[:1] != [cancel_event]:
|
||||
self._executing_cancel_events[:] = [cancel_event]
|
||||
|
||||
def _release_worker(self, cancel_event) -> None:
|
||||
with self._active_cancel_lock:
|
||||
for bucket in (self._active_cancel_events, self._executing_cancel_events):
|
||||
try:
|
||||
bucket.remove(cancel_event)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
def _owns_worker(self, cancel_event) -> bool:
|
||||
"""Whether a reset from this request may signal the shared cancel event.
|
||||
|
||||
True when it is one of the EXECUTING generations, and when nothing is in
|
||||
flight at all: an error path that resets before anything started has no
|
||||
one else to interrupt, so it must not become a silent no-op. Claimed but
|
||||
queued does not count, or a Stop on a queued request would end the
|
||||
running one, including during the prefill before any response arrives.
|
||||
"""
|
||||
with self._active_cancel_lock:
|
||||
if not self._active_cancel_events:
|
||||
# Nothing in flight at all, so there is no one to protect.
|
||||
return True
|
||||
if self._executing_cancel_events:
|
||||
return any(ev is cancel_event for ev in self._executing_cancel_events)
|
||||
# Claimed but nothing has answered yet (A is in prefill). The worker takes commands
|
||||
# in order, so the oldest claim is the executor; anyone else here is queued behind it.
|
||||
return self._active_cancel_events[0] is cancel_event
|
||||
|
||||
def reset_generation_state(self, caller_cancel_event = None):
|
||||
"""Cancel any ongoing generation and reset state.
|
||||
|
||||
``caller_cancel_event`` scopes the reset to one request. The worker has a
|
||||
single cancel event and generation is serialized on _gen_lock, so a chat
|
||||
that is still queued has no generation of its own to reset: calling this
|
||||
from its Stop handler would kill whichever chat currently holds the lock.
|
||||
Pass the request's own event and the reset is dropped unless that request
|
||||
is the one running. Omit it for genuinely global resets (unload, switch).
|
||||
"""
|
||||
if caller_cancel_event is not None and not self._owns_worker(caller_cancel_event):
|
||||
return
|
||||
self._cancel_generation()
|
||||
if not self._ensure_subprocess_alive():
|
||||
return
|
||||
|
|
@ -1673,35 +1888,40 @@ class InferenceOrchestrator:
|
|||
if use_adapter is not None:
|
||||
cmd["use_adapter"] = use_adapter
|
||||
|
||||
self._send_cmd(cmd)
|
||||
# Same shared-queue hazard as _generate_inner: see _direct_reader.
|
||||
read_one, _drain, release_mailbox = self._direct_reader(request_id)
|
||||
try:
|
||||
self._send_cmd(cmd)
|
||||
|
||||
deadline = time.monotonic() + 120.0
|
||||
while time.monotonic() < deadline:
|
||||
remaining = max(0.1, deadline - time.monotonic())
|
||||
resp = self._read_resp(timeout = min(remaining, 1.0))
|
||||
deadline = time.monotonic() + 120.0
|
||||
while time.monotonic() < deadline:
|
||||
remaining = max(0.1, deadline - time.monotonic())
|
||||
resp = read_one(timeout = min(remaining, 1.0))
|
||||
|
||||
if resp is None:
|
||||
if not self._ensure_subprocess_alive():
|
||||
raise RuntimeError(self._subprocess_crash_message("audio generation"))
|
||||
continue
|
||||
if resp is None:
|
||||
if not self._ensure_subprocess_alive():
|
||||
raise RuntimeError(self._subprocess_crash_message("audio generation"))
|
||||
continue
|
||||
|
||||
rtype = resp.get("type", "")
|
||||
rtype = resp.get("type", "")
|
||||
|
||||
if rtype == "audio_done":
|
||||
wav_bytes = base64.b64decode(resp["wav_base64"])
|
||||
sample_rate = resp["sample_rate"]
|
||||
return wav_bytes, sample_rate
|
||||
if rtype == "audio_done":
|
||||
wav_bytes = base64.b64decode(resp["wav_base64"])
|
||||
sample_rate = resp["sample_rate"]
|
||||
return wav_bytes, sample_rate
|
||||
|
||||
if rtype == "audio_error":
|
||||
raise RuntimeError(resp.get("error", "Audio generation failed"))
|
||||
if rtype == "audio_error":
|
||||
raise RuntimeError(resp.get("error", "Audio generation failed"))
|
||||
|
||||
if rtype == "error":
|
||||
raise RuntimeError(resp.get("error", "Unknown error"))
|
||||
if rtype == "error":
|
||||
raise RuntimeError(resp.get("error", "Unknown error"))
|
||||
|
||||
if rtype == "status":
|
||||
continue
|
||||
if rtype == "status":
|
||||
continue
|
||||
|
||||
raise RuntimeError("Timeout waiting for audio generation (120s)")
|
||||
raise RuntimeError("Timeout waiting for audio generation (120s)")
|
||||
finally:
|
||||
release_mailbox()
|
||||
|
||||
def generate_whisper_response(
|
||||
self,
|
||||
|
|
@ -1775,6 +1995,9 @@ class InferenceOrchestrator:
|
|||
# Won the lock handoff during a switch; don't start on the outgoing model.
|
||||
yield GenStreamError("Error: model is being unloaded", public = True)
|
||||
return
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
# Stopped while queued on the lock, same as _generate_inner.
|
||||
return
|
||||
request_id = str(uuid.uuid4())
|
||||
|
||||
# numpy array -> list for mp.Queue serialization
|
||||
|
|
@ -1797,18 +2020,28 @@ class InferenceOrchestrator:
|
|||
"repetition_penalty": repetition_penalty,
|
||||
}
|
||||
|
||||
# Same shared-queue hazard as _generate_inner: see _direct_reader.
|
||||
read_one, drain, release_mailbox = self._direct_reader(request_id)
|
||||
try:
|
||||
self._send_cmd(cmd)
|
||||
except RuntimeError as exc:
|
||||
yield GenStreamError(f"Error: {exc}")
|
||||
return
|
||||
try:
|
||||
# Claim under the send lock, like _generate_inner: unclaimed, a compare request queued
|
||||
# behind this looked like the oldest owner, so stopping it killed this one.
|
||||
with self._send_order_lock:
|
||||
self._claim_worker(cancel_event)
|
||||
self._send_cmd(cmd)
|
||||
except RuntimeError as exc:
|
||||
yield GenStreamError(f"Error: {exc}")
|
||||
return
|
||||
|
||||
yield from self._consume_token_stream(
|
||||
self._read_resp,
|
||||
lambda: self._drain_until_gen_done(timeout = 5.0),
|
||||
crash_context = "audio input generation",
|
||||
cancel_event = cancel_event,
|
||||
)
|
||||
yield from self._consume_token_stream(
|
||||
read_one,
|
||||
lambda: drain(timeout = 5.0),
|
||||
crash_context = "audio input generation",
|
||||
cancel_event = cancel_event,
|
||||
)
|
||||
finally:
|
||||
self._release_worker(cancel_event)
|
||||
release_mailbox()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Local helpers (no subprocess needed)
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ from core.tool_healing import (
|
|||
from core.inference.tool_loop_controller import (
|
||||
ToolLoopController,
|
||||
append_deferred_nudges,
|
||||
awaiting_approval_status,
|
||||
coerce_tool_arguments,
|
||||
status_for_tool,
|
||||
tool_event_provenance,
|
||||
|
|
@ -1209,18 +1210,30 @@ def run_safetensors_tool_loop(
|
|||
start_event["awaiting_confirmation"] = needs_confirm
|
||||
|
||||
try:
|
||||
yield {"type": "status", "text": decision.status_text}
|
||||
# A gated call has not started: say waiting, not "Running" (GGUF parity).
|
||||
yield {
|
||||
"type": "status",
|
||||
"text": (
|
||||
awaiting_approval_status(decision.tool_name)
|
||||
if needs_confirm
|
||||
else decision.status_text
|
||||
),
|
||||
}
|
||||
yield start_event
|
||||
|
||||
if (
|
||||
decision_slot is not None
|
||||
and wait_tool_decision(
|
||||
_decision = (
|
||||
wait_tool_decision(
|
||||
decision_slot,
|
||||
approval_id,
|
||||
cancel_event = cancel_event,
|
||||
)
|
||||
== "deny"
|
||||
):
|
||||
if decision_slot is not None
|
||||
else None
|
||||
)
|
||||
if _decision is not None and _decision != "deny":
|
||||
# Approved: now it really is running.
|
||||
yield {"type": "status", "text": decision.status_text}
|
||||
if _decision == "deny":
|
||||
decision_slot = None
|
||||
if provisional_match:
|
||||
provisional_resolved = True
|
||||
|
|
|
|||
|
|
@ -238,6 +238,19 @@ def status_for_tool(tool_name: str, arguments: Mapping[str, Any]) -> str:
|
|||
return f"Calling: {tool_name}"
|
||||
|
||||
|
||||
def awaiting_approval_status(tool_name: str) -> str:
|
||||
"""Status text for a call parked on the approval prompt.
|
||||
|
||||
It has not started, so reporting "Running ..." with a climbing timer reads
|
||||
as a hang.
|
||||
"""
|
||||
if tool_name == "python":
|
||||
return "Waiting for approval: Python"
|
||||
if tool_name == "terminal":
|
||||
return "Waiting for approval: command"
|
||||
return f"Waiting for approval: {tool_name}"
|
||||
|
||||
|
||||
def is_tool_error(result: str) -> bool:
|
||||
return isinstance(result, str) and result.lstrip().startswith(TOOL_ERROR_PREFIXES)
|
||||
|
||||
|
|
|
|||
|
|
@ -3105,6 +3105,22 @@ def is_always_safe_tool(name: str) -> bool:
|
|||
return name in _ALWAYS_SAFE_TOOLS
|
||||
|
||||
|
||||
# Tools whose provisional card is only a text preview of the arguments, so it can stream
|
||||
# while awaiting approval.
|
||||
_TEXT_PREVIEW_TOOLS = frozenset({"python", "terminal"})
|
||||
|
||||
|
||||
def has_text_only_provisional_card(name: str) -> bool:
|
||||
"""True when streaming this tool's arguments before approval shows only text.
|
||||
|
||||
A large code payload takes a minute or more to write, and suppressing the
|
||||
card until the call completes leaves the chat blank the whole time. Nothing
|
||||
runs before the decision either way, and you have to read the code to make
|
||||
it.
|
||||
"""
|
||||
return name in _TEXT_PREVIEW_TOOLS
|
||||
|
||||
|
||||
def is_potentially_unsafe_tool_call(name: str, arguments: dict) -> bool:
|
||||
"""Whether a tool call must still pause for approval in auto mode.
|
||||
|
||||
|
|
|
|||
|
|
@ -191,12 +191,26 @@ class LoadRequest(BaseModel):
|
|||
"auth, UI/server mode) are rejected. Ignored for non-GGUF models."
|
||||
),
|
||||
)
|
||||
force_cancel_active: bool = Field(
|
||||
False,
|
||||
description = (
|
||||
"Stop chats still generating instead of refusing with 409. A load "
|
||||
"replaces the llama-server every open conversation decodes on."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class UnloadRequest(BaseModel):
|
||||
"""Request to unload a model"""
|
||||
|
||||
model_path: str = Field(..., description = "Model identifier to unload")
|
||||
force_cancel_active: bool = Field(
|
||||
False,
|
||||
description = (
|
||||
"Stop chats still generating instead of refusing with 409. An "
|
||||
"unload takes away the llama-server they are decoding on."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class TranscribeRequest(BaseModel):
|
||||
|
|
@ -350,6 +364,14 @@ class InstallLatestTransformersRequest(BaseModel):
|
|||
description = "Exact transformers version to install; must match the current "
|
||||
"latest PyPI release reported by /validate.",
|
||||
)
|
||||
force_cancel_active: bool = Field(
|
||||
False,
|
||||
description = (
|
||||
"Stop chats still generating instead of refusing with 409. The install "
|
||||
"is a step of the model swap that raised the same prompt, so a client "
|
||||
"that already got consent for that swap can carry it through here."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class InstallLatestTransformersResponse(BaseModel):
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -1377,13 +1377,21 @@ def _apply_cli_tool_policy(enable_tools: "Optional[bool]") -> None:
|
|||
set_tool_policy(enable_tools)
|
||||
|
||||
|
||||
# Mirror unsloth_cli/commands/studio.py's _PARALLEL_*: the admission queue caps concurrent
|
||||
# chats at the slot count, so a direct launch matches the CLI (VRAM fit may still cut it
|
||||
# back). Defined above run_server() so embedders that omit it do not serialise every chat.
|
||||
_PARALLEL_MIN = 1
|
||||
_PARALLEL_MAX = 64
|
||||
_PARALLEL_DEFAULT_PLAIN = 4
|
||||
|
||||
|
||||
def run_server(
|
||||
host: str = "127.0.0.1",
|
||||
port: int = 8888,
|
||||
frontend_path: Path = _DEFAULT_FRONTEND_PATH,
|
||||
silent: bool = False,
|
||||
api_only: bool = False,
|
||||
llama_parallel_slots: int = 1,
|
||||
llama_parallel_slots: int = _PARALLEL_DEFAULT_PLAIN,
|
||||
cloudflare: "Optional[bool]" = None,
|
||||
secure: bool = False,
|
||||
enable_tools: "Optional[bool]" = None,
|
||||
|
|
@ -1399,7 +1407,8 @@ def run_server(
|
|||
frontend_path: Path to frontend build directory (optional)
|
||||
silent: Suppress startup messages
|
||||
api_only: API server only, no frontend (for Tauri desktop app)
|
||||
llama_parallel_slots: parallel slots for llama-server
|
||||
llama_parallel_slots: parallel slots for llama-server (default
|
||||
_PARALLEL_DEFAULT_PLAIN, matching the CLI entry points)
|
||||
cloudflare: opt in to the public Cloudflare HTTPS tunnel for a wildcard
|
||||
bind. Tri-state: None (unset) and False both mean off; True enables it.
|
||||
--secure implies it (True) and rejects an explicit False.
|
||||
|
|
@ -1817,13 +1826,6 @@ def run_server(
|
|||
return app
|
||||
|
||||
|
||||
# Mirror unsloth_cli/commands/studio.py's _PARALLEL_*. Default 1 is for direct
|
||||
# backend launches; `unsloth studio run` always passes its own value (4).
|
||||
_PARALLEL_MIN = 1
|
||||
_PARALLEL_MAX = 64
|
||||
_PARALLEL_DEFAULT_PLAIN = 1
|
||||
|
||||
|
||||
def _build_arg_parser():
|
||||
"""Build the backend CLI argument parser.
|
||||
|
||||
|
|
@ -1918,7 +1920,7 @@ def _build_arg_parser():
|
|||
default = _PARALLEL_DEFAULT_PLAIN,
|
||||
help = (
|
||||
f"llama-server parallel decode slots ({_PARALLEL_MIN}..{_PARALLEL_MAX}). "
|
||||
f"Default {_PARALLEL_DEFAULT_PLAIN}; `unsloth studio run` uses 4."
|
||||
f"Default {_PARALLEL_DEFAULT_PLAIN}."
|
||||
),
|
||||
)
|
||||
return parser
|
||||
|
|
|
|||
146
studio/backend/state/active_generations.py
Normal file
146
studio/backend/state/active_generations.py
Normal file
|
|
@ -0,0 +1,146 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Registry of in-flight chat generations, keyed by conversation.
|
||||
|
||||
New Chat leaves the previous conversation streaming, so /load and /unload need
|
||||
to know which chats a reload would interrupt: they refuse with 409 unless the
|
||||
caller opts in to cancelling them, and GET /inference/active-generations lets
|
||||
the UI name them. A frontend guard alone would miss a second tab or a REST call.
|
||||
|
||||
Entries hold the same threading.Event as the per-run cancel registry in
|
||||
routes/inference.py, so cancel_all() closes each generation's own upstream
|
||||
stream and never signals llama-server itself.
|
||||
|
||||
A plain dict plus a threading.Lock: no signals, no process groups, no event loop
|
||||
affinity, so it behaves identically on Linux, macOS, Windows and WSL.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Optional
|
||||
|
||||
# handle id -> entry. Keyed by handle, not thread_id: a tool continuation can register
|
||||
# before the previous leg unregisters, and one key would drop the other.
|
||||
_ACTIVE: dict[str, dict[str, Any]] = {}
|
||||
_LOCK = threading.Lock()
|
||||
|
||||
|
||||
class ActiveGeneration:
|
||||
"""Registers one in-flight generation for the duration of the block.
|
||||
|
||||
Each __enter__ mints its own handle, so overlapping uses never clobber.
|
||||
"""
|
||||
|
||||
__slots__ = ("thread_id", "cancel_event", "model", "kind", "_handle")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cancel_event: threading.Event,
|
||||
*,
|
||||
thread_id: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
kind: str = "chat",
|
||||
):
|
||||
self.thread_id = thread_id or None
|
||||
self.cancel_event = cancel_event
|
||||
self.model = model or None
|
||||
self.kind = kind
|
||||
self._handle: Optional[str] = None
|
||||
|
||||
def __enter__(self) -> "ActiveGeneration":
|
||||
self._handle = uuid.uuid4().hex
|
||||
with _LOCK:
|
||||
_ACTIVE[self._handle] = {
|
||||
"handle": self._handle,
|
||||
"thread_id": self.thread_id,
|
||||
"model": self.model,
|
||||
"kind": self.kind,
|
||||
"started_at": time.time(),
|
||||
"event": self.cancel_event,
|
||||
}
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc) -> bool:
|
||||
handle, self._handle = self._handle, None
|
||||
if handle is not None:
|
||||
with _LOCK:
|
||||
_ACTIVE.pop(handle, None)
|
||||
return False
|
||||
|
||||
|
||||
def snapshot() -> list[dict[str, Any]]:
|
||||
"""In-flight generations, newest last. Drops the Event: this is a response."""
|
||||
with _LOCK:
|
||||
entries = list(_ACTIVE.values())
|
||||
entries.sort(key = lambda e: e["started_at"])
|
||||
return [
|
||||
{
|
||||
"handle": e["handle"],
|
||||
"thread_id": e["thread_id"],
|
||||
"model": e["model"],
|
||||
"kind": e["kind"],
|
||||
"started_at": e["started_at"],
|
||||
}
|
||||
for e in entries
|
||||
]
|
||||
|
||||
|
||||
def active_thread_ids() -> list[str]:
|
||||
"""Distinct conversation ids with a generation in flight, in start order.
|
||||
|
||||
A first turn that races persistence has no thread id yet: count() sees it,
|
||||
this cannot name it.
|
||||
"""
|
||||
seen: list[str] = []
|
||||
for e in snapshot():
|
||||
tid = e["thread_id"]
|
||||
if tid and tid not in seen:
|
||||
seen.append(tid)
|
||||
return seen
|
||||
|
||||
|
||||
def count() -> int:
|
||||
"""Number of generations currently in flight."""
|
||||
with _LOCK:
|
||||
return len(_ACTIVE)
|
||||
|
||||
|
||||
def cancel_all() -> int:
|
||||
"""Signal every in-flight generation to stop. Returns how many were signalled.
|
||||
|
||||
Only sets the cancel events; each stream tears itself down. Entries are
|
||||
removed by their own __exit__, so one mid-cleanup is neither lost nor double
|
||||
counted.
|
||||
"""
|
||||
with _LOCK:
|
||||
events = [e["event"] for e in _ACTIVE.values()]
|
||||
for ev in events:
|
||||
try:
|
||||
ev.set()
|
||||
except Exception:
|
||||
pass
|
||||
return len(events)
|
||||
|
||||
|
||||
def cancel_thread(thread_id: str) -> int:
|
||||
"""Signal only the generations belonging to ``thread_id``."""
|
||||
if not thread_id:
|
||||
return 0
|
||||
with _LOCK:
|
||||
events = [e["event"] for e in _ACTIVE.values() if e["thread_id"] == thread_id]
|
||||
for ev in events:
|
||||
try:
|
||||
ev.set()
|
||||
except Exception:
|
||||
pass
|
||||
return len(events)
|
||||
|
||||
|
||||
def reset_for_tests() -> None:
|
||||
"""Drop every entry. Test-only; never called from request paths."""
|
||||
with _LOCK:
|
||||
_ACTIVE.clear()
|
||||
2635
studio/backend/tests/test_active_generations.py
Normal file
2635
studio/backend/tests/test_active_generations.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -641,12 +641,13 @@ def test_every_dispatch_site_goes_through_admission():
|
|||
for node in ast.walk(tree)
|
||||
if isinstance(node, ast.AsyncFunctionDef) and node.name == "anthropic_messages"
|
||||
)
|
||||
# The wrappers themselves call _monitored_anthropic; only the dispatch sites count.
|
||||
# The wrappers themselves call _monitored_anthropic (the non-streaming one
|
||||
# through the swap-gate tracker); only the dispatch sites count.
|
||||
nested = {
|
||||
node
|
||||
for node in ast.walk(handler)
|
||||
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
|
||||
and node.name.startswith("_admitted_anthropic")
|
||||
and node.name.startswith(("_admitted_anthropic", "_tracked_anthropic"))
|
||||
}
|
||||
inner = {id(n) for wrapper in nested for n in ast.walk(wrapper)}
|
||||
|
||||
|
|
@ -763,12 +764,13 @@ def _passthrough_payload(**fields):
|
|||
return _payload(tools = _CLIENT_TOOLS, enable_tools = False, **fields)
|
||||
|
||||
|
||||
def test_response_pre_start_cleanup_exits_the_passthrough_tracker(monkeypatch):
|
||||
"""A disconnect before the body starts must still exit the cancel tracker.
|
||||
def test_response_pre_start_cleanup_leaves_no_passthrough_tracker(monkeypatch):
|
||||
"""A disconnect before the body starts must leave no tracker and no slot.
|
||||
|
||||
The wrapper replaces the response's own pre-start hook, so it has to chain to
|
||||
it. Asserting through _CANCEL_REGISTRY rather than the wiring, because the
|
||||
hook can be present and still be a no-op.
|
||||
The passthrough registers from inside its body rather than eagerly, so a
|
||||
generator that never runs registers nothing; the hook still has to hand the
|
||||
admission slot back. Asserting through _CANCEL_REGISTRY and the pool rather
|
||||
than the wiring, because the hook can be present and still be a no-op.
|
||||
"""
|
||||
backend = _install_backend(monkeypatch, slots = 1)
|
||||
backend.supports_tool_passthrough = True
|
||||
|
|
@ -778,7 +780,7 @@ def test_response_pre_start_cleanup_exits_the_passthrough_tracker(monkeypatch):
|
|||
response = await anthropic_messages(
|
||||
_passthrough_payload(stream = True), request = _Request(), current_subject = "t"
|
||||
)
|
||||
assert inf_mod._CANCEL_REGISTRY, "passthrough should have registered a tracker"
|
||||
assert inf_mod._CANCEL_REGISTRY == {}, "nothing runs the body's exit for it yet"
|
||||
|
||||
cleanup = getattr(response, "_unstarted_cleanup", None)
|
||||
assert cleanup is not None
|
||||
|
|
|
|||
|
|
@ -74,6 +74,10 @@ class _Request:
|
|||
class _FakeNonStreamingClient:
|
||||
def __init__(self):
|
||||
self.urls = []
|
||||
self.closed = False
|
||||
|
||||
async def aclose(self):
|
||||
self.closed = True
|
||||
|
||||
async def post(self, url, **_kwargs):
|
||||
self.urls.append(url)
|
||||
|
|
@ -189,7 +193,7 @@ def test_retry_url_tolerates_a_backend_without_respawn_hooks():
|
|||
|
||||
def test_non_streaming_retries_against_the_new_port(monkeypatch):
|
||||
client = _FakeNonStreamingClient()
|
||||
monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client)
|
||||
monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: client)
|
||||
backend = _Backend()
|
||||
|
||||
response = asyncio.run(_run_non_streaming(backend))
|
||||
|
|
@ -201,7 +205,7 @@ def test_non_streaming_retries_against_the_new_port(monkeypatch):
|
|||
|
||||
def test_non_streaming_raises_when_the_server_stays_dead(monkeypatch):
|
||||
client = _FakeNonStreamingClient()
|
||||
monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client)
|
||||
monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: client)
|
||||
backend = _Backend(respawn_ok = False)
|
||||
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
|
|
@ -212,7 +216,7 @@ def test_non_streaming_raises_when_the_server_stays_dead(monkeypatch):
|
|||
|
||||
def test_non_streaming_does_not_retry_an_mtp_crash(monkeypatch):
|
||||
client = _FakeNonStreamingClient()
|
||||
monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client)
|
||||
monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: client)
|
||||
backend = _Backend(mtp_handled = True)
|
||||
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ def _dispatcher():
|
|||
o._dispatcher_stop = threading.Event()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
return o
|
||||
|
||||
|
||||
|
|
@ -118,3 +119,68 @@ def test_route_llama_streaming_async_clients_disable_proxy_env():
|
|||
kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False
|
||||
for kw in call.keywords
|
||||
), f"httpx.AsyncClient at line {call.lineno} must set trust_env=False"
|
||||
|
||||
|
||||
def _direct_reader_host():
|
||||
"""Orchestrator with only what _direct_reader and the ownership helpers touch."""
|
||||
o = InferenceOrchestrator.__new__(InferenceOrchestrator)
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._direct_mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._active_cancel_lock = threading.Lock()
|
||||
o._active_cancel_events = []
|
||||
o._executing_cancel_events = []
|
||||
o._dispatcher_thread = None
|
||||
return o
|
||||
|
||||
|
||||
def test_rerouting_a_foreign_response_moves_worker_ownership():
|
||||
# A _gen_lock reader already blocked on resp_queue can beat the compare dispatcher to
|
||||
# that request's first response. The compare consumer passes mark_started=False, so if
|
||||
# this path does not promote it nothing does: the direct request stays recorded as the
|
||||
# executor, so the compare chat's Stop is ignored and a late reset from the direct one
|
||||
# cancels the compare generation instead.
|
||||
o = _direct_reader_host()
|
||||
mine, theirs = threading.Event(), threading.Event()
|
||||
o._request_cancel_events = {"mine": mine, "theirs": theirs}
|
||||
o._claim_worker(mine)
|
||||
o._mark_worker_started(mine)
|
||||
o._claim_worker(theirs)
|
||||
compare_mailbox = queue.Queue()
|
||||
o._mailboxes["theirs"] = compare_mailbox
|
||||
|
||||
read_one, _drain, release = _direct_reader_calls(o, "mine")
|
||||
o._scripted = [{"request_id": "theirs", "type": "token", "text": "hi"}]
|
||||
|
||||
assert read_one(timeout = 0.1) is None, "a foreign response is routed, not returned"
|
||||
assert compare_mailbox.get_nowait()["text"] == "hi"
|
||||
assert o._owns_worker(theirs), "the compare request is the one the worker answered"
|
||||
assert not o._owns_worker(mine), "so a late reset from the direct request must not fire"
|
||||
release()
|
||||
|
||||
|
||||
def test_rerouting_a_foreign_gen_done_retires_that_request():
|
||||
# The other half of the dispatcher's move: once its last response is routed, the
|
||||
# request no longer owns the worker, or a Stop for it would end whatever starts next.
|
||||
o = _direct_reader_host()
|
||||
mine, theirs = threading.Event(), threading.Event()
|
||||
o._request_cancel_events = {"mine": mine, "theirs": theirs}
|
||||
o._claim_worker(theirs)
|
||||
o._mark_worker_started(theirs)
|
||||
o._claim_worker(mine)
|
||||
o._mailboxes["theirs"] = queue.Queue()
|
||||
|
||||
read_one, _drain, release = _direct_reader_calls(o, "mine")
|
||||
o._scripted = [{"request_id": "theirs", "type": "gen_done"}]
|
||||
|
||||
assert read_one(timeout = 0.1) is None
|
||||
assert not o._owns_worker(theirs), "retired once its last response was routed"
|
||||
assert o._owns_worker(mine), "the next claim takes over"
|
||||
release()
|
||||
|
||||
|
||||
def _direct_reader_calls(o, request_id):
|
||||
"""_direct_reader wired to a scripted _read_resp (o._scripted, popped in order)."""
|
||||
o._read_resp = lambda timeout = 1.0: o._scripted.pop(0) if o._scripted else None
|
||||
return o._direct_reader(request_id)
|
||||
|
|
|
|||
|
|
@ -847,3 +847,222 @@ def test_dead_waiters_stop_counting_against_the_queue_limit():
|
|||
assert queue.is_idle()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_parking_frees_the_slot_for_a_waiter():
|
||||
"""A holder waiting on a tool approval must not hold a decode slot.
|
||||
|
||||
It is not generating, and with several prompts unanswered every slot would
|
||||
be held by a run parked on a human while llama-server sits idle.
|
||||
"""
|
||||
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
first = queue.reserve(capacity = 1, config = config)
|
||||
second = queue.reserve(capacity = 1, config = config)
|
||||
first_lease = first.lease_nowait()
|
||||
assert first_lease is not None
|
||||
assert second.lease_nowait() is None
|
||||
|
||||
first_lease.park()
|
||||
assert first_lease.slot is None, "the slot went back to the pool"
|
||||
second_lease = await second.wait(0.1)
|
||||
assert second_lease is not None, "parking did not free the slot"
|
||||
|
||||
# The parked holder keeps its lease, so releasing it is still correct.
|
||||
first_lease.unpark()
|
||||
first_lease.release()
|
||||
second_lease.release()
|
||||
assert queue.snapshot().active == 0
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_unpark_without_park_is_a_no_op():
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
first = queue.reserve(capacity = 1, config = config)
|
||||
first_lease = first.lease_nowait()
|
||||
assert first_lease is not None
|
||||
first_lease.unpark()
|
||||
first_lease.unpark()
|
||||
|
||||
second = queue.reserve(capacity = 1, config = config)
|
||||
assert second.lease_nowait() is None, "capacity leaked past the limit"
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_releasing_a_parked_lease_leaves_the_queue_evictable():
|
||||
# is_idle() drives registry eviction, and a parked holder owns no slot, so
|
||||
# nothing but the parked count keeps its queue alive. A stuck count would
|
||||
# pin every dead queue for the life of the process.
|
||||
async def _run():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
lease = queue.reserve(capacity = 1, config = config).lease_nowait()
|
||||
lease.park()
|
||||
assert not queue.is_idle(), "a parked holder is coming back to this queue"
|
||||
lease.release()
|
||||
assert queue.is_idle()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_unpark_waits_instead_of_putting_two_holders_on_one_slot():
|
||||
# park() hands the freed slot to a waiter, so by the time the user answers an approval
|
||||
# prompt someone else may be decoding in it. Resuming regardless left two holders
|
||||
# against capacity 1, and the resumed tool loop went past the admission limit.
|
||||
async def scenario():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
a = queue.reserve(capacity = 1, config = config)
|
||||
a_lease = a.lease_nowait()
|
||||
assert a_lease is not None, "A takes the only slot"
|
||||
b = queue.reserve(capacity = 1, config = config)
|
||||
assert b.lease_nowait() is None, "B waits behind A"
|
||||
|
||||
a_lease.park() # A parks on an approval prompt; its slot goes to B
|
||||
b_lease = await asyncio.wait_for(b.wait(timeout_s = 1), timeout = 2)
|
||||
assert b_lease is not None, "B was granted the parked slot"
|
||||
|
||||
# A answers the prompt while B is still decoding: it must WAIT.
|
||||
resumed = asyncio.ensure_future(a_lease.unpark_async(poll_s = 0.01))
|
||||
await asyncio.sleep(0.05)
|
||||
assert not resumed.done(), "A must not resume while B holds the slot"
|
||||
assert queue.snapshot().active <= 1, "never over capacity while waiting"
|
||||
|
||||
b_lease.release()
|
||||
await asyncio.wait_for(resumed, timeout = 2)
|
||||
assert a_lease.slot is not None, "A took a real slot back"
|
||||
assert queue.snapshot().active <= 1, "still within capacity after resuming"
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_unpark_gives_up_when_the_caller_is_cancelled():
|
||||
# A holder being torn down must not sit in the wait loop.
|
||||
async def scenario():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
a = queue.reserve(capacity = 1, config = config)
|
||||
a_lease = a.lease_nowait()
|
||||
assert a_lease is not None
|
||||
b = queue.reserve(capacity = 1, config = config)
|
||||
a_lease.park()
|
||||
assert await asyncio.wait_for(b.wait(timeout_s = 1), timeout = 2) is not None
|
||||
|
||||
ev = threading.Event()
|
||||
waiting = asyncio.ensure_future(a_lease.unpark_async(cancel_event = ev, poll_s = 0.01))
|
||||
await asyncio.sleep(0.03)
|
||||
assert not waiting.done()
|
||||
ev.set()
|
||||
await asyncio.wait_for(waiting, timeout = 2)
|
||||
assert a_lease.slot is None, "gave up without a slot rather than over-admitting"
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_an_approved_chat_is_not_overtaken_by_later_arrivals():
|
||||
# A parks on an approval prompt, B takes the slot, C arrives afterwards. release() grants
|
||||
# under the same lock, so a plain poll in unpark_async never saw a free slot: A waited
|
||||
# behind every later arrival and starved.
|
||||
async def scenario():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
a = queue.reserve(capacity = 1, config = config)
|
||||
a_lease = a.lease_nowait()
|
||||
assert a_lease is not None
|
||||
b = queue.reserve(capacity = 1, config = config)
|
||||
a_lease.park() # A's slot goes to B
|
||||
b_lease = await asyncio.wait_for(b.wait(timeout_s = 1), timeout = 2)
|
||||
assert b_lease is not None
|
||||
|
||||
# A is approved and starts waiting; C arrives only after that.
|
||||
resumed = asyncio.ensure_future(a_lease.unpark_async(poll_s = 0.01))
|
||||
await asyncio.sleep(0.03)
|
||||
c = queue.reserve(capacity = 1, config = config)
|
||||
assert c.lease_nowait() is None
|
||||
|
||||
b_lease.release() # the slot frees exactly once
|
||||
await asyncio.wait_for(resumed, timeout = 2)
|
||||
# A resumed; C is still queued behind it rather than having overtaken it.
|
||||
assert c.lease_nowait() is None
|
||||
assert queue.snapshot().active <= 1
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_two_approved_chats_do_not_block_each_other():
|
||||
# A bare pending-count made every approved holder count against every other: park A, admit
|
||||
# and park B, admit C, approve both, and once C released the predicate stayed false forever.
|
||||
async def scenario():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
a = queue.reserve(capacity = 1, config = config)
|
||||
a_lease = a.lease_nowait()
|
||||
assert a_lease is not None
|
||||
b = queue.reserve(capacity = 1, config = config)
|
||||
a_lease.park() # A parks; B is admitted
|
||||
b_lease = await asyncio.wait_for(b.wait(timeout_s = 1), timeout = 2)
|
||||
assert b_lease is not None
|
||||
|
||||
c = queue.reserve(capacity = 1, config = config)
|
||||
b_lease.park() # B parks too; C is admitted
|
||||
c_lease = await asyncio.wait_for(c.wait(timeout_s = 1), timeout = 2)
|
||||
assert c_lease is not None
|
||||
|
||||
# Both approvals come back while C is still decoding.
|
||||
first = asyncio.ensure_future(a_lease.unpark_async(poll_s = 0.01))
|
||||
await asyncio.sleep(0.02)
|
||||
second = asyncio.ensure_future(b_lease.unpark_async(poll_s = 0.01))
|
||||
await asyncio.sleep(0.02)
|
||||
assert not first.done() and not second.done()
|
||||
|
||||
c_lease.release()
|
||||
# The earlier approval goes first; the other follows once it releases.
|
||||
await asyncio.wait_for(first, timeout = 2)
|
||||
assert not second.done(), "the second approval waits its turn, not forever"
|
||||
a_lease.release()
|
||||
await asyncio.wait_for(second, timeout = 2)
|
||||
assert queue.snapshot().active <= 1
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_an_immediate_arrival_cannot_take_an_approved_chats_slot():
|
||||
# The fairness reservation lived only in _grant_waiters_locked. reserve()'s fast path
|
||||
# ignored it, so a request arriving in the window between the slot freeing and the
|
||||
# approved chat's next poll took the slot straight off the top.
|
||||
async def scenario():
|
||||
queue = get_llama_admission_queue("http://llama.test")
|
||||
config = LlamaAdmissionConfig()
|
||||
|
||||
a = queue.reserve(capacity = 1, config = config)
|
||||
a_lease = a.lease_nowait()
|
||||
assert a_lease is not None
|
||||
a_lease.park() # A is on an approval prompt; its slot is up for grabs
|
||||
b = queue.reserve(capacity = 1, config = config)
|
||||
b_lease = b.lease_nowait()
|
||||
assert b_lease is not None
|
||||
|
||||
resumed = asyncio.ensure_future(a_lease.unpark_async(poll_s = 0.01))
|
||||
await asyncio.sleep(0.03) # A is approved and now holds a ticket
|
||||
|
||||
# No await between these two: C arrives before A's poll can run again.
|
||||
b_lease.release()
|
||||
c = queue.reserve(capacity = 1, config = config)
|
||||
assert c.lease_nowait() is None, "the freed slot is reserved for the approved chat"
|
||||
|
||||
await asyncio.wait_for(resumed, timeout = 2)
|
||||
assert queue.snapshot().active <= 1
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
|
|
|||
|
|
@ -2076,6 +2076,50 @@ def test_large_python_tool_call_emits_early_provisional_start(monkeypatch):
|
|||
assert any(e.get("type") == "tool_end" and e.get("tool_name") == "python" for e in events)
|
||||
|
||||
|
||||
def test_gated_python_call_still_streams_its_arguments(monkeypatch):
|
||||
"""A call awaiting approval still streams its code into the card.
|
||||
|
||||
Suppressing it left the chat completely blank for as long as the model took
|
||||
to write the payload, which for a large file is minutes. Nothing runs before
|
||||
the decision either way, and the code is what the user is approving.
|
||||
"""
|
||||
|
||||
big_code = "total = 0\n" + "\n".join(f"total += {i}" for i in range(120))
|
||||
assert len(json.dumps({"code": big_code})) > _PROVISIONAL_ARGS_MIN_CHARS
|
||||
|
||||
first_stream = _streamed_structured_tool_call("python", {"code": big_code}, "call_gated")
|
||||
final_stream = [_sse({"content": "Done."}), _done()]
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
|
||||
|
||||
monkeypatch.setattr("core.inference.tools.execute_tool", lambda name, arguments, **_k: "OK")
|
||||
monkeypatch.setattr("core.inference.llama_cpp.wait_tool_decision", lambda *_a, **_k: "allow")
|
||||
|
||||
events = list(
|
||||
backend.generate_chat_completion_with_tools(
|
||||
messages = [{"role": "user", "content": "write code"}],
|
||||
tools = [{"type": "function", "function": {"name": "python"}}],
|
||||
confirm_tool_calls = True,
|
||||
permission_mode = "ask",
|
||||
max_tool_iterations = 1,
|
||||
)
|
||||
)
|
||||
|
||||
tool_starts = [e for e in events if e.get("type") == "tool_start"]
|
||||
provisional = [e for e in tool_starts if not e.get("arguments")]
|
||||
assert len(provisional) == 1, tool_starts
|
||||
assert provisional[0]["tool_call_id"] == "call_gated"
|
||||
|
||||
args_events = [e for e in events if e.get("type") == "tool_args"]
|
||||
assert args_events, "gated call streamed no arguments"
|
||||
assert "total += 119" in "".join(e["text"] for e in args_events)
|
||||
|
||||
# The approval prompt still fires, and it comes after the code is on screen.
|
||||
gated = [e for e in tool_starts if e.get("awaiting_confirmation")]
|
||||
assert gated, tool_starts
|
||||
assert events.index(provisional[0]) < events.index(gated[0])
|
||||
|
||||
|
||||
def test_auto_mode_render_html_suppresses_provisional_card_under_confirm(monkeypatch):
|
||||
"""render_html is no longer unconditionally safe (a networked canvas asks), so
|
||||
with confirm_tool_calls set under permission_mode="auto" its early provisional
|
||||
|
|
|
|||
|
|
@ -1606,22 +1606,28 @@ def test_load_route_holds_lifecycle_gate(monkeypatch):
|
|||
|
||||
|
||||
def test_model_replacements_recheck_sidecar_swap_before_either_backend_is_unloaded():
|
||||
# Both replacement directions drain active inference, then recheck whether a
|
||||
# sidecar install reserved the lifecycle gate during that wait. Exact-model
|
||||
# reuse exits earlier, so an already-loaded model never waits on unrelated inference.
|
||||
# Both replacement directions drain, then recheck whether a sidecar install reserved the
|
||||
# gate meanwhile. That recheck is the last thing that can reject the load, so the
|
||||
# destructive cancel must follow it. Exact-model reuse exits earlier and never waits.
|
||||
import inspect
|
||||
|
||||
src = inspect.getsource(inference_route._load_model_impl)
|
||||
already_loaded = src.index('status = "already_loaded"')
|
||||
standard_branch = src.index("# ── Standard path")
|
||||
|
||||
gguf_wait = src.index("await _wait_for_model_switch_idle", src.index("if config.is_gguf:"))
|
||||
gguf_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", gguf_wait)
|
||||
gguf_cancel = src.index("on_reload_confirmed(cancel = True)", gguf_wait)
|
||||
unload_unsloth = src.index("unsloth_backend.unload_model", gguf_wait)
|
||||
standard_wait = src.index("await _wait_for_model_switch_idle", gguf_wait + 1)
|
||||
standard_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", standard_wait)
|
||||
unload_gguf = src.index("llama_backend.unload_model()", standard_wait)
|
||||
already_loaded = src.index('status = "already_loaded"')
|
||||
|
||||
assert already_loaded < gguf_wait < gguf_sidecar_check < unload_unsloth
|
||||
assert standard_wait < standard_sidecar_check < unload_gguf
|
||||
standard_wait = src.index("await _wait_for_model_switch_idle", standard_branch)
|
||||
standard_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", standard_wait)
|
||||
standard_cancel = src.index("on_reload_confirmed(cancel = True)", standard_wait)
|
||||
unload_gguf = src.index("llama_backend.unload_model()", standard_wait)
|
||||
|
||||
assert already_loaded < gguf_wait < gguf_sidecar_check < gguf_cancel < unload_unsloth
|
||||
assert standard_branch < standard_wait < standard_sidecar_check
|
||||
assert standard_sidecar_check < standard_cancel < unload_gguf
|
||||
|
||||
|
||||
def test_switch_waiter_deregisters_before_swap_gate_release():
|
||||
|
|
|
|||
|
|
@ -4615,6 +4615,9 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
async def json(self):
|
||||
return {"prompt": "hi", "stream": False}
|
||||
|
||||
async def is_disconnected(self):
|
||||
return False
|
||||
|
||||
class FailingAsyncClient:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
|
@ -4622,14 +4625,18 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
async def __aexit__(self, *_args):
|
||||
return False
|
||||
|
||||
async def aclose(self):
|
||||
return None
|
||||
|
||||
async def post(self, *_args, **_kwargs):
|
||||
raise httpx.ConnectError("llama down")
|
||||
|
||||
monitor = ApiMonitor(max_entries = 3)
|
||||
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
|
||||
# Per-request client so a forced swap can close it mid-call; the pooled one is shared.
|
||||
monkeypatch.setattr(
|
||||
inf_mod,
|
||||
"nonstreaming_client",
|
||||
"_cancelable_nonstreaming_client",
|
||||
lambda: FailingAsyncClient(),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
|
|
@ -4667,9 +4674,15 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
async def json(self):
|
||||
return {"prompt": "hi", "stream": False}
|
||||
|
||||
async def is_disconnected(self):
|
||||
return False
|
||||
|
||||
captured = []
|
||||
|
||||
class CapturingClient:
|
||||
async def aclose(self):
|
||||
return None
|
||||
|
||||
async def post(self, _url, *, json, **_kwargs):
|
||||
captured.append(dict(json))
|
||||
return httpx.Response(
|
||||
|
|
@ -4687,7 +4700,9 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
|
||||
monitor = ApiMonitor(max_entries = 3)
|
||||
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
|
||||
monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: CapturingClient())
|
||||
monkeypatch.setattr(
|
||||
inf_mod, "_cancelable_nonstreaming_client", lambda: CapturingClient()
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
inf_mod,
|
||||
"get_llama_cpp_backend",
|
||||
|
|
@ -4718,9 +4733,15 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
async def json(self):
|
||||
return {"prompt": "hi", "stream": False, "max_tokens": 0}
|
||||
|
||||
async def is_disconnected(self):
|
||||
return False
|
||||
|
||||
captured = []
|
||||
|
||||
class CapturingClient:
|
||||
async def aclose(self):
|
||||
return None
|
||||
|
||||
async def post(self, _url, *, json, **_kwargs):
|
||||
captured.append(dict(json))
|
||||
return httpx.Response(
|
||||
|
|
@ -4738,7 +4759,9 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
|
||||
monitor = ApiMonitor(max_entries = 3)
|
||||
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
|
||||
monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: CapturingClient())
|
||||
monkeypatch.setattr(
|
||||
inf_mod, "_cancelable_nonstreaming_client", lambda: CapturingClient()
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
inf_mod,
|
||||
"get_llama_cpp_backend",
|
||||
|
|
@ -4776,6 +4799,7 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
monitor = ApiMonitor(max_entries = 3)
|
||||
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
|
||||
monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: UnusedClient())
|
||||
monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: UnusedClient())
|
||||
monkeypatch.setattr(
|
||||
inf_mod,
|
||||
"get_llama_cpp_backend",
|
||||
|
|
@ -4880,6 +4904,9 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
async def json(self):
|
||||
return {"input": ["alpha", "beta"], "model": "embed"}
|
||||
|
||||
async def is_disconnected(self):
|
||||
return False
|
||||
|
||||
class FakeAsyncClient:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
|
@ -4887,6 +4914,9 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
async def __aexit__(self, *_args):
|
||||
return False
|
||||
|
||||
async def aclose(self):
|
||||
return None
|
||||
|
||||
async def post(self, *_args, **_kwargs):
|
||||
assert monitor.active_count() == 1
|
||||
return httpx.Response(
|
||||
|
|
@ -4899,9 +4929,10 @@ class TestApiMonitorProviderAndCompletionStreams:
|
|||
|
||||
monitor = ApiMonitor(max_entries = 3)
|
||||
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
|
||||
# Per-request client so a forced swap can close it mid-call; the pooled one is shared.
|
||||
monkeypatch.setattr(
|
||||
inf_mod,
|
||||
"nonstreaming_client",
|
||||
"_cancelable_nonstreaming_client",
|
||||
lambda: FakeAsyncClient(),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
|
|
@ -6372,7 +6403,7 @@ class TestApiMonitorSafetensorsUsage:
|
|||
}
|
||||
yield "safe reply"
|
||||
|
||||
def reset_generation_state(self):
|
||||
def reset_generation_state(self, caller_cancel_event = None):
|
||||
pass
|
||||
|
||||
monitor = ApiMonitor(max_entries = 3)
|
||||
|
|
@ -6443,7 +6474,7 @@ class TestApiMonitorSafetensorsUsage:
|
|||
cancel_event.set()
|
||||
yield {"type": "content", "text": "ignored"}
|
||||
|
||||
def reset_generation_state(self):
|
||||
def reset_generation_state(self, caller_cancel_event = None):
|
||||
pass
|
||||
|
||||
monitor = ApiMonitor(max_entries = 3)
|
||||
|
|
@ -6504,7 +6535,7 @@ class TestApiMonitorSafetensorsUsage:
|
|||
def generate_chat_completion_with_tools(self, **_kwargs):
|
||||
yield {"type": "content", "text": "unused"}
|
||||
|
||||
def reset_generation_state(self):
|
||||
def reset_generation_state(self, caller_cancel_event = None):
|
||||
nonlocal reset_called
|
||||
reset_called = True
|
||||
|
||||
|
|
|
|||
|
|
@ -19,6 +19,10 @@ def _bare_orchestrator():
|
|||
"""An orchestrator without the real __init__ subprocess/network."""
|
||||
o = InferenceOrchestrator.__new__(InferenceOrchestrator)
|
||||
o._gen_lock = threading.Lock()
|
||||
o._send_order_lock = threading.Lock()
|
||||
o._active_cancel_lock = threading.Lock()
|
||||
o._active_cancel_events = []
|
||||
o._executing_cancel_events = []
|
||||
o._cancel_event = threading.Event() # stands in for the mp.Event
|
||||
o._drain_event = threading.Event() # stands in for the unload-drain mp.Event
|
||||
o._proc = object() # truthy so _ensure_subprocess_alive reports alive
|
||||
|
|
@ -775,6 +779,7 @@ def test_dispatched_bails_when_unload_flips_before_mailbox_registration(monkeypa
|
|||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._unload_pending = False
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
monkeypatch.setattr(o, "_start_dispatcher", lambda: None)
|
||||
|
|
@ -817,6 +822,7 @@ def test_dispatched_bails_when_model_swapped_before_mailbox_registration(monkeyp
|
|||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._unload_pending = False
|
||||
o._dispatcher_thread = _AliveDispatcher()
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
|
|
@ -846,6 +852,7 @@ def test_dispatched_bails_when_dispatcher_stopped_before_mailbox_registration(mo
|
|||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._unload_pending = False
|
||||
o._dispatcher_thread = _AliveDispatcher()
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
|
|
@ -872,6 +879,7 @@ def test_dispatched_happy_path_registers_and_sends(monkeypatch):
|
|||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._unload_pending = False
|
||||
o._dispatcher_thread = _AliveDispatcher()
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
|
|
@ -1338,6 +1346,7 @@ def test_dispatched_bail_stops_orphan_dispatcher_it_started(monkeypatch):
|
|||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._unload_pending = False
|
||||
o._dispatcher_thread = None # none running -> this call starts it
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
|
|
@ -1382,6 +1391,7 @@ def test_dispatched_bail_keeps_dispatcher_with_other_active_mailbox(monkeypatch)
|
|||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._unload_pending = False
|
||||
o._dispatcher_thread = None
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
|
|
@ -1419,6 +1429,7 @@ def test_dispatched_bail_keeps_preexisting_dispatcher(monkeypatch):
|
|||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._unload_pending = False
|
||||
o._dispatcher_thread = _AliveDispatcher() # already running
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
|
|
@ -1545,6 +1556,7 @@ def test_concurrent_start_dispatcher_spawns_exactly_one():
|
|||
o._resp_queue = _queue.Queue() # real queue so the dispatcher loop blocks and stays alive
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._dispatcher_thread = None
|
||||
o._dispatcher_stop = threading.Event()
|
||||
o._dispatcher_lifecycle_lock = threading.Lock()
|
||||
|
|
@ -1660,6 +1672,7 @@ def test_queued_start_behind_unload_stop_spawns_no_dispatcher():
|
|||
o._resp_queue = _queue.Queue() # a spawned dispatcher would block-read here and stay alive
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._dispatcher_stop = threading.Event()
|
||||
o._dispatcher_lifecycle_lock = threading.Lock()
|
||||
o._unload_pending = False
|
||||
|
|
@ -1713,3 +1726,310 @@ def test_queued_start_behind_unload_stop_spawns_no_dispatcher():
|
|||
assert o._dispatcher_thread is None, "the stop cleared it and the queued start spawned nothing"
|
||||
live = [t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive()]
|
||||
assert live == [], "no fresh dispatcher may be left to consume the unloaded reply"
|
||||
|
||||
|
||||
def _dispatch(o, resps):
|
||||
"""Run the dispatcher over a fixed response list and stop it."""
|
||||
import queue as _queue
|
||||
|
||||
o._resp_queue = _queue.Queue()
|
||||
for r in resps:
|
||||
o._resp_queue.put(r)
|
||||
o._dispatcher_stop = threading.Event()
|
||||
t = threading.Thread(target = o._dispatcher_loop, daemon = True)
|
||||
t.start()
|
||||
deadline = time.monotonic() + 5.0
|
||||
while not o._resp_queue.empty() and time.monotonic() < deadline:
|
||||
time.sleep(0.01)
|
||||
o._dispatcher_stop.set()
|
||||
t.join(timeout = 5.0)
|
||||
|
||||
|
||||
def test_worker_ownership_follows_the_worker_not_the_consumer():
|
||||
# The subprocess runs one generation at a time and can start B while A's consumer has yet to
|
||||
# drain its mailbox. A must stop owning the worker the moment its gen_done is routed, else
|
||||
# a late Stop for A cancels B.
|
||||
import queue as _queue
|
||||
|
||||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
a_cancel, b_cancel = threading.Event(), threading.Event()
|
||||
o._mailboxes = {"a": _queue.Queue(), "b": _queue.Queue()}
|
||||
o._request_cancel_events = {"a": a_cancel, "b": b_cancel}
|
||||
o._claim_worker(a_cancel)
|
||||
o._claim_worker(b_cancel)
|
||||
|
||||
_dispatch(o, [{"type": "token", "request_id": "a", "token": "hi"}])
|
||||
assert o._owns_worker(a_cancel), "the request the worker is answering owns it"
|
||||
assert not o._owns_worker(b_cancel), "a queued request does not"
|
||||
|
||||
# A finishes. B has been sent but has not answered yet (it is prefilling), so the gap
|
||||
# between the two is the window a late Stop for A used to fire into.
|
||||
_dispatch(o, [{"type": "gen_done", "request_id": "a"}])
|
||||
assert not o._owns_worker(a_cancel), "a finished request stops owning the worker"
|
||||
assert o._owns_worker(b_cancel), "the next queued request is the one prefilling"
|
||||
|
||||
# Worker moves on to B, still before A's consumer reads anything.
|
||||
_dispatch(o, [{"type": "token", "request_id": "b", "token": "yo"}])
|
||||
assert not o._owns_worker(a_cancel), "a finished request must not cancel its successor"
|
||||
assert o._owns_worker(b_cancel), "the worker moved on to B, so B owns it"
|
||||
|
||||
# A's own stream unwinding afterwards must not disturb B.
|
||||
o._release_worker(a_cancel)
|
||||
assert o._owns_worker(b_cancel)
|
||||
|
||||
|
||||
def test_status_responses_do_not_transfer_worker_ownership():
|
||||
# Status lines are not an answer to any request; the dispatcher drops them before routing.
|
||||
import queue as _queue
|
||||
|
||||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
a_cancel, b_cancel = threading.Event(), threading.Event()
|
||||
o._mailboxes = {"a": _queue.Queue(), "b": _queue.Queue()}
|
||||
o._request_cancel_events = {"a": a_cancel, "b": b_cancel}
|
||||
o._claim_worker(a_cancel)
|
||||
o._claim_worker(b_cancel)
|
||||
|
||||
_dispatch(o, [{"type": "status", "request_id": "b", "message": "loading"}])
|
||||
# Nothing has answered, so the oldest claim is still the one prefilling.
|
||||
assert o._owns_worker(a_cancel)
|
||||
assert not o._owns_worker(b_cancel)
|
||||
|
||||
|
||||
def test_only_the_latest_responder_executes():
|
||||
# The subprocess runs one generation at a time, so answering B means it has left A.
|
||||
# _generate_inner promotes from its own consumer and can share the worker with a
|
||||
# dispatched request, so the two must not both count as executing.
|
||||
o = _bare_orchestrator()
|
||||
a_cancel, b_cancel = threading.Event(), threading.Event()
|
||||
o._claim_worker(a_cancel)
|
||||
o._claim_worker(b_cancel)
|
||||
|
||||
o._mark_worker_started(a_cancel)
|
||||
assert o._owns_worker(a_cancel)
|
||||
o._mark_worker_started(b_cancel)
|
||||
assert o._owns_worker(b_cancel), "the latest responder is the one executing"
|
||||
assert not o._owns_worker(a_cancel), "and it is the only one"
|
||||
# Idempotent: more of B's own tokens must not disturb it.
|
||||
o._mark_worker_started(b_cancel)
|
||||
assert o._owns_worker(b_cancel)
|
||||
|
||||
|
||||
def test_a_stale_mailbox_read_does_not_cancel_the_running_generation():
|
||||
# A dispatched consumer can still be draining tokens after the dispatcher retired its request
|
||||
# and started the next one. Stopping it then must tear down only its own stream: signalling
|
||||
# the shared worker event would end its successor.
|
||||
import queue as _queue
|
||||
|
||||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
a_cancel, b_cancel = threading.Event(), threading.Event()
|
||||
o._mailboxes = {"a": _queue.Queue(), "b": _queue.Queue()}
|
||||
o._request_cancel_events = {"a": a_cancel, "b": b_cancel}
|
||||
o._claim_worker(a_cancel)
|
||||
o._claim_worker(b_cancel)
|
||||
# Worker finished A and moved on to B.
|
||||
_dispatch(
|
||||
o,
|
||||
[
|
||||
{"type": "gen_done", "request_id": "a"},
|
||||
{"type": "token", "request_id": "b", "token": "yo"},
|
||||
],
|
||||
)
|
||||
assert o._owns_worker(b_cancel) and not o._owns_worker(a_cancel)
|
||||
|
||||
# A's consumer now reads a token buffered before that, with A stopped.
|
||||
a_cancel.set()
|
||||
stale = [{"type": "token", "request_id": "a", "text": "late"}]
|
||||
drained = []
|
||||
list(
|
||||
o._consume_token_stream(
|
||||
lambda timeout: stale.pop(0) if stale else None,
|
||||
lambda: drained.append(True),
|
||||
crash_context = "generation",
|
||||
cancel_event = a_cancel,
|
||||
mark_started = False,
|
||||
)
|
||||
)
|
||||
assert drained, "the stopped stream still tears itself down"
|
||||
assert not o._cancel_event.is_set(), "a retired request must not signal the shared worker event"
|
||||
|
||||
# The generation that does own the worker still can.
|
||||
b_cancel.set()
|
||||
stale_b = [{"type": "token", "request_id": "b", "text": "live"}]
|
||||
list(
|
||||
o._consume_token_stream(
|
||||
lambda timeout: stale_b.pop(0) if stale_b else None,
|
||||
lambda: None,
|
||||
crash_context = "generation",
|
||||
cancel_event = b_cancel,
|
||||
mark_started = False,
|
||||
)
|
||||
)
|
||||
assert o._cancel_event.is_set(), "the running generation's own Stop must reach the worker"
|
||||
|
||||
|
||||
def test_a_dispatcher_started_mid_stream_still_reaches_the_direct_reader():
|
||||
# A compare request can start the dispatcher while an ordinary chat is streaming. The
|
||||
# dispatcher then owns resp_queue, and without a mailbox for the direct reader it dropped
|
||||
# that chat's tokens and its gen_done as unaddressed, hanging it.
|
||||
import queue as _queue
|
||||
|
||||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._direct_mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
|
||||
read_one, _drain, release = o._direct_reader("direct-1")
|
||||
try:
|
||||
_dispatch(
|
||||
o,
|
||||
[
|
||||
{"type": "token", "request_id": "direct-1", "text": "hi"},
|
||||
{"type": "gen_done", "request_id": "direct-1"},
|
||||
],
|
||||
)
|
||||
assert read_one(timeout = 0.1) == {
|
||||
"type": "token",
|
||||
"request_id": "direct-1",
|
||||
"text": "hi",
|
||||
}, "the dispatcher must route to the direct reader, not drop"
|
||||
assert read_one(timeout = 0.1)["type"] == "gen_done"
|
||||
finally:
|
||||
release()
|
||||
assert o._direct_mailboxes == {}, "the mailbox is dropped when the stream ends"
|
||||
|
||||
|
||||
def test_the_direct_reader_hands_back_a_compare_response_it_took():
|
||||
# The mirror race: this reader is already blocked on resp_queue when a compare request's
|
||||
# dispatcher starts, so it can take that request's response first. Consuming it would
|
||||
# corrupt this chat and hang the compare pane.
|
||||
import queue as _queue
|
||||
|
||||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
compare_box: _queue.Queue = _queue.Queue()
|
||||
o._mailboxes = {"compare-1": compare_box}
|
||||
o._direct_mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
o._resp_queue = _queue.Queue()
|
||||
o._dispatcher_thread = None # no dispatcher yet: this reader owns the queue
|
||||
|
||||
read_one, _drain, release = o._direct_reader("direct-1")
|
||||
try:
|
||||
o._resp_queue.put({"type": "token", "request_id": "compare-1", "text": "theirs"})
|
||||
o._resp_queue.put({"type": "token", "request_id": "direct-1", "text": "mine"})
|
||||
assert read_one(timeout = 0.1) is None, "a foreign response is not ours to yield"
|
||||
assert compare_box.get_nowait()["text"] == "theirs", "it goes to its own mailbox"
|
||||
assert read_one(timeout = 0.1)["text"] == "mine"
|
||||
finally:
|
||||
release()
|
||||
|
||||
|
||||
def test_a_direct_mailbox_is_not_mistaken_for_compare_activity():
|
||||
# _mailboxes means "compare requests are in flight" to the unload and distributed paths,
|
||||
# so an ordinary chat's mailbox must live somewhere else.
|
||||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._direct_mailboxes = {}
|
||||
_read_one, _drain, release = o._direct_reader("direct-1")
|
||||
try:
|
||||
assert o._mailboxes == {}
|
||||
assert "direct-1" in o._direct_mailboxes
|
||||
finally:
|
||||
release()
|
||||
|
||||
|
||||
def test_replacing_the_subprocess_clears_worker_scoped_state():
|
||||
# Ownership is keyed only by cancel-event identity, so a consumer still blocked on its
|
||||
# mailbox when the worker was replaced stayed recorded as the executor. A generation on
|
||||
# the fresh worker then failed _owns_worker and could not be stopped.
|
||||
import queue as _queue
|
||||
|
||||
o = _bare_orchestrator()
|
||||
o._mailbox_lock = threading.Lock()
|
||||
dead = threading.Event()
|
||||
o._mailboxes = {"compare-1": _queue.Queue()}
|
||||
o._direct_mailboxes = {"direct-1": _queue.Queue()}
|
||||
o._request_cancel_events = {"compare-1": dead}
|
||||
o._claim_worker(dead)
|
||||
o._mark_worker_started(dead)
|
||||
assert o._owns_worker(dead)
|
||||
|
||||
o._reset_worker_scoped_state()
|
||||
|
||||
assert o._mailboxes == {} and o._direct_mailboxes == {}
|
||||
assert o._request_cancel_events == {}
|
||||
assert o._active_cancel_events == [] and o._executing_cancel_events == []
|
||||
# A generation on the fresh worker owns it rather than being refused by a ghost.
|
||||
fresh = threading.Event()
|
||||
o._claim_worker(fresh)
|
||||
assert o._owns_worker(fresh), "the dead worker's request must not outrank a live one"
|
||||
|
||||
|
||||
def test_audio_input_claims_the_worker_before_sending():
|
||||
# Unclaimed, a compare request queued behind an audio-input generation looked like the
|
||||
# oldest owner, so stopping that queued request signalled the worker and killed this.
|
||||
import ast
|
||||
import pathlib
|
||||
|
||||
src = pathlib.Path(orch_mod.__file__).read_text(encoding = "utf-8")
|
||||
tree = ast.parse(src)
|
||||
fn = next(
|
||||
n
|
||||
for n in ast.walk(tree)
|
||||
if isinstance(n, ast.FunctionDef) and n.name == "_generate_audio_input_inner"
|
||||
)
|
||||
body = ast.get_source_segment(src, fn) or ""
|
||||
claim = body.find("self._claim_worker(cancel_event)")
|
||||
send = body.find("self._send_cmd(cmd)")
|
||||
assert claim != -1, "_generate_audio_input_inner must claim the worker"
|
||||
assert send != -1
|
||||
assert claim < send, "the claim has to happen before the command is enqueued"
|
||||
assert "with self._send_order_lock:" in body, "claim and send must be one critical section"
|
||||
assert "self._release_worker(cancel_event)" in body
|
||||
|
||||
|
||||
def test_generation_stopped_while_queued_is_never_sent(monkeypatch):
|
||||
# Two chats on the serialized backend: the second blocks on _gen_lock, and Stop sets its
|
||||
# event while it waits. Sending anyway occupied the worker with a run the user ended --
|
||||
# the cancel is only checked on a token, so a long prefill (or a generation that reaches
|
||||
# gen_done without one) still held up its siblings.
|
||||
o = _bare_orchestrator()
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
monkeypatch.setattr(o, "_wait_dispatcher_idle", lambda *a, **k: None)
|
||||
monkeypatch.setattr(
|
||||
o, "_send_cmd", lambda cmd: pytest.fail("must not send a generation already stopped")
|
||||
)
|
||||
stopped = threading.Event()
|
||||
stopped.set()
|
||||
|
||||
out = list(
|
||||
o._generate_inner(messages = [{"role": "user", "content": "hi"}], cancel_event = stopped)
|
||||
)
|
||||
|
||||
assert out == [], "a stopped request yields nothing rather than an error banner"
|
||||
assert o._active_cancel_events == [], "it must not claim the worker either"
|
||||
assert o._gen_lock.acquire(blocking = False)
|
||||
o._gen_lock.release()
|
||||
|
||||
|
||||
def test_audio_input_stopped_while_queued_is_never_sent(monkeypatch):
|
||||
# Same lock, same hole.
|
||||
o = _bare_orchestrator()
|
||||
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
|
||||
monkeypatch.setattr(
|
||||
o, "_send_cmd", lambda cmd: pytest.fail("must not send a generation already stopped")
|
||||
)
|
||||
stopped = threading.Event()
|
||||
stopped.set()
|
||||
|
||||
out = list(o._generate_audio_input_inner(audio_array = [0.0, 0.1], cancel_event = stopped))
|
||||
|
||||
assert out == []
|
||||
assert o._active_cancel_events == []
|
||||
assert o._gen_lock.acquire(blocking = False)
|
||||
o._gen_lock.release()
|
||||
|
|
|
|||
|
|
@ -504,11 +504,12 @@ def _upstream_message(
|
|||
|
||||
|
||||
class ScriptedClient:
|
||||
"""Fake nonstreaming_client() returning scripted JSON bodies, counting POSTs."""
|
||||
"""Fake upstream client returning scripted JSON bodies, counting POSTs."""
|
||||
|
||||
def __init__(self, bodies):
|
||||
self.bodies = list(bodies)
|
||||
self.posts = []
|
||||
self.closed = False
|
||||
|
||||
async def post(
|
||||
self,
|
||||
|
|
@ -520,6 +521,10 @@ class ScriptedClient:
|
|||
self.posts.append(json)
|
||||
return httpx.Response(200, json = self.bodies[min(len(self.posts) - 1, len(self.bodies) - 1)])
|
||||
|
||||
async def aclose(self):
|
||||
# The Anthropic pass-through owns its client and closes it in a finally.
|
||||
self.closed = True
|
||||
|
||||
|
||||
async def _drive_non_streaming(monkeypatch, payload, bodies):
|
||||
import routes.inference as inf_mod
|
||||
|
|
@ -867,7 +872,7 @@ class TestNudgeRetryAnthropic:
|
|||
from routes.inference import _anthropic_passthrough_non_streaming
|
||||
|
||||
client = ScriptedClient(bodies)
|
||||
monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client)
|
||||
monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: client)
|
||||
response = await _anthropic_passthrough_non_streaming(
|
||||
_llama_backend(),
|
||||
[{"role": "user", "content": "hi"}],
|
||||
|
|
@ -925,7 +930,7 @@ class TestAnthropicPassthroughHealingText:
|
|||
from routes.inference import _anthropic_passthrough_non_streaming
|
||||
|
||||
client = ScriptedClient([upstream])
|
||||
monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client)
|
||||
monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: client)
|
||||
response = await _anthropic_passthrough_non_streaming(
|
||||
_llama_backend(),
|
||||
[{"role": "user", "content": "hi"}],
|
||||
|
|
@ -1171,7 +1176,7 @@ class TestAnthropicNonStreamingRoute:
|
|||
from routes.inference import _anthropic_passthrough_non_streaming
|
||||
|
||||
client = ScriptedClient(bodies)
|
||||
monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client)
|
||||
monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: client)
|
||||
response = await _anthropic_passthrough_non_streaming(
|
||||
_llama_backend(),
|
||||
[{"role": "user", "content": "hi"}],
|
||||
|
|
|
|||
|
|
@ -5051,3 +5051,27 @@ class TestFalseAlarmMarkerProse:
|
|||
assert [c[0] for c in exec_fn.calls] == ["web_search", "python"]
|
||||
assistant = next(m for m in convs[1] if m["role"] == "assistant")
|
||||
assert '"python"' not in (assistant.get("content") or "")
|
||||
|
||||
|
||||
def test_both_tool_loops_say_they_are_waiting_for_approval():
|
||||
"""A gated call must not report "Running" in either loop.
|
||||
|
||||
The GGUF loop was fixed first and the safetensors one was missed, so the
|
||||
badge counted up "Running ..." against a prompt nobody had answered yet.
|
||||
Asserted on the source so the two paths cannot drift apart again.
|
||||
"""
|
||||
import ast
|
||||
import os
|
||||
|
||||
backend = os.path.join(os.path.dirname(__file__), "..")
|
||||
for name in ("core/inference/safetensors_agentic.py", "core/inference/llama_cpp.py"):
|
||||
with open(os.path.join(backend, name), encoding = "utf-8") as f:
|
||||
tree = ast.parse(f.read())
|
||||
calls = [
|
||||
node
|
||||
for node in ast.walk(tree)
|
||||
if isinstance(node, ast.Call)
|
||||
and isinstance(node.func, ast.Name)
|
||||
and node.func.id == "awaiting_approval_status"
|
||||
]
|
||||
assert calls, f"{name} still announces a gated tool call as running"
|
||||
|
|
|
|||
|
|
@ -95,7 +95,7 @@ class _ScriptedBackend:
|
|||
for snap in snapshots:
|
||||
yield snap
|
||||
|
||||
def reset_generation_state(self):
|
||||
def reset_generation_state(self, caller_cancel_event = None):
|
||||
self.reset_count += 1
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -9,6 +9,8 @@ holds sidecar transformers modules (breaking the rename on Windows). The methods
|
|||
the handle and return False so callers can refuse the swap.
|
||||
"""
|
||||
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
from core.export.orchestrator import ExportOrchestrator
|
||||
|
|
@ -52,6 +54,14 @@ def _bare_inference():
|
|||
o._resp_queue = _Q()
|
||||
o._cancel_event = None
|
||||
o._drain_event = None
|
||||
# Worker-scoped bookkeeping the teardown clears (see _reset_worker_scoped_state).
|
||||
o._active_cancel_lock = threading.Lock()
|
||||
o._active_cancel_events = []
|
||||
o._executing_cancel_events = []
|
||||
o._mailbox_lock = threading.Lock()
|
||||
o._mailboxes = {}
|
||||
o._direct_mailboxes = {}
|
||||
o._request_cancel_events = {}
|
||||
return o
|
||||
|
||||
|
||||
|
|
|
|||
80
studio/backend/tests/test_tool_sandbox_per_thread.py
Normal file
80
studio/backend/tests/test_tool_sandbox_per_thread.py
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Every conversation runs its tools in its own sandbox directory.
|
||||
|
||||
Parallel chats lean on this: two conversations can be mid tool call at the same
|
||||
time, so a shared working directory would let one overwrite the other's files.
|
||||
The session id is the chat's thread id (or project-<id> for project chats), and
|
||||
the dir is derived from it here.
|
||||
|
||||
HOME is redirected at import time, so nothing touches the real ~/studio_sandbox.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
_backend = os.path.join(os.path.dirname(__file__), "..")
|
||||
sys.path.insert(0, _backend)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def workdir(tmp_path, monkeypatch):
|
||||
"""_get_workdir with HOME pointed at tmp_path and its cache cleared."""
|
||||
from core.inference import tools
|
||||
|
||||
monkeypatch.setattr(os.path, "expanduser", lambda path: str(tmp_path))
|
||||
monkeypatch.setattr(tools, "_workdirs", {})
|
||||
return tools._get_workdir
|
||||
|
||||
|
||||
def test_two_conversations_get_two_directories(workdir, tmp_path):
|
||||
a = workdir("thread-alpha")
|
||||
b = workdir("thread-beta")
|
||||
assert a != b
|
||||
assert os.path.basename(a) == "thread-alpha"
|
||||
assert os.path.basename(b) == "thread-beta"
|
||||
assert os.path.isdir(a) and os.path.isdir(b)
|
||||
assert os.path.dirname(a) == os.path.dirname(b) == str(tmp_path / "studio_sandbox")
|
||||
|
||||
|
||||
def test_the_same_conversation_keeps_its_directory(workdir):
|
||||
# A later turn, or a tool continuation, must land back in the same place.
|
||||
assert workdir("thread-alpha") == workdir("thread-alpha")
|
||||
|
||||
|
||||
def test_a_directory_is_private_to_its_conversation(workdir):
|
||||
a = workdir("thread-alpha")
|
||||
b = workdir("thread-beta")
|
||||
with open(os.path.join(a, "secret.txt"), "w", encoding = "utf-8") as f:
|
||||
f.write("alpha")
|
||||
assert os.listdir(b) == []
|
||||
|
||||
|
||||
def test_project_chats_deliberately_share_one_workspace(workdir, monkeypatch):
|
||||
# Chats in a project are meant to see each other's files.
|
||||
from core.inference import tools
|
||||
monkeypatch.setattr(tools, "_get_project_workdir", lambda sid: "/tmp/project-ws")
|
||||
assert tools._get_workdir("project-abc") == "/tmp/project-ws"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"session_id",
|
||||
["../escape", "a/b", "", " ", "x" * 65],
|
||||
)
|
||||
def test_a_session_id_cannot_escape_the_sandbox_root(workdir, tmp_path, session_id):
|
||||
resolved = workdir(session_id) if session_id else workdir(None)
|
||||
root = os.path.realpath(str(tmp_path / "studio_sandbox"))
|
||||
assert os.path.realpath(resolved).startswith(root + os.sep)
|
||||
assert os.path.basename(resolved) in {"_invalid", "_default"}
|
||||
|
||||
|
||||
def test_no_session_id_falls_back_to_default(workdir):
|
||||
assert os.path.basename(workdir(None)) == "_default"
|
||||
|
||||
|
||||
@pytest.mark.skipif(sys.platform == "win32", reason = "POSIX permission bits")
|
||||
def test_directories_are_private_to_the_user(workdir):
|
||||
assert os.stat(workdir("thread-alpha")).st_mode & 0o777 == 0o700
|
||||
|
|
@ -668,7 +668,8 @@ def test_route_history_and_passthrough_forward_the_display_gate():
|
|||
blocks = {
|
||||
"safetensors history": r"Strip stale tool-call XML from prior assistant turns.*?\.strip\(\)",
|
||||
"anthropic history": r"Strip stale tool-call XML via the protected display helper.*?\.strip\(\)",
|
||||
"anthropic passthrough": r"gated on the declared tools so an\n.*?\.strip\(\)",
|
||||
# Anchored on the code, not the comment above it, so rewrapping prose cannot break this.
|
||||
"anthropic passthrough": r"if not healing_active:.*?\.strip\(\)",
|
||||
}
|
||||
for label, pat in blocks.items():
|
||||
m = _re.search(pat, _src, _re.DOTALL)
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import {
|
|||
import {
|
||||
ChatPage,
|
||||
clearNewChatDraft,
|
||||
StopRunningChatsDialog,
|
||||
useChatRuntimeStore,
|
||||
type ChatSearch,
|
||||
} from "@/features/chat";
|
||||
|
|
@ -227,6 +228,8 @@ function RootLayout() {
|
|||
<HfTokenWarningDialog />
|
||||
<RemoteCodeConsentDialog />
|
||||
<TransformersUpgradeDialog />
|
||||
{/* At the root, not under /chat: a swap can start from the Hub too. */}
|
||||
<StopRunningChatsDialog />
|
||||
{hideNavbar ? (
|
||||
<main className="flex-1 pt-[var(--studio-hidden-route-top-inset,0px)] [--studio-titlebar-height:var(--studio-hidden-route-top-inset,0px)]">
|
||||
<Suspense fallback={<RouteFallback />}>
|
||||
|
|
|
|||
|
|
@ -520,15 +520,46 @@ export function AppSidebar() {
|
|||
});
|
||||
const storeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
|
||||
const setActiveThreadId = useChatRuntimeStore((s) => s.setActiveThreadId);
|
||||
const anyChatRunning = useChatRuntimeStore((s) =>
|
||||
Object.values(s.runningByThreadId).some(Boolean),
|
||||
);
|
||||
// The thread currently generating (if any), so "Return to Chat" lands on the
|
||||
// live chat rather than an empty new-chat draft left active after New Chat.
|
||||
const runningThreadId = useChatRuntimeStore((s) => {
|
||||
const entry = Object.entries(s.runningByThreadId).find(([, on]) => on);
|
||||
return entry ? entry[0] : null;
|
||||
});
|
||||
// The whole map, so each row can show its own spinner.
|
||||
const runningThreadIds = useChatRuntimeStore((s) => s.runningByThreadId);
|
||||
// Rows, not raw thread ids: a compare conversation runs two pane threads but is one chat
|
||||
// in the sidebar, so counting the map said "2 Chats" for a single compare row.
|
||||
const runningChatCount = useMemo(() => {
|
||||
const running = new Set(
|
||||
Object.entries(runningThreadIds)
|
||||
.filter(([, on]) => on)
|
||||
.map(([id]) => id),
|
||||
);
|
||||
if (running.size === 0) return 0;
|
||||
let rows = 0;
|
||||
for (const item of allChatItems) {
|
||||
const ids = item.type === "compare" ? (item.threadIds ?? []) : [item.id];
|
||||
let claimed = false;
|
||||
for (const id of ids) {
|
||||
if (running.delete(id)) claimed = true;
|
||||
}
|
||||
if (claimed) rows += 1;
|
||||
}
|
||||
// Anything left belongs to no known row (a first turn mid-persist); count it as one.
|
||||
return rows + running.size;
|
||||
}, [runningThreadIds, allChatItems]);
|
||||
const anyChatRunning = runningChatCount > 0;
|
||||
// Where "Return to Chat" lands: the newest running chat, not the empty draft New Chat left
|
||||
// active (map insertion order is start order). A compare row runs pane threads that /chat
|
||||
// cannot address, so resolve those back to the pair id the route expects.
|
||||
const runningTarget = useMemo(() => {
|
||||
const ids = Object.entries(runningThreadIds)
|
||||
.filter(([, on]) => on)
|
||||
.map(([id]) => id);
|
||||
const id = ids.length > 0 ? ids[ids.length - 1] : null;
|
||||
if (!id) return null;
|
||||
const pair = allChatItems.find(
|
||||
(item) => item.type === "compare" && (item.threadIds ?? []).includes(id),
|
||||
);
|
||||
return pair
|
||||
? { id: pair.id, compare: true as const }
|
||||
: { id, compare: false as const };
|
||||
}, [runningThreadIds, allChatItems]);
|
||||
const activeThreadId = isChatRoute
|
||||
? (search.thread as string | undefined) ??
|
||||
(search.compare as string | undefined) ??
|
||||
|
|
@ -892,6 +923,12 @@ export function AppSidebar() {
|
|||
variant: "project" | "recent",
|
||||
) {
|
||||
const isPinned = pinnedIdSet.has(item.id);
|
||||
// A compare row's id is the pair id while runningByThreadId is keyed per pane thread,
|
||||
// so aggregate its member threads instead.
|
||||
const isGenerating =
|
||||
item.type === "compare"
|
||||
? (item.threadIds ?? []).some((id) => Boolean(runningThreadIds[id]))
|
||||
: Boolean(runningThreadIds[item.id]);
|
||||
const itemClass =
|
||||
variant === "project"
|
||||
? "group/project-chat-item relative"
|
||||
|
|
@ -951,6 +988,8 @@ export function AppSidebar() {
|
|||
data-testid="recent-thread"
|
||||
data-thread-type={item.type}
|
||||
data-thread-id={item.id}
|
||||
data-generating={isGenerating ? "true" : undefined}
|
||||
aria-busy={isGenerating || undefined}
|
||||
isActive={activeThreadId === item.id}
|
||||
className={buttonClass}
|
||||
onClick={() => {
|
||||
|
|
@ -976,6 +1015,14 @@ export function AppSidebar() {
|
|||
<span className="truncate">
|
||||
{pendingRename?.id === item.id ? pendingRename.title : item.title}
|
||||
</span>
|
||||
{isGenerating && (
|
||||
<Spinner
|
||||
data-testid="chat-row-spinner"
|
||||
// role="status" + label: announced, not motion-only.
|
||||
label={translate("shell.navigation.chatGenerating")}
|
||||
className="ml-auto size-3.5 shrink-0 text-muted-foreground"
|
||||
/>
|
||||
)}
|
||||
</SidebarMenuButton>
|
||||
{variant === "project" && (
|
||||
<button
|
||||
|
|
@ -1288,9 +1335,16 @@ export function AppSidebar() {
|
|||
icon={PencilEdit02Icon}
|
||||
label={
|
||||
showReturnToChat
|
||||
? t("shell.navigation.returnToChat")
|
||||
? runningChatCount > 1
|
||||
// Name the count rather than imply a single live chat.
|
||||
? t("shell.navigation.returnToChats", {
|
||||
count: runningChatCount,
|
||||
})
|
||||
: t("shell.navigation.returnToChat")
|
||||
: t("shell.navigation.newChat")
|
||||
}
|
||||
// Off-route this row is the only sign chats are still running.
|
||||
spinner={anyChatRunning && !isChatRoute}
|
||||
active={
|
||||
isChatRoute &&
|
||||
!search.thread &&
|
||||
|
|
@ -1301,8 +1355,13 @@ export function AppSidebar() {
|
|||
if (showReturnToChat) {
|
||||
// Prefer the running thread so we return to the live generation,
|
||||
// not the empty new chat that became active after New Chat.
|
||||
if (runningThreadId && runningThreadId !== storeThreadId) {
|
||||
navigate({ to: "/chat", search: { thread: runningThreadId } });
|
||||
if (runningTarget && runningTarget.id !== storeThreadId) {
|
||||
navigate({
|
||||
to: "/chat",
|
||||
search: runningTarget.compare
|
||||
? { compare: runningTarget.id }
|
||||
: { thread: runningTarget.id },
|
||||
});
|
||||
} else {
|
||||
navigate({ to: "/chat" });
|
||||
}
|
||||
|
|
|
|||
|
|
@ -371,6 +371,9 @@ function useRafCoalescedText(text: string, isStreaming: boolean): string {
|
|||
|
||||
const MarkdownTextImpl = () => {
|
||||
const { text, status } = useMessagePartText();
|
||||
// Parts are keyed by index, so switching conversations hands this instance a different
|
||||
// message, and Streamdown only extends its parsed blocks: key it per message instead.
|
||||
const messageId = useAuiState(({ message }) => message.id);
|
||||
const displayText = useRafCoalescedText(text, status.type === "running");
|
||||
const processedText = useMemo(
|
||||
() => preprocessLaTeX(displayText),
|
||||
|
|
@ -385,6 +388,7 @@ const MarkdownTextImpl = () => {
|
|||
return (
|
||||
<div data-status={status.type} className="min-w-0 max-w-full">
|
||||
<Streamdown
|
||||
key={messageId}
|
||||
mode="streaming"
|
||||
isAnimating={status.type === "running"}
|
||||
plugins={{ code, math, mermaid }}
|
||||
|
|
|
|||
|
|
@ -724,10 +724,10 @@ function startPromptQueue(
|
|||
}
|
||||
}
|
||||
|
||||
function stopPromptQueueRun() {
|
||||
function stopPromptQueueRun(cancelActiveRun = true) {
|
||||
const activeItem = promptQueueItems[Math.max(promptQueueIndex, 0)];
|
||||
const activeTarget = activeItem?.target;
|
||||
const shouldCancelActiveRun = Boolean(activeItem?.dispatched);
|
||||
const shouldCancelActiveRun = cancelActiveRun && Boolean(activeItem?.dispatched);
|
||||
resetPromptQueue();
|
||||
if (!shouldCancelActiveRun) {
|
||||
return;
|
||||
|
|
@ -740,7 +740,11 @@ function stopPromptQueueRun() {
|
|||
}
|
||||
|
||||
if (typeof window !== "undefined") {
|
||||
window.addEventListener(PROMPT_QUEUE_STOP_EVENT, () => stopPromptQueueRun());
|
||||
window.addEventListener(PROMPT_QUEUE_STOP_EVENT, (event) => {
|
||||
// Navigation leaves the dispatched prompt streaming; an explicit stop cancels it too.
|
||||
const detail = (event as CustomEvent<{ cancelActiveRun?: boolean }>).detail;
|
||||
stopPromptQueueRun(detail?.cancelActiveRun ?? true);
|
||||
});
|
||||
}
|
||||
|
||||
interface PromptQueueCallbacks {
|
||||
|
|
@ -2760,9 +2764,28 @@ const ArtifactsToggle: FC = () => {
|
|||
};
|
||||
|
||||
const ToolStatusDisplay: FC = () => {
|
||||
const toolStatus = useChatRuntimeStore((s) => s.toolStatus);
|
||||
// This conversation's tool call only: a global status would put one chat's "Running
|
||||
// Python..." above every composer. remoteId, not id: the adapter keys this map by
|
||||
// unstable_threadId, so reading id lost the status of every restored chat.
|
||||
const threadListItemId = useAuiState(
|
||||
({ threadListItem }) => threadListItem.remoteId,
|
||||
);
|
||||
const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
|
||||
const [elapsed, setElapsed] = useState(0);
|
||||
const entry = useChatRuntimeStore((s) => {
|
||||
// A first turn starts before its id is persisted, so the adapter files it under
|
||||
// "__default"; only this thread's own run may claim it. Two first turns share that key
|
||||
// with nothing to tell them apart, so claim it only when it holds one run.
|
||||
const unresolved = s.toolStatusByThreadId.__default;
|
||||
const own =
|
||||
s.toolStatusByThreadId[threadListItemId ?? ""] ??
|
||||
(isThreadRunning && unresolved?.length === 1 ? unresolved : undefined);
|
||||
// Newest of the runs behind this key: separate entries, so one finishing cannot blank
|
||||
// a sibling still running a tool.
|
||||
return own?.[own.length - 1];
|
||||
});
|
||||
const toolStatus = entry?.status ?? null;
|
||||
const startedAt = entry?.startedAt ?? null;
|
||||
const [now, setNow] = useState(() => Date.now());
|
||||
const [visible, setVisible] = useState(false);
|
||||
const visibleRef = useRef(false);
|
||||
|
||||
|
|
@ -2771,15 +2794,14 @@ const ToolStatusDisplay: FC = () => {
|
|||
}, [visible]);
|
||||
|
||||
useEffect(() => {
|
||||
if (!toolStatus) {
|
||||
setElapsed(0);
|
||||
if (!startedAt) {
|
||||
if (!isThreadRunning) {
|
||||
setVisible(false);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
setElapsed(0);
|
||||
setNow(Date.now());
|
||||
|
||||
// Debounce visibility by 300ms when the badge isn't already on screen.
|
||||
// Once visible from a prior tool, later tools show immediately so it
|
||||
|
|
@ -2789,24 +2811,27 @@ const ToolStatusDisplay: FC = () => {
|
|||
showTimer = setTimeout(() => setVisible(true), 300);
|
||||
}
|
||||
|
||||
const interval = setInterval(() => {
|
||||
setElapsed((prev) => prev + 1);
|
||||
}, 1000);
|
||||
const interval = setInterval(() => setNow(Date.now()), 1000);
|
||||
return () => {
|
||||
clearInterval(interval);
|
||||
if (showTimer) {
|
||||
clearTimeout(showTimer);
|
||||
}
|
||||
};
|
||||
}, [toolStatus, isThreadRunning]);
|
||||
}, [startedAt, isThreadRunning]);
|
||||
|
||||
if (!(toolStatus && visible)) {
|
||||
if (!(toolStatus && startedAt && visible)) {
|
||||
return null;
|
||||
}
|
||||
// From the store's start time, so returning to the conversation resumes rather than restarting.
|
||||
const elapsed = Math.max(0, Math.floor((now - startedAt) / 1000));
|
||||
const isRunning = toolStatus.startsWith("Running");
|
||||
const StatusIcon = isRunning ? TerminalIcon : GlobeIcon;
|
||||
return (
|
||||
<div className="mb-2 flex w-full flex-row items-center gap-2 px-1.5 pt-0.5 pb-1">
|
||||
<div
|
||||
data-testid="composer-tool-status"
|
||||
className="mb-2 flex w-full flex-row items-center gap-2 px-1.5 pt-0.5 pb-1"
|
||||
>
|
||||
<div className="flex animate-pulse items-center gap-2 rounded-full border border-primary/20 bg-primary/5 px-3 py-1.5 text-xs text-primary">
|
||||
<StatusIcon className="size-3.5" />
|
||||
<span>{toolStatus}</span>
|
||||
|
|
@ -3769,9 +3794,15 @@ const DiffusionCanvas: FC = () => {
|
|||
const isRunning = useAuiState(
|
||||
({ message }) => message.status?.type === "running",
|
||||
);
|
||||
// A non-null canvas is set only by diffusion_frame events (diffusion models only),
|
||||
// so it is a sufficient gate; loadedIsDiffusion can lag the first frame on a fresh load.
|
||||
const canvas = useChatRuntimeStore((s) => s.activeDiffusionCanvas);
|
||||
// Only this conversation's own frames render here; a first turn has no id yet, so it reads
|
||||
// "__default", which is where its run files them until the thread persists.
|
||||
const threadKey =
|
||||
useAuiState(({ threadListItem }) => threadListItem.remoteId) ?? "__default";
|
||||
// A canvas is set only by diffusion_frame events, so its presence is a sufficient gate;
|
||||
// loadedIsDiffusion can lag the first frame on a fresh load.
|
||||
const canvas = useChatRuntimeStore(
|
||||
(s) => s.activeDiffusionCanvasByThreadId[threadKey],
|
||||
);
|
||||
if (!isRunning || !canvas) {
|
||||
return null;
|
||||
}
|
||||
|
|
|
|||
219
studio/frontend/src/components/assistant-ui/tool-code-cell.tsx
Normal file
219
studio/frontend/src/components/assistant-ui/tool-code-cell.tsx
Normal file
|
|
@ -0,0 +1,219 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"use client";
|
||||
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
import { code as codePlugin } from "@streamdown/code";
|
||||
import { CopyIcon, DownloadIcon } from "lucide-react";
|
||||
import { Tick02Icon } from "@/lib/tick-icon";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||||
import { Streamdown } from "streamdown";
|
||||
|
||||
const COPY_RESET_MS = 2000;
|
||||
const SHIKI_THEME = ["github-light", "github-dark"] as [
|
||||
"github-light",
|
||||
"github-dark",
|
||||
];
|
||||
/** Past this the block stays plain monospace: shiki is not worth the main-thread time. */
|
||||
const MAX_HIGHLIGHT_CHARS = 20_000;
|
||||
/** Within this many px of the bottom counts as following the stream. */
|
||||
const PIN_SLACK_PX = 40;
|
||||
|
||||
export function CopyBtn({ text }: { text: string }) {
|
||||
const [copied, setCopied] = useState(false);
|
||||
const timer = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
|
||||
useEffect(() => {
|
||||
return () => {
|
||||
if (timer.current) {
|
||||
clearTimeout(timer.current);
|
||||
}
|
||||
};
|
||||
}, []);
|
||||
|
||||
const copy = useCallback(async () => {
|
||||
if (await copyToClipboard(text)) {
|
||||
setCopied(true);
|
||||
if (timer.current) {
|
||||
clearTimeout(timer.current);
|
||||
}
|
||||
timer.current = setTimeout(() => setCopied(false), COPY_RESET_MS);
|
||||
}
|
||||
}, [text]);
|
||||
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
onClick={copy}
|
||||
className="inline-flex items-center gap-1 rounded px-1.5 py-0.5 text-xs text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
|
||||
aria-label="Copy to clipboard"
|
||||
>
|
||||
{copied ? (
|
||||
<HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="size-3" />
|
||||
) : (
|
||||
<CopyIcon className="size-3" />
|
||||
)}
|
||||
{copied ? "Copied" : "Copy"}
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
function DownloadBtn({ code, name }: { code: string; name: string }) {
|
||||
const download = useCallback(() => {
|
||||
if (typeof document === "undefined") {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
const blob = new Blob([code], { type: "text/plain;charset=utf-8" });
|
||||
const url = URL.createObjectURL(blob);
|
||||
const anchor = document.createElement("a");
|
||||
anchor.href = url;
|
||||
anchor.download = name;
|
||||
document.body.appendChild(anchor);
|
||||
anchor.click();
|
||||
anchor.remove();
|
||||
// Revoke next tick, after the click consumes the URL.
|
||||
setTimeout(() => URL.revokeObjectURL(url), 0);
|
||||
} catch {
|
||||
// Never break the transcript over a download.
|
||||
}
|
||||
}, [code, name]);
|
||||
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
onClick={download}
|
||||
className="inline-flex items-center gap-1 rounded px-1.5 py-0.5 text-xs text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
|
||||
aria-label="Download"
|
||||
>
|
||||
<DownloadIcon className="size-3" />
|
||||
Download
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
/** A fence longer than any backtick run in the code, so a script containing ``` cannot end it early. */
|
||||
function fenceFor(source: string): string {
|
||||
const longest = (source.match(/`+/g) ?? []).reduce(
|
||||
(max, run) => Math.max(max, run.length),
|
||||
0,
|
||||
);
|
||||
return "`".repeat(Math.max(3, longest + 1));
|
||||
}
|
||||
|
||||
/** Syntax-highlighted code via Streamdown + shiki. Always in the DOM as plain monospace, but
|
||||
* shiki only tokenizes once the block nears the viewport, so a long transcript does not
|
||||
* highlight every script up front. Immediate where IntersectionObserver is missing. */
|
||||
function HighlightedCode({
|
||||
code: source,
|
||||
language,
|
||||
plain = false,
|
||||
}: {
|
||||
code: string;
|
||||
language: string;
|
||||
plain?: boolean;
|
||||
}) {
|
||||
const markdown = useMemo(() => {
|
||||
const fence = fenceFor(source);
|
||||
return `${fence}${language}\n${source}\n${fence}`;
|
||||
}, [source, language]);
|
||||
const containerRef = useRef<HTMLDivElement>(null);
|
||||
const [nearViewport, setNearViewport] = useState(
|
||||
() => typeof IntersectionObserver === "undefined",
|
||||
);
|
||||
// Pinned to the bottom until the reader scrolls up, so a streaming payload visibly grows.
|
||||
const pinnedToBottom = useRef(true);
|
||||
useEffect(() => {
|
||||
if (nearViewport) return;
|
||||
const el = containerRef.current;
|
||||
if (!el) return;
|
||||
const io = new IntersectionObserver(
|
||||
(entries) => {
|
||||
if (entries.some((entry) => entry.isIntersecting)) {
|
||||
setNearViewport(true);
|
||||
io.disconnect();
|
||||
}
|
||||
},
|
||||
// Highlight just before the block enters view, so it is ready on arrival.
|
||||
{ rootMargin: "200px" },
|
||||
);
|
||||
io.observe(el);
|
||||
return () => io.disconnect();
|
||||
}, [nearViewport]);
|
||||
|
||||
useEffect(() => {
|
||||
const el = containerRef.current;
|
||||
if (plain && el && pinnedToBottom.current) {
|
||||
el.scrollTop = el.scrollHeight;
|
||||
}
|
||||
}, [plain, source]);
|
||||
|
||||
const handleScroll = () => {
|
||||
const el = containerRef.current;
|
||||
if (el) {
|
||||
pinnedToBottom.current =
|
||||
el.scrollHeight - el.scrollTop - el.clientHeight < PIN_SLACK_PX;
|
||||
}
|
||||
};
|
||||
|
||||
// Skip shiki while the model is writing (it re-tokenizes every fragment) and on payloads too big.
|
||||
const highlight =
|
||||
nearViewport && !plain && source.length <= MAX_HIGHLIGHT_CHARS;
|
||||
|
||||
return (
|
||||
<div
|
||||
ref={containerRef}
|
||||
onScroll={handleScroll}
|
||||
className="max-h-48 overflow-auto text-xs [&_pre]:!m-0 [&_pre]:!bg-transparent [&_pre]:!p-0 [&_pre]:!text-xs [&_[data-streamdown=code-block]]:!my-0 [&_[data-streamdown=code-block]]:!p-3 [&_[data-streamdown=code-block]]:!border-0"
|
||||
>
|
||||
{highlight ? (
|
||||
<Streamdown
|
||||
mode="static"
|
||||
plugins={{ code: codePlugin }}
|
||||
controls={{ code: false }}
|
||||
shikiTheme={SHIKI_THEME}
|
||||
>
|
||||
{markdown}
|
||||
</Streamdown>
|
||||
) : (
|
||||
// A div, not a <pre>: the container's [&_pre]:!p-0 would strip the padding and shift
|
||||
// the content when shiki swaps in. whitespace-pre so long lines scroll.
|
||||
<div className="whitespace-pre p-3 font-mono text-xs text-muted-foreground">
|
||||
{source}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
/** The code a tool is about to run, in the card's collapsible content so the chevron hides code and output together. */
|
||||
export function ToolCodeCell({
|
||||
label,
|
||||
code,
|
||||
language,
|
||||
downloadName,
|
||||
streaming = false,
|
||||
}: {
|
||||
label: string;
|
||||
code: string;
|
||||
language: string;
|
||||
downloadName: string;
|
||||
streaming?: boolean;
|
||||
}) {
|
||||
return (
|
||||
<div className="border-l-2 border-muted-foreground/20 pl-2">
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="text-xs font-medium text-muted-foreground">
|
||||
{label}
|
||||
</span>
|
||||
<div className="flex items-center gap-1">
|
||||
<CopyBtn text={code} />
|
||||
<DownloadBtn code={code} name={downloadName} />
|
||||
</div>
|
||||
</div>
|
||||
<HighlightedCode code={code} language={language} plain={streaming} />
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
@ -10,7 +10,11 @@ import {
|
|||
} from "react";
|
||||
import { useAuiState } from "@assistant-ui/react";
|
||||
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
|
||||
import { toolOutputKey, useToolPaneScope } from "@/features/chat";
|
||||
import {
|
||||
toolOutputKey,
|
||||
useToolPaneScope,
|
||||
useUnresolvedToolPaneScope,
|
||||
} from "@/features/chat";
|
||||
import { ChevronDownIcon } from "lucide-react";
|
||||
import { Wrench01Icon } from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
|
|
@ -239,16 +243,23 @@ const ToolGroupImpl: FC<
|
|||
// Force the group open when any call is receiving tool_output events.
|
||||
const toolLiveOutput = useChatRuntimeStore((s) => s.toolLiveOutput);
|
||||
const paneScope = useToolPaneScope();
|
||||
const unresolvedScope = useUnresolvedToolPaneScope();
|
||||
const hasLiveOutput = useAuiState(({ message }) =>
|
||||
message.parts
|
||||
.slice(startIndex, endIndex + 1)
|
||||
.some(
|
||||
(part) =>
|
||||
part.type === "tool-call" &&
|
||||
Object.prototype.hasOwnProperty.call(
|
||||
// Either scope: a first turn writes under the unresolved one for its whole
|
||||
// life, even after the autosave assigns the id (see useToolOutputFor).
|
||||
(Object.prototype.hasOwnProperty.call(
|
||||
toolLiveOutput,
|
||||
toolOutputKey(paneScope, part.toolCallId),
|
||||
),
|
||||
) ||
|
||||
Object.prototype.hasOwnProperty.call(
|
||||
toolLiveOutput,
|
||||
toolOutputKey(unresolvedScope, part.toolCallId),
|
||||
)),
|
||||
),
|
||||
);
|
||||
// Keep the group open once a confirmation or live output forced it (so an
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@
|
|||
"use client";
|
||||
|
||||
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
|
||||
import { toolOutputKey, useToolPaneScope } from "@/features/chat";
|
||||
import { useToolOutputFor, useToolPaneScope } from "@/features/chat";
|
||||
import { useEffect, useMemo, useRef } from "react";
|
||||
import { tailText } from "./tool-result-output";
|
||||
|
||||
|
|
@ -16,8 +16,10 @@ import { tailText } from "./tool-result-output";
|
|||
*/
|
||||
export function ToolLiveOutput({ toolCallId }: { toolCallId: string }) {
|
||||
const paneScope = useToolPaneScope();
|
||||
const output = useChatRuntimeStore(
|
||||
(s) => s.toolLiveOutput[toolOutputKey(paneScope, toolCallId)] ?? "",
|
||||
const output = useToolOutputFor(
|
||||
useChatRuntimeStore((s) => s.toolLiveOutput),
|
||||
paneScope,
|
||||
toolCallId,
|
||||
);
|
||||
const scrollRef = useRef<HTMLPreElement>(null);
|
||||
// Pinned to the bottom until the user scrolls up (handler below), so
|
||||
|
|
|
|||
|
|
@ -3,28 +3,25 @@
|
|||
|
||||
"use client";
|
||||
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
import { getAuthToken } from "@/features/auth/session";
|
||||
import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
|
||||
import { useToolArgsStatus } from "@assistant-ui/react";
|
||||
import { code as codePlugin } from "@streamdown/code";
|
||||
import { CodeIcon, CopyIcon, DownloadIcon } from "lucide-react";
|
||||
import { Tick02Icon } from "@/lib/tick-icon";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { CodeIcon } from "lucide-react";
|
||||
import { Spinner } from "@/components/ui/spinner";
|
||||
import { memo, useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||||
import { Streamdown } from "streamdown";
|
||||
import { memo } from "react";
|
||||
import {
|
||||
ToolFallbackContent,
|
||||
ToolFallbackRoot,
|
||||
ToolFallbackTrigger,
|
||||
} from "./tool-fallback";
|
||||
import { CopyBtn, ToolCodeCell } from "./tool-code-cell";
|
||||
import { ToolLiveOutput } from "./tool-live-output";
|
||||
import { ToolResultOutput } from "./tool-result-output";
|
||||
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
|
||||
import {
|
||||
preferFullToolOutput,
|
||||
toolOutputKey,
|
||||
useToolAwaitingApproval,
|
||||
useToolOutputFor,
|
||||
useToolPaneScope,
|
||||
} from "@/features/chat";
|
||||
|
||||
|
|
@ -34,151 +31,6 @@ interface StructuredResult {
|
|||
sessionId: string;
|
||||
}
|
||||
|
||||
const MAX_DISPLAY = 10_000;
|
||||
const COPY_RESET_MS = 2000;
|
||||
const SHIKI_THEME = ["github-light", "github-dark"] as ["github-light", "github-dark"];
|
||||
|
||||
function truncate(text: string): string {
|
||||
return text.length <= MAX_DISPLAY
|
||||
? text
|
||||
: `${text.slice(0, MAX_DISPLAY)}\n... (truncated)`;
|
||||
}
|
||||
|
||||
function CopyBtn({ text }: { text: string }) {
|
||||
const [copied, setCopied] = useState(false);
|
||||
const timer = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
|
||||
useEffect(() => {
|
||||
return () => {
|
||||
if (timer.current) {
|
||||
clearTimeout(timer.current);
|
||||
}
|
||||
};
|
||||
}, []);
|
||||
|
||||
const copy = useCallback(async () => {
|
||||
if (await copyToClipboard(text)) {
|
||||
setCopied(true);
|
||||
if (timer.current) {
|
||||
clearTimeout(timer.current);
|
||||
}
|
||||
timer.current = setTimeout(() => setCopied(false), COPY_RESET_MS);
|
||||
}
|
||||
}, [text]);
|
||||
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
onClick={copy}
|
||||
className="inline-flex items-center gap-1 rounded px-1.5 py-0.5 text-xs text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
|
||||
aria-label="Copy to clipboard"
|
||||
>
|
||||
{copied ? (
|
||||
<HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="size-3" />
|
||||
) : (
|
||||
<CopyIcon className="size-3" />
|
||||
)}
|
||||
{copied ? "Copied" : "Copy"}
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
/** Save the script as a .py file via a client-side Blob. */
|
||||
function DownloadBtn({ code, name = "script.py" }: { code: string; name?: string }) {
|
||||
const download = useCallback(() => {
|
||||
if (typeof document === "undefined") {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
const blob = new Blob([code], { type: "text/x-python" });
|
||||
const url = URL.createObjectURL(blob);
|
||||
const anchor = document.createElement("a");
|
||||
anchor.href = url;
|
||||
anchor.download = name;
|
||||
document.body.appendChild(anchor);
|
||||
anchor.click();
|
||||
anchor.remove();
|
||||
// Revoke next tick, after the click consumes the URL.
|
||||
setTimeout(() => URL.revokeObjectURL(url), 0);
|
||||
} catch {
|
||||
// Best-effort: never break the transcript over a download.
|
||||
}
|
||||
}, [code, name]);
|
||||
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
onClick={download}
|
||||
className="inline-flex items-center gap-1 rounded px-1.5 py-0.5 text-xs text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
|
||||
aria-label="Download script"
|
||||
>
|
||||
<DownloadIcon className="size-3" />
|
||||
Download
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
/** Syntax-highlighted code via Streamdown + shiki; inherits parent container.
|
||||
* The script is always in the DOM (a plain monospace placeholder), but shiki
|
||||
* only tokenizes once the block scrolls near the viewport, so a long transcript
|
||||
* with many scripts doesn't highlight every one up front. Falls back to
|
||||
* immediate highlight when IntersectionObserver is unavailable (SSR / tests). */
|
||||
function HighlightedCode({ code: source, language }: { code: string; language: string }) {
|
||||
const display = useMemo(() => truncate(source), [source]);
|
||||
const markdown = useMemo(
|
||||
() => `\`\`\`${language}\n${display}\n\`\`\``,
|
||||
[display, language],
|
||||
);
|
||||
const containerRef = useRef<HTMLDivElement>(null);
|
||||
const [highlight, setHighlight] = useState(
|
||||
() => typeof IntersectionObserver === "undefined",
|
||||
);
|
||||
useEffect(() => {
|
||||
if (highlight) return;
|
||||
const el = containerRef.current;
|
||||
if (!el) return;
|
||||
const io = new IntersectionObserver(
|
||||
(entries) => {
|
||||
if (entries.some((entry) => entry.isIntersecting)) {
|
||||
setHighlight(true);
|
||||
io.disconnect();
|
||||
}
|
||||
},
|
||||
// Highlight just before the block enters view so it's colorized by the
|
||||
// time the user reaches it, without tokenizing off-screen scripts.
|
||||
{ rootMargin: "200px" },
|
||||
);
|
||||
io.observe(el);
|
||||
return () => io.disconnect();
|
||||
}, [highlight]);
|
||||
return (
|
||||
<div
|
||||
ref={containerRef}
|
||||
className="max-h-48 overflow-auto text-xs [&_pre]:!m-0 [&_pre]:!bg-transparent [&_pre]:!p-0 [&_pre]:!text-xs [&_[data-streamdown=code-block]]:!my-0 [&_[data-streamdown=code-block]]:!p-3 [&_[data-streamdown=code-block]]:!border-0"
|
||||
>
|
||||
{highlight ? (
|
||||
<Streamdown
|
||||
mode="static"
|
||||
plugins={{ code: codePlugin }}
|
||||
controls={{ code: false }}
|
||||
shikiTheme={SHIKI_THEME}
|
||||
>
|
||||
{markdown}
|
||||
</Streamdown>
|
||||
) : (
|
||||
// A div, not a <pre>: the container's [&_pre]:!p-0 would override a
|
||||
// <pre>'s padding and shift the content by p-3 when shiki swaps in. Keep
|
||||
// the same p-3, and whitespace-pre (not pre-wrap) so long lines scroll in
|
||||
// the container's overflow-auto exactly like the highlighted <pre>, rather
|
||||
// than wrapping taller and then collapsing when shiki swaps in.
|
||||
<div className="whitespace-pre p-3 font-mono text-xs text-muted-foreground">
|
||||
{display}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function isStructuredResult(val: unknown): val is StructuredResult {
|
||||
return (
|
||||
typeof val === "object" &&
|
||||
|
|
@ -221,46 +73,50 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({
|
|||
// Show the fuller live stream over a truncated result, keeping its exit
|
||||
// status. Session-transient: after a reload only the result remains.
|
||||
const paneScope = useToolPaneScope();
|
||||
const fullOutput = useChatRuntimeStore(
|
||||
(s) => s.toolFullOutput[toolOutputKey(paneScope, toolCallId)] ?? "",
|
||||
const fullOutput = useToolOutputFor(
|
||||
useChatRuntimeStore((s) => s.toolFullOutput),
|
||||
paneScope,
|
||||
toolCallId,
|
||||
);
|
||||
const displayOutput = preferFullToolOutput(fullOutput, output);
|
||||
|
||||
const authToken = getAuthToken();
|
||||
// The gate only opens once the call parsed, so a pending approval means the script is
|
||||
// written even while the args status still reads as streaming.
|
||||
const awaitingApproval = useToolAwaitingApproval(toolCallId);
|
||||
const isWriting = isWritingCode && !awaitingApproval;
|
||||
|
||||
return (
|
||||
// Status/output collapse from history; the script source renders outside
|
||||
// ToolFallbackContent so it stays visible on reopen (#7165).
|
||||
// Script, status and output all collapse behind the one chevron.
|
||||
<ToolFallbackRoot defaultOpen={isRunning}>
|
||||
<ToolFallbackTrigger
|
||||
toolName={firstLine ? `Python: ${firstLine}` : "Python"}
|
||||
status={status}
|
||||
icon={CodeIcon}
|
||||
/>
|
||||
{code && (
|
||||
<div className="mt-1 pl-5">
|
||||
<div className="border-l-2 border-muted-foreground/20 pl-2">
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="text-xs font-medium text-muted-foreground">
|
||||
script
|
||||
</span>
|
||||
<div className="flex items-center gap-1">
|
||||
<CopyBtn text={code} />
|
||||
<DownloadBtn code={code} />
|
||||
</div>
|
||||
</div>
|
||||
<HighlightedCode code={code} language="python" />
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
<ToolFallbackContent>
|
||||
{code && (
|
||||
<ToolCodeCell
|
||||
label="script"
|
||||
code={code}
|
||||
language="python"
|
||||
downloadName="script.py"
|
||||
streaming={isWriting}
|
||||
/>
|
||||
)}
|
||||
<div className="border-l-2 border-muted-foreground/20 pl-2">
|
||||
{/* Output */}
|
||||
{isRunning ? (
|
||||
<>
|
||||
<div className="mt-2 flex items-center gap-2 text-sm text-muted-foreground">
|
||||
<Spinner className="size-3.5" />
|
||||
<span>{isWritingCode ? "Writing code…" : "Running…"}</span>
|
||||
<span>
|
||||
{awaitingApproval
|
||||
? "Waiting for approval…"
|
||||
: isWriting
|
||||
? "Writing code…"
|
||||
: "Running…"}
|
||||
</span>
|
||||
</div>
|
||||
{/* Live stdout streamed via tool_output SSE events. */}
|
||||
<ToolLiveOutput toolCallId={toolCallId} />
|
||||
|
|
|
|||
|
|
@ -3,69 +3,27 @@
|
|||
|
||||
"use client";
|
||||
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
|
||||
import { useToolArgsStatus } from "@assistant-ui/react";
|
||||
import { CopyIcon, TerminalIcon } from "lucide-react";
|
||||
import { Tick02Icon } from "@/lib/tick-icon";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { TerminalIcon } from "lucide-react";
|
||||
import { Spinner } from "@/components/ui/spinner";
|
||||
import { memo, useCallback, useEffect, useRef, useState } from "react";
|
||||
import { memo } from "react";
|
||||
import {
|
||||
ToolFallbackContent,
|
||||
ToolFallbackRoot,
|
||||
ToolFallbackTrigger,
|
||||
} from "./tool-fallback";
|
||||
import { CopyBtn, ToolCodeCell } from "./tool-code-cell";
|
||||
import { ToolLiveOutput } from "./tool-live-output";
|
||||
import { ToolResultOutput } from "./tool-result-output";
|
||||
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
|
||||
import {
|
||||
preferFullToolOutput,
|
||||
toolOutputKey,
|
||||
useToolAwaitingApproval,
|
||||
useToolOutputFor,
|
||||
useToolPaneScope,
|
||||
} from "@/features/chat";
|
||||
|
||||
const COPY_RESET_MS = 2000;
|
||||
|
||||
function CopyBtn({ text }: { text: string }) {
|
||||
const [copied, setCopied] = useState(false);
|
||||
const timer = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
|
||||
useEffect(() => {
|
||||
return () => {
|
||||
if (timer.current) {
|
||||
clearTimeout(timer.current);
|
||||
}
|
||||
};
|
||||
}, []);
|
||||
|
||||
const copy = useCallback(async () => {
|
||||
if (await copyToClipboard(text)) {
|
||||
setCopied(true);
|
||||
if (timer.current) {
|
||||
clearTimeout(timer.current);
|
||||
}
|
||||
timer.current = setTimeout(() => setCopied(false), COPY_RESET_MS);
|
||||
}
|
||||
}, [text]);
|
||||
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
onClick={copy}
|
||||
className="inline-flex items-center gap-1 rounded px-1.5 py-0.5 text-xs text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
|
||||
aria-label="Copy to clipboard"
|
||||
>
|
||||
{copied ? (
|
||||
<HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="size-3" />
|
||||
) : (
|
||||
<CopyIcon className="size-3" />
|
||||
)}
|
||||
{copied ? "Copied" : "Copy"}
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
const TerminalToolUIImpl: ToolCallMessagePartComponent = ({
|
||||
toolCallId,
|
||||
args,
|
||||
|
|
@ -87,13 +45,19 @@ const TerminalToolUIImpl: ToolCallMessagePartComponent = ({
|
|||
// Show the fuller live stream over a truncated result, keeping its exit
|
||||
// status. Session-transient: after a reload only the result remains.
|
||||
const paneScope = useToolPaneScope();
|
||||
const fullOutput = useChatRuntimeStore(
|
||||
(s) => s.toolFullOutput[toolOutputKey(paneScope, toolCallId)] ?? "",
|
||||
const fullOutput = useToolOutputFor(
|
||||
useChatRuntimeStore((s) => s.toolFullOutput),
|
||||
paneScope,
|
||||
toolCallId,
|
||||
);
|
||||
const displayOutput = preferFullToolOutput(fullOutput, output);
|
||||
// The gate only opens once the call parsed, so a pending approval means the command is
|
||||
// written even while the args status still reads as streaming.
|
||||
const awaitingApproval = useToolAwaitingApproval(toolCallId);
|
||||
const isWriting = isWritingCommand && !awaitingApproval;
|
||||
|
||||
return (
|
||||
// Open when mounted mid-run so live output shows; collapsed from history.
|
||||
// Open mid-run so command and live output show, collapsed from history.
|
||||
<ToolFallbackRoot defaultOpen={isRunning}>
|
||||
<ToolFallbackTrigger
|
||||
toolName={command ? `$ ${command.slice(0, 60)}` : "Terminal"}
|
||||
|
|
@ -101,12 +65,27 @@ const TerminalToolUIImpl: ToolCallMessagePartComponent = ({
|
|||
icon={TerminalIcon}
|
||||
/>
|
||||
<ToolFallbackContent>
|
||||
{command && (
|
||||
<ToolCodeCell
|
||||
label="command"
|
||||
code={command}
|
||||
language="bash"
|
||||
downloadName="command.sh"
|
||||
streaming={isWriting}
|
||||
/>
|
||||
)}
|
||||
<div className="border-l-2 border-muted-foreground/20 pl-2">
|
||||
{isRunning ? (
|
||||
<>
|
||||
<div className="flex items-center gap-2 text-sm text-muted-foreground">
|
||||
<Spinner className="size-3.5" />
|
||||
<span>{isWritingCommand ? "Writing command…" : "Running…"}</span>
|
||||
<span>
|
||||
{awaitingApproval
|
||||
? "Waiting for approval…"
|
||||
: isWriting
|
||||
? "Writing command…"
|
||||
: "Running…"}
|
||||
</span>
|
||||
</div>
|
||||
{/* Live stdout streamed via tool_output SSE events. */}
|
||||
<ToolLiveOutput toolCallId={toolCallId} />
|
||||
|
|
|
|||
|
|
@ -6,15 +6,22 @@
|
|||
import { Loader2Icon } from "lucide-react";
|
||||
import { cn } from "@/lib/utils";
|
||||
|
||||
/**
|
||||
* App-wide spinner: a clean circular arc with a rounded cap (lucide
|
||||
* Loader2 / LoaderCircle), animated, inheriting the current text color.
|
||||
*/
|
||||
function Spinner({ className }: { className?: string }) {
|
||||
/** App-wide spinner inheriting the current text color. `label` overrides the announcement
|
||||
* where "loading" is not what it means (a sidebar chat is generating). */
|
||||
function Spinner({
|
||||
className,
|
||||
label = "Loading",
|
||||
"data-testid": dataTestId,
|
||||
}: {
|
||||
className?: string;
|
||||
label?: string;
|
||||
"data-testid"?: string;
|
||||
}) {
|
||||
return (
|
||||
<Loader2Icon
|
||||
role="status"
|
||||
aria-label="Loading"
|
||||
aria-label={label}
|
||||
data-testid={dataTestId}
|
||||
className={cn("size-4 shrink-0 animate-spin", className)}
|
||||
/>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ import {
|
|||
shouldPreserveFullOutput,
|
||||
toolOutputKey,
|
||||
toolPaneScope,
|
||||
toolThreadScope,
|
||||
} from "../tool-output-scope";
|
||||
import type { ModelType } from "../types";
|
||||
import { isMultimodalResponse } from "../types/api";
|
||||
|
|
@ -2232,7 +2233,20 @@ export function createOpenAIStreamAdapter(
|
|||
: undefined;
|
||||
|
||||
const threadKey = resolvedThreadId;
|
||||
runtime.setThreadRunning(threadKey, true);
|
||||
// The run is durable on the server, but Stop, archive and delete reach a background
|
||||
// thread only through this map: without a handle the supervisor kept planning against
|
||||
// a deleted conversation. Registered before the run exists, since the thread can be
|
||||
// stopped while createResearchRun is still in flight.
|
||||
let researchRunId: string | null = null;
|
||||
let researchStopRequested = false;
|
||||
const researchServerCancel = () => {
|
||||
researchStopRequested = true;
|
||||
if (researchRunId) {
|
||||
void cancelResearchRun(researchRunId).catch(() => {});
|
||||
}
|
||||
};
|
||||
runtime.registerThreadServerCancel(threadKey, researchServerCancel);
|
||||
runtime.setThreadRunning(threadKey, true, { owner: researchServerCancel });
|
||||
let report = "";
|
||||
let releaseResearchFollow: (() => void) | null = null;
|
||||
const researchFollowController = new AbortController();
|
||||
|
|
@ -2272,6 +2286,13 @@ export function createOpenAIStreamAdapter(
|
|||
blockedDomains: [...runtime.researchWebsitePolicy.blockedDomains],
|
||||
},
|
||||
});
|
||||
researchRunId = createdRun.id;
|
||||
if (researchStopRequested) {
|
||||
// Stopped while createResearchRun was still in flight, so the handle had no
|
||||
// id to act on. Replay it rather than following a run the user already ended.
|
||||
void cancelResearchRun(createdRun.id).catch(() => {});
|
||||
return;
|
||||
}
|
||||
releaseResearchFollow = beginExternalResearchFollow(
|
||||
createdRun,
|
||||
detachResearchFollow,
|
||||
|
|
@ -2330,7 +2351,8 @@ export function createOpenAIStreamAdapter(
|
|||
} finally {
|
||||
abortSignal.removeEventListener("abort", forwardAdapterAbort);
|
||||
releaseResearchFollow?.();
|
||||
runtime.setThreadRunning(threadKey, false);
|
||||
runtime.clearThreadServerCancel(threadKey, researchServerCancel);
|
||||
runtime.setThreadRunning(threadKey, false, { owner: researchServerCancel });
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
|
@ -2339,17 +2361,21 @@ export function createOpenAIStreamAdapter(
|
|||
? `${sandboxSessionId || "_default"}:${resolvedThreadId}`
|
||||
: sandboxSessionId || "_default";
|
||||
const toolConfirmationIdsByBackendId = new Map<string, string>();
|
||||
// Store keys are pane-scoped since local tool ids ("call_0") repeat across
|
||||
// turns and concurrent panes (compare mode). Track this run's keys so
|
||||
// cleanup can't wipe another pane's.
|
||||
const toolOutputPaneScope = toolPaneScope(
|
||||
options.modelType,
|
||||
options.pairId,
|
||||
// Local tool ids ("call_0") repeat across turns, panes and conversations, so scope by pane
|
||||
// AND thread. unstable_threadId alone, no activeThreadId fallback: the reader has only
|
||||
// threadListItem.remoteId, which is exactly this value.
|
||||
const toolOutputPaneScope = toolThreadScope(
|
||||
toolPaneScope(options.modelType, options.pairId),
|
||||
unstable_threadId,
|
||||
);
|
||||
const scopedToolOutputKey = (id: string) =>
|
||||
toolOutputKey(toolOutputPaneScope, id);
|
||||
const runToolLiveOutputKeys = new Set<string>();
|
||||
const resolvedThreadKey = resolvedThreadId ?? null;
|
||||
// Which conversation was on screen when this run started. A first turn has no id yet, so
|
||||
// this is the only way to tell later whether the user has switched away from it.
|
||||
const activeThreadIdAtRunStart =
|
||||
useChatRuntimeStore.getState().activeThreadId ?? null;
|
||||
const pendingImageEditReferenceForRun = runtime.pendingImageEditReference;
|
||||
const selectedImageEditReference =
|
||||
(pendingImageEditReferenceForRun?.threadId ?? null) ===
|
||||
|
|
@ -2755,8 +2781,11 @@ export function createOpenAIStreamAdapter(
|
|||
// waitForRunEnd resolves instead of hanging: this gate fires
|
||||
// before the streaming path's setThreadRunning(true).
|
||||
const gatedThreadKey = resolvedThreadId || "__default";
|
||||
runtime.setThreadRunning(gatedThreadKey, true);
|
||||
runtime.setThreadRunning(gatedThreadKey, false);
|
||||
// Own token: siblings share "__default", so an ownerless clear would drop their
|
||||
// entries while they are still generating.
|
||||
const gateOwner = () => {};
|
||||
runtime.setThreadRunning(gatedThreadKey, true, { owner: gateOwner });
|
||||
runtime.setThreadRunning(gatedThreadKey, false, { owner: gateOwner });
|
||||
clearSelectedImageEditReference();
|
||||
throw new Error(imageGateReason);
|
||||
}
|
||||
|
|
@ -2774,13 +2803,44 @@ export function createOpenAIStreamAdapter(
|
|||
}
|
||||
const useAdapter = await resolveUseAdapter(resolvedThreadId, options);
|
||||
|
||||
const threadKey = resolvedThreadId || "__default";
|
||||
// A first turn files its handles under "__default"; autosave then assigns a real id and
|
||||
// adoptDefaultThreadRun re-keys them mid-run. Resolve per use so later writes and the
|
||||
// final clear follow the run instead of stranding entries behind.
|
||||
const liveThreadKey = (owner: () => void) =>
|
||||
threadKey === "__default"
|
||||
? useChatRuntimeStore.getState().runKeyForOwner(threadKey, owner)
|
||||
: threadKey;
|
||||
|
||||
// Per-run token so a delayed stop POST can't match the next run.
|
||||
const cancelId =
|
||||
typeof crypto !== "undefined" && "randomUUID" in crypto
|
||||
? crypto.randomUUID()
|
||||
: `${Date.now()}-${Math.random().toString(36).slice(2)}`;
|
||||
|
||||
// Per-run abort, chained to assistant-ui's signal. cancelByThreadId only holds the visible
|
||||
// thread's cancelRun(), so this controller is the only way to end a backgrounded chat's
|
||||
// request; the cancel POST below reaches llama-server only.
|
||||
const runAbort = new AbortController();
|
||||
const runSignal = runAbort.signal;
|
||||
const forwardAbort = () => runAbort.abort(abortSignal.reason);
|
||||
// Declared here, not at its registration below: it doubles as this run's identity token
|
||||
// on the per-thread maps (see registerThreadServerCancel).
|
||||
const serverCancel = () => runAbort.abort();
|
||||
if (abortSignal.aborted) {
|
||||
forwardAbort();
|
||||
} else {
|
||||
abortSignal.addEventListener("abort", forwardAbort, { once: true });
|
||||
}
|
||||
|
||||
// ── Audio model path (non-streaming) ─────────────────────
|
||||
const activeModel = runtime.models.find(
|
||||
(m) => m.id === params.checkpoint,
|
||||
);
|
||||
if (activeModel?.isAudio && !activeModel?.hasAudioInput) {
|
||||
const threadKey = resolvedThreadId || "__default";
|
||||
runtime.setThreadRunning(threadKey, true);
|
||||
const audioCancel = () => runAbort.abort();
|
||||
runtime.registerThreadServerCancel(threadKey, audioCancel);
|
||||
runtime.setThreadRunning(threadKey, true, { owner: audioCancel });
|
||||
try {
|
||||
yield {
|
||||
content: [{ type: "text" as const, text: "Generating audio..." }],
|
||||
|
|
@ -2790,6 +2850,10 @@ export function createOpenAIStreamAdapter(
|
|||
{
|
||||
model: params.checkpoint,
|
||||
messages: outboundMessages,
|
||||
// Same run in both registries: without it the backend files this under no
|
||||
// thread, and the stop-chats prompt counts the named local run and the
|
||||
// unnamed backend one as two.
|
||||
...(resolvedThreadId ? { thread_id: resolvedThreadId } : {}),
|
||||
stream: false,
|
||||
temperature: params.temperature,
|
||||
top_p: params.topP,
|
||||
|
|
@ -2800,7 +2864,7 @@ export function createOpenAIStreamAdapter(
|
|||
presence_penalty: params.presencePenalty,
|
||||
...(useAdapter === undefined ? {} : { use_adapter: useAdapter }),
|
||||
},
|
||||
abortSignal,
|
||||
runSignal,
|
||||
);
|
||||
|
||||
const audioUrl = `data:audio/wav;base64,${result.audio.data}`;
|
||||
|
|
@ -2813,19 +2877,21 @@ export function createOpenAIStreamAdapter(
|
|||
],
|
||||
};
|
||||
} catch (err) {
|
||||
if (!abortSignal.aborted) {
|
||||
if (!runSignal.aborted) {
|
||||
toast.error("Audio generation failed", {
|
||||
description: err instanceof Error ? err.message : "Unknown error",
|
||||
});
|
||||
}
|
||||
throw err;
|
||||
} finally {
|
||||
runtime.setThreadRunning(threadKey, false);
|
||||
abortSignal.removeEventListener("abort", forwardAbort);
|
||||
const audioKey = liveThreadKey(audioCancel);
|
||||
runtime.setThreadRunning(audioKey, false, { owner: audioCancel });
|
||||
runtime.clearThreadServerCancel(audioKey, audioCancel);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const threadKey = resolvedThreadId || "__default";
|
||||
let waitingFirstChunk = true;
|
||||
let firstTokenSettled = false;
|
||||
const streamStartTime = Date.now();
|
||||
|
|
@ -2856,10 +2922,15 @@ export function createOpenAIStreamAdapter(
|
|||
const warmupDelayMs = 450;
|
||||
const warmupTimer = setTimeout(() => {
|
||||
if (!waitingFirstChunk) return;
|
||||
if (abortSignal.aborted) return;
|
||||
if (runSignal.aborted) return;
|
||||
runtime.setGeneratingStatus("waiting");
|
||||
}, warmupDelayMs);
|
||||
runtime.setThreadRunning(threadKey, true);
|
||||
// Flagged local/external so the model-swap gate only counts the chats a reload ends; the
|
||||
// backend leaves external-provider runs out of active_generations for the same reason.
|
||||
runtime.setThreadRunning(threadKey, true, {
|
||||
local: !isExternalRequest,
|
||||
owner: serverCancel,
|
||||
});
|
||||
let cumulativeText = "";
|
||||
let reasoningStartAt: number | null = null;
|
||||
let reasoningDuration = 0;
|
||||
|
|
@ -3025,21 +3096,12 @@ export function createOpenAIStreamAdapter(
|
|||
timings?: ServerTimings;
|
||||
} | null = null;
|
||||
|
||||
// Per-run cancellation token so a delayed stop POST can't match
|
||||
// the next run on the same thread.
|
||||
const cancelId =
|
||||
typeof crypto !== "undefined" && "randomUUID" in crypto
|
||||
? crypto.randomUUID()
|
||||
: `${Date.now()}-${Math.random().toString(36).slice(2)}`;
|
||||
|
||||
// Colab-style proxies can swallow fetch aborts, so also POST
|
||||
// /inference/cancel explicitly on abort.
|
||||
const onAbortCancel = () => {
|
||||
// assistant-ui aborts with AbortError(detach=true) when a thread's runtime
|
||||
// unmounts (navigation / background thread switch) and detach=false for an
|
||||
// explicit Stop. Only a real Stop cancels the backend run; a detach must
|
||||
// leave a backgrounded generation streaming.
|
||||
if ((abortSignal.reason as { detach?: boolean } | undefined)?.detach) {
|
||||
// assistant-ui aborts with detach=true when a runtime unmounts and detach=false for an
|
||||
// explicit Stop. Only a real Stop cancels the backend run; runSignal forwards the reason.
|
||||
if ((runSignal.reason as { detach?: boolean } | undefined)?.detach) {
|
||||
return;
|
||||
}
|
||||
const body: Record<string, string> = { cancel_id: cancelId };
|
||||
|
|
@ -3060,11 +3122,17 @@ export function createOpenAIStreamAdapter(
|
|||
keepalive: true,
|
||||
}).catch(() => {});
|
||||
};
|
||||
|
||||
// Stop handle for when this conversation is not the visible one, which cancelByThreadId
|
||||
// cannot reach. Aborting this run's own controller closes just its request, and the
|
||||
// listener above posts its cancel_id so llama-server stops decoding too. For an
|
||||
// external provider the abort is the stop, since its cancel_id is never registered.
|
||||
runtime.registerThreadServerCancel(threadKey, serverCancel);
|
||||
try {
|
||||
if (abortSignal.aborted) {
|
||||
if (runSignal.aborted) {
|
||||
onAbortCancel();
|
||||
} else {
|
||||
abortSignal.addEventListener("abort", onAbortCancel, { once: true });
|
||||
runSignal.addEventListener("abort", onAbortCancel, { once: true });
|
||||
}
|
||||
|
||||
const {
|
||||
|
|
@ -3536,7 +3604,7 @@ export function createOpenAIStreamAdapter(
|
|||
}
|
||||
clearSelectedImageEditReference();
|
||||
await ThreadAutosaveHandle.awaitFirstSave(resolvedThreadId);
|
||||
const stream = streamChatCompletions(requestPayload, abortSignal);
|
||||
const stream = streamChatCompletions(requestPayload, runSignal);
|
||||
|
||||
for await (const chunk of stream) {
|
||||
const chunkModel = (chunk as { model?: unknown }).model;
|
||||
|
|
@ -3549,7 +3617,11 @@ export function createOpenAIStreamAdapter(
|
|||
chunk as unknown as { _toolStatus?: string }
|
||||
)._toolStatus;
|
||||
if (toolStatusText !== undefined) {
|
||||
runtime.setToolStatus(toolStatusText || null);
|
||||
runtime.setToolStatus(
|
||||
liveThreadKey(serverCancel),
|
||||
toolStatusText || null,
|
||||
serverCancel,
|
||||
);
|
||||
continue;
|
||||
}
|
||||
|
||||
|
|
@ -3578,7 +3650,9 @@ export function createOpenAIStreamAdapter(
|
|||
}
|
||||
)._diffusionFrame;
|
||||
if (diffusionFrame !== undefined) {
|
||||
runtime.setActiveDiffusionCanvas({
|
||||
// Keyed by thread so a background run's frames stay out of the visible chat
|
||||
// instead of overwriting the frame it is painting.
|
||||
runtime.setActiveDiffusionCanvas(liveThreadKey(serverCancel), {
|
||||
block: diffusionFrame.block ?? 0,
|
||||
step: diffusionFrame.step ?? 0,
|
||||
total: diffusionFrame.total ?? 0,
|
||||
|
|
@ -3719,8 +3793,16 @@ export function createOpenAIStreamAdapter(
|
|||
const approvalId = (toolEvent.approval_id as string) || "";
|
||||
const awaitingConfirmation =
|
||||
toolEvent.awaiting_confirmation === true;
|
||||
// Reuse a provisional card's part id, else the confirmation-scoped id
|
||||
// opens a second card and the first spins "Running" forever.
|
||||
const openPartId = backendToolCallId
|
||||
? toolPartIdByBackendId.get(backendToolCallId)
|
||||
: undefined;
|
||||
const reuseOpenPart =
|
||||
!!openPartId &&
|
||||
toolCallParts.some((p) => p.toolCallId === openPartId);
|
||||
const id =
|
||||
awaitingConfirmation && approvalId
|
||||
awaitingConfirmation && approvalId && !reuseOpenPart
|
||||
? `${toolConfirmationScopeId}:${approvalId}`
|
||||
: backendToolCallId
|
||||
? resolveToolPartId(backendToolCallId)
|
||||
|
|
@ -4299,9 +4381,17 @@ export function createOpenAIStreamAdapter(
|
|||
// Anthropic-only (billed at the write premium).
|
||||
const cacheWriteTokens = meta?.usage?.cache_creation_input_tokens ?? 0;
|
||||
|
||||
// Gate on the captured checkpoint still being active so a late
|
||||
// completion from provider A doesn't populate the bar after a
|
||||
// mid-stream switch to provider B.
|
||||
// Gate on the captured checkpoint so a late completion from provider A cannot populate
|
||||
// the bar after a mid-stream switch to B, and on the captured thread so a background
|
||||
// run finishing after New Chat cannot repaint another chat's usage. An unresolved run
|
||||
// has no id to compare, so compare what was on screen when it started. A first turn is
|
||||
// adopted onto an id mid-run and autosave moves activeThreadId with it, so read the
|
||||
// adopted key, or the run stays "unresolved" for life and the bar stays blank.
|
||||
const usageKey = liveThreadKey(serverCancel);
|
||||
const usageThreadKey = usageKey === "__default" ? null : usageKey;
|
||||
const usageThreadIsVisible =
|
||||
useChatRuntimeStore.getState().activeThreadId ===
|
||||
(usageThreadKey ?? activeThreadIdAtRunStart);
|
||||
if (
|
||||
meta?.usage &&
|
||||
typeof meta.usage.prompt_tokens === "number" &&
|
||||
|
|
@ -4309,13 +4399,23 @@ export function createOpenAIStreamAdapter(
|
|||
typeof meta.usage.total_tokens === "number" &&
|
||||
useChatRuntimeStore.getState().params.checkpoint === params.checkpoint
|
||||
) {
|
||||
useChatRuntimeStore.getState().setContextUsage({
|
||||
const usage = {
|
||||
promptTokens: meta.usage.prompt_tokens,
|
||||
completionTokens: meta.usage.completion_tokens,
|
||||
totalTokens: meta.usage.total_tokens,
|
||||
cachedTokens,
|
||||
cacheWriteTokens,
|
||||
});
|
||||
};
|
||||
// File it under this run's own thread even when the gate below blocks the visible
|
||||
// write, so switching back re-applies it.
|
||||
if (usageThreadKey !== null) {
|
||||
useChatRuntimeStore
|
||||
.getState()
|
||||
.setThreadContextUsage(usageThreadKey, usage);
|
||||
}
|
||||
if (usageThreadIsVisible) {
|
||||
useChatRuntimeStore.getState().setContextUsage(usage);
|
||||
}
|
||||
}
|
||||
|
||||
const finishedAt = Date.now();
|
||||
|
|
@ -4368,7 +4468,7 @@ export function createOpenAIStreamAdapter(
|
|||
settleFirstTokenErr(
|
||||
err instanceof Error ? err : new Error("Generation failed"),
|
||||
);
|
||||
if (!abortSignal.aborted) {
|
||||
if (!runSignal.aborted) {
|
||||
const msg = err instanceof Error ? err.message : String(err);
|
||||
if (err instanceof GenerationLengthError) {
|
||||
toast.error("Response ran out of tokens", {
|
||||
|
|
@ -4406,13 +4506,18 @@ export function createOpenAIStreamAdapter(
|
|||
}
|
||||
throw err;
|
||||
} finally {
|
||||
abortSignal.removeEventListener("abort", onAbortCancel);
|
||||
runSignal.removeEventListener("abort", onAbortCancel);
|
||||
abortSignal.removeEventListener("abort", forwardAbort);
|
||||
// Resolve once: the clears below drop the owner the lookup keys on.
|
||||
const cleanupKey = liveThreadKey(serverCancel);
|
||||
const confirmStore = useChatRuntimeStore.getState();
|
||||
for (const part of toolCallParts) {
|
||||
confirmStore.clearToolConfirmation(part.toolCallId);
|
||||
}
|
||||
runtime.setGeneratingStatus(null);
|
||||
runtime.setToolStatus(null);
|
||||
// Scoped by thread AND by run: a global clear wiped every other running chat's badge,
|
||||
// and an unowned one wiped a concurrent run's badge behind the same key.
|
||||
runtime.setToolStatus(cleanupKey, null, serverCancel);
|
||||
// Clear only this run's live keys (a concurrent pane owns its own). A
|
||||
// key still here streamed stdout but never reached tool_end (SSE drop or
|
||||
// cancel), so promote it to full output first, else the partial
|
||||
|
|
@ -4426,20 +4531,23 @@ export function createOpenAIStreamAdapter(
|
|||
store.clearToolLiveOutput(liveKey);
|
||||
}
|
||||
runToolLiveOutputKeys.clear();
|
||||
// Drop the transient denoising canvas so the finished bubble shows only
|
||||
// the committed markdown answer (cancellation/error included).
|
||||
runtime.setActiveDiffusionCanvas(null);
|
||||
// Drop the transient denoising canvas so the finished bubble shows only the committed
|
||||
// answer. Scoped: a global clear wiped another denoising chat's frame.
|
||||
runtime.clearActiveDiffusionCanvasForThread(cleanupKey);
|
||||
clearTimeout(warmupTimer);
|
||||
if (waitingFirstChunk) {
|
||||
if (firstTokenSettled) {
|
||||
settleFirstTokenOk();
|
||||
} else if (abortSignal.aborted) {
|
||||
} else if (runSignal.aborted) {
|
||||
settleFirstTokenErr(new Error("Cancelled"));
|
||||
} else {
|
||||
settleFirstTokenErr(new Error("No tokens received"));
|
||||
}
|
||||
}
|
||||
runtime.setThreadRunning(threadKey, false);
|
||||
// serverCancel narrows both clears: runs with no resolved thread id share the "__default"
|
||||
// key, so a blind clear could drop a sibling's entry.
|
||||
runtime.setThreadRunning(cleanupKey, false, { owner: serverCancel });
|
||||
runtime.clearThreadServerCancel(cleanupKey, serverCancel);
|
||||
}
|
||||
},
|
||||
};
|
||||
|
|
|
|||
|
|
@ -129,6 +129,27 @@ export async function getApiMonitorEntry(id: string): Promise<ApiMonitorEntry> {
|
|||
return parseJsonOrThrow<ApiMonitorEntry>(response);
|
||||
}
|
||||
|
||||
export interface ActiveGenerationsResponse {
|
||||
count: number;
|
||||
/** Conversations with a generation in flight. Shorter than `count` when a
|
||||
* first turn started before its thread id was persisted. */
|
||||
thread_ids: string[];
|
||||
/** One entry per in-flight request. `kind` is "chat" unless it is an
|
||||
* embeddings / completions / audio call, which has no conversation. */
|
||||
active?: { thread_id: string | null; kind?: string }[];
|
||||
parallel_slots: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Chats generating on the backend right now. Authoritative where `runningByThreadId` is not:
|
||||
* that map is per-tab, empty after a reload and blind to a second tab, and /load and /unload
|
||||
* 409 on these.
|
||||
*/
|
||||
export async function getActiveGenerations(): Promise<ActiveGenerationsResponse> {
|
||||
const response = await authFetch("/api/inference/active-generations");
|
||||
return parseJsonOrThrow<ActiveGenerationsResponse>(response);
|
||||
}
|
||||
|
||||
export async function loadModel(
|
||||
payload: LoadModelRequest,
|
||||
): Promise<LoadModelResponse> {
|
||||
|
|
|
|||
|
|
@ -2682,9 +2682,10 @@ export function ChatPage({
|
|||
ggufNativeContextLength: null,
|
||||
activeNativePathToken: null,
|
||||
activeNativePathExpiresAtMs: null,
|
||||
// Clear previous-model counters, else the relaxed external-provider
|
||||
// render gate shows stale stats until the next completion.
|
||||
// Clear previous-model counters, else the relaxed external-provider render gate shows
|
||||
// stale stats. The per-thread copies go too, so a switch back cannot re-apply.
|
||||
contextUsage: null,
|
||||
contextUsageByThreadId: {},
|
||||
supportsReasoning: reasoningCaps.supportsReasoning,
|
||||
reasoningAlwaysOn: reasoningCaps.reasoningAlwaysOn,
|
||||
reasoningStyle: reasoningCaps.reasoningStyle,
|
||||
|
|
@ -2906,7 +2907,13 @@ export function ChatPage({
|
|||
) {
|
||||
return;
|
||||
}
|
||||
store.setContextUsage(usage);
|
||||
// Key by the thread this restore read, like the history loader: the await above can
|
||||
// outlast a switch away, and an unkeyed write would file this thread's usage under
|
||||
// the incoming one.
|
||||
store.setThreadContextUsage(threadId, usage);
|
||||
if (store.activeThreadId === threadId) {
|
||||
store.setContextUsage(usage);
|
||||
}
|
||||
})
|
||||
.catch((error) => {
|
||||
if (!isExpectedBackgroundChatStorageError(error)) {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,92 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import {
|
||||
AlertDialog,
|
||||
AlertDialogAction,
|
||||
AlertDialogCancel,
|
||||
AlertDialogContent,
|
||||
AlertDialogDescription,
|
||||
AlertDialogFooter,
|
||||
AlertDialogHeader,
|
||||
AlertDialogTitle,
|
||||
} from "@/components/ui/alert-dialog";
|
||||
import { useStopRunningChatsDialogStore } from "../stores/stop-running-chats-dialog-store";
|
||||
|
||||
/**
|
||||
* Confirmation for applying a model or reload-required setting while chats are generating.
|
||||
* They share one llama-server, so the swap ends all of them: name them and make the user
|
||||
* opt in rather than truncating silently.
|
||||
*/
|
||||
export function StopRunningChatsDialog() {
|
||||
const open = useStopRunningChatsDialogStore((s) => s.open);
|
||||
const count = useStopRunningChatsDialogStore((s) => s.count);
|
||||
const titles = useStopRunningChatsDialogStore((s) => s.titles);
|
||||
const action = useStopRunningChatsDialogStore((s) => s.action);
|
||||
const hasNonChat = useStopRunningChatsDialogStore((s) => s.hasNonChat);
|
||||
const effect = useStopRunningChatsDialogStore((s) => s.effect);
|
||||
const resolve = useStopRunningChatsDialogStore((s) => s.resolve);
|
||||
|
||||
// Embeddings, raw completions and audio share the model but are not conversations,
|
||||
// so name them generically rather than offering to stop chats that do not exist.
|
||||
const noun = hasNonChat
|
||||
? count === 1
|
||||
? "request"
|
||||
: "requests"
|
||||
: count === 1
|
||||
? "chat"
|
||||
: "chats";
|
||||
const sharer = hasNonChat ? "request" : "conversation";
|
||||
// Ejecting leaves no model loaded. Saying it "reloads the model" and offering "Stop and
|
||||
// reload" promised the opposite of what confirming does, for the destructive one.
|
||||
const unloads = effect === "unload";
|
||||
const lead = unloads
|
||||
? `${action || "Unloading the model"} leaves no model loaded, and every open ${sharer} shares it, `
|
||||
: `${action ? `${action} reloads the model, ` : "Reloading the model "}which every open ${sharer} shares, `;
|
||||
const shown = titles.slice(0, 5);
|
||||
const remaining = Math.max(0, titles.length - shown.length);
|
||||
|
||||
return (
|
||||
<AlertDialog
|
||||
open={open}
|
||||
onOpenChange={(next) => {
|
||||
// Escape / overlay click must resolve, or the caller's await hangs.
|
||||
if (!next) resolve(false);
|
||||
}}
|
||||
>
|
||||
<AlertDialogContent>
|
||||
<AlertDialogHeader>
|
||||
<AlertDialogTitle>
|
||||
Stop {count} running {noun}?
|
||||
</AlertDialogTitle>
|
||||
<AlertDialogDescription>
|
||||
{lead}so {count === 1 ? "this" : "these"} {noun} will stop
|
||||
{hasNonChat ? "" : " generating"}. Work produced so far is kept.
|
||||
</AlertDialogDescription>
|
||||
</AlertDialogHeader>
|
||||
{shown.length > 0 && (
|
||||
<ul className="max-h-40 overflow-y-auto rounded-md border bg-muted/40 px-3 py-2 text-sm">
|
||||
{shown.map((title) => (
|
||||
<li key={title} className="truncate py-0.5">
|
||||
{title}
|
||||
</li>
|
||||
))}
|
||||
{remaining > 0 && (
|
||||
<li className="py-0.5 text-muted-foreground">
|
||||
and {remaining} more
|
||||
</li>
|
||||
)}
|
||||
</ul>
|
||||
)}
|
||||
<AlertDialogFooter>
|
||||
<AlertDialogCancel onClick={() => resolve(false)}>
|
||||
Keep generating
|
||||
</AlertDialogCancel>
|
||||
<AlertDialogAction onClick={() => resolve(true)}>
|
||||
{unloads ? "Stop and unload" : "Stop and reload"}
|
||||
</AlertDialogAction>
|
||||
</AlertDialogFooter>
|
||||
</AlertDialogContent>
|
||||
</AlertDialog>
|
||||
);
|
||||
}
|
||||
|
|
@ -28,6 +28,7 @@ import {
|
|||
validateModel,
|
||||
} from "../api/chat-api";
|
||||
import { formatEta, formatRate } from "../utils/format-transfer";
|
||||
import { confirmStopRunningChatsIfNeeded } from "../utils/confirm-stop-running-chats";
|
||||
import {
|
||||
GPU_LAYERS_AUTO,
|
||||
isLocalModelPath,
|
||||
|
|
@ -463,7 +464,14 @@ export function useChatModelRuntime() {
|
|||
useChatRuntimeStore.getState().setModelLoading(true);
|
||||
void (async () => {
|
||||
try {
|
||||
// Unforced on purpose: a chat may stream on the PREVIOUS model and must not be killed by
|
||||
// cancelling this load. Nothing to report, since the route runs its stop-loading fast
|
||||
// path ahead of the active-chat refusal.
|
||||
await unloadModel({ model_path: model.id }).catch(() => {});
|
||||
// clearCheckpoint above assumed nothing was left loaded, but a forced switch keeps the
|
||||
// previous model resident until /load's teardown, and the stop-loading fast path leaves
|
||||
// it there. Take the answer from the backend, which reports none once it was evicted.
|
||||
await syncInferenceStatusToStore().catch(() => {});
|
||||
} finally {
|
||||
cancelUnloadPendingRef.current = false;
|
||||
if (!loadingModelRef.current) {
|
||||
|
|
@ -505,10 +513,11 @@ export function useChatModelRuntime() {
|
|||
// as a duplicate), don't start a second concurrent load and don't swallow the
|
||||
// request: surface it so the user waits or cancels. Centralized here so every
|
||||
// entry point is covered, not just the staged Load button.
|
||||
const inFlightLoad =
|
||||
loadingModelRef.current ??
|
||||
useChatRuntimeStore.getState().loadingModelPick;
|
||||
if (inFlightLoad) {
|
||||
const bailIfLoadInFlight = (): boolean => {
|
||||
const inFlightLoad =
|
||||
loadingModelRef.current ??
|
||||
useChatRuntimeStore.getState().loadingModelPick;
|
||||
if (!inFlightLoad) return false;
|
||||
if (typeof selection !== "string" && selection.previousConfig) {
|
||||
applyPerModelConfigToRuntime(selection.previousConfig);
|
||||
}
|
||||
|
|
@ -516,7 +525,7 @@ export function useChatModelRuntime() {
|
|||
inFlightLoad.id === modelId &&
|
||||
(inFlightLoad.ggufVariant ?? null) === (ggufVariant ?? null) &&
|
||||
(inFlightLoad.nativePathToken ?? null) === (nativePathToken ?? null);
|
||||
if (loadingSamePick) return;
|
||||
if (loadingSamePick) return true;
|
||||
const message =
|
||||
"Another model is already loading. Wait for it to finish or cancel it first.";
|
||||
setModelsError(message);
|
||||
|
|
@ -524,8 +533,61 @@ export function useChatModelRuntime() {
|
|||
toast.info("Another model is already loading", {
|
||||
description: "Wait for it to finish or cancel it first.",
|
||||
});
|
||||
return true;
|
||||
};
|
||||
if (bailIfLoadInFlight()) return;
|
||||
|
||||
// Picking an external provider leaves the local model resident and stops the status poll
|
||||
// mirroring it, so params.checkpoint cannot tell whether this pick is that same model.
|
||||
// Ask the backend before prompting: /load answers already_loaded ahead of its cancel
|
||||
// hook, so the dialog would promise to stop chats this pick never interrupts. A staged
|
||||
// config always carries forceReload, so Apply still reloads and prompts.
|
||||
const selectedCheckpoint =
|
||||
useChatRuntimeStore.getState().params.checkpoint;
|
||||
if (!forceReload && isExternalModelId(selectedCheckpoint)) {
|
||||
const residentStatus = await getInferenceStatus().catch(() => null);
|
||||
if (
|
||||
residentStatus &&
|
||||
resolveInferenceCheckpointId(residentStatus) === modelId &&
|
||||
(residentStatus.gguf_variant ?? null) === (ggufVariant ?? null)
|
||||
) {
|
||||
// Same window as the confirm below: a rival load may have started during that GET,
|
||||
// and it owns the resident model now.
|
||||
if (bailIfLoadInFlight()) return;
|
||||
// Roll back the config pre-applied for the load that is not happening BEFORE hydrating,
|
||||
// so the resident model's status wins over the staged snapshot.
|
||||
if (typeof selection !== "string" && selection.previousConfig) {
|
||||
applyPerModelConfigToRuntime(selection.previousConfig);
|
||||
}
|
||||
const previousGgufVariant =
|
||||
useChatRuntimeStore.getState().activeGgufVariant;
|
||||
useChatRuntimeStore
|
||||
.getState()
|
||||
.setCheckpoint(modelId, residentStatus.gguf_variant);
|
||||
applyActiveModelStatusToStore(residentStatus, {
|
||||
previousCheckpoint: selectedCheckpoint,
|
||||
previousGgufVariant,
|
||||
});
|
||||
syncModelCapabilities(modelId, residentStatus);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Every chat decodes on the llama-server this load replaces, so ask first, then allow the
|
||||
// cancel; the 409 gate stays armed for callers that never confirmed.
|
||||
const stopDecision = await confirmStopRunningChatsIfNeeded(
|
||||
forceReload ? "Applying these settings" : "Loading a different model",
|
||||
);
|
||||
if (!stopDecision.proceed) {
|
||||
if (typeof selection !== "string" && selection.previousConfig) {
|
||||
applyPerModelConfigToRuntime(selection.previousConfig);
|
||||
}
|
||||
return;
|
||||
}
|
||||
// Re-check: the confirm above awaits a GET, so a pick in that window would start a rival
|
||||
// load over the same refs. Nothing awaits before the reservation below.
|
||||
if (bailIfLoadInFlight()) return;
|
||||
const forceCancelActive = stopDecision.forceCancelActive;
|
||||
|
||||
const explicitIsLora =
|
||||
typeof selection === "string" ? undefined : selection.isLora;
|
||||
|
|
@ -765,6 +827,10 @@ export function useChatModelRuntime() {
|
|||
upgrade: validation.transformers_upgrade,
|
||||
// No installable release: custom-code models may fall back to the trust_remote_code gate below.
|
||||
trustRemoteCodeFallback: validation.requires_trust_remote_code,
|
||||
// The install refuses while chats generate and takes no force flag of its own, so
|
||||
// without this the "Stop and reload" the user just confirmed dies here: Retry hits
|
||||
// the same 409, and this path leaves chats running.
|
||||
forceCancelActive,
|
||||
});
|
||||
// The install unloads the previous model before the swap (even when
|
||||
// the swap then fails), so any exit after this point must roll back.
|
||||
|
|
@ -808,7 +874,14 @@ export function useChatModelRuntime() {
|
|||
: undefined;
|
||||
|
||||
if (currentCheckpoint) {
|
||||
await unloadModel({ model_path: currentCheckpoint });
|
||||
// With chats generating, skip this preliminary unload: it cancels them ahead of /load's
|
||||
// preflight, so a rejected target truncates replies for a model that never loads
|
||||
// (/load evicts past those checks itself). Idle, unload first and free VRAM early.
|
||||
if (!forceCancelActive) {
|
||||
await unloadModel({ model_path: currentCheckpoint });
|
||||
}
|
||||
// Set either way: /load can still leave no model resident, and an unneeded rollback
|
||||
// hits already_loaded before the gate.
|
||||
previousWasUnloaded = true;
|
||||
}
|
||||
if (abortCtrl.signal.aborted) throw new Error("Cancelled");
|
||||
|
|
@ -915,6 +988,7 @@ export function useChatModelRuntime() {
|
|||
n_cpu_moe: loadNCpuMoe,
|
||||
tensor_split: loadSplitRatio ?? undefined,
|
||||
gpu_ids: loadSelectedGpuIds ?? undefined,
|
||||
force_cancel_active: forceCancelActive,
|
||||
});
|
||||
|
||||
// If cancelled while loading, don't update UI to show
|
||||
|
|
@ -1144,6 +1218,8 @@ export function useChatModelRuntime() {
|
|||
n_cpu_moe: stateBeforeUnload.loadedNCpuMoe ?? 0,
|
||||
tensor_split: stateBeforeUnload.loadedSplitRatio ?? undefined,
|
||||
gpu_ids: stateBeforeUnload.loadedGpuIds ?? undefined,
|
||||
// The failed swap already unloaded the server those runs used.
|
||||
force_cancel_active: true,
|
||||
});
|
||||
const rollbackSpeculativeType = normalizeSpeculativeType(
|
||||
rollbackResponse.speculative_type,
|
||||
|
|
@ -1540,13 +1616,15 @@ export function useChatModelRuntime() {
|
|||
if (!params.checkpoint) {
|
||||
return false;
|
||||
}
|
||||
const runtime = useChatRuntimeStore.getState();
|
||||
if (runtime.modelLoading || runtime.loadingModelPick) {
|
||||
const bailIfLoading = (): boolean => {
|
||||
const runtime = useChatRuntimeStore.getState();
|
||||
if (!runtime.modelLoading && !runtime.loadingModelPick) return false;
|
||||
toast.info("A model is loading", {
|
||||
description: "Wait for it to finish or cancel it first.",
|
||||
});
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
};
|
||||
if (bailIfLoading()) return false;
|
||||
setModelsError(null);
|
||||
if (isExternalModelId(params.checkpoint)) {
|
||||
clearCheckpoint();
|
||||
|
|
@ -1554,8 +1632,21 @@ export function useChatModelRuntime() {
|
|||
return true;
|
||||
}
|
||||
try {
|
||||
// Ejecting tears down llama-server, so every chat stops. Same prompt, but it
|
||||
// leaves no model loaded, so it must not be worded as a reload.
|
||||
const stopDecision = await confirmStopRunningChatsIfNeeded(
|
||||
"Unloading the model",
|
||||
"unload",
|
||||
);
|
||||
if (!stopDecision.proceed) return false;
|
||||
// Same window as selectModel: a load may have started during the confirm.
|
||||
if (bailIfLoading()) return false;
|
||||
|
||||
async function performUnload(): Promise<void> {
|
||||
await unloadModel({ model_path: params.checkpoint });
|
||||
await unloadModel({
|
||||
model_path: params.checkpoint,
|
||||
force_cancel_active: stopDecision.forceCancelActive,
|
||||
});
|
||||
clearCheckpoint();
|
||||
await refresh();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ import {
|
|||
updateStoredChatThread,
|
||||
} from "../utils/chat-history-storage";
|
||||
import { clearComposerDraft } from "../utils/composer-draft";
|
||||
import { stopChatThread } from "../utils/stop-chat-thread";
|
||||
import {
|
||||
markChatThreadsDeleted,
|
||||
removeChatThreadTombstones,
|
||||
|
|
@ -25,6 +26,8 @@ import {
|
|||
export interface SidebarItem {
|
||||
type: "single" | "compare";
|
||||
id: string;
|
||||
/** The pane threads behind this row id; `runningByThreadId` is keyed per pane thread. */
|
||||
threadIds?: string[];
|
||||
title: string;
|
||||
createdAt: number;
|
||||
updatedAt: number;
|
||||
|
|
@ -56,11 +59,13 @@ export function groupThreads(
|
|||
const existing = pairItems.get(t.pairId);
|
||||
if (existing) {
|
||||
existing.updatedAt = Math.max(existing.updatedAt, lastActivityAt(t));
|
||||
existing.threadIds?.push(t.id);
|
||||
continue;
|
||||
}
|
||||
const item: SidebarItem = {
|
||||
type: "compare",
|
||||
id: t.pairId,
|
||||
threadIds: [t.id],
|
||||
title: t.title,
|
||||
createdAt: t.createdAt,
|
||||
updatedAt: lastActivityAt(t),
|
||||
|
|
@ -160,10 +165,9 @@ export function useChatSidebarItems(options?: {
|
|||
}
|
||||
|
||||
function cancelIfRunning(threadId: string): void {
|
||||
const { runningByThreadId, cancelByThreadId } =
|
||||
useChatRuntimeStore.getState();
|
||||
if (!runningByThreadId[threadId]) return;
|
||||
cancelByThreadId[threadId]?.();
|
||||
// Reaches a background thread, which cancelByThreadId cannot: a deleted chat must stop,
|
||||
// or the run keeps writing to a conversation that is gone.
|
||||
stopChatThread(threadId);
|
||||
}
|
||||
|
||||
export async function renameChatItem(
|
||||
|
|
|
|||
|
|
@ -55,8 +55,12 @@ export {
|
|||
export {
|
||||
preferFullToolOutput,
|
||||
toolOutputKey,
|
||||
toolThreadScope,
|
||||
useToolOutputFor,
|
||||
useUnresolvedToolPaneScope,
|
||||
useToolPaneScope,
|
||||
} from "./tool-output-scope";
|
||||
export { useToolAwaitingApproval } from "./tool-approval";
|
||||
export { PermissionModeDropdown } from "./permission-mode-select";
|
||||
export { useChatSearchStore } from "./stores/chat-search-store";
|
||||
export { usePinnedChatsStore } from "./stores/pinned-chats-store";
|
||||
|
|
@ -80,6 +84,7 @@ export {
|
|||
export { ApiProviderLogo } from "./api-provider-logo";
|
||||
export { useExternalProvidersStore } from "./stores/external-providers-store";
|
||||
export { ChatSearchDialog } from "./components/chat-search-dialog";
|
||||
export { StopRunningChatsDialog } from "./components/stop-running-chats-dialog";
|
||||
export { setTrainingCompareHandoff } from "./lib/training-compare-handoff";
|
||||
export type { ProjectRecord } from "./types";
|
||||
export { clearAllChats, countAllChats } from "./utils/clear-all-chats";
|
||||
|
|
|
|||
|
|
@ -85,7 +85,11 @@ import { requestPromptQueueStop } from "./utils/prompt-queue-boundary";
|
|||
import { isAssistantLocalThreadId } from "./utils/thread-ids";
|
||||
|
||||
const pendingHistoryAppendByMessageId = new Map<string, Promise<void>>();
|
||||
const pendingRunStartReadyByMessageId = new Map<string, Promise<void>>();
|
||||
// Resolves to the thread id assigned when this message's chat was first persisted.
|
||||
const pendingRunStartReadyByMessageId = new Map<
|
||||
string,
|
||||
Promise<string | undefined>
|
||||
>();
|
||||
|
||||
type TitleResponse = {
|
||||
choices?: Array<{
|
||||
|
|
@ -699,6 +703,10 @@ function createStudioDbAdapter(
|
|||
|
||||
async initialize(threadId: string) {
|
||||
await ensureThreadRecord({ threadId, modelType, pairId, projectId });
|
||||
// A run already streaming on this thread filed its handles under "__default" because
|
||||
// the id did not exist yet. Re-key them now, or the sidebar row and Stop look up an
|
||||
// id nothing is registered against.
|
||||
useChatRuntimeStore.getState().adoptDefaultThreadRun(threadId);
|
||||
return { remoteId: threadId, externalId: undefined };
|
||||
},
|
||||
|
||||
|
|
@ -835,8 +843,8 @@ function trackHistoryAppend(
|
|||
|
||||
function trackRunStartReady(
|
||||
messageId: string,
|
||||
ready: Promise<void>,
|
||||
): Promise<void> {
|
||||
ready: Promise<string | undefined>,
|
||||
): Promise<string | undefined> {
|
||||
pendingRunStartReadyByMessageId.set(messageId, ready);
|
||||
const cleanup = () => {
|
||||
setTimeout(() => {
|
||||
|
|
@ -851,7 +859,7 @@ function trackRunStartReady(
|
|||
|
||||
async function waitForRunStartHistoryAppend(
|
||||
messages: Parameters<ChatModelAdapter["run"]>[0]["messages"],
|
||||
): Promise<void> {
|
||||
): Promise<string | undefined> {
|
||||
// Deep Research reserves an assistant placeholder before invoking the model
|
||||
// adapter, so the user message is not necessarily the final entry here.
|
||||
const userMessage = [...messages]
|
||||
|
|
@ -862,15 +870,16 @@ async function waitForRunStartHistoryAppend(
|
|||
}
|
||||
const runStartReady = pendingRunStartReadyByMessageId.get(userMessage.id);
|
||||
const historyAppendReady = pendingHistoryAppendByMessageId.get(userMessage.id);
|
||||
const pending = [runStartReady, historyAppendReady].filter(
|
||||
(ready): ready is Promise<void> => ready !== undefined,
|
||||
);
|
||||
if (pending.length === 0) {
|
||||
return;
|
||||
if (runStartReady === undefined && historyAppendReady === undefined) {
|
||||
return undefined;
|
||||
}
|
||||
let didBecomeReady = false;
|
||||
let adoptedThreadId: string | undefined;
|
||||
try {
|
||||
await Promise.all(pending);
|
||||
[adoptedThreadId] = await Promise.all([
|
||||
runStartReady ?? Promise.resolve(undefined),
|
||||
historyAppendReady?.then(() => undefined),
|
||||
]);
|
||||
didBecomeReady = true;
|
||||
} finally {
|
||||
if (
|
||||
|
|
@ -881,14 +890,22 @@ async function waitForRunStartHistoryAppend(
|
|||
pendingRunStartReadyByMessageId.delete(userMessage.id);
|
||||
}
|
||||
}
|
||||
return adoptedThreadId;
|
||||
}
|
||||
|
||||
function createPersistedRunAdapter(adapter: ChatModelAdapter): ChatModelAdapter {
|
||||
return {
|
||||
...adapter,
|
||||
async *run(options) {
|
||||
await waitForRunStartHistoryAppend(options.messages);
|
||||
const result = adapter.run(options);
|
||||
const adoptedThreadId = await waitForRunStartHistoryAppend(options.messages);
|
||||
// The thread has an id by the time that resolves, but assistant-ui bound unstable_threadId
|
||||
// before the await. Hand the run its real id so a first turn never files its handles
|
||||
// under the unresolved key that concurrent runs share.
|
||||
const result = adapter.run(
|
||||
!options.unstable_threadId && adoptedThreadId
|
||||
? { ...options, unstable_threadId: adoptedThreadId }
|
||||
: options,
|
||||
);
|
||||
if (!result) {
|
||||
return;
|
||||
}
|
||||
|
|
@ -1153,7 +1170,13 @@ function useStudioRuntimeAdapters(
|
|||
: typeof store.ggufContextLength === "number" &&
|
||||
store.ggufContextLength > 0;
|
||||
if (savedUsage && withinLocalLimit && modelMatches) {
|
||||
store.setContextUsage(savedUsage);
|
||||
// Key by the thread this loader read, not whichever is active when the await resolves:
|
||||
// a switch inside it would file this thread's usage under the incoming one. Same rule
|
||||
// the adapter's end-of-run write follows.
|
||||
store.setThreadContextUsage(remoteId, savedUsage);
|
||||
if (store.activeThreadId === remoteId) {
|
||||
store.setContextUsage(savedUsage);
|
||||
}
|
||||
}
|
||||
|
||||
// If any message has a stored parentId, reconstruct the tree so
|
||||
|
|
@ -1179,7 +1202,10 @@ function useStudioRuntimeAdapters(
|
|||
|
||||
append({ parentId, message }: ExportedMessageRepositoryItem) {
|
||||
const initializeThread = aui.threadListItem().initialize();
|
||||
trackRunStartReady(message.id, initializeThread.then(() => undefined));
|
||||
trackRunStartReady(
|
||||
message.id,
|
||||
initializeThread.then(({ remoteId }) => remoteId),
|
||||
);
|
||||
const write = (async () => {
|
||||
const { remoteId } = await initializeThread;
|
||||
if (isChatThreadDeleted(remoteId)) {
|
||||
|
|
@ -1308,17 +1334,6 @@ function createRuntimeHook(modelType: ModelType, pairId?: string) {
|
|||
};
|
||||
}
|
||||
|
||||
function stopChatRun(threadId: string | null | undefined) {
|
||||
if (!threadId) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
useChatRuntimeStore.getState().cancelByThreadId[threadId]?.();
|
||||
} catch {
|
||||
// The run may have ended while navigation was mounting.
|
||||
}
|
||||
}
|
||||
|
||||
function ThreadAutoSwitch({
|
||||
threadId,
|
||||
syncActiveThreadId = true,
|
||||
|
|
@ -1333,8 +1348,9 @@ function ThreadAutoSwitch({
|
|||
useEffect(() => {
|
||||
if (!isLoading && mainThreadId !== threadId) {
|
||||
if (syncActiveThreadId) {
|
||||
requestPromptQueueStop();
|
||||
stopChatRun(mainThreadId);
|
||||
// Stop queueing prompts to the outgoing thread but leave its run alone: its runtime
|
||||
// stays mounted and keeps streaming. Only an explicit Stop cancels one.
|
||||
requestPromptQueueStop({ cancelActiveRun: false });
|
||||
}
|
||||
const switchResult = aui.threads().switchToThread(threadId) as unknown;
|
||||
if (
|
||||
|
|
@ -1365,16 +1381,14 @@ function ThreadNewChatSwitch({
|
|||
}: { nonce: string }): ReactElement | null {
|
||||
const aui = useAui();
|
||||
const isLoading = useAuiState(({ threads }) => threads.isLoading);
|
||||
const mainThreadId = useAuiState(({ threads }) => threads.mainThreadId);
|
||||
const mainThreadIdRef = useRef(mainThreadId);
|
||||
mainThreadIdRef.current = mainThreadId;
|
||||
|
||||
// The outgoing thread is not read here: New Chat leaves it running.
|
||||
useEffect(() => {
|
||||
if (isLoading) {
|
||||
return;
|
||||
}
|
||||
requestPromptQueueStop();
|
||||
stopChatRun(mainThreadIdRef.current);
|
||||
// New Chat leaves the previous conversation generating: its runtime stays mounted and
|
||||
// the sidebar spins. Stopping it is its own Stop button's job.
|
||||
requestPromptQueueStop({ cancelActiveRun: false });
|
||||
// Switch to a fresh local thread without persisting it yet; persistence
|
||||
// still happens on first message append.
|
||||
void aui.threads().switchToNewThread();
|
||||
|
|
|
|||
|
|
@ -762,6 +762,30 @@ export function isDownloadableHubRepo(x: {
|
|||
);
|
||||
}
|
||||
|
||||
type ContextUsageSnapshot = {
|
||||
promptTokens: number;
|
||||
completionTokens: number;
|
||||
totalTokens: number;
|
||||
cachedTokens: number;
|
||||
// Anthropic-only; optional so pre-cache-stats persisted entries load.
|
||||
cacheWriteTokens?: number;
|
||||
};
|
||||
|
||||
/**
|
||||
* One live run behind `runningByThreadId[id]`, with the `local` flag it started with so the
|
||||
* model-swap gate can tell llama-server runs from external ones when runs share a key.
|
||||
*/
|
||||
type ThreadRunOwner = {
|
||||
owner: () => void;
|
||||
local: boolean;
|
||||
};
|
||||
|
||||
type ToolStatusEntry = {
|
||||
status: string;
|
||||
startedAt: number;
|
||||
owner?: () => void;
|
||||
};
|
||||
|
||||
type ChatRuntimeStore = {
|
||||
settingsHydrated: boolean;
|
||||
params: InferenceParams;
|
||||
|
|
@ -771,7 +795,25 @@ type ChatRuntimeStore = {
|
|||
models: ChatModelSummary[];
|
||||
loras: ChatLoraSummary[];
|
||||
runningByThreadId: Record<string, boolean>;
|
||||
/**
|
||||
* The subset of `runningByThreadId` decoding on the local llama-server. Swapping the local
|
||||
* model neither interrupts an external-provider chat nor needs its consent, which is why
|
||||
* the backend keeps those out of `active_generations` too.
|
||||
*/
|
||||
localRunByThreadId: Record<string, boolean>;
|
||||
/**
|
||||
* Which runs set `runningByThreadId[id]`; see `setThreadRunning`'s `owner`. A list, not one
|
||||
* entry: runs without a resolved thread id share the "__default" key, so one entry would let
|
||||
* a newer run's clear delete an older run's flag while it still generates.
|
||||
*/
|
||||
runOwnerByThreadId: Record<string, ThreadRunOwner[]>;
|
||||
cancelByThreadId: Record<string, () => void>;
|
||||
/**
|
||||
* Backend cancels for the threads generating in the background. `cancelByThreadId` only holds
|
||||
* the visible thread's `cancelRun()`, so the adapter parks a closure here that POSTs that
|
||||
* run's own cancel_id. A list for the same reason as `runOwnerByThreadId`: "__default" is shared.
|
||||
*/
|
||||
serverCancelByThreadId: Record<string, (() => void)[]>;
|
||||
autoTitle: boolean;
|
||||
hfToken: string;
|
||||
modelsError: string | null;
|
||||
|
|
@ -892,7 +934,16 @@ type ChatRuntimeStore = {
|
|||
* consulted when `providerSupportsBuiltinWebFetch` is true.
|
||||
*/
|
||||
webFetchToolsEnabled: boolean;
|
||||
toolStatus: string | null;
|
||||
/**
|
||||
* Live tool status per conversation ("Running Python: ...") with its start time. Keyed by
|
||||
* thread, or one chat's tool call shows above every other composer; the timestamp keeps the
|
||||
* counter running across a thread switch.
|
||||
*/
|
||||
/**
|
||||
* Per-run entries, newest last. Unresolved threads share "__default", so one scalar per key
|
||||
* meant a finishing run's clear removed a sibling's status while its tool was still running.
|
||||
*/
|
||||
toolStatusByThreadId: Record<string, ToolStatusEntry[]>;
|
||||
/** Live stdout/stderr from running tools, keyed by toolCallId. Transient:
|
||||
* appended by tool_output, cleared on tool_end or run end. */
|
||||
toolLiveOutput: Record<string, string>;
|
||||
|
|
@ -959,9 +1010,12 @@ type ChatRuntimeStore = {
|
|||
/** Active model is a block-diffusion model (DiffusionGemma): drives the
|
||||
* denoising-canvas artifact auto-render. */
|
||||
loadedIsDiffusion: boolean;
|
||||
/** Live denoising frame for the in-progress diffusion message. Transient: set
|
||||
* per step, cleared when the run ends, never persisted into the transcript. */
|
||||
activeDiffusionCanvas: DiffusionCanvasFrame | null;
|
||||
/**
|
||||
* Live denoising frame per conversation ("__default" until the id exists). Transient: set per
|
||||
* step, cleared when the run ends, never persisted. Keyed, not global: two denoising chats
|
||||
* overwrote each other's frame, so the visible preview flickered or vanished.
|
||||
*/
|
||||
activeDiffusionCanvasByThreadId: Record<string, DiffusionCanvasFrame>;
|
||||
customContextLength: number | null;
|
||||
/** The pinned context the loaded model used (null = Auto), so dirty-tracking
|
||||
* and a later fit Apply can tell an explicit pin apart from Auto. */
|
||||
|
|
@ -984,14 +1038,13 @@ type ChatRuntimeStore = {
|
|||
pendingAudioBase64: string | null;
|
||||
pendingAudioName: string | null;
|
||||
pendingImageEditReference: PendingImageEditReference | null;
|
||||
contextUsage: {
|
||||
promptTokens: number;
|
||||
completionTokens: number;
|
||||
totalTokens: number;
|
||||
cachedTokens: number;
|
||||
// Anthropic-only; optional so pre-cache-stats persisted entries load.
|
||||
cacheWriteTokens?: number;
|
||||
} | null;
|
||||
contextUsage: ContextUsageSnapshot | null;
|
||||
/**
|
||||
* Per-thread copy of the above, so the bar survives a switch away and back. `contextUsage` is
|
||||
* the VISIBLE conversation's usage and a background run may not write it, so without this a
|
||||
* run finishing off-screen leaves nothing to restore.
|
||||
*/
|
||||
contextUsageByThreadId: Record<string, ContextUsageSnapshot>;
|
||||
modelLoading: boolean;
|
||||
loadingModelPick: LoadingModelPick | null;
|
||||
activeNativePathToken: string | null;
|
||||
|
|
@ -1010,9 +1063,35 @@ type ChatRuntimeStore = {
|
|||
setActivePresetSource: (source: ChatPresetSource) => void;
|
||||
setModels: (models: ChatModelSummary[]) => void;
|
||||
setLoras: (loras: ChatLoraSummary[]) => void;
|
||||
setThreadRunning: (threadId: string, running: boolean) => void;
|
||||
/**
|
||||
* `local` defaults to true, so an unqualified caller still counts for the model-swap gate.
|
||||
* `owner` narrows the clear to the run that set the flag: unresolved thread ids share the
|
||||
* "__default" key, so a blind delete would drop a sibling's live entry. Owners accumulate,
|
||||
* so the flag survives until the last one clears.
|
||||
*/
|
||||
setThreadRunning: (
|
||||
threadId: string,
|
||||
running: boolean,
|
||||
options?: { local?: boolean; owner?: () => void },
|
||||
) => void;
|
||||
/**
|
||||
* Re-key a first turn's run handles once its thread is persisted.
|
||||
*
|
||||
* A run that starts before its id exists files everything under "__default". Nothing moved it
|
||||
* afterwards, so once the user navigated away the sidebar found no run and showed no spinner;
|
||||
* stopChatThread had no handle either and the generation carried on holding a slot.
|
||||
*/
|
||||
adoptDefaultThreadRun: (threadId: string) => void;
|
||||
/**
|
||||
* Which key this run's handles live under now. `adoptDefaultThreadRun` re-keys them mid-run,
|
||||
* so a run that started under "__default" must look its owner up instead of reusing the key
|
||||
* it captured, or its writes and its final clear miss the entries.
|
||||
*/
|
||||
runKeyForOwner: (fallbackKey: string, owner: () => void) => string;
|
||||
registerThreadCancel: (threadId: string, cancel: () => void) => void;
|
||||
clearThreadCancel: (threadId: string) => void;
|
||||
registerThreadServerCancel: (threadId: string, cancel: () => void) => void;
|
||||
clearThreadServerCancel: (threadId: string, cancel?: () => void) => void;
|
||||
setAutoTitle: (enabled: boolean) => void;
|
||||
setHfToken: (token: string) => void;
|
||||
setModelsError: (error: string | null) => void;
|
||||
|
|
@ -1066,7 +1145,15 @@ type ChatRuntimeStore = {
|
|||
setRagAutoInjectMinScore: (score: number) => void;
|
||||
setRagOcrScanned: (enabled: boolean) => void;
|
||||
setRagCaptionFigures: (enabled: boolean) => void;
|
||||
setToolStatus: (status: string | null) => void;
|
||||
/**
|
||||
* `owner` is the run's identity token, as for `setThreadRunning`: unresolved threads share
|
||||
* "__default", so without it one run's cleanup clears a concurrent run's status.
|
||||
*/
|
||||
setToolStatus: (
|
||||
threadId: string,
|
||||
status: string | null,
|
||||
owner?: () => void,
|
||||
) => void;
|
||||
appendToolLiveOutput: (toolCallId: string, text: string) => void;
|
||||
/** Clear one tool's live output, or all when no id is given. */
|
||||
clearToolLiveOutput: (toolCallId?: string) => void;
|
||||
|
|
@ -1075,7 +1162,13 @@ type ChatRuntimeStore = {
|
|||
/** Drop a stale preserved full output (a new run is reusing the id). */
|
||||
clearToolFullOutput: (toolCallId: string) => void;
|
||||
setGeneratingStatus: (status: string | null) => void;
|
||||
setActiveDiffusionCanvas: (canvas: DiffusionCanvasFrame | null) => void;
|
||||
setActiveDiffusionCanvas: (
|
||||
threadId: string | null,
|
||||
canvas: DiffusionCanvasFrame,
|
||||
) => void;
|
||||
/** Drop only `threadId`'s canvas: a run ending in a background chat must not wipe the
|
||||
* frame another chat is still painting. */
|
||||
clearActiveDiffusionCanvasForThread: (threadId: string | null) => void;
|
||||
setAutoHealToolCalls: (enabled: boolean) => void;
|
||||
setNudgeToolCalls: (enabled: boolean) => void;
|
||||
setMaxToolCallsPerMessage: (value: number) => void;
|
||||
|
|
@ -1095,6 +1188,11 @@ type ChatRuntimeStore = {
|
|||
) => void;
|
||||
clearPendingImageEditReference: () => void;
|
||||
setContextUsage: (usage: ChatRuntimeStore["contextUsage"]) => void;
|
||||
/** A finished run's usage, kept per thread so switching back re-applies it. */
|
||||
setThreadContextUsage: (
|
||||
threadId: string,
|
||||
usage: ContextUsageSnapshot,
|
||||
) => void;
|
||||
};
|
||||
|
||||
type PersistedChatSettings = Awaited<
|
||||
|
|
@ -1310,7 +1408,10 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
models: [],
|
||||
loras: [],
|
||||
runningByThreadId: {},
|
||||
localRunByThreadId: {},
|
||||
runOwnerByThreadId: {},
|
||||
cancelByThreadId: {},
|
||||
serverCancelByThreadId: {},
|
||||
autoTitle: false,
|
||||
hfToken: useHfTokenStore.getState().token,
|
||||
modelsError: null,
|
||||
|
|
@ -1374,11 +1475,11 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
),
|
||||
ragOcrScanned: loadBool(CHAT_RAG_OCR_KEY, DEFAULT_RAG_OCR),
|
||||
ragCaptionFigures: loadBool(CHAT_RAG_CAPTION_KEY, DEFAULT_RAG_CAPTION),
|
||||
toolStatus: null,
|
||||
toolStatusByThreadId: {},
|
||||
toolLiveOutput: {},
|
||||
toolFullOutput: {},
|
||||
generatingStatus: null,
|
||||
activeDiffusionCanvas: null,
|
||||
activeDiffusionCanvasByThreadId: {},
|
||||
autoHealToolCalls: true,
|
||||
nudgeToolCalls: true,
|
||||
maxToolCallsPerMessage: 25,
|
||||
|
|
@ -1423,6 +1524,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
pendingAudioName: null,
|
||||
pendingImageEditReference: null,
|
||||
contextUsage: null,
|
||||
contextUsageByThreadId: {},
|
||||
modelLoading: false,
|
||||
loadingModelPick: null,
|
||||
activeNativePathToken: null,
|
||||
|
|
@ -1495,7 +1597,9 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
const checkpointChanged = state.params.checkpoint !== params.checkpoint;
|
||||
return {
|
||||
params,
|
||||
...(checkpointChanged ? { contextUsage: null } : {}),
|
||||
...(checkpointChanged
|
||||
? { contextUsage: null, contextUsageByThreadId: {} }
|
||||
: {}),
|
||||
};
|
||||
}),
|
||||
setCustomPresets: (customPresets) =>
|
||||
|
|
@ -1518,16 +1622,94 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
}),
|
||||
setModels: (models) => set({ models }),
|
||||
setLoras: (loras) => set({ loras }),
|
||||
setThreadRunning: (threadId, running) =>
|
||||
setThreadRunning: (threadId, running, options) =>
|
||||
set((state) => {
|
||||
const next = { ...state.runningByThreadId };
|
||||
const nextLocal = { ...state.localRunByThreadId };
|
||||
const nextOwner = { ...state.runOwnerByThreadId };
|
||||
const owners = state.runOwnerByThreadId[threadId] ?? [];
|
||||
const local = options?.local !== false;
|
||||
if (running) {
|
||||
next[threadId] = true;
|
||||
if (options?.owner) {
|
||||
nextOwner[threadId] = [...owners, { owner: options.owner, local }];
|
||||
}
|
||||
// Any local owner keeps the key counted by the model-swap gate, so an external run
|
||||
// joining a shared key must not clear a sibling's flag.
|
||||
if (local) {
|
||||
nextLocal[threadId] = true;
|
||||
} else if (!owners.some((o) => o.local)) {
|
||||
delete nextLocal[threadId];
|
||||
}
|
||||
} else {
|
||||
delete next[threadId];
|
||||
const remaining = options?.owner
|
||||
? owners.filter((o) => o.owner !== options.owner)
|
||||
: [];
|
||||
// An owner missing from the list was already cleared, or the key belongs to siblings
|
||||
// only: either way this run must change nothing.
|
||||
if (options?.owner && remaining.length === owners.length) return state;
|
||||
// An ownerless clear predates per-run tracking, so it must not speak for runs that
|
||||
// own the key: leave them to clear themselves.
|
||||
if (!options?.owner && owners.length > 0) return state;
|
||||
if (remaining.length > 0) {
|
||||
nextOwner[threadId] = remaining;
|
||||
if (remaining.some((o) => o.local)) {
|
||||
nextLocal[threadId] = true;
|
||||
} else {
|
||||
delete nextLocal[threadId];
|
||||
}
|
||||
} else {
|
||||
delete next[threadId];
|
||||
delete nextLocal[threadId];
|
||||
delete nextOwner[threadId];
|
||||
}
|
||||
}
|
||||
return { runningByThreadId: next };
|
||||
return {
|
||||
runningByThreadId: next,
|
||||
localRunByThreadId: nextLocal,
|
||||
runOwnerByThreadId: nextOwner,
|
||||
};
|
||||
}),
|
||||
adoptDefaultThreadRun: (threadId) =>
|
||||
set((state) => {
|
||||
const key = "__default";
|
||||
if (!threadId || threadId === key) return state;
|
||||
// Two first turns can share "__default", and nothing links a run there to the thread being
|
||||
// persisted. Moving the arrays wholesale handed this thread the sibling's owner and stop
|
||||
// handle too, so stopping one aborted both. Adopt only when the key holds a single run.
|
||||
if ((state.runOwnerByThreadId[key]?.length ?? 0) > 1) return state;
|
||||
// Only the transient run maps move. Anything already filed under the real id wins,
|
||||
// since that is a later, better-identified run.
|
||||
const moved: Partial<ChatRuntimeStore> = {};
|
||||
const move = <T,>(
|
||||
map: Record<string, T>,
|
||||
name: keyof ChatRuntimeStore,
|
||||
) => {
|
||||
const entry = map[key];
|
||||
if (entry === undefined || map[threadId] !== undefined) return;
|
||||
const next = { ...map };
|
||||
delete next[key];
|
||||
next[threadId] = entry;
|
||||
(moved as Record<string, unknown>)[name as string] = next;
|
||||
};
|
||||
move(state.runningByThreadId, "runningByThreadId");
|
||||
move(state.localRunByThreadId, "localRunByThreadId");
|
||||
move(state.runOwnerByThreadId, "runOwnerByThreadId");
|
||||
move(state.cancelByThreadId, "cancelByThreadId");
|
||||
move(state.serverCancelByThreadId, "serverCancelByThreadId");
|
||||
move(state.toolStatusByThreadId, "toolStatusByThreadId");
|
||||
move(
|
||||
state.activeDiffusionCanvasByThreadId,
|
||||
"activeDiffusionCanvasByThreadId",
|
||||
);
|
||||
return Object.keys(moved).length > 0 ? moved : state;
|
||||
}),
|
||||
runKeyForOwner: (fallbackKey, owner) => {
|
||||
for (const [key, entries] of Object.entries(get().runOwnerByThreadId)) {
|
||||
if (entries.some((e) => e.owner === owner)) return key;
|
||||
}
|
||||
return fallbackKey;
|
||||
},
|
||||
registerThreadCancel: (threadId, cancel) =>
|
||||
set((state) => {
|
||||
const next = { ...state.cancelByThreadId };
|
||||
|
|
@ -1541,6 +1723,29 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
delete next[threadId];
|
||||
return { cancelByThreadId: next };
|
||||
}),
|
||||
registerThreadServerCancel: (threadId, cancel) =>
|
||||
set((state) => {
|
||||
const next = { ...state.serverCancelByThreadId };
|
||||
next[threadId] = [...(state.serverCancelByThreadId[threadId] ?? []), cancel];
|
||||
return { serverCancelByThreadId: next };
|
||||
}),
|
||||
// `cancel` narrows removal to the run that registered it: unresolved thread ids share the
|
||||
// "__default" key, so a blind delete would drop a live sibling.
|
||||
clearThreadServerCancel: (threadId, cancel) =>
|
||||
set((state) => {
|
||||
const current = state.serverCancelByThreadId[threadId];
|
||||
if (current === undefined) return state;
|
||||
const remaining =
|
||||
cancel === undefined ? [] : current.filter((c) => c !== cancel);
|
||||
if (remaining.length === current.length) return state;
|
||||
const next = { ...state.serverCancelByThreadId };
|
||||
if (remaining.length > 0) {
|
||||
next[threadId] = remaining;
|
||||
} else {
|
||||
delete next[threadId];
|
||||
}
|
||||
return { serverCancelByThreadId: next };
|
||||
}),
|
||||
setAutoTitle: (autoTitle) =>
|
||||
set((state) => {
|
||||
setScalarSettingVersion("autoTitle", autoTitle, state.autoTitle);
|
||||
|
|
@ -1588,14 +1793,24 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
maxTokens: nextMaxTokens,
|
||||
},
|
||||
activeGgufVariant: ggufVariant ?? null,
|
||||
...(checkpointChanged ? { contextUsage: null } : {}),
|
||||
...(checkpointChanged
|
||||
? { contextUsage: null, contextUsageByThreadId: {} }
|
||||
: {}),
|
||||
// Switching to an external provider disables Deep Research, which only
|
||||
// applies to the local base model.
|
||||
...(isExternalModelId(modelId) ? { deepResearchEnabled: false } : {}),
|
||||
};
|
||||
}),
|
||||
// Re-apply the incoming thread's own usage rather than blanking the bar: a run that finished
|
||||
// in the background never wrote the visible value, and a still-mounted runtime skips the
|
||||
// history loader on the way back.
|
||||
setActiveThreadId: (activeThreadId) =>
|
||||
set({ activeThreadId, contextUsage: null }),
|
||||
set((state) => ({
|
||||
activeThreadId,
|
||||
contextUsage: activeThreadId
|
||||
? (state.contextUsageByThreadId[activeThreadId] ?? null)
|
||||
: null,
|
||||
})),
|
||||
setActiveProjectId: (activeProjectId) => set({ activeProjectId }),
|
||||
setIncognito: (incognito) => {
|
||||
if (incognito) saveBool(CHAT_DEEP_RESEARCH_ENABLED_KEY, false);
|
||||
|
|
@ -1626,6 +1841,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
ggufNativeContextLength: null,
|
||||
modelRequiresTrustRemoteCode: false,
|
||||
contextUsage: null,
|
||||
contextUsageByThreadId: {},
|
||||
supportsReasoning: false,
|
||||
reasoningAlwaysOn: false,
|
||||
reasoningEnabled: true,
|
||||
|
|
@ -1647,10 +1863,10 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
webFetchToolsEnabled: false,
|
||||
// Only the per-session enable pill resets; source/mode/top_k persist.
|
||||
ragEnabled: false,
|
||||
toolStatus: null,
|
||||
toolStatusByThreadId: {},
|
||||
toolLiveOutput: {},
|
||||
toolFullOutput: {},
|
||||
activeDiffusionCanvas: null,
|
||||
activeDiffusionCanvasByThreadId: {},
|
||||
kvCacheDtype: null,
|
||||
loadedKvCacheDtype: null,
|
||||
speculativeType: readPersistedSpeculativeType(),
|
||||
|
|
@ -1945,7 +2161,31 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
saveBool(CHAT_RAG_CAPTION_KEY, ragCaptionFigures);
|
||||
return { ragCaptionFigures };
|
||||
}),
|
||||
setToolStatus: (toolStatus) => set({ toolStatus }),
|
||||
setToolStatus: (threadId, status, owner) =>
|
||||
set((state) => {
|
||||
const next = { ...state.toolStatusByThreadId };
|
||||
const entries = state.toolStatusByThreadId[threadId] ?? [];
|
||||
const mine = entries.find((e) => e.owner === owner);
|
||||
if (!status) {
|
||||
// Drop only this run's entry: a sibling behind the same key may still be running a tool,
|
||||
// and its status has to survive this clear.
|
||||
if (mine === undefined) return state;
|
||||
const rest = entries.filter((e) => e !== mine);
|
||||
if (rest.length > 0) {
|
||||
next[threadId] = rest;
|
||||
} else {
|
||||
delete next[threadId];
|
||||
}
|
||||
} else {
|
||||
// Same text from the same run means the same call, so keep startedAt: only a new tool restarts it.
|
||||
if (mine?.status === status) return state;
|
||||
const entry = { status, startedAt: Date.now(), owner };
|
||||
next[threadId] = mine
|
||||
? entries.map((e) => (e === mine ? entry : e))
|
||||
: [...entries, entry];
|
||||
}
|
||||
return { toolStatusByThreadId: next };
|
||||
}),
|
||||
appendToolLiveOutput: (toolCallId, text) =>
|
||||
set((state) => ({
|
||||
toolLiveOutput: {
|
||||
|
|
@ -1983,8 +2223,21 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
delete next[toolCallId];
|
||||
return { toolLiveOutput: next };
|
||||
}),
|
||||
setActiveDiffusionCanvas: (activeDiffusionCanvas) =>
|
||||
set({ activeDiffusionCanvas }),
|
||||
setActiveDiffusionCanvas: (threadId, canvas) =>
|
||||
set((state) => ({
|
||||
activeDiffusionCanvasByThreadId: {
|
||||
...state.activeDiffusionCanvasByThreadId,
|
||||
[threadId || "__default"]: canvas,
|
||||
},
|
||||
})),
|
||||
clearActiveDiffusionCanvasForThread: (threadId) =>
|
||||
set((state) => {
|
||||
const key = threadId || "__default";
|
||||
if (state.activeDiffusionCanvasByThreadId[key] === undefined) return state;
|
||||
const next = { ...state.activeDiffusionCanvasByThreadId };
|
||||
delete next[key];
|
||||
return { activeDiffusionCanvasByThreadId: next };
|
||||
}),
|
||||
setGeneratingStatus: (generatingStatus) => set({ generatingStatus }),
|
||||
setAutoHealToolCalls: (autoHealToolCalls) =>
|
||||
set((state) => {
|
||||
|
|
@ -2050,7 +2303,27 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
set({ pendingImageEditReference }),
|
||||
clearPendingImageEditReference: () =>
|
||||
set({ pendingImageEditReference: null }),
|
||||
setContextUsage: (contextUsage) => set({ contextUsage }),
|
||||
// Write through to the visible thread's own entry, so a value restored by the history loader
|
||||
// survives a switch away and back: that loader runs once per mount and setActiveThreadId
|
||||
// reads the map, so without this the bar goes blank on return.
|
||||
setContextUsage: (contextUsage) =>
|
||||
set((state) => {
|
||||
if (!state.activeThreadId) return { contextUsage };
|
||||
const next = { ...state.contextUsageByThreadId };
|
||||
if (contextUsage) {
|
||||
next[state.activeThreadId] = contextUsage;
|
||||
} else {
|
||||
delete next[state.activeThreadId];
|
||||
}
|
||||
return { contextUsage, contextUsageByThreadId: next };
|
||||
}),
|
||||
setThreadContextUsage: (threadId, usage) =>
|
||||
set((state) => ({
|
||||
contextUsageByThreadId: {
|
||||
...state.contextUsageByThreadId,
|
||||
[threadId]: usage,
|
||||
},
|
||||
})),
|
||||
}));
|
||||
|
||||
// Mirror token edits made through the shared store (e.g. Unsloth's field).
|
||||
|
|
|
|||
|
|
@ -0,0 +1,69 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { create } from "zustand";
|
||||
|
||||
type Resolver = (confirmed: boolean) => void;
|
||||
|
||||
/** What confirming does to the model: reload it, or leave none loaded. */
|
||||
export type StopRunningChatsEffect = "reload" | "unload";
|
||||
|
||||
// One at a time: a new request declines any pending one so no promise leaks.
|
||||
let pendingResolver: Resolver | null = null;
|
||||
|
||||
interface StopRunningChatsDialogStore {
|
||||
open: boolean;
|
||||
/** How many conversations the pending action would stop. */
|
||||
count: number;
|
||||
/** Titles of those conversations, when known, for the dialog body. */
|
||||
titles: string[];
|
||||
/** What the user is about to do, e.g. "Loading a different model". */
|
||||
action: string;
|
||||
/** The set includes an embeddings/completions/audio request, which is not a chat. */
|
||||
hasNonChat: boolean;
|
||||
/** Ejecting leaves no model loaded, so it must not be described as a reload. */
|
||||
effect: StopRunningChatsEffect;
|
||||
requestConfirm: (args: {
|
||||
count: number;
|
||||
titles?: string[];
|
||||
action?: string;
|
||||
hasNonChat?: boolean;
|
||||
effect?: StopRunningChatsEffect;
|
||||
}) => Promise<boolean>;
|
||||
resolve: (confirmed: boolean) => void;
|
||||
}
|
||||
|
||||
export const useStopRunningChatsDialogStore =
|
||||
create<StopRunningChatsDialogStore>()((set) => ({
|
||||
open: false,
|
||||
count: 0,
|
||||
titles: [],
|
||||
action: "",
|
||||
hasNonChat: false,
|
||||
effect: "reload",
|
||||
requestConfirm: ({
|
||||
count,
|
||||
titles = [],
|
||||
action = "",
|
||||
hasNonChat = false,
|
||||
effect = "reload",
|
||||
}) =>
|
||||
new Promise<boolean>((resolve) => {
|
||||
pendingResolver?.(false);
|
||||
pendingResolver = resolve;
|
||||
set({ open: true, count, titles, action, hasNonChat, effect });
|
||||
}),
|
||||
resolve: (confirmed) => {
|
||||
const resolver = pendingResolver;
|
||||
pendingResolver = null;
|
||||
set({
|
||||
open: false,
|
||||
count: 0,
|
||||
titles: [],
|
||||
action: "",
|
||||
hasNonChat: false,
|
||||
effect: "reload",
|
||||
});
|
||||
resolver?.(confirmed);
|
||||
},
|
||||
}));
|
||||
18
studio/frontend/src/features/chat/tool-approval.ts
Normal file
18
studio/frontend/src/features/chat/tool-approval.ts
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"use client";
|
||||
|
||||
import { useChatRuntimeStore } from "./stores/chat-runtime-store";
|
||||
|
||||
/**
|
||||
* True while this card's call is parked on the Allow / Deny prompt, so it can say it is
|
||||
* waiting rather than counting up "Running". Set when the backend gates the call.
|
||||
*/
|
||||
export function useToolAwaitingApproval(toolCallId?: string): boolean {
|
||||
return useChatRuntimeStore(
|
||||
(s) =>
|
||||
!!toolCallId &&
|
||||
Object.prototype.hasOwnProperty.call(s.toolConfirmations, toolCallId),
|
||||
);
|
||||
}
|
||||
|
|
@ -3,6 +3,7 @@
|
|||
|
||||
"use client";
|
||||
|
||||
import { useAuiState } from "@assistant-ui/react";
|
||||
import { createContext, useContext } from "react";
|
||||
import type { ModelType } from "./types";
|
||||
|
||||
|
|
@ -20,10 +21,58 @@ export function toolPaneScope(modelType?: ModelType, pairId?: string): string {
|
|||
return `${modelType ?? "base"}\u0000${pairId ?? ""}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Narrow a pane scope to one conversation: two threads in a pane can both be mid "call_0",
|
||||
* so without the thread in the key they share a store entry and swap outputs.
|
||||
*/
|
||||
export function toolThreadScope(paneScope: string, threadId?: string): string {
|
||||
return `${paneScope}\u0000${threadId ?? ""}`;
|
||||
}
|
||||
|
||||
export const ToolPaneScopeContext = createContext<string>(toolPaneScope());
|
||||
|
||||
/**
|
||||
* Store-key scope for the conversation this component renders in, taken from the surrounding
|
||||
* runtime so reader and writer agree without a prop.
|
||||
*
|
||||
* `remoteId`, not `id`: the adapter gets `unstable_threadId`, which assistant-ui sources from
|
||||
* `remoteId`, and an uninitialized thread has `id` but no `remoteId`. Reading `id` split the
|
||||
* keys apart for the first turn of every New Chat, so live tool output never reached the card.
|
||||
*/
|
||||
export function useToolPaneScope(): string {
|
||||
return useContext(ToolPaneScopeContext);
|
||||
const paneScope = useContext(ToolPaneScopeContext);
|
||||
const threadId = useAuiState(({ threadListItem }) => threadListItem.remoteId);
|
||||
return toolThreadScope(paneScope, threadId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Read a tool-output map for one call, tolerating a run that started before its thread had an id.
|
||||
*
|
||||
* The adapter captures its scope once at run start, so a first turn writes under the unresolved
|
||||
* scope for its whole life. The autosave can assign `remoteId` mid-run, which moves this
|
||||
* component's key but not the writer's, and the card went blank. Falling back to the pane-wide
|
||||
* scope keeps those entries reachable; only an unpersisted first turn can be filed there.
|
||||
*/
|
||||
/** The scope a run that started before its thread had an id writes under. */
|
||||
export function useUnresolvedToolPaneScope(): string {
|
||||
return toolThreadScope(useContext(ToolPaneScopeContext), undefined);
|
||||
}
|
||||
|
||||
export function useToolOutputFor(
|
||||
map: Record<string, string>,
|
||||
paneScope: string,
|
||||
toolCallId: string,
|
||||
): string {
|
||||
// Unconditional: hooks cannot sit behind the early return below.
|
||||
const unresolvedScope = useUnresolvedToolPaneScope();
|
||||
// Only a thread mid-run can be the one that just gained its id. Local ids repeat
|
||||
// ("call_0"), so an unconditional fallback showed a live first turn's stdout in every
|
||||
// older conversation whose own entry had been cleared.
|
||||
const isRunning = useAuiState(({ thread }) => thread.isRunning);
|
||||
const own = map[toolOutputKey(paneScope, toolCallId)];
|
||||
if (own !== undefined) return own;
|
||||
if (!isRunning) return "";
|
||||
return map[toolOutputKey(unresolvedScope, toolCallId)] ?? "";
|
||||
}
|
||||
|
||||
/** Store key for the live/full tool output maps: pane scope + tool call id. */
|
||||
|
|
|
|||
|
|
@ -35,6 +35,11 @@ export interface ListLorasResponse {
|
|||
|
||||
export interface LoadModelRequest {
|
||||
model_path: string;
|
||||
/**
|
||||
* Stop any chats still generating instead of getting a 409: a load replaces the single
|
||||
* llama-server they all decode on. Set only after the user confirms.
|
||||
*/
|
||||
force_cancel_active?: boolean;
|
||||
nativePathLease?: string | null;
|
||||
hf_token: string | null;
|
||||
max_seq_length: number;
|
||||
|
|
@ -201,6 +206,9 @@ export interface LoadModelResponse {
|
|||
|
||||
export interface UnloadModelRequest {
|
||||
model_path: string;
|
||||
/** Stop any chats still generating instead of getting a 409: the unload takes down the
|
||||
* llama-server they all decode on. */
|
||||
force_cancel_active?: boolean;
|
||||
}
|
||||
|
||||
export interface InferenceStatusResponse {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,100 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { getActiveGenerations } from "../api/chat-api";
|
||||
import { useChatRuntimeStore } from "../stores/chat-runtime-store";
|
||||
import {
|
||||
type StopRunningChatsEffect,
|
||||
useStopRunningChatsDialogStore,
|
||||
} from "../stores/stop-running-chats-dialog-store";
|
||||
import { listStoredChatThreads } from "./chat-history-storage";
|
||||
|
||||
export interface StopRunningChatsDecision {
|
||||
/** False when the user chose to keep generating; the caller must not load. */
|
||||
proceed: boolean;
|
||||
/** Pass as `force_cancel_active`. True only after an explicit confirmation, so the backend's 409 still guards every other caller. */
|
||||
forceCancelActive: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Gate a model load / reload on the chats still generating: they share one llama-server,
|
||||
* so a reload ends all of them. Ask first, then let the backend cancel them once the load
|
||||
* is past preflight. External-provider chats are left out of both.
|
||||
*/
|
||||
export async function confirmStopRunningChatsIfNeeded(
|
||||
action = "Loading a different model",
|
||||
effect: StopRunningChatsEffect = "reload",
|
||||
): Promise<StopRunningChatsDecision> {
|
||||
// Local runs only: an external-provider chat is not stopped by the swap, so counting it
|
||||
// would block a safe load behind a dialog. The backend excludes them for the same reason.
|
||||
const { runningByThreadId, localRunByThreadId } =
|
||||
useChatRuntimeStore.getState();
|
||||
let running = Object.entries(runningByThreadId)
|
||||
.filter(([threadId, on]) => on && localRunByThreadId[threadId])
|
||||
.map(([threadId]) => threadId);
|
||||
let count = running.length;
|
||||
let hasNonChat = false;
|
||||
|
||||
// Always merge the backend snapshot: runningByThreadId is this tab's memory, empty after a
|
||||
// reload and blind to a second tab, while force_cancel_active cancels every backend run.
|
||||
// The union stays local-only, since external-provider runs are never in it.
|
||||
try {
|
||||
const active = await getActiveGenerations();
|
||||
const entries = active.active ?? [];
|
||||
const merged = new Set(running);
|
||||
for (const threadId of active.thread_ids ?? []) {
|
||||
merged.add(threadId);
|
||||
}
|
||||
running = [...merged];
|
||||
// Count conversations, not handles: one chat holds several at once while a tool
|
||||
// continuation registers its next leg before the previous unwinds, and active.count
|
||||
// counts those separately. A first turn started before its id was persisted has no
|
||||
// id to merge, so add those back or the prompt names fewer chats than will stop.
|
||||
const unnamed = entries.filter((entry) => !entry.thread_id).length;
|
||||
count = entries.length
|
||||
? running.length + unnamed
|
||||
: Math.max(active.count ?? 0, running.length);
|
||||
// Embeddings / completions / audio share the model but are not conversations, so the
|
||||
// prompt must not offer to stop chats that do not exist.
|
||||
hasNonChat = entries.some((entry) => (entry.kind ?? "chat") !== "chat");
|
||||
} catch {
|
||||
// Backend unreachable / older build: fall back to the local map only.
|
||||
}
|
||||
|
||||
if (count === 0) {
|
||||
return { proceed: true, forceCancelActive: false };
|
||||
}
|
||||
|
||||
let titles: string[] = [];
|
||||
try {
|
||||
const threads = await listStoredChatThreads();
|
||||
const byId = new Map(threads.map((t) => [t.id, t]));
|
||||
// A compare conversation runs two pane threads, and the sidebar and the route both treat
|
||||
// it as one chat. Counting the raw ids asked to stop two and listed its title twice. Fold
|
||||
// panes onto their pairId, keeping the backend's count when it is higher.
|
||||
const seen = new Set<string>();
|
||||
for (const id of running) {
|
||||
const thread = byId.get(id);
|
||||
const key = thread?.pairId ?? id;
|
||||
if (seen.has(key)) continue;
|
||||
seen.add(key);
|
||||
titles.push(thread?.title || "Untitled chat");
|
||||
}
|
||||
count = Math.max(seen.size, count - (running.length - seen.size));
|
||||
} catch {
|
||||
// Titles are decoration; the count alone is enough to make the choice.
|
||||
titles = [];
|
||||
}
|
||||
|
||||
const confirmed = await useStopRunningChatsDialogStore
|
||||
.getState()
|
||||
.requestConfirm({ count, titles, action, hasNonChat, effect });
|
||||
|
||||
if (!confirmed) {
|
||||
return { proceed: false, forceCancelActive: false };
|
||||
}
|
||||
|
||||
// Deliberately no local stop: the backend holds the cancel until the load clears preflight,
|
||||
// so stopping now would truncate every chat even for a rejected load.
|
||||
return { proceed: true, forceCancelActive: true };
|
||||
}
|
||||
|
|
@ -1,8 +1,17 @@
|
|||
export const PROMPT_QUEUE_STOP_EVENT = "unsloth:prompt-queue-stop";
|
||||
|
||||
export function requestPromptQueueStop() {
|
||||
export interface PromptQueueStopOptions {
|
||||
/** Also cancel the prompt the queue already dispatched. Navigation passes `false` to
|
||||
* leave it generating; an explicit stop passes `true` (the default). */
|
||||
cancelActiveRun?: boolean;
|
||||
}
|
||||
|
||||
export function requestPromptQueueStop(options: PromptQueueStopOptions = {}) {
|
||||
if (typeof window === "undefined") {
|
||||
return;
|
||||
}
|
||||
window.dispatchEvent(new Event(PROMPT_QUEUE_STOP_EVENT));
|
||||
const { cancelActiveRun = true } = options;
|
||||
window.dispatchEvent(
|
||||
new CustomEvent(PROMPT_QUEUE_STOP_EVENT, { detail: { cancelActiveRun } }),
|
||||
);
|
||||
}
|
||||
|
|
|
|||
39
studio/frontend/src/features/chat/utils/stop-chat-thread.ts
Normal file
39
studio/frontend/src/features/chat/utils/stop-chat-thread.ts
Normal file
|
|
@ -0,0 +1,39 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { useChatRuntimeStore } from "../stores/chat-runtime-store";
|
||||
|
||||
/**
|
||||
* Stop one conversation's generation, visible or not. Returns true if a stop was dispatched.
|
||||
*
|
||||
* `cancelByThreadId` is assistant-ui's `cancelRun()`, registered only for the thread on screen;
|
||||
* `serverCancelByThreadId` is registered for every run and POSTs that run's own `cancel_id`, so
|
||||
* it is the only handle a background conversation has. Both are per-run. Runs with an unresolved
|
||||
* thread id share the "__default" key, so stop every handle filed under it.
|
||||
*/
|
||||
export function stopChatThread(threadId: string | null | undefined): boolean {
|
||||
if (!threadId) return false;
|
||||
const { runningByThreadId, cancelByThreadId, serverCancelByThreadId } =
|
||||
useChatRuntimeStore.getState();
|
||||
if (!runningByThreadId[threadId]) return false;
|
||||
let stopped = false;
|
||||
try {
|
||||
const cancel = cancelByThreadId[threadId];
|
||||
if (cancel) {
|
||||
cancel();
|
||||
stopped = true;
|
||||
}
|
||||
} catch {
|
||||
// The run may have ended between the read above and this call.
|
||||
}
|
||||
// Also after cancelRun(): a proxy that swallows the fetch abort leaves the backend decoding.
|
||||
for (const serverCancel of serverCancelByThreadId[threadId] ?? []) {
|
||||
try {
|
||||
serverCancel();
|
||||
stopped = true;
|
||||
} catch {
|
||||
// Same as above.
|
||||
}
|
||||
}
|
||||
return stopped;
|
||||
}
|
||||
|
|
@ -16,14 +16,19 @@ interface InstallLatestTransformersResponse {
|
|||
latest_version?: string | null;
|
||||
}
|
||||
|
||||
/** Consented install of the latest transformers into the sidecar; synchronous, can take minutes. */
|
||||
/** Consented install of the latest transformers into the sidecar; synchronous, can take minutes.
|
||||
*
|
||||
* `forceCancelActive` carries the answer the user already gave the model swap's "stop N
|
||||
* chats" prompt: without it the install 409s while those chats run, and nothing between the
|
||||
* two dialogs stops them. Only ever true after that confirmation. */
|
||||
export async function installLatestTransformers(
|
||||
version: string,
|
||||
forceCancelActive = false,
|
||||
): Promise<InstallLatestTransformersResponse> {
|
||||
const response = await authFetch("/api/inference/install-latest-transformers", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ version }),
|
||||
body: JSON.stringify({ version, force_cancel_active: forceCancelActive }),
|
||||
});
|
||||
if (!response.ok) {
|
||||
throw new Error(await readFastApiError(response));
|
||||
|
|
|
|||
|
|
@ -10,6 +10,9 @@ interface ConfirmArgs {
|
|||
upgrade: TransformersUpgradeInfo | null | undefined;
|
||||
/** When no release is installable, offer continuing into the caller's custom-code gate. */
|
||||
trustRemoteCodeFallback?: boolean;
|
||||
/** The caller already confirmed the swap's "stop N chats" prompt: carry it into
|
||||
* the install, which otherwise 409s on those same chats with no way forward. */
|
||||
forceCancelActive?: boolean;
|
||||
}
|
||||
|
||||
/** Pause a load needing a newer transformers on the consent dialog and run the install.
|
||||
|
|
@ -18,11 +21,13 @@ export async function confirmTransformersUpgradeIfNeeded({
|
|||
modelName,
|
||||
upgrade,
|
||||
trustRemoteCodeFallback,
|
||||
forceCancelActive,
|
||||
}: ConfirmArgs): Promise<boolean> {
|
||||
if (!upgrade) return true;
|
||||
return useTransformersUpgradeDialogStore
|
||||
.getState()
|
||||
.requestConsent(modelName, upgrade, {
|
||||
trustRemoteCodeFallback: Boolean(trustRemoteCodeFallback),
|
||||
forceCancelActive: Boolean(forceCancelActive),
|
||||
});
|
||||
}
|
||||
|
|
|
|||
|
|
@ -18,6 +18,9 @@ interface TransformersUpgradeDialogStore {
|
|||
errorMessage: string | null;
|
||||
/** Model ships custom code; without a PyPI install the load may fall back to trust_remote_code. */
|
||||
trustRemoteCodeFallback: boolean;
|
||||
/** The caller already confirmed the model swap's "stop N chats" prompt, so the install
|
||||
* may stop them too; without it the install 409s and Retry can never succeed. */
|
||||
forceCancelActive: boolean;
|
||||
/** True once this consent's install completed. The install unloads the previous
|
||||
* model before swapping, so the caller must treat it as already unloaded; the
|
||||
* custom-code fallback resolves true without installing and leaves it loaded. */
|
||||
|
|
@ -34,7 +37,7 @@ interface TransformersUpgradeDialogStore {
|
|||
requestConsent: (
|
||||
modelName: string,
|
||||
upgrade: TransformersUpgradeInfo,
|
||||
options?: { trustRemoteCodeFallback?: boolean },
|
||||
options?: { trustRemoteCodeFallback?: boolean; forceCancelActive?: boolean },
|
||||
) => Promise<boolean>;
|
||||
/** Accept/Retry: run the install; on success resolve(true) and close. */
|
||||
install: () => Promise<void>;
|
||||
|
|
@ -49,6 +52,7 @@ export const useTransformersUpgradeDialogStore =
|
|||
phase: "consent",
|
||||
errorMessage: null,
|
||||
trustRemoteCodeFallback: false,
|
||||
forceCancelActive: false,
|
||||
installRan: false,
|
||||
serverUnloadedChat: false,
|
||||
requestConsent: (modelName, upgrade, options) =>
|
||||
|
|
@ -62,6 +66,7 @@ export const useTransformersUpgradeDialogStore =
|
|||
phase: "consent",
|
||||
errorMessage: null,
|
||||
trustRemoteCodeFallback: Boolean(options?.trustRemoteCodeFallback),
|
||||
forceCancelActive: Boolean(options?.forceCancelActive),
|
||||
installRan: false,
|
||||
});
|
||||
}),
|
||||
|
|
@ -71,14 +76,14 @@ export const useTransformersUpgradeDialogStore =
|
|||
return value;
|
||||
},
|
||||
install: async () => {
|
||||
const { upgrade, phase } = get();
|
||||
const { upgrade, phase, forceCancelActive } = get();
|
||||
const version = upgrade?.pypi_version;
|
||||
if (!version || phase === "installing") return;
|
||||
const requestResolver = pendingResolver;
|
||||
set({ phase: "installing", errorMessage: null });
|
||||
let result: Awaited<ReturnType<typeof installLatestTransformers>>;
|
||||
try {
|
||||
result = await installLatestTransformers(version);
|
||||
result = await installLatestTransformers(version, forceCancelActive);
|
||||
// Latch the server-side unload IMMEDIATELY, before any resolver-identity
|
||||
// guard: even a superseded consent's install may have unloaded the chat
|
||||
// model, and the signal must survive for whichever load consumes it next.
|
||||
|
|
@ -133,6 +138,7 @@ export const useTransformersUpgradeDialogStore =
|
|||
phase: "consent",
|
||||
errorMessage: null,
|
||||
trustRemoteCodeFallback: false,
|
||||
forceCancelActive: false,
|
||||
});
|
||||
resolver?.(installed);
|
||||
},
|
||||
|
|
|
|||
|
|
@ -37,6 +37,8 @@ export const ar = {
|
|||
navigation: {
|
||||
newChat: "محادثة جديدة",
|
||||
returnToChat: "العودة إلى المحادثة",
|
||||
returnToChats: "العودة إلى {count} محادثات",
|
||||
chatGenerating: "جارٍ الإنشاء",
|
||||
compare: "مقارنة",
|
||||
search: "بحث",
|
||||
hub: "مركز النماذج",
|
||||
|
|
|
|||
|
|
@ -37,6 +37,8 @@ export const de = {
|
|||
navigation: {
|
||||
newChat: "Neuer Chat",
|
||||
returnToChat: "Zurück zum Chat",
|
||||
returnToChats: "Zurück zu {count} Chats",
|
||||
chatGenerating: "Wird generiert",
|
||||
compare: "Vergleichen",
|
||||
search: "Suchen",
|
||||
hub: "Modell-Hub",
|
||||
|
|
|
|||
|
|
@ -34,6 +34,8 @@ export const en = {
|
|||
navigation: {
|
||||
newChat: "New chat",
|
||||
returnToChat: "Return to Chat",
|
||||
returnToChats: "Return to {count} Chats",
|
||||
chatGenerating: "Generating",
|
||||
compare: "Compare",
|
||||
search: "Search",
|
||||
hub: "Model hub",
|
||||
|
|
|
|||
|
|
@ -37,6 +37,8 @@ export const es = {
|
|||
navigation: {
|
||||
newChat: "Nuevo chat",
|
||||
returnToChat: "Volver al chat",
|
||||
returnToChats: "Volver a {count} chats",
|
||||
chatGenerating: "Generando",
|
||||
compare: "Comparar",
|
||||
search: "Buscar",
|
||||
hub: "Centro de modelos",
|
||||
|
|
|
|||
|
|
@ -37,6 +37,8 @@ export const fr = {
|
|||
navigation: {
|
||||
newChat: "Nouvelle discussion",
|
||||
returnToChat: "Retour à la discussion",
|
||||
returnToChats: "Retour à {count} discussions",
|
||||
chatGenerating: "Génération en cours",
|
||||
compare: "Comparer",
|
||||
search: "Rechercher",
|
||||
hub: "Hub de modèles",
|
||||
|
|
|
|||
|
|
@ -37,6 +37,8 @@ export const hi = {
|
|||
navigation: {
|
||||
newChat: "नई चैट",
|
||||
returnToChat: "चैट पर लौटें",
|
||||
returnToChats: "{count} चैट पर लौटें",
|
||||
chatGenerating: "जनरेट हो रहा है",
|
||||
compare: "तुलना करें",
|
||||
search: "खोजें",
|
||||
hub: "मॉडल हब",
|
||||
|
|
|
|||
|
|
@ -38,6 +38,8 @@ export const ja = {
|
|||
navigation: {
|
||||
newChat: "新規チャット",
|
||||
returnToChat: "チャットに戻る",
|
||||
returnToChats: "{count} 件のチャットに戻る",
|
||||
chatGenerating: "生成中",
|
||||
compare: "比較",
|
||||
search: "検索",
|
||||
hub: "モデルハブ",
|
||||
|
|
|
|||
|
|
@ -37,6 +37,8 @@ export const ko = {
|
|||
navigation: {
|
||||
newChat: "새 채팅",
|
||||
returnToChat: "채팅으로 돌아가기",
|
||||
returnToChats: "채팅 {count}개로 돌아가기",
|
||||
chatGenerating: "생성 중",
|
||||
compare: "비교",
|
||||
search: "검색",
|
||||
hub: "모델 허브",
|
||||
|
|
|
|||
|
|
@ -37,6 +37,8 @@ export const ptBR = {
|
|||
navigation: {
|
||||
newChat: "Novo Chat",
|
||||
returnToChat: "Retornar ao Chat",
|
||||
returnToChats: "Retornar a {count} chats",
|
||||
chatGenerating: "Gerando",
|
||||
compare: "Comparar",
|
||||
search: "Buscar",
|
||||
hub: "Hub de modelos",
|
||||
|
|
|
|||
|
|
@ -37,6 +37,8 @@ export const ru = {
|
|||
navigation: {
|
||||
newChat: "Новый чат",
|
||||
returnToChat: "Вернуться к чату",
|
||||
returnToChats: "Вернуться к {count} чатам",
|
||||
chatGenerating: "Генерация",
|
||||
compare: "Сравнить",
|
||||
search: "Поиск",
|
||||
hub: "Хаб моделей",
|
||||
|
|
|
|||
|
|
@ -37,6 +37,8 @@ export const zhCN = {
|
|||
navigation: {
|
||||
newChat: "新聊天",
|
||||
returnToChat: "返回聊天",
|
||||
returnToChats: "返回 {count} 个聊天",
|
||||
chatGenerating: "生成中",
|
||||
compare: "对比",
|
||||
search: "搜索",
|
||||
hub: "模型中心",
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue