fix(go/internal): refine JSON markdown detection

go/ai/format_array.go

@@ -91,7 +91,7 @@ func (a arrayHandler) ParseMessage(m *Message) (*Message, error) {
 		}
 
 		var newParts []*Part
-		lines := base.GetJsonObjectLines(accumulatedText.String())
+		lines := base.GetJSONObjectLines(accumulatedText.String())
 		for _, line := range lines {
 			var schemaBytes []byte
 			schemaBytes, err := json.Marshal(a.config.Schema["items"])

go/ai/format_jsonl.go

@@ -92,7 +92,7 @@ func (j jsonlHandler) ParseMessage(m *Message) (*Message, error) {
 		}
 
 		var newParts []*Part
-		lines := base.GetJsonObjectLines(accumulatedText.String())
+		lines := base.GetJSONObjectLines(accumulatedText.String())
 		for _, line := range lines {
 			if j.config.Schema != nil {
 				var schemaBytes []byte

go/internal/base/json.go

@@ -119,22 +119,23 @@ func SchemaAsMap(s *jsonschema.Schema) map[string]any {
 	return m
 }
 
-var jsonMarkdownRegex = regexp.MustCompile("```(json)?((\n|.)*?)```")
+// jsonMarkdownRegex specifically looks for "json" language identifier
+var jsonMarkdownRegex = regexp.MustCompile("(?s)```json(.*?)```")
 
 // ExtractJSONFromMarkdown returns the contents of the first fenced code block in
 // the markdown text md. If there is none, it returns md.
 func ExtractJSONFromMarkdown(md string) string {
-	// TODO: improve this
 	matches := jsonMarkdownRegex.FindStringSubmatch(md)
-	if matches == nil {
+	if len(matches) < 2 {
 		return md
 	}
-	return matches[2]
+	// capture group 1 matches the actual fenced JSON block
+	return matches[1]
 }
 
-// GetJsonObjectLines splits a string by newlines, trims whitespace from each line,
+// GetJSONObjectLines splits a string by newlines, trims whitespace from each line,
 // and returns a slice containing only the lines that start with '{'.
-func GetJsonObjectLines(text string) []string {
+func GetJSONObjectLines(text string) []string {
 	jsonText := ExtractJSONFromMarkdown(text)
 
 	// Handle both actual "\n" newline strings, as well as newline bytes

go/internal/base/json_test.go

@@ -41,15 +41,20 @@ func TestExtractJSONFromMarkdown(t *testing.T) {
 		{
 			desc: "simple markdown",
 			in:   "```foo bar```",
-			want: "foo bar",
+			want: "```foo bar```",
+		},
+		{
+			desc: "empty markdown",
+			in:   "``` ```",
+			want: "``` ```",
 		},
 		{
 			desc: "json markdown",
 			in:   "```json{\"a\":1}```",
 			want: "{\"a\":1}",
 		},
 		{
-			desc: "json multipline markdown",
+			desc: "json multiple line markdown",
 			in:   "```json\n{\"a\": 1}\n```",
 			want: "\n{\"a\": 1}\n",
 		},
@@ -58,6 +63,21 @@ func TestExtractJSONFromMarkdown(t *testing.T) {
 			in:   "```json{\"a\":\n1}```\n```json\n{\"b\":\n1}```",
 			want: "{\"a\":\n1}",
 		},
+		{
+			desc: "yaml markdown",
+			in:   "```yaml\nkey: 1\nanother-key: 2```",
+			want: "```yaml\nkey: 1\nanother-key: 2```",
+		},
+		{
+			desc: "yaml + json markdown",
+			in:   "```yaml\nkey: 1\nanother-key: 2``` ```json\n{\"a\": 1}\n```",
+			want: "\n{\"a\": 1}\n",
+		},
+		{
+			desc: "json + yaml markdown",
+			in:   "```json\n{\"a\": 1}\n``` ```yaml\nkey: 1\nanother-key: 2```",
+			want: "\n{\"a\": 1}\n",
+		},
 	}
 	for _, tc := range tests {
 		t.Run(tc.desc, func(t *testing.T) {