diff --git a/src/google/adk/agents/invocation_context.py b/src/google/adk/agents/invocation_context.py index a4bc8955bd..754c1e4c61 100644 --- a/src/google/adk/agents/invocation_context.py +++ b/src/google/adk/agents/invocation_context.py @@ -206,6 +206,9 @@ class InvocationContext(BaseModel): Set to True in callbacks or tools to terminate this invocation.""" + _cancel_event: asyncio.Event = PrivateAttr(default_factory=asyncio.Event) + """Set when a caller requests this invocation to stop (Runner.cancel_async).""" + live_request_queue: LiveRequestQueue | None = None """The queue to receive live requests.""" @@ -314,6 +317,15 @@ async def _enqueue_event(self, event: Event) -> None: await self._event_queue.put((event, processed)) await processed.wait() + def request_cancel(self) -> None: + """Asks this invocation to stop as soon as possible. + + Sets ``end_invocation`` so LLM loops exit at the next step, and trips + ``_cancel_event`` so the runner can abort a blocked tool or model call. + """ + self.end_invocation = True + self._cancel_event.set() + def set_agent_state( self, agent_name: str, diff --git a/src/google/adk/cli/api_server.py b/src/google/adk/cli/api_server.py index 8285af05a9..2661e79737 100644 --- a/src/google/adk/cli/api_server.py +++ b/src/google/adk/cli/api_server.py @@ -532,6 +532,26 @@ class RunAgentRequest(common.BaseModel): custom_metadata: Optional[dict[str, Any]] = None +class CancelAgentRequest(common.BaseModel): + """Stops a live invocation on a session. + + If ``invocation_id`` is omitted, every live invocation on the session is + cancelled. + """ + + invocation_id: Optional[str] = None + + +class EditMessageRequest(common.BaseModel): + """Replaces a previous user prompt and regenerates from that turn.""" + + invocation_id: str + new_message: types.Content + streaming: bool = False + state_delta: Optional[dict[str, Any]] = None + custom_metadata: Optional[dict[str, Any]] = None + + class CreateSessionRequest(common.BaseModel): session_id: Optional[str] = Field( default=None, @@ -1862,7 +1882,9 @@ async def monitor(): monitor_task.cancel() @app.post("/run_sse") - async def run_agent_sse(req: RunAgentRequest) -> StreamingResponse: + async def run_agent_sse( + req: RunAgentRequest, request: Request + ) -> StreamingResponse: app_name = req.app_name or self.default_app_name if not app_name: raise HTTPException( @@ -1897,6 +1919,28 @@ async def run_agent_sse(req: RunAgentRequest) -> StreamingResponse: async def event_generator(): is_closing = False original_exc = None + + async def watch_disconnect(): + try: + while True: + message = await request.receive() + if message.get("type") == "http.disconnect": + logger.warning( + "Client disconnected. Cancelling agent run for session %s.", + req.session_id, + ) + await runner.cancel_async( + user_id=req.user_id, session_id=req.session_id + ) + break + except asyncio.CancelledError: + pass + except Exception as e: # pylint: disable=broad-exception-caught + logger.error( + "Exception in disconnect monitor: %s", e, exc_info=True + ) + + monitor_task = asyncio.create_task(watch_disconnect()) try: async with Aclosing( runner.run_async( @@ -1975,6 +2019,8 @@ async def event_generator(): "Error during generator cleanup after completion: %s", e ) raise e + finally: + monitor_task.cancel() # Returns a streaming response with the proper media type for SSE return StreamingResponse( @@ -1982,6 +2028,93 @@ async def event_generator(): media_type="text/event-stream", ) + @app.post( + "/apps/{app_name}/users/{user_id}/sessions/{session_id}/cancel", + response_model_exclude_none=True, + ) + async def cancel_agent_run( + app_name: str, + user_id: str, + session_id: str, + req: CancelAgentRequest = CancelAgentRequest(), + ) -> dict[str, Any]: + """Stops a live agent invocation on this session.""" + self.current_app_name_ref.value = app_name + runner = await self.get_runner_async(app_name) + _set_telemetry_context_if_needed(runner) + invocation_id = req.invocation_id + try: + cancelled_ids = await runner.cancel_async( + user_id=user_id, + session_id=session_id, + invocation_id=invocation_id, + ) + except SessionNotFoundError as e: + raise HTTPException(status_code=404, detail=str(e)) from e + return { + "cancelled": bool(cancelled_ids), + "invocationIds": cancelled_ids, + } + + @app.post( + "/apps/{app_name}/users/{user_id}/sessions/{session_id}/edit", + response_model_exclude_none=True, + ) + async def edit_agent_message( + app_name: str, + user_id: str, + session_id: str, + req: EditMessageRequest, + ): + """Rewinds to a previous user turn and regenerates with a new prompt.""" + self.current_app_name_ref.value = app_name + runner = await self.get_runner_async(app_name) + _set_telemetry_context_if_needed(runner) + run_config = ( + RunConfig( + streaming_mode=( + StreamingMode.SSE if req.streaming else StreamingMode.NONE + ), + custom_metadata=req.custom_metadata, + ) + if req.custom_metadata or req.streaming + else None + ) + + async def _edit_events(): + try: + async with Aclosing( + runner.edit_message_async( + user_id=user_id, + session_id=session_id, + invocation_id=req.invocation_id, + new_message=req.new_message, + state_delta=req.state_delta, + run_config=run_config, + ) + ) as agen: + async for event in agen: + yield event + except SessionNotFoundError as e: + raise HTTPException(status_code=404, detail=str(e)) from e + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) from e + + if req.streaming: + + async def event_generator(): + async for event in _edit_events(): + sse_event = event.model_dump_json(exclude_none=True, by_alias=True) + yield f"data: {sse_event}\n\n" + + return StreamingResponse( + event_generator(), + media_type="text/event-stream", + ) + + events = [event async for event in _edit_events()] + return events + @app.websocket("/run_live") async def run_agent_live( websocket: WebSocket, diff --git a/src/google/adk/runners.py b/src/google/adk/runners.py index 6d9f70072e..ab39aa91eb 100644 --- a/src/google/adk/runners.py +++ b/src/google/adk/runners.py @@ -16,6 +16,8 @@ import asyncio from contextlib import aclosing +from dataclasses import dataclass +from dataclasses import field import inspect import logging from pathlib import Path @@ -78,6 +80,24 @@ _EventQueueItem = tuple[object, asyncio.Event | None] +# Put on InvocationContext._event_queue by cancel_async so the consume loop +# can yield the already-persisted interrupted event and exit cleanly. +_CANCEL_SENTINEL = object() + + +@dataclass +class _ActiveInvocation: + """A live run_async invocation that cancel_async can stop.""" + + user_id: str + session_id: str + invocation_id: str + context: InvocationContext + consumer_task: asyncio.Task[Any] | None = None + root_task: asyncio.Task[None] | None = None + interrupted_event: Event | None = field(default=None) + + # Silence unused warning. # tracer is imported for backwards compatibility, to avoid breaking change in the API. _ = tracer @@ -324,6 +344,8 @@ def __init__( self._app_name_alignment_hint: Optional[str] = None self._enforce_app_name_alignment() self._warn_uncached_agent_transfer() + # Live run_async invocations, keyed by (user_id, session_id, invocation_id). + self._active_invocations: dict[tuple[str, str, str], _ActiveInvocation] = {} def _require_root_agent(self) -> BaseAgent: """Returns the root as an agent for agent-only execution paths.""" @@ -632,6 +654,12 @@ async def _run() -> AsyncGenerator[Event, None]: invocation_id=invocation_id, ) ic._event_queue = asyncio.Queue() + rec = self._register_active_invocation( + user_id=user_id, + session_id=session_id, + invocation_context=ic, + consumer_task=asyncio.current_task(), + ) # 2. Append user message to session and resolve node_input node_input = None @@ -743,6 +771,7 @@ async def _drive_root_node() -> None: await ic._event_queue.put((done_sentinel, None)) task = asyncio.create_task(_drive_root_node()) + rec.root_task = task # 4. Main loop: consume events, persist, yield try: @@ -770,18 +799,23 @@ async def _drive_root_node() -> None: # unhandled runner error, so notify on_run_error_callback once and # re-raise. on_run_error is notification-only and never raises, so # there is no recursive notification. - if run_error is None: - try: - await ic.plugin_manager.run_after_run_callback( - invocation_context=ic - ) - await self._run_post_invocation_compaction( - session=session, - skip_token_compaction=ic.token_compaction_checked, - ) - except Exception as e: - await _notify_run_error(ic.plugin_manager, ic, e) - raise + try: + if run_error is None: + try: + await ic.plugin_manager.run_after_run_callback( + invocation_context=ic + ) + await self._run_post_invocation_compaction( + session=session, + skip_token_compaction=ic.token_compaction_checked, + ) + except Exception as e: + await _notify_run_error(ic.plugin_manager, ic, e) + raise + finally: + self._unregister_active_invocation( + user_id, session_id, ic.invocation_id + ) async with aclosing(_with_caller_context(_run(), caller_ctx)) as agen: async for event in agen: @@ -980,6 +1014,12 @@ async def _consume_event_queue( event_or_done, processed_signal = await event_queue.get() if event_or_done is done_sentinel: break + if event_or_done is _CANCEL_SENTINEL: + key = (ic.session.user_id, ic.session.id, ic.invocation_id) + rec = self._active_invocations.get(key) + if rec is not None and rec.interrupted_event is not None: + yield rec.interrupted_event + break if not isinstance(event_or_done, Event): raise TypeError( f'Unexpected node event queue item: {type(event_or_done).__name__}' @@ -1432,36 +1472,50 @@ async def _run_with_trace( # already final. return - async def execute( - ctx: InvocationContext, - ) -> AsyncGenerator[Event, None]: - active_agent = ctx.agent - if not isinstance(active_agent, BaseAgent): - raise RuntimeError('Agent execution has no active BaseAgent.') - async with aclosing(active_agent.run_async(ctx)) as agen: + self._register_active_invocation( + user_id=user_id, + session_id=session_id, + invocation_context=invocation_context, + consumer_task=asyncio.current_task(), + ) + try: + + async def execute( + ctx: InvocationContext, + ) -> AsyncGenerator[Event, None]: + active_agent = ctx.agent + if not isinstance(active_agent, BaseAgent): + raise RuntimeError('Agent execution has no active BaseAgent.') + async with aclosing(active_agent.run_async(ctx)) as agen: + async for event in agen: + yield event + + async with aclosing( + _with_caller_context( + self._exec_with_plugin( + invocation_context=invocation_context, + session=invocation_context.session, + execute_fn=execute, + is_live_call=False, + ), + caller_ctx_trace, + ) + ) as agen: async for event in agen: yield event - - async with aclosing( - _with_caller_context( - self._exec_with_plugin( - invocation_context=invocation_context, - session=invocation_context.session, - execute_fn=execute, - is_live_call=False, - ), - caller_ctx_trace, - ) - ) as agen: - async for event in agen: - yield event - # Run compaction after all events are yielded from the agent. - # (We don't compact in the middle of an invocation, we only compact at - # the end of an invocation.) - await self._run_post_invocation_compaction( - session=invocation_context.session, - skip_token_compaction=(invocation_context.token_compaction_checked), - ) + # Run compaction after all events are yielded from the agent. + # (We don't compact in the middle of an invocation, we only compact at + # the end of an invocation.) + await self._run_post_invocation_compaction( + session=invocation_context.session, + skip_token_compaction=( + invocation_context.token_compaction_checked + ), + ) + finally: + self._unregister_active_invocation( + user_id, session_id, invocation_context.invocation_id + ) async with aclosing(_run_with_trace(new_message, invocation_id)) as agen: async for event in agen: @@ -1518,6 +1572,218 @@ async def rewind_async( await self.session_service.append_event(session=session, event=rewind_event) + def _active_invocation_key( + self, user_id: str, session_id: str, invocation_id: str + ) -> tuple[str, str, str]: + return (user_id, session_id, invocation_id) + + def _register_active_invocation( + self, + *, + user_id: str, + session_id: str, + invocation_context: InvocationContext, + consumer_task: asyncio.Task[Any] | None, + ) -> _ActiveInvocation: + rec = _ActiveInvocation( + user_id=user_id, + session_id=session_id, + invocation_id=invocation_context.invocation_id, + context=invocation_context, + consumer_task=consumer_task, + ) + self._active_invocations[ + self._active_invocation_key( + user_id, session_id, invocation_context.invocation_id + ) + ] = rec + return rec + + def _unregister_active_invocation( + self, user_id: str, session_id: str, invocation_id: str + ) -> None: + self._active_invocations.pop( + self._active_invocation_key(user_id, session_id, invocation_id), None + ) + + def _matching_active_invocations( + self, + *, + user_id: str, + session_id: str, + invocation_id: Optional[str], + ) -> list[_ActiveInvocation]: + if invocation_id is not None: + rec = self._active_invocations.get( + self._active_invocation_key(user_id, session_id, invocation_id) + ) + return [rec] if rec is not None else [] + return [ + rec + for rec in self._active_invocations.values() + if rec.user_id == user_id and rec.session_id == session_id + ] + + def _interrupted_event(self, ic: InvocationContext) -> Event: + agent = ic.agent + author = ( + agent.name if agent is not None and hasattr(agent, 'name') else 'model' + ) + event = Event( + invocation_id=ic.invocation_id, + author=author, + interrupted=True, + ) + _apply_run_config_custom_metadata(event, ic.run_config) + ic.stamp_event_branch_context(event) + return event + + async def _wait_for_invocations_to_finish( + self, keys: list[tuple[str, str, str]], timeout: float = 10.0 + ) -> None: + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + pending = list(keys) + while pending and loop.time() < deadline: + pending = [key for key in pending if key in self._active_invocations] + if not pending: + return + await asyncio.sleep(0.01) + if pending: + logger.warning( + 'Timed out waiting for cancelled invocations to finish: %s', pending + ) + + async def cancel_async( + self, + *, + user_id: str, + session_id: str, + invocation_id: Optional[str] = None, + ) -> list[str]: + """Stops a live agent invocation. + + Cooperative LLM loops see ``InvocationContext.end_invocation``. Blocked + tool/model calls are cancelled via the root node task (workflow path) or + the consumer task (legacy agents). An ``interrupted=True`` event is + appended so the session records the stop. + + Args: + user_id: The user ID of the session. + session_id: The session ID. + invocation_id: If set, only that invocation is cancelled. If omitted, + every live invocation on the session is cancelled. + + Returns: + The invocation IDs that were actually running and are now stopping. + Empty if nothing was in flight (idempotent). + """ + records = self._matching_active_invocations( + user_id=user_id, session_id=session_id, invocation_id=invocation_id + ) + if not records: + return [] + + cancelled_ids: list[str] = [] + caller_task = asyncio.current_task() + for rec in records: + rec.context.request_cancel() + interrupted = self._interrupted_event(rec.context) + rec.interrupted_event = interrupted + await self.session_service.append_event( + session=rec.context.session, event=interrupted + ) + event_queue = rec.context._event_queue + if event_queue is not None: + await event_queue.put((_CANCEL_SENTINEL, None)) + if rec.root_task is not None and not rec.root_task.done(): + rec.root_task.cancel() + elif ( + rec.consumer_task is not None + and rec.consumer_task is not caller_task + and not rec.consumer_task.done() + ): + rec.consumer_task.cancel() + cancelled_ids.append(rec.invocation_id) + logger.info( + 'Cancelled invocation %s for session %s', + rec.invocation_id, + session_id, + ) + return cancelled_ids + + async def edit_message_async( + self, + *, + user_id: str, + session_id: str, + invocation_id: str, + new_message: types.Content, + state_delta: Optional[dict[str, Any]] = None, + run_config: Optional[RunConfig] = None, + ) -> AsyncGenerator[Event, None]: + """Replaces a previous user prompt and regenerates from that turn. + + Stops any live run on the session, rewinds history to before + ``invocation_id``, then runs the agent with ``new_message``. + + Args: + user_id: The user ID of the session. + session_id: The session ID. + invocation_id: The user turn to replace (the invocation that user + message started). + new_message: The edited prompt. + state_delta: Optional state changes applied with the new message. + run_config: The run config for the regenerated turn. + + Yields: + Events from the regenerated invocation. + + Raises: + ValueError: If ``invocation_id`` is not a user turn in the session. + """ + if new_message and not new_message.role: + new_message.role = 'user' + + run_config = run_config or RunConfig() + session = await self._get_or_create_session( + user_id=user_id, + session_id=session_id, + get_session_config=run_config.get_session_config, + ) + original = self._find_original_user_content(session, invocation_id) + if original is None: + raise ValueError( + f'No user message found for invocation ID: {invocation_id}' + ) + + cancelled_ids = await self.cancel_async( + user_id=user_id, session_id=session_id + ) + if cancelled_ids: + await self._wait_for_invocations_to_finish([ + self._active_invocation_key(user_id, session_id, cancelled_id) + for cancelled_id in cancelled_ids + ]) + + await self.rewind_async( + user_id=user_id, + session_id=session_id, + rewind_before_invocation_id=invocation_id, + run_config=run_config, + ) + async with aclosing( + self.run_async( + user_id=user_id, + session_id=session_id, + new_message=new_message, + state_delta=state_delta, + run_config=run_config, + ) + ) as agen: + async for event in agen: + yield event + async def _compute_state_delta_for_rewind( self, session: Session, rewind_event_index: int ) -> dict[str, Any]: diff --git a/tests/unittests/cli/test_fast_api.py b/tests/unittests/cli/test_fast_api.py index a3fe28d35a..ac18dfa8cf 100644 --- a/tests/unittests/cli/test_fast_api.py +++ b/tests/unittests/cli/test_fast_api.py @@ -153,6 +153,34 @@ async def dummy_run_async( yield _event_state_delta(state_delta) +async def dummy_cancel_async( + self, + *, + user_id, + session_id, + invocation_id=None, +): + del self, user_id, session_id + return [invocation_id] if invocation_id else ["inv-live"] + + +async def dummy_edit_message_async( + self, + *, + user_id, + session_id, + invocation_id, + new_message, + state_delta=None, + run_config: Optional[RunConfig] = None, +): + del self, user_id, session_id, invocation_id, new_message + del state_delta, run_config + yield _event_1() + await asyncio.sleep(0) + yield _event_3() + + # Define a local mock for EvalCaseResult specific to fast_api tests class _MockEvalCaseResult(BaseModel): eval_set_id: str @@ -176,6 +204,8 @@ def patch_runner(monkeypatch): """Patch the Runner methods to use our dummy implementations.""" monkeypatch.setattr(Runner, "run_live", dummy_run_live) monkeypatch.setattr(Runner, "run_async", dummy_run_async) + monkeypatch.setattr(Runner, "cancel_async", dummy_cancel_async) + monkeypatch.setattr(Runner, "edit_message_async", dummy_edit_message_async) @pytest.fixture @@ -2093,7 +2123,13 @@ def run_async_mock(self, **kwargs): ) # Call handler - response = await handler(req) + async def _pending_receive(): + await asyncio.sleep(3600) + return {"type": "http.request", "body": b""} + + mock_request = MagicMock() + mock_request.receive = _pending_receive + response = await handler(req, mock_request) assert response.status_code == 200 # Iterate generator and close it early @@ -2166,7 +2202,13 @@ def run_async_mock(self, **kwargs): ) # Call handler - response = await handler(req) + async def _pending_receive(): + await asyncio.sleep(3600) + return {"type": "http.request", "body": b""} + + mock_request = MagicMock() + mock_request.receive = _pending_receive + response = await handler(req, mock_request) assert response.status_code == 200 # Iterate generator @@ -4785,5 +4827,104 @@ def test_create_eval_set_legacy_route_creates_eval_set( ) +def test_cancel_agent_run_endpoint(test_app, create_test_session, monkeypatch): + """POST /sessions/{id}/cancel stops the live invocation via Runner.""" + captured = {} + + async def capturing_cancel_async( + self, *, user_id, session_id, invocation_id=None + ): + captured["user_id"] = user_id + captured["session_id"] = session_id + captured["invocation_id"] = invocation_id + return ["inv-live"] + + monkeypatch.setattr(Runner, "cancel_async", capturing_cancel_async) + + info = create_test_session + url = ( + f"/apps/{info['app_name']}/users/{info['user_id']}" + f"/sessions/{info['session_id']}/cancel" + ) + response = test_app.post(url, json={"invocationId": "inv-live"}) + + assert response.status_code == 200 + data = response.json() + assert data["cancelled"] is True + assert data["invocationIds"] == ["inv-live"] + assert captured["user_id"] == info["user_id"] + assert captured["session_id"] == info["session_id"] + assert captured["invocation_id"] == "inv-live" + + +def test_cancel_agent_run_endpoint_without_invocation_id( + test_app, create_test_session +): + """Omitting invocation_id cancels every live run on the session.""" + info = create_test_session + url = ( + f"/apps/{info['app_name']}/users/{info['user_id']}" + f"/sessions/{info['session_id']}/cancel" + ) + response = test_app.post(url, json={}) + + assert response.status_code == 200 + data = response.json() + assert data["cancelled"] is True + assert data["invocationIds"] == ["inv-live"] + + +def test_edit_agent_message_endpoint( + test_app, create_test_session, monkeypatch +): + """POST /sessions/{id}/edit regenerates from the targeted user turn.""" + captured = {} + + async def capturing_edit_message_async( + self, + *, + user_id, + session_id, + invocation_id, + new_message, + state_delta=None, + run_config=None, + ): + captured["user_id"] = user_id + captured["session_id"] = session_id + captured["invocation_id"] = invocation_id + captured["new_message"] = new_message + yield _event_1() + yield _event_3() + + monkeypatch.setattr( + Runner, "edit_message_async", capturing_edit_message_async + ) + + info = create_test_session + url = ( + f"/apps/{info['app_name']}/users/{info['user_id']}" + f"/sessions/{info['session_id']}/edit" + ) + response = test_app.post( + url, + json={ + "invocationId": "inv-to-edit", + "newMessage": { + "role": "user", + "parts": [{"text": "edited prompt"}], + }, + }, + ) + + assert response.status_code == 200 + data = response.json() + assert isinstance(data, list) + assert len(data) == 2 + assert data[0]["author"] == "dummy agent" + assert captured["invocation_id"] == "inv-to-edit" + assert captured["new_message"].parts[0].text == "edited prompt" + + if __name__ == "__main__": pytest.main(["-xvs", __file__]) diff --git a/tests/unittests/test_runners_stop_and_edit.py b/tests/unittests/test_runners_stop_and_edit.py new file mode 100644 index 0000000000..91c0a29612 --- /dev/null +++ b/tests/unittests/test_runners_stop_and_edit.py @@ -0,0 +1,372 @@ +# 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. + +"""Stop a running agent and edit a previous prompt (#3849).""" + +from __future__ import annotations + +import asyncio +from typing import AsyncGenerator + +from google.adk.agents.base_agent import BaseAgent +from google.adk.agents.invocation_context import InvocationContext +from google.adk.agents.llm_agent import LlmAgent +from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService +from google.adk.events.event import Event +from google.adk.runners import Runner +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.genai import types +import pytest + +TEST_APP_ID = "test_app" +TEST_USER_ID = "test_user" +TEST_SESSION_ID = "test_session" + + +def _user_message(text: str) -> types.Content: + return types.Content(role="user", parts=[types.Part(text=text)]) + + +def _user_text(invocation_context: InvocationContext) -> str: + if invocation_context.user_content and invocation_context.user_content.parts: + return invocation_context.user_content.parts[0].text or "" + return "" + + +def _make_slow_agent(kind: str) -> tuple[BaseAgent, asyncio.Event]: + """kind is 'base' (legacy path) or 'llm' (node path).""" + started = asyncio.Event() + + class SlowBaseAgent(BaseAgent): + + 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="thinking")] + ), + ) + started.set() + try: + await asyncio.sleep(30) + yield Event( + invocation_id=invocation_context.invocation_id, + author=self.name, + content=types.Content( + role="model", parts=[types.Part(text="should not appear")] + ), + ) + except (asyncio.CancelledError, GeneratorExit): + raise + + class SlowLlmAgent(LlmAgent): + + 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="thinking")] + ), + ) + started.set() + try: + await asyncio.sleep(30) + yield Event( + invocation_id=invocation_context.invocation_id, + author=self.name, + content=types.Content( + role="model", parts=[types.Part(text="should not appear")] + ), + ) + except (asyncio.CancelledError, GeneratorExit): + raise + + if kind == "llm": + return SlowLlmAgent(name="slow_agent", model="gemini-1.5-pro"), started + return SlowBaseAgent(name="slow_agent"), started + + +def _make_slow_then_echo_agent() -> tuple[BaseAgent, asyncio.Event]: + started = asyncio.Event() + calls = {"n": 0} + + class SlowThenEchoAgent(BaseAgent): + + async def _run_async_impl( + self, invocation_context: InvocationContext + ) -> AsyncGenerator[Event, None]: + calls["n"] += 1 + text = _user_text(invocation_context) + if calls["n"] == 1: + yield Event( + invocation_id=invocation_context.invocation_id, + author=self.name, + content=types.Content( + role="model", parts=[types.Part(text="thinking")] + ), + ) + started.set() + await asyncio.sleep(30) + return + yield Event( + invocation_id=invocation_context.invocation_id, + author=self.name, + content=types.Content( + role="model", parts=[types.Part(text=f"echo:{text}")] + ), + ) + + return SlowThenEchoAgent(name="slow_echo"), started + + +class EchoAgent(BaseAgent): + """Replies with the latest user text so edits are observable.""" + + def __init__(self, name: str): + super().__init__(name=name, sub_agents=[]) + + async def _run_async_impl( + self, invocation_context: InvocationContext + ) -> AsyncGenerator[Event, None]: + text = "" + if ( + invocation_context.user_content + and invocation_context.user_content.parts + ): + text = invocation_context.user_content.parts[0].text or "" + yield Event( + invocation_id=invocation_context.invocation_id, + author=self.name, + content=types.Content( + role="model", parts=[types.Part(text=f"echo:{text}")] + ), + ) + + +def _runner(agent: BaseAgent) -> tuple[Runner, InMemorySessionService]: + session_service = InMemorySessionService() + runner = Runner( + app_name=TEST_APP_ID, + agent=agent, + session_service=session_service, + artifact_service=InMemoryArtifactService(), + auto_create_session=True, + ) + return runner, session_service + + +async def _drain(agen) -> list[Event]: + events: list[Event] = [] + try: + async for event in agen: + events.append(event) + except asyncio.CancelledError: + pass + return events + + +@pytest.mark.parametrize("kind", ["base", "llm"]) +async def test_cancel_async_stops_a_running_agent(kind): + """cancel_async aborts a blocked agent and records interrupted=True.""" + agent, started = _make_slow_agent(kind) + runner, session_service = _runner(agent) + first_invocation = asyncio.get_running_loop().create_future() + + async def _run() -> list[Event]: + events: list[Event] = [] + async for event in runner.run_async( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + new_message=_user_message("go"), + ): + events.append(event) + if not first_invocation.done() and event.invocation_id: + first_invocation.set_result(event.invocation_id) + return events + + run_task = asyncio.create_task(_run()) + await asyncio.wait_for(started.wait(), timeout=5) + cancelled_ids = await runner.cancel_async( + user_id=TEST_USER_ID, session_id=TEST_SESSION_ID + ) + try: + run_events = await asyncio.wait_for(run_task, timeout=5) + except asyncio.CancelledError: + run_events = [] + + assert cancelled_ids + assert first_invocation.result() in cancelled_ids + 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 any(event.interrupted for event in session.events) + assert all( + event.content is None + or not event.content.parts + or event.content.parts[0].text != "should not appear" + for event in run_events + session.events + ) + + +async def test_cancel_async_is_idempotent_when_nothing_is_running(): + """Stopping a session with no live run returns an empty list.""" + runner, _ = _runner(EchoAgent("echo")) + async for _ in runner.run_async( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + new_message=_user_message("done"), + ): + pass + + cancelled_ids = await runner.cancel_async( + user_id=TEST_USER_ID, session_id=TEST_SESSION_ID + ) + assert cancelled_ids == [] + + +async def test_edit_message_async_regenerates_from_the_edited_prompt(): + """Editing a prior user turn rewinds history and reruns with the new text.""" + runner, session_service = _runner(EchoAgent("echo")) + + async for _ in runner.run_async( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + new_message=_user_message("first"), + ): + pass + 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 + first_invocation_id = session.events[0].invocation_id + + async for _ in runner.run_async( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + new_message=_user_message("second"), + ): + pass + + edited = [] + async for event in runner.edit_message_async( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + invocation_id=first_invocation_id, + new_message=_user_message("first edited"), + ): + edited.append(event) + + assert any( + event.content + and event.content.parts + and event.content.parts[0].text == "echo:first edited" + for event in edited + ) + + 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 any( + event.actions.rewind_before_invocation_id == first_invocation_id + for event in session.events + ) + user_texts = [ + event.content.parts[0].text + for event in session.events + if event.author == "user" + and event.content + and event.content.parts + and event.content.parts[0].text + ] + assert "first edited" in user_texts + + +async def test_edit_message_async_cancels_a_live_turn_then_regenerates(): + """Editing the in-flight prompt stops it, then reruns with the new text.""" + agent, started = _make_slow_then_echo_agent() + runner, session_service = _runner(agent) + first_invocation = asyncio.get_running_loop().create_future() + + async def _run() -> None: + async for event in runner.run_async( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + new_message=_user_message("original"), + ): + if not first_invocation.done() and event.invocation_id: + first_invocation.set_result(event.invocation_id) + + run_task = asyncio.create_task(_run()) + await asyncio.wait_for(started.wait(), timeout=5) + live_invocation_id = await first_invocation + + edited_events = await asyncio.wait_for( + _drain( + runner.edit_message_async( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + invocation_id=live_invocation_id, + new_message=_user_message("edited live"), + ) + ), + timeout=5, + ) + try: + await asyncio.wait_for(run_task, timeout=5) + except asyncio.CancelledError: + pass + + 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 any(event.interrupted for event in session.events) + assert any( + event.actions.rewind_before_invocation_id == live_invocation_id + for event in session.events + ) + assert any( + event.content + and event.content.parts + and event.content.parts[0].text == "echo:edited live" + for event in edited_events + ) + + +async def test_edit_message_async_rejects_unknown_invocation(): + """Editing a turn that does not exist raises ValueError.""" + runner, _ = _runner(EchoAgent("echo")) + async for _ in runner.run_async( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + new_message=_user_message("hello"), + ): + pass + + with pytest.raises(ValueError, match="No user message found"): + agen = runner.edit_message_async( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + invocation_id="missing-inv", + new_message=_user_message("nope"), + ) + await agen.__anext__()