# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio
from contextlib import aclosing
import importlib
from pathlib import Path
import sys
import textwrap
from typing import AsyncGenerator
from typing import Optional
from unittest.mock import AsyncMock

from google.adk.agents.base_agent import BaseAgent
from google.adk.agents.context_cache_config import ContextCacheConfig
from google.adk.agents.invocation_context import InvocationContext
from google.adk.agents.llm_agent import LlmAgent
from google.adk.agents.run_config import RunConfig
from google.adk.apps.app import App
from google.adk.apps.app import ResumabilityConfig
from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
from google.adk.cli.utils.agent_loader import AgentLoader
from google.adk.errors.session_not_found_error import SessionNotFoundError
from google.adk.events.event import Event
from google.adk.plugins.base_plugin import BasePlugin
from google.adk.runners import Runner
from google.adk.sessions.base_session_service import GetSessionConfig
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.adk.sessions.session import Session
from google.adk.tools.base_toolset import BaseToolset
from google.genai import types
import pytest

TEST_APP_ID = "test_app"
TEST_USER_ID = "test_user"
TEST_SESSION_ID = "test_session"


class MockAgent(BaseAgent):
  """Mock agent for unit testing."""

  def __init__(
      self,
      name: str,
      parent_agent: Optional[BaseAgent] = None,
  ):
    super().__init__(name=name, sub_agents=[])
    # BaseAgent doesn't have disallow_transfer_to_parent field
    # This is intentional as we want to test non-LLM agents
    if parent_agent:
      self.parent_agent = parent_agent

  async def _run_async_impl(
      self, invocation_context: InvocationContext
  ) -> AsyncGenerator[Event, None]:
    yield Event(
        invocation_id=invocation_context.invocation_id,
        author=self.name,
        content=types.Content(
            role="model", parts=[types.Part(text="Test response")]
        ),
    )


class MockLiveAgent(BaseAgent):
  """Mock live agent for unit testing."""

  def __init__(self, name: str):
    super().__init__(name=name, sub_agents=[])

  async def _run_live_impl(
      self, invocation_context: InvocationContext
  ) -> AsyncGenerator[Event, None]:
    yield Event(
        invocation_id=invocation_context.invocation_id,
        author=self.name,
        content=types.Content(
            role="model", parts=[types.Part(text="live hello")]
        ),
    )


class MockLlmAgent(LlmAgent):
  """Mock LLM agent for unit testing."""

  def __init__(
      self,
      name: str,
      disallow_transfer_to_parent: bool = False,
      parent_agent: Optional[BaseAgent] = None,
  ):
    # Use a string model instead of mock
    super().__init__(name=name, model="gemini-1.5-pro", sub_agents=[])
    self.disallow_transfer_to_parent = disallow_transfer_to_parent
    self.parent_agent = parent_agent

  async def _run_async_impl(
      self, invocation_context: InvocationContext
  ) -> AsyncGenerator[Event, None]:
    yield Event(
        invocation_id=invocation_context.invocation_id,
        author=self.name,
        content=types.Content(
            role="model", parts=[types.Part(text="Test LLM response")]
        ),
    )


class MockAgentWithMetadata(BaseAgent):
  """Mock agent that returns event-level custom metadata."""

  def __init__(self, name: str):
    super().__init__(name=name, sub_agents=[])

  async def _run_async_impl(
      self, invocation_context: InvocationContext
  ) -> AsyncGenerator[Event, None]:
    yield Event(
        invocation_id=invocation_context.invocation_id,
        author=self.name,
        content=types.Content(
            role="model", parts=[types.Part(text="Test response")]
        ),
        custom_metadata={"event_key": "event_value"},
    )


class MockPlugin(BasePlugin):
  """Mock plugin for unit testing."""

  ON_USER_CALLBACK_MSG = (
      "Modified user message ON_USER_CALLBACK_MSG from MockPlugin"
  )
  ON_EVENT_CALLBACK_MSG = "Modified event ON_EVENT_CALLBACK_MSG from MockPlugin"
  ON_EVENT_CALLBACK_METADATA = {"plugin_key": "plugin_value"}

  def __init__(self):
    super().__init__(name="mock_plugin")
    self.enable_user_message_callback = False
    self.enable_event_callback = False
    self.user_content_seen_in_before_run_callback = None

  async def on_user_message_callback(
      self,
      *,
      invocation_context: InvocationContext,
      user_message: types.Content,
  ) -> Optional[types.Content]:
    if not self.enable_user_message_callback:
      return None
    return types.Content(
        role="model",
        parts=[types.Part(text=self.ON_USER_CALLBACK_MSG)],
    )

  async def before_run_callback(
      self,
      *,
      invocation_context: InvocationContext,
  ) -> None:
    self.user_content_seen_in_before_run_callback = (
        invocation_context.user_content
    )

  async def on_event_callback(
      self, *, invocation_context: InvocationContext, event: Event
  ) -> Optional[Event]:
    if not self.enable_event_callback:
      return None
    return Event(
        invocation_id="",
        author="",
        content=types.Content(
            parts=[
                types.Part(
                    text=self.ON_EVENT_CALLBACK_MSG,
                )
            ],
            role=event.content.role,
        ),
        custom_metadata=self.ON_EVENT_CALLBACK_METADATA,
    )


