feat: leverage session state to manage the status of client tool call and response

Author: ksachdeva

adk-chatkit/src/adk_chatkit/_client_tool_call.py

@@ -1,15 +1,23 @@
 from typing import Any, Literal
+from uuid import uuid4
 
-from pydantic import BaseModel
+from google.adk.tools import ToolContext
+from pydantic import BaseModel, Field
 
-from ._constants import CLIENT_TOOL_KEY_IN_TOOL_RESPONSE
+from ._constants import CHATKIT_THREAD_METADTA_KEY, CLIENT_TOOL_KEY_IN_TOOL_RESPONSE
+from ._thread_utils import (
+    add_client_tool_status,
+    serialize_thread_metadata,
+)
 
 
 class ClientToolCallState(BaseModel):
     """
     Returned from tool methods to indicate a client-side tool call.
     """
 
+    id: str = Field(default_factory=lambda: uuid4().hex)
+
     name: str
     arguments: dict[str, Any]
     status: Literal["pending", "completed"] = "pending"
@@ -18,11 +26,22 @@ class ClientToolCallState(BaseModel):
 def add_client_tool_call_to_tool_response(
     response: dict[str, Any],
     client_tool_call: ClientToolCallState,
+    tool_context: ToolContext,
 ) -> None:
     """Add a client tool call to a tool response dictionary.
 
     Args:
         response: The tool response dictionary to modify.
         client_tool_call: The client tool call state to add.
     """
+
+    thread_metadata = add_client_tool_status(
+        tool_context.state,
+        client_tool_call.id,
+        client_tool_call.status,
+    )
+
+    # update the state
+    tool_context.state[CHATKIT_THREAD_METADTA_KEY] = serialize_thread_metadata(thread_metadata)
+
     response[CLIENT_TOOL_KEY_IN_TOOL_RESPONSE] = client_tool_call

adk-chatkit/src/adk_chatkit/_response.py

@@ -40,7 +40,7 @@ async def stream_agent_response(
                     id=response_id,
                     content=[],
                     thread_id=thread.id,
-                    created_at=datetime.now(),
+                    created_at=datetime.fromtimestamp(event.timestamp),
                 )
             )
 
@@ -76,13 +76,13 @@ async def stream_agent_response(
                     if adk_client_tool:
                         yield ThreadItemDoneEvent(
                             item=ClientToolCallItem(
-                                id=str(uuid.uuid4()),
+                                id=event.id,
                                 thread_id=thread.id,
                                 name=adk_client_tool.name,
                                 arguments=adk_client_tool.arguments,
                                 status=adk_client_tool.status,
                                 created_at=datetime.fromtimestamp(event.timestamp),
-                                call_id=fn_response.id,  # type: ignore
+                                call_id=adk_client_tool.id,
                             ),
                         )
 
@@ -113,6 +113,6 @@ async def stream_agent_response(
                         id=response_id,
                         content=[AssistantMessageContent(text=text_from_final_update)],
                         thread_id=thread.id,
-                        created_at=datetime.now(),
+                        created_at=datetime.fromtimestamp(event.timestamp),
                     )
                 )

adk-chatkit/src/adk_chatkit/_store.py

@@ -24,7 +24,12 @@
 from ._client_tool_call import ClientToolCallState
 from ._constants import CHATKIT_THREAD_METADTA_KEY, CLIENT_TOOL_KEY_IN_TOOL_RESPONSE, WIDGET_KEY_IN_TOOL_RESPONSE
 from ._context import ADKContext
-from ._thread_utils import serialize_thread_metadata
+from ._thread_utils import (
+    add_client_tool_status,
+    get_client_tool_status,
+    get_thread_metadata_from_state,
+    serialize_thread_metadata,
+)
 
 
 def _to_user_message_content(event: Event) -> list[UserMessageContent]:
@@ -67,7 +72,7 @@ async def load_thread(self, thread_id: str, context: ADKContext) -> ThreadMetada
                 f"Session with id {thread_id} not found for user {context['user_id']} in app {context['app_name']}"
             )
 
