diff --git a/server/src/__tests__/subscription-refresh.test.ts b/server/src/__tests__/subscription-refresh.test.ts index cb2f4bddc49..dc59f6a2fce 100644 --- a/server/src/__tests__/subscription-refresh.test.ts +++ b/server/src/__tests__/subscription-refresh.test.ts @@ -6,7 +6,7 @@ import { idleWorkSnapshot } from "../services/task-admission.js"; const mocks = vi.hoisted(() => ({ accounts: vi.fn(), credential: vi.fn(), register: vi.fn(), probe: vi.fn() })); vi.mock("../services/ai-connections.js", () => ({ aiConnectionService: () => ({ subscriptionAccounts: mocks.accounts, credential: mocks.credential }) })); -vi.mock("../services/subscriptions.js", () => ({ subscriptionService: () => ({ register: mocks.register }) })); +vi.mock("../services/subscriptions.js", () => ({ subscriptionService: () => ({ register: mocks.register }), subscriptionSession: async () => null })); vi.mock("../services/subscription-identity.js", () => ({ probeSubscriptionIdentity: mocks.probe })); function deferred() { diff --git a/server/src/__tests__/subscriptions.test.ts b/server/src/__tests__/subscriptions.test.ts index 1b8f391dbc9..99d1ba61517 100644 --- a/server/src/__tests__/subscriptions.test.ts +++ b/server/src/__tests__/subscriptions.test.ts @@ -2,7 +2,7 @@ import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; import { randomUUID } from "node:crypto"; import { readFile } from "node:fs/promises"; import { and, eq, sql } from "drizzle-orm"; -import { createDb, companies, agents, heartbeatRuns, costEvents, activityLog, aiSubscriptions, aiSubscriptionPrices, aiSubscriptionConnections, toolApplications, toolConnections, connectionGrants, connectionGrantMembers } from "@paperclipai/db"; +import { createDb, companies, agents, heartbeatRuns, costEvents, activityLog, aiSubscriptions, aiSubscriptionPrices, aiSubscriptionConnections, companySecrets, toolApplications, toolConnections, connectionGrants, connectionGrantMembers } from "@paperclipai/db"; import { startEmbeddedPostgresTestDatabase } from "@paperclipai/db/test-embedded-postgres"; import { monthlySubscriptionCents, subscriptionPlan, subscriptionPriceSchema } from "@paperclipai/shared"; import { subscriptionService, type SubscriptionConnection } from "../services/subscriptions.js"; @@ -475,6 +475,39 @@ describe("durable subscription reporting", () => { await refreshSubscriptionConnection(db, input, id, fixtureRequest(usage("pro"))); expect((await subscriptionCostReport(db, c, owner)).accounts[0].price.amountCents).toBe("1750.0000000"); }); + it("keeps one unverified account through OAuth refreshes and separates a reconnect", async () => { + const c = await company(), connected = await connection(c, "grok-first", "xai"), service = subscriptionService(db); + const secretId = randomUUID(); + await db.insert(companySecrets).values({ id: secretId, companyId: c, key: `ai-${connected.grantId}`, name: `ai-${connected.grantId}` }); + await db.update(connectionGrants).set({ credentialSecretRefs: [{ secretId, configPath: "ai.credential" }] }).where(eq(connectionGrants.id, connected.grantId)); + const input = { ...connected, session: { secretId, epoch: 0 } }; + const id = await service.register(input); + await price(c, id, "3000"); + for (const credential of ["grok-refreshed", "grok-refreshed-again"]) expect(await service.register({ ...input, credential })).toBe(id); + await refreshSubscriptionConnection(db, { ...input, credential: "grok-refreshed-again" }, id, fixtureRequest(usage())); + let report = await subscriptionCostReport(db, c, owner); + expect(report.accounts.map(account => account.id)).toEqual([id]); + expect(report.accounts[0]).toMatchObject({ identityVerified: false, price: { amountCents: "3000.0000000" } }); + await db.update(companySecrets).set({ aiSessionEpoch: 1 }).where(eq(companySecrets.id, secretId)); + const reconnected = await service.register({ ...input, credential: "grok-reconnected", session: { secretId, epoch: 1 } }); + expect(reconnected).not.toBe(id); + expect(await service.register({ ...input, credential: "grok-reconnected-refreshed", session: { secretId, epoch: 1 } })).toBe(reconnected); + report = await subscriptionCostReport(db, c, owner); + expect(report.accounts.map(account => account.id).sort()).toEqual([id, reconnected].sort()); + }); + it("rejects an unverified registration from a session older than the latest reconnect", async () => { + const c = await company(), connected = await connection(c, "grok-old", "xai"), service = subscriptionService(db); + const secretId = randomUUID(); + await db.insert(companySecrets).values({ id: secretId, companyId: c, key: `ai-${connected.grantId}`, name: `ai-${connected.grantId}` }); + await db.update(connectionGrants).set({ credentialSecretRefs: [{ secretId, configPath: "ai.credential" }] }).where(eq(connectionGrants.id, connected.grantId)); + await db.update(companySecrets).set({ aiSessionEpoch: 1 }).where(eq(companySecrets.id, secretId)); + const current = await service.register({ ...connected, credential: "grok-new", session: { secretId, epoch: 1 } }); + await expect(service.register({ ...connected, credential: "grok-old-refreshed", session: { secretId, epoch: 0 } })).rejects.toMatchObject({ status: 409 }); + await expect(service.register({ ...connected, credential: "grok-old-refreshed" })).rejects.toMatchObject({ status: 409 }); + const [binding] = await db.select().from(aiSubscriptionConnections).where(eq(aiSubscriptionConnections.grantId, connected.grantId)); + expect(binding.subscriptionId).toBe(current); + expect((await subscriptionCostReport(db, c, owner)).accounts.map(account => account.id)).toEqual([current]); + }); it("ignores a delayed provider observation after reconnecting another account", async () => { const c = await company(), input = await connection(c), service = subscriptionService(db), id = await service.register(input); let finish!: (response: Response) => void; diff --git a/server/src/services/ai-connection-runtime.ts b/server/src/services/ai-connection-runtime.ts index 2620e8cbb53..54fe1ded35b 100644 --- a/server/src/services/ai-connection-runtime.ts +++ b/server/src/services/ai-connection-runtime.ts @@ -384,6 +384,7 @@ export async function prepareManagedAiRuntime( companyId: input.companyId, connectionId: selection.connection.id, grantId: selection.grant.id, provider: selection.attribution.provider, name: selection.connection.name, ownerUserId: selection.grant.subjectUserId, credential: value, + session: credentialRef && freshness ? { secretId: credentialRef.secretId, epoch: freshness.epoch } : null, }; attribution.subscriptionId = await subscriptionService(db).register(subscription); void refreshSubscriptionConnection(db, subscription, attribution.subscriptionId); diff --git a/server/src/services/subscription-refresh.ts b/server/src/services/subscription-refresh.ts index 24322d099ee..6cec8ee524f 100644 --- a/server/src/services/subscription-refresh.ts +++ b/server/src/services/subscription-refresh.ts @@ -1,7 +1,7 @@ import { and, eq, lt, or, isNull } from "drizzle-orm"; import { aiSubscriptions, aiSubscriptionConnections, type Db } from "@paperclipai/db"; import { aiConnectionService } from "./ai-connections.js"; -import { subscriptionService, type SubscriptionConnection } from "./subscriptions.js"; +import { subscriptionService, subscriptionSession, type SubscriptionConnection } from "./subscriptions.js"; import { probeSubscriptionIdentity } from "./subscription-identity.js"; import { logger } from "../middleware/logger.js"; import { tooManyRequests } from "../errors.js"; @@ -58,9 +58,10 @@ export function refreshCompanySubscriptions(db: Db, companyId: string, userId: s for (const row of await service.subscriptionAccounts(companyId, userId)) { if (row.summary.status !== "connected") continue; try { + const session = await subscriptionSession(db, companyId, row.grant.credentialSecretRefs); const credential = await service.credential(row); const input: SubscriptionConnection = { companyId, connectionId: row.connection.id, grantId: row.grant.id, - provider: row.summary.provider, name: row.connection.name, ownerUserId: row.grant.subjectUserId, credential }; + provider: row.summary.provider, name: row.connection.name, ownerUserId: row.grant.subjectUserId, credential, session }; const id = await subscriptionService(db).register(input); // Accepted discovery must not silently skip an account. At most twenty // discovery tasks can wait here; never create an unbounded provider queue. diff --git a/server/src/services/subscriptions.ts b/server/src/services/subscriptions.ts index e1d80430e02..45a3fc3b71a 100644 --- a/server/src/services/subscriptions.ts +++ b/server/src/services/subscriptions.ts @@ -1,5 +1,5 @@ import { and, eq, inArray, sql } from "drizzle-orm"; -import { aiSubscriptions, aiSubscriptionPrices, aiSubscriptionConnections, connectionGrants, toolConnections, type Db } from "@paperclipai/db"; +import { aiSubscriptions, aiSubscriptionPrices, aiSubscriptionConnections, companySecrets, connectionGrants, toolConnections, type Db } from "@paperclipai/db"; import { aiConnectionMetadataSchema, normalizeCents, subscriptionPlan, type AiProvider, type UpdateSubscriptionPrice } from "@paperclipai/shared"; import { conflict, forbidden, notFound } from "../errors.js"; import { logActivity, type ActivityPublication } from "./activity-log.js"; @@ -12,6 +12,16 @@ export interface SubscriptionActor { userId: string; canManage: boolean; readOnl export interface SubscriptionConnection { companyId: string; connectionId: string; grantId: string; provider: AiProvider; name: string; ownerUserId: string | null; credential: string; + session?: SubscriptionSession | null; +} +export interface SubscriptionSession { secretId: string; epoch: number } + +export async function subscriptionSession(db: Db, companyId: string, refs: { secretId: string; configPath: string }[]): Promise { + const ref = refs.find(r => r.configPath === "ai.credential"); + if (!ref) return null; + const [secret] = await db.select({ epoch: companySecrets.aiSessionEpoch }).from(companySecrets) + .where(and(eq(companySecrets.companyId, companyId), eq(companySecrets.id, ref.secretId))); + return secret ? { secretId: ref.secretId, epoch: secret.epoch } : null; } export function canonicalSubscriptionId(id: string, accounts: Pick[] | Map): string { @@ -88,10 +98,6 @@ async function mergeInTransaction(db: Db, companyId: string, source: Account, ta export function subscriptionService(db: Db) { async function register(input: SubscriptionConnection, observed?: SubscriptionIdentity) { const identity = { ...(observed ?? subscriptionCredentialIdentity(input.companyId, input.provider, input.credential)) }; - // Only a successful provider observation may match another grant's billing - // identity or add its editors. Provisional rows are grant-scoped, even when - // two callers supply the same unverified credential or forged ID claims. - if (!observed?.verified) identity.accountKey = subscriptionIdentityKey(input.companyId, input.provider, ["grant", input.grantId, identity.credentialKey]); if (!observed) { const [cached] = await db.select().from(aiSubscriptionConnections).where(and(eq(aiSubscriptionConnections.companyId, input.companyId), eq(aiSubscriptionConnections.grantId, input.grantId))); if (cached?.credentialKey === identity.credentialKey && cached.connectionId === input.connectionId) return cached.subscriptionId; @@ -107,6 +113,14 @@ export function subscriptionService(db: Db) { || !metadata.success || metadata.data.method !== "subscription" || metadata.data.provider !== input.provider) { throw notFound("Subscription connection not found"); } + if (!observed?.verified) { + const session = await subscriptionSession(tx, input.companyId, grant.grant.credentialSecretRefs); + if (session && (input.session?.secretId !== session.secretId || input.session.epoch !== session.epoch)) { + throw conflict("The AI credential changed. Retry this execution."); + } + identity.accountKey = subscriptionIdentityKey(input.companyId, input.provider, ["grant", input.grantId, + ...(session ? ["session", session.secretId, String(session.epoch)] : [identity.credentialKey])]); + } const [binding] = await tx.select().from(aiSubscriptionConnections).where(and(eq(aiSubscriptionConnections.companyId, input.companyId), eq(aiSubscriptionConnections.grantId, input.grantId))); const accounts = await tx.select().from(aiSubscriptions).where(eq(aiSubscriptions.companyId, input.companyId)); if (observed && binding && binding.credentialKey !== identity.credentialKey) return canonicalSubscriptionId(binding.subscriptionId, accounts);