diff --git a/apps/cloud/package.json b/apps/cloud/package.json index 0240c02cf..bf13d2819 100644 --- a/apps/cloud/package.json +++ b/apps/cloud/package.json @@ -69,6 +69,7 @@ "autumn-js": "^1.2.8", "drizzle-orm": "catalog:", "effect": "catalog:", + "iron-webcrypto": "2.0.0", "jose": "^5.6.3", "postgres": "^3.4.9", "posthog-js": "^1.372.5", diff --git a/apps/cloud/src/auth/workos.node.test.ts b/apps/cloud/src/auth/workos.node.test.ts index 82cbc572c..ca81b7717 100644 --- a/apps/cloud/src/auth/workos.node.test.ts +++ b/apps/cloud/src/auth/workos.node.test.ts @@ -1,6 +1,287 @@ +import { createServer, type IncomingMessage, type ServerResponse } from "node:http"; +import type { AddressInfo } from "node:net"; + import { describe, expect, it } from "@effect/vitest"; +import { env } from "cloudflare:workers"; +import { Effect, Schema } from "effect"; +import { defaults as ironDefaults, seal as sealIron } from "iron-webcrypto"; +import { + SignJWT, + exportJWK, + generateKeyPair, + type JSONWebKeySet, + type JWK, + type KeyLike, +} from "jose"; + +import { WorkOSClient, collectRawWorkOSList, collectWorkOSList } from "./workos"; + +const COOKIE_PASSWORD = "test_cookie_password_at_least_32_chars!"; +const CLIENT_ID = "client_test"; +const API_KEY = "sk_test"; +const USER = { + id: "user_test", + email: "test@example.com", + firstName: "Test", + lastName: "User", + profilePictureUrl: "https://example.com/avatar.png", +} as const; + +interface Keypair { + readonly kid: string; + readonly publicJwk: JWK; + readonly privateKey: KeyLike; +} + +interface RecordedRequest { + readonly method: string; + readonly path: string; + readonly body: unknown; +} + +interface WorkOSStub { + readonly baseUrl: string; + readonly requests: () => readonly RecordedRequest[]; + readonly setKeys: (keys: ReadonlyArray) => void; +} + +const generateKeypair = async (kid: string): Promise => { + const { publicKey, privateKey } = await generateKeyPair("RS256"); + const jwk = await exportJWK(publicKey); + return { kid, publicJwk: { ...jwk, kid, alg: "RS256" }, privateKey }; +}; + +const signAccessToken = ( + keypair: Keypair, + claims: { + readonly subject?: string; + readonly organizationId?: string; + readonly sessionId?: string; + readonly expiresIn?: string | number; + } = {}, +) => { + const jwt = new SignJWT({ + org_id: claims.organizationId ?? "org_test", + sid: claims.sessionId ?? "session_test", + }) + .setProtectedHeader({ alg: "RS256", kid: keypair.kid }) + .setSubject(claims.subject ?? USER.id) + .setIssuedAt(); + + return ( + typeof claims.expiresIn === "number" + ? jwt.setExpirationTime(claims.expiresIn) + : jwt.setExpirationTime(claims.expiresIn ?? "5m") + ).sign(keypair.privateKey); +}; + +const sealSession = async (accessToken: string, refreshToken = "refresh_test"): Promise => + `${await sealIron( + { + accessToken, + refreshToken, + user: USER, + }, + { id: "1", secret: COOKIE_PASSWORD }, + { ...ironDefaults, ttl: 0, encode: JSON.stringify }, + )}~2`; + +const workosUserResponse = { + object: "user", + id: USER.id, + email: USER.email, + email_verified: true, + profile_picture_url: USER.profilePictureUrl, + first_name: USER.firstName, + last_name: USER.lastName, + last_sign_in_at: null, + locale: null, + created_at: "2026-07-09T00:00:00.000Z", + updated_at: "2026-07-09T00:00:00.000Z", + external_id: null, + metadata: {}, +} as const; + +const decodeJsonBody = Schema.decodeUnknownSync(Schema.fromJsonString(Schema.Unknown)); + +const readJsonBody = async (request: IncomingMessage): Promise => { + const chunks: Buffer[] = []; + for await (const chunk of request) { + chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)); + } + const body = Buffer.concat(chunks).toString("utf8"); + return body ? decodeJsonBody(body) : null; +}; + +const writeJson = (response: ServerResponse, status: number, body: unknown): void => { + response.writeHead(status, { "content-type": "application/json" }); + response.end(JSON.stringify(body)); +}; + +const withWorkOSStub = async ( + keypair: Keypair, + use: (stub: WorkOSStub) => Promise, +): Promise => { + let keys: ReadonlyArray = [keypair.publicJwk]; + const requests: RecordedRequest[] = []; + const server = createServer((request, response) => { + // oxlint-disable-next-line executor/no-promise-catch -- boundary: Node's request callback cannot await, so the fixture converts async failures into a stable HTTP 500 + void (async () => { + const path = request.url ?? "/"; + const body = request.method === "POST" ? await readJsonBody(request) : null; + requests.push({ method: request.method ?? "GET", path, body }); + + if (request.method === "GET" && path === `/sso/jwks/${CLIENT_ID}`) { + writeJson(response, 200, { keys: keys.map((key) => ({ ...key })) } satisfies JSONWebKeySet); + return; + } + + if (request.method === "POST" && path === "/user_management/authenticate") { + const refreshed = await signAccessToken(keypair, { + sessionId: "session_refreshed", + }); + writeJson(response, 200, { + user: workosUserResponse, + organization_id: "org_test", + access_token: refreshed, + refresh_token: "refresh_next", + authentication_method: "Password", + }); + return; + } + + writeJson(response, 404, { error: "unexpected request" }); + })().catch((_error: unknown) => { + writeJson(response, 500, { + error: "fixture server failed", + }); + }); + }); + + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const address = server.address() as AddressInfo; + const stub: WorkOSStub = { + baseUrl: `http://127.0.0.1:${address.port}`, + requests: () => requests, + setKeys: (next) => { + keys = next; + }, + }; + + return use(stub).finally(() => { + server.closeAllConnections(); + server.close(); + }); +}; -import { collectRawWorkOSList, collectWorkOSList } from "./workos"; +const runAuthenticate = (sessionData: string, baseUrl: string) => { + Object.assign(env, { + WORKOS_API_KEY: API_KEY, + WORKOS_CLIENT_ID: CLIENT_ID, + WORKOS_COOKIE_PASSWORD: COOKIE_PASSWORD, + WORKOS_API_URL: baseUrl, + }); + + return Effect.runPromise( + Effect.gen(function* () { + const workos = yield* WorkOSClient; + return yield* workos.authenticateSealedSession(sessionData); + }).pipe(Effect.provide(WorkOSClient.Default)), + ); +}; + +describe("authenticateSealedSession", () => { + it("validates a sealed session locally with the cached JWKS", async () => { + const keypair = await generateKeypair("k_valid"); + await withWorkOSStub(keypair, async (stub) => { + const token = await signAccessToken(keypair, { + organizationId: "org_test", + sessionId: "session_valid", + }); + const session = await sealSession(token); + + const result = await runAuthenticate(session, stub.baseUrl); + + expect(result).toEqual({ + userId: USER.id, + email: USER.email, + firstName: USER.firstName, + lastName: USER.lastName, + avatarUrl: USER.profilePictureUrl, + organizationId: "org_test", + sessionId: "session_valid", + refreshedSession: undefined, + }); + expect(stub.requests()).toEqual([ + { method: "GET", path: `/sso/jwks/${CLIENT_ID}`, body: null }, + ]); + }); + }); + + it("reuses the module cached JWKS for a second sealed session", async () => { + const keypair = await generateKeypair("k_cached"); + await withWorkOSStub(keypair, async (stub) => { + const first = await sealSession( + await signAccessToken(keypair, { + sessionId: "session_first", + }), + ); + const second = await sealSession( + await signAccessToken(keypair, { + sessionId: "session_second", + }), + ); + + await runAuthenticate(first, stub.baseUrl); + const afterFirst = stub.requests().length; + const result = await runAuthenticate(second, stub.baseUrl); + + expect(afterFirst).toBe(1); + expect(result?.sessionId).toBe("session_second"); + expect(stub.requests()).toEqual([ + { method: "GET", path: `/sso/jwks/${CLIENT_ID}`, body: null }, + ]); + }); + }); + + it("falls through to SDK refresh when the access token is expired", async () => { + const keypair = await generateKeypair("k_expired"); + await withWorkOSStub(keypair, async (stub) => { + const expiredAt = Math.floor(Date.now() / 1000) - 60; + const expired = await signAccessToken(keypair, { + expiresIn: expiredAt, + sessionId: "session_expired", + }); + const session = await sealSession(expired, "refresh_expired"); + + const result = await runAuthenticate(session, stub.baseUrl); + + expect(result?.sessionId).toBe("session_refreshed"); + expect(result?.refreshedSession).toEqual(expect.any(String)); + expect(stub.requests().map((request) => [request.method, request.path])).toEqual([ + ["GET", `/sso/jwks/${CLIENT_ID}`], + ["POST", "/user_management/authenticate"], + ]); + expect(stub.requests()[1]?.body).toMatchObject({ + grant_type: "refresh_token", + client_id: CLIENT_ID, + client_secret: API_KEY, + refresh_token: "refresh_expired", + organization_id: "org_test", + }); + }); + }); + + it("returns null for garbage session data", async () => { + const keypair = await generateKeypair("k_garbage"); + await withWorkOSStub(keypair, async (stub) => { + const result = await runAuthenticate("not-a-sealed-session", stub.baseUrl); + + expect(result).toBeNull(); + expect(stub.requests()).toEqual([]); + }); + }); +}); describe("collectWorkOSList", () => { it("collects memberships beyond the first WorkOS page", async () => { diff --git a/apps/cloud/src/auth/workos.ts b/apps/cloud/src/auth/workos.ts index 0f9765288..c9b0809cb 100644 --- a/apps/cloud/src/auth/workos.ts +++ b/apps/cloud/src/auth/workos.ts @@ -3,13 +3,23 @@ // --------------------------------------------------------------------------- import { env } from "cloudflare:workers"; -import { Context, Data, Effect, Layer, Option, Schema } from "effect"; +import { Context, Data, Effect, Layer, Option, Predicate, Schema } from "effect"; import { GeneratePortalLinkIntent, WorkOS } from "@workos-inc/node/worker"; +import { defaults as ironDefaults, unseal as unsealIron } from "iron-webcrypto"; +import { decodeJwt, jwtVerify } from "jose"; +import { JWKSInvalid, JWKSNoMatchingKey, JWKSTimeout } from "jose/errors"; import { parseCookie } from "./cookies"; -import { tryPromiseService, withServiceLogging, workosErrorFromFailure } from "./errors"; +import { createCachedRemoteJWKSet, type CachedRemoteJWKSet } from "./jwks-cache"; +import { + ServiceAdapterError, + tryPromiseService, + withServiceLogging, + workosErrorFromFailure, +} from "./errors"; const COOKIE_NAME = "wos-session"; const INVALID_COOKIE_PASSWORD_MESSAGE = "WORKOS_COOKIE_PASSWORD must be at least 32 characters"; +const WORKOS_SEAL_VERSION_DELIMITER = "~"; type RawWorkOS = WorkOS & { readonly get: ( @@ -59,6 +69,178 @@ const RawWorkOSListResponse = Schema.Struct({ const decodeRawWorkOSListResponse = Schema.decodeUnknownOption(RawWorkOSListResponse); +const SealedSessionUser = Schema.Struct({ + id: Schema.String, + email: Schema.String, + firstName: Schema.NullOr(Schema.String), + lastName: Schema.NullOr(Schema.String), + profilePictureUrl: Schema.NullOr(Schema.String), +}); + +const SealedSessionPayload = Schema.Struct({ + accessToken: Schema.String, + refreshToken: Schema.String, + user: SealedSessionUser, + authenticationMethod: Schema.optional(Schema.Unknown), + impersonator: Schema.optional(Schema.Unknown), +}); + +const decodeSealedSessionPayload = Schema.decodeUnknownOption(SealedSessionPayload); + +const JwtClaims = Schema.Struct({ + sid: Schema.String, + org_id: Schema.optional(Schema.String), +}); + +const decodeJwtClaims = Schema.decodeUnknownOption(JwtClaims); + +const LegacySealedSessionPayload = Schema.Struct({ + persistent: Schema.Unknown, +}); + +const decodeLegacySealedSessionPayload = Schema.decodeUnknownOption(LegacySealedSessionPayload); + +type SealedSessionPayload = typeof SealedSessionPayload.Type; + +type LocalSessionVerification = + | { + readonly _tag: "Valid"; + readonly session: SealedSessionPayload; + readonly organizationId: string | undefined; + readonly sessionId: string; + } + | { readonly _tag: "Refresh" } + | { readonly _tag: "InvalidCookie" }; + +class LocalSessionCookieError extends Data.TaggedError("LocalSessionCookieError")<{ + readonly cause: unknown; +}> {} + +const isLocalSessionValid = Predicate.isTagged("Valid") as ( + value: LocalSessionVerification, +) => value is Extract; + +const isLocalSessionInvalidCookie = Predicate.isTagged("InvalidCookie") as ( + value: LocalSessionVerification, +) => value is Extract; + +const ErrorWithCode = Schema.Struct({ code: Schema.String }); +const isErrorWithCode = Schema.is(ErrorWithCode); + +const isJwtVerificationFailure = (cause: unknown): boolean => + isErrorWithCode(cause) && + (cause.code.startsWith("ERR_JWT_") || + cause.code.startsWith("ERR_JWS_") || + cause.code === JWKSNoMatchingKey.code); + +const isJwksSystemFailure = (cause: unknown): boolean => + isErrorWithCode(cause) && (cause.code === JWKSTimeout.code || cause.code === JWKSInvalid.code); + +const parseWorkOSSeal = ( + seal: string, +): { readonly sealWithoutVersion: string; readonly tokenVersion: number | null } => { + const [sealWithoutVersion = "", tokenVersionAsString] = seal.split(WORKOS_SEAL_VERSION_DELIMITER); + return { + sealWithoutVersion, + tokenVersion: + tokenVersionAsString === undefined ? null : Number.parseInt(tokenVersionAsString, 10), + }; +}; + +const unsealWorkOSSession = async ( + sessionData: string, + cookiePassword: string, +): Promise => { + const { sealWithoutVersion, tokenVersion } = parseWorkOSSeal(sessionData); + const data = + (await unsealIron(sealWithoutVersion, { 1: cookiePassword }, { ...ironDefaults, ttl: 0 })) ?? + {}; + // Mirrors the SDK's unsealData version handling: current (v2) seals hold the + // payload directly, OTHER versioned seals nest it under `persistent`, and an + // unversioned seal is the payload itself. + if (tokenVersion === 2 || tokenVersion === null) return data; + return Option.match(decodeLegacySealedSessionPayload(data), { + onNone: () => data, + onSome: (legacy) => legacy.persistent, + }); +}; + +const getWorkOSSessionJwks = (() => { + const resolvers = new Map(); + return (url: URL): CachedRemoteJWKSet => { + const key = url.toString(); + const existing = resolvers.get(key); + if (existing) return existing; + const created = createCachedRemoteJWKSet(url); + resolvers.set(key, created); + return created; + }; +})(); + +const verifyJwtOnce = (accessToken: string, jwks: CachedRemoteJWKSet) => + Effect.tryPromise({ + try: () => jwtVerify(accessToken, jwks), + catch: (cause) => new ServiceAdapterError({ cause }), + }); + +const verifyJwtWithRefreshRetry = ( + accessToken: string, + jwks: CachedRemoteJWKSet, +): Effect.Effect => + verifyJwtOnce(accessToken, jwks).pipe( + Effect.as(true), + Effect.catchTag("ServiceAdapterError", (error) => { + if (isErrorWithCode(error.cause) && error.cause.code === JWKSNoMatchingKey.code) { + return Effect.sync(() => jwks.forceRefresh()).pipe( + Effect.flatMap(() => verifyJwtOnce(accessToken, jwks)), + Effect.as(true), + Effect.catchTag("ServiceAdapterError", (retryError) => + isJwtVerificationFailure(retryError.cause) + ? Effect.succeed(false) + : Effect.fail(retryError), + ), + ); + } + if (isJwtVerificationFailure(error.cause)) return Effect.succeed(false); + if (isJwksSystemFailure(error.cause)) return Effect.fail(error); + return Effect.fail(error); + }), + ); + +const verifySealedSessionLocally = ( + sessionData: string, + cookiePassword: string, + jwks: CachedRemoteJWKSet, +): Effect.Effect => + Effect.gen(function* () { + const unsealed = yield* Effect.tryPromise({ + try: () => unsealWorkOSSession(sessionData, cookiePassword), + catch: (cause) => new LocalSessionCookieError({ cause }), + }).pipe( + Effect.catchTag("LocalSessionCookieError", () => Effect.succeed(null as unknown | null)), + ); + if (!unsealed) return { _tag: "InvalidCookie" }; + + const session = Option.match(decodeSealedSessionPayload(unsealed), { + onNone: (): SealedSessionPayload | null => null, + onSome: (payload) => payload, + }); + if (!session) return { _tag: "InvalidCookie" }; + + const verified = yield* verifyJwtWithRefreshRetry(session.accessToken, jwks); + if (!verified) return { _tag: "Refresh" }; + + const claims = Option.getOrNull(decodeJwtClaims(decodeJwt(session.accessToken))); + if (!claims) return { _tag: "Refresh" }; + + return { + _tag: "Valid", + session, + organizationId: claims.org_id, + sessionId: claims.sid, + }; + }); + const completedListMetadata = { before: null, after: null, @@ -150,6 +332,7 @@ const make = Effect.gen(function* () { clientId, ...workosApiUrlOptions(env.WORKOS_API_URL), }); + const sessionJwks = getWorkOSSessionJwks(new URL(workos.userManagement.getJwksUrl(clientId))); // The public `WorkOSError` carries the upstream HTTP status when the SDK // exception had one (all its typed exceptions do), so consumers can tell a @@ -171,22 +354,26 @@ const make = Effect.gen(function* () { cookiePassword, }); - const result = yield* use(() => session.authenticate()); + const local = yield* withServiceLogging( + "workos.session.local_verify", + workosErrorFromFailure, + verifySealedSessionLocally(sessionData, cookiePassword, sessionJwks), + ); - if (result.authenticated) { + if (isLocalSessionValid(local)) { return { - userId: result.user.id, - email: result.user.email, - firstName: result.user.firstName, - lastName: result.user.lastName, - avatarUrl: result.user.profilePictureUrl, - organizationId: result.organizationId, - sessionId: result.sessionId, + userId: local.session.user.id, + email: local.session.user.email, + firstName: local.session.user.firstName, + lastName: local.session.user.lastName, + avatarUrl: local.session.user.profilePictureUrl, + organizationId: local.organizationId, + sessionId: local.sessionId, refreshedSession: undefined as string | undefined, }; } - if (result.reason === "no_session_cookie_provided") return null; + if (isLocalSessionInvalidCookie(local)) return null; // Try refreshing const refreshed = yield* use(() => session.refresh()).pipe( diff --git a/apps/cloud/src/mcp/session-durable-object.ts b/apps/cloud/src/mcp/session-durable-object.ts index 987b8055e..b6b9ccb52 100644 --- a/apps/cloud/src/mcp/session-durable-object.ts +++ b/apps/cloud/src/mcp/session-durable-object.ts @@ -271,7 +271,9 @@ export class McpSessionDOSqlite extends McpAgentSessionDOBase the service stays unprovided and // the URL falls back to its bare, client-canonicalized form. const { executor, engine } = yield* resolved.organizationSlug !== undefined - ? stack.pipe(Effect.provideService(RequestOrgSlug, { slug: resolved.organizationSlug })) - : stack; + ? stack.pipe( + Effect.provideService(RequestOrgSlug, { slug: resolved.organizationSlug }), + Effect.withSpan("executor.stack.http.resolve"), + ) + : stack.pipe(Effect.withSpan("executor.stack.http.resolve")); return yield* httpEffect.pipe( Effect.provideService(AuthContext, auth), Effect.provideService(ExecutorService, executor), diff --git a/packages/core/api/src/server/execution-stack.ts b/packages/core/api/src/server/execution-stack.ts index 78cfd1ca6..40e9d921d 100644 --- a/packages/core/api/src/server/execution-stack.ts +++ b/packages/core/api/src/server/execution-stack.ts @@ -108,13 +108,19 @@ export const makeExecutionStack = < organizationId, organizationName, { plugins: { mcpResource: options?.mcpResource } }, + ).pipe(Effect.withSpan("executor.stack.scoped_executor")); + const codeExecutor = yield* CodeExecutorProvider.asEffect().pipe( + Effect.withSpan("executor.stack.code_executor"), ); - const codeExecutor = yield* CodeExecutorProvider; - const { decorate } = yield* EngineDecorator; - const engine = decorate(createExecutionEngine({ executor, codeExecutor }), { - accountId, - organizationId, - organizationName, - }); + const { decorate } = yield* EngineDecorator.asEffect().pipe( + Effect.withSpan("executor.stack.decorator"), + ); + const engine = yield* Effect.sync(() => + decorate(createExecutionEngine({ executor, codeExecutor }), { + accountId, + organizationId, + organizationName, + }), + ).pipe(Effect.withSpan("executor.stack.engine.init")); return { executor, engine }; - }); + }).pipe(Effect.withSpan("executor.stack.build")); diff --git a/packages/core/api/src/server/mcp-build.ts b/packages/core/api/src/server/mcp-build.ts index c274cd623..275185754 100644 --- a/packages/core/api/src/server/mcp-build.ts +++ b/packages/core/api/src/server/mcp-build.ts @@ -42,6 +42,7 @@ export const makeMcpBuildServer = makeExecutionStack(principal.accountId, principal.organizationId, principal.organizationName, { mcpResource: options?.resource, }).pipe( + Effect.withSpan("mcp.execution_stack.build"), Effect.map(({ engine }) => engine), // Pin browser-handoff URLs to the principal's org slug when present; // absent slug leaves the service unprovided and the URL stays bare. @@ -52,6 +53,7 @@ export const makeMcpBuildServer = Effect.mapError((cause) => new McpEngineBuildError({ cause })), Effect.flatMap((engine) => createExecutorMcpServer({ engine, ...(options ?? {}) }).pipe( + Effect.withSpan("mcp.server.create"), Effect.map((mcpServer) => ({ mcpServer, engine })), ), ), diff --git a/packages/core/api/src/server/scoped-executor.ts b/packages/core/api/src/server/scoped-executor.ts index 4164bc7d3..c119cb0a2 100644 --- a/packages/core/api/src/server/scoped-executor.ts +++ b/packages/core/api/src/server/scoped-executor.ts @@ -210,9 +210,13 @@ export const makeScopedExecutor = < options?: { readonly plugins?: PluginsProviderContext }, ): Effect.Effect, StorageFailure, DbProvider | PluginsProvider | HostConfig> => Effect.gen(function* () { - const { db, blobs } = yield* DbProvider; - const { plugins: pluginsFactory } = yield* PluginsProvider; - const config = yield* HostConfig; + const { db, blobs } = yield* DbProvider.asEffect().pipe( + Effect.withSpan("executor.stack.db_provider"), + ); + const { plugins: pluginsFactory } = yield* PluginsProvider.asEffect().pipe( + Effect.withSpan("executor.stack.plugins_provider"), + ); + const config = yield* HostConfig.asEffect().pipe(Effect.withSpan("executor.stack.host_config")); // Explicit config wins; otherwise fall back to the request origin if a host // provided one (HTTP middleware / MCP session DO). Stays `undefined` for // non-request callers — `coreTools.webBaseUrl` is optional and only the @@ -248,7 +252,9 @@ export const makeScopedExecutor = < oauthCallbackPath: config.oauthCallbackPath, }); - const plugins = pluginsFactory(options?.plugins); + const plugins = yield* Effect.sync(() => pluginsFactory(options?.plugins)).pipe( + Effect.withSpan("executor.plugins.init"), + ); const hostedHttpOptions = { allowLocalNetwork: config.allowLocalNetwork, }; @@ -274,7 +280,7 @@ export const makeScopedExecutor = < orgSlug, includeProviders: config.exposeCredentialProviders ?? true, }, - }); + }).pipe(Effect.withSpan("executor.stack.create_executor")); // The seam erases the plugin tuple type; the caller re-narrows via the // `TPlugins` phantom. Runtime shape is identical to a typed // `createExecutor({ plugins })` call.