Commit graph

6 commits

Author SHA1 Message Date
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
oobabooga
522f308be2
Studio: cache MCP tool discovery instead of re-probing every chat send (#5828)
* Studio: cache MCP tool discovery instead of re-probing every chat send

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

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

* Stop re-probing offline/down MCP servers every time

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

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

* Add tests for mid-probe delete and OAuth cool-off paths

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

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

* Don't cool-off a server edited or deleted mid-probe

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

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

* Guard MCP refresh cache writes

---------

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>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>
2026-06-12 14:18:06 +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
8292e699e4
Studio: make code comments and docstrings more succinct (#6029)
Trim and tighten code comments and docstrings across studio/ Python. Comment-only: every changed file verified code-identical to main via AST/token comparison.
2026-06-08 23:07:28 -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
Nilay
9a907a8acb
Studio: add remote MCP server support (#5750)
* added remote MCP server support

* trim

* added tests

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

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

* increased timeout

* disabling MCP chat toggle

* Fix MCP OpenAI function-name validation + cancel propagation for PR #5750

OpenAI requires function.name to match ^[a-zA-Z0-9_-]{1,64}$ before
streaming starts. The existing 64-char length check is necessary but
not sufficient: MCP servers can return tool names containing '.', '/',
spaces, etc. that would 400 the whole chat request. Validate the
composed mcp__<server_id>__<tool> name against the regex, skip + warn
on miss, and drop duplicate tool names from the same server (which
would also 400 the request as "duplicates").

Also propagate the agentic-loop cancel_event into MCP tool execution
so a /cancel POST during a long-running MCP call (e.g. GitHub MCP
search across a large repo) actually interrupts the in-flight HTTP
call instead of waiting out the 300 s timeout. The watcher polls the
threading.Event at 50 ms cadence inside the asyncio loop (matches
routes/inference.py's existing cancel-watcher cadence) and races
against the call task with asyncio.wait FIRST_COMPLETED.

Tests added:
  - test_mcp_specs_skip_invalid_openai_function_names: drops bad chars
  - test_mcp_specs_skip_empty_tool_name
  - test_mcp_specs_drops_duplicate_names
  - test_call_tool_sync_respects_pre_set_cancel_event

Also fix test_desktop_auth.py's router stub that listed every existing
router but missed mcp_servers_router, so importing main.py fails after
this PR adds it to routes/__init__.py.

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

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

* PR #5750 round 2: OAuth cleanup on delete/url-change + mcp_enabled standalone

Round 2 of cross-platform validation surfaced two more P1 findings:

1. OAuth tokens never get cleared. fastmcp keys tokens by MCP URL, not by
   server row, and delete / URL change / use_oauth toggle only updated
   the SQLite row. Re-registering the same URL would silently reuse the
   old account's credentials. Adds clear_oauth_tokens_async() in
   mcp_client.py and calls it from the delete + put route handlers when
   the row had use_oauth=True and either the URL changes or OAuth is
   turned off.

2. mcp_enabled=true was ignored unless the caller also sent
   enable_tools=true. The frontend always sends both together so the UI
   path was fine, but a direct API caller sending only mcp_enabled would
   silently get no MCP tools, which contradicts the field's documented
   "append tools from every enabled MCP server" behavior. Loosens the
   use_tools gate in both the GGUF and safetensors paths so mcp_enabled
   opens the tool loop on its own; when the caller did not also opt
   into built-ins, the built-in list starts empty.

Tests added:
  - test_clear_oauth_tokens_async_no_op_safe
  - test_delete_server_calls_oauth_cleanup_when_oauth_was_on
  - test_delete_server_skips_oauth_cleanup_when_oauth_off
  - test_update_server_clears_oauth_on_url_change
  - test_update_server_clears_oauth_when_oauth_disabled

26 backend MCP tests pass; full studio/backend suite 1710 passed locally.
Cross-platform CI (Linux, macOS, Windows) green on staging fork.

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

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

* PR #5750 round 3: reject null bool updates + /test surfaces 400

Round 3 of cross-platform validation:

1. PUT /api/mcp/servers/<id> would 500 with TypeError when the body
   explicitly set is_enabled or use_oauth to null. Pydantic accepts
   None for an Optional[bool] and _changes_from_payload then passed
   None into mcp_servers_db.update_server, which int(None)d. Reject
   explicit null at the validation layer with 400 instead.

2. POST /api/mcp/servers/test caught HTTPException under
   "except Exception", so an invalid URL came back as HTTP 200 with
   {"ok": false, "error": "400: ..."} instead of a real 400. The
   create + update paths return 400 for the same input. Move
   validation outside the transport try/except so it surfaces 400.

Tests added:
  - test_changes_from_payload_rejects_null_is_enabled
  - test_changes_from_payload_rejects_null_use_oauth
  - test_test_endpoint_surfaces_url_validation_as_400

* PR #5750 round 4: hyphenated MCP tool names + empty-tool-list gate

Round 4 surfaces two more interaction bugs between the new MCP path
and existing safetensors tool plumbing:

1. OpenAI accepts ^[a-zA-Z0-9_-]{1,64}$ for function.name, and round 1
   widened the MCP regex to that set, so MCP tools can now be advertised
   as `mcp__srv__list-issues`. But the XML tool-call parser in
   tool_call_parser.py used `\w+` (no hyphen), so the model could call
   the tool but Studio could not parse the call. Same in
   routes/inference.py's `_TOOL_XML_RE` stripper, which would leave
   hyphenated tool-call XML in the visible content. Both regexes now
   use `[\w-]+`.

2. safetensors_agentic treats `tools=[]` as "allow all" (documented
   contract, exercised by test_empty_tools_list_does_not_enforce_allowlist).
   When a caller sends `enable_tools=true` + `enabled_tools=[]` +
   `mcp_enabled=true` and MCP discovery returns 0, the resolved tool
   list is genuinely empty and built-in tools (web_search / python /
   terminal) could execute via the model's emitted call. Fix at the
   route gate instead of breaking the documented contract: set
   `use_tools=False` when the resolved list is empty, in both GGUF and
   safetensors paths. Existing callers who omit `enabled_tools` still
   get ALL_TOOLS and are unaffected.

Tests added (32 total):
  - test_tool_xml_parser_handles_hyphenated_function_names
  - test_tool_xml_strip_handles_hyphenated_function_names
  - test_safetensors_agentic_empty_allowlist_still_means_allow_all
    (documents the contract round 4 preserved)

1716 passed locally; cross-platform CI on staging fork still green.

* PR #5750 round 5: GGUF allow-list + CLI policy + hyphenated params + cancel race

Round 5 of parallel-reviewer aggregation surfaced six additional
findings; five are real and fixed here:

1. Hyphenated MCP parameter names (`<parameter=issue-number>`) were
   dropped by the XML parser's `\w+` regex. Extended to `[\w-]+` in
   both core/inference/tool_call_parser.py and core/tool_healing.py.
   The latter is GGUF's own copy of the parser/strip patterns and was
   missed by round 4.

2. core/tool_healing.py's `strip_tool_call_markup` still used
   `<function=\w+>` so hyphenated MCP tool-call XML leaked into the
   GGUF visible content even after round 4 fixed the shared parser.

3+4. `mcp_enabled` re-opened the tool loop even when the operator
   passed `unsloth run --disable-tools` (CLI policy False). Round 2's
   `(_tools_on or payload.mcp_enabled)` gate ignored the raw process
   policy. Now reads `state.tool_policy.get_tool_policy()` and gates
   mcp_enabled on `_cli_policy is not False`. Applied to both GGUF
   and safetensors paths.

5. GGUF's agentic loop called `execute_tool(tool_name, ...)` without
   checking the model-emitted name against the per-request tool list,
   while the safetensors loop already enforces this. Added the same
   allow-list check so a model that hallucinates a filtered MCP name
   or a built-in the caller opted out of returns "not enabled" instead
   of executing.

Bonus P2 fixes:
  - `call_tool_sync` now checks `cancel_event.is_set()` BEFORE
    creating the call task, so a pre-set cancellation does not open
    the HTTP transport.
  - `clear_oauth_tokens_async` moved the OAuth import + construction
    inside the protected try block; a fastmcp.client.auth load error
    used to escape and 500 the delete / update route.

NOT fixed (verified false or out of scope):
  - finding #10 "structured_content vs structuredContent": fastmcp's
    CallToolResult dataclass uses snake_case (verified live against
    structured-only tool result; fields are
    `dict_keys(['content', 'structured_content', 'meta', 'data', 'is_error'])`).
  - finding #11 "asyncio.run from running loop": call_tool_sync is
    invoked from `asyncio.to_thread` worker threads which have no
    event loop; asyncio.run() is safe there.

Tests added (37 total): hyphenated param names, tool_healing strip,
GGUF allow-list gate, cancel pre-set short-circuit, OAuth cleanup
constructor-error swallowing. 1721 passed locally, no regressions.

* [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: danielhanchen <danielhanchen@gmail.com>
2026-05-27 07:01:11 -07:00