Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -11,14 +11,16 @@ import { useRouter } from 'next/navigation'
import { CHAT_PLAN_REQUIRED_UPGRADE } from '@/lib/billing/chat-plan'
import { rememberSettingsReturnUrl } from '@/lib/navigation/settings-return'
import { usePlanUpgradeLink } from '@/app/workspace/[workspaceId]/home/hooks/use-plan-upgrade-link'
import { useOptionalWorkspaceHostContext } from '@/app/workspace/[workspaceId]/providers/workspace-host-provider'
import { useChatPlanModalStore } from '@/stores/modals/chat-plan/store'

/**
* Answers a send the workspace's plan refuses: Chat is on paid plans, with the way to upgrade
* for a viewer who may, and who to ask for one who may not. Render it once per workspace shell.
*/
export function ChatPlanModal() {
const isOpen = useChatPlanModalStore((state) => state.isOpen)
const workspaceId = useOptionalWorkspaceHostContext()?.workspace.id
const isOpen = useChatPlanModalStore((state) => state.workspaceId === workspaceId)
const close = useChatPlanModalStore((state) => state.close)

return (
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -533,7 +533,7 @@ const UserInputImpl = forwardRef<UserInputHandle, UserInputProps>(function UserI

const handleSubmit = useCallback(() => {
if (chatPlanRequired) {
useChatPlanModalStore.getState().open()
useChatPlanModalStore.getState().open(workspaceId)
return
}
const currentFiles = filesRef.current
Expand Down Expand Up @@ -561,7 +561,7 @@ const UserInputImpl = forwardRef<UserInputHandle, UserInputProps>(function UserI
)
if (accepted === false) return
clearComposer()
}, [chatPlanRequired, onSubmit, clearComposer])
}, [chatPlanRequired, workspaceId, onSubmit, clearComposer])

/**
* Enter policy for the editor: mirror canSubmit's uploading guard (Enter
Expand Down
4 changes: 2 additions & 2 deletions apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4443,7 +4443,7 @@ export function useChat(
) => {
if ((!message.trim() && !fileAttachments?.length) || !scopeKey) return
if (chatPlanRequired) {
useChatPlanModalStore.getState().open()
useChatPlanModalStore.getState().open(workspaceId)
return
}

Expand Down Expand Up @@ -5307,7 +5307,7 @@ export function useChat(
if (!msg || queueState.editing[chatKey] === msg.id) return
if (queuedMessageDispatchIds.has(msg.id)) return
if (chatPlanRequired) {
useChatPlanModalStore.getState().open()
useChatPlanModalStore.getState().open(workspaceId)
return
}
/* Sent by hand, it still must not run a second turn: if the history shows it
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,9 @@ export function usePlanUpgradeLink(data: PlanUpgradeLinkInput): PlanUpgradeLink
? organizationContext.viewer.isAdmin
: Boolean(hostContext && canManageWorkspaceBilling(hostContext, session?.user?.id)))
const action = planRequired
? 'upgrade this workspace’s plan'
? organizationContext
? 'upgrade the organization’s plan'
: 'upgrade this workspace’s plan'
: 'manage this workspace’s usage limits'
const unavailableMessage =
organizationContext && !planRequired
Expand Down
30 changes: 23 additions & 7 deletions apps/sim/lib/execution/remote-sandbox/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -327,10 +327,15 @@ async function throwIfRunTimedOut(
if (!result.timedOut) return
if (!run.lease.session) throw new DOMException('timeout', 'AbortError')
const { outputSandboxDir, sandbox } = run
// Half of what the call has left, so the report lands before the call's deadline.
const listingMs = Math.min(
TIMED_OUT_OUTPUT_LISTING_MS,
(getRemainingExecutionMs(run.signal) ?? Number.POSITIVE_INFINITY) / 2
)
const listed = outputSandboxDir
? await withinDeadline(
() => sandbox.listFiles(outputSandboxDir, { depth: SANDBOX_OUTPUT_DIR_MAX_DEPTH }),
Date.now() + TIMED_OUT_OUTPUT_LISTING_MS,
Date.now() + listingMs,
run.signal
).catch(() => [])
: []
Expand Down Expand Up @@ -908,21 +913,32 @@ function installBudgetMs(timeoutMs: number): number {
*/
const MIN_EXPORT_BUDGET_MS = 10_000

/**
* Held back from a workbench run without an export reserve, capped at a quarter of its budget,
* so the run's own timeout answers and reports its progress before the call's deadline.
*/
const MAX_TIMEOUT_REPORT_BUDGET_MS = 4_000

/**
* The budget handed to user code, less an export reserve when this request will
* read files back. Short budgets are left alone: taking the reserve out of one
* would starve the code to buy time for an export it never reaches.
* would starve the code to buy time for an export it never reaches. A workbench
* run, whose timeout reports progress, still keeps a smaller report reserve.
*/
function codeBudgetMs(
req: { outputSandboxPath?: string; outputSandboxPaths?: string[]; outputSandboxDir?: string },
signal: AbortSignal
signal: AbortSignal,
reportsTimeout: boolean
): number {
const remainingMs = remainingSandboxBudgetMs(signal)
const exportsFiles = Boolean(
req.outputSandboxDir || req.outputSandboxPath || req.outputSandboxPaths?.length
)
if (!exportsFiles || remainingMs <= MIN_EXPORT_BUDGET_MS * 2) return remainingMs
return remainingMs - MIN_EXPORT_BUDGET_MS
if (exportsFiles && remainingMs > MIN_EXPORT_BUDGET_MS * 2) {
return remainingMs - MIN_EXPORT_BUDGET_MS
}
if (!reportsTimeout) return remainingMs
return remainingMs - Math.min(MAX_TIMEOUT_REPORT_BUDGET_MS, Math.floor(remainingMs / 4))
}

/**
Expand Down Expand Up @@ -1060,7 +1076,7 @@ async function executeInSandboxWithinBudget(
signal,
phases,
(verifyCli) => {
runTimeoutMs = codeBudgetMs(req, signal)
runTimeoutMs = codeBudgetMs(req, signal, Boolean(lease.session))
return sandbox.runCode(code, {
timeoutMs: runTimeoutMs,
javascriptPreload: buildJavaScriptRuntimeBindingsSource(req.runtimeBindings ?? []),
Expand Down Expand Up @@ -1283,7 +1299,7 @@ async function executeShellInSandboxWithinBudget(
signal,
phases,
(verifyCli) => {
runTimeoutMs = codeBudgetMs(req, signal)
runTimeoutMs = codeBudgetMs(req, signal, Boolean(lease.session))
return sandbox.runCommand(code, {
envs: commandEnvironment,
timeoutMs: runTimeoutMs,
Expand Down
37 changes: 37 additions & 0 deletions apps/sim/lib/execution/remote-sandbox/session-sandbox.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1130,6 +1130,43 @@ describe('session sandbox lease', () => {
}
})

it('reports the output of a short run whose own timeout answers after a round trip', async () => {
vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout', 'Date'] })
try {
const { handle } = fakeSandbox('timeout-short')
mockFindSessionSandbox.mockResolvedValue(handle)
// The provider answers its own timeout one round trip late, and stops on the call's abort.
handle.runCommand = (_command, options) =>
new Promise((resolve, reject) => {
const timer = setTimeout(
() => resolve({ stdout: 'step 1', stderr: '', exitCode: 124, timedOut: true }),
(options?.timeoutMs ?? 0) + 300
)
options?.signal?.addEventListener('abort', () => {
clearTimeout(timer)
reject(options.signal?.reason)
})
})
handle.listFiles = () => new Promise<never>(() => {})
const pending = executeShellInSandbox({
...CODE_REQUEST,
timeoutMs: 10_000,
envs: {},
session: { key: 'timeout-short' },
outputSandboxDir: OUTPUT_DIR,
}).then(
() => undefined,
(rejection: unknown) => rejection
)
await vi.runAllTimersAsync()
const error = await pending
if (!(error instanceof SandboxExecutionTimeoutError)) throw error
expect(error.progress.output).toBe('step 1')
} finally {
vi.useRealTimers()
}
})

it('keeps a cancellation that arrives while listing the outputs a cancellation', async () => {
const { handle } = fakeSandbox('timeout-cancelled')
mockFindSessionSandbox.mockResolvedValue(handle)
Expand Down
23 changes: 23 additions & 0 deletions apps/sim/lib/mothership/tools/server/settings.test.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import type { OrganizationDelegatedPrincipal } from '@sim/auth/principal'
import { featureFlagsMock, featureFlagsMockFns } from '@sim/testing/mocks/feature-flags.mock'
import {
knowledgeAvailabilityMock,
knowledgeAvailabilityMockFns,
Expand All @@ -11,6 +12,7 @@ vi.mock('@/lib/settings/application/organization-section-access', () => ({
authorizeOrganizationSettingsSection: sectionAccess,
}))
vi.mock('@/lib/knowledge/access/availability', () => knowledgeAvailabilityMock)
vi.mock('@/lib/core/config/feature-flags', () => featureFlagsMock)
vi.mock('@/lib/mothership/application/settings-context', () => ({
resolveSettingsContext: resolve,
}))
Expand Down Expand Up @@ -98,6 +100,27 @@ describe('settings tool dispatch', () => {
})
})

it('offers Sim links only to an organization that has them enabled', async () => {
const simLinks = expect.objectContaining({ id: 'sim-links' })
featureFlagsMockFns.mockIsFeatureEnabled.mockResolvedValue(false)
expect(
await settingsServerTool.execute({ scope: 'organization', action: 'list' })
).not.toMatchObject({ sections: expect.arrayContaining([simLinks]) })
for (const action of ['open', 'get'] as const) {
await expect(
settingsServerTool.execute({ scope: 'organization', action, section: 'sim-links' })
).rejects.toThrow('unavailable')
}

featureFlagsMockFns.mockIsFeatureEnabled.mockImplementation(
async (flag, context) =>
flag === 'sim-links' && (context as { orgId?: string })?.orgId === 'org'
)
expect(
await settingsServerTool.execute({ scope: 'organization', action: 'list' })
).toMatchObject({ sections: expect.arrayContaining([simLinks]) })
})

it('discovers names without eagerly reading every setting or dumping every schema', async () => {
const read = vi.spyOn(readOrganizationSettings, 'execute')
const result = await settingsServerTool.execute({ scope: 'organization', action: 'list' })
Expand Down
9 changes: 8 additions & 1 deletion apps/sim/lib/mothership/tools/server/settings.ts
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ import {
} from '@/lib/organizations/application/settings'
import { authorizeOrganizationSettingsSection } from '@/lib/settings/application/organization-section-access'
import { type SettingsSection, settingsSections } from '@/lib/settings/sections'
import { isSimLinksEnabled } from '@/lib/sim-links/feature-flag'
import {
delegatedAccountPreferencesSchema,
updateCurrentUserPreferences,
Expand All @@ -86,6 +87,8 @@ interface SettingAdapter {
get(context: SettingsContext): Promise<unknown>
schema?: z.ZodType
update?(context: SettingsContext, changes: Record<string, unknown>): Promise<unknown>
/** Whether the section exists for this target at all; a section that does not is never listed. */
available?(context: SettingsContext): Promise<boolean>
}

