feat(system-agent): constrain planner JSON output at generation time (#113482)

* feat(system-agent): constrain planner JSON output

* fix(ai): align response format request typing

* fix(ai): preserve response format backend contracts

extensions/llama-cpp/src/inference-provider.test.ts

@@ -24,6 +24,8 @@ const mocks = vi.hoisted(() => {
   };
   const llama = {
     loadModel: vi.fn(async () => model),
+    createGrammarForJsonSchema: vi.fn(async (schema: unknown) => ({ schema })),
+    getGrammarFor: vi.fn(async (type: string) => ({ type })),
     dispose: llamaDispose,
   };
   return {
@@ -235,6 +237,96 @@ describe("llama.cpp inference provider", () => {
     });
   });
 
+  it("builds a JSON Schema grammar for tool-free responseFormat requests", async () => {
+    const schema = {
+      type: "object",
+      properties: { reply: { type: "string" } },
+      required: ["reply"],
+      additionalProperties: false,
+    };
+    const stream = await createLlamaCppStreamFn({})(
+      model,
+      { messages: [{ role: "user", content: "Hi", timestamp: 1 }] },
+      { responseFormat: schema },
+    );
+
+    await collectEvents(stream);
+
+    expect(mocks.llama.createGrammarForJsonSchema).toHaveBeenCalledWith(schema);
+    expect(mocks.generateResponse.mock.calls[0]?.[1]).toMatchObject({
+      grammar: { schema },
+    });
+    expect(mocks.generateResponse.mock.calls[0]?.[1]).not.toHaveProperty("functions");
+  });
+
+  it("unwraps provider-shaped json_schema response formats", async () => {
+    const schema = {
+      type: "object",
+      properties: { reply: { type: "string" } },
+      required: ["reply"],
+      additionalProperties: false,
+    };
+    const stream = await createLlamaCppStreamFn({})(
+      model,
+      { messages: [{ role: "user", content: "Hi", timestamp: 1 }] },
+      {
+        responseFormat: {
+          type: "json_schema",
+          json_schema: { name: "planner", schema },
+        },
+      },
+    );
+
+    await collectEvents(stream);
+
+    expect(mocks.llama.createGrammarForJsonSchema).toHaveBeenCalledWith(schema);
+    expect(mocks.generateResponse.mock.calls[0]?.[1]).toMatchObject({ grammar: { schema } });
+  });
+
+  it("maps provider-shaped json_object response formats to the JSON grammar", async () => {
+    const stream = await createLlamaCppStreamFn({})(
+      model,
+      { messages: [{ role: "user", content: "Hi", timestamp: 1 }] },
+      { responseFormat: { type: "json_object" } },
+    );
+
+    await collectEvents(stream);
+
+    expect(mocks.llama.getGrammarFor).toHaveBeenCalledWith("json");
+    expect(mocks.generateResponse.mock.calls[0]?.[1]).toMatchObject({
+      grammar: { type: "json" },
+    });
+  });
+
+  it("maps an empty JSON Schema to the generic JSON grammar", async () => {
+    const stream = await createLlamaCppStreamFn({})(
+      model,
+      { messages: [{ role: "user", content: "Hi", timestamp: 1 }] },
+      { responseFormat: {} },
+    );
+
+    await collectEvents(stream);
+
+    expect(mocks.llama.getGrammarFor).toHaveBeenCalledWith("json");
+    expect(mocks.generateResponse.mock.calls[0]?.[1]).toMatchObject({
+      grammar: { type: "json" },
+    });
+  });
+
+  it("keeps provider-shaped text response formats unconstrained", async () => {
+    const stream = await createLlamaCppStreamFn({})(
+      model,
+      { messages: [{ role: "user", content: "Hi", timestamp: 1 }] },
+      { responseFormat: { type: "text" } },
+    );
+
+    await collectEvents(stream);
+
+    expect(mocks.llama.getGrammarFor).not.toHaveBeenCalled();
+    expect(mocks.llama.createGrammarForJsonSchema).not.toHaveBeenCalled();
+    expect(mocks.generateResponse.mock.calls[0]?.[1]).not.toHaveProperty("grammar");
+  });
+
   it("emits native function calls in the final assistant message", async () => {
     mocks.generateResponse.mockResolvedValueOnce({
       response: "",
@@ -269,6 +361,40 @@ describe("llama.cpp inference provider", () => {
         ],
       },
     });
+    expect(mocks.llama.createGrammarForJsonSchema).not.toHaveBeenCalled();
+  });
+
+  it("lets tools win when responseFormat is also present", async () => {
+    const stream = await createLlamaCppStreamFn({})(
+      model,
+      {
+        messages: [{ role: "user", content: "Weather?", timestamp: 1 }],
+        tools: [
+          {
+            name: "weather",
+            description: "Get weather",
+            parameters: { type: "object", properties: { city: { type: "string" } } },
+          },
+        ],
+      },
+      {
+        responseFormat: {
+          type: "object",
+          properties: { reply: { type: "string" } },
+          required: ["reply"],
+          additionalProperties: false,
+        },
+      },
+    );
+
+    await collectEvents(stream);
+
+    expect(mocks.llama.createGrammarForJsonSchema).not.toHaveBeenCalled();
+    expect(mocks.generateResponse.mock.calls[0]?.[1]).toMatchObject({
+      functions: { weather: expect.any(Object) },
+      documentFunctionParams: true,
+    });
+    expect(mocks.generateResponse.mock.calls[0]?.[1]).not.toHaveProperty("grammar");
   });
 
   it("disposes the previous model and context when the model changes", async () => {

extensions/llama-cpp/src/inference-provider.ts

@@ -30,11 +30,14 @@ import {
 
 type LoadedModel = {
   key: string;
+  llama: Llama;
   model: LlamaModel;
   context: LlamaContext;
   sequence: LlamaContextSequence;
 };
 
+type LlamaJsonSchemaInput = Parameters<Llama["createGrammarForJsonSchema"]>[0];
+
 // Process-owned, single-slot cache. A model/context pair lives until another
 // model replaces it or the process exits, bounding resident model memory.
 let loadedModel: LoadedModel | undefined;
@@ -94,6 +97,33 @@ function normalizeArguments(value: unknown): Record<string, unknown> {
     : {};
 }
 
+async function resolveLlamaCppResponseGrammar(params: {
+  llama: Llama;
+  responseFormat: Record<string, unknown> | undefined;
+}) {
+  const responseFormat = params.responseFormat;
+  if (!responseFormat) {
+    return undefined;
+  }
+  if (Object.keys(responseFormat).length === 0) {
+    return await params.llama.getGrammarFor("json");
+  }
+  if (responseFormat.type === "json_object") {
+    return await params.llama.getGrammarFor("json");
+  }
+  if (responseFormat.type === "text") {
+    return undefined;
+  }
+  if (responseFormat.type === "json_schema") {
+    const envelope = normalizeArguments(responseFormat.json_schema);
+    const schema = normalizeArguments(envelope.schema);
+    return Object.keys(schema).length > 0
+      ? await params.llama.createGrammarForJsonSchema(schema as LlamaJsonSchemaInput)
+      : await params.llama.getGrammarFor("json");
+  }
+  return await params.llama.createGrammarForJsonSchema(responseFormat as LlamaJsonSchemaInput);
+}
+
 function mapContextToLlamaChatHistory(context: Context): ChatHistoryItem[] {
   const history: ChatHistoryItem[] = [];
   if (context.systemPrompt?.trim()) {
@@ -237,7 +267,7 @@ async function getLoadedModel(params: {
     // Serialized requests reuse this one sequence. Disposing/reallocating it per
     // turn races node-llama-cpp's asynchronous sequence-id reclamation.
     const sequence = context.getSequence();
-    loadedModel = { key, model, context, sequence };
+    loadedModel = { key, llama, model, context, sequence };
     return loadedModel;
   } catch (error) {
     await context?.dispose();
@@ -312,6 +342,7 @@ export function createLlamaCppStreamFn(params: { providerConfig?: ModelProviderC
           autoDisposeSequence: false,
         });
         const before = sequence.tokenMeter.getState();
+        const functions = mapToolsToLlamaFunctions(context);
         let textStarted = false;
         const partial = () =>
           buildMessage({
@@ -332,15 +363,31 @@ export function createLlamaCppStreamFn(params: { providerConfig?: ModelProviderC
           stream.push({ type: "text_delta", contentIndex: 0, delta });
         };
         try {
-          const result = await chat.generateResponse(mapContextToLlamaChatHistory(context), {
-            functions: mapToolsToLlamaFunctions(context),
-            documentFunctionParams: true,
+          // node-llama-cpp makes grammar and functions mutually exclusive. Tool
+          // turns keep function calling; constrained decoding is for tool-free turns.
+          const grammar =
+            functions || !options?.responseFormat
+              ? undefined
+              : await resolveLlamaCppResponseGrammar({
+                  llama: loaded.llama,
+                  responseFormat: options.responseFormat,
+                });
+          const generationOptions = {
             signal: options?.signal,
             maxTokens: options?.maxTokens ?? model.maxTokens,
             temperature: options?.temperature,
             customStopTriggers: options?.stop,
             onTextChunk: appendTextDelta,
-          });
+            ...(functions
+              ? { functions, documentFunctionParams: true as const }
+              : grammar
+                ? { grammar }
+                : {}),
+          };
+          const result = await chat.generateResponse(
+            mapContextToLlamaChatHistory(context),
+            generationOptions,
+          );
           if (result.metadata.stopReason === "abort" || signal?.aborted) {
             generationAborted = true;
             throw signal?.reason ?? new Error("Request was aborted");