From e1fdb140b4bad80767c2e74c8db5c12130bb65a1 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 22 Jun 2026 15:37:00 +0000 Subject: [PATCH] Watch disconnects on Anthropic streams; keep timestamps in Gemma values Two follow-ups on the streaming and tool-parse paths: - _anthropic_tool_stream and _anthropic_plain_stream drove generation through asyncio.to_thread(next, gen, ...) and only polled is_disconnected() between events, so a client disconnect during prefill or a long generation/tool step held the decode slot until the next event or a failed send. Both now run the _await_disconnect_then_cancel watcher used by the other local streams, stop it in finally, and break promptly when cancel_event is set. - _GEMMA_NEXT_KEY_RE treated any comma followed by word-chars-then-colon as the next key, so a bare value such as "meet at 10:00, 11:00 tomorrow" was split into bogus keys. The next-key token must now be identifier-shaped (start with a letter or underscore), so a comma before a timestamp, ratio, or other numeric-then-colon text stays part of the value. Adds a timestamp-in-bare-value regression test. --- studio/backend/core/tool_healing.py | 7 ++++-- studio/backend/routes/inference.py | 24 +++++++++++++++++-- .../tests/test_gemma_tool_parse_edge_cases.py | 13 ++++++++++ 3 files changed, 40 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 7ef1870519..c647738a3a 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -39,8 +39,11 @@ _PARAM_CLOSE_TAG = "" _FUNC_CLOSE_TAG = "" # A bare (unquoted) Gemma value ends at `}` or at a comma that begins the next # `key:` pair. A comma NOT followed by a key token is part of the value (e.g. -# `location:New York, NY`), so it must not terminate the value. -_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[\w-]+\s*:") +# `location:New York, NY`), so it must not terminate the value. The key token +# must be identifier-shaped (start with a letter or underscore); a comma +# followed by digits-then-colon is value text such as a timestamp or ratio +# (`meet at 10:00, 11:00 tomorrow`), not a new key. +_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w-]*\s*:") def _balanced_brace_end( diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1f50f2815a..6b31662885 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -8471,9 +8471,17 @@ async def _anthropic_tool_stream( drop_until_tool_end = False gen = run_gen() + # Concurrent disconnect watcher: the loop only polls is_disconnected() + # between events, so a client disconnect during a long prefill or + # generation step would otherwise hold the decode slot until the next + # event or a failed send. The watcher sets cancel_event so the backend + # stops promptly. + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: while True: - if await request.is_disconnected(): + if cancel_event.is_set() or await request.is_disconnected(): cancel_event.set() return event = await asyncio.to_thread(next, gen, _sentinel) @@ -8521,6 +8529,8 @@ async def _anthropic_tool_stream( if _error_event is not None: yield _error_event return + finally: + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) stop_reason = openai_finish_to_anthropic_stop( captured_finish_reason, had_tool_calls = ends_on_tool_use @@ -8557,9 +8567,17 @@ async def _anthropic_plain_stream( captured_finish_reason = None gen = run_gen() + # Concurrent disconnect watcher: the loop only polls is_disconnected() + # between chunks, so a client disconnect during a long prefill or + # generation step would otherwise hold the decode slot until the next + # chunk or a failed send. The watcher sets cancel_event so the backend + # stops promptly. + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: while True: - if await request.is_disconnected(): + if cancel_event.is_set() or await request.is_disconnected(): cancel_event.set() return cumulative = await asyncio.to_thread(next, gen, _sentinel) @@ -8582,6 +8600,8 @@ async def _anthropic_plain_stream( if _error_event is not None: yield _error_event return + finally: + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) 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): diff --git a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py index b198f426df..15ca81cb5c 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -45,6 +45,19 @@ def test_normal_multi_key_arguments_still_split(): assert _args(calls[0]) == {"a": 1, "b": "hello", "c": "x,y"} +def test_bare_value_with_timestamps_after_comma_is_kept(): + # A comma followed by digits-then-colon (a timestamp/ratio) is value text, + # not a new key, so the whole query must be preserved as one argument. + calls = parse_tool_calls_from_text( + "<|tool_call>call:remind{query:meet at 10:00, 11:00 tomorrow,priority:high}" + ) + assert len(calls) == 1, calls + assert _args(calls[0]) == { + "query": "meet at 10:00, 11:00 tomorrow", + "priority": "high", + } + + def test_marker_inside_json_argument_is_not_a_second_call(): # A python call whose `code` argument contains a Gemma marker string. The # marker is data and must not execute as a second `terminal` call.