From 9bafc23a410ca2fbcbfbb6ec71a84b285d78c227 Mon Sep 17 00:00:00 2001 From: Tanish Manku Date: Mon, 8 Jun 2026 12:01:24 -0700 Subject: [PATCH] artcraft-mcp: MVP stdio server with generate_image tool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Standalone crate at crates/cli/artcraft-mcp/ exposing a single MCP tool that submits an image generation via /v1/omni_gen/generate/image, polls /v1/jobs/session until the job completes, and returns the CDN URL. Reads Storyteller cookies from ~/Artcraft/credentials/ written by the desktop app — no separate sign-in needed. 90s poll timeout. Kept as its own [workspace] during the spike so the parent workspace Cargo.toml stays untouched and parallel work on main doesn't conflict. Workspace registration deferred to a final commit once the spike is validated end-to-end against Claude Desktop. --- crates/cli/artcraft-mcp/.gitignore | 2 + crates/cli/artcraft-mcp/Cargo.toml | 42 ++++++ crates/cli/artcraft-mcp/src/main.rs | 223 ++++++++++++++++++++++++++++ 3 files changed, 267 insertions(+) create mode 100644 crates/cli/artcraft-mcp/.gitignore create mode 100644 crates/cli/artcraft-mcp/Cargo.toml create mode 100644 crates/cli/artcraft-mcp/src/main.rs diff --git a/crates/cli/artcraft-mcp/.gitignore b/crates/cli/artcraft-mcp/.gitignore new file mode 100644 index 0000000000..4fffb2f89c --- /dev/null +++ b/crates/cli/artcraft-mcp/.gitignore @@ -0,0 +1,2 @@ +/target +/Cargo.lock diff --git a/crates/cli/artcraft-mcp/Cargo.toml b/crates/cli/artcraft-mcp/Cargo.toml new file mode 100644 index 0000000000..a35f20417b --- /dev/null +++ b/crates/cli/artcraft-mcp/Cargo.toml @@ -0,0 +1,42 @@ +[workspace] +# Standalone workspace during the MVP spike — keeps us out of the root +# workspace's Cargo.toml until the final commit so other agents on `main` +# don't collide with this work. Will be merged into the parent workspace +# as the last step (see docs/mcp-integration-plan.md, Phase 6). + +[package] +name = "artcraft-mcp" +version = "0.0.1" +edition = "2021" +publish = false + +[[bin]] +name = "artcraft-mcp" +path = "src/main.rs" + +[dependencies] +# MCP SDK +rmcp = { version = "1.7.0", features = ["server", "macros", "transport-io", "schemars"] } +schemars = "1.0" + +# Async + logging +tokio = { version = "1.50.0", features = ["macros", "rt-multi-thread", "io-std", "time"] } +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] } + +# Serde + errors +serde = { version = "1.0.228", features = ["derive"] } +serde_json = "1.0.149" +anyhow = "1.0.102" + +# Home directory resolution (matches desktop crate) +directories = "6.0.0" + +# Internal crates — path deps into the parent workspace. +# Cargo resolves their own `.workspace = true` deps against the root +# workspace at storyteller-rust/Cargo.toml (Cargo walks up from each +# path-dep target), so we don't have to mirror those. +artcraft_client = { path = "../../api_clients/artcraft/artcraft_client" } +artcraft_api_defs = { path = "../../api_clients/artcraft/artcraft_api_defs" } +enums = { path = "../../schema/public/enums" } +tokens = { path = "../../schema/public/tokens" } diff --git a/crates/cli/artcraft-mcp/src/main.rs b/crates/cli/artcraft-mcp/src/main.rs new file mode 100644 index 0000000000..dc05675555 --- /dev/null +++ b/crates/cli/artcraft-mcp/src/main.rs @@ -0,0 +1,223 @@ +use std::collections::HashSet; +use std::path::PathBuf; +use std::time::Duration; + +use anyhow::{anyhow, Result}; +use rmcp::handler::server::router::tool::ToolRouter; +use rmcp::handler::server::wrapper::Parameters; +use rmcp::model::{ + CallToolResult, Content, Implementation, ProtocolVersion, ServerCapabilities, ServerInfo, +}; +use rmcp::transport::stdio; +use rmcp::{schemars, tool, tool_handler, tool_router}; +use rmcp::{ErrorData as McpError, ServerHandler, ServiceExt}; +use serde::Deserialize; +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_avt_cookie::StorytellerAvtCookie; +use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet; +use artcraft_client::credentials::storyteller_session_cookie::StorytellerSessionCookie; +use artcraft_client::endpoints::jobs::list_session_jobs::{list_session_jobs, States}; +use artcraft_client::endpoints::omni_gen::generate::image::omni_gen_image::omni_gen_image_generate; +use artcraft_client::utils::api_host::ApiHost; +use enums::common::generation::common_image_model::CommonImageModel; +use enums::common::job_status_plus::JobStatusPlus; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; + +const POLL_INTERVAL: Duration = Duration::from_secs(3); +const POLL_TIMEOUT: Duration = Duration::from_secs(90); +const DEFAULT_MODEL: CommonImageModel = CommonImageModel::NanoBananaPro; + +#[derive(Debug, Deserialize, schemars::JsonSchema)] +struct GenerateImageArgs { + /// Text prompt describing the image to generate. + prompt: String, +} + +#[derive(Clone)] +struct ArtcraftServer { + tool_router: ToolRouter, +} + +#[tool_router] +impl ArtcraftServer { + fn new() -> Self { + Self { + tool_router: Self::tool_router(), + } + } + + #[tool( + description = "Generate an image with Artcraft from a text prompt. Blocks up to 90 seconds for completion. Returns the CDN URL of the generated image." + )] + async fn generate_image( + &self, + Parameters(GenerateImageArgs { prompt }): Parameters, + ) -> Result { + match run_generate_image(&prompt).await { + Ok(url) => Ok(CallToolResult::success(vec![Content::text(url)])), + Err(err) => Ok(CallToolResult::error(vec![Content::text(format!( + "{:#}", + err + ))])), + } + } +} + +#[tool_handler] +impl ServerHandler for ArtcraftServer { + fn get_info(&self) -> ServerInfo { + ServerInfo::new( + ServerCapabilities::builder().enable_tools().build(), + ) + .with_server_info(Implementation::from_build_env()) + .with_protocol_version(ProtocolVersion::V_2024_11_05) + .with_instructions( + "Artcraft MCP server. Use generate_image with a text prompt to create an image.".to_string(), + ) + } +} + +async fn run_generate_image(prompt: &str) -> Result { + let trimmed = prompt.trim(); + if trimmed.is_empty() { + return Err(anyhow!("prompt is empty")); + } + + let api_host = ApiHost::Storyteller; + let creds = load_storyteller_credentials()?; + if creds.is_empty() { + return Err(anyhow!( + "No Storyteller credentials found at ~/Artcraft/credentials/. \ + Sign in to the Artcraft desktop app first." + )); + } + + let request = OmniGenImageCostAndGenerateRequest { + idempotency_token: None, + model: Some(DEFAULT_MODEL), + prompt: Some(trimmed.to_string()), + image_media_tokens: None, + resolution: None, + aspect_ratio: None, + quality: None, + image_batch_count: Some(1), + adjust_horizontal_angle: None, + adjust_vertical_angle: None, + adjust_zoom: None, + }; + + tracing::info!("submitting omni_gen image request, model={:?}", DEFAULT_MODEL); + + let submit = omni_gen_image_generate(&api_host, Some(&creds), request) + .await + .map_err(|e| anyhow!("submit failed: {:?}", e))?; + + if !submit.success { + return Err(anyhow!("submit returned success=false")); + } + + tracing::info!("submitted; inference_job_token={:?}", submit.inference_job_token); + + poll_for_image_url(&api_host, &creds, &submit.inference_job_token).await +} + +async fn poll_for_image_url( + api_host: &ApiHost, + creds: &StorytellerCredentialSet, + target: &InferenceJobToken, +) -> Result { + let deadline = Instant::now() + POLL_TIMEOUT; + + let mut include_states = HashSet::new(); + include_states.insert(JobStatusPlus::CompleteSuccess); + include_states.insert(JobStatusPlus::CompleteFailure); + include_states.insert(JobStatusPlus::Dead); + + loop { + if Instant::now() >= deadline { + return Err(anyhow!( + "Generation did not complete within {}s. Job token: {:?}", + POLL_TIMEOUT.as_secs(), + target + )); + } + + let response = list_session_jobs( + api_host, + Some(creds), + States::Include(include_states.clone()), + ) + .await + .map_err(|e| anyhow!("poll failed: {:?}", e))?; + + if let Some(job) = response.jobs.iter().find(|j| &j.job_token == target) { + match job.status.status { + JobStatusPlus::CompleteSuccess => { + let result = job + .maybe_result + .as_ref() + .ok_or_else(|| anyhow!("job marked complete but result missing"))?; + return Ok(result.media_links.cdn_url.to_string()); + } + JobStatusPlus::CompleteFailure | JobStatusPlus::Dead => { + let msg = job + .status + .maybe_failure_message + .clone() + .unwrap_or_else(|| format!("{:?}", job.status.status)); + return Err(anyhow!("generation failed: {}", msg)); + } + _ => {} + } + } + + sleep(POLL_INTERVAL).await; + } +} + +fn load_storyteller_credentials() -> Result { + let creds_dir = home_dir()?.join("Artcraft").join("credentials"); + let session = read_trimmed(&creds_dir.join("artcraft_session.txt"))? + .map(StorytellerSessionCookie::new); + let avt = read_trimmed(&creds_dir.join("artcraft_avt.txt"))? + .map(StorytellerAvtCookie::new); + Ok(StorytellerCredentialSet::initialize(avt, session)) +} + +fn read_trimmed(path: &PathBuf) -> Result> { + if !path.exists() { + return Ok(None); + } + let raw = std::fs::read_to_string(path)?; + let trimmed = raw.trim(); + if trimmed.is_empty() { + Ok(None) + } else { + Ok(Some(trimmed.to_string())) + } +} + +fn home_dir() -> Result { + directories::UserDirs::new() + .map(|d| d.home_dir().to_path_buf()) + .ok_or_else(|| anyhow!("could not determine home directory")) +} + +#[tokio::main] +async fn main() -> Result<()> { + // Logs go to stderr; stdout belongs to the MCP transport. + tracing_subscriber::fmt() + .with_writer(std::io::stderr) + .with_ansi(false) + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_env("ARTCRAFT_MCP_LOG") + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")), + ) + .init(); + + let service = ArtcraftServer::new().serve(stdio()).await?; + service.waiting().await?; + Ok(()) +}