chore: Add session patch endpoint to api server for state update

This is allow user to update session state without running the agent. e.g. if I want to test some case when session has certain state on adk web.

PiperOrigin-RevId: 814252851
This commit is contained in:
Xiang (Sean) Zhou
2025-10-02 13:42:49 -07:00
committed by Copybara-Service
parent 822efe0065
commit c46308b7cf
4 changed files with 214 additions and 0 deletions
+57
View File
@@ -220,6 +220,13 @@ class UpdateMemoryRequest(common.BaseModel):
"""The ID of the session to add to memory."""
class UpdateSessionRequest(common.BaseModel):
"""Request to update session state without running the agent."""
state_delta: dict[str, Any]
"""The state changes to apply to the session."""
class RunEvalResult(common.BaseModel):
eval_set_file: str
eval_set_id: str
@@ -767,6 +774,56 @@ class AdkWebServer:
app_name=app_name, user_id=user_id, session_id=session_id
)
@app.patch(
"/apps/{app_name}/users/{user_id}/sessions/{session_id}",
response_model_exclude_none=True,
)
async def update_session(
app_name: str,
user_id: str,
session_id: str,
req: UpdateSessionRequest,
) -> Session:
"""Updates session state without running the agent.
Args:
app_name: The name of the application.
user_id: The ID of the user.
session_id: The ID of the session to update.
req: The patch request containing state changes.
Returns:
The updated session.
Raises:
HTTPException: If the session is not found.
"""
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")
# Create an event to record the state change
import uuid
from ..events.event import Event
from ..events.event import EventActions
state_update_event = Event(
invocation_id="p-" + str(uuid.uuid4()),
author="user",
actions=EventActions(state_delta=req.state_delta),
)
# Append the event to the session
# This will automatically update the session state through __update_session_state
await self.session_service.append_event(
session=session, event=state_update_event
)
return session
@app.post(
"/apps/{app_name}/eval-sets",
response_model_exclude_none=True,
@@ -176,6 +176,36 @@ class AdkWebServerClient:
)
response.raise_for_status()
async def update_session(
self,
*,
app_name: str,
user_id: str,
session_id: str,
state_delta: Dict[str, Any],
) -> Session:
"""Update session state without running the agent.
Args:
app_name: Name of the application
user_id: User identifier
session_id: Session identifier to update
state_delta: The state changes to apply to the session
Returns:
The updated Session object
Raises:
httpx.HTTPStatusError: If the request fails or session not found
"""
async with self._get_client() as client:
response = await client.patch(
f"/apps/{app_name}/users/{user_id}/sessions/{session_id}",
json={"state_delta": state_delta},
)
response.raise_for_status()
return Session.model_validate(response.json())
async def run_agent(
self,
request: RunAgentRequest,