LCORE-2914: A2A client support: enable LCS to discover and call external A2A agents#2162
LCORE-2914: A2A client support: enable LCS to discover and call external A2A agents#2162are-ces wants to merge 1 commit into
Conversation
WalkthroughAdds external A2A agent configuration, startup discovery and caching, Pydantic-AI delegation tools, lifecycle cleanup, OpenAPI schemas, tests, an example configuration, and updated A2A protocol documentation. ChangesA2A delegation
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant FastAPI lifespan
participant A2AClientManager
participant Remote A2A agent
participant Pydantic-AI agent
FastAPI lifespan->>A2AClientManager: Initialize configured agents
A2AClientManager->>Remote A2A agent: Connect and fetch AgentCard
Pydantic-AI agent->>A2AClientManager: Call list_agents or delegate_to_agent
A2AClientManager->>Remote A2A agent: Send delegated user message
Remote A2A agent-->>Pydantic-AI agent: Stream response events
Possibly related PRs
Suggested reviewers: Important Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional. ❌ Failed checks (1 error)
✅ 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 |
Enable LCS to discover and delegate tasks to external A2A-compliant agents. The delegation is exposed as a pydantic-ai capability, so any agent-powered endpoint (query, streaming, A2A) can use it. - A2AAgentsConfiguration: YAML config for remote agent endpoints with timeout, max_retries, and optional auth_token - A2AClientManager: singleton that resolves agent cards at startup and holds Client instances with httpx timeout/retry transport - A2ADelegationCapability: pydantic-ai tools (list_agents, delegate_to_agent) auto-injected into build_agent() - Startup/shutdown lifecycle wired into FastAPI lifespan - A2A protocol docs updated with client documentation - OpenAPI schema regenerated - Example config: examples/lightspeed-stack-a2a-agents.yaml Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
71f1074 to
60029d6
Compare
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| httpx_client = httpx.AsyncClient( | ||
| timeout=httpx.Timeout(agent_config.timeout), |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/app/main.py (1)
145-153: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winGuarantee A2A cleanup when another shutdown step fails.
If either preceding cleanup raises, Line 149 is skipped and remote connections remain open. Put
A2AClientManager().close()in an independentfinallyblock.🤖 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/main.py` around lines 145 - 153, Update the shutdown cleanup sequence around shutdown_background_topic_summary_tasks, A2AStorageFactory.cleanup, and A2AClientManager().close so the client close operation runs in its own finally block even when an earlier cleanup step raises. Preserve the outer Sentry flush behavior after all cleanup attempts.
🤖 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/a2a_client/__init__.py`:
- Around line 3-9: Update the package initializer by removing the
A2ADelegationCapability and A2AClientManager imports and deleting __all__; move
those imports into each consuming module so this __init__.py contains only a
brief package description.
In `@src/a2a_client/capability.py`:
- Around line 87-89: Update the documentation section header in the affected
function docstring from “Args:” to “Parameters:”, while preserving the existing
agent_name and task descriptions.
- Around line 130-142: Update the agent-list prompt construction in the
capability method around self._manager.list_agents() so remote card.description
and skill names are not inserted as trusted system-instruction prose. Restrict
displayed metadata to validated identifiers, or clearly encode the description
and skills as untrusted data while preserving the delegation list and agent
names.
- Around line 51-52: Replace the Any annotation on the _manager field with
A2AClientManager, and update _send_and_collect.client to use the concrete Client
interface. Add or reuse the appropriate imports, preserving the existing runtime
behavior while ensuring both class attributes and parameters are statically
checked against the SDK contracts.
- Around line 101-103: Update the delegation debug log in the try block around
create_text_message_object to remove the user-provided task from logger.debug,
retaining only the target agent name and non-sensitive diagnostic context.
- Around line 170-181: Update the terminal-state handling in the capability
delegation logic around the existing failed-state branch to treat
TaskState.failed, TaskState.rejected, and TaskState.canceled uniformly,
preserving the existing failure-message extraction and delegation error
response. In tests/unit/a2a_client/test_capability.py, extend the relevant unit
test coverage to verify all three terminal states produce the delegation failure
result.
In `@src/a2a_client/manager.py`:
- Around line 128-148: Update get_client and get_card in
src/a2a_client/manager.py to use modern union return annotations with | None
instead of Optional. Apply the same annotation change to the helper at
src/utils/pydantic_ai.py line 193; preserve each function’s existing return
behavior.
- Around line 97-113: Update the exception path surrounding
ClientFactory.connect and client.get_card in the connection manager to close any
partially initialized client and its httpx_client before handling or logging the
error. Ensure cleanup is conditional and safe when either resource was never
created, while preserving caching only after successful card retrieval.
- Around line 55-58: Update the request-kwargs preparation around the headers
assignment to copy the incoming mapping and its headers mapping before adding
the Bearer Authorization value. Preserve all existing request options while
ensuring the original input and caller-owned headers remain unchanged, and
assign the copied headers back to the copied kwargs before returning.
In `@src/configuration.py`:
- Around line 431-437: Update the a2a_agents property docstring to use the
Google-style format, adding a Returns section for its A2AAgentsConfiguration |
None result and a Raises section documenting LogicError when _configuration is
not loaded.
---
Outside diff comments:
In `@src/app/main.py`:
- Around line 145-153: Update the shutdown cleanup sequence around
shutdown_background_topic_summary_tasks, A2AStorageFactory.cleanup, and
A2AClientManager().close so the client close operation runs in its own finally
block even when an earlier cleanup step raises. Preserve the outer Sentry flush
behavior after all cleanup attempts.
🪄 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: e733f60c-2ce1-4f05-898b-ecf8346ba5f3
📒 Files selected for processing (14)
docs/a2a_protocol.mddocs/openapi.jsonexamples/lightspeed-stack-a2a-agents.yamlsrc/a2a_client/__init__.pysrc/a2a_client/capability.pysrc/a2a_client/manager.pysrc/app/main.pysrc/configuration.pysrc/models/config.pysrc/utils/pydantic_ai.pytests/unit/a2a_client/__init__.pytests/unit/a2a_client/test_capability.pytests/unit/a2a_client/test_manager.pytests/unit/models/config/test_dump_configuration.py
📜 Review details
⏰ Context from checks skipped due to timeout. (7)
- GitHub Check: E2E Tests for Lightspeed Evaluation job
- GitHub Check: E2E: server mode / ci / group 2
- 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: server mode / ci / group 3
- GitHub Check: E2E: library mode / ci / group 3
🧰 Additional context used
📓 Path-based instructions (8)
**/*
📄 CodeRabbit inference engine (Custom checks)
**/*: Flag meaningful O(n^2)+ algorithms on non-trivial inputs, including handlers and Kubernetes list operations.
Flag N+1 patterns that list items and then query once per item, including Kubernetes API and database access.
Flag expensive work inside loops, including API calls, JSON parsing, and regex compilation.
Flag unbounded growth in caches, watchers, or buffers when eviction or limits are missing.
Flag missing pagination or limits on list operations and API endpoints.
Flag secrets or tokens logged in plaintext or hardcoded in source.
Flag API endpoints missing authentication or authorization.
Flag injection vulnerabilities, including SQL injection, command injection, and path traversal.
Flag sensitive data leaked in API responses, WebSocket messages, or logs.
Flag Kubernetes Secrets and Red Hat secrets missing OwnerReferences.
Files:
tests/unit/a2a_client/__init__.pysrc/a2a_client/__init__.pyexamples/lightspeed-stack-a2a-agents.yamlsrc/a2a_client/capability.pydocs/openapi.jsonsrc/configuration.pysrc/a2a_client/manager.pysrc/app/main.pysrc/utils/pydantic_ai.pysrc/models/config.pytests/unit/models/config/test_dump_configuration.pytests/unit/a2a_client/test_manager.pytests/unit/a2a_client/test_capability.pydocs/a2a_protocol.md
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
**/*.py: Use absolute imports for internal Python modules.
Start every module with a descriptive docstring and uselogger = get_logger(__name__)for module logging.
Define shared constants in the centralconstants.pymodule, use descriptive comments, and annotate constants withFinal[type].
Use complete type annotations for function parameters and return values, modern union syntax, andtyping_extensions.Selffor model validators.
Document all modules, classes, and functions using Google-style Python docstrings, including applicable Parameters, Returns, Raises, and Attributes sections.
Use descriptive snake_case, action-oriented function names such asget_,validate_, andcheck_.
Avoid modifying input parameters in place; return a newly created data structure instead.
Useasync deffor I/O operations and external API calls.
Use standard logger levels appropriately: debug for diagnostics, info for normal execution, warning for unexpected conditions, and error for serious failures.
Name classes with PascalCase and use descriptive suffixes such asConfiguration,Error/Exception,Resolver, andInterfacewhere applicable.
Use ABC and@abstractmethodfor abstract interfaces; provide complete, specific type annotations for class attributes and avoidAny.
Files:
tests/unit/a2a_client/__init__.pysrc/a2a_client/__init__.pysrc/a2a_client/capability.pysrc/configuration.pysrc/a2a_client/manager.pysrc/app/main.pysrc/utils/pydantic_ai.pysrc/models/config.pytests/unit/models/config/test_dump_configuration.pytests/unit/a2a_client/test_manager.pytests/unit/a2a_client/test_capability.py
**/__init__.py
📄 CodeRabbit inference engine (AGENTS.md)
Keep package
__init__.pyfiles limited to brief package descriptions.
Files:
tests/unit/a2a_client/__init__.pysrc/a2a_client/__init__.py
tests/unit/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use pytest for unit tests, shared fixtures in
conftest.py,pytest-mockfor mocks, andpytest.mark.asynciofor asynchronous tests; do not use unittest.
Files:
tests/unit/a2a_client/__init__.pytests/unit/models/config/test_dump_configuration.pytests/unit/a2a_client/test_manager.pytests/unit/a2a_client/test_capability.py
tests/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Maintain at least 60% unit-test coverage and 10% integration-test coverage.
Files:
tests/unit/a2a_client/__init__.pytests/unit/models/config/test_dump_configuration.pytests/unit/a2a_client/test_manager.pytests/unit/a2a_client/test_capability.py
**/*.{py,yaml,yml,json,toml}
📄 CodeRabbit inference engine (AGENTS.md)
Never commit secrets or keys; use environment variables for sensitive data.
Files:
tests/unit/a2a_client/__init__.pysrc/a2a_client/__init__.pyexamples/lightspeed-stack-a2a-agents.yamlsrc/a2a_client/capability.pydocs/openapi.jsonsrc/configuration.pysrc/a2a_client/manager.pysrc/app/main.pysrc/utils/pydantic_ai.pysrc/models/config.pytests/unit/models/config/test_dump_configuration.pytests/unit/a2a_client/test_manager.pytests/unit/a2a_client/test_capability.py
src/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Configuration models must extend
ConfigurationBase, whoseextra="forbid"behavior rejects unknown fields; use types such asOptional[FilePath],PositiveInt, andSecretStrwhere appropriate.
Files:
src/a2a_client/__init__.pysrc/a2a_client/capability.pysrc/configuration.pysrc/a2a_client/manager.pysrc/app/main.pysrc/utils/pydantic_ai.pysrc/models/config.py
src/models/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use
ConfigurationBasefor configuration Pydantic models andBaseModelfor data models; use@model_validatorand@field_validatorfor validation.
Files:
src/models/config.py
🧠 Learnings (5)
📚 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/a2a_client/__init__.pysrc/a2a_client/__init__.pysrc/a2a_client/capability.pysrc/configuration.pysrc/a2a_client/manager.pysrc/app/main.pysrc/utils/pydantic_ai.pysrc/models/config.pytests/unit/models/config/test_dump_configuration.pytests/unit/a2a_client/test_manager.pytests/unit/a2a_client/test_capability.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/a2a_client/__init__.pysrc/a2a_client/capability.pysrc/configuration.pysrc/a2a_client/manager.pysrc/app/main.pysrc/utils/pydantic_ai.pysrc/models/config.py
📚 Learning: 2026-05-20T08:09:30.641Z
Learnt from: max-svistunov
Repo: lightspeed-core/lightspeed-stack PR: 1580
File: docs/design/llama-stack-config-merge/poc-results/library-mode/synthesized-run.yaml:107-110
Timestamp: 2026-05-20T08:09:30.641Z
Learning: In Llama-stack config YAMLs, when defining a Llama Guard safety shield entry, set `provider_shield_id` to the *guard model identifier* (e.g., `meta-llama/Llama-Guard-3-8B`). Do not use a chat/generative model id (e.g., `openai/gpt-4o-mini`): a chat-model id (or `native_override`) indicates only an override landed and does **not** mean the safety shield is actually gating queries. Ensure any E2E coverage for the related implementation (JIRA/E2E tests) exercises a real Llama Guard model to verify that the shield is effective.
Applied to files:
examples/lightspeed-stack-a2a-agents.yaml
📚 Learning: 2026-01-12T10:58:40.230Z
Learnt from: blublinsky
Repo: lightspeed-core/lightspeed-stack PR: 972
File: src/models/config.py:459-513
Timestamp: 2026-01-12T10:58:40.230Z
Learning: In lightspeed-core/lightspeed-stack, for Python files under src/models, when a user claims a fix is done but the issue persists, verify the current code state before accepting the fix. Steps: review the diff, fetch the latest changes, run relevant tests, reproduce the issue, search the codebase for lingering references to the original problem, confirm the fix is applied and not undone by subsequent commits, and validate with local checks to ensure the issue is resolved.
Applied to files:
src/models/config.py
📚 Learning: 2026-02-25T07:46:33.545Z
Learnt from: asimurka
Repo: lightspeed-core/lightspeed-stack PR: 1211
File: src/models/responses.py:8-16
Timestamp: 2026-02-25T07:46:33.545Z
Learning: In the Python codebase, requests.py should use OpenAIResponseInputTool as Tool while responses.py uses OpenAIResponseTool as Tool. This difference is intentional due to differing schemas for input vs output tools in llama-stack-api. Apply this distinction consistently to other models under src/models (e.g., ensure request-related tools use the InputTool variant and response-related tools use the ResponseTool variant). If adding new tools, choose the corresponding InputTool or Tool class based on whether the tool represents input or output, and document the rationale in code comments.
Applied to files:
src/models/config.py
🪛 Checkov (3.3.8)
docs/openapi.json
[medium] 10768-10775: Ensure that arrays have a maximum number of items
(CKV_OPENAPI_21)
🔇 Additional comments (11)
tests/unit/a2a_client/__init__.py (1)
1-1: LGTM!src/models/config.py (1)
1981-2034: LGTM!Also applies to: 2727-2732
docs/openapi.json (2)
10714-10781: LGTM!
12230-12241: LGTM!examples/lightspeed-stack-a2a-agents.yaml (1)
1-42: LGTM!tests/unit/models/config/test_dump_configuration.py (1)
212-212: LGTM!Also applies to: 436-436, 811-811, 1070-1070, 1309-1309, 1528-1528, 1907-1907, 2132-2132, 2357-2357, 2589-2589
docs/a2a_protocol.md (1)
7-10: LGTM!Also applies to: 316-317, 792-819
src/app/main.py (1)
17-17: LGTM!Also applies to: 132-136
src/utils/pydantic_ai.py (1)
15-16: LGTM!Also applies to: 199-202, 219-222
src/a2a_client/manager.py (1)
102-108: 🎯 Functional Correctness
ClientFactory.connectis part of the supporteda2a-sdkAPI.> Likely an incorrect or invalid review comment.tests/unit/a2a_client/test_manager.py (1)
72-75: 🎯 Functional CorrectnessNo change needed for
ClientFactory.connectThis test matchessrc/a2a_client/manager.py's use ofClientFactory.connect, so the API mismatch is not present.> Likely an incorrect or invalid review comment.
| from a2a_client.capability import A2ADelegationCapability | ||
| from a2a_client.manager import A2AClientManager | ||
|
|
||
| __all__ = [ | ||
| "A2AClientManager", | ||
| "A2ADelegationCapability", | ||
| ] |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Keep this package initializer description-only.
Move these imports to their consumers and remove __all__.
As per coding guidelines, package __init__.py files must be limited to brief package descriptions.
🤖 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/a2a_client/__init__.py` around lines 3 - 9, Update the package
initializer by removing the A2ADelegationCapability and A2AClientManager imports
and deleting __all__; move those imports into each consuming module so this
__init__.py contains only a brief package description.
Source: Coding guidelines
| _manager: Any = field(repr=False) | ||
| _toolset: FunctionToolset[object] = field(init=False, repr=False) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Replace Any with the concrete manager and client interfaces.
Type _manager as A2AClientManager and _send_and_collect.client as Client so SDK contract changes are caught statically.
As per coding guidelines, class attributes and function parameters require specific annotations and should avoid Any.
Also applies to: 150-150
🤖 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/a2a_client/capability.py` around lines 51 - 52, Replace the Any
annotation on the _manager field with A2AClientManager, and update
_send_and_collect.client to use the concrete Client interface. Add or reuse the
appropriate imports, preserving the existing runtime behavior while ensuring
both class attributes and parameters are statically checked against the SDK
contracts.
Source: Coding guidelines
| Args: | ||
| agent_name: Name of the agent to delegate to (from list_agents). | ||
| task: The task description or question to send to the agent. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Use Parameters: instead of Args:.
Based on learnings, this repository requires the section header Parameters: for documented function arguments.
🤖 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/a2a_client/capability.py` around lines 87 - 89, Update the documentation
section header in the affected function docstring from “Args:” to “Parameters:”,
while preserving the existing agent_name and task descriptions.
Source: Learnings
| try: | ||
| logger.debug("Delegating to agent '%s': %s", agent_name, task) | ||
| message = create_text_message_object(role=Role.user, content=task) |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Do not log the delegated user task.
task may contain prompts, credentials, or personal data. Log only the target agent and non-sensitive diagnostics.
Proposed fix
- logger.debug("Delegating to agent '%s': %s", agent_name, task)
+ logger.debug("Delegating task to agent '%s'", agent_name)As per coding guidelines, sensitive data must not be logged in plaintext.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| try: | |
| logger.debug("Delegating to agent '%s': %s", agent_name, task) | |
| message = create_text_message_object(role=Role.user, content=task) | |
| try: | |
| logger.debug("Delegating task to agent '%s'", agent_name) | |
| message = create_text_message_object(role=Role.user, content=task) |
🤖 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/a2a_client/capability.py` around lines 101 - 103, Update the delegation
debug log in the try block around create_text_message_object to remove the
user-provided task from logger.debug, retaining only the target agent name and
non-sensitive diagnostic context.
Source: Coding guidelines
| agents = self._manager.list_agents() | ||
| if not agents: | ||
| return None | ||
| lines = ["You can delegate tasks to these external agents when appropriate:"] | ||
| for name, card in agents.items(): | ||
| skills_text = "" | ||
| if card.skills: | ||
| skill_names = [s.name for s in card.skills if s.name] | ||
| if skill_names: | ||
| skills_text = f" (skills: {', '.join(skill_names)})" | ||
| lines.append( | ||
| f"- {name}: {card.description or 'No description'}{skills_text}" | ||
| ) |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | 🏗️ Heavy lift
Do not place untrusted agent-card prose directly in system instructions.
A remote agent controls its description and skill names, allowing a compromised endpoint to inject model instructions. Restrict this prompt to validated identifiers or encode remote metadata as explicitly untrusted data.
🤖 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/a2a_client/capability.py` around lines 130 - 142, Update the agent-list
prompt construction in the capability method around self._manager.list_agents()
so remote card.description and skill names are not inserted as trusted
system-instruction prose. Restrict displayed metadata to validated identifiers,
or clearly encode the description and skills as untrusted data while preserving
the delegation list and agent names.
| elif ( | ||
| update_event is not None | ||
| and hasattr(update_event, "status") | ||
| and update_event.status.state == TaskState.failed | ||
| ): | ||
| fail_msg = "Agent task failed" | ||
| if update_event.status.message: | ||
| for part in update_event.status.message.parts: | ||
| if isinstance(part.root, TextPart): | ||
| fail_msg = part.root.text | ||
| break | ||
| return f"Delegation failed: {fail_msg}" |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== capability.py slice ==\n'
sed -n '130,220p' src/a2a_client/capability.py
printf '\n== test_capability.py slice ==\n'
sed -n '150,260p' tests/unit/a2a_client/test_capability.py
printf '\n== TaskState definitions/usages ==\n'
rg -n "class TaskState|TaskState\.(failed|rejected|canceled)|rejected|canceled" src tests -g '!**/__pycache__/**'Repository: lightspeed-core/lightspeed-stack
Length of output: 50388
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== TaskState definition ==\n'
rg -n "class TaskState|TaskState =|enum.*TaskState" src tests -g '*.py' | head -n 40
printf '\n== a2a endpoint handling slice ==\n'
sed -n '130,210p' src/app/endpoints/a2a.py
printf '\n== a2a client capability surrounding file top ==\n'
sed -n '1,240p' src/a2a_client/capability.pyRepository: lightspeed-core/lightspeed-stack
Length of output: 9923
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== TaskState definition ==\n'
rg -n "class TaskState|TaskState =" src tests -g '*.py' | head -n 20
printf '\n== a2a endpoint handling slice ==\n'
sed -n '130,210p' src/app/endpoints/a2a.py
printf '\n== capability file with line numbers ==\n'
cat -n src/a2a_client/capability.py | sed -n '1,260p'Repository: lightspeed-core/lightspeed-stack
Length of output: 11188
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== capability.py imports ==\n'
sed -n '1,120p' src/a2a_client/capability.py
printf '\n== app/endpoints/a2a.py imports and handler slice ==\n'
sed -n '1,120p' src/app/endpoints/a2a.py
sed -n '150,190p' src/app/endpoints/a2a.py
printf '\n== files mentioning TaskState in source tree (limited) ==\n'
rg -n "TaskState" src -g '*.py' | head -n 80Repository: lightspeed-core/lightspeed-stack
Length of output: 12098
🌐 Web query:
A2A protocol TaskState failed rejected canceled task status update enum
💡 Result:
In the A2A protocol, the TaskState enum defines the lifecycle states of a task, and the states failed, rejected, and canceled are officially classified as terminal states [1][2][3]. TaskState Enum Definitions The TaskState enum includes the following terminal states [1][2][3]: TASK_STATE_FAILED (4): Indicates the task finished with an error [1][2]. TASK_STATE_CANCELED (5): Indicates the task was canceled before completion (often by the client) [1][2][4]. TASK_STATE_REJECTED (7): Indicates the agent decided not to perform the task, which may occur at initial creation or later [1][2]. Task Status Updates When a task transitions into any of these terminal states, the protocol typically requires a final status update [5][3]. These updates are handled via the TaskStatusUpdateEvent message [1][6][7]. Key components of the update process include: TaskStatusUpdateEvent: This event notifies the client of a status change [1][6]. It contains the task ID, context ID, and the new status [1][6]. Terminal Marker: The TaskStatusUpdateEvent includes a boolean field (final) that is set to true when the task reaches a terminal state, such as failed, rejected, or canceled [6][8][7]. Implementation: SDKs often provide helper methods (e.g., cancel, failed, reject) within a task_updater module that automatically mark the task as having reached a terminal state and publish the necessary final status update [5][7]. Once a task reaches a terminal state, subsequent status updates are generally prohibited or will raise an error [7].
Citations:
- 1: https://github.com/a2aproject/A2A/blob/629190ae/specification/a2a.proto
- 2: https://docs.rs/a2a-protocol-types/latest/a2a_protocol_types/task/enum.TaskState.html
- 3: https://github.com/a2aproject/A2A/blob/main/docs/specification.md
- 4: https://docs.rs/xa2a/latest/xa2a/types/enum.TaskState.html
- 5: https://a2a-protocol.org/latest/sdk/python/api/a2a.server.tasks.task_updater.html
- 6: https://github.com/google-a2a/a2a-python/blob/fa14dbf4/src/a2a/grpc/a2a_pb2.pyi
- 7: https://github.com/google-a2a/a2a-python/blob/fa14dbf4/src/a2a/server/tasks/task_updater.py
- 8: https://a2aprotocol.ai/docs/guide/a2a-protocol-specification-python
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== capability.py imports ==\n'
sed -n '1,120p' src/a2a_client/capability.py
printf '\n== a2a.py imports and status handling ==\n'
sed -n '1,190p' src/app/endpoints/a2a.py
printf '\n== TaskState mentions in src ==\n'
rg -n "TaskState" src -g '*.py' | head -n 120Repository: lightspeed-core/lightspeed-stack
Length of output: 13048
🌐 Web query:
A2A protocol TaskState values rejected canceled failed terminal states
💡 Result:
In the A2A protocol, the TaskState values rejected, canceled, failed, and completed are all defined as terminal states [1][2][3]. A terminal state indicates that a task has permanently ceased execution and cannot be restarted [4][5][6]. The specific terminal TaskState values are: - TASK_STATE_COMPLETED: Indicates the task finished successfully [1][2]. - TASK_STATE_FAILED: Indicates the task finished with an error [1][2]. - TASK_STATE_CANCELED: Indicates the task was canceled before completion [1][2]. - TASK_STATE_REJECTED: Indicates the agent decided not to perform the task, which may occur at creation or later during processing [1][2]. Once a task reaches any of these terminal states, it is no longer capable of further progression [4][6]. Attempting to send messages to a task in a terminal state will result in an UnsupportedOperationError [1][7][8]. Additionally, stream subscriptions for tasks will terminate upon reaching these states [1][7][8]. Other states, such as TASK_STATE_INPUT_REQUIRED and TASK_STATE_AUTH_REQUIRED, are classified as interrupted states rather than terminal states [1][2][3].
Citations:
- 1: https://a2a-protocol.org/v1.0.0/specification/
- 2: https://github.com/a2aproject/A2A/blob/629190ae/specification/a2a.proto
- 3: https://docs.rs/a2a-protocol-types/latest/a2a_protocol_types/task/enum.TaskState.html
- 4: https://a2a-protocol.org/dev/topics/life-of-a-task/
- 5: https://github.com/google/A2A/blob/7b900e77/docs/topics/life-of-a-task.md
- 6: https://a2a-protocol.org/latest/topics/life-of-a-task/
- 7: https://github.com/google/A2A/blob/7b900e77/docs/specification.md
- 8: https://github.com/a2aproject/A2A/blob/main/docs/specification.md
Handle all terminal task states from A2A. TaskState.rejected and TaskState.canceled are terminal too; they currently fall through to "Agent returned no response." instead of surfacing a delegation error. Add those states here and cover failed/rejected/canceled in the unit test.
📍 Affects 2 files
src/a2a_client/capability.py#L170-L181(this comment)tests/unit/a2a_client/test_capability.py#L178-L210
🤖 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/a2a_client/capability.py` around lines 170 - 181, Update the
terminal-state handling in the capability delegation logic around the existing
failed-state branch to treat TaskState.failed, TaskState.rejected, and
TaskState.canceled uniformly, preserving the existing failure-message extraction
and delegation error response. In tests/unit/a2a_client/test_capability.py,
extend the relevant unit test coverage to verify all three terminal states
produce the delegation failure result.
| headers = http_kwargs.get("headers", {}) | ||
| headers["Authorization"] = f"Bearer {self._token}" | ||
| http_kwargs["headers"] = headers | ||
| return request_payload, http_kwargs |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Copy request kwargs before injecting the credential.
This mutates both the input mapping and potentially its caller-owned headers mapping. Build new dictionaries so authorization state cannot leak across reused request options.
Proposed fix
- headers = http_kwargs.get("headers", {})
+ http_kwargs = dict(http_kwargs)
+ headers = dict(http_kwargs.get("headers", {}))
headers["Authorization"] = f"Bearer {self._token}"
http_kwargs["headers"] = headersAs per coding guidelines, input parameters must not be modified in place.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| headers = http_kwargs.get("headers", {}) | |
| headers["Authorization"] = f"Bearer {self._token}" | |
| http_kwargs["headers"] = headers | |
| return request_payload, http_kwargs | |
| http_kwargs = dict(http_kwargs) | |
| headers = dict(http_kwargs.get("headers", {})) | |
| headers["Authorization"] = f"Bearer {self._token}" | |
| http_kwargs["headers"] = headers | |
| return request_payload, http_kwargs |
🤖 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/a2a_client/manager.py` around lines 55 - 58, Update the request-kwargs
preparation around the headers assignment to copy the incoming mapping and its
headers mapping before adding the Bearer Authorization value. Preserve all
existing request options while ensuring the original input and caller-owned
headers remain unchanged, and assign the copied headers back to the copied
kwargs before returning.
Source: Coding guidelines
| transport = httpx.AsyncHTTPTransport(retries=agent_config.max_retries) | ||
| httpx_client = httpx.AsyncClient( | ||
| timeout=httpx.Timeout(agent_config.timeout), | ||
| transport=transport, | ||
| ) | ||
| client = await ClientFactory.connect( | ||
| url, | ||
| client_config=ClientConfig( | ||
| streaming=True, httpx_client=httpx_client | ||
| ), | ||
| interceptors=interceptors or None, | ||
| ) | ||
| card = await client.get_card() | ||
| self._clients[name] = client | ||
| self._cards[name] = card | ||
| logger.info("Connected to A2A agent '%s' at %s", name, url) | ||
| except (A2AClientError, httpx.HTTPError, OSError) as e: |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Close partially created connections when discovery fails.
If connection or card retrieval raises, httpx_client and possibly client are neither cached nor closed. Explicitly close the partially initialized resource in the exception path.
🤖 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/a2a_client/manager.py` around lines 97 - 113, Update the exception path
surrounding ClientFactory.connect and client.get_card in the connection manager
to close any partially initialized client and its httpx_client before handling
or logging the error. Ensure cleanup is conditional and safe when either
resource was never created, while preserving caching only after successful card
retrieval.
| def get_client(self, agent_name: str) -> Optional[Client]: | ||
| """Return the Client for a named agent. | ||
|
|
||
| Parameters: | ||
| agent_name: Name of the agent as configured. | ||
|
|
||
| Returns: | ||
| Client instance, or None if the agent is not available. | ||
| """ | ||
| return self._clients.get(agent_name) | ||
|
|
||
| def get_card(self, agent_name: str) -> Optional[AgentCard]: | ||
| """Return the cached AgentCard for a named agent. | ||
|
|
||
| Parameters: | ||
| agent_name: Name of the agent as configured. | ||
|
|
||
| Returns: | ||
| AgentCard, or None if the agent is not available. | ||
| """ | ||
| return self._cards.get(agent_name) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Use modern union syntax consistently in the new A2A surface.
src/a2a_client/manager.py#L128-L148: replace bothOptional[...]return annotations with... | None.src/utils/pydantic_ai.py#L193-L193: replace the helper’sOptional[...]return annotation with... | None.
As per coding guidelines, functions must use modern union syntax.
📍 Affects 2 files
src/a2a_client/manager.py#L128-L148(this comment)src/utils/pydantic_ai.py#L193-L193
🤖 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/a2a_client/manager.py` around lines 128 - 148, Update get_client and
get_card in src/a2a_client/manager.py to use modern union return annotations
with | None instead of Optional. Apply the same annotation change to the helper
at src/utils/pydantic_ai.py line 193; preserve each function’s existing return
behavior.
Source: Coding guidelines
| @property | ||
| def a2a_agents(self) -> A2AAgentsConfiguration | None: | ||
| """Return A2A clients configuration.""" | ||
| if self._configuration is None: | ||
| raise LogicError("logic error: configuration is not loaded") | ||
| return self._configuration.a2a_agents | ||
|
|
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Include Returns: and Raises: in the property docstring.
As per coding guidelines, all functions and properties must be documented using Google-style Python docstrings, including applicable Returns: and Raises: sections. Since this property dynamically raises a LogicError and returns a specific type, these details should be explicitly documented to match the surrounding properties (like database_configuration).
📝 Proposed docstring fix
`@property`
def a2a_agents(self) -> A2AAgentsConfiguration | None:
- """Return A2A clients configuration."""
+ """Return A2A clients configuration.
+
+ Returns:
+ A2AAgentsConfiguration | None: The configured A2A agents or None if not set.
+
+ Raises:
+ LogicError: If the configuration has not been loaded.
+ """
if self._configuration is None:
raise LogicError("logic error: configuration is not loaded")
return self._configuration.a2a_agents📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| @property | |
| def a2a_agents(self) -> A2AAgentsConfiguration | None: | |
| """Return A2A clients configuration.""" | |
| if self._configuration is None: | |
| raise LogicError("logic error: configuration is not loaded") | |
| return self._configuration.a2a_agents | |
| `@property` | |
| def a2a_agents(self) -> A2AAgentsConfiguration | None: | |
| """Return A2A clients configuration. | |
| Returns: | |
| A2AAgentsConfiguration | None: The configured A2A agents or None if not set. | |
| Raises: | |
| LogicError: If the configuration has not been loaded. | |
| """ | |
| if self._configuration is None: | |
| raise LogicError("logic error: configuration is not loaded") | |
| return self._configuration.a2a_agents |
🤖 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/configuration.py` around lines 431 - 437, Update the a2a_agents property
docstring to use the Google-style format, adding a Returns section for its
A2AAgentsConfiguration | None result and a Raises section documenting LogicError
when _configuration is not loaded.
Source: Coding guidelines
2 participants
Description
Enable LCS to act as an A2A client — discovering and delegating tasks to external A2A-compliant agents during inference. Previously LCS only supported the A2A protocol as a server. This PR adds both the client-side infrastructure and refactors the A2A server to use pydantic-ai agents.
A2A Server Refactor
A2AAgentExecutorto usebuild_agent()+agent.run_stream_events()instead of rawclient.responses.create(), aligning the A2A path with query/streaming endpoints_build_a2a_parts_from_agent_result()and_dispatch_agent_event()for pydantic-ai → A2A event mappingA2A Client
A2AAgentsConfiguration: YAML config for remote agent endpoints withtimeout,max_retries, and optionalauth_tokenA2AClientManager: singleton that resolves agent cards at startup, holdsClientinstances with httpx timeout/retry transportA2ADelegationCapability: pydantic-ai capability withlist_agentsanddelegate_to_agenttools, auto-injected intobuild_agent()when agents are configuredDocs & Config
docs/a2a_protocol.mdupdated with A2A client documentationexamples/lightspeed-stack-a2a-agents.yamlexample configType 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— 2992 tests pass, 0 failuresSummary by CodeRabbit
New Features
Documentation
Tests