Add A2A agent access token auth
This commit is contained in:
@@ -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()
|
||||
|
||||
@@ -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. 请求级资源覆盖示例
|
||||
|
||||
如果你不想在启动时固定资源,可以在具体任务里传:
|
||||
|
||||
@@ -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 复用同一套解析逻辑。
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user