Add notifications via websocket.

This commit is contained in:
Robin Huang
2025-06-30 13:03:24 -07:00
committed by Jennifer Weber
parent 23e881e220
commit c27edb7e94
3 changed files with 137 additions and 1 deletions

View File

@@ -5,6 +5,7 @@ import { LGraphNode } from '@/lib/litegraph/src/litegraph'
import { IWidget } from '@/lib/litegraph/src/litegraph'
import type { RemoteWidgetConfig } from '@/schemas/nodeDefSchema'
import { api } from '@/scripts/api'
import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
const MAX_RETRIES = 5
const TIMEOUT = 4096
@@ -58,10 +59,21 @@ const fetchData = async (
controller: AbortController
) => {
const { route, response_key, query_params, timeout = TIMEOUT } = config
// Get auth header from Firebase
const authStore = useFirebaseAuthStore()
const authHeader = await authStore.getAuthHeader()
const headers: Record<string, string> = {}
if (authHeader) {
Object.assign(headers, authHeader)
}
const res = await axios.get(route, {
params: query_params,
signal: controller.signal,
timeout
timeout,
headers
})
return response_key ? res.data[response_key] : res.data
}

View File

@@ -112,6 +112,11 @@ const zDisplayComponentWsMessage = z.object({
props: z.record(z.string(), z.any()).optional()
})
const zNotificationWsMessage = z.object({
value: z.string(),
id: z.string().optional()
})
const zTerminalSize = z.object({
cols: z.number(),
row: z.number()
@@ -153,6 +158,7 @@ export type DisplayComponentWsMessage = z.infer<
export type NodeProgressState = z.infer<typeof zNodeProgressState>
export type ProgressStateWsMessage = z.infer<typeof zProgressStateWsMessage>
export type FeatureFlagsWsMessage = z.infer<typeof zFeatureFlagsWsMessage>
export type NotificationWsMessage = z.infer<typeof zNotificationWsMessage>
// End of ws messages
const zPromptInputItem = z.object({

View File

@@ -1,4 +1,5 @@
import axios from 'axios'
import { debounce } from 'lodash'
import defaultClientFeatureFlags from '@/config/clientFeatureFlags.json'
import type {
@@ -16,6 +17,7 @@ import type {
HistoryTaskItem,
LogsRawResponse,
LogsWsMessage,
NotificationWsMessage,
PendingTaskItem,
ProgressStateWsMessage,
ProgressTextWsMessage,
@@ -37,6 +39,7 @@ import type {
import type { ComfyNodeDef } from '@/schemas/nodeDefSchema'
import type { NodeExecutionId } from '@/types/nodeIdentification'
import { useFirebaseAuthStore } from '@/stores/firebaseAuthStore'
import { useToastStore } from '@/stores/toastStore'
import { WorkflowTemplates } from '@/types/workflowTemplateTypes'
interface QueuePromptRequestBody {
@@ -131,6 +134,7 @@ interface BackendApiCalls {
progress_state: ProgressStateWsMessage
display_component: DisplayComponentWsMessage
feature_flags: FeatureFlagsWsMessage
notification: NotificationWsMessage
}
/** Dictionary of all api calls */
@@ -272,6 +276,81 @@ export class ComfyApi extends EventTarget {
* Feature flags received from the backend server.
*/
serverFeatureFlags: Record<string, unknown> = {}
/**
* Map of notification toasts by ID
*/
#notificationToasts = new Map<string, any>()
/**
* Map of timers for auto-hiding notifications by ID
*/
#notificationTimers = new Map<string, number>()
/**
* Handle notification messages (with optional ID for multiple parallel notifications)
*/
#handleNotification(value: string, id?: string) {
try {
const toastStore = useToastStore()
const notificationId = id || 'default'
console.log(`Updating notification (${notificationId}):`, value)
// Get existing toast for this ID
const existingToast = this.#notificationToasts.get(notificationId)
if (existingToast) {
// Update existing toast by removing and re-adding with new content
console.log(`Updating existing notification toast: ${notificationId}`)
toastStore.remove(existingToast)
// Update the detail text
existingToast.detail = value
toastStore.add(existingToast)
} else {
// Create new persistent notification toast
console.log(`Creating new notification toast: ${notificationId}`)
const newToast = {
severity: 'info' as const,
summary: 'Notification',
detail: value,
closable: true
// No 'life' property means it won't auto-hide
}
this.#notificationToasts.set(notificationId, newToast)
toastStore.add(newToast)
}
// Clear existing timer for this ID and set new one
const existingTimer = this.#notificationTimers.get(notificationId)
if (existingTimer) {
clearTimeout(existingTimer)
}
const timer = window.setTimeout(() => {
const toast = this.#notificationToasts.get(notificationId)
if (toast) {
console.log(`Auto-hiding notification toast: ${notificationId}`)
toastStore.remove(toast)
this.#notificationToasts.delete(notificationId)
this.#notificationTimers.delete(notificationId)
}
}, 3000)
this.#notificationTimers.set(notificationId, timer)
console.log('Toast updated successfully')
} catch (error) {
console.error('Error handling notification:', error)
}
}
/**
* Debounced notification handler to avoid rapid toast updates
*/
#debouncedNotificationHandler = debounce((value: string, id?: string) => {
this.#handleNotification(value, id)
}, 300) // 300ms debounce delay
/**
* The auth token for the comfy org account if the user is logged in.
@@ -596,6 +675,16 @@ export class ComfyApi extends EventTarget {
this.serverFeatureFlags
)
break
case 'notification':
// Display notification in toast with debouncing
console.log(
'Received notification message:',
msg.data.value,
msg.data.id ? `(ID: ${msg.data.id})` : ''
)
this.#debouncedNotificationHandler(msg.data.value, msg.data.id)
this.dispatchCustomEvent(msg.type, msg.data)
break
default:
if (this.#registered.has(msg.type)) {
// Fallback for custom types - calls super direct.
@@ -621,6 +710,35 @@ export class ComfyApi extends EventTarget {
this.#createSocket()
}
/**
* Test method to simulate a notification message (for development/testing)
*/
testNotification(message: string = 'Test notification message', id?: string) {
console.log(
'Testing notification with message:',
message,
id ? `(ID: ${id})` : ''
)
const mockEvent = {
data: JSON.stringify({
type: 'notification',
data: { value: message, id }
})
}
// Simulate the websocket message handler
const msg = JSON.parse(mockEvent.data)
if (msg.type === 'notification') {
console.log(
'Received notification message:',
msg.data.value,
msg.data.id ? `(ID: ${msg.data.id})` : ''
)
this.#debouncedNotificationHandler(msg.data.value, msg.data.id)
this.dispatchCustomEvent(msg.type, msg.data)
}
}
/**
* Gets a list of extension urls
*/