Agent Zero AI framework
Enable A2A streaming capability by default
Wrap FastA2A agent card generation so Agent Zero advertises streaming support on A2A endpoints by default. Add focused regression coverage for the agent-card capability rewrite and proxy wiring, and update the helper DOX profile for the new wrapper.
Alessandro committed
Jun 12, 2026 at 03:10 UTC
b8c390f50ddde6e61caf079d967db73d2d5596c3
3 files changed
+134
-4
helpers/fasta2a_server.py
+24
-1
@@ -2,11 +2,13 @@
2
import asyncio
3
import uuid
4
import atexit
5
+import json
6
from typing import Any, List
7
import contextlib
8
import threading
9
10
from helpers import settings, projects
11
+from starlette.responses import Response as StarletteResponse
12
from starlette.requests import Request
13
14
# Local imports
@@ -61,6 +63,27 @@ except ImportError: # pragma: no cover – library not installed
63
_PRINTER = PrintStyle(italic=True, font_color="purple", padding=False)
64
65
66
+def _enable_streaming_capability(agent_card_body: bytes) -> bytes:
67
+ """Return an agent-card JSON body with A2A streaming enabled."""
68
+ agent_card = json.loads(agent_card_body)
69
+ capabilities = agent_card.get("capabilities")
70
+ if not isinstance(capabilities, dict):
71
+ capabilities = {}
72
+ agent_card["capabilities"] = capabilities
73
+ capabilities["streaming"] = True
74
+ return json.dumps(agent_card, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
75
+
76
+
77
+class AgentZeroFastA2A(FastA2A): # type: ignore[misc]
78
+ """FastA2A app with Agent Zero defaults layered over library defaults."""
79
+
80
+ async def _agent_card_endpoint(self, request: Request) -> StarletteResponse:
81
+ response = await super()._agent_card_endpoint(request)
82
+ body = _enable_streaming_capability(response.body)
83
+ self._agent_card_json_schema = body
84
+ return StarletteResponse(content=body, media_type="application/json")
85
+
86
+
87
class AgentZeroWorker(Worker): # type: ignore[misc]
88
"""Agent Zero implementation of FastA2A Worker."""
89
@@ -242,7 +265,7 @@ class DynamicA2AProxy:
265
}
266
267
# Create new FastA2A app with proper thread safety
245
- new_app = FastA2A( # type: ignore
268
+ new_app = AgentZeroFastA2A( # type: ignore
269
storage=storage,
270
broker=broker,
271
name="Agent Zero",
helpers/fasta2a_server.py.dox.md
+8
-3
@@ -11,6 +11,8 @@
11
- `fasta2a_server.py` owns the runtime implementation.
12
- `fasta2a_server.py.dox.md` owns durable notes about responsibilities, contracts, side effects, and verification for that implementation.
13
- Classes:
14
+- `AgentZeroFastA2A` (`FastA2A`)
15
+ - `async _agent_card_endpoint(self, request: Request) -> StarletteResponse`
16
- `AgentZeroWorker` (`Worker`)
17
- `async run_task(self, params: Any) -> None`
18
- `async cancel_task(self, params: Any) -> None`
@@ -20,6 +22,7 @@
22
- `get_instance()`
23
- `reconfigure(self, token: str)`
24
- Top-level functions:
25
+- `_enable_streaming_capability(agent_card_body: bytes) -> bytes`: Ensure the generated A2A agent card advertises streaming support.
26
- `is_available()`: Check if FastA2A is available and properly configured.
27
- `get_proxy()`: Get the FastA2A proxy instance.
28
- Notable constants/configuration names: `_PRINTER`.
@@ -28,12 +31,13 @@
31
32
- Helper modules own reusable framework APIs and must preserve public callers unless all callers, tests, and docs are updated together.
33
- Update this file whenever public functions, classes, persistence behavior, path/security assumptions, side effects, or cross-module contracts change.
34
+- Agent Zero wraps `FastA2A` with `AgentZeroFastA2A` so the generated agent card sets `capabilities.streaming` to `true` by default.
35
- Observed side-effect areas: filesystem writes, filesystem deletion, settings/state persistence, secret handling, scheduler state.
32
-- Imported dependency areas include: `agent`, `asyncio`, `atexit`, `contextlib`, `helpers`, `helpers.persist_chat`, `helpers.print_style`, `initialize`, `starlette.requests`, `threading`, `typing`, `uuid`.
36
+- Imported dependency areas include: `agent`, `asyncio`, `atexit`, `contextlib`, `helpers`, `helpers.persist_chat`, `helpers.print_style`, `initialize`, `json`, `starlette.requests`, `starlette.responses`, `threading`, `typing`, `uuid`.
37
38
## Key Concepts
39
36
-- Important called helpers/classes observed in the source: `PrintStyle`, `DynamicA2AProxy.get_instance`, `super.__init__`, `join`, `UserMessage`, `threading.Lock`, `atexit.register`, `self._configure`, `settings.get_settings`, `path.startswith`, `self._convert_message`, `initialize_agent`, `AgentContext`, `context.log.log`, `context.communicate`, `context.reset`, `AgentContext.remove`, `remove_chat`, `self.storage.update_task`, `self._register_shutdown`.
40
+- Important called helpers/classes observed in the source: `PrintStyle`, `DynamicA2AProxy.get_instance`, `json.loads`, `json.dumps`, `super.__init__`, `super._agent_card_endpoint`, `join`, `UserMessage`, `threading.Lock`, `atexit.register`, `self._configure`, `settings.get_settings`, `path.startswith`, `self._convert_message`, `initialize_agent`, `AgentContext`, `context.log.log`, `context.communicate`, `context.reset`, `AgentContext.remove`, `remove_chat`, `self.storage.update_task`, `self._register_shutdown`.
41
- Keep request/response, tool, or helper semantics documented here at the same time as source changes.
42
43
## Work Guidance
@@ -45,7 +49,8 @@
49
## Verification
50
51
- Run targeted tests for changed helper behavior; run security regressions for auth, filesystem, WebSocket, tunnel, upload, or secret-handling helpers.
48
-- No direct test reference was found by name search; choose the nearest behavioral test or perform a focused smoke check.
52
+- Related tests:
53
+ - `tests/test_fasta2a_server.py`
54
55
## Child DOX Index
56
tests/test_fasta2a_server.py
new
+102
@@ -0,0 +1,102 @@
1
+from __future__ import annotations
2
+
3
+import importlib
4
+import json
5
+import sys
6
+import types
7
+from pathlib import Path
8
+
9
+
10
+ROOT = Path(__file__).resolve().parents[1]
11
+if str(ROOT) not in sys.path:
12
+ sys.path.insert(0, str(ROOT))
13
+
14
+
15
+def _load_fasta2a_server(monkeypatch):
16
+ settings_stub = types.ModuleType("helpers.settings")
17
+ settings_stub.get_settings = lambda: {
18
+ "a2a_server_enabled": True,
19
+ "mcp_server_token": "test-token",
20
+ }
21
+ monkeypatch.setitem(sys.modules, "helpers.settings", settings_stub)
22
+
23
+ projects_stub = types.ModuleType("helpers.projects")
24
+ projects_stub.activate_project = lambda *args, **kwargs: None
25
+ monkeypatch.setitem(sys.modules, "helpers.projects", projects_stub)
26
+
27
+ print_style_stub = types.ModuleType("helpers.print_style")
28
+
29
+ class _PrintStyle:
30
+ def __init__(self, *args, **kwargs):
31
+ pass
32
+
33
+ def print(self, *args, **kwargs):
34
+ pass
35
+
36
+ print_style_stub.PrintStyle = _PrintStyle
37
+ monkeypatch.setitem(sys.modules, "helpers.print_style", print_style_stub)
38
+
39
+ starlette_stub = types.ModuleType("starlette")
40
+ starlette_responses_stub = types.ModuleType("starlette.responses")
41
+
42
+ class _Response:
43
+ def __init__(self, content=b"", media_type=None, *args, **kwargs):
44
+ self.body = content if isinstance(content, bytes) else str(content).encode()
45
+ self.media_type = media_type
46
+
47
+ starlette_responses_stub.Response = _Response
48
+ starlette_requests_stub = types.ModuleType("starlette.requests")
49
+ starlette_requests_stub.Request = object
50
+ monkeypatch.setitem(sys.modules, "starlette", starlette_stub)
51
+ monkeypatch.setitem(sys.modules, "starlette.responses", starlette_responses_stub)
52
+ monkeypatch.setitem(sys.modules, "starlette.requests", starlette_requests_stub)
53
+
54
+ agent_stub = types.ModuleType("agent")
55
+ agent_stub.AgentContext = type(
56
+ "AgentContext",
57
+ (),
58
+ {"remove": staticmethod(lambda *args, **kwargs: None)},
59
+ )
60
+ agent_stub.UserMessage = lambda **kwargs: types.SimpleNamespace(**kwargs)
61
+ agent_stub.AgentContextType = types.SimpleNamespace(BACKGROUND="background")
62
+ monkeypatch.setitem(sys.modules, "agent", agent_stub)
63
+
64
+ initialize_stub = types.ModuleType("initialize")
65
+ initialize_stub.initialize_agent = lambda: {}
66
+ monkeypatch.setitem(sys.modules, "initialize", initialize_stub)
67
+
68
+ persist_chat_stub = types.ModuleType("helpers.persist_chat")
69
+ persist_chat_stub.remove_chat = lambda *args, **kwargs: None
70
+ monkeypatch.setitem(sys.modules, "helpers.persist_chat", persist_chat_stub)
71
+
72
+ sys.modules.pop("helpers.fasta2a_server", None)
73
+ return importlib.import_module("helpers.fasta2a_server")
74
+
75
+
76
+def test_a2a_agent_card_streaming_capability_is_enabled_by_default(monkeypatch):
77
+ module = _load_fasta2a_server(monkeypatch)
78
+
79
+ updated = module._enable_streaming_capability(
80
+ b'{"name":"Agent Zero","capabilities":{"streaming":false,"pushNotifications":false}}'
81
+ )
82
+
83
+ agent_card = json.loads(updated)
84
+ assert agent_card["capabilities"]["streaming"] is True
85
+ assert agent_card["capabilities"]["pushNotifications"] is False
86
+
87
+
88
+def test_a2a_agent_card_streaming_capability_creates_missing_block(monkeypatch):
89
+ module = _load_fasta2a_server(monkeypatch)
90
+
91
+ updated = module._enable_streaming_capability(b'{"name":"Agent Zero"}')
92
+
93
+ assert json.loads(updated)["capabilities"] == {"streaming": True}
94
+
95
+
96
+def test_a2a_proxy_uses_streaming_enabled_fast_a2a_wrapper(monkeypatch):
97
+ module = _load_fasta2a_server(monkeypatch)
98
+ proxy = object.__new__(module.DynamicA2AProxy)
99
+
100
+ proxy._configure()
101
+
102
+ assert isinstance(proxy.app, module.AgentZeroFastA2A)