Skip to content

Commit 3d46ab8

Browse files
committed
feat(a2a): add credential service and auto auth token injection
feat(a2a): add credential service support for authenticated agent communication This commit introduces comprehensive authentication support for A2A (Agent-to-Agent) communication, enabling secure credential management and automatic token injection. Key Changes: 1. **VeCredentialStore** (veadk/a2a/credentials.py) - Implement custom credential store with user ID and session ID support - Support both synchronous and asynchronous credential operations - Prioritize user ID over session ID for credential retrieval 2. **AuthenticatedA2ARequestConverter** (veadk/a2a/ve_request_converter.py) - Extract JWT tokens from Authorization headers - Parse user ID from JWT payload (sub field) 3. **RemoteVeAgent** (veadk/a2a/remote_ve_agent.py) - Add credential_service parameter to constructor - Implement _run_async_impl with automatic auth token injection - Inject Bearer tokens into httpx client headers before requests - Add comprehensive error handling and logging 4. **VeA2AServer** (veadk/a2a/ve_a2a_server.py) - Add credential_service parameter to constructor - Integrate with AuthenticatedA2ARequestConverter 5. **Unit Tests** (tests/) - Add comprehensive test coverage for VeCredentialStore - Add tests for AuthenticatedA2ARequestConverter - All tests passing with proper fixtures and mocking Benefits: - Seamless authentication for remote agent calls - Automatic credential propagation across agent boundaries - Support for both session-based and user-based authentication - Clean separation of concerns with dedicated credential service Breaking Changes: - None (backward compatible - credential_service is optional) Related: A2A authentication and secure agent communication
1 parent c7fd907 commit 3d46ab8

File tree

6 files changed

+1166
-4
lines changed

6 files changed

+1166
-4
lines changed

tests/test_a2a_credentials.py

