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
29 changes: 28 additions & 1 deletion apps/sim/app/api/billing/update-cost/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ import { BILLING_CALLBACK_OUTCOME } from '@/lib/mothership/generated/billing-pro
import { BillingRouteOutcome } from '@/lib/mothership/generated/trace-attribute-values-v1'
import { TraceAttr } from '@/lib/mothership/generated/trace-attributes-v1'
import { TraceSpan } from '@/lib/mothership/generated/trace-spans-v1'
import { recordJobRunSpend } from '@/lib/mothership/jobs/budget'
import { checkInternalApiKey } from '@/lib/mothership/request/http'
import { withIncomingGoSpan } from '@/lib/mothership/request/otel'

Expand Down Expand Up @@ -122,6 +123,28 @@ async function readUsageStanding(
}
}

/**
* Adds a scheduled job run's own spend cap to the payer's standing. A payer over its plan limit
* keeps the plan card; otherwise a job run at its cap pauses with the job's card. A cap that
* cannot be read answers as the payer's standing alone, like a standing read that fails.
*/
async function withJobRunBudget(
verdict: BillingUsageVerdict,
billingRequestId: string | undefined,
cumulativeCostUsd: number
): Promise<BillingUsageVerdict> {
if (verdict.usageExceeded || !billingRequestId) return verdict
try {
const card = await recordJobRunSpend(billingRequestId, cumulativeCostUsd)
return card ? { usageExceeded: true, usageUpgrade: card } : verdict
} catch (error) {
logger.warn('Job run spend could not be recorded; answering the payer standing alone', {
error: toError(error).message,
})
return verdict
}
}

function getBillingResolution(
isMarkerlessLegacy: boolean,
billingAttribution: BillingAttributionSnapshot | undefined
Expand Down Expand Up @@ -442,7 +465,11 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
await checkAndBillOverageThreshold(userId, undefined, { onError: 'throw' })
}

const usageVerdict = await readUsageStanding(userId, billingAttribution, accountDecision)
const usageVerdict = await withJobRunBudget(
await readUsageStanding(userId, billingAttribution, accountDecision),
billingRequestId,
cost
)
const duration = Date.now() - startTime

