diff --git a/backend/cmd/ftester/mocks/tools.go b/backend/cmd/ftester/mocks/tools.go index e0f6f561..6d95c01c 100644 --- a/backend/cmd/ftester/mocks/tools.go +++ b/backend/cmd/ftester/mocks/tools.go @@ -40,11 +40,15 @@ func MockResponse(funcName string, args json.RawMessage) (string, error) { terminal.PrintMock("File operation:") terminal.PrintKeyValue("Operation", string(fileArgs.Action)) - terminal.PrintKeyValue("Path", fileArgs.Path) + terminal.PrintKeyValue("Path", fileArgs.Path.String()) - if fileArgs.Action == tools.ReadFile { + switch fileArgs.Action { + case tools.ReadFile: resultObj = fmt.Sprintf("Mock content of file: %s\nThis is a sample content that would be read from the file.\nIt contains multiple lines to simulate a real file.", fileArgs.Path) - } else { + case tools.EditFile: + terminal.PrintKeyValue("Diff", fileArgs.Diff.String()) + resultObj = fmt.Sprintf("Applied 1 diff hunk(s) to %s (mock)", fileArgs.Path) + default: resultObj = fmt.Sprintf("file %s written successfully", fileArgs.Path) } @@ -211,8 +215,8 @@ func MockResponse(funcName string, args json.RawMessage) (string, error) { terminal.PrintMock("Sploitus search:") terminal.PrintKeyValue("Query", sploitusArgs.Query) - terminal.PrintKeyValue("Exploit type", exploitType) - terminal.PrintKeyValue("Sort", sploitusArgs.Sort) + terminal.PrintKeyValue("Exploit type", exploitType.String()) + terminal.PrintKeyValue("Sort", sploitusArgs.Sort.String()) terminal.PrintKeyValueFormat("Max results", "%d", sploitusArgs.MaxResults.Int()) var builder strings.Builder @@ -414,7 +418,7 @@ func MockResponse(funcName string, args json.RawMessage) (string, error) { for i, q := range searchGuideArgs.Questions { terminal.PrintKeyValueFormat(fmt.Sprintf("Question %d", i+1), "%s", q) } - terminal.PrintKeyValue("Guide type", searchGuideArgs.Type) + terminal.PrintKeyValue("Guide type", searchGuideArgs.Type.String()) questionsText := strings.Join(searchGuideArgs.Questions, " | ") @@ -433,7 +437,7 @@ func MockResponse(funcName string, args json.RawMessage) (string, error) { } terminal.PrintMock("Store guide:") - terminal.PrintKeyValue("Type", storeGuideArgs.Type) + terminal.PrintKeyValue("Type", storeGuideArgs.Type.String()) terminal.PrintKeyValueFormat("Guide length", "%d chars", len(storeGuideArgs.Guide)) terminal.PrintKeyValue("Guide question", storeGuideArgs.Question) @@ -450,7 +454,7 @@ func MockResponse(funcName string, args json.RawMessage) (string, error) { for i, q := range searchAnswerArgs.Questions { terminal.PrintKeyValueFormat(fmt.Sprintf("Question %d", i+1), "%s", q) } - terminal.PrintKeyValue("Answer type", searchAnswerArgs.Type) + terminal.PrintKeyValue("Answer type", searchAnswerArgs.Type.String()) questionsText := strings.Join(searchAnswerArgs.Questions, " | ") @@ -467,7 +471,7 @@ func MockResponse(funcName string, args json.RawMessage) (string, error) { } terminal.PrintMock("Store answer:") - terminal.PrintKeyValue("Type", storeAnswerArgs.Type) + terminal.PrintKeyValue("Type", storeAnswerArgs.Type.String()) terminal.PrintKeyValueFormat("Answer length", "%d chars", len(storeAnswerArgs.Answer)) terminal.PrintKeyValue("Question", storeAnswerArgs.Question) @@ -520,7 +524,7 @@ func MockResponse(funcName string, args json.RawMessage) (string, error) { } terminal.PrintMock("Graphiti Search:") - terminal.PrintKeyValue("Search Type", searchArgs.SearchType) + terminal.PrintKeyValue("Search Type", searchArgs.SearchType.String()) terminal.PrintKeyValue("Query", searchArgs.Query) var builder strings.Builder diff --git a/backend/go.mod b/backend/go.mod index 7301cacc..36370949 100644 --- a/backend/go.mod +++ b/backend/go.mod @@ -39,10 +39,11 @@ require ( github.com/joho/godotenv v1.5.1 github.com/lib/pq v1.10.9 github.com/mattn/go-runewidth v0.0.16 - github.com/ollama/ollama v0.18.0 + github.com/ollama/ollama v0.23.0 github.com/pgvector/pgvector-go v0.1.1 github.com/pressly/goose/v3 v3.19.2 github.com/rivo/uniseg v0.4.7 + github.com/sergi/go-diff v1.3.1 github.com/shirou/gopsutil/v3 v3.24.5 github.com/sirupsen/logrus v1.9.3 github.com/sqlc-dev/pqtype v0.3.0 @@ -52,7 +53,7 @@ require ( github.com/vektah/gqlparser/v2 v2.5.19 github.com/vxcontrol/cloud v0.9.0 github.com/vxcontrol/graphiti-go-client v0.9.0 - github.com/vxcontrol/langchaingo v0.1.15-0.20260717093433-cf5f5f41f4eb + github.com/vxcontrol/langchaingo v0.1.15-0.20260723091023-35da5c0f0620 github.com/wasilibs/go-re2 v1.10.0 github.com/xeipuuv/gojsonschema v1.2.0 go.opentelemetry.io/otel v1.39.0 diff --git a/backend/go.sum b/backend/go.sum index 46df62b9..58709846 100644 --- a/backend/go.sum +++ b/backend/go.sum @@ -483,8 +483,8 @@ github.com/ncruces/go-strftime v0.1.9 h1:bY0MQC28UADQmHmaF5dgpLmImcShSi2kHU9XLdh github.com/ncruces/go-strftime v0.1.9/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= github.com/nikolalohinski/gonja v1.5.3 h1:GsA+EEaZDZPGJ8JtpeGN78jidhOlxeJROpqMT9fTj9c= github.com/nikolalohinski/gonja v1.5.3/go.mod h1:RmjwxNiXAEqcq1HeK5SSMmqFJvKOfTfXhkJv6YBtPa4= -github.com/ollama/ollama v0.18.0 h1:loPvswLB07Cn3SnRy5E9tZziGS4nqfnoVllSKO68vX8= -github.com/ollama/ollama v0.18.0/go.mod h1:tCX4IMV8DHjl3zY0THxuEkpWDZSOchJpzTuLACpMwFw= +github.com/ollama/ollama v0.23.0 h1:13V15B9Pkwl+WAaaU90NW6vik54dYgPhVrdjtw+dqck= +github.com/ollama/ollama v0.23.0/go.mod h1:274niu48upWz/M7vL53i1WFe+TJRRw5oo4GiacbIYrA= github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U= github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM= github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJwooC2xJA040= @@ -632,8 +632,8 @@ github.com/vxcontrol/cloud v0.9.0 h1:p7xYTgUctbY8w6YfhugNzvfi3/0EQoZGumMe67keAng github.com/vxcontrol/cloud v0.9.0/go.mod h1:AeiQFqiMgJJAXy6FYXtDS2a3P/PMB56iiBNY2vGrZhQ= github.com/vxcontrol/graphiti-go-client v0.9.0 h1:3GxpFmQoHmz/d7/9tyEqD8+S99v2cuqG1UEmrbAFrLU= github.com/vxcontrol/graphiti-go-client v0.9.0/go.mod h1:6UHL5uqAKp4KAdziva4qgcAxFtBzU05Hm/BAo4NkAuo= -github.com/vxcontrol/langchaingo v0.1.15-0.20260717093433-cf5f5f41f4eb h1:rufkyNpHdmdXXw2YcBZRLocgNMAHrdDDrQyosGubhr4= -github.com/vxcontrol/langchaingo v0.1.15-0.20260717093433-cf5f5f41f4eb/go.mod h1:VcRflCqD/8y/qPEZaqdtDFCMX378SZwChm5XXuV8xQk= +github.com/vxcontrol/langchaingo v0.1.15-0.20260723091023-35da5c0f0620 h1:k03Hr4LKY09OnwUaXXKlXR0rZfJ3TVilSClY60Ch/Nk= +github.com/vxcontrol/langchaingo v0.1.15-0.20260723091023-35da5c0f0620/go.mod h1:Lpk9Nc9N7B3/CuMy2Ay9c68Xh/AQhPABG1NH3cFI+sE= github.com/wasilibs/go-re2 v1.10.0 h1:vQZEBYZOCA9jdBMmrO4+CvqyCj0x4OomXTJ4a5/urQ0= github.com/wasilibs/go-re2 v1.10.0/go.mod h1:k+5XqO2bCJS+QpGOnqugyfwC04nw0jaglmjrrkG8U6o= github.com/wasilibs/wazero-helpers v0.0.0-20240620070341-3dff1577cd52 h1:OvLBa8SqJnZ6P+mjlzc2K7PM22rRUPE1x32G9DTPrC4= diff --git a/backend/pkg/observability/langfuse/observer.go b/backend/pkg/observability/langfuse/observer.go index 179a5673..f54eb955 100644 --- a/backend/pkg/observability/langfuse/observer.go +++ b/backend/pkg/observability/langfuse/observer.go @@ -2,11 +2,15 @@ package langfuse import ( "context" + "encoding/json" + "errors" "fmt" + "net/http" "sync" "time" "pentagi/pkg/observability/langfuse/api" + "pentagi/pkg/observability/langfuse/api/core" "github.com/sirupsen/logrus" ) @@ -197,6 +201,54 @@ func (o *observer) flush(ctx context.Context, batch []*api.IngestionEvent) error return nil } +// flushWithSplit flushes the batch and, if the request fails only because +// the serialized payload exceeded Langfuse's body-size limit (413), splits +// the batch in half and retries each half recursively instead of discarding +// every event in the batch on a single oversized flush. Only an event that +// still doesn't fit the limit on its own is ever dropped, and that is logged +// with its approximate size so the loss is visible instead of silent. +func (o *observer) flushWithSplit(ctx context.Context, batch []*api.IngestionEvent) error { + if len(batch) == 0 { + return nil + } + + err := o.flush(ctx, batch) + if err == nil { + return nil + } + + var apiErr *core.APIError + if !errors.As(err, &apiErr) || apiErr.StatusCode != http.StatusRequestEntityTooLarge { + // Not a body-size problem (network error, auth failure, a 5xx from + // Langfuse itself, ...) - splitting would not help here and would + // just multiply requests, so surface the error as-is. + return err + } + + if len(batch) == 1 { + logrus.WithContext(ctx).WithFields(logrus.Fields{ + "event_size_bytes": approxIngestionEventSize(batch[0]), + }).Error("dropping a single telemetry event that exceeds Langfuse's body size limit even alone") + return nil + } + + mid := len(batch) / 2 + errFirst := o.flushWithSplit(ctx, batch[:mid]) + errSecond := o.flushWithSplit(ctx, batch[mid:]) + return errors.Join(errFirst, errSecond) +} + +// approxIngestionEventSize returns the serialized size of a single event, or +// -1 if it can't be marshaled - used only for the drop-log message above, so +// a marshal failure there must not itself fail or panic. +func approxIngestionEventSize(event *api.IngestionEvent) int { + data, err := json.Marshal(event) + if err != nil { + return -1 + } + return len(data) +} + func (o *observer) sender() { batch := make([]*api.IngestionEvent, 0, o.queueSize) ticker := time.NewTicker(o.interval) @@ -207,11 +259,11 @@ func (o *observer) sender() { case <-o.ctx.Done(): return case ch := <-o.flusher: - ch <- o.flush(o.ctx, batch) + ch <- o.flushWithSplit(o.ctx, batch) batch = batch[:0] ticker.Reset(o.interval) case <-ticker.C: - if err := o.flush(o.ctx, batch); err != nil { + if err := o.flushWithSplit(o.ctx, batch); err != nil { logrus.WithContext(o.ctx).WithError(err).Error("failed to flush events by interval") } batch = batch[:0] @@ -219,7 +271,7 @@ func (o *observer) sender() { case event := <-o.queue: batch = append(batch, event) if len(batch) >= o.queueSize { - if err := o.flush(o.ctx, batch); err != nil { + if err := o.flushWithSplit(o.ctx, batch); err != nil { logrus.WithContext(o.ctx).WithError(err).Error("failed to flush events by queue size") } batch = batch[:0] diff --git a/backend/pkg/observability/langfuse/observer_test.go b/backend/pkg/observability/langfuse/observer_test.go index 95835470..5cda9950 100644 --- a/backend/pkg/observability/langfuse/observer_test.go +++ b/backend/pkg/observability/langfuse/observer_test.go @@ -2,8 +2,11 @@ package langfuse import ( "context" + "encoding/json" + "io" "net/http" "net/http/httptest" + "sync" "sync/atomic" "testing" "time" @@ -101,3 +104,168 @@ func TestObserverShutdownDropsBatch_ForceFlushDrains(t *testing.T) { _ = o.Shutdown(context.Background()) }) } + +// TestObserverFlushWithSplit_413SplitsInstead OfDroppingBatch reproduces the +// production failure: Langfuse rejects an oversized batch with 413 "Body +// exceeded ... limit". Previously the whole batch was discarded on any flush +// error; flushWithSplit must instead split the batch and retry each half so +// every event still lands, as long as no single event alone is oversized. +func TestObserverFlushWithSplit_413SplitsInsteadOfDroppingBatch(t *testing.T) { + const maxEventsPerRequest = 2 + + var ( + mu sync.Mutex + received = map[string]bool{} + postCount int32 + rejectCount int32 + ) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.Method == http.MethodGet { + _, _ = w.Write([]byte(`{"data":[{"id":"proj","name":"test"}]}`)) + return + } + + atomic.AddInt32(&postCount, 1) + + body, err := io.ReadAll(r.Body) + if err != nil { + t.Errorf("failed to read request body: %v", err) + w.WriteHeader(http.StatusInternalServerError) + return + } + var req api.IngestionBatchRequest + if err := json.Unmarshal(body, &req); err != nil { + t.Errorf("failed to unmarshal batch request: %v", err) + w.WriteHeader(http.StatusInternalServerError) + return + } + + if len(req.Batch) > maxEventsPerRequest { + atomic.AddInt32(&rejectCount, 1) + w.WriteHeader(http.StatusRequestEntityTooLarge) + _, _ = w.Write([]byte("Body exceeded 4.5mb limit")) + return + } + + mu.Lock() + for _, event := range req.Batch { + if zero := event.GetIngestionEventZero(); zero != nil { + received[zero.ID] = true + } + } + mu.Unlock() + + _, _ = w.Write([]byte(`{"successes":[],"errors":[]}`)) + })) + defer srv.Close() + + client, err := NewClient( + WithBaseURL(srv.URL), + WithPublicKey("pk"), + WithSecretKey("sk"), + WithProjectID("proj"), + ) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + + o := NewObserver(client, + WithSendInterval(10*time.Minute), + WithSendTimeout(2*time.Second), + WithQueueSize(100), + ).(*observer) + + const numEvents = 5 + ids := make([]string, 0, numEvents) + for i := 0; i < numEvents; i++ { + id := newSpanID() + ids = append(ids, id) + o.enqueue(&api.IngestionEvent{IngestionEventZero: &api.IngestionEventZero{ + ID: id, + Timestamp: getCurrentTimeString(), + }}) + } + + for i := 0; i < 400; i++ { + if len(o.queue) == 0 { + break + } + time.Sleep(5 * time.Millisecond) + } + + if err := o.ForceFlush(context.Background()); err != nil { + t.Fatalf("ForceFlush() error = %v, want nil (splitting should recover from 413)", err) + } + _ = o.Shutdown(context.Background()) + + if got := atomic.LoadInt32(&rejectCount); got == 0 { + t.Fatal("expected at least one 413 rejection to exercise the split path") + } + if got := atomic.LoadInt32(&postCount); got <= 1 { + t.Fatalf("expected multiple POST requests from splitting, got %d", got) + } + + mu.Lock() + defer mu.Unlock() + for _, id := range ids { + if !received[id] { + t.Errorf("event %s was never received by the sink - it was dropped instead of split", id) + } + } +} + +// TestObserverFlushWithSplit_OversizedSingleEventIsDroppedNotRetriedForever +// checks that when even a single event alone still exceeds the body-size +// limit, flushWithSplit gives up after splitting down to it (returning nil, +// logging the drop) instead of recursing forever. +func TestObserverFlushWithSplit_OversizedSingleEventIsDroppedNotRetriedForever(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.Method == http.MethodGet { + _, _ = w.Write([]byte(`{"data":[{"id":"proj","name":"test"}]}`)) + return + } + w.WriteHeader(http.StatusRequestEntityTooLarge) + _, _ = w.Write([]byte("Body exceeded 4.5mb limit")) + })) + defer srv.Close() + + client, err := NewClient( + WithBaseURL(srv.URL), + WithPublicKey("pk"), + WithSecretKey("sk"), + WithProjectID("proj"), + ) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + + o := NewObserver(client, + WithSendInterval(10*time.Minute), + WithSendTimeout(2*time.Second), + WithQueueSize(100), + ).(*observer) + + batch := []*api.IngestionEvent{ + {IngestionEventZero: &api.IngestionEventZero{ID: newSpanID(), Timestamp: getCurrentTimeString()}}, + {IngestionEventZero: &api.IngestionEventZero{ID: newSpanID(), Timestamp: getCurrentTimeString()}}, + } + + done := make(chan error, 1) + go func() { + done <- o.flushWithSplit(context.Background(), batch) + }() + + select { + case err := <-done: + if err != nil { + t.Fatalf("flushWithSplit() error = %v, want nil (an oversized single event must be dropped, not surfaced)", err) + } + case <-time.After(5 * time.Second): + t.Fatal("flushWithSplit did not return - it likely recursed forever on an always-413 sink") + } + + _ = o.Shutdown(context.Background()) +} diff --git a/backend/pkg/providers/tester/file_edit.go b/backend/pkg/providers/tester/file_edit.go new file mode 100644 index 00000000..f10b0117 --- /dev/null +++ b/backend/pkg/providers/tester/file_edit.go @@ -0,0 +1,238 @@ +package tester + +import ( + "encoding/json" + "fmt" + "strings" + "sync" + "time" + + "pentagi/pkg/providers/tester/testdata" + "pentagi/pkg/tools" + + "github.com/vxcontrol/langchaingo/llms" + "github.com/vxcontrol/langchaingo/llms/streaming" +) + +// The fixture fileEditTestCase's handler simulates: read_file for +// FileEditTestPath must return FileEditTestContent, and edit_file's diff +// must turn FileEditTestOldLine into FileEditTestNewLine somewhere in it. +// Exported so callers driving this scenario end to end through a mock +// provider (e.g. cmd/ctester's own tests) can build a matching response +// sequence without duplicating the fixture. +const ( + FileEditTestPath = "/work/report.txt" + FileEditTestContent = "Status: draft\nOwner: alice\nPriority: low\n" + FileEditTestOldLine = "Priority: low" + FileEditTestNewLine = "Priority: high" + + fileEditID = "file_read_then_edit" + fileEditTestName = "Read a file, then edit it via unified diff" +) + +// newFileEditTestCase builds a two-turn scenario exercising PentAGI's real +// file tool end to end: read a file, then change one line in it. Unlike +// every YAML-driven TestCase, this one is hand-built in Go (not tests.yml) +// because its whole point is a genuinely dynamic exchange - the harness +// doesn't know what the model will send until it sends it, so it can't be +// expressed as a fixed list of messages the way testdata.TestDefinition is. +// +// The tool declaration is not hand-copied JSON: it comes straight from +// tools.GetRegistryDefinitions(), the exact schema real PentAGI agents get, +// so this test also catches accidental schema regressions in pkg/tools. The +// read/edit "handler" below is this package's own in-memory simulation (no +// Docker container is involved), deliberately separate from - but for the +// actual diff-merge step, reusing - tools.ApplyUnifiedDiff, the same +// function terminal.EditFile calls in production. +func newFileEditTestCase() (testdata.TestCase, error) { + def, ok := tools.GetRegistryDefinitions()[tools.FileToolName] + if !ok { + return nil, fmt.Errorf("tools.GetRegistryDefinitions() has no definition for %q", tools.FileToolName) + } + defCopy := def + + prompt := fmt.Sprintf( + "Read the file at %q. Then change the line %q to exactly %q and save the file.", + FileEditTestPath, FileEditTestOldLine, FileEditTestNewLine, + ) + + return &fileEditTestCase{ + messages: []llms.MessageContent{llms.TextParts(llms.ChatMessageTypeHuman, prompt)}, + tools: []llms.Tool{{ + Type: "function", + Function: &defCopy, + }}, + }, nil +} + +// fileEditTestCase implements testdata.TestCase and testdata.MultiTurnTestCase. +type fileEditTestCase struct { + mu sync.Mutex + messages []llms.MessageContent + tools []llms.Tool + + readFileSeen bool // the model has made its first (expected read_file) call + editFileSeen bool // the model has made its second (expected edit_file) call + editApplied bool // edit_file's diff actually produced the requested change + failure string +} + +func (f *fileEditTestCase) ID() string { return fileEditID } +func (f *fileEditTestCase) Name() string { return fileEditTestName } +func (f *fileEditTestCase) Type() testdata.TestType { return testdata.TestTypeFileEdit } +func (f *fileEditTestCase) Group() testdata.TestGroup { return testdata.TestGroupAdvanced } +func (f *fileEditTestCase) Streaming() bool { return false } +func (f *fileEditTestCase) Prompt() string { return "" } +func (f *fileEditTestCase) Tools() []llms.Tool { return f.tools } +func (f *fileEditTestCase) Capability() testdata.TestCapability { return testdata.CapabilityNone } +func (f *fileEditTestCase) ExtraOptions() []llms.CallOption { return nil } +func (f *fileEditTestCase) StreamingCallback() streaming.Callback { return nil } + +func (f *fileEditTestCase) Messages() []llms.MessageContent { + f.mu.Lock() + defer f.mu.Unlock() + + out := make([]llms.MessageContent, len(f.messages)) + copy(out, f.messages) + return out +} + +// HandleToolResponse implements testdata.MultiTurnTestCase. It expects, +// across up to two calls, first a read_file call for FileEditTestPath +// (answered with FileEditTestContent) and then an edit_file call whose diff is applied +// in-memory via tools.ApplyUnifiedDiff; any other shape ends the exchange +// immediately with a recorded failure reason for Execute to report. +func (f *fileEditTestCase) HandleToolResponse(resp *llms.ContentResponse) bool { + f.mu.Lock() + defer f.mu.Unlock() + + call, args, ok := firstFileToolCall(resp) + if !ok { + f.failure = fmt.Sprintf("model did not call the %q tool", tools.FileToolName) + return false + } + + action, _ := args["action"].(string) + path, _ := args["path"].(string) + + if !f.readFileSeen { + f.readFileSeen = true + + if action != "" && action != string(tools.ReadFile) { + f.failure = fmt.Sprintf("expected the first call to be read_file, got action=%q", action) + return false + } + if path != FileEditTestPath { + f.failure = fmt.Sprintf("expected read_file to target %q, got %q", FileEditTestPath, path) + return false + } + + f.appendToolExchange(call, FileEditTestContent) + return true + } + + f.editFileSeen = true + + if action != string(tools.EditFile) { + f.failure = fmt.Sprintf("expected the second call to be edit_file, got action=%q", action) + f.appendToolExchange(call, "unexpected action for this scenario; ending the test") + return false + } + if path != FileEditTestPath { + f.failure = fmt.Sprintf("expected edit_file to target %q, got %q", FileEditTestPath, path) + return false + } + + diff, _ := args["diff"].(string) + newContent, _, err := tools.ApplyUnifiedDiff(FileEditTestContent, diff) + if err != nil { + f.failure = fmt.Sprintf("edit_file's diff did not apply: %v", err) + f.appendToolExchange(call, fmt.Sprintf("failed to apply diff: %v", err)) + return false + } + + f.editApplied = strings.Contains(newContent, FileEditTestNewLine) && !strings.Contains(newContent, FileEditTestOldLine) + if !f.editApplied { + f.failure = fmt.Sprintf("edit_file's diff applied but did not produce %q (result: %q)", FileEditTestNewLine, newContent) + } + f.appendToolExchange(call, fmt.Sprintf("Applied 1 diff hunk(s) to %s", FileEditTestPath)) + return false +} + +// appendToolExchange records the assistant's tool call and a synthesized +// tool result, so the next Messages() call reflects them for the model's +// next turn. Caller must hold f.mu. +func (f *fileEditTestCase) appendToolExchange(call llms.ToolCall, result string) { + f.messages = append(f.messages, + llms.MessageContent{ + Role: llms.ChatMessageTypeAI, + Parts: []llms.ContentPart{call}, + }, + llms.MessageContent{ + Role: llms.ChatMessageTypeTool, + Parts: []llms.ContentPart{ + llms.ToolCallResponse{ + ToolCallID: call.ID, + Name: call.FunctionCall.Name, + Content: result, + }, + }, + }, + ) +} + +// Execute implements testdata.TestCase. By the time the runner calls this, +// HandleToolResponse has already driven the exchange to completion (or to +// the point where it gave up); Execute only needs to report that outcome. +func (f *fileEditTestCase) Execute(response any, latency time.Duration) testdata.TestResult { + f.mu.Lock() + defer f.mu.Unlock() + + result := testdata.TestResult{ + ID: f.ID(), + Name: f.Name(), + Type: f.Type(), + Group: f.Group(), + Latency: latency, + } + + if _, ok := response.(*llms.ContentResponse); !ok { + result.Error = fmt.Errorf("expected *llms.ContentResponse, got %T", response) + return result + } + + switch { + case f.failure != "": + result.Error = fmt.Errorf("%s", f.failure) + case !f.readFileSeen: + result.Error = fmt.Errorf("model never called the %q tool", tools.FileToolName) + case !f.editFileSeen: + result.Error = fmt.Errorf("model called read_file but never followed up with edit_file") + case !f.editApplied: + result.Error = fmt.Errorf("edit_file was called but did not produce the requested change") + default: + result.Success = true + } + + return result +} + +// firstFileToolCall returns the first tool call in resp targeting PentAGI's +// file tool (tools.FileToolName) along with its decoded arguments. +func firstFileToolCall(resp *llms.ContentResponse) (llms.ToolCall, map[string]any, bool) { + for _, choice := range resp.Choices { + for _, call := range choice.ToolCalls { + if call.FunctionCall == nil || call.FunctionCall.Name != tools.FileToolName { + continue + } + + var args map[string]any + if err := json.Unmarshal([]byte(call.FunctionCall.Arguments), &args); err != nil { + continue + } + + return call, args, true + } + } + return llms.ToolCall{}, nil, false +} diff --git a/backend/pkg/providers/tester/file_edit_test.go b/backend/pkg/providers/tester/file_edit_test.go new file mode 100644 index 00000000..8e7b7eff --- /dev/null +++ b/backend/pkg/providers/tester/file_edit_test.go @@ -0,0 +1,333 @@ +package tester + +import ( + "encoding/json" + "strings" + "testing" + "time" + + "pentagi/pkg/providers/pconfig" + "pentagi/pkg/providers/provider" + "pentagi/pkg/providers/tester/mock" + "pentagi/pkg/providers/tester/testdata" + "pentagi/pkg/tools" + + "github.com/vxcontrol/langchaingo/llms" +) + +// fileToolCallResponse builds a *llms.ContentResponse whose single choice +// calls PentAGI's file tool with the given arguments (empty strings are +// omitted, mirroring how a real model would only send fields it needs). +func fileToolCallResponse(id string, args map[string]string) *llms.ContentResponse { + argMap := make(map[string]any, len(args)) + for k, v := range args { + if v != "" { + argMap[k] = v + } + } + argsJSON, _ := json.Marshal(argMap) + + return &llms.ContentResponse{ + Choices: []*llms.ContentChoice{{ + ToolCalls: []llms.ToolCall{{ + ID: id, + Type: "function", + FunctionCall: &llms.FunctionCall{ + Name: tools.FileToolName, + Arguments: string(argsJSON), + }, + }}, + }}, + } +} + +// correctFileEditDiff is a unified diff that turns FileEditTestOldLine into +// FileEditTestNewLine in FileEditTestContent (Priority is its 3rd line). +const correctFileEditDiff = "@@ -3,1 +3,1 @@\n-" + FileEditTestOldLine + "\n+" + FileEditTestNewLine + "\n" + +func TestNewFileEditTestCase(t *testing.T) { + t.Parallel() + + tc, err := newFileEditTestCase() + if err != nil { + t.Fatalf("newFileEditTestCase() error = %v", err) + } + + if tc.ID() == "" || tc.Name() == "" { + t.Error("expected non-empty ID and Name") + } + if tc.Group() != testdata.TestGroupAdvanced { + t.Errorf("Group() = %q, want %q", tc.Group(), testdata.TestGroupAdvanced) + } + if tc.Type() != testdata.TestTypeFileEdit { + t.Errorf("Type() = %q, want %q", tc.Type(), testdata.TestTypeFileEdit) + } + if tc.Capability() != testdata.CapabilityNone { + t.Errorf("Capability() = %q, want CapabilityNone", tc.Capability()) + } + + // the tool declaration must be PentAGI's real one, not a hand-copied schema + toolList := tc.Tools() + if len(toolList) != 1 || toolList[0].Function == nil || toolList[0].Function.Name != tools.FileToolName { + t.Fatalf("expected exactly the %q tool declaration, got %+v", tools.FileToolName, toolList) + } + realDef := tools.GetRegistryDefinitions()[tools.FileToolName] + if toolList[0].Function.Description != realDef.Description { + t.Error("tool description does not match tools.GetRegistryDefinitions() - declarations must come from pkg/tools") + } + + if len(tc.Messages()) != 1 { + t.Fatalf("expected exactly 1 initial message, got %d", len(tc.Messages())) + } + + if _, ok := tc.(testdata.MultiTurnTestCase); !ok { + t.Fatal("fileEditTestCase must implement testdata.MultiTurnTestCase") + } +} + +// TestFileEditTestCase_HandleToolResponse covers the read_file -> edit_file +// exchange with both a positive (correct sequence and diff) and negative +// (wrong tool, wrong action, wrong path, wrong order, bad diff, diff that +// applies but doesn't produce the requested change) cases. +func TestFileEditTestCase_HandleToolResponse(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + firstResp *llms.ContentResponse + secondResp *llms.ContentResponse // nil: the exchange must end after firstResp + wantSuccess bool + wantErr string + }{ + { + name: "happy path: read_file then edit_file with a correct diff", + firstResp: fileToolCallResponse("c1", map[string]string{"action": "read_file", "path": FileEditTestPath}), + secondResp: fileToolCallResponse("c2", map[string]string{"action": "edit_file", "path": FileEditTestPath, "diff": correctFileEditDiff}), + wantSuccess: true, + }, + { + name: "action omitted on first call still infers read_file", + firstResp: fileToolCallResponse("c1", map[string]string{"path": FileEditTestPath}), + secondResp: fileToolCallResponse("c2", map[string]string{"action": "edit_file", "path": FileEditTestPath, "diff": correctFileEditDiff}), + wantSuccess: true, + }, + { + name: "no tool call at all", + firstResp: &llms.ContentResponse{Choices: []*llms.ContentChoice{{Content: "I can't help with that."}}}, + wantErr: `did not call the "file" tool`, + }, + { + name: "calls a different tool entirely", + firstResp: &llms.ContentResponse{Choices: []*llms.ContentChoice{{ToolCalls: []llms.ToolCall{{FunctionCall: &llms.FunctionCall{Name: "terminal", Arguments: `{}`}}}}}}, + wantErr: `did not call the "file" tool`, + }, + { + name: "first call is write_file instead of read_file", + firstResp: fileToolCallResponse("c1", map[string]string{"action": "write_file", "path": FileEditTestPath, "content": "whatever"}), + wantErr: "expected the first call to be read_file", + }, + { + name: "read_file targets the wrong path", + firstResp: fileToolCallResponse("c1", map[string]string{"action": "read_file", "path": "/etc/passwd"}), + wantErr: "expected read_file to target", + }, + { + name: "second call is write_file instead of edit_file", + firstResp: fileToolCallResponse("c1", map[string]string{"action": "read_file", "path": FileEditTestPath}), + secondResp: fileToolCallResponse("c2", map[string]string{"action": "write_file", "path": FileEditTestPath, "content": "whatever"}), + wantErr: "expected the second call to be edit_file", + }, + { + name: "edit_file targets the wrong path", + firstResp: fileToolCallResponse("c1", map[string]string{"action": "read_file", "path": FileEditTestPath}), + secondResp: fileToolCallResponse("c2", map[string]string{"action": "edit_file", "path": "/etc/passwd", "diff": correctFileEditDiff}), + wantErr: "expected edit_file to target", + }, + { + name: "edit_file diff doesn't match the file content", + firstResp: fileToolCallResponse("c1", map[string]string{"action": "read_file", "path": FileEditTestPath}), + secondResp: fileToolCallResponse("c2", map[string]string{"action": "edit_file", "path": FileEditTestPath, "diff": "@@ -1,1 +1,1 @@\n-this line does not exist\n+replacement\n"}), + wantErr: "did not apply", + }, + { + name: "edit_file diff applies but produces the wrong content", + firstResp: fileToolCallResponse("c1", map[string]string{"action": "read_file", "path": FileEditTestPath}), + secondResp: fileToolCallResponse("c2", map[string]string{"action": "edit_file", "path": FileEditTestPath, "diff": "@@ -1,1 +1,1 @@\n-Status: draft\n+Status: final\n"}), + wantErr: "did not produce", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + tc, err := newFileEditTestCase() + if err != nil { + t.Fatalf("newFileEditTestCase() error = %v", err) + } + mt := tc.(testdata.MultiTurnTestCase) + + finalResp := tt.firstResp + hasMore := mt.HandleToolResponse(tt.firstResp) + + if tt.secondResp != nil { + if !hasMore { + t.Fatal("expected HandleToolResponse to request a second round after the first response") + } + finalResp = tt.secondResp + hasMore = mt.HandleToolResponse(tt.secondResp) + } + if hasMore { + t.Fatal("expected the exchange to end, but HandleToolResponse asked for another round") + } + + result := tc.Execute(finalResp, 5*time.Millisecond) + + if tt.wantSuccess { + if !result.Success { + t.Fatalf("expected success, got error: %v", result.Error) + } + return + } + + if result.Success { + t.Fatal("expected failure, got success") + } + if result.Error == nil || !strings.Contains(result.Error.Error(), tt.wantErr) { + t.Fatalf("error = %v, want it to contain %q", result.Error, tt.wantErr) + } + }) + } +} + +// TestFileEditTestCase_MessagesAccumulate checks that each answered tool +// call is appended to the conversation, so the model's next Messages() call +// actually sees the growing history (assistant tool-call + tool result per turn). +func TestFileEditTestCase_MessagesAccumulate(t *testing.T) { + t.Parallel() + + tc, err := newFileEditTestCase() + if err != nil { + t.Fatalf("newFileEditTestCase() error = %v", err) + } + mt := tc.(testdata.MultiTurnTestCase) + + if got := len(tc.Messages()); got != 1 { + t.Fatalf("expected 1 message before any turn, got %d", got) + } + + mt.HandleToolResponse(fileToolCallResponse("c1", map[string]string{"action": "read_file", "path": FileEditTestPath})) + if got := len(tc.Messages()); got != 3 { + t.Fatalf("expected 3 messages after round 1 (initial + assistant tool-call + tool result), got %d", got) + } + + mt.HandleToolResponse(fileToolCallResponse("c2", map[string]string{"action": "edit_file", "path": FileEditTestPath, "diff": correctFileEditDiff})) + if got := len(tc.Messages()); got != 5 { + t.Fatalf("expected 5 messages after round 2, got %d", got) + } +} + +// TestFileEditTestCase_Integration_ViaExecuteTest drives the scenario +// through tester.executeTest itself (not just the TestCase's own methods), +// using a mock provider configured to answer in sequence - proving the +// runner's multi-turn loop (see executeTest's MultiTurnTestCase branch) +// actually re-calls the provider with the growing conversation and stops at +// the right point. +func TestFileEditTestCase_Integration_ViaExecuteTest(t *testing.T) { + tc, err := newFileEditTestCase() + if err != nil { + t.Fatalf("newFileEditTestCase() error = %v", err) + } + + mockProvider := mock.NewProvider(provider.ProviderCustom, provider.DefaultProviderNameCustom, "test-model") + mockProvider.SetSequentialResponses( + fileToolCallResponse("call_1", map[string]string{"action": "read_file", "path": FileEditTestPath}), + fileToolCallResponse("call_2", map[string]string{"action": "edit_file", "path": FileEditTestPath, "diff": correctFileEditDiff}), + ) + + result, err := executeTest(t.Context(), testRequest{ + agentType: pconfig.OptionsTypePrimaryAgent, + testCase: tc, + provider: mockProvider, + }) + if err != nil { + t.Fatalf("executeTest() error = %v", err) + } + if !result.Success { + t.Fatalf("expected success, got error: %v", result.Error) + } + if result.Type != testdata.TestTypeFileEdit { + t.Errorf("result.Type = %q, want %q", result.Type, testdata.TestTypeFileEdit) + } +} + +// TestCollectTestRequests_FileEditTest_GetsIndependentInstancePerAgent is a +// regression test for a real bug found by running ctester live against +// Ollama: collectTestRequests used to build ONE fileEditTestCase per group +// and reuse that same pointer for every agentType's testRequest. Every other +// TestCase is immutable once built, so sharing is harmless there, but +// fileEditTestCase mutates itself turn by turn (see HandleToolResponse) - +// sharing it meant one agent's conversation, and even its pass/fail outcome +// (f.failure persists until overwritten), leaked into the next agent's +// "independent" run. Each agentType must get its own instance. +func TestCollectTestRequests_FileEditTest_GetsIndependentInstancePerAgent(t *testing.T) { + emptyRegistry, err := testdata.LoadRegistryFromYAML([]byte("[]")) + if err != nil { + t.Fatalf("LoadRegistryFromYAML() error = %v", err) + } + + config := &testConfig{ + agentTypes: []pconfig.ProviderOptionsType{pconfig.OptionsTypePrimaryAgent, pconfig.OptionsTypeCoder, pconfig.OptionsTypePentester}, + groups: []testdata.TestGroup{testdata.TestGroupAdvanced}, + customRegistry: emptyRegistry, + } + + requests := collectTestRequests(emptyRegistry, mock.NewProvider(provider.ProviderCustom, provider.DefaultProviderNameCustom, "test-model"), config) + + if len(requests) != len(config.agentTypes) { + t.Fatalf("expected %d file_edit requests (one per agent type), got %d", len(config.agentTypes), len(requests)) + } + + seen := make(map[testdata.TestCase]pconfig.ProviderOptionsType, len(requests)) + for _, req := range requests { + if req.testCase.Type() != testdata.TestTypeFileEdit { + t.Fatalf("expected only file_edit requests, got type %q", req.testCase.Type()) + } + if owner, dup := seen[req.testCase]; dup { + t.Fatalf("agent types %q and %q were given the SAME fileEditTestCase instance - state from one run would leak into the other", owner, req.agentType) + } + seen[req.testCase] = req.agentType + + if _, ok := req.testCase.(testdata.MultiTurnTestCase); !ok { + t.Fatalf("file_edit test case for %q does not implement MultiTurnTestCase", req.agentType) + } + } +} + +// TestFileEditTestCase_Integration_StopsAfterOneRoundOnFailure checks that a +// bad first call doesn't spend a second provider request: the mock's +// sequential responses only ever get consumed once here. +func TestFileEditTestCase_Integration_StopsAfterOneRoundOnFailure(t *testing.T) { + tc, err := newFileEditTestCase() + if err != nil { + t.Fatalf("newFileEditTestCase() error = %v", err) + } + + mockProvider := mock.NewProvider(provider.ProviderCustom, provider.DefaultProviderNameCustom, "test-model") + mockProvider.SetSequentialResponses( + fileToolCallResponse("call_1", map[string]string{"action": "write_file", "path": FileEditTestPath, "content": "oops"}), + fileToolCallResponse("call_2", map[string]string{"action": "edit_file", "path": FileEditTestPath, "diff": correctFileEditDiff}), + ) + + result, err := executeTest(t.Context(), testRequest{ + agentType: pconfig.OptionsTypePrimaryAgent, + testCase: tc, + provider: mockProvider, + }) + if err != nil { + t.Fatalf("executeTest() error = %v", err) + } + if result.Success { + t.Fatal("expected failure since the model called write_file instead of read_file first") + } +} diff --git a/backend/pkg/providers/tester/mock/provider.go b/backend/pkg/providers/tester/mock/provider.go index d020a718..e2e8ca2e 100644 --- a/backend/pkg/providers/tester/mock/provider.go +++ b/backend/pkg/providers/tester/mock/provider.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "strings" + "sync/atomic" "time" "pentagi/pkg/providers/pconfig" @@ -20,17 +21,26 @@ type Provider struct { providerType provider.ProviderType providerName provider.ProviderName modelName string - responses map[string]interface{} // key -> response mapping + responses map[string]any // key -> response mapping defaultResp string streamingDelay time.Duration providerConfig *pconfig.ProviderConfig models pconfig.ModelsConfig + + // sequence, when set via SetSequentialResponses, makes CallWithTools + // ignore content-based matching and return each response strictly in + // call order instead. Needed for multi-turn tool-calling scenarios + // where a tool call is answered and the model is called again with the + // same TextContent (a tool response carries no text), so content-based + // matching alone can't tell the calls apart. + sequence []any + sequenceCalls atomic.Int32 } // ResponseConfig configures mock responses type ResponseConfig struct { - Key string // Request identifier (prompt/message content) - Response interface{} // Response (string, *llms.ContentResponse, or error) + Key string // Request identifier (prompt/message content) + Response any // Response (string, *llms.ContentResponse, or error) } // NewProvider creates a new mock provider @@ -39,7 +49,7 @@ func NewProvider(providerType provider.ProviderType, providerName provider.Provi providerType: providerType, providerName: providerName, modelName: modelName, - responses: make(map[string]interface{}), + responses: make(map[string]any), defaultResp: "Mock response", streamingDelay: time.Millisecond * 10, } @@ -57,6 +67,16 @@ func (p *Provider) SetDefaultResponse(response string) { p.defaultResp = response } +// SetSequentialResponses configures CallWithTools to return responses +// strictly in order, one per call, bypassing content-based matching +// entirely. The last response repeats for any call beyond len(responses). +// Each element is handled exactly like a ResponseConfig.Response value +// (string, *llms.ContentResponse, or error). +func (p *Provider) SetSequentialResponses(responses ...any) { + p.sequence = responses + p.sequenceCalls.Store(0) +} + // SetStreamingDelay configures delay between streaming chunks func (p *Provider) SetStreamingDelay(delay time.Duration) { p.streamingDelay = delay @@ -151,7 +171,7 @@ func (p *Provider) CallEx( content = strings.TrimSpace(content) // Look for response - var respInterface interface{} + var respInterface any if resp, ok := p.responses[content]; ok { respInterface = resp } else { @@ -184,6 +204,18 @@ func (p *Provider) CallWithTools( tools []llms.Tool, streamCb streaming.Callback, ) (*llms.ContentResponse, error) { + if p.sequence != nil { + idx := int(p.sequenceCalls.Add(1)) - 1 + if idx >= len(p.sequence) { + idx = len(p.sequence) - 1 // repeat the last configured response past the end + } + + if streamCb != nil { + return p.handleStreamingResponse(ctx, p.sequence[idx], streamCb) + } + return p.handleContentResponse(p.sequence[idx]) + } + // Extract content for matching var content string for _, msg := range chain { @@ -196,7 +228,7 @@ func (p *Provider) CallWithTools( content = strings.TrimSpace(content) // Look for tool-specific response - var respInterface interface{} + var respInterface any toolKey := fmt.Sprintf("tools:%s", content) if resp, ok := p.responses[toolKey]; ok { respInterface = resp @@ -264,7 +296,7 @@ func (p *Provider) CallWithExtraOptions( } content = strings.TrimSpace(content) - var respInterface interface{} + var respInterface any if resp, ok := p.responses[content]; ok { respInterface = resp } else { @@ -326,7 +358,7 @@ func (p *Provider) GetPriceInfo(opt pconfig.ProviderOptionsType) *pconfig.PriceI } // handleResponse processes different response types for Call method -func (p *Provider) handleResponse(resp interface{}) (string, error) { +func (p *Provider) handleResponse(resp any) (string, error) { switch r := resp.(type) { case string: return r, nil @@ -343,7 +375,7 @@ func (p *Provider) handleResponse(resp interface{}) (string, error) { } // handleContentResponse processes responses for CallEx/CallWithTools -func (p *Provider) handleContentResponse(resp interface{}) (*llms.ContentResponse, error) { +func (p *Provider) handleContentResponse(resp any) (*llms.ContentResponse, error) { switch r := resp.(type) { case error: return nil, r @@ -371,7 +403,7 @@ func (p *Provider) handleContentResponse(resp interface{}) (*llms.ContentRespons // handleStreamingResponse simulates streaming behavior func (p *Provider) handleStreamingResponse( ctx context.Context, - resp interface{}, + resp any, streamCb streaming.Callback, ) (*llms.ContentResponse, error) { contentResp, err := p.handleContentResponse(resp) diff --git a/backend/pkg/providers/tester/mock/provider_test.go b/backend/pkg/providers/tester/mock/provider_test.go new file mode 100644 index 00000000..631f5b70 --- /dev/null +++ b/backend/pkg/providers/tester/mock/provider_test.go @@ -0,0 +1,69 @@ +package mock + +import ( + "testing" + + "pentagi/pkg/providers/pconfig" + "pentagi/pkg/providers/provider" + + "github.com/vxcontrol/langchaingo/llms" +) + +// TestProvider_SetSequentialResponses covers the sequential-response mode +// CallWithTools needs for multi-turn scenarios (see tester.fileEditTestCase): +// each call gets the next configured response regardless of content, and the +// last configured response repeats for any call beyond the configured count. +func TestProvider_SetSequentialResponses(t *testing.T) { + t.Parallel() + + p := NewProvider(provider.ProviderCustom, provider.DefaultProviderNameCustom, "test-model") + p.SetSequentialResponses("first", "second") + + chain := []llms.MessageContent{llms.TextParts(llms.ChatMessageTypeHuman, "same prompt every time")} + tool := []llms.Tool{{Type: "function", Function: &llms.FunctionDefinition{Name: "noop"}}} + + tests := []struct { + name string + want string + }{ + {name: "first call returns the first response", want: "first"}, + {name: "second call returns the second response", want: "second"}, + {name: "third call repeats the last response", want: "second"}, + {name: "fourth call still repeats the last response", want: "second"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + resp, err := p.CallWithTools(t.Context(), pconfig.OptionsTypePrimaryAgent, chain, tool, nil) + if err != nil { + t.Fatalf("CallWithTools() error = %v", err) + } + if len(resp.Choices) == 0 || resp.Choices[0].Content != tt.want { + t.Fatalf("CallWithTools() content = %+v, want %q", resp.Choices, tt.want) + } + }) + } +} + +// TestProvider_SetSequentialResponses_IgnoresContentMatching checks that, +// once configured, sequential mode takes priority over the normal +// content-keyed response lookup (SetResponses) - the two are mutually +// exclusive per test, by design. +func TestProvider_SetSequentialResponses_IgnoresContentMatching(t *testing.T) { + t.Parallel() + + p := NewProvider(provider.ProviderCustom, provider.DefaultProviderNameCustom, "test-model") + p.SetResponses([]ResponseConfig{{Key: "same prompt", Response: "should not be used"}}) + p.SetSequentialResponses("sequential response") + + chain := []llms.MessageContent{llms.TextParts(llms.ChatMessageTypeHuman, "same prompt")} + tool := []llms.Tool{{Type: "function", Function: &llms.FunctionDefinition{Name: "noop"}}} + + resp, err := p.CallWithTools(t.Context(), pconfig.OptionsTypePrimaryAgent, chain, tool, nil) + if err != nil { + t.Fatalf("CallWithTools() error = %v", err) + } + if resp.Choices[0].Content != "sequential response" { + t.Fatalf("CallWithTools() content = %q, want %q (content-matching must not take priority)", resp.Choices[0].Content, "sequential response") + } +} diff --git a/backend/pkg/providers/tester/runner.go b/backend/pkg/providers/tester/runner.go index 8d14a91f..ba705cf5 100644 --- a/backend/pkg/providers/tester/runner.go +++ b/backend/pkg/providers/tester/runner.go @@ -97,13 +97,13 @@ func collectTestRequests(registry *testdata.TestRegistry, prv provider.Provider, continue } - // skip capability-gated tests (adaptive thinking, reasoning - // off, structured output) whose wire behavior this agent's - // ACTUAL config wouldn't produce in a real PentAGI flow — see - // capabilitySupported for why. - if !capabilitySupported(prv, agentType, testCase.Capability()) { - continue - } + // skip capability-gated tests (adaptive thinking, reasoning + // off, structured output) whose wire behavior this agent's + // ACTUAL config wouldn't produce in a real PentAGI flow — see + // capabilitySupported for why. + if !capabilitySupported(prv, agentType, testCase.Capability()) { + continue + } requests = append(requests, testRequest{ agentType: agentType, @@ -112,6 +112,50 @@ func collectTestRequests(registry *testdata.TestRegistry, prv provider.Provider, }) } } + + if group != testdata.TestGroupAdvanced { + continue + } + + // fileEditTestCase is hand-built in Go, not tests.yml (see + // newFileEditTestCase) - its whole point is a genuinely dynamic + // exchange that can't be expressed as a fixed message list, so it + // can't come from the YAML-driven registry like every other test + // here. It's added for TestGroupAdvanced so it still runs by default + // without a special opt-in. + // + // Unlike every YAML-driven TestCase above (immutable once built, so + // sharing one instance across every agentType request is harmless), + // fileEditTestCase mutates itself turn by turn (see + // HandleToolResponse): reusing a single instance across multiple + // agentType requests would leak one agent's conversation and outcome + // into the next agent's run. A fresh instance per agentType is + // required, not just a style preference. + for _, agentType := range config.agentTypes { + if len(agentFilter) > 0 && !agentFilter[agentType] { + continue + } + if !isTestCompatibleWithAgent(testdata.TestTypeFileEdit, agentType) { + continue + } + if !capabilitySupported(prv, agentType, testdata.CapabilityNone) { + continue + } + + fileEdit, err := newFileEditTestCase() + if err != nil { + if config.verbose { + log.Printf("Warning: failed to build file_edit test case: %v", err) + } + break // same failure for every agentType - no point retrying + } + + requests = append(requests, testRequest{ + agentType: agentType, + testCase: fileEdit, + provider: prv, + }) + } } return requests @@ -188,7 +232,7 @@ func testWorker(ctx context.Context, requests <-chan testRequest, responses chan func executeTest(ctx context.Context, req testRequest) (testdata.TestResult, error) { startTime := time.Now() - var response interface{} + var response any var err error extra := req.testCase.ExtraOptions() @@ -217,6 +261,31 @@ func executeTest(ctx context.Context, req testRequest) (testdata.TestResult, err req.testCase.Tools(), req.testCase.StreamingCallback(), ) + + // MultiTurnTestCase (e.g. fileEditTestCase's read_file -> edit_file + // exchange) answers each tool call itself and asks for another round + // by returning true; every other TestCase leaves this a no-op. + if err == nil { + if multiTurn, ok := req.testCase.(testdata.MultiTurnTestCase); ok { + for { + contentResp, isContentResp := response.(*llms.ContentResponse) + if !isContentResp || !multiTurn.HandleToolResponse(contentResp) { + break + } + + response, err = req.provider.CallWithTools( + ctx, + req.agentType, + req.testCase.Messages(), + req.testCase.Tools(), + req.testCase.StreamingCallback(), + ) + if err != nil { + break + } + } + } + } case len(req.testCase.Messages()) > 0: // messages without tools response, err = req.provider.CallEx( diff --git a/backend/pkg/providers/tester/testdata/models.go b/backend/pkg/providers/tester/testdata/models.go index fb3183fd..da44c562 100644 --- a/backend/pkg/providers/tester/testdata/models.go +++ b/backend/pkg/providers/tester/testdata/models.go @@ -16,6 +16,9 @@ const ( TestTypeCompletion TestType = "completion" TestTypeJSON TestType = "json" TestTypeTool TestType = "tool" + // TestTypeFileEdit is a MultiTurnTestCase: it isn't built from tests.yml + // (see tester.newFileEditTestCase), only used to label its TestResult. + TestTypeFileEdit TestType = "file_edit" ) type TestGroup string @@ -223,6 +226,24 @@ type TestCase interface { Execute(response any, latency time.Duration) TestResult } +// MultiTurnTestCase is an optional extension of TestCase for scenarios that +// need more than one round-trip to the provider before Execute can judge the +// outcome - e.g. a tool call whose result must be answered before the model +// makes its next call. The runner (tester.executeTest) detects it via a type +// assertion; every TestCase that doesn't implement it keeps going through +// the plain single-call path unmodified. +type MultiTurnTestCase interface { + TestCase + + // HandleToolResponse receives the latest provider response. If it + // recognizes something it needs to answer, it records the exchange + // internally (so the next Messages() call reflects it) and returns + // true, asking the runner for another round. Returning false ends the + // exchange: the runner calls Execute with this same response, exactly + // as it would for a plain TestCase. + HandleToolResponse(resp *llms.ContentResponse) bool +} + // TestSuite contains stateful test cases for execution type TestSuite struct { Group TestGroup diff --git a/backend/pkg/tools/args.go b/backend/pkg/tools/args.go index 0052b138..4f080b56 100644 --- a/backend/pkg/tools/args.go +++ b/backend/pkg/tools/args.go @@ -7,21 +7,36 @@ import ( "strings" ) -type FileOp string +// FileOp is a type alias (not a distinct type) for String: it shares the +// exact same lenient UnmarshalJSON, so every FileOp value - including the +// ReadFile/WriteFile/EditFile constants below - automatically benefits from +// the quote-unwrapping behaviour documented on String, with zero changes +// needed at any existing call site (switches, comparisons, assignments) +// elsewhere in the codebase. +type FileOp = String const ( ReadFile FileOp = "read_file" WriteFile FileOp = "write_file" + EditFile FileOp = "edit_file" ) +// FileAction's field descriptions are layered on purpose, each one adding +// only what the level above doesn't already say: the tool description (see +// registry.go) gives the read/write/edit decision; Action's enum values name +// which payload field each one consumes; Content/Diff each spec out the +// format of their own payload. Path/Message never vary by action, so they're +// described once, action-agnostically. type FileAction struct { - Action FileOp `json:"action" jsonschema:"required,enum=read_file,enum=write_file" jsonschema_description:"Action to perform with the code. 'read_file' - Returns the content of the file. 'write_file' - Writes or updates the content of the file"` - Content string `json:"content" jsonschema_description:"Content to write to the file (raw file content, not a localized message)"` - Path string `json:"path" jsonschema:"required" jsonschema_description:"Absolute path to the file to read or write"` - Message string `json:"message" jsonschema:"required,title=File action message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary describing what you are reading from or writing to the file and why. Written in the engagement language declared by your system prompt."` + Action FileOp `json:"action" jsonschema:"required,type=string,enum=read_file,enum=write_file,enum=edit_file" jsonschema_description:"'read_file' reads the file (no other field needed). 'write_file' overwrites it with 'content' (the whole file). 'edit_file' applies the patch in 'diff' (existing content elsewhere is untouched)."` + Content string `json:"content,omitempty" jsonschema_description:"write_file only: the complete new file content (not a diff, not a partial update)."` + Diff String `json:"diff,omitempty" jsonschema:"type=string" jsonschema_description:"edit_file only: unified-diff hunk(s) - '@@ -old +new @@' header, then ' '/'-'/'+' lines. Always keep at least one unchanged context line so the location is unambiguous; context/removed lines must match the file's current content verbatim (read_file first). Header line numbers are only a hint - the line text is what must match."` + Path String `json:"path" jsonschema:"required,type=string" jsonschema_description:"Absolute path to the file"` + Message string `json:"message" jsonschema:"required,title=File action message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary describing what you are reading, writing, or editing and why. Written in the engagement language declared by your system prompt."` } -type BrowserAction string +// BrowserAction is a type alias for String - see the FileOp comment above. +type BrowserAction = String const ( Markdown BrowserAction = "markdown" @@ -31,7 +46,7 @@ const ( type Browser struct { Url string `json:"url" jsonschema:"required" jsonschema_description:"URL to open in the browser"` - Action BrowserAction `json:"action" jsonschema:"required,enum=markdown,enum=html,enum=links" jsonschema_description:"Action to perform in the browser. 'markdown' - Returns the content of the page in markdown format. 'html' - Returns the content of the page in html format. 'links' - Get the list of all URLs on the page to be used in later calls (e.g., open search results after the initial search lookup)."` + Action BrowserAction `json:"action" jsonschema:"required,type=string,enum=markdown,enum=html,enum=links" jsonschema_description:"Action to perform in the browser. 'markdown' - Returns the content of the page in markdown format. 'html' - Returns the content of the page in html format. 'links' - Get the list of all URLs on the page to be used in later calls (e.g., open search results after the initial search lookup)."` Message string `json:"message" jsonschema:"required,title=Browser action message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary describing what content you are fetching, in which format, and why. Written in the engagement language declared by your system prompt."` } @@ -45,8 +60,9 @@ type SubtaskList struct { Message string `json:"message" jsonschema:"required,title=Subtask generation result" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary on the generation result and the main goal of the plan. Written in the engagement language declared by your system prompt."` } -// SubtaskOperationType defines the type of operation to perform on a subtask -type SubtaskOperationType string +// SubtaskOperationType defines the type of operation to perform on a subtask. +// It is a type alias for String - see the FileOp comment above. +type SubtaskOperationType = String const ( SubtaskOpAdd SubtaskOperationType = "add" @@ -57,7 +73,7 @@ const ( // SubtaskOperation defines a single operation on the subtask list for delta-based refinement type SubtaskOperation struct { - Op SubtaskOperationType `json:"op" jsonschema:"required,enum=add,enum=remove,enum=modify,enum=reorder" jsonschema_description:"Operation type: 'add' creates a new subtask, 'remove' deletes a subtask by ID, 'modify' updates title/description of existing subtask, 'reorder' moves a subtask to a different position"` + Op SubtaskOperationType `json:"op" jsonschema:"required,type=string,enum=add,enum=remove,enum=modify,enum=reorder" jsonschema_description:"Operation type: 'add' creates a new subtask, 'remove' deletes a subtask by ID, 'modify' updates title/description of existing subtask, 'reorder' moves a subtask to a different position"` ID *int64 `json:"id,omitempty" jsonschema:"title=Subtask ID" jsonschema_description:"ID of existing subtask (required for remove/modify/reorder operations)"` AfterID *int64 `json:"after_id,omitempty" jsonschema:"title=Insert after ID" jsonschema_description:"For add/reorder: insert after this subtask ID (null/0 = insert at beginning)"` Title string `json:"title,omitempty" jsonschema:"title=New title" jsonschema_description:"Engagement-log plan entry — new subtask title (required for add, optional for modify). Written in the engagement language declared by your system prompt."` @@ -124,14 +140,14 @@ type SearchResult struct { type SploitusAction struct { Query string `json:"query" jsonschema:"required" jsonschema_description:"Technical-channel payload — search query for Sploitus (e.g. 'ssh', 'apache 2.4', 'CVE-2021-44228'). ALWAYS written in English; the Sploitus index is English-only. Short and precise queries return the best results."` - ExploitType string `json:"exploit_type,omitempty" jsonschema:"enum=exploits,enum=tools" jsonschema_description:"What to search for: 'exploits' (default) for exploit code and PoCs, 'tools' for offensive security tools"` - Sort string `json:"sort,omitempty" jsonschema:"enum=default,enum=date,enum=score" jsonschema_description:"Result ordering: 'default' (relevance), 'date' (newest first), 'score' (highest CVSS first)"` + ExploitType String `json:"exploit_type,omitempty" jsonschema:"type=string,enum=exploits,enum=tools" jsonschema_description:"What to search for: 'exploits' (default) for exploit code and PoCs, 'tools' for offensive security tools"` + Sort String `json:"sort,omitempty" jsonschema:"type=string,enum=default,enum=date,enum=score" jsonschema_description:"Result ordering: 'default' (relevance), 'date' (newest first), 'score' (highest CVSS first)"` MaxResults Int64 `json:"max_results" jsonschema:"required,type=integer" jsonschema_description:"Maximum number of results to return (minimum 1; maximum 25; default 10)"` Message string `json:"message" jsonschema:"required,title=Search query message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary explaining the expected result and how it advances the goal. Written in the engagement language declared by your system prompt."` } type GraphitiSearchAction struct { - SearchType string `json:"search_type" jsonschema:"required,enum=temporal_window,enum=entity_relationships,enum=diverse_results,enum=episode_context,enum=successful_tools,enum=recent_context,enum=entity_by_label" jsonschema_description:"Type of search to perform: temporal_window (time-bounded search), entity_relationships (graph traversal from an entity), diverse_results (anti-redundancy search), episode_context (full agent reasoning and tool outputs), successful_tools (proven techniques), recent_context (latest findings), entity_by_label (type-specific entity search)"` + SearchType String `json:"search_type" jsonschema:"required,type=string,enum=temporal_window,enum=entity_relationships,enum=diverse_results,enum=episode_context,enum=successful_tools,enum=recent_context,enum=entity_by_label" jsonschema_description:"Type of search to perform: temporal_window (time-bounded search), entity_relationships (graph traversal from an entity), diverse_results (anti-redundancy search), episode_context (full agent reasoning and tool outputs), successful_tools (proven techniques), recent_context (latest findings), entity_by_label (type-specific entity search)"` Query string `json:"query" jsonschema:"required" jsonschema_description:"Technical-channel payload — natural language query against the team's temporal knowledge graph. ALWAYS written in English regardless of the engagement language: the graph is indexed in English and shared across all engagements; non-English queries will fail to retrieve stored episodic memory."` MaxResults *Int64 `json:"max_results,omitempty" jsonschema:"title=Maximum Results,type=integer" jsonschema_description:"Maximum number of results to return (default varies by search type)"` TimeStart string `json:"time_start,omitempty" jsonschema_description:"Start of time window (ISO 8601 format, required for temporal_window)"` @@ -140,9 +156,9 @@ type GraphitiSearchAction struct { MaxDepth *Int64 `json:"max_depth,omitempty" jsonschema:"title=Maximum Depth,type=integer" jsonschema_description:"Maximum graph traversal depth (default: 2, max: 3, for entity_relationships)"` NodeLabels []string `json:"node_labels,omitempty" jsonschema_description:"Filter to specific node types (e.g., ['IP_ADDRESS', 'SERVICE', 'VULNERABILITY'])"` EdgeTypes []string `json:"edge_types,omitempty" jsonschema_description:"Filter to specific relationship types (e.g., ['HAS_PORT', 'EXPLOITS'])"` - DiversityLevel string `json:"diversity_level,omitempty" jsonschema:"enum=low,enum=medium,enum=high" jsonschema_description:"How much diversity to prioritize (default: medium, for diverse_results)"` + DiversityLevel String `json:"diversity_level,omitempty" jsonschema:"type=string,enum=low,enum=medium,enum=high" jsonschema_description:"How much diversity to prioritize (default: medium, for diverse_results)"` MinMentions *Int64 `json:"min_mentions,omitempty" jsonschema:"title=Minimum Mentions,type=integer" jsonschema_description:"Minimum episode mentions (default: 2, for successful_tools)"` - RecencyWindow string `json:"recency_window,omitempty" jsonschema:"enum=1h,enum=6h,enum=24h,enum=7d" jsonschema_description:"How far back to search (default: 24h, for recent_context)"` + RecencyWindow String `json:"recency_window,omitempty" jsonschema:"type=string,enum=1h,enum=6h,enum=24h,enum=7d" jsonschema_description:"How far back to search (default: 24h, for recent_context)"` Message string `json:"message" jsonschema:"required,title=Search message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary summarizing the search query and expected results. Written in the engagement language declared by your system prompt."` } @@ -172,27 +188,27 @@ type SearchInMemoryAction struct { type SearchGuideAction struct { Questions Strings `json:"questions" jsonschema:"required,type=array,minItems=1,maxItems=5" jsonschema_description:"Technical-channel payload — 1 to 5 detailed, context-rich semantic queries for the team's guide vector store. Must be a real JSON array of strings, e.g. [\"query 1\",\"query 2\"] - NOT a JSON-encoded string containing an array. ALWAYS written in English regardless of the engagement language: the store is indexed in English and shared across all engagements, so non-English queries will fail to retrieve relevant guides. Each query should include scenario context, objectives, and specific intent. Note: The 'Type' field acts as a strict filter."` - Type string `json:"type" jsonschema:"required,enum=install,enum=configure,enum=use,enum=pentest,enum=development,enum=other" jsonschema_description:"The specific type of guide you need. This required field acts as a strict filter to enhance the relevance of search results by narrowing down the scope to the specified guide type."` + Type String `json:"type" jsonschema:"required,type=string,enum=install,enum=configure,enum=use,enum=pentest,enum=development,enum=other" jsonschema_description:"The specific type of guide you need. This required field acts as a strict filter to enhance the relevance of search results by narrowing down the scope to the specified guide type."` Message string `json:"message" jsonschema:"required,title=Guide search message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary summarizing the queries and the type of guide needed. Written in the engagement language declared by your system prompt."` } type StoreGuideAction struct { Guide string `json:"guide" jsonschema:"required" jsonschema_description:"Technical-channel payload — ready guide in markdown format that will be stored in the team's vector store for future retrieval. ALWAYS written in English regardless of the engagement language: the store is indexed in English and shared across all engagements; non-English content becomes unreachable to future searches. Anonymize all sensitive data (IPs, domains, credentials, paths) using descriptive placeholders."` Question string `json:"question" jsonschema:"required" jsonschema_description:"Technical-channel payload — question that was used to prepare this guide; co-indexed with the guide. Always written in English; never translated."` - Type string `json:"type" jsonschema:"required,enum=install,enum=configure,enum=use,enum=pentest,enum=development,enum=other" jsonschema_description:"Type of the guide to store; it will be used as a hard filter for search"` + Type String `json:"type" jsonschema:"required,type=string,enum=install,enum=configure,enum=use,enum=pentest,enum=development,enum=other" jsonschema_description:"Type of the guide to store; it will be used as a hard filter for search"` Message string `json:"message" jsonschema:"required,title=Store guide message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary summarizing the guide. Written in the engagement language declared by your system prompt."` } type SearchAnswerAction struct { Questions Strings `json:"questions" jsonschema:"required,type=array,minItems=1,maxItems=5" jsonschema_description:"Technical-channel payload — 1 to 5 detailed, context-rich semantic queries for the team's answer vector store. Must be a real JSON array of strings, e.g. [\"query 1\",\"query 2\"] - NOT a JSON-encoded string containing an array. ALWAYS written in English regardless of the engagement language: the store is indexed in English and shared across all engagements, so non-English queries will fail to retrieve relevant answers. Each query should include the context, what you want to find, what you intend to do with the information, and why you need it. Note: The 'Type' field acts as a strict filter."` - Type string `json:"type" jsonschema:"required,enum=guide,enum=vulnerability,enum=code,enum=tool,enum=other" jsonschema_description:"The specific type of information or answer you are seeking. This required field acts as a strict filter to enhance the relevance of search results by narrowing down the scope to the specified type."` + Type String `json:"type" jsonschema:"required,type=string,enum=guide,enum=vulnerability,enum=code,enum=tool,enum=other" jsonschema_description:"The specific type of information or answer you are seeking. This required field acts as a strict filter to enhance the relevance of search results by narrowing down the scope to the specified type."` Message string `json:"message" jsonschema:"required,title=Answer search message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary summarizing the queries and the type of answer needed. Written in the engagement language declared by your system prompt."` } type StoreAnswerAction struct { Answer string `json:"answer" jsonschema:"required" jsonschema_description:"Technical-channel payload — ready answer in markdown format that will be stored in the team's vector store for future retrieval. ALWAYS written in English regardless of the engagement language: the store is indexed in English and shared across all engagements; non-English content becomes unreachable to future searches. Anonymize all sensitive data (IPs, domains, credentials) using descriptive placeholders."` Question string `json:"question" jsonschema:"required" jsonschema_description:"Technical-channel payload — question that was used to prepare this answer; co-indexed with the answer. Always written in English; never translated."` - Type string `json:"type" jsonschema:"required,enum=guide,enum=vulnerability,enum=code,enum=tool,enum=other" jsonschema_description:"Type of the search query and answer to store; it will be used as a hard filter for search"` + Type String `json:"type" jsonschema:"required,type=string,enum=guide,enum=vulnerability,enum=code,enum=tool,enum=other" jsonschema_description:"Type of the search query and answer to store; it will be used as a hard filter for search"` Message string `json:"message" jsonschema:"required,title=Store answer message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary summarizing the answer. Written in the engagement language declared by your system prompt."` } @@ -242,7 +258,8 @@ type HackResult struct { } // FlowStatusDetail controls the level of detail returned by get_flow_status. -type FlowStatusDetail string +// It is a type alias for String - see the FileOp comment above. +type FlowStatusDetail = String const ( FlowStatusDetailSummary FlowStatusDetail = "summary" @@ -254,7 +271,7 @@ const ( // GetFlowStatusAction defines arguments for the get_flow_status tool. type GetFlowStatusAction struct { - Detail FlowStatusDetail `json:"detail" jsonschema:"required,enum=summary,enum=tasks,enum=subtasks,enum=running,enum=planned" jsonschema_description:"Level of detail: 'summary' - flow health snapshot with status and counts; 'tasks' - all tasks with ID/status/title; 'subtasks' - all subtasks optionally filtered by task_id; 'running' - full Task→Subtask execution chain including task input and recent agent messages; 'planned' - only subtasks with status 'created' (not yet started), optionally filtered by task_id"` + Detail FlowStatusDetail `json:"detail" jsonschema:"required,type=string,enum=summary,enum=tasks,enum=subtasks,enum=running,enum=planned" jsonschema_description:"Level of detail: 'summary' - flow health snapshot with status and counts; 'tasks' - all tasks with ID/status/title; 'subtasks' - all subtasks optionally filtered by task_id; 'running' - full Task→Subtask execution chain including task input and recent agent messages; 'planned' - only subtasks with status 'created' (not yet started), optionally filtered by task_id"` TaskID *Int64 `json:"task_id,omitempty" jsonschema:"title=Task ID,type=integer" jsonschema_description:"Optional task ID filter. Applies to detail=subtasks and detail=planned to narrow results to a specific task."` Verbose Bool `json:"verbose,omitempty" jsonschema:"type=boolean" jsonschema_description:"Set to true for deeper investigation: includes descriptions, inputs, results, and execution context per entry; shows up to 50 recent agent messages instead of 10."` Message string `json:"message" jsonschema:"required,title=Status message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary describing what status information you are requesting. Written in the engagement language declared by your system prompt."` @@ -402,6 +419,66 @@ func (i *Int64) String() string { return strconv.FormatInt(int64(*i), 10) } +// String is a lenient string for LLM-generated tool-call arguments, in +// particular enum-like fields (action/type/search_type/... - anywhere a +// value is compared or switched on rather than read as free text). Models +// occasionally wrap the value in an extra, literal pair of quote characters +// - e.g. the JSON string "\"write_file\"" decodes via a plain string +// UnmarshalJSON into the 12-character Go string `"write_file"` (quotes +// included), which then matches no known enum value, fails validation, and +// forces an unnecessary tool-call-fixer round-trip. UnmarshalJSON unwraps +// that redundant quoting instead of failing outright. +// +// FileOp, BrowserAction, SubtaskOperationType, and FlowStatusDetail are type +// aliases for String (not distinct types), so every field declared with one +// of those names gets this leniency automatically, with zero changes needed +// at any existing call site: switches, `==` comparisons, and assignments +// against their constants all keep working exactly as before, because the +// alias makes them the exact same type as String, not merely a similar one. +type String string + +func (s *String) UnmarshalJSON(data []byte) error { + // A bare JSON "null" unmarshals into an empty string with no error by + // default, which would silently mask a required field being omitted - + // treat it as invalid instead, consistent with Bool/Int64/Strings above. + if trimmed := strings.TrimSpace(string(data)); trimmed == "null" { + return fmt.Errorf("invalid string value: got null") + } + + var raw string + if err := json.Unmarshal(data, &raw); err != nil { + return err + } + *s = String(unwrapRedundantQuotes(raw)) + return nil +} + +func (s String) MarshalJSON() ([]byte, error) { + return json.Marshal(string(s)) +} + +// String implements fmt.Stringer, so %s/%v formatting and logrus fields show +// the unwrapped value, and callers that need a plain string (map keys, +// strings.* helpers, external struct fields, typed casts) can call it explicitly. +func (s String) String() string { + return string(s) +} + +// unwrapRedundantQuotes strips at most a few layers of a matching leading+ +// trailing '"' pair from a decoded JSON string value. Bounded to guard +// against pathological input; the observed corruption is a single extra +// layer, but a small bound costs nothing and covers repeated wrapping too. +func unwrapRedundantQuotes(s string) string { + for range 3 { + trimmed := strings.TrimSpace(s) + if len(trimmed) < 2 || trimmed[0] != '"' || trimmed[len(trimmed)-1] != '"' { + return s + } + s = trimmed[1 : len(trimmed)-1] + } + return s +} + // Strings is a lenient []string for LLM-generated tool-call arguments (e.g. // the "questions" field of vector-store search tools). Models occasionally // double-encode the array as a JSON string containing the array literal diff --git a/backend/pkg/tools/args_test.go b/backend/pkg/tools/args_test.go index 1ff2070d..c7a68458 100644 --- a/backend/pkg/tools/args_test.go +++ b/backend/pkg/tools/args_test.go @@ -2,6 +2,7 @@ package tools import ( "encoding/json" + "fmt" "testing" ) @@ -744,3 +745,197 @@ func TestSearchInMemoryAction_QuestionsDoubleEncoded(t *testing.T) { t.Errorf("Questions[0] = %q, unexpected value", action.Questions[0]) } } + +func TestStringUnmarshalJSON(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + input string + want string + wantErr bool + }{ + {name: "plain string", input: `"write_file"`, want: "write_file"}, + {name: "empty string", input: `""`, want: ""}, + { + name: "single extra-quoted (production bug case)", + input: `"\"write_file\""`, + want: "write_file", + }, + { + name: "single extra-quoted with whitespace", + input: `" \"read_file\" "`, + want: "read_file", + }, + {name: "number is an error", input: `42`, wantErr: true}, + {name: "null is an error", input: `null`, wantErr: true}, + {name: "bool is an error", input: `true`, wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var s String + err := s.UnmarshalJSON([]byte(tt.input)) + if (err != nil) != tt.wantErr { + t.Errorf("UnmarshalJSON(%s) error = %v, wantErr %v", tt.input, err, tt.wantErr) + return + } + if !tt.wantErr && string(s) != tt.want { + t.Errorf("UnmarshalJSON(%s) = %q, want %q", tt.input, string(s), tt.want) + } + }) + } +} + +// TestStringUnmarshalJSON_DoubleWrapped checks a value wrapped in two extra +// bare quote pairs (the JSON text `"\"\"markdown\"\""` decodes, via a single +// json.Unmarshal into a plain string, to the 12-character Go string +// `""markdown""`). The observed production bug is a single extra layer; +// unwrapRedundantQuotes is bounded to a few iterations defensively, so this +// checks it also handles two. +func TestStringUnmarshalJSON_DoubleWrapped(t *testing.T) { + t.Parallel() + + // Backtick (raw) string literal: no Go escape processing, so this is + // exactly the 18 literal JSON bytes intended, with no risk of the source + // escaping and the intended JSON escaping being conflated. + input := `"\"\"markdown\"\""` + + var s String + if err := s.UnmarshalJSON([]byte(input)); err != nil { + t.Fatalf("UnmarshalJSON(%s) unexpected error: %v", input, err) + } + if string(s) != "markdown" { + t.Errorf("UnmarshalJSON(%s) = %q, want %q", input, string(s), "markdown") + } +} + +func TestStringMarshalJSON(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + s String + want string + }{ + {name: "plain value", s: String("write_file"), want: `"write_file"`}, + {name: "empty value", s: String(""), want: `""`}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got, err := tt.s.MarshalJSON() + if err != nil { + t.Fatalf("MarshalJSON() unexpected error: %v", err) + } + if string(got) != tt.want { + t.Errorf("MarshalJSON() = %s, want %s", got, tt.want) + } + }) + } +} + +func TestStringStringerMethod(t *testing.T) { + t.Parallel() + + var s fmt.Stringer = String("write_file") + if s.String() != "write_file" { + t.Errorf("String() = %q, want %q", s.String(), "write_file") + } +} + +func TestStringJSONRoundTrip(t *testing.T) { + t.Parallel() + + type container struct { + Action String `json:"action"` + } + + data, err := json.Marshal(container{Action: String("write_file")}) + if err != nil { + t.Fatalf("Marshal() error = %v", err) + } + + var c2 container + if err := json.Unmarshal(data, &c2); err != nil { + t.Fatalf("round-trip Unmarshal() error = %v", err) + } + if c2.Action != "write_file" { + t.Errorf("round-trip Action = %q, want %q", c2.Action, "write_file") + } +} + +// TestFileAction_ActionAndPathExtraQuoted reproduces the exact production +// failure: the LLM sent both "action" and "path" wrapped in an extra literal +// pair of quotes (action="\"write_file\"", path="\"/some/path\""), which used +// to fail with "unknown file action" because the corrupted value matched no +// enum case. Both fields must now unmarshal cleanly via the String type. +func TestFileAction_ActionAndPathExtraQuoted(t *testing.T) { + t.Parallel() + + raw := `{"action": "\"write_file\"", "path": "\"/home/evidence/inject.py\"", "content": "print(1)", "message": "m"}` + + var action FileAction + if err := json.Unmarshal([]byte(raw), &action); err != nil { + t.Fatalf("Unmarshal() unexpected error: %v", err) + } + if action.Action != WriteFile { + t.Errorf("Action = %q, want %q", action.Action, WriteFile) + } + if action.Path.String() != "/home/evidence/inject.py" { + t.Errorf("Path = %q, want %q", action.Path.String(), "/home/evidence/inject.py") + } +} + +// TestFileAction_ReadActionExtraQuoted mirrors the read_file variant of the +// same production bug. +func TestFileAction_ReadActionExtraQuoted(t *testing.T) { + t.Parallel() + + raw := `{"action": "\"read_file\"", "path": "/home/evidence/inject.py", "message": "m"}` + + var action FileAction + if err := json.Unmarshal([]byte(raw), &action); err != nil { + t.Fatalf("Unmarshal() unexpected error: %v", err) + } + if action.Action != ReadFile { + t.Errorf("Action = %q, want %q", action.Action, ReadFile) + } +} + +// TestSubtaskOperation_OpExtraQuoted checks the same leniency for +// SubtaskOperationType (a String alias), confirming the fix generalizes to +// enum fields beyond the file/browser tools where it was first observed. +func TestSubtaskOperation_OpExtraQuoted(t *testing.T) { + t.Parallel() + + raw := `{"op": "\"add\"", "title": "t", "description": "d"}` + + var op SubtaskOperation + if err := json.Unmarshal([]byte(raw), &op); err != nil { + t.Fatalf("Unmarshal() unexpected error: %v", err) + } + if op.Op != SubtaskOpAdd { + t.Errorf("Op = %q, want %q", op.Op, SubtaskOpAdd) + } +} + +// TestGetFlowStatusAction_DetailExtraQuoted checks the same leniency for +// FlowStatusDetail (a String alias). +func TestGetFlowStatusAction_DetailExtraQuoted(t *testing.T) { + t.Parallel() + + raw := `{"detail": "\"summary\"", "message": "m"}` + + var action GetFlowStatusAction + if err := json.Unmarshal([]byte(raw), &action); err != nil { + t.Fatalf("Unmarshal() unexpected error: %v", err) + } + if action.Detail != FlowStatusDetailSummary { + t.Errorf("Detail = %q, want %q", action.Detail, FlowStatusDetailSummary) + } +} diff --git a/backend/pkg/tools/file_diff.go b/backend/pkg/tools/file_diff.go new file mode 100644 index 00000000..9e42e0d1 --- /dev/null +++ b/backend/pkg/tools/file_diff.go @@ -0,0 +1,322 @@ +package tools + +import ( + "fmt" + "net/url" + "regexp" + "strconv" + "strings" + + "github.com/sergi/go-diff/diffmatchpatch" +) + +// maxDiffHunkPreviewBytes bounds how much of a failed hunk's "old content" +// is ever echoed back to the LLM in an error message. +const maxDiffHunkPreviewBytes = 200 + +// unifiedDiffHunkHeaderRe matches a unified-diff hunk header. The strict +// form is "@@ -12,3 +12,4 @@"; the old/new line counts are optional (default +// to 1, per the unified diff spec) and, like the line numbers, are treated +// only as hints - see applyUnifiedDiff. The entire "-old +new" position +// clause is ALSO optional, tolerating a bare "@@" some models emit when +// they're unsure of exact line numbers: see hasPosition on diffHunk for how +// that case is handled downstream. +var unifiedDiffHunkHeaderRe = regexp.MustCompile(`^@@(?:\s+-(\d+)(?:,(\d+))?\s+\+(\d+)(?:,(\d+))?)?\s*(?:@@)?`) + +// diffHunkLine is one line of a hunk body: sign is ' ' (context), '-' +// (removed), or '+' (added); text excludes the sign and any line terminator. +type diffHunkLine struct { + sign byte + text string +} + +// diffHunk is one parsed "@@ ... @@" section of a unified diff. +type diffHunk struct { + header string // original header line, kept only for error messages + // oldStart is the 1-based line number in the current file where the + // hunk begins - defaults to 1 when the header omitted it (hasPosition + // false), in which case it is NOT a real hint and callers must not + // derive anything positional (e.g. "the line before/after this hunk") + // from it - only diffmatchpatch's own content-based fuzzy search can + // locate such a hunk. + oldStart int + hasPosition bool + lines []diffHunkLine +} + +// parseUnifiedDiff parses a unified diff into hunks. It tolerates the parts +// of the format that vary across generators without affecting correctness: +// optional "--- a/file" / "+++ b/file" headers (the path is already a +// separate argument), "\ No newline at end of file" markers, and a fully +// blank line standing in for a one-space (empty) context line. +func parseUnifiedDiff(diffText string) ([]diffHunk, error) { + normalized := strings.TrimSuffix(strings.ReplaceAll(diffText, "\r\n", "\n"), "\n") + if normalized == "" { + return nil, fmt.Errorf("diff is empty") + } + lines := strings.Split(normalized, "\n") + + i := 0 + for i < len(lines) && !strings.HasPrefix(lines[i], "@@") { + trimmed := strings.TrimSpace(lines[i]) + if trimmed != "" && !strings.HasPrefix(trimmed, "---") && !strings.HasPrefix(trimmed, "+++") { + return nil, fmt.Errorf("expected a hunk header (\"@@ -old +new @@\") but found: %q", lines[i]) + } + i++ + } + if i >= len(lines) { + return nil, fmt.Errorf(`diff contains no hunks (no "@@ ... @@" header found)`) + } + + var hunks []diffHunk + for i < len(lines) { + header := lines[i] + m := unifiedDiffHunkHeaderRe.FindStringSubmatch(header) + if m == nil { + return nil, fmt.Errorf("invalid hunk header: %q", header) + } + hunk := diffHunk{header: header, oldStart: 1} + if m[1] != "" { + oldStart, err := strconv.Atoi(m[1]) + if err != nil { + return nil, fmt.Errorf("invalid hunk header %q: %w", header, err) + } + hunk.oldStart = oldStart + hunk.hasPosition = true + } + i++ + + for i < len(lines) && !strings.HasPrefix(lines[i], "@@") { + line := lines[i] + i++ + + if strings.HasPrefix(line, `\`) { + // e.g. "\ No newline at end of file" - not a content line. + continue + } + if line == "" { + hunk.lines = append(hunk.lines, diffHunkLine{sign: ' ', text: ""}) + continue + } + + sign := line[0] + if sign != ' ' && sign != '-' && sign != '+' { + return nil, fmt.Errorf("invalid diff line (must start with ' ', '-', or '+'): %q", line) + } + hunk.lines = append(hunk.lines, diffHunkLine{sign: sign, text: line[1:]}) + } + + if len(hunk.lines) == 0 { + return nil, fmt.Errorf("hunk %q has no content lines", header) + } + hunks = append(hunks, hunk) + } + + return hunks, nil +} + +// buildLineOffsets returns, for content, the byte offset at which each +// 1-based line begins: offsets[0] is the (always 0) offset of line 1, +// offsets[1] of line 2, and so on. +func buildLineOffsets(content string) []int { + offsets := make([]int, 1, strings.Count(content, "\n")+1) + offsets[0] = 0 + for i := 0; i < len(content); i++ { + if content[i] == '\n' { + offsets = append(offsets, i+1) + } + } + return offsets +} + +// lineOffset returns the byte offset where 1-based lineNum begins, clamping +// to the start/end of content for out-of-range line numbers rather than +// panicking - an inaccurate hunk header should surface as a clear "hunk +// didn't apply" error, not a crash. +func lineOffset(offsets []int, content string, lineNum int) int { + if lineNum < 1 { + return 0 + } + if lineNum-1 < len(offsets) { + return offsets[lineNum-1] + } + return len(content) +} + +// encodePatchTextLine renders one hunk line in the wire format +// diffmatchpatch.PatchFromText expects: a sign byte followed by the line +// text URL-escaped, with the escaper's space-as-'+' swapped back to a +// literal space (mirroring Patch.String in the go-diff source exactly, so +// PatchFromText's decoder - which reverses precisely that transform - +// reconstructs the original text, including any literal '+' or '%'). +func encodePatchTextLine(sign byte, text string) string { + escaped := strings.ReplaceAll(url.QueryEscape(text), "+", " ") + return string(sign) + escaped +} + +// ensureContextBoundaries returns hunks where the first hunk (if missing +// leading context) and/or the last hunk (if missing trailing context) gets +// one synthesized from the file's actual adjacent line, taken verbatim. +// Hunks in between are left exactly as parsed - they don't need this (see +// below) and, unlike the first/last hunk, a wrong hint there could inject +// misleading context instead of merely being redundant. +// +// diffmatchpatch's PatchApply pads every patch list (see PatchAddPadding in +// the go-diff source) by: (a) treating the *first* patch as if it started at +// position 0 when its first diff isn't DiffEqual, and (b) appending its own +// filler bytes as a fake trailing DiffEqual onto the *last* patch when its +// last diff isn't one either. Both corrupt that hunk's search pattern +// unless it is genuinely at the very start/end of the file - the common +// case is a single-line change with no context on one side, not actually at +// either file boundary. A real leading/trailing context line - which +// unified diffs otherwise tend to include anyway - sidesteps both special +// cases rather than relying on them being correct for hand-built patches. +func ensureContextBoundaries(hunks []diffHunk, content string) []diffHunk { + if len(hunks) == 0 { + return hunks + } + + var fileLines []string // split lazily; most diffs don't need it + linesOf := func() []string { + if fileLines == nil { + fileLines = strings.Split(content, "\n") + } + return fileLines + } + + out := make([]diffHunk, len(hunks)) + copy(out, hunks) + + // Both fixes below derive "the file line right before/after this hunk" + // from oldStart - meaningless, and actively misleading, for a hunk whose + // header omitted its position entirely (hasPosition false, oldStart just + // defaults to 1). Such a hunk relies purely on diffmatchpatch's own + // content-based fuzzy search, which already handles the Start1==0 case + // (the only one this function exists to work around) correctly on its + // own - nothing to fix up here. + if h := out[0]; h.hasPosition && len(h.lines) > 0 && h.lines[0].sign != ' ' && h.oldStart > 1 { + lines := linesOf() + precedingIdx := h.oldStart - 2 // 0-based index of file line (oldStart-1) + if precedingIdx >= 0 && precedingIdx < len(lines) { + h.oldStart-- + h.lines = append( + []diffHunkLine{{sign: ' ', text: lines[precedingIdx]}}, + h.lines..., + ) + out[0] = h + } + } + + lastIdx := len(out) - 1 + if h := out[lastIdx]; h.hasPosition && len(h.lines) > 0 && h.lines[len(h.lines)-1].sign != ' ' { + lines := linesOf() + oldLineCount := 0 + for _, l := range h.lines { + if l.sign != '+' { + oldLineCount++ + } + } + followingIdx := h.oldStart - 1 + oldLineCount // 0-based index of the line right after the hunk + // The last split element past a trailing "\n" is a Split artifact + // (nothing actually follows it), not a real line to match against. + isEOFArtifact := followingIdx == len(lines)-1 && strings.HasSuffix(content, "\n") + if followingIdx >= 0 && followingIdx < len(lines) && !isEOFArtifact { + h.lines = append(h.lines, diffHunkLine{sign: ' ', text: lines[followingIdx]}) + out[lastIdx] = h + } + } + + return out +} + +// buildGoDiffPatchText renders parsed hunks into diffmatchpatch's patch text +// format. The "@@ -start,len +start,len @@" header uses byte offsets +// computed against the real, current file content - not the line numbers +// the LLM supplied, which only ever serve as a locate-nearby hint once the +// patch is applied (diffmatchpatch.PatchApply falls back to fuzzy matching +// keyed on the header's position when the exact offset misses). +func buildGoDiffPatchText(hunks []diffHunk, content string) string { + offsets := buildLineOffsets(content) + + var b strings.Builder + for _, h := range hunks { + start := lineOffset(offsets, content, h.oldStart) + + var oldLen, newLen int + for _, l := range h.lines { + switch l.sign { + case ' ': + oldLen += len(l.text) + 1 + newLen += len(l.text) + 1 + case '-': + oldLen += len(l.text) + 1 + case '+': + newLen += len(l.text) + 1 + } + } + + fmt.Fprintf(&b, "@@ -%d,%d +%d,%d @@\n", start+1, oldLen, start+1, newLen) + for _, l := range h.lines { + b.WriteString(encodePatchTextLine(l.sign, l.text+"\n")) + b.WriteByte('\n') + } + } + return b.String() +} + +// hunkOldPreview renders the pre-patch text a hunk searched for (context and +// removed lines only), for use in a "hunk didn't apply" error message. +func hunkOldPreview(h diffHunk) string { + var b strings.Builder + for _, l := range h.lines { + if l.sign != '+' { + b.WriteString(l.text) + b.WriteByte('\n') + } + } + return truncateString(b.String(), maxDiffHunkPreviewBytes) +} + +// ApplyUnifiedDiff applies a unified diff to content entirely in memory, +// using diffmatchpatch.PatchApply for the actual merge: an exact match at +// the hunk's expected position is tried first, falling back to a +// fuzzy/context-based search nearby (tolerating minor line-number drift) +// before a hunk is considered unappliable. It returns the patched content +// and the number of hunks applied, or a descriptive error naming every hunk +// that failed to apply and a preview of the content it looked for - +// content is returned unchanged (empty) on error, so a partial/bad diff +// never corrupts the file. Exported so other packages (e.g. the provider +// tester) can exercise the exact production diff-merge semantics without +// going through EditFile's Docker-backed read/write. +func ApplyUnifiedDiff(content, diffText string) (string, int, error) { + hunks, err := parseUnifiedDiff(diffText) + if err != nil { + return "", 0, err + } + hunks = ensureContextBoundaries(hunks, content) + + patchText := buildGoDiffPatchText(hunks, content) + + dmp := diffmatchpatch.New() + patches, err := dmp.PatchFromText(patchText) + if err != nil { + return "", 0, fmt.Errorf("internal error building patch: %w", err) + } + + newContent, applied := dmp.PatchApply(patches, content) + + var failed []string + for i, ok := range applied { + if !ok && i < len(hunks) { + failed = append(failed, fmt.Sprintf("%s (not found in the file, looked for: %q)", hunks[i].header, hunkOldPreview(hunks[i]))) + } + } + if len(failed) > 0 { + return "", 0, fmt.Errorf( + "%d of %d hunk(s) could not be applied - read the file again and retry with context that matches its current content exactly:\n%s", + len(failed), len(hunks), strings.Join(failed, "\n"), + ) + } + + return newContent, len(hunks), nil +} diff --git a/backend/pkg/tools/file_diff_test.go b/backend/pkg/tools/file_diff_test.go new file mode 100644 index 00000000..ae916d62 --- /dev/null +++ b/backend/pkg/tools/file_diff_test.go @@ -0,0 +1,299 @@ +package tools + +import ( + "strings" + "testing" +) + +// TestApplyUnifiedDiff covers ApplyUnifiedDiff end to end (parse + patch + +// apply) with both positive (diff applies, exact content produced) and +// negative (diff rejected, or doesn't match, with content left untouched) +// cases, using table-driven test cases as sub-tests. +func TestApplyUnifiedDiff(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + content string + diff string + want string // expected new content, when wantErr is false + wantErr string // substring expected in the error, when non-empty + }{ + { + name: "replace a single line", + content: "line1\nline2\nline3\n", + diff: "@@ -1,2 +1,2 @@\n" + + " line1\n" + + "-line2\n" + + "+line2 changed\n", + want: "line1\nline2 changed\nline3\n", + }, + { + name: "insert lines with pure addition", + content: "line1\nline2\nline3\n", + diff: "@@ -1,2 +1,3 @@\n" + + " line1\n" + + "+inserted\n" + + " line2\n", + want: "line1\ninserted\nline2\nline3\n", + }, + { + name: "delete a line", + content: "line1\nline2\nline3\n", + diff: "@@ -1,3 +1,2 @@\n" + + " line1\n" + + "-line2\n" + + " line3\n", + want: "line1\nline3\n", + }, + { + name: "multiple non-overlapping hunks applied together", + content: "a\nb\nc\nd\ne\nf\ng\n", + diff: "@@ -1,2 +1,2 @@\n" + + " a\n" + + "-b\n" + + "+B\n" + + "@@ -6,2 +6,2 @@\n" + + " f\n" + + "-g\n" + + "+G\n", + want: "a\nB\nc\nd\ne\nf\nG\n", + }, + { + name: "tolerates minor line-number drift when content matches", + content: "line1\nline2\nline3\nline4\n", + // Header claims the hunk starts at line 2; its own context lines + // show it's really at line 1 (as if a line had been + // inserted/removed elsewhere since the diff was written) - + // PatchApply's fuzzy search must find the real position from + // the content, not the (slightly wrong) header line number. + diff: "@@ -2,3 +2,3 @@\n" + + " line1\n" + + "-line2\n" + + "+line2 changed\n" + + " line3\n", + want: "line1\nline2 changed\nline3\nline4\n", + }, + { + name: "tolerates '--- a/file' / '+++ b/file' headers", + content: "line1\nline2\n", + diff: "--- a/file.txt\n" + + "+++ b/file.txt\n" + + "@@ -2,1 +2,1 @@\n" + + "-line2\n" + + "+line2 changed\n", + want: "line1\nline2 changed\n", + }, + { + name: "tolerates a fully blank line as an empty context line", + content: "line1\n\nline3\n", + diff: "@@ -1,3 +1,3 @@\n" + + " line1\n" + + "\n" + + "-line3\n" + + "+line3 changed\n", + want: "line1\n\nline3 changed\n", + }, + { + name: "preserves special characters requiring URL-escaping", + content: "a = 1 + 2 % 3 & done\n", + diff: "@@ -1,1 +1,1 @@\n" + + "-a = 1 + 2 % 3 & done\n" + + "+a = 4 + 5 % 6 & done\n", + want: "a = 4 + 5 % 6 & done\n", + }, + { + name: "multi-hunk: middle hunk has no context on either side", + content: "a\nb\nc\nd\ne\nf\ng\n", + // Only the first and last hunk need synthesized context (see + // ensureContextBoundaries); the middle hunk works fine without it. + diff: "@@ -1,2 +1,2 @@\n" + + " a\n" + + "-b\n" + + "+B\n" + + "@@ -4,1 +4,1 @@\n" + + "-d\n" + + "+D\n" + + "@@ -6,2 +6,2 @@\n" + + "-f\n" + + "+F\n" + + " g\n", + want: "a\nB\nc\nD\ne\nF\ng\n", + }, + { + name: "hunk header without explicit line counts (implicit 1)", + content: "line1\nline2\nline3\n", + diff: "@@ -2 +2 @@\n" + + "-line2\n" + + "+line2 changed\n", + want: "line1\nline2 changed\nline3\n", + }, + { + name: "negative: empty diff is rejected", + content: "line1\n", + diff: "", + wantErr: "diff is empty", + }, + { + name: "negative: whitespace-only diff is rejected", + content: "line1\n", + diff: " \n\n ", + wantErr: "no hunks", + }, + { + name: "negative: no hunk header at all", + content: "line1\n", + diff: "just some text\nwith no diff markers\n", + wantErr: "expected a hunk header", + }, + { + // Observed from a real model (gpt-oss:120b via Ollama): it + // sometimes emits a bare "@@" with no position info at all when + // unsure of exact line numbers. diffmatchpatch's content-based + // fuzzy search must still locate the hunk from context alone - + // which requires the target line to be textually distinct from + // the file's other lines (unlike near-duplicates such as + // "line1"/"line2"/"line3", fuzzy search without any position + // hint can genuinely prefer a closer-but-wrong near-duplicate). + name: "tolerates a hunk header with no position info at all (bare '@@')", + content: "Status: draft\nOwner: alice\nPriority: low\n", + diff: "@@\n-Priority: low\n+Priority: high\n", + want: "Status: draft\nOwner: alice\nPriority: high\n", + }, + { + name: "tolerates unparseable text after a bare '@@' header", + content: "line1\n", + diff: "@@ not a real header @@\n-line1\n+line2\n", + want: "line2\n", + }, + { + name: "negative: hunk body line with invalid prefix", + content: "line1\nline2\n", + diff: "@@ -1,2 +1,2 @@\n line1\n*line2\n", + wantErr: "invalid diff line", + }, + { + name: "negative: hunk header with no content lines", + content: "line1\nline2\n", + diff: "@@ -1,1 +1,1 @@\n@@ -2,1 +2,1 @@\n-line2\n+line2 changed\n", + wantErr: "no content lines", + }, + { + name: "negative: context does not match file content", + content: "line1\nline2\nline3\n", + diff: "@@ -2,1 +2,1 @@\n" + + "-this line does not exist in the file\n" + + "+replacement\n", + wantErr: "could not be applied", + }, + { + name: "negative: content matches nowhere near the claimed location", + content: strings.Repeat("filler line\n", 50) + "target line\n" + strings.Repeat("filler line\n", 50), + diff: "@@ -1,1 +1,1 @@\n" + + "-completely different text that is nowhere in the file\n" + + "+replacement\n", + wantErr: "could not be applied", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got, hunks, err := ApplyUnifiedDiff(tt.content, tt.diff) + + if tt.wantErr != "" { + if err == nil { + t.Fatalf("ApplyUnifiedDiff() expected error containing %q, got nil (result: %q)", tt.wantErr, got) + } + if !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("ApplyUnifiedDiff() error = %q, want it to contain %q", err.Error(), tt.wantErr) + } + if got != "" { + t.Errorf("ApplyUnifiedDiff() on error must return empty content, got %q", got) + } + return + } + + if err != nil { + t.Fatalf("ApplyUnifiedDiff() unexpected error: %v", err) + } + if got != tt.want { + t.Errorf("ApplyUnifiedDiff() content = %q, want %q", got, tt.want) + } + if hunks <= 0 { + t.Errorf("ApplyUnifiedDiff() hunks applied = %d, want > 0", hunks) + } + }) + } +} + +// TestParseUnifiedDiff checks the parser in isolation, independent of +// go-diff, so a bug in hunk parsing and a bug in patch application are +// distinguishable from their respective test failures. +func TestParseUnifiedDiff(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + diff string + wantHunks int + wantLines int // lines in the first hunk, when wantHunks > 0 + wantErr string + }{ + { + name: "single hunk", + diff: "@@ -1,2 +1,2 @@\n line1\n-line2\n+line2 changed\n", + wantHunks: 1, + wantLines: 3, + }, + { + name: "two hunks", + diff: "@@ -1,1 +1,1 @@\n-a\n+A\n@@ -3,1 +3,1 @@\n-c\n+C\n", + wantHunks: 2, + }, + { + name: "empty", + diff: "", + wantErr: "diff is empty", + }, + { + name: "garbage before first hunk", + diff: "not a diff at all", + wantErr: "expected a hunk header", + }, + { + name: "garbage between header lines and hunk", + diff: "--- a/file\nsome garbage line\n@@ -1,1 +1,1 @@\n-a\n+A\n", + wantErr: "expected a hunk header", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + hunks, err := parseUnifiedDiff(tt.diff) + + if tt.wantErr != "" { + if err == nil { + t.Fatalf("parseUnifiedDiff() expected error containing %q, got nil", tt.wantErr) + } + if !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("parseUnifiedDiff() error = %q, want it to contain %q", err.Error(), tt.wantErr) + } + return + } + + if err != nil { + t.Fatalf("parseUnifiedDiff() unexpected error: %v", err) + } + if len(hunks) != tt.wantHunks { + t.Fatalf("parseUnifiedDiff() hunks = %d, want %d", len(hunks), tt.wantHunks) + } + if tt.wantLines > 0 && len(hunks[0].lines) != tt.wantLines { + t.Errorf("parseUnifiedDiff() first hunk lines = %d, want %d", len(hunks[0].lines), tt.wantLines) + } + }) + } +} diff --git a/backend/pkg/tools/graphiti_search.go b/backend/pkg/tools/graphiti_search.go index fd70414b..99f048e2 100644 --- a/backend/pkg/tools/graphiti_search.go +++ b/backend/pkg/tools/graphiti_search.go @@ -6,6 +6,8 @@ import ( "errors" "fmt" "net/url" + "regexp" + "strconv" "strings" "time" @@ -16,6 +18,18 @@ import ( "github.com/sirupsen/logrus" ) +// maxGraphitiResponsePreviewBytes bounds how much of a non-2xx Graphiti +// response body is ever shown to the LLM (via truncateString), so a verbose +// error page or stack trace from the Graphiti backend can't blow up the +// agent's context window. +const maxGraphitiResponsePreviewBytes = 512 + +// graphitiAPIStatusErrorRe matches the one fixed error shape the vendored +// graphiti-go-client emits for a non-2xx HTTP response: a plain fmt.Errorf +// (see client.go's `do`) with no typed status code, so this is the only way +// to recover the status and body without forking the dependency. +var graphitiAPIStatusErrorRe = regexp.MustCompile(`(?s)API request failed with status (\d+): (.*)`) + type GraphitiSearcher interface { IsEnabled() bool TemporalWindowSearch(ctx context.Context, req graphiti.TemporalSearchRequest) (*graphiti.TemporalSearchResponse, error) @@ -130,7 +144,7 @@ func (t *graphitiSearchTool) Handle(ctx context.Context, name string, args json. ctx, observation := obs.Observer.NewObservation(ctx) - retrieverTitle, ok := graphitiRetrieverTitles[searchArgs.SearchType] + retrieverTitle, ok := graphitiRetrieverTitles[searchArgs.SearchType.String()] if !ok { retrieverTitle = "retrieve context from graphiti knowledge graph" } @@ -189,8 +203,8 @@ func (t *graphitiSearchTool) Handle(ctx context.Context, name string, args json. var urlErr *url.Error if errors.As(err, &urlErr) { softMsg := fmt.Sprintf( - "Graphiti knowledge graph is temporarily unavailable (%v); continuing without historical context.", - err, + "Graphiti knowledge graph is temporarily unavailable (%s); continuing without historical context.", + truncateString(err.Error(), maxGraphitiResponsePreviewBytes), ) retriever.End( langfuse.WithRetrieverStatus(softMsg), @@ -200,6 +214,37 @@ func (t *graphitiSearchTool) Handle(ctx context.Context, name string, args json. return softMsg, nil } + // The Graphiti server responded (connection succeeded), but with a + // non-2xx status. graphiti-go-client has no typed error for this - it + // embeds the raw, unbounded response body in a plain fmt.Errorf - so + // recover the status/body via the one fixed message shape it emits. + if m := graphitiAPIStatusErrorRe.FindStringSubmatch(err.Error()); m != nil { + statusCode, convErr := strconv.Atoi(m[1]) + body := truncateString(strings.TrimSpace(m[2]), maxGraphitiResponsePreviewBytes) + + if convErr == nil && statusCode >= 500 { + // A 5xx from Graphiti's own backend is the same class of + // problem as a transport failure: not fixable by editing + // arguments, so degrade gracefully and show the LLM whatever + // of the response body we could read. + softMsg := fmt.Sprintf( + "Graphiti knowledge graph returned HTTP %d and is likely temporarily unavailable; continuing without historical context. Response: %s", + statusCode, body, + ) + retriever.End( + langfuse.WithRetrieverStatus(softMsg), + langfuse.WithRetrieverLevel(langfuse.ObservationLevelWarning), + ) + logger.WithError(err).Warnf("graphiti search '%s' returned HTTP %d, degrading gracefully", searchArgs.SearchType, statusCode) + return softMsg, nil + } + + // 4xx (or an unparseable status) likely means our own request was + // malformed - stays a hard failure so the tool-call fixer can act + // on it, but capped so a verbose error page can't blow up its context. + err = fmt.Errorf("graphiti API request failed with status %s: %s", m[1], body) + } + retriever.End( langfuse.WithRetrieverStatus(err.Error()), langfuse.WithRetrieverLevel(langfuse.ObservationLevelError), @@ -371,7 +416,7 @@ func (t *graphitiSearchTool) handleDiverseResultsSearch( maxResults = DefaultDiverseMaxResults } - diversityLevel := args.DiversityLevel + diversityLevel := args.DiversityLevel.String() if diversityLevel == "" { diversityLevel = DefaultDiversityLevel } @@ -467,7 +512,7 @@ func (t *graphitiSearchTool) handleRecentContextSearch( maxResults = DefaultRecentMaxResults } - recencyWindow := args.RecencyWindow + recencyWindow := args.RecencyWindow.String() if recencyWindow == "" { recencyWindow = DefaultRecencyWindow } diff --git a/backend/pkg/tools/graphiti_search_test.go b/backend/pkg/tools/graphiti_search_test.go index f6dbd8e8..9b47b8a6 100644 --- a/backend/pkg/tools/graphiti_search_test.go +++ b/backend/pkg/tools/graphiti_search_test.go @@ -75,6 +75,78 @@ func fakeNetError() error { ) } +// fakeStatusError mimics the exact error shape graphiti-go-client's `do` +// produces for a non-2xx HTTP response: a plain fmt.Errorf embedding the raw +// response body, further wrapped by the handle*Search method's own %w wrap. +func fakeStatusError(statusCode int, body string) error { + return fmt.Errorf( + "recent context search failed: API request failed with status %d: %s", + statusCode, body, + ) +} + +func TestGraphitiSearchTool_Handle_ServerError5xx_DegradesGracefullyWithBody(t *testing.T) { + tool := NewGraphitiSearchTool(1, nil, nil, &stubGraphitiSearcher{ + enabled: true, + err: fakeStatusError(502, "502 Bad Gateway"), + }) + + args := []byte(`{"search_type":"recent_context","query":"test query","message":"m"}`) + result, err := tool.Handle(t.Context(), GraphitiSearchToolName, args) + + if err != nil { + t.Fatalf("expected graceful degradation (nil error) on 5xx, got error: %v", err) + } + if !strings.Contains(result, "HTTP 502") { + t.Fatalf("expected result to mention the status code, got: %q", result) + } + if !strings.Contains(result, "502 Bad Gateway") { + t.Fatalf("expected result to include the response body, got: %q", result) + } +} + +func TestGraphitiSearchTool_Handle_ServerError5xx_BodyTruncatedAt512Bytes(t *testing.T) { + hugeBody := strings.Repeat("x", 2000) + tool := NewGraphitiSearchTool(1, nil, nil, &stubGraphitiSearcher{ + enabled: true, + err: fakeStatusError(500, hugeBody), + }) + + args := []byte(`{"search_type":"recent_context","query":"test query","message":"m"}`) + result, err := tool.Handle(t.Context(), GraphitiSearchToolName, args) + + if err != nil { + t.Fatalf("expected graceful degradation (nil error) on 5xx, got error: %v", err) + } + if strings.Count(result, "x") >= 2000 { + t.Fatalf("expected the 2000-byte body to be truncated to the 512-byte cap, got a result of length %d", len(result)) + } + if !strings.Contains(result, "truncated") { + t.Fatalf("expected result to indicate truncation, got: %q", result) + } +} + +func TestGraphitiSearchTool_Handle_ClientError4xx_StaysHardWithTruncatedBody(t *testing.T) { + hugeBody := strings.Repeat("y", 2000) + tool := NewGraphitiSearchTool(1, nil, nil, &stubGraphitiSearcher{ + enabled: true, + err: fakeStatusError(400, hugeBody), + }) + + args := []byte(`{"search_type":"recent_context","query":"test query","message":"m"}`) + _, err := tool.Handle(t.Context(), GraphitiSearchToolName, args) + + if err == nil { + t.Fatal("expected a hard failure for a 4xx status, got nil error") + } + if !strings.Contains(err.Error(), "status 400") { + t.Fatalf("expected error to mention the status code, got: %v", err) + } + if strings.Count(err.Error(), "y") >= 2000 { + t.Fatalf("expected the 2000-byte body to be truncated to the 512-byte cap even on the hard-fail path, got error of length %d", len(err.Error())) + } +} + func TestGraphitiSearchTool_Handle_NetworkFailure_DegradesGracefully(t *testing.T) { tool := NewGraphitiSearchTool(1, nil, nil, &stubGraphitiSearcher{enabled: true, err: fakeNetError()}) diff --git a/backend/pkg/tools/registry.go b/backend/pkg/tools/registry.go index 877c58e4..dfeba43c 100644 --- a/backend/pkg/tools/registry.go +++ b/backend/pkg/tools/registry.go @@ -176,9 +176,10 @@ var registryDefinitions = map[string]llms.FunctionDefinition{ Parameters: reflector.Reflect(&TerminalAction{}), }, FileToolName: { - Name: FileToolName, - Description: "Modifies or reads local files", - Parameters: reflector.Reflect(&FileAction{}), + Name: FileToolName, + Description: "Reads, writes, or edits local files (see 'action'). " + + "Prefer edit_file for targeted changes to an existing file; use write_file only for a new file or a full rewrite.", + Parameters: reflector.Reflect(&FileAction{}), }, ReportResultToolName: { Name: ReportResultToolName, diff --git a/backend/pkg/tools/sploitus.go b/backend/pkg/tools/sploitus.go index 28a229b2..7ab35270 100644 --- a/backend/pkg/tools/sploitus.go +++ b/backend/pkg/tools/sploitus.go @@ -77,13 +77,13 @@ func (s *sploitus) Handle(ctx context.Context, name string, args json.RawMessage } // Normalise exploit type - exploitType := strings.ToLower(strings.TrimSpace(action.ExploitType)) + exploitType := strings.ToLower(strings.TrimSpace(action.ExploitType.String())) if exploitType == "" { exploitType = defaultSploitusType } // Normalise sort order - sort := strings.ToLower(strings.TrimSpace(action.Sort)) + sort := strings.ToLower(strings.TrimSpace(action.Sort.String())) if sort == "" { sort = sploitusDefaultSort } diff --git a/backend/pkg/tools/terminal.go b/backend/pkg/tools/terminal.go index b47d2122..c1d2e8dc 100644 --- a/backend/pkg/tools/terminal.go +++ b/backend/pkg/tools/terminal.go @@ -150,9 +150,12 @@ func (t *terminal) Handle(ctx context.Context, name string, args json.RawMessage // from the other fields present, so infer it instead of failing the call // outright and burning a tool-call-fixer round-trip on something that // doesn't need one. - if action.Content != "" { + switch { + case action.Diff != "": + action.Action = EditFile + case action.Content != "": action.Action = WriteFile - } else { + default: action.Action = ReadFile } } @@ -164,10 +167,13 @@ func (t *terminal) Handle(ctx context.Context, name string, args json.RawMessage switch action.Action { case ReadFile: - result, err := t.ReadFile(ctx, t.flowID, action.Path) + result, err := t.ReadFile(ctx, t.flowID, action.Path.String()) return t.wrapCommandResult(ctx, args, name, result, err) case WriteFile: - result, err := t.WriteFile(ctx, t.flowID, action.Content, action.Path) + result, err := t.WriteFile(ctx, t.flowID, action.Content, action.Path.String()) + return t.wrapCommandResult(ctx, args, name, result, err) + case EditFile: + result, err := t.EditFile(ctx, t.flowID, action.Path.String(), action.Diff.String()) return t.wrapCommandResult(ctx, args, name, result, err) default: logger.Error("unknown file action") @@ -325,6 +331,36 @@ func (t *terminal) ReadFile(ctx context.Context, flowID int64, path string) (str return "", fmt.Errorf("path is required and cannot be empty") } + cwd := docker.WorkFolderPathInContainer + escapedPath := strings.ReplaceAll(path, "'", "'\"'\"'") + catCommand := fmt.Sprintf("cat '%s'", escapedPath) + // Format read file command with styling + styledCommand := fmt.Sprintf("%s $ %s%s%s%s", cwd, ansiColorInputCmd, catCommand, ansiColorReset, ansiLineTerminator) + _, err := t.tlp.PutMsg(ctx, database.TermlogTypeStdin, styledCommand, t.containerID, t.taskID, t.subtaskID) + if err != nil { + return "", fmt.Errorf("failed to put terminal log (read file cmd): %w", err) + } + + content, err := t.readFileFromContainer(ctx, flowID, path) + if err != nil { + return "", err + } + + // Style file content output + styledContent := fmt.Sprintf("%s%s%s%s", ansiColorSystemMsg, content, ansiColorReset, ansiLineTerminator) + _, err = t.tlp.PutMsg(ctx, database.TermlogTypeStdout, styledContent, t.containerID, t.taskID, t.subtaskID) + if err != nil { + return "", fmt.Errorf("failed to put terminal log (read file content): %w", err) + } + + return content, nil +} + +// readFileFromContainer copies path out of the flow's container and returns +// its content. It performs no terminal-log writes, so callers that need the +// content only as an intermediate step (e.g. EditFile, before reapplying a +// diff and writing back) don't echo a spurious "cat" transcript entry. +func (t *terminal) readFileFromContainer(ctx context.Context, flowID int64, path string) (string, error) { containerName := PrimaryTerminalName(flowID) isRunning, err := t.dockerClient.IsContainerRunning(ctx, t.containerLID) @@ -335,16 +371,6 @@ func (t *terminal) ReadFile(ctx context.Context, flowID int64, path string) (str return "", fmt.Errorf("container runtime is not operational") } - cwd := docker.WorkFolderPathInContainer - escapedPath := strings.ReplaceAll(path, "'", "'\"'\"'") - catCommand := fmt.Sprintf("cat '%s'", escapedPath) - // Format read file command with styling - styledCommand := fmt.Sprintf("%s $ %s%s%s%s", cwd, ansiColorInputCmd, catCommand, ansiColorReset, ansiLineTerminator) - _, err = t.tlp.PutMsg(ctx, database.TermlogTypeStdin, styledCommand, t.containerID, t.taskID, t.subtaskID) - if err != nil { - return "", fmt.Errorf("failed to put terminal log (read file cmd): %w", err) - } - reader, stats, err := t.dockerClient.CopyFromContainer(ctx, containerName, path) if err != nil { return "", fmt.Errorf("failed to copy file: %w", err) @@ -395,15 +421,7 @@ func (t *terminal) ReadFile(ctx context.Context, flowID int64, path string) (str } } - content := buffer.String() - // Style file content output - styledContent := fmt.Sprintf("%s%s%s%s", ansiColorSystemMsg, content, ansiColorReset, ansiLineTerminator) - _, err = t.tlp.PutMsg(ctx, database.TermlogTypeStdout, styledContent, t.containerID, t.taskID, t.subtaskID) - if err != nil { - return "", fmt.Errorf("failed to put terminal log (read file content): %w", err) - } - - return content, nil + return buffer.String(), nil } func (t *terminal) WriteFile(ctx context.Context, flowID int64, content string, path string) (string, error) { @@ -411,14 +429,33 @@ func (t *terminal) WriteFile(ctx context.Context, flowID int64, content string, return "", fmt.Errorf("path is required and cannot be empty") } + if err := t.writeFileToContainer(ctx, flowID, path, content); err != nil { + return "", err + } + + // Format success message with styling + successMsg := fmt.Sprintf("File successfully saved to %s", path) + styledMsg := fmt.Sprintf("%s%s%s%s", ansiColorSystemMsg, successMsg, ansiColorReset, ansiLineTerminator) + _, err := t.tlp.PutMsg(ctx, database.TermlogTypeStdin, styledMsg, t.containerID, t.taskID, t.subtaskID) + if err != nil { + return "", fmt.Errorf("failed to put terminal log (write file cmd): %w", err) + } + + return fmt.Sprintf("Successfully wrote %d bytes to %s", len(content), path), nil +} + +// writeFileToContainer copies content into the flow's container at path, +// overwriting it. It performs no terminal-log writes; WriteFile and EditFile +// each log their own, differently-worded, success message. +func (t *terminal) writeFileToContainer(ctx context.Context, flowID int64, path, content string) error { containerName := PrimaryTerminalName(flowID) isRunning, err := t.dockerClient.IsContainerRunning(ctx, t.containerLID) if err != nil { - return "", fmt.Errorf("container runtime check failed: %w", err) + return fmt.Errorf("container runtime check failed: %w", err) } if !isRunning { - return "", fmt.Errorf("target container is not operational") + return fmt.Errorf("target container is not operational") } // Docker SDK requires TAR format for file transfer @@ -434,17 +471,17 @@ func (t *terminal) WriteFile(ctx context.Context, flowID int64, content string, } err = archiveWriter.WriteHeader(fileDescriptor) if err != nil { - return "", fmt.Errorf("tar archive header generation failed: %w", err) + return fmt.Errorf("tar archive header generation failed: %w", err) } _, err = archiveWriter.Write([]byte(content)) if err != nil { - return "", fmt.Errorf("tar archive content serialization failed: %w", err) + return fmt.Errorf("tar archive content serialization failed: %w", err) } err = archiveWriter.Close() if err != nil { - return "", fmt.Errorf("failed to close tar writer: %w", err) + return fmt.Errorf("failed to close tar writer: %w", err) } dir := filepath.Dir(path) @@ -452,18 +489,45 @@ func (t *terminal) WriteFile(ctx context.Context, flowID int64, content string, AllowOverwriteDirWithFile: true, }) if err != nil { - return "", fmt.Errorf("container file transfer failed: %w", err) + return fmt.Errorf("container file transfer failed: %w", err) } - // Format success message with styling - successMsg := fmt.Sprintf("File successfully saved to %s", path) - styledMsg := fmt.Sprintf("%s%s%s%s", ansiColorSystemMsg, successMsg, ansiColorReset, ansiLineTerminator) - _, err = t.tlp.PutMsg(ctx, database.TermlogTypeStdin, styledMsg, t.containerID, t.taskID, t.subtaskID) + return nil +} + +// EditFile applies a unified diff to the file at path: it reads the current +// content, applies the diff to it entirely in memory (see applyUnifiedDiff), +// and only if every hunk applied cleanly writes the result back - a diff +// that doesn't fully apply leaves the file untouched. +func (t *terminal) EditFile(ctx context.Context, flowID int64, path, diffText string) (string, error) { + if path == "" { + return "", fmt.Errorf("path is required and cannot be empty") + } + if strings.TrimSpace(diffText) == "" { + return "", fmt.Errorf("diff is required and cannot be empty") + } + + current, err := t.readFileFromContainer(ctx, flowID, path) if err != nil { - return "", fmt.Errorf("failed to put terminal log (write file cmd): %w", err) + return "", fmt.Errorf("failed to read current content of %s before editing: %w", path, err) } - return fmt.Sprintf("Successfully wrote %d bytes to %s", len(content), path), nil + newContent, hunksApplied, err := ApplyUnifiedDiff(current, diffText) + if err != nil { + return "", fmt.Errorf("failed to apply diff to %s: %w", path, err) + } + + if err := t.writeFileToContainer(ctx, flowID, path, newContent); err != nil { + return "", fmt.Errorf("failed to write edited content of %s: %w", path, err) + } + + successMsg := fmt.Sprintf("Applied %d diff hunk(s) to %s (%d -> %d bytes)", hunksApplied, path, len(current), len(newContent)) + styledMsg := fmt.Sprintf("%s%s%s%s", ansiColorSystemMsg, successMsg, ansiColorReset, ansiLineTerminator) + if _, err := t.tlp.PutMsg(ctx, database.TermlogTypeStdin, styledMsg, t.containerID, t.taskID, t.subtaskID); err != nil { + return "", fmt.Errorf("failed to put terminal log (edit file cmd): %w", err) + } + + return successMsg, nil } func PrimaryTerminalName(flowID int64) string { diff --git a/backend/pkg/tools/terminal_test.go b/backend/pkg/tools/terminal_test.go index cf116e6e..fe71570c 100644 --- a/backend/pkg/tools/terminal_test.go +++ b/backend/pkg/tools/terminal_test.go @@ -52,6 +52,15 @@ type contextAwareMockDockerClient struct { // Set by CopyFromContainer/CopyToContainer to track which file operation ran copyFromCalled bool copyToCalled bool + + // readFileContent, when non-empty, is what CopyFromContainer returns as + // the "current" file content (wrapped in a single-file tar archive) - + // used by EditFile tests that need to read-then-patch existing content. + readFileContent string + // writtenContent captures the last file content CopyToContainer received + // (unwrapped from its tar archive), so tests can assert on the result of + // an edit_file/write_file operation. + writtenContent string } func (m *contextAwareMockDockerClient) RunContainer(_ context.Context, _ string, _ database.ContainerType, @@ -111,13 +120,36 @@ func (m *contextAwareMockDockerClient) ListContainerDir(_ context.Context, _ str func (m *contextAwareMockDockerClient) ContainerExecInspect(_ context.Context, _ string) (container.ExecInspect, error) { return m.inspectResp, nil } -func (m *contextAwareMockDockerClient) CopyToContainer(_ context.Context, _ string, _ string, _ io.Reader, _ container.CopyToContainerOptions) error { +func (m *contextAwareMockDockerClient) CopyToContainer(_ context.Context, _ string, _ string, src io.Reader, _ container.CopyToContainerOptions) error { m.copyToCalled = true + + tarReader := tar.NewReader(src) + if hdr, err := tarReader.Next(); err == nil { + buf := make([]byte, hdr.Size) + _, _ = io.ReadFull(tarReader, buf) + m.writtenContent = string(buf) + } + return nil } func (m *contextAwareMockDockerClient) CopyFromContainer(_ context.Context, _ string, _ string) (io.ReadCloser, container.PathStat, error) { m.copyFromCalled = true - return io.NopCloser(bytes.NewReader(nil)), container.PathStat{}, nil + + if m.readFileContent == "" { + return io.NopCloser(bytes.NewReader(nil)), container.PathStat{}, nil + } + + var tarBuffer bytes.Buffer + tarWriter := tar.NewWriter(&tarBuffer) + _ = tarWriter.WriteHeader(&tar.Header{ + Name: "file", + Mode: 0600, + Size: int64(len(m.readFileContent)), + }) + _, _ = tarWriter.Write([]byte(m.readFileContent)) + _ = tarWriter.Close() + + return io.NopCloser(&tarBuffer), container.PathStat{}, nil } func (m *contextAwareMockDockerClient) Cleanup(_ context.Context) error { return nil } func (m *contextAwareMockDockerClient) GetDefaultImage() string { return "test-image" } @@ -233,6 +265,179 @@ func TestTerminalHandle_FileAction_DefaultsToWriteFile_WhenContentPresent(t *tes } } +// TestTerminalHandle_FileAction_ExtraQuotedAction_StillDispatches reproduces +// the exact production failure: the LLM sent action/path wrapped in an extra +// literal pair of quotes (e.g. `"action": "\"write_file\""`), which used to +// fail with "unknown file action" since the corrupted value matched no case. +// The String type now unwraps this at unmarshal time, so dispatch succeeds. +func TestTerminalHandle_FileAction_ExtraQuotedAction_StillDispatches(t *testing.T) { + mock := &contextAwareMockDockerClient{isRunning: true} + term := &terminal{ + flowID: 1, + containerID: 1, + containerLID: "test-container", + dockerClient: mock, + tlp: &contextTestTermLogProvider{}, + } + + args := json.RawMessage(`{"action": "\"write_file\"", "path": "\"/home/evidence/inject.py\"", "content": "print(1)", "message": "m"}`) + _, err := term.Handle(t.Context(), FileToolName, args) + + if err != nil { + t.Fatalf("expected extra-quoted write_file to still dispatch, got error: %v", err) + } + if !mock.copyToCalled || mock.copyFromCalled { + t.Fatalf("expected CopyToContainer (write_file) to be called, copyTo=%v copyFrom=%v", + mock.copyToCalled, mock.copyFromCalled) + } +} + +// TestTerminalHandle_FileAction_EditFile_AppliesDiff exercises edit_file end +// to end through Handle(): it reads the mock's current content, applies the +// diff in memory, and writes the result back - all via CopyFromContainer / +// CopyToContainer, exactly like write_file, but without resending unchanged +// content. +func TestTerminalHandle_FileAction_EditFile_AppliesDiff(t *testing.T) { + mock := &contextAwareMockDockerClient{ + isRunning: true, + readFileContent: "line1\nline2\nline3\n", + } + term := &terminal{ + flowID: 1, + containerID: 1, + containerLID: "test-container", + dockerClient: mock, + tlp: &contextTestTermLogProvider{}, + } + + diff := "@@ -1,2 +1,2 @@\n line1\n-line2\n+line2 changed\n" + args := json.RawMessage(fmt.Sprintf( + `{"action":"edit_file","path":"/work/test.py","diff":%s,"message":"m"}`, + mustJSONString(t, diff), + )) + result, err := term.Handle(t.Context(), FileToolName, args) + + if err != nil { + t.Fatalf("expected edit_file to succeed, got error: %v", err) + } + if !mock.copyFromCalled || !mock.copyToCalled { + t.Fatalf("expected both CopyFromContainer (read) and CopyToContainer (write), got from=%v to=%v", + mock.copyFromCalled, mock.copyToCalled) + } + want := "line1\nline2 changed\nline3\n" + if mock.writtenContent != want { + t.Errorf("written content = %q, want %q", mock.writtenContent, want) + } + if !strings.Contains(result, "1 diff hunk") { + t.Errorf("result = %q, want it to mention the number of hunks applied", result) + } +} + +// TestTerminalHandle_FileAction_DefaultsToEditFile_WhenDiffPresent mirrors +// the write_file/read_file action-inference tests: when 'action' is omitted +// but 'diff' is present, the intent is unambiguous. +func TestTerminalHandle_FileAction_DefaultsToEditFile_WhenDiffPresent(t *testing.T) { + mock := &contextAwareMockDockerClient{ + isRunning: true, + readFileContent: "line1\nline2\n", + } + term := &terminal{ + flowID: 1, + containerID: 1, + containerLID: "test-container", + dockerClient: mock, + tlp: &contextTestTermLogProvider{}, + } + + diff := "@@ -2,1 +2,1 @@\n-line2\n+line2 changed\n" + args := json.RawMessage(fmt.Sprintf( + `{"path":"/work/test.py","diff":%s,"message":"m"}`, + mustJSONString(t, diff), + )) + _, err := term.Handle(t.Context(), FileToolName, args) + + if err != nil { + t.Fatalf("expected inferred edit_file to succeed, got error: %v", err) + } + if !mock.copyFromCalled || !mock.copyToCalled { + t.Fatalf("expected inferred edit_file to read then write, got from=%v to=%v", + mock.copyFromCalled, mock.copyToCalled) + } +} + +// TestTerminalHandle_FileAction_EditFile_NoMatch_LeavesFileUntouched checks +// that a diff whose context doesn't match the current content fails clearly +// and never reaches CopyToContainer - a bad edit must not corrupt the file. +func TestTerminalHandle_FileAction_EditFile_NoMatch_LeavesFileUntouched(t *testing.T) { + mock := &contextAwareMockDockerClient{ + isRunning: true, + readFileContent: "line1\nline2\nline3\n", + } + term := &terminal{ + flowID: 1, + containerID: 1, + containerLID: "test-container", + dockerClient: mock, + tlp: &contextTestTermLogProvider{}, + } + + diff := "@@ -2,1 +2,1 @@\n-this text is not in the file\n+replacement\n" + args := json.RawMessage(fmt.Sprintf( + `{"action":"edit_file","path":"/work/test.py","diff":%s,"message":"m"}`, + mustJSONString(t, diff), + )) + result, err := term.Handle(t.Context(), FileToolName, args) + + // Handle() wraps tool errors into a successful-looking result string (see + // wrapCommandResult) rather than a Go error, so assert on the message. + if err != nil { + t.Fatalf("Handle() should swallow tool errors via wrapCommandResult, got error: %v", err) + } + if !strings.Contains(result, "could not be applied") { + t.Errorf("result = %q, want it to mention the hunk could not be applied", result) + } + if mock.copyToCalled { + t.Fatal("expected CopyToContainer to NOT be called when the diff doesn't apply") + } +} + +// TestTerminalHandle_FileAction_EmptyDiff_ReturnsClearError checks edit_file +// with an empty diff fails with a clear message instead of a Docker-level error. +func TestTerminalHandle_FileAction_EmptyDiff_ReturnsClearError(t *testing.T) { + mock := &contextAwareMockDockerClient{isRunning: true} + term := &terminal{ + flowID: 1, + containerID: 1, + containerLID: "test-container", + dockerClient: mock, + tlp: &contextTestTermLogProvider{}, + } + + args := json.RawMessage(`{"action":"edit_file","path":"/work/test.py","diff":"","message":"m"}`) + result, err := term.Handle(t.Context(), FileToolName, args) + + if err != nil { + t.Fatalf("Handle() should swallow tool errors via wrapCommandResult, got error: %v", err) + } + if !strings.Contains(result, "diff is required") { + t.Errorf("result = %q, want it to mention that diff is required", result) + } + if mock.copyFromCalled || mock.copyToCalled { + t.Fatal("expected neither CopyFromContainer nor CopyToContainer to be called for an empty diff") + } +} + +// mustJSONString marshals s as a JSON string literal, for embedding +// multi-line diff text into a hand-written JSON args payload in tests. +func mustJSONString(t *testing.T, s string) string { + t.Helper() + b, err := json.Marshal(s) + if err != nil { + t.Fatalf("json.Marshal(%q) error: %v", s, err) + } + return string(b) +} + func TestTerminalHandle_FileAction_DefaultsToReadFile_WhenContentAbsent(t *testing.T) { mock := &contextAwareMockDockerClient{isRunning: true} term := &terminal{