use std::convert::TryFrom; use anyhow::anyhow; use async_trait::async_trait; use chrono::Utc; use cynic::{MutationBuilder, QueryBuilder}; #[cfg(test)] use mockall::automock; use warp_core::channel::{Channel, ChannelState}; use warp_graphql::mutations::share_block::{ BlockInput, ShareBlock, ShareBlockResult, ShareBlockVariables, }; use warp_graphql::mutations::unshare_block::{ UnshareBlock, UnshareBlockInput, UnshareBlockResult, UnshareBlockVariables, }; use warp_graphql::queries::get_blocks_for_user::{ Block as GqlBlock, GetBlocksForUser, GetBlocksForUserVariables, }; use super::ServerApi; use crate::ai::generate_block_title::api::{GenerateBlockTitleRequest, GenerateBlockTitleResponse}; use crate::server::block::{Block, DisplaySetting}; use crate::server::graphql::{get_request_context, get_user_facing_error_message}; #[cfg_attr(test, automock)] #[cfg_attr(not(target_family = "wasm"), async_trait)] #[cfg_attr(target_family = "wasm", async_trait(?Send))] pub trait BlockClient: 'static + Send + Sync { /// Unshares a block identified at `block_id`. async fn unshare_block(&self, block_id: String) -> Result<(), anyhow::Error>; /// Uploads a given block to the server via the /share_block endpoint. async fn save_block( &self, block: &Block, title: Option, show_prompt: bool, display_setting: DisplaySetting, ) -> Result; async fn blocks_owned_by_user(&self) -> Result, anyhow::Error>; async fn generate_shared_block_title( &self, request: GenerateBlockTitleRequest, ) -> Result; } #[cfg_attr(not(target_family = "wasm"), async_trait)] #[cfg_attr(target_family = "wasm", async_trait(?Send))] impl BlockClient for ServerApi { async fn unshare_block(&self, block_uid: String) -> Result<(), anyhow::Error> { let variables = UnshareBlockVariables { input: UnshareBlockInput { block_uid }, request_context: get_request_context(), }; let operation = UnshareBlock::build(variables); let response = self.send_graphql_request(operation, None).await?; match response.unshare_block { UnshareBlockResult::UnshareBlockOutput(output) => { if output.success { Ok(()) } else { Err(anyhow!("Failed to unshare block")) } } UnshareBlockResult::UserFacingError(error) => { Err(anyhow!(get_user_facing_error_message(error))) } UnshareBlockResult::Unknown => Err(anyhow!("Failed to unshare block")), } } async fn save_block( &self, block: &Block, title: Option, show_prompt: bool, display_setting: DisplaySetting, ) -> Result { let variables = ShareBlockVariables { block: BlockInput { command: block.command.as_deref(), embed_display_setting: display_setting.into(), output: block.output.as_deref(), show_prompt, stylized_command: block.stylized_command.as_deref(), stylized_output: block.stylized_output.as_deref(), stylized_prompt: block.stylized_prompt.as_deref(), stylized_prompt_and_command: block.stylized_prompt_and_command.as_deref(), time_started_term: Some(block.time_started_term.with_timezone(&Utc).into()), title: title.as_deref(), }, request_context: get_request_context(), }; let operation = ShareBlock::build(variables); let response = self.send_graphql_request(operation, None).await?; match response.share_block { ShareBlockResult::ShareBlockOutput(output) => { let mut created_url = format!("{}{}", ChannelState::server_root_url(), output.url_ending); // If this is a preview build, ensure the link routes to a preview build. if matches!(ChannelState::channel(), Channel::Preview) { created_url.push_str("?preview=true"); } Ok(created_url) } ShareBlockResult::UserFacingError(error) => { Err(anyhow!(get_user_facing_error_message(error))) } ShareBlockResult::Unknown => Err(anyhow!("Failed to share block")), } } async fn blocks_owned_by_user(&self) -> Result, anyhow::Error> { let variables = GetBlocksForUserVariables { request_context: get_request_context(), }; let operation = GetBlocksForUser::build(variables); let response = self.send_graphql_request(operation, None).await?; match response.user { warp_graphql::queries::get_blocks_for_user::UserResult::UserOutput(user_output) => { Ok(user_output .user .blocks .into_iter() .filter_map(|block| block.try_into().ok()) .collect()) } warp_graphql::queries::get_blocks_for_user::UserResult::Unknown => { Err(anyhow!("Unable to fetch blocks")) } } } async fn generate_shared_block_title( &self, request: GenerateBlockTitleRequest, ) -> Result { let auth_token = self.get_or_refresh_access_token().await?; let request_builder = self.base_client.http_client().post(format!( "{}/ai/generate_block_title", ChannelState::server_root_url() )); let response = if let Some(token) = auth_token.as_bearer_token() { request_builder.bearer_auth(token) } else { request_builder } .json(&request) .send() .await? .error_for_status()? .json() .await?; Ok(response) } } impl TryFrom for Block { type Error = anyhow::Error; fn try_from(value: GqlBlock) -> Result { match (value.uid, value.time_started_term) { (uid, Some(time_started_term)) => { Ok(Block { id: Some(uid.into_inner()), command: value.command, output: None, stylized_command: None, stylized_output: None, pwd: None, time_started_term: time_started_term.utc().into(), // This is a dummy value - we are no longer using time_completed_term, // and GqlBlock does not have a time_completed_term field. time_completed_term: time_started_term.utc().into(), stylized_prompt: None, stylized_prompt_and_command: None, }) } _ => Err(anyhow!("missing id or time_started_term")), } } }