mech.app
AI Agents

TRACE-ROUTER: How Task-Consistent Routing Keeps Multi-Step Agent Workflows from Falling Apart

Why per-call LLM routing breaks agentic workflows, and how task-aware routing preserves state and quality across multi-step executions.

Source: arxiv.org
TRACE-ROUTER: How Task-Consistent Routing Keeps Multi-Step Agent Workflows from Falling Apart

Most LLM routers optimize each call independently. You send a query, the router picks GPT-4 or Claude or Llama based on predicted difficulty, and the call completes. This works fine for single-shot inference. It breaks down when you chain multiple LLM calls into a workflow.

Agentic applications are long-horizon pipelines. A research agent might call an LLM to plan subtasks, then call it again to execute each subtask, then call it a third time to synthesize results. If your router switches models mid-workflow (GPT-4 for planning, Claude for execution, Llama for synthesis), you lose context continuity, introduce formatting inconsistencies, and make it impossible to attribute success or failure to specific routing decisions.

TRACE-ROUTER addresses this by treating the task, not the call, as the unit of routing. It assigns a model once at the start of a workflow, pins all subsequent calls to that backend, and updates its policy using the terminal reward from the completed task. This aligns routing with the actual supervision signal and prevents the compounding errors that occur when you switch models mid-execution.

Why Per-Call Routing Fails for Workflows

Traditional routers evaluate each LLM call in isolation. They predict difficulty, estimate cost-quality trade-offs, and select a model. But in multi-step workflows, the quality of intermediate outputs only matters if they contribute to a successful final outcome.

Credit assignment becomes impossible. If your workflow fails after five LLM calls, which routing decision caused the failure? The router that picked a weak model for step three? Or the router that picked an expensive model for step one, leaving no budget for step four? Per-call routers cannot answer this question because they never see the task-level outcome.

Context drift accumulates. Different models have different instruction-following behaviors, output formats, and reasoning styles. Switching from GPT-4 to Claude mid-workflow means the second model receives context generated by a different model. Prompt engineering that works for one model may not work for another. The workflow becomes brittle.

Latency budgets fragment. If you have a 10-second SLA for a five-step workflow, you need to allocate latency across all five calls. Per-call routers do not coordinate. They might burn 8 seconds on the first call, leaving 2 seconds for the remaining four, which forces you to use weak models and degrades quality.

Task-Consistent Routing Architecture

TRACE-ROUTER treats each incoming task as a contextual bandit problem. At admission, it observes task features (query complexity, domain, expected number of steps), selects a model, and pins all subsequent calls in that workflow to the chosen backend.

Admission-time routing. The router makes one decision per task, not per call. This decision is based on:

  • Task metadata (domain, expected complexity)
  • Current workload (queue depth, latency distribution)
  • Historical performance (which models succeed on similar tasks)

Pinning enforcement. Once a model is assigned, all LLM calls in that workflow route to the same backend. This preserves context, eliminates format drift, and ensures consistent reasoning style.

Delayed feedback integration. After the workflow completes, the router receives a terminal reward (accuracy score, user satisfaction, task success) and a total latency measurement. It uses this signal to update its policy via contextual bandit learning (Thompson Sampling or Upper Confidence Bound).

The key insight is that task-level feedback is the only reliable signal for multi-step workflows. Intermediate call quality does not matter if the final outcome fails.

Implementation Details

TRACE-ROUTER uses a contextual bandit with Thompson Sampling. For each task, it maintains a posterior distribution over model performance given task features. At admission, it samples from this distribution and selects the model with the highest sampled reward.

class TraceRouter:
    def __init__(self, models, feature_dim):
        self.models = models
        # Bayesian linear regression per model
        self.posteriors = {
            m: BayesianLinearRegression(feature_dim)
            for m in models
        }
    
    def route(self, task_features):
        # Sample expected reward for each model
        samples = {
            m: self.posteriors[m].sample(task_features)
            for m in self.models
        }
        # Select model with highest sampled reward
        selected = max(samples, key=samples.get)
        return selected
    
    def update(self, task_features, model, reward, latency):
        # Composite reward: accuracy + latency penalty
        composite = reward - (latency / self.latency_budget)
        # Update posterior for selected model
        self.posteriors[model].update(task_features, composite)

The reward function combines accuracy and latency. You can weight these based on your SLA requirements. For latency-sensitive applications, increase the latency penalty. For accuracy-critical tasks, reduce it.

Feature engineering matters. Task features should capture:

  • Query length and complexity (token count, syntactic depth)
  • Domain or category (code generation, summarization, reasoning)
  • Expected workflow depth (number of anticipated LLM calls)
  • Historical success rate for similar tasks

The bandit learns which models perform well for which task types. Over time, it adapts to workload shifts without explicit retraining.

Observability and Debugging

Multi-step workflows require different observability primitives than single-shot inference. You need to track:

