fix: bound Ollama prompts to prevent silent context truncation

apps/webapp/app/jobs/labels/label-assignment.logic.ts

@@ -6,7 +6,11 @@
  */
 
 import { z } from "zod";
-import { makeStructuredModelCall, getEmbedding } from "~/lib/model.server";
+import {
+  makeStructuredModelCall,
+  getEmbedding,
+  resolveProfileForCall,
+} from "~/lib/model.server";
 import { logger } from "~/services/logger.service";
 import { prisma } from "~/db.server";
 import { LabelService } from "~/services/label.server";
@@ -15,9 +19,55 @@ import { generateOklchColor } from "~/components/ui/color-utils";
 import { type ModelMessage } from "ai";
 import { ProviderFactory, VECTOR_NAMESPACES } from "@core/providers";
 import { countTokens } from "~/services/search/tokenBudget";
+import {
+  OLLAMA_NUM_CTX,
+  assertPromptWithinBudget,
+  capToTokenBudget,
+  capToTokenBudgetFromEnd,
+} from "~/services/prompts/promptBudget";
+import {
+  type PromptProfile,
+} from "~/services/prompts/normalizeProfile";
 
 const MAX_CONTENT_TOKENS = 20000;
 
+/**
+ * Label extraction budget for small-context providers.
+ *
+ * `sessionContext` here is the ENTIRE document body (see processLabelAssignment,
+ * where it is set from `document.content`). The only previous guard was
+ * MAX_CONTENT_TOKENS at 20000, which is roughly five times a whole Ollama
+ * context window — so against Ollama this prompt was silently truncated rather
+ * than bounded.
+ *
+ * Measured with the o200k_base tokenizer, the static system prompt is 611
+ * tokens and the user-prompt scaffolding is 21. Reserving 512 tokens of output
+ * leaves 4096 - 512 = 3584 for input. Allowing 400 tokens for the existing
+ * labels list:
+ *   611 static + 21 glue + 400 labels + 2400 content = 3432, inside the 3500
+ *   assertion budget, and 3500 + 512 output = 4012 < 4096.
+ * Pinned by test rather than assumed.
+ */
+export const LABEL_OUTPUT_TOKEN_RESERVE = 512;
+export const LABEL_PROMPT_TOKEN_BUDGET = Math.min(
+  3500,
+  OLLAMA_NUM_CTX - LABEL_OUTPUT_TOKEN_RESERVE,
+);
+export const LABEL_CONTENT_TOKEN_BUDGET_OLLAMA = 2400;
+
+/**
+ * Budget for the existing-labels list.
+ *
+ * Previously this was an unenforced claim in a docstring while the list itself
+ * rendered every workspace label. Since this job creates labels, the list only
+ * grows, so the unenforced version fails progressively: measured on Ollama it
+ * threw at ~29 labels with session context (134 without). The hosted budget is
+ * generous because hosted windows are not the constraint; it exists so the list
+ * cannot grow without any bound at all.
+ */
+export const LABEL_LIST_TOKEN_BUDGET_OLLAMA = 400;
+export const LABEL_LIST_TOKEN_BUDGET_HOSTED = 4000;
+
 // Similarity threshold for matching labels (higher = stricter matching)
 const LABEL_SIMILARITY_THRESHOLD = 0.85;
 
