support jwt token to invoke a2a agent

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(
         {

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

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)