Better context management across sub-tasks for ReAct Agent

[marutilai]

Problem

Currently, Katalyst's ReAct agent carries all context between subtasks within a planning session:

• Full action_trace from all previous tasks
• All chat history
• Complete content_store
• No isolation between subtasks

This leads to:

• Context bloat: Later tasks have less available context window for actual work
• Increased costs: Every LLM call includes unnecessary historical context
• Performance degradation: LLMs perform worse with excessive, irrelevant context
• Debugging difficulty: Hard to trace which context influenced a decision

Current Implementation

Current: Everything carries forward
task_1: read_file("config.json") -> 500 tokens added to action_trace
task_2: write_to_file("api.py") -> Previous 500 tokens + new 300 tokens
task_3: run_tests() -> Previous 800 tokens + new 400 tokens

By task 5, we might have 3000+ tokens of mostly irrelevant context

Proposed Solution: Hybrid Context Management

Core Concept

Implement a two-tier context system that intelligently manages what information flows between tasks:

1. Local Context (per-task isolation)
2. Global Context (shared knowledge)

Architecture

  class TaskContext:
      # Local to current task
      local_action_trace: List[Tuple[AgentAction, str]] = []

      # Shared across tasks
      global_knowledge: Dict[str, Any] = {
          "files_modified": Set[str],
          "files_read": Dict[str, str],  # path -> summary
          "discovered_patterns": List[str],
          "error_solutions": Dict[str, str]
      }

      # Always preserved
      content_store: Dict[str, Tuple[str, str]]  # For content_ref
      completed_tasks: List[Tuple[str, str]]  # (task, summary)

Key Features

1. Smart Context Summarization

After each task completes:

  def summarize_task_execution(action_trace):
      return {
          "files_touched": extract_file_operations(action_trace),
          "key_findings": extract_important_observations(action_trace),
          "errors_encountered": extract_errors_and_solutions(action_trace),
          "tools_used": extract_tool_patterns(action_trace)
      }

2. Context Budget Management

  MAX_CONTEXT_TOKENS = 50_000
  TASK_CONTEXT_BUDGET = 10_000

  def manage_context_budget(state):
      current_size = estimate_tokens(state)
      if current_size > MAX_CONTEXT_TOKENS:
          # Compress older task traces into summaries
          state = compress_old_contexts(state)
      return state

3. Relevance-Based Context Injection

Before starting a new task:

  def prepare_task_context(new_task, global_knowledge):
      relevant_context = {
          "related_files": find_relevant_files(new_task, global_knowledge),
          "similar_operations": find_similar_patterns(new_task, global_knowledge),
          "potential_errors": predict_errors(new_task, global_knowledge)
      }
      return relevant_context

Benefits Over Current Approach

1. Performance Improvement

• Current: O(n) context growth where n = number of tasks
• Proposed: O(1) context per task + O(log n) summarized knowledge

2. Cost Reduction

• Current: Each task pays for all previous context
• Proposed: Each task only pays for relevant context

3. Better LLM Performance

• Current: LLM must parse through irrelevant information
• Proposed: LLM sees only pertinent context

4. Improved Debugging

• Current: Difficult to isolate task-specific issues
• Proposed: Clear separation of local vs global context

Implementation Example

  # Task 1: Read configuration
  task_1_context = TaskContext()
  # ... execute task 1 ...
  # After completion:
  global_knowledge.update({
      "files_read": {"config.json": "Contains API_KEY, DATABASE_URL"},
      "discovered_patterns": ["Uses environment variables for secrets"]
  })

  # Task 2: Modify API endpoint
  task_2_context = TaskContext()
  # Inject only relevant context:
  task_2_context.relevant_prior = {
      "config_location": "config.json has API settings",
      # NOT injecting unrelated file operations from task 1
  }

Open Questions for Team Discussion

1. What should be the optimal context budget per task?
2. Which information is always valuable across tasks?
3. Should we allow users to configure context management strategy?
4. How do we handle tasks that explicitly need full history?

Metrics to Track

• Average context size per task
• LLM token costs per session
• Task success rate with/without full context
• Time to completion for multi-task sessions

This proposal maintains the benefits of shared learning while dramatically reducing the downsides of context bloat. The hybrid approach gives us the best of both worlds: isolation where helpful, sharing where valuable.

[shobhi1310]

First need to understand what are the various tasks that katalyst is supposed to perform. Based on that we can disect point 4 and point 2 under Open Questions for Team Discussion. The metric Task success rate with/without full context will automatically help us gauge the optimal contect budget per task.

If there is noticable token reduction with this new context window, then we should allow users to configure because at the end its their call on how much they want to bill themselves and with how much success rate and accomplishment of the task

[marutilai]

I'd prefer for us to standardize the context window and not expose this to end users. It is not a straightforward parameter for them to tweak.