diff --git a/.memory/mobile-stream-recovery.md b/.memory/mobile-stream-recovery.md index 949e79827..57ad8bc7a 100644 --- a/.memory/mobile-stream-recovery.md +++ b/.memory/mobile-stream-recovery.md @@ -174,3 +174,9 @@ Cache model after the merge: `saveChats` keeps #242's committed-write ordering ( Dropped as duplicates of #242: 37b5920a's load() stream-before-draft reorder and the create() `!saveChats` fallback (#242's older-list rule admits the create's own committed row); d41ad06c's cancelledControls enum case, URLProtocol branch and restorePendingApproval delivery. The carried cancelled-Stop test now follows #242's silent-cancellation rule. Verified: iOS build-for-testing after every pick; Android testDebugUnitTest 206/206 and lint; test:ci-policy 45/45. Simulator XCTest pending (Mac authorization prompt). + +# PR #278 iOS removal-test synchronization follow-up (2026-09-27) + +The CI failure in `testRemovalCancelsAdmittedConsumerBeforeHeldEventsPublish` was not accompanied by an XCTest result bundle. Inspection found a test-harness race: its `/events` suffix gate could hold either the intended `/streams/stream-recovery/events` consumer or the unrelated `/progress/events` observer, depending on request timing. The test now holds the unique stream-consumer path and deliberately runs the denied progress observer while that hold is armed, waiting for it to finish before sending. This verifies that unrelated progress traffic cannot steal the held consumer request; no production behavior changed. + +On iPhone 17 Pro Max / iOS 27 simulator, the focused test passed 10 consecutive repetitions. The six neighboring accepted-turn/removal lifecycle tests also passed. Evidence: `/Users/sambitbiswas/Library/Developer/Xcode/DerivedData/AidenOnTheGo-dsoibpylxjwthkgezllxaeiietwn/Logs/Test/Test-AidenOnTheGo-2026.09.27_16-47-08--0400.xcresult` and `/Users/sambitbiswas/Library/Developer/Xcode/DerivedData/AidenOnTheGo-dsoibpylxjwthkgezllxaeiietwn/Logs/Test/Test-AidenOnTheGo-2026.09.27_16-48-31--0400.xcresult`. diff --git a/.memory/web-search-key-pool.md b/.memory/web-search-key-pool.md new file mode 100644 index 000000000..42a4cae23 --- /dev/null +++ b/.memory/web-search-key-pool.md @@ -0,0 +1,32 @@ +# Web Search API key pool — 2026-09-27 + +- Plan: `docs/plans/web-search-key-pool-plan.md`. Source idea: pi-web-access #453. No code was copied. +- Pool-capable providers are listed in `WEB_SEARCH_KEY_POOL_PROVIDER_IDS` in `renderer/shared/web-search-key-pool.ts`. Only `tavily` is listed today. + - Add a provider only after its adapter maps 401/403 to `auth` and its quota statuses to `quota`. +- Storage lives in `main/services/web-search-credential-core.ts`: + - The `primary` entry is the legacy single-key slot. + - Other keys are stored in `…:pool:`. + - An encrypted index at `…:pool-index` holds the order, labels and strategy. + - With no index, a saved key is treated as `[primary]`. + - `read()` returns the first pool key, so `has()` and readiness already understand pools. + - `remove()` wipes every slot. + - Pool mutations are serialized through a promise chain. +- Provider-wide removal deletes the credential's colon-delimited secret family in one encrypted-map write; cleanup does not rely on a readable index, so orphaned slots from a corrupt index or interrupted add/remove are erased. +- Runtime lives in `main/services/web-search-key-pool-core.ts`: + - `runWithWebSearchKeyPool` plus `WebSearchKeyPoolTracker`. The tracker's cooldowns and round-robin cursor are in memory only. + - The service reads `getCredentialPool` (from `web-search-main.ts`) and uses the shared `webSearchKeyPoolTracker` singleton in `web-search-credentials.ts`. + - `beforeProviderAttempt` runs once per keyed request. + - When every key is cooling, no request is sent: the pool throws `quota` if any key was quota-limited, otherwise `auth`. +- IPC lives in `main/handlers/web-search-key-pool.ts`: + - It has injected dependencies and is registered from `phase2.ts`. + - Tests register it and invoke it directly (register-and-invoke). + - `webSearch:removeCredential` also clears the tracker. +- UI: `renderer/components/settings/web-search-key-pool.tsx` replaces the single API-key field in `ProviderSetupDialog` for pool providers. + - It adds a `warning` Badge color that uses the status-warning tokens. +- `renderer/components/settings/web-search-settings.test.tsx` greps the source. Do not grow it. New UI coverage goes in `web-search-key-pool.test.tsx`, which uses renderToStaticMarkup. +- Formatting: the repo uses `oxfmt`, not prettier. Running `oxfmt` on `ui.tsx`, `ipc.ts` or `package.json` reformats unrelated code, so format only new files. +- Validation: + - `npm run test:web-search`: 175 tests pass. + - `npm run test:settings-design`: 60 tests pass. + - `npm run type-check` and scoped ESLint pass. + - No live provider requests were made. diff --git a/.papercuts/troubleshooting.md b/.papercuts/troubleshooting.md index 0c114abfa..34320dfd2 100644 --- a/.papercuts/troubleshooting.md +++ b/.papercuts/troubleshooting.md @@ -1407,6 +1407,11 @@ because their native file-mutator test binary had not been built. Run ## 2026-09-26 PR #121 merge of #251 (Remote contract revision 14) - A PR that adds to the Remote contract has to renumber when main bumps `contractRevision`. The conflicts show up in 7 files: both fixtures, the TS/iOS/Android fixture assertions and the iOS fixture CodingKeys. After resolving, `cmp` the Android copy against the shared fixture. Plan docs that name the revision also go stale. +## 2026-09-27 Web Search key pool (feature/web-search-key-pool) +- The worktree guard is inconsistent about heredocs. `cd … && python3 - <> file <> <:api-key`), which the pool calls the `primary` + entry. + - Other keys use `…:pool:`. + - An encrypted index at `…:pool-index` stores the order, labels, `addedAt` + and the strategy. It never holds key material. + - A key saved before this change, when no index exists yet, is read as a + one-entry pool. No migration step is needed. + - Provider-wide removal deletes the credential's colon-delimited secret + family in one encrypted-map write, including secondary slots omitted from + a corrupt or stale index. +- **Selection.** Two strategies: + - `ordered` always starts with the first key. + - `round-robin` starts each search one key further along. + - With either strategy, keys that are cooling down are skipped. +- **Failover.** What happens depends on the error: + + | Error | Result | + | --- | --- | + | `auth` (401/403) | The key cools down for 15 minutes, doubling on each repeat up to 24 hours. The next key is tried. | + | `quota` (429/432/433) | The key cools down for 1 minute, doubling on each repeat up to 1 hour. The next key is tried. | + | Any other error, including cancellation and timeout | The search stops immediately. | + + - Cooldowns live only in memory, so restarting Aiden clears them. + - If every key is cooling down, no request is sent. The search fails with + `quota` if any key hit a quota limit, and with `auth` otherwise. Automatic + routing then falls back to the next provider as usual. +- **Budget.** The subagent `beforeProviderAttempt` fence runs before every + keyed request, so each key tried counts against the child's network budget. +- **IPC.** The `webSearch:keyPool:get|add|remove|reorder|setStrategy|resetCooldown` + channels: + - check the owner document; + - enforce the rollout mutation fence; + - return only IDs, labels, order, strategy and cooldown state. + + Keys can be written but never read back from the renderer. +- **Settings.** In Settings → Web Search → provider setup, Tavily replaces its + single API-key field with a pool editor. The editor offers: + - an ordered key list with a position, label and status badge; + - move up and down buttons, plus Alt+Arrow keyboard reordering; + - Remove, and Retry now for a key that is cooling down; + - an optional label and a password field for adding a key; + - a strategy radio group with no card borders. +- **Logging.** Keys never appear in errors, logs or the renderer state. + +## Tests + +| Test file | Covers | +| --- | --- | +| `main/services/web-search-key-pool-core.test.ts` | Rotation and failover through the real Tavily adapter with a fake `fetch`, cooldown timing, and the index parser. | +| `main/services/web-search-credential-core.test.ts` | Pool storage, legacy primary synthesis, duplicate keys and the maximum pool size, serialization, corrupt-index cleanup, and removal of orphaned secondary slots. | +| `main/services/web-search.test.ts` | Service-level failover, charging each key attempt, and falling back to the next route. | +| `main/handlers/web-search-key-pool.test.ts` | Register-and-invoke IPC. | +| `renderer/components/settings/web-search-key-pool.test.tsx` | The rendered list's status, order and disabled states. | + +## Follow-ups + +- Extend the pool to other keyed providers once their adapters map + quota/auth statuses reliably. +- Optionally persist cooldowns across restarts. +- CLI parity for `packages/cli`, which still uses one key per provider. diff --git a/ios/AidenOnTheGoTests/AidenChatTests.swift b/ios/AidenOnTheGoTests/AidenChatTests.swift index ac7241bf5..948669c9a 100644 --- a/ios/AidenOnTheGoTests/AidenChatTests.swift +++ b/ios/AidenOnTheGoTests/AidenChatTests.swift @@ -1172,12 +1172,19 @@ final class AidenChatTests: XCTestCase { } return fixture.response(request) } - await model.load() + await model.load(observeProgress: false) model.draft = "Hello" XCTAssertTrue(model.canSend) let arrived = expectation(description: "consumer events held") - AidenChatProgressLifecycleURLProtocol.holdNextRequest(endingIn: "/events") { arrived.fulfill() } + AidenChatProgressLifecycleURLProtocol.holdNextRequest(endingIn: "/streams/stream-recovery/events") { arrived.fulfill() } defer { AidenChatProgressLifecycleURLProtocol.releaseHeldRequest() } + // The chat's unrelated progress stream also ends in `/events`. Exercise + // it while the stream-consumer hold is armed to prove it cannot steal + // the intended gate. + model.startProgressObservation() + try await waitForProgressRequestCount(1) + try await waitForProgressObservationToStop(model) + XCTAssertEqual(AidenChatProgressLifecycleURLProtocol.progressRequestCount, 1) await model.send() await fulfillment(of: [arrived], timeout: 2) let admitted = await cache.loadChat(instanceId: "instance-progress-lifecycle", chatId: model.chat.id) diff --git a/main/handlers/phase2.ts b/main/handlers/phase2.ts index ed48ab3f9..1f4cbaeb2 100644 --- a/main/handlers/phase2.ts +++ b/main/handlers/phase2.ts @@ -39,7 +39,11 @@ import { rendererDocumentOwner } from "../services/renderer-document-owner.js"; import type { RendererDocumentOwner } from "../services/renderer-document-owner.js"; import { mutatePortableConfigAndSync } from "../services/portable-credential-snapshot.js"; import { withMcpConfigurationPublication } from "../services/mcp-config-lease.js"; -import { webSearchCredentials } from "../services/web-search-credentials.js"; +import { + webSearchCredentials, + webSearchKeyPoolTracker, +} from "../services/web-search-credentials.js"; +import { registerWebSearchKeyPoolHandlers } from "./web-search-key-pool.js"; import { webSearchExistingAuthReuse } from "../services/web-search-auth-reuse-main.js"; import { DEFAULT_WEB_SEARCH_FALLBACK_ON, @@ -549,9 +553,20 @@ export function registerPhase2Handlers(): void { settings.providerConfig[providerId], ); await webSearchCredentials.remove(reference, () => !owner.isDestroyed()); + webSearchKeyPoolTracker.clear(providerId); if (owner.isDestroyed()) throw new Error("The renderer document is no longer active."); return readWebSearchSnapshot(); }); + registerWebSearchKeyPoolHandlers({ + handle: (channel, handler) => ipcMain.handle(channel, handler), + credentials: webSearchCredentials, + tracker: webSearchKeyPoolTracker, + providerConfig: async (providerId) => + (await configStore.getWebSearchSettings()).providerConfig[providerId], + owner: webSearchMutationOwner, + assertMutationAllowed: (providerId) => + assertWebSearchRolloutMutationAllowed("set-credential", providerId, webSearchRollout), + }); // Legacy Exa aliases remain for one rollback window. They use the same // fenced v2 credential path and never expose the plaintext key. diff --git a/main/handlers/web-search-key-pool.test.ts b/main/handlers/web-search-key-pool.test.ts new file mode 100644 index 000000000..b802082b5 --- /dev/null +++ b/main/handlers/web-search-key-pool.test.ts @@ -0,0 +1,173 @@ +import assert from "node:assert/strict"; +import test from "node:test"; + +import { + WEB_SEARCH_KEY_POOL_CHANNELS, + registerWebSearchKeyPoolHandlers, +} from "./web-search-key-pool.js"; +import { + createWebSearchCredentialAccess, + type WebSearchEncryptedSecretPort, +} from "../services/web-search-credential-core.js"; +import { WebSearchKeyPoolTracker } from "../services/web-search-key-pool-core.js"; +import type { WebSearchKeyPoolRendererState } from "../../renderer/shared/web-search-key-pool.js"; + +interface FakeEvent { + readonly sender: "settings" | "foreign"; +} + +function harness(options: { mutationsAllowed?: boolean } = {}) { + const bound = new Map(); + const port: WebSearchEncryptedSecretPort = { + async getProviderKey(id, binding) { + const entry = bound.get(id); + return entry?.binding === binding ? entry.key : null; + }, + async getOrBindLegacyProviderKey() { + return null; + }, + async setProviderKey(id, key, binding, isCurrent = () => true) { + if (!isCurrent()) throw new Error("stale mutation"); + bound.set(id, { key, binding }); + }, + async deleteKey(id, isCurrent = () => true) { + if (!isCurrent()) throw new Error("stale mutation"); + bound.delete(id); + }, + async deleteKeyFamily(id, isCurrent = () => true) { + if (!isCurrent()) throw new Error("stale mutation"); + for (const key of bound.keys()) { + if (key === id || key.startsWith(`${id}:`)) bound.delete(key); + } + }, + }; + let nextId = 0; + const credentials = createWebSearchCredentialAccess(port, { + createEntryId: () => `team${(nextId += 1)}`, + }); + const tracker = new WebSearchKeyPoolTracker({ now: () => 5_000 }); + const handlers = new Map unknown>(); + let destroyed = false; + registerWebSearchKeyPoolHandlers({ + handle: (channel, handler) => { + assert.equal(handlers.has(channel), false, `duplicate handler for ${channel}`); + handlers.set(channel, handler); + }, + credentials, + tracker, + providerConfig: async () => undefined, + owner: (event) => { + if (event.sender !== "settings") throw new Error("Untrusted sender."); + return { isDestroyed: () => destroyed }; + }, + assertMutationAllowed: () => { + if (options.mutationsAllowed === false) throw new Error("Credential changes are paused."); + }, + }); + const invoke = async (channel: string, ...args: unknown[]) => { + const handler = handlers.get(channel); + assert.ok(handler, `no handler for ${channel}`); + return (await handler({ sender: "settings" }, ...args)) as WebSearchKeyPoolRendererState; + }; + return { + bound, + tracker, + handlers, + invoke, + foreign: (channel: string, ...args: unknown[]) => + handlers.get(channel)!({ sender: "foreign" }, ...args), + destroy: () => { + destroyed = true; + }, + }; +} + +test("every key-pool channel is registered once", () => { + const h = harness(); + assert.deepEqual([...h.handlers.keys()].sort(), [...WEB_SEARCH_KEY_POOL_CHANNELS].sort()); +}); + +test("renderer round trip adds, reorders, and removes keys and only ever returns redacted state", async () => { + const h = harness(); + assert.deepEqual((await h.invoke("webSearch:keyPool:get", "tavily")).entries, []); + + await h.invoke("webSearch:keyPool:add", "tavily", "tvly-secret-one", "Personal"); + const added = await h.invoke("webSearch:keyPool:add", "tavily", "tvly-secret-two"); + assert.deepEqual( + added.entries.map((entry) => [entry.id, entry.label]), + [ + ["primary", "Personal"], + ["team1", "Key 2"], + ], + ); + assert.equal(added.maxEntries, 8); + + const reordered = await h.invoke("webSearch:keyPool:reorder", "tavily", ["team1", "primary"]); + assert.deepEqual( + reordered.entries.map((entry) => entry.id), + ["team1", "primary"], + ); + const strategy = await h.invoke("webSearch:keyPool:setStrategy", "tavily", "round-robin"); + assert.equal(strategy.strategy, "round-robin"); + + const removed = await h.invoke("webSearch:keyPool:remove", "tavily", "team1"); + assert.deepEqual( + removed.entries.map((entry) => entry.id), + ["primary"], + ); + + for (const state of [added, reordered, strategy, removed]) { + assert.doesNotMatch(JSON.stringify(state), /tvly-secret/u); + } +}); + +test("cooldowns are visible to the renderer and can be cleared manually", async () => { + const h = harness(); + await h.invoke("webSearch:keyPool:add", "tavily", "tvly-a"); + await h.invoke("webSearch:keyPool:add", "tavily", "tvly-b"); + h.tracker.recordFailure("tavily", "team1", "quota"); + + const cooling = await h.invoke("webSearch:keyPool:get", "tavily"); + assert.deepEqual( + cooling.entries.map((entry) => entry.cooldown?.reason ?? null), + [null, "quota"], + ); + assert.ok(cooling.entries[1]!.cooldown!.until > cooling.observedAt); + + const reset = await h.invoke("webSearch:keyPool:resetCooldown", "tavily", "team1"); + assert.equal(reset.entries[1]!.cooldown, null); + + // Removing a cooling key forgets its cooldown so a reused ID starts clean. + h.tracker.recordFailure("tavily", "team1", "auth"); + await h.invoke("webSearch:keyPool:remove", "tavily", "team1"); + assert.equal(h.tracker.cooldown("tavily", "team1"), undefined); +}); + +test("foreign senders, unsupported providers, rollout fences, and closed documents are refused", async () => { + const h = harness(); + await assert.rejects( + Promise.resolve().then(() => h.foreign("webSearch:keyPool:get", "tavily")), + /Untrusted sender/u, + ); + await assert.rejects(h.invoke("webSearch:keyPool:get", "brave"), /does not support/u); + await assert.rejects( + h.invoke("webSearch:keyPool:resetCooldown", "tavily", "../escape"), + /Invalid Web Search key entry/u, + ); + + const fenced = harness({ mutationsAllowed: false }); + await assert.rejects( + fenced.invoke("webSearch:keyPool:add", "tavily", "tvly-a"), + /Credential changes are paused/u, + ); + assert.equal(fenced.bound.size, 0); + // Reads stay available while mutations are fenced. + assert.deepEqual((await fenced.invoke("webSearch:keyPool:get", "tavily")).entries, []); + + h.destroy(); + await assert.rejects( + h.invoke("webSearch:keyPool:add", "tavily", "tvly-late"), + /no longer active/u, + ); + assert.equal(h.bound.size, 0); +}); diff --git a/main/handlers/web-search-key-pool.ts b/main/handlers/web-search-key-pool.ts new file mode 100644 index 000000000..a9d1a961c --- /dev/null +++ b/main/handlers/web-search-key-pool.ts @@ -0,0 +1,131 @@ +/** + * Web Search API-key pool IPC. + * + * Every channel returns only the redacted pool projection: opaque entry IDs, + * labels, order, strategy, and cooldown state. Key material is write-only. + * Dependencies are injected so the channel surface can be registered and + * invoked in tests without Electron. + */ + +import type { WebSearchCredentialAccess } from "../services/web-search-credential-core.js"; +import { + isWebSearchKeyPoolEntryId, + webSearchKeyPoolRendererState, + webSearchKeyPoolSupported, + type WebSearchKeyPoolDocument, + type WebSearchKeyPoolRendererState, + type WebSearchKeyPoolTracker, +} from "../services/web-search-key-pool-core.js"; +import type { + BoundedNonSecretProviderConfig, + WebSearchProviderId, +} from "../services/web-search-provider-registry-core.js"; + +export const WEB_SEARCH_KEY_POOL_CHANNELS = Object.freeze([ + "webSearch:keyPool:get", + "webSearch:keyPool:add", + "webSearch:keyPool:remove", + "webSearch:keyPool:reorder", + "webSearch:keyPool:setStrategy", + "webSearch:keyPool:resetCooldown", +] as const); + +export interface WebSearchKeyPoolHandlerDependencies { + readonly handle: ( + channel: string, + handler: (event: Event, ...args: unknown[]) => unknown, + ) => void; + readonly credentials: WebSearchCredentialAccess; + readonly tracker: WebSearchKeyPoolTracker; + readonly providerConfig: ( + providerId: WebSearchProviderId, + ) => Promise; + /** Returns a liveness check for the sender document; throws for a foreign sender. */ + readonly owner: (event: Event) => { isDestroyed(): boolean }; + /** Throws while a rollout fence forbids credential changes for this provider. */ + readonly assertMutationAllowed: (providerId: WebSearchProviderId) => void; +} + +const INACTIVE = "The renderer document is no longer active."; + +export function registerWebSearchKeyPoolHandlers( + dependencies: WebSearchKeyPoolHandlerDependencies, +): void { + const { credentials, tracker } = dependencies; + + const resolve = async (event: Event, providerId: unknown) => { + const owner = dependencies.owner(event); + if (owner.isDestroyed()) throw new Error(INACTIVE); + if (!webSearchKeyPoolSupported(providerId)) { + throw new Error("This Web Search provider does not support multiple API keys."); + } + const reference = credentials.reference( + providerId, + await dependencies.providerConfig(providerId), + ); + return { owner, providerId, reference, isCurrent: () => !owner.isDestroyed() }; + }; + + const project = ( + providerId: Parameters[0], + document: WebSearchKeyPoolDocument, + owner: { isDestroyed(): boolean }, + ): WebSearchKeyPoolRendererState => { + if (owner.isDestroyed()) throw new Error(INACTIVE); + // Removed entries must not leave cooldown state behind for a reused ID. + tracker.retain( + providerId, + document.entries.map((entry) => entry.id), + ); + return webSearchKeyPoolRendererState(providerId, document, tracker); + }; + + const mutation = ( + channel: (typeof WEB_SEARCH_KEY_POOL_CHANNELS)[number], + run: ( + context: Awaited>, + ...args: unknown[] + ) => Promise, + ) => { + dependencies.handle(channel, async (event: Event, ...args: unknown[]) => { + const [providerId, ...rest] = args; + const context = await resolve(event, providerId); + dependencies.assertMutationAllowed(context.providerId); + const document = await run(context, ...rest); + return project(context.providerId, document, context.owner); + }); + }; + + dependencies.handle("webSearch:keyPool:get", async (event: Event, ...args: unknown[]) => { + const context = await resolve(event, args[0]); + return project( + context.providerId, + await credentials.listPool(context.reference), + context.owner, + ); + }); + + mutation("webSearch:keyPool:add", (context, secret, label) => + credentials.addPoolKey(context.reference, secret, label, context.isCurrent), + ); + + mutation("webSearch:keyPool:remove", async (context, entryId) => { + const document = await credentials.removePoolKey(context.reference, entryId, context.isCurrent); + if (isWebSearchKeyPoolEntryId(entryId)) tracker.clear(context.providerId, entryId); + return document; + }); + + mutation("webSearch:keyPool:reorder", (context, orderedIds) => + credentials.reorderPool(context.reference, orderedIds, context.isCurrent), + ); + + mutation("webSearch:keyPool:setStrategy", (context, strategy) => + credentials.setPoolStrategy(context.reference, strategy, context.isCurrent), + ); + + mutation("webSearch:keyPool:resetCooldown", async (context, entryId) => { + if (!isWebSearchKeyPoolEntryId(entryId)) throw new Error("Invalid Web Search key entry."); + tracker.clear(context.providerId, entryId); + return credentials.listPool(context.reference); + }); +} diff --git a/main/services/secret-map-core.test.ts b/main/services/secret-map-core.test.ts index 233c4bdb6..4b459c6f5 100644 --- a/main/services/secret-map-core.test.ts +++ b/main/services/secret-map-core.test.ts @@ -2,6 +2,7 @@ import assert from "node:assert/strict"; import test from "node:test"; import { bindSecretEntryIfUnbound, + deleteSecretKeyFamily, deleteSecretKeyEntry, moveSecretEntryPairIfVacant, moveSecretEntryWithBindingIfVacant, @@ -26,6 +27,35 @@ test("prototype-sensitive provider IDs round-trip as own secret-map entries", () } }); +test("deleting a secret family removes descendant values and bindings only", () => { + const root = "web-search:tavily:api-key"; + const bindingPrefix = "__binding__:"; + const poolIndex = `${root}:pool-index`; + const orphanedPoolKey = `${root}:pool:orphaned-entry`; + const bindingOnlyOrphan = `${root}:pool:binding-only-orphan`; + const similarlyNamedKey = `${root}-backup`; + const map = normalizeSecretKeyMap({ + [root]: "encrypted-primary", + [`${bindingPrefix}${root}`]: "encrypted-primary-binding", + [poolIndex]: "encrypted-index", + [`${bindingPrefix}${poolIndex}`]: "encrypted-index-binding", + [orphanedPoolKey]: "encrypted-orphaned-key", + [`${bindingPrefix}${orphanedPoolKey}`]: "encrypted-orphaned-binding", + [`${bindingPrefix}${bindingOnlyOrphan}`]: "encrypted-binding-only-orphan", + [similarlyNamedKey]: "encrypted-backup", + [`${bindingPrefix}${similarlyNamedKey}`]: "encrypted-backup-binding", + unrelated: "encrypted-unrelated", + }); + + assert.equal(deleteSecretKeyFamily(map, root, bindingPrefix), true); + assert.deepEqual(map, { + [similarlyNamedKey]: "encrypted-backup", + [`${bindingPrefix}${similarlyNamedKey}`]: "encrypted-backup-binding", + unrelated: "encrypted-unrelated", + }); + assert.equal(deleteSecretKeyFamily(map, root, bindingPrefix), false); +}); + test("strict secret-map parsing rejects roots that a write must never replace", () => { for (const value of [null, [], "ciphertext", 1]) { assert.throws(() => parseSecretKeyMap(value), /Invalid encrypted secret map/u); diff --git a/main/services/secret-map-core.ts b/main/services/secret-map-core.ts index 1f5acc210..d8f8b2386 100644 --- a/main/services/secret-map-core.ts +++ b/main/services/secret-map-core.ts @@ -35,6 +35,30 @@ export function deleteSecretKeyEntry(map: SecretKeyMap, providerId: string): boo return true; } +/** Delete one secret ID and its colon-delimited descendants, including bindings. */ +export function deleteSecretKeyFamily( + map: SecretKeyMap, + rootId: string, + bindingPrefix: string, +): boolean { + if (!rootId) throw new Error("A secret family root is required."); + if (!bindingPrefix) throw new Error("A secret binding prefix is required."); + const isFamilyId = (id: string) => id === rootId || id.startsWith(`${rootId}:`); + const ids = new Set(Object.keys(map).filter(isFamilyId)); + for (const key of Object.keys(map)) { + if (key.startsWith(bindingPrefix)) { + const id = key.slice(bindingPrefix.length); + if (isFamilyId(id)) ids.add(id); + } + } + let changed = false; + for (const id of ids) { + changed = deleteSecretKeyEntry(map, id) || changed; + changed = deleteSecretKeyEntry(map, `${bindingPrefix}${id}`) || changed; + } + return changed; +} + export interface SecretEntryPair { valueId: string; bindingId: string; diff --git a/main/services/secrets.ts b/main/services/secrets.ts index 1a5ea1e84..ccf320868 100644 --- a/main/services/secrets.ts +++ b/main/services/secrets.ts @@ -8,6 +8,7 @@ import { app, logger } from "../platform.js"; import { secureStorage } from "./secure-storage.js"; import { bindSecretEntryIfUnbound, + deleteSecretKeyFamily, deleteSecretKeyEntry, moveSecretEntryPairIfVacant, moveSecretEntryWithBindingIfVacant, @@ -407,6 +408,18 @@ export const secrets = { }); }, + /** Remove one namespaced credential and every colon-delimited child in one map write. */ + async deleteKeyFamily(providerId: string, isCurrent: MutationGuard = () => true): Promise { + return serialized(async () => { + assertMutationCurrent(isCurrent); + const map = await readMap(); + assertMutationCurrent(isCurrent); + if (deleteSecretKeyFamily(map, providerId, PROVIDER_BINDING_PREFIX)) { + await writeMap(map, isCurrent); + } + }); + }, + async clearAll(isCurrent: MutationGuard = () => true): Promise { return serialized(async () => { assertMutationCurrent(isCurrent); diff --git a/main/services/web-search-credential-core.test.ts b/main/services/web-search-credential-core.test.ts index 08f90bbb1..7ec7e28ef 100644 --- a/main/services/web-search-credential-core.test.ts +++ b/main/services/web-search-credential-core.test.ts @@ -9,6 +9,7 @@ import { webSearchCredentialReference, type WebSearchEncryptedSecretPort, } from "./web-search-credential-core.js"; +import { webSearchKeyPoolIndexSecretId } from "./web-search-key-pool-core.js"; function fakeSecrets(initial: { legacy?: Record } = {}) { const bound = new Map(); @@ -39,6 +40,13 @@ function fakeSecrets(initial: { legacy?: Record } = {}) { bound.delete(providerId); legacy.delete(providerId); }, + async deleteKeyFamily(providerId, isCurrent = () => true) { + if (!isCurrent()) throw new Error("stale mutation"); + calls.push({ operation: "delete-family", providerId }); + for (const key of bound.keys()) { + if (key === providerId || key.startsWith(`${providerId}:`)) bound.delete(key); + } + }, }; return { bound, legacy, calls, access: createWebSearchCredentialAccess(port) }; } @@ -170,3 +178,176 @@ test("credential mutation guards are forwarded before publication", async () => assert.equal(h.bound.size, 0); assert.equal(h.calls.length, 0); }); + +function poolSecrets() { + const bound = new Map(); + let nextId = 0; + let clock = 100; + const port: WebSearchEncryptedSecretPort = { + async getProviderKey(providerId, binding) { + const entry = bound.get(providerId); + return entry?.binding === binding ? entry.key : null; + }, + async getOrBindLegacyProviderKey() { + return null; + }, + async setProviderKey(providerId, key, binding, isCurrent = () => true) { + if (!isCurrent()) throw new Error("stale mutation"); + bound.set(providerId, { key, binding }); + }, + async deleteKey(providerId, isCurrent = () => true) { + if (!isCurrent()) throw new Error("stale mutation"); + bound.delete(providerId); + }, + async deleteKeyFamily(providerId, isCurrent = () => true) { + if (!isCurrent()) throw new Error("stale mutation"); + for (const key of bound.keys()) { + if (key === providerId || key.startsWith(`${providerId}:`)) bound.delete(key); + } + }, + }; + const access = createWebSearchCredentialAccess(port, { + createEntryId: () => `entry${(nextId += 1)}`, + now: () => (clock += 1), + }); + return { bound, access }; +} + +test("a key saved before pools existed becomes the first pool entry", async () => { + const h = poolSecrets(); + const tavily = h.access.reference("tavily"); + await h.access.set(tavily, "tvly-legacy"); + + const listed = await h.access.listPool(tavily); + assert.deepEqual( + listed.entries.map((entry) => entry.id), + ["primary"], + ); + assert.deepEqual(await h.access.readPool(tavily), { + strategy: "ordered", + keys: [{ id: "primary", key: "tvly-legacy" }], + }); + assert.equal(await h.access.read(tavily), "tvly-legacy"); +}); + +test("provider removal deletes orphaned pool secrets without trusting the index", async () => { + const h = poolSecrets(); + const tavily = h.access.reference("tavily"); + const indexId = webSearchKeyPoolIndexSecretId(tavily.secretId); + const orphanedId = `${tavily.secretId}:pool:orphaned-entry`; + + h.bound.set(tavily.secretId, { key: "primary-key", binding: tavily.binding }); + h.bound.set(indexId, { key: "{corrupt", binding: tavily.binding }); + h.bound.set(orphanedId, { key: "unindexed-secondary-key", binding: tavily.binding }); + + await h.access.remove(tavily); + + assert.equal(h.bound.has(tavily.secretId), false); + assert.equal(h.bound.has(indexId), false); + assert.equal(h.bound.has(orphanedId), false); +}); + +test("provider pool removal leaves secrets intact when its mutation owner is stale", async () => { + const h = poolSecrets(); + const tavily = h.access.reference("tavily"); + const orphanedId = `${tavily.secretId}:pool:orphaned-entry`; + h.bound.set(tavily.secretId, { key: "primary-key", binding: tavily.binding }); + h.bound.set(orphanedId, { key: "secondary-key", binding: tavily.binding }); + + await assert.rejects( + h.access.remove(tavily, () => false), + /stale mutation/u, + ); + + assert.equal(h.bound.has(tavily.secretId), true); + assert.equal(h.bound.has(orphanedId), true); +}); + +test("pool keys can be added, reordered, re-strategized, and removed without exposing keys", async () => { + const h = poolSecrets(); + const tavily = h.access.reference("tavily"); + + await h.access.addPoolKey(tavily, "tvly-one"); + await h.access.addPoolKey(tavily, "tvly-two", "Team"); + const added = await h.access.addPoolKey(tavily, "tvly-three"); + assert.deepEqual( + added.entries.map((entry) => [entry.id, entry.label]), + [ + ["primary", "Key 1"], + ["entry1", "Team"], + ["entry2", "Key 3"], + ], + ); + assert.doesNotMatch(JSON.stringify(added), /tvly-/u); + // The first pooled key lives in the pre-pool slot for single-key readers. + assert.equal(await h.access.read(tavily), "tvly-one"); + + await assert.rejects(h.access.addPoolKey(tavily, " tvly-two "), /already in the pool/u); + + await h.access.reorderPool(tavily, ["entry2", "primary", "entry1"]); + assert.deepEqual( + (await h.access.readPool(tavily)).keys.map((entry) => entry.key), + ["tvly-three", "tvly-one", "tvly-two"], + ); + assert.equal(await h.access.read(tavily), "tvly-three"); + await assert.rejects(h.access.reorderPool(tavily, ["entry2", "primary"]), /no longer matches/u); + + const strategy = await h.access.setPoolStrategy(tavily, "round-robin"); + assert.equal(strategy.strategy, "round-robin"); + await assert.rejects(h.access.setPoolStrategy(tavily, "random")); + + const removed = await h.access.removePoolKey(tavily, "entry2"); + assert.deepEqual( + removed.entries.map((entry) => entry.id), + ["primary", "entry1"], + ); + assert.equal(await h.access.read(tavily), "tvly-one"); + assert.ok(![...h.bound.values()].some((entry) => entry.key === "tvly-three")); + await assert.rejects(h.access.removePoolKey(tavily, "entry2"), /no longer in the pool/u); + + // Removing the provider credential clears every pooled slot and the index. + await h.access.remove(tavily); + assert.equal(await h.access.has(tavily), false); + assert.deepEqual(await h.access.readPool(tavily), { strategy: "ordered", keys: [] }); + assert.equal(h.bound.size, 0); +}); + +test("the pool refuses a ninth key and non-pool providers keep single-key behavior", async () => { + const h = poolSecrets(); + const tavily = h.access.reference("tavily"); + for (let index = 0; index < 8; index += 1) await h.access.addPoolKey(tavily, `tvly-${index}`); + await assert.rejects(h.access.addPoolKey(tavily, "tvly-9"), /at most 8 keys/u); + + const brave = h.access.reference("brave"); + await assert.rejects(h.access.listPool(brave), /does not support multiple API keys/u); + await h.access.set(brave, "brave-key"); + assert.deepEqual(await h.access.readPool(brave), { + strategy: "ordered", + keys: [{ id: "primary", key: "brave-key" }], + }); +}); + +test("concurrent pool edits are serialized so neither is lost", async () => { + const h = poolSecrets(); + const tavily = h.access.reference("tavily"); + await Promise.all([ + h.access.addPoolKey(tavily, "tvly-a"), + h.access.addPoolKey(tavily, "tvly-b"), + h.access.addPoolKey(tavily, "tvly-c"), + ]); + assert.deepEqual((await h.access.readPool(tavily)).keys.map((entry) => entry.key).sort(), [ + "tvly-a", + "tvly-b", + "tvly-c", + ]); +}); + +test("a corrupt pool index falls back to the saved primary key", async () => { + const h = poolSecrets(); + const tavily = h.access.reference("tavily"); + await h.access.addPoolKey(tavily, "tvly-a"); + await h.access.addPoolKey(tavily, "tvly-b"); + const indexSlot = [...h.bound.keys()].find((id) => id.endsWith(":pool-index"))!; + h.bound.set(indexSlot, { key: "{not json", binding: tavily.binding }); + assert.deepEqual((await h.access.readPool(tavily)).keys, [{ id: "primary", key: "tvly-a" }]); +}); diff --git a/main/services/web-search-credential-core.ts b/main/services/web-search-credential-core.ts index b1d204028..c43f06bbb 100644 --- a/main/services/web-search-credential-core.ts +++ b/main/services/web-search-credential-core.ts @@ -5,6 +5,23 @@ import { type BoundedNonSecretProviderConfig, type WebSearchProviderId, } from "./web-search-provider-registry-core.js"; +import { + EMPTY_WEB_SEARCH_KEY_POOL_DOCUMENT, + MAX_WEB_SEARCH_KEY_POOL_ENTRIES, + WEB_SEARCH_KEY_POOL_PRIMARY_ENTRY_ID, + isWebSearchKeyPoolEntryId, + isWebSearchKeyPoolStrategy, + normalizeWebSearchKeyPoolLabel, + parseWebSearchKeyPoolDocument, + reorderWebSearchKeyPoolEntries, + serializeWebSearchKeyPoolDocument, + webSearchKeyPoolEntrySecretId, + webSearchKeyPoolIndexSecretId, + webSearchKeyPoolSupported, + type WebSearchKeyPool, + type WebSearchKeyPoolDocument, + type WebSearchPooledKey, +} from "./web-search-key-pool-core.js"; /** Stable namespace for Web Search credentials in the encrypted key store. */ export const WEB_SEARCH_CREDENTIAL_PREFIX = "web-search:"; @@ -33,6 +50,7 @@ export interface WebSearchEncryptedSecretPort { isCurrent?: () => boolean, ): Promise; deleteKey(providerId: string, isCurrent?: () => boolean): Promise; + deleteKeyFamily(providerId: string, isCurrent?: () => boolean): Promise; } export interface WebSearchCredentialAccess { @@ -48,6 +66,44 @@ export interface WebSearchCredentialAccess { isCurrent?: () => boolean, ): Promise; remove(reference: WebSearchCredentialReference, isCurrent?: () => boolean): Promise; + /** + * Non-secret pool index for a pool-capable provider: order, labels, and + * strategy of the keys that currently exist. Never includes key material. + */ + listPool(reference: WebSearchCredentialReference): Promise; + /** Main-only: every existing pooled key, in index order. */ + readPool(reference: WebSearchCredentialReference): Promise; + addPoolKey( + reference: WebSearchCredentialReference, + key: unknown, + label?: unknown, + isCurrent?: () => boolean, + ): Promise; + removePoolKey( + reference: WebSearchCredentialReference, + entryId: unknown, + isCurrent?: () => boolean, + ): Promise; + reorderPool( + reference: WebSearchCredentialReference, + orderedIds: unknown, + isCurrent?: () => boolean, + ): Promise; + setPoolStrategy( + reference: WebSearchCredentialReference, + strategy: unknown, + isCurrent?: () => boolean, + ): Promise; +} + +export interface WebSearchCredentialAccessOptions { + /** Opaque pool entry IDs; injectable for deterministic tests. */ + readonly createEntryId?: () => string; + readonly now?: () => number; +} + +function defaultEntryId(): string { + return globalThis.crypto.randomUUID().replace(/-/gu, "").slice(0, 16); } function hasControlCharacter(value: string): boolean { @@ -168,33 +224,265 @@ export function webSearchCredentialReference( /** Main-only credential access with exact compatibility for the legacy Exa slot. */ export function createWebSearchCredentialAccess( secrets: WebSearchEncryptedSecretPort, + options: WebSearchCredentialAccessOptions = {}, ): WebSearchCredentialAccess { - return { + const createEntryId = options.createEntryId ?? defaultEntryId; + const now = options.now ?? Date.now; + // Pool index writes are read-modify-write; serialize them per process so + // two quick Settings actions cannot drop each other's change. + let poolTail: Promise = Promise.resolve(); + const serializedPool = (operation: () => Promise): Promise => { + const result = poolTail.then(operation, operation); + poolTail = result.catch(() => undefined); + return result; + }; + + const readPrimary = async (reference: WebSearchCredentialReference) => { + const current = await secrets.getProviderKey(reference.secretId, reference.binding); + if (current !== null) return current; + if (reference.legacySecretId) { + return secrets.getOrBindLegacyProviderKey(reference.legacySecretId, reference.binding); + } + return null; + }; + + const readEntryKey = (reference: WebSearchCredentialReference, entryId: string) => + entryId === WEB_SEARCH_KEY_POOL_PRIMARY_ENTRY_ID + ? readPrimary(reference) + : secrets.getProviderKey( + webSearchKeyPoolEntrySecretId(reference.secretId, entryId), + reference.binding, + ); + + /** + * Resolve the index against the keys that actually exist. A key saved + * before pools existed (no index) becomes the single primary entry. + */ + const loadPool = async ( + reference: WebSearchCredentialReference, + ): Promise<{ document: WebSearchKeyPoolDocument; keys: WebSearchPooledKey[] }> => { + const stored = parseWebSearchKeyPoolDocument( + await secrets.getProviderKey( + webSearchKeyPoolIndexSecretId(reference.secretId), + reference.binding, + ), + ); + const document: WebSearchKeyPoolDocument = stored ?? { + ...EMPTY_WEB_SEARCH_KEY_POOL_DOCUMENT, + entries: [{ id: WEB_SEARCH_KEY_POOL_PRIMARY_ENTRY_ID, label: "Key 1", addedAt: 0 }], + }; + const loaded = await Promise.all( + document.entries.map(async (entry) => ({ + entry, + key: await readEntryKey(reference, entry.id), + })), + ); + const present = loaded.flatMap((item) => + item.key === null ? [] : [{ entry: item.entry, key: item.key }], + ); + return { + document: { ...document, entries: present.map((item) => item.entry) }, + keys: present.map((item) => ({ id: item.entry.id, key: item.key })), + }; + }; + + const assertPool = (reference: WebSearchCredentialReference) => { + if (!webSearchKeyPoolSupported(reference.providerId)) { + throw new Error("This Web Search provider does not support multiple API keys."); + } + }; + + const writeIndex = ( + reference: WebSearchCredentialReference, + document: WebSearchKeyPoolDocument, + isCurrent: () => boolean, + ) => + secrets.setProviderKey( + webSearchKeyPoolIndexSecretId(reference.secretId), + serializeWebSearchKeyPoolDocument(document), + reference.binding, + isCurrent, + ); + + const access: WebSearchCredentialAccess = { reference: webSearchCredentialReference, async read(reference) { - const current = await secrets.getProviderKey(reference.secretId, reference.binding); - if (current !== null) return current; - if (reference.legacySecretId) { - return secrets.getOrBindLegacyProviderKey(reference.legacySecretId, reference.binding); - } - return null; + if (!webSearchKeyPoolSupported(reference.providerId)) return readPrimary(reference); + // A pool-capable provider's single-key view is its first existing key. + await poolTail; + const { keys } = await loadPool(reference); + return keys[0]?.key ?? null; }, async has(reference) { - return (await this.read(reference)) !== null; + return (await access.read(reference)) !== null; }, async set(reference, key, isCurrent = () => true) { const normalized = normalizeWebSearchCredential(key); - await secrets.setProviderKey(reference.secretId, normalized, reference.binding, isCurrent); + if (!webSearchKeyPoolSupported(reference.providerId)) { + await secrets.setProviderKey(reference.secretId, normalized, reference.binding, isCurrent); + return; + } + await serializedPool(async () => { + const { document } = await loadPool(reference); + await secrets.setProviderKey(reference.secretId, normalized, reference.binding, isCurrent); + if (document.entries.some((entry) => entry.id === WEB_SEARCH_KEY_POOL_PRIMARY_ENTRY_ID)) { + return; + } + // Replacing through the single-key path re-adds the primary slot at + // the front so it is used first, matching the pre-pool behavior. + await writeIndex( + reference, + { + ...document, + entries: [ + { id: WEB_SEARCH_KEY_POOL_PRIMARY_ENTRY_ID, label: "Key 1", addedAt: now() }, + ...document.entries, + ].slice(0, MAX_WEB_SEARCH_KEY_POOL_ENTRIES), + }, + isCurrent, + ); + }); }, async remove(reference, isCurrent = () => true) { - await secrets.deleteKey(reference.secretId, isCurrent); - // Removal is provider-scoped. Clear the old Exa slot too so a removed - // credential cannot silently reappear through the compatibility path. - if (reference.legacySecretId) await secrets.deleteKey(reference.legacySecretId, isCurrent); + const removeAll = async () => { + if (webSearchKeyPoolSupported(reference.providerId)) { + // Provider removal is an explicit request to erase every pooled slot. + // Do not trust the index: it may be corrupt or may omit a slot left + // by an interrupted add/remove mutation. + await secrets.deleteKeyFamily(reference.secretId, isCurrent); + } else { + await secrets.deleteKey(reference.secretId, isCurrent); + } + // Removal is provider-scoped. Clear the old Exa slot too so a removed + // credential cannot silently reappear through the compatibility path. + if (reference.legacySecretId) await secrets.deleteKey(reference.legacySecretId, isCurrent); + }; + if (webSearchKeyPoolSupported(reference.providerId)) await serializedPool(removeAll); + else await removeAll(); + }, + + async listPool(reference) { + assertPool(reference); + await poolTail; + return (await loadPool(reference)).document; + }, + + async readPool(reference) { + if (!webSearchKeyPoolSupported(reference.providerId)) { + const key = await readPrimary(reference); + return { + strategy: "ordered", + keys: key === null ? [] : [{ id: WEB_SEARCH_KEY_POOL_PRIMARY_ENTRY_ID, key }], + }; + } + await poolTail; + const { document, keys } = await loadPool(reference); + return { strategy: document.strategy, keys }; + }, + + async addPoolKey(reference, key, label, isCurrent = () => true) { + assertPool(reference); + const normalized = normalizeWebSearchCredential(key); + return serializedPool(async () => { + const { document, keys } = await loadPool(reference); + if (document.entries.length >= MAX_WEB_SEARCH_KEY_POOL_ENTRIES) { + throw new Error( + `A Web Search key pool holds at most ${MAX_WEB_SEARCH_KEY_POOL_ENTRIES} keys.`, + ); + } + if (keys.some((existing) => existing.key === normalized)) { + throw new Error("This API key is already in the pool."); + } + // Use the pre-pool slot first so older single-key readers keep working. + const primaryFree = !document.entries.some( + (entry) => entry.id === WEB_SEARCH_KEY_POOL_PRIMARY_ENTRY_ID, + ); + let entryId = primaryFree ? WEB_SEARCH_KEY_POOL_PRIMARY_ENTRY_ID : createEntryId(); + if ( + !isWebSearchKeyPoolEntryId(entryId) || + document.entries.some((entry) => entry.id === entryId) + ) { + entryId = defaultEntryId(); + } + const next: WebSearchKeyPoolDocument = { + ...document, + entries: [ + ...document.entries, + { + id: entryId, + label: normalizeWebSearchKeyPoolLabel(label, `Key ${document.entries.length + 1}`), + addedAt: now(), + }, + ], + }; + await secrets.setProviderKey( + webSearchKeyPoolEntrySecretId(reference.secretId, entryId), + normalized, + reference.binding, + isCurrent, + ); + await writeIndex(reference, next, isCurrent); + return next; + }); + }, + + async removePoolKey(reference, entryId, isCurrent = () => true) { + assertPool(reference); + return serializedPool(async () => { + const { document } = await loadPool(reference); + if ( + !isWebSearchKeyPoolEntryId(entryId) || + !document.entries.some((entry) => entry.id === entryId) + ) { + throw new Error("That API key is no longer in the pool."); + } + const next: WebSearchKeyPoolDocument = { + ...document, + entries: document.entries.filter((entry) => entry.id !== entryId), + }; + // Publish the smaller index before deleting the key so an interrupted + // removal leaves an unreferenced secret, never a dangling entry. + await writeIndex(reference, next, isCurrent); + await secrets.deleteKey( + webSearchKeyPoolEntrySecretId(reference.secretId, entryId), + isCurrent, + ); + if (entryId === WEB_SEARCH_KEY_POOL_PRIMARY_ENTRY_ID && reference.legacySecretId) { + await secrets.deleteKey(reference.legacySecretId, isCurrent); + } + return next; + }); + }, + + async reorderPool(reference, orderedIds, isCurrent = () => true) { + assertPool(reference); + return serializedPool(async () => { + const { document } = await loadPool(reference); + const next: WebSearchKeyPoolDocument = { + ...document, + entries: reorderWebSearchKeyPoolEntries(document.entries, orderedIds), + }; + await writeIndex(reference, next, isCurrent); + return next; + }); + }, + + async setPoolStrategy(reference, strategy, isCurrent = () => true) { + assertPool(reference); + if (!isWebSearchKeyPoolStrategy(strategy)) { + throw new Error("Unknown Web Search key rotation strategy."); + } + return serializedPool(async () => { + const { document } = await loadPool(reference); + const next: WebSearchKeyPoolDocument = { ...document, strategy }; + await writeIndex(reference, next, isCurrent); + return next; + }); }, }; + return access; } diff --git a/main/services/web-search-credentials.ts b/main/services/web-search-credentials.ts index 5aaff7648..d9d94971c 100644 --- a/main/services/web-search-credentials.ts +++ b/main/services/web-search-credentials.ts @@ -3,7 +3,14 @@ import { createWebSearchCredentialAccess, type WebSearchCredentialAccess, } from "./web-search-credential-core.js"; +import { WebSearchKeyPoolTracker } from "./web-search-key-pool-core.js"; /** Electron-bound Web Search credential access; plaintext stays in main. */ export const webSearchCredentials: WebSearchCredentialAccess = createWebSearchCredentialAccess(secrets); + +/** + * Process-local key-pool cooldown and rotation state. The search service + * records failures here; Settings reads the redacted cooldown projection. + */ +export const webSearchKeyPoolTracker = new WebSearchKeyPoolTracker(); diff --git a/main/services/web-search-key-pool-core.test.ts b/main/services/web-search-key-pool-core.test.ts new file mode 100644 index 000000000..16847967c --- /dev/null +++ b/main/services/web-search-key-pool-core.test.ts @@ -0,0 +1,363 @@ +import assert from "node:assert/strict"; +import { readFileSync } from "node:fs"; +import test from "node:test"; + +import { WebSearchError } from "./web-search-core.js"; +import { + EMPTY_WEB_SEARCH_KEY_POOL_DOCUMENT, + WebSearchKeyPoolTracker, + formatWebSearchKeyPoolCooldown, + moveWebSearchKeyPoolEntry, + normalizeWebSearchKeyPoolLabel, + parseWebSearchKeyPoolDocument, + reorderWebSearchKeyPoolEntries, + runWithWebSearchKeyPool, + serializeWebSearchKeyPoolDocument, + webSearchKeyPoolRendererState, + type WebSearchKeyPool, +} from "./web-search-key-pool-core.js"; +import { createTavilyWebSearchAdapter } from "./web-search-tavily-adapter.js"; + +const SUCCESS_BODY = readFileSync( + new URL("./fixtures/tavily-search/json-success.json", import.meta.url), + "utf8", +); + +const MINUTE = 60_000; + +function fakeClock(start = 1_000_000) { + let now = start; + return { + now: () => now, + advance: (ms: number) => { + now += ms; + }, + }; +} + +/** + * Fake Tavily endpoint: each key has a scripted status. Records the bearer + * key of every request actually sent. + */ +function fakeTavily(statusByKey: Record) { + const sent: string[] = []; + const fetch = (async (_url: string | URL | Request, init?: RequestInit) => { + const auth = new Headers(init?.headers).get("authorization") ?? ""; + const key = auth.replace(/^Bearer /u, ""); + sent.push(key); + const status = statusByKey[key] ?? 500; + return status === 200 + ? new Response(SUCCESS_BODY, { + status: 200, + headers: { "content-type": "application/json" }, + }) + : new Response(`{"detail":"rejected ${key}"}`, { status }); + }) as typeof globalThis.fetch; + return { sent, adapter: createTavilyWebSearchAdapter({ fetch }) }; +} + +function pool(strategy: WebSearchKeyPool["strategy"], ...keys: string[]): WebSearchKeyPool { + return { strategy, keys: keys.map((key, index) => ({ id: `k${index + 1}`, key })) }; +} + +function search(adapter: ReturnType["adapter"], key: string) { + return adapter.search({ + query: "pool failover", + numResults: 2, + credentialMode: "api-key", + credential: key, + signal: new AbortController().signal, + }); +} + +test("ordered pools fail over past rejected and rate-limited keys to a healthy one", async () => { + const clock = fakeClock(); + const tracker = new WebSearchKeyPoolTracker({ now: clock.now }); + const tavily = fakeTavily({ "tvly-a": 401, "tvly-b": 429, "tvly-c": 200 }); + const attempts: string[] = []; + + const result = await runWithWebSearchKeyPool({ + providerId: "tavily", + pool: pool("ordered", "tvly-a", "tvly-b", "tvly-c"), + tracker, + beforeKeyAttempt: (entryId) => { + attempts.push(entryId); + }, + run: (key) => search(tavily.adapter, key), + }); + + assert.ok(result.results.length > 0); + assert.deepEqual(tavily.sent, ["tvly-a", "tvly-b", "tvly-c"]); + assert.deepEqual(attempts, ["k1", "k2", "k3"]); + assert.equal(tracker.cooldown("tavily", "k1")?.reason, "auth"); + assert.equal(tracker.cooldown("tavily", "k2")?.reason, "quota"); + assert.equal(tracker.cooldown("tavily", "k3"), undefined); + + // The next search skips both cooling keys and goes straight to the healthy one. + tavily.sent.length = 0; + await runWithWebSearchKeyPool({ + providerId: "tavily", + pool: pool("ordered", "tvly-a", "tvly-b", "tvly-c"), + tracker, + run: (key) => search(tavily.adapter, key), + }); + assert.deepEqual(tavily.sent, ["tvly-c"]); +}); + +test("Tavily's plan-quota status codes also rotate to the next key", async () => { + const tracker = new WebSearchKeyPoolTracker({ now: fakeClock().now }); + const tavily = fakeTavily({ "tvly-a": 432, "tvly-b": 433, "tvly-c": 200 }); + await runWithWebSearchKeyPool({ + providerId: "tavily", + pool: pool("ordered", "tvly-a", "tvly-b", "tvly-c"), + tracker, + run: (key) => search(tavily.adapter, key), + }); + assert.deepEqual(tavily.sent, ["tvly-a", "tvly-b", "tvly-c"]); + assert.equal(tracker.cooldown("tavily", "k1")?.reason, "quota"); + assert.equal(tracker.cooldown("tavily", "k2")?.reason, "quota"); +}); + +test("non-credential failures stop immediately instead of burning other keys", async () => { + const tracker = new WebSearchKeyPoolTracker({ now: fakeClock().now }); + const tavily = fakeTavily({ "tvly-a": 500, "tvly-b": 200 }); + await assert.rejects( + runWithWebSearchKeyPool({ + providerId: "tavily", + pool: pool("ordered", "tvly-a", "tvly-b"), + tracker, + run: (key) => search(tavily.adapter, key), + }), + (error: unknown) => error instanceof WebSearchError && error.kind !== "auth", + ); + assert.deepEqual(tavily.sent, ["tvly-a"]); + assert.equal(tracker.cooldown("tavily", "k1"), undefined); +}); + +test("when every key is cooling, no request is sent and the error names the recoverable kind", async () => { + const clock = fakeClock(); + const tracker = new WebSearchKeyPoolTracker({ now: clock.now }); + const tavily = fakeTavily({ "tvly-a": 401, "tvly-b": 429 }); + const keys = pool("ordered", "tvly-a", "tvly-b"); + + await assert.rejects( + runWithWebSearchKeyPool({ + providerId: "tavily", + pool: keys, + tracker, + run: (key) => search(tavily.adapter, key), + }), + (error: unknown) => error instanceof WebSearchError && error.kind === "quota", + ); + assert.equal(tavily.sent.length, 2); + + tavily.sent.length = 0; + await assert.rejects( + runWithWebSearchKeyPool({ + providerId: "tavily", + pool: keys, + tracker, + run: (key) => search(tavily.adapter, key), + }), + (error: unknown) => + error instanceof WebSearchError && + error.kind === "quota" && + !error.message.includes("tvly-a") && + !error.message.includes("tvly-b"), + ); + assert.deepEqual(tavily.sent, [], "a fully cooling pool must not touch the network"); + + // A pool whose only keys were rejected reports an auth failure instead. + const authOnly = new WebSearchKeyPoolTracker({ now: clock.now }); + authOnly.recordFailure("tavily", "k1", "auth"); + await assert.rejects( + runWithWebSearchKeyPool({ + providerId: "tavily", + pool: pool("ordered", "tvly-a"), + tracker: authOnly, + run: (key) => search(tavily.adapter, key), + }), + (error: unknown) => error instanceof WebSearchError && error.kind === "auth", + ); + assert.deepEqual(tavily.sent, []); +}); + +test("cooldowns expire on the clock and escalate on repeated failures of the same kind", () => { + const clock = fakeClock(); + const tracker = new WebSearchKeyPoolTracker({ now: clock.now }); + + const first = tracker.recordFailure("tavily", "k1", "quota"); + assert.equal(first.until - clock.now(), MINUTE); + const second = tracker.recordFailure("tavily", "k1", "quota"); + assert.equal(second.until - clock.now(), 2 * MINUTE); + + clock.advance(2 * MINUTE - 1); + assert.ok(tracker.cooldown("tavily", "k1")); + clock.advance(1); + assert.equal(tracker.cooldown("tavily", "k1"), undefined); + + // Quota escalation is capped at one hour. + for (let index = 0; index < 12; index += 1) tracker.recordFailure("tavily", "k2", "quota"); + assert.equal(tracker.cooldown("tavily", "k2")!.until - clock.now(), 60 * MINUTE); + + // A different failure kind restarts from that kind's base, and auth cools longer. + const auth = tracker.recordFailure("tavily", "k2", "auth"); + assert.equal(auth.consecutiveFailures, 1); + assert.equal(auth.until - clock.now(), 15 * MINUTE); + + // Success clears the record so the next failure starts from base again. + tracker.recordSuccess("tavily", "k2"); + assert.equal(tracker.recordFailure("tavily", "k2", "quota").until - clock.now(), MINUTE); +}); + +test("round-robin starts each search on the next key while ordered always starts first", async () => { + const tracker = new WebSearchKeyPoolTracker({ now: fakeClock().now }); + const tavily = fakeTavily({ "tvly-a": 200, "tvly-b": 200, "tvly-c": 200 }); + for (let index = 0; index < 4; index += 1) { + await runWithWebSearchKeyPool({ + providerId: "tavily", + pool: pool("round-robin", "tvly-a", "tvly-b", "tvly-c"), + tracker, + run: (key) => search(tavily.adapter, key), + }); + } + assert.deepEqual(tavily.sent, ["tvly-a", "tvly-b", "tvly-c", "tvly-a"]); + + tavily.sent.length = 0; + for (let index = 0; index < 3; index += 1) { + await runWithWebSearchKeyPool({ + providerId: "tavily", + pool: pool("ordered", "tvly-a", "tvly-b", "tvly-c"), + tracker: new WebSearchKeyPoolTracker(), + run: (key) => search(tavily.adapter, key), + }); + } + assert.deepEqual(tavily.sent, ["tvly-a", "tvly-a", "tvly-a"]); +}); + +test("round-robin skips a cooling key and cancellation stops failover", async () => { + const tracker = new WebSearchKeyPoolTracker({ now: fakeClock().now }); + tracker.recordFailure("tavily", "k2", "quota"); + const tavily = fakeTavily({ "tvly-a": 200, "tvly-b": 200, "tvly-c": 200 }); + const keys = pool("round-robin", "tvly-a", "tvly-b", "tvly-c"); + for (let index = 0; index < 3; index += 1) { + await runWithWebSearchKeyPool({ + providerId: "tavily", + pool: keys, + tracker, + run: (key) => search(tavily.adapter, key), + }); + } + assert.deepEqual(tavily.sent, ["tvly-a", "tvly-c", "tvly-c"]); + + const controller = new AbortController(); + const aborting = fakeTavily({ "tvly-a": 429, "tvly-b": 200 }); + await assert.rejects( + runWithWebSearchKeyPool({ + providerId: "tavily", + pool: pool("ordered", "tvly-a", "tvly-b"), + tracker: new WebSearchKeyPoolTracker(), + signal: controller.signal, + run: async (key) => { + try { + return await search(aborting.adapter, key); + } finally { + controller.abort(); + } + }, + }), + (error: unknown) => error instanceof WebSearchError && error.kind === "cancelled", + ); + assert.deepEqual(aborting.sent, ["tvly-a"]); +}); + +test("the renderer projection carries labels and cooldowns but never key material", () => { + const clock = fakeClock(); + const tracker = new WebSearchKeyPoolTracker({ now: clock.now }); + tracker.recordFailure("tavily", "primary", "quota"); + const state = webSearchKeyPoolRendererState( + "tavily", + { + version: 1, + strategy: "round-robin", + entries: [ + { id: "primary", label: "Personal", addedAt: 10 }, + { id: "abc123", label: "Team", addedAt: 20 }, + ], + }, + tracker, + ); + assert.equal(state.strategy, "round-robin"); + assert.deepEqual( + state.entries.map((entry) => [entry.id, entry.label, entry.cooldown?.reason ?? null]), + [ + ["primary", "Personal", "quota"], + ["abc123", "Team", null], + ], + ); + assert.equal(state.entries[0]!.cooldown!.until, clock.now() + MINUTE); + for (const entry of state.entries) { + assert.deepEqual(Object.keys(entry).sort(), ["addedAt", "cooldown", "id", "label"]); + } +}); + +test("the encrypted pool index round-trips and rejects malformed or oversized documents", () => { + const document = { + version: 1 as const, + strategy: "ordered" as const, + entries: [ + { id: "primary", label: "Key 1", addedAt: 1 }, + { id: "a1b2", label: "Key 2", addedAt: 2 }, + ], + }; + assert.deepEqual( + parseWebSearchKeyPoolDocument(serializeWebSearchKeyPoolDocument(document)), + document, + ); + assert.deepEqual( + parseWebSearchKeyPoolDocument( + serializeWebSearchKeyPoolDocument(EMPTY_WEB_SEARCH_KEY_POOL_DOCUMENT), + ), + EMPTY_WEB_SEARCH_KEY_POOL_DOCUMENT, + ); + for (const raw of [ + "not json", + JSON.stringify({ ...document, version: 2 }), + JSON.stringify({ ...document, strategy: "random" }), + JSON.stringify({ ...document, entries: [document.entries[0], document.entries[0]] }), + JSON.stringify({ ...document, entries: [{ id: "../x", label: "x", addedAt: 1 }] }), + JSON.stringify({ + ...document, + entries: Array.from({ length: 9 }, (_, index) => ({ + id: `k${index}`, + label: "x", + addedAt: 1, + })), + }), + ]) { + assert.equal(parseWebSearchKeyPoolDocument(raw), undefined, raw.slice(0, 60)); + } +}); + +test("shared list helpers reorder exact permutations and format remaining cooldown", () => { + const entries = [{ id: "a" }, { id: "b" }, { id: "c" }]; + assert.deepEqual(moveWebSearchKeyPoolEntry(entries, 2, -1), ["a", "c", "b"]); + assert.deepEqual(moveWebSearchKeyPoolEntry(entries, 0, -1), ["a", "b", "c"]); + assert.deepEqual( + reorderWebSearchKeyPoolEntries(entries, ["c", "a", "b"]).map((entry) => entry.id), + ["c", "a", "b"], + ); + assert.throws(() => reorderWebSearchKeyPoolEntries(entries, ["a", "b"]), /no longer matches/u); + assert.throws( + () => reorderWebSearchKeyPoolEntries(entries, ["a", "a", "b"]), + /no longer matches/u, + ); + + assert.equal(formatWebSearchKeyPoolCooldown(45_000, 0), "45s"); + assert.equal(formatWebSearchKeyPoolCooldown(12 * MINUTE, 0), "12m"); + assert.equal(formatWebSearchKeyPoolCooldown(185 * MINUTE, 0), "3h 5m"); + + assert.equal(normalizeWebSearchKeyPoolLabel(" Team\u0000 key ", "Key 2"), "Team key"); + assert.equal(normalizeWebSearchKeyPoolLabel(" ", "Key 2"), "Key 2"); + assert.equal(normalizeWebSearchKeyPoolLabel("x".repeat(80), "Key 2").length, 48); +}); diff --git a/main/services/web-search-key-pool-core.ts b/main/services/web-search-key-pool-core.ts new file mode 100644 index 000000000..785afb7e9 --- /dev/null +++ b/main/services/web-search-key-pool-core.ts @@ -0,0 +1,345 @@ +/** + * Main-only Web Search API-key pool: rotation, failover, and cooldowns. + * + * The pool is pure: no Electron, no secret store, and no clock unless one is + * injected. Keys pass through `runWithWebSearchKeyPool` only as opaque values + * handed to the caller's request function. They are never stored in tracker + * state, errors, or renderer projections. + */ + +import { WebSearchError, webSearchError } from "./web-search-core.js"; +import type { WebSearchProviderId } from "./web-search-provider-registry-core.js"; +import { + MAX_WEB_SEARCH_KEY_POOL_ENTRIES, + WEB_SEARCH_KEY_POOL_PRIMARY_ENTRY_ID, + isWebSearchKeyPoolEntryId, + isWebSearchKeyPoolProviderId, + isWebSearchKeyPoolStrategy, + normalizeWebSearchKeyPoolLabel, + type WebSearchKeyPoolCooldownReason, + type WebSearchKeyPoolEntryMetadata, + type WebSearchKeyPoolProviderId, + type WebSearchKeyPoolRendererState, + type WebSearchKeyPoolStrategy, +} from "../../renderer/shared/web-search-key-pool.js"; + +export * from "../../renderer/shared/web-search-key-pool.js"; + +// ── Durable index document ───────────────────────────────────────────── + +/** Non-secret pool index. It is stored encrypted beside the keys it orders. */ +export interface WebSearchKeyPoolDocument { + readonly version: 1; + readonly strategy: WebSearchKeyPoolStrategy; + readonly entries: readonly WebSearchKeyPoolEntryMetadata[]; +} + +export const EMPTY_WEB_SEARCH_KEY_POOL_DOCUMENT: WebSearchKeyPoolDocument = Object.freeze({ + version: 1, + strategy: "ordered", + entries: Object.freeze([]), +}); + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +/** + * Parse a stored index strictly. Anything malformed, duplicated, oversized, + * or future-versioned is rejected as a whole so an unreadable index never + * silently reorders or drops keys; callers treat `undefined` as "no index". + */ +export function parseWebSearchKeyPoolDocument( + raw: string | null | undefined, +): WebSearchKeyPoolDocument | undefined { + if (typeof raw !== "string" || raw.length === 0 || raw.length > 64 * 1_024) return undefined; + let value: unknown; + try { + value = JSON.parse(raw); + } catch { + return undefined; + } + if (!isRecord(value) || value.version !== 1 || !isWebSearchKeyPoolStrategy(value.strategy)) { + return undefined; + } + if (!Array.isArray(value.entries) || value.entries.length > MAX_WEB_SEARCH_KEY_POOL_ENTRIES) { + return undefined; + } + const seen = new Set(); + const entries: WebSearchKeyPoolEntryMetadata[] = []; + for (const [index, candidate] of value.entries.entries()) { + if (!isRecord(candidate) || !isWebSearchKeyPoolEntryId(candidate.id)) return undefined; + if (seen.has(candidate.id)) return undefined; + const addedAt = candidate.addedAt; + if (typeof addedAt !== "number" || !Number.isSafeInteger(addedAt) || addedAt < 0) { + return undefined; + } + seen.add(candidate.id); + entries.push({ + id: candidate.id, + label: normalizeWebSearchKeyPoolLabel(candidate.label, `Key ${index + 1}`), + addedAt, + }); + } + return { version: 1, strategy: value.strategy, entries }; +} + +export function serializeWebSearchKeyPoolDocument(document: WebSearchKeyPoolDocument): string { + return JSON.stringify({ + version: 1, + strategy: document.strategy, + entries: document.entries.map(({ id, label, addedAt }) => ({ id, label, addedAt })), + }); +} + +/** Secret-store slot for one pool entry. The primary entry keeps the pre-pool slot. */ +export function webSearchKeyPoolEntrySecretId(baseSecretId: string, entryId: string): string { + if (!isWebSearchKeyPoolEntryId(entryId)) throw new Error("Invalid Web Search key entry."); + return entryId === WEB_SEARCH_KEY_POOL_PRIMARY_ENTRY_ID + ? baseSecretId + : `${baseSecretId}:pool:${entryId}`; +} + +export function webSearchKeyPoolIndexSecretId(baseSecretId: string): string { + return `${baseSecretId}:pool-index`; +} + +export function webSearchKeyPoolSupported( + providerId: unknown, +): providerId is WebSearchKeyPoolProviderId & WebSearchProviderId { + return isWebSearchKeyPoolProviderId(providerId); +} + +// ── Cooldown tracker ─────────────────────────────────────────────────── + +export interface WebSearchKeyPoolCooldownPolicy { + /** First quota/rate-limit cooldown; doubles per consecutive failure. */ + readonly quotaBaseMs: number; + readonly quotaMaxMs: number; + /** First rejected-credential cooldown; doubles per consecutive failure. */ + readonly authBaseMs: number; + readonly authMaxMs: number; +} + +export const DEFAULT_WEB_SEARCH_KEY_POOL_COOLDOWN_POLICY: WebSearchKeyPoolCooldownPolicy = + Object.freeze({ + quotaBaseMs: 60_000, + quotaMaxMs: 60 * 60_000, + authBaseMs: 15 * 60_000, + authMaxMs: 24 * 60 * 60_000, + }); + +export interface WebSearchKeyPoolCooldown { + readonly reason: WebSearchKeyPoolCooldownReason; + readonly until: number; + readonly consecutiveFailures: number; +} + +interface CooldownRecord { + reason: WebSearchKeyPoolCooldownReason; + until: number; + consecutiveFailures: number; +} + +/** + * Process-local cooldown and round-robin state. It is deliberately not + * persisted: a restart gives every key a fresh chance, and nothing about + * provider failures is written to disk. + */ +export class WebSearchKeyPoolTracker { + private readonly records = new Map(); + private readonly cursors = new Map(); + private readonly now: () => number; + private readonly policy: WebSearchKeyPoolCooldownPolicy; + + constructor( + options: { + readonly now?: () => number; + readonly policy?: Partial; + } = {}, + ) { + this.now = options.now ?? Date.now; + this.policy = { ...DEFAULT_WEB_SEARCH_KEY_POOL_COOLDOWN_POLICY, ...options.policy }; + } + + currentTime(): number { + return this.now(); + } + + /** + * The attempt order for one search. Ordered pools always start at the + * first key; round-robin pools start one position later on every call so + * successive searches spread across healthy keys. + */ + order( + providerId: string, + entryIds: readonly string[], + strategy: WebSearchKeyPoolStrategy, + ): string[] { + if (strategy === "ordered" || entryIds.length < 2) return [...entryIds]; + const cursor = this.cursors.get(providerId) ?? 0; + this.cursors.set(providerId, (cursor + 1) % entryIds.length); + const start = cursor % entryIds.length; + return [...entryIds.slice(start), ...entryIds.slice(0, start)]; + } + + /** Active cooldown, or undefined when the key is eligible. */ + cooldown(providerId: string, entryId: string): WebSearchKeyPoolCooldown | undefined { + const record = this.records.get(recordKey(providerId, entryId)); + if (!record || record.until <= this.now()) return undefined; + return { ...record }; + } + + recordFailure( + providerId: string, + entryId: string, + reason: WebSearchKeyPoolCooldownReason, + ): WebSearchKeyPoolCooldown { + const key = recordKey(providerId, entryId); + const previous = this.records.get(key); + // Escalate only a repeated failure of the same kind. A quota failure after + // an auth cooldown (or vice versa) starts that kind's backoff from base. + const consecutiveFailures = + previous && previous.reason === reason ? previous.consecutiveFailures + 1 : 1; + const base = reason === "auth" ? this.policy.authBaseMs : this.policy.quotaBaseMs; + const max = reason === "auth" ? this.policy.authMaxMs : this.policy.quotaMaxMs; + const duration = Math.min(max, base * 2 ** Math.min(consecutiveFailures - 1, 20)); + const record = { reason, until: this.now() + duration, consecutiveFailures }; + this.records.set(key, record); + return { ...record }; + } + + recordSuccess(providerId: string, entryId: string): void { + this.records.delete(recordKey(providerId, entryId)); + } + + /** Manually make one key (or a provider's whole pool) eligible again. */ + clear(providerId: string, entryId?: string): void { + if (entryId !== undefined) { + this.records.delete(recordKey(providerId, entryId)); + return; + } + const prefix = `${providerId}\u0000`; + for (const key of [...this.records.keys()]) { + if (key.startsWith(prefix)) this.records.delete(key); + } + this.cursors.delete(providerId); + } + + /** Drop state for entries that no longer exist, so reused IDs start clean. */ + retain(providerId: string, entryIds: readonly string[]): void { + const keep = new Set(entryIds.map((entryId) => recordKey(providerId, entryId))); + const prefix = `${providerId}\u0000`; + for (const key of [...this.records.keys()]) { + if (key.startsWith(prefix) && !keep.has(key)) this.records.delete(key); + } + } +} + +function recordKey(providerId: string, entryId: string): string { + return `${providerId}\u0000${entryId}`; +} + +/** Redacted renderer projection of one provider's pool. */ +export function webSearchKeyPoolRendererState( + providerId: WebSearchKeyPoolProviderId, + document: WebSearchKeyPoolDocument, + tracker: WebSearchKeyPoolTracker, +): WebSearchKeyPoolRendererState { + return { + providerId, + strategy: document.strategy, + maxEntries: MAX_WEB_SEARCH_KEY_POOL_ENTRIES, + observedAt: tracker.currentTime(), + entries: document.entries.map((entry) => { + const cooldown = tracker.cooldown(providerId, entry.id); + return { + id: entry.id, + label: entry.label, + addedAt: entry.addedAt, + cooldown: cooldown ? { reason: cooldown.reason, until: cooldown.until } : null, + }; + }), + }; +} + +// ── Failover runner ──────────────────────────────────────────────────── + +export interface WebSearchPooledKey { + readonly id: string; + readonly key: string; +} + +export interface WebSearchKeyPool { + readonly strategy: WebSearchKeyPoolStrategy; + readonly keys: readonly WebSearchPooledKey[]; +} + +function cooldownReasonFor(error: unknown): WebSearchKeyPoolCooldownReason | undefined { + if (!(error instanceof WebSearchError)) return undefined; + if (error.kind === "auth") return "auth"; + if (error.kind === "quota") return "quota"; + return undefined; +} + +/** + * Run one provider request against a key pool. + * + * Keys are tried in pool order (rotated for round-robin), skipping keys that + * are cooling down. A rejected credential (401/403) or quota/rate-limit + * response (429 and provider quota codes) puts that key on cooldown and + * moves to the next key. Any other failure, including cancellation and + * timeout, stops immediately: another key would not fix it. + * + * When every key is cooling down, no request is sent. The pool fails with + * `quota` if any key is only rate-limited (so automatic routing can fall + * back and a later search can retry), otherwise with `auth`. + */ +export async function runWithWebSearchKeyPool(options: { + readonly providerId: WebSearchProviderId; + readonly pool: WebSearchKeyPool; + readonly tracker: WebSearchKeyPoolTracker; + readonly signal?: AbortSignal; + /** Called immediately before each keyed request is sent. */ + readonly beforeKeyAttempt?: (entryId: string) => void | Promise; + readonly run: (key: string) => Promise; +}): Promise { + const { providerId, pool, tracker } = options; + if (pool.keys.length === 0) throw webSearchError("auth", providerId); + const byId = new Map(pool.keys.map((entry) => [entry.id, entry])); + const order = tracker.order( + providerId, + pool.keys.map((entry) => entry.id), + pool.strategy, + ); + + let lastError: unknown; + let sawQuota = false; + for (const entryId of order) { + const entry = byId.get(entryId); + if (!entry) continue; + const cooling = tracker.cooldown(providerId, entryId); + if (cooling) { + if (cooling.reason === "quota") sawQuota = true; + continue; + } + if (options.signal?.aborted) throw webSearchError("cancelled", providerId); + if (options.beforeKeyAttempt) await options.beforeKeyAttempt(entryId); + try { + const result = await options.run(entry.key); + tracker.recordSuccess(providerId, entryId); + return result; + } catch (error) { + const reason = cooldownReasonFor(error); + if (!reason) throw error; + tracker.recordFailure(providerId, entryId, reason); + if (reason === "quota") sawQuota = true; + lastError = error; + if (options.signal?.aborted) throw webSearchError("cancelled", providerId); + } + } + if (lastError instanceof WebSearchError && (lastError.kind === "quota" || !sawQuota)) { + throw lastError; + } + throw webSearchError(sawQuota ? "quota" : "auth", providerId); +} diff --git a/main/services/web-search-main.ts b/main/services/web-search-main.ts index 50b3a768d..22448824e 100644 --- a/main/services/web-search-main.ts +++ b/main/services/web-search-main.ts @@ -2,7 +2,7 @@ import { configStore } from "./config-store.js"; import { configStores } from "./portable-config.js"; -import { webSearchCredentials } from "./web-search-credentials.js"; +import { webSearchCredentials, webSearchKeyPoolTracker } from "./web-search-credentials.js"; import { webSearchExistingAuthReuse } from "./web-search-auth-reuse-main.js"; import { webSearchRollout } from "./web-search-rollout.js"; import { WebSearchService } from "./web-search.js"; @@ -41,6 +41,20 @@ export const webSearchService = new WebSearchService({ return null; } }, + getCredentialPool: async (providerId) => { + try { + const settings = await configStore.getWebSearchSettings(); + return await webSearchCredentials.readPool( + webSearchCredentials.reference( + providerId, + settings.providerConfig[providerId as keyof typeof settings.providerConfig], + ), + ); + } catch { + return null; + } + }, + keyPoolTracker: webSearchKeyPoolTracker, getExistingAuthStatus: () => webSearchExistingAuthReuse.status(), resolveExistingAuth: () => webSearchExistingAuthReuse.resolve(), persistSettings: (patch) => configStore.setSettings(patch), diff --git a/main/services/web-search.test.ts b/main/services/web-search.test.ts index 069fb21f9..c22d60831 100644 --- a/main/services/web-search.test.ts +++ b/main/services/web-search.test.ts @@ -15,6 +15,8 @@ import { } from "./web-search-auth-reuse-core.js"; import type { WebSearchResolvedExistingAuth } from "./web-search-auth-reuse.js"; import type { AppSettings } from "./types.js"; +import { WebSearchKeyPoolTracker } from "./web-search-key-pool-core.js"; +import { createTavilyWebSearchAdapter } from "./web-search-tavily-adapter.js"; const PRIVATE_KEY = "exa-key-private-7dbfe9"; const QUERY = "current Exa documentation"; @@ -621,3 +623,68 @@ test("revocation during deferred existing-auth I/O fences evidence before public ); assert.equal(adapterCalls, 1); }); + +test("a pooled Tavily route fails over between keys, charges each keyed request, and then falls back", async () => { + const sentKeys: string[] = []; + const statusByKey: Record = { "tvly-first": 429, "tvly-second": 200 }; + const tavilyBody = JSON.stringify({ + results: [{ title: "Pooled", url: "https://example.test/pooled", content: "Evidence" }], + }); + const tracker = new WebSearchKeyPoolTracker({ now: () => 1_000 }); + const settings = baseSettings({ + mode: "automatic", + route: [ + { providerId: "tavily", credentialMode: "api-key" }, + { providerId: "parallel-mcp", credentialMode: "anonymous" }, + ], + fallbackOn: ["quota"], + }); + const service = new WebSearchService({ + getSettings: async () => settings, + getCredential: async () => "tvly-first", + getCredentialPool: async (providerId) => + providerId === "tavily" + ? { + strategy: "ordered", + keys: [ + { id: "primary", key: "tvly-first" }, + { id: "second", key: "tvly-second" }, + ], + } + : null, + keyPoolTracker: tracker, + fetch: async (_input, init) => { + const key = new Headers(init?.headers).get("authorization")?.replace(/^Bearer /u, "") ?? ""; + sentKeys.push(key); + const status = statusByKey[key] ?? 500; + return response(status === 200 ? tavilyBody : "{}", status); + }, + adapterFactories: { + tavily: createTavilyWebSearchAdapter, + "parallel-mcp": adapter("parallel-mcp", async () => evidence("parallel-mcp")), + }, + }); + + const charged: string[] = []; + const charge = { beforeProviderAttempt: (providerId: string) => void charged.push(providerId) }; + const first = await service.search({ query: QUERY }, charge); + assert.equal(first.providerId, "tavily"); + assert.equal(first.results[0]?.url, "https://example.test/pooled"); + assert.deepEqual(sentKeys, ["tvly-first", "tvly-second"]); + assert.deepEqual(charged, ["tavily", "tavily"]); + assert.equal(tracker.cooldown("tavily", "primary")?.reason, "quota"); + + // Once every pooled key is limited, the route falls back to the next provider. + statusByKey["tvly-second"] = 432; + sentKeys.length = 0; + charged.length = 0; + const second = await service.search({ query: QUERY }, charge); + assert.deepEqual(sentKeys, ["tvly-second"]); + assert.equal(second.providerId, "parallel-mcp"); + assert.deepEqual(charged, ["tavily", "parallel-mcp"]); + + sentKeys.length = 0; + const third = await service.search({ query: QUERY }); + assert.equal(third.providerId, "parallel-mcp"); + assert.deepEqual(sentKeys, [], "a fully cooling pool sends no Tavily request"); +}); diff --git a/main/services/web-search.ts b/main/services/web-search.ts index 4a7e4d2d5..37d8007e1 100644 --- a/main/services/web-search.ts +++ b/main/services/web-search.ts @@ -46,6 +46,12 @@ import { type WebSearchRolloutPolicy, } from "./web-search-rollout.js"; import type { WebSearchResolvedExistingAuth } from "./web-search-auth-reuse.js"; +import { + runWithWebSearchKeyPool, + webSearchKeyPoolSupported, + WebSearchKeyPoolTracker, + type WebSearchKeyPool, +} from "./web-search-key-pool-core.js"; import type { WebSearchExistingAuthRendererStatus } from "./web-search-auth-reuse-core.js"; import { declarePiRuntimeReplay } from "./pi-runtime-tool.js"; import type { AppSettings } from "./types.js"; @@ -57,6 +63,14 @@ export interface WebSearchServiceDependencies { getSettings: () => Promise; /** Reads one main-owned credential; null means absent or unavailable. */ getCredential?: (providerId: string) => Promise; + /** + * Reads every pooled key for a pool-capable provider (main-only). When + * present, it replaces `getCredential` for those providers and enables + * per-key failover with cooldowns. + */ + getCredentialPool?: (providerId: string) => Promise; + /** Process-local cooldown/rotation state shared with Settings status. */ + keyPoolTracker?: WebSearchKeyPoolTracker; /** Reads only the redacted status of the explicit existing-auth binding. */ getExistingAuthStatus?: () => Promise; /** Re-verifies and resolves the explicit binding immediately before I/O. */ @@ -108,6 +122,8 @@ interface PreparedAttempt { readonly entry: WebSearchRouteEntry; readonly adapter?: WebSearchAdapter; readonly credential?: string; + /** Pooled keys; when present the attempt fails over across them. */ + readonly keyPool?: WebSearchKeyPool; readonly readiness: WebSearchProviderReadiness; readonly ready: boolean; } @@ -171,9 +187,11 @@ export class WebSearchService { Partial> >; private readonly rollout: WebSearchRolloutPolicy; + private readonly keyPoolTracker: WebSearchKeyPoolTracker; constructor(private readonly dependencies: WebSearchServiceDependencies) { this.adapterFactories = dependencies.adapterFactories ?? WEB_SEARCH_ADAPTER_FACTORIES; + this.keyPoolTracker = dependencies.keyPoolTracker ?? new WebSearchKeyPoolTracker(); // The default policy is captured when the main service is constructed. // It is intentionally not re-read per request or generation. this.rollout = pinWebSearchRolloutPolicy(dependencies.rollout ?? webSearchRollout); @@ -305,6 +323,20 @@ export class WebSearchService { } } + private async readCredentialPool(providerId: string): Promise { + if (!this.dependencies.getCredentialPool) return undefined; + try { + const pool = await this.dependencies.getCredentialPool(providerId); + if (!pool) return undefined; + const keys = pool.keys.filter( + (entry) => typeof entry.key === "string" && entry.key.trim().length > 0, + ); + return keys.length > 0 ? { strategy: pool.strategy, keys } : undefined; + } catch { + return undefined; + } + } + private async readExistingAuthStatus(): Promise { if (!this.dependencies.getExistingAuthStatus) return undefined; try { @@ -369,6 +401,7 @@ export class WebSearchService { const routeForCredential = settings.selection.mode === "fixed" ? [settings.selection] : settings.selection.route; const credentials = new Map(); + const keyPools = new Map(); if (legacyCredentialRead) { // The legacy Exa read above is a migration discriminator and remains // the credential for a legacy keyed Exa selection. Do not read it twice. @@ -378,13 +411,21 @@ export class WebSearchService { (entry) => entry.credentialMode === "api-key" && !credentials.has(entry.providerId), ); const routedCredentials = await Promise.all( - routedApiKeyEntries.map(async (entry) => ({ - providerId: entry.providerId, - credential: await this.readCredential(entry.providerId), - })), + routedApiKeyEntries.map(async (entry) => { + if (this.dependencies.getCredentialPool && webSearchKeyPoolSupported(entry.providerId)) { + const keyPool = await this.readCredentialPool(entry.providerId); + return { providerId: entry.providerId, credential: keyPool?.keys[0]?.key, keyPool }; + } + return { + providerId: entry.providerId, + credential: await this.readCredential(entry.providerId), + keyPool: undefined, + }; + }), ); - for (const { providerId, credential } of routedCredentials) { + for (const { providerId, credential, keyPool } of routedCredentials) { credentials.set(providerId, credential); + if (keyPool) keyPools.set(providerId, keyPool); } const route = snapshotWebSearchRoute(settings); @@ -420,7 +461,16 @@ export class WebSearchService { settings.enabled && adapter !== undefined && webSearchRouteEntryReady(entry, settings.providerConfig[entry.providerId], readiness); - return { entry, adapter, credential, readiness, ready }; + const keyPool = + entry.credentialMode === "api-key" ? keyPools.get(entry.providerId) : undefined; + return { + entry, + adapter, + credential, + ...(keyPool === undefined ? {} : { keyPool }), + readiness, + ready, + }; }), ); return { settings, route, attempts }; @@ -495,23 +545,42 @@ export class WebSearchService { // the hook is the final per-attempt fence immediately before adapter I/O. // A failed binding must therefore consume neither a child budget nor a // network attempt. - if (options.beforeProviderAttempt) { - await options.beforeProviderAttempt(attempt.entry.providerId); - } - if (callerSignal.aborted) throw webSearchError("cancelled", attempt.entry.providerId); - if (controller.signal.aborted) { - throw webSearchError(timedOut ? "timeout" : "cancelled", attempt.entry.providerId); - } - const adapterRequest: WebSearchAdapterRequest = { - query: request.query, - numResults: request.numResults, - credentialMode: attempt.entry.credentialMode, - ...(attempt.credential === undefined ? {} : { credential: attempt.credential }), - ...(existingAuth === undefined ? {} : { existingAuth }), - signal: controller.signal, - timedOut: () => timedOut, + const adapter = attempt.adapter; + const beforeRequest = async () => { + if (options.beforeProviderAttempt) { + await options.beforeProviderAttempt(attempt.entry.providerId); + } + if (callerSignal.aborted) throw webSearchError("cancelled", attempt.entry.providerId); + if (controller.signal.aborted) { + throw webSearchError(timedOut ? "timeout" : "cancelled", attempt.entry.providerId); + } }; - const result = await attempt.adapter.search(adapterRequest); + const send = (credential: string | undefined) => + adapter.search({ + query: request.query, + numResults: request.numResults, + credentialMode: attempt.entry.credentialMode, + ...(credential === undefined ? {} : { credential }), + ...(existingAuth === undefined ? {} : { existingAuth }), + signal: controller.signal, + timedOut: () => timedOut, + } satisfies WebSearchAdapterRequest); + let result: WebSearchResultSet; + if (attempt.keyPool) { + // Each keyed request is its own provider request, so the per-request + // fence (budget, authority, cancellation) runs before every key. + result = await runWithWebSearchKeyPool({ + providerId: attempt.entry.providerId, + pool: attempt.keyPool, + tracker: this.keyPoolTracker, + signal: controller.signal, + beforeKeyAttempt: beforeRequest, + run: send, + }); + } else { + await beforeRequest(); + result = await send(attempt.credential); + } if (callerSignal.aborted) throw webSearchError("cancelled", attempt.entry.providerId); if ( existingAuth !== undefined && diff --git a/package.json b/package.json index c506e83a3..7e63a438c 100644 --- a/package.json +++ b/package.json @@ -115,7 +115,7 @@ "test:terminal:coverage": "tsx --test --experimental-test-coverage --test-coverage-include=main/services/terminal-spawn-helper.ts --test-coverage-lines=100 --test-coverage-branches=100 --test-coverage-functions=100 main/services/terminal.test.ts && tsx --test --experimental-test-coverage --test-coverage-include=main/services/terminal.ts --test-coverage-lines=95 --test-coverage-branches=80 --test-coverage-functions=90 main/services/terminal.test.ts", "test:ghostty-terminal": "tsx --test renderer/lib/ghostty-terminal/core.test.ts renderer/lib/ghostty-terminal/keyCodes.test.ts renderer/lib/ghostty-terminal/runtime.test.ts renderer/lib/ghostty-terminal/surface.test.ts renderer/components/terminal-drawer.test.tsx", "test:provider-failure": "tsx --test main/services/provider-failure.test.ts", - "test:web-search": "tsx --test main/services/web-search.test.ts main/services/web-search-core.test.ts main/services/web-search-exa-core.test.ts main/services/web-search-provider-registry-core.test.ts main/services/web-search-provider-registry.test.ts main/services/web-search-credential-core.test.ts main/services/web-search-auth-reuse.test.ts main/services/web-search-rollout.test.ts main/services/web-search-wave1-adapters.test.ts main/services/web-search-wave1-ai.test.ts main/services/web-search-wave2-batch-a.test.ts main/services/web-search-wave2-batch-b.test.ts main/services/web-search-wave4-batch-a.test.ts main/services/web-search-wave4-batch-b.test.ts main/handlers/web-search-contract.test.ts renderer/components/settings/web-search-settings.test.tsx renderer/lib/settings-section.test.ts", + "test:web-search": "tsx --test main/services/web-search.test.ts main/services/web-search-core.test.ts main/services/web-search-exa-core.test.ts main/services/web-search-provider-registry-core.test.ts main/services/web-search-provider-registry.test.ts main/services/web-search-credential-core.test.ts main/services/web-search-auth-reuse.test.ts main/services/web-search-rollout.test.ts main/services/web-search-wave1-adapters.test.ts main/services/web-search-wave1-ai.test.ts main/services/web-search-wave2-batch-a.test.ts main/services/web-search-wave2-batch-b.test.ts main/services/web-search-wave4-batch-a.test.ts main/services/web-search-wave4-batch-b.test.ts main/handlers/web-search-contract.test.ts renderer/components/settings/web-search-settings.test.tsx renderer/lib/settings-section.test.ts main/services/web-search-key-pool-core.test.ts main/handlers/web-search-key-pool.test.ts renderer/components/settings/web-search-key-pool.test.tsx", "test:concentrate": "tsx --test main/services/concentrate-provider.test.ts", "test:compaction": "npm run test:vcc && tsx --test main/services/memory-shared.test.ts main/services/pi-runtime-events.test.ts main/services/pi-agent-runtime-harness.test.ts main/services/pi-compaction-core.test.ts main/services/pi-legacy-session.test.ts main/services/pi-session-migration.test.ts main/services/pi-session-v4-upgrade.test.ts main/services/pi-session-port.test.ts main/services/context-lifecycle-service.test.ts main/services/context-lifecycle-surface-matrix.test.ts main/services/telegram/telegram-session.test.ts main/services/generation-emergency-outcome.test.ts main/services/pi-runtime-effect-store.test.ts main/services/generation-context.test.ts main/services/memory-store.test.ts main/services/memory-context.test.ts main/services/memory-surface-matrix.test.ts main/services/pi-upgrade-evaluation.test.ts main/services/subagents/agent-compatibility.test.ts renderer/lib/ipc-stream.test.ts renderer/main/chat-transition.test.tsx", "test:compaction:packaged": "node --import tsx scripts/pi-session-packaged-restart-acceptance.mjs", diff --git a/renderer/components/settings/web-search-key-pool.test.tsx b/renderer/components/settings/web-search-key-pool.test.tsx new file mode 100644 index 000000000..cab132f7b --- /dev/null +++ b/renderer/components/settings/web-search-key-pool.test.tsx @@ -0,0 +1,98 @@ +import assert from "node:assert/strict"; +import test from "node:test"; +import { renderToStaticMarkup } from "react-dom/server"; + +import { WebSearchKeyPoolList, webSearchKeyPoolEntryStatus } from "./web-search-key-pool"; +import type { + WebSearchKeyPoolRendererEntry, + WebSearchKeyPoolRendererState, +} from "../../shared/web-search-key-pool"; + +const NOW = 1_000_000; +const noop = () => undefined; + +function entry( + id: string, + label: string, + cooldown: WebSearchKeyPoolRendererEntry["cooldown"] = null, +): WebSearchKeyPoolRendererEntry { + return { id, label, addedAt: 1, cooldown }; +} + +function state(entries: WebSearchKeyPoolRendererEntry[]): WebSearchKeyPoolRendererState { + return { providerId: "tavily", strategy: "ordered", maxEntries: 8, observedAt: NOW, entries }; +} + +function render(entries: WebSearchKeyPoolRendererEntry[]) { + const html = renderToStaticMarkup( + , + ); + return html; +} + +function entryHtml(html: string, id: string): string { + const start = html.indexOf(`data-entry-id="${id}"`); + assert.ok(start >= 0, `entry ${id} is rendered`); + const end = html.indexOf("", start); + return html.slice(start, end); +} + +test("pooled keys render in use order with their cooldown state and a retry action only while cooling", () => { + const html = render([ + entry("primary", "Personal"), + entry("team1", "Team", { reason: "quota", until: NOW + 45_000 }), + entry("team2", "Old", { reason: "auth", until: NOW + 15 * 60_000 }), + ]); + + assert.ok(html.indexOf("Personal") < html.indexOf("Team")); + assert.ok(html.indexOf("Team") < html.indexOf("Old")); + assert.match(html, /aria-label="Tavily API keys, in use order"/u); + + const personal = entryHtml(html, "primary"); + assert.match(personal, /data-cooling="false"/u); + assert.match(personal, /Active/u); + assert.doesNotMatch(personal, /Retry Personal now/u); + + const team = entryHtml(html, "team1"); + assert.match(team, /data-cooling="true"/u); + assert.match(team, /Rate limited · retry in 45s/u); + assert.match(team, /aria-label="Retry Team now"/u); + + const old = entryHtml(html, "team2"); + assert.match(old, /Rejected · retry in 15m/u); + assert.match(old, /aria-label="Retry Old now"/u); +}); + +function buttonDisabled(html: string, label: string): boolean { + const at = html.indexOf(`aria-label="${label}"`); + assert.ok(at >= 0, `button ${label} is rendered`); + const tag = html.slice(html.lastIndexOf("", at) + 1); + return /\sdisabled=""/u.test(tag); +} + +test("move controls are disabled at the ends of the list", () => { + const html = render([entry("a", "First"), entry("b", "Last")]); + assert.equal(buttonDisabled(html, "Move First up"), true); + assert.equal(buttonDisabled(html, "Move First down"), false); + assert.equal(buttonDisabled(html, "Move Last up"), false); + assert.equal(buttonDisabled(html, "Move Last down"), true); + assert.equal(buttonDisabled(html, "Remove First"), false); +}); + +test("an empty pool explains how to add a key", () => { + const html = render([]); + assert.match(html, /No Tavily keys saved yet/u); + assert.doesNotMatch(html, /
    { + const expired = entry("a", "Key", { reason: "quota", until: NOW - 1 }); + assert.deepEqual(webSearchKeyPoolEntryStatus(expired, NOW), { label: "Active", color: "green" }); +}); diff --git a/renderer/components/settings/web-search-key-pool.tsx b/renderer/components/settings/web-search-key-pool.tsx new file mode 100644 index 000000000..1698f4392 --- /dev/null +++ b/renderer/components/settings/web-search-key-pool.tsx @@ -0,0 +1,414 @@ +// Web Search API-key pool editor. +// +// Keys are write-only: the renderer sends a new key once and afterwards sees +// only an opaque entry ID, the user's label, its position, and redacted +// cooldown state. Rotation, failover, and cooldowns run in the main process. + +import * as React from "react"; +import { ArrowDown, ArrowUp, KeyRound, Plus, RotateCcw, Trash2 } from "lucide-react"; +import { Badge, Button, Field, Input, RadioGroup, RadioGroupItem, Text } from "../ui"; +import { webSearchApi } from "../../lib/ipc"; +import { + formatWebSearchKeyPoolCooldown, + moveWebSearchKeyPoolEntry, + type WebSearchKeyPoolProviderId, + type WebSearchKeyPoolRendererEntry, + type WebSearchKeyPoolRendererState, + type WebSearchKeyPoolStrategy, +} from "../../shared/web-search-key-pool"; + +const STRATEGY_OPTIONS: ReadonlyArray<{ + value: WebSearchKeyPoolStrategy; + title: string; + description: string; +}> = [ + { + value: "ordered", + title: "In order", + description: + "Always start with the first key. Later keys are used only when earlier ones fail.", + }, + { + value: "round-robin", + title: "Round robin", + description: "Start each search on the next key to spread usage evenly across the pool.", + }, +]; + +/** Status label for one pooled key; exported for rendering tests. */ +export function webSearchKeyPoolEntryStatus( + entry: WebSearchKeyPoolRendererEntry, + now: number, +): { label: string; color: "green" | "warning" | "red" } { + if (!entry.cooldown || entry.cooldown.until <= now) return { label: "Active", color: "green" }; + const remaining = formatWebSearchKeyPoolCooldown(entry.cooldown.until, now); + return entry.cooldown.reason === "auth" + ? { label: `Rejected · retry in ${remaining}`, color: "red" } + : { label: `Rate limited · retry in ${remaining}`, color: "warning" }; +} + +export interface WebSearchKeyPoolListProps { + readonly providerLabel: string; + readonly state: WebSearchKeyPoolRendererState; + readonly now: number; + readonly disabled?: boolean; + readonly onMove: (index: number, direction: -1 | 1) => void; + readonly onRemove: (entry: WebSearchKeyPoolRendererEntry) => void; + readonly onResetCooldown: (entry: WebSearchKeyPoolRendererEntry) => void; +} + +/** Presentational, ordered list of pooled keys. */ +export function WebSearchKeyPoolList({ + providerLabel, + state, + now, + disabled, + onMove, + onRemove, + onResetCooldown, +}: WebSearchKeyPoolListProps) { + if (state.entries.length === 0) { + return ( + + No {providerLabel} keys saved yet. Add one below. + + ); + } + const total = state.entries.length; + return ( +
      + {state.entries.map((entry, index) => { + const status = webSearchKeyPoolEntryStatus(entry, now); + const cooling = status.color !== "green"; + return ( +
    1. { + if (disabled) return; + if (!event.altKey || (event.key !== "ArrowUp" && event.key !== "ArrowDown")) return; + event.preventDefault(); + onMove(index, event.key === "ArrowUp" ? -1 : 1); + }} + className="group flex min-w-0 items-center gap-2 rounded-control bg-well px-2.5 py-2 outline-none transition-colors duration-150 ease-out hover:bg-list-hover focus-visible:bg-input focus-visible:shadow-control motion-reduce:transition-none" + > + +
      2. + ); + })} +
    + ); +} + +function errorMessage(error: unknown, fallback: string): string { + if (!(error instanceof Error) || !error.message) return fallback; + // Electron prefixes invoke failures with the channel; keep only the reason. + return error.message.replace(/^Error invoking remote method '[^']+': (?:Error: )?/u, ""); +} + +function hasActiveCooldown(state: WebSearchKeyPoolRendererState | null, now: number): boolean { + return state?.entries.some((entry) => entry.cooldown && entry.cooldown.until > now) === true; +} + +/** Stateful pool editor used by the provider setup dialog. */ +export function WebSearchKeyPoolEditor({ + providerId, + providerLabel, + onChanged, +}: { + providerId: WebSearchKeyPoolProviderId; + providerLabel: string; + /** Called after a successful change so the provider snapshot can refresh. */ + onChanged?: () => void; +}) { + const [state, setState] = React.useState(null); + const [now, setNow] = React.useState(() => Date.now()); + const [keyDraft, setKeyDraft] = React.useState(""); + const [labelDraft, setLabelDraft] = React.useState(""); + const [busy, setBusy] = React.useState(false); + const [message, setMessage] = React.useState(null); + const [error, setError] = React.useState(null); + + React.useEffect(() => { + let active = true; + webSearchApi.keyPool + .get(providerId) + .then((next) => { + if (!active) return; + setState(next); + setNow(Date.now()); + }) + .catch((caught: unknown) => { + if (active) setError(errorMessage(caught, "Couldn’t load saved keys.")); + }); + return () => { + active = false; + }; + }, [providerId]); + + // Keep the remaining-cooldown labels current only while one is visible. + const ticking = hasActiveCooldown(state, now); + React.useEffect(() => { + if (!ticking) return; + const timer = window.setInterval(() => setNow(Date.now()), 1_000); + return () => window.clearInterval(timer); + }, [ticking]); + + const run = async ( + operation: () => Promise, + success: string, + failure: string, + ): Promise => { + if (busy) return false; + setBusy(true); + setMessage(null); + setError(null); + try { + const next = await operation(); + setState(next); + setNow(Date.now()); + setMessage(success); + onChanged?.(); + return true; + } catch (caught) { + setError(errorMessage(caught, failure)); + return false; + } finally { + setBusy(false); + } + }; + + const addKey = async () => { + const value = keyDraft.trim(); + if (!value) return; + const added = await run( + () => webSearchApi.keyPool.add(providerId, value, labelDraft.trim() || undefined), + "Key added to the pool. Saving does not select the provider or send a test request.", + "Couldn’t add the API key.", + ); + if (added) { + setKeyDraft(""); + setLabelDraft(""); + } + }; + + const full = state !== null && state.entries.length >= state.maxEntries; + + return ( +
    + + {state ? ( + { + const order = moveWebSearchKeyPoolEntry(state.entries, index, direction); + void run( + () => webSearchApi.keyPool.reorder(providerId, order), + "Key order saved.", + "Couldn’t reorder the keys.", + ); + }} + onRemove={(entry) => + void run( + () => webSearchApi.keyPool.remove(providerId, entry.id), + `${entry.label} removed.`, + "Couldn’t remove the API key.", + ) + } + onResetCooldown={(entry) => + void run( + () => webSearchApi.keyPool.resetCooldown(providerId, entry.id), + `${entry.label} will be tried on the next search.`, + "Couldn’t clear the cooldown.", + ) + } + /> + ) : error ? null : ( + + Loading saved keys… + + )} +
    { + event.preventDefault(); + void addKey(); + }} + > + setLabelDraft(event.target.value)} + placeholder="Label (optional)" + autoComplete="off" + spellCheck={false} + maxLength={48} + aria-label={`${providerLabel} key label`} + disabled={busy || full} + className="h-10 w-full" + /> + setKeyDraft(event.target.value)} + placeholder={full ? "The pool is full" : "Paste an API key"} + autoComplete="new-password" + spellCheck={false} + aria-label={`${providerLabel} API key to add`} + disabled={busy || full} + className="h-10 w-full" + /> + +
    + {state ? ( + + {state.entries.length} of {state.maxEntries} keys. Cooldowns reset when Aiden restarts. + + ) : null} +     + + {state && state.entries.length > 1 ? ( + + + void run( + () => + webSearchApi.keyPool.setStrategy(providerId, value as WebSearchKeyPoolStrategy), + "Rotation saved.", + "Couldn’t change the rotation.", + ) + } + aria-label={`${providerLabel} key rotation`} + className="grid gap-2" + disabled={busy} + > + {STRATEGY_OPTIONS.map((option) => ( + + ))} + + + ) : null} + + {message ? ( + + {message} + + ) : null} + {error ? ( + + {error} + + ) : null} +
    + ); +} diff --git a/renderer/components/settings/web-search-settings.tsx b/renderer/components/settings/web-search-settings.tsx index d319456cd..6f47f66bd 100644 --- a/renderer/components/settings/web-search-settings.tsx +++ b/renderer/components/settings/web-search-settings.tsx @@ -58,6 +58,8 @@ import type { WebSearchSelection, } from "../../lib/types"; import { useAppCapabilities } from "../../lib/app-capabilities"; +import { isWebSearchKeyPoolProviderId } from "../../shared/web-search-key-pool"; +import { WebSearchKeyPoolEditor } from "./web-search-key-pool"; type WebSearchProvider = WebSearchProviderRendererMetadata; type CredentialMode = WebSearchRouteEntry["credentialMode"]; @@ -901,7 +903,21 @@ function ProviderSetupDialog({ {providerNeedsApiKey(provider) && (!supportsExistingAuth || routeMode === "api-key") && - (provider.credentialKind !== "optional-api-key" || routeMode === "api-key") ? ( + (provider.credentialKind !== "optional-api-key" || routeMode === "api-key") && + isWebSearchKeyPoolProviderId(provider.id) ? ( + { + webSearchApi + .get() + .then(onSnapshot) + .catch(() => undefined); + }} + /> + ) : providerNeedsApiKey(provider) && + (!supportsExistingAuth || routeMode === "api-key") && + (provider.credentialKind !== "optional-api-key" || routeMode === "api-key") ? ( ("webSearch:setCredential", providerId, key), removeCredential: (providerId: WebSearchProviderId) => invoke("webSearch:removeCredential", providerId), + keyPool: { + get: (providerId: WebSearchKeyPoolProviderId) => + invoke("webSearch:keyPool:get", providerId), + add: (providerId: WebSearchKeyPoolProviderId, key: string, label?: string) => + invoke("webSearch:keyPool:add", providerId, key, label), + remove: (providerId: WebSearchKeyPoolProviderId, entryId: string) => + invoke("webSearch:keyPool:remove", providerId, entryId), + reorder: (providerId: WebSearchKeyPoolProviderId, entryIds: string[]) => + invoke("webSearch:keyPool:reorder", providerId, entryIds), + setStrategy: (providerId: WebSearchKeyPoolProviderId, strategy: WebSearchKeyPoolStrategy) => + invoke("webSearch:keyPool:setStrategy", providerId, strategy), + resetCooldown: (providerId: WebSearchKeyPoolProviderId, entryId: string) => + invoke("webSearch:keyPool:resetCooldown", providerId, entryId), + }, }; // ── Telegram remote control ────────────────────────────────────────── diff --git a/renderer/shared/web-search-key-pool.ts b/renderer/shared/web-search-key-pool.ts new file mode 100644 index 000000000..a199d7c27 --- /dev/null +++ b/renderer/shared/web-search-key-pool.ts @@ -0,0 +1,132 @@ +/** + * Renderer-safe Web Search API-key pool contract. + * + * Key material never appears in this module. The renderer sees only opaque + * entry IDs, user-chosen labels, order, rotation strategy, and redacted + * cooldown state. Plaintext keys stay in the main process. + */ + +/** Providers whose adapter can rotate through several keys. */ +export const WEB_SEARCH_KEY_POOL_PROVIDER_IDS = Object.freeze(["tavily"] as const); +export type WebSearchKeyPoolProviderId = (typeof WEB_SEARCH_KEY_POOL_PROVIDER_IDS)[number]; + +export const MAX_WEB_SEARCH_KEY_POOL_ENTRIES = 8; +export const MAX_WEB_SEARCH_KEY_POOL_LABEL_CHARS = 48; +/** The pre-pool single credential slot is always addressed by this entry ID. */ +export const WEB_SEARCH_KEY_POOL_PRIMARY_ENTRY_ID = "primary"; + +export type WebSearchKeyPoolStrategy = "ordered" | "round-robin"; +export const WEB_SEARCH_KEY_POOL_STRATEGIES: readonly WebSearchKeyPoolStrategy[] = Object.freeze([ + "ordered", + "round-robin", +]); + +/** Why a key is temporarily skipped. */ +export type WebSearchKeyPoolCooldownReason = "auth" | "quota"; + +export interface WebSearchKeyPoolEntryMetadata { + readonly id: string; + readonly label: string; + /** Epoch milliseconds; zero for a key saved before pools existed. */ + readonly addedAt: number; +} + +export interface WebSearchKeyPoolRendererEntry extends WebSearchKeyPoolEntryMetadata { + readonly cooldown: null | { + readonly reason: WebSearchKeyPoolCooldownReason; + /** Epoch milliseconds when the key becomes eligible again. */ + readonly until: number; + }; +} + +export interface WebSearchKeyPoolRendererState { + readonly providerId: WebSearchKeyPoolProviderId; + readonly strategy: WebSearchKeyPoolStrategy; + readonly maxEntries: number; + readonly entries: readonly WebSearchKeyPoolRendererEntry[]; + /** Epoch milliseconds when main produced this projection. */ + readonly observedAt: number; +} + +const ENTRY_ID_PATTERN = /^[A-Za-z0-9_-]{1,64}$/u; + +export function isWebSearchKeyPoolProviderId(value: unknown): value is WebSearchKeyPoolProviderId { + return ( + typeof value === "string" && + (WEB_SEARCH_KEY_POOL_PROVIDER_IDS as readonly string[]).includes(value) + ); +} + +export function isWebSearchKeyPoolStrategy(value: unknown): value is WebSearchKeyPoolStrategy { + return value === "ordered" || value === "round-robin"; +} + +export function isWebSearchKeyPoolEntryId(value: unknown): value is string { + return typeof value === "string" && ENTRY_ID_PATTERN.test(value); +} + +/** + * Bound a display label. Control characters are collapsed to spaces and an + * empty label becomes the supplied fallback, so a label can never carry a + * line break into logs or UI. + */ +export function normalizeWebSearchKeyPoolLabel(value: unknown, fallback: string): string { + if (typeof value !== "string") return fallback; + // eslint-disable-next-line no-control-regex + const cleaned = value.replace(/[\u0000-\u001f\u007f-\u009f]+/gu, " ").replace(/\s+/gu, " "); + const trimmed = Array.from(cleaned.trim()) + .slice(0, MAX_WEB_SEARCH_KEY_POOL_LABEL_CHARS) + .join("") + .trim(); + return trimmed || fallback; +} + +/** + * Apply a caller-provided order. The order must be an exact permutation of + * the current entry IDs so a stale renderer cannot drop or resurrect a key. + */ +export function reorderWebSearchKeyPoolEntries( + entries: readonly T[], + orderedIds: unknown, +): T[] { + if (!Array.isArray(orderedIds) || orderedIds.length !== entries.length) { + throw new Error("The key order no longer matches the saved keys."); + } + const byId = new Map(entries.map((entry) => [entry.id, entry])); + const seen = new Set(); + const next: T[] = []; + for (const id of orderedIds) { + if (typeof id !== "string" || seen.has(id)) { + throw new Error("The key order no longer matches the saved keys."); + } + const entry = byId.get(id); + if (!entry) throw new Error("The key order no longer matches the saved keys."); + seen.add(id); + next.push(entry); + } + return next; +} + +/** Move one entry up or down, returning the resulting ID order. */ +export function moveWebSearchKeyPoolEntry( + entries: readonly { readonly id: string }[], + index: number, + direction: -1 | 1, +): string[] { + const ids = entries.map((entry) => entry.id); + const target = index + direction; + if (index < 0 || index >= ids.length || target < 0 || target >= ids.length) return ids; + [ids[index], ids[target]] = [ids[target]!, ids[index]!]; + return ids; +} + +/** Short, human-readable remaining cooldown, e.g. "45s", "12m", "3h 5m". */ +export function formatWebSearchKeyPoolCooldown(until: number, now: number): string { + const remainingSeconds = Math.max(1, Math.ceil((until - now) / 1_000)); + if (remainingSeconds < 60) return `${remainingSeconds}s`; + const minutes = Math.ceil(remainingSeconds / 60); + if (minutes < 60) return `${minutes}m`; + const hours = Math.floor(minutes / 60); + const rest = minutes % 60; + return rest === 0 ? `${hours}h` : `${hours}h ${rest}m`; +} diff --git a/scripts/ci-test-registry.json b/scripts/ci-test-registry.json index ccc802f40..bcbfd411a 100644 --- a/scripts/ci-test-registry.json +++ b/scripts/ci-test-registry.json @@ -27,6 +27,7 @@ "main/handlers/phase2-parse.test.ts", "main/handlers/voice-codec.test.ts", "main/handlers/web-search-contract.test.ts", + "main/handlers/web-search-key-pool.test.ts", "main/handlers/worktree-create-params.test.ts", "main/services/advisor-attempt-store.test.ts", "main/services/advisor-context.test.ts", @@ -85,8 +86,8 @@ "main/services/git-executable.test.ts", "main/services/git.test.ts", "main/services/github-pull-request.test.ts", - "main/services/local-runtime-status.test.ts", "main/services/linux-dictation-portal.test.ts", + "main/services/local-runtime-status.test.ts", "main/services/main-window-state-core.test.ts", "main/services/managed-worktree-admission.test.ts", "main/services/managed-worktree-creation-core.test.ts", @@ -120,8 +121,8 @@ "main/services/pi-runtime-effect-store.test.ts", "main/services/pi-runtime-events.test.ts", "main/services/pi-session-migration.test.ts", - "main/services/pi-session-v4-upgrade.test.ts", "main/services/pi-session-port.test.ts", + "main/services/pi-session-v4-upgrade.test.ts", "main/services/pi-upgrade-evaluation.test.ts", "main/services/pi-vcc/vcc.test.ts", "main/services/portable-config-core.roundtrip.test.ts", @@ -173,6 +174,7 @@ "main/services/web-search-core.test.ts", "main/services/web-search-credential-core.test.ts", "main/services/web-search-exa-core.test.ts", + "main/services/web-search-key-pool-core.test.ts", "main/services/web-search-provider-registry-core.test.ts", "main/services/web-search-provider-registry.test.ts", "main/services/web-search-rollout.test.ts", @@ -482,6 +484,7 @@ "renderer/components/settings/simulator-settings.test.tsx", "renderer/components/settings/skills-settings.test.tsx", "renderer/components/settings/tts-settings.test.tsx", + "renderer/components/settings/web-search-key-pool.test.tsx", "renderer/components/settings/web-search-settings.test.tsx", "renderer/components/subagent-shell-approval.test.tsx", "renderer/components/subagents-panel.test.tsx",