class TestRunnerFindAgentToRun:
  """Tests for Runner._find_agent_to_run method."""

  def setup_method(self):
    """Set up test fixtures."""
    self.session_service = InMemorySessionService()
    self.artifact_service = InMemoryArtifactService()

    # Create test agents
    self.root_agent = MockLlmAgent("root_agent")
    self.sub_agent1 = MockLlmAgent("sub_agent1", parent_agent=self.root_agent)
    self.sub_agent2 = MockLlmAgent("sub_agent2", parent_agent=self.root_agent)
    self.non_transferable_agent = MockLlmAgent(
        "non_transferable",
        disallow_transfer_to_parent=True,
        parent_agent=self.root_agent,
    )

    self.root_agent.sub_agents = [
        self.sub_agent1,
        self.sub_agent2,
        self.non_transferable_agent,
    ]

    self.runner = Runner(
        app_name="test_app",
        agent=self.root_agent,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )

  def test_find_agent_to_run_with_function_response_scenario(self):
    """Test finding agent when last event is function response."""
    # Create a function call from sub_agent1
    function_call = types.FunctionCall(id="func_123", name="test_func", args={})
    function_response = types.FunctionResponse(
        id="func_123", name="test_func", response={}
    )

    call_event = Event(
        invocation_id="inv1",
        author="sub_agent1",
        content=types.Content(
            role="model", parts=[types.Part(function_call=function_call)]
        ),
    )

    response_event = Event(
        invocation_id="inv2",
        author="user",
        content=types.Content(
            role="user", parts=[types.Part(function_response=function_response)]
        ),
    )

    session = Session(
        id="test_session",
        user_id="test_user",
        app_name="test_app",
        events=[call_event, response_event],
    )

    result = self.runner._find_agent_to_run(session, self.root_agent)
    assert result == self.sub_agent1

  def test_find_agent_to_run_returns_root_agent_when_no_events(self):
    """Test that root agent is returned when session has no non-user events."""
    session = Session(
        id="test_session",
        user_id="test_user",
        app_name="test_app",
        events=[
            Event(
                invocation_id="inv1",
                author="user",
                content=types.Content(
                    role="user", parts=[types.Part(text="Hello")]
                ),
            )
        ],
    )

    result = self.runner._find_agent_to_run(session, self.root_agent)
    assert result == self.root_agent

  def test_find_agent_to_run_returns_root_agent_when_found_in_events(self):
    """Test that root agent is returned when it's found in session events."""
    session = Session(
        id="test_session",
        user_id="test_user",
        app_name="test_app",
        events=[
            Event(
                invocation_id="inv1",
                author="root_agent",
                content=types.Content(
                    role="model", parts=[types.Part(text="Root response")]
                ),
            )
        ],
    )

    result = self.runner._find_agent_to_run(session, self.root_agent)
    assert result == self.root_agent

  def test_find_agent_to_run_returns_transferable_sub_agent(self):
    """Test that transferable sub agent is returned when found."""
    session = Session(
        id="test_session",
        user_id="test_user",
        app_name="test_app",
        events=[
            Event(
                invocation_id="inv1",
                author="sub_agent1",
                content=types.Content(
                    role="model", parts=[types.Part(text="Sub agent response")]
                ),
            )
        ],
    )

    result = self.runner._find_agent_to_run(session, self.root_agent)
    assert result == self.sub_agent1

  def test_find_agent_to_run_skips_non_transferable_agent(self):
    """Test that non-transferable agent is skipped and root agent is returned."""
    session = Session(
        id="test_session",
        user_id="test_user",
        app_name="test_app",
        events=[
            Event(
                invocation_id="inv1",
                author="non_transferable",
                content=types.Content(
                    role="model",
                    parts=[types.Part(text="Non-transferable response")],
                ),
            )
        ],
    )

    result = self.runner._find_agent_to_run(session, self.root_agent)
    assert result == self.root_agent

  def test_find_agent_to_run_skips_unknown_agent(self):
    """Test that unknown agent is skipped and root agent is returned."""
    session = Session(
        id="test_session",
        user_id="test_user",
        app_name="test_app",
        events=[
            Event(
                invocation_id="inv1",
                author="unknown_agent",
                content=types.Content(
                    role="model",
                    parts=[types.Part(text="Unknown agent response")],
                ),
            ),
            Event(
                invocation_id="inv2",
                author="root_agent",
                content=types.Content(
                    role="model", parts=[types.Part(text="Root response")]
                ),
            ),
        ],
    )

    result = self.runner._find_agent_to_run(session, self.root_agent)
    assert result == self.root_agent

  def test_find_agent_to_run_function_response_takes_precedence(self):
    """Test that function response scenario takes precedence over other logic."""
    # Create a function call from sub_agent2
    function_call = types.FunctionCall(id="func_456", name="test_func", args={})
    function_response = types.FunctionResponse(
        id="func_456", name="test_func", response={}
    )

    call_event = Event(
        invocation_id="inv1",
        author="sub_agent2",
        content=types.Content(
            role="model", parts=[types.Part(function_call=function_call)]
        ),
    )

    # Add another event from root_agent
    root_event = Event(
        invocation_id="inv2",
        author="root_agent",
        content=types.Content(
            role="model", parts=[types.Part(text="Root response")]
        ),
    )

    response_event = Event(
        invocation_id="inv3",
        author="user",
        content=types.Content(
            role="user", parts=[types.Part(function_response=function_response)]
        ),
    )

    session = Session(
        id="test_session",
        user_id="test_user",
        app_name="test_app",
        events=[call_event, root_event, response_event],
    )

    # Function-response routing only applies when resumability is enabled.
    self.runner.resumability_config = ResumabilityConfig(is_resumable=True)

    # Should return sub_agent2 due to function response, not root_agent
    result = self.runner._find_agent_to_run(session, self.root_agent)
    assert result == self.sub_agent2

  def test_find_agent_to_run_skips_function_response_when_not_resumable(self):
    """Test that function response scenario is skipped when not resumable."""
    function_call = types.FunctionCall(id="func_456", name="test_func", args={})
    function_response = types.FunctionResponse(
        id="func_456", name="test_func", response={}
    )

    call_event = Event(
        invocation_id="inv1",
        author="non_transferable",
        content=types.Content(
            role="model", parts=[types.Part(function_call=function_call)]
        ),
    )

    response_event = Event(
        invocation_id="inv2",
        author="user",
        content=types.Content(
            role="user", parts=[types.Part(function_response=function_response)]
        ),
    )

    session = Session(
        id="test_session",
        user_id="test_user",
        app_name="test_app",
        events=[call_event, response_event],
    )

    self.runner.resumability_config = ResumabilityConfig(is_resumable=False)

    result = self.runner._find_agent_to_run(session, self.root_agent)
    assert result == self.root_agent

  def test_find_agent_to_run_uses_function_response_when_resumable(self):
    """Test that function response scenario is used when resumable."""
    function_call = types.FunctionCall(id="func_456", name="test_func", args={})
    function_response = types.FunctionResponse(
        id="func_456", name="test_func", response={}
    )

    call_event = Event(
        invocation_id="inv1",
        author="non_transferable",
        content=types.Content(
            role="model", parts=[types.Part(function_call=function_call)]
        ),
    )

    response_event = Event(
        invocation_id="inv2",
        author="user",
        content=types.Content(
            role="user", parts=[types.Part(function_response=function_response)]
        ),
    )

    session = Session(
        id="test_session",
        user_id="test_user",
        app_name="test_app",
        events=[call_event, response_event],
    )

    self.runner.resumability_config = ResumabilityConfig(is_resumable=True)

    result = self.runner._find_agent_to_run(session, self.root_agent)
    assert result == self.non_transferable_agent

  def test_find_agent_to_run_resumable_unknown_function_call_author_falls_back(
      self,
  ):
    """Resumable routing must not return None for an unknown call author.

    When the matching function-call event is authored by something that is not
    an agent in the current hierarchy (e.g. "user", or a stale/foreign agent
    name carried over from a previous turn/session), `find_agent` returns None.
    Previously this None propagated to `build_node`, raising a confusing
    "Invalid node type: <class 'NoneType'>" error. We now fall through to the
    root agent instead.
    """
    function_call = types.FunctionCall(id="func_456", name="test_func", args={})
    function_response = types.FunctionResponse(
        id="func_456", name="test_func", response={}
    )

    # The function call is authored by "user", which is not an agent name.
    call_event = Event(
        invocation_id="inv1",
        author="user",
        content=types.Content(
            role="model", parts=[types.Part(function_call=function_call)]
        ),
    )

    response_event = Event(
        invocation_id="inv2",
        author="user",
        content=types.Content(
            role="user", parts=[types.Part(function_response=function_response)]
        ),
    )

    session = Session(
        id="test_session",
        user_id="test_user",
        app_name="test_app",
        events=[call_event, response_event],
    )

    self.runner.resumability_config = ResumabilityConfig(is_resumable=True)

    result = self.runner._find_agent_to_run(session, self.root_agent)
    assert result == self.root_agent

  def test_find_agent_to_run_resumable_stale_function_call_author_falls_back(
      self,
  ):
    """Resumable routing falls back to root for a stale/foreign call author."""
    function_call = types.FunctionCall(id="func_789", name="test_func", args={})
    function_response = types.FunctionResponse(
        id="func_789", name="test_func", response={}
    )

    # The function call is authored by an agent that is not in the hierarchy.
    call_event = Event(
        invocation_id="inv1",
        author="agent_from_a_previous_session",
        content=types.Content(
            role="model", parts=[types.Part(function_call=function_call)]
        ),
    )

    response_event = Event(
        invocation_id="inv2",
        author="user",
        content=types.Content(
            role="user", parts=[types.Part(function_response=function_response)]
        ),
    )

    session = Session(
        id="test_session",
        user_id="test_user",
        app_name="test_app",
        events=[call_event, response_event],
    )

    self.runner.resumability_config = ResumabilityConfig(is_resumable=True)

    result = self.runner._find_agent_to_run(session, self.root_agent)
    assert result == self.root_agent

  def test_is_transferable_across_agent_tree_with_llm_agent(self):
    """Test _is_transferable_across_agent_tree with LLM agent."""
    result = self.runner._is_transferable_across_agent_tree(self.sub_agent1)
    assert result is True

  def test_is_transferable_across_agent_tree_with_non_transferable_agent(self):
    """Test _is_transferable_across_agent_tree with non-transferable agent."""
    result = self.runner._is_transferable_across_agent_tree(
        self.non_transferable_agent
    )
    assert result is False

  def test_is_transferable_across_agent_tree_with_non_llm_agent(self):
    """Test _is_transferable_across_agent_tree with non-LLM agent."""
    non_llm_agent = MockAgent("non_llm_agent")
    # MockAgent inherits from BaseAgent, not LlmAgent, so it should return False
    result = self.runner._is_transferable_across_agent_tree(non_llm_agent)
    assert result is False


