From 6d5f2aeb61d4e8f04c7bed30d0f96e3777c0ea6d Mon Sep 17 00:00:00 2001
From: Brad P <brad@ad-astra.video>
Date: Tue, 18 Feb 2025 17:34:45 -0600
Subject: [PATCH] add sessions on hold to /getAISessionPoolsInfo data

---
 server/handlers.go | 33 +++++++++++++++++++++++++++------
 1 file changed, 27 insertions(+), 6 deletions(-)

diff --git a/server/handlers.go b/server/handlers.go
index 578284abef..16819ff533 100644
--- a/server/handlers.go
+++ b/server/handlers.go
@@ -305,9 +305,10 @@ type poolOrchestrator struct {
 
 // aiPoolInfo contains information about an AI pool.
 type aiPoolInfo struct {
-	Size          int                `json:"size"`
-	InUse         int                `json:"in_use"`
-	Orchestrators []poolOrchestrator `json:"orchestrators"`
+	Size          int                           `json:"size"`
+	InUse         int                           `json:"in_use"`
+	Orchestrators []poolOrchestrator            `json:"orchestrators"`
+	OnHold        map[string][]poolOrchestrator `json:"on_hold"`
 }
 
 // suspendedInfo contains information about suspended orchestrators.
@@ -342,8 +343,9 @@ func (s *LivepeerServer) getAIPoolsInfoHandler() http.Handler {
 		// Loop through selectors and get pools info.
 		for cap, pool := range s.AISessionManager.selectors {
 			warmPool := aiPoolInfo{
-				Size:  pool.warmPool.Size(),
-				InUse: len(pool.warmPool.inUseSess),
+				Size:   pool.warmPool.Size(),
+				InUse:  len(pool.warmPool.inUseSess),
+				OnHold: make(map[string][]poolOrchestrator),
 			}
 			for _, sess := range pool.warmPool.sessMap {
 				poolOrchestrator := poolOrchestrator{
@@ -353,6 +355,16 @@ func (s *LivepeerServer) getAIPoolsInfoHandler() http.Handler {
 				}
 				warmPool.Orchestrators = append(warmPool.Orchestrators, poolOrchestrator)
 			}
+			for id, sessions := range pool.warmPool.sessionsOnHold {
+				for _, sess := range sessions {
+					poolOrchestrator := poolOrchestrator{
+						Url:          sess.Transcoder(),
+						LatencyScore: sess.LatencyScore,
+						InFlight:     len(sess.SegsInFlight),
+					}
+					warmPool.OnHold[id] = append(warmPool.OnHold[id], poolOrchestrator)
+				}
+			}
 
 			coldPool := aiPoolInfo{
 				Size:  pool.coldPool.Size(),
@@ -365,7 +377,16 @@ func (s *LivepeerServer) getAIPoolsInfoHandler() http.Handler {
 					InFlight:     len(sess.SegsInFlight),
 				})
 			}
-
+			for id, sessions := range pool.coldPool.sessionsOnHold {
+				for _, sess := range sessions {
+					poolOrchestrator := poolOrchestrator{
+						Url:          sess.Transcoder(),
+						LatencyScore: sess.LatencyScore,
+						InFlight:     len(sess.SegsInFlight),
+					}
+					coldPool.OnHold[id] = append(coldPool.OnHold[id], poolOrchestrator)
+				}
+			}
 			aiPoolsInfoResp[cap] = aiOrchestratorPools{
 				Cold:        coldPool,
 				Warm:        warmPool,