Skip to content

Commit d4ff60a

Browse files
authored
test(api_server): caller principal on /run and /run_sse
Unit tests: the routes pass CallerPrincipal(authenticated=False, source="api_server") by default, use a principal set on request.state by middleware, and omit the kwarg for a runner whose run_async does not accept it. End to end over /run: strict mode refuses a confirmation from an unauthenticated caller, honors one from an authenticated operator set by middleware, and default mode honors it with a warning.
1 parent 9631358 commit d4ff60a

1 file changed

Lines changed: 318 additions & 0 deletions

File tree

‎tests/unittests/cli/test_fast_api.py‎

Lines changed: 318 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
from fastapi import HTTPException
3131
from fastapi.testclient import TestClient
3232
from google.adk.a2a import _compat
33+
from google.adk.agents._caller_principal import CallerPrincipal
3334
from google.adk.agents.base_agent import BaseAgent
3435
from google.adk.agents.llm_agent import LlmAgent
3536
from google.adk.agents.run_config import RunConfig
@@ -64,7 +65,9 @@
6465
from google.adk.sessions.base_session_service import ListSessionsResponse
6566
from google.adk.sessions.in_memory_session_service import InMemorySessionService
6667
from google.adk.sessions.session import Session
68+
from google.adk.tools.function_tool import FunctionTool
6769
from google.adk.tools.tool_confirmation import ToolConfirmation
70+
from google.adk.utils._function_call_names import REQUEST_CONFIRMATION_FUNCTION_CALL_NAME
6871
from google.api_core.exceptions import GoogleAPICallError
6972
from google.api_core.exceptions import InvalidArgument
7073
from google.genai import types
@@ -74,6 +77,8 @@
7477
import starlette.requests
7578
from starlette.routing import Mount
7679

