How to ensure that worker agent - custom BaseAgent wrapped in an AgentTool - writes to the same session.state as the Orchestrator Agent?

[vivekkrd]

Hi all: I have been working on a typical architecture with 1 Orchestrator agent and multiple worker agents. However, I want to use a few worker agents which don't use LLM as Custom BaseAgents wrapped in an AgentTool. However, when I want to write the output of these BaseAgent-based worker agent-tools to the session.state, I end up have difficulty accessing the Orchestrator session.state. How do I fix this?

Worker-agent Code for reference:

"""
AI Brief Generator V5 - Deduplication Agent
BaseAgent implementation - deterministic, no LLM
Embedding-based cosine similarity only
"""
from google.adk.agents.base_agent import BaseAgent
from google.adk.events import Event
from google.genai import types
from typing import List, Dict, AsyncGenerator
import numpy as np

import config
from tools.dedup_tools import (
get_embedding_model,
create_embedding_input,
)

def _make_event(message: str) -> Event:
"""Helper to create an Event with a text message."""
content = types.Content(parts=[types.Part(text=message)])
return Event(author="deduplication", content=content)

=============================================================================

DeduplicationAgent (BaseAgent, deterministic, no LLM)

=============================================================================

class DeduplicationAgent(BaseAgent):
"""
Deterministic semantic deduplication agent.

- No LLM
- No instructions
- No judgment
- Embedding-based cosine similarity only
"""

