AkurAI Build
Menu

popagent

public

Latest change c5c74dbdbdce2026200ca34f315ffbf4df7da967 - Fail model gateway requests fast when no response arrives within 90s by AkurAI Build

import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
import { appLogger } from "./observability";
import { FREE_MODEL_ROUTE, LOCAL_MODEL_ROUTE, parseModelSpec, ROUTING_MODEL_ID, runtimeModelId, type ModelCatalog, type ModelProfile } from "./api-types";

const BASE_URL = process.env.NINEROUTER_URL ?? "http://127.0.0.1:20128/v1";
const API_KEY = process.env.NINEROUTER_API_KEY ?? "";
const EMBEDDING_BASE_URL = process.env.POPAGENT_EMBEDDING_URL ?? BASE_URL;
const EMBEDDING_API_KEY = process.env.POPAGENT_EMBEDDING_API_KEY ?? API_KEY;
const configuredLocalModelConcurrency = Number.parseInt(
  process.env.POPAGENT_LOCAL_MODEL_CONCURRENCY ?? "2",
  10,
);
const LOCAL_MODEL_CONCURRENCY =
  Number.isSafeInteger(configuredLocalModelConcurrency) && configuredLocalModelConcurrency > 0
    ? configuredLocalModelConcurrency
    : 2;
const ORNITH_MODEL_PATTERN = /\/ornith-1\.0-9b-mtp-q4_k_m$/i;

/**
 * FIFO admission control for local generations. A lease remains active until
 * the response stream ends or is cancelled, not merely until headers arrive.
 */
export function createLocalModelFetch(
  upstream: typeof fetch = globalThis.fetch,
  concurrency = LOCAL_MODEL_CONCURRENCY,
): typeof fetch {
  type Release = () => void;
  type Waiter = {
    resolve: (release: Release) => void;
    reject: (reason?: unknown) => void;
    signal?: AbortSignal;
    onAbort?: () => void;
  };
  const waiting: Waiter[] = [];
  let active = 0;

  const releaseLease = (): Release => {
    let released = false;
    return () => {
      if (released) return;
      released = true;
      const next = waiting.shift();
      if (!next) {
        active--;
        return;
      }
      if (next.signal && next.onAbort) next.signal.removeEventListener("abort", next.onAbort);
      next.resolve(releaseLease());
    };
  };

  const acquire = (signal?: AbortSignal): Promise<Release> => {
    if (signal?.aborted) return Promise.reject(signal.reason);
    if (active < concurrency) {
      active++;
      return Promise.resolve(releaseLease());
    }
    return new Promise<Release>((resolve, reject) => {
      const waiter: Waiter = { resolve, reject, signal };
      if (signal) {
        waiter.onAbort = () => {
          const index = waiting.indexOf(waiter);
          if (index >= 0) waiting.splice(index, 1);
          reject(signal.reason);
        };
        signal.addEventListener("abort", waiter.onAbort, { once: true });
      }
      waiting.push(waiter);
    });
  };

  const queuedFetch = async (input: RequestInfo | URL, init?: RequestInit) => {
    let queued = false;
    if (typeof init?.body === "string") {
      try {
        const body = JSON.parse(init.body) as { model?: unknown };
        queued = typeof body.model === "string" && ORNITH_MODEL_PATTERN.test(`/${body.model}`);
      } catch {
        // The upstream provider owns malformed request handling.
      }
    }
    if (!queued) return upstream(input, init);

    const release = await acquire(init?.signal ?? undefined);
    try {
      const response = await upstream(input, init);
      if (!response.body) {
        release();
        return response;
      }
      const reader = response.body.getReader();
      let finished = false;
      const finish = () => {
        if (finished) return;
        finished = true;
        release();
      };
      const body = new ReadableStream<Uint8Array>({
        async pull(controller) {
          try {
            const chunk = await reader.read();
            if (chunk.done) {
              finish();
              controller.close();
            } else {
              controller.enqueue(chunk.value);
            }
          } catch (error) {
            finish();
            controller.error(error);
          }
        },
        async cancel(reason) {
          try {
            await reader.cancel(reason);
          } finally {
            finish();
          }
        },
      });
      return new Response(body, {
        status: response.status,
        statusText: response.statusText,
        headers: response.headers,
      });
    } catch (error) {
      release();
      throw error;
    }
  };
  return queuedFetch as typeof fetch;
}

