Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 20 additions & 4 deletions agentex/src/domain/services/agent_acp_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -271,20 +271,33 @@ def get_delegation_headers(self, agent: AgentEntity) -> dict[str, str]:
agent_identity=getattr(state, "agent_identity", None),
)

async def get_headers(
def get_request_context_headers(
self,
agent: AgentEntity,
request_headers: dict[str, str] | None = None,
) -> dict[str, str]:
"""Headers safe to expose to agent handler request context."""
filtered_request_headers = filter_request_headers(request_headers)
delegation_headers = self.get_delegation_headers(agent)
return {
**filtered_request_headers,
**delegation_headers,
}

async def get_headers(
self,
agent: AgentEntity,
request_headers: dict[str, str] | None = None,
) -> dict[str, str]:
request_context_headers = self.get_request_context_headers(
agent, request_headers
)
auth_headers = await self.get_agent_auth_headers(agent)
request_id = ctx_var_request_id.get(uuid4().hex)

# Later keys win. Client passthrough and delegation first; agent auth last.
return {
**filtered_request_headers,
**delegation_headers,
**request_context_headers,
**auth_headers,
"x-request-id": request_id,
}
Expand Down Expand Up @@ -420,12 +433,15 @@ async def send_event(
request_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
"""Send an event to a running task"""
request_context_headers = self.get_request_context_headers(
agent, request_headers
)

params = SendEventParams(
agent=agent,
task=task,
event=event,
request=None,
request={"headers": request_context_headers},
)

headers = await self.get_headers(agent, request_headers)
Expand Down
89 changes: 83 additions & 6 deletions agentex/tests/unit/services/test_agent_acp_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -333,11 +333,85 @@ async def test_send_event_delegation_not_raw_api_key_passthrough(
},
)

http_headers = mock_http_gateway.async_call.call_args[1]["default_headers"]
call_args = mock_http_gateway.async_call.call_args
payload = call_args[1]["payload"]
request_context_headers = payload["params"]["request"]["headers"]
assert request_context_headers["x-acting-user-api-key"] == "user-delegation-key"
assert request_context_headers["x-trace-id"] == "trace-456"
assert "x-api-key" not in request_context_headers
assert "x-agent-api-key" not in request_context_headers

http_headers = call_args[1]["default_headers"]
assert http_headers["x-acting-user-api-key"] == "user-delegation-key"
assert http_headers["x-trace-id"] == "trace-456"
assert "x-api-key" not in http_headers

async def test_send_event_request_context_exposes_delegated_headers_not_agent_auth(
self,
mock_http_gateway,
mock_request,
sample_agent,
sample_task,
sample_event,
):
"""EVENT_SEND params.request exposes user delegation, not agent auth."""
mock_request.state.principal_context = type(
"Principal",
(),
{"user_id": "user-1", "account_id": "acct-1"},
)()
mock_request.state.agent_identity = None
mock_request.headers = {
"x-api-key": "user-delegation-key",
"x-selected-account-id": "acct-1",
}
mock_agent_api_key_repository = MagicMock()
mock_agent_api_key_repository.get_internal_api_key_by_agent_id = AsyncMock(
return_value=MagicMock(api_key="agent-internal-key")
)
service = AgentACPService(
agent_repository=MagicMock(),
agent_api_key_repository=mock_agent_api_key_repository,
http_gateway=mock_http_gateway,
request=mock_request,
)

from src.domain.entities.agents_rpc import AgentRPCMethod

expected_request_id = f"{AgentRPCMethod.EVENT_SEND}-{sample_task.id}"
mock_http_gateway.async_call.return_value = {
"jsonrpc": "2.0",
"result": {"status": "event_sent", "event_id": sample_event.id},
"id": expected_request_id,
}

await service.send_event(
agent=sample_agent,
event=sample_event,
task=sample_task,
acp_url="http://test-acp.example.com",
request_headers={
"x-api-key": "must-not-forward",
"x-trace-id": "trace-456",
},
)

call_args = mock_http_gateway.async_call.call_args
request_context_headers = call_args[1]["payload"]["params"]["request"][
"headers"
]
assert request_context_headers == {
"x-trace-id": "trace-456",
"x-acting-user-api-key": "user-delegation-key",
"x-selected-account-id": "acct-1",
}
assert "x-api-key" not in request_context_headers
assert "x-agent-api-key" not in request_context_headers

http_headers = call_args[1]["default_headers"]
assert http_headers["x-agent-api-key"] == "agent-internal-key"
assert http_headers["x-acting-user-api-key"] == "user-delegation-key"

async def test_send_message_includes_cookie_delegation_headers(
self,
agent_acp_service,
Expand Down Expand Up @@ -719,12 +793,17 @@ async def test_send_event_with_request_headers(

call_args = mock_http_gateway.async_call.call_args
payload = call_args[1]["payload"]
assert payload["params"]["request"] is None
request_context_headers = payload["params"]["request"]["headers"]
assert request_context_headers == {
"x-user-id": "user-123",
"x-trace-id": "trace-456",
}

http_headers = call_args[1]["default_headers"]
assert http_headers["x-user-id"] == "user-123"
assert http_headers["x-trace-id"] == "trace-456"
assert http_headers["x-agent-api-key"] == "test-api-key"
assert "x-agent-api-key" not in request_context_headers
assert "x-api-key" not in http_headers
assert "user-agent" not in http_headers
assert "authorization" not in http_headers
Expand Down Expand Up @@ -765,14 +844,12 @@ async def test_send_event_without_request_headers(
assert result["status"] == "event_sent"
assert result["event_id"] == sample_event.id

# Verify call was made and request field is None when headers not provided
# Verify request context is present but empty when headers are not provided.
mock_http_gateway.async_call.assert_called_once()
call_args = mock_http_gateway.async_call.call_args
payload = call_args[1]["payload"]
assert "params" in payload
# Request field should be None or not included when headers not provided
if "request" in payload["params"]:
assert payload["params"]["request"] is None
assert payload["params"]["request"] == {"headers": {}}

async def test_jsonrpc_error_handling(
self,
Expand Down
Loading