mirror of
https://github.com/simstudioai/sim.git
synced 2026-09-24 15:45:35 +08:00
perf(server): stop calling our own API over HTTP during render, execution, and tool runs (#6660)
* fix(prefetch): parse the file-folder seed through its contract The audit found this key was the workspaceFilesKeys bug waiting to recur. The manager's record type and workspaceFileFolderSchema are two independent declarations that agree today by coincidence; the seed had no parse, so adding a column to one would have silently cached a shape a client fetch strips — and three of its fields are z.coerce.date(), the exact divergence that put ISO strings under the file-list key. Its sibling is immune because listWorkspaceFilesWithShares parses at the data layer. This does the same at the seed, and adds the shape-parity assertion the key never had. Verified falsifiable: removing the parse turns it red. Doing so also exposed the existing folder test as fixture-thin — a folder with only an id, which the contract rightly rejects — so it now uses a real row. Also points the credential block's fetchQuery at the exported staleTime constant instead of restating 60 * 1000; it was a fifth producer on that key free to drift from the four that share it. * fix(queries): stop a table cell edit throwing, and an upload gate failing shut Two functional bugs found auditing the query layer. patchCachedRows walked tableKeys.rowsRoot non-exact, but rowsRoot is a prefix: the find (search results) and write (pending writes) subtrees hang off it with non-paged shapes, and the updater's old.pages.map threw on them. It runs inside onMutate, so the whole cell edit rejected before reaching the server — reachable as soon as a find entry exists, i.e. after the user searches the table once. The sibling isDefaultOrderRowsQuery already excluded those subtrees and its docstring claimed they "never match"; that was only true of the sibling. Both now share one isRowListQueryKey helper so they cannot drift apart again. useCloudStorageConfigured combined staleTime: Infinity, retry: false, and the global retryOnMount: false on a workspace-independent key, so one transient failure left it errored for the tab's lifetime with no way back — navigating or switching workspace cannot change the key, and the upload path fails closed, so cloud-backed uploads stayed disabled until a full reload. useVoiceSettings carries the same three options and already escapes this with retryOnMount: true; this one now matches. Note: hooks/queries/workspace-files.test.tsx cannot load in a git worktree (pre-existing postcss resolution failure), so CI is the first place that file runs against this change. * perf(server): memoize the request-scoped workspace and entitlement reads The workspace row was read ~3x per workspace route and ~5x on settings, and the same Max-tier entitlement was resolved twice on one render. Memoization is deliberately partial. getWorkspaceWithOwner accepts a transaction and forUpdate, and live callers use both, so only the plain no-options read routes through the memo; a row read inside one caller's transaction or under a lock it alone holds can never be served to a later caller. includeArchived is part of the key so the two variants cannot alias. Three substitutions were considered and rejected as behavior changes, not optimizations: hostContext.ownerBilling resolves subscriptions differently from hasWorkspaceTierAccess and exposes no Max tier, so it cannot answer the Inbox/Sandbox gates; isOrganizationOnEnterprisePlan carries self-host short-circuits ownerBilling has no equivalent for; and widening WorkspaceHostContext to carry the full row would push owner and org ids onto the wire for every viewer to save a server-side read, since that type is a response contract rather than an internal struct. * fix(selectors): key CloudWatch lists by search, and stop a caller erasing the credential gate The CloudWatch log-group and log-stream selectors forwarded `search` into the request as `prefix` but left it out of the query key, so every keystroke resolved to the same fresh entry and no refetch fired. Server-side filtering was dead: a log group outside the first page could not be reached. An audit of all 69 selector definitions found these two and no others. useSelectorOptions resolved `args.enabled ?? definition.enabled(...)`, so a caller supplying its own gate replaced the definition's precondition rather than narrowing it. useSelectorDisplayName knows nothing about credentials, so a card holding a saved value with no credential context ran a query that could only reject. The two are now conjoined. The detail hooks keep the override deliberately — resolving one known id needs less context than listing, which their TSDoc already documents. The list-key fix has a test, proven to fail without it. The `enabled` change has none: loading use-selector-query pulls the selector registry and emcn CSS, which cannot resolve in a git worktree. * fix(queries): give optimistic rows collision-free ids generateTempId used Date.now(), so two rows created in the same millisecond shared an id and the first server response overwrote both — leaving one row duplicated and the other's real id lost until a refetch. Now uses generateId(), matching what the workflow mutations already do. Reachable by double-clicking create, or by any scripted or bulk create. Also documents the contract of fetchOAuthConnections, which reports an unknown connection state as disconnected. No consumer reads that field today — both read names and icons, and connection state comes from useWorkspaceCredentials — so letting the query reject would blank the suggested-action rows and drop the credential page to raw provider ids. The note is what stops a future consumer branching on it silently. * fix(queries): close six stale-data gaps found auditing the query layer Each was verified against the mutation that changes the data and the keys that expose it, not taken on report. - Workspace usage/credits were invalidated nowhere in the app. Six sites already refreshed subscriptionKeys after credits moved — post-run, post-wand, limit edits, upgrades, top-ups — and none touched workspace usage, so the credits chip and the run gate held their page-load values until a reload. Adds one shared invalidateWorkspaceUsage and calls it from all six. - Knowledge-base list doc counts went stale: document upload, delete, and bulk delete invalidated only the detail key, though the list carries docCount. - Plan switches that do not redirect refreshed only the host context, leaving subscription and credit state showing the previous plan. - The copilot tool-event handler invalidated a raw workflowKeys.list, which covers only the active scope and skips the selector prefix; it now uses the shared invalidateWorkflowLists like the other thirteen call sites. - scheduleKeys.byId was a strict prefix of scheduleKeys.schedule, so the two addressings aliased, and nothing invalidated byId. De-aliased and invalidated. Not changed: the CSV preview key already folds in the file version and storage key, so a content update addresses a different cache entry — version-in-key is the mechanism there, not a missing invalidation. Tests added for the usage and knowledge fixes, both proven to fail without them. The other four live in files that cannot load in a git worktree (pre-existing postcss resolution failure), so CI is where they first run. * improvement(queries): make the row-list prefix non-collidable, and share the usage refresh Two corrections from reviewing the previous commits. patchCachedRows was fixed with a predicate naming the sibling subtrees to skip — a denylist that rots the moment a fifth subtree is added under rowsRoot. The key factory already separated row lists under an 'infinite' segment; it just had no prefix accessor, so every caller reached for the parent and subtracted. Adding infiniteRowsRoot lets the walk be an allowlist by construction and deletes the predicate, the helper, and both docblocks explaining the subtraction. The searched-rows view is consequently no longer patched by a cell edit and is left to its own refetch — it holds a flat result, not pages. That is recorded on the function rather than left to be rediscovered. The delayed usage refresh was written out three times across two files, a duplication the previous commit enlarged rather than introduced. It is now one scheduleUsageRefresh beside the keys it invalidates, which also gives the bare 1000ms a name and one place to change it. * improvement(prefetch): budget the workspace file seed and narrow its columns * improvement(api): resolve credentials and checkpoint reverts in-process * improvement(executor): run router and evaluator provider calls in-process * improvement(queries): fix a second row-cache collision, and make the memoized workspace read actually dedupe * chore(test): type the evaluator provider-request helper instead of using any * improvement(perf): restore the parallel file read, and close the gaps a diff audit surfaced * improvement(perf): drop a duplicate authorization, parallelize the credential reads, and trim the comments
This commit is contained in:
@@ -24,6 +24,7 @@ vi.mock('@/lib/oauth/credential-service', () => ({
|
||||
|
||||
vi.mock('@/lib/auth/credential-access', () => ({
|
||||
authorizeCredentialUse: mockAuthorizeCredentialUse,
|
||||
authorizeCredentialUseForAuth: mockAuthorizeCredentialUse,
|
||||
}))
|
||||
|
||||
import { TokenServiceAccountValidationError } from '@/lib/credentials/token-service-accounts/errors'
|
||||
|
||||
@@ -11,17 +11,9 @@ import { authorizeCredentialUse } from '@/lib/auth/credential-access'
|
||||
import { AuthType, checkSessionOrInternalAuth } from '@/lib/auth/hybrid'
|
||||
import { generateRequestId } from '@/lib/core/utils/request'
|
||||
import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
|
||||
import { TokenServiceAccountValidationError } from '@/lib/credentials/token-service-accounts/errors'
|
||||
import {
|
||||
getCredential,
|
||||
getOAuthToken,
|
||||
refreshTokenIfNeeded,
|
||||
resolveOAuthAccountId,
|
||||
resolveServiceAccountToken,
|
||||
} from '@/lib/oauth/credential-service'
|
||||
import { extractSalesforceInstanceUrl, isSalesforceOAuthProviderId } from '@/lib/oauth/salesforce'
|
||||
import { getCredential, getOAuthToken } from '@/lib/oauth/credential-service'
|
||||
import { completeOAuthCredentialToken, resolveCredentialToken } from '@/lib/oauth/token-resolution'
|
||||
import { captureServerEvent } from '@/lib/posthog/server'
|
||||
import { extractZohoDeskBaseFromScope } from '@/tools/zoho_desk/host-allowlist'
|
||||
|
||||
export const dynamic = 'force-dynamic'
|
||||
|
||||
@@ -123,194 +115,25 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
|
||||
}
|
||||
}
|
||||
|
||||
if (!credentialId) {
|
||||
return NextResponse.json({ error: 'Credential ID is required' }, { status: 400 })
|
||||
}
|
||||
|
||||
const resolved = await resolveOAuthAccountId(credentialId)
|
||||
if (resolved?.credentialType === 'service_account' && resolved.credentialId) {
|
||||
const authz = await authorizeCredentialUse(request, {
|
||||
credentialId,
|
||||
workflowId: workflowId ?? undefined,
|
||||
requireWorkflowIdForInternal: false,
|
||||
callerUserId,
|
||||
})
|
||||
if (!authz.ok) {
|
||||
return NextResponse.json({ error: authz.error || 'Unauthorized' }, { status: 403 })
|
||||
}
|
||||
|
||||
const saActorId = authz.requesterUserId
|
||||
const saWorkspaceId = resolved.workspaceId ?? authz.workspaceId ?? null
|
||||
const emitServiceAccountAccess = () => {
|
||||
if (!saActorId) return
|
||||
recordAudit({
|
||||
workspaceId: saWorkspaceId,
|
||||
actorId: saActorId,
|
||||
action: AuditAction.CREDENTIAL_ACCESSED,
|
||||
resourceType: AuditResourceType.CREDENTIAL,
|
||||
resourceId: resolved.credentialId ?? credentialId,
|
||||
description: `Accessed service account credential for provider ${resolved.providerId ?? 'unknown'}`,
|
||||
metadata: {
|
||||
provider: resolved.providerId,
|
||||
credentialType: 'service_account',
|
||||
},
|
||||
request,
|
||||
})
|
||||
captureServerEvent(
|
||||
saActorId,
|
||||
'credential_used',
|
||||
{
|
||||
credential_type: 'service_account',
|
||||
provider_id: resolved.providerId ?? 'unknown',
|
||||
...(saWorkspaceId ? { workspace_id: saWorkspaceId } : {}),
|
||||
},
|
||||
saWorkspaceId ? { groups: { workspace: saWorkspaceId } } : undefined
|
||||
)
|
||||
}
|
||||
|
||||
try {
|
||||
const result = await resolveServiceAccountToken(
|
||||
resolved.credentialId,
|
||||
resolved.providerId,
|
||||
scopes ?? [],
|
||||
impersonateEmail
|
||||
)
|
||||
emitServiceAccountAccess()
|
||||
return NextResponse.json(
|
||||
{
|
||||
accessToken: result.accessToken,
|
||||
cloudId: result.cloudId,
|
||||
domain: result.domain,
|
||||
instanceUrl: result.instanceUrl,
|
||||
apiDomain: result.apiDomain,
|
||||
authStyle: result.authStyle,
|
||||
},
|
||||
{ status: 200 }
|
||||
)
|
||||
} catch (error) {
|
||||
logger.error(`[${requestId}] Service account token error:`, error)
|
||||
if (error instanceof TokenServiceAccountValidationError) {
|
||||
// Classified provider outages are infra failures, not bad credentials.
|
||||
if (error.code === 'provider_unavailable') {
|
||||
return NextResponse.json(
|
||||
{ error: 'Credential provider is temporarily unavailable' },
|
||||
{ status: 502 }
|
||||
)
|
||||
}
|
||||
// A stored host that no longer resolves is a configuration failure —
|
||||
// surface the code so runtime consumers can say "check the host"
|
||||
// instead of a generic auth error.
|
||||
if (error.code === 'site_not_found') {
|
||||
return NextResponse.json(
|
||||
{
|
||||
code: error.code,
|
||||
error: 'Credential host not found — reconnect the credential with a valid host',
|
||||
},
|
||||
{ status: 400 }
|
||||
)
|
||||
}
|
||||
// A revoked/rotated-away or misconfigured stored secret — surface the
|
||||
// code so runtime consumers can prompt to reconnect the credential
|
||||
// rather than showing a generic auth failure.
|
||||
if (error.code === 'invalid_credentials') {
|
||||
return NextResponse.json(
|
||||
{
|
||||
code: error.code,
|
||||
error: 'Credential rejected by the provider — reconnect the credential',
|
||||
},
|
||||
{ status: 401 }
|
||||
)
|
||||
}
|
||||
}
|
||||
return NextResponse.json({ error: 'Failed to get service account token' }, { status: 401 })
|
||||
}
|
||||
}
|
||||
|
||||
const authz = await authorizeCredentialUse(request, {
|
||||
const auth = await checkSessionOrInternalAuth(request, { requireWorkflowId: false })
|
||||
const result = await resolveCredentialToken(auth, {
|
||||
requestId,
|
||||
credentialId,
|
||||
workflowId: workflowId ?? undefined,
|
||||
requireWorkflowIdForInternal: false,
|
||||
scopes,
|
||||
impersonateEmail,
|
||||
callerUserId,
|
||||
auditRequest: request,
|
||||
})
|
||||
if (!authz.ok || !authz.credentialOwnerUserId) {
|
||||
return NextResponse.json({ error: authz.error || 'Unauthorized' }, { status: 403 })
|
||||
}
|
||||
|
||||
const resolvedCredentialId = authz.resolvedCredentialId || credentialId
|
||||
const credential = await getCredential(
|
||||
requestId,
|
||||
resolvedCredentialId,
|
||||
authz.credentialOwnerUserId
|
||||
)
|
||||
|
||||
if (!credential) {
|
||||
return NextResponse.json({ error: 'Credential not found' }, { status: 404 })
|
||||
}
|
||||
|
||||
const oauthActorId = authz.requesterUserId
|
||||
const oauthWorkspaceId = authz.workspaceId ?? null
|
||||
|
||||
try {
|
||||
const { accessToken } = await refreshTokenIfNeeded(
|
||||
requestId,
|
||||
credential,
|
||||
resolvedCredentialId
|
||||
)
|
||||
|
||||
if (oauthActorId) {
|
||||
recordAudit({
|
||||
workspaceId: oauthWorkspaceId,
|
||||
actorId: oauthActorId,
|
||||
action: AuditAction.CREDENTIAL_ACCESSED,
|
||||
resourceType: AuditResourceType.CREDENTIAL,
|
||||
resourceId: resolvedCredentialId,
|
||||
description: `Accessed OAuth credential for provider ${credential.providerId}`,
|
||||
metadata: {
|
||||
provider: credential.providerId,
|
||||
credentialType: 'oauth',
|
||||
},
|
||||
request,
|
||||
})
|
||||
captureServerEvent(
|
||||
oauthActorId,
|
||||
'credential_used',
|
||||
{
|
||||
credential_type: 'oauth',
|
||||
provider_id: credential.providerId,
|
||||
...(oauthWorkspaceId ? { workspace_id: oauthWorkspaceId } : {}),
|
||||
},
|
||||
oauthWorkspaceId ? { groups: { workspace: oauthWorkspaceId } } : undefined
|
||||
)
|
||||
}
|
||||
|
||||
const instanceUrl = isSalesforceOAuthProviderId(credential.providerId)
|
||||
? extractSalesforceInstanceUrl(credential.scope)
|
||||
: undefined
|
||||
|
||||
// Zoho Desk persists its data-center-specific REST base URL in the scope
|
||||
// string (derived from the token response api_domain) so callers never
|
||||
// assume a host. Surface it as apiDomain for tool param injection.
|
||||
let apiDomain: string | undefined
|
||||
if (credential.providerId === 'zoho-desk' && credential.scope) {
|
||||
// Use the shared extractor, not a local regex: it also enforces https +
|
||||
// the Zoho apex allowlist. This value is injected into EVERY tool call,
|
||||
// so an unvalidated host here would receive the OAuth token.
|
||||
apiDomain = extractZohoDeskBaseFromScope(credential.scope)
|
||||
}
|
||||
|
||||
if (!result.ok) {
|
||||
return NextResponse.json(
|
||||
{
|
||||
accessToken,
|
||||
idToken: credential.idToken || undefined,
|
||||
...(instanceUrl && { instanceUrl }),
|
||||
...(apiDomain && { apiDomain }),
|
||||
},
|
||||
{ status: 200 }
|
||||
{ ...(result.code ? { code: result.code } : {}), error: result.error },
|
||||
{ status: result.status }
|
||||
)
|
||||
} catch (error) {
|
||||
logger.error(`[${requestId}] Failed to refresh access token:`, error)
|
||||
return NextResponse.json({ error: 'Failed to refresh access token' }, { status: 401 })
|
||||
}
|
||||
|
||||
return NextResponse.json(result.token, { status: 200 })
|
||||
} catch (error) {
|
||||
logger.error(`[${requestId}] Error getting access token`, error)
|
||||
return NextResponse.json({ error: 'Internal server error' }, { status: 500 })
|
||||
@@ -366,70 +189,20 @@ export const GET = withRouteHandler(async (request: NextRequest) => {
|
||||
return NextResponse.json({ error: 'No access token available' }, { status: 400 })
|
||||
}
|
||||
|
||||
const actorId = authz.requesterUserId
|
||||
const workspaceId = authz.workspaceId ?? null
|
||||
const result = await completeOAuthCredentialToken({
|
||||
requestId,
|
||||
credential,
|
||||
resolvedCredentialId,
|
||||
actorId: authz.requesterUserId,
|
||||
workspaceId: authz.workspaceId ?? null,
|
||||
auditRequest: request,
|
||||
})
|
||||
|
||||
try {
|
||||
const { accessToken } = await refreshTokenIfNeeded(
|
||||
requestId,
|
||||
credential,
|
||||
resolvedCredentialId
|
||||
)
|
||||
|
||||
if (actorId) {
|
||||
recordAudit({
|
||||
workspaceId,
|
||||
actorId,
|
||||
action: AuditAction.CREDENTIAL_ACCESSED,
|
||||
resourceType: AuditResourceType.CREDENTIAL,
|
||||
resourceId: resolvedCredentialId,
|
||||
description: `Accessed OAuth credential for provider ${credential.providerId}`,
|
||||
metadata: {
|
||||
provider: credential.providerId,
|
||||
credentialType: 'oauth',
|
||||
},
|
||||
request,
|
||||
})
|
||||
captureServerEvent(
|
||||
actorId,
|
||||
'credential_used',
|
||||
{
|
||||
credential_type: 'oauth',
|
||||
provider_id: credential.providerId,
|
||||
...(workspaceId ? { workspace_id: workspaceId } : {}),
|
||||
},
|
||||
workspaceId ? { groups: { workspace: workspaceId } } : undefined
|
||||
)
|
||||
}
|
||||
|
||||
const instanceUrl = isSalesforceOAuthProviderId(credential.providerId)
|
||||
? extractSalesforceInstanceUrl(credential.scope)
|
||||
: undefined
|
||||
|
||||
// Zoho Desk persists its data-center-specific REST base URL in the scope
|
||||
// string (derived from the token response api_domain) so callers never
|
||||
// assume a host. Surface it as apiDomain for tool param injection.
|
||||
let apiDomain: string | undefined
|
||||
if (credential.providerId === 'zoho-desk' && credential.scope) {
|
||||
// Use the shared extractor, not a local regex: it also enforces https +
|
||||
// the Zoho apex allowlist. This value is injected into EVERY tool call,
|
||||
// so an unvalidated host here would receive the OAuth token.
|
||||
apiDomain = extractZohoDeskBaseFromScope(credential.scope)
|
||||
}
|
||||
|
||||
return NextResponse.json(
|
||||
{
|
||||
accessToken,
|
||||
idToken: credential.idToken || undefined,
|
||||
...(instanceUrl && { instanceUrl }),
|
||||
...(apiDomain && { apiDomain }),
|
||||
},
|
||||
{ status: 200 }
|
||||
)
|
||||
} catch (error) {
|
||||
logger.error(`[${requestId}] Failed to refresh access token:`, error)
|
||||
return NextResponse.json({ error: 'Failed to refresh access token' }, { status: 401 })
|
||||
if (!result.ok) {
|
||||
return NextResponse.json({ error: result.error }, { status: result.status })
|
||||
}
|
||||
|
||||
return NextResponse.json(result.token, { status: 200 })
|
||||
} catch (error) {
|
||||
logger.error(`[${requestId}] Error fetching access token`, error)
|
||||
return NextResponse.json({ error: 'Internal server error' }, { status: 500 })
|
||||
|
||||
@@ -17,12 +17,23 @@ import {
|
||||
import { NextRequest } from 'next/server'
|
||||
import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
|
||||
const { mockGetAccessibleCopilotChat } = vi.hoisted(() => ({
|
||||
const {
|
||||
mockGetAccessibleCopilotChat,
|
||||
mockParseWorkflowStateForPersistence,
|
||||
mockSaveWorkflowNormalizedState,
|
||||
} = vi.hoisted(() => ({
|
||||
mockGetAccessibleCopilotChat: vi.fn(),
|
||||
mockParseWorkflowStateForPersistence: vi.fn(),
|
||||
mockSaveWorkflowNormalizedState: vi.fn(),
|
||||
}))
|
||||
|
||||
vi.mock('@/lib/workflows/utils', () => workflowsUtilsMock)
|
||||
|
||||
vi.mock('@/lib/workflows/persistence/save-normalized-state', () => ({
|
||||
parseWorkflowStateForPersistence: mockParseWorkflowStateForPersistence,
|
||||
saveWorkflowNormalizedState: mockSaveWorkflowNormalizedState,
|
||||
}))
|
||||
|
||||
vi.mock('@/lib/copilot/chat/lifecycle', () => ({
|
||||
getAccessibleCopilotChat: mockGetAccessibleCopilotChat,
|
||||
getAccessibleCopilotChatAuth: mockGetAccessibleCopilotChat,
|
||||
@@ -38,13 +49,21 @@ describe('Copilot Checkpoints Revert API Route', () => {
|
||||
|
||||
authMockFns.mockGetSession.mockResolvedValue(null)
|
||||
|
||||
/** Authorization is the route's workflow read, so an allowed result always carries one. */
|
||||
workflowAuthzMockFns.mockAuthorizeWorkflowByWorkspacePermission.mockResolvedValue({
|
||||
allowed: true,
|
||||
status: 200,
|
||||
workflow: { id: 'b2c3d4e5-f6a7-4b89-a0d1-e2f3a4b5c6d7', workspaceId: 'ws-123' },
|
||||
})
|
||||
|
||||
mockGetAccessibleCopilotChat.mockResolvedValue({ id: 'chat-123', userId: 'user-123' })
|
||||
|
||||
mockParseWorkflowStateForPersistence.mockImplementation((value: unknown) => ({
|
||||
success: true,
|
||||
data: value,
|
||||
}))
|
||||
mockSaveWorkflowNormalizedState.mockResolvedValue({ success: true, warnings: [] })
|
||||
|
||||
global.fetch = vi.fn()
|
||||
|
||||
vi.spyOn(Date, 'now').mockReturnValue(1640995200000)
|
||||
@@ -184,7 +203,12 @@ describe('Copilot Checkpoints Revert API Route', () => {
|
||||
}
|
||||
|
||||
queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint])
|
||||
queueTableRows(schemaMock.workflow, [])
|
||||
/** Authorization performs the workflow read, so a missing workflow surfaces through it. */
|
||||
workflowAuthzMockFns.mockAuthorizeWorkflowByWorkspacePermission.mockResolvedValueOnce({
|
||||
allowed: false,
|
||||
status: 404,
|
||||
workflow: null,
|
||||
})
|
||||
|
||||
const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', {
|
||||
method: 'POST',
|
||||
@@ -197,6 +221,7 @@ describe('Copilot Checkpoints Revert API Route', () => {
|
||||
expect(response.status).toBe(404)
|
||||
const responseData = await response.json()
|
||||
expect(responseData.error).toBe('Workflow not found')
|
||||
expect(mockSaveWorkflowNormalizedState).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('should return 401 when workflow belongs to different user', async () => {
|
||||
@@ -220,6 +245,7 @@ describe('Copilot Checkpoints Revert API Route', () => {
|
||||
workflowAuthzMockFns.mockAuthorizeWorkflowByWorkspacePermission.mockResolvedValueOnce({
|
||||
allowed: false,
|
||||
status: 403,
|
||||
workflow: { id: 'b2c3d4e5-f6a7-4b89-a0d1-e2f3a4b5c6d7', workspaceId: 'ws-123' },
|
||||
})
|
||||
|
||||
const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', {
|
||||
@@ -297,24 +323,19 @@ describe('Copilot Checkpoints Revert API Route', () => {
|
||||
},
|
||||
})
|
||||
|
||||
// Verify fetch was called with correct parameters
|
||||
expect(global.fetch).toHaveBeenCalledWith(
|
||||
'http://localhost:3000/api/workflows/c3d4e5f6-a7b8-4c09-a1e2-f3a4b5c6d7e8/state',
|
||||
{
|
||||
method: 'PUT',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
Cookie: 'session=test-session',
|
||||
},
|
||||
body: JSON.stringify({
|
||||
expect(mockSaveWorkflowNormalizedState).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
workflowId: 'c3d4e5f6-a7b8-4c09-a1e2-f3a4b5c6d7e8',
|
||||
userId: 'user-123',
|
||||
state: {
|
||||
blocks: { block1: { type: 'start' } },
|
||||
edges: [{ from: 'block1', to: 'block2' }],
|
||||
loops: {},
|
||||
parallels: {},
|
||||
isDeployed: true,
|
||||
lastSaved: 1640995200000,
|
||||
}),
|
||||
}
|
||||
},
|
||||
})
|
||||
)
|
||||
})
|
||||
|
||||
@@ -452,7 +473,7 @@ describe('Copilot Checkpoints Revert API Route', () => {
|
||||
})
|
||||
})
|
||||
|
||||
it('should return 500 when state API call fails', async () => {
|
||||
it('should return 500 when the state write fails', async () => {
|
||||
setAuthenticated()
|
||||
|
||||
const mockCheckpoint = {
|
||||
@@ -470,9 +491,10 @@ describe('Copilot Checkpoints Revert API Route', () => {
|
||||
queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint])
|
||||
queueTableRows(schemaMock.workflow, [mockWorkflow])
|
||||
|
||||
;(global.fetch as any).mockResolvedValue({
|
||||
ok: false,
|
||||
text: () => Promise.resolve('State validation failed'),
|
||||
mockSaveWorkflowNormalizedState.mockResolvedValueOnce({
|
||||
success: false,
|
||||
status: 500,
|
||||
error: 'Failed to save workflow state',
|
||||
})
|
||||
|
||||
const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', {
|
||||
@@ -488,6 +510,36 @@ describe('Copilot Checkpoints Revert API Route', () => {
|
||||
expect(responseData.error).toBe('Failed to revert workflow to checkpoint')
|
||||
})
|
||||
|
||||
it('should return 500 when the checkpoint state fails validation', async () => {
|
||||
setAuthenticated()
|
||||
|
||||
const mockCheckpoint = {
|
||||
id: 'checkpoint-123',
|
||||
workflowId: 'a7b8c9d0-e1f2-4a34-b5c6-d7e8f9a0b1c2',
|
||||
userId: 'user-123',
|
||||
workflowState: { blocks: {}, edges: [] },
|
||||
}
|
||||
|
||||
queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint])
|
||||
queueTableRows(schemaMock.workflow, [{ id: mockCheckpoint.workflowId, userId: 'user-123' }])
|
||||
|
||||
mockParseWorkflowStateForPersistence.mockReturnValueOnce({
|
||||
success: false,
|
||||
error: { issues: [{ message: 'blocks: invalid' }] },
|
||||
})
|
||||
|
||||
const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ checkpointId: 'checkpoint-123' }),
|
||||
})
|
||||
|
||||
const response = await POST(req)
|
||||
|
||||
expect(response.status).toBe(500)
|
||||
expect(mockSaveWorkflowNormalizedState).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('should handle database errors during checkpoint lookup', async () => {
|
||||
setAuthenticated()
|
||||
|
||||
@@ -519,8 +571,9 @@ describe('Copilot Checkpoints Revert API Route', () => {
|
||||
}
|
||||
|
||||
dbChainMockFns.where.mockReturnValueOnce(Promise.resolve([mockCheckpoint]))
|
||||
dbChainMockFns.where.mockReturnValueOnce(
|
||||
Promise.reject(new Error('Database error during workflow lookup'))
|
||||
/** Authorization performs the workflow read, so a failed lookup surfaces through it. */
|
||||
workflowAuthzMockFns.mockAuthorizeWorkflowByWorkspacePermission.mockRejectedValueOnce(
|
||||
new Error('Database error during workflow lookup')
|
||||
)
|
||||
|
||||
const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', {
|
||||
@@ -536,7 +589,7 @@ describe('Copilot Checkpoints Revert API Route', () => {
|
||||
expect(responseData.error).toBe('Failed to revert to checkpoint')
|
||||
})
|
||||
|
||||
it('should handle fetch network errors', async () => {
|
||||
it('should handle unexpected errors from the state write', async () => {
|
||||
setAuthenticated()
|
||||
|
||||
const mockCheckpoint = {
|
||||
@@ -554,7 +607,7 @@ describe('Copilot Checkpoints Revert API Route', () => {
|
||||
queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint])
|
||||
queueTableRows(schemaMock.workflow, [mockWorkflow])
|
||||
|
||||
;(global.fetch as any).mockRejectedValue(new Error('Network error'))
|
||||
mockSaveWorkflowNormalizedState.mockRejectedValueOnce(new Error('Network error'))
|
||||
|
||||
const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', {
|
||||
method: 'POST',
|
||||
@@ -587,7 +640,7 @@ describe('Copilot Checkpoints Revert API Route', () => {
|
||||
expect(responseData.error).toBe('Failed to revert to checkpoint')
|
||||
})
|
||||
|
||||
it('should forward cookies to state API call', async () => {
|
||||
it('should apply the state in-process instead of re-authenticating over HTTP', async () => {
|
||||
setAuthenticated()
|
||||
|
||||
const mockCheckpoint = {
|
||||
@@ -623,17 +676,15 @@ describe('Copilot Checkpoints Revert API Route', () => {
|
||||
|
||||
await POST(req)
|
||||
|
||||
expect(global.fetch).toHaveBeenCalledWith(
|
||||
'http://localhost:3000/api/workflows/d0e1f2a3-b4c5-4d67-a8f9-a0b1c2d3e4f5/state',
|
||||
{
|
||||
method: 'PUT',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
Cookie: 'session=test-session; auth=token123',
|
||||
},
|
||||
body: expect.any(String),
|
||||
}
|
||||
expect(mockSaveWorkflowNormalizedState).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
workflowId: 'd0e1f2a3-b4c5-4d67-a8f9-a0b1c2d3e4f5',
|
||||
userId: 'user-123',
|
||||
})
|
||||
)
|
||||
for (const call of (global.fetch as any).mock.calls) {
|
||||
expect(String(call[0])).not.toContain('/state')
|
||||
}
|
||||
})
|
||||
|
||||
it('should handle missing cookies gracefully', async () => {
|
||||
@@ -673,16 +724,11 @@ describe('Copilot Checkpoints Revert API Route', () => {
|
||||
const response = await POST(req)
|
||||
|
||||
expect(response.status).toBe(200)
|
||||
expect(global.fetch).toHaveBeenCalledWith(
|
||||
'http://localhost:3000/api/workflows/e1f2a3b4-c5d6-4e78-a9a0-b1c2d3e4f5a6/state',
|
||||
{
|
||||
method: 'PUT',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
Cookie: '', // Empty string when no cookies
|
||||
},
|
||||
body: expect.any(String),
|
||||
}
|
||||
expect(mockSaveWorkflowNormalizedState).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
workflowId: 'e1f2a3b4-c5d6-4e78-a9a0-b1c2d3e4f5a6',
|
||||
userId: 'user-123',
|
||||
})
|
||||
)
|
||||
})
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import { db } from '@sim/db'
|
||||
import { workflowCheckpoints, workflow as workflowTable } from '@sim/db/schema'
|
||||
import { workflowCheckpoints } from '@sim/db/schema'
|
||||
import { createLogger } from '@sim/logger'
|
||||
import { authorizeWorkflowByWorkspacePermission } from '@sim/platform-authz/workflow'
|
||||
import { and, eq } from 'drizzle-orm'
|
||||
@@ -15,8 +15,11 @@ import {
|
||||
createRequestTracker,
|
||||
createUnauthorizedResponse,
|
||||
} from '@/lib/copilot/request/http'
|
||||
import { getInternalApiBaseUrl } from '@/lib/core/utils/urls'
|
||||
import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
|
||||
import {
|
||||
parseWorkflowStateForPersistence,
|
||||
saveWorkflowNormalizedState,
|
||||
} from '@/lib/workflows/persistence/save-normalized-state'
|
||||
import { isUuidV4 } from '@/executor/constants'
|
||||
|
||||
const logger = createLogger('CheckpointRevertAPI')
|
||||
@@ -62,21 +65,15 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
|
||||
return createNotFoundResponse('Checkpoint not found or access denied')
|
||||
}
|
||||
|
||||
const workflowData = await db
|
||||
.select()
|
||||
.from(workflowTable)
|
||||
.where(eq(workflowTable.id, checkpoint.workflowId))
|
||||
.then((rows) => rows[0])
|
||||
|
||||
if (!workflowData) {
|
||||
return createNotFoundResponse('Workflow not found')
|
||||
}
|
||||
|
||||
/** Authorization already loads the workflow, so its absence is the not-found signal. */
|
||||
const authorization = await authorizeWorkflowByWorkspacePermission({
|
||||
workflowId: checkpoint.workflowId,
|
||||
userId,
|
||||
action: 'write',
|
||||
})
|
||||
if (!authorization.workflow) {
|
||||
return createNotFoundResponse('Workflow not found')
|
||||
}
|
||||
if (!authorization.allowed) {
|
||||
return createUnauthorizedResponse()
|
||||
}
|
||||
@@ -121,28 +118,40 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
|
||||
return NextResponse.json({ error: 'Invalid workflow ID format' }, { status: 400 })
|
||||
}
|
||||
|
||||
const stateResponse = await fetch(
|
||||
`${getInternalApiBaseUrl()}/api/workflows/${checkpoint.workflowId}/state`,
|
||||
{
|
||||
method: 'PUT',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
Cookie: request.headers.get('Cookie') || '',
|
||||
},
|
||||
body: JSON.stringify(cleanedState),
|
||||
}
|
||||
)
|
||||
/**
|
||||
* The checkpoint blob is persisted JSONB, so it goes through the same
|
||||
* schema the PUT state contract applies before it is written back — the
|
||||
* validation the removed HTTP hop used to provide.
|
||||
*/
|
||||
const parsedState = parseWorkflowStateForPersistence(cleanedState)
|
||||
if (!parsedState.success) {
|
||||
logger.error(
|
||||
`[${tracker.requestId}] Checkpoint state failed validation`,
|
||||
parsedState.error.issues
|
||||
)
|
||||
return NextResponse.json(
|
||||
{ error: 'Failed to revert workflow to checkpoint' },
|
||||
{ status: 500 }
|
||||
)
|
||||
}
|
||||
|
||||
if (!stateResponse.ok) {
|
||||
const errorData = await stateResponse.text()
|
||||
logger.error(`[${tracker.requestId}] Failed to apply checkpoint state: ${errorData}`)
|
||||
const saveResult = await saveWorkflowNormalizedState({
|
||||
requestId: tracker.requestId,
|
||||
workflowId: checkpoint.workflowId,
|
||||
userId,
|
||||
state: parsedState.data,
|
||||
/** Already resolved above; re-deriving it would repeat 2-3 sequential reads. */
|
||||
authorization,
|
||||
})
|
||||
|
||||
if (!saveResult.success) {
|
||||
logger.error(`[${tracker.requestId}] Failed to apply checkpoint state: ${saveResult.error}`)
|
||||
return NextResponse.json(
|
||||
{ error: 'Failed to revert workflow to checkpoint' },
|
||||
{ status: 500 }
|
||||
)
|
||||
}
|
||||
|
||||
const result = await stateResponse.json()
|
||||
logger.info(
|
||||
`[${tracker.requestId}] Successfully reverted workflow ${checkpoint.workflowId} to checkpoint ${checkpointId}`
|
||||
)
|
||||
|
||||
@@ -1,28 +1,17 @@
|
||||
import { db } from '@sim/db'
|
||||
import { workflow } from '@sim/db/schema'
|
||||
import { createLogger } from '@sim/logger'
|
||||
import {
|
||||
assertWorkflowMutable,
|
||||
authorizeWorkflowByWorkspacePermission,
|
||||
WorkflowLockedError,
|
||||
} from '@sim/platform-authz/workflow'
|
||||
import { authorizeWorkflowByWorkspacePermission } from '@sim/platform-authz/workflow'
|
||||
import { toError } from '@sim/utils/errors'
|
||||
import { eq, sql } from 'drizzle-orm'
|
||||
import { type NextRequest, NextResponse } from 'next/server'
|
||||
import { putWorkflowNormalizedStateContract } from '@/lib/api/contracts/workflows'
|
||||
import { parseRequest } from '@/lib/api/server'
|
||||
import { checkSessionOrInternalAuth } from '@/lib/auth/hybrid'
|
||||
import { env } from '@/lib/core/config/env'
|
||||
import { generateRequestId } from '@/lib/core/utils/request'
|
||||
import { getSocketServerUrl } from '@/lib/core/utils/urls'
|
||||
import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
|
||||
import { extractAndPersistCustomTools } from '@/lib/workflows/persistence/custom-tools-persistence'
|
||||
import { prepareWorkflowStateForPersistence } from '@/lib/workflows/persistence/prepare-state'
|
||||
import {
|
||||
loadWorkflowFromNormalizedTables,
|
||||
saveWorkflowToNormalizedTables,
|
||||
} from '@/lib/workflows/persistence/utils'
|
||||
import type { BlockState, WorkflowState } from '@/stores/workflows/workflow/types'
|
||||
import { saveWorkflowNormalizedState } from '@/lib/workflows/persistence/save-normalized-state'
|
||||
import { loadWorkflowFromNormalizedTables } from '@/lib/workflows/persistence/utils'
|
||||
|
||||
const logger = createLogger('WorkflowStateAPI')
|
||||
|
||||
@@ -117,164 +106,29 @@ export const PUT = withRouteHandler(
|
||||
|
||||
const parsed = await parseRequest(putWorkflowNormalizedStateContract, request, context)
|
||||
if (!parsed.success) return parsed.response
|
||||
const state = parsed.data.body
|
||||
|
||||
const authorization = await authorizeWorkflowByWorkspacePermission({
|
||||
const result = await saveWorkflowNormalizedState({
|
||||
requestId,
|
||||
workflowId,
|
||||
userId,
|
||||
action: 'write',
|
||||
})
|
||||
const workflowData = authorization.workflow
|
||||
|
||||
if (!workflowData) {
|
||||
logger.warn(`[${requestId}] Workflow ${workflowId} not found for state update`)
|
||||
return NextResponse.json({ error: 'Workflow not found' }, { status: 404 })
|
||||
}
|
||||
|
||||
const canUpdate = authorization.allowed
|
||||
|
||||
if (!canUpdate) {
|
||||
logger.warn(
|
||||
`[${requestId}] User ${userId} denied permission to update workflow state ${workflowId}`
|
||||
)
|
||||
return NextResponse.json(
|
||||
{ error: authorization.message || 'Access denied' },
|
||||
{ status: authorization.status || 403 }
|
||||
)
|
||||
}
|
||||
|
||||
await assertWorkflowMutable(workflowId)
|
||||
|
||||
// Note: prior versions cross-checked that each variable's `workflowId`
|
||||
// equalled the path param. The write contract does not carry `workflowId`
|
||||
// per variable (the path param is the source of truth), so the check
|
||||
// is unreachable and was removed.
|
||||
|
||||
const { state: preparedState, warnings: preparationWarnings } =
|
||||
prepareWorkflowStateForPersistence({
|
||||
blocks: state.blocks as Record<string, BlockState>,
|
||||
edges: state.edges as WorkflowState['edges'],
|
||||
})
|
||||
|
||||
const workflowState = {
|
||||
...preparedState,
|
||||
lastSaved: state.lastSaved || Date.now(),
|
||||
isDeployed: state.isDeployed || false,
|
||||
deployedAt: state.deployedAt,
|
||||
}
|
||||
|
||||
const saveResult = await db.transaction(async (tx) => {
|
||||
await tx
|
||||
.select({ id: workflow.id })
|
||||
.from(workflow)
|
||||
.where(eq(workflow.id, workflowId))
|
||||
.limit(1)
|
||||
.for('update')
|
||||
|
||||
const result = await saveWorkflowToNormalizedTables(
|
||||
workflowId,
|
||||
workflowState as WorkflowState,
|
||||
tx
|
||||
)
|
||||
|
||||
if (!result.success) return result
|
||||
|
||||
// Update workflow's lastSynced timestamp and variables if provided
|
||||
const updateData: {
|
||||
lastSynced: Date
|
||||
updatedAt: Date
|
||||
variables?: typeof state.variables
|
||||
} = {
|
||||
lastSynced: new Date(),
|
||||
updatedAt: new Date(),
|
||||
}
|
||||
|
||||
// If variables are provided in the state, update them in the workflow record
|
||||
if (state.variables !== undefined) {
|
||||
updateData.variables = state.variables
|
||||
}
|
||||
|
||||
await tx.update(workflow).set(updateData).where(eq(workflow.id, workflowId))
|
||||
|
||||
return result
|
||||
state: parsed.data.body,
|
||||
})
|
||||
|
||||
if (!saveResult.success) {
|
||||
logger.error(
|
||||
`[${requestId}] Failed to save workflow ${workflowId} state:`,
|
||||
saveResult.error
|
||||
)
|
||||
if (!result.success) {
|
||||
return NextResponse.json(
|
||||
{ error: 'Failed to save workflow state', details: saveResult.error },
|
||||
{ status: 500 }
|
||||
{
|
||||
error: result.error,
|
||||
...(result.details !== undefined ? { details: result.details } : {}),
|
||||
},
|
||||
{ status: result.status }
|
||||
)
|
||||
}
|
||||
|
||||
// Extract and persist custom tools to database
|
||||
try {
|
||||
const workspaceId = workflowData.workspaceId
|
||||
if (workspaceId) {
|
||||
const { saved, errors } = await extractAndPersistCustomTools(
|
||||
workflowState,
|
||||
workspaceId,
|
||||
userId
|
||||
)
|
||||
|
||||
if (saved > 0) {
|
||||
logger.info(`[${requestId}] Persisted ${saved} custom tool(s) to database`, {
|
||||
workflowId,
|
||||
})
|
||||
}
|
||||
|
||||
if (errors.length > 0) {
|
||||
logger.warn(`[${requestId}] Some custom tools failed to persist`, {
|
||||
errors,
|
||||
workflowId,
|
||||
})
|
||||
}
|
||||
} else {
|
||||
logger.warn(
|
||||
`[${requestId}] Workflow has no workspaceId, skipping custom tools persistence`,
|
||||
{
|
||||
workflowId,
|
||||
}
|
||||
)
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error(`[${requestId}] Failed to persist custom tools`, { error, workflowId })
|
||||
}
|
||||
|
||||
const elapsed = Date.now() - startTime
|
||||
logger.info(`[${requestId}] Successfully saved workflow ${workflowId} state in ${elapsed}ms`)
|
||||
|
||||
try {
|
||||
const notifyResponse = await fetch(`${getSocketServerUrl()}/api/workflow-updated`, {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
'x-api-key': env.INTERNAL_API_SECRET,
|
||||
},
|
||||
body: JSON.stringify({ workflowId }),
|
||||
})
|
||||
|
||||
if (!notifyResponse.ok) {
|
||||
logger.warn(
|
||||
`[${requestId}] Failed to notify Socket.IO server about workflow ${workflowId} update`
|
||||
)
|
||||
}
|
||||
} catch (notificationError) {
|
||||
logger.warn(
|
||||
`[${requestId}] Error notifying Socket.IO server about workflow ${workflowId} update`,
|
||||
notificationError
|
||||
)
|
||||
}
|
||||
|
||||
return NextResponse.json({ success: true, warnings: preparationWarnings }, { status: 200 })
|
||||
return NextResponse.json({ success: true, warnings: result.warnings }, { status: 200 })
|
||||
} catch (error: any) {
|
||||
if (error instanceof WorkflowLockedError) {
|
||||
return NextResponse.json({ error: error.message }, { status: error.status })
|
||||
}
|
||||
|
||||
const elapsed = Date.now() - startTime
|
||||
logger.error(
|
||||
`[${requestId}] Error saving workflow ${workflowId} state after ${elapsed}ms`,
|
||||
|
||||
@@ -15,7 +15,7 @@ import { InviteLayout, InviteStatusCard } from '@/app/invite/components'
|
||||
import { useInvitationDetails } from '@/hooks/queries/invitations'
|
||||
import { organizationKeys } from '@/hooks/queries/organization'
|
||||
import { refreshSessionQuery } from '@/hooks/queries/session'
|
||||
import { subscriptionKeys } from '@/hooks/queries/subscription'
|
||||
import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys'
|
||||
import { workspaceKeys } from '@/hooks/queries/workspace'
|
||||
|
||||
const logger = createLogger('InviteById')
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import type { QueryClient } from '@tanstack/react-query'
|
||||
import { listWorkspaceFileFoldersContract } from '@/lib/api/contracts/workspace-file-folders'
|
||||
import { listWorkspaceFileFolders } from '@/lib/uploads/contexts/workspace/workspace-file-folder-manager'
|
||||
import { getWorkspaceHostContextForViewer } from '@/lib/workspaces/host-context'
|
||||
import { prefetchResourceListChrome } from '@/app/workspace/[workspaceId]/lib/prefetch-resource-list-chrome'
|
||||
@@ -15,10 +16,13 @@ import {
|
||||
* (scope `active`), so the browser paints populated on first render.
|
||||
*
|
||||
* The FILE LIST itself is deliberately not here: the sidebar reads it on every workspace route, so
|
||||
* it is prefetched by `prefetchWorkspaceSidebar` in the layout — the only boundary that renders
|
||||
* it is seeded by `prefetchWorkspaceSidebar` in the layout — the only boundary that renders
|
||||
* before the sidebar registers the query. Prefetching it again here would re-read it per request
|
||||
* and still not reach the server render (`HydrationBoundary` defers an already-seen query to an
|
||||
* effect, which SSR never runs). See the note on that entry.
|
||||
* effect, which SSR never runs). See the note on that entry. The layout declines to seed a
|
||||
* workspace whose file list exceeds its payload budget; recovering those here would mean
|
||||
* mirroring that budget check inversely, since an unconditional prefetch would re-read and
|
||||
* duplicate the entry for every workspace under the budget.
|
||||
*
|
||||
* Folders and the chrome reads all go through the data layer, shaped to their route contracts so a
|
||||
* hydrated entry matches a client fetch.
|
||||
@@ -40,7 +44,17 @@ export async function prefetchFilesBrowser(
|
||||
await Promise.all([
|
||||
queryClient.prefetchQuery({
|
||||
queryKey: workspaceFileFolderKeys.list(workspaceId, 'active'),
|
||||
queryFn: () => listWorkspaceFileFolders(workspaceId, { scope: 'active' }),
|
||||
/**
|
||||
* Parsed through the route's own response schema rather than seeded raw. The
|
||||
* manager's record type and `workspaceFileFolderSchema` are two independent
|
||||
* declarations that happen to agree today; without this parse, adding a column
|
||||
* to one silently seeds a shape a client fetch would have stripped — the exact
|
||||
* divergence that put ISO strings under `workspaceFilesKeys.list`.
|
||||
*/
|
||||
queryFn: async () => {
|
||||
const folders = await listWorkspaceFileFolders(workspaceId, { scope: 'active' })
|
||||
return listWorkspaceFileFoldersContract.response.schema.shape.folders.parse(folders)
|
||||
},
|
||||
staleTime: WORKSPACE_FILE_FOLDERS_STALE_TIME,
|
||||
}),
|
||||
prefetchResourceListChrome(queryClient, workspaceId, 'file', userId),
|
||||
|
||||
@@ -27,7 +27,7 @@ import {
|
||||
} from '@/app/workspace/[workspaceId]/home/hooks/stream/turn-model'
|
||||
import { deploymentKeys } from '@/hooks/queries/deployments'
|
||||
import { folderKeys } from '@/hooks/queries/utils/folder-keys'
|
||||
import { workflowKeys } from '@/hooks/queries/workflows'
|
||||
import { invalidateWorkflowLists } from '@/hooks/queries/utils/invalidate-workflow-lists'
|
||||
|
||||
type ToolEvent = Extract<PersistedStreamEventEnvelope, { type: 'tool' }>
|
||||
|
||||
@@ -58,7 +58,7 @@ function runToolResultSideEffects(ctx: StreamLoopContext, node: ToolNode): void
|
||||
if (deployedWorkflowId && typeof out?.isDeployed === 'boolean') {
|
||||
deps.queryClient.invalidateQueries({ queryKey: deploymentKeys.info(deployedWorkflowId) })
|
||||
deps.queryClient.invalidateQueries({ queryKey: deploymentKeys.versions(deployedWorkflowId) })
|
||||
deps.queryClient.invalidateQueries({ queryKey: workflowKeys.list(deps.workspaceId) })
|
||||
void invalidateWorkflowLists(deps.queryClient, deps.workspaceId)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -66,7 +66,9 @@ function runToolResultSideEffects(ctx: StreamLoopContext, node: ToolNode): void
|
||||
deps.queryClient.invalidateQueries({ queryKey: folderKeys.list(deps.workspaceId) })
|
||||
}
|
||||
if (WORKFLOW_MUTATION_TOOL_NAMES.has(name) && isSuccess) {
|
||||
deps.queryClient.invalidateQueries({ queryKey: workflowKeys.list(deps.workspaceId) })
|
||||
// `rm` archives, so the archived list moves too — and the shared helper also
|
||||
// refreshes the workflow selector lists that `@`-mentions and pickers read.
|
||||
void invalidateWorkflowLists(deps.queryClient, deps.workspaceId, ['active', 'archived'])
|
||||
}
|
||||
|
||||
const extractedResources =
|
||||
|
||||
@@ -209,7 +209,11 @@ export function useKnowledgeUpload(options: UseKnowledgeUploadOptions = {}) {
|
||||
setUploadProgress((prev) => ({ ...prev, stage: 'processing' }))
|
||||
logger.info(`Successfully started processing ${uploadedDocuments.length} documents`)
|
||||
|
||||
await queryClient.invalidateQueries({ queryKey: knowledgeKeys.detail(knowledgeBaseId) })
|
||||
await Promise.all([
|
||||
queryClient.invalidateQueries({ queryKey: knowledgeKeys.detail(knowledgeBaseId) }),
|
||||
/** The knowledge-base list rows carry `docCount`, so an upload changes them too. */
|
||||
queryClient.invalidateQueries({ queryKey: knowledgeKeys.lists() }),
|
||||
])
|
||||
|
||||
return uploadedDocuments
|
||||
} catch (err) {
|
||||
|
||||
@@ -97,7 +97,10 @@ vi.mock('@sim/emcn', () => ({
|
||||
|
||||
import { prefetchFilesBrowser } from '@/app/workspace/[workspaceId]/files/prefetch'
|
||||
import { prefetchKnowledgeBases } from '@/app/workspace/[workspaceId]/knowledge/prefetch'
|
||||
import { prefetchWorkspaceSidebar } from '@/app/workspace/[workspaceId]/prefetch'
|
||||
import {
|
||||
prefetchWorkspaceSidebar,
|
||||
WORKSPACE_FILE_SEED_MAX,
|
||||
} from '@/app/workspace/[workspaceId]/prefetch'
|
||||
import { prefetchTables } from '@/app/workspace/[workspaceId]/tables/prefetch'
|
||||
import { folderKeys } from '@/hooks/queries/utils/folder-keys'
|
||||
import { knowledgeKeys } from '@/hooks/queries/utils/knowledge-keys'
|
||||
@@ -290,17 +293,67 @@ describe('workspace list prefetches', () => {
|
||||
})
|
||||
})
|
||||
describe('prefetchFilesBrowser', () => {
|
||||
/**
|
||||
* The sibling `workspaceFilesKeys.list` once held ISO strings from one producer and
|
||||
* `Date`s from another because a seed skipped the contract parse. This key is fed by
|
||||
* a manager whose record type and the contract schema are independent declarations,
|
||||
* so the parse — and this assertion — are what stop that recurring here.
|
||||
*/
|
||||
it('seeds the shape a client fetch caches, not the raw manager row', async () => {
|
||||
mockListWorkspaceFileFolders.mockResolvedValue([
|
||||
{
|
||||
id: 'folder-1',
|
||||
workspaceId: WORKSPACE_ID,
|
||||
userId: USER_ID,
|
||||
name: 'Docs',
|
||||
parentId: null,
|
||||
path: '/Docs',
|
||||
sortOrder: 0,
|
||||
deletedAt: null,
|
||||
createdAt: '2026-01-01T00:00:00.000Z',
|
||||
updatedAt: '2026-01-02T00:00:00.000Z',
|
||||
serverOnlyColumn: 'should-be-stripped',
|
||||
},
|
||||
])
|
||||
const client = makeClient()
|
||||
|
||||
await prefetchFilesBrowser(client, WORKSPACE_ID, USER_ID)
|
||||
|
||||
const [cached] = client.getQueryData(
|
||||
workspaceFileFolderKeys.list(WORKSPACE_ID, 'active')
|
||||
) as Array<Record<string, unknown>>
|
||||
expect(cached.createdAt).toBeInstanceOf(Date)
|
||||
expect(cached.updatedAt).toBeInstanceOf(Date)
|
||||
expect(cached).not.toHaveProperty('serverOnlyColumn')
|
||||
})
|
||||
|
||||
it('primes the folder key the client hook reads', async () => {
|
||||
const folders = [{ id: 'folder-1' }]
|
||||
const folders = [
|
||||
{
|
||||
id: 'folder-1',
|
||||
workspaceId: WORKSPACE_ID,
|
||||
userId: USER_ID,
|
||||
name: 'Docs',
|
||||
parentId: null,
|
||||
path: '/Docs',
|
||||
sortOrder: 0,
|
||||
deletedAt: null,
|
||||
createdAt: '2026-01-01T00:00:00.000Z',
|
||||
updatedAt: '2026-01-02T00:00:00.000Z',
|
||||
},
|
||||
]
|
||||
mockListWorkspaceFileFolders.mockResolvedValue(folders)
|
||||
const client = makeClient()
|
||||
|
||||
await prefetchFilesBrowser(client, WORKSPACE_ID, USER_ID)
|
||||
|
||||
expect(mockListWorkspaceFileFolders).toHaveBeenCalledWith(WORKSPACE_ID, { scope: 'active' })
|
||||
expect(client.getQueryData(workspaceFileFolderKeys.list(WORKSPACE_ID, 'active'))).toEqual(
|
||||
folders
|
||||
)
|
||||
/** Shape parity is asserted by the sibling test; this one pins the key and the args. */
|
||||
expect(
|
||||
client.getQueryData(workspaceFileFolderKeys.list(WORKSPACE_ID, 'active')) as Array<{
|
||||
id: string
|
||||
}>
|
||||
).toHaveLength(folders.length)
|
||||
})
|
||||
|
||||
/**
|
||||
@@ -460,6 +513,52 @@ describe('workspace list prefetches', () => {
|
||||
expect(client.getQueryData(workspaceKeys.list('active'))).toBeUndefined()
|
||||
})
|
||||
|
||||
/**
|
||||
* The file list is seeded on every workspace route, so it is the one entry whose size
|
||||
* scales with a workspace's content on routes that never read it. The budget is passed
|
||||
* down rather than applied here, so the read can stop before the share join.
|
||||
*/
|
||||
it('seeds the file list, bounded by the document payload budget', async () => {
|
||||
const files = [{ id: 'file-1', name: 'a.txt' }]
|
||||
mockListWorkspaceFilesWithShares.mockResolvedValue(files)
|
||||
const client = makeClient()
|
||||
|
||||
await prefetchWorkspaceSidebar(client, WORKSPACE_ID, USER_ID, HOST_CONTEXT, null)
|
||||
|
||||
expect(mockListWorkspaceFilesWithShares).toHaveBeenCalledWith(WORKSPACE_ID, 'active', {
|
||||
maxRows: WORKSPACE_FILE_SEED_MAX,
|
||||
/** A failed read must reach the catch, not degrade to a cached empty list. */
|
||||
throwOnError: true,
|
||||
})
|
||||
expect(client.getQueryData(workspaceFilesKeys.list(WORKSPACE_ID, 'active'))).toEqual(files)
|
||||
})
|
||||
|
||||
/**
|
||||
* The load-bearing half of the budget: a workspace over it seeds NOTHING rather than the
|
||||
* prefix that was read. The sidebar search filters this list client-side and the Files
|
||||
* browser renders it as the workspace's files, so a truncated seed would silently hide
|
||||
* files — the client fetch must reach the route for the complete list instead.
|
||||
*/
|
||||
it('seeds nothing when the workspace exceeds the budget', async () => {
|
||||
mockListWorkspaceFilesWithShares.mockResolvedValue(null)
|
||||
const client = makeClient()
|
||||
|
||||
await prefetchWorkspaceSidebar(client, WORKSPACE_ID, USER_ID, HOST_CONTEXT, null)
|
||||
|
||||
expect(client.getQueryData(workspaceFilesKeys.list(WORKSPACE_ID, 'active'))).toBeUndefined()
|
||||
})
|
||||
|
||||
/** A failed file read is an optimization loss, not a render failure. */
|
||||
it('does not throw when the file read rejects, and seeds no files', async () => {
|
||||
mockListWorkspaceFilesWithShares.mockRejectedValue(new Error('500'))
|
||||
const client = makeClient()
|
||||
|
||||
await expect(
|
||||
prefetchWorkspaceSidebar(client, WORKSPACE_ID, USER_ID, HOST_CONTEXT, null)
|
||||
).resolves.toBeUndefined()
|
||||
expect(client.getQueryData(workspaceFilesKeys.list(WORKSPACE_ID, 'active'))).toBeUndefined()
|
||||
})
|
||||
|
||||
/** Guards the mismatch check that keeps one workspace's data out of another's cache. */
|
||||
it('seeds nothing when the host context is for a different workspace', async () => {
|
||||
mockListWorkspacesForViewer.mockResolvedValue(LIST_PAYLOAD)
|
||||
|
||||
@@ -25,10 +25,7 @@ import { workflowKeys } from '@/hooks/queries/utils/workflow-keys'
|
||||
import { mapWorkflow, WORKFLOW_LIST_STALE_TIME } from '@/hooks/queries/utils/workflow-list-query'
|
||||
import { normalizeWorkspacesResponse } from '@/hooks/queries/utils/workspace-list-query'
|
||||
import { WORKSPACE_PERMISSIONS_STALE_TIME, workspaceKeys } from '@/hooks/queries/workspace'
|
||||
import {
|
||||
WORKSPACE_FILES_LIST_STALE_TIME,
|
||||
workspaceFilesKeys,
|
||||
} from '@/hooks/queries/workspace-files'
|
||||
import { workspaceFilesKeys } from '@/hooks/queries/workspace-files'
|
||||
import {
|
||||
WORKSPACE_HOST_CONTEXT_STALE_TIME,
|
||||
workspaceHostKeys,
|
||||
@@ -95,6 +92,46 @@ async function seedWorkspaceList(
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* How many files the layout is willing to inline into the document. Seeded on EVERY
|
||||
* workspace route, so at ~500 bytes of JSON per file this budgets the entry at ~150 KB.
|
||||
*
|
||||
* A workspace above the budget seeds NOTHING rather than a prefix: the sidebar filters
|
||||
* this list client-side, so a truncated seed would silently hide files.
|
||||
*/
|
||||
export const WORKSPACE_FILE_SEED_MAX = 300
|
||||
|
||||
/**
|
||||
* Seeds the workspace's file list, which sidebar chrome registers on EVERY workspace
|
||||
* route. It must be seeded HERE, not by the Files pages: `HydrationBoundary` defers a
|
||||
* query the cache has already seen to a `useEffect`, which SSR never runs.
|
||||
*
|
||||
* Seeded rather than prefetched so it can decline to create an entry at all above
|
||||
* {@link WORKSPACE_FILE_SEED_MAX} — `prefetchQuery` always creates one, and a partial
|
||||
* entry would be read as the whole list. Parsed through the route's response contract.
|
||||
*/
|
||||
async function seedWorkspaceFiles(queryClient: QueryClient, workspaceId: string): Promise<void> {
|
||||
try {
|
||||
const files = await listWorkspaceFilesWithShares(workspaceId, 'active', {
|
||||
maxRows: WORKSPACE_FILE_SEED_MAX,
|
||||
/**
|
||||
* A failed read must reach the catch below, not degrade to an empty list: seeding
|
||||
* `[]` would cache "this workspace has no files" as authoritative for the entry's
|
||||
* lifetime, which is worse than seeding nothing and letting the client fetch.
|
||||
*/
|
||||
throwOnError: true,
|
||||
})
|
||||
if (!files) return
|
||||
queryClient.setQueryData(workspaceFilesKeys.list(workspaceId, 'active'), files)
|
||||
} catch (error) {
|
||||
/** Optimization only: the client fetch reaches the route instead. Logged so drift between
|
||||
* this read and the contract's response schema doesn't degrade silently into a waterfall. */
|
||||
logger.warn('Workspace file list seed failed; client will fetch', {
|
||||
error: getErrorMessage(error),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Prefetches the sidebar's workflow, chat, folder, workspace-permissions,
|
||||
* workspace, and viewer-profile reads for a workspace and stores them under the
|
||||
@@ -153,22 +190,7 @@ export async function prefetchWorkspaceSidebar(
|
||||
]
|
||||
: []),
|
||||
prefetchResourceFolders(queryClient, workspaceId, 'workflow', userId),
|
||||
/**
|
||||
* The sidebar reads the workspace's files for its search modal, on EVERY workspace route — so this
|
||||
* query is registered by sidebar chrome before any page renders. That ordering is why it has to be
|
||||
* prefetched HERE and not only by the Files pages: `HydrationBoundary` hydrates a query the cache
|
||||
* has already seen from a `useEffect`, which never runs during SSR, so a page-level boundary can
|
||||
* only ever hand this entry to the client. Seeding it with the layout's own boundary — the first
|
||||
* one to render — is what lets the server paint the Files browser and the open file's header
|
||||
* populated instead of shipping a spinner and resolving it a beat later on the client.
|
||||
*
|
||||
* Same key + shape as {@link prefetchFilesBrowser}, so whichever runs is a no-op for the other.
|
||||
*/
|
||||
queryClient.prefetchQuery({
|
||||
queryKey: workspaceFilesKeys.list(workspaceId, 'active'),
|
||||
queryFn: () => listWorkspaceFilesWithShares(workspaceId, 'active'),
|
||||
staleTime: WORKSPACE_FILES_LIST_STALE_TIME,
|
||||
}),
|
||||
seedWorkspaceFiles(queryClient, workspaceId),
|
||||
queryClient.prefetchQuery({
|
||||
queryKey: workspaceKeys.permissions(workspaceId),
|
||||
queryFn: () =>
|
||||
|
||||
@@ -9,6 +9,8 @@ import type { WorkspaceHostContext } from '@/lib/api/contracts/workspaces'
|
||||
import { useSubscriptionUpgrade } from '@/lib/billing/client/upgrade'
|
||||
import { CREDIT_TIERS } from '@/lib/billing/constants'
|
||||
import { getPlanTierCredits, isEnterprise, isFree, isPro, isTeam } from '@/lib/billing/plan-helpers'
|
||||
import { invalidateWorkspaceUsage } from '@/hooks/queries/utils/invalidate-usage'
|
||||
import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys'
|
||||
import { workspaceHostKeys } from '@/hooks/queries/workspace-host'
|
||||
|
||||
const PRO_TIER = CREDIT_TIERS[0]
|
||||
@@ -89,8 +91,22 @@ export function useUpgradeState({
|
||||
}
|
||||
}, [ownerBilling.billingInterval, subscription.isPaid])
|
||||
|
||||
const refreshHostContext = useCallback(
|
||||
() => queryClient.invalidateQueries({ queryKey: workspaceHostKeys.detail(workspaceId) }),
|
||||
/**
|
||||
* A non-redirect plan switch settles server-side immediately, so every read that
|
||||
* describes the plan has to be refetched — the host context the page renders from,
|
||||
* the subscription/usage reads the billing surfaces share, the proration invoice the
|
||||
* switch just produced, and the workspace credit availability that drives the credits
|
||||
* chip and the run gate.
|
||||
*/
|
||||
const refreshBillingState = useCallback(
|
||||
() =>
|
||||
Promise.all([
|
||||
queryClient.invalidateQueries({ queryKey: workspaceHostKeys.detail(workspaceId) }),
|
||||
queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }),
|
||||
queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() }),
|
||||
queryClient.invalidateQueries({ queryKey: subscriptionKeys.invoicesAll() }),
|
||||
invalidateWorkspaceUsage(queryClient),
|
||||
]),
|
||||
[queryClient, workspaceId]
|
||||
)
|
||||
|
||||
@@ -123,9 +139,9 @@ export function useUpgradeState({
|
||||
await requestJson(billingSwitchPlanContract, {
|
||||
body: { targetPlanName: subscription.plan, interval, workspaceId },
|
||||
})
|
||||
await refreshHostContext()
|
||||
await refreshBillingState()
|
||||
},
|
||||
[isLegacyPlan, refreshHostContext, subscription.plan, workspaceId]
|
||||
[isLegacyPlan, refreshBillingState, subscription.plan, workspaceId]
|
||||
)
|
||||
|
||||
const currentCredits = getPlanTierCredits(subscription.plan)
|
||||
@@ -154,11 +170,11 @@ export function useUpgradeState({
|
||||
workspaceId,
|
||||
},
|
||||
})
|
||||
await refreshHostContext()
|
||||
await refreshBillingState()
|
||||
} catch (e) {
|
||||
toast.error(getErrorMessage(e, 'Failed to upgrade'))
|
||||
}
|
||||
}, [subscription.isTeam, isAnnual, refreshHostContext, workspaceId])
|
||||
}, [subscription.isTeam, isAnnual, refreshBillingState, workspaceId])
|
||||
|
||||
const onUpgradeToOtherTier = useCallback(async () => {
|
||||
const onMax =
|
||||
@@ -170,11 +186,11 @@ export function useUpgradeState({
|
||||
await requestJson(billingSwitchPlanContract, {
|
||||
body: { targetPlanName, workspaceId },
|
||||
})
|
||||
await refreshHostContext()
|
||||
await refreshBillingState()
|
||||
} catch (e) {
|
||||
toast.error(getErrorMessage(e, 'Failed to switch plan'))
|
||||
}
|
||||
}, [subscription.plan, subscription.isTeam, refreshHostContext, workspaceId])
|
||||
}, [subscription.plan, subscription.isTeam, refreshBillingState, workspaceId])
|
||||
|
||||
return {
|
||||
isLoading: false,
|
||||
|
||||
@@ -10,7 +10,7 @@ import { wandGenerateStreamContract } from '@/lib/api/contracts'
|
||||
import { readSSEStream } from '@/lib/core/utils/sse'
|
||||
import { shouldStripCodeFences, stripCodeFences } from '@/lib/wand/strip-code-fences'
|
||||
import type { GenerationType } from '@/blocks/types'
|
||||
import { subscriptionKeys } from '@/hooks/queries/subscription'
|
||||
import { scheduleUsageRefresh } from '@/hooks/queries/utils/invalidate-usage'
|
||||
import { useSettingsNavigation } from '@/hooks/use-settings-navigation'
|
||||
import { useWorkflowRegistry } from '@/stores/workflows/registry/store'
|
||||
|
||||
@@ -300,9 +300,7 @@ export function useWand({
|
||||
strippedFences: generatedContent !== accumulatedContent,
|
||||
})
|
||||
|
||||
setTimeout(() => {
|
||||
queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() })
|
||||
}, 1000)
|
||||
scheduleUsageRefresh(queryClient)
|
||||
} catch (error: any) {
|
||||
if (error.name === 'AbortError') {
|
||||
logger.debug('Wand generation cancelled')
|
||||
|
||||
-4
@@ -191,10 +191,6 @@ vi.mock('@/executor/utils/start-block', () => ({
|
||||
coerceValue: (_type: string, value: unknown) => value,
|
||||
}))
|
||||
|
||||
vi.mock('@/hooks/queries/subscription', () => ({
|
||||
subscriptionKeys: { users: () => ['subscription', 'users'] },
|
||||
}))
|
||||
|
||||
vi.mock('@/hooks/queries/utils/workflow-cache', () => ({
|
||||
getWorkflows: () => [],
|
||||
}))
|
||||
|
||||
@@ -61,7 +61,7 @@ import type { SerializableExecutionState } from '@/executor/execution/types'
|
||||
import type { BlockLog, BlockState, ExecutionResult, StreamingExecution } from '@/executor/types'
|
||||
import { hasExecutionResult } from '@/executor/utils/errors'
|
||||
import { coerceValue } from '@/executor/utils/start-block'
|
||||
import { subscriptionKeys } from '@/hooks/queries/subscription'
|
||||
import { scheduleUsageRefresh } from '@/hooks/queries/utils/invalidate-usage'
|
||||
import { getWorkflows } from '@/hooks/queries/utils/workflow-cache'
|
||||
import {
|
||||
isExecutionStreamHttpError,
|
||||
@@ -922,9 +922,7 @@ export function useWorkflowExecution() {
|
||||
}
|
||||
|
||||
// Invalidate subscription queries to update usage
|
||||
setTimeout(() => {
|
||||
queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() })
|
||||
}, 1000)
|
||||
scheduleUsageRefresh(queryClient)
|
||||
|
||||
safeEnqueue(encodeSSE({ event: 'final', data: result }))
|
||||
// Note: Logs are already persisted server-side via execution-core.ts
|
||||
@@ -1458,9 +1456,7 @@ export function useWorkflowExecution() {
|
||||
setIsExecuting(activeWorkflowId, false)
|
||||
setActiveBlocks(activeWorkflowId, new Set())
|
||||
}
|
||||
setTimeout(() => {
|
||||
queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() })
|
||||
}, 1000)
|
||||
scheduleUsageRefresh(queryClient)
|
||||
}
|
||||
},
|
||||
|
||||
|
||||
@@ -3,7 +3,10 @@ import { getServiceConfigByProviderId } from '@/lib/oauth/utils'
|
||||
import { getQueryClient } from '@/app/_shell/providers/get-query-client'
|
||||
import type { BlockConfig } from '@/blocks/types'
|
||||
import { workspaceCredentialKeys } from '@/hooks/queries/utils/credential-keys'
|
||||
import { fetchWorkspaceCredentialList } from '@/hooks/queries/utils/fetch-workspace-credentials'
|
||||
import {
|
||||
fetchWorkspaceCredentialList,
|
||||
WORKSPACE_CREDENTIAL_LIST_STALE_TIME,
|
||||
} from '@/hooks/queries/utils/fetch-workspace-credentials'
|
||||
import { useWorkflowRegistry } from '@/stores/workflows/registry/store'
|
||||
|
||||
interface CredentialBlockOutput {
|
||||
@@ -72,7 +75,7 @@ export const CredentialBlock: BlockConfig<CredentialBlockOutput> = {
|
||||
const credentials = await getQueryClient().fetchQuery({
|
||||
queryKey: workspaceCredentialKeys.list(workspaceId),
|
||||
queryFn: () => fetchWorkspaceCredentialList(workspaceId),
|
||||
staleTime: 60 * 1000,
|
||||
staleTime: WORKSPACE_CREDENTIAL_LIST_STALE_TIME,
|
||||
})
|
||||
|
||||
const seen = new Set<string>()
|
||||
|
||||
@@ -4,10 +4,15 @@ import { createLogger } from '@sim/logger'
|
||||
import { authOAuthUtilsMock, authOAuthUtilsMockFns } from '@sim/testing'
|
||||
import { beforeEach, describe, expect, it, type Mock, vi } from 'vitest'
|
||||
|
||||
const { mockResolveAutoModel } = vi.hoisted(() => ({
|
||||
const { mockResolveAutoModel, mockCheckWorkspaceAccess } = vi.hoisted(() => ({
|
||||
mockCheckWorkspaceAccess: vi.fn(),
|
||||
mockResolveAutoModel: vi.fn(),
|
||||
}))
|
||||
|
||||
vi.mock('@/lib/workspaces/permissions/utils', () => ({
|
||||
checkWorkspaceAccess: mockCheckWorkspaceAccess,
|
||||
}))
|
||||
|
||||
vi.mock('@/lib/oauth/credential-service', () => authOAuthUtilsMock)
|
||||
|
||||
vi.mock('@/lib/credentials/access', () => ({
|
||||
@@ -34,24 +39,30 @@ vi.mock('@/lib/model-router/resolve', () => ({
|
||||
SIM_AUTO_SYSTEM_PREAMBLE: 'Sim auto system preamble',
|
||||
}))
|
||||
|
||||
import {
|
||||
PRIVATE_MODEL_INPUT_PROVENANCE_HEADER,
|
||||
PRIVATE_MODEL_INPUT_STATE_HEADER,
|
||||
PROJECTED_MODEL_INPUT_PATHS_V1,
|
||||
} from '@/lib/execution/model-input-provenance'
|
||||
import {
|
||||
RESOLVED_SECRET_PROVENANCE_FIELD,
|
||||
RESOLVED_SECRET_PROVENANCE_METADATA_V1,
|
||||
} from '@/lib/execution/private-tool-metadata'
|
||||
import { BlockType } from '@/executor/constants'
|
||||
import { EvaluatorBlockHandler } from '@/executor/handlers/evaluator/evaluator-handler'
|
||||
import type { ExecutionContext } from '@/executor/types'
|
||||
import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
|
||||
import { executeProviderRequest } from '@/providers'
|
||||
import type { ProviderRequest } from '@/providers/types'
|
||||
import { getProviderFromModel } from '@/providers/utils'
|
||||
import type { SerializedBlock } from '@/serializer/types'
|
||||
|
||||
const mockGetProviderFromModel = getProviderFromModel as Mock
|
||||
const mockFetch = vi.fn()
|
||||
const mockExecuteProviderRequest = executeProviderRequest as Mock
|
||||
|
||||
/** The provider request the handler built, keyed the way the old wire body was. */
|
||||
function providerRequestBody(index = 0): ProviderRequest & { provider: string } {
|
||||
const [provider, request] = mockExecuteProviderRequest.mock.calls[index] as [
|
||||
string,
|
||||
ProviderRequest,
|
||||
]
|
||||
return { provider, ...request }
|
||||
}
|
||||
|
||||
function providerRuntimeRegistry(index = 0): ResolvedSecretTraceRegistry | undefined {
|
||||
return mockExecuteProviderRequest.mock.calls[index][2]?.resolvedSecretTraceRegistry
|
||||
}
|
||||
const mockLogger =
|
||||
vi.mocked(createLogger).mock.results[
|
||||
vi.mocked(createLogger).mock.calls.findIndex(([name]) => name === 'EvaluatorBlockHandler')
|
||||
@@ -97,8 +108,7 @@ describe('EvaluatorBlockHandler', () => {
|
||||
// Reset mocks using vi
|
||||
vi.clearAllMocks()
|
||||
|
||||
// unstubGlobals removes any module-scope fetch stub before each test, so re-stub here
|
||||
vi.stubGlobal('fetch', mockFetch)
|
||||
mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: true })
|
||||
|
||||
// Default mock implementations
|
||||
authOAuthUtilsMockFns.mockResolveOAuthAccountId.mockResolvedValue({
|
||||
@@ -117,19 +127,12 @@ describe('EvaluatorBlockHandler', () => {
|
||||
billableRoutingCost: 0.002,
|
||||
})
|
||||
|
||||
// Set up fetch mock to return a successful response
|
||||
mockFetch.mockImplementation(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({ score1: 5, score2: 8 }),
|
||||
model: 'mock-model',
|
||||
tokens: { input: 50, output: 10, total: 60 },
|
||||
cost: 0.002,
|
||||
timing: { total: 200 },
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValue({
|
||||
content: JSON.stringify({ score1: 5, score2: 8 }),
|
||||
model: 'mock-model',
|
||||
tokens: { input: 50, output: 10, total: 60 },
|
||||
cost: 0.002,
|
||||
timing: { total: 200 },
|
||||
})
|
||||
})
|
||||
|
||||
@@ -142,6 +145,37 @@ describe('EvaluatorBlockHandler', () => {
|
||||
expect(handler.canHandle(nonEvalBlock)).toBe(false)
|
||||
})
|
||||
|
||||
/**
|
||||
* The admission checks the removed `/api/providers` hop owned. Mirrors the router's
|
||||
* coverage — both handlers reach the provider through the same shared entry point.
|
||||
*/
|
||||
const admissionInputs = {
|
||||
content: 'Evaluate this.',
|
||||
metrics: [{ name: 'score1', description: 'First score', range: { min: 0, max: 10 } }],
|
||||
model: 'gpt-4o',
|
||||
apiKey: 'test-api-key',
|
||||
}
|
||||
|
||||
it('refuses to reach the provider without an execution subject', async () => {
|
||||
mockContext.userId = undefined
|
||||
|
||||
await expect(handler.execute(mockContext, mockBlock, admissionInputs)).rejects.toThrow(
|
||||
'Unauthorized'
|
||||
)
|
||||
expect(mockExecuteProviderRequest).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('refuses to reach the provider when the subject lost workspace access', async () => {
|
||||
mockContext.workspaceId = 'test-workspace'
|
||||
mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: false })
|
||||
|
||||
await expect(handler.execute(mockContext, mockBlock, admissionInputs)).rejects.toThrow(
|
||||
'Forbidden'
|
||||
)
|
||||
expect(mockCheckWorkspaceAccess).toHaveBeenCalledWith('test-workspace', 'test-user')
|
||||
expect(mockExecuteProviderRequest).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('should execute evaluator block correctly with basic inputs', async () => {
|
||||
const inputs = {
|
||||
content: 'This is the content to evaluate.',
|
||||
@@ -165,17 +199,9 @@ describe('EvaluatorBlockHandler', () => {
|
||||
const result = await handler.execute(mockContext, mockBlock, inputs)
|
||||
|
||||
expect(mockGetProviderFromModel).toHaveBeenCalledWith('gpt-4o')
|
||||
expect(mockFetch).toHaveBeenCalledWith(
|
||||
expect.any(String),
|
||||
expect.objectContaining({
|
||||
method: 'POST',
|
||||
headers: expect.any(Object),
|
||||
body: expect.any(String),
|
||||
})
|
||||
)
|
||||
expect(mockExecuteProviderRequest).toHaveBeenCalledTimes(1)
|
||||
|
||||
const fetchCallArgs = mockFetch.mock.calls[0]
|
||||
const requestBody = JSON.parse(fetchCallArgs[1].body)
|
||||
const requestBody = providerRequestBody()
|
||||
expect(requestBody).toMatchObject({
|
||||
provider: 'openai',
|
||||
model: 'gpt-4o',
|
||||
@@ -257,15 +283,8 @@ describe('EvaluatorBlockHandler', () => {
|
||||
apiKey: credentialSecret,
|
||||
})
|
||||
|
||||
const request = mockFetch.mock.calls[0][1]
|
||||
const requestBody = JSON.parse(request.body)
|
||||
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBe(
|
||||
RESOLVED_SECRET_PROVENANCE_METADATA_V1
|
||||
)
|
||||
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBe(
|
||||
PROJECTED_MODEL_INPUT_PATHS_V1
|
||||
)
|
||||
expect(requestBody[RESOLVED_SECRET_PROVENANCE_FIELD]).toEqual({
|
||||
const requestBody = providerRequestBody()
|
||||
expect(providerRuntimeRegistry()?.exportProvenance()).toEqual({
|
||||
version: 1,
|
||||
complete: true,
|
||||
entries: [
|
||||
@@ -336,15 +355,11 @@ describe('EvaluatorBlockHandler', () => {
|
||||
registry.recordResolvedInputProjection(secret.path, secret.plaintext, secret.projected)
|
||||
}
|
||||
mockContext.resolvedSecretTraceRegistry = registry
|
||||
mockFetch.mockResolvedValueOnce({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({ [projectedMetric.name.toLowerCase()]: 7 }),
|
||||
model: 'mock-model',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
}),
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: JSON.stringify({ [projectedMetric.name.toLowerCase()]: 7 }),
|
||||
model: 'mock-model',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
})
|
||||
|
||||
const result = await handler.execute(mockContext, mockBlock, {
|
||||
@@ -354,7 +369,7 @@ describe('EvaluatorBlockHandler', () => {
|
||||
apiKey: 'test-api-key',
|
||||
})
|
||||
|
||||
const requestBody = JSON.parse(mockFetch.mock.calls[0][1].body)
|
||||
const requestBody = providerRequestBody()
|
||||
const serializedRequest = JSON.stringify(requestBody)
|
||||
for (const secret of secrets) {
|
||||
expect(serializedRequest).not.toContain(secret.plaintext)
|
||||
@@ -365,8 +380,9 @@ describe('EvaluatorBlockHandler', () => {
|
||||
[projectedMetric.name.toLowerCase()]: { type: 'number' },
|
||||
})
|
||||
expect(
|
||||
requestBody[RESOLVED_SECRET_PROVENANCE_FIELD].entries
|
||||
.map((entry: { name: string }) => entry.name)
|
||||
providerRuntimeRegistry()
|
||||
?.exportProvenance()
|
||||
.entries.map((entry: { name: string }) => entry.name)
|
||||
.sort()
|
||||
).toEqual(secrets.map((secret) => secret.name).sort())
|
||||
expect(result).toMatchObject({ [rawMetric.name.toLowerCase()]: 7 })
|
||||
@@ -380,11 +396,7 @@ describe('EvaluatorBlockHandler', () => {
|
||||
apiKey: 'test-api-key',
|
||||
})
|
||||
|
||||
const request = mockFetch.mock.calls[0][1]
|
||||
const requestBody = JSON.parse(request.body)
|
||||
expect(Object.hasOwn(requestBody, RESOLVED_SECRET_PROVENANCE_FIELD)).toBe(false)
|
||||
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBeNull()
|
||||
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBeNull()
|
||||
expect(providerRuntimeRegistry()).toBeUndefined()
|
||||
})
|
||||
|
||||
it('resolves sim-auto before executing evaluator and preserves its public identity', async () => {
|
||||
@@ -400,15 +412,11 @@ describe('EvaluatorBlockHandler', () => {
|
||||
model: 'sim-auto',
|
||||
}
|
||||
|
||||
mockFetch.mockResolvedValueOnce({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({ quality: 5 }),
|
||||
model: 'fireworks/glm-5.2',
|
||||
tokens: { input: 80, output: 10, total: 90 },
|
||||
cost: { input: 0.001, output: 0.0005, total: 0.0015 },
|
||||
}),
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: JSON.stringify({ quality: 5 }),
|
||||
model: 'fireworks/glm-5.2',
|
||||
tokens: { input: 80, output: 10, total: 90 },
|
||||
cost: { input: 0.001, output: 0.0005, total: 0.0015 },
|
||||
})
|
||||
|
||||
const result = await handler.execute(mockContext, mockBlock, inputs)
|
||||
@@ -427,7 +435,7 @@ describe('EvaluatorBlockHandler', () => {
|
||||
})
|
||||
expect(mockGetProviderFromModel).toHaveBeenCalledWith('fireworks/glm-5.2')
|
||||
|
||||
const requestBody = JSON.parse(mockFetch.mock.calls[0][1].body)
|
||||
const requestBody = providerRequestBody()
|
||||
expect(requestBody).toMatchObject({
|
||||
provider: 'openai',
|
||||
model: 'fireworks/glm-5.2',
|
||||
@@ -448,19 +456,13 @@ describe('EvaluatorBlockHandler', () => {
|
||||
it('bills the cost the provider proxy decided rather than recomputing it', async () => {
|
||||
// The proxy already resolved key provenance and the margin; recomputing
|
||||
// here would re-charge a BYOK caller the proxy correctly zeroed.
|
||||
mockFetch.mockImplementation(() =>
|
||||
Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({ score1: 5, score2: 8 }),
|
||||
model: 'mock-model',
|
||||
tokens: { input: 50, output: 10, total: 60 },
|
||||
cost: { input: 0.001, output: 0.0005, total: 0.0015 },
|
||||
timing: { total: 200 },
|
||||
}),
|
||||
})
|
||||
)
|
||||
mockExecuteProviderRequest.mockResolvedValue({
|
||||
content: JSON.stringify({ score1: 5, score2: 8 }),
|
||||
model: 'mock-model',
|
||||
tokens: { input: 50, output: 10, total: 60 },
|
||||
cost: { input: 0.001, output: 0.0005, total: 0.0015 },
|
||||
timing: { total: 200 },
|
||||
})
|
||||
|
||||
const result = await handler.execute(mockContext, mockBlock, {
|
||||
content: 'This is the content to evaluate.',
|
||||
@@ -487,24 +489,17 @@ describe('EvaluatorBlockHandler', () => {
|
||||
apiKey: 'test-api-key',
|
||||
}
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({ clarity: 4 }),
|
||||
model: 'm',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: JSON.stringify({ clarity: 4 }),
|
||||
model: 'm',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
})
|
||||
|
||||
await handler.execute(mockContext, mockBlock, inputs)
|
||||
|
||||
const fetchCallArgs = mockFetch.mock.calls[0]
|
||||
const requestBody = JSON.parse(fetchCallArgs[1].body)
|
||||
const requestBody = providerRequestBody()
|
||||
expect(requestBody).toMatchObject({
|
||||
systemPrompt: expect.stringContaining(JSON.stringify(contentObj, null, 2)),
|
||||
})
|
||||
@@ -524,24 +519,17 @@ describe('EvaluatorBlockHandler', () => {
|
||||
apiKey: 'test-api-key',
|
||||
}
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({ completeness: 1 }),
|
||||
model: 'm',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: JSON.stringify({ completeness: 1 }),
|
||||
model: 'm',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
})
|
||||
|
||||
await handler.execute(mockContext, mockBlock, inputs)
|
||||
|
||||
const fetchCallArgs = mockFetch.mock.calls[0]
|
||||
const requestBody = JSON.parse(fetchCallArgs[1].body)
|
||||
const requestBody = providerRequestBody()
|
||||
expect(requestBody).toMatchObject({
|
||||
systemPrompt: expect.stringContaining(JSON.stringify(contentObj, null, 2)),
|
||||
})
|
||||
@@ -560,18 +548,12 @@ describe('EvaluatorBlockHandler', () => {
|
||||
apiKey: 'test-api-key',
|
||||
}
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: '```json\n{ "quality": 9 }\n```',
|
||||
model: 'm',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: '```json\n{ "quality": 9 }\n```',
|
||||
model: 'm',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
})
|
||||
|
||||
const result = await handler.execute(mockContext, mockBlock, inputs)
|
||||
@@ -586,18 +568,12 @@ describe('EvaluatorBlockHandler', () => {
|
||||
apiKey: 'test-api-key',
|
||||
}
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: 'Sorry, I cannot provide a score.',
|
||||
model: 'm',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: 'Sorry, I cannot provide a score.',
|
||||
model: 'm',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
})
|
||||
|
||||
const result = await handler.execute(mockContext, mockBlock, inputs)
|
||||
@@ -615,18 +591,12 @@ describe('EvaluatorBlockHandler', () => {
|
||||
apiKey: 'test-api-key',
|
||||
}
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: '{ "accuracy": 1, "fluency": invalid }',
|
||||
model: 'm',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: '{ "accuracy": 1, "fluency": invalid }',
|
||||
model: 'm',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
})
|
||||
|
||||
const result = await handler.execute(mockContext, mockBlock, inputs)
|
||||
@@ -647,18 +617,12 @@ describe('EvaluatorBlockHandler', () => {
|
||||
apiKey: 'test-api-key',
|
||||
}
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({ camelcasescore: 7 }),
|
||||
model: 'm',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: JSON.stringify({ camelcasescore: 7 }),
|
||||
model: 'm',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
})
|
||||
|
||||
const result = await handler.execute(mockContext, mockBlock, inputs)
|
||||
@@ -684,18 +648,12 @@ describe('EvaluatorBlockHandler', () => {
|
||||
apiKey: 'test-api-key',
|
||||
}
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({ presentScore: 4 }),
|
||||
model: 'm',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: JSON.stringify({ presentScore: 4 }),
|
||||
model: 'm',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
})
|
||||
|
||||
const result = await handler.execute(mockContext, mockBlock, inputs)
|
||||
@@ -708,13 +666,7 @@ describe('EvaluatorBlockHandler', () => {
|
||||
const inputs = { content: 'Test error handling.', apiKey: 'test-api-key' }
|
||||
|
||||
// Override fetch mock to return an error
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: false,
|
||||
status: 500,
|
||||
json: () => Promise.resolve({ error: 'Server error' }),
|
||||
})
|
||||
})
|
||||
mockExecuteProviderRequest.mockRejectedValueOnce(new Error('Server error'))
|
||||
|
||||
await expect(handler.execute(mockContext, mockBlock, inputs)).rejects.toThrow('Server error')
|
||||
})
|
||||
@@ -730,11 +682,7 @@ describe('EvaluatorBlockHandler', () => {
|
||||
])
|
||||
registry.recordResolved('CONTENT_SECRET', 'resolved-evaluator-secret')
|
||||
mockContext.resolvedSecretTraceRegistry = registry
|
||||
mockFetch.mockResolvedValueOnce({
|
||||
ok: false,
|
||||
status: 500,
|
||||
json: () => Promise.resolve({ error: providerError }),
|
||||
})
|
||||
mockExecuteProviderRequest.mockRejectedValueOnce(new Error(providerError))
|
||||
|
||||
await expect(
|
||||
handler.execute(mockContext, mockBlock, {
|
||||
@@ -767,24 +715,17 @@ describe('EvaluatorBlockHandler', () => {
|
||||
|
||||
mockGetProviderFromModel.mockReturnValue('azure-openai')
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({ quality: 8 }),
|
||||
model: 'gpt-4o',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: JSON.stringify({ quality: 8 }),
|
||||
model: 'gpt-4o',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
})
|
||||
|
||||
await handler.execute(mockContext, mockBlock, inputs)
|
||||
|
||||
const fetchCallArgs = mockFetch.mock.calls[0]
|
||||
const requestBody = JSON.parse(fetchCallArgs[1].body)
|
||||
const requestBody = providerRequestBody()
|
||||
|
||||
expect(requestBody).toMatchObject({
|
||||
provider: 'azure-openai',
|
||||
@@ -824,24 +765,17 @@ describe('EvaluatorBlockHandler', () => {
|
||||
;(mockDb.db.query as any).account = { findFirst: vi.fn() }
|
||||
vi.spyOn(mockDb.db.query.account, 'findFirst').mockResolvedValue(mockAccount as any)
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({ quality: 9 }),
|
||||
model: 'gemini-2.0-flash-exp',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: JSON.stringify({ quality: 9 }),
|
||||
model: 'gemini-2.0-flash-exp',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
})
|
||||
|
||||
await handler.execute(mockContext, mockBlock, inputs)
|
||||
|
||||
const fetchCallArgs = mockFetch.mock.calls[0]
|
||||
const requestBody = JSON.parse(fetchCallArgs[1].body)
|
||||
const requestBody = providerRequestBody()
|
||||
|
||||
expect(requestBody).toMatchObject({
|
||||
provider: 'vertex',
|
||||
@@ -860,24 +794,17 @@ describe('EvaluatorBlockHandler', () => {
|
||||
// No model provided - should use default
|
||||
}
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({ score: 7 }),
|
||||
model: 'claude-sonnet-5',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: JSON.stringify({ score: 7 }),
|
||||
model: 'claude-sonnet-5',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
})
|
||||
|
||||
await handler.execute(mockContext, mockBlock, inputs)
|
||||
|
||||
const fetchCallArgs = mockFetch.mock.calls[0]
|
||||
const requestBody = JSON.parse(fetchCallArgs[1].body)
|
||||
const requestBody = providerRequestBody()
|
||||
|
||||
expect(requestBody.model).toBe('claude-sonnet-5')
|
||||
})
|
||||
|
||||
@@ -1,10 +1,5 @@
|
||||
import { createLogger } from '@sim/logger'
|
||||
import {
|
||||
addModelInputProvenanceToRequest,
|
||||
createModelInputProvenanceRequestMetadata,
|
||||
markModelInputProjected,
|
||||
projectResolvedModelInput,
|
||||
} from '@/lib/execution/model-input-provenance'
|
||||
import { projectResolvedModelInput } from '@/lib/execution/model-input-provenance'
|
||||
import {
|
||||
type AutoRoutingResult,
|
||||
addAutoRoutingCost,
|
||||
@@ -15,8 +10,8 @@ import type { BlockOutput } from '@/blocks/types'
|
||||
import { validateModelProvider } from '@/ee/access-control/utils/permission-check'
|
||||
import { BlockType, DEFAULTS, EVALUATOR } from '@/executor/constants'
|
||||
import type { BlockHandler, ExecutionContext } from '@/executor/types'
|
||||
import { buildAPIUrl, buildAuthHeaders, extractAPIErrorMessage } from '@/executor/utils/http'
|
||||
import { isJSONString, parseJSON, stringifyJSON } from '@/executor/utils/json'
|
||||
import { executeBlockProviderRequest } from '@/executor/utils/provider-request'
|
||||
import { projectResolvedSecretDiagnosticError } from '@/executor/utils/resolved-secret-content-projection'
|
||||
import { refuseResolvedSecretProjection } from '@/executor/utils/resolved-secret-projection-refusal'
|
||||
import type {
|
||||
@@ -192,8 +187,6 @@ export class EvaluatorBlockHandler implements BlockHandler {
|
||||
}
|
||||
|
||||
try {
|
||||
const url = buildAPIUrl('/api/providers', ctx.userId ? { userId: ctx.userId } : {})
|
||||
|
||||
const providerRequest: ProviderRequest = {
|
||||
model,
|
||||
systemPrompt: systemPromptObj.systemPrompt,
|
||||
@@ -219,30 +212,13 @@ export class EvaluatorBlockHandler implements BlockHandler {
|
||||
workspaceId: ctx.workspaceId,
|
||||
}
|
||||
|
||||
const headers = new Headers(await buildAuthHeaders(ctx.userId))
|
||||
const modelInputMetadata = createModelInputProvenanceRequestMetadata(
|
||||
modelInputProjection.registry,
|
||||
modelInputPaths
|
||||
)
|
||||
const requestBody = addModelInputProvenanceToRequest(
|
||||
{ provider: providerId, ...providerRequest },
|
||||
headers,
|
||||
modelInputMetadata
|
||||
)
|
||||
if (modelInputMetadata) markModelInputProjected(headers)
|
||||
const response = await fetch(url.toString(), {
|
||||
method: 'POST',
|
||||
headers,
|
||||
body: stringifyJSON(requestBody),
|
||||
const result = await executeBlockProviderRequest({
|
||||
ctx,
|
||||
providerId,
|
||||
request: providerRequest,
|
||||
resolvedSecretTraceRegistry: modelInputProjection.registry,
|
||||
})
|
||||
|
||||
if (!response.ok) {
|
||||
const errorMessage = await extractAPIErrorMessage(response)
|
||||
throw new Error(errorMessage)
|
||||
}
|
||||
|
||||
const result = await response.json()
|
||||
|
||||
const parsedContent = this.extractJSONFromResponse(
|
||||
result.content,
|
||||
ctx.resolvedSecretTraceRegistry
|
||||
@@ -250,9 +226,8 @@ export class EvaluatorBlockHandler implements BlockHandler {
|
||||
|
||||
const metricScores = this.extractMetricScores(parsedContent, metrics, projectedMetrics)
|
||||
|
||||
const inputTokens = result.tokens?.input || result.tokens?.prompt || DEFAULTS.TOKENS.PROMPT
|
||||
const outputTokens =
|
||||
result.tokens?.output || result.tokens?.completion || DEFAULTS.TOKENS.COMPLETION
|
||||
const inputTokens = result.tokens?.input || DEFAULTS.TOKENS.PROMPT
|
||||
const outputTokens = result.tokens?.output || DEFAULTS.TOKENS.COMPLETION
|
||||
|
||||
const cost = addAutoRoutingCost(
|
||||
resolveProxiedModelCost(result.cost),
|
||||
|
||||
@@ -9,8 +9,13 @@ import {
|
||||
} from '@sim/testing'
|
||||
import { beforeEach, describe, expect, it, type Mock, vi } from 'vitest'
|
||||
|
||||
const { mockResolveAutoModel } = vi.hoisted(() => ({
|
||||
const { mockResolveAutoModel, mockCheckWorkspaceAccess } = vi.hoisted(() => ({
|
||||
mockResolveAutoModel: vi.fn(),
|
||||
mockCheckWorkspaceAccess: vi.fn(),
|
||||
}))
|
||||
|
||||
vi.mock('@/lib/workspaces/permissions/utils', () => ({
|
||||
checkWorkspaceAccess: mockCheckWorkspaceAccess,
|
||||
}))
|
||||
|
||||
vi.mock('@/lib/oauth/credential-service', () => authOAuthUtilsMock)
|
||||
@@ -40,27 +45,30 @@ vi.mock('@/lib/model-router/resolve', () => ({
|
||||
SIM_AUTO_SYSTEM_PREAMBLE: 'Sim auto system preamble',
|
||||
}))
|
||||
|
||||
import {
|
||||
PRIVATE_MODEL_INPUT_PROVENANCE_HEADER,
|
||||
PRIVATE_MODEL_INPUT_STATE_HEADER,
|
||||
PROJECTED_MODEL_INPUT_PATHS_V1,
|
||||
} from '@/lib/execution/model-input-provenance'
|
||||
import {
|
||||
RESOLVED_SECRET_PROVENANCE_FIELD,
|
||||
RESOLVED_SECRET_PROVENANCE_METADATA_V1,
|
||||
} from '@/lib/execution/private-tool-metadata'
|
||||
import { generateRouterPrompt, generateRouterV2Prompt } from '@/blocks/blocks/router'
|
||||
import { BlockType } from '@/executor/constants'
|
||||
import { RouterBlockHandler } from '@/executor/handlers/router/router-handler'
|
||||
import type { ExecutionContext } from '@/executor/types'
|
||||
import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
|
||||
import { executeProviderRequest } from '@/providers'
|
||||
import { getProviderFromModel } from '@/providers/utils'
|
||||
import type { SerializedBlock, SerializedWorkflow } from '@/serializer/types'
|
||||
|
||||
const mockGenerateRouterPrompt = generateRouterPrompt as Mock
|
||||
const mockGenerateRouterV2Prompt = generateRouterV2Prompt as Mock
|
||||
const mockGetProviderFromModel = getProviderFromModel as Mock
|
||||
const mockFetch = vi.fn()
|
||||
const mockExecuteProviderRequest = executeProviderRequest as Mock
|
||||
|
||||
/** The provider request the handler built, keyed the way the old wire body was. */
|
||||
function providerRequestBody(index = 0): Record<string, unknown> {
|
||||
const [provider, request] = mockExecuteProviderRequest.mock.calls[index]
|
||||
return { provider, ...request }
|
||||
}
|
||||
|
||||
function providerRuntimeRegistry(index = 0): ResolvedSecretTraceRegistry | undefined {
|
||||
return mockExecuteProviderRequest.mock.calls[index][2]?.resolvedSecretTraceRegistry
|
||||
}
|
||||
|
||||
const mockLogger =
|
||||
vi.mocked(createLogger).mock.results[
|
||||
vi.mocked(createLogger).mock.calls.findIndex(([name]) => name === 'RouterBlockHandler')
|
||||
@@ -138,8 +146,7 @@ describe('RouterBlockHandler', () => {
|
||||
vi.clearAllMocks()
|
||||
encryptionMockFns.mockDecryptSecret.mockResolvedValue({ decrypted: 'test-decrypted' })
|
||||
|
||||
// unstubGlobals removes any module-scope fetch stub before each test, so re-stub here
|
||||
vi.stubGlobal('fetch', mockFetch)
|
||||
mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: true })
|
||||
|
||||
authOAuthUtilsMockFns.mockResolveOAuthAccountId.mockResolvedValue({
|
||||
accountId: 'test-vertex-credential-id',
|
||||
@@ -158,18 +165,12 @@ describe('RouterBlockHandler', () => {
|
||||
billableRoutingCost: 0.002,
|
||||
})
|
||||
|
||||
mockFetch.mockImplementation(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: 'target-block-1',
|
||||
model: 'mock-model',
|
||||
tokens: { input: 100, output: 5, total: 105 },
|
||||
cost: 0.003,
|
||||
timing: { total: 300 },
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValue({
|
||||
content: 'target-block-1',
|
||||
model: 'mock-model',
|
||||
tokens: { input: 100, output: 5, total: 105 },
|
||||
cost: 0.003,
|
||||
timing: { total: 300 },
|
||||
})
|
||||
})
|
||||
|
||||
@@ -219,17 +220,9 @@ describe('RouterBlockHandler', () => {
|
||||
|
||||
expect(mockGenerateRouterPrompt).toHaveBeenCalledWith(inputs.prompt, expectedTargetBlocks)
|
||||
expect(mockGetProviderFromModel).toHaveBeenCalledWith('gpt-4o')
|
||||
expect(mockFetch).toHaveBeenCalledWith(
|
||||
expect.any(String),
|
||||
expect.objectContaining({
|
||||
method: 'POST',
|
||||
headers: expect.any(Object),
|
||||
body: expect.any(String),
|
||||
})
|
||||
)
|
||||
expect(mockExecuteProviderRequest).toHaveBeenCalledTimes(1)
|
||||
|
||||
const fetchCallArgs = mockFetch.mock.calls[0]
|
||||
const requestBody = JSON.parse(fetchCallArgs[1].body)
|
||||
const requestBody = providerRequestBody()
|
||||
expect(requestBody).toMatchObject({
|
||||
provider: 'openai',
|
||||
model: 'gpt-4o',
|
||||
@@ -282,15 +275,8 @@ describe('RouterBlockHandler', () => {
|
||||
apiKey: credentialSecret,
|
||||
})
|
||||
|
||||
const request = mockFetch.mock.calls[0][1]
|
||||
const requestBody = JSON.parse(request.body)
|
||||
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBe(
|
||||
RESOLVED_SECRET_PROVENANCE_METADATA_V1
|
||||
)
|
||||
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBe(
|
||||
PROJECTED_MODEL_INPUT_PATHS_V1
|
||||
)
|
||||
expect(requestBody[RESOLVED_SECRET_PROVENANCE_FIELD]).toEqual({
|
||||
const requestBody = providerRequestBody()
|
||||
expect(providerRuntimeRegistry()?.exportProvenance()).toEqual({
|
||||
version: 1,
|
||||
complete: true,
|
||||
entries: [
|
||||
@@ -353,8 +339,7 @@ describe('RouterBlockHandler', () => {
|
||||
expect(rawState).toEqual({ result: stateSecret, ordinary: 'Box remains raw state' })
|
||||
expect(mockTargetBlock1.config.params).toEqual({ p: 'a' })
|
||||
|
||||
const requestBody = JSON.parse(mockFetch.mock.calls[0][1].body)
|
||||
expect(requestBody[RESOLVED_SECRET_PROVENANCE_FIELD]).toEqual({
|
||||
expect(providerRuntimeRegistry()?.exportProvenance()).toEqual({
|
||||
version: 1,
|
||||
complete: true,
|
||||
entries: [],
|
||||
@@ -404,29 +389,19 @@ describe('RouterBlockHandler', () => {
|
||||
apiKey: 'test-api-key',
|
||||
})
|
||||
|
||||
const request = mockFetch.mock.calls[0][1]
|
||||
const requestBody = JSON.parse(request.body)
|
||||
expect(Object.hasOwn(requestBody, RESOLVED_SECRET_PROVENANCE_FIELD)).toBe(false)
|
||||
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBeNull()
|
||||
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBeNull()
|
||||
expect(providerRuntimeRegistry()).toBeUndefined()
|
||||
})
|
||||
|
||||
it('bills the cost the provider proxy decided rather than recomputing it', async () => {
|
||||
// The proxy already resolved key provenance and the margin; recomputing
|
||||
// here would re-charge a BYOK caller the proxy correctly zeroed.
|
||||
mockFetch.mockImplementation(() =>
|
||||
Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: 'target-block-1',
|
||||
model: 'mock-model',
|
||||
tokens: { input: 100, output: 5, total: 105 },
|
||||
cost: { input: 0.004, output: 0.002, total: 0.006 },
|
||||
timing: { total: 300 },
|
||||
}),
|
||||
})
|
||||
)
|
||||
mockExecuteProviderRequest.mockResolvedValue({
|
||||
content: 'target-block-1',
|
||||
model: 'mock-model',
|
||||
tokens: { input: 100, output: 5, total: 105 },
|
||||
cost: { input: 0.004, output: 0.002, total: 0.006 },
|
||||
timing: { total: 300 },
|
||||
})
|
||||
|
||||
const result = await handler.execute(mockContext, mockBlock, {
|
||||
prompt: 'Choose the best option.',
|
||||
@@ -439,6 +414,26 @@ describe('RouterBlockHandler', () => {
|
||||
})
|
||||
})
|
||||
|
||||
it('refuses to reach the provider without an execution subject', async () => {
|
||||
mockContext.userId = undefined
|
||||
|
||||
await expect(
|
||||
handler.execute(mockContext, mockBlock, { prompt: 'Choose the best option.' })
|
||||
).rejects.toThrow('Unauthorized')
|
||||
expect(mockExecuteProviderRequest).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('refuses to reach the provider when the subject lost workspace access', async () => {
|
||||
mockContext.workspaceId = 'test-workspace'
|
||||
mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: false })
|
||||
|
||||
await expect(
|
||||
handler.execute(mockContext, mockBlock, { prompt: 'Choose the best option.' })
|
||||
).rejects.toThrow('Forbidden')
|
||||
expect(mockCheckWorkspaceAccess).toHaveBeenCalledWith('test-workspace', 'test-user')
|
||||
expect(mockExecuteProviderRequest).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('should throw error if target block is missing', async () => {
|
||||
const inputs = { prompt: 'Test' }
|
||||
mockContext.workflow!.blocks = [mockBlock, mockTargetBlock2]
|
||||
@@ -446,24 +441,18 @@ describe('RouterBlockHandler', () => {
|
||||
await expect(handler.execute(mockContext, mockBlock, inputs)).rejects.toThrow(
|
||||
'Target block target-block-1 not found'
|
||||
)
|
||||
expect(mockFetch).not.toHaveBeenCalled()
|
||||
expect(mockExecuteProviderRequest).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('should throw error if LLM response is not a valid target block ID', async () => {
|
||||
const inputs = { prompt: 'Test', apiKey: 'test-api-key' }
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: 'invalid-block-id',
|
||||
model: 'mock-model',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: 'invalid-block-id',
|
||||
model: 'mock-model',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
})
|
||||
|
||||
await expect(handler.execute(mockContext, mockBlock, inputs)).rejects.toThrow(
|
||||
@@ -475,16 +464,12 @@ describe('RouterBlockHandler', () => {
|
||||
const plaintext = 'router-provider-plaintext-secret'
|
||||
const content = `${plaintext} __var_API_KEY __sim_runtime`
|
||||
|
||||
mockFetch.mockResolvedValueOnce({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content,
|
||||
model: 'mock-model',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
}),
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content,
|
||||
model: 'mock-model',
|
||||
tokens: {},
|
||||
cost: 0,
|
||||
timing: {},
|
||||
})
|
||||
|
||||
await expect(handler.execute(mockContext, mockBlock, { prompt: 'Test' })).rejects.toThrow(
|
||||
@@ -512,8 +497,7 @@ describe('RouterBlockHandler', () => {
|
||||
|
||||
expect(mockGetProviderFromModel).toHaveBeenCalledWith('claude-sonnet-5')
|
||||
|
||||
const fetchCallArgs = mockFetch.mock.calls[0]
|
||||
const requestBody = JSON.parse(fetchCallArgs[1].body)
|
||||
const requestBody = providerRequestBody()
|
||||
expect(requestBody).toMatchObject({
|
||||
model: 'claude-sonnet-5',
|
||||
temperature: 0.1,
|
||||
@@ -523,13 +507,7 @@ describe('RouterBlockHandler', () => {
|
||||
it('should handle server error responses', async () => {
|
||||
const inputs = { prompt: 'Test error handling.', apiKey: 'test-api-key' }
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: false,
|
||||
status: 500,
|
||||
json: () => Promise.resolve({ error: 'Server error' }),
|
||||
})
|
||||
})
|
||||
mockExecuteProviderRequest.mockRejectedValueOnce(new Error('Server error'))
|
||||
|
||||
await expect(handler.execute(mockContext, mockBlock, inputs)).rejects.toThrow('Server error')
|
||||
})
|
||||
@@ -537,11 +515,7 @@ describe('RouterBlockHandler', () => {
|
||||
it('does not log sensitive provider errors while preserving the thrown error', async () => {
|
||||
const providerError = 'provider-plaintext-secret __var_API_KEY __sim_runtime'
|
||||
|
||||
mockFetch.mockResolvedValueOnce({
|
||||
ok: false,
|
||||
status: 500,
|
||||
json: () => Promise.resolve({ error: providerError }),
|
||||
})
|
||||
mockExecuteProviderRequest.mockRejectedValueOnce(new Error(providerError))
|
||||
|
||||
await expect(handler.execute(mockContext, mockBlock, { prompt: 'Test' })).rejects.toThrow(
|
||||
providerError
|
||||
@@ -569,8 +543,7 @@ describe('RouterBlockHandler', () => {
|
||||
|
||||
await handler.execute(mockContext, mockBlock, inputs)
|
||||
|
||||
const fetchCallArgs = mockFetch.mock.calls[0]
|
||||
const requestBody = JSON.parse(fetchCallArgs[1].body)
|
||||
const requestBody = providerRequestBody()
|
||||
|
||||
expect(requestBody).toMatchObject({
|
||||
provider: 'azure-openai',
|
||||
@@ -604,8 +577,7 @@ describe('RouterBlockHandler', () => {
|
||||
|
||||
await handler.execute(mockContext, mockBlock, inputs)
|
||||
|
||||
const fetchCallArgs = mockFetch.mock.calls[0]
|
||||
const requestBody = JSON.parse(fetchCallArgs[1].body)
|
||||
const requestBody = providerRequestBody()
|
||||
|
||||
expect(requestBody).toMatchObject({
|
||||
provider: 'vertex',
|
||||
@@ -688,8 +660,7 @@ describe('RouterBlockHandler V2', () => {
|
||||
|
||||
vi.clearAllMocks()
|
||||
|
||||
// unstubGlobals removes any module-scope fetch stub before each test, so re-stub here
|
||||
vi.stubGlobal('fetch', mockFetch)
|
||||
mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: true })
|
||||
|
||||
authOAuthUtilsMockFns.mockResolveOAuthAccountId.mockResolvedValue({
|
||||
accountId: 'test-vertex-credential-id',
|
||||
@@ -732,19 +703,13 @@ describe('RouterBlockHandler V2', () => {
|
||||
]),
|
||||
}
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({
|
||||
route: 'route-support',
|
||||
reasoning: 'The user mentioned a billing issue which is a customer support matter.',
|
||||
}),
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 150, output: 25, total: 175 },
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: JSON.stringify({
|
||||
route: 'route-support',
|
||||
reasoning: 'The user mentioned a billing issue which is a customer support matter.',
|
||||
}),
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 150, output: 25, total: 175 },
|
||||
})
|
||||
|
||||
const result = await handler.execute(mockContext, mockRouterV2Block, inputs)
|
||||
@@ -781,14 +746,10 @@ describe('RouterBlockHandler V2', () => {
|
||||
registry.recordResolvedInputProjection(['context'], contextSecret, '{{CONTEXT_SECRET}}')
|
||||
registry.recordResolved('API_KEY', credentialSecret)
|
||||
mockContext.resolvedSecretTraceRegistry = registry
|
||||
mockFetch.mockResolvedValueOnce({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({ route: 'route-support', reasoning: 'Matched support.' }),
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 10, output: 5, total: 15 },
|
||||
}),
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: JSON.stringify({ route: 'route-support', reasoning: 'Matched support.' }),
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 10, output: 5, total: 15 },
|
||||
})
|
||||
|
||||
await handler.execute(mockContext, mockRouterV2Block, {
|
||||
@@ -798,15 +759,8 @@ describe('RouterBlockHandler V2', () => {
|
||||
routes: [{ id: 'route-support', title: 'Support', value: 'Support requests' }],
|
||||
})
|
||||
|
||||
const request = mockFetch.mock.calls[0][1]
|
||||
const requestBody = JSON.parse(request.body)
|
||||
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBe(
|
||||
RESOLVED_SECRET_PROVENANCE_METADATA_V1
|
||||
)
|
||||
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBe(
|
||||
PROJECTED_MODEL_INPUT_PATHS_V1
|
||||
)
|
||||
expect(requestBody[RESOLVED_SECRET_PROVENANCE_FIELD]).toEqual({
|
||||
const requestBody = providerRequestBody()
|
||||
expect(providerRuntimeRegistry()?.exportProvenance()).toEqual({
|
||||
version: 1,
|
||||
complete: true,
|
||||
entries: [
|
||||
@@ -821,14 +775,10 @@ describe('RouterBlockHandler V2', () => {
|
||||
})
|
||||
|
||||
it('keeps the router V2 request shape when no provenance registry exists', async () => {
|
||||
mockFetch.mockResolvedValueOnce({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({ route: 'route-support', reasoning: 'Matched support.' }),
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 10, output: 5, total: 15 },
|
||||
}),
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: JSON.stringify({ route: 'route-support', reasoning: 'Matched support.' }),
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 10, output: 5, total: 15 },
|
||||
})
|
||||
|
||||
await handler.execute(mockContext, mockRouterV2Block, {
|
||||
@@ -838,11 +788,7 @@ describe('RouterBlockHandler V2', () => {
|
||||
routes: [{ id: 'route-support', title: 'Support', value: 'Support requests' }],
|
||||
})
|
||||
|
||||
const request = mockFetch.mock.calls[0][1]
|
||||
const requestBody = JSON.parse(request.body)
|
||||
expect(Object.hasOwn(requestBody, RESOLVED_SECRET_PROVENANCE_FIELD)).toBe(false)
|
||||
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBeNull()
|
||||
expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBeNull()
|
||||
expect(providerRuntimeRegistry()).toBeUndefined()
|
||||
})
|
||||
|
||||
it('resolves sim-auto before executing router V2 and preserves its public identity', async () => {
|
||||
@@ -859,18 +805,14 @@ describe('RouterBlockHandler V2', () => {
|
||||
],
|
||||
}
|
||||
|
||||
mockFetch.mockResolvedValueOnce({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({
|
||||
route: 'route-sales',
|
||||
reasoning: 'This is a new request.',
|
||||
}),
|
||||
model: 'fireworks/glm-5.2',
|
||||
tokens: { input: 100, output: 20, total: 120 },
|
||||
cost: { input: 0.001, output: 0.0005, total: 0.0015 },
|
||||
}),
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: JSON.stringify({
|
||||
route: 'route-sales',
|
||||
reasoning: 'This is a new request.',
|
||||
}),
|
||||
model: 'fireworks/glm-5.2',
|
||||
tokens: { input: 100, output: 20, total: 120 },
|
||||
cost: { input: 0.001, output: 0.0005, total: 0.0015 },
|
||||
})
|
||||
|
||||
const result = await handler.execute(mockContext, mockRouterV2Block, inputs)
|
||||
@@ -889,7 +831,7 @@ describe('RouterBlockHandler V2', () => {
|
||||
})
|
||||
expect(mockGetProviderFromModel).toHaveBeenCalledWith('fireworks/glm-5.2')
|
||||
|
||||
const requestBody = JSON.parse(mockFetch.mock.calls[0][1].body)
|
||||
const requestBody = providerRequestBody()
|
||||
expect(requestBody).toMatchObject({
|
||||
provider: 'openai',
|
||||
model: 'fireworks/glm-5.2',
|
||||
@@ -915,25 +857,18 @@ describe('RouterBlockHandler V2', () => {
|
||||
routes: JSON.stringify([{ id: 'route-1', title: 'Route 1', value: 'Description 1' }]),
|
||||
}
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({
|
||||
route: 'route-1',
|
||||
reasoning: 'Test reasoning',
|
||||
}),
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 100, output: 20, total: 120 },
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: JSON.stringify({
|
||||
route: 'route-1',
|
||||
reasoning: 'Test reasoning',
|
||||
}),
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 100, output: 20, total: 120 },
|
||||
})
|
||||
|
||||
await handler.execute(mockContext, mockRouterV2Block, inputs)
|
||||
|
||||
const fetchCallArgs = mockFetch.mock.calls[0]
|
||||
const requestBody = JSON.parse(fetchCallArgs[1].body)
|
||||
const requestBody = providerRequestBody()
|
||||
|
||||
expect(requestBody.responseFormat).toEqual({
|
||||
name: 'router_response',
|
||||
@@ -964,19 +899,13 @@ describe('RouterBlockHandler V2', () => {
|
||||
routes: JSON.stringify([{ id: 'route-1', title: 'Route 1', value: 'Specific topic' }]),
|
||||
}
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({
|
||||
route: 'NO_MATCH',
|
||||
reasoning: 'The query does not relate to any available route.',
|
||||
}),
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 100, output: 20, total: 120 },
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: JSON.stringify({
|
||||
route: 'NO_MATCH',
|
||||
reasoning: 'The query does not relate to any available route.',
|
||||
}),
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 100, output: 20, total: 120 },
|
||||
})
|
||||
|
||||
await expect(handler.execute(mockContext, mockRouterV2Block, inputs)).rejects.toThrow(
|
||||
@@ -992,19 +921,13 @@ describe('RouterBlockHandler V2', () => {
|
||||
routes: JSON.stringify([{ id: 'route-1', title: 'Route 1', value: 'Description' }]),
|
||||
}
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({
|
||||
route: 'invalid-route',
|
||||
reasoning: 'Some reasoning',
|
||||
}),
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 100, output: 20, total: 120 },
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: JSON.stringify({
|
||||
route: 'invalid-route',
|
||||
reasoning: 'Some reasoning',
|
||||
}),
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 100, output: 20, total: 120 },
|
||||
})
|
||||
|
||||
await expect(handler.execute(mockContext, mockRouterV2Block, inputs)).rejects.toThrow(
|
||||
@@ -1020,19 +943,13 @@ describe('RouterBlockHandler V2', () => {
|
||||
routes: [{ id: 'route-1', title: 'Route 1', value: 'Description' }],
|
||||
}
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: JSON.stringify({
|
||||
route: 'route-1',
|
||||
reasoning: 'Matched route 1',
|
||||
}),
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 100, output: 20, total: 120 },
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: JSON.stringify({
|
||||
route: 'route-1',
|
||||
reasoning: 'Matched route 1',
|
||||
}),
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 100, output: 20, total: 120 },
|
||||
})
|
||||
|
||||
const result = await handler.execute(mockContext, mockRouterV2Block, inputs)
|
||||
@@ -1084,16 +1001,10 @@ describe('RouterBlockHandler V2', () => {
|
||||
routes: JSON.stringify([{ id: 'route-1', title: 'Route 1', value: 'Description' }]),
|
||||
}
|
||||
|
||||
mockFetch.mockImplementationOnce(() => {
|
||||
return Promise.resolve({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content: 'route-1',
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 100, output: 5, total: 105 },
|
||||
}),
|
||||
})
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content: 'route-1',
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 100, output: 5, total: 105 },
|
||||
})
|
||||
|
||||
const result = await handler.execute(mockContext, mockRouterV2Block, inputs)
|
||||
@@ -1111,14 +1022,10 @@ describe('RouterBlockHandler V2', () => {
|
||||
routes: [{ id: 'route-1', title: 'Route 1', value: 'Description' }],
|
||||
}
|
||||
|
||||
mockFetch.mockResolvedValueOnce({
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
content,
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 100, output: 5, total: 105 },
|
||||
}),
|
||||
mockExecuteProviderRequest.mockResolvedValueOnce({
|
||||
content,
|
||||
model: 'gpt-4o',
|
||||
tokens: { input: 100, output: 5, total: 105 },
|
||||
})
|
||||
|
||||
await expect(handler.execute(mockContext, mockRouterV2Block, inputs)).rejects.toThrow(content)
|
||||
|
||||
@@ -1,11 +1,5 @@
|
||||
import { createLogger } from '@sim/logger'
|
||||
import { getInternalApiBaseUrl } from '@/lib/core/utils/urls'
|
||||
import {
|
||||
addModelInputProvenanceToRequest,
|
||||
createModelInputProvenanceRequestMetadata,
|
||||
markModelInputProjected,
|
||||
projectResolvedModelInput,
|
||||
} from '@/lib/execution/model-input-provenance'
|
||||
import { projectResolvedModelInput } from '@/lib/execution/model-input-provenance'
|
||||
import {
|
||||
type AutoRoutingResult,
|
||||
addAutoRoutingCost,
|
||||
@@ -23,7 +17,7 @@ import {
|
||||
ROUTER,
|
||||
} from '@/executor/constants'
|
||||
import type { BlockHandler, ExecutionContext } from '@/executor/types'
|
||||
import { buildAuthHeaders } from '@/executor/utils/http'
|
||||
import { executeBlockProviderRequest } from '@/executor/utils/provider-request'
|
||||
import { refuseResolvedSecretProjection } from '@/executor/utils/resolved-secret-projection-refusal'
|
||||
import type { ResolvedSecretInputPath } from '@/executor/utils/resolved-secret-trace-registry'
|
||||
import { resolveVertexCredential } from '@/executor/utils/vertex-credential'
|
||||
@@ -103,9 +97,6 @@ export class RouterBlockHandler implements BlockHandler {
|
||||
}
|
||||
|
||||
try {
|
||||
const url = new URL('/api/providers', getInternalApiBaseUrl())
|
||||
if (ctx.userId) url.searchParams.set('userId', ctx.userId)
|
||||
|
||||
const messages = [{ role: 'user', content: routerConfig.prompt }]
|
||||
const systemPrompt = generateRouterPrompt(routerConfig.prompt, targetBlocks)
|
||||
const resolved = await this.resolveModel(
|
||||
@@ -147,36 +138,13 @@ export class RouterBlockHandler implements BlockHandler {
|
||||
workspaceId: ctx.workspaceId,
|
||||
}
|
||||
|
||||
const headers = new Headers(await buildAuthHeaders(ctx.userId))
|
||||
const modelInputMetadata = createModelInputProvenanceRequestMetadata(
|
||||
modelInputProjection.registry,
|
||||
promptModelInputPaths
|
||||
)
|
||||
const requestBody = addModelInputProvenanceToRequest(
|
||||
{ provider: providerId, ...providerRequest },
|
||||
headers,
|
||||
modelInputMetadata
|
||||
)
|
||||
if (modelInputMetadata) markModelInputProjected(headers)
|
||||
const response = await fetch(url.toString(), {
|
||||
method: 'POST',
|
||||
headers,
|
||||
body: JSON.stringify(requestBody),
|
||||
const result = await executeBlockProviderRequest({
|
||||
ctx,
|
||||
providerId,
|
||||
request: providerRequest,
|
||||
resolvedSecretTraceRegistry: modelInputProjection.registry,
|
||||
})
|
||||
|
||||
if (!response.ok) {
|
||||
let errorMessage = `Provider API request failed with status ${response.status}`
|
||||
try {
|
||||
const errorData = await response.json()
|
||||
if (errorData.error) {
|
||||
errorMessage = errorData.error
|
||||
}
|
||||
} catch (_e) {}
|
||||
throw new Error(errorMessage)
|
||||
}
|
||||
|
||||
const result = await response.json()
|
||||
|
||||
const chosenBlockId = result.content.trim().toLowerCase()
|
||||
const chosenBlock = targetBlocks?.find((b) => b.id === chosenBlockId)
|
||||
|
||||
@@ -291,9 +259,6 @@ export class RouterBlockHandler implements BlockHandler {
|
||||
}
|
||||
|
||||
try {
|
||||
const url = new URL('/api/providers', getInternalApiBaseUrl())
|
||||
if (ctx.userId) url.searchParams.set('userId', ctx.userId)
|
||||
|
||||
const messages = [{ role: 'user', content: routerConfig.context }]
|
||||
const systemPrompt = generateRouterV2Prompt(routerConfig.context, modelRoutes)
|
||||
const resolved = await this.resolveModel(
|
||||
@@ -354,36 +319,13 @@ export class RouterBlockHandler implements BlockHandler {
|
||||
},
|
||||
}
|
||||
|
||||
const headers = new Headers(await buildAuthHeaders(ctx.userId))
|
||||
const modelInputMetadata = createModelInputProvenanceRequestMetadata(
|
||||
modelInputProjection.registry,
|
||||
modelInputPaths
|
||||
)
|
||||
const requestBody = addModelInputProvenanceToRequest(
|
||||
{ provider: providerId, ...providerRequest },
|
||||
headers,
|
||||
modelInputMetadata
|
||||
)
|
||||
if (modelInputMetadata) markModelInputProjected(headers)
|
||||
const response = await fetch(url.toString(), {
|
||||
method: 'POST',
|
||||
headers,
|
||||
body: JSON.stringify(requestBody),
|
||||
const result = await executeBlockProviderRequest({
|
||||
ctx,
|
||||
providerId,
|
||||
request: providerRequest,
|
||||
resolvedSecretTraceRegistry: modelInputProjection.registry,
|
||||
})
|
||||
|
||||
if (!response.ok) {
|
||||
let errorMessage = `Provider API request failed with status ${response.status}`
|
||||
try {
|
||||
const errorData = await response.json()
|
||||
if (errorData.error) {
|
||||
errorMessage = errorData.error
|
||||
}
|
||||
} catch (_e) {}
|
||||
throw new Error(errorMessage)
|
||||
}
|
||||
|
||||
const result = await response.json()
|
||||
|
||||
let chosenRouteId: string
|
||||
let reasoning = ''
|
||||
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
import { createLogger } from '@sim/logger'
|
||||
import { checkWorkspaceAccess } from '@/lib/workspaces/permissions/utils'
|
||||
import type { ExecutionContext } from '@/executor/types'
|
||||
import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
|
||||
import { executeProviderRequest } from '@/providers'
|
||||
import type { ProviderRequest, ProviderResponse } from '@/providers/types'
|
||||
|
||||
const logger = createLogger('ExecutorProviderRequest')
|
||||
|
||||
interface ExecuteBlockProviderRequestInput {
|
||||
ctx: ExecutionContext
|
||||
providerId: string
|
||||
request: ProviderRequest
|
||||
/** Supplied in place of the provenance envelope the HTTP boundary serialized and re-imported. */
|
||||
resolvedSecretTraceRegistry: ResolvedSecretTraceRegistry | undefined
|
||||
}
|
||||
|
||||
/**
|
||||
* Runs one non-streaming provider request for a block handler in-process, replacing the
|
||||
* executor's `POST /api/providers` round trip. The route's two admission checks are
|
||||
* reproduced so the outcome is unchanged: an internal token with no user is rejected (the
|
||||
* executor mints it from `ctx.userId`), and an execution subject who has left the billed
|
||||
* workspace is rejected. The route's remaining work is already done by the caller or lives
|
||||
* inside `executeProviderRequest`.
|
||||
*/
|
||||
export async function executeBlockProviderRequest({
|
||||
ctx,
|
||||
providerId,
|
||||
request,
|
||||
resolvedSecretTraceRegistry,
|
||||
}: ExecuteBlockProviderRequestInput): Promise<ProviderResponse> {
|
||||
if (!ctx.userId) {
|
||||
throw new Error('Unauthorized')
|
||||
}
|
||||
|
||||
if (request.workspaceId) {
|
||||
const workspaceAccess = await checkWorkspaceAccess(request.workspaceId, ctx.userId)
|
||||
if (!workspaceAccess.hasAccess) {
|
||||
throw new Error('Forbidden')
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* No `executionContext`: it is only inherited by model-emitted tool calls, and the route
|
||||
* this replaces never carried one. The whole context is omitted when there is no registry
|
||||
* rather than passed carrying `undefined` — `executeProviderTool` reads that as missing
|
||||
* provenance and fails the call closed with no error text.
|
||||
*/
|
||||
const response = await executeProviderRequest(
|
||||
providerId,
|
||||
{ ...request, userId: ctx.userId },
|
||||
resolvedSecretTraceRegistry ? { resolvedSecretTraceRegistry } : undefined
|
||||
)
|
||||
|
||||
if (
|
||||
response instanceof ReadableStream ||
|
||||
(typeof response === 'object' && response !== null && 'stream' in response)
|
||||
) {
|
||||
logger.error('Provider returned a stream for a non-streaming block request', { providerId })
|
||||
throw new Error('Provider returned a streaming response for a non-streaming request')
|
||||
}
|
||||
|
||||
logger.info('Provider request completed', {
|
||||
providerId,
|
||||
model: request.model,
|
||||
workflowId: ctx.workflowId,
|
||||
})
|
||||
|
||||
return response
|
||||
}
|
||||
@@ -22,8 +22,8 @@ import {
|
||||
import { updateWorkspacePermissionsContract } from '@/lib/api/contracts/workspaces'
|
||||
import { organizationKeys } from '@/hooks/queries/organization'
|
||||
import { refreshSessionQuery } from '@/hooks/queries/session'
|
||||
import { subscriptionKeys } from '@/hooks/queries/subscription'
|
||||
import { workspaceCredentialKeys } from '@/hooks/queries/utils/credential-keys'
|
||||
import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys'
|
||||
import { workspaceKeys } from '@/hooks/queries/workspace'
|
||||
|
||||
export const invitationKeys = {
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
/**
|
||||
* @vitest-environment node
|
||||
*/
|
||||
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
|
||||
const mocks = vi.hoisted(() => ({
|
||||
requestJson: vi.fn(),
|
||||
useMutation: vi.fn(),
|
||||
invalidateQueries: vi.fn(),
|
||||
}))
|
||||
|
||||
vi.mock('@tanstack/react-query', () => ({
|
||||
keepPreviousData: Symbol('keepPreviousData'),
|
||||
useInfiniteQuery: vi.fn(),
|
||||
useMutation: mocks.useMutation,
|
||||
useQuery: vi.fn(),
|
||||
useQueryClient: vi.fn(() => ({ invalidateQueries: mocks.invalidateQueries })),
|
||||
}))
|
||||
|
||||
vi.mock('@sim/emcn', () => ({
|
||||
toast: { error: vi.fn(), success: vi.fn() },
|
||||
}))
|
||||
|
||||
vi.mock('@/lib/api/client/request', () => ({
|
||||
requestJson: mocks.requestJson,
|
||||
}))
|
||||
|
||||
import { useBulkDocumentOperation, useDeleteDocument } from '@/hooks/queries/kb/knowledge'
|
||||
import { knowledgeKeys } from '@/hooks/queries/utils/knowledge-keys'
|
||||
|
||||
interface CapturedMutation {
|
||||
onSettled: (data: unknown, error: unknown, variables: Record<string, unknown>) => void
|
||||
}
|
||||
|
||||
function captureMutation(build: () => unknown): CapturedMutation {
|
||||
let captured: CapturedMutation | undefined
|
||||
mocks.useMutation.mockImplementation((options: CapturedMutation) => {
|
||||
captured = options
|
||||
return {}
|
||||
})
|
||||
build()
|
||||
if (!captured) throw new Error('useMutation was not called')
|
||||
return captured
|
||||
}
|
||||
|
||||
describe('knowledge document mutations', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
})
|
||||
|
||||
it('invalidates the knowledge-base lists when a document is deleted', () => {
|
||||
const mutation = captureMutation(() => useDeleteDocument())
|
||||
|
||||
mutation.onSettled(undefined, undefined, { knowledgeBaseId: 'kb-1', documentId: 'doc-1' })
|
||||
|
||||
expect(mocks.invalidateQueries).toHaveBeenCalledWith({ queryKey: knowledgeKeys.lists() })
|
||||
})
|
||||
|
||||
it('invalidates the knowledge-base lists on a bulk delete', () => {
|
||||
const mutation = captureMutation(() => useBulkDocumentOperation())
|
||||
|
||||
mutation.onSettled(undefined, undefined, { knowledgeBaseId: 'kb-1', operation: 'delete' })
|
||||
|
||||
expect(mocks.invalidateQueries).toHaveBeenCalledWith({ queryKey: knowledgeKeys.lists() })
|
||||
})
|
||||
|
||||
it('leaves the knowledge-base lists alone on a bulk enable', () => {
|
||||
const mutation = captureMutation(() => useBulkDocumentOperation())
|
||||
|
||||
mutation.onSettled(undefined, undefined, { knowledgeBaseId: 'kb-1', operation: 'enable' })
|
||||
|
||||
expect(mocks.invalidateQueries).not.toHaveBeenCalledWith({ queryKey: knowledgeKeys.lists() })
|
||||
})
|
||||
})
|
||||
@@ -539,6 +539,10 @@ export function useDeleteDocument() {
|
||||
queryClient.invalidateQueries({
|
||||
queryKey: knowledgeKeys.detail(knowledgeBaseId),
|
||||
})
|
||||
/** The knowledge-base list rows carry `docCount`, so removing a document changes them too. */
|
||||
queryClient.invalidateQueries({
|
||||
queryKey: knowledgeKeys.lists(),
|
||||
})
|
||||
},
|
||||
})
|
||||
}
|
||||
@@ -573,10 +577,16 @@ export function useBulkDocumentOperation() {
|
||||
|
||||
return useMutation({
|
||||
mutationFn: bulkDocumentOperation,
|
||||
onSettled: (_data, _error, { knowledgeBaseId }) => {
|
||||
onSettled: (_data, _error, { knowledgeBaseId, operation }) => {
|
||||
queryClient.invalidateQueries({
|
||||
queryKey: knowledgeKeys.detail(knowledgeBaseId),
|
||||
})
|
||||
/** Only a bulk delete changes the `docCount` the knowledge-base list rows render. */
|
||||
if (operation === 'delete') {
|
||||
queryClient.invalidateQueries({
|
||||
queryKey: knowledgeKeys.lists(),
|
||||
})
|
||||
}
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
@@ -56,6 +56,18 @@ function defineServices(): ServiceInfo[] {
|
||||
return servicesList
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolves the service catalog merged with the caller's connections.
|
||||
*
|
||||
* A failed request resolves with the bare catalog rather than rejecting, so
|
||||
* consumers keep correct service names and ids when the merge data is
|
||||
* unavailable. The cost is that `isConnected`/`accounts` then report *unknown*
|
||||
* as *disconnected*, which the result cannot distinguish. Read connection
|
||||
* state from the workspace credentials query (`useWorkspaceCredentials`), which
|
||||
* surfaces its own errors; a consumer that must branch on `isConnected` here
|
||||
* needs this fallback removed first, or it will tell a connected user they are
|
||||
* not.
|
||||
*/
|
||||
async function fetchOAuthConnections(signal?: AbortSignal): Promise<ServiceInfo[]> {
|
||||
try {
|
||||
const serviceDefinitions = defineServices()
|
||||
|
||||
@@ -41,8 +41,8 @@ import {
|
||||
import { client } from '@/lib/auth/auth-client'
|
||||
import { isEnterprise, isPaid, isTeam } from '@/lib/billing/plan-helpers'
|
||||
import { hasPaidSubscriptionStatus } from '@/lib/billing/subscriptions/utils'
|
||||
import { subscriptionKeys } from '@/hooks/queries/subscription'
|
||||
import { workspaceCredentialKeys } from '@/hooks/queries/utils/credential-keys'
|
||||
import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys'
|
||||
import { workspaceKeys } from '@/hooks/queries/workspace'
|
||||
|
||||
const logger = createLogger('OrganizationQueries')
|
||||
|
||||
@@ -27,7 +27,13 @@ export const scheduleKeys = {
|
||||
details: () => [...scheduleKeys.all, 'detail'] as const,
|
||||
schedule: (workflowId: string, blockId: string) =>
|
||||
[...scheduleKeys.details(), workflowId, blockId] as const,
|
||||
byId: (scheduleId: string) => [...scheduleKeys.details(), scheduleId] as const,
|
||||
/**
|
||||
* By-id reads sit under their own segment rather than directly under `details()`:
|
||||
* a bare `[...details(), scheduleId]` is a prefix of `schedule(scheduleId, blockId)`,
|
||||
* so the two addressings of the same schedule would alias in the cache.
|
||||
*/
|
||||
byIds: () => [...scheduleKeys.details(), 'by-id'] as const,
|
||||
byId: (scheduleId: string) => [...scheduleKeys.byIds(), scheduleId] as const,
|
||||
}
|
||||
|
||||
export type ScheduleData = WorkflowScheduleRow
|
||||
@@ -196,7 +202,7 @@ export function useReactivateSchedule() {
|
||||
body: { action: 'reactivate' },
|
||||
})
|
||||
|
||||
return { workflowId, blockId, workspaceId }
|
||||
return { scheduleId, workflowId, blockId, workspaceId }
|
||||
},
|
||||
onSuccess: ({ workflowId, blockId }) => {
|
||||
logger.info('Schedule reactivated', { workflowId, blockId })
|
||||
@@ -206,9 +212,10 @@ export function useReactivateSchedule() {
|
||||
},
|
||||
onSettled: async (data) => {
|
||||
if (!data) return
|
||||
const { workflowId, blockId, workspaceId } = data
|
||||
const { scheduleId, workflowId, blockId, workspaceId } = data
|
||||
await Promise.all([
|
||||
queryClient.invalidateQueries({ queryKey: scheduleKeys.schedule(workflowId, blockId) }),
|
||||
queryClient.invalidateQueries({ queryKey: scheduleKeys.byId(scheduleId) }),
|
||||
workspaceId
|
||||
? queryClient.invalidateQueries({ queryKey: scheduleKeys.list(workspaceId) })
|
||||
: Promise.resolve(),
|
||||
@@ -242,6 +249,8 @@ export function useRedeployWorkflowSchedule() {
|
||||
const { workflowId, blockId } = data
|
||||
await Promise.all([
|
||||
queryClient.invalidateQueries({ queryKey: scheduleKeys.schedule(workflowId, blockId) }),
|
||||
/** A redeploy recreates the schedule; the id-keyed reads are a separate subtree. */
|
||||
queryClient.invalidateQueries({ queryKey: scheduleKeys.byIds() }),
|
||||
queryClient.invalidateQueries({ queryKey: deploymentKeys.info(workflowId) }),
|
||||
queryClient.invalidateQueries({ queryKey: deploymentKeys.versions(workflowId) }),
|
||||
])
|
||||
|
||||
@@ -13,6 +13,8 @@ import {
|
||||
updateUsageLimitContract,
|
||||
} from '@/lib/api/contracts/subscription'
|
||||
import { organizationKeys } from '@/hooks/queries/organization'
|
||||
import { invalidateWorkspaceUsage } from '@/hooks/queries/utils/invalidate-usage'
|
||||
import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys'
|
||||
import { workspaceKeys } from '@/hooks/queries/workspace'
|
||||
|
||||
export type { SubscriptionApiResponse }
|
||||
@@ -21,19 +23,6 @@ export const SUBSCRIPTION_DATA_STALE_TIME = 5 * 60 * 1000
|
||||
export const USAGE_LIMIT_STALE_TIME = 30 * 1000
|
||||
export const INVOICES_STALE_TIME = 5 * 60 * 1000
|
||||
|
||||
/**
|
||||
* Query key factories for subscription-related queries
|
||||
*/
|
||||
export const subscriptionKeys = {
|
||||
all: ['subscription'] as const,
|
||||
users: () => [...subscriptionKeys.all, 'user'] as const,
|
||||
user: (includeOrg?: boolean) => [...subscriptionKeys.users(), { includeOrg }] as const,
|
||||
usage: () => [...subscriptionKeys.all, 'usage'] as const,
|
||||
invoicesAll: () => [...subscriptionKeys.all, 'invoices'] as const,
|
||||
invoices: (context: 'user' | 'organization' = 'user', organizationId?: string) =>
|
||||
[...subscriptionKeys.invoicesAll(), context, organizationId ?? ''] as const,
|
||||
}
|
||||
|
||||
/**
|
||||
* Fetch user subscription data
|
||||
* @param includeOrg - Whether to include organization role data
|
||||
@@ -265,6 +254,7 @@ export function useUpdateUsageLimit() {
|
||||
return Promise.all([
|
||||
queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }),
|
||||
queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() }),
|
||||
invalidateWorkspaceUsage(queryClient),
|
||||
])
|
||||
},
|
||||
})
|
||||
@@ -291,6 +281,7 @@ export function useUpgradeSubscription() {
|
||||
queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() }),
|
||||
queryClient.invalidateQueries({ queryKey: subscriptionKeys.invoicesAll() }),
|
||||
queryClient.invalidateQueries({ queryKey: workspaceKeys.lists() }),
|
||||
invalidateWorkspaceUsage(queryClient),
|
||||
...(variables.orgId
|
||||
? [
|
||||
queryClient.invalidateQueries({
|
||||
@@ -328,6 +319,7 @@ export function usePurchaseCredits() {
|
||||
return Promise.all([
|
||||
queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }),
|
||||
queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() }),
|
||||
invalidateWorkspaceUsage(queryClient),
|
||||
...(variables.orgId
|
||||
? [
|
||||
queryClient.invalidateQueries({
|
||||
|
||||
@@ -69,6 +69,13 @@ import { tableKeys } from '@/hooks/queries/utils/table-keys'
|
||||
const TABLE_ID = 'tbl-1'
|
||||
const WORKSPACE_ID = 'ws-1'
|
||||
|
||||
/**
|
||||
* Where a paged row list actually lives. Seeding at the bare `rowsRoot` prefix would
|
||||
* exercise a key no hook writes, and would keep matching a cache walk that has been
|
||||
* narrowed away from the `find` sibling hanging off the same parent.
|
||||
*/
|
||||
const ROWS_KEY = tableKeys.infiniteRows(TABLE_ID, tableRowsParamsKey({ pageSize: 1000 }))
|
||||
|
||||
function setCache(key: readonly unknown[], value: unknown) {
|
||||
cacheStore.set(JSON.stringify(key), value)
|
||||
}
|
||||
@@ -96,7 +103,7 @@ describe('useDeleteColumn optimistic update', () => {
|
||||
columnWidths: { name: 200, age: 100 },
|
||||
},
|
||||
})
|
||||
setCache(tableKeys.rowsRoot(TABLE_ID), {
|
||||
setCache(ROWS_KEY, {
|
||||
rows: [
|
||||
{ id: 'r1', data: { name: 'a', age: 1 } },
|
||||
{ id: 'r2', data: { name: 'b', age: 2 } },
|
||||
@@ -114,9 +121,7 @@ describe('useDeleteColumn optimistic update', () => {
|
||||
expect(detail?.schema.columns.map((c) => c.name)).toEqual(['name'])
|
||||
expect(detail?.metadata.columnWidths).toEqual({ name: 200 })
|
||||
|
||||
const rows = getCache<{ rows: Array<{ data: Record<string, unknown> }> }>(
|
||||
tableKeys.rowsRoot(TABLE_ID)
|
||||
)
|
||||
const rows = getCache<{ rows: Array<{ data: Record<string, unknown> }> }>(ROWS_KEY)
|
||||
expect(rows?.rows.every((r) => !('age' in r.data))).toBe(true)
|
||||
expect(rows?.rows[0]?.data).toEqual({ name: 'a' })
|
||||
|
||||
@@ -124,6 +129,33 @@ describe('useDeleteColumn optimistic update', () => {
|
||||
expect(ctx?.rowSnapshots?.length).toBeGreaterThan(0)
|
||||
})
|
||||
|
||||
/**
|
||||
* The `find` cache hangs off the same `rowsRoot` parent as the paged rows but holds
|
||||
* `{matches, truncated}` — no `pages`, no `rows`. A cache walk starting at the shared
|
||||
* parent reaches it and throws inside `onMutate`, rejecting the mutation before it ever
|
||||
* reaches the server: search a table, dismiss the search, then edit a cell.
|
||||
*/
|
||||
it('survives a cached search result hanging off the shared rows prefix', async () => {
|
||||
setCache(tableKeys.detail(TABLE_ID), {
|
||||
id: TABLE_ID,
|
||||
schema: { columns: [{ name: 'age', type: 'number' }] },
|
||||
})
|
||||
setCache(ROWS_KEY, {
|
||||
rows: [{ id: 'r1', data: { age: 1 } }],
|
||||
totalCount: 1,
|
||||
})
|
||||
setCache(tableKeys.find(TABLE_ID, 'q'), { matches: [{ rowId: 'r1', column: 'age' }] })
|
||||
|
||||
const hook = useDeleteColumn({ workspaceId: WORKSPACE_ID, tableId: TABLE_ID })
|
||||
|
||||
await expect(hook.onMutate?.('age')).resolves.toBeDefined()
|
||||
|
||||
const rows = getCache<{ rows: Array<{ data: Record<string, unknown> }> }>(ROWS_KEY)
|
||||
expect(rows?.rows[0]?.data).toEqual({})
|
||||
/** The find entry is match coordinates, not row values — it must be left untouched. */
|
||||
expect(getCache<{ matches: unknown[] }>(tableKeys.find(TABLE_ID, 'q'))?.matches).toHaveLength(1)
|
||||
})
|
||||
|
||||
it('rolls back schema and rows on error using snapshots', async () => {
|
||||
const originalDetail = {
|
||||
id: TABLE_ID,
|
||||
@@ -135,7 +167,7 @@ describe('useDeleteColumn optimistic update', () => {
|
||||
totalCount: 1,
|
||||
}
|
||||
setCache(tableKeys.detail(TABLE_ID), originalDetail)
|
||||
setCache(tableKeys.rowsRoot(TABLE_ID), originalRows)
|
||||
setCache(ROWS_KEY, originalRows)
|
||||
|
||||
const hook = useDeleteColumn({ workspaceId: WORKSPACE_ID, tableId: TABLE_ID })
|
||||
const ctx = await hook.onMutate?.('age')
|
||||
@@ -145,7 +177,7 @@ describe('useDeleteColumn optimistic update', () => {
|
||||
hook.onError?.(new Error('boom'), 'age', ctx)
|
||||
|
||||
expect(getCache(tableKeys.detail(TABLE_ID))).toEqual(originalDetail)
|
||||
expect(getCache(tableKeys.rowsRoot(TABLE_ID))).toEqual(originalRows)
|
||||
expect(getCache(ROWS_KEY)).toEqual(originalRows)
|
||||
})
|
||||
|
||||
it('invalidates schema, rows, and lists in onSettled', () => {
|
||||
@@ -194,7 +226,7 @@ describe('useUpdateColumn optimistic update', () => {
|
||||
id: TABLE_ID,
|
||||
schema: { columns: [{ name: 'age', type: 'number' }] },
|
||||
})
|
||||
setCache(tableKeys.rowsRoot(TABLE_ID), {
|
||||
setCache(ROWS_KEY, {
|
||||
rows: [
|
||||
{ id: 'r1', data: { age: 30 } },
|
||||
{ id: 'r2', data: { age: 40 } },
|
||||
@@ -207,9 +239,7 @@ describe('useUpdateColumn optimistic update', () => {
|
||||
|
||||
// Row data is id-keyed; a rename never moves it. The stored key (`age`)
|
||||
// becomes the column's stamped id, so cells stay reachable via getColumnId.
|
||||
const rows = getCache<{ rows: Array<{ data: Record<string, unknown> }> }>(
|
||||
tableKeys.rowsRoot(TABLE_ID)
|
||||
)
|
||||
const rows = getCache<{ rows: Array<{ data: Record<string, unknown> }> }>(ROWS_KEY)
|
||||
expect(rows?.rows[0]?.data).toEqual({ age: 30 })
|
||||
expect(rows?.rows[1]?.data).toEqual({ age: 40 })
|
||||
|
||||
@@ -265,7 +295,7 @@ describe('useDeleteColumn case-insensitive row cleanup', () => {
|
||||
id: TABLE_ID,
|
||||
schema: { columns: [{ name: 'Age', type: 'number' }] },
|
||||
})
|
||||
setCache(tableKeys.rowsRoot(TABLE_ID), {
|
||||
setCache(ROWS_KEY, {
|
||||
rows: [{ id: 'r1', data: { Age: 30, name: 'a' } }],
|
||||
totalCount: 1,
|
||||
})
|
||||
@@ -273,9 +303,7 @@ describe('useDeleteColumn case-insensitive row cleanup', () => {
|
||||
const hook = useDeleteColumn({ workspaceId: WORKSPACE_ID, tableId: TABLE_ID })
|
||||
await hook.onMutate?.('age')
|
||||
|
||||
const rows = getCache<{ rows: Array<{ data: Record<string, unknown> }> }>(
|
||||
tableKeys.rowsRoot(TABLE_ID)
|
||||
)
|
||||
const rows = getCache<{ rows: Array<{ data: Record<string, unknown> }> }>(ROWS_KEY)
|
||||
expect(rows?.rows[0]?.data).toEqual({ name: 'a' })
|
||||
})
|
||||
})
|
||||
|
||||
@@ -832,6 +832,11 @@ function withOptimisticAutoFireExec(groups: WorkflowGroup[], row: TableRow): Tab
|
||||
/**
|
||||
* Apply a row-level transformation to all cached infinite row queries for this
|
||||
* table. Used for cell edits where positions don't change.
|
||||
*
|
||||
* Walks {@link tableKeys.infiniteRowsRoot} rather than `rowsRoot`: the latter is a
|
||||
* shared parent, and handing this updater a `find` entry — a flat
|
||||
* {@link TableFindResult}, not pages — throws on `old.pages` inside `onMutate`, so
|
||||
* the whole cell edit would reject before reaching the server.
|
||||
*/
|
||||
function patchCachedRows(
|
||||
queryClient: ReturnType<typeof useQueryClient>,
|
||||
@@ -839,7 +844,7 @@ function patchCachedRows(
|
||||
patchRow: (row: TableRow) => TableRow
|
||||
) {
|
||||
queryClient.setQueriesData<InfiniteData<TableRowsResponse, TableRowsPageParam>>(
|
||||
{ queryKey: tableKeys.rowsRoot(tableId), exact: false },
|
||||
{ queryKey: tableKeys.infiniteRowsRoot(tableId), exact: false },
|
||||
(old) => {
|
||||
if (!old) return old
|
||||
return {
|
||||
@@ -2187,7 +2192,7 @@ export async function snapshotAndMutateRows(
|
||||
): Promise<RowsCacheSnapshots> {
|
||||
const scope = options?.onlyKey
|
||||
? ({ queryKey: options.onlyKey, exact: true } as const)
|
||||
: ({ queryKey: tableKeys.rowsRoot(tableId) } as const)
|
||||
: ({ queryKey: tableKeys.infiniteRowsRoot(tableId) } as const)
|
||||
if (options?.cancelInFlight !== false) {
|
||||
await queryClient.cancelQueries(scope)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
import type { QueryClient } from '@tanstack/react-query'
|
||||
import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys'
|
||||
import { workspaceUsageKeys } from '@/hooks/queries/utils/workspace-usage-keys'
|
||||
|
||||
/**
|
||||
* Usage is written asynchronously as a run settles, so a refetch fired on completion
|
||||
* races the write and re-reads the old balance.
|
||||
*/
|
||||
const USAGE_SETTLE_DELAY_MS = 1000
|
||||
|
||||
/**
|
||||
* Invalidates the workspace credit/usage reads after anything that moves the balance.
|
||||
* Both families are keyed per workspace but derive from one billing account, so the
|
||||
* family prefixes are what must refetch.
|
||||
*/
|
||||
export function invalidateWorkspaceUsage(queryClient: QueryClient) {
|
||||
return Promise.all([
|
||||
queryClient.invalidateQueries({ queryKey: workspaceUsageKeys.creditAvailabilities() }),
|
||||
queryClient.invalidateQueries({ queryKey: workspaceUsageKeys.gates() }),
|
||||
])
|
||||
}
|
||||
|
||||
/**
|
||||
* Refreshes the billing reads a run touches, after {@link USAGE_SETTLE_DELAY_MS}.
|
||||
* Shared by workflow execution and wand generation.
|
||||
*/
|
||||
export function scheduleUsageRefresh(queryClient: QueryClient) {
|
||||
setTimeout(() => {
|
||||
void queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() })
|
||||
void invalidateWorkspaceUsage(queryClient)
|
||||
}, USAGE_SETTLE_DELAY_MS)
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
/**
|
||||
* @vitest-environment node
|
||||
*/
|
||||
import { afterEach, describe, expect, it, vi } from 'vitest'
|
||||
import { generateTempId } from '@/hooks/queries/utils/optimistic-mutation'
|
||||
|
||||
describe('generateTempId', () => {
|
||||
afterEach(() => {
|
||||
vi.useRealTimers()
|
||||
})
|
||||
|
||||
it('is unique for ids created within the same millisecond', () => {
|
||||
vi.useFakeTimers()
|
||||
vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z'))
|
||||
|
||||
const ids = new Set(Array.from({ length: 100 }, () => generateTempId('temp-folder')))
|
||||
|
||||
expect(ids.size).toBe(100)
|
||||
})
|
||||
|
||||
it('keeps the prefix so callers can identify optimistic rows', () => {
|
||||
expect(generateTempId('temp-folder').startsWith('temp-folder-')).toBe(true)
|
||||
})
|
||||
})
|
||||
@@ -1,4 +1,5 @@
|
||||
import { createLogger } from '@sim/logger'
|
||||
import { generateId } from '@sim/utils/id'
|
||||
import type { QueryClient } from '@tanstack/react-query'
|
||||
|
||||
const logger = createLogger('OptimisticMutation')
|
||||
@@ -72,6 +73,12 @@ export function createOptimisticMutationHandlers<TData, TVariables, TItem>(
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Placeholder id for an optimistic row, held only until the server response
|
||||
* replaces it. Uses `generateId()` rather than a timestamp so two rows created
|
||||
* in the same millisecond cannot collide — a collision would make
|
||||
* `replaceOptimisticEntry` overwrite both entries with one server row.
|
||||
*/
|
||||
export function generateTempId(prefix: string): string {
|
||||
return `${prefix}-${Date.now()}`
|
||||
return `${prefix}-${generateId()}`
|
||||
}
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
/**
|
||||
* React Query key factory for subscription and billing reads. Standalone so the shared
|
||||
* billing invalidations can use it without closing an import cycle through the hook module.
|
||||
*/
|
||||
export const subscriptionKeys = {
|
||||
all: ['subscription'] as const,
|
||||
users: () => [...subscriptionKeys.all, 'user'] as const,
|
||||
user: (includeOrg?: boolean) => [...subscriptionKeys.users(), { includeOrg }] as const,
|
||||
usage: () => [...subscriptionKeys.all, 'usage'] as const,
|
||||
invoicesAll: () => [...subscriptionKeys.all, 'invoices'] as const,
|
||||
invoices: (context: 'user' | 'organization' = 'user', organizationId?: string) =>
|
||||
[...subscriptionKeys.invoicesAll(), context, organizationId ?? ''] as const,
|
||||
}
|
||||
@@ -25,8 +25,14 @@ export const tableKeys = {
|
||||
exportJobs: (workspaceId?: string) =>
|
||||
[...tableKeys.all, 'export-jobs', workspaceId ?? ''] as const,
|
||||
rowsRoot: (tableId: string) => [...tableKeys.detail(tableId), 'rows'] as const,
|
||||
/**
|
||||
* Prefix covering only the paged row lists. `rowsRoot` is a shared parent — `find`
|
||||
* hangs off it holding a different shape — so anything walking the cache for row
|
||||
* pages must start here.
|
||||
*/
|
||||
infiniteRowsRoot: (tableId: string) => [...tableKeys.rowsRoot(tableId), 'infinite'] as const,
|
||||
infiniteRows: (tableId: string, paramsKey: string) =>
|
||||
[...tableKeys.rowsRoot(tableId), 'infinite', paramsKey] as const,
|
||||
[...tableKeys.infiniteRowsRoot(tableId), paramsKey] as const,
|
||||
rowWrites: (tableId: string) => [...tableKeys.rowsRoot(tableId), 'write'] as const,
|
||||
find: (tableId: string, paramsKey: string) =>
|
||||
[...tableKeys.rowsRoot(tableId), 'find', paramsKey] as const,
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
/**
|
||||
* React Query key factory for the per-workspace credit and usage-gate reads. Standalone
|
||||
* for the same import-cycle reason as {@link file://./subscription-keys.ts}.
|
||||
*/
|
||||
export const workspaceUsageKeys = {
|
||||
all: ['workspace-usage'] as const,
|
||||
creditAvailabilities: () => [...workspaceUsageKeys.all, 'credit-availability'] as const,
|
||||
creditAvailability: (workspaceId: string) =>
|
||||
[...workspaceUsageKeys.creditAvailabilities(), workspaceId] as const,
|
||||
gates: () => [...workspaceUsageKeys.all, 'gate'] as const,
|
||||
gate: (workspaceId: string) => [...workspaceUsageKeys.gates(), workspaceId] as const,
|
||||
}
|
||||
@@ -521,6 +521,13 @@ export function useCloudStorageConfigured(enabled = true) {
|
||||
enabled,
|
||||
retry: false,
|
||||
staleTime: CLOUD_STORAGE_CONFIGURED_STALE_TIME,
|
||||
/**
|
||||
* Escapes the global `retryOnMount: false`: with an infinite `staleTime` and
|
||||
* `retry: false`, one transient error leaves this query errored for the tab's lifetime,
|
||||
* and the upload path reads "unknown" as "not configured" — disabling cloud uploads
|
||||
* until a full reload. The key is global, so navigation cannot recover it.
|
||||
*/
|
||||
retryOnMount: true,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -15,12 +15,13 @@ import {
|
||||
getWorkspaceCreditAvailabilityContract,
|
||||
getWorkspaceUsageGateContract,
|
||||
} from '@/lib/api/contracts/workspaces'
|
||||
import { invalidateWorkspaceUsage } from '@/hooks/queries/utils/invalidate-usage'
|
||||
import { workspaceUsageKeys } from '@/hooks/queries/utils/workspace-usage-keys'
|
||||
import {
|
||||
fetchWorkspaceCreditAvailability,
|
||||
fetchWorkspaceUsageGate,
|
||||
WORKSPACE_CREDIT_AVAILABILITY_STALE_TIME,
|
||||
WORKSPACE_USAGE_GATE_STALE_TIME,
|
||||
workspaceUsageKeys,
|
||||
} from '@/hooks/queries/workspace-usage'
|
||||
|
||||
describe('workspace usage gate query', () => {
|
||||
@@ -65,4 +66,18 @@ describe('workspace usage gate query', () => {
|
||||
signal,
|
||||
})
|
||||
})
|
||||
|
||||
it('invalidates both usage families so a spend or top-up refetches them', async () => {
|
||||
const invalidateQueries = vi.fn().mockResolvedValue(undefined)
|
||||
const queryClient = { invalidateQueries } as unknown as Parameters<
|
||||
typeof invalidateWorkspaceUsage
|
||||
>[0]
|
||||
|
||||
await invalidateWorkspaceUsage(queryClient)
|
||||
|
||||
expect(invalidateQueries).toHaveBeenCalledWith({
|
||||
queryKey: workspaceUsageKeys.creditAvailabilities(),
|
||||
})
|
||||
expect(invalidateQueries).toHaveBeenCalledWith({ queryKey: workspaceUsageKeys.gates() })
|
||||
})
|
||||
})
|
||||
|
||||
@@ -6,15 +6,7 @@ import {
|
||||
type WorkspaceCreditAvailability,
|
||||
type WorkspaceUsageGate,
|
||||
} from '@/lib/api/contracts/workspaces'
|
||||
|
||||
export const workspaceUsageKeys = {
|
||||
all: ['workspace-usage'] as const,
|
||||
creditAvailabilities: () => [...workspaceUsageKeys.all, 'credit-availability'] as const,
|
||||
creditAvailability: (workspaceId: string) =>
|
||||
[...workspaceUsageKeys.creditAvailabilities(), workspaceId] as const,
|
||||
gates: () => [...workspaceUsageKeys.all, 'gate'] as const,
|
||||
gate: (workspaceId: string) => [...workspaceUsageKeys.gates(), workspaceId] as const,
|
||||
}
|
||||
import { workspaceUsageKeys } from '@/hooks/queries/utils/workspace-usage-keys'
|
||||
|
||||
export const WORKSPACE_CREDIT_AVAILABILITY_STALE_TIME = 30 * 1000
|
||||
export const WORKSPACE_USAGE_GATE_STALE_TIME = 30 * 1000
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
/**
|
||||
* @vitest-environment node
|
||||
*/
|
||||
import { describe, expect, it } from 'vitest'
|
||||
import { cloudwatchSelectors } from '@/hooks/selectors/providers/cloudwatch/selectors'
|
||||
import type { SelectorQueryArgs } from '@/hooks/selectors/types'
|
||||
|
||||
const AWS_CONTEXT: SelectorQueryArgs['context'] = {
|
||||
awsAccessKeyId: 'AKIA',
|
||||
awsSecretAccessKey: 'secret',
|
||||
awsRegion: 'us-east-1',
|
||||
logGroupName: '/aws/lambda/fn',
|
||||
}
|
||||
|
||||
describe('cloudwatch selector query keys', () => {
|
||||
it.each([['cloudwatch.logGroups' as const], ['cloudwatch.logStreams' as const]])(
|
||||
'%s scopes its key by search, which fetchList forwards as `prefix`',
|
||||
(key) => {
|
||||
const definition = cloudwatchSelectors[key]
|
||||
const base: SelectorQueryArgs = { key, context: AWS_CONTEXT }
|
||||
|
||||
const noSearch = definition.getQueryKey(base)
|
||||
const withSearch = definition.getQueryKey({ ...base, search: 'api' })
|
||||
const otherSearch = definition.getQueryKey({ ...base, search: 'worker' })
|
||||
|
||||
expect(withSearch).not.toEqual(noSearch)
|
||||
expect(withSearch).not.toEqual(otherSearch)
|
||||
}
|
||||
)
|
||||
})
|
||||
@@ -20,11 +20,12 @@ export const cloudwatchSelectors = {
|
||||
key: 'cloudwatch.logGroups',
|
||||
contracts: [selectorContracts.cloudwatchLogGroupsSelectorContract],
|
||||
staleTime: SELECTOR_STALE,
|
||||
getQueryKey: ({ context }: SelectorQueryArgs) => [
|
||||
getQueryKey: ({ context, search }: SelectorQueryArgs) => [
|
||||
'selectors',
|
||||
'cloudwatch.logGroups',
|
||||
context.awsAccessKeyId ?? 'none',
|
||||
context.awsRegion ?? 'none',
|
||||
search ?? '',
|
||||
],
|
||||
enabled: ({ context }) =>
|
||||
Boolean(context.awsAccessKeyId && context.awsSecretAccessKey && context.awsRegion),
|
||||
@@ -51,12 +52,13 @@ export const cloudwatchSelectors = {
|
||||
key: 'cloudwatch.logStreams',
|
||||
contracts: [selectorContracts.cloudwatchLogStreamsSelectorContract],
|
||||
staleTime: SELECTOR_STALE,
|
||||
getQueryKey: ({ context }: SelectorQueryArgs) => [
|
||||
getQueryKey: ({ context, search }: SelectorQueryArgs) => [
|
||||
'selectors',
|
||||
'cloudwatch.logStreams',
|
||||
context.awsAccessKeyId ?? 'none',
|
||||
context.awsRegion ?? 'none',
|
||||
context.logGroupName ?? 'none',
|
||||
search ?? '',
|
||||
],
|
||||
enabled: ({ context }) =>
|
||||
Boolean(
|
||||
|
||||
@@ -70,7 +70,12 @@ export function useSelectorOptions(
|
||||
context: args.context,
|
||||
search: args.search,
|
||||
}
|
||||
const isEnabled = args.enabled ?? (definition.enabled ? definition.enabled(queryArgs) : true)
|
||||
/**
|
||||
* `definition.enabled` mirrors the preconditions the definition's own fetchers assert, so
|
||||
* it is a hard precondition for the list, not a default a caller may replace. A caller's
|
||||
* `enabled` only narrows — widening would run a fetch guaranteed to reject and cache it.
|
||||
*/
|
||||
const isEnabled = args.enabled !== false && (definition.enabled?.(queryArgs) ?? true)
|
||||
const supportsPagination = Boolean(definition.fetchPage)
|
||||
|
||||
const flatQuery = useQuery<SelectorOption[]>({
|
||||
|
||||
@@ -62,6 +62,9 @@ const trelloCallbackQuerySchema = z
|
||||
})
|
||||
.passthrough()
|
||||
|
||||
/** Google domain-wide-delegation subject. Also applied by in-process credential callers. */
|
||||
export const impersonateEmailSchema = z.string().email()
|
||||
|
||||
export const oauthTokenRequestBodySchema = z
|
||||
.object({
|
||||
credentialId: z.string().min(1).optional(),
|
||||
@@ -69,7 +72,7 @@ export const oauthTokenRequestBodySchema = z
|
||||
providerId: z.string().min(1).optional(),
|
||||
workflowId: z.string().min(1).nullish(),
|
||||
scopes: z.array(z.string()).optional(),
|
||||
impersonateEmail: z.string().email().optional(),
|
||||
impersonateEmail: impersonateEmailSchema.optional(),
|
||||
})
|
||||
.refine(
|
||||
(data) => data.credentialId || (data.credentialAccountUserId && data.providerId),
|
||||
@@ -99,6 +102,9 @@ const oauthTokenResponseSchema = z.object({
|
||||
authStyle: z.enum(['x-api-token']).optional(),
|
||||
})
|
||||
|
||||
/** Token material a resolved credential yields, on the wire and in-process alike. */
|
||||
export type OAuthTokenResponse = z.output<typeof oauthTokenResponseSchema>
|
||||
|
||||
export const oauthTokenGetContract = defineRouteContract({
|
||||
method: 'GET',
|
||||
path: '/api/auth/oauth/token',
|
||||
|
||||
@@ -23,7 +23,7 @@ vi.mock('@/lib/workspaces/permissions/utils', () => ({
|
||||
resolveWorkspaceAccess: mockResolveWorkspaceAccess,
|
||||
}))
|
||||
|
||||
import { authorizeCredentialUse } from '@/lib/auth/credential-access'
|
||||
import { authorizeCredentialUse, authorizeCredentialUseForAuth } from '@/lib/auth/credential-access'
|
||||
|
||||
afterAll(resetDbChainMock)
|
||||
|
||||
@@ -235,4 +235,41 @@ describe('authorizeCredentialUse', () => {
|
||||
expect(result.error).toBe('Credential not found')
|
||||
})
|
||||
})
|
||||
|
||||
/**
|
||||
* The in-process tool executor synthesizes the AuthResult an internal JWT
|
||||
* would have produced instead of minting one and POSTing to ourselves, so the
|
||||
* subject-less case must still fail closed here.
|
||||
*/
|
||||
describe('authorizeCredentialUseForAuth', () => {
|
||||
it('fails closed when the authenticated caller carries no user id', async () => {
|
||||
const result = await authorizeCredentialUseForAuth(
|
||||
{ success: true, authType: 'internal_jwt' },
|
||||
{ credentialId: ACCOUNT_ID }
|
||||
)
|
||||
|
||||
expect(result.ok).toBe(false)
|
||||
expect(result.error).toBe('Authentication required')
|
||||
})
|
||||
|
||||
it('fails closed when authentication did not succeed', async () => {
|
||||
const result = await authorizeCredentialUseForAuth(
|
||||
{ success: false, error: 'Unauthorized' },
|
||||
{ credentialId: ACCOUNT_ID }
|
||||
)
|
||||
|
||||
expect(result.ok).toBe(false)
|
||||
expect(result.error).toBe('Unauthorized')
|
||||
})
|
||||
|
||||
it('rejects an asserted caller that does not match the internal token subject', async () => {
|
||||
const result = await authorizeCredentialUseForAuth(
|
||||
{ success: true, userId: OWNER, authType: 'internal_jwt' },
|
||||
{ credentialId: ACCOUNT_ID, callerUserId: 'someone-else' }
|
||||
)
|
||||
|
||||
expect(result.ok).toBe(false)
|
||||
expect(result.error).toBe('Caller user does not match internal token subject')
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
@@ -2,7 +2,7 @@ import { db } from '@sim/db'
|
||||
import { account, credential, workflow as workflowTable } from '@sim/db/schema'
|
||||
import { and, asc, eq } from 'drizzle-orm'
|
||||
import type { NextRequest } from 'next/server'
|
||||
import { AuthType, checkSessionOrInternalAuth } from '@/lib/auth/hybrid'
|
||||
import { type AuthResult, AuthType, checkSessionOrInternalAuth } from '@/lib/auth/hybrid'
|
||||
import {
|
||||
type CredentialActorContext,
|
||||
canUseCredential,
|
||||
@@ -59,11 +59,27 @@ export async function authorizeCredentialUse(
|
||||
callerUserId?: string
|
||||
}
|
||||
): Promise<CredentialAccessResult> {
|
||||
const { credentialId, workflowId, requireWorkflowIdForInternal = true, callerUserId } = params
|
||||
|
||||
const auth = await checkSessionOrInternalAuth(request, {
|
||||
requireWorkflowId: requireWorkflowIdForInternal,
|
||||
requireWorkflowId: params.requireWorkflowIdForInternal ?? true,
|
||||
})
|
||||
return authorizeCredentialUseForAuth(auth, params)
|
||||
}
|
||||
|
||||
/**
|
||||
* Credential authorization for an already-authenticated caller.
|
||||
* {@link authorizeCredentialUse} is the HTTP wrapper; in-process callers build the same
|
||||
* {@link AuthResult} directly, so both paths run one identical rule.
|
||||
*/
|
||||
export async function authorizeCredentialUseForAuth(
|
||||
auth: AuthResult,
|
||||
params: {
|
||||
credentialId: string
|
||||
workflowId?: string
|
||||
callerUserId?: string
|
||||
}
|
||||
): Promise<CredentialAccessResult> {
|
||||
const { credentialId, workflowId, callerUserId } = params
|
||||
|
||||
if (!auth.success || !auth.userId) {
|
||||
return { ok: false, error: auth.error || 'Authentication required' }
|
||||
}
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import { cache } from 'react'
|
||||
import { db } from '@sim/db'
|
||||
import { member, organization, subscription, user } from '@sim/db/schema'
|
||||
import { createLogger } from '@sim/logger'
|
||||
@@ -429,11 +430,7 @@ export async function isEnterpriseOrgAdminOrOwner(userId: string): Promise<boole
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if an organization has an enterprise plan
|
||||
* Used for Access Control (Permission Groups) feature gating
|
||||
*/
|
||||
export async function isOrganizationOnEnterprisePlan(organizationId: string): Promise<boolean> {
|
||||
async function resolveOrganizationEnterprisePlan(organizationId: string): Promise<boolean> {
|
||||
try {
|
||||
if (!isBillingEnabled) {
|
||||
return true
|
||||
@@ -456,6 +453,15 @@ export async function isOrganizationOnEnterprisePlan(organizationId: string): Pr
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if an organization has an enterprise plan
|
||||
* Used for Access Control (Permission Groups) feature gating
|
||||
*
|
||||
* Request-memoized: a settings render gates several sections on the same
|
||||
* organization's plan, and it cannot change mid-render.
|
||||
*/
|
||||
export const isOrganizationOnEnterprisePlan = cache(resolveOrganizationEnterprisePlan)
|
||||
|
||||
/**
|
||||
* Entitlement for a single org-scoped enterprise feature.
|
||||
*
|
||||
@@ -612,10 +618,15 @@ async function hasWorkspaceTierAccess(
|
||||
* Whether the workspace's payer is on a usable Max-or-Enterprise subscription.
|
||||
* Shared by the inbox (Sim Mailer), live sync, and custom sandboxes, which all
|
||||
* sit on the same entitlement tier.
|
||||
*
|
||||
* Request-memoized: these features are gated side by side on one settings render,
|
||||
* each otherwise repeating the identical workspace and subscription reads. The
|
||||
* per-feature deployment and env short-circuits live in the wrappers and still run
|
||||
* per call.
|
||||
*/
|
||||
async function hasMaxTierWorkspaceAccess(workspaceId: string): Promise<boolean> {
|
||||
return hasWorkspaceTierAccess(workspaceId, isMaxTier)
|
||||
}
|
||||
const hasMaxTierWorkspaceAccess = cache(
|
||||
(workspaceId: string): Promise<boolean> => hasWorkspaceTierAccess(workspaceId, isMaxTier)
|
||||
)
|
||||
|
||||
/**
|
||||
* Check whether a workspace is entitled to the inbox (Sim Mailer) feature.
|
||||
|
||||
@@ -0,0 +1,214 @@
|
||||
/**
|
||||
* @vitest-environment node
|
||||
*/
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
|
||||
const {
|
||||
mockAuthorizeCredentialUseForAuth,
|
||||
mockGetCredential,
|
||||
mockRecordAudit,
|
||||
mockRefreshTokenIfNeeded,
|
||||
mockResolveOAuthAccountId,
|
||||
mockResolveServiceAccountToken,
|
||||
} = vi.hoisted(() => ({
|
||||
mockAuthorizeCredentialUseForAuth: vi.fn(),
|
||||
mockGetCredential: vi.fn(),
|
||||
mockRecordAudit: vi.fn(),
|
||||
mockRefreshTokenIfNeeded: vi.fn(),
|
||||
mockResolveOAuthAccountId: vi.fn(),
|
||||
mockResolveServiceAccountToken: vi.fn(),
|
||||
}))
|
||||
|
||||
vi.mock('@sim/audit', () => ({
|
||||
AuditAction: { CREDENTIAL_ACCESSED: 'credential.accessed' },
|
||||
AuditResourceType: { CREDENTIAL: 'credential' },
|
||||
recordAudit: mockRecordAudit,
|
||||
}))
|
||||
|
||||
vi.mock('@/lib/auth/credential-access', () => ({
|
||||
authorizeCredentialUseForAuth: mockAuthorizeCredentialUseForAuth,
|
||||
}))
|
||||
|
||||
vi.mock('@/lib/oauth/credential-service', () => ({
|
||||
getCredential: mockGetCredential,
|
||||
refreshTokenIfNeeded: mockRefreshTokenIfNeeded,
|
||||
resolveOAuthAccountId: mockResolveOAuthAccountId,
|
||||
resolveServiceAccountToken: mockResolveServiceAccountToken,
|
||||
}))
|
||||
|
||||
vi.mock('@/lib/posthog/server', () => ({
|
||||
captureServerEvent: vi.fn(),
|
||||
}))
|
||||
|
||||
import { TokenServiceAccountValidationError } from '@/lib/credentials/token-service-accounts/errors'
|
||||
import { resolveCredentialToken } from '@/lib/oauth/token-resolution'
|
||||
|
||||
const INTERNAL_AUTH = { success: true, userId: 'user-1', authType: 'internal_jwt' } as const
|
||||
|
||||
describe('resolveCredentialToken', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
mockResolveOAuthAccountId.mockResolvedValue(null)
|
||||
})
|
||||
|
||||
it('fails closed when the credential is not authorized', async () => {
|
||||
mockAuthorizeCredentialUseForAuth.mockResolvedValue({
|
||||
ok: false,
|
||||
error: 'You do not have access to this credential.',
|
||||
})
|
||||
|
||||
const result = await resolveCredentialToken(INTERNAL_AUTH, {
|
||||
requestId: 'req-1',
|
||||
credentialId: 'cred-1',
|
||||
})
|
||||
|
||||
expect(result).toEqual({
|
||||
ok: false,
|
||||
status: 403,
|
||||
error: 'You do not have access to this credential.',
|
||||
})
|
||||
expect(mockGetCredential).not.toHaveBeenCalled()
|
||||
expect(mockRefreshTokenIfNeeded).not.toHaveBeenCalled()
|
||||
expect(mockRecordAudit).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('fails closed when the caller carries no user id', async () => {
|
||||
mockAuthorizeCredentialUseForAuth.mockResolvedValue({
|
||||
ok: false,
|
||||
error: 'Authentication required',
|
||||
})
|
||||
|
||||
const result = await resolveCredentialToken(
|
||||
{ success: true, authType: 'internal_jwt' },
|
||||
{ requestId: 'req-1', credentialId: 'cred-1' }
|
||||
)
|
||||
|
||||
expect(result).toEqual({ ok: false, status: 403, error: 'Authentication required' })
|
||||
})
|
||||
|
||||
it('refreshes the token, records the access trail, and returns the payload', async () => {
|
||||
mockAuthorizeCredentialUseForAuth.mockResolvedValue({
|
||||
ok: true,
|
||||
requesterUserId: 'user-1',
|
||||
credentialOwnerUserId: 'owner-1',
|
||||
workspaceId: 'ws-1',
|
||||
resolvedCredentialId: 'account-1',
|
||||
})
|
||||
mockGetCredential.mockResolvedValue({
|
||||
providerId: 'google',
|
||||
idToken: 'id-token',
|
||||
scope: 'https://www.googleapis.com/auth/gmail.send',
|
||||
})
|
||||
mockRefreshTokenIfNeeded.mockResolvedValue({ accessToken: 'fresh', refreshed: true })
|
||||
|
||||
const result = await resolveCredentialToken(INTERNAL_AUTH, {
|
||||
requestId: 'req-1',
|
||||
credentialId: 'cred-1',
|
||||
workflowId: 'wf-1',
|
||||
})
|
||||
|
||||
expect(result).toEqual({ ok: true, token: { accessToken: 'fresh', idToken: 'id-token' } })
|
||||
expect(mockGetCredential).toHaveBeenCalledWith('req-1', 'account-1', 'owner-1')
|
||||
expect(mockRefreshTokenIfNeeded).toHaveBeenCalled()
|
||||
expect(mockRecordAudit).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
actorId: 'user-1',
|
||||
workspaceId: 'ws-1',
|
||||
resourceId: 'account-1',
|
||||
action: 'credential.accessed',
|
||||
})
|
||||
)
|
||||
})
|
||||
|
||||
it('returns 404 when the authorized credential is missing', async () => {
|
||||
mockAuthorizeCredentialUseForAuth.mockResolvedValue({
|
||||
ok: true,
|
||||
requesterUserId: 'user-1',
|
||||
credentialOwnerUserId: 'owner-1',
|
||||
})
|
||||
mockGetCredential.mockResolvedValue(undefined)
|
||||
|
||||
const result = await resolveCredentialToken(INTERNAL_AUTH, {
|
||||
requestId: 'req-1',
|
||||
credentialId: 'cred-1',
|
||||
})
|
||||
|
||||
expect(result).toEqual({ ok: false, status: 404, error: 'Credential not found' })
|
||||
})
|
||||
|
||||
it('reports a failed refresh as 401 without recording access', async () => {
|
||||
mockAuthorizeCredentialUseForAuth.mockResolvedValue({
|
||||
ok: true,
|
||||
requesterUserId: 'user-1',
|
||||
credentialOwnerUserId: 'owner-1',
|
||||
})
|
||||
mockGetCredential.mockResolvedValue({ providerId: 'google' })
|
||||
mockRefreshTokenIfNeeded.mockRejectedValue(new Error('refresh token revoked'))
|
||||
|
||||
const result = await resolveCredentialToken(INTERNAL_AUTH, {
|
||||
requestId: 'req-1',
|
||||
credentialId: 'cred-1',
|
||||
})
|
||||
|
||||
expect(result).toEqual({ ok: false, status: 401, error: 'Failed to refresh access token' })
|
||||
expect(mockRecordAudit).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('authorizes service-account credentials before minting a token', async () => {
|
||||
mockResolveOAuthAccountId.mockResolvedValue({
|
||||
credentialType: 'service_account',
|
||||
credentialId: 'sa-1',
|
||||
providerId: 'google',
|
||||
workspaceId: 'ws-1',
|
||||
accountId: '',
|
||||
usedCredentialTable: true,
|
||||
})
|
||||
mockAuthorizeCredentialUseForAuth.mockResolvedValue({ ok: false, error: 'Unauthorized' })
|
||||
|
||||
const result = await resolveCredentialToken(INTERNAL_AUTH, {
|
||||
requestId: 'req-1',
|
||||
credentialId: 'cred-1',
|
||||
})
|
||||
|
||||
expect(result).toEqual({ ok: false, status: 403, error: 'Unauthorized' })
|
||||
expect(mockResolveServiceAccountToken).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('surfaces the classified service-account failure code', async () => {
|
||||
mockResolveOAuthAccountId.mockResolvedValue({
|
||||
credentialType: 'service_account',
|
||||
credentialId: 'sa-1',
|
||||
providerId: 'atlassian',
|
||||
workspaceId: 'ws-1',
|
||||
accountId: '',
|
||||
usedCredentialTable: true,
|
||||
})
|
||||
mockAuthorizeCredentialUseForAuth.mockResolvedValue({ ok: true, requesterUserId: 'user-1' })
|
||||
mockResolveServiceAccountToken.mockRejectedValue(
|
||||
new TokenServiceAccountValidationError('invalid_credentials', 401)
|
||||
)
|
||||
|
||||
const result = await resolveCredentialToken(INTERNAL_AUTH, {
|
||||
requestId: 'req-1',
|
||||
credentialId: 'cred-1',
|
||||
})
|
||||
|
||||
expect(result).toEqual({
|
||||
ok: false,
|
||||
status: 401,
|
||||
code: 'invalid_credentials',
|
||||
error: 'Credential rejected by the provider — reconnect the credential',
|
||||
})
|
||||
})
|
||||
|
||||
it('rejects a malformed impersonation subject before touching the credential', async () => {
|
||||
const result = await resolveCredentialToken(INTERNAL_AUTH, {
|
||||
requestId: 'req-1',
|
||||
credentialId: 'cred-1',
|
||||
impersonateEmail: 'not-an-email',
|
||||
})
|
||||
|
||||
expect(result.ok).toBe(false)
|
||||
expect(mockAuthorizeCredentialUseForAuth).not.toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,297 @@
|
||||
import { AuditAction, AuditResourceType, recordAudit } from '@sim/audit'
|
||||
import { createLogger } from '@sim/logger'
|
||||
import {
|
||||
impersonateEmailSchema,
|
||||
type OAuthTokenResponse,
|
||||
} from '@/lib/api/contracts/oauth-connections'
|
||||
import { authorizeCredentialUseForAuth } from '@/lib/auth/credential-access'
|
||||
import type { AuthResult } from '@/lib/auth/hybrid'
|
||||
import { TokenServiceAccountValidationError } from '@/lib/credentials/token-service-accounts/errors'
|
||||
import {
|
||||
getCredential,
|
||||
refreshTokenIfNeeded,
|
||||
resolveOAuthAccountId,
|
||||
resolveServiceAccountToken,
|
||||
} from '@/lib/oauth/credential-service'
|
||||
import { extractSalesforceInstanceUrl, isSalesforceOAuthProviderId } from '@/lib/oauth/salesforce'
|
||||
import { captureServerEvent } from '@/lib/posthog/server'
|
||||
import { extractZohoDeskBaseFromScope } from '@/tools/zoho_desk/host-allowlist'
|
||||
|
||||
const logger = createLogger('OAuthTokenResolution')
|
||||
|
||||
/**
|
||||
* Duck type of the inbound request, used only so audit rows record IP and user agent.
|
||||
* In-process callers omit it.
|
||||
*/
|
||||
export interface CredentialAuditRequest {
|
||||
headers: { get(name: string): string | null }
|
||||
}
|
||||
|
||||
/** Token material a resolved credential yields; taken from the contract so it cannot drift. */
|
||||
export type CredentialTokenPayload = OAuthTokenResponse
|
||||
|
||||
export interface ResolveCredentialTokenInput {
|
||||
/** Correlation id used by the credential service's own logging. */
|
||||
requestId: string
|
||||
credentialId?: string
|
||||
workflowId?: string
|
||||
/** Canonical provider scopes, used only by service-account token minting. */
|
||||
scopes?: string[]
|
||||
/** Google domain-wide-delegation subject for service-account credentials. */
|
||||
impersonateEmail?: string
|
||||
/**
|
||||
* Asserted acting user. When the caller authenticated with an internal JWT it
|
||||
* must equal the token subject, so a forged assertion cannot widen access.
|
||||
*/
|
||||
callerUserId?: string
|
||||
auditRequest?: CredentialAuditRequest
|
||||
}
|
||||
|
||||
export type ResolveCredentialTokenResult =
|
||||
| { ok: true; token: CredentialTokenPayload }
|
||||
| { ok: false; status: number; error: string; code?: string }
|
||||
|
||||
/**
|
||||
* Emits the semantic "credential used" trail for one resolved credential.
|
||||
* Both the audit row and the analytics event are fire-and-forget.
|
||||
*/
|
||||
function recordCredentialAccess(params: {
|
||||
actorId: string
|
||||
workspaceId: string | null
|
||||
resourceId: string
|
||||
providerId: string | null | undefined
|
||||
credentialType: 'oauth' | 'service_account'
|
||||
auditRequest?: CredentialAuditRequest
|
||||
}): void {
|
||||
const { actorId, workspaceId, resourceId, providerId, credentialType } = params
|
||||
recordAudit({
|
||||
workspaceId,
|
||||
actorId,
|
||||
action: AuditAction.CREDENTIAL_ACCESSED,
|
||||
resourceType: AuditResourceType.CREDENTIAL,
|
||||
resourceId,
|
||||
description: `Accessed ${credentialType === 'oauth' ? 'OAuth' : 'service account'} credential for provider ${providerId ?? 'unknown'}`,
|
||||
metadata: {
|
||||
provider: providerId,
|
||||
credentialType,
|
||||
},
|
||||
request: params.auditRequest,
|
||||
})
|
||||
captureServerEvent(
|
||||
actorId,
|
||||
'credential_used',
|
||||
{
|
||||
credential_type: credentialType,
|
||||
provider_id: providerId ?? 'unknown',
|
||||
...(workspaceId ? { workspace_id: workspaceId } : {}),
|
||||
},
|
||||
workspaceId ? { groups: { workspace: workspaceId } } : undefined
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Projects a stored OAuth credential plus its access token into the wire payload.
|
||||
* Provider hosts come out of the scope string through shared allowlisted helpers, never a
|
||||
* local regex — these values are injected into tool calls that carry the token.
|
||||
*/
|
||||
function buildOAuthTokenPayload(
|
||||
credential: { providerId: string; scope?: string | null; idToken?: string | null },
|
||||
accessToken: string
|
||||
): CredentialTokenPayload {
|
||||
const instanceUrl = isSalesforceOAuthProviderId(credential.providerId)
|
||||
? extractSalesforceInstanceUrl(credential.scope ?? undefined)
|
||||
: undefined
|
||||
|
||||
let apiDomain: string | undefined
|
||||
if (credential.providerId === 'zoho-desk' && credential.scope) {
|
||||
apiDomain = extractZohoDeskBaseFromScope(credential.scope)
|
||||
}
|
||||
|
||||
return {
|
||||
accessToken,
|
||||
idToken: credential.idToken || undefined,
|
||||
...(instanceUrl && { instanceUrl }),
|
||||
...(apiDomain && { apiDomain }),
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Refreshes an authorized OAuth credential, records its access trail, and
|
||||
* projects the token payload. Shared by every surface that has already
|
||||
* authorized the credential and loaded it.
|
||||
*/
|
||||
export async function completeOAuthCredentialToken(params: {
|
||||
requestId: string
|
||||
credential: { providerId: string; scope?: string | null; idToken?: string | null }
|
||||
resolvedCredentialId: string
|
||||
actorId?: string
|
||||
workspaceId: string | null
|
||||
auditRequest?: CredentialAuditRequest
|
||||
}): Promise<ResolveCredentialTokenResult> {
|
||||
const { requestId, credential, resolvedCredentialId, actorId, workspaceId, auditRequest } = params
|
||||
try {
|
||||
const { accessToken } = await refreshTokenIfNeeded(requestId, credential, resolvedCredentialId)
|
||||
|
||||
if (actorId) {
|
||||
recordCredentialAccess({
|
||||
actorId,
|
||||
workspaceId,
|
||||
resourceId: resolvedCredentialId,
|
||||
providerId: credential.providerId,
|
||||
credentialType: 'oauth',
|
||||
auditRequest,
|
||||
})
|
||||
}
|
||||
|
||||
return { ok: true, token: buildOAuthTokenPayload(credential, accessToken) }
|
||||
} catch (error) {
|
||||
logger.error(`[${requestId}] Failed to refresh access token:`, error)
|
||||
return { ok: false, status: 401, error: 'Failed to refresh access token' }
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Authorized application operation behind `POST /api/auth/oauth/token`. Every surface that
|
||||
* needs a credential token — the route and the in-process tool executor — goes through
|
||||
* here, so authorization, refresh, and audit cannot drift between them.
|
||||
*
|
||||
* @param auth Result of authenticating the caller (session or internal JWT).
|
||||
*/
|
||||
export async function resolveCredentialToken(
|
||||
auth: AuthResult,
|
||||
input: ResolveCredentialTokenInput
|
||||
): Promise<ResolveCredentialTokenResult> {
|
||||
const {
|
||||
requestId,
|
||||
credentialId,
|
||||
workflowId,
|
||||
scopes,
|
||||
impersonateEmail,
|
||||
callerUserId,
|
||||
auditRequest,
|
||||
} = input
|
||||
|
||||
try {
|
||||
if (!credentialId) {
|
||||
return { ok: false, status: 400, error: 'Credential ID is required' }
|
||||
}
|
||||
if (
|
||||
impersonateEmail !== undefined &&
|
||||
!impersonateEmailSchema.safeParse(impersonateEmail).success
|
||||
) {
|
||||
return { ok: false, status: 400, error: 'impersonateEmail must be a valid email address' }
|
||||
}
|
||||
|
||||
/**
|
||||
* Both branches below authorize with the same arguments, and neither read depends
|
||||
* on the other, so they resolve together — this runs per credentialed tool call.
|
||||
*/
|
||||
const [resolved, authz] = await Promise.all([
|
||||
resolveOAuthAccountId(credentialId),
|
||||
authorizeCredentialUseForAuth(auth, { credentialId, workflowId, callerUserId }),
|
||||
])
|
||||
|
||||
if (resolved?.credentialType === 'service_account' && resolved.credentialId) {
|
||||
if (!authz.ok) {
|
||||
return { ok: false, status: 403, error: authz.error || 'Unauthorized' }
|
||||
}
|
||||
|
||||
const saActorId = authz.requesterUserId
|
||||
const saWorkspaceId = resolved.workspaceId ?? authz.workspaceId ?? null
|
||||
|
||||
try {
|
||||
const result = await resolveServiceAccountToken(
|
||||
resolved.credentialId,
|
||||
resolved.providerId,
|
||||
scopes ?? [],
|
||||
impersonateEmail
|
||||
)
|
||||
|
||||
if (saActorId) {
|
||||
recordCredentialAccess({
|
||||
actorId: saActorId,
|
||||
workspaceId: saWorkspaceId,
|
||||
resourceId: resolved.credentialId,
|
||||
providerId: resolved.providerId,
|
||||
credentialType: 'service_account',
|
||||
auditRequest,
|
||||
})
|
||||
}
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
token: {
|
||||
accessToken: result.accessToken,
|
||||
cloudId: result.cloudId,
|
||||
domain: result.domain,
|
||||
instanceUrl: result.instanceUrl,
|
||||
apiDomain: result.apiDomain,
|
||||
authStyle: result.authStyle,
|
||||
},
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error(`[${requestId}] Service account token error:`, error)
|
||||
if (error instanceof TokenServiceAccountValidationError) {
|
||||
// Classified provider outages are infra failures, not bad credentials.
|
||||
if (error.code === 'provider_unavailable') {
|
||||
return {
|
||||
ok: false,
|
||||
status: 502,
|
||||
error: 'Credential provider is temporarily unavailable',
|
||||
}
|
||||
}
|
||||
// A stored host that no longer resolves is a configuration failure —
|
||||
// surface the code so runtime consumers can say "check the host"
|
||||
// instead of a generic auth error.
|
||||
if (error.code === 'site_not_found') {
|
||||
return {
|
||||
ok: false,
|
||||
status: 400,
|
||||
code: error.code,
|
||||
error: 'Credential host not found — reconnect the credential with a valid host',
|
||||
}
|
||||
}
|
||||
// A revoked/rotated-away or misconfigured stored secret — surface the
|
||||
// code so runtime consumers can prompt to reconnect the credential
|
||||
// rather than showing a generic auth failure.
|
||||
if (error.code === 'invalid_credentials') {
|
||||
return {
|
||||
ok: false,
|
||||
status: 401,
|
||||
code: error.code,
|
||||
error: 'Credential rejected by the provider — reconnect the credential',
|
||||
}
|
||||
}
|
||||
}
|
||||
return { ok: false, status: 401, error: 'Failed to get service account token' }
|
||||
}
|
||||
}
|
||||
|
||||
if (!authz.ok || !authz.credentialOwnerUserId) {
|
||||
return { ok: false, status: 403, error: authz.error || 'Unauthorized' }
|
||||
}
|
||||
|
||||
const resolvedCredentialId = authz.resolvedCredentialId || credentialId
|
||||
const credential = await getCredential(
|
||||
requestId,
|
||||
resolvedCredentialId,
|
||||
authz.credentialOwnerUserId
|
||||
)
|
||||
|
||||
if (!credential) {
|
||||
return { ok: false, status: 404, error: 'Credential not found' }
|
||||
}
|
||||
|
||||
return completeOAuthCredentialToken({
|
||||
requestId,
|
||||
credential,
|
||||
resolvedCredentialId,
|
||||
actorId: authz.requesterUserId,
|
||||
workspaceId: authz.workspaceId ?? null,
|
||||
auditRequest,
|
||||
})
|
||||
} catch (error) {
|
||||
logger.error(`[${requestId}] Error getting access token`, error)
|
||||
return { ok: false, status: 500, error: 'Internal server error' }
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
import { cache } from 'react'
|
||||
import { db, dbReplica } from '@sim/db'
|
||||
import { settings, user } from '@sim/db/schema'
|
||||
import { eq } from 'drizzle-orm'
|
||||
@@ -41,8 +42,11 @@ export async function verifyEffectiveSuperUser(userId: string): Promise<{
|
||||
* served from the replica: this gates features, not security-critical auth, so it
|
||||
* tolerates the replica's bounded staleness (admin role rarely changes). Falls back
|
||||
* to the primary when no replica is configured.
|
||||
*
|
||||
* Request-memoized: an account-settings render checks the same viewer in both the
|
||||
* layout and the page.
|
||||
*/
|
||||
export async function isPlatformAdmin(userId: string): Promise<boolean> {
|
||||
export const isPlatformAdmin = cache(async (userId: string): Promise<boolean> => {
|
||||
const [row] = await dbReplica
|
||||
.select({ role: user.role })
|
||||
.from(user)
|
||||
@@ -50,4 +54,4 @@ export async function isPlatformAdmin(userId: string): Promise<boolean> {
|
||||
.limit(1)
|
||||
|
||||
return row?.role === 'admin'
|
||||
}
|
||||
})
|
||||
|
||||
@@ -161,6 +161,12 @@ interface ListWorkspaceFilesOptions {
|
||||
hydrateFolderPaths?: boolean
|
||||
/** Propagate storage errors when an incomplete list would be unsafe. */
|
||||
throwOnError?: boolean
|
||||
/**
|
||||
* Row cap for callers that only need to know whether the workspace fits a budget.
|
||||
* The result is a prefix of the full list, so a caller that reads it as "the
|
||||
* workspace's files" must not set this.
|
||||
*/
|
||||
limit?: number
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -223,7 +229,9 @@ interface WorkspaceFileMetadataInsert {
|
||||
size: number
|
||||
}
|
||||
|
||||
function workspaceFileSize(file: typeof workspaceFiles.$inferSelect): number {
|
||||
function workspaceFileSize(
|
||||
file: Pick<typeof workspaceFiles.$inferSelect, 'size' | 'sizeBytes'>
|
||||
): number {
|
||||
return file.sizeBytes ?? file.size
|
||||
}
|
||||
|
||||
@@ -1007,7 +1015,7 @@ export async function fileExistsInWorkspace(
|
||||
}
|
||||
|
||||
function mapWorkspaceFileRecord(
|
||||
file: typeof workspaceFiles.$inferSelect,
|
||||
file: WorkspaceFileListRow,
|
||||
workspaceId: string,
|
||||
folderPaths: Map<string, string>
|
||||
): WorkspaceFileRecord {
|
||||
@@ -1144,9 +1152,37 @@ function workspaceFileScopeCondition(workspaceId: string, scope: WorkspaceFileSc
|
||||
: and(...base, isNull(workspaceFiles.deletedAt))
|
||||
}
|
||||
|
||||
/**
|
||||
* The columns {@link mapWorkspaceFileRecord} reads. These list reads are workspace-wide,
|
||||
* so `select()` would ship five unprojected columns for every row of the scan.
|
||||
*/
|
||||
const workspaceFileListColumns = {
|
||||
id: workspaceFiles.id,
|
||||
key: workspaceFiles.key,
|
||||
userId: workspaceFiles.userId,
|
||||
workspaceId: workspaceFiles.workspaceId,
|
||||
folderId: workspaceFiles.folderId,
|
||||
originalName: workspaceFiles.originalName,
|
||||
contentType: workspaceFiles.contentType,
|
||||
size: workspaceFiles.size,
|
||||
sizeBytes: workspaceFiles.sizeBytes,
|
||||
width: workspaceFiles.width,
|
||||
height: workspaceFiles.height,
|
||||
deletedAt: workspaceFiles.deletedAt,
|
||||
uploadedAt: workspaceFiles.uploadedAt,
|
||||
updatedAt: workspaceFiles.updatedAt,
|
||||
contentUpdatedAt: workspaceFiles.contentUpdatedAt,
|
||||
} as const
|
||||
|
||||
/** A row carrying exactly the columns {@link mapWorkspaceFileRecord} needs; a full row satisfies it. */
|
||||
type WorkspaceFileListRow = Pick<
|
||||
typeof workspaceFiles.$inferSelect,
|
||||
keyof typeof workspaceFileListColumns
|
||||
>
|
||||
|
||||
/** Resolves `folderPath` for a page of rows, reading the folder tree only if any row needs it. */
|
||||
async function hydrateWorkspaceFilePaths(
|
||||
files: (typeof workspaceFiles.$inferSelect)[],
|
||||
files: WorkspaceFileListRow[],
|
||||
workspaceId: string,
|
||||
options?: { folders?: WorkspaceFileFolderRecord[]; hydrateFolderPaths?: boolean }
|
||||
): Promise<WorkspaceFileRecord[]> {
|
||||
@@ -1167,12 +1203,13 @@ export async function listWorkspaceFiles(
|
||||
options?: ListWorkspaceFilesOptions
|
||||
): Promise<WorkspaceFileRecord[]> {
|
||||
try {
|
||||
const { scope = 'active' } = options ?? {}
|
||||
const files = await db
|
||||
.select()
|
||||
const { scope = 'active', limit } = options ?? {}
|
||||
const query = db
|
||||
.select(workspaceFileListColumns)
|
||||
.from(workspaceFiles)
|
||||
.where(workspaceFileScopeCondition(workspaceId, scope))
|
||||
.orderBy(workspaceFiles.uploadedAt)
|
||||
const files = await (limit === undefined ? query : query.limit(limit))
|
||||
|
||||
return hydrateWorkspaceFilePaths(files, workspaceId, options)
|
||||
} catch (error) {
|
||||
@@ -1265,7 +1302,7 @@ export async function queryWorkspaceFiles(
|
||||
]
|
||||
|
||||
const rows = await db
|
||||
.select()
|
||||
.select(workspaceFileListColumns)
|
||||
.from(workspaceFiles)
|
||||
.where(and(...conditions))
|
||||
.orderBy(...listOrderBy(keysetColumns(keys), sortOrder))
|
||||
|
||||
@@ -46,7 +46,10 @@ vi.mock('@/lib/uploads/contexts/workspace/workspace-file-folder-manager', () =>
|
||||
resolveWorkspaceFileFolderTarget: vi.fn(async () => null),
|
||||
}))
|
||||
|
||||
import { queryWorkspaceFiles } from '@/lib/uploads/contexts/workspace/workspace-file-manager'
|
||||
import {
|
||||
listWorkspaceFiles,
|
||||
queryWorkspaceFiles,
|
||||
} from '@/lib/uploads/contexts/workspace/workspace-file-manager'
|
||||
|
||||
const WS = 'workspace-1'
|
||||
|
||||
@@ -205,3 +208,52 @@ describe('queryWorkspaceFiles', () => {
|
||||
).rejects.toMatchObject({ code: 'validation' })
|
||||
})
|
||||
})
|
||||
|
||||
/**
|
||||
* `listWorkspaceFiles` materializes a whole scope, so what it reads per row is
|
||||
* multiplied by the size of the workspace. These assertions pin the two ways that
|
||||
* stays bounded: the projection, and the optional row cap its one budgeted caller
|
||||
* (the workspace layout's server seed) uses to detect an oversized workspace.
|
||||
*/
|
||||
describe('listWorkspaceFiles', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
resetDbChainMock()
|
||||
})
|
||||
|
||||
const lastProjection = () =>
|
||||
Object.keys((dbChainMockFns.select.mock.calls.at(-1)?.[0] ?? {}) as Record<string, unknown>)
|
||||
|
||||
it('projects only the columns the record mapper reads', async () => {
|
||||
queueTableRows(schemaMock.workspaceFiles, [buildRow()])
|
||||
|
||||
await listWorkspaceFiles(WS)
|
||||
|
||||
const projected = lastProjection()
|
||||
expect(projected).toEqual(
|
||||
expect.arrayContaining(['id', 'key', 'originalName', 'sizeBytes', 'contentUpdatedAt'])
|
||||
)
|
||||
/** Columns no reader projects; `select()` would ship them for every row of the scan. */
|
||||
expect(projected).not.toContain('context')
|
||||
expect(projected).not.toContain('chatId')
|
||||
expect(projected).not.toContain('messageId')
|
||||
expect(projected).not.toContain('displayName')
|
||||
expect(projected).not.toContain('secretProvenanceVersion')
|
||||
})
|
||||
|
||||
it('reads the whole scope when no cap is given', async () => {
|
||||
queueTableRows(schemaMock.workspaceFiles, [buildRow()])
|
||||
|
||||
await listWorkspaceFiles(WS)
|
||||
|
||||
expect(dbChainMockFns.limit).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('caps the rows read when the caller only needs to fit a budget', async () => {
|
||||
queueTableRows(schemaMock.workspaceFiles, [buildRow()])
|
||||
|
||||
await listWorkspaceFiles(WS, { limit: 2 })
|
||||
|
||||
expect(dbChainMockFns.limit).toHaveBeenCalledWith(2)
|
||||
})
|
||||
})
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
/**
|
||||
* @vitest-environment node
|
||||
*/
|
||||
import { describe, expect, it } from 'vitest'
|
||||
import { parseWorkflowStateForPersistence } from '@/lib/workflows/persistence/save-normalized-state'
|
||||
|
||||
/**
|
||||
* A checkpoint blob as the revert route builds it: JSONB-derived blocks and edges, plus a
|
||||
* real `Date` for `deployedAt`.
|
||||
*/
|
||||
function checkpointState(overrides?: Record<string, unknown>) {
|
||||
return {
|
||||
blocks: {
|
||||
'block-1': {
|
||||
id: 'block-1',
|
||||
type: 'starter',
|
||||
name: 'Start',
|
||||
position: { x: 0, y: 0 },
|
||||
subBlocks: {},
|
||||
outputs: {},
|
||||
enabled: true,
|
||||
},
|
||||
},
|
||||
edges: [],
|
||||
loops: {},
|
||||
parallels: {},
|
||||
isDeployed: false,
|
||||
lastSaved: 1_754_000_000_000,
|
||||
...overrides,
|
||||
}
|
||||
}
|
||||
|
||||
describe('parseWorkflowStateForPersistence', () => {
|
||||
/**
|
||||
* The revert used to reach this schema by POSTing the blob over HTTP, so every value
|
||||
* arrived JSON-serialized. In-process the blob keeps its runtime types. Both forms must
|
||||
* parse identically, or a checkpoint that reverted before would start failing.
|
||||
*/
|
||||
it('accepts a Date for deployedAt exactly as it accepted the serialized string', () => {
|
||||
const deployedAt = new Date('2026-01-02T03:04:05.678Z')
|
||||
|
||||
const fromDate = parseWorkflowStateForPersistence(checkpointState({ deployedAt }))
|
||||
const overTheWire = parseWorkflowStateForPersistence(
|
||||
JSON.parse(JSON.stringify(checkpointState({ deployedAt })))
|
||||
)
|
||||
|
||||
expect(fromDate.success).toBe(true)
|
||||
expect(overTheWire.success).toBe(true)
|
||||
expect(fromDate.data?.deployedAt).toEqual(deployedAt)
|
||||
expect(overTheWire.data?.deployedAt).toEqual(fromDate.data?.deployedAt)
|
||||
})
|
||||
|
||||
it('round-trips a JSONB-shaped blob without dropping blocks or edges', () => {
|
||||
const state = checkpointState()
|
||||
|
||||
const parsed = parseWorkflowStateForPersistence(state)
|
||||
|
||||
expect(parsed.success).toBe(true)
|
||||
expect(Object.keys(parsed.data?.blocks ?? {})).toEqual(['block-1'])
|
||||
expect(parsed.data?.lastSaved).toBe(1_754_000_000_000)
|
||||
})
|
||||
|
||||
it('accepts a null deployedAt, which the revert passes for a never-deployed checkpoint', () => {
|
||||
const parsed = parseWorkflowStateForPersistence(checkpointState({ deployedAt: null }))
|
||||
|
||||
expect(parsed.success).toBe(true)
|
||||
expect(parsed.data?.deployedAt).toBeNull()
|
||||
})
|
||||
|
||||
/** The validation the removed HTTP hop used to provide: a malformed blob must not be written. */
|
||||
it('rejects a blob whose blocks are malformed', () => {
|
||||
const parsed = parseWorkflowStateForPersistence({
|
||||
blocks: { 'block-1': { id: 'block-1' } },
|
||||
edges: [],
|
||||
})
|
||||
|
||||
expect(parsed.success).toBe(false)
|
||||
})
|
||||
|
||||
it('rejects a blob missing blocks entirely', () => {
|
||||
expect(parseWorkflowStateForPersistence({ edges: [] }).success).toBe(false)
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,176 @@
|
||||
import { db } from '@sim/db'
|
||||
import { workflow } from '@sim/db/schema'
|
||||
import { createLogger } from '@sim/logger'
|
||||
import {
|
||||
assertWorkflowMutable,
|
||||
authorizeWorkflowByWorkspacePermission,
|
||||
WorkflowLockedError,
|
||||
type WorkflowWorkspaceAuthorizationResult,
|
||||
} from '@sim/platform-authz/workflow'
|
||||
import { eq } from 'drizzle-orm'
|
||||
import type { z } from 'zod'
|
||||
import {
|
||||
type WorkflowStateContractOutput,
|
||||
workflowStateSchema,
|
||||
} from '@/lib/api/contracts/workflows'
|
||||
import { notifyWorkflowUpdated } from '@/lib/realtime/notify'
|
||||
import { extractAndPersistCustomTools } from '@/lib/workflows/persistence/custom-tools-persistence'
|
||||
import { prepareWorkflowStateForPersistence } from '@/lib/workflows/persistence/prepare-state'
|
||||
import { saveWorkflowToNormalizedTables } from '@/lib/workflows/persistence/utils'
|
||||
import type { BlockState, WorkflowState } from '@/stores/workflows/workflow/types'
|
||||
|
||||
const logger = createLogger('WorkflowStatePersistence')
|
||||
|
||||
export type SaveWorkflowNormalizedStateResult =
|
||||
| { success: true; warnings: string[] }
|
||||
| { success: false; status: number; error: string; details?: string }
|
||||
|
||||
/**
|
||||
* Validates an untrusted state blob against the same schema `PUT /api/workflows/[id]/state`
|
||||
* applies, so in-process callers get the coercion and rejection the HTTP hop gave them.
|
||||
*/
|
||||
export function parseWorkflowStateForPersistence(
|
||||
value: unknown
|
||||
): z.ZodSafeParseResult<WorkflowStateContractOutput> {
|
||||
return workflowStateSchema.safeParse(value)
|
||||
}
|
||||
|
||||
/**
|
||||
* Writes a complete workflow state to the normalized tables: write authorization,
|
||||
* the lock check, block/edge preparation, the row-locked save transaction,
|
||||
* `lastSynced`/variables, custom-tool extraction, and the socket notification.
|
||||
* Every surface that replaces a workflow's state calls this, so no step can be
|
||||
* skipped by going through a different door.
|
||||
*
|
||||
* Every refusal, the lock included, comes back as a failure result — callers need
|
||||
* one branch.
|
||||
*
|
||||
* `authorization` lets a caller that already resolved the same decision hand it in
|
||||
* rather than pay for it twice; it must be the `write` decision for this workflow
|
||||
* and user.
|
||||
*/
|
||||
export async function saveWorkflowNormalizedState(params: {
|
||||
requestId: string
|
||||
workflowId: string
|
||||
userId: string
|
||||
state: WorkflowStateContractOutput
|
||||
authorization?: WorkflowWorkspaceAuthorizationResult
|
||||
}): Promise<SaveWorkflowNormalizedStateResult> {
|
||||
const { requestId, workflowId, userId, state } = params
|
||||
|
||||
const authorization =
|
||||
params.authorization ??
|
||||
(await authorizeWorkflowByWorkspacePermission({ workflowId, userId, action: 'write' }))
|
||||
const workflowData = authorization.workflow
|
||||
|
||||
if (!workflowData) {
|
||||
logger.warn(`[${requestId}] Workflow ${workflowId} not found for state update`)
|
||||
return { success: false, status: 404, error: 'Workflow not found' }
|
||||
}
|
||||
|
||||
if (!authorization.allowed) {
|
||||
logger.warn(
|
||||
`[${requestId}] User ${userId} denied permission to update workflow state ${workflowId}`
|
||||
)
|
||||
return {
|
||||
success: false,
|
||||
status: authorization.status || 403,
|
||||
error: authorization.message || 'Access denied',
|
||||
}
|
||||
}
|
||||
|
||||
try {
|
||||
await assertWorkflowMutable(workflowId)
|
||||
} catch (error) {
|
||||
if (error instanceof WorkflowLockedError) {
|
||||
return { success: false, status: error.status, error: error.message }
|
||||
}
|
||||
throw error
|
||||
}
|
||||
|
||||
const { state: preparedState, warnings: preparationWarnings } =
|
||||
prepareWorkflowStateForPersistence({
|
||||
blocks: state.blocks as Record<string, BlockState>,
|
||||
edges: state.edges as WorkflowState['edges'],
|
||||
})
|
||||
|
||||
const workflowState = {
|
||||
...preparedState,
|
||||
lastSaved: state.lastSaved || Date.now(),
|
||||
isDeployed: state.isDeployed || false,
|
||||
deployedAt: state.deployedAt,
|
||||
}
|
||||
|
||||
const saveResult = await db.transaction(async (tx) => {
|
||||
await tx
|
||||
.select({ id: workflow.id })
|
||||
.from(workflow)
|
||||
.where(eq(workflow.id, workflowId))
|
||||
.limit(1)
|
||||
.for('update')
|
||||
|
||||
const result = await saveWorkflowToNormalizedTables(
|
||||
workflowId,
|
||||
workflowState as WorkflowState,
|
||||
tx
|
||||
)
|
||||
|
||||
if (!result.success) return result
|
||||
|
||||
const updateData: {
|
||||
lastSynced: Date
|
||||
updatedAt: Date
|
||||
variables?: typeof state.variables
|
||||
} = {
|
||||
lastSynced: new Date(),
|
||||
updatedAt: new Date(),
|
||||
}
|
||||
|
||||
if (state.variables !== undefined) {
|
||||
updateData.variables = state.variables
|
||||
}
|
||||
|
||||
await tx.update(workflow).set(updateData).where(eq(workflow.id, workflowId))
|
||||
|
||||
return result
|
||||
})
|
||||
|
||||
if (!saveResult.success) {
|
||||
logger.error(`[${requestId}] Failed to save workflow ${workflowId} state:`, saveResult.error)
|
||||
return {
|
||||
success: false,
|
||||
status: 500,
|
||||
error: 'Failed to save workflow state',
|
||||
details: saveResult.error,
|
||||
}
|
||||
}
|
||||
|
||||
try {
|
||||
const workspaceId = workflowData.workspaceId
|
||||
if (workspaceId) {
|
||||
const { saved, errors } = await extractAndPersistCustomTools(
|
||||
workflowState,
|
||||
workspaceId,
|
||||
userId
|
||||
)
|
||||
|
||||
if (saved > 0) {
|
||||
logger.info(`[${requestId}] Persisted ${saved} custom tool(s) to database`, { workflowId })
|
||||
}
|
||||
|
||||
if (errors.length > 0) {
|
||||
logger.warn(`[${requestId}] Some custom tools failed to persist`, { errors, workflowId })
|
||||
}
|
||||
} else {
|
||||
logger.warn(`[${requestId}] Workflow has no workspaceId, skipping custom tools persistence`, {
|
||||
workflowId,
|
||||
})
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error(`[${requestId}] Failed to persist custom tools`, { error, workflowId })
|
||||
}
|
||||
|
||||
await notifyWorkflowUpdated(workflowId)
|
||||
|
||||
return { success: true, warnings: preparationWarnings }
|
||||
}
|
||||
@@ -53,6 +53,62 @@ describe('listWorkspaceFilesWithShares', () => {
|
||||
expect(file.uploadedAt).toEqual(new Date('2026-01-01T00:00:00.000Z'))
|
||||
})
|
||||
|
||||
/**
|
||||
* `maxRows` exists for a caller that will only use the list if the whole workspace fits
|
||||
* its payload budget, so overflow must be reported as `null` — a prefix returned here
|
||||
* would be presented as the workspace's complete file list. The share read still runs
|
||||
* concurrently and is discarded: the under-budget workspaces are the common case, and
|
||||
* serializing the two reads to save this one would tax every normal request.
|
||||
*/
|
||||
it('returns null when the workspace exceeds maxRows', async () => {
|
||||
mockListWorkspaceFiles.mockResolvedValue([STORED_FILE, STORED_FILE, STORED_FILE])
|
||||
|
||||
const result = await listWorkspaceFilesWithShares('ws-1', 'active', { maxRows: 2 })
|
||||
|
||||
expect(result).toBeNull()
|
||||
expect(mockListWorkspaceFiles).toHaveBeenCalledWith('ws-1', { scope: 'active', limit: 3 })
|
||||
})
|
||||
|
||||
it('returns the list when it fits maxRows', async () => {
|
||||
mockListWorkspaceFiles.mockResolvedValue([STORED_FILE])
|
||||
|
||||
const result = await listWorkspaceFilesWithShares('ws-1', 'active', { maxRows: 2 })
|
||||
|
||||
expect(result).toHaveLength(1)
|
||||
expect(mockGetWorkspaceShares).toHaveBeenCalledWith('file', 'ws-1')
|
||||
})
|
||||
|
||||
/** The boundary the `>` comparison turns on: exactly maxRows must still be the list. */
|
||||
it('returns the list when it sits exactly on maxRows', async () => {
|
||||
mockListWorkspaceFiles.mockResolvedValue([STORED_FILE, STORED_FILE])
|
||||
|
||||
const result = await listWorkspaceFilesWithShares('ws-1', 'active', { maxRows: 2 })
|
||||
|
||||
expect(result).toHaveLength(2)
|
||||
})
|
||||
|
||||
/**
|
||||
* The file read swallows errors and returns `[]` by default. A caller seeding a cache
|
||||
* must not receive that: an empty list would be cached as "this workspace has no files".
|
||||
*/
|
||||
it('propagates a failed read instead of degrading to an empty list', async () => {
|
||||
await listWorkspaceFilesWithShares('ws-1', 'active', { throwOnError: true })
|
||||
|
||||
expect(mockListWorkspaceFiles).toHaveBeenCalledWith(
|
||||
'ws-1',
|
||||
expect.objectContaining({ throwOnError: true })
|
||||
)
|
||||
})
|
||||
|
||||
it('does not ask the file read to throw unless the caller opts in', async () => {
|
||||
await listWorkspaceFilesWithShares('ws-1', 'active')
|
||||
|
||||
expect(mockListWorkspaceFiles).toHaveBeenCalledWith(
|
||||
'ws-1',
|
||||
expect.not.objectContaining({ throwOnError: true })
|
||||
)
|
||||
})
|
||||
|
||||
it('joins each file public share onto its row', async () => {
|
||||
const share = {
|
||||
id: 'share-1',
|
||||
|
||||
@@ -6,21 +6,40 @@ import {
|
||||
} from '@/lib/uploads/contexts/workspace/workspace-file-manager'
|
||||
|
||||
/**
|
||||
* Lists a workspace's files with each file's public share joined on — shared by
|
||||
* `GET /api/workspaces/[id]/files` and the Files browser's server prefetch so both cache
|
||||
* one shape.
|
||||
* Lists a workspace's files with each file's public share joined on, parsed through the
|
||||
* `GET /api/workspaces/[id]/files` response contract so the workspace layout's server seed
|
||||
* caches exactly the shape that route returns.
|
||||
*
|
||||
* Parsing through the route contract's response schema strips the server-only fields
|
||||
* `requestJson` strips on the client (`contentUpdatedAt`), so a prefetched entry is identical
|
||||
* to a client fetch rather than carrying a field that vanishes on the next refetch.
|
||||
*
|
||||
* Callers authorize the viewer against `workspaceId` first.
|
||||
*
|
||||
* `maxRows` bounds the result for a caller that only uses the list if the whole workspace
|
||||
* fits a payload budget: on overflow it returns `null` rather than the prefix, so a
|
||||
* truncated read is never presented as the workspace's files. The two reads still run
|
||||
* concurrently, since under-budget workspaces are the common case.
|
||||
*
|
||||
* `throwOnError` propagates a failed file read instead of degrading to an empty list,
|
||||
* which a cache seed would store as authoritative.
|
||||
*/
|
||||
export async function listWorkspaceFilesWithShares(workspaceId: string, scope: WorkspaceFileScope) {
|
||||
export async function listWorkspaceFilesWithShares(
|
||||
workspaceId: string,
|
||||
scope: WorkspaceFileScope,
|
||||
options?: { maxRows?: number; throwOnError?: boolean }
|
||||
) {
|
||||
const maxRows = options?.maxRows
|
||||
const [files, shares] = await Promise.all([
|
||||
listWorkspaceFiles(workspaceId, { scope }),
|
||||
listWorkspaceFiles(workspaceId, {
|
||||
scope,
|
||||
limit: maxRows === undefined ? undefined : maxRows + 1,
|
||||
throwOnError: options?.throwOnError,
|
||||
}),
|
||||
getWorkspaceShares('file', workspaceId),
|
||||
])
|
||||
if (maxRows !== undefined && files.length > maxRows) return null
|
||||
|
||||
const withShares = files.map((file) => ({ ...file, share: shares.get(file.id) ?? null }))
|
||||
return listWorkspaceFilesContract.response.schema.shape.files.parse(withShares)
|
||||
}
|
||||
|
||||
@@ -745,6 +745,33 @@ describe('Permission Utils', () => {
|
||||
|
||||
expect(result).toEqual({ id: 'workspace123', ownerId: null })
|
||||
})
|
||||
|
||||
/**
|
||||
* Archived visibility is applied in JS, not SQL, so the read can be shared by the
|
||||
* gates that disagree about it. That makes these two the boundary worth pinning:
|
||||
* if the filter ever stops matching the old `archived_at IS NULL` predicate, archived
|
||||
* workspaces silently become visible to callers that asked not to see them.
|
||||
*/
|
||||
it.concurrent('should hide an archived workspace by default', async () => {
|
||||
const chain = createMockChain([
|
||||
{ id: 'workspace123', ownerId: 'owner456', archivedAt: new Date('2026-01-01') },
|
||||
])
|
||||
mockDb.select.mockReturnValue(chain)
|
||||
|
||||
const result = await getWorkspaceWithOwner('workspace123')
|
||||
|
||||
expect(result).toBeNull()
|
||||
})
|
||||
|
||||
it.concurrent('should return an archived workspace when asked to include them', async () => {
|
||||
const archivedAt = new Date('2026-01-01')
|
||||
const chain = createMockChain([{ id: 'workspace123', ownerId: 'owner456', archivedAt }])
|
||||
mockDb.select.mockReturnValue(chain)
|
||||
|
||||
const result = await getWorkspaceWithOwner('workspace123', { includeArchived: true })
|
||||
|
||||
expect(result).toEqual({ id: 'workspace123', ownerId: 'owner456', archivedAt })
|
||||
})
|
||||
})
|
||||
|
||||
describe('workspaceExists', () => {
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import { cache } from 'react'
|
||||
import { db } from '@sim/db'
|
||||
import { member, permissions, user, type WorkspaceMode, workspace } from '@sim/db/schema'
|
||||
import {
|
||||
@@ -77,17 +78,12 @@ export async function getWorkspaceById(
|
||||
return exists ? { id: workspaceId } : null
|
||||
}
|
||||
|
||||
/**
|
||||
* Get a workspace with owner info by ID
|
||||
*
|
||||
* @param workspaceId - The workspace ID to look up
|
||||
* @returns The workspace with owner info if found, null otherwise
|
||||
*/
|
||||
export async function getWorkspaceWithOwner(
|
||||
async function selectWorkspaceWithOwner(
|
||||
workspaceId: string,
|
||||
options?: { includeArchived?: boolean; executor?: DbOrTx; forUpdate?: boolean }
|
||||
includeArchived: boolean,
|
||||
executor: DbOrTx,
|
||||
forUpdate: boolean
|
||||
): Promise<WorkspaceWithOwner | null> {
|
||||
const { includeArchived = false, executor = db, forUpdate = false } = options ?? {}
|
||||
const query = executor
|
||||
.select({
|
||||
id: workspace.id,
|
||||
@@ -110,6 +106,46 @@ export async function getWorkspaceWithOwner(
|
||||
return ws || null
|
||||
}
|
||||
|
||||
/**
|
||||
* Request-memoized plain workspace read, keyed by id alone. One render pass resolves the
|
||||
* same row through several independent gates.
|
||||
*
|
||||
* Keyed on the id and NOT on archived visibility: the gates disagree about archived
|
||||
* workspaces, so memoizing that argument would give each answer its own entry and dedupe
|
||||
* nothing.
|
||||
*
|
||||
* The returned row is SHARED by every consumer in the render. Treat it as immutable — an
|
||||
* in-place edit poisons every gate that reads it for the rest of the pass.
|
||||
*/
|
||||
const readWorkspaceWithOwner = cache(
|
||||
(workspaceId: string): Promise<WorkspaceWithOwner | null> =>
|
||||
selectWorkspaceWithOwner(workspaceId, true, db, false)
|
||||
)
|
||||
|
||||
/**
|
||||
* Get a workspace with owner info by ID
|
||||
*
|
||||
* Transaction-scoped (`executor`) and lock-acquiring (`forUpdate`) reads
|
||||
* deliberately bypass {@link readWorkspaceWithOwner}: a row read inside one
|
||||
* caller's transaction, or under a row lock only that caller holds, must never
|
||||
* be handed to a later caller that took neither.
|
||||
*
|
||||
* @param workspaceId - The workspace ID to look up
|
||||
* @returns The workspace with owner info if found, null otherwise
|
||||
*/
|
||||
export async function getWorkspaceWithOwner(
|
||||
workspaceId: string,
|
||||
options?: { includeArchived?: boolean; executor?: DbOrTx; forUpdate?: boolean }
|
||||
): Promise<WorkspaceWithOwner | null> {
|
||||
const { includeArchived = false, executor, forUpdate = false } = options ?? {}
|
||||
if (executor || forUpdate) {
|
||||
return selectWorkspaceWithOwner(workspaceId, includeArchived, executor ?? db, forUpdate)
|
||||
}
|
||||
const ws = await readWorkspaceWithOwner(workspaceId)
|
||||
if (!ws) return null
|
||||
return includeArchived || !ws.archivedAt ? ws : null
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve the effective workspace permission for a user under the governance
|
||||
* inheritance model: the owners/admins of the organization that owns the
|
||||
@@ -131,18 +167,7 @@ export async function getEffectiveWorkspacePermission(
|
||||
return resolveEffectiveWorkspacePermission(userId, ws.id, ws.organizationId, executor)
|
||||
}
|
||||
|
||||
/**
|
||||
* Check workspace access for a user
|
||||
*
|
||||
* Verifies the workspace exists and the user has access to it.
|
||||
* Returns access level (read/write) based on ownership, explicit permissions,
|
||||
* and organization-admin inheritance.
|
||||
*
|
||||
* @param workspaceId - The workspace ID to check
|
||||
* @param userId - The user ID to check access for
|
||||
* @returns WorkspaceAccess object with exists, hasAccess, canWrite, and workspace data
|
||||
*/
|
||||
export async function checkWorkspaceAccess(
|
||||
async function resolveWorkspaceAccessForUser(
|
||||
workspaceId: string,
|
||||
userId: string
|
||||
): Promise<WorkspaceAccess> {
|
||||
@@ -167,6 +192,25 @@ export async function checkWorkspaceAccess(
|
||||
return { exists: true, hasAccess, canWrite, canAdmin, workspace: ws, permission }
|
||||
}
|
||||
|
||||
/**
|
||||
* Check workspace access for a user
|
||||
*
|
||||
* Verifies the workspace exists and the user has access to it.
|
||||
* Returns access level (read/write) based on ownership, explicit permissions,
|
||||
* and organization-admin inheritance.
|
||||
*
|
||||
* Request-memoized: a Server Component render pass authorizes the same
|
||||
* (workspace, viewer) pair through several gates, and the answer cannot change
|
||||
* mid-render. Outside a render React evaluates the resolver normally, so route
|
||||
* handlers and background work re-read exactly as before. Takes no executor, so
|
||||
* no transaction-scoped or locked read can ever be memoized here.
|
||||
*
|
||||
* @param workspaceId - The workspace ID to check
|
||||
* @param userId - The user ID to check access for
|
||||
* @returns WorkspaceAccess object with exists, hasAccess, canWrite, and workspace data
|
||||
*/
|
||||
export const checkWorkspaceAccess = cache(resolveWorkspaceAccessForUser)
|
||||
|
||||
/**
|
||||
* Returns `provided` when it was resolved for this exact workspace, otherwise
|
||||
* resolves fresh. The id match is what keeps a caller from authorizing against
|
||||
|
||||
+71
-37
@@ -56,6 +56,7 @@ import {
|
||||
} from '@/lib/execution/private-tool-metadata'
|
||||
import { parseMcpToolId } from '@/lib/mcp/utils'
|
||||
import { hostedKeyMetrics } from '@/lib/monitoring/metrics'
|
||||
import type { CredentialTokenPayload } from '@/lib/oauth/token-resolution'
|
||||
import { resolveWorkspaceFileReference } from '@/lib/uploads/contexts/workspace/workspace-file-manager'
|
||||
import { markWorkspaceFileSecretProvenanceUnknown } from '@/lib/uploads/contexts/workspace/workspace-file-secret-provenance'
|
||||
import { assertPermissionsAllowed } from '@/ee/access-control/utils/permission-check'
|
||||
@@ -1702,8 +1703,6 @@ async function executeToolImplementation(
|
||||
`[${requestId}] Tool ${toolId} needs access token for credential: ${contextParams.credential}`
|
||||
)
|
||||
try {
|
||||
const baseUrl = getInternalApiBaseUrl()
|
||||
|
||||
const workflowId = contextParams._context?.workflowId
|
||||
const userId = contextParams._context?.userId
|
||||
|
||||
@@ -1724,51 +1723,86 @@ async function executeToolImplementation(
|
||||
}
|
||||
}
|
||||
|
||||
logger.info(`[${requestId}] Fetching access token from ${baseUrl}/api/auth/oauth/token`)
|
||||
/**
|
||||
* The acting user asserted alongside an internal token. Only sent when the
|
||||
* run enforces credential access, matching the `userId` query param the HTTP
|
||||
* surface accepted — it never widens access, it only pins the assertion to
|
||||
* the token subject.
|
||||
*/
|
||||
const callerUserId =
|
||||
userId && contextParams._context?.enforceCredentialAccess ? userId : undefined
|
||||
|
||||
const tokenUrlObj = new URL('/api/auth/oauth/token', baseUrl)
|
||||
if (workflowId) {
|
||||
tokenUrlObj.searchParams.set('workflowId', workflowId)
|
||||
}
|
||||
if (userId && contextParams._context?.enforceCredentialAccess) {
|
||||
tokenUrlObj.searchParams.set('userId', userId)
|
||||
}
|
||||
let data: CredentialTokenPayload
|
||||
|
||||
// Always send Content-Type; add internal auth on server-side runs
|
||||
const tokenHeaders: Record<string, string> = { 'Content-Type': 'application/json' }
|
||||
if (typeof window === 'undefined') {
|
||||
try {
|
||||
const internalToken = await generateInternalToken(userId)
|
||||
tokenHeaders.Authorization = `Bearer ${internalToken}`
|
||||
} catch (_e) {
|
||||
// Swallow token generation errors; the request will fail and be reported upstream
|
||||
// Server-side runs resolve the credential through the same application
|
||||
// operation the route calls, rather than minting an internal JWT and
|
||||
// POSTing to ourselves through the load balancer. The synthesized
|
||||
// `AuthResult` is exactly what verifying that self-issued token would
|
||||
// have produced, so authorization, refresh, and audit are unchanged —
|
||||
// including failing closed when the run carries no user id.
|
||||
const { resolveCredentialToken } = await import('@/lib/oauth/token-resolution')
|
||||
const result = await resolveCredentialToken(
|
||||
{ success: true, authType: 'internal_jwt', userId },
|
||||
{
|
||||
requestId,
|
||||
credentialId: contextParams.credential as string,
|
||||
workflowId,
|
||||
scopes: tokenPayload.scopes,
|
||||
impersonateEmail: tokenPayload.impersonateEmail,
|
||||
callerUserId,
|
||||
}
|
||||
)
|
||||
|
||||
if (!result.ok) {
|
||||
logger.error(`[${requestId}] Token fetch failed for ${toolId}:`, {
|
||||
status: result.status,
|
||||
error: result.error,
|
||||
})
|
||||
const toolLabel = tool?.name || toolId
|
||||
throw new Error(`Failed to obtain credential for ${toolLabel}: ${result.error}`)
|
||||
}
|
||||
}
|
||||
|
||||
const response = await fetch(tokenUrlObj.toString(), {
|
||||
method: 'POST',
|
||||
headers: tokenHeaders,
|
||||
body: JSON.stringify(tokenPayload),
|
||||
})
|
||||
data = result.token
|
||||
} else {
|
||||
const baseUrl = getInternalApiBaseUrl()
|
||||
logger.info(`[${requestId}] Fetching access token from ${baseUrl}/api/auth/oauth/token`)
|
||||
|
||||
if (!response.ok) {
|
||||
const errorText = await response.text()
|
||||
logger.error(`[${requestId}] Token fetch failed for ${toolId}:`, {
|
||||
status: response.status,
|
||||
error: errorText,
|
||||
const tokenUrlObj = new URL('/api/auth/oauth/token', baseUrl)
|
||||
if (workflowId) {
|
||||
tokenUrlObj.searchParams.set('workflowId', workflowId)
|
||||
}
|
||||
if (callerUserId) {
|
||||
tokenUrlObj.searchParams.set('userId', callerUserId)
|
||||
}
|
||||
|
||||
// boundary-raw-fetch: browser-side tool runs authenticate with the session cookie against the same-origin token route
|
||||
const response = await fetch(tokenUrlObj.toString(), {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify(tokenPayload),
|
||||
})
|
||||
let parsedError = errorText
|
||||
try {
|
||||
const parsed = JSON.parse(errorText)
|
||||
if (parsed.error) parsedError = parsed.error
|
||||
} catch {
|
||||
// Use raw text
|
||||
|
||||
if (!response.ok) {
|
||||
const errorText = await response.text()
|
||||
logger.error(`[${requestId}] Token fetch failed for ${toolId}:`, {
|
||||
status: response.status,
|
||||
error: errorText,
|
||||
})
|
||||
let parsedError = errorText
|
||||
try {
|
||||
const parsed = JSON.parse(errorText)
|
||||
if (parsed.error) parsedError = parsed.error
|
||||
} catch {
|
||||
// Use raw text
|
||||
}
|
||||
const toolLabel = tool?.name || toolId
|
||||
throw new Error(`Failed to obtain credential for ${toolLabel}: ${parsedError}`)
|
||||
}
|
||||
const toolLabel = tool?.name || toolId
|
||||
throw new Error(`Failed to obtain credential for ${toolLabel}: ${parsedError}`)
|
||||
|
||||
data = (await response.json()) as CredentialTokenPayload
|
||||
}
|
||||
|
||||
const data = await response.json()
|
||||
contextParams.accessToken = data.accessToken
|
||||
if (data.idToken) {
|
||||
contextParams.idToken = data.idToken
|
||||
|
||||
Reference in New Issue
Block a user