fix(stream): tolerate partial response snapshots
Guard partially parsed tool arguments while streaming and render both text and legacy message response fields. Keep live response logs synchronized without allowing incomplete code arguments to raise.
Alessandro committed
Jul 25, 2026 at 00:45 UTC
87bafe10e02c36d936b634c21762799b85abfc02
4 files changed
+90
-12
extensions/python/response_stream/AGENTS.md
+2
@@ -11,6 +11,8 @@
11
## Local Contracts
12
13
- Keep streaming output synchronized with UI log items.
14
+- Treat parsed stream snapshots as partial data; nested tool fields may be `None`
15
+ until their values arrive.
16
- Preserve include-alias replacement semantics where prompts/tools rely on them.
17
- Do not expose unmasked secrets in live responses.
18
extensions/python/response_stream/_10_log_from_stream.py
+9
-7
@@ -53,20 +53,22 @@ class LogFromStream(Extension):
53
if parsed is not None and "tool_name" in parsed and parsed["tool_name"]:
54
kvps["step"] = f"Using {parsed['tool_name']}..." # using tool XY
55
if parsed["tool_name"]=="code_execution_tool":
56
- if "tool_args" in parsed and "runtime" in parsed["tool_args"]:
56
+ tool_args = parsed.get("tool_args")
57
+ if isinstance(tool_args, dict) and "runtime" in tool_args:
58
length = ""
58
- if "code" in parsed["tool_args"]:
59
- length = f"({len(parsed['tool_args']['code'])})"
59
+ code = tool_args.get("code")
60
+ if isinstance(code, str):
61
+ length = f"({len(code)})"
62
kvps["step"] = f"Writing code... {length}"
61
- if parsed["tool_args"]["runtime"] == "python":
63
+ if tool_args["runtime"] == "python":
64
kvps["step"] = f"Writing Python code... {length}"
63
- elif parsed["tool_args"]["runtime"] == "nodejs":
65
+ elif tool_args["runtime"] == "nodejs":
66
kvps["step"] = f"Writing Node.js code... {length}"
65
- elif parsed["tool_args"]["runtime"] == "terminal":
67
+ elif tool_args["runtime"] == "terminal":
68
kvps["step"] = f"Writing terminal command... {length}"
69
kvps.update(parsed)
70
71
72
73
# update the log item
72
- log_item.update(heading=heading, content=text, kvps=kvps)
\ No newline at end of file
74
+ log_item.update(heading=heading, content=text, kvps=kvps)
extensions/python/response_stream/_20_live_response.py
+9
-2
@@ -21,7 +21,14 @@ class LiveResponse(Extension):
21
22
try:
23
tool_name, tool_args = extract_tools.normalize_tool_request(parsed)
24
- if tool_name != "response" or not tool_args.get("text"):
24
+ message = tool_args.get("text")
25
+ if not isinstance(message, str) or not message.strip():
26
+ message = tool_args.get("message")
27
+ if (
28
+ tool_name != "response"
29
+ or not isinstance(message, str)
30
+ or not message.strip()
31
+ ):
32
return # not a response
33
34
# create log message and store it in loop data temporary params
@@ -39,6 +46,6 @@ class LiveResponse(Extension):
46
47
# update log message
48
log_item = loop_data.params_temporary["log_item_response"]
42
- log_item.update(content=tool_args["text"])
49
+ log_item.update(content=message)
50
except Exception as e:
51
pass
tests/test_plain_response_logging.py
+70
-3
@@ -1,11 +1,17 @@
1
-import pytest
1
from types import SimpleNamespace
2
4
-from helpers.log import Log
5
-from extensions.python.response_stream._20_live_response import LiveResponse
3
+import pytest
4
+
5
+from agent import LoopData
6
from extensions.python._functions.agent.Agent.hist_add_ai_response.end._10_log_plain_responses import (
7
LogPlainResponses,
8
)
9
+from extensions.python.response_stream._10_log_from_stream import (
10
+ LogFromStream as StreamLog,
11
+)
12
+from extensions.python.response_stream._20_live_response import LiveResponse
13
+from helpers.dirty_json import DirtyJson
14
+from helpers.log import Log
15
16
17
def _agent_with_generating_log():
@@ -105,3 +111,64 @@ async def test_live_response_renders_single_action_wrapper():
111
assert response.type == "response"
112
assert response.content == "wrapper works"
113
assert response.id == "msg-1"
114
+
115
+
116
+@pytest.mark.asyncio
117
+async def test_live_response_renders_legacy_message_argument():
118
+ log = Log()
119
+ generating = log.log(type="agent", id="msg-1")
120
+ loop_data = SimpleNamespace(params_temporary={"log_item_generating": generating})
121
+ agent = SimpleNamespace(
122
+ context=SimpleNamespace(log=log),
123
+ agent_name="A0",
124
+ )
125
+
126
+ await LiveResponse(agent=agent).execute(
127
+ loop_data=loop_data,
128
+ parsed={"tool_name": "response", "tool_args": {"message": "legacy works"}},
129
+ )
130
+
131
+ response = loop_data.params_temporary["log_item_response"]
132
+ assert response.content == "legacy works"
133
+ assert response.id == "msg-1"
134
+
135
+
136
+@pytest.mark.asyncio
137
+@pytest.mark.parametrize(
138
+ ("stream", "expected_step"),
139
+ [
140
+ (
141
+ '{"tool_name":"code_execution_tool","tool_args":',
142
+ "Using code_execution_tool...",
143
+ ),
144
+ (
145
+ '{"tool_name":"code_execution_tool","tool_args":'
146
+ '{"runtime":"python","code":',
147
+ "Writing Python code... ",
148
+ ),
149
+ (
150
+ '{"tool_name":"code_execution_tool","tool_args":'
151
+ '{"runtime":"python","code":"pri',
152
+ "Writing Python code... (3)",
153
+ ),
154
+ ],
155
+)
156
+async def test_stream_log_tolerates_partial_tool_arguments(
157
+ stream: str,
158
+ expected_step: str,
159
+):
160
+ log = Log()
161
+ loop_data = LoopData()
162
+ agent = SimpleNamespace(
163
+ context=SimpleNamespace(log=log),
164
+ agent_name="A0",
165
+ )
166
+
167
+ await StreamLog(agent=agent).execute(
168
+ loop_data=loop_data,
169
+ text=stream,
170
+ parsed=DirtyJson.parse_string(stream),
171
+ )
172
+
173
+ item = loop_data.params_temporary["log_item_generating"]
174
+ assert item.kvps["step"] == expected_step