alien-sdk 2.1.1

use alien_bindings::error::{ErrorData, Result};
use alien_error::{AlienError, Context, IntoAlienError};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::{
    collections::HashMap,
    future::Future,
    sync::{
        atomic::{AtomicU32, Ordering},
        Arc,
    },
    task::{Context as TaskContext, Poll},
    time::Duration,
};
use tokio::{sync::Mutex, task::JoinHandle, time::timeout};
use tonic::{
    codegen::{
        http::{Request as HttpRequest, Response as HttpResponse, Uri},
        Service,
    },
    transport::Channel,
};
use tracing::{debug, error, info, warn};
use uuid::Uuid;

#[cfg(feature = "openapi")]
use utoipa::ToSchema;

use alien_worker_protocol::wait_until::{
    wait_until_service_client::WaitUntilServiceClient, NotifyDrainCompleteRequest,
    NotifyDrainCompleteResponse, NotifyTaskRegisteredRequest, NotifyTaskRegisteredResponse,
    WaitForDrainSignalRequest, WaitForDrainSignalResponse,
};

pub(crate) const ENV_ALIEN_BINDINGS_GRPC_ADDRESS: &str = "ALIEN_BINDINGS_GRPC_ADDRESS";

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum WorkerProtocolGeneration {
    Current,
    Legacy,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) struct WorkerProtocolEndpoint<'a> {
    pub address: &'a str,
    pub generation: WorkerProtocolGeneration,
}

pub(crate) fn worker_protocol_endpoint<'a>(
    env_vars: &'a HashMap<String, String>,
) -> Option<WorkerProtocolEndpoint<'a>> {
    env_vars
        .get(alien_core::ENV_ALIEN_WORKER_GRPC_ADDRESS)
        .map(|address| WorkerProtocolEndpoint {
            address,
            generation: WorkerProtocolGeneration::Current,
        })
        .or_else(|| {
            env_vars
                .get(ENV_ALIEN_BINDINGS_GRPC_ADDRESS)
                .map(|address| WorkerProtocolEndpoint {
                    address,
                    generation: WorkerProtocolGeneration::Legacy,
                })
        })
}

#[derive(Clone, Debug)]
pub(crate) struct WorkerProtocolChannel {
    inner: Channel,
    generation: WorkerProtocolGeneration,
}

impl WorkerProtocolChannel {
    pub fn new(inner: Channel, generation: WorkerProtocolGeneration) -> Self {
        Self { inner, generation }
    }
}

impl Service<HttpRequest<tonic::body::Body>> for WorkerProtocolChannel {
    type Response = HttpResponse<tonic::body::Body>;
    type Error = tonic::transport::Error;
    type Future = <Channel as Service<HttpRequest<tonic::body::Body>>>::Future;

    fn poll_ready(
        &mut self,
        context: &mut TaskContext<'_>,
    ) -> Poll<std::result::Result<(), Self::Error>> {
        self.inner.poll_ready(context)
    }

    fn call(&mut self, mut request: HttpRequest<tonic::body::Body>) -> Self::Future {
        if self.generation == WorkerProtocolGeneration::Legacy {
            if let Some(uri) = legacy_worker_protocol_uri(request.uri().path()) {
                *request.uri_mut() = uri;
            }
        }
        self.inner.call(request)
    }
}

fn legacy_worker_protocol_uri(path: &str) -> Option<Uri> {
    match path {
        "/alien_worker.control.ControlService/RegisterHttpServer" => Some(Uri::from_static(
            "/alien_bindings.control.ControlService/RegisterHttpServer",
        )),
        "/alien_worker.control.ControlService/RegisterEventHandler" => Some(Uri::from_static(
            "/alien_bindings.control.ControlService/RegisterEventHandler",
        )),
        "/alien_worker.control.ControlService/WaitForTasks" => Some(Uri::from_static(
            "/alien_bindings.control.ControlService/WaitForTasks",
        )),
        "/alien_worker.control.ControlService/SendTaskResult" => Some(Uri::from_static(
            "/alien_bindings.control.ControlService/SendTaskResult",
        )),
        "/alien_worker.wait_until.WaitUntilService/NotifyTaskRegistered" => Some(Uri::from_static(
            "/alien_bindings.wait_until.WaitUntilService/NotifyTaskRegistered",
        )),
        "/alien_worker.wait_until.WaitUntilService/WaitForDrainSignal" => Some(Uri::from_static(
            "/alien_bindings.wait_until.WaitUntilService/WaitForDrainSignal",
        )),
        "/alien_worker.wait_until.WaitUntilService/NotifyDrainComplete" => Some(Uri::from_static(
            "/alien_bindings.wait_until.WaitUntilService/NotifyDrainComplete",
        )),
        "/alien_worker.wait_until.WaitUntilService/GetTaskCount" => Some(Uri::from_static(
            "/alien_bindings.wait_until.WaitUntilService/GetTaskCount",
        )),
        _ => None,
    }
}

