Better fuzzy json parsing

pkg/tui/components/toolcommon/common.go

@@ -23,7 +23,7 @@ func ParseArgs[T any](args string) (T, error) {
 		return result, nil
 	}
 
-	for _, fixed := range tryFixPartialJSON(args) {
+	if fixed, ok := tryFixPartialJSON(args); ok {
 		if partialErr := json.Unmarshal([]byte(fixed), &result); partialErr == nil {
 			return result, nil
 		}
@@ -32,11 +32,59 @@ func ParseArgs[T any](args string) (T, error) {
 	return result, err
 }
 
-func tryFixPartialJSON(s string) []string {
-	if s == "" {
-		return nil
+// tryFixPartialJSON attempts to complete a partial JSON object by closing
+// any unclosed strings, arrays, and objects. Returns the fixed JSON and
+// true if a fix was attempted, or the original string and false if input
+// is empty or not a valid JSON object start.
+func tryFixPartialJSON(s string) (string, bool) {
+	if s == "" || s[0] != '{' {
+		return s, false
 	}
-	return []string{s + "\"}", s + "}"}
+
+	var result strings.Builder
+	result.WriteString(s)
+
+	inString := false
+	escaped := false
+	var stack []byte
+
+	for _, r := range s {
+		if escaped {
+			escaped = false
+			continue
+		}
+		if r == '\\' && inString {
+			escaped = true
+			continue
+		}
+		if r == '"' {
+			inString = !inString
+			continue
+		}
+		if inString {
+			continue
+		}
+		switch r {
+		case '{':
+			stack = append(stack, '}')
+		case '[':
+			stack = append(stack, ']')
+		case '}', ']':
+			if len(stack) > 0 {
+				stack = stack[:len(stack)-1]
+			}
+		}
+	}
+
+	if inString {
+		result.WriteByte('"')
+	}
+
+	for i := len(stack) - 1; i >= 0; i-- {
+		result.WriteByte(stack[i])
+	}
+
+	return result.String(), true
 }
 
 // ExtractField creates an argument extractor function that parses JSON and extracts a field.

pkg/tui/components/toolcommon/common_test.go

@@ -9,25 +9,89 @@ import (
 
 func TestTryFixPartialJSON(t *testing.T) {
 	tests := []struct {
-		name     string
-		input    string
-		expected []string
+		name      string
+		input     string
+		expected  string
+		shouldFix bool
 	}{
 		{
-			name:     "empty string",
-			input:    "",
-			expected: nil,
+			name:      "empty string",
+			input:     "",
+			expected:  "",
+			shouldFix: false,
+		},
+		{
+			name:      "not json object",
+			input:     "hello",
+			expected:  "hello",
+			shouldFix: false,
+		},
+		{
+			name:      "just opening brace",
+			input:     `{`,
+			expected:  `{}`,
+			shouldFix: true,
+		},
+		{
+			name:      "partial key",
+			input:     `{"path`,
+			expected:  `{"path"}`,
+			shouldFix: true,
+		},
+		{
+			name:      "key with colon",
+			input:     `{"path":`,
+			expected:  `{"path":}`,
+			shouldFix: true,
+		},
+		{
+			name:      "incomplete string value",
+			input:     `{"path": "/tmp/fi`,
+			expected:  `{"path": "/tmp/fi"}`,
+			shouldFix: true,
+		},
+		{
+			name:      "complete string missing brace",
+			input:     `{"path": "/tmp/file"`,
+			expected:  `{"path": "/tmp/file"}`,
+			shouldFix: true,
 		},
 		{
-			name:     "partial JSON",
-			input:    `{"key": "val`,
-			expected: []string{`{"key": "val"}`, `{"key": "val}`},
+			name:      "trailing comma",
+			input:     `{"path": "/tmp/file",`,
+			expected:  `{"path": "/tmp/file",}`,
+			shouldFix: true,
+		},
+		{
+			name:      "nested object incomplete",
+			input:     `{"outer": {"inner": "val`,
+			expected:  `{"outer": {"inner": "val"}}`,
+			shouldFix: true,
+		},
+		{
+			name:      "array incomplete",
+			input:     `{"paths": ["/tmp/a", "/tmp/b`,
+			expected:  `{"paths": ["/tmp/a", "/tmp/b"]}`,
+			shouldFix: true,
+		},
+		{
+			name:      "escaped quote in string",
+			input:     `{"msg": "hello \"world`,
+			expected:  `{"msg": "hello \"world"}`,
+			shouldFix: true,
+		},
+		{
+			name:      "complete json",
+			input:     `{"path": "/tmp/file"}`,
+			expected:  `{"path": "/tmp/file"}`,
+			shouldFix: true,
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			result := tryFixPartialJSON(tt.input)
+			result, ok := tryFixPartialJSON(tt.input)
+			assert.Equal(t, tt.shouldFix, ok)
 			assert.Equal(t, tt.expected, result)
 		})
 	}
@@ -88,6 +152,20 @@ func TestParsePartialArgs(t *testing.T) {
 			wantCmd:  "",
 			wantErr:  false,
 		},
+		{
+			name:     "nested object in progress",
+			input:    `{"path": "/tmp", "nested": {"key": "val`,
+			wantPath: "/tmp",
+			wantCmd:  "",
+			wantErr:  false,
+		},
+		{
+			name:     "array value in progress",
+			input:    `{"path": "/tmp", "items": ["a", "b`,
+			wantPath: "/tmp",
+			wantCmd:  "",
+			wantErr:  false,
+		},
 	}
 
 	for _, tt := range tests {