[Bug] Multi User Setup - Data Security Issues

[flobaader]

We need to have the user_id in a toolkit function in order to get the right data when a user calls an agent.
We have a multi-user setup for agent system to handle permissions and data access.
So basically we create one agent per user:

main function

# Cache to store agent instances
_agent_cache = {}

def get_agent_for_user(user: User):
    # Check if agent exists in cache
    if user.id in _agent_cache:
        return _agent_cache[user.id]

    # Create toolkit instances with user metadata
    personal_toolkit = PersonalTools()
    # Create new agent
    agent = Agent(
        session_id=session_id,
        user_id=user_id,
        tools=[personal_toolkit]
    )

    # Cache the agent instance
    _agent_cache[user.id] = agent
    
    return agent

toolkit

class PersonalToolkit(Toolkit)
...
  async def get_personal_data(self,  agent: Agent) -> str:
     user_id = agent.user_id
      # performs db lookup for the data with the user_id
     return personal_data
...

But we are facing issues with that topic: A multi user setup is not document and it seems like not supported as a first class citizen. We did find the agent parameter in the toolkit only by looking at the code.

We are now facing the issue, that when we have many users, we get the wrong user ids by the agent.

This caused us significant data privacy leaks and security issues.

How can we develop a secure multi user setup using your framework?

[flobaader]

@dirkbrnd could you help me with that? We are currently evaluating our agent frameworks, because we have the feeling that there are still many bugs and issues in this library for a production setup, because we really faced some nasty bugs during the last 4 weeks of production.
On the other side I see that you guys are focusing on improving the stability and testing of the library

[ItsRoy69]

I guess for handling this multi user secure setup setting up user isolation where each user gets their own agent instance and adding a context validation before accessing the data is needed

[dirkbrnd]

@flobaader It is an issue with state management in the agent. Your use-case is of course valid. I'll replicate from my side and see what we can do.

[dirkbrnd]

@flobaader Are you using the same session ID across the agents?

[flobaader]

We use session_id = user_id + platform_id

Platform can be WhatsApp or Teams

@dirkbrnd

[flobaader]

@ItsRoy69 yes that's exactly what we did here

[dirkbrnd]

Ah ok that makes sense. Looking into it.

[dirkbrnd]

@flobaader
I made this toy example locally

from typing import Dict

from pydantic import BaseModel
from agno.agent.agent import Agent
from agno.models.openai.chat import OpenAIChat
from agno.tools.duckduckgo import DuckDuckGoTools
from agno.tools.toolkit import Toolkit

class User(BaseModel):
    id: str
    session_id: str
    personal_data: str = ""

users: Dict[str, User] = {}


class PersonalToolkit(Toolkit):

    def __init__(self):
        super().__init__("personal_tools")

        self.register(self.get_user_data)

    def get_user_data(self, agent: Agent) -> str:
        """
        Get the user data for the current user
        """
        user_id = agent.user_id
        # performs db lookup for the data with the user_id
        return users[user_id].personal_data


_agent_cache = {}

def get_agent_for_user(user: User):
    # Check if agent exists in cache
    if user.id in _agent_cache:
        return _agent_cache[user.id]

    # Create new agent
    agent = Agent(
        model=OpenAIChat("gpt-4o-mini"),
        session_id=user.session_id,
        user_id=user.id,
        tools=[PersonalToolkit(),
               DuckDuckGoTools()],
        debug_mode=True 
    )

    # Cache the agent instance
    _agent_cache[user.id] = agent
    
    return agent


for user_id in ["user_1", "user_2", "user_3"]:
    users[user_id] = User(id=user_id, session_id=user_id, personal_data="personal data for user " + user_id)


# User 1 makes a request
agent = get_agent_for_user(users["user_1"])
agent.print_response("What is the weather in Tokyo?")


# User 2 makes a request
agent = get_agent_for_user(users["user_2"])
agent.print_response("What user are you talking to?")


# User 1 makes a request
agent = get_agent_for_user(users["user_1"])
agent.print_response("What user are you talking to?")

# User 3 makes a request
agent = get_agent_for_user(users["user_3"])
agent.print_response("What user are you talking to?")

# User 1 makes a request
agent = get_agent_for_user(users["user_1"])
agent.print_response("What user are you talking to?")

In this case it worked correctly each time. Correct me if this example is missing something or too simple.

[flobaader]

Hi @dirkbrnd the example is correct, we had some issues after running the agent after multiple days and many users.
Users from multiple orgs reported that they got the wrong data back.
Is there any sort of caching or maybe race condition in the code for that?
Maybe we can find the issue with more iterations/parallelism on the code. We use arun for the agent.

[dirkbrnd]

Hmm the async use shouldn't have that impact, but something to look into. And no, there isn't any caching involved, however individual agents have state maintained, something we want to minimise over time.

[dirkbrnd]

To help further I would need to get a more representative example then, since this must be happening in a very specific case it seems.