server: generate asks the runner for pass-one metrics too

docs/maxusai/adr/0010-transition-flow-metrics-reconstruction.md

@@ -3,7 +3,12 @@
 - **Status:** accepted, implemented + live-validated 2026-08-08 on fork `main`
   lineage. Amends the transition-flow metrics behaviour of
   [ADR 0004](0004-routes-layer-think-format-double-request.md); 0004's marker
-  flow is unchanged.
+  flow is unchanged. **2026-09-04 (v0.33.3 fold):** both handlers now prefer
+  the runner's own reported pass-one metrics (`reportedPassMetrics()`, fed by
+  upstream's `IncludeIntermediateMetrics`) at the transition site; the textual
+  reconstruction described below is the fallback for the runners and paths
+  that report nothing, and `transitionPromptDelta` still refines a
+  reconstructed pass only.
 - **Date:** 2026-08-08
 - **Deciders:** MaxusAI fork maintainers
 

server/routes.go

@@ -711,6 +711,10 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 		transitionPromptDelta := -1
 		var continueViaMarker bool
 		var contentStarted bool
+		// firstPassMetrics is the runner's own report for pass one, kept from
+		// the last chunk that carried metrics (upstream's
+		// IncludeIntermediateMetrics makes every chunk carry them).
+		var firstPassMetrics api.Metrics
 		// contextPrompt is what the final Context field is tokenized against;
 		// it stays the original prompt for the textual continuation and is
 		// retargeted to the re-rendered prompt in the transition flow.
@@ -739,6 +743,7 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 					passOpts = &o
 				}
 			}
+			includeIntermediateMetrics := req.Format != nil && currentFormat == nil
 
 			ctx, cancel := context.WithCancel(c.Request.Context())
 
