From 09ba1d63e2fe3a60db0ee085b878d3317cfa5580 Mon Sep 17 00:00:00 2001 From: lizhixuan Date: Mon, 13 Jul 2026 13:18:43 +0800 Subject: [PATCH] fix: session isolation, Gemini tool-call correlation, compression failure hardening MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes from a full architecture review of the three packages: core: - fromGemini correlates functionResponses via per-name FIFO queues spanning the conversation; repeat calls of the same tool no longer produce duplicate IDs or misattributed results - compression-model failures no longer inject the raw error into the LLM-bound fallback summary (logged via `logger` instead) - onCompress and onBeforeCompress share one degradation stance: a throwing hook is caught and logged, never fails compile() - new SessionPool (LRU, validated maxSize) + normalizeSessionKey / DEFAULT_SESSION_KEY / dedupeConstructionWarnings - new flattenForCompression export — canonical role-flattening for the summarizeHistory compress-callback contract - OpenAIAdapter.compile clones without the JSON round-trip (honors toJSON) - placeholder vocabulary centralized in Prompts ai-sdk-middleware / tanstack-ai: - Janitor state is now per-session (providerOptions.contextChef.sessionId / ctx.conversationId) instead of shared across every conversation the middleware serves; invalid session keys warn once and route to the default slot with unified semantics across both packages - construction-time config nags fire once per middleware, not per session - non-text tool-result parts flatten to typed placeholders instead of being silently dropped - compression flattening reuses core's flattenForCompression - peerDependencies tightened to ^4.0.0 / ^7.0.0 (ai-sdk-middleware) --- .changeset/great-wolves-shake.md | 11 ++ .changeset/lucky-pianos-guess.md | 7 + .changeset/tidy-moons-search.md | 14 ++ packages/ai-sdk-middleware/README.md | 20 +++ packages/ai-sdk-middleware/package.json | 4 +- packages/ai-sdk-middleware/src/adapter.ts | 12 +- packages/ai-sdk-middleware/src/middleware.ts | 91 ++++++----- packages/ai-sdk-middleware/src/truncator.ts | 10 +- packages/ai-sdk-middleware/src/types.ts | 10 ++ .../ai-sdk-middleware/tests/adapter.test.ts | 28 +++- .../tests/middleware.test.ts | 121 +++++++++++++++ .../core/src/adapters/geminiAdapter.test.ts | 58 +++++++ packages/core/src/adapters/geminiAdapter.ts | 29 +++- .../core/src/adapters/openAIAdapter.test.ts | 14 ++ packages/core/src/adapters/openAIAdapter.ts | 35 ++++- packages/core/src/index.ts | 7 + .../core/src/modules/janitor/index.test.ts | 135 ++++++++++++++++- packages/core/src/modules/janitor/index.ts | 142 ++++++++++++------ packages/core/src/prompts.ts | 43 ++++++ packages/core/src/utils/sessionPool.test.ts | 97 ++++++++++++ packages/core/src/utils/sessionPool.ts | 106 +++++++++++++ packages/tanstack-ai/README.md | 4 + packages/tanstack-ai/src/middleware.ts | 101 +++++++------ packages/tanstack-ai/src/types.ts | 9 ++ packages/tanstack-ai/tests/middleware.test.ts | 72 +++++++++ 25 files changed, 1029 insertions(+), 151 deletions(-) create mode 100644 .changeset/great-wolves-shake.md create mode 100644 .changeset/lucky-pianos-guess.md create mode 100644 .changeset/tidy-moons-search.md create mode 100644 packages/core/src/utils/sessionPool.test.ts create mode 100644 packages/core/src/utils/sessionPool.ts diff --git a/.changeset/great-wolves-shake.md b/.changeset/great-wolves-shake.md new file mode 100644 index 0000000..f27aaec --- /dev/null +++ b/.changeset/great-wolves-shake.md @@ -0,0 +1,11 @@ +--- +"@context-chef/ai-sdk-middleware": minor +--- + +Session isolation, lossless-er tool output flattening, and tighter peer ranges. + +- **Session isolation**: compression state (fed token usage, compression suppression, the failure circuit breaker) is now tracked per session instead of per middleware instance. Pass `providerOptions: { contextChef: { sessionId } }` on each call; calls without a `sessionId` share one default session (prior behavior). New `maxSessions` option caps concurrently tracked sessions (default 256, LRU-evicted). Previously a middleware created at module scope leaked Janitor state across every conversation it served. +- `stringifyToolOutput` no longer silently drops non-text content parts: file/media parts flatten to a `[tool result file: ]` placeholder so compression and truncation see that they exist. The truncator's duplicate text extraction now delegates to the same implementation. +- Tighten peerDependency ranges from `>=4` / `>=7` to `^4.0.0` / `^7.0.0`. The middleware depends on `LanguageModelV4*` type shapes from `@ai-sdk/provider` v4 — an unbounded range would let a future major install silently and fail at runtime. +- Compression flattening now reuses `flattenForCompression` from `@context-chef/core` (behavior unchanged). +- An invalid `sessionId` (empty string, non-string, or a malformed `contextChef` namespace) now warns once and routes to the default session instead of failing silently. The Janitor missing-compression-config nag fires once per middleware instead of once per session. diff --git a/.changeset/lucky-pianos-guess.md b/.changeset/lucky-pianos-guess.md new file mode 100644 index 0000000..f0537d0 --- /dev/null +++ b/.changeset/lucky-pianos-guess.md @@ -0,0 +1,7 @@ +--- +"@context-chef/tanstack-ai": minor +--- + +Conversation isolation: compression state (fed token usage, compression suppression, the failure circuit breaker) is now tracked per `ctx.conversationId` instead of per middleware instance. Calls without a `conversationId` share one default slot (prior behavior). New `maxSessions` option caps concurrently tracked conversations (default 256, LRU-evicted). Previously a middleware instance reused across chat() calls leaked Janitor state across every conversation it served. Compression flattening now reuses `flattenForCompression` from `@context-chef/core` (behavior unchanged). + +An empty-string `conversationId` now warns once and routes to the default slot (previously it was silently treated as a distinct conversation, diverging from the AI SDK middleware's semantics). The Janitor missing-compression-config nag fires once per middleware instead of once per conversation. diff --git a/.changeset/tidy-moons-search.md b/.changeset/tidy-moons-search.md new file mode 100644 index 0000000..847f1b3 --- /dev/null +++ b/.changeset/tidy-moons-search.md @@ -0,0 +1,14 @@ +--- +"@context-chef/core": minor +--- + +Fix Gemini tool call correlation, harden Janitor compression failure paths, and add shared middleware infrastructure. + +- `fromGemini` now synthesizes tool call IDs with per-name counters spanning the whole conversation and correlates each `functionResponse` to the oldest unconsumed call of the same name (FIFO). Previously the correlation state was scoped to a single message, so responses arriving in a later content entry always fell back to the `-0` suffix — repeat calls of the same tool produced duplicate IDs and misattributed results. +- When the compression model throws, the raw error is no longer appended to the LLM-bound fallback summary. It now goes to the configured `logger` (default `console`) instead, keeping stack traces out of model context. +- A throwing/rejecting `onCompress` hook no longer aborts `compile()`. It is caught and logged via `logger`; the compression result is kept. Hooks that follow the documented "must not throw" contract see no behavior change. +- `onBeforeCompress` (and the deprecated `onBudgetExceeded`) now degrade the same way: a throwing hook is caught, logged via `logger`, and treated as if it returned `null` — default compression proceeds. Both hooks now share one failure stance: a broken hook never fails `compile()`. +- The text-placeholder vocabulary is centralized in `Prompts`: new `getAttachmentPlaceholder` (previously janitor-internal), `getToolResultFilePlaceholder`, and `getToolResultPartPlaceholder`. Formats are unchanged — this gives each convention a single source instead of scattered string literals. +- New export `SessionPool` — a keyed instance pool with LRU eviction, used by the middleware packages to hold one Janitor per conversation. `maxSize` is validated (throws `RangeError` unless a positive integer — a non-positive cap would silently evict every entry on insert, disabling pooling). Companion exports `DEFAULT_SESSION_KEY`, `normalizeSessionKey`, and `dedupeConstructionWarnings` back the middlewares' shared session-key normalization and construction-nag dedupe. +- New export `flattenForCompression(messages)` — the canonical role-flattening implementation required by the `summarizeHistory` compress-callback contract (tool results → user messages, tool calls → assistant text). +- `OpenAIAdapter.compile` replaces its per-message `JSON.parse(JSON.stringify(...))` deep clone with a direct clone that skips `undefined` properties — same output, no string serialization detour (noticeable with base64 attachments). The clone honors `toJSON`, so a `Date` still lands as its ISO string. diff --git a/packages/ai-sdk-middleware/README.md b/packages/ai-sdk-middleware/README.md index 51e3878..655aaa5 100644 --- a/packages/ai-sdk-middleware/README.md +++ b/packages/ai-sdk-middleware/README.md @@ -126,6 +126,26 @@ Tools not listed fall back to the top-level defaults. The lookup key is `tool-re The middleware automatically extracts token usage from `generateText` and `streamText` responses and feeds it back to the compression engine. No manual `reportTokenUsage()` calls needed. +### Session Isolation (Multi-User Servers) + +A wrapped model is usually created once at module scope and reused across requests. Compression state (fed token usage, compression suppression, the failure circuit breaker) is tracked per **session** — pass a `sessionId` through `providerOptions` on each call so concurrent conversations never share state: + +```typescript +const model = withContextChef(openai('gpt-4o'), { + contextWindow: 128_000, + compress: { model: openai('gpt-4o-mini') }, +}); + +// Per request: +const result = await generateText({ + model, + messages, + providerOptions: { contextChef: { sessionId: conversationId } }, +}); +``` + +Calls without a `sessionId` share one default session — fine for single-conversation processes (a CLI, a notebook, one agent loop), wrong for multi-user servers: one user's over-budget conversation would trigger or suppress compression for everyone else. Up to `maxSessions` sessions are tracked concurrently (default 256, LRU-evicted); an evicted session is transparently recreated on next access, losing only its fed token usage. + ### Compact (Mechanical Pruning) Zero-LLM-cost message pruning via AI SDK's `pruneMessages` — removes reasoning, tool calls, and empty messages: diff --git a/packages/ai-sdk-middleware/package.json b/packages/ai-sdk-middleware/package.json index 7e12695..e0a0351 100644 --- a/packages/ai-sdk-middleware/package.json +++ b/packages/ai-sdk-middleware/package.json @@ -47,8 +47,8 @@ "@context-chef/core": "workspace:*" }, "peerDependencies": { - "@ai-sdk/provider": ">=4", - "ai": ">=7" + "@ai-sdk/provider": "^4.0.0", + "ai": "^7.0.0" }, "devDependencies": { "@ai-sdk/provider": "^4.0.0", diff --git a/packages/ai-sdk-middleware/src/adapter.ts b/packages/ai-sdk-middleware/src/adapter.ts index b35c7fc..680b73a 100644 --- a/packages/ai-sdk-middleware/src/adapter.ts +++ b/packages/ai-sdk-middleware/src/adapter.ts @@ -10,6 +10,7 @@ import { type Attachment, ensureValidHistory, type Message, + Prompts, type ToolCall, } from '@context-chef/core'; @@ -314,7 +315,16 @@ export function stringifyToolOutput(output: LanguageModelV4ToolResultOutput): st return JSON.stringify(output.value); case 'content': return output.value - .map((v) => (v.type === 'text' ? v.text : '')) + .map((v) => { + if (v.type === 'text') return v.text; + // Non-text parts (files/media) must leave a trace in the flattened + // text — silently dropping them would break the round-trip claim + // and hide the part from compression/truncation entirely. + if (v.type === 'file') { + return Prompts.getToolResultFilePlaceholder(v.mediaType, v.filename); + } + return Prompts.getToolResultPartPlaceholder((v as { type: string }).type); + }) .filter(Boolean) .join('\n'); default: diff --git a/packages/ai-sdk-middleware/src/middleware.ts b/packages/ai-sdk-middleware/src/middleware.ts index 8d75220..f608bc2 100644 --- a/packages/ai-sdk-middleware/src/middleware.ts +++ b/packages/ai-sdk-middleware/src/middleware.ts @@ -8,9 +8,14 @@ import { type ChefLogger, type CompressionDetails, compactMessages, + DEFAULT_SESSION_KEY, + dedupeConstructionWarnings, + flattenForCompression, Janitor, type Message, + normalizeSessionKey, Prompts, + SessionPool, type SummarizeHistoryOptions, summarizeHistory, XmlGenerator, @@ -28,8 +33,6 @@ import { fromModelMessages } from './modelMessageAdapter'; import { truncateToolResults } from './truncator'; import type { ContextChefOptions, DynamicStateConfig } from './types'; -type CompressRole = 'system' | 'user' | 'assistant'; - /** * After this many compressions fire without an `onCompress` persistence hook, * warn once. The middleware compresses in-flight only — it never mutates the @@ -92,10 +95,50 @@ export function createMiddleware(options: ContextChefOptions): LanguageModelMidd ); }; - const janitor = budgeting - ? createJanitor(options, options.contextWindow as number, logger, onCompressionFired) + // One Janitor per session. A middleware instance is usually created once at + // module scope (`const model = withContextChef(...)`) but serves many + // conversations — sharing one Janitor would leak token-usage feeds, + // compression suppression, and circuit-breaker counts across them. Callers + // opt in per call via `providerOptions: { contextChef: { sessionId } }`; + // calls without a sessionId share the default session (prior behavior). + // Construction-time config nags are deduped across sessions — the config + // is identical for every pooled Janitor, so once is enough. + const janitors = budgeting + ? new SessionPool( + dedupeConstructionWarnings(logger, (constructionLogger) => + createJanitor( + options, + options.contextWindow as number, + constructionLogger, + onCompressionFired, + ), + ), + { maxSize: options.maxSessions }, + ) : null; + let invalidSessionKeyWarned = false; + const flagInvalidSessionKey = (raw: unknown) => { + if (invalidSessionKeyWarned) return; + invalidSessionKeyWarned = true; + logger.warn( + '[context-chef] Invalid providerOptions.contextChef sessionId (expected a non-empty ' + + `string, got ${raw === '' ? 'empty string' : typeof raw}); routing to the default session.`, + ); + }; + + const janitorFor = (params: { providerOptions?: Record }): Janitor | null => { + if (!janitors) return null; + const ns = params.providerOptions?.contextChef; + if (ns != null && typeof ns !== 'object') { + // Malformed namespace (e.g. contextChef: 'abc') — never a session key. + flagInvalidSessionKey(ns); + return janitors.get(DEFAULT_SESSION_KEY); + } + const raw = ns ? (ns as Record).sessionId : undefined; + return janitors.get(normalizeSessionKey(raw, flagInvalidSessionKey)); + }; + const clearsToolResults = !!options.clear?.some( (t) => t === 'tool-result' || (typeof t === 'object' && t.target === 'tool-result'), ); @@ -116,6 +159,7 @@ export function createMiddleware(options: ContextChefOptions): LanguageModelMidd specificationVersion: 'v4', transformParams: async ({ params }) => { + const janitor = janitorFor(params); let { prompt } = params; // 1. Truncate large tool results @@ -176,9 +220,10 @@ export function createMiddleware(options: ContextChefOptions): LanguageModelMidd return { ...params, prompt }; }, - wrapGenerate: async ({ doGenerate }) => { + wrapGenerate: async ({ doGenerate, params }) => { const result = await doGenerate(); + const janitor = janitorFor(params); if (!janitor) return result; if (result.usage?.inputTokens?.total != null) { @@ -195,7 +240,8 @@ export function createMiddleware(options: ContextChefOptions): LanguageModelMidd return result; }, - wrapStream: async ({ doStream }) => { + wrapStream: async ({ doStream, params }) => { + const janitor = janitorFor(params); if (!janitor) return doStream(); const { stream, ...rest } = await doStream(); @@ -371,15 +417,6 @@ async function injectDynamicState( return result; } -/** - * Maps an IR role to a role accepted by generateText. - * Tool messages are handled separately before this is called. - */ -function toCompressRole(role: string): CompressRole { - if (role === 'system' || role === 'user' || role === 'assistant') return role; - return 'user'; -} - /** * Adapts an AI SDK LanguageModelV4 into the compressionModel callback * that Janitor expects: (messages: Message[]) => Promise @@ -391,31 +428,9 @@ export function createCompressionAdapter( model: LanguageModel, ): (messages: Message[]) => Promise { return async (messages: Message[]): Promise => { - const formatted = messages.map((m): { role: CompressRole; content: string } => { - if (m.role === 'tool') { - return { - role: 'user' satisfies CompressRole, - content: `[Tool result${m.tool_call_id ? ` (${m.tool_call_id})` : ''}: ${m.content}]`, - }; - } - if (m.role === 'assistant' && m.tool_calls?.length) { - const toolCallsDesc = m.tool_calls - .map((tc) => `[Called tool: ${tc.function.name}(${tc.function.arguments})]`) - .join('\n'); - return { - role: 'assistant' satisfies CompressRole, - content: m.content ? `${m.content}\n${toolCallsDesc}` : toolCallsDesc, - }; - } - return { - role: toCompressRole(m.role), - content: m.content, - }; - }); - const { text } = await generateText({ model, - messages: formatted, + messages: flattenForCompression(messages), maxOutputTokens: 2048, }); diff --git a/packages/ai-sdk-middleware/src/truncator.ts b/packages/ai-sdk-middleware/src/truncator.ts index d0c0218..fb406c6 100644 --- a/packages/ai-sdk-middleware/src/truncator.ts +++ b/packages/ai-sdk-middleware/src/truncator.ts @@ -4,6 +4,7 @@ import type { LanguageModelV4ToolResultPart, } from '@ai-sdk/provider'; import { type ChefLogger, Offloader } from '@context-chef/core'; +import { stringifyToolOutput } from './adapter'; import type { TruncateOptions } from './types'; /** @@ -141,16 +142,13 @@ function extractText(output: LanguageModelV4ToolResultOutput): string { switch (output.type) { case 'text': case 'error-text': - return output.value; case 'json': case 'error-json': - return JSON.stringify(output.value); case 'content': - return output.value - .map((v: { type: string; text?: string }) => (v.type === 'text' ? (v.text ?? '') : '')) - .filter(Boolean) - .join('\n'); + return stringifyToolOutput(output); default: + // Marker outputs (e.g. execution-denied) measure as empty — they must + // never be rewritten by truncation. return ''; } } diff --git a/packages/ai-sdk-middleware/src/types.ts b/packages/ai-sdk-middleware/src/types.ts index 7ed27d6..fa744b0 100644 --- a/packages/ai-sdk-middleware/src/types.ts +++ b/packages/ai-sdk-middleware/src/types.ts @@ -257,4 +257,14 @@ export interface ContextChefOptions { * underlying Janitor and Offloader. */ logger?: ChefLogger; + /** + * Cap on concurrently tracked sessions when session isolation is used + * (see the `providerOptions.contextChef.sessionId` docs on the package + * README). Each session gets its own Janitor so token-usage feeds and + * compression state never leak across conversations. Least-recently-used + * sessions beyond the cap are dropped and transparently recreated on next + * access. Must be a positive integer — the pool throws a RangeError + * otherwise. Default: 256. + */ + maxSessions?: number; } diff --git a/packages/ai-sdk-middleware/tests/adapter.test.ts b/packages/ai-sdk-middleware/tests/adapter.test.ts index d762eff..febcbdc 100644 --- a/packages/ai-sdk-middleware/tests/adapter.test.ts +++ b/packages/ai-sdk-middleware/tests/adapter.test.ts @@ -1,7 +1,31 @@ -import type { LanguageModelV4Prompt } from '@ai-sdk/provider'; +import type { LanguageModelV4Prompt, LanguageModelV4ToolResultOutput } from '@ai-sdk/provider'; import type { Message } from '@context-chef/core'; import { describe, expect, it } from 'vitest'; -import { fromAISDK, toAISDK } from '../src/adapter'; +import { fromAISDK, stringifyToolOutput, toAISDK } from '../src/adapter'; + +describe('stringifyToolOutput', () => { + it('keeps a typed placeholder for non-text content parts instead of dropping them', () => { + const output: LanguageModelV4ToolResultOutput = { + type: 'content', + value: [ + { type: 'text', text: 'Chart generated.' }, + { + type: 'file', + data: { type: 'url', url: new URL('https://example.com/chart.png') }, + mediaType: 'image/png', + filename: 'chart.png', + }, + ], + }; + + const text = stringifyToolOutput(output); + + expect(text).toContain('Chart generated.'); + // The file part must leave a trace — not vanish silently. + expect(text).toContain('image/png'); + expect(text).toContain('chart.png'); + }); +}); describe('fromAISDK', () => { it('converts system messages', () => { diff --git a/packages/ai-sdk-middleware/tests/middleware.test.ts b/packages/ai-sdk-middleware/tests/middleware.test.ts index 05d69fa..6e0cee3 100644 --- a/packages/ai-sdk-middleware/tests/middleware.test.ts +++ b/packages/ai-sdk-middleware/tests/middleware.test.ts @@ -280,6 +280,127 @@ describe('createMiddleware', () => { expect(result.prompt.length).toBeLessThan(longPrompt.length); }); + it('isolates janitor state between sessions via providerOptions.contextChef.sessionId', async () => { + // Window large enough that the heuristic estimate of the short prompt + // stays under budget — only the FED usage (session A's) exceeds it. + const middleware = createMiddleware({ + contextWindow: 10_000, + onCompress: vi.fn(), + }); + + const model = createMockModel({ inputTokens: 20_000 }); + const doGenerate = (): PromiseLike => + model.doGenerate({ prompt: [] }); + const doStream = (): PromiseLike => model.doStream({ prompt: [] }); + + const sessionA = { contextChef: { sessionId: 'user-a' } }; + const sessionB = { contextChef: { sessionId: 'user-b' } }; + + // Session A exceeds the budget (20_000 > 10_000). + await assertDefined( + middleware.wrapGenerate, + 'wrapGenerate', + )({ + doGenerate, + doStream, + params: { prompt: [], providerOptions: sessionA }, + model, + }); + + const longPrompt = makeConversation(10); + + // Session B must NOT inherit A's fed token usage — no compression. + const resultB = await assertDefined( + middleware.transformParams, + 'transformParams', + )({ + params: { prompt: longPrompt, providerOptions: sessionB }, + type: 'generate', + model, + }); + expect(resultB.prompt.length).toBe(longPrompt.length); + + // Session A compresses on its own next call. + const resultA = await assertDefined( + middleware.transformParams, + 'transformParams', + )({ + params: { prompt: longPrompt, providerOptions: sessionA }, + type: 'generate', + model, + }); + expect(resultA.prompt.length).toBeLessThan(longPrompt.length); + }); + + it('routes an invalid sessionId to the default session and warns once', async () => { + const warn = vi.fn(); + const middleware = createMiddleware({ + contextWindow: 10_000, + onCompress: vi.fn(), + logger: { warn }, + }); + + const model = createMockModel({ inputTokens: 20_000 }); + const doGenerate = (): PromiseLike => + model.doGenerate({ prompt: [] }); + const doStream = (): PromiseLike => model.doStream({ prompt: [] }); + + // Over-budget usage fed under an EMPTY-string sessionId (invalid). + await assertDefined( + middleware.wrapGenerate, + 'wrapGenerate', + )({ + doGenerate, + doStream, + params: { prompt: [], providerOptions: { contextChef: { sessionId: '' } } }, + model, + }); + + // A call with NO sessionId shares the default session → sees the fed usage. + const longPrompt = makeConversation(10); + const result = await assertDefined( + middleware.transformParams, + 'transformParams', + )({ + params: { prompt: longPrompt }, + type: 'generate', + model, + }); + expect(result.prompt.length).toBeLessThan(longPrompt.length); + + const invalidWarns = warn.mock.calls.filter((c) => String(c[0]).includes('sessionId')); + expect(invalidWarns).toHaveLength(1); + }); + + it('emits the missing-compression-config nag once across sessions', async () => { + const warn = vi.fn(); + const middleware = createMiddleware({ + contextWindow: 10_000, + onBeforeCompress: () => undefined, + logger: { warn }, + }); + const model = createMockModel(); + + for (const sessionId of ['user-a', 'user-b']) { + await assertDefined( + middleware.transformParams, + 'transformParams', + )({ + params: { + prompt: [{ role: 'user', content: [{ type: 'text', text: 'hi' }] }], + providerOptions: { contextChef: { sessionId } }, + }, + type: 'generate', + model, + }); + } + + const nags = warn.mock.calls.filter((c) => + String(c[0]).includes('No tokenizer and no compressionModel'), + ); + expect(nags).toHaveLength(1); + }); + it('onCompress receives the compressed slice in AI SDK format', async () => { const onCompressSpy = vi.fn(); const middleware = createMiddleware({ diff --git a/packages/core/src/adapters/geminiAdapter.test.ts b/packages/core/src/adapters/geminiAdapter.test.ts index 01f5837..1aa1b3d 100644 --- a/packages/core/src/adapters/geminiAdapter.test.ts +++ b/packages/core/src/adapters/geminiAdapter.test.ts @@ -725,6 +725,64 @@ describe('fromGemini', () => { expect(toolMsg?.content).toBe('{"temp":15}'); }); + it('correlates parallel same-name functionResponses to distinct call IDs', () => { + const contents: Content[] = [ + { role: 'user', parts: [{ text: 'weather in SF and NY?' }] }, + { + role: 'model', + parts: [ + { functionCall: { name: 'get_weather', args: { city: 'SF' } } }, + { functionCall: { name: 'get_weather', args: { city: 'NY' } } }, + ], + }, + { + role: 'user', + parts: [ + { functionResponse: { name: 'get_weather', response: { temp: 15 } } }, + { functionResponse: { name: 'get_weather', response: { temp: 8 } } }, + ], + }, + ]; + const { history } = fromGemini(contents); + + const callIds = history.find((m) => m.role === 'assistant')?.tool_calls?.map((tc) => tc.id); + expect(callIds).toEqual(['gemini-fc-get_weather-0', 'gemini-fc-get_weather-1']); + + const toolMsgs = history.filter((m) => m.role === 'tool'); + expect(toolMsgs.map((m) => m.tool_call_id)).toEqual([ + 'gemini-fc-get_weather-0', + 'gemini-fc-get_weather-1', + ]); + expect(toolMsgs.map((m) => m.content)).toEqual(['{"temp":15}', '{"temp":8}']); + }); + + it('generates unique call IDs for repeat calls of the same tool across turns', () => { + const contents: Content[] = [ + { role: 'user', parts: [{ text: 'ls' }] }, + { role: 'model', parts: [{ functionCall: { name: 'run_bash', args: { cmd: 'ls' } } }] }, + { + role: 'user', + parts: [{ functionResponse: { name: 'run_bash', response: { out: 'a.ts' } } }], + }, + { role: 'model', parts: [{ functionCall: { name: 'run_bash', args: { cmd: 'pwd' } } }] }, + { + role: 'user', + parts: [{ functionResponse: { name: 'run_bash', response: { out: '/repo' } } }], + }, + ]; + const { history } = fromGemini(contents); + + const assistantMsgs = history.filter((m) => m.role === 'assistant'); + expect(assistantMsgs[0]?.tool_calls?.[0].id).toBe('gemini-fc-run_bash-0'); + expect(assistantMsgs[1]?.tool_calls?.[0].id).toBe('gemini-fc-run_bash-1'); + + const toolMsgs = history.filter((m) => m.role === 'tool'); + expect(toolMsgs.map((m) => m.tool_call_id)).toEqual([ + 'gemini-fc-run_bash-0', + 'gemini-fc-run_bash-1', + ]); + }); + it('handles image-only messages (no text)', () => { const contents: Content[] = [ { diff --git a/packages/core/src/adapters/geminiAdapter.ts b/packages/core/src/adapters/geminiAdapter.ts index 0e2a91e..1d04575 100644 --- a/packages/core/src/adapters/geminiAdapter.ts +++ b/packages/core/src/adapters/geminiAdapter.ts @@ -42,6 +42,17 @@ export function fromGemini( const system: Message[] = []; const history: HistoryMessage[] = []; + // Gemini has no native tool call IDs. Synthesize `gemini-fc--` with + // a per-name counter spanning the whole conversation so repeat calls stay + // unique, and correlate each functionResponse to the oldest unconsumed call + // of the same name (FIFO) — responses normally arrive in a later content + // entry, so the correlation state must outlive a single message. + // FIFO assumes same-name responses arrive in call order (Gemini emits them + // in order); out-of-order same-name responses would swap attributions — an + // inherent limit of name-only correlation. + const callCounters = new Map(); + const pendingCallIds = new Map(); + if (systemInstruction) { for (const part of systemInstruction.parts) { system.push({ role: 'system', content: part.text }); @@ -72,26 +83,30 @@ export function fromGemini( data: part.fileData.fileUri, }); } else if ('functionCall' in part && part.functionCall) { - const id = `gemini-fc-${part.functionCall.name}-${toolCalls.length}`; + const name = part.functionCall.name; + const seq = callCounters.get(name) ?? 0; + callCounters.set(name, seq + 1); + const id = `gemini-fc-${name}-${seq}`; + const pending = pendingCallIds.get(name); + if (pending) pending.push(id); + else pendingCallIds.set(name, [id]); toolCalls.push({ id, type: 'function', function: { - name: part.functionCall.name, + name, arguments: JSON.stringify(part.functionCall.args), }, }); } else if ('functionResponse' in part && part.functionResponse) { - // Gemini has no native tool call IDs. Use a synthetic ID derived from - // the function name to enable cross-provider compilation (e.g. → OpenAI). const name = part.functionResponse.name; - // Try to match a preceding functionCall by name for ID correlation - const matchingCall = toolCalls.find((tc) => tc.function.name === name); + // Orphan responses (no unconsumed call of this name) fall back to `-0` + // and are cleaned up by ensureValidHistory below. history.push({ role: 'tool', content: JSON.stringify(part.functionResponse.response), name, - tool_call_id: matchingCall?.id ?? `gemini-fc-${name}-0`, + tool_call_id: pendingCallIds.get(name)?.shift() ?? `gemini-fc-${name}-0`, }); } } diff --git a/packages/core/src/adapters/openAIAdapter.test.ts b/packages/core/src/adapters/openAIAdapter.test.ts index c21b515..1bbccc2 100644 --- a/packages/core/src/adapters/openAIAdapter.test.ts +++ b/packages/core/src/adapters/openAIAdapter.test.ts @@ -395,4 +395,18 @@ describe('OpenAIAdapter — attachments output', () => { expect(Array.isArray(msg.content)).toBe(true); expect((msg.content as unknown as unknown[]).length).toBe(3); }); + + it('serializes toJSON-bearing values (e.g. Date) like the JSON round-trip did', () => { + const messages: Message[] = [ + { + role: 'user', + content: 'hi', + meta: { at: new Date('2026-01-02T03:04:05.000Z') }, + } as unknown as Message, + ]; + const result = adapter.compile(messages); + + const cloned = result.messages[0] as unknown as { meta: { at: unknown } }; + expect(cloned.meta.at).toBe('2026-01-02T03:04:05.000Z'); + }); }); diff --git a/packages/core/src/adapters/openAIAdapter.ts b/packages/core/src/adapters/openAIAdapter.ts index 16d81f3..2d97172 100644 --- a/packages/core/src/adapters/openAIAdapter.ts +++ b/packages/core/src/adapters/openAIAdapter.ts @@ -138,6 +138,34 @@ export function fromOpenAI(messages: SDKMessageParam[]): ParsedMessages { // ─── Output: IR → OpenAI ─── +/** + * Deep-clones plain message data, dropping `undefined`-valued properties. + * Equivalent to a `JSON.parse(JSON.stringify(...))` round-trip for this data + * domain (plain objects/arrays/primitives) without serializing — the string + * detour is measurably expensive for messages carrying base64 attachments. + * + * Matches JSON semantics for `toJSON`-bearing values (a Date lands as its + * ISO string, not `{}`). Assumes acyclic data: a circular reference + * overflows the stack, where the old round-trip threw a TypeError. + */ +function cloneWithoutUndefined(value: T): T { + if (Array.isArray(value)) { + return value.map(cloneWithoutUndefined) as T; + } + if (value !== null && typeof value === 'object') { + const toJSON = (value as { toJSON?: () => unknown }).toJSON; + if (typeof toJSON === 'function') { + return cloneWithoutUndefined(toJSON.call(value)) as T; + } + const out: Record = {}; + for (const [k, v] of Object.entries(value as Record)) { + if (v !== undefined) out[k] = cloneWithoutUndefined(v); + } + return out as T; + } + return value; +} + export class OpenAIAdapter implements ITargetAdapter { compile(messages: Message[]): OpenAIPayload { const formattedMessages: SDKMessageParam[] = messages.map((msg) => { @@ -163,10 +191,13 @@ export class OpenAIAdapter implements ITargetAdapter { }); } } - return JSON.parse(JSON.stringify({ ...cleanMsg, content: contentParts })); + return cloneWithoutUndefined({ + ...cleanMsg, + content: contentParts, + }) as unknown as SDKMessageParam; } - return JSON.parse(JSON.stringify(cleanMsg)); + return cloneWithoutUndefined(cleanMsg) as SDKMessageParam; }); if (formattedMessages.length > 0) { diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index cd7c6b8..e502605 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -57,6 +57,7 @@ export { Guardrail } from './modules/guardrail'; export { type CompressionDetails, compactMessages, + flattenForCompression, groupIntoTurns, Janitor, type JanitorConfig, @@ -114,6 +115,12 @@ export type { ChefLogger, ClearTarget, CompactOptions } from './types'; export * from './types'; export { ensureValidHistory } from './utils/ensureValidHistory'; export { type EventHandler, TypedEventEmitter } from './utils/eventEmitter'; +export { + DEFAULT_SESSION_KEY, + dedupeConstructionWarnings, + normalizeSessionKey, + SessionPool, +} from './utils/sessionPool'; export { TokenUtils } from './utils/tokenUtils'; export { XmlGenerator } from './utils/xmlGenerator'; diff --git a/packages/core/src/modules/janitor/index.test.ts b/packages/core/src/modules/janitor/index.test.ts index 1879a4c..bf34b15 100644 --- a/packages/core/src/modules/janitor/index.test.ts +++ b/packages/core/src/modules/janitor/index.test.ts @@ -2,7 +2,7 @@ import { describe, expect, it, vi } from 'vitest'; import { ContextChef } from '../../index'; import { Prompts } from '../../prompts'; import type { Message } from '../../types'; -import { compactMessages, groupIntoTurns, Janitor } from '.'; +import { compactMessages, flattenForCompression, groupIntoTurns, Janitor } from '.'; // ─── Helpers ─── @@ -1121,6 +1121,139 @@ describe('Janitor — compression circuit breaker', () => { }); }); +// ═══════════════════════════════════════════════════════ +// flattenForCompression — role-flattening for text-only compression models +// ═══════════════════════════════════════════════════════ + +describe('flattenForCompression', () => { + it('flattens tool results into user messages and tool calls into assistant text', () => { + const messages: Message[] = [ + { role: 'system', content: 'sys' }, + { role: 'user', content: 'do it' }, + { + role: 'assistant', + content: 'ok', + tool_calls: [ + { id: 'c1', type: 'function', function: { name: 'run', arguments: '{"a":1}' } }, + ], + }, + { role: 'tool', content: 'output', tool_call_id: 'c1' }, + ]; + + expect(flattenForCompression(messages)).toEqual([ + { role: 'system', content: 'sys' }, + { role: 'user', content: 'do it' }, + { role: 'assistant', content: 'ok\n[Called tool: run({"a":1})]' }, + { role: 'user', content: '[Tool result (c1): output]' }, + ]); + }); + + it('describes tool calls alone when the assistant message has no text', () => { + const messages: Message[] = [ + { + role: 'assistant', + content: '', + tool_calls: [{ id: 'c1', type: 'function', function: { name: 'run', arguments: '{}' } }], + }, + { role: 'tool', content: 'done', tool_call_id: undefined }, + ]; + + expect(flattenForCompression(messages)).toEqual([ + { role: 'assistant', content: '[Called tool: run({})]' }, + { role: 'user', content: '[Tool result: done]' }, + ]); + }); +}); + +// ═══════════════════════════════════════════════════════ +// Compression failure degradation +// ═══════════════════════════════════════════════════════ + +describe('Janitor — compression failure degradation', () => { + it('keeps the raw error out of the LLM-bound fallback summary and logs it instead', async () => { + const compressionModel = vi.fn().mockRejectedValue(new Error('ECONNRESET at http.Agent')); + const logger = { warn: vi.fn() }; + const janitor = new Janitor({ + contextWindow: 30, + tokenizer: makeTokenizer(10), + compressionModel, + logger, + }); + + const result = await janitor.compress(buildHistory(5)); + + const summary = result[0]; + expect(summary.content).toContain('older messages were truncated'); + expect(summary.content).not.toContain('ECONNRESET'); + expect(logger.warn).toHaveBeenCalledWith( + expect.stringContaining('compression model failed'), + expect.any(Error), + ); + }); + + it('does not abort compression when onCompress throws', async () => { + const compressionModel = vi.fn().mockResolvedValue('ok'); + const logger = { warn: vi.fn() }; + const onCompress = vi.fn().mockRejectedValue(new Error('db down')); + const janitor = new Janitor({ + contextWindow: 30, + tokenizer: makeTokenizer(10), + compressionModel, + onCompress, + logger, + }); + + const result = await janitor.compress(buildHistory(5)); + + expect(onCompress).toHaveBeenCalled(); + expect(result.length).toBeLessThan(5); + expect(logger.warn).toHaveBeenCalledWith( + expect.stringContaining('onCompress'), + expect.any(Error), + ); + }); + + it('falls back to default compression when onBeforeCompress throws', async () => { + const compressionModel = vi.fn().mockResolvedValue('ok'); + const logger = { warn: vi.fn() }; + const janitor = new Janitor({ + contextWindow: 30, + tokenizer: makeTokenizer(10), + compressionModel, + onBeforeCompress: () => { + throw new Error('hook boom'); + }, + logger, + }); + + const result = await janitor.compress(buildHistory(5)); + + // The throw is treated as "return null": default compression proceeds. + expect(result.length).toBeLessThan(5); + expect(compressionModel).toHaveBeenCalled(); + expect(logger.warn).toHaveBeenCalledWith( + expect.stringContaining('onBeforeCompress'), + expect.any(Error), + ); + }); + + it('does not abort the no-compressionModel fallback when onCompress throws', async () => { + const logger = { warn: vi.fn() }; + const janitor = new Janitor({ + contextWindow: 30, + tokenizer: makeTokenizer(10), + onCompress: vi.fn().mockRejectedValue(new Error('db down')), + logger, + }); + + await expect(janitor.compress(buildHistory(5))).resolves.toBeDefined(); + expect(logger.warn).toHaveBeenCalledWith( + expect.stringContaining('onCompress'), + expect.any(Error), + ); + }); +}); + // ═══════════════════════════════════════════════════════ // Media-aware compression: attachments → text placeholders // ═══════════════════════════════════════════════════════ diff --git a/packages/core/src/modules/janitor/index.ts b/packages/core/src/modules/janitor/index.ts index 5eccdf2..10bea34 100644 --- a/packages/core/src/modules/janitor/index.ts +++ b/packages/core/src/modules/janitor/index.ts @@ -1,11 +1,44 @@ import { Prompts } from '../../prompts'; -import type { Attachment, ChefLogger, CompactOptions, Message } from '../../types'; +import type { ChefLogger, CompactOptions, Message } from '../../types'; import { estimateObject } from '../../utils/tokenUtils'; const DEFAULT_PRESERVE_RATIO = 0.8; const DEFAULT_PRESERVE_RECENT_MESSAGES = 1; const MAX_CONSECUTIVE_COMPRESSION_FAILURES = 3; +/** + * Role-flattens history for a text-only compression model: tool results + * become user messages describing the result, and assistant tool calls are + * appended to the assistant text. This is the canonical implementation of + * the role-flattening contract documented on {@link summarizeHistory} — + * plain chat-completion endpoints reject `tool` roles and `tool_calls` + * fields, so a `compress` callback must flatten before forwarding. + */ +export function flattenForCompression( + messages: Message[], +): Array<{ role: 'system' | 'user' | 'assistant'; content: string }> { + return messages.map((m) => { + if (m.role === 'tool') { + return { + role: 'user' as const, + content: `[Tool result${m.tool_call_id ? ` (${m.tool_call_id})` : ''}: ${m.content}]`, + }; + } + if (m.role === 'assistant' && m.tool_calls?.length) { + const toolCallsDesc = m.tool_calls + .map((tc) => `[Called tool: ${tc.function.name}(${tc.function.arguments})]`) + .join('\n'); + return { + role: 'assistant' as const, + content: m.content ? `${m.content}\n${toolCallsDesc}` : toolCallsDesc, + }; + } + const role = + m.role === 'system' || m.role === 'user' || m.role === 'assistant' ? m.role : 'user'; + return { role, content: m.content }; + }); +} + // ─── Turn-based grouping ─── export interface Turn { @@ -52,25 +85,6 @@ export function groupIntoTurns(history: Message[]): Turn[] { // ─── Attachment stripping for compression ─── -/** - * Builds a single-line text placeholder for an attachment. - * Includes the filename when available so the summary can reference it by name. - * - * { mediaType: 'image/png', filename: 'photo.png' } → '[image: photo.png]' - * { mediaType: 'image/png' } → '[image]' - * { mediaType: 'application/pdf', filename: 'r.pdf' } → '[document: r.pdf]' - * { mediaType: 'application/pdf' } → '[document]' - * { mediaType: '' } → '[attachment]' - * - * Categorization mirrors Claude Code's binary image-vs-document split — keeping - * the placeholder vocabulary small reduces surprises for the compression model. - */ -function attachmentToPlaceholder(att: Attachment): string { - const mt = att.mediaType.toLowerCase(); - const kind = mt.startsWith('image/') ? 'image' : mt ? 'document' : 'attachment'; - return att.filename ? `[${kind}: ${att.filename}]` : `[${kind}]`; -} - /** * Replaces media attachments with text placeholders for the compression model. * @@ -90,7 +104,9 @@ function stripAttachmentsForCompression(messages: Message[]): Message[] { return messages.map((msg) => { if (!msg.attachments?.length) return msg; - const placeholders = msg.attachments.map(attachmentToPlaceholder).join('\n'); + const placeholders = msg.attachments + .map((att) => Prompts.getAttachmentPlaceholder(att.mediaType, att.filename)) + .join('\n'); const newContent = msg.content ? `${placeholders}\n${msg.content}` : placeholders; const { attachments: _attachments, ...rest } = msg; @@ -264,8 +280,9 @@ interface JanitorConfigBase { * @param details - Boundary metadata: the exact messages that were replaced, * useful for persistence layers that need to map the summary back to their store. * - * Contract: must not throw or reject. Errors propagate out of compile() — there - * is no fallback path. Wrap your logic in try/catch if it can fail. + * Contract: should not throw or reject. As a safety net, a throwing hook is + * caught and logged via `logger` — the compression result is kept and + * compile() continues, but your sink may have missed the summary. */ onCompress?: ( summaryMessage: Message, @@ -278,8 +295,10 @@ interface JanitorConfigBase { * Return a modified Message[] to replace the history before compression proceeds, * or return null/undefined to let the default compression handle it. * - * Contract: must not throw or reject. Errors propagate out of compile() — return - * null on failure to fall back to default LLM compression rather than throwing. + * Contract: should not throw or reject. As a safety net, a throwing hook is + * caught and logged via `logger`, then treated as if it returned null — + * default compression proceeds (the same failure recipe the return + * contract documents). Mirrors the `onCompress` degradation stance. */ onBeforeCompress?: ( history: Message[], @@ -545,10 +564,23 @@ export class Janitor { // Fire onBeforeCompress hook — developer gets a chance to intervene const hook = this.config.onBeforeCompress ?? this.config.onBudgetExceeded; if (hook) { - const modified = await hook(history, { - currentTokens, - limit: this.config.contextWindow, - }); + let modified: Message[] | null | undefined; + try { + modified = await hook(history, { + currentTokens, + limit: this.config.contextWindow, + }); + } catch (error) { + // Same degradation stance as onCompress: a broken hook must not fail + // compile(). A throw is treated as "return null" — the failure recipe + // the return contract already documents — so default compression + // proceeds. + (this.config.logger ?? console).warn( + '[context-chef] onBeforeCompress hook threw — proceeding with default compression', + error, + ); + modified = null; + } if (modified != null) { // Re-evaluate with the developer-modified history @@ -702,13 +734,11 @@ export class Janitor { const toKeep = history.slice(splitIndex); if (!this.config.compressionModel) { - if (this.config.onCompress) { - await this.config.onCompress( - { role: 'system', content: Prompts.getFallbackCompressionSummary(toCompress.length) }, - toCompress.length, - { compressedMessages: toCompress }, - ); - } + await this._fireOnCompress( + { role: 'system', content: Prompts.getFallbackCompressionSummary(toCompress.length) }, + toCompress.length, + { compressedMessages: toCompress }, + ); this._suppressNextCompression = true; return [...toKeep]; } @@ -731,9 +761,14 @@ export class Janitor { // Increment circuit breaker. After MAX_CONSECUTIVE_COMPRESSION_FAILURES, // compress() will short-circuit to avoid futile retries. this._consecutiveFailures++; - summaryText = - Prompts.getFallbackCompressionSummary(toCompress.length) + - `\n(Compression failed: ${error})`; + // The raw error belongs in the logger, not in the LLM-bound summary. + // The compressed slice is still dropped — recover it via onCompress + // details if you need durable history. + (this.config.logger ?? console).warn( + '[context-chef] compression model failed — compressed messages replaced by a bare truncation notice', + error, + ); + summaryText = Prompts.getFallbackCompressionSummary(toCompress.length); } const summaryMessage: Message = { @@ -741,15 +776,34 @@ export class Janitor { content: Prompts.getCompactSummaryWrapper(summaryText), }; - if (this.config.onCompress) { - await this.config.onCompress(summaryMessage, toCompress.length, { - compressedMessages: toCompress, - }); - } + await this._fireOnCompress(summaryMessage, toCompress.length, { + compressedMessages: toCompress, + }); // E10: Suppress the immediate next compression check. this._suppressNextCompression = true; return [summaryMessage, ...toKeep]; } + + /** + * Invokes the `onCompress` hook, downgrading a throwing/rejecting hook to a + * logger warning. The hook is an observation/persistence sink — its failure + * must not discard an already-computed compression or fail compile(). + */ + private async _fireOnCompress( + summaryMessage: Message, + truncatedCount: number, + details: CompressionDetails, + ): Promise { + if (!this.config.onCompress) return; + try { + await this.config.onCompress(summaryMessage, truncatedCount, details); + } catch (error) { + (this.config.logger ?? console).warn( + '[context-chef] onCompress hook threw — compression result is kept, but your sink may have missed this summary', + error, + ); + } + } } diff --git a/packages/core/src/prompts.ts b/packages/core/src/prompts.ts index 1be7c81..c85ceb6 100644 --- a/packages/core/src/prompts.ts +++ b/packages/core/src/prompts.ts @@ -220,6 +220,49 @@ Do not output any introductory text or acknowledgement. Start directly with the `.trim(), + // ─── Placeholder vocabulary ─── + // Three placeholder families coexist, each serving a distinct purpose: + // 1. getAttachmentPlaceholder — stands in for binary attachments shown to + // the COMPRESSION model (janitor strips media before summarizing). + // 2. getToolResultFilePlaceholder / getToolResultPartPlaceholder — stand + // in for non-text tool-result content parts flattened to text by the + // middleware adapters (dropping them silently would hide the part + // from compression and truncation entirely). + // 3. getVFSOffloadReminder (above) — a dereferenceable `context://` + // pointer, not a placeholder: the full content is retrievable. + // Formats are frozen conventions — changing one reshapes what compression + // models see mid-conversation. + + /** + * Single-line text placeholder for a media attachment, shown to the + * compression model in place of binary data. Includes the filename when + * available so the summary can reference it by name. + * + * ('image/png', 'photo.png') → '[image: photo.png]' + * ('image/png') → '[image]' + * ('application/pdf', 'r.pdf') → '[document: r.pdf]' + * ('') → '[attachment]' + * + * Categorization mirrors Claude Code's binary image-vs-document split — + * keeping the placeholder vocabulary small reduces surprises for the + * compression model. + */ + getAttachmentPlaceholder: (mediaType: string, filename?: string): string => { + const mt = mediaType.toLowerCase(); + const kind = mt.startsWith('image/') ? 'image' : mt ? 'document' : 'attachment'; + return filename ? `[${kind}: ${filename}]` : `[${kind}]`; + }, + + /** + * Placeholder for a file/media part inside a tool result that is being + * flattened to text — the part must leave a trace, not vanish silently. + */ + getToolResultFilePlaceholder: (mediaType: string, filename?: string): string => + `[tool result file: ${mediaType}${filename ? ` (${filename})` : ''}]`, + + /** Placeholder for any other non-text tool-result content part. */ + getToolResultPartPlaceholder: (partType: string): string => `[tool result part: ${partType}]`, + /** * Static instruction injected into the system prompt when memory is enabled. * Guides the LLM to use memory tools for persistence. diff --git a/packages/core/src/utils/sessionPool.test.ts b/packages/core/src/utils/sessionPool.test.ts new file mode 100644 index 0000000..bd914e6 --- /dev/null +++ b/packages/core/src/utils/sessionPool.test.ts @@ -0,0 +1,97 @@ +import { describe, expect, it, vi } from 'vitest'; +import { + DEFAULT_SESSION_KEY, + dedupeConstructionWarnings, + normalizeSessionKey, + SessionPool, +} from './sessionPool'; + +describe('SessionPool', () => { + it('returns the same instance for the same key', () => { + let n = 0; + const pool = new SessionPool(() => ({ id: n++ })); + + const first = pool.get('a'); + expect(pool.get('a')).toBe(first); + expect(first.id).toBe(0); + }); + + it('creates distinct instances per key', () => { + let n = 0; + const pool = new SessionPool(() => ({ id: n++ })); + + expect(pool.get('a').id).toBe(0); + expect(pool.get('b').id).toBe(1); + }); + + it('evicts the least-recently-used entry beyond maxSize', () => { + let n = 0; + const pool = new SessionPool(() => ({ id: n++ }), { maxSize: 2 }); + + const a = pool.get('a'); // id 0 + pool.get('b'); // id 1 + pool.get('a'); // touch a — b becomes LRU + pool.get('c'); // id 2 — evicts b + + expect(pool.get('a')).toBe(a); // still cached + expect(pool.get('b').id).toBe(3); // was evicted, recreated + }); + + it('exposes the current entry count', () => { + const pool = new SessionPool(() => ({}), { maxSize: 10 }); + pool.get('a'); + pool.get('b'); + expect(pool.size).toBe(2); + }); + + it('rejects zero, negative, and non-integer maxSize', () => { + expect(() => new SessionPool(() => ({}), { maxSize: 0 })).toThrow(RangeError); + expect(() => new SessionPool(() => ({}), { maxSize: -1 })).toThrow(RangeError); + expect(() => new SessionPool(() => ({}), { maxSize: 1.5 })).toThrow(RangeError); + expect(() => new SessionPool(() => ({}), { maxSize: Number.NaN })).toThrow(RangeError); + }); + + it('accepts maxSize 1 and still caches the single entry', () => { + const pool = new SessionPool(() => ({}), { maxSize: 1 }); + expect(pool.get('a')).toBe(pool.get('a')); + }); +}); + +describe('normalizeSessionKey', () => { + it('passes non-empty strings through', () => { + expect(normalizeSessionKey('user-1')).toBe('user-1'); + }); + + it('maps absent values to the default key without flagging', () => { + const onInvalid = vi.fn(); + expect(normalizeSessionKey(undefined, onInvalid)).toBe(DEFAULT_SESSION_KEY); + expect(normalizeSessionKey(null, onInvalid)).toBe(DEFAULT_SESSION_KEY); + expect(onInvalid).not.toHaveBeenCalled(); + }); + + it('maps present-but-unusable values to the default key and flags them', () => { + const onInvalid = vi.fn(); + expect(normalizeSessionKey('', onInvalid)).toBe(DEFAULT_SESSION_KEY); + expect(normalizeSessionKey(42, onInvalid)).toBe(DEFAULT_SESSION_KEY); + expect(normalizeSessionKey({ id: 'x' }, onInvalid)).toBe(DEFAULT_SESSION_KEY); + expect(onInvalid).toHaveBeenCalledTimes(3); + }); +}); + +describe('dedupeConstructionWarnings', () => { + it('dedupes identical construction-time warnings across instances but passes runtime warnings through', () => { + const warn = vi.fn(); + const factory = dedupeConstructionWarnings({ warn }, (constructionLogger) => { + constructionLogger.warn('config nag'); + return { logger: constructionLogger }; + }); + + const first = factory(); + factory(); + expect(warn).toHaveBeenCalledTimes(1); // nag fired once despite two constructions + + first.logger.warn('runtime issue'); + first.logger.warn('runtime issue'); + expect(warn).toHaveBeenCalledTimes(3); // post-construction warnings never deduped + }); +}); diff --git a/packages/core/src/utils/sessionPool.ts b/packages/core/src/utils/sessionPool.ts new file mode 100644 index 0000000..8334d03 --- /dev/null +++ b/packages/core/src/utils/sessionPool.ts @@ -0,0 +1,106 @@ +import type { ChefLogger } from '../types'; + +/** Pool key shared by every call that carries no usable session identity. */ +export const DEFAULT_SESSION_KEY = '__default__'; + +/** + * Normalizes a caller-supplied session identifier to a pool key. + * + * Absent values (`undefined` / `null`) map to {@link DEFAULT_SESSION_KEY} + * silently — sharing one default session is the documented behavior for + * callers that never opt into isolation. A value that is PRESENT but + * unusable (empty string, non-string) also maps to the default key, but + * fires `onInvalid` first: a misconfigured session id silently merging + * unrelated conversations is exactly the failure mode this guards against. + */ +export function normalizeSessionKey(raw: unknown, onInvalid?: (raw: unknown) => void): string { + if (typeof raw === 'string' && raw) return raw; + if (raw !== undefined && raw !== null) onInvalid?.(raw); + return DEFAULT_SESSION_KEY; +} + +/** + * Wraps a pool factory so logger warnings emitted DURING construction are + * deduped by message across instances, while warnings after construction + * pass through untouched. Pooled factories construct many identically + * configured instances — a configuration nag repeated per session is log + * noise, but runtime warnings (e.g. a compression failure) must stay + * per-occurrence. + */ +export function dedupeConstructionWarnings( + logger: ChefLogger, + build: (constructionLogger: ChefLogger) => T, +): () => T { + const seen = new Set(); + return () => { + let constructing = true; + const constructionLogger: ChefLogger = { + warn(message, ...args) { + if (constructing) { + if (seen.has(message)) return; + seen.add(message); + } + logger.warn(message, ...args); + }, + }; + const instance = build(constructionLogger); + constructing = false; + return instance; + }; +} + +/** + * Keyed instance pool with LRU eviction. + * + * Built for the middleware packages: a middleware instance is typically + * created once at module scope but serves many concurrent conversations, so + * per-conversation state (a stateful Janitor) must be keyed by session — not + * shared — or token-usage feeds, compression suppression, and circuit-breaker + * counts leak across callers. + * + * `get()` refreshes the entry's LRU position. When the pool exceeds + * `maxSize`, the least-recently-used entries are dropped; a dropped session + * is transparently recreated on next access (losing only its fed token + * usage — the next over-budget call re-triggers compression). + */ +export class SessionPool { + private readonly entries = new Map(); + private readonly create: (key: string) => T; + private readonly maxSize: number; + + constructor(create: (key: string) => T, options?: { maxSize?: number }) { + const maxSize = options?.maxSize ?? 256; + // A non-positive cap would evict every entry the moment it is inserted, + // silently disabling pooling (and with it, fed-usage compression). + if (!Number.isInteger(maxSize) || maxSize < 1) { + throw new RangeError( + `[context-chef] SessionPool maxSize must be a positive integer, got ${maxSize}`, + ); + } + this.create = create; + this.maxSize = maxSize; + } + + /** Returns the instance for `key`, creating it on first access. */ + get(key: string): T { + const existing = this.entries.get(key); + if (existing !== undefined) { + // Re-insert to refresh LRU order (Map preserves insertion order). + this.entries.delete(key); + this.entries.set(key, existing); + return existing; + } + const created = this.create(key); + this.entries.set(key, created); + while (this.entries.size > this.maxSize) { + const oldest = this.entries.keys().next().value; + if (oldest === undefined) break; + this.entries.delete(oldest); + } + return created; + } + + get size(): number { + return this.entries.size; + } +} diff --git a/packages/tanstack-ai/README.md b/packages/tanstack-ai/README.md index 5db23fd..4772259 100644 --- a/packages/tanstack-ai/README.md +++ b/packages/tanstack-ai/README.md @@ -118,6 +118,10 @@ The lookup key is the tool's name — read from `ModelMessage.name` when set, ot The middleware automatically extracts token usage from `onUsage` callbacks and feeds it back to the compression engine. No manual tracking needed. +### Conversation Isolation + +Compression state (fed token usage, compression suppression, the failure circuit breaker) is tracked per `ctx.conversationId`, so one middleware instance safely serves many conversations. Pass a `conversationId` to `chat()`; calls without one share a default slot — fine for single-conversation processes, wrong for multi-user servers. Up to `maxSessions` conversations are tracked concurrently (default 256, LRU-evicted); an evicted conversation is transparently recreated on next access, losing only its fed token usage. + ### Compact (Mechanical Pruning) Zero-LLM-cost message pruning — removes tool call/result pairs and empty messages before compression: diff --git a/packages/tanstack-ai/src/middleware.ts b/packages/tanstack-ai/src/middleware.ts index bc1ce92..56e5b6d 100644 --- a/packages/tanstack-ai/src/middleware.ts +++ b/packages/tanstack-ai/src/middleware.ts @@ -2,9 +2,13 @@ import { type ChefLogger, type CompressionDetails, compactMessages as clearMessages, + dedupeConstructionWarnings, + flattenForCompression, Janitor, type Message, + normalizeSessionKey, Prompts, + SessionPool, XmlGenerator, } from '@context-chef/core'; import type { AnyTextAdapter, ChatMiddleware, ModelMessage } from '@tanstack/ai'; @@ -14,8 +18,6 @@ import { compactMessages } from './compact'; import { truncateToolResults } from './truncator'; import type { ContextChefOptions, DynamicStateConfig } from './types'; -type CompressRole = 'system' | 'user' | 'assistant'; - /** * Creates a TanStack AI ChatMiddleware that transparently applies * context-chef compression and truncation to chat() calls. @@ -90,24 +92,52 @@ export function contextChefMiddleware(options: ContextChefOptions): ChatMiddlewa usagePreference = 'max'; } - const janitor = options.tokenizer - ? new Janitor({ - ...sharedJanitorConfig, - tokenizer: (msgs: Message[]) => options.tokenizer?.(msgs) ?? 0, - preserveRatio: options.compress?.preserveRatio ?? 0.8, - usagePreference, - }) - : new Janitor({ - ...sharedJanitorConfig, - // 'tokenizerFirst' has been sanitized above; the cast narrows the - // remaining values to the no-tokenizer branch. - usagePreference: usagePreference as 'max' | 'feedFirst' | undefined, - }); + // One Janitor per conversation. A middleware instance is usually created + // once and reused across chat() calls for many conversations — sharing one + // Janitor would leak token-usage feeds, compression suppression, and + // circuit-breaker counts across them. TanStack AI supplies the conversation + // identity via ctx.conversationId; calls without one share the default + // session (prior behavior). Construction-time config nags are deduped + // across conversations — the config is identical for every pooled Janitor. + const janitors = new SessionPool( + dedupeConstructionWarnings(logger, (constructionLogger) => + options.tokenizer + ? new Janitor({ + ...sharedJanitorConfig, + logger: constructionLogger, + tokenizer: (msgs: Message[]) => options.tokenizer?.(msgs) ?? 0, + preserveRatio: options.compress?.preserveRatio ?? 0.8, + usagePreference, + }) + : new Janitor({ + ...sharedJanitorConfig, + logger: constructionLogger, + // 'tokenizerFirst' has been sanitized above; the cast narrows the + // remaining values to the no-tokenizer branch. + usagePreference: usagePreference as 'max' | 'feedFirst' | undefined, + }), + ), + { maxSize: options.maxSessions }, + ); + + let invalidConversationIdWarned = false; + const flagInvalidConversationId = (raw: unknown) => { + if (invalidConversationIdWarned) return; + invalidConversationIdWarned = true; + logger.warn( + '[context-chef] Invalid ctx.conversationId (expected a non-empty string, got ' + + `${raw === '' ? 'empty string' : typeof raw}); routing to the default conversation slot.`, + ); + }; + + const janitorFor = (ctx: { conversationId?: string }): Janitor => + janitors.get(normalizeSessionKey(ctx.conversationId, flagInvalidConversationId)); return { name: 'context-chef', - onConfig: async (_ctx, config) => { + onConfig: async (ctx, config) => { + const janitor = janitorFor(ctx); let { messages } = config; let systemPrompts = [...config.systemPrompts]; @@ -166,9 +196,9 @@ export function contextChefMiddleware(options: ContextChefOptions): ChatMiddlewa return { messages, systemPrompts }; }, - onUsage: (_ctx, usage) => { + onUsage: (ctx, usage) => { if (usage.promptTokens != null) { - janitor.feedTokenUsage(usage.promptTokens); + janitorFor(ctx).feedTokenUsage(usage.promptTokens); } else if (!usageWarned && !options.tokenizer) { usageWarned = true; logger.warn( @@ -193,31 +223,11 @@ function createCompressionAdapter( adapter: AnyTextAdapter, ): (messages: Message[]) => Promise { return async (messages: Message[]): Promise => { - const formatted = messages.map((m): { role: CompressRole; content: string } => { - if (m.role === 'tool') { - return { - role: 'user' satisfies CompressRole, - content: `[Tool result${m.tool_call_id ? ` (${m.tool_call_id})` : ''}: ${m.content}]`, - }; - } - if (m.role === 'assistant' && m.tool_calls?.length) { - const toolCallsDesc = m.tool_calls - .map((tc) => `[Called tool: ${tc.function.name}(${tc.function.arguments})]`) - .join('\n'); - return { - role: 'assistant' satisfies CompressRole, - content: m.content ? `${m.content}\n${toolCallsDesc}` : toolCallsDesc, - }; - } - return { - role: toCompressRole(m.role), - content: m.content, - }; - }); - - // Convert to ModelMessage format for chatStream - const modelMessages: ModelMessage[] = formatted.map((m) => ({ - role: m.role === 'system' ? ('user' as const) : (m.role as 'user' | 'assistant'), + // Role-flatten via the shared core helper, then convert to ModelMessage + // for chatStream — providers reject a `system` role mid-conversation + // here, so it degrades to `user`. + const modelMessages: ModelMessage[] = flattenForCompression(messages).map((m) => ({ + role: m.role === 'system' ? ('user' as const) : m.role, content: m.content, })); @@ -289,11 +299,6 @@ async function injectDynamicState( }; } -function toCompressRole(role: string): CompressRole { - if (role === 'system' || role === 'user' || role === 'assistant') return role; - return 'user'; -} - /** * Resolves the `skill` option to its instructions string. * Returns the instructions when a Skill is active and has non-empty diff --git a/packages/tanstack-ai/src/types.ts b/packages/tanstack-ai/src/types.ts index 9802e3a..7284ed9 100644 --- a/packages/tanstack-ai/src/types.ts +++ b/packages/tanstack-ai/src/types.ts @@ -197,6 +197,15 @@ export interface ContextChefOptions { * underlying Janitor and Offloader. */ logger?: ChefLogger; + /** + * Cap on concurrently tracked conversations. Each `ctx.conversationId` + * gets its own Janitor so token-usage feeds and compression state never + * leak across conversations sharing one middleware instance. + * Least-recently-used conversations beyond the cap are dropped and + * transparently recreated on next access. Must be a positive integer — + * the pool throws a RangeError otherwise. Default: 256. + */ + maxSessions?: number; /** * Hook called after compression occurs. * diff --git a/packages/tanstack-ai/tests/middleware.test.ts b/packages/tanstack-ai/tests/middleware.test.ts index 8606f94..1748590 100644 --- a/packages/tanstack-ai/tests/middleware.test.ts +++ b/packages/tanstack-ai/tests/middleware.test.ts @@ -242,6 +242,78 @@ describe('contextChefMiddleware', () => { }); }); + it('isolates janitor state between conversations via ctx.conversationId', async () => { + // Window large enough that the heuristic estimate of the short history + // stays under budget — only the FED usage (conversation A's) exceeds it. + const mw = contextChefMiddleware({ contextWindow: 10_000 }); + const messages: ModelMessage[] = [ + { role: 'user', content: 'q1' }, + { role: 'assistant', content: 'a1' }, + { role: 'user', content: 'q2' }, + { role: 'assistant', content: 'a2' }, + { role: 'user', content: 'q3' }, + { role: 'assistant', content: 'a3' }, + ]; + + // Conversation A exceeds the budget (20_000 > 10_000). + mw.onUsage?.(createMockCtx({ conversationId: 'conv-a' }), { + promptTokens: 20_000, + completionTokens: 10, + totalTokens: 20_010, + }); + + // Conversation B must NOT inherit A's fed token usage — no compression. + const resB = await mw.onConfig?.( + createMockCtx({ conversationId: 'conv-b' }), + createMockConfig(messages), + ); + expect(getResult(resB ?? {}).messages).toHaveLength(messages.length); + + // Conversation A compresses on its own next call. + const resA = await mw.onConfig?.( + createMockCtx({ conversationId: 'conv-a' }), + createMockConfig(messages), + ); + expect(getResult(resA ?? {}).messages.length).toBeLessThan(messages.length); + }); + + it('routes an empty conversationId to the default slot and warns once', async () => { + const warn = vi.fn(); + const mw = contextChefMiddleware({ contextWindow: 10_000, logger: { warn } }); + const messages: ModelMessage[] = [ + { role: 'user', content: 'q1' }, + { role: 'assistant', content: 'a1' }, + { role: 'user', content: 'q2' }, + { role: 'assistant', content: 'a2' }, + ]; + + // Over-budget usage reported under an EMPTY conversationId (invalid). + mw.onUsage?.(createMockCtx({ conversationId: '' }), { + promptTokens: 20_000, + completionTokens: 10, + totalTokens: 20_010, + }); + + // A call with NO conversationId shares the default slot → sees the fed usage. + const res = await mw.onConfig?.(createMockCtx(), createMockConfig(messages)); + expect(getResult(res ?? {}).messages.length).toBeLessThan(messages.length); + expect(warn).toHaveBeenCalledWith(expect.stringContaining('conversationId')); + }); + + it('emits the missing-compression-config nag once across conversations', async () => { + const warn = vi.fn(); + const mw = contextChefMiddleware({ contextWindow: 100_000, logger: { warn } }); + const messages: ModelMessage[] = [{ role: 'user', content: 'hi' }]; + + await mw.onConfig?.(createMockCtx({ conversationId: 'conv-a' }), createMockConfig(messages)); + await mw.onConfig?.(createMockCtx({ conversationId: 'conv-b' }), createMockConfig(messages)); + + const nags = warn.mock.calls.filter((c) => + String(c[0]).includes('No tokenizer and no compressionModel'), + ); + expect(nags).toHaveLength(1); + }); + it('warns once when promptTokens is missing', () => { const localWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {}); const mw = contextChefMiddleware({ contextWindow: 100_000 });