Skip to content
This repository was archived by the owner on May 15, 2026. It is now read-only.

Commit 2644752

Browse files
author
James Mtendamema
committed
feat: subscription-gated LLM tool result compression
1 parent 2e3ec13 commit 2644752

30 files changed

Lines changed: 2983 additions & 24 deletions

packages/types/src/global-settings.ts

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -232,10 +232,45 @@ export const globalSettingsSchema = z.object({
232232
* Tools in this list will be excluded from prompt generation and rejected at execution time.
233233
*/
234234
disabledTools: z.array(toolNamesSchema).optional(),
235+
236+
/**
237+
* Settings for the ToolResultProcessor (tool output compression).
238+
* Controls thresholds for when compression kicks in.
239+
*/
240+
toolResultProcessorSettings: z
241+
.object({
242+
/** Master switch — false disables all compression */
243+
enabled: z.boolean(),
244+
/** Compress read_file results above this many characters (default: 1500) */
245+
readFileCharsAbove: z.number(),
246+
/** Compress search_files results above this many matches (default: 20) */
247+
searchMatchesAbove: z.number(),
248+
/** Compress list_files results above this many paths (default: 100) */
249+
listFilesCountAbove: z.number(),
250+
})
251+
.optional(),
252+
253+
/**
254+
* Zoo Code API key for subscription features (smart compression).
255+
* Stored in VS Code SecretStorage — never in global state.
256+
* Generate at https://zoocode.dev/dashboard/api-tokens
257+
*/
258+
zooCodeApiKey: z.string().optional(),
259+
260+
/**
261+
* Zoo Code API base URL for subscription features (smart compression).
262+
* Defaults to "https://zoocode.dev". Override for dev/staging.
263+
*/
264+
zooCodeBaseUrl: z.string().optional(),
235265
})
236266

237267
export type GlobalSettings = z.infer<typeof globalSettingsSchema>
238268

269+
/**
270+
* Settings for the ToolResultProcessor, extracted from GlobalSettings for convenience.
271+
*/
272+
export type ToolResultProcessorSettings = NonNullable<GlobalSettings["toolResultProcessorSettings"]>
273+
239274
export const GLOBAL_SETTINGS_KEYS = globalSettingsSchema.keyof().options
240275

241276
/**
@@ -285,6 +320,7 @@ export const SECRET_STATE_KEYS = [
285320
// Global secrets that are part of GlobalSettings (not ProviderSettings)
286321
export const GLOBAL_SECRET_KEYS = [
287322
"openRouterImageApiKey", // For image generation
323+
"zooCodeApiKey", // Zoo Code subscription API key
288324
] as const
289325

290326
// Type for the actual secret storage keys

src/api/index.ts

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,8 +46,9 @@ export interface ApiHandlerCreateMessageMetadata {
4646
* Task ID used for tracking and provider-specific features:
4747
* - Roo: Sent as X-Roo-Task-ID header
4848
* - Requesty: Sent as trace_id
49+
* Optional — not required for internal compression calls (ZooGatewayApiHandler).
4950
*/
50-
taskId: string
51+
taskId?: string
5152
/**
5253
* Current mode slug for provider-specific tracking:
5354
* - Requesty: Sent in extra metadata
@@ -87,6 +88,12 @@ export interface ApiHandlerCreateMessageMetadata {
8788
* Only applies to providers that support function calling restrictions (e.g., Gemini).
8889
*/
8990
allowedFunctionNames?: string[]
91+
/**
92+
* Name of the tool that produced the result being compressed.
93+
* Used by ZooGatewayApiHandler to pass the originating tool name
94+
* to the compression endpoint for logging and prompt tuning.
95+
*/
96+
toolName?: string
9097
}
9198