@pytest.mark.asyncio
async def test_session_not_found_message_includes_alignment_hint():

  class RunnerWithMismatch(Runner):

    def _infer_agent_origin(
        self, agent: BaseAgent
    ) -> tuple[Optional[str], Optional[Path]]:
      del agent
      return "expected_app", Path("/workspace/agents/expected_app")

  session_service = InMemorySessionService()
  runner = RunnerWithMismatch(
      app_name="configured_app",
      agent=MockLlmAgent("root_agent"),
      session_service=session_service,
      artifact_service=InMemoryArtifactService(),
  )

  agen = runner.run_async(
      user_id="user",
      session_id="missing",
      new_message=types.Content(role="user", parts=[]),
  )

  with pytest.raises(SessionNotFoundError) as excinfo:
    await agen.__anext__()

  await agen.aclose()

  message = str(excinfo.value)
  assert "Session not found" in message
  assert "configured_app" in message
  assert "expected_app" in message
  assert "Ensure the runner app_name matches" in message


@pytest.mark.asyncio
async def test_session_auto_creation():

  class RunnerWithMismatch(Runner):

    def _infer_agent_origin(
        self, agent: BaseAgent
    ) -> tuple[Optional[str], Optional[Path]]:
      del agent
      return "expected_app", Path("/workspace/agents/expected_app")

  session_service = InMemorySessionService()
  runner = RunnerWithMismatch(
      app_name="expected_app",
      agent=MockLlmAgent("test_agent"),
      session_service=session_service,
      artifact_service=InMemoryArtifactService(),
      auto_create_session=True,
  )

  agen = runner.run_async(
      user_id="user",
      session_id="missing",
      new_message=types.Content(role="user", parts=[types.Part(text="hi")]),
  )

  event = await agen.__anext__()
  await agen.aclose()

  # Verify that session_id="missing" doesn't error out - session is auto-created
  assert event.author == "test_agent"
  assert event.content.parts[0].text == "Test LLM response"


@pytest.mark.asyncio
async def test_rewind_auto_create_session_on_missing_session():
  """When auto_create_session=True, rewind should create session if missing.

  The newly created session won't contain the target invocation, so
  `rewind_async` should raise an Invocation ID not found error (rather than
  a session not found error), demonstrating auto-creation occurred.
  """
  session_service = InMemorySessionService()
  runner = Runner(
      app_name="auto_create_app",
      agent=MockLlmAgent("agent_for_rewind"),
      session_service=session_service,
      artifact_service=InMemoryArtifactService(),
      auto_create_session=True,
  )

  with pytest.raises(ValueError, match=r"Invocation ID not found: inv_missing"):
    await runner.rewind_async(
        user_id="user",
        session_id="missing",
        rewind_before_invocation_id="inv_missing",
    )

  # Verify the session actually exists now due to auto-creation.
  session = await session_service.get_session(
      app_name="auto_create_app", user_id="user", session_id="missing"
  )
  assert session is not None
  assert session.app_name == "auto_create_app"


@pytest.mark.asyncio
async def test_run_live_auto_create_session():
  """run_live should auto-create session when missing and yield events."""
  session_service = InMemorySessionService()
  artifact_service = InMemoryArtifactService()
  runner = Runner(
      app_name="live_app",
      agent=MockLiveAgent("live_agent"),
      session_service=session_service,
      artifact_service=artifact_service,
      auto_create_session=True,
  )

  # An empty LiveRequestQueue is sufficient for our mock agent.
  from google.adk.agents.live_request_queue import LiveRequestQueue

  live_queue = LiveRequestQueue()

  agen = runner.run_live(
      user_id="user",
      session_id="missing",
      live_request_queue=live_queue,
  )

  event = await agen.__anext__()
  await agen.aclose()

  assert event.author == "live_agent"
  assert event.content.parts[0].text == "live hello"

  # Session should have been created automatically.
  session = await session_service.get_session(
      app_name="live_app", user_id="user", session_id="missing"
  )
  assert session is not None


def test_run_passes_state_delta():
  """run should forward state_delta down to run_async."""
  import asyncio

  session_service = InMemorySessionService()
  runner = Runner(
      app_name=TEST_APP_ID,
      agent=MockAgent("test_agent"),
      session_service=session_service,
      artifact_service=InMemoryArtifactService(),
      auto_create_session=True,
  )

  state_delta = {"test_key": "test_value"}

  events = list(
      runner.run(
          user_id=TEST_USER_ID,
          session_id=TEST_SESSION_ID,
          new_message=types.Content(
              role="user", parts=[types.Part(text="hello")]
          ),
          state_delta=state_delta,
      )
  )

  assert len(events) >= 1

  session = asyncio.run(
      session_service.get_session(
          app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
      )
  )
  session_events = session.events

  user_event = next(e for e in session_events if e.author == "user")
  assert user_event.actions.state_delta == state_delta


@pytest.mark.asyncio
async def test_run_async_propagates_invocation_id():
  """run_async should propagate invocation_id to the invocation context and events."""

  session_service = InMemorySessionService()
  runner = Runner(
      app_name=TEST_APP_ID,
      agent=MockAgent("test_agent"),
      session_service=session_service,
      artifact_service=InMemoryArtifactService(),
      auto_create_session=True,
  )

  custom_invocation_id = "my_custom_invocation_id"

  agen = runner.run_async(
      user_id=TEST_USER_ID,
      session_id=TEST_SESSION_ID,
      new_message=types.Content(role="user", parts=[types.Part(text="hello")]),
      invocation_id=custom_invocation_id,
  )

  events = []
  async with aclosing(agen) as a:
    async for event in a:
      events.append(event)

  assert len(events) >= 1
  # Verify yielded events have the custom invocation ID
  for event in events:
    assert event.invocation_id == custom_invocation_id

  # Verify the session has the custom invocation ID in its events
  session = await session_service.get_session(
      app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
  )
  assert session is not None
  assert len(session.events) == 2
  for event in session.events:
    assert event.invocation_id == custom_invocation_id


