mirror of
https://github.com/storytold/artcraft.git
synced 2026-10-20 17:02:53 -04:00
tauri providers - fal api key support (#1502)
This commit is contained in:
-12
@@ -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"
|
||||
}
|
||||
+12
@@ -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"
|
||||
}
|
||||
+16
-4
@@ -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
@@ -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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+3
@@ -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)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+2
@@ -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),
|
||||
}))
|
||||
}
|
||||
|
||||
+2
@@ -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),
|
||||
}))
|
||||
}
|
||||
|
||||
+2
@@ -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),
|
||||
}))
|
||||
}
|
||||
|
||||
+2
@@ -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),
|
||||
}))
|
||||
}
|
||||
|
||||
+2
@@ -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),
|
||||
}))
|
||||
}
|
||||
|
||||
+2
@@ -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),
|
||||
}))
|
||||
}
|
||||
|
||||
+2
@@ -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),
|
||||
}))
|
||||
}
|
||||
|
||||
+2
@@ -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),
|
||||
}))
|
||||
}
|
||||
|
||||
+2
@@ -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),
|
||||
}))
|
||||
}
|
||||
|
||||
+2
@@ -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),
|
||||
}))
|
||||
}
|
||||
|
||||
+2
@@ -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),
|
||||
}))
|
||||
}
|
||||
|
||||
+2
@@ -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),
|
||||
}))
|
||||
}
|
||||
|
||||
+2
@@ -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),
|
||||
}))
|
||||
}
|
||||
|
||||
+6
@@ -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.
|
||||
|
||||
+12
-4
@@ -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),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
+12
-4
@@ -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),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
+12
-4
@@ -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),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
+12
-4
@@ -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),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+3
-1
@@ -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"),
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+2
@@ -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,
|
||||
|
||||
+31
-22
@@ -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,
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+2
@@ -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;
|
||||
+159
@@ -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,
|
||||
})
|
||||
}
|
||||
+66
@@ -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,
|
||||
}
|
||||
}
|
||||
+17
@@ -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")
|
||||
}
|
||||
+4
@@ -0,0 +1,4 @@
|
||||
pub mod handle_router;
|
||||
mod handle_api_providers;
|
||||
mod handle_web_login_providers;
|
||||
mod utils;
|
||||
+50
@@ -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,
|
||||
}
|
||||
}
|
||||
+27
@@ -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)
|
||||
}
|
||||
+2
@@ -0,0 +1,2 @@
|
||||
pub mod convert_enums_to_router;
|
||||
pub mod map_media_files_to_urls;
|
||||
+2
@@ -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
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+2
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+2
@@ -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);
|
||||
|
||||
@@ -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;
|
||||
+10
@@ -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",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+187
@@ -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
Reference in New Issue
Block a user