refactor(pbl-v2): route runtime agents through the shared LLM entry point

The five PBL v2 runtime call sites (instructor x2, evaluator, simulator x2)
invoked the AI SDK directly while the other 21 call sites in lib/ and app/ go
through callLLM / streamLLM, so they silently opted out of everything the shared
entry point applies:

- Usage accounting. recordUsage is reached only from the wrappers, so no
  instructor turn, evaluator pass, simulator line or narrator pass was ever
  accounted for - the highest-frequency, longest-context LLM traffic in the
  product. The comment at llm.ts:275, that every server-side call funnels
  through the wrappers, did not hold.
- The LLM_THINKING_DISABLED kill switch, read inside the wrappers only.
- One thinking resolution path. runtime-thinking.ts existed precisely because
  the agents bypassed the wrapper, and it was wired into instructor (x2) and
  evaluator but neither simulator site.

Also fix callLLM recording result.usage, which on a multi-step tool run
(stopWhen) is the final step alone; totalUsage is the aggregate, and the
instructor turn is exactly such a run. streamLLM already preferred the
aggregate.

Thinking semantics are preserved per call site: the teaching turns keep
force-disabling thinking, now through the wrapper's own `thinking` argument,
and keep their existing providerOptions spread, which injectProviderOptions
yields to; the simulator keeps honouring a per-request config, as before. The
one behavioural difference is on native adapters when no thinking config
arrives: the policy now resolves to an explicit lowest/disabled provider option
instead of leaving the model's default in place - which is what the policy
always claimed to do, and what OpenAI-compatible providers already got.
Unifying the wider divergence is a product decision and stays out of scope;
runtime-thinking.ts now documents it, along with the fact that the original
tool_choice justification no longer matches the tree.

A lint guard restricts importing generateText / streamText to lib/ai/llm.ts so
the split cannot grow back. It uses the @typescript-eslint variant of
no-restricted-imports deliberately: flat config replaces a rule's options per
key rather than merging them, and the base rule is already configured for the
lib/choreography and lib/video-export module boundaries.

The runtime-thinking test moves from asserting the retired helper to asserting
what the call site is responsible for, at the provider boundary: an evaluator
turn reaches the provider with a disabled config in the thinking store, does
not leak it, and is accounted for under its own usage source.

