Skip to content

feat(ai-chat): route Auto model by task signals #599

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
v-rdyy wants to merge 5 commits into
base: main
Choose a base branch
from
Open

feat(ai-chat): route Auto model by task signals #599

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
91f3abb
feat(ai-chat): add auto model routing policy and router
v-rdy Jul 7, 2026
2cefa29
feat(ai-chat): extract routing signals from turn context
v-rdy Jul 7, 2026
9fef534
feat(ai-chat): route auto turns through the task router
v-rdy Jul 7, 2026
b8ed446
feat(ai-chat): record requested and routed model in telemetry
v-rdy Jul 7, 2026
944f556
fix(ai-chat): align Auto picker copy with routed behavior
v-rdy Jul 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 28 additions & 0 deletions src/features/workspaces/ai/ai-inspector-view-model.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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");
});
});
2 changes: 2 additions & 0 deletions src/features/workspaces/ai/ai-inspector-view-model.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
2 changes: 2 additions & 0 deletions src/features/workspaces/ai/ai-inspector-view-types.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
4 changes: 4 additions & 0 deletions src/features/workspaces/ai/ai-thread-inspector-recorder.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,8 @@ export class AIThreadInspectorRecorder {
async recordTurnStarted(input: {
ctx: TurnContext;
modelId: string;
requestedModelId: string;
routingReason?: string;
system: string;
thread: AIThreadContext;
tools: unknown;
Expand All @@ -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),
Expand Down
4 changes: 4 additions & 0 deletions src/features/workspaces/ai/ai-thread-posthog-recorder.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,8 @@ export class AIThreadPostHogRecorder {
recordTurnStarted(input: {
ctx: TurnContext;
modelId: WorkspaceAiChatModelId;
requestedModelId: WorkspaceAiChatModelId;
routingReason?: string;
thread: AIThreadContext;
tools?: unknown;
}) {
Expand Down Expand Up @@ -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,
Expand Down
8 changes: 8 additions & 0 deletions src/features/workspaces/ai/ai-thread-tcc-recorder.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,8 @@ export class AIThreadTccRecorder {
ctx: TurnContext;
env: Cloudflare.Env;
modelId: WorkspaceAiChatModelId;
requestedModelId: WorkspaceAiChatModelId;
routingReason?: string;
system: string;
thread: AIThreadContext;
tools?: unknown;
Expand All @@ -94,6 +96,8 @@ export class AIThreadTccRecorder {
const metadata = createTccMetadata({
gatewayModel,
modelId: input.modelId,
requestedModelId: input.requestedModelId,
routingReason: input.routingReason,
thread: input.thread,
});

Expand Down Expand Up @@ -454,6 +458,8 @@ function getTccApiKey(env: Cloudflare.Env) {
function createTccMetadata(input: {
gatewayModel: string;
modelId: WorkspaceAiChatModelId;
requestedModelId: WorkspaceAiChatModelId;
routingReason?: string;
thread: AIThreadContext;
}) {
return {
Expand All @@ -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<string, string>;
}
Expand Down
4 changes: 4 additions & 0 deletions src/features/workspaces/ai/ai-thread-telemetry-recorder.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
52 changes: 51 additions & 1 deletion src/features/workspaces/ai/ai-thread.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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;
}

Expand Down Expand Up @@ -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,
Expand All @@ -188,6 +201,8 @@ export function createAIThreadClass(getUserAIStore: () => typeof UserAIStore) {

this.activeUsageContext = {
modelId,
requestedModelId,
routingReason: route?.reason,
thread,
};
}
Expand All @@ -213,6 +228,8 @@ export function createAIThreadClass(getUserAIStore: () => typeof UserAIStore) {
await this.telemetry.recordTurnStarted({
ctx,
modelId,
requestedModelId,
routingReason: route?.reason,
system,
thread,
tools: activeTools,
Expand Down Expand Up @@ -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;
}
Comment on lines +318 to +332

@coderabbitai coderabbitai Bot Jul 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Explicit-picker continuation branch still switches models mid-flight.

When a continuation's activeContext.requestedModelId !== "auto" (original turn used an explicit model, e.g. claude-sonnet), this branch returns undefined. Since beforeTurn computes modelId = route?.modelId ?? requestedModelId, and requestedModelId for this continuation call is "auto" (the body lost the explicit selection), the effective model silently becomes Auto's default (moonshotai/kimi-k2.6) instead of the originally selected explicit model. This directly contradicts the function's own header comment that "a run never switches models mid-flight" — continuity is only preserved for the auto-originated branch (lines 334-339), not for the explicit-picker branch.

🐛 Proposed fix to preserve the original explicit model 
 			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;
+				// The run started from an explicit picker choice; a continuation
+				// whose body lost the selection must reuse that model, not silently
+				// fall back to Auto's default.
+				return { modelId: activeContext.modelId, reason: "turn-continuation" };
 			}
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
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;
}
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 reuse that model, not silently
// fall back to Auto's default.
return { modelId: activeContext.modelId, reason: "turn-continuation" };
}
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/features/workspaces/ai/ai-thread.ts` around lines 318 - 332, The
explicit-picker continuation path in _resolveAutoModelRoute is returning
undefined, which lets beforeTurn fall back to the Auto default instead of
preserving the original model. Update _resolveAutoModelRoute so that when
ctx.continuation is true and this.activeUsageContext.requestedModelId is not
"auto", it routes back to that original explicit modelId with an appropriate
reason. Keep the existing auto-originated continuation handling intact, and make
sure beforeTurn still uses route?.modelId to prevent mid-flight model switches.


if (activeContext) {
return {
modelId: activeContext.modelId,
reason: activeContext.routingReason ?? "turn-continuation",
};
}

const routed = routeWorkspaceAiAutoModel(extractWorkspaceAiRoutingSignals(ctx));

return { modelId: routed.modelId, reason: routed.reason };
}

Comment on lines +333 to +345

@coderabbitai coderabbitai Bot Jul 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Check if activeUsageContext (or equivalent) is persisted to durable storage anywhere,
# and how "continuation" is defined/triggered relative to DO hibernation.
rg -n "activeUsageContext" -C3 --type=ts
rg -n "continuation" -C3 --type=ts src/features/workspaces/ai/ai-thread.ts

Repository: ThinkEx-OSS/thinkex

Length of output: 157

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Map the target file and read the relevant sections.
wc -l src/features/workspaces/ai/ai-thread.ts
sed -n '160,460p' src/features/workspaces/ai/ai-thread.ts

# Search for persistence/restoration and continuation-related paths.
rg -n "activeUsageContext|usageContext|routeWorkspaceAiAutoModel|extractWorkspaceAiRoutingSignals|beforeTurn|_trackCompletedMessageUsage|continuation|durable|storage|hibernate|hibernation|restore|recover" src/features/workspaces/ai -C 3
rg -n "activeUsageContext|usageContext|continuation|hibernate|hibernation|durable|storage|restore|recover" src -C 2

Repository: ThinkEx-OSS/thinkex

Length of output: 50377

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Find durable-object / hibernation / state-restoration references near AIThread.
rg -n "hibernate|hibernation|durable object|Durable Object|storage|persist|restore|recover|recovery" src/features/workspaces/ai -C 2

# Show the AIThread class definition and any storage-related fields/methods.
sed -n '1,160p' src/features/workspaces/ai/ai-thread.ts

Repository: ThinkEx-OSS/thinkex

Length of output: 13924

🌐 Web query:

Cloudflare Think Durable Object continuation recovery instance fields hibernation activeUsageContext beforeTurn continuation onChatResponse

💡 Result:

In the context of Cloudflare Agents and Durable Objects, the terms you mentioned pertain to the lifecycle, lifecycle hooks, and operational behavior of agentic applications built on the Cloudflare platform [1][2]. ActiveUsageContext and BeforeTurn BeforeTurn is a lifecycle hook in the Cloudflare Agents framework that fires before streamText is executed [3][2]. It receives a TurnContext object containing the assembled system prompt, messages, tools, and model. This hook allows developers to return a TurnConfig object to override parts of the turn (e.g., model, system prompt, tools, maxSteps) [3][2]. Continuation Continuation refers to a specific type of turn in an agent's lifecycle, often used for auto-continuing a conversation after tool results are coalesced [3][2]. In the TurnContext provided to lifecycle hooks like beforeTurn, the continuation field is a boolean indicating whether the current turn is an auto-continuation [2]. Hibernation and Durable Objects Durable Objects support WebSocket Hibernation, which allows objects to be removed from memory to save costs while maintaining active WebSocket connections [4][5]. When an agent relies on Durable Objects for state, hibernation can impact certain behaviors: - Programmatic turns and sub-agent RPCs may be affected, as certain signals (like AbortSignal) do not persist across hibernation boundaries [6]. - Lifecycle hooks and responses: Some responses triggered from within hooks like onChatResponse do not recursively fire onChatResponse again, which is a consideration when designing agent communication flows [7][8]. Recovery Durable recovery refers to mechanisms designed to manage stability and handle interruptions (such as stream-stall watchdogs) [1]. The chatStreamStallTimeoutMs field in TurnConfig is specifically mentioned as a way to override the watchdog for a specific turn, such as when executing a known-slow tool, and it is explicitly noted to refer to Durable recovery documentation [2][9].

Citations:

🌐 Web query:

Cloudflare Durable Objects instance fields lost on hibernation state.storage persistence docs

💡 Result:

When a Cloudflare Durable Object enters a hibernation state—or is evicted from memory due to inactivity or system events like code deployments—all in-memory instance fields and variables are discarded [1][2][3][4]. This occurs because the object is removed from the active memory of the runtime [1][5]. To maintain state across these events, you must explicitly persist any necessary information to the Durable Object's durable storage using the Storage API [1][2][3]. The Durable Object Storage API provides transactional, strongly consistent, and private storage (including SQL and key-value options) that survives hibernation and eviction [1][6][5][7]. Special handling is required for WebSocket Hibernation: * While general in-memory state is lost, you can use the Hibernation WebSocket API methods serializeAttachment and deserializeAttachment to persist per-connection data [8][9]. * These attachments survive hibernation as long as the WebSocket connection remains healthy [8]. * If you need to store larger amounts of data or information that must persist beyond the lifetime of a specific WebSocket connection, you should use the standard Storage API and store only the reference key in the WebSocket attachment [8]. In summary, you should treat all instance variables as ephemeral and assume they will be re-initialized from persistent storage upon the next request [1][4].

Citations:

Persist activeUsageContext for resumed continuations. A fresh instance can reroute the turn from current signals, but _trackCompletedMessageUsage will then skip billing because activeUsageContext is still unset. Store or rehydrate the routed model/thread context so hibernation or recovery doesn’t drop usage tracking for completed runs.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/features/workspaces/ai/ai-thread.ts` around lines 333 - 345, Persist or
rehydrate the routed context in ai-thread.ts so resumed continuations keep an
activeUsageContext instead of only rerouting from current signals. Update the
logic around the existing activeContext check and the
routeWorkspaceAiAutoModel/extractWorkspaceAiRoutingSignals path so the returned
model/thread routing data is also stored on the thread state and available to
_trackCompletedMessageUsage. This should ensure hibernation or recovery restores
the same usage context before completed-message billing runs.

private async _getThreadContext() {
const directory = await this.parentAgent(getUserAIStore());
return directory.getThreadContext(this.name);
Expand Down
Loading
Loading