fix(inference): harden Ollama restart recovery

src/lib/actions/sandbox/agent/ollama-restart-recovery.test.ts

@@ -3,7 +3,6 @@
 
 import { describe, expect, it, vi } from "vitest";
 import { OLLAMA_PORT, OLLAMA_PROXY_PORT } from "../../../core/ports";
-import { prepareOllamaApiExecution } from "../../../inference/local";
 import {
   maybeWarmOllamaAfterDaemonRestart,
   type OllamaRestartRecoveryDeps,
@@ -76,11 +75,6 @@ describe("maybeWarmOllamaAfterDaemonRestart", () => {
           runCaptureImpl,
           runCaptureExImpl,
           prepareDockerEnvironment,
-          prepareOllamaApiExecution: (command, host, options) =>
-            prepareOllamaApiExecution(command, host, {
-              ...options,
-              prepareDockerEnvironment,
-            }),
         },
       ),
     ).toEqual({ kind: "warmed", ok: true, timedOut: false });
@@ -199,7 +193,11 @@ describe("maybeWarmOllamaAfterDaemonRestart", () => {
         { provider: "ollama-local", model: "qwen3.6:35b" },
         { runCaptureImpl: () => "", runCaptureExImpl },
       ),
-    ).toEqual({ kind: "skipped", reason: "unreachable" });
+    ).toEqual({
+      kind: "skipped",
+      reason: "unreachable",
+      endpoint: `http://127.0.0.1:${OLLAMA_PORT}`,
+    });
     expect(runCaptureExImpl).not.toHaveBeenCalled();
   });
 
@@ -226,6 +224,52 @@ describe("maybeWarmOllamaAfterDaemonRestart", () => {
     });
   });
 
