Skip to content

Commit 863c3b6

Browse files
authored
fix(core): mark prompt-cache breakpoints for Claude models (#349)
1 parent 64ed3dc commit 863c3b6

8 files changed

Lines changed: 197 additions & 24 deletions

File tree

‎packages/core/src/conversations/turns/context.test.ts‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,12 +21,17 @@ describe("modelPrompt", () => {
2121
"You are Host Agent (@host-agent), an agent in pod Release of workspace Suga",
2222
);
2323
expect(prompt.system).toContain("Check facts carefully.");
24-
// The history, then one trailing instruction with everything that varies
25-
// between turns, so the prefix the provider caches never changes shape.
24+
// The history, marked where it ends, then one trailing instruction with
25+
// everything that varies between turns, so the prefix the provider caches
26+
// never changes shape.
2627
expect(prompt.messages.slice(0, -1)).toEqual([
2728
{ role: "user", content: "Sam (@sam): Check the release" },
2829
{ role: "assistant", content: "I am checking it." },
29-
{ role: "user", content: "Reviewer (@reviewer) (agent): Migration is present." },
30+
{
31+
role: "user",
32+
content: "Reviewer (@reviewer) (agent): Migration is present.",
33+
cacheBreakpoint: true,
34+
},
3035
]);
3136
const instruction = prompt.messages.at(-1);
3237
expect(instruction?.role).toBe("user");

‎packages/core/src/conversations/turns/context.ts‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,9 @@ export interface TurnEnvironment {
3636
* thread now, why this agent has the floor, the turn's environment) goes in one
3737
* trailing instruction after the history, where a change costs nothing. A
3838
* compacted thread starts its history with the Compaction agent's summary, which
39-
* changes only when the thread is compacted again.
39+
* changes only when the thread is compacted again. The history's last message
40+
* is marked as the end of that shared prefix, for providers that cache only
41+
* what a request marks.
4042
*/
4143
export function modelPrompt(
4244
context: TurnContext,
@@ -65,13 +67,19 @@ export function modelPrompt(
6567
.filter(Boolean)
6668
.join("\n\n");
6769

70+
const history = [
71+
...(context.compaction ? [compactionSummary(context.compaction)] : []),
72+
...context.messages
73+
.filter((message) => message.status === "complete")
74+
.flatMap((message) => promptMessagesFor(message, context.agent.id, resultFiles)),
75+
];
76+
const lastInHistory = history.at(-1);
6877
return {
6978
system,
7079
messages: [
71-
...(context.compaction ? [compactionSummary(context.compaction)] : []),
72-
...context.messages
73-
.filter((message) => message.status === "complete")
74-
.flatMap((message) => promptMessagesFor(message, context.agent.id, resultFiles)),
80+
...(lastInHistory
81+
? history.with(-1, { ...lastInHistory, cacheBreakpoint: true as const })
82+
: history),
7583
{ role: "user" as const, content: turnInstruction(context, environment) },
7684
],
7785
};

‎packages/core/src/conversations/turns/repository.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -810,7 +810,11 @@ export const TurnCheckpoint = Schema.Struct({
810810
model: Schema.String,
811811
system: Schema.String,
812812
messages: Schema.Array(
813-
Schema.Struct({ role: Schema.Literals(["user", "assistant"]), content: Schema.String }),
813+
Schema.Struct({
814+
role: Schema.Literals(["user", "assistant"]),
815+
content: Schema.String,
816+
cacheBreakpoint: Schema.optional(Schema.Literal(true)),
817+
}),
814818
),
815819
}),
816820
reply: ReplyDraft,

‎packages/core/src/providers/model-providers/dialects/anthropic.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,15 @@ import {
99
/**
1010
* Anthropic keys with a header of its own and pages its listing, twenty models
1111
* at a time unless asked for more. The listing has no capabilities, so the
12-
* registry supplies them.
12+
* registry supplies them. Its models cache only the prompt prefixes a
13+
* request marks.
1314
*/
1415
export const anthropic: ProviderDialect = {
1516
name: "Anthropic",
1617
discoveryRoot: (baseUrl) => (baseUrl.endsWith("/v1") ? baseUrl : `${baseUrl}/v1`),
1718
listingUrl: (root) => `${root}/models?limit=1000`,
1819
authorization: (apiKey) => ({ "x-api-key": apiKey ?? "", "anthropic-version": "2023-06-01" }),
20+
cacheBreakpoint: () => ({ anthropic: { cacheControl: { type: "ephemeral" } } }),
1921
listing: dataListing,
2022
model: Schema.Struct({
2123
id: Schema.String,

‎packages/core/src/providers/model-providers/dialects/dialect.ts‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import {
55
type ProviderPresetId,
66
providerModelCapabilitySchema,
77
} from "@sugabots/contracts";
8+
import type { ModelMessage } from "ai";
89
import { Data, Effect, Schema, SchemaGetter } from "effect";
910
import type { EgressHttpClient } from "../../network/egress.ts";
1011

@@ -19,8 +20,9 @@ import type { EgressHttpClient } from "../../network/egress.ts";
1920
* to a provider never asks which one it is.
2021
*
2122
* Discovery, how a provider lists its models and what it says about them, is
22-
* what a dialect covers so far. A header a provider wants on every request,
23-
* or an option its completions take, belongs here too as the need arrives.
23+
* most of what a dialect covers; the rest is the options its requests take. A
24+
* header a provider wants on every request belongs here too as the need
25+
* arrives.
2426
*/
2527

2628
export interface DiscoveredModel {
@@ -41,6 +43,9 @@ export interface ProviderIdentity {
4143
apiFormat: ProviderApiFormat;
4244
}
4345

46+
/** Options for one provider's SDK, keyed by the provider they are for. */
47+
export type ProviderOptions = NonNullable<ModelMessage["providerOptions"]>;
48+
4449
export interface ProviderDialect {
4550
/** Named in tests and in what a person is told. */
4651
readonly name: string;
@@ -69,6 +74,13 @@ export interface ProviderDialect {
6974
model: DiscoveredModel,
7075
http: EgressHttpClient,
7176
): Effect.Effect<DiscoveredModel, ModelInspectionFailed>;
77+
/**
78+
* The options that mark the end of a prompt prefix later requests repeat,
79+
* for a model that caches only the prefixes a request marks. Absent, or
80+
* undefined for a model, where the provider caches prefixes by itself or
81+
* not at all.
82+
*/
83+
cacheBreakpoint?(modelId: string): ProviderOptions | undefined;
7284
}
7385

7486
export class ModelInspectionFailed extends Data.TaggedError("ModelInspectionFailed")<{

‎packages/core/src/providers/model-providers/dialects/openrouter.ts‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,15 +13,24 @@ const names = Schema.Array(Schema.String).pipe(
1313
Schema.catchDecoding(() => Effect.succeedSome([])),
1414
);
1515

16+
/** How OpenRouter's ids for Anthropic's models start, as in `anthropic/claude-sonnet-4.5`. */
17+
const ANTHROPIC_MODELS = "anthropic/";
18+
1619
/**
1720
* OpenRouter's listing says what each model takes in, gives out, and accepts
18-
* as parameters, so nothing is left to the registry.
21+
* as parameters, so nothing is left to the registry. Its Anthropic models,
22+
* like Anthropic's own, cache only the prompt prefixes a request marks;
23+
* the rest cache by themselves or not at all.
1924
*/
2025
export const openrouter: ProviderDialect = {
2126
name: "OpenRouter",
2227
// The public /models endpoint accepts invalid keys; /models/user requires authentication.
2328
listingUrl: (root) => `${root}/models/user`,
2429
authorization: openaiCompatible.authorization,
30+
cacheBreakpoint: (modelId) =>
31+
modelId.startsWith(ANTHROPIC_MODELS)
32+
? { openrouter: { cacheControl: { type: "ephemeral" } } }
33+
: undefined,
2534
listing: dataListing,
2635
model: Schema.Struct({
2736
id: Schema.String,

‎packages/core/src/providers/models/models.recording.test.ts‎

Lines changed: 85 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import { describe, expect, it } from "vitest";
44
import type { ModelRequests } from "../../accounting/model-requests.ts";
55
import { effectRunner } from "../../database/database.ts";
66
import { noDatabase } from "../../database/testing.ts";
7+
import type { ProviderIdentity } from "../model-providers/dialects/index.ts";
78
import { Models } from "./models.ts";
89

910
/** The fake provider never queries, so nothing here reaches the database. */
@@ -56,8 +57,17 @@ const reply = () =>
5657
type Recorded = ModelRequests.Started &
5758
Partial<Omit<ModelRequests.Ending, "outcome">> & { outcome: ModelRequests.Outcome };
5859

59-
/** A model over an OpenAI provider that answers with `responses`, one per request. */
60-
function modelAnswering(responses: Array<(recorded: readonly Recorded[]) => Response>) {
60+
const OPENAI = {
61+
preset: "openai",
62+
baseUrl: "https://api.openai.com/v1",
63+
apiFormat: "openai",
64+
} satisfies ProviderIdentity;
65+
66+
/** A model over `provider`, OpenAI's unless given, that answers with `responses`, one per request. */
67+
function modelAnswering(
68+
responses: Array<(recorded: readonly Recorded[]) => Response>,
69+
provider: ProviderIdentity = OPENAI,
70+
) {
6171
const recorded: Recorded[] = [];
6272
/** Each request's body, as the provider was sent it. */
6373
const sent: Array<Record<string, unknown>> = [];
@@ -68,9 +78,7 @@ function modelAnswering(responses: Array<(recorded: readonly Recorded[]) => Resp
6878
resolve: () =>
6979
Effect.succeed({
7080
providerId: PROVIDER_ID,
71-
preset: "openai",
72-
baseUrl: "https://api.openai.com/v1",
73-
apiFormat: "openai",
81+
...provider,
7482
apiKey: "secret",
7583
headers: {},
7684
configurationUpdatedAt: new Date(0),
@@ -208,3 +216,75 @@ describe("a response at its step limit", () => {
208216
]);
209217
});
210218
});
219+
220+
describe("prompt caching", () => {
221+
const marked = { type: "ephemeral" };
222+
const turn: Models.StreamRequest = {
223+
...input,
224+
messages: [
225+
{ role: "user", content: "Earlier." },
226+
{ role: "assistant", content: "Noted.", cacheBreakpoint: true },
227+
{ role: "user", content: "[Turn] Look it up." },
228+
],
229+
};
230+
231+
it("marks the system text, the history's end and each step's newest message for Anthropic", async () => {
232+
const { model, sent } = modelAnswering([() => new Response("Unavailable", { status: 503 })], {
233+
preset: "anthropic",
234+
baseUrl: "https://api.anthropic.com",
235+
apiFormat: "anthropic",
236+
});
237+
238+
await answer(model, { ...turn, model: "claude-test" });
239+
240+
expect(sent[0]).toMatchObject({
241+
system: [{ text: input.system, cache_control: marked }],
242+
messages: [
243+
{ role: "user", content: [{ text: "Earlier." }] },
244+
{ role: "assistant", content: [{ text: "Noted.", cache_control: marked }] },
245+
{ role: "user", content: [{ text: "[Turn] Look it up.", cache_control: marked }] },
246+
],
247+
});
248+
});
249+
250+
it("moves the newest message's mark on each step, keeping the history's", async () => {
251+
const { model, sent } = modelAnswering([toolCall, reply], {
252+
preset: "openrouter",
253+
baseUrl: "https://openrouter.ai/api/v1",
254+
apiFormat: "openai",
255+
});
256+
257+
await answer(model, { ...turn, model: "anthropic/claude-test" });
258+
259+
/** Which of a request's messages are marked, on the message or on one of its parts. */
260+
const markedMessages = (body: Record<string, unknown> = {}) =>
261+
(body.messages as Array<{ content: unknown; cache_control?: unknown }>).map(
262+
(message) =>
263+
message.cache_control !== undefined ||
264+
(Array.isArray(message.content) &&
265+
message.content.some((part) => part.cache_control !== undefined)),
266+
);
267+
// System text, history, turn.
268+
expect(markedMessages(sent[0])).toEqual([true, false, true, true]);
269+
// The turn's mark has moved to the tool result.
270+
expect(markedMessages(sent[1])).toEqual([true, false, true, false, false, true]);
271+
});
272+
273+
it.each([
274+
["OpenAI", OPENAI, "gpt-test"],
275+
[
276+
"OpenRouter's other models",
277+
{ preset: "openrouter", baseUrl: "https://openrouter.ai/api/v1", apiFormat: "openai" },
278+
"openai/gpt-test",
279+
],
280+
] satisfies Array<[string, ProviderIdentity, string]>)(
281+
"leaves requests to %s unmarked, since they cache by themselves",
282+
async (_name, provider, modelId) => {
283+
const { model, sent } = modelAnswering([toolCall, reply], provider);
284+
285+
await answer(model, { ...turn, model: modelId });
286+
287+
expect(JSON.stringify(sent)).not.toContain("cache_control");
288+
},
289+
);
290+
});

‎packages/core/src/providers/models/models.ts‎

Lines changed: 59 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,8 @@ import { ModelRequests } from "../../accounting/model-requests.ts";
2727
import { streamLedger } from "../../accounting/stream-ledger.ts";
2828
import type { Database } from "../../database/database.ts";
2929
import { type UserFacing, UserMessage } from "../../user-message.ts";
30-
import { type ModelRegistry, modelsDev } from "../model-providers/dialects/index.ts";
30+
import type { ProviderOptions } from "../model-providers/dialects/dialect.ts";
31+
import { dialectFor, type ModelRegistry, modelsDev } from "../model-providers/dialects/index.ts";
3132
import { ModelProviderRepository } from "../model-providers/model-provider-repository.ts";
3233
import { withSignInAccess } from "../model-providers/sign-in/sign-in.ts";
3334
import { Egress, type EgressHttpClients } from "../network/egress.ts";
@@ -43,6 +44,13 @@ export interface Prompt {
4344
export interface PromptMessage {
4445
role: "user" | "assistant";
4546
content: string;
47+
/**
48+
* Whether later requests repeat the prompt up to and including this
49+
* message, as a turn's history is repeated by the thread's next turn. A
50+
* provider that caches only the prefixes a request marks is told to cache
51+
* it; one that caches by itself needs no telling.
52+
*/
53+
cacheBreakpoint?: true;
4654
}
4755

4856
/** A response streamed as it arrives, with tools the model may call along the way. */
@@ -326,6 +334,8 @@ export function make({ modelProviders, httpClients, requests, registry }: Option
326334
connection,
327335
});
328336
const codex = connection.preset === "chatgpt";
337+
const breakpoint = dialectFor(connection).cacheBreakpoint?.(input.model);
338+
const markNewest = breakpoint ? newestMessageBreakpoint(breakpoint) : undefined;
329339
// The SDK does not throw a provider's error into the text stream: it
330340
// reports it here and ends the stream, and whatever is asked of the
331341
// result afterwards fails with "No output generated". Keeping the
@@ -350,8 +360,17 @@ export function make({ modelProviders, httpClients, requests, registry }: Option
350360
},
351361
},
352362
}
353-
: { system: input.system }),
354-
messages: [...input.messages, ...(input.continuationMessages ?? [])],
363+
: {
364+
system: breakpoint
365+
? withBreakpoint({ role: "system", content: input.system }, breakpoint)
366+
: input.system,
367+
}),
368+
messages: [
369+
...input.messages.map(({ cacheBreakpoint, ...message }) =>
370+
cacheBreakpoint && breakpoint ? withBreakpoint(message, breakpoint) : message,
371+
),
372+
...(input.continuationMessages ?? []),
373+
],
355374
tools: input.tools,
356375
toolApproval: input.toolApproval,
357376
abortSignal: signal,
@@ -379,9 +398,13 @@ export function make({ modelProviders, httpClients, requests, registry }: Option
379398
// with nothing said, so the last round is told to answer. The note
380399
// rides on the newest tool result, which no provider has cached yet:
381400
// changing the tools or system text instead would cost the cache
382-
// every earlier call built.
383-
prepareStep: ({ stepNumber, messages }) =>
384-
stepNumber === input.maxSteps - 1 ? { messages: withLastStepNote(messages) } : undefined,
401+
// every earlier call built. Every other step marks its newest message
402+
// for the next step to read up to from the cache; the last has no
403+
// next step to read it.
404+
prepareStep: ({ stepNumber, messages }) => {
405+
if (stepNumber === input.maxSteps - 1) return { messages: withLastStepNote(messages) };
406+
return markNewest ? { messages: markNewest(messages) } : undefined;
407+
},
385408
stopWhen: stepCountIs(input.maxSteps),
386409
maxRetries: 0,
387410
});
@@ -489,6 +512,36 @@ function noted(output: ToolResultPart["output"]): ToolResultPart["output"] {
489512
}
490513
}
491514

515+
/**
516+
* Marks each step's newest message with `breakpoint`, taking it off the
517+
* message the step before marked. The SDK sends each step the messages the
518+
* step before was sent, so a mark left behind would count against the four
519+
* Anthropic allows a request.
520+
*/
521+
function newestMessageBreakpoint(breakpoint: ProviderOptions) {
522+
let marked: { index: number; unmarked: ModelMessage } | undefined;
523+
return (messages: ModelMessage[]): ModelMessage[] => {
524+
const restored = marked ? messages.with(marked.index, marked.unmarked) : messages;
525+
const index = restored.length - 1;
526+
const newest = restored[index];
527+
if (!newest) return restored;
528+
marked = { index, unmarked: newest };
529+
return restored.with(index, withBreakpoint(newest, breakpoint));
530+
};
531+
}
532+
533+
/** `message` with `breakpoint`'s options added to whatever options it has for the same provider. */
534+
function withBreakpoint<Message extends ModelMessage>(
535+
message: Message,
536+
breakpoint: ProviderOptions,
537+
): Message {
538+
const providerOptions = { ...message.providerOptions };
539+
for (const [provider, options] of Object.entries(breakpoint)) {
540+
providerOptions[provider] = { ...providerOptions[provider], ...options };
541+
}
542+
return { ...message, providerOptions };
543+
}
544+
492545
function languageModel(
493546
connection: ModelProviderRepository.ProviderEndpoint,
494547
modelId: string,

0 commit comments

Comments
 (0)