Skip to content
Merged
Show file tree
Hide file tree
Changes from 6 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 @@ -135,18 +135,22 @@ export function TraceSpans({ traceSpans, totalDuration = 0, onExpansionChange }:
.filter(([, v]) => v)
.map(([k]) => k)
)
const filterTree = (spans: TraceSpan[]): TraceSpan[] =>
const filterTree = (spans: TraceSpan[], parentIsWorkflow = false): TraceSpan[] =>
spans
.map((s) => ({ ...s }))
.map((s) => normalizeChildWorkflowSpan(s))
.filter((s) => {
const tl = s.type?.toLowerCase?.() || ''
if (tl === 'workflow') return true
if (parentIsWorkflow) return true
return allowed.has(tl)
})
.map((s) => ({
...s,
children: s.children ? filterTree(s.children) : undefined,
}))
.map((s) => {
const tl = s.type?.toLowerCase?.() || ''
return {
...s,
children: s.children ? filterTree(s.children, tl === 'workflow') : undefined,
}
})
return traceSpans ? filterTree(traceSpans) : []
}, [traceSpans, effectiveTypeFilters])

Expand Down Expand Up @@ -181,7 +185,6 @@ export function TraceSpans({ traceSpans, totalDuration = 0, onExpansionChange }:
return () => ro.disconnect()
}, [])