9299
export interface ApiHandler {
Lines changed: 320 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,320 @@
1+
// npx vitest run api/providers/__tests__/prompt-caching.spec.ts
2+
3+
/**
4+
* Regression tests for prompt caching across providers.
5+
*
6+
* Findings summary (audited 2026-05-02):
7+
*
8+
* ANTHROPIC:
9+
* - `cache_control: { type: 'ephemeral' }` is set on the system block for all
10+
* cache-capable models (line 123 of anthropic.ts).
11+
* - The `prompt-caching-2024-07-31` beta header is added via the inner IIFE.
12+
* - `cache_creation_input_tokens` and `cache_read_input_tokens` are extracted
13+
* from the `message_start` event and yielded as `cacheWriteTokens` /
14+
* `cacheReadTokens` in the `ApiStreamUsageChunk`.
15+
*
16+
* BEDROCK:
17+
* - Uses AWS-native cachePoint blocks (not `cache_control`) via MultiPointStrategy.
18+
* - `supportsAwsPromptCache()` gates caching on `supportsPromptCache` AND
19+
* non-empty `cachableFields` in model info.
20+
* - Cache token fields (`cacheReadInputTokens`, `cacheWriteInputTokens`, and
21+
* their `*TokenCount` aliases) are captured from the `metadata.usage` stream
22+
* event and yielded as `cacheReadTokens` / `cacheWriteTokens`.
23+
*
24+
* ROO (OpenAI-compatible proxy):
25+
* - No `cache_control` headers — caching is handled server-side by the proxy.
26+
* - Cache metrics surface as `prompt_tokens_details.cached_tokens` (read) and
27+
* `cache_creation_input_tokens` (write) in the final usage chunk.
28+
* - Both are correctly mapped to `cacheReadTokens` / `cacheWriteTokens` in the
29+
* yielded `ApiStreamUsageChunk`.
30+
*
31+
* STREAM TYPE:
32+
* - `ApiStreamUsageChunk` declares `cacheWriteTokens?: number` and
33+
* `cacheReadTokens?: number` — all providers use these fields consistently.
34+
*/
35+
36+
import { AnthropicHandler } from "../anthropic"
37+
import { ApiHandlerOptions } from "../../../shared/api"
38+
39+
// ---------------------------------------------------------------------------
40+
// Shared mock infrastructure
41+
// ---------------------------------------------------------------------------
42+
43+
vitest.mock("@roo-code/telemetry", () => ({
44+
TelemetryService: {
45+
instance: {
46+
captureException: vitest.fn(),
47+
},
48+
},
49+
}))
50+
51+
// ---------------------------------------------------------------------------
52+
// Anthropic SDK mock
53+
// ---------------------------------------------------------------------------
54+
55+
/** Capture what was passed to `messages.create` so tests can assert on it. */
56+
let lastCreateCall: any = undefined
57+
58+
const mockCreate = vitest.fn()
59+
60+
vitest.mock("@anthropic-ai/sdk", () => {
61+
const mockAnthropicConstructor = vitest.fn().mockImplementation(() => ({
62+
messages: {
63+
create: mockCreate,
64+
},
65+
}))
66+
return { Anthropic: mockAnthropicConstructor }
67+
})
68+
69+
// ---------------------------------------------------------------------------
70+
// Helper: build a minimal streaming response
71+
// ---------------------------------------------------------------------------
72+
73+
function makeAnthropicStream(cacheCreationTokens: number | undefined, cacheReadTokens: number | undefined) {
74+
return {
75+
async *[Symbol.asyncIterator]() {
76+
yield {
77+
type: "message_start",
78+
message: {
79+
usage: {
80+
input_tokens: 100,
81+
output_tokens: 0,
82+
cache_creation_input_tokens: cacheCreationTokens,
83+
cache_read_input_tokens: cacheReadTokens,
84+
},
85+
},
86+
}
87+
yield {
88+
type: "content_block_start",
89+
index: 0,
90+
content_block: { type: "text", text: "hi" },
91+
}
92+
yield {
93+
type: "message_delta",
94+
usage: { output_tokens: 5 },
95+
}
96+
yield { type: "message_stop" }
97+
},
98+
}
99+
}
100+
101+
// ---------------------------------------------------------------------------
102+
// Tests: Anthropic provider – system prompt cache_control
103+
// ---------------------------------------------------------------------------
104+
105+
describe("AnthropicHandler – prompt caching", () => {
106+
const baseOptions: ApiHandlerOptions = {
107+
apiKey: "test-key",
108+
apiModelId: "claude-3-5-sonnet-20241022",
109+
}
110+
111+
beforeEach(() => {
112+
vitest.clearAllMocks()
113+
lastCreateCall = undefined
114+
})
115+
116+
describe("system prompt cache_control", () => {
117+
it("sets cache_control: { type: 'ephemeral' } on the system block for cache-capable models", async () => {
118+
mockCreate.mockReturnValue(makeAnthropicStream(500, 0))
119+
120+
const handler = new AnthropicHandler(baseOptions)
121+
const gen = handler.createMessage("System prompt here", [{ role: "user", content: "Hello" }])
122+
123+
// Drain the generator.
124+
for await (const _ of gen) {
125+
// consume
126+
}
127+
128+
expect(mockCreate).toHaveBeenCalledOnce()
129+
const callArgs = mockCreate.mock.calls[0][0]
130+
131+
// System must be an array with exactly one block.
132+
expect(Array.isArray(callArgs.system)).toBe(true)
133+
expect(callArgs.system).toHaveLength(1)
134+
135+
const systemBlock = callArgs.system[0]
136+
expect(systemBlock.type).toBe("text")
137+
expect(systemBlock.text).toBe("System prompt here")
138+
// THE KEY ASSERTION: cache_control must be ephemeral.
139+
expect(systemBlock.cache_control).toEqual({ type: "ephemeral" })
140+
})
141+
142+
it("includes the prompt-caching beta header for cache-capable models", async () => {
143+
mockCreate.mockReturnValue(makeAnthropicStream(0, 0))
144+
145+
const handler = new AnthropicHandler(baseOptions)
146+
const gen = handler.createMessage("System prompt", [{ role: "user", content: "Hi" }])
147+
for await (const _ of gen) {
148+
/* drain */
149+
}
150+
151+
expect(mockCreate).toHaveBeenCalledOnce()
152+
// The second argument to create() is the request options (headers).
153+
const requestOptions = mockCreate.mock.calls[0][1]
154+
const betaHeader: string = requestOptions?.headers?.["anthropic-beta"] ?? ""
155+
expect(betaHeader).toContain("prompt-caching-2024-07-31")
156+
})
157+
158+
it("falls back to the default model (which is cache-capable) when an unknown model ID is supplied", async () => {
159+
// `getModel()` maps unknown IDs to `anthropicDefaultModelId`, which IS in the
160+
// cache-capable switch list. Therefore cache_control is always applied even for
161+
// unknown/legacy model strings, because the fallback model supports caching.
162+
const unknownModelOptions: ApiHandlerOptions = {
163+
...baseOptions,
164+
apiModelId: "claude-ancient-unsupported-model" as any,
165+
}
166+
167+
mockCreate.mockReturnValue(makeAnthropicStream(0, 0))
168+
169+
const handler = new AnthropicHandler(unknownModelOptions)
170+
for await (const _ of handler.createMessage("System", [{ role: "user", content: "Hi" }])) {
171+
/* drain */
172+
}
173+
174+
const callArgs = mockCreate.mock.calls[0][0]
175+
const systemBlock = callArgs.system[0]
176+
177+
// The fallback model is cache-capable, so cache_control IS set.
178+
// This is the correct / expected behaviour – the default branch is only
179+
// reached when the model ID exactly matches a known non-cached model,
180+
// which currently does not exist in the active anthropicModels list.
181+
expect(systemBlock.type).toBe("text")
182+
expect(systemBlock.text).toBe("System")
183+
// cache_control is present because the fallback model supports caching.
184+
expect(systemBlock.cache_control).toEqual({ type: "ephemeral" })
185+
})
186+
})
187+
188+
// ---------------------------------------------------------------------------
189+
// Stream processing: cache metric capture
190+
// ---------------------------------------------------------------------------
191+
192+
describe("stream processing – cache metric capture", () => {
193+
it("yields cacheWriteTokens from cache_creation_input_tokens in message_start", async () => {
194+
mockCreate.mockReturnValue(makeAnthropicStream(1234, 0))
195+
196+
const handler = new AnthropicHandler(baseOptions)
197+
const chunks: any[] = []
198+
for await (const chunk of handler.createMessage("Sys", [{ role: "user", content: "Hi" }])) {
199+
chunks.push(chunk)
200+
}
201+
202+
const usageChunks = chunks.filter((c) => c.type === "usage")
203+
expect(usageChunks.length).toBeGreaterThan(0)
204+
205+
// The first usage chunk (from message_start) carries cacheWriteTokens.
206+
const firstUsage = usageChunks[0]
207+
expect(firstUsage.cacheWriteTokens).toBe(1234)
208+
})
209+
210+
it("yields cacheReadTokens from cache_read_input_tokens in message_start", async () => {
211+
mockCreate.mockReturnValue(makeAnthropicStream(0, 567))
212+
213+
const handler = new AnthropicHandler(baseOptions)
214+
const chunks: any[] = []
215+
for await (const chunk of handler.createMessage("Sys", [{ role: "user", content: "Hi" }])) {
216+
chunks.push(chunk)
217+
}
218+
219+
const usageChunks = chunks.filter((c) => c.type === "usage")
220+
expect(usageChunks.length).toBeGreaterThan(0)
221+
222+
const firstUsage = usageChunks[0]
223+
expect(firstUsage.cacheReadTokens).toBe(567)
224+
})
225+
226+
it("yields both cacheWriteTokens and cacheReadTokens when both are present", async () => {
227+
mockCreate.mockReturnValue(makeAnthropicStream(800, 200))
228+
229+
const handler = new AnthropicHandler(baseOptions)
230+
const chunks: any[] = []
231+
for await (const chunk of handler.createMessage("Sys", [{ role: "user", content: "Hi" }])) {
232+
chunks.push(chunk)
233+
}
234+
235+
const firstUsage = chunks.find((c) => c.type === "usage" && c.cacheWriteTokens !== undefined)
236+
expect(firstUsage).toBeDefined()
237+
expect(firstUsage.cacheWriteTokens).toBe(800)
238+
expect(firstUsage.cacheReadTokens).toBe(200)
239+
})
240+
241+
it("omits cacheWriteTokens when cache_creation_input_tokens is 0 (falsy → undefined)", async () => {
242+
mockCreate.mockReturnValue(makeAnthropicStream(0, 0))
243+
244+
const handler = new AnthropicHandler(baseOptions)
245+
const chunks: any[] = []
246+
for await (const chunk of handler.createMessage("Sys", [{ role: "user", content: "Hi" }])) {
247+
chunks.push(chunk)
248+
}
249+
250+
const firstUsage = chunks.find((c) => c.type === "usage")
251+
// 0 is falsy so the handler maps it to `undefined`.
252+
expect(firstUsage?.cacheWriteTokens).toBeUndefined()
253+
expect(firstUsage?.cacheReadTokens).toBeUndefined()
254+
})
255+
})
256+
257+
// ---------------------------------------------------------------------------
258+
// cache_control on user messages
259+
// ---------------------------------------------------------------------------
260+
261+
describe("user message cache markers", () => {
262+
it("attaches cache_control to the last and second-to-last user messages", async () => {
263+
mockCreate.mockReturnValue(makeAnthropicStream(0, 0))
264+
265+
const handler = new AnthropicHandler(baseOptions)
266+
267+
const messages: any[] = [
268+
{ role: "user", content: "First user message" },
269+
{ role: "assistant", content: "First assistant reply" },
270+
{ role: "user", content: "Second user message" },
271+
]
272+
273+
for await (const _ of handler.createMessage("System", messages)) {
274+
/* drain */
275+
}
276+
277+
const callArgs = mockCreate.mock.calls[0][0]
278+
const sentMessages: any[] = callArgs.messages
279+
280+
// Message indices: 0 = user, 1 = assistant, 2 = user
281+
// userMsgIndices = [0, 2] → last=2, secondLast=0
282+
const lastUserMsg = sentMessages[2]
283+
const secondLastUserMsg = sentMessages[0]
284+
285+
// Last user message content should be an array with cache_control on the last block.
286+
const lastContent = Array.isArray(lastUserMsg.content) ? lastUserMsg.content : null
287+
expect(lastContent).not.toBeNull()
288+
const lastBlock = lastContent![lastContent!.length - 1]
289+
expect(lastBlock.cache_control).toEqual({ type: "ephemeral" })
290+
291+
// Second-to-last user message should also carry cache_control.
292+
const secondLastContent = Array.isArray(secondLastUserMsg.content) ? secondLastUserMsg.content : null
293+
expect(secondLastContent).not.toBeNull()
294+
const secondLastBlock = secondLastContent![secondLastContent!.length - 1]
295+
expect(secondLastBlock.cache_control).toEqual({ type: "ephemeral" })
296+
})
297+
})
298+
})
299+
300+
// ---------------------------------------------------------------------------
301+
// Tests: ApiStreamUsageChunk type completeness (static/compile-time guard)
302+
// ---------------------------------------------------------------------------
303+
304+
describe("ApiStreamUsageChunk – type completeness", () => {
305+
it("declares cacheWriteTokens and cacheReadTokens fields", () => {
306+
// This is a compile-time / shape test. If the fields are removed from the
307+
// type definition the TypeScript compilation will fail here.
308+
const chunk = {
309+
type: "usage" as const,
310+
inputTokens: 100,
311+
outputTokens: 50,
312+
cacheWriteTokens: 20,
313+
cacheReadTokens: 10,
314+
}
315+
316+
// Runtime assertion as a belt-and-suspenders guard.
317+
expect(chunk).toHaveProperty("cacheWriteTokens", 20)
318+
expect(chunk).toHaveProperty("cacheReadTokens", 10)
319+
})
320+
})

0 commit comments

Comments
 (0)