feat: split classify and LLM timings in request_trace

Add llm.WithTimings context to capture classifier duration inside
router.pick. After Chat/ChatStream returns, iter_N_classify_ms and
iter_N_llm_ms are logged separately — llm_ms is net model time only
(total minus classify).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: Alexey Panfilov

internal/agent/agent.go

@@ -164,9 +164,16 @@ func (a *Agent) Process(ctx context.Context, chatID int64, userMsg llm.Message,
 			sysPrompt += "\n\n" + crossSessionCtx
 		}
 
-		endLLM := tr.begin(fmt.Sprintf("iter%d_llm", i))
-		resp, err := a.router.Chat(ctx, history, sysPrompt, tools)
-		endLLM()
+		iterCtx, timing := llm.WithTimings(ctx)
+		llmStart := time.Now()
+		resp, err := a.router.Chat(iterCtx, history, sysPrompt, tools)
+		llmTotal := time.Since(llmStart)
+		if timing.ClassifyRan {
+			tr.addSpan(fmt.Sprintf("iter%d_classify", i), timing.ClassifyDur)
+			tr.addSpan(fmt.Sprintf("iter%d_llm", i), llmTotal-timing.ClassifyDur)
+		} else {
+			tr.addSpan(fmt.Sprintf("iter%d_llm", i), llmTotal)
+		}
 		if err != nil {
 			tr.log(a.logger, chatID, "llm_error")
 			return "", fmt.Errorf("llm: %w", err)
@@ -256,10 +263,17 @@ func (a *Agent) ProcessStream(ctx context.Context, chatID int64, userMsg llm.Mes
 			sysPrompt += "\n\n" + crossSessionCtx
 		}
 
-		endLLM := tr.begin(fmt.Sprintf("iter%d_llm", i))
-		ch, err := a.router.ChatStream(ctx, history, sysPrompt, tools)
+		iterCtx, timing := llm.WithTimings(ctx)
+		llmStart := time.Now()
+		ch, err := a.router.ChatStream(iterCtx, history, sysPrompt, tools)
 		if err != nil {
-			endLLM()
+			llmTotal := time.Since(llmStart)
+			if timing.ClassifyRan {
+				tr.addSpan(fmt.Sprintf("iter%d_classify", i), timing.ClassifyDur)
+				tr.addSpan(fmt.Sprintf("iter%d_llm", i), llmTotal-timing.ClassifyDur)
+			} else {
+				tr.addSpan(fmt.Sprintf("iter%d_llm", i), llmTotal)
+			}
 			tr.log(a.logger, chatID, "llm_error")
 			return "", fmt.Errorf("llm: %w", err)
 		}
@@ -270,7 +284,13 @@ func (a *Agent) ProcessStream(ctx context.Context, chatID int64, userMsg llm.Mes
 
 		for chunk := range ch {
 			if chunk.Err != nil {
-				endLLM()
+				llmTotal := time.Since(llmStart)
+				if timing.ClassifyRan {
+					tr.addSpan(fmt.Sprintf("iter%d_classify", i), timing.ClassifyDur)
+					tr.addSpan(fmt.Sprintf("iter%d_llm", i), llmTotal-timing.ClassifyDur)
+				} else {
+					tr.addSpan(fmt.Sprintf("iter%d_llm", i), llmTotal)
+				}
 				tr.log(a.logger, chatID, "stream_error")
 				return "", fmt.Errorf("llm stream: %w", chunk.Err)
 			}
@@ -287,7 +307,13 @@ func (a *Agent) ProcessStream(ctx context.Context, chatID int64, userMsg llm.Mes
 				}
 			}
 		}
-		endLLM()
+		llmTotal := time.Since(llmStart)
+		if timing.ClassifyRan {
+			tr.addSpan(fmt.Sprintf("iter%d_classify", i), timing.ClassifyDur)
+			tr.addSpan(fmt.Sprintf("iter%d_llm", i), llmTotal-timing.ClassifyDur)
+		} else {
+			tr.addSpan(fmt.Sprintf("iter%d_llm", i), llmTotal)
+		}
 
 		if len(toolCalls) == 0 {
 			a.store.AddMessage(chatID, llm.Message{Role: "assistant", Content: accumulated})

internal/agent/trace.go

@@ -36,6 +36,11 @@ func (t *requestTrace) begin(name string) func() {
 	}
 }
 
+// addSpan records a pre-measured span (e.g. extracted from context timings).
+func (t *requestTrace) addSpan(name string, dur time.Duration) {
+	t.spans = append(t.spans, traceSpan{name: name, dur: dur})
+}
+
 // log emits a single structured log line with all span durations in milliseconds.
 func (t *requestTrace) log(logger *slog.Logger, chatID int64, status string) {
 	args := make([]any, 0, 3+len(t.spans)*2)

internal/llm/router.go

@@ -537,7 +537,13 @@ func (r *Router) classify(ctx context.Context, text string) int {
 		prompt = defaultClassifierPrompt
 	}
 	msgs := []Message{{Role: "user", Content: classifierText}}
+	classifyStart := time.Now()
 	resp, err := provider.Chat(classifierCtx, msgs, prompt, nil)
+	classifyDur := time.Since(classifyStart)
+	if t := timingsFrom(ctx); t != nil {
+		t.ClassifyDur = classifyDur
+		t.ClassifyRan = true
+	}
 	if err != nil {
 		r.logger.Warn("classifier error, falling back to primary", "classifier", classifierKey, "err", err)
 		return 2

internal/llm/timing.go

@@ -0,0 +1,27 @@
+package llm
+
+import (
+	"context"
+	"time"
+)
+
+// RequestTimings collects per-phase durations for a single router.Chat call.
+// Attach to a context via WithTimings; the router populates it during the call.
+type RequestTimings struct {
+	ClassifyDur time.Duration
+	ClassifyRan bool
+}
+
+type timingCtxKey struct{}
+
+// WithTimings returns a child context carrying a *RequestTimings that the router
+// will populate. The caller reads timings after Chat/ChatStream returns.
+func WithTimings(ctx context.Context) (context.Context, *RequestTimings) {
+	t := &RequestTimings{}
+	return context.WithValue(ctx, timingCtxKey{}, t), t
+}
+
+func timingsFrom(ctx context.Context) *RequestTimings {
+	t, _ := ctx.Value(timingCtxKey{}).(*RequestTimings)
+	return t
+}