vecna/pkg/server/trace.go

package server

import (
	"context"
	"time"
)

type contextKey int

const traceKey contextKey = iota

// RequestTrace holds per-request timing data populated by handlers and middleware.
type RequestTrace struct {
	Start             time.Time
	ForwardDuration   time.Duration
	TranslateDuration time.Duration
	ForwardTarget     string
	ForwardURL        string
	ForwardModel      string
	AdapterType       string
	PromptTokens      int
	TotalTokens       int
}

// WithTrace injects a new *RequestTrace into ctx.
func WithTrace(ctx context.Context) context.Context {
	return context.WithValue(ctx, traceKey, &RequestTrace{Start: time.Now()})
}

// TraceFromContext retrieves the *RequestTrace from ctx.
// Returns a zero-value trace (non-nil) if none was set.
func TraceFromContext(ctx context.Context) *RequestTrace {
	if t, ok := ctx.Value(traceKey).(*RequestTrace); ok && t != nil {
		return t
	}
	return &RequestTrace{Start: time.Now()}
}