Commit graph

8 commits

Author SHA1 Message Date
Daniel Han
7a25266203
Fix _SameTaskStreamingResponse disconnect test bypassing __init__ (#6627)
* Fix _SameTaskStreamingResponse disconnect test bypassing __init__

test_same_task_response_closes_body_iterator_on_send_disconnect builds the
response via __new__ to skip Starlette's __init__, then wires body_iterator,
background, and stream_response by hand. It never set _unstarted_cleanup, so the
disconnect-before-first-chunk branch of __call__ raised AttributeError instead of
ClientDisconnect, failing the Backend CI "Repo tests (CPU)" job on main.

Set response._unstarted_cleanup = None in the manual construction, matching the
default __init__ assigns.

* Shorten the _unstarted_cleanup comment to one line
2026-06-24 03:47:35 -07:00
Wasim Yousef Said
37166efcfc
Fix Gemma 4 GGUF OpenAI API streams (#6476)
* Fix Gemma 4 GGUF OpenAI API streams

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Avoid duplicate Responses stream disconnect watcher

* Keep reasoning-only Responses output hidden

* Address Gemma stream review comments

* Avoid Responses stream task-group cleanup

* Harden OpenAI chat completion streams

* Address OpenAI stream review issues

* Clean up Studio OpenAI stream helpers

* Fix Studio passthrough cold stream timeout

* Fix tool parser compatibility exports lint

* Preserve audio stream disconnect cancellation

* Avoid synthetic finish after passthrough errors

* Address stream cleanup and Gemma parser reviews

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Gemma 4: parse bare-string tool args and keep safetensors tools for native <|tool_call>

- Quote bare unquoted string values in Gemma native tool-call args (e.g.
  {location:Tokyo,unit:celsius}) so they parse; JSON scalars stay typed.
- Stop _detect_safetensors_features from suppressing supports_tools for
  templates that emit Gemma native <|tool_call>, which the shared parser
  now reads.
- Add tests for both.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Harden Gemma tool-call parsing and stream-error detection

Address three issues in the Gemma-native tool-call path:

- _quote_gemma_object_keys stopped a bare (unquoted) string value at the
  first comma, so an argument like `location:New York, NY` was split
  mid-value and the synthesized JSON failed to parse, dropping the whole
  tool call. A bare value now ends only at `}` or a comma that begins the
  next `key:` pair.

- parse_tool_calls_from_text scanned the entire response for Gemma markers
  even inside a tool call already parsed from a `<tool_call>{...}` JSON
  block, so a marker-like string inside an argument (data) was promoted to
  a second, unintended tool call. Matches inside an already-consumed call
  span are now skipped.

- _openai_passthrough_stream relied on _monitor_openai_sse_line to flag a
  stream error, which returns early when monitor_id is None
  (skip_api_monitor), so an upstream error chunk left saw_stream_error
  unset and the synthetic-finish guard emitted a successful finish_reason
  after a failed stream. Error chunks are now detected independently of API
  monitoring.

Adds tests/test_gemma_tool_parse_edge_cases.py covering the comma and
marker-injection cases.

* Emit the terminal finish_reason chunk in GGUF streams

The OpenAI chat-completions GGUF tool stream and plain stream both built a
final ChatCompletionChunk carrying finish_reason but never yielded it, so
clients received the optional usage chunk and [DONE] with no chunk carrying
finish_reason. OpenAI-compatible consumers rely on that terminal choice to
distinguish stop/length/tool_calls. Yield it before the usage chunk and
[DONE], matching the other streaming paths.

* Parse tool calls in document order and skip nested markers both ways

Unify the JSON- and Gemma-format tool-call passes into a single
position-ordered scan:

- Calls are now emitted in byte order across both formats, so a mixed
  output like `<|tool_call>call:create{...}<tool_call|> ... <tool_call>
  {"name":"read",...}</tool_call>` executes create before read, matching
  the order they appear in (tools run in returned order).

- A candidate that starts inside an already-accepted call's span is
  skipped, in both directions: a JSON marker inside a Gemma argument and a
  Gemma marker inside a JSON argument are treated as data, not promoted to
  a second executable tool call.

Extends tests/test_gemma_tool_parse_edge_cases.py with the ordering and
JSON-in-Gemma nesting cases.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Quote bare Gemma array elements; order finish before trailing usage

- _quote_gemma_object_keys skipped array values, so a Gemma call with a
  bare-string array argument like labels:[bug,ui] produced invalid JSON and
  the whole tool call was dropped. Array values are now scanned and bare
  string elements quoted, while numbers, quoted strings, and JSON literals
  are preserved.

- In the OpenAI passthrough stream, a trailing usage-only chunk
  (stream_options.include_usage) that arrived before any finish chunk was
  relayed before the synthetic finish, producing usage -> finish -> [DONE].
  Emit the synthetic finish before that usage chunk so the order matches the
  other streams (finish -> usage -> [DONE]).

Extends tests/test_gemma_tool_parse_edge_cases.py with the bare-array cases.

* Harden Gemma array parsing, XML-parameter guard, and stream teardown

Address five review findings on the Gemma tool-call and OpenAI passthrough
streaming paths:

- parse_tool_calls_from_text collected JSON and Gemma markers without the
  _inside_open_parameter guard, so a marker embedded in an existing
  <function=...><parameter=...> value was promoted to a separate tool call.
  Candidates that start inside an open XML parameter are now skipped, matching
  the guard the XML-style parser already applies.

- _quote_gemma_array_elements preserved array elements starting with { or [
  verbatim, so an array of objects (items:[{path:a}]) or a nested array failed
  json.loads and the whole call was dropped. Object and nested-array elements
  are now normalised recursively.

- _openai_passthrough_stream synthesized a finish chunk before a trailing
  usage-only chunk and set saw_finish_reason, which made the EOF guard skip the
  [DONE] sentinel. The EOF path now emits [DONE] whenever the upstream omitted
  it, even after a finish chunk was already synthesized.

- /generate/stream drove generation through asyncio.to_thread with no
  disconnect watcher, so a client disconnect during a long generation went
  unnoticed until the next send. It now runs _await_disconnect_then_cancel
  against the request, matching the other local streaming endpoints.

- _SameTaskStreamingResponse closed the body iterator with aclose() on a
  send-side disconnect, raising GeneratorExit so the generators' cancellation
  handlers (which finish the api_monitor entry) never ran. It now throws
  CancelledError, falling back to aclose() when athrow is unavailable.

Extends tests/test_gemma_tool_parse_edge_cases.py with array-of-objects,
nested-array, and marker-inside-XML-parameter cases.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Watch disconnects on Anthropic streams; keep timestamps in Gemma values

Two follow-ups on the streaming and tool-parse paths:

- _anthropic_tool_stream and _anthropic_plain_stream drove generation through
  asyncio.to_thread(next, gen, ...) and only polled is_disconnected() between
  events, so a client disconnect during prefill or a long generation/tool step
  held the decode slot until the next event or a failed send. Both now run the
  _await_disconnect_then_cancel watcher used by the other local streams, stop it
  in finally, and break promptly when cancel_event is set.

- _GEMMA_NEXT_KEY_RE treated any comma followed by word-chars-then-colon as the
  next key, so a bare value such as "meet at 10:00, 11:00 tomorrow" was split
  into bogus keys. The next-key token must now be identifier-shaped (start with
  a letter or underscore), so a comma before a timestamp, ratio, or other
  numeric-then-colon text stays part of the value.

Adds a timestamp-in-bare-value regression test.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Guard nested markers, reset on disconnect, clean unstarted streams

Three follow-ups on the tool-parse and streaming paths:

- parse_tool_calls_from_text only skipped markers that fell inside a span it
  had already parsed successfully, so when an unquoted Gemma argument contained
  a literal marker (code:<|tool_call>call:terminal{...}<tool_call|>) the outer
  object failed to normalize, its span was never recorded, and the inner marker
  was promoted to a standalone terminal call. Candidates nested inside any other
  candidate's brace span are now skipped regardless of whether the enclosing
  candidate parsed, so a marker in malformed outer data is never executed.

- /generate/stream skipped backend.reset_generation_state() when the disconnect
  watcher set cancel_event between chunks: the loop broke and the finally's reset
  is guarded on cancel_event being unset. A subprocess backend kept decoding
  after the client left. The cancel-break path now resets the backend.

- _SameTaskStreamingResponse threw CancelledError / called aclose() on the body
  iterator on a send-side disconnect, but neither runs the try/finally of a
  generator that never started (early disconnect on http.response.start), so the
  passthrough's eagerly-opened upstream httpx stream and cancel-registry entry
  leaked. It now tracks whether the body started and, when it did not, runs an
  optional unstarted_cleanup hook; the OpenAI passthrough wires it to close the
  upstream resp/client and exit the cancel tracker.

Adds a nested-unquoted-marker regression test.

* [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>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-06-23 06:13:56 -07:00
Parvesh Saini
17e9714a98
studio: run /generate/stream's sync generator off the event loop to avoid blocking it (#6466)
* studio: run /generate/stream's sync generator off the event loop to avoid blocking it

* fix: close generator in finally on client disconnect in generate_stream

* Fix/adjust generate stream test for PR #6466

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix/adjust generate stream cancellation for PR #6466

* Fix/adjust generate stream cleanup for PR #6466

* fix: cancel incomplete generate stream cleanup

---------

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>
2026-06-19 14:19:04 +01:00
Daniel Han
a6dc10dad2
Reduce and tighten comments and docstrings across the test suite (#6429)
* Reduce and tighten comments and docstrings in tests

Shorten verbose comments and docstrings across the test suite without
changing any test logic. Remove narration that restates the next line,
collapse long module and test docstrings to a single line, and drop banner
separators. Keep regression context (issue and PR references, run ids),
skip reasons, mocking and timing rationale, license headers, lint and type
directives, and commented-out code.

Comments and docstrings only: an AST signature check confirms no code,
assertions, or string literals changed, and the suite byte-compiles cleanly.

* [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>
2026-06-18 01:07:09 -07:00
Anmol Mishra
9f694ab750
fix(studio): Windows GGUF cancel hang + CPU spinlock overhead (#5692) (#5749)
* fix(studio): Windows GGUF cancel hang + CPU spinlock overhead (#5692)

Two fixes for Windows-native GGUF inference via llama-server:

**Issue 1 — GPU/CUDA Hang on Stream Cancellation:**
- Add `Connection: close` header to all httpx requests proxying to
  llama-server, preventing Keep-Alive from masking downstream socket
  closure.
- Introduce `_await_disconnect_then_close` background watcher that
  polls `request.is_disconnected()` every 100ms and calls
  `resp.aclose()` immediately when the client disconnects. This runs
  alongside the existing cancel-POST watcher and covers client aborts
  that never reach the /cancel endpoint (tab close, proxy aborts,
  Colab, mobile navigation, etc.).
- Change all StreamingResponse `Connection: keep-alive` headers to
  `Connection: close`.

**Issue 2 — High CPU Spinlock & KV Cache Backup Overhead:**
- Set OMP_WAIT_POLICY=PASSIVE and OMP_NUM_THREADS=2 in the
  llama-server subprocess environment on Windows to prevent OpenMP
  from spin-waiting on all logical cores while the GPU decodes.
- Limit `--threads` to 2 on Windows when the model is fully
  GPU-offloaded (`-ngl -1`). Auto-detect otherwise.
- Pass `--cache-ram 0 --ctx-checkpoints 0 --no-cache-prompt
  --checkpoint-every-n-tokens -1` on Windows to disable prompt-cache
  snapshots that copy KV cache to system RAM over the WDDM/PCI-E bus.

Closes #5692.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: use local import to avoid ruff F823 (sys used before assignment)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* review: address gemini review feedback

- Simplify _fully_gpu_offloaded init: default to False, only set True
  in the gpu_indices branch, drop redundant else.
- Log exceptions in _await_disconnect_then_close at debug level instead
  of silent pass, per review suggestion.

* Adjust review feedback for PR #5749

- _await_disconnect_then_close: set cancel_event before resp.aclose() so
  the streamer's RemoteProtocolError handler treats the watcher-driven
  close as cancellation, not an upstream error. Both call sites pass
  cancel_event through.
- Windows --cache-ram / --no-cache-prompt / --ctx-checkpoints block: gate
  on _fully_gpu_offloaded so CPU and partial-offload Windows runs keep
  prompt-cache reuse across turns.
- Windows OMP_WAIT_POLICY / OMP_NUM_THREADS env: same gate so CPU and
  partial-offload Windows runs keep default OpenMP parallelism.

* Shorten code comments touched by PR #5749

* Clean up local imports and rename underscore locals in PR #5749

- Drop the function-local `import sys as _sys` introduced as an F823
  workaround; remove the redundant in-function `import os`/`import sys`
  block so module-level imports resolve sys/os instead. F823 no longer
  triggers because no shadowing import remains inside load_model.
- Rename `_fully_gpu_offloaded` and `_t` to `fully_gpu_offloaded` and
  `threads_arg`. Underscore-prefixed names usually mean private/module-
  level; plain locals match Python style for in-function temporaries.

No behavior change. ruff clean, py_compile clean, 35 studio cancel-
infra tests + 13 launch-gating AST locks + 6 disconnect-watcher locks
+ 4 spoof live-import tests all pass.

* Fix Windows GGUF follow-ups for PR #5749

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix cache flag gating for PR #5749

* Fix Python 3.9 annotations for PR #5749

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Anmol Mishra <anmolx.work@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
2026-06-15 10:32:10 +01:00
Daniel Han
187144d4e7
Reduce and tighten code comments and docstrings repo-wide (#6095)
Trim and tighten code comments and docstrings across the repository. Comment-only: every changed file verified code-identical to main via AST/token comparison.
2026-06-08 23:09:51 -07:00
Daniel Han
3ce187da02
Formatting: ruff line-length 100, kwarg-spacing passes, drop blank after short local imports (#6079)
Raise ruff line-length to 100 and extend the local pre-commit format pipeline (def-signature magic-comma normalization, short multi-line assert collapse, kwarg '=' spacing, blank-line-after-short-import removal, adjacent string-literal / f-string+plain merge, redundant-pass pruning). Every transform re-checks the file AST and is dropped if it would differ; the whole-repo reformat is verified AST-identical per file and idempotent.
2026-06-08 04:24:13 -07:00
Daniel Han
eb8b0dee2e
Studio: make stop button actually stop generation (#5069)
* Studio: make stop button actually stop generation

The UI stop button routes through assistant-ui's cancelRun, which aborts
the frontend fetch. Four issues combined to let llama-server keep decoding
long after the user clicked stop:

1. request.is_disconnected() does not fire reliably behind proxies
   (e.g. Colab) that don't propagate fetch aborts.
2. llama-server defaults n_predict to n_ctx when max_tokens is not sent,
   so a cancelled request keeps producing tokens up to 262144.
3. The httpx.Client pool keeps TCP keep-alive, so even a cleanly closed
   stream reuses the same connection and llama-server's liveness poll
   never sees a disconnect.
4. No explicit backend route to cancel - every cancel path relied on
   is_disconnected.

Changes:
- Add POST /api/inference/cancel keyed by session_id/completion_id, with
  a registry populated for the lifetime of each streaming response.
- Have the frontend (chat-adapter.ts) POST /inference/cancel on
  AbortController abort, alongside the existing fetch teardown.
- Send max_tokens=4096 + t_max_predict_ms=120000 as defaults on every
  outbound chat completion to llama-server; honoured by user overrides.
- Disable httpx keep-alive on the streaming client so connection close
  reaches llama-server and its 1s liveness check fires.

No behaviour changes for non-streaming paths or for existing callers
that already pass max_tokens/session_id.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: harden stop-button cancel path and scope cancel route

- Require at least one identifier for /api/inference/cancel so a missing
  thread id cannot silently cancel every in-flight generation.
- Scope /cancel to a dedicated studio_router so it is not exposed under
  the /v1 OpenAI-compat prefix as a surprise endpoint.
- Store a set of cancel events per key in _CANCEL_REGISTRY so concurrent
  requests on the same session_id do not overwrite each other, and
  deduplicate in _cancel_by_keys so the cancelled count reflects unique
  requests.
- Always send session_id with chat completions (not only when tools are
  enabled) so non-tool GGUF streams register under it and are reachable
  from /cancel.
- Register the non-GGUF stream_chunks path in the cancel registry too,
  so transformers-based stop-button works behind proxies that swallow
  fetch aborts.
- Only apply the 2-minute t_max_predict_ms wall-clock cap when the
  caller did not pass max_tokens, so legitimate long generations on
  slow CPU/macOS/Windows supported installs are not silently truncated.
- Remove the abort listener on normal stream completion so reused
  AbortSignals cannot fire a spurious cancel POST after the fact.

* studio: close cancel-race and stale-cancel gaps in stop path

- Register the cancel tracker before returning StreamingResponse so a
  stop POST that arrives during prefill / warmup / proxy buffering
  finds an entry in _CANCEL_REGISTRY. Cleanup now runs via a Starlette
  BackgroundTask instead of a finally inside the async generator body.
- Add a per-run cancel_id on the frontend (crypto.randomUUID) and in
  ChatCompletionRequest so /api/inference/cancel matches one specific
  generation. Removes the stale-cancel bug where pressing stop then
  starting a new run in the same thread would cancel the retry.
- Apply t_max_predict_ms unconditionally in all three llama-server
  payload builders (previously gated on max_tokens=None, which made it
  dead code for UI callers that always send params.maxTokens). Raise
  the default to 10 minutes so slow CPU / macOS / Windows installs are
  not cut off mid-generation.
- Make _cancel_by_keys refuse empty input (return 0) so a future
  internal caller can not accidentally mass-cancel every in-flight
  request.
- Accept cancel_id (primary), session_id, and completion_id on the
  /api/inference/cancel route. Unify the three streaming sites on the
  same _cancel_keys / _tracker variable names.
- Annotate _CANCEL_REGISTRY as dict[str, set[threading.Event]].

* Add review tests for PR #5069

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: harden stop-button cancel semantics and wall-clock cap

- Make /inference/cancel match cancel_id EXCLUSIVELY when supplied.
  Previously the handler iterated ('cancel_id','session_id','completion_id')
  and unioned matches, so a stale cancel POST carrying {cancel_id:old,
  session_id:thr} would still cancel a later run on the same thread via
  the shared session_id. cancel_id is now a per-run exclusive key;
  session_id / completion_id are only used as fallbacks when cancel_id
  is absent.

- Close the early-cancel race. If /inference/cancel lands before the
  streaming handler reaches _TrackedCancel.__enter__() (stop clicked
  during prefill / warmup / proxy buffering), the cancel was silently
  dropped. Stash unmatched cancel_ids in _PENDING_CANCELS with a 30 s
  TTL; _TrackedCancel.__enter__() now replays any matching pending
  cancel by set()-ing the event immediately after registration.

- Make t_max_predict_ms = _DEFAULT_T_MAX_PREDICT_MS conditional on
  max_tokens is None at all three llama-server payload sites. The cap
  is a safety net for callers who leave max_tokens unset (otherwise
  llama-server defaults n_predict to n_ctx, up to 262144). Callers who
  set an explicit max_tokens are already self-limiting and must not be
  silently truncated at 10 minutes on slow CPU / macOS / Windows
  legitimate long generations.

- Guard each StreamingResponse return with try/except BaseException so
  _tracker.__exit__ runs even if StreamingResponse construction or any
  preceding statement raises between _tracker.__enter__() and the
  BackgroundTask attachment. Prevents a registry leak on that narrow
  window.

* studio: close TOCTOU race and restore wall-clock backstop on UI path

- Close TOCTOU race in the pending-cancel mechanism. The previous fix
  split cancel_inference's (cancel_by_keys + remember_pending_cancel)
  and _TrackedCancel.__enter__'s (register + consume_pending) into
  four separate lock acquisitions. Under contention a cancel POST
  could acquire-then-release the lock, find the registry empty, and
  stash ONLY AFTER __enter__ had already registered and consumed an
  empty pending map -- silently dropping the cancel. Both call sites
  now do their work inside a single _CANCEL_LOCK critical section, via
  the new atomic helper _cancel_by_cancel_id_or_stash() and an
  inlined consume-pending step in __enter__. Reproduced the race under
  forced interleaving pre-fix; 0/2000 drops post-fix under parallel
  stress.

- Apply t_max_predict_ms UNCONDITIONALLY at all three llama-server
  payload sites. The previous iteration gated the cap on
  `max_tokens is None`, which turned out to be dead code on the
  primary Studio UI path: chat-adapter.ts sets
  maxTokens=loadResp.context_length after every model load, so every
  chat request carries an explicit max_tokens and the wall-clock
  safety net never fired. The cap's original purpose is to bound
  stuck decodes regardless of the token budget; it must always apply.

- Raise _DEFAULT_T_MAX_PREDICT_MS from 10 minutes to 1 hour. 10
  minutes was too aggressive for legitimate slow-CPU chat responses
  (a 4096-token reply at 2 tok/s takes ~34 min); 1 hour accommodates
  that and still catches genuine zombie decodes.

- Prune _PENDING_CANCELS inside _cancel_by_keys as well, so stashed
  entries expire proportionally to overall cancel traffic rather than
  only to cancel_id-specific POSTs.

* studio: trim verbose comments and docstrings in cancel path

* studio/llama_cpp: drop upstream PR hashes from benchmark comment

* Add review tests for Studio stop button

* Consolidate review tests for Studio stop button

* Align cancel-route test with exclusive cancel_id semantics

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: move cancel cleanup to generator finally; drop dead helper

- Move _tracker.__exit__ from Starlette BackgroundTask into each
  streaming generator's finally block. Starlette skips the background
  callback when stream_response raises (OSError / ClientDisconnect),
  which leaked _CANCEL_REGISTRY entries on abrupt disconnect.
- Check cancel_event.is_set() at the top of each GGUF while loop so a
  pending-replay cancel falls through to final_chunk + [DONE] instead
  of propagating GeneratorExit out of _stream_with_retry.
- Remove unused _remember_pending_cancel; _cancel_by_cancel_id_or_stash
  superseded it.

* Add review tests for Studio stop-button

* studio: wire audio-input stream into cancel registry

- Register cancel_event with _TrackedCancel on the audio-input streaming
  path so POST /api/inference/cancel can stop whisper / audio-input GGUF
  runs. Previously the registry stayed empty on this branch, so the stop
  button returned {"cancelled":0} and the decode ran to completion.
- Apply the same finally-based cleanup and pre-iteration cancel-event
  check used on the other three streaming paths.
- Update the _CANCEL_REGISTRY block comment to list cancel_id as the
  primary key (was stale "session_id preferred").

* Consolidate review tests for Studio stop-button cancel flow

- Merge the 6 behavioral tests from test_stream_cleanup_on_disconnect.py
  (finally cleanup on normal/exception/aclose, pre-set cancel_event
  pattern, and its regressions) into test_stream_cancel_registration_timing.py,
  which is the PR's existing file covering the same area.
- Extend structural invariants to include audio_input_stream alongside the
  three GGUF / Unsloth streaming generators: no _tracker.__enter__ inside
  the async gen body, cleanup via try/finally, no background= on
  StreamingResponse.
- Delete test_stream_cleanup_on_disconnect.py (now empty).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: make cancel-via-POST interrupt Unsloth and audio-input streams

Close two remaining gaps in the stop-button cancellation wiring:

- stream_chunks (Unsloth path): add a top-of-loop cancel_event check and
  call backend.reset_generation_state() so cancel POSTs flush GPU state
  and close the SSE cleanly instead of relying on request.is_disconnected
  (which does not fire through proxies like Colab's).
- audio_input_stream: run the synchronous audio_input_generate() via
  asyncio.to_thread so blocking whisper chunks do not freeze the event
  loop, matching the pattern already used by the GGUF streaming paths.

* Add review tests for Studio stop-button cancel flow

* Consolidate review tests for Studio stop-button cancel flow

- Delete standalone test_cancel_registry.py at repo root: tests duplicated
  test_cancel_atomicity.py / test_cancel_id_wiring.py and re-implemented
  registry primitives inline (scaffolding).
- Extend tests/studio/test_stream_cancel_registration_timing.py with
  regression guards for the iter-1 cancel-loop fixes:
    structural: each streaming generator checks cancel_event in its loop;
                audio_input_stream offloads next() via asyncio.to_thread;
                stream_chunks cancel branch calls reset_generation_state().
    runtime:    Unsloth loop breaks on external cancel and resets state;
                audio loop stays responsive under blocking next();
                both loops emit zero tokens on pre-set cancel (replay path).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: extend stop-path to passthrough streams; tighten wall-clock cap

- Lower _DEFAULT_T_MAX_PREDICT_MS from 1 hour to 10 minutes so the
  wall-clock backstop actually bounds runaway decodes when cancel
  signaling fails.
- Wire _TrackedCancel and cancel_event.is_set() into
  _openai_passthrough_stream and _anthropic_passthrough_stream and
  disable httpx keepalive so stop requests from /v1 and /v1/messages
  tool-calling clients reach llama-server.
- Apply t_max_predict_ms to the tool-passthrough request body so the
  backstop covers passthrough paths as well.
- Symmetric pre-registration stash for session_id/completion_id
  cancels (_cancel_by_keys_or_stash) so early cancels by those keys
  replay on later registration like cancel_id.
- Drop dead except BaseException guards around StreamingResponse()
  at four streaming sites; cleanup lives in the generator's finally.

* studio: harden cancel registry against ghost-cancel and leak paths

- Revert the session_id/completion_id stash in the fallback cancel
  helper. session_id is thread-scoped and reused across runs, so
  stashing it on an unmatched POST would fire cancel_event for the
  user's next unrelated request via _TrackedCancel.__enter__.
  cancel_id remains the only per-run unique key that gets stashed.
- Default max_tokens to _DEFAULT_MAX_TOKENS in the tool-passthrough
  body. Mirror the direct GGUF path so OpenAI/Anthropic passthrough
  callers who omit max_tokens get the same zombie-decode cap instead
  of relying on the wall-clock backstop alone.
- Wrap _openai_passthrough_stream setup with an outer try/except
  BaseException. The inner except httpx.RequestError does not catch
  asyncio.CancelledError at await client.send, which would otherwise
  leave _tracker registered in _CANCEL_REGISTRY indefinitely.
- Frontend stop POST uses plain fetch + manual Authorization header
  instead of authFetch. A 401 on the cancel POST no longer refreshes
  tokens or redirects the user to the login page mid-stop.

* Add review tests for Studio stop-button cancel flow

* studio: trim comments on stop-button review changes

Collapse multi-paragraph rationale blocks on the cancel registry,
_openai_passthrough_stream, and the frontend onAbortCancel handler
into one-line explanations of why the non-obvious behaviour exists.
Drop authFetch import that became unused when the cancel POST
switched to plain fetch.

* Consolidate review tests for Studio stop-button cancel flow

Move review-added tests out of test_cancel_dispatch_edges.py into the
existing PR test files that already cover the same areas:
- backend registry fan-out / exclusivity / idempotency / falsy-keys
  edge cases moved into tests/studio/test_cancel_atomicity.py
- frontend plain-fetch (not authFetch) + manual Authorization header
  moved into tests/studio/test_cancel_id_wiring.py
Delete the now-empty test_cancel_dispatch_edges.py.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: stop default-capping responses at 4096 tokens (follow-up to #5069) (#5174)

* Studio: stop default-capping responses at 4096 tokens

Follow-up to #5069. The 4096 default introduced for runaway-decode
defense silently truncates any caller that omits max_tokens. The
Studio chat UI sets params.maxTokens = loadResp.context_length after
a GGUF load, so it's fine, but every other consumer is not:

- OpenAI-API direct callers (/v1/chat/completions, /v1/responses,
  /v1/messages, /v1/completions) where the OpenAI default is
  effectively unlimited per response. langchain, llama-index, raw
  curl, and the openai SDK all rely on that.
- Reasoning models. Qwen3 / gpt-oss reasoning traces routinely exceed
  4096 tokens before the model emits a single visible content token.
  The user sees the trace cut off mid-thought.
- Long-form generation ("write a chapter", "produce a full SVG").

Reproduced on this branch: gemma-4-E2B-it-GGUF Q8_0, prompt asking
for a 10000-word story, no max_tokens in the request:

    finish_reason: stop  (misleading -- should be 'length')
    content_chars: 19772
    content_tail: ...'a comforting, yet immense, pressure.\n\n*"'

Body ended mid-sentence on a stray opening quote, right at the 4096
token mark.

After this patch the same request returns 38357 chars ending with
'...held in a perfect, dynamic equilibrium.' -- a natural stop, not
a truncation.

Implementation: rename the constant to _DEFAULT_MAX_TOKENS_FLOOR and
set it to 32768. Each call site now uses the model's effective
context length when known, falling back to the floor:

    default_cap = self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR

The 10-minute t_max_predict_ms wall-clock backstop from #5069 is
preserved as the second line of defense.

Plumbed _build_passthrough_payload + _build_openai_passthrough_body
through the routes layer so the Anthropic and OpenAI passthrough
paths also respect the model's context length.

* [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>

* Studio: cancel passthrough streams during llama-server prefill + route through apiUrl for Tauri

Three reviewer-flagged correctness gaps in the stop-button mechanism.

1) `_openai_passthrough_stream` could not honor cancel during prefill.
   The cancel check ran inside the `async for raw_line in lines_iter`
   body, so a cancel POST that arrived before llama-server emitted the
   first SSE line was unobservable until prefill completed. With a long
   prompt under proxy/Colab conditions -- the exact target scenario for
   this PR -- that left the model decoding for a long time after the
   user clicked Stop. Add an asyncio watcher task that closes `resp` as
   soon as `cancel_event` is set, raising in `aiter_lines` so the
   generator can exit. The watcher polls a threading.Event because the
   cancel registry is keyed by threading.Event for the synchronous
   /cancel handler.

2) `_anthropic_passthrough_stream` had the same blocking-prefill pattern.
   Same fix.

3) The frontend's stop-button cancel POST used a bare relative
   `fetch("/api/inference/cancel", ...)`, which targets the webview
   origin in Tauri production builds (where the backend is at
   `http://127.0.0.1:8888`). Route through the existing `apiUrl()`
   helper from `lib/api-base.ts` to match every other Studio call.
   Browser/dev builds get the empty base, so behavior is unchanged
   there.

Verified via temp/pr_simulation/sim_5069_prefill_cancel.py: cancel
during prefill terminates within ~250ms on both passthrough paths
(was 145s+ on the Anthropic path before this change), and the standard
non-passthrough chat path still cancels with no regression.

* Studio: log cancel-body parse errors instead of silently swallowing

Reviewer-flagged defensive logging gap. The bare `except Exception: pass`
in `cancel_inference` would mask malformed payloads that hint at a buggy
client or a transport issue. Log at debug so future investigation isn't
left guessing whether `body={}` came from a missing body or a parse
failure. Behavior is unchanged: an unparseable body still falls through
to the empty-dict path and the cancel call returns `{"cancelled": 0}`.

* Studio: Anthropic passthrough cancel parity with OpenAI passthrough

Two reviewer-flagged consistency gaps in the cancel surface for
/v1/messages.

1) Anthropic passthrough did not register cancel_id, so a per-run cancel
   POST (the cleanest Studio-style cancel path) silently missed when
   the route hit `_anthropic_passthrough_stream`. The OpenAI passthrough
   has registered (cancel_id, session_id, completion_id) since this PR
   was first opened; mirror that here. Also add `cancel_id` to
   `AnthropicMessagesRequest` so the route handler can plumb it through.

2) The cancel handler's fallback key list checked only completion_id
   and session_id, never message_id. Anthropic clients that send their
   native `id` (returned in the SSE message_start event) for cancel had
   no way to hit the registry. Add message_id to the fallback list.

Verified via temp/pr_simulation/sim_5069_prefill_cancel.py: P2 now
cancels by cancel_id in 137ms (was hanging pre-fix), and the new P2b
case cancels by message_id in 77ms. P1 (OpenAI) and P3 (standard chat)
still pass with no regression.

---------

Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
2026-04-24 10:09:25 -07:00