Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
5 changes: 5 additions & 0 deletions .changeset/prompt-prefix-stability.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
'@nanocollective/nanocoder': patch
---

The per-response usage indicator now shows what share of the prompt the provider served from its cache (`9.8k cached (82%)`), and debug logs report prompt prefix stability per request: whether a turn only appended messages, or whether the system prompt, tool schemas or earlier history changed — and for history, the index of the first rewritten message. Local inference backends only reuse their KV cache for the byte-identical prefix of a request, so a prefix break near the start of a 30k-token context means paying full prefill again; these two measurements make that visible instead of showing up only as slow responses. Groundwork for #1574.
6 changes: 6 additions & 0 deletions source/ai-sdk-client/ai-sdk-client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import type {
import {getLogger} from '@/utils/logging';
import {isLocalURL} from '@/utils/url-utils';
import {handleChat} from './chat/chat-handler.js';
import {PrefixTracker} from './chat/prefix-tracker.js';
import {normalizeModelIdForRequest} from './model-id.js';
import {
createProvider,
Expand All @@ -44,6 +45,10 @@ export class AISDKClient implements LLMClient {
current: {},
};

// Per client, matching what a single-slot local server caches: see
// PrefixTracker for why helper calls on this client are tracked too.
private readonly prefixTracker = new PrefixTracker();

constructor(providerConfig: AIProviderConfig) {
const logger = getLogger();

Expand Down Expand Up @@ -208,6 +213,7 @@ export class AISDKClient implements LLMClient {
modeOverrides?.privacySessionMapRef ?? this.privacySessionMapRef,
privacyEnabled: modeOverrides?.privacyEnabled ?? getPrivacyPreference(),
onPrivacyEvent: callbacks.onPrivacyEvent,
prefixTracker: this.prefixTracker,
});
}

Expand Down
36 changes: 36 additions & 0 deletions source/ai-sdk-client/chat/chat-handler.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import type {
import type {LanguageModel} from 'ai';
import {handleChat, stopOnUnknownTool} from './chat-handler.js';
import type {ChatHandlerParams} from './chat-handler.js';
import {PrefixTracker} from './prefix-tracker.js';
import {rehydrateResponse} from './privacy.js';

// Note: This file contains basic structure tests
Expand Down Expand Up @@ -978,6 +979,41 @@ test('handleChat keeps the system string for non-anthropic providers', async t =
t.is(prompt[prompt.length - 1]?.providerOptions, undefined);
});

test('handleChat records each request with the prefix tracker', async t => {
const tracker = new PrefixTracker();
const send = (messages: Message[]) =>
handleChat({
model: capturingModel({}),
currentModel: 'test-model',
providerConfig: {
name: 'TestProvider',
type: 'openai',
models: ['test-model'],
config: {baseURL: 'https://api.test.com'},
},
messages,
tools: {},
callbacks: {},
maxRetries: 0,
prefixTracker: tracker,
});
const system: Message = {role: 'system', content: 'sys'};
await send([system, {role: 'user', content: 'hello'}]);
await send([
system,
{role: 'user', content: 'hello'},
{role: 'assistant', content: 'hi'},
{role: 'user', content: 'more'},
]);

// A third request that drops the first message must diverge at index 0,
// which proves both earlier requests were recorded in send order.
t.deepEqual(
await tracker.record('sys', undefined, [{role: 'user', content: 'more'}]),
{kind: 'history-changed', index: 0, previousMessages: 3},
);
});

// --- Output-token ceiling ---
//
// @ai-sdk/anthropic derives max_tokens from the model id and falls back to
Expand Down
29 changes: 29 additions & 0 deletions source/ai-sdk-client/chat/chat-handler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ import {convertAISDKToolCalls} from '../converters/tool-converter.js';
import {extractRootError} from '../error-handling/error-extractor.js';
import {parseAPIError} from '../error-handling/error-parser.js';
import {isToolSupportError} from '../error-handling/tool-error-detector.js';
import type {PrefixTracker} from './prefix-tracker.js';
import {rehydrateResponse, scrubOutgoing} from './privacy.js';
import {
buildProviderOptions,
Expand All @@ -85,6 +86,8 @@ export interface ChatHandlerParams {
privacySessionMapRef?: React.MutableRefObject<Record<string, string>>;
privacyEnabled?: boolean;
onPrivacyEvent?: (scrubbedDelta: number) => void;
/** Compares each request with the previous one to log prompt-prefix breaks. */
prefixTracker?: PrefixTracker;
}

/**
Expand All @@ -107,6 +110,7 @@ export async function handleChat(
privacySessionMapRef,
privacyEnabled,
onPrivacyEvent,
prefixTracker,
} = params;
const logger = getLogger();

Expand Down Expand Up @@ -206,6 +210,31 @@ export async function handleChat(
? withCacheBreakpoints(convertedMessages, finalSystemContent)
: convertedMessages;

if (prefixTracker) {
try {
// Fingerprint the logical prompt (system + tools + converted
// history), not `modelMessages`: the Anthropic cache breakpoints
// move to the newest message by design, so including them would
// report a break on every single turn.
const divergence = await prefixTracker.record(
finalSystemContent,
aiTools,
convertedMessages,
);
logger.debug('Prompt prefix stability', {
...divergence,
model: currentModel,
correlationId,
provider: providerConfig.name,
});
} catch (error) {
// Diagnostics only: never let fingerprinting fail a request.
logger.debug('Prompt prefix tracking failed', {
error: error instanceof Error ? error.message : String(error),
});
}
}

logger.debug('AI SDK request prepared', {
messageCount: modelMessages.length,
hasSystem: systemContent.length > 0,
Expand Down
146 changes: 146 additions & 0 deletions source/ai-sdk-client/chat/prefix-tracker.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,146 @@
import test from 'ava';
import type {ModelMessage} from 'ai';
import {jsonSchema, tool} from 'ai';
import type {AISDKCoreTool} from '@/types/index';
import {
comparePromptFingerprints,
fingerprintPrompt,
PrefixTracker,
} from './prefix-tracker.js';

console.log('\nprefix-tracker.spec.ts');

function makeTool(description: string): AISDKCoreTool {
return tool({
description,
inputSchema: jsonSchema({
type: 'object',
properties: {path: {type: 'string'}},
}),
});
}

const readFile = makeTool('Read a file');
const writeFile = makeTool('Write a file');

const user = (text: string): ModelMessage => ({role: 'user', content: text});
const assistant = (text: string): ModelMessage => ({
role: 'assistant',
content: text,
});

test('first request has nothing to compare against', async t => {
const tracker = new PrefixTracker();
t.deepEqual(await tracker.record('sys', {readFile}, [user('hi')]), {
kind: 'first',
});
});

test('appending messages keeps the prefix', async t => {
const tracker = new PrefixTracker();
await tracker.record('sys', {readFile}, [user('hi')]);
t.deepEqual(
await tracker.record('sys', {readFile}, [
user('hi'),
assistant('hello'),
user('next'),
]),
{kind: 'append-only', reusedMessages: 1},
);
});

test('a changed system prompt is reported before anything else', async t => {
const tracker = new PrefixTracker();
await tracker.record('sys', {readFile}, [user('hi')]);
t.deepEqual(
await tracker.record('sys + per-turn skills', {writeFile}, [user('bye')]),
{kind: 'system-changed'},
);
});

test('added and removed tools are named', async t => {
const tracker = new PrefixTracker();
await tracker.record('sys', {readFile}, [user('hi')]);
t.deepEqual(await tracker.record('sys', {writeFile}, [user('hi')]), {
kind: 'tools-changed',
added: ['writeFile'],
removed: ['readFile'],
reordered: false,
});
});

test('reordering tools breaks the prefix', async t => {
const tracker = new PrefixTracker();
await tracker.record('sys', {readFile, writeFile}, [user('hi')]);
t.deepEqual(
await tracker.record('sys', {writeFile, readFile}, [user('hi')]),
{kind: 'tools-changed', added: [], removed: [], reordered: true},
);
});

test('a changed tool schema breaks the prefix', async t => {
const tracker = new PrefixTracker();
await tracker.record('sys', {readFile}, [user('hi')]);
t.deepEqual(
await tracker.record('sys', {readFile: makeTool('Read a file v2')}, [
user('hi'),
]),
{kind: 'tools-changed', added: [], removed: [], reordered: false},
);
});

test('rewritten history reports the first changed index', async t => {
const tracker = new PrefixTracker();
await tracker.record('sys', {readFile}, [
user('a'),
assistant('b'),
user('c'),
]);
t.deepEqual(
await tracker.record('sys', {readFile}, [
user('a'),
user('<conversation-summary>b</conversation-summary>'),
user('c'),
]),
{kind: 'history-changed', index: 1, previousMessages: 3},
);
});

test('a dropped leading message (sliding window) diverges at index 0', async t => {
const tracker = new PrefixTracker();
await tracker.record('sys', undefined, [user('a'), user('b')]);
t.deepEqual(await tracker.record('sys', undefined, [user('b'), user('c')]), {
kind: 'history-changed',
index: 0,
previousMessages: 2,
});
});

test('reset forgets the previous request', async t => {
const tracker = new PrefixTracker();
await tracker.record('sys', undefined, [user('a')]);
tracker.reset();
t.deepEqual(await tracker.record('other', undefined, [user('b')]), {
kind: 'first',
});
});

test('fingerprints are deterministic for identical input', async t => {
const a = await fingerprintPrompt('sys', {readFile, writeFile}, [user('x')]);
const b = await fingerprintPrompt('sys', {readFile, writeFile}, [user('x')]);
t.deepEqual(a, b);
t.deepEqual(comparePromptFingerprints(a, b), {
kind: 'append-only',
reusedMessages: 1,
});
});

test('overlapping records are compared in call order', async t => {
const tracker = new PrefixTracker();
const [first, second] = await Promise.all([
tracker.record('sys', {readFile}, [user('a')]),
tracker.record('sys', {readFile}, [user('a'), assistant('b')]),
]);
t.deepEqual(first, {kind: 'first'});
t.deepEqual(second, {kind: 'append-only', reusedMessages: 1});
});
Loading
Loading