fix: repair malformed JSON in edit_file tool call arguments

LLMs occasionally generate structurally invalid JSON for edit_file
arguments when the text payload contains characters that confuse brace
counting (YAML, Dockerfiles, shell scripts). The most common pattern is
an extra closing brace or bracket near the end of the arguments.

Go's json.Unmarshal scanner rejects this before any custom UnmarshalJSON
method runs, so the existing double-serialization fix never fires for
these cases.

Add tryRepairEditFileJSON which iteratively removes the offending
character at the json.SyntaxError offset (up to 3 rounds), handling
extra }, extra ], stray backslash-n, and stray backslash before quotes.

Replace NewHandler with a custom editFileHandler that calls
ParseEditFileArgs (which does repair then unmarshal) before delegating
to handleEditFile. Update all other call sites (ACP filesystem handler,
TUI rendering) to use ParseEditFileArgs as well.

Author: trungutt

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
 	}