#[derive(Clone, Debug)]
enum WaitUntilGrpcClient {
    Direct(WaitUntilServiceClient<Channel>),
    Runtime(WaitUntilServiceClient<WorkerProtocolChannel>),
}

impl WaitUntilGrpcClient {
    async fn notify_task_registered(
        &mut self,
        request: NotifyTaskRegisteredRequest,
    ) -> std::result::Result<tonic::Response<NotifyTaskRegisteredResponse>, tonic::Status> {
        match self {
            Self::Direct(client) => client.notify_task_registered(request).await,
            Self::Runtime(client) => client.notify_task_registered(request).await,
        }
    }

    async fn wait_for_drain_signal(
        &mut self,
        request: WaitForDrainSignalRequest,
    ) -> std::result::Result<tonic::Response<WaitForDrainSignalResponse>, tonic::Status> {
        match self {
            Self::Direct(client) => client.wait_for_drain_signal(request).await,
            Self::Runtime(client) => client.wait_for_drain_signal(request).await,
        }
    }

    async fn notify_drain_complete(
        &mut self,
        request: NotifyDrainCompleteRequest,
    ) -> std::result::Result<tonic::Response<NotifyDrainCompleteResponse>, tonic::Status> {
        match self {
            Self::Direct(client) => client.notify_drain_complete(request).await,
            Self::Runtime(client) => client.notify_drain_complete(request).await,
        }
    }
}

/// Response from drain operations.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "openapi", derive(ToSchema))]
pub struct DrainResponse {
    /// Number of tasks that were drained.
    pub tasks_drained: u32,
    /// Whether all tasks completed successfully.
    pub success: bool,
    /// Optional error message if draining failed.
    pub error_message: Option<String>,
}

/// Configuration for wait_until drain behavior.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "openapi", derive(ToSchema))]
pub struct DrainConfig {
    /// Maximum time to wait for all tasks to complete.
    pub timeout: Duration,
    /// Reason for the drain request.
    pub reason: String,
}

/// Worker task coordination and drain operations.
///
/// This trait is not object-safe due to its generic task method; callers use
/// [`WaitUntilContext`] directly.
#[async_trait]
pub trait WaitUntil: Send + Sync + std::fmt::Debug {
    /// Waits for a drain signal from the Worker runtime.
    /// This is a blocking call that returns when the Worker runtime decides it's time to drain.
    async fn wait_for_drain_signal(&self, timeout: Option<Duration>) -> Result<DrainConfig>;

    /// Drains all currently registered tasks.
    /// This waits for all tasks to complete or timeout.
    async fn drain_all(&self, config: DrainConfig) -> Result<DrainResponse>;

    /// Gets the current number of registered tasks.
    async fn get_task_count(&self) -> Result<u32>;

    /// Notifies the Worker runtime that draining is complete.
    async fn notify_drain_complete(&self, response: DrainResponse) -> Result<()>;
}

/// A context for managing wait_until tasks within an application.
/// This handles local task execution and coordinates with the Worker runtime via gRPC.
#[derive(Debug)]
pub struct WaitUntilContext {
    /// Unique identifier for this application instance.
    application_id: String,
    /// Currently running tasks.
    tasks: Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
    /// Task counter for generating unique task IDs.
    task_counter: AtomicU32,
    /// gRPC client for communicating with the Worker runtime.
    grpc_client: Option<WaitUntilGrpcClient>,
    /// Whether we're currently draining tasks.
    draining: Arc<Mutex<bool>>,
}

impl WaitUntilContext {
    /// Creates a new WaitUntilContext.
    pub fn new(application_id: Option<String>) -> Self {
        let app_id = application_id.unwrap_or_else(|| Uuid::new_v4().to_string());

        Self {
            application_id: app_id,
            tasks: Arc::new(Mutex::new(HashMap::new())),
            task_counter: AtomicU32::new(0),
            grpc_client: None,
            draining: Arc::new(Mutex::new(false)),
        }
    }

    /// Creates a new WaitUntilContext and connects to gRPC endpoint from environment variables.
    /// This is the recommended way to create a WaitUntilContext in production.
    pub async fn from_env(application_id: Option<String>) -> Result<Self> {
        let env_vars: std::collections::HashMap<String, String> = std::env::vars().collect();
        Self::from_env_with_vars(application_id, &env_vars).await
    }

    /// Creates a new WaitUntilContext and connects to gRPC endpoint from provided environment variables.
    pub async fn from_env_with_vars(
        application_id: Option<String>,
        env_vars: &std::collections::HashMap<String, String>,
    ) -> Result<Self> {
        let app_id = application_id.unwrap_or_else(|| Uuid::new_v4().to_string());

        // Current runtimes inject both names. Prefer the current name while
        // retaining compatibility with runtimes released before the rename.
        if let Some(endpoint) = worker_protocol_endpoint(env_vars) {
            let channel = Self::create_grpc_channel(endpoint.address).await?;
            let grpc_client = WaitUntilGrpcClient::Runtime(WaitUntilServiceClient::new(
                WorkerProtocolChannel::new(channel, endpoint.generation),
            ));

            return Ok(Self {
                application_id: app_id,
                tasks: Arc::new(Mutex::new(HashMap::new())),
                task_counter: AtomicU32::new(0),
                grpc_client: Some(grpc_client),
                draining: Arc::new(Mutex::new(false)),
            });
        }

        // No gRPC needed - run in-process
        Ok(Self::new(Some(app_id)))
    }

    /// Creates a gRPC channel from an address string.
    /// This creates a dedicated channel for wait_until with proper timeout and keep-alive configuration.
    async fn create_grpc_channel(grpc_address: &str) -> Result<Channel> {
        use std::time::Duration;

        // Ensure the address has a scheme, default to http if not present
        let endpoint_uri = if grpc_address.contains("://") {
            grpc_address.to_string()
        } else {
            format!("http://{}", grpc_address)
        };

        let endpoint = Channel::from_shared(endpoint_uri.clone())
            .into_alien_error()
            .context(ErrorData::GrpcConnectionFailed {
                endpoint: endpoint_uri.clone(),
                reason: "Invalid gRPC endpoint URI format".to_string(),
            })?
            .timeout(Duration::from_secs(300)) // 5 min timeout for long-lived drain signal RPC
            .connect_timeout(Duration::from_secs(5)) // Connection establishment timeout
            .http2_keep_alive_interval(Duration::from_secs(30)) // Keep connection alive
            .keep_alive_timeout(Duration::from_secs(10))
            .keep_alive_while_idle(true); // Keep alive even when idle (important for drain listener)

        let channel = endpoint.connect().await.into_alien_error().context(
            ErrorData::GrpcConnectionFailed {
                endpoint: grpc_address.to_string(),
                reason: "Failed to establish gRPC connection".to_string(),
            },
        )?;

        Ok(channel)
    }

    /// Creates a new WaitUntilContext with a gRPC client.
    pub fn new_with_grpc_client(
        application_id: Option<String>,
        grpc_client: WaitUntilServiceClient<Channel>,
    ) -> Self {
        let app_id = application_id.unwrap_or_else(|| Uuid::new_v4().to_string());

        Self {
            application_id: app_id,
            tasks: Arc::new(Mutex::new(HashMap::new())),
            task_counter: AtomicU32::new(0),
            grpc_client: Some(WaitUntilGrpcClient::Direct(grpc_client)),
            draining: Arc::new(Mutex::new(false)),
        }
    }

    /// Sets the gRPC client for communicating with the Worker runtime.
    pub fn set_grpc_client(&mut self, client: WaitUntilServiceClient<Channel>) {
        self.grpc_client = Some(WaitUntilGrpcClient::Direct(client));
    }

    /// Gets the application ID.
    pub fn application_id(&self) -> &str {
        &self.application_id
    }

