diff --git a/src/providers/key-store.ts b/src/providers/key-store.ts index 12e4ce6cb7..614fd3372f 100644 --- a/src/providers/key-store.ts +++ b/src/providers/key-store.ts @@ -64,6 +64,16 @@ function keychainAccount(reference: string): string { return reference.slice(KEYCHAIN_REFERENCE_PREFIX.length); } +/** + * A reference belongs to `name` only when its account is that provider's own active account + * or one of its pool accounts. `storeProviderKeyInKeychain` writes exactly those two shapes, + * so anything else in a provider's config names another provider's secret. + */ +function keychainReferenceBelongsToProvider(reference: string, name: string): boolean { + const account = keychainAccount(reference); + return account === name || account.startsWith(`${name}/`); +} + function readKeychain(account: string): string | undefined { const cached = resolvedCache.get(account); if (cached !== undefined) return cached; @@ -185,6 +195,18 @@ export function restoreProviderKeyFromKeychain(config: OcxConfig, name: string): const pool = provider.apiKeyPool ?? []; const resolved = new Map(); const refs = [provider.apiKey, ...pool.map(e => e.key)].filter(isKeychainReference); + // Restore reads a secret out of the keychain, writes it back to config as plaintext, and then + // DELETES the keychain item. Following a reference to another provider's account would both + // disclose that secret through this provider's config and destroy the real owner's credential, + // so refuse before anything is read or removed. + const foreign = refs.filter(ref => !keychainReferenceBelongsToProvider(ref, name)); + if (foreign.length > 0) { + return { + ok: false, + error: `provider "${name}" references a keychain account it does not own (${foreign.length} reference(s)); config left unchanged`, + status: 400, + }; + } for (const ref of refs) { const account = keychainAccount(ref); if (resolved.has(account)) continue; diff --git a/src/responses/citation-markers.ts b/src/responses/citation-markers.ts index 5fe58142cf..3e3943cad5 100644 --- a/src/responses/citation-markers.ts +++ b/src/responses/citation-markers.ts @@ -68,6 +68,15 @@ export interface CitationMarkerFilter { flush(): string; } +/** + * Upper bound on the text withheld for one unterminated START. + * + * A real span is `cite` plus a few turn-scoped ids, so it is far under this. Without a + * bound, a backend that emits a START and never terminates it makes `held` grow for the + * whole response, and every later delta re-scans that accumulated prefix. + */ +const MAX_STREAMING_MARKER_SPAN_LENGTH = 4_096; + /** * Streaming filter. * @@ -75,6 +84,9 @@ export interface CitationMarkerFilter { * next — so a stateless per-delta strip would emit the tail of a span it never recognized. * This holds back the text from an unterminated START and releases it once the END arrives * (removed) or the stream ends (verbatim, so nothing the model actually said is lost). + * + * A span that grows past `MAX_STREAMING_MARKER_SPAN_LENGTH` is malformed ordinary text, so + * it is released verbatim instead of withheld; a later START can still open a valid span. */ export function createCitationMarkerFilter(): CitationMarkerFilter { // Text from an open START that has not been terminated yet. @@ -83,13 +95,29 @@ export function createCitationMarkerFilter(): CitationMarkerFilter { push(delta: string): string { const combined = held + delta; held = ""; - const start = combined.lastIndexOf(CITATION_MARKER_START); - if (start === -1) return stripCitationMarkers(combined); - const endAfterStart = combined.indexOf(CITATION_MARKER_END, start + 1); - if (endAfterStart !== -1) return stripCitationMarkers(combined); - // The trailing span is still open: emit everything before it, hold the rest. - held = combined.slice(start); - return stripCitationMarkers(combined.slice(0, start)); + let start = combined.indexOf(CITATION_MARKER_START); + if (start === -1) return combined; + let out = combined.slice(0, start); + // Walk START-delimited segments independently so an earlier malformed START is never + // paired with a later span's END (the whole-string strip would do exactly that). + while (start !== -1) { + const nextStart = combined.indexOf(CITATION_MARKER_START, start + 1); + const segment = combined.slice(start, nextStart === -1 ? combined.length : nextStart); + const end = segment.indexOf(CITATION_MARKER_END, 1); + if (end !== -1) { + // A complete span: drop it, keep whatever trails it inside this segment. + out += segment.slice(end + 1); + } else if (nextStart === -1 && segment.length <= MAX_STREAMING_MARKER_SPAN_LENGTH) { + // Only a bounded trailing span can still be completed by a later delta. + held = segment; + } else { + // Superseded by a later START, or over the bound: ordinary text, emitted verbatim + // so neither the retained text nor the per-delta rescan grows without limit. + out += segment; + } + start = nextStart; + } + return out; }, flush(): string { const rest = held; @@ -98,4 +126,3 @@ export function createCitationMarkerFilter(): CitationMarkerFilter { }, }; } - diff --git a/tests/providers/provider-key-store.test.ts b/tests/providers/provider-key-store.test.ts index 645920e32c..1197a1fea9 100644 --- a/tests/providers/provider-key-store.test.ts +++ b/tests/providers/provider-key-store.test.ts @@ -156,6 +156,37 @@ describe("store / restore", () => { expect(probeProviderKeychain().available).toBe(false); }); + test("restore refuses a reference to another provider's keychain account", () => { + const { store, factory } = fakeKeychain(); + setProviderKeychainEntryFactoryForTests(factory); + const config = loadConfig(); + config.providers.other = { adapter: "openai-chat", baseUrl: "https://other.example/v1", apiKey: POOL_SECRET }; + expect(storeProviderKeyInKeychain(config, "other")).toEqual({ ok: true, moved: 1 }); + expect(config.providers.other!.apiKey).toBe("keychain:other"); + + // Point "relay" at the account "other" owns. Restore would otherwise read that secret, + // write it into relay's config as plaintext, and delete the owner's keychain item. + config.providers.relay!.apiKey = "keychain:other"; + const result = restoreProviderKeyFromKeychain(config, "relay"); + expect(result.ok).toBe(false); + if (!result.ok) expect(result.status).toBe(400); + + expect(config.providers.relay!.apiKey).toBe("keychain:other"); + expect(readFileSync(join(testDir, "config.json"), "utf8")).not.toContain(POOL_SECRET); + // The real owner's secret is still in the keychain and still resolves for that provider. + expect(store.size).toBe(1); + expect(resolveProviderApiKey(config.providers.other!.apiKey)).toBe(POOL_SECRET); + }); + + test("restore still accepts a provider's own active and pool accounts", () => { + const { factory } = fakeKeychain(); + setProviderKeychainEntryFactoryForTests(factory); + const config = loadConfig(); + config.providers.relay!.apiKeyPool = [{ id: "a1", key: SECRET }, { id: "b2", key: POOL_SECRET }]; + expect(storeProviderKeyInKeychain(config, "relay")).toEqual({ ok: true, moved: 2 }); + expect(restoreProviderKeyFromKeychain(config, "relay")).toEqual({ ok: true, restored: 2 }); + }); + test("management route: GET reports store kind, POST store/restore round-trips", async () => { const { factory } = fakeKeychain(); setProviderKeychainEntryFactoryForTests(factory); @@ -192,4 +223,3 @@ describe("store / restore", () => { } }); }); - diff --git a/tests/responses/citation-markers.test.ts b/tests/responses/citation-markers.test.ts index 0c1921750c..dcc7f9abc6 100644 --- a/tests/responses/citation-markers.test.ts +++ b/tests/responses/citation-markers.test.ts @@ -87,4 +87,34 @@ describe("streaming citation marker filter (#3150)", () => { const filter = createCitationMarkerFilter(); expect(filter.push(`visible now ${S}cite`)).toBe("visible now "); }); + + test("an unterminated span past the bound is released instead of retained", () => { + // A backend that opens a span and never closes it must not make the filter accumulate + // the rest of the response, which every later delta would then re-scan. + const filter = createCitationMarkerFilter(); + let out = filter.push(`kept ${S}cite`); + expect(out).toBe("kept "); + for (let i = 0; i < 5_000; i += 1) out += filter.push("x"); + + // Everything after the malformed START is emitted verbatim, so nothing is lost, and + // flush() has nothing left to release. + expect(out).toBe(`kept ${S}cite${"x".repeat(5_000)}`); + expect(filter.flush()).toBe(""); + }); + + test("a later START still opens a valid span after a released malformed one", () => { + const filter = createCitationMarkerFilter(); + let out = filter.push(`a${S}${"y".repeat(5_000)}`); + out += filter.push(`${S}cite${P}turn1view0${E} tail`); + expect(out).toBe(`a${S}${"y".repeat(5_000)} tail`); + expect(filter.flush()).toBe(""); + }); + + test("an oversized malformed span survives a later valid marker in the same delta", () => { + const filter = createCitationMarkerFilter(); + const malformed = `${S}${"y".repeat(5_000)}`; + expect(filter.push(`a${span}${malformed}${S}cite${P}turn1view0${E} tail`)) + .toBe(`a${malformed} tail`); + expect(filter.flush()).toBe(""); + }); });