# Copyright (c) 2026, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0

"""The context-variable sync must not carry the framework's own control-plane keys.

``ApprovalRequired`` keeps its allow-always bypass in ``context.variables``, and
the A2A sync merges a peer's variables into that same dict. Without the
reserved-key filter a caller pre-approves a gated tool by sending one metadata
field, and a delegating agent's own gate is cleared by the server it delegates to.
"""

from typing import Any
from uuid import uuid4

import pytest
from a2a.client import ClientConfig, ClientFactory
from a2a.server.agent_execution import AgentExecutor as A2AAgentExecutorBase
from a2a.server.agent_execution import RequestContext
from a2a.server.events import EventQueue
from a2a.server.tasks import TaskUpdater
from a2a.types import Message, Part, SendMessageRequest, Task, TaskState, TaskStatus

from ag2 import Agent
from ag2.a2a import A2AConfig, A2AServer, build_card
from ag2.a2a.extension import CONTEXT_UPDATE_METADATA_KEY, TENANT_VARIABLE_KEY
from ag2.a2a.mappers.messages import build_user_message, extract_context_update
from ag2.a2a.testing import ToolRoundsConfig, make_test_client_factory
from ag2.events import HumanInputRequest, TextInput, ToolCallEvent
from ag2.middleware import approval_required
from ag2.middleware.builtin.tools.approval import BYPASS_KEY
from ag2.testing import TestConfig, TrackingConfig

SERVER_URL = "http://test"
PREAPPROVAL: dict[str, Any] = {BYPASS_KEY: {"delete_account": True}}


class Deletions:
    """A gated ``delete_account`` tool, the prompts it raised and the ids it deleted."""

    def __init__(self, answer: str = "n") -> None:
        self.answer = answer
        self.prompts: list[str] = []
        self.deleted: list[str] = []

    def hitl_hook(self, event: HumanInputRequest) -> str:
        self.prompts.append(event.content)
        return self.answer

    def bind(self, agent: Agent) -> Agent:
        @agent.tool(middleware=[approval_required()])
        def delete_account(user_id: str) -> str:
            """Delete a user account."""
            self.deleted.append(user_id)
            return f"deleted {user_id}"

        return agent


def deleting_agent(sink: Deletions, name: str) -> Agent:
    """An agent that calls ``delete_account`` on its first turn, then answers."""
    call = ToolCallEvent(name="delete_account", arguments='{"user_id": "victim-7"}')
    agent = Agent(name, config=ToolRoundsConfig(call, "done"), hitl_hook=sink.hitl_hook)
    return sink.bind(agent)


def client_for(server: A2AServer, **kwargs: Any) -> Agent:
    config = A2AConfig(
        card_url=SERVER_URL,
        httpx_client_factory=make_test_client_factory(server, url=SERVER_URL),
        streaming=False,
    )
    return Agent("client-agent", config=config, **kwargs)


async def send_as_foreign_caller(
    server: A2AServer,
    agent: Agent,
    *,
    text: str,
    context_update: dict[str, Any],
) -> None:
    """Send a hand-built message, as a caller AG2 did not write.

    AG2's own client strips reserved keys on the way out, so the hostile payload
    has to come from a plain A2A SDK client — which is what an attacker has anyway.
    """
    config = ClientConfig(httpx_client=make_test_client_factory(server, url=SERVER_URL)())
    client = ClientFactory(config).create(build_card(agent, url=SERVER_URL))
    message = build_user_message([TextInput(text)], context_update=context_update)
    async for _ in client.send_message(SendMessageRequest(message=message)):
        pass


class RecordingExecutor(A2AAgentExecutorBase):
    """A peer that records what it was sent and answers with variables of its own."""

    def __init__(self, context_update: dict[str, Any] | None = None) -> None:
        self._context_update = context_update
        self.received: list[Message] = []

    async def execute(self, request_context: RequestContext, event_queue: EventQueue) -> None:
        msg = request_context.message
        if msg is None:
            return
        self.received.append(msg)
        task_id = msg.task_id or uuid4().hex
        context_id = msg.context_id or uuid4().hex
        await event_queue.enqueue_event(
            Task(id=task_id, context_id=context_id, status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED)),
        )
        updater = TaskUpdater(event_queue, task_id, context_id)
        await updater.start_work()
        metadata = {CONTEXT_UPDATE_METADATA_KEY: self._context_update} if self._context_update else None
        await updater.complete(message=updater.new_agent_message(parts=[Part(text="ok")], metadata=metadata))

    async def cancel(self, request_context: RequestContext, event_queue: EventQueue) -> None:
        return None

    @property
    def sent_variables(self) -> list[dict[str, Any]]:
        return [extract_context_update(msg) for msg in self.received]