@pytest.mark.asyncio
async def test_run_live_persists_event_callback_modifications():
  """run_live should persist the same event it streams after callback changes."""
  session_service = InMemorySessionService()
  artifact_service = InMemoryArtifactService()
  plugin = MockPlugin()
  plugin.enable_event_callback = True
  runner = Runner(
      app_name="live_app",
      agent=MockLiveAgent("live_agent"),
      session_service=session_service,
      artifact_service=artifact_service,
      plugins=[plugin],
  )
  await session_service.create_session(
      app_name="live_app", user_id="user", session_id="live_session"
  )

  from google.adk.agents.live_request_queue import LiveRequestQueue

  live_queue = LiveRequestQueue()
  agen = runner.run_live(
      user_id="user",
      session_id="live_session",
      live_request_queue=live_queue,
  )

  streamed_event = await agen.__anext__()
  await agen.aclose()

  session = await session_service.get_session(
      app_name="live_app", user_id="user", session_id="live_session"
  )
  persisted_event = session.events[0]

  assert streamed_event.author == "live_agent"
  assert streamed_event.invocation_id
  assert streamed_event.content.parts[0].text == (
      MockPlugin.ON_EVENT_CALLBACK_MSG
  )
  assert streamed_event.custom_metadata == MockPlugin.ON_EVENT_CALLBACK_METADATA

  assert persisted_event.id == streamed_event.id
  assert persisted_event.timestamp == streamed_event.timestamp
  assert persisted_event.author == streamed_event.author
  assert persisted_event.invocation_id == streamed_event.invocation_id
  assert persisted_event.content.parts[0].text == (
      MockPlugin.ON_EVENT_CALLBACK_MSG
  )
  assert (
      persisted_event.custom_metadata == MockPlugin.ON_EVENT_CALLBACK_METADATA
  )


@pytest.mark.asyncio
async def test_runner_allows_nested_agent_directories(tmp_path, monkeypatch):
  project_root = tmp_path / "workspace"
  agent_dir = project_root / "agents" / "examples" / "hello_world"
  agent_dir.mkdir(parents=True)
  # Make package structure importable.
  for pkg_dir in [
      project_root / "agents",
      project_root / "agents" / "examples",
      agent_dir,
  ]:
    (pkg_dir / "__init__.py").write_text("", encoding="utf-8")
  # Extra directories that previously confused origin inference, e.g. virtualenv.
  (project_root / "agents" / ".venv").mkdir()

  agent_source = textwrap.dedent("""\
      from google.adk.events.event import Event
      from google.adk.agents.base_agent import BaseAgent
      from google.genai import types


      class SimpleAgent(BaseAgent):

        def __init__(self):
          super().__init__(name='simplest_agent', sub_agents=[])

        async def _run_async_impl(self, invocation_context):
          yield Event(
              invocation_id=invocation_context.invocation_id,
              author=self.name,
              content=types.Content(
                  role='model',
                  parts=[types.Part(text='hello from nested')],
              ),
          )


      root_agent = SimpleAgent()
      """)
  (agent_dir / "agent.py").write_text(agent_source, encoding="utf-8")

  monkeypatch.chdir(project_root)
  loader = AgentLoader(agents_dir="agents/examples")
  loaded_agent = loader.load_agent("hello_world")

  assert isinstance(loaded_agent, BaseAgent)
  session_service = InMemorySessionService()
  artifact_service = InMemoryArtifactService()
  runner = Runner(
      app_name="hello_world",
      agent=loaded_agent,
      session_service=session_service,
      artifact_service=artifact_service,
  )
  assert runner._app_name_alignment_hint is None

  session = await session_service.create_session(
      app_name="hello_world",
      user_id="user",
  )
  agen = runner.run_async(
      user_id=session.user_id,
      session_id=session.id,
      new_message=types.Content(
          role="user",
          parts=[types.Part(text="hi")],
      ),
  )
  event = await agen.__anext__()
  await agen.aclose()

  assert event.author == "simplest_agent"
  assert event.content
  assert event.content.parts
  assert event.content.parts[0].text == "hello from nested"


@pytest.mark.asyncio
async def test_run_config_custom_metadata_propagates_to_events():
  session_service = InMemorySessionService()
  runner = Runner(
      app_name=TEST_APP_ID,
      agent=MockAgentWithMetadata("metadata_agent"),
      session_service=session_service,
      artifact_service=InMemoryArtifactService(),
  )
  await session_service.create_session(
      app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
  )

  run_config = RunConfig(custom_metadata={"request_id": "req-1"})
  events = [
      event
      async for event in runner.run_async(
          user_id=TEST_USER_ID,
          session_id=TEST_SESSION_ID,
          new_message=types.Content(role="user", parts=[types.Part(text="hi")]),
          run_config=run_config,
      )
  ]

  assert events[0].custom_metadata is not None
  assert events[0].custom_metadata["request_id"] == "req-1"
  assert events[0].custom_metadata["event_key"] == "event_value"

  session = await session_service.get_session(
      app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
  )
  user_event = next(event for event in session.events if event.author == "user")
  assert user_event.custom_metadata == {"request_id": "req-1"}


@pytest.mark.asyncio
async def test_run_config_custom_metadata_stamps_user_event_in_chat_mode():
  """LlmAgent chat path stamps the user event with run-level custom_metadata."""
  session_service = InMemorySessionService()

  def _before_agent_callback(callback_context) -> types.Content:
    del callback_context  # Unused; short-circuits the model call.
    return types.Content(role="model", parts=[types.Part(text="hi back")])

  agent = LlmAgent(
      name="chat_agent", before_agent_callback=_before_agent_callback
  )
  runner = Runner(
      app_name=TEST_APP_ID, agent=agent, session_service=session_service
  )
  await session_service.create_session(
      app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
  )

  run_config = RunConfig(custom_metadata={"turn_id": "t-1"})
  async for _ in runner.run_async(
      user_id=TEST_USER_ID,
      session_id=TEST_SESSION_ID,
      new_message=types.Content(role="user", parts=[types.Part(text="hi")]),
      run_config=run_config,
  ):
    pass

  session = await session_service.get_session(
      app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
  )
  user_event = next(event for event in session.events if event.author == "user")
  assert user_event.custom_metadata == {"turn_id": "t-1"}


@pytest.mark.asyncio
async def test_chat_mode_fetches_session_once_per_turn():
  """Root LlmAgent chat path reuses the prologue fetch inside the node run."""
  session_service = InMemorySessionService()

  def _before_agent_callback(callback_context) -> types.Content:
    del callback_context  # Unused; short-circuits the model call.
    return types.Content(role="model", parts=[types.Part(text="hi back")])

  agent = LlmAgent(
      name="chat_agent", before_agent_callback=_before_agent_callback
  )
  runner = Runner(
      app_name=TEST_APP_ID, agent=agent, session_service=session_service
  )
  original_get_session = session_service.get_session
  await session_service.create_session(
      app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
  )

  spy = AsyncMock(wraps=session_service.get_session)
  session_service.get_session = spy

  async for _ in runner.run_async(
      user_id=TEST_USER_ID,
      session_id=TEST_SESSION_ID,
      new_message=types.Content(role="user", parts=[types.Part(text="hi")]),
  ):
    pass

  assert spy.call_count == 1

  # Correctness: the user message is still persisted despite the single fetch.
  session = await original_get_session(
      app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
  )
  assert any(event.author == "user" for event in session.events)


@pytest.mark.asyncio
async def test_chat_mode_honors_get_session_config():
  """Root LlmAgent chat path threads get_session_config into the fetch."""
  session_service = InMemorySessionService()

  def _before_agent_callback(callback_context) -> types.Content:
    del callback_context  # Unused; short-circuits the model call.
    return types.Content(role="model", parts=[types.Part(text="hi back")])

  agent = LlmAgent(
      name="chat_agent", before_agent_callback=_before_agent_callback
  )
  runner = Runner(
      app_name=TEST_APP_ID, agent=agent, session_service=session_service
  )
  session = await session_service.create_session(
      app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
  )
  for i in range(3):
    await session_service.append_event(
        session=session,
        event=Event(
            invocation_id=f"seed-{i}",
            author="user",
            content=types.Content(
                role="user", parts=[types.Part(text=f"seed-{i}")]
            ),
        ),
    )

  seen_configs = []
  seen_event_counts = []
  original_get_session = session_service.get_session

  async def _spy_get_session(*args, **kwargs):
    fetched = await original_get_session(*args, **kwargs)
    seen_configs.append(kwargs.get("config"))
    seen_event_counts.append(None if fetched is None else len(fetched.events))
    return fetched

  session_service.get_session = _spy_get_session

  run_config = RunConfig(
      get_session_config=GetSessionConfig(num_recent_events=1)
  )
  async for _ in runner.run_async(
      user_id=TEST_USER_ID,
      session_id=TEST_SESSION_ID,
      new_message=types.Content(role="user", parts=[types.Part(text="hi")]),
      run_config=run_config,
  ):
    pass

  assert seen_configs
  assert all(
      config == GetSessionConfig(num_recent_events=1) for config in seen_configs
  )
  # num_recent_events=1 bounds the fetched history to the single latest event.
  assert all(count == 1 for count in seen_event_counts)


