Address more probe comments

Author: imogenkraak

gateway/mw_organisation_activity.go

@@ -394,12 +394,35 @@ func (k *OrganizationMonitor) refreshOrgSession() {
 
 	k.Logger().Debug("Background refresh started for org session")
 
-	session, found := k.OrgSession(k.Spec.OrgID)
-	if found {
-		k.cacheOrgSession(session)
-		k.Logger().Debug("Background refresh completed successfully")
-	} else {
-		k.Logger().Debug("Background refresh failed")
+	// Use timeout to prevent hanging indefinitely
+	timeoutCtx, cancel := context.WithTimeout(context.Background(), orgSessionFetchTimeout)
+	defer cancel()
+
+	type result struct {
+		session user.SessionState
+		found   bool
+	}
+	resultChan := make(chan result, 1)
+
+	go func() {
+		session, found := k.OrgSession(k.Spec.OrgID)
+		select {
+		case resultChan <- result{session: session, found: found}:
+		case <-timeoutCtx.Done():
+			return
+		}
+	}()
+
+	select {
+	case res := <-resultChan:
+		if res.found {
+			k.cacheOrgSession(res.session)
+			k.Logger().Debug("Background refresh completed successfully")
+		} else {
+			k.Logger().Debug("Background refresh failed, org session not found")
+		}
+	case <-timeoutCtx.Done():
+		k.Logger().Warning("Background refresh timed out after 2s")
 	}
 }
 

gateway/mw_organisation_activity_test.go

@@ -539,12 +539,12 @@ func TestOrganizationMonitorStaleWhileRevalidate(t *testing.T) {
 		}
 	})
 
-	t.Run("should respect timeout on cold start", func(t *testing.T) {
+	t.Run("should handle cold start gracefully", func(t *testing.T) {
 		// clear cache to simulate cold start
 		orgSessionCache.Delete(orgID)
 		orgRefreshInProgress.Delete(orgID)
 
-		// use non-existent org to trigger timeout
+		// use non-existent org
 		nonExistentOrgID := "non-existent-org-" + uuid.New()
 
 		spec := ts.Gw.apisByID[ts.Gw.apiSpecs[0].APIID]
@@ -558,7 +558,7 @@ func TestOrganizationMonitorStaleWhileRevalidate(t *testing.T) {
 			},
 		}
 
-		// cold start with non-existent org should time out at 2 seconds
+		// cold start with non-existent org should return quickly (not found)
 		start := time.Now()
 		_, found := monitor.getOrgSessionWithStaleWhileRevalidate()
 		duration := time.Since(start)
@@ -567,11 +567,9 @@ func TestOrganizationMonitorStaleWhileRevalidate(t *testing.T) {
 			t.Error("Should not find non-existent org session")
 		}
 
-		if duration < 1900*time.Millisecond {
-			t.Errorf("Timeout happened too quickly: %v", duration)
-		}
-		if duration > 2500*time.Millisecond {
-			t.Errorf("Timeout took too long: %v", duration)
+		// should return quickly since org doesn't exist (not wait for full timeout)
+		if duration > 500*time.Millisecond {
+			t.Errorf("Should return quickly when org not found, took %v", duration)
 		}
 	})
 }