80+
from tests.unittests import testing_utils
81+
7782
# Configure logging to help diagnose server startup issues
7883
logging.basicConfig(
7984
level=logging.INFO,
@@ -2643,6 +2648,319 @@ async def run_async_capture(
26432648
assert captured["run_config"].max_llm_calls == 37
26442649

26452650

2651+
class _CallerPrincipalInjector:
2652+
"""ASGI wrapper standing in for a deployment's own authentication.
2653+
2654+
Sets ``request.state.caller_principal`` the way an auth middleware would,
2655+
before the request reaches the ADK routes.
2656+
"""
2657+
2658+
def __init__(self, app, principal: CallerPrincipal):
2659+
self._app = app
2660+
self._principal = principal
2661+
2662+
async def __call__(self, scope, receive, send):
2663+
if scope["type"] == "http":
2664+
scope.setdefault("state", {})["caller_principal"] = self._principal
2665+
await self._app(scope, receive, send)
2666+
2667+
2668+
def _capture_caller_principal(monkeypatch) -> dict[str, Any]:
2669+
"""Patches Runner.run_async to record the caller_principal it is given."""
2670+
captured: dict[str, Any] = {"calls": 0, "caller_principal": "unset"}
2671+
2672+
async def run_async_capture(
2673+
self,
2674+
*,
2675+
user_id: str,
2676+
session_id: str,
2677+
invocation_id: Optional[str] = None,
2678+
new_message: Optional[types.Content] = None,
2679+
state_delta: Optional[dict[str, Any]] = None,
2680+
run_config: Optional[RunConfig] = None,
2681+
caller_principal: Optional[CallerPrincipal] = None,
2682+
**kwargs,
2683+
):
2684+
del self, user_id, session_id, invocation_id, new_message, state_delta
2685+
del run_config, kwargs
2686+
captured["calls"] += 1
2687+
captured["caller_principal"] = caller_principal
2688+
yield _event_1()
2689+
2690+
monkeypatch.setattr(Runner, "run_async", run_async_capture)
2691+
return captured
2692+
2693+
2694+
@pytest.mark.parametrize("route", ["/run", "/run_sse"])
2695+
def test_agent_run_records_unauthenticated_caller_by_default(
2696+
test_app, create_test_session, monkeypatch, route
2697+
):
2698+
"""Without authentication in front of the server, a run is unauthenticated.
2699+
2700+
The user_id in the body is a claim the caller made, so it must not turn
2701+
into an authenticated principal.
2702+
"""
2703+
info = create_test_session
2704+
captured = _capture_caller_principal(monkeypatch)
2705+
2706+
payload = {
2707+
"app_name": info["app_name"],
2708+
"user_id": info["user_id"],
2709+
"session_id": info["session_id"],
2710+
"new_message": {"role": "user", "parts": [{"text": "Hello"}]},
2711+
"streaming": False,
2712+
}
2713+
2714+
response = test_app.post(route, json=payload)
2715+
2716+
assert response.status_code == 200
2717+
assert captured["calls"] == 1
2718+
assert captured["caller_principal"] == CallerPrincipal(
2719+
authenticated=False, source="api_server"
2720+
)
2721+
2722+
2723+
@pytest.mark.parametrize("route", ["/run", "/run_sse"])
2724+
def test_agent_run_uses_caller_principal_set_by_middleware(
2725+
test_app, create_test_session, monkeypatch, route
2726+
):
2727+
"""A principal the deployment's own authentication set is used as is."""
2728+
info = create_test_session
2729+
captured = _capture_caller_principal(monkeypatch)
2730+
principal = CallerPrincipal(
2731+
authenticated=True, user_name="operator", source="gateway"
2732+
)
2733+
client = TestClient(_CallerPrincipalInjector(test_app.app, principal))
2734+
2735+
payload = {
2736+
"app_name": info["app_name"],
2737+
"user_id": info["user_id"],
2738+
"session_id": info["session_id"],
2739+
"new_message": {"role": "user", "parts": [{"text": "Hello"}]},
2740+
"streaming": False,
2741+
}
2742+
2743+
response = client.post(route, json=payload)
2744+
2745+
assert response.status_code == 200
2746+
assert captured["caller_principal"] == principal
2747+
2748+
2749+
def test_agent_run_omits_caller_principal_for_runner_without_the_parameter(
2750+
test_app, create_test_session, monkeypatch
2751+
):
2752+
"""A Runner subclass whose run_async predates the parameter still works."""
2753+
info = create_test_session
2754+
captured: dict[str, Any] = {"kwargs": None}
2755+
2756+
async def run_async_capture(
2757+
self,
2758+
*,
2759+
user_id: str,
2760+
session_id: str,
2761+
invocation_id: Optional[str] = None,
2762+
new_message: Optional[types.Content] = None,
2763+
state_delta: Optional[dict[str, Any]] = None,
2764+
run_config: Optional[RunConfig] = None,
2765+
):
2766+
del self, user_id, session_id, new_message, state_delta, run_config
2767+
captured["kwargs"] = {"invocation_id": invocation_id}
2768+
yield _event_1()
2769+
2770+
monkeypatch.setattr(Runner, "run_async", run_async_capture)
2771+
2772+
payload = {
2773+
"app_name": info["app_name"],
2774+
"user_id": info["user_id"],
2775+
"session_id": info["session_id"],
2776+
"new_message": {"role": "user", "parts": [{"text": "Hello"}]},
2777+
"streaming": False,
2778+
}
2779+
2780+
response = test_app.post("/run", json=payload)
2781+
2782+
assert response.status_code == 200
2783+
assert captured["kwargs"] == {"invocation_id": None}
2784+
2785+
2786+
def _confirmation_gated_app(info, calls: list[str]) -> App:
2787+
"""An App whose agent calls a confirmation-gated tool on the first turn."""
2788+
2789+
def rm_all() -> dict:
2790+
calls.append("rm_all")
2791+
return {"status": "deleted"}
2792+
2793+
model = testing_utils.MockModel.create(
2794+
responses=[
2795+
types.Part.from_function_call(name="rm_all", args={}),
2796+
"after tool",
2797+
"after tool 2",
2798+
]
2799+
)
2800+
return App(
2801+
name=info["app_name"],
2802+
root_agent=LlmAgent(
2803+
name="ops_agent",
2804+
model=model,
2805+
tools=[FunctionTool(func=rm_all, require_confirmation=True)],
2806+
),
2807+
)
2808+
2809+
2810+
def _run_payload(info, parts):
2811+
return {
2812+
"app_name": info["app_name"],
2813+
"user_id": info["user_id"],
2814+
"session_id": info["session_id"],
2815+
"new_message": {"role": "user", "parts": parts},
2816+
"streaming": False,
2817+
}
2818+
2819+
2820+
def _approve_through_run(client, info) -> list[dict]:
2821+
"""Triggers the gated tool over /run, then answers its confirmation."""
2822+
first = client.post("/run", json=_run_payload(info, [{"text": "delete"}]))
2823+
assert first.status_code == 200
2824+
confirmation_id = next(
2825+
part["functionCall"]["id"]
2826+
for event in first.json()
2827+
for part in (event.get("content") or {}).get("parts", [])
2828+
if part.get("functionCall", {}).get("name")
2829+
== REQUEST_CONFIRMATION_FUNCTION_CALL_NAME
2830+
)
2831+
second = client.post(
2832+
"/run",
2833+
json=_run_payload(
2834+
info,
2835+
[{
2836+
"functionResponse": {
2837+
"id": confirmation_id,
2838+
"name": REQUEST_CONFIRMATION_FUNCTION_CALL_NAME,
2839+
"response": {"confirmed": True},
2840+
}
2841+
}],
2842+
),
2843+
)
2844+
assert second.status_code == 200
2845+
return [
2846+
part["functionResponse"]["response"]
2847+
for event in second.json()
2848+
for part in (event.get("content") or {}).get("parts", [])
2849+
if part.get("functionResponse", {}).get("name") == "rm_all"
2850+
]
2851+
2852+
2853+
def test_agent_run_strict_mode_refuses_confirmation_without_authentication(
2854+
create_test_session,
2855+
mock_session_service,
2856+
mock_agent_loader,
2857+
mock_eval_sets_manager,
2858+
mock_eval_set_results_manager,
2859+
monkeypatch,
2860+
):
2861+
"""A /run caller cannot approve its own gated tool call in strict mode."""
2862+
info = create_test_session
2863+
monkeypatch.setenv("ADK_ENABLE_STRICT_CALLER_PRINCIPAL", "1")
2864+
monkeypatch.setattr(Runner, "run_async", _ORIGINAL_RUNNER_RUN_ASYNC)
2865+
calls: list[str] = []
2866+
loaded_app = _confirmation_gated_app(info, calls)
2867+
monkeypatch.setattr(
2868+
mock_agent_loader, "load_agent", lambda app_name: loaded_app
2869+
)
2870+
client = _create_test_client(
2871+
mock_session_service,
2872+
InMemoryArtifactService(),
2873+
InMemoryMemoryService(),
2874+
mock_agent_loader,
2875+
mock_eval_sets_manager,
2876+
mock_eval_set_results_manager,
2877+
)
2878+
2879+
results = _approve_through_run(client, info)
2880+
2881+
assert calls == []
2882+
assert results == [{"error": "This tool call is rejected."}]
2883+
2884+
2885+
def test_agent_run_strict_mode_honors_confirmation_from_authenticated_operator(
2886+
create_test_session,
2887+
mock_session_service,
2888+
mock_agent_loader,
2889+
mock_eval_sets_manager,
2890+
mock_eval_set_results_manager,
2891+
monkeypatch,
2892+
):
2893+
"""A principal set by the deployment's authentication is honored."""
2894+
info = create_test_session
2895+
monkeypatch.setenv("ADK_ENABLE_STRICT_CALLER_PRINCIPAL", "1")
2896+
monkeypatch.setattr(Runner, "run_async", _ORIGINAL_RUNNER_RUN_ASYNC)
2897+
calls: list[str] = []
2898+
loaded_app = _confirmation_gated_app(info, calls)
2899+
monkeypatch.setattr(
2900+
mock_agent_loader, "load_agent", lambda app_name: loaded_app
2901+
)
2902+
base = _create_test_client(
2903+
mock_session_service,
2904+
InMemoryArtifactService(),
2905+
InMemoryMemoryService(),
2906+
mock_agent_loader,
2907+
mock_eval_sets_manager,
2908+
mock_eval_set_results_manager,
2909+
)
2910+
client = TestClient(
2911+
_CallerPrincipalInjector(
2912+
base.app,
2913+
CallerPrincipal(
2914+
authenticated=True, user_name="operator", source="gateway"
2915+
),
2916+
)
2917+
)
2918+
2919+
results = _approve_through_run(client, info)
2920+
2921+
assert calls == ["rm_all"]
2922+
assert results == [{"status": "deleted"}]
2923+
2924+
2925+
def test_agent_run_default_mode_honors_confirmation_with_a_warning(
2926+
create_test_session,
2927+
mock_session_service,
2928+
mock_agent_loader,
2929+
mock_eval_sets_manager,
2930+
mock_eval_set_results_manager,
2931+
monkeypatch,
2932+
caplog,
2933+
):
2934+
"""Without strict mode the behavior is unchanged apart from the warning."""
2935+
info = create_test_session
2936+
monkeypatch.delenv("ADK_ENABLE_STRICT_CALLER_PRINCIPAL", raising=False)
2937+
monkeypatch.setattr(Runner, "run_async", _ORIGINAL_RUNNER_RUN_ASYNC)
2938+
calls: list[str] = []
2939+
loaded_app = _confirmation_gated_app(info, calls)
2940+
monkeypatch.setattr(
2941+
mock_agent_loader, "load_agent", lambda app_name: loaded_app
2942+
)
2943+
client = _create_test_client(
2944+
mock_session_service,
2945+
InMemoryArtifactService(),
2946+
InMemoryMemoryService(),
2947+
mock_agent_loader,
2948+
mock_eval_sets_manager,
2949+
mock_eval_set_results_manager,
2950+
)
2951+
2952+
with caplog.at_level("WARNING", logger="google_adk"):
2953+
results = _approve_through_run(client, info)
2954+
2955+
assert calls == ["rm_all"]
2956+
assert results == [{"status": "deleted"}]
2957+
assert any(
2958+
"unauthenticated caller" in record.getMessage()
2959+
and "api_server" in record.getMessage()
2960+
for record in caplog.records
2961+
)
2962+
2963+
26462964
def test_agent_run_sse_splits_artifact_delta(
26472965
test_app, create_test_session, monkeypatch
26482966
):

0 commit comments

Comments
 (0)