diff --git a/doc/connections/AI-CONNECTIONS.md b/doc/connections/AI-CONNECTIONS.md index 4e2c0494c33..96289490ede 100644 --- a/doc/connections/AI-CONNECTIONS.md +++ b/doc/connections/AI-CONNECTIONS.md @@ -834,3 +834,17 @@ updates, account changes, duplicate identity resolution, ownership, monetary receipt replay, unknown prices, and failed provider observations. Provider payloads are fixtures; live provider access still depends on the deployment's credentials and scopes. + +### Login completion ownership + +The internal connection save service takes an internal auth-session row ID for +every login flow. The setup-token owner resolves its public session ID within +the company, user, and adapter scope before calling save. Local and device +logins already use internal IDs. Provider names do not select the ID column. + +A login flow can supply a completion callback that runs inside the connection +save transaction, after the session binding is written. Setup-token completion +uses its existing conditional transition to `stored`; a rejected transition +rolls back the credentials and connection binding. Already-bound sessions return +the existing connection without repeating the callback. Local and device login +services continue to own their final states. diff --git a/server/src/__tests__/ai-connections.test.ts b/server/src/__tests__/ai-connections.test.ts index 183052c5105..49b9cd7d81d 100644 --- a/server/src/__tests__/ai-connections.test.ts +++ b/server/src/__tests__/ai-connections.test.ts @@ -1,3 +1,4 @@ +import { completeSetupTokenAiConnection } from "../services/setup-token-ai-connection.js"; import { connectionIntentService } from "../services/connection-intents.js"; import { connectionIntentDeliveryService } from "../services/connection-intent-delivery.js"; import { issueRecoveryActionService } from "../services/issue-recovery-actions.js"; @@ -844,7 +845,7 @@ describe("managed AI connections", () => { lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock") .mockRejectedValueOnce(Object.assign(new Error("lock timeout"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE })) .mockRejectedValueOnce(Object.assign(new Error("lock timeout"), { code: WORKSPACE_RESTORE_LOCK_TIMEOUT_CODE })); - saving = service.save(companyId, loginOwner, { ...intent, agentIds: [] }, credential, sessionId); + saving = service.save(companyId, loginOwner, { ...intent, agentIds: [] }, credential, { id: sessionId }); void saving.catch(() => {}); await vi.waitFor(() => expect(lock).toHaveBeenCalledTimes(3)); expect((await db.select().from(adapterAuthSessions).where(eq(adapterAuthSessions.id, sessionId)))[0].connectionId).toBeNull(); @@ -852,7 +853,7 @@ describe("managed AI connections", () => { const [, saved] = await Promise.all([refreshing, saving]); const [row] = (await service.quotaAccounts(companyId, loginOwner)).filter(row => row.connection.id === saved.connectionId); expect(await service.credential(row)).toBe(credential); - expect(await service.save(companyId, loginOwner, { ...intent, agentIds: [] }, credential, sessionId)).toEqual(saved); + expect(await service.save(companyId, loginOwner, { ...intent, agentIds: [] }, credential, { id: sessionId })).toEqual(saved); if (kind === "new") expect(JSON.parse(await service.credential(quotaRow))).toMatchObject({ tokens: { access_token: "quota-fresh", refresh_token: "quota-replacement" } }); } finally { release.resolve(); await Promise.allSettled([refreshing, saving]); lock?.mockRestore(); request.mockRestore(); @@ -1667,11 +1668,57 @@ console.log(JSON.stringify({ type: "result", subtype: "success", result: "hello" const intent = { provider: "anthropic", method: "subscription", ownership: "personal", name: "Claude subscription", agentIds: [], allAgents: true } as const; const sessionId = randomUUID(); await db.insert(adapterAuthSessions).values({ companyId, environmentId: environment.id, adapterType: "claude_local", startedByUserId: "alice", publicSessionId: sessionId, status: "submitting", aiConnection: { ...intent, agentIds: [] }, expiresAt: new Date(Date.now() + 60000) }); - const first = await service.save(companyId, "alice", { ...intent, agentIds: [] }, "fixture-subscription", sessionId); - expect(await service.save(companyId, "alice", { ...intent, agentIds: [] }, "fixture-subscription", sessionId)).toEqual(first); + const first = await completeSetupTokenAiConnection(db, { scope: { companyId, ownerUserId: "alice", environmentId: environment.id, adapterType: "claude_local", aiConnection: { ...intent, agentIds: [] } }, token: "fixture-subscription", sessionId }); + expect(await completeSetupTokenAiConnection(db, { scope: { companyId, ownerUserId: "alice", environmentId: environment.id, adapterType: "claude_local", aiConnection: { ...intent, agentIds: [] } }, token: "fixture-subscription", sessionId })).toEqual(first); + expect((await db.select().from(adapterAuthSessions).where(eq(adapterAuthSessions.publicSessionId, sessionId)))[0]).toMatchObject({ status: "stored", connectionId: first.connectionId, connectionGrantId: first.grantId }); const cancelled = randomUUID(); await db.insert(adapterAuthSessions).values({ companyId, environmentId: environment.id, adapterType: "claude_local", startedByUserId: "alice", publicSessionId: cancelled, status: "cancelled", expiresAt: new Date(Date.now() + 60000) }); - await expect(service.save(companyId, "alice", { ...intent, agentIds: [] }, "fixture-never-save", cancelled)).rejects.toThrow("no longer active"); + await expect(completeSetupTokenAiConnection(db, { scope: { companyId, ownerUserId: "alice", environmentId: environment.id, adapterType: "claude_local", aiConnection: { ...intent, agentIds: [] } }, token: "fixture-never-save", sessionId: cancelled })).rejects.toThrow("no longer active"); + }); + it.each(["anthropic", "openai"] as const)("binds %s by internal session ID without choosing its completion state", async provider => { + const owner = `session-owner-${provider}`; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const [environment] = await db.insert(environments).values({ name: `Session identity ${provider}`, driver: "sandbox" }).returning(); + const intent = { provider, method: "subscription", ownership: "personal", name: `Session identity ${provider}`, allAgents: false, agentIds: [] } as const; + const [session] = await db.insert(adapterAuthSessions).values({ companyId, environmentId: environment.id, + adapterType: provider === "anthropic" ? "claude_local" : "codex_local", startedByUserId: owner, + publicSessionId: randomUUID(), status: "promoting", aiConnection: { ...intent, agentIds: [] }, + expiresAt: new Date(Date.now() + 60_000), + }).returning(); + const complete = vi.fn(async () => {}); + try { + await expect(service.save(companyId, owner, { ...intent, agentIds: [] }, "fixture", { id: session.publicSessionId, complete })).rejects.toThrow("ownership mismatch"); + await expect(service.save(companyId, "bob", { ...intent, agentIds: [] }, "fixture", { id: session.id, complete })).rejects.toThrow("ownership mismatch"); + await expect(service.save(companyId, owner, { ...intent, agentIds: [], ownership: "shared" }, "fixture", { id: session.id, complete })).rejects.toThrow("target mismatch"); + expect(complete).not.toHaveBeenCalled(); + const saved = await service.save(companyId, owner, { ...intent, agentIds: [] }, "fixture", { id: session.id, complete }); + expect((await db.select().from(adapterAuthSessions).where(eq(adapterAuthSessions.id, session.id)))[0]).toMatchObject({ + connectionId: saved.connectionId, connectionGrantId: saved.grantId, status: "promoting", + }); + expect(await service.save(companyId, owner, { ...intent, agentIds: [] }, "replacement", { id: session.id, complete })).toEqual(saved); + expect(complete).toHaveBeenCalledTimes(1); + } finally { await db.delete(adapterAuthSessions).where(eq(adapterAuthSessions.id, session.id)); } + }); + it("rolls back credential and connection writes if setup-token completion loses its claim", async () => { + const owner = "lost-setup-token-claim"; + await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" }); + const [environment] = await db.insert(environments).values({ name: "Lost setup claim", driver: "sandbox" }).returning(); + const intent = { provider: "anthropic", method: "subscription", ownership: "personal", name: "Lost setup claim", allAgents: false, agentIds: [] } as const; + const [session] = await db.insert(adapterAuthSessions).values({ companyId, environmentId: environment.id, + adapterType: "claude_local", startedByUserId: owner, publicSessionId: randomUUID(), status: "submitting", + aiConnection: { ...intent, agentIds: [] }, expiresAt: new Date(Date.now() + 60_000), boundAt: new Date(), + }).returning(); + try { + await expect(completeSetupTokenAiConnection(db, { scope: { companyId, ownerUserId: owner, + environmentId: environment.id, adapterType: "claude_local", aiConnection: { ...intent, agentIds: [] } }, + sessionId: session.publicSessionId, token: "must-rollback", + })).rejects.toThrow(); + expect((await db.select().from(adapterAuthSessions).where(eq(adapterAuthSessions.id, session.id)))[0]).toMatchObject({ + connectionId: null, connectionGrantId: null, status: "submitting", + }); + expect(await db.select().from(toolConnections).where(and(eq(toolConnections.companyId, companyId), eq(toolConnections.name, intent.name)))).toEqual([]); + expect(await db.select().from(companySecrets).where(eq(companySecrets.ownerUserId, owner))).toEqual([]); + } finally { await db.delete(adapterAuthSessions).where(eq(adapterAuthSessions.id, session.id)); } }); it("preserves connection identity and defaults through reconnect; revocation wins over older attempts", async () => { const current = await service.select({ ...input, userId: "bob" }); @@ -1919,7 +1966,7 @@ console.log(JSON.stringify({ type: "result", subtype: "success", result: "hello" const sessionId = randomUUID(); const intent = { provider: "anthropic", method: "subscription", ownership: "personal", name: "Expired", agentIds: [], allAgents: true } as const; await db.insert(adapterAuthSessions).values({ companyId, environmentId: environment.id, adapterType: "claude_local", startedByUserId: "alice", publicSessionId: sessionId, status: "submitting", aiConnection: { ...intent, agentIds: [] }, expiresAt: new Date(Date.now() - 1000) }); - await expect(service.save(companyId, "alice", { ...intent, agentIds: [] }, "fixture-never-save", sessionId)).rejects.toThrow("no longer active"); + await expect(completeSetupTokenAiConnection(db, { scope: { companyId, ownerUserId: "alice", environmentId: environment.id, adapterType: "claude_local", aiConnection: { ...intent, agentIds: [] } }, token: "fixture-never-save", sessionId })).rejects.toThrow("no longer active"); expect((await service.list(companyId, "alice")).some(account => account.name === "Expired")).toBe(false); }); it("authorizes account creation, reconnect and defaults at the HTTP boundary before provider calls", async () => { @@ -2040,6 +2087,7 @@ console.log(JSON.stringify({ type: "result", subtype: "success", result: "hello" expect(started.headers["cache-control"]).toBe("no-store"); expect(started.body.command).toBeUndefined(); expect((await request(app).post(`${base}/attempts`).send(intent)).body).toEqual(started.body); + await db.update(adapterAuthSessions).set({ publicSessionId: randomUUID() }).where(eq(adapterAuthSessions.id, started.body.sessionId)); const input = { ...intent, localSessionId: started.body.sessionId }; for (const endpoint of [base, `${base}/check`]) { expect((await request(app).post(endpoint).set("x-test-user", "bob").send(input)).status).toBe(404); diff --git a/server/src/routes/agents.ts b/server/src/routes/agents.ts index fefd2afed44..a89ed24a218 100644 --- a/server/src/routes/agents.ts +++ b/server/src/routes/agents.ts @@ -1,3 +1,4 @@ +import { completeSetupTokenAiConnection } from "../services/setup-token-ai-connection.js"; import { cancellationRequestId } from "../services/native-runtime/native-cancellation-request.js"; import { agentHarnessType, type AgentRunnerChoice } from "@paperclipai/shared"; import { resolveNewAgentRunnerForCompany } from "../services/agent-runner-selection.js"; @@ -743,7 +744,7 @@ export function agentRoutes( store: setupTokenCleanupStore, completeCredential: async (input) => { if (!input.scope.aiConnection) return setupTokenSecretWriter(input); - await aiConnectionService(db).save(input.scope.companyId, input.scope.ownerUserId, input.scope.aiConnection, input.token, input.sessionId); + await completeSetupTokenAiConnection(db, input); }, rateLimiter: setupTokenRateLimiter, }); @@ -880,7 +881,7 @@ export function agentRoutes( if (managedSession?.aiConnection) { await adapterLoginStore.withCompanyAdapterPromotionLock(context.companyId, context.startedByUserId, context.adapterType, async () => { if (!(await checkStagedCredentialReadiness(authBytes)).ready) throw new Error("Provider credential is not ready"); - await aiConnectionService(db).save(context.companyId, context.startedByUserId, managedSession.aiConnection!, authBytes.toString("utf8"), context.sessionId); + await aiConnectionService(db).save(context.companyId, context.startedByUserId, managedSession.aiConnection!, authBytes.toString("utf8"), { id: context.sessionId }); }); return; } @@ -1125,7 +1126,7 @@ export function agentRoutes( if (managedSession?.aiConnection) { await adapterLoginStore.withCompanyAdapterPromotionLock(context.companyId, context.startedByUserId, context.adapterType, async () => { if (!(await checkStagedGrokCredentialReadiness(authBytes)).ready) throw new Error("Provider credential is not ready"); - await aiConnectionService(db).save(context.companyId, context.startedByUserId, managedSession.aiConnection!, authBytes.toString("utf8"), context.sessionId); + await aiConnectionService(db).save(context.companyId, context.startedByUserId, managedSession.aiConnection!, authBytes.toString("utf8"), { id: context.sessionId }); }); return; } diff --git a/server/src/services/ai-connections.ts b/server/src/services/ai-connections.ts index a12fbd744b6..0177f6352ac 100644 --- a/server/src/services/ai-connections.ts +++ b/server/src/services/ai-connections.ts @@ -569,7 +569,7 @@ export function aiConnectionService(db: Db) { userId: string, input: CreateAiConnection | AiConnectionLoginIntent, verifiedCredential: string, - sessionId?: string, + loginSession?: { id: string; complete?: (tx: Db) => Promise }, attemptStartedAt = new Date(), ) { const routing = "routing" in input ? input.routing : undefined; @@ -610,15 +610,13 @@ export function aiConnectionService(db: Db) { const grantId = reconnect?.grant.id ?? randomUUID(); const persist = () => withAccountHomeSecretMutationLock(undefined, companyId, () => db.transaction(async (tx) => { const secrets = secretService(tx); - if (sessionId) { + if (loginSession) { const [session] = await tx .select() .from(adapterAuthSessions) .where( and( - input.provider === "anthropic" - ? eq(adapterAuthSessions.publicSessionId, sessionId) - : eq(adapterAuthSessions.id, sessionId), + eq(adapterAuthSessions.id, loginSession.id), eq(adapterAuthSessions.companyId, companyId), eq(adapterAuthSessions.startedByUserId, userId), ), @@ -888,22 +886,18 @@ export function aiConnectionService(db: Db) { })), ); } - if (sessionId) + if (loginSession) await tx .update(adapterAuthSessions) .set({ connectionId: id, connectionGrantId: grantId, connectionMethod: input.method, - ...(input.provider === "anthropic" - ? { status: "stored" as const } - : {}), }) .where( - input.provider === "anthropic" - ? eq(adapterAuthSessions.publicSessionId, sessionId) - : eq(adapterAuthSessions.id, sessionId), + eq(adapterAuthSessions.id, loginSession.id), ); + await loginSession?.complete?.(tx as unknown as Db); await logActivity(tx as unknown as Db, { companyId, actorType: "user", diff --git a/server/src/services/local-ai-login.ts b/server/src/services/local-ai-login.ts index 2108a916baf..3da1fa60ec7 100644 --- a/server/src/services/local-ai-login.ts +++ b/server/src/services/local-ai-login.ts @@ -222,7 +222,7 @@ export function localAiLoginService(db: Db) { // Nested transaction is a savepoint on this same connection. Holding the // attempt lock makes completion/cancellation/restart retries idempotent. const saved = await aiConnectionService(tx as unknown as Db) - .save(companyId, userId, intent, credential, id, session.createdAt); + .save(companyId, userId, intent, credential, { id }, session.createdAt); await tx.update(adapterAuthSessions).set({ status: "authenticated", connectionMethod: LOCAL_LOGIN_METHOD, finishedAt: new Date(), updatedAt: new Date(), diff --git a/server/src/services/setup-token-ai-connection.ts b/server/src/services/setup-token-ai-connection.ts new file mode 100644 index 00000000000..4b3547816fa --- /dev/null +++ b/server/src/services/setup-token-ai-connection.ts @@ -0,0 +1,34 @@ +import { and, eq } from "drizzle-orm"; +import { adapterAuthSessions, type Db } from "@paperclipai/db"; +import type { AgentAdapterType } from "@paperclipai/shared"; +import { forbidden } from "../errors.js"; +import { aiConnectionService } from "./ai-connections.js"; +import { + SETUP_TOKEN_STORAGE_FAILED, + SetupTokenSessionError, + transitionSetupTokenSessionToStored, + type SetupTokenSecretWriter, +} from "./setup-token-session.js"; + +export async function completeSetupTokenAiConnection( + db: Db, + { scope, sessionId, token }: Parameters[0], +) { + if (!scope.aiConnection) throw new SetupTokenSessionError(500, SETUP_TOKEN_STORAGE_FAILED); + const [session] = await db.select({ id: adapterAuthSessions.id }).from(adapterAuthSessions).where(and( + eq(adapterAuthSessions.publicSessionId, sessionId), + eq(adapterAuthSessions.companyId, scope.companyId), + eq(adapterAuthSessions.startedByUserId, scope.ownerUserId), + eq(adapterAuthSessions.adapterType, scope.adapterType as AgentAdapterType), + )); + if (!session) throw forbidden("Login session ownership mismatch"); + return aiConnectionService(db).save(scope.companyId, scope.ownerUserId, scope.aiConnection, token, { + id: session.id, + complete: async (tx) => { + const transitioned = await transitionSetupTokenSessionToStored(tx, { + sessionId, companyId: scope.companyId, ownerUserId: scope.ownerUserId, adapterType: scope.adapterType, + }); + if (!transitioned) throw new SetupTokenSessionError(500, SETUP_TOKEN_STORAGE_FAILED); + }, + }); +}