From cdff0e5fa8436cffa1a04ebce704adf190c87472 Mon Sep 17 00:00:00 2001 From: Waleed Date: Sun, 1 Nov 2026 11:28:26 -0800 Subject: [PATCH 1/4] fix(sandbox): initialize workbench history before honoring a stop during creation (#9549) * fix(sandbox): initialize workbench history before honoring a stop during creation A chat workbench created while the caller stopped kept running for the next call to reuse, but the abort check ran before its history was initialized. The reused machine then had no history, so it could never be certified and every later code result in that chat was withheld. * fix(sandbox): type the certification run helper to the fields both boundaries accept --- ...session-input-certification.integration.ts | 110 ++++++++++++++---- .../lib/execution/remote-sandbox/session.ts | 4 +- 2 files changed, 88 insertions(+), 26 deletions(-) diff --git a/apps/sim/lib/execution/remote-sandbox/session-input-certification.integration.ts b/apps/sim/lib/execution/remote-sandbox/session-input-certification.integration.ts index 07b3311f34b..dc0ab982d79 100644 --- a/apps/sim/lib/execution/remote-sandbox/session-input-certification.integration.ts +++ b/apps/sim/lib/execution/remote-sandbox/session-input-certification.integration.ts @@ -2,7 +2,7 @@ * The persistent workbench's input history is what lets a scratch file reach the model. This runs * the real code boundary against the real history script in a disposable Redis, so a mount the * caller could not classify has to leave the machine uncertified. Only the sandbox provider is a - * stand-in: it hands back an existing machine that executes nothing. + * stand-in: it hands back a machine that executes nothing. * * Set TEST_REDIS_URL to an isolated local Redis service. */ @@ -14,14 +14,23 @@ import { import { generateShortId } from '@sim/utils/id' import { afterAll, describe, expect, it, vi } from 'vitest' -const { redisUrl, inheritedRedisUrl, mockFindSessionSandbox } = await vi.hoisted(async () => { - const { readTestRedisUrl } = await import('@sim/db/testing/test-infrastructure') - const url = readTestRedisUrl() - const inheritedRedisUrl = process.env.REDIS_URL - /** The real Redis module reads this at import. */ - if (url) process.env.REDIS_URL = url - return { redisUrl: url, inheritedRedisUrl, mockFindSessionSandbox: vi.fn() } -}) +const { redisUrl, inheritedRedisUrl, mockCreate, mockFindSessionSandbox } = await vi.hoisted( + async () => { + const { readTestRedisUrl } = await import('@sim/db/testing/test-infrastructure') + const url = readTestRedisUrl() + const inheritedRedisUrl = process.env.REDIS_URL + /** The real Redis module reads this at import. */ + if (url) process.env.REDIS_URL = url + return { + redisUrl: url, + inheritedRedisUrl, + mockCreate: vi.fn(async () => { + throw new Error('The certification fixture only reuses an existing machine') + }), + mockFindSessionSandbox: vi.fn>(), + } + } +) vi.mock('@/lib/execution/remote-sandbox/provider', () => remoteSandboxProviderMock) vi.mock('@/lib/execution/remote-sandbox/resolve', () => ({ @@ -43,7 +52,11 @@ import { initializeSessionFileProvenance, isSessionFileProvenanceClean, } from '@/lib/execution/remote-sandbox/session-file-provenance' -import type { SandboxHandle, SandboxProvider } from '@/lib/execution/remote-sandbox/types' +import type { + SandboxExecutionRequest, + SandboxHandle, + SandboxProvider, +} from '@/lib/execution/remote-sandbox/types' remoteSandboxProviderMockFns.mockResolveProvider.mockImplementation( () => @@ -51,9 +64,7 @@ remoteSandboxProviderMockFns.mockResolveProvider.mockImplementation( id: 'e2b', dependencyStrategy: 'prebuilt', resolveLifetimeMs: (ms: number) => ms, - create: async () => { - throw new Error('The certification fixture only reuses an existing machine') - }, + create: mockCreate, findSessionSandbox: mockFindSessionSandbox, }) satisfies SandboxProvider ) @@ -77,6 +88,28 @@ function machine(sandboxId: string): SandboxHandle { /** Machine-history keys this suite created, so cleanup never touches another suite's state. */ const createdKeys: string[] = [] +function trackHistoryKey(key: string, sandboxId: string) { + createdKeys.push( + `mothership:workbench-provenance:v2:${createHash('sha256') + .update(JSON.stringify([key, 'e2b', sandboxId])) + .digest('hex')}` + ) +} + +/** The fields both the code and the shell boundary accept. */ +type SessionRunRequest = Pick< + SandboxExecutionRequest, + 'code' | 'language' | 'timeoutMs' | 'session' | 'signal' +> + +function run(kind: 'code' | 'shell', request: SessionRunRequest) { + return observeSandboxSessionInputs( + () => true, + () => + kind === 'code' ? executeInSandbox(request) : executeShellInSandbox({ ...request, envs: {} }) + ) +} + afterAll(async () => { if (redisUrl && createdKeys.length) { const redis = getRedisClient() @@ -99,11 +132,7 @@ describe.skipIf(!redisUrl)('workbench certification at the code boundary', () => const sandboxId = `machine-${generateShortId(12)}` const key = `certification-${generateShortId(12)}` const identity = { providerId: 'e2b', sandboxId } as const - createdKeys.push( - `mothership:workbench-provenance:v2:${createHash('sha256') - .update(JSON.stringify([key, identity.providerId, sandboxId])) - .digest('hex')}` - ) + trackHistoryKey(key, sandboxId) mockFindSessionSandbox.mockResolvedValue(machine(sandboxId)) await initializeSessionFileProvenance(key, identity) expect(await isSessionFileProvenanceClean(key, identity)).toBe(true) @@ -114,15 +143,46 @@ describe.skipIf(!redisUrl)('workbench certification at the code boundary', () => timeoutMs: 30_000, session: { key, ...(unprovenanced ? { unprovenancedInputs: true } : {}) }, } - const execution = observeSandboxSessionInputs( - () => true, - () => - kind === 'code' - ? executeInSandbox(request) - : executeShellInSandbox({ ...request, envs: {} }) - ) + const execution = run(kind, request) if (unprovenanced) await expect(execution).rejects.toThrow('Workbench output withheld') else await expect(execution).resolves.toMatchObject({ sandboxId }) expect(await isSessionFileProvenanceClean(key, identity)).toBe(!unprovenanced) }) + + /** + * The provider finishes creating a session machine even when the caller stops mid-create, and + * the next call reuses that machine. Its history must already exist then: a machine reused with + * no history can never be certified, so every later result in the chat would be withheld. + */ + it.each(['code', 'shell'] as const)( + '%s certifies a workbench whose creation outlived a stopped caller', + async (kind) => { + const sandboxId = `machine-${generateShortId(12)}` + const key = `stopped-create-${generateShortId(12)}` + trackHistoryKey(key, sandboxId) + mockCreate.mockClear() + const controller = new AbortController() + mockFindSessionSandbox.mockResolvedValueOnce(null) + mockCreate.mockImplementationOnce(async () => { + controller.abort(new Error('User Stop')) + return machine(sandboxId) + }) + const request = { + code: 'print(1)', + language: CodeLanguage.Python, + timeoutMs: 30_000, + session: { key }, + } + + await expect(run(kind, { ...request, signal: controller.signal })).rejects.toThrow( + 'User Stop' + ) + expect(mockCreate).toHaveBeenCalledTimes(1) + + mockFindSessionSandbox.mockResolvedValue(machine(sandboxId)) + await expect(run(kind, request)).resolves.toMatchObject({ sandboxId }) + expect(mockCreate).toHaveBeenCalledTimes(1) + expect(await isSessionFileProvenanceClean(key, { providerId: 'e2b', sandboxId })).toBe(true) + } + ) }) diff --git a/apps/sim/lib/execution/remote-sandbox/session.ts b/apps/sim/lib/execution/remote-sandbox/session.ts index 6b3ac9dff86..aba642ac799 100644 --- a/apps/sim/lib/execution/remote-sandbox/session.ts +++ b/apps/sim/lib/execution/remote-sandbox/session.ts @@ -83,12 +83,14 @@ export async function ensureSessionSandbox(args: { true, provider ) - signal.throwIfAborted() + // Before the abort check: a stopped caller still leaves this machine for the next call to + // reuse, and a reused machine without history can never be certified. if (!existing) await initializeSessionFileProvenance(session.key, { providerId: provider.id, sandboxId: created.sandbox.sandboxId, }) + signal.throwIfAborted() // Measured from before acquisition: any lease the lookup or create granted starts no earlier. if (!created.sandbox.outlives?.(lifetimeMs, requestedAtMs)) await created.sandbox.extendLifetime?.(lifetimeMs) From 1defff3078227715a62ea06725b90b4a49f2c487 Mon Sep 17 00:00:00 2001 From: Waleed Date: Sun, 1 Nov 2026 11:41:04 -0800 Subject: [PATCH 2/4] fix(public-shares): stream shared files with byte ranges instead of buffering them under 100 MB (#9550) * fix(public-shares): stream shared files with byte ranges instead of buffering them under 100 MB * fix(public-shares): size a streamed share from storage, audit every delivered response, cover both read paths --- .../api/files/public/[token]/content/route.ts | 11 +- apps/sim/hooks/use-file-content-source.tsx | 17 +- .../content-sharing.integration.ts | 183 ++++++++++++++---- .../sim/lib/public-shares/application/read.ts | 151 ++++++++++++--- .../scripts/test-project-file-public-e2e.ts | 17 +- 5 files changed, 306 insertions(+), 73 deletions(-) diff --git a/apps/sim/app/api/files/public/[token]/content/route.ts b/apps/sim/app/api/files/public/[token]/content/route.ts index b85fe6df606..e665b3a474c 100644 --- a/apps/sim/app/api/files/public/[token]/content/route.ts +++ b/apps/sim/app/api/files/public/[token]/content/route.ts @@ -17,7 +17,10 @@ import { } from '@/lib/public-shares/application' import { enforcePublicFileRateLimit } from '@/lib/public-shares/rate-limit' import { FILE_CACHE_CONTROL } from '@/lib/uploads/server/delivery' -import { createConditionalFileResponse } from '@/app/api/files/utils' +import { + createConditionalFileResponse, + createStoredFileStreamResponse, +} from '@/app/api/files/utils' export const dynamic = 'force-dynamic' const logger = createLogger('PublicFileContentAPI') @@ -49,6 +52,12 @@ export const GET = withRouteHandler( preview: parsed.data.query.preview === '1', request, }) + if (result.kind === 'stream') { + return await createStoredFileStreamResponse( + { ...result, cacheControl: FILE_CACHE_CONTROL.revalidate }, + request.headers + ) + } return createConditionalFileResponse( { buffer: result.buffer, diff --git a/apps/sim/hooks/use-file-content-source.tsx b/apps/sim/hooks/use-file-content-source.tsx index 9dd838cc9b6..803e54dcee3 100644 --- a/apps/sim/hooks/use-file-content-source.tsx +++ b/apps/sim/hooks/use-file-content-source.tsx @@ -146,12 +146,17 @@ export function createPublicFileContentSource( contentUrl: string, owner?: EditableFileOwner ): FileContentSource { - return inlineImageSource( - (_key, opts) => - opts?.preview ? `${contentUrl}${contentUrl.includes('?') ? '&' : '?'}preview=1` : contentUrl, - `/api/files/public/${token}/inline`, - owner - ) + return { + ...inlineImageSource( + (_key, opts) => + opts?.preview + ? `${contentUrl}${contentUrl.includes('?') ? '&' : '?'}preview=1` + : contentUrl, + `/api/files/public/${token}/inline`, + owner + ), + servesByteRanges: true, + } } /** diff --git a/apps/sim/lib/projects/files/__integration__/content-sharing.integration.ts b/apps/sim/lib/projects/files/__integration__/content-sharing.integration.ts index 4dc2fbaf745..841573b0f67 100644 --- a/apps/sim/lib/projects/files/__integration__/content-sharing.integration.ts +++ b/apps/sim/lib/projects/files/__integration__/content-sharing.integration.ts @@ -1,4 +1,4 @@ -import { mkdtempSync } from 'node:fs' +import { existsSync, mkdtempSync } from 'node:fs' import { mkdir, rm, writeFile } from 'node:fs/promises' import { tmpdir } from 'node:os' import { dirname, join } from 'node:path' @@ -39,8 +39,11 @@ import { createProjectFile, updateProjectFileContent, } from '@/lib/projects/files/application/content' +import type { readPublicFileShareContent } from '@/lib/public-shares/application' import * as shareManager from '@/lib/public-shares/manager' import * as storage from '@/lib/uploads/core/storage-service' +import { MAX_BUFFERED_TRANSFER_BYTES } from '@/lib/uploads/shared/types' +import { sanitizeFileKey } from '@/lib/uploads/utils/file' vi.mock('@/lib/core/config/feature-flags', () => featureFlagsMock) @@ -129,6 +132,16 @@ async function fixture() { return { ...ids, principal: createSessionPrincipal({ userId: editorId }) } } +type PublicShareContent = Awaited> + +/** The bytes a public read delivers, whether streamed from storage or built whole. */ +async function sharedBytes(read: PublicShareContent | Promise) { + const result = await read + return result.kind === 'stream' + ? Buffer.from(await new Response(await result.open(undefined)).arrayBuffer()) + : result.buffer +} + function createInput(projectId: string, content = 'Project architecture') { return { projectId, @@ -410,9 +423,9 @@ describe('Public file shares against PostgreSQL and private storage', () => { credential: { method: 'POST', password: 'correct-password' }, }) if (!accepted.authorized) throw new Error('Correct share password rejected') - expect((await readPublicFileShareContent({ grant: accepted.grant })).buffer.toString()).toBe( - 'Shared architecture' - ) + expect( + (await sharedBytes(readPublicFileShareContent({ grant: accepted.grant }))).toString() + ).toBe('Shared architecture') const [passwordShare] = await db .select() .from(publicShare) @@ -653,6 +666,74 @@ describe('Public file shares against PostgreSQL and private storage', () => { } ) + check( + 'public share streams stored bytes past the buffered ceiling and refuses once the share is disabled', + async () => { + const { updateProjectFileShare } = await import('@/lib/projects/files/application/shares') + const { authorizePublicFileShare, readPublicFileShareContent } = await import( + '@/lib/public-shares/application' + ) + const f = await fixture() + const share = async (name: string, contentType: string) => { + const { file } = await createProjectFile.execute({ + principal: f.principal, + input: { ...createInput(f.projectId, 'placeholder'), name, contentType }, + }) + const { share } = await updateProjectFileShare.execute({ + principal: f.principal, + input: { projectId: f.projectId, fileId: file.id, isActive: true }, + }) + const access = await authorizePublicFileShare({ + token: share.token, + credential: { method: 'GET' }, + }) + if (!access.authorized) throw new Error(`Public ${name} fixture rejected`) + return { file, grant: access.grant } + } + + const recording = await share('recording.mp4', 'video/mp4') + const size = MAX_BUFFERED_TRANSFER_BYTES + 1024 + const bytes = Buffer.alloc(size) + bytes.write('tail', size - 4) + const [row] = await db + .select() + .from(workspaceFiles) + .where(eq(workspaceFiles.id, recording.file.id)) + const objectPath = join(localStorageRoot, sanitizeFileKey(row.key)) + expect(existsSync(objectPath)).toBe(true) + await writeFile(objectPath, bytes) + await db + .update(workspaceFiles) + .set({ sizeBytes: size }) + .where(eq(workspaceFiles.id, recording.file.id)) + + const stored = await readPublicFileShareContent({ grant: recording.grant }) + if (stored.kind !== 'stream') throw new Error('A stored recording was buffered') + expect(stored.size).toBe(size) + const whole = Buffer.from(await new Response(await stored.open(undefined)).arrayBuffer()) + expect(whole.length).toBe(size) + expect(whole.subarray(size - 4).toString()).toBe('tail') + const span = await stored.open({ start: size - 4, end: size - 1 }) + expect(Buffer.from(await new Response(span).arrayBuffer()).toString()).toBe('tail') + + await updateProjectFileShare.execute({ + principal: f.principal, + input: { projectId: f.projectId, fileId: recording.file.id, isActive: false }, + }) + await expect(stored.open({ start: 0, end: 3 })).rejects.toMatchObject({ code: 'not_found' }) + + const page = await share('notes.html', 'text/html') + expect(await readPublicFileShareContent({ grant: page.grant })).toMatchObject({ + kind: 'buffer', + }) + // A generation source is compiled, never streamed: with no cached artifact it refuses. + const source = await share('report.pptx', 'text/x-pptxgenjs') + await expect(readPublicFileShareContent({ grant: source.grant })).rejects.toMatchObject({ + code: 'conflict', + }) + } + ) + check( 'public share Office artifacts require the current owner-qualified dependency cache without compilation', async () => { @@ -693,7 +774,7 @@ describe('Public file shares against PostgreSQL and private storage', () => { principal: f.principal, input: { projectId: f.projectId, fileId: source.file.id, maxBytes: 1024 }, }) - expect((await readPublicFileShareContent({ grant: access.grant })).buffer).toEqual( + expect(await sharedBytes(readPublicFileShareContent({ grant: access.grant }))).toEqual( rendered.buffer ) await updateProjectFileContent.execute({ @@ -719,31 +800,38 @@ describe('Public file shares against PostgreSQL and private storage', () => { '@/lib/public-shares/application' ) const { encryptSecret } = await import('@/lib/core/security/encryption') - for (const change of ['disabled', 'password', 'archived', 'provenance'] as const) { - const f = await fixture() - const { file } = await createProjectFile.execute({ - principal: f.principal, - input: createInput(f.projectId), - }) - const { share } = await updateProjectFileShare.execute({ - principal: f.principal, - input: { - projectId: f.projectId, - fileId: file.id, - isActive: true, - authType: 'password', - password: 'initial-password', - }, - }) - const access = await authorizePublicFileShare({ - token: share.token, - credential: { method: 'POST', password: 'initial-password' }, - }) - if (!access.authorized) throw new Error('Initial credential rejected') - const download = storage.downloadFile - const spy = vi.spyOn(storage, 'downloadFile').mockImplementation(async (options) => { - const result = await download(options) - if (options.key === file.key) { + // A Markdown file streams from storage; a page is rendered whole. Both reads must refuse. + const shared = { + streamed: (projectId: string) => createInput(projectId), + rendered: (projectId: string) => ({ + ...createInput(projectId, '---\ntitle: Architecture\n---\n\nShared page'), + name: 'Architecture', + contentType: 'text/x-sim-page', + }), + } + for (const [read, input] of Object.entries(shared)) + for (const change of ['disabled', 'password', 'archived', 'provenance'] as const) { + const f = await fixture() + const { file } = await createProjectFile.execute({ + principal: f.principal, + input: input(f.projectId), + }) + const { share } = await updateProjectFileShare.execute({ + principal: f.principal, + input: { + projectId: f.projectId, + fileId: file.id, + isActive: true, + authType: 'password', + password: 'initial-password', + }, + }) + const access = await authorizePublicFileShare({ + token: share.token, + credential: { method: 'POST', password: 'initial-password' }, + }) + if (!access.authorized) throw new Error('Initial credential rejected') + const applyChange = async () => { if (change === 'disabled') await db .update(publicShare) @@ -765,13 +853,30 @@ describe('Public file shares against PostgreSQL and private storage', () => { .set({ deletedAt: new Date() }) .where(eq(workspaceFiles.id, file.id)) } - return result - }) - await expect(readPublicFileShareContent({ grant: access.grant })).rejects.toMatchObject({ - code: change === 'provenance' ? 'forbidden' : 'not_found', - }) - spy.mockRestore() - } + // The change lands while the object is being read, buffered or streamed alike. + const download = storage.downloadFile + const downloadStream = storage.downloadFileStream + const spies = [ + vi.spyOn(storage, 'downloadFile').mockImplementation(async (options) => { + const result = await download(options) + if (options.key === file.key) await applyChange() + return result + }), + vi.spyOn(storage, 'downloadFileStream').mockImplementation(async (options) => { + const result = await downloadStream(options) + if (options.key === file.key) await applyChange() + return result + }), + ] + const delivered = readPublicFileShareContent({ grant: access.grant }) + await expect(sharedBytes(delivered)).rejects.toMatchObject({ + code: change === 'provenance' ? 'forbidden' : 'not_found', + }) + // Each fixture took the path it names: the stream refuses on open, the page as it renders. + if (read === 'streamed') expect(await delivered).toMatchObject({ kind: 'stream' }) + else await expect(delivered).rejects.toBeDefined() + for (const spy of spies) spy.mockRestore() + } } ) @@ -836,7 +941,7 @@ describe('Public file shares against PostgreSQL and private storage', () => { }) if (!legacyAccess.authorized) throw new Error('Existing workspace share rejected') expect( - (await readPublicFileShareContent({ grant: legacyAccess.grant })).buffer.toString() + (await sharedBytes(readPublicFileShareContent({ grant: legacyAccess.grant }))).toString() ).toBe('Workspace link') } ) @@ -892,7 +997,7 @@ describe('Public file shares against PostgreSQL and private storage', () => { if (!access.authorized) throw new Error('Public page rejected') const rendered = await readPublicFileShareContent({ grant: access.grant }) expect(rendered.contentType).toBe('text/html') - expect(rendered.buffer.toString()).toContain( + expect((await sharedBytes(rendered)).toString()).toContain( `data:image/png;base64,${image.toString('base64')}` ) expect( diff --git a/apps/sim/lib/public-shares/application/read.ts b/apps/sim/lib/public-shares/application/read.ts index 9b881e76820..cd35d316e83 100644 --- a/apps/sim/lib/public-shares/application/read.ts +++ b/apps/sim/lib/public-shares/application/read.ts @@ -7,6 +7,7 @@ import { OrchestrationError, type OrchestrationRequestContext, } from '@/lib/core/orchestration/types' +import { nodeReadableToWebStream } from '@/lib/core/utils/node-stream' import { assertKnownSizeWithinLimit, isPayloadSizeLimitError } from '@/lib/core/utils/stream-limits' import type { DbTransaction } from '@/lib/db/types' import { @@ -23,16 +24,17 @@ import { mergeWorkspaceFileSecretProvenance, type WorkspaceFileSecretProvenance, } from '@/lib/uploads/contexts/workspace/file-secret-provenance' -import { downloadFile, headObject } from '@/lib/uploads/core/storage-service' +import { downloadFile, downloadFileStream, headObject } from '@/lib/uploads/core/storage-service' import { collectReferencedFileIds, getDocumentSourceLanguage, getE2BDocFormat, isCompiledDocumentBuffer, } from '@/lib/uploads/documents/compile' +import { storedObjectEtag } from '@/lib/uploads/server/delivery' import { hasEmbeddedFileRef } from '@/lib/uploads/server/embedded-image-refs' import { resolveServableImageBytes } from '@/lib/uploads/server/image-derivative' -import { MAX_BUFFERED_TRANSFER_BYTES } from '@/lib/uploads/shared/types' +import { type ByteRange, MAX_BUFFERED_TRANSFER_BYTES } from '@/lib/uploads/shared/types' import { resolveEffectiveMimeType } from '@/lib/uploads/utils/file' import { sniffImageContentType } from '@/lib/uploads/utils/validation' import { reportWorkspaceFileDelivery } from '@/lib/workspace-files/application/file-delivery-observer' @@ -44,6 +46,7 @@ import { collectSimPageFileReferences, inlineSimPageImages, } from '@/lib/workspace-files/page-document.server' +import { servesStoredBytes } from '@/lib/workspace-files/stored-bytes' interface PublicReadInput { grant: VerifiedPublicFileShareGrant @@ -127,6 +130,42 @@ async function dependencies( await admitPublicFiles(tx, rows) return rows } +/** Audits one anonymous delivery of a shared file, or of a document input it embedded. */ +function recordPublicShareDownload(input: { + workspaceId: string | null + organizationId: string | null + owner: PublicFileShareSnapshot['owner'] + sharedFile: WorkspaceFileRow + servedFile: WorkspaceFileRow + operation: PublicFileReadOperation + bytes: number + /** The inclusive span a ranged response carried; absent for the whole file. */ + range?: ByteRange + request?: OrchestrationRequestContext +}) { + const { servedFile } = input + recordAudit({ + actorId: null, + workspaceId: input.workspaceId, + action: AuditAction.FILE_DOWNLOADED, + resourceType: AuditResourceType.FILE, + resourceId: servedFile.id, + resourceName: servedFile.originalName, + description: `Public share download of "${servedFile.originalName}"`, + request: input.request, + metadata: { + organizationId: input.organizationId, + operation: input.operation.id, + sharedDocumentId: input.sharedFile.id, + sharedByUserId: input.sharedFile.userId, + access: 'public_share', + anonymous: true, + owner: input.owner, + bytes: input.bytes, + ...(input.range ? { range: `${input.range.start}-${input.range.end}` } : {}), + }, + }) +} async function finish( grant: VerifiedPublicFileShareGrant, manifest: PublicManifest, @@ -165,25 +204,15 @@ async function finish( } }) await reportWorkspaceFileDelivery(result.secretProvenance, servedFile.id) - recordAudit({ - actorId: null, + recordPublicShareDownload({ workspaceId: manifest.snapshot.workspaceId, - action: AuditAction.FILE_DOWNLOADED, - resourceType: AuditResourceType.FILE, - resourceId: servedFile.id, - resourceName: servedFile.originalName, - description: `Public share download of "${servedFile.originalName}"`, + organizationId: result.organizationId, + owner: result.owner, + sharedFile: manifest.snapshot.file, + servedFile, + operation, + bytes: buffer.length, request, - metadata: { - organizationId: result.organizationId, - operation: operation.id, - sharedDocumentId: result.file.id, - sharedByUserId: manifest.snapshot.file.userId, - access: 'public_share', - anonymous: true, - owner: result.owner, - bytes: buffer.length, - }, }) return result } @@ -207,8 +236,10 @@ export async function checkPublicFileShareContent( }) const stored = await headObject(initial.file.key, initial.adapter.storageContext) if (!stored) throw new OrchestrationError('not_found', 'File not found') - assertKnownSizeWithinLimit(stored.size, MAX_BUFFERED_TRANSFER_BYTES, 'File') - const format = await getE2BDocFormat(initial.file.originalName) + // Stored bytes stream at any size; only a representation built from them is held whole. + const streamed = servesStoredBytes(initial.file.contentType, initial.file.originalName) + if (!streamed) assertKnownSizeWithinLimit(stored.size, MAX_BUFFERED_TRANSFER_BYTES, 'File') + const format = streamed ? null : await getE2BDocFormat(initial.file.originalName) let inputs: WorkspaceFileRow[] = [] if (format) { const source = await downloadFile({ @@ -264,7 +295,75 @@ export async function checkPublicFileShareContent( } } -/** A public read never executes generated document code; it serves current cached artifacts only. */ +/** + * A shared file served as its stored bytes, read from storage only as far as a response needs: + * the whole object, or one inclusive span. The bytes are never held, so a file of any size is + * served and a byte range is read on its own. Storage is asked for the object first, so a missing + * object is refused before any header is sent, and every length is measured against the object. + * + * `open` rechecks the grant, the revision, and the file's provenance once storage has opened and + * before the stream is handed out, as a buffered read does after rendering. Each delivered + * response is reported and audited, with the span it carried. + */ +async function streamSharedFile( + grant: VerifiedPublicFileShareGrant, + initial: PublicFileShareSnapshot, + operation: PublicFileReadOperation, + request?: OrchestrationRequestContext +) { + const { file } = initial + const context = initial.adapter.storageContext + const stored = await headObject(file.key, context) + if (!stored) throw new OrchestrationError('not_found', 'File not found') + const { size } = stored + return { + kind: 'stream' as const, + filename: file.originalName, + contentType: resolveEffectiveMimeType(undefined, file.originalName), + size, + /** A content write uploads under a fresh key and stamps `contentUpdatedAt`. */ + etag: storedObjectEtag([file.key, size, file.contentUpdatedAt.getTime(), stored.version ?? '']), + async open(range: ByteRange | undefined): Promise> { + const stream = nodeReadableToWebStream( + await downloadFileStream({ key: file.key, context, range }) + ) + try { + const provenance = await withPublicFileShareGrant( + grant, + operation, + async (tx, snapshot) => { + requireRevision(snapshot.file, file) + return admitPublicFiles(tx, [snapshot.file]) + } + ) + await reportWorkspaceFileDelivery(provenance, file.id) + } catch (error) { + await stream.cancel(error).catch(() => undefined) + throw asOrchestrationError(error) ?? error + } + recordPublicShareDownload({ + workspaceId: initial.workspaceId, + organizationId: initial.organizationId, + owner: initial.owner, + sharedFile: file, + servedFile: file, + operation, + bytes: range ? range.end - range.start + 1 : size, + range, + request, + }) + return stream + }, + } +} + +/** + * Reads a shared file for delivery. A file served as its stored bytes comes back as a stream of any + * size (see {@link servesStoredBytes}); a page or generated document is rendered, and a `preview` + * may be a HEIC transcode sniffed from the bytes, so those are built whole under the buffered + * ceiling. A public read never executes generated document code; it serves current cached + * artifacts only. + */ export async function readPublicFileShareContent({ grant, preview, @@ -276,6 +375,9 @@ export async function readPublicFileShareContent({ await admitPublicFiles(tx, [snapshot.file]) return snapshot }) + if (!preview && servesStoredBytes(initial.file.contentType, initial.file.originalName)) { + return await streamSharedFile(grant, initial, operation, request) + } const raw = await downloadFile({ key: initial.file.key, context: initial.adapter.storageContext, @@ -361,7 +463,10 @@ export async function readPublicFileShareContent({ const image = await resolveServableImageBytes(raw, initial.file.key, { owner: initial.owner }) if (image) ({ buffer, contentType } = image) } - return await finish(grant, manifest, operation, buffer, contentType, request) + return { + kind: 'buffer' as const, + ...(await finish(grant, manifest, operation, buffer, contentType, request)), + } } catch (error) { throw asOrchestrationError(error) ?? error } diff --git a/apps/sim/scripts/test-project-file-public-e2e.ts b/apps/sim/scripts/test-project-file-public-e2e.ts index 12db759cc67..fb67a147adb 100644 --- a/apps/sim/scripts/test-project-file-public-e2e.ts +++ b/apps/sim/scripts/test-project-file-public-e2e.ts @@ -60,12 +60,18 @@ const image = await sharp({ async function request( path: string, - options: { method?: string; body?: Record; cookie?: string } = {} + options: { + method?: string + body?: Record + cookie?: string + headers?: Record + } = {} ) { const headers = new Headers({ Origin: base.origin, 'Content-Type': 'application/json', 'X-Forwarded-For': proofIp, + ...options.headers, }) if (options.cookie) headers.set('Cookie', options.cookie) // boundary-raw-fetch: this proof exercises actual public binary, RSC, and credential HTTP surfaces. @@ -281,10 +287,13 @@ try { ) }) await check( - `${owner}: source and inline scan enforce independent real byte ceilings`, + `${owner}: stored bytes stream past the buffered ceiling while the inline scan keeps its own`, async () => { - await withSparseFile(required(doc.id), 100 * 1024 * 1024 + 1, async () => { - assertBinaryError(await request(`${path}/content`), 413) + const size = 100 * 1024 * 1024 + 1 + await withSparseFile(required(doc.id), size, async () => { + const head = await request(`${path}/content`, { headers: { range: 'bytes=0-0' } }) + assert.equal(head.status, 206) + assert.equal(head.headers.get('content-range'), `bytes 0-0/${size}`) assertBinaryError(await request(`${path}/inline?fileId=${required(photo.id)}`), 404) }) await withSparseFile(required(photo.id), 100 * 1024 * 1024 + 1, async () => { From 9491d3fc3962b05ce98546158e30636753bf3e76 Mon Sep 17 00:00:00 2001 From: Waleed Date: Sun, 1 Nov 2026 11:41:30 -0800 Subject: [PATCH 3/4] feat(billing): require a paid plan or prepaid credits for Chat (#9551) --- .../copilot/api-keys/validate/route.test.ts | 2 +- .../api/copilot/api-keys/validate/route.ts | 24 ++++++++++- .../components/special-tags/special-tags.tsx | 33 ++++++++++---- .../home/components/user-input/user-input.tsx | 7 ++- .../home/hooks/use-chat-plan-gate.ts | 43 +++++++++++++++++++ .../[workspaceId]/home/hooks/use-chat.ts | 4 ++ .../comparison-table/comparison-data.ts | 4 ++ .../[workspaceId]/upgrade/plan-configs.ts | 2 + apps/sim/lib/api/contracts/workspaces.ts | 5 +++ apps/sim/lib/billing/chat-plan.ts | 5 +++ apps/sim/lib/billing/core/chat-access.ts | 40 +++++++++++++++++ apps/sim/lib/billing/usage-upgrade.ts | 17 +++++++- apps/sim/lib/mothership/jobs/dispatch.ts | 7 +++ .../mothership/request/lifecycle/run.test.ts | 18 ++++++-- .../lib/mothership/request/lifecycle/run.ts | 24 ++++++++--- .../lib/mothership/request/lifecycle/start.ts | 3 ++ .../lib/mothership/request/tools/billing.ts | 25 ++++++++++- apps/sim/lib/workspaces/host-context.test.ts | 1 + apps/sim/lib/workspaces/host-context.ts | 10 ++++- 19 files changed, 249 insertions(+), 25 deletions(-) create mode 100644 apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat-plan-gate.ts create mode 100644 apps/sim/lib/billing/chat-plan.ts create mode 100644 apps/sim/lib/billing/core/chat-access.ts diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts index 1a5f8d643be..e04b0aac04b 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts @@ -292,7 +292,7 @@ describe('POST /api/copilot/api-keys/validate billing protocols', () => { const res = await POST(request(SELF_HOSTED_WORKSPACELESS_VALIDATE_BODY)) expect(res.status).toBe(200) - expect(mockCheckServerSideUsageLimits).toHaveBeenCalledWith('user-1') + expect(mockCheckServerSideUsageLimits).toHaveBeenCalledWith('user-1', ACCOUNT_SUBSCRIPTION) expect(mockResolveLegacyV0BillingAttribution).not.toHaveBeenCalled() expect(mockCheckAttributedUsageLimits).not.toHaveBeenCalled() }) diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.ts b/apps/sim/app/api/copilot/api-keys/validate/route.ts index e52ba5e5dde..1a159072395 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.ts @@ -26,6 +26,7 @@ import { serializeAccountBillingDecisionHeader, serializeBillingAttributionHeader, } from '@/lib/billing/core/attribution' +import { payerHasChatAccess } from '@/lib/billing/core/chat-access' import { isFreeRunPayer, readMidRunAccountUsageVerdict, @@ -34,7 +35,7 @@ import { import { getHighestPrioritySubscription } from '@/lib/billing/core/plan' import { isEnterprisePlan } from '@/lib/billing/core/subscription' import { deriveBillingContext } from '@/lib/billing/core/usage-log' -import { resolveUsageUpgradePayload } from '@/lib/billing/usage-upgrade' +import { resolveUsageUpgradePayload, type UsageUpgradePayer } from '@/lib/billing/usage-upgrade' import { isBillingEnabled, isHosted } from '@/lib/core/config/env-flags' import { asOrchestrationError } from '@/lib/core/orchestration/types' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' @@ -246,11 +247,13 @@ function resolveContinuationBilling( } } +/** The admitted payer's usage standing, and the payer admission reads to decide Chat access. */ async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise<{ isExceeded: boolean currentUsage: number limit: number scope: string + payer: UsageUpgradePayer accountBillingDecision?: AccountBillingDecision }> { if (admission.kind === 'attributed' || admission.kind === 'legacy-scoped') { @@ -262,6 +265,7 @@ async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise currentUsage: enforcedUsage?.currentUsage ?? 0, limit: enforcedUsage?.limit ?? 0, scope: usage.scope ?? 'payer', + payer: admission.attribution, } } @@ -276,6 +280,7 @@ async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise currentUsage: usage.currentUsage, limit: usage.limit, scope: 'account', + payer: { billingEntity: billingContext.billingEntity, payerSubscription: subscription }, accountBillingDecision: { userId: admission.userId, billingEntity: billingContext.billingEntity, @@ -291,12 +296,17 @@ async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise } } - const usage = await checkServerSideUsageLimits(admission.userId) + const subscription = await getHighestPrioritySubscription(admission.userId) + const usage = await checkServerSideUsageLimits(admission.userId, subscription) return { isExceeded: usage.isExceeded, currentUsage: usage.currentUsage, limit: usage.limit, scope: 'account', + payer: { + billingEntity: deriveBillingContext(admission.userId, subscription).billingEntity, + payerSubscription: subscription, + }, } } @@ -528,6 +538,16 @@ export const POST = withRouteHandler((req: NextRequest) => : (usage.accountBillingDecision?.billingEntity ?? { type: 'account', id: userId }), }) + // A self-hosted or key-based turn is billed to this hosted payer, whose free plan does + // not include Chat. The refusal is the admission 402 every caller already renders as + // the upgrade card. + if (!(await payerHasChatAccess(usage.payer))) { + logger.info('[API VALIDATION] Free plan does not include Chat', { userId }) + span.setAttribute(TraceAttr.CopilotValidateOutcome, CopilotValidateOutcome.UsageExceeded) + span.setAttribute(TraceAttr.HttpStatusCode, 402) + return new NextResponse(null, { status: 402 }) + } + if (usage.isExceeded) { logger.info('[API VALIDATION] Usage exceeded', { userId, diff --git a/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.tsx b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.tsx index 7d55d929cf4..1e78416c8f0 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.tsx +++ b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/special-tags.tsx @@ -33,7 +33,12 @@ import { useParams } from 'next/navigation' import { getOrganizationSettingsHref } from '@/components/settings/navigation' import { useSettingsNavigationTarget } from '@/components/settings/settings-navigation-provider' import { useSession } from '@/lib/auth/client' -import { buildHostedUpgradeUrl, HOSTED_BILLING_SETTINGS_URL } from '@/lib/billing/upgrade-reasons' +import { CHAT_PLAN_REQUIRED_REASON } from '@/lib/billing/chat-plan' +import { + buildHostedUpgradeUrl, + buildUpgradeHref, + HOSTED_BILLING_SETTINGS_URL, +} from '@/lib/billing/upgrade-reasons' import { canManageWorkspaceBilling } from '@/lib/billing/workspace-permissions' import { isBrowserAgentAvailable, sendBrowserPanelAction } from '@/lib/browser-agent/transport' import { useDeploymentShape } from '@/lib/core/config/deployment-shape' @@ -3491,7 +3496,10 @@ interface PlanUsageLimitDisplayProps { data: UsageUpgradeTagData } -/** The payer's plan limit: the card links to where the plan or the limit can change. */ +/** + * The payer's plan limit, or a free plan that does not include Chat: the card links to where + * the plan or the limit can change, and a workspace's plan to its upgrade page. + */ function PlanUsageLimitDisplay({ data }: PlanUsageLimitDisplayProps) { const { data: session } = useSession() const hostContext = useOptionalWorkspaceHostContext() @@ -3500,6 +3508,7 @@ function PlanUsageLimitDisplay({ data }: PlanUsageLimitDisplayProps) { const { hosted } = useDeploymentShape() const navigationTarget = useSettingsNavigationTarget() const opensInNewTab = !hosted || navigationTarget === '_blank' + const planRequired = data.reason === CHAT_PLAN_REQUIRED_REASON const buttonLabel = data.action === 'upgrade_plan' ? 'Upgrade Plan' : 'Increase Limit' // Self-hosted plan and limit both live on the hosted account, so local @@ -3507,7 +3516,9 @@ function PlanUsageLimitDisplay({ data }: PlanUsageLimitDisplayProps) { const href = hosted ? organizationContext ? getOrganizationSettingsHref(organizationContext.organization.id, 'billing') - : getSettingsHref({ section: 'billing' }) + : planRequired && hostContext + ? buildUpgradeHref(hostContext.workspace.id) + : getSettingsHref({ section: 'billing' }) : data.action === 'upgrade_plan' ? buildHostedUpgradeUrl() : HOSTED_BILLING_SETTINGS_URL @@ -3519,11 +3530,15 @@ function PlanUsageLimitDisplay({ data }: PlanUsageLimitDisplayProps) { const usageGate = useWorkspaceUsageGate( data.action === 'increase_limit' && !canManageBilling ? hostContext?.workspace.id : undefined ) - const unavailableMessage = organizationContext - ? 'Contact an organization admin to manage usage limits.' - : hostContext?.hostOrganizationId - ? 'Contact an organization admin to manage this workspace’s usage limits.' - : 'Only the workspace owner can manage this workspace’s usage limits.' + const action = planRequired + ? 'upgrade this workspace’s plan' + : 'manage this workspace’s usage limits' + const unavailableMessage = + organizationContext && !planRequired + ? 'Contact an organization admin to manage usage limits.' + : organizationContext || hostContext?.hostOrganizationId + ? `Contact an organization admin to ${action}.` + : `Only the workspace owner can ${action}.` return ( - Usage Limit Reached + {planRequired ? 'Paid Plan Required' : 'Usage Limit Reached'} } > diff --git a/apps/sim/app/workspace/[workspaceId]/home/components/user-input/user-input.tsx b/apps/sim/app/workspace/[workspaceId]/home/components/user-input/user-input.tsx index 34d9d043b6a..3595beb2823 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/components/user-input/user-input.tsx +++ b/apps/sim/app/workspace/[workspaceId]/home/components/user-input/user-input.tsx @@ -35,6 +35,7 @@ import { InputToolbar } from '@/app/workspace/[workspaceId]/home/components/user import type { ModelSelectorOptions } from '@/app/workspace/[workspaceId]/home/components/user-input/components/model-selector' import { usePromptHistory } from '@/app/workspace/[workspaceId]/home/components/user-input/hooks/use-prompt-history' import { handleMothershipAddContextEvent } from '@/app/workspace/[workspaceId]/home/components/user-input/mothership-context-event' +import { useChatPlanGate } from '@/app/workspace/[workspaceId]/home/hooks/use-chat-plan-gate' import type { ChatMessage, FileAttachmentForApi, @@ -507,6 +508,9 @@ const UserInputImpl = forwardRef(function UserI textareaRef.current?.focus() } + /** Refuses a send the workspace plan excludes here, before `onSubmit`, so the draft survives. */ + const stopSendWithoutChatPlan = useChatPlanGate() + /** Empties the text, chips, attachments, transcript, and the saved draft in one step. */ const clearComposer = useCallback(() => { history.reset() @@ -527,6 +531,7 @@ const UserInputImpl = forwardRef(function UserI }, [resetTranscript, history.reset]) const handleSubmit = useCallback(() => { + if (stopSendWithoutChatPlan()) return const currentFiles = filesRef.current const currentEditor = editorRef.current @@ -552,7 +557,7 @@ const UserInputImpl = forwardRef(function UserI ) if (accepted === false) return clearComposer() - }, [onSubmit, clearComposer]) + }, [stopSendWithoutChatPlan, onSubmit, clearComposer]) /** * Enter policy for the editor: mirror canSubmit's uploading guard (Enter diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat-plan-gate.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat-plan-gate.ts new file mode 100644 index 00000000000..93c444584e3 --- /dev/null +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat-plan-gate.ts @@ -0,0 +1,43 @@ +'use client' + +import { useCallback } from 'react' +import { toast } from '@sim/emcn' +import { useRouter } from 'next/navigation' +import { useSession } from '@/lib/auth/client' +import { CHAT_PLAN_REQUIRED_MESSAGE } from '@/lib/billing/chat-plan' +import { buildUpgradeHref } from '@/lib/billing/upgrade-reasons' +import { canManageWorkspaceBilling } from '@/lib/billing/workspace-permissions' +import { useOptionalWorkspaceHostContext } from '@/app/workspace/[workspaceId]/providers/workspace-host-provider' + +/** + * Stops a send from a workspace whose payer may not use Chat before it starts a turn that run + * admission would refuse. A member who can manage the workspace's billing goes to its upgrade + * page; anyone else is told who can upgrade. The returned callback reports whether the send was + * stopped. + * + * The verdict is the host context's `features.chat`, which the server resolves with the same + * check admission runs. A surface without a workspace host (an organization conversation), or a + * context that predates the flag, is not stopped here; admission still answers it with the + * upgrade card. + */ +export function useChatPlanGate(): () => boolean { + const hostContext = useOptionalWorkspaceHostContext() + const { data: session } = useSession() + const viewerId = session?.user?.id + const router = useRouter() + + return useCallback(() => { + if (!hostContext || hostContext.features?.chat !== false) return false + + if (canManageWorkspaceBilling(hostContext, viewerId)) { + router.push(buildUpgradeHref(hostContext.workspace.id)) + } else { + toast.error( + hostContext.hostOrganizationId + ? `${CHAT_PLAN_REQUIRED_MESSAGE} Contact an organization admin to upgrade this workspace’s plan.` + : `${CHAT_PLAN_REQUIRED_MESSAGE} Ask the workspace owner to upgrade this workspace’s plan.` + ) + } + return true + }, [hostContext, viewerId, router]) +} diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts index 5577b79af38..d21ab236c1f 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts @@ -119,6 +119,7 @@ import { dispatchStreamEvent, finalizeResidualToolCalls, } from '@/app/workspace/[workspaceId]/home/hooks/stream' +import { useChatPlanGate } from '@/app/workspace/[workspaceId]/home/hooks/use-chat-plan-gate' import { useNativeActiveTabIds } from '@/app/workspace/[workspaceId]/home/hooks/use-desktop-tab-resources' import { useMothershipChatOlderMessages } from '@/app/workspace/[workspaceId]/home/hooks/use-mothership-chat-older-messages' import { resolveEffectiveResourceId } from '@/app/workspace/[workspaceId]/home/resource-view-policy' @@ -779,6 +780,7 @@ export function useChat( const pathname = usePathname() const router = useRouter() const queryClient = useQueryClient() + const stopSendWithoutChatPlan = useChatPlanGate() const [pendingMessages, setPendingMessages] = useState([]) const [completedTurn, setCompletedTurn] = useState(null) const completedTurnRef = useRef(null) @@ -4437,6 +4439,7 @@ export function useChat( options?: SendMessageOptions ) => { if ((!message.trim() && !fileAttachments?.length) || !scopeKey) return + if (stopSendWithoutChatPlan()) return const queueStore = useMothershipQueueStore.getState() const activeChatKey = chatKeyRef.current @@ -4524,6 +4527,7 @@ export function useChat( }, [ workspaceId, + stopSendWithoutChatPlan, createQueuedMessage, startSendMessage, handOffWithdrawnSend, diff --git a/apps/sim/app/workspace/[workspaceId]/upgrade/components/comparison-table/comparison-data.ts b/apps/sim/app/workspace/[workspaceId]/upgrade/components/comparison-table/comparison-data.ts index 0eb4e368a5e..5f443636eb3 100644 --- a/apps/sim/app/workspace/[workspaceId]/upgrade/components/comparison-table/comparison-data.ts +++ b/apps/sim/app/workspace/[workspaceId]/upgrade/components/comparison-table/comparison-data.ts @@ -176,6 +176,10 @@ export const COMPARISON_SECTIONS: ComparisonSection[] = [ { title: 'Features', rows: [ + { + label: 'Chat with Sim', + values: [false, true, true, true], + }, { label: 'Sim Mailer (Inbox)', values: [false, false, true, true], diff --git a/apps/sim/app/workspace/[workspaceId]/upgrade/plan-configs.ts b/apps/sim/app/workspace/[workspaceId]/upgrade/plan-configs.ts index 03b87e22618..9b48675f6eb 100644 --- a/apps/sim/app/workspace/[workspaceId]/upgrade/plan-configs.ts +++ b/apps/sim/app/workspace/[workspaceId]/upgrade/plan-configs.ts @@ -27,6 +27,7 @@ export const ENTERPRISE_PLAN_CREDITS: PlanCredits = { } export const PRO_PLAN_FEATURES: readonly string[] = [ + 'Chat with Sim', `${DEFAULT_BILLING_CONCURRENCY_LIMITS.pro.toLocaleString('en-US')} concurrent executions`, 'Invite teammates', 'Higher rate limits', @@ -35,6 +36,7 @@ export const PRO_PLAN_FEATURES: readonly string[] = [ ] export const MAX_PLAN_FEATURES: readonly string[] = [ + 'Chat with Sim', `${DEFAULT_BILLING_CONCURRENCY_LIMITS.team.toLocaleString('en-US')} concurrent executions`, 'Invite teammates', 'Sim Mailer, KB Live Sync & Sandboxes', diff --git a/apps/sim/lib/api/contracts/workspaces.ts b/apps/sim/lib/api/contracts/workspaces.ts index bbfc2b2e924..71bb7ae8efb 100644 --- a/apps/sim/lib/api/contracts/workspaces.ts +++ b/apps/sim/lib/api/contracts/workspaces.ts @@ -288,6 +288,11 @@ export const workspaceHostContextSchema = z.object({ knowledgeMemberAccess: z.boolean().optional(), /** Optional for rolling compatibility with app versions that predate administrator mode. */ knowledgeSourceMirroredAccess: z.boolean().optional(), + /** + * Whether the workspace payer may use Chat. Optional for rolling compatibility; absent + * leaves the composer open, and run admission still decides. + */ + chat: z.boolean().optional(), }) .optional(), /** Optional for rolling compatibility with app versions that predate deployment projection. */ diff --git a/apps/sim/lib/billing/chat-plan.ts b/apps/sim/lib/billing/chat-plan.ts new file mode 100644 index 00000000000..a1153b6c037 --- /dev/null +++ b/apps/sim/lib/billing/chat-plan.ts @@ -0,0 +1,5 @@ +/** The upgrade card's `reason` for a payer whose plan does not include Chat. */ +export const CHAT_PLAN_REQUIRED_REASON = 'plan_required' as const + +/** Lead sentence of every refusal that withholds Chat from a plan. */ +export const CHAT_PLAN_REQUIRED_MESSAGE = 'Chat is included on paid plans.' diff --git a/apps/sim/lib/billing/core/chat-access.ts b/apps/sim/lib/billing/core/chat-access.ts new file mode 100644 index 00000000000..72645f0afd8 --- /dev/null +++ b/apps/sim/lib/billing/core/chat-access.ts @@ -0,0 +1,40 @@ +import { createLogger } from '@sim/logger' +import { getErrorMessage } from '@sim/utils/errors' +import { getCreditBalanceForEntity } from '@/lib/billing/credits/balance' +import { isFree } from '@/lib/billing/plan-helpers' +import type { UsageUpgradePayer } from '@/lib/billing/usage-upgrade' +import { isBillingEnabled, isHosted } from '@/lib/core/config/env-flags' + +const logger = createLogger('ChatAccess') + +/** + * Whether a payer may use Chat (talking to Sim). Every turn spends hosted model credits, so a + * hosted deployment that bills withholds Chat from a payer on the free plan — unless that payer + * holds prepaid credits, which an admin grants to customers billed outside a subscription. + * + * The plan is the payer's entitled subscription (active or past due), so a canceled + * plan reads as free; a blocked payer is refused by the usage check every caller runs next. + * + * This is the single verdict behind every Chat plan gate: run admission, chat titles, scheduled + * job dispatch, Copilot key admission, and the workspace host context the composer reads, so the + * composer never offers a send that admission will refuse. A balance that cannot be read admits + * the payer: a database blip must not lock a paying customer out, and a free payer stays bounded + * by its usage limit. + */ +export async function payerHasChatAccess(payer: UsageUpgradePayer): Promise { + if (!isHosted || !isBillingEnabled) return true + if (!isFree(payer.payerSubscription?.plan)) return true + try { + const balance = await getCreditBalanceForEntity( + payer.billingEntity.type, + payer.billingEntity.id + ) + return balance > 0 + } catch (error) { + logger.warn('Credit balance unavailable; admitting the payer to Chat', { + billingEntity: payer.billingEntity, + error: getErrorMessage(error), + }) + return true + } +} diff --git a/apps/sim/lib/billing/usage-upgrade.ts b/apps/sim/lib/billing/usage-upgrade.ts index 37fe1e6dbae..10f129f5927 100644 --- a/apps/sim/lib/billing/usage-upgrade.ts +++ b/apps/sim/lib/billing/usage-upgrade.ts @@ -1,6 +1,7 @@ import { createLogger } from '@sim/logger' import { getErrorMessage } from '@sim/utils/errors' import type { UsageUpgradePayload } from '@/lib/api/contracts/subscription' +import { CHAT_PLAN_REQUIRED_MESSAGE, CHAT_PLAN_REQUIRED_REASON } from '@/lib/billing/chat-plan' import type { AttributedUsageLimitsResult } from '@/lib/billing/core/attribution' import { getHighestPrioritySubscription } from '@/lib/billing/core/plan' import type { BillingEntity } from '@/lib/billing/core/usage-log' @@ -68,7 +69,21 @@ export async function resolveUsageUpgradePayload( return { reason: 'usage_limit', action: 'increase_limit', message } } +/** + * The card for a payer whose plan does not include Chat. Sim writes it at run admission, before + * any worker run exists, so it never crosses the worker billing protocol: its reason widens only + * the tag Sim renders, not the worker-shared card. + */ +export const CHAT_PLAN_REQUIRED_UPGRADE = { + reason: CHAT_PLAN_REQUIRED_REASON, + action: 'upgrade_plan', + message: `${CHAT_PLAN_REQUIRED_MESSAGE} Upgrade your plan to talk to Sim.`, +} as const + +/** Every card body the usage tag carries: the worker-shared card or Sim's Chat plan card. */ +export type UsageUpgradeTagPayload = UsageUpgradePayload | typeof CHAT_PLAN_REQUIRED_UPGRADE + /** The assistant text that renders {@link payload} as the usage card. */ -export function formatUsageUpgradeTag(payload: UsageUpgradePayload): string { +export function formatUsageUpgradeTag(payload: UsageUpgradeTagPayload): string { return `${JSON.stringify(payload)}` } diff --git a/apps/sim/lib/mothership/jobs/dispatch.ts b/apps/sim/lib/mothership/jobs/dispatch.ts index 90c37bf4448..91c02e34a6f 100644 --- a/apps/sim/lib/mothership/jobs/dispatch.ts +++ b/apps/sim/lib/mothership/jobs/dispatch.ts @@ -5,6 +5,8 @@ import { checkAttributedUsageLimits, resolveBillingAttribution, } from '@/lib/billing/core/attribution' +import { payerHasChatAccess } from '@/lib/billing/core/chat-access' +import { CHAT_PLAN_REQUIRED_UPGRADE } from '@/lib/billing/usage-upgrade' import { isHosted } from '@/lib/core/config/env-flags' import { acquireLock, releaseLock } from '@/lib/core/config/redis' import { asOrchestrationError } from '@/lib/core/orchestration/types' @@ -104,6 +106,11 @@ async function dispatchQueuedOccurrence(executionId: string): Promise ({ vi.mock('@/lib/mothership/request/tools/billing', () => ({ handleBillingLimitResponse: vi.fn(), + handleChatPlanRequiredResponse: vi.fn(), })) const mockRequestExplicitStreamAbort = vi.hoisted(() => vi.fn()) @@ -2310,6 +2311,17 @@ describe('runCopilotLifecycle', () => { expect(capturedExecContext?.userPermission).toBe('read') }) + /** A paid organization payer, which the free-plan Chat gate admits. */ + const paidOrganizationPayer = { + id: 'sub-org-1', + referenceId: 'org-1', + plan: 'team', + status: 'active', + seats: 2, + periodStart: null, + periodEnd: null, + } + it.each([ { workspaceId: 'ws-1', organizationId: undefined }, { workspaceId: undefined, organizationId: 'org-1' }, @@ -2324,7 +2336,7 @@ describe('runCopilotLifecycle', () => { start: '2026-07-01T00:00:00.000Z', end: '2026-08-01T00:00:00.000Z', }, - payerSubscription: null, + payerSubscription: paidOrganizationPayer, } setEnvFlags({ isHosted: true, isBillingEnabled: true }) mockEnv.COPILOT_API_KEY = 'sim-agent-key' @@ -2464,7 +2476,7 @@ describe('runCopilotLifecycle', () => { start: '2026-07-01T00:00:00.000Z', end: '2026-08-01T00:00:00.000Z', }, - payerSubscription: null, + payerSubscription: paidOrganizationPayer, } setEnvFlags({ isHosted: true, isBillingEnabled: true }) mockCheckAttributedUsageLimits.mockResolvedValue({ @@ -2517,7 +2529,7 @@ describe('runCopilotLifecycle', () => { billedAccountUserId: 'user-1', billingEntity: { type: 'organization' as const, id: 'org-1' }, billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2099-01-01T00:00:00.000Z' }, - payerSubscription: null, + payerSubscription: paidOrganizationPayer, } setEnvFlags({ isHosted: true, isBillingEnabled: true }) resetUsageGateCache() diff --git a/apps/sim/lib/mothership/request/lifecycle/run.ts b/apps/sim/lib/mothership/request/lifecycle/run.ts index 52db59650a2..135b1bc3f0c 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.ts @@ -14,6 +14,7 @@ import { checkAttributedUsageLimits, createAttributedBillingRequestEnvelope, } from '@/lib/billing/core/attribution' +import { payerHasChatAccess } from '@/lib/billing/core/chat-access' import { env } from '@/lib/core/config/env' import { isCopilotToolPermissionsEnabled, isHosted } from '@/lib/core/config/env-flags' import { withFileOwnerContext } from '@/lib/mothership/application/file-owner-context' @@ -71,7 +72,10 @@ import { requireToolCallStateResult, setTerminalToolCallState, } from '@/lib/mothership/request/tool-call-state' -import { handleBillingLimitResponse } from '@/lib/mothership/request/tools/billing' +import { + handleBillingLimitResponse, + handleChatPlanRequiredResponse, +} from '@/lib/mothership/request/tools/billing' import { executeToolAndReport, failPendingToolCall, @@ -572,22 +576,30 @@ export async function runCopilotLifecycle( let onCompleteStarted = false try { - // Hosted admission (usage limits) belongs HERE, at dispatch, with the attribution + // Hosted admission (plan and usage limits) belongs HERE, at dispatch, with the attribution // already in hand. The old path outsourced it to the Go backend's api-keys/validate // callback — a call the TS worker rightly never makes, so without this gate the // limit check simply never runs on the new path. On exceeded, the same synthetic // 402 UX as the mid-stream path renders the upgrade prompt, the backend is never // dispatched, and the shared verdict assembly below runs exactly as after a // mid-stream billing break. - const admission = + const admittedAttribution = isHosted && !isContinuation && execContext.billingAttribution - ? await checkAttributedUsageLimits( - assertBillingAttributionSnapshot(execContext.billingAttribution) - ) + ? assertBillingAttributionSnapshot(execContext.billingAttribution) + : undefined + // The free plan does not include Chat, on any surface that starts a run: such a payer's + // turn ends here with the upgrade card, before usage is read or the backend dispatched. + const chatPlanRequired = + admittedAttribution !== undefined && !(await payerHasChatAccess(admittedAttribution)) + const admission = + admittedAttribution && !chatPlanRequired + ? await checkAttributedUsageLimits(admittedAttribution) : { isExceeded: false as const } if (options.recovery && context.completionStatus) { // The worker terminal was already delivered before the relay died. // Rebuild persistence from that receipt without charging its usage twice. + } else if (chatPlanRequired) { + await handleChatPlanRequiredResponse(context, execContext, lifecycleOptions) } else if (admission.isExceeded) { await handleBillingLimitResponse( execContext.userId, diff --git a/apps/sim/lib/mothership/request/lifecycle/start.ts b/apps/sim/lib/mothership/request/lifecycle/start.ts index 62f87601a6b..d499c7b136f 100644 --- a/apps/sim/lib/mothership/request/lifecycle/start.ts +++ b/apps/sim/lib/mothership/request/lifecycle/start.ts @@ -13,6 +13,7 @@ import { resolveBillingAttribution, resolveOrganizationBillingAttribution, } from '@/lib/billing/core/attribution' +import { payerHasChatAccess } from '@/lib/billing/core/chat-access' import { isHosted } from '@/lib/core/config/env-flags' import { createRunSegment, recordRunBillingAdmission } from '@/lib/mothership/async-runs/repository' import { buildChatTitleContext } from '@/lib/mothership/chat/title-context' @@ -771,6 +772,8 @@ export async function requestChatTitle(params: { throw new Error('Title billing attribution does not match its actor and workspace') } + // A title is a Chat model call; a payer whose plan excludes Chat gets none. + if (!(await payerHasChatAccess(attribution))) return null const admission = await checkAttributedUsageLimits(attribution) if (admission.isExceeded) return null diff --git a/apps/sim/lib/mothership/request/tools/billing.ts b/apps/sim/lib/mothership/request/tools/billing.ts index 295bf69ddcc..93fd55d4b9b 100644 --- a/apps/sim/lib/mothership/request/tools/billing.ts +++ b/apps/sim/lib/mothership/request/tools/billing.ts @@ -1,5 +1,10 @@ import { createLogger } from '@sim/logger' -import { formatUsageUpgradeTag, resolveUsageUpgradePayload } from '@/lib/billing/usage-upgrade' +import { + CHAT_PLAN_REQUIRED_UPGRADE, + formatUsageUpgradeTag, + resolveUsageUpgradePayload, + type UsageUpgradeTagPayload, +} from '@/lib/billing/usage-upgrade' import { MothershipStreamV1CompletionStatus, MothershipStreamV1EventType, @@ -31,6 +36,24 @@ export async function handleBillingLimitResponse( scope?: 'actor' | 'payer' | 'member' ): Promise { const payload = await resolveUsageUpgradePayload(userId, execContext.billingAttribution, scope) + await endTurnWithUpgradeCard(payload, context, execContext, options) +} + +/** Ends a free payer's turn at admission with the card offering a plan that includes Chat. */ +export async function handleChatPlanRequiredResponse( + context: StreamingContext, + execContext: ExecutionContext, + options: OrchestratorOptions +): Promise { + await endTurnWithUpgradeCard(CHAT_PLAN_REQUIRED_UPGRADE, context, execContext, options) +} + +async function endTurnWithUpgradeCard( + payload: UsageUpgradeTagPayload, + context: StreamingContext, + execContext: ExecutionContext, + options: OrchestratorOptions +): Promise { const syntheticContent = formatUsageUpgradeTag(payload) // The card is this turn's terminal even when the refused leg follows one that already ended. context.streamComplete = false diff --git a/apps/sim/lib/workspaces/host-context.test.ts b/apps/sim/lib/workspaces/host-context.test.ts index 48314ac75d2..3f73b965fcf 100644 --- a/apps/sim/lib/workspaces/host-context.test.ts +++ b/apps/sim/lib/workspaces/host-context.test.ts @@ -104,6 +104,7 @@ describe('getWorkspaceHostContextForViewer', () => { organizationSearch, knowledgeMemberAccess: workspaceKnowledge, knowledgeSourceMirroredAccess: workspaceKnowledge, + chat: true, }) expect(mockIsKnowledgeMemberAccessAvailable).toHaveBeenCalledExactlyOnceWith({ organizationId: 'org-host', diff --git a/apps/sim/lib/workspaces/host-context.ts b/apps/sim/lib/workspaces/host-context.ts index df82e3d4829..a91ba7c8edf 100644 --- a/apps/sim/lib/workspaces/host-context.ts +++ b/apps/sim/lib/workspaces/host-context.ts @@ -1,5 +1,6 @@ import { cache } from 'react' import type { WorkspaceHostContext } from '@/lib/api/contracts/workspaces' +import { payerHasChatAccess } from '@/lib/billing/core/chat-access' import { getWorkspaceOwnerSubscriptionAccess } from '@/lib/billing/core/workspace-access' import { resolveDeploymentShape } from '@/lib/core/config/deployment-shape' import { isScopedCredentialGroupsAvailable } from '@/lib/credential-groups/scoped-availability' @@ -33,7 +34,7 @@ async function resolveWorkspaceHostContextForViewer( ? getOrganizationSettingsAccess(hostOrganizationId, userId) : Promise.resolve({ role: null, isMember: false, isAdmin: false }), ]) - const [credentialGroupsAvailable, knowledgeAccess, organizationSearch] = await Promise.all([ + const [credentialGroupsAvailable, knowledgeAccess, organizationSearch, chat] = await Promise.all([ hostOrganizationId ? isScopedCredentialGroupsAvailable({ kind: 'organization', @@ -44,6 +45,12 @@ async function resolveWorkspaceHostContextForViewer( hostOrganizationId && hostOrganizationAccess.isMember ? isKnowledgeMemberAccessAvailable({ organizationId: hostOrganizationId }) : Promise.resolve(false), + payerHasChatAccess({ + billingEntity: ownerBilling.organizationId + ? { type: 'organization', id: ownerBilling.organizationId } + : { type: 'user', id: access.workspace.billedAccountUserId }, + payerSubscription: { plan: ownerBilling.plan }, + }), ]) return { @@ -67,6 +74,7 @@ async function resolveWorkspaceHostContextForViewer( organizationSearch, knowledgeMemberAccess: knowledgeAccess.memberScoped, knowledgeSourceMirroredAccess: knowledgeAccess.sourceMirrored, + chat, }, deployment: resolveDeploymentShape(), } From f8a561d20a487ad69a39d92126e3d1c61f2ddab7 Mon Sep 17 00:00:00 2001 From: Waleed Date: Sun, 1 Nov 2026 12:36:19 -0800 Subject: [PATCH 4/4] fix(billing): fail closed on legacy Chat payer reads and cover the Chat plan gate (#9553) * fix(billing): fail closed on legacy Chat payer reads and cover the Chat plan gate * test(billing): cover the legacy-account subscription read failing closed * test(billing): pass the route context in the legacy-account failure case --- .../copilot/api-keys/validate/route.test.ts | 14 +++ .../api/copilot/api-keys/validate/route.ts | 4 +- .../billing/core/chat-access.integration.ts | 95 +++++++++++++++++++ .../mothership/request/lifecycle/run.test.ts | 65 ++++++++++--- 4 files changed, 162 insertions(+), 16 deletions(-) create mode 100644 apps/sim/lib/billing/core/chat-access.integration.ts diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts index e04b0aac04b..c8b4943b687 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts @@ -297,6 +297,20 @@ describe('POST /api/copilot/api-keys/validate billing protocols', () => { expect(mockCheckAttributedUsageLimits).not.toHaveBeenCalled() }) + it('fails account admission closed when its subscription cannot be read', async () => { + mockGetHighestPrioritySubscription.mockImplementationOnce( + async (_userId: string, options?: { onError?: 'throw' | 'return-null' }) => { + if (options?.onError === 'throw') throw new Error('database unavailable') + return null + } + ) + + const res = await POST(request(SELF_HOSTED_WORKSPACELESS_VALIDATE_BODY), createRouteContext({})) + + expect(res.status).toBe(500) + expect(mockCheckServerSideUsageLimits).not.toHaveBeenCalled() + }) + it('rejects markerless admission on hosted Sim', async () => { setEnvFlags({ isHosted: true }) const res = await POST(request(SELF_HOSTED_VALIDATE_BODY)) diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.ts b/apps/sim/app/api/copilot/api-keys/validate/route.ts index 1a159072395..77dccc60f97 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.ts @@ -296,7 +296,9 @@ async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise } } - const subscription = await getHighestPrioritySubscription(admission.userId) + const subscription = await getHighestPrioritySubscription(admission.userId, { + onError: 'throw', + }) const usage = await checkServerSideUsageLimits(admission.userId, subscription) return { isExceeded: usage.isExceeded, diff --git a/apps/sim/lib/billing/core/chat-access.integration.ts b/apps/sim/lib/billing/core/chat-access.integration.ts new file mode 100644 index 00000000000..804597d16d2 --- /dev/null +++ b/apps/sim/lib/billing/core/chat-access.integration.ts @@ -0,0 +1,95 @@ +/** + * The Chat plan verdict against real PostgreSQL credit balances: a hosted free payer is refused + * Chat until it holds prepaid credits, on both the organization and personal balance paths, and a + * paid payer is admitted without a balance. + */ +import { db } from '@sim/db' +import { organization, user, userStats } from '@sim/db/schema' +import { envFlagsMock } from '@sim/testing/mocks/env-flags.mock' +import { generateId } from '@sim/utils/id' +import { eq, inArray } from 'drizzle-orm' +import { afterAll, describe, expect, it, vi } from 'vitest' +import { payerHasChatAccess } from '@/lib/billing/core/chat-access' + +vi.mock('@/lib/core/config/env-flags', () => ({ + ...envFlagsMock, + isHosted: true, + isBillingEnabled: true, +})) + +const userIds: string[] = [] +const organizationIds: string[] = [] + +async function createUser(): Promise { + const id = `chat-access-user-${generateId()}` + const now = new Date() + await db.insert(user).values({ + id, + name: 'Chat Access Test', + email: `${id}@chat-access.test`, + emailVerified: true, + createdAt: now, + updatedAt: now, + }) + await db.insert(userStats).values({ id: generateId(), userId: id }) + userIds.push(id) + return id +} + +async function createOrganization(): Promise { + const id = `chat-access-org-${generateId()}` + await db.insert(organization).values({ id, name: 'Chat Access Org', slug: id }) + organizationIds.push(id) + return id +} + +afterAll(async () => { + if (organizationIds.length > 0) { + await db.delete(organization).where(inArray(organization.id, organizationIds)) + } + if (userIds.length > 0) await db.delete(user).where(inArray(user.id, userIds)) +}) + +describe('payerHasChatAccess', () => { + it('refuses a free organization payer until it holds prepaid credits', async () => { + const organizationId = await createOrganization() + const payer = { + billingEntity: { type: 'organization' as const, id: organizationId }, + payerSubscription: null, + } + + expect(await payerHasChatAccess(payer)).toBe(false) + + await db + .update(organization) + .set({ creditBalance: '25' }) + .where(eq(organization.id, organizationId)) + + expect(await payerHasChatAccess(payer)).toBe(true) + }) + + it('refuses a free personal payer until it holds prepaid credits', async () => { + const userId = await createUser() + const payer = { + billingEntity: { type: 'user' as const, id: userId }, + payerSubscription: { plan: 'free' }, + } + + expect(await payerHasChatAccess(payer)).toBe(false) + + await db.update(userStats).set({ creditBalance: '10' }).where(eq(userStats.userId, userId)) + + expect(await payerHasChatAccess(payer)).toBe(true) + }) + + it('admits a paid payer that holds no credits', async () => { + const userId = await createUser() + + expect( + await payerHasChatAccess({ + billingEntity: { type: 'user', id: userId }, + payerSubscription: { plan: 'pro_6000' }, + }) + ).toBe(true) + }) +}) diff --git a/apps/sim/lib/mothership/request/lifecycle/run.test.ts b/apps/sim/lib/mothership/request/lifecycle/run.test.ts index 46bab93eb0b..62d2ad4fe1c 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.test.ts @@ -24,7 +24,10 @@ import { import { generateId } from '@sim/utils/id' import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { scopeProviderToolCallId } from '@/lib/mothership/request/go/tool-call-identity' -import { handleBillingLimitResponse } from '@/lib/mothership/request/tools/billing' +import { + handleBillingLimitResponse, + handleChatPlanRequiredResponse, +} from '@/lib/mothership/request/tools/billing' import type { ExecutionContext, StreamingContext } from '@/lib/mothership/request/types' import { executeFunctionExecute } from '@/lib/mothership/tools/handlers/function-execute' import { executeRunCode } from '@/lib/mothership/tools/handlers/run-code' @@ -204,6 +207,11 @@ vi.mock('@/lib/mothership/request/tools/billing', () => ({ handleChatPlanRequiredResponse: vi.fn(), })) +const mockPayerHasChatAccess = vi.hoisted(() => vi.fn()) +vi.mock('@/lib/billing/core/chat-access', () => ({ + payerHasChatAccess: mockPayerHasChatAccess, +})) + const mockRequestExplicitStreamAbort = vi.hoisted(() => vi.fn()) vi.mock('@/lib/mothership/request/session/explicit-abort', () => ({ requestExplicitStreamAbort: mockRequestExplicitStreamAbort, @@ -276,6 +284,7 @@ describe('runCopilotLifecycle', () => { }) mockCreateRunSegment.mockResolvedValue({ status: 'active' }) mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false }) + mockPayerHasChatAccess.mockResolvedValue(true) continuationAuth.mockResolvedValue(undefined) mockEnv.COPILOT_API_KEY = undefined mockEnv.MSHIP_SYSPROMPT_OVERRIDE = undefined @@ -2311,17 +2320,6 @@ describe('runCopilotLifecycle', () => { expect(capturedExecContext?.userPermission).toBe('read') }) - /** A paid organization payer, which the free-plan Chat gate admits. */ - const paidOrganizationPayer = { - id: 'sub-org-1', - referenceId: 'org-1', - plan: 'team', - status: 'active', - seats: 2, - periodStart: null, - periodEnd: null, - } - it.each([ { workspaceId: 'ws-1', organizationId: undefined }, { workspaceId: undefined, organizationId: 'org-1' }, @@ -2336,7 +2334,7 @@ describe('runCopilotLifecycle', () => { start: '2026-07-01T00:00:00.000Z', end: '2026-08-01T00:00:00.000Z', }, - payerSubscription: paidOrganizationPayer, + payerSubscription: null, } setEnvFlags({ isHosted: true, isBillingEnabled: true }) mockEnv.COPILOT_API_KEY = 'sim-agent-key' @@ -2476,7 +2474,7 @@ describe('runCopilotLifecycle', () => { start: '2026-07-01T00:00:00.000Z', end: '2026-08-01T00:00:00.000Z', }, - payerSubscription: paidOrganizationPayer, + payerSubscription: null, } setEnvFlags({ isHosted: true, isBillingEnabled: true }) mockCheckAttributedUsageLimits.mockResolvedValue({ @@ -2505,6 +2503,43 @@ describe('runCopilotLifecycle', () => { expect(result.cancelled).not.toBe(true) }) + it('refuses dispatch at admission when the payer plan excludes Chat', async () => { + const billingAttribution = { + actorUserId: 'user-1', + workspaceId: 'ws-1', + organizationId: 'org-1', + billedAccountUserId: 'user-1', + billingEntity: { type: 'organization' as const, id: 'org-1' }, + billingPeriod: { + start: '2026-07-01T00:00:00.000Z', + end: '2026-08-01T00:00:00.000Z', + }, + payerSubscription: null, + } + setEnvFlags({ isHosted: true, isBillingEnabled: true }) + mockPayerHasChatAccess.mockResolvedValue(false) + + const result = await runCopilotLifecycle( + { message: 'hello', messageId: 'message-1' }, + { + userId: 'user-1', + workspaceId: 'ws-1', + chatId: 'chat-1', + executionId: 'execution-1', + runId: 'run-1', + simRequestId: 'request-1', + billingAttribution, + } + ) + + expect(mockPayerHasChatAccess).toHaveBeenCalledWith(billingAttribution) + expect(handleChatPlanRequiredResponse).toHaveBeenCalledOnce() + expect(mockCheckAttributedUsageLimits).not.toHaveBeenCalled() + expect(handleBillingLimitResponse).not.toHaveBeenCalled() + expect(mockRunStreamLoop).not.toHaveBeenCalled() + expect(result.cancelled).not.toBe(true) + }) + it('stops the worker run when the worker itself refuses a leg at the usage limit', async () => { mockRequestExplicitStreamAbort.mockResolvedValue({ settled: true }) mockRunStreamLoop.mockRejectedValueOnce(new BillingLimitError('user-1')) @@ -2529,7 +2564,7 @@ describe('runCopilotLifecycle', () => { billedAccountUserId: 'user-1', billingEntity: { type: 'organization' as const, id: 'org-1' }, billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2099-01-01T00:00:00.000Z' }, - payerSubscription: paidOrganizationPayer, + payerSubscription: null, } setEnvFlags({ isHosted: true, isBillingEnabled: true }) resetUsageGateCache()