Agent Workflow Designer

Use a typed handoff dataclass — never pass raw strings between stages.

from dataclasses import dataclass
import anthropic

@dataclass
class ReviewHandoff:
    pr_diff: str
    stage1_findings: str = ""
    stage2_security: str = ""
    token_budget_used: int = 0

def run_pipeline(diff: str) -> str:
    client = anthropic.Anthropic()
    ctx = ReviewHandoff(pr_diff=diff)

    # Stage 1: Analyze
    r1 = client.messages.create(
        model="claude-3-5-sonnet-20241022", max_tokens=1024,
        system="Senior code reviewer. Identify logic errors, performance issues, missing tests. Be specific with line numbers.",
        messages=[{"role": "user", "content": f"Analyze this diff:\n{ctx.pr_diff[:4000]}"}]
    )
    ctx.stage1_findings = r1.content[0].text
    ctx.token_budget_used += r1.usage.input_tokens + r1.usage.output_tokens

    # Stage 2: Security — receives SUMMARY of stage 1, not full diff again
    r2 = client.messages.create(
        model="claude-3-5-sonnet-20241022", max_tokens=512,
        system="Security engineer. Check for SQL injection, XSS, hardcoded secrets, insecure deserialization.",
        messages=[{"role": "user", "content": f"Findings:\n{ctx.stage1_findings}\n\nDiff (first 2000 chars):\n{ctx.pr_diff[:2000]}"}]
    )
    ctx.stage2_security = r2.content[0].text

    # Stage 3: Write review — synthesize both
    r3 = client.messages.create(
        model="claude-3-5-sonnet-20241022", max_tokens=800,
        system="Write a constructive GitHub PR review. Lead with positives, group issues by severity (blocking/non-blocking).",
        messages=[{"role": "user", "content": f"Findings:\n{ctx.stage1_findings}\n\nSecurity:\n{ctx.stage2_security}"}]
    )
    return r3.content[0].text

Key pattern: Stage 2 receives a summary of stage 1, not the full diff again — prevents context bleed and cuts token cost ~40%.

Use a ContextBudget tracker — allocate at most 60% of remaining budget per step.

CONTEXT_LIMITS = {"claude-3-5-sonnet-20241022": 200_000}

class ContextBudget:
    def __init__(self, model: str, reserve_pct=0.15):
        total = CONTEXT_LIMITS.get(model, 128_000)
        self.reserve = int(total * reserve_pct)
        self.used = 0
        self.total = total

    @property
    def remaining(self):
        return self.total - self.reserve - self.used

    def allocate(self, step: str, requested: int) -> int:
        allocated = min(requested, int(self.remaining * 0.60))
        print(f"[{step}] allocated {allocated:,} / {self.remaining:,} remaining")
        return allocated

    def consume(self, tokens: int):
        self.used += tokens

# Usage in pipeline:
budget = ContextBudget("claude-3-5-sonnet-20241022")
for stage in stages:
    allowed = budget.allocate(stage.name, stage.requested_tokens)
    input_text = stage_input[:int(allowed * 3.8)]  # ~3.8 chars/token
    result = call_agent(stage, input_text, max_tokens=min(2048, allowed // 4))
    budget.consume(result.usage.input_tokens + result.usage.output_tokens)

For stage 4 specifically: stages 1-3 are likely passing their full output forward. Fix: ask each stage to end with a 200-token summary. Stages 3+ forward only the summary chain — not cumulative full outputs. With this pattern, stage 4 consistently has 80K+ tokens available regardless of earlier stage output sizes.