async def _run_async_impl(self, ctx) -> AsyncGenerator:
    """Execute deduplication pipeline."""
    
    # DEBUG: Log the FULL session object details
    print(f"[DEDUP DEBUG] === SESSION OBJECT DETAILS ===")
    print(f"[DEDUP DEBUG] ctx type: {type(ctx)}")
    print(f"[DEDUP DEBUG] ctx.session type: {type(ctx.session)}")
    print(f"[DEDUP DEBUG] ctx.session id(): {id(ctx.session)}")  # Python object ID
    print(f"[DEDUP DEBUG] ctx.session.id: {getattr(ctx.session, 'id', 'NO ID')}")
    print(f"[DEDUP DEBUG] ctx.session.user_id: {getattr(ctx.session, 'user_id', 'NO USER_ID')}")
    print(f"[DEDUP DEBUG] ctx.session.app_name: {getattr(ctx.session, 'app_name', 'NO APP_NAME')}")
    print(f"[DEDUP DEBUG] ctx.session.state id(): {id(ctx.session.state) if hasattr(ctx.session, 'state') else 'NO STATE'}")
    print(f"[DEDUP DEBUG] ctx.session.state keys: {list(ctx.session.state.keys()) if hasattr(ctx.session, 'state') else 'NO STATE'}")
    
    # -------------------------------------------------
    # STEP 0: Read vetted articles
    # -------------------------------------------------
    vetted_articles: List[Dict] = ctx.session.state.get("vetted_articles", [])
    print(f"[DEDUP DEBUG] vetted_articles count from ctx.session.state: {len(vetted_articles)}")
    
    if not vetted_articles:
        print("[DEDUP] No vetted articles to deduplicate")
        ctx.session.state["unique_articles"] = []
        yield _make_event("No vetted articles to deduplicate")
        return
    
    # Only deduplicate accepted articles
    accepted = [a for a in vetted_articles if a.get("is_accepted", False)]
    
    if not accepted:
        print("[DEDUP] No accepted articles to deduplicate")
        ctx.session.state["unique_articles"] = []
        yield _make_event("No accepted articles to deduplicate")
        return
    
    print(f"[DEDUP] Deduplicating {len(accepted)} accepted articles")
    
    # -------------------------------------------------
    # STEP 1: Create embedding inputs
    # -------------------------------------------------
    texts: List[str] = []
    article_urls: List[str] = []
    
    for article in accepted:
        article_url = article.get("url", "")
        if not article_url:
            continue
        
        text = create_embedding_input(article)
        texts.append(text)
        article_urls.append(article_url)
    
    if not texts:
        print("[DEDUP] No valid articles with URLs to deduplicate")
        ctx.session.state["unique_articles"] = accepted
        yield _make_event("No valid articles with URLs")
        return
    
    # -------------------------------------------------
    # STEP 2: Compute embeddings (batch)
    # -------------------------------------------------
    model = get_embedding_model()
    if model is None or model == "fallback":
        print("[DEDUP] Embedding model unavailable, skipping deduplication")
        ctx.session.state["unique_articles"] = accepted
        yield _make_event("Embedding model unavailable")
        return
    
    embeddings = model.encode(texts, convert_to_numpy=True)
    
    # Normalize vectors (safety)
    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
    norms[norms == 0] = 1  # Avoid division by zero
    embeddings = embeddings / norms
    
    embedding_map: Dict[str, np.ndarray] = {
        article_urls[i]: embeddings[i]
        for i in range(len(article_urls))
    }
    
    # -------------------------------------------------
    # STEP 3: Pairwise similarity comparison
    # -------------------------------------------------
    threshold = config.SIMILARITY_THRESHOLD  # e.g. 0.85
    duplicates: Dict[str, str] = {}  # duplicate_url -> canonical_url
    
    for i, url1 in enumerate(article_urls):
        if url1 in duplicates:
            continue
        
        emb1 = embedding_map[url1]
        art1 = next(a for a in accepted if a.get("url") == url1)
        
        for j in range(i + 1, len(article_urls)):
            url2 = article_urls[j]
            if url2 in duplicates:
                continue
            
            emb2 = embedding_map[url2]
            similarity = float(np.dot(emb1, emb2))
            
            if similarity < threshold:
                continue
            
            art2 = next(a for a in accepted if a.get("url") == url2)
            
            # STEP 4: Resolve which article to keep
            score1 = art1.get("total_score", 0)
            score2 = art2.get("total_score", 0)
            date1 = art1.get("published_at", "")
            date2 = art2.get("published_at", "")
            
            if score1 > score2 or (score1 == score2 and date1 <= date2):
                duplicates[url2] = url1
                print(f"[DEDUP] Duplicate: '{art2.get('title', '')[:40]}...' (sim={similarity:.3f})")
            else:
                duplicates[url1] = url2
                print(f"[DEDUP] Duplicate: '{art1.get('title', '')[:40]}...' (sim={similarity:.3f})")
                break
    
    # -------------------------------------------------
    # STEP 5: Build unique_articles list
    # -------------------------------------------------
    unique_articles = []
    
    for article in accepted:
        url = article.get("url", "")
        article_copy = dict(article)
        
        if url in duplicates:
            article_copy["is_duplicate"] = True
            article_copy["duplicate_of"] = duplicates[url]
        else:
            article_copy["is_duplicate"] = False
            article_copy["duplicate_of"] = None
            unique_articles.append(article_copy)
    
    # Set the state
    ctx.session.state["unique_articles"] = unique_articles
    
    # DEBUG: Verify state was set
    print(f"[DEDUP DEBUG] Set unique_articles in ctx.session.state: {len(unique_articles)}")
    print(f"[DEDUP DEBUG] Verify read back: {len(ctx.session.state.get('unique_articles', []))}")
    
    result = f"Deduplication complete: {len(duplicates)} duplicates removed, {len(unique_articles)} unique remain"
    print(f"[DEDUP] {result}")
    yield _make_event(result)

Create agent instance and FunctionTool

async def dedup_call_func(tool_context: ToolContext):
"""
This function is called by the orchestrator to run the deduplication agent.
"""
dedup_agent = DeduplicationAgent(
name="deduplication_agent",
description="Removes semantic duplicates using embedding similarity (threshold: 0.85). Deterministic, no LLM."
)

async for event in dedup_agent.run_async(
    input_str="Deduplicating articles...",
    session_id = tool_context.session.id   
):
    yield event

For orchestrator compatibility - wrap as FunctionTool