+  it("limits warm-up to the command timeout budget remaining after the probe", () => {
+    const runCaptureExImpl = vi.fn(
+      (_command: string[], _options?: { env?: NodeJS.ProcessEnv; timeout?: number }) =>
+        successfulWarmResult(),
+    );
+    const now = vi.fn().mockReturnValueOnce(1_000).mockReturnValueOnce(6_000);
+
+    expect(
+      maybeWarmOllamaAfterDaemonRestart(
+        { provider: "ollama-local", model: "qwen3.6:35b" },
+        {
+          probeRuntimeModelStatus: () => unloadedStatus,
+          runCaptureExImpl,
+          timeoutSeconds: 30,
+          now,
+        },
+      ),
+    ).toEqual({ kind: "warmed", ok: true, timedOut: false });
+
+    const warmCommand = runCaptureExImpl.mock.calls[0][0];
+    expect(warmCommand[warmCommand.indexOf("--max-time") + 1]).toBe("25");
+    expect(runCaptureExImpl.mock.calls[0][1]?.timeout).toBe(25_000);
+  });
+
+  it("skips warm-up when the probe consumes the command timeout budget", () => {
+    const runCaptureExImpl = vi.fn(() => successfulWarmResult());
+    const now = vi.fn().mockReturnValueOnce(1_000).mockReturnValueOnce(31_000);
+
+    expect(
+      maybeWarmOllamaAfterDaemonRestart(
+        { provider: "ollama-local", model: "qwen3.6:35b" },
+        {
+          probeRuntimeModelStatus: () => unloadedStatus,
+          runCaptureExImpl,
+          timeoutSeconds: 30,
+          now,
+        },
+      ),
+    ).toEqual({
+      kind: "skipped",
+      reason: "deadline-exhausted",
+      endpoint: `http://127.0.0.1:${OLLAMA_PORT}`,
+    });
+    expect(runCaptureExImpl).not.toHaveBeenCalled();
+  });
+
   it("does not treat an exit-zero Ollama error body as a successful warm-up", () => {
     expect(
       maybeWarmOllamaAfterDaemonRestart(

src/lib/actions/sandbox/agent/ollama-restart-recovery.ts

@@ -19,12 +19,12 @@ import { OLLAMA_PORT, OLLAMA_PROXY_PORT } from "../../../core/ports";
 import {
   describeModelInventory,
   createOllamaApiCapture,
+  createOllamaApiCaptureEx,
   getOllamaApiCommand,
   getResolvedOllamaHost,
   ollamaInventoryContainsModel,
   OLLAMA_HOST_DOCKER_INTERNAL,
   OLLAMA_LOCALHOST,
-  prepareOllamaApiExecution,
   probeOllamaEndpointInventory,
   type RunCaptureFn,
   type RunCaptureExFn,
@@ -33,15 +33,19 @@ import {
   type OllamaRuntimeModelStatus,
   probeOllamaRuntimeModelStatus,
 } from "../../../inference/ollama-runtime-context";
-import { runCaptureEx } from "../../../runner";
+import { redact, redactFull } from "../../../runner";
 
 export interface OllamaRestartRecoveryRoute {
   provider?: string | null;
   model?: string | null;
   endpointUrl?: string | null;
 }
 
-export interface OllamaRestartRecoveryDeps {
+export interface OllamaRestartRecoveryOptions {
+  timeoutSeconds?: number;
+}
+
+export interface OllamaRestartRecoveryDeps extends OllamaRestartRecoveryOptions {
   probeRuntimeModelStatus?: (
     model: string,
     getOllamaHost: () => string,
@@ -52,7 +56,7 @@ export interface OllamaRestartRecoveryDeps {
   getOllamaHost?: () => string;
   runCaptureImpl?: RunCaptureFn;
   prepareDockerEnvironment?: Parameters<typeof createOllamaApiCapture>[2];
-  prepareOllamaApiExecution?: typeof prepareOllamaApiExecution;
+  now?: () => number;
 }
 
 export type OllamaRestartRecoveryFailureReason =
@@ -63,7 +67,9 @@ export type OllamaRestartRecoveryFailureReason =
   | "spawn-failed";
 
 export type OllamaRestartRecoveryResult =
-  | { kind: "skipped"; reason: "not-ollama" | "missing-model" | "already-loaded" | "unreachable" }
+  | { kind: "skipped"; reason: "not-ollama" | "missing-model" | "already-loaded" }
+  | { kind: "skipped"; reason: "unreachable"; endpoint: string }
+  | { kind: "skipped"; reason: "deadline-exhausted"; endpoint: string }
   | { kind: "skipped"; reason: "model-absent"; endpoint: string; inventoryLabel: string }
   | { kind: "warmed"; ok: true; timedOut: false }
   | {
@@ -145,7 +151,7 @@ function resolveRawOllamaHost(
   return getAllowedFallbackHost(getOllamaHost);
 }
 
-function buildWarmCommand(model: string, hostname: string): string[] {
+function buildWarmCommand(model: string, hostname: string, maxTimeSeconds: number): string[] {
   const body = JSON.stringify({
     model,
     prompt: "Hello, reply in less than 5 words",
@@ -160,7 +166,7 @@ function buildWarmCommand(model: string, hostname: string): string[] {
       "--connect-timeout",
       "3",
       "--max-time",
-      String(OLLAMA_RESTART_RECOVERY_TIMEOUT_SECONDS),
+      String(maxTimeSeconds),
       "-H",
       "Content-Type: application/json",
       "-d",
@@ -171,6 +177,24 @@ function buildWarmCommand(model: string, hostname: string): string[] {
   );
 }
 
+function recoveryDeadlineMilliseconds(
+  timeoutSeconds: number | undefined,
+  now: () => number,
+): number | null {
+  if (timeoutSeconds === undefined) return null;
+  const boundedSeconds =
+    Number.isFinite(timeoutSeconds) && timeoutSeconds > 0
+      ? Math.min(timeoutSeconds, OLLAMA_RESTART_RECOVERY_TIMEOUT_SECONDS)
+      : 0;
+  return now() + boundedSeconds * 1000;
+}
+
+function remainingRecoveryMilliseconds(deadline: number | null, now: () => number): number {
+  return deadline === null
+    ? OLLAMA_RESTART_RECOVERY_TIMEOUT_SECONDS * 1000
+    : Math.max(0, Math.floor(deadline - now()));
+}
+
 function validateWarmResponse(stdout: string): "ok" | "ollama-error" | "invalid-response" {
   try {
     const parsed = JSON.parse(stdout) as {
@@ -193,9 +217,11 @@ function validateWarmResponse(stdout: string): "ok" | "ollama-error" | "invalid-
   }
 }
 
-function boundedWarmFailureDetail(value: unknown, fallback: string): string {
-  const detail = String(value ?? "")
-    .replace(/\s+/g, " ")
+export function boundedOllamaRestartRecoveryDetail(value: unknown, fallback: string): string {
+  const raw = value instanceof Error ? value.message : String(value ?? "");
+  const detail = redactFull(redact(raw))
+    .replace(/[\u0000-\u001f\u007f-\u009f\u061c\u200e\u200f\u2028-\u202e\u2066-\u2069]/gu, " ")
+    .replace(/\s+/gu, " ")
     .trim();
   return (detail || fallback).slice(0, 300);
 }
@@ -217,6 +243,9 @@ export function maybeWarmOllamaAfterDaemonRestart(
     return { kind: "skipped", reason: "missing-model" };
   }
 
+  const now = deps.now ?? Date.now;
+  const recoveryDeadline = recoveryDeadlineMilliseconds(deps.timeoutSeconds, now);
+
   const getOllamaHost = deps.getOllamaHost ?? getResolvedOllamaHost;
   const rawHost = resolveRawOllamaHost(route.endpointUrl, getOllamaHost);
   const rawEndpoint = `http://${rawHost}:${OLLAMA_PORT}`;
@@ -226,44 +255,44 @@ export function maybeWarmOllamaAfterDaemonRestart(
     rawHost,
     deps.prepareDockerEnvironment,
   );
+  const rawCaptureEx = createOllamaApiCaptureEx(
+    deps.runCaptureExImpl,
+    rawHost,
+    deps.prepareDockerEnvironment,
+  );
   let status: OllamaRuntimeModelStatus;
   try {
     status = probe(model, () => rawHost, rawCapture);
   } catch {
-    return { kind: "skipped", reason: "unreachable" };
+    return { kind: "skipped", reason: "unreachable", endpoint: rawEndpoint };
   }
   if (!status.probed) {
-    return { kind: "skipped", reason: "unreachable" };
+    return { kind: "skipped", reason: "unreachable", endpoint: rawEndpoint };
   }
   if (status.loaded) {
     return { kind: "skipped", reason: "already-loaded" };
   }
 
-  const captureEx = deps.runCaptureExImpl ?? runCaptureEx;
+  const warmupTimeoutMilliseconds = remainingRecoveryMilliseconds(recoveryDeadline, now);
+  if (warmupTimeoutMilliseconds === 0) {
+    return { kind: "skipped", reason: "deadline-exhausted", endpoint: rawEndpoint };
+  }
+  const warmupTimeoutSeconds = warmupTimeoutMilliseconds / 1000;
+
   try {
-    const execution = (deps.prepareOllamaApiExecution ?? prepareOllamaApiExecution)(
-      buildWarmCommand(model, rawHost),
-      rawHost,
-      { operation: `Ollama restart warm-up for '${model}'` },
-    );
-    let result;
-    try {
-      result = captureEx(execution.command, {
-        ...(execution.env === undefined ? {} : { env: execution.env }),
-      });
-    } finally {
-      execution.cleanup();
-    }
+    const result = rawCaptureEx(buildWarmCommand(model, rawHost, warmupTimeoutSeconds), {
+      timeout: warmupTimeoutMilliseconds,
+    });
     if (result.timedOut) {
       return {
         kind: "warmed",
         ok: false,
         timedOut: true,
         reason: "timeout",
         endpoint: rawEndpoint,
-        detail: boundedWarmFailureDetail(
+        detail: boundedOllamaRestartRecoveryDetail(
           result.stderr,
-          `warm-up exceeded ${OLLAMA_RESTART_RECOVERY_TIMEOUT_SECONDS} seconds`,
+          `warm-up exceeded ${String(warmupTimeoutSeconds)} seconds`,
         ),
       };
     }
@@ -274,7 +303,7 @@ export function maybeWarmOllamaAfterDaemonRestart(
         timedOut: false,
         reason: "command-failed",
         endpoint: rawEndpoint,
-        detail: boundedWarmFailureDetail(
+        detail: boundedOllamaRestartRecoveryDetail(
           result.stderr || result.stdout,
           `warm-up exited ${String(result.exitCode)}`,
         ),
@@ -305,7 +334,7 @@ export function maybeWarmOllamaAfterDaemonRestart(
         timedOut: false,
         reason: response,
         endpoint: rawEndpoint,
-        detail: boundedWarmFailureDetail(result.stdout, `Ollama returned ${response}`),
+        detail: boundedOllamaRestartRecoveryDetail(result.stdout, `Ollama returned ${response}`),
       };
     }
     return { kind: "warmed", ok: true, timedOut: false };
@@ -316,7 +345,7 @@ export function maybeWarmOllamaAfterDaemonRestart(
       timedOut: false,
       reason: "spawn-failed",
       endpoint: rawEndpoint,
-      detail: boundedWarmFailureDetail(
+      detail: boundedOllamaRestartRecoveryDetail(
         error instanceof Error ? error.message : error,
         "warm-up process could not start",
       ),