diff --git a/packages/shared/src/validators/tool-access.test.ts b/packages/shared/src/validators/tool-access.test.ts index 3ed4ed968aa..1ccae506753 100644 --- a/packages/shared/src/validators/tool-access.test.ts +++ b/packages/shared/src/validators/tool-access.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it } from "vitest"; import { + connectToolAppSchema, connectionTokenRequestSchema, createToolConnectionSchema, startConnectionAuthorizationSchema, @@ -66,6 +67,19 @@ describe("tool access validators", () => { expect(parsed.success).toBe(true); }); + it("accepts a stable connection id when saving an edited app draft", () => { + const connectionId = "33333333-3333-4333-8333-333333333333"; + expect(connectToolAppSchema.parse({ + link: "https://mcp.example.test/project/edited", + name: "Edited project", + connectionId, + })).toMatchObject({ connectionId }); + expect(connectToolAppSchema.safeParse({ + link: "https://mcp.example.test/project/edited", + connectionId: "not-a-uuid", + }).success).toBe(false); + }); + it("keeps invocation payload summaries redacted and bounded", () => { const parsed = toolRedactedValueSummarySchema.parse({ summary: "Redacted arguments: 2 fields omitted.", diff --git a/packages/shared/src/validators/tool-access.ts b/packages/shared/src/validators/tool-access.ts index 01071fff7c9..30a97c490ef 100644 --- a/packages/shared/src/validators/tool-access.ts +++ b/packages/shared/src/validators/tool-access.ts @@ -265,6 +265,7 @@ export const connectToolAppSchema = z.object({ credentialValues: z.record(z.string().trim().min(1).max(200), z.string().min(1)).optional(), configValues: z.record(z.string().trim().min(1).max(200), z.unknown()).optional(), applicationId: z.string().uuid().optional(), + connectionId: z.string().uuid().optional(), }).refine( (value) => Boolean(value.galleryKey) !== Boolean(value.link), { message: "Provide exactly one of galleryKey or link" }, diff --git a/server/src/__tests__/tool-access-service.test.ts b/server/src/__tests__/tool-access-service.test.ts index 291cec311a7..b8b8ae03776 100644 --- a/server/src/__tests__/tool-access-service.test.ts +++ b/server/src/__tests__/tool-access-service.test.ts @@ -3076,7 +3076,12 @@ describeEmbeddedPostgres("tool access service", () => { credentialSecretRefs: [], config: expect.objectContaining({ sourceTemplateKey: "slack" }), }); - const startUrl = new URL(connectRes.body.auth.startUrl); + expect(connectRes.body.auth).toEqual({ kind: "oauth", startUrl: null }); + const startRes = await request(app) + .post(`/api/tools/oauth/${connectRes.body.connectionId}/start`) + .send({}) + .expect(200); + const startUrl = new URL(startRes.body.authorizationUrl); expect(`${startUrl.origin}${startUrl.pathname}`).toBe("https://slack.com/oauth/v2/authorize"); expect(startUrl.searchParams.get("client_id")).toBe("slack-client-id"); expect(startUrl.searchParams.get("code_challenge_method")).toBe("S256"); @@ -3158,7 +3163,12 @@ describeEmbeddedPostgres("tool access service", () => { .post(`/api/companies/${company.id}/tools/apps/connect`) .send({ galleryKey: "slack", name: "Slack redirect" }) .expect(201); - const redirectState = new URL(redirectConnectRes.body.auth.startUrl).searchParams.get("state"); + expect(redirectConnectRes.body.auth).toEqual({ kind: "oauth", startUrl: null }); + const redirectStartRes = await request(app) + .post(`/api/tools/oauth/${redirectConnectRes.body.connectionId}/start`) + .send({}) + .expect(200); + const redirectState = new URL(redirectStartRes.body.authorizationUrl).searchParams.get("state"); expect(redirectState).toBeTruthy(); const redirectCallbackRes = await request(app) .get("/api/tools/oauth/callback") @@ -4538,6 +4548,175 @@ describeEmbeddedPostgres("tool access service", () => { expect(clearedRes.body.apps).toEqual([]); }); + it("persists an OAuth draft before startup and updates that draft in place", async () => { + const company = await createCompany(db); + const app = createRouteApp(db); + + const first = await request(app) + .post(`/api/companies/${company.id}/tools/apps/connect`) + .send({ galleryKey: "notion", name: "PostHog project" }) + .expect(201); + + expect(first.body.auth).toEqual({ kind: "oauth", startUrl: null }); + vi.spyOn(globalThis, "fetch").mockRejectedValue(new Error("provider unavailable")); + const start = await request(app) + .post(`/api/tools/oauth/${first.body.connectionId}/start`) + .send({}); + expect(start.status).toBeGreaterThanOrEqual(400); + + const second = await request(app) + .post(`/api/companies/${company.id}/tools/apps/connect`) + .send({ + galleryKey: "notion", + name: "PostHog production project", + connectionId: first.body.connectionId, + }) + .expect(200); + + expect(second.body.connectionId).toBe(first.body.connectionId); + expect(second.body.application.id).toBe(first.body.application.id); + expect(second.body.application.name).toBe("PostHog production project"); + expect(second.body.connection.name).toBe("PostHog production project"); + await expect(db.select().from(toolApplications)).resolves.toHaveLength(1); + await expect(db.select().from(toolConnections)).resolves.toHaveLength(1); + + const activities = await db + .select() + .from(activityLog) + .where(eq(activityLog.entityId, first.body.connectionId)); + expect(activities.map((row) => row.details)).toEqual(expect.arrayContaining([ + expect.objectContaining({ operation: "created" }), + expect.objectContaining({ operation: "updated" }), + ])); + + const collision = await request(app) + .post(`/api/companies/${company.id}/tools/apps/connect`) + .send({ galleryKey: "slack", name: "Reserved app name" }) + .expect(201); + expect(collision.body.application.id).not.toBe(first.body.application.id); + + await request(app) + .post(`/api/companies/${company.id}/tools/apps/connect`) + .send({ + galleryKey: "notion", + name: "Reserved app name", + connectionId: first.body.connectionId, + }) + .expect(409); + const [unchanged] = await db + .select() + .from(toolApplications) + .where(eq(toolApplications.id, first.body.application.id)); + expect(unchanged.name).toBe("PostHog production project"); + }); + + it("fails closed when a draft update crosses company, source, application, or status boundaries", async () => { + const company = await createCompany(db); + const otherCompany = await createCompany(db); + const service = toolAccessService(db); + const draft = await service.connectGalleryApp(company.id, { + galleryKey: "slack", + name: "Guarded draft", + }); + + await expect(service.connectGalleryApp(otherCompany.id, { + galleryKey: "slack", + connectionId: draft.connectionId, + })).rejects.toMatchObject({ status: 404 }); + await expect(service.connectGalleryApp(company.id, { + galleryKey: "notion", + connectionId: draft.connectionId, + })).rejects.toMatchObject({ + status: 409, + details: { code: "tool_connection_draft_source_mismatch" }, + }); + await expect(service.connectGalleryApp(company.id, { + galleryKey: "slack", + applicationId: randomUUID(), + connectionId: draft.connectionId, + })).rejects.toMatchObject({ + status: 409, + details: { code: "tool_connection_draft_application_mismatch" }, + }); + + await db.update(toolConnections) + .set({ status: "active", enabled: true }) + .where(eq(toolConnections.id, draft.connectionId)); + await expect(service.connectGalleryApp(company.id, { + galleryKey: "slack", + connectionId: draft.connectionId, + })).rejects.toMatchObject({ + status: 409, + details: { code: "tool_connection_draft_not_editable" }, + }); + }); + + it("replaces draft credentials after validation and restores the prior draft when validation fails", async () => { + const company = await createCompany(db); + const service = toolAccessService(db); + const fetchMock = mockToolsList([{ name: "read_project", annotations: { readOnlyHint: true } }]); + + const first = await service.connectGalleryApp(company.id, { + link: "https://posthog.example.test/project/one", + name: "PostHog project", + credentialValues: { "credentials.authorization": "old-secret" }, + }, { actorType: "user", actorId: "board" }); + const oldSecretId = first.connection.credentialSecretRefs[0]!.secretId; + + const updated = await service.connectGalleryApp(company.id, { + link: "https://posthog.example.test/project/two", + name: "PostHog production", + credentialValues: { "credentials.authorization": "new-secret" }, + connectionId: first.connectionId, + }, { actorType: "user", actorId: "board" }); + const retainedSecretId = updated.connection.credentialSecretRefs[0]!.secretId; + + expect(updated.connectionId).toBe(first.connectionId); + expect(updated.application.id).toBe(first.application.id); + expect(updated.application.name).toBe("PostHog production"); + expect(updated.connection.config.url).toBe("https://posthog.example.test/project/two"); + expect(retainedSecretId).not.toBe(oldSecretId); + await expect(db.select().from(companySecrets)).resolves.toEqual([ + expect.objectContaining({ id: retainedSecretId }), + ]); + const updatedBindings = await db + .select() + .from(companySecretBindings) + .where(eq(companySecretBindings.targetId, first.connectionId)); + expect(new Set(updatedBindings.map((binding) => binding.secretId))).toEqual(new Set([retainedSecretId])); + + fetchMock.mockRejectedValue(new Error("validation unavailable")); + await expect(service.connectGalleryApp(company.id, { + link: "https://posthog.example.test/project/broken", + name: "Should roll back", + credentialValues: { "credentials.authorization": "do-not-keep" }, + connectionId: first.connectionId, + }, { actorType: "user", actorId: "board" })).rejects.toMatchObject({ status: 502 }); + + const [applicationAfterRollback] = await db + .select() + .from(toolApplications) + .where(eq(toolApplications.id, first.application.id)); + const [connectionAfterRollback] = await db + .select() + .from(toolConnections) + .where(eq(toolConnections.id, first.connectionId)); + expect(applicationAfterRollback.name).toBe("PostHog production"); + expect(connectionAfterRollback).toMatchObject({ + name: "PostHog production", + config: { url: "https://posthog.example.test/project/two", quarantineNewEntries: true }, + credentialSecretRefs: [expect.objectContaining({ secretId: retainedSecretId })], + }); + await expect(db.select().from(companySecrets)).resolves.toEqual([ + expect.objectContaining({ id: retainedSecretId }), + ]); + const restoredBindings = await db + .select() + .from(companySecretBindings) + .where(eq(companySecretBindings.targetId, first.connectionId)); + expect(new Set(restoredBindings.map((binding) => binding.secretId))).toEqual(new Set([retainedSecretId])); + }); + it("rolls back app connect drafts when health check fails", async () => { const company = await createCompany(db); const service = toolAccessService(db); @@ -4827,8 +5006,12 @@ describeEmbeddedPostgres("tool access service", () => { .send({ link: "https://generic.example.test/mcp", name: "Generic OAuth MCP" }); expect(connectRes.status).toBe(201); - expect(connectRes.body.auth).toMatchObject({ kind: "oauth" }); - const startUrl = new URL(connectRes.body.auth.startUrl); + expect(connectRes.body.auth).toEqual({ kind: "oauth", startUrl: null }); + const startRes = await request(app) + .post(`/api/tools/oauth/${connectRes.body.connectionId}/start`) + .send({}) + .expect(200); + const startUrl = new URL(startRes.body.authorizationUrl); expect(`${startUrl.origin}${startUrl.pathname}`).toBe("https://generic.example.test/oauth/authorize"); expect(startUrl.searchParams.get("client_id")).toBe("generic-client-id"); expect(startUrl.searchParams.get("scope")).toBe("tools.read tools.write"); diff --git a/server/src/routes/tool-access.ts b/server/src/routes/tool-access.ts index 0c4bef458e2..f14ebc2cb88 100644 --- a/server/src/routes/tool-access.ts +++ b/server/src/routes/tool-access.ts @@ -257,13 +257,7 @@ export function toolAccessRoutes( assertToolAppMutationAccess(req, companyId); try { const result = await svc.connectGalleryApp(companyId, req.body, getActorInfo(req)); - if (result.auth?.kind === "oauth") { - const start = await svc.startOAuth(companyId, result.connectionId, { - redirectUri: oauthRedirectUri(), - actor: getActorInfo(req), - }); - result.auth.startUrl = start.authorizationUrl; - } + const operation = req.body.connectionId ? "updated" : "created"; await logActivity(db, { companyId, actorType: "user", @@ -275,12 +269,13 @@ export function toolAccessRoutes( galleryKey: req.body.galleryKey ?? null, link: req.body.link ?? null, applicationId: result.application.id, + operation, catalogEntryCount: result.catalog.length, readOnlyActionCount: result.actions.readOnly.length, canMakeChangesActionCount: result.actions.canMakeChanges.length, }, }); - res.status(201).json(result); + res.status(operation === "updated" ? 200 : 201).json(result); } catch (error) { svc.ensureNoDuplicateNameError(error); } diff --git a/server/src/services/tool-access.ts b/server/src/services/tool-access.ts index 90e0bab1449..cdd58902e37 100644 --- a/server/src/services/tool-access.ts +++ b/server/src/services/tool-access.ts @@ -4827,9 +4827,57 @@ export function toolAccessService(db: Db, options: ToolAccessServiceOptions = {} ): Promise { const galleryEntry = input.galleryKey ? getConnectableAppDefinition(input.galleryKey) : null; if (input.galleryKey && !galleryEntry) throw notFound("Tool app gallery entry not found"); + const method = galleryEntry ? connectionMethodFor(galleryEntry) : null; let existingApplication: typeof toolApplications.$inferSelect | null = null; - if (input.applicationId) { + let editableConnectionPrevious: typeof toolConnections.$inferSelect | null = null; + if (input.connectionId) { + const [connection] = await db.select().from(toolConnections).where(and( + eq(toolConnections.id, input.connectionId), + eq(toolConnections.companyId, companyId), + )); + if (!connection) throw notFound("App connection not found"); + const [application] = await db.select().from(toolApplications).where(and( + eq(toolApplications.id, connection.applicationId), + eq(toolApplications.companyId, companyId), + )); + if (!application) throw notFound("App not found"); + if (connection.status !== "draft" || connection.enabled || application.status === "archived") { + throw conflict("Only editable app connection drafts can be updated", { + code: "tool_connection_draft_not_editable", + }); + } + if (input.applicationId && input.applicationId !== application.id) { + throw conflict("The app connection draft does not belong to that app", { + code: "tool_connection_draft_application_mismatch", + }); + } + + const connectionSource = typeof connection.config.sourceTemplateKey === "string" + ? connection.config.sourceTemplateKey + : null; + const applicationMetadata = asRecord(application.metadata); + const applicationSource = typeof applicationMetadata.sourceTemplateKey === "string" + ? applicationMetadata.sourceTemplateKey + : typeof applicationMetadata.galleryKey === "string" + ? applicationMetadata.galleryKey + : null; + const sourceMatches = galleryEntry + ? connectionSource === galleryEntry.slug + && applicationSource === galleryEntry.slug + && connection.transport === method?.transport + && connection.authKind === method?.auth + : applicationMetadata.source === "link" + && connectionSource === null + && connection.transport === "mcp_remote"; + if (!sourceMatches) { + throw conflict("The app connection draft was created from a different setup source", { + code: "tool_connection_draft_source_mismatch", + }); + } + existingApplication = application; + editableConnectionPrevious = connection; + } else if (input.applicationId) { const [row] = await db.select().from(toolApplications).where(and( eq(toolApplications.id, input.applicationId), eq(toolApplications.companyId, companyId), @@ -4839,7 +4887,6 @@ export function toolAccessService(db: Db, options: ToolAccessServiceOptions = {} } const name = input.name ?? existingApplication?.name ?? galleryEntry?.name ?? defaultLinkName(input.link ?? ""); - const method = galleryEntry ? connectionMethodFor(galleryEntry) : null; const transport = method?.transport ?? "mcp_remote"; const baseConfig = transport === "mcp_remote" ? { url: method?.defaults?.serverUrl ?? input.link ?? "" } @@ -4875,6 +4922,7 @@ export function toolAccessService(db: Db, options: ToolAccessServiceOptions = {} let applicationRow: typeof toolApplications.$inferSelect | null = null; let connectionRow: typeof toolConnections.$inferSelect | null = null; let revivedConnectionPrevious: typeof toolConnections.$inferSelect | null = null; + let result: ConnectToolAppResult | null = null; try { const credentialFields = galleryEntry ? credentialFieldsFor(galleryEntry) : linkCredentialFields(credentialValues); @@ -4911,7 +4959,29 @@ export function toolAccessService(db: Db, options: ToolAccessServiceOptions = {} } } - if (existingApplication) { + if (editableConnectionPrevious && existingApplication?.status === "draft") { + [applicationRow] = await db.update(toolApplications) + .set({ + name, + description: galleryEntry?.description ?? `Connected app at ${input.link}`, + type: transport === "mcp_remote" ? "mcp_http" : "mcp_stdio", + metadata: galleryEntry + ? { sourceTemplateKey: galleryEntry.slug, galleryKey: galleryEntry.slug } + : { source: "link" }, + updatedAt: new Date(), + }) + .where(and( + eq(toolApplications.id, existingApplication.id), + eq(toolApplications.companyId, companyId), + eq(toolApplications.status, "draft"), + )) + .returning(); + if (!applicationRow) { + throw conflict("The app connection draft changed while it was being updated", { + code: "tool_connection_draft_changed", + }); + } + } else if (existingApplication) { if (existingApplication.status !== "active") { [applicationRow] = await db.update(toolApplications) .set({ status: "draft", archivedAt: null, updatedAt: new Date() }) @@ -4936,7 +5006,7 @@ export function toolAccessService(db: Db, options: ToolAccessServiceOptions = {} // Reconnecting an app revives its most recent archived connection instead // of inserting a fresh row: keeps the connection id (and its activity // history) stable and avoids the unique (company, name) constraint. - if (existingApplication) { + if (existingApplication && !editableConnectionPrevious) { const [archived] = await db .select() .from(toolConnections) @@ -4949,7 +5019,36 @@ export function toolAccessService(db: Db, options: ToolAccessServiceOptions = {} .limit(1); revivedConnectionPrevious = archived ?? null; } - if (revivedConnectionPrevious) { + if (editableConnectionPrevious) { + [connectionRow] = await db.update(toolConnections).set({ + name, + authKind: galleryEntry ? method!.auth : editableConnectionPrevious.authKind, + transport, + status: "draft", + enabled: false, + config, + transportConfig: config, + credentialRefs, + credentialSecretRefs, + healthStatus: "unchecked", + healthMessage: null, + healthCheckedAt: null, + lastHealthAt: null, + lastCatalogRefreshAt: null, + lastError: null, + updatedAt: new Date(), + }).where(and( + eq(toolConnections.id, editableConnectionPrevious.id), + eq(toolConnections.companyId, companyId), + eq(toolConnections.applicationId, applicationRow.id), + eq(toolConnections.status, "draft"), + )).returning(); + if (!connectionRow) { + throw conflict("The app connection draft changed while it was being updated", { + code: "tool_connection_draft_changed", + }); + } + } else if (revivedConnectionPrevious) { [connectionRow] = await db.update(toolConnections).set({ name, transport, @@ -4986,7 +5085,7 @@ export function toolAccessService(db: Db, options: ToolAccessServiceOptions = {} await ensureRuntimeSlot(connectionRow); if (galleryEntry && connectionMethodFor(galleryEntry).auth === "oauth") { - return { + result = { connectionId: connectionRow.id, application: toApplication(applicationRow), connection: toConnection(connectionRow), @@ -4995,16 +5094,23 @@ export function toolAccessService(db: Db, options: ToolAccessServiceOptions = {} suggestedDefaults: recommendedDefaultsForApp(galleryEntry), auth: { kind: "oauth", startUrl: null }, }; - } - - try { - await checkConnectionHealth(connectionRow.id, actor); - } catch (error) { - if (!galleryEntry && error instanceof HttpError && asRecord(error.details).code === "oauth_challenge") { - const [oauthConnection] = await db.select().from(toolConnections).where(eq(toolConnections.id, connectionRow.id)); - const endpoints = await discoverOAuthEndpoints(oauthConnection).catch(() => null); - if (!endpoints) throw error; - return { + } else { + let oauthConnection: typeof toolConnections.$inferSelect | null = null; + try { + await checkConnectionHealth(connectionRow.id, actor); + } catch (error) { + if (!galleryEntry && error instanceof HttpError && asRecord(error.details).code === "oauth_challenge") { + [oauthConnection] = await db.select().from(toolConnections).where(eq(toolConnections.id, connectionRow.id)); + const endpoints = oauthConnection + ? await discoverOAuthEndpoints(oauthConnection).catch(() => null) + : null; + if (!endpoints) throw error; + } else { + throw error; + } + } + if (oauthConnection) { + result = { connectionId: oauthConnection.id, application: toApplication(applicationRow), connection: toConnection(oauthConnection), @@ -5016,43 +5122,64 @@ export function toolAccessService(db: Db, options: ToolAccessServiceOptions = {} }, auth: { kind: "oauth", startUrl: null }, }; + } else { + const refresh = await refreshCatalog(connectionRow.id, actor); + const [application] = await db.select().from(toolApplications).where(eq(toolApplications.id, applicationRow.id)); + result = { + connectionId: refresh.connection.id, + application: toApplication(application), + connection: refresh.connection, + catalog: refresh.catalog, + actions: groupedActions(refresh.catalog), + suggestedDefaults: galleryEntry ? recommendedDefaultsForApp(galleryEntry) : { + access: "all_agents", + askFirstRiskLevels: ["write", "destructive"], + }, + }; } - throw error; } - const refresh = await refreshCatalog(connectionRow.id, actor); - const [application] = await db.select().from(toolApplications).where(eq(toolApplications.id, applicationRow.id)); - return { - connectionId: refresh.connection.id, - application: toApplication(application), - connection: refresh.connection, - catalog: refresh.catalog, - actions: groupedActions(refresh.catalog), - suggestedDefaults: galleryEntry ? recommendedDefaultsForApp(galleryEntry) : { - access: "all_agents", - askFirstRiskLevels: ["write", "destructive"], - }, - }; } catch (error) { - if (connectionRow && revivedConnectionPrevious) { - await db.update(toolConnections).set({ - name: revivedConnectionPrevious.name, - transport: revivedConnectionPrevious.transport, - status: revivedConnectionPrevious.status, - enabled: revivedConnectionPrevious.enabled, - config: revivedConnectionPrevious.config, - transportConfig: revivedConnectionPrevious.transportConfig, - credentialRefs: revivedConnectionPrevious.credentialRefs, - credentialSecretRefs: revivedConnectionPrevious.credentialSecretRefs, - updatedAt: new Date(), - }).where(eq(toolConnections.id, revivedConnectionPrevious.id)).catch(() => undefined); + const connectionPrevious = editableConnectionPrevious ?? revivedConnectionPrevious; + if (connectionRow && connectionPrevious) { + const [restored] = await db.update(toolConnections).set({ + name: connectionPrevious.name, + authKind: connectionPrevious.authKind, + transport: connectionPrevious.transport, + status: connectionPrevious.status, + enabled: connectionPrevious.enabled, + config: connectionPrevious.config, + transportConfig: connectionPrevious.transportConfig, + credentialRefs: connectionPrevious.credentialRefs, + credentialSecretRefs: connectionPrevious.credentialSecretRefs, + healthStatus: connectionPrevious.healthStatus, + healthMessage: connectionPrevious.healthMessage, + healthCheckedAt: connectionPrevious.healthCheckedAt, + lastHealthAt: connectionPrevious.lastHealthAt, + lastCatalogRefreshAt: connectionPrevious.lastCatalogRefreshAt, + lastError: connectionPrevious.lastError, + updatedAt: connectionPrevious.updatedAt, + }).where(eq(toolConnections.id, connectionPrevious.id)).returning().catch(() => []); + if (restored) await syncCredentialBindings(restored).catch(() => undefined); } else if (connectionRow) { await db.delete(toolConnections).where(eq(toolConnections.id, connectionRow.id)).catch(() => undefined); } if (applicationRow && !existingApplication) { await db.delete(toolApplications).where(eq(toolApplications.id, applicationRow.id)).catch(() => undefined); - } else if (existingApplication && applicationRow && applicationRow.status !== existingApplication.status) { + } else if (existingApplication && applicationRow && ( + applicationRow.name !== existingApplication.name + || applicationRow.status !== existingApplication.status + || applicationRow.description !== existingApplication.description + )) { await db.update(toolApplications) - .set({ status: existingApplication.status, archivedAt: existingApplication.archivedAt, updatedAt: new Date() }) + .set({ + name: existingApplication.name, + description: existingApplication.description, + type: existingApplication.type, + status: existingApplication.status, + metadata: existingApplication.metadata, + archivedAt: existingApplication.archivedAt, + updatedAt: existingApplication.updatedAt, + }) .where(eq(toolApplications.id, existingApplication.id)) .catch(() => undefined); } @@ -5061,6 +5188,20 @@ export function toolAccessService(db: Db, options: ToolAccessServiceOptions = {} } throw error; } + + if (!result) throw new Error("App connection setup completed without a result"); + if (editableConnectionPrevious) { + const retainedSecretIds = new Set([ + ...credentialRefs.map((ref) => ref.secretId), + ...credentialSecretRefs.map((ref) => ref.secretId), + ]); + const supersededSecretIds = new Set([ + ...editableConnectionPrevious.credentialRefs.map((ref) => ref.secretId), + ...editableConnectionPrevious.credentialSecretRefs.map((ref) => ref.secretId), + ].filter((secretId) => !retainedSecretIds.has(secretId))); + for (const secretId of supersededSecretIds) await secrets.remove(secretId); + } + return result; } async function assertCatalogEntriesForConnection( diff --git a/ui/src/api/tools.ts b/ui/src/api/tools.ts index 47692aec3dd..3ca06c11c26 100644 --- a/ui/src/api/tools.ts +++ b/ui/src/api/tools.ts @@ -241,6 +241,7 @@ export const toolsApi = { credentialValues?: Record; configValues?: Record; applicationId?: string; + connectionId?: string; }) => api.post(`/companies/${companyId}/tools/apps/connect`, input), startOAuth: (connectionId: string) => diff --git a/ui/src/pages/apps/AppsConnect.test.tsx b/ui/src/pages/apps/AppsConnect.test.tsx index b05e2c25bf6..0d711f5499d 100644 --- a/ui/src/pages/apps/AppsConnect.test.tsx +++ b/ui/src/pages/apps/AppsConnect.test.tsx @@ -482,6 +482,65 @@ describe("AppsConnect — Connect with a link (M4 frame)", () => { ); }); + it("saves edited link settings into the prepared OAuth draft before retrying sign-in", async () => { + const draftResult = { + connectionId: "conn-posthog", + application: { id: "app-posthog", name: "PostHog project" }, + connection: { id: "conn-posthog", name: "PostHog project" }, + actions: { readOnly: [], canMakeChanges: [] }, + catalog: [], + suggestedDefaults: {}, + auth: { kind: "oauth", startUrl: null }, + }; + connectAppMock + .mockResolvedValueOnce(draftResult) + .mockResolvedValueOnce({ + ...draftResult, + application: { id: "app-posthog", name: "PostHog production" }, + connection: { id: "conn-posthog", name: "PostHog production" }, + }); + startOAuthMock + .mockRejectedValueOnce(new Error("Provider unavailable")) + .mockResolvedValueOnce({ + connectionId: "conn-posthog", + provider: "posthog", + authorizationUrl: "https://posthog.example.test/oauth?state=retry", + expiresAt: "2099-01-01T00:00:00.000Z", + }); + + await render(); + await gotoLinkFrame(container, "https://mcp.posthog.example/project/one"); + await act(async () => { + buttonByText("Check link")?.dispatchEvent(new MouseEvent("click", { bubbles: true })); + }); + await flushReact(); + + expect(container.textContent).toContain("Provider unavailable"); + expect(buttonByText("Continue to sign in")).toBeTruthy(); + const nameInput = Array.from(container.querySelectorAll("input")).find( + (input) => input.getAttribute("placeholder") === "My app", + ); + await act(async () => setInputValue(nameInput!, "PostHog production")); + await act(async () => { + buttonByText("Continue to sign in")?.dispatchEvent(new MouseEvent("click", { bubbles: true })); + }); + await flushReact(); + + expect(connectAppMock).toHaveBeenCalledTimes(2); + expect(connectAppMock).toHaveBeenLastCalledWith("company-1", { + link: "https://mcp.posthog.example/project/one", + name: "PostHog production", + credentialValues: undefined, + applicationId: undefined, + connectionId: "conn-posthog", + }); + expect(startOAuthMock).toHaveBeenCalledTimes(2); + expect(startOAuthMock).toHaveBeenLastCalledWith("conn-posthog"); + expect(navigateTopLevelMock).toHaveBeenCalledWith( + "https://posthog.example.test/oauth?state=retry", + ); + }); + it("recovers a response-lost Notion draft before retrying creation", async () => { mockSearch.value = "source=notion"; listGalleryMock.mockResolvedValueOnce({ apps: [NOTION] }); diff --git a/ui/src/pages/apps/AppsConnect.tsx b/ui/src/pages/apps/AppsConnect.tsx index 0f2791cd53d..7f95e5b1819 100644 --- a/ui/src/pages/apps/AppsConnect.tsx +++ b/ui/src/pages/apps/AppsConnect.tsx @@ -1,5 +1,5 @@ import { useEffect, useMemo, useRef, useState } from "react"; -import { useMutation, useQuery } from "@tanstack/react-query"; +import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; import { ArrowUpRight, Check, @@ -134,6 +134,7 @@ function reusableOAuthConnection( export function AppsConnect() { const navigate = useNavigate(); + const queryClient = useQueryClient(); const routeParams = useParams<{ appKey?: string }>(); const { selectedCompany, selectedCompanyId } = useCompany(); const { setBreadcrumbs } = useBreadcrumbs(); @@ -278,6 +279,7 @@ export function AppsConnect() { credentialValues: credentials, configValues: isGoogleSheetsEntry(connectEntry) ? { allowedSpreadsheetIds: sheetIds } : undefined, applicationId: prefill.applicationId, + ...(connectResult?.connectionId ? { connectionId: connectResult.connectionId } : {}), }); } const trimmedKey = linkNeedsKey ? linkKey.trim() : ""; @@ -287,11 +289,17 @@ export function AppsConnect() { name: trimmedName || undefined, credentialValues: trimmedKey ? { [LINK_CREDENTIAL_CONFIG_PATH]: trimmedKey } : undefined, applicationId: prefill.applicationId, + ...(connectResult?.connectionId ? { connectionId: connectResult.connectionId } : {}), }); }, onSuccess: (result) => { + void Promise.all([ + queryClient.invalidateQueries({ queryKey: queryKeys.tools.applications(selectedCompanyId!) }), + queryClient.invalidateQueries({ queryKey: queryKeys.tools.connections(selectedCompanyId!) }), + ]); if (result.auth?.kind === "oauth") { setConnectResult(result); + setOAuthError(null); const startUrl = result.auth.startUrl?.trim(); if (!startUrl) { setOAuthPhase("starting"); @@ -630,9 +638,14 @@ export function AppsConnect() { }} keyValue={linkKey} onKeyChange={setLinkKey} - submitting={connectMutation.isPending} - onBack={() => setStep("gallery")} - onConnect={() => connectMutation.mutate(undefined)} + submitting={connectMutation.isPending || oauthStartMutation.isPending} + oauthError={oauthPhase === "error" ? oauthError : null} + hasOAuthDraft={connectResult?.auth?.kind === "oauth"} + onBack={openGallery} + onConnect={() => { + setOAuthError(null); + connectMutation.mutate(undefined); + }} /> )} @@ -1175,6 +1188,8 @@ function LinkConnectStep({ keyValue, onKeyChange, submitting, + oauthError, + hasOAuthDraft, onBack, onConnect, }: { @@ -1186,6 +1201,8 @@ function LinkConnectStep({ keyValue: string; onKeyChange: (next: string) => void; submitting: boolean; + oauthError: string | null; + hasOAuthDraft: boolean; onBack: () => void; onConnect: () => void; }) { @@ -1202,6 +1219,11 @@ function LinkConnectStep({
+ {oauthError && ( + + {oauthError} Edit these settings if needed, then continue to sign in. + + )}