Lines changed: 279 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,279 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Unit tests for VeCredentialStore."""
16+
17+
import pytest
18+
19+
from a2a.client.base_client import ClientCallContext
20+
21+
from veadk.a2a.credentials import VeCredentialStore
22+
23+
24+
class TestVeCredentialStore:
25+
"""Test suite for VeCredentialStore."""
26+
27+
@pytest.mark.asyncio
28+
async def test_set_and_get_by_session_id(self):
29+
"""Test setting and getting credentials by session ID."""
30+
store = VeCredentialStore()
31+
32+
# Set credentials synchronously
33+
store.set_credentials(
34+
session_id="session_123",
35+
security_scheme_name="inbound_auth",
36+
credential="bearer_token_xyz",
37+
)
38+
39+
# Get credentials
40+
context = ClientCallContext(state={"sessionId": "session_123"})
41+
token = await store.get_credentials(
42+
security_scheme_name="inbound_auth",
43+
context=context,
44+
)
45+
46+
assert token == "bearer_token_xyz"
47+
48+
@pytest.mark.asyncio
49+
async def test_set_and_get_by_user_id(self):
50+
"""Test setting and getting credentials by user ID."""
51+
store = VeCredentialStore()
52+
53+
# Set credentials synchronously
54+
store.set_credentials(
55+
user_id="user_456",
56+
security_scheme_name="inbound_auth",
57+
credential="bearer_token_abc",
58+
)
59+
60+
# Get credentials
61+
context = ClientCallContext(state={"userId": "user_456"})
62+
token = await store.get_credentials(
63+
security_scheme_name="inbound_auth",
64+
context=context,
65+
)
66+
67+
assert token == "bearer_token_abc"
68+
69+
@pytest.mark.asyncio
70+
async def test_fallback_to_user_id(self):
71+
"""Test fallback from session ID to user ID."""
72+
store = VeCredentialStore()
73+
74+
# Set credentials by user ID
75+
store.set_credentials(
76+
user_id="user_789",
77+
security_scheme_name="inbound_auth",
78+
credential="bearer_token_fallback",
79+
)
80+
81+
# Try to get with session ID (should fail) then fallback to user ID
82+
context = ClientCallContext(
83+
state={"sessionId": "nonexistent_session", "userId": "user_789"}
84+
)
85+
token = await store.get_credentials(
86+
security_scheme_name="inbound_auth",
87+
context=context,
88+
)
89+
90+
assert token == "bearer_token_fallback"
91+
92+
@pytest.mark.asyncio
93+
async def test_user_id_priority(self):
94+
"""Test that user ID takes priority over user ID."""
95+
store = VeCredentialStore()
96+
97+
# Set credentials for both session and user
98+
store.set_credentials(
99+
session_id="session_priority",
100+
security_scheme_name="inbound_auth",
101+
credential="session_token",
102+
)
103+
store.set_credentials(
104+
user_id="user_priority",
105+
security_scheme_name="inbound_auth",
106+
credential="user_token",
107+
)
108+
109+
# Get with both session ID and user ID - should return session token
110+
context = ClientCallContext(
111+
state={"sessionId": "session_priority", "userId": "user_priority"}
112+
)
113+
token = await store.get_credentials(
114+
security_scheme_name="inbound_auth",
115+
context=context,
116+
)
117+
118+
assert token == "user_token"
119+
120+
@pytest.mark.asyncio
121+
async def test_async_set_credentials(self):
122+
"""Test async version of set_credentials."""
123+
store = VeCredentialStore()
124+
125+
# Set credentials asynchronously
126+
await store.set_credentials_async(
127+
session_id="async_session",
128+
security_scheme_name="inbound_auth",
129+
credential="async_token",
130+
)
131+
132+
# Get credentials
133+
context = ClientCallContext(state={"sessionId": "async_session"})
134+
token = await store.get_credentials(
135+
security_scheme_name="inbound_auth",
136+
context=context,
137+
)
138+
139+
assert token == "async_token"
140+
141+
@pytest.mark.asyncio
142+
async def test_multiple_security_schemes(self):
143+
"""Test storing multiple security schemes for the same session."""
144+
store = VeCredentialStore()
145+
146+
# Set multiple credentials for the same session
147+
store.set_credentials(
148+
session_id="multi_session",
149+
security_scheme_name="inbound_auth",
150+
credential="inbound_token",
151+
)
152+
store.set_credentials(
153+
session_id="multi_session",
154+
security_scheme_name="outbound_auth",
155+
credential="outbound_token",
156+
)
157+
158+
# Get both credentials
159+
context = ClientCallContext(state={"sessionId": "multi_session"})
160+
161+
inbound_token = await store.get_credentials(
162+
security_scheme_name="inbound_auth",
163+
context=context,
164+
)
165+
outbound_token = await store.get_credentials(
166+
security_scheme_name="outbound_auth",
167+
context=context,
168+
)
169+
170+
assert inbound_token == "inbound_token"
171+
assert outbound_token == "outbound_token"
172+
173+
@pytest.mark.asyncio
174+
async def test_clear_specific_context(self):
175+
"""Test clearing credentials for a specific context."""
176+
store = VeCredentialStore()
177+
178+
# Set credentials for two sessions
179+
store.set_credentials(
180+
session_id="session_1",
181+
security_scheme_name="inbound_auth",
182+
credential="token_1",
183+
)
184+
store.set_credentials(
185+
session_id="session_2",
186+
security_scheme_name="inbound_auth",
187+
credential="token_2",
188+
)
189+
190+
# Clear session_1
191+
store.clear("session_1")
192+
193+
# Verify session_1 is cleared
194+
context1 = ClientCallContext(state={"sessionId": "session_1"})
195+
token1 = await store.get_credentials(
196+
security_scheme_name="inbound_auth",
197+
context=context1,
198+
)
199+
assert token1 is None
200+
201+
# Verify session_2 still exists
202+
context2 = ClientCallContext(state={"sessionId": "session_2"})
203+
token2 = await store.get_credentials(
204+
security_scheme_name="inbound_auth",
205+
context=context2,
206+
)
207+
assert token2 == "token_2"
208+
209+
@pytest.mark.asyncio
210+
async def test_clear_all(self):
211+
"""Test clearing all credentials."""
212+
store = VeCredentialStore()
213+
214+
# Set multiple credentials
215+
store.set_credentials(
216+
session_id="session_a",
217+
security_scheme_name="inbound_auth",
218+
credential="token_a",
219+
)
220+
store.set_credentials(
221+
user_id="user_b",
222+
security_scheme_name="inbound_auth",
223+
credential="token_b",
224+
)
225+
226+
# Clear all
227+
store.clear()
228+
229+
# Verify all are cleared
230+
context_a = ClientCallContext(state={"sessionId": "session_a"})
231+
token_a = await store.get_credentials(
232+
security_scheme_name="inbound_auth",
233+
context=context_a,
234+
)
235+
assert token_a is None
236+
237+
context_b = ClientCallContext(state={"userId": "user_b"})
238+
token_b = await store.get_credentials(
239+
security_scheme_name="inbound_auth",
240+
context=context_b,
241+
)
242+
assert token_b is None
243+
244+
@pytest.mark.asyncio
245+
async def test_no_context(self):
246+
"""Test behavior when no context is provided."""
247+
store = VeCredentialStore()
248+
249+
token = await store.get_credentials(
250+
security_scheme_name="inbound_auth",
251+
context=None,
252+
)
253+
254+
assert token is None
255+
256+
@pytest.mark.asyncio
257+
async def test_missing_session_and_user(self):
258+
"""Test behavior when neither session ID nor user ID is in context."""
259+
store = VeCredentialStore()
260+
261+
context = ClientCallContext(state={"someOtherKey": "value"})
262+
token = await store.get_credentials(
263+
security_scheme_name="inbound_auth",
264+
context=context,
265+
)
266+
267+
assert token is None
268+
269+
def test_set_credentials_without_ids_raises_error(self):
270+
"""Test that set_credentials raises error when neither session_id nor user_id is provided."""
271+
store = VeCredentialStore()
272+
273+
with pytest.raises(
274+
ValueError, match="Either session_id or user_id must be provided"
275+
):
276+
store.set_credentials(
277+
security_scheme_name="inbound_auth",
278+
credential="token",
279+
)

0 commit comments

Comments
 (0)