// Early return after all hooks are declared to comply with React's Rules of Hooks
if (!traceSpans || traceSpans.length === 0) {
return <div className='text-muted-foreground text-sm'>No trace data available</div>
}
Expand Down Expand Up @@ -217,21 +220,19 @@ export function TraceSpans({ traceSpans, totalDuration = 0, onExpansionChange }:
</div>
<div ref={containerRef} className='relative w-full overflow-hidden border shadow-sm'>
{filtered.map((span, index) => {
const normalizedSpan = normalizeChildWorkflowSpan(span)
const hasSubItems = Boolean(
(normalizedSpan.children && normalizedSpan.children.length > 0) ||
(normalizedSpan.toolCalls && normalizedSpan.toolCalls.length > 0) ||
normalizedSpan.input ||
normalizedSpan.output
(span.children && span.children.length > 0) ||
(span.toolCalls && span.toolCalls.length > 0) ||
span.input ||
span.output
)

// Calculate gap from previous span (for sequential execution visualization)
let gapMs = 0
let gapPercent = 0
if (index > 0) {
const prevSpan = filtered[index - 1]
const prevEndTime = new Date(prevSpan.endTime).getTime()
const currentStartTime = new Date(normalizedSpan.startTime).getTime()
const currentStartTime = new Date(span.startTime).getTime()
gapMs = currentStartTime - prevEndTime
if (gapMs > 0 && actualTotalDuration > 0) {
gapPercent = (gapMs / actualTotalDuration) * 100
Expand All @@ -241,13 +242,13 @@ export function TraceSpans({ traceSpans, totalDuration = 0, onExpansionChange }:
return (
<TraceSpanItem
key={index}
span={normalizedSpan}
span={span}
depth={0}
totalDuration={
actualTotalDuration !== undefined ? actualTotalDuration : totalDuration
}
isLast={index === traceSpans.length - 1}
parentStartTime={new Date(normalizedSpan.startTime).getTime()}
parentStartTime={new Date(span.startTime).getTime()}
workflowStartTime={workflowStartTime}
onToggle={handleSpanToggle}
expandedSpans={expandedSpans}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,239 +1,23 @@
import { useCallback, useEffect, useState } from 'react'
import { createLogger } from '@/lib/logs/console/logger'
import type { UsageData as StoreUsageData } from '@/lib/subscription/types'

const logger = createLogger('useUsageLimits')

/**
* Extended usage data structure that combines API response and store data.
* Supports both 'current' (from store) and 'currentUsage' (from API) for compatibility.
*/
export interface UsageData {
percentUsed: number
isWarning: boolean
isExceeded: boolean
current: number
currentUsage: number
limit: number
}

/**
* Normalizes usage data to ensure both 'current' and 'currentUsage' fields exist
*/
function normalizeUsageData(data: StoreUsageData | any): UsageData {
return {
percentUsed: data.percentUsed,
isWarning: data.isWarning,
isExceeded: data.isExceeded,
current: data.current ?? data.currentUsage ?? 0,
currentUsage: data.currentUsage ?? data.current ?? 0,
limit: data.limit,
}
}
import { useSubscriptionData } from '@/hooks/queries/subscription'

/**
* Cache for usage data to prevent excessive API calls
* Simplified hook that uses React Query for usage limits.
* Provides usage exceeded status from existing subscription data.
*/
let usageDataCache: {
data: UsageData | null
timestamp: number
expirationMs: number
} = {
data: null,
timestamp: 0,
expirationMs: 60 * 1000, // Cache expires after 1 minute
}

/**
* Custom hook to manage usage limits with caching and automatic refresh.
* Provides usage data, exceeded status, and methods to check, refresh, and update limits.
*
* Features:
* - Automatic caching with 60-second expiration
* - Fallback to subscription store if API unavailable
* - Auto-refresh on mount
* - Manual refresh capability
* - Update limit functionality (for user and organization contexts)
* - Integration with usage-limit.tsx component
*
* @param options - Configuration options
* @param options.context - Context for usage check ('user' or 'organization')
* @param options.organizationId - Required when context is 'organization'
* @param options.autoRefresh - Whether to automatically check on mount (default: true)
* @returns Usage state and helper methods
*/
export function useUsageLimits(options?: {
context?: 'user' | 'organization'
organizationId?: string
autoRefresh?: boolean
}) {
Comment thread
icecrasher321 marked this conversation as resolved.
const { context = 'user', organizationId, autoRefresh = true } = options || {}

const [usageData, setUsageData] = useState<UsageData | null>(null)
const [usageExceeded, setUsageExceeded] = useState(false)
const [isLoading, setIsLoading] = useState(false)
const [isUpdating, setIsUpdating] = useState(false)
const [error, setError] = useState<Error | null>(null)

/**
* Check user/organization usage limits with caching
*/
const checkUsage = useCallback(
async (forceRefresh = false): Promise<UsageData | null> => {
const now = Date.now()
const cacheAge = now - usageDataCache.timestamp

// Return cached data if still valid and not forcing refresh
if (!forceRefresh && usageDataCache.data && cacheAge < usageDataCache.expirationMs) {
logger.info('Using cached usage data', {
cacheAge: `${Math.round(cacheAge / 1000)}s`,
})
return usageDataCache.data
}

setIsLoading(true)
setError(null)

try {
// Build query params
const params = new URLSearchParams({ context })
if (context === 'organization' && organizationId) {
params.append('organizationId', organizationId)
}

// Primary: call server-side usage check to mirror backend enforcement
const res = await fetch(`/api/usage?${params.toString()}`, { cache: 'no-store' })
if (res.ok) {
const payload = await res.json()
const usage = normalizeUsageData(payload?.data)

// Update cache
usageDataCache = {
data: usage,
timestamp: now,
expirationMs: usageDataCache.expirationMs,
}

setUsageData(usage)
setUsageExceeded(usage?.isExceeded || false)
return usage
}

// No fallback available - React Query handles this globally
throw new Error('Failed to fetch usage data')
} catch (err) {
const error = err instanceof Error ? err : new Error('Failed to check usage limits')
logger.error('Error checking usage limits:', { error })
setError(error)
return null
} finally {
setIsLoading(false)
}
},
[context, organizationId]
)

/**
* Update usage limit for user or organization
*/
const updateLimit = useCallback(
async (newLimit: number): Promise<{ success: boolean; error?: string }> => {
setIsUpdating(true)
setError(null)

try {
if (context === 'organization') {
if (!organizationId) {
throw new Error('Organization ID is required')
}

const response = await fetch('/api/usage', {
method: 'PUT',
headers: {
'Content-Type': 'application/json',
},
body: JSON.stringify({ context: 'organization', organizationId, limit: newLimit }),
})

const data = await response.json()
if (!response.ok) {
throw new Error(data.error || 'Failed to update limit')
}

// Clear cache and refresh
clearCache()
await checkUsage(true)

return { success: true }
}

// User context - use API directly
const response = await fetch('/api/usage', {
method: 'PUT',
headers: {
'Content-Type': 'application/json',
},
body: JSON.stringify({ context: 'user', limit: newLimit }),
})

const data = await response.json()
if (!response.ok) {
throw new Error(data.error || 'Failed to update limit')
}

// Clear cache and refresh
clearCache()
await checkUsage(true)

return { success: true }
} catch (err) {
const errorMessage = err instanceof Error ? err.message : 'Failed to update usage limit'
logger.error('Failed to update usage limit', { error: err })
setError(err instanceof Error ? err : new Error(errorMessage))
return { success: false, error: errorMessage }
} finally {
setIsUpdating(false)
}
},
[context, organizationId, checkUsage]
)

/**
* Refresh usage data, bypassing cache
*/
const refresh = useCallback(async () => {
return checkUsage(true)
}, [checkUsage])

/**
* Clear the cache (useful for testing or forced refresh)
*/
const clearCache = useCallback(() => {
usageDataCache = {
data: null,
timestamp: 0,
expirationMs: usageDataCache.expirationMs,
}
}, [])
// For now, we only support user context via React Query
// Organization context should use useOrganizationBilling directly
const { data: subscriptionData, isLoading } = useSubscriptionData()

/**
* Auto-refresh on mount if enabled
*/
useEffect(() => {
if (autoRefresh) {
checkUsage()
}
}, [autoRefresh, checkUsage])
const usageExceeded = subscriptionData?.data?.usage?.isExceeded || false
Comment thread
icecrasher321 marked this conversation as resolved.

return {
usageData,
usageExceeded,
isLoading,
isUpdating,
error,
checkUsage,
refresh,
updateLimit,
clearCache,
}
}
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { useCallback, useState } from 'react'
import { useQueryClient } from '@tanstack/react-query'
import { v4 as uuidv4 } from 'uuid'
import { shallow } from 'zustand/shallow'
import { createLogger } from '@/lib/logs/console/logger'
Expand All @@ -11,6 +12,7 @@ import {
} from '@/lib/workflows/trigger-utils'
import { resolveStartCandidates, StartBlockPath, TriggerUtils } from '@/lib/workflows/triggers'
import type { BlockLog, ExecutionResult, StreamingExecution } from '@/executor/types'
import { subscriptionKeys } from '@/hooks/queries/subscription'
import { useExecutionStream } from '@/hooks/use-execution-stream'
import { WorkflowValidationError } from '@/serializer'
import { useExecutionStore } from '@/stores/execution/store'
Expand Down Expand Up @@ -84,6 +86,7 @@ function extractExecutionResult(error: unknown): ExecutionResult | null {
export function useWorkflowExecution() {
const currentWorkflow = useCurrentWorkflow()
const { activeWorkflowId, workflows } = useWorkflowRegistry()
const queryClient = useQueryClient()
const { toggleConsole, addConsole } = useTerminalConsoleStore()
const { getAllVariables } = useEnvironmentStore()
const { getVariablesByWorkflowId, variables } = useVariablesStore()
Expand Down Expand Up @@ -416,7 +419,7 @@ export function useWorkflowExecution() {
if (!streamingExecution.stream) return
const reader = streamingExecution.stream.getReader()
const blockId = (streamingExecution.execution as any)?.blockId
const streamStartTime = Date.now()

let isFirstChunk = true

if (blockId) {
Expand Down Expand Up @@ -577,6 +580,10 @@ export function useWorkflowExecution() {
logger.info(`Processed ${processedCount} blocks for streaming tokenization`)
}

// Invalidate subscription query to update usage
queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() })
queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() })

const { encodeSSE } = await import('@/lib/utils')
controller.enqueue(encodeSSE({ event: 'final', data: result }))
// Note: Logs are already persisted server-side via execution-core.ts
Expand Down Expand Up @@ -639,6 +646,10 @@ export function useWorkflowExecution() {
}
;(result.metadata as any).source = 'chat'
}

// Invalidate subscription query to update usage
queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() })
queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() })
}
return result
} catch (error: any) {
Expand Down
Loading
Loading