Merge pull request #2452 from trungutt/fix/edit-file-malformed-json

fix: repair malformed JSON in edit_file tool call arguments

Author: dgageot

pkg/acp/filesystem.go

@@ -188,8 +188,12 @@ func (t *FilesystemToolset) handleWriteFile(ctx context.Context, toolCall tools.
 }
 
 func (t *FilesystemToolset) handleEditFile(ctx context.Context, toolCall tools.ToolCall) (*tools.ToolCallResult, error) {
-	var args builtin.EditFileArgs
-	if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &args); err != nil {
+	data := toolCall.Function.Arguments
+	if data == "" {
+		data = "{}"
+	}
+	args, err := builtin.ParseEditFileArgs([]byte(data))
+	if err != nil {
 		return nil, fmt.Errorf("failed to parse arguments: %w", err)
 	}
 

pkg/tools/builtin/filesystem.go

@@ -166,40 +166,109 @@ type EditFileArgs struct {
 	Edits []Edit `json:"edits" jsonschema:"Array of edit operations"`
 }
 
-// UnmarshalJSON handles LLM-generated arguments where "edits" may be
-// a JSON string instead of a JSON array (double-serialized).
-func (a *EditFileArgs) UnmarshalJSON(data []byte) error {
+// ParseEditFileArgs parses LLM-generated edit_file arguments, handling two
+// common failure modes:
+//  1. The outer JSON itself is malformed — typically extra closing braces/brackets
+//     or stray escape sequences caused by the model losing track of nesting depth
+//     when the text payload contains structural characters (e.g. YAML, Dockerfiles).
+//  2. The "edits" field is double-serialized (a JSON string instead of an array).
+func ParseEditFileArgs(data []byte) (EditFileArgs, error) {
 	var raw struct {
 		Path  string          `json:"path"`
 		Edits json.RawMessage `json:"edits"`
 	}
+
 	if err := json.Unmarshal(data, &raw); err != nil {
-		return fmt.Errorf("failed to parse edit_file arguments: %w", err)
+		repaired, ok := tryRepairEditFileJSON(data)
+		if !ok {
+			return EditFileArgs{}, fmt.Errorf("failed to parse edit_file arguments: %w", err)
+		}
+		if err := json.Unmarshal(repaired, &raw); err != nil {
+			return EditFileArgs{}, fmt.Errorf("failed to parse edit_file arguments after repair: %w", err)
+		}
+		slog.Debug("Repaired malformed edit_file JSON arguments")
 	}
 
-	a.Path = raw.Path
+	args := EditFileArgs{Path: raw.Path}
 
 	// When edits is missing or null (e.g. during argument streaming in
 	// the TUI, or partial tool calls), accept the partial result.
 	if len(raw.Edits) == 0 || string(raw.Edits) == "null" {
-		return nil
+		return args, nil
 	}
 
 	// Try parsing edits as an array first (normal case).
-	if err := json.Unmarshal(raw.Edits, &a.Edits); err == nil {
-		return nil
+	if err := json.Unmarshal(raw.Edits, &args.Edits); err == nil {
+		return args, nil
 	}
 
 	// Try unwrapping a double-serialized JSON string.
 	var editsStr string
 	if err := json.Unmarshal(raw.Edits, &editsStr); err != nil {
-		return fmt.Errorf("edits field is neither an array nor a JSON string: %w", err)
-	}
-	if err := json.Unmarshal([]byte(editsStr), &a.Edits); err != nil {
-		return fmt.Errorf("failed to parse double-serialized edits string: %w", err)
+		return EditFileArgs{}, fmt.Errorf("edits field is neither an array nor a JSON string: %w", err)
+	}
+	if err := json.Unmarshal([]byte(editsStr), &args.Edits); err != nil {
+		return EditFileArgs{}, fmt.Errorf("failed to parse double-serialized edits string: %w", err)
+	}
+
+	return args, nil
+}
+
+// tryRepairEditFileJSON attempts to fix common LLM JSON malformations by
+// iteratively removing the offending character(s) at each json.SyntaxError
+// offset. Observed failure modes from production sessions:
+//
+//   - Extra '}' — model loses brace count (e.g. "}}]}" instead of "}]}")
+//   - Extra ']' — model adds a spurious array wrapper
+//   - Stray '\' — model emits an escape sequence outside of a string value
+//     (e.g. literal \n between tokens, or \" where " is expected)
+func tryRepairEditFileJSON(data []byte) ([]byte, bool) {
+	current := append([]byte(nil), data...) // defensive copy
+	for range 3 {
+		var synErr *json.SyntaxError
+		if err := json.Unmarshal(current, &json.RawMessage{}); err == nil {
+			return current, true
+		} else if !errors.As(err, &synErr) {
+			return nil, false
+		}
+
+		// json.SyntaxError.Offset is 1-based.
+		offset := int(synErr.Offset) - 1
+		if offset < 0 || offset >= len(current) {
+			return nil, false
+		}
+
+		ch := current[offset]
+		removeCount := 1
+
+		switch ch {
+		case '}', ']':
+			// Extra closing delimiter — just remove it.
+		case '\\':
+			// Stray escape sequence outside a string value. For \n, \t, \r
+			// both characters are garbage so remove them. For \" the quote
+			// is a valid structural character (string delimiter), so only
+			// strip the backslash.
+			if offset+1 < len(current) {
+				switch current[offset+1] {
+				case 'n', 't', 'r':
+					removeCount = 2
+				}
+			}
+		default:
+			return nil, false
+		}
+
+		repaired := make([]byte, 0, len(current)-removeCount)
+		repaired = append(repaired, current[:offset]...)
+		repaired = append(repaired, current[offset+removeCount:]...)
+		current = repaired
 	}
 
-	return nil
+	if json.Valid(current) {
+		return current, true
+	}
+	return nil, false
 }
 
 func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
@@ -245,7 +314,7 @@ func (t *FilesystemTool) Tools(context.Context) ([]tools.Tool, error) {
 			Description:  "Make line-based edits to a text file. Each edit replaces exact line sequences with new content.",
 			Parameters:   tools.MustSchemaFor[EditFileArgs](),
 			OutputSchema: tools.MustSchemaFor[string](),
-			Handler:      tools.NewHandler(t.handleEditFile),
+			Handler:      t.editFileHandler(),
 			Annotations: tools.ToolAnnotations{
 				Title: "Edit",
 			},
@@ -466,6 +535,24 @@ func countTreeNodes(node *fsx.TreeNode) (files, dirs int) {
 	return files, dirs
 }
 
+// editFileHandler returns a ToolHandler that parses edit_file arguments with
+// repair logic for malformed JSON, then delegates to handleEditFile.
+// This bypasses tools.NewHandler because Go's json.Unmarshal scanner rejects
+// structurally invalid JSON before calling any custom UnmarshalJSON method.
+func (t *FilesystemTool) editFileHandler() tools.ToolHandler {
+	return func(ctx context.Context, toolCall tools.ToolCall) (*tools.ToolCallResult, error) {
+		data := toolCall.Function.Arguments
+		if data == "" {
+			data = "{}"
+		}
+		args, err := ParseEditFileArgs([]byte(data))
+		if err != nil {
+			return nil, err
+		}
+		return t.handleEditFile(ctx, args)
+	}
+}
+
 func (t *FilesystemTool) handleEditFile(ctx context.Context, args EditFileArgs) (*tools.ToolCallResult, error) {
 	resolvedPath := t.resolvePath(args.Path)
 

pkg/tools/builtin/filesystem_test.go

@@ -260,7 +260,7 @@ func TestFilesystemTool_EditFile(t *testing.T) {
 	assert.Contains(t, result.Output, "old text not found")
 }
 
-func TestEditFileArgs_UnmarshalJSON(t *testing.T) {
+func TestParseEditFileArgs(t *testing.T) {
 	t.Parallel()
 
 	tests := []struct {
@@ -338,13 +338,61 @@ func TestEditFileArgs_UnmarshalJSON(t *testing.T) {
 				{OldText: "a", NewText: "b"},
 			},
 		},
+
+		// Malformed outer JSON — LLM brace/bracket counting errors.
+		{
+			name:     "repair: extra closing brace before array close",
+			input:    `{"path": "ci.yml", "edits": [{"oldText": "old", "newText": "new"}}]}`,
+			wantPath: "ci.yml",
+			wantEdits: []Edit{
+				{OldText: "old", NewText: "new"},
+			},
+		},
+		{
+			name:     "repair: extra closing brace with trailing newline",
+			input:    "{\"path\": \"ci.yml\", \"edits\": [{\"oldText\": \"old\", \"newText\": \"new\"}}]\n}",
+			wantPath: "ci.yml",
+			wantEdits: []Edit{
+				{OldText: "old", NewText: "new"},
+			},
+		},
+		{
+			name:     "repair: extra closing bracket (spurious array wrapper)",
+			input:    `{"path": "build.sh", "edits": [{"oldText": "a", "newText": "b"}]]}`,
+			wantPath: "build.sh",
+			wantEdits: []Edit{
+				{OldText: "a", NewText: "b"},
+			},
+		},
+		{
+			name:     "repair: stray backslash-n between tokens",
+			input:    "{\"path\": \"Dockerfile\", \"edits\": [{\"oldText\": \"a\", \"newText\": \"b\"}\\n]}",
+			wantPath: "Dockerfile",
+			wantEdits: []Edit{
+				{OldText: "a", NewText: "b"},
+			},
+		},
+		{
+			name:     "repair: stray backslash before property name",
+			input:    `{"path": "f.go", "edits": [{"oldText": "a", "newText": "b"},{\"oldText": "c", "newText": "d"}]}`,
+			wantPath: "f.go",
+			wantEdits: []Edit{
+				{OldText: "a", NewText: "b"},
+				{OldText: "c", NewText: "d"},
+			},
+		},
+		{
+			name:       "unrepairable garbage",
+			input:      `{totally broken <<<>>>`,
+			wantErr:    true,
+			wantErrMsg: "failed to parse edit_file arguments",
+		},
 	}
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
-			var args EditFileArgs
-			err := json.Unmarshal([]byte(tc.input), &args)
+			args, err := ParseEditFileArgs([]byte(tc.input))
 			if tc.wantErr {
 				require.Error(t, err)
 				assert.Contains(t, err.Error(), tc.wantErrMsg)

pkg/tui/components/tool/editfile/editfile.go

@@ -1,7 +1,6 @@
 package editfile
 
 import (
-	"encoding/json"
 	"fmt"
 
 	"github.com/docker/docker-agent/pkg/tools/builtin"
@@ -32,8 +31,8 @@ func render(
 	_ int,
 ) string {
 	// Parse tool arguments to extract the file path for display.
-	var args builtin.EditFileArgs
-	if err := json.Unmarshal([]byte(msg.ToolCall.Function.Arguments), &args); err != nil {
+	args, err := builtin.ParseEditFileArgs([]byte(msg.ToolCall.Function.Arguments))
+	if err != nil {
 		// If arguments cannot be parsed, fail silently to avoid breaking the TUI.
 		return ""
 	}
@@ -111,8 +110,8 @@ func renderCollapsed(
 	width,
 	_ int,
 ) string {
-	var args builtin.EditFileArgs
-	if err := json.Unmarshal([]byte(msg.ToolCall.Function.Arguments), &args); err != nil {
+	args, err := builtin.ParseEditFileArgs([]byte(msg.ToolCall.Function.Arguments))
+	if err != nil {
 		return ""
 	}
 

pkg/tui/components/tool/editfile/render.go

@@ -1,7 +1,6 @@
 package editfile
 
 import (
-	"encoding/json"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -118,8 +117,8 @@ func renderEditFile(toolCall tools.ToolCall, width int, splitView bool, toolStat
 }
 
 func renderEditFileUncached(toolCall tools.ToolCall, width int, splitView bool, toolStatus types.ToolStatus) string {
-	var args builtin.EditFileArgs
-	if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &args); err != nil {
+	args, err := builtin.ParseEditFileArgs([]byte(toolCall.Function.Arguments))
+	if err != nil {
 		return ""
 	}
 
@@ -169,8 +168,8 @@ func countDiffLines(toolCall tools.ToolCall, _ types.ToolStatus) (added, removed
 }
 
 func countDiffLinesUncached(toolCall tools.ToolCall) (added, removed int) {
-	var args builtin.EditFileArgs
-	if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &args); err != nil {
+	args, err := builtin.ParseEditFileArgs([]byte(toolCall.Function.Arguments))
+	if err != nil {
 		return 0, 0
 	}