from google.adk.tools import FunctionTool
deduplication_tool = FunctionTool(
name="deduplication_tool",
description="Removes semantic duplicates using embedding similarity. Deterministic, no LLM.",
func=dedup_call_func

[adk-bot]

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

TLDR: Your DeduplicationAgent, when called via run_async inside the tool, creates a new, separate session. Changes to its session.state are not reflected back to the orchestrator's in-memory session in the same turn. The solution is to move the agent's logic into a regular function that the tool can call directly, using the tool_context to access and modify the session state. [1]

---

Hello! I can help with this. The issue you're encountering is due to how sessions are managed when you invoke an agent from within a tool.

The Problem: Separate Sessions

When your dedup_call_func tool calls dedup_agent.run_async(session_id=tool_context.session.id), the ADK framework starts a completely new agent execution. Even though you pass the same session_id, it loads that session from the database into a new in-memory session object for the DeduplicationAgent.

Your DeduplicationAgent then correctly modifies its own session state (ctx.session.state["unique_articles"] = ...) and saves it back to the database. However, the Orchestrator agent that called the tool is still operating with its own, separate in-memory session object, which remains unchanged for the duration of its current turn. The Orchestrator would only see the changes made by your tool if it were to re-read the session from the database in a subsequent turn.

The Solution: Use ToolContext Directly

The most effective way to share state between an orchestrator and its tools is to use the tool_context object that is passed to every tool function. [1] This context provides direct access to the orchestrator's live session state. Instead of wrapping a BaseAgent in a tool, you should place your business logic into a function that the tool can call directly.

Here is a refactored version of your code demonstrating this pattern:

# No longer a BaseAgent, but a regular Python module/function with the core logic.
# This makes the code simpler and more direct.

import numpy as np
from google.adk.tools import ToolContext

# All your helper functions like create_embedding_input, get_embedding_model, etc.
# remain the same.

def run_deduplication_logic(vetted_articles: list, similarity_threshold: float) -> tuple[list, int]:
    """Contains the core deduplication logic, refactored from the BaseAgent."""
    
    if not vetted_articles:
        print("[DEDUP] No vetted articles to deduplicate")
        return [], 0

    accepted = [a for a in vetted_articles if a.get("is_accepted", False)]
    if not accepted:
        print("[DEDUP] No accepted articles to deduplicate")
        return [], 0

    print(f"[DEDUP] Deduplicating {len(accepted)} accepted articles")

    # ... (Steps 1 & 2: Create embeddings as before)
    # texts, article_urls = ...
    # embeddings = ...

    # For brevity, assuming `embeddings` and `article_urls` are created correctly.
    # Let's imagine the rest of your logic from steps 1-4 is here.
    
    # A placeholder for the complex logic to produce the final unique list
    # and the count of duplicates.
    unique_articles_list = accepted # Replace with your actual unique list
    duplicates_count = 0

    # ... (Step 5: Build unique_articles list logic would be here)
    # Your logic for pairwise comparison and building `unique_articles` goes here.
    # For this example, let's assume it populates `unique_articles_list`
    # and `duplicates_count`.
    
    print(f"[DEDUP] Deduplication complete: {duplicates_count} duplicates removed, {len(unique_articles_list)} unique remain")
    
    return unique_articles_list, duplicates_count


# This is the function that the orchestrator will call via the FunctionTool.
# It now accepts ToolContext.
async def dedup_call_func(tool_context: ToolContext) -> str:
    """
    This function is called by the orchestrator to run the deduplication logic.
    It directly reads from and writes to the orchestrator's session state.
    """
    print("[DEDUP TOOL] Running deduplication function.")
    
    # 1. Read directly from the tool_context's state
    vetted_articles = tool_context.state.get("vetted_articles", [])
    
    # 2. Run the business logic
    # You can get the threshold from config or pass it in.
    similarity_threshold = 0.85 # Example value
    unique_articles, duplicates_removed = run_deduplication_logic(vetted_articles, similarity_threshold)
    
    # 3. Write the result directly back to the tool_context's state
    tool_context.state["unique_articles"] = unique_articles
    
    print(f"[DEDUP TOOL DEBUG] Set unique_articles in state: {len(tool_context.state.get('unique_articles', []))}")
    
    # 4. Return a result message to the orchestrator
    result_message = f"Deduplication complete: {duplicates_removed} duplicates removed, {len(unique_articles)} unique remain."
    return result_message


# The FunctionTool definition remains largely the same, just ensure it's pointing
# to the new `dedup_call_func`.
from google.adk.tools import FunctionTool

deduplication_tool = FunctionTool(
    name="deduplication_tool",
    description="Removes semantic duplicates from a list of articles provided in the session state.",
    func=dedup_call_func
)

By following this pattern, your tool operates within the same session context as the orchestrator, allowing for seamless and immediate state sharing. For more complex scenarios, you might consider using an AgentTool for true agent-to-agent delegation, but for sharing and modifying state, this direct approach is recommended. [2]

I hope this helps clarify the behavior and gets you on the right track!

---

[1] Custom Tools
[2] Multi-Agent Architectures

[vivekkrd]

While converting to a simpler deterministic FunctionTool works, I wanted to understand how to make something work for a case that involves Custom BaseAgent as one of the worker agents