Self-hosting ADK agent over websocket server with multiple clients

[jacetest]

Currently, I am deploying a single Websocket server that can handle multiple connections, and creates a separate agent runner for each connection. However, some of my tool calls involve communicating with the client's local data, which means it sends commands over the websocket connection, but needs to know which one to use since multiple clients may be connected at the same time. I need to access the userID of the session - is this provided via Tool Context? Is my architecture the correct method to go about this?

[adk-bot]

Response from ADK Answering Agent (experimental, answer may be inaccurate)

TLDR: Your architecture is on the right track. You should manage userIDs when you create sessions for each client. Within your tools, you can then access session-specific data using the ToolContext object, which holds the session's state. For a production environment, it is recommended to use a distributed session store instead of the default in-memory one.

---

Hello! Based on the ADK documentation, here is some information that might help you with your questions.

Accessing Client-Specific Data

You are correct to think about the context of a tool call. In ADK, this is handled by the ToolContext. While the userID is not passed directly in the ToolContext, the context itself is tied to a session that was created using that userID.

Here's the recommended approach:

1. Create a Session with a userID: When a new client connects via WebSocket, your application should be responsible for creating a session for them, identified by a unique userID.

2. Use ToolContext to Access Session State: Your tool functions can accept a tool_context: ToolContext parameter. [1] This object has a state attribute (tool_context.state), which is a dictionary that persists across all turns in a single user session. [2, 3] You can use this state dictionary to store and retrieve any client-specific information you need, such as which WebSocket connection to use.

   For example, you could store the WebSocket connection object or a client identifier in the session state when the connection is established.

   from adk.context import ToolContext
   from adk.tools import tool
   
   @tool
   def my_tool(query: str, tool_context: ToolContext):
     # Access session-specific state
     client_identifier = tool_context.state.get("client_id")
     print(f"Running tool for client: {client_identifier}")
     # ... your tool logic

Architecture for Multiple Clients

Your architecture of a single WebSocket server creating an agent runner for each connection is a valid approach and is similar to patterns shown in the documentation. [4]

For scaling to a production environment, the documentation provides a few key recommendations:

• Stateless Application Instances: To handle a large number of users, you can run multiple instances of your agent application behind a load balancer.
• Distributed Session Store: When running multiple instances, the default InMemorySessionService will not work, as it stores session data only on the instance that created the session. You should switch to a distributed session store (like one using Redis or a database) that all your application instances can connect to. [5] This ensures that any instance can handle a request for any session.

By managing sessions with unique userIDs and leveraging the tool_context.state to store connection-specific information, you can build a robust multi-client agent.

---

[1] https://google.github.io/adk-docs/tools/
[2] https://google.github.io/adk-docs/context/
[3] https://google.github.io/adk-docs/sessions/session/
[4] https://google.github.io/adk-docs/streaming/custom-streaming-ws/
[5] https://google.github.io/adk-docs/sessions/