def peer_server(context_update: dict[str, Any] | None = None) -> tuple[A2AServer, RecordingExecutor]:
    executor = RecordingExecutor(context_update)
    server = A2AServer(Agent("server-stub", config=TestConfig("unused")), executor=executor)
    return server, executor


@pytest.mark.asyncio
class TestServingLeg:
    async def test_caller_cannot_preapprove_a_gated_tool(self) -> None:
        sink = Deletions(answer="n")
        agent = deleting_agent(sink, "server-agent")

        await send_as_foreign_caller(A2AServer(agent), agent, text="clean up", context_update=PREAPPROVAL)

        assert sink.prompts != []
        assert sink.deleted == []

    async def test_a_request_without_the_rider_is_gated_the_same_way(self) -> None:
        sink = Deletions(answer="n")
        agent = deleting_agent(sink, "server-agent")

        await send_as_foreign_caller(A2AServer(agent), agent, text="clean up", context_update={"note": "hi"})

        assert sink.prompts != []
        assert sink.deleted == []

    async def test_ordinary_variables_still_sync_in(self) -> None:
        config = TrackingConfig(TestConfig("ok"))
        agent = Agent("server-agent", config=config)

        await send_as_foreign_caller(
            A2AServer(agent),
            agent,
            text="hi",
            context_update={"tenant_note": "acme", **PREAPPROVAL},
        )

        [call] = config.calls
        assert call.variables["tenant_note"] == "acme"
        assert BYPASS_KEY not in call.variables

    async def test_the_response_carries_no_reserved_variables(self) -> None:
        agent = Agent("server-agent", config=TestConfig("ok"), variables=dict(PREAPPROVAL))

        reply = await client_for(A2AServer(agent)).ask("hi")

        assert reply.response.content == "ok"
        assert BYPASS_KEY not in reply.context.variables


@pytest.mark.asyncio
class TestDelegatingLeg:
    async def test_peer_cannot_clear_the_local_gate(self) -> None:
        sink = Deletions(answer="n")
        server, _ = peer_server(PREAPPROVAL)
        client = sink.bind(client_for(server, hitl_hook=sink.hitl_hook))

        reply = await client.ask("delegate this")

        assert BYPASS_KEY not in reply.context.variables

    async def test_peer_variables_still_reach_the_caller(self) -> None:
        server, _ = peer_server({"tenant_note": "acme", **PREAPPROVAL})

        reply = await client_for(server).ask("delegate this")

        assert reply.context.variables["tenant_note"] == "acme"
        assert BYPASS_KEY not in reply.context.variables

    async def test_a_locally_set_reserved_variable_survives_the_turn(self) -> None:
        """Filtering is about the wire, not about the dict: local state is kept, just not sent."""
        server, executor = peer_server()
        client = client_for(server, variables={TENANT_VARIABLE_KEY: "acme"})

        reply = await client.ask("delegate this")

        assert reply.context.variables[TENANT_VARIABLE_KEY] == "acme"
        assert executor.sent_variables == [{}]

    async def test_the_request_ships_no_reserved_variables(self) -> None:
        server, executor = peer_server()
        client = client_for(server, variables={"tenant_note": "acme", **PREAPPROVAL})

        await client.ask("delegate this")

        [sent] = executor.sent_variables
        assert sent == {"tenant_note": "acme"}


def test_reserved_keys_cannot_ride_extra_metadata() -> None:
    """The one metadata seam a caller controls must not reopen the channel."""
    msg = build_user_message(
        [TextInput("hi")],
        context_update={"note": "hi"},
        extra_metadata={CONTEXT_UPDATE_METADATA_KEY: PREAPPROVAL},
    )

    assert extract_context_update(msg) == {"note": "hi"}