diff --git a/tests/a2a/test_registry_client.py b/tests/a2a/test_registry_client.py index 32cf0ad5b..d3ed8d448 100644 --- a/tests/a2a/test_registry_client.py +++ b/tests/a2a/test_registry_client.py @@ -108,6 +108,86 @@ def _oauth_agent_card() -> dict: } +def _agent_card_v1_managed_api_key() -> dict: + return { + "name": "Weather-A2A-Agent", + "description": "Weather agent", + "version": "1.0.0", + "supportedInterfaces": [ + { + "url": " `https://example.test/a2a/v1` ", + "protocolBinding": "JSONRPC", + "protocolVersion": "1.0", + } + ], + "capabilities": { + "streaming": False, + "pushNotifications": False, + "extendedAgentCard": False, + "extensions": [ + { + "params": { + "credentialProviderName": "sycTest", + "poolName": "default", + } + } + ], + }, + "securitySchemes": { + "apiKeyAuth": { + "apiKeySecurityScheme": { + "description": "ApiKey client credentials", + "name": "Authorization", + "location": "Header", + } + } + }, + "securityRequirements": [ + { + "schemes": { + "apiKeyAuth": { + "list": [], + } + } + } + ], + "defaultInputModes": ["text/plain"], + "defaultOutputModes": ["text/plain"], + "skills": [ + { + "id": "weather", + "name": "Weather", + "description": "Query weather", + "tags": ["weather"], + } + ], + } + + +def _api_key_credential_provider_response() -> dict: + return { + "ResponseMetadata": { + "RequestId": "credential-provider-req", + "Action": "GetApiKeyCredentialProvider", + "Version": "2025-10-30", + "Service": "id", + "Region": "cn-beijing", + }, + "Result": { + "Name": "sycTest", + "PoolName": "default", + "ApiKeyMetadata": [ + { + "Location": "Header", + "ParameterName": "Authorization", + "Prefix": "Bearer", + } + ], + "ApiKey": "123456", + }, + } + + @patch.dict( "os.environ", { @@ -251,6 +331,76 @@ def test_create_task_gets_agent_and_sends_message(post: Mock): assert "Authorization" 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_supports_v1_agent_card_with_managed_api_key(post: Mock): + card = _agent_card_v1_managed_api_key() + post.side_effect = [ + _mock_response( + { + "ResponseMetadata": {"RequestId": "get-req"}, + "Result": { + "Id": "agent-id", + "Status": "running", + "AgentCard": json.dumps(card), + }, + } + ), + _mock_response(_api_key_credential_provider_response()), + _mock_response( + { + "result": { + "kind": "message", + "parts": [{"kind": "text", "text": "今天北京晴。"}], + } + } + ), + ] + + result = create_task( + "Weather-A2A-Agent", + "北京天气", + config=AgentKitA2ARegistryConfig( + space_id="space-test", + endpoint="https://agentkit.cn-shanghai.volcengineapi.com/", + region="cn-shanghai", + ), + ) + + assert result["outcome"] == "success" + assert post.call_args_list[0].kwargs["params"]["Action"] == "GetA2aAgent" + assert post.call_args_list[1].args[0] == ( + "https://id.cn-shanghai.volcengineapi.com/" + ) + assert post.call_args_list[1].kwargs["params"] == { + "Action": "GetApiKeyCredentialProvider", + "Version": "2025-10-30", + } + assert json.loads(post.call_args_list[1].kwargs["data"].decode("utf-8")) == { + "Name": "sycTest", + "PoolName": "default", + } + assert ( + "/cn-shanghai/id/request" + in post.call_args_list[1].kwargs["headers"]["Authorization"] + ) + assert post.call_args_list[2].args[0] == "https://example.test/a2a/v1" + assert post.call_args_list[2].kwargs["headers"]["Authorization"] == ( + "Bearer 123456" + ) + + serialized = json.dumps(result, ensure_ascii=False) + assert "123456" not in serialized + assert "Authorization" not in serialized + + @patch.dict( "os.environ", { @@ -553,6 +703,58 @@ def test_poll_task_returns_terminal_without_sleep(post: Mock, sleep: Mock): sleep.assert_not_called() +@patch.dict( + "os.environ", + { + "AGENTKIT_ACCESS_KEY": "ak-test", + "AGENTKIT_SECRET_KEY": "sk-test", + }, + clear=False, +) +@patch("veadk.a2a.registry_client.time.sleep") +@patch("veadk.a2a.registry_client.requests.post") +def test_poll_task_supports_v1_agent_card_with_managed_api_key(post: Mock, sleep: Mock): + card = _agent_card_v1_managed_api_key() + post.side_effect = [ + _mock_response( + { + "ResponseMetadata": {"RequestId": "get-req"}, + "Result": { + "Id": "agent-id", + "Status": "running", + "AgentCard": json.dumps(card), + }, + } + ), + _mock_response(_api_key_credential_provider_response()), + _mock_response( + { + "result": { + "id": "task-1", + "status": {"state": "completed"}, + } + } + ), + ] + + result = poll_task( + "Weather-A2A-Agent", + "task-1", + config=AgentKitA2ARegistryConfig(space_id="space-test"), + ) + + assert result["outcome"] == "success" + assert post.call_args_list[1].kwargs["params"]["Action"] == ( + "GetApiKeyCredentialProvider" + ) + assert post.call_args_list[2].args[0] == "https://example.test/a2a/v1" + assert post.call_args_list[2].kwargs["headers"]["Authorization"] == ( + "Bearer 123456" + ) + assert post.call_args_list[2].kwargs["json"]["method"] == "tasks/get" + sleep.assert_not_called() + + @patch.dict( "os.environ", { @@ -741,6 +943,64 @@ def test_build_remote_a2a_agent_tools_searches_gets_and_sends(post: Mock): assert post.call_args_list[2].args[0] == "https://example.test/a2a" +@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_dynamic_remote_a2a_tool_supports_v1_agent_card_with_managed_api_key( + post: Mock, +): + card = _agent_card_v1_managed_api_key() + post.side_effect = [ + _mock_response( + { + "ResponseMetadata": {"RequestId": "search-req"}, + "Result": {"AgentCards": [json.dumps(card)], "TotalCount": 1}, + } + ), + _mock_response( + { + "ResponseMetadata": {"RequestId": "get-req"}, + "Result": { + "Id": "agent-id", + "Status": "running", + "AgentCard": json.dumps(card), + }, + } + ), + _mock_response(_api_key_credential_provider_response()), + _mock_response( + { + "result": { + "kind": "message", + "parts": [{"kind": "text", "text": "今天北京晴。"}], + } + } + ), + ] + + tools = build_remote_a2a_agent_tools( + "北京天气", AgentKitA2ARegistryConfig(space_id="space-test") + ) + result = tools[0](input="北京天气") + + assert result["outcome"] == "success" + assert post.call_args_list[0].kwargs["params"]["Action"] == "SearchAgentCards" + assert post.call_args_list[1].kwargs["params"]["Action"] == "GetA2aAgent" + assert post.call_args_list[2].kwargs["params"]["Action"] == ( + "GetApiKeyCredentialProvider" + ) + assert post.call_args_list[3].args[0] == "https://example.test/a2a/v1" + assert post.call_args_list[3].kwargs["headers"]["Authorization"] == ( + "Bearer 123456" + ) + + @patch.dict( "os.environ", { diff --git a/veadk/a2a/registry_client.py b/veadk/a2a/registry_client.py index d0a384dba..d87c55221 100644 --- a/veadk/a2a/registry_client.py +++ b/veadk/a2a/registry_client.py @@ -37,6 +37,7 @@ DEFAULT_VERSION = "2025-10-30" IDENTITY_VERSION = "2025-10-30" DEFAULT_SERVICE_NAME = "agentkit" +IDENTITY_SERVICE_NAME = "id" DEFAULT_REGION = "cn-beijing" DEFAULT_TOP_K = 3 DEFAULT_TIMEOUT_MS = 60000 @@ -79,6 +80,12 @@ class _RegistryCredentials: session_token: str = "" +@dataclass(frozen=True) +class _ApiKeyCredentialProviderRef: + name: str + pool_name: str + + def registry_config_from_env() -> AgentKitA2ARegistryConfig: """Read AgentKit A2A registry config from Harness-compatible env vars.""" @@ -401,7 +408,7 @@ def _identity_post( return _signed_openapi_post( config=config, endpoint=_identity_endpoint(config), - service_name="id", + service_name=IDENTITY_SERVICE_NAME, action=action, version=IDENTITY_VERSION, body=body, @@ -559,9 +566,10 @@ def _get_a2a_agent( card = _parse_json_object( result.get("AgentCard"), "AGENT_CARD_PARSE_FAILED", "Result.AgentCard" ) - if "url" in card: - card["url"] = _clean_config_url(card.get("url", "")) - if not card.get("url"): + card_url = _agent_card_url(card) + if card_url: + card["url"] = card_url + else: raise RegistryError( "AGENT_URL_MISSING", f"Agent {agent_name} AgentCard missing url" ) @@ -885,6 +893,118 @@ def _sanitize_get_agent_result( } +def _agent_card_url(card: dict[str, Any]) -> str: + url = _clean_config_url(card.get("url", "")) + if url: + return url + + interfaces = card.get("supportedInterfaces") or [] + if not isinstance(interfaces, list): + return "" + + candidates: list[tuple[int, str]] = [] + for interface in interfaces: + if not isinstance(interface, dict): + continue + interface_url = _clean_config_url(interface.get("url", "")) + if not interface_url: + continue + + score = 0 + protocol_binding = str(interface.get("protocolBinding") or "").lower() + if protocol_binding in {"jsonrpc", "json-rpc"}: + score += 2 + protocol_version = str(interface.get("protocolVersion") or "").strip() + if protocol_version in {"1.0", "1.0.0"}: + score += 1 + candidates.append((score, interface_url)) + + if not candidates: + return "" + return max(candidates, key=lambda candidate: candidate[0])[1] + + +def _api_key_credential_provider_ref( + card: dict[str, Any], +) -> _ApiKeyCredentialProviderRef | None: + capabilities = card.get("capabilities") or {} + if not isinstance(capabilities, dict): + return None + + extensions = capabilities.get("extensions") or [] + if not isinstance(extensions, list): + return None + + for extension in extensions: + if not isinstance(extension, dict): + continue + params = extension.get("params") or {} + if not isinstance(params, dict): + continue + + provider_name = _clean_config_url( + params.get("credentialProviderName") + or params.get("CredentialProviderName") + or "" + ) + pool_name = _clean_config_url(params.get("poolName") or params.get("PoolName")) + if provider_name and pool_name: + return _ApiKeyCredentialProviderRef( + name=provider_name, + pool_name=pool_name, + ) + if provider_name or pool_name: + raise RegistryError( + "AGENT_API_KEY_PROVIDER_INVALID", + "AgentCard credential provider extension requires credentialProviderName and poolName", + ) + + return None + + +def _managed_api_key_headers( + provider_ref: _ApiKeyCredentialProviderRef, + config: AgentKitA2ARegistryConfig, +) -> dict[str, str]: + response, _ = _identity_post( + config, + "GetApiKeyCredentialProvider", + {"Name": provider_ref.name, "PoolName": provider_ref.pool_name}, + ) + result = response.get("Result") or {} + api_key = str(result.get("ApiKey") or "").strip() + if not api_key: + raise RegistryError( + "AGENT_API_KEY_PROVIDER_INVALID", + "GetApiKeyCredentialProvider response missing ApiKey", + ) + + metadata = result.get("ApiKeyMetadata") or [] + if not isinstance(metadata, list): + metadata = [] + + headers: dict[str, str] = {} + for item in metadata: + if not isinstance(item, dict): + continue + location = str(item.get("Location") or item.get("location") or "").lower() + parameter_name = str( + item.get("ParameterName") or item.get("parameterName") or "" + ).strip() + if location != "header" or not parameter_name: + continue + + prefix = str(item.get("Prefix") or item.get("prefix") or "").strip() + headers[parameter_name] = f"{prefix} {api_key}" if prefix else api_key + + if not headers: + raise RegistryError( + "AGENT_API_KEY_PROVIDER_INVALID", + "GetApiKeyCredentialProvider response missing header ApiKeyMetadata", + ) + return headers + + def _agent_auth_headers( card: dict[str, Any], config: AgentKitA2ARegistryConfig | None = None, @@ -923,6 +1043,11 @@ def _agent_auth_headers( + _oauth2_client_credentials_token(scheme, resolved_config) ) + if not headers: + provider_ref = _api_key_credential_provider_ref(card) + if provider_ref is not None: + headers.update(_managed_api_key_headers(provider_ref, resolved_config)) + tip_token = resolved_config.upstream_tip_token if tip_token: headers[VE_TIP_TOKEN_HEADER] = tip_token