From efed68ca2c8fbf148a2c40d61906612a9cc006f1 Mon Sep 17 00:00:00 2001 From: weichen Date: Wed, 14 Oct 2026 02:57:17 +0000 Subject: [PATCH] Stop evaluation after an incomplete provider terminal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Why: the controlled A01 wrapper still allowed another paid request after a malformed tool terminal with returned Usage. Evidence: local SSE RED→GREEN tests cover invalid JSON, missing DONE, late errors, missing Usage, rejected checkpoints, cross-run stop and record-write failure; live-tag fake tests and full Backend vet/tests pass. Impact: diagnostic metadata and batch-wide fail-closed stop only; no new paid authorization or production route change. Refs: U3 r0.4 design; internal/server/work_checkpoint_eval_test.go --- .../work_checkpoint_eval_budget_test.go | 181 ++++++++++++++++++ internal/server/work_checkpoint_eval_test.go | 38 +++- 2 files changed, 216 insertions(+), 3 deletions(-) diff --git a/internal/server/work_checkpoint_eval_budget_test.go b/internal/server/work_checkpoint_eval_budget_test.go index d4edac5..1c1da0b 100644 --- a/internal/server/work_checkpoint_eval_budget_test.go +++ b/internal/server/work_checkpoint_eval_budget_test.go @@ -15,11 +15,192 @@ import ( "os" "path/filepath" "strings" + "sync/atomic" "testing" agentllm "agent-native/runtime/llm" ) +func TestEvaluationStopsAfterMalformedToolTerminal(t *testing.T) { + t.Setenv("WORK_CHECKPOINT_REMAINING_TOKENS", "60000") + for _, tc := range []struct { + name, arguments, trailer string + complete bool + argumentsValid bool + }{ + {"invalid-json", `{"findings":`, "data: [DONE]\n\n", true, false}, + {"missing-done", `{}`, "", false, true}, + } { + t.Run(tc.name, func(t *testing.T) { + var providerCalls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + providerCalls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + args, _ := json.Marshal(tc.arguments) + fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"function\":{\"name\":\"publish_work_checkpoint\",\"arguments\":%s}}]}}]}\n\n", args) + fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"tool_calls\"}],\"usage\":{\"prompt_tokens\":12,\"completion_tokens\":5,\"total_tokens\":17}}\n\n") + fmt.Fprint(w, tc.trailer) + })) + defer srv.Close() + dir, _ := newEvaluationEventStore(t) + stream := evalProviderStream(map[string]any{"baseUrl": srv.URL, "apiKey": "fixture", "model": "fixture"}, dir) + var terminal agentllm.Event + for event := range stream(evalTestContext(t), "fixture", nil, nil, agentllm.StreamOptions{MaxTokens: 4000}) { + if event.Kind == agentllm.KindMessageEnd { + terminal = event + } + } + if terminal.Kind != agentllm.KindMessageEnd || terminal.Err == nil || terminal.Usage.TotalTokens != 17 || + terminal.StreamComplete != tc.complete || len(terminal.Message.ToolCalls) != 0 { + t.Fatalf("unsafe provider terminal: %#v", terminal) + } + blocked := false + otherRun := bruntime.WithProviderRunID(context.Background(), t.Name()+"-other-run") + for event := range stream(otherRun, "fixture", nil, nil, agentllm.StreamOptions{MaxTokens: 4000}) { + blocked = blocked || event.Kind == agentllm.KindError + } + if !blocked || providerCalls.Load() != 1 { + t.Fatalf("terminal failure allowed another provider call: blocked=%v calls=%d", blocked, providerCalls.Load()) + } + raw, err := os.ReadFile(filepath.Join(dir, "provider-records.json")) + if err != nil { + t.Fatal(err) + } + var record struct { + ReportedTokens int `json:"reportedTokens"` + UsageUnknown bool `json:"usageUnknown"` + TerminalFault bool `json:"terminalFault"` + Calls []struct { + ReportedTokens int `json:"reportedTokens"` + OutputCap int `json:"outputCap"` + FinishReason string `json:"finishReason"` + StreamComplete bool `json:"streamComplete"` + ToolArgumentBytes int `json:"toolArgumentBytes"` + ArgumentsValid bool `json:"toolArgumentsValid"` + UsagePresent bool `json:"usagePresent"` + } `json:"calls"` + } + if err = json.Unmarshal(raw, &record); err != nil { + t.Fatal(err) + } + if record.ReportedTokens != 17 || record.UsageUnknown || !record.TerminalFault || len(record.Calls) != 1 || + record.Calls[0].ReportedTokens != 17 || record.Calls[0].OutputCap != 4000 || + record.Calls[0].FinishReason != "tool_calls" || record.Calls[0].StreamComplete != tc.complete || + record.Calls[0].ToolArgumentBytes == 0 || record.Calls[0].ArgumentsValid != tc.argumentsValid || + !record.Calls[0].UsagePresent { + t.Fatalf("incomplete diagnostic or lost Usage: %+v", record) + } + }) + } +} + +func TestEvaluationKeepsUsageOnTerminalFailure(t *testing.T) { + t.Setenv("WORK_CHECKPOINT_REMAINING_TOKENS", "60000") + var providerCalls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + providerCalls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"function\":{\"name\":\"publish_work_checkpoint\",\"arguments\":\"{}\"}}]}}],\"usage\":{\"prompt_tokens\":12,\"completion_tokens\":5,\"total_tokens\":17}}\n\n") + fmt.Fprint(w, "data: {not-json}\n\n") + })) + defer srv.Close() + dir, _ := newEvaluationEventStore(t) + stream := evalProviderStream(map[string]any{"baseUrl": srv.URL, "apiKey": "fixture", "model": "fixture"}, dir) + var terminal agentllm.Event + for ev := range stream(evalTestContext(t), "fixture", nil, nil, agentllm.StreamOptions{MaxTokens: 4000}) { + if ev.Kind == agentllm.KindMessageEnd { + terminal = ev + } + } + if terminal.Err == nil || terminal.Usage.TotalTokens != 17 { + t.Fatalf("late stream failure lost returned Usage: %#v", terminal) + } + blocked := false + otherRun := bruntime.WithProviderRunID(context.Background(), t.Name()+"-other-run") + for ev := range stream(otherRun, "fixture", nil, nil, agentllm.StreamOptions{MaxTokens: 4000}) { + blocked = blocked || ev.Kind == agentllm.KindError + } + raw, err := os.ReadFile(filepath.Join(dir, "provider-records.json")) + if err != nil { + t.Fatal(err) + } + var record struct { + ReportedTokens int `json:"reportedTokens"` + UsageUnknown bool `json:"usageUnknown"` + TerminalFault bool `json:"terminalFault"` + } + if err = json.Unmarshal(raw, &record); err != nil { + t.Fatal(err) + } + if !blocked || providerCalls.Load() != 1 || record.ReportedTokens != 17 || record.UsageUnknown || !record.TerminalFault { + t.Fatalf("late terminal was retried or Usage lost: blocked=%v calls=%d record=%+v", blocked, providerCalls.Load(), record) + } +} + +func TestEvaluationStopsWhenUsageMissing(t *testing.T) { + t.Setenv("WORK_CHECKPOINT_REMAINING_TOKENS", "60000") + var providerCalls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + providerCalls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"done\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n") + })) + defer srv.Close() + dir, _ := newEvaluationEventStore(t) + stream := evalProviderStream(map[string]any{"baseUrl": srv.URL, "apiKey": "fixture", "model": "fixture"}, dir) + for range stream(evalTestContext(t), "fixture", nil, nil, agentllm.StreamOptions{MaxTokens: 4000}) { + } + blocked := false + for ev := range stream(evalTestContext(t), "fixture", nil, nil, agentllm.StreamOptions{MaxTokens: 4000}) { + blocked = blocked || ev.Kind == agentllm.KindError + } + if !blocked || providerCalls.Load() != 1 { + t.Fatalf("missing Usage allowed a second paid request: blocked=%v calls=%d", blocked, providerCalls.Load()) + } + raw, err := os.ReadFile(filepath.Join(dir, "provider-records.json")) + if err != nil { + t.Fatal(err) + } + var record struct { + UsageUnknown bool `json:"usageUnknown"` + } + if err = json.Unmarshal(raw, &record); err != nil || !record.UsageUnknown { + t.Fatalf("missing Usage was not recorded: %v %+v", err, record) + } +} + +func TestEvaluationStopsWhenProviderRecordUnavailable(t *testing.T) { + t.Setenv("WORK_CHECKPOINT_REMAINING_TOKENS", "60000") + var providerCalls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + providerCalls.Add(1) + 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) + if err := os.Mkdir(filepath.Join(dir, "provider-records.json"), 0700); err != nil { + t.Fatal(err) + } + stream := evalProviderStream(map[string]any{"baseUrl": srv.URL, "apiKey": "fixture", "model": "fixture"}, dir) + var terminal agentllm.Event + for ev := range stream(evalTestContext(t), "fixture", nil, nil, agentllm.StreamOptions{MaxTokens: 4000}) { + if ev.Kind == agentllm.KindMessageEnd { + terminal = ev + } + } + if terminal.Usage.TotalTokens != 2 { + t.Fatalf("control provider Usage missing: %#v", terminal) + } + blocked := false + for ev := range stream(evalTestContext(t), "fixture", nil, nil, agentllm.StreamOptions{MaxTokens: 4000}) { + blocked = blocked || ev.Kind == agentllm.KindError + } + if !blocked || providerCalls.Load() != 1 { + t.Fatalf("missing durable record allowed another request: blocked=%v calls=%d", blocked, providerCalls.Load()) + } +} + func TestEvaluationGroupsChangingPromptsByActualRunID(t *testing.T) { t.Setenv("WORK_CHECKPOINT_REMAINING_TOKENS", "60000") providerCalls := 0 diff --git a/internal/server/work_checkpoint_eval_test.go b/internal/server/work_checkpoint_eval_test.go index 311613b..f85b78f 100644 --- a/internal/server/work_checkpoint_eval_test.go +++ b/internal/server/work_checkpoint_eval_test.go @@ -366,6 +366,7 @@ func evalProviderStream(route map[string]any, dir string) agentllm.StreamFunc { total := 0 unknown := false checkpointRejected := false + terminalFault := false records := []map[string]any{} return func(ctx context.Context, m string, msgs []agentllm.Message, decls []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event { runID := runtime.ProviderRunID(ctx) @@ -398,6 +399,8 @@ func evalProviderStream(route map[string]any, dir string) agentllm.StreamFunc { reason = "checkpoint stop state unavailable" } else if checkpointRejected { reason = "rejected work checkpoint" + } else if terminalFault { + reason = "incomplete or failed provider terminal" } else if unknown { reason = "provider usage unavailable" } else if total >= batchLimit { @@ -428,7 +431,13 @@ func evalProviderStream(route map[string]any, dir string) agentllm.StreamFunc { defer cancel() usage := 0 var terminal *agentllm.Event + arguments := map[int]string{} + argumentBytes := 0 for event := range in { + if event.Kind == agentllm.KindToolCallDelta { + arguments[event.ToolIndex] += event.ArgumentsDelta + argumentBytes += len(event.ArgumentsDelta) + } if event.Kind == agentllm.KindMessageEnd { copy := event terminal = © @@ -443,15 +452,38 @@ func evalProviderStream(route map[string]any, dir string) agentllm.StreamFunc { return } } + argumentsValid := true + for _, raw := range arguments { + var object map[string]json.RawMessage + if err := json.Unmarshal([]byte(raw), &object); err != nil || object == nil { + argumentsValid = false + break + } + } + fault := terminal == nil || terminal.Err != nil || terminal.Message.StopReason == agentllm.StopError + finishReason, streamComplete := "", false + if terminal != nil { + finishReason, streamComplete = terminal.FinishReason, terminal.StreamComplete + } mu.Lock() total += usage if usage == 0 { unknown = true } - record := map[string]any{"runId": runID, "runKey": runKey, "call": a.Calls, "model": m, "messages": msgs, "toolDeclarations": decls, "outputCap": opts.MaxTokens, "reportedTokens": usage, "terminal": terminal} + if fault { + terminalFault = true + } + record := map[string]any{"runId": runID, "runKey": runKey, "call": a.Calls, "model": m, + "messages": msgs, "toolDeclarations": decls, "outputCap": opts.MaxTokens, + "reportedTokens": usage, "usagePresent": usage > 0, "terminal": terminal, + "finishReason": finishReason, "streamComplete": streamComplete, + "toolArgumentBytes": argumentBytes, "toolArgumentsValid": argumentsValid} records = append(records, record) - raw, _ := json.MarshalIndent(map[string]any{"calls": records, "reportedTokens": total, "usageUnknown": unknown, "modelRuns": len(groups)}, "", " ") - _ = os.WriteFile(filepath.Join(dir, "provider-records.json"), raw, 0600) + raw, _ := json.MarshalIndent(map[string]any{"calls": records, "reportedTokens": total, + "usageUnknown": unknown, "terminalFault": terminalFault, "modelRuns": len(groups)}, "", " ") + if err := os.WriteFile(filepath.Join(dir, "provider-records.json"), raw, 0600); err != nil { + terminalFault = true // no durable accounting: never issue another paid request + } mu.Unlock() }() return out