// Same-or-lower cumulative than already recorded: nothing new to bill.
Expand Down
5 changes: 5 additions & 0 deletions apps/sim/app/api/schedules/execute/route.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,11 @@ vi.mock('@/lib/execution/manual-cancellation', () => ({

vi.mock('@/lib/core/async-jobs', () => asyncJobsMock)

/** The job lane has its own integration suite. */
vi.mock('@/lib/mothership/jobs/lane', () => ({
runJobLane: vi.fn(async () => ({ claimed: 0 })),
}))

vi.mock('drizzle-orm', () => ({
and: vi.fn((...conditions: unknown[]) => ({ type: 'and', conditions })),
eq: vi.fn((field: unknown, value: unknown) => ({ field, value, type: 'eq' })),
Expand Down
8 changes: 8 additions & 0 deletions apps/sim/app/api/schedules/execute/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ import {
registerManualExecutionAborter,
unregisterManualExecutionAborter,
} from '@/lib/execution/manual-cancellation'
import { runJobLane } from '@/lib/mothership/jobs/lane'
import {
buildCarrierReconciledMetadata,
carrierNotReconciledSql,
Expand Down Expand Up @@ -1788,6 +1789,13 @@ export async function runScheduleTick(requestId: string): Promise<ScheduleTickRe
await Promise.allSettled(schedulePromises)
}

try {
const { claimed } = await runJobLane()
totalSchedules += claimed
} catch (error) {
logger.error(`[${requestId}] Scheduled job lane failed`, { error: toError(error).message })
}

const totalCount = totalSchedules
const durationMs = Date.now() - tickStart
logger.info(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -523,3 +523,28 @@ describe('CredentialDisplay link tag', () => {
act(() => root.unmount())
})
})

describe('UsageUpgradeDisplay', () => {
it("shows a job run's own spend cap without an account billing action", () => {
;(globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true
const container = document.createElement('div')
const root = createRoot(container)
act(() => {
root.render(
<SpecialTags
segment={{
type: 'usage_upgrade',
data: {
reason: 'job_run_budget',
action: 'increase_limit',
message: 'This scheduled job run reached its $5.00 spend limit.',
},
}}
/>
)
})

expect(container.querySelector('a')).toBeNull()
act(() => root.unmount())
})
})
Original file line number Diff line number Diff line change
Expand Up @@ -3443,6 +3443,43 @@ function MothershipErrorDisplay({ data }: { data: MothershipErrorTagData }) {
}

function UsageUpgradeDisplay({ data }: { data: UsageUpgradeTagData }) {
return data.reason === 'job_run_budget' ? (
<JobRunBudgetDisplay message={data.message} />
) : (
<PlanUsageLimitDisplay data={data} />
)
}

/**
* A scheduled job's run reached its own spend cap. The cap belongs to the job, not to the
* account, so the card offers no billing action: the next run starts with a fresh cap.
*/
interface JobRunBudgetDisplayProps {
message: string
}

function JobRunBudgetDisplay({ message }: JobRunBudgetDisplayProps) {
return (
<InteractionCard
title={
<span className='inline-flex items-center gap-2'>
<TriangleAlert
aria-hidden
className='size-[14px] shrink-0 text-[var(--badge-amber-text)]'
/>
Run Spend Limit Reached
</span>
}
>
<div className='px-2 pb-2'>
<p className='text-[var(--text-body)] text-small leading-5'>{message}</p>
</div>
</InteractionCard>
)
}

/** The payer's plan limit: the card links to where the plan or the limit can change. */
function PlanUsageLimitDisplay({ data }: { data: UsageUpgradeTagData }) {
const { data: session } = useSession()
const hostContext = useOptionalWorkspaceHostContext()
const organizationContext = useOptionalOrganizationContext()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import { useQueryClient } from '@tanstack/react-query'
import { defaultRangeExtractor, type Range, useVirtualizer } from '@tanstack/react-virtual'
import { ComputerUseActivity } from '@/components/computer-use/activity'
import { SMOOTH_CHASE_RATE } from '@/lib/core/utils/smooth-bottom-chase'
import type { TurnOrigin } from '@/lib/mothership/generated/protocol'
import type { WorkspaceFileRecord } from '@/lib/uploads/contexts/workspace'
import { inter } from '@/app/_styles/fonts/inter/inter'
import { FindBar } from '@/app/workspace/[workspaceId]/components/find-bar'
Expand Down Expand Up @@ -187,24 +188,36 @@ interface UserMessageRowProps {
}

/**
* A background-task notification that opened a turn (mothership 21-background-tasks.md
* §6.4): a muted system chip, never a user bubble — the user did not type it. Shows the
* outcome line; the provenance header and the agent's own note stay out of the way.
* A turn automation opened, never a user bubble — the user did not type it: a muted system chip
* with one line. A background-task notification (mothership 21-background-tasks.md §6.4) shows
* its outcome line; a scheduled job's run shows the brief's first line (the job and its run).
*/
const TaskNotificationRow = memo(function TaskNotificationRow({
content,
rowClassName,
}: {
interface AutomatedTurnRowProps {
origin: TurnOrigin
content: string
rowClassName: string
}) {
}

const AutomatedTurnRow = memo(function AutomatedTurnRow({
origin,
content,
rowClassName,
}: AutomatedTurnRowProps) {
const lines = content.split('\n').filter((line) => line.length > 0)
const outcome = lines.find((line) => line.startsWith('Task ')) ?? lines[lines.length - 1] ?? ''
const summary =
origin === 'task'
? (lines.find((line) => line.startsWith('Task ')) ?? lines[lines.length - 1] ?? '').replace(
/^Task [0-9a-f-]+ /,
''
)
: (lines[0] ?? '').replace(/^Job: /, '')
return (
<div className={cn('flex w-full justify-center', rowClassName)}>
<div className='max-w-[85%] rounded-full border border-[var(--border)] bg-[var(--bg)] px-3 py-1 text-[12px] text-[var(--text-secondary)]'>
<span className='mr-1.5 font-medium text-[var(--text-primary)]'>Background task</span>
{outcome.replace(/^Task [0-9a-f-]+ /, '')}
<span className='mr-1.5 font-medium text-[var(--text-primary)]'>
{origin === 'task' ? 'Background task' : 'Scheduled job'}
</span>
{summary}
</div>
</div>
)
Expand Down Expand Up @@ -930,8 +943,12 @@ export function MothershipChat({
style={{ top: virtualItem.start }}
>
{msg.role === 'user' ? (
interactionPairing.hiddenUserByIndex[index] ? null : msg.origin === 'task' ? (
<TaskNotificationRow content={msg.content} rowClassName={styles.rowGap} />
interactionPairing.hiddenUserByIndex[index] ? null : msg.origin ? (
<AutomatedTurnRow
origin={msg.origin}
content={msg.content}
rowClassName={styles.rowGap}
/>
) : (
<UserMessageRow
content={msg.content}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,18 @@ describe('useChatFind', () => {
expect(find.current).toMatchObject({ isOpen: false, query: '', count: 0 })
})

it('leaves turns automation opened out of the search, whatever opened them', () => {
const find = renderFind({
messages: [
{ id: 't', role: 'user', origin: 'task', content: 'Task abc finished: triage done.' },
{ id: 'j', role: 'user', origin: 'job', content: 'Job: Nightly triage (Run 1)' },
assistant('a', 'The triage is done.'),
],
})
find.search('triage')
expect(find.current.count).toBe(1)
})

it('counts the same footnoted matches in the index as it highlights in the rendered message', async () => {
const find = renderFind({
chatId: 'chat',
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,7 @@ export function useChatFind({
if (term) {
for (let messageIndex = messages.length - 1; messageIndex >= 0; messageIndex--) {
const message = messages[messageIndex]
if (!message || hiddenUserByIndex[messageIndex] || message.origin === 'task') continue
if (!message || hiddenUserByIndex[messageIndex] || message.origin) continue
const text = getMessageSearchText(message, textCache)
const occurrences: ChatFindMatch[] = []
let found = 0
Expand Down
6 changes: 3 additions & 3 deletions apps/sim/app/workspace/[workspaceId]/home/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ import type { WorkspaceSearchFilters } from '@/lib/api/contracts/knowledge/searc
import type { MothershipChat } from '@/lib/api/contracts/mothership-chats'
import type { ManagedMcpConnectorId } from '@/lib/credential-groups/managed-mcp-connectors'
import type { AssistantSearchLevel } from '@/lib/mothership/generated/assistant'
import type { ToolActivity } from '@/lib/mothership/generated/protocol'
import type { ToolActivity, TurnOrigin } from '@/lib/mothership/generated/protocol'
import type { ChatContext } from '@/stores/panel'
import type { BrowserTextSelection, TerminalTextSelection } from '@/stores/panel/types'

Expand Down Expand Up @@ -197,8 +197,8 @@ export interface ChatMessage {
role: 'user' | 'assistant'
requestMode?: 'agent' | 'assistant' | 'plan'
content: string
/** "task": a background-task notification opened this turn, not the user (a system chip, not a bubble). */
origin?: 'task'
/** Automation opened this turn, not the user: a system chip, not a bubble. */
origin?: TurnOrigin
contentBlocks?: ContentBlock[]
attachments?: ChatMessageAttachment[]
fileAttachments?: FileAttachmentForApi[]
Expand Down
11 changes: 8 additions & 3 deletions apps/sim/lib/api/contracts/subscription.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,11 @@ import {
usageLimitReadSchema,
usageLimitUpdateSchema,
} from '@/lib/billing/application/usage-limit-validation'
import { BillingCallbackBody, BillingCallbackHeaders } from '@/lib/mothership/generated/billing'
import {
BillingCallbackBody,
BillingCallbackHeaders,
UsageUpgrade,
} from '@/lib/mothership/generated/billing'

const booleanQueryParamSchema = z
.preprocess((value) => {
Expand Down Expand Up @@ -315,8 +319,9 @@ export const billingSwitchPlanResponseSchema = z.object({
* limit writes the tag with this payload verbatim into its durable log.
*/
export const usageUpgradePayloadSchema = z.object({
reason: z.literal('usage_limit'),
action: z.enum(['upgrade_plan', 'increase_limit']),
/** `job_run_budget`: a scheduled job's run reached its own spend cap, not the payer's plan. */
reason: UsageUpgrade.shape.reason,
action: UsageUpgrade.shape.action,
message: z.string(),
})
export type UsageUpgradePayload = z.infer<typeof usageUpgradePayloadSchema>
Expand Down
3 changes: 2 additions & 1 deletion apps/sim/lib/mothership/async-runs/orphaned-runs.ts
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ import { getRedisClient } from '@/lib/core/config/redis'
import type { DbTransaction } from '@/lib/db/types'
import { SIM_TOOL_EXECUTION_VERSION } from '@/lib/mothership/async-runs/lifecycle'
import { publishChatStatusChanged } from '@/lib/mothership/chat-status'
import { isDurableStreamRun } from '@/lib/mothership/request/lifecycle/durable-run'
import {
acquirePendingChatStream,
getChatStreamLockOwners,
Expand Down Expand Up @@ -519,7 +520,7 @@ export async function findAbandonedRecoverableRuns(
and(
inArray(copilotRuns.status, UNFINISHED_RUN_STATUSES),
isNotNull(controllerToken),
sql`${copilotRuns.requestContext}->'recovery'->>'kind' = 'interactive_stream'`,
isDurableStreamRun(copilotRuns.requestContext),
idleFor(ABANDONED_RUN_IDLE_MS),
not(idleFor(ORPHANED_RUN_GRACE_MS)),
not(stopRequested),
Expand Down
21 changes: 8 additions & 13 deletions apps/sim/lib/mothership/async-runs/repository.ts
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,6 @@ import {
type SQL,
sql,
} from 'drizzle-orm'
import { z } from 'zod'
import { secretMountPolicySchema } from '@/lib/api/contracts/secret-mount-policy'
import { type ResourceOwner, resourceScopeFromOwner } from '@/lib/core/resource-scope'
import { acquireAdvisoryXactLock } from '@/lib/db/advisory-locks'
import type { DbOrTx, DbTransaction } from '@/lib/db/types'
Expand Down Expand Up @@ -64,6 +62,10 @@ import {
traceMothershipQuery,
traceMothershipTransaction,
} from '@/lib/mothership/observability/database'
import {
isDurableStreamRun,
parseRecordedSecretAuthority,
} from '@/lib/mothership/request/lifecycle/durable-run'
import type { BillingAdmission } from '@/lib/mothership/request/lifecycle/recovery-config'
import { markSpanForError } from '@/lib/mothership/request/otel'
import { sealClientToolSettlement } from '@/lib/mothership/request/tools/client-completion-seal.server'
Expand All @@ -82,7 +84,7 @@ import { chatSandboxSessionKey } from '@/lib/mothership/tools/sandbox-session-ke

const logger = createLogger('CopilotAsyncRunsRepo')
const WORKFLOW_EXECUTION_CLAIM_PREFIX = 'workflow:'
const TERMINAL_RUN_STATUSES: CopilotRunStatus[] = ['complete', 'error', 'cancelled']
export const TERMINAL_RUN_STATUSES: CopilotRunStatus[] = ['complete', 'error', 'cancelled']
// Resolve the tracer lazily per-call to avoid capturing the NoOp tracer
// before NodeSDK installs the global TracerProvider (Next.js 16/Turbopack
// can evaluate modules before instrumentation-node.ts finishes).
Expand Down Expand Up @@ -571,7 +573,7 @@ export async function recordRunBillingAdmission(
eq(copilotRuns.id, runId),
notInArray(copilotRuns.status, TERMINAL_RUN_STATUSES),
sql`${copilotRuns.requestContext}->>'controllerToken' = ${controllerToken}`,
sql`${copilotRuns.requestContext}->'recovery'->>'kind' = 'interactive_stream'`
isDurableStreamRun(copilotRuns.requestContext)
)
)
.returning({ id: copilotRuns.id, status: copilotRuns.status })
Expand All @@ -580,11 +582,6 @@ export async function recordRunBillingAdmission(
)
}

const recordedSecretAuthoritySchema = z.object({
secretMountPolicy: secretMountPolicySchema,
secretActorUserId: z.string().nullable(),
})

/**
* The secret authority a turn was admitted with, for a lane that rebuilds its tool context from
* Go's request instead of inheriting the turn's. Headless runs record it at admission; an
Expand All @@ -610,10 +607,8 @@ export async function getRunSecretAuthority(input: {
.limit(1)
if (!run || run.chatId !== input.chatId) return NO_SECRET_AUTHORITY

if (run.recorded !== null) {
const recorded = recordedSecretAuthoritySchema.safeParse(run.recorded)
return recorded.success ? recorded.data : NO_SECRET_AUTHORITY
}
if (run.recorded !== null)
return parseRecordedSecretAuthority(run.recorded) ?? NO_SECRET_AUTHORITY
if (run.recoveryKind === 'interactive_stream' && run.requestMode !== 'assistant') {
return {
secretMountPolicy: { ...DEFAULT_SECRET_MOUNT_POLICY },
Expand Down
Loading
Loading