feat(extract): add brace-balanced JSON object/array extractor

Port the rtk extract_json_object brace-balancer (src/parser/mod.rs) to
native Go. The extractor finds the first balanced JSON object in text
that may contain surrounding prose, markdown code fences, or other LLM
output artifacts.

Algorithm:
  1. Scan for the first '{' (or '[') that is not inside a double-quoted
     string.
  2. Brace-balance, supporting nested objects, arrays, and strings.
  3. Skip string contents (respecting backslash escapes).
  4. Stop when depth returns to zero.

Apostrophes in English prose (e.g. "Here's the JSON:") are NOT
treated as string delimiters — only standard JSON double-quotes are.
This avoids false positives when LLM output is wrapped in natural
language.

Functions:
  - ExtractJSONObject(s) Result — first balanced {…} in s
  - ExtractJSONArray(s) Result — first balanced […] in s
  - ExtractAllJSONObjects(s) []Result — every top-level object

Exposed at the top-level tok API as tok.ExtractJSON, tok.ExtractJSONArray,
and tok.ExtractAllJSON.

Source: rtk-ai/rtk, src/parser/mod.rs (extract_json_object).
Tests: 11 cases including LLM output with prose, markdown fences,
nested objects, escaped quotes, and unterminated inputs.

Author: Patel230

extract.go

@@ -0,0 +1,40 @@
+// JSON extraction from prose — public API.
+//
+// Wraps internal/extract to expose brace-balanced JSON object and
+// array extraction as a top-level tok API. Useful for parsing LLM
+// output that contains JSON embedded in surrounding prose.
+package tok
+
+import "github.com/GrayCodeAI/tok/internal/extract"
+
+// JSONExtract is the result of an extraction attempt.
+type JSONExtract = extract.Result
+
+// ExtractJSON returns the first complete JSON object found in text,
+// or JSONExtract{Found: false} if no balanced object exists.
+//
+// Handles nested objects, arrays, escaped quotes, and double-quoted
+// strings. Skips preceding and trailing prose. Apostrophes in English
+// prose (e.g. "Here's the JSON:") are NOT treated as string delimiters
+// (only standard JSON double-quotes are).
+//
+// Example:
+//
+//	out := tok.ExtractJSON(`Sure! Here's the data: {"name": "alice"}`)
+//	// out.JSON = `{"name": "alice"}`, out.Found = true
+func ExtractJSON(text string) JSONExtract {
+	return extract.ExtractJSONObject(text)
+}
+
+// ExtractJSONArray is the array counterpart of ExtractJSON. It looks
+// for a balanced `[...]` instead of `{...}`.
+func ExtractJSONArray(text string) JSONExtract {
+	return extract.ExtractJSONArray(text)
+}
+
+// ExtractAllJSON returns every top-level balanced JSON object in
+// text, in source order. Useful for parsing LLM output that contains
+// multiple JSON snippets.
+func ExtractAllJSON(text string) []JSONExtract {
+	return extract.ExtractAllJSONObjects(text)
+}

extract_test.go

@@ -0,0 +1,63 @@
+package tok_test
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/GrayCodeAI/tok"
+)
+
+func TestExtractJSON_TopLevel(t *testing.T) {
+	tests := []struct {
+		name  string
+		input string
+		want  string
+		found bool
+	}{
+		{"bare", `{"a": 1}`, `{"a": 1}`, true},
+		{"with prose", `Sure! {"a": 1} done`, `{"a": 1}`, true},
+		{"not found", "no braces", "", false},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			r := tok.ExtractJSON(tc.input)
+			if r.Found != tc.found {
+				t.Errorf("found mismatch: got %v want %v", r.Found, tc.found)
+			}
+			if r.JSON != tc.want {
+				t.Errorf("JSON mismatch: got %q want %q", r.JSON, tc.want)
+			}
+		})
+	}
+}
+
+func TestExtractJSON_ApostropheInProse(t *testing.T) {
+	// Regression: apostrophes in English prose must not be confused
+	// with string delimiters.
+	in := "Here's the JSON: {\"a\": 1}"
+	r := tok.ExtractJSON(in)
+	if !r.Found {
+		t.Fatal("expected to find JSON")
+	}
+	if !strings.Contains(r.JSON, `"a"`) {
+		t.Errorf("unexpected extraction: %q", r.JSON)
+	}
+}
+
+func TestExtractJSONArray_TopLevel(t *testing.T) {
+	r := tok.ExtractJSONArray(`Result: [1, 2, 3]`)
+	if !r.Found {
+		t.Fatal("expected to find array")
+	}
+	if r.JSON != "[1, 2, 3]" {
+		t.Errorf("unexpected array: %q", r.JSON)
+	}
+}
+
+func TestExtractAllJSON_TopLevel(t *testing.T) {
+	in := `First: {"a": 1} middle {"b": 2} tail`
+	out := tok.ExtractAllJSON(in)
+	if len(out) != 2 {
+		t.Fatalf("expected 2 objects, got %d", len(out))
+	}
+}

internal/extract/extract.go

