# How Request Tracing Works with Thread-Aware Observability in AxonHub

> Learn how AxonHub uses request tracing and thread-aware observability to correlate data across threads and services. Understand trace and thread ID propagation for enhanced debugging.

- Repository: [Loop/axonhub](https://github.com/looplj/axonhub)
- Tags: internals
- Published: 2026-03-06

---

**AxonHub implements request tracing with thread-aware observability by extracting trace and thread identifiers from HTTP headers, storing them in the request context via `middleware.WithTrace` and `middleware.Thread`, and propagating this data to load-balancers, loggers, and metrics collectors to enable per-thread routing and correlation.**

Request tracing with thread-aware observability allows distributed systems to correlate related requests across service boundaries while maintaining affinity to specific execution threads. In the `looplj/axonhub` repository, this capability is implemented through a coordinated middleware pipeline that extracts identifiers from incoming HTTP requests and injects them into Go contexts for downstream consumption.

## Core Components of Thread-Aware Request Tracing

The system relies on three tightly-coupled components that share configuration and context state.

### middleware.WithTrace: The Entry Point

Located in [`internal/server/middleware/trace.go`](https://github.com/looplj/axonhub/blob/main/internal/server/middleware/trace.go), the `WithTrace` middleware reads the trace identifier from incoming HTTP requests. By default, it looks for the `AH-Trace-Id` header, but this is configurable via `tracing.Config`. The middleware creates or fetches a `Trace` entity from the database and injects it into the request's `context` using `contexts.WithTrace`.

### middleware.Thread: Thread Context Injection

While not explicitly detailed in the trace middleware, the analogous `middleware.Thread` (implied by the analysis and analogous patterns in [`internal/server/middleware/thread.go`](https://github.com/looplj/axonhub/blob/main/internal/server/middleware/thread.go)) extracts a **thread identifier** from a configurable header (`tracing.Config.ThreadHeader`). It stores a `Thread` object in the same request context, making it available for downstream middleware.

### internal/tracing.Config: Unified Configuration

The [`internal/tracing/config.go`](https://github.com/looplj/axonhub/blob/main/internal/tracing/config.go) file defines the centralized configuration that both middlewares consume. It specifies which header names to inspect (`TraceHeader`, `ThreadHeader`, `ExtraTraceHeaders`), which upstream providers need special handling (Claude Code, Codex), and body extraction rules (`ExtraTraceBodyFields`). This ensures both trace and thread identifiers are extracted consistently across the application.

## Request Flow: From Headers to Context

When a request enters the AxonHub server, it passes through a five-stage extraction and enrichment pipeline.

### Step 1: Header and Body Extraction

The `WithTrace` middleware first checks the primary trace header (`AH-Trace-Id` or the custom one defined in `tracing.Config`). If the header is missing, it falls back to any `ExtraTraceHeaders` configured. If still empty and `ExtraTraceBodyFields` are configured, the middleware reads the request body once, searches for the specified JSON fields using **gjson**, and restores the body for downstream handlers.

### Step 2: Provider-Specific Trace Resolution

When Claude Code or Codex requests are detected, the middleware runs special helper functions (`tryExtractTraceIDFromClaudeCodeRequest`, `tryExtractTraceIDFromCodexRequest`). These extract trace IDs from provider-specific payloads, ensuring compatibility with external AI coding assistants while maintaining observability.

### Step 3: Thread Association and Context Enrichment

After a trace ID is found, the middleware queries `contexts.GetThread` for a thread object that may have been placed there by the thread middleware. If present, the thread's numeric ID is attached to the trace record via `traceService.GetOrCreateTrace`. The resulting `Trace` entity and the raw trace ID are stored back into the request's context using `contexts.WithTrace` and `shared.WithSessionID`.

## Thread-Aware Observability in Action

Once the context contains both trace and thread identifiers, downstream components leverage this data for intelligent routing and monitoring.

### Load-Balancing with Trace-Aware Strategy

The **Trace-Aware strategy** in [`internal/server/orchestrator/trace_aware_strategy.go`](https://github.com/looplj/axonhub/blob/main/internal/server/orchestrator/trace_aware_strategy.go) queries the stored trace to find the "last successful channel" for that trace (and thus for that thread). It then **boosts** the score of that channel by `1000`, making it the preferred route for the remainder of the thread's life. This ensures that related requests within the same thread hit the same backend channel, improving cache locality and reducing connection overhead.

### Correlated Logging and Metrics

In [`internal/tracing/log.go`](https://github.com/looplj/axonhub/blob/main/internal/tracing/log.go), the logging system adds `trace_id` (and `request_id`) fields to every log entry when they exist in the context, allowing logs to be correlated per thread. Similarly, the metrics middleware in [`internal/server/middleware/metrics.go`](https://github.com/looplj/axonhub/blob/main/internal/server/middleware/metrics.go) reads the trace and thread IDs to tag Prometheus metrics, enabling dashboards that show per-thread request latency, error rates, and channel utilization.

## Implementation Examples

### Adding Middleware to the Gin Router

```go
// internal/server/server.go (excerpt)
router := gin.New()

// Load tracing configuration (from config file / env)
tracingCfg := tracing.LoadConfig()

// Attach thread-aware middleware first (sets thread in context)
router.Use(middleware.Thread(tracingCfg))

// Attach trace middleware (reads trace, links thread, stores both)
router.Use(middleware.WithTrace(tracingCfg, biz.NewTraceService(db)))

router.Use(middleware.Logging())   // now logs include trace_id & thread_id
router.Use(middleware.Metrics())   // metrics are tagged with trace_id & thread_id

```

### Retrieving Trace and Thread Information in Handlers

```go
func MyHandler(c *gin.Context) {
    // Grab the trace ID (string) – useful for downstream provider calls
    traceID, _ := tracing.GetTraceID(c.Request.Context())

    // Grab the full Trace entity (contains DB ID, ProjectID, ThreadID)
    trace, _ := contexts.GetTrace(c.Request.Context())

    // Grab the Thread ID (numeric) if the request is part of a thread
    threadID := ""
    if trace != nil && trace.ThreadID != nil {
        threadID = fmt.Sprintf("%d", *trace.ThreadID)
    }

    log.Info(c.Request.Context(),
        "handling request",
        log.String("trace_id", traceID),
        log.String("thread_id", threadID),
    )

    // … business logic …
}

```

### Trace-Aware Load-Balancer Strategy

```go
// internal/server/orchestrator/trace_aware_strategy.go (simplified)
func (s *TraceAwareStrategy) RankChannels(ctx context.Context, candidates []Channel) []Channel {
    trace, ok := contexts.GetTrace(ctx)
    if !ok || trace.LastSuccessfulChannelID == nil {
        return candidates // no trace info → fall back to other strategies
    }

    // Boost the channel that was last successful for this trace/thread
    for i, ch := range candidates {
        if ch.ID == *trace.LastSuccessfulChannelID {
            candidates[i].Score += 1000 // massive boost
            log.Debug(ctx, "Trace-aware boost applied", log.Any("channel", ch.ID))
        }
    }
    // sort by Score descending …
    return candidates
}

```

## Key Source Files

| File | Role |
|------|------|
| [[`internal/server/middleware/trace.go`](https://github.com/looplj/axonhub/blob/main/internal/server/middleware/trace.go)](https://github.com/looplj/axonhub/blob/unstable/internal/server/middleware/trace.go) | Core tracing middleware – extracts trace IDs, ties them to projects and threads, stores them in context. |
| [[`internal/server/middleware/thread.go`](https://github.com/looplj/axonhub/blob/main/internal/server/middleware/thread.go)](https://github.com/looplj/axonhub/blob/unstable/internal/server/middleware/thread.go) | Extracts thread identifiers from headers and stores `Thread` objects in context. |
| [[`internal/tracing/config.go`](https://github.com/looplj/axonhub/blob/main/internal/tracing/config.go)](https://github.com/looplj/axonhub/blob/unstable/internal/tracing/config.go) | Configuration struct (`TraceHeader`, `ThreadHeader`, `ExtraTraceHeaders`, etc.) shared by both middlewares. |
| [[`internal/contexts/context.go`](https://github.com/looplj/axonhub/blob/main/internal/contexts/context.go)](https://github.com/looplj/axonhub/blob/unstable/internal/contexts/context.go) | Helpers `WithTrace`, `GetTrace`, `WithThread`, `GetThread` that embed trace/thread data into Go contexts. |
| [[`internal/server/orchestrator/trace_aware_strategy.go`](https://github.com/looplj/axonhub/blob/main/internal/server/orchestrator/trace_aware_strategy.go)](https://github.com/looplj/axonhub/blob/unstable/internal/server/orchestrator/trace_aware_strategy.go) | Load-balancer strategy that uses stored trace data to bias channel selection. |
| [[`internal/tracing/log.go`](https://github.com/looplj/axonhub/blob/main/internal/tracing/log.go)](https://github.com/looplj/axonhub/blob/unstable/internal/tracing/log.go) | Adds `trace_id` and `request_id` fields to log entries when present in context. |
| [[`internal/server/middleware/metrics.go`](https://github.com/looplj/axonhub/blob/main/internal/server/middleware/metrics.go)](https://github.com/looplj/axonhub/blob/unstable/internal/server/middleware/metrics.go) | Tags Prometheus metrics with trace and thread identifiers. |

## Summary

- **Three core components**—`middleware.WithTrace`, `middleware.Thread`, and `internal/tracing.Config`—work together to extract and link trace and thread identifiers.
- **Multi-stage extraction** supports headers, request bodies, and provider-specific payloads (Claude Code, Codex) to ensure trace continuity.
- **Context propagation** stores `Trace` and `Thread` entities in the Go context, making identifiers available to load-balancers, loggers, and metrics collectors without manual plumbing.
- **Thread-aware routing** uses the `TraceAwareStrategy` to boost the score of previously successful channels by 1000 points, ensuring related requests hit the same backend.
- **Observability integration** automatically injects `trace_id` and `thread_id` into logs and Prometheus metrics, enabling per-thread latency and error rate analysis.

## Frequently Asked Questions

### What is thread-aware observability?

Thread-aware observability is the practice of tracking a sequence of related requests—grouped under a single thread identifier—across distributed services. In AxonHub, this allows the system to correlate logs, metrics, and routing decisions for all requests belonging to the same conversation or session, rather than treating each request in isolation.

### How does AxonHub extract trace IDs from request bodies?

When the `AH-Trace-Id` header and configured `ExtraTraceHeaders` are absent, the `WithTrace` middleware in [`internal/server/middleware/trace.go`](https://github.com/looplj/axonhub/blob/main/internal/server/middleware/trace.go) checks `ExtraTraceBodyFields` from the tracing configuration. It reads the request body once, uses **gjson** to extract the specified JSON fields, and then restores the body so downstream handlers can read it normally.

### Can I customize the header names for tracing and threading?

Yes. The [`internal/tracing/config.go`](https://github.com/looplj/axonhub/blob/main/internal/tracing/config.go) file defines a centralized configuration where you can set custom header names via `TraceHeader` and `ThreadHeader`. You can also specify `ExtraTraceHeaders` for fallback extraction and `ExtraTraceBodyFields` for JSON body parsing, allowing the middleware to adapt to various client conventions without code changes.

### How does thread-aware routing improve performance?

The `TraceAwareStrategy` in [`internal/server/orchestrator/trace_aware_strategy.go`](https://github.com/looplj/axonhub/blob/main/internal/server/orchestrator/trace_aware_strategy.go) improves performance by adding a **1000-point boost** to the channel that previously succeeded for a given trace/thread. This ensures that subsequent requests in the same thread are routed to the same backend channel, improving cache locality, reducing connection overhead, and minimizing latency for conversational or stateful workloads.