const configuredFirstByteTimeout = Number.parseInt(process.env.POPAGENT_MODEL_FIRST_BYTE_TIMEOUT_MS ?? "90000", 10);
export const MODEL_FIRST_BYTE_TIMEOUT_MS = Number.isSafeInteger(configuredFirstByteTimeout) && configuredFirstByteTimeout > 0
  ? configuredFirstByteTimeout
  : 90_000;

/**
 * Fail a gateway request that has not returned response headers within the
 * budget. 9Router's Claude route sometimes hangs on a stale pooled connection
 * for ~249 s before reporting "fetch connect timeout"; every task waited that
 * long before the routing policy could fall over. Streaming bodies are not
 * bounded here — only the time to first byte.
 */
export function createFirstByteTimeoutFetch(
  upstream: typeof fetch = globalThis.fetch,
  timeoutMs = MODEL_FIRST_BYTE_TIMEOUT_MS,
): typeof fetch {
  const timedFetch = async (input: RequestInfo | URL, init?: RequestInit) => {
    const controller = new AbortController();
    const onAbort = () => controller.abort(init?.signal?.reason);
    if (init?.signal?.aborted) onAbort();
    else init?.signal?.addEventListener("abort", onAbort, { once: true });
    const timer = setTimeout(() => {
      const url = typeof input === "string" ? input : input instanceof URL ? input.href : input.url;
      appLogger().warn("model.request.first_byte_timeout", { url, timeoutMs });
      controller.abort(new Error(`Model gateway returned no response within ${Math.round(timeoutMs / 1000)}s (first-byte timeout)`));
    }, timeoutMs);
    try {
      return await upstream(input, { ...init, signal: controller.signal });
    } finally {
      clearTimeout(timer);
      init?.signal?.removeEventListener("abort", onAbort);
    }
  };
  return timedFetch as typeof fetch;
}

export { FREE_MODEL_ROUTE, LOCAL_MODEL_ROUTE, runtimeModelId };

/**
 * Some strict providers (Claude-family via 9Router) reject a chat-completions
 * request whose message array ends on an "assistant" turn as invalid prefill,
 * even when the assistant message only contains pending tool calls. Mastra's
 * own internal continuation loop can legitimately produce such a request
 * (e.g. when a near-maxSteps iteration guard interrupts before a tool call is
 * resolved). This is the last point before the request is serialized for the
 * wire, so it is the correct, provider-agnostic place to guarantee every
 * outgoing request ends on a non-assistant message rather than crashing the
 * whole run. Trimming a dangling assistant tail here only ever removes an
 * already-invalid, already-about-to-fail turn.
 */
export function trimDanglingAssistantWireTail(request: Record<string, unknown>): Record<string, unknown> {
  if (!Array.isArray(request.messages)) return request;
  const messages = request.messages as Array<{ role?: unknown }>;
  let end = messages.length;
  while (end > 0 && messages[end - 1]?.role === "assistant") end--;
  if (end === messages.length) return request;
  appLogger().warn("model.request.trimmed_dangling_assistant_tail", {
    model: typeof request.model === "string" ? request.model : undefined,
    droppedCount: messages.length - end,
  });
  return { ...request, messages: messages.slice(0, end) };
}

/**
 * `model@low|medium|high` specs (see `modelSpec`) carry the reasoning effort in
 * the model id so every caller can keep passing one string; the suffix is
 * stripped here and sent as OpenAI-compatible `reasoning_effort`.
 */
export function applyReasoningEffort(request: Record<string, unknown>): Record<string, unknown> {
  if (typeof request.model !== "string") return request;
  const { model, reasoningEffort } = parseModelSpec(request.model);
  if (reasoningEffort === "default") return request;
  return { ...request, model, reasoning_effort: reasoningEffort };
}

/** The local Ornith llama.cpp template accepts exactly one leading system message. */
export function normalizeChatRequest(request: Record<string, unknown>): Record<string, unknown> {
  const trimmed = trimDanglingAssistantWireTail(applyReasoningEffort(request));
  if (typeof trimmed.model !== "string" || !ORNITH_MODEL_PATTERN.test(`/${trimmed.model}`) || !Array.isArray(trimmed.messages)) return trimmed;
  const messages = trimmed.messages as Array<{ role?: unknown; content?: unknown }>;
  let systemCount = 0;
  while (messages[systemCount]?.role === "system") systemCount++;
  if (systemCount < 2 || messages.slice(0, systemCount).some((message) => typeof message.content !== "string")) return trimmed;
  return {
    ...trimmed,
    messages: [
      { ...messages[0], content: messages.slice(0, systemCount).map((message) => message.content).join("\n\n") },
      ...messages.slice(systemCount),
    ],
  };
}

