agent-exec transport bf732506478b02c63a7b4fe251f0ef0f/ba449849829e

internal/runtime/material_final_reserve_test.go

@@ -1,6 +1,9 @@
 package runtime
 
-import "testing"
+import (
+	core "agent-system-core"
+	"testing"
+)
 
 func TestMaterialFinalReserveDecision(t *testing.T) {
 	tests := []struct {
@@ -36,3 +39,37 @@ func TestMaterialFinalReserveDecision(t *testing.T) {
 		})
 	}
 }
+
+// TestSubRunMaterialDecisionIncludesLastCompletion locks the subRun-level
+// invariant the spec demands: materialDecision must sum LastPromptTokens and
+// LastCompletionTokens when estimating the next request size, so a continuation
+// after a truncated assistant turn is still caught by the reserve. Reverting
+// the summation to LastPromptTokens only would let the second case below pass
+// through and break the third continuation under the 30000-token label.
+func TestSubRunMaterialDecisionIncludesLastCompletion(t *testing.T) {
+	stWith := func(lastCompletion int) *subRun {
+		st := &subRun{
+			materialLimit:     30000,
+			materialOutputCap: 4000,
+			responses:         3,
+			usageKnown:        true,
+			usageTotal: core.Usage{
+				PromptTokens:         0,
+				CompletionTokens:     0,
+				TotalTokens:          22070,
+				LastPromptTokens:     5810,
+				LastCompletionTokens: lastCompletion,
+			},
+			savedCheckpointIDs: map[string]struct{}{"cp1": {}},
+		}
+		return st
+	}
+	_, stop := stWith(4000).materialDecision()
+	if !stop {
+		t.Fatal("subRun.materialDecision did not stop after a 4000-token previous completion over the reserve")
+	}
+	_, stop = stWith(0).materialDecision()
+	if stop {
+		t.Fatal("subRun.materialDecision stopped without the previous completion in the estimate; reserve estimate regressed")
+	}
+}

internal/server/work_checkpoint_eval_budget_test.go

@@ -789,6 +789,46 @@ func TestEvaluationPreservesProductionOutputCap(t *testing.T) {
 	}
 }
 
+// TestEvaluationFinalLabelCapHonouredByVariant: the 6000-token cap the final-label
+// run in model mode demands only reaches the provider through the parameterized
+// variant. The httptest server reads the request body so the assertion covers the
+// exact max_tokens the host sent downstream, mirroring TestEvaluationPreservesProductionOutputCap.
+func TestEvaluationFinalLabelCapHonouredByVariant(t *testing.T) {
+	for _, tc := range []struct {
+		name     string
+		supplied int
+		want     int
+	}{
+		{"supplied equals cap", 6000, 6000},
+		{"supplied zero clamps up", 0, 6000},
+		{"supplied above cap clamps down", 9000, 6000},
+		{"supplied below cap passes through", 780, 780},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			got := make(chan int, 1)
+			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				var body struct {
+					MaxTokens int `json:"max_tokens"`
+				}
+				if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+					t.Error(err)
+				}
+				got <- body.MaxTokens
+				w.Header().Set("Content-Type", "text/event-stream")
+				fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"done\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":1,\"total_tokens\":2}}\n\ndata: [DONE]\n\n")
+			}))
+			defer srv.Close()
+			dir, _ := newEvaluationEventStore(t)
+			stream := evalProviderStreamWithProviderCap(evalProviderClient(map[string]any{"baseUrl": srv.URL, "apiKey": "fixture", "model": "fixture"}, core.EffectiveSettings{}).Stream, dir, 6000)
+			for range stream(evalTestContext(t), "fixture", []agentllm.Message{{Role: agentllm.RoleUser, Content: "check cap"}}, nil, agentllm.StreamOptions{MaxTokens: tc.supplied}) {
+			}
+			if actual := <-got; actual != tc.want {
+				t.Fatalf("variant did not send the label-driven cap: got %d want %d", actual, tc.want)
+			}
+		})
+	}
+}
+
 func TestEvaluationStopsPaidRequestsAfterRejectedCheckpoint(t *testing.T) {
 	t.Setenv("WORK_CHECKPOINT_REMAINING_TOKENS", "60000")
 	providerCalls := 0

internal/server/work_checkpoint_eval_test.go

@@ -264,7 +264,7 @@ func TestServeWorkCheckpointEval(t *testing.T) {
 		if !evalRouteContractAttached(modelSettings, client) {
 			t.Fatal("route contract configured but not attached")
 		}
-		stream := evalProviderStreamWithProvider(client.Stream, dir)
+		stream := evalProviderStreamWithProviderCap(client.Stream, dir, evalA01MaxOutput(os.Getenv("WORK_CHECKPOINT_EXTRA_A01")))
 		if guardedA01 {
 			stream = limitExtraA01Runs(st, stream)
 		}
@@ -610,7 +610,15 @@ func evalTerminalObservation(observation *agentllm.TerminalObservation) *agentll
 	return clean
 }
 
+// evalProviderStreamWithProvider keeps the production 4000-token cap and is
+// the surface every existing call site and test asserts on. The actual clamp
+// and request logging live in evalProviderStreamWithProviderCap so the
+// final-label run in model mode can pass its 6000-token cap through.
 func evalProviderStreamWithProvider(provider agentllm.StreamFunc, dir string) agentllm.StreamFunc {
+	return evalProviderStreamWithProviderCap(provider, dir, 4000)
+}
+
+func evalProviderStreamWithProviderCap(provider agentllm.StreamFunc, dir string, outputCap int) agentllm.StreamFunc {
 	batchLimit := 120000
 	if raw := os.Getenv("WORK_CHECKPOINT_REMAINING_TOKENS"); raw != "" {
 		limit, err := strconv.Atoi(raw)
@@ -685,8 +693,8 @@ func evalProviderStreamWithProvider(provider agentllm.StreamFunc, dir string) ag
 		a.Calls++
 		groups[runKey] = a
 		mu.Unlock()
-		if opts.MaxTokens <= 0 || opts.MaxTokens > 4000 {
-			opts.MaxTokens = 4000
+		if opts.MaxTokens <= 0 || opts.MaxTokens > outputCap {
+			opts.MaxTokens = outputCap
 		}
 		reqCtx, cancel := context.WithDeadline(ctx, a.Started.Add(300*time.Second))
 		in := provider(reqCtx, m, msgs, decls, opts)