stop tool dispatch at first completed json object
Tool execution no longer waits for the full streamed assistant text. We now detect the first explicitly closed top-level JSON object, freeze that snapshot as the canonical tool request, and stop the model stream there for dispatch. To make that safe, DirtyJson completion semantics are tightened so completed=true only means the root object was explicitly closed, not that parsing hit end of file. I also restricted the new extraction path to object roots only, since tool calls are always brace-delimited objects, and added tests for parser completion and early stream stop.
Alessandro committed
Apr 3, 2026 at 16:56 UTC
5a2223596a4aaea93d24e628d15c9d2d30b8ec35
6 files changed
+256
-40
agent.py
+25
-1
@@ -388,6 +388,7 @@ class Agent:
388
self.context.streaming_agent = self # mark self as current streamer
389
self.loop_data.iteration += 1
390
self.loop_data.params_temporary = {} # clear temporary params
391
+ last_response_stream_full = ""
392
393
# call message_loop_start extensions
394
await extension.call_extensions_async(
@@ -425,12 +426,32 @@ class Agent:
426
await self.handle_reasoning_stream(stream_data["full"])
427
428
async def stream_callback(chunk: str, full: str):
429
+ nonlocal last_response_stream_full
430
await self.handle_intervention()
431
# output the agent response stream
432
if chunk == full:
433
printer.print("Response: ") # start of response
434
# Pass chunk and full data to extensions for processing
435
stream_data = {"chunk": chunk, "full": full}
436
+ stop_response: str | None = None
437
+
438
+ snapshot = extract_tools.extract_json_root_string(full)
439
+ if snapshot:
440
+ parsed_snapshot = extract_tools.json_parse_dirty(snapshot)
441
+ if parsed_snapshot is not None:
442
+ try:
443
+ await self.validate_tool_request(parsed_snapshot)
444
+ except Exception:
445
+ pass
446
+ else:
447
+ previous_full = last_response_stream_full
448
+ stream_data["full"] = snapshot
449
+ if snapshot.startswith(previous_full):
450
+ stream_data["chunk"] = snapshot[len(previous_full) :]
451
+ else:
452
+ stream_data["chunk"] = snapshot
453
+ stop_response = snapshot
454
+
455
await extension.call_extensions_async(
456
"response_stream_chunk",
457
self,
@@ -442,6 +463,9 @@ class Agent:
463
printer.stream(stream_data["chunk"])
464
# Use the potentially modified full text for downstream processing
465
await self.handle_response_stream(stream_data["full"])
466
+ last_response_stream_full = stream_data["full"]
467
+ if stop_response is not None:
468
+ return stop_response
469
470
# call main LLM
471
agent_response, _reasoning = await self.call_chat_model(
@@ -770,7 +794,7 @@ class Agent:
794
async def call_chat_model(
795
self,
796
messages: list[BaseMessage],
773
- response_callback: Callable[[str, str], Awaitable[None]] | None = None,
797
+ response_callback: Callable[[str, str], Awaitable[str | None]] | None = None,
798
reasoning_callback: Callable[[str, str], Awaitable[None]] | None = None,
799
background: bool = False,
800
explicit_caching: bool = True,
helpers/dirty_json.py
+16
-8
@@ -28,10 +28,10 @@ class DirtyJson:
28
self.completed = False
29
self._parsing_started = False
30
31
- def _pop_stack(self):
32
- """Pop from the parsing stack and mark completed if root structure is closed."""
31
+ def _pop_stack(self, root_closed: bool = False):
32
+ """Pop from the parsing stack and mark completed only on an explicit root close."""
33
self.stack.pop()
34
- if self._parsing_started and not self.stack:
34
+ if root_closed and self._parsing_started and not self.stack:
35
self.completed = True
36
37
@staticmethod
@@ -103,6 +103,8 @@ class DirtyJson:
103
self._advance()
104
105
def _parse(self):
106
+ if self.completed and not self.stack:
107
+ return
108
if self.result is None:
109
self.result = self._parse_value()
110
else:
@@ -110,6 +112,8 @@ class DirtyJson:
112
113
def _continue_parsing(self):
114
while self.current_char is not None:
115
+ if self.completed and not self.stack:
116
+ return
117
if isinstance(self.result, dict):
118
self._parse_object_content()
119
elif isinstance(self.result, list):
@@ -122,7 +126,9 @@ class DirtyJson:
126
def _parse_value(self):
127
self._skip_whitespace()
128
if self.current_char == "{":
125
- if self._peek(1) == "{": # Handle {{
129
+ # Only treat doubled braces as a wrapper at the root; nested objects
130
+ # must keep their closing braces paired correctly.
131
+ if not self.stack and self._peek(1) == "{": # Handle {{
132
self._advance(2)
133
return self._parse_object()
134
elif self.current_char == "[":
@@ -169,11 +175,13 @@ class DirtyJson:
175
while self.current_char is not None:
176
self._skip_whitespace()
177
if self.current_char == "}":
172
- if self._peek(1) == "}": # Handle }}
178
+ # Root-level wrapper outputs may end in "}}"; nested objects must
179
+ # still close one brace at a time.
180
+ if len(self.stack) == 1 and self._peek(1) == "}": # Handle }}
181
self._advance(2)
182
else:
183
self._advance()
176
- self._pop_stack()
184
+ self._pop_stack(root_closed=True)
185
return
186
if self.current_char is None:
187
self._pop_stack()
@@ -234,7 +242,7 @@ class DirtyJson:
242
self._skip_whitespace()
243
if self.current_char == "]":
244
self._advance()
237
- self._pop_stack()
245
+ self._pop_stack(root_closed=True)
246
return
247
value = self._parse_value()
248
self.stack[-1].append(value)
@@ -246,7 +254,7 @@ class DirtyJson:
254
if self.current_char is None or self.current_char == "]":
255
if self.current_char == "]":
256
self._advance()
249
- self._pop_stack()
257
+ self._pop_stack(root_closed=True)
258
return
259
elif self.current_char != "]":
260
self._pop_stack()
helpers/extract_tools.py
+22
@@ -19,6 +19,28 @@ def json_parse_dirty(json: str) -> dict[str, Any] | None:
19
return None
20
return None
21
22
+def extract_json_root_string(content: str) -> str | None:
23
+ if not content or not isinstance(content, str):
24
+ return None
25
+
26
+ start = content.find("{")
27
+ if start == -1:
28
+ return None
29
+ first_array = content.find("[")
30
+ if first_array != -1 and first_array < start:
31
+ return None
32
+
33
+ parser = DirtyJson()
34
+ try:
35
+ parser.parse(content[start:])
36
+ except Exception:
37
+ return None
38
+
39
+ if not parser.completed:
40
+ return None
41
+
42
+ return content[start : start + parser.index]
43
+
44
45
def extract_json_object_string(content):
46
start = content.find("{")
models.py
+41
-31
@@ -475,7 +475,7 @@ class LiteLLMChatWrapper(SimpleChatModel):
475
system_message="",
476
user_message="",
477
messages: List[BaseMessage] | None = None,
478
- response_callback: Callable[[str, str], Awaitable[None]] | None = None,
478
+ response_callback: Callable[[str, str], Awaitable[str | None]] | None = None,
479
reasoning_callback: Callable[[str, str], Awaitable[None]] | None = None,
480
tokens_callback: Callable[[str, int], Awaitable[None]] | None = None,
481
rate_limiter_callback: (
@@ -526,36 +526,46 @@ class LiteLLMChatWrapper(SimpleChatModel):
526
527
if stream:
528
# iterate over chunks
529
- async for chunk in _completion: # type: ignore
530
- got_any_chunk = True
531
- # parse chunk
532
- parsed = _parse_chunk(chunk)
533
- output = result.add_chunk(parsed)
534
-
535
- # collect reasoning delta and call callbacks
536
- if output["reasoning_delta"]:
537
- if reasoning_callback:
538
- await reasoning_callback(output["reasoning_delta"], result.reasoning)
539
- if tokens_callback:
540
- await tokens_callback(
541
- output["reasoning_delta"],
542
- approximate_tokens(output["reasoning_delta"]),
543
- )
544
- # Add output tokens to rate limiter if configured
545
- if limiter:
546
- limiter.add(output=approximate_tokens(output["reasoning_delta"]))
547
- # collect response delta and call callbacks
548
- if output["response_delta"]:
549
- if response_callback:
550
- await response_callback(output["response_delta"], result.response)
551
- if tokens_callback:
552
- await tokens_callback(
553
- output["response_delta"],
554
- approximate_tokens(output["response_delta"]),
555
- )
556
- # Add output tokens to rate limiter if configured
557
- if limiter:
558
- limiter.add(output=approximate_tokens(output["response_delta"]))
529
+ stop_response: str | None = None
530
+ try:
531
+ async for chunk in _completion: # type: ignore
532
+ got_any_chunk = True
533
+ # parse chunk
534
+ parsed = _parse_chunk(chunk)
535
+ output = result.add_chunk(parsed)
536
+
537
+ # collect reasoning delta and call callbacks
538
+ if output["reasoning_delta"]:
539
+ if reasoning_callback:
540
+ await reasoning_callback(output["reasoning_delta"], result.reasoning)
541
+ if tokens_callback:
542
+ await tokens_callback(
543
+ output["reasoning_delta"],
544
+ approximate_tokens(output["reasoning_delta"]),
545
+ )
546
+ # Add output tokens to rate limiter if configured
547
+ if limiter:
548
+ limiter.add(output=approximate_tokens(output["reasoning_delta"]))
549
+ # collect response delta and call callbacks
550
+ if output["response_delta"]:
551
+ if response_callback:
552
+ stop_response = await response_callback(
553
+ output["response_delta"], result.response
554
+ )
555
+ if tokens_callback:
556
+ await tokens_callback(
557
+ output["response_delta"],
558
+ approximate_tokens(output["response_delta"]),
559
+ )
560
+ # Add output tokens to rate limiter if configured
561
+ if limiter:
562
+ limiter.add(output=approximate_tokens(output["response_delta"]))
563
+ if stop_response is not None:
564
+ result.response = stop_response
565
+ break
566
+ finally:
567
+ if stop_response is not None and hasattr(_completion, "aclose"):
568
+ await _completion.aclose() # type: ignore[attr-defined]
569
570
# non-stream response
571
else:
tests/test_dirty_json.py
new
+56
@@ -0,0 +1,56 @@
1
+from __future__ import annotations
2
+
3
+import sys
4
+from pathlib import Path
5
+
6
+import pytest
7
+
8
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
9
+if str(PROJECT_ROOT) not in sys.path:
10
+ sys.path.insert(0, str(PROJECT_ROOT))
11
+
12
+from helpers.dirty_json import DirtyJson
13
+
14
+
15
+@pytest.mark.parametrize(
16
+ ("payload", "expected"),
17
+ [
18
+ (
19
+ '{"tool_name":"x","tool_args":{}}',
20
+ {"tool_name": "x", "tool_args": {}},
21
+ ),
22
+ ("[1, 2, 3]", [1, 2, 3]),
23
+ ],
24
+)
25
+def test_completed_true_when_root_is_explicitly_closed(payload, expected) -> None:
26
+ parser = DirtyJson()
27
+
28
+ assert parser.parse(payload) == expected
29
+ assert parser.completed is True
30
+
31
+
32
+def test_completed_false_when_root_hits_eof_before_closing() -> None:
33
+ parser = DirtyJson()
34
+
35
+ assert parser.parse('{"tool_name":"x","tool_args":{}') == {
36
+ "tool_name": "x",
37
+ "tool_args": {},
38
+ }
39
+ assert parser.completed is False
40
+
41
+
42
+def test_completed_remains_true_after_trailing_content() -> None:
43
+ parser = DirtyJson()
44
+
45
+ assert parser.feed('{"tool_name":"x","tool_args":{}}') == {
46
+ "tool_name": "x",
47
+ "tool_args": {},
48
+ }
49
+ assert parser.completed is True
50
+
51
+ assert parser.feed(" trailing noise") == {
52
+ "tool_name": "x",
53
+ "tool_args": {},
54
+ }
55
+
56
+ assert parser.completed is True
tests/test_stream_tool_early_stop.py
new
+96
@@ -0,0 +1,96 @@
1
+import sys
2
+from pathlib import Path
3
+
4
+import pytest
5
+
6
+
7
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
8
+if str(PROJECT_ROOT) not in sys.path:
9
+ sys.path.insert(0, str(PROJECT_ROOT))
10
+
11
+import models
12
+from helpers import extract_tools
13
+
14
+
15
+def _chunk(content: str) -> dict:
16
+ return {"choices": [{"delta": {"content": content}, "message": {}}]}
17
+
18
+
19
+class _AsyncChunkStream:
20
+ def __init__(self, chunks: list[dict]):
21
+ self._chunks = chunks
22
+ self.index = 0
23
+
24
+ def __aiter__(self):
25
+ return self
26
+
27
+ async def __anext__(self):
28
+ if self.index >= len(self._chunks):
29
+ raise StopAsyncIteration
30
+ chunk = self._chunks[self.index]
31
+ self.index += 1
32
+ return chunk
33
+
34
+
35
+def test_extract_json_root_string_returns_canonical_snapshot():
36
+ text = (
37
+ 'prefix {"tool_name":"response","tool_args":{"text":"brace } inside"}} '
38
+ "trailing noise"
39
+ )
40
+
41
+ root = extract_tools.extract_json_root_string(text)
42
+
43
+ assert root == '{"tool_name":"response","tool_args":{"text":"brace } inside"}}'
44
+ assert extract_tools.json_parse_dirty(root)["tool_args"]["text"] == "brace } inside"
45
+ assert extract_tools.extract_json_root_string(
46
+ '{"tool_name":"response","tool_args":{"text":"missing"'
47
+ ) is None
48
+ assert extract_tools.extract_json_root_string('[{"tool_name":"response"}]') is None
49
+
50
+
51
+@pytest.mark.asyncio
52
+async def test_unified_call_stops_after_canonical_root_snapshot(monkeypatch):
53
+ stream = _AsyncChunkStream(
54
+ [
55
+ _chunk(
56
+ '{"tool_name":"response","tool_args":{"text":"hello"}} trailing text'
57
+ ),
58
+ _chunk(" unreachable"),
59
+ ]
60
+ )
61
+
62
+ async def fake_acompletion(*args, **kwargs):
63
+ assert kwargs["stream"] is True
64
+ return stream
65
+
66
+ async def fake_rate_limiter(*args, **kwargs):
67
+ return None
68
+
69
+ monkeypatch.setattr(models, "acompletion", fake_acompletion)
70
+ monkeypatch.setattr(models, "apply_rate_limiter", fake_rate_limiter)
71
+
72
+ wrapper = models.LiteLLMChatWrapper(
73
+ model="test-model",
74
+ provider="openai",
75
+ model_config=None,
76
+ )
77
+
78
+ seen: list[tuple[str, str]] = []
79
+
80
+ async def response_callback(chunk: str, full: str):
81
+ seen.append((chunk, full))
82
+ snapshot = extract_tools.extract_json_root_string(full)
83
+ if snapshot:
84
+ return snapshot
85
+ return None
86
+
87
+ response, reasoning = await wrapper.unified_call(
88
+ messages=[],
89
+ response_callback=response_callback,
90
+ )
91
+
92
+ assert response == '{"tool_name":"response","tool_args":{"text":"hello"}}'
93
+ assert reasoning == ""
94
+ assert stream.index == 1
95
+ assert len(seen) == 1
96
+ assert seen[0][1] == '{"tool_name":"response","tool_args":{"text":"hello"}} trailing text'