Closes #1003

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
wyuc
2026-07-28 22:59:20 -04:00
co-authored by Claude Opus 5
parent 0cf2a33041
commit cc48adae91
8 changed files with 346 additions and 77 deletions
+45
View File
@@ -362,6 +362,51 @@ const eslintConfig = defineConfig([
],
},
},
// Single LLM entry point (machine-enforced): server-side model calls go through
// `callLLM` / `streamLLM` in lib/ai/llm.ts. That wrapper is where usage
// accounting (`recordUsage`), the `LLM_THINKING_DISABLED` kill switch, and
// per-provider thinking resolution live — a direct `generateText` / `streamText`
// silently opts out of all three, and the opt-out is invisible at the call site.
// The PBL v2 runtime drifted exactly this way (#1003): five direct calls meant
// zero usage records for the busiest traffic in the product, plus three
// different meanings for one thinking config.
//
// The rule uses the @typescript-eslint variant deliberately: the base
// `no-restricted-imports` is already configured for lib/choreography and
// lib/video-export, and flat config REPLACES a rule's options per key rather
// than merging them, so reusing that key here would silently drop those
// module-boundary bans. Different key, no interference.
//
// Out of scope (same spirit as the boundaries above): `require('ai')` and
// dynamic `import('ai')`. Nothing in the tree reaches the SDK that way, and the
// static import is the shape a hurried call site actually reaches for.
{
files: ['**/*.{ts,tsx}'],
ignores: [
// The entry point itself.
'lib/ai/llm.ts',
// Offline harnesses and a test whose subject IS the SDK: not server request
// paths, nothing to account for, and the integration test must be able to
// call the raw SDK to assert what the wrapper is built on.
'eval/**',
'tests/**',
],
rules: {
'@typescript-eslint/no-restricted-imports': [
'error',
{
paths: [
{
name: 'ai',
importNames: ['generateText', 'streamText'],
message:
'Call the model through callLLM / streamLLM in @/lib/ai/llm instead of the AI SDK directly — that is where usage accounting, the LLM_THINKING_DISABLED kill switch, and per-provider thinking resolution are applied. Both wrappers pass every SDK option straight through, and take an optional per-call thinking config.',
},
],
},
],
},
},
]);
export default eslintConfig;
+5 -1
View File
@@ -352,7 +352,11 @@ export async function callLLM<T extends GenerateTextParams>(
continue;
}
recordUsageSafe(result.usage, buildUsageMeta(params, source));
// `usage` is the LAST step only; on a multi-step tool run (`stopWhen`)
// every earlier step would go unaccounted. `totalUsage` aggregates across
// steps and equals `usage` for a single-step call. Mirrors streamLLM,
// which already prefers the aggregate.
recordUsageSafe(result.totalUsage ?? result.usage, buildUsageMeta(params, source));
return result;
} catch (error) {
lastError = error;
+7 -6
View File
@@ -41,12 +41,11 @@
* persist).
*/
import { streamText } from 'ai';
import type { LanguageModel } from 'ai';
import { createLogger } from '@/lib/logger';
import { resolveThinkingProviderOptions } from '@/lib/ai/llm';
import { withThinkingDisabled } from './runtime-thinking';
import { resolveThinkingProviderOptions, streamLLM } from '@/lib/ai/llm';
import { PBL_V2_TEACHING_THINKING } from './runtime-thinking';
import { buildVisionUserContent } from '@/lib/generation/prompt-formatters';
import type { ThinkingConfig } from '@/lib/types/provider';
@@ -154,8 +153,8 @@ async function* runShared(args: RunSharedArgs): AsyncGenerator<PBLSSEEvent, void
const isMilestone = kind === 'milestone';
let lastSanitizedLength = 0;
try {
const result = withThinkingDisabled(() =>
streamText({
const result = streamLLM(
{
model: languageModel,
system: systemPrompt,
// Image submission on a vision-capable model → send the picture as a
@@ -177,7 +176,9 @@ async function* runShared(args: RunSharedArgs): AsyncGenerator<PBLSSEEvent, void
? { providerOptions: resolveThinkingProviderOptions(languageModel, thinkingConfig) }
: {}),
...(signal ? { abortSignal: signal } : {}),
}),
},
`pbl-v2-evaluator-${kind}`,
PBL_V2_TEACHING_THINKING,
);
for await (const part of result.fullStream) {
switch (part.type) {
+13 -9
View File
@@ -11,16 +11,16 @@
* owned by the right-side submission/evaluation flow, not by chat.
*/
import { streamText, tool, stepCountIs } from 'ai';
import { tool, stepCountIs } from 'ai';
import type { LanguageModel } from 'ai';
import { createLogger } from '@/lib/logger';
import { resolveThinkingProviderOptions } from '@/lib/ai/llm';
import { resolveThinkingProviderOptions, streamLLM } from '@/lib/ai/llm';
import { loadPBLV2Prompt } from '../prompts/loader';
import type { ThinkingConfig } from '@/lib/types/provider';
import { tierGuidanceBlock } from './tier-guidance';
import { compressIfNeeded } from './instructor-memory';
import { withThinkingDisabled } from './runtime-thinking';
import { PBL_V2_TEACHING_THINKING } from './runtime-thinking';
import type {
PBLProjectV2,
@@ -1497,8 +1497,8 @@ export async function* runInstructorTurn(
// speaking, leaving the learner with an empty chat. The instructing path
// exposes only the two non-advance tools (record_observation /
// adjust_difficulty).
const result = withThinkingDisabled(() =>
streamText({
const result = streamLLM(
{
model: languageModel,
system: systemPrompt,
messages: finalMessages,
@@ -1512,7 +1512,9 @@ export async function* runInstructorTurn(
? { providerOptions: resolveThinkingProviderOptions(languageModel, thinkingConfig) }
: {}),
...(signal ? { abortSignal: signal } : {}),
}),
},
'pbl-v2-instructor',
PBL_V2_TEACHING_THINKING,
);
for await (const part of result.fullStream) {
@@ -1770,8 +1772,8 @@ async function* runSetupFollowup(args: SetupFollowupArgs): AsyncGenerator<PBLSSE
const historyMessages = buildSetupHistoryMessages(instructorThread);
try {
const result = withThinkingDisabled(() =>
streamText({
const result = streamLLM(
{
model: languageModel,
system: systemPrompt,
messages: [
@@ -1785,7 +1787,9 @@ async function* runSetupFollowup(args: SetupFollowupArgs): AsyncGenerator<PBLSSE
? { providerOptions: resolveThinkingProviderOptions(languageModel, thinkingConfig) }
: {}),
...(signal ? { abortSignal: signal } : {}),
}),
},
'pbl-v2-instructor-setup-followup',
PBL_V2_TEACHING_THINKING,
);
let rawAssistantText = '';
+27 -21
View File
@@ -1,28 +1,34 @@
import { thinkingContext } from '@/lib/ai/thinking-context';
import type { ThinkingConfig } from '@/lib/types/provider';
/**
* PBL v2 runtime LLM calls force-disable thinking.
* Thinking policy for the PBL v2 teaching turns (instructor / evaluator):
* thinking is force-disabled. Generation (planner) is intentionally untouched,
* and so is the simulator — see the caveats below.
*
* The instructing turn forces `begin_turn` via `tool_choice`, which several
* providers reject when thinking is on — DeepSeek returns 400 "Thinking mode
* does not support this tool_choice". We never intentionally enabled thinking
* on these turns (the PBL v2 client sends no `thinkingConfig`); some pinned
* models just default it on. Disabling it removes the incompatibility without
* losing any behavior we relied on.
* We never intentionally enabled thinking on these turns (the PBL v2 client
* sends no `thinkingConfig`); some pinned models just default it on, and the
* teaching turns gain nothing from it.
*
* For OpenAI-compatible providers (e.g. DeepSeek) thinking is injected by the
* fetch wrapper in `providers.ts`, which reads the per-request config from the
* `thinkingContext` AsyncLocalStorage. The agents call the AI SDK directly
* (not via `callLLM`/`streamLLM`), so nothing seeds that store — we do it here.
* The SDK call must be started INSIDE `run` (its consumption can be outside) so
* the lazily-issued fetch inherits the context, matching `streamLLM`.
* This module used to also carry the *mechanism*: a `withThinkingDisabled`
* helper that seeded the `thinkingContext` AsyncLocalStorage by hand, because
* the agents called the AI SDK directly and nothing else seeded it. The agents
* now go through `streamLLM` / `callLLM`, which take a `thinking` argument and
* do the seeding (plus provider-option resolution) themselves — so only the
* policy value is left here.
*
* Scope: PBL v2 runtime only (instructor / evaluator). Generation (planner) is
* intentionally untouched.
* Two caveats a future reader should know before treating this as settled:
*
* 1. The original justification has expired. It read: "the instructing turn
* forces `begin_turn` via `tool_choice`, which several providers reject when
* thinking is on — DeepSeek returns 400 'Thinking mode does not support this
* tool_choice'". `begin_turn` no longer exists, and the instructing turn now
* passes `tools` + `stopWhen` and forces nothing. If no provider still
* rejects these calls, the policy can be dropped outright rather than
* applied — it is kept only because the absence of a 400 is harder to
* confirm than its presence.
* 2. The simulator does NOT apply this policy, and did not before either. Its
* turns therefore honour a per-request / stage-route `thinkingConfig` where
* the teaching turns override it. That divergence is pre-existing and
* deliberate here; unifying it is a product decision, not a refactor.
*/
const PBL_V2_THINKING_DISABLED: ThinkingConfig = { mode: 'disabled', enabled: false };
export function withThinkingDisabled<T>(startCall: () => T): T {
return thinkingContext.run(PBL_V2_THINKING_DISABLED, startCall);
}
export const PBL_V2_TEACHING_THINKING: ThinkingConfig = { mode: 'disabled', enabled: false };
+27 -17
View File
@@ -23,11 +23,10 @@
* intentionally NOT here — they land in increment 4.
*/
import { streamText, generateText } from 'ai';
import type { LanguageModel } from 'ai';
import { createLogger } from '@/lib/logger';
import { resolveThinkingProviderOptions } from '@/lib/ai/llm';
import { callLLM, resolveThinkingProviderOptions, streamLLM } from '@/lib/ai/llm';
import type { ThinkingConfig } from '@/lib/types/provider';
import { loadPBLV2Prompt } from '../prompts/loader';
@@ -278,12 +277,17 @@ async function runDirectorNarratorPass(args: {
const messages = [...history, { role: 'user' as const, content: nudge }];
try {
const result = await generateText({
model: languageModel,
system,
messages,
...(signal ? { abortSignal: signal } : {}),
});
// No thinking argument: the narrator pass never resolved a thinking config
// and still doesn't — see the caveats in `runtime-thinking.ts`.
const result = await callLLM(
{
model: languageModel,
system,
messages,
...(signal ? { abortSignal: signal } : {}),
},
'pbl-v2-simulator-narrator',
);
const text = (result.text ?? '').trim();
if (!text) return [];
// Sentinel: model says nothing happened → no narration this turn.
@@ -500,15 +504,21 @@ export async function* runSimulatorTurn(
// so the caller can decide retry vs surface.
async function* streamCharacterLine(): AsyncGenerator<PBLSSEEvent, string, void> {
let acc = '';
const stream = streamText({
model: languageModel,
system,
messages,
...(thinkingConfig
? { providerOptions: resolveThinkingProviderOptions(languageModel, thinkingConfig) }
: {}),
...(signal ? { abortSignal: signal } : {}),
});
// No thinking argument: unlike the teaching turns the simulator has never
// force-disabled thinking, so a per-request / stage-route config keeps
// applying — see the caveats in `runtime-thinking.ts`.
const stream = streamLLM(
{
model: languageModel,
system,
messages,
...(thinkingConfig
? { providerOptions: resolveThinkingProviderOptions(languageModel, thinkingConfig) }
: {}),
...(signal ? { abortSignal: signal } : {}),
},
'pbl-v2-simulator',
);
for await (const part of stream.fullStream) {
if (part.type === 'text-delta') {
const delta =
+63 -5
View File
@@ -1,11 +1,16 @@
import { beforeEach, describe, expect, it, vi } from 'vitest';
const aiMock = vi.hoisted(() => ({
generateText: vi.fn(async (params: unknown) => ({
text: 'ok',
params,
usage: undefined as unknown,
})),
// `totalUsage` is optional here so a single-step case can omit it: multi-step
// runs report the aggregate there and callLLM prefers it over `usage`.
generateText: vi.fn(
async (
params: unknown,
): Promise<{ text: string; params: unknown; usage?: unknown; totalUsage?: unknown }> => ({
text: 'ok',
params,
}),
),
streamText: vi.fn(),
}));
@@ -116,6 +121,59 @@ describe('LLM thinking provider options', () => {
});
});
it('records the aggregate usage of a multi-step tool run, not the last step', async () => {
// `usage` on a multi-step run (`stopWhen`) is the final step alone; every
// earlier step's tokens live only in `totalUsage`.
aiMock.generateText.mockResolvedValueOnce({
text: 'ok',
params: undefined,
usage: { inputTokens: 3, outputTokens: 4 },
totalUsage: { inputTokens: 30, outputTokens: 40 },
});
await callLLM(
{
model: {
provider: 'openai.responses',
modelId: 'gpt-5.6',
},
prompt: 'hi',
} as Parameters<typeof callLLM>[0],
'test',
);
await vi.waitFor(() => {
expect(usageMock.recordUsage).toHaveBeenCalledWith(
expect.objectContaining({ usage: { inputTokens: 30, outputTokens: 40 } }),
);
});
});
it('falls back to the single-step usage when no aggregate is reported', async () => {
aiMock.generateText.mockResolvedValueOnce({
text: 'ok',
params: undefined,
usage: { inputTokens: 3, outputTokens: 4 },
});
await callLLM(
{
model: {
provider: 'openai.responses',
modelId: 'gpt-5.6',
},
prompt: 'hi',
} as Parameters<typeof callLLM>[0],
'test',
);
await vi.waitFor(() => {
expect(usageMock.recordUsage).toHaveBeenCalledWith(
expect.objectContaining({ usage: { inputTokens: 3, outputTokens: 4 } }),
);
});
});
it('sends Claude Haiku 4.5 thinking budget without effort', async () => {
await callLLM(
{
+159 -18
View File
@@ -1,28 +1,169 @@
import { describe, it, expect } from 'vitest';
import { thinkingContext } from '@/lib/ai/thinking-context';
import { withThinkingDisabled } from '@/lib/pbl/v2/agents/runtime-thinking';
import { beforeEach, describe, expect, it, vi } from 'vitest';
import { MockLanguageModelV3, convertArrayToReadableStream } from 'ai/test';
/**
* #669: PBL v2 runtime LLM calls must run with thinking disabled so a pinned
* model whose thinking defaults on (e.g. DeepSeek) doesn't reject the forced
* `begin_turn` tool_choice. The provider fetch wrapper reads the disabled
* config from this AsyncLocalStorage, so the contract is "the store reads
* disabled inside the wrapped call".
* #669 / #1003: the PBL v2 teaching turns must reach the model through
* `streamLLM` with thinking force-disabled.
*
* The original test asserted the mechanism — a `withThinkingDisabled` helper
* that seeded the thinking AsyncLocalStorage by hand. That helper is gone: the
* agents now go through the shared entry point, which seeds the store itself.
* So this asserts the two things the call site is actually responsible for, at
* the provider boundary where a regression would bite:
*
* 1. the store reads a disabled config by the time the provider is invoked
* (the OpenAI-compatible fetch wrapper reads exactly this), and
* 2. the turn is accounted for — going through the wrapper is what makes the
* call visible to usage recording at all.
*
* Both break silently if someone drops the arguments at the call site, which is
* how the direct-SDK drift in #1003 went unnoticed for so long.
*/
describe('withThinkingDisabled (#669)', () => {
it('seeds the thinking AsyncLocalStorage with a disabled config', () => {
const seen = withThinkingDisabled(() => thinkingContext.getStore());
expect(seen).toEqual({ mode: 'disabled', enabled: false });
const usageMock = vi.hoisted(() => ({
normalizeUsage: vi.fn((usage: unknown) => usage),
recordUsage: vi.fn(async () => undefined),
}));
vi.mock('@/lib/usage/normalize', () => ({
normalizeUsage: usageMock.normalizeUsage,
}));
vi.mock('@/lib/server/usage-storage', () => ({
recordUsage: usageMock.recordUsage,
}));
import { thinkingContext } from '@/lib/ai/thinking-context';
import { PBL_V2_TEACHING_THINKING } from '@/lib/pbl/v2/agents/runtime-thinking';
import { runTaskEvaluation } from '@/lib/pbl/v2/agents/evaluator';
import { addSubmission } from '@/lib/pbl/v2/operations/submission';
import type { PBLProjectV2 } from '@/lib/pbl/v2/types';
import type { PBLSSEEvent } from '@/lib/pbl/v2/api/sse';
type DoStreamConfig = NonNullable<
NonNullable<ConstructorParameters<typeof MockLanguageModelV3>[0]>['doStream']
>;
type StreamResult = Extract<DoStreamConfig, { stream: unknown }>;
type StreamPart = StreamResult['stream'] extends ReadableStream<infer P> ? P : never;
const USAGE = {
inputTokens: { total: 12, noCache: 12, cacheRead: 0, cacheWrite: 0 },
outputTokens: { total: 8, text: 8, reasoning: 0 },
};
function textStep(text: string): StreamPart[] {
return [
{ type: 'stream-start', warnings: [] },
{ type: 'text-start', id: 'p1' },
{ type: 'text-delta', id: 'p1', delta: text },
{ type: 'text-end', id: 'p1' },
{ type: 'finish', finishReason: { unified: 'stop', raw: 'stop' }, usage: USAGE },
];
}
function mkProject(): PBLProjectV2 {
return {
uiPhase: 'workspace',
title: 't',
description: 'd',
proficiency: 'intermediate',
language: 'zh-CN',
tags: [],
status: 'active',
roles: [],
milestones: [
{
id: 'ms1',
title: 'M1',
status: 'completed',
order: 0,
microtasks: [
{
id: 't1',
title: 'T1',
status: 'completed',
assignee: 'user',
hints: [],
order: 0,
},
],
documents: [],
},
],
submissions: [],
evaluations: [],
threads: [],
engagementEvents: [],
createdAt: 'ts',
updatedAt: 'ts',
};
}
/** Run one task evaluation against a scripted model, capturing what the
* provider saw in the thinking store at the moment it was invoked. */
async function runEvaluationCapturingThinking(): Promise<{ seenThinking: unknown }> {
const project = mkProject();
addSubmission(project, {
microtaskId: 't1',
milestoneId: 'ms1',
kind: 'text',
content: 'submission',
});
it('returns the wrapped callable result (so streamText/generateText pass through)', () => {
const result = withThinkingDisabled(() => 'sentinel');
expect(result).toBe('sentinel');
const NOT_CAPTURED = Symbol('provider was never invoked');
let seenThinking: unknown = NOT_CAPTURED;
const model = new MockLanguageModelV3({
doStream: async () => {
seenThinking = thinkingContext.getStore();
return {
stream: convertArrayToReadableStream(textStep('{"feedback":"ok","score":80}')),
};
},
});
it('does not leak the disabled config outside the wrapped call', () => {
withThinkingDisabled(() => undefined);
const events: PBLSSEEvent[] = [];
for await (const ev of runTaskEvaluation({
project,
milestoneId: 'ms1',
microtaskId: 't1',
languageModel: model,
})) {
events.push(ev);
}
// Guard the guard: an evaluation that errored out before calling the model
// would make every assertion below vacuous.
expect(seenThinking).not.toBe(NOT_CAPTURED);
expect(events.at(-1)?.type).toBe('done');
return { seenThinking };
}
describe('PBL v2 teaching-turn thinking policy (#669, #1003)', () => {
beforeEach(() => {
usageMock.normalizeUsage.mockClear();
usageMock.recordUsage.mockClear();
});
it('is a disabled thinking config', () => {
expect(PBL_V2_TEACHING_THINKING).toEqual({ mode: 'disabled', enabled: false });
});
it('reaches the provider as a disabled config on an evaluator turn', async () => {
const { seenThinking } = await runEvaluationCapturingThinking();
expect(seenThinking).toEqual(PBL_V2_TEACHING_THINKING);
});
it('does not leak the disabled config past the turn', async () => {
await runEvaluationCapturingThinking();
expect(thinkingContext.getStore()).toBeUndefined();
});
it('accounts for the turn under its own usage source', async () => {
await runEvaluationCapturingThinking();
await vi.waitFor(() => {
expect(usageMock.recordUsage).toHaveBeenCalledWith(
expect.objectContaining({ kind: 'llm', source: 'pbl-v2-evaluator-task' }),
);
});
});
});