diff --git a/typescript/server/package.json b/typescript/server/package.json index 7df77b6..d1b31e4 100644 --- a/typescript/server/package.json +++ b/typescript/server/package.json @@ -28,7 +28,7 @@ "start": "node dist/main.js" }, "dependencies": { - "@smooai/smooth-operator-core": "^0.1.1", + "@smooai/smooth-operator-core": "^0.20.1", "ws": "^8.18.0" }, "devDependencies": { diff --git a/typescript/server/src/frameDispatcher.ts b/typescript/server/src/frameDispatcher.ts index 935b143..ff898cd 100644 --- a/typescript/server/src/frameDispatcher.ts +++ b/typescript/server/src/frameDispatcher.ts @@ -15,6 +15,7 @@ import { randomUUID } from 'node:crypto'; import { ANONYMOUS_ACCESS, type AccessContext } from './auth.js'; import { ConfirmationRegistry } from './confirmation.js'; +import type { ModelCeilingResolver } from './modelCeiling.js'; import * as protocol from './protocol.js'; import type { Frame } from './protocol.js'; import type { SessionStore } from './sessionStore.js'; @@ -52,6 +53,10 @@ export interface FrameDispatcherOptions { * connection). Created on demand if not supplied. */ confirmations?: ConfirmationRegistry; + /** Model id for turns (default {@link DEFAULT_MODEL}); forwarded to the {@link TurnRunner}. */ + model?: string; + /** Best-effort per-model output-ceiling resolver; forwarded to the {@link TurnRunner} (EPIC th-1cc9fa). */ + modelCeiling?: ModelCeilingResolver; } export class FrameDispatcher { @@ -63,6 +68,8 @@ export class FrameDispatcher { private readonly tools: Tool[]; private readonly confirmTools: string[]; private readonly confirmations: ConfirmationRegistry; + private readonly model?: string; + private readonly modelCeiling?: ModelCeilingResolver; /** In-flight spawned `send_message` turns, tracked so teardown can await them. */ private readonly turns = new Set>(); @@ -75,6 +82,8 @@ export class FrameDispatcher { this.tools = options.tools ?? []; this.confirmTools = options.confirmTools ?? []; this.confirmations = options.confirmations ?? new ConfirmationRegistry(); + this.model = options.model; + this.modelCeiling = options.modelCeiling; } /** @@ -213,6 +222,8 @@ export class FrameDispatcher { confirmTools: this.confirmTools, confirmations: this.confirmations, sessionId, + model: this.model, + modelCeiling: this.modelCeiling, }); // Run the turn as a background task, NOT awaited inline. A turn that calls a diff --git a/typescript/server/src/index.ts b/typescript/server/src/index.ts index 236d984..3c93ba9 100644 --- a/typescript/server/src/index.ts +++ b/typescript/server/src/index.ts @@ -27,9 +27,12 @@ export type { AccessKnowledge, FrameDispatcherOptions } from './frameDispatcher. export { ConfirmationRegistry } from './confirmation.js'; -export { TurnRunner } from './turnRunner.js'; +export { DEFAULT_MAX_ITERATIONS, DEFAULT_MAX_TOKENS, DEFAULT_MODEL, TurnRunner } from './turnRunner.js'; export type { Sink, TurnResult, TurnRunnerOptions } from './turnRunner.js'; +export { createGatewayModelCeilingResolver, extractModelCeilings } from './modelCeiling.js'; +export type { FetchLike, ModelCeilingResolver } from './modelCeiling.js'; + export { InMemorySessionStore } from './sessionStore.js'; export type { MessageDirection, SessionStore, StoredMessage, StoredSession } from './sessionStore.js'; diff --git a/typescript/server/src/main.ts b/typescript/server/src/main.ts index bd848fe..9362b36 100644 --- a/typescript/server/src/main.ts +++ b/typescript/server/src/main.ts @@ -13,12 +13,31 @@ * SMOOTH_OPERATOR_PORT bind port (default 8787) * SMOOAI_GATEWAY_URL OpenAI-compatible base URL (enables live turns with a key) * SMOOAI_GATEWAY_KEY gateway API key - * SMOOAI_MODEL model id (default gpt-4o-mini) + * SMOOAI_MODEL model id (default claude-haiku-4-5) */ import type { ChatClientLike } from '@smooai/smooth-operator-core'; +import { createGatewayModelCeilingResolver, type ModelCeilingResolver } from './modelCeiling.js'; import { serveLocal } from './server.js'; +/** The model id turns run against — SMOOAI_MODEL, else the engine's default. */ +function resolveModel(): string { + return process.env.SMOOAI_MODEL ?? 'claude-haiku-4-5'; +} + +/** + * A per-model output-ceiling resolver backed by the gateway's `/model/info`, so each + * turn clamps `max_tokens` to what the model can physically emit (EPIC th-1cc9fa). Only + * built when a gateway url+key are configured; otherwise `undefined` ⇒ turns run + * unclamped (behaviour unchanged on the keyless local path). + */ +function buildModelCeiling(): ModelCeilingResolver | undefined { + const url = process.env.SMOOAI_GATEWAY_URL; + const key = process.env.SMOOAI_GATEWAY_KEY; + if (!url || !key) return undefined; + return createGatewayModelCeilingResolver(url, key); +} + /** * A keyless client: every model call rejects, so `send_message` surfaces a clean * protocol error (the dispatcher's catch → INTERNAL_ERROR) instead of hanging. The @@ -53,8 +72,10 @@ async function buildChatClient(): Promise { // keyless local flavor needs no LLM SDK). Production installs it alongside. const openaiModule = 'openai'; const mod = (await import(openaiModule)) as { default: new (opts: { apiKey: string; baseURL: string }) => ChatClientLike }; - const model = process.env.SMOOAI_MODEL ?? 'gpt-4o-mini'; - process.env.SMOOAI_MODEL = model; + // Pin the resolved model into the env so the turn runner and the ceiling lookup + // agree on which model is in play (the request model and its /model/info ceiling + // must be the same model). + process.env.SMOOAI_MODEL = resolveModel(); return new mod.default({ apiKey: key, baseURL: url }); } catch { // The `openai` package isn't installed — fall back to the keyless client so @@ -68,7 +89,7 @@ async function main(): Promise { const port = Number(process.env.SMOOTH_OPERATOR_PORT ?? '8787'); const chatClient = await buildChatClient(); - const server = await serveLocal({ chatClient, host, port }); + const server = await serveLocal({ chatClient, host, port, model: resolveModel(), modelCeiling: buildModelCeiling() }); // eslint-disable-next-line no-console console.log(`smooth-operator-server (TypeScript, local flavor) listening on ${server.url}`); // serveLocal already wires SIGTERM/SIGINT → graceful drain + close. diff --git a/typescript/server/src/modelCeiling.ts b/typescript/server/src/modelCeiling.ts new file mode 100644 index 0000000..196112e --- /dev/null +++ b/typescript/server/src/modelCeiling.ts @@ -0,0 +1,93 @@ +/** + * Per-model output-ceiling lookup from the LLM gateway's `/model/info` — the + * consumer half of the model-output clamp (EPIC th-1cc9fa). + * + * The engine (`@smooai/smooth-operator-core`) clamps a turn's `max_tokens` to + * `min(maxTokens, modelMaxOutput)` but takes no LiteLLM-specific HTTP itself. The + * server sources the ceiling here and passes it in via `AgentOptions.modelMaxOutput`, + * mirroring the Rust server's `admin::fetch_model_costs` / `model_output_ceiling` / + * `map_model_info` split (`rust/smooth-operator-server/src/admin.rs`). + * + * **Best-effort**: any gateway/transport/decode error, an unknown model, or a model + * whose gateway entry has no ceiling ⇒ `undefined` ⇒ the engine leaves `max_tokens` + * unclamped (graceful, no behaviour change). The `/model/info` map is fetched at most + * once per process (cached on first success; failures are not cached, so the next turn + * retries). + */ + +/** Resolve a model id to its hard output ceiling (`max_output_tokens`), or `undefined`. */ +export type ModelCeilingResolver = (model: string) => Promise; + +/** The subset of `fetch` this module needs — injectable so the resolver is unit-testable offline. */ +export type FetchLike = (url: string, init?: { headers?: Record }) => Promise<{ ok: boolean; json: () => Promise }>; + +/** + * Map a gateway `/model/info` payload + * (`{ data: [{ model_name, model_info: { max_output_tokens, … } }] }`) to a + * `model_name → max_output_tokens` map. The TS analog of the Rust `map_model_info`, + * narrowed to the one field the clamp needs. Entries without a `model_name`, without a + * positive integer `max_output_tokens`, are skipped. Pure + network-free so it's + * unit-testable on a sample payload. + */ +export function extractModelCeilings(payload: unknown): Map { + const out = new Map(); + const data = (payload as { data?: unknown })?.data; + if (!Array.isArray(data)) return out; + for (const entry of data) { + if (typeof entry !== 'object' || entry === null) continue; + const name = (entry as { model_name?: unknown }).model_name; + if (typeof name !== 'string' || name.length === 0) continue; + const info = (entry as { model_info?: unknown }).model_info; + const max = typeof info === 'object' && info !== null ? (info as { max_output_tokens?: unknown }).max_output_tokens : undefined; + if (typeof max === 'number' && Number.isInteger(max) && max > 0) { + out.set(name, max); + } + } + return out; +} + +/** + * Build a {@link ModelCeilingResolver} backed by the gateway's `/model/info`. + * + * `gatewayUrl` is the OpenAI-compatible base url (e.g. `https://llm.smoo.ai/v1`); the + * model-info endpoint is `{gatewayUrl}/model/info`. `gatewayKey`, when present, is + * sent as a bearer token (the same creds the turns use). The whole model map is + * fetched at most once per process on the first lookup and cached; a lost race just + * recomputes the same stable map. Any error ⇒ every lookup returns `undefined` (and + * the failure is not cached, so a later turn retries). + */ +export function createGatewayModelCeilingResolver(gatewayUrl: string, gatewayKey?: string, fetchImpl: FetchLike = fetch as unknown as FetchLike): ModelCeilingResolver { + const url = `${gatewayUrl.replace(/\/+$/, '')}/model/info`; + let cached: Map | undefined; + let inflight: Promise> | undefined; + + const load = async (): Promise> => { + if (cached) return cached; + // Single-flight: concurrent first turns share one fetch rather than stampeding + // the gateway. On failure we return an empty map WITHOUT caching, so the next + // turn retries (matches the Rust best-effort-not-cached-on-error behaviour). + if (!inflight) { + inflight = (async () => { + try { + const res = await fetchImpl(url, gatewayKey ? { headers: { authorization: `Bearer ${gatewayKey}` } } : undefined); + if (!res.ok) throw new Error(`gateway /model/info returned non-ok`); + const map = extractModelCeilings(await res.json()); + cached = map; + return map; + } finally { + inflight = undefined; + } + })(); + } + try { + return await inflight; + } catch { + return new Map(); + } + }; + + return async (model: string): Promise => { + const ceiling = (await load()).get(model); + return ceiling !== undefined && ceiling > 0 ? ceiling : undefined; + }; +} diff --git a/typescript/server/src/server.ts b/typescript/server/src/server.ts index 238206c..40fe2ab 100644 --- a/typescript/server/src/server.ts +++ b/typescript/server/src/server.ts @@ -24,6 +24,7 @@ import { WebSocketServer, type WebSocket } from 'ws'; import { ANONYMOUS_ACCESS } from './auth.js'; import { Backplane, InMemoryBackplane } from './backplane.js'; import { type AccessKnowledge, FrameDispatcher } from './frameDispatcher.js'; +import type { ModelCeilingResolver } from './modelCeiling.js'; import type { Frame } from './protocol.js'; import type { ChatClientLike, Tool } from '@smooai/smooth-operator-core'; import type { AuthVerifier } from './auth.js'; @@ -51,6 +52,14 @@ export interface ServerOptions { * client replies with `confirm_tool_action`. */ confirmTools?: string[]; + /** Model id for turns (default the engine's own default); forwarded to each connection's dispatcher. */ + model?: string; + /** + * Best-effort per-model output-ceiling resolver (from the gateway's `/model/info`). + * When set, each turn clamps `max_tokens` to what the model can physically emit + * (EPIC th-1cc9fa). Absent (tests, keyless local) ⇒ unclamped, behaviour unchanged. + */ + modelCeiling?: ModelCeilingResolver; /** WS path to mount on (default `/ws`). */ path?: string; } @@ -108,6 +117,8 @@ export function buildServer(options: ServerOptions): { systemPrompt: options.systemPrompt, tools: options.tools, confirmTools: options.confirmTools, + model: options.model, + modelCeiling: options.modelCeiling, }); // Fire-and-forget the per-connection loop; it owns the socket's lifecycle. void runConnection(socket, dispatcher, backplane, drain.signal); @@ -167,13 +178,22 @@ export async function serve(options: ServerOptions & { host?: string; port?: num * embeddable in-process. The TS analog of the Rust `serve_local` / `LocalServer` * and the C# in-memory host. Everything in memory, loopback bind, no auth. */ -export async function serveLocal(options: { chatClient: ChatClientLike; host?: string; port?: number; knowledge?: AccessKnowledge }): Promise { +export async function serveLocal(options: { + chatClient: ChatClientLike; + host?: string; + port?: number; + knowledge?: AccessKnowledge; + model?: string; + modelCeiling?: ModelCeilingResolver; +}): Promise { return serve({ chatClient: options.chatClient, store: new InMemorySessionStore(), auth: new NoAuthVerifier(), backplane: new InMemoryBackplane(), knowledge: options.knowledge, + model: options.model, + modelCeiling: options.modelCeiling, host: options.host ?? '127.0.0.1', port: options.port ?? 0, }); diff --git a/typescript/server/src/turnRunner.ts b/typescript/server/src/turnRunner.ts index 082206c..5e485da 100644 --- a/typescript/server/src/turnRunner.ts +++ b/typescript/server/src/turnRunner.ts @@ -15,6 +15,7 @@ import { approve, deny, SmoothAgent } from '@smooai/smooth-operator-core'; import type { AgentOptions, ChatClientLike, HumanApprovalRequest, HumanApprovalResponse, Knowledge, StreamEvent, Tool } from '@smooai/smooth-operator-core'; import type { ConfirmationRegistry } from './confirmation.js'; +import type { ModelCeilingResolver } from './modelCeiling.js'; import * as protocol from './protocol.js'; import type { Citation, Frame } from './protocol.js'; import type { SessionStore } from './sessionStore.js'; @@ -30,6 +31,25 @@ const AUTO_CONTEXT_LIMIT = 3; const MAX_PRIOR_MESSAGES = 50; const CITATION_SNIPPET_MAX_CHARS = 280; +/** + * Default model when the host doesn't set one. Matches the engine's own default so + * behaviour is unchanged; also the model the per-turn output ceiling is looked up for. + */ +export const DEFAULT_MODEL = 'claude-haiku-4-5'; +/** + * Default agent-loop iteration cap. Was the engine's chat-widget-sized 6 — too tight + * for any multi-step turn. Raised to 20 for agentic use (EPIC th-1cc9fa). + */ +export const DEFAULT_MAX_ITERATIONS = 20; +/** + * Default `max_tokens` per LLM call. Was 512 (chat-widget sizing), which STARVES + * reasoning models — they spend the whole budget on `reasoning_content` and return + * empty `content`. Raised to 8192 (EPIC th-1cc9fa). Safe now that the per-model output + * ceiling clamps this down per request: a cap only bounds runaway output, it doesn't + * lengthen concise answers. + */ +export const DEFAULT_MAX_TOKENS = 8192; + const DEFAULT_SYSTEM_PROMPT = 'You are a helpful customer support agent. Answer using only the knowledge provided to you; if it is not there, say you don\'t know.'; @@ -55,6 +75,14 @@ export interface TurnRunnerOptions { confirmations?: ConfirmationRegistry; /** The session id a parked confirmation is keyed by (so a `confirm_tool_action` frame routes here). */ sessionId?: string; + /** Model id for the turn (default {@link DEFAULT_MODEL}); also the model whose output ceiling is looked up. */ + model?: string; + /** + * Best-effort per-model output-ceiling resolver (from the gateway's `/model/info`). + * When set, each turn clamps `max_tokens` to `min(DEFAULT_MAX_TOKENS, ceiling)` via + * the engine's `modelMaxOutput`. Absent (tests, keyless local) ⇒ unclamped (EPIC th-1cc9fa). + */ + modelCeiling?: ModelCeilingResolver; } export class TurnRunner { @@ -66,6 +94,8 @@ export class TurnRunner { private readonly confirmTools: string[]; private readonly confirmations?: ConfirmationRegistry; private readonly sessionId?: string; + private readonly model: string; + private readonly modelCeiling?: ModelCeilingResolver; constructor(options: TurnRunnerOptions) { this.chatClient = options.chatClient; @@ -76,6 +106,8 @@ export class TurnRunner { this.confirmTools = options.confirmTools ?? []; this.confirmations = options.confirmations; this.sessionId = options.sessionId; + this.model = options.model ?? DEFAULT_MODEL; + this.modelCeiling = options.modelCeiling; } /** True when `name` matches a confirmation-gated pattern (substring, like the Rust hook). */ @@ -112,10 +144,23 @@ export class TurnRunner { // 2. Build the agent + replay prior history as the thread (before persisting // this turn's inbound message). The engine consumes history as OpenAI-format // messages passed to runStream. - const agentOptions: AgentOptions = { instructions: this.systemPrompt }; + const agentOptions: AgentOptions = { + instructions: this.systemPrompt, + model: this.model, + maxTokens: DEFAULT_MAX_TOKENS, + maxIterations: DEFAULT_MAX_ITERATIONS, + }; if (this.knowledge) agentOptions.knowledge = this.knowledge; if (this.tools.length > 0) agentOptions.tools = this.tools; + // Clamp max_tokens to the resolved model's output ceiling (best-effort; a + // missing/unknown ceiling ⇒ unclamped). Reuses the cached /model/info fetch. + // EPIC th-1cc9fa — the consumer half of the engine's model-output clamp. + if (this.modelCeiling) { + const ceiling = await this.modelCeiling(this.model); + if (ceiling !== undefined) agentOptions.modelMaxOutput = ceiling; + } + // Write-confirmation HITL: when configured with tool patterns AND a registry // is present, install a HumanGate that parks the turn before a gated tool runs // (emit `write_confirmation_required`, await the client's verdict via the diff --git a/typescript/server/test/modelCeiling.test.ts b/typescript/server/test/modelCeiling.test.ts new file mode 100644 index 0000000..7868537 --- /dev/null +++ b/typescript/server/test/modelCeiling.test.ts @@ -0,0 +1,140 @@ +/** + * Unit tests for the gateway model-output-ceiling lookup (EPIC th-1cc9fa) — the + * consumer half of the engine's `max_tokens` clamp. All offline: the `/model/info` + * fetch is injected as a {@link FetchLike} fake, so no gateway is contacted. + */ +import { describe, expect, it } from 'vitest'; + +import { createGatewayModelCeilingResolver, extractModelCeilings, type FetchLike } from '../src/modelCeiling.js'; + +/** A representative LiteLLM `/model/info` payload. */ +const SAMPLE = { + data: [ + { model_name: 'claude-haiku-4-5', model_info: { max_output_tokens: 8192, input_cost_per_token: 0.000001 } }, + { model_name: 'claude-opus-4', model_info: { max_output_tokens: 65536 } }, + { model_name: 'groq-compound', model_info: { max_output_tokens: 8192 } }, + ], +}; + +/** A `FetchLike` that always resolves the given payload with `ok: true`, counting calls. */ +function okFetch(payload: unknown): { fetch: FetchLike; calls: () => number; lastInit: () => unknown } { + let count = 0; + let last: unknown; + const fetch: FetchLike = (_url, init) => { + count += 1; + last = init; + return Promise.resolve({ ok: true, json: () => Promise.resolve(payload) }); + }; + return { fetch, calls: () => count, lastInit: () => last }; +} + +describe('extractModelCeilings', () => { + it('maps model_name -> max_output_tokens from a sample payload', () => { + const map = extractModelCeilings(SAMPLE); + expect(map.get('claude-haiku-4-5')).toBe(8192); + expect(map.get('claude-opus-4')).toBe(65536); + expect(map.get('groq-compound')).toBe(8192); + expect(map.size).toBe(3); + }); + + it('skips entries with no model_name, no model_info, or a non-positive/non-integer ceiling', () => { + const map = extractModelCeilings({ + data: [ + { model_info: { max_output_tokens: 4096 } }, // no model_name + { model_name: 'bare' }, // no model_info + { model_name: 'zero', model_info: { max_output_tokens: 0 } }, // non-positive + { model_name: 'floaty', model_info: { max_output_tokens: 1024.5 } }, // non-integer + { model_name: 'good', model_info: { max_output_tokens: 2048 } }, + ], + }); + expect(map.size).toBe(1); + expect(map.get('good')).toBe(2048); + }); + + it('returns an empty map when data is missing or not an array', () => { + expect(extractModelCeilings({}).size).toBe(0); + expect(extractModelCeilings({ data: 'nope' }).size).toBe(0); + expect(extractModelCeilings(null).size).toBe(0); + }); +}); + +describe('createGatewayModelCeilingResolver', () => { + it('resolves a known model to its ceiling', async () => { + const { fetch } = okFetch(SAMPLE); + const resolve = createGatewayModelCeilingResolver('https://llm.smoo.ai/v1', 'k', fetch); + expect(await resolve('claude-haiku-4-5')).toBe(8192); + expect(await resolve('groq-compound')).toBe(8192); + }); + + it('returns undefined for an unknown model (best-effort, no clamp)', async () => { + const { fetch } = okFetch(SAMPLE); + const resolve = createGatewayModelCeilingResolver('https://llm.smoo.ai/v1', 'k', fetch); + expect(await resolve('who-dis')).toBeUndefined(); + }); + + it('fetches /model/info at most once and caches the map', async () => { + const { fetch, calls } = okFetch(SAMPLE); + const resolve = createGatewayModelCeilingResolver('https://llm.smoo.ai/v1', 'k', fetch); + await resolve('claude-haiku-4-5'); + await resolve('claude-opus-4'); + await resolve('groq-compound'); + expect(calls()).toBe(1); + }); + + it('hits {gatewayUrl}/model/info with a trimmed base and a bearer header when a key is set', async () => { + const seen: { url?: string; init?: unknown } = {}; + const fetch: FetchLike = (url, init) => { + seen.url = url; + seen.init = init; + return Promise.resolve({ ok: true, json: () => Promise.resolve(SAMPLE) }); + }; + const resolve = createGatewayModelCeilingResolver('https://llm.smoo.ai/v1/', 'secret', fetch); + await resolve('claude-haiku-4-5'); + expect(seen.url).toBe('https://llm.smoo.ai/v1/model/info'); + expect(seen.init).toEqual({ headers: { authorization: 'Bearer secret' } }); + }); + + it('sends no auth header when no key is provided', async () => { + const seen: { init?: unknown } = {}; + const fetch: FetchLike = (_url, init) => { + seen.init = init; + return Promise.resolve({ ok: true, json: () => Promise.resolve(SAMPLE) }); + }; + const resolve = createGatewayModelCeilingResolver('https://llm.smoo.ai/v1', undefined, fetch); + await resolve('claude-haiku-4-5'); + expect(seen.init).toBeUndefined(); + }); + + it('is best-effort: a non-ok response yields undefined and is not cached (retries next call)', async () => { + let ok = false; + let count = 0; + const fetch: FetchLike = () => { + count += 1; + return Promise.resolve({ ok, json: () => Promise.resolve(SAMPLE) }); + }; + const resolve = createGatewayModelCeilingResolver('https://llm.smoo.ai/v1', 'k', fetch); + expect(await resolve('claude-haiku-4-5')).toBeUndefined(); // first fetch fails + ok = true; // gateway recovers + expect(await resolve('claude-haiku-4-5')).toBe(8192); // retried, not stuck on the failure + expect(count).toBe(2); + }); + + it('is best-effort: a thrown fetch yields undefined', async () => { + const fetch: FetchLike = () => Promise.reject(new Error('connreset')); + const resolve = createGatewayModelCeilingResolver('https://llm.smoo.ai/v1', 'k', fetch); + expect(await resolve('claude-haiku-4-5')).toBeUndefined(); + }); + + it('shares a single in-flight fetch across concurrent first lookups', async () => { + let count = 0; + const fetch: FetchLike = () => { + count += 1; + return new Promise((resolve) => setTimeout(() => resolve({ ok: true, json: () => Promise.resolve(SAMPLE) }), 5)); + }; + const resolve = createGatewayModelCeilingResolver('https://llm.smoo.ai/v1', 'k', fetch); + const [a, b] = await Promise.all([resolve('claude-haiku-4-5'), resolve('claude-opus-4')]); + expect(a).toBe(8192); + expect(b).toBe(65536); + expect(count).toBe(1); + }); +}); diff --git a/typescript/server/test/turn-max-tokens.test.ts b/typescript/server/test/turn-max-tokens.test.ts new file mode 100644 index 0000000..56adaa4 --- /dev/null +++ b/typescript/server/test/turn-max-tokens.test.ts @@ -0,0 +1,67 @@ +/** + * The turn runner threads the raised starvation defaults and the per-turn model-output + * ceiling into the engine (EPIC th-1cc9fa). Driven offline with the shared + * `MockLlmProvider`, which records each request body so we can assert on `max_tokens`. + */ +import { MockLlmProvider } from '@smooai/smooth-operator-core'; +import { describe, expect, it } from 'vitest'; + +import { InMemorySessionStore } from '../src/sessionStore.js'; +import { DEFAULT_MAX_TOKENS, TurnRunner } from '../src/turnRunner.js'; +import type { Frame } from '../src/protocol.js'; + +/** Run one turn against a fresh conversation and return the recorded request body. */ +async function runTurn(mock: MockLlmProvider, runnerOptions: Partial[0]> = {}): Promise> { + const store = new InMemorySessionStore(); + const session = await store.createSession('agent-1'); + const runner = new TurnRunner({ chatClient: mock, store, ...runnerOptions }); + const sink = (_frame: Frame): void => {}; + await runner.run(session.conversationId, 'req-1', 'hello', sink); + return mock.calls[0].body; +} + +describe('TurnRunner max_tokens clamp + defaults', () => { + it('sends the raised DEFAULT_MAX_TOKENS when no ceiling resolver is set', async () => { + const body = await runTurn(new MockLlmProvider().pushText('hi')); + expect(body.max_tokens).toBe(DEFAULT_MAX_TOKENS); + expect(DEFAULT_MAX_TOKENS).toBe(8192); + }); + + it('clamps max_tokens down to the model ceiling when it is below the budget', async () => { + // A ceiling below DEFAULT_MAX_TOKENS (8192) must win. + const body = await runTurn(new MockLlmProvider().pushText('hi'), { + model: 'tiny-model', + modelCeiling: async (model) => (model === 'tiny-model' ? 4096 : undefined), + }); + expect(body.max_tokens).toBe(4096); + }); + + it('leaves max_tokens at the budget when the ceiling is >= the budget', async () => { + const body = await runTurn(new MockLlmProvider().pushText('hi'), { + model: 'big-model', + modelCeiling: async () => 65536, + }); + expect(body.max_tokens).toBe(DEFAULT_MAX_TOKENS); + }); + + it('leaves max_tokens unclamped when the resolver returns undefined (unknown model)', async () => { + const body = await runTurn(new MockLlmProvider().pushText('hi'), { + model: 'mystery', + modelCeiling: async () => undefined, + }); + expect(body.max_tokens).toBe(DEFAULT_MAX_TOKENS); + }); + + it('resolves the ceiling for the model the turn actually uses', async () => { + let askedFor: string | undefined; + const body = await runTurn(new MockLlmProvider().pushText('hi'), { + model: 'claude-haiku-4-5', + modelCeiling: async (model) => { + askedFor = model; + return 8192; + }, + }); + expect(askedFor).toBe('claude-haiku-4-5'); + expect(body.model).toBe('claude-haiku-4-5'); + }); +});