-        return ThreadMetadata.model_validate(session.state[CHATKIT_THREAD_METADTA_KEY])
+        return get_thread_metadata_from_state(session.state)
 
     async def save_thread(self, thread: ThreadMetadata, context: ADKContext) -> None:
         session = await self._session_service.get_session(
@@ -159,15 +164,20 @@ async def load_thread_items(
                             adk_client_tool = fn_response.response.get(CLIENT_TOOL_KEY_IN_TOOL_RESPONSE, None)
                             if adk_client_tool:
                                 adk_client_tool = ClientToolCallState.model_validate(adk_client_tool)
-                                an_item = ClientToolCallItem(
-                                    id=event.id,
-                                    thread_id=thread_id,
-                                    name=adk_client_tool.name,
-                                    arguments=adk_client_tool.arguments,
-                                    status=adk_client_tool.status,
-                                    created_at=datetime.fromtimestamp(event.timestamp),
-                                    call_id=fn_response.id,
+                                status = get_client_tool_status(
+                                    session.state,
+                                    adk_client_tool.id,
                                 )
+                                if status:
+                                    an_item = ClientToolCallItem(
+                                        id=event.id,
+                                        thread_id=thread_id,
+                                        name=adk_client_tool.name,
+                                        arguments=adk_client_tool.arguments,
+                                        status=status,  # type: ignore
+                                        created_at=datetime.fromtimestamp(event.timestamp),
+                                        call_id=adk_client_tool.id,
+                                    )
 
             if an_item:
                 thread_items.append(an_item)
@@ -188,8 +198,7 @@ async def delete_attachment(self, attachment_id: str, context: ADKContext) -> No
         raise NotImplementedError()
 
     async def delete_thread_item(self, thread_id: str, item_id: str, context: ADKContext) -> None:
-        # deletion is called primarily to remove the ClientToolCallItem calls
-        # we simply ignore them here as they are not stored separately
+        # simply ignoring it for now (ClientToolCallItem is typically not deleted because of this)
         pass
 
     async def delete_thread(self, thread_id: str, context: ADKContext) -> None:
@@ -198,7 +207,32 @@ async def delete_thread(self, thread_id: str, context: ADKContext) -> None:
     async def save_item(self, thread_id: str, item: ThreadItem, context: ADKContext) -> None:
         # we will only handle specify types of items here
         # as quite many are automatically handled by runner
-        pass
+        if isinstance(item, ClientToolCallItem):
+            session = await self._session_service.get_session(
+                app_name=context["app_name"],
+                user_id=context["user_id"],
+                session_id=thread_id,
+            )
+
+            if not session:
+                raise ValueError(
+                    f"Session with id {thread_id} not found for user {context['user_id']} in app {context['app_name']}"
+                )
+
+            thread_metadata = add_client_tool_status(session.state, item.call_id, item.status)
+
+            state_delta = {
+                CHATKIT_THREAD_METADTA_KEY: serialize_thread_metadata(thread_metadata),
+            }
+
+            actions_with_update = EventActions(state_delta=state_delta)
+            system_event = Event(
+                invocation_id=uuid4().hex,
+                author="system",
+                actions=actions_with_update,
+                timestamp=datetime.now().timestamp(),
+            )
+            await self._session_service.append_event(session, system_event)
 
     async def load_item(self, thread_id: str, item_id: str, context: ADKContext) -> ThreadItem:
         raise NotImplementedError()
@@ -218,7 +252,7 @@ async def load_threads(
         items: list[ThreadMetadata] = []
 
         for session in sessions_response.sessions:
-            thread_metatdata_item = ThreadMetadata.model_validate(session.state[CHATKIT_THREAD_METADTA_KEY])
-            items.append(thread_metatdata_item)
+            thread_metadata = get_thread_metadata_from_state(session.state)
+            items.append(thread_metadata)
 
         return Page(data=items)

adk-chatkit/src/adk_chatkit/_thread_utils.py

@@ -2,8 +2,36 @@
 from typing import Any
 
 from chatkit.types import ThreadMetadata
+from google.adk.sessions.state import State
+
+from ._constants import CHATKIT_THREAD_METADTA_KEY
+
+_CLIENT_TOOL_PREFIX = "client-tool"
 
 
 def serialize_thread_metadata(thread: ThreadMetadata) -> dict[str, Any]:
     json_dump = thread.model_dump_json(exclude_none=True, exclude={"items"})
     return json.loads(json_dump)  # type: ignore
+
+
+def get_thread_metadata_from_state(state: State | dict[str, Any]) -> ThreadMetadata:
+    thread_metadata_dict = state[CHATKIT_THREAD_METADTA_KEY]
+    return ThreadMetadata.model_validate(thread_metadata_dict)
+
+
+def add_client_tool_status(
+    state: State | dict[str, Any],
+    client_tool_id: str,
+    status: str,
+) -> ThreadMetadata:
+    thread_metadata = get_thread_metadata_from_state(state)
+    thread_metadata.metadata[f"{_CLIENT_TOOL_PREFIX}-{client_tool_id}"] = status
+    return thread_metadata
+
+
+def get_client_tool_status(
+    state: State | dict[str, Any],
+    client_tool_id: str,
+) -> str | None:
+    thread_metadata = get_thread_metadata_from_state(state)
+    return thread_metadata.metadata.get(f"{_CLIENT_TOOL_PREFIX}-{client_tool_id}", None)

examples/backend/src/backend/agents/facts/_tools.py

@@ -50,12 +50,12 @@ async def save_fact(
 
     result = {"fact_id": confirmed.id, "status": "saved"}
 
-    add_client_tool_call_to_tool_response(result, client_tool_call)
+    add_client_tool_call_to_tool_response(result, client_tool_call, tool_context)
 
     return result
 
 
-async def switch_theme(theme: str) -> dict[str, str]:
+async def switch_theme(theme: str, tool_context: ToolContext) -> dict[str, str]:
     """Switch the chat interface between light and dark color schemes.
 
     Args:
@@ -73,7 +73,7 @@ async def switch_theme(theme: str) -> dict[str, str]:
 
     result = {"theme": requested}
 
-    add_client_tool_call_to_tool_response(result, client_tool_call)
+    add_client_tool_call_to_tool_response(result, client_tool_call, tool_context)
 
     return result