refactor: let login flows own session completion (#17129)

## Thinking Path

> - Paperclip manages AI agents and their connected accounts.
> - Login flows bind verified credentials to an owned connection.
> - The shared save service currently uses the provider name to choose a
session ID column and completion state.
> - A provider can support more than one login flow. Provider identity
does not define session identity or state.
> - This PR gives every flow the same internal session ID contract and
lets the setup-token owner complete its own claim.
> - Credential writes, session binding, and completion remain in one
transaction.

## Linked Issues or Issue Description

Refs #17106. This is an independent prerequisite for that sleep cleanup
PR. It contains no scheduler or cleanup-worker changes.

**Current behavior**

The connection save service looks up Anthropic sessions by public ID and
all other sessions by internal ID. It also marks every Anthropic session
as stored. Local Anthropic login works because its public and internal
IDs happen to match, and its caller overwrites that intermediate state.

**Proposed behavior**

Every caller passes an internal session ID. A small setup-token
completion service resolves the public ID within its owner scope and
supplies a transaction callback. That callback uses the existing
conditional stored transition. Device and local login flows keep their
own final states.

**Reason and benefit**

Adding or changing a login flow no longer requires provider branches in
the shared session binding path. The service still owns membership
checks, row locking, target validation, credential writes, and
idempotent binding. The flow owns only public-ID resolution and its
completion transition.

**Breaking changes**

The internal TypeScript save signature changes. All production callers
are updated. Public HTTP payloads, database schema, and local session
URLs stay unchanged.

## What Changed

- Replace the ambiguous session ID argument with an internal session
reference and an optional completion callback.
- Run the callback inside the existing save transaction, only for the
first successful binding.
- Move public-ID resolution and the stored transition to
setup-token-ai-connection.ts.
- Update setup-token, device, and local callers.
- Add real PostgreSQL tests for ID selection, owner and target checks,
one-time completion, and rollback after a rejected completion.
- Exercise local Anthropic and OpenAI HTTP completion with distinct
public and internal IDs.
- Document the ownership contract in AI-CONNECTIONS.md.

## Verification

- Passed: pnpm -r typecheck.
- Passed: pnpm exec vitest run
server/src/__tests__/ai-connections.test.ts
server/src/__tests__/agent-device-login-routes.test.ts
server/src/__tests__/agents-claude-oauth-binding.test.ts — 233 tests.
- Passed: pnpm build.
- Passed: full CI, including test shards, typecheck, build, and E2E, on
d98ce6e504:
https://github.com/paperclipai/paperclip/actions/runs/39654210168.
- The additional local pnpm test:run was stopped after the full CI pass
to avoid duplicating the completed run. It is not reported as a
completed local full-suite pass.
- Greptile: 5/5 on d98ce6e504, with no
actionable findings or unresolved threads.
- The first AI suite run had two new fixture failures: a duplicate
environment name across provider cases, and an assertion that ignored
existing shared connections. The fixtures now use distinct names and
query the attempted connection directly. The corrected affected suites
passed.

The first CI run passed all 1,420 assertions in server shard 11/12 but
failed on an unhandled child-process rejection in
`workspace-runtime-exposure.test.ts` (the synthetic EADDRINUSE case).
Both the test and runtime implementation are unchanged from this PR's
base. The two affected parameter cases passed locally unchanged (2
passed, 29 skipped). Inspection shows a possible timing window before
the early-exit promise joins the readiness race; that cause has not been
reproduced conclusively. The single CI retry passed on the same commit.
This demonstrates recovery; the original timing cause was not
conclusively reproduced. Original job:
https://github.com/paperclipai/paperclip/actions/runs/39654210168/job/119410225203.

## Risks

- A completion callback failure must roll back the credential and
session binding. The real-database regression test proves this for a
rejected setup-token claim.
- The callback is internal, awaited inside the existing transaction, and
must use the supplied transaction handle. It adds no network work or
provider SDK dependency.
- Setup-token resolution adds one indexed lookup per completion. The
shared save then locks and validates the row again. No recurring work is
added.
- The setup-token owner now uses the existing conditional completion
guard, including the unconsumed-claim and commit-time expiry checks.
Invalid claims fail closed.
- Live provider sign-in was not exercised; tests use fixture credentials
and real PostgreSQL. No schema migration is needed.

## Model Used

OpenAI GPT-6 through Codex, with reasoning, source inspection, code
editing, and test execution. The exact serving model ID and
context-window size are not exposed in this session.

## Checklist

- [x] I have included a thinking path that traces from project context
to this change
- [x] I have specified the model used (with version and capability
details)
- [x] I have checked ROADMAP.md and confirmed this PR does not duplicate
planned core work
- [x] I have searched GitHub for duplicate or related PRs and linked
them above
- [x] I have either (a) linked existing issues with `Fixes: #` / `Closes
#` / `Refs #` OR (b) described the issue in-PR following the relevant
issue template
- [x] I have not referenced internal/instance-local Paperclip issues or
links (only public GitHub `#NNN` / `github.com/paperclipai/paperclip`
URLs)
- [x] My branch name describes the change (e.g. `docs/...`, `fix/...`)
and contains no internal Paperclip ticket id or instance-derived details
- [x] I have run tests locally and they pass
- [x] I have added or updated tests where applicable
- [x] I have updated relevant documentation to reflect my changes
- [x] I have considered and documented any risks above
- [x] All Paperclip CI gates are green
- [x] Greptile is 5/5 with no open P2s, recommendations, or follow-ups
- [x] I will address all Greptile and reviewer comments before
requesting merge

Co-authored-by: Paperclip <noreply@paperclip.ing>

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.

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);

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;
           }

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<void> },
     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",

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(),

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<SetupTokenSecretWriter>[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);
+    },
+  });
+}