From eff94be8c0b040f12893283c654f3fee2aafbc92 Mon Sep 17 00:00:00 2001 From: Tanish Manku Date: Mon, 15 Jun 2026 16:56:13 -0700 Subject: [PATCH] artcraft-mcp: include realized cost on generate_image responses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Lets the LLM truthfully answer 'what did that just cost?' without remembering the pre-flight estimate. The backend has no per-job cost endpoint or credit-transactions API; the only way to know the realized cost is to diff the credit balance before and after. generate_image::run now: - Captures get_session_credits before submit. - Runs generation + poll as before. - Captures get_session_credits after success. - Computes credits_charged = balance_before - balance_after. - Attaches { credits_charged, credits_balance_after } to the response. Best-effort: if either credit lookup fails the generation still succeeds and the cost just isn't reported. Caveat noted in comments: concurrent generations settling inside this window can over-count the diff. The tool's text content gains a 'Cost: N credits charged. Balance after: M credits.' line when cost info is available — picked up by both the user reading the chat and the LLM answering follow-up questions about cost. --- crates/cli/artcraft-mcp/src/server.rs | 12 ++++- .../artcraft-mcp/src/tools/generate_image.rs | 54 ++++++++++++++++++- 2 files changed, 62 insertions(+), 4 deletions(-) diff --git a/crates/cli/artcraft-mcp/src/server.rs b/crates/cli/artcraft-mcp/src/server.rs index 955ef7f377..048cfb9d76 100644 --- a/crates/cli/artcraft-mcp/src/server.rs +++ b/crates/cli/artcraft-mcp/src/server.rs @@ -99,7 +99,7 @@ impl ArtcraftServer { } #[tool( - description = "Generate an image with Artcraft. Required: `prompt`. Optional: `model` (snake_case id — call list_image_models for valid values; defaults to nano_banana_pro), `aspect_ratio`, `resolution`, `quality`, `num_images`, `reference_image_urls` (https URLs for image-edit / img2img — only valid on models with image_refs_supported=true). Each parameter must satisfy the constraints reported by list_image_models for the chosen model. Blocks up to 90 seconds for completion. Returns the generated image inline plus its full-resolution CDN URL." + description = "Generate an image with Artcraft. Required: `prompt`. Optional: `model` (snake_case id — call list_image_models for valid values; defaults to nano_banana_pro), `aspect_ratio`, `resolution`, `quality`, `num_images`, `reference_image_urls` (https URLs for image-edit / img2img — only valid on models with image_refs_supported=true). Each parameter must satisfy the constraints reported by list_image_models for the chosen model. Blocks up to 90 seconds for completion. Returns the generated images inline, plus their full-resolution CDN URLs and the realized cost (credits charged + remaining balance). The cost is derived by diffing the credit balance before and after generation." )] async fn generate_image( &self, @@ -157,7 +157,15 @@ async fn success_images_result(images: generate_image::GeneratedImages) -> CallT .collect::>() .join("\n") }; - content.push(Content::text(urls_text)); + + let text = match &images.cost { + Some(cost) => format!( + "{}\n\nCost: {} credits charged. Balance after: {} credits.", + urls_text, cost.credits_charged, cost.credits_balance_after + ), + None => urls_text, + }; + content.push(Content::text(text)); CallToolResult::success(content) } diff --git a/crates/cli/artcraft-mcp/src/tools/generate_image.rs b/crates/cli/artcraft-mcp/src/tools/generate_image.rs index 555e54c978..cdfe7fb032 100644 --- a/crates/cli/artcraft-mcp/src/tools/generate_image.rs +++ b/crates/cli/artcraft-mcp/src/tools/generate_image.rs @@ -7,12 +7,14 @@ use tokio::time::{sleep, Instant}; use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet; +use artcraft_client::endpoints::credits::get_session_credits::get_session_credits; use artcraft_client::endpoints::jobs::list_session_jobs::{list_session_jobs, States}; use artcraft_client::endpoints::media_files::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files; use artcraft_client::endpoints::media_files::upload_image_media_file_from_bytes::{ upload_image_media_file_from_bytes, ImageType, UploadImageBytesArgs, }; use artcraft_client::endpoints::omni_gen::generate::image::omni_gen_image::omni_gen_image_generate; +use enums::common::payments_namespace::PaymentsNamespace; use artcraft_client::utils::api_host::ApiHost; use enums::common::generation::common_aspect_ratio::CommonAspectRatio; use enums::common::generation::common_image_model::CommonImageModel; @@ -38,6 +40,15 @@ pub struct GeneratedImage { pub struct GeneratedImages { pub images: Vec, + /// None if either credit lookup failed; otherwise the realized cost + /// of this generation derived by diffing the user's credit balance + /// from just-before-submit to just-after-success. + pub cost: Option, +} + +pub struct CostInfo { + pub credits_charged: u64, + pub credits_balance_after: u64, } #[derive(Debug, Deserialize, schemars::JsonSchema)] @@ -113,6 +124,16 @@ pub async fn run(args: Args) -> Result { adjust_zoom: None, }; + // Snapshot credits before submit so we can compute the realized cost + // by diffing against the balance after success. Best-effort — if the + // credits API hiccups we still complete the generation, we just don't + // report a cost. + let credits_before = get_session_credits(&api_host, Some(&creds), PaymentsNamespace::Artcraft) + .await + .map(|c| c.sum_total_credits) + .map_err(|e| tracing::warn!("pre-submit credits lookup failed: {:?}", e)) + .ok(); + tracing::info!("submitting omni_gen image request, model={:?}", model); let submit = omni_gen_image_generate(&api_host, Some(&creds), request) @@ -125,7 +146,34 @@ pub async fn run(args: Args) -> Result { tracing::info!("submitted; inference_job_token={:?}", submit.inference_job_token); - poll_for_image_url(&api_host, &creds, &submit.inference_job_token).await + let mut result = poll_for_image_url(&api_host, &creds, &submit.inference_job_token).await?; + + // Compute realized cost. If concurrent generations on other tools + // settle in this window the diff can over-count, so we annotate + // the value in the log; the LLM gets the number either way. + let credits_after = get_session_credits(&api_host, Some(&creds), PaymentsNamespace::Artcraft) + .await + .map(|c| c.sum_total_credits) + .map_err(|e| { + tracing::warn!("post-success credits lookup failed: {:?}", e); + }) + .ok(); + + if let (Some(before), Some(after)) = (credits_before, credits_after) { + let charged = before.saturating_sub(after); + tracing::info!( + "realized cost: {} credits (balance {} -> {})", + charged, + before, + after + ); + result.cost = Some(CostInfo { + credits_charged: charged, + credits_balance_after: after, + }); + } + + Ok(result) } async fn poll_for_image_url( @@ -194,9 +242,10 @@ async fn poll_for_image_url( .maybe_thumbnail_template .clone(), }], + cost: None, }); } - return Ok(GeneratedImages { images }); + return Ok(GeneratedImages { images, cost: None }); } return Ok(GeneratedImages { @@ -204,6 +253,7 @@ async fn poll_for_image_url( cdn_url: result.media_links.cdn_url.to_string(), maybe_thumbnail_template: result.media_links.maybe_thumbnail_template.clone(), }], + cost: None, }); } JobStatusPlus::CompleteFailure | JobStatusPlus::Dead => {