Memory, context cache/compaction, latest ADK#1323
Memory, context cache/compaction, latest ADK#1323dobesv wants to merge 1 commit intokagent-dev:mainfrom
Conversation
dobesv
requested review from
EItanya,
ilackarms,
peterj and
yuval-k
as code owners
There was a problem hiding this comment.
Pull request overview
This pull request updates the codebase to support the latest Google ADK (v1.25.0) and adds comprehensive configuration for new ADK features including context management (compaction and caching), memory services (InMemory, VertexAI, and MCP-based), and agent resumability.
Changes:
- Added TypeScript type definitions and React UI components for configuring context management, memory, and resumability features in the agent creation interface
- Extended Go CRDs and API types to support context configuration (compaction/caching), memory configuration (InMemory/VertexAI/MCP), and resumability settings
- Implemented Python integration with ADK for context configs, memory services (including new McpMemoryService), and resumability, along with comprehensive unit tests for context configuration
- Included a sample Redis Agent Memory Server MCP implementation demonstrating how to integrate external memory services
- Updated dependency versions across Python packages (google-adk 1.25.0, google-genai 1.63.0, mcp 1.26.0, litellm 1.81.12, etc.)
Reviewed changes
Copilot reviewed 28 out of 30 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| ui/src/types/index.ts | Added TypeScript type definitions for ContextConfig, MemoryConfig, and ResumabilityConfig |
| ui/src/components/ui/switch.tsx | New Switch component for toggling features |
| ui/src/components/create/ResumabilitySection.tsx | UI section for enabling agent resumability |
| ui/src/components/create/MemorySection.tsx | UI section for configuring agent memory (InMemory/VertexAI/MCP) |
| ui/src/components/create/ContextSection.tsx | UI section for context management (compaction and caching) |
| ui/src/components/AgentsProvider.tsx | Extended validation and form data types for new configurations |
| ui/src/app/agents/new/page.tsx | Integrated new configuration sections into agent creation flow |
| ui/src/app/actions/agents.ts | Pass-through of new configurations to agent spec |
| ui/package.json | Added @radix-ui/react-switch dependency |
| go/api/v1alpha2/agent_types.go | Added CRD types for ContextConfig, MemoryConfig, and ResumabilityConfig |
| go/api/v1alpha2/zz_generated.deepcopy.go | Auto-generated DeepCopy methods for new types |
| go/config/crd/bases/kagent.dev_agents.yaml | CRD schema definitions with validation and defaults |
| go/pkg/adk/types.go | Go types for ADK agent configuration including context, memory, and resumability |
| go/internal/controller/translator/agent/adk_api_translator.go | Translation logic from K8s CRDs to ADK configuration, including memory MCP server resolution |
| go/internal/controller/translator/agent/adk_api_translator_test.go | Tests for context configuration translation |
| python/packages/kagent-adk/src/kagent/adk/types.py | Python types and builders for context/memory/resumability configurations |
| python/packages/kagent-adk/src/kagent/adk/memory.py | New McpMemoryService for delegating memory to MCP servers |
| python/packages/kagent-adk/src/kagent/adk/cli.py | CLI integration for passing context configs to KAgentApp |
| python/packages/kagent-adk/src/kagent/adk/_a2a.py | A2A app integration with memory service and context configurations |
| python/packages/kagent-adk/tests/unittests/test_context_config.py | Comprehensive tests for context configuration parsing and building |
| python/packages/kagent-adk/pyproject.toml | Updated dependencies (google-adk 1.25.0, mcp 1.26.0, etc.) |
| python/samples/adk/basic/pyproject.toml | Updated google-adk dependency to 1.25.0 |
| python/samples/memory/redis-agent-memory-server/* | Complete sample MCP server for Redis Agent Memory integration |
| python/uv.lock | Updated lock file with new package versions |
| helm/kagent-crds/templates/kagent.dev_agents.yaml | Helm template with CRD updates |
Files not reviewed (1)
- ui/package-lock.json: Language not supported
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| now = datetime.utcnow() | ||
| iso_utc = now.replace(microsecond=0).isoformat() + "Z" |
There was a problem hiding this comment.
datetime.utcnow() is deprecated in Python 3.12+ in favor of datetime.now(timezone.utc). Consider updating to use the recommended approach to avoid deprecation warnings and ensure compatibility with newer Python versions.
| name: config?.mcp?.name || "", | ||
| ...config?.mcp, |
There was a problem hiding this comment.
When spreading config.mcp on line 55, if the user changes a field and then that field is being spread again, it could lead to unexpected behavior. Consider restructuring this to avoid spreading the entire mcp object after already setting name. The current pattern (line 54 sets name, line 55 spreads mcp which might overwrite it, line 56 applies the new value) works but is confusing. Better to spread first, then apply specific values.
| name: config?.mcp?.name || "", | |
| ...config?.mcp, | |
| ...config?.mcp, | |
| name: config?.mcp?.name || "", |
| return &adk.McpMemoryConfig{ | ||
| BaseMemoryConfig: adk.BaseMemoryConfig{Type: "mcp"}, | ||
| Name: memConfig.Mcp.Name, | ||
| Kind: memConfig.Mcp.Kind, | ||
| ApiGroup: memConfig.Mcp.ApiGroup, | ||
| ServerConfig: serverConfig, | ||
| }, nil |
There was a problem hiding this comment.
The Kind and ApiGroup fields of McpMemoryConfig have kubebuilder defaults defined in the CRD ("MCPServer" and "kagent.dev" respectively), but when these are used in the translator, the empty string values may be passed if not set by the user. Consider explicitly setting defaults here if memConfig.Mcp.Kind or memConfig.Mcp.ApiGroup are empty strings, to ensure consistent behavior even if the CRD defaults aren't applied.
| func (a *adkApiTranslator) translateMemory(ctx context.Context, agent *v1alpha2.Agent, memConfig *v1alpha2.MemoryConfig) (any, error) { | ||
| if memConfig == nil { | ||
| return nil, nil | ||
| } | ||
|
|
||
| switch memConfig.Type { | ||
| case v1alpha2.MemoryTypeInMemory: | ||
| return &adk.InMemoryConfig{BaseMemoryConfig: adk.BaseMemoryConfig{Type: "in_memory"}}, nil | ||
| case v1alpha2.MemoryTypeVertexAI: | ||
| var projectID *string | ||
| var location *string | ||
| if memConfig.VertexAI != nil { | ||
| if memConfig.VertexAI.ProjectID != "" { | ||
| projectID = &memConfig.VertexAI.ProjectID | ||
| } | ||
| if memConfig.VertexAI.Location != "" { | ||
| location = &memConfig.VertexAI.Location | ||
| } | ||
| } | ||
| return &adk.VertexAIMemoryConfig{ | ||
| BaseMemoryConfig: adk.BaseMemoryConfig{Type: "vertex_ai"}, | ||
| ProjectID: projectID, | ||
| Location: location, | ||
| }, nil | ||
| case v1alpha2.MemoryTypeMcp: | ||
| if memConfig.Mcp == nil { | ||
| return nil, fmt.Errorf("mcp configuration is required for type Mcp") | ||
| } | ||
|
|
||
| // Use the existing translation logic for MCP servers, but capture the result | ||
| // instead of adding it to the agent's tool list. | ||
| tempAgentConfig := &adk.AgentConfig{} | ||
|
|
||
| // Create a synthetic tool definition for the memory server | ||
| // We don't specify any tool names because we want the memory service | ||
| // to discover/use the memory-related tools, not expose them as general agent tools. | ||
| toolDef := &v1alpha2.McpServerTool{ | ||
| TypedLocalReference: v1alpha2.TypedLocalReference{ | ||
| Name: memConfig.Mcp.Name, | ||
| Kind: memConfig.Mcp.Kind, | ||
| ApiGroup: memConfig.Mcp.ApiGroup, | ||
| }, | ||
| } | ||
|
|
||
| // Reuse the translation logic | ||
| err := a.translateMCPServerTarget(ctx, tempAgentConfig, agent.Namespace, toolDef, nil, a.globalProxyURL) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to translate memory MCP server: %w", err) | ||
| } | ||
|
|
||
| // Extract the resulting configuration | ||
| var serverConfig any | ||
| if len(tempAgentConfig.HttpTools) > 0 { | ||
| serverConfig = tempAgentConfig.HttpTools[0] | ||
| } else if len(tempAgentConfig.SseTools) > 0 { | ||
| serverConfig = tempAgentConfig.SseTools[0] | ||
| } else { | ||
| return nil, fmt.Errorf("failed to resolve configuration for memory MCP server %s", memConfig.Mcp.Name) | ||
| } | ||
|
|
||
| return &adk.McpMemoryConfig{ | ||
| BaseMemoryConfig: adk.BaseMemoryConfig{Type: "mcp"}, | ||
| Name: memConfig.Mcp.Name, | ||
| Kind: memConfig.Mcp.Kind, | ||
| ApiGroup: memConfig.Mcp.ApiGroup, | ||
| ServerConfig: serverConfig, | ||
| }, nil | ||
| default: | ||
| return nil, fmt.Errorf("unknown memory type: %s", memConfig.Type) | ||
| } | ||
| } |
There was a problem hiding this comment.
The new memory configuration translation functionality (translateMemory function) lacks test coverage. The test file includes tests for context configuration but doesn't cover memory configuration scenarios such as InMemory, VertexAI, and MCP memory types. Consider adding tests to verify correct translation of all memory types and edge cases.
| class McpMemoryService(BaseMemoryService): | ||
| """Memory service that delegates to an MCP server.""" | ||
|
|
||
| def __init__(self, connection_params: SseConnectionParams | StreamableHTTPConnectionParams): | ||
| super().__init__() | ||
| # Use KAgentMcpToolset for consistent error handling and features | ||
| self.toolset = KAgentMcpToolset(connection_params=connection_params) | ||
| self._tools = None | ||
|
|
||
| async def _ensure_tools(self): | ||
| if self._tools is None: | ||
| tools = await self.toolset.get_tools() | ||
| self._tools = {t.name: t for t in tools} | ||
|
|
||
| async def _call_tool(self, name: str, **kwargs) -> Any: | ||
| await self._ensure_tools() | ||
| if name not in self._tools: | ||
| # Try to find a tool that matches the name loosely or assume the server provides it? | ||
| # For now, strict match. | ||
| logger.warning(f"Tool '{name}' not found in MCP server. Available tools: {list(self._tools.keys())}") | ||
| raise ValueError(f"Tool '{name}' not found in MCP server.") | ||
|
|
||
| tool = self._tools[name] | ||
| logger.debug(f"Calling memory tool '{name}' with kwargs: {kwargs.keys()}") | ||
| return await tool.run(**kwargs) | ||
|
|
||
| async def add_session_to_memory(self, session: Session) -> None: | ||
| """Adds a session to the memory store.""" | ||
| try: | ||
| # We pass the session as a dict | ||
| # session.model_dump(mode='json') creates a dict with JSON-compatible types | ||
| session_data = session.model_dump(mode="json") | ||
| await self._call_tool("add_session_to_memory", session=session_data) | ||
| except Exception as e: | ||
| logger.error(f"Failed to add session to memory: {e}") | ||
| raise | ||
|
|
||
| async def search_memory(self, *, app_name: str, user_id: str, query: str) -> SearchMemoryResponse: | ||
| """Searches the memory store.""" | ||
| try: | ||
| result = await self._call_tool("search_memory", app_name=app_name, user_id=user_id, query=query) | ||
|
|
||
| # The result should be compatible with SearchMemoryResponse | ||
| if isinstance(result, dict): | ||
| # Resilience: convert "text" or string "content" to types.Content structure | ||
| if "memories" in result and isinstance(result["memories"], list): | ||
| for m in result["memories"]: | ||
| if not isinstance(m, dict): | ||
| continue | ||
|
|
||
| # If it has "text" but no "content", map it | ||
| if "text" in m and "content" not in m: | ||
| m["content"] = {"parts": [{"text": m["text"]}], "role": "user"} | ||
|
|
||
| # If "content" is a raw string, wrap it | ||
| if "content" in m and isinstance(m["content"], str): | ||
| m["content"] = {"parts": [{"text": m["content"]}], "role": "user"} | ||
|
|
||
| return SearchMemoryResponse.model_validate(result) | ||
| elif isinstance(result, str): | ||
| try: | ||
| # Try to parse json if string | ||
| data = json.loads(result) | ||
| if isinstance(data, dict): | ||
| return SearchMemoryResponse.model_validate(data) | ||
| except json.JSONDecodeError: | ||
| pass | ||
|
|
||
| # Maybe the string itself is content? Unlikely for SearchMemoryResponse. | ||
|
|
||
| # If result is an object (like a Pydantic model from ADK tools?), try model_dump | ||
| if hasattr(result, "model_dump"): | ||
| return SearchMemoryResponse.model_validate(result.model_dump()) | ||
|
|
||
| raise ValueError(f"Unexpected result type from search_memory: {type(result)}") | ||
| except Exception as e: | ||
| logger.error(f"Failed to search memory: {e}") | ||
| raise |
There was a problem hiding this comment.
The new McpMemoryService class in memory.py lacks unit test coverage. While there are tests for context configuration, there are no tests for the memory service implementation, including tool discovery, session addition, and memory search functionality. Consider adding comprehensive tests for this new functionality.
| mcp | ||
| agent-memory-client | ||
| python-dotenv | ||
| starlette | ||
| uvicorn |
There was a problem hiding this comment.
This requirements.txt installs third-party packages without any version pinning (e.g., mcp, agent-memory-client, starlette, uvicorn), so Docker builds will always pull whatever the latest versions are on PyPI. If one of these packages or its transitive dependencies is compromised upstream, a simple rebuild of this image could introduce malicious code into your cluster without any source change. Use pinned versions and/or a lock file with hashes so that dependencies are reproducible and only change through explicit, reviewed updates.
EItanya
left a comment
EItanya
left a comment
There was a problem hiding this comment.
Hey there @dobesv, thanks so much for the PR. To be honest this PR adds too many features at once. Each of these deserves their own conversation, in fact Memory already has a dedicated issue with a design doc. The best first step with work like this is to attend the community meetings and discuss them with the core community. Looking forward to discussing more :)
|
Most of this is just allowing configuration of built-in ADK features so it doesn't seem like too much innovation. The memory proxy maybe is a bit more of a design question. Context compression was my main pain point with kagent so far, it's pretty useless without it. Do you think you would accept a pull request with just the context options? Or maybe just rip out the memory proxy from this but leave all the built-in ADK features? When are the meetings? |
Update to the latest ADK and support configuration of new ADK features: - context compaction: automatically based on the number of events or tokens - context cache config - memory, either using provided in-memory or vertex memory, or delegating to an MCP server - resumability - ability to enable resumability After this I am hoping get these working: - Show in the UI if compaction is running (similar UI as "thinking...") ideally - Add a button to request manual compaction - Attempt to resume sessions automatically if the pod was restarted (and maybe in some other cases? Not sure)
dobesv
force-pushed
the
context-mem-updates
branch
from
ba1ee24 to
36d4e04
Compare
|
Will open a new pull request with only context configs (cache and compression) and revisit the resumability and memory features later. |
2 participants
Update to the latest ADK and support configuration of new ADK features:
After this I am hoping get these working: