diff --git a/deps/agent-squad b/deps/agent-squad index 490b0b2..3b03a04 160000 --- a/deps/agent-squad +++ b/deps/agent-squad @@ -1 +1 @@ -Subproject commit 490b0b2dfa1567cb16ff7a7230fd78e52bc2938e +Subproject commit 3b03a04b4bfc54ed8a7873eea98ce6d8394018c5 diff --git a/internal/server/work_checkpoint_eval_budget_test.go b/internal/server/work_checkpoint_eval_budget_test.go index 8f0a931..5bc35f9 100644 --- a/internal/server/work_checkpoint_eval_budget_test.go +++ b/internal/server/work_checkpoint_eval_budget_test.go @@ -247,6 +247,166 @@ func TestEvaluationStopsAfterMalformedToolTerminal(t *testing.T) { } } +func TestEvaluationTransportObservationSanitized(t *testing.T) { + const maxTail, maxBytes = 10000, int64(32*1024*1024 + 1) + for _, tc := range []struct { + name string + input agentllm.TerminalObservation + wantTail int + wantBytes int64 + wantKind string + wantFields [3]string + }{ + {"negative-tail-oversized-body", agentllm.TerminalObservation{ReadOutcome: "done_marker", + TailEventKind: "PRIVATE_TAIL_R06", TailEventCount: -5, + BodyBoundary: &agentllm.BodyBoundaryObservation{ProtocolClass: "PRIVATE_HEADER_R06", FramingClass: "PRIVATE_FRAMING_R06", + LengthRelation: "PRIVATE_LENGTH_R06", BytesRead: 1 << 40, DeclaredLengthKnown: true, BytesSaturated: true, EOFSeen: true, Uncompressed: true}}, + 0, maxBytes, "unknown", [3]string{"other", "unknown", "unknown"}}, + {"oversized-tail-negative-body", agentllm.TerminalObservation{ReadOutcome: "done_marker", + TailEventKind: "PRIVATE_TAIL_R06", TailEventCount: 1 << 20, + BodyBoundary: &agentllm.BodyBoundaryObservation{ProtocolClass: "PRIVATE_HEADER_R06", FramingClass: "PRIVATE_FRAMING_R06", + LengthRelation: "PRIVATE_LENGTH_R06", BytesRead: -7, DeclaredLengthKnown: true, BytesSaturated: true, EOFSeen: true, Uncompressed: true}}, + maxTail, 0, "unknown", [3]string{"other", "unknown", "unknown"}}, + {"known-values-kept", agentllm.TerminalObservation{ReadOutcome: "done_marker", + TailEventKind: "comment", TailEventCount: 3, + BodyBoundary: &agentllm.BodyBoundaryObservation{ProtocolClass: "http2", FramingClass: "chunked", + LengthRelation: "match", BytesRead: 128, DeclaredLengthKnown: true, BytesSaturated: true, EOFSeen: true, Uncompressed: true}}, + 3, 128, "comment", [3]string{"http2", "chunked", "match"}}, + } { + t.Run(tc.name, func(t *testing.T) { + input := tc.input + inputBody := *input.BodyBoundary + input.BodyBoundary = &inputBody + clean := evalTerminalObservation(&input) + if clean == nil || clean.BodyBoundary == nil || clean.BodyBoundary == input.BodyBoundary { + t.Fatalf("observation was not deep-copied: %+v", clean) + } + body := clean.BodyBoundary + if clean.ReadOutcome != "done_marker" || clean.TailEventKind != tc.wantKind || clean.TailEventCount != tc.wantTail || + body.ProtocolClass != tc.wantFields[0] || body.FramingClass != tc.wantFields[1] || body.LengthRelation != tc.wantFields[2] || + body.BytesRead != tc.wantBytes || !body.DeclaredLengthKnown || !body.BytesSaturated || !body.EOFSeen || !body.Uncompressed { + t.Fatalf("unbounded observation: %+v body=%+v", clean, body) + } + input.BodyBoundary.ProtocolClass = "PRIVATE_LATE_R06" + if body.ProtocolClass == "PRIVATE_LATE_R06" { + t.Fatal("sanitized observation shares the provider body boundary") + } + + dir, _ := newEvaluationEventStore(t) + observed := tc.input + observedBody := *observed.BodyBoundary + observed.BodyBoundary = &observedBody + provider := func(context.Context, string, []agentllm.Message, []agentllm.ToolDecl, agentllm.StreamOptions) <-chan agentllm.Event { + out := make(chan agentllm.Event, 1) + out <- agentllm.Event{Kind: agentllm.KindMessageEnd, + Message: agentllm.Message{Role: agentllm.RoleAssistant, Content: "PRIVATE_BODY_R06", StopReason: agentllm.StopEnd}, + Usage: agentllm.Usage{PromptTokens: 12, CompletionTokens: 5, TotalTokens: 17}, + FinishReason: "stop", StreamComplete: true, HTTPStatus: http.StatusOK, TerminalObservation: &observed, + } + close(out) + return out + } + for range evalProviderStreamWithProvider(provider, dir)(evalTestContext(t), "fixture", + []agentllm.Message{{Role: agentllm.RoleUser, Content: "PRIVATE_PROMPT_R06"}}, nil, agentllm.StreamOptions{MaxTokens: 4000}) { + } + raw, err := os.ReadFile(filepath.Join(dir, "provider-records.json")) + if err != nil { + t.Fatal(err) + } + for _, forbidden := range []string{"PRIVATE_TAIL_R06", "PRIVATE_HEADER_R06", "PRIVATE_FRAMING_R06", "PRIVATE_LENGTH_R06", + "PRIVATE_BODY_R06", "PRIVATE_PROMPT_R06", `"messages"`, `"terminal"`} { + if strings.Contains(string(raw), forbidden) { + t.Fatalf("provider record retained %q", forbidden) + } + } + var record struct { + Calls []struct { + RunID string `json:"runId"` + Usage agentllm.Usage `json:"usage"` + Observation *agentllm.TerminalObservation `json:"terminalObservation"` + } `json:"calls"` + } + if err := json.Unmarshal(raw, &record); err != nil || len(record.Calls) != 1 || record.Calls[0].RunID != t.Name() || + record.Calls[0].Usage.TotalTokens != 17 || record.Calls[0].Observation == nil || record.Calls[0].Observation.BodyBoundary == nil { + t.Fatalf("sanitized record lost accounting: %s err=%v", raw, err) + } + got := record.Calls[0].Observation + if got.TailEventKind != tc.wantKind || got.TailEventCount != tc.wantTail || got.BodyBoundary.ProtocolClass != tc.wantFields[0] || + got.BodyBoundary.FramingClass != tc.wantFields[1] || got.BodyBoundary.LengthRelation != tc.wantFields[2] || + got.BodyBoundary.BytesRead != tc.wantBytes { + t.Fatalf("unbounded persisted observation: %s", raw) + } + }) + } +} + +func TestEvaluationTransportEvidenceStopsAfterUnmarkedToolEOF(t *testing.T) { + t.Setenv("WORK_CHECKPOINT_REMAINING_TOKENS", "60000") + dir, _ := newEvaluationEventStore(t) + var providerCalls atomic.Int32 + provider := func(context.Context, string, []agentllm.Message, []agentllm.ToolDecl, agentllm.StreamOptions) <-chan agentllm.Event { + providerCalls.Add(1) + out := make(chan agentllm.Event, 2) + out <- agentllm.Event{Kind: agentllm.KindToolCallDelta, ToolIndex: 0, ToolID: "call_1", ToolName: "publish_work_checkpoint", ArgumentsDelta: `{}`} + out <- agentllm.Event{Kind: agentllm.KindMessageEnd, + Message: agentllm.Message{Role: agentllm.RoleAssistant, StopReason: agentllm.StopError}, + Usage: agentllm.Usage{PromptTokens: 12, CompletionTokens: 5, TotalTokens: 17}, + Err: errors.New("PRIVATE_ERROR_R06"), ErrorMessage: "PRIVATE_ERROR_R06", + FinishReason: "tool_calls", StreamComplete: false, HTTPStatus: http.StatusOK, + TerminalObservation: &agentllm.TerminalObservation{ReadOutcome: "clean_eof", LastFrameJSONComplete: true, UsagePresent: true, + ToolFragments: 1, ToolArgumentBytes: 2, ToolArgumentsValid: true, TailEventKind: "json_frame", TailEventCount: 2, + BodyBoundary: &agentllm.BodyBoundaryObservation{ProtocolClass: "http1", FramingClass: "chunked", BytesRead: 256, + EOFSeen: true, LengthRelation: "unknown"}}, + } + close(out) + return out + } + stream := evalProviderStreamWithProvider(provider, 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.StreamComplete || + terminal.Usage.TotalTokens != 17 || len(terminal.Message.ToolCalls) != 0 { + t.Fatalf("unmarked tool EOF escaped the controlled terminal: %#v", terminal) + } + blocked := false + for event := range stream(bruntime.WithProviderRunID(context.Background(), t.Name()+"-next"), "fixture", nil, nil, agentllm.StreamOptions{MaxTokens: 4000}) { + blocked = blocked || event.Kind == agentllm.KindError + } + if !blocked || providerCalls.Load() != 1 { + t.Fatalf("transport EOF 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) + } + if strings.Contains(string(raw), "PRIVATE_ERROR_R06") || strings.Contains(string(raw), `"terminal"`) { + t.Fatal("transport evidence record retained raw provider content") + } + var record struct { + ReportedTokens int `json:"reportedTokens"` + TerminalFault bool `json:"terminalFault"` + Calls []struct { + Usage agentllm.Usage `json:"usage"` + StreamComplete bool `json:"streamComplete"` + FailureClass string `json:"failureClass"` + Observation *agentllm.TerminalObservation `json:"terminalObservation"` + } `json:"calls"` + } + if err := json.Unmarshal(raw, &record); err != nil || record.ReportedTokens != 17 || !record.TerminalFault || len(record.Calls) != 1 || + record.Calls[0].Usage.TotalTokens != 17 || record.Calls[0].StreamComplete || record.Calls[0].FailureClass != "terminal_error" { + t.Fatalf("transport EOF accounting missing: %s err=%v", raw, err) + } + got := record.Calls[0].Observation + if got == nil || got.ReadOutcome != "clean_eof" || got.TailEventKind != "json_frame" || got.TailEventCount != 2 || + got.BodyBoundary == nil || !got.BodyBoundary.EOFSeen || got.BodyBoundary.FramingClass != "chunked" { + t.Fatalf("transport evidence missing from failed terminal record: %s", raw) + } +} + func TestEvaluationKeepsUsageOnTerminalFailure(t *testing.T) { t.Setenv("WORK_CHECKPOINT_REMAINING_TOKENS", "60000") var providerCalls atomic.Int32 diff --git a/internal/server/work_checkpoint_eval_test.go b/internal/server/work_checkpoint_eval_test.go index 7ee1e4f..71bd1c6 100644 --- a/internal/server/work_checkpoint_eval_test.go +++ b/internal/server/work_checkpoint_eval_test.go @@ -462,17 +462,60 @@ func evalFinishReasonClass(reason string) string { } } +// Limits mirror Agent-Native's saturating observation counters. +const ( + evalMaxTailEvents = 10000 + evalMaxBodyBytes = 32*1024*1024 + 1 +) + +func evalLabel(value, fallback string, allowed ...string) string { + for _, label := range allowed { + if value == label { + return value + } + } + return fallback +} + +func clampCount(value, limit int64) int64 { + if value < 0 { + return 0 + } + if value > limit { + return limit + } + return value +} + func evalTerminalObservation(observation *agentllm.TerminalObservation) *agentllm.TerminalObservation { if observation == nil { return nil } - copy := *observation - switch copy.ReadOutcome { - case "done_marker", "clean_eof", "read_error", "cancelled", "invalid_frame", "provider_error", "transport_error": - default: - copy.ReadOutcome = "unknown" + clean := &agentllm.TerminalObservation{ + ReadOutcome: evalLabel(observation.ReadOutcome, "unknown", + "done_marker", "clean_eof", "read_error", "cancelled", "invalid_frame", "provider_error", "transport_error"), + LastFrameJSONComplete: observation.LastFrameJSONComplete, + UsagePresent: observation.UsagePresent, + ToolFragments: observation.ToolFragments, + ToolArgumentBytes: observation.ToolArgumentBytes, + ToolArgumentsValid: observation.ToolArgumentsValid, + TailEventKind: evalLabel(observation.TailEventKind, "unknown", + "done_literal", "json_frame", "comment", "empty_data", "partial_frame", "unknown_non_json", "unknown"), + TailEventCount: int(clampCount(int64(observation.TailEventCount), evalMaxTailEvents)), } - return © + if body := observation.BodyBoundary; body != nil { + clean.BodyBoundary = &agentllm.BodyBoundaryObservation{ + ProtocolClass: evalLabel(body.ProtocolClass, "other", "http1", "http2", "other"), + FramingClass: evalLabel(body.FramingClass, "unknown", "chunked", "known_length", "close_delimited", "http2_stream", "unknown"), + DeclaredLengthKnown: body.DeclaredLengthKnown, + BytesRead: clampCount(body.BytesRead, evalMaxBodyBytes), + BytesSaturated: body.BytesSaturated, + EOFSeen: body.EOFSeen, + LengthRelation: evalLabel(body.LengthRelation, "unknown", "match", "mismatch", "unknown"), + Uncompressed: body.Uncompressed, + } + } + return clean } func evalProviderStreamWithProvider(provider agentllm.StreamFunc, dir string) agentllm.StreamFunc {