diff --git a/src/features/workspaces/ai/ai-inspector-view-model.test.ts b/src/features/workspaces/ai/ai-inspector-view-model.test.ts new file mode 100644 index 00000000..1e0c5bf1 --- /dev/null +++ b/src/features/workspaces/ai/ai-inspector-view-model.test.ts @@ -0,0 +1,28 @@ +import { describe, expect, it } from "vitest"; + +import type { AIInspectorEvent } from "#/features/workspaces/ai/ai-inspector"; +import { getAIInspectorRunViews } from "#/features/workspaces/ai/ai-inspector-view-model"; + +describe("AI inspector run views", () => { + it("surfaces the requested and routed model for auto turns", () => { + const [run] = getAIInspectorRunViews([ + { + id: "event-1", + runId: "run-1", + sequence: 1, + createdAt: 1, + type: "turn.started", + payload: { + modelId: "claude-haiku", + requestedModelId: "auto", + routingReason: "simple-chat", + system: "system prompt", + }, + } satisfies AIInspectorEvent, + ]); + + expect(run.modelId).toBe("claude-haiku"); + expect(run.requestedModelId).toBe("auto"); + expect(run.routingReason).toBe("simple-chat"); + }); +}); diff --git a/src/features/workspaces/ai/ai-inspector-view-model.ts b/src/features/workspaces/ai/ai-inspector-view-model.ts index 0363463d..0d4adc6e 100644 --- a/src/features/workspaces/ai/ai-inspector-view-model.ts +++ b/src/features/workspaces/ai/ai-inspector-view-model.ts @@ -60,6 +60,8 @@ function buildRunView(runId: string, events: AIInspectorEvent[]) { case "turn.started": run.startedAt = event.createdAt; run.modelId = getString(payload.modelId); + run.requestedModelId = getString(payload.requestedModelId); + run.routingReason = getString(payload.routingReason); run.system = getString(payload.system); run.thread = payload.thread; run.body = payload.body; diff --git a/src/features/workspaces/ai/ai-inspector-view-types.ts b/src/features/workspaces/ai/ai-inspector-view-types.ts index 5d9716a6..e1980acc 100644 --- a/src/features/workspaces/ai/ai-inspector-view-types.ts +++ b/src/features/workspaces/ai/ai-inspector-view-types.ts @@ -64,6 +64,8 @@ export interface AIInspectorRunView { finishedAt?: number; status: "running" | "completed" | "failed"; modelId?: string; + requestedModelId?: string; + routingReason?: string; system?: string; thread?: unknown; body?: unknown; diff --git a/src/features/workspaces/ai/ai-thread-inspector-recorder.ts b/src/features/workspaces/ai/ai-thread-inspector-recorder.ts index 8dc9cab7..64843a2f 100644 --- a/src/features/workspaces/ai/ai-thread-inspector-recorder.ts +++ b/src/features/workspaces/ai/ai-thread-inspector-recorder.ts @@ -45,6 +45,8 @@ export class AIThreadInspectorRecorder { async recordTurnStarted(input: { ctx: TurnContext; modelId: string; + requestedModelId: string; + routingReason?: string; system: string; thread: AIThreadContext; tools: unknown; @@ -56,6 +58,8 @@ export class AIThreadInspectorRecorder { continuation: input.ctx.continuation, messages: summarizeInspectorMessages(input.ctx.messages), modelId: input.modelId, + requestedModelId: input.requestedModelId, + routingReason: input.routingReason, system: input.system, thread: input.thread, tools: await summarizeInspectorTools(input.tools), diff --git a/src/features/workspaces/ai/ai-thread-posthog-recorder.ts b/src/features/workspaces/ai/ai-thread-posthog-recorder.ts index d30f3172..65327bf1 100644 --- a/src/features/workspaces/ai/ai-thread-posthog-recorder.ts +++ b/src/features/workspaces/ai/ai-thread-posthog-recorder.ts @@ -138,6 +138,8 @@ export class AIThreadPostHogRecorder { recordTurnStarted(input: { ctx: TurnContext; modelId: WorkspaceAiChatModelId; + requestedModelId: WorkspaceAiChatModelId; + routingReason?: string; thread: AIThreadContext; tools?: unknown; }) { @@ -166,6 +168,8 @@ export class AIThreadPostHogRecorder { properties: { ...turnTelemetryProperties(turn), model_id: input.modelId, + requested_model_id: input.requestedModelId, + routing_reason: input.routingReason ?? null, continuation: Boolean(input.ctx.continuation), }, ...this.serverEventRuntime, diff --git a/src/features/workspaces/ai/ai-thread-tcc-recorder.ts b/src/features/workspaces/ai/ai-thread-tcc-recorder.ts index 28371d3f..1e4206f2 100644 --- a/src/features/workspaces/ai/ai-thread-tcc-recorder.ts +++ b/src/features/workspaces/ai/ai-thread-tcc-recorder.ts @@ -78,6 +78,8 @@ export class AIThreadTccRecorder { ctx: TurnContext; env: Cloudflare.Env; modelId: WorkspaceAiChatModelId; + requestedModelId: WorkspaceAiChatModelId; + routingReason?: string; system: string; thread: AIThreadContext; tools?: unknown; @@ -94,6 +96,8 @@ export class AIThreadTccRecorder { const metadata = createTccMetadata({ gatewayModel, modelId: input.modelId, + requestedModelId: input.requestedModelId, + routingReason: input.routingReason, thread: input.thread, }); @@ -454,6 +458,8 @@ function getTccApiKey(env: Cloudflare.Env) { function createTccMetadata(input: { gatewayModel: string; modelId: WorkspaceAiChatModelId; + requestedModelId: WorkspaceAiChatModelId; + routingReason?: string; thread: AIThreadContext; }) { return { @@ -466,6 +472,8 @@ function createTccMetadata(input: { gateway_model: input.gatewayModel, model_id: input.modelId, mutation_mode: input.thread.promptScope.canMutate ? "mutate" : "view", + requested_model_id: input.requestedModelId, + ...(input.routingReason ? { routing_reason: input.routingReason } : {}), workspace_id: input.thread.workspaceId, } satisfies Record; } diff --git a/src/features/workspaces/ai/ai-thread-telemetry-recorder.ts b/src/features/workspaces/ai/ai-thread-telemetry-recorder.ts index 0aa27f96..43f8075f 100644 --- a/src/features/workspaces/ai/ai-thread-telemetry-recorder.ts +++ b/src/features/workspaces/ai/ai-thread-telemetry-recorder.ts @@ -40,7 +40,11 @@ export class AIThreadTelemetryRecorder { async recordTurnStarted(input: { ctx: TurnContext; + /** Model the turn actually runs on (the routed model for Auto turns). */ modelId: WorkspaceAiChatModelId; + /** Model the user picked; "auto" when the router chose modelId. */ + requestedModelId: WorkspaceAiChatModelId; + routingReason?: string; system: string; thread: AIThreadContext; tools: unknown; diff --git a/src/features/workspaces/ai/ai-thread.ts b/src/features/workspaces/ai/ai-thread.ts index 2995e1a5..0500577c 100644 --- a/src/features/workspaces/ai/ai-thread.ts +++ b/src/features/workspaces/ai/ai-thread.ts @@ -31,6 +31,10 @@ import { getWorkspaceAiGatewayProviderOptions, getWorkspaceAiLanguageModel, } from "#/features/workspaces/ai/ai-thread-runtime"; +import { + extractWorkspaceAiRoutingSignals, + routeWorkspaceAiAutoModel, +} from "#/features/workspaces/ai/model-router"; import { DEFAULT_WORKSPACE_AI_CHAT_MODEL_ID, getWorkspaceAiChatModel, @@ -69,7 +73,11 @@ type AIThreadRunSettlement = }; interface AIThreadUsageContext { + /** Model actually run and metered; equals requestedModelId unless Auto routed. */ modelId: WorkspaceAiChatModelId; + /** Model the user picked; "auto" when the router chose modelId. */ + requestedModelId: WorkspaceAiChatModelId; + routingReason?: string; thread: AIThreadContext; } @@ -173,9 +181,14 @@ export function createAIThreadClass(getUserAIStore: () => typeof UserAIStore) { void this.keepAliveWhile(() => this._maybeGenerateThreadTitle()); } - const modelId = resolveWorkspaceAiChatModelId(ctx.body?.modelId); + const requestedModelId = resolveWorkspaceAiChatModelId(ctx.body?.modelId); + const route = this._resolveAutoModelRoute(requestedModelId, ctx); + const modelId = route?.modelId ?? requestedModelId; if (!ctx.continuation) { + // Enforcement is currently disabled inside this check. When it lands + // it must gate on the routed model's billing tier (modelId here), not + // the requested picker id, so Auto turns cannot dodge premium gating. const access = await checkWorkspaceAiMessageAccess({ env: this.env, modelId, @@ -188,6 +201,8 @@ export function createAIThreadClass(getUserAIStore: () => typeof UserAIStore) { this.activeUsageContext = { modelId, + requestedModelId, + routingReason: route?.reason, thread, }; } @@ -213,6 +228,8 @@ export function createAIThreadClass(getUserAIStore: () => typeof UserAIStore) { await this.telemetry.recordTurnStarted({ ctx, modelId, + requestedModelId, + routingReason: route?.reason, system, thread, tools: activeTools, @@ -293,6 +310,39 @@ export function createAIThreadClass(getUserAIStore: () => typeof UserAIStore) { return this.telemetry.getInspectorSnapshot(this.name); } + // Auto resolves to a concrete standard-tier model at turn time; every + // other picker id passes through untouched. Continuation turns reuse the + // id routed when the run started so a run never switches models + // mid-flight; if that state is gone (e.g. recovery), we route again from + // the current turn's signals. + private _resolveAutoModelRoute( + requestedModelId: WorkspaceAiChatModelId, + ctx: TurnContext, + ): { modelId: WorkspaceAiChatModelId; reason: string } | undefined { + if (requestedModelId !== "auto") { + return undefined; + } + + const activeContext = ctx.continuation ? this.activeUsageContext : undefined; + + if (activeContext && activeContext.requestedModelId !== "auto") { + // The run started from an explicit picker choice; a continuation + // whose body lost the selection must not be re-routed. + return undefined; + } + + if (activeContext) { + return { + modelId: activeContext.modelId, + reason: activeContext.routingReason ?? "turn-continuation", + }; + } + + const routed = routeWorkspaceAiAutoModel(extractWorkspaceAiRoutingSignals(ctx)); + + return { modelId: routed.modelId, reason: routed.reason }; + } + private async _getThreadContext() { const directory = await this.parentAgent(getUserAIStore()); return directory.getThreadContext(this.name); diff --git a/src/features/workspaces/ai/model-router.test.ts b/src/features/workspaces/ai/model-router.test.ts new file mode 100644 index 00000000..656fd17a --- /dev/null +++ b/src/features/workspaces/ai/model-router.test.ts @@ -0,0 +1,279 @@ +import type { TurnContext } from "@cloudflare/think"; +import { describe, expect, it } from "vitest"; + +import { + extractWorkspaceAiRoutingSignals, + routeWorkspaceAiAutoModel, + WORKSPACE_AI_AUTO_ROUTING_THRESHOLDS, + type WorkspaceAiAutoRoutingRule, + type WorkspaceAiRoutingSignals, +} from "#/features/workspaces/ai/model-router"; +import type { WorkspaceAiChatModelId } from "#/features/workspaces/ai/models"; + +function buildTurnContext(overrides: Partial = {}): TurnContext { + return { + continuation: false, + messages: [], + model: "test-model", + system: "", + tools: {}, + ...overrides, + } as TurnContext; +} + +function buildSignals( + overrides: Partial = {}, +): WorkspaceAiRoutingSignals { + return { + attachmentCount: 0, + hasCodeSignals: false, + lastUserTextChars: 0, + messageCount: 0, + priorToolCallCount: 0, + totalTextChars: 0, + workspaceReferenceCount: 0, + ...overrides, + }; +} + +describe("workspace AI auto model routing", () => { + it("routes short simple chat to the fast standard model", () => { + expect( + routeWorkspaceAiAutoModel( + buildSignals({ lastUserTextChars: 40, messageCount: 3, totalTextChars: 40 }), + ), + ).toEqual({ + gatewayModel: "anthropic/claude-haiku-4.5", + modelId: "claude-haiku", + reason: "simple-chat", + }); + }); + + it("routes attachment-heavy turns to the large-context model", () => { + expect( + routeWorkspaceAiAutoModel( + buildSignals({ attachmentCount: 2, lastUserTextChars: 40, messageCount: 3 }), + ), + ).toEqual({ + gatewayModel: "google/gemini-3-flash", + modelId: "gemini", + reason: "attachment-heavy", + }); + }); + + it("routes long-context turns to the large-context model", () => { + expect( + routeWorkspaceAiAutoModel( + buildSignals({ lastUserTextChars: 120, messageCount: 8, totalTextChars: 60_000 }), + ), + ).toEqual({ + gatewayModel: "google/gemini-3-flash", + modelId: "gemini", + reason: "long-context", + }); + + expect(routeWorkspaceAiAutoModel(buildSignals({ workspaceReferenceCount: 6 })).reason).toBe( + "long-context", + ); + }); + + it("routes tool-heavy and code turns to the tool-following model", () => { + expect( + routeWorkspaceAiAutoModel( + buildSignals({ hasCodeSignals: true, lastUserTextChars: 90, messageCount: 2 }), + ), + ).toEqual({ + gatewayModel: "openai/gpt-5.4-mini", + modelId: "chatgpt-mini", + reason: "tool-heavy", + }); + + expect( + routeWorkspaceAiAutoModel( + buildSignals({ lastUserTextChars: 90, messageCount: 9, priorToolCallCount: 5 }), + ).reason, + ).toBe("tool-heavy"); + }); + + it("falls back to the default model for high-reasoning prompts", () => { + expect( + routeWorkspaceAiAutoModel( + buildSignals({ lastUserTextChars: 2_000, messageCount: 4, totalTextChars: 2_400 }), + ), + ).toEqual({ + gatewayModel: "moonshotai/kimi-k2.6", + modelId: "auto", + reason: "high-reasoning-default", + }); + }); + + it("prefers capability rules over the simple-chat rule", () => { + const route = routeWorkspaceAiAutoModel( + buildSignals({ attachmentCount: 1, lastUserTextChars: 20, messageCount: 1 }), + ); + + expect(route.reason).toBe("attachment-heavy"); + }); + + it("treats threshold boundary values consistently", () => { + const thresholds = WORKSPACE_AI_AUTO_ROUTING_THRESHOLDS; + + expect( + routeWorkspaceAiAutoModel( + buildSignals({ + lastUserTextChars: thresholds.simpleChatMaxLastUserTextChars, + messageCount: thresholds.simpleChatMaxMessageCount, + }), + ).reason, + ).toBe("simple-chat"); + + expect( + routeWorkspaceAiAutoModel( + buildSignals({ + lastUserTextChars: thresholds.simpleChatMaxLastUserTextChars + 1, + messageCount: 1, + }), + ).reason, + ).toBe("high-reasoning-default"); + + expect( + routeWorkspaceAiAutoModel( + buildSignals({ + lastUserTextChars: 100, + totalTextChars: thresholds.longContextTotalTextChars, + }), + ).reason, + ).toBe("long-context"); + + expect( + routeWorkspaceAiAutoModel( + buildSignals({ + lastUserTextChars: 100, + messageCount: 2, + totalTextChars: thresholds.longContextTotalTextChars - 1, + }), + ).reason, + ).toBe("simple-chat"); + }); + + it("never routes into a premium billing tier", () => { + const premiumRule: WorkspaceAiAutoRoutingRule = { + matches: () => true, + reason: "tool-heavy", + targetModelId: "claude-sonnet", + }; + + expect(routeWorkspaceAiAutoModel(buildSignals(), [premiumRule])).toEqual({ + gatewayModel: "moonshotai/kimi-k2.6", + modelId: "auto", + reason: "fallback-unavailable", + }); + }); + + it("falls back safely when a routed model has been removed", () => { + const staleRule: WorkspaceAiAutoRoutingRule = { + matches: () => true, + reason: "simple-chat", + targetModelId: "removed-model" as WorkspaceAiChatModelId, + }; + + expect(routeWorkspaceAiAutoModel(buildSignals(), [staleRule])).toEqual({ + gatewayModel: "moonshotai/kimi-k2.6", + modelId: "auto", + reason: "fallback-unavailable", + }); + }); +}); + +describe("workspace AI routing signal extraction", () => { + it("extracts text, attachment, and tool-call signals from turn messages", () => { + const signals = extractWorkspaceAiRoutingSignals( + buildTurnContext({ + messages: [ + { role: "user", content: "Summarize this photo" }, + { + role: "assistant", + content: [ + { type: "text", text: "Sure." }, + { type: "tool-call", toolCallId: "call-1", toolName: "compute", input: {} }, + ], + }, + { + role: "user", + content: [ + { type: "text", text: "What about this one?" }, + { type: "image", image: "data:image/png;base64,AAAA" }, + { type: "file", data: "AAAA", mediaType: "image/png" }, + ], + }, + ], + } as Partial), + ); + + expect(signals).toEqual({ + attachmentCount: 2, + hasCodeSignals: false, + lastUserTextChars: "What about this one?".length, + messageCount: 3, + priorToolCallCount: 1, + totalTextChars: ("Summarize this photo" + "Sure." + "What about this one?").length, + workspaceReferenceCount: 0, + }); + }); + + it("detects code signals in the latest user message", () => { + const signals = extractWorkspaceAiRoutingSignals( + buildTurnContext({ + messages: [{ role: "user", content: "Fix this:\n```ts\nconst a = 1\n```" }], + } as Partial), + ); + + expect(signals.hasCodeSignals).toBe(true); + }); + + it("counts workspace references from the request body", () => { + const signals = extractWorkspaceAiRoutingSignals( + buildTurnContext({ + body: { + workspaceAiContext: { + selectedItems: [{ id: "a" }, { id: "b" }], + selectedQuotes: [{ text: "quoted" }], + }, + }, + }), + ); + + expect(signals.workspaceReferenceCount).toBe(3); + }); + + it("returns zero signals for missing or malformed turn data", () => { + expect( + extractWorkspaceAiRoutingSignals( + buildTurnContext({ + body: { workspaceAiContext: "garbage" }, + messages: [null, 42, { role: "user", content: { nested: true } }], + } as unknown as Partial), + ), + ).toEqual({ + attachmentCount: 0, + hasCodeSignals: false, + lastUserTextChars: 0, + messageCount: 3, + priorToolCallCount: 0, + totalTextChars: 0, + workspaceReferenceCount: 0, + }); + + expect( + extractWorkspaceAiRoutingSignals({ messages: undefined } as unknown as TurnContext), + ).toEqual({ + attachmentCount: 0, + hasCodeSignals: false, + lastUserTextChars: 0, + messageCount: 0, + priorToolCallCount: 0, + totalTextChars: 0, + workspaceReferenceCount: 0, + }); + }); +}); diff --git a/src/features/workspaces/ai/model-router.ts b/src/features/workspaces/ai/model-router.ts new file mode 100644 index 00000000..37617623 --- /dev/null +++ b/src/features/workspaces/ai/model-router.ts @@ -0,0 +1,263 @@ +import type { TurnContext } from "@cloudflare/think"; + +import { + DEFAULT_WORKSPACE_AI_CHAT_MODEL_ID, + getWorkspaceAiChatModelById, + type WorkspaceAiChatModelId, +} from "#/features/workspaces/ai/models"; + +// Task signals the Auto router reads for one turn. Extraction is defensive: +// anything missing or malformed becomes the zero value, never an error. +export interface WorkspaceAiRoutingSignals { + /** Image/file parts across user messages; the routed model must accept them. */ + attachmentCount: number; + /** Code fences or code-like tokens in the latest user message. */ + hasCodeSignals: boolean; + /** Text characters in the latest user message. */ + lastUserTextChars: number; + messageCount: number; + /** Tool-call parts already present in the conversation history. */ + priorToolCallCount: number; + /** Text characters across every message in the assembled turn. */ + totalTextChars: number; + /** Workspace items and quotes the user attached via the context picker. */ + workspaceReferenceCount: number; +} + +export type WorkspaceAiAutoRouteReason = + | "attachment-heavy" + | "long-context" + | "tool-heavy" + | "simple-chat" + | "high-reasoning-default" + | "fallback-unavailable"; + +export interface WorkspaceAiAutoRoute { + gatewayModel: string; + modelId: WorkspaceAiChatModelId; + reason: WorkspaceAiAutoRouteReason; +} + +// Tuning knobs for the routing rules below, kept as data so policy changes are +// threshold edits rather than runtime rewrites. +export const WORKSPACE_AI_AUTO_ROUTING_THRESHOLDS = { + /** Roughly 6k tokens of conversation, where large-context models earn their keep. */ + longContextTotalTextChars: 24_000, + longContextMessageCount: 40, + longContextWorkspaceReferenceCount: 4, + toolHeavyPriorToolCallCount: 3, + /** A couple of sentences at most. */ + simpleChatMaxLastUserTextChars: 280, + simpleChatMaxMessageCount: 12, +} as const; + +export type WorkspaceAiAutoRoutingThresholds = typeof WORKSPACE_AI_AUTO_ROUTING_THRESHOLDS; + +export interface WorkspaceAiAutoRoutingRule { + matches: ( + signals: WorkspaceAiRoutingSignals, + thresholds: WorkspaceAiAutoRoutingThresholds, + ) => boolean; + reason: Exclude; + targetModelId: WorkspaceAiChatModelId; +} + +// First match wins. Capability constraints (attachments, context size) come +// before cost preferences so a short message with an image still lands on a +// model that can read it. Targets must stay standard billing tier: a user who +// picked Auto must never be routed into premium pricing. +export const WORKSPACE_AI_AUTO_ROUTING_RULES: readonly WorkspaceAiAutoRoutingRule[] = [ + { + matches: (signals) => signals.attachmentCount > 0, + reason: "attachment-heavy", + targetModelId: "gemini", + }, + { + matches: (signals, thresholds) => + signals.totalTextChars >= thresholds.longContextTotalTextChars || + signals.messageCount >= thresholds.longContextMessageCount || + signals.workspaceReferenceCount >= thresholds.longContextWorkspaceReferenceCount, + reason: "long-context", + targetModelId: "gemini", + }, + { + matches: (signals, thresholds) => + signals.hasCodeSignals || + signals.priorToolCallCount >= thresholds.toolHeavyPriorToolCallCount, + reason: "tool-heavy", + targetModelId: "chatgpt-mini", + }, + { + matches: (signals, thresholds) => + signals.lastUserTextChars > 0 && + signals.lastUserTextChars <= thresholds.simpleChatMaxLastUserTextChars && + signals.messageCount <= thresholds.simpleChatMaxMessageCount, + reason: "simple-chat", + targetModelId: "claude-haiku", + }, +]; + +export function routeWorkspaceAiAutoModel( + signals: WorkspaceAiRoutingSignals, + rules: readonly WorkspaceAiAutoRoutingRule[] = WORKSPACE_AI_AUTO_ROUTING_RULES, +): WorkspaceAiAutoRoute { + for (const rule of rules) { + if (!rule.matches(signals, WORKSPACE_AI_AUTO_ROUTING_THRESHOLDS)) { + continue; + } + + const target = getStandardTierRouteTarget(rule.targetModelId); + + if (!target) { + // A rule pointing at a removed or re-tiered model degrades to the + // default instead of failing the turn or leaking into premium pricing. + return getDefaultWorkspaceAiAutoRoute("fallback-unavailable"); + } + + return { + gatewayModel: target.gatewayModel, + modelId: target.id, + reason: rule.reason, + }; + } + + return getDefaultWorkspaceAiAutoRoute("high-reasoning-default"); +} + +function getDefaultWorkspaceAiAutoRoute( + reason: Extract, +): WorkspaceAiAutoRoute { + const model = getWorkspaceAiChatModelById(DEFAULT_WORKSPACE_AI_CHAT_MODEL_ID); + + return { + gatewayModel: model.gatewayModel, + modelId: model.id, + reason, + }; +} + +function getStandardTierRouteTarget(modelId: WorkspaceAiChatModelId) { + const model = getWorkspaceAiChatModelById(modelId); + + // getWorkspaceAiChatModelById falls back to the default entry for unknown + // ids; treat that as "unavailable" rather than silently routing there. + if (model.id !== modelId || model.billingTier !== "standard") { + return undefined; + } + + return model; +} + +const EMPTY_WORKSPACE_AI_ROUTING_SIGNALS: WorkspaceAiRoutingSignals = { + attachmentCount: 0, + hasCodeSignals: false, + lastUserTextChars: 0, + messageCount: 0, + priorToolCallCount: 0, + totalTextChars: 0, + workspaceReferenceCount: 0, +}; + +// Lightweight markers that the user is doing code or data work. These bias +// toward the tool-following model; misses simply keep the default route, so +// the list favors precision over recall. +const WORKSPACE_AI_CODE_SIGNAL_PATTERNS: readonly RegExp[] = [ + /```/, + /\b(?:function|const|class |def |import |export |return|async|await)\b/, + /\b(?:SELECT|INSERT|UPDATE|DELETE)\b.+\b(?:FROM|INTO|SET|WHERE)\b/, + /\.(?:py|ts|tsx|js|jsx|json|csv|sql|sh|ya?ml)\b/i, + /\b(?:stack ?trace|traceback|regex|refactor|typescript|javascript|python)\b/i, +]; + +export function extractWorkspaceAiRoutingSignals(ctx: TurnContext): WorkspaceAiRoutingSignals { + try { + const messages: unknown[] = Array.isArray(ctx.messages) ? ctx.messages : []; + let attachmentCount = 0; + let priorToolCallCount = 0; + let totalTextChars = 0; + let lastUserText = ""; + + for (const message of messages) { + const record = asRecord(message); + + if (!record) { + continue; + } + + const text = getMessageText(record.content); + totalTextChars += text.length; + + if (record.role === "user") { + attachmentCount += countParts(record.content, ["file", "image"]); + lastUserText = text; + } else if (record.role === "assistant") { + priorToolCallCount += countParts(record.content, ["tool-call"]); + } + } + + return { + attachmentCount, + hasCodeSignals: WORKSPACE_AI_CODE_SIGNAL_PATTERNS.some((pattern) => + pattern.test(lastUserText), + ), + lastUserTextChars: lastUserText.length, + messageCount: messages.length, + priorToolCallCount, + totalTextChars, + workspaceReferenceCount: countWorkspaceReferences(ctx.body?.workspaceAiContext), + }; + } catch { + // Routing signals are best-effort: a malformed turn falls back to the + // default route instead of failing the turn. + return EMPTY_WORKSPACE_AI_ROUTING_SIGNALS; + } +} + +function asRecord(value: unknown): Record | undefined { + return typeof value === "object" && value !== null + ? (value as Record) + : undefined; +} + +function getMessageText(content: unknown): string { + if (typeof content === "string") { + return content; + } + + if (!Array.isArray(content)) { + return ""; + } + + return content + .map((part) => { + const record = asRecord(part); + return record?.type === "text" && typeof record.text === "string" ? record.text : ""; + }) + .join(""); +} + +function countParts(content: unknown, types: readonly string[]): number { + if (!Array.isArray(content)) { + return 0; + } + + return content.filter((part) => { + const type = asRecord(part)?.type; + return typeof type === "string" && types.includes(type); + }).length; +} + +function countWorkspaceReferences(workspaceAiContext: unknown): number { + const snapshot = asRecord(workspaceAiContext); + + if (!snapshot) { + return 0; + } + + const selectedItems = Array.isArray(snapshot.selectedItems) ? snapshot.selectedItems.length : 0; + const selectedQuotes = Array.isArray(snapshot.selectedQuotes) + ? snapshot.selectedQuotes.length + : 0; + + return selectedItems + selectedQuotes; +} diff --git a/src/features/workspaces/ai/models.ts b/src/features/workspaces/ai/models.ts index 8e4ce839..baca5159 100644 --- a/src/features/workspaces/ai/models.ts +++ b/src/features/workspaces/ai/models.ts @@ -24,14 +24,15 @@ export const WORKSPACE_AI_CHAT_MODELS = [ { id: "auto", name: "Auto", - // ThinkEx's own "let us pick for you" option. For now "Auto" is Kimi K2.6 - // under the hood until we build or adopt a real router; the slug can - // change without affecting the user-facing choice. + // ThinkEx's own "let us pick for you" option. Each turn is routed to a + // standard-tier model by the task router (see model-router.ts); this + // gateway slug is the router's default and fallback, and also serves + // non-chat paths (base model, compaction) that resolve the "auto" id. gatewayModel: "moonshotai/kimi-k2.6", provider: "auto", - tagline: "Picks a good fit for you", + tagline: "Picks a model per message", description: - "ThinkEx picks the model for you. It stays quick for simple things and uses a stronger model when the task is harder.", + "ThinkEx picks a model for each message. Quick questions get a fast model, big documents and images get one that reads a lot, and harder tasks get a stronger one.", bestFor: "Most tasks", intelligence: 3, speed: 3, diff --git a/src/features/workspaces/components/ai-chat/AiChatInspectorViews.tsx b/src/features/workspaces/components/ai-chat/AiChatInspectorViews.tsx index 2bedf4f9..09acfa23 100644 --- a/src/features/workspaces/components/ai-chat/AiChatInspectorViews.tsx +++ b/src/features/workspaces/components/ai-chat/AiChatInspectorViews.tsx @@ -62,7 +62,14 @@ function RunSummary({ run }: { run: AIInspectorRunView }) { {run.modelId ? ( - {run.modelId} + {run.requestedModelId && run.requestedModelId !== run.modelId + ? `${run.requestedModelId} → ${run.modelId}` + : run.modelId} + + ) : null} + {run.routingReason ? ( + + {run.routingReason} ) : null} diff --git a/src/integrations/posthog/events.ts b/src/integrations/posthog/events.ts index c268f228..242a19d0 100644 --- a/src/integrations/posthog/events.ts +++ b/src/integrations/posthog/events.ts @@ -45,6 +45,8 @@ export interface PostHogEventPropertiesByName { workspace_id: string; trace_id: string; model_id: string; + requested_model_id: string; + routing_reason: string | null; continuation: boolean; }; ai_turn_completed: {