Studio: stream live tool output with SSE heartbeats, fix web page extraction, and surface interrupted turns (#7083)
* Studio: stream live tool output with SSE heartbeats and fix web page extraction Server-side python/terminal tools now stream incremental stdout to the chat UI while running (new tool_output SSE event), and every blocking tool execution emits heartbeat keepalives so reverse proxies (Cloudflare tunnels cap idle streams at ~100s) cannot drop the connection mid-turn. The tool loop routes also emit a stall keepalive during silent prompt prefill between tool iterations. The final role=tool message the model sees is byte-identical to before, so tool-call parsing, nudging, and healing are untouched. web_search page fetches now extract main content: GitHub repo root pages are rewritten to the README API (with HTML fallback), hidden/aria-hidden client error placeholders are dropped, conversion scopes to article/main, and known boilerplate fragments are stripped. Non-HTML responses are returned raw instead of being run through the HTML converter. The frontend renders live-scrolling tool output inside running python and terminal cards, and a chat stream that ends without a terminal signal now surfaces an explicit interrupted state with a Retry action instead of silently ending the turn. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: fix content-type sniffing, unlimited-timeout drain, and env parity in tool streaming Content-Type sniffing: get_content_type() defaults to text/plain when the header is absent, so the sniffing fallback never fired and header-less HTML came back as raw markup. Report an empty type for a missing header and sniff the body whenever the declared type is not HTML, so mislabeled text/plain HTML pages are converted like before the extraction change. Unlimited timeout drain: with tool_call_timeout disabled the old path used communicate(timeout=None) and waited for EOF, but the streaming drain capped the post-exit drain at a 5 second join, truncating output from a grandchild that holds stdout open. When timeout is None, drain until EOF or the cancel event fires; finite timeouts keep the bounded remaining-budget join. Env parity: drop the PYTHONUNBUFFERED=1 injection on the streaming path so the child invocation is byte-identical with and without streaming (the env var was model-visible via os.getenv). Live streaming granularity now depends on the child flushing; unflushed output arrives in ~8 KB chunks or at exit and the final result is unchanged, with SSE heartbeats covering the gaps. * Studio: stream tool-call arguments while the model writes them A model writing a large tool call (a full python game is minutes of generation) produced nothing on the stream: the structured path accumulated delta.tool_calls fragments silently after the provisional card, and the text path's DRAINING state consumed everything until stream end. The user saw a dead Running spinner while the model was in fact writing code, and the byte-silent SSE segment was also the window where proxies drop the connection. New tool_args SSE events stream the arguments as they generate. The structured path forwards each fragment once a provisional card exists (backlog first, so the card starts from the top of the call). The text path sniffs the drained call for an enabled tool name and streams the raw call text under the id the stream-end parser assigns its first call (call_0), so the final tool_start reconciles the same card; the sniff is gated on enabled names plus the provisional size floor, and prose or ordinary JSON answers never spawn a card. The safetensors loop streams the drained render_html call to its existing provisional card the same way. The chat adapter accumulates the raw stream per card and feeds a partial JSON parse (call envelopes and stringified arguments unwrapped) into the part's args, so the python and terminal cards render the code live and the render_html canvas builds while streaming; both cards now say Writing code / Writing command during this phase via useToolArgsStatus. Display only: the parser input, the executed call, and the conversation the model sees are byte-identical, covered by new loop-level tests for the structured path, the text path, and the no-tool JSON answer. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: keep full tool output visible past the model cap; heal /mnt/data habits Live testing surfaced two issues in the tool streaming UX. First, a long python stdout ended in '... (truncated' in the finished card: the model-visible result is capped by tools._truncate (_MAX_OUTPUT_CHARS, previously 8000 chars) to protect the context window, and the card rendered that capped text even though the live stream had already shown everything. The cap stays (raised to 16000, overridable via UNSLOTH_TOOL_RESULT_MAX_CHARS) but display and model concerns are now split: the adapter preserves the accumulated live stream on tool_end whenever it captured more than the final result, and the finished python/terminal cards prefer it. The live-stream ceiling rises from 16 KB to 400 KB (chunks batch per poll, so SSE stays cheap), and both the live pane and the finished card render only the last 2000 lines with a Show all control so a huge output cannot jank the DOM. The truncation notice now tells the model the user saw the full output and that written files persist in the working directory. The final result string remains byte-identical with and without streaming. Second, models trained on ChatGPT code-interpreter transcripts write to /mnt/data, which does not exist here (the sandbox CWD is a per-thread persistent dir). Three layers, all identical across streaming and non-streaming paths: the python/terminal tool descriptions gain one sentence saying to use relative paths in the persistent CWD; a failed execution whose output shows a missing-file error on a known code-interpreter prefix (/mnt/data, /mnt/outputs, /home/sandbox, /workspace) gets a model-visible retry hint appended after truncation so it always survives; and a sitecustomize shim on the sandbox PYTHONPATH remaps those prefixes onto the CWD in open()/os.makedirs() with a one-line stderr notice, covering the python tool and any Python launched from the terminal tool without touching the exec wrapper (so tracebacks keep their line numbers). Bash-level file operations cannot be redirected without root or mount namespaces, so they rely on the description and the hint. * Studio: fix hidden-element parsing, heartbeat gaps, and tool output id collisions Review follow-ups on the tool streaming work: - _html_to_md: treat any present hidden attribute value as hidden (it is an enumerated attribute whose invalid value default is the Hidden state, so hidden="false" is still hidden), and implement HTML5 optional end tags so an unclosed <p hidden> or <li hidden> ends at the next sibling start tag instead of swallowing every following sibling until the parent closes - tool_stream_exec: keep heartbeats flowing after the live-output cap; a tool that keeps printing past the cap kept the queue non-empty, so neither tool_output nor heartbeat events were emitted and the SSE stream went silent past proxy idle timeouts - routes/inference: forward tool heartbeats before the disable_parallel_tool_use drop window swallows events, so a dropped call that executes server-side cannot leave the Anthropic stream silent - llama_cpp: close the provisional text tool card with a tool_end when the drained call fails to parse (DRAINING false-positive path), so the card cannot spin forever while the text is delivered as content - tools: decode terminal output as utf-8 with errors=replace like the python tool; invalid bytes used to raise UnicodeDecodeError from communicate() on the non-streaming path and silently truncate the streaming reader, so the two paths diverged - sitecustomize: patch io.open alongside builtins.open; pathlib Path.open, read_text and write_text call io.open directly and bypassed the remap - frontend: scope the toolLiveOutput/toolFullOutput store keys by pane (modelType and pairId) and clear stale entries on tool_start; backend ids like call_0 repeat across turns and across concurrently streaming panes (compare mode), so a later turn or another pane could display the wrong preserved output, and run-end cleanup now clears only its own keys Each backend fix carries a regression test that fails on the previous code; the byte-identity tests between streaming and non-streaming stay green. * Studio: keep tool failure status visible and truncation/remap notices truthful Finished python/terminal cards preferred the fuller live stream by length alone, so a tool that printed a lot then timed out or exited non-zero showed the captured stdout but dropped the final result's status (timeout notice, Exit code N). preferFullToolOutput now shows the stream when the result is just its truncated prefix, and appends the result otherwise so the failure tail always survives and the copy button copies both. The result truncation notice claimed the user was shown the full output, but the same wrapper serves non-streaming chat/API and direct execute_tool() callers where nothing is streamed to anyone. The notice is now mode-neutral and stays byte-identical with and without an output_callback, keeping the streaming vs non-streaming invariant intact. The sandbox sitecustomize shim now remaps /tmp/outputs into the working directory only while it does not already exist, so a real /tmp/outputs the user's own code created is never shadowed; /tmp/outputs also joins the missing-path retry-hint list. * Studio: suppress hidden void elements and keep live output scroll pinned only when at bottom * Studio: drop capped tool output without concatenating; remap pathlib mkdir Past the live-output cap stream_tool_execution built item + _drain_pending() (the current chunk joined with every queued sibling) only to discard it in the capped branch, so a chatty tool (yes, a tight print loop) could enqueue far more than one poll interval of text and blow past the memory/CPU ceiling the cap exists to enforce. Drain and drop queued items without building a combined string, still counting each drain toward the heartbeat cadence so the SSE keepalive survives. Generated code often prepares code-interpreter paths with Path('/mnt/data').mkdir(parents=True, exist_ok=True); pathlib drives that through os.mkdir (not the patched os.makedirs) per component and, on FileExistsError, probes the unpatched os.stat via Path.is_dir(), so the setup raised before open() ever ran. Patch os.mkdir with the same remap and patch Path.mkdir so the whole parents/exist_ok dance lands on the mapped working directory and stays idempotent; real paths still pass through. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: generalize sandbox write remap and hint to any hallucinated absolute path Models invent absolute paths from seeing their CWD (a DeepSeek run did open('/home/ubuntu/Sandbox/flappy_bird.html', 'w') and died with FileNotFoundError). A prefix list cannot enumerate these, so the sitecustomize shim gains a write-mode fallback in open()/io.open(): when a write/create-mode open targets an absolute path outside the CWD whose parent directory does not exist, redirect it to the basename in the CWD and emit the same one-line stderr notice, echoing the original path. The prefix remaps still run first (they cover reads and preserve subpaths); read modes never hit the fallback so real system files fail or succeed truthfully; bytes paths pass through. The fallback is not applied to mkdir/makedirs/Path.mkdir, since creating an arbitrary absolute directory can legitimately succeed on the host, and that decision is documented in a comment. The model-visible retry hint now echoes the real failing path (parsed from the traceback tail) instead of the canned /mnt/data example, and fires for any absolute path outside the working directory, not just the enumerated prefixes, while a relative miss still gets no hint. The shim wrapper still adds one frame to tracebacks that surface open() errors; suppressing only our frame has no clean standard mechanism (a wrapper always adds a frame), so the frame is left as an accepted compromise. Tests: hallucinated absolute write remaps to the CWD basename across w/a/x/w+; reads of a missing absolute path pass through untouched; writes to an existing external dir pass through; prefix subpaths still preserved; end-to-end write fallback lands the file in the sandbox workdir identically with and without streaming; the hint echoes the actual path for convention and non-convention absolute paths alike. * Studio: kill exited process groups on drain; bound the over-cap output batch _drain_process_output killed the process only via _kill_process_tree, which short-circuits once the parent has exited, so a grandchild that inherited stdout and outlived the parent was never signaled: a finite-timeout run could return while it kept holding the pipe, and a timeout=None cancel left it behind. Capture the setsid process group before waiting and SIGKILL that group at both give-up points so the whole tree is torn down. The streaming wrapper's first over-cap batch joined the current chunk with the entire pending backlog before enforcing the live-output cap, so a chatty tool could allocate far past the cap on the crossing batch. Bound the drain to the remaining budget and drop the surplus in place, keeping the truncated output byte-identical to joining everything. * Studio: harden sandbox path healing and process/generator cleanup Sandbox sitecustomize shim: - Make the generalized write fallback collision-safe: never redirect an invented absolute path onto an already-present CWD file (refuse and let the original open raise FileNotFoundError, preserving the workspace file). - Only w/a/x create a file; r+/rb+ are read-update modes that require the target to exist, so a bare + no longer trips the write fallback. - Gate every convention-prefix remap (/mnt/data, /mnt/outputs, /home/sandbox, /workspace) on the prefix root being absent, so a real host mount is never shadowed; a miss under an existing real prefix passes through. - Patch os.open so Path.touch and other low-level creators heal convention paths too, matching the Path.mkdir patch. Local code execution (tools.py): - Capture the setsid process group right after Popen (before any watcher can poll/reap the leader) and thread it through the cancel watcher and drain. - Kill the captured group in the non-streaming python/terminal timeout branch so an exited leader no longer leaks a stdout-holding grandchild (matches the streaming drain path). - Guard os.getpgid/os.killpg by platform so streamed execution no longer raises on Windows; fall back to single-pid kill. - Judge missing-path hints against the executor's real workdir so a legitimate miss inside a project workspace outside the sandbox root is not mislabeled. Tool streaming routes (routes/inference.py): - Drain a pending next(gen) worker before closing the generator in the safetensors and Anthropic tool streams, so a disconnect no longer races gen.close() (generator already executing) or leaks the thread/generator. HTML to markdown: - Only drop boilerplate lines composed entirely of known furniture phrases so real prose that merely quotes one (for example "we use cookies to authenticate requests") is preserved. Adds hermetic tests for each change. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: keep aside callouts, contain sandbox path remaps, and keepalive dropped Anthropic tool events _html_to_md: stop dropping <aside> unconditionally. Documentation pages render notes/warnings/examples as aside admonition callouts; those inside the selected article/main scope are real content. A furniture aside outside the scope is already excluded by the main-content pass. sitecustomize: contain the code-interpreter path remap under the sandbox CWD. A hallucinated habit path such as /mnt/data/../other_session/file no longer escapes the per-conversation workdir; parent-traversal components in the suffix are dropped and a '.'/'..' write-fallback basename is refused. routes/inference: emit a rate-limited comment keepalive when the Anthropic Messages stream drops tool_output/tool_args events. A chatty tool keeps the generator busy so the stall keepalive never fires and the tool wrapper emits heartbeats only while idle, which left the SSE stream silent past proxy idle caps; the OpenAI passthrough paths forward these events, this path now keeps the connection alive. * Studio: bound the tool-output chunk that first crosses the live cap _drain_queue joined the entire chunk that first crossed the live-output cap before dropping the rest, so a single multi-megabyte line (or any chunk dequeued once the budget was already met at max_chars <= 0) was materialized in full only to be truncated away, defeating the memory ceiling the cap enforces. Slice the crossing chunk to one character past the budget: that preserves the caller's overflow signal and its byte-identical truncation while dropping the arbitrarily large remainder in place. * Studio: scope missing-path hint to the failing line, keepalive dropped-call output, and preserve truncated tool streams over byte length - tools._missing_path_hint: the code-interpreter convention-prefix trigger scanned the whole output, so a convention prefix mentioned only in a traceback frame (a /workspace project root) or printed by the user's code would add a misleading 'use a relative path' hint even when the actual FileNotFoundError was a relative or in-workdir path. Scope the convention test to the failing-path error line(s), matching _extract_missing_abs_path. - _anthropic_tool_stream: the tool_output/tool_args rate-limited keepalive sat after the drop_until_tool_end skip, so under disable_parallel_tool_use a chatty second-or-later tool call was dropped whole with no keepalive, letting an idle proxy kill the SSE stream. Check the keepalive branch before the drop skip (like the heartbeat branch) so dropped-call output keeps the stream alive. - preferFullToolOutput / chat-adapter: a truncated result can be longer than the live stream by byte count once its footer, an 'Exit code N:' notice, or an __IMAGES__ base64 tail is appended, so the length-only gate discarded the full stream and the finished card fell back to the truncated text. Add a shared truncation-aware shouldPreserveFullOutput used by both the write and read sites: preserve the stream whenever the result carries the truncation footer. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: skip the habit-path hint for real project paths under a convention prefix * Kill captured process group on streamed wait-timeout The streamed drain path's proc.wait() timeout branch only called _kill_process_tree(proc). If the leader exits in the narrow window between the wait timing out and _kill_process_tree sampling its pgid, that helper short-circuits on the reaped leader and a stdout-holding grandchild in the same group survives. Also kill the captured pgid there, matching the non-streaming communicate() timeout path. Adds a hermetic regression test that models the reaped-leader race by stubbing _kill_process_tree. * Fix 3.10 pathlib write_text remap and honor cancel in finite drain On Python < 3.11 pathlib routes Path.open / read_text / write_text through a module-level accessor singleton whose open attribute captured the original io.open at import time (_NormalAccessor.open = io.open). Patching io.open in the sandbox shim therefore never reached that captured reference, so a Path('/mnt/data/x').write_text(...) raised FileNotFoundError on 3.10 while passing on 3.11+ (which dropped the accessor and calls io.open at call time). Repoint _NormalAccessor.open at the same io.open wrapper via a staticmethod, guarded so it is an idempotent no-op on 3.11+. Keep the test save/restore helpers symmetric so the accessor is restored too, and add a hermetic write_text/read_text remap test that covers every version. Also honor cancellation while draining inherited stdout after the leader exits. Once the leader is reaped the cancel watcher returns (its loop is while proc.poll() is None), so the finite-timeout drain did one blocking reader.join(timeout=remaining) that ignored cancel_event and kept draining a chatty grandchild for the whole budget after a disconnect/Stop. Poll cancel_event in 0.5s slices against a deadline like the timeout=None branch and kill the captured process group promptly on cancel. The normal path still reaches EOF on its own, so the streamed vs non-streamed result is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: port no-tool stream keepalive/drain and fix subprocess/queue/extraction asymmetries Streaming no-tool paths now match their tool twins: - _anthropic_plain_stream, safetensors/MLX no-tool stream, and standard GGUF no-tool stream run next(gen) in a worker with a timed SSE keepalive loop so a long prompt prefill cannot leave the stream idle past a proxy cap. - The Anthropic plain and safetensors/MLX no-tool teardowns now drain the pending next(gen) worker and close the generator on disconnect instead of leaking the suspended generator. Other asymmetries: - Non-streaming _python_exec/_bash_exec always drain via _drain_process_output (output_callback may be None) so a cancelled run reaps a stdout-holding grandchild that outlived the leader instead of blocking in communicate(). The joined bytes are identical to communicate(), so streamed vs non-streamed results stay byte-identical. - _build_bypass_env installs the sitecustomize path shim on PYTHONPATH (prepend, keeping the operator's entries) so /mnt/data remap works in bypass mode too. - GGUF forwards output_callback to execute_tool only when the callable accepts it (shared accepts_output_callback), matching safetensors and preserving legacy monkey-patched signatures. - tool_stream_exec bounds accepted live output at the producer boundary so a chatty tool cannot grow the queue without limit under consumer backpressure and cannot keep the drain spinning and starve heartbeats. - html_to_md implicit-close now searches past unclosed inline descendants so a hidden <p>/<li> is closed by a following block; main-content scoping gates on the largest single <article>/<main> so a swarm of tiny cards cannot pass the threshold in aggregate and displace the real main. - preferFullToolOutput re-attaches the "Exit code N:" prefix to the fuller stream instead of appending the still-prefixed result, so a failed truncated tool no longer duplicates its stdout in the finished card. Adds hermetic tests for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Preserve short live output on timed-out tools; strip inline-CSS-hidden subtrees and score truncated main-content scopes * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten chat tool streaming comments and docstrings * Keep HTML READMEs from the GitHub API and preserve interrupted tool output Convert a 200 HTML README body from the GitHub README API to Markdown instead of discarding it and falling back to the repo page chrome, and promote captured live stdout to full output when a tool never reaches tool_end (stream interrupted or cancelled) so the partial diagnostics stay on the finished card. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: anchor HTML sniff, keep repeated sandbox writes, reuse textual tool ids Anchor _looks_like_html to the leading doctype/tag so a Markdown README that opens with a fenced HTML example stays Markdown (no html_to_markdown corruption), while bare HTML fragments (<body>/<article>/<section>) are still detected and converted on a missing/wrong Content-Type. Let the sandbox write fallback re-serve a target it already healed for the same invented absolute path, so iterative overwrites of a generated artifact stop failing with FileNotFoundError while the anti-clobber guard still refuses unrelated same-basename files. Reconcile the first textual tool call carrying an explicit id onto the open provisional TEXT card instead of spawning a duplicate card under that id. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: run implicit-close before skipping tags and keep leading README tables as Markdown A skipped block (<nav>/<footer>) is an HTML5 optional-end-tag closer of an open <p>, but handle_starttag returned before the implicit-close bookkeeping, so a never-closed <p hidden> kept its hidden mark and swallowed every following sibling. Run _close_implicit before the skip decision so the hidden mark is released and trailing content renders. Drop <table> (and its <thead>/<tbody>/<tr>/<td>/<th> children) from the _looks_like_html leading set: Markdown READMEs routinely open with a raw HTML <table> badge/layout row, and sniffing that as HTML collapsed the whole Markdown body through html_to_markdown, exactly like the already-excluded <div align>/<p align> layout headers. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: make bypass-permissions Popen double faithful to the unified drain path The non-streaming _python_exec/_bash_exec now share _drain_process_output, which reads proc.stdout in a reader thread and calls proc.wait(); the test double only implemented communicate(), so bypass-mode bash returned an AttributeError instead of the faked output. Give _FakeProc a readable stdout pipe (yields the fake line then EOF), wait()/poll()/pid, so the test exercises the real drain path on both the python and bash bypass branches. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: persist sandbox path heals across runs and suppress nested hidden lists * Studio: run tool Python child unbuffered (-u) so unflushed prints stream live A long-running snippet doing bare print() without flush=True never reached the live-output pane: CPython block-buffers stdout when writing to a pipe, so _drain_process_output's readline() saw nothing until the buffer filled or the process exited. Launch the child with the interpreter -u flag so stdout is unbuffered and each print streams as it is produced. -u is applied unconditionally on both the streaming and non-streaming path, so the child invocation stays byte-identical with and without streaming and the final joined result is unchanged (buffering/timing only). Unlike the earlier PYTHONUNBUFFERED=1 env injection that was removed, -u does not pollute the child's os.environ and is not visible via os.getenv. * Render only the selected main-content subtree in html_to_markdown The main-content heuristic sized each <article>/<main> candidate individually to pick the largest subtree, but then rendered every matching tag in the document. A page with one real article plus sibling related-post cards or comment threads passed the size gate on the real article yet still emitted the unrelated siblings. Size and render the same chosen subtree so only the selected main-content subtree reaches the output. * Studio: tighten chat-tool-streaming fix comments * Studio: store tool-output-scope separators as unicode escapes The pane-scope and tool-output-key separators were literal NUL (0x00) bytes, which made git treat the file as binary and hide its diff and blame. Write them as \u0000 escapes instead; the runtime key value is unchanged. * Studio: bound tool-stream teardown when the client disconnects stream_tool_execution ran its yield loop with no try/finally, so a gen.close() on client disconnect (GeneratorExit at a yield) skipped the worker join and never signalled cancellation. A tool that does not poll cancel_event mid-flight (web_search, MCP, search_knowledge_base) then kept request teardown blocked until the tool's own timeout. Thread the request cancel_event into the wrapper, set it only on the abnormal-exit path so a clean multi-tool turn is unaffected, and bound the worker join to a few seconds; the daemon worker cannot outlive the process. * Studio: sandbox path remap no longer masks missing reads The sandbox sitecustomize shim remapped code-interpreter prefixes (/mnt/data, /workspace, ...) onto the working directory for every open mode, including reads. A read of a path that truly did not exist was silently redirected onto a same-basename workdir file instead of raising on the path the model used, hiding real missing-input errors. Remap writes and creates as before, but remap a read only when the mapped workdir target already exists (re-reading a just-written artifact); otherwise keep the original absolute path so the failure stays truthful. * Studio: bound web fetch with one overall deadline and cancellation The web fetch applied timeouts per network operation, so a GitHub README API attempt plus its HTML fallback plus up to five redirect hops could run well past the tool timeout, and nothing aborted once the client had disconnected. Add a single wall-clock deadline shared across the API attempt, the fallback, every redirect hop and the body read, cap each hop's socket timeout at the time left on the budget, and poll cancel_event. SSRF host pinning, per-hop redirect revalidation, the five-hop cap and the size cap are unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: keep tool-stream teardown off the event loop on disconnect The bounded worker join added for disconnect safety still ran on an abnormal close, so a client disconnect could wait the full join timeout; and the safetensors and Anthropic tool streams closed their generator synchronously on the event loop, unlike the GGUF path. On abnormal exit the daemon worker is abandoned, so join with a zero timeout instead of waiting; offload the safetensors and Anthropic gen.close to a thread to match GGUF; and surface a heartbeat as soon as cancel_event is set while the worker is silent so the route regains control at once instead of after a heartbeat interval. * Studio: extend the web-fetch deadline to DNS, the body read, and search The overall fetch deadline did not cover host resolution or the response body read, and query-mode web_search ignored cancellation. Resolve hosts (initial and every redirect) on a budget-polled helper so a slow or pre-cancelled getaddrinfo aborts on time; read the capped body in chunks with the budget re-checked between them so a slow-drip server cannot stretch a single read past the deadline; and gate the blocking DDGS query on cancel_event on both sides. SSRF host pinning, per-hop redirect revalidation, the five-hop cap and the size cap are unchanged. * Studio: defer the sandbox remap notice and tighten os.open create flags The one-shot remap notice fired while computing the mapping, so a read that kept its original path emitted a false notice and spent the notice a later genuine remap needed. Only emit it once _remap_open commits to the redirect. Separately, os.open classified O_TRUNC / O_APPEND without O_CREAT as creating, but those cannot create a missing file, so a missing target now stays truthful (only O_CREAT maps to the creating mode). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: convert only genuine HTML README bodies, not Markdown with a leading block tag The GitHub README API returns the raw file, almost always Markdown. _looks_like_html classified a Markdown README opening with a block tag (<ul>, <ol>, <dl>, <pre>, <blockquote>) as HTML, so _fetch_page_text ran it through html_to_markdown and collapsed its headings, lists and fenced code into a single line. Sniff the README body with a stricter document-level check (doctype or a leading <html>/<head>/<body>) so only a real .html README is converted; the general page path is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Surface unclassified mid-stream Anthropic errors as SSE error events The local Anthropic tool-stream and plain-stream paths called _anthropic_stream_error_event(e) with force defaulting to False, so an unclassified mid-stream failure (llama-server crash, decode OOM, a dropped upstream socket) returned no event. The except block then fell through to emitter.finish(), emitting a normal message_delta and message_stop that masked a truncated turn as a clean finish. Pass force = True at both fall-through sites so an unclassified failure emits a 500 SSE error event and returns, matching the Anthropic passthrough path that already forces it. Add regression tests covering both stream paths. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: give each tool run a unique part id so finished cards keep their own output Backend tool ids restart at call_0 every assistant response, and the transient toolLiveOutput/toolFullOutput store maps were keyed by pane scope plus that bare backend id. Two turns in the same pane therefore shared one key: the stale-clear at tool_start only guards the forward direction, so when a later call_0 finished and wrote its preserved full output, every earlier still-mounted finished card reading the same key re-rendered and displayed the newer tool's output instead of its own. Mint one per-run-unique part id per backend id (call_0:<uuid>) and route tool_start/output/args/end through a single resolver so all events for a call resolve the same id. The durable part carries the unique id, so the finished-card readers derive a collision-free key with no change, and the awaiting-confirmation path keeps its own synthesized id. Outbound replay stays paired (the assistant tool_call id and the role=tool result tool_call_id both come from the part id) and gains unique ids across turns, which strict providers require. * Studio: tighten PR comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
This commit is contained in:
Daniel Han
GitHub
parent
9de84888cb
commit
73af334d11
28 changed files with 6558 additions and 262 deletions
|
|
@ -7,6 +7,11 @@ Minimal HTML-to-Markdown converter using only the standard library.
|
|||
Replaces the external ``html2text`` (GPL-3.0) dependency with a ~250-line
|
||||
``html.parser.HTMLParser`` subclass. Covers headings, links, bold/italic,
|
||||
lists, tables, blockquotes, code blocks, and entity decoding.
|
||||
|
||||
``main_content=True`` also applies a readability-style heuristic: scope
|
||||
conversion to the page's ``<article>`` (else ``<main>``) subtree when it
|
||||
carries substantial text, and strip known boilerplate fragments (skip-links,
|
||||
error placeholders, session banners, cookie prompts) from the result.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -27,8 +32,138 @@ _SKIP_TAGS = frozenset(
|
|||
"math",
|
||||
"nav",
|
||||
"footer",
|
||||
# Never-rendered / form-chrome elements, not page content.
|
||||
"template",
|
||||
"dialog",
|
||||
"button",
|
||||
"select",
|
||||
"datalist",
|
||||
}
|
||||
)
|
||||
# <aside> is NOT skipped: docs use it for admonition callouts (real content);
|
||||
# page-furniture asides are excluded by the main-content scoping pass instead.
|
||||
|
||||
# Void elements never produce an end tag, so they must not join the
|
||||
# open-element stack used to bound hidden subtrees.
|
||||
_VOID_TAGS = frozenset(
|
||||
{
|
||||
"area",
|
||||
"base",
|
||||
"br",
|
||||
"col",
|
||||
"embed",
|
||||
"hr",
|
||||
"img",
|
||||
"input",
|
||||
"link",
|
||||
"meta",
|
||||
"param",
|
||||
"source",
|
||||
"track",
|
||||
"wbr",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _style_hides_element(style: str) -> bool:
|
||||
"""True when an inline ``style`` sets ``display:none`` / ``visibility:hidden``.
|
||||
|
||||
Parsed per property so an unrelated value that merely contains ``none`` is
|
||||
not misread as hidden."""
|
||||
lowered = style.lower()
|
||||
if "none" not in lowered and "hidden" not in lowered:
|
||||
return False
|
||||
for declaration in style.split(";"):
|
||||
prop, sep, value = declaration.partition(":")
|
||||
if not sep:
|
||||
continue
|
||||
prop = prop.strip().lower()
|
||||
# Drop any !important flag and keep the first token of the value.
|
||||
value = value.split("!", 1)[0].strip().lower()
|
||||
if prop == "display" and value == "none":
|
||||
return True
|
||||
if prop == "visibility" and value == "hidden":
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _is_hidden_element(attr_dict: dict) -> bool:
|
||||
"""True when the element is not rendered: ``hidden`` attribute,
|
||||
``aria-hidden="true"``, or an inline ``style`` hiding it. Such JS-only
|
||||
placeholders ship in the HTML but must not reach the output. ``hidden`` is
|
||||
enumerated: any present value (even ``hidden="false"``) means not rendered."""
|
||||
if "hidden" in attr_dict:
|
||||
return True
|
||||
if (attr_dict.get("aria-hidden") or "").strip().lower() == "true":
|
||||
return True
|
||||
return _style_hides_element(attr_dict.get("style") or "")
|
||||
|
||||
|
||||
# HTML5 optional end tags: a listed start tag implicitly closes an open element
|
||||
# of the key type (as browsers do), else an unclosed ``<p hidden>``/``<li hidden>``
|
||||
# swallows every following sibling. Keys: closable elements; values: closers.
|
||||
_P_CLOSING_TAGS = frozenset(
|
||||
{
|
||||
"address",
|
||||
"article",
|
||||
"aside",
|
||||
"blockquote",
|
||||
"details",
|
||||
"div",
|
||||
"dl",
|
||||
"fieldset",
|
||||
"figcaption",
|
||||
"figure",
|
||||
"footer",
|
||||
"form",
|
||||
"h1",
|
||||
"h2",
|
||||
"h3",
|
||||
"h4",
|
||||
"h5",
|
||||
"h6",
|
||||
"header",
|
||||
"hgroup",
|
||||
"hr",
|
||||
"main",
|
||||
"menu",
|
||||
"nav",
|
||||
"ol",
|
||||
"p",
|
||||
"pre",
|
||||
"section",
|
||||
"table",
|
||||
"ul",
|
||||
}
|
||||
)
|
||||
_IMPLICIT_CLOSERS: dict = {
|
||||
"p": _P_CLOSING_TAGS,
|
||||
"li": frozenset({"li"}),
|
||||
"dt": frozenset({"dt", "dd"}),
|
||||
"dd": frozenset({"dt", "dd"}),
|
||||
"tr": frozenset({"tr"}),
|
||||
"td": frozenset({"td", "th", "tr"}),
|
||||
"th": frozenset({"td", "th", "tr"}),
|
||||
"option": frozenset({"option", "optgroup"}),
|
||||
"optgroup": frozenset({"optgroup"}),
|
||||
}
|
||||
|
||||
|
||||
# Item tag -> container tags that re-scope it: a nested container makes an inner
|
||||
# item a descendant, not an optional-close sibling, so recovery must stop there
|
||||
# rather than close (and un-hide) the outer item and leak its nested content.
|
||||
_CLOSE_BARRIERS: dict = {
|
||||
"li": frozenset({"ul", "ol", "menu"}),
|
||||
"dt": frozenset({"dl"}),
|
||||
"dd": frozenset({"dl"}),
|
||||
"tr": frozenset({"table"}),
|
||||
"td": frozenset({"table"}),
|
||||
"th": frozenset({"table"}),
|
||||
"option": frozenset({"select", "datalist"}),
|
||||
"optgroup": frozenset({"select", "datalist"}),
|
||||
}
|
||||
|
||||
|
||||
_BLOCK_TAGS = frozenset(
|
||||
{
|
||||
"p",
|
||||
|
|
@ -51,13 +186,33 @@ _INLINE_EMPHASIS = {"strong": "**", "b": "**", "em": "*", "i": "*"}
|
|||
|
||||
|
||||
class _MarkdownRenderer(HTMLParser):
|
||||
"""HTMLParser subclass that emits Markdown tokens into a list."""
|
||||
"""HTMLParser subclass that emits Markdown tokens into a list.
|
||||
|
||||
def __init__(self):
|
||||
``scope_tags`` restricts emission to the subtree(s) of the given tags
|
||||
(e.g. ``{"article"}``): outside them every handler is a no-op, which is
|
||||
how the readability-style main-content pass drops page furniture.
|
||||
"""
|
||||
|
||||
def __init__(self, scope_tags: frozenset[str] | None = None):
|
||||
super().__init__(convert_charrefs = False)
|
||||
self._out: list[str] = []
|
||||
self._skip_depth: int = 0
|
||||
|
||||
# Main-content scoping: emit only while inside a scope tag.
|
||||
self._scope_tags = scope_tags
|
||||
self._scope_depth: int = 0
|
||||
|
||||
# Output boundaries per top-level scope element, so a caller can size each
|
||||
# candidate alone and a swarm of tiny sibling cards can't clear the threshold.
|
||||
self.scope_segments: list[str] = []
|
||||
self._scope_seg_start: int | None = None
|
||||
|
||||
# Hidden-subtree tracking: stack of open non-void tags plus the indices
|
||||
# where a hidden element started. End tags pop to the matching tag, so
|
||||
# an omitted </p>/<li> close cannot leave the renderer stuck hidden.
|
||||
self._open_tags: list[str] = []
|
||||
self._hidden_marks: list[int] = []
|
||||
|
||||
# Link state
|
||||
self._link_href: str | None = None
|
||||
self._link_text_parts: list[str] = []
|
||||
|
|
@ -150,16 +305,95 @@ class _MarkdownRenderer(HTMLParser):
|
|||
# ------------------------------------------------------------------
|
||||
# Tag handlers
|
||||
# ------------------------------------------------------------------
|
||||
# Structural bookkeeping shared by every start tag (skip/hidden/scope).
|
||||
def _close_implicit(self, tag: str) -> None:
|
||||
"""HTML5 optional-end-tag recovery for a start tag about to open.
|
||||
|
||||
Pops each implicitly-closed ancestor (and its hidden marks), scanning the
|
||||
whole stack so an open ``<p>``/``<li>`` still closes under an unclosed inline
|
||||
``<span>``. Stops at a ``_CLOSE_BARRIERS`` container so recovery never crosses
|
||||
a nested list/table/dl and leaks the outer item's hidden content. Runs even
|
||||
for skipped ``<nav>``/``<footer>``, which also close ``<p>``."""
|
||||
barriers = _CLOSE_BARRIERS.get(tag, ())
|
||||
while True:
|
||||
close_at = None
|
||||
for i in range(len(self._open_tags) - 1, -1, -1):
|
||||
name = self._open_tags[i]
|
||||
if tag in _IMPLICIT_CLOSERS.get(name, ()):
|
||||
close_at = i
|
||||
break
|
||||
# A barrier container re-scopes the item; stop before it.
|
||||
if name in barriers:
|
||||
break
|
||||
if close_at is None:
|
||||
break
|
||||
del self._open_tags[close_at:]
|
||||
while self._hidden_marks and self._hidden_marks[-1] >= close_at:
|
||||
self._hidden_marks.pop()
|
||||
|
||||
def _enter_tag(self, tag: str, attr_dict: dict) -> bool:
|
||||
"""Track open/hidden/scope state; return True when the tag's content
|
||||
should be rendered (False = suppressed). Caller runs ``_close_implicit``
|
||||
first so recovery also fires for skipped tags."""
|
||||
if tag not in _VOID_TAGS:
|
||||
self._open_tags.append(tag)
|
||||
if _is_hidden_element(attr_dict):
|
||||
self._hidden_marks.append(len(self._open_tags) - 1)
|
||||
elif _is_hidden_element(attr_dict):
|
||||
# Void elements never join the stack, so suppress a hidden one inline.
|
||||
return False
|
||||
if self._scope_tags is not None and tag in self._scope_tags:
|
||||
if self._scope_depth == 0:
|
||||
self._scope_seg_start = len(self._out)
|
||||
self._scope_depth += 1
|
||||
if self._hidden_marks:
|
||||
return False
|
||||
if self._scope_tags is not None and self._scope_depth == 0:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _exit_tag(self, tag: str) -> bool:
|
||||
"""Pop to the matching open tag; return True when the end tag should
|
||||
be rendered (False = it closed inside a hidden / out-of-scope region)."""
|
||||
suppressed = bool(self._hidden_marks) or (
|
||||
self._scope_tags is not None and self._scope_depth == 0
|
||||
)
|
||||
if tag not in _VOID_TAGS:
|
||||
# Pop to the innermost matching open tag (recovers omitted closes).
|
||||
for i in range(len(self._open_tags) - 1, -1, -1):
|
||||
if self._open_tags[i] == tag:
|
||||
del self._open_tags[i:]
|
||||
while self._hidden_marks and self._hidden_marks[-1] >= i:
|
||||
self._hidden_marks.pop()
|
||||
break
|
||||
if self._scope_tags is not None and tag in self._scope_tags and self._scope_depth > 0:
|
||||
self._scope_depth -= 1
|
||||
if self._scope_depth == 0 and self._scope_seg_start is not None:
|
||||
self.scope_segments.append("".join(self._out[self._scope_seg_start :]))
|
||||
self._scope_seg_start = None
|
||||
return not suppressed
|
||||
|
||||
def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
|
||||
tag = tag.lower()
|
||||
|
||||
if self._skip_depth:
|
||||
# Inside a skipped subtree: only track nested skip depth.
|
||||
if tag in _SKIP_TAGS:
|
||||
self._skip_depth += 1
|
||||
return
|
||||
|
||||
# Recover optional end tags before the skip decision: a skipped
|
||||
# <nav>/<footer> still implicitly closes an open <p>, releasing its
|
||||
# hidden mark so following siblings render.
|
||||
self._close_implicit(tag)
|
||||
|
||||
if tag in _SKIP_TAGS:
|
||||
self._skip_depth += 1
|
||||
return
|
||||
if self._skip_depth:
|
||||
return
|
||||
|
||||
attr_dict = dict(attrs)
|
||||
if not self._enter_tag(tag, attr_dict):
|
||||
return
|
||||
|
||||
if tag in _HEADING_TAGS:
|
||||
level = int(tag[1])
|
||||
|
|
@ -250,6 +484,9 @@ class _MarkdownRenderer(HTMLParser):
|
|||
if self._skip_depth:
|
||||
return
|
||||
|
||||
if not self._exit_tag(tag):
|
||||
return
|
||||
|
||||
if tag in _HEADING_TAGS:
|
||||
self._emit("\n\n")
|
||||
|
||||
|
|
@ -308,8 +545,13 @@ class _MarkdownRenderer(HTMLParser):
|
|||
# ------------------------------------------------------------------
|
||||
# Text / entity handlers
|
||||
# ------------------------------------------------------------------
|
||||
def _text_suppressed(self) -> bool:
|
||||
if self._skip_depth or self._hidden_marks:
|
||||
return True
|
||||
return self._scope_tags is not None and self._scope_depth == 0
|
||||
|
||||
def handle_data(self, data: str) -> None:
|
||||
if self._skip_depth:
|
||||
if self._text_suppressed():
|
||||
return
|
||||
if self._in_pre:
|
||||
self._pre_parts.append(data)
|
||||
|
|
@ -326,12 +568,12 @@ class _MarkdownRenderer(HTMLParser):
|
|||
self._emit(text)
|
||||
|
||||
def handle_entityref(self, name: str) -> None:
|
||||
if self._skip_depth:
|
||||
if self._text_suppressed():
|
||||
return
|
||||
self._emit(html.unescape(f"&{name};"))
|
||||
|
||||
def handle_charref(self, name: str) -> None:
|
||||
if self._skip_depth:
|
||||
if self._text_suppressed():
|
||||
return
|
||||
self._emit(html.unescape(f"&#{name};"))
|
||||
|
||||
|
|
@ -366,6 +608,14 @@ class _MarkdownRenderer(HTMLParser):
|
|||
else:
|
||||
self._out.append("\n\n" + prefixed + "\n\n")
|
||||
|
||||
# A scope left open by truncated HTML never reached _exit_tag, so its output
|
||||
# never joined scope_segments and would score 0. Flush the still-open segment
|
||||
# here (after the side-buffers) so a truncated main-content page is scored.
|
||||
if self._scope_seg_start is not None:
|
||||
self.scope_segments.append("".join(self._out[self._scope_seg_start :]))
|
||||
self._scope_seg_start = None
|
||||
self._scope_depth = 0
|
||||
|
||||
|
||||
# Post-processing
|
||||
def _cleanup(text: str) -> str:
|
||||
|
|
@ -399,17 +649,124 @@ def _cleanup(text: str) -> str:
|
|||
return "\n".join(out).strip()
|
||||
|
||||
|
||||
# Public API
|
||||
def html_to_markdown(source_html: str) -> str:
|
||||
"""Convert HTML to Markdown (headings, links, emphasis, lists, tables, blockquotes, code, entities).
|
||||
# Known boilerplate fragments stripped from main-content conversions, matched
|
||||
# only against short lines. Sources: GitHub page furniture / client-side error
|
||||
# placeholders, skip-links, cookie banners.
|
||||
_BOILERPLATE_FRAGMENTS = (
|
||||
"skip to content",
|
||||
"skip to main content",
|
||||
"there was an error while loading",
|
||||
"please reload this page",
|
||||
"you can't perform that action at this time",
|
||||
"you signed in with another tab or window",
|
||||
"you signed out in another tab or window",
|
||||
"you switched accounts on another tab or window",
|
||||
"reload to refresh your session",
|
||||
"you must be signed in to change notification settings",
|
||||
"uh oh!",
|
||||
"{{ message }}",
|
||||
"this website uses cookies",
|
||||
"we use cookies",
|
||||
"accept all cookies",
|
||||
"manage cookie preferences",
|
||||
)
|
||||
# Only shorter lines are eligible for boilerplate dropping; real content
|
||||
# sentences quoting a fragment run longer.
|
||||
_BOILERPLATE_MAX_LINE_CHARS = 300
|
||||
|
||||
``<script>``, ``<style>``, and ``<head>`` are stripped entirely.
|
||||
"""
|
||||
# Normalize line endings before parsing.
|
||||
source_html = source_html.replace("\r\n", "\n").replace("\r", "\n")
|
||||
renderer = _MarkdownRenderer()
|
||||
# Normalized furniture phrases for whole-segment matching. See _line_is_boilerplate.
|
||||
_BOILERPLATE_NORMALIZED = frozenset(
|
||||
re.sub(r"\s+", " ", fragment).strip().casefold().rstrip(".!:")
|
||||
for fragment in _BOILERPLATE_FRAGMENTS
|
||||
)
|
||||
|
||||
|
||||
def _line_is_boilerplate(line: str) -> bool:
|
||||
"""True only when a whole line is composed of known furniture phrases.
|
||||
|
||||
Splits on sentence terminators and requires every segment to be furniture, so a
|
||||
line stacking several phrases is dropped while prose that merely quotes one is
|
||||
kept (its other words leave a non-furniture segment)."""
|
||||
normalized = re.sub(r"\s+", " ", line).strip().casefold()
|
||||
if not normalized:
|
||||
return False
|
||||
segments = [segment.strip().rstrip(".!:") for segment in re.split(r"[.!]", normalized)]
|
||||
segments = [segment for segment in segments if segment]
|
||||
return bool(segments) and all(segment in _BOILERPLATE_NORMALIZED for segment in segments)
|
||||
|
||||
|
||||
def _strip_boilerplate_lines(text: str) -> str:
|
||||
"""Drop short lines that consist entirely of known page-furniture phrases.
|
||||
|
||||
Fenced code blocks are preserved verbatim: boilerplate never renders
|
||||
inside ``<pre>``, while READMEs legitimately quote error strings."""
|
||||
out: list[str] = []
|
||||
in_fence = False
|
||||
for line in text.split("\n"):
|
||||
if line.lstrip().startswith("```"):
|
||||
in_fence = not in_fence
|
||||
out.append(line)
|
||||
continue
|
||||
if not in_fence and len(line) <= _BOILERPLATE_MAX_LINE_CHARS and _line_is_boilerplate(line):
|
||||
continue
|
||||
out.append(line)
|
||||
# Collapse blank runs the dropped lines may have left behind.
|
||||
return re.sub(r"\n{3,}", "\n\n", "\n".join(out)).strip()
|
||||
|
||||
|
||||
def _render(source_html: str, scope_tags: frozenset[str] | None) -> str:
|
||||
renderer = _MarkdownRenderer(scope_tags = scope_tags)
|
||||
renderer.feed(source_html)
|
||||
renderer.close()
|
||||
renderer.flush_pending()
|
||||
raw = "".join(renderer._out)
|
||||
return _cleanup(raw)
|
||||
|
||||
|
||||
def _select_main_scope_render(source_html: str, tag: str) -> tuple[int, str]:
|
||||
"""Length and boilerplate-stripped render of the largest single ``<tag>``
|
||||
subtree. Sizing candidates one at a time stops many tiny sibling cards from
|
||||
clearing the threshold together, and returning that one subtree keeps
|
||||
unrelated siblings (related cards, comment threads) out of the output."""
|
||||
renderer = _MarkdownRenderer(scope_tags = frozenset({tag}))
|
||||
renderer.feed(source_html)
|
||||
renderer.close()
|
||||
renderer.flush_pending()
|
||||
best_len = 0
|
||||
best_render = ""
|
||||
for seg in renderer.scope_segments:
|
||||
rendered = _strip_boilerplate_lines(_cleanup(seg))
|
||||
if len(rendered) > best_len:
|
||||
best_len = len(rendered)
|
||||
best_render = rendered
|
||||
return best_len, best_render
|
||||
|
||||
|
||||
# A scoped conversion below this size is judged not to be the page's main
|
||||
# content (e.g. an empty <article> stub) and the next candidate is tried.
|
||||
_MIN_MAIN_CONTENT_CHARS = 200
|
||||
|
||||
|
||||
# Public API
|
||||
def html_to_markdown(source_html: str, *, main_content: bool = False) -> str:
|
||||
"""Convert HTML to Markdown (headings, links, emphasis, lists, tables, blockquotes, code, entities).
|
||||
|
||||
``<script>``, ``<style>``, and ``<head>`` are stripped entirely, as are
|
||||
subtrees hidden from rendering (``hidden`` / ``aria-hidden="true"``).
|
||||
|
||||
``main_content=True`` applies a readability-style heuristic for page
|
||||
fetches: prefer the ``<article>`` subtree (GitHub renders READMEs there),
|
||||
then ``<main>``, falling back to the whole document, and strip known
|
||||
boilerplate fragments from the result.
|
||||
"""
|
||||
# Normalize line endings before parsing.
|
||||
source_html = source_html.replace("\r\n", "\n").replace("\r", "\n")
|
||||
if main_content:
|
||||
for scope_tag in ("article", "main"):
|
||||
# Render only the chosen subtree so sibling <article>/<main>
|
||||
# elements do not leak in once the largest passes the size gate.
|
||||
length, rendered = _select_main_scope_render(source_html, scope_tag)
|
||||
if length >= _MIN_MAIN_CONTENT_CHARS:
|
||||
return rendered
|
||||
return _strip_boilerplate_lines(_render(source_html, None))
|
||||
return _render(source_html, None)
|
||||
|
|
|
|||
|
|
@ -301,6 +301,27 @@ def _gguf_has_genuine_tool_signal(text: str, signals, active_tools: list[dict])
|
|||
return False
|
||||
|
||||
|
||||
_TEXT_TOOL_NAME_RE = re.compile(r'"name"\s*:\s*"([\w.\-]+)"')
|
||||
_TEXT_TOOL_GEMMA_RE = re.compile(r"\s*call:([\w.\-]+)")
|
||||
_TEXT_TOOL_REHEARSAL_RE = re.compile(r"\s*([\w.\-]+)\s*\[ARGS\]")
|
||||
|
||||
|
||||
def _sniff_text_tool_name(text: str, enabled_names: set) -> str:
|
||||
"""Best-effort tool name from a partially drained TEXT tool call, gated on
|
||||
enabled names so prose can never spawn a card. Used only to open the live
|
||||
argument pane early; the authoritative parse still happens at stream end."""
|
||||
m = _TEXT_TOOL_NAME_RE.search(text[:4096])
|
||||
if m and m.group(1) in enabled_names:
|
||||
return m.group(1)
|
||||
m = _TEXT_TOOL_GEMMA_RE.match(text[:256])
|
||||
if m and m.group(1) in enabled_names:
|
||||
return m.group(1)
|
||||
m = _TEXT_TOOL_REHEARSAL_RE.match(text[:256])
|
||||
if m and m.group(1) in enabled_names:
|
||||
return m.group(1)
|
||||
return ""
|
||||
|
||||
|
||||
def _is_rehearsal_prefix(stripped: str, active_tools: list[dict]) -> bool:
|
||||
"""True if ``stripped`` is a (possibly partial) prefix of ``NAME[ARGS]`` for an
|
||||
active tool -- the bare tool name arriving in its own chunk before ``[ARGS]{...}``.
|
||||
|
|
@ -8984,6 +9005,7 @@ class LlamaCppBackend:
|
|||
{"type": "content", "text": "token"} -- streamed content tokens (cumulative)
|
||||
{"type": "reasoning", "text": "token"} -- streamed reasoning tokens (cumulative)
|
||||
"""
|
||||
from core.inference.tool_stream_exec import accepts_output_callback, stream_tool_execution
|
||||
from core.inference.tools import (
|
||||
build_rag_autoinject,
|
||||
execute_tool,
|
||||
|
|
@ -9268,6 +9290,17 @@ class LlamaCppBackend:
|
|||
provisional_started_tool_calls: dict[str, str] = {}
|
||||
resolved_provisional_tool_call_ids: set[str] = set()
|
||||
_suppress_visible_output = _forced_tool_call_pending
|
||||
# Cards that already got their first tool_args event; later
|
||||
# fragments stream individually so a big payload isn't a dead spinner.
|
||||
arg_streamed_tool_call_ids: set[str] = set()
|
||||
# TEXT tool-call path: a committed call's raw text streams to a
|
||||
# provisional card under the parser's first-call id ("call_0"), so
|
||||
# the final tool_start reconciles in place.
|
||||
_text_args_call_start = -1
|
||||
_text_args_streamed_upto = -1
|
||||
_text_args_id = ""
|
||||
_text_args_name = ""
|
||||
_confirm_gated_iteration = bool(confirm_tool_calls) and not bypass_permissions
|
||||
|
||||
with self._open_stream(url, payload, cancel_event) as (
|
||||
response,
|
||||
|
|
@ -9430,6 +9463,29 @@ class LlamaCppBackend:
|
|||
provisional = True,
|
||||
),
|
||||
}
|
||||
# Stream argument text so the UI shows the code being
|
||||
# written: first event the backlog, later the fragment.
|
||||
# Display only; accumulator untouched.
|
||||
if current_id in provisional_started_tool_calls:
|
||||
if current_id not in arg_streamed_tool_call_ids:
|
||||
arg_streamed_tool_call_ids.add(current_id)
|
||||
_args_backlog = tool_calls_acc[idx]["function"].get(
|
||||
"arguments", ""
|
||||
)
|
||||
if _args_backlog:
|
||||
yield {
|
||||
"type": "tool_args",
|
||||
"tool_call_id": current_id,
|
||||
"tool_name": current_name,
|
||||
"text": _args_backlog,
|
||||
}
|
||||
elif func.get("arguments"):
|
||||
yield {
|
||||
"type": "tool_args",
|
||||
"tool_call_id": current_id,
|
||||
"tool_name": current_name,
|
||||
"text": func["arguments"],
|
||||
}
|
||||
continue
|
||||
|
||||
# ── Reasoning tokens ──
|
||||
|
|
@ -9468,7 +9524,60 @@ class LlamaCppBackend:
|
|||
content_accum += token
|
||||
|
||||
if detect_state == _S_DRAINING:
|
||||
pass # accumulate silently
|
||||
# Accumulate silently for parsing, but stream the drained
|
||||
# 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 _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
|
||||
):
|
||||
_text_args_id = "call_0"
|
||||
_text_args_name = _sniffed
|
||||
if (
|
||||
_text_args_id
|
||||
not in provisional_started_tool_calls
|
||||
):
|
||||
provisional_started_tool_calls[
|
||||
_text_args_id
|
||||
] = _sniffed
|
||||
yield {
|
||||
"type": "tool_start",
|
||||
"tool_name": _sniffed,
|
||||
"tool_call_id": _text_args_id,
|
||||
"arguments": {},
|
||||
"provenance": tool_event_provenance(
|
||||
provisional = True,
|
||||
),
|
||||
}
|
||||
yield {
|
||||
"type": "tool_args",
|
||||
"tool_call_id": _text_args_id,
|
||||
"tool_name": _sniffed,
|
||||
"text": _call_text,
|
||||
}
|
||||
_text_args_streamed_upto = len(content_accum)
|
||||
elif len(content_accum) > _text_args_streamed_upto:
|
||||
yield {
|
||||
"type": "tool_args",
|
||||
"tool_call_id": _text_args_id,
|
||||
"tool_name": _text_args_name,
|
||||
"text": content_accum[
|
||||
_text_args_streamed_upto:
|
||||
],
|
||||
}
|
||||
_text_args_streamed_upto = len(content_accum)
|
||||
|
||||
elif detect_state == _S_STREAMING:
|
||||
if in_thinking:
|
||||
|
|
@ -9575,6 +9684,11 @@ class LlamaCppBackend:
|
|||
# without yielding. A live <think> prefix is
|
||||
# separate from it -- close that.
|
||||
detect_state = _S_DRAINING
|
||||
# Call text begins at the held buffer
|
||||
# (live arg display only; UI extracts the code).
|
||||
_text_args_call_start = len(content_accum) - len(
|
||||
content_buffer
|
||||
)
|
||||
if _close_streamed_think():
|
||||
yield {
|
||||
"type": "content",
|
||||
|
|
@ -9602,6 +9716,11 @@ class LlamaCppBackend:
|
|||
"text": cleaned,
|
||||
}
|
||||
detect_state = _S_DRAINING
|
||||
# Live-arg display starts at the held buffer
|
||||
# (visible prefix flushed above; UI extracts the code).
|
||||
_text_args_call_start = len(content_accum) - len(
|
||||
content_buffer
|
||||
)
|
||||
elif _hold_buffer or (
|
||||
is_prefix
|
||||
and (
|
||||
|
|
@ -9849,9 +9968,20 @@ class LlamaCppBackend:
|
|||
f"{'structured delta' if has_structured_tc else 'content text'}"
|
||||
)
|
||||
if not tool_calls:
|
||||
# DRAINING but no tool calls (false positive). Merge
|
||||
# accumulated metrics from prior tool iterations so
|
||||
# they aren't silently dropped.
|
||||
# DRAINING but no tool calls (false positive): close any
|
||||
# provisional cards (a sniff can open one whose call never
|
||||
# parses); without a tool_end the card spins forever.
|
||||
for _pid, _pname in provisional_started_tool_calls.items():
|
||||
if _pid not in resolved_provisional_tool_call_ids:
|
||||
resolved_provisional_tool_call_ids.add(_pid)
|
||||
yield {
|
||||
"type": "tool_end",
|
||||
"tool_name": _pname,
|
||||
"tool_call_id": _pid,
|
||||
"result": "",
|
||||
"provenance": tool_event_provenance(provisional = True),
|
||||
}
|
||||
# Merge metrics from prior tool iterations so they aren't dropped.
|
||||
yield {"type": "status", "text": ""}
|
||||
if content_accum:
|
||||
# Strip leaked tool-call XML before yielding.
|
||||
|
|
@ -9912,9 +10042,22 @@ class LlamaCppBackend:
|
|||
assistant_msg: dict = {"role": "assistant", "content": content_text}
|
||||
assistant_appended = False
|
||||
|
||||
# The text-path provisional card uses the parser's default id ("call_0");
|
||||
# a Mistral-style call carries its own id and would open a duplicate. Reuse
|
||||
# the card's id for the first matching call (same tool name) to reconcile.
|
||||
_text_provisional_id = _text_args_id if not has_structured_tc else ""
|
||||
|
||||
for tc in tool_calls or []:
|
||||
func = tc.get("function", {})
|
||||
tool_name = func.get("name", "")
|
||||
if (
|
||||
_text_provisional_id
|
||||
and _text_provisional_id in provisional_started_tool_calls
|
||||
and _text_provisional_id not in resolved_provisional_tool_call_ids
|
||||
and tc.get("id") not in provisional_started_tool_calls
|
||||
and provisional_started_tool_calls[_text_provisional_id] == tool_name
|
||||
):
|
||||
tc = {**tc, "id": _text_provisional_id}
|
||||
provisional_match = tc.get("id") in provisional_started_tool_calls
|
||||
decision = tool_controller.prepare_call(
|
||||
tc,
|
||||
|
|
@ -10026,15 +10169,33 @@ class LlamaCppBackend:
|
|||
):
|
||||
result = RAG_SEARCH_CAP_NUDGE
|
||||
else:
|
||||
result = execute_tool(
|
||||
decision.tool_name,
|
||||
decision.arguments,
|
||||
# Execute in a worker thread so live stdout chunks and heartbeats
|
||||
# stream while the tool blocks (the SSE route turns heartbeats into
|
||||
# keepalives). Result is byte-identical to a direct call.
|
||||
def _invoke_tool(_output_callback, _decision = decision):
|
||||
# execute_tool is injectable and may be monkey-patched with the
|
||||
# pre-PR signature; forward output_callback only if it's accepted.
|
||||
kwargs = dict(
|
||||
cancel_event = cancel_event,
|
||||
timeout = _effective_timeout,
|
||||
session_id = session_id,
|
||||
thread_id = thread_id,
|
||||
rag_scope = rag_scope,
|
||||
disable_sandbox = bypass_permissions,
|
||||
)
|
||||
if accepts_output_callback(execute_tool):
|
||||
kwargs["output_callback"] = _output_callback
|
||||
return execute_tool(
|
||||
_decision.tool_name,
|
||||
_decision.arguments,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
result = yield from stream_tool_execution(
|
||||
_invoke_tool,
|
||||
tool_name = decision.tool_name,
|
||||
tool_call_id = decision.tool_call_id,
|
||||
cancel_event = cancel_event,
|
||||
timeout = _effective_timeout,
|
||||
session_id = session_id,
|
||||
thread_id = thread_id,
|
||||
rag_scope = rag_scope,
|
||||
disable_sandbox = bypass_permissions,
|
||||
)
|
||||
if decision.tool_name == "search_knowledge_base":
|
||||
_kb_search_count += 1
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ parses tool calls from the cumulative text and dispatches via
|
|||
"""
|
||||
|
||||
import bisect
|
||||
import inspect
|
||||
import re
|
||||
import threading
|
||||
from typing import Callable, Generator, Optional
|
||||
|
|
@ -60,6 +61,7 @@ from core.inference.tool_loop_controller import (
|
|||
status_for_tool,
|
||||
tool_event_provenance,
|
||||
)
|
||||
from core.inference.tool_stream_exec import stream_tool_execution
|
||||
from state.tool_approvals import (
|
||||
TOOL_REJECTED_MESSAGE,
|
||||
abort_tool_decision,
|
||||
|
|
@ -402,6 +404,22 @@ def _tool_event_provenance(**flags: object) -> dict[str, object]:
|
|||
return tool_event_provenance(**flags)
|
||||
|
||||
|
||||
def _accepts_output_callback(func: Callable[..., str]) -> bool:
|
||||
"""Whether an injectable ``execute_tool`` supports ``output_callback``.
|
||||
|
||||
The loop's ``execute_tool`` is a parameter (tests inject fakes), so forward
|
||||
the live-output kwarg only when the callable declares it or takes ``**kwargs``.
|
||||
"""
|
||||
try:
|
||||
sig = inspect.signature(func)
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
params = sig.parameters
|
||||
if "output_callback" in params:
|
||||
return True
|
||||
return any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
|
||||
|
||||
|
||||
def _call_single_turn(single_turn, conversation: list, active_tools: list[dict]):
|
||||
"""Call a single-turn generator with active tool schemas when supported."""
|
||||
try:
|
||||
|
|
@ -552,6 +570,9 @@ def run_safetensors_tool_loop(
|
|||
provisional_render_html_started = False
|
||||
provisional_resolved = False
|
||||
provisional_render_html_id = f"call_{next_call_id}"
|
||||
# Live-args offset for the provisional render_html card: the drained call
|
||||
# text streams as tool_args so the canvas shows the HTML being written.
|
||||
_live_args_streamed_upto = -1
|
||||
# When a human confirmation gate is active the real tool_start is keyed
|
||||
# by an approval id and carries awaiting_confirmation, so an early
|
||||
# provisional card (keyed by tool_call_id, no approval) would show the
|
||||
|
|
@ -623,6 +644,30 @@ def run_safetensors_tool_loop(
|
|||
"arguments": {},
|
||||
"provenance": _tool_event_provenance(provisional = True),
|
||||
}
|
||||
# Backlog first: everything drained so far.
|
||||
yield {
|
||||
"type": "tool_args",
|
||||
"tool_call_id": provisional_render_html_id,
|
||||
"tool_name": "render_html",
|
||||
"text": content_accum,
|
||||
}
|
||||
_live_args_streamed_upto = len(content_accum)
|
||||
elif (
|
||||
provisional_render_html_started
|
||||
and not provisional_resolved
|
||||
and _live_args_streamed_upto >= 0
|
||||
and len(content_accum) > _live_args_streamed_upto
|
||||
):
|
||||
# Still writing the call: stream the fragment so the canvas
|
||||
# renders live. Display only; content_accum still feeds the
|
||||
# stream-end parser verbatim.
|
||||
yield {
|
||||
"type": "tool_args",
|
||||
"tool_call_id": provisional_render_html_id,
|
||||
"tool_name": "render_html",
|
||||
"text": content_accum[_live_args_streamed_upto:],
|
||||
}
|
||||
_live_args_streamed_upto = len(content_accum)
|
||||
continue
|
||||
|
||||
if detect_state == _state_streaming:
|
||||
|
|
@ -663,6 +708,13 @@ def run_safetensors_tool_loop(
|
|||
"arguments": {},
|
||||
"provenance": _tool_event_provenance(provisional = True),
|
||||
}
|
||||
yield {
|
||||
"type": "tool_args",
|
||||
"tool_call_id": provisional_render_html_id,
|
||||
"tool_name": "render_html",
|
||||
"text": content_accum,
|
||||
}
|
||||
_live_args_streamed_upto = len(content_accum)
|
||||
continue
|
||||
cumulative_display = candidate
|
||||
cleaned = strip_tool_markup_streaming(
|
||||
|
|
@ -819,6 +871,13 @@ def run_safetensors_tool_loop(
|
|||
"arguments": {},
|
||||
"provenance": _tool_event_provenance(provisional = True),
|
||||
}
|
||||
yield {
|
||||
"type": "tool_args",
|
||||
"tool_call_id": provisional_render_html_id,
|
||||
"tool_name": "render_html",
|
||||
"text": content_accum,
|
||||
}
|
||||
_live_args_streamed_upto = len(content_accum)
|
||||
elif is_prefix and (is_rehearsal_prefix or len(stripped) < _MAX_BUFFER_CHARS):
|
||||
# A rehearsal prefix is self-bounded; the buffer cap must not cut long MCP names short.
|
||||
continue
|
||||
|
|
@ -1146,10 +1205,12 @@ def run_safetensors_tool_loop(
|
|||
):
|
||||
result = RAG_SEARCH_CAP_NUDGE
|
||||
else:
|
||||
try:
|
||||
result = execute_tool(
|
||||
decision.tool_name,
|
||||
decision.arguments,
|
||||
# Execute in a worker thread so live stdout chunks and heartbeats
|
||||
# stream while the tool blocks (the SSE route turns heartbeats into
|
||||
# keepalives). execute_tool is injectable; pass output_callback
|
||||
# only when it accepts it.
|
||||
def _invoke_tool(_output_callback, _decision = decision):
|
||||
kwargs = dict(
|
||||
cancel_event = cancel_event,
|
||||
timeout = eff_timeout,
|
||||
session_id = session_id,
|
||||
|
|
@ -1157,6 +1218,17 @@ def run_safetensors_tool_loop(
|
|||
rag_scope = rag_scope,
|
||||
disable_sandbox = bypass_permissions,
|
||||
)
|
||||
if _accepts_output_callback(execute_tool):
|
||||
kwargs["output_callback"] = _output_callback
|
||||
return execute_tool(_decision.tool_name, _decision.arguments, **kwargs)
|
||||
|
||||
try:
|
||||
result = yield from stream_tool_execution(
|
||||
_invoke_tool,
|
||||
tool_name = decision.tool_name,
|
||||
tool_call_id = decision.tool_call_id,
|
||||
cancel_event = cancel_event,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.exception("Tool %s raised: %s", decision.tool_name, exc)
|
||||
result = f"Error: tool raised an exception: {exc}"
|
||||
|
|
|
|||
6
studio/backend/core/inference/sandbox_site/__init__.py
Normal file
6
studio/backend/core/inference/sandbox_site/__init__.py
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
# Package marker only, so wheel builds ship this directory. It goes on the
|
||||
# sandbox PYTHONPATH so site machinery imports the sibling ``sitecustomize`` at
|
||||
# startup; nothing in the backend imports it directly.
|
||||
313
studio/backend/core/inference/sandbox_site/sitecustomize.py
Normal file
313
studio/backend/core/inference/sandbox_site/sitecustomize.py
Normal file
|
|
@ -0,0 +1,313 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Sandbox-side compatibility shim for ChatGPT code-interpreter paths.
|
||||
|
||||
Models habitually write to /mnt/data (or /mnt/outputs, /home/sandbox,
|
||||
/workspace), none of which exist in the Studio sandbox. This module sits on the
|
||||
sandbox subprocess PYTHONPATH (see ``tools._build_safe_env``), so it loads at
|
||||
interpreter startup in every sandboxed ``python`` run and any Python the
|
||||
``terminal`` tool launches.
|
||||
|
||||
It remaps those prefixes onto the CWD in ``open`` / ``io.open``, ``os.open``,
|
||||
``os.makedirs`` / ``os.mkdir`` and ``pathlib.Path.mkdir``. A write/create to a
|
||||
convention prefix always heals onto the CWD; a READ heals only when the mapped
|
||||
target already exists (re-reading an earlier write), so a genuinely missing
|
||||
input stays truthful on the path the model used instead of silently reading a
|
||||
same-basename workdir file. Since prefix lists cannot cover every invented path,
|
||||
``open`` / ``io.open`` also get a create-mode fallback: an absolute path outside
|
||||
the CWD whose parent is missing is redirected to the basename in the CWD. Reads
|
||||
and mkdir never use the fallback (an arbitrary absolute directory can legitimately
|
||||
succeed). It is collision-safe: it refuses to redirect onto an existing CWD file
|
||||
(letting open raise). The patch set (io.open, os.open, os.mkdir, Path.mkdir, and
|
||||
the <3.11 ``_NormalAccessor.open``) covers the low-level entry points pathlib
|
||||
routes through. A one-line stderr notice fires on the first remap, and everything
|
||||
is wrapped in try/except so a failure never breaks the interpreter.
|
||||
|
||||
Identical with and without output streaming because the child env is.
|
||||
"""
|
||||
|
||||
import builtins
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Code-interpreter convention prefixes. Remapping is gated on the prefix being
|
||||
# ABSENT (see _remap) so a genuine host mount / user dir is never shadowed.
|
||||
_PREFIXES = ("/mnt/data", "/mnt/outputs", "/home/sandbox", "/workspace")
|
||||
# /tmp exists on the host; separate only to note that. The absence gate applies alike.
|
||||
_CONDITIONAL_PREFIXES = ("/tmp/outputs",)
|
||||
_notified = False
|
||||
# Invented absolute write path -> healed CWD target, so re-writing the same
|
||||
# artifact re-serves it instead of tripping the anti-clobber guard.
|
||||
_remapped_writes: dict = {}
|
||||
# Each tool call is a fresh subprocess (in-process map starts empty), so this
|
||||
# on-disk sidecar carries the map across runs. It records only sources the
|
||||
# fallback healed, so an unrelated same-basename file is never adopted.
|
||||
_REMAP_SIDECAR = ".unsloth_sandbox_remap.json"
|
||||
|
||||
|
||||
def _note(subject, original, mapped):
|
||||
"""Print the one-shot stderr notice so the model learns the real location.
|
||||
|
||||
``subject`` is what "does not exist" (the prefix, or the whole invented
|
||||
path); ``original`` is echoed in the ``(original -> mapped)`` tail.
|
||||
"""
|
||||
global _notified
|
||||
if _notified:
|
||||
return
|
||||
_notified = True
|
||||
print(
|
||||
f"note: {subject} does not exist in this sandbox; "
|
||||
f"using the working directory instead ({original} -> {mapped})",
|
||||
file = sys.stderr,
|
||||
)
|
||||
|
||||
|
||||
def _contained_join(cwd, rel):
|
||||
"""Join ``rel`` onto ``cwd`` so the result can never escape ``cwd``.
|
||||
|
||||
A habit path can carry ``..`` segments; joining verbatim would let the target
|
||||
climb above the sandbox. ``..`` components are dropped and empty / ``.`` ones
|
||||
ignored, keeping the result under ``cwd``.
|
||||
"""
|
||||
parts = []
|
||||
for part in rel.split("/"):
|
||||
if part == "" or part == ".":
|
||||
continue
|
||||
if part == "..":
|
||||
if parts:
|
||||
parts.pop()
|
||||
continue
|
||||
parts.append(part)
|
||||
return os.path.join(cwd, *parts) if parts else cwd
|
||||
|
||||
|
||||
def _map_onto_cwd(
|
||||
prefix,
|
||||
text,
|
||||
notify = True,
|
||||
):
|
||||
"""Map ``<prefix>/rest`` onto ``./rest`` in the CWD, noting it once.
|
||||
|
||||
The suffix is contained under the CWD (see ``_contained_join``) so a path
|
||||
like ``/mnt/data/../other_session/file`` cannot escape the workdir.
|
||||
``notify`` is False when the caller may keep the original path (a read), so
|
||||
the one-shot notice is not spent on a remap that never happens.
|
||||
"""
|
||||
rel = text[len(prefix) :].lstrip("/")
|
||||
mapped = _contained_join(os.getcwd(), rel)
|
||||
if notify:
|
||||
_note(prefix, text, mapped)
|
||||
return mapped
|
||||
|
||||
|
||||
def _sidecar_path(cwd):
|
||||
return os.path.join(cwd, _REMAP_SIDECAR)
|
||||
|
||||
|
||||
def _load_sidecar(cwd):
|
||||
"""Return the persisted ``source -> healed target`` map, or {} on any error
|
||||
(missing/corrupt/foreign sidecar degrades to in-process-only behaviour)."""
|
||||
try:
|
||||
with open(_sidecar_path(cwd)) as fh:
|
||||
data = json.load(fh)
|
||||
except Exception: # noqa: BLE001 - a bad sidecar must never break user code
|
||||
return {}
|
||||
return data if isinstance(data, dict) else {}
|
||||
|
||||
|
||||
def _record_sidecar(cwd, source, target):
|
||||
"""Persist ``source -> target`` so the next run re-serves it.
|
||||
|
||||
Written atomically (temp + ``os.replace``) and wrapped so a read-only/full
|
||||
filesystem never breaks the interpreter. The path is inside the CWD, so the
|
||||
patched ``open`` leaves it untouched (no remap, no recursion).
|
||||
"""
|
||||
try:
|
||||
data = _load_sidecar(cwd)
|
||||
if data.get(source) == target:
|
||||
return
|
||||
data[source] = target
|
||||
tmp = _sidecar_path(cwd) + ".tmp"
|
||||
with open(tmp, "w") as fh:
|
||||
json.dump(data, fh)
|
||||
os.replace(tmp, _sidecar_path(cwd))
|
||||
except Exception: # noqa: BLE001 - persistence is best effort only
|
||||
pass
|
||||
|
||||
|
||||
def _is_creating_mode(mode):
|
||||
"""True only when an ``open()`` mode string can CREATE a missing file.
|
||||
|
||||
Only ``w`` / ``a`` / ``x`` create. ``r+`` / ``rb+`` require the path to exist,
|
||||
so they must not trip the write fallback (which would corrupt an unrelated
|
||||
same-basename file); ``w+`` / ``a+`` / ``x+`` still match.
|
||||
"""
|
||||
return isinstance(mode, str) and any(c in mode for c in ("w", "a", "x"))
|
||||
|
||||
|
||||
def _remap_open(file, mode):
|
||||
"""Remap for ``open()`` / ``io.open()``.
|
||||
|
||||
A prefix remap runs first: a write/create heals onto the CWD; a READ heals
|
||||
only when the mapped target already exists (re-reading an earlier write),
|
||||
else the original path is kept so a genuine missing input fails truthfully
|
||||
instead of silently reading a same-basename workdir file. Only if no prefix
|
||||
matched and the call creates does the fallback kick in: an absolute target
|
||||
outside the CWD whose parent is missing is redirected to the basename in the
|
||||
CWD, unless ``CWD/<basename>`` already exists (an unrelated file), in which
|
||||
case the original path is kept so open raises.
|
||||
"""
|
||||
creating = _is_creating_mode(mode)
|
||||
# notify=False: emit the notice only once we commit to the mapping below.
|
||||
mapped = _remap(file, notify = False)
|
||||
if mapped is not file:
|
||||
# Write always heals; a read only when the mapped target exists (else keep
|
||||
# the original path so a missing input stays truthful).
|
||||
if creating or os.path.exists(mapped):
|
||||
# Commit: emit the notice now (the notify=False peek above deferred it).
|
||||
_remap(file, notify = True)
|
||||
return mapped
|
||||
return file
|
||||
if not creating:
|
||||
return file
|
||||
try:
|
||||
text = os.fspath(file)
|
||||
except TypeError:
|
||||
return file
|
||||
# bytes paths left untouched (str-only, matching the prefix remaps).
|
||||
if not isinstance(text, str) or not os.path.isabs(text):
|
||||
return file
|
||||
cwd = os.getcwd()
|
||||
# Already inside the CWD: a real target the model meant; leave it alone.
|
||||
if text == cwd or text.startswith(cwd + os.sep):
|
||||
return file
|
||||
parent = os.path.dirname(text)
|
||||
# Redirect only when the parent is missing; an existing external directory is
|
||||
# a deliberate target and stays truthful (os.path.exists follows symlinks).
|
||||
if parent and os.path.exists(parent):
|
||||
return file
|
||||
base = os.path.basename(text)
|
||||
# A trailing sep or '.'/'..' basename would redirect onto the CWD or its
|
||||
# parent; refuse and let open raise.
|
||||
if base in ("", ".", ".."):
|
||||
return file
|
||||
remapped = os.path.join(cwd, base)
|
||||
# Never clobber an unrelated file sharing this basename (lexists catches
|
||||
# dangling symlinks). But a target this fallback already healed for the same
|
||||
# invented path (in-process map or cross-run sidecar) is the artifact being
|
||||
# re-written, so re-serve it instead of raising on every overwrite.
|
||||
if os.path.lexists(remapped) and remapped not in (
|
||||
_remapped_writes.get(text),
|
||||
_load_sidecar(cwd).get(text),
|
||||
):
|
||||
return file
|
||||
_remapped_writes[text] = remapped
|
||||
_record_sidecar(cwd, text, remapped)
|
||||
_note(text, text, remapped)
|
||||
return remapped
|
||||
|
||||
|
||||
def _remap(path, notify = True):
|
||||
"""Map ``<prefix>/rest`` onto ``./rest`` in the CWD; other paths pass through.
|
||||
|
||||
``notify`` is forwarded to ``_map_onto_cwd``; ``_remap_open`` passes False so
|
||||
a read that keeps its original path emits no false notice.
|
||||
"""
|
||||
try:
|
||||
text = os.fspath(path)
|
||||
except TypeError:
|
||||
return path
|
||||
if not isinstance(text, str):
|
||||
return path
|
||||
for prefix in _PREFIXES + _CONDITIONAL_PREFIXES:
|
||||
# Heal only while the real prefix directory is absent, so a genuine host
|
||||
# mount / user directory at that prefix is never shadowed.
|
||||
if (text == prefix or text.startswith(prefix + "/")) and not os.path.exists(prefix):
|
||||
return _map_onto_cwd(prefix, text, notify = notify)
|
||||
return path
|
||||
|
||||
|
||||
def _install():
|
||||
import pathlib
|
||||
|
||||
original_open = builtins.open
|
||||
original_io_open = io.open
|
||||
original_os_open = os.open
|
||||
original_makedirs = os.makedirs
|
||||
original_mkdir = os.mkdir
|
||||
original_path_mkdir = pathlib.Path.mkdir
|
||||
|
||||
def _open(
|
||||
file,
|
||||
mode = "r",
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
return original_open(_remap_open(file, mode), mode, *args, **kwargs)
|
||||
|
||||
def _io_open(
|
||||
file,
|
||||
mode = "r",
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
return original_io_open(_remap_open(file, mode), mode, *args, **kwargs)
|
||||
|
||||
# mkdir/makedirs get only the prefix remap, never the write-mode fallback:
|
||||
# an arbitrary absolute directory can legitimately succeed on the host.
|
||||
def _makedirs(name, *args, **kwargs):
|
||||
return original_makedirs(_remap(name), *args, **kwargs)
|
||||
|
||||
def _mkdir(path, *args, **kwargs):
|
||||
return original_mkdir(_remap(path), *args, **kwargs)
|
||||
|
||||
def _os_open(
|
||||
path,
|
||||
flags,
|
||||
mode = 0o777,
|
||||
*,
|
||||
dir_fd = None,
|
||||
):
|
||||
# Path.touch() etc. go through os.open, not builtins.open. Only O_CREAT
|
||||
# can create, so only it maps to "creating" mode; O_TRUNC / O_APPEND
|
||||
# without O_CREAT still require the file to exist, so behave as a read.
|
||||
logical_mode = "w" if (flags & os.O_CREAT) else "r"
|
||||
mapped = _remap_open(path, logical_mode)
|
||||
if dir_fd is None:
|
||||
return original_os_open(mapped, flags, mode)
|
||||
return original_os_open(mapped, flags, mode, dir_fd = dir_fd)
|
||||
|
||||
def _path_mkdir(self, *args, **kwargs):
|
||||
# pathlib probes Path.is_dir()/os.stat (unpatched) on FileExistsError, so
|
||||
# a bare os.mkdir remap would still raise when the target exists. Remap
|
||||
# the receiver up front so parents/exist_ok stays idempotent.
|
||||
mapped = _remap(self)
|
||||
target = self if mapped is self else self.__class__(mapped)
|
||||
return original_path_mkdir(target, *args, **kwargs)
|
||||
|
||||
builtins.open = _open
|
||||
# pathlib.Path.open / write_text / read_text call io.open directly, so patch both.
|
||||
io.open = _io_open
|
||||
# Python < 3.11 only: pathlib's accessor captured the ORIGINAL io.open at
|
||||
# import (``_NormalAccessor.open = io.open``), so the io.open patch misses it.
|
||||
# Repoint it at the same wrapper (staticmethod to stay unbound); 3.11+ dropped
|
||||
# the accessor, so this is a no-op there.
|
||||
accessor = getattr(pathlib, "_NormalAccessor", None)
|
||||
if accessor is not None and hasattr(accessor, "open"):
|
||||
accessor.open = staticmethod(_io_open)
|
||||
# Path.touch() and other low-level opens call os.open directly, so patch it too.
|
||||
os.open = _os_open
|
||||
os.makedirs = _makedirs
|
||||
# Path.mkdir(parents=True) calls os.mkdir per component, so patch os.mkdir;
|
||||
# patch Path.mkdir itself too so exist_ok/parents land on the mapped path.
|
||||
os.mkdir = _mkdir
|
||||
pathlib.Path.mkdir = _path_mkdir
|
||||
|
||||
|
||||
try:
|
||||
_install()
|
||||
except Exception: # noqa: BLE001 - a broken shim must never break user code
|
||||
pass
|
||||
284
studio/backend/core/inference/tool_stream_exec.py
Normal file
284
studio/backend/core/inference/tool_stream_exec.py
Normal file
|
|
@ -0,0 +1,284 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Streaming wrapper around blocking server-side tool execution.
|
||||
|
||||
``stream_tool_execution`` runs a blocking tool call in a worker thread and
|
||||
turns it into a generator that yields:
|
||||
|
||||
* ``{"type": "tool_output", "tool_name", "tool_call_id", "text"}`` -- an
|
||||
incremental stdout/stderr chunk (python/terminal tools) for live UI output;
|
||||
* ``{"type": "heartbeat"}`` -- emitted whenever nothing else has been yielded
|
||||
for ``heartbeat_interval_s`` seconds, so the SSE route can write a
|
||||
keepalive and reverse proxies (Cloudflare tunnels cap idle streams at
|
||||
~100 s) never see a silent connection while a tool runs;
|
||||
|
||||
and *returns* the tool's final result string via ``StopIteration.value``
|
||||
(``result = yield from stream_tool_execution(...)``). The returned result is
|
||||
byte-identical to calling the tool directly, so tool-result parsing, nudging,
|
||||
and healing downstream are untouched.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Callable, Generator
|
||||
|
||||
from loggers import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def accepts_output_callback(func: Callable[..., str]) -> bool:
|
||||
"""Whether an injectable ``execute_tool`` supports ``output_callback``.
|
||||
|
||||
``execute_tool`` is replaceable (tests inject fakes / the pre-PR signature),
|
||||
so forward the kwarg only when the callable declares it or takes ``**kwargs``
|
||||
(passing it unconditionally would ``TypeError`` on an old signature).
|
||||
"""
|
||||
try:
|
||||
params = inspect.signature(func).parameters
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
if "output_callback" in params:
|
||||
return True
|
||||
return any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
|
||||
|
||||
|
||||
# Cadence of heartbeat events while a tool blocks with no output. Well under
|
||||
# common proxy idle caps (Cloudflare ~100 s, nginx default 60 s).
|
||||
TOOL_HEARTBEAT_INTERVAL_S = 10.0
|
||||
|
||||
# How often the wrapper wakes to poll for output / completion / cancellation.
|
||||
_POLL_INTERVAL_S = 0.25
|
||||
|
||||
# Upper bound on how long teardown waits for the worker once the stream is
|
||||
# closed or errors. A cancel-observing tool returns within this after
|
||||
# ``cancel_event`` is set; a cancel-ignoring one is a daemon left to finish on
|
||||
# its own rather than blocking teardown for the tool's full timeout.
|
||||
_WORKER_JOIN_TIMEOUT_S = 5.0
|
||||
|
||||
# Cap on total streamed live-output characters per tool call, bounding the
|
||||
# transient UI stream so a tight print loop cannot flood the SSE channel. Much
|
||||
# higher than the model-visible result cap (tools._MAX_OUTPUT_CHARS) since the
|
||||
# UI keeps the live stream as the displayed output when the result is truncated.
|
||||
TOOL_OUTPUT_STREAM_MAX_CHARS = 400_000
|
||||
|
||||
_STREAM_CAPPED_NOTICE = "\n... (further live output not streamed)\n"
|
||||
|
||||
|
||||
def _drain_queue(q: "queue.Queue", sentinel: object, max_chars: int | None) -> tuple[str, bool]:
|
||||
"""Pull every currently-queued item, joining chunks in FIFO order.
|
||||
|
||||
With ``max_chars`` set, stop concatenating at the budget and discard the
|
||||
remaining chunks in place, bounding peak allocation when a chatty tool queues
|
||||
far more than the cap before the consumer wakes. The crossing chunk is sliced
|
||||
to one char past the budget, enough for the caller's truncation to stay
|
||||
byte-identical. Returns ``(joined_text, hit_sentinel)``; the surplus is still
|
||||
scanned so completion is detected promptly.
|
||||
"""
|
||||
parts: list[str] = []
|
||||
total = 0
|
||||
dropping = False
|
||||
hit_sentinel = False
|
||||
while True:
|
||||
try:
|
||||
item = q.get_nowait()
|
||||
except queue.Empty:
|
||||
break
|
||||
if item is sentinel:
|
||||
hit_sentinel = True
|
||||
break
|
||||
if dropping:
|
||||
continue
|
||||
if max_chars is not None and total + len(item) > max_chars:
|
||||
# Keep one char past the budget as the overflow signal; drop the rest.
|
||||
parts.append(item[: max(0, max_chars - total) + 1])
|
||||
dropping = True
|
||||
continue
|
||||
parts.append(item)
|
||||
total += len(item)
|
||||
return "".join(parts), hit_sentinel
|
||||
|
||||
|
||||
def stream_tool_execution(
|
||||
invoke: Callable[[Callable[[str], None]], str],
|
||||
*,
|
||||
tool_name: str,
|
||||
tool_call_id: str = "",
|
||||
cancel_event: Any = None,
|
||||
heartbeat_interval_s: float = TOOL_HEARTBEAT_INTERVAL_S,
|
||||
poll_interval_s: float = _POLL_INTERVAL_S,
|
||||
) -> Generator[dict, None, str]:
|
||||
"""Run ``invoke(output_callback)`` in a thread; yield live events; return the result.
|
||||
|
||||
``invoke`` receives a thread-safe ``callable(str)`` it may call with
|
||||
incremental output chunks (or ignore entirely). Exceptions raised by the
|
||||
tool propagate to the caller unchanged after the worker thread finishes.
|
||||
|
||||
``cancel_event`` is the request-level cancellation signal already handed to
|
||||
the tool. If the consumer closes this generator early (an SSE disconnect
|
||||
calls ``gen.close()``, raising ``GeneratorExit`` at a ``yield``), the wrapper
|
||||
sets it so a cancel-observing tool stops, then joins the worker with a bounded
|
||||
timeout. Set ONLY on that abnormal-exit path, never on a clean finish, because
|
||||
the event is shared across a turn's tool calls and setting it early would
|
||||
abort the next tool.
|
||||
"""
|
||||
output_queue: queue.Queue[Any] = queue.Queue()
|
||||
done_sentinel = object()
|
||||
outcome: dict[str, Any] = {}
|
||||
|
||||
# Bound accepted output at the PRODUCER boundary: the consumer-side cap alone
|
||||
# wouldn't stop a fast worker enqueuing unboundedly while a slow SSE client
|
||||
# backpressures. Accept at most one char past the cap (so the consumer still
|
||||
# emits the capped notice) and drop the rest. The final result is captured
|
||||
# independently, so this never changes the byte-identical result.
|
||||
accepted_output_chars = 0
|
||||
accepted_output_lock = threading.Lock()
|
||||
|
||||
def _on_output(text: str) -> None:
|
||||
nonlocal accepted_output_chars
|
||||
if not text:
|
||||
return
|
||||
with accepted_output_lock:
|
||||
remaining = TOOL_OUTPUT_STREAM_MAX_CHARS + 1 - accepted_output_chars
|
||||
if remaining <= 0:
|
||||
return
|
||||
accepted = text[:remaining]
|
||||
accepted_output_chars += len(accepted)
|
||||
output_queue.put(accepted)
|
||||
|
||||
def _run() -> None:
|
||||
try:
|
||||
outcome["result"] = invoke(_on_output)
|
||||
except BaseException as exc: # noqa: BLE001 - re-raised on the caller side
|
||||
outcome["error"] = exc
|
||||
finally:
|
||||
# Posted after the result/error is recorded; wakes the consumer
|
||||
# immediately so fast tools pay no poll-interval latency.
|
||||
output_queue.put(done_sentinel)
|
||||
|
||||
worker = threading.Thread(
|
||||
target = _run,
|
||||
daemon = True,
|
||||
name = f"tool-exec-{tool_name or 'unknown'}",
|
||||
)
|
||||
worker.start()
|
||||
|
||||
# Heartbeats are paced by counting idle queue polls rather than a wall clock
|
||||
# (tests patch ``time.monotonic`` globally, so the wrapper must not read it).
|
||||
idle_polls_per_heartbeat = max(1, int(round(heartbeat_interval_s / poll_interval_s)))
|
||||
idle_polls = 0
|
||||
streamed_chars = 0
|
||||
stream_capped = False
|
||||
finished = False
|
||||
|
||||
def _drain_pending(max_chars: int | None = None) -> str:
|
||||
nonlocal finished
|
||||
text, hit_sentinel = _drain_queue(output_queue, done_sentinel, max_chars)
|
||||
if hit_sentinel:
|
||||
finished = True
|
||||
return text
|
||||
|
||||
def _drain_and_drop() -> None:
|
||||
"""Discard the current and every queued chunk without concatenating.
|
||||
|
||||
Past the cap every chunk is dropped, so don't pay to build a combined
|
||||
string only to drop it. Still detect completion so the loop can exit.
|
||||
"""
|
||||
nonlocal finished
|
||||
while True:
|
||||
try:
|
||||
item = output_queue.get_nowait()
|
||||
except queue.Empty:
|
||||
return
|
||||
if item is done_sentinel:
|
||||
finished = True
|
||||
return
|
||||
|
||||
abnormal_exit = False
|
||||
try:
|
||||
while not finished:
|
||||
try:
|
||||
item = output_queue.get(timeout = poll_interval_s)
|
||||
except queue.Empty:
|
||||
# A disconnect sets cancel_event while the worker is silent;
|
||||
# surface a heartbeat this poll so the route regains control and
|
||||
# tears down at once, not after a full heartbeat interval.
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
yield {"type": "heartbeat"}
|
||||
continue
|
||||
idle_polls += 1
|
||||
if idle_polls >= idle_polls_per_heartbeat:
|
||||
idle_polls = 0
|
||||
yield {"type": "heartbeat"}
|
||||
continue
|
||||
|
||||
if item is done_sentinel:
|
||||
break
|
||||
|
||||
if stream_capped:
|
||||
# Past the cap: drop this chunk and every queued sibling (see
|
||||
# _drain_and_drop). Pace with one time.sleep per poll (not
|
||||
# time.monotonic -- tests patch the clock), counted as an idle
|
||||
# poll so heartbeats keep flowing while the queue stays non-empty.
|
||||
_drain_and_drop()
|
||||
if finished:
|
||||
break
|
||||
time.sleep(poll_interval_s)
|
||||
idle_polls += 1
|
||||
if idle_polls >= idle_polls_per_heartbeat:
|
||||
idle_polls = 0
|
||||
yield {"type": "heartbeat"}
|
||||
continue
|
||||
|
||||
# Bound the join to the remaining budget so the crossing batch can't
|
||||
# allocate far past the cap (surplus is truncated below anyway); the
|
||||
# prefix is long enough that truncation stays byte-identical.
|
||||
budget = TOOL_OUTPUT_STREAM_MAX_CHARS - streamed_chars
|
||||
chunk = item + _drain_pending(max_chars = budget - len(item))
|
||||
idle_polls = 0
|
||||
if streamed_chars + len(chunk) > TOOL_OUTPUT_STREAM_MAX_CHARS:
|
||||
chunk = chunk[: max(0, TOOL_OUTPUT_STREAM_MAX_CHARS - streamed_chars)]
|
||||
chunk += _STREAM_CAPPED_NOTICE
|
||||
stream_capped = True
|
||||
streamed_chars += len(chunk)
|
||||
if chunk:
|
||||
yield {
|
||||
"type": "tool_output",
|
||||
"tool_name": tool_name,
|
||||
"tool_call_id": tool_call_id,
|
||||
"text": chunk,
|
||||
}
|
||||
except BaseException:
|
||||
# The loop only raises when the consumer closes us early: an SSE
|
||||
# disconnect calls gen.close() (GeneratorExit at the yield) or the route
|
||||
# throws in. Signal cancellation so a cancel-observing tool returns; the
|
||||
# daemon worker is then abandoned (see finally). Re-raise so the caller
|
||||
# sees the real cause (GeneratorExit must not be swallowed). Runs ONLY on
|
||||
# abnormal exit, so the shared cancel_event is never set out from under
|
||||
# the next tool in a clean multi-tool turn.
|
||||
abnormal_exit = True
|
||||
if cancel_event is not None:
|
||||
try:
|
||||
cancel_event.set()
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
finally:
|
||||
# Clean finish: the worker already recorded its result and queued the
|
||||
# sentinel we consumed, so this join returns at once. Abnormal exit:
|
||||
# cancel_event is set and the daemon worker abandoned, so join with a zero
|
||||
# timeout -- teardown never blocks the caller (the route may close this
|
||||
# generator on the event loop), and the daemon cannot outlive the process.
|
||||
worker.join(timeout = 0 if abnormal_exit else _WORKER_JOIN_TIMEOUT_S)
|
||||
|
||||
error = outcome.get("error")
|
||||
if error is not None:
|
||||
raise error
|
||||
# Returned verbatim (the loop's record_result handles non-str), so the
|
||||
# final tool result is byte-identical to a direct execute_tool call.
|
||||
return outcome.get("result")
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -1083,6 +1083,10 @@ _OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S = 0.1
|
|||
_OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S = 5.0
|
||||
_OPENAI_PASSTHROUGH_SSE_KEEPALIVE = ": keep-alive\n\n"
|
||||
_OPENAI_LLAMA_ADMISSION_POLL_S = 0.25
|
||||
# Idle window before a local tool-loop stream emits an SSE keepalive comment
|
||||
# (e.g. prompt prefill between tool iterations). A second layer atop the
|
||||
# tool_stream_exec heartbeats, keeping proxies (Cloudflare drops idle at ~100s).
|
||||
_LOCAL_TOOL_STREAM_STALL_KEEPALIVE_S = 15.0
|
||||
|
||||
|
||||
def _openai_llama_admission_capacity(request: Optional[Request], llama_backend = None) -> int:
|
||||
|
|
@ -2298,6 +2302,36 @@ async def _stop_local_disconnect_cancel_watcher(watcher) -> None:
|
|||
pass
|
||||
|
||||
|
||||
async def _drain_pending_next_task(task, cancel_event) -> None:
|
||||
"""Wait for a pending ``asyncio.to_thread(next, gen, ...)`` task to finish
|
||||
before its generator is closed.
|
||||
|
||||
On disconnect a ``next(gen)`` call may still run in a worker thread;
|
||||
cancelling the awaiting task does NOT stop it, and ``gen.close()`` mid-
|
||||
``next(gen)`` raises ``ValueError: generator already executing``, leaking the
|
||||
generator's cleanup. So re-set the cancel flag (the generator polls it) and
|
||||
shield the task until the worker returns. No-op when there is no pending task.
|
||||
"""
|
||||
if task is None:
|
||||
return
|
||||
if cancel_event is not None:
|
||||
cancel_event.set()
|
||||
while not task.done():
|
||||
try:
|
||||
await asyncio.shield(task)
|
||||
except asyncio.CancelledError:
|
||||
if cancel_event is not None:
|
||||
cancel_event.set()
|
||||
continue
|
||||
except Exception:
|
||||
break
|
||||
if task.done():
|
||||
try:
|
||||
task.exception()
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass
|
||||
|
||||
|
||||
# Centralized local/server tool nudge. Keep render_html guidance gated to turns
|
||||
# where the canvas tool is actually present in the tool schema; otherwise
|
||||
# small local models can hallucinate a missing tool call instead of following
|
||||
|
|
@ -7462,13 +7496,35 @@ async def openai_chat_completions(
|
|||
asyncio.to_thread(next, gen, _tool_sentinel)
|
||||
)
|
||||
try:
|
||||
event = await asyncio.shield(next_task)
|
||||
# Stall-timeout wait: keepalive while the generator stays
|
||||
# silent (e.g. prefill between tool iterations). asyncio.wait
|
||||
# never cancels next_task, matching the finally-drain shield.
|
||||
while True:
|
||||
done_tasks, _ = await asyncio.wait(
|
||||
{next_task},
|
||||
timeout = _LOCAL_TOOL_STREAM_STALL_KEEPALIVE_S,
|
||||
)
|
||||
if done_tasks:
|
||||
break
|
||||
yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE
|
||||
event = next_task.result()
|
||||
finally:
|
||||
if next_task.done():
|
||||
next_task = None
|
||||
if event is _tool_sentinel:
|
||||
break
|
||||
|
||||
if event["type"] == "heartbeat":
|
||||
# Tool-wrapper heartbeat while a server-side tool blocks; keeps SSE alive.
|
||||
yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE
|
||||
continue
|
||||
|
||||
if event["type"] in ("tool_output", "tool_args"):
|
||||
# Live stdout/stderr or tool-call arguments, forwarded
|
||||
# verbatim for the UI. Final result still arrives in tool_end.
|
||||
yield f"data: {json.dumps(event)}\n\n"
|
||||
continue
|
||||
|
||||
if event["type"] == "status":
|
||||
# Empty status marks an iteration boundary in the
|
||||
# GGUF tool loop (e.g. after a re-prompt). Reset the
|
||||
|
|
@ -8018,7 +8074,18 @@ async def openai_chat_completions(
|
|||
asyncio.to_thread(next, gen, _gguf_sentinel)
|
||||
)
|
||||
try:
|
||||
cumulative = await asyncio.shield(next_task)
|
||||
# Stall-timeout wait: keepalive while the generator stays
|
||||
# silent (e.g. no-tool prefill). asyncio.wait never cancels
|
||||
# next_task, matching the finally-drain shield (see GGUF stream).
|
||||
while True:
|
||||
done_tasks, _ = await asyncio.wait(
|
||||
{next_task},
|
||||
timeout = _LOCAL_TOOL_STREAM_STALL_KEEPALIVE_S,
|
||||
)
|
||||
if done_tasks:
|
||||
break
|
||||
yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE
|
||||
cumulative = next_task.result()
|
||||
finally:
|
||||
if next_task.done():
|
||||
next_task = None
|
||||
|
|
@ -8709,6 +8776,7 @@ async def openai_chat_completions(
|
|||
|
||||
async def sf_tool_stream():
|
||||
gen = None
|
||||
_sf_next_task = None
|
||||
disconnect_watcher = asyncio.create_task(
|
||||
_await_disconnect_then_cancel(request, cancel_event)
|
||||
)
|
||||
|
|
@ -8740,10 +8808,35 @@ async def openai_chat_completions(
|
|||
api_monitor.finish(monitor_id, "cancelled")
|
||||
return
|
||||
|
||||
event = await asyncio.to_thread(next, gen, _sf_tool_sentinel)
|
||||
# Stall keepalive (see GGUF tool stream): silent backend segments
|
||||
# must not leave the SSE stream idle past proxy timeouts.
|
||||
_sf_next_task = asyncio.create_task(
|
||||
asyncio.to_thread(next, gen, _sf_tool_sentinel)
|
||||
)
|
||||
while True:
|
||||
_sf_done, _ = await asyncio.wait(
|
||||
{_sf_next_task},
|
||||
timeout = _LOCAL_TOOL_STREAM_STALL_KEEPALIVE_S,
|
||||
)
|
||||
if _sf_done:
|
||||
break
|
||||
yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE
|
||||
event = _sf_next_task.result()
|
||||
# Done; drop the reference so the finally-block drain no-ops.
|
||||
_sf_next_task = None
|
||||
if event is _sf_tool_sentinel:
|
||||
break
|
||||
|
||||
if event["type"] == "heartbeat":
|
||||
# Tool-execution wrapper heartbeat -> SSE keepalive.
|
||||
yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE
|
||||
continue
|
||||
|
||||
if event["type"] in ("tool_output", "tool_args"):
|
||||
# Live stdout/stderr, or tool-call arguments as the model writes them.
|
||||
yield f"data: {json.dumps(event)}\n\n"
|
||||
continue
|
||||
|
||||
if event["type"] == "status":
|
||||
if not event["text"]:
|
||||
# Iteration boundary: flush reasoning, then a fresh prefilled extractor for the next turn.
|
||||
|
|
@ -8835,9 +8928,15 @@ async def openai_chat_completions(
|
|||
yield _openai_stream_error_sse(error_chunk)
|
||||
finally:
|
||||
await _stop_local_disconnect_cancel_watcher(disconnect_watcher)
|
||||
# Drain a still-running next(gen) worker before closing: closing
|
||||
# mid-next(gen) raises ValueError('generator already executing') and
|
||||
# skips the generator's cleanup finally. Matches the GGUF tool stream.
|
||||
await _drain_pending_next_task(_sf_next_task, cancel_event)
|
||||
if gen is not None:
|
||||
try:
|
||||
gen.close()
|
||||
# Offload the close so the generator's cleanup runs off the event
|
||||
# loop (matches the GGUF SSE path); a disconnect can't stall the loop.
|
||||
await asyncio.to_thread(gen.close)
|
||||
except (RuntimeError, ValueError):
|
||||
pass
|
||||
_sf_tracker.__exit__(None, None, None)
|
||||
|
|
@ -9027,6 +9126,8 @@ async def openai_chat_completions(
|
|||
_tracker.__enter__()
|
||||
|
||||
async def stream_chunks():
|
||||
gen = None
|
||||
_next_task = None
|
||||
disconnect_watcher = asyncio.create_task(
|
||||
_await_disconnect_then_cancel(request, cancel_event)
|
||||
)
|
||||
|
|
@ -9041,23 +9142,30 @@ async def openai_chat_completions(
|
|||
prev_text = ""
|
||||
# Split prefilled <think> into reasoning_content deltas (GGUF parity); single turn, serves MLX.
|
||||
reasoning_extractor = _new_sf_reasoning_extractor()
|
||||
# Run the sync generator in a thread pool to avoid blocking the
|
||||
# event loop. Critical for compare mode: two SSE requests arrive
|
||||
# concurrently but the orchestrator serializes them via
|
||||
# _gen_lock; without run_in_executor the second request's
|
||||
# blocking lock acquisition would freeze the entire event loop,
|
||||
# stalling both streams.
|
||||
# Run the sync generator in a worker thread so it can't block the event
|
||||
# loop. Critical for compare mode: a second request's blocking _gen_lock
|
||||
# acquisition would otherwise freeze the loop and stall both streams.
|
||||
_DONE = object() # sentinel for generator exhaustion
|
||||
loop = asyncio.get_event_loop()
|
||||
gen = generate()
|
||||
while True:
|
||||
if cancel_event.is_set():
|
||||
backend.reset_generation_state()
|
||||
break
|
||||
# next(gen, _DONE) returns _DONE instead of raising
|
||||
# StopIteration -- StopIteration can't propagate through
|
||||
# asyncio futures (Python limitation).
|
||||
cumulative = await loop.run_in_executor(None, next, gen, _DONE)
|
||||
# Stall keepalive (see safetensors tool stream) each window while
|
||||
# next(gen) runs in a worker. next(gen, _DONE) returns _DONE rather
|
||||
# than raising StopIteration (which can't cross asyncio futures).
|
||||
_next_task = asyncio.create_task(asyncio.to_thread(next, gen, _DONE))
|
||||
while True:
|
||||
_done_tasks, _ = await asyncio.wait(
|
||||
{_next_task},
|
||||
timeout = _LOCAL_TOOL_STREAM_STALL_KEEPALIVE_S,
|
||||
)
|
||||
if _done_tasks:
|
||||
break
|
||||
yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE
|
||||
cumulative = _next_task.result()
|
||||
# Done; drop the reference so the finally-block drain no-ops.
|
||||
_next_task = None
|
||||
if cumulative is _DONE:
|
||||
break
|
||||
if await request.is_disconnected():
|
||||
|
|
@ -9180,6 +9288,17 @@ async def openai_chat_completions(
|
|||
yield _openai_stream_error_sse(error_chunk)
|
||||
finally:
|
||||
await _stop_local_disconnect_cancel_watcher(disconnect_watcher)
|
||||
# Drain a still-running next(gen) worker before closing: closing
|
||||
# mid-next(gen) raises ValueError('generator already executing') and
|
||||
# skips the generator's cleanup finally. Matches the safetensors stream.
|
||||
await _drain_pending_next_task(_next_task, cancel_event)
|
||||
if gen is not None:
|
||||
try:
|
||||
# Offload the close so the generator's cleanup runs off the event
|
||||
# loop (matches the GGUF SSE path); a disconnect can't stall the loop.
|
||||
await asyncio.to_thread(gen.close)
|
||||
except (RuntimeError, ValueError):
|
||||
pass
|
||||
_tracker.__exit__(None, None, None)
|
||||
|
||||
return _SameTaskStreamingResponse(
|
||||
|
|
@ -12389,8 +12508,12 @@ async def _anthropic_tool_stream(
|
|||
ends_on_tool_use = False
|
||||
tool_blocks_emitted = 0
|
||||
drop_until_tool_end = False
|
||||
# Last drop-branch keepalive, seeded to stream start so a chatty tool busy
|
||||
# past the stall window still gets a keepalive though its events are dropped.
|
||||
_last_drop_keepalive = time.monotonic()
|
||||
|
||||
gen = run_gen()
|
||||
_next_task = None
|
||||
# Watcher to cancel on disconnect: the in-loop poll fires only between
|
||||
# events, so a mid-prefill disconnect would otherwise hold the decode slot.
|
||||
disconnect_watcher = asyncio.create_task(
|
||||
|
|
@ -12401,13 +12524,42 @@ async def _anthropic_tool_stream(
|
|||
if cancel_event.is_set() or await request.is_disconnected():
|
||||
cancel_event.set()
|
||||
return
|
||||
event = await asyncio.to_thread(next, gen, _sentinel)
|
||||
# Stall keepalive (see GGUF tool stream): silent backend segments
|
||||
# must not leave the SSE stream idle past proxy timeouts.
|
||||
_next_task = asyncio.create_task(asyncio.to_thread(next, gen, _sentinel))
|
||||
while True:
|
||||
_done_tasks, _ = await asyncio.wait(
|
||||
{_next_task},
|
||||
timeout = _LOCAL_TOOL_STREAM_STALL_KEEPALIVE_S,
|
||||
)
|
||||
if _done_tasks:
|
||||
break
|
||||
yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE
|
||||
event = _next_task.result()
|
||||
# Done; drop the reference so the finally-block drain no-ops.
|
||||
_next_task = None
|
||||
if event is _sentinel:
|
||||
break
|
||||
etype = event.get("type")
|
||||
if etype == "heartbeat":
|
||||
# Tool-wrapper heartbeat -> SSE keepalive, checked BEFORE the drop
|
||||
# skip: a dropped tool still runs server-side and its events keep the
|
||||
# stall keepalive from firing, so dropping heartbeats would go silent.
|
||||
yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE
|
||||
continue
|
||||
if etype in ("tool_output", "tool_args"):
|
||||
# Live stdout / arg streaming have no Anthropic Messages equivalent
|
||||
# (the full call/result follow in tool_use / tool_result), so drop them.
|
||||
# They keep the stall keepalive from firing, so a chatty tool would go
|
||||
# silent past the ~100s proxy cap; emit a rate-limited keepalive instead.
|
||||
_now = time.monotonic()
|
||||
if _now - _last_drop_keepalive >= _LOCAL_TOOL_STREAM_STALL_KEEPALIVE_S:
|
||||
_last_drop_keepalive = _now
|
||||
yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE
|
||||
continue
|
||||
if drop_until_tool_end:
|
||||
# disable_parallel_tool_use: a later tool call is being
|
||||
# dropped — skip every event until (and including) its tool_end.
|
||||
# disable_parallel_tool_use: skip every event until (and
|
||||
# including) this dropped tool call's tool_end.
|
||||
if etype == "tool_end":
|
||||
drop_until_tool_end = False
|
||||
continue
|
||||
|
|
@ -12448,12 +12600,24 @@ async def _anthropic_tool_stream(
|
|||
yield line
|
||||
except Exception as e:
|
||||
logger.error("anthropic_messages stream error: %s", e)
|
||||
_error_event = _anthropic_stream_error_event(e)
|
||||
# force = True so an unclassified mid-stream failure (llama-server crash,
|
||||
# decode OOM, dropped socket) still emits an SSE error and returns, instead
|
||||
# of a normal message_stop that masks a truncated turn as a clean finish.
|
||||
_error_event = _anthropic_stream_error_event(e, force = True)
|
||||
if _error_event is not None:
|
||||
yield _error_event
|
||||
return
|
||||
finally:
|
||||
await _stop_local_disconnect_cancel_watcher(disconnect_watcher)
|
||||
# Drain a still-running next(gen) worker before closing, so a mid-prefill
|
||||
# disconnect releases the thread/generator/tool resources. Closing first
|
||||
# would race into ValueError('generator already executing').
|
||||
await _drain_pending_next_task(_next_task, cancel_event)
|
||||
if gen is not None:
|
||||
try:
|
||||
await asyncio.to_thread(gen.close)
|
||||
except (RuntimeError, ValueError):
|
||||
pass
|
||||
|
||||
stop_reason = openai_finish_to_anthropic_stop(
|
||||
captured_finish_reason, had_tool_calls = ends_on_tool_use
|
||||
|
|
@ -12490,6 +12654,7 @@ async def _anthropic_plain_stream(
|
|||
captured_finish_reason = None
|
||||
|
||||
gen = run_gen()
|
||||
_next_task = None
|
||||
# Watcher to cancel on disconnect: the in-loop poll fires only between
|
||||
# chunks, so a mid-prefill disconnect would otherwise hold the decode slot.
|
||||
disconnect_watcher = asyncio.create_task(
|
||||
|
|
@ -12500,7 +12665,20 @@ async def _anthropic_plain_stream(
|
|||
if cancel_event.is_set() or await request.is_disconnected():
|
||||
cancel_event.set()
|
||||
return
|
||||
cumulative = await asyncio.to_thread(next, gen, _sentinel)
|
||||
# Stall keepalive (see Anthropic tool stream) each window while
|
||||
# next(gen) runs in a worker.
|
||||
_next_task = asyncio.create_task(asyncio.to_thread(next, gen, _sentinel))
|
||||
while True:
|
||||
_done_tasks, _ = await asyncio.wait(
|
||||
{_next_task},
|
||||
timeout = _LOCAL_TOOL_STREAM_STALL_KEEPALIVE_S,
|
||||
)
|
||||
if _done_tasks:
|
||||
break
|
||||
yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE
|
||||
cumulative = _next_task.result()
|
||||
# Done; drop the reference so the finally-block drain no-ops.
|
||||
_next_task = None
|
||||
if cumulative is _sentinel:
|
||||
break
|
||||
if isinstance(cumulative, dict):
|
||||
|
|
@ -12516,12 +12694,24 @@ async def _anthropic_plain_stream(
|
|||
yield line
|
||||
except Exception as e:
|
||||
logger.error("anthropic_messages stream error: %s", e)
|
||||
_error_event = _anthropic_stream_error_event(e)
|
||||
# force = True so an unclassified mid-stream failure (llama-server crash,
|
||||
# decode OOM, dropped socket) still emits an SSE error and returns, instead
|
||||
# of a normal message_stop that masks a truncated turn as a clean finish.
|
||||
_error_event = _anthropic_stream_error_event(e, force = True)
|
||||
if _error_event is not None:
|
||||
yield _error_event
|
||||
return
|
||||
finally:
|
||||
await _stop_local_disconnect_cancel_watcher(disconnect_watcher)
|
||||
# Drain a still-running next(gen) worker before closing, so a mid-prefill
|
||||
# disconnect releases the thread/generator/model resources. Closing first
|
||||
# would race into ValueError('generator already executing').
|
||||
await _drain_pending_next_task(_next_task, cancel_event)
|
||||
if gen is not None:
|
||||
try:
|
||||
await asyncio.to_thread(gen.close)
|
||||
except (RuntimeError, ValueError):
|
||||
pass
|
||||
|
||||
stop_reason = openai_finish_to_anthropic_stop(captured_finish_reason, had_tool_calls = False)
|
||||
for line in emitter.finish(stop_reason = stop_reason, stop_sequence = None):
|
||||
|
|
|
|||
|
|
@ -1631,6 +1631,48 @@ class TestAnthropicMessagesToolRouting:
|
|||
assert entry["status"] == "cancelled"
|
||||
assert monitor.active_count() == 0
|
||||
|
||||
@staticmethod
|
||||
def _sse_blob(chunks):
|
||||
# StreamingResponse may hand back str or already-encoded bytes.
|
||||
return "".join(c.decode() if isinstance(c, (bytes, bytearray)) else c for c in chunks)
|
||||
|
||||
def test_plain_streaming_unclassified_error_emits_error_event(self, monkeypatch):
|
||||
# An unclassified mid-stream failure must surface as an SSE `error` event
|
||||
# and stop, not a message_stop that masks a truncated turn as clean.
|
||||
def _gen_boom(**_kwargs):
|
||||
yield "partial"
|
||||
raise RuntimeError("llama-server crashed mid-decode")
|
||||
|
||||
_mock_backend(monkeypatch, generate_chat_completion = _gen_boom)
|
||||
payload = _basic_payload(stream = True)
|
||||
|
||||
response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t"))
|
||||
blob = self._sse_blob(self._consume_response(response))
|
||||
|
||||
assert "event: error" in blob
|
||||
assert '"type": "error"' in blob
|
||||
assert "event: message_stop" not in blob
|
||||
|
||||
def test_tool_streaming_unclassified_error_emits_error_event(self, monkeypatch):
|
||||
# Same guarantee on the tool-calling stream path.
|
||||
def _gen_tools_boom(**_kwargs):
|
||||
yield {"type": "content", "text": "partial"}
|
||||
raise RuntimeError("llama-server crashed mid-decode")
|
||||
|
||||
_mock_backend(monkeypatch, generate_chat_completion_with_tools = _gen_tools_boom)
|
||||
payload = _basic_payload(
|
||||
stream = True,
|
||||
enable_tools = True,
|
||||
tools = [{"type": "web_search_20250305", "name": "web_search"}],
|
||||
)
|
||||
|
||||
response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t"))
|
||||
blob = self._sse_blob(self._consume_response(response))
|
||||
|
||||
assert "event: error" in blob
|
||||
assert '"type": "error"' in blob
|
||||
assert "event: message_stop" not in blob
|
||||
|
||||
def test_mixed_server_and_client_tools_rejected_with_400(self, monkeypatch):
|
||||
_mock_backend(monkeypatch)
|
||||
payload = _basic_payload(
|
||||
|
|
@ -2076,3 +2118,339 @@ def test_resumed_null_assistant_between_users_coalesced_on_messages_route(monkey
|
|||
if isinstance(merged, list):
|
||||
merged = " ".join(p.get("text", "") for p in merged if isinstance(p, dict))
|
||||
assert "first question" in merged and "please continue" in merged
|
||||
|
||||
|
||||
def test_disable_parallel_tool_use_forwards_heartbeats_while_dropping():
|
||||
"""Heartbeats from a parallel-disabled, dropped tool call must still reach
|
||||
the client as SSE keepalives: the dropped call runs server-side and the
|
||||
stall keepalive never fires while the generator keeps producing events, so
|
||||
swallowing them recreates the silent window keepalives exist to prevent."""
|
||||
import threading as _threading
|
||||
|
||||
from routes.inference import (
|
||||
_OPENAI_PASSTHROUGH_SSE_KEEPALIVE,
|
||||
_anthropic_tool_stream,
|
||||
)
|
||||
|
||||
def run_gen():
|
||||
def gen():
|
||||
yield {
|
||||
"type": "tool_start",
|
||||
"tool_name": "python",
|
||||
"tool_call_id": "call_0",
|
||||
"arguments": {},
|
||||
}
|
||||
yield {"type": "heartbeat"}
|
||||
yield {
|
||||
"type": "tool_end",
|
||||
"tool_name": "python",
|
||||
"tool_call_id": "call_0",
|
||||
"result": "r1",
|
||||
}
|
||||
# Second call: dropped by disable_parallel_tool_use, still executed
|
||||
# server-side (heartbeats + live output).
|
||||
yield {
|
||||
"type": "tool_start",
|
||||
"tool_name": "python",
|
||||
"tool_call_id": "call_1",
|
||||
"arguments": {},
|
||||
}
|
||||
yield {"type": "heartbeat"}
|
||||
yield {
|
||||
"type": "tool_output",
|
||||
"tool_name": "python",
|
||||
"tool_call_id": "call_1",
|
||||
"text": "x",
|
||||
}
|
||||
yield {"type": "heartbeat"}
|
||||
yield {
|
||||
"type": "tool_end",
|
||||
"tool_name": "python",
|
||||
"tool_call_id": "call_1",
|
||||
"result": "r2",
|
||||
}
|
||||
yield {"type": "content", "text": "final answer"}
|
||||
|
||||
return gen()
|
||||
|
||||
async def _drive():
|
||||
async def _is_disconnected():
|
||||
return False
|
||||
|
||||
request = SimpleNamespace(is_disconnected = _is_disconnected)
|
||||
resp = await _anthropic_tool_stream(
|
||||
request,
|
||||
_threading.Event(),
|
||||
run_gen,
|
||||
"msg_hb",
|
||||
"m",
|
||||
disable_parallel_tool_use = True,
|
||||
)
|
||||
return [chunk async for chunk in resp.body_iterator]
|
||||
|
||||
chunks = asyncio.run(_drive())
|
||||
keepalives = [c for c in chunks if c == _OPENAI_PASSTHROUGH_SSE_KEEPALIVE]
|
||||
# One heartbeat inside the kept call, two inside the dropped window.
|
||||
assert len(keepalives) >= 3
|
||||
# The dropped call must not surface as a second tool_use block.
|
||||
tool_use_starts = [c for c in chunks if "content_block_start" in c and '"tool_use"' in c]
|
||||
assert len(tool_use_starts) == 1
|
||||
|
||||
|
||||
def test_dropped_tool_output_events_emit_rate_limited_keepalives(monkeypatch):
|
||||
"""A chatty tool streaming tool_output/tool_args with no heartbeats keeps the
|
||||
generator busy (stall keepalive never fires); the Anthropic path can't
|
||||
translate those events and drops them. Dropping silently would let an idle
|
||||
proxy kill the stream, so the drop branch emits a rate-limited keepalive."""
|
||||
import threading as _threading
|
||||
|
||||
import routes.inference as inf_mod
|
||||
from routes.inference import (
|
||||
_OPENAI_PASSTHROUGH_SSE_KEEPALIVE,
|
||||
_anthropic_tool_stream,
|
||||
)
|
||||
|
||||
# Deterministic clock: only the drop-branch keepalive uses time.monotonic
|
||||
# here, so jumping past the stall window per call makes each dropped event
|
||||
# cross the rate-limit threshold. asyncio.wait uses the loop clock and
|
||||
# next(gen) returns promptly, so the outer stall keepalive never fires --
|
||||
# every keepalive here is from the drop branch.
|
||||
_real_time = inf_mod.time
|
||||
_tick = {"v": 0.0}
|
||||
|
||||
def _fast_monotonic():
|
||||
_tick["v"] += 100.0
|
||||
return _tick["v"]
|
||||
|
||||
fake_time = SimpleNamespace(
|
||||
monotonic = _fast_monotonic,
|
||||
sleep = _real_time.sleep,
|
||||
time = _real_time.time,
|
||||
perf_counter = _real_time.perf_counter,
|
||||
)
|
||||
monkeypatch.setattr(inf_mod, "time", fake_time)
|
||||
|
||||
n_output = 4
|
||||
|
||||
def run_gen():
|
||||
def gen():
|
||||
yield {
|
||||
"type": "tool_start",
|
||||
"tool_name": "python",
|
||||
"tool_call_id": "call_0",
|
||||
"arguments": {},
|
||||
}
|
||||
# Chatty streamed stdout, no heartbeats.
|
||||
for i in range(n_output):
|
||||
yield {
|
||||
"type": "tool_output",
|
||||
"tool_name": "python",
|
||||
"tool_call_id": "call_0",
|
||||
"text": f"line {i}\n",
|
||||
}
|
||||
yield {
|
||||
"type": "tool_end",
|
||||
"tool_name": "python",
|
||||
"tool_call_id": "call_0",
|
||||
"result": "done",
|
||||
}
|
||||
yield {"type": "content", "text": "final answer"}
|
||||
|
||||
return gen()
|
||||
|
||||
async def _drive():
|
||||
async def _is_disconnected():
|
||||
return False
|
||||
|
||||
request = SimpleNamespace(is_disconnected = _is_disconnected)
|
||||
resp = await _anthropic_tool_stream(
|
||||
request,
|
||||
_threading.Event(),
|
||||
run_gen,
|
||||
"msg_drop_ka",
|
||||
"m",
|
||||
)
|
||||
return [chunk async for chunk in resp.body_iterator]
|
||||
|
||||
chunks = asyncio.run(_drive())
|
||||
keepalives = [c for c in chunks if c == _OPENAI_PASSTHROUGH_SSE_KEEPALIVE]
|
||||
assert len(keepalives) == n_output
|
||||
# Final answer still reaches the client (drop is transport-only).
|
||||
assert any("final answer" in c for c in chunks)
|
||||
|
||||
|
||||
def test_parallel_disabled_dropped_call_output_emits_rate_limited_keepalives(monkeypatch):
|
||||
"""Under disable_parallel_tool_use a chatty second call is dropped whole
|
||||
(drop_until_tool_end). Its tool_output/tool_args events must still emit
|
||||
rate-limited keepalives: the drop window can last minutes with no heartbeats
|
||||
and no stall keepalive, so swallowing them silently would let an idle proxy
|
||||
kill the stream. The keepalive branch runs before the drop skip."""
|
||||
import threading as _threading
|
||||
|
||||
import routes.inference as inf_mod
|
||||
from routes.inference import (
|
||||
_OPENAI_PASSTHROUGH_SSE_KEEPALIVE,
|
||||
_anthropic_tool_stream,
|
||||
)
|
||||
|
||||
# Deterministic clock: jumps past the stall window per call (see sibling test).
|
||||
_real_time = inf_mod.time
|
||||
_tick = {"v": 0.0}
|
||||
|
||||
def _fast_monotonic():
|
||||
_tick["v"] += 100.0
|
||||
return _tick["v"]
|
||||
|
||||
fake_time = SimpleNamespace(
|
||||
monotonic = _fast_monotonic,
|
||||
sleep = _real_time.sleep,
|
||||
time = _real_time.time,
|
||||
perf_counter = _real_time.perf_counter,
|
||||
)
|
||||
monkeypatch.setattr(inf_mod, "time", fake_time)
|
||||
|
||||
n_output = 4
|
||||
|
||||
def run_gen():
|
||||
def gen():
|
||||
# First (kept) call.
|
||||
yield {
|
||||
"type": "tool_start",
|
||||
"tool_name": "python",
|
||||
"tool_call_id": "call_0",
|
||||
"arguments": {},
|
||||
}
|
||||
yield {
|
||||
"type": "tool_end",
|
||||
"tool_name": "python",
|
||||
"tool_call_id": "call_0",
|
||||
"result": "r1",
|
||||
}
|
||||
# Second call: dropped whole by disable_parallel_tool_use but still
|
||||
# executed server-side, streaming chatty stdout with no heartbeats.
|
||||
yield {
|
||||
"type": "tool_start",
|
||||
"tool_name": "python",
|
||||
"tool_call_id": "call_1",
|
||||
"arguments": {},
|
||||
}
|
||||
for i in range(n_output):
|
||||
yield {
|
||||
"type": "tool_output",
|
||||
"tool_name": "python",
|
||||
"tool_call_id": "call_1",
|
||||
"text": f"line {i}\n",
|
||||
}
|
||||
yield {
|
||||
"type": "tool_end",
|
||||
"tool_name": "python",
|
||||
"tool_call_id": "call_1",
|
||||
"result": "r2",
|
||||
}
|
||||
yield {"type": "content", "text": "final answer"}
|
||||
|
||||
return gen()
|
||||
|
||||
async def _drive():
|
||||
async def _is_disconnected():
|
||||
return False
|
||||
|
||||
request = SimpleNamespace(is_disconnected = _is_disconnected)
|
||||
resp = await _anthropic_tool_stream(
|
||||
request,
|
||||
_threading.Event(),
|
||||
run_gen,
|
||||
"msg_drop_ka2",
|
||||
"m",
|
||||
disable_parallel_tool_use = True,
|
||||
)
|
||||
return [chunk async for chunk in resp.body_iterator]
|
||||
|
||||
chunks = asyncio.run(_drive())
|
||||
keepalives = [c for c in chunks if c == _OPENAI_PASSTHROUGH_SSE_KEEPALIVE]
|
||||
assert len(keepalives) == n_output
|
||||
# The dropped call must not surface as a second tool_use block.
|
||||
tool_use_starts = [c for c in chunks if "content_block_start" in c and '"tool_use"' in c]
|
||||
assert len(tool_use_starts) == 1
|
||||
assert any("final answer" in c for c in chunks)
|
||||
|
||||
|
||||
def test_plain_stream_emits_keepalive_during_prompt_stall(monkeypatch):
|
||||
"""No-tool Anthropic stream must emit SSE keepalives while a long prompt
|
||||
prefill blocks next(gen), matching the tool stream (finding 5). The old
|
||||
single unbounded to_thread(next, ...) could sit silent past a proxy idle cap."""
|
||||
import threading as _threading
|
||||
import time as _time
|
||||
|
||||
from routes import inference as inf_mod
|
||||
from routes.inference import _OPENAI_PASSTHROUGH_SSE_KEEPALIVE, _anthropic_plain_stream
|
||||
|
||||
monkeypatch.setattr(inf_mod, "_LOCAL_TOOL_STREAM_STALL_KEEPALIVE_S", 0.05)
|
||||
|
||||
def run_gen():
|
||||
def gen():
|
||||
_time.sleep(0.24) # stall past several shortened keepalive windows
|
||||
yield "hello world"
|
||||
|
||||
return gen()
|
||||
|
||||
async def _drive():
|
||||
async def _is_disconnected():
|
||||
return False
|
||||
|
||||
request = SimpleNamespace(is_disconnected = _is_disconnected)
|
||||
resp = await _anthropic_plain_stream(
|
||||
request, _threading.Event(), run_gen, "msg_plain_ka", "m"
|
||||
)
|
||||
return [chunk async for chunk in resp.body_iterator]
|
||||
|
||||
chunks = asyncio.run(_drive())
|
||||
keepalives = [c for c in chunks if c == _OPENAI_PASSTHROUGH_SSE_KEEPALIVE]
|
||||
assert len(keepalives) >= 2
|
||||
assert any("hello world" in c for c in chunks)
|
||||
|
||||
|
||||
def test_plain_stream_closes_generator_on_disconnect():
|
||||
"""On disconnect the no-tool teardown must drain any pending worker and close
|
||||
the generator (finding 6). The old finally only stopped the disconnect
|
||||
watcher, leaking the generator. A fake generator records close() so the
|
||||
teardown is asserted deterministically, not via GC."""
|
||||
import threading as _threading
|
||||
|
||||
from routes.inference import _anthropic_plain_stream
|
||||
|
||||
closed = _threading.Event()
|
||||
|
||||
class _FakeGen:
|
||||
def __init__(self):
|
||||
self._items = iter(["tok0", "tok1", "tok2", "tok3"])
|
||||
|
||||
def __next__(self):
|
||||
return next(self._items)
|
||||
|
||||
def close(self):
|
||||
closed.set()
|
||||
|
||||
def run_gen():
|
||||
return _FakeGen()
|
||||
|
||||
state = {"disconnected": False}
|
||||
|
||||
async def _drive():
|
||||
async def _is_disconnected():
|
||||
return state["disconnected"]
|
||||
|
||||
request = SimpleNamespace(is_disconnected = _is_disconnected)
|
||||
resp = await _anthropic_plain_stream(
|
||||
request, _threading.Event(), run_gen, "msg_plain_close", "m"
|
||||
)
|
||||
out = []
|
||||
async for chunk in resp.body_iterator:
|
||||
out.append(chunk)
|
||||
if "tok0" in chunk:
|
||||
# Client drops after the first token; the next loop turn tears down.
|
||||
state["disconnected"] = True
|
||||
return out
|
||||
|
||||
asyncio.run(_drive())
|
||||
assert closed.is_set()
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ never gating under bypass.
|
|||
Run with: ``PYTHONPATH=studio/backend python -m pytest studio/backend/tests/test_bypass_permissions.py -q``
|
||||
"""
|
||||
|
||||
import io
|
||||
import os
|
||||
import sys
|
||||
|
||||
|
|
@ -94,10 +95,23 @@ def test_safe_env_excludes_host_and_secret(monkeypatch, tmp_path):
|
|||
|
||||
|
||||
class _FakeProc:
|
||||
returncode = 0
|
||||
"""A subprocess.Popen double for the drain path (``tools._drain_process_output``):
|
||||
a readable ``stdout`` pipe yielding the fake output then EOF, plus
|
||||
``wait()`` / ``poll()`` / ``pid``. The pid is non-existent so
|
||||
``_capture_process_group``'s ``os.getpgid`` returns None; ``wait`` returns
|
||||
immediately so the drain never kills.
|
||||
"""
|
||||
|
||||
def communicate(self, timeout = None):
|
||||
return ("FAKEOUT", None)
|
||||
returncode = 0
|
||||
# Unlikely-to-exist pid: os.getpgid raises ProcessLookupError (caught) -> None.
|
||||
pid = 2**22
|
||||
|
||||
def __init__(self):
|
||||
# Readable stdout: iter(readline, "") yields "FAKEOUT" then hits EOF.
|
||||
self.stdout = io.StringIO("FAKEOUT")
|
||||
|
||||
def wait(self, timeout = None):
|
||||
return 0
|
||||
|
||||
def poll(self):
|
||||
return 0
|
||||
|
|
|
|||
|
|
@ -1498,6 +1498,59 @@ def test_textual_mistral_marker_not_leaked_when_inline_with_preface(monkeypatch)
|
|||
assert any("Let me search." in t for t in content_texts)
|
||||
|
||||
|
||||
def test_textual_explicit_id_reuses_provisional_card(monkeypatch):
|
||||
# A textual Mistral-style call with an explicit ``id`` must reconcile onto the
|
||||
# open provisional TEXT card (keyed "call_0"), not spawn a duplicate under the
|
||||
# explicit id (which the parser keeps for execution).
|
||||
big_query = "cats " * 80 # push the drained call past the provisional floor
|
||||
call = "[TOOL_CALLS]" + json.dumps(
|
||||
[{"name": "web_search", "arguments": {"query": big_query}, "id": "explicit-42"}]
|
||||
)
|
||||
assert len(call) > 256
|
||||
# Small chunks so the provisional card opens mid-generation (a single-shot
|
||||
# delta parses instantly and never shows a provisional to exercise).
|
||||
chunks = [call[i : i + 24] for i in range(0, len(call), 24)]
|
||||
streams = [
|
||||
[_sse({"content": c}) for c in chunks] + [_done()],
|
||||
[_sse({"content": "done"}), _done()],
|
||||
]
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(monkeypatch, streams, payloads)
|
||||
calls: list[tuple[str, dict]] = []
|
||||
|
||||
def fake_execute_tool(name, arguments, **_kwargs):
|
||||
calls.append((name, arguments))
|
||||
return "result"
|
||||
|
||||
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
|
||||
|
||||
events = list(
|
||||
backend.generate_chat_completion_with_tools(
|
||||
messages = [{"role": "user", "content": "search"}],
|
||||
tools = [{"type": "function", "function": {"name": "web_search"}}],
|
||||
max_tool_iterations = 1,
|
||||
)
|
||||
)
|
||||
|
||||
assert calls == [("web_search", {"query": big_query})]
|
||||
tool_starts = [e for e in events if e.get("type") == "tool_start"]
|
||||
# Empty-args card = provisional open; full-args card = reconciled real start.
|
||||
provisional = [e for e in tool_starts if not e.get("arguments")]
|
||||
real = [e for e in tool_starts if e.get("arguments", {}).get("query")]
|
||||
assert len(provisional) == 1, tool_starts # provisional actually opened
|
||||
prov_id = provisional[0]["tool_call_id"]
|
||||
# Exactly one real card, sharing the provisional id, not a duplicate under
|
||||
# the explicit "explicit-42" id.
|
||||
assert len(real) == 1, tool_starts
|
||||
assert real[0]["tool_call_id"] == prov_id
|
||||
assert real[0]["tool_name"] == "web_search"
|
||||
assert {e["tool_call_id"] for e in tool_starts} == {prov_id}
|
||||
# A single tool_end reconciles the card; no stale empty-result close.
|
||||
ends = [e for e in events if e.get("type") == "tool_end"]
|
||||
assert [e["tool_call_id"] for e in ends] == [prov_id]
|
||||
assert ends[0]["result"] == "result"
|
||||
|
||||
|
||||
def test_textual_llama_python_tag_marker_not_leaked(monkeypatch):
|
||||
# Same leak class for the Llama-3 built-in ``<|python_tag|>NAME.call(...)`` form.
|
||||
streams = [
|
||||
|
|
@ -2916,3 +2969,204 @@ def test_gguf_valid_tool_calls_respect_max_tool_iterations(monkeypatch):
|
|||
m.get("role") == "user" and "used all available tool calls" in m.get("content", "")
|
||||
for m in payloads[2]["messages"]
|
||||
), payloads[2]["messages"]
|
||||
|
||||
|
||||
# ── Live tool-call argument streaming (tool_args events) ─────────────────────
|
||||
|
||||
|
||||
def _python_tool_schema() -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "python",
|
||||
"description": "Run python code.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"code": {"type": "string"}},
|
||||
"required": ["code"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_structured_tool_args_stream_to_provisional_card(monkeypatch):
|
||||
"""A large structured tool call must stream its arguments as tool_args events
|
||||
to the provisional card (backlog that triggered the card, then each
|
||||
fragment), while the executed call and the model's view stay exactly what the
|
||||
accumulator built."""
|
||||
|
||||
code = "print('x')\n" + ("# pad\n" * 80)
|
||||
args_json = json.dumps({"code": code})
|
||||
call_id = "call_live_args"
|
||||
split = _PROVISIONAL_ARGS_MIN_CHARS + 16
|
||||
frag1, frag2, frag3 = (
|
||||
args_json[:split],
|
||||
args_json[split : split + 40],
|
||||
args_json[split + 40 :],
|
||||
)
|
||||
|
||||
def _tc_delta(fragment: str, with_header: bool) -> str:
|
||||
entry: dict = {"index": 0, "function": {"arguments": fragment}}
|
||||
if with_header:
|
||||
entry.update({"id": call_id, "type": "function"})
|
||||
entry["function"]["name"] = "python"
|
||||
return _sse({"tool_calls": [entry]})
|
||||
|
||||
first_stream = [
|
||||
_tc_delta(frag1, with_header = True),
|
||||
_tc_delta(frag2, with_header = False),
|
||||
_tc_delta(frag3, with_header = False),
|
||||
_done(),
|
||||
]
|
||||
second_stream = [_sse({"content": "Done."}), _done()]
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(monkeypatch, [first_stream, second_stream], payloads)
|
||||
|
||||
executed: list[tuple[str, dict]] = []
|
||||
|
||||
def fake_execute_tool(name, arguments, **_kwargs):
|
||||
executed.append((name, arguments))
|
||||
return "ok"
|
||||
|
||||
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
|
||||
|
||||
events = list(
|
||||
backend.generate_chat_completion_with_tools(
|
||||
messages = [{"role": "user", "content": "run it"}],
|
||||
tools = _python_tool_schema(),
|
||||
max_tool_iterations = 1,
|
||||
)
|
||||
)
|
||||
|
||||
starts = [e for e in events if e.get("type") == "tool_start"]
|
||||
assert starts and starts[0]["tool_call_id"] == call_id
|
||||
|
||||
args_events = [e for e in events if e.get("type") == "tool_args"]
|
||||
assert args_events, "no tool_args events were streamed"
|
||||
assert all(e["tool_call_id"] == call_id for e in args_events)
|
||||
# First event is the backlog, the rest raw fragments; together the args JSON.
|
||||
assert args_events[0]["text"] == frag1
|
||||
assert "".join(e["text"] for e in args_events) == args_json
|
||||
|
||||
# The streamed display path must not perturb execution or the model view.
|
||||
assert executed == [("python", {"code": code})]
|
||||
assistant_messages = [m for m in payloads[1]["messages"] if m.get("role") == "assistant"]
|
||||
tc = assistant_messages[-1]["tool_calls"][0]
|
||||
assert tc["id"] == call_id
|
||||
# Controller re-serializes args (normalized JSON); parsed payload unchanged.
|
||||
assert json.loads(tc["function"]["arguments"]) == {"code": code}
|
||||
|
||||
|
||||
def test_text_tool_call_streams_args_and_reconciles_card(monkeypatch):
|
||||
"""A TEXT (XML) tool call must stream its raw call text as tool_args under the
|
||||
id the stream-end parser assigns ("call_0"), so the provisional card and the
|
||||
final tool_start reconcile."""
|
||||
|
||||
code = "print('hello')\n" + ("# filler\n" * 60)
|
||||
call_json = json.dumps({"name": "python", "arguments": {"code": code}})
|
||||
call_text = f"<tool_call>{call_json}</tool_call>"
|
||||
chunks = [call_text[i : i + 48] for i in range(0, len(call_text), 48)]
|
||||
first_stream = [_sse({"content": chunk}) for chunk in chunks] + [_done()]
|
||||
second_stream = [_sse({"content": "Done."}), _done()]
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(monkeypatch, [first_stream, second_stream], payloads)
|
||||
|
||||
executed: list[tuple[str, dict]] = []
|
||||
|
||||
def fake_execute_tool(name, arguments, **_kwargs):
|
||||
executed.append((name, arguments))
|
||||
return "ok"
|
||||
|
||||
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
|
||||
|
||||
events = list(
|
||||
backend.generate_chat_completion_with_tools(
|
||||
messages = [{"role": "user", "content": "run it"}],
|
||||
tools = _python_tool_schema(),
|
||||
max_tool_iterations = 1,
|
||||
)
|
||||
)
|
||||
|
||||
starts = [e for e in events if e.get("type") == "tool_start"]
|
||||
assert starts, "no tool_start emitted"
|
||||
# Provisional card first (parser's first-call id), then the reconciling start.
|
||||
assert starts[0]["tool_call_id"] == "call_0"
|
||||
assert starts[0]["arguments"] == {}
|
||||
assert starts[-1]["tool_call_id"] == "call_0"
|
||||
|
||||
args_events = [e for e in events if e.get("type") == "tool_args"]
|
||||
assert args_events, "no tool_args events for the text call"
|
||||
assert all(e["tool_call_id"] == "call_0" for e in args_events)
|
||||
streamed = "".join(e["text"] for e in args_events)
|
||||
# Streamed text is the drained call (display only); it must never leak into
|
||||
# content events.
|
||||
assert '"name": "python"' in streamed
|
||||
assert executed == [("python", {"code": code})]
|
||||
content_events = [e for e in events if e.get("type") == "content"]
|
||||
assert not any("<tool_call>" in e["text"] for e in content_events)
|
||||
|
||||
|
||||
def test_ordinary_json_answer_streams_no_tool_args(monkeypatch):
|
||||
"""A large ordinary JSON answer (no enabled tool name) must not spawn a
|
||||
provisional card or tool_args events; it stays a normal content answer."""
|
||||
|
||||
answer = json.dumps({"result": "fine", "data": ["x" * 40] * 12, "note": "not a tool call"})
|
||||
chunks = [answer[i : i + 64] for i in range(0, len(answer), 64)]
|
||||
stream = [_sse({"content": chunk}) for chunk in chunks] + [_done()]
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(monkeypatch, [stream], payloads)
|
||||
|
||||
events = list(
|
||||
backend.generate_chat_completion_with_tools(
|
||||
messages = [{"role": "user", "content": "give me json"}],
|
||||
tools = _python_tool_schema(),
|
||||
max_tool_iterations = 1,
|
||||
)
|
||||
)
|
||||
|
||||
assert not [e for e in events if e.get("type") == "tool_args"]
|
||||
assert not [e for e in events if e.get("type") == "tool_start"]
|
||||
content_events = [e for e in events if e.get("type") == "content"]
|
||||
assert content_events and answer in content_events[-1]["text"]
|
||||
|
||||
|
||||
def test_provisional_text_card_closed_when_parse_fails(monkeypatch):
|
||||
"""A >=256-char enabled-name text sniff opens a provisional card; if the
|
||||
drained text then fails to parse (auto-heal off, truncated call), the
|
||||
DRAINING false-positive path must close the card with a tool_end instead of
|
||||
leaving it spinning forever."""
|
||||
|
||||
# Truncated mid-arguments and never closed: unparseable without healing.
|
||||
call_text = '<tool_call>{"name": "python", "arguments": {"code": "' + "x" * (
|
||||
_PROVISIONAL_ARGS_MIN_CHARS + 64
|
||||
)
|
||||
chunks = [call_text[i : i + 48] for i in range(0, len(call_text), 48)]
|
||||
stream = [_sse({"content": chunk}) for chunk in chunks] + [_done()]
|
||||
payloads: list[dict] = []
|
||||
backend = _make_backend(monkeypatch, [stream], payloads)
|
||||
|
||||
executed: list[tuple[str, dict]] = []
|
||||
|
||||
def fake_execute_tool(name, arguments, **_kwargs):
|
||||
executed.append((name, arguments))
|
||||
return "ok"
|
||||
|
||||
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
|
||||
|
||||
events = list(
|
||||
backend.generate_chat_completion_with_tools(
|
||||
messages = [{"role": "user", "content": "run it"}],
|
||||
tools = _python_tool_schema(),
|
||||
max_tool_iterations = 1,
|
||||
auto_heal_tool_calls = False,
|
||||
)
|
||||
)
|
||||
|
||||
starts = [e for e in events if e.get("type") == "tool_start"]
|
||||
ends = [e for e in events if e.get("type") == "tool_end"]
|
||||
assert starts and starts[0]["tool_call_id"] == "call_0"
|
||||
assert executed == [] # nothing parsed, nothing ran
|
||||
assert ends, "provisional card left dangling (no tool_end)"
|
||||
assert ends[-1]["tool_call_id"] == "call_0"
|
||||
|
|
|
|||
522
studio/backend/tests/test_sandbox_sitecustomize.py
Normal file
522
studio/backend/tests/test_sandbox_sitecustomize.py
Normal file
|
|
@ -0,0 +1,522 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Hermetic tests for the sandbox sitecustomize path-remap shim.
|
||||
|
||||
The shim (``core/inference/sandbox_site/sitecustomize.py``) runs at interpreter
|
||||
startup inside every sandboxed tool subprocess and remaps ChatGPT
|
||||
code-interpreter habit paths (``/mnt/data`` etc.) onto the per-conversation
|
||||
working directory. Importing it calls ``_install()``, which monkeypatches
|
||||
``builtins.open`` / ``io.open`` / ``os.makedirs`` / ``os.mkdir`` /
|
||||
``pathlib.Path.mkdir`` process-wide, so these tests
|
||||
load it into a throwaway module and restore those globals immediately, then
|
||||
exercise the pure ``_remap()`` function directly -- no subprocess, and no real
|
||||
``/mnt`` or ``/tmp`` writes. The mkdir test keeps the patch installed under a
|
||||
``chdir`` into ``tmp_path`` so the only real writes land in that temp dir.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import builtins
|
||||
import importlib.util
|
||||
import io
|
||||
import os
|
||||
import pathlib
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
_SHIM = (
|
||||
Path(__file__).resolve().parent.parent
|
||||
/ "core"
|
||||
/ "inference"
|
||||
/ "sandbox_site"
|
||||
/ "sitecustomize.py"
|
||||
)
|
||||
|
||||
|
||||
def _save_patch_targets():
|
||||
"""Snapshot every global the shim patches, so tests can restore them.
|
||||
|
||||
On Python < 3.11 the shim also repoints ``pathlib._NormalAccessor.open``
|
||||
(pathlib captured the original io.open at import there); the accessor is
|
||||
absent on 3.11+, so the snapshot skips it.
|
||||
"""
|
||||
accessor = getattr(pathlib, "_NormalAccessor", None)
|
||||
return (
|
||||
(builtins.open, io.open, os.open, os.makedirs, os.mkdir, pathlib.Path.mkdir),
|
||||
accessor,
|
||||
accessor.open if accessor is not None else None,
|
||||
)
|
||||
|
||||
|
||||
def _restore_patch_targets(saved):
|
||||
"""Undo _save_patch_targets so the test process stays clean."""
|
||||
globals_tuple, accessor, accessor_open = saved
|
||||
(builtins.open, io.open, os.open, os.makedirs, os.mkdir, pathlib.Path.mkdir) = globals_tuple
|
||||
if accessor is not None:
|
||||
accessor.open = accessor_open
|
||||
|
||||
|
||||
def _load_shim():
|
||||
"""Import the shim without leaving its open()/mkdir patches installed."""
|
||||
saved = _save_patch_targets()
|
||||
spec = importlib.util.spec_from_file_location("_sandbox_sitecustomize_under_test", _SHIM)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
try:
|
||||
spec.loader.exec_module(mod) # runs _install(), patching the globals
|
||||
finally:
|
||||
# Undo the process-wide patch so the test process stays clean.
|
||||
_restore_patch_targets(saved)
|
||||
mod._notified = True # silence the one-shot stderr notice in tests
|
||||
return mod
|
||||
|
||||
|
||||
def test_always_remap_prefixes_map_into_cwd(monkeypatch, tmp_path):
|
||||
mod = _load_shim()
|
||||
monkeypatch.chdir(tmp_path)
|
||||
cwd = os.getcwd()
|
||||
assert mod._remap("/mnt/data/out.txt") == os.path.join(cwd, "out.txt")
|
||||
assert mod._remap("/mnt/data") == cwd
|
||||
# Unrelated absolute and relative paths pass straight through.
|
||||
assert mod._remap("/etc/passwd") == "/etc/passwd"
|
||||
assert mod._remap("relative.txt") == "relative.txt"
|
||||
|
||||
|
||||
def test_prefix_remap_contains_parent_traversal_inside_cwd(monkeypatch, tmp_path):
|
||||
# A hallucinated habit path can carry '..' in its suffix. The remapped target
|
||||
# must stay under the per-conversation CWD, never climbing into a sibling
|
||||
# session's directory: '..' components are dropped, the rest of the subpath kept.
|
||||
mod = _load_shim()
|
||||
workdir = tmp_path / "session_current" / "work"
|
||||
workdir.mkdir(parents = True)
|
||||
monkeypatch.chdir(workdir)
|
||||
cwd = os.getcwd()
|
||||
|
||||
for escaping in (
|
||||
"/mnt/data/../other_session/file",
|
||||
"/mnt/data/../../secrets.txt",
|
||||
"/mnt/data/a/../../b/c.txt",
|
||||
"/mnt/data/./sub/./x.txt",
|
||||
):
|
||||
mapped = mod._remap(escaping)
|
||||
# Never escapes the CWD subtree.
|
||||
assert mapped == cwd or mapped.startswith(cwd + os.sep), (escaping, mapped)
|
||||
assert os.path.realpath(mapped).startswith(os.path.realpath(cwd))
|
||||
# '../other_session/file' collapses to CWD/other_session/file.
|
||||
assert mod._remap("/mnt/data/../other_session/file") == os.path.join(
|
||||
cwd, "other_session", "file"
|
||||
)
|
||||
# A bare '/mnt/data/..' with nothing left maps onto the CWD itself.
|
||||
assert mod._remap("/mnt/data/..") == cwd
|
||||
|
||||
|
||||
def test_write_fallback_refuses_dotdot_basename(monkeypatch, tmp_path):
|
||||
# basename('/no/such/tree/..') == '..'; joining that onto the CWD would target
|
||||
# its parent (outside the sandbox). The fallback must refuse such non-filename
|
||||
# basenames and return the path unchanged so the real open raises.
|
||||
mod = _load_shim()
|
||||
workdir = tmp_path / "work"
|
||||
workdir.mkdir()
|
||||
monkeypatch.chdir(workdir)
|
||||
for escaping in ("/no/such/tree/..", "/no/such/tree/.", "/no/such/tree/"):
|
||||
assert mod._remap_open(escaping, "w") == escaping
|
||||
|
||||
|
||||
def test_write_fallback_remaps_hallucinated_absolute_path(monkeypatch, tmp_path):
|
||||
# Models invent absolute paths from their CWD (e.g. /home/ubuntu/Sandbox/x.html),
|
||||
# which prefix lists cannot enumerate. A write/create-mode open on an absolute
|
||||
# path outside the CWD whose parent is missing is redirected to the basename in the CWD.
|
||||
mod = _load_shim()
|
||||
workdir = tmp_path / "workdir"
|
||||
workdir.mkdir()
|
||||
monkeypatch.chdir(workdir)
|
||||
cwd = os.getcwd()
|
||||
hallucinated = "/home/ubuntu/Sandbox/flappy_bird.html"
|
||||
for mode in ("w", "a", "x", "w+"):
|
||||
assert mod._remap_open(hallucinated, mode) == os.path.join(cwd, "flappy_bird.html")
|
||||
# A nested missing tree collapses to just the basename in the CWD.
|
||||
assert mod._remap_open("/no/such/tree/report.txt", "w") == os.path.join(cwd, "report.txt")
|
||||
|
||||
|
||||
def test_write_fallback_never_touches_read_modes(monkeypatch, tmp_path):
|
||||
# Reading a real (or genuinely missing) file must succeed/fail truthfully --
|
||||
# the fallback is write-only.
|
||||
mod = _load_shim()
|
||||
monkeypatch.chdir(tmp_path)
|
||||
for mode in ("r", "rb", "r+"):
|
||||
assert mod._remap_open("/etc/definitely_missing_xyz.conf", mode) == (
|
||||
"/etc/definitely_missing_xyz.conf"
|
||||
)
|
||||
|
||||
|
||||
def test_write_fallback_passes_through_existing_external_dir(monkeypatch, tmp_path):
|
||||
# A write to an absolute path whose parent dir exists is a deliberate, working
|
||||
# target and must NOT be redirected.
|
||||
mod = _load_shim()
|
||||
external = tmp_path / "external"
|
||||
external.mkdir()
|
||||
workdir = tmp_path / "workdir"
|
||||
workdir.mkdir()
|
||||
monkeypatch.chdir(workdir)
|
||||
target = str(external / "out.txt")
|
||||
assert mod._remap_open(target, "w") is target
|
||||
|
||||
|
||||
def test_write_fallback_never_clobbers_same_basename(monkeypatch, tmp_path):
|
||||
# A same-named CWD file is an unrelated persistent conversation file.
|
||||
# Redirecting an invented absolute path (missing parent) onto it would clobber
|
||||
# data the model never asked to touch, so the fallback refuses on collision for
|
||||
# every create mode: it returns the original path and the real open() raises.
|
||||
mod = _load_shim()
|
||||
monkeypatch.chdir(tmp_path)
|
||||
|
||||
existing = tmp_path / "report.txt"
|
||||
existing.write_text("KEEP-ME")
|
||||
|
||||
requested = "/definitely_missing_parent_7083/report.txt"
|
||||
for mode in ("w", "a", "x", "w+", "a+"):
|
||||
# Refused: returns the original absolute path unchanged (no redirect).
|
||||
assert mod._remap_open(requested, mode) == requested
|
||||
|
||||
# And opening the refused path really does raise, leaving the file intact.
|
||||
with pytest.raises(FileNotFoundError):
|
||||
open(mod._remap_open(requested, "w"), "w")
|
||||
assert existing.read_text() == "KEEP-ME"
|
||||
|
||||
# No collision -> still healed into the working directory as before.
|
||||
fresh = "/definitely_missing_parent_7083/brand_new.txt"
|
||||
assert mod._remap_open(fresh, "w") == os.path.join(os.getcwd(), "brand_new.txt")
|
||||
|
||||
|
||||
def test_write_fallback_reserves_same_target_on_repeated_writes(monkeypatch, tmp_path):
|
||||
# Iterative overwrite of the SAME invented path must keep landing on the CWD
|
||||
# target the fallback first healed it to. Once ./app.html exists, a naive
|
||||
# anti-clobber guard would return the original (parent-missing) path and every
|
||||
# regenerate would raise; the fallback must recognise its own prior remap and re-serve it.
|
||||
mod = _load_shim()
|
||||
monkeypatch.chdir(tmp_path)
|
||||
cwd = os.getcwd()
|
||||
invented = "/home/ubuntu/Sandbox/app.html"
|
||||
target = os.path.join(cwd, "app.html")
|
||||
|
||||
# First write: healed into the CWD, and create the file so the collision guard
|
||||
# would trigger on the next call without the fix.
|
||||
assert mod._remap_open(invented, "w") == target
|
||||
with open(mod._remap_open(invented, "w"), "w") as fh:
|
||||
fh.write("v1")
|
||||
|
||||
# Repeated overwrites of the same invented path stay on the same target.
|
||||
for _ in range(3):
|
||||
assert mod._remap_open(invented, "w") == target
|
||||
with open(mod._remap_open(invented, "w"), "w") as fh:
|
||||
fh.write("v2")
|
||||
assert Path(target).read_text() == "v2"
|
||||
|
||||
# A DIFFERENT invented source colliding on basename is still refused, so it can
|
||||
# never clobber the artifact the first path owns.
|
||||
other = "/opt/other/app.html"
|
||||
assert mod._remap_open(other, "w") == other
|
||||
|
||||
|
||||
def test_write_fallback_reserves_healed_target_across_separate_runs(monkeypatch, tmp_path):
|
||||
# Each tool call is a FRESH subprocess, so the in-process remap map is empty on
|
||||
# the next run while the healed file persists in the working directory. A second
|
||||
# run overwriting the SAME invented path (whose healed basename now exists) must
|
||||
# still re-serve that target via the on-disk sidecar, else the model could never
|
||||
# overwrite last turn's artifact. Each _load_shim() simulates a brand-new interpreter.
|
||||
monkeypatch.chdir(tmp_path)
|
||||
cwd = os.getcwd()
|
||||
invented = "/home/ubuntu/Sandbox/app.html"
|
||||
target = os.path.join(cwd, "app.html")
|
||||
|
||||
# Run 1: heal the invented path, create the file, persist source->target to the sidecar.
|
||||
run1 = _load_shim()
|
||||
assert run1._remap_open(invented, "w") == target
|
||||
with open(run1._remap_open(invented, "w"), "w") as fh:
|
||||
fh.write("v1")
|
||||
|
||||
# Run 2: brand-new interpreter, nothing in memory -- still recognises its prior
|
||||
# heal from the sidecar and re-serves it, even though ./app.html now exists
|
||||
# (which without the sidecar would trip the anti-clobber guard and raise).
|
||||
run2 = _load_shim()
|
||||
assert run2._remapped_writes == {}
|
||||
assert run2._remap_open(invented, "w") == target
|
||||
with open(run2._remap_open(invented, "w"), "w") as fh:
|
||||
fh.write("v2")
|
||||
assert Path(target).read_text() == "v2"
|
||||
|
||||
# A DIFFERENT invented source colliding only on basename is still refused across
|
||||
# runs: the sidecar records solely the source it healed, so an unrelated path
|
||||
# can never adopt/clobber the artifact.
|
||||
other = "/opt/other/app.html"
|
||||
assert run2._remap_open(other, "w") == other
|
||||
|
||||
# A foreign CWD file (created directly, never healed) stays protected in a later
|
||||
# run from an invented path sharing its basename.
|
||||
(tmp_path / "notes.txt").write_text("KEEP-ME")
|
||||
run3 = _load_shim()
|
||||
assert run3._remap_open("/some/missing/notes.txt", "w") == "/some/missing/notes.txt"
|
||||
with pytest.raises(FileNotFoundError):
|
||||
open(run3._remap_open("/some/missing/notes.txt", "w"), "w")
|
||||
assert (tmp_path / "notes.txt").read_text() == "KEEP-ME"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["r+", "rb+"])
|
||||
def test_read_update_modes_never_redirected_even_with_missing_parent(monkeypatch, tmp_path, mode):
|
||||
# r+ / rb+ REQUIRE the target to exist and never create; a "+" must not qualify
|
||||
# as creation, or a missing absolute path would be redirected onto a same-basename
|
||||
# workspace file and corrupt it. The parent is missing, so only the mode predicate
|
||||
# protects the victim.
|
||||
mod = _load_shim()
|
||||
monkeypatch.chdir(tmp_path)
|
||||
|
||||
victim = tmp_path / "victim.txt"
|
||||
victim.write_text("original")
|
||||
|
||||
requested = "/definitely_missing_parent_xyz/victim.txt"
|
||||
assert mod._remap_open(requested, mode) == requested
|
||||
with pytest.raises(FileNotFoundError):
|
||||
open(mod._remap_open(requested, mode), mode)
|
||||
assert victim.read_text() == "original"
|
||||
|
||||
|
||||
def test_existing_convention_prefix_is_not_shadowed(monkeypatch, tmp_path):
|
||||
# A convention prefix (/mnt/data etc.) is remapped ONLY while absent. If a real
|
||||
# host directory exists there it must pass through so its own filesystem semantics
|
||||
# apply: a real read succeeds, and a missing file under it is created there by a
|
||||
# write, never shadowed by a CWD file.
|
||||
mod = _load_shim()
|
||||
external = tmp_path / "real_prefix"
|
||||
external.mkdir()
|
||||
(external / "data.txt").write_text("real external content")
|
||||
|
||||
workdir = tmp_path / "conversation"
|
||||
workdir.mkdir()
|
||||
monkeypatch.chdir(workdir)
|
||||
monkeypatch.setattr(mod, "_PREFIXES", (str(external),))
|
||||
monkeypatch.setattr(mod, "_CONDITIONAL_PREFIXES", ())
|
||||
|
||||
target = str(external / "data.txt")
|
||||
# Prefix exists -> pass through for read and write.
|
||||
assert mod._remap(target) == target
|
||||
assert mod._remap_open(target, "r") == target
|
||||
assert mod._remap_open(target, "w") == target
|
||||
# A missing file under the EXISTING real prefix is left alone (parent exists),
|
||||
# so the real directory creates it -- not a CWD shadow.
|
||||
missing = str(external / "new.txt")
|
||||
assert mod._remap_open(missing, "w") == missing
|
||||
|
||||
# Remove the prefix directory -> healing resumes (absent prefix).
|
||||
(external / "data.txt").unlink()
|
||||
external.rmdir()
|
||||
assert mod._remap(target) == os.path.join(os.getcwd(), "data.txt")
|
||||
|
||||
|
||||
def test_os_open_and_path_touch_remap_convention_path(monkeypatch, tmp_path):
|
||||
# Path.touch() and other low-level creators go through os.open, not builtins/io.open.
|
||||
# Keep the shim's patches installed under a chdir into tmp_path so os.open is
|
||||
# patched, and confirm a convention path is healed into the CWD instead of raising.
|
||||
saved = _save_patch_targets()
|
||||
spec = importlib.util.spec_from_file_location("_sandbox_sitecustomize_osopen", _SHIM)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.chdir(tmp_path)
|
||||
cwd = os.getcwd()
|
||||
try:
|
||||
spec.loader.exec_module(mod) # installs the os.open patch
|
||||
mod._notified = True
|
||||
pathlib.Path("/mnt/data/touched.txt").touch()
|
||||
assert os.path.isfile(os.path.join(cwd, "touched.txt"))
|
||||
# Direct os.open with create flags is healed too.
|
||||
fd = os.open("/mnt/data/via_os_open.txt", os.O_CREAT | os.O_WRONLY, 0o600)
|
||||
os.close(fd)
|
||||
assert os.path.isfile(os.path.join(cwd, "via_os_open.txt"))
|
||||
finally:
|
||||
_restore_patch_targets(saved)
|
||||
|
||||
|
||||
def test_path_write_read_text_remap_convention_path(monkeypatch, tmp_path):
|
||||
# Path.open / write_text / read_text route through io.open (3.11+) or the captured
|
||||
# accessor open (< 3.11). Keep the patches installed under a chdir into tmp_path
|
||||
# and confirm a convention path is healed into the CWD on every version. This is
|
||||
# the hermetic guard for the 3.10 accessor path a plain io.open patch misses.
|
||||
saved = _save_patch_targets()
|
||||
spec = importlib.util.spec_from_file_location("_sandbox_sitecustomize_writetext", _SHIM)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.chdir(tmp_path)
|
||||
cwd = os.getcwd()
|
||||
try:
|
||||
spec.loader.exec_module(mod) # installs the io.open / accessor patch
|
||||
mod._notified = True
|
||||
pathlib.Path("/mnt/data/note.txt").write_text("pathlib remap")
|
||||
assert os.path.isfile(os.path.join(cwd, "note.txt"))
|
||||
# read_text goes through the same mapped path and sees what was written.
|
||||
assert pathlib.Path("/mnt/data/note.txt").read_text() == "pathlib remap"
|
||||
# A real absolute path passes through both patches untouched.
|
||||
real = tmp_path / "real.txt"
|
||||
pathlib.Path(str(real)).write_text("verbatim")
|
||||
assert real.read_text() == "verbatim"
|
||||
finally:
|
||||
_restore_patch_targets(saved)
|
||||
|
||||
|
||||
def test_write_fallback_leaves_relative_and_bytes_paths(monkeypatch, tmp_path):
|
||||
mod = _load_shim()
|
||||
monkeypatch.chdir(tmp_path)
|
||||
# Relative paths are already inside the CWD.
|
||||
assert mod._remap_open("out.txt", "w") == "out.txt"
|
||||
# Bytes paths are left untouched (prefix remap skips non-str).
|
||||
assert mod._remap_open(b"/no/such/tree/x.bin", "w") == b"/no/such/tree/x.bin"
|
||||
|
||||
|
||||
def test_remap_open_still_applies_prefix_remaps(monkeypatch, tmp_path):
|
||||
# The prefix remap runs first and preserves subpaths. A write heals onto the CWD
|
||||
# unconditionally; the write-mode fallback is only the last resort.
|
||||
mod = _load_shim()
|
||||
monkeypatch.chdir(tmp_path)
|
||||
cwd = os.getcwd()
|
||||
assert mod._remap_open("/mnt/data/sub/out.txt", "w") == os.path.join(cwd, "sub", "out.txt")
|
||||
# A read whose mapped target does NOT exist keeps the original path: a missing
|
||||
# input stays truthful, not silently redirected into the CWD.
|
||||
assert mod._remap_open("/mnt/data/sub/out.txt", "r") == "/mnt/data/sub/out.txt"
|
||||
|
||||
|
||||
def test_prefix_read_heals_only_when_mapped_target_exists(monkeypatch, tmp_path):
|
||||
# A convention-prefix READ must not redirect onto the CWD when the mapped target
|
||||
# is absent -- that masks a genuine missing-input error and could serve an
|
||||
# unrelated same-basename workdir file. It heals only when the mapped CWD target
|
||||
# exists, so re-reading an artifact an earlier write produced still works.
|
||||
mod = _load_shim()
|
||||
monkeypatch.chdir(tmp_path)
|
||||
cwd = os.getcwd()
|
||||
|
||||
# Mapped target absent: read keeps the original path (truthful miss).
|
||||
assert mod._remap_open("/mnt/data/input.csv", "r") == "/mnt/data/input.csv"
|
||||
with pytest.raises(FileNotFoundError):
|
||||
open(mod._remap_open("/mnt/data/input.csv", "r"))
|
||||
|
||||
# r+ (never creates) behaves the same: no redirect while absent.
|
||||
assert mod._remap_open("/mnt/data/input.csv", "r+") == "/mnt/data/input.csv"
|
||||
|
||||
# A write heals onto the CWD and creates the artifact...
|
||||
mapped = mod._remap_open("/mnt/data/input.csv", "w")
|
||||
assert mapped == os.path.join(cwd, "input.csv")
|
||||
with open(mapped, "w") as fh:
|
||||
fh.write("col\n1\n")
|
||||
|
||||
# ...and now a READ of the same convention path heals onto that existing artifact.
|
||||
read_target = mod._remap_open("/mnt/data/input.csv", "r")
|
||||
assert read_target == os.path.join(cwd, "input.csv")
|
||||
with open(read_target) as fh:
|
||||
assert fh.read() == "col\n1\n"
|
||||
|
||||
|
||||
def test_prefix_boundary_not_matched_by_similar_paths(monkeypatch, tmp_path):
|
||||
# The prefix match is anchored on a segment boundary (prefix or prefix + '/'), so
|
||||
# a sibling merely sharing the textual prefix must NOT be remapped: /workspace2
|
||||
# is not /workspace, /mnt/database is not /mnt/data.
|
||||
mod = _load_shim()
|
||||
monkeypatch.chdir(tmp_path)
|
||||
for unrelated in ("/workspace2/file.txt", "/mnt/database/x", "/home/sandboxed/y"):
|
||||
assert mod._remap(unrelated) == unrelated
|
||||
# And through open() for a read too (no silent redirect).
|
||||
assert mod._remap_open(unrelated, "r") == unrelated
|
||||
|
||||
|
||||
def test_tmp_outputs_is_a_conditional_prefix():
|
||||
mod = _load_shim()
|
||||
assert "/tmp/outputs" in mod._CONDITIONAL_PREFIXES
|
||||
# NOT in the always-remap set: /tmp exists on the host, so an unconditional remap
|
||||
# could shadow a real /tmp/outputs the user code made.
|
||||
assert "/tmp/outputs" not in mod._PREFIXES
|
||||
|
||||
|
||||
def test_tmp_outputs_remapped_only_while_absent(monkeypatch, tmp_path):
|
||||
mod = _load_shim()
|
||||
monkeypatch.chdir(tmp_path)
|
||||
cwd = os.getcwd()
|
||||
# Point the conditional prefix at a real temp location so we can toggle its
|
||||
# existence on disk instead of mocking os.path.exists.
|
||||
cond = str(tmp_path / "cond_outputs")
|
||||
monkeypatch.setattr(mod, "_CONDITIONAL_PREFIXES", (cond,))
|
||||
|
||||
# Absent: heal the habit path into the working directory (preserved/served).
|
||||
assert not os.path.exists(cond)
|
||||
assert mod._remap(cond + "/plot.png") == os.path.join(cwd, "plot.png")
|
||||
assert mod._remap(cond) == cwd
|
||||
|
||||
# Present (the user's own code created it): pass through, never shadowed.
|
||||
os.makedirs(cond)
|
||||
assert mod._remap(cond + "/plot.png") == cond + "/plot.png"
|
||||
assert mod._remap(cond) == cond
|
||||
|
||||
|
||||
def test_pathlib_mkdir_parents_remaps_convention_path(monkeypatch, tmp_path):
|
||||
# `Path('/mnt/data').mkdir(parents=True, exist_ok=True)` is a stock setup line.
|
||||
# pathlib drives it through os.mkdir per component and Path.is_dir()/os.stat on
|
||||
# FileExistsError, so the shim must patch os.mkdir AND Path.mkdir for the whole
|
||||
# parents/exist_ok dance to land in the CWD instead of raising. Keeps the mkdir
|
||||
# patches installed under a chdir into tmp_path and restores them in finally.
|
||||
saved = _save_patch_targets()
|
||||
spec = importlib.util.spec_from_file_location("_sandbox_sitecustomize_mkdir", _SHIM)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.chdir(tmp_path)
|
||||
cwd = os.getcwd()
|
||||
try:
|
||||
spec.loader.exec_module(mod) # installs the os.mkdir / Path.mkdir patches
|
||||
mod._notified = True
|
||||
# Bare convention path maps onto the CWD, which already exists: exist_ok=True
|
||||
# must be honoured against the mapped location, not raise.
|
||||
pathlib.Path("/mnt/data").mkdir(parents = True, exist_ok = True)
|
||||
# A nested convention path is created inside the CWD, parents and all.
|
||||
pathlib.Path("/mnt/data/plots/run1").mkdir(parents = True, exist_ok = True)
|
||||
assert os.path.isdir(os.path.join(cwd, "plots", "run1"))
|
||||
# Idempotent: exist_ok is evaluated on the mapped path (which now exists),
|
||||
# not the never-present /mnt/data.
|
||||
pathlib.Path("/mnt/data/plots/run1").mkdir(parents = True, exist_ok = True)
|
||||
|
||||
# Passthrough: real paths are created verbatim through both patches,
|
||||
# never remapped into the CWD.
|
||||
real_dir = tmp_path / "real_via_path"
|
||||
pathlib.Path(str(real_dir)).mkdir()
|
||||
assert real_dir.is_dir()
|
||||
real_os = tmp_path / "real_via_os"
|
||||
os.mkdir(str(real_os))
|
||||
assert real_os.is_dir()
|
||||
finally:
|
||||
_restore_patch_targets(saved)
|
||||
|
||||
|
||||
def test_read_of_missing_prefix_path_emits_no_notice(monkeypatch, tmp_path, capsys):
|
||||
# A read of a missing convention path keeps the original path and must not spend
|
||||
# the one-shot notice; a genuine remap afterward still notifies.
|
||||
mod = _load_shim()
|
||||
monkeypatch.chdir(tmp_path)
|
||||
mod._notified = False # re-arm the one-shot notice for this test
|
||||
# Read of a missing prefixed path: original kept, no notice, flag unspent.
|
||||
assert mod._remap_open("/mnt/data/missing.csv", "r") == "/mnt/data/missing.csv"
|
||||
assert mod._notified is False
|
||||
assert "does not exist" not in capsys.readouterr().err
|
||||
# A committed write then heals and fires the notice exactly once.
|
||||
assert mod._remap_open("/mnt/data/out.txt", "w") == os.path.join(os.getcwd(), "out.txt")
|
||||
assert mod._notified is True
|
||||
assert "/mnt/data does not exist in this sandbox" in capsys.readouterr().err
|
||||
|
||||
|
||||
def test_os_open_trunc_without_creat_missing_stays_truthful(monkeypatch, tmp_path):
|
||||
# O_TRUNC / O_APPEND without O_CREAT cannot create a missing file, so the shim
|
||||
# treats them as a read: a missing convention path stays truthful (the error
|
||||
# names the caller's path) and nothing is created in the CWD.
|
||||
saved = _save_patch_targets()
|
||||
spec = importlib.util.spec_from_file_location("_sandbox_sitecustomize_trunc", _SHIM)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.chdir(tmp_path)
|
||||
try:
|
||||
spec.loader.exec_module(mod)
|
||||
mod._notified = True
|
||||
with pytest.raises(FileNotFoundError) as exc:
|
||||
os.open("/mnt/data/missing_xyz.bin", os.O_WRONLY | os.O_TRUNC)
|
||||
assert exc.value.filename == "/mnt/data/missing_xyz.bin"
|
||||
assert not os.path.exists(os.path.join(os.getcwd(), "missing_xyz.bin"))
|
||||
finally:
|
||||
_restore_patch_targets(saved)
|
||||
|
|
@ -294,11 +294,16 @@ class TestSandboxEnvIsolation:
|
|||
"LANG",
|
||||
"TERM",
|
||||
"PYTHONIOENCODING",
|
||||
"PYTHONPATH",
|
||||
"VIRTUAL_ENV",
|
||||
"SystemRoot",
|
||||
}
|
||||
extras = set(env.keys()) - allowed
|
||||
assert not extras, f"sandbox env added unexpected keys: {extras}"
|
||||
# PYTHONPATH is whitelist-built, never inherited: only the sandbox
|
||||
# sitecustomize shim dir (code-interpreter path remap).
|
||||
assert env["PYTHONPATH"].endswith("sandbox_site")
|
||||
assert "leak-me" not in env["PYTHONPATH"]
|
||||
|
||||
def test_home_points_at_sandbox_workdir(self, tmp_path):
|
||||
from core.inference.tools import _build_safe_env
|
||||
|
|
@ -315,6 +320,23 @@ class TestSandboxEnvIsolation:
|
|||
env = _build_safe_env(str(tmp_path))
|
||||
assert env["TERM"] == "dumb"
|
||||
|
||||
def test_bypass_env_installs_sitecustomize_path_shim(self, tmp_path):
|
||||
# Bypass mode must install the same /mnt/data path-remap shim as the safe
|
||||
# env (finding 17), else /mnt/data writes work only in normal mode.
|
||||
from core.inference.tools import _SANDBOX_SITE_DIR, _build_bypass_env
|
||||
env = _build_bypass_env(str(tmp_path))
|
||||
assert _SANDBOX_SITE_DIR in env["PYTHONPATH"].split(os.pathsep)
|
||||
|
||||
def test_bypass_env_prepends_shim_and_keeps_inherited_pythonpath(self, monkeypatch, tmp_path):
|
||||
from core.inference.tools import _SANDBOX_SITE_DIR, _build_bypass_env
|
||||
|
||||
monkeypatch.setenv("PYTHONPATH", "/operator/libs")
|
||||
env = _build_bypass_env(str(tmp_path))
|
||||
parts = env["PYTHONPATH"].split(os.pathsep)
|
||||
# Shim first so its open()/makedirs remap wins, operator entries kept.
|
||||
assert parts[0] == _SANDBOX_SITE_DIR
|
||||
assert "/operator/libs" in parts
|
||||
|
||||
|
||||
class TestSandboxCpuRlimitDefault:
|
||||
"""Pin the default so a regression below 600s without opt-in is caught."""
|
||||
|
|
|
|||
1234
studio/backend/tests/test_tool_output_streaming.py
Normal file
1234
studio/backend/tests/test_tool_output_streaming.py
Normal file
File diff suppressed because it is too large
Load diff
83
studio/backend/tests/test_tool_stream_generator_drain.py
Normal file
83
studio/backend/tests/test_tool_stream_generator_drain.py
Normal file
|
|
@ -0,0 +1,83 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Regression tests for generator-close cleanup in the tool-streaming routes.
|
||||
|
||||
Tool streams run ``next(gen)`` in an ``asyncio.to_thread`` worker. Closing the
|
||||
generator while that worker is still inside ``next`` raises ``ValueError:
|
||||
generator already executing`` and skips the generator's ``finally`` (tool
|
||||
cleanup); the routes drain the pending task first (``_drain_pending_next_task``),
|
||||
which these tests exercise.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
from routes.inference import _drain_pending_next_task
|
||||
|
||||
|
||||
def test_drain_before_close_avoids_generator_already_executing():
|
||||
cancel_event = threading.Event()
|
||||
entered = threading.Event()
|
||||
finally_ran = threading.Event()
|
||||
|
||||
def blocking_gen():
|
||||
try:
|
||||
entered.set()
|
||||
# Blocking call inside next(gen) that respects the cancel flag.
|
||||
cancel_event.wait()
|
||||
yield "value"
|
||||
finally:
|
||||
finally_ran.set()
|
||||
|
||||
async def scenario():
|
||||
gen = blocking_gen()
|
||||
next_task = asyncio.create_task(asyncio.to_thread(next, gen, object()))
|
||||
await asyncio.to_thread(entered.wait) # worker now inside next(gen)
|
||||
|
||||
# Closing mid-next races and raises, leaving the finally unrun.
|
||||
with pytest.raises(ValueError):
|
||||
gen.close()
|
||||
assert not finally_ran.is_set()
|
||||
|
||||
# Draining sets the cancel flag so the worker returns; then close is
|
||||
# clean and the generator's finally runs.
|
||||
await _drain_pending_next_task(next_task, cancel_event)
|
||||
gen.close()
|
||||
return
|
||||
|
||||
asyncio.run(scenario())
|
||||
assert finally_ran.is_set()
|
||||
|
||||
|
||||
def test_drain_pending_next_task_is_noop_without_task():
|
||||
# None (task already consumed): draining is a no-op, cancel flag untouched.
|
||||
cancel_event = threading.Event()
|
||||
|
||||
asyncio.run(_drain_pending_next_task(None, cancel_event))
|
||||
assert not cancel_event.is_set()
|
||||
|
||||
|
||||
def test_drain_pending_next_task_returns_when_worker_finishes():
|
||||
# A worker finishing on its own drains without error; the cancel flag stays
|
||||
# set (the caller is tearing the stream down).
|
||||
cancel_event = threading.Event()
|
||||
release = threading.Event()
|
||||
|
||||
def gen():
|
||||
release.wait()
|
||||
yield "done"
|
||||
|
||||
async def scenario():
|
||||
g = gen()
|
||||
task = asyncio.create_task(asyncio.to_thread(next, g, object()))
|
||||
release.set() # let the worker complete before draining
|
||||
await _drain_pending_next_task(task, cancel_event)
|
||||
assert task.done()
|
||||
|
||||
asyncio.run(scenario())
|
||||
assert cancel_event.is_set()
|
||||
1196
studio/backend/tests/test_web_fetch_extraction.py
Normal file
1196
studio/backend/tests/test_web_fetch_extraction.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -3444,8 +3444,18 @@ const ComposerRightControls: FC<{
|
|||
const MessageError: FC = () => {
|
||||
return (
|
||||
<MessagePrimitive.Error>
|
||||
<ErrorPrimitive.Root className="aui-message-error-root mt-2 rounded-md bg-destructive/10 p-3 text-destructive text-sm dark:bg-destructive/5 dark:text-red-200">
|
||||
<ErrorPrimitive.Message className="aui-message-error-message line-clamp-2" />
|
||||
<ErrorPrimitive.Root className="aui-message-error-root mt-2 flex flex-wrap items-center gap-x-3 gap-y-2 rounded-md bg-destructive/10 p-3 text-destructive text-sm dark:bg-destructive/5 dark:text-red-200">
|
||||
<ErrorPrimitive.Message className="aui-message-error-message line-clamp-2 min-w-0 flex-1" />
|
||||
{/* Recovery path for interrupted/failed turns: regenerate in place. */}
|
||||
<ActionBarPrimitive.Reload asChild={true}>
|
||||
<button
|
||||
type="button"
|
||||
className="aui-message-error-retry inline-flex shrink-0 items-center gap-1.5 rounded-md border border-destructive/40 px-2.5 py-1 text-xs font-medium transition-colors hover:bg-destructive/15"
|
||||
>
|
||||
<RefreshCwIcon strokeWidth={1.75} className="size-3.5" />
|
||||
Retry
|
||||
</button>
|
||||
</ActionBarPrimitive.Reload>
|
||||
</ErrorPrimitive.Root>
|
||||
</MessagePrimitive.Error>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ 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 { ChevronDownIcon } from "lucide-react";
|
||||
import { Wrench01Icon } from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
|
|
@ -235,13 +236,29 @@ const ToolGroupImpl: FC<
|
|||
const messageRunning = useAuiState(
|
||||
({ message }) => message.status?.type === "running",
|
||||
);
|
||||
// Keep the group open once a confirmation forced it open, so answering an
|
||||
// allow/deny doesn't snap it shut between sequential tool calls. It reverts
|
||||
// to the default collapsed state once the turn finishes.
|
||||
// Force the group open when any call is receiving tool_output events.
|
||||
const toolLiveOutput = useChatRuntimeStore((s) => s.toolLiveOutput);
|
||||
const paneScope = useToolPaneScope();
|
||||
const hasLiveOutput = useAuiState(({ message }) =>
|
||||
message.parts
|
||||
.slice(startIndex, endIndex + 1)
|
||||
.some(
|
||||
(part) =>
|
||||
part.type === "tool-call" &&
|
||||
Object.prototype.hasOwnProperty.call(
|
||||
toolLiveOutput,
|
||||
toolOutputKey(paneScope, part.toolCallId),
|
||||
),
|
||||
),
|
||||
);
|
||||
// Keep the group open once a confirmation or live output forced it (so an
|
||||
// allow/deny doesn't snap it shut between calls); reverts once the turn ends.
|
||||
const forcedOpenRef = useRef(false);
|
||||
if (hasPendingConfirmation) forcedOpenRef.current = true;
|
||||
if (hasPendingConfirmation || hasLiveOutput) forcedOpenRef.current = true;
|
||||
const forceOpen =
|
||||
hasPendingConfirmation || (forcedOpenRef.current && messageRunning);
|
||||
hasPendingConfirmation ||
|
||||
(hasLiveOutput && messageRunning) ||
|
||||
(forcedOpenRef.current && messageRunning);
|
||||
|
||||
// Render single tool calls and canvases directly so cards never hide in a
|
||||
// collapsed group.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,63 @@
|
|||
// 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 "@/features/chat/stores/chat-runtime-store";
|
||||
import { toolOutputKey, useToolPaneScope } from "@/features/chat";
|
||||
import { useEffect, useMemo, useRef } from "react";
|
||||
import { tailText } from "./tool-result-output";
|
||||
|
||||
/**
|
||||
* Live-scrolling stdout/stderr pane for a running server-side tool, backed by
|
||||
* the transient `toolLiveOutput` map fed by `tool_output` SSE events. Renders
|
||||
* nothing until the first chunk, then follows the tail. Mounted only while
|
||||
* running; the finished card shows the persisted result instead.
|
||||
*/
|
||||
export function ToolLiveOutput({ toolCallId }: { toolCallId: string }) {
|
||||
const paneScope = useToolPaneScope();
|
||||
const output = useChatRuntimeStore(
|
||||
(s) => s.toolLiveOutput[toolOutputKey(paneScope, toolCallId)] ?? "",
|
||||
);
|
||||
const scrollRef = useRef<HTMLPreElement>(null);
|
||||
// Pinned to the bottom until the user scrolls up (handler below), so
|
||||
// streaming chunks no longer yank them down.
|
||||
const pinnedToBottom = useRef(true);
|
||||
|
||||
// The stream can reach hundreds of KB; render only the tail while live.
|
||||
const visible = useMemo(() => tailText(output).visible, [output]);
|
||||
|
||||
const handleScroll = () => {
|
||||
const el = scrollRef.current;
|
||||
if (!el) {
|
||||
return;
|
||||
}
|
||||
// Within 40px of the bottom counts as pinned (tolerates small nudges).
|
||||
pinnedToBottom.current =
|
||||
el.scrollHeight - el.scrollTop - el.clientHeight < 40;
|
||||
};
|
||||
|
||||
useEffect(() => {
|
||||
const el = scrollRef.current;
|
||||
if (el && pinnedToBottom.current) {
|
||||
el.scrollTop = el.scrollHeight;
|
||||
}
|
||||
}, [visible]);
|
||||
|
||||
if (!output) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="aui-tool-live-output mt-2 border-t border-dashed pt-2">
|
||||
<span className="text-xs font-medium text-muted-foreground">output</span>
|
||||
<pre
|
||||
ref={scrollRef}
|
||||
onScroll={handleScroll}
|
||||
className="mt-1 max-h-60 overflow-auto whitespace-pre-wrap break-words font-mono text-xs"
|
||||
>
|
||||
{visible}
|
||||
</pre>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
@ -0,0 +1,63 @@
|
|||
// 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 { useMemo, useState } from "react";
|
||||
|
||||
/** Tail-line cap so a huge output never mounts a megabyte <pre> block. */
|
||||
const TAIL_LINES = 2000;
|
||||
/** Char backstop for pathological single-line outputs. */
|
||||
const TAIL_CHARS = 200_000;
|
||||
|
||||
interface Tail {
|
||||
visible: string;
|
||||
hiddenLines: number;
|
||||
hiddenChars: number;
|
||||
}
|
||||
|
||||
export function tailText(text: string): Tail {
|
||||
let visible = text;
|
||||
let hiddenLines = 0;
|
||||
let hiddenChars = 0;
|
||||
const lines = visible.split("\n");
|
||||
if (lines.length > TAIL_LINES) {
|
||||
hiddenLines = lines.length - TAIL_LINES;
|
||||
visible = lines.slice(hiddenLines).join("\n");
|
||||
}
|
||||
if (visible.length > TAIL_CHARS) {
|
||||
hiddenChars = visible.length - TAIL_CHARS;
|
||||
visible = visible.slice(hiddenChars);
|
||||
}
|
||||
return { visible, hiddenLines, hiddenChars };
|
||||
}
|
||||
|
||||
/**
|
||||
* Finished-tool output pane: renders the tail (~2000 lines) with a "Show all"
|
||||
* toggle so a large output stays scrollable without janking the DOM. Copy
|
||||
* buttons still copy the FULL text (owned by the caller), not the tail.
|
||||
*/
|
||||
export function ToolResultOutput({ text }: { text: string }) {
|
||||
const [showAll, setShowAll] = useState(false);
|
||||
const tail = useMemo(() => tailText(text), [text]);
|
||||
const truncated = !showAll && (tail.hiddenLines > 0 || tail.hiddenChars > 0);
|
||||
|
||||
return (
|
||||
<>
|
||||
{truncated && (
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => setShowAll(true)}
|
||||
className="mt-1 rounded px-1.5 py-0.5 text-xs text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
|
||||
>
|
||||
{tail.hiddenLines > 0
|
||||
? `Show all (${tail.hiddenLines.toLocaleString()} earlier lines hidden)`
|
||||
: `Show all (${tail.hiddenChars.toLocaleString()} earlier chars hidden)`}
|
||||
</button>
|
||||
)}
|
||||
<pre className="mt-1 max-h-60 overflow-auto whitespace-pre-wrap break-words font-mono text-xs">
|
||||
{showAll ? text : tail.visible}
|
||||
</pre>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
|
@ -6,6 +6,7 @@
|
|||
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 } from "lucide-react";
|
||||
import { Tick02Icon } from "@/lib/tick-icon";
|
||||
|
|
@ -18,6 +19,14 @@ import {
|
|||
ToolFallbackRoot,
|
||||
ToolFallbackTrigger,
|
||||
} from "./tool-fallback";
|
||||
import { ToolLiveOutput } from "./tool-live-output";
|
||||
import { ToolResultOutput } from "./tool-result-output";
|
||||
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
|
||||
import {
|
||||
preferFullToolOutput,
|
||||
toolOutputKey,
|
||||
useToolPaneScope,
|
||||
} from "@/features/chat";
|
||||
|
||||
interface StructuredResult {
|
||||
text: string;
|
||||
|
|
@ -105,6 +114,7 @@ function isStructuredResult(val: unknown): val is StructuredResult {
|
|||
}
|
||||
|
||||
const PythonToolUIImpl: ToolCallMessagePartComponent = ({
|
||||
toolCallId,
|
||||
args,
|
||||
result,
|
||||
status,
|
||||
|
|
@ -112,6 +122,9 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({
|
|||
const code = (args as { code?: string })?.code ?? "";
|
||||
const firstLine = code.split("\n")[0]?.slice(0, 60) ?? "";
|
||||
const isRunning = status?.type === "running";
|
||||
// Args still streaming = the model is WRITING the code, not running it yet.
|
||||
const { propStatus } = useToolArgsStatus();
|
||||
const isWritingCode = isRunning && propStatus.code === "streaming";
|
||||
|
||||
let output: string;
|
||||
let images: string[] = [];
|
||||
|
|
@ -129,10 +142,19 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({
|
|||
output = "";
|
||||
}
|
||||
|
||||
// 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 displayOutput = preferFullToolOutput(fullOutput, output);
|
||||
|
||||
const authToken = getAuthToken();
|
||||
|
||||
return (
|
||||
<ToolFallbackRoot>
|
||||
// Open when mounted mid-run so live output shows; collapsed from history.
|
||||
<ToolFallbackRoot defaultOpen={isRunning}>
|
||||
<ToolFallbackTrigger
|
||||
toolName={firstLine ? `Python: ${firstLine}` : "Python"}
|
||||
status={status}
|
||||
|
|
@ -150,19 +172,21 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({
|
|||
|
||||
{/* Output */}
|
||||
{isRunning ? (
|
||||
<div className="mt-2 flex items-center gap-2 text-sm text-muted-foreground">
|
||||
<Spinner className="size-3.5" />
|
||||
<span>Running…</span>
|
||||
</div>
|
||||
) : output ? (
|
||||
<>
|
||||
<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>
|
||||
</div>
|
||||
{/* Live stdout streamed via tool_output SSE events. */}
|
||||
<ToolLiveOutput toolCallId={toolCallId} />
|
||||
</>
|
||||
) : displayOutput ? (
|
||||
<div className="mt-2 border-t border-dashed pt-2">
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="text-xs font-medium text-muted-foreground">output</span>
|
||||
<CopyBtn text={output} />
|
||||
<CopyBtn text={displayOutput} />
|
||||
</div>
|
||||
<pre className="mt-1 max-h-60 overflow-auto whitespace-pre-wrap break-words font-mono text-xs">
|
||||
{truncate(output)}
|
||||
</pre>
|
||||
<ToolResultOutput text={displayOutput} />
|
||||
</div>
|
||||
) : null}
|
||||
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@
|
|||
|
||||
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";
|
||||
|
|
@ -15,16 +16,17 @@ import {
|
|||
ToolFallbackRoot,
|
||||
ToolFallbackTrigger,
|
||||
} from "./tool-fallback";
|
||||
import { ToolLiveOutput } from "./tool-live-output";
|
||||
import { ToolResultOutput } from "./tool-result-output";
|
||||
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
|
||||
import {
|
||||
preferFullToolOutput,
|
||||
toolOutputKey,
|
||||
useToolPaneScope,
|
||||
} from "@/features/chat";
|
||||
|
||||
const MAX_DISPLAY = 10_000;
|
||||
const COPY_RESET_MS = 2000;
|
||||
|
||||
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);
|
||||
|
|
@ -65,12 +67,16 @@ function CopyBtn({ text }: { text: string }) {
|
|||
}
|
||||
|
||||
const TerminalToolUIImpl: ToolCallMessagePartComponent = ({
|
||||
toolCallId,
|
||||
args,
|
||||
result,
|
||||
status,
|
||||
}) => {
|
||||
const command = (args as { command?: string })?.command ?? "";
|
||||
const isRunning = status?.type === "running";
|
||||
// Args still streaming = the model is WRITING the command, not running it yet.
|
||||
const { propStatus } = useToolArgsStatus();
|
||||
const isWritingCommand = isRunning && propStatus.command === "streaming";
|
||||
const output =
|
||||
typeof result === "string"
|
||||
? result
|
||||
|
|
@ -78,8 +84,17 @@ const TerminalToolUIImpl: ToolCallMessagePartComponent = ({
|
|||
? JSON.stringify(result, null, 2)
|
||||
: "";
|
||||
|
||||
// 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 displayOutput = preferFullToolOutput(fullOutput, output);
|
||||
|
||||
return (
|
||||
<ToolFallbackRoot>
|
||||
// Open when mounted mid-run so live output shows; collapsed from history.
|
||||
<ToolFallbackRoot defaultOpen={isRunning}>
|
||||
<ToolFallbackTrigger
|
||||
toolName={command ? `$ ${command.slice(0, 60)}` : "Terminal"}
|
||||
status={status}
|
||||
|
|
@ -88,19 +103,21 @@ const TerminalToolUIImpl: ToolCallMessagePartComponent = ({
|
|||
<ToolFallbackContent>
|
||||
<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>Running…</span>
|
||||
</div>
|
||||
) : output ? (
|
||||
<>
|
||||
<div className="flex items-center gap-2 text-sm text-muted-foreground">
|
||||
<Spinner className="size-3.5" />
|
||||
<span>{isWritingCommand ? "Writing command…" : "Running…"}</span>
|
||||
</div>
|
||||
{/* Live stdout streamed via tool_output SSE events. */}
|
||||
<ToolLiveOutput toolCallId={toolCallId} />
|
||||
</>
|
||||
) : displayOutput ? (
|
||||
<div>
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="text-xs font-medium text-muted-foreground">output</span>
|
||||
<CopyBtn text={output} />
|
||||
<CopyBtn text={displayOutput} />
|
||||
</div>
|
||||
<pre className="mt-1 max-h-60 overflow-auto whitespace-pre-wrap break-words font-mono text-xs">
|
||||
{truncate(output)}
|
||||
</pre>
|
||||
<ToolResultOutput text={displayOutput} />
|
||||
</div>
|
||||
) : null}
|
||||
</div>
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import { parseParamCountB } from "@/lib/model-size";
|
|||
import { toast } from "@/lib/toast";
|
||||
import type { MessageTiming, ToolCallMessagePart } from "@assistant-ui/core";
|
||||
import type { ChatModelAdapter } from "@assistant-ui/react";
|
||||
import { parsePartialJsonObject } from "assistant-stream/utils";
|
||||
import {
|
||||
getExternalProviderApiKey,
|
||||
isCustomProviderType,
|
||||
|
|
@ -51,6 +52,11 @@ import {
|
|||
useChatRuntimeStore,
|
||||
} from "../stores/chat-runtime-store";
|
||||
import { useExternalProvidersStore } from "../stores/external-providers-store";
|
||||
import {
|
||||
shouldPreserveFullOutput,
|
||||
toolOutputKey,
|
||||
toolPaneScope,
|
||||
} from "../tool-output-scope";
|
||||
import type { ModelType } from "../types";
|
||||
import { isMultimodalResponse } from "../types/api";
|
||||
import type {
|
||||
|
|
@ -85,6 +91,7 @@ import {
|
|||
listGgufVariants,
|
||||
loadModel,
|
||||
streamChatCompletions,
|
||||
StreamInterruptedError,
|
||||
validateModel,
|
||||
} from "./chat-api";
|
||||
import {
|
||||
|
|
@ -228,6 +235,49 @@ function wait(ms: number): Promise<void> {
|
|||
return new Promise((resolve) => setTimeout(resolve, ms));
|
||||
}
|
||||
|
||||
/**
|
||||
* Best-effort partial parse of a live tool_args stream into a tool part's
|
||||
* `args`, so cards render the payload while the model is still writing it. The
|
||||
* structured path streams raw arguments JSON; the text path wraps it in call
|
||||
* markup, unwrapped here. Returns null until something parses; never throws.
|
||||
*/
|
||||
function parseLiveToolArgs(
|
||||
raw: string,
|
||||
): { args: Record<string, unknown>; argsText: string } | null {
|
||||
let candidate = raw.trimStart();
|
||||
if (!candidate.startsWith("{")) {
|
||||
const brace = candidate.indexOf("{");
|
||||
if (brace < 0) return null;
|
||||
candidate = candidate.slice(brace);
|
||||
}
|
||||
const parsed = parsePartialJsonObject(candidate) as
|
||||
| Record<string, unknown>
|
||||
| undefined;
|
||||
if (!parsed || typeof parsed !== "object") return null;
|
||||
// Call envelope from the text path: unwrap to the arguments payload.
|
||||
const inner = parsed.arguments ?? parsed.parameters;
|
||||
if (typeof parsed.name === "string" && inner !== undefined) {
|
||||
if (typeof inner === "string") {
|
||||
// Stringified arguments: partial-parse the inner JSON string.
|
||||
const innerParsed = parsePartialJsonObject(inner) as
|
||||
| Record<string, unknown>
|
||||
| undefined;
|
||||
if (innerParsed && typeof innerParsed === "object") {
|
||||
return { args: innerParsed, argsText: inner };
|
||||
}
|
||||
return null;
|
||||
}
|
||||
if (inner && typeof inner === "object" && !Array.isArray(inner)) {
|
||||
return {
|
||||
args: inner as Record<string, unknown>,
|
||||
argsText: JSON.stringify(inner),
|
||||
};
|
||||
}
|
||||
return null;
|
||||
}
|
||||
return { args: parsed, argsText: candidate };
|
||||
}
|
||||
|
||||
function parseSystemVariablesMap(raw: string): Record<string, unknown> {
|
||||
if (!raw.trim()) {
|
||||
return {};
|
||||
|
|
@ -1877,6 +1927,16 @@ 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,
|
||||
);
|
||||
const scopedToolOutputKey = (id: string) =>
|
||||
toolOutputKey(toolOutputPaneScope, id);
|
||||
const runToolLiveOutputKeys = new Set<string>();
|
||||
const resolvedThreadKey = resolvedThreadId ?? null;
|
||||
const pendingImageEditReferenceForRun = runtime.pendingImageEditReference;
|
||||
const selectedImageEditReference =
|
||||
|
|
@ -2426,6 +2486,34 @@ export function createOpenAIStreamAdapter(
|
|||
};
|
||||
// Tool call parts, cumulative; result lands on tool_end.
|
||||
const toolCallParts: PositionedToolCallPart[] = [];
|
||||
// Raw tool_args accumulator per card: the backend forwards arguments while
|
||||
// the model is still WRITING them, and the partial parse below feeds the
|
||||
// card's args so the code renders live.
|
||||
const liveArgsTextById = new Map<string, string>();
|
||||
// Backend tool ids ("call_0", ...) restart every response, so a bare id as
|
||||
// store key lets a later turn's stream overwrite the preserved output an
|
||||
// earlier still-mounted finished card reads (the tool_start stale-clear
|
||||
// only guards the forward direction). Mint one per-run-unique part id per
|
||||
// backend id (confirmation ids already synthesize their own) so each card
|
||||
// key is unique; every tool_start/output/args/end resolves the same id via
|
||||
// this map, dropped at tool_end.
|
||||
const toolPartIdByBackendId = new Map<string, string>();
|
||||
const resolveToolPartId = (backendToolCallId: string): string => {
|
||||
if (!backendToolCallId) {
|
||||
return toolCallParts[toolCallParts.length - 1]?.toolCallId ?? "";
|
||||
}
|
||||
const confirmationId =
|
||||
toolConfirmationIdsByBackendId.get(backendToolCallId);
|
||||
if (confirmationId) {
|
||||
return confirmationId;
|
||||
}
|
||||
let partId = toolPartIdByBackendId.get(backendToolCallId);
|
||||
if (!partId) {
|
||||
partId = `${backendToolCallId}:${crypto.randomUUID()}`;
|
||||
toolPartIdByBackendId.set(backendToolCallId, partId);
|
||||
}
|
||||
return partId;
|
||||
};
|
||||
// Latest Gemini text-part thoughtSignature; pinned onto the final
|
||||
// text MessagePart so next-turn replay carries it.
|
||||
let latestTextThoughtSignature: string | undefined;
|
||||
|
|
@ -3160,6 +3248,66 @@ export function createOpenAIStreamAdapter(
|
|||
anthropicRefusalSeen = true;
|
||||
continue;
|
||||
}
|
||||
if (toolEvent.type === "tool_output") {
|
||||
// Incremental stdout from a running tool: append to the live
|
||||
// store so the card renders it while the spinner runs. Final
|
||||
// result arrives via tool_end.
|
||||
const backendToolCallId =
|
||||
(toolEvent.tool_call_id as string) || "";
|
||||
const liveId = resolveToolPartId(backendToolCallId);
|
||||
const liveText =
|
||||
typeof toolEvent.text === "string" ? toolEvent.text : "";
|
||||
if (liveId && liveText) {
|
||||
const liveKey = scopedToolOutputKey(liveId);
|
||||
runToolLiveOutputKeys.add(liveKey);
|
||||
useChatRuntimeStore
|
||||
.getState()
|
||||
.appendToolLiveOutput(liveKey, liveText);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if (toolEvent.type === "tool_args") {
|
||||
// The model is still WRITING this call's arguments: accumulate
|
||||
// the raw stream and feed a partial parse into the part's args
|
||||
// so the card shows the code live. tool_start later replaces
|
||||
// args with the authoritative parse.
|
||||
const backendToolCallId =
|
||||
(toolEvent.tool_call_id as string) || "";
|
||||
const liveId = resolveToolPartId(backendToolCallId);
|
||||
const fragment =
|
||||
typeof toolEvent.text === "string" ? toolEvent.text : "";
|
||||
if (liveId && fragment) {
|
||||
const accum =
|
||||
(liveArgsTextById.get(liveId) ?? "") + fragment;
|
||||
liveArgsTextById.set(liveId, accum);
|
||||
const partial = parseLiveToolArgs(accum);
|
||||
const idx = toolCallParts.findIndex(
|
||||
(p) => p.toolCallId === liveId,
|
||||
);
|
||||
if (partial && idx !== -1) {
|
||||
const existing = toolCallParts[
|
||||
idx
|
||||
] as PositionedToolCallPart;
|
||||
toolCallParts[idx] = {
|
||||
...existing,
|
||||
args: partial.args as ToolCallMessagePart["args"],
|
||||
argsText: partial.argsText,
|
||||
};
|
||||
yield {
|
||||
content: buildAssistantContent(cumulativeText),
|
||||
metadata: {
|
||||
timing: buildTiming(
|
||||
streamStartTime,
|
||||
totalChunks,
|
||||
firstTokenTime,
|
||||
),
|
||||
custom: { reasoningDuration },
|
||||
},
|
||||
};
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
closeReasoningContent();
|
||||
const toolProvenance = parseToolProvenance(
|
||||
toolEvent.provenance,
|
||||
|
|
@ -3173,12 +3321,18 @@ export function createOpenAIStreamAdapter(
|
|||
const id =
|
||||
awaitingConfirmation && approvalId
|
||||
? `${toolConfirmationScopeId}:${approvalId}`
|
||||
: backendToolCallId ||
|
||||
approvalId ||
|
||||
`${toolEvent.tool_name}_${Date.now()}`;
|
||||
: backendToolCallId
|
||||
? resolveToolPartId(backendToolCallId)
|
||||
: approvalId ||
|
||||
`${toolEvent.tool_name}_${Date.now()}`;
|
||||
if (awaitingConfirmation && backendToolCallId) {
|
||||
toolConfirmationIdsByBackendId.set(backendToolCallId, id);
|
||||
}
|
||||
// "call_0" restarts every response: drop stale live/preserved
|
||||
// output under this key, else the card shows the previous call's.
|
||||
const staleKey = scopedToolOutputKey(id);
|
||||
useChatRuntimeStore.getState().clearToolLiveOutput(staleKey);
|
||||
useChatRuntimeStore.getState().clearToolFullOutput(staleKey);
|
||||
const toolArgs = (toolEvent.arguments ??
|
||||
{}) as ToolCallMessagePart["args"];
|
||||
const idx = toolCallParts.findIndex(
|
||||
|
|
@ -3222,17 +3376,35 @@ export function createOpenAIStreamAdapter(
|
|||
} else if (toolEvent.type === "tool_end") {
|
||||
const backendToolCallId =
|
||||
(toolEvent.tool_call_id as string) || "";
|
||||
const id =
|
||||
(backendToolCallId
|
||||
? toolConfirmationIdsByBackendId.get(backendToolCallId)
|
||||
: undefined) ||
|
||||
backendToolCallId ||
|
||||
toolCallParts[toolCallParts.length - 1]?.toolCallId ||
|
||||
"";
|
||||
const id = resolveToolPartId(backendToolCallId);
|
||||
if (backendToolCallId) {
|
||||
toolConfirmationIdsByBackendId.delete(backendToolCallId);
|
||||
toolPartIdByBackendId.delete(backendToolCallId);
|
||||
}
|
||||
useChatRuntimeStore.getState().clearToolConfirmation(id);
|
||||
// The result replaces the live output, but if the stream
|
||||
// captured MORE than the truncated result, preserve it so the
|
||||
// finished card keeps everything. Uses the shared predicate,
|
||||
// not a length compare (footer / "Exit code N:" / __IMAGES__
|
||||
// tail can make the result longer by byte).
|
||||
const liveKey = scopedToolOutputKey(id);
|
||||
const liveOutput =
|
||||
useChatRuntimeStore.getState().toolLiveOutput[liveKey] ??
|
||||
"";
|
||||
if (
|
||||
id &&
|
||||
shouldPreserveFullOutput(
|
||||
liveOutput,
|
||||
(toolEvent.result as string) ?? "",
|
||||
)
|
||||
) {
|
||||
useChatRuntimeStore
|
||||
.getState()
|
||||
.setToolFullOutput(liveKey, liveOutput);
|
||||
}
|
||||
useChatRuntimeStore.getState().clearToolLiveOutput(liveKey);
|
||||
runToolLiveOutputKeys.delete(liveKey);
|
||||
liveArgsTextById.delete(id);
|
||||
const idx = toolCallParts.findIndex(
|
||||
(p) => p.toolCallId === id,
|
||||
);
|
||||
|
|
@ -3797,7 +3969,16 @@ export function createOpenAIStreamAdapter(
|
|||
);
|
||||
if (!abortSignal.aborted) {
|
||||
const msg = err instanceof Error ? err.message : String(err);
|
||||
if (isContextLimitError(msg)) {
|
||||
if (err instanceof StreamInterruptedError) {
|
||||
// Connection dropped mid-turn: surface it explicitly (the rethrow
|
||||
// below also marks the message with an inline error + Retry).
|
||||
toast.error("Response interrupted", {
|
||||
description:
|
||||
"The connection dropped before the model finished. " +
|
||||
"The partial answer is kept. Use Retry to regenerate.",
|
||||
duration: 8000,
|
||||
});
|
||||
} else if (isContextLimitError(msg)) {
|
||||
// llama-server runs with --no-context-shift, returning a hard
|
||||
// error instead of silently dropping old KV-cache turns. Point
|
||||
// the user at the control that raises the ceiling.
|
||||
|
|
@ -3823,6 +4004,19 @@ export function createOpenAIStreamAdapter(
|
|||
}
|
||||
runtime.setGeneratingStatus(null);
|
||||
runtime.setToolStatus(null);
|
||||
// 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
|
||||
// diagnostics the user was watching vanish from the card.
|
||||
for (const liveKey of runToolLiveOutputKeys) {
|
||||
const store = useChatRuntimeStore.getState();
|
||||
const liveOutput = store.toolLiveOutput[liveKey] ?? "";
|
||||
if (liveOutput) {
|
||||
store.setToolFullOutput(liveKey, liveOutput);
|
||||
}
|
||||
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);
|
||||
|
|
|
|||
|
|
@ -29,6 +29,21 @@ import type {
|
|||
|
||||
export const CHAT_HISTORY_UPDATED_EVENT = "unsloth-chat-history-updated";
|
||||
|
||||
/**
|
||||
* Thrown when the chat SSE stream ends without a terminal signal (`[DONE]` or a
|
||||
* finish_reason chunk): the connection dropped mid-generation. The adapter
|
||||
* surfaces it as an explicit interrupted state instead of ending the turn.
|
||||
*/
|
||||
export class StreamInterruptedError extends Error {
|
||||
constructor() {
|
||||
super(
|
||||
"Response interrupted: the connection dropped before the model finished. " +
|
||||
"Use Retry to regenerate.",
|
||||
);
|
||||
this.name = "StreamInterruptedError";
|
||||
}
|
||||
}
|
||||
|
||||
export function notifyChatHistoryUpdated(): void {
|
||||
if (typeof window !== "undefined") {
|
||||
window.dispatchEvent(new Event(CHAT_HISTORY_UPDATED_EVENT));
|
||||
|
|
@ -851,12 +866,18 @@ export async function* streamChatCompletions(
|
|||
const decoder = new TextDecoder();
|
||||
let buffer = "";
|
||||
let completed = false;
|
||||
// EOF without `[DONE]` or a finish_reason chunk means the stream was cut
|
||||
// mid-generation: surface as interrupted, not silent success.
|
||||
let sawTerminalSignal = false;
|
||||
|
||||
try {
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) {
|
||||
completed = true;
|
||||
if (!sawTerminalSignal) {
|
||||
throw new StreamInterruptedError();
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
|
|
@ -877,6 +898,7 @@ export async function* streamChatCompletions(
|
|||
const dataText = dataLines.join("\n");
|
||||
if (dataText === "[DONE]") {
|
||||
completed = true;
|
||||
sawTerminalSignal = true;
|
||||
return;
|
||||
}
|
||||
|
||||
|
|
@ -903,10 +925,14 @@ export async function* streamChatCompletions(
|
|||
separatorIndex = buffer.search(/\r?\n\r?\n/);
|
||||
continue;
|
||||
}
|
||||
// Tool start/end events carry full input/output for the tool outputs panel
|
||||
// tool_start/end carry full input/output; tool_output streams
|
||||
// incremental stdout and tool_args streams the call arguments live.
|
||||
if (
|
||||
"type" in parsed &&
|
||||
(parsed.type === "tool_start" || parsed.type === "tool_end")
|
||||
(parsed.type === "tool_start" ||
|
||||
parsed.type === "tool_end" ||
|
||||
parsed.type === "tool_output" ||
|
||||
parsed.type === "tool_args")
|
||||
) {
|
||||
yield { _toolEvent: parsed } as unknown as OpenAIChatChunk;
|
||||
separatorIndex = buffer.search(/\r?\n\r?\n/);
|
||||
|
|
@ -925,6 +951,16 @@ export async function* streamChatCompletions(
|
|||
separatorIndex = buffer.search(/\r?\n\r?\n/);
|
||||
continue;
|
||||
}
|
||||
// finish_reason is a valid terminal signal for providers that close
|
||||
// the stream without an explicit [DONE] sentinel.
|
||||
const finishReason = (
|
||||
parsed as {
|
||||
choices?: Array<{ finish_reason?: string | null }>;
|
||||
}
|
||||
).choices?.[0]?.finish_reason;
|
||||
if (finishReason) {
|
||||
sawTerminalSignal = true;
|
||||
}
|
||||
yield parsed as OpenAIChatChunk;
|
||||
separatorIndex = buffer.search(/\r?\n\r?\n/);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -17,6 +17,11 @@ export {
|
|||
type Preset,
|
||||
} from "./chat-settings-sheet";
|
||||
export { useChatRuntimeStore } from "./stores/chat-runtime-store";
|
||||
export {
|
||||
preferFullToolOutput,
|
||||
toolOutputKey,
|
||||
useToolPaneScope,
|
||||
} from "./tool-output-scope";
|
||||
export { PermissionModeDropdown } from "./permission-mode-select";
|
||||
export { useChatSearchStore } from "./stores/chat-search-store";
|
||||
export { usePinnedChatsStore } from "./stores/pinned-chats-store";
|
||||
|
|
|
|||
|
|
@ -54,6 +54,7 @@ import {
|
|||
} from "./open-document";
|
||||
import { AudioAttachmentAdapter } from "./audio-attachment-adapter";
|
||||
import { useChatRuntimeStore } from "./stores/chat-runtime-store";
|
||||
import { ToolPaneScopeContext, toolPaneScope } from "./tool-output-scope";
|
||||
import type { MessageRecord, ModelType, ThreadRecord } from "./types";
|
||||
import {
|
||||
deleteStoredChatThreads,
|
||||
|
|
@ -1347,26 +1348,33 @@ export function ChatRuntimeProvider({
|
|||
|
||||
return (
|
||||
<AssistantRuntimeProvider runtime={runtime} aui={aui}>
|
||||
<ActiveThreadSync
|
||||
enabled={
|
||||
modelType === "base" && !pairId && !newThreadNonce && !initialThreadId
|
||||
}
|
||||
/>
|
||||
<ThreadBackendAutosave modelType={modelType} pairId={pairId} />
|
||||
<CancelRegistrar />
|
||||
{initialThreadId && (
|
||||
<ThreadAutoSwitch
|
||||
threadId={initialThreadId}
|
||||
syncActiveThreadId={syncActiveThreadId}
|
||||
{/* Pane identity for the tool-output store maps: the adapter prefixes its
|
||||
keys with this scope so concurrent panes with colliding tool ids
|
||||
("call_0") can't bleed live output into each other's cards. */}
|
||||
<ToolPaneScopeContext.Provider value={toolPaneScope(modelType, pairId)}>
|
||||
<ActiveThreadSync
|
||||
enabled={
|
||||
modelType === "base" &&
|
||||
!pairId &&
|
||||
!newThreadNonce &&
|
||||
!initialThreadId
|
||||
}
|
||||
/>
|
||||
)}
|
||||
{!initialThreadId && newThreadNonce && (
|
||||
<ThreadNewChatSwitch nonce={newThreadNonce} />
|
||||
)}
|
||||
{/* The view stays mounted (only CSS-hidden by RootLayout) while off-route
|
||||
so assistant-ui keeps the run attached and the stream alive. Unmounting
|
||||
it here aborts the in-flight generation. */}
|
||||
{children}
|
||||
<ThreadBackendAutosave modelType={modelType} pairId={pairId} />
|
||||
<CancelRegistrar />
|
||||
{initialThreadId && (
|
||||
<ThreadAutoSwitch
|
||||
threadId={initialThreadId}
|
||||
syncActiveThreadId={syncActiveThreadId}
|
||||
/>
|
||||
)}
|
||||
{!initialThreadId && newThreadNonce && (
|
||||
<ThreadNewChatSwitch nonce={newThreadNonce} />
|
||||
)}
|
||||
{/* The view stays mounted (only CSS-hidden) while off-route so the run
|
||||
stays attached and the stream alive; unmounting aborts generation. */}
|
||||
{children}
|
||||
</ToolPaneScopeContext.Provider>
|
||||
</AssistantRuntimeProvider>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -701,6 +701,13 @@ type ChatRuntimeStore = {
|
|||
*/
|
||||
webFetchToolsEnabled: boolean;
|
||||
toolStatus: string | null;
|
||||
/** 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>;
|
||||
/** Full live output of finished tools whose result was truncated for the
|
||||
* model, keyed by toolCallId. Set from tool_end; finished cards prefer it
|
||||
* over the truncated result. Session-transient. */
|
||||
toolFullOutput: Record<string, string>;
|
||||
generatingStatus: string | null;
|
||||
autoHealToolCalls: boolean;
|
||||
nudgeToolCalls: boolean;
|
||||
|
|
@ -836,6 +843,13 @@ type ChatRuntimeStore = {
|
|||
setRagOcrScanned: (enabled: boolean) => void;
|
||||
setRagCaptionFigures: (enabled: boolean) => void;
|
||||
setToolStatus: (status: string | null) => void;
|
||||
appendToolLiveOutput: (toolCallId: string, text: string) => void;
|
||||
/** Clear one tool's live output, or all when no id is given. */
|
||||
clearToolLiveOutput: (toolCallId?: string) => void;
|
||||
/** Preserve a finished tool's full live-streamed output for display. */
|
||||
setToolFullOutput: (toolCallId: string, text: string) => void;
|
||||
/** 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;
|
||||
setAutoHealToolCalls: (enabled: boolean) => void;
|
||||
|
|
@ -1166,6 +1180,8 @@ 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,
|
||||
toolLiveOutput: {},
|
||||
toolFullOutput: {},
|
||||
generatingStatus: null,
|
||||
activeDiffusionCanvas: null,
|
||||
autoHealToolCalls: true,
|
||||
|
|
@ -1415,6 +1431,8 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
// Only the per-session enable pill resets; source/mode/top_k persist.
|
||||
ragEnabled: false,
|
||||
toolStatus: null,
|
||||
toolLiveOutput: {},
|
||||
toolFullOutput: {},
|
||||
activeDiffusionCanvas: null,
|
||||
kvCacheDtype: null,
|
||||
loadedKvCacheDtype: null,
|
||||
|
|
@ -1638,6 +1656,43 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
|
|||
return { ragCaptionFigures };
|
||||
}),
|
||||
setToolStatus: (toolStatus) => set({ toolStatus }),
|
||||
appendToolLiveOutput: (toolCallId, text) =>
|
||||
set((state) => ({
|
||||
toolLiveOutput: {
|
||||
...state.toolLiveOutput,
|
||||
[toolCallId]: (state.toolLiveOutput[toolCallId] ?? "") + text,
|
||||
},
|
||||
})),
|
||||
setToolFullOutput: (toolCallId, text) =>
|
||||
set((state) => ({
|
||||
toolFullOutput: {
|
||||
...state.toolFullOutput,
|
||||
[toolCallId]: text,
|
||||
},
|
||||
})),
|
||||
clearToolFullOutput: (toolCallId) =>
|
||||
set((state) => {
|
||||
if (!(toolCallId in state.toolFullOutput)) {
|
||||
return {};
|
||||
}
|
||||
const next = { ...state.toolFullOutput };
|
||||
delete next[toolCallId];
|
||||
return { toolFullOutput: next };
|
||||
}),
|
||||
clearToolLiveOutput: (toolCallId) =>
|
||||
set((state) => {
|
||||
if (toolCallId === undefined) {
|
||||
return Object.keys(state.toolLiveOutput).length
|
||||
? { toolLiveOutput: {} }
|
||||
: {};
|
||||
}
|
||||
if (!(toolCallId in state.toolLiveOutput)) {
|
||||
return {};
|
||||
}
|
||||
const next = { ...state.toolLiveOutput };
|
||||
delete next[toolCallId];
|
||||
return { toolLiveOutput: next };
|
||||
}),
|
||||
setActiveDiffusionCanvas: (activeDiffusionCanvas) =>
|
||||
set({ activeDiffusionCanvas }),
|
||||
setGeneratingStatus: (generatingStatus) => set({ generatingStatus }),
|
||||
|
|
|
|||
95
studio/frontend/src/features/chat/tool-output-scope.ts
Normal file
95
studio/frontend/src/features/chat/tool-output-scope.ts
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
// 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 { createContext, useContext } from "react";
|
||||
import type { ModelType } from "./types";
|
||||
|
||||
/**
|
||||
* Pane scope prefix for the transient tool-output store keys.
|
||||
*
|
||||
* Local GGUF tool ids are only unique within one response ("call_0", "call_1",
|
||||
* ...), and panes stream concurrently (compare mode mounts two runtimes; the
|
||||
* main view stays CSS-hidden off-route), so a bare id would let one pane's
|
||||
* stdout bleed into another's same-id card. Derived from static props
|
||||
* (`modelType` + `pairId`) shared by writer (adapter) and reader (components)
|
||||
* via one `ChatRuntimeProvider`, so they can never disagree.
|
||||
*/
|
||||
export function toolPaneScope(modelType?: ModelType, pairId?: string): string {
|
||||
return `${modelType ?? "base"}\u0000${pairId ?? ""}`;
|
||||
}
|
||||
|
||||
export const ToolPaneScopeContext = createContext<string>(toolPaneScope());
|
||||
|
||||
export function useToolPaneScope(): string {
|
||||
return useContext(ToolPaneScopeContext);
|
||||
}
|
||||
|
||||
/** Store key for the live/full tool output maps: pane scope + tool call id. */
|
||||
export function toolOutputKey(paneScope: string, toolCallId: string): string {
|
||||
return `${paneScope}\u0000${toolCallId}`;
|
||||
}
|
||||
|
||||
// Footer the backend appends when it truncates a result to protect the context
|
||||
// window (see backend tools._truncate). Marks where the result stops being a
|
||||
// copy of the stream, so it distinguishes "just truncated" from "carries
|
||||
// failure/exit status the stream never produced".
|
||||
const TRUNCATION_FOOTER_MARKER = "\n\n... (truncated";
|
||||
|
||||
/**
|
||||
* Whether the live stdout holds more real output than the model-visible
|
||||
* `result` and should be preserved for the finished card. Shared by writer
|
||||
* (retain?) and reader (display?) so they agree.
|
||||
*
|
||||
* True when the result is truncated, OR the stream is longer. Truncation can't
|
||||
* fall back to length: a truncated result may be longer by byte count once its
|
||||
* footer / `Exit code N:` / `__IMAGES__` blob is appended, yet the stream still
|
||||
* holds more stdout. Also true when a short stream is absent from the result: a
|
||||
* timed-out/cancelled tool returns only a status line, so length alone would
|
||||
* drop the partial stdout the stream captured.
|
||||
*/
|
||||
export function shouldPreserveFullOutput(full: string, result: string): boolean {
|
||||
if (!full) {
|
||||
return false;
|
||||
}
|
||||
if (result.includes(TRUNCATION_FOOTER_MARKER)) {
|
||||
return true;
|
||||
}
|
||||
if (full.length > result.length) {
|
||||
return true;
|
||||
}
|
||||
// Stream no longer than the result, but a timed-out/cancelled tool's status
|
||||
// line never echoes the captured stdout: preserve the stream whenever its
|
||||
// content is absent from the result (trimmed to ignore trailing-newline drift).
|
||||
const core = full.trim();
|
||||
return core.length > 0 && !result.includes(core);
|
||||
}
|
||||
|
||||
/**
|
||||
* Pick what a finished python/terminal card shows. Prefer the fuller live
|
||||
* stream over the truncated `result`, but the result can carry failure/exit
|
||||
* text that never reached stdout ("Exit code N: ...", timeouts), so show the
|
||||
* stream when the result is just a truncated prefix of it, else append the
|
||||
* result so its status survives (and the copy button copies both).
|
||||
*/
|
||||
export function preferFullToolOutput(full: string, result: string): string {
|
||||
if (!shouldPreserveFullOutput(full, result)) {
|
||||
return result;
|
||||
}
|
||||
const marker = result.indexOf(TRUNCATION_FOOTER_MARKER);
|
||||
const core = marker === -1 ? result : result.slice(0, marker);
|
||||
if (!core || full === result || full.startsWith(core)) {
|
||||
return full;
|
||||
}
|
||||
// Failed executions prefix the result (not the stream) with "Exit code N:\n",
|
||||
// so `full.startsWith(core)` above misses and a plain append would duplicate
|
||||
// the stdout. Re-attach just the exit prefix (and any missing-path hint) to
|
||||
// the fuller stream so the status survives without duplicating stdout.
|
||||
const exitMatch = core.match(/^(Exit code -?\d+:\n)([\s\S]*)$/);
|
||||
if (exitMatch && full.startsWith(exitMatch[2])) {
|
||||
const hint = result.match(/\nHint:[\s\S]*$/)?.[0] ?? "";
|
||||
return `${exitMatch[1]}${full}${hint}`;
|
||||
}
|
||||
return `${full.replace(/\s+$/, "")}\n\n${result}`;
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue