|
30 | 30 | from fastapi import HTTPException |
31 | 31 | from fastapi.testclient import TestClient |
32 | 32 | from google.adk.a2a import _compat |
| 33 | +from google.adk.agents._caller_principal import CallerPrincipal |
33 | 34 | from google.adk.agents.base_agent import BaseAgent |
34 | 35 | from google.adk.agents.llm_agent import LlmAgent |
35 | 36 | from google.adk.agents.run_config import RunConfig |
|
64 | 65 | from google.adk.sessions.base_session_service import ListSessionsResponse |
65 | 66 | from google.adk.sessions.in_memory_session_service import InMemorySessionService |
66 | 67 | from google.adk.sessions.session import Session |
| 68 | +from google.adk.tools.function_tool import FunctionTool |
67 | 69 | from google.adk.tools.tool_confirmation import ToolConfirmation |
| 70 | +from google.adk.utils._function_call_names import REQUEST_CONFIRMATION_FUNCTION_CALL_NAME |
68 | 71 | from google.api_core.exceptions import GoogleAPICallError |
69 | 72 | from google.api_core.exceptions import InvalidArgument |
70 | 73 | from google.genai import types |
|
74 | 77 | import starlette.requests |
75 | 78 | from starlette.routing import Mount |
76 | 79 |
|
| 80 | +from tests.unittests import testing_utils |
| 81 | + |
77 | 82 | # Configure logging to help diagnose server startup issues |
78 | 83 | logging.basicConfig( |
79 | 84 | level=logging.INFO, |
@@ -2643,6 +2648,319 @@ async def run_async_capture( |
2643 | 2648 | assert captured["run_config"].max_llm_calls == 37 |
2644 | 2649 |
|
2645 | 2650 |
|
| 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 | + |
2646 | 2964 | def test_agent_run_sse_splits_artifact_delta( |
2647 | 2965 | test_app, create_test_session, monkeypatch |
2648 | 2966 | ): |
|
0 commit comments