const ninerouter = createOpenAICompatible({
  name: "9router",
  baseURL: BASE_URL,
  apiKey: API_KEY,
  fetch: createLocalModelFetch(createFirstByteTimeoutFetch()),
  transformRequestBody: normalizeChatRequest,
});
const embeddings = createOpenAICompatible({
  name: "embeddings",
  baseURL: EMBEDDING_BASE_URL,
  apiKey: EMBEDDING_API_KEY,
});

/** Language model for one explicit 9Router model id. */
export function resolveModel(modelId: string) {
  return ninerouter.chatModel(modelId);
}

/** Embedding model for the configured OpenAI-compatible embedding service. */
export function resolveEmbeddingModel(
  modelId = process.env.POPAGENT_EMBEDDING_MODEL ?? "text-embedding-3-small",
) {
  return embeddings.embeddingModel(modelId);
}

const MODEL_PROFILES: Record<string, ModelProfile> = {
  [ROUTING_MODEL_ID]: {
    label: "Model routing policy",
    role: "orchestrator",
    summary: "Resolves to the routes configured under Settings → Runtime → Model routing (fallback or round-robin, with reasoning effort).",
    members: [],
  },
  "best-orchestrator": {
    label: "Fusion orchestrator",
    role: "orchestrator",
    summary: "Sol Max and Fable reason in parallel; OmniRoute synthesizes one final direction.",
    members: ["GPT-5.6 Sol Max", "Claude Fable 5"],
  },
  "fast-orchestrator": {
    label: "Fast orchestrator",
    role: "fast",
    summary: "Alternates Sol Max and Fable for routine work with one upstream call.",
    members: ["GPT-5.6 Sol Max", "Claude Fable 5"],
  },
  "titan/ornith-1.0-9b-mtp-q4_k_m": {
    label: "Local build model",
    role: "local",
    summary: "Titan-resident Ornith executes repository work without consuming cloud-model quota.",
    members: ["Ornith 1.0 9B MTP"],
  },
};

function availableModelProfiles(models: string[]): Record<string, ModelProfile> {
  return Object.fromEntries(
    models.flatMap((model) => MODEL_PROFILES[model] ? [[model, MODEL_PROFILES[model]]] : []),
  );
}


/** Selectable model ids and their context limits from the 9router gateway. */
export async function listModelCatalog(): Promise<ModelCatalog> {
  const res = await fetch(`${BASE_URL}/models`, {
    headers: { authorization: `Bearer ${API_KEY}` },
    signal: AbortSignal.timeout(5000),
  });
  if (!res.ok) throw new Error(`9router /models failed: ${res.status}`);
  const { data } = (await res.json()) as {
    data: {
      id: string;
      context_length?: number;
      max_input_tokens?: number;
      capabilities?: { contextWindow?: number };
    }[];
  };
  const contextWindows: Record<string, number> = {};
  for (const model of data) {
    const limit =
      model.context_length ??
      model.max_input_tokens ??
      model.capabilities?.contextWindow ??
      0;
    contextWindows[model.id] = Math.max(contextWindows[model.id] ?? 0, limit);
  }
  const models = [ROUTING_MODEL_ID, ...new Set(data.map((model) => model.id))];
  return {
    models,
    contextWindows,
    profiles: availableModelProfiles(models),
  };
}

let catalogCache: { catalog: ModelCatalog; expires: number } | undefined;
export async function isKnownModel(spec: string): Promise<boolean> {
  if (spec === ROUTING_MODEL_ID) return true;
  const { model } = parseModelSpec(spec);
  if (catalogCache && catalogCache.expires > Date.now()) return catalogCache.catalog.models.includes(model);
  try {
    const catalog = await listModelCatalog();
    catalogCache = { catalog, expires: Date.now() + 60_000 };
    return catalog.models.includes(model);
  } catch (error) {
    appLogger().warn("model.catalog.unavailable", { error });
    return true;
  }
}