Task-level traces. Each workflow gets a unique task ID. All LLM calls, tool invocations, and state transitions are tagged with this ID. When a workflow fails, you can reconstruct the entire execution path.

Routing decision logs. Record the task features, selected model, and sampled reward at admission time. This lets you audit why the router made a specific decision.

Reward attribution. After each task completes, log the terminal reward, total latency, and per-call latency breakdown. This shows whether latency bottlenecks occur in specific steps or across the entire workflow.

Posterior drift monitoring. Track how the bandit’s posterior distributions change over time. Sudden shifts indicate workload changes or model performance degradation.

A useful debugging pattern is to compare task-level routing decisions against a baseline per-call router. Run both in shadow mode and measure how often they disagree. High disagreement rates indicate that task-level context is influencing routing in ways per-call routers cannot capture.

Failure Modes and Mitigation

Cold start for new task types. If the router encounters a task type it has never seen, it has no historical data to guide selection. Mitigation: use a prior based on aggregate model performance, or route new task types to a strong default model until you collect feedback.

Reward sparsity. Some workflows take minutes or hours to complete. Delayed feedback slows learning. Mitigation: use intermediate checkpoints (partial task completion) as auxiliary rewards, or batch updates to amortize learning overhead.

Model availability changes. If a model becomes unavailable mid-workflow, pinning breaks. Mitigation: implement fallback routing with a secondary model, and mark the task for policy update with a penalty.

Adversarial workloads. If task features are user-controlled, adversaries can game the router by submitting tasks with misleading features. Mitigation: validate features server-side, or use a separate classifier to detect anomalous tasks.

Comparison: Task-Level vs. Per-Call Routing

DimensionPer-Call RoutingTask-Level Routing
Routing granularityEvery LLM callOnce per workflow
Context consistencyNo guaranteesEnforced by pinning
Credit assignmentImpossible for workflowsDirect via terminal reward
Latency coordinationIndependent per callCoordinated across workflow
Learning signalPer-call accuracyTask-level outcome
OverheadHigh (route every call)Low (route once at admission)
Adaptation speedFast for single-shotSlower (delayed feedback)

Task-level routing trades faster adaptation for better workflow coherence. If your workload is dominated by single-shot queries, per-call routing is fine. If you run multi-step agents, task-level routing prevents the compounding errors that occur when you switch models mid-execution.

Deployment Considerations

Stateful orchestration required. You need a workflow engine that tracks task IDs and enforces pinning. Temporal, Prefect, or Langraph can handle this. Stateless API gateways cannot.

Latency budget allocation. Set a task-level SLA, not a per-call SLA. The router needs to know the total budget to make informed trade-offs. If your SLA is 10 seconds and the workflow has five steps, the router can allocate latency dynamically (e.g., spend more on planning, less on execution).

Model inventory management. Task-level routing assumes you have multiple models available. If you only have one model, routing is trivial. If you have dozens, feature engineering becomes critical to distinguish which models excel at which tasks.

Cost tracking. Pin task IDs to billing records. This lets you analyze cost per task, not just cost per call. You can identify which task types are expensive and adjust routing policies accordingly.

Benchmark Results

On tau2-Bench (a multi-step reasoning benchmark), TRACE-ROUTER outperforms latency-matched interpolation between individual models by 7 to 8 accuracy points. This means that for the same latency budget, task-consistent routing produces more accurate outcomes than per-call routing.

On Terminal-Bench (a long-horizon agentic benchmark), TRACE-ROUTER achieves 7.1 higher accuracy points than the strongest single model baseline while using 36% less latency. This demonstrates that adaptive routing can both improve quality and reduce cost when it accounts for task-level feedback.

The Pareto frontier analysis shows that TRACE-ROUTER consistently dominates per-call routers across the accuracy-latency trade-off space. For any given latency budget, task-consistent routing produces higher accuracy. For any given accuracy target, it uses less latency.

Technical Verdict

Use task-consistent routing when:

  • Your workload consists of multi-step agent workflows with more than three LLM calls per task.
  • You have a diverse model inventory with different cost-quality-latency profiles.
  • You can tolerate delayed feedback (minutes to hours) for policy updates.
  • You need to enforce context consistency across a workflow.
  • You have observability infrastructure to track task-level traces and terminal rewards.

Avoid task-consistent routing when:

  • Your workload is dominated by single-shot queries with no follow-up calls.
  • You need sub-second adaptation to workload changes (per-call routers adapt faster).
  • You cannot implement stateful orchestration or task pinning.
  • Your models have identical performance profiles (routing becomes trivial).
  • You lack the instrumentation to measure task-level outcomes.

Task-consistent routing is not a replacement for per-call routing. It is a specialized tool for long-horizon workflows where context continuity and credit assignment matter more than instant adaptation. If you are building production agent systems that chain multiple LLM calls, this is the routing layer you need.

Tags

agentic-ai orchestration infrastructure

Primary Source

arxiv.org