Multi-Agent Tool Orchestration and Context Partitioning

graph TD
    UserPrompt[User Goal: Investigate Incident & Fix Regression] --> Supervisor[Supervisor Agent: LangGraph Coordinator]
    
    subgraph Context Partition 1: SRE Agent
        Supervisor -->|Handoff: Triage Logs| SREAgent[SRE Investigator Agent]
        SREAgent <-->|Scoped MCP: k8s, cloudwatch| MCP_SRE[DevOps MCP Server]
    end

    subgraph Context Partition 2: Data Agent
        Supervisor -->|Handoff: Check Data Corruption| DataAgent[Data Analyst Agent]
        DataAgent <-->|Scoped MCP: postgres_ro| MCP_DB[Database MCP Server]
    end

    subgraph Context Partition 3: Dev Agent
        Supervisor -->|Handoff: Patch Bug & Test| DevAgent[Developer Agent]
        DevAgent <-->|Scoped MCP: git, filesystem, tests| MCP_Dev[Code Execution MCP Server]
    end

    SREAgent -->|Condensed Synthesis| Supervisor
    DataAgent -->|Condensed Synthesis| Supervisor
    DevAgent -->|PR Link & Test Results| Supervisor

sequenceDiagram
    participant Sup as Supervisor
    participant Data as Data Specialist
    participant MCP as Database MCP Server

    Sup->>Data: Handoff Task: "Audit table accounts for duplicate IDs"
    Note over Data: Context Partition Active (0 tokens history)
    Data->>MCP: tools/call [execute_sql: SELECT ...]
    MCP-->>Data: 8,500 Tokens Raw SQL Rows
    Note over Data: Synthesizes findings internally
    Data-->>Sup: Handoff Return: "Found 2 duplicate rows: IDs 991, 992. Billing impact: $0."
    Note over Sup: Ingests only 45 tokens into main state

{
  "origin_agent": "database_specialist",
  "target_agent": "supervisor",
  "task_id": "audit-billing-dup-44",
  "status": "COMPLETED",
  "artifact": {
    "summary": "Identified 2 orphaned duplicate accounts in table 'billing_accounts'.",
    "affected_keys": ["acc_991", "acc_992"],
    "requires_code_patch": true
  },
  "raw_token_expenditure": 9410,
  "compressed_token_payload": 78
}

import os
import asyncio
from typing import TypedDict, Annotated, Sequence, List
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

# 1. State Definition
class MultiAgentState(TypedDict):
    messages: Sequence[BaseMessage]
    next_step: str
    db_findings: str
    sre_findings: str

# 2. Initialize Language Models
llm = ChatOpenAI(model="gpt-4o", temperature=0)

# 3. Helper: Execute Scoped MCP Tool
async def run_mcp_tool(server_script: str, tool_name: str, arguments: dict):
    server_params = StdioServerParameters(
        command="python",
        args=[server_script]
    )
    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:
            await session.initialize()
            result = await session.call_tool(tool_name, arguments)
            return result.content[0].text

# 4. Agent Nodes with Scoped Tool Access

async def supervisor_node(state: MultiAgentState):
    system_prompt = SystemMessage(content="""
    You are the Lead Systems Supervisor. Decide which specialist should act next.
    Options:
    - 'data_specialist': To query and inspect database state.
    - 'sre_specialist': To inspect cloud logs and pod status.
    - 'FINISH': When sufficient evidence is collected to conclude the task.
    """)
    messages = [system_prompt] + list(state["messages"])
    response = await llm.ainvoke(messages)
    
    # Route based on LLM response
    choice = "FINISH"
    if "data_specialist" in response.content.lower():
        choice = "data_specialist"
    elif "sre_specialist" in response.content.lower():
        choice = "sre_specialist"

    return {"next_step": choice, "messages": [response]}

async def data_specialist_node(state: MultiAgentState):
    # Context Partition: Specialist only sees its specific instruction
    query = "SELECT count(*) FROM error_events WHERE resolved = false;"
    
    # Invoke Database MCP Server
    mcp_result = await run_mcp_tool("servers/db_mcp.py", "execute_readonly_sql", {"query": query})
    
    # Compress finding before returning to global state
    synthesis = f"Data Audit: Found unresolved errors. Result: {mcp_result.strip()}"
    return {
        "db_findings": synthesis,
        "messages": [AIMessage(content=synthesis)]
    }

async def sre_specialist_node(state: MultiAgentState):
    # Invoke SRE MCP Server
    mcp_result = await run_mcp_tool("servers/sre_mcp.py", "get_pod_health", {"namespace": "prod"})
    
    synthesis = f"SRE Audit: Checked pod health. Status: {mcp_result.strip()}"
    return {
        "sre_findings": synthesis,
        "messages": [AIMessage(content=synthesis)]
    }

# 5. Build StateGraph Workflow
workflow = StateGraph(MultiAgentState)

workflow.add_node("supervisor", supervisor_node)
workflow.add_node("data_specialist", data_specialist_node)
workflow.add_node("sre_specialist", sre_specialist_node)

workflow.set_entry_point("supervisor")

workflow.add_conditional_edges(
    "supervisor",
    lambda state: state["next_step"],
    {
        "data_specialist": "data_specialist",
        "sre_specialist": "sre_specialist",
        "FINISH": END
    }
)

# After workers execute, hand control back to supervisor
workflow.add_edge("data_specialist", "supervisor")
workflow.add_edge("sre_specialist", "supervisor")

app = workflow.compile()

GET /sse HTTP/1.1
Host: mcp-gateway.internal.enterprise.com
X-Agent-Role: database_specialist
Authorization: Bearer <JWT_DATA_SPECIALIST>

{
  "jsonrpc": "2.0",
  "id": 1,
  "result": {
    "tools": [
      {
        "name": "describe_table_schema",
        "description": "Inspect column names and constraints."
      },
      {
        "name": "execute_readonly_sql",
        "description": "Run read-only analytical queries."
      }
    ]
  }
}