class TestRunnerWithPlugins:
  """Tests for Runner with plugins."""

  def setup_method(self):
    self.plugin = MockPlugin()
    self.session_service = InMemorySessionService()
    self.artifact_service = InMemoryArtifactService()
    self.root_agent = MockLlmAgent("root_agent")
    self.runner = Runner(
        app_name="test_app",
        agent=MockLlmAgent("test_agent"),
        session_service=self.session_service,
        artifact_service=self.artifact_service,
        plugins=[self.plugin],
    )

  async def run_test(self, original_user_input="Hello") -> list[Event]:
    """Prepares the test by creating a session and running the runner."""
    await self.session_service.create_session(
        app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
    )
    events = []
    async for event in self.runner.run_async(
        user_id=TEST_USER_ID,
        session_id=TEST_SESSION_ID,
        new_message=types.Content(
            role="user", parts=[types.Part(text=original_user_input)]
        ),
    ):
      events.append(event)
    return events

  @pytest.mark.asyncio
  async def test_runner_is_initialized_with_plugins(self):
    """Test that the runner is initialized with plugins."""
    await self.run_test()

    assert self.runner.plugin_manager is not None

  @pytest.mark.asyncio
  async def test_runner_modifies_user_message_before_execution(self):
    """Test that the runner modifies the user message before execution."""
    original_user_input = "original_input"
    self.plugin.enable_user_message_callback = True

    await self.run_test(original_user_input=original_user_input)
    session = await self.session_service.get_session(
        app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
    )
    generated_event = session.events[0]
    modified_user_message = generated_event.content.parts[0].text

    assert modified_user_message == MockPlugin.ON_USER_CALLBACK_MSG
    assert self.plugin.user_content_seen_in_before_run_callback is not None
    assert (
        self.plugin.user_content_seen_in_before_run_callback.parts[0].text
        == MockPlugin.ON_USER_CALLBACK_MSG
    )

  @pytest.mark.asyncio
  async def test_runner_modifies_event_after_execution(self):
    """Test that the runner modifies the event after execution."""
    self.plugin.enable_event_callback = True

    events = await self.run_test()
    generated_event = events[0]
    modified_event_message = generated_event.content.parts[0].text

    assert modified_event_message == MockPlugin.ON_EVENT_CALLBACK_MSG

  @pytest.mark.asyncio
  async def test_runner_persists_event_callback_modifications(self):
    """Event callback output should be persisted, not only streamed."""
    self.plugin.enable_event_callback = True

    events = await self.run_test()
    streamed_event = events[0]

    session = await self.session_service.get_session(
        app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
    )
    persisted_event = session.events[1]

    assert streamed_event.author == "test_agent"
    assert streamed_event.invocation_id
    assert streamed_event.content.parts[0].text == (
        MockPlugin.ON_EVENT_CALLBACK_MSG
    )
    assert (
        streamed_event.custom_metadata == MockPlugin.ON_EVENT_CALLBACK_METADATA
    )

    assert persisted_event.id == streamed_event.id
    assert persisted_event.timestamp == streamed_event.timestamp
    assert persisted_event.author == streamed_event.author
    assert persisted_event.invocation_id == streamed_event.invocation_id
    assert persisted_event.content.parts[0].text == (
        MockPlugin.ON_EVENT_CALLBACK_MSG
    )
    assert (
        persisted_event.custom_metadata == MockPlugin.ON_EVENT_CALLBACK_METADATA
    )

  @pytest.mark.asyncio
  async def test_runner_close_calls_plugin_close(self):
    """Test that runner.close() calls plugin manager close."""
    # Mock the plugin manager's close method
    self.runner.plugin_manager.close = AsyncMock()

    await self.runner.close()

    self.runner.plugin_manager.close.assert_awaited_once()

  @pytest.mark.asyncio
  async def test_runner_close_does_not_cancel_toolset_cleanup(self):
    """Caller cancellation should not cancel an in-flight toolset close."""

    class SlowCloseToolset(BaseToolset):

      def __init__(self):
        super().__init__()
        self.close_started = asyncio.Event()
        self.close_finished = asyncio.Event()
        self.close_cancelled = False

      async def get_tools(self, readonly_context=None):
        del readonly_context
        return []

      async def close(self) -> None:
        self.close_started.set()
        try:
          await asyncio.sleep(0.05)
          self.close_finished.set()
        except asyncio.CancelledError:
          self.close_cancelled = True
          raise

    toolset = SlowCloseToolset()
    runner = Runner(
        app_name="test_app",
        agent=LlmAgent(
            name="test_agent", model="gemini-1.5-pro", tools=[toolset]
        ),
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )

    close_task = asyncio.create_task(runner.close())
    await toolset.close_started.wait()
    close_task.cancel()

    with pytest.raises(asyncio.CancelledError):
      await close_task

    assert close_task.cancelled() is True
    assert toolset.close_cancelled is False
    assert toolset.close_finished.is_set()

  @pytest.mark.asyncio
  async def test_runner_passes_plugin_close_timeout(self):
    """Test that runner passes plugin_close_timeout to PluginManager."""
    runner = Runner(
        app_name="test_app",
        agent=MockLlmAgent("test_agent"),
        session_service=self.session_service,
        artifact_service=self.artifact_service,
        plugins=[self.plugin],
        plugin_close_timeout=10.0,
    )
    assert runner.plugin_manager._close_timeout == 10.0

  @pytest.mark.filterwarnings(
      "ignore:The `plugins` argument is deprecated:DeprecationWarning"
  )
  def test_runner_init_raises_error_with_app_and_agent(self):
    """Test that ValueError is raised when app and agent are provided."""
    with pytest.raises(
        ValueError,
        match="Only one of app, agent, or node may be provided.",
    ):
      Runner(
          app=App(name="test_app", root_agent=self.root_agent),
          agent=self.root_agent,
          session_service=self.session_service,
          artifact_service=self.artifact_service,
      )

  @pytest.mark.filterwarnings(
      "ignore:The `plugins` argument is deprecated:DeprecationWarning"
  )
  def test_runner_init_allows_app_name_override_with_app(self):
    """Test that app_name can override app.name when both are provided."""
    app = App(name="test_app", root_agent=self.root_agent)
    runner = Runner(
        app=app,
        app_name="override_name",
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )
    assert runner.app_name == "override_name"
    assert runner.agent == self.root_agent
    assert runner.app == app

  def test_runner_init_raises_error_without_app_and_app_name(self):
    """Test ValueError is raised when app is not provided and app_name is missing."""
    with pytest.raises(
        ValueError,
        match=(
            "app_name is required when agent is provided|One of app, agent, or"
            " node must be provided"
        ),
    ):
      Runner(
          agent=self.root_agent,
          session_service=self.session_service,
          artifact_service=self.artifact_service,
      )

  def test_runner_init_raises_error_without_app_and_agent(self):
    """Test ValueError is raised when app is not provided and agent is missing."""
    with pytest.raises(
        ValueError,
        match=(
            "app_name is required when agent is provided|One of app, agent, or"
            " node must be provided"
        ),
    ):
      Runner(
          app_name="test_app",
          session_service=self.session_service,
          artifact_service=self.artifact_service,
      )


