diff --git a/tests/a2a/test_registry_client.py b/tests/a2a/test_registry_client.py index a6531136..2beb03c5 100644 --- a/tests/a2a/test_registry_client.py +++ b/tests/a2a/test_registry_client.py @@ -27,6 +27,7 @@ _volc_sign_v4, create_task, poll_task, + registry_authorization_from_headers, registry_tip_token_from_headers, search_agent_cards, truncate_utf8_bytes, @@ -306,6 +307,130 @@ def test_create_task_gets_oauth_agent_token_and_sends_message(post: Mock): assert "m2m-client-secret" not in serialized +@patch.dict( + "os.environ", + { + "AGENTKIT_ACCESS_KEY": "ak-test", + "AGENTKIT_SECRET_KEY": "sk-test", + }, + clear=False, +) +@patch("veadk.a2a.registry_client.requests.post") +def test_create_task_prefers_upstream_authorization_for_oauth_agent(post: Mock): + _OAUTH_TOKEN_CACHE.clear() + card = _oauth_agent_card() + post.side_effect = [ + _mock_response( + { + "ResponseMetadata": {"RequestId": "get-req"}, + "Result": { + "Id": "agent-id", + "Status": "running", + "AgentCard": json.dumps(card), + }, + } + ), + _mock_response( + { + "result": { + "kind": "message", + "parts": [{"kind": "text", "text": "用户 JWT 已通过。"}], + } + } + ), + ] + + result = create_task( + "Finance Policy Remote Agent", + "这笔支出是否需要审批？", + config=AgentKitA2ARegistryConfig( + space_id="space-test", + endpoint="https://open.volcengineapi.com/", + upstream_authorization="Bearer user-jwt", + ), + ) + + assert result["outcome"] == "success" + assert post.call_count == 2 + assert post.call_args_list[1].kwargs["headers"]["Authorization"] == ( + "Bearer user-jwt" + ) + + +@patch.dict( + "os.environ", + { + "AGENTKIT_ACCESS_KEY": "ak-test", + "AGENTKIT_SECRET_KEY": "sk-test", + }, + clear=False, +) +@patch("veadk.a2a.registry_client.requests.post") +def test_create_task_falls_back_to_m2m_when_upstream_authorization_is_401( + post: Mock, +): + _OAUTH_TOKEN_CACHE.clear() + card = _oauth_agent_card() + unauthorized = _mock_response({"detail": "unauthorized"}, status_code=401) + unauthorized.raise_for_status.side_effect = requests.HTTPError( + "401 Client Error", response=unauthorized + ) + post.side_effect = [ + _mock_response( + { + "ResponseMetadata": {"RequestId": "get-req"}, + "Result": { + "Id": "agent-id", + "Status": "running", + "AgentCard": json.dumps(card), + }, + } + ), + unauthorized, + _mock_response( + { + "ResponseMetadata": {"RequestId": "list-client-req"}, + "Result": {"Data": [{"Uid": "m2m-client-id"}]}, + } + ), + _mock_response( + { + "ResponseMetadata": {"RequestId": "get-client-req"}, + "Result": {"ClientSecret": "m2m-client-secret"}, + } + ), + _mock_response({"access_token": "oauth-access-token", "expires_in": 3600}), + _mock_response( + { + "result": { + "kind": "message", + "parts": [{"kind": "text", "text": "M2M fallback 已通过。"}], + } + } + ), + ] + + result = create_task( + "Finance Policy Remote Agent", + "这笔支出是否需要审批？", + config=AgentKitA2ARegistryConfig( + space_id="space-test", + endpoint="https://open.volcengineapi.com/", + upstream_authorization="Bearer user-jwt", + ), + ) + + assert result["outcome"] == "success" + assert result["response"]["text"] == "M2M fallback 已通过。" + assert post.call_args_list[1].kwargs["headers"]["Authorization"] == ( + "Bearer user-jwt" + ) + assert post.call_args_list[2].kwargs["params"]["Action"] == "ListUserPoolClients" + assert post.call_args_list[5].kwargs["headers"]["Authorization"] == ( + "Bearer oauth-access-token" + ) + + @patch.dict( "os.environ", { @@ -690,6 +815,16 @@ def test_agent_auth_headers_extracts_api_key_header(): } +@patch.dict("os.environ", {}, clear=True) +def test_agent_auth_headers_prefers_upstream_authorization_for_oauth(): + assert _agent_auth_headers( + _oauth_agent_card(), + AgentKitA2ARegistryConfig( + space_id="space-test", upstream_authorization="Bearer user-jwt" + ), + ) == {"Authorization": "Bearer user-jwt"} + + @patch.dict("os.environ", {}, clear=True) def test_agent_auth_headers_rejects_unusable_security(): with pytest.raises(RegistryError) as ctx: @@ -715,6 +850,13 @@ def test_registry_tip_token_from_headers_is_case_insensitive(): ) +def test_registry_authorization_from_headers_accepts_bearer_only(): + assert registry_authorization_from_headers( + {"Authorization": " bearer user-jwt "} + ) == ("Bearer user-jwt") + assert registry_authorization_from_headers({"Authorization": "Basic abc"}) == "" + + def test_agentkit_http_error_uses_safe_diagnostics(): response = _mock_response( { diff --git a/tests/cloud/test_harness_app_contract.py b/tests/cloud/test_harness_app_contract.py index b1690d7c..8e2f2677 100644 --- a/tests/cloud/test_harness_app_contract.py +++ b/tests/cloud/test_harness_app_contract.py @@ -201,14 +201,18 @@ def test_registry_dynamic_tools_are_added_per_run(self): assert "has_a2a_registry_config(self.agent)" in app_source assert "spawn_harness_run_agent(" in app_source - def test_registry_tip_token_is_bound_to_run_agent_config(self): + def test_registry_request_auth_is_bound_to_run_agent_config(self): registry_source = Path("veadk/a2a/registry_client.py").read_text() utils_source = Path("veadk/cloud/harness_app/utils.py").read_text() app_source = Path("veadk/cloud/harness_app/app.py").read_text() - assert "_apply_registry_tip_token(" in utils_source + assert "_apply_registry_request_auth(" in utils_source assert "upstream_tip_token=cleaned_tip_token" in utils_source + assert "upstream_authorization=cleaned_authorization" in utils_source assert "registry_tip_token=tip_token" in app_source + assert "registry_authorization=auth_header" in app_source + assert "registry_authorization_from_headers" in app_source + assert "registry_authorization_from_headers(" in registry_source assert "ContextVar" not in registry_source assert "use_registry_tip_token" not in registry_source assert "use_registry_tip_token" not in app_source diff --git a/veadk/a2a/registry_client.py b/veadk/a2a/registry_client.py index 2508899a..79f2ac82 100644 --- a/veadk/a2a/registry_client.py +++ b/veadk/a2a/registry_client.py @@ -69,6 +69,7 @@ class AgentKitA2ARegistryConfig: timeout_ms: int = DEFAULT_TIMEOUT_MS poll_interval_ms: int = DEFAULT_POLL_INTERVAL_MS upstream_tip_token: str = "" + upstream_authorization: str = "" @dataclass(frozen=True) @@ -120,6 +121,21 @@ def registry_tip_token_from_headers(headers: Mapping[str, str]) -> str: return normalized.get(VE_TIP_TOKEN_HEADER.lower(), "").strip() +def registry_authorization_from_headers(headers: Mapping[str, str]) -> str: + """Extract a bearer Authorization header for downstream A2A calls.""" + + normalized = {str(key).lower(): str(value) for key, value in headers.items()} + for header_name in ( + "authorization", + "x-forwarded-authorization", + "x-original-authorization", + ): + value = _normalize_bearer_authorization(normalized.get(header_name, "")) + if value: + return value + return "" + + def search_agent_cards( prompt: str, top_k: int | None = None, @@ -273,6 +289,9 @@ def _resolve_config( upstream_tip_token=( config.upstream_tip_token or env_config.upstream_tip_token ).strip(), + upstream_authorization=_normalize_bearer_authorization( + config.upstream_authorization or env_config.upstream_authorization + ), ) @@ -557,15 +576,25 @@ def _send_message( if task_id: message["taskId"] = task_id + headers: dict[str, str] = {} try: + headers = _agent_auth_headers(card, config) return _a2a_jsonrpc( card["url"], "message/send", {"message": message, "configuration": {"blocking": False}}, - _agent_auth_headers(card, config), + headers, config, ) except RegistryError as exc: + if _should_retry_with_m2m(exc, card, config, headers): + return _a2a_jsonrpc( + card["url"], + "message/send", + {"message": message, "configuration": {"blocking": False}}, + _agent_auth_headers(card, config, force_m2m=True), + config, + ) if exc.code in { "A2A_HTTP_FAILED", "A2A_RESPONSE_PARSE_FAILED", @@ -586,13 +615,25 @@ def _poll_card( started: float | None = None, ) -> dict[str, Any]: started = started or time.monotonic() - a2a_result = _a2a_jsonrpc( - card["url"], - "tasks/get", - {"id": task_id.strip(), "historyLength": max(0, int(history_length))}, - _agent_auth_headers(card, config), - config, - ) + headers = _agent_auth_headers(card, config) + try: + a2a_result = _a2a_jsonrpc( + card["url"], + "tasks/get", + {"id": task_id.strip(), "historyLength": max(0, int(history_length))}, + headers, + config, + ) + except RegistryError as exc: + if not _should_retry_with_m2m(exc, card, config, headers): + raise + a2a_result = _a2a_jsonrpc( + card["url"], + "tasks/get", + {"id": task_id.strip(), "historyLength": max(0, int(history_length))}, + _agent_auth_headers(card, config, force_m2m=True), + config, + ) state = _task_state(a2a_result) is_terminal = state in TERMINAL_STATES payload: dict[str, Any] = { @@ -838,12 +879,18 @@ def _sanitize_get_agent_result( def _agent_auth_headers( - card: dict[str, Any], config: AgentKitA2ARegistryConfig | None = None + card: dict[str, Any], + config: AgentKitA2ARegistryConfig | None = None, + *, + force_m2m: bool = False, ) -> dict[str, str]: resolved_config = _resolve_config(config) security = card.get("security") or [] schemes = card.get("securitySchemes") or {} headers: dict[str, str] = {} + upstream_authorization = _normalize_bearer_authorization( + resolved_config.upstream_authorization + ) for requirement in security: if not isinstance(requirement, dict): @@ -861,9 +908,13 @@ def _agent_auth_headers( if isinstance(token, str) and token: headers[header_name] = token elif scheme_type == "oauth2": - headers["Authorization"] = "Bearer " + _oauth2_client_credentials_token( - scheme, resolved_config - ) + if upstream_authorization and not force_m2m: + headers["Authorization"] = upstream_authorization + else: + headers["Authorization"] = ( + "Bearer " + + _oauth2_client_credentials_token(scheme, resolved_config) + ) tip_token = resolved_config.upstream_tip_token if tip_token: @@ -877,6 +928,38 @@ def _agent_auth_headers( return headers +def _should_retry_with_m2m( + exc: RegistryError, + card: dict[str, Any], + config: AgentKitA2ARegistryConfig | None, + headers: dict[str, str], +) -> bool: + if exc.code != "A2A_HTTP_FAILED" or exc.diagnostics.get("status_code") != 401: + return False + + resolved_config = _resolve_config(config) + upstream_authorization = _normalize_bearer_authorization( + resolved_config.upstream_authorization + ) + if not upstream_authorization: + return False + if headers.get("Authorization") != upstream_authorization: + return False + return _has_oauth2_security(card) + + +def _has_oauth2_security(card: dict[str, Any]) -> bool: + schemes = card.get("securitySchemes") or {} + for requirement in card.get("security") or []: + if not isinstance(requirement, dict): + continue + for scheme_name in requirement: + scheme = schemes.get(scheme_name) or {} + if str(scheme.get("type") or "").lower() == "oauth2": + return True + return False + + def _oauth2_client_credentials_token( scheme: dict[str, Any], config: AgentKitA2ARegistryConfig ) -> str: @@ -985,6 +1068,16 @@ def _clean_config_url(value: Any) -> str: return cleaned +def _normalize_bearer_authorization(value: Any) -> str: + cleaned = _clean_config_url(value) + if not cleaned: + return "" + parts = cleaned.split(None, 1) + if len(parts) == 2 and parts[0].lower() == "bearer" and parts[1].strip(): + return f"Bearer {parts[1].strip()}" + return "" + + def _user_pool_id_from_token_url(token_url: str) -> str: host = urlparse(token_url).netloc match = re.match(r"userpool-([^.]+)\.userpool\.auth\.id\.", host) diff --git a/veadk/cloud/harness_app/app.py b/veadk/cloud/harness_app/app.py index b858a539..e42ca569 100644 --- a/veadk/cloud/harness_app/app.py +++ b/veadk/cloud/harness_app/app.py @@ -62,7 +62,10 @@ from typing_extensions import override from veadk import Agent -from veadk.a2a.registry_client import registry_tip_token_from_headers +from veadk.a2a.registry_client import ( + registry_authorization_from_headers, + registry_tip_token_from_headers, +) from veadk.a2a.utils.agent_to_a2a import to_a2a from veadk.cloud.harness_app.agent import agent, short_term_memory from veadk.cloud.harness_app.harness_plugins import ( @@ -231,6 +234,7 @@ async def invoke_harness( try: tip_token = registry_tip_token_from_headers(http_request.headers) + auth_header = registry_authorization_from_headers(http_request.headers) header_plugins = build_harness_plugins_from_headers( http_request.headers ) @@ -272,6 +276,7 @@ async def invoke_harness( request.harness, download_dir=Path(work_dir), registry_tip_token=tip_token, + registry_authorization=auth_header, ) runner = Runner( agent=agent, @@ -291,6 +296,7 @@ async def invoke_harness( self.agent, request.prompt, registry_tip_token=tip_token, + registry_authorization=auth_header, ) else: run_agent = self.agent @@ -374,8 +380,10 @@ async def run_sse(req: HarnessRunAgentRequest, http_request: Request): # No override -> exactly ADK's default /run_sse. return await adk_run_sse(req) tip_token = registry_tip_token_from_headers(http_request.headers) + auth_header = registry_authorization_from_headers(http_request.headers) return StreamingResponse( - self._run_sse_events(req, tip_token), media_type="text/event-stream" + self._run_sse_events(req, tip_token, auth_header), + media_type="text/event-stream", ) # Move ours to the front so it wins (Starlette matches the first route), @@ -388,7 +396,12 @@ async def run_sse(req: HarnessRunAgentRequest, http_request: Request): routes.insert(0, routes.pop(i)) break - async def _run_sse_events(self, req: "HarnessRunAgentRequest", tip_token: str = ""): + async def _run_sse_events( + self, + req: "HarnessRunAgentRequest", + tip_token: str = "", + auth_header: str = "", + ): """Yield SSE ``data:`` lines for a run, spawning the agent on override.""" run_config = RunConfig( streaming_mode=StreamingMode.SSE if req.streaming else StreamingMode.NONE @@ -408,6 +421,7 @@ async def _run_sse_events(self, req: "HarnessRunAgentRequest", tip_token: str = req.harness, download_dir=Path(work_dir_ctx.name), registry_tip_token=tip_token, + registry_authorization=auth_header, ) except (SkillLoadError, ToolLoadError) as e: logger.error(f"Once-time override failed to load: {e}") @@ -418,6 +432,7 @@ async def _run_sse_events(self, req: "HarnessRunAgentRequest", tip_token: str = self.agent, prompt, registry_tip_token=tip_token, + registry_authorization=auth_header, ) else: agent = self.agent diff --git a/veadk/cloud/harness_app/utils.py b/veadk/cloud/harness_app/utils.py index f124e31f..cb0991c8 100644 --- a/veadk/cloud/harness_app/utils.py +++ b/veadk/cloud/harness_app/utils.py @@ -527,9 +527,12 @@ def _apply_registry_overrides( setattr(agent, _REGISTRY_CONFIG_ATTR, overridden_config) -def _apply_registry_tip_token(agent: Agent, tip_token: str = "") -> None: +def _apply_registry_request_auth( + agent: Agent, tip_token: str = "", authorization: str = "" +) -> None: cleaned_tip_token = (tip_token or "").strip() - if not cleaned_tip_token: + cleaned_authorization = (authorization or "").strip() + if not cleaned_tip_token and not cleaned_authorization: return from veadk.tools.builtin_tools.a2a_registry import build_a2a_registry_tools @@ -538,7 +541,11 @@ def _apply_registry_tip_token(agent: Agent, tip_token: str = "") -> None: if config is None: return - updated_config = replace(config, upstream_tip_token=cleaned_tip_token) + updated_config = replace( + config, + upstream_tip_token=cleaned_tip_token or config.upstream_tip_token, + upstream_authorization=cleaned_authorization or config.upstream_authorization, + ) _remove_a2a_registry_tools(agent) agent.tools.extend(build_a2a_registry_tools(updated_config)) setattr(agent, _REGISTRY_CONFIG_ATTR, updated_config) @@ -622,6 +629,7 @@ def spawn_harness_run_agent( overrides: HarnessOverrides | None = None, download_dir: Path | None = None, registry_tip_token: str = "", + registry_authorization: str = "", ) -> Agent: """Clone a harness agent for one run and attach per-turn dynamic tools.""" @@ -630,6 +638,6 @@ def spawn_harness_run_agent( else: cloned = base_agent.clone(update={}) - _apply_registry_tip_token(cloned, registry_tip_token) + _apply_registry_request_auth(cloned, registry_tip_token, registry_authorization) _add_dynamic_a2a_agent_tools(cloned, prompt) return cloned