tauri providers - fal api key support (#1502)

This commit is contained in:
Brandon Thomas
2026-05-12 00:02:41 -04:00
committed by GitHub
parent 0efafe8043
commit d2dbef249c
126 changed files with 2503 additions and 72 deletions
@@ -1,12 +0,0 @@
{
"db_name": "SQLite",
"query": "\n INSERT INTO tasks (\n id,\n task_status,\n task_type,\n model_type,\n provider,\n provider_job_id,\n frontend_caller,\n frontend_subscriber_id,\n frontend_subscriber_payload\n )\n VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)\n ",
"describe": {
"columns": [],
"parameters": {
"Right": 9
},
"nullable": []
},
"hash": "01a0bf33e8331c19dc73f064c54eaa863bec158ef7713450a06c75e246f5f08c"
}
@@ -0,0 +1,12 @@
{
"db_name": "SQLite",
"query": "\n INSERT INTO tasks (\n id,\n task_status,\n task_type,\n model_type,\n provider,\n provider_job_id,\n queue_status_url,\n queue_response_url,\n frontend_caller,\n frontend_subscriber_id,\n frontend_subscriber_payload\n )\n VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)\n ",
"describe": {
"columns": [],
"parameters": {
"Right": 11
},
"nullable": []
},
"hash": "15a2f8712c80b72c86d9e77eb590d74ea21dd818a424267bcf5b88f9c993c300"
}
@@ -1,6 +1,6 @@
{
"db_name": "SQLite",
"query": "\n SELECT\n id,\n task_status,\n task_type,\n model_type,\n provider,\n provider_job_id,\n frontend_caller,\n frontend_subscriber_id,\n frontend_subscriber_payload\n FROM tasks\n WHERE\n provider = ?\n AND provider_job_id = ?\n ",
"query": "\n SELECT\n id,\n task_status,\n task_type,\n model_type,\n provider,\n provider_job_id,\n frontend_caller,\n queue_status_url,\n queue_response_url,\n frontend_subscriber_id,\n frontend_subscriber_payload\n FROM tasks\n WHERE\n provider = ?\n AND provider_job_id = ?\n ",
"describe": {
"columns": [
{
@@ -39,14 +39,24 @@
"type_info": "Text"
},
{
"name": "frontend_subscriber_id",
"name": "queue_status_url",
"ordinal": 7,
"type_info": "Text"
},
{
"name": "frontend_subscriber_payload",
"name": "queue_response_url",
"ordinal": 8,
"type_info": "Text"
},
{
"name": "frontend_subscriber_id",
"ordinal": 9,
"type_info": "Text"
},
{
"name": "frontend_subscriber_payload",
"ordinal": 10,
"type_info": "Text"
}
],
"parameters": {
@@ -61,8 +71,10 @@
true,
true,
true,
true,
true,
true
]
},
"hash": "38ff2df237b16f077865711dd81261de1b5fa05a9787679ce9ccd4ac576b1069"
"hash": "68f9fed45fc55aa2ec22dc1ed22795d0c99882917abb9b52b78735b8194c2624"
}
Generated
+1
View File
@@ -679,6 +679,7 @@ dependencies = [
"discord-rich-presence",
"enums",
"errors",
"fal_client",
"filesys",
"futures 0.3.32",
"grok_client",
@@ -19,6 +19,9 @@
-- tasks_v6.sqlite - added two fields:
-- on_failure_type
-- on_failure_message
-- tasks_v7.sqlite - added two fields (for fal):
-- queue_status_url
-- queue_response_url
CREATE TABLE tasks (
-- Task auto-incrementing primary key.
@@ -63,6 +66,16 @@ CREATE TABLE tasks (
-- Whether the user has dismissed the task from view.
is_dismissed_by_user INTEGER NOT NULL DEFAULT 0,
-- OPTIONAL.
-- For some third party systems (eg. FAL), the URL to check the status of the
-- job in the queue.
queue_status_url TEXT,
-- OPTIONAL.
-- For some third party systems (eg. FAL), the URL to check the
-- results (when the job completes)
queue_response_url TEXT,
-- OPTIONAL.
-- When the generation completes, the batch token (if any).
-- This will point to a collection of multiple files.
@@ -9,4 +9,11 @@ impl RouterFalClient {
pub fn new(api_key: FalApiKey, webhook_url: String) -> Self {
Self { api_key, webhook_url }
}
pub fn new_from_raw_key(api_key: &str, webhook_url: String) -> Self {
Self {
api_key: FalApiKey::from_str(api_key),
webhook_url,
}
}
}
@@ -9,6 +9,9 @@ impl RouterFalWebhookOptionalClient {
pub fn new(api_key: FalApiKey) -> Self {
Self { api_key, webhook_url: None }
}
pub fn from_str(api_key: &str) -> Self {
Self::new(FalApiKey::from_str(api_key))
}
pub fn new_with_webhook(api_key: FalApiKey, webhook_url: String) -> Self {
Self {
@@ -1,8 +1,13 @@
use crate::errors::client_error::ClientError;
use crate::errors::download_error::DownloadError;
use crate::errors::provider_error::ProviderError;
use fal_client::error::fal_error_plus::FalErrorPlus;
use seedance2pro_client::error::seedance2pro_error::Seedance2ProError;
use seedance2pro_client::error::seedance2pro_specific_api_error::Seedance2ProSpecificApiError;
use std::error::Error;
use std::fmt::{Display, Formatter};
use artcraft_client::error::api_error::ApiError;
use artcraft_client::error::storyteller_error::StorytellerError;
#[derive(Debug)]
pub enum ArtcraftRouterError {
@@ -23,6 +28,9 @@ pub enum ArtcraftRouterError {
/// An error from an underlying provider.
Provider(ProviderError),
/// A billing error from an underlying provider.
ProviderBillingError(ProviderError),
}
impl Error for ArtcraftRouterError {}
@@ -36,6 +44,7 @@ impl Display for ArtcraftRouterError {
Self::UnsupportedProviderAndModelForNewApi(msg) => write!(f, "Unsupported provider/model (for new API during migration): {}", msg),
Self::InvalidInput(msg) => write!(f, "Invalid input: {}", msg),
Self::Provider(e) => write!(f, "Provider error: {}", e),
Self::ProviderBillingError(e) => write!(f, "Provider billing error: {}", e),
}
}
}
@@ -54,6 +63,25 @@ impl From<DownloadError> for ArtcraftRouterError {
impl From<ProviderError> for ArtcraftRouterError {
fn from(error: ProviderError) -> Self {
Self::Provider(error)
let is_billing_error = match &error {
ProviderError::Fal(FalErrorPlus::FalBillingError(_)) => true,
ProviderError::Seedance2Pro(Seedance2ProError::ApiSpecific(Seedance2ProSpecificApiError::BillingError { .. })) => true,
ProviderError::Storyteller(StorytellerError::Api(ApiError::PaymentRequired(_))) => true,
_ => false,
};
if is_billing_error {
Self::ProviderBillingError(error)
} else {
Self::Provider(error)
}
}
}
impl From<FalErrorPlus> for ArtcraftRouterError {
fn from(error: FalErrorPlus) -> Self {
match &error {
FalErrorPlus::FalBillingError(e) => ArtcraftRouterError::ProviderBillingError(ProviderError::Fal(error)),
_ => ArtcraftRouterError::Provider(ProviderError::Fal(error)),
}
}
}
@@ -58,6 +58,8 @@ pub async fn execute_fal_flux_1_dev(
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: webhook_response.request_id,
gateway_request_id: webhook_response.gateway_request_id,
maybe_status_url: None,
maybe_response_url: None,
maybe_outbound_request: Some(outbound_request),
}))
}
@@ -58,6 +58,8 @@ pub async fn execute_fal_flux_1_schnell(
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: webhook_response.request_id,
gateway_request_id: webhook_response.gateway_request_id,
maybe_status_url: None,
maybe_response_url: None,
maybe_outbound_request: Some(outbound_request),
}))
}
@@ -35,6 +35,8 @@ pub async fn execute_fal_flux_pro_1p1(
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: webhook_response.request_id,
gateway_request_id: webhook_response.gateway_request_id,
maybe_status_url: None,
maybe_response_url: None,
maybe_outbound_request: Some(outbound_request),
}))
}
@@ -35,6 +35,8 @@ pub async fn execute_fal_flux_pro_1p1_ultra(
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: webhook_response.request_id,
gateway_request_id: webhook_response.gateway_request_id,
maybe_status_url: None,
maybe_response_url: None,
maybe_outbound_request: Some(outbound_request),
}))
}
@@ -67,6 +67,8 @@ pub async fn execute_fal_gpt_image_1(
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: webhook_response.request_id,
gateway_request_id: webhook_response.gateway_request_id,
maybe_status_url: None,
maybe_response_url: None,
maybe_outbound_request: Some(outbound_request),
}))
}
@@ -67,6 +67,8 @@ pub async fn execute_fal_gpt_image_1p5(
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: webhook_response.request_id,
gateway_request_id: webhook_response.gateway_request_id,
maybe_status_url: None,
maybe_response_url: None,
maybe_outbound_request: Some(outbound_request),
}))
}
@@ -64,6 +64,8 @@ pub async fn execute_fal_gpt_image_2(
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: webhook_response.request_id,
gateway_request_id: webhook_response.gateway_request_id,
maybe_status_url: None,
maybe_response_url: None,
maybe_outbound_request: Some(outbound_request),
}))
}
@@ -63,6 +63,8 @@ pub async fn execute_fal_nano_banana(
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: webhook_response.request_id,
gateway_request_id: webhook_response.gateway_request_id,
maybe_status_url: None,
maybe_response_url: None,
maybe_outbound_request: Some(outbound_request),
}))
}
@@ -67,6 +67,8 @@ pub async fn execute_fal_nano_banana_2(
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: webhook_response.request_id,
gateway_request_id: webhook_response.gateway_request_id,
maybe_status_url: None,
maybe_response_url: None,
maybe_outbound_request: Some(outbound_request),
}))
}
@@ -66,6 +66,8 @@ pub async fn execute_fal_nano_banana_pro(
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: webhook_response.request_id,
gateway_request_id: webhook_response.gateway_request_id,
maybe_status_url: None,
maybe_response_url: None,
maybe_outbound_request: Some(outbound_request),
}))
}
@@ -61,6 +61,8 @@ pub async fn execute_fal_seedream_4(
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: webhook_response.request_id,
gateway_request_id: webhook_response.gateway_request_id,
maybe_status_url: None,
maybe_response_url: None,
maybe_outbound_request: Some(outbound_request),
}))
}
@@ -61,6 +61,8 @@ pub async fn execute_fal_seedream_4p5(
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: webhook_response.request_id,
gateway_request_id: webhook_response.gateway_request_id,
maybe_status_url: None,
maybe_response_url: None,
maybe_outbound_request: Some(outbound_request),
}))
}
@@ -61,6 +61,8 @@ pub async fn execute_fal_seedream_5_lite(
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: webhook_response.request_id,
gateway_request_id: webhook_response.gateway_request_id,
maybe_status_url: None,
maybe_response_url: None,
maybe_outbound_request: Some(outbound_request),
}))
}
@@ -13,6 +13,12 @@ pub struct FalImageResponsePayload {
pub request_id: Option<String>,
pub gateway_request_id: Option<String>,
/// The queue status URL (for polling job progress).
pub maybe_status_url: Option<String>,
/// The queue response URL (for fetching completed results).
pub maybe_response_url: Option<String>,
/// The outbound request that was sent to Fal.
/// Stored as a trait object so any Request type can be captured.
/// Use `format!("{:?}", ...)` or `format!("{:#?}", ...)` to print.
@@ -27,6 +27,8 @@ impl FalFlux1DevRequestState {
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: payload.request_id,
gateway_request_id: payload.gateway_request_id,
maybe_status_url: payload.status_url,
maybe_response_url: payload.response_url,
maybe_outbound_request: Some(outbound),
}))
}
@@ -36,6 +38,8 @@ impl FalFlux1DevRequestState {
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: payload.request_id,
gateway_request_id: payload.gateway_request_id,
maybe_status_url: payload.status_url,
maybe_response_url: payload.response_url,
maybe_outbound_request: Some(outbound),
}))
}
@@ -48,6 +52,8 @@ impl FalFlux1DevRequestState {
struct FalResponseIds {
request_id: Option<String>,
gateway_request_id: Option<String>,
status_url: Option<String>,
response_url: Option<String>,
}
async fn send_fal_request<T: FalEndpoint>(
@@ -57,20 +63,22 @@ async fn send_fal_request<T: FalEndpoint>(
if let Some(webhook_url) = &client.webhook_url {
let response = request
.send_webhook_request(&client.api_key, webhook_url)
.await
.map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?;
.await?;
Ok(FalResponseIds {
request_id: response.request_id,
gateway_request_id: response.gateway_request_id,
status_url: None,
response_url: None,
})
} else {
let response = request
.send_queue_request(&client.api_key)
.await
.map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?;
.await?;
Ok(FalResponseIds {
request_id: Some(response.request_id),
gateway_request_id: None,
status_url: Some(response.status_url),
response_url: Some(response.response_url),
})
}
}
@@ -27,6 +27,8 @@ impl FalFlux1SchnellRequestState {
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: payload.request_id,
gateway_request_id: payload.gateway_request_id,
maybe_status_url: payload.status_url,
maybe_response_url: payload.response_url,
maybe_outbound_request: Some(outbound),
}))
}
@@ -36,6 +38,8 @@ impl FalFlux1SchnellRequestState {
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: payload.request_id,
gateway_request_id: payload.gateway_request_id,
maybe_status_url: payload.status_url,
maybe_response_url: payload.response_url,
maybe_outbound_request: Some(outbound),
}))
}
@@ -48,6 +52,8 @@ impl FalFlux1SchnellRequestState {
struct FalResponseIds {
request_id: Option<String>,
gateway_request_id: Option<String>,
status_url: Option<String>,
response_url: Option<String>,
}
async fn send_fal_request<T: FalEndpoint>(
@@ -57,20 +63,22 @@ async fn send_fal_request<T: FalEndpoint>(
if let Some(webhook_url) = &client.webhook_url {
let response = request
.send_webhook_request(&client.api_key, webhook_url)
.await
.map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?;
.await?;
Ok(FalResponseIds {
request_id: response.request_id,
gateway_request_id: response.gateway_request_id,
status_url: None,
response_url: None,
})
} else {
let response = request
.send_queue_request(&client.api_key)
.await
.map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?;
.await?;
Ok(FalResponseIds {
request_id: Some(response.request_id),
gateway_request_id: None,
status_url: Some(response.status_url),
response_url: Some(response.response_url),
})
}
}
@@ -27,6 +27,8 @@ impl FalNanoBanana2RequestState {
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: payload.request_id,
gateway_request_id: payload.gateway_request_id,
maybe_status_url: payload.status_url,
maybe_response_url: payload.response_url,
maybe_outbound_request: Some(outbound),
}))
}
@@ -36,6 +38,8 @@ impl FalNanoBanana2RequestState {
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: payload.request_id,
gateway_request_id: payload.gateway_request_id,
maybe_status_url: payload.status_url,
maybe_response_url: payload.response_url,
maybe_outbound_request: Some(outbound),
}))
}
@@ -48,6 +52,8 @@ impl FalNanoBanana2RequestState {
struct FalResponseIds {
request_id: Option<String>,
gateway_request_id: Option<String>,
status_url: Option<String>,
response_url: Option<String>,
}
async fn send_fal_request<T: FalEndpoint>(
@@ -57,20 +63,22 @@ async fn send_fal_request<T: FalEndpoint>(
if let Some(webhook_url) = &client.webhook_url {
let response = request
.send_webhook_request(&client.api_key, webhook_url)
.await
.map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?;
.await?;
Ok(FalResponseIds {
request_id: response.request_id,
gateway_request_id: response.gateway_request_id,
status_url: None,
response_url: None,
})
} else {
let response = request
.send_queue_request(&client.api_key)
.await
.map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?;
.await?;
Ok(FalResponseIds {
request_id: Some(response.request_id),
gateway_request_id: None,
status_url: Some(response.status_url),
response_url: Some(response.response_url),
})
}
}
@@ -27,6 +27,8 @@ impl FalNanoBananaProRequestState {
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: payload.request_id,
gateway_request_id: payload.gateway_request_id,
maybe_status_url: payload.status_url,
maybe_response_url: payload.response_url,
maybe_outbound_request: Some(outbound),
}))
}
@@ -36,6 +38,8 @@ impl FalNanoBananaProRequestState {
Ok(GenerateImageResponse::Fal(FalImageResponsePayload {
request_id: payload.request_id,
gateway_request_id: payload.gateway_request_id,
maybe_status_url: payload.status_url,
maybe_response_url: payload.response_url,
maybe_outbound_request: Some(outbound),
}))
}
@@ -48,6 +52,8 @@ impl FalNanoBananaProRequestState {
struct FalResponseIds {
request_id: Option<String>,
gateway_request_id: Option<String>,
status_url: Option<String>,
response_url: Option<String>,
}
/// Send a FAL request via webhook (if URL present) or queue (if not).
@@ -58,20 +64,22 @@ async fn send_fal_request<T: FalEndpoint>(
if let Some(webhook_url) = &client.webhook_url {
let response = request
.send_webhook_request(&client.api_key, webhook_url)
.await
.map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?;
.await?;
Ok(FalResponseIds {
request_id: response.request_id,
gateway_request_id: response.gateway_request_id,
status_url: None,
response_url: None,
})
} else {
let response = request
.send_queue_request(&client.api_key)
.await
.map_err(|e| ArtcraftRouterError::Provider(ProviderError::Fal(e)))?;
.await?;
Ok(FalResponseIds {
request_id: Some(response.request_id),
gateway_request_id: None,
status_url: Some(response.status_url),
response_url: Some(response.response_url),
})
}
}
+1
View File
@@ -33,6 +33,7 @@ build_metadata.workspace = true
cookie_store.workspace = true
enums.workspace = true
errors.workspace = true
fal_client.workspace = true
filesys.workspace = true
grok_client.workspace = true
uuid_utils.workspace = true
@@ -124,5 +124,7 @@ pub async fn handle_artcraft_flux_kontext_edit(
model: Some(GenerationModel::FluxProKontextMax),
provider_job_id: Some(job_id.to_string()),
task_type: TaskType::ImageGeneration,
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -122,5 +122,7 @@ pub async fn handle_artcraft_gpt_image_1_edit(
model: Some(GenerationModel::GptImage1),
provider_job_id: Some(job_id.to_string()),
task_type: TaskType::ImageGeneration,
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -109,6 +109,8 @@ pub(super) async fn handle_image_edit_artcraft_via_router(
model: Some(generation_model),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_id),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -176,5 +176,7 @@ pub async fn handle_sora_gpt_image_1_edit(
model: Some(GenerationModel::GptImage1),
provider: GenerationProvider::Sora,
provider_job_id: Some(response.task_id.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -113,6 +113,8 @@ pub async fn handle_artcraft_flux_dev_juggernaut_inpaint(
model: Some(GenerationModel::FluxDevJuggernaut),
provider_job_id: Some(job_id.to_string()),
task_type: TaskType::ImageInpaintEdit,
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -112,6 +112,8 @@ pub async fn handle_artcraft_flux_pro_1_inpaint(
model: Some(GenerationModel::FluxPro1),
provider_job_id: Some(job_id.to_string()),
task_type: TaskType::ImageInpaintEdit,
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -125,5 +125,7 @@ pub async fn handle_artcraft_flux_pro_kontext_inpaint(
model: Some(GenerationModel::FluxProKontextMax),
provider_job_id: Some(job_id.to_string()),
task_type: TaskType::ImageInpaintEdit,
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -110,5 +110,7 @@ pub async fn handle_artcraft_nano_banana_inpaint(
model: Some(GenerationModel::Gemini25Flash),
provider_job_id: Some(job_id.to_string()),
task_type: TaskType::ImageInpaintEdit,
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -108,5 +108,7 @@ pub async fn handle_artcraft_gpt_image_1_text_to_image(
model: Some(GenerationModel::GptImage1),
provider_job_id: Some(job_id.to_string()),
task_type: TaskType::ImageGeneration,
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -87,6 +87,8 @@ pub(super) async fn handle_text_to_image_artcraft_via_router(
model: Some(generation_model),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_id),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -68,6 +68,8 @@ pub async fn handle_grok(
model: Some(GenerationModel::GrokImage),
provider_job_id: Some(job_id),
task_type: TaskType::ImageGeneration,
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -111,6 +111,8 @@ pub async fn handle_midjourney(
model: Some(GenerationModel::Midjourney),
provider_job_id: Some(job_id),
task_type: TaskType::ImageGeneration,
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -111,6 +111,8 @@ pub async fn handle_gpt_image_1_sora_text_to_image(
task_type: TaskType::ImageGeneration,
model: Some(GenerationModel::GptImage1),
provider_job_id: Some(response.task_id.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -28,10 +28,12 @@ fn billing_error(
reason: &BillingIssueReason,
) {
let provider = match reason.provider {
BillingProvider::Artcraft => GenerationProvider::Artcraft,
BillingProvider::Fal => GenerationProvider::Fal,
BillingProvider::Kinovi => GenerationProvider::Artcraft, // NB: We don't support Kinovi yet.
BillingProvider::Midjourney => GenerationProvider::Midjourney,
BillingProvider::Muapi=> GenerationProvider::Artcraft, // NB: We don't support Muapi yet.
BillingProvider::Sora => GenerationProvider::Sora,
BillingProvider::Storyteller => GenerationProvider::Artcraft,
};
warn!("Billing issue with: {:?}", provider);
ShowProviderBillingModalEvent::send_for_provider(provider, app);
@@ -93,10 +93,12 @@ pub struct BillingIssueReason {
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum BillingProvider {
Artcraft,
Fal,
Kinovi,
Midjourney,
Muapi,
Sora,
Storyteller,
}
#[derive(Debug)]
@@ -230,6 +232,15 @@ impl From<ArtcraftRouterError> for GenerateError {
ArtcraftRouterError::Client(e) => Self::BadInput(BadInputReason::WrongImageArguments(e.to_string())),
ArtcraftRouterError::Download(e) => Self::ArtcraftRouterDownloadError(e),
ArtcraftRouterError::InvalidInput(msg) => Self::BadInput(BadInputReason::WrongImageArguments(msg)),
ArtcraftRouterError::ProviderBillingError(err) => {
let provider = match err {
ProviderError::Fal(_) => BillingProvider::Fal,
ProviderError::Muapi(_) => BillingProvider::Muapi,
ProviderError::Seedance2Pro(_) => BillingProvider::Kinovi,
ProviderError::Storyteller(_) => BillingProvider::Artcraft,
};
Self::BillingIssue(BillingIssueReason { provider })
},
ArtcraftRouterError::Provider(ProviderError::Storyteller(e)) => Self::ProviderFailure(ProviderFailureReason::StorytellerError(e)),
ArtcraftRouterError::Provider(ProviderError::Fal(_)) => Self::FalNoLongerSupported,
ArtcraftRouterError::Provider(ProviderError::Muapi(_)) => Self::ArtcraftRouterNotYetSupportedProvider("muapi"),
@@ -79,6 +79,8 @@ pub async fn handle_generic_bg_removal_artcraft(
task_type: TaskType::BackgroundRemoval,
provider_job_id: Some(success_result.inference_job_token.to_string()),
model: None,
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -58,5 +58,7 @@ pub(super) async fn handle_artcraft_splat_via_router(
model: Some(generation_model),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_id),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -102,6 +102,8 @@ pub async fn handle_worldlabs_marble(
model: Some(GenerationModel::WorldlabsMarble),
provider_job_id: Some(world_id.0),
task_type: TaskType::ImageGeneration,
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -152,5 +152,7 @@ pub async fn handle_artcraft_object(
model: used_model,
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_token.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -14,7 +14,8 @@ pub struct TaskEnqueueSuccess {
pub model: Option<GenerationModel>,
pub provider: GenerationProvider,
pub provider_job_id: Option<String>,
// TODO: We may want to change the `model` type - this has weird ownership and semantics
pub maybe_queue_status_url: Option<String>,
pub maybe_queue_response_url: Option<String>,
}
impl TaskEnqueueSuccess{
@@ -121,6 +122,8 @@ impl TaskEnqueueSuccess{
model_type,
provider: self.provider,
provider_job_id: self.provider_job_id.as_deref(),
queue_status_url: self.maybe_queue_status_url.as_deref(),
queue_response_url: self.maybe_queue_response_url.as_deref(),
frontend_caller,
frontend_subscriber_id,
frontend_subscriber_payload,
@@ -1,7 +1,8 @@
use crate::core::commands::enqueue::common::notify_frontend_of_errors::notify_frontend_of_errors;
use crate::core::commands::enqueue::generate_error::GenerateError;
use crate::core::commands::enqueue::generate_error::{GenerateError, MissingCredentialsReason};
use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
use crate::core::commands::generate::generate_image::providers::artcraft::handle_artcraft;
use crate::core::commands::generate::generate_image::providers::router::handle_router::handle_router;
use crate::core::commands::generate::generate_image::tauri_generate_image_request::{
TauriGenerateImageErrorType, TauriGenerateImageRequest, TauriGenerateImageResponse,
};
@@ -10,6 +11,7 @@ use crate::core::commands::response::shorthand::Response;
use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
use crate::core::events::functional_events::credits_balance_changed_event::CreditsBalanceChangedEvent;
use crate::core::events::generation_events::generation_enqueue_success_event::GenerationEnqueueSuccessEvent;
use crate::core::providers::credentials::provider_credential_loading_cache::ProviderCredentialLoadingCache;
use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
use crate::core::state::task_database::TaskDatabase;
use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
@@ -24,32 +26,37 @@ pub async fn generate_image_command(
app_env_configs: State<'_, AppEnvConfigs>,
task_database: State<'_, TaskDatabase>,
storyteller_creds_manager: State<'_, StorytellerCredentialManager>,
credential_cache: State<'_, ProviderCredentialLoadingCache>,
) -> Response<TauriGenerateImageResponse, TauriGenerateImageErrorType, ()> {
info!("generate_image_command called, request: {:?}", request);
// Only Artcraft provider is supported via this command.
// Other providers (Grok, Sora, Midjourney) should use the legacy commands.
let provider = request.provider.unwrap_or(GenerationProvider::Artcraft);
let is_artcraft_provider = matches!(
request.provider,
None | Some(GenerationProvider::Artcraft)
);
if !is_artcraft_provider {
return Err(CommandErrorResponseWrapper {
status: CommandErrorStatus::BadRequest,
error_message: Some("Only Artcraft is supported.".to_string()),
error_type: Some(TauriGenerateImageErrorType::NoProviderAvailable),
error_details: None,
});
}
let result = handle_artcraft(
&request,
&app_env_configs,
&storyteller_creds_manager,
).await;
let result = match provider {
GenerationProvider::Artcraft => {
handle_artcraft(
&request,
&app_env_configs,
&storyteller_creds_manager,
).await
}
// Midjourney uses its own legacy command path, not this one.
GenerationProvider::Midjourney => {
Err(GenerateError::NotYetImplemented(
"Midjourney should use its dedicated command".to_string(),
))
}
// All other providers go through the router.
other => {
handle_router(
&request,
other,
&app_env_configs,
&credential_cache,
).await
}
};
match result {
Ok(success) => handle_success_behavior(&app, &task_database, &request, success).await,
@@ -119,6 +126,8 @@ async fn handle_error_behavior(
let error_type = match &err {
GenerateError::BadInput(_) => TauriGenerateImageErrorType::BadInput,
GenerateError::MissingCredentials(MissingCredentialsReason::NeedsFalApiKey) => TauriGenerateImageErrorType::NeedsFalApiKey,
GenerateError::MissingCredentials(MissingCredentialsReason::NeedsGrokCredentials) => TauriGenerateImageErrorType::NeedsGrokCredentials,
GenerateError::MissingCredentials(_) => TauriGenerateImageErrorType::NeedsStorytellerCredentials,
GenerateError::NoProviderAvailable => TauriGenerateImageErrorType::NoProviderAvailable,
GenerateError::BillingIssue(_) => TauriGenerateImageErrorType::BillingIssue,
@@ -65,6 +65,8 @@ pub async fn handle_artcraft_via_omni_endpoint(
model: Some(generation_model),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(response.inference_job_token.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -70,6 +70,8 @@ pub async fn handle_flux_2_lora_angles(
model: Some(GenerationModel::Flux2LoraAngles),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_id),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -76,6 +76,8 @@ pub async fn handle_flux_dev_juggernaut_inpaint(
model: Some(GenerationModel::FluxDevJuggernaut),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(response.inference_job_token.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -74,6 +74,8 @@ pub async fn handle_flux_pro_kontext_edit(
model: Some(GenerationModel::FluxProKontextMax),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(response.inference_job_token.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -70,6 +70,8 @@ pub async fn handle_qwen_edit_2511_angles(
model: Some(GenerationModel::QwenEdit2511Angles),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_id),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -1 +1,2 @@
pub mod artcraft;
pub mod artcraft;
pub mod router;
@@ -0,0 +1,159 @@
use artcraft_client::utils::api_host::ApiHost;
use artcraft_router::api::image_list_ref::ImageListRef;
use artcraft_router::api::provider::Provider;
use artcraft_router::client::generation_mode_mismatch_strategy::GenerationModeMismatchStrategy;
use artcraft_router::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy;
use artcraft_router::client::router_client::RouterClient;
use artcraft_router::client::router_fal_client::RouterFalClient;
use artcraft_router::client::router_fal_webhook_optional_client::RouterFalWebhookOptionalClient;
use artcraft_router::generate::generate_image::generate_image_request_builder::GenerateImageRequestBuilder;
use artcraft_router::generate::generate_image::generate_image_response::GenerateImageResponse;
use artcraft_router::generate::generate_image_v2::image_generation_draft_or_request::ImageGenerationDraftOrRequest;
use enums::common::generation_provider::GenerationProvider;
use enums::tauri::tasks::task_type::TaskType;
use log::{info, warn};
use tokens::tokens::media_files::MediaFileToken;
use crate::core::api_adapters::models::image::tauri_image_model_to_generation_model::tauri_image_model_to_generation_model;
use crate::core::api_adapters::models::image::tauri_image_model_to_router_model::tauri_image_model_to_router_model;
use crate::core::commands::enqueue::generate_error::GenerateError;
use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
use crate::core::commands::generate::generate_image::providers::router::utils::convert_enums_to_router::{convert_aspect_ratio, convert_quality, convert_resolution};
use crate::core::commands::generate::generate_image::providers::router::utils::map_media_files_to_urls::map_media_file_tokens_to_cdn_urls;
use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest;
use crate::core::commands::generate::generate_image::tauri_image_model::TauriImageModel;
use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
/// Handle image generation for providers that authenticate via API key.
pub async fn handle_api_key_provider(
request: &TauriGenerateImageRequest,
provider: GenerationProvider,
api_key: &str,
app_env_configs: &AppEnvConfigs,
) -> Result<TaskEnqueueSuccess, GenerateError> {
match provider {
GenerationProvider::Fal => {
handle_fal(request, api_key, &app_env_configs.storyteller_host).await
}
_ => {
Err(GenerateError::NotYetImplemented(
format!("API key provider {:?} is not yet supported via the router path", provider),
))
}
}
}
// ── FAL ──
async fn handle_fal(
request: &TauriGenerateImageRequest,
api_key: &str,
api_host: &ApiHost,
) -> Result<TaskEnqueueSuccess, GenerateError> {
let tauri_model = request.model.ok_or(GenerateError::no_model_specified())?;
let router_model = tauri_image_model_to_router_model(tauri_model)
.ok_or(GenerateError::NotYetImplemented(
format!("Model {:?} is not supported via the FAL router path", tauri_model),
))?;
// Collect all media file tokens that need resolving.
let image_inputs = resolve_image_inputs(request, api_host).await?;
let router_request = GenerateImageRequestBuilder {
model: router_model,
provider: Provider::Fal,
prompt: request.prompt.clone(),
image_inputs,
resolution: request.resolution.map(convert_resolution),
aspect_ratio: request.aspect_ratio.map(convert_aspect_ratio),
quality: request.quality.map(convert_quality),
image_batch_count: request.batch_size.map(|n| n as u16),
horizontal_angle: request.adjust_horizontal_angle,
vertical_angle: request.adjust_vertical_angle,
zoom: request.adjust_zoom,
request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::PayMoreUpgrade,
generation_mode_mismatch_strategy: Some(GenerationModeMismatchStrategy::GenerateAnyway),
idempotency_token: None,
};
let fal_client = RouterFalWebhookOptionalClient::from_str(api_key);
let client = RouterClient::FalWebhookOptional(fal_client);
info!("Building FAL image generation plan: model={:?}", router_model);
let request = match router_request.build2() {
Ok(ImageGenerationDraftOrRequest::Request(request)) => request,
Ok(ImageGenerationDraftOrRequest::Draft(draft)) => {
warn!("Fal is trying to send draft request: {:?}", draft);
return Err(GenerateError::NotYetImplemented("Fal should not be sending draft requests".to_string()));
},
Err(err) => {
warn!("Could not use FAL: {:?}", err);
return Err(GenerateError::NotYetImplemented("Error Message: TODO".to_string()));
}
};
info!("Executing FAL image generation. Request: {:?}", request);
match request.send_request(&client).await {
Ok(response) => {
build_task_enqueue_success(tauri_model, response)
},
Err(err) => {
warn!("Fal image generation failed: {:?}", err);
Err(err.into())
}
}
}
// ── Helpers ──
async fn resolve_image_inputs(
request: &TauriGenerateImageRequest,
api_host: &ApiHost,
) -> Result<Option<ImageListRef>, GenerateError> {
let mut tokens: Vec<MediaFileToken> = Vec::new();
if let Some(canvas_token) = &request.canvas_image_media_token {
tokens.push(canvas_token.clone());
}
if let Some(scene_token) = &request.scene_image_media_token {
tokens.push(scene_token.clone());
}
if let Some(media_tokens) = &request.image_media_tokens {
tokens.extend(media_tokens.clone());
}
if tokens.is_empty() {
return Ok(None);
}
let urls = map_media_file_tokens_to_cdn_urls(&tokens, api_host).await?;
Ok(Some(ImageListRef::Urls(urls)))
}
fn build_task_enqueue_success(
tauri_model: TauriImageModel,
response: GenerateImageResponse,
) -> Result<TaskEnqueueSuccess, GenerateError> {
let fal_payload = response.get_fal_payload()
.ok_or(GenerateError::ResponseHadNoJobTokens)?;
let provider_job_id = fal_payload.request_id
.or(fal_payload.gateway_request_id)
.ok_or(GenerateError::ResponseHadNoJobTokens)?;
let generation_model = tauri_image_model_to_generation_model(tauri_model);
Ok(TaskEnqueueSuccess {
task_type: TaskType::ImageGeneration,
model: Some(generation_model),
provider: GenerationProvider::Fal,
provider_job_id: Some(provider_job_id),
maybe_queue_status_url: fal_payload.maybe_status_url,
maybe_queue_response_url: fal_payload.maybe_response_url,
})
}
@@ -0,0 +1,66 @@
use enums::common::generation_provider::GenerationProvider;
use log::info;
use crate::core::commands::enqueue::generate_error::{GenerateError, MissingCredentialsReason};
use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
use crate::core::commands::generate::generate_image::providers::router::handle_api_providers::handle_api_key_provider;
use crate::core::commands::generate::generate_image::providers::router::handle_web_login_providers::handle_web_login_provider;
use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest;
use crate::core::providers::credentials::payload::provider_credential_payload::ProviderCredentialPayload;
use crate::core::providers::credentials::provider_credential_key::ProviderCredentialKey;
use crate::core::providers::credentials::provider_credential_loading_cache::ProviderCredentialLoadingCache;
use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
/// Dispatch an image generation request to a third-party provider via the router.
///
/// This handles providers that authenticate with the user's own credentials
/// (API key or web login) rather than going through the Artcraft backend.
pub async fn handle_router(
request: &TauriGenerateImageRequest,
provider: GenerationProvider,
app_env_configs: &AppEnvConfigs,
credential_cache: &ProviderCredentialLoadingCache,
) -> Result<TaskEnqueueSuccess, GenerateError> {
let credential_key = map_provider_to_credential_key(provider)?;
info!("handle_router: provider={:?}, credential_key={:?}", provider, credential_key);
let payload = credential_cache.get_credentials(credential_key)
.map_err(|err| {
GenerateError::AnyhowError(anyhow::anyhow!("Failed to load credentials: {:?}", err))
})?
.ok_or_else(|| map_provider_to_missing_credentials_error(provider))?;
match payload {
ProviderCredentialPayload::ApiKey(api_key_data) => {
handle_api_key_provider(request, provider, api_key_data.as_str(), app_env_configs).await
}
ProviderCredentialPayload::WebLogin(web_login_data) => {
handle_web_login_provider(request, provider, &web_login_data).await
}
}
}
// ── Helpers ──
fn map_provider_to_credential_key(
provider: GenerationProvider,
) -> Result<ProviderCredentialKey, GenerateError> {
match provider {
GenerationProvider::Fal => Ok(ProviderCredentialKey::FalApiKey),
_ => Err(GenerateError::NotYetImplemented(
format!("Provider {:?} does not have a mapped credential key yet", provider),
)),
}
}
fn map_provider_to_missing_credentials_error(provider: GenerationProvider) -> GenerateError {
match provider {
GenerationProvider::Fal => GenerateError::MissingCredentials(MissingCredentialsReason::NeedsFalApiKey),
GenerationProvider::Grok => GenerateError::MissingCredentials(MissingCredentialsReason::NeedsGrokCredentials),
GenerationProvider::Midjourney => GenerateError::MissingCredentials(MissingCredentialsReason::NeedsMidjourneyCredentials),
GenerationProvider::Sora => GenerateError::MissingCredentials(MissingCredentialsReason::NeedsSoraCredentials),
GenerationProvider::WorldLabs => GenerateError::MissingCredentials(MissingCredentialsReason::NeedsWorldLabsCredentials),
_ => GenerateError::NoProviderAvailable,
}
}
@@ -0,0 +1,17 @@
use enums::common::generation_provider::GenerationProvider;
use crate::core::commands::enqueue::generate_error::GenerateError;
use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest;
use crate::core::providers::credentials::payload::web_login::WebLoginData;
/// Handle image generation for providers that authenticate via web login.
///
/// Not yet implemented — placeholder for Grok, Runway, Higgsfield, etc.
pub async fn handle_web_login_provider(
_request: &TauriGenerateImageRequest,
_provider: GenerationProvider,
_web_login: &WebLoginData,
) -> Result<TaskEnqueueSuccess, GenerateError> {
unimplemented!("Web login provider image generation is not yet supported via the router path")
}
@@ -0,0 +1,4 @@
pub mod handle_router;
mod handle_api_providers;
mod handle_web_login_providers;
mod utils;
@@ -0,0 +1,50 @@
use artcraft_router::api::common_aspect_ratio::CommonAspectRatio as RouterAspectRatio;
use artcraft_router::api::common_quality::CommonQuality as RouterQuality;
use artcraft_router::api::common_resolution::CommonResolution as RouterResolution;
use enums::common::generation::common_aspect_ratio::CommonAspectRatio as EnumsAspectRatio;
use enums::common::generation::common_quality::CommonQuality as EnumsQuality;
use enums::common::generation::common_resolution::CommonResolution as EnumsResolution;
pub fn convert_aspect_ratio(value: EnumsAspectRatio) -> RouterAspectRatio {
match value {
EnumsAspectRatio::Auto => RouterAspectRatio::Auto,
EnumsAspectRatio::Square => RouterAspectRatio::Square,
EnumsAspectRatio::WideThreeByTwo => RouterAspectRatio::WideThreeByTwo,
EnumsAspectRatio::WideFourByThree => RouterAspectRatio::WideFourByThree,
EnumsAspectRatio::WideFiveByFour => RouterAspectRatio::WideFiveByFour,
EnumsAspectRatio::WideSixteenByNine => RouterAspectRatio::WideSixteenByNine,
EnumsAspectRatio::WideTwentyOneByNine => RouterAspectRatio::WideTwentyOneByNine,
EnumsAspectRatio::TallTwoByThree => RouterAspectRatio::TallTwoByThree,
EnumsAspectRatio::TallThreeByFour => RouterAspectRatio::TallThreeByFour,
EnumsAspectRatio::TallFourByFive => RouterAspectRatio::TallFourByFive,
EnumsAspectRatio::TallNineBySixteen => RouterAspectRatio::TallNineBySixteen,
EnumsAspectRatio::TallNineByTwentyOne => RouterAspectRatio::TallNineByTwentyOne,
EnumsAspectRatio::Wide => RouterAspectRatio::Wide,
EnumsAspectRatio::Tall => RouterAspectRatio::Tall,
EnumsAspectRatio::Auto2k => RouterAspectRatio::Auto2k,
EnumsAspectRatio::Auto3k => RouterAspectRatio::Auto3k,
EnumsAspectRatio::Auto4k => RouterAspectRatio::Auto4k,
EnumsAspectRatio::SquareHd => RouterAspectRatio::SquareHd,
}
}
pub fn convert_quality(value: EnumsQuality) -> RouterQuality {
match value {
EnumsQuality::High => RouterQuality::High,
EnumsQuality::Medium => RouterQuality::Medium,
EnumsQuality::Low => RouterQuality::Low,
}
}
pub fn convert_resolution(value: EnumsResolution) -> RouterResolution {
match value {
EnumsResolution::OneK => RouterResolution::OneK,
EnumsResolution::TwoK => RouterResolution::TwoK,
EnumsResolution::ThreeK => RouterResolution::ThreeK,
EnumsResolution::FourK => RouterResolution::FourK,
EnumsResolution::HalfK => RouterResolution::HalfK,
EnumsResolution::FourEightyP => RouterResolution::FourEightyP,
EnumsResolution::SevenTwentyP => RouterResolution::SevenTwentyP,
EnumsResolution::TenEightyP => RouterResolution::TenEightyP,
}
}
@@ -0,0 +1,27 @@
use artcraft_client::endpoints::media_files::get_media_file::get_media_file;
use artcraft_client::utils::api_host::ApiHost;
use log::info;
use tokens::tokens::media_files::MediaFileToken;
use crate::core::commands::enqueue::generate_error::GenerateError;
/// Resolve a list of media file tokens to their CDN URLs.
///
/// FAL (and other third-party providers) can fetch images by URL directly,
/// so we just need the CDN link — no need to download the bytes locally.
pub async fn map_media_file_tokens_to_cdn_urls(
tokens: &[MediaFileToken],
api_host: &ApiHost,
) -> Result<Vec<String>, GenerateError> {
let mut urls = Vec::with_capacity(tokens.len());
for token in tokens {
info!("Resolving media file token {} to CDN URL...", token.as_str());
let response = get_media_file(api_host, token).await?;
let cdn_url = response.media_file.media_links.cdn_url.to_string();
info!("Resolved {} -> {}", token.as_str(), &cdn_url);
urls.push(cdn_url);
}
Ok(urls)
}
@@ -0,0 +1,2 @@
pub mod convert_enums_to_router;
pub mod map_media_files_to_urls;
@@ -116,6 +116,8 @@ pub enum TauriGenerateImageErrorType {
ServerError,
/// Needs to be logged into Artcraft
NeedsStorytellerCredentials,
/// Needs a FAL API key
NeedsFalApiKey,
/// Needs Grok credentials
NeedsGrokCredentials,
/// Billing issue
@@ -71,5 +71,7 @@ pub (super) async fn handle_artcraft_kling_1p6_pro(
model: Some(GenerationModel::Kling1_6),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_token.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -70,5 +70,7 @@ pub (super) async fn handle_artcraft_kling_2p1_master(
model: Some(GenerationModel::Kling21Master),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_token.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -70,5 +70,7 @@ pub (super) async fn handle_artcraft_kling_2p1_pro(
model: Some(GenerationModel::Kling21Pro),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_token.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -70,5 +70,7 @@ pub (super) async fn handle_artcraft_kling_2p5_turbo_pro(
model: Some(GenerationModel::Kling2p5TurboPro),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_token.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -72,5 +72,7 @@ pub (super) async fn handle_artcraft_kling_2p6_pro(
model: Some(GenerationModel::Kling2p6Pro),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_token.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -66,5 +66,7 @@ pub (super) async fn handle_artcraft_seedance_1_lite(
model: Some(GenerationModel::Kling1_6),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_token.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -72,5 +72,7 @@ pub (super) async fn handle_artcraft_sora2(
model: Some(GenerationModel::Sora2),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_token.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -71,5 +71,7 @@ pub (super) async fn handle_artcraft_sora2_pro(
model: Some(GenerationModel::Sora2Pro),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_token.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -69,5 +69,7 @@ pub (super) async fn handle_artcraft_veo2(
model: Some(GenerationModel::Veo2),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_token.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -71,5 +71,7 @@ pub (super) async fn handle_artcraft_veo3(
model: Some(GenerationModel::Veo3),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_token.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -66,5 +66,7 @@ pub (super) async fn handle_artcraft_veo3_fast(
model: Some(GenerationModel::Veo3Fast),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_token.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -74,5 +74,7 @@ pub (super) async fn handle_artcraft_veo3p1(
model: Some(GenerationModel::Veo3p1),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_token.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -74,5 +74,7 @@ pub (super) async fn handle_artcraft_veo3p1_fast(
model: Some(GenerationModel::Veo3p1Fast),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_token.to_string()),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -79,6 +79,8 @@ pub(super) async fn handle_artcraft_video_via_router(
model: Some(generation_model),
provider: GenerationProvider::Artcraft,
provider_job_id: Some(job_id),
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -96,6 +96,8 @@ pub async fn handle_grok_video(
model: Some(GenerationModel::GrokVideo),
provider_job_id: Some(post_id.to_string()),
task_type: TaskType::VideoGeneration,
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -98,5 +98,7 @@ pub async fn handle_sora_sora2(
model: Some(GenerationModel::Sora2),
provider_job_id: Some(job_id.to_string()),
task_type: TaskType::VideoGeneration,
maybe_queue_status_url: None,
maybe_queue_response_url: None,
})
}
@@ -1,2 +1,3 @@
pub mod generate_image;
pub mod generate_video;
pub mod common;
@@ -0,0 +1,2 @@
pub mod get_provider_order_command;
pub mod set_provider_order_command;
@@ -1,2 +1,4 @@
pub mod get_provider_order_command;
pub mod set_provider_order_command;
pub mod deprecated;
pub mod provider_list_command;
pub mod provider_set_api_key_command;
pub mod provider_clear_command;
@@ -0,0 +1,33 @@
use crate::core::commands::response::shorthand::ResponseOrErrorMessage;
use crate::core::commands::response::success_response_wrapper::SerializeMarker;
use crate::core::providers::credentials::provider_credential_key::ProviderCredentialKey;
use crate::core::providers::credentials::provider_credential_loading_cache::ProviderCredentialLoadingCache;
use log::{error, info};
use serde_derive::{Deserialize, Serialize};
use tauri::State;
#[derive(Deserialize)]
pub struct ProviderClearRequest {
pub provider_credential: ProviderCredentialKey,
}
#[derive(Serialize)]
pub struct ProviderClearResponse {}
impl SerializeMarker for ProviderClearResponse {}
#[tauri::command]
pub async fn provider_clear_command(
request: ProviderClearRequest,
credential_cache: State<'_, ProviderCredentialLoadingCache>,
) -> ResponseOrErrorMessage<ProviderClearResponse> {
info!("provider_clear_command called for key: {:?}", request.provider_credential);
credential_cache.delete_credentials(request.provider_credential)
.map_err(|err| {
error!("Failed to clear provider credential: {:?}", err);
"Failed to clear provider credential"
})?;
Ok(ProviderClearResponse {}.into())
}
@@ -0,0 +1,106 @@
use crate::core::commands::response::shorthand::ResponseOrErrorMessage;
use crate::core::commands::response::success_response_wrapper::SerializeMarker;
use crate::core::providers::credentials::payload::provider_credential_payload::ProviderCredentialPayload;
use crate::core::providers::credentials::provider_credential_key::ProviderCredentialKey;
use crate::core::providers::credentials::provider_credential_loading_cache::ProviderCredentialLoadingCache;
use crate::core::providers::credentials::provider_credential_type::ProviderCredentialType;
use log::{info, warn};
use serde_derive::Serialize;
use tauri::State;
/// All known credential keys. Add new ones here as providers are added.
const ALL_KEYS: &[ProviderCredentialKey] = &[
ProviderCredentialKey::FalApiKey,
ProviderCredentialKey::ReplicateApiKey,
ProviderCredentialKey::GrokWebLogin,
ProviderCredentialKey::HiggsfieldWebLogin,
ProviderCredentialKey::MidjourneyLogin,
ProviderCredentialKey::RunwayWebLogin,
];
const REDACTED_KEY_VISIBLE_CHARS: usize = 6;
#[derive(Serialize)]
pub struct ProviderListResponse {
pub providers: Vec<ProviderListEntry>,
}
impl SerializeMarker for ProviderListResponse {}
#[derive(Serialize)]
pub struct ProviderListEntry {
pub provider_credential: ProviderCredentialKey,
pub credential_type: ProviderCredentialType,
pub has_credentials: bool,
pub maybe_details: Option<ProviderCredentialDetails>,
}
#[derive(Serialize)]
pub struct ProviderCredentialDetails {
/// For API keys: the first few characters followed by asterisks.
pub maybe_key_start: Option<String>,
/// For API keys: the full key value.
pub maybe_full_key: Option<String>,
/// For web logins: the email address if available.
pub maybe_email_address: Option<String>,
/// For web logins: the username if available.
pub maybe_username: Option<String>,
}
#[tauri::command]
pub async fn provider_list_command(
credential_cache: State<'_, ProviderCredentialLoadingCache>,
) -> ResponseOrErrorMessage<ProviderListResponse> {
info!("provider_list_command called");
let mut providers = Vec::with_capacity(ALL_KEYS.len());
for &key in ALL_KEYS {
let maybe_payload = match credential_cache.get_credentials(key) {
Ok(payload) => payload,
Err(err) => {
warn!("Error checking credential for {:?}: {:?}", key, err);
None
}
};
let has_credentials = maybe_payload.is_some();
let maybe_details = maybe_payload.map(|payload| match payload {
ProviderCredentialPayload::ApiKey(data) => {
ProviderCredentialDetails {
maybe_key_start: Some(redact_key(data.as_str())),
maybe_full_key: Some(data.as_str().to_string()),
maybe_email_address: None,
maybe_username: None,
}
}
ProviderCredentialPayload::WebLogin(data) => {
ProviderCredentialDetails {
maybe_key_start: None,
maybe_full_key: None,
maybe_email_address: data.email_address,
maybe_username: data.username,
}
}
});
providers.push(ProviderListEntry {
provider_credential: key,
credential_type: key.get_type(),
has_credentials,
maybe_details,
});
}
Ok(ProviderListResponse { providers }.into())
}
fn redact_key(key: &str) -> String {
if key.len() <= REDACTED_KEY_VISIBLE_CHARS {
"********".to_string()
} else {
let visible = &key[..REDACTED_KEY_VISIBLE_CHARS];
format!("{}********", visible)
}
}
@@ -0,0 +1,39 @@
use crate::core::commands::response::shorthand::ResponseOrErrorMessage;
use crate::core::commands::response::success_response_wrapper::SerializeMarker;
use crate::core::providers::credentials::payload::api_key::ApiKeyData;
use crate::core::providers::credentials::payload::provider_credential_payload::ProviderCredentialPayload;
use crate::core::providers::credentials::provider_credential_key::ProviderCredentialKey;
use crate::core::providers::credentials::provider_credential_loading_cache::ProviderCredentialLoadingCache;
use log::{error, info};
use serde_derive::{Deserialize, Serialize};
use tauri::State;
#[derive(Deserialize)]
pub struct ProviderSetApiKeyRequest {
pub provider_credential: ProviderCredentialKey,
pub api_key: String,
}
#[derive(Serialize)]
pub struct ProviderSetApiKeyResponse {}
impl SerializeMarker for ProviderSetApiKeyResponse {}
#[tauri::command]
pub async fn provider_set_api_key_command(
request: ProviderSetApiKeyRequest,
credential_cache: State<'_, ProviderCredentialLoadingCache>,
) -> ResponseOrErrorMessage<ProviderSetApiKeyResponse> {
info!("provider_set_api_key_command called for provider: {:?}", request.provider_credential);
let api_key_data = ApiKeyData::from_str(&request.api_key);
let payload = ProviderCredentialPayload::ApiKey(api_key_data);
credential_cache.save_credentials(request.provider_credential, payload)
.map_err(|err| {
error!("Failed to save API key: {:?}", err);
"Failed to save API key"
})?;
Ok(ProviderSetApiKeyResponse {}.into())
}
@@ -21,6 +21,8 @@ use crate::services::sora::state::sora_task_queue::SoraTaskQueue;
use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
use crate::services::worldlabs::state::worldlabs_bearer_bridge::WorldlabsBearerBridge;
use crate::services::worldlabs::state::worldlabs_credential_manager::WorldlabsCredentialManager;
use crate::core::providers::credentials::provider_credential_loading_cache::ProviderCredentialLoadingCache;
use crate::core::threads::third_party_task_polling_thread::third_party_task_polling_thread::third_party_task_polling_thread;
use crate::services::worldlabs::threads::worldlabs_marble_task_polling::worldlabs_marble_task_polling;
use errors::AnyhowResult;
use tauri::{AppHandle, Manager};
@@ -39,6 +41,7 @@ pub async fn handle_tauri_startup(
grok_image_prompt_queue: GrokImagePromptQueue,
worldlabs_bearer_bridge: WorldlabsBearerBridge,
worldlabs_creds_manager: WorldlabsCredentialManager,
credential_cache: ProviderCredentialLoadingCache,
) -> AnyhowResult<()> {
set_app_log_level(
@@ -116,6 +119,15 @@ pub async fn handle_tauri_startup(
storyteller_creds_manager.clone(),
));
tauri::async_runtime::spawn(third_party_task_polling_thread(
app.clone(),
app_env_configs.clone(),
root.clone(),
task_database.clone(),
storyteller_creds_manager.clone(),
credential_cache,
));
spawn_discord_presence_thread()?;
initially_size_and_position_windows(&app, &root);
+1
View File
@@ -3,6 +3,7 @@ pub mod artcraft_error;
pub mod commands;
pub mod events;
pub mod lifecycle;
pub mod providers;
pub mod state;
pub mod threads;
pub mod utils;
@@ -0,0 +1,4 @@
pub mod payload;
pub mod provider_credential_key;
pub mod provider_credential_loading_cache;
pub mod provider_credential_type;
@@ -0,0 +1,100 @@
use std::error::Error;
use std::fmt::Display;
use std::path::Path;
#[derive(Clone)]
pub struct ApiKeyData(String);
impl ApiKeyData {
pub fn from_str(s: &str) -> Self {
Self(s.trim().to_string())
}
pub fn as_str(&self) -> &str {
self.0.as_str()
}
pub fn load_from_file<P: AsRef<Path>>(file_path: P) -> Result<Self, ApiKeyDataError> {
let contents = std::fs::read_to_string(file_path)
.map_err(ApiKeyDataError::IoError)?;
Ok(Self(contents.trim().to_string()))
}
pub fn save_to_file<P: AsRef<Path>>(&self, file_path: P) -> Result<(), ApiKeyDataError> {
std::fs::write(file_path, self.0.trim())
.map_err(ApiKeyDataError::IoError)?;
Ok(())
}
}
#[derive(Debug)]
pub enum ApiKeyDataError {
IoError(std::io::Error),
}
impl Error for ApiKeyDataError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
ApiKeyDataError::IoError(e) => Some(e),
}
}
}
impl Display for ApiKeyDataError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ApiKeyDataError::IoError(e) => write!(f, "IO error: {}", e),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::NamedTempFile;
use std::io::Write;
#[test]
fn load_trims_whitespace() {
let mut file = NamedTempFile::new().unwrap();
write!(file, " sk-abc123 \n").unwrap();
let data = ApiKeyData::load_from_file(file.path()).unwrap();
assert_eq!(data.as_str(), "sk-abc123");
}
#[test]
fn load_trims_newlines() {
let mut file = NamedTempFile::new().unwrap();
write!(file, "\nmy_secret_key\n\n").unwrap();
let data = ApiKeyData::load_from_file(file.path()).unwrap();
assert_eq!(data.as_str(), "my_secret_key");
}
#[test]
fn save_trims_whitespace() {
let file = NamedTempFile::new().unwrap();
let data = ApiKeyData::from_str(" sk-padded ");
data.save_to_file(file.path()).unwrap();
let contents = std::fs::read_to_string(file.path()).unwrap();
assert_eq!(contents, "sk-padded");
}
#[test]
fn round_trip() {
let file = NamedTempFile::new().unwrap();
let original = ApiKeyData::from_str("sk-round-trip-test-key");
original.save_to_file(file.path()).unwrap();
let loaded = ApiKeyData::load_from_file(file.path()).unwrap();
assert_eq!(loaded.as_str(), "sk-round-trip-test-key");
}
#[test]
fn load_nonexistent_file_returns_io_error() {
let result = ApiKeyData::load_from_file("/tmp/nonexistent_api_key_file_12345.txt");
assert!(result.is_err());
}
}
@@ -0,0 +1,3 @@
pub mod provider_credential_payload;
pub mod api_key;
pub mod web_login;
@@ -0,0 +1,10 @@
use crate::core::providers::credentials::payload::api_key::ApiKeyData;
use crate::core::providers::credentials::payload::web_login::WebLoginData;
#[derive(Clone)]
pub enum ProviderCredentialPayload {
ApiKey(ApiKeyData),
// TODO: There might be logins in the future that use weird header states, etc.
WebLogin(WebLoginData),
}
@@ -0,0 +1,251 @@
use chrono::{DateTime, Utc};
use serde_derive::{Deserialize, Serialize};
use std::collections::HashMap;
use std::error::Error;
use std::fmt::Display;
use std::path::Path;
const CURRENT_VERSION: u8 = 1;
#[derive(Clone)]
pub struct WebLoginData {
pub cookies_header: Option<String>,
pub additional_headers: Option<HashMap<String, String>>,
pub username: Option<String>,
pub email_address: Option<String>,
pub created_at: Option<DateTime<Utc>>,
pub updated_at: Option<DateTime<Utc>>,
}
impl WebLoginData {
pub fn new() -> Self {
Self {
cookies_header: None,
additional_headers: None,
username: None,
email_address: None,
created_at: None,
updated_at: None,
}
}
pub fn load_from_file<P: AsRef<Path>>(file_path: P) -> Result<Self, WebLoginDataError> {
let contents = std::fs::read_to_string(file_path)
.map_err(WebLoginDataError::IoError)?;
let serializable: WebLoginDataSerializable = serde_json::from_str(&contents)
.map_err(WebLoginDataError::DeserializeError)?;
Ok(Self {
cookies_header: serializable.cookies_header,
additional_headers: serializable.additional_headers,
username: serializable.username,
email_address: serializable.email_address,
created_at: serializable.created_at,
updated_at: serializable.updated_at,
})
}
pub fn save_to_file<P: AsRef<Path>>(&self, file_path: P) -> Result<(), WebLoginDataError> {
let serializable = WebLoginDataSerializable {
version: CURRENT_VERSION,
cookies_header: self.cookies_header.clone(),
additional_headers: self.additional_headers.clone(),
username: self.username.clone(),
email_address: self.email_address.clone(),
created_at: self.created_at,
updated_at: self.updated_at,
};
let contents = serde_json::to_string_pretty(&serializable)
.map_err(WebLoginDataError::SerializeError)?;
std::fs::write(file_path, contents)
.map_err(WebLoginDataError::IoError)?;
Ok(())
}
}
#[derive(Debug)]
pub enum WebLoginDataError {
IoError(std::io::Error),
DeserializeError(serde_json::Error),
SerializeError(serde_json::Error),
}
impl Error for WebLoginDataError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
WebLoginDataError::IoError(e) => Some(e),
WebLoginDataError::DeserializeError(e) => Some(e),
WebLoginDataError::SerializeError(e) => Some(e),
_ => None,
}
}
}
impl Display for WebLoginDataError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
WebLoginDataError::IoError(e) => write!(f, "IO error: {}", e),
WebLoginDataError::DeserializeError(e) => write!(f, "Deserialization error: {}", e),
WebLoginDataError::SerializeError(e) => write!(f, "Serialization error: {}", e),
}
}
}
#[derive(Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
struct WebLoginDataSerializable {
version: u8,
#[serde(skip_serializing_if = "Option::is_none")]
cookies_header: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
additional_headers: Option<HashMap<String, String>>,
#[serde(skip_serializing_if = "Option::is_none")]
username: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
email_address: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
created_at: Option<DateTime<Utc>>,
#[serde(skip_serializing_if = "Option::is_none")]
updated_at: Option<DateTime<Utc>>,
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::NamedTempFile;
use std::io::Write;
#[test]
fn save_writes_version() {
let file = NamedTempFile::new().unwrap();
let data = WebLoginData::new();
data.save_to_file(file.path()).unwrap();
let contents = std::fs::read_to_string(file.path()).unwrap();
assert!(contents.contains("\"version\": 1"));
}
#[test]
fn load_reads_all_fields() {
let mut file = NamedTempFile::new().unwrap();
write!(file, r#"{{
"version": 1,
"cookies_header": "session=abc123; visitor=xyz",
"additional_headers": {{"Authorization": "Bearer tok"}},
"username": "testuser",
"email_address": "test@example.com",
"created_at": "2026-01-01T00:00:00Z",
"updated_at": "2026-05-10T12:00:00Z"
}}"#).unwrap();
let data = WebLoginData::load_from_file(file.path()).unwrap();
assert_eq!(data.cookies_header.as_deref(), Some("session=abc123; visitor=xyz"));
assert_eq!(data.username.as_deref(), Some("testuser"));
assert_eq!(data.email_address.as_deref(), Some("test@example.com"));
assert!(data.additional_headers.is_some());
let headers = data.additional_headers.unwrap();
assert_eq!(headers.get("Authorization").unwrap(), "Bearer tok");
assert!(data.created_at.is_some());
assert!(data.updated_at.is_some());
}
#[test]
fn load_tolerates_absent_optional_fields() {
let mut file = NamedTempFile::new().unwrap();
write!(file, r#"{{"version": 1}}"#).unwrap();
let data = WebLoginData::load_from_file(file.path()).unwrap();
assert!(data.cookies_header.is_none());
assert!(data.additional_headers.is_none());
assert!(data.username.is_none());
assert!(data.email_address.is_none());
assert!(data.created_at.is_none());
assert!(data.updated_at.is_none());
}
#[test]
fn load_tolerates_null_fields() {
let mut file = NamedTempFile::new().unwrap();
write!(file, r#"{{
"version": 1,
"cookies_header": null,
"username": null,
"email_address": null
}}"#).unwrap();
let data = WebLoginData::load_from_file(file.path()).unwrap();
assert!(data.cookies_header.is_none());
assert!(data.username.is_none());
assert!(data.email_address.is_none());
}
#[test]
fn load_invalid_json_returns_error() {
let mut file = NamedTempFile::new().unwrap();
write!(file, "not valid json at all").unwrap();
let result = WebLoginData::load_from_file(file.path());
assert!(result.is_err());
}
#[test]
fn load_nonexistent_file_returns_io_error() {
let result = WebLoginData::load_from_file("/tmp/nonexistent_web_login_12345.json");
assert!(result.is_err());
}
#[test]
fn round_trip_full() {
let file = NamedTempFile::new().unwrap();
let mut headers = HashMap::new();
headers.insert("X-Custom".to_string(), "value".to_string());
let now = Utc::now();
let original = WebLoginData {
cookies_header: Some("session=roundtrip".to_string()),
additional_headers: Some(headers),
username: Some("round_trip_user".to_string()),
email_address: Some("rt@example.com".to_string()),
created_at: Some(now),
updated_at: Some(now),
};
original.save_to_file(file.path()).unwrap();
let loaded = WebLoginData::load_from_file(file.path()).unwrap();
assert_eq!(loaded.cookies_header.as_deref(), Some("session=roundtrip"));
assert_eq!(loaded.username.as_deref(), Some("round_trip_user"));
assert_eq!(loaded.email_address.as_deref(), Some("rt@example.com"));
assert_eq!(loaded.additional_headers.as_ref().unwrap().get("X-Custom").unwrap(), "value");
assert!(loaded.created_at.is_some());
assert!(loaded.updated_at.is_some());
}
#[test]
fn round_trip_minimal() {
let file = NamedTempFile::new().unwrap();
let original = WebLoginData::new();
original.save_to_file(file.path()).unwrap();
let loaded = WebLoginData::load_from_file(file.path()).unwrap();
assert!(loaded.cookies_header.is_none());
assert!(loaded.additional_headers.is_none());
assert!(loaded.username.is_none());
assert!(loaded.email_address.is_none());
}
}
@@ -0,0 +1,50 @@
use serde_derive::{Deserialize, Serialize};
use crate::core::providers::credentials::provider_credential_type::ProviderCredentialType;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ProviderCredentialKey {
// ========== API KEYS==========
FalApiKey,
ReplicateApiKey,
// ========== WEB LOGINS ==========
GrokWebLogin,
HiggsfieldWebLogin,
MidjourneyLogin,
RunwayWebLogin,
}
impl ProviderCredentialKey {
pub fn get_type(&self) -> ProviderCredentialType {
match self {
// Api keys
Self::FalApiKey => ProviderCredentialType::ApiKey,
Self::ReplicateApiKey => ProviderCredentialType::ApiKey,
// Web logins
Self::GrokWebLogin => ProviderCredentialType::WebLogin,
Self::HiggsfieldWebLogin => ProviderCredentialType::WebLogin,
Self::MidjourneyLogin => ProviderCredentialType::WebLogin,
Self::RunwayWebLogin => ProviderCredentialType::WebLogin,
}
}
pub fn get_filename(&self) -> &'static str {
match self {
// Api keys
ProviderCredentialKey::FalApiKey => "fal.api_key.txt",
ProviderCredentialKey::ReplicateApiKey => "replicate.api_key.txt",
// Web logins
ProviderCredentialKey::GrokWebLogin => "grok.web_login.toml",
ProviderCredentialKey::HiggsfieldWebLogin => "higgsfield.web_login.toml",
ProviderCredentialKey::MidjourneyLogin => "midjourney.web_login.toml",
ProviderCredentialKey::RunwayWebLogin => "runway.web_login.toml",
}
}
}
@@ -0,0 +1,187 @@
use std::collections::HashMap;
use std::error::Error;
use std::fmt::Display;
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};
use log::{info, warn};
use crate::core::providers::credentials::payload::api_key::{ApiKeyData, ApiKeyDataError};
use crate::core::providers::credentials::payload::provider_credential_payload::ProviderCredentialPayload;
use crate::core::providers::credentials::payload::web_login::{WebLoginData, WebLoginDataError};
use crate::core::providers::credentials::provider_credential_key::ProviderCredentialKey;
use crate::core::providers::credentials::provider_credential_type::ProviderCredentialType;
use crate::core::state::data_dir::app_data_root::AppDataRoot;
use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
struct CacheEntry {
payload: ProviderCredentialPayload,
loaded_at: Instant,
}
#[derive(Clone)]
pub struct ProviderCredentialLoadingCache {
inner: Arc<RwLock<HashMap<ProviderCredentialKey, CacheEntry>>>,
app_data_root: AppDataRoot,
ttl: Duration,
}
impl ProviderCredentialLoadingCache {
pub fn new(app_data_root: AppDataRoot) -> Self {
Self {
inner: Arc::new(RwLock::new(HashMap::new())),
app_data_root,
ttl: Duration::from_secs(300), // 5 minutes
}
}
pub fn new_with_ttl(app_data_root: AppDataRoot, ttl: Duration) -> Self {
Self {
inner: Arc::new(RwLock::new(HashMap::new())),
app_data_root,
ttl,
}
}
pub fn get_credentials(
&self,
key: ProviderCredentialKey,
) -> Result<Option<ProviderCredentialPayload>, ProviderCredentialLoadingCacheError> {
// Check cache first.
{
let cache = self.inner.read()
.map_err(|_| ProviderCredentialLoadingCacheError::LockPoisoned)?;
if let Some(entry) = cache.get(&key) {
if entry.loaded_at.elapsed() < self.ttl {
return Ok(Some(entry.payload.clone()));
}
}
}
// Cache miss or expired — load from disk.
let file_path = self.app_data_root.credentials_dir().path().join(key.get_filename());
if !file_path.exists() {
return Ok(None);
}
info!("Loading credential from disk: {:?}", file_path);
let payload = match key.get_type() {
ProviderCredentialType::ApiKey => {
let data = ApiKeyData::load_from_file(&file_path)
.map_err(ProviderCredentialLoadingCacheError::ApiKeyError)?;
ProviderCredentialPayload::ApiKey(data)
}
ProviderCredentialType::WebLogin => {
let data = WebLoginData::load_from_file(&file_path)
.map_err(ProviderCredentialLoadingCacheError::WebLoginError)?;
ProviderCredentialPayload::WebLogin(data)
}
};
// Store in cache.
{
let mut cache = self.inner.write()
.map_err(|_| ProviderCredentialLoadingCacheError::LockPoisoned)?;
cache.insert(key, CacheEntry {
payload: payload.clone(),
loaded_at: Instant::now(),
});
}
Ok(Some(payload))
}
pub fn save_credentials(
&self,
key: ProviderCredentialKey,
payload: ProviderCredentialPayload,
) -> Result<(), ProviderCredentialLoadingCacheError> {
let file_path = self.app_data_root.credentials_dir().path().join(key.get_filename());
info!("Saving credential to disk: {:?}", file_path);
match (&payload, key.get_type()) {
(ProviderCredentialPayload::ApiKey(data), ProviderCredentialType::ApiKey) => {
data.save_to_file(&file_path)
.map_err(ProviderCredentialLoadingCacheError::ApiKeyError)?;
}
(ProviderCredentialPayload::WebLogin(data), ProviderCredentialType::WebLogin) => {
data.save_to_file(&file_path)
.map_err(ProviderCredentialLoadingCacheError::WebLoginError)?;
}
_ => {
warn!("Credential type mismatch for key {:?}", key);
return Err(ProviderCredentialLoadingCacheError::TypeMismatch);
}
}
// Update cache.
{
let mut cache = self.inner.write()
.map_err(|_| ProviderCredentialLoadingCacheError::LockPoisoned)?;
cache.insert(key, CacheEntry {
payload,
loaded_at: Instant::now(),
});
}
Ok(())
}
/// Invalidate a cached entry, forcing a re-read from disk on next access.
pub fn invalidate(&self, key: ProviderCredentialKey) {
if let Ok(mut cache) = self.inner.write() {
cache.remove(&key);
}
}
pub fn delete_credentials(&self, key: ProviderCredentialKey) -> Result<(), ProviderCredentialLoadingCacheError> {
let file_path = self.app_data_root.credentials_dir().path().join(key.get_filename());
if file_path.exists() {
std::fs::remove_file(file_path)
.map_err(|e| ProviderCredentialLoadingCacheError::IoError(e))?;
}
self.invalidate(key);
Ok(())
}
}
#[derive(Debug)]
pub enum ProviderCredentialLoadingCacheError {
IoError(std::io::Error),
ApiKeyError(ApiKeyDataError),
WebLoginError(WebLoginDataError),
TypeMismatch,
LockPoisoned,
}
impl Error for ProviderCredentialLoadingCacheError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
Self::IoError(err) => Some(err),
Self::ApiKeyError(err) => Some(err),
Self::WebLoginError(err) => Some(err),
_ => None,
}
}
}
impl Display for ProviderCredentialLoadingCacheError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::IoError(e) => write!(f, "IO error: {}", e),
Self::ApiKeyError(e) => write!(f, "API key error: {}", e),
Self::WebLoginError(e) => write!(f, "Web login error: {}", e),
Self::TypeMismatch => write!(f, "Credential type does not match key type"),
Self::LockPoisoned => write!(f, "Internal lock poisoned"),
}
}
}
@@ -0,0 +1,9 @@
use serde_derive::{Deserialize, Serialize};
#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ProviderCredentialType {
ApiKey,
WebLogin,
}
@@ -0,0 +1 @@
pub mod credentials;
@@ -7,4 +7,4 @@
/// We increment this whenever we make a change to the database schema.
/// This prevents deadlock on startup if the schema has changed from the
/// previous version.
pub const TASK_DATABASE_VERSION: u32 = 6;
pub const TASK_DATABASE_VERSION: u32 = 7;
@@ -1,2 +1,3 @@
pub mod discord_presence_thread;
pub mod main_window_thread;
pub mod main_window_thread;
pub mod third_party_task_polling_thread;

Some files were not shown because too many files have changed in this diff Show More