    /// Starts a background task that waits for drain signals from the Worker runtime.
    /// This should be called once when the application starts.
    pub async fn start_drain_listener(&self) -> Result<()> {
        if let Some(mut client) = self.grpc_client.clone() {
            let app_id = self.application_id.clone();
            let context = self.clone_for_background();

            tokio::spawn(async move {
                loop {
                    debug!(app_id = %app_id, "Waiting for drain signal from Worker runtime");

                    let request = WaitForDrainSignalRequest {
                        application_id: app_id.clone(),
                        timeout: Some(prost_types::Duration {
                            seconds: 300, // 5 minute timeout
                            nanos: 0,
                        }),
                    };

                    match client.wait_for_drain_signal(request).await {
                        Ok(response) => {
                            let resp = response.into_inner();
                            if resp.should_drain {
                                info!(
                                    app_id = %app_id,
                                    reason = %resp.drain_reason,
                                    "Received drain signal from Worker runtime"
                                );

                                let drain_timeout = resp
                                    .drain_timeout
                                    .map(|d| Duration::from_secs(d.seconds as u64))
                                    .unwrap_or(Duration::from_secs(10));

                                let config = DrainConfig {
                                    timeout: drain_timeout,
                                    reason: resp.drain_reason,
                                };

                                // Drain all tasks
                                match context.drain_all(config).await {
                                    Ok(drain_response) => {
                                        // Notify the Worker runtime that draining is complete.
                                        let complete_request = NotifyDrainCompleteRequest {
                                            application_id: app_id.clone(),
                                            tasks_drained: drain_response.tasks_drained,
                                            success: drain_response.success,
                                            error_message: drain_response.error_message,
                                        };

                                        if let Err(e) =
                                            client.notify_drain_complete(complete_request).await
                                        {
                                            error!(app_id = %app_id, error = %e, "Failed to notify Worker runtime of drain completion");
                                        } else {
                                            info!(app_id = %app_id, "Successfully notified Worker runtime of drain completion");
                                        }
                                    }
                                    Err(e) => {
                                        error!(app_id = %app_id, error = %e, "Failed to drain tasks");
                                        // Still notify the Worker runtime of the failure.
                                        let complete_request = NotifyDrainCompleteRequest {
                                            application_id: app_id.clone(),
                                            tasks_drained: 0,
                                            success: false,
                                            error_message: Some(e.to_string()),
                                        };
                                        let _ =
                                            client.notify_drain_complete(complete_request).await;
                                    }
                                }
                            }
                        }
                        Err(e) => {
                            warn!(app_id = %app_id, error = %e, "Failed to wait for drain signal, retrying in 5 seconds");
                            tokio::time::sleep(Duration::from_secs(5)).await;
                        }
                    }
                }
            });
        }

        Ok(())
    }

    /// Creates a clone suitable for background tasks.
    fn clone_for_background(&self) -> Self {
        Self {
            application_id: self.application_id.clone(),
            tasks: Arc::clone(&self.tasks),
            task_counter: AtomicU32::new(self.task_counter.load(Ordering::Relaxed)),
            grpc_client: self.grpc_client.clone(),
            draining: Arc::clone(&self.draining),
        }
    }

    /// Notifies the Worker runtime that a task has been registered (if a gRPC client is available).
    async fn notify_task_registered(&self, task_description: String) -> Result<()> {
        if let Some(mut client) = self.grpc_client.clone() {
            let request = NotifyTaskRegisteredRequest {
                application_id: self.application_id.clone(),
                task_description: Some(task_description),
            };

            client
                .notify_task_registered(request)
                .await
                .into_alien_error()
                .context(ErrorData::HttpRequestFailed {
                    url: "grpc://wait_until_service".to_string(),
                    method: "notify_task_registered".to_string(),
                })?;
        }

        Ok(())
    }
}

impl WaitUntilContext {
    /// Registers a new wait_until task that will be executed immediately.
    /// The task runs in the application process but is tracked by the Worker runtime.
    pub fn wait_until<F, Fut>(&self, task_fn: F) -> Result<()>
    where
        F: FnOnce() -> Fut + Send + 'static,
        Fut: Future<Output = ()> + Send + 'static,
    {
        let task_id = self.task_counter.fetch_add(1, Ordering::Relaxed);
        let task_key = format!("task_{}", task_id);
        let task_description = format!("wait_until_task_{}", task_id);

        // Check if we're currently draining - if so, reject new tasks
        let draining = self.draining.clone();
        let tasks = self.tasks.clone();
        let app_id = self.application_id.clone();
        let task_key_clone = task_key.clone();

        // Start the task immediately
        let handle = tokio::spawn(async move {
            // Double-check if we're draining
            if *draining.lock().await {
                warn!(app_id = %app_id, task_id = %task_key_clone, "Rejecting new task - currently draining");
                return;
            }

            debug!(app_id = %app_id, task_id = %task_key_clone, "Starting wait_until task");

            let future = task_fn();
            future.await;

            debug!(app_id = %app_id, task_id = %task_key_clone, "Completed wait_until task");

            // Remove ourselves from the tasks map when done
            tasks.lock().await.remove(&task_key_clone);
        });

        // Store the task handle
        {
            let mut tasks_guard = futures::executor::block_on(self.tasks.lock());
            tasks_guard.insert(task_key.clone(), handle);
        }

        // Notify the Worker runtime in a background task (non-blocking).
        let context_clone = self.clone_for_background();
        tokio::spawn(async move {
            if let Err(e) = context_clone.notify_task_registered(task_description).await {
                warn!(app_id = %context_clone.application_id, task_id = %task_key, error = %e, "Failed to notify Worker runtime of task registration");
            }
        });

        Ok(())
    }
}

#[async_trait]
impl WaitUntil for WaitUntilContext {
    async fn wait_for_drain_signal(
        &self,
        timeout_duration: Option<Duration>,
    ) -> Result<DrainConfig> {
        if let Some(mut client) = self.grpc_client.clone() {
            let timeout_proto = timeout_duration.map(|d| prost_types::Duration {
                seconds: d.as_secs() as i64,
                nanos: d.subsec_nanos() as i32,
            });

            let request = WaitForDrainSignalRequest {
                application_id: self.application_id.clone(),
                timeout: timeout_proto,
            };

            let response = client
                .wait_for_drain_signal(request)
                .await
                .into_alien_error()
                .context(ErrorData::HttpRequestFailed {
                    url: "grpc://wait_until_service".to_string(),
                    method: "wait_for_drain_signal".to_string(),
                })?;

            let resp = response.into_inner();
            if resp.should_drain {
                let drain_timeout = resp
                    .drain_timeout
                    .map(|d| Duration::from_secs(d.seconds as u64))
                    .unwrap_or(Duration::from_secs(10));

                return Ok(DrainConfig {
                    timeout: drain_timeout,
                    reason: resp.drain_reason,
                });
            }
        }

        // If no gRPC client or no drain signal, return a default config
        Err(AlienError::new(ErrorData::Other {
            message: "No drain signal received or gRPC client not available".to_string(),
        }))
    }

    async fn drain_all(&self, config: DrainConfig) -> Result<DrainResponse> {
        info!(
            app_id = %self.application_id,
            reason = %config.reason,
            timeout_secs = config.timeout.as_secs(),
            "Starting to drain all wait_until tasks"
        );

        // Mark that we're draining to prevent new tasks
        {
            let mut draining_guard = self.draining.lock().await;
            *draining_guard = true;
        }

        let tasks_to_drain = {
            let mut tasks_guard = self.tasks.lock().await;
            std::mem::take(&mut *tasks_guard) // Take all tasks out of the map
        };

        let task_count = tasks_to_drain.len() as u32;
        info!(app_id = %self.application_id, task_count = task_count, "Draining tasks");

        let mut success = true;
        let mut error_messages = Vec::new();

        // Wait for all tasks to complete or timeout
        let drain_result = timeout(config.timeout, async {
            for (task_id, handle) in tasks_to_drain {
                match handle.await {
                    Ok(_) => {
                        debug!(app_id = %self.application_id, task_id = %task_id, "Task completed successfully");
                    }
                    Err(e) => {
                        warn!(app_id = %self.application_id, task_id = %task_id, error = %e, "Task failed");
                        success = false;
                        error_messages.push(format!("Task {} failed: {}", task_id, e));
                    }
                }
            }
        })
        .await;

        match drain_result {
            Ok(_) => {
                info!(app_id = %self.application_id, "All tasks drained successfully");
            }
            Err(_) => {
                warn!(app_id = %self.application_id, "Drain timeout exceeded");
                success = false;
                error_messages.push("Drain timeout exceeded".to_string());
            }
        }

        // Reset draining flag
        {
            let mut draining_guard = self.draining.lock().await;
            *draining_guard = false;
        }

        let error_message = if error_messages.is_empty() {
            None
        } else {
            Some(error_messages.join("; "))
        };

        Ok(DrainResponse {
            tasks_drained: task_count,
            success,
            error_message,
        })
    }

    async fn get_task_count(&self) -> Result<u32> {
        let tasks_guard = self.tasks.lock().await;
        Ok(tasks_guard.len() as u32)
    }

    async fn notify_drain_complete(&self, response: DrainResponse) -> Result<()> {
        if let Some(mut client) = self.grpc_client.clone() {
            let request = NotifyDrainCompleteRequest {
                application_id: self.application_id.clone(),
                tasks_drained: response.tasks_drained,
                success: response.success,
                error_message: response.error_message,
            };

            client
                .notify_drain_complete(request)
                .await
                .into_alien_error()
                .context(ErrorData::HttpRequestFailed {
                    url: "grpc://wait_until_service".to_string(),
                    method: "notify_drain_complete".to_string(),
                })?;
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn worker_protocol_endpoint_accepts_legacy_runtime_address() {
        let env_vars = HashMap::from([(
            ENV_ALIEN_BINDINGS_GRPC_ADDRESS.to_string(),
            "127.0.0.1:51000".to_string(),
        )]);

        assert_eq!(
            worker_protocol_endpoint(&env_vars),
            Some(WorkerProtocolEndpoint {
                address: "127.0.0.1:51000",
                generation: WorkerProtocolGeneration::Legacy,
            })
        );
    }

    #[test]
    fn worker_protocol_endpoint_prefers_current_runtime_address() {
        let env_vars = HashMap::from([
            (
                alien_core::ENV_ALIEN_WORKER_GRPC_ADDRESS.to_string(),
                "127.0.0.1:52000".to_string(),
            ),
            (
                ENV_ALIEN_BINDINGS_GRPC_ADDRESS.to_string(),
                "127.0.0.1:51000".to_string(),
            ),
        ]);

        assert_eq!(
            worker_protocol_endpoint(&env_vars),
            Some(WorkerProtocolEndpoint {
                address: "127.0.0.1:52000",
                generation: WorkerProtocolGeneration::Current,
            })
        );
    }

    #[test]
    fn legacy_worker_protocol_uri_rewrites_all_generated_rpc_paths() {
        for (current, legacy) in [
            (
                "/alien_worker.control.ControlService/RegisterHttpServer",
                "/alien_bindings.control.ControlService/RegisterHttpServer",
            ),
            (
                "/alien_worker.control.ControlService/RegisterEventHandler",
                "/alien_bindings.control.ControlService/RegisterEventHandler",
            ),
            (
                "/alien_worker.control.ControlService/WaitForTasks",
                "/alien_bindings.control.ControlService/WaitForTasks",
            ),
            (
                "/alien_worker.control.ControlService/SendTaskResult",
                "/alien_bindings.control.ControlService/SendTaskResult",
            ),
            (
                "/alien_worker.wait_until.WaitUntilService/NotifyTaskRegistered",
                "/alien_bindings.wait_until.WaitUntilService/NotifyTaskRegistered",
            ),
            (
                "/alien_worker.wait_until.WaitUntilService/WaitForDrainSignal",
                "/alien_bindings.wait_until.WaitUntilService/WaitForDrainSignal",
            ),
            (
                "/alien_worker.wait_until.WaitUntilService/NotifyDrainComplete",
                "/alien_bindings.wait_until.WaitUntilService/NotifyDrainComplete",
            ),
            (
                "/alien_worker.wait_until.WaitUntilService/GetTaskCount",
                "/alien_bindings.wait_until.WaitUntilService/GetTaskCount",
            ),
        ] {
            assert_eq!(
                legacy_worker_protocol_uri(current).as_ref().map(Uri::path),
                Some(legacy)
            );
        }
    }
}