diff --git a/doc/idle-tool-connection-health.md b/doc/idle-tool-connection-health.md index 6434343c244..795dc386020 100644 --- a/doc/idle-tool-connection-health.md +++ b/doc/idle-tool-connection-health.md @@ -19,8 +19,11 @@ no deadline until the settings producer enables them. Application status and company status are not new filters: this migration preserves the old policy. An unchecked row is due immediately. Otherwise its deadline is `healthCheckedAt` -plus 15 minutes. A company hint refreshes that company's cache. Startup and an -instance feature-settings change can refresh the whole instance. Refreshes are +plus 15 minutes. Individual writes carry a connection ID and reload only that +connection. Application mutations refresh only that application's connections; +company deletion invalidates the company's cache. Startup and instance feature +settings can refresh the whole instance. Inventory reads use UUID keyset pages +of at most 100 rows and observe shutdown cancellation between pages. Refreshes are coalesced with a 30-second minimum spacing when no due backlog remains. Batches retain at most 25 checks and continue after one second while already due rows remain, so the batch limit does not cap dispatch at 50 checks per minute. There is no scheduled refresh when @@ -32,8 +35,8 @@ rotates across passes when more companies are due than fit in one batch. Before a provider check, the service reloads current eligibility and deadline. A saved health result, including a failed health observation, owns the next 15-minute deadline. A check that cannot persist an observation retains recovery -debt with exponential backoff from 30 seconds to 15 minutes. A failed company -refresh does not prevent another company's checks. Dirty scopes arriving during +debt with exponential backoff from 30 seconds to 15 minutes. A failed scoped +refresh does not prevent checks outside that connection, application, or company. Dirty scopes arriving during a refresh remain pending. Drain retains them without passive database reads. Shutdown joins admitted work; it does not abandon a provider call midway. @@ -134,8 +137,8 @@ selection. The two-company regression failed before the repair; 19 health-worker and three shared batching cases pass afterward. Each probe's pre/post eligibility query already includes the connection ID, so those indexed current-state checks remain. A scheduled observation still emits the normal transaction signal and -may cause one coalesced company refresh; this deliberately preserves uncertainty -and concurrent-writer recovery rather than introducing a silent-write exception. +now refreshes only that connection following the security repair below. This +preserves uncertainty and concurrent-writer recovery without a silent-write exception. The real-service producer suite now proves AgentMail activation triggers its first health probe, pause cancels the next deadline, GitHub installation events update @@ -164,3 +167,44 @@ fully future-dated inventory has no continuation. This is scheduling concrete backlog, not a periodic idle scan. The initial deadline-only assertion passed incorrectly because a pending transaction notification temporarily reported an immediate coordinator wake; acceptance measures actual probe timestamps instead. + + +## Security review: bounded refresh scope + +Security review identified amplification in the original company hints. Five +separate real `updateConnection` calls in a company with 205 eligible connections +returned 1,025 inventory rows. A regression reproduced that result. The repaired +producer contract returns exactly five rows across those updates. Saved probe +observations use the same single-connection scope, so draining a due backlog no +longer reloads its whole company after each batch. + +The typed hint accepts an existing connection row (only company and connection +IDs are encoded), or an explicit application target. Company/global scopes remain +for company deletion and administrative feature changes. Smoke Lab allocates a +new connection ID before its insert so the transaction intent is registered before +the write. No schema, outer transaction ownership, or provider policy changes. +Application eligibility changes retain their own scoped refresh. Single-connection +refreshes remove an old cached deadline even when the row was deleted, disabled, +or moved to an excluded application; unrelated deadlines remain intact. + +Inventory queries use 100-row UUID keyset pages. A 205-row regression initially +observed one unbounded query and now requires three pages with all 205 distinct +IDs retained. Cancellation is checked between pages. Results replace a scope only +after its complete load succeeds; failed loads retain recovery and other scopes +can proceed. Malformed hint shapes request full reconstruction instead of dropping +later scopes. + +The first amplification fixture lacked its required remote endpoint and failed +validation before reaching its assertion. After adding a controlled endpoint +lookup, the baseline reproduced 1,025 versus five. The first typecheck found a +partial access-service projection without companyId; the target now supplies the +already-owned companyId explicitly. These failures are retained as evidence, not +counted as passes. Final verification and remote review are recorded below. + + +Final scoped-refresh verification: 26 health/batching cases, 63 email/access +cases, 56 GitHub/Smoke Lab/startup cases, 469 tool-access cases, and 92 gateway +cases pass (706 targeted tests on this repair). Server typecheck and build pass. +The full local test matrix was not rerun; prior remote Actions passed at +f8ad26f0ff, which still had the security finding. Fresh remote checks/security +review on the repaired head are required and are not inferred from those results. diff --git a/doc/plans/2026-10-24-recurring-server-sleep-stack.md b/doc/plans/2026-10-24-recurring-server-sleep-stack.md index 16e679ac57c..207052d0546 100644 --- a/doc/plans/2026-10-24-recurring-server-sleep-stack.md +++ b/doc/plans/2026-10-24-recurring-server-sleep-stack.md @@ -1169,3 +1169,14 @@ before the earlier cleanup. Seven producer integrations and all six cleanup case pass on the merged branch. The 1,413 deselected broad chat cases are not passes. Health `f8ad26f0ff` now has fresh Greptile 5/5 and no failed checks; Paperclip Review was still queued at the last inspection. The full goal remains active. + + +Health security repair: five real connection API updates in a 205-connection +company reproduced 1,025 inventory rows. Individual writers now signal their +connection IDs; application writers signal their application IDs. Genuine broad +loads use 100-row UUID pages. The regression now reads exactly five rows, with +application invalidation and unrelated-deadline preservation verified. All 706 +targeted health/producer/consumer cases and server typecheck/build pass. Publish +the repair to #17554, propagate through company-deletion #17561 into Discord +#17557, and require fresh remote gates. No PR in this stack is merged. The next +new owner remains chat runtime discovery; the full remaining ledger is unchanged. diff --git a/doc/plans/2026-10-24-tool-connection-health-sleep-audit.md b/doc/plans/2026-10-24-tool-connection-health-sleep-audit.md index 04abe345156..b06ba9070dc 100644 --- a/doc/plans/2026-10-24-tool-connection-health-sleep-audit.md +++ b/doc/plans/2026-10-24-tool-connection-health-sleep-audit.md @@ -163,3 +163,19 @@ the intended regression and passes. Producer/chat consumers, query-plan evidence build, and remote gates remain pending. Master was refreshed to `8415f4a679`; this uncommitted health branch remains based on the open migration stack, so final stack reconciliation with master is still required. No merge is authorized. + + +## Security follow-up: individual refreshes + +The company-hint design amplified ordinary connection writes into full company +reads. The current contract scopes individual producers (API/example creation, +setup/rollback, OAuth, gateway, health/catalog, email, voice, GitHub events, access +removal, and Smoke Lab) by company plus connection ID. Existing application +updates in example setup and Smoke Lab use application scopes. Company deletion +and instance feature-settings changes retain broad invalidation. Startup and +broad loads use bounded UUID pages. The health owner still defines eligibility +and provider policy; the shared notification/coordinator layers remain generic. + +See `doc/idle-tool-connection-health.md` for the failing amplification baseline, +changed contract, scope-cancellation tests, and verification results. The older +company-hint observations in this audit describe the superseded implementation. diff --git a/server/src/__tests__/helpers/tool-health-worker.ts b/server/src/__tests__/helpers/tool-health-worker.ts index 7247d5566ab..d51c18c221d 100644 --- a/server/src/__tests__/helpers/tool-health-worker.ts +++ b/server/src/__tests__/helpers/tool-health-worker.ts @@ -16,7 +16,7 @@ export async function startToolHealthWorker(db: Db, companyId: string, now = new }); try { await registerToolConnectionHealthWork(coordinator, { - load: scope => schedule.load(scope ?? companyId), check: schedule.check, + load: (scope, connectionId, applicationId, signal) => schedule.load(scope ?? companyId, connectionId, applicationId, signal), check: schedule.check, }, () => true).ready; } catch (error) { await coordinator.stop(); diff --git a/server/src/__tests__/tool-connection-health-work.test.ts b/server/src/__tests__/tool-connection-health-work.test.ts index 5f2d1caa4b2..247f3c36873 100644 --- a/server/src/__tests__/tool-connection-health-work.test.ts +++ b/server/src/__tests__/tool-connection-health-work.test.ts @@ -35,7 +35,7 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu applicationId = app.id; probe.mockReset(); probe.mockImplementation(async (id: string) => { const [connection] = await db.select().from(toolConnections).where(eq(toolConnections.id, id)); - await writeToolConnectionHealthChange(db, connection.companyId, async tx => { + await writeToolConnectionHealthChange(db, connection, async tx => { await tx.update(toolConnections).set({ healthCheckedAt: new Date(), healthStatus: "ok" }).where(eq(toolConnections.id, id)); }); return {}; @@ -51,10 +51,11 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu }); afterAll(async () => { await db?.$client.end({ timeout: 1 }); await temporary?.cleanup(); }); function connection(patch: Partial = {}, signal = true) { + const id = patch.id ?? randomUUID(); const insert = async (writer: typeof db) => (await writer.insert(toolConnections).values({ - companyId, applicationId, name: "Check", uid: randomUUID(), transport: "mcp_remote", status: "active", enabled: true, ...patch, + id, companyId, applicationId, name: "Check", uid: randomUUID(), transport: "mcp_remote", status: "active", enabled: true, ...patch, }).returning())[0]; - return signal ? writeToolConnectionHealthChange(db, companyId, insert) : insert(db); + return signal ? writeToolConnectionHealthChange(db, { companyId, id }, insert) : insert(db); } async function start(service = connectionHealthSchedule(db, probe)) { coordinator = createDeliveryWorkCoordinator({ owner: db, canRun: () => !draining, canReconcile: () => true, onError: error => errors.push(error) }); @@ -88,6 +89,71 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu reads.mockRestore(); writes.mockRestore(); await vi.advanceTimersByTimeAsync(due - Date.now()); await checked(); }); + it("refreshes only the changed connection after real API writes in a large company", async () => { + const checkedAt = new Date(); + const rows = await db.insert(toolConnections).values(Array.from({ length: 205 }, (_, index) => ({ + companyId, applicationId, name: `Inventory ${index}`, uid: randomUUID(), transport: "mcp_remote" as const, + config: { url: "https://health.example.test/mcp" }, + status: "active" as const, enabled: true, healthCheckedAt: checkedAt, + }))).returning(); + const tools = toolAccessService(db, { remoteHttpEndpointLookup: async () => [{ address: "8.8.8.8", family: 4 }] }); + const service = tools.scheduledHealth, load = service.load; + let returnedRows = 0; + service.load = async (...args) => { const result = await load(...args); returnedRows += result.length; return result; }; + await start(service); await idleAt(checkedAt.getTime() + 900_000); + returnedRows = 0; + for (const row of rows.slice(0, 5)) { + await tools.updateConnection(row.id, { name: `${row.name} changed` }); + await idleAt(checkedAt.getTime() + 900_000); + } + expect(returnedRows).toBe(5); + expect(probe).not.toHaveBeenCalled(); + }); + it("pages startup inventory while retaining every eligible deadline", async () => { + const checkedAt = new Date(); + await db.insert(toolConnections).values(Array.from({ length: 205 }, (_, index) => ({ + companyId, applicationId, name: `Page ${index}`, uid: randomUUID(), transport: "mcp_remote" as const, + status: "active" as const, enabled: true, healthCheckedAt: checkedAt, + }))); + const reads = vi.spyOn(db, "select"); + const rows = await connectionHealthSchedule(db, probe).load(companyId); + expect(rows).toHaveLength(205); + expect(new Set(rows.map(row => row.id)).size).toBe(205); + expect(reads).toHaveBeenCalledTimes(3); + reads.mockRestore(); + }); + it("invalidates an application's cached deadlines without removing another application's work", async () => { + const checkedAt = new Date(); + const [otherApp] = await db.insert(toolApplications).values({ companyId, name: "Other app", type: "mcp_http" }).returning(); + await connection({ healthCheckedAt: checkedAt }); + const other = await connection({ applicationId: otherApp.id, healthCheckedAt: new Date(checkedAt.getTime() + 60_000) }); + const service = connectionHealthSchedule(db, probe), load = vi.spyOn(service, "load"); + await start(service); await idleAt(checkedAt.getTime() + 900_000); load.mockClear(); + await writeToolConnectionHealthChange(db, { companyId, applicationId }, async tx => { + await tx.update(toolApplications).set({ type: "paperclip_plugin" }).where(eq(toolApplications.id, applicationId)); + }); + await idleAt(other.healthCheckedAt!.getTime() + 900_000); + expect(load).toHaveBeenCalledWith(companyId, undefined, applicationId, expect.any(AbortSignal)); + expect(load).toHaveBeenCalledTimes(1); + await writeToolConnectionHealthChange(db, { companyId, applicationId }, async tx => { + await tx.update(toolApplications).set({ type: "mcp_http" }).where(eq(toolApplications.id, applicationId)); + }); + await idleAt(checkedAt.getTime() + 900_000); + expect(probe).not.toHaveBeenCalled(); + }); + it("keeps unrelated deadlines when a connection changes application or is deleted", async () => { + const checkedAt = new Date(); + const first = await connection({ healthCheckedAt: checkedAt }); + const other = await connection({ healthCheckedAt: new Date(checkedAt.getTime() + 60_000) }); + const [plugin] = await db.insert(toolApplications).values({ companyId, name: "Excluded app", type: "paperclip_plugin" }).returning(); + await start(); await idleAt(checkedAt.getTime() + 900_000); + await writeToolConnectionHealthChange(db, first, async tx => { + await tx.update(toolConnections).set({ applicationId: plugin.id }).where(eq(toolConnections.id, first.id)); + }); + await idleAt(other.healthCheckedAt!.getTime() + 900_000); + await writeToolConnectionHealthChange(db, other, async tx => { await tx.delete(toolConnections).where(eq(toolConnections.id, other.id)); }); + await quiet(); + }); it("keeps disabled, archived, and excluded transports quiet", async () => { await connection({ enabled: false }); await connection({ status: "archived" }); await connection({ transport: "chat_sdk", connectionPurpose: "channel" }); @@ -189,7 +255,7 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu it("continues another company's checks when one scope cannot refresh", async () => { const service = connectionHealthSchedule(db, probe), load = service.load; await start(service); await quiet(); - service.load = async scope => { if (scope === companyId) throw new Error("scope unavailable"); return load(scope); }; + service.load = async (...args) => { if (args[0] === companyId) throw new Error("scope unavailable"); return load(...args); }; const otherId = randomUUID(); await db.insert(companies).values({ id: otherId, name: "Other", issuePrefix: `O${otherId.slice(0, 6)}` }); const [app] = await db.insert(toolApplications).values({ companyId: otherId, name: "Other health", type: "mcp_http" }).returning(); @@ -225,10 +291,10 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu it("keeps SDK-only chat mutations quiet and wakes the existing voice obligation", async () => { const voice = await connection({ transport: "voice", connectionPurpose: "channel", status: "draft", enabled: false }); await start(); await quiet(); - await db.transaction(async tx => { await notifyChatConnectionHealth(tx, { companyId, provider: "slack" }); }); + await db.transaction(async tx => { await notifyChatConnectionHealth(tx, { companyId, connectionId: voice.id, provider: "slack" }); }); await quiet(); await db.transaction(async tx => { - await notifyChatConnectionHealth(tx, { companyId, provider: "speko" }); + await notifyChatConnectionHealth(tx, { companyId, connectionId: voice.id, provider: "speko" }); await tx.update(toolConnections).set({ status: "active", enabled: true }).where(eq(toolConnections.id, voice.id)); }); await checked(); expect(probe.mock.calls[0][0]).toBe(voice.id); @@ -237,9 +303,9 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu const service = connectionHealthSchedule(db, probe), load = service.load; await start(service); await quiet(); let inserted = false; - service.load = async scope => { - const rows = await load(scope); - if (scope === companyId && !inserted) { inserted = true; await connection(); } + service.load = async (...args) => { + const rows = await load(...args); + if (args[0] === companyId && !inserted) { inserted = true; await connection(); } return rows; }; await connection(); await checked(2); diff --git a/server/src/services/access.ts b/server/src/services/access.ts index 279a0753e22..0414ee007dd 100644 --- a/server/src/services/access.ts +++ b/server/src/services/access.ts @@ -328,7 +328,7 @@ export function accessService(db: Db) { && grant.credentialSecretRefs.every((ref) => !connectionSecretIdsToRemove.has(ref.secretId)), ); if (connection.transport === "chat_sdk") await notifyChatEndpointWork(tx, connection.id); - if (!hasUnaffectedActiveGrant) await notifyToolConnectionHealth(tx, companyId); + if (!hasUnaffectedActiveGrant) await notifyToolConnectionHealth(tx, { companyId, id: connection.id }); await tx.update(toolConnections).set({ credentialRefs, credentialSecretRefs, diff --git a/server/src/services/chat-channels/connection-health-notifications.ts b/server/src/services/chat-channels/connection-health-notifications.ts index c2709c339e0..a29302ff3b8 100644 --- a/server/src/services/chat-channels/connection-health-notifications.ts +++ b/server/src/services/chat-channels/connection-health-notifications.ts @@ -2,7 +2,7 @@ import type { ChatProvider } from "@paperclipai/shared"; import { notifyToolConnectionHealth } from "../tool-connection-health-notifications.js"; import { chatConnectionTransport } from "./provider-catalog.js"; -export async function notifyChatConnectionHealth(transaction: object, endpoint: { companyId: string; provider: ChatProvider }) { +export async function notifyChatConnectionHealth(transaction: object, endpoint: { companyId: string; connectionId: string; provider: ChatProvider }) { if (chatConnectionTransport(endpoint.provider) === "chat_sdk") return; - await notifyToolConnectionHealth(transaction, endpoint.companyId); + await notifyToolConnectionHealth(transaction, { companyId: endpoint.companyId, id: endpoint.connectionId }); } diff --git a/server/src/services/email-channels.ts b/server/src/services/email-channels.ts index ab01aaf96f6..31b5854f87d 100644 --- a/server/src/services/email-channels.ts +++ b/server/src/services/email-channels.ts @@ -914,7 +914,7 @@ export function emailChannelService(db: Db, options: EmailChannelOptions) { updatedAt: now, }) .where(eq(chatEndpoints.id, endpoint.id)); - await notifyToolConnectionHealth(tx, endpoint.companyId); + await notifyToolConnectionHealth(tx, { companyId: endpoint.companyId, id: endpoint.connectionId }); await tx .update(toolConnections) .set({ status: "active", enabled: true, healthStatus: "ok" }) @@ -1986,7 +1986,7 @@ export function emailChannelService(db: Db, options: EmailChannelOptions) { .update(chatEndpoints) .set({ lastError: null, healthMessage: "Connected" }) .where(eq(chatEndpoints.id, endpoint.id)); - await writeToolConnectionHealthChange(db, endpoint.companyId, tx => tx + await writeToolConnectionHealthChange(db, { companyId: endpoint.companyId, id: endpoint.connectionId }, tx => tx .update(toolConnections) .set({ healthStatus: "ok", @@ -2094,7 +2094,7 @@ export function emailChannelService(db: Db, options: EmailChannelOptions) { }); } async function markError(endpoint: Endpoint, error: string) { - await writeToolConnectionHealthChange(db, endpoint.companyId, tx => tx + await writeToolConnectionHealthChange(db, { companyId: endpoint.companyId, id: endpoint.connectionId }, tx => tx .update(toolConnections) .set({ healthStatus: "degraded", @@ -2364,7 +2364,7 @@ export function emailChannelService(db: Db, options: EmailChannelOptions) { .update(chatEndpoints) .set({ status: "paused" }) .where(eq(chatEndpoints.id, endpoint.id)); - await notifyToolConnectionHealth(tx, endpoint.companyId); + await notifyToolConnectionHealth(tx, { companyId: endpoint.companyId, id: endpoint.connectionId }); await tx .update(toolConnections) .set({ enabled: false, status: "disabled" }) @@ -2458,7 +2458,7 @@ export function emailChannelService(db: Db, options: EmailChannelOptions) { updatedAt: new Date(), }) .where(eq(chatEndpoints.id, id)); - await notifyToolConnectionHealth(tx, endpoint.companyId); + await notifyToolConnectionHealth(tx, { companyId: endpoint.companyId, id: endpoint.connectionId }); await tx .update(toolConnections) .set({ diff --git a/server/src/services/email-connections.ts b/server/src/services/email-connections.ts index 264b920e971..74dc6d65a46 100644 --- a/server/src/services/email-connections.ts +++ b/server/src/services/email-connections.ts @@ -258,7 +258,7 @@ export function emailConnectionService( actorSource: actor.localImplicit ? "local_implicit" : "session", }, ); - await writeToolConnectionHealthChange(db, companyId, tx => tx + await writeToolConnectionHealthChange(db, connection, tx => tx .update(toolConnections) .set({ uid: `agentmail-account-${input.idempotencyKey}`, diff --git a/server/src/services/github-connection-events.ts b/server/src/services/github-connection-events.ts index d4786c34ab8..378f7f86341 100644 --- a/server/src/services/github-connection-events.ts +++ b/server/src/services/github-connection-events.ts @@ -327,7 +327,7 @@ export function githubConnectionEventService( }; await database.update(connectionGrants).set({ providerTenant, updatedAt: now() }) .where(and(eq(connectionGrants.id, binding.grantId), eq(connectionGrants.companyId, binding.companyId))); - await notifyToolConnectionHealth(database, binding.companyId); + await notifyToolConnectionHealth(database, { companyId: binding.companyId, id: binding.connectionId }); await database.update(toolConnections).set({ healthStatus: unavailable ? "failed" : "ok", healthMessage: unavailable diff --git a/server/src/services/smoke-lab.ts b/server/src/services/smoke-lab.ts index 2ae60f5dd8b..4bdd7532f61 100644 --- a/server/src/services/smoke-lab.ts +++ b/server/src/services/smoke-lab.ts @@ -691,7 +691,7 @@ export function smokeLabService(db: Db, options: { )); const now = new Date(); if (existing) { - const [updated] = await writeToolConnectionHealthChange(db, input.companyId, tx => tx.update(toolApplications).set({ + const [updated] = await writeToolConnectionHealthChange(db, { companyId: input.companyId, applicationId: existing.id }, tx => tx.update(toolApplications).set({ name: input.name, description: input.description, type: input.type, @@ -743,7 +743,8 @@ export function smokeLabService(db: Db, options: { lastHealthAt: now, updatedAt: now, }; - return writeToolConnectionHealthChange(db, input.companyId, async tx => { + const connectionId = existing?.id ?? randomUUID(); + return writeToolConnectionHealthChange(db, { companyId: input.companyId, id: connectionId }, async tx => { const ensureDefaultOrganizationGrant = async (connectionId: string) => { const [existingGrant] = await tx.select({ id: connectionGrants.id }).from(connectionGrants).where(and( eq(connectionGrants.companyId, input.companyId), @@ -769,6 +770,7 @@ export function smokeLabService(db: Db, options: { return { row: updated ?? existing, created: false }; } const [created] = await tx.insert(toolConnections).values({ + id: connectionId, companyId: input.companyId, name: input.name, uid: `smoke-lab/${input.name.toLowerCase().replace(/[^a-z0-9]+/g, "-").replace(/^-|-$/g, "")}`, @@ -968,7 +970,7 @@ export function smokeLabService(db: Db, options: { if (!connection || !httpSidecar) return; const sidecarUrl = httpSidecar.url; const now = new Date(); - await writeToolConnectionHealthChange(db, companyId, tx => tx.update(toolConnections).set({ + await writeToolConnectionHealthChange(db, connection, tx => tx.update(toolConnections).set({ config: { ...asRecord(connection.config), url: `${sidecarUrl}/mcp`, diff --git a/server/src/services/tool-access.ts b/server/src/services/tool-access.ts index e3d4c33feb9..6038816c6c5 100644 --- a/server/src/services/tool-access.ts +++ b/server/src/services/tool-access.ts @@ -1219,7 +1219,7 @@ export function toolAccessService( safeDefault: true, }; if (existing) { - const [updated] = await writeToolConnectionHealthChange(db, companyId, tx => tx + const [updated] = await writeToolConnectionHealthChange(db, { companyId, applicationId: existing.id }, tx => tx .update(toolApplications) .set({ name: definition.applicationName, @@ -1262,7 +1262,7 @@ export function toolAccessService( quarantineNewEntries: true, }; if (existing) { - const updated = await writeToolConnectionHealthChange(db, companyId, async tx => { + const updated = await writeToolConnectionHealthChange(db, existing, async tx => { const [updated] = await tx .update(toolConnections) .set({ @@ -1288,7 +1288,7 @@ export function toolAccessService( return { row: updated, created: false }; } const connectionId = randomUUID(); - const created = await writeToolConnectionHealthChange(db, companyId, async tx => { + const created = await writeToolConnectionHealthChange(db, { companyId, id: connectionId }, async tx => { const [created] = await tx .insert(toolConnections) .values({ @@ -2486,7 +2486,7 @@ export function toolAccessService( personalIdentityUserId || dedicatedAgentId ? [] : credentialSecretRefs; if (revivedConnectionPrevious) { forgetMcpHttpSessions(revivedConnectionPrevious.id); - [connectionRow] = await writeToolConnectionHealthChange(db, companyId, tx => tx + [connectionRow] = await writeToolConnectionHealthChange(db, revivedConnectionPrevious, tx => tx .update(toolConnections) .set({ name, @@ -2931,7 +2931,7 @@ export function toolAccessService( connectionSetupMutationFingerprint(attemptedConnection), ); if (connectionMutationIsStillCurrent) { - await notifyToolConnectionHealth(tx, companyId); + await notifyToolConnectionHealth(tx, revivedConnectionPrevious); await tx .update(toolConnections) .set({ @@ -3007,7 +3007,7 @@ export function toolAccessService( // surface the restoration failure instead of returning only the // original provider error. try { - await writeToolConnectionHealthChange(db, companyId, tx => tx + await writeToolConnectionHealthChange(db, revivedConnectionPrevious, tx => tx .update(toolConnections) .set({ status: "draft", @@ -3033,7 +3033,7 @@ export function toolAccessService( } } else if (connectionRow) { const failedConnectionId = connectionRow.id; - await writeToolConnectionHealthChange(db, companyId, tx => tx + await writeToolConnectionHealthChange(db, { companyId, id: failedConnectionId }, tx => tx .delete(toolConnections) .where(eq(toolConnections.id, failedConnectionId))) .catch(() => undefined); @@ -3148,7 +3148,7 @@ export function toolAccessService( await tx.select({ id: toolApplications.id }).from(toolApplications) .where(and(eq(toolApplications.id, connection.applicationId), eq(toolApplications.companyId, companyId))) .for("update"); - await notifyToolConnectionHealth(tx, companyId); + await notifyToolConnectionHealth(tx, connection); const [row] = await tx.update(toolConnections).set({ config, transportConfig: config, status: "active", enabled: true, ...(input.agentInstructions !== undefined ? { agentInstructions: input.agentInstructions } : {}), updatedAt: new Date() }) .where(and(eq(toolConnections.id, connection.id), eq(toolConnections.companyId, companyId))).returning(); await tx.update(toolApplications).set({ status: "active", updatedAt: new Date() }).where(and(eq(toolApplications.id, connection.applicationId), eq(toolApplications.companyId, companyId))); @@ -3505,7 +3505,7 @@ export function toolAccessService( }, tx, ); - await notifyToolConnectionHealth(tx, companyId); + await notifyToolConnectionHealth(tx, connection); const [updatedConnection] = await tx .update(toolConnections) .set({ status: "active", enabled: true, ...(input.agentInstructions !== undefined ? { agentInstructions: input.agentInstructions } : {}), updatedAt: new Date() }) @@ -4526,7 +4526,7 @@ export function toolAccessService( } const connectionId = randomUUID(); const binding = actorBinding(actor); - const row = await writeToolConnectionHealthChange(db, companyId, async tx => { + const row = await writeToolConnectionHealthChange(db, { companyId, id: connectionId }, async tx => { const [row] = await tx .insert(toolConnections) .values({ @@ -4648,7 +4648,7 @@ export function toolAccessService( await syncCredentialBindings(rows[0]!, [], tx); return rows; }) - : await writeToolConnectionHealthChange(db, existing.companyId, async tx => { + : await writeToolConnectionHealthChange(db, existing, async tx => { const rows = await write(tx); await syncCredentialBindings(rows[0]!, [], tx); await createToolRuntimeSlots(tx, runtimeSupervisor).ensureRuntimeSlot(rows[0]!); diff --git a/server/src/services/tool-access/catalog/health-schedule.ts b/server/src/services/tool-access/catalog/health-schedule.ts index 06a2781121d..130e0d0ee53 100644 --- a/server/src/services/tool-access/catalog/health-schedule.ts +++ b/server/src/services/tool-access/catalog/health-schedule.ts @@ -1,38 +1,48 @@ import { logger } from "../../../middleware/logger.js"; -import { and, eq, ne, sql } from "drizzle-orm"; +import { and, asc, eq, gt, ne, sql } from "drizzle-orm"; import { toolApplications, toolConnections, type Db } from "@paperclipai/db"; import { isConnectorFeatureEnabled } from "@paperclipai/shared"; import { instanceSettingsService } from "../../instance-settings.js"; import { connectionSetupApp } from "../connections/setup-policy.js"; import type { ConnectionHealth } from "./health.js"; -export type ConnectionHealthDeadline = { id: string; companyId: string; dueAt: number }; +export type ConnectionHealthDeadline = { id: string; companyId: string; applicationId: string; dueAt: number }; const HEALTH_FRESHNESS_MS = 15 * 60_000; export function connectionHealthSchedule(db: Db, checkHealth: ConnectionHealth["checkConnectionHealth"]) { - async function load(companyId?: string, connectionId?: string): Promise { - const rows = await db.select({ - id: toolConnections.id, - companyId: toolConnections.companyId, - checkedAt: toolConnections.healthCheckedAt, - applicationKey: toolApplications.applicationKey, - sourceTemplateKey: sql`coalesce(${toolConnections.config}->>'sourceTemplateKey', ${toolConnections.transportConfig}->>'sourceTemplateKey')`, - }).from(toolConnections).innerJoin(toolApplications, and( - eq(toolApplications.id, toolConnections.applicationId), - eq(toolApplications.companyId, toolConnections.companyId), - )).where(and( - eq(toolConnections.enabled, true), eq(toolConnections.status, "active"), - ne(toolConnections.transport, "chat_sdk"), ne(toolConnections.transport, "runtime_auth"), - ne(toolApplications.type, "paperclip_plugin"), - companyId === undefined ? undefined : eq(toolConnections.companyId, companyId), - connectionId === undefined ? undefined : eq(toolConnections.id, connectionId), - )); - const definitions = rows.map(row => ({ row, app: connectionSetupApp(row.applicationKey, row.sourceTemplateKey) })); - const settings = definitions.some(({ app }) => app?.experimentalFeatureKey) - ? await instanceSettingsService(db).getExperimental() : null; - return definitions.filter(({ app }) => !app || !settings || isConnectorFeatureEnabled(app, settings)) - .map(({ row }) => ({ id: row.id, companyId: row.companyId, - dueAt: row.checkedAt === null ? 0 : row.checkedAt.getTime() + HEALTH_FRESHNESS_MS })); + async function load(companyId?: string, connectionId?: string, applicationId?: string, signal?: AbortSignal): Promise { + const deadlines: ConnectionHealthDeadline[] = []; + let cursor: string | undefined; + while (true) { + signal?.throwIfAborted(); + const rows = await db.select({ + id: toolConnections.id, + companyId: toolConnections.companyId, + applicationId: toolConnections.applicationId, + checkedAt: toolConnections.healthCheckedAt, + applicationKey: toolApplications.applicationKey, + sourceTemplateKey: sql`coalesce(${toolConnections.config}->>'sourceTemplateKey', ${toolConnections.transportConfig}->>'sourceTemplateKey')`, + }).from(toolConnections).innerJoin(toolApplications, and( + eq(toolApplications.id, toolConnections.applicationId), + eq(toolApplications.companyId, toolConnections.companyId), + )).where(and( + eq(toolConnections.enabled, true), eq(toolConnections.status, "active"), + ne(toolConnections.transport, "chat_sdk"), ne(toolConnections.transport, "runtime_auth"), + ne(toolApplications.type, "paperclip_plugin"), + companyId === undefined ? undefined : eq(toolConnections.companyId, companyId), + connectionId === undefined ? undefined : eq(toolConnections.id, connectionId), + applicationId === undefined ? undefined : eq(toolConnections.applicationId, applicationId), + cursor === undefined ? undefined : gt(toolConnections.id, cursor), + )).orderBy(asc(toolConnections.id)).limit(100); + const definitions = rows.map(row => ({ row, app: connectionSetupApp(row.applicationKey, row.sourceTemplateKey) })); + const settings = definitions.some(({ app }) => app?.experimentalFeatureKey) + ? await instanceSettingsService(db).getExperimental() : null; + deadlines.push(...definitions.filter(({ app }) => !app || !settings || isConnectorFeatureEnabled(app, settings)) + .map(({ row }) => ({ id: row.id, companyId: row.companyId, applicationId: row.applicationId, + dueAt: row.checkedAt === null ? 0 : row.checkedAt.getTime() + HEALTH_FRESHNESS_MS }))); + if (rows.length < 100) return deadlines; + cursor = rows[rows.length - 1].id; + } } async function check(id: string): Promise { diff --git a/server/src/services/tool-access/catalog/health.ts b/server/src/services/tool-access/catalog/health.ts index c06ef95e324..86ef894a54b 100644 --- a/server/src/services/tool-access/catalog/health.ts +++ b/server/src/services/tool-access/catalog/health.ts @@ -48,7 +48,7 @@ export function createConnectionHealth(db: Db, { now, audit, getConnectionRow, d message: string | null, ) { const now = new Date(); - const [updated] = await writeToolConnectionHealthChange(db, connection.companyId, async tx => { + const [updated] = await writeToolConnectionHealthChange(db, connection, async tx => { await notifyChatConnectionCommandRegistration(tx, connection.id, connection.transport); return tx.update(toolConnections) .set({ diff --git a/server/src/services/tool-access/catalog/publication.ts b/server/src/services/tool-access/catalog/publication.ts index 5257cbb3bfe..77fcb16ba2d 100644 --- a/server/src/services/tool-access/catalog/publication.ts +++ b/server/src/services/tool-access/catalog/publication.ts @@ -269,7 +269,7 @@ export function createToolCatalogPublication(db: Db, { now, audit, getConnection : refreshOptions.enableAllByDefault ? { ...connection.transportConfig, quarantineNewEntries: false } : connection.transportConfig; - const [updatedConnection] = await writeToolConnectionHealthChange(db, connection.companyId, tx => tx + const [updatedConnection] = await writeToolConnectionHealthChange(db, connection, tx => tx .update(toolConnections) .set({ config: normalizedConfig, diff --git a/server/src/services/tool-access/connections/removal.ts b/server/src/services/tool-access/connections/removal.ts index 9be31478434..e8dca87f4a3 100644 --- a/server/src/services/tool-access/connections/removal.ts +++ b/server/src/services/tool-access/connections/removal.ts @@ -539,7 +539,7 @@ async function removeConnection( const archived = await db.transaction(async (tx) => { if (connection.transport === "chat_sdk") await notifyChatEndpointWork(tx, connection.id); - await notifyToolConnectionHealth(tx, connection.companyId); + await notifyToolConnectionHealth(tx, connection); await tx.select({ id: toolApplications.id }).from(toolApplications) .where(and(eq(toolApplications.id, connection.applicationId), eq(toolApplications.companyId, connection.companyId))) .for("update"); diff --git a/server/src/services/tool-access/oauth/access-finalization.ts b/server/src/services/tool-access/oauth/access-finalization.ts index 85569d99d3a..764187e755f 100644 --- a/server/src/services/tool-access/oauth/access-finalization.ts +++ b/server/src/services/tool-access/oauth/access-finalization.ts @@ -113,7 +113,7 @@ async function finalizeOAuthAccess( ) { throw conflict("This connection already uses a company identity"); } - connection = await writeToolConnectionHealthChange(db, companyId, async tx => { + connection = await writeToolConnectionHealthChange(db, connection, async tx => { const [updated] = await tx .update(toolConnections) .set({ @@ -328,7 +328,7 @@ async function finalizeOAuthAccess( updatedAt: nowAt, }) .where(eq(connectionGrants.id, personalGrant.id)); - await notifyToolConnectionHealth(tx, companyId); + await notifyToolConnectionHealth(tx, connection); [connection] = await tx .update(toolConnections) .set({ diff --git a/server/src/services/tool-access/oauth/cloud-callback.ts b/server/src/services/tool-access/oauth/cloud-callback.ts index a61c3b203f3..82c727df61f 100644 --- a/server/src/services/tool-access/oauth/cloud-callback.ts +++ b/server/src/services/tool-access/oauth/cloud-callback.ts @@ -356,7 +356,7 @@ export function createCloudOAuthCallback(db: Db, dependencies: Dependencies) { scopes: [...profile.scopes], }, }; - await notifyToolConnectionHealth(tx, connection.companyId); + await notifyToolConnectionHealth(tx, connection); [connection] = await tx .update(toolConnections) .set({ diff --git a/server/src/services/tool-access/oauth/connection-refresh.ts b/server/src/services/tool-access/oauth/connection-refresh.ts index fd3ae9b7aff..1e0d70445ae 100644 --- a/server/src/services/tool-access/oauth/connection-refresh.ts +++ b/server/src/services/tool-access/oauth/connection-refresh.ts @@ -60,7 +60,7 @@ const oauthRefreshFlights = new Map>(); reauthorizationRequiredAt: now().toISOString(), }, }; - const [updated] = await writeToolConnectionHealthChange(db, connection.companyId, tx => tx + const [updated] = await writeToolConnectionHealthChange(db, connection, tx => tx .update(toolConnections) .set({ status: "draft", diff --git a/server/src/services/tool-access/oauth/direct-callback.ts b/server/src/services/tool-access/oauth/direct-callback.ts index cd9b0723def..eb4f53022ad 100644 --- a/server/src/services/tool-access/oauth/direct-callback.ts +++ b/server/src/services/tool-access/oauth/direct-callback.ts @@ -332,7 +332,7 @@ export function createDirectOAuthCallback(db: Db, dependencies: Dependencies) { }, }, }; - await notifyToolConnectionHealth(tx, connection.companyId); + await notifyToolConnectionHealth(tx, connection); const [updatedConnection] = await tx .update(toolConnections) .set({ @@ -612,7 +612,7 @@ export function createDirectOAuthCallback(db: Db, dependencies: Dependencies) { }, }, }; - await notifyToolConnectionHealth(tx, connection.companyId); + await notifyToolConnectionHealth(tx, connection); const [updatedConnection] = await tx .update(toolConnections) .set({ diff --git a/server/src/services/tool-access/oauth/grant-refresh.ts b/server/src/services/tool-access/oauth/grant-refresh.ts index 6a1657d8482..48e09cf27b9 100644 --- a/server/src/services/tool-access/oauth/grant-refresh.ts +++ b/server/src/services/tool-access/oauth/grant-refresh.ts @@ -394,7 +394,7 @@ export function createOAuthGrantRefresh(db: Db, { now, store, leases, currentClo connection.id, connection.companyId, ); - const [reauthorizationRequired] = await writeToolConnectionHealthChange(db, latestConnection.companyId, tx => tx + const [reauthorizationRequired] = await writeToolConnectionHealthChange(db, latestConnection, tx => tx .update(toolConnections) .set({ status: "draft", diff --git a/server/src/services/tool-access/oauth/vercel-callback.ts b/server/src/services/tool-access/oauth/vercel-callback.ts index f8391eda77c..111e3bdb171 100644 --- a/server/src/services/tool-access/oauth/vercel-callback.ts +++ b/server/src/services/tool-access/oauth/vercel-callback.ts @@ -184,7 +184,7 @@ export function createVercelOAuthCallback(db: Db, dependencies: Dependencies) { status: connection.status, enabled: connection.enabled, }; - [connection] = await writeToolConnectionHealthChange(db, connection.companyId, tx => tx + [connection] = await writeToolConnectionHealthChange(db, connection, tx => tx .update(toolConnections) .set({ status: "active", diff --git a/server/src/services/tool-connection-health-notifications.ts b/server/src/services/tool-connection-health-notifications.ts index 163023d674a..075983d2972 100644 --- a/server/src/services/tool-connection-health-notifications.ts +++ b/server/src/services/tool-connection-health-notifications.ts @@ -1,12 +1,27 @@ import type { Db } from "@paperclipai/db"; import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js"; -export const notifyToolConnectionHealth = (transaction: object, companyId: string) => - notifyDeliveryWork(transaction, DELIVERY_QUEUES.toolConnectionHealth, companyId); +type HealthChange = string | { companyId: string; id: string } | { companyId: string; applicationId: string }; +export type HealthScope = { companyId?: string; connectionId?: string; applicationId?: string }; -export function writeToolConnectionHealthChange(db: Db, companyId: string, write: (transaction: Db) => Promise): Promise { +export function healthScope(hint: string): HealthScope { + if (hint === "*") return {}; + const [companyId, kind, id, extra] = hint.split(":"); + if (companyId && kind === undefined) return { companyId }; + if (companyId && id && extra === undefined) { + if (kind === "connection") return { companyId, connectionId: id }; + if (kind === "application") return { companyId, applicationId: id }; + } + throw new Error("Invalid tool connection health scope"); +} + +export const notifyToolConnectionHealth = (transaction: object, target: HealthChange) => + notifyDeliveryWork(transaction, DELIVERY_QUEUES.toolConnectionHealth, typeof target === "string" ? target + : "id" in target ? `${target.companyId}:connection:${target.id}` : `${target.companyId}:application:${target.applicationId}`); + +export function writeToolConnectionHealthChange(db: Db, target: HealthChange, write: (transaction: Db) => Promise): Promise { return db.transaction(async tx => { - await notifyToolConnectionHealth(tx, companyId); + await notifyToolConnectionHealth(tx, target); return write(tx as unknown as Db); }); } diff --git a/server/src/services/tool-connection-health-work.ts b/server/src/services/tool-connection-health-work.ts index 6743dd78dfe..6d684ba4a42 100644 --- a/server/src/services/tool-connection-health-work.ts +++ b/server/src/services/tool-connection-health-work.ts @@ -1,3 +1,4 @@ +import { healthScope, type HealthScope } from "./tool-connection-health-notifications.js"; import { takeCompanyWorkBatch } from "./company-work-batch.js"; import type { createDeliveryWorkCoordinator } from "./delivery-work-coordinator.js"; import { DELIVERY_QUEUES } from "./delivery-work-notifications.js"; @@ -6,7 +7,7 @@ import type { ConnectionHealthDeadline } from "./tool-access/catalog/health-sche export function registerToolConnectionHealthWork( coordinator: ReturnType, service: { - load(companyId?: string): Promise; + load(companyId?: string, connectionId?: string, applicationId?: string, signal?: AbortSignal): Promise; check(id: string): Promise; }, canRun: () => boolean, @@ -14,6 +15,10 @@ export function registerToolConnectionHealthWork( const dirty = new Set(["*"]); const deadlines = new Map(); const failures = new Map(); + const matches = (entry: ConnectionHealthDeadline, scope: HealthScope) => + (scope.companyId === undefined || entry.companyId === scope.companyId) + && (scope.connectionId === undefined || entry.id === scope.connectionId) + && (scope.applicationId === undefined || entry.applicationId === scope.applicationId); const retryMs = 30_000; let earliest = 0; let lastCompany: string | null = null; @@ -38,21 +43,26 @@ export function registerToolConnectionHealthWork( const scopes = dirty.has("*") ? ["*"] : [...dirty]; dirty.clear(); const errors: unknown[] = []; - const unavailable = new Set(); + const unavailable: HealthScope[] = []; for (let index = 0; index < scopes.length; index++) { const scope = scopes[index]; if (signal.aborted || !canRun()) { for (const remaining of scopes.slice(index)) dirty.add(remaining); return; } + let target: HealthScope | undefined; try { - const rows = await service.load(scope === "*" ? undefined : scope); - for (const entry of deadlines.values()) if (scope === "*" || entry.companyId === scope) deadlines.delete(entry.id); + target = healthScope(scope); + const rows = await service.load(target.companyId, target.connectionId, target.applicationId, signal); + if (target.connectionId) { + const current = deadlines.get(target.connectionId); + if (current && matches(current, target)) deadlines.delete(current.id); + } else for (const entry of deadlines.values()) if (matches(entry, target)) deadlines.delete(entry.id); for (const entry of rows) deadlines.set(entry.id, entry); - for (const id of failures.keys()) { - const current = deadlines.get(id); - if (!current || current.dueAt > Date.now()) failures.delete(id); - } - } catch (error) { dirty.add(scope); unavailable.add(scope); errors.push(error); } + } catch (error) { dirty.add(target ? scope : "*"); unavailable.push(target ?? {}); errors.push(error); } } - const due = [...deadlines.values()].filter(entry => !unavailable.has("*") && !unavailable.has(entry.companyId) + for (const id of failures.keys()) { + const current = deadlines.get(id); + if (!current || current.dueAt > Date.now()) failures.delete(id); + } + const due = [...deadlines.values()].filter(entry => !unavailable.some(scope => matches(entry, scope)) && Math.max(entry.dueAt, failures.get(entry.id)?.retryAt ?? 0) <= Date.now()) .sort((a, b) => a.dueAt - b.dueAt || a.id.localeCompare(b.id)); const batch = takeCompanyWorkBatch(due, 25, lastCompany); diff --git a/server/src/services/tool-gateway.ts b/server/src/services/tool-gateway.ts index 6a0b8fb60bb..4bf6f93f8a7 100644 --- a/server/src/services/tool-gateway.ts +++ b/server/src/services/tool-gateway.ts @@ -3805,7 +3805,7 @@ export function createToolGatewayService( message: string | null, ) { const now = new Date(); - await writeToolConnectionHealthChange(db, connection.companyId, tx => tx + await writeToolConnectionHealthChange(db, connection, tx => tx .update(toolConnections) .set({ healthStatus: status,