refactor(async-tool-messages): unify on AsyncToolMessagePayload, fix JSON shape on the wire

Reshape the helper module so AsyncToolMessagePayload is the canonical
in-memory form and the on-the-wire JSON is always derived from it
(never stored). This eliminates a drift risk that came with caching
the JSON in raw_content, and it lets prepare_message_payload_for_realtime
edit the payload (graft the re-invocation reminder into 'description')
and then serialize cleanly — which fixes a 'Tool Response parsing error'
from AWS Nova Sonic that was caused by wrapping the JSON with extra
prose.

Other changes:

- Builders construct an AsyncToolMessagePayload internally and convert
  via shared private _payload_to_message and _payload_to_json helpers
  (centralizing field-omission rules, e.g. no 'result' on 'started').
- prepare_message_payload_for_realtime replaces format_text_for_provider,
  dispatching to per-kind helpers. Reminder is now appended after the
  canonical description so the model reads the protocol explanation
  first and the directive flows from it.
- Final-result payloads are pass-through; the task is done at that
  point and re-invocation is no longer a mistake.
- Stream-tool example: lengthen intermediate sleeps 10s → 20s for more
  interesting empirical testing.
This commit is contained in:
Paul Kompfner
2026-05-06 13:55:56 -04:00
parent 7d3726a74b
commit 1bb0dc1d4f
4 changed files with 238 additions and 76 deletions

View File

@@ -58,20 +58,6 @@ class TestParseMessage(unittest.TestCase):
assert info.status == "finished"
assert info.result == '"done"'
def test_raw_content_preserves_original_payload(self):
# raw_content should round-trip the source message's `content` field so
# services can forward the full payload to providers.
msg = _final_message("abc", '"done"')
info = async_tool_messages.parse_message(msg)
assert info is not None
assert info.raw_content == msg["content"]
# Sanity: it should parse back to the original payload dict.
payload = json.loads(info.raw_content)
assert payload["type"] == "async_tool"
assert payload["tool_call_id"] == "abc"
assert payload["status"] == "finished"
assert payload["result"] == '"done"'
def test_parses_completed_sentinel_result(self):
# When a function returns no value, the aggregator sets the result to
# the literal "COMPLETED" — same convention used for synchronous tool
@@ -226,7 +212,6 @@ class TestBuilders(unittest.TestCase):
assert info.tool_call_id == "call_x"
assert info.status == "running"
assert info.result is None
assert info.raw_content == msg["content"]
def test_intermediate_round_trip(self):
msg = async_tool_messages.build_intermediate_result_message("call_x", '{"step": 1}')
@@ -247,5 +232,55 @@ class TestBuilders(unittest.TestCase):
assert info.result == '{"answer": 42}'
class TestPrepareMessagePayloadForRealtime(unittest.TestCase):
def test_started_grafts_reminder_into_description(self):
msg = async_tool_messages.build_started_message("call_42")
info = async_tool_messages.parse_message(msg)
assert info is not None
text = async_tool_messages.prepare_message_payload_for_realtime(info)
# The output is well-formed JSON (the formal tool-result channel
# requires it).
decoded = json.loads(text)
# The reminder lives inside the description field, not outside the
# JSON envelope.
assert "do not call the same tool" in decoded["description"]
assert "duplicate task" in decoded["description"]
# And the original description text is still present after the reminder.
assert "asynchronous task" in decoded["description"]
# Other payload fields are preserved.
assert decoded["type"] == "async_tool"
assert decoded["tool_call_id"] == "call_42"
assert decoded["status"] == "running"
# Started payloads have no result field.
assert "result" not in decoded
def test_intermediate_grafts_reminder_into_description(self):
msg = async_tool_messages.build_intermediate_result_message("call_42", '"step-1"')
info = async_tool_messages.parse_message(msg)
assert info is not None
text = async_tool_messages.prepare_message_payload_for_realtime(info)
decoded = json.loads(text)
assert "do not call the same tool" in decoded["description"]
assert decoded["type"] == "async_tool"
assert decoded["tool_call_id"] == "call_42"
assert decoded["status"] == "running"
assert decoded["result"] == '"step-1"'
def test_final_is_pass_through(self):
# The task is done at this point; the re-invocation reminder no
# longer applies, so the final payload is forwarded as-is (no
# reminder grafted onto the description).
msg = async_tool_messages.build_final_result_message("call_42", '"the answer"')
info = async_tool_messages.parse_message(msg)
assert info is not None
text = async_tool_messages.prepare_message_payload_for_realtime(info)
decoded = json.loads(text)
assert "do not call the same tool" not in decoded["description"]
assert decoded["type"] == "async_tool"
assert decoded["tool_call_id"] == "call_42"
assert decoded["status"] == "finished"
assert decoded["result"] == '"the answer"'
if __name__ == "__main__":
unittest.main()