test/a2a/test_events_mapping.py
Rama main · réplica de lectura del repositorio original.
# Copyright (c) 2026, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0
import json
from datetime import UTC, datetime
import pytest
from a2a.types import (
Artifact,
Message,
Part,
Role,
StreamResponse,
Task,
TaskArtifactUpdateEvent,
TaskState,
TaskStatus,
TaskStatusUpdateEvent,
)
from ag2.a2a.events import (
A2AMessage,
A2ATaskArtifactUpdate,
A2ATaskSnapshot,
A2ATaskStatusUpdate,
A2ATextArtifact,
A2AToolCallArtifact,
)
from ag2.a2a.extension import MIME_TOOL_CALL
from ag2.a2a.mappers.events import (
a2a_event_to_sdk,
chunk_to_text_artifact,
client_call_to_artifact,
parse_artifact_update,
parse_stream_response,
parse_task_artifact,
task_state_to_status_update,
)
from ag2.a2a.mappers.parts import data_part, struct_to_dict
from ag2.a2a.mappers.tools import call_to_payload
from ag2.events import ClientToolCallEvent, ModelMessageChunk, ToolCallEvent
class TestAg2ToA2A:
def test_chunk_to_text_artifact(self) -> None:
ev = chunk_to_text_artifact(
ModelMessageChunk("hello"),
artifact_id="art-1",
task_id="t-1",
context_id="c-1",
)
assert ev == A2ATextArtifact(
update=TaskArtifactUpdateEvent(
task_id="t-1",
context_id="c-1",
artifact=Artifact(artifact_id="art-1", parts=[Part(text="hello")]),
append=True,
last_chunk=False,
),
append=True,
last_chunk=False,
text="hello",
)
def test_client_call_to_artifact(self) -> None:
call = ToolCallEvent(id="call-1", name="lookup", arguments='{"q":"x"}')
ev = client_call_to_artifact(
ClientToolCallEvent(id="call-1", name="lookup", arguments='{"q":"x"}'),
task_id="t-1",
context_id="c-1",
)
assert ev == A2AToolCallArtifact(
update=TaskArtifactUpdateEvent(
task_id="t-1",
context_id="c-1",
artifact=Artifact(
artifact_id="call-1",
name="lookup",
parts=[data_part(call_to_payload(call), media_type=MIME_TOOL_CALL)],
),
append=False,
last_chunk=True,
),
append=False,
last_chunk=True,
call=call,
)
def test_task_state_to_status_update(self) -> None:
ev = task_state_to_status_update(
TaskState.TASK_STATE_COMPLETED,
task_id="t-1",
context_id="c-1",
)
assert ev == A2ATaskStatusUpdate(
update=TaskStatusUpdateEvent(
task_id="t-1",
context_id="c-1",
status=TaskStatus(state=TaskState.TASK_STATE_COMPLETED),
),
state=TaskState.TASK_STATE_COMPLETED,
)
class TestA2AToSdk:
def test_artifact_update_unwrap(self) -> None:
ev = chunk_to_text_artifact(
ModelMessageChunk("hi"),
artifact_id="a",
task_id="t",
context_id="c",
)
assert a2a_event_to_sdk(ev) is ev.update
def test_status_update_unwrap(self) -> None:
ev = task_state_to_status_update(TaskState.TASK_STATE_FAILED, task_id="t", context_id="c")
proto = a2a_event_to_sdk(ev)
assert proto == TaskStatusUpdateEvent(
task_id="t",
context_id="c",
status=TaskStatus(state=TaskState.TASK_STATE_FAILED),
)
def test_message_unwrap(self) -> None:
msg = Message(message_id="m-1", role=Role.ROLE_AGENT, parts=[Part(text="ok")])
assert a2a_event_to_sdk(A2AMessage(message=msg)) is msg
def test_task_snapshot_unwrap(self) -> None:
task = Task(id="t-1", context_id="c-1", status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED))
assert a2a_event_to_sdk(A2ATaskSnapshot(task=task)) is task
class TestSdkToA2A:
def test_text_artifact(self) -> None:
update = TaskArtifactUpdateEvent(
task_id="t",
context_id="c",
artifact=Artifact(artifact_id="a", parts=[Part(text="streamed")]),
append=True,
last_chunk=False,
)
ev = parse_stream_response(StreamResponse(artifact_update=update))
assert ev == A2ATextArtifact(update=update, append=True, last_chunk=False, text="streamed")
@pytest.mark.parametrize(
("parts", "text"),
[
([Part(text="")], ""),
([Part(text=""), Part(text="body")], "body"),
([Part(text="body"), Part(text="")], "body"),
([Part(text=""), Part(text="body"), Part(text="")], "body"),
],
)
def test_empty_text_parts_are_still_text_artifacts(self, parts: list[Part], text: str) -> None:
update = TaskArtifactUpdateEvent(
task_id="t",
context_id="c",
artifact=Artifact(artifact_id="a", parts=parts),
)
ev = parse_stream_response(StreamResponse(artifact_update=update))
assert ev == A2ATextArtifact(update=update, append=False, last_chunk=False, text=text)
def test_tool_call_artifact(self) -> None:
payload = {"id": "c-1", "name": "lookup", "arguments": json.dumps({"q": "x"})}
update = TaskArtifactUpdateEvent(
task_id="t",
context_id="c",
artifact=Artifact(
artifact_id="c-1",
parts=[data_part(payload, media_type=MIME_TOOL_CALL)],
),
append=False,
last_chunk=True,
)
ev = parse_stream_response(StreamResponse(artifact_update=update))
assert ev == A2AToolCallArtifact(
update=update,
append=False,
last_chunk=True,
call=ToolCallEvent(id="c-1", name="lookup", arguments=json.dumps({"q": "x"})),
)
def test_status_update(self) -> None:
update = TaskStatusUpdateEvent(
task_id="t",
context_id="c",
status=TaskStatus(state=TaskState.TASK_STATE_INPUT_REQUIRED),
)
ev = parse_stream_response(StreamResponse(status_update=update))
assert ev == A2ATaskStatusUpdate(update=update, state=TaskState.TASK_STATE_INPUT_REQUIRED)
def test_task_snapshot(self) -> None:
task = Task(id="t", context_id="c", status=TaskStatus(state=TaskState.TASK_STATE_WORKING))
ev = parse_stream_response(StreamResponse(task=task))
assert ev == A2ATaskSnapshot(task=task)
def test_message(self) -> None:
msg = Message(message_id="m", role=Role.ROLE_AGENT, parts=[Part(text="reply")])
ev = parse_stream_response(StreamResponse(message=msg))
assert ev == A2AMessage(message=msg)
def test_unknown_artifact_falls_back_to_base(self) -> None:
# Mixed text + opaque data: not text-only, not single tool-call+json — falls
# through to plain A2ATaskArtifactUpdate so unknown extensions still surface.
update = TaskArtifactUpdateEvent(
task_id="t",
context_id="c",
artifact=Artifact(
artifact_id="a",
parts=[
Part(text="head"),
data_part({"x": 1}, media_type="application/vnd.example+json"),
],
),
)
ev = parse_stream_response(StreamResponse(artifact_update=update))
assert type(ev) is A2ATaskArtifactUpdate
assert ev == A2ATaskArtifactUpdate(update=update, append=False, last_chunk=False)
class TestParseTaskArtifact:
def test_text_artifact_from_polling_snapshot(self) -> None:
artifact = Artifact(artifact_id="a", parts=[Part(text="final text")])
ev = parse_task_artifact(artifact, task_id="t", context_id="c")
# Polling treats every artifact as final, non-appended.
assert ev == A2ATextArtifact(
update=TaskArtifactUpdateEvent(
task_id="t",
context_id="c",
artifact=artifact,
append=False,
last_chunk=False,
),
append=False,
last_chunk=False,
text="final text",
)
def test_tool_call_artifact_from_polling_snapshot(self) -> None:
payload = {"id": "c-1", "name": "f", "arguments": "{}"}
artifact = Artifact(
artifact_id="c-1",
parts=[data_part(payload, media_type=MIME_TOOL_CALL)],
)
ev = parse_task_artifact(artifact, task_id="t", context_id="c")
assert ev == A2AToolCallArtifact(
update=TaskArtifactUpdateEvent(
task_id="t",
context_id="c",
artifact=artifact,
append=False,
last_chunk=False,
),
append=False,
last_chunk=False,
call=ToolCallEvent(id="c-1", name="f", arguments="{}"),
)
def test_routes_through_same_classifier_as_streaming(self) -> None:
# Confirms both transport paths converge on the same typed view.
artifact = Artifact(artifact_id="a", parts=[Part(text="x")])
polled = parse_task_artifact(artifact, task_id="t", context_id="c")
streamed = parse_artifact_update(polled.update)
assert type(polled) is type(streamed)
def test_chunk_artifact_roundtrip() -> None:
original = chunk_to_text_artifact(
ModelMessageChunk("hello"),
artifact_id="a-1",
task_id="t-1",
context_id="c-1",
append=True,
last_chunk=True,
)
reparsed = parse_stream_response(StreamResponse(artifact_update=a2a_event_to_sdk(original)))
assert reparsed == original
class TestWireMetadata:
def test_status_update_timestamp_propagates(self) -> None:
ts = datetime(2026, 5, 11, 12, 30, 45, tzinfo=UTC)
ev = task_state_to_status_update(
TaskState.TASK_STATE_COMPLETED,
task_id="t",
context_id="c",
timestamp=ts,
)
assert ev.update.status.timestamp.seconds == int(ts.timestamp())
def test_status_update_omits_timestamp_by_default(self) -> None:
ev = task_state_to_status_update(TaskState.TASK_STATE_WORKING, task_id="t", context_id="c")
assert ev.update.status.timestamp.seconds == 0
def test_chunk_artifact_optional_name_and_description(self) -> None:
ev = chunk_to_text_artifact(
ModelMessageChunk("hi"),
artifact_id="a",
task_id="t",
context_id="c",
name="streamed-output",
description="model text stream",
)
assert ev.update.artifact.name == "streamed-output"
assert ev.update.artifact.description == "model text stream"
def test_chunk_artifact_metadata_round_trips_through_struct(self) -> None:
ev = chunk_to_text_artifact(
ModelMessageChunk("hi"),
artifact_id="a",
task_id="t",
context_id="c",
artifact_metadata={"source": "llm", "chunk_index": 3},
)
assert struct_to_dict(ev.update.artifact.metadata) == {"source": "llm", "chunk_index": 3}
def test_client_call_artifact_uses_tool_name_by_default(self) -> None:
ev = client_call_to_artifact(
ClientToolCallEvent(id="c-1", name="search", arguments="{}"),
task_id="t",
context_id="c",
)
assert ev.update.artifact.name == "search"
def test_client_call_artifact_name_override(self) -> None:
ev = client_call_to_artifact(
ClientToolCallEvent(id="c-1", name="search", arguments="{}"),
task_id="t",
context_id="c",
name="custom-label",
)
assert ev.update.artifact.name == "custom-label"