diff --git a/docs/concepts/memory.md b/docs/concepts/memory.md index 38903a1..f548a24 100644 --- a/docs/concepts/memory.md +++ b/docs/concepts/memory.md @@ -141,6 +141,11 @@ agent = Agent( | `DeleteMemoryTool` | Remove memories by ID | | `MemoryPromptTool` | Enrich a system prompt with relevant memories | +When a tool runs inside an ADK agent loop, it resolves the user in this order: +an explicit `user_id` argument, the invocation user from the ADK tool context, +`default_owner_id`, then `default_user_id`. Configured defaults are only used +when no per-call or context user is available. + ## MCP vs SDK Decision | | MCP | SDK Tools | @@ -160,6 +165,8 @@ agent = Agent( | `store_id` | `None` | Redis Agent Memory store ID | | `timeout` | `30` | HTTP timeout in seconds | | `default_namespace` | `default` | Namespace for memory isolation | +| `default_owner_id` | `None` | Default owner ID used when no per-call or context user is available | +| `default_user_id` | `None` | Legacy default user ID fallback after `default_owner_id` | | `search_top_k` | `10` | Default max search results | | `distance_threshold` | `None` | Compatibility alias for search threshold | | `deduplicate` | `True` | Deduplicate when creating memories | diff --git a/src/adk_redis/tools/memory/_base.py b/src/adk_redis/tools/memory/_base.py index 74da8a1..c3c3dbd 100644 --- a/src/adk_redis/tools/memory/_base.py +++ b/src/adk_redis/tools/memory/_base.py @@ -25,6 +25,7 @@ from adk_redis.memory._backends import OPENSOURCE_AGENT_MEMORY_BACKEND from adk_redis.memory._backends import REDIS_AGENT_MEMORY_BACKEND +from adk_redis.memory._utils import read_field from adk_redis.memory._utils import sanitize_managed_identifier from adk_redis.tools.memory._config import MemoryToolConfig @@ -184,15 +185,66 @@ def _get_namespace(self, namespace: str | None = None) -> str: return sanitize_managed_identifier(resolved) return resolved - def _get_user_id(self, user_id: str | None = None) -> str | None: + def _get_user_id( + self, + user_id: str | None = None, + tool_context: Any = None, + ) -> str | None: """Get the user ID to use for operations. + Resolution order: + + 1. Explicit ``user_id`` argument. + 2. The invocation user from the ADK ``tool_context`` (its public + ``user_id`` property), when the tool runs inside an ADK agent loop. + 3. ``config.default_owner_id``. + 4. ``config.default_user_id``. + Args: user_id: Optional user ID override. + tool_context: Optional ADK ToolContext for the current invocation. + Read defensively; a missing or broken context never raises. Returns: - The user ID to use (override or default). + The resolved user ID, or None if no source provides one. """ - return ( - user_id or self._config.default_owner_id or self._config.default_user_id - ) + if user_id: + return user_id + + try: + context_user_id = getattr(tool_context, "user_id", None) + except Exception: + context_user_id = None + if isinstance(context_user_id, str) and context_user_id: + return context_user_id + + return self._config.default_owner_id or self._config.default_user_id + + def _require_memory_scope( + self, + memory: object, + *, + namespace: str, + user_id: str | None, + ) -> None: + """Require a memory record to belong to the resolved operation scope. + + Args: + memory: Backend memory record to validate. + namespace: Resolved namespace for the current invocation. + user_id: Resolved user for the current invocation, if any. + + Raises: + PermissionError: If the memory is outside the resolved scope. + """ + memory_namespace = read_field(memory, "namespace") + memory_user_id = read_field(memory, "owner_id") + if memory_user_id is None: + memory_user_id = read_field(memory, "user_id") + + if memory_namespace != namespace or ( + user_id is not None and memory_user_id != user_id + ): + raise PermissionError( + "Memory is outside the resolved namespace or user scope" + ) diff --git a/src/adk_redis/tools/memory/create.py b/src/adk_redis/tools/memory/create.py index 8ffc4f0..df11423 100644 --- a/src/adk_redis/tools/memory/create.py +++ b/src/adk_redis/tools/memory/create.py @@ -123,19 +123,22 @@ async def run_async(self, **kwargs: Any) -> dict[str, Any]: topics: Optional list of topics/tags. memory_type: Type of memory (semantic, episodic, message). namespace: Optional namespace override. - user_id: Optional user ID override. + user_id: Optional user ID override. When omitted, the user is + resolved from the ADK tool_context invocation user, then the + configured defaults. Returns: A dictionary with status and memory_id. """ # ADK passes parameters in kwargs['args'] args = kwargs.get("args", kwargs) + tool_context = kwargs.get("tool_context") content = args.get("content") topics = args.get("topics", []) memory_type_raw = args.get("memory_type", "semantic") namespace = self._get_namespace(args.get("namespace")) - user_id = self._get_user_id(args.get("user_id")) + user_id = self._get_user_id(args.get("user_id"), tool_context=tool_context) if not content: return {"status": "error", "message": "content is required"} diff --git a/src/adk_redis/tools/memory/delete.py b/src/adk_redis/tools/memory/delete.py index fc443c9..f889bb0 100644 --- a/src/adk_redis/tools/memory/delete.py +++ b/src/adk_redis/tools/memory/delete.py @@ -16,9 +16,10 @@ from __future__ import annotations +import asyncio import logging import re -from typing import Any +from typing import Any, Awaitable, Callable from google.genai import types @@ -29,6 +30,8 @@ logger = logging.getLogger("adk_redis." + __name__) +_PREFLIGHT_BATCH_SIZE = 10 + class DeleteMemoryTool(BaseMemoryTool): """Tool for deleting long-term memories. @@ -103,32 +106,61 @@ def _get_declaration(self) -> types.FunctionDeclaration: ), ) + async def _validate_memory_scopes( + self, + *, + memory_ids: list[str], + get_memory: Callable[..., Awaitable[object]], + namespace: str, + user_id: str | None, + ) -> None: + """Validate memory scopes with bounded concurrent backend reads.""" + for start in range(0, len(memory_ids), _PREFLIGHT_BATCH_SIZE): + batch = memory_ids[start : start + _PREFLIGHT_BATCH_SIZE] + memories = await asyncio.gather( + *(get_memory(memory_id=memory_id) for memory_id in batch) + ) + for memory in memories: + self._require_memory_scope( + memory, + namespace=namespace, + user_id=user_id, + ) + async def run_async(self, **kwargs: Any) -> dict[str, Any]: """Delete long-term memories by ID. Args: memory_ids: List of memory IDs to delete. namespace: Optional namespace override. - user_id: Optional user ID override. + user_id: Optional user ID override. When omitted, the user is + resolved from the ADK tool_context invocation user, then the + configured defaults. Returns: A dictionary with status and deleted_count. """ # ADK passes parameters in kwargs['args'] args = kwargs.get("args", kwargs) + tool_context = kwargs.get("tool_context") memory_ids = args.get("memory_ids", []) - self._get_namespace(args.get("namespace")) - self._get_user_id(args.get("user_id")) + namespace = self._get_namespace(args.get("namespace")) + user_id = self._get_user_id(args.get("user_id"), tool_context=tool_context) if not memory_ids: return {"status": "error", "message": "memory_ids is required"} try: if self._config.backend == OPENSOURCE_AGENT_MEMORY_BACKEND: - response = await self._get_agent_memory_server_client().delete_long_term_memories( + client = self._get_agent_memory_server_client() + await self._validate_memory_scopes( memory_ids=memory_ids, + get_memory=client.get_long_term_memory, + namespace=namespace, + user_id=user_id, ) + response = await client.delete_long_term_memories(memory_ids=memory_ids) status_msg = response.status match = re.search(r"deleted (\d+)", status_msg) deleted_count = int(match.group(1)) if match else 0 @@ -141,6 +173,12 @@ async def run_async(self, **kwargs: Any) -> dict[str, Any]: } async with self._agent_memory() as agent_memory: + await self._validate_memory_scopes( + memory_ids=memory_ids, + get_memory=agent_memory.get_long_term_memory_async, + namespace=namespace, + user_id=user_id, + ) response = await agent_memory.bulk_delete_long_term_memories_async( memory_ids=memory_ids, ) diff --git a/src/adk_redis/tools/memory/prompt.py b/src/adk_redis/tools/memory/prompt.py index ace4a5f..2cc22b9 100644 --- a/src/adk_redis/tools/memory/prompt.py +++ b/src/adk_redis/tools/memory/prompt.py @@ -115,18 +115,21 @@ async def run_async(self, **kwargs: Any) -> dict[str, Any]: query: The query to search for relevant memories. system_prompt: Optional base system prompt to enrich. namespace: Optional namespace override. - user_id: Optional user ID override. + user_id: Optional user ID override. When omitted, the user is + resolved from the ADK tool_context invocation user, then the + configured defaults. Returns: A dictionary with status and enriched_prompt. """ # ADK passes parameters in kwargs['args'] args = kwargs.get("args", kwargs) + tool_context = kwargs.get("tool_context") query = args.get("query") system_prompt = args.get("system_prompt", "") namespace = self._get_namespace(args.get("namespace")) - user_id = self._get_user_id(args.get("user_id")) + user_id = self._get_user_id(args.get("user_id"), tool_context=tool_context) if not query: return {"status": "error", "message": "query is required"} diff --git a/src/adk_redis/tools/memory/search.py b/src/adk_redis/tools/memory/search.py index 421e512..9d15022 100644 --- a/src/adk_redis/tools/memory/search.py +++ b/src/adk_redis/tools/memory/search.py @@ -124,18 +124,21 @@ async def run_async(self, **kwargs: Any) -> dict[str, Any]: query: The search query. limit: Maximum number of memories to return. namespace: Optional namespace override. - user_id: Optional user ID override. + user_id: Optional user ID override. When omitted, the user is + resolved from the ADK tool_context invocation user, then the + configured defaults. Returns: A dictionary with status and list of memories. """ # ADK passes parameters in kwargs['args'] args = kwargs.get("args", kwargs) + tool_context = kwargs.get("tool_context") query = args.get("query") limit = args.get("limit", self._config.search_top_k) namespace = self._get_namespace(args.get("namespace")) - user_id = self._get_user_id(args.get("user_id")) + user_id = self._get_user_id(args.get("user_id"), tool_context=tool_context) if not query: return {"status": "error", "message": "query is required"} diff --git a/src/adk_redis/tools/memory/update.py b/src/adk_redis/tools/memory/update.py index 352c60f..8e107ed 100644 --- a/src/adk_redis/tools/memory/update.py +++ b/src/adk_redis/tools/memory/update.py @@ -118,19 +118,22 @@ async def run_async(self, **kwargs: Any) -> dict[str, Any]: content: New content for the memory. topics: New list of topics/tags. namespace: Optional namespace override. - user_id: Optional user ID override. + user_id: Optional user ID override. When omitted, the user is + resolved from the ADK tool_context invocation user, then the + configured defaults. Returns: A dictionary with status and updated memory info. """ # ADK passes parameters in kwargs['args'] args = kwargs.get("args", kwargs) + tool_context = kwargs.get("tool_context") memory_id = args.get("memory_id") content = args.get("content") topics = args.get("topics") namespace = self._get_namespace(args.get("namespace")) - user_id = self._get_user_id(args.get("user_id")) + user_id = self._get_user_id(args.get("user_id"), tool_context=tool_context) if not memory_id: return {"status": "error", "message": "memory_id is required"} @@ -150,6 +153,12 @@ async def run_async(self, **kwargs: Any) -> dict[str, Any]: updates["topics"] = topics client = self._get_agent_memory_server_client() + memory = await client.get_long_term_memory(memory_id=memory_id) + self._require_memory_scope( + memory, + namespace=namespace, + user_id=user_id, + ) response = await client.edit_long_term_memory( memory_id=memory_id, updates=updates, @@ -171,6 +180,14 @@ async def run_async(self, **kwargs: Any) -> dict[str, Any]: update_kwargs["owner_id"] = user_id async with self._agent_memory() as agent_memory: + memory = await agent_memory.get_long_term_memory_async( + memory_id=memory_id + ) + self._require_memory_scope( + memory, + namespace=namespace, + user_id=user_id, + ) response = await agent_memory.update_long_term_memory_async( memory_id=memory_id, **update_kwargs, diff --git a/tests/tools/test_memory_tools.py b/tests/tools/test_memory_tools.py index db8a33c..f87d7e6 100644 --- a/tests/tools/test_memory_tools.py +++ b/tests/tools/test_memory_tools.py @@ -14,6 +14,7 @@ """Tests for Redis Agent Memory tools.""" +import asyncio from types import SimpleNamespace from unittest.mock import patch @@ -23,6 +24,7 @@ from adk_redis import REDIS_AGENT_MEMORY_BACKEND from adk_redis.tools.memory import CreateMemoryTool from adk_redis.tools.memory import DeleteMemoryTool +from adk_redis.tools.memory import MemoryPromptTool from adk_redis.tools.memory import MemoryToolConfig from adk_redis.tools.memory import SearchMemoryTool from adk_redis.tools.memory import UpdateMemoryTool @@ -36,6 +38,10 @@ def __init__(self): self.deleted_ids = [] self.search_request = None self.update_kwargs = None + self.memory_namespace = "test-ns" + self.memory_owner_id = "alice" + self.active_gets = 0 + self.max_concurrent_gets = 0 async def __aenter__(self): return self @@ -68,6 +74,19 @@ async def update_long_term_memory_async(self, **kwargs): self.update_kwargs = kwargs return SimpleNamespace(id=kwargs["memory_id"]) + async def get_long_term_memory_async(self, *, memory_id): + self.active_gets += 1 + self.max_concurrent_gets = max(self.max_concurrent_gets, self.active_gets) + try: + await asyncio.sleep(0) + return SimpleNamespace( + id=memory_id, + namespace=self.memory_namespace, + owner_id=self.memory_owner_id, + ) + finally: + self.active_gets -= 1 + async def bulk_delete_long_term_memories_async(self, *, memory_ids): self.deleted_ids.extend(memory_ids) return SimpleNamespace(deleted=memory_ids, errors=None) @@ -78,6 +97,12 @@ class FakeAgentMemoryServerClient: def __init__(self): self.add_memory_kwargs = None + self.edit_memory_kwargs = None + self.deleted_ids = [] + self.memory_namespace = "test_ns" + self.memory_user_id = "alice" + self.active_gets = 0 + self.max_concurrent_gets = 0 async def add_memory_tool(self, **kwargs): self.add_memory_kwargs = kwargs @@ -87,6 +112,27 @@ async def add_memory_tool(self, **kwargs): "summary": "Memory created successfully", } + async def get_long_term_memory(self, *, memory_id): + self.active_gets += 1 + self.max_concurrent_gets = max(self.max_concurrent_gets, self.active_gets) + try: + await asyncio.sleep(0) + return SimpleNamespace( + id=memory_id, + namespace=self.memory_namespace, + user_id=self.memory_user_id, + ) + finally: + self.active_gets -= 1 + + async def edit_long_term_memory(self, **kwargs): + self.edit_memory_kwargs = kwargs + return SimpleNamespace(id=kwargs["memory_id"]) + + async def delete_long_term_memories(self, *, memory_ids): + self.deleted_ids.extend(memory_ids) + return SimpleNamespace(status=f"ok: deleted {len(memory_ids)}") + @pytest.fixture def config(): @@ -111,6 +157,70 @@ def test_memory_tool_config_accepts_opensource_backend(): ) +def test_get_user_id_resolution_order(config): + """_get_user_id resolves explicit arg, context user, then defaults.""" + tool = SearchMemoryTool(config=config) + tool_context = SimpleNamespace(user_id="context-user") + + # Explicit user_id beats the tool_context user. + assert ( + tool._get_user_id("explicit-user", tool_context=tool_context) + == "explicit-user" + ) + + # The tool_context user beats configured defaults. + assert tool._get_user_id(None, tool_context=tool_context) == "context-user" + + # Configured defaults still apply with no tool_context. + assert tool._get_user_id(None) == "alice" + + # An empty context user counts as absent. + empty_context = SimpleNamespace(user_id="") + assert tool._get_user_id(None, tool_context=empty_context) == "alice" + + # A context without a user_id attribute falls through safely. + assert tool._get_user_id(None, tool_context=object()) == "alice" + + +def test_get_user_id_without_defaults_returns_none(): + """_get_user_id returns None when no source provides a user.""" + tool = SearchMemoryTool(config=MemoryToolConfig()) + assert tool._get_user_id(None, tool_context=object()) is None + + +@pytest.mark.asyncio +async def test_search_memory_tool_scopes_to_tool_context_user( + config, fake_client +): + """SearchMemoryTool scopes search to the ADK tool_context user.""" + tool = SearchMemoryTool(config=config) + tool_context = SimpleNamespace(user_id="bob") + with patch.object(tool, "_get_client", return_value=fake_client): + result = await tool.run_async( + args={"query": "seat"}, tool_context=tool_context + ) + + assert result["status"] == "success" + assert fake_client.search_request["filter"] == { + "namespace": {"eq": "test-ns"}, + "ownerId": {"eq": "bob"}, + } + + +@pytest.mark.asyncio +async def test_create_memory_tool_uses_tool_context_user(config, fake_client): + """CreateMemoryTool stamps records with the ADK tool_context user.""" + tool = CreateMemoryTool(config=config) + tool_context = SimpleNamespace(user_id="bob") + with patch.object(tool, "_get_client", return_value=fake_client): + result = await tool.run_async( + args={"content": "User likes tea."}, tool_context=tool_context + ) + + assert result["status"] == "success" + assert fake_client.created_records[0]["ownerId"] == "bob" + + @pytest.mark.asyncio async def test_create_memory_tool_writes_record(config, fake_client): """CreateMemoryTool writes a Redis Agent Memory record.""" @@ -142,6 +252,26 @@ async def test_search_memory_tool_uses_owner_and_namespace_filter( } +@pytest.mark.asyncio +async def test_memory_prompt_tool_scopes_to_tool_context_user( + config, fake_client +): + """MemoryPromptTool scopes its request to the ADK invocation user.""" + tool = MemoryPromptTool(config=config) + + with patch.object(tool, "_get_client", return_value=fake_client): + result = await tool.run_async( + args={"query": "seat"}, + tool_context=SimpleNamespace(user_id="bob"), + ) + + assert result["status"] == "success" + assert fake_client.search_request["filter"] == { + "namespace": {"eq": "test-ns"}, + "ownerId": {"eq": "bob"}, + } + + @pytest.mark.asyncio async def test_update_memory_tool_calls_update(config, fake_client): """UpdateMemoryTool calls Redis Agent Memory update.""" @@ -160,6 +290,50 @@ async def test_update_memory_tool_calls_update(config, fake_client): } +@pytest.mark.asyncio +async def test_update_memory_tool_scopes_to_tool_context_user( + config, fake_client +): + """UpdateMemoryTool verifies and applies the ADK invocation user.""" + tool = UpdateMemoryTool(config=config) + fake_client.memory_owner_id = "bob" + + with patch.object(tool, "_get_client", return_value=fake_client): + result = await tool.run_async( + args={"memory_id": "memory-1", "content": "Updated"}, + tool_context=SimpleNamespace(user_id="bob"), + ) + + assert result["status"] == "success" + assert fake_client.update_kwargs["owner_id"] == "bob" + + +@pytest.mark.asyncio +async def test_opensource_update_rejects_memory_owned_by_another_user(): + """Self-hosted updates cannot mutate another invocation user's memory.""" + fake_client = FakeAgentMemoryServerClient() + fake_client.memory_user_id = "alice" + config = MemoryToolConfig( + backend=OPENSOURCE_AGENT_MEMORY_BACKEND, + default_namespace="test_ns", + ) + tool = UpdateMemoryTool(config=config) + + with patch.object( + tool, + "_get_agent_memory_server_client", + return_value=fake_client, + ): + result = await tool.run_async( + args={"memory_id": "memory-1", "content": "Updated"}, + tool_context=SimpleNamespace(user_id="bob"), + ) + + assert result["status"] == "error" + assert "outside the resolved" in result["message"] + assert fake_client.edit_memory_kwargs is None + + @pytest.mark.asyncio async def test_delete_memory_tool_calls_bulk_delete(config, fake_client): """DeleteMemoryTool deletes memory IDs.""" @@ -172,6 +346,63 @@ async def test_delete_memory_tool_calls_bulk_delete(config, fake_client): assert fake_client.deleted_ids == ["memory-1"] +@pytest.mark.asyncio +async def test_delete_memory_tool_bounds_preflight_concurrency( + config, fake_client +): + """Managed delete preflight reads use bounded concurrency.""" + tool = DeleteMemoryTool(config=config) + memory_ids = [f"memory-{index}" for index in range(25)] + + with patch.object(tool, "_get_client", return_value=fake_client): + result = await tool.run_async(args={"memory_ids": memory_ids}) + + assert result["status"] == "success" + assert fake_client.max_concurrent_gets == 10 + assert fake_client.deleted_ids == memory_ids + + +@pytest.mark.asyncio +async def test_opensource_delete_bounds_preflight_concurrency(): + """Self-hosted delete preflight reads use bounded concurrency.""" + fake_client = FakeAgentMemoryServerClient() + config = MemoryToolConfig( + backend=OPENSOURCE_AGENT_MEMORY_BACKEND, + default_namespace="test_ns", + default_user_id="alice", + ) + tool = DeleteMemoryTool(config=config) + memory_ids = [f"memory-{index}" for index in range(25)] + + with patch.object( + tool, + "_get_agent_memory_server_client", + return_value=fake_client, + ): + result = await tool.run_async(args={"memory_ids": memory_ids}) + + assert result["status"] == "success" + assert fake_client.max_concurrent_gets == 10 + assert fake_client.deleted_ids == memory_ids + + +@pytest.mark.asyncio +async def test_delete_memory_tool_rejects_cross_user_ids(config, fake_client): + """DeleteMemoryTool validates ownership before deleting any memory.""" + tool = DeleteMemoryTool(config=config) + fake_client.memory_owner_id = "alice" + + with patch.object(tool, "_get_client", return_value=fake_client): + result = await tool.run_async( + args={"memory_ids": ["memory-1"]}, + tool_context=SimpleNamespace(user_id="bob"), + ) + + assert result["status"] == "error" + assert "outside the resolved" in result["message"] + assert fake_client.deleted_ids == [] + + @pytest.mark.asyncio async def test_create_memory_tool_can_use_agent_memory_server_backend(): """CreateMemoryTool can write through the self-hosted backend."""