Skip to content

Memory Bank not persisting memories when Agent is called via Gemini Enterprise (works in Playground) #4463

@lucasarlego

Description

@lucasarlego

Environment:

  • ADK Version: 1.25.0 (also reproduced in 1.24.1)
  • Service: Agent Engine (Reasoning Engine)
  • Memory Services: VertexAiMemoryBankService and VertexAiSessionService
  • Environment: Gemini Enterprise

Description

I have deployed an agent using the Google ADK on Agent Engine. The agent is configured with VertexAiMemoryBankService and VertexAiSessionService.

The issue: When testing the agent directly through the Agent Engine Playground, the memory bank works as expected, and memories are correctly persisted. However, when the same deployed agent is accessed via Gemini Enterprise, the after_agent_callback (intended to save the session to memory) fails to generate/persist memories.

It appears that the invocation_context or the memory service state might not be correctly handled or triggered during the hand-off between Gemini Enterprise and the Reasoning Engine.


Code Snippets

1. Deployment Configuration: I am using factories to inject the services into AdkApp:

Python

def setup_memory_bank(env_vars: dict):
    project = env_vars.get("GOOGLE_CLOUD_PROJECT")
    location = env_vars.get("GOOGLE_CLOUD_LOCATION")
    engine_id = env_vars.get('REASONING_ENGINE_ID')

    return lambda: VertexAiMemoryBankService(
        project=project,
        location=location,
        agent_engine_id=engine_id
    )

def setup_vertexai_sessions(env_vars: dict):
    project = env_vars.get("GOOGLE_CLOUD_PROJECT")
    location = env_vars.get("GOOGLE_CLOUD_LOCATION")
    engine_id = env_vars.get('REASONING_ENGINE_ID')

    return lambda: VertexAiSessionService(
        project=project,
        location=location,
        agent_engine_id=engine_id
    )

# Deployment logic
adk_app = AdkApp(
    agent=root_agent,
    enable_tracing=True,
    memory_service_builder=setup_memory_bank(env_vars),
    session_service_builder=setup_vertexai_sessions(env_vars)
)

2. Memory Callback: This callback is registered as an after_agent_callback:

Python

async def add_session_to_memory(
        callback_context: CallbackContext
) -> Optional[types.Content]:
    """Automatically save completed sessions to memory bank"""
    if hasattr(callback_context, "_invocation_context"):
        invocation_context = callback_context._invocation_context
        if invocation_context.memory_service:
            # This part seems to be bypassed or fails silently in Gemini Enterprise
            await callback_context.add_events_to_memory(
                events=invocation_context.session.events
            )

3. Agent Initialization:

Python

root_agent = Agent(
    model=ROOT_AGENT_MODEL,
    name="my_agent",
    tools=[PreloadMemoryTool(), question_analyser],
    before_agent_callback=setup_before_agent_call,
    after_agent_callback=add_session_to_memory,
    generate_content_config=types.GenerateContentConfig(temperature=0.01),
)

Steps to Reproduce

  1. Deploy the agent to Agent Engine using ADK v1.25.0.
  2. Interact with the agent via Agent Engine Playground -> Result: Memory is saved.
  3. Register the agent and interact with it via Gemini Enterprise.
  4. Check the Memory Bank -> Result: No new memories are created/updated.

Expected Behavior

The after_agent_callback should consistently trigger add_events_to_memory regardless of whether the agent is called from the Playground or an external integration like Gemini Enterprise.

Additional Context

  • Is there any specific requirement for Gemini Enterprise to pass down the session/invocation context so the ADK can identify the memory service?
  • Tracing is enabled, but no explicit errors are shown regarding the memory write operation when coming from Gemini Enterprise.

Metadata

Metadata

Assignees

No one assigned

    Labels

    agent engine[Component] This issue is related to Vertex AI Agent Engine

    Type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions