Skip to content

Latest commit

 

History

History
232 lines (190 loc) · 9.26 KB

File metadata and controls

232 lines (190 loc) · 9.26 KB

Architecture Plan: Context-Isolated Agent Pipeline

Problem

The current SequentialAgent pipeline passes the same session to every sub-agent. ADK's ContentsRequestProcessor walks all session events and includes them in req.Contents for every LLM call. By the time the 5th agent (PR creator) runs, the request contains the full conversation history from all 4 prior agents — easily exceeding Gemini's 1M token/minute quota.

triage → search → plan → fix → PR
         ↑ accumulates ALL prior events into req.Contents

Root Cause in ADK

// google.golang.org/adk/internal/llminternal/contents_processor.go
for e := range ctx.Session().Events().All() {
    events = append(events, e)   // ALL events from the session
}

The SequentialAgent is literally LoopAgent{MaxIterations:1} — same InvocationContext (same Session) is passed to every sub-agent. No branching, no isolation.

Solution: Two-Layer Fix

Layer 1: IncludeContents: "none" + OutputKey (immediate, high-impact)

ADK v0.5.0 exposes llmagent.IncludeContentsNone. When set, the agent only sees its current turn — not the accumulated history from prior agents. Context is passed via {key} template injection in the Instruction field (resolved from session state).

triageAgent, _ := llmagent.New(llmagent.Config{
    Name:            "triage",
    Model:           mdl,
    IncludeContents: llmagent.IncludeContentsNone, // ← ADD THIS
    Instruction:     `Triage Sentry issue {sentry_issue_id} in {repo}...`,
    Tools:           []tool.Tool{getSentryEventTool, updateStatusTool},
    OutputKey:       "triage",
})

searchAgent, _ := llmagent.New(llmagent.Config{
    Name:            "code_search",
    Model:           mdl,
    IncludeContents: llmagent.IncludeContentsNone, // ← ADD THIS
    Instruction:     `Search codebase using triage: {triage}...`,
    Tools:           codeSearchTools,
    OutputKey:       "code_context",
})

planAgent, _ := llmagent.New(llmagent.Config{
    Name:            "planner",
    Model:           mdl,
    IncludeContents: llmagent.IncludeContentsNone, // ← ADD THIS
    Instruction:     `Plan fix from:\nTriage: {triage}\nCode: {code_context}`,
    OutputKey:       "fix_plan",
})

fixAgent, _ := llmagent.New(llmagent.Config{
    Name:            "fixer",
    Model:           mdl,
    IncludeContents: llmagent.IncludeContentsNone, // ← ADD THIS
    Instruction:     `Implement fix:\n{fix_plan}\nRepo: {repo} Branch: {branch}`,
    Tools:           fixerTools,
    OutputKey:       "fix_result",
})

prAgent, _ := llmagent.New(llmagent.Config{
    Name:            "pr_creator",
    Model:           mdl,
    IncludeContents: llmagent.IncludeContentsNone, // ← ADD THIS
    Instruction:     `Create PR:\nFix: {fix_result}\nPlan: {fix_plan}\nTriage: {triage}...`,
    Tools:           prTools,
    OutputKey:       "pr_result",
})

Why this works: Each agent's Instruction is injected into the system prompt (not conversation history). The {key} placeholders are resolved from session state by InjectSessionState. The agent only gets:

  1. Its system prompt (with injected summaries from prior stages)
  2. The single user message triggering this turn
  3. Its own tool calls/responses within the current turn

Impact: Eliminates ~80-90% of token accumulation. Each agent processes only what it needs, not the full history of every prior agent's tool calls and responses.

Layer 2: BeforeModelCallback for tool output compression

Even with IncludeContentsNone, individual agents (especially code_search and fixer) can accumulate large tool outputs within their own turn. Add a callback to cap per-agent token usage:

trimContentsCallback := func(ctx agent.CallbackContext, req *model.LLMRequest) (*model.LLMResponse, error) {
    // Keep only the last N messages to prevent intra-agent context bloat
    const maxMessages = 20
    if len(req.Contents) > maxMessages {
        req.Contents = req.Contents[len(req.Contents)-maxMessages:]
    }
    return nil, nil // nil = proceed with actual model call
}

searchAgent, _ := llmagent.New(llmagent.Config{
    // ...
    BeforeModelCallbacks: []llmagent.BeforeModelCallback{trimContentsCallback},
})

Layer 3: temp: prefix for intermediate state

Large intermediate outputs that only the next agent needs should use temp: prefix so they're discarded from the persisted session:

searchAgent, _ := llmagent.New(llmagent.Config{
    OutputKey: "temp:code_context", // discarded after invocation
})

// The planner reads it via {temp:code_context} in its Instruction
planAgent, _ := llmagent.New(llmagent.Config{
    Instruction: `Plan fix:\nTriage: {triage}\nCode: {temp:code_context}`,
    OutputKey:   "fix_plan", // persisted — needed by both fixer and PR
})

Layer 4: Tool output size limits (defense in depth)

Limit the output of verbose tools at the tool level:

func SearchCode(deps *Dependencies) func(tool.Context, SearchCodeInput) (SearchCodeOutput, error) {
    return func(ctx tool.Context, in SearchCodeInput) (SearchCodeOutput, error) {
        maxResults := in.MaxResults
        if maxResults == 0 {
            maxResults = 10 // reduce from 20
        }
        results, err := deps.Sandbox.SearchCode(ctx, in.Query, maxResults)
        if err != nil {
            return SearchCodeOutput{}, err
        }
        // Truncate if too long
        if len(results) > 8000 {
            results = results[:8000] + "\n... (truncated)"
        }
        return SearchCodeOutput{Results: results}, nil
    }
}

Layer 5: Rate limit retry with backoff

The 429 error includes retryDelay: 31s. Add retry logic in the pipeline runner:

// In ProcessIssue, wrap the runner with retry logic
for event, err := range p.runner.Run(ctx, ...) {
    if err != nil {
        if isRateLimitError(err) {
            delay := extractRetryDelay(err) // parse from error message
            p.logger.Warn("rate limited, waiting", "delay", delay)
            time.Sleep(delay)
            continue
        }
        // ...
    }
}

Data Flow After Changes

┌─────────────────────────────────────────────────────────────┐
│ Session State (key-value store)                              │
│                                                              │
│  sentry_issue_id: "12345"                                    │
│  repo: "my-backend"                                          │
│  branch: "sentry-fix/12345"                                  │
│  triage: "<structured summary>"        ← persisted           │
│  temp:code_context: "<search results>" ← ephemeral           │
│  fix_plan: "<fix details>"             ← persisted           │
│  fix_result: "<what was changed>"      ← persisted           │
│  pr_result: "<PR URL>"                 ← persisted           │
└─────────────────────────────────────────────────────────────┘

  Agent          Sees in req.Contents    Reads from State
  ─────          ────────────────────    ────────────────
  triage         [user_msg]             {sentry_issue_id}, {repo}
  code_search    [user_msg]             {triage}
  planner        [user_msg]             {triage}, {temp:code_context}
  fixer          [user_msg]             {fix_plan}, {repo}, {branch}
  pr_creator     [user_msg]             {fix_result}, {fix_plan}, {triage}

Each agent operates with a clean context window containing only:

  • System prompt (~500-1000 tokens)
  • Injected state values (~2000-5000 tokens each)
  • Single user message (~50 tokens)
  • Its own tool calls within the current turn

Estimated per-agent context: 10-50K tokens (vs current ~200K-1M accumulated).

Implementation Priority

  1. Add IncludeContents: llmagent.IncludeContentsNone to all 5 agents in pipeline.go — this alone fixes the root cause.
  2. Add temp: prefix to code_context OutputKey.
  3. Add tool output truncation in tools.go for search/read tools.
  4. Add BeforeModelCallback to code_search and fixer agents.
  5. Add retry with backoff for 429 errors in ProcessIssue.

Reference: How Amp/Claude Code Solve This

  • Amp's Task tool: Spawns each sub-agent in a completely isolated thread with its own context window. Parent passes a goal string + file paths. Child returns only a compressed summary (~1-2K tokens). File system is the communication bus.

  • Claude Code compaction: When nearing context limits, summarizes the full conversation into a compressed form, then reinitializes with the summary + 5 most recent files. Preserves architectural decisions, drops redundant tool outputs.

  • ADK's approach (what we use): IncludeContents: "none" + OutputKey state injection achieves the same isolation without needing separate sessions or threads. Each agent gets a clean context window with only the structured outputs from prior stages injected via its system prompt.

Our implementation uses ADK's native approach (cheapest change, best fit for the sequential pipeline pattern).