mlxrunner: compile structured output as xgrammar structural tags

The runner compiled a format as a JSON Schema, the one grammar kind its
xgrammar binding exposed. A structural tag holds a schema as one node of
a larger tree and also expresses what a schema cannot: free text around
constrained spans, a thinking region that closes before constrained
content, tool calls pinned to their schemas.

The runner now compiles structural tags only; its client wraps the API's
formats into one, which compiles to the same grammar as before. The JSON
token and vocabulary caps go with it: neither bounds compile cost, which
follows the grammar's state count. The byte and nesting caps stay.

x/mlxrunner/client.go

@@ -155,12 +155,29 @@ func (c *Client) Close() error {
 	return nil
 }
 
+// requestGrammar returns the structural tag the runner decodes under: the
+// API's format wrapped into a json_schema tag.
+func requestGrammar(req llm.CompletionRequest) json.RawMessage {
+	schema := req.Format
+	switch string(schema) {
+	case ``, `null`, `""`:
+		return nil
+	case `"json"`:
+		// The API documents "json" as producing a JSON object.
+		schema = json.RawMessage(`{"type":"object"}`)
+	}
+	tag := make(json.RawMessage, 0, len(schema)+64)
+	tag = append(tag, `{"type":"structural_tag","format":{"type":"json_schema","json_schema":`...)
+	tag = append(tag, schema...)
+	return append(tag, `}}`...)
+}
+
 // Completion implements llm.LlamaServer.
 func (c *Client) Completion(ctx context.Context, req llm.CompletionRequest, fn func(llm.CompletionResponse)) error {
 	creq := CompletionRequest{
 		Prompt:                     req.Prompt,
 		Media:                      req.Media,
-		Format:                     req.Format,
+		Format:                     requestGrammar(req),
 		Logprobs:                   req.Logprobs,
 		TopLogprobs:                req.TopLogprobs,
 		IncludeIntermediateMetrics: req.IncludeIntermediateMetrics,

x/mlxrunner/client_test.go

@@ -16,6 +16,32 @@ func testIntPtr(v int) *int {
 	return &v
 }
 
+func TestRequestGrammar(t *testing.T) {
+	schema := `{"type":"object","properties":{"answer":{"type":"string"}}}`
+	tag := `{"type":"structural_tag","format":{"type":"json_schema","json_schema":` + schema + `}}`
+	for _, tt := range []struct {
+		name string
+		req  llm.CompletionRequest
+		want string
+	}{
+		{name: "unset"},
+		{name: "null", req: llm.CompletionRequest{Format: json.RawMessage(`null`)}},
+		{name: "empty", req: llm.CompletionRequest{Format: json.RawMessage(`""`)}},
+		{
+			name: "json",
+			req:  llm.CompletionRequest{Format: json.RawMessage(`"json"`)},
+			want: `{"type":"structural_tag","format":{"type":"json_schema","json_schema":{"type":"object"}}}`,
+		},
+		{name: "schema", req: llm.CompletionRequest{Format: json.RawMessage(schema)}, want: tag},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := string(requestGrammar(tt.req)); got != tt.want {
+				t.Fatalf("requestGrammar = %s, want %s", got, tt.want)
+			}
+		})
+	}
+}
+
 func TestClientCompletionRequestsIntermediateMetrics(t *testing.T) {
 	var request CompletionRequest
 	want := CompletionResponse{

x/mlxrunner/grammar.go

@@ -21,14 +21,9 @@ import (
 	"github.com/ollama/ollama/x/tokenizer"
 )
 
-const maxGrammarVocabSize = 1 << 20
-
 const (
-	maxGrammarSchemaBytes = 1 << 20
-	maxGrammarSchemaDepth = 128
-	// The token cap bounds grammar compile cost, which the serial runner
-	// pays as head-of-line blocking of the request queue.
-	maxGrammarSchemaTokens = 1 << 14
+	maxGrammarBytes = 1 << 20
+	maxGrammarDepth = 128
 )
 
 const (
@@ -92,9 +87,6 @@ func validateGrammarVocab(logitsWidth, tokenizerSize int) error {
 	if logitsWidth <= 0 {
 		return fmt.Errorf("invalid model logits width %d", logitsWidth)
 	}
-	if logitsWidth > maxGrammarVocabSize {
-		return fmt.Errorf("model logits width %d exceeds structured output limit %d", logitsWidth, maxGrammarVocabSize)
-	}
 	return nil
 }
 
@@ -133,81 +125,64 @@ func (e *grammarEngine) close() {
 // returns nil. Safe on a nil subsystem, which reports structured output
 // unavailable.
 func (e *grammarEngine) prepare(format json.RawMessage) (*grammarCompilation, error) {
-	spec, err := parseGrammar(format)
-	if err != nil || spec == nil {
+	source, err := parseGrammar(format)
+	if err != nil || source == "" {
 		return nil, err
 	}
 	if e == nil {
 		return nil, api.StatusError{StatusCode: http.StatusNotImplemented, ErrorMessage: "structured output is unavailable"}
 	}
-	return e.compile(spec), nil
+	return e.compile(source), nil
 }
 
-type grammarSpec struct {
-	kind   xgrammar.Kind
-	source string
-}
-
-func parseGrammar(format json.RawMessage) (*grammarSpec, error) {
-	if len(format) > 0 {
-		switch string(format) {
-		case `null`, `""`:
-			return nil, nil
-		case `"json"`:
-			// The API documents "json" as producing a JSON object; the engine's
-			// builtin JSON grammar would also admit arrays and bare values.
-			return &grammarSpec{kind: xgrammar.JSONSchema, source: `{"type":"object"}`}, nil
-		default:
-			if format[0] != '{' {
-				return nil, errors.New("invalid format: expected \"json\" or a valid JSON Schema object")
-			}
-			if err := validateGrammarSchema(format); err != nil {
-				return nil, fmt.Errorf("invalid JSON Schema: %w", err)
-			}
-			return &grammarSpec{kind: xgrammar.JSONSchema, source: string(format)}, nil
-		}
+// parseGrammar returns the format's structural tag, or "" when the format
+// asks for no structured output.
+func parseGrammar(format json.RawMessage) (string, error) {
+	switch string(format) {
+	case ``, `null`, `""`:
+		return "", nil
 	}
-	return nil, nil
-}
-
-func validateGrammarSchema(schema []byte) error {
-	if len(schema) > maxGrammarSchemaBytes {
-		return fmt.Errorf("schema is %d bytes; limit is %d", len(schema), maxGrammarSchemaBytes)
+	if len(format) > maxGrammarBytes {
+		return "", fmt.Errorf("invalid format: grammar is %d bytes; limit is %d", len(format), maxGrammarBytes)
 	}
-	if !utf8.Valid(schema) {
-		return errors.New("schema is not valid UTF-8")
+	if !utf8.Valid(format) {
+		return "", errors.New("invalid format: grammar is not valid UTF-8")
+	}
+	if format[0] != '{' {
+		return "", errors.New("invalid format: expected a structural tag")
 	}
 
-	decoder := json.NewDecoder(bytes.NewReader(schema))
+	decoder := json.NewDecoder(bytes.NewReader(format))
 	decoder.UseNumber()
-	first, err := decoder.Token()
-	if err != nil {
-		return fmt.Errorf("invalid JSON: %w", err)
-	}
-	if first != json.Delim('{') {
-		return errors.New("schema must be a JSON object")
+	if _, err := decoder.Token(); err != nil {
+		return "", fmt.Errorf("invalid format: %w", err)
 	}
 
-	tokens, depth := 1, 1
+	depth, wantKey, key, structuralTag := 1, true, "", false
 	for depth > 0 {
 		token, err := decoder.Token()
 		if err != nil {
 			if err == io.EOF {
-				return errors.New("unexpected end of JSON")
+				return "", errors.New("invalid format: unexpected end of JSON")
 			}
-			return fmt.Errorf("invalid JSON: %w", err)
+			return "", fmt.Errorf("invalid format: %w", err)
 		}
-		tokens++
-		if tokens > maxGrammarSchemaTokens {
-			return fmt.Errorf("schema contains more than %d JSON tokens", maxGrammarSchemaTokens)
+		delim, isDelim := token.(json.Delim)
+		if depth == 1 && !(isDelim && (delim == '}' || delim == ']')) {
+			if wantKey {
+				key, _ = token.(string)
+			} else if key == "type" {
+				value, _ := token.(string)
+				structuralTag = value == "structural_tag"
+			}
+			wantKey = !wantKey
 		}
-
-		if delim, ok := token.(json.Delim); ok {
+		if isDelim {
 			switch delim {
 			case '{', '[':
 				depth++
-				if depth > maxGrammarSchemaDepth {
-					return fmt.Errorf("schema nesting exceeds %d levels", maxGrammarSchemaDepth)
+				if depth > maxGrammarDepth {
+					return "", fmt.Errorf("invalid format: grammar nesting exceeds %d levels", maxGrammarDepth)
 				}
 			case '}', ']':
 				depth--
@@ -217,14 +192,17 @@ func validateGrammarSchema(schema []byte) error {
 
 	if _, err := decoder.Token(); err != io.EOF {
 		if err != nil {
-			return fmt.Errorf("invalid JSON after schema object: %w", err)
+			return "", fmt.Errorf("invalid format: %w", err)
 		}
-		return errors.New("schema contains more than one JSON value")
+		return "", errors.New("invalid format: grammar contains more than one JSON value")
 	}
-	return nil
+	if !structuralTag {
+		return "", errors.New("invalid format: expected a structural tag")
+	}
+	return string(format), nil
 }
 
-func (e *grammarEngine) compile(spec *grammarSpec) *grammarCompilation {
+func (e *grammarEngine) compile(source string) *grammarCompilation {
 	c := &grammarCompilation{done: make(chan struct{})}
 	go func() {
 		defer close(c.done)
@@ -242,10 +220,10 @@ func (e *grammarEngine) compile(spec *grammarSpec) *grammarCompilation {
 			c.err = errors.New("grammar engine closed")
 			return
 		}
-		matcher, err := e.compiler.Compile(spec.kind, spec.source)
+		matcher, err := e.compiler.Compile(source)
 		if err != nil {
-			// A schema can pass validateGrammarSchema yet be rejected by
-			// the engine (e.g. an empty enum); that is still a request error.
+			// A grammar can pass parseGrammar yet be rejected by the engine
+			// (e.g. an empty enum); that is still a request error.
 			c.err = api.StatusError{
 				StatusCode:   http.StatusBadRequest,
 				ErrorMessage: fmt.Sprintf("invalid structured output grammar: %v", err),

x/mlxrunner/grammar_mask_test.go

@@ -102,7 +102,7 @@ func testDraftGrammar(t *mlxtest.T, schema string) (*grammarEngine, *grammar) {
 		t.Fatal(err)
 	}
 	t.Cleanup(compiler.Close)
-	m, err := compiler.Compile(xgrammar.JSONSchema, schema)
+	m, err := compiler.Compile(schemaTag(schema))
 	if err != nil {
 		t.Fatal(err)
 	}

x/mlxrunner/grammar_test.go

@@ -20,50 +20,67 @@ import (
 	"github.com/ollama/ollama/x/mlxrunner/xgrammar"
 )
 
+// schemaTag wraps a JSON Schema into the structural tag the MLX client
+// sends for a schema format.
+func schemaTag(schema string) string {
+	return `{"type":"structural_tag","format":{"type":"json_schema","json_schema":` + schema + `}}`
+}
+
+func nestedGrammar(depth int) string {
+	return `{"type":"structural_tag","format":` + strings.Repeat("[", depth-1) + `0` + strings.Repeat("]", depth-1) + `}`
+}
+
 func TestParseGrammar(t *testing.T) {
+	invalidUTF8 := `{"type":"structural_tag","value":"` + "\xff" + `"}`
 	tests := []struct {
 		name    string
 		format  string
-		kind    xgrammar.Kind
-		source  string
-		want    bool
-		wantErr bool
+		want    string
+		wantErr string
 	}{
 		{name: "unset"},
 		{name: "null", format: `null`},
 		{name: "empty", format: `""`},
-		{name: "json", format: `"json"`, kind: xgrammar.JSONSchema, source: `{"type":"object"}`, want: true},
-		{name: "schema", format: `{"type":"integer"}`, kind: xgrammar.JSONSchema, source: `{"type":"integer"}`, want: true},
-		{name: "unsupported string", format: `"xml"`, wantErr: true},
-		{name: "whitespace JSON", format: ` "json" `, wantErr: true},
-		{name: "whitespace schema", format: ` {"type":"integer"} `, wantErr: true},
-		{name: "array", format: `[]`, wantErr: true},
-		{name: "invalid json", format: `{`, wantErr: true},
+		{name: "structural tag", format: schemaTag(`{"type":"integer"}`), want: schemaTag(`{"type":"integer"}`)},
+		{name: "type after nested members", format: `{"format":{"type":"any_text","excludes":["type"]},"type":"structural_tag"}`, want: `{"format":{"type":"any_text","excludes":["type"]},"type":"structural_tag"}`},
+		{name: "maximum depth", format: nestedGrammar(maxGrammarDepth), want: nestedGrammar(maxGrammarDepth)},
+		{name: "too deep", format: nestedGrammar(maxGrammarDepth + 1), wantErr: "nesting exceeds"},
+		{name: "too large", format: `{"value":"` + strings.Repeat("x", maxGrammarBytes) + `"}`, wantErr: "limit is 1048576"},
+		{name: "invalid UTF-8", format: invalidUTF8, wantErr: "not valid UTF-8"},
+		{name: "json", format: `"json"`, wantErr: "expected a structural tag"},
+		{name: "schema", format: `{"type":"integer"}`, wantErr: "expected a structural tag"},
+		{name: "untyped object", format: `{"format":{}}`, wantErr: "expected a structural tag"},
+		{name: "nested type only", format: `{"format":{"type":"structural_tag"}}`, wantErr: "expected a structural tag"},
+		{name: "whitespace", format: ` ` + schemaTag(`{}`) + ` `, wantErr: "expected a structural tag"},
+		{name: "array", format: `[]`, wantErr: "expected a structural tag"},
+		{name: "trailing value", format: `{"type":"structural_tag"} {}`, wantErr: "more than one JSON value"},
+		{name: "malformed", format: `{"type":`, wantErr: "unexpected end"},
+		{name: "invalid json", format: `{`, wantErr: "unexpected end"},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			got, err := parseGrammar(json.RawMessage(tt.format))
-			if (err != nil) != tt.wantErr {
-				t.Fatalf("parseGrammar error = %v, wantErr %v", err, tt.wantErr)
-			}
-			if tt.wantErr {
+			if tt.wantErr != "" {
+				if err == nil || !strings.Contains(err.Error(), tt.wantErr) {
+					t.Fatalf("parseGrammar error = %v, want containing %q", err, tt.wantErr)
+				}
 				return
 			}
-			if (got != nil) != tt.want {
-				t.Fatalf("parseGrammar = %#v, want present %v", got, tt.want)
+			if err != nil {
+				t.Fatal(err)
 			}
-			if got != nil && (got.kind != tt.kind || got.source != tt.source) {
-				t.Errorf("parseGrammar = %#v, want kind %v source %q", got, tt.kind, tt.source)
+			if got != tt.want {
+				t.Errorf("parseGrammar = %q, want %q", got, tt.want)
 			}
 		})
 	}
 }
 
 func TestParseGrammarDoesNotEchoOversizedInput(t *testing.T) {
-	format := json.RawMessage("{" + strings.Repeat("x", maxGrammarSchemaBytes))
+	format := json.RawMessage("{" + strings.Repeat("x", maxGrammarBytes))
 	_, err := parseGrammar(format)
-	if err == nil || !strings.Contains(err.Error(), "schema is 1048577 bytes") {
+	if err == nil || !strings.Contains(err.Error(), "grammar is 1048577 bytes") {
 		t.Fatalf("parseGrammar error = %v, want bounded size error", err)
 	}
 	if len(err.Error()) > 256 {
@@ -71,62 +88,11 @@ func TestParseGrammarDoesNotEchoOversizedInput(t *testing.T) {
 	}
 }
 
-func nestedGrammarSchema(depth int) []byte {
-	return []byte(`{"value":` + strings.Repeat("[", depth-1) + `0` + strings.Repeat("]", depth-1) + `}`)
-}
-
-func grammarSchemaArray(entries int) []byte {
-	var b strings.Builder
-	b.WriteString(`{"enum":[`)
-	for i := range entries {
-		if i > 0 {
-			b.WriteByte(',')
-		}
-		b.WriteByte('0')
-	}
-	b.WriteString(`]}`)
-	return []byte(b.String())
-}
-
-func TestValidateGrammarSchema(t *testing.T) {
-	invalidUTF8 := append([]byte(`{"value":"`), 0xff)
-	invalidUTF8 = append(invalidUTF8, []byte(`"}`)...)
-	tests := []struct {
-		name    string
-		schema  []byte
-		wantErr string
-	}{
-		{name: "object", schema: []byte(`{"type":"object","properties":{"answer":{"type":"string"}}}`)},
-		{name: "maximum depth", schema: nestedGrammarSchema(maxGrammarSchemaDepth)},
-		{name: "too deep", schema: nestedGrammarSchema(maxGrammarSchemaDepth + 1), wantErr: "nesting exceeds"},
-		{name: "too many tokens", schema: grammarSchemaArray(maxGrammarSchemaTokens), wantErr: "more than 16384 JSON tokens"},
-		{name: "too large", schema: []byte(`{"value":"` + strings.Repeat("x", maxGrammarSchemaBytes) + `"}`), wantErr: "limit is 1048576"},
-		{name: "invalid UTF-8", schema: invalidUTF8, wantErr: "not valid UTF-8"},
-		{name: "trailing value", schema: []byte(`{} {}`), wantErr: "more than one JSON value"},
-		{name: "malformed", schema: []byte(`{"type":`), wantErr: "unexpected end"},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			err := validateGrammarSchema(tt.schema)
-			if tt.wantErr == "" {
-				if err != nil {
-					t.Fatal(err)
-				}
-				return
-			}
-			if err == nil || !strings.Contains(err.Error(), tt.wantErr) {
-				t.Fatalf("validateGrammarSchema error = %v, want containing %q", err, tt.wantErr)
-			}
-		})
-	}
-}
-
-func BenchmarkValidateGrammarSchema(b *testing.B) {
-	schema := []byte(`{"type":"object","properties":{"answer":{"type":"string","enum":["ok"]}},"required":["answer"],"additionalProperties":false}`)
+func BenchmarkParseGrammar(b *testing.B) {
+	grammar := json.RawMessage(schemaTag(`{"type":"object","properties":{"answer":{"type":"string","enum":["ok"]}},"required":["answer"],"additionalProperties":false}`))
 	b.ReportAllocs()
 	for b.Loop() {
-		if err := validateGrammarSchema(schema); err != nil {
+		if _, err := parseGrammar(grammar); err != nil {
 			b.Fatal(err)
 		}
 	}
@@ -137,20 +103,16 @@ func FuzzParseGrammar(f *testing.F) {
 		"",
 		`"json"`,
 		`{"type":"integer"}`,
+		schemaTag(`{"type":"integer"}`),
 		`{`,
 	} {
 		f.Add(format)
 	}
 
 	f.Fuzz(func(t *testing.T, format string) {
-		spec, err := parseGrammar(json.RawMessage(format))
-		if err != nil || spec == nil {
-			return
-		}
-		switch spec.kind {
-		case xgrammar.JSONSchema:
-		default:
-			t.Fatalf("parseGrammar kind = %d", spec.kind)
+		source, err := parseGrammar(json.RawMessage(format))
+		if err == nil && source != "" && source != format {
+			t.Fatalf("parseGrammar = %q, want the format verbatim", source)
 		}
 	})
 }
@@ -167,7 +129,6 @@ func TestValidateGrammarVocab(t *testing.T) {
 		{name: "input-only tokens past the head", logits: 31, tokenizer: 32},
 		{name: "invalid tokenizer", logits: 32, tokenizer: -1, wantErr: true},
 		{name: "invalid logits width", logits: 0, tokenizer: 32, wantErr: true},
-		{name: "allocation bound", logits: maxGrammarVocabSize + 1, tokenizer: 32, wantErr: true},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
 			err := validateGrammarVocab(tt.logits, tt.tokenizer)
@@ -194,7 +155,7 @@ func TestPrepareGrammarUnavailable(t *testing.T) {
 	}
 	request := &Request{CompletionRequest: CompletionRequest{
 		Prompt: "0",
-		Format: json.RawMessage(`"json"`),
+		Format: json.RawMessage(schemaTag(`{"type":"object"}`)),
 	}}
 	err := r.Prepare(request)
 	var statusErr api.StatusError
@@ -260,7 +221,7 @@ func testDigitGrammar(t *mlxtest.T, schema string) (*grammarEngine, *grammar) {
 		t.Fatal(err)
 	}
 	t.Cleanup(compiler.Close)
-	m, err := compiler.Compile(xgrammar.JSONSchema, schema)
+	m, err := compiler.Compile(schemaTag(schema))
 	if err != nil {
 		t.Fatal(err)
 	}

x/mlxrunner/xgrammar/dynamic.c

@@ -49,7 +49,7 @@ static const char* (*version_fn)(void);
 static const char* (*last_error_fn)(void);
 static int (*compiler_new_fn)(const char*, size_t, const uint64_t*, size_t, int32_t, const int32_t*, size_t, int32_t, int64_t, ollama_xgrammar_compiler**);
 static void (*compiler_free_fn)(ollama_xgrammar_compiler*);
-static int (*matcher_new_fn)(ollama_xgrammar_compiler*, ollama_xgrammar_kind, const char*, size_t, ollama_xgrammar_matcher**);
+static int (*matcher_new_fn)(ollama_xgrammar_compiler*, const char*, size_t, ollama_xgrammar_matcher**);
 static void (*matcher_free_fn)(ollama_xgrammar_matcher*);
 static int (*matcher_fill_fn)(ollama_xgrammar_matcher*, int32_t*, size_t, int*);
 static int (*matcher_accept_fn)(ollama_xgrammar_matcher*, int32_t, int*);
@@ -140,8 +140,8 @@ int ollama_xgrammar_dynamic_compiler_new(const char* d, size_t ds, const uint64_
     return capture_error(compiler_new_fn(d, ds, o, n, v, s, ns, mt, cb, m), error);
 }
 void ollama_xgrammar_dynamic_compiler_free(ollama_xgrammar_compiler* m) { compiler_free_fn(m); }
-int ollama_xgrammar_dynamic_matcher_new(ollama_xgrammar_compiler* m, ollama_xgrammar_kind k, const char* s, size_t n, ollama_xgrammar_matcher** out, char** error) {
-    return capture_error(matcher_new_fn(m, k, s, n, out), error);
+int ollama_xgrammar_dynamic_matcher_new(ollama_xgrammar_compiler* m, const char* s, size_t n, ollama_xgrammar_matcher** out, char** error) {
+    return capture_error(matcher_new_fn(m, s, n, out), error);
 }
 void ollama_xgrammar_dynamic_matcher_free(ollama_xgrammar_matcher* m) { matcher_free_fn(m); }
 int ollama_xgrammar_dynamic_matcher_fill(ollama_xgrammar_matcher* m, int32_t* b, size_t n, int* a, char** error) {

x/mlxrunner/xgrammar/dynamic.h

@@ -16,8 +16,7 @@ int ollama_xgrammar_dynamic_compiler_new(
     const int32_t*, size_t, int32_t, int64_t, ollama_xgrammar_compiler**, char**);
 void ollama_xgrammar_dynamic_compiler_free(ollama_xgrammar_compiler*);
 int ollama_xgrammar_dynamic_matcher_new(
-    ollama_xgrammar_compiler*, ollama_xgrammar_kind, const char*, size_t,
-    ollama_xgrammar_matcher**, char**);
+    ollama_xgrammar_compiler*, const char*, size_t, ollama_xgrammar_matcher**, char**);
 void ollama_xgrammar_dynamic_matcher_free(ollama_xgrammar_matcher*);
 int ollama_xgrammar_dynamic_matcher_fill(ollama_xgrammar_matcher*, int32_t*, size_t, int*, char**);
 int ollama_xgrammar_dynamic_matcher_accept(ollama_xgrammar_matcher*, int32_t, int*, char**);

x/mlxrunner/xgrammar/native/xgrammar.cpp

@@ -152,7 +152,6 @@ int ollama_xgrammar_compiler_new(
 
 int ollama_xgrammar_matcher_new(
     ollama_xgrammar_compiler* compiler,
-    ollama_xgrammar_kind kind,
     const char* source,
     size_t source_size,
     ollama_xgrammar_matcher** matcher) {
@@ -166,14 +165,8 @@ int ollama_xgrammar_matcher_new(
         }
         const char* source_data = source == nullptr ? "" : source;
 
-        xgrammar::CompiledGrammar compiled = [&]() -> xgrammar::CompiledGrammar {
-            switch (kind) {
-            case OLLAMA_XGRAMMAR_JSON_SCHEMA:
-                return compiler->compiler.CompileJSONSchema(std::string(source_data, source_size));
-            default:
-                throw std::invalid_argument("unknown grammar kind");
-            }
-        }();
+        xgrammar::CompiledGrammar compiled =
+            compiler->compiler.CompileStructuralTag(std::string(source_data, source_size));
         *matcher = new ollama_xgrammar_matcher(compiler->vocab_size, std::move(compiled));
     });
 }

x/mlxrunner/xgrammar/native/xgrammar.h

@@ -21,10 +21,6 @@ extern "C" {
 typedef struct ollama_xgrammar_compiler ollama_xgrammar_compiler;
 typedef struct ollama_xgrammar_matcher ollama_xgrammar_matcher;
 
-typedef enum ollama_xgrammar_kind {
-    OLLAMA_XGRAMMAR_JSON_SCHEMA = 0,
-} ollama_xgrammar_kind;
-
 // The pinned xgrammar release the library was built from.
 OLLAMA_XGRAMMAR_API const char* ollama_xgrammar_version(void);
 
@@ -43,9 +39,9 @@ OLLAMA_XGRAMMAR_API int ollama_xgrammar_compiler_new(
     ollama_xgrammar_compiler** compiler);
 OLLAMA_XGRAMMAR_API void ollama_xgrammar_compiler_free(ollama_xgrammar_compiler* compiler);
 
+// Compiles a structural tag into a matcher.
 OLLAMA_XGRAMMAR_API int ollama_xgrammar_matcher_new(
     ollama_xgrammar_compiler* compiler,
-    ollama_xgrammar_kind kind,
     const char* source,
     size_t source_size,
     ollama_xgrammar_matcher** matcher);

x/mlxrunner/xgrammar/xgrammar.go

@@ -15,12 +15,6 @@ import (
 	"unsafe"
 )
 
-type Kind int
-
-const (
-	JSONSchema Kind = C.OLLAMA_XGRAMMAR_JSON_SCHEMA
-)
-
 // The native library is loaded once per process and never unloaded.
 var (
 	loadOnce sync.Once
@@ -138,7 +132,8 @@ func (c *Compiler) Version() string {
 	return C.GoString(C.ollama_xgrammar_dynamic_version())
 }
 
-func (c *Compiler) Compile(kind Kind, source string) (*Matcher, error) {
+// Compile compiles a structural tag, given as its JSON text.
+func (c *Compiler) Compile(source string) (*Matcher, error) {
 	if c == nil {
 		return nil, errors.New("grammar compiler is unavailable")
 	}
@@ -153,7 +148,7 @@ func (c *Compiler) Compile(kind Kind, source string) (*Matcher, error) {
 	var ctx *C.ollama_xgrammar_matcher
 	var cError *C.char
 	if C.ollama_xgrammar_dynamic_matcher_new(
-		c.ctx, C.ollama_xgrammar_kind(kind), sourcePtr, C.size_t(len(source)), &ctx, &cError,
+		c.ctx, sourcePtr, C.size_t(len(source)), &ctx, &cError,
 	) != 0 {
 		return nil, nativeError("compile grammar", cError)
 	}

x/mlxrunner/xgrammar/xgrammar_test.go

@@ -39,6 +39,12 @@ func testLibraryDir(t testing.TB) string {
 	return filepath.Dir(path)
 }
 
+// schemaTag wraps a JSON Schema into the structural tag the MLX client
+// sends for a schema format.
+func schemaTag(schema string) string {
+	return `{"type":"structural_tag","format":{"type":"json_schema","json_schema":` + schema + `}}`
+}
+
 func testGrammarCompiler(t testing.TB) *xgrammar.Compiler {
 	t.Helper()
 	compiler, err := xgrammar.New(testLibraryDir(t), testVocabulary(), testVocabSize, []int32{testEOS}, 8, 128<<20)
@@ -72,7 +78,7 @@ func FuzzJSONSchemaMatcher(f *testing.F) {
 		if len(schema) > 1<<20 || len(tokens) > 256 {
 			return
 		}
-		matcher, err := compiler.Compile(xgrammar.JSONSchema, schema)
+		matcher, err := compiler.Compile(schemaTag(schema))
 		if err != nil {
 			return
 		}
@@ -127,7 +133,7 @@ func acceptPieces(t *testing.T, matcher *xgrammar.Matcher, pieces ...string) {
 func TestJSONSchemaMatcher(t *testing.T) {
 	compiler := testGrammarCompiler(t)
 	schema := `{"type":"object","properties":{"answer":{"type":"string","enum":["ok"]}},"required":["answer"],"additionalProperties":false}`
-	matcher, err := compiler.Compile(xgrammar.JSONSchema, schema)
+	matcher, err := compiler.Compile(schemaTag(schema))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -156,7 +162,7 @@ func TestJSONSchemaMatcher(t *testing.T) {
 // a decoder can sample past the grammar's end.
 func TestFillAfterTermination(t *testing.T) {
 	compiler := testGrammarCompiler(t)
-	matcher, err := compiler.Compile(xgrammar.JSONSchema, "{}")
+	matcher, err := compiler.Compile(schemaTag("{}"))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -173,7 +179,7 @@ func TestFillAfterTermination(t *testing.T) {
 func TestRollbackRestoresState(t *testing.T) {
 	compiler := testGrammarCompiler(t)
 	schema := `{"type":"object","properties":{"answer":{"type":"string","enum":["ok"]}},"required":["answer"],"additionalProperties":false}`
-	matcher, err := compiler.Compile(xgrammar.JSONSchema, schema)
+	matcher, err := compiler.Compile(schemaTag(schema))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -225,9 +231,11 @@ func TestInvalidGrammarReturnsError(t *testing.T) {
 	}{
 		{name: "empty", source: ""},
 		{name: "malformed", source: `{"type":`},
+		{name: "bare schema", source: `{"type":"object"}`},
+		{name: "empty enum", source: schemaTag(`{"enum":[]}`)},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
-			matcher, err := compiler.Compile(xgrammar.JSONSchema, tt.source)
+			matcher, err := compiler.Compile(tt.source)
 			if matcher != nil {
 				matcher.Close()
 			}
@@ -265,7 +273,7 @@ func TestCompilerValidationAndClosedState(t *testing.T) {
 		})
 	}
 
-	matcher, err := loaded.Compile(xgrammar.JSONSchema, "{}")
+	matcher, err := loaded.Compile(schemaTag("{}"))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -279,7 +287,45 @@ func TestCompilerValidationAndClosedState(t *testing.T) {
 	}
 	loaded.Close()
 	loaded.Close()
-	if matcher, err := loaded.Compile(xgrammar.JSONSchema, "{}"); err == nil || matcher != nil {
+	if matcher, err := loaded.Compile(schemaTag("{}")); err == nil || matcher != nil {
 		t.Fatalf("Compile on closed compiler = %v, %v; want error", matcher, err)
 	}
 }
+
+// A structural-tag grammar leaves text free until a trigger opens a tag,
+// constrains the tag body to its schema (EOS masked until the tag closes),
+// and frees the text again after the end tag.
+func TestStructuralTagMatcher(t *testing.T) {
+	compiler := testGrammarCompiler(t)
+	tag := `{"type":"structural_tag","format":{"type":"triggered_tags","triggers":["["],` +
+		`"tags":[{"begin":"[","content":{"type":"json_schema","json_schema":` +
+		`{"type":"object","properties":{"answer":{"type":"string","enum":["ok"]}},"required":["answer"],"additionalProperties":false}},"end":"]"}]}}`
+	matcher, err := compiler.Compile(tag)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer matcher.Close()
+
+	if mask, constrained := fillMask(t, matcher); constrained && !allowed(mask, testEOS) {
+		t.Fatal("EOS is masked in free text before any tag")
+	}
+	acceptPieces(t, matcher, "answer", " ")
+
+	acceptPieces(t, matcher, "[")
+	mask, constrained := fillMask(t, matcher)
+	if !constrained {
+		t.Fatal("an open tag does not constrain sampling")
+	}
+	if allowed(mask, testEOS) {
+		t.Fatal("EOS is allowed inside an open tag")
+	}
+	acceptPieces(t, matcher, "{", `"`, "answer", `"`, ":", `"`, "ok", `"`, "}", "]")
+
+	if mask, constrained := fillMask(t, matcher); constrained && !allowed(mask, testEOS) {
+		t.Fatal("EOS is masked in free text after the tag closes")
+	}
+	acceptPieces(t, matcher, "ok", "<eos>")
+	if !matcher.Terminated() {
+		t.Fatal("matcher not terminated after EOS")
+	}
+}