class TestRunnerCacheConfig:
  """Tests for Runner cache config extraction and handling."""

  def setup_method(self):
    """Set up test fixtures."""
    self.session_service = InMemorySessionService()
    self.artifact_service = InMemoryArtifactService()
    self.root_agent = MockLlmAgent("root_agent")

  def test_runner_extracts_cache_config_from_app(self):
    """Test that Runner extracts cache config from App."""
    cache_config = ContextCacheConfig(
        cache_intervals=15, ttl_seconds=3600, min_tokens=1024
    )

    app = App(
        name="test_app",
        root_agent=self.root_agent,
        context_cache_config=cache_config,
    )

    runner = Runner(
        app=app,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )

    assert runner.context_cache_config == cache_config
    assert runner.context_cache_config.cache_intervals == 15
    assert runner.context_cache_config.ttl_seconds == 3600
    assert runner.context_cache_config.min_tokens == 1024

  def test_runner_with_app_without_cache_config(self):
    """Test Runner with App that has no cache config."""
    app = App(
        name="test_app", root_agent=self.root_agent, context_cache_config=None
    )

    runner = Runner(
        app=app,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )

    assert runner.context_cache_config is None

  def test_runner_without_app_has_no_cache_config(self):
    """Test Runner created without App has no cache config."""
    runner = Runner(
        app_name="test_app",
        agent=self.root_agent,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )

    assert runner.context_cache_config is None

  def test_runner_cache_config_passed_to_invocation_context(self):
    """Test that cache config is passed to InvocationContext."""
    cache_config = ContextCacheConfig(
        cache_intervals=20, ttl_seconds=7200, min_tokens=2048
    )

    app = App(
        name="test_app",
        root_agent=self.root_agent,
        context_cache_config=cache_config,
    )

    runner = Runner(
        app=app,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )

    # Create a mock session
    mock_session = Session(
        id=TEST_SESSION_ID,
        app_name=TEST_APP_ID,
        user_id=TEST_USER_ID,
        events=[],
    )

    # Create invocation context using runner's method
    invocation_context = runner._new_invocation_context(mock_session)

    assert invocation_context.context_cache_config == cache_config
    assert invocation_context.context_cache_config.cache_intervals == 20

  def test_runner_validate_params_return_order(self):
    """Test that _validate_runner_params returns values in correct order."""
    cache_config = ContextCacheConfig(cache_intervals=25)

    app = App(
        name="order_test_app",
        root_agent=self.root_agent,
        context_cache_config=cache_config,
        resumability_config=ResumabilityConfig(is_resumable=True),
    )

    runner = Runner(
        app=app,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )

    # Test the validation method directly
    app_name, agent, context_cache_config, resumability_config, plugins = (
        runner._validate_runner_params(app, None, None, None)
    )

    assert app_name == "order_test_app"
    assert agent == self.root_agent
    assert context_cache_config == cache_config
    assert context_cache_config.cache_intervals == 25
    assert resumability_config == app.resumability_config
    assert plugins == []

  def test_runner_validate_params_without_app(self):
    """Test _validate_runner_params without App returns None for cache config."""
    runner = Runner(
        app_name="test_app",
        agent=self.root_agent,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )

    app_name, agent, context_cache_config, resumability_config, plugins = (
        runner._validate_runner_params(None, "test_app", self.root_agent, None)
    )

    assert app_name == "test_app"
    assert agent == self.root_agent
    assert context_cache_config is None
    assert resumability_config is None
    assert plugins is None

  def test_runner_app_name_and_agent_extracted_correctly(self):
    """Test that app_name and agent are correctly extracted from App."""
    cache_config = ContextCacheConfig()

    app = App(
        name="extracted_app",
        root_agent=self.root_agent,
        context_cache_config=cache_config,
    )

    runner = Runner(
        app=app,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )

    assert runner.app_name == "extracted_app"
    assert runner.agent == self.root_agent
    assert runner.context_cache_config == cache_config

  def test_runner_realistic_cache_config_scenario(self):
    """Test realistic scenario with production-like cache config."""
    # Production cache config
    production_cache_config = ContextCacheConfig(
        cache_intervals=30,
        ttl_seconds=14400,
        min_tokens=4096,  # 4 hours
    )

    app = App(
        name="production_app",
        root_agent=self.root_agent,
        context_cache_config=production_cache_config,
    )

    runner = Runner(
        app=app,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )

    # Verify all settings are preserved
    assert runner.context_cache_config.cache_intervals == 30
    assert runner.context_cache_config.ttl_seconds == 14400
    assert runner.context_cache_config.ttl_string == "14400s"
    assert runner.context_cache_config.min_tokens == 4096

    # Verify string representation
    expected_str = (
        "ContextCacheConfig(cache_intervals=30, ttl=14400s, min_tokens=4096, "
        "create_http_options=None)"
    )
    assert str(runner.context_cache_config) == expected_str


class TestRunnerResolveApp:
  """Tests for Runner._resolve_app and node support."""

  def setup_method(self):
    self.session_service = InMemorySessionService()
    self.artifact_service = InMemoryArtifactService()
    self.root_agent = MockLlmAgent("root_agent")

  def test_resolve_app_with_agent_wraps_in_app(self):
    """Test that a bare agent is wrapped into an App."""
    runner = Runner(
        app_name="test_app",
        agent=self.root_agent,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )
    assert runner.app is not None
    assert runner.app.root_agent is self.root_agent
    assert runner.app_name == "test_app"
    assert runner.agent is self.root_agent

  def test_resolve_app_with_node_wraps_in_app(self):
    """Test that a bare node is wrapped into an App."""
    from google.adk.workflow._base_node import BaseNode

    node = BaseNode(name="test_node")
    runner = Runner(
        node=node,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )
    assert runner.app is not None
    assert runner.app.root_agent is node
    assert runner.app_name == "test_node"
    assert runner.agent is node

  def test_resolve_app_with_node_and_app_name(self):
    """Test that app_name overrides node.name."""
    from google.adk.workflow._base_node import BaseNode

    node = BaseNode(name="node_name")
    runner = Runner(
        app_name="custom_name",
        node=node,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )
    assert runner.app_name == "custom_name"

  def test_resolve_app_rejects_app_and_agent(self):
    """Test that providing both app and agent raises."""
    app = App(name="test_app", root_agent=self.root_agent)
    with pytest.raises(ValueError, match="Only one of app, agent, or node"):
      Runner(
          app=app,
          agent=self.root_agent,
          session_service=self.session_service,
      )

  def test_resolve_app_rejects_app_and_node(self):
    """Test that providing both app and node raises."""
    from google.adk.workflow._base_node import BaseNode

    app = App(name="test_app", root_agent=self.root_agent)
    node = BaseNode(name="test_node")
    with pytest.raises(ValueError, match="Only one of app, agent, or node"):
      Runner(
          app=app,
          node=node,
          session_service=self.session_service,
      )

  def test_resolve_app_rejects_agent_and_node(self):
    """Test that providing both agent and node raises."""
    from google.adk.workflow._base_node import BaseNode

    node = BaseNode(name="test_node")
    with pytest.raises(ValueError, match="Only one of app, agent, or node"):
      Runner(
          app_name="test_app",
          agent=self.root_agent,
          node=node,
          session_service=self.session_service,
      )

  def test_resolve_app_rejects_none(self):
    """Test that providing no app, agent, or node raises."""
    with pytest.raises(
        ValueError, match="One of app, agent, or node must be provided"
    ):
      Runner(
          app_name="test_app",
          session_service=self.session_service,
      )

  def test_resolve_app_extracts_node_from_app(self):
    """Test that Runner extracts node from App into agent field."""
    from google.adk.workflow._base_node import BaseNode

    node = BaseNode(name="test_node")
    app = App(name="test_app", root_agent=node)
    runner = Runner(
        app=app,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )
    assert runner.agent is node
    assert runner.app_name == "test_app"
    assert runner.context_cache_config is None
    assert runner.resumability_config is None


