improvement(monitoring): add SSE metering to wand, execution-stream, and a2a-message endpoints

Author: waleedlatif1

apps/sim/app/api/a2a/serve/[agentId]/route.ts

@@ -19,7 +19,10 @@ import { validateUrlWithDNS } from '@/lib/core/security/input-validation.server'
 import { SSE_HEADERS } from '@/lib/core/utils/sse'
 import { getBaseUrl } from '@/lib/core/utils/urls'
 import { markExecutionCancelled } from '@/lib/execution/cancellation'
-import { decrementSSEConnections, incrementSSEConnections } from '@/lib/monitoring/sse-connections'
+import {
+  decrementSSEConnections,
+  incrementSSEConnections,
+} from '@/lib/monitoring/sse-connections'
 import { checkWorkspaceAccess } from '@/lib/workspaces/permissions/utils'
 import { getWorkspaceBilledAccountUserId } from '@/lib/workspaces/utils'
 import {
@@ -631,9 +634,11 @@ async function handleMessageStream(
   }
 
   const encoder = new TextEncoder()
+  let messageStreamDecremented = false
 
   const stream = new ReadableStream({
     async start(controller) {
+      incrementSSEConnections('a2a-message')
       const sendEvent = (event: string, data: unknown) => {
         try {
           const jsonRpcResponse = {
@@ -842,9 +847,19 @@ async function handleMessageStream(
         })
       } finally {
         await releaseLock(lockKey, lockValue)
+        if (!messageStreamDecremented) {
+          messageStreamDecremented = true
+          decrementSSEConnections('a2a-message')
+        }
         controller.close()
       }
     },
+    cancel() {
+      if (!messageStreamDecremented) {
+        messageStreamDecremented = true
+        decrementSSEConnections('a2a-message')
+      }
+    },
   })
 
   return new NextResponse(stream, {

apps/sim/app/api/wand/route.ts

@@ -10,6 +10,7 @@ import { checkAndBillOverageThreshold } from '@/lib/billing/threshold-billing'
 import { env } from '@/lib/core/config/env'
 import { getCostMultiplier, isBillingEnabled } from '@/lib/core/config/feature-flags'
 import { generateRequestId } from '@/lib/core/utils/request'
+import { decrementSSEConnections, incrementSSEConnections } from '@/lib/monitoring/sse-connections'
 import { enrichTableSchema } from '@/lib/table/llm/wand'
 import { verifyWorkspaceMembership } from '@/app/api/workflows/utils'
 import { extractResponseText, parseResponsesUsage } from '@/providers/openai/utils'
@@ -330,10 +331,13 @@ export async function POST(req: NextRequest) {
         const encoder = new TextEncoder()
         const decoder = new TextDecoder()
 
+        let wandStreamClosed = false
         const readable = new ReadableStream({
           async start(controller) {
+            incrementSSEConnections('wand')
             const reader = response.body?.getReader()
             if (!reader) {
+              decrementSSEConnections('wand')
               controller.close()
               return
             }
@@ -478,6 +482,16 @@ export async function POST(req: NextRequest) {
               controller.close()
             } finally {
               reader.releaseLock()
+              if (!wandStreamClosed) {
+                wandStreamClosed = true
+                decrementSSEConnections('wand')
+              }
+            }
+          },
+          cancel() {
+            if (!wandStreamClosed) {
+              wandStreamClosed = true
+              decrementSSEConnections('wand')
             }
           },
         })

apps/sim/app/api/workflows/[id]/executions/[executionId]/stream/route.ts

@@ -7,6 +7,7 @@ import {
   getExecutionMeta,
   readExecutionEvents,
 } from '@/lib/execution/event-buffer'
+import { decrementSSEConnections, incrementSSEConnections } from '@/lib/monitoring/sse-connections'
 import { formatSSEEvent } from '@/lib/workflows/executor/execution-events'
 import { authorizeWorkflowByWorkspacePermission } from '@/lib/workflows/utils'
 
@@ -73,8 +74,10 @@ export async function GET(
 
     let closed = false
 
+    let sseDecremented = false
     const stream = new ReadableStream<Uint8Array>({
       async start(controller) {
+        incrementSSEConnections('execution-stream-reconnect')
         let lastEventId = fromEventId
         const pollDeadline = Date.now() + MAX_POLL_DURATION_MS
 
@@ -142,11 +145,20 @@ export async function GET(
               controller.close()
             } catch {}
           }
+        } finally {
+          if (!sseDecremented) {
+            sseDecremented = true
+            decrementSSEConnections('execution-stream-reconnect')
+          }
         }
       },
       cancel() {
         closed = true
         logger.info('Client disconnected from reconnection stream', { executionId })
+        if (!sseDecremented) {
+          sseDecremented = true
+          decrementSSEConnections('execution-stream-reconnect')
+        }
       },
     })