LCORE-2914: refactor A2A executor to use pydantic-ai agents#2092
LCORE-2914: refactor A2A executor to use pydantic-ai agents#2092are-ces wants to merge 1 commit into
Conversation
WalkthroughThe A2A executor now streams from a Pydantic-AI agent, converts agent events into A2A task status and artifact updates, and uses agent results to build final parts. Error handling now includes ChangesA2A agent streaming migration
Estimated code review effort: 4 (Complex) | ~60 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
✨ Simplify code
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
@coderabbitai full review |
✅ Action performedFull review finished. |
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/app/endpoints/a2a.py`:
- Around line 119-121: The new helper signatures and docstrings in the A2A
endpoint module need to match repo conventions. Update the affected helpers such
as the ones around the text-to-parts conversion and related endpoint helpers to
use modern union syntax instead of Optional[...] (for example,
AgentRunResult[str] | None and similar) and replace any docstring section
headers from Args: to Parameters:. Keep the changes consistent across the
referenced helpers so the annotations and Google-style docstrings all follow the
same convention.
- Around line 546-578: The _dispatch_agent_event() method is mutating the
caller-owned text_parts list via append, which violates the no in-place
parameter modification guideline. Refactor _dispatch_agent_event() so it only
derives the event payload and does not update text_parts directly; instead,
return the text delta or updated text alongside the TaskStatusUpdateEvent and
let _convert_stream_to_events() handle accumulation. Keep the change localized
around _dispatch_agent_event(), PartStartEvent, and PartDeltaEvent handling.
In `@tests/unit/app/endpoints/test_a2a.py`:
- Around line 885-892: The stream mock setup in the test is attaching a final
result to mock_stream.result, but the production path reads the final output
from AgentRunResultEvent.result instead. Update the test around mock_run_result
and mock_stream to yield an AgentRunResultEvent through the async iterator so
the final artifact/result path is actually exercised, and remove the unused
mock_stream.result assignment.
- Around line 740-745: The test name and setup in
test_process_task_streaming_handles_agent_run_error do not match the injected
failure, since it currently triggers APIConnectionError instead of the
AgentRunError path. Update this test to either explicitly inject and assert the
AgentRunError mapping in _process_task_streaming, or rename the test to reflect
the APIConnectionError behavior; if needed, add a separate case so the
AgentRunError branch remains covered.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: 06d76901-367a-4e71-885f-702e1e288fd9
📒 Files selected for processing (2)
src/app/endpoints/a2a.pytests/unit/app/endpoints/test_a2a.py
📜 Review details
⚠️ CI failures not shown inline (1)
GitHub Actions: Integration tests / 1_integration_tests (3.13).txt: LCORE-2914: refactor A2A executor to use pydantic-ai agents
Conclusion: failure
dict])
v = handler(item, index)
tests/integration/endpoints/test_responses_integration.py::test_streaming_blocked_returns_sse_and_persists_turn
/home/runner/work/lightspeed-stack/lightspeed-stack/.venv/lib/python3.13/site-packages/pydantic/main.py:475: UserWarning: Pydantic serializer warnings:
PydanticSerializationUnexpectedValue(Expected `OpenAIResponseToolMCP` - serialized value may not be as expected [field_name='tools', input_value={'type': 'mcp', 'server_l..., 'allowed_tools': None}, input_type=dict])
PydanticSerializationUnexpectedValue(Expected `OpenAIResponseToolMCP` - serialized value may not be as expected [field_name='tools', input_value={'type': 'mcp', 'server_l..., 'allowed_tools': None}, input_type=dict])
PydanticSerializationUnexpectedValue(Expected `OpenAIResponseToolMCP` - serialized value may not be as expected [field_name='tools', input_value={'type': 'mcp', 'server_l..., 'allowed_tools': None}, input_type=dict])
return self.__pydantic_serializer__.to_python(
-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
================================ tests coverage ================================
_______________ coverage: platform linux, python 3.13.14-final-0 _______________
Name Stmts Miss Cover Missing
--------------------------------------------------------------------------------------------------------
src/__init__.py 0 0 100%
src/a2a_storage/__init__.py 6 0 100%
src/a2a_storage/context_store.py 13 0 100%
src/a2a_storage/in_memory_context_store.py 32 20 38% 25-28, 39-47, 56-58, 70-75, 85, 93
src/a2a_storage/postgres_context_store.py 54 37 31% 47-51,...
🧰 Additional context used
📓 Path-based instructions (3)
tests/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
tests/**/*.py: Use pytest for all unit and integration tests; do not use unittest
Usepytest.mark.asynciomarker for async tests
Files:
tests/unit/app/endpoints/test_a2a.py
src/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
src/**/*.py: Use absolute imports for internal modules:from authentication import get_auth_dependency
Llama Stack imports: Usefrom llama_stack_client import AsyncLlamaStackClient
Checkconstants.pyfor shared constants before defining new ones
All modules must start with descriptive docstrings explaining purpose
Uselogger = get_logger(__name__)fromlog.pyfor module logging
All functions must have complete type annotations for parameters and return types, use modern syntax (str | int), and include descriptive docstrings
Use snake_case with descriptive, action-oriented names for functions (get_, validate_, check_)
Avoid in-place parameter modification anti-patterns; return new data structures instead of modifying function parameters
Useasync deffor I/O operations and external API calls
Use standard log levels with clear purposes:debug()for diagnostic info,info()for program execution,warning()for unexpected events,error()for serious problems
All classes must have descriptive docstrings explaining purpose and use PascalCase with standard suffixes:Configuration,Error/Exception,Resolver,Interface
Abstract classes must use ABC with@abstractmethoddecorators
Follow Google Python docstring conventions with required sections: Parameters, Returns, Raises, and Attributes for classes
Files:
src/app/endpoints/a2a.py
src/app/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
src/app/**/*.py: FastAPI dependencies: Import fromfastapimodule forAPIRouter,HTTPException,Request,status,Depends
Use FastAPIHTTPExceptionwith appropriate status codes for API endpoints and handleAPIConnectionErrorfrom Llama Stack
Files:
src/app/endpoints/a2a.py
🧠 Learnings (3)
📚 Learning: 2026-06-24T13:45:37.249Z
Learnt from: Jdubrick
Repo: lightspeed-core/lightspeed-stack PR: 1971
File: src/utils/markdown_repair.py:31-36
Timestamp: 2026-06-24T13:45:37.249Z
Learning: In the lightspeed-stack repository, docstrings must use the section header name "Parameters:" (not "Args:") for function arguments, even if the project references Google Python docstring conventions. Ensure docstrings follow the project’s established "Parameters:" header format for any documented function parameters.
Applied to files:
tests/unit/app/endpoints/test_a2a.pysrc/app/endpoints/a2a.py
📚 Learning: 2026-04-06T20:18:07.852Z
Learnt from: major
Repo: lightspeed-core/lightspeed-stack PR: 1463
File: src/app/endpoints/rlsapi_v1.py:266-271
Timestamp: 2026-04-06T20:18:07.852Z
Learning: In the lightspeed-stack codebase, within `src/app/endpoints/` inference/MCP endpoints, treat `tools: Optional[list[Any]]` in MCP tool definitions as an intentional, consistent typing pattern (used across `query`, `responses`, `streaming_query`, `rlsapi_v1`). Do not raise or suggest this as a typing issue during code review; changing it in isolation could break endpoint typing consistency across the codebase.
Applied to files:
src/app/endpoints/a2a.py
📚 Learning: 2026-07-06T15:26:18.398Z
Learnt from: Jdubrick
Repo: lightspeed-core/lightspeed-stack PR: 2071
File: src/models/config.py:2416-2422
Timestamp: 2026-07-06T15:26:18.398Z
Learning: In this repo’s Python code under src/**, don’t treat differences in string concatenation style as a style inconsistency when Black has effectively forced (or made clearer) use of explicit `+` string concatenation in multi-line logger/string expressions. If adjacent-literal implicit concatenation is avoided/changed specifically to accommodate Black’s formatting in these call sites, accept the `+` usage and don’t recommend converting it solely for consistency with nearby blocks that use implicit concatenation.
Applied to files:
src/app/endpoints/a2a.py
🪛 ast-grep (0.44.1)
tests/unit/app/endpoints/test_a2a.py
[warning] 803-803: Do not make http calls without encryption
Context: "http://test-llama-stack/responses"
Note: [CWE-319] Cleartext Transmission of Sensitive Information.
(requests-http)
50b2a99 to
31a6cad
Compare
Replace direct client.responses.create() calls with build_agent() + agent.run_stream_events(), aligning the A2A server path with the query and streaming query endpoints. This enables pydantic-ai skills and capabilities in A2A-served requests. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
31a6cad to
29bd62d
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/app/endpoints/a2a.py`:
- Around line 489-496: The unused executor state in _convert_stream_to_events()
should be removed: self._run_result is assigned when handling
AgentRunResultEvent but never read anywhere. Update the A2A executor flow by
deleting the self._run_result assignment and removing the attribute from the
class if it exists, keeping the run_result local to the streaming logic unless
another method truly needs it.
- Around line 130-131: The `a2a` endpoint is reading the wrong field from
`AgentRunResult`, so the final text assignment will break at runtime in the
pinned `pydantic-ai` version. Update the `run_result` handling in the endpoint
logic to use `run_result.output` instead of `run_result.response.text`, and make
sure the corresponding code in `src/utils/agents/streaming.py` uses the same
`AgentRunResult` output field for consistency.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: f47ca9b2-cc6d-449a-92ec-315cf42ae90a
📒 Files selected for processing (2)
src/app/endpoints/a2a.pytests/unit/app/endpoints/test_a2a.py
📜 Review details
⏰ Context from checks skipped due to timeout. (8)
- GitHub Check: Konflux kflux-prd-rh02 / lightspeed-stack-0-7-on-pull-request
- GitHub Check: E2E: server mode / ci / group 2
- GitHub Check: E2E Tests for Lightspeed Evaluation job
- GitHub Check: E2E: library mode / ci / group 2
- GitHub Check: E2E: library mode / ci / group 1
- GitHub Check: E2E: server mode / ci / group 1
- GitHub Check: E2E: library mode / ci / group 3
- GitHub Check: E2E: server mode / ci / group 3
🧰 Additional context used
📓 Path-based instructions (3)
src/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
src/**/*.py: Use absolute imports for internal modules:from authentication import get_auth_dependency
Llama Stack imports: Usefrom llama_stack_client import AsyncLlamaStackClient
Checkconstants.pyfor shared constants before defining new ones
All modules must start with descriptive docstrings explaining purpose
Uselogger = get_logger(__name__)fromlog.pyfor module logging
All functions must have complete type annotations for parameters and return types, use modern syntax (str | int), and include descriptive docstrings
Use snake_case with descriptive, action-oriented names for functions (get_, validate_, check_)
Avoid in-place parameter modification anti-patterns; return new data structures instead of modifying function parameters
Useasync deffor I/O operations and external API calls
Use standard log levels with clear purposes:debug()for diagnostic info,info()for program execution,warning()for unexpected events,error()for serious problems
All classes must have descriptive docstrings explaining purpose and use PascalCase with standard suffixes:Configuration,Error/Exception,Resolver,Interface
Abstract classes must use ABC with@abstractmethoddecorators
Follow Google Python docstring conventions with required sections: Parameters, Returns, Raises, and Attributes for classes
Files:
src/app/endpoints/a2a.py
src/app/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
src/app/**/*.py: FastAPI dependencies: Import fromfastapimodule forAPIRouter,HTTPException,Request,status,Depends
Use FastAPIHTTPExceptionwith appropriate status codes for API endpoints and handleAPIConnectionErrorfrom Llama Stack
Files:
src/app/endpoints/a2a.py
tests/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
tests/**/*.py: Use pytest for all unit and integration tests; do not use unittest
Usepytest.mark.asynciomarker for async tests
Files:
tests/unit/app/endpoints/test_a2a.py
🧠 Learnings (3)
📚 Learning: 2026-04-06T20:18:07.852Z
Learnt from: major
Repo: lightspeed-core/lightspeed-stack PR: 1463
File: src/app/endpoints/rlsapi_v1.py:266-271
Timestamp: 2026-04-06T20:18:07.852Z
Learning: In the lightspeed-stack codebase, within `src/app/endpoints/` inference/MCP endpoints, treat `tools: Optional[list[Any]]` in MCP tool definitions as an intentional, consistent typing pattern (used across `query`, `responses`, `streaming_query`, `rlsapi_v1`). Do not raise or suggest this as a typing issue during code review; changing it in isolation could break endpoint typing consistency across the codebase.
Applied to files:
src/app/endpoints/a2a.py
📚 Learning: 2026-06-24T13:45:37.249Z
Learnt from: Jdubrick
Repo: lightspeed-core/lightspeed-stack PR: 1971
File: src/utils/markdown_repair.py:31-36
Timestamp: 2026-06-24T13:45:37.249Z
Learning: In the lightspeed-stack repository, docstrings must use the section header name "Parameters:" (not "Args:") for function arguments, even if the project references Google Python docstring conventions. Ensure docstrings follow the project’s established "Parameters:" header format for any documented function parameters.
Applied to files:
src/app/endpoints/a2a.pytests/unit/app/endpoints/test_a2a.py
📚 Learning: 2026-07-06T15:26:18.398Z
Learnt from: Jdubrick
Repo: lightspeed-core/lightspeed-stack PR: 2071
File: src/models/config.py:2416-2422
Timestamp: 2026-07-06T15:26:18.398Z
Learning: In this repo’s Python code under src/**, don’t treat differences in string concatenation style as a style inconsistency when Black has effectively forced (or made clearer) use of explicit `+` string concatenation in multi-line logger/string expressions. If adjacent-literal implicit concatenation is avoided/changed specifically to accommodate Black’s formatting in these call sites, accept the `+` usage and don’t recommend converting it solely for consistency with nearby blocks that use implicit concatenation.
Applied to files:
src/app/endpoints/a2a.py
🪛 ast-grep (0.44.1)
tests/unit/app/endpoints/test_a2a.py
[warning] 804-804: Do not make http calls without encryption
Context: "http://test-llama-stack/responses"
Note: [CWE-319] Cleartext Transmission of Sensitive Information.
(requests-http)
🔇 Additional comments (9)
src/app/endpoints/a2a.py (5)
538-548: 📐 Maintainability & Code Quality
_dispatch_agent_event()mutates the caller-ownedtext_partslist.This was previously flagged: the method appends to the
text_partsparameter, which violates the repo guideline against in-place parameter modification. Consider returning the delta alongside the event and accumulating in_convert_stream_to_events().Source: Coding guidelines
34-49: LGTM!Also applies to: 63-68
83-84: LGTM!
374-388: LGTM!Also applies to: 423-447
586-615: LGTM!tests/unit/app/endpoints/test_a2a.py (4)
26-43: LGTM!Also applies to: 155-217
741-828: LGTM!Also applies to: 830-904
906-985: LGTM!Also applies to: 987-1020
1078-1179: LGTM!
| run_result: AgentRunResult[str] | None = None | ||
|
|
||
| async with agent.run_stream_events(prompt) as stream: | ||
| async for event in stream: | ||
| if isinstance(event, AgentRunResultEvent): | ||
| run_result = event.result | ||
| self._run_result = run_result | ||
| continue |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check for any reader of self._run_result / executor reuse
rg -nP '_run_result' src tests
echo "--- how executor is instantiated (per-request vs shared) ---"
rg -nP -C3 'A2AAgentExecutor\(' srcRepository: lightspeed-core/lightspeed-stack
Length of output: 7522
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "--- surrounding A2AAgentExecutor state and usage ---"
sed -n '202,280p' src/app/endpoints/a2a.py
echo "--- _run_result references in the file ---"
rg -n '_run_result|run_result' src/app/endpoints/a2a.py
echo "--- where the A2A agent executor is created and how long it lives ---"
sed -n '740,790p' src/app/endpoints/a2a.py
echo "--- any other A2AAgentExecutor instantiation sites ---"
rg -n 'A2AAgentExecutor\(' srcRepository: lightspeed-core/lightspeed-stack
Length of output: 5677
Remove the unused executor state
self._run_result is assigned in _convert_stream_to_events() but never read. Drop the attribute and the assignment unless a later step needs per-run state.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/app/endpoints/a2a.py` around lines 489 - 496, The unused executor state
in _convert_stream_to_events() should be removed: self._run_result is assigned
when handling AgentRunResultEvent but never read anywhere. Update the A2A
executor flow by deleting the self._run_result assignment and removing the
attribute from the class if it exists, keeping the run_result local to the
streaming logic unless another method truly needs it.
asimurka
left a comment
There was a problem hiding this comment.
LGTM in overall, I pointed out some things worth checking
| _TASK_STORE: Optional[TaskStore] = None | ||
| _CONTEXT_STORE: Optional[A2AContextStore] = None | ||
| _TASK_STORE: TaskStore | None = None | ||
| _CONTEXT_STORE: A2AContextStore | None = None |
There was a problem hiding this comment.
nit: The original version with using Optional aligns with contribution guides (occurs also below sometimes)
| state=TaskState.working, | ||
| message=new_agent_text_message( | ||
| f"MCP call: {item_id}", | ||
| f"MCP call: {event.part.tool_name}", |
There was a problem hiding this comment.
There is an equivalent attribute event.part.tool_call_id
| _ = artifact_id | ||
|
|
||
| if isinstance(event, PartStartEvent): | ||
| if isinstance(event.part, PydanticTextPart): |
There was a problem hiding this comment.
This should be only TextPart there is doubled import at the top of this module.
| item_id = getattr(chunk, "item_id", "") | ||
| yield TaskStatusUpdateEvent( | ||
| elif isinstance(event, PartEndEvent): | ||
| if isinstance(event.part, NativeToolCallPart): |
There was a problem hiding this comment.
Just note: NativeToolCallPart can contain also file search, mcp tools list and web search (not only mcp call as in the original implementation)
Description
Prerequisite for A2A client support (LCORE-2914). The A2A server endpoint currently bypasses pydantic-ai entirely — it calls
client.responses.create()directly against Llama Stack. This PR refactorsA2AAgentExecutorto usebuild_agent()+agent.run_stream_events(), aligning it with the query and streaming query endpoints.This enables pydantic-ai skills and capabilities in A2A-served requests, which is required before external agent delegation tools can be registered as pydantic-ai capabilities in a follow-up PR.
Changes
client.responses.create()withbuild_agent()+agent.run_stream_events()inA2AAgentExecutor._process_task_streaming()_convert_stream_to_events()to consume pydantic-aiAgentStreamEvent/AgentRunResultEventinstead ofOpenAIResponseObjectStreamchunks_build_a2a_parts_from_agent_result()to convert agent results to A2A Parts (replaces_convert_responses_content_to_a2a_parts())_dispatch_agent_event()to map pydantic-ai events to A2ATaskStatusUpdateEventAgentRunErrorusingmap_agent_inference_error()from existing utilsType of change
Tools used to create PR
Related Tickets & Documents
Checklist before requesting a review
Testing
uv run make verify— all linters passuv run make test-unit— 2973 tests pass, 0 failuresbuild_agent()instead ofclient.responses.create()TestDispatchAgentEventandTestBuildA2APartsFromAgentResulttest classes cover the pydantic-ai event mappingSummary by CodeRabbit
New Features
Bug Fixes
Tests