From 88c12a9e5e4f375452027690d6f7472e4bc83c57 Mon Sep 17 00:00:00 2001 From: 2862282695gjh-afk <2862282695gjh@gmail.com> Date: Tue, 11 Aug 2026 10:32:41 +0800 Subject: [PATCH] fix: pass security credentials to external A2A sub-agents Co-Authored-By: Claude --- backend/agents/create_agent_info.py | 65 ++++++++ backend/database/a2a_agent_db.py | 3 + sdk/nexent/core/agents/a2a_agent_proxy.py | 3 + sdk/nexent/core/agents/agent_model.py | 7 +- test/backend/agents/test_create_agent_info.py | 144 ++++++++++++++++++ test/backend/database/test_a2a_agent_db.py | 18 +++ test/sdk/core/agents/test_a2a_agent_proxy.py | 16 ++ test/sdk/core/agents/test_agent_model.py | 20 ++- 8 files changed, 274 insertions(+), 2 deletions(-) diff --git a/backend/agents/create_agent_info.py b/backend/agents/create_agent_info.py index e97c3c2afe..b583c0a972 100644 --- a/backend/agents/create_agent_info.py +++ b/backend/agents/create_agent_info.py @@ -515,6 +515,70 @@ def _extract_url_from_card(raw_card: Optional[dict]) -> str: return raw_card.get("url", "") +def _resolve_scheme_field(scheme: dict, wrapper_key: str) -> Optional[dict]: + """Get a security scheme field from wrapper or flat format.""" + field = scheme.get(wrapper_key) + if isinstance(field, dict) and field: + return field + # Flat format fallback: scheme itself has the fields + if wrapper_key == "httpAuthSecurityScheme" and isinstance(scheme.get("scheme"), str): + return scheme if scheme["scheme"].strip() else None + if wrapper_key == "apiKeySecurityScheme" and scheme.get("name") and scheme.get("location"): + return scheme + return None + + +def _build_auth_header_for_scheme(scheme: dict, credential: str) -> Optional[tuple]: + """Build a single (header_name, header_value) pair from a security scheme. + + Supports httpAuth (Bearer/basic) and apiKey (header location). + """ + # HTTP auth (bearer, basic) + http_auth = _resolve_scheme_field(scheme, "httpAuthSecurityScheme") + if http_auth: + auth_scheme = http_auth.get("scheme", "") + if http_auth.get("bearerFormat", "").lower() == "jwt": + auth_scheme = "Bearer" + return ("Authorization", f"{auth_scheme} {credential}") if auth_scheme else None + + # API key in header + api_key = _resolve_scheme_field(scheme, "apiKeySecurityScheme") + if api_key: + location = (api_key.get("location") or "").lower() + name = api_key.get("name") + if location == "header" and name: + return (name, credential) + + return None + + +def _collect_auth_headers(requirements, schemes, credentials): + """Collect (header_name, value) pairs from security requirements.""" + pairs = [] + for req in requirements: + if not isinstance(req, dict): + continue + for scheme_id in req.get("schemes", {}): + credential = credentials.get(scheme_id) + scheme = schemes.get(scheme_id) + if credential and isinstance(scheme, dict): + pair = _build_auth_header_for_scheme(scheme, credential) + if pair: + pairs.append(pair) + return pairs + + +def _build_security_headers(agent: dict) -> dict: + """Build auth headers from securitySchemes + security_credentials.""" + schemes = agent.get("security_schemes") or {} + requirements = agent.get("security_requirements") or [] + credentials = agent.get("security_credentials") or {} + if not requirements or not credentials: + return {} + return dict(_collect_auth_headers(requirements, schemes, credentials)) + + + def _build_external_agent_config(agent: dict, agent_url: str) -> ExternalA2AAgentConfig: """Build an ExternalA2AAgentConfig from agent data.""" return ExternalA2AAgentConfig( @@ -528,6 +592,7 @@ def _build_external_agent_config(agent: dict, agent_url: str) -> ExternalA2AAgen protocol_type=agent.get("protocol_type", PROTOCOL_JSONRPC), timeout=300.0, raw_card=agent.get("raw_card"), + custom_headers=_build_security_headers(agent) or None, ) diff --git a/backend/database/a2a_agent_db.py b/backend/database/a2a_agent_db.py index 982be32cd2..f91b049149 100644 --- a/backend/database/a2a_agent_db.py +++ b/backend/database/a2a_agent_db.py @@ -957,6 +957,9 @@ def query_external_sub_agents( "streaming": agent.streaming, "supported_interfaces": agent.supported_interfaces, "raw_card": agent.raw_card, + "security_schemes": agent.security_schemes, + "security_requirements": agent.security_requirements, + "security_credentials": agent.security_credentials, "is_enabled": relation.is_enabled, } for relation, agent in results diff --git a/sdk/nexent/core/agents/a2a_agent_proxy.py b/sdk/nexent/core/agents/a2a_agent_proxy.py index bd7651dd01..7ab2e2aa02 100644 --- a/sdk/nexent/core/agents/a2a_agent_proxy.py +++ b/sdk/nexent/core/agents/a2a_agent_proxy.py @@ -34,6 +34,7 @@ class A2AAgentInfo: protocol_type: str = PROTOCOL_JSONRPC timeout: float = 300.0 raw_card: Optional[Dict[str, Any]] = None + custom_headers: Optional[Dict[str, str]] = None def get_protocol_type(self) -> str: """Get the protocol type for calling this agent. @@ -120,6 +121,8 @@ def _build_headers(self) -> Dict[str, str]: } if self.agent_info.api_key: headers["Authorization"] = f"Bearer {self.agent_info.api_key}" + if self.agent_info.custom_headers: + headers.update(self.agent_info.custom_headers) return headers def _build_message_payload( diff --git a/sdk/nexent/core/agents/agent_model.py b/sdk/nexent/core/agents/agent_model.py index 8d68992682..2c34d4bece 100644 --- a/sdk/nexent/core/agents/agent_model.py +++ b/sdk/nexent/core/agents/agent_model.py @@ -444,6 +444,10 @@ class ExternalA2AAgentConfig(BaseModel): description="Raw Agent Card containing skills and capabilities", default=None ) + custom_headers: Optional[Dict[str, str]] = Field( + description="Pre-built auth headers from securitySchemes + credentials", + default=None + ) def model_post_init(self, __context) -> None: """Auto-enhance description with skills info from raw_card.""" @@ -496,7 +500,8 @@ def to_a2a_agent_info(self) -> "A2AAgentInfo": protocol_version=self.protocol_version, protocol_type=self.protocol_type, timeout=self.timeout, - raw_card=self.raw_card + raw_card=self.raw_card, + custom_headers=self.custom_headers ) diff --git a/test/backend/agents/test_create_agent_info.py b/test/backend/agents/test_create_agent_info.py index c142990fad..f07dbd5a71 100644 --- a/test/backend/agents/test_create_agent_info.py +++ b/test/backend/agents/test_create_agent_info.py @@ -502,6 +502,9 @@ class MockUncertaintyReserveBasisUnknown(Exception): _get_skill_script_tools, _extract_url_from_card, _build_external_agent_config, + _build_security_headers, + _resolve_scheme_field, + _build_auth_header_for_scheme, _get_external_a2a_agents, _build_internal_s3_url, _format_minio_files_for_content, @@ -5199,6 +5202,7 @@ def test_build_external_agent_config_basic(self): protocol_type="JSONRPC", timeout=300.0, raw_card=None, + custom_headers=None, ) assert result == MockConfig.return_value @@ -5223,6 +5227,7 @@ def test_build_external_agent_config_defaults(self): protocol_type="JSONRPC", timeout=300.0, raw_card=None, + custom_headers=None, ) assert result == MockConfig.return_value @@ -7343,3 +7348,142 @@ def capture_and_return(**kwargs): assert "langchain_tool" in mock_tc_instance.metadata assert mock_tc_instance.metadata["langchain_tool"] is matching_langchain_tool assert "allowed_kds_set" in mock_tc_instance.metadata + + +class TestBuildSecurityHeaders: + """Tests for _build_security_headers and related functions.""" + + def test_build_security_headers_apikey(self): + """Two apiKey-header schemes.""" + agent = { + "security_schemes": { + "scheme_a": {"apiKeySecurityScheme": {"name": "X-Custom-Id", "location": "header"}}, + "scheme_b": {"apiKeySecurityScheme": {"name": "X-Custom-Key", "location": "header"}}, + }, + "security_requirements": [{"schemes": {"scheme_a": {}, "scheme_b": {}}}], + "security_credentials": {"scheme_a": "id_value", "scheme_b": "key_value"}, + } + assert _build_security_headers(agent) == {"X-Custom-Id": "id_value", "X-Custom-Key": "key_value"} + + def test_build_security_headers_http_bearer(self): + """HTTP Bearer JWT.""" + agent = { + "security_schemes": {"jwt": {"httpAuthSecurityScheme": {"scheme": "bearer", "bearerFormat": "JWT"}}}, + "security_requirements": [{"schemes": {"jwt": {}}}], + "security_credentials": {"jwt": "token123"}, + } + assert _build_security_headers(agent) == {"Authorization": "Bearer token123"} + + def test_build_security_headers_no_credentials(self): + """No credentials -> empty.""" + agent = { + "security_schemes": {"k": {"apiKeySecurityScheme": {"name": "X-Key", "location": "header"}}}, + "security_requirements": [{"schemes": {"k": {}}}], + "security_credentials": {}, + } + assert _build_security_headers(agent) == {} + + def test_build_security_headers_mixed(self): + """Mixed apiKey + httpAuth.""" + agent = { + "security_schemes": { + "k": {"apiKeySecurityScheme": {"name": "X-Custom", "location": "header"}}, + "j": {"httpAuthSecurityScheme": {"scheme": "bearer"}}, + }, + "security_requirements": [{"schemes": {"k": {}, "j": {}}}], + "security_credentials": {"k": "key-val", "j": "jwt-val"}, + } + assert _build_security_headers(agent) == {"X-Custom": "key-val", "Authorization": "bearer jwt-val"} + + def test_resolve_scheme_field_wrapper(self): + scheme = {"apiKeySecurityScheme": {"name": "X-Key", "location": "header"}} + assert _resolve_scheme_field(scheme, "apiKeySecurityScheme") == {"name": "X-Key", "location": "header"} + + def test_resolve_scheme_field_flat_apikey(self): + scheme = {"name": "X-Key", "location": "header"} + assert _resolve_scheme_field(scheme, "apiKeySecurityScheme") == scheme + + def test_resolve_scheme_field_flat_http(self): + scheme = {"scheme": "bearer", "bearerFormat": "JWT"} + assert _resolve_scheme_field(scheme, "httpAuthSecurityScheme") == scheme + + def test_resolve_scheme_field_none(self): + assert _resolve_scheme_field({}, "apiKeySecurityScheme") is None + assert _resolve_scheme_field({"foo": "bar"}, "httpAuthSecurityScheme") is None + + def test_build_auth_header_apikey(self): + scheme = {"apiKeySecurityScheme": {"name": "X-Custom", "location": "header"}} + assert _build_auth_header_for_scheme(scheme, "secret") == ("X-Custom", "secret") + + def test_build_auth_header_http_bearer(self): + scheme = {"httpAuthSecurityScheme": {"scheme": "bearer"}} + assert _build_auth_header_for_scheme(scheme, "tok") == ("Authorization", "bearer tok") + + def test_build_auth_header_http_jwt(self): + scheme = {"httpAuthSecurityScheme": {"scheme": "bearer", "bearerFormat": "JWT"}} + assert _build_auth_header_for_scheme(scheme, "tok") == ("Authorization", "Bearer tok") + + def test_build_auth_header_apikey_query(self): + scheme = {"apiKeySecurityScheme": {"name": "key", "location": "query"}} + assert _build_auth_header_for_scheme(scheme, "val") is None + + def test_build_auth_header_no_match(self): + assert _build_auth_header_for_scheme({"foo": "bar"}, "val") is None + + def test_build_external_agent_config_with_security(self): + """_build_external_agent_config passes custom_headers from security.""" + agent = { + "external_agent_id": "ext_sec", + "name": "Secured Agent", + "security_schemes": {"k": {"apiKeySecurityScheme": {"name": "X-Token", "location": "header"}}}, + "security_requirements": [{"schemes": {"k": {}}}], + "security_credentials": {"k": "secret-token"}, + } + with patch('backend.agents.create_agent_info.ExternalA2AAgentConfig') as MockConfig: + _build_external_agent_config(agent, "http://test/a2a") + call_kwargs = MockConfig.call_args[1] + assert call_kwargs["custom_headers"] == {"X-Token": "secret-token"} + + def test_resolve_scheme_field_empty_scheme_string(self): + """scheme['scheme'] is empty string -> None.""" + from backend.agents.create_agent_info import _resolve_scheme_field + scheme = {"scheme": ""} + assert _resolve_scheme_field(scheme, "httpAuthSecurityScheme") is None + + def test_build_auth_header_empty_auth_scheme(self): + """httpAuth with empty scheme -> None.""" + from backend.agents.create_agent_info import _build_auth_header_for_scheme + scheme = {"httpAuthSecurityScheme": {"scheme": ""}} + assert _build_auth_header_for_scheme(scheme, "cred") is None + + def test_collect_auth_headers_non_dict_req(self): + """Requirement is not a dict -> skipped.""" + from backend.agents.create_agent_info import _build_security_headers + agent = { + "security_schemes": {"k": {"apiKeySecurityScheme": {"name": "X", "location": "header"}}}, + "security_requirements": ["not_a_dict"], + "security_credentials": {"k": "v"}, + } + assert _build_security_headers(agent) == {} + + def test_build_security_headers_missing_credential(self): + """scheme_id in requirements but not in credentials -> skipped.""" + from backend.agents.create_agent_info import _build_security_headers + agent = { + "security_schemes": {"k": {"apiKeySecurityScheme": {"name": "X", "location": "header"}}}, + "security_requirements": [{"schemes": {"k": {}, "missing": {}}}], + "security_credentials": {"k": "v"}, + } + headers = _build_security_headers(agent) + assert headers == {"X": "v"} + + def test_build_security_headers_scheme_builds_none(self): + """Credential present but header build returns None (e.g. query location) -> skipped.""" + from backend.agents.create_agent_info import _build_security_headers + agent = { + "security_schemes": {"k": {"apiKeySecurityScheme": {"name": "key", "location": "query"}}}, + "security_requirements": [{"schemes": {"k": {}}}], + "security_credentials": {"k": "v"}, + } + assert _build_security_headers(agent) == {} + diff --git a/test/backend/database/test_a2a_agent_db.py b/test/backend/database/test_a2a_agent_db.py index b06c5d06b9..11a10f58b6 100644 --- a/test/backend/database/test_a2a_agent_db.py +++ b/test/backend/database/test_a2a_agent_db.py @@ -1211,6 +1211,24 @@ def test_returns_results_with_joined_data(self, external_relation): # The join returns tuples; MockJoinQuery.all() returns raw list assert isinstance(result, list) + def test_returns_security_fields(self, external_relation): + """query_external_sub_agents returns security_schemes/requirements/credentials.""" + rel, agent = external_relation + agent.security_schemes = {"k": {"apiKeySecurityScheme": {"name": "X-Token", "location": "header"}}} + agent.security_requirements = [{"schemes": {"k": {}}}] + agent.security_credentials = {"k": "secret"} + with patch.object(a2a_db, '_get_db_session') as mk: + mk.return_value = MockSession({ + db_models_mock.A2AExternalAgentRelation: [rel], + db_models_mock.A2AExternalAgent: [agent], + }) + result = a2a_db.query_external_sub_agents(100, 'tenant-1') + assert len(result) > 0 + entry = result[0] + assert entry["security_schemes"] == {"k": {"apiKeySecurityScheme": {"name": "X-Token", "location": "header"}}} + assert entry["security_requirements"] == [{"schemes": {"k": {}}}] + assert entry["security_credentials"] == {"k": "secret"} + class TestListExternalRelationsByLocalAgent: def test_returns_empty_when_no_relations(self): diff --git a/test/sdk/core/agents/test_a2a_agent_proxy.py b/test/sdk/core/agents/test_a2a_agent_proxy.py index 50fc4176a8..f73858cbd2 100644 --- a/test/sdk/core/agents/test_a2a_agent_proxy.py +++ b/test/sdk/core/agents/test_a2a_agent_proxy.py @@ -424,6 +424,22 @@ def test_build_headers_with_api_key(self): headers = proxy._build_headers() assert headers["Authorization"] == "Bearer my-secret" + def test_build_headers_with_custom_headers(self): + """Test _build_headers merges custom_headers.""" + proxy = ExternalA2AAgentProxy(self._make_info( + custom_headers={"X-Custom": "val", "Authorization": "Bearer tok"} + )) + headers = proxy._build_headers() + assert headers["X-Custom"] == "val" + assert headers["Authorization"] == "Bearer tok" + assert headers["Content-Type"] == "application/json" + + def test_build_headers_without_custom_headers(self): + """Test _build_headers works without custom_headers (default None).""" + proxy = ExternalA2AAgentProxy(self._make_info()) + headers = proxy._build_headers() + assert "X-Custom" not in headers + def test_build_message_payload_query_only(self): """Test _build_message_payload builds correct structure with only query.""" proxy = ExternalA2AAgentProxy(self._make_info()) diff --git a/test/sdk/core/agents/test_agent_model.py b/test/sdk/core/agents/test_agent_model.py index dc903e8dc2..841341ffce 100644 --- a/test/sdk/core/agents/test_agent_model.py +++ b/test/sdk/core/agents/test_agent_model.py @@ -939,6 +939,22 @@ def test_external_a2a_agent_config_default_values(self): assert config.protocol_type == agent_model_module.PROTOCOL_JSONRPC assert config.timeout == 300.0 + def test_external_a2a_agent_config_custom_headers_default(self): + """Test custom_headers defaults to None.""" + config = agent_model_module.ExternalA2AAgentConfig( + agent_id="test", name="Test", description="", url="https://test.com" + ) + assert config.custom_headers is None + + def test_external_a2a_agent_config_with_custom_headers(self): + """Test custom_headers can be set.""" + config = agent_model_module.ExternalA2AAgentConfig( + agent_id="test", name="Test", description="", url="https://test.com", + custom_headers={"X-Custom": "val", "Authorization": "Bearer tok"} + ) + assert config.custom_headers == {"X-Custom": "val", "Authorization": "Bearer tok"} + + def test_external_a2a_agent_config_with_raw_card_skills(self): """Test ExternalA2AAgentConfig auto-enhances description from raw_card skills.""" config = agent_model_module.ExternalA2AAgentConfig( @@ -1018,7 +1034,8 @@ def test_external_a2a_agent_config_to_a2a_agent_info(self): protocol_version="1.5", protocol_type=agent_model_module.PROTOCOL_HTTP_JSON, timeout=450.0, - raw_card={"test": "data"} + raw_card={"test": "data"}, + custom_headers={"X-Token": "secret"} ) agent_info = config.to_a2a_agent_info() @@ -1036,6 +1053,7 @@ def test_external_a2a_agent_config_to_a2a_agent_info(self): assert call_kwargs["protocol_type"] == agent_model_module.PROTOCOL_HTTP_JSON assert call_kwargs["timeout"] == 450.0 assert call_kwargs["raw_card"] == {"test": "data"} + assert call_kwargs.get("custom_headers") == {"X-Token": "secret"} def test_external_a2a_agent_config_multiple_skills_examples(self): """Test ExternalA2AAgentConfig handles multiple skills with many examples."""