diff --git a/internal/server/work_checkpoint_eval_budget_test.go b/internal/server/work_checkpoint_eval_budget_test.go index 3d2b982..7df8f8b 100644 --- a/internal/server/work_checkpoint_eval_budget_test.go +++ b/internal/server/work_checkpoint_eval_budget_test.go @@ -6,15 +6,32 @@ import ( "agent-system-backend/internal/store" core "agent-system-core" "context" + "database/sql" "encoding/json" "fmt" "net/http" "net/http/httptest" + "path/filepath" "testing" agentllm "agent-native/runtime/llm" ) +func newEvaluationEventStore(t *testing.T) (string, *sql.DB) { + t.Helper() + dir := t.TempDir() + db, err := sql.Open("sqlite", filepath.Join(dir, "agent-system.db")) + if err != nil { + t.Fatal(err) + } + if _, err = db.Exec(`CREATE TABLE events (conversation_id TEXT, run_id TEXT, seq INTEGER, type TEXT, payload TEXT)`); err != nil { + db.Close() + t.Fatal(err) + } + t.Cleanup(func() { db.Close() }) + return dir, db +} + func TestExtraA01LimitsActualRunIDs(t *testing.T) { _, st := newTestServer(t) p, err := st.CreateProject(store.OwnerDefault, "allowance") @@ -91,7 +108,8 @@ func TestEvaluationPreservesProductionOutputCap(t *testing.T) { 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() - stream := evalProviderStream(map[string]any{"baseUrl": srv.URL, "apiKey": "fixture", "model": "fixture"}, t.TempDir()) + dir, _ := newEvaluationEventStore(t) + stream := evalProviderStream(map[string]any{"baseUrl": srv.URL, "apiKey": "fixture", "model": "fixture"}, dir) for range stream(context.Background(), "fixture", []agentllm.Message{{Role: agentllm.RoleUser, Content: "check cap"}}, nil, agentllm.StreamOptions{MaxTokens: tc.supplied}) { } if actual := <-got; actual != tc.want { @@ -110,7 +128,8 @@ func TestEvaluationStopsPaidRequestsAfterRejectedCheckpoint(t *testing.T) { 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() - stream := evalProviderStream(map[string]any{"baseUrl": srv.URL, "apiKey": "fixture", "model": "fixture"}, t.TempDir()) + dir, _ := newEvaluationEventStore(t) + stream := evalProviderStream(map[string]any{"baseUrl": srv.URL, "apiKey": "fixture", "model": "fixture"}, dir) for _, messages := range [][]agentllm.Message{ {{Role: agentllm.RoleTool, Name: "publish_work_checkpoint", IsError: true, Content: "阶段发现未发布：checkpoint.source: quote not in provided block"}}, {{Role: agentllm.RoleUser, Content: "another request after rejection"}}, @@ -129,3 +148,83 @@ func TestEvaluationStopsPaidRequestsAfterRejectedCheckpoint(t *testing.T) { t.Fatalf("%d paid provider calls started after checkpoint rejection", providerCalls) } } + +func TestEvaluationStopsAcrossRunsAfterPersistedCheckpointRejection(t *testing.T) { + t.Setenv("WORK_CHECKPOINT_REMAINING_TOKENS", "60000") + dir, db := newEvaluationEventStore(t) + _, err := db.Exec(`INSERT INTO events(conversation_id,run_id,seq,type,payload) VALUES(?,?,?,?,?)`, "background", "child", 1, "tool.call", `{"callId":"rejected-call","name":"publish_work_checkpoint"}`) + if err != nil { + t.Fatal(err) + } + _, err = db.Exec(`INSERT INTO events(conversation_id,run_id,seq,type,payload) VALUES(?,?,?,?,?)`, "background", "child", 2, "tool.result", `{"callId":"rejected-call","status":"error","summary":"json: unknown field \"$text\""}`) + if err != nil { + t.Fatal(err) + } + providerCalls := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + providerCalls++ + 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() + stream := evalProviderStream(map[string]any{"baseUrl": srv.URL, "apiKey": "fixture", "model": "fixture"}, dir) + rejected := false + for event := range stream(context.Background(), "fixture", []agentllm.Message{{Role: agentllm.RoleUser, Content: "unrelated foreground history"}}, nil, agentllm.StreamOptions{MaxTokens: 4000}) { + if event.Kind == agentllm.KindError { + rejected = true + } + } + if !rejected || providerCalls != 0 { + t.Fatalf("persisted background rejection allowed foreground provider request: error=%v calls=%d", rejected, providerCalls) + } +} + +func TestEvaluationFailsClosedWhenCheckpointEventsUnavailable(t *testing.T) { + providerCalls := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + providerCalls++ + })) + defer srv.Close() + stream := evalProviderStream(map[string]any{"baseUrl": srv.URL, "apiKey": "fixture", "model": "fixture"}, t.TempDir()) + rejected := false + for event := range stream(context.Background(), "fixture", []agentllm.Message{{Role: agentllm.RoleUser, Content: "request without the event store"}}, nil, agentllm.StreamOptions{MaxTokens: 4000}) { + if event.Kind == agentllm.KindError { + rejected = true + } + } + if !rejected || providerCalls != 0 { + t.Fatalf("missing checkpoint event store allowed provider request: error=%v calls=%d", rejected, providerCalls) + } +} + +func TestEvaluationMatchesCheckpointRejectionWithinItsRun(t *testing.T) { + for _, tc := range []struct { + name, tool, resultRun string + }{ + {"different tool", "read_file", "child"}, + {"different run", "publish_work_checkpoint", "other-run"}, + } { + t.Run(tc.name, func(t *testing.T) { + dir, db := newEvaluationEventStore(t) + if _, err := db.Exec(`INSERT INTO events(conversation_id,run_id,seq,type,payload) VALUES(?,?,?,?,?)`, "background", "child", 1, "tool.call", fmt.Sprintf(`{"callId":"shared","name":%q}`, tc.tool)); err != nil { + t.Fatal(err) + } + if _, err := db.Exec(`INSERT INTO events(conversation_id,run_id,seq,type,payload) VALUES(?,?,?,?,?)`, "background", tc.resultRun, 2, "tool.result", `{"callId":"shared","status":"error"}`); err != nil { + t.Fatal(err) + } + providerCalls := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + providerCalls++ + 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() + stream := evalProviderStream(map[string]any{"baseUrl": srv.URL, "apiKey": "fixture", "model": "fixture"}, dir) + for range stream(context.Background(), "fixture", []agentllm.Message{{Role: agentllm.RoleUser, Content: "unrelated request"}}, nil, agentllm.StreamOptions{MaxTokens: 4000}) { + } + if providerCalls != 1 { + t.Fatalf("unrelated tool error blocked the provider: calls=%d", providerCalls) + } + }) + } +} diff --git a/internal/server/work_checkpoint_eval_test.go b/internal/server/work_checkpoint_eval_test.go index a47d8b2..e279571 100644 --- a/internal/server/work_checkpoint_eval_test.go +++ b/internal/server/work_checkpoint_eval_test.go @@ -5,6 +5,7 @@ package server import ( "context" "crypto/sha256" + "database/sql" "encoding/hex" "encoding/json" "fmt" @@ -284,6 +285,55 @@ func limitExtraA01Runs(st *store.Store, next agentllm.StreamFunc) agentllm.Strea // Fixed shared provider limits apply equally to baseline and candidate, including // versions that did not yet honor the child profile budget. +func evalRejectedCheckpointPersisted(dir string) (bool, error) { + db, err := sql.Open("sqlite", "file:"+filepath.Join(dir, "agent-system.db")+"?mode=ro") + if err != nil { + return false, err + } + defer db.Close() + rows, err := db.Query(`SELECT conversation_id, run_id, type, payload FROM events WHERE type IN ('tool.call', 'tool.result')`) + if err != nil { + return false, err + } + defer rows.Close() + checkpointCalls := map[string]bool{} + rejectedResults := map[string]bool{} + for rows.Next() { + var conversationID, runID, eventType string + var raw []byte + if err = rows.Scan(&conversationID, &runID, &eventType, &raw); err != nil { + return false, err + } + var event struct { + CallID string `json:"callId"` + Name string `json:"name"` + Status string `json:"status"` + } + if err = json.Unmarshal(raw, &event); err != nil { + return false, err + } + if event.CallID == "" { + return false, fmt.Errorf("checkpoint stop event missing call ID") + } + key := conversationID + "\x00" + runID + "\x00" + event.CallID + if eventType == "tool.call" && event.Name == "publish_work_checkpoint" { + checkpointCalls[key] = true + } + if eventType == "tool.result" && event.Status == "error" { + rejectedResults[key] = true + } + } + if err = rows.Err(); err != nil { + return false, err + } + for key := range rejectedResults { + if checkpointCalls[key] { + return true, nil + } + } + return false, nil +} + func evalProviderStream(route map[string]any, dir string) agentllm.StreamFunc { batchLimit := 120000 if raw := os.Getenv("WORK_CHECKPOINT_REMAINING_TOKENS"); raw != "" { @@ -308,6 +358,7 @@ func evalProviderStream(route map[string]any, dir string) agentllm.StreamFunc { checkpointRejected := false records := []map[string]any{} return func(ctx context.Context, m string, msgs []agentllm.Message, decls []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event { + persistedRejected, persistenceErr := evalRejectedCheckpointPersisted(dir) rejectedInput := false for _, msg := range msgs { if msg.Role == agentllm.RoleTool && msg.Name == "publish_work_checkpoint" && msg.IsError { @@ -329,7 +380,7 @@ func evalProviderStream(route map[string]any, dir string) agentllm.StreamFunc { hash := sha256.Sum256(fp) runKey := hex.EncodeToString(hash[:]) mu.Lock() - if rejectedInput { + if rejectedInput || persistedRejected { checkpointRejected = true } a := groups[runKey] @@ -337,7 +388,9 @@ func evalProviderStream(route map[string]any, dir string) agentllm.StreamFunc { a.Started = time.Now() } reason := "" - if checkpointRejected { + if persistenceErr != nil { + reason = "checkpoint stop state unavailable" + } else if checkpointRejected { reason = "rejected work checkpoint" } else if unknown { reason = "provider usage unavailable"