@@ -0,0 +1,218 @@
+// Package extract provides utilities for extracting structured data
+// (JSON objects, arrays) from text that may contain prose around it.
+//
+// Common use case: an LLM returns something like
+//
+//	"Sure! Here's the JSON you asked for:
+//	{\"name\": \"alice\", \"age\": 30}
+//	Hope that helps!"
+//
+// and the caller wants just the {"name": "alice", "age": 30} part. The
+// ExtractJSONObject function performs brace-balanced extraction that
+// handles nested objects, arrays, escaped quotes, and strings.
+//
+// Source: rtk-ai/rtk, src/parser/mod.rs (extract_json_object).
+// Ported to native Go.
+package extract
+
+// Result describes a successful extraction. JSON is the extracted
+// substring (no surrounding prose). Start is the byte offset of the
+// opening brace in the input. End is the byte offset one past the
+// closing brace. Found is true when extraction succeeded.
+type Result struct {
+	JSON  string
+	Start int
+	End   int
+	Found bool
+}
+
+// ExtractJSONObject returns the first complete JSON object found in s,
+// or Result{Found: false} if no balanced object exists.
+//
+// The function:
+//
+//  1. Scans for the first '{' that is not inside a string literal.
+//  2. Tracks brace depth, supporting nested objects and arrays.
+//  3. Skips over string contents (respecting both single and double
+//     quotes, with backslash escape sequences).
+//  4. Stops when depth returns to zero.
+//
+// The output preserves the exact bytes from the input (no
+// normalization). If the input is a valid bare JSON object, the output
+// equals the input.
+func ExtractJSONObject(s string) Result {
+	return extractFirst(s, '{', '}')
+}
+
+// ExtractJSONArray is the array counterpart of ExtractJSONObject.
+// It looks for a balanced '[...]' instead of '{...}'.
+func ExtractJSONArray(s string) Result {
+	return extractFirst(s, '[', ']')
+}
+
+// extractFirst is the common implementation for both object and array
+// extraction. The opener and closer parameters are byte literals.
+//
+// Only double-quote strings are recognized (standard JSON). Single
+// quotes / apostrophes in English prose are ignored, which is what
+// callers want for LLM output where "Here's the JSON:" is common
+// prose around the actual JSON.
+func extractFirst(s string, opener, closer byte) Result {
+	// Phase 1: find the first opener that isn't inside a string.
+	start := findUnquotedOpener(s, opener)
+	if start < 0 {
+		return Result{Found: false}
+	}
+
+	// Phase 2: brace-balance, skipping string contents.
+	depth := 0
+	i := start
+	for i < len(s) {
+		c := s[i]
+		switch c {
+		case '\\':
+			// Skip the escaped character. This handles \" inside a
+			// string, \\ for a literal backslash, and \n, \t, etc.
+			i += 2
+			continue
+		case '"':
+			// Skip the entire double-quoted string literal.
+			i = skipDoubleQuotedString(s, i)
+			continue
+		case opener:
+			depth++
+			i++
+		case closer:
+			depth--
+			i++
+			if depth == 0 {
+				return Result{
+					JSON:  s[start:i],
+					Start: start,
+					End:   i,
+					Found: true,
+				}
+			}
+		default:
+			i++
+		}
+	}
+	// Unterminated; return what we have so far if depth > 0.
+	if depth > 0 {
+		return Result{
+			JSON:  s[start:],
+			Start: start,
+			End:   len(s),
+			Found: true,
+		}
+	}
+	return Result{Found: false}
+}
+
+// findUnquotedOpener returns the index of the first byte equal to c
+// that is not inside a double-quoted string literal, or -1 if none
+// exists. Single quotes / apostrophes in English prose are NOT treated
+// as string delimiters.
+func findUnquotedOpener(s string, c byte) int {
+	i := 0
+	for i < len(s) {
+		switch s[i] {
+		case '\\':
+			i += 2
+			continue
+		case '"':
+			i = skipDoubleQuotedString(s, i)
+			continue
+		}
+		if s[i] == c {
+			return i
+		}
+		i++
+	}
+	return -1
+}
+
+// skipDoubleQuotedString advances past a double-quoted string literal
+// starting at s[i] (which must be a `"`). Returns the index one past
+// the closing `"` (or len(s) if unterminated).
+func skipDoubleQuotedString(s string, i int) int {
+	i++ // skip opening quote
+	for i < len(s) {
+		switch s[i] {
+		case '\\':
+			i += 2 // skip escaped char
+		case '"':
+			return i + 1
+		default:
+			i++
+		}
+	}
+	return i
+}
+
+// ExtractAllJSONObjects returns every top-level balanced JSON object
+// in s, in source order. Useful when parsing LLM output that may
+// contain several JSON snippets (e.g. a list of objects in prose).
+func ExtractAllJSONObjects(s string) []Result {
+	var out []Result
+	i := 0
+	for i < len(s) {
+		// Find next opener from position i
+		start := findUnquotedOpener(s[i:], '{')
+		if start < 0 {
+			break
+		}
+		start += i
+		// Try to extract starting from `start`
+		r := extractFrom(s, start, '{', '}')
+		if !r.Found {
+			break
+		}
+		out = append(out, r)
+		i = r.End
+	}
+	return out
+}
+
+// extractFrom is a helper for ExtractAllJSONObjects that starts scanning
+// from a known opener position (no opener-search needed).
+func extractFrom(s string, start int, opener, closer byte) Result {
+	depth := 0
+	i := start
+	for i < len(s) {
+		c := s[i]
+		switch c {
+		case '\\':
+			i += 2
+			continue
+		case '"':
+			i = skipDoubleQuotedString(s, i)
+			continue
+		case opener:
+			depth++
+			i++
+		case closer:
+			depth--
+			i++
+			if depth == 0 {
+				return Result{
+					JSON:  s[start:i],
+					Start: start,
+					End:   i,
+					Found: true,
+				}
+			}
+		default:
+			i++
+		}
+	}
+	if depth > 0 {
+		return Result{
+			JSON:  s[start:],
+			Start: start,
+			End:   len(s),
+			Found: true,
+		}
+	}
+	return Result{Found: false}
+}