function organizationInput(context: SettingsContext) {
Expand Down Expand Up @@ -254,7 +257,10 @@ const adapters: Record<string, SettingAdapter> = {
},
'organization/members': { get: readSettingsRoster },
'organization/requests': { get: (context) => readSettingsAccessRequests(context) },
'organization/sim-links': { get: (context) => readSettingsSimLinks(context) },
'organization/sim-links': {
available: (context) => isSimLinksEnabled(organizationInput(context).organizationId),
get: (context) => readSettingsSimLinks(context),
},
'organization/whitelabeling': {
schema: updateOrganizationWhitelabelBodySchema.strict(),
get: async (context) =>
Expand Down Expand Up @@ -315,6 +321,7 @@ const adapters: Record<string, SettingAdapter> = {

async function describeSection(context: SettingsContext, section: SettingsSection) {
const adapter = adapters[`${context.scope}/${section.id}`]
if (adapter?.available && !(await adapter.available(context))) return null
const uiSection =
context.scope === 'organization' &&
section.id === 'connected-accounts' &&
Expand Down
17 changes: 9 additions & 8 deletions apps/sim/stores/modals/chat-plan/store.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,22 +2,23 @@ import { create } from 'zustand'
import { devtools } from 'zustand/middleware'

interface ChatPlanModalState {
isOpen: boolean
open: () => void
/** The workspace whose refused send opened the modal; another workspace never shows it. */
workspaceId: string | null
open: (workspaceId: string | undefined) => void
close: () => void
}

/**
* Whether the Chat plan modal is showing. A send the payer's plan refuses opens it from wherever
* the send started (the composer, a queued message, a handoff from another surface), and the
* workspace shell renders the one modal.
* Which workspace's Chat plan modal is showing. A send the payer's plan refuses opens it from
* wherever the send started (the composer, a queued message, a handoff from another surface), and
* the workspace shell renders the one modal, only for the workspace that opened it.
*/
export const useChatPlanModalStore = create<ChatPlanModalState>()(
devtools(
(set) => ({
isOpen: false,
open: () => set({ isOpen: true }),
close: () => set({ isOpen: false }),
workspaceId: null,
open: (workspaceId) => set({ workspaceId: workspaceId ?? null }),
Comment thread
waleedlatif1 marked this conversation as resolved.
Resolved
close: () => set({ workspaceId: null }),
}),
{ name: 'chat-plan-modal-store' }
)
Expand Down
Loading