Add A2A agent access token auth

This commit is contained in:
elipitc
2026-06-04 23:06:45 +08:00
parent e0bf45db2f
commit f6851c9680
9 changed files with 197 additions and 8 deletions
@@ -8,6 +8,7 @@ import asyncio
import json
import uuid
import os
import secrets
from typing import Optional, Dict, Any, AsyncGenerator
from datetime import datetime
from contextlib import asynccontextmanager
@@ -38,6 +39,9 @@ SERVICE_PORT = int(os.getenv("SERVICE_PORT", "8000"))
POD_NAME = os.getenv("POD_NAME", "a2a-litellm-agent")
TEMPLATE_TYPE = os.getenv("TEMPLATE_TYPE", "a2a_litellm_agent")
USER_ID = os.getenv("USER_ID", "")
HEICODE_AGENT_ID = os.getenv("HEICODE_AGENT_ID", "")
AGENT_ACCESS_TOKEN = os.getenv("AGENT_ACCESS_TOKEN", "")
AGENT_ACCESS_HEADER = "X-Agent-Access-Token"
# ============== A2A 协议数据模型 ==============
@@ -182,6 +186,39 @@ class A2AAgentServer:
# 创建FastAPI应用
self.app = self._create_app()
def _agent_access_required(self) -> bool:
return bool(AGENT_ACCESS_TOKEN)
def _agent_authentication_card(self) -> Optional[Dict[str, Any]]:
if not self._agent_access_required():
return None
return {
"type": "header",
"header": AGENT_ACCESS_HEADER,
"required": True,
"agent_id": HEICODE_AGENT_ID or POD_NAME,
}
def _authorize_agent_request(self, request: Request) -> Optional[JSONResponse]:
expected_token = AGENT_ACCESS_TOKEN
if not expected_token:
return None
provided_token = request.headers.get(AGENT_ACCESS_HEADER, "")
if not provided_token:
return JSONResponse(
status_code=401,
content={"detail": f"missing {AGENT_ACCESS_HEADER}"},
)
if not secrets.compare_digest(expected_token, provided_token):
return JSONResponse(
status_code=403,
content={"detail": "agent access denied"},
)
return None
def _create_app(self) -> FastAPI:
"""创建FastAPI应用"""
@@ -238,7 +275,9 @@ class A2AAgentServer:
"protocol": "A2A",
"status": "running",
"pod_name": POD_NAME,
"template_type": TEMPLATE_TYPE
"template_type": TEMPLATE_TYPE,
"auth_required": self._agent_access_required(),
"agent_id": HEICODE_AGENT_ID or POD_NAME,
}
@app.get("/health")
@@ -249,6 +288,8 @@ class A2AAgentServer:
"pod_name": POD_NAME,
"template_type": TEMPLATE_TYPE,
"configured": self.llm_config.api_key is not None,
"auth_required": self._agent_access_required(),
"agent_id": HEICODE_AGENT_ID or POD_NAME,
"timestamp": datetime.utcnow().isoformat()
}
@@ -267,6 +308,7 @@ class A2AAgentServer:
streaming=self.agent_config.enable_streaming,
push_notifications=False
),
authentication=self._agent_authentication_card(),
skills=[
AgentSkill(
id="general-assistant",
@@ -285,6 +327,9 @@ class A2AAgentServer:
@app.post("/message/send")
async def send_message(request: Request):
"""A2A message/send 端点"""
auth_error = self._authorize_agent_request(request)
if auth_error:
return auth_error
body = await request.json()
# 解析JSON-RPC请求
@@ -318,6 +363,9 @@ class A2AAgentServer:
@app.post("/message/stream")
async def stream_message(request: Request):
"""A2A message/stream 端点 (SSE流式响应)"""
auth_error = self._authorize_agent_request(request)
if auth_error:
return auth_error
body = await request.json()
try:
@@ -335,8 +383,11 @@ class A2AAgentServer:
return await self._handle_message_stream(rpc_request)
@app.get("/tasks/{task_id}")
async def get_task(task_id: str):
async def get_task(task_id: str, request: Request):
"""获取任务状态"""
auth_error = self._authorize_agent_request(request)
if auth_error:
return auth_error
if task_id not in self.tasks:
raise HTTPException(status_code=404, detail="Task not found")
return self.tasks[task_id].model_dump()
@@ -25,6 +25,8 @@
| `AGENT_ROLE_NAME` | 启动时指定角色名,例如 `backend`、`reviewer` |
| `AGENT_INSTRUCTION_TEXT` | 启动时直接注入角色/行为说明文本,支持类似 `AGENTS.md` / `claude.md` 内容 |
| `AGENT_INSTRUCTION_FILE` | 启动时读取角色说明文件路径,文件内容会并入系统提示词 |
| `AGENT_ACCESS_TOKEN` | 可选。若设置,则 A2A 请求必须携带 `X-Agent-Access-Token` 且与其完全匹配 |
| `HEICODE_AGENT_ID` | 可选。用于在健康检查和 agent card 中暴露上层分配的 agent 标识 |
| `SERVICE_PORT` | 服务端口,默认 `8000` |
## 动态资源工具
@@ -130,6 +132,14 @@ export AGENT_INSTRUCTION_FILE=/workspace/AGENTS.md
## A2A 示例
如果设置了 `AGENT_ACCESS_TOKEN`,调用 `/message/send`、`/message/stream`、`/tasks/{task_id}` 时需要带:
```http
X-Agent-Access-Token: <AGENT_ACCESS_TOKEN>
```
服务端会使用常量时间比较校验请求头与环境变量值;未设置 `AGENT_ACCESS_TOKEN` 的旧实例继续兼容放行。
`POST /message/send`
```json
@@ -5,6 +5,7 @@ from __future__ import annotations
import json
import os
import secrets
import uuid
from contextlib import asynccontextmanager
from datetime import datetime
@@ -25,6 +26,9 @@ SERVICE_PORT = int(os.getenv("SERVICE_PORT", "8000"))
POD_NAME = os.getenv("POD_NAME", "coding-a2a-agent")
TEMPLATE_TYPE = os.getenv("TEMPLATE_TYPE", "coding_a2a_agent")
USER_ID = os.getenv("USER_ID", "")
HEICODE_AGENT_ID = os.getenv("HEICODE_AGENT_ID", "")
AGENT_ACCESS_TOKEN = os.getenv("AGENT_ACCESS_TOKEN", "")
AGENT_ACCESS_HEADER = "X-Agent-Access-Token"
class A2APart(BaseModel):
@@ -89,6 +93,7 @@ class AgentCard(BaseModel):
url: str
capabilities: AgentCapabilities
skills: list[AgentSkill]
authentication: Optional[dict[str, Any]] = None
class CodingA2AServer:
@@ -100,6 +105,39 @@ class CodingA2AServer:
self.tasks: dict[str, A2ATask] = {}
self.app = self._create_app()
def _agent_access_required(self) -> bool:
return bool(AGENT_ACCESS_TOKEN)
def _agent_authentication_card(self) -> Optional[dict[str, Any]]:
if not self._agent_access_required():
return None
return {
"type": "header",
"header": AGENT_ACCESS_HEADER,
"required": True,
"agent_id": HEICODE_AGENT_ID or POD_NAME,
}
def _authorize_agent_request(self, request: Request) -> Optional[JSONResponse]:
expected_token = AGENT_ACCESS_TOKEN
if not expected_token:
return None
provided_token = request.headers.get(AGENT_ACCESS_HEADER, "")
if not provided_token:
return JSONResponse(
status_code=401,
content={"detail": f"missing {AGENT_ACCESS_HEADER}"},
)
if not secrets.compare_digest(expected_token, provided_token):
return JSONResponse(
status_code=403,
content={"detail": "agent access denied"},
)
return None
def _create_app(self) -> FastAPI:
@asynccontextmanager
async def lifespan(app: FastAPI):
@@ -134,6 +172,8 @@ class CodingA2AServer:
"role_name": self.metadata.role_name,
"instruction_source": self.metadata.instruction_source,
"enabled_resources": runtime_defaults.resources.enabled_resource_names,
"auth_required": self._agent_access_required(),
"agent_id": HEICODE_AGENT_ID or POD_NAME,
}
@app.get("/health")
@@ -145,6 +185,8 @@ class CodingA2AServer:
"role_name": self.metadata.role_name,
"instruction_source": self.metadata.instruction_source,
"enabled_resources": runtime_defaults.resources.enabled_resource_names,
"auth_required": self._agent_access_required(),
"agent_id": HEICODE_AGENT_ID or POD_NAME,
"timestamp": datetime.utcnow().isoformat() + "Z",
}
@@ -157,6 +199,7 @@ class CodingA2AServer:
version=self.metadata.version,
url=base_url,
capabilities=AgentCapabilities(streaming=self.metadata.enable_streaming),
authentication=self._agent_authentication_card(),
skills=[
AgentSkill(id="coding", name="Coding", description="Inspect, edit, and verify repositories like a Claude Code style coding agent."),
AgentSkill(id="git", name="Git", description="Prepare workspaces, inspect git state, branch, commit, and push for Gitea, GitHub, and GitLab."),
@@ -167,6 +210,9 @@ class CodingA2AServer:
@app.post("/message/send")
async def message_send(request: Request):
auth_error = self._authorize_agent_request(request)
if auth_error:
return auth_error
rpc_request = A2ARequest(**(await request.json()))
if rpc_request.method != "message/send":
return JSONResponse(
@@ -176,11 +222,17 @@ class CodingA2AServer:
@app.post("/message/stream")
async def message_stream(request: Request):
auth_error = self._authorize_agent_request(request)
if auth_error:
return auth_error
rpc_request = A2ARequest(**(await request.json()))
return await self._handle_message_stream(rpc_request)
@app.get("/tasks/{task_id}")
async def get_task(task_id: str):
async def get_task(task_id: str, request: Request):
auth_error = self._authorize_agent_request(request)
if auth_error:
return auth_error
if task_id not in self.tasks:
raise HTTPException(status_code=404, detail="Task not found")
return self.tasks[task_id].model_dump()
+33 -1
View File
@@ -67,7 +67,9 @@ POST /agents
"env": {
"OPENAI_BASE_URL": "https://code.xinghanlab.com/v1",
"OPENAI_API_KEY": "sk-xxxx",
"MODEL_NAME": "gpt-5.4"
"MODEL_NAME": "gpt-5.4",
"AGENT_ACCESS_TOKEN": "550e8400-e29b-41d4-a716-446655440000",
"HEICODE_AGENT_ID": "dep-b5fab27e9255"
}
}
```
@@ -77,6 +79,7 @@ POST /agents
- `OPENAI_API_KEY` 当前建议在启动时传入
- 当前实测可用模型示例是 `gpt-5.4`
- 返回中会带 `namespace`、`pod_ip`、`access_info.external_ip`、`access_info.domain`
- 如果上层已注入 `AGENT_ACCESS_TOKEN`,A2A 请求入口会要求请求头 `X-Agent-Access-Token`
- 为兼容 HM 模板 Agent Runtime 契约,响应同时补充:
- `runtime_id` / `agent_id` / `id` = agent 名称
- `runtime_status` / `state` = 规范化后的生命周期状态
@@ -267,12 +270,29 @@ GET /.well-known/agent.json
"role_name": "backend",
"instruction_source": "env_text",
"enabled_resources": ["git", "azure_blob"],
"auth_required": true,
"timestamp": "2026-06-04T05:04:22.760314Z"
}
```
## 7. A2A 调用方式
### 7.0 访问鉴权
当前模板 Agent 支持 HM 约定的本地访问鉴权:
- 如果实例环境变量里存在 `AGENT_ACCESS_TOKEN`,则 `POST /message/send`、`POST /message/stream`、`GET /tasks/{task_id}` 必须带请求头 `X-Agent-Access-Token`
- 服务端使用常量时间比较校验 `X-Agent-Access-Token == AGENT_ACCESS_TOKEN`
- 缺少请求头时返回 `401`
- 请求头不匹配时返回 `403`
- 如果实例没有注入 `AGENT_ACCESS_TOKEN`,则继续兼容放行
注意:
- `X-Agent-Access-Token` 负责“谁有权访问这个 agent”
- A2A body 里的 `api_key` 负责“本次请求用谁的模型额度”
- 两者职责分离,不互相替代
### 7.1 同步调用
```http
@@ -281,6 +301,12 @@ POST /message/send
最小调用示例:
如果实例启用了访问鉴权,请附带请求头:
```http
X-Agent-Access-Token: 550e8400-e29b-41d4-a716-446655440000
```
```json
{
"jsonrpc": "2.0",
@@ -316,6 +342,12 @@ POST /message/stream
返回为 `text/event-stream`。
如果实例开启了访问鉴权,流式调用同样需要带:
```http
X-Agent-Access-Token: 550e8400-e29b-41d4-a716-446655440000
```
## 8. 请求级资源覆盖示例
如果你不想在启动时固定资源,可以在具体任务里传:
+18 -1
View File
@@ -1250,7 +1250,9 @@ curl -L \
"AGENT_INSTRUCTION_TEXT": "---\nname: architect\n---\n<Agent_Prompt>...</Agent_Prompt>",
"OPENAI_BASE_URL": "https://code.xinghanlab.com/v1",
"OPENAI_API_KEY": "sk-xxxx",
"MODEL_NAME": "gpt-5.4"
"MODEL_NAME": "gpt-5.4",
"AGENT_ACCESS_TOKEN": "550e8400-e29b-41d4-a716-446655440000",
"HEICODE_AGENT_ID": "dep-b5fab27e9255"
}
}
```
@@ -1284,6 +1286,21 @@ curl -L \
- `subdomain` 取自 `access_info.domain`,若 DNS 尚未就绪则回退到 `access_info.external_ip`。
- `runtime_status` / `state` 是对 Pod 生命周期的兼容投影;当前可能值为 `pending`、`running`、`stopped`、`failed`。
#### 客户端直连鉴权
模板 Agent 当前支持 HM 约定的本地访问鉴权:
- 当实例环境变量存在 `AGENT_ACCESS_TOKEN` 时,`POST /message/send`、`POST /message/stream`、`GET /tasks/{task_id}` 必须携带请求头 `X-Agent-Access-Token`
- 服务端使用常量时间比较校验 `X-Agent-Access-Token == AGENT_ACCESS_TOKEN`
- 请求头缺失时返回 `401`
- 请求头不匹配时返回 `403`
- 若实例未注入 `AGENT_ACCESS_TOKEN`,则继续兼容放行
职责边界:
- `X-Agent-Access-Token` 用于“谁有权访问这个 agent”
- A2A body 中的 `api_key` 仍用于“本次请求走谁的模型额度”
#### `GET /agents/{agent_name}`
用于 HM 轮询模板 Agent 生命周期。返回体与 `POST /agents` 的核心生命周期字段保持一致,便于 HM 复用同一套解析逻辑。
+1 -1
View File
@@ -31,7 +31,7 @@ spec:
containers:
- name: agent-manager
image: agnettaiji.azurecr.io/ai-agents/agent-manager:heicode-v2-runtime-20260604014451-arm64
image: agnettaiji.azurecr.io/ai-agents/agent-manager:heicode-v2-runtime-20260604230600-arm64
imagePullPolicy: Always
ports:
+1 -1
View File
@@ -22,7 +22,7 @@ spec:
- name: acr-secret
containers:
- name: agent-manager
image: agnettaiji.azurecr.io/ai-agents/agent-manager:heicode-v2-runtime-20260604014451-arm64
image: agnettaiji.azurecr.io/ai-agents/agent-manager:heicode-v2-runtime-20260604230600-arm64
imagePullPolicy: Always
ports:
- containerPort: 8000
+1 -1
View File
@@ -20,7 +20,7 @@ spec:
- name: acr-secret
containers:
- name: agent-manager
image: agnettaiji.azurecr.io/ai-agents/agent-manager:heicode-v2-runtime-20260604014451-arm64
image: agnettaiji.azurecr.io/ai-agents/agent-manager:heicode-v2-runtime-20260604230600-arm64
imagePullPolicy: Always
ports:
- containerPort: 8000
+27
View File
@@ -66,3 +66,30 @@ def test_integration_doc_mentions_template_agent_runtime_contract():
assert "| `GET` | `/agents/{agent_name}` | 查询模板 Agent 生命周期状态;返回平铺 `status` / `runtime_status` / `state` |" in doc_source
assert "| `POST` | `/agents/{agent_name}/stop` | 幂等停止模板 Agent;停止运行 Pod,但保留数据库记录 |" in doc_source
assert "删除接口当前已修复模板 Agent 场景下的数据库变量引用问题" in doc_source
def test_a2a_servers_enforce_agent_access_token_contract():
"""A2A servers should support local X-Agent-Access-Token validation."""
coding_server_source = _read("agent_templates/agents/coding_a2a_agent/a2a_server.py")
litellm_server_source = _read("agent_templates/agents/a2a_litellm_agent/a2a_server.py")
for source in (coding_server_source, litellm_server_source):
assert 'AGENT_ACCESS_TOKEN = os.getenv("AGENT_ACCESS_TOKEN", "")' in source
assert 'AGENT_ACCESS_HEADER = "X-Agent-Access-Token"' in source
assert "secrets.compare_digest" in source
assert 'status_code=401' in source
assert 'status_code=403' in source
assert 'request.headers.get(AGENT_ACCESS_HEADER, "")' in source
def test_a2a_docs_describe_agent_access_token_header():
"""Template-agent docs should explain the local access-token authentication flow."""
coding_doc_source = _read("docs/CODING_A2A_AGENT_CREATE_AND_INVOKE.md")
integration_doc_source = _read("docs/HEICODE_API_INTEGRATION.md")
assert "AGENT_ACCESS_TOKEN" in coding_doc_source
assert "X-Agent-Access-Token" in coding_doc_source
assert "缺少请求头时返回 `401`" in coding_doc_source
assert "请求头不匹配时返回 `403`" in coding_doc_source
assert "客户端直连鉴权" in integration_doc_source
assert "X-Agent-Access-Token == AGENT_ACCESS_TOKEN" in integration_doc_source