diff --git a/doc/idle-discord-command-registration.md b/doc/idle-discord-command-registration.md index bbbf7cf1ae0..c5ebe702299 100644 --- a/doc/idle-discord-command-registration.md +++ b/doc/idle-discord-command-registration.md @@ -26,6 +26,20 @@ connection and a known bot/application ID and guild ID. Attention plus failed connection health remains excluded. These are the old runtime discovery filters; company status and unrelated endpoint metadata do not add new policy. +## Receiver construction and ownership + +The shared leased receiver owner may renew during construction only while the +exact runtime context remains in the service's pending initialization map. +Renewal still checks the company, endpoint, generation, credentials and original +lease token. A missing listener alone grants no ownership. Once attached, the +registered runtime must match that context. The service verifies durable ownership +before it calls listener initialization and retains its checks after that await. +That verification waits for an older in-flight renewal, rechecks local ownership, +and then performs a fresh durable renewal. The older result alone is insufficient. +Token loss, pause, removal, +replacement and shutdown retain their existing fences and listener stop order. +This also applies to the Photon receiver through the same ownership module. + ## Scheduling and lifecycle Startup reconstructs eligible state in UUID keyset pages of 100. Committed hints diff --git a/server/src/__tests__/chat-channels.integration.test.ts b/server/src/__tests__/chat-channels.integration.test.ts index c79878aaa01..7dc31fab4cb 100644 --- a/server/src/__tests__/chat-channels.integration.test.ts +++ b/server/src/__tests__/chat-channels.integration.test.ts @@ -8800,6 +8800,96 @@ describeEmbeddedPostgres("chat channel control-plane integration", () => { } }); + it.each(["owned", "slow-renewal", "lost-token", "foreign-token", "foreign-token-before-renewal", "paused", "removed", "replacement", "credential-change", "foreign-company", "shutdown"] as const)("fences receiver construction across renewal: %s", async fault => { + const f = await registrationFixture("success", { discordGatewayLeaseTtlMs: 60_000, discordGatewayLeaseRenewalIntervalMs: fault === "slow-renewal" || fault === "foreign-token-before-renewal" ? 5000 : 50 }); + const listenerStarted = vi.fn(async () => {}); + f.runtime.initializeHook = listenerStarted; + let release!: () => void; + let entered!: () => void; + const gate = new Promise(resolve => { release = resolve; }); + const constructing = new Promise(resolve => { entered = resolve; }); + const replace = f.runtime.replaceEndpoint.bind(f.runtime); + const replacement = vi.spyOn(f.runtime, "replaceEndpoint").mockImplementation(async options => { + entered(); await gate; return replace(options); + }); + const configured = f.configure().then(value => ({ value }), error => ({ error })); + const foreignToken = randomUUID(); + let shutdown: Promise | undefined; + try { + await constructing; + if (fault === "lost-token") await db.delete(chatEndpointLeases).where(and(eq(chatEndpointLeases.endpointId, f.endpoint.id), eq(chatEndpointLeases.leaseKey, "discord_gateway_runtime"))); + if (fault === "foreign-token" || fault === "foreign-token-before-renewal") await db.update(chatEndpointLeases).set({ token: foreignToken }).where(and(eq(chatEndpointLeases.endpointId, f.endpoint.id), eq(chatEndpointLeases.leaseKey, "discord_gateway_runtime"))); + if (fault === "paused") await db.update(chatEndpoints).set({ status: "paused" }).where(eq(chatEndpoints.id, f.endpoint.id)); + if (fault === "removed") await db.delete(chatEndpoints).where(eq(chatEndpoints.id, f.endpoint.id)); + if (fault === "replacement") await db.update(chatEndpoints).set({ setup: sql`${chatEndpoints.setup} || jsonb_build_object('runtimeGeneration', (${chatEndpoints.setup}->>'runtimeGeneration')::integer + 1)` }).where(eq(chatEndpoints.id, f.endpoint.id)); + if (fault === "foreign-company") { + const foreign = await seedCompany(); + await expect(db.update(chatEndpoints).set({ companyId: foreign.companyId }).where(eq(chatEndpoints.id, f.endpoint.id))) + .rejects.toMatchObject({ cause: { code: "23503" } }); + } + if (fault === "credential-change") await db.update(toolConnections).set({ credentialSecretRefs: [] }).where(eq(toolConnections.id, f.endpoint.connectionId)); + if (fault === "shutdown") shutdown = f.service.shutdown(); + await new Promise(resolve => setTimeout(resolve, 200)); + release(); + const result = await configured; + const [lease] = await db.select().from(chatEndpointLeases).where(and(eq(chatEndpointLeases.endpointId, f.endpoint.id), eq(chatEndpointLeases.leaseKey, "discord_gateway_runtime"))); + if (fault === "owned" || fault === "slow-renewal" || fault === "foreign-company") { + expect(result).not.toHaveProperty("error"); + expect(listenerStarted).toHaveBeenCalledOnce(); + expect(f.runtime.get(f.endpoint.id)).not.toBeNull(); + expect(lease).toBeDefined(); + } else { + expect(result).toHaveProperty("error"); + expect(listenerStarted).not.toHaveBeenCalled(); + expect(f.runtime.get(f.endpoint.id)).toBeNull(); + if (fault === "foreign-token" || fault === "foreign-token-before-renewal") expect(lease?.token).toBe(foreignToken); + else expect(lease).toBeUndefined(); + } + } finally { + release(); await configured; await shutdown; replacement.mockRestore(); + if (fault === "removed") await f.service.shutdown(); + else await retireRegistrationFixture(f.service, f.endpoint.id); + } + }); + + it("rechecks durable ownership after an earlier renewal finishes before listener start", async () => { + const f = await registrationFixture("success", { discordGatewayLeaseTtlMs: 60_000, discordGatewayLeaseRenewalIntervalMs: 50 }); + const listenerStarted = vi.fn(async () => {}); f.runtime.initializeHook = listenerStarted; + let releaseConstruction!: () => void, constructionEntered!: () => void, releaseRenewal!: () => void, renewalEntered!: () => void; + const construction = new Promise(resolve => { releaseConstruction = resolve; }); + const enteredConstruction = new Promise(resolve => { constructionEntered = resolve; }); + const renewal = new Promise(resolve => { releaseRenewal = resolve; }); + const enteredRenewal = new Promise(resolve => { renewalEntered = resolve; }); + let constructionPending = false; + const replace = f.runtime.replaceEndpoint.bind(f.runtime); + const replacement = vi.spyOn(f.runtime, "replaceEndpoint").mockImplementation(async options => { + constructionPending = true; constructionEntered(); await construction; return replace(options); + }); + let gated = false; + const transact = db.transaction.bind(db); + const transaction = vi.spyOn(db, "transaction").mockImplementation((async (...args: Parameters) => { + const result = await transact(...args); + if (result === true && constructionPending && !gated) { gated = true; renewalEntered(); await renewal; } + return result; + }) as typeof db.transaction); + const configured = f.configure().then(value => ({ value }), error => ({ error })); + const foreignToken = randomUUID(); + try { + await enteredConstruction; await enteredRenewal; + await db.update(chatEndpointLeases).set({ token: foreignToken }).where(and(eq(chatEndpointLeases.endpointId, f.endpoint.id), eq(chatEndpointLeases.leaseKey, "discord_gateway_runtime"))); + releaseConstruction(); + await vi.waitFor(() => expect(f.runtime.get(f.endpoint.id)).not.toBeNull()); + releaseRenewal(); + expect(await configured).toHaveProperty("error"); + expect(listenerStarted).not.toHaveBeenCalled(); + const [lease] = await db.select().from(chatEndpointLeases).where(and(eq(chatEndpointLeases.endpointId, f.endpoint.id), eq(chatEndpointLeases.leaseKey, "discord_gateway_runtime"))); + expect(lease?.token).toBe(foreignToken); + } finally { + releaseConstruction(); releaseRenewal(); await configured; transaction.mockRestore(); replacement.mockRestore(); + await retireRegistrationFixture(f.service, f.endpoint.id); + } + }); + it("restarts a receiver after token loss only after the old listener has stopped", async () => { const f = await registrationFixture("success", { discordGatewayLeaseTtlMs: 60_000, discordGatewayLeaseRenewalIntervalMs: 50 }); const onError = vi.fn(); diff --git a/server/src/services/chat-channels.ts b/server/src/services/chat-channels.ts index 010e9670202..e6cdc79894a 100644 --- a/server/src/services/chat-channels.ts +++ b/server/src/services/chat-channels.ts @@ -513,6 +513,7 @@ export function chatChannelService(db: Db, options: ChatChannelServiceOptions) { const runtimeInitializations = new Map< string, { + context: RuntimeContext; discordGatewayOwned: boolean; promise: Promise; version: string; @@ -574,7 +575,7 @@ export function chatChannelService(db: Db, options: ChatChannelServiceOptions) { const discordGatewayStopTasks = new Set>(); let shuttingDown = false; let receiverWork: ReturnType | undefined; - const { discordGatewayNotOwnedError, isDiscordGatewayNotOwnedError, releaseDiscordGatewayLeaseRow, stopDiscordGatewayOwnership, renewDiscordGatewayOwnershipForMessageAdmission, ensureDiscordGatewayRuntimeIsCurrent, acquireDiscordGatewayOwnership, discordGatewayRuntimeIsCurrent } = createChatReceiverOwnership(db, { runtime, runtimeVersions, runtimeCallbackEndpoint, redactError, leasedChatProvider, discordGatewayOwnerships, discordGatewayStopTasks, isShuttingDown: () => shuttingDown, onStopped: id => receiverWork?.changed(id), get discordGatewayLeaseTtlMs() { return options.discordGatewayLeaseTtlMs; }, get discordGatewayLeaseRenewalIntervalMs() { return options.discordGatewayLeaseRenewalIntervalMs; }, get discordGatewayLeaseWaitMs() { return options.discordGatewayLeaseWaitMs; }, get renewDiscordGatewayLease() { return options.renewDiscordGatewayLease; } }); + const { discordGatewayNotOwnedError, isDiscordGatewayNotOwnedError, releaseDiscordGatewayLeaseRow, stopDiscordGatewayOwnership, renewDiscordGatewayOwnershipForMessageAdmission, ensureDiscordGatewayRuntimeIsCurrent, acquireDiscordGatewayOwnership, discordGatewayRuntimeIsCurrent } = createChatReceiverOwnership(db, { runtime, runtimeVersions, runtimeCallbackEndpoint, redactError, leasedChatProvider, discordGatewayOwnerships, discordGatewayStopTasks, isShuttingDown: () => shuttingDown, isRuntimeInitializing: (endpointId, context) => { const pending = runtimeInitializations.get(endpointId); return pending?.context === context && pending.discordGatewayOwned; }, onStopped: id => receiverWork?.changed(id), get discordGatewayLeaseTtlMs() { return options.discordGatewayLeaseTtlMs; }, get discordGatewayLeaseRenewalIntervalMs() { return options.discordGatewayLeaseRenewalIntervalMs; }, get discordGatewayLeaseWaitMs() { return options.discordGatewayLeaseWaitMs; }, get renewDiscordGatewayLease() { return options.renewDiscordGatewayLease; } }); const { receiveSdkMessages } = createChatSdkMessageReceiver({ endpointRecord, processMessage, ingressOnly: (provider) => provider === "imessage-photon" @@ -985,10 +986,12 @@ export function chatChannelService(db: Db, options: ChatChannelServiceOptions) { } const initialization: { + context: RuntimeContext; discordGatewayOwned: boolean; promise: Promise; version: string; } = { + context, discordGatewayOwned: false, promise: undefined as unknown as Promise, version: context.version, @@ -1091,6 +1094,9 @@ export function chatChannelService(db: Db, options: ChatChannelServiceOptions) { }); context.endpointRuntime = instance; runtimeContexts.set(instance as object, context); + if (discordOwnership && !(await ensureDiscordGatewayRuntimeIsCurrent(endpoint.id, context, { renewLease: true }))) { + throw discordGatewayNotOwnedError(); + } await instance.initialize(); const latest = await endpointRecord(endpoint.id); const stillRegistered = runtime.get(endpoint.id) === instance; diff --git a/server/src/services/chat-channels/runtime/ownership.ts b/server/src/services/chat-channels/runtime/ownership.ts index 4a673d41505..1956a55beda 100644 --- a/server/src/services/chat-channels/runtime/ownership.ts +++ b/server/src/services/chat-channels/runtime/ownership.ts @@ -18,6 +18,7 @@ type OwnershipOptions = { discordGatewayOwnerships: Map; discordGatewayStopTasks: Set>; isShuttingDown: () => boolean; + isRuntimeInitializing: (endpointId: string, context: RuntimeContext) => boolean; onStopped?: (endpointId: string) => void; readonly discordGatewayLeaseTtlMs?: number; readonly discordGatewayLeaseRenewalIntervalMs?: number; @@ -116,19 +117,24 @@ export function createChatReceiverOwnership(db: Db, options: OwnershipOptions) { return ownership.stopPromise; } - async function renewDiscordGatewayOwnership( - ownership: DiscordGatewayOwnership, - ): Promise { + function canRenewLocalOwnership(ownership: DiscordGatewayOwnership): boolean { if ( options.isShuttingDown() || ownership.stopping || discordGatewayOwnerships.get(ownership.endpointId) !== ownership || - ownership.context.discordGatewayOwned !== true || - ownership.context.endpointRuntime === undefined || - runtime.get(ownership.endpointId) !== ownership.context.endpointRuntime + ownership.context.discordGatewayOwned !== true ) { return false; } + return ownership.context.endpointRuntime === undefined + ? options.isRuntimeInitializing(ownership.endpointId, ownership.context) + : runtime.get(ownership.endpointId) === ownership.context.endpointRuntime; + } + + async function renewDiscordGatewayOwnership( + ownership: DiscordGatewayOwnership, + ): Promise { + if (!canRenewLocalOwnership(ownership)) return false; const proposedExpiresAt = new Date( Date.now() + (options.discordGatewayLeaseTtlMs ?? DISCORD_GATEWAY_LEASE_TTL_MS), @@ -158,7 +164,8 @@ export function createChatReceiverOwnership(db: Db, options: OwnershipOptions) { }, ["verifying", "active", "attention"], ); - if (!current || !leasedChatProvider(current.provider)) return false; + if (!current || current.companyId !== ownership.companyId || !leasedChatProvider(current.provider) + || !canRenewLocalOwnership(ownership)) return false; const expiresAt = new Date( now.getTime() + (options.discordGatewayLeaseTtlMs ?? DISCORD_GATEWAY_LEASE_TTL_MS), @@ -265,6 +272,7 @@ export function createChatReceiverOwnership(db: Db, options: OwnershipOptions) { async function ensureDiscordGatewayRuntimeIsCurrent( endpointId: string, context: RuntimeContext, + verification?: { renewLease: true }, ): Promise { const ownership = discordGatewayOwnerships.get(endpointId); if ( @@ -277,7 +285,12 @@ export function createChatReceiverOwnership(db: Db, options: OwnershipOptions) { ) { return false; } - if (ownership.expiresAt.getTime() <= Date.now()) { + if (verification?.renewLease) { + if (ownership.renewal) await ownership.renewal; + if (!canRenewLocalOwnership(ownership)) return false; + await refreshDiscordGatewayOwnership(ownership); + if (!canRenewLocalOwnership(ownership)) return false; + } else if (ownership.expiresAt.getTime() <= Date.now()) { // A host can resume after the local TTL while its durable token is still // authoritative. Reclaim that exact token before accepting buffered // callbacks. A standby takeover changes the token, so the CAS fails and