feat: add comfyui sdk

This commit is contained in:
imeepos
2025-08-08 13:48:38 +08:00
parent 763b4a975c
commit 5f6a302dfd
27 changed files with 4713 additions and 0 deletions

View File

@@ -0,0 +1,216 @@
//! Error handling utilities
use crate::error::{ComfyUIError, Result};
use std::time::Duration;
use tokio::time::sleep;
/// Retry configuration
#[derive(Debug, Clone)]
pub struct RetryConfig {
pub max_attempts: u32,
pub initial_delay: Duration,
pub max_delay: Duration,
pub backoff_multiplier: f64,
}
impl Default for RetryConfig {
fn default() -> Self {
Self {
max_attempts: 3,
initial_delay: Duration::from_millis(1000),
max_delay: Duration::from_secs(30),
backoff_multiplier: 2.0,
}
}
}
/// Retries an async operation with exponential backoff
pub async fn retry_with_backoff<F, Fut, T>(
operation: F,
config: RetryConfig,
) -> Result<T>
where
F: Fn() -> Fut,
Fut: std::future::Future<Output = Result<T>>,
{
let mut last_error = None;
let mut delay = config.initial_delay;
for attempt in 1..=config.max_attempts {
match operation().await {
Ok(result) => return Ok(result),
Err(error) => {
last_error = Some(error);
if attempt < config.max_attempts {
log::warn!("Attempt {} failed, retrying in {:?}", attempt, delay);
sleep(delay).await;
// Exponential backoff
delay = std::cmp::min(
Duration::from_millis(
(delay.as_millis() as f64 * config.backoff_multiplier) as u64
),
config.max_delay,
);
}
}
}
}
Err(last_error.unwrap_or_else(|| ComfyUIError::new("Retry failed with no error")))
}
/// Checks if an error is retryable
pub fn is_retryable_error(error: &ComfyUIError) -> bool {
match error {
ComfyUIError::Http(reqwest_error) => {
// Retry on network errors, timeouts, and 5xx status codes
reqwest_error.is_timeout() ||
reqwest_error.is_connect() ||
reqwest_error.status().map_or(false, |status| status.is_server_error())
}
ComfyUIError::WebSocket(_) => true, // Most WebSocket errors are retryable
ComfyUIError::Connection(_) => true,
ComfyUIError::Timeout(_) => true,
ComfyUIError::Io(_) => true,
_ => false, // Don't retry validation errors, etc.
}
}
/// Retries an operation only if the error is retryable
pub async fn retry_if_retryable<F, Fut, T>(
operation: F,
config: RetryConfig,
) -> Result<T>
where
F: Fn() -> Fut,
Fut: std::future::Future<Output = Result<T>>,
{
let mut last_error = None;
let mut delay = config.initial_delay;
for attempt in 1..=config.max_attempts {
match operation().await {
Ok(result) => return Ok(result),
Err(error) => {
if !is_retryable_error(&error) {
return Err(error);
}
last_error = Some(error);
if attempt < config.max_attempts {
log::warn!("Attempt {} failed with retryable error, retrying in {:?}", attempt, delay);
sleep(delay).await;
// Exponential backoff
delay = std::cmp::min(
Duration::from_millis(
(delay.as_millis() as f64 * config.backoff_multiplier) as u64
),
config.max_delay,
);
}
}
}
}
Err(last_error.unwrap_or_else(|| ComfyUIError::new("Retry failed with no error")))
}
/// Wraps an operation with timeout
pub async fn with_timeout<F, T>(
operation: F,
timeout: Duration,
) -> Result<T>
where
F: std::future::Future<Output = Result<T>>,
{
match tokio::time::timeout(timeout, operation).await {
Ok(result) => result,
Err(_) => Err(ComfyUIError::timeout(format!(
"Operation timed out after {:?}",
timeout
))),
}
}
/// Error context helper
pub trait ErrorContext<T> {
fn with_context(self, context: &str) -> Result<T>;
}
impl<T> ErrorContext<T> for Result<T> {
fn with_context(self, context: &str) -> Result<T> {
self.map_err(|error| {
ComfyUIError::new(format!("{}: {}", context, error))
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
#[tokio::test]
async fn test_retry_success_on_second_attempt() {
let counter = Arc::new(AtomicU32::new(0));
let counter_clone = counter.clone();
let config = RetryConfig {
max_attempts: 3,
initial_delay: Duration::from_millis(10),
max_delay: Duration::from_millis(100),
backoff_multiplier: 2.0,
};
let result = retry_with_backoff(
|| {
let counter = counter_clone.clone();
async move {
let count = counter.fetch_add(1, Ordering::SeqCst);
if count == 0 {
Err(ComfyUIError::connection("First attempt fails"))
} else {
Ok("Success")
}
}
},
config,
).await;
assert!(result.is_ok());
assert_eq!(result.unwrap(), "Success");
assert_eq!(counter.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn test_retry_exhausted() {
let counter = Arc::new(AtomicU32::new(0));
let counter_clone = counter.clone();
let config = RetryConfig {
max_attempts: 2,
initial_delay: Duration::from_millis(10),
max_delay: Duration::from_millis(100),
backoff_multiplier: 2.0,
};
let result: Result<&str> = retry_with_backoff(
|| {
let counter = counter_clone.clone();
async move {
counter.fetch_add(1, Ordering::SeqCst);
Err(ComfyUIError::connection("Always fails"))
}
},
config,
).await;
assert!(result.is_err());
assert_eq!(counter.load(Ordering::SeqCst), 2);
}
}

View File

@@ -0,0 +1,171 @@
//! Event emitter utilities
use std::sync::Arc;
use tokio::sync::RwLock;
use std::collections::HashMap;
use crate::types::{ExecutionProgress, ExecutionResult, ExecutionError, ExecutionCallbacks};
/// Event emitter for execution events
#[derive(Clone)]
pub struct EventEmitter {
callbacks: Arc<RwLock<HashMap<String, Arc<dyn ExecutionCallbacks>>>>,
}
impl EventEmitter {
/// Creates a new event emitter
pub fn new() -> Self {
Self {
callbacks: Arc::new(RwLock::new(HashMap::new())),
}
}
/// Registers callbacks with a unique ID
pub async fn register_callbacks(&self, id: String, callbacks: Arc<dyn ExecutionCallbacks>) {
let mut callbacks_map = self.callbacks.write().await;
callbacks_map.insert(id, callbacks);
}
/// Unregisters callbacks by ID
pub async fn unregister_callbacks(&self, id: &str) {
let mut callbacks_map = self.callbacks.write().await;
callbacks_map.remove(id);
}
/// Emits a progress event
pub async fn emit_progress(&self, progress: ExecutionProgress) {
let callbacks_map = self.callbacks.read().await;
for callback in callbacks_map.values() {
callback.on_progress(progress.clone());
}
}
/// Emits an executing event
pub async fn emit_executing(&self, node_id: String) {
let callbacks_map = self.callbacks.read().await;
for callback in callbacks_map.values() {
callback.on_executing(node_id.clone());
}
}
/// Emits an executed event
pub async fn emit_executed(&self, result: ExecutionResult) {
let callbacks_map = self.callbacks.read().await;
for callback in callbacks_map.values() {
callback.on_executed(result.clone());
}
}
/// Emits an error event
pub async fn emit_error(&self, error: ExecutionError) {
let callbacks_map = self.callbacks.read().await;
for callback in callbacks_map.values() {
callback.on_error(error.clone());
}
}
/// Gets the number of registered callbacks
pub async fn callback_count(&self) -> usize {
let callbacks_map = self.callbacks.read().await;
callbacks_map.len()
}
/// Clears all callbacks
pub async fn clear(&self) {
let mut callbacks_map = self.callbacks.write().await;
callbacks_map.clear();
}
}
impl Default for EventEmitter {
fn default() -> Self {
Self::new()
}
}
/// Simple callback implementation for testing
pub struct SimpleCallbacks {
pub on_progress_fn: Option<Box<dyn Fn(ExecutionProgress) + Send + Sync>>,
pub on_executing_fn: Option<Box<dyn Fn(String) + Send + Sync>>,
pub on_executed_fn: Option<Box<dyn Fn(ExecutionResult) + Send + Sync>>,
pub on_error_fn: Option<Box<dyn Fn(ExecutionError) + Send + Sync>>,
}
impl SimpleCallbacks {
/// Creates new simple callbacks
pub fn new() -> Self {
Self {
on_progress_fn: None,
on_executing_fn: None,
on_executed_fn: None,
on_error_fn: None,
}
}
/// Sets the progress callback
pub fn with_progress<F>(mut self, f: F) -> Self
where
F: Fn(ExecutionProgress) + Send + Sync + 'static,
{
self.on_progress_fn = Some(Box::new(f));
self
}
/// Sets the executing callback
pub fn with_executing<F>(mut self, f: F) -> Self
where
F: Fn(String) + Send + Sync + 'static,
{
self.on_executing_fn = Some(Box::new(f));
self
}
/// Sets the executed callback
pub fn with_executed<F>(mut self, f: F) -> Self
where
F: Fn(ExecutionResult) + Send + Sync + 'static,
{
self.on_executed_fn = Some(Box::new(f));
self
}
/// Sets the error callback
pub fn with_error<F>(mut self, f: F) -> Self
where
F: Fn(ExecutionError) + Send + Sync + 'static,
{
self.on_error_fn = Some(Box::new(f));
self
}
}
impl ExecutionCallbacks for SimpleCallbacks {
fn on_progress(&self, progress: ExecutionProgress) {
if let Some(ref f) = self.on_progress_fn {
f(progress);
}
}
fn on_executing(&self, node_id: String) {
if let Some(ref f) = self.on_executing_fn {
f(node_id);
}
}
fn on_executed(&self, result: ExecutionResult) {
if let Some(ref f) = self.on_executed_fn {
f(result);
}
}
fn on_error(&self, error: ExecutionError) {
if let Some(ref f) = self.on_error_fn {
f(error);
}
}
}
impl Default for SimpleCallbacks {
fn default() -> Self {
Self::new()
}
}

View File

@@ -0,0 +1,12 @@
//! Utility functions for ComfyUI SDK
pub mod validation;
pub mod template_parser;
pub mod event_emitter;
pub mod error_handling;
// Re-export main utilities
pub use validation::*;
pub use template_parser::*;
pub use event_emitter::*;
pub use error_handling::*;

View File

@@ -0,0 +1,256 @@
//! Template parsing utilities for parameter substitution
use std::collections::HashMap;
use regex::Regex;
use crate::types::{ComfyUIWorkflow, ParameterValues, ComfyUINode, VariableSyntax};
use crate::error::{ComfyUIError, Result};
/// Applies parameters to a workflow template
pub fn apply_parameters(
workflow: &ComfyUIWorkflow,
parameters: &ParameterValues,
) -> Result<ComfyUIWorkflow> {
let mut resolved_workflow = HashMap::new();
for (node_id, node) in workflow {
let resolved_node = apply_parameters_to_node(node, parameters)?;
resolved_workflow.insert(node_id.clone(), resolved_node);
}
Ok(resolved_workflow)
}
/// Applies parameters to a single node
fn apply_parameters_to_node(
node: &ComfyUINode,
parameters: &ParameterValues,
) -> Result<ComfyUINode> {
let mut resolved_inputs = HashMap::new();
for (input_name, input_value) in &node.inputs {
let resolved_value = apply_parameters_to_value(input_value, parameters)?;
resolved_inputs.insert(input_name.clone(), resolved_value);
}
Ok(ComfyUINode {
class_type: node.class_type.clone(),
inputs: resolved_inputs,
_meta: node._meta.clone(),
})
}
/// Applies parameters to a JSON value
fn apply_parameters_to_value(
value: &serde_json::Value,
parameters: &ParameterValues,
) -> Result<serde_json::Value> {
match value {
serde_json::Value::String(s) => {
let resolved_string = substitute_variables(s, parameters, VariableSyntax::DoubleBrace)?;
Ok(serde_json::Value::String(resolved_string))
}
serde_json::Value::Array(arr) => {
let mut resolved_array = Vec::new();
for item in arr {
resolved_array.push(apply_parameters_to_value(item, parameters)?);
}
Ok(serde_json::Value::Array(resolved_array))
}
serde_json::Value::Object(obj) => {
let mut resolved_object = serde_json::Map::new();
for (key, val) in obj {
let resolved_key = substitute_variables(key, parameters, VariableSyntax::DoubleBrace)?;
let resolved_val = apply_parameters_to_value(val, parameters)?;
resolved_object.insert(resolved_key, resolved_val);
}
Ok(serde_json::Value::Object(resolved_object))
}
_ => Ok(value.clone()),
}
}
/// Substitutes variables in a string using the specified syntax
pub fn substitute_variables(
template: &str,
parameters: &ParameterValues,
syntax: VariableSyntax,
) -> Result<String> {
let pattern = match syntax {
VariableSyntax::DoubleBrace => r"\{\{([^}]+)\}\}",
VariableSyntax::DollarBrace => r"\$\{([^}]+)\}",
VariableSyntax::AtBrace => r"@\{([^}]+)\}",
};
let regex = Regex::new(pattern)
.map_err(|e| ComfyUIError::new(format!("Invalid regex pattern: {}", e)))?;
let mut result = template.to_string();
let mut offset = 0i32;
for captures in regex.captures_iter(template) {
let full_match = captures.get(0).unwrap();
let var_name = captures.get(1).unwrap().as_str().trim();
// Get parameter value
let replacement = match parameters.get(var_name) {
Some(value) => value_to_string(value)?,
None => {
return Err(ComfyUIError::template_validation(
format!("Parameter '{}' not found", var_name)
));
}
};
// Calculate positions with offset
let start = (full_match.start() as i32 + offset) as usize;
let end = (full_match.end() as i32 + offset) as usize;
// Replace the variable
result.replace_range(start..end, &replacement);
// Update offset
offset += replacement.len() as i32 - full_match.len() as i32;
}
Ok(result)
}
/// Converts a JSON value to string for substitution
fn value_to_string(value: &serde_json::Value) -> Result<String> {
match value {
serde_json::Value::String(s) => Ok(s.clone()),
serde_json::Value::Number(n) => Ok(n.to_string()),
serde_json::Value::Bool(b) => Ok(b.to_string()),
serde_json::Value::Null => Ok("null".to_string()),
_ => {
// For complex types, serialize to JSON
serde_json::to_string(value)
.map_err(ComfyUIError::from)
}
}
}
/// Extracts variable names from a template string
pub fn extract_variables(template: &str, syntax: VariableSyntax) -> Result<Vec<String>> {
let pattern = match syntax {
VariableSyntax::DoubleBrace => r"\{\{([^}]+)\}\}",
VariableSyntax::DollarBrace => r"\$\{([^}]+)\}",
VariableSyntax::AtBrace => r"@\{([^}]+)\}",
};
let regex = Regex::new(pattern)
.map_err(|e| ComfyUIError::new(format!("Invalid regex pattern: {}", e)))?;
let mut variables = Vec::new();
for captures in regex.captures_iter(template) {
if let Some(var_match) = captures.get(1) {
let var_name = var_match.as_str().trim().to_string();
if !variables.contains(&var_name) {
variables.push(var_name);
}
}
}
Ok(variables)
}
/// Extracts all variables from a workflow
pub fn extract_workflow_variables(workflow: &ComfyUIWorkflow) -> Result<Vec<String>> {
let mut all_variables = Vec::new();
for node in workflow.values() {
let node_variables = extract_node_variables(node)?;
for var in node_variables {
if !all_variables.contains(&var) {
all_variables.push(var);
}
}
}
Ok(all_variables)
}
/// Extracts variables from a single node
fn extract_node_variables(node: &ComfyUINode) -> Result<Vec<String>> {
let mut variables = Vec::new();
for input_value in node.inputs.values() {
let value_variables = extract_value_variables(input_value)?;
for var in value_variables {
if !variables.contains(&var) {
variables.push(var);
}
}
}
Ok(variables)
}
/// Extracts variables from a JSON value
fn extract_value_variables(value: &serde_json::Value) -> Result<Vec<String>> {
let mut variables = Vec::new();
match value {
serde_json::Value::String(s) => {
let string_vars = extract_variables(s, VariableSyntax::DoubleBrace)?;
variables.extend(string_vars);
}
serde_json::Value::Array(arr) => {
for item in arr {
let item_vars = extract_value_variables(item)?;
variables.extend(item_vars);
}
}
serde_json::Value::Object(obj) => {
for (key, val) in obj {
let key_vars = extract_variables(key, VariableSyntax::DoubleBrace)?;
variables.extend(key_vars);
let val_vars = extract_value_variables(val)?;
variables.extend(val_vars);
}
}
_ => {}
}
// Remove duplicates
variables.sort();
variables.dedup();
Ok(variables)
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn test_substitute_variables() {
let mut parameters = HashMap::new();
parameters.insert("name".to_string(), json!("test"));
parameters.insert("value".to_string(), json!(42));
let template = "Hello {{name}}, value is {{value}}";
let result = substitute_variables(template, &parameters, VariableSyntax::DoubleBrace).unwrap();
assert_eq!(result, "Hello test, value is 42");
}
#[test]
fn test_extract_variables() {
let template = "{{var1}} and {{var2}} and {{var1}} again";
let variables = extract_variables(template, VariableSyntax::DoubleBrace).unwrap();
assert_eq!(variables, vec!["var1", "var2"]);
}
#[test]
fn test_missing_parameter() {
let parameters = HashMap::new();
let template = "Hello {{missing}}";
let result = substitute_variables(template, &parameters, VariableSyntax::DoubleBrace);
assert!(result.is_err());
}
}

View File

@@ -0,0 +1,276 @@
//! Parameter validation utilities
use std::collections::HashMap;
use crate::types::{ParameterSchema, ParameterValues, ValidationResult, ValidationError, ParameterType};
use regex::Regex;
/// Validates parameters against their schemas
pub fn validate_parameters(
parameters: &ParameterValues,
schemas: &HashMap<String, ParameterSchema>,
) -> ValidationResult {
let mut errors = Vec::new();
// Check required parameters
for (name, schema) in schemas {
if schema.required.unwrap_or(false) {
if !parameters.contains_key(name) {
errors.push(ValidationError::new(
name,
"Required parameter is missing",
));
continue;
}
}
// Validate parameter if present
if let Some(value) = parameters.get(name) {
if let Err(error) = validate_parameter_value(name, value, schema) {
errors.push(error);
}
}
}
// Check for unknown parameters
for name in parameters.keys() {
if !schemas.contains_key(name) {
errors.push(ValidationError::new(
name,
"Unknown parameter",
));
}
}
if errors.is_empty() {
ValidationResult::success()
} else {
ValidationResult::failure(errors)
}
}
/// Validates a single parameter value against its schema
pub fn validate_parameter_value(
name: &str,
value: &serde_json::Value,
schema: &ParameterSchema,
) -> Result<(), ValidationError> {
// Check type
match (&schema.param_type, value) {
(ParameterType::String, serde_json::Value::String(s)) => {
validate_string_constraints(name, s, schema)?;
}
(ParameterType::Number, serde_json::Value::Number(n)) => {
validate_number_constraints(name, n, schema)?;
}
(ParameterType::Boolean, serde_json::Value::Bool(_)) => {
// Boolean values are always valid
}
(ParameterType::Array, serde_json::Value::Array(arr)) => {
validate_array_constraints(name, arr, schema)?;
}
(ParameterType::Object, serde_json::Value::Object(_)) => {
// Object validation could be more complex, but for now just check type
}
_ => {
return Err(ValidationError::with_value(
name,
format!("Expected type {:?}, got {:?}", schema.param_type, get_value_type(value)),
value.clone(),
));
}
}
// Check enum constraints
if let Some(enum_values) = &schema.r#enum {
if !enum_values.contains(value) {
return Err(ValidationError::with_value(
name,
format!("Value must be one of: {:?}", enum_values),
value.clone(),
));
}
}
Ok(())
}
/// Validates string constraints
fn validate_string_constraints(
name: &str,
value: &str,
schema: &ParameterSchema,
) -> Result<(), ValidationError> {
// Check pattern
if let Some(pattern) = &schema.pattern {
let regex = Regex::new(pattern).map_err(|_| {
ValidationError::new(name, "Invalid regex pattern in schema")
})?;
if !regex.is_match(value) {
return Err(ValidationError::with_value(
name,
format!("String does not match pattern: {}", pattern),
serde_json::Value::String(value.to_string()),
));
}
}
// Check length constraints (using min/max as length bounds for strings)
if let Some(min) = schema.min {
if (value.len() as f64) < min {
return Err(ValidationError::with_value(
name,
format!("String length must be at least {}", min),
serde_json::Value::String(value.to_string()),
));
}
}
if let Some(max) = schema.max {
if (value.len() as f64) > max {
return Err(ValidationError::with_value(
name,
format!("String length must be at most {}", max),
serde_json::Value::String(value.to_string()),
));
}
}
Ok(())
}
/// Validates number constraints
fn validate_number_constraints(
name: &str,
value: &serde_json::Number,
schema: &ParameterSchema,
) -> Result<(), ValidationError> {
let num_value = value.as_f64().unwrap_or(0.0);
if let Some(min) = schema.min {
if num_value < min {
return Err(ValidationError::with_value(
name,
format!("Number must be at least {}", min),
serde_json::Value::Number(value.clone()),
));
}
}
if let Some(max) = schema.max {
if num_value > max {
return Err(ValidationError::with_value(
name,
format!("Number must be at most {}", max),
serde_json::Value::Number(value.clone()),
));
}
}
Ok(())
}
/// Validates array constraints
fn validate_array_constraints(
name: &str,
value: &[serde_json::Value],
schema: &ParameterSchema,
) -> Result<(), ValidationError> {
// Check length constraints
if let Some(min) = schema.min {
if (value.len() as f64) < min {
return Err(ValidationError::with_value(
name,
format!("Array length must be at least {}", min),
serde_json::Value::Array(value.to_vec()),
));
}
}
if let Some(max) = schema.max {
if (value.len() as f64) > max {
return Err(ValidationError::with_value(
name,
format!("Array length must be at most {}", max),
serde_json::Value::Array(value.to_vec()),
));
}
}
// Validate items if schema is provided
if let Some(items_schema) = &schema.items {
for (index, item) in value.iter().enumerate() {
let item_name = format!("{}[{}]", name, index);
if let Err(error) = validate_parameter_value(&item_name, item, items_schema) {
return Err(error);
}
}
}
Ok(())
}
/// Gets the type of a JSON value
fn get_value_type(value: &serde_json::Value) -> &'static str {
match value {
serde_json::Value::Null => "null",
serde_json::Value::Bool(_) => "boolean",
serde_json::Value::Number(_) => "number",
serde_json::Value::String(_) => "string",
serde_json::Value::Array(_) => "array",
serde_json::Value::Object(_) => "object",
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn test_validate_required_parameter() {
let mut schemas = HashMap::new();
schemas.insert("test".to_string(), ParameterSchema {
param_type: ParameterType::String,
required: Some(true),
default: None,
description: None,
r#enum: None,
min: None,
max: None,
pattern: None,
items: None,
properties: None,
});
let parameters = HashMap::new();
let result = validate_parameters(&parameters, &schemas);
assert!(!result.valid);
assert_eq!(result.errors.len(), 1);
assert_eq!(result.errors[0].path, "test");
}
#[test]
fn test_validate_string_parameter() {
let mut schemas = HashMap::new();
schemas.insert("test".to_string(), ParameterSchema {
param_type: ParameterType::String,
required: Some(true),
default: None,
description: None,
r#enum: None,
min: Some(3.0),
max: Some(10.0),
pattern: None,
items: None,
properties: None,
});
let mut parameters = HashMap::new();
parameters.insert("test".to_string(), json!("hello"));
let result = validate_parameters(&parameters, &schemas);
assert!(result.valid);
}
}