implement 'manages_sender_collection' in generate_tokens_sender_collection

Author: mcharytoniuk

src/balancer/agent_controller.rs

@@ -24,6 +24,7 @@ use crate::balancer::agent_controller_update_result::AgentControllerUpdateResult
 use crate::balancer::chat_template_override_sender_collection::ChatTemplateOverrideSenderCollection;
 use crate::balancer::embedding_sender_collection::EmbeddingSenderCollection;
 use crate::balancer::generate_tokens_sender_collection::GenerateTokensSenderCollection;
+use crate::balancer::manages_senders::ManagesSenders as _;
 use crate::balancer::manages_senders_controller::ManagesSendersController;
 use crate::balancer::model_metadata_sender_collection::ModelMetadataSenderCollection;
 use crate::balancer::receive_tokens_controller::ReceiveTokensController;

src/balancer/generate_tokens_sender_collection.rs

@@ -1,56 +1,27 @@
-use anyhow::anyhow;
-use anyhow::Result;
+use async_trait::async_trait;
 use dashmap::DashMap;
 use tokio::sync::mpsc;
 
+use crate::balancer::manages_senders::ManagesSenders;
 use crate::generated_token_envelope::GeneratedTokenEnvelope;
 
 pub struct GenerateTokensSenderCollection {
-    generate_tokens_senders: DashMap<String, mpsc::UnboundedSender<GeneratedTokenEnvelope>>,
+    senders: DashMap<String, mpsc::UnboundedSender<GeneratedTokenEnvelope>>,
 }
 
 impl GenerateTokensSenderCollection {
     pub fn new() -> Self {
         Self {
-            generate_tokens_senders: DashMap::new(),
+            senders: DashMap::new(),
         }
     }
+}
 
-    pub fn deregister_sender(&self, request_id: String) -> Result<()> {
-        if let Some(sender) = self.generate_tokens_senders.remove(&request_id) {
-            drop(sender);
-
-            Ok(())
-        } else {
-            Err(anyhow!("No sender found for request_id {request_id}"))
-        }
-    }
-
-    pub async fn forward_generated_token(
-        &self,
-        request_id: String,
-        generated_token_envelope: GeneratedTokenEnvelope,
-    ) -> Result<()> {
-        if let Some(sender) = self.generate_tokens_senders.get(&request_id) {
-            sender.send(generated_token_envelope)?;
-
-            Ok(())
-        } else {
-            Err(anyhow!("No sender found for request_id {request_id}"))
-        }
-    }
-
-    pub fn register_sender(
-        &self,
-        request_id: String,
-        sender: mpsc::UnboundedSender<GeneratedTokenEnvelope>,
-    ) -> Result<()> {
-        if self.generate_tokens_senders.contains_key(&request_id) {
-            return Err(anyhow!("Sender for request_id {request_id} already exists"));
-        }
-
-        self.generate_tokens_senders.insert(request_id, sender);
+#[async_trait]
+impl ManagesSenders for GenerateTokensSenderCollection {
+    type Value = GeneratedTokenEnvelope;
 
-        Ok(())
+    fn get_sender_collection(&self) -> &DashMap<String, mpsc::UnboundedSender<Self::Value>> {
+        &self.senders
     }
 }

src/balancer/management_service/http_route/api/ws_agent_socket/mod.rs

@@ -22,7 +22,6 @@ use anyhow::Result;
 use async_trait::async_trait;
 use log::error;
 use log::info;
-use log::warn;
 use serde::Deserialize;
 use tokio::sync::broadcast;
 use tokio::sync::mpsc;
@@ -260,14 +259,10 @@ impl ControlsWebSocketEndpoint for AgentSocketController {
                 request_id,
                 response: AgentJsonRpcResponse::GeneratedToken(generated_token_envelope),
             }) => {
-                if let Err(err) = context
+                context
                     .generate_tokens_sender_collection
-                    .forward_generated_token(request_id, generated_token_envelope)
-                    .await
-                {
-                    // Token might come in after the sender was deregistered
-                    warn!("Error forwarding generated token: {err}");
-                }
+                    .forward_response_safe(request_id, generated_token_envelope)
+                    .await;
 
                 Ok(ContinuationDecision::Continue)
             }

src/balancer/receive_tokens_controller.rs

@@ -4,6 +4,7 @@ use log::error;
 use tokio::sync::mpsc;
 
 use crate::balancer::generate_tokens_sender_collection::GenerateTokensSenderCollection;
+use crate::balancer::manages_senders::ManagesSenders as _;
 use crate::generated_token_envelope::GeneratedTokenEnvelope;
 
 pub struct ReceiveTokensController {