class TestRunnerShouldAppendEvent:
  """Tests for Runner._should_append_event method."""

  def setup_method(self):
    """Set up test fixtures."""
    self.session_service = InMemorySessionService()
    self.artifact_service = InMemoryArtifactService()
    self.root_agent = MockLlmAgent("root_agent")
    self.runner = Runner(
        app_name="test_app",
        agent=self.root_agent,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )

  def test_should_append_event_finished_input_transcription(self):
    event = Event(
        invocation_id="inv1",
        author="user",
        input_transcription=types.Transcription(text="hello", finished=True),
    )
    assert self.runner._should_append_event(event, is_live_call=True) is True

  def test_should_append_event_unfinished_input_transcription(self):
    event = Event(
        invocation_id="inv1",
        author="user",
        input_transcription=types.Transcription(text="hello", finished=False),
    )
    assert self.runner._should_append_event(event, is_live_call=True) is True

  def test_should_append_event_finished_output_transcription(self):
    event = Event(
        invocation_id="inv1",
        author="model",
        output_transcription=types.Transcription(text="world", finished=True),
    )
    assert self.runner._should_append_event(event, is_live_call=True) is True

  def test_should_append_event_unfinished_output_transcription(self):
    event = Event(
        invocation_id="inv1",
        author="model",
        output_transcription=types.Transcription(text="world", finished=False),
    )
    assert self.runner._should_append_event(event, is_live_call=True) is True

  def test_should_not_append_event_live_model_audio(self):
    event = Event(
        invocation_id="inv1",
        author="model",
        content=types.Content(
            parts=[
                types.Part(
                    inline_data=types.Blob(data=b"123", mime_type="audio/pcm")
                )
            ]
        ),
    )
    assert self.runner._should_append_event(event, is_live_call=True) is False

  def test_should_append_event_non_live_model_audio(self):
    event = Event(
        invocation_id="inv1",
        author="model",
        content=types.Content(
            parts=[
                types.Part(
                    inline_data=types.Blob(data=b"123", mime_type="audio/pcm")
                )
            ]
        ),
    )
    assert self.runner._should_append_event(event, is_live_call=False) is True

  def test_should_append_event_other_event(self):
    event = Event(
        invocation_id="inv1",
        author="model",
        content=types.Content(parts=[types.Part(text="text")]),
    )
    assert self.runner._should_append_event(event, is_live_call=True) is True

  def test_should_not_append_event_live_model_video(self):
    event = Event(
        invocation_id="inv1",
        author="model",
        content=types.Content(
            parts=[
                types.Part(
                    inline_data=types.Blob(data=b"123", mime_type="video/mp4")
                )
            ]
        ),
    )
    assert self.runner._should_append_event(event, is_live_call=True) is False

  def test_should_append_event_non_live_model_video(self):
    event = Event(
        invocation_id="inv1",
        author="model",
        content=types.Content(
            parts=[
                types.Part(
                    inline_data=types.Blob(data=b"123", mime_type="video/mp4")
                )
            ]
        ),
    )
    assert self.runner._should_append_event(event, is_live_call=False) is True


@pytest.fixture
def user_agent_module(tmp_path, monkeypatch):
  """Fixture that creates a temporary user agent module for testing.

  Yields a callable that creates an agent module with the given name and
  returns the loaded agent.
  """
  created_modules = []
  original_path = None

  def _create_agent(agent_dir_name: str):
    nonlocal original_path
    agent_dir = tmp_path / "agents" / agent_dir_name
    agent_dir.mkdir(parents=True, exist_ok=True)
    (tmp_path / "agents" / "__init__.py").write_text("", encoding="utf-8")
    (agent_dir / "__init__.py").write_text("", encoding="utf-8")

    agent_source = f"""\
from google.adk.agents.llm_agent import LlmAgent

class MyAgent(LlmAgent):
    pass

root_agent = MyAgent(name="{agent_dir_name}", model="gemini-2.5-flash")
"""
    (agent_dir / "agent.py").write_text(agent_source, encoding="utf-8")

    monkeypatch.chdir(tmp_path)
    if original_path is None:
      original_path = str(tmp_path)
      sys.path.insert(0, original_path)

    module_name = f"agents.{agent_dir_name}.agent"
    module = importlib.import_module(module_name)
    created_modules.append(module_name)
    return module.root_agent

  yield _create_agent

  # Cleanup
  if original_path and original_path in sys.path:
    sys.path.remove(original_path)
  for mod_name in list(sys.modules.keys()):
    if mod_name.startswith("agents"):
      del sys.modules[mod_name]


class TestRunnerInferAgentOrigin:
  """Tests for Runner._infer_agent_origin method."""

  def setup_method(self):
    """Set up test fixtures."""
    self.session_service = InMemorySessionService()
    self.artifact_service = InMemoryArtifactService()

  def test_infer_agent_origin_uses_adk_metadata_when_available(self):
    """Test that _infer_agent_origin uses _adk_origin_* metadata when set."""
    agent = MockLlmAgent("test_agent")
    # Simulate metadata set by AgentLoader
    agent._adk_origin_app_name = "my_app"
    agent._adk_origin_path = Path("/workspace/agents/my_app")

    runner = Runner(
        app_name="my_app",
        agent=agent,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )

    origin_name, origin_path = runner._infer_agent_origin(agent)
    assert origin_name == "my_app"
    assert origin_path == Path("/workspace/agents/my_app")

  def test_infer_agent_origin_no_false_positive_for_direct_llm_agent(self):
    """Test that using LlmAgent directly doesn't trigger mismatch warning.

    Regression test for GitHub issue #3143: Users who instantiate LlmAgent
    directly and run from a directory that is a parent of the ADK installation
    were getting false positive 'App name mismatch' warnings.

    This also verifies that _infer_agent_origin returns None for ADK internal
    modules (google.adk.*).
    """
    agent = LlmAgent(
        name="my_custom_agent",
        model="gemini-2.5-flash",
    )

    runner = Runner(
        app_name="my_custom_agent",
        agent=agent,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )

    # Should return None for ADK internal modules
    origin_name, _ = runner._infer_agent_origin(agent)
    assert origin_name is None
    # No mismatch warning should be generated
    assert runner._app_name_alignment_hint is None

  def test_infer_agent_origin_with_subclassed_agent_in_user_code(
      self, user_agent_module
  ):
    """Test that subclassed agents in user code still trigger origin inference."""
    agent = user_agent_module("my_agent")

    runner = Runner(
        app_name="my_agent",
        agent=agent,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )

    # Should infer origin correctly from user's code
    origin_name, origin_path = runner._infer_agent_origin(agent)
    assert origin_name == "my_agent"
    assert runner._app_name_alignment_hint is None

  def test_infer_agent_origin_detects_mismatch_for_user_agent(
      self, user_agent_module
  ):
    """Test that mismatched app_name is detected for user-defined agents."""
    agent = user_agent_module("actual_name")

    runner = Runner(
        app_name="wrong_name",  # Intentionally wrong
        agent=agent,
        session_service=self.session_service,
        artifact_service=self.artifact_service,
    )

    # Should detect the mismatch
    assert runner._app_name_alignment_hint is not None
    assert "wrong_name" in runner._app_name_alignment_hint
    assert "actual_name" in runner._app_name_alignment_hint


