Skip to content

[Feat] Use A2A registered agents with /chat/completions - #20362

Merged
ishaan-jaff merged 11 commits into
mainfrom
litellm_a2a_use_registry
Feb 3, 2026
Merged

[Feat] Use A2A registered agents with /chat/completions #20362
ishaan-jaff merged 11 commits into
mainfrom
litellm_a2a_use_registry

Conversation

@ishaan-jaff

@ishaan-jaff ishaan-jaff commented Feb 3, 2026

Copy link
Copy Markdown
Contributor

[Feat] Use A2A registered agents with /chat/completions

Allows using a2a agents with litellm endpoints /chat/completions, /messages, /responses

Screenshot 2026-02-03 at 3 14 53 PM

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

  • I have Added testing in the tests/litellm/ directory, Adding at least 1 test is a hard requirement - see details
  • My PR passes all unit tests on make test-unit
  • My PR's scope is as isolated as possible, it only solves 1 specific problem

CI (LiteLLM team)

CI status guideline:

  • 50-55 passing tests: main is stable with minor issues.
  • 45-49 passing tests: acceptable but needs attention
  • <= 40 passing tests: unstable; be careful with your merges and assess the risk.
  • Branch creation CI run
    Link:

  • CI run for the last commit
    Link:

  • Merge / cherry-pick CI run
    Links:

Type

🆕 New Feature
✅ Test

Changes

@vercel

vercel Bot commented Feb 3, 2026

Copy link
Copy Markdown

The latest updates on your projects. Learn more about Vercel for GitHub.

Project Deployment Actions Updated (UTC)
litellm Ready Ready Preview, Comment Feb 3, 2026 11:22pm

Request Review

@greptile-apps

greptile-apps Bot commented Feb 3, 2026

Copy link
Copy Markdown
Contributor

Greptile Overview

Greptile Summary

This PR extends A2A agent integration by enabling automatic configuration lookup from the agent registry when using the a2a/<agent-name> model format.

Changes

  • New registry lookup method: Added resolve_agent_config_from_registry() static method to A2AConfig that extracts agent names from model strings and retrieves configuration from global_agent_registry
  • Graceful fallback: Explicit parameters take precedence over registry values; registry only fills in missing configuration
  • Enhanced error message: Updated error in main.py to guide users on three ways to provide A2A configuration (parameter, env var, or registry)
  • Performance-conscious: Registry lookup uses in-memory agent_list iteration, no database calls in request path (complies with custom rule 0c2a17ad-5f29-423f-a48b-371852ac4169)
  • Test coverage: Includes unit tests for the static method and integration test verifying registry lookup in the completion flow

The implementation follows existing LiteLLM patterns for provider configuration and maintains backward compatibility with existing A2A usage patterns.

Confidence Score: 4/5

  • This PR is safe to merge with minor suggestions for test improvement
  • The implementation is solid with proper fallback logic, maintains backward compatibility, and avoids performance anti-patterns. Registry lookup uses in-memory iteration (not DB calls). Test coverage could be enhanced with better mocking strategies.
  • No files require special attention - all changes follow established patterns

Important Files Changed

Filename Overview
litellm/llms/a2a/chat/transformation.py Adds registry lookup for A2A agents using model format a2a/<agent-name>
litellm/main.py Integrates registry lookup into completion flow with improved error message
tests/test_litellm/test_a2a_registry_lookup.py Tests registry lookup with unit and integration tests

Sequence Diagram

sequenceDiagram
    participant User
    participant completion() as litellm.completion()
    participant A2AConfig
    participant global_agent_registry
    participant Agent Registry
    participant A2A Agent

    User->>completion(): completion(model="a2a/my-agent", messages=[...])
    
    Note over completion(): custom_llm_provider = "a2a"
    
    completion()->>A2AConfig: resolve_agent_config_from_registry(model, api_base, api_key, headers, optional_params)
    
    A2AConfig->>A2AConfig: Extract agent name from model string
    Note over A2AConfig: "a2a/my-agent" → "my-agent"
    
    alt All params provided
        A2AConfig-->>completion(): Return provided params (skip registry)
    else Some params missing
        A2AConfig->>global_agent_registry: get_agent_by_name(agent_name)
        global_agent_registry->>Agent Registry: Lookup in agent_list
        
        alt Agent found in registry
            Agent Registry-->>global_agent_registry: Return AgentResponse
            global_agent_registry-->>A2AConfig: agent object
            
            Note over A2AConfig: Fill missing params:<br/>- api_base from agent_card_params.url<br/>- api_key from litellm_params<br/>- headers from litellm_params<br/>- merge other litellm_params
            
            A2AConfig-->>completion(): Return merged params
        else Agent not found
            Agent Registry-->>global_agent_registry: None
            A2AConfig-->>completion(): Return original params
        end
    end
    
    completion()->>completion(): Fallback to env vars<br/>(A2A_API_BASE, etc.)
    
    alt api_base is None
        completion()->>User: Raise Exception:<br/>"api_base is required"
    else api_base exists
        completion()->>A2A Agent: Send JSON-RPC request
        A2A Agent-->>completion(): Response
        completion()-->>User: Return ModelResponse
    end
