diff --git a/mcp/src/aieo/src/provider.ts b/mcp/src/aieo/src/provider.ts index 72c4aa533..3f38c60e8 100644 --- a/mcp/src/aieo/src/provider.ts +++ b/mcp/src/aieo/src/provider.ts @@ -197,7 +197,7 @@ export interface TokenUsageForCost { export function computeSessionCost(provider: Provider, usage: TokenUsageForCost, modelId?: string, actualCost?: number): number { if (typeof actualCost === "number" && actualCost >= 0) return actualCost; - const pricing = (modelId ? getModelPricing(modelId) : undefined) ?? TOKEN_PRICING[provider]; + const pricing = (modelId ? getModelPricing(modelId, provider) : undefined) ?? TOKEN_PRICING[provider]; if (!pricing) return 0; const inputCost = (usage.input / 1_000_000) * pricing.inputTokenPrice; const cacheReadCost = (usage.cache_read / 1_000_000) * (pricing.cacheReadPrice ?? pricing.inputTokenPrice); @@ -546,38 +546,98 @@ export function getTokenPricing(provider: Provider): TokenPricing { return TOKEN_PRICING[provider]; } +// Per-model pricing (per 1M tokens), sourced live from the OpenRouter catalog. const modelPricing: Record = {}; -export async function loadModelPricing(): Promise { - try { - const res = await fetch("https://openrouter.ai/api/v1/models"); - const { data } = (await res.json()) as { data: any[] }; - for (const m of data) { - const p = m.pricing || {}; - const inputTokenPrice = parseFloat(p.prompt) * 1_000_000; - const outputTokenPrice = parseFloat(p.completion) * 1_000_000; - if (!inputTokenPrice && !outputTokenPrice) continue; - modelPricing[m.id] = { - inputTokenPrice, - outputTokenPrice, - ...(parseFloat(p.input_cache_read) > 0 && { - cacheReadPrice: parseFloat(p.input_cache_read) * 1_000_000, - }), - ...(parseFloat(p.input_cache_write) > 0 && { - cacheWritePrice: parseFloat(p.input_cache_write) * 1_000_000, - }), - }; +const PRICING_URL = "https://openrouter.ai/api/v1/models"; +const PRICING_REFRESH_MS = 6 * 60 * 60 * 1000; // 6h +let pricingTimer: ReturnType | undefined; + +// Native provider model ids are bare (e.g. "claude-opus-4-6"); the OpenRouter +// catalog keys them vendor-namespaced with dotted versions +// (e.g. "anthropic/claude-opus-4.6"). Map a bare id to the catalog scheme so +// native Anthropic/OpenAI/Google sessions get per-model pricing instead of the +// flat provider default. Returns the id unchanged when it is already +// vendor/model form. Unknown models simply miss the table and fall back. +const PROVIDER_VENDOR: Partial> = { + anthropic: "anthropic", + openai: "openai", + google: "google", +}; + +function toCatalogId(modelId: string, provider?: Provider): string { + let key = modelId.startsWith("openrouter/") + ? modelId.slice("openrouter/".length) + : modelId; + if (key.includes("/")) return key; // already vendor/model (OpenRouter form) + // Normalize a trailing version separator: "opus-4-6" -> "opus-4.6". + key = key.replace(/-(\d+)-(\d+)$/, "-$1.$2"); + const vendor = provider ? PROVIDER_VENDOR[provider] : undefined; + return vendor ? `${vendor}/${key}` : key; +} + +async function fetchModelPricing(): Promise { + const res = await fetch(PRICING_URL); + if (!res.ok) throw new Error(`OpenRouter models fetch failed: ${res.status}`); + const { data } = (await res.json()) as { data: any[] }; + const next: Record = {}; + for (const m of data) { + const p = m.pricing || {}; + const inputTokenPrice = parseFloat(p.prompt) * 1_000_000; + const outputTokenPrice = parseFloat(p.completion) * 1_000_000; + const input = Number.isFinite(inputTokenPrice) ? inputTokenPrice : 0; + const output = Number.isFinite(outputTokenPrice) ? outputTokenPrice : 0; + if (!input && !output) continue; + const cacheRead = parseFloat(p.input_cache_read) * 1_000_000; + const cacheWrite = parseFloat(p.input_cache_write) * 1_000_000; + next[m.id] = { + inputTokenPrice: input, + outputTokenPrice: output, + ...(Number.isFinite(cacheRead) && cacheRead > 0 && { cacheReadPrice: cacheRead }), + ...(Number.isFinite(cacheWrite) && cacheWrite > 0 && { cacheWritePrice: cacheWrite }), + }; + } + // Merge in place (built fully before assign, so readers never see a partial + // table); keeps prior prices if a later refresh returns fewer models. + Object.assign(modelPricing, next); +} + +export async function loadModelPricing(opts?: { + retries?: number; + refreshMs?: number; +}): Promise { + const retries = opts?.retries ?? 3; + for (let attempt = 1; attempt <= retries; attempt++) { + try { + await fetchModelPricing(); + console.log(`[pricing] loaded ${Object.keys(modelPricing).length} model prices`); + break; + } catch (e) { + if (attempt === retries) { + console.warn( + `[pricing] failed after ${retries} attempts; using provider defaults`, + e, + ); + } else { + await new Promise((r) => setTimeout(r, attempt * 1000)); + } } - } catch { - console.warn("Failed to load model pricing; using provider defaults"); + } + // Refresh periodically so long-running processes don't drift on stale prices. + const refreshMs = opts?.refreshMs ?? PRICING_REFRESH_MS; + if (!pricingTimer && refreshMs > 0) { + pricingTimer = setInterval(() => { + fetchModelPricing().catch(() => {}); + }, refreshMs); + if (typeof pricingTimer.unref === "function") pricingTimer.unref(); } } -function getModelPricing(modelId: string): TokenPricing | undefined { - const key = modelId.startsWith("openrouter/") - ? modelId.slice("openrouter/".length) - : modelId; - return modelPricing[key]; +function getModelPricing( + modelId: string, + provider?: Provider, +): TokenPricing | undefined { + return modelPricing[toCatalogId(modelId, provider)]; } export type ThinkingSpeed = "thinking" | "fast"; diff --git a/mcp/src/repo/descriptions.ts b/mcp/src/repo/descriptions.ts index 96ab0da2a..de2c08aab 100644 --- a/mcp/src/repo/descriptions.ts +++ b/mcp/src/repo/descriptions.ts @@ -199,6 +199,11 @@ ${content.slice(0, 2000)}`; const usage = normalizeUsage(rawUsage); const actualCost = (providerMetadata as any)?.openrouter?.usage?.cost; const cost = computeSessionCost(llm.provider, usage, llm.modelName, actualCost); + // Accumulate immediately so the in-flight cost guard above is live + // and concurrent overshoot is bounded to ~concurrency, not a full + // batch. Safe: no await between read and write of totalCost. + totalCost += cost; + totalUsage = addUsage(totalUsage, usage); console.log( `[describe_nodes] LLM done: ${name} ($${cost.toFixed(6)})`, ); @@ -215,10 +220,16 @@ ${content.slice(0, 2000)}`; }), ); - // Accumulate costs - for (const r of results) { - totalCost += r.cost; - totalUsage = addUsage(totalUsage, r.usage); + // Cost/usage already accumulated per-task above. + + // Guard against an infinite loop: if every LLM call in this batch failed, + // no cost accrues and the same undescribed nodes are re-fetched forever. + // Stop making zero progress rather than spin. + if (results.length === 0) { + console.warn( + `[describe_nodes] Zero progress on batch of ${nodes.length} nodes (all calls failed); aborting.`, + ); + break; } // Bulk write to Neo4j diff --git a/mcp/src/repo/tools.ts b/mcp/src/repo/tools.ts index 4dda268ee..4a5af5293 100644 --- a/mcp/src/repo/tools.ts +++ b/mcp/src/repo/tools.ts @@ -317,6 +317,11 @@ export async function get_tools( // For single repo, extract owner/name from path. For multi-repo, these will be empty. const repoOwner = isMultiRepo ? "" : repoArr[repoArr.length - 2]; const repoName = isMultiRepo ? "" : repoArr[repoArr.length - 1]; + // Default content-search scope: constrain graph search to the session's repos + // (owner/repo form) so results don't leak across repos on a multi-repo graph. + // Returns undefined when no repos are known (single-repo deployments / no scope). + const repoScopePatterns = (rs?: string[]): string[] | undefined => + rs && rs.length > 0 ? rs.map((r) => `${r}/**`) : undefined; // Resolve a repo owner/name pair from an optional "owner/name" string, // falling back to the single-repo values or the first entry in repos. function resolveRepo(repo?: string): { owner: string; name: string } | null { @@ -624,6 +629,10 @@ export async function get_tools( query, limit || 10, codeNodeTypes as any, + [], + 0, + undefined, + repoScopePatterns(repos), ); return results.map((node) => ({ name: node.properties.name, @@ -667,7 +676,10 @@ export async function get_tools( [], args.language, "relevance", - args.include_patterns, + // Default to the session's repo scope so search doesn't leak across + // repos on a multi-repo graph; the model may still override with a + // narrower pattern. + args.include_patterns ?? repoScopePatterns(repos), args.exclude_patterns, ); if (provenanceCollector) { @@ -695,7 +707,14 @@ export async function get_tools( description: stak.GetMapTool.description || defaultDescriptions.stakgraph_map, inputSchema: stak.GetMapSchema, execute: async (args: z.infer) => { - const result = await stak.getMap(args); + // When exactly one repo is in scope and the model resolved the node by + // name (no ref_id/file), constrain node resolution to that repo so a + // same-named node in another repo isn't picked. + const scoped = + !args.ref_id && !args.file && repos && repos.length === 1 + ? { ...args, file: repos[0] } + : args; + const result = await stak.getMap(scoped); return result.content?.[0]?.text ?? ""; }, });