feat(runners): stop a live agent and edit a previous prompt

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,

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,13 +2019,102 @@ 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(
           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,