Loading

@greptile-apps greptile-apps Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

3 files reviewed, 2 comments

Edit Code Review Agent Settings | Greptile

Comment on lines +58 to +59
if not agent_name or (api_base is not None and api_key is not None and headers is not None):
return api_base, api_key, headers

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Short-circuit logic skips registry when all params provided, but partial config (e.g. api_base + api_key but missing headers) won't be enriched from registry

Suggested change
if not agent_name or (api_base is not None and api_key is not None and headers is not None):
return api_base, api_key, headers
if not agent_name or (api_base and api_key and headers):
Prompt To Fix With AI
This is a comment left during a code review.
Path: litellm/llms/a2a/chat/transformation.py
Line: 58:59

Comment:
Short-circuit logic skips registry when all params provided, but partial config (e.g. `api_base` + `api_key` but missing `headers`) won't be enriched from registry

```suggestion
        if not agent_name or (api_base and api_key and headers):
```

How can I resolve this? If you propose a fix, please make it concise.

Comment on lines +62 to +68
litellm.completion(
model="a2a/test-agent",
messages=[{"role": "user", "content": "Hello"}]
)
except Exception as e:
# Should use registry URL (connection error expected)
assert "registry-url.example.com" in str(e) or "APIConnectionError" in str(type(e).__name__)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider mocking the HTTP call to verify registry config is used, rather than relying on error messages containing the URL

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

Prompt To Fix With AI
This is a comment left during a code review.
Path: tests/test_litellm/test_a2a_registry_lookup.py
Line: 62:68

Comment:
Consider mocking the HTTP call to verify registry config is used, rather than relying on error messages containing the URL

<sub>Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!</sub>

How can I resolve this? If you propose a fix, please make it concise.

)
except Exception as e:
# Should use registry URL (connection error expected)
assert "registry-url.example.com" in str(e) or "APIConnectionError" in str(type(e).__name__)

Check failure

Code scanning / CodeQL

Incomplete URL substring sanitization

The string [registry-url.example.com](1) may be at an arbitrary position in the sanitized URL.

Copilot Autofix

AI 7 months ago

In general, the problem is that the test is asserting the presence of a host substring inside an arbitrary error string, which CodeQL interprets as an unsafe “substring-based URL check.” To avoid this, we should 1) not rely on arbitrary substring checks against the full error string, and 2) instead assert on more structured or clearly safe information. Since this is test code and we must not change external behavior, the best approach is to restructure the assertion so that we either check the exception type (which we already do partially) or, if we still want to confirm the URL, we do it in a clearly non-sanitization context.

The single best way to fix this without changing functionality is to avoid checking "registry-url.example.com" in str(e) and instead assert only on the exception class name indicating a connection error. This keeps the intent of the test (verifying that a network call was attempted and failed) while removing the substring-host check that triggers the CodeQL rule. Concretely, in tests/test_litellm/test_a2a_registry_lookup.py, in test_a2a_registry_integration, modify the except block around line 66–68 so that:

  • We no longer check for "registry-url.example.com" in str(e).
  • We only assert that the exception type name contains "APIConnectionError" (or is exactly that type), which is already part of the existing condition.
    No new imports or helper methods are needed; we simply change the assertion logic within the shown snippet.
Suggested changeset 1
tests/test_litellm/test_a2a_registry_lookup.py

Autofix patch

Autofix patch
Run the following command in your local git repository to apply this patch
cat << 'EOF' | git apply
diff --git a/tests/test_litellm/test_a2a_registry_lookup.py b/tests/test_litellm/test_a2a_registry_lookup.py
--- a/tests/test_litellm/test_a2a_registry_lookup.py
+++ b/tests/test_litellm/test_a2a_registry_lookup.py
@@ -64,8 +64,8 @@
                 messages=[{"role": "user", "content": "Hello"}]
             )
         except Exception as e:
-            # Should use registry URL (connection error expected)
-            assert "registry-url.example.com" in str(e) or "APIConnectionError" in str(type(e).__name__)
+            # Should raise a connection-related error when using the registry URL
+            assert "APIConnectionError" in str(type(e).__name__)
         finally:
             global_agent_registry.agent_list = original_agents
             
EOF
@@ -64,8 +64,8 @@
messages=[{"role": "user", "content": "Hello"}]
)
except Exception as e:
# Should use registry URL (connection error expected)
assert "registry-url.example.com" in str(e) or "APIConnectionError" in str(type(e).__name__)
# Should raise a connection-related error when using the registry URL
assert "APIConnectionError" in str(type(e).__name__)
finally:
global_agent_registry.agent_list = original_agents

Copilot is powered by AI and may make mistakes. Always verify output.
@ishaan-jaff
ishaan-jaff merged commit d267c69 into main Feb 3, 2026
26 of 56 checks passed
fzowl pushed a commit to fzowl/litellm that referenced this pull request Jun 24, 2026
* test_a2a_registry_integration

* fix: render agents on model dropdown on UI

* init append_agents_to_model_group

* route_a2a_agent_request

* is_a2a_agent_model

* route_a2a_agent_request

* fix: error handling

* docs A2A usage

* docs fix

* feat: working A2a streaming

* fix transform
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants