comfyui-client 0.1.2

#![warn(rust_2018_idioms, missing_docs)]
#![warn(clippy::dbg_macro, clippy::print_stdout)]
#![doc = include_str!("../README.md")]

/// Module containing error definitions.
pub mod errors;
/// Module containing metadata such as prompt and file information.
pub mod meta;

pub use crate::errors::{ClientError, ClientResult};
use crate::meta::{FileInfo, PromptInfo};
use bytes::Bytes;
use cfg_if::cfg_if;
use errors::{ApiBody, ApiError};
use futures_util::StreamExt;
use meta::{Event, History, Prompt};
use reqwest::{
    Body, IntoUrl, Response,
    multipart::{self},
};
use serde::Serialize;
use serde_json::{Value, json};
use std::{
    collections::HashMap,
    ops::{Deref, DerefMut},
};
use tokio::{sync::mpsc, task::JoinHandle};
use tokio_stream::wrappers::ReceiverStream;
use tokio_tungstenite::{
    connect_async,
    tungstenite::{self, Message},
};
use url::Url;
use uuid::Uuid;

/// A builder for creating a [`ComfyUIClient`] instance.
///
/// This builder helps initialize the client with the provided base URL and sets
/// up a websocket connection to stream events.
pub struct ClientBuilder {
    base_url: Url,
    channel_bound: usize,
}

impl ClientBuilder {
    /// Creates a new [`ClientBuilder`] instance.
    ///
    /// # Parameters
    ///
    /// - `base_url`: The base URL of the ComfyUI service.
    ///
    /// # Returns
    ///
    /// A new instance of [`ClientBuilder`] wrapped in a `ClientResult`, or an
    /// error if the URL is invalid.
    pub fn new(base_url: impl IntoUrl) -> ClientResult<Self> {
        Ok(Self {
            base_url: base_url.into_url()?,
            channel_bound: 100,
        })
    }

    /// Builds the [`ComfyUIClient`] along with an associated [`EventStream`].
    ///
    /// This method establishes a websocket connection and spawns an
    /// asynchronous task to process incoming messages.
    ///
    /// # Returns
    ///
    /// A tuple containing the [`ComfyUIClient`] and [`EventStream`] on success,
    /// or an error.
    pub async fn build(self) -> ClientResult<(ComfyUIClient, EventStream)> {
        let base_url = self.base_url;
        let http_client = reqwest::Client::new();
        let client_id = Uuid::new_v4().to_string();

        let (ev_tx, ev_rx) = mpsc::channel(self.channel_bound);

        let ws_url = Self::generate_websocket_url(base_url.clone(), &client_id)?;
        let (stream, _) = if ws_url.scheme() == "wss" {
            cfg_if! {
                if #[cfg(feature = "rustls")] {
                    let root_store = rustls::RootCertStore {
                        roots: webpki_roots::TLS_SERVER_ROOTS.into(),
                    };
                    let config = rustls::ClientConfig::builder()
                        .with_root_certificates(root_store)
                        .with_no_client_auth();

                    tokio_tungstenite::connect_async_tls_with_config(
                        ws_url,
                        None,
                        false,
                        Some(tokio_tungstenite::Connector::Rustls(std::sync::Arc::new(config))),
                    )
                    .await?
                } else {
                    connect_async(ws_url).await?
                }
            }
        } else {
            connect_async(ws_url).await?
        };

        let stream_handle = tokio::spawn(async move {
            let (_, mut read_stream) = stream.split();
            while let Some(msg) = read_stream.next().await {
                let ev = EventStream::handle_message(msg);
                let Some(ev) = ev.transpose() else {
                    continue;
                };
                if ev_tx.send(ev).await.is_err() {
                    break;
                }
            }
        });

        let rx_stream = ReceiverStream::new(ev_rx);

        let client = ComfyUIClient {
            base_url,
            http_client,
            client_id,
        };

        let stream = EventStream {
            stream_handle,
            rx_stream,
        };

        Ok((client, stream))
    }

    /// Builds a [`ComfyUIClient`] instance configured for HTTP-only
    /// communication.
    ///
    /// This method initializes the client without establishing a websocket
    /// connection, enabling you to interact with the ComfyUI service using
    /// only HTTP (REST) requests.
    ///
    /// # Returns
    ///
    /// A [`ComfyUIClient`] instance on success, or an error.
    pub async fn build_only_http(self) -> ClientResult<ComfyUIClient> {
        let base_url = self.base_url;
        let http_client = reqwest::Client::new();
        let client_id = Uuid::new_v4().to_string();

        Ok(ComfyUIClient {
            base_url,
            http_client,
            client_id,
        })
    }

    /// Generates the websocket URL based on the base URL and client ID.
    ///
    /// This method changes the URL scheme to `wss` if the base URL uses HTTPS,
    /// or `ws` otherwise, appends the `ws` path, and adds a query parameter
    /// for `clientId`.
    ///
    /// # Parameters
    ///
    /// - `base_url`: The base URL of the ComfyUI service.
    /// - `client_id`: The unique identifier for the client.
    ///
    /// # Returns
    ///
    /// The generated websocket URL on success, or an error if the URL cannot be
    /// modified.
    fn generate_websocket_url(base_url: Url, client_id: &str) -> ClientResult<Url> {
        let mut ws_url = base_url;
        let scheme = if ws_url.scheme() == "https" {
            "wss"
        } else {
            "ws"
        };
        ws_url
            .set_scheme(scheme)
            .map_err(|_| ClientError::SetWsScheme)?;
        ws_url = ws_url.join("ws")?;
        ws_url.query_pairs_mut().append_pair("clientId", client_id);
        Ok(ws_url)
    }
}

/// A client for interacting with the ComfyUI service.
///
/// This client provides methods to fetch history, prompts, views, and to upload
/// images.
pub struct ComfyUIClient {
    client_id: String,
    base_url: Url,
    http_client: reqwest::Client,
}

impl ComfyUIClient {
    /// Retrieves the history for a specified prompt.
    ///
    /// Sends a GET request to the `history/{prompt_id}` endpoint and parses the
    /// returned history data.
    ///
    /// # Parameters
    ///
    /// - `prompt_id`: The ID of the prompt whose history is being requested.
    ///
    /// # Returns
    ///
    /// An optional [`History`] object wrapped in a `ClientResult`. Returns
    /// `None` if the history is not found.
    pub async fn get_history(&self, prompt_id: &str) -> ClientResult<Option<History>> {
        let resp = self
            .http_client
            .get(self.base_url.join(&format!("history/{prompt_id}"))?)
            .send()
            .await?;
        let resp = Self::error_for_status(resp).await?;
        let mut histories = resp.json::<HashMap<String, History>>().await?;
        Ok(histories.remove(prompt_id))
    }

    /// Retrieves the current prompt information.
    ///
    /// Sends a GET request to the `prompt` endpoint and returns the parsed
    /// [`PromptInfo`] data.
    ///
    /// # Returns
    ///
    /// A [`PromptInfo`] object on success, or an error.
    pub async fn get_prompt(&self) -> ClientResult<PromptInfo> {
        let resp = self
            .http_client
            .get(self.base_url.join("prompt")?)
            .send()
            .await?;
        let resp = Self::error_for_status(resp).await?;
        Ok(resp.json().await?)
    }

    /// Retrieves view data corresponding to the provided file information.
    ///
    /// Sends a GET request to the `view` endpoint, including the file
    /// information as query parameters.
    ///
    /// # Parameters
    ///
    /// - `file_info`: A [`FileInfo`] object containing details about the file.
    ///
    /// # Returns
    ///
    /// The response as a [`Bytes`] object on success, or an error.
    pub async fn get_view(&self, file_info: &FileInfo) -> ClientResult<Bytes> {
        let resp = self
            .http_client
            .get(self.base_url.join("view")?)
            .query(file_info)
            .send()
            .await?;
        let resp = Self::error_for_status(resp).await?;
        Ok(resp.bytes().await?)
    }

    /// Sends a prompt in string format.
    ///
    /// Parses the input string as JSON and calls [`Self::post_prompt_value`] to
    /// send the prompt.
    ///
    /// # Parameters
    ///
    /// - `prompt`: A string slice representing the prompt in JSON format.
    ///
    /// # Returns
    ///
    /// A [`Prompt`] object on success, or an error.
    pub async fn post_prompt_str(&self, prompt: &str) -> ClientResult<Prompt> {
        let prompt = serde_json::from_str::<Value>(prompt)?;
        self.post_prompt_value(&prompt).await
    }

    /// Sends a prompt from any serializable data.
    ///
    /// Converts the provided data into JSON and calls
    /// [`Self::post_prompt_value`] to send the prompt.
    ///
    /// # Parameters
    ///
    /// - `prompt`: A reference to any data that implements [`Serialize`].
    ///
    /// # Returns
    ///
    /// A [`Prompt`] object on success, or an error.
    pub async fn post_prompt<T: Serialize>(&self, prompt: &T) -> ClientResult<Prompt> {
        let prompt = serde_json::to_value(prompt)?;
        self.post_prompt_value(&prompt).await
    }