@pytest.mark.asyncio
async def test_run_async_passes_get_session_config():
  """run_async should forward RunConfig.get_session_config to get_session."""
  from google.adk.sessions.base_session_service import GetSessionConfig

  session_service = InMemorySessionService()

  # Pre-create a session with multiple events.
  session = await session_service.create_session(
      app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
  )
  for i in range(10):
    await session_service.append_event(
        session=session,
        event=Event(
            invocation_id=f"inv_{i}",
            author="user",
            content=types.Content(
                role="user", parts=[types.Part(text=f"message {i}")]
            ),
        ),
    )

  runner = Runner(
      app_name=TEST_APP_ID,
      agent=MockAgent("test_agent"),
      session_service=session_service,
      artifact_service=InMemoryArtifactService(),
  )

  # Run with num_recent_events=3 to only load recent events.
  config = RunConfig(
      get_session_config=GetSessionConfig(num_recent_events=3),
  )

  events = []
  async for event in runner.run_async(
      user_id=TEST_USER_ID,
      session_id=TEST_SESSION_ID,
      new_message=types.Content(role="user", parts=[types.Part(text="hello")]),
      run_config=config,
  ):
    events.append(event)

  # Agent should still produce output (session was found).
  assert len(events) >= 1
  assert events[0].author == "test_agent"


@pytest.mark.asyncio
async def test_run_async_teardown_on_aclose():
  """Closing run_async generator using aclose() should abort and cancel the running agent task."""
  import asyncio

  session_service = InMemorySessionService()
  artifact_service = InMemoryArtifactService()

  was_cancelled = {"value": False}

  class CancellingAgent(BaseAgent):

    def __init__(self, name: str):
      super().__init__(name=name, sub_agents=[])

    async def _run_async_impl(
        self, invocation_context: InvocationContext
    ) -> AsyncGenerator[Event, None]:
      try:
        yield Event(
            invocation_id=invocation_context.invocation_id,
            author=self.name,
            content=types.Content(
                role="model", parts=[types.Part(text="First response")]
            ),
        )
        # Block simulating slow ongoing task
        await asyncio.sleep(5.0)
        yield Event(
            invocation_id=invocation_context.invocation_id,
            author=self.name,
            content=types.Content(
                role="model", parts=[types.Part(text="Second response")]
            ),
        )
      except (asyncio.CancelledError, GeneratorExit):
        was_cancelled["value"] = True
        raise

  runner = Runner(
      app_name=TEST_APP_ID,
      agent=CancellingAgent("cancel_agent"),
      session_service=session_service,
      artifact_service=artifact_service,
      auto_create_session=True,
  )

  # Given a run session
  agen = runner.run_async(
      user_id=TEST_USER_ID,
      session_id=TEST_SESSION_ID,
      new_message=types.Content(role="user", parts=[types.Part(text="hello")]),
  )

  # When the client reads the first event and then calls aclose()
  event = await agen.__anext__()
  assert event.content.parts[0].text == "First response"

  await agen.aclose()

  # Then the running agent was immediately aborted and cancelled
  assert was_cancelled["value"] is True


@pytest.mark.asyncio
async def test_run_live_passes_get_session_config():
  """run_live should forward RunConfig.get_session_config to get_session."""
  from google.adk.agents.live_request_queue import LiveRequestQueue
  from google.adk.sessions.base_session_service import GetSessionConfig

  session_service = InMemorySessionService()

  # Pre-create session.
  await session_service.create_session(
      app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
  )

  runner = Runner(
      app_name=TEST_APP_ID,
      agent=MockLiveAgent("live_agent"),
      session_service=session_service,
      artifact_service=InMemoryArtifactService(),
  )

  config = RunConfig(
      get_session_config=GetSessionConfig(num_recent_events=5),
  )

  live_queue = LiveRequestQueue()
  agen = runner.run_live(
      user_id=TEST_USER_ID,
      session_id=TEST_SESSION_ID,
      live_request_queue=live_queue,
      run_config=config,
  )

  event = await agen.__anext__()
  await agen.aclose()

  assert event.author == "live_agent"
  assert event.content.parts[0].text == "live hello"


@pytest.mark.asyncio
async def test_rewind_async_passes_get_session_config():
  """rewind_async should forward RunConfig.get_session_config to get_session."""
  from google.adk.sessions.base_session_service import GetSessionConfig

  session_service = InMemorySessionService()

  runner = Runner(
      app_name=TEST_APP_ID,
      agent=MockAgent("test_agent"),
      session_service=session_service,
      artifact_service=InMemoryArtifactService(),
      auto_create_session=True,
  )

  config = RunConfig(
      get_session_config=GetSessionConfig(num_recent_events=5),
  )

  # rewind_async on a fresh session will raise because the invocation_id
  # doesn't exist, but it demonstrates that the config path works.
  with pytest.raises(ValueError, match=r"Invocation ID not found"):
    await runner.rewind_async(
        user_id=TEST_USER_ID,
        session_id="new_session",
        rewind_before_invocation_id="inv_missing",
        run_config=config,
    )


@pytest.mark.asyncio
async def test_run_debug_passes_get_session_config():
  """run_debug should forward RunConfig.get_session_config to get_session."""
  from google.adk.sessions.base_session_service import GetSessionConfig

  session_service = InMemorySessionService()

  runner = Runner(
      app_name=TEST_APP_ID,
      agent=MockAgent("test_agent"),
      session_service=session_service,
      artifact_service=InMemoryArtifactService(),
  )

  config = RunConfig(
      get_session_config=GetSessionConfig(num_recent_events=5),
  )

  events = await runner.run_debug(
      "hello",
      run_config=config,
      quiet=True,
  )

  assert len(events) >= 1
  assert events[0].author == "test_agent"


@pytest.mark.asyncio
async def test_get_session_config_limits_events():
  """Verify that num_recent_events actually limits loaded events."""
  from google.adk.sessions.base_session_service import GetSessionConfig

  session_service = InMemorySessionService()

  # Create session and add events.
  session = await session_service.create_session(
      app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
  )
  for i in range(10):
    await session_service.append_event(
        session=session,
        event=Event(
            invocation_id=f"inv_{i}",
            author="user",
            content=types.Content(
                role="user", parts=[types.Part(text=f"message {i}")]
            ),
        ),
    )

  # Without config: should load all events.
  full_session = await session_service.get_session(
      app_name=TEST_APP_ID, user_id=TEST_USER_ID, session_id=TEST_SESSION_ID
  )
  assert len(full_session.events) == 10

  # With config: should limit events.
  limited_session = await session_service.get_session(
      app_name=TEST_APP_ID,
      user_id=TEST_USER_ID,
      session_id=TEST_SESSION_ID,
      config=GetSessionConfig(num_recent_events=3),
  )
  assert len(limited_session.events) == 3


@pytest.mark.asyncio
async def test_run_async_rejects_user_function_call():
  """Verify that runner rejects user-authored messages with function calls."""
  session_service = InMemorySessionService()
  runner = Runner(
      app_name=TEST_APP_ID,
      agent=MockAgent("test_agent"),
      session_service=session_service,
      artifact_service=InMemoryArtifactService(),
      auto_create_session=True,
  )

  malicious_message = types.Content(
      role="user",
      parts=[
          types.Part(
              function_call=types.FunctionCall(
                  name="some_tool",
                  args={"key": "value"},
              )
          )
      ],
  )

  agen = runner.run_async(
      user_id=TEST_USER_ID,
      session_id=TEST_SESSION_ID,
      new_message=malicious_message,
  )

  with pytest.raises(ValueError, match="cannot contain function calls"):
    async with aclosing(agen) as a:
      async for _ in a:
        pass


if __name__ == "__main__":
  pytest.main([__file__])