@@ -749,33 +754,46 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 			var firstChunkAt time.Time
 
 			err := r.Completion(ctx, llm.CompletionRequest{
-				Prompt:          prompt,
-				Media:           media,
-				Format:          currentFormat,
-				Options:         passOpts,
-				Shift:           req.Shift == nil || *req.Shift,
-				Truncate:        truncate,
-				Logprobs:        req.Logprobs,
-				TopLogprobs:     req.TopLogprobs,
-				PreservedTokens: preservedTokensForCompletion(builtinParser),
-				LeadingBOS:      leadingBOS,
+				Prompt:                     prompt,
+				Media:                      media,
+				Format:                     currentFormat,
+				Options:                    passOpts,
+				Shift:                      req.Shift == nil || *req.Shift,
+				Truncate:                   truncate,
+				Logprobs:                   req.Logprobs,
+				TopLogprobs:                req.TopLogprobs,
+				PreservedTokens:            preservedTokensForCompletion(builtinParser),
+				LeadingBOS:                 leadingBOS,
+				IncludeIntermediateMetrics: includeIntermediateMetrics,
 			}, func(cr llm.CompletionResponse) {
 				if firstChunkAt.IsZero() {
 					firstChunkAt = time.Now()
 				}
+				metrics := api.Metrics{
+					PromptEvalCount:       cr.PromptEvalCount,
+					PromptEvalCachedCount: cr.PromptEvalCachedCount,
+					PromptEvalDuration:    cr.PromptEvalDuration,
+					EvalCount:             cr.EvalCount,
+					EvalDuration:          cr.EvalDuration,
+				}
+				if includeIntermediateMetrics {
+					// A deferring pass asks the runner for per-chunk timings,
+					// so keep the latest as pass one's own report and blank the
+					// intermediate copies: a mid-stream counter must not leak
+					// into every streamed API chunk (ADR 0010's reason for
+					// rejecting a running-counter protocol in the first place).
+					firstPassMetrics = metrics
+					if !cr.Done {
+						metrics = api.Metrics{}
+					}
+				}
 				res := api.GenerateResponse{
 					Model:     req.Model,
 					CreatedAt: time.Now().UTC(),
 					Response:  cr.Content,
 					Done:      cr.Done,
-					Metrics: api.Metrics{
-						PromptEvalCount:       cr.PromptEvalCount,
-						PromptEvalCachedCount: cr.PromptEvalCachedCount,
-						PromptEvalDuration:    cr.PromptEvalDuration,
-						EvalCount:             cr.EvalCount,
-						EvalDuration:          cr.EvalDuration,
-					},
-					Logprobs: toAPILogprobs(cr.Logprobs),
+					Metrics:   metrics,
+					Logprobs:  toAPILogprobs(cr.Logprobs),
 				}
 
 				if _, err := sb.WriteString(cr.Content); err != nil {
@@ -937,10 +955,21 @@ func (s *Server) GenerateHandler(c *gin.Context) {
 				} else {
 					// The transition cancelled pass one mid-stream before its
 					// final metrics arrived; reconstruct them so the final
-					// response still counts every pass-one token (R6).
+					// response still counts every pass-one token (R6). The
+					// prompt count is refined once pass two reports its prefill.
 					if pass1 == nil {
-						pass1 = transitionPassMetrics(c.Request.Context(), r, contextPrompt, sb.String(), passStart, firstChunkAt)
-						reconstructed = pass1 != nil
+						// Upstream's per-chunk metrics carry pass one's own
+						// report even though the cancel discarded its final
+						// chunk, and that report is the runner's cache-inclusive
+						// prefill — image-embedding tokens included — which is
+						// exactly what ADR 0010's subtraction exists to recover.
+						// Prefer it, and keep the textual reconstruction for the
+						// runners and paths that report nothing (ADR 0010 is
+						// best-effort).
+						if pass1 = reportedPassMetrics(firstPassMetrics); pass1 == nil {
+							pass1 = transitionPassMetrics(c.Request.Context(), r, contextPrompt, sb.String(), passStart, firstChunkAt)
+							reconstructed = pass1 != nil
+						}
 					}
 
 					msg := api.Message{

server/routes_generate_test.go

@@ -4469,6 +4469,181 @@ func TestGenerateThinkFormatTransitionMetrics(t *testing.T) {
 	}
 }
 
+// TestGenerateThinkFormatTransitionMetricsReportedPassOne is /api/generate's
+// twin of TestChatThinkFormatTransitionMetricsReportedPassOne: the same
+// think+format transition flow, except pass one actually reports. Upstream's
+// IncludeIntermediateMetrics makes every chunk of a deferring pass carry the
+// runner's own metrics, so a pass cancelled at the thinking-to-content
+// transition still has a report of its own — and that report is the runner's
+// cache-inclusive prefill, image-embedding tokens included.
+//
+// The assertion that matters: the folded prompt count carries the image
+// surplus. Nothing in the textual reconstruction can see those tokens (ADR
+// 0010 exists to recover them by subtraction), so a fold that fell back to
+// tokenizing the prompt would understate this by imageTokens and the test
+// would catch it. TestGenerateThinkFormatTransitionMetrics pins that fallback,
+// where the runner reports nothing.
+//
+// Streamed, because asking for per-chunk metrics has a second half: the
+// intermediate copies must be blanked, so no chunk before done exposes a
+// mid-stream counter.
+func TestGenerateThinkFormatTransitionMetricsReportedPassOne(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	mock := &mockRunner{}
+	s := setupTransitionThinkingModel(t, mock, "test-transition-thinking-vision-generate")
+
+	format := json.RawMessage(`{"type":"object","properties":{"answer":{"type":"string"}}}`)
+	passOneRaw := `<think> Reasoning deeply about the answer. </think> {"answer`
+
+	// One nemotron3 image's worth of soft tokens, present in the runner's
+	// prefill counts and invisible to the whitespace-field mock Tokenize.
+	const imageTokens = 2042
+
+	var (
+		requestsMu sync.Mutex
+		requests   []llm.CompletionRequest
+	)
+
+	var passOnePrompt int
+	mock.CompletionFn = func(ctx context.Context, r llm.CompletionRequest, fn func(r llm.CompletionResponse)) error {
+		requestsMu.Lock()
+		requests = append(requests, r)
+		callNum := len(requests)
+		requestsMu.Unlock()
+
+		switch callNum {
+		case 1:
+			passOnePrompt = len(strings.Fields(r.Prompt)) + imageTokens
+			if !r.IncludeIntermediateMetrics {
+				t.Error("expected the deferring pass to request per-token metrics")
+			}
+			// A streamed, non-terminal chunk carrying pass one's metrics:
+			// what upstream's timings_per_token and the MLX pipeline's
+			// per-chunk reporting deliver, and what the cancel would
+			// otherwise discard.
+			fn(llm.CompletionResponse{
+				Content:               passOneRaw,
+				Done:                  false,
+				PromptEvalCount:       passOnePrompt,
+				PromptEvalCachedCount: testIntPtr(7),
+				PromptEvalDuration:    11,
+				EvalCount:             12,
+				EvalDuration:          13,
+			})
+			// the transition cancels this pass: the final chunk that would
+			// carry the pass-one metrics never arrives
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			case <-time.After(time.Second):
+				t.Errorf("timeout waiting for structured outputs cancellation")
+				return nil
+			}
+		case 2:
+			if r.IncludeIntermediateMetrics {
+				t.Error("expected the constrained pass to use terminal metrics")
+			}
+			fn(llm.CompletionResponse{
+				Content:    `{"answer":"42"}`,
+				Done:       true,
+				DoneReason: llm.DoneReasonStop,
+				// The continuation re-counts the reasoning as prompt; the
+				// fold must not report this.
+				PromptEvalCount:       len(strings.Fields(r.Prompt)) + imageTokens,
+				PromptEvalCachedCount: testIntPtr(19),
+				PromptEvalDuration:    30,
+				EvalCount:             5,
+				EvalDuration:          50,
+			})
+			return nil
+		default:
+			t.Errorf("unexpected number of completion calls: %d", callNum)
+			return nil
+		}
+	}
+
+	think := true
+	w := createRequest(t, s.GenerateHandler, api.GenerateRequest{
+		Model:  "test-transition-thinking-vision-generate",
+		Prompt: "Why is the sky blue?",
+		Think:  &api.ThinkValue{Value: think},
+		Format: format,
+	})
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("expected status 200, got %d: %s", w.Code, w.Body.String())
+	}
+	if len(requests) != 2 {
+		t.Fatalf("expected two completion calls, got %d", len(requests))
+	}
+	if requests[0].Format != nil {
+		t.Errorf("expected first completion format to be nil, got %q", requests[0].Format)
+	}
+	if !bytes.Equal([]byte(format), []byte(requests[1].Format)) {
+		t.Errorf("expected second completion format to match original format")
+	}
+
+	var thinking, response strings.Builder
+	var final api.GenerateResponse
+	dec := json.NewDecoder(w.Body)
+	for {
+		var chunk api.GenerateResponse
+		if err := dec.Decode(&chunk); err == io.EOF {
+			break
+		} else if err != nil {
+			t.Fatal(err)
+		}
+		thinking.WriteString(chunk.Thinking)
+		response.WriteString(chunk.Response)
+		if chunk.Done {
+			final = chunk
+			continue
+		}
+		// The per-chunk metrics are pass one's own report, not a running
+		// total for the client: they must not leak into a streamed chunk
+		// (the reason ADR 0010 rejected a running-counter protocol).
+		if chunk.PromptEvalCount != 0 || chunk.PromptEvalCachedCount != nil ||
+			chunk.PromptEvalDuration != 0 || chunk.EvalCount != 0 || chunk.EvalDuration != 0 {
+			t.Errorf("non-terminal chunk unexpectedly exposed metrics: %+v", chunk.Metrics)
+		}
+	}
+
+	if got := thinking.String(); got != "Reasoning deeply about the answer. " {
+		t.Errorf("thinking = %q, want %q", got, "Reasoning deeply about the answer. ")
+	}
+	if got := response.String(); got != `{"answer":"42"}` {
+		t.Errorf("response = %q, want %q", got, `{"answer":"42"}`)
+	}
+	if !final.Done || final.DoneReason != "stop" {
+		t.Fatalf("expected a final stop chunk, got done=%v reason=%q", final.Done, final.DoneReason)
+	}
+
+	// R6: the request's own prompt cost, image tokens included.
+	if final.PromptEvalCount != passOnePrompt {
+		t.Errorf("prompt eval count = %d, want %d (pass one's own cache-inclusive prefill, image tokens included)", final.PromptEvalCount, passOnePrompt)
+	}
+	if final.PromptEvalCount <= len(strings.Fields(requests[0].Prompt)) {
+		t.Errorf("prompt eval count %d does not carry the image surplus", final.PromptEvalCount)
+	}
+	// Every generated token counted once across both passes.
+	if want := 12 + 5; final.EvalCount != want {
+		t.Errorf("eval count = %d, want %d (both passes)", final.EvalCount, want)
+	}
+	// ADR 0004's summing: each pass's prefill is prefill, each pass's decode
+	// is decode (see wantMetrics in TestChatWithPromptEndingInThinkTag).
+	if want := time.Duration(11 + 30); final.PromptEvalDuration != want {
+		t.Errorf("prompt eval duration = %d, want %d", final.PromptEvalDuration, want)
+	}
+	if want := time.Duration(13 + 50); final.EvalDuration != want {
+		t.Errorf("eval duration = %d, want %d", final.EvalDuration, want)
+	}
+	// prompt_eval_cached_count is deliberately not asserted: unlike chat's,
+	// generate's fold does not re-attribute it to pass one, so the response
+	// still carries the continuation's. Changing that is a fold change, not
+	// part of wiring the flag.
+}
+
 // TestChatThinkFormatTransitionMetrics is the /api/chat side of R6 on the
 // transition flow: pass one is cancelled at the thinking→content transition
 // with no runner-reported metrics, and the final response must still total