    /// Sends a prompt in JSON format.
    ///
    /// Constructs the request payload (including the client ID and prompt data)
    /// and sends a POST request to the `prompt` endpoint.
    ///
    /// # Parameters
    ///
    /// - `prompt`: A JSON value representing the prompt data.
    ///
    /// # Returns
    ///
    /// A [`Prompt`] object on success, or an error.
    pub async fn post_prompt_value(&self, prompt: &Value) -> ClientResult<Prompt> {
        let data = json!({"client_id": &self.client_id, "prompt": prompt});
        let resp = self
            .http_client
            .post(self.base_url.join("prompt")?)
            .json(&data)
            .send()
            .await?;
        let resp = Self::error_for_status(resp).await?;
        Ok(resp.json().await?)
    }

    /// Uploads an image.
    ///
    /// Constructs a multipart form containing the image data and file
    /// information, then sends a POST request to the `upload/image` endpoint.
    ///
    /// # Parameters
    ///
    /// - `body`: The image data, convertible into a [`Body`].
    /// - `info`: A [`FileInfo`] object containing details about the image file.
    /// - `overwrite`: A boolean indicating whether to overwrite an existing
    ///   file.
    ///
    /// # Returns
    ///
    /// An updated [`FileInfo`] object on success, or an error.
    pub async fn upload_image(
        &self, body: impl Into<Body>, info: &FileInfo, overwrite: bool,
    ) -> ClientResult<FileInfo> {
        let part = multipart::Part::stream(body).file_name(info.filename.to_string());
        let mut form = multipart::Form::new()
            .part("image", part)
            .text("overwrite", overwrite.to_string())
            .text("type", info.r#type.to_string());
        if !info.subfolder.is_empty() {
            form = form.text("subfolder", info.subfolder.to_string());
        }

        let resp = self
            .http_client
            .post(self.base_url.join("upload/image")?)
            .multipart(form)
            .send()
            .await?;

        let resp = Self::error_for_status(resp).await?;
        Ok(resp.json().await?)
    }

    /// Checks the HTTP response status code and returns an error if it
    /// indicates failure.
    ///
    /// If the response status is a client or server error, this method attempts
    /// to parse the response body as JSON. If parsing fails, it returns the
    /// body as text.
    ///
    /// # Parameters
    ///
    /// - `resp`: The HTTP response to evaluate.
    ///
    /// # Returns
    ///
    /// The original response if the status is successful, or an error if the
    /// status indicates a failure.
    async fn error_for_status(resp: Response) -> ClientResult<Response> {
        let status = resp.status();
        if status.is_client_error() || status.is_server_error() {
            let body = resp.text().await?;
            let body = match serde_json::from_str::<Value>(&body) {
                Ok(value) => ApiBody::Json(value),
                Err(_) => ApiBody::Text(body),
            };
            Err(ApiError { status, body }.into())
        } else {
            Ok(resp)
        }
    }
}

/// A structure representing the event stream received via a websocket
/// connection.
///
/// This stream continuously processes events from the ComfyUI service.
pub struct EventStream {
    stream_handle: JoinHandle<()>,
    rx_stream: ReceiverStream<ClientResult<Event>>,
}

impl EventStream {
    /// Handles a single websocket message and attempts to parse it as an
    /// [`Event`].
    ///
    /// For text messages, it tries to deserialize the message into an
    /// [`Event`]. If the deserialization fails, it wraps the message as
    /// [`Event::Unknown`]. Other message types are ignored.
    ///
    /// # Parameters
    ///
    /// - `msg`: A result containing a [`Message`] from the websocket.
    ///
    /// # Returns
    ///
    /// An `Option<Event>` wrapped in a `ClientResult`. Returns `None` for
    /// unsupported message types.
    fn handle_message(msg: tungstenite::Result<Message>) -> ClientResult<Option<Event>> {
        let msg = msg?;
        match msg {
            Message::Text(b) => {
                let value = serde_json::from_slice::<Value>(b.as_bytes())?;
                match serde_json::from_value::<Event>(value.clone()) {
                    Ok(ev) => Ok(Some(ev)),
                    Err(_) => Ok(Some(Event::Unknown(value))),
                }
            }
            _ => Ok(None),
        }
    }
}

impl Drop for EventStream {
    /// When the [`EventStream`] is dropped, abort the associated websocket
    /// handling task.
    fn drop(&mut self) {
        self.stream_handle.abort();
    }
}

impl Deref for EventStream {
    type Target = ReceiverStream<ClientResult<Event>>;

    /// Allows access to the inner [`ReceiverStream`] containing the events.
    fn deref(&self) -> &Self::Target {
        &self.rx_stream
    }
}

impl DerefMut for EventStream {
    /// Allows mutable access to the inner [`ReceiverStream`].
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.rx_stream
    }
}