@steel-dev/atlas 0.1.0

package/dist/model.js

@@ -0,0 +1,630 @@
+import { createHash } from "node:crypto";
+import { wrapLanguageModel } from "ai";
+import { sleep } from "./async.js";
+import { addTokenUsage, emptyTokenUsage, resolvePricing, usageCostUSD, } from "./budget.js";
+import { ECONOMY } from "./economy.js";
+import { AtlasError, errorMessage } from "./errors.js";
+import { currentFrame } from "./trace.js";
+export function createRunUsage() {
+    return { byRole: new Map(), replayedUSD: 0 };
+}
+function trackUsage(runUsage, role, usage) {
+    const existing = runUsage.byRole.get(role) ?? emptyTokenUsage();
+    addTokenUsage(existing, usage);
+    runUsage.byRole.set(role, existing);
+}
+export function totalFreshTokens(usage) {
+    let total = 0;
+    for (const roleUsage of usage.byRole.values()) {
+        total += roleUsage.input + roleUsage.output + roleUsage.cacheWrite;
+    }
+    return total;
+}
+export function tokenUsageFromV3(usage) {
+    const cacheRead = usage.inputTokens.cacheRead ?? 0;
+    const cacheWrite = usage.inputTokens.cacheWrite ?? 0;
+    const input = usage.inputTokens.noCache ??
+        Math.max(0, (usage.inputTokens.total ?? 0) - cacheRead - cacheWrite);
+    return {
+        input,
+        output: usage.outputTokens.total ?? 0,
+        cacheRead,
+        cacheWrite,
+    };
+}
+function nextOrdinal(map, key) {
+    if (!map)
+        return 0;
+    const n = map.get(key) ?? 0;
+    map.set(key, n + 1);
+    return n;
+}
+const VOLATILE_BUDGET_PATTERNS = [
+    /(budget: )≈\$-?\d[\d,]*(?:\.\d+)?/g,
+    /(Remaining research budget: )≈\$-?\d[\d,]*(?:\.\d+)?/g,
+];
+export function normalizeForCacheKey(serialized) {
+    let out = serialized;
+    for (const pattern of VOLATILE_BUDGET_PATTERNS) {
+        out = out.replace(pattern, (_match, label) => `${label}≈$`);
+    }
+    return out;
+}
+function callKey(model, params, role) {
+    const material = {
+        role,
+        provider: model.provider,
+        modelId: model.modelId,
+        prompt: params.prompt,
+        tools: params.tools?.map((tool) => ({
+            type: tool.type,
+            name: "name" in tool ? tool.name : undefined,
+            description: "description" in tool ? tool.description : undefined,
+            inputSchema: "inputSchema" in tool ? tool.inputSchema : undefined,
+        })),
+        toolChoice: params.toolChoice,
+        responseFormat: params.responseFormat,
+        maxOutputTokens: params.maxOutputTokens,
+    };
+    return createHash("sha256")
+        .update(normalizeForCacheKey(JSON.stringify(material)))
+        .digest("hex")
+        .slice(0, 40);
+}
+export function createModelCallCache() {
+    return new Map();
+}
+function serializeResult(result) {
+    const record = {
+        content: result.content,
+        finishReason: result.finishReason,
+        usage: result.usage,
+        ...(result.providerMetadata
+            ? { providerMetadata: result.providerMetadata }
+            : {}),
+    };
+    try {
+        JSON.stringify(record);
+        return record;
+    }
+    catch {
+        return null;
+    }
+}
+function streamFromJournaledCall(record) {
+    const parts = [
+        { type: "stream-start", warnings: [] },
+    ];
+    let nextId = 0;
+    for (const item of record.content) {
+        if (item.type === "text") {
+            const id = `replay_${nextId++}`;
+            parts.push({ type: "text-start", id }, { type: "text-delta", id, delta: item.text }, { type: "text-end", id });
+        }
+        else if (item.type === "reasoning") {
+            const id = `replay_${nextId++}`;
+            parts.push({ type: "reasoning-start", id }, { type: "reasoning-delta", id, delta: item.text }, { type: "reasoning-end", id });
+        }
+        else {
+            parts.push(item);
+        }
+    }
+    parts.push({
+        type: "finish",
+        usage: record.usage,
+        finishReason: record.finishReason,
+        ...(record.providerMetadata
+            ? { providerMetadata: record.providerMetadata }
+            : {}),
+    });
+    return new ReadableStream({
+        start(controller) {
+            for (const part of parts)
+                controller.enqueue(part);
+            controller.close();
+        },
+    });
+}
+function collectStreamPart(state, part) {
+    switch (part.type) {
+        case "text-start":
+            state.openText.set(part.id, state.content.push({ type: "text", text: "" }) - 1);
+            break;
+        case "text-delta": {
+            const index = state.openText.get(part.id);
+            const entry = index === undefined ? undefined : state.content[index];
+            if (entry && entry.type === "text")
+                entry.text += part.delta;
+            break;
+        }
+        case "reasoning-start":
+            state.openReasoning.set(part.id, state.content.push({ type: "reasoning", text: "" }) - 1);
+            break;
+        case "reasoning-delta": {
+            const index = state.openReasoning.get(part.id);
+            const entry = index === undefined ? undefined : state.content[index];
+            if (entry && entry.type === "reasoning")
+                entry.text += part.delta;
+            break;
+        }
+        case "tool-call":
+        case "tool-result":
+        case "file":
+        case "source":
+            state.content.push(part);
+            break;
+        case "finish":
+            state.finish = {
+                usage: part.usage,
+                finishReason: part.finishReason,
+                ...(part.providerMetadata
+                    ? { providerMetadata: part.providerMetadata }
+                    : {}),
+            };
+            break;
+        default:
+            break;
+    }
+}
+function isAnthropicModel(model) {
+    return model.provider.toLowerCase().includes("anthropic");
+}
+function isOpenAIModel(model) {
+    return model.provider.toLowerCase().includes("openai");
+}
+function withOpenAIDefaults(params) {
+    const openai = {
+        ...params.providerOptions?.openai,
+        strictJsonSchema: false,
+    };
+    return {
+        ...params,
+        providerOptions: { ...params.providerOptions, openai },
+    };
+}
+function withCacheBreakpoint(params) {
+    if (params.prompt.length === 0)
+        return params;
+    const prompt = [...params.prompt];
+    const last = prompt[prompt.length - 1];
+    prompt[prompt.length - 1] = {
+        ...last,
+        providerOptions: {
+            ...last.providerOptions,
+            anthropic: {
+                ...last.providerOptions?.anthropic,
+                cacheControl: { type: "ephemeral" },
+            },
+        },
+    };
+    return { ...params, prompt };
+}
+const RETRY_MAX_ATTEMPTS = 6; // 1 initial try + up to 5 retries
+const RETRY_BASE_DELAY_MS = 1_000;
+const RETRY_MAX_DELAY_MS = 30_000; // cap so a sustained limit can't hang minutes per attempt
+function isAbortError(err) {
+    return (err instanceof Error &&
+        (err.name === "AbortError" || err.name === "TimeoutError"));
+}
+function parseRetryAfterMs(headers) {
+    const raw = headers?.["retry-after"] ?? headers?.["Retry-After"];
+    if (!raw)
+        return undefined;
+    const seconds = Number(raw);
+    if (Number.isFinite(seconds))
+        return Math.max(0, seconds * 1_000);
+    const at = Date.parse(raw);
+    if (Number.isFinite(at))
+        return Math.max(0, at - Date.now());
+    return undefined;
+}
+function unwrapErrors(err) {
+    const seen = new Set();
+    const queue = [err];
+    const out = [];
+    while (queue.length > 0 && out.length < 8) {
+        const current = queue.shift();
+        if (!current || typeof current !== "object" || seen.has(current))
+            continue;
+        seen.add(current);
+        const shape = current;
+        out.push(shape);
+        if (shape.cause)
+            queue.push(shape.cause);
+        if (shape.lastError)
+            queue.push(shape.lastError);
+        if (Array.isArray(shape.errors))
+            queue.push(...shape.errors.slice(0, 4));
+    }
+    return out;
+}
+function classifyRetry(err) {
+    if (isAbortError(err))
+        return { retryable: false };
+    const shapes = unwrapErrors(err);
+    const retryAfterMs = shapes
+        .map((shape) => parseRetryAfterMs(shape.responseHeaders))
+        .find((ms) => ms !== undefined);
+    const message = shapes
+        .map((shape) => shape.message ?? "")
+        .join("\n")
+        .toLowerCase();
+    const rateLimit = shapes.some((shape) => shape.statusCode === 429) ||
+        /rate limit|too many requests|overloaded|concurrent connections/.test(message);
+    const withDelay = () => ({
+        retryable: true,
+        ...(retryAfterMs !== undefined ? { retryAfterMs } : {}),
+        ...(rateLimit ? { rateLimit: true } : {}),
+    });
+    for (const shape of shapes) {
+        if (typeof shape.isRetryable === "boolean") {
+            return shape.isRetryable ? withDelay() : { retryable: false };
+        }
+    }
+    const retryableStatus = shapes.some((shape) => {
+        const status = typeof shape.statusCode === "number" ? shape.statusCode : undefined;
+        return (status === 408 ||
+            status === 409 ||
+            status === 429 ||
+            (status !== undefined && status >= 500));
+    });
+    const retryableMessage = /rate limit|too many requests|overloaded|concurrent connections|timeout|timed out|econnreset|etimedout|eai_again|socket hang up|fetch failed|network error/.test(message);
+    return retryableStatus || retryableMessage
+        ? withDelay()
+        : { retryable: false };
+}
+function backoffDelayMs(attempt, retryAfterMs) {
+    const exponential = Math.min(RETRY_MAX_DELAY_MS, RETRY_BASE_DELAY_MS * 2 ** (attempt - 1));
+    const jittered = exponential / 2 + Math.random() * (exponential / 2);
+    return Math.min(RETRY_MAX_DELAY_MS, Math.max(retryAfterMs ?? 0, jittered));
+}
+async function callWithRetry(op, gate, signal, onRateLimit) {
+    let tries = 0;
+    let release = null;
+    try {
+        for (;;) {
+            tries++;
+            if (gate && !release)
+                release = await gate.acquire();
+            try {
+                return await op();
+            }
+            catch (err) {
+                const { retryable, retryAfterMs, rateLimit } = classifyRetry(err);
+                if (!retryable || tries >= RETRY_MAX_ATTEMPTS || signal?.aborted)
+                    throw err;
+                const delayMs = backoffDelayMs(tries, retryAfterMs);
+                onRateLimit?.({ attempt: tries, delayMs, error: err });
+                if (release && !rateLimit) {
+                    release();
+                    release = null;
+                }
+                await sleep(delayMs, signal);
+            }
+        }
+    }
+    finally {
+        release?.();
+    }
+}
+export function engineModel(model, hooks) {
+    const inner = model;
+    const recorder = hooks.recorder;
+    const settle = (usage, hold) => {
+        const tokens = tokenUsageFromV3(usage);
+        const { pricing, known } = resolvePricing(inner.modelId, hooks.pricing);
+        if (!known)
+            hooks.onUnknownModel?.(inner.modelId);
+        const cost = usageCostUSD(tokens, pricing);
+        if (hold)
+            hold.settle(cost);
+        else
+            hooks.grant.charge(cost);
+        trackUsage(hooks.usage, hooks.role, tokens);
+        hooks.onCost?.(cost);
+        return cost;
+    };
+    const settleReplay = (usage) => {
+        const tokens = tokenUsageFromV3(usage);
+        const { pricing } = resolvePricing(inner.modelId, hooks.pricing);
+        const cost = usageCostUSD(tokens, pricing);
+        hooks.grant.charge(cost);
+        trackUsage(hooks.usage, hooks.role, tokens);
+        hooks.usage.replayedUSD += cost;
+        hooks.onCost?.(cost);
+        return cost;
+    };
+    const recordReplayStep = (key, params, cached) => {
+        if (!recorder)
+            return;
+        const at = recorder.now();
+        recorder.recordModelCall({
+            callKey: key,
+            role: hooks.role,
+            provider: inner.provider,
+            modelId: inner.modelId,
+            t0: at,
+            t1: at,
+            waitMs: 0,
+            computeMs: 0,
+            tokens: tokenUsageFromV3(cached.usage),
+            finishReason: cached.finishReason.unified,
+            status: "replayed",
+            replayed: true,
+            params,
+            content: cached.content,
+        }, currentFrame());
+    };
+    const recordCall = (key, params, timing, outcome) => {
+        if (!recorder)
+            return;
+        const tEnd = recorder.now();
+        const waitMs = Math.max(0, timing.firstWork - timing.tEnqueue);
+        const retryDelayMs = Math.max(0, tEnd - timing.firstWork - timing.computeMs);
+        recorder.recordModelCall({
+            callKey: key,
+            role: hooks.role,
+            provider: inner.provider,
+            modelId: inner.modelId,
+            t0: timing.tEnqueue,
+            t1: tEnd,
+            waitMs,
+            computeMs: timing.computeMs,
+            ...(retryDelayMs ? { retryDelayMs } : {}),
+            ...(outcome.status === "ok"
+                ? {
+                    tokens: tokenUsageFromV3(outcome.usage),
+                    costUSD: outcome.costUSD,
+                    finishReason: outcome.finishReason,
+                    status: "ok",
+                    ...(outcome.content ? { content: outcome.content } : {}),
+                }
+                : { status: outcome.status, error: outcome.error }),
+            params,
+        }, currentFrame());
+    };
+    const reserveFor = (params) => {
+        const { pricing } = resolvePricing(inner.modelId, hooks.pricing);
+        let promptChars = 0;
+        try {
+            promptChars = JSON.stringify(params.prompt).length;
+        }
+        catch {
+            promptChars = 0;
+        }
+        const estimateUSD = usageCostUSD({
+            input: promptChars / ECONOMY.callReserve.promptCharsPerToken,
+            output: params.maxOutputTokens ?? ECONOMY.callReserve.assumedOutputTokens,
+            cacheRead: 0,
+            cacheWrite: 0,
+        }, pricing);
+        const hold = hooks.grant.reserve(estimateUSD);
+        return hold ? { hold, estimateUSD } : null;
+    };
+    const settleFailure = (reservation, err, abortSignal) => {
+        if (!reservation)
+            return;
+        if (isAbortError(err) || abortSignal?.aborted) {
+            reservation.hold.settle(reservation.estimateUSD);
+        }
+        else {
+            reservation.hold.release();
+        }
+    };
+    const middleware = {
+        specificationVersion: "v3",
+        transformParams: async ({ params }) => {
+            if (isAnthropicModel(inner))
+                return withCacheBreakpoint(params);
+            if (isOpenAIModel(inner))
+                return withOpenAIDefaults(params);
+            return params;
+        },
+        wrapGenerate: async ({ doGenerate, params }) => {
+            const key = callKey(inner, params, hooks.role);
+            const journalKey = `${key}#${nextOrdinal(hooks.callOrdinals, key)}`;
+            const cached = hooks.replay?.take(journalKey);
+            if (cached) {
+                settleReplay(cached.usage);
+                recordReplayStep(key, params, cached);
+                hooks.modelCache?.set(key, cached);
+                return {
+                    content: cached.content,
+                    finishReason: cached.finishReason,
+                    usage: cached.usage,
+                    ...(cached.providerMetadata
+                        ? { providerMetadata: cached.providerMetadata }
+                        : {}),
+                    warnings: [],
+                };
+            }
+            const reused = hooks.modelCache?.get(key);
+            if (reused) {
+                recordReplayStep(key, params, reused);
+                hooks.onCacheHit?.();
+                return {
+                    content: reused.content,
+                    finishReason: reused.finishReason,
+                    usage: reused.usage,
+                    ...(reused.providerMetadata
+                        ? { providerMetadata: reused.providerMetadata }
+                        : {}),
+                    warnings: [],
+                };
+            }
+            if (hooks.budgetExhausted?.()) {
+                throw new AtlasError("budget exhausted: total run budget reached", "budget");
+            }
+            const reservation = reserveFor(params);
+            const timing = {
+                tEnqueue: recorder ? recorder.now() : 0,
+                firstWork: 0,
+                computeMs: 0,
+            };
+            let acquired = false;
+            let result;
+            try {
+                result = await callWithRetry(() => {
+                    if (!recorder)
+                        return Promise.resolve(doGenerate());
+                    const s = recorder.now();
+                    if (!acquired) {
+                        timing.firstWork = s;
+                        acquired = true;
+                    }
+                    return Promise.resolve(doGenerate()).then((r) => {
+                        timing.computeMs += recorder.now() - s;
+                        return r;
+                    }, (e) => {
+                        timing.computeMs += recorder.now() - s;
+                        throw e;
+                    });
+                }, hooks.gate, params.abortSignal, hooks.onRateLimit);
+            }
+            catch (err) {
+                settleFailure(reservation, err, params.abortSignal);
+                recordCall(key, params, timing, {
+                    status: isAbortError(err) || params.abortSignal?.aborted
+                        ? "aborted"
+                        : "error",
+                    error: errorMessage(err),
+                });
+                throw err;
+            }
+            const costUSD = settle(result.usage, reservation?.hold ?? null);
+            if (hooks.journal || hooks.modelCache) {
+                const record = serializeResult(result);
+                if (record) {
+                    hooks.journal?.call(journalKey, record);
+                    hooks.modelCache?.set(key, record);
+                }
+            }
+            recordCall(key, params, timing, {
+                status: "ok",
+                usage: result.usage,
+                finishReason: result.finishReason.unified,
+                costUSD,
+                content: result.content,
+            });
+            return result;
+        },
+        wrapStream: async ({ doStream, params }) => {
+            const key = callKey(inner, params, hooks.role);
+            const journalKey = `${key}#${nextOrdinal(hooks.callOrdinals, key)}`;
+            const cached = hooks.replay?.take(journalKey);
+            if (cached) {
+                settleReplay(cached.usage);
+                recordReplayStep(key, params, cached);
+                hooks.modelCache?.set(key, cached);
+                return { stream: streamFromJournaledCall(cached) };
+            }
+            if (hooks.budgetExhausted?.()) {
+                throw new AtlasError("budget exhausted: total run budget reached", "budget");
+            }
+            const reservation = reserveFor(params);
+            const tEnqueue = recorder ? recorder.now() : 0;
+            let firstWork = tEnqueue;
+            const releaseSlot = await hooks.gate.acquire();
+            let result;
+            try {
+                result = await callWithRetry(() => {
+                    if (recorder)
+                        firstWork = recorder.now();
+                    return Promise.resolve(doStream());
+                }, undefined, params.abortSignal, hooks.onRateLimit);
+            }
+            catch (err) {
+                releaseSlot();
+                settleFailure(reservation, err, params.abortSignal);
+                if (recorder) {
+                    recordCall(key, params, {
+                        tEnqueue,
+                        firstWork,
+                        computeMs: Math.max(0, recorder.now() - firstWork),
+                    }, {
+                        status: isAbortError(err) || params.abortSignal?.aborted
+                            ? "aborted"
+                            : "error",
+                        error: errorMessage(err),
+                    });
+                }
+                throw err;
+            }
+            const state = {
+                content: [],
+                openText: new Map(),
+                openReasoning: new Map(),
+            };
+            let lastCost;
+            let recorded = false;
+            const recordStream = (status) => {
+                if (!recorder || recorded)
+                    return;
+                recorded = true;
+                const tEnd = recorder.now();
+                recorder.recordModelCall({
+                    callKey: key,
+                    role: hooks.role,
+                    provider: inner.provider,
+                    modelId: inner.modelId,
+                    t0: tEnqueue,
+                    t1: tEnd,
+                    waitMs: Math.max(0, firstWork - tEnqueue),
+                    computeMs: Math.max(0, tEnd - firstWork),
+                    ...(state.finish
+                        ? { tokens: tokenUsageFromV3(state.finish.usage) }
+                        : {}),
+                    ...(lastCost !== undefined ? { costUSD: lastCost } : {}),
+                    ...(state.finish?.finishReason
+                        ? { finishReason: state.finish.finishReason.unified }
+                        : {}),
+                    status,
+                    params,
+                    content: state.content,
+                }, currentFrame());
+            };
+            const metered = result.stream.pipeThrough(new TransformStream({
+                transform(part, controller) {
+                    collectStreamPart(state, part);
+                    if (part.type === "finish") {
+                        lastCost = settle(part.usage, reservation?.hold ?? null);
+                        releaseSlot();
+                    }
+                    controller.enqueue(part);
+                },
+                cancel() {
+                    if (reservation && !state.finish) {
+                        reservation.hold.settle(reservation.estimateUSD);
+                    }
+                    reservation?.hold.release();
+                    releaseSlot();
+                    recordStream("aborted");
+                },
+                flush() {
+                    if (reservation && !state.finish) {
+                        reservation.hold.settle(reservation.estimateUSD);
+                    }
+                    reservation?.hold.release();
+                    releaseSlot();
+                    if (state.finish && hooks.journal) {
+                        const record = serializeResult({
+                            content: state.content,
+                            finishReason: state.finish.finishReason,
+                            usage: state.finish.usage,
+                            ...(state.finish.providerMetadata
+                                ? { providerMetadata: state.finish.providerMetadata }
+                                : {}),
+                            warnings: [],
+                        });
+                        if (record)
+                            hooks.journal.call(journalKey, record);
+                    }
+                    recordStream(state.finish ? "ok" : "aborted");
+                },
+            }));
+            return { ...result, stream: metered };
+        },
+    };
+    return wrapLanguageModel({ model: inner, middleware });
+}
+export const MODEL_CALL_MAX_RETRIES = 0;

package/dist/orchestrate.d.ts

@@ -0,0 +1,5 @@
+import type { Researcher } from "./researcher.js";
+import { type SpineOutput } from "./spine.js";
+import type { RunCtx } from "./state.js";
+export declare const ATLAS_KEY = "atlas";
+export declare function runOrchestrated(rctx: RunCtx, researchers: Record<string, Researcher>): Promise<SpineOutput>;