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
42 changes: 38 additions & 4 deletions apps/sim/app/api/billing/update-cost/route.test.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
/**
* @vitest-environment node
*/
import { createMockRequest } from '@sim/testing'
import { createMockRequest, dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing'
import { beforeEach, describe, expect, it, vi } from 'vitest'

const {
Expand All @@ -16,6 +16,8 @@ const {
mockCheckAndBillOverageThreshold: vi.fn(),
}))

vi.mock('@sim/db', () => dbChainMock)

vi.mock('@/lib/copilot/request/http', () => ({
checkInternalApiKey: mockCheckInternalApiKey,
}))
Expand Down Expand Up @@ -47,10 +49,12 @@ import { POST } from '@/app/api/billing/update-cost/route'
describe('POST /api/billing/update-cost — workspaceId attribution', () => {
beforeEach(() => {
vi.clearAllMocks()
resetDbChainMock()
mockCheckInternalApiKey.mockReturnValue({ success: true })
mockRecordUsage.mockResolvedValue(undefined)
mockRecordCumulativeUsage.mockResolvedValue({ billed: true, delta: 0.5, total: 0.5 })
mockCheckAndBillOverageThreshold.mockResolvedValue(undefined)
dbChainMockFns.limit.mockResolvedValue([{ id: 'ws-1' }])
})

it('stamps workspaceId onto recorded usage when provided (no idempotency key)', async () => {
Expand Down Expand Up @@ -120,15 +124,45 @@ describe('POST /api/billing/update-cost — workspaceId attribution', () => {
expect(mockCheckAndBillOverageThreshold).not.toHaveBeenCalled()
})

it('rejects with 400 when workspaceId is omitted (contract-required, fail loud)', async () => {
it('records unattributed when workspaceId is omitted (headless client)', async () => {
const res = await POST(
createMockRequest(
'POST',
{ userId: 'user-1', cost: 0.5, model: 'gpt', source: 'copilot' },
{ 'x-api-key': 'internal' }
)
)
expect(res.status).toBe(400)
expect(mockRecordUsage).not.toHaveBeenCalled()
expect(res.status).toBe(200)
expect(dbChainMockFns.limit).not.toHaveBeenCalled()
expect(mockRecordUsage).toHaveBeenCalledTimes(1)
expect(mockRecordUsage.mock.calls[0][0]).toMatchObject({
userId: 'user-1',
workspaceId: undefined,
})
})

it('records unattributed when the workspace does not exist in this deployment (self-hosted client)', async () => {
dbChainMockFns.limit.mockResolvedValue([])
const res = await POST(
createMockRequest(
'POST',
{
userId: 'user-1',
cost: 0.5,
model: 'claude-opus-4.8',
source: 'workspace-chat',
workspaceId: 'self-hosted-ws',
idempotencyKey: 'msg-1-billing',
},
{ 'x-api-key': 'internal' }
)
)
expect(res.status).toBe(200)
expect(mockRecordCumulativeUsage).toHaveBeenCalledTimes(1)
expect(mockRecordCumulativeUsage.mock.calls[0][0]).toMatchObject({
userId: 'user-1',
workspaceId: undefined,
eventKey: 'update-cost:msg-1-billing',
})
})
})
48 changes: 45 additions & 3 deletions apps/sim/app/api/billing/update-cost/route.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
import type { Span } from '@opentelemetry/api'
import { db } from '@sim/db'
import { workspace } from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { getPostgresConstraintName, getPostgresErrorCode, toError } from '@sim/utils/errors'
import { eq } from 'drizzle-orm'
import { type NextRequest, NextResponse } from 'next/server'
import { billingUpdateCostContract } from '@/lib/api/contracts/subscription'
import { parseRequest } from '@/lib/api/server'
Expand All @@ -17,6 +20,35 @@ import { withRouteHandler } from '@/lib/core/utils/with-route-handler'

const logger = createLogger('BillingUpdateCostAPI')

/**
* Resolves the request-supplied workspace to one that exists in this
* deployment. Workspace attribution on the usage ledger is best-effort:
* self-hosted and headless clients bill through this endpoint with workspace
* IDs from their own databases, and `usage_log.workspace_id` carries an FK to
* `workspace`, so stamping a foreign ID would fail the entire flush with an
* FK violation and strand real cost in the caller's dead-letter queue.
* Unknown workspaces are recorded unattributed instead — billing is keyed on
* the user's billing entity and never depends on the workspace.
*/
async function resolveAttributableWorkspaceId(
requestId: string,
workspaceId: string | undefined
): Promise<string | undefined> {
if (!workspaceId) return undefined

const [row] = await db
.select({ id: workspace.id })
.from(workspace)
.where(eq(workspace.id, workspaceId))
.limit(1)
if (row) return row.id

logger.warn(`[${requestId}] Workspace not found in this deployment; recording unattributed`, {
workspaceId,
})
return undefined
Comment thread
icecrasher321 marked this conversation as resolved.
}

/**
* POST /api/billing/update-cost
* Update user cost with a pre-calculated cost value (internal API key auth required)
Expand Down Expand Up @@ -129,6 +161,8 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
source,
})

const attributedWorkspaceId = await resolveAttributableWorkspaceId(requestId, workspaceId)

// Go sends the request's CUMULATIVE cost, possibly more than once (a
// mid-loop provider-error flush, then the recovered terminal flush, plus
// abort-race duplicates). Record it as a monotonic top-up: one ledger row
Expand All @@ -141,7 +175,7 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
if (idempotencyKey) {
const result = await recordCumulativeUsage({
userId,
workspaceId,
workspaceId: attributedWorkspaceId,
source,
model,
cost,
Expand All @@ -160,7 +194,7 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
} else {
await recordUsage({
userId,
workspaceId,
workspaceId: attributedWorkspaceId,
entries: [
{
category: 'model',
Expand Down Expand Up @@ -229,8 +263,16 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
} catch (error) {
const duration = Date.now() - startTime

// Surface the underlying Postgres failure (e.g. 23503 FK violation vs a
// lock timeout) — Drizzle's "Failed query" wrapper alone cannot
// distinguish them, which made the dead-workspace incident undiagnosable
// from logs.
const pgCode = getPostgresErrorCode(error)
const pgConstraint = getPostgresConstraintName(error)
logger.error(`[${requestId}] Cost update failed`, {
error: toError(error).message,
...(pgCode && { pgCode }),
...(pgConstraint && { pgConstraint }),
stack: error instanceof Error ? error.stack : undefined,
duration,
})
Expand Down
3 changes: 2 additions & 1 deletion apps/sim/app/api/logs/export/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import { and, desc, eq, sql } from 'drizzle-orm'
import { type NextRequest, NextResponse } from 'next/server'
import { getSession } from '@/lib/auth'
import { MATERIALIZE_CONCURRENCY, mapWithConcurrency } from '@/lib/core/utils/concurrency'
import { neutralizeCsvFormula } from '@/lib/core/utils/csv'
import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
import { materializeExecutionData } from '@/lib/logs/execution/trace-store'
import { buildFilterConditions, LogFilterParamsSchema } from '@/lib/logs/filters'
Expand All @@ -16,7 +17,7 @@ export const revalidate = 0

function escapeCsv(value: any): string {
if (value === null || value === undefined) return ''
const str = String(value)
const str = typeof value === 'string' ? neutralizeCsvFormula(value) : String(value)
if (/[",\n]/.test(str)) {
return `"${str.replace(/"/g, '""')}"`
}
Expand Down
9 changes: 1 addition & 8 deletions apps/sim/app/api/table/[tableId]/export/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ import { type NextRequest, NextResponse } from 'next/server'
import { tableExportFormatSchema, tableIdParamsSchema } from '@/lib/api/contracts/tables'
import { getValidationErrorMessage } from '@/lib/api/server'
import { checkSessionOrInternalAuth } from '@/lib/auth/hybrid'
import { neutralizeCsvFormula } from '@/lib/core/utils/csv'
import { generateRequestId } from '@/lib/core/utils/request'
import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
import { buildNameById, getColumnId, rowDataIdToName } from '@/lib/table/column-keys'
Expand Down Expand Up @@ -119,14 +120,6 @@ function sanitizeFilename(name: string): string {
return cleaned || 'table'
}

/**
* Prefixes a single quote to values starting with a spreadsheet formula trigger
* (`=`, `+`, `-`, `@`, tab, CR), neutralizing CSV injection in Excel/Sheets.
*/
function neutralizeCsvFormula(value: string): string {
return /^[=+\-@\t\r]/.test(value) ? `'${value}` : value
}

/**
* Serializes a cell for CSV. Only string cells are formula-neutralized; numbers,
* booleans, dates, and JSON objects can never form a trigger and pass through verbatim.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ import {
toast,
} from '@/components/emcn'
import { cn } from '@/lib/core/utils/cn'
import { useChatSurface } from '@/app/workspace/[workspaceId]/home/components/chat-surface-context'
import { useSubmitCopilotFeedback } from '@/hooks/queries/copilot-feedback'
import { useForkMothershipChat } from '@/hooks/queries/mothership-chats'
import { useFolderStore } from '@/stores/folders/store'
Expand Down Expand Up @@ -49,21 +50,20 @@ const BUTTON_CLASS =

interface MessageActionsProps {
content: string
chatId?: string
userQuery?: string
requestId?: string
messageId?: string
}

export const MessageActions = memo(function MessageActions({
content,
chatId,
userQuery,
requestId,
messageId,
}: MessageActionsProps) {
const router = useRouter()
const params = useParams<{ workspaceId: string }>()
const { chatId } = useChatSurface()
const [copied, setCopied] = useState(false)
const [copiedRequestId, setCopiedRequestId] = useState(false)
const [pendingFeedback, setPendingFeedback] = useState<'up' | 'down' | null>(null)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
'use client'

import {
createContext,
type ReactNode,
useCallback,
useContext,
useLayoutEffect,
useMemo,
useRef,
} from 'react'
import type { MothershipResource } from '@/app/workspace/[workspaceId]/home/types'
import type { ChatContext } from '@/stores/panel'

/**
* Identity and interaction callbacks shared across a Mothership chat surface
* (home conversation view, home initial view, copilot panel). Carried via
* context so leaf components (UserInput, MessageContent, MessageActions) can
* consume them without relaying through every intermediate component.
*/
interface ChatSurfaceContextValue {
/** Resolved id of the chat backing this surface, if one exists yet. */
chatId?: string
/** Id of the user interacting with this surface. */
userId?: string
/** Notifies the surface owner that a context chip was added to the input. */
onContextAdd: (context: ChatContext) => void
/** Notifies the surface owner that a context chip was removed from the input. */
onContextRemove: (context: ChatContext) => void
/** Opens a workspace resource referenced from rendered message content. */
onWorkspaceResourceSelect: (resource: MothershipResource) => void
}

const noop = () => {}

const ChatSurfaceContext = createContext<ChatSurfaceContextValue>({
onContextAdd: noop,
onContextRemove: noop,
onWorkspaceResourceSelect: noop,
})

interface ChatSurfaceProviderProps {
chatId?: string
userId?: string
onContextAdd?: (context: ChatContext) => void
onContextRemove?: (context: ChatContext) => void
onWorkspaceResourceSelect?: (resource: MothershipResource) => void
children: ReactNode
}

/**
* Provides the chat-surface identity and interaction callbacks to descendants.
* Callbacks are latched in refs and exposed as stable wrappers so the memoized
* context value only changes when `chatId` or `userId` change — consumers do
* not re-render when a parent re-creates a handler.
*/
export function ChatSurfaceProvider({
chatId,
userId,
onContextAdd,
onContextRemove,
onWorkspaceResourceSelect,
children,
}: ChatSurfaceProviderProps) {
const onContextAddRef = useRef(onContextAdd)
const onContextRemoveRef = useRef(onContextRemove)
const onWorkspaceResourceSelectRef = useRef(onWorkspaceResourceSelect)

useLayoutEffect(() => {
onContextAddRef.current = onContextAdd
onContextRemoveRef.current = onContextRemove
onWorkspaceResourceSelectRef.current = onWorkspaceResourceSelect
})

const stableOnContextAdd = useCallback((context: ChatContext) => {
onContextAddRef.current?.(context)
}, [])
const stableOnContextRemove = useCallback((context: ChatContext) => {
onContextRemoveRef.current?.(context)
}, [])
const stableOnWorkspaceResourceSelect = useCallback((resource: MothershipResource) => {
onWorkspaceResourceSelectRef.current?.(resource)
}, [])

const value = useMemo<ChatSurfaceContextValue>(
() => ({
chatId,
userId,
onContextAdd: stableOnContextAdd,
onContextRemove: stableOnContextRemove,
onWorkspaceResourceSelect: stableOnWorkspaceResourceSelect,
}),
[chatId, userId, stableOnContextAdd, stableOnContextRemove, stableOnWorkspaceResourceSelect]
)

return <ChatSurfaceContext.Provider value={value}>{children}</ChatSurfaceContext.Provider>
}

/**
* Reads the surrounding chat surface. Outside a provider this returns no-op
* callbacks and undefined identity, matching the previous optional-prop
* behavior.
*/
export function useChatSurface(): ChatSurfaceContextValue {
return useContext(ChatSurfaceContext)
}
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
export { ChatSurfaceProvider, useChatSurface } from './chat-surface-context'
5 changes: 5 additions & 0 deletions apps/sim/app/workspace/[workspaceId]/home/components/index.ts
Original file line number Diff line number Diff line change
@@ -1,11 +1,16 @@
export { ChatMessageAttachments } from './chat-message-attachments'
export { ChatSurfaceProvider, useChatSurface } from './chat-surface-context'
export { ContextMentionIcon } from './context-mention-icon'
export { CreditsChip } from './credits-chip'
export {
assistantMessageHasRenderableContent,
MessageContent,
} from './message-content'
export { MothershipChat } from './mothership-chat'
export {
MothershipResourcesProvider,
useMothershipResources,
} from './mothership-resources-context'
export { MothershipView } from './mothership-view'
export { QueuedMessages } from './queued-messages'
export { SuggestedActions } from './suggested-actions'
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,8 @@ import { Read as ReadTool, WorkspaceFile } from '@/lib/copilot/generated/tool-ca
import { isToolHiddenInUi } from '@/lib/copilot/tools/client/hidden-tools'
import { resolveToolDisplay } from '@/lib/copilot/tools/client/store-utils'
import { ClientToolCallState } from '@/lib/copilot/tools/client/tool-call-state'
import type { ContentBlock, MothershipResource, OptionItem, ToolCallData } from '../../types'
import { useChatSurface } from '@/app/workspace/[workspaceId]/home/components/chat-surface-context'
import type { ContentBlock, OptionItem, ToolCallData } from '../../types'
import { SUBAGENT_LABELS, TOOL_UI_METADATA } from '../../types'
import type { AgentGroupItem } from './components'
import {
Expand Down Expand Up @@ -676,16 +677,15 @@ interface MessageContentProps {
fallbackContent: string
isStreaming: boolean
onOptionSelect?: (id: string) => void
onWorkspaceResourceSelect?: (resource: MothershipResource) => void
}

function MessageContentInner({
blocks,
fallbackContent,
isStreaming = false,
onOptionSelect,
onWorkspaceResourceSelect,
}: MessageContentProps) {
const { onWorkspaceResourceSelect } = useChatSurface()
const parsed = useMemo(() => (blocks.length > 0 ? parseBlocks(blocks) : []), [blocks])

const segments: MessageSegment[] =
Expand Down
Loading
Loading