diff --git a/src/google/adk/cli/api_server.py b/src/google/adk/cli/api_server.py index c68624bba0..4e028a3e81 100644 --- a/src/google/adk/cli/api_server.py +++ b/src/google/adk/cli/api_server.py @@ -622,6 +622,13 @@ class UpdateSessionRequest(common.BaseModel): """The state changes to apply to the session.""" +class RewindSessionRequest(common.BaseModel): + """Request to rewind a session to before a given invocation.""" + + rewind_before_invocation_id: str + """The invocation ID to rewind the session to before.""" + + class FinalizeAgentIdentityCredentialsRequest(common.BaseModel): """Request to finalize a 3LO consent for an Agent Identity connector.""" @@ -1590,6 +1597,47 @@ async def update_session( return session + @app.post( + "/apps/{app_name}/users/{user_id}/sessions/{session_id}/rewind", + response_model_exclude_none=True, + ) + async def rewind_session( + app_name: str, + user_id: str, + session_id: str, + req: RewindSessionRequest, + ) -> Session: + """Rewinds a session to before the specified invocation. + + Args: + app_name: The name of the application. + user_id: The ID of the user. + session_id: The ID of the session to rewind. + req: The rewind request identifying the invocation to rewind to. + + Returns: + The rewound session. + + Raises: + HTTPException: If the session or invocation is not found. + """ + runner = await self.get_runner_async(app_name) + try: + await runner.rewind_async( + user_id=user_id, + session_id=session_id, + rewind_before_invocation_id=req.rewind_before_invocation_id, + ) + except ValueError as ve: + raise HTTPException(status_code=404, detail=str(ve)) from ve + + session = await self.session_service.get_session( + app_name=app_name, user_id=user_id, session_id=session_id + ) + if not session: + raise HTTPException(status_code=404, detail="Session not found") + return session + @app.get( "/apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{artifact_name:path}/versions/{version_id}/metadata", response_model=ArtifactVersion, diff --git a/tests/unittests/cli/test_fast_api.py b/tests/unittests/cli/test_fast_api.py index a3fe28d35a..0c2f50ed71 100644 --- a/tests/unittests/cli/test_fast_api.py +++ b/tests/unittests/cli/test_fast_api.py @@ -1737,6 +1737,60 @@ def test_patch_session_not_found(test_app, test_session_info): logger.info("Patch session not found test passed") +def test_rewind_session(test_app, create_test_session, monkeypatch): + """Test rewinding a session to before a given invocation.""" + info = create_test_session + captured: dict[str, Any] = {} + + async def rewind_async_capture( + self, + *, + user_id: str, + session_id: str, + rewind_before_invocation_id: str, + run_config: Optional[RunConfig] = None, + ): + del self, run_config + captured["user_id"] = user_id + captured["session_id"] = session_id + captured["rewind_before_invocation_id"] = rewind_before_invocation_id + + monkeypatch.setattr(Runner, "rewind_async", rewind_async_capture) + + url = f"/apps/{info['app_name']}/users/{info['user_id']}/sessions/{info['session_id']}/rewind" + response = test_app.post( + url, json={"rewind_before_invocation_id": "some-invocation-id"} + ) + + assert response.status_code == 200 + assert captured["user_id"] == info["user_id"] + assert captured["session_id"] == info["session_id"] + assert captured["rewind_before_invocation_id"] == "some-invocation-id" + data = response.json() + assert data["id"] == info["session_id"] + + +def test_rewind_session_invocation_not_found( + test_app, create_test_session, monkeypatch +): + """Test rewinding to an unknown invocation ID returns 404.""" + info = create_test_session + + async def rewind_async_raise(self, **kwargs): + del self, kwargs + raise ValueError("Invocation ID not found: missing-invocation-id") + + monkeypatch.setattr(Runner, "rewind_async", rewind_async_raise) + + url = f"/apps/{info['app_name']}/users/{info['user_id']}/sessions/{info['session_id']}/rewind" + response = test_app.post( + url, json={"rewind_before_invocation_id": "missing-invocation-id"} + ) + + assert response.status_code == 404 + assert "missing-invocation-id" in response.json()["detail"] + + def test_agent_run(test_app, create_test_session): """Test running an agent with a message.""" info = create_test_session