@@ -301,7 +351,17 @@ export async function extractLabelsFromEpisode(
   workspaceId: string,
   sessionContext?: string,
 ): Promise<ExtractedLabel[]> {
-  const messages = buildLabelExtractionMessages(episodeBody, availableLabels, sessionContext);
+  // Resolved from the model this call actually resolves to, not the global env
+  // var: a workspace override can land on Ollama while CHAT_PROVIDER says
+  // otherwise, which would silently select the uncapped hosted profile.
+  const profile = await resolveProfileForCall(workspaceId, "memory", "medium");
+
+  const messages = buildLabelExtractionMessages(
+    episodeBody,
+    availableLabels,
+    sessionContext,
+    profile,
+  );
 
   logger.info("Extracting labels from episode", {
     episodeTokens: countTokens(episodeBody),
@@ -317,6 +377,8 @@ export async function extractLabelsFromEpisode(
     0.3, // Low temperature for consistent label extraction
     workspaceId,
     "memory",
+    undefined,
+    profile === "ollama" ? LABEL_OUTPUT_TOKEN_RESERVE : undefined,
   );
 
   // Create lookup map for existing labels (case-insensitive) for exact matching
@@ -426,43 +488,87 @@ export function buildLabelExtractionMessages(
     description: string | null;
   }>,
   sessionContext?: string,
+  profile: PromptProfile = "hosted",
 ): ModelMessage[] {
+  // Ollama silently drops anything past num_ctx, so the generous hosted budget
+  // has to shrink to something that actually fits the window.
+  const contentBudget =
+    profile === "ollama" ? LABEL_CONTENT_TOKEN_BUDGET_OLLAMA : MAX_CONTENT_TOKENS;
+  const labelListBudget =
+    profile === "ollama" ? LABEL_LIST_TOKEN_BUDGET_OLLAMA : LABEL_LIST_TOKEN_BUDGET_HOSTED;
+
   // Token-aware truncation: prioritise current episode, fill remainder with session context
   const episodeTokens = countTokens(episodeBody);
   let truncatedEpisode = episodeBody;
   let truncatedContext: string | undefined;
 
-  if (episodeTokens > MAX_CONTENT_TOKENS) {
-    // Edge case: episode alone exceeds budget — hard-trim from the end
-    const chars = Math.floor((MAX_CONTENT_TOKENS / episodeTokens) * episodeBody.length);
-    truncatedEpisode = episodeBody.substring(0, chars) + "...[truncated]";
+  if (episodeTokens > contentBudget) {
+    // Edge case: episode alone exceeds budget. This used to scale by character
+    // ratio, which only estimates the resulting token count and can overshoot;
+    // capToTokenBudget converges on the real count, which matters now that
+    // going over the budget throws.
+    truncatedEpisode = capToTokenBudget(episodeBody, contentBudget);
   }
 
   if (sessionContext) {
-    const remaining = MAX_CONTENT_TOKENS - countTokens(truncatedEpisode);
+    const remaining = contentBudget - countTokens(truncatedEpisode);
     if (remaining > 200) {
       const contextTokens = countTokens(sessionContext);
       if (contextTokens <= remaining) {
         truncatedContext = sessionContext;
       } else {
-        // Keep the most recent part of the session (tail), drop oldest
-        const ratio = remaining / contextTokens;
-        const startChar = Math.floor((1 - ratio) * sessionContext.length);
-        truncatedContext =
-          "...[earlier context omitted]\n" + sessionContext.substring(startChar);
+        // Keep the most recent part of the session (tail), drop oldest.
+        // Token-exact rather than scaled by character ratio: the ratio only
+        // estimates the result, and a document with a sparse-ASCII head and a
+        // dense CJK/emoji tail overshoots enough to blow the budget and throw.
+        truncatedContext = capToTokenBudgetFromEnd(sessionContext, remaining);
       }
     }
   }
 
+  // The label list is the one injection that grows on its own: this job creates
+  // labels, so every run can enlarge the next run's prompt. Left unbounded it
+  // does not degrade, it hits the assertion and the job starts failing outright
+  // once a workspace accumulates enough labels. Measured breaking point on
+  // Ollama was ~29 labels with session context.
+  //
+  // KNOWN LIMITATION — the drop order is alphabetical, not by relevance or
+  // recency. LabelService.getWorkspaceLabels orders by `name: "asc"`, so this
+  // truncates from the end of the alphabet: a workspace with 300 labels shows
+  // the model roughly "000".."020" and permanently hides everything later in
+  // the alphabet. Because hidden labels cannot be matched, the model proposes
+  // new ones instead, so those workspaces will accumulate near-duplicate labels
+  // over time.
+  //
+  // This bounds the crash, it does not solve label selection. Doing that
+  // properly means selecting candidates by embedding similarity to the episode
+  // (the machinery already exists here for dedup) rather than taking a
+  // prefix of an alphabetical list. Filed as follow-up rather than folded in,
+  // because it changes which labels the model can see and deserves its own
+  // evaluation.
+  const renderLabel = (l: { name: string; description: string | null }) =>
+    `  <label name="${l.name}"${l.description ? ` description="${l.description}"` : ""} />`;
+
+  const fittedLabels = [...availableLabels];
+  while (
+    fittedLabels.length > 0 &&
+    countTokens(fittedLabels.map(renderLabel).join("\n")) > labelListBudget
+  ) {
+    fittedLabels.pop();
+  }
+
+  if (fittedLabels.length < availableLabels.length) {
+    logger.warn("Label list truncated to fit the prompt budget", {
+      total: availableLabels.length,
+      kept: fittedLabels.length,
+      budget: labelListBudget,
+    });
+  }
+
   const existingLabelsXml =
-    availableLabels.length > 0
+    fittedLabels.length > 0
       ? `<existing_labels>
-${availableLabels
-  .map(
-    (l) =>
-      `  <label name="${l.name}"${l.description ? ` description="${l.description}"` : ""} />`,
-  )
-  .join("\n")}
+${fittedLabels.map(renderLabel).join("\n")}
 </existing_labels>`
       : "<existing_labels />";
 
@@ -476,7 +582,7 @@ ${truncatedContext}
 ${truncatedEpisode}
 </current_episode>`;
 
-  return [
+  const messages: ModelMessage[] = [
     {
       role: "system",
       content: `You extract LABELS from episodes for a USER'S PERSONAL KNOWLEDGE SYSTEM.
@@ -547,4 +653,15 @@ ${sessionContextXml}
 ${currentEpisodeXml}`,
     },
   ];
+
+  if (profile === "ollama") {
+    // Fail loudly rather than let Ollama silently drop the overflow.
+    assertPromptWithinBudget({
+      label: "label extraction prompt (ollama profile)",
+      text: messages.map((m) => m.content as string).join("\n"),
+      budget: LABEL_PROMPT_TOKEN_BUDGET,
+    });
+  }
+
+  return messages;
 }

apps/webapp/app/lib/__tests__/model-boundary-invariant.test.ts

@@ -0,0 +1,154 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { z } from "zod";
+
+/**
+ * Integration coverage for the boundary invariant's actual wiring, not just
+ * the standalone checkPromptBoundary function (already covered in
+ * promptBudget.test.ts). A cross-vendor review found the earlier version of
+ * this diff had zero coverage proving the check fires from inside
+ * makeModelCall / structuredCallWithTolerantParsing themselves — deleting the
+ * wiring left every test green. These tests mock only the network-facing
+ * edge (Mastra's Agent class) so the real makeModelCall /
+ * makeStructuredModelCall / structuredCallWithTolerantParsing logic runs,
+ * including checkPromptBoundary.
+ *
+ * model.server pulls in prisma/env through llm-provider.server and
+ * tokenUsage.server, so both are stubbed — same approach as
+ * ollama-model-detection.test.ts and rerank-budget.test.ts.
+ */
+
+const getDefaultChatProviderType = vi.fn<() => string>();
+const resolveModelForWorkspace = vi.fn();
+const getProviderConfig = vi.fn((provider: string) =>
+  provider === "ollama" ? { baseUrl: "http://localhost:11434" } : {},
+);
+
+vi.mock("~/services/llm-provider.server", () => ({
+  getDefaultChatProviderType: () => getDefaultChatProviderType(),
+  resolveModelForWorkspace: (...args: unknown[]) => resolveModelForWorkspace(...args),
+  getDefaultChatModelId: vi.fn(() => "qwen3:8b"),
+  getDefaultEmbeddingInfo: vi.fn(),
+  getProviderConfig: (...args: [string]) => getProviderConfig(...args),
+  getEmbeddingDimensions: vi.fn(),
+  resolveApiKey: vi.fn(),
+  resolveApiKeyForWorkspace: vi.fn(),
+}));
+
+vi.mock("~/services/tokenUsage.server", () => ({ recordTokenUsage: vi.fn() }));
+vi.mock("~/services/localEmbeddings.server", () => ({ embedLocal: vi.fn() }));
+
+const generateMock = vi.fn();
+const streamMock = vi.fn();
+
+vi.mock("@mastra/core/agent", () => ({
+  Agent: vi.fn().mockImplementation(() => ({
+    generate: (...args: unknown[]) => generateMock(...args),
+    stream: (...args: unknown[]) => streamMock(...args),
+  })),
+}));
+
+vi.mock("~/services/logger.service", () => ({
+  logger: { info: vi.fn(), warn: vi.fn(), error: vi.fn(), log: vi.fn(), debug: vi.fn() },
+}));
+
+import { logger } from "~/services/logger.service";
+import { makeModelCall, makeStructuredModelCall } from "../model.server";
+
+const OVER_BUDGET_TEXT = "a ".repeat(4000);
+
+beforeEach(() => {
+  vi.clearAllMocks();
+  delete process.env.OLLAMA_PROMPT_BOUNDARY_MODE; // warn-only for all these tests
+  getDefaultChatProviderType.mockReturnValue("openai");
+  generateMock.mockResolvedValue({ text: "ok", usage: undefined });
+  streamMock.mockResolvedValue({
+    text: Promise.resolve("ok"),
+    usage: Promise.resolve(undefined),
+  });
+});
+
+function errorMessages(): string[] {
+  return (logger.error as unknown as { mock: { calls: unknown[][] } }).mock.calls.map(
+    (call) => call[0] as string,
+  );
+}
+
+describe("makeModelCall boundary wiring", () => {
+  it("fires the boundary check for an over-budget prompt resolved to Ollama", async () => {
+    resolveModelForWorkspace.mockResolvedValue({ modelId: "ollama/qwen3:8b" });
+
+    await makeModelCall(false, [{ role: "user", content: OVER_BUDGET_TEXT }], () => {});
+
+    expect(errorMessages().some((m) => m.includes("PromptBudget:boundary"))).toBe(true);
+  });
+
+  it("stays silent for the same oversized prompt when the model resolves to a hosted provider", async () => {
+    resolveModelForWorkspace.mockResolvedValue({ modelId: "openai/gpt-5" });
+
+    await makeModelCall(false, [{ role: "user", content: OVER_BUDGET_TEXT }], () => {});
+
+    expect(errorMessages()).toEqual([]);
+  });
+
+  it("does not throw in default (warn-only) mode even when over budget", async () => {
+    resolveModelForWorkspace.mockResolvedValue({ modelId: "ollama/qwen3:8b" });
+
+    await expect(
+      makeModelCall(false, [{ role: "user", content: OVER_BUDGET_TEXT }], () => {}),
+    ).resolves.toBeDefined();
+  });
+
+  it("throws once OLLAMA_PROMPT_BOUNDARY_MODE=throw is set", async () => {
+    process.env.OLLAMA_PROMPT_BOUNDARY_MODE = "throw";
+    resolveModelForWorkspace.mockResolvedValue({ modelId: "ollama/qwen3:8b" });
+
+    await expect(
+      makeModelCall(false, [{ role: "user", content: OVER_BUDGET_TEXT }], () => {}),
+    ).rejects.toThrow(/PromptBudget:boundary/);
+  });
+});
+
+describe("makeStructuredModelCall boundary wiring (tolerant-parsing path)", () => {
+  const schema = z.object({ ok: z.boolean() });
+
+  it("fires the boundary check, counting the JSON preamble alongside the message content", async () => {
+    resolveModelForWorkspace.mockResolvedValue({ modelId: "ollama/qwen3:8b" });
+    generateMock.mockResolvedValue({ text: JSON.stringify({ ok: true }), usage: undefined });
+
+    // Small message content alone would fit; only content + jsonPreamble + schema
+    // together breach the budget, so this proves the preamble is really counted.
+    await makeStructuredModelCall(schema, [{ role: "user", content: OVER_BUDGET_TEXT }]);
+
+    expect(
+      errorMessages().some(
+        (m) => m.includes("PromptBudget:boundary") && m.includes("makeStructuredModelCall:chat/medium"),
+      ),
+    ).toBe(true);
+  });
+
+  it("stays silent when the resolved model is hosted", async () => {
+    resolveModelForWorkspace.mockResolvedValue({ modelId: "openai/gpt-5" });
+    generateMock.mockResolvedValue({
+      object: { ok: true },
+      usage: undefined,
+    });
+
+    await makeStructuredModelCall(schema, [{ role: "user", content: OVER_BUDGET_TEXT }]);
+
+    expect(errorMessages()).toEqual([]);
+  });
+
+  it("also fires the boundary check on the repair-retry path, under a distinct label", async () => {
+    resolveModelForWorkspace.mockResolvedValue({ modelId: "ollama/qwen3:8b" });
+    // First call: invalid JSON, forces the repair path. Second call (repair):
+    // succeeds. The repair call's own input (this huge first-pass output) is
+    // what should breach the repair check's budget.
+    generateMock
+      .mockResolvedValueOnce({ text: OVER_BUDGET_TEXT, usage: undefined })
+      .mockResolvedValueOnce({ text: JSON.stringify({ ok: true }), usage: undefined });
+
+    await makeStructuredModelCall(schema, [{ role: "user", content: "short" }]);
+
+    expect(errorMessages().some((m) => m.includes(":repair"))).toBe(true);
+  });
+});