Initial public release of Warp.
Repo-Sync-Origin: warpdotdev/warp-internal@12af1d983b
This commit is contained in:
David Stern
@@ -0,0 +1,275 @@
|
||||
use crate::terminal::model::{
|
||||
block::{Block as ClientBlock, BlockTime},
|
||||
grid::Dimensions as _,
|
||||
ObfuscateSecrets,
|
||||
};
|
||||
use chrono::{DateTime, FixedOffset, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use warp_graphql::mutations::share_block::DisplaySetting as GqlDisplaySetting;
|
||||
|
||||
// These are pixel heights of various parts of an embedded block.
|
||||
pub const TITLE_HEIGHT: u32 = 34;
|
||||
pub const HEADER_PADDING: u32 = 30;
|
||||
pub const OUTPUT_PADDING: u32 = 32;
|
||||
pub const LINE_HEIGHT: u32 = 19;
|
||||
pub const PROMPT_LINE_HEIGHT: u32 = 16;
|
||||
pub const OUTPUT_CELL_WIDTH: u32 = 10;
|
||||
pub const EMBED_FOOTER_HEIGHT: u32 = 38;
|
||||
pub const EXTRA_PADDING: u32 = 35;
|
||||
|
||||
/// This enum is a replica of the `share_block::DisplaySetting` struct auto-generated from the GraphQL Schema.
|
||||
/// We cannot derive traits on the auto-generated structs because any rust attributes
|
||||
/// will be rewritten.
|
||||
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
|
||||
pub enum DisplaySetting {
|
||||
Command,
|
||||
Output,
|
||||
CommandAndOutput,
|
||||
Other(String),
|
||||
}
|
||||
|
||||
impl From<DisplaySetting> for GqlDisplaySetting {
|
||||
fn from(value: DisplaySetting) -> Self {
|
||||
match value {
|
||||
DisplaySetting::Command => GqlDisplaySetting::Command,
|
||||
DisplaySetting::Output => GqlDisplaySetting::Output,
|
||||
DisplaySetting::CommandAndOutput => GqlDisplaySetting::CommandAndOutput,
|
||||
DisplaySetting::Other(s) => GqlDisplaySetting::Other(s),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A representation of a Block for the server.
|
||||
#[derive(Serialize, Deserialize, Debug)]
|
||||
#[serde(rename_all = "PascalCase")]
|
||||
pub struct Block {
|
||||
pub id: Option<String>,
|
||||
|
||||
/// The input lines for a block.
|
||||
pub command: Option<String>,
|
||||
|
||||
/// The output lines for a block.
|
||||
pub output: Option<String>,
|
||||
|
||||
/// The input lines with their corresponding escape sequences so it can be rendered outside of
|
||||
/// the terminal.
|
||||
pub stylized_command: Option<String>,
|
||||
|
||||
/// The output lines with their corresponding escape sequences so it can be rendered outside of
|
||||
/// the terminal.
|
||||
pub stylized_output: Option<String>,
|
||||
|
||||
/// The prompt lines with their corresponding escape sequences so it can be rendered outside of
|
||||
/// the terminal.
|
||||
pub stylized_prompt: Option<String>,
|
||||
|
||||
/// The prompt and command (combined) lines with their corresponding escape sequences so it can
|
||||
/// be rendered outside of the terminal. Only non-null if using PS1 with the combined grid.
|
||||
pub stylized_prompt_and_command: Option<String>,
|
||||
|
||||
/// The current working directory of the block.
|
||||
pub pwd: Option<String>,
|
||||
|
||||
/// The terminal's timestamp of block completion.
|
||||
pub time_started_term: DateTime<FixedOffset>,
|
||||
|
||||
/// The terminal's timestamp of block completion.
|
||||
pub time_completed_term: DateTime<FixedOffset>,
|
||||
}
|
||||
|
||||
/// A helper struct to organize the block's contents.
|
||||
struct BlockContents {
|
||||
command: Option<String>,
|
||||
stylized_command: Option<String>,
|
||||
output: Option<String>,
|
||||
stylized_output: Option<String>,
|
||||
stylized_prompt: Option<String>,
|
||||
stylized_prompt_and_command: Option<String>,
|
||||
}
|
||||
|
||||
impl Block {
|
||||
pub fn new(
|
||||
block: &ClientBlock,
|
||||
show_prompt: bool,
|
||||
display_setting: &DisplaySetting,
|
||||
obfuscate_secrets: ObfuscateSecrets,
|
||||
) -> Self {
|
||||
let block_time = BlockTime::new(DateTime::from(Utc::now()), DateTime::from(Utc::now()));
|
||||
|
||||
let block_contents =
|
||||
if obfuscate_secrets.is_visually_obfuscated() {
|
||||
let (command, stylized_command) = match display_setting {
|
||||
DisplaySetting::Command | DisplaySetting::CommandAndOutput => (
|
||||
Some(block.command_with_secrets_obfuscated(
|
||||
false, /*include_escape_sequences*/
|
||||
)),
|
||||
Some(block.command_with_secrets_obfuscated(
|
||||
true, /*include_escape_sequences*/
|
||||
)),
|
||||
),
|
||||
_ => (None, None),
|
||||
};
|
||||
let (output, stylized_output) = match display_setting {
|
||||
DisplaySetting::Output | DisplaySetting::CommandAndOutput => (
|
||||
Some(
|
||||
block
|
||||
.output_grid()
|
||||
.contents_to_string_force_secrets_obfuscated(
|
||||
false, /*include_escape_sequences*/
|
||||
None, /*max_rows*/
|
||||
),
|
||||
),
|
||||
Some(
|
||||
block
|
||||
.output_grid()
|
||||
.contents_to_string_force_secrets_obfuscated(
|
||||
true, /*include_escape_sequences*/
|
||||
None, /*max_rows*/
|
||||
),
|
||||
),
|
||||
),
|
||||
_ => (None, None),
|
||||
};
|
||||
|
||||
let stylized_prompt = show_prompt.then_some(if block.honor_ps1() {
|
||||
block.prompt_with_secrets_obfuscated(true)
|
||||
} else {
|
||||
Self::native_prompt_for_server(block)
|
||||
});
|
||||
let stylized_prompt_and_command = (show_prompt && block.honor_ps1())
|
||||
.then(|| block.prompt_and_command_with_secrets_obfuscated(true));
|
||||
BlockContents {
|
||||
command,
|
||||
stylized_command,
|
||||
output,
|
||||
stylized_output,
|
||||
stylized_prompt,
|
||||
stylized_prompt_and_command,
|
||||
}
|
||||
} else {
|
||||
let (command, stylized_command) = match display_setting {
|
||||
DisplaySetting::Command | DisplaySetting::CommandAndOutput => (
|
||||
Some(block.command_with_secrets_unobfuscated(
|
||||
false, /*include_escape_sequences*/
|
||||
)),
|
||||
Some(block.command_with_secrets_unobfuscated(
|
||||
true, /*include_escape_sequences*/
|
||||
)),
|
||||
),
|
||||
_ => (None, None),
|
||||
};
|
||||
let (output, stylized_output) = match display_setting {
|
||||
DisplaySetting::Output | DisplaySetting::CommandAndOutput => (
|
||||
Some(
|
||||
block
|
||||
.output_grid()
|
||||
.contents_to_string_with_secrets_unobfuscated(
|
||||
false, /*include_escape_sequences*/
|
||||
None, /*max_rows*/
|
||||
),
|
||||
),
|
||||
Some(
|
||||
block
|
||||
.output_grid()
|
||||
.contents_to_string_with_secrets_unobfuscated(
|
||||
true, /*include_escape_sequences*/
|
||||
None, /*max_rows*/
|
||||
),
|
||||
),
|
||||
),
|
||||
_ => (None, None),
|
||||
};
|
||||
|
||||
let stylized_prompt = show_prompt.then_some(if block.honor_ps1() {
|
||||
block.prompt_with_secrets_unobfuscated(true)
|
||||
} else {
|
||||
Self::native_prompt_for_server(block)
|
||||
});
|
||||
let stylized_prompt_and_command = (show_prompt && block.honor_ps1())
|
||||
.then(|| block.prompt_and_command_with_secrets_unobfuscated(true));
|
||||
BlockContents {
|
||||
command,
|
||||
stylized_command,
|
||||
output,
|
||||
stylized_output,
|
||||
stylized_prompt,
|
||||
stylized_prompt_and_command,
|
||||
}
|
||||
};
|
||||
|
||||
Block {
|
||||
id: None,
|
||||
command: block_contents.command,
|
||||
output: block_contents.output,
|
||||
stylized_command: block_contents.stylized_command,
|
||||
stylized_output: block_contents.stylized_output,
|
||||
stylized_prompt: block_contents.stylized_prompt,
|
||||
stylized_prompt_and_command: block_contents.stylized_prompt_and_command,
|
||||
pwd: block.pwd().map(String::from),
|
||||
time_started_term: block_time.time_started_term,
|
||||
time_completed_term: block_time.time_completed_term,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn native_prompt_for_server(block: &ClientBlock) -> String {
|
||||
if let Some(prompt_snapshot) = block.prompt_snapshot() {
|
||||
prompt_snapshot.to_string()
|
||||
} else {
|
||||
let mut stylized_prompt = String::new();
|
||||
if let Some(conda_env) = block.conda_env() {
|
||||
stylized_prompt.push_str(format!("({conda_env}) ").as_str());
|
||||
}
|
||||
if let Some(virtual_env) = block.virtual_env_short_name() {
|
||||
stylized_prompt.push_str(format!("({virtual_env}) ").as_str());
|
||||
}
|
||||
if let Some(pwd) = block.server_pwd().to_owned() {
|
||||
stylized_prompt.push_str(format!("{pwd} ").as_str());
|
||||
}
|
||||
if let Some(git_branch) = block.git_branch() {
|
||||
stylized_prompt.push_str(format!("git:({git_branch})").as_str());
|
||||
}
|
||||
stylized_prompt
|
||||
}
|
||||
}
|
||||
|
||||
pub fn embed_pixel_height(
|
||||
block: &ClientBlock,
|
||||
show_prompt: bool,
|
||||
display_setting: &DisplaySetting,
|
||||
) -> u32 {
|
||||
let mut height = TITLE_HEIGHT;
|
||||
height += HEADER_PADDING;
|
||||
|
||||
if show_prompt {
|
||||
if block.honor_ps1() && !block.render_prompt_on_same_line() {
|
||||
height += block.prompt_number_of_rows() as u32 * PROMPT_LINE_HEIGHT;
|
||||
} else if !block.honor_ps1() {
|
||||
height += PROMPT_LINE_HEIGHT;
|
||||
}
|
||||
}
|
||||
|
||||
match display_setting {
|
||||
DisplaySetting::Command => {
|
||||
height += block.prompt_and_command_number_of_rows() as u32 * LINE_HEIGHT;
|
||||
}
|
||||
DisplaySetting::Output => {
|
||||
height += LINE_HEIGHT; // The command is blank, but space is still rendered inside the sticky header.
|
||||
height += block.output_grid().len() as u32 * LINE_HEIGHT;
|
||||
height += OUTPUT_PADDING;
|
||||
}
|
||||
_ => {
|
||||
height += block.prompt_and_command_number_of_rows() as u32 * LINE_HEIGHT;
|
||||
height += block.output_grid().len() as u32 * LINE_HEIGHT;
|
||||
height += OUTPUT_PADDING;
|
||||
}
|
||||
}
|
||||
|
||||
height += EMBED_FOOTER_HEIGHT;
|
||||
height += EXTRA_PADDING;
|
||||
height
|
||||
}
|
||||
|
||||
pub fn embed_pixel_width(block: &ClientBlock) -> u32 {
|
||||
(block.output_grid().grid_handler().columns() as u32 * OUTPUT_CELL_WIDTH) + OUTPUT_PADDING
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,474 @@
|
||||
//! Stateful fake implementation of [`ObjectClient`] for end-to-end tests
|
||||
//! of cloud preferences sync.
|
||||
//!
|
||||
//! Unlike the auto-generated [`MockObjectClient`] (which requires each test
|
||||
//! to script per-method `expect_*()` calls up front), this fake maintains
|
||||
//! an in-memory store of preferences that can be read back after the
|
||||
//! syncer runs. Tests use [`FakeObjectClient::seed_preference`] to add
|
||||
//! cloud-side state and [`FakeObjectClient::cloud_value`] to assert on
|
||||
//! the result after a sync roundtrip.
|
||||
//!
|
||||
//! Only the subset of [`ObjectClient`] methods actually called by the
|
||||
//! cloud preferences syncer has a real implementation — every other
|
||||
//! method panics with `unimplemented!()`. This is deliberate: callers
|
||||
//! that hit one of these panics have likely wired a non-preferences
|
||||
//! code path through the fake by mistake.
|
||||
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use async_channel::Sender;
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use warp_graphql::object_permissions::AccessLevel;
|
||||
|
||||
use crate::{
|
||||
cloud_object::{
|
||||
model::{
|
||||
actions::{ObjectActionHistory, ObjectActionType},
|
||||
generic_string_model::GenericStringObjectId,
|
||||
},
|
||||
BulkCreateCloudObjectResult, BulkCreateGenericStringObjectsRequest,
|
||||
CreateCloudObjectResult, CreateObjectRequest, CreatedCloudObject,
|
||||
GenericStringObjectFormat, GenericStringObjectUniqueKey, JsonObjectType,
|
||||
ObjectDeleteResult, ObjectIdType, ObjectMetadataUpdateResult, ObjectPermissionUpdateResult,
|
||||
ObjectPermissionsUpdateData, ObjectType, ObjectsToUpdate, Owner, Revision,
|
||||
RevisionAndLastEditor, ServerFolder, ServerMetadata, ServerNotebook, ServerObject,
|
||||
ServerPermissions, ServerPreference, ServerWorkflow, UpdateCloudObjectResult,
|
||||
},
|
||||
drive::{folders::FolderId, sharing::SharingAccessLevel},
|
||||
notebooks::NotebookId,
|
||||
server::{
|
||||
cloud_objects::{
|
||||
listener::ObjectUpdateMessage,
|
||||
update_manager::{GetCloudObjectResponse, InitialLoadResponse},
|
||||
},
|
||||
ids::{ServerId, ServerIdAndType, SyncId},
|
||||
server_api::object::{GuestIdentifier, ObjectClient},
|
||||
sync_queue::SerializedModel,
|
||||
},
|
||||
settings::cloud_preferences::{CloudPreferenceModel, Platform, Preference},
|
||||
workflows::WorkflowId,
|
||||
};
|
||||
|
||||
/// A stateful fake cloud preferences backend.
|
||||
#[derive(Clone, Default)]
|
||||
pub struct FakeObjectClient {
|
||||
state: Arc<Mutex<FakeCloudState>>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct FakeCloudState {
|
||||
/// Stored preferences keyed by `GenericStringObjectId`. The
|
||||
/// preferences syncer addresses updates and deletes by id, so this
|
||||
/// is the authoritative lookup.
|
||||
preferences: HashMap<GenericStringObjectId, StoredPreference>,
|
||||
/// Monotonically increasing counter used to allocate fresh
|
||||
/// [`ServerId`]s for newly created preferences.
|
||||
next_server_id: i64,
|
||||
}
|
||||
|
||||
struct StoredPreference {
|
||||
model: CloudPreferenceModel,
|
||||
revision: Revision,
|
||||
metadata_ts: chrono::DateTime<Utc>,
|
||||
}
|
||||
|
||||
impl FakeObjectClient {
|
||||
/// Seeds a preference into the cloud store as if another client had
|
||||
/// already written it. Returns the allocated [`ServerId`] so tests
|
||||
/// can reference it later if needed.
|
||||
///
|
||||
/// The `value_json` argument is the JSON-serialized setting value
|
||||
/// (e.g. `"14.0"` for a float, `"true"` for a bool, `"\"Hack\""`
|
||||
/// for a string). This matches the format used by
|
||||
/// [`Preference::new`].
|
||||
pub fn seed_preference(
|
||||
&self,
|
||||
storage_key: &str,
|
||||
value_json: &str,
|
||||
platform: Platform,
|
||||
) -> ServerId {
|
||||
let preference = build_preference(storage_key, value_json, platform);
|
||||
let model = CloudPreferenceModel::new(preference);
|
||||
|
||||
let mut state = self.state.lock().unwrap();
|
||||
let server_id = state.alloc_server_id();
|
||||
state.preferences.insert(
|
||||
GenericStringObjectId::from(server_id),
|
||||
StoredPreference {
|
||||
model,
|
||||
revision: Revision::now(),
|
||||
metadata_ts: Utc::now(),
|
||||
},
|
||||
);
|
||||
server_id
|
||||
}
|
||||
|
||||
/// Returns the current cloud value for the preference with the
|
||||
/// given storage key and platform, or `None` if the fake has no
|
||||
/// such preference.
|
||||
///
|
||||
/// The returned string is the JSON-serialized setting value (same
|
||||
/// shape as [`Self::seed_preference`]).
|
||||
pub fn cloud_value(&self, storage_key: &str, platform: Platform) -> Option<String> {
|
||||
let state = self.state.lock().unwrap();
|
||||
state
|
||||
.preferences
|
||||
.values()
|
||||
.find(|stored| {
|
||||
stored.model.string_model.storage_key == storage_key
|
||||
&& stored.model.string_model.platform == platform
|
||||
})
|
||||
.map(|stored| stored.model.string_model.value.to_string())
|
||||
}
|
||||
|
||||
/// Builds an [`InitialLoadResponse`] that reflects the current
|
||||
/// state of the fake. Tests pass the return value to
|
||||
/// [`UpdateManager::mock_initial_load`] to simulate the syncer's
|
||||
/// initial load seeing the seeded cloud state.
|
||||
pub fn snapshot_as_initial_load_response(&self) -> InitialLoadResponse {
|
||||
let state = self.state.lock().unwrap();
|
||||
let server_objects: Vec<Box<dyn ServerObject>> = state
|
||||
.preferences
|
||||
.iter()
|
||||
.map(|(id, stored)| {
|
||||
let metadata = ServerMetadata {
|
||||
uid: ServerId::from(*id),
|
||||
revision: stored.revision.clone(),
|
||||
metadata_last_updated_ts: stored.metadata_ts.into(),
|
||||
trashed_ts: None,
|
||||
folder_id: None,
|
||||
is_welcome_object: false,
|
||||
creator_uid: None,
|
||||
last_editor_uid: None,
|
||||
current_editor_uid: None,
|
||||
};
|
||||
let permissions = ServerPermissions {
|
||||
space: Owner::mock_current_user(),
|
||||
guests: Vec::new(),
|
||||
anyone_link_sharing: None,
|
||||
permissions_last_updated_ts: stored.metadata_ts.into(),
|
||||
};
|
||||
let server_pref = ServerPreference {
|
||||
id: SyncId::ServerId(ServerId::from(*id)),
|
||||
model: stored.model.clone(),
|
||||
metadata,
|
||||
permissions,
|
||||
};
|
||||
Box::new(server_pref) as Box<dyn ServerObject>
|
||||
})
|
||||
.collect();
|
||||
|
||||
let mut response = InitialLoadResponse::default();
|
||||
if !server_objects.is_empty() {
|
||||
response.updated_generic_string_objects.insert(
|
||||
GenericStringObjectFormat::Json(JsonObjectType::Preference),
|
||||
server_objects,
|
||||
);
|
||||
}
|
||||
response
|
||||
}
|
||||
}
|
||||
|
||||
impl FakeCloudState {
|
||||
fn alloc_server_id(&mut self) -> ServerId {
|
||||
self.next_server_id += 1;
|
||||
ServerId::from(self.next_server_id)
|
||||
}
|
||||
}
|
||||
|
||||
/// Builds a [`Preference`] struct from a storage key, JSON value, and
|
||||
/// explicit platform. This bypasses [`Preference::new`]'s syncing-mode
|
||||
/// inference so tests can seed preferences with an arbitrary platform.
|
||||
fn build_preference(storage_key: &str, value_json: &str, platform: Platform) -> Preference {
|
||||
let value: serde_json::Value = serde_json::from_str(value_json)
|
||||
.unwrap_or_else(|err| panic!("invalid JSON value {value_json:?}: {err}"));
|
||||
Preference {
|
||||
storage_key: storage_key.to_owned(),
|
||||
value,
|
||||
platform,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
impl ObjectClient for FakeObjectClient {
|
||||
async fn bulk_create_generic_string_objects(
|
||||
&self,
|
||||
_owner: Owner,
|
||||
objects: &[BulkCreateGenericStringObjectsRequest],
|
||||
) -> Result<BulkCreateCloudObjectResult> {
|
||||
let mut state = self.state.lock().unwrap();
|
||||
let mut created = Vec::with_capacity(objects.len());
|
||||
for request in objects {
|
||||
let serialized = request.serialized_model.model_as_str();
|
||||
let model = CloudPreferenceModel::deserialize_owned(serialized)
|
||||
.map_err(|err| anyhow!("fake cloud failed to deserialize preference: {err}"))?;
|
||||
let server_id = state.alloc_server_id();
|
||||
let now = Utc::now();
|
||||
state.preferences.insert(
|
||||
GenericStringObjectId::from(server_id),
|
||||
StoredPreference {
|
||||
model,
|
||||
revision: Revision::now(),
|
||||
metadata_ts: now,
|
||||
},
|
||||
);
|
||||
created.push(CreatedCloudObject {
|
||||
client_id: request.id,
|
||||
revision_and_editor: RevisionAndLastEditor {
|
||||
revision: Revision::now(),
|
||||
last_editor_uid: None,
|
||||
},
|
||||
metadata_ts: now.into(),
|
||||
server_id_and_type: ServerIdAndType {
|
||||
id: server_id,
|
||||
id_type: ObjectIdType::GenericStringObject,
|
||||
},
|
||||
creator_uid: None,
|
||||
permissions: ServerPermissions::mock_personal(),
|
||||
});
|
||||
}
|
||||
Ok(BulkCreateCloudObjectResult::Success {
|
||||
created_cloud_objects: created,
|
||||
})
|
||||
}
|
||||
|
||||
async fn update_generic_string_object(
|
||||
&self,
|
||||
object_id: GenericStringObjectId,
|
||||
model: SerializedModel,
|
||||
_revision: Option<Revision>,
|
||||
) -> Result<UpdateCloudObjectResult<Box<dyn ServerObject>>> {
|
||||
let mut state = self.state.lock().unwrap();
|
||||
let stored = state
|
||||
.preferences
|
||||
.get_mut(&object_id)
|
||||
.ok_or_else(|| anyhow!("fake cloud: no preference with id {object_id:?}"))?;
|
||||
stored.model = CloudPreferenceModel::deserialize_owned(model.model_as_str())
|
||||
.map_err(|err| anyhow!("fake cloud failed to deserialize preference: {err}"))?;
|
||||
stored.revision = Revision::now();
|
||||
stored.metadata_ts = Utc::now();
|
||||
Ok(UpdateCloudObjectResult::Success {
|
||||
revision_and_editor: RevisionAndLastEditor {
|
||||
revision: stored.revision.clone(),
|
||||
last_editor_uid: None,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
async fn delete_object(&self, id: ServerId) -> Result<ObjectDeleteResult> {
|
||||
let mut state = self.state.lock().unwrap();
|
||||
state.preferences.remove(&GenericStringObjectId::from(id));
|
||||
Ok(ObjectDeleteResult::Success {
|
||||
deleted_ids: vec![SyncId::ServerId(id)],
|
||||
})
|
||||
}
|
||||
|
||||
async fn fetch_changed_objects(
|
||||
&self,
|
||||
_objects_to_update: ObjectsToUpdate,
|
||||
_force_refresh: bool,
|
||||
) -> Result<InitialLoadResponse> {
|
||||
Ok(self.snapshot_as_initial_load_response())
|
||||
}
|
||||
|
||||
async fn fetch_environment_last_task_run_timestamps(
|
||||
&self,
|
||||
) -> Result<HashMap<String, DateTime<Utc>>> {
|
||||
Ok(HashMap::new())
|
||||
}
|
||||
|
||||
// ───────────────────────────────────────────────────────────────
|
||||
// The methods below are not exercised by CloudPreferencesSyncer,
|
||||
// so they intentionally panic. If a future change to the syncer
|
||||
// starts calling one of these, the test that triggered the call
|
||||
// will fail loudly with a clear message rather than silently
|
||||
// misbehave.
|
||||
// ───────────────────────────────────────────────────────────────
|
||||
|
||||
async fn create_workflow(
|
||||
&self,
|
||||
_request: CreateObjectRequest,
|
||||
) -> Result<CreateCloudObjectResult> {
|
||||
unimplemented!("FakeObjectClient::create_workflow")
|
||||
}
|
||||
|
||||
async fn update_workflow(
|
||||
&self,
|
||||
_workflow_id: WorkflowId,
|
||||
_data: SerializedModel,
|
||||
_revision: Option<Revision>,
|
||||
) -> Result<UpdateCloudObjectResult<ServerWorkflow>> {
|
||||
unimplemented!("FakeObjectClient::update_workflow")
|
||||
}
|
||||
|
||||
async fn create_generic_string_object(
|
||||
&self,
|
||||
_format: GenericStringObjectFormat,
|
||||
_uniqueness_key: Option<GenericStringObjectUniqueKey>,
|
||||
_request: CreateObjectRequest,
|
||||
) -> Result<CreateCloudObjectResult> {
|
||||
unimplemented!("FakeObjectClient::create_generic_string_object")
|
||||
}
|
||||
|
||||
async fn create_notebook(
|
||||
&self,
|
||||
_request: CreateObjectRequest,
|
||||
) -> Result<CreateCloudObjectResult> {
|
||||
unimplemented!("FakeObjectClient::create_notebook")
|
||||
}
|
||||
|
||||
async fn update_notebook(
|
||||
&self,
|
||||
_notebook_id: NotebookId,
|
||||
_title: Option<String>,
|
||||
_data: Option<SerializedModel>,
|
||||
_revision: Option<Revision>,
|
||||
) -> Result<UpdateCloudObjectResult<ServerNotebook>> {
|
||||
unimplemented!("FakeObjectClient::update_notebook")
|
||||
}
|
||||
|
||||
async fn create_folder(
|
||||
&self,
|
||||
_request: CreateObjectRequest,
|
||||
) -> Result<CreateCloudObjectResult> {
|
||||
unimplemented!("FakeObjectClient::create_folder")
|
||||
}
|
||||
|
||||
async fn update_folder(
|
||||
&self,
|
||||
_folder_id: FolderId,
|
||||
_name: SerializedModel,
|
||||
) -> Result<UpdateCloudObjectResult<ServerFolder>> {
|
||||
unimplemented!("FakeObjectClient::update_folder")
|
||||
}
|
||||
|
||||
async fn grab_notebook_edit_access(&self, _notebook_id: NotebookId) -> Result<ServerMetadata> {
|
||||
unimplemented!("FakeObjectClient::grab_notebook_edit_access")
|
||||
}
|
||||
|
||||
async fn give_up_notebook_edit_access(
|
||||
&self,
|
||||
_notebook_id: NotebookId,
|
||||
) -> Result<ServerMetadata> {
|
||||
unimplemented!("FakeObjectClient::give_up_notebook_edit_access")
|
||||
}
|
||||
|
||||
async fn get_warp_drive_updates(
|
||||
&self,
|
||||
_message_sender: Sender<ObjectUpdateMessage>,
|
||||
_stream_ready_sender: Sender<()>,
|
||||
) -> Result<()> {
|
||||
unimplemented!("FakeObjectClient::get_warp_drive_updates")
|
||||
}
|
||||
|
||||
async fn fetch_single_cloud_object(&self, _id: ServerId) -> Result<GetCloudObjectResponse> {
|
||||
unimplemented!("FakeObjectClient::fetch_single_cloud_object")
|
||||
}
|
||||
|
||||
async fn transfer_notebook_owner(
|
||||
&self,
|
||||
_notebook_id: NotebookId,
|
||||
_owner: Owner,
|
||||
) -> Result<bool> {
|
||||
unimplemented!("FakeObjectClient::transfer_notebook_owner")
|
||||
}
|
||||
|
||||
async fn transfer_workflow_owner(
|
||||
&self,
|
||||
_workflow_id: WorkflowId,
|
||||
_owner: Owner,
|
||||
) -> Result<bool> {
|
||||
unimplemented!("FakeObjectClient::transfer_workflow_owner")
|
||||
}
|
||||
|
||||
async fn transfer_generic_string_object_owner(
|
||||
&self,
|
||||
_id: GenericStringObjectId,
|
||||
_owner: Owner,
|
||||
) -> Result<bool> {
|
||||
unimplemented!("FakeObjectClient::transfer_generic_string_object_owner")
|
||||
}
|
||||
|
||||
async fn trash_object(&self, _id: ServerId) -> Result<bool> {
|
||||
unimplemented!("FakeObjectClient::trash_object")
|
||||
}
|
||||
|
||||
async fn untrash_object(&self, _id: ServerId) -> Result<ObjectMetadataUpdateResult> {
|
||||
unimplemented!("FakeObjectClient::untrash_object")
|
||||
}
|
||||
|
||||
async fn empty_trash(&self, _owner: Owner) -> Result<ObjectDeleteResult> {
|
||||
unimplemented!("FakeObjectClient::empty_trash")
|
||||
}
|
||||
|
||||
async fn move_object(
|
||||
&self,
|
||||
_id: ServerId,
|
||||
_folder_id: Option<FolderId>,
|
||||
_owner: Owner,
|
||||
_object_type: ObjectType,
|
||||
) -> Result<bool> {
|
||||
unimplemented!("FakeObjectClient::move_object")
|
||||
}
|
||||
|
||||
async fn record_object_action(
|
||||
&self,
|
||||
_id: ServerId,
|
||||
_action_type: ObjectActionType,
|
||||
_timestamp: DateTime<Utc>,
|
||||
_data: Option<String>,
|
||||
) -> Result<ObjectActionHistory> {
|
||||
unimplemented!("FakeObjectClient::record_object_action")
|
||||
}
|
||||
|
||||
async fn leave_object(&self, _id: ServerId) -> Result<ObjectDeleteResult> {
|
||||
unimplemented!("FakeObjectClient::leave_object")
|
||||
}
|
||||
|
||||
async fn set_object_link_permissions(
|
||||
&self,
|
||||
_object_id: ServerId,
|
||||
_access_level: SharingAccessLevel,
|
||||
) -> Result<ObjectPermissionUpdateResult> {
|
||||
unimplemented!("FakeObjectClient::set_object_link_permissions")
|
||||
}
|
||||
|
||||
async fn remove_object_link_permissions(
|
||||
&self,
|
||||
_object_id: ServerId,
|
||||
) -> Result<ObjectPermissionUpdateResult> {
|
||||
unimplemented!("FakeObjectClient::remove_object_link_permissions")
|
||||
}
|
||||
|
||||
async fn add_object_guests(
|
||||
&self,
|
||||
_object_id: ServerId,
|
||||
_guest_emails: Vec<String>,
|
||||
_access_level: AccessLevel,
|
||||
) -> Result<ObjectPermissionsUpdateData> {
|
||||
unimplemented!("FakeObjectClient::add_object_guests")
|
||||
}
|
||||
|
||||
async fn update_object_guests(
|
||||
&self,
|
||||
_object_id: ServerId,
|
||||
_guest_emails: Vec<String>,
|
||||
_access_level: AccessLevel,
|
||||
) -> Result<ServerPermissions> {
|
||||
unimplemented!("FakeObjectClient::update_object_guests")
|
||||
}
|
||||
|
||||
async fn remove_object_guest(
|
||||
&self,
|
||||
_object_id: ServerId,
|
||||
_guest: GuestIdentifier,
|
||||
) -> Result<ServerPermissions> {
|
||||
unimplemented!("FakeObjectClient::remove_object_guest")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,438 @@
|
||||
use crate::network::{NetworkStatus, NetworkStatusEvent, NetworkStatusKind};
|
||||
use crate::report_error;
|
||||
use crate::server::{ids::ServerId, retry_strategies::LISTENER_RETRY_STRATEGY};
|
||||
use crate::system::{SystemStats, SystemStatsEvent};
|
||||
use crate::workspaces::{
|
||||
user_profiles::UserProfileWithUID,
|
||||
user_workspaces::{UserWorkspaces, UserWorkspacesEvent},
|
||||
};
|
||||
use crate::{
|
||||
cloud_object::{
|
||||
model::{
|
||||
actions::ObjectActionHistory,
|
||||
persistence::{CloudModel, CloudModelEvent},
|
||||
},
|
||||
ServerCloudObject, ServerMetadata, ServerPermissions,
|
||||
},
|
||||
server::server_api::object::ObjectClient,
|
||||
};
|
||||
|
||||
use super::update_manager::UpdateManager;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use futures_util::stream::AbortHandle;
|
||||
use std::time::Duration;
|
||||
use warpui::r#async::Timer;
|
||||
|
||||
use async_channel::Sender;
|
||||
|
||||
use std::sync::Arc;
|
||||
use warpui::{Entity, ModelContext, RequestState, SingletonEntity};
|
||||
|
||||
use instant::Instant;
|
||||
|
||||
lazy_static::lazy_static! {
|
||||
/// Between successful websocket connections, we ensured at least this amount of time
|
||||
/// has elapsed so that we aren't spamming the websocket server (e.g. if connections are being
|
||||
/// closed quickly for any reason).
|
||||
static ref WAIT_PERIOD_BETWEEN_SUCCESSFUL_RECONNECTS: Duration = Duration::from_secs(30);
|
||||
}
|
||||
|
||||
/// If the websocket reconnects within this duration of the last disconnection, skip the
|
||||
/// out-of-band refresh of cloud objects. The periodic poll will catch any missed updates.
|
||||
///
|
||||
/// This needs to be relatively large, due to the retry policy we use on websocket disconnection,
|
||||
/// which waits between 10-40s between retries. At the very least, this should always be slightly
|
||||
/// larger than the upper end of that range.
|
||||
const RECONNECTION_REFRESH_THRESHOLD: Duration = Duration::from_secs(60);
|
||||
|
||||
/// Maximum random delay added before making an out-of-band refresh after a longer reconnection.
|
||||
/// Spreading out requests across this window helps avoid a thundering herd when many clients
|
||||
/// reconnect simultaneously (e.g. after a server release).
|
||||
const MAX_RECONNECTION_REFRESH_DELAY: Duration = Duration::from_secs(30);
|
||||
|
||||
/// Describes the type of websocket connection that was just established.
|
||||
enum ConnectionEvent {
|
||||
/// The very first websocket connection after application startup.
|
||||
InitialConnection,
|
||||
/// A reconnection after a previous disconnection.
|
||||
Reconnection {
|
||||
/// The duration since the last websocket disconnection.
|
||||
time_since_disconnection: Duration,
|
||||
},
|
||||
}
|
||||
|
||||
pub enum ListenerEvent {}
|
||||
|
||||
/// The Listener is responsible for listening to updates from
|
||||
/// the server for cloud-object related things (e.g. a notebook was changed,
|
||||
/// or edit access was taken for a workflow, etc.)
|
||||
pub struct Listener {
|
||||
cloud_objects_client: Arc<dyn ObjectClient>,
|
||||
/// Since we only want to start websocket connections if we know the user is
|
||||
/// on a team or has access to cloud objects, we keep track of whether
|
||||
/// or not we should be subscribing for updates. Once we start websockets, we don't stop
|
||||
/// so that the user gets a snappier experience once they start using Warp Drive.
|
||||
should_subscribe_to_updates: bool,
|
||||
/// Abort handle for the (retried) future that resolves when the subscription is done.
|
||||
current_subscription_abort_handle: Option<AbortHandle>,
|
||||
/// Channel that we send a message over each time we've successfully established a subscription.
|
||||
subscription_ready_tx: Sender<()>,
|
||||
/// The time at which the last websocket disconnection occurred. `None` if no disconnection
|
||||
/// has occurred yet (i.e., this is the first connection attempt).
|
||||
last_disconnected_at: Option<Instant>,
|
||||
/// Abort handle for a pending delayed refresh spawned after a long reconnection. Tracked so
|
||||
/// that it can be cancelled if the websocket disconnects again before the refresh fires.
|
||||
pending_refresh_abort_handle: Option<AbortHandle>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
#[allow(clippy::enum_variant_names)]
|
||||
pub enum ObjectUpdateMessage {
|
||||
ObjectMetadataChanged {
|
||||
metadata: ServerMetadata,
|
||||
},
|
||||
ObjectPermissionsChanged,
|
||||
// TODO(CLD-2425): Replace `ObjectPermissionsChanged` with this.
|
||||
ObjectPermissionsChangedV2 {
|
||||
object_uid: ServerId,
|
||||
permissions: ServerPermissions,
|
||||
user_profiles: Vec<UserProfileWithUID>,
|
||||
},
|
||||
ObjectContentChanged {
|
||||
server_object: Box<ServerCloudObject>,
|
||||
last_editor: Option<UserProfileWithUID>,
|
||||
},
|
||||
ObjectDeleted {
|
||||
object_uid: ServerId,
|
||||
},
|
||||
ObjectActionOccurred {
|
||||
history: ObjectActionHistory,
|
||||
},
|
||||
TeamMembershipsChanged,
|
||||
AmbientTaskUpdated {
|
||||
task_id: String,
|
||||
timestamp: DateTime<Utc>,
|
||||
},
|
||||
}
|
||||
|
||||
impl ObjectUpdateMessage {
|
||||
fn as_str(&self) -> &'static str {
|
||||
use ObjectUpdateMessage::*;
|
||||
match self {
|
||||
ObjectMetadataChanged { .. } => "ObjectMetadataChanged",
|
||||
ObjectPermissionsChanged => "ObjectPermissionsChanged",
|
||||
ObjectPermissionsChangedV2 { .. } => "ObjectPermissionsChanged (V2)",
|
||||
ObjectContentChanged { .. } => "ObjectContentChanged",
|
||||
ObjectDeleted { .. } => "ObjectDeleted",
|
||||
ObjectActionOccurred { .. } => "ObjectActionOccurred",
|
||||
TeamMembershipsChanged => "TeamMembershipsChanged",
|
||||
AmbientTaskUpdated { .. } => "AmbientTaskUpdated",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Listener {
|
||||
pub fn new(cloud_objects_client: Arc<dyn ObjectClient>, ctx: &mut ModelContext<Self>) -> Self {
|
||||
let (subscription_ready_tx, subscription_ready_rx) = async_channel::unbounded();
|
||||
let mut listener = Self {
|
||||
cloud_objects_client,
|
||||
should_subscribe_to_updates: false,
|
||||
current_subscription_abort_handle: None,
|
||||
subscription_ready_tx,
|
||||
last_disconnected_at: None,
|
||||
pending_refresh_abort_handle: None,
|
||||
};
|
||||
|
||||
// When the websocket signals readiness, decide whether to refresh cloud objects
|
||||
// based on how long the connection was down.
|
||||
let _ = ctx.spawn_stream_local(
|
||||
subscription_ready_rx,
|
||||
Self::on_subscription_ready,
|
||||
|_, _| {},
|
||||
);
|
||||
|
||||
ctx.subscribe_to_model(&SystemStats::handle(ctx), Self::handle_cpu_event);
|
||||
|
||||
ctx.subscribe_to_model(
|
||||
&NetworkStatus::handle(ctx),
|
||||
Self::handle_network_status_changed_event,
|
||||
);
|
||||
|
||||
// To prevent creating unnecessary websockets, we only open a websocket if
|
||||
// - a user is known to be part of a team
|
||||
// - or a user has access to >= 1 cloud object
|
||||
// In either of these cases, it's worth creating a websocket for cloud object updates.
|
||||
//
|
||||
// Note that we also want a websocket for CloudPreferences, but this is handled via listening
|
||||
// to the cloud model for the creation of cloud preferences objects (which happens when settings sync
|
||||
// is enabled for the first time).
|
||||
ctx.subscribe_to_model(
|
||||
&UserWorkspaces::handle(ctx),
|
||||
Self::handle_user_workspaces_event,
|
||||
);
|
||||
ctx.subscribe_to_model(&CloudModel::handle(ctx), Self::handle_cloud_model_event);
|
||||
|
||||
// We need to do a one-time check of cloud objects when starting
|
||||
// because the Cloud Model was initialized before this model and we could have populated
|
||||
// its object cache with objects from sqlite.
|
||||
if listener.has_non_welcome_cloud_objects(ctx) {
|
||||
listener.start_listener(ctx);
|
||||
}
|
||||
|
||||
listener
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn mock(ctx: &mut ModelContext<Self>) -> Self {
|
||||
use crate::server::server_api::ServerApiProvider;
|
||||
|
||||
Self::new(ServerApiProvider::new_for_test().get(), ctx)
|
||||
}
|
||||
|
||||
fn is_part_of_some_team(&self, ctx: &ModelContext<Self>) -> bool {
|
||||
UserWorkspaces::as_ref(ctx).has_teams()
|
||||
}
|
||||
|
||||
// If the user is part of a team, we should start subscribing for updates.
|
||||
fn handle_user_workspaces_event(
|
||||
&mut self,
|
||||
event: &UserWorkspacesEvent,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
if let UserWorkspacesEvent::TeamsChanged = event {
|
||||
if self.is_part_of_some_team(ctx) {
|
||||
self.start_listener(ctx);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns true if the user has any object that is not a welcome object. If the user only has objects
|
||||
/// that are welcome objects, returns false.
|
||||
fn has_non_welcome_cloud_objects(&self, ctx: &ModelContext<Self>) -> bool {
|
||||
CloudModel::as_ref(ctx).has_non_welcome_objects()
|
||||
}
|
||||
|
||||
// If the user has access to >= 1 cloud objects, we should subscribe for updates.
|
||||
fn handle_cloud_model_event(&mut self, _event: &CloudModelEvent, ctx: &mut ModelContext<Self>) {
|
||||
if self.has_non_welcome_cloud_objects(ctx) {
|
||||
self.start_listener(ctx);
|
||||
}
|
||||
}
|
||||
|
||||
// This is a workaround for an issue where the future that should resolve when the websocket
|
||||
// is finished is _not_ polled when the websocket is closed by the server and the CPU is asleep.
|
||||
// To get around this, we manually abort the future (effectively closing the websocket)
|
||||
// when the CPU goes to sleep and restart it when it's awakened.
|
||||
// https://linear.app/warpdotdev/issue/CLD-172/websocket-hangs-when-closed-during-cpu-sleep
|
||||
fn handle_cpu_event(&mut self, event: &SystemStatsEvent, ctx: &mut ModelContext<Self>) {
|
||||
match event {
|
||||
SystemStatsEvent::CpuWasAwakened => {
|
||||
if let Some(abort_handle) = self.current_subscription_abort_handle.take() {
|
||||
abort_handle.abort();
|
||||
}
|
||||
|
||||
// We intentionally do not update `last_disconnected_at` or cancel pending
|
||||
// refreshes here. The paired `CpuWillSleep` event already handled both;
|
||||
// this handler just restarts the websocket so that `on_subscription_ready`
|
||||
// can decide whether to refresh based on the sleep-time gap.
|
||||
if self.should_subscribe_to_updates {
|
||||
self.get_warp_drive_updates(ctx);
|
||||
}
|
||||
}
|
||||
SystemStatsEvent::CpuWillSleep => {
|
||||
if let Some(abort_handle) = self.current_subscription_abort_handle.take() {
|
||||
abort_handle.abort();
|
||||
self.last_disconnected_at = Some(Instant::now());
|
||||
self.cancel_pending_refresh();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_network_status_changed_event(
|
||||
&mut self,
|
||||
event: &NetworkStatusEvent,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
match event {
|
||||
NetworkStatusEvent::NetworkStatusChanged { new_status } => match new_status {
|
||||
// When coming back online, restart a websocket.
|
||||
NetworkStatusKind::Online => {
|
||||
if let Some(abort_handle) = self.current_subscription_abort_handle.take() {
|
||||
abort_handle.abort();
|
||||
}
|
||||
|
||||
if self.should_subscribe_to_updates {
|
||||
self.get_warp_drive_updates(ctx);
|
||||
}
|
||||
}
|
||||
|
||||
// When losing connection, abort the current subscription to avoid a lingering future that doesn't resolve.
|
||||
NetworkStatusKind::Offline => {
|
||||
if let Some(abort_handle) = self.current_subscription_abort_handle.take() {
|
||||
abort_handle.abort();
|
||||
self.last_disconnected_at = Some(Instant::now());
|
||||
self.cancel_pending_refresh();
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn start_listener(&mut self, ctx: &mut ModelContext<Self>) {
|
||||
if !self.should_subscribe_to_updates {
|
||||
self.should_subscribe_to_updates = true;
|
||||
self.get_warp_drive_updates(ctx);
|
||||
}
|
||||
}
|
||||
|
||||
/// Cancels any pending delayed refresh that was scheduled after a reconnection.
|
||||
fn cancel_pending_refresh(&mut self) {
|
||||
if let Some(abort_handle) = self.pending_refresh_abort_handle.take() {
|
||||
abort_handle.abort();
|
||||
}
|
||||
}
|
||||
|
||||
/// Called each time the websocket signals readiness. Decides whether to trigger an
|
||||
/// out-of-band refresh of cloud objects based on how long the connection was down.
|
||||
fn on_subscription_ready(&mut self, _: (), ctx: &mut ModelContext<Self>) {
|
||||
// Cancel any pending refresh from a previous reconnection to avoid accumulating
|
||||
// stale refresh requests if the websocket is rapidly cycling.
|
||||
self.cancel_pending_refresh();
|
||||
|
||||
let connection_event = match self.last_disconnected_at {
|
||||
None => ConnectionEvent::InitialConnection,
|
||||
Some(disconnected_at) => ConnectionEvent::Reconnection {
|
||||
time_since_disconnection: disconnected_at.elapsed(),
|
||||
},
|
||||
};
|
||||
|
||||
match connection_event {
|
||||
ConnectionEvent::InitialConnection => {
|
||||
// No out-of-band refresh needed for the initial connection. The periodic
|
||||
// poll (started by TeamTesterStatus) already fetches cloud objects at
|
||||
// startup, so an additional request here would be duplicative.
|
||||
log::info!(
|
||||
"Initial websocket connection established; skipping out-of-band refresh."
|
||||
);
|
||||
}
|
||||
ConnectionEvent::Reconnection {
|
||||
time_since_disconnection,
|
||||
} if time_since_disconnection < RECONNECTION_REFRESH_THRESHOLD => {
|
||||
log::info!(
|
||||
"Websocket reconnected after {time_since_disconnection:?}, within the \
|
||||
{RECONNECTION_REFRESH_THRESHOLD:?} refresh threshold; \
|
||||
skipping out-of-band refresh.",
|
||||
);
|
||||
}
|
||||
ConnectionEvent::Reconnection {
|
||||
time_since_disconnection,
|
||||
} => {
|
||||
// Add a random delay to avoid a thundering herd when many clients reconnect
|
||||
// simultaneously (e.g. after a server release).
|
||||
let delay = MAX_RECONNECTION_REFRESH_DELAY.mul_f32(rand::random::<f32>());
|
||||
log::info!(
|
||||
"Websocket reconnected after {time_since_disconnection:?}; \
|
||||
refreshing objects after {delay:?} delay.",
|
||||
);
|
||||
|
||||
let handle = ctx.spawn(async move { Timer::after(delay).await }, |me, _, ctx| {
|
||||
me.pending_refresh_abort_handle = None;
|
||||
UpdateManager::handle(ctx).update(ctx, |update_manager, ctx| {
|
||||
update_manager.refresh_updated_objects(ctx);
|
||||
});
|
||||
});
|
||||
self.pending_refresh_abort_handle = Some(handle.abort_handle());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn get_warp_drive_updates(&mut self, ctx: &mut ModelContext<Self>) {
|
||||
let object_client = self.cloud_objects_client.clone();
|
||||
let (message_sender, message_receiver) = async_channel::unbounded();
|
||||
let subscription_ready_tx = self.subscription_ready_tx.clone();
|
||||
|
||||
// On every message we receive (over the message_receiver), send it
|
||||
// to the UpdateManager.
|
||||
let _ = ctx.spawn_stream_local(
|
||||
message_receiver,
|
||||
|_me, item: ObjectUpdateMessage, ctx| {
|
||||
log::info!(
|
||||
"Received {} message in CloudObjects::Listener",
|
||||
item.as_str()
|
||||
);
|
||||
UpdateManager::handle(ctx).update(ctx, move |update_manager, ctx| {
|
||||
update_manager.received_message_from_server(item, ctx);
|
||||
})
|
||||
},
|
||||
|_, _| {},
|
||||
);
|
||||
|
||||
// Start the future that sends messages over the message_sender stream.
|
||||
// TODO: we should investigate having get_warp_drive_updates (and in turn,
|
||||
// start_graphql_streaming_operation) return an `impl Stream` so that we don't
|
||||
// need to spawn and then spawn_stream_local. For this, we'll need an equivalent
|
||||
// spawn_stream (which is like spawn_stream_local but polls the futures in the stream
|
||||
// on a background thread).
|
||||
let spawn_handle = ctx.spawn_with_retry_on_error(
|
||||
move || {
|
||||
let object_client = object_client.clone();
|
||||
let message_sender = message_sender.clone();
|
||||
let subscription_ready_tx = subscription_ready_tx.clone();
|
||||
async move {
|
||||
let start_time = Instant::now();
|
||||
log::info!("Attempting to start websocket connection in CloudObjects::Listener");
|
||||
let res = object_client
|
||||
.get_warp_drive_updates(
|
||||
message_sender,
|
||||
subscription_ready_tx,
|
||||
).await;
|
||||
res.map(|_| start_time.elapsed())
|
||||
}
|
||||
},
|
||||
LISTENER_RETRY_STRATEGY,
|
||||
|me, req_state, ctx| {
|
||||
match req_state {
|
||||
RequestState::RequestSucceeded(elapsed_time) => {
|
||||
// Record the disconnection time now that a live connection has ended.
|
||||
// Only set this here (not on failed retries) so that the elapsed time
|
||||
// accurately reflects the full duration since the real disconnection.
|
||||
me.last_disconnected_at = Some(Instant::now());
|
||||
me.cancel_pending_refresh();
|
||||
// The future only resolves once the stream is done, so
|
||||
// at that point, we should restart the stream.
|
||||
// In case the websocket was closed quickly by the server,
|
||||
// let's ensure at least some time has passed before we restart the connection.
|
||||
let time_to_wait = (*WAIT_PERIOD_BETWEEN_SUCCESSFUL_RECONNECTS).saturating_sub(elapsed_time);
|
||||
log::info!("Websocket for CloudObjects::Listener is done; restarting after {}s.", time_to_wait.as_secs());
|
||||
ctx.spawn(async move {
|
||||
Timer::after(time_to_wait).await
|
||||
}, |me, _, ctx| {
|
||||
me.get_warp_drive_updates(ctx);
|
||||
});
|
||||
}
|
||||
RequestState::RequestFailedRetryPending(e) => {
|
||||
log::warn!("CloudObjects::Listener: websocket connection failed to connect or finished with an error; trying again: {e:#}");
|
||||
}
|
||||
RequestState::RequestFailed(e) => {
|
||||
report_error!(e.context("CloudObjects::Listener websocket connection failed"));
|
||||
}
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
self.current_subscription_abort_handle = Some(spawn_handle.abort_handle());
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn has_current_subscription_abort_handle(&self) -> bool {
|
||||
self.current_subscription_abort_handle.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
impl Entity for Listener {
|
||||
type Event = ListenerEvent;
|
||||
}
|
||||
|
||||
impl SingletonEntity for Listener {}
|
||||
@@ -0,0 +1,6 @@
|
||||
#[cfg(test)]
|
||||
pub mod fake_object_client;
|
||||
pub mod listener;
|
||||
#[cfg(test)]
|
||||
pub mod test_utils;
|
||||
pub mod update_manager;
|
||||
@@ -0,0 +1,123 @@
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
sync::{
|
||||
mpsc::{sync_channel, Receiver},
|
||||
Arc,
|
||||
},
|
||||
};
|
||||
|
||||
use settings::manager::SettingsManager;
|
||||
use warp_core::execution_mode::{AppExecutionMode, ExecutionMode};
|
||||
use warpui::{App, ModelHandle, SingletonEntity};
|
||||
|
||||
use crate::{
|
||||
auth::{auth_manager::AuthManager, AuthStateProvider},
|
||||
cloud_object::model::{
|
||||
actions::ObjectActions,
|
||||
persistence::{CloudModel, CloudModelEvent},
|
||||
},
|
||||
network::NetworkStatus,
|
||||
persistence::ModelEvent,
|
||||
server::{
|
||||
server_api::{
|
||||
object::{MockObjectClient, ObjectClient},
|
||||
ServerApiProvider,
|
||||
},
|
||||
sync_queue::SyncQueue,
|
||||
telemetry::context_provider::AppTelemetryContextProvider,
|
||||
},
|
||||
settings::{PrivacySettings, WarpDrivePrivacySettings},
|
||||
workspaces::{
|
||||
team_tester::TeamTesterStatus, update_manager::TeamUpdateManager,
|
||||
user_profiles::UserProfiles, user_workspaces::UserWorkspaces,
|
||||
},
|
||||
};
|
||||
|
||||
use super::update_manager::UpdateManager;
|
||||
|
||||
/// The size of the bounded channel that we use to queue persistence/sqlite-related events.
|
||||
const CHANNEL_SIZE: usize = 128;
|
||||
|
||||
pub struct UpdateManagerStruct {
|
||||
pub update_manager: ModelHandle<UpdateManager>,
|
||||
pub receiver: Receiver<ModelEvent>,
|
||||
pub cloud_model_events: async_channel::Receiver<CloudModelEvent>,
|
||||
}
|
||||
|
||||
pub fn initialize_app(app: &mut App) {
|
||||
app.add_singleton_model(|ctx| AppExecutionMode::new(ExecutionMode::App, false, ctx));
|
||||
|
||||
app.add_singleton_model(|_| NetworkStatus::new());
|
||||
app.add_singleton_model(|_| SettingsManager::default());
|
||||
app.add_singleton_model(TeamTesterStatus::mock);
|
||||
app.update(crate::settings::init_and_register_user_preferences);
|
||||
// This ServerApiProvider is used for the PrivacySettings model, but not the UpdateManager
|
||||
// under test.
|
||||
app.add_singleton_model(|_| ServerApiProvider::new_for_test());
|
||||
app.add_singleton_model(|_| AuthStateProvider::new_for_test());
|
||||
app.add_singleton_model(AppTelemetryContextProvider::new_context_provider);
|
||||
app.add_singleton_model(AuthManager::new_for_test);
|
||||
WarpDrivePrivacySettings::register(app);
|
||||
app.update(PrivacySettings::register_singleton);
|
||||
app.add_singleton_model(CloudModel::mock);
|
||||
app.add_singleton_model(UserWorkspaces::default_mock);
|
||||
app.add_singleton_model(TeamUpdateManager::mock);
|
||||
app.add_singleton_model(|_| ObjectActions::new(Vec::new()));
|
||||
}
|
||||
|
||||
pub fn create_update_manager_struct(
|
||||
app: &mut App,
|
||||
server_api: Arc<dyn ObjectClient>,
|
||||
) -> UpdateManagerStruct {
|
||||
let (sender, receiver) = sync_channel(CHANNEL_SIZE);
|
||||
|
||||
// the sync queue can't be mocked; needs to use the same server_api as the update_manager
|
||||
app.add_singleton_model(|ctx| SyncQueue::new(Default::default(), server_api.clone(), ctx));
|
||||
let update_manager =
|
||||
app.add_singleton_model(|ctx| UpdateManager::new(Some(sender.clone()), server_api, ctx));
|
||||
|
||||
app.add_singleton_model(|_| UserProfiles::new(Vec::new()));
|
||||
|
||||
// set up the sync queue in a dequeueing state
|
||||
SyncQueue::handle(app).update(app, |sync_queue, ctx| {
|
||||
sync_queue.start_dequeueing(ctx);
|
||||
});
|
||||
|
||||
let cloud_model_events = app.update(|ctx| {
|
||||
let (tx, rx) = async_channel::unbounded();
|
||||
ctx.subscribe_to_model(&CloudModel::handle(ctx), move |_, event, _| {
|
||||
let _ = tx.try_send(event.clone());
|
||||
});
|
||||
rx
|
||||
});
|
||||
|
||||
// The start of polling is normally triggered by authentication completion, but
|
||||
// we need to do it manually for tests. We do this AFTER UpdateManager is created
|
||||
// so the polling uses the correct mock.
|
||||
TeamTesterStatus::handle(app).update(app, |team_tester, ctx| {
|
||||
team_tester.initiate_data_pollers(false, ctx);
|
||||
});
|
||||
|
||||
UpdateManagerStruct {
|
||||
update_manager,
|
||||
receiver,
|
||||
cloud_model_events,
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates a baseline [`MockObjectClient`] with common mocks like:
|
||||
/// * The logged-in user
|
||||
/// * Background polling for updated objects
|
||||
pub fn mock_server_api() -> MockObjectClient {
|
||||
let mut mock_object_client = MockObjectClient::new();
|
||||
// Mock *failures* for background fetches. This prevents `UpdateManager` clearing out any
|
||||
// objects that tests manually add to `CloudModel`.
|
||||
mock_object_client
|
||||
.expect_fetch_changed_objects()
|
||||
.returning(|_, _| Err(anyhow::anyhow!("Ignoring background refresh in tests")));
|
||||
// Mock environment timestamps fetch - return empty by default.
|
||||
mock_object_client
|
||||
.expect_fetch_environment_last_task_run_timestamps()
|
||||
.returning(|| Ok(HashMap::new()));
|
||||
mock_object_client
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,13 @@
|
||||
use chrono::{DateTime, FixedOffset, Local};
|
||||
|
||||
pub trait DateTimeExt {
|
||||
fn now() -> DateTime<FixedOffset>;
|
||||
}
|
||||
|
||||
impl DateTimeExt for DateTime<FixedOffset> {
|
||||
/// Gets current date and time and timezone in DateTime<FixedOffset>.
|
||||
fn now() -> DateTime<FixedOffset> {
|
||||
let local_time = Local::now();
|
||||
local_time.with_timezone(local_time.offset())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
//! Logic to convert from / to [`ServerExperiment`].
|
||||
|
||||
use std::fmt::{Display, Formatter};
|
||||
|
||||
use anyhow::{Ok, Result};
|
||||
use warp_graphql::experiment::Experiment;
|
||||
|
||||
use super::ServerExperiment;
|
||||
|
||||
impl Display for ServerExperiment {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
let str = match self {
|
||||
Self::SessionSharingControl => "SESSION_SHARING_CONTROL",
|
||||
Self::SessionSharingExperiment => "SESSION_SHARING_EXPERIMENT",
|
||||
Self::DisableAgentModeExperiment => "DISABLE_AGENT_MODE_EXPERIMENT",
|
||||
Self::EnvVarsEarlyAccessExperiment => "ENV_VARS_EARLY_ACCESS_EXPERIMENT",
|
||||
Self::AgentModeAnalyticsExperiment => "AGENT_MODE_ANALYTICS_EXPERIMENT",
|
||||
Self::WindowsLaunchExperiment => "WINDOWS_LAUNCH_EXPERIMENT",
|
||||
Self::TmuxSshWarpificationControl => "TMUX_SSH_WARPIFICATION_CONTROL",
|
||||
Self::TmuxSshWarpificationExperiment => "TMUX_SSH_WARPIFICATION_EXPERIMENT",
|
||||
Self::CodebaseContextControl => "CODEBASE_CONTEXT_CONTROL",
|
||||
Self::CodebaseContextExperiment => "CODEBASE_CONTEXT_EXPERIMENT",
|
||||
Self::SuggestedCodeDiffsControl => "SUGGESTED_CODE_DIFFS_CONTROL",
|
||||
Self::SuggestedCodeDiffsExperiment => "SUGGESTED_CODE_DIFFS_EXPERIMENT",
|
||||
Self::BuildPlanAutoReloadControl => "BUILD_PLAN_AUTO_RELOAD_CONTROL",
|
||||
Self::BuildPlanAutoReloadBannerToggle => "BUILD_PLAN_AUTO_RELOAD_BANNER_TOGGLE",
|
||||
Self::BuildPlanAutoReloadPostPurchaseModal => {
|
||||
"BUILD_PLAN_AUTO_RELOAD_POST_PURCHASE_MODAL"
|
||||
}
|
||||
Self::PromptSuggestionsViaMaaControl => "PROMPT_SUGGESTIONS_VIA_MAA_CONTROL",
|
||||
Self::PromptSuggestionsViaMaaExperiment => "PROMPT_SUGGESTIONS_VIA_MAA_EXPERIMENT",
|
||||
Self::PromptSuggestionsViaMaaOutOfBandExperiment => {
|
||||
"PROMPT_SUGGESTIONS_VIA_MAA_OOB_EXPERIMENT"
|
||||
}
|
||||
Self::FreeUserNoAiControl => "FREE_USER_NO_AI_CONTROL",
|
||||
Self::FreeUserNoAiExperiment => "FREE_USER_NO_AI_EXPERIMENT",
|
||||
Self::OzMultiHarnessControl => "OZ_MULTI_HARNESS_CONTROL",
|
||||
Self::OzMultiHarnessExperiment => "OZ_MULTI_HARNESS_EXPERIMENT",
|
||||
#[cfg(test)]
|
||||
Self::TestExperiment => "TEST_EXPERIMENT",
|
||||
};
|
||||
write!(f, "{str}")
|
||||
}
|
||||
}
|
||||
|
||||
impl ServerExperiment {
|
||||
pub fn from_string(s: String) -> Result<Self> {
|
||||
match s.as_str() {
|
||||
"SESSION_SHARING_CONTROL" => Ok(Self::SessionSharingControl),
|
||||
"SESSION_SHARING_EXPERIMENT" => Ok(Self::SessionSharingExperiment),
|
||||
"DISABLE_AGENT_MODE_EXPERIMENT" => Ok(Self::DisableAgentModeExperiment),
|
||||
"ENV_VARS_EARLY_ACCESS_EXPERIMENT" => Ok(Self::EnvVarsEarlyAccessExperiment),
|
||||
"AGENT_MODE_ANALYTICS_EXPERIMENT" => Ok(Self::AgentModeAnalyticsExperiment),
|
||||
"WINDOWS_LAUNCH_EXPERIMENT" => Ok(Self::WindowsLaunchExperiment),
|
||||
"TMUX_SSH_WARPIFICATION_CONTROL" => Ok(Self::TmuxSshWarpificationControl),
|
||||
"TMUX_SSH_WARPIFICATION_EXPERIMENT" => Ok(Self::TmuxSshWarpificationExperiment),
|
||||
"CODEBASE_CONTEXT_EXPERIMENT" => Ok(Self::CodebaseContextExperiment),
|
||||
"CODEBASE_CONTEXT_CONTROL" => Ok(Self::CodebaseContextControl),
|
||||
"SUGGESTED_CODE_DIFFS_CONTROL" => Ok(Self::SuggestedCodeDiffsControl),
|
||||
"SUGGESTED_CODE_DIFFS_EXPERIMENT" => Ok(Self::SuggestedCodeDiffsExperiment),
|
||||
"BUILD_PLAN_AUTO_RELOAD_CONTROL" => Ok(Self::BuildPlanAutoReloadControl),
|
||||
"BUILD_PLAN_AUTO_RELOAD_BANNER_TOGGLE" => Ok(Self::BuildPlanAutoReloadBannerToggle),
|
||||
"BUILD_PLAN_AUTO_RELOAD_POST_PURCHASE_MODAL" => {
|
||||
Ok(Self::BuildPlanAutoReloadPostPurchaseModal)
|
||||
}
|
||||
"PROMPT_SUGGESTIONS_VIA_MAA_CONTROL" => Ok(Self::PromptSuggestionsViaMaaControl),
|
||||
"PROMPT_SUGGESTIONS_VIA_MAA_EXPERIMENT" => Ok(Self::PromptSuggestionsViaMaaExperiment),
|
||||
"FREE_USER_NO_AI_CONTROL" => Ok(Self::FreeUserNoAiControl),
|
||||
"FREE_USER_NO_AI_EXPERIMENT" => Ok(Self::FreeUserNoAiExperiment),
|
||||
"OZ_MULTI_HARNESS_CONTROL" => Ok(Self::OzMultiHarnessControl),
|
||||
"OZ_MULTI_HARNESS_EXPERIMENT" => Ok(Self::OzMultiHarnessExperiment),
|
||||
s => Err(anyhow::anyhow!(
|
||||
"String doesn't match any server experiment variant {s}"
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<Experiment> for ServerExperiment {
|
||||
type Error = anyhow::Error;
|
||||
|
||||
fn try_from(value: Experiment) -> Result<Self, Self::Error> {
|
||||
match value {
|
||||
Experiment::SessionSharingExperiment => Ok(Self::SessionSharingExperiment),
|
||||
Experiment::SessionSharingControl => Ok(Self::SessionSharingControl),
|
||||
Experiment::BuildPlanAutoReloadControl => Ok(Self::BuildPlanAutoReloadControl),
|
||||
Experiment::BuildPlanAutoReloadBannerToggle => {
|
||||
Ok(Self::BuildPlanAutoReloadBannerToggle)
|
||||
}
|
||||
Experiment::BuildPlanAutoReloadPostPurchaseModal => {
|
||||
Ok(Self::BuildPlanAutoReloadPostPurchaseModal)
|
||||
}
|
||||
Experiment::DisableAgentModeExperiment => Ok(Self::DisableAgentModeExperiment),
|
||||
Experiment::EnvVarsEarlyAccessExperiment => Ok(Self::EnvVarsEarlyAccessExperiment),
|
||||
Experiment::AgentModeAnalyticsExperiment => Ok(Self::AgentModeAnalyticsExperiment),
|
||||
Experiment::TmuxSshWarpificationControl => Ok(Self::TmuxSshWarpificationControl),
|
||||
Experiment::TmuxSshWarpificationExperiment => Ok(Self::TmuxSshWarpificationExperiment),
|
||||
Experiment::WindowsLaunchExperiment => Ok(Self::WindowsLaunchExperiment),
|
||||
Experiment::CodebaseContextControl => Ok(Self::CodebaseContextControl),
|
||||
Experiment::CodebaseContextExperiment => Ok(Self::CodebaseContextExperiment),
|
||||
Experiment::SuggestedCodeDiffsControl => Ok(Self::SuggestedCodeDiffsControl),
|
||||
Experiment::SuggestedCodeDiffsExperiment => Ok(Self::SuggestedCodeDiffsExperiment),
|
||||
Experiment::PromptSuggestionsViaMaaControl => Ok(Self::PromptSuggestionsViaMaaControl),
|
||||
Experiment::PromptSuggestionsViaMaaOob => {
|
||||
Ok(Self::PromptSuggestionsViaMaaOutOfBandExperiment)
|
||||
}
|
||||
Experiment::FreeUserNoAiControl => Ok(Self::FreeUserNoAiControl),
|
||||
Experiment::FreeUserNoAiExperiment => Ok(Self::FreeUserNoAiExperiment),
|
||||
Experiment::OzMultiHarnessControl => Ok(Self::OzMultiHarnessControl),
|
||||
Experiment::OzMultiHarnessExperiment => Ok(Self::OzMultiHarnessExperiment),
|
||||
// Experiments that we no longer support on the client.
|
||||
e => Err(anyhow::anyhow!(
|
||||
"Server-side enabled experiment '{e:?}' is no longer supported by the client."
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! convert_to_server_experiment {
|
||||
($gql_type:expr) => {{
|
||||
let mut acc = Vec::new();
|
||||
for a in $gql_type {
|
||||
// Note for server experiments we don't currently track on the client.
|
||||
// This could be because the client is out of date and we should still
|
||||
// apply the experiments the client does track.
|
||||
if let Ok(b) = ServerExperiment::try_from(a) {
|
||||
acc.push(b);
|
||||
}
|
||||
}
|
||||
Some(acc)
|
||||
}};
|
||||
}
|
||||
@@ -0,0 +1,173 @@
|
||||
//! This module is responsible for applying
|
||||
//! server-side experiment state to the client.
|
||||
//!
|
||||
//! After adding support for your experiment on the server,
|
||||
//! you need to define what effect the experiment will
|
||||
//! have on the client (i.e. see [`ServerExperiment::on_added_to`]).
|
||||
//!
|
||||
//! Then, you can use the global [`ServerExperiments`] model
|
||||
//! to update and query the latest experiment state.
|
||||
//!
|
||||
//! See [here](https://www.notion.so/warpdev/Server-side-experiments-dynamic-feature-enablement-c0fb9aed695d4178a19b8830e3269094)
|
||||
//! for a full guide on the server-side experiment framework.
|
||||
|
||||
use crate::features::FeatureFlag;
|
||||
use crate::workspaces::user_workspaces::UserWorkspaces;
|
||||
use crate::workspaces::workspace::CustomerType;
|
||||
use warpui::AppContext;
|
||||
#[cfg(not(test))]
|
||||
use warpui::SingletonEntity as _;
|
||||
#[cfg(test)]
|
||||
use warpui::SingletonEntity;
|
||||
|
||||
mod convert;
|
||||
mod model;
|
||||
|
||||
pub use model::{Event as ServerExperimentsEvent, ServerExperiments};
|
||||
|
||||
/// The known server-side experiments.
|
||||
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
|
||||
pub enum ServerExperiment {
|
||||
SessionSharingExperiment,
|
||||
SessionSharingControl,
|
||||
DisableAgentModeExperiment,
|
||||
EnvVarsEarlyAccessExperiment,
|
||||
AgentModeAnalyticsExperiment,
|
||||
WindowsLaunchExperiment,
|
||||
TmuxSshWarpificationControl,
|
||||
TmuxSshWarpificationExperiment,
|
||||
CodebaseContextExperiment,
|
||||
CodebaseContextControl,
|
||||
SuggestedCodeDiffsControl,
|
||||
SuggestedCodeDiffsExperiment,
|
||||
BuildPlanAutoReloadControl,
|
||||
BuildPlanAutoReloadBannerToggle,
|
||||
BuildPlanAutoReloadPostPurchaseModal,
|
||||
PromptSuggestionsViaMaaControl,
|
||||
PromptSuggestionsViaMaaExperiment,
|
||||
PromptSuggestionsViaMaaOutOfBandExperiment,
|
||||
FreeUserNoAiControl,
|
||||
FreeUserNoAiExperiment,
|
||||
OzMultiHarnessControl,
|
||||
OzMultiHarnessExperiment,
|
||||
/// A test-only experiment.
|
||||
/// Does not correspond to a real server-side experiment.
|
||||
#[cfg(test)]
|
||||
TestExperiment,
|
||||
}
|
||||
|
||||
impl ServerExperiment {
|
||||
/// When the client is added to an experiment.
|
||||
///
|
||||
// TODO: currently, it isn't possible to interact with other
|
||||
// application singletons in this function because [`ServerExperiments`]
|
||||
// is initialized before most other singletons during app init. We should either:
|
||||
// a) remove the `ctx` from this function to prevent reading / updating
|
||||
// other application entities here, which forces those other entities
|
||||
// to subscribe to [`ServerExperiments`] updates instead, or
|
||||
// b) continue initializing [`ServerExperiments`] as one of the first
|
||||
// singletons but apply the cached state after initializing all other singletons.
|
||||
// That way, this method would not be called until after all other singletons
|
||||
// have been initialized and can thus be referenced.
|
||||
fn on_added_to(&self, _ctx: &mut AppContext) {
|
||||
match self {
|
||||
Self::SessionSharingExperiment => {
|
||||
FeatureFlag::CreatingSharedSessions.set_enabled(true);
|
||||
}
|
||||
Self::SessionSharingControl => {
|
||||
FeatureFlag::CreatingSharedSessions.set_enabled(false);
|
||||
}
|
||||
Self::DisableAgentModeExperiment => {
|
||||
FeatureFlag::AgentMode.set_enabled(false);
|
||||
}
|
||||
Self::EnvVarsEarlyAccessExperiment => {
|
||||
// EnvVars is now always enabled; no-op.
|
||||
}
|
||||
Self::AgentModeAnalyticsExperiment => {
|
||||
FeatureFlag::AgentModeAnalytics.set_enabled(true);
|
||||
FeatureFlag::AIRules.set_enabled(true);
|
||||
FeatureFlag::SuggestedRules.set_enabled(true);
|
||||
}
|
||||
Self::WindowsLaunchExperiment => {
|
||||
// TODO(alokedesai): Clean this up now that we no longer gate access to the Windows
|
||||
// build on an allowlist.
|
||||
}
|
||||
Self::TmuxSshWarpificationControl => FeatureFlag::SSHTmuxWrapper.set_enabled(false),
|
||||
Self::TmuxSshWarpificationExperiment => {
|
||||
// Only enable the TMUX-based experience if not on windows. ConPTY doesn't support
|
||||
// DCS, which we need in order to use tmux control mode.
|
||||
if cfg!(not(windows)) {
|
||||
FeatureFlag::SSHTmuxWrapper.set_enabled(true)
|
||||
}
|
||||
}
|
||||
Self::CodebaseContextExperiment => {
|
||||
FeatureFlag::FullSourceCodeEmbedding.set_enabled(true);
|
||||
FeatureFlag::CodebaseIndexPersistence.set_enabled(true);
|
||||
FeatureFlag::CodebaseIndexSpeedbump.set_enabled(true);
|
||||
FeatureFlag::CrossRepoContext.set_enabled(true);
|
||||
}
|
||||
Self::CodebaseContextControl => {
|
||||
FeatureFlag::FullSourceCodeEmbedding.set_enabled(false);
|
||||
FeatureFlag::CodebaseIndexPersistence.set_enabled(false);
|
||||
FeatureFlag::CodebaseIndexSpeedbump.set_enabled(false);
|
||||
FeatureFlag::CrossRepoContext.set_enabled(false);
|
||||
}
|
||||
Self::SuggestedCodeDiffsExperiment => {}
|
||||
Self::SuggestedCodeDiffsControl => {}
|
||||
Self::BuildPlanAutoReloadControl => {
|
||||
// Control group - disable both experiment flags
|
||||
FeatureFlag::BuildPlanAutoReloadBannerToggle.set_enabled(false);
|
||||
FeatureFlag::BuildPlanAutoReloadPostPurchaseModal.set_enabled(false);
|
||||
}
|
||||
Self::BuildPlanAutoReloadBannerToggle => {
|
||||
// Experiment variant 1 - enable banner toggle modal
|
||||
FeatureFlag::BuildPlanAutoReloadBannerToggle.set_enabled(true);
|
||||
FeatureFlag::BuildPlanAutoReloadPostPurchaseModal.set_enabled(false);
|
||||
}
|
||||
Self::BuildPlanAutoReloadPostPurchaseModal => {
|
||||
// Experiment variant 2 - enable post-purchase modal
|
||||
FeatureFlag::BuildPlanAutoReloadBannerToggle.set_enabled(false);
|
||||
FeatureFlag::BuildPlanAutoReloadPostPurchaseModal.set_enabled(true);
|
||||
}
|
||||
Self::PromptSuggestionsViaMaaControl => {
|
||||
FeatureFlag::PromptSuggestionsViaMAA.set_enabled(false);
|
||||
}
|
||||
Self::PromptSuggestionsViaMaaOutOfBandExperiment => {
|
||||
FeatureFlag::PromptSuggestionsViaMAA.set_enabled(true);
|
||||
}
|
||||
// The normal experiment arm is no longer used.
|
||||
Self::PromptSuggestionsViaMaaExperiment => {}
|
||||
Self::FreeUserNoAiControl => {
|
||||
FeatureFlag::FreeUserNoAi.set_enabled(false);
|
||||
}
|
||||
Self::FreeUserNoAiExperiment => {
|
||||
FeatureFlag::FreeUserNoAi.set_enabled(true);
|
||||
}
|
||||
Self::OzMultiHarnessControl => {
|
||||
FeatureFlag::AgentHarness.set_enabled(false);
|
||||
}
|
||||
Self::OzMultiHarnessExperiment => {
|
||||
FeatureFlag::AgentHarness.set_enabled(true);
|
||||
}
|
||||
#[cfg(test)]
|
||||
Self::TestExperiment => {
|
||||
model::TestModel::handle(_ctx).update(_ctx, |model, _| {
|
||||
model.0 += 1;
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns `true` when the user is in the `FreeUserNoAiExperiment` arm **and** is on the
|
||||
/// free tier. This is the single source of truth for gating any client-side behaviour
|
||||
/// that should be locked/disabled for users without AI credits.
|
||||
pub fn is_free_user_no_ai_experiment_active(ctx: &AppContext) -> bool {
|
||||
let in_experiment = FeatureFlag::FreeUserNoAi.is_enabled();
|
||||
let is_free_tier = UserWorkspaces::handle(ctx)
|
||||
.as_ref(ctx)
|
||||
.current_team()
|
||||
.map(|team| team.billing_metadata.customer_type == CustomerType::Free)
|
||||
.unwrap_or(true); // no team = solo free user
|
||||
in_experiment && is_free_tier
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
//! The model for maintaining global experiment state.
|
||||
|
||||
use std::collections::HashSet;
|
||||
|
||||
use super::ServerExperiment;
|
||||
use crate::{persistence::ModelEvent, report_if_error, GlobalResourceHandlesProvider};
|
||||
use anyhow::Context;
|
||||
use warpui::{Entity, ModelContext, SingletonEntity};
|
||||
|
||||
#[cfg(test)]
|
||||
pub use tests::TestModel;
|
||||
|
||||
/// A global model for maintaining server-side experiment state.
|
||||
pub struct ServerExperiments {
|
||||
/// The latest-known set of server-side enabled experiments.
|
||||
latest: HashSet<ServerExperiment>,
|
||||
}
|
||||
|
||||
impl ServerExperiments {
|
||||
/// Creates a new [`ServerExperiments`] model and seeds it with
|
||||
/// the provided `cached` experiment state.
|
||||
pub fn new_from_cache(cached: Vec<ServerExperiment>, ctx: &mut ModelContext<Self>) -> Self {
|
||||
let mut model = Self {
|
||||
latest: HashSet::new(),
|
||||
};
|
||||
model.apply_latest_state(cached, ctx);
|
||||
model
|
||||
}
|
||||
|
||||
/// Updates the model with the latest server-side state.
|
||||
///
|
||||
/// Assumes the set of proivded [`ServerExperiment`]s are unambiguous;
|
||||
/// that is, there are not two arms enabled for the same experiment group.
|
||||
pub fn apply_latest_state(
|
||||
&mut self,
|
||||
incoming: Vec<ServerExperiment>,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
// Dedup the set of experiments.
|
||||
let incoming = HashSet::from_iter(incoming);
|
||||
|
||||
// For every experiment that the client isn't already part of,
|
||||
// perform the necessary logic to add them.
|
||||
for experiment in incoming.difference(&self.latest) {
|
||||
experiment.on_added_to(ctx);
|
||||
}
|
||||
|
||||
self.cache_latest_state(incoming, ctx);
|
||||
ctx.emit(Event::ExperimentsUpdated);
|
||||
}
|
||||
|
||||
/// Returns true iff the `experiment` is enabled.
|
||||
pub fn is_experiment_enabled(&self, experiment: &ServerExperiment) -> bool {
|
||||
self.latest.contains(experiment)
|
||||
}
|
||||
|
||||
/// Saves the latest experiment state in-memory and to the local cache.
|
||||
fn cache_latest_state(
|
||||
&mut self,
|
||||
latest: HashSet<ServerExperiment>,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
self.latest = latest;
|
||||
|
||||
if let Some(model_event_sender) = GlobalResourceHandlesProvider::as_ref(ctx)
|
||||
.get()
|
||||
.model_event_sender
|
||||
.as_ref()
|
||||
{
|
||||
let event = ModelEvent::SaveExperiments {
|
||||
experiments: self.latest.iter().copied().collect(),
|
||||
};
|
||||
report_if_error!(model_event_sender
|
||||
.send(event)
|
||||
.context("Unable to save experiments to sqlite"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub enum Event {
|
||||
ExperimentsUpdated,
|
||||
}
|
||||
|
||||
impl Entity for ServerExperiments {
|
||||
type Event = Event;
|
||||
}
|
||||
|
||||
impl SingletonEntity for ServerExperiments {}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "model_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,65 @@
|
||||
use super::{ServerExperiment, ServerExperiments};
|
||||
use crate::{GlobalResourceHandles, GlobalResourceHandlesProvider};
|
||||
use warpui::{App, Entity, SingletonEntity};
|
||||
|
||||
/// A model for testing purposes only.
|
||||
///
|
||||
/// We use it to demonstrate how client-side
|
||||
/// models can be mutated to reflect server
|
||||
/// experiment state changes.
|
||||
pub struct TestModel(pub usize);
|
||||
|
||||
impl Entity for TestModel {
|
||||
type Event = ();
|
||||
}
|
||||
impl SingletonEntity for TestModel {}
|
||||
|
||||
fn initialize_app(app: &mut App) {
|
||||
app.update(crate::settings::init_and_register_user_preferences);
|
||||
|
||||
let global_resources = GlobalResourceHandles::mock(app);
|
||||
app.add_singleton_model(|_| GlobalResourceHandlesProvider::new(global_resources));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_new_from_cached() {
|
||||
App::test((), |mut app| async move {
|
||||
initialize_app(&mut app);
|
||||
|
||||
let model = app.add_singleton_model(|_| TestModel(0));
|
||||
let cache = vec![ServerExperiment::TestExperiment];
|
||||
app.add_singleton_model(|ctx| ServerExperiments::new_from_cache(cache, ctx));
|
||||
|
||||
// The experiment should have been enabled.
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(model.0, 1);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_apply_latest_state() {
|
||||
App::test((), |mut app| async move {
|
||||
initialize_app(&mut app);
|
||||
|
||||
let model = app.add_singleton_model(|_| TestModel(0));
|
||||
let experiments =
|
||||
app.add_singleton_model(|ctx| ServerExperiments::new_from_cache(vec![], ctx));
|
||||
|
||||
// Enable the experiment.
|
||||
experiments.update(&mut app, |experiments, ctx| {
|
||||
experiments.apply_latest_state(vec![ServerExperiment::TestExperiment], ctx);
|
||||
});
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(model.0, 1);
|
||||
});
|
||||
|
||||
// Redundant experiment state should be a no-op.
|
||||
experiments.update(&mut app, |experiments, ctx| {
|
||||
experiments.apply_latest_state(vec![ServerExperiment::TestExperiment], ctx);
|
||||
});
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(model.0, 1);
|
||||
});
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
pub mod schema;
|
||||
|
||||
use warp_graphql::client::RequestOptions;
|
||||
pub use warp_graphql::client::{get_request_context, get_user_facing_error_message, GraphQLError};
|
||||
|
||||
/// Returns the default [`RequestOptions`] that should be used for a GraphQL request.
|
||||
pub fn default_request_options() -> RequestOptions {
|
||||
RequestOptions {
|
||||
#[cfg(feature = "agent_mode_evals")]
|
||||
path_prefix: Some("/agent-mode-evals".to_string()),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,322 @@
|
||||
pub mod util;
|
||||
|
||||
use crate::{
|
||||
ai::cloud_environments::CloudAmbientAgentEnvironmentModel,
|
||||
ai::{
|
||||
ambient_agents::scheduled::CloudScheduledAmbientAgentModel,
|
||||
execution_profiles::CloudAIExecutionProfileModel,
|
||||
facts::CloudAIFactModel,
|
||||
mcp::{templatable::CloudTemplatableMCPServerModel, CloudMCPServerModel},
|
||||
},
|
||||
cloud_object::{
|
||||
model::generic_string_model::GenericStringObjectId, GenericServerObject,
|
||||
RevisionAndLastEditor, ServerFolder, ServerObject, UpdateCloudObjectResult,
|
||||
},
|
||||
env_vars::CloudEnvVarCollectionModel,
|
||||
server::{graphql::get_user_facing_error_message, ids::ServerId},
|
||||
settings::cloud_preferences::CloudPreferenceModel,
|
||||
workflows::workflow_enum::CloudWorkflowEnumModel,
|
||||
};
|
||||
|
||||
use anyhow::{bail, Result};
|
||||
use warp_graphql::{
|
||||
generic_string_object::GenericStringObjectFormat,
|
||||
mutations::update_generic_string_object::{
|
||||
GenericStringObjectUpdate, UpdateGenericStringObjectResult,
|
||||
},
|
||||
object::ObjectUpdateSuccess,
|
||||
};
|
||||
|
||||
impl TryFrom<UpdateGenericStringObjectResult> for UpdateCloudObjectResult<Box<dyn ServerObject>> {
|
||||
type Error = anyhow::Error;
|
||||
|
||||
fn try_from(value: UpdateGenericStringObjectResult) -> std::result::Result<Self, Self::Error> {
|
||||
match value {
|
||||
UpdateGenericStringObjectResult::UpdateGenericStringObjectOutput(output) => {
|
||||
match output.update {
|
||||
GenericStringObjectUpdate::ObjectUpdateSuccess(success) => {
|
||||
Ok(UpdateCloudObjectResult::Success {
|
||||
revision_and_editor: RevisionAndLastEditor {
|
||||
revision: success.revision_ts.into(),
|
||||
last_editor_uid: Some(success.last_editor_uid.into_inner()),
|
||||
},
|
||||
})
|
||||
}
|
||||
GenericStringObjectUpdate::GenericStringObjectUpdateRejected(rejected) => {
|
||||
let boxed: Box<dyn ServerObject> = match rejected
|
||||
.conflicting_generic_string_object
|
||||
.format
|
||||
{
|
||||
GenericStringObjectFormat::JsonEnvVarCollection => {
|
||||
let gso = GenericServerObject::<
|
||||
GenericStringObjectId,
|
||||
CloudEnvVarCollectionModel,
|
||||
>::try_from_graphql_fields(
|
||||
ServerId::from_string_lossy(
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.uid
|
||||
.inner(),
|
||||
),
|
||||
Some(
|
||||
rejected.conflicting_generic_string_object.serialized_model,
|
||||
),
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.try_into()?,
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.permissions
|
||||
.try_into()?,
|
||||
)?;
|
||||
let boxed: Box<dyn ServerObject> = Box::new(gso);
|
||||
boxed
|
||||
}
|
||||
GenericStringObjectFormat::JsonPreference => {
|
||||
let gso = GenericServerObject::<
|
||||
GenericStringObjectId,
|
||||
CloudPreferenceModel,
|
||||
>::try_from_graphql_fields(
|
||||
ServerId::from_string_lossy(
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.uid
|
||||
.inner(),
|
||||
),
|
||||
Some(
|
||||
rejected.conflicting_generic_string_object.serialized_model,
|
||||
),
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.try_into()?,
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.permissions
|
||||
.try_into()?,
|
||||
)?;
|
||||
let boxed: Box<dyn ServerObject> = Box::new(gso);
|
||||
boxed
|
||||
}
|
||||
GenericStringObjectFormat::JsonWorkflowEnum => {
|
||||
let gso = GenericServerObject::<
|
||||
GenericStringObjectId,
|
||||
CloudWorkflowEnumModel,
|
||||
>::try_from_graphql_fields(
|
||||
ServerId::from_string_lossy(
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.uid
|
||||
.inner(),
|
||||
),
|
||||
Some(
|
||||
rejected.conflicting_generic_string_object.serialized_model,
|
||||
),
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.try_into()?,
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.permissions
|
||||
.try_into()?,
|
||||
)?;
|
||||
let boxed: Box<dyn ServerObject> = Box::new(gso);
|
||||
boxed
|
||||
}
|
||||
GenericStringObjectFormat::JsonAIFact => {
|
||||
let gso = GenericServerObject::<
|
||||
GenericStringObjectId,
|
||||
CloudAIFactModel,
|
||||
>::try_from_graphql_fields(
|
||||
ServerId::from_string_lossy(
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.uid
|
||||
.inner(),
|
||||
),
|
||||
Some(
|
||||
rejected.conflicting_generic_string_object.serialized_model,
|
||||
),
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.try_into()?,
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.permissions
|
||||
.try_into()?,
|
||||
)?;
|
||||
let boxed: Box<dyn ServerObject> = Box::new(gso);
|
||||
boxed
|
||||
}
|
||||
GenericStringObjectFormat::JsonAIExecutionProfile => {
|
||||
let gso = GenericServerObject::<
|
||||
GenericStringObjectId,
|
||||
CloudAIExecutionProfileModel,
|
||||
>::try_from_graphql_fields(
|
||||
ServerId::from_string_lossy(
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.uid
|
||||
.inner(),
|
||||
),
|
||||
Some(
|
||||
rejected.conflicting_generic_string_object.serialized_model,
|
||||
),
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.try_into()?,
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.permissions
|
||||
.try_into()?,
|
||||
)?;
|
||||
let boxed: Box<dyn ServerObject> = Box::new(gso);
|
||||
boxed
|
||||
}
|
||||
GenericStringObjectFormat::JsonMCPServer => {
|
||||
let gso = GenericServerObject::<
|
||||
GenericStringObjectId,
|
||||
CloudMCPServerModel,
|
||||
>::try_from_graphql_fields(
|
||||
ServerId::from_string_lossy(
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.uid
|
||||
.inner(),
|
||||
),
|
||||
Some(
|
||||
rejected.conflicting_generic_string_object.serialized_model,
|
||||
),
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.try_into()?,
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.permissions
|
||||
.try_into()?,
|
||||
)?;
|
||||
let boxed: Box<dyn ServerObject> = Box::new(gso);
|
||||
boxed
|
||||
}
|
||||
GenericStringObjectFormat::JsonTemplatableMCPServer => {
|
||||
let gso = GenericServerObject::<
|
||||
GenericStringObjectId,
|
||||
CloudTemplatableMCPServerModel,
|
||||
>::try_from_graphql_fields(
|
||||
ServerId::from_string_lossy(
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.uid
|
||||
.inner(),
|
||||
),
|
||||
Some(
|
||||
rejected.conflicting_generic_string_object.serialized_model,
|
||||
),
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.try_into()?,
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.permissions
|
||||
.try_into()?,
|
||||
)?;
|
||||
let boxed: Box<dyn ServerObject> = Box::new(gso);
|
||||
boxed
|
||||
}
|
||||
GenericStringObjectFormat::JsonCloudEnvironment => {
|
||||
let gso = GenericServerObject::<
|
||||
GenericStringObjectId,
|
||||
CloudAmbientAgentEnvironmentModel,
|
||||
>::try_from_graphql_fields(
|
||||
ServerId::from_string_lossy(
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.uid
|
||||
.inner(),
|
||||
),
|
||||
Some(
|
||||
rejected.conflicting_generic_string_object.serialized_model,
|
||||
),
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.try_into()?,
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.permissions
|
||||
.try_into()?,
|
||||
)?;
|
||||
let boxed: Box<dyn ServerObject> = Box::new(gso);
|
||||
boxed
|
||||
}
|
||||
GenericStringObjectFormat::JsonScheduledAmbientAgent => {
|
||||
let gso = GenericServerObject::<
|
||||
GenericStringObjectId,
|
||||
CloudScheduledAmbientAgentModel,
|
||||
>::try_from_graphql_fields(
|
||||
ServerId::from_string_lossy(
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.uid
|
||||
.inner(),
|
||||
),
|
||||
Some(
|
||||
rejected.conflicting_generic_string_object.serialized_model,
|
||||
),
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.metadata
|
||||
.try_into()?,
|
||||
rejected
|
||||
.conflicting_generic_string_object
|
||||
.permissions
|
||||
.try_into()?,
|
||||
)?;
|
||||
let boxed: Box<dyn ServerObject> = Box::new(gso);
|
||||
boxed
|
||||
}
|
||||
};
|
||||
Ok(UpdateCloudObjectResult::Rejected { object: boxed })
|
||||
}
|
||||
GenericStringObjectUpdate::Unknown => {
|
||||
bail!("update generic string object response has unknown variant")
|
||||
}
|
||||
}
|
||||
}
|
||||
UpdateGenericStringObjectResult::UserFacingError(e) => {
|
||||
bail!(get_user_facing_error_message(e))
|
||||
}
|
||||
UpdateGenericStringObjectResult::Unknown => {
|
||||
bail!("update generic string object response has unknown variant")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<ObjectUpdateSuccess> for UpdateCloudObjectResult<ServerFolder> {
|
||||
type Error = anyhow::Error;
|
||||
|
||||
fn try_from(value: ObjectUpdateSuccess) -> Result<Self, Self::Error> {
|
||||
Ok(UpdateCloudObjectResult::Success {
|
||||
revision_and_editor: RevisionAndLastEditor {
|
||||
revision: value.revision_ts.into(),
|
||||
last_editor_uid: Some(value.last_editor_uid.into_inner()),
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
use crate::anyhow;
|
||||
|
||||
use crate::cloud_object::model::actions::ObjectActionHistory;
|
||||
use crate::cloud_object::model::actions::ObjectActionType;
|
||||
use crate::cloud_object::model::actions::{ObjectAction, ObjectActionSubtype};
|
||||
use crate::cloud_object::{GenericStringObjectUniqueKey, UniquePer};
|
||||
use crate::server::ids::{HashedSqliteId, ObjectUid, ServerId, SyncId};
|
||||
|
||||
impl From<GenericStringObjectUniqueKey>
|
||||
for warp_graphql::generic_string_object::GenericStringObjectUniqueKey
|
||||
{
|
||||
fn from(key: GenericStringObjectUniqueKey) -> Self {
|
||||
use warp_graphql::generic_string_object::GenericStringObjectUniqueKey as GraphQLFormat;
|
||||
GraphQLFormat {
|
||||
key: key.key,
|
||||
unique_per: key.unique_per.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<UniquePer> for warp_graphql::generic_string_object::UniquePer {
|
||||
fn from(unique_per: UniquePer) -> Self {
|
||||
use warp_graphql::generic_string_object::UniquePer as GraphQLUniquePer;
|
||||
match unique_per {
|
||||
UniquePer::User => GraphQLUniquePer::User,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ObjectActionType> for warp_graphql::object_actions::ActionType {
|
||||
fn from(action: ObjectActionType) -> Self {
|
||||
match action {
|
||||
ObjectActionType::Execute => warp_graphql::object_actions::ActionType::Executed,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Converts the graphql action type ("EXECUTED", etc) to ObjectActionType.
|
||||
fn try_into_object_action_type(
|
||||
action_type: warp_graphql::object_actions::ActionType,
|
||||
) -> Result<ObjectActionType, anyhow::Error> {
|
||||
match action_type {
|
||||
warp_graphql::object_actions::ActionType::Executed => Ok(ObjectActionType::Execute),
|
||||
}
|
||||
}
|
||||
|
||||
/// Converts the graphql action entry (SingleAction, BundledActions) into its ObjectAction corollary.
|
||||
fn try_into_object_action(
|
||||
record: &warp_graphql::object_actions::ActionRecord,
|
||||
uid: ObjectUid,
|
||||
hashed_sqlite_id: HashedSqliteId,
|
||||
) -> Result<ObjectAction, anyhow::Error> {
|
||||
match record {
|
||||
warp_graphql::object_actions::ActionRecord::SingleAction(s) => Ok(ObjectAction {
|
||||
action_type: try_into_object_action_type(s.action_type)?,
|
||||
action_subtype: ObjectActionSubtype::SingleAction {
|
||||
timestamp: s.timestamp.utc(),
|
||||
processed_at_timestamp: Some(s.processed_at_timestamp.utc()),
|
||||
data: None, // The server doesn't send data for actions, although it could in the future.
|
||||
pending: false, // Actions received from the server always have pending=false.
|
||||
},
|
||||
uid,
|
||||
hashed_sqlite_id,
|
||||
}),
|
||||
warp_graphql::object_actions::ActionRecord::BundledActions(b) => Ok(ObjectAction {
|
||||
action_type: try_into_object_action_type(b.action_type)?,
|
||||
action_subtype: ObjectActionSubtype::BundledActions {
|
||||
count: b.count,
|
||||
oldest_timestamp: b.oldest_timestamp.utc(),
|
||||
latest_timestamp: b.latest_timestamp.utc(),
|
||||
latest_processed_at_timestamp: b.latest_processed_at_timestamp.utc(),
|
||||
},
|
||||
uid,
|
||||
hashed_sqlite_id,
|
||||
}),
|
||||
warp_graphql::object_actions::ActionRecord::Unknown => {
|
||||
Err(anyhow!("Unknown object action subtype"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Converts the graphql action history type into an ObjectActionHistory, requires converting
|
||||
/// the individual actions, action types, and action subtypes.
|
||||
impl TryInto<ObjectActionHistory> for warp_graphql::object_actions::ObjectActionHistory {
|
||||
type Error = anyhow::Error;
|
||||
fn try_into(self) -> Result<ObjectActionHistory, Self::Error> {
|
||||
let uid: ObjectUid = self.uid.into_inner();
|
||||
let sync_id = SyncId::ServerId(ServerId::from_string_lossy(&uid));
|
||||
let hashed_sqlite_id = sync_id.sqlite_uid_hash(self.object_type.try_into()?);
|
||||
|
||||
let actions = self
|
||||
.actions
|
||||
.map(|actions| {
|
||||
actions
|
||||
.iter()
|
||||
.filter_map(|action| {
|
||||
try_into_object_action(action, uid.clone(), hashed_sqlite_id.clone()).ok()
|
||||
})
|
||||
.collect::<Vec<ObjectAction>>()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
Ok(ObjectActionHistory {
|
||||
uid,
|
||||
hashed_sqlite_id,
|
||||
latest_processed_at_timestamp: self
|
||||
.latest_processed_at_timestamp
|
||||
.ok_or(anyhow!(
|
||||
"Parsing error: latest processed at timestamp did not exist."
|
||||
))?
|
||||
.utc(),
|
||||
actions,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
// Re-export types from warp_server_client.
|
||||
pub use warp_server_client::ids::{
|
||||
parse_sqlite_id_to_uid, ApiKeyUid, ClientId, HashableId, HashedSqliteId, ObjectUid, ServerId,
|
||||
ServerIdAndType, SyncId, ToServerId,
|
||||
};
|
||||
|
||||
/// server_id_traits is a macro used for generating implementations for the type aliases on
|
||||
/// ServerId. It implements different To/From and Display, and HashableId traits.
|
||||
/// Takes type and desired prefix for HashableId.
|
||||
///
|
||||
/// Note: This macro uses `$crate::server::ids::*` paths, so it only works within the warp crate.
|
||||
/// For types defined in warp_server_client, use `warp_server_client::server_id_traits!` instead.
|
||||
#[macro_export]
|
||||
macro_rules! server_id_traits {
|
||||
($t:ty, $prefix:literal) => {
|
||||
#[cfg(test)]
|
||||
impl From<i64> for $t {
|
||||
fn from(id: i64) -> Self {
|
||||
Self(id.into())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<String> for $t {
|
||||
fn from(id: String) -> Self {
|
||||
Self($crate::server::ids::ServerId::from_string_lossy(id))
|
||||
}
|
||||
}
|
||||
|
||||
impl From<$t> for String {
|
||||
fn from(id: $t) -> String {
|
||||
id.0.into()
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for $t {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
|
||||
write!(f, "{}", self.0)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<$t> for $crate::server::ids::ServerId {
|
||||
fn from(id: $t) -> Self {
|
||||
id.0
|
||||
}
|
||||
}
|
||||
|
||||
impl $crate::server::ids::HashableId for $t {
|
||||
fn to_hash(&self) -> String {
|
||||
format!("{}-{}", $prefix, self)
|
||||
}
|
||||
|
||||
fn from_hash(hash: &str) -> Option<$t> {
|
||||
hash.strip_prefix(&format!("{}-", $prefix))
|
||||
.map(|s| s.to_string().into())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<$crate::server::ids::ServerId> for $t {
|
||||
fn from(id: $crate::server::ids::ServerId) -> Self {
|
||||
Self(id)
|
||||
}
|
||||
}
|
||||
|
||||
impl $crate::server::ids::ToServerId for $t {
|
||||
fn to_server_id(&self) -> $crate::server::ids::ServerId {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "ids_test.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,36 @@
|
||||
use crate::{notebooks::NotebookId, workflows::WorkflowId};
|
||||
|
||||
use super::{ClientId, ServerId, SyncId};
|
||||
|
||||
#[test]
|
||||
pub fn test_client_sync_id_serialization() {
|
||||
let id: SyncId = SyncId::ClientId(ClientId::new());
|
||||
let serialized = serde_json::to_string(&id).expect("failed to serialize");
|
||||
assert_eq!(serialized, format!("\"{}\"", id.uid()));
|
||||
let deserialized: SyncId =
|
||||
serde_json::from_str(serialized.as_str()).expect("failed to deserialize");
|
||||
assert_eq!(id, deserialized);
|
||||
}
|
||||
|
||||
#[test]
|
||||
pub fn test_server_sync_id_serialization() {
|
||||
let id = SyncId::ServerId(WorkflowId::from(ServerId::from(123)).into());
|
||||
let serialized = serde_json::to_string(&id).expect("failed to serialize");
|
||||
assert_eq!(serialized, format!("\"{}\"", ServerId::from(123)));
|
||||
let deserialized: SyncId =
|
||||
serde_json::from_str(serialized.as_str()).expect("failed to deserialize");
|
||||
assert_eq!(id, deserialized);
|
||||
}
|
||||
|
||||
#[test]
|
||||
pub fn test_server_sync_id_uid_serialization() {
|
||||
let id = SyncId::ServerId(NotebookId::from(String::from("Ymgrzu0nh2HwDNeYEtXF1x")).into());
|
||||
let serialized = serde_json::to_string(&id).expect("failed to serialize");
|
||||
assert_eq!(
|
||||
serialized,
|
||||
format!("\"{}\"", String::from("Ymgrzu0nh2HwDNeYEtXF1x"))
|
||||
);
|
||||
let deserialized: SyncId =
|
||||
serde_json::from_str(serialized.as_str()).expect("failed to deserialize");
|
||||
assert_eq!(id, deserialized);
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
pub mod block;
|
||||
pub mod cloud_objects;
|
||||
pub mod datetime_ext;
|
||||
pub mod experiments;
|
||||
pub mod graphql;
|
||||
pub mod ids;
|
||||
pub mod network_log_pane_manager;
|
||||
pub mod network_log_view;
|
||||
pub mod network_logging;
|
||||
pub mod retry_strategies;
|
||||
pub mod server_api;
|
||||
pub mod sync_queue;
|
||||
pub mod telemetry;
|
||||
pub(crate) mod telemetry_ext;
|
||||
pub mod voice_transcriber;
|
||||
|
||||
pub use warp_core::operating_system_info::OperatingSystemInfo;
|
||||
@@ -0,0 +1,36 @@
|
||||
//! Tracks open [`NetworkLogPane`]s across windows so that we show at most one
|
||||
//! per window and can focus the existing one when reopened.
|
||||
//!
|
||||
//! Mirrors the pattern used by [`crate::ai::execution_profiles::editor::manager::ExecutionProfileEditorManager`].
|
||||
use std::collections::HashMap;
|
||||
|
||||
use warpui::{Entity, SingletonEntity, WindowId};
|
||||
|
||||
use crate::workspace::PaneViewLocator;
|
||||
|
||||
/// Singleton that maintains a map of `WindowId -> PaneViewLocator` for any open
|
||||
/// network log panes.
|
||||
#[derive(Default)]
|
||||
pub struct NetworkLogPaneManager {
|
||||
panes: HashMap<WindowId, PaneViewLocator>,
|
||||
}
|
||||
|
||||
impl NetworkLogPaneManager {
|
||||
pub fn find_pane(&self, window_id: WindowId) -> Option<PaneViewLocator> {
|
||||
self.panes.get(&window_id).copied()
|
||||
}
|
||||
|
||||
pub fn register_pane(&mut self, window_id: WindowId, locator: PaneViewLocator) {
|
||||
self.panes.insert(window_id, locator);
|
||||
}
|
||||
|
||||
pub fn deregister_pane(&mut self, window_id: &WindowId) {
|
||||
self.panes.remove(window_id);
|
||||
}
|
||||
}
|
||||
|
||||
impl Entity for NetworkLogPaneManager {
|
||||
type Event = ();
|
||||
}
|
||||
|
||||
impl SingletonEntity for NetworkLogPaneManager {}
|
||||
@@ -0,0 +1,256 @@
|
||||
//! A read-only pane that shows a one-shot snapshot of the in-memory
|
||||
//! [`NetworkLogModel`].
|
||||
//!
|
||||
//! The pane is opened via [`Workspace::open_network_log_pane`]. It seeds a
|
||||
//! `CodeEditorView` with the snapshot text at open time and does not live-
|
||||
//! update as new requests arrive. Re-triggering the open action while the
|
||||
//! pane is already open reloads the snapshot via [`Self::reload_snapshot`]
|
||||
//! so the user can pick up items captured since the pane was opened. The
|
||||
//! pane header also exposes a refresh icon that reloads the snapshot in
|
||||
//! place.
|
||||
use warp_editor::content::buffer::InitialBufferState;
|
||||
use warp_editor::render::element::VerticalExpansionBehavior;
|
||||
use warp_util::path::LineAndColumnArg;
|
||||
use warpui::{
|
||||
elements::{ChildView, MouseStateHandle},
|
||||
text_layout::ClipConfig,
|
||||
ui_components::components::UiComponent,
|
||||
AppContext, Element, Entity, ModelHandle, SingletonEntity, TypedActionView, View, ViewContext,
|
||||
ViewHandle,
|
||||
};
|
||||
|
||||
use crate::appearance::Appearance;
|
||||
use crate::code::editor::scroll::{ScrollPosition, ScrollTrigger};
|
||||
use crate::code::editor::view::{CodeEditorRenderOptions, CodeEditorView};
|
||||
use crate::editor::InteractionState;
|
||||
use crate::pane_group::focus_state::PaneFocusHandle;
|
||||
use crate::pane_group::{
|
||||
pane::view::{self, HeaderContent, StandardHeader, StandardHeaderOptions},
|
||||
BackingView, PaneConfiguration, PaneEvent, PaneHeaderAction,
|
||||
};
|
||||
use crate::server::network_logging::NetworkLogModel;
|
||||
use crate::ui_components::blended_colors;
|
||||
use crate::ui_components::buttons::icon_button_with_color;
|
||||
use crate::ui_components::icons;
|
||||
|
||||
/// Header text for the network log pane.
|
||||
pub const NETWORK_LOG_HEADER_TEXT: &str = "Network log";
|
||||
|
||||
/// Tooltip shown on hover over the refresh button in the pane header.
|
||||
const REFRESH_TOOLTIP: &str = "Refresh";
|
||||
|
||||
/// Event emitted by the [`NetworkLogView`].
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum NetworkLogViewEvent {
|
||||
Pane(PaneEvent),
|
||||
}
|
||||
|
||||
/// Actions supported by the pane header's overflow menu (currently none).
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum NetworkLogViewAction {}
|
||||
|
||||
/// Custom actions dispatched by elements that the [`NetworkLogView`] renders
|
||||
/// inside its pane header (e.g. the refresh button).
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum NetworkLogViewCustomAction {
|
||||
Refresh,
|
||||
}
|
||||
|
||||
/// A pane view backed by a read-only [`CodeEditorView`] displaying a snapshot
|
||||
/// of the current in-memory network log.
|
||||
pub struct NetworkLogView {
|
||||
editor: ViewHandle<CodeEditorView>,
|
||||
pane_configuration: ModelHandle<PaneConfiguration>,
|
||||
focus_handle: Option<PaneFocusHandle>,
|
||||
refresh_button_mouse_state: MouseStateHandle,
|
||||
}
|
||||
|
||||
impl NetworkLogView {
|
||||
pub fn new(ctx: &mut ViewContext<Self>) -> Self {
|
||||
let pane_configuration =
|
||||
ctx.add_model(|_ctx| PaneConfiguration::new(NETWORK_LOG_HEADER_TEXT));
|
||||
|
||||
// Capture a one-shot snapshot of the model. We intentionally do not
|
||||
// subscribe to the model: new items that arrive after the pane is
|
||||
// opened are not reflected until the pane is explicitly reopened
|
||||
// (see `reload_snapshot`).
|
||||
let snapshot = NetworkLogModel::as_ref(ctx).snapshot_text();
|
||||
|
||||
let editor = ctx.add_typed_action_view(|ctx| {
|
||||
let mut view = CodeEditorView::new(
|
||||
None,
|
||||
None,
|
||||
CodeEditorRenderOptions::new(VerticalExpansionBehavior::FillMaxHeight),
|
||||
ctx,
|
||||
);
|
||||
Self::apply_snapshot_to_editor(&mut view, &snapshot, ctx);
|
||||
// Read-only pane: disallow editing but keep selection/copy/find
|
||||
// available.
|
||||
view.set_interaction_state(InteractionState::Selectable, ctx);
|
||||
view
|
||||
});
|
||||
|
||||
Self {
|
||||
editor,
|
||||
pane_configuration,
|
||||
focus_handle: None,
|
||||
refresh_button_mouse_state: MouseStateHandle::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn pane_configuration(&self) -> ModelHandle<PaneConfiguration> {
|
||||
self.pane_configuration.clone()
|
||||
}
|
||||
|
||||
pub fn focus(&mut self, ctx: &mut ViewContext<Self>) {
|
||||
ctx.focus(&self.editor);
|
||||
}
|
||||
|
||||
/// Re-seed the editor with a fresh snapshot from [`NetworkLogModel`] and
|
||||
/// scroll back to the top. Called when the user re-triggers the
|
||||
/// open-network-log-pane action while the pane is already open, or when
|
||||
/// the user clicks the refresh icon in the pane header, so they can see
|
||||
/// items captured since the pane was opened.
|
||||
pub fn reload_snapshot(&self, ctx: &mut ViewContext<Self>) {
|
||||
let snapshot = NetworkLogModel::as_ref(ctx).snapshot_text();
|
||||
self.editor.update(ctx, |view, ctx| {
|
||||
Self::apply_snapshot_to_editor(view, &snapshot, ctx);
|
||||
});
|
||||
}
|
||||
|
||||
/// Resets the editor buffer with the given snapshot text and queues a
|
||||
/// pending scroll-to-top once layout completes. `reset` places the
|
||||
/// cursor at the end of the buffer by default, which would scroll the
|
||||
/// viewport to the bottom when the pane renders.
|
||||
fn apply_snapshot_to_editor(
|
||||
view: &mut CodeEditorView,
|
||||
snapshot: &str,
|
||||
ctx: &mut ViewContext<CodeEditorView>,
|
||||
) {
|
||||
let state = InitialBufferState::plain_text(snapshot);
|
||||
view.reset(state, ctx);
|
||||
let version = view.buffer_version(ctx);
|
||||
view.set_pending_scroll(ScrollTrigger::new(
|
||||
ScrollPosition::LineAndColumn(LineAndColumnArg {
|
||||
line_num: 1,
|
||||
column_num: Some(0),
|
||||
}),
|
||||
version,
|
||||
));
|
||||
}
|
||||
|
||||
/// Renders the refresh icon button for the pane header. Clicking the
|
||||
/// button dispatches [`NetworkLogViewCustomAction::Refresh`], which the
|
||||
/// pane header forwards back to [`Self::handle_custom_action`].
|
||||
fn render_refresh_button(&self, app: &AppContext) -> Box<dyn Element> {
|
||||
let appearance = Appearance::as_ref(app);
|
||||
let theme = appearance.theme();
|
||||
let ui_builder = appearance.ui_builder().clone();
|
||||
|
||||
icon_button_with_color(
|
||||
appearance,
|
||||
icons::Icon::Refresh,
|
||||
false, /* active */
|
||||
self.refresh_button_mouse_state.clone(),
|
||||
blended_colors::text_sub(theme, theme.background()).into(),
|
||||
)
|
||||
.with_tooltip(move || {
|
||||
ui_builder
|
||||
.tool_tip(REFRESH_TOOLTIP.to_string())
|
||||
.build()
|
||||
.finish()
|
||||
})
|
||||
.build()
|
||||
.on_click(|ctx, _, _| {
|
||||
ctx.dispatch_typed_action::<PaneHeaderAction<
|
||||
NetworkLogViewAction,
|
||||
NetworkLogViewCustomAction,
|
||||
>>(PaneHeaderAction::CustomAction(
|
||||
NetworkLogViewCustomAction::Refresh,
|
||||
));
|
||||
})
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl Entity for NetworkLogView {
|
||||
type Event = NetworkLogViewEvent;
|
||||
}
|
||||
|
||||
impl View for NetworkLogView {
|
||||
fn ui_name() -> &'static str {
|
||||
"NetworkLogView"
|
||||
}
|
||||
|
||||
fn render(&self, _app: &AppContext) -> Box<dyn Element> {
|
||||
ChildView::new(&self.editor).finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl TypedActionView for NetworkLogView {
|
||||
type Action = NetworkLogViewAction;
|
||||
|
||||
fn handle_action(&mut self, _action: &Self::Action, _ctx: &mut ViewContext<Self>) {
|
||||
// NetworkLogViewAction is currently uninhabited.
|
||||
}
|
||||
}
|
||||
|
||||
impl BackingView for NetworkLogView {
|
||||
type PaneHeaderOverflowMenuAction = NetworkLogViewAction;
|
||||
type CustomAction = NetworkLogViewCustomAction;
|
||||
type AssociatedData = ();
|
||||
|
||||
fn handle_pane_header_overflow_menu_action(
|
||||
&mut self,
|
||||
_action: &Self::PaneHeaderOverflowMenuAction,
|
||||
_ctx: &mut ViewContext<Self>,
|
||||
) {
|
||||
// No overflow menu items are registered.
|
||||
}
|
||||
|
||||
fn handle_custom_action(
|
||||
&mut self,
|
||||
custom_action: &Self::CustomAction,
|
||||
ctx: &mut ViewContext<Self>,
|
||||
) {
|
||||
match custom_action {
|
||||
NetworkLogViewCustomAction::Refresh => self.reload_snapshot(ctx),
|
||||
}
|
||||
}
|
||||
|
||||
fn close(&mut self, ctx: &mut ViewContext<Self>) {
|
||||
ctx.emit(NetworkLogViewEvent::Pane(PaneEvent::Close));
|
||||
}
|
||||
|
||||
fn focus_contents(&mut self, ctx: &mut ViewContext<Self>) {
|
||||
self.focus(ctx);
|
||||
}
|
||||
|
||||
fn render_header_content(
|
||||
&self,
|
||||
_ctx: &view::HeaderRenderContext<'_>,
|
||||
app: &AppContext,
|
||||
) -> HeaderContent {
|
||||
HeaderContent::Standard(StandardHeader {
|
||||
title: NETWORK_LOG_HEADER_TEXT.to_string(),
|
||||
title_secondary: None,
|
||||
title_style: None,
|
||||
title_clip_config: ClipConfig::start(),
|
||||
title_max_width: None,
|
||||
left_of_title: None,
|
||||
right_of_title: None,
|
||||
left_of_overflow: Some(self.render_refresh_button(app)),
|
||||
// Keep the close button always visible so hovering the header
|
||||
// doesn't cause the refresh button to shift horizontally as the
|
||||
// close button appears.
|
||||
options: StandardHeaderOptions {
|
||||
always_show_icons: true,
|
||||
..StandardHeaderOptions::default()
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
fn set_focus_handle(&mut self, focus_handle: PaneFocusHandle, _ctx: &mut ViewContext<Self>) {
|
||||
self.focus_handle = Some(focus_handle);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
use std::fmt;
|
||||
|
||||
use bounded_vec_deque::BoundedVecDeque;
|
||||
use chrono::{DateTime, FixedOffset};
|
||||
use enclose::enclose;
|
||||
use warpui::{Entity, ModelContext, SingletonEntity};
|
||||
|
||||
use crate::server::datetime_ext::DateTimeExt;
|
||||
use crate::server::server_api::ServerApiProvider;
|
||||
|
||||
/// Maximum number of network log items retained in memory. Matches the
|
||||
/// previous file-rotation threshold so the pane surface behaves consistently
|
||||
/// with historical expectations.
|
||||
const NETWORK_LOGGING_MAX_ITEMS: usize = 50;
|
||||
|
||||
/// Upper bound on the bounded async channel between the HTTP client hooks and
|
||||
/// the in-memory model. Keeps a small backlog to tolerate bursts without
|
||||
/// blocking the request thread.
|
||||
const NETWORK_LOGGING_MAX_QUEUE_SIZE: usize = 100;
|
||||
|
||||
/// In-memory store of the most recent network log items. Populated by
|
||||
/// [`init`] and read by the network log pane. Holds at most
|
||||
/// [`NETWORK_LOGGING_MAX_ITEMS`] entries; older entries are dropped when new
|
||||
/// ones arrive.
|
||||
pub struct NetworkLogModel {
|
||||
items: BoundedVecDeque<NetworkLogItem>,
|
||||
}
|
||||
|
||||
impl Default for NetworkLogModel {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
items: BoundedVecDeque::new(NETWORK_LOGGING_MAX_ITEMS),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl NetworkLogModel {
|
||||
/// Appends a new log item, evicting the oldest if at capacity.
|
||||
pub fn push(&mut self, item: NetworkLogItem, ctx: &mut ModelContext<Self>) {
|
||||
// `BoundedVecDeque::push_back` returns the evicted item when the
|
||||
// store is at capacity; we discard it since the pane only needs the
|
||||
// most recent entries.
|
||||
let _evicted = self.items.push_back(item);
|
||||
ctx.notify();
|
||||
}
|
||||
|
||||
/// Returns the current snapshot as a single string with one item per line,
|
||||
/// in chronological order. Returns an empty string when no items have been
|
||||
/// captured.
|
||||
pub fn snapshot_text(&self) -> String {
|
||||
let mut out = String::new();
|
||||
for (i, item) in self.items.iter().enumerate() {
|
||||
if i > 0 {
|
||||
out.push('\n');
|
||||
}
|
||||
out.push_str(&item.0);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Number of items currently retained. Exposed for tests.
|
||||
#[cfg(test)]
|
||||
pub fn len(&self) -> usize {
|
||||
self.items.len()
|
||||
}
|
||||
}
|
||||
|
||||
impl Entity for NetworkLogModel {
|
||||
type Event = ();
|
||||
}
|
||||
|
||||
impl SingletonEntity for NetworkLogModel {}
|
||||
|
||||
/// Initializes a network logging task that listens for requests that pass
|
||||
/// through the provided HTTP clients and forwards them to the in-memory
|
||||
/// [`NetworkLogModel`].
|
||||
///
|
||||
/// The logging happens via an async channel so that request hooks never block
|
||||
/// on the main thread. Items are delivered to the model on the main thread via
|
||||
/// [`ModelContext::spawn_stream_local`], mirroring how `ServerApiProvider`
|
||||
/// consumes its own event stream.
|
||||
pub(super) fn init<'a>(
|
||||
http_clients: impl IntoIterator<Item = &'a mut http_client::Client>,
|
||||
ctx: &mut ModelContext<ServerApiProvider>,
|
||||
) {
|
||||
let (tx, rx) = async_channel::bounded::<NetworkLogItem>(NETWORK_LOGGING_MAX_QUEUE_SIZE);
|
||||
|
||||
ctx.spawn_stream_local(
|
||||
rx,
|
||||
move |_, item, ctx| {
|
||||
NetworkLogModel::handle(ctx).update(ctx, |model, ctx| {
|
||||
model.push(item, ctx);
|
||||
});
|
||||
},
|
||||
|_, _| {},
|
||||
);
|
||||
|
||||
for client in http_clients.into_iter() {
|
||||
client.set_before_request_fn(Box::new(enclose!((tx) move |request, serialized_payload| {
|
||||
if !tx.is_closed() {
|
||||
if let Err(e) = tx.try_send(NetworkLogItem::request(
|
||||
request,
|
||||
serialized_payload.clone(),
|
||||
DateTime::now(),
|
||||
)) {
|
||||
log::error!(
|
||||
"Error sending request from http client to logging task: {e}"
|
||||
);
|
||||
}
|
||||
}
|
||||
})));
|
||||
|
||||
client.set_after_response_fn(Box::new(enclose!((tx) move |response| {
|
||||
if !tx.is_closed() {
|
||||
if let Err(e) = tx.try_send(NetworkLogItem::response(response, DateTime::now())) {
|
||||
log::error!("Error sending request from http client to logging task: {e}");
|
||||
}
|
||||
}
|
||||
})));
|
||||
}
|
||||
}
|
||||
|
||||
/// Represents an item (either a request or response) captured for the network
|
||||
/// activity log. The inner string contains a timestamp and the
|
||||
/// [`Debug`]-formatted representation of the request or response, matching the
|
||||
/// format previously written to `warp_network.log`.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct NetworkLogItem(String);
|
||||
|
||||
impl NetworkLogItem {
|
||||
pub fn request(
|
||||
request: &reqwest::Request,
|
||||
serialized_payload: Option<String>,
|
||||
timestamp: DateTime<FixedOffset>,
|
||||
) -> Self {
|
||||
Self(format!(
|
||||
"[{}]: {:?}{}",
|
||||
timestamp.format("%Y-%m-%d %H:%M:%S,%3f"),
|
||||
request,
|
||||
serialized_payload.map_or("".to_owned(), |payload| format!("\nBody {payload}"))
|
||||
))
|
||||
}
|
||||
|
||||
pub fn response(response: &reqwest::Response, timestamp: DateTime<FixedOffset>) -> Self {
|
||||
Self(format!(
|
||||
"[{}]: {:?}",
|
||||
timestamp.format("%Y-%m-%d %H:%M:%S,%3f"),
|
||||
response
|
||||
))
|
||||
}
|
||||
|
||||
/// Constructs a log item directly from a pre-formatted string. Used in
|
||||
/// tests where we don't have a real `reqwest` request/response handy.
|
||||
#[cfg(test)]
|
||||
pub fn from_string(s: impl Into<String>) -> Self {
|
||||
Self(s.into())
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for NetworkLogItem {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "{}", self.0)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "network_logging_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,70 @@
|
||||
use super::{NetworkLogItem, NetworkLogModel, NETWORK_LOGGING_MAX_ITEMS};
|
||||
use warpui::App;
|
||||
|
||||
#[test]
|
||||
fn empty_snapshot_is_empty_string() {
|
||||
App::test((), |app| async move {
|
||||
let model = app.add_singleton_model(|_| NetworkLogModel::default());
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(model.snapshot_text(), "");
|
||||
assert_eq!(model.len(), 0);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_joins_items_with_newlines() {
|
||||
App::test((), |mut app| async move {
|
||||
let model = app.add_singleton_model(|_| NetworkLogModel::default());
|
||||
model.update(&mut app, |model, ctx| {
|
||||
model.push(NetworkLogItem::from_string("first"), ctx);
|
||||
model.push(NetworkLogItem::from_string("second"), ctx);
|
||||
model.push(NetworkLogItem::from_string("third"), ctx);
|
||||
});
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(model.snapshot_text(), "first\nsecond\nthird");
|
||||
assert_eq!(model.len(), 3);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn push_beyond_capacity_drops_oldest() {
|
||||
App::test((), |mut app| async move {
|
||||
let model = app.add_singleton_model(|_| NetworkLogModel::default());
|
||||
// Push exactly the capacity; the snapshot should contain all items
|
||||
// and the count should equal the capacity.
|
||||
model.update(&mut app, |model, ctx| {
|
||||
for i in 0..NETWORK_LOGGING_MAX_ITEMS {
|
||||
model.push(NetworkLogItem::from_string(format!("item-{i}")), ctx);
|
||||
}
|
||||
});
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(model.len(), NETWORK_LOGGING_MAX_ITEMS);
|
||||
// Oldest is still present when we're exactly at capacity.
|
||||
assert!(model.snapshot_text().starts_with("item-0\n"));
|
||||
});
|
||||
|
||||
// Push one more: the oldest item should be evicted so the store stays
|
||||
// at capacity, and the snapshot should start at item-1 now.
|
||||
model.update(&mut app, |model, ctx| {
|
||||
model.push(NetworkLogItem::from_string("overflow"), ctx);
|
||||
});
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(model.len(), NETWORK_LOGGING_MAX_ITEMS);
|
||||
assert!(!model.snapshot_text().contains("item-0\n"));
|
||||
assert!(model.snapshot_text().starts_with("item-1\n"));
|
||||
assert!(model.snapshot_text().ends_with("\noverflow"));
|
||||
});
|
||||
|
||||
// Pushing many additional items keeps the store at capacity.
|
||||
model.update(&mut app, |model, ctx| {
|
||||
for i in 0..10 {
|
||||
model.push(NetworkLogItem::from_string(format!("extra-{i}")), ctx);
|
||||
}
|
||||
});
|
||||
model.read(&app, |model, _| {
|
||||
assert_eq!(model.len(), NETWORK_LOGGING_MAX_ITEMS);
|
||||
});
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
use std::time::Duration;
|
||||
use warpui::RetryOption;
|
||||
|
||||
use crate::server::server_api::presigned_upload::HttpStatusError;
|
||||
|
||||
/// Common duration for a periodic poll. In our app, we generally have the following to update the same data:
|
||||
/// - RTC messages
|
||||
/// - Out-of-band queries based on user actions (i.e. fetch team info when user opens the settings page, user
|
||||
/// starts the app)
|
||||
/// However, we also periodically poll for updates in case RTC is down, the user's websocket
|
||||
/// is borked, etc.
|
||||
/// For team memberships, we also don't yet process messages for joining or leaving a team, so the user would see these
|
||||
/// updates only after a periodic poll.
|
||||
pub const PERIODIC_POLL: Duration = Duration::from_secs(60 * 10);
|
||||
|
||||
/// For a periodic poll, it's fine to wait for longer period of time between retries. However, we don't want this to be so
|
||||
/// long that it's around the same as the overall periodic poll interval.
|
||||
pub const PERIODIC_POLL_RETRY_STRATEGY: RetryOption = RetryOption::exponential(
|
||||
Duration::from_secs(2), /* interval */
|
||||
2., /* exponential factor */
|
||||
3, /* max retry count */
|
||||
)
|
||||
.with_jitter(0.2 /* max_jitter_percentage */);
|
||||
|
||||
/// When there's an out-of-band request for a periodic poll, we want to retry quickly, because the UI is depending on the
|
||||
/// request succeeding in a timely way. These are things like loading all object updates upon startup, checking the team
|
||||
/// metadata when we visit the team page, etc.
|
||||
pub const OUT_OF_BAND_REQUEST_RETRY_STRATEGY: RetryOption = RetryOption::exponential(
|
||||
Duration::from_millis(100), /* interval */
|
||||
5., /* exponential factor */
|
||||
3, /* max retry count */
|
||||
)
|
||||
.with_jitter(0.5 /* max_jitter_percentage */);
|
||||
|
||||
// For listeners, retry up to 5 times, waiting between 10-40 seconds between retries.
|
||||
pub const LISTENER_RETRY_STRATEGY: RetryOption = RetryOption::linear(
|
||||
Duration::from_secs(25), /* interval */
|
||||
5, /* max retry count */
|
||||
)
|
||||
.with_jitter(0.6 /* max_jitter_multiplier */);
|
||||
|
||||
/// Classify an HTTP-backed error as transient (worth retrying) or permanent (fail fast).
|
||||
///
|
||||
/// Transient: 5xx responses, 408, 429, or any error whose chain does not carry an
|
||||
/// [`HttpStatusError`] (connection reset, timeout, DNS failure, etc.).
|
||||
/// Permanent: other 4xx responses (bad signature, 404, 403, etc.).
|
||||
pub(crate) fn is_transient_http_error(e: &anyhow::Error) -> bool {
|
||||
// Callers typically wrap an `HttpStatusError` cause with a `.context(...)` message for
|
||||
// human-friendly Display, so the typed error sits somewhere in the chain rather than as
|
||||
// the top-level error object — walk the chain.
|
||||
for cause in e.chain() {
|
||||
if let Some(http_err) = cause.downcast_ref::<HttpStatusError>() {
|
||||
return matches!(http_err.status, 408 | 429 | 500..=599);
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,974 @@
|
||||
use chrono::TimeZone;
|
||||
use chrono::Utc;
|
||||
|
||||
use super::{
|
||||
build_list_agent_runs_url, AgentMessageHeader, AgentRunEvent, AgentSource,
|
||||
AmbientAgentTaskState, Artifact, ArtifactDownloadResponse, ArtifactType, ExecutionLocation,
|
||||
ListRunsResponse, ReadAgentMessageResponse, RunSortBy, RunSortOrder, TaskListFilter,
|
||||
};
|
||||
use crate::notebooks::NotebookId;
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_file_artifact_download_response() {
|
||||
let json = r#"{
|
||||
"artifact_uid": "artifact-123",
|
||||
"artifact_type": "FILE",
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"data": {
|
||||
"download_url": "https://storage.example.com/report.txt",
|
||||
"expires_at": "2024-01-15T11:30:00Z",
|
||||
"content_type": "text/plain",
|
||||
"filepath": "outputs/report.txt",
|
||||
"filename": "report.txt",
|
||||
"description": "daily summary",
|
||||
"size_bytes": 42
|
||||
}
|
||||
}"#;
|
||||
|
||||
let artifact: ArtifactDownloadResponse = serde_json::from_str(json).unwrap();
|
||||
|
||||
let ArtifactDownloadResponse::File { common, data } = artifact else {
|
||||
panic!("expected File artifact download response");
|
||||
};
|
||||
assert_eq!(common.artifact_uid, "artifact-123");
|
||||
assert_eq!(common.created_at.to_rfc3339(), "2024-01-15T10:30:00+00:00");
|
||||
assert_eq!(data.download_url, "https://storage.example.com/report.txt");
|
||||
assert_eq!(data.expires_at.to_rfc3339(), "2024-01-15T11:30:00+00:00");
|
||||
assert_eq!(data.content_type, "text/plain");
|
||||
assert_eq!(data.filepath, "outputs/report.txt");
|
||||
assert_eq!(data.filename, "report.txt");
|
||||
assert_eq!(data.description.as_deref(), Some("daily summary"));
|
||||
assert_eq!(data.size_bytes, Some(42));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_screenshot_artifact_download_response() {
|
||||
let json = r#"{
|
||||
"artifact_uid": "screenshot-123",
|
||||
"artifact_type": "SCREENSHOT",
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"data": {
|
||||
"download_url": "https://storage.example.com/screenshot.png",
|
||||
"expires_at": "2024-01-15T11:30:00Z",
|
||||
"content_type": "image/png",
|
||||
"description": "dashboard screenshot"
|
||||
}
|
||||
}"#;
|
||||
|
||||
let artifact: ArtifactDownloadResponse = serde_json::from_str(json).unwrap();
|
||||
|
||||
let ArtifactDownloadResponse::Screenshot { common, data } = artifact else {
|
||||
panic!("expected Screenshot artifact download response");
|
||||
};
|
||||
assert_eq!(common.artifact_uid, "screenshot-123");
|
||||
assert_eq!(common.created_at.to_rfc3339(), "2024-01-15T10:30:00+00:00");
|
||||
assert_eq!(
|
||||
data.download_url,
|
||||
"https://storage.example.com/screenshot.png"
|
||||
);
|
||||
assert_eq!(data.expires_at.to_rfc3339(), "2024-01-15T11:30:00+00:00");
|
||||
assert_eq!(data.content_type, "image/png");
|
||||
assert_eq!(data.description.as_deref(), Some("dashboard screenshot"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_plan_artifact() {
|
||||
let json = r#"{
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"artifact_type": "PLAN",
|
||||
"data": {
|
||||
"document_uid": "doc-uid-123",
|
||||
"notebook_uid": "1234567890123456789012",
|
||||
"title": "My Plan"
|
||||
}
|
||||
}"#;
|
||||
|
||||
let artifact: Artifact = serde_json::from_str(json).unwrap();
|
||||
|
||||
let Artifact::Plan {
|
||||
document_uid,
|
||||
notebook_uid,
|
||||
title,
|
||||
} = &artifact
|
||||
else {
|
||||
panic!("expected Plan artifact");
|
||||
};
|
||||
assert_eq!(document_uid, "doc-uid-123");
|
||||
assert_eq!(
|
||||
notebook_uid.as_ref().map(|n| n.to_string()),
|
||||
Some("1234567890123456789012".to_string())
|
||||
);
|
||||
assert_eq!(*title, Some("My Plan".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_pull_request_artifact() {
|
||||
let json = r#"{
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"artifact_type": "PULL_REQUEST",
|
||||
"data": {
|
||||
"url": "https://github.com/org/repo/pull/42",
|
||||
"branch": "feature-branch"
|
||||
}
|
||||
}"#;
|
||||
|
||||
let artifact: Artifact = serde_json::from_str(json).unwrap();
|
||||
|
||||
let Artifact::PullRequest {
|
||||
url,
|
||||
branch,
|
||||
repo,
|
||||
number,
|
||||
} = &artifact
|
||||
else {
|
||||
panic!("expected PullRequest artifact");
|
||||
};
|
||||
assert_eq!(url, "https://github.com/org/repo/pull/42");
|
||||
assert_eq!(branch, "feature-branch");
|
||||
assert_eq!(*repo, Some("repo".to_string()));
|
||||
assert_eq!(*number, Some(42));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_pull_request_non_github_url() {
|
||||
let json = r#"{
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"artifact_type": "PULL_REQUEST",
|
||||
"data": {
|
||||
"url": "https://gitlab.com/org/repo/merge_requests/42",
|
||||
"branch": "feature-branch"
|
||||
}
|
||||
}"#;
|
||||
|
||||
let artifact: Artifact = serde_json::from_str(json).unwrap();
|
||||
|
||||
let Artifact::PullRequest { repo, number, .. } = &artifact else {
|
||||
panic!("expected PullRequest artifact");
|
||||
};
|
||||
assert_eq!(*repo, None);
|
||||
assert_eq!(*number, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_plan_artifact_with_optional_fields_missing() {
|
||||
let json = r#"{
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"artifact_type": "PLAN",
|
||||
"data": {
|
||||
"document_uid": "doc-uid-123",
|
||||
"notebook_uid": "abcdefghijklmnopqrstuv"
|
||||
}
|
||||
}"#;
|
||||
|
||||
let artifact: Artifact = serde_json::from_str(json).unwrap();
|
||||
|
||||
let Artifact::Plan {
|
||||
document_uid,
|
||||
notebook_uid,
|
||||
title,
|
||||
} = &artifact
|
||||
else {
|
||||
panic!("expected Plan artifact");
|
||||
};
|
||||
assert_eq!(document_uid, "doc-uid-123");
|
||||
assert_eq!(
|
||||
notebook_uid.as_ref().map(|n| n.to_string()),
|
||||
Some("abcdefghijklmnopqrstuv".to_string())
|
||||
);
|
||||
assert!(title.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_list_tasks_response_with_artifacts() {
|
||||
let json = r#"{
|
||||
"runs": [
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440000",
|
||||
"title": "Test Task",
|
||||
"state": "SUCCEEDED",
|
||||
"prompt": "test prompt",
|
||||
"created_at": "2024-01-15T10:00:00Z",
|
||||
"updated_at": "2024-01-15T10:30:00Z",
|
||||
"is_sandbox_running": true,
|
||||
"artifacts": [
|
||||
{
|
||||
"created_at": "2024-01-15T10:20:00Z",
|
||||
"artifact_type": "PLAN",
|
||||
"data": {
|
||||
"document_uid": "doc-1",
|
||||
"notebook_uid": "xyz1234567890123456789",
|
||||
"title": "Plan Title"
|
||||
}
|
||||
},
|
||||
{
|
||||
"created_at": "2024-01-15T10:25:00Z",
|
||||
"artifact_type": "PULL_REQUEST",
|
||||
"data": {
|
||||
"url": "https://github.com/org/repo/pull/1",
|
||||
"branch": "main"
|
||||
}
|
||||
},
|
||||
{
|
||||
"created_at": "2024-01-15T10:27:00Z",
|
||||
"artifact_type": "FILE",
|
||||
"data": {
|
||||
"artifact_uid": "artifact-file-1",
|
||||
"filepath": "outputs/report.txt",
|
||||
"filename": "report.txt",
|
||||
"mime_type": "text/plain",
|
||||
"description": "Daily summary",
|
||||
"size_bytes": 42
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}"#;
|
||||
|
||||
let response: ListRunsResponse = serde_json::from_str(json).unwrap();
|
||||
|
||||
assert_eq!(response.runs.len(), 1);
|
||||
let task = &response.runs[0];
|
||||
assert_eq!(
|
||||
task.task_id.to_string(),
|
||||
"550e8400-e29b-41d4-a716-446655440000"
|
||||
);
|
||||
assert_eq!(task.artifacts.len(), 3);
|
||||
|
||||
// Check first artifact (Plan)
|
||||
let Artifact::Plan {
|
||||
document_uid,
|
||||
title,
|
||||
..
|
||||
} = &task.artifacts[0]
|
||||
else {
|
||||
panic!("expected Plan artifact");
|
||||
};
|
||||
assert_eq!(document_uid, "doc-1");
|
||||
assert_eq!(*title, Some("Plan Title".to_string()));
|
||||
|
||||
// Check second artifact (PullRequest)
|
||||
let Artifact::PullRequest {
|
||||
url,
|
||||
branch,
|
||||
repo,
|
||||
number,
|
||||
..
|
||||
} = &task.artifacts[1]
|
||||
else {
|
||||
panic!("expected PullRequest artifact");
|
||||
};
|
||||
assert_eq!(url, "https://github.com/org/repo/pull/1");
|
||||
assert_eq!(branch, "main");
|
||||
assert_eq!(*repo, Some("repo".to_string()));
|
||||
assert_eq!(*number, Some(1));
|
||||
|
||||
let Artifact::File {
|
||||
artifact_uid,
|
||||
filepath,
|
||||
filename,
|
||||
mime_type,
|
||||
description,
|
||||
size_bytes,
|
||||
} = &task.artifacts[2]
|
||||
else {
|
||||
panic!("expected File artifact");
|
||||
};
|
||||
assert_eq!(artifact_uid, "artifact-file-1");
|
||||
assert_eq!(filepath, "outputs/report.txt");
|
||||
assert_eq!(filename, "report.txt");
|
||||
assert_eq!(mime_type, "text/plain");
|
||||
assert_eq!(*description, Some("Daily summary".to_string()));
|
||||
assert_eq!(*size_bytes, Some(42));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_list_tasks_response_empty_artifacts() {
|
||||
let json = r#"{
|
||||
"runs": [
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440001",
|
||||
"title": "Test Task",
|
||||
"state": "INPROGRESS",
|
||||
"prompt": "test prompt",
|
||||
"created_at": "2024-01-15T10:00:00Z",
|
||||
"updated_at": "2024-01-15T10:30:00Z",
|
||||
"is_sandbox_running": true,
|
||||
"artifacts": []
|
||||
}
|
||||
]
|
||||
}"#;
|
||||
|
||||
let response: ListRunsResponse = serde_json::from_str(json).unwrap();
|
||||
|
||||
assert_eq!(response.runs.len(), 1);
|
||||
assert!(response.runs[0].artifacts.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_list_tasks_response_missing_artifacts_field() {
|
||||
// Server may not include artifacts field at all for older responses
|
||||
let json = r#"{
|
||||
"runs": [
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440002",
|
||||
"title": "Test Task",
|
||||
"state": "QUEUED",
|
||||
"prompt": "test prompt",
|
||||
"created_at": "2024-01-15T10:00:00Z",
|
||||
"updated_at": "2024-01-15T10:30:00Z",
|
||||
"is_sandbox_running": true
|
||||
}
|
||||
]
|
||||
}"#;
|
||||
|
||||
let response: ListRunsResponse = serde_json::from_str(json).unwrap();
|
||||
|
||||
assert_eq!(response.runs.len(), 1);
|
||||
assert!(response.runs[0].artifacts.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_artifacts_skips_invalid_items() {
|
||||
// deserialize_artifacts should skip invalid items and keep valid ones
|
||||
let json = r#"{
|
||||
"runs": [
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440000",
|
||||
"title": "Test Task",
|
||||
"state": "SUCCEEDED",
|
||||
"prompt": "test prompt",
|
||||
"created_at": "2024-01-15T10:00:00Z",
|
||||
"updated_at": "2024-01-15T10:30:00Z",
|
||||
"is_sandbox_running": true,
|
||||
"artifacts": [
|
||||
{
|
||||
"created_at": "2024-01-15T10:20:00Z",
|
||||
"artifact_type": "PLAN",
|
||||
"data": {
|
||||
"document_uid": "valid-doc",
|
||||
"notebook_uid": "validnotebook123456789",
|
||||
"title": "Valid Plan"
|
||||
}
|
||||
},
|
||||
{
|
||||
"created_at": "2024-01-15T10:25:00Z",
|
||||
"artifact_type": "UNKNOWN_TYPE",
|
||||
"data": {
|
||||
"some_field": "value"
|
||||
}
|
||||
},
|
||||
{
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"artifact_type": "PULL_REQUEST",
|
||||
"data": {
|
||||
"url": "https://github.com/org/repo/pull/1",
|
||||
"branch": "main"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}"#;
|
||||
|
||||
let response: ListRunsResponse = serde_json::from_str(json).unwrap();
|
||||
|
||||
assert_eq!(response.runs.len(), 1);
|
||||
// Invalid artifact skipped, valid ones kept
|
||||
assert_eq!(response.runs[0].artifacts.len(), 2);
|
||||
assert!(matches!(
|
||||
response.runs[0].artifacts[0],
|
||||
Artifact::Plan { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
response.runs[0].artifacts[1],
|
||||
Artifact::PullRequest { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_artifacts_all_invalid_returns_empty() {
|
||||
// When all artifacts are invalid, result should be empty vec
|
||||
let json = r#"{
|
||||
"runs": [
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440000",
|
||||
"title": "Test Task",
|
||||
"state": "SUCCEEDED",
|
||||
"prompt": "test prompt",
|
||||
"created_at": "2024-01-15T10:00:00Z",
|
||||
"updated_at": "2024-01-15T10:30:00Z",
|
||||
"is_sandbox_running": true,
|
||||
"artifacts": [
|
||||
{
|
||||
"created_at": "2024-01-15T10:20:00Z",
|
||||
"artifact_type": "UNKNOWN_TYPE",
|
||||
"data": {}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}"#;
|
||||
|
||||
let response: ListRunsResponse = serde_json::from_str(json).unwrap();
|
||||
|
||||
assert_eq!(response.runs.len(), 1);
|
||||
assert!(response.runs[0].artifacts.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_artifact_missing_data_field() {
|
||||
let json = r#"{
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"artifact_type": "PLAN"
|
||||
}"#;
|
||||
|
||||
let result = serde_json::from_str::<Artifact>(json);
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().to_string().contains("missing field"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_artifact_invalid_plan_data() {
|
||||
// Missing required `document_uid` field should fail deserialization
|
||||
let json = r#"{
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"artifact_type": "PLAN",
|
||||
"data": {
|
||||
"title": "Only title, no document_uid"
|
||||
}
|
||||
}"#;
|
||||
|
||||
let result = serde_json::from_str::<Artifact>(json);
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().to_string().contains("missing field"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_artifact_invalid_pr_data() {
|
||||
let json = r#"{
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"artifact_type": "PULL_REQUEST",
|
||||
"data": {
|
||||
"url": "https://github.com/org/repo/pull/1"
|
||||
}
|
||||
}"#;
|
||||
|
||||
let result = serde_json::from_str::<Artifact>(json);
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().to_string().contains("missing field"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_artifact_unknown_variant() {
|
||||
let json = r#"{
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"artifact_type": "UNKNOWN_TYPE",
|
||||
"data": {
|
||||
"some_field": "value"
|
||||
}
|
||||
}"#;
|
||||
|
||||
let result = serde_json::from_str::<Artifact>(json);
|
||||
assert!(result.is_err());
|
||||
let error_msg = result.unwrap_err().to_string();
|
||||
assert!(error_msg.contains("unknown variant"));
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------------------------------------------------
|
||||
// Tests for resilient task list deserialization (skipping malformed tasks while tolerating unknown states)
|
||||
// ---------------------------------------------------------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_list_tasks_skips_invalid_task() {
|
||||
// One valid task and one invalid task (missing required field)
|
||||
let json = r#"{
|
||||
"runs": [
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440000",
|
||||
"title": "Valid Task",
|
||||
"state": "SUCCEEDED",
|
||||
"prompt": "test prompt",
|
||||
"created_at": "2024-01-15T10:00:00Z",
|
||||
"updated_at": "2024-01-15T10:30:00Z",
|
||||
"is_sandbox_running": true
|
||||
},
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440001",
|
||||
"title": "Invalid Task",
|
||||
"state": "INPROGRESS"
|
||||
}
|
||||
]
|
||||
}"#;
|
||||
|
||||
let response: ListRunsResponse = serde_json::from_str(json).unwrap();
|
||||
|
||||
// Should only have the valid task
|
||||
assert_eq!(response.runs.len(), 1);
|
||||
assert_eq!(
|
||||
response.runs[0].task_id.to_string(),
|
||||
"550e8400-e29b-41d4-a716-446655440000"
|
||||
);
|
||||
assert_eq!(response.runs[0].title, "Valid Task");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_list_tasks_error_and_blocked_states() {
|
||||
let json = r#"{
|
||||
"runs": [
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440000",
|
||||
"title": "Errored Task",
|
||||
"state": "ERROR",
|
||||
"prompt": "test prompt",
|
||||
"created_at": "2024-01-15T10:00:00Z",
|
||||
"updated_at": "2024-01-15T10:30:00Z",
|
||||
"is_sandbox_running": false
|
||||
},
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440001",
|
||||
"title": "Blocked Task",
|
||||
"state": "BLOCKED",
|
||||
"prompt": "test prompt",
|
||||
"created_at": "2024-01-15T10:00:00Z",
|
||||
"updated_at": "2024-01-15T10:30:00Z",
|
||||
"is_sandbox_running": false
|
||||
}
|
||||
]
|
||||
}"#;
|
||||
|
||||
let response: ListRunsResponse = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(response.runs.len(), 2);
|
||||
assert_eq!(response.runs[0].state, AmbientAgentTaskState::Error);
|
||||
assert_eq!(response.runs[1].state, AmbientAgentTaskState::Blocked);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_list_tasks_all_tasks_invalid_returns_empty() {
|
||||
// All tasks are missing required fields
|
||||
let json = r#"{
|
||||
"runs": [
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440000",
|
||||
"title": "Missing State"
|
||||
},
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440001",
|
||||
"state": "SUCCEEDED"
|
||||
}
|
||||
]
|
||||
}"#;
|
||||
|
||||
let response: ListRunsResponse = serde_json::from_str(json).unwrap();
|
||||
|
||||
// Should return empty list, not fail
|
||||
assert_eq!(response.runs.len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_list_tasks_invalid_state_enum() {
|
||||
// Task with an unknown state enum value
|
||||
let json = r#"{
|
||||
"runs": [
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440000",
|
||||
"title": "Valid Task",
|
||||
"state": "SUCCEEDED",
|
||||
"prompt": "test prompt",
|
||||
"created_at": "2024-01-15T10:00:00Z",
|
||||
"updated_at": "2024-01-15T10:30:00Z",
|
||||
"is_sandbox_running": true
|
||||
},
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440001",
|
||||
"title": "Task with Invalid State",
|
||||
"state": "INVALID_STATE",
|
||||
"prompt": "test prompt",
|
||||
"created_at": "2024-01-15T10:00:00Z",
|
||||
"updated_at": "2024-01-15T10:30:00Z",
|
||||
"is_sandbox_running": true
|
||||
}
|
||||
]
|
||||
}"#;
|
||||
|
||||
let response: ListRunsResponse = serde_json::from_str(json).unwrap();
|
||||
|
||||
// Unknown states should deserialize to AmbientAgentTaskState::Unknown.
|
||||
assert_eq!(response.runs.len(), 2);
|
||||
assert_eq!(response.runs[0].title, "Valid Task");
|
||||
assert_eq!(response.runs[1].title, "Task with Invalid State");
|
||||
assert_eq!(response.runs[1].state, AmbientAgentTaskState::Unknown);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_list_tasks_corrupted_json_in_middle() {
|
||||
// Mix of valid and completely malformed JSON
|
||||
let json = r#"{
|
||||
"runs": [
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440000",
|
||||
"title": "First Valid Task",
|
||||
"state": "SUCCEEDED",
|
||||
"prompt": "test prompt",
|
||||
"created_at": "2024-01-15T10:00:00Z",
|
||||
"updated_at": "2024-01-15T10:30:00Z",
|
||||
"is_sandbox_running": true
|
||||
},
|
||||
{
|
||||
"task_id": 12345,
|
||||
"title": 999,
|
||||
"state": true
|
||||
},
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440002",
|
||||
"title": "Second Valid Task",
|
||||
"state": "INPROGRESS",
|
||||
"prompt": "test prompt 2",
|
||||
"created_at": "2024-01-15T11:00:00Z",
|
||||
"updated_at": "2024-01-15T11:30:00Z",
|
||||
"is_sandbox_running": false
|
||||
}
|
||||
]
|
||||
}"#;
|
||||
|
||||
let response: ListRunsResponse = serde_json::from_str(json).unwrap();
|
||||
|
||||
// Should have both valid tasks, malformed one skipped
|
||||
assert_eq!(response.runs.len(), 2);
|
||||
assert_eq!(response.runs[0].title, "First Valid Task");
|
||||
assert_eq!(response.runs[1].title, "Second Valid Task");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_list_tasks_empty_tasks_array() {
|
||||
// Empty tasks array should work fine
|
||||
let json = r#"{
|
||||
"runs": []
|
||||
}"#;
|
||||
|
||||
let response: ListRunsResponse = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(response.runs.len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_list_tasks_all_tasks_valid() {
|
||||
// Ensure we don't break the happy path
|
||||
let json = r#"{
|
||||
"runs": [
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440000",
|
||||
"title": "Task 1",
|
||||
"state": "SUCCEEDED",
|
||||
"prompt": "test prompt",
|
||||
"created_at": "2024-01-15T10:00:00Z",
|
||||
"updated_at": "2024-01-15T10:30:00Z",
|
||||
"is_sandbox_running": true
|
||||
},
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440001",
|
||||
"title": "Task 2",
|
||||
"state": "INPROGRESS",
|
||||
"prompt": "test prompt 2",
|
||||
"created_at": "2024-01-15T11:00:00Z",
|
||||
"updated_at": "2024-01-15T11:30:00Z",
|
||||
"is_sandbox_running": false
|
||||
},
|
||||
{
|
||||
"task_id": "550e8400-e29b-41d4-a716-446655440002",
|
||||
"title": "Task 3",
|
||||
"state": "FAILED",
|
||||
"prompt": "test prompt 3",
|
||||
"created_at": "2024-01-15T12:00:00Z",
|
||||
"updated_at": "2024-01-15T12:30:00Z",
|
||||
"is_sandbox_running": false
|
||||
}
|
||||
]
|
||||
}"#;
|
||||
|
||||
let response: ListRunsResponse = serde_json::from_str(json).unwrap();
|
||||
|
||||
// All tasks should be present
|
||||
assert_eq!(response.runs.len(), 3);
|
||||
assert_eq!(response.runs[0].title, "Task 1");
|
||||
assert_eq!(response.runs[1].title, "Task 2");
|
||||
assert_eq!(response.runs[2].title, "Task 3");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------------------------------------------------
|
||||
// We test roundtripping serialize and deserialize since we use this for persisting artifacts for local conversations.
|
||||
// ---------------------------------------------------------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn test_artifact_plan_serialize_deserialize_roundtrip() {
|
||||
let original = Artifact::Plan {
|
||||
document_uid: "doc-123".to_string(),
|
||||
notebook_uid: Some(NotebookId::from("notebook12345678901234".to_string())),
|
||||
title: Some("My Plan".to_string()),
|
||||
};
|
||||
|
||||
let serialized = serde_json::to_string(&original).unwrap();
|
||||
let deserialized: Artifact = serde_json::from_str(&serialized).unwrap();
|
||||
|
||||
assert_eq!(original, deserialized);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_agent_message_headers() {
|
||||
let json = r#"[
|
||||
{
|
||||
"message_id": "message-1",
|
||||
"sender_run_id": "run-1",
|
||||
"subject": "Build finished",
|
||||
"sent_at": "2026-04-09T20:00:00Z",
|
||||
"delivered_at": "2026-04-09T20:01:00Z",
|
||||
"read_at": null
|
||||
}
|
||||
]"#;
|
||||
|
||||
let headers: Vec<AgentMessageHeader> = serde_json::from_str(json).unwrap();
|
||||
|
||||
assert_eq!(headers.len(), 1);
|
||||
assert_eq!(headers[0].message_id, "message-1");
|
||||
assert_eq!(headers[0].sender_run_id, "run-1");
|
||||
assert_eq!(headers[0].subject, "Build finished");
|
||||
assert_eq!(headers[0].sent_at, "2026-04-09T20:00:00Z");
|
||||
assert_eq!(
|
||||
headers[0].delivered_at.as_deref(),
|
||||
Some("2026-04-09T20:01:00Z")
|
||||
);
|
||||
assert_eq!(headers[0].read_at, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_read_agent_message_response_with_timestamps() {
|
||||
let json = r#"{
|
||||
"message_id": "message-1",
|
||||
"sender_run_id": "run-1",
|
||||
"subject": "Build finished",
|
||||
"body": "Everything passed.",
|
||||
"sent_at": "2026-04-09T20:00:00Z",
|
||||
"delivered_at": "2026-04-09T20:01:00Z",
|
||||
"read_at": "2026-04-09T20:02:00Z"
|
||||
}"#;
|
||||
|
||||
let response: ReadAgentMessageResponse = serde_json::from_str(json).unwrap();
|
||||
|
||||
assert_eq!(response.message_id, "message-1");
|
||||
assert_eq!(response.sender_run_id, "run-1");
|
||||
assert_eq!(response.subject, "Build finished");
|
||||
assert_eq!(response.body, "Everything passed.");
|
||||
assert_eq!(response.sent_at, "2026-04-09T20:00:00Z");
|
||||
assert_eq!(
|
||||
response.delivered_at.as_deref(),
|
||||
Some("2026-04-09T20:01:00Z")
|
||||
);
|
||||
assert_eq!(response.read_at.as_deref(), Some("2026-04-09T20:02:00Z"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_agent_run_events_with_optional_fields() {
|
||||
let json = r#"[
|
||||
{
|
||||
"event_type": "run_started",
|
||||
"run_id": "run-1",
|
||||
"ref_id": null,
|
||||
"execution_id": "exec-1",
|
||||
"occurred_at": "2026-04-09T20:00:00Z",
|
||||
"sequence": 7
|
||||
},
|
||||
{
|
||||
"event_type": "new_message",
|
||||
"run_id": "run-2",
|
||||
"ref_id": "message-9",
|
||||
"execution_id": null,
|
||||
"occurred_at": "2026-04-09T20:05:00Z",
|
||||
"sequence": 8
|
||||
}
|
||||
]"#;
|
||||
|
||||
let events: Vec<AgentRunEvent> = serde_json::from_str(json).unwrap();
|
||||
|
||||
assert_eq!(events.len(), 2);
|
||||
assert_eq!(events[0].event_type, "run_started");
|
||||
assert_eq!(events[0].execution_id.as_deref(), Some("exec-1"));
|
||||
assert_eq!(events[0].ref_id, None);
|
||||
assert_eq!(events[0].sequence, 7);
|
||||
assert_eq!(events[1].event_type, "new_message");
|
||||
assert_eq!(events[1].ref_id.as_deref(), Some("message-9"));
|
||||
assert_eq!(events[1].execution_id, None);
|
||||
assert_eq!(events[1].sequence, 8);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_artifact_plan_serialize_deserialize_roundtrip_no_notebook_uid() {
|
||||
let original = Artifact::Plan {
|
||||
document_uid: "doc-123".to_string(),
|
||||
notebook_uid: None,
|
||||
title: Some("My Plan".to_string()),
|
||||
};
|
||||
|
||||
let serialized = serde_json::to_string(&original).unwrap();
|
||||
let deserialized: Artifact = serde_json::from_str(&serialized).unwrap();
|
||||
|
||||
assert_eq!(original, deserialized);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_artifact_pr_serialize_deserialize_roundtrip() {
|
||||
let original = Artifact::PullRequest {
|
||||
url: "https://github.com/org/repo/pull/42".to_string(),
|
||||
branch: "feature-branch".to_string(),
|
||||
repo: Some("repo".to_string()),
|
||||
number: Some(42),
|
||||
};
|
||||
|
||||
let serialized = serde_json::to_string(&original).unwrap();
|
||||
let deserialized: Artifact = serde_json::from_str(&serialized).unwrap();
|
||||
|
||||
// repo/number are re-derived from URL on deserialize, so should match
|
||||
assert_eq!(original, deserialized);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_artifact_file_serialize_deserialize_roundtrip() {
|
||||
let original = Artifact::File {
|
||||
artifact_uid: "artifact-file-1".to_string(),
|
||||
filepath: "outputs/report.txt".to_string(),
|
||||
filename: "report.txt".to_string(),
|
||||
mime_type: "text/plain".to_string(),
|
||||
description: Some("Daily summary".to_string()),
|
||||
size_bytes: Some(42),
|
||||
};
|
||||
|
||||
let serialized = serde_json::to_string(&original).unwrap();
|
||||
let deserialized: Artifact = serde_json::from_str(&serialized).unwrap();
|
||||
|
||||
assert_eq!(original, deserialized);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_artifact_vec_serialize_deserialize_roundtrip() {
|
||||
let original = vec![
|
||||
Artifact::Plan {
|
||||
document_uid: "doc-1".to_string(),
|
||||
notebook_uid: None,
|
||||
title: Some("Plan 1".to_string()),
|
||||
},
|
||||
Artifact::PullRequest {
|
||||
url: "https://github.com/org/repo/pull/1".to_string(),
|
||||
branch: "main".to_string(),
|
||||
repo: Some("repo".to_string()),
|
||||
number: Some(1),
|
||||
},
|
||||
Artifact::File {
|
||||
artifact_uid: "artifact-file-1".to_string(),
|
||||
filepath: "outputs/report.txt".to_string(),
|
||||
filename: "report.txt".to_string(),
|
||||
mime_type: "text/plain".to_string(),
|
||||
description: Some("Daily summary".to_string()),
|
||||
size_bytes: Some(42),
|
||||
},
|
||||
];
|
||||
|
||||
let serialized = serde_json::to_string(&original).unwrap();
|
||||
let deserialized: Vec<Artifact> = serde_json::from_str(&serialized).unwrap();
|
||||
|
||||
assert_eq!(original, deserialized);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_list_agent_runs_url_empty_filter() {
|
||||
let url = build_list_agent_runs_url(10, &TaskListFilter::default());
|
||||
assert_eq!(url, "agent/runs?limit=10");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_list_agent_runs_url_all_fields() {
|
||||
let filter = TaskListFilter {
|
||||
creator_uid: Some("user-uid".to_string()),
|
||||
updated_after: Some(Utc.with_ymd_and_hms(2026, 4, 3, 12, 30, 0).unwrap()),
|
||||
created_after: Some(Utc.with_ymd_and_hms(2026, 4, 1, 0, 0, 0).unwrap()),
|
||||
created_before: Some(Utc.with_ymd_and_hms(2026, 4, 2, 0, 0, 0).unwrap()),
|
||||
states: Some(vec![
|
||||
AmbientAgentTaskState::Failed,
|
||||
AmbientAgentTaskState::Error,
|
||||
]),
|
||||
source: Some(AgentSource::AgentWebhook),
|
||||
execution_location: Some(ExecutionLocation::Remote),
|
||||
environment_id: Some("env-123".to_string()),
|
||||
skill_spec: Some("owner/repo:SKILL.md".to_string()),
|
||||
schedule_id: Some("sched-1".to_string()),
|
||||
ancestor_run_id: Some("run-parent".to_string()),
|
||||
config_name: Some("nightly".to_string()),
|
||||
model_id: Some("claude-4-5".to_string()),
|
||||
artifact_type: Some(ArtifactType::PullRequest),
|
||||
search_query: Some("oz run".to_string()),
|
||||
sort_by: Some(RunSortBy::CreatedAt),
|
||||
sort_order: Some(RunSortOrder::Asc),
|
||||
cursor: Some("abcd==".to_string()),
|
||||
};
|
||||
|
||||
let url = build_list_agent_runs_url(42, &filter);
|
||||
assert_eq!(
|
||||
url,
|
||||
"agent/runs?limit=42\
|
||||
&creator=user-uid\
|
||||
&updated_after=2026-04-03T12%3A30%3A00%2B00%3A00\
|
||||
&created_after=2026-04-01T00%3A00%3A00%2B00%3A00\
|
||||
&created_before=2026-04-02T00%3A00%3A00%2B00%3A00\
|
||||
&state=FAILED\
|
||||
&state=ERROR\
|
||||
&source=API\
|
||||
&execution_location=REMOTE\
|
||||
&environment_id=env-123\
|
||||
&skill_spec=owner%2Frepo%3ASKILL.md\
|
||||
&schedule_id=sched-1\
|
||||
&ancestor_run_id=run-parent\
|
||||
&name=nightly\
|
||||
&model_id=claude-4-5\
|
||||
&artifact_type=PULL_REQUEST\
|
||||
&q=oz%20run\
|
||||
&sort_by=created_at\
|
||||
&sort_order=asc\
|
||||
&cursor=abcd%3D%3D"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_list_agent_runs_url_repeats_state_filter() {
|
||||
let filter = TaskListFilter {
|
||||
states: Some(vec![
|
||||
AmbientAgentTaskState::Queued,
|
||||
AmbientAgentTaskState::InProgress,
|
||||
AmbientAgentTaskState::Succeeded,
|
||||
]),
|
||||
..TaskListFilter::default()
|
||||
};
|
||||
let url = build_list_agent_runs_url(5, &filter);
|
||||
assert_eq!(
|
||||
url,
|
||||
"agent/runs?limit=5&state=QUEUED&state=INPROGRESS&state=SUCCEEDED"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_list_agent_runs_url_skips_unknown_state() {
|
||||
// The deserializer keeps `Unknown` for forward compatibility, but we shouldn't send it to
|
||||
// the server as a filter value.
|
||||
let filter = TaskListFilter {
|
||||
states: Some(vec![
|
||||
AmbientAgentTaskState::Unknown,
|
||||
AmbientAgentTaskState::Succeeded,
|
||||
]),
|
||||
..TaskListFilter::default()
|
||||
};
|
||||
let url = build_list_agent_runs_url(1, &filter);
|
||||
assert_eq!(url, "agent/runs?limit=1&state=SUCCEEDED");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_list_agent_runs_url_routes_to_runs_not_tasks() {
|
||||
let url = build_list_agent_runs_url(10, &TaskListFilter::default());
|
||||
assert!(url.starts_with("agent/runs?"));
|
||||
assert!(!url.starts_with("agent/tasks"));
|
||||
}
|
||||
@@ -0,0 +1,921 @@
|
||||
use std::{result::Result as StdResult, sync::Arc};
|
||||
|
||||
use anyhow::{anyhow, bail, Context as _, Result};
|
||||
use async_trait::async_trait;
|
||||
use cynic::{MutationBuilder, QueryBuilder};
|
||||
use firebase::{FetchAccessTokenResponse, FirebaseError};
|
||||
use futures::FutureExt;
|
||||
use instant::Duration;
|
||||
#[cfg(test)]
|
||||
use mockall::{automock, predicate::*};
|
||||
use oauth2::TokenResponse;
|
||||
use thiserror::Error;
|
||||
use warp_core::errors::{AnyhowErrorExt, ErrorExt};
|
||||
use warp_graphql::client::Operation;
|
||||
use warp_graphql::mutations::expire_api_key::{
|
||||
ExpireApiKey, ExpireApiKeyResult, ExpireApiKeyVariables,
|
||||
};
|
||||
use warp_graphql::queries::get_conversation_usage::{
|
||||
ConversationUsage, GetConversationUsage, GetConversationUsageVariables, UserResult,
|
||||
};
|
||||
|
||||
use warp_graphql::mutations::set_user_is_onboarded::{
|
||||
SetUserIsOnboarded, SetUserIsOnboardedResult, SetUserIsOnboardedVariables,
|
||||
};
|
||||
use warp_graphql::mutations::update_user_settings::{
|
||||
UpdateUserSettings, UpdateUserSettingsInput, UpdateUserSettingsResult,
|
||||
UpdateUserSettingsVariables,
|
||||
};
|
||||
use warp_graphql::mutations::{
|
||||
create_anonymous_user::{
|
||||
AnonymousUserType, CreateAnonymousUser, CreateAnonymousUserResult,
|
||||
CreateAnonymousUserVariables,
|
||||
},
|
||||
generate_api_key::{
|
||||
GenerateApiKey, GenerateApiKeyInput, GenerateApiKeyResult, GenerateApiKeyVariables,
|
||||
},
|
||||
mint_custom_token::{MintCustomTokenResult, MintCustomTokenVariables},
|
||||
};
|
||||
use warp_graphql::object_permissions::OwnerType;
|
||||
use warp_graphql::queries::api_keys::{
|
||||
ApiKeyProperties, ApiKeyPropertiesResult, ApiKeys, ApiKeysVariables,
|
||||
};
|
||||
use warp_graphql::queries::get_user::{GetUser, GetUserVariables, UserOutput as GqlUserOutput};
|
||||
use warp_graphql::queries::get_user_settings::{GetUserSettings, GetUserSettingsVariables};
|
||||
use warpui::r#async::BoxFuture;
|
||||
|
||||
use crate::auth::UserUid;
|
||||
use crate::server::graphql::{default_request_options, get_user_facing_error_message};
|
||||
use crate::server::ids::ApiKeyUid;
|
||||
use crate::server::server_api::register_error;
|
||||
use crate::server::server_api::EXPERIMENT_ID_HEADER;
|
||||
use crate::settings::PrivacySettingsSnapshot;
|
||||
use crate::{
|
||||
auth::{
|
||||
credentials::{AuthToken, Credentials, FirebaseToken, LoginToken, RefreshToken},
|
||||
user::FirebaseAuthTokens,
|
||||
user::User,
|
||||
},
|
||||
channel::ChannelState,
|
||||
convert_to_server_experiment,
|
||||
server::{
|
||||
datetime_ext::DateTimeExt as _, experiments::ServerExperiment,
|
||||
graphql::get_request_context, server_api::ServerApiEvent,
|
||||
},
|
||||
};
|
||||
|
||||
use super::ServerApi;
|
||||
|
||||
/// Error messages returned from the Firebase REST API when attempting to convert a refresh token
|
||||
/// into an access token that indicate the user's token is in an errored state.
|
||||
/// These are "soft" errors because the user likely just needs to log in again.
|
||||
/// See https://firebase.google.com/docs/reference/rest/auth#section-refresh-token.
|
||||
static FETCH_ACCESS_TOKEN_SOFT_ERROR_MESSAGES: &[&str] = &[
|
||||
"TOKEN_EXPIRED",
|
||||
"INVALID_REFRESH_TOKEN",
|
||||
"MISSING_REFRESH_TOKEN",
|
||||
];
|
||||
|
||||
/// Error messages returned from the Firebase REST API when attempting to convert a refresh token
|
||||
/// into an access token that indicate the user's account is in an errored state.
|
||||
/// These are "hard" errors because the user likely can no longer sign in with their account,
|
||||
/// for example if it were disabled or deleted.
|
||||
/// See https://firebase.google.com/docs/reference/rest/auth#section-refresh-token.
|
||||
static FETCH_ACCESS_TOKEN_HARD_ERROR_MESSAGES: &[&str] = &["USER_DISABLED", "USER_NOT_FOUND"];
|
||||
|
||||
const FETCH_ACCESS_TOKEN_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
|
||||
/// Header key for the ambient workload token attached to multi-agent requests.
|
||||
pub const AMBIENT_WORKLOAD_TOKEN_HEADER: &str = "X-Warp-Ambient-Workload-Token";
|
||||
|
||||
/// Header key for the cloud agent task ID attached to requests from ambient agents.
|
||||
pub const CLOUD_AGENT_ID_HEADER: &str = "X-Warp-Cloud-Agent-ID";
|
||||
|
||||
/// Duration for which the ambient workload token is valid (3 hours).
|
||||
const AMBIENT_WORKLOAD_TOKEN_DURATION: Duration = Duration::from_secs(3 * 60 * 60);
|
||||
|
||||
/// User settings that are currently 'synced' (e.g. stored server-side) on a per-user basis.
|
||||
#[derive(Copy, Clone, Debug, Default)]
|
||||
pub struct SyncedUserSettings {
|
||||
pub is_cloud_conversation_storage_enabled: bool,
|
||||
pub is_crash_reporting_enabled: bool,
|
||||
pub is_telemetry_enabled: bool,
|
||||
}
|
||||
|
||||
/// Results of an attempt to fetch the current user.
|
||||
pub struct FetchUserResult {
|
||||
pub user: User,
|
||||
/// The credentials used to authenticate this user.
|
||||
pub credentials: Credentials,
|
||||
pub server_experiments: Vec<ServerExperiment>,
|
||||
/// Whether this attempt to fetch the user was for refreshing an existing logged-in user.
|
||||
pub from_refresh: bool,
|
||||
/// LLM model choices for this user.
|
||||
pub llms: crate::ai::llms::ModelsByFeature,
|
||||
}
|
||||
|
||||
#[cfg_attr(test, automock)]
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
pub trait AuthClient: 'static + Send + Sync {
|
||||
/// Creates an anonymous user, who is allowed to use Warp but may lack the ability
|
||||
/// to interact with particular features.
|
||||
async fn create_anonymous_user(
|
||||
&self,
|
||||
referral_code: Option<String>,
|
||||
anonymous_user_type: AnonymousUserType,
|
||||
) -> Result<CreateAnonymousUserResult>;
|
||||
|
||||
/// Returns the cached access token, if it is still valid. If it has expired, fetches a new
|
||||
/// access token using the user's refresh token, caches it, and the returns it.
|
||||
/// Returns an auth mode that may not require an Authorization header (e.g. session cookies or
|
||||
/// test credentials).
|
||||
async fn get_or_refresh_access_token(&self) -> Result<AuthToken>;
|
||||
|
||||
/// Fetches data required to construct the [`User`] object. This includes the user's metadata
|
||||
/// and authentication tokens.
|
||||
async fn fetch_user(
|
||||
&self,
|
||||
token: LoginToken,
|
||||
for_refresh: bool,
|
||||
) -> StdResult<FetchUserResult, UserAuthenticationError>;
|
||||
|
||||
/// Creates and fetches an new custom token for the current user from Firebase.
|
||||
/// This only works for anonymous users, and will surface an error if the user is not anonymous.
|
||||
async fn fetch_new_custom_token(&self) -> Result<MintCustomTokenResult>;
|
||||
|
||||
/// Handles the response from [`Self::fetch_new_custom_token`], returning the newly-minted custom token.
|
||||
fn on_custom_token_fetched(
|
||||
&self,
|
||||
response: Result<MintCustomTokenResult>,
|
||||
) -> Result<String, MintCustomTokenError>;
|
||||
|
||||
/// Queries warp-server for a set of the currently logged-in user's fields.
|
||||
async fn fetch_user_properties<'a>(&self, auth_token: Option<&'a str>)
|
||||
-> Result<GqlUserOutput>;
|
||||
|
||||
/// Upon success, returns an `Option` containing the user's settings retrieved from the server,
|
||||
/// if any. The user may not have server-side settings if they onboarded prior to the launch
|
||||
/// of telemetry opt-out, have not logged in since the launch, and have never changed defaults
|
||||
/// for any of the settings in [`SyncedUserSettings`]. If the fetched settings object exists
|
||||
/// but is missing required fields, or if the request itself failed, returns an error.
|
||||
async fn get_user_settings(&self) -> Result<Option<SyncedUserSettings>>;
|
||||
|
||||
/// Returns conversation usage history for the current user over the past n days.
|
||||
/// If last_updated_end_timestamp is provided, only conversations with
|
||||
/// lastUpdated earlier than this timestamp are returned.
|
||||
async fn get_conversation_usage_history(
|
||||
&self,
|
||||
days: Option<i32>,
|
||||
limit: Option<i32>,
|
||||
last_updated_end_timestamp: Option<warp_graphql::scalars::Time>,
|
||||
) -> Result<Vec<ConversationUsage>>;
|
||||
|
||||
async fn set_is_telemetry_enabled(&self, value: bool) -> Result<()>;
|
||||
|
||||
async fn set_is_crash_reporting_enabled(&self, value: bool) -> Result<()>;
|
||||
|
||||
async fn set_is_cloud_conversation_storage_enabled(&self, value: bool) -> Result<()>;
|
||||
|
||||
/// Sends a request to update the user's settings on the server with values contained in the
|
||||
/// given `settings_snapshot`.
|
||||
async fn update_user_settings(&self, settings_snapshot: PrivacySettingsSnapshot) -> Result<()>;
|
||||
|
||||
async fn set_user_is_onboarded(&self) -> Result<bool>;
|
||||
|
||||
/// Requests a device authorization code from the server. This is only used for headless CLI/SDK authentication.
|
||||
async fn request_device_code(
|
||||
&self,
|
||||
) -> StdResult<oauth2::StandardDeviceAuthorizationResponse, UserAuthenticationError>;
|
||||
|
||||
/// Wait for the request to be approved or rejected and exchange it for a short-lived custom access token.
|
||||
async fn exchange_device_access_token(
|
||||
&self,
|
||||
details: &oauth2::StandardDeviceAuthorizationResponse,
|
||||
timeout: Duration,
|
||||
) -> StdResult<FirebaseToken, UserAuthenticationError>;
|
||||
// API Keys
|
||||
async fn list_api_keys(&self) -> Result<Vec<ApiKeyProperties>>;
|
||||
|
||||
async fn create_api_key(
|
||||
&self,
|
||||
name: String,
|
||||
team_id: Option<cynic::Id>,
|
||||
expires_at: Option<warp_graphql::scalars::Time>,
|
||||
) -> Result<GenerateApiKeyResult>;
|
||||
|
||||
async fn expire_api_key(&self, key_uid: &ApiKeyUid) -> Result<ExpireApiKeyResult>;
|
||||
|
||||
/// Returns a cached ambient workload token, or issues a new one if not present or expired.
|
||||
///
|
||||
/// Returns `Ok(None)` if not running in an isolation platform (e.g., Namespace) or on WASM.
|
||||
async fn get_or_create_ambient_workload_token(&self) -> Result<Option<String>>;
|
||||
}
|
||||
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
impl AuthClient for ServerApi {
|
||||
async fn create_anonymous_user(
|
||||
&self,
|
||||
referral_code: Option<String>,
|
||||
anonymous_user_type: AnonymousUserType,
|
||||
) -> Result<CreateAnonymousUserResult> {
|
||||
let variables = CreateAnonymousUserVariables {
|
||||
input: warp_graphql::mutations::create_anonymous_user::CreateAnonymousUserInput {
|
||||
anonymous_user_type,
|
||||
expiration_type: warp_graphql::mutations::create_anonymous_user::AnonymousUserExpirationType::NoExpiration,
|
||||
referral_code,
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = CreateAnonymousUser::build(variables);
|
||||
let response = operation
|
||||
.send_request(self.client.clone(), default_request_options())
|
||||
.await?;
|
||||
|
||||
Ok(response
|
||||
.data
|
||||
.ok_or_else(|| anyhow!("missing data in response"))?
|
||||
.create_anonymous_user)
|
||||
}
|
||||
|
||||
async fn get_or_refresh_access_token(&self) -> Result<AuthToken> {
|
||||
if cfg!(feature = "skip_login") {
|
||||
bail!("skip_login enabled; failing all authenticated requests");
|
||||
}
|
||||
|
||||
let Some(credentials) = self.auth_state.credentials() else {
|
||||
bail!("Attempted to retrieve access token when user is logged out");
|
||||
};
|
||||
|
||||
match credentials {
|
||||
Credentials::ApiKey { key, .. } => Ok(AuthToken::ApiKey(key)),
|
||||
Credentials::Firebase(auth_tokens) => {
|
||||
let expiration_time = auth_tokens.expiration_time;
|
||||
|
||||
// Generate a new ID token if the token has expired or will expire in the
|
||||
// next five minutes. This matches the behavior of the Firebase Auth SDK.
|
||||
if chrono::DateTime::now() + chrono::Duration::minutes(5) >= expiration_time {
|
||||
let refresh_token = auth_tokens.refresh_token.clone();
|
||||
let firebase_token = FirebaseToken::Refresh(RefreshToken::new(refresh_token));
|
||||
|
||||
let result = fetch_auth_tokens(self.client.clone(), firebase_token).await;
|
||||
|
||||
if let Err(UserAuthenticationError::DeniedAccessToken(_)) = result {
|
||||
let _ = self.event_sender.send(ServerApiEvent::NeedsReauth).await;
|
||||
}
|
||||
let new_firebase_token_info = result?;
|
||||
self.auth_state
|
||||
.update_firebase_tokens(new_firebase_token_info.clone());
|
||||
return Ok(AuthToken::Firebase(new_firebase_token_info.id_token));
|
||||
}
|
||||
|
||||
Ok(AuthToken::Firebase(auth_tokens.id_token))
|
||||
}
|
||||
Credentials::SessionCookie => Ok(AuthToken::NoAuth),
|
||||
#[cfg(any(test, feature = "integration_tests", feature = "skip_login"))]
|
||||
Credentials::Test => Ok(AuthToken::NoAuth),
|
||||
}
|
||||
}
|
||||
|
||||
async fn fetch_user(
|
||||
&self,
|
||||
token: LoginToken,
|
||||
for_refresh: bool,
|
||||
) -> StdResult<FetchUserResult, UserAuthenticationError> {
|
||||
let new_credentials = exchange_credentials(self.client.clone(), token).await?;
|
||||
let auth_token = new_credentials.bearer_token();
|
||||
let user_output = self
|
||||
.fetch_user_properties(auth_token.as_bearer_token())
|
||||
.await
|
||||
.context("Failed to fetch user response data")
|
||||
.map_err(UserAuthenticationError::Unexpected)?;
|
||||
|
||||
let UserProperties {
|
||||
user,
|
||||
server_experiments,
|
||||
llms,
|
||||
api_key_owner_type,
|
||||
} = user_output.into();
|
||||
|
||||
// Store the owner type if using an API key.
|
||||
let new_credentials = match new_credentials {
|
||||
Credentials::ApiKey { key, .. } => Credentials::ApiKey {
|
||||
key,
|
||||
owner_type: api_key_owner_type,
|
||||
},
|
||||
other => other,
|
||||
};
|
||||
|
||||
Ok(FetchUserResult {
|
||||
user,
|
||||
credentials: new_credentials,
|
||||
server_experiments,
|
||||
from_refresh: for_refresh,
|
||||
llms,
|
||||
})
|
||||
}
|
||||
|
||||
async fn fetch_new_custom_token(&self) -> Result<MintCustomTokenResult> {
|
||||
let variables = MintCustomTokenVariables {
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation =
|
||||
warp_graphql::mutations::mint_custom_token::MintCustomToken::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
Ok(response.mint_custom_token)
|
||||
}
|
||||
|
||||
fn on_custom_token_fetched(
|
||||
&self,
|
||||
response: Result<MintCustomTokenResult>,
|
||||
) -> Result<String, MintCustomTokenError> {
|
||||
match response {
|
||||
Ok(response_data) => match response_data {
|
||||
MintCustomTokenResult::MintCustomTokenOutput(output) => Ok(output.custom_token),
|
||||
MintCustomTokenResult::UserFacingError(user_facing_error) => {
|
||||
Err(MintCustomTokenError::UserFacingError(
|
||||
get_user_facing_error_message(user_facing_error),
|
||||
))
|
||||
}
|
||||
MintCustomTokenResult::Unknown => Err(MintCustomTokenError::Unknown),
|
||||
},
|
||||
Err(_) => Err(MintCustomTokenError::Unknown),
|
||||
}
|
||||
}
|
||||
|
||||
async fn fetch_user_properties<'a>(
|
||||
&self,
|
||||
auth_token: Option<&'a str>,
|
||||
) -> Result<GqlUserOutput> {
|
||||
let variables = GetUserVariables {
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = GetUser::build(variables);
|
||||
let response = operation
|
||||
.send_request(
|
||||
self.client.clone(),
|
||||
warp_graphql::client::RequestOptions {
|
||||
auth_token: auth_token.map(ToOwned::to_owned),
|
||||
headers: std::collections::HashMap::from([(
|
||||
EXPERIMENT_ID_HEADER.to_string(),
|
||||
self.auth_state.anonymous_id(),
|
||||
)]),
|
||||
..default_request_options()
|
||||
},
|
||||
)
|
||||
.await?
|
||||
.data
|
||||
.ok_or_else(|| anyhow!("Expected valid response.data"))?;
|
||||
|
||||
match response.user {
|
||||
warp_graphql::queries::get_user::UserResult::UserOutput(user_output) => Ok(user_output),
|
||||
warp_graphql::queries::get_user::UserResult::Unknown => {
|
||||
Err(anyhow!("Unable to fetch user"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_user_settings(&self) -> Result<Option<SyncedUserSettings>> {
|
||||
let variables = GetUserSettingsVariables {
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = GetUserSettings::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.user {
|
||||
warp_graphql::queries::get_user_settings::UserResult::UserOutput(user_output) => {
|
||||
match user_output.user.settings {
|
||||
Some(user_settings) => Ok(Some(SyncedUserSettings {
|
||||
is_cloud_conversation_storage_enabled: user_settings
|
||||
.is_cloud_conversation_storage_enabled,
|
||||
is_crash_reporting_enabled: user_settings.is_crash_reporting_enabled,
|
||||
is_telemetry_enabled: user_settings.is_telemetry_enabled,
|
||||
})),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
warp_graphql::queries::get_user_settings::UserResult::Unknown => {
|
||||
Err(anyhow!("Unable to fetch user settings"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Returns a history of the current user's conversation usage over the past n days.
|
||||
async fn get_conversation_usage_history(
|
||||
&self,
|
||||
days: Option<i32>,
|
||||
limit: Option<i32>,
|
||||
last_updated_end_timestamp: Option<warp_graphql::scalars::Time>,
|
||||
) -> Result<Vec<ConversationUsage>> {
|
||||
let operation = GetConversationUsage::build(GetConversationUsageVariables {
|
||||
request_context: get_request_context(),
|
||||
days,
|
||||
limit,
|
||||
last_updated_end_timestamp,
|
||||
});
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
match response.user {
|
||||
UserResult::UserOutput(out) => Ok(out.user.conversation_usage),
|
||||
UserResult::Unknown => Err(anyhow!("Unable to fetch conversation usage")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn set_is_telemetry_enabled(&self, value: bool) -> Result<()> {
|
||||
let variables = UpdateUserSettingsVariables {
|
||||
input: UpdateUserSettingsInput {
|
||||
telemetry_enabled: Some(value),
|
||||
..Default::default()
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = UpdateUserSettings::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.update_user_settings;
|
||||
|
||||
match result {
|
||||
UpdateUserSettingsResult::UpdateUserSettingsOutput(_) => Ok(()),
|
||||
UpdateUserSettingsResult::UserFacingError(user_facing_error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
UpdateUserSettingsResult::Unknown => Err(anyhow!("failed to set telemetry enabled")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn set_is_crash_reporting_enabled(&self, value: bool) -> Result<()> {
|
||||
let variables = UpdateUserSettingsVariables {
|
||||
input: UpdateUserSettingsInput {
|
||||
crash_reporting_enabled: Some(value),
|
||||
..Default::default()
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = UpdateUserSettings::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.update_user_settings;
|
||||
|
||||
match result {
|
||||
UpdateUserSettingsResult::UpdateUserSettingsOutput(_) => Ok(()),
|
||||
UpdateUserSettingsResult::UserFacingError(user_facing_error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
UpdateUserSettingsResult::Unknown => {
|
||||
Err(anyhow!("failed to set crash reporting enabled"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn set_is_cloud_conversation_storage_enabled(&self, value: bool) -> Result<()> {
|
||||
let variables = UpdateUserSettingsVariables {
|
||||
input: UpdateUserSettingsInput {
|
||||
cloud_conversation_storage_enabled: Some(value),
|
||||
..Default::default()
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = UpdateUserSettings::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.update_user_settings;
|
||||
|
||||
match result {
|
||||
UpdateUserSettingsResult::UpdateUserSettingsOutput(_) => Ok(()),
|
||||
UpdateUserSettingsResult::UserFacingError(user_facing_error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
UpdateUserSettingsResult::Unknown => {
|
||||
Err(anyhow!("failed to set cloud conversation storage enabled"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn update_user_settings(&self, settings_snapshot: PrivacySettingsSnapshot) -> Result<()> {
|
||||
let variables = UpdateUserSettingsVariables {
|
||||
input: UpdateUserSettingsInput {
|
||||
telemetry_enabled: Some(settings_snapshot.is_telemetry_enabled()),
|
||||
crash_reporting_enabled: Some(settings_snapshot.is_crash_reporting_enabled()),
|
||||
cloud_conversation_storage_enabled: settings_snapshot
|
||||
.cloud_conversation_storage_enabled(),
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = UpdateUserSettings::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.update_user_settings;
|
||||
|
||||
match result {
|
||||
UpdateUserSettingsResult::UpdateUserSettingsOutput(_) => Ok(()),
|
||||
UpdateUserSettingsResult::UserFacingError(user_facing_error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
UpdateUserSettingsResult::Unknown => Err(anyhow!("failed to update user settings")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn set_user_is_onboarded(&self) -> Result<bool> {
|
||||
let variables = SetUserIsOnboardedVariables {
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = SetUserIsOnboarded::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.set_user_is_onboarded;
|
||||
|
||||
match result {
|
||||
SetUserIsOnboardedResult::SetUserIsOnboardedOutput(_) => Ok(true),
|
||||
SetUserIsOnboardedResult::UserFacingError(user_facing_error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
SetUserIsOnboardedResult::Unknown => Err(anyhow!("failed to set user is onboarded")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn request_device_code(
|
||||
&self,
|
||||
) -> StdResult<oauth2::StandardDeviceAuthorizationResponse, UserAuthenticationError> {
|
||||
self.oauth_client
|
||||
.exchange_device_code()
|
||||
.request_async(self.client.as_ref())
|
||||
.await
|
||||
.context("Failed to generate device code")
|
||||
.map_err(UserAuthenticationError::Unexpected)
|
||||
}
|
||||
|
||||
async fn exchange_device_access_token(
|
||||
&self,
|
||||
details: &oauth2::StandardDeviceAuthorizationResponse,
|
||||
timeout: Duration,
|
||||
) -> StdResult<FirebaseToken, UserAuthenticationError> {
|
||||
let result = self
|
||||
.oauth_client
|
||||
.exchange_device_access_token(details)
|
||||
.request_async(
|
||||
self.client.as_ref(),
|
||||
|delay| warpui::r#async::Timer::after(delay).map(|_| ()),
|
||||
Some(timeout),
|
||||
)
|
||||
.await
|
||||
.context("Unable to obtain access token")
|
||||
.map_err(UserAuthenticationError::Unexpected)?;
|
||||
|
||||
// Firebase doesn't directly support the device flow. Instead, the server mints a short-lived
|
||||
// custom access token, which we can then exchange for a refresh token.
|
||||
Ok(FirebaseToken::Custom(
|
||||
result.access_token().secret().to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
// API Keys
|
||||
async fn list_api_keys(&self) -> Result<Vec<ApiKeyProperties>> {
|
||||
let variables = ApiKeysVariables {
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = ApiKeys::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
match response.api_keys {
|
||||
ApiKeyPropertiesResult::ApiKeyPropertiesOutput(output) => Ok(output.api_keys),
|
||||
ApiKeyPropertiesResult::UserFacingError(e) => {
|
||||
Err(anyhow!(get_user_facing_error_message(e)))
|
||||
}
|
||||
ApiKeyPropertiesResult::Unknown => Err(anyhow!("failed to fetch API keys")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn create_api_key(
|
||||
&self,
|
||||
name: String,
|
||||
team_id: Option<cynic::Id>,
|
||||
expires_at: Option<warp_graphql::scalars::Time>,
|
||||
) -> Result<GenerateApiKeyResult> {
|
||||
let variables = GenerateApiKeyVariables {
|
||||
input: GenerateApiKeyInput {
|
||||
name,
|
||||
team_id,
|
||||
expires_at,
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = GenerateApiKey::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
Ok(response.generate_api_key)
|
||||
}
|
||||
async fn expire_api_key(&self, key_uid: &ApiKeyUid) -> Result<ExpireApiKeyResult> {
|
||||
let variables = ExpireApiKeyVariables {
|
||||
key_uid: key_uid.into(),
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let op = ExpireApiKey::build(variables);
|
||||
let res = self.send_graphql_request(op, None).await?;
|
||||
Ok(res.expire_api_key)
|
||||
}
|
||||
|
||||
async fn get_or_create_ambient_workload_token(&self) -> Result<Option<String>> {
|
||||
if cfg!(target_family = "wasm") {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// Check if we have a cached token that's still valid (with 5 minute buffer).
|
||||
// Tokens without an expiration time are always considered valid.
|
||||
{
|
||||
let cached = self.ambient_workload_token.lock();
|
||||
if let Some(ref token) = *cached {
|
||||
let is_valid = token.expires_at.is_none_or(|expires_at| {
|
||||
chrono::Utc::now() + chrono::Duration::minutes(5) < expires_at
|
||||
});
|
||||
if is_valid {
|
||||
return Ok(Some(token.token.clone()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Issue a new token.
|
||||
let workload_token = match warp_isolation_platform::issue_workload_token(Some(
|
||||
AMBIENT_WORKLOAD_TOKEN_DURATION,
|
||||
))
|
||||
.await
|
||||
{
|
||||
Ok(token) => token,
|
||||
Err(warp_isolation_platform::IsolationPlatformError::NoIsolationPlatformDetected) => {
|
||||
return Ok(None);
|
||||
}
|
||||
Err(e) => return Err(e.into()),
|
||||
};
|
||||
|
||||
let token_str = workload_token.token.clone();
|
||||
|
||||
{
|
||||
let mut cached = self.ambient_workload_token.lock();
|
||||
*cached = Some(workload_token);
|
||||
}
|
||||
|
||||
Ok(Some(token_str))
|
||||
}
|
||||
}
|
||||
|
||||
/// Exchange a long-lived token for fresh [`Credentials`].
|
||||
async fn exchange_credentials(
|
||||
client: Arc<http_client::Client>,
|
||||
token: LoginToken,
|
||||
) -> StdResult<Credentials, UserAuthenticationError> {
|
||||
match token {
|
||||
LoginToken::Firebase(firebase_token) => {
|
||||
let tokens = fetch_auth_tokens(client, firebase_token).await?;
|
||||
Ok(Credentials::Firebase(tokens))
|
||||
}
|
||||
LoginToken::ApiKey(key) => Ok(Credentials::ApiKey {
|
||||
key,
|
||||
owner_type: None,
|
||||
}),
|
||||
LoginToken::SessionCookie => Ok(Credentials::SessionCookie),
|
||||
}
|
||||
}
|
||||
|
||||
fn fetch_auth_tokens(
|
||||
client: Arc<http_client::Client>,
|
||||
token: FirebaseToken,
|
||||
) -> BoxFuture<'static, StdResult<FirebaseAuthTokens, UserAuthenticationError>> {
|
||||
Box::pin(async move {
|
||||
let firebase_api_key = ChannelState::firebase_api_key();
|
||||
let url = token.access_token_url(&firebase_api_key);
|
||||
let request_body = token.access_token_request_body();
|
||||
let proxy_url = token.proxy_url(&ChannelState::server_root_url(), &firebase_api_key);
|
||||
let response = match client
|
||||
.post(&url)
|
||||
.form(&request_body)
|
||||
.timeout(FETCH_ACCESS_TOKEN_TIMEOUT)
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(response) => match response.error_for_status_ref() {
|
||||
Ok(_) => Ok(response),
|
||||
Err(error) => {
|
||||
log::warn!(
|
||||
"Request to firebase to fetch access token completed, but was unsuccessful: {error:?}"
|
||||
);
|
||||
|
||||
fetch_access_token_via_proxy(client, &request_body, proxy_url).await
|
||||
}
|
||||
},
|
||||
Err(error) => {
|
||||
log::warn!("Failed to make response to firebase to fetch access token: {error:?}");
|
||||
|
||||
fetch_access_token_via_proxy(client, &request_body, proxy_url).await
|
||||
}
|
||||
}?;
|
||||
|
||||
let response = response
|
||||
.json::<FetchAccessTokenResponse>()
|
||||
.await
|
||||
.map_err(anyhow::Error::from)?;
|
||||
match response {
|
||||
FetchAccessTokenResponse::Success {
|
||||
id_token,
|
||||
expires_in,
|
||||
refresh_token,
|
||||
} => Ok(FirebaseAuthTokens::from_response(
|
||||
id_token,
|
||||
refresh_token,
|
||||
expires_in,
|
||||
)?),
|
||||
FetchAccessTokenResponse::Error { error } => Err(error.into()),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn fetch_access_token_via_proxy<'a>(
|
||||
client: Arc<http_client::Client>,
|
||||
request_body: &'a [(&'a str, &'a str)],
|
||||
proxy_url: String,
|
||||
) -> BoxFuture<'a, Result<http_client::Response>> {
|
||||
Box::pin(async move {
|
||||
client
|
||||
.post(&proxy_url)
|
||||
.form(request_body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(anyhow::Error::from)
|
||||
})
|
||||
}
|
||||
|
||||
/// The [`oauth2::Client`] type, specialized to the endpoints that we require.
|
||||
pub type OAuth2Client = oauth2::basic::BasicClient<
|
||||
oauth2::EndpointNotSet, // HasAuthUrl
|
||||
oauth2::EndpointSet, // HasDeviceAuthUrl
|
||||
oauth2::EndpointNotSet, // HasIntrospectionUrl
|
||||
oauth2::EndpointNotSet, // HasRevocationUrl
|
||||
oauth2::EndpointSet, // HasTokenUrl
|
||||
>;
|
||||
|
||||
/// Intermediate type produced by converting a [`GqlUserOutput`] from the server.
|
||||
struct UserProperties {
|
||||
user: User,
|
||||
server_experiments: Vec<ServerExperiment>,
|
||||
llms: crate::ai::llms::ModelsByFeature,
|
||||
api_key_owner_type: Option<OwnerType>,
|
||||
}
|
||||
|
||||
impl From<GqlUserOutput> for UserProperties {
|
||||
fn from(user_output: GqlUserOutput) -> Self {
|
||||
let principal_type = user_output
|
||||
.principal_type
|
||||
.map(|pt| pt.into())
|
||||
.unwrap_or_default();
|
||||
let user_properties = user_output.user;
|
||||
|
||||
let is_on_work_domain = user_properties.is_on_work_domain;
|
||||
let is_onboarded = user_properties.is_onboarded;
|
||||
let api_key_owner_type = user_output.api_key_owner_type;
|
||||
|
||||
let linked_at = user_properties
|
||||
.anonymous_user_info
|
||||
.as_ref()
|
||||
.and_then(|info| info.linked_at);
|
||||
|
||||
let anonymous_user_type = user_properties
|
||||
.anonymous_user_info
|
||||
.as_ref()
|
||||
.map(|info| info.anonymous_user_type.clone());
|
||||
let personal_object_limits = user_properties
|
||||
.anonymous_user_info
|
||||
.and_then(|info| info.personal_object_limits.clone());
|
||||
let user_profile = user_properties.profile;
|
||||
let local_id = UserUid::new(user_profile.uid.as_str());
|
||||
let needs_sso_link = user_profile.needs_sso_link;
|
||||
|
||||
let server_experiments: Vec<ServerExperiment> = user_properties
|
||||
.experiments
|
||||
.and_then(|experiments| convert_to_server_experiment!(experiments))
|
||||
.unwrap_or_default();
|
||||
|
||||
// Convert LLM model choices from GraphQL response
|
||||
let llms = user_properties.llms.try_into().unwrap_or_default();
|
||||
|
||||
let user = User {
|
||||
is_onboarded,
|
||||
local_id,
|
||||
metadata: user_profile.into(),
|
||||
needs_sso_link,
|
||||
anonymous_user_type: anonymous_user_type.and_then(|t| t.try_into().ok()),
|
||||
is_on_work_domain,
|
||||
linked_at,
|
||||
personal_object_limits: personal_object_limits.and_then(|t| t.try_into().ok()),
|
||||
principal_type,
|
||||
};
|
||||
|
||||
UserProperties {
|
||||
user,
|
||||
server_experiments,
|
||||
llms,
|
||||
api_key_owner_type,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
/// Error type when retrieving a user and validating it against Firebase.
|
||||
pub enum UserAuthenticationError {
|
||||
/// The user's refresh token is invalid. This could occur if the user authed through
|
||||
/// e.g. Google/GitHub and changed their password.
|
||||
#[error("Firebase returned a token error when fetching an ID token")]
|
||||
DeniedAccessToken(FirebaseError),
|
||||
/// The user's account is invalid. This could occur if the user requested their account
|
||||
/// be deleted per their GDPR/CCPA rights.
|
||||
#[error("Firebase returned a user error when fetching an ID token")]
|
||||
UserAccountDisabled(FirebaseError),
|
||||
#[error("Invalid state parameter in auth redirect")]
|
||||
InvalidStateParameter,
|
||||
#[error("Missing state parameter in auth redirect")]
|
||||
MissingStateParameter,
|
||||
#[error("unexpected error occurred when fetching an ID token: {0:#}")]
|
||||
Unexpected(#[from] anyhow::Error),
|
||||
}
|
||||
|
||||
impl ErrorExt for UserAuthenticationError {
|
||||
fn is_actionable(&self) -> bool {
|
||||
match self {
|
||||
UserAuthenticationError::DeniedAccessToken(err) => {
|
||||
// If a request to our server failed because the user's refresh token
|
||||
// has expired, they should re-auth, but there's no value in reporting
|
||||
// this back to us.
|
||||
log::info!("ignoring denied access token error: {err:#}");
|
||||
false
|
||||
}
|
||||
UserAuthenticationError::UserAccountDisabled(err) => {
|
||||
// Similarly, if their account is disabled, they can't make requests.
|
||||
log::info!("ignoring user account disabled error: {err:#}");
|
||||
false
|
||||
}
|
||||
UserAuthenticationError::Unexpected(err) => err.is_actionable(),
|
||||
UserAuthenticationError::InvalidStateParameter
|
||||
| UserAuthenticationError::MissingStateParameter => {
|
||||
// For now, we're marking these as actionable, since a surplus of these errors
|
||||
// could mean that something is wrong in our login flow (e.g. we're not properly
|
||||
// passing the `state` variable back to the desktop client).
|
||||
// But in general, someone attempting to trick another into logging into their
|
||||
// account with a spoofed `state` variable is not actionable.
|
||||
true
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
register_error!(UserAuthenticationError);
|
||||
|
||||
impl From<FirebaseError> for UserAuthenticationError {
|
||||
fn from(error: FirebaseError) -> Self {
|
||||
if FETCH_ACCESS_TOKEN_SOFT_ERROR_MESSAGES.contains(&error.message.as_str()) {
|
||||
UserAuthenticationError::DeniedAccessToken(error)
|
||||
} else if FETCH_ACCESS_TOKEN_HARD_ERROR_MESSAGES.contains(&error.message.as_str()) {
|
||||
UserAuthenticationError::UserAccountDisabled(error)
|
||||
} else {
|
||||
UserAuthenticationError::Unexpected(
|
||||
anyhow::Error::from(error)
|
||||
.context("Failed to exchange refresh token with access token."),
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
/// Error type when creating anonymous users
|
||||
pub enum AnonymousUserCreationError {
|
||||
#[error("The network request to create the anonymous user failed")]
|
||||
CreationFailed,
|
||||
|
||||
#[error("Received a user facing error: {0}")]
|
||||
UserFacingError(String),
|
||||
|
||||
/// Failure that occurs after the user is created, but the ID token could not be fetched.
|
||||
#[error("The user was created, but the ID token could not be fetched")]
|
||||
UserAuthenticationFailed(#[from] UserAuthenticationError),
|
||||
|
||||
#[error("Failed to create anonymous user with unknown error")]
|
||||
Unknown,
|
||||
}
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
/// Error type when minting a new custom token for an anonymous user
|
||||
pub enum MintCustomTokenError {
|
||||
#[error("Received a user facing error: {0}")]
|
||||
UserFacingError(String),
|
||||
#[error("Failed to create new custom token with unknown error")]
|
||||
Unknown,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "auth_test.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,36 @@
|
||||
use crate::auth::credentials::{FirebaseToken, RefreshToken};
|
||||
use anyhow::Result;
|
||||
|
||||
#[test]
|
||||
fn test_firebase_token_urls() -> Result<()> {
|
||||
let custom_token = FirebaseToken::Custom("ct".to_string());
|
||||
let refresh_token = FirebaseToken::Refresh(RefreshToken::new("rt".to_string()));
|
||||
|
||||
assert_eq!(
|
||||
custom_token.access_token_url("api_key"),
|
||||
"https://identitytoolkit.googleapis.com/v1/accounts:signInWithCustomToken?key=api_key"
|
||||
);
|
||||
assert_eq!(
|
||||
refresh_token.access_token_url("api_key"),
|
||||
"https://securetoken.googleapis.com/v1/token?key=api_key"
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
custom_token.access_token_request_body(),
|
||||
vec![("returnSecureToken", "true"), ("token", "ct")]
|
||||
);
|
||||
assert_eq!(
|
||||
refresh_token.access_token_request_body(),
|
||||
vec![("grant_type", "refresh_token"), ("refresh_token", "rt")],
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
custom_token.proxy_url("https://staging.warp.dev", "api_key"),
|
||||
"https://staging.warp.dev/proxy/customToken?key=api_key"
|
||||
);
|
||||
assert_eq!(
|
||||
refresh_token.proxy_url("https://staging.warp.dev", "api_key"),
|
||||
"https://staging.warp.dev/proxy/token?key=api_key"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,192 @@
|
||||
use super::auth::AuthClient;
|
||||
use super::ServerApi;
|
||||
use crate::ai::generate_block_title::api::{GenerateBlockTitleRequest, GenerateBlockTitleResponse};
|
||||
use crate::server::{
|
||||
block::{Block, DisplaySetting},
|
||||
graphql::{get_request_context, get_user_facing_error_message},
|
||||
};
|
||||
use anyhow::anyhow;
|
||||
use async_trait::async_trait;
|
||||
use chrono::Utc;
|
||||
use cynic::{MutationBuilder, QueryBuilder};
|
||||
#[cfg(test)]
|
||||
use mockall::automock;
|
||||
use std::convert::TryFrom;
|
||||
use warp_core::channel::{Channel, ChannelState};
|
||||
use warp_graphql::{
|
||||
mutations::{
|
||||
share_block::{BlockInput, ShareBlock, ShareBlockResult, ShareBlockVariables},
|
||||
unshare_block::{
|
||||
UnshareBlock, UnshareBlockInput, UnshareBlockResult, UnshareBlockVariables,
|
||||
},
|
||||
},
|
||||
queries::get_blocks_for_user::{
|
||||
Block as GqlBlock, GetBlocksForUser, GetBlocksForUserVariables,
|
||||
},
|
||||
};
|
||||
|
||||
#[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<String>,
|
||||
show_prompt: bool,
|
||||
display_setting: DisplaySetting,
|
||||
) -> Result<String, anyhow::Error>;
|
||||
|
||||
async fn blocks_owned_by_user(&self) -> Result<Vec<Block>, anyhow::Error>;
|
||||
|
||||
async fn generate_shared_block_title(
|
||||
&self,
|
||||
request: GenerateBlockTitleRequest,
|
||||
) -> Result<GenerateBlockTitleResponse, anyhow::Error>;
|
||||
}
|
||||
|
||||
#[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<String>,
|
||||
show_prompt: bool,
|
||||
display_setting: DisplaySetting,
|
||||
) -> Result<String, anyhow::Error> {
|
||||
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<Vec<Block>, 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<GenerateBlockTitleResponse, anyhow::Error> {
|
||||
let auth_token = self.get_or_refresh_access_token().await?;
|
||||
let request_builder = self.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<GqlBlock> for Block {
|
||||
type Error = anyhow::Error;
|
||||
|
||||
fn try_from(value: GqlBlock) -> Result<Self, Self::Error> {
|
||||
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")),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,293 @@
|
||||
// We don't directly run agent harnesses on WASM, so this code is unused.
|
||||
#![cfg_attr(target_family = "wasm", expect(dead_code))]
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use async_trait::async_trait;
|
||||
#[cfg(test)]
|
||||
use mockall::automock;
|
||||
|
||||
use super::ServerApi;
|
||||
use crate::ai::agent::conversation::AIConversationId;
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use crate::ai::agent_sdk::retry::with_bounded_retry;
|
||||
use crate::ai::artifacts::Artifact;
|
||||
|
||||
/// A presigned upload target returned by the server.
|
||||
#[derive(Debug, Clone, serde::Deserialize)]
|
||||
pub struct UploadTarget {
|
||||
pub url: String,
|
||||
pub method: String,
|
||||
pub headers: HashMap<String, String>,
|
||||
}
|
||||
|
||||
/// Request body for upload-snapshot upload targets.
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct SnapshotUploadRequest {
|
||||
pub files: Vec<SnapshotFileInfo>,
|
||||
}
|
||||
|
||||
/// Describes a single file in a snapshot upload request.
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct SnapshotFileInfo {
|
||||
pub filename: String,
|
||||
pub mime_type: String,
|
||||
}
|
||||
|
||||
/// Response from the upload-snapshot endpoint.
|
||||
///
|
||||
/// The `uploads` list is aligned by index with the [`SnapshotUploadRequest::files`]
|
||||
/// list in the request, so callers match each upload target back to the filename
|
||||
/// they requested by position. The server does not include filenames on the
|
||||
/// response entries — see the `UploadSnapshotResponse` schema in
|
||||
/// `warp-server`'s `public_api/openapi.yaml`.
|
||||
#[derive(Debug, Clone, serde::Deserialize)]
|
||||
pub struct SnapshotUploadResponse {
|
||||
pub uploads: Vec<UploadTarget>,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct CreateExternalConversationRequest {
|
||||
format: String,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct CreateExternalConversationResponse {
|
||||
conversation_id: String,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct GetUploadTargetRequest {
|
||||
conversation_id: String,
|
||||
}
|
||||
|
||||
/// Skill attached to a resolve-prompt request,
|
||||
/// used when invoking a third-party harness with a skill
|
||||
/// via the CLI.
|
||||
#[derive(serde::Serialize)]
|
||||
pub struct ResolvePromptAttachedSkill {
|
||||
pub name: String,
|
||||
pub content: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub path: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
pub struct ResolvePromptRequest {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub skill: Option<ResolvePromptAttachedSkill>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub attachments_dir: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
pub struct ResolvedHarnessPrompt {
|
||||
pub prompt: String,
|
||||
#[serde(default)]
|
||||
pub system_prompt: Option<String>,
|
||||
/// Optional user-turn preamble for resumed third-party harness sessions. The harness
|
||||
/// decides how to surface this — Claude Code prepends it to the user-turn prompt fed
|
||||
/// into the CLI so the agent treats it as immediate intent rather than background
|
||||
/// system context. Empty when no resumption is in effect.
|
||||
#[serde(default)]
|
||||
pub resumption_prompt: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize, serde::Serialize)]
|
||||
pub struct ReportArtifactResponse {
|
||||
pub artifact_uid: String,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct NotifyUserRequest {
|
||||
message: String,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct FinishTaskRequest {
|
||||
success: bool,
|
||||
summary: String,
|
||||
}
|
||||
|
||||
/// Trait for API endpoints used to support third-party agent harnesses in Oz.
|
||||
#[cfg_attr(test, automock)]
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
pub trait HarnessSupportClient: 'static + Send + Sync {
|
||||
/// Create a new external conversation for a third-party harness.
|
||||
async fn create_external_conversation(&self, format: &str) -> Result<AIConversationId>;
|
||||
|
||||
/// Get a presigned upload target for the conversation's raw transcript.
|
||||
async fn get_transcript_upload_target(
|
||||
&self,
|
||||
conversation_id: &AIConversationId,
|
||||
) -> Result<UploadTarget>;
|
||||
|
||||
/// Get a presigned upload target for the conversation's block snapshot.
|
||||
async fn get_block_snapshot_upload_target(
|
||||
&self,
|
||||
conversation_id: &AIConversationId,
|
||||
) -> Result<UploadTarget>;
|
||||
|
||||
/// Resolve the prompt for a third-party harness run for a task stored on the server.
|
||||
async fn resolve_prompt(&self, request: ResolvePromptRequest) -> Result<ResolvedHarnessPrompt>;
|
||||
|
||||
/// Report an artifact created by a third-party harness back to the Oz platform.
|
||||
async fn report_artifact(&self, artifact: &Artifact) -> Result<ReportArtifactResponse>;
|
||||
|
||||
/// Send a progress notification to the task's originating platform.
|
||||
async fn notify_user(&self, message: &str) -> Result<()>;
|
||||
|
||||
/// Report task completion or failure. The server derives PR links/branches from
|
||||
/// artifacts already reported via `report_artifact`.
|
||||
async fn finish_task(&self, success: bool, summary: &str) -> Result<()>;
|
||||
|
||||
/// Get presigned upload targets for a workspace state snapshot.
|
||||
///
|
||||
/// The returned list is aligned by index with `request.files`. See
|
||||
/// [`SnapshotUploadResponse`] for details on the server contract.
|
||||
async fn get_snapshot_upload_targets(
|
||||
&self,
|
||||
request: &SnapshotUploadRequest,
|
||||
) -> Result<Vec<UploadTarget>>;
|
||||
|
||||
/// Download the raw third-party harness transcript bytes for the current task's
|
||||
/// conversation.
|
||||
///
|
||||
/// Hits `GET /harness-support/transcript`, which redirects to a signed GCS URL.
|
||||
/// The conversation is resolved from the task's `agent_conversation_id` server-side,
|
||||
/// so callers do not pass a conversation id. Each harness deserializes the returned
|
||||
/// bytes into its own envelope shape (e.g. Claude Code parses
|
||||
/// `ClaudeTranscriptEnvelope`). Transient failures retry with bounded exponential
|
||||
/// backoff; permanent 4xx (e.g. 404 "no transcript") fail fast so the caller can
|
||||
/// surface a resume-specific error.
|
||||
async fn fetch_transcript(&self) -> Result<bytes::Bytes>;
|
||||
|
||||
/// Get an HTTP client to use with [`UploadTarget`]s for saving blobs.
|
||||
fn http_client(&self) -> &http_client::Client;
|
||||
}
|
||||
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
impl HarnessSupportClient for ServerApi {
|
||||
async fn create_external_conversation(&self, format: &str) -> Result<AIConversationId> {
|
||||
let response: CreateExternalConversationResponse = self
|
||||
.post_public_api(
|
||||
"harness-support/external-conversation",
|
||||
&CreateExternalConversationRequest {
|
||||
format: format.to_string(),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
|
||||
AIConversationId::try_from(response.conversation_id)
|
||||
.context("Server returned an invalid conversation ID")
|
||||
}
|
||||
|
||||
async fn get_transcript_upload_target(
|
||||
&self,
|
||||
conversation_id: &AIConversationId,
|
||||
) -> Result<UploadTarget> {
|
||||
self.post_public_api(
|
||||
"harness-support/transcript",
|
||||
&GetUploadTargetRequest {
|
||||
conversation_id: conversation_id.to_string(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn get_block_snapshot_upload_target(
|
||||
&self,
|
||||
conversation_id: &AIConversationId,
|
||||
) -> Result<UploadTarget> {
|
||||
self.post_public_api(
|
||||
"harness-support/block-snapshot",
|
||||
&GetUploadTargetRequest {
|
||||
conversation_id: conversation_id.to_string(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn resolve_prompt(&self, request: ResolvePromptRequest) -> Result<ResolvedHarnessPrompt> {
|
||||
self.post_public_api("harness-support/resolve-prompt", &request)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn report_artifact(&self, artifact: &Artifact) -> Result<ReportArtifactResponse> {
|
||||
self.post_public_api("harness-support/report-artifact", artifact)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn notify_user(&self, message: &str) -> Result<()> {
|
||||
self.post_public_api_unit(
|
||||
"harness-support/notify-user",
|
||||
&NotifyUserRequest {
|
||||
message: message.to_string(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn finish_task(&self, success: bool, summary: &str) -> Result<()> {
|
||||
self.post_public_api_unit(
|
||||
"harness-support/finish-task",
|
||||
&FinishTaskRequest {
|
||||
success,
|
||||
summary: summary.to_string(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn get_snapshot_upload_targets(
|
||||
&self,
|
||||
request: &SnapshotUploadRequest,
|
||||
) -> Result<Vec<UploadTarget>> {
|
||||
let response: SnapshotUploadResponse = self
|
||||
.post_public_api("harness-support/upload-snapshot", request)
|
||||
.await?;
|
||||
Ok(response.uploads)
|
||||
}
|
||||
|
||||
async fn fetch_transcript(&self) -> Result<bytes::Bytes> {
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
{
|
||||
with_bounded_retry("fetch harness-support transcript", || async {
|
||||
let response = self
|
||||
.get_public_api_response("harness-support/transcript")
|
||||
.await?;
|
||||
response
|
||||
.bytes()
|
||||
.await
|
||||
.context("Failed to read harness-support transcript body")
|
||||
})
|
||||
.await
|
||||
}
|
||||
#[cfg(target_family = "wasm")]
|
||||
{
|
||||
unreachable!(
|
||||
"fetch_transcript is not supported on wasm; agent_sdk is not built on this target"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn http_client(&self) -> &http_client::Client {
|
||||
&self.client
|
||||
}
|
||||
}
|
||||
|
||||
/// Upload a blob to a presigned upload target.
|
||||
pub async fn upload_to_target(
|
||||
http_client: &http_client::Client,
|
||||
target: &UploadTarget,
|
||||
body: impl Into<reqwest::Body>,
|
||||
) -> Result<()> {
|
||||
super::presigned_upload::upload_to_target(http_client, target, body).await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "harness_support_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,25 @@
|
||||
use crate::ai::artifacts::Artifact;
|
||||
|
||||
/// Assert that `Artifact`s serialize to the expected format for the /harness-support/report-artifact
|
||||
/// endpoint.
|
||||
/// If `Artifact` serialization changes, this test will catch it.
|
||||
#[test]
|
||||
fn pull_request_artifact_serializes_to_expected_wire_format() {
|
||||
let artifact = Artifact::PullRequest {
|
||||
url: "https://github.com/org/repo/pull/42".to_string(),
|
||||
branch: "feature-branch".to_string(),
|
||||
repo: Some("repo".to_string()),
|
||||
number: Some(42),
|
||||
};
|
||||
let json = serde_json::to_value(&artifact).unwrap();
|
||||
assert_eq!(
|
||||
json,
|
||||
serde_json::json!({
|
||||
"artifact_type": "PULL_REQUEST",
|
||||
"data": {
|
||||
"url": "https://github.com/org/repo/pull/42",
|
||||
"branch": "feature-branch"
|
||||
}
|
||||
})
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,350 @@
|
||||
use super::ServerApi;
|
||||
use anyhow::{anyhow, Result};
|
||||
use async_trait::async_trait;
|
||||
use cynic::{MutationBuilder, QueryBuilder};
|
||||
|
||||
use crate::channel::ChannelState;
|
||||
use crate::features::FeatureFlag;
|
||||
#[cfg(test)]
|
||||
use mockall::automock;
|
||||
|
||||
use crate::server::graphql::{get_request_context, get_user_facing_error_message};
|
||||
use warp_graphql::mutations::create_simple_integration::{
|
||||
CreateSimpleIntegration, CreateSimpleIntegrationOutput, CreateSimpleIntegrationResult,
|
||||
CreateSimpleIntegrationVariables, SimpleIntegrationConfig,
|
||||
};
|
||||
use warp_graphql::queries::get_integrations_using_environment::{
|
||||
GetIntegrationsUsingEnvironment, GetIntegrationsUsingEnvironmentInput,
|
||||
GetIntegrationsUsingEnvironmentOutput, GetIntegrationsUsingEnvironmentResult,
|
||||
GetIntegrationsUsingEnvironmentVariables,
|
||||
};
|
||||
use warp_graphql::queries::get_oauth_connect_tx_status::{
|
||||
GetOAuthConnectTxStatus, GetOAuthConnectTxStatusInput, GetOAuthConnectTxStatusResult,
|
||||
GetOAuthConnectTxStatusVariables, OauthConnectTxStatus,
|
||||
};
|
||||
use warp_graphql::queries::get_simple_integrations::{
|
||||
SimpleIntegrations, SimpleIntegrationsInput, SimpleIntegrationsOutput,
|
||||
SimpleIntegrationsResult, SimpleIntegrationsVariables,
|
||||
};
|
||||
use warp_graphql::queries::suggest_cloud_environment_image::{
|
||||
RepoInput as SuggestCloudEnvironmentImageRepoInput, SuggestCloudEnvironmentImage,
|
||||
SuggestCloudEnvironmentImageInput, SuggestCloudEnvironmentImageResult,
|
||||
SuggestCloudEnvironmentImageVariables,
|
||||
};
|
||||
use warp_graphql::queries::user_github_info::{
|
||||
GithubAuthRequiredOutput, UserGithubInfo, UserGithubInfoResult, UserGithubInfoVariables,
|
||||
};
|
||||
use warp_graphql::queries::user_repo_auth_status::{
|
||||
RepoInput as UserRepoAuthStatusRepoInput, UserRepoAuthStatus, UserRepoAuthStatusInput,
|
||||
UserRepoAuthStatusOutput, UserRepoAuthStatusResult, UserRepoAuthStatusVariables,
|
||||
};
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
pub trait IntegrationsClientBounds: Send + Sync {}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
impl<T: 'static + Send + Sync> IntegrationsClientBounds for T {}
|
||||
|
||||
#[cfg(target_family = "wasm")]
|
||||
pub trait IntegrationsClientBounds {}
|
||||
|
||||
#[cfg(target_family = "wasm")]
|
||||
impl<T: 'static> IntegrationsClientBounds for T {}
|
||||
|
||||
#[cfg_attr(test, automock)]
|
||||
#[cfg_attr(target_family = "wasm", allow(dead_code))]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
pub trait IntegrationsClient: 'static + IntegrationsClientBounds {
|
||||
/// Checks the user's GitHub authorization status for the given repositories.
|
||||
///
|
||||
/// Returns a list of statuses for each repo, indicating whether the user has
|
||||
/// access to the repo, and an optional auth URL for the user to authorize.
|
||||
async fn check_user_repo_auth_status(
|
||||
&self,
|
||||
repos: Vec<(String, String)>,
|
||||
) -> Result<UserRepoAuthStatusOutput>;
|
||||
|
||||
/// Creates or updates a simple integration on the server.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `integration_type` - The type of integration (e.g. "github", "linear", "slack")
|
||||
/// * `is_update` - Whether this is an update to an existing integration
|
||||
/// * `environment_uid` - The UID of the environment to associate with this integration
|
||||
/// * `base_prompt` - Optional base prompt for the integration
|
||||
/// * `model_id` - Optional model ID for the integration
|
||||
/// * `mcp_servers_json` - Optional JSON string encoding a map[string]MCPServerConfig (ambient agent spec)
|
||||
/// * `remove_mcp_server_names` - Optional list of MCP server names to remove (applies on update)
|
||||
/// * `worker_host` - Optional worker host ID for self-hosted workers
|
||||
/// * `enabled` - Whether the integration should be enabled on creation
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn create_or_update_simple_integration(
|
||||
&self,
|
||||
integration_type: String,
|
||||
is_update: bool,
|
||||
environment_uid: Option<String>,
|
||||
base_prompt: Option<String>,
|
||||
model_id: Option<String>,
|
||||
mcp_servers_json: Option<String>,
|
||||
remove_mcp_server_names: Option<Vec<String>>,
|
||||
worker_host: Option<String>,
|
||||
enabled: bool,
|
||||
) -> Result<CreateSimpleIntegrationOutput>;
|
||||
|
||||
/// Lists simple integrations for a fixed set of provider slugs.
|
||||
///
|
||||
/// The server will return one SimpleIntegration entry per requested provider,
|
||||
/// regardless of whether the connection or integration currently exists.
|
||||
async fn list_simple_integrations(
|
||||
&self,
|
||||
providers: Vec<String>,
|
||||
) -> Result<SimpleIntegrationsOutput>;
|
||||
|
||||
/// Polls the status of an OAuth connect transaction.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `tx_id` - The transaction ID returned from create_simple_integration
|
||||
///
|
||||
/// # Returns
|
||||
/// * `Ok(OauthConnectTxStatus)` - The current status of the transaction
|
||||
/// * `Err` - If the transaction is not found or polling fails
|
||||
async fn poll_oauth_connect_status(&self, tx_id: String) -> Result<OauthConnectTxStatus>;
|
||||
|
||||
/// Gets the list of integration provider names that are using the specified environment.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `environment_id` - The ID of the environment to check
|
||||
///
|
||||
/// # Returns
|
||||
/// * `Ok(Vec<String>)` - List of provider names (e.g., ["linear", "slack"]) using this environment
|
||||
/// * `Err` - If the query fails
|
||||
async fn get_integrations_using_environment(
|
||||
&self,
|
||||
environment_id: String,
|
||||
) -> Result<GetIntegrationsUsingEnvironmentOutput>;
|
||||
|
||||
/// Gets the user's GitHub connection info, including accessible repos.
|
||||
///
|
||||
/// # Returns
|
||||
/// * `Ok(UserGithubInfoResult)` - Either connected with repos, or auth required
|
||||
/// * `Err` - If the query fails
|
||||
async fn get_user_github_info(&self) -> Result<UserGithubInfoResult>;
|
||||
|
||||
/// Suggests a Docker image for a cloud environment based on the provided repos.
|
||||
async fn suggest_cloud_environment_image(
|
||||
&self,
|
||||
repos: Vec<(String, String)>,
|
||||
) -> Result<SuggestCloudEnvironmentImageResult>;
|
||||
}
|
||||
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
impl IntegrationsClient for ServerApi {
|
||||
async fn check_user_repo_auth_status(
|
||||
&self,
|
||||
repos: Vec<(String, String)>,
|
||||
) -> Result<UserRepoAuthStatusOutput> {
|
||||
let repo_inputs: Vec<UserRepoAuthStatusRepoInput> = repos
|
||||
.into_iter()
|
||||
.map(|(owner, repo)| UserRepoAuthStatusRepoInput { owner, repo })
|
||||
.collect();
|
||||
|
||||
let variables = UserRepoAuthStatusVariables {
|
||||
request_context: get_request_context(),
|
||||
input: UserRepoAuthStatusInput { repos: repo_inputs },
|
||||
};
|
||||
|
||||
let operation = UserRepoAuthStatus::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.user_repo_auth_status {
|
||||
UserRepoAuthStatusResult::UserRepoAuthStatusOutput(output) => Ok(output),
|
||||
UserRepoAuthStatusResult::Unknown => Err(anyhow::anyhow!(
|
||||
"Failed to check GitHub auth status: unknown response"
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn create_or_update_simple_integration(
|
||||
&self,
|
||||
integration_type: String,
|
||||
is_update: bool,
|
||||
environment_uid: Option<String>,
|
||||
base_prompt: Option<String>,
|
||||
model_id: Option<String>,
|
||||
mcp_servers_json: Option<String>,
|
||||
remove_mcp_server_names: Option<Vec<String>>,
|
||||
worker_host: Option<String>,
|
||||
enabled: bool,
|
||||
) -> Result<CreateSimpleIntegrationOutput> {
|
||||
let variables = CreateSimpleIntegrationVariables {
|
||||
config: SimpleIntegrationConfig {
|
||||
base_prompt,
|
||||
environment_uid,
|
||||
model_id,
|
||||
mcp_servers_json,
|
||||
remove_mcp_server_names,
|
||||
worker_host,
|
||||
},
|
||||
enabled,
|
||||
integration_type,
|
||||
is_update,
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = CreateSimpleIntegration::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
match response.create_simple_integration {
|
||||
CreateSimpleIntegrationResult::CreateSimpleIntegrationOutput(output) => Ok(output),
|
||||
CreateSimpleIntegrationResult::UserFacingError(error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(error)))
|
||||
}
|
||||
CreateSimpleIntegrationResult::Unknown => {
|
||||
Err(anyhow!("Unknown error while creating integration"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_integrations_using_environment(
|
||||
&self,
|
||||
environment_id: String,
|
||||
) -> Result<GetIntegrationsUsingEnvironmentOutput> {
|
||||
let variables = GetIntegrationsUsingEnvironmentVariables {
|
||||
request_context: get_request_context(),
|
||||
input: GetIntegrationsUsingEnvironmentInput { environment_id },
|
||||
};
|
||||
|
||||
let operation = GetIntegrationsUsingEnvironment::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.get_integrations_using_environment {
|
||||
GetIntegrationsUsingEnvironmentResult::GetIntegrationsUsingEnvironmentOutput(
|
||||
output,
|
||||
) => Ok(output),
|
||||
GetIntegrationsUsingEnvironmentResult::UserFacingError(error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(error)))
|
||||
}
|
||||
GetIntegrationsUsingEnvironmentResult::Unknown => Err(anyhow!(
|
||||
"Unknown error while getting integrations using environment"
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_simple_integrations(
|
||||
&self,
|
||||
providers: Vec<String>,
|
||||
) -> Result<SimpleIntegrationsOutput> {
|
||||
let variables = SimpleIntegrationsVariables {
|
||||
request_context: get_request_context(),
|
||||
input: SimpleIntegrationsInput { providers },
|
||||
};
|
||||
|
||||
let operation = SimpleIntegrations::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.simple_integrations {
|
||||
SimpleIntegrationsResult::SimpleIntegrationsOutput(output) => Ok(output),
|
||||
SimpleIntegrationsResult::UserFacingError(error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(error)))
|
||||
}
|
||||
SimpleIntegrationsResult::Unknown => {
|
||||
Err(anyhow!("Unknown error while listing simple integrations"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn poll_oauth_connect_status(&self, tx_id: String) -> Result<OauthConnectTxStatus> {
|
||||
let variables = GetOAuthConnectTxStatusVariables {
|
||||
request_context: get_request_context(),
|
||||
input: GetOAuthConnectTxStatusInput {
|
||||
tx_id: cynic::Id::new(tx_id),
|
||||
},
|
||||
};
|
||||
|
||||
let operation = GetOAuthConnectTxStatus::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.get_oauth_connect_tx_status {
|
||||
GetOAuthConnectTxStatusResult::GetOAuthConnectTxStatusOutput(output) => {
|
||||
Ok(output.status)
|
||||
}
|
||||
GetOAuthConnectTxStatusResult::UserFacingError(error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(error)))
|
||||
}
|
||||
GetOAuthConnectTxStatusResult::Unknown => {
|
||||
Err(anyhow!("Unknown error while polling OAuth status"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_user_github_info(&self) -> Result<UserGithubInfoResult> {
|
||||
let variables = UserGithubInfoVariables {
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = UserGithubInfo::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
let result = response.user_github_info;
|
||||
|
||||
// Dev-only helper for testing GitHub-unauthed flows.
|
||||
//
|
||||
// Important: this runs after the network request completes so the UI can still
|
||||
// show the loading state.
|
||||
if FeatureFlag::SimulateGithubUnauthed.is_enabled() {
|
||||
if let UserGithubInfoResult::GithubConnectedOutput(connected) = &result {
|
||||
let auth_url = format!("{}/oauth/connect/github", ChannelState::server_root_url());
|
||||
return Ok(UserGithubInfoResult::GithubAuthRequiredOutput(
|
||||
GithubAuthRequiredOutput {
|
||||
auth_url,
|
||||
// This value is unused by the app UI; it exists in the schema for
|
||||
// tx-bound flows. We intentionally omit txId from the auth URL so
|
||||
// the web flow can proceed without a server-created tx.
|
||||
tx_id: cynic::Id::new("simulated"),
|
||||
app_install_link: connected.app_install_link.clone(),
|
||||
},
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn suggest_cloud_environment_image(
|
||||
&self,
|
||||
repos: Vec<(String, String)>,
|
||||
) -> Result<SuggestCloudEnvironmentImageResult> {
|
||||
let repo_inputs: Vec<SuggestCloudEnvironmentImageRepoInput> = repos
|
||||
.into_iter()
|
||||
.map(|(owner, repo)| SuggestCloudEnvironmentImageRepoInput { owner, repo })
|
||||
.collect();
|
||||
|
||||
let variables = SuggestCloudEnvironmentImageVariables {
|
||||
request_context: get_request_context(),
|
||||
input: SuggestCloudEnvironmentImageInput { repos: repo_inputs },
|
||||
};
|
||||
|
||||
let operation = SuggestCloudEnvironmentImage::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.suggest_cloud_environment_image {
|
||||
SuggestCloudEnvironmentImageResult::SuggestCloudEnvironmentImageAuthRequiredOutput(
|
||||
output,
|
||||
) => Ok(
|
||||
SuggestCloudEnvironmentImageResult::SuggestCloudEnvironmentImageAuthRequiredOutput(
|
||||
output,
|
||||
),
|
||||
),
|
||||
SuggestCloudEnvironmentImageResult::SuggestCloudEnvironmentImageOutput(output) => {
|
||||
Ok(SuggestCloudEnvironmentImageResult::SuggestCloudEnvironmentImageOutput(output))
|
||||
}
|
||||
SuggestCloudEnvironmentImageResult::UserFacingError(error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(error)))
|
||||
}
|
||||
SuggestCloudEnvironmentImageResult::Unknown => Err(anyhow!(
|
||||
"Unknown response from suggestCloudEnvironmentImage query"
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,291 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
use async_trait::async_trait;
|
||||
use cynic::{MutationBuilder, QueryBuilder};
|
||||
use warp_graphql::mutations::issue_task_identity_token::{
|
||||
IssueTaskIdentityToken, IssueTaskIdentityTokenInput, IssueTaskIdentityTokenResult,
|
||||
IssueTaskIdentityTokenVariables,
|
||||
};
|
||||
use warp_graphql::object_permissions::OwnerType;
|
||||
use warp_graphql::queries::list_managed_secrets::{
|
||||
ListManagedSecrets, ListManagedSecretsVariables, ManagedSecretsInput, ManagedSecretsResult,
|
||||
};
|
||||
use warp_graphql::queries::managed_secret_config::{
|
||||
GetManagedSecretConfig, GetManagedSecretConfigVariables, UserResult,
|
||||
};
|
||||
use warp_graphql::queries::task_secrets::{
|
||||
ManagedSecretValue, TaskSecrets, TaskSecretsInput, TaskSecretsResult, TaskSecretsVariables,
|
||||
};
|
||||
use warp_graphql::{
|
||||
managed_secrets::{ManagedSecret, ManagedSecretType},
|
||||
mutations::{
|
||||
create_managed_secret::{
|
||||
CreateManagedSecret, CreateManagedSecretInput, CreateManagedSecretResult,
|
||||
CreateManagedSecretVariables,
|
||||
},
|
||||
delete_managed_secret::{
|
||||
DeleteManagedSecret, DeleteManagedSecretInput, DeleteManagedSecretResult,
|
||||
DeleteManagedSecretVariables,
|
||||
},
|
||||
update_managed_secret::{
|
||||
UpdateManagedSecret, UpdateManagedSecretInput, UpdateManagedSecretResult,
|
||||
UpdateManagedSecretVariables,
|
||||
},
|
||||
},
|
||||
object_permissions::Owner,
|
||||
};
|
||||
use warp_managed_secrets::client::{SecretOwner, TaskIdentityToken};
|
||||
|
||||
use super::ServerApi;
|
||||
use crate::server::graphql::{get_request_context, get_user_facing_error_message};
|
||||
|
||||
pub use warp_managed_secrets::client::{ManagedSecretConfigs, ManagedSecretsClient};
|
||||
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
impl ManagedSecretsClient for ServerApi {
|
||||
async fn get_managed_secret_configs(&self) -> Result<ManagedSecretConfigs> {
|
||||
let variables = GetManagedSecretConfigVariables {
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = GetManagedSecretConfig::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.user {
|
||||
UserResult::UserOutput(output) => {
|
||||
let mut team_configs = HashMap::new();
|
||||
for workspace in output.user.workspaces {
|
||||
for team in workspace.teams {
|
||||
if let Some(config) = team.managed_secrets {
|
||||
// DO NOT inline the `insert` call into the `debug_assert!` macro. It will get compiled out in release builds.
|
||||
let prior_config = team_configs.insert(team.uid.into_inner(), config);
|
||||
debug_assert!(
|
||||
prior_config.is_none(),
|
||||
"Duplicate team UID returned from server"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(ManagedSecretConfigs {
|
||||
user_secrets: output.user.managed_secrets,
|
||||
team_secrets: team_configs,
|
||||
})
|
||||
}
|
||||
UserResult::UserFacingError(error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(error)))
|
||||
}
|
||||
UserResult::Unknown => Err(anyhow!(
|
||||
"Unknown error while getting managed secret configs"
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn create_managed_secret(
|
||||
&self,
|
||||
owner: SecretOwner,
|
||||
name: String,
|
||||
secret_type: ManagedSecretType,
|
||||
encrypted_value: String,
|
||||
description: Option<String>,
|
||||
) -> Result<ManagedSecret> {
|
||||
let graphql_owner = match owner {
|
||||
SecretOwner::CurrentUser => Owner {
|
||||
type_: OwnerType::User,
|
||||
uid: None,
|
||||
},
|
||||
SecretOwner::Team { team_uid } => Owner {
|
||||
type_: OwnerType::Team,
|
||||
uid: Some(cynic::Id::new(team_uid)),
|
||||
},
|
||||
};
|
||||
|
||||
let variables = CreateManagedSecretVariables {
|
||||
input: CreateManagedSecretInput {
|
||||
description,
|
||||
encrypted_value,
|
||||
name,
|
||||
owner: graphql_owner,
|
||||
type_: secret_type,
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = CreateManagedSecret::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.create_managed_secret {
|
||||
CreateManagedSecretResult::CreateManagedSecretOutput(output) => {
|
||||
Ok(output.managed_secret)
|
||||
}
|
||||
CreateManagedSecretResult::UserFacingError(error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(error)))
|
||||
}
|
||||
CreateManagedSecretResult::Unknown => {
|
||||
Err(anyhow!("Unknown error while creating managed secret"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_managed_secret(&self, owner: SecretOwner, name: String) -> Result<()> {
|
||||
let graphql_owner = match owner {
|
||||
SecretOwner::CurrentUser => Owner {
|
||||
type_: OwnerType::User,
|
||||
uid: None,
|
||||
},
|
||||
SecretOwner::Team { team_uid } => Owner {
|
||||
type_: OwnerType::Team,
|
||||
uid: Some(cynic::Id::new(team_uid)),
|
||||
},
|
||||
};
|
||||
|
||||
let variables = DeleteManagedSecretVariables {
|
||||
input: DeleteManagedSecretInput {
|
||||
name,
|
||||
owner: graphql_owner,
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = DeleteManagedSecret::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.delete_managed_secret {
|
||||
DeleteManagedSecretResult::DeleteManagedSecretOutput(_) => Ok(()),
|
||||
DeleteManagedSecretResult::UserFacingError(error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(error)))
|
||||
}
|
||||
DeleteManagedSecretResult::Unknown => {
|
||||
Err(anyhow!("Unknown error while deleting managed secret"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn update_managed_secret(
|
||||
&self,
|
||||
owner: SecretOwner,
|
||||
name: String,
|
||||
encrypted_value: Option<String>,
|
||||
description: Option<String>,
|
||||
) -> Result<ManagedSecret> {
|
||||
let graphql_owner = match owner {
|
||||
SecretOwner::CurrentUser => Owner {
|
||||
type_: OwnerType::User,
|
||||
uid: None,
|
||||
},
|
||||
SecretOwner::Team { team_uid } => Owner {
|
||||
type_: OwnerType::Team,
|
||||
uid: Some(cynic::Id::new(team_uid)),
|
||||
},
|
||||
};
|
||||
|
||||
let variables = UpdateManagedSecretVariables {
|
||||
input: UpdateManagedSecretInput {
|
||||
name,
|
||||
owner: graphql_owner,
|
||||
encrypted_value,
|
||||
description,
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = UpdateManagedSecret::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.update_managed_secret {
|
||||
UpdateManagedSecretResult::UpdateManagedSecretOutput(output) => {
|
||||
Ok(output.managed_secret)
|
||||
}
|
||||
UpdateManagedSecretResult::UserFacingError(error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(error)))
|
||||
}
|
||||
UpdateManagedSecretResult::Unknown => {
|
||||
Err(anyhow!("Unknown error while updating managed secret"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_secrets(&self) -> Result<Vec<ManagedSecret>> {
|
||||
let variables = ListManagedSecretsVariables {
|
||||
// Pagination over managed secrets is not yet supported.
|
||||
input: ManagedSecretsInput { cursor: None },
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = ListManagedSecrets::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.managed_secrets {
|
||||
ManagedSecretsResult::ManagedSecretsOutput(output) => Ok(output.managed_secrets),
|
||||
ManagedSecretsResult::UserFacingError(error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(error)))
|
||||
}
|
||||
ManagedSecretsResult::Unknown => {
|
||||
Err(anyhow!("Unknown error while listing managed secrets"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_task_secrets(
|
||||
&self,
|
||||
task_id: String,
|
||||
workload_token: String,
|
||||
) -> Result<HashMap<String, ManagedSecretValue>> {
|
||||
let variables = TaskSecretsVariables {
|
||||
input: TaskSecretsInput {
|
||||
task_id: cynic::Id::new(task_id),
|
||||
workload_token,
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = TaskSecrets::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.task_secrets {
|
||||
TaskSecretsResult::TaskSecretsOutput(output) => {
|
||||
let mut secrets = HashMap::new();
|
||||
for entry in output.secrets {
|
||||
secrets.insert(entry.name, entry.value);
|
||||
}
|
||||
Ok(secrets)
|
||||
}
|
||||
TaskSecretsResult::UserFacingError(error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(error)))
|
||||
}
|
||||
TaskSecretsResult::Unknown => Err(anyhow!("Unknown error while getting task secrets")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn issue_task_identity_token(
|
||||
&self,
|
||||
options: warp_managed_secrets::client::IdentityTokenOptions,
|
||||
) -> Result<TaskIdentityToken> {
|
||||
let requested_duration_seconds = options
|
||||
.requested_duration
|
||||
.as_secs()
|
||||
.try_into()
|
||||
.context("Requested duration out of bounds")?;
|
||||
let variables = IssueTaskIdentityTokenVariables {
|
||||
input: IssueTaskIdentityTokenInput {
|
||||
audience: options.audience,
|
||||
requested_duration_seconds,
|
||||
subject_template: Some(options.subject_template.into_vec()),
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = IssueTaskIdentityToken::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.issue_task_identity_token {
|
||||
IssueTaskIdentityTokenResult::IssueTaskIdentityTokenOutput(output) => {
|
||||
Ok(TaskIdentityToken {
|
||||
token: output.token,
|
||||
expires_at: output.expires_at.utc(),
|
||||
issuer: output.issuer,
|
||||
})
|
||||
}
|
||||
IssueTaskIdentityTokenResult::UserFacingError(error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(error)))
|
||||
}
|
||||
IssueTaskIdentityTokenResult::Unknown => {
|
||||
Err(anyhow!("Unknown error while issuing task identity token"))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,268 @@
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use std::path::{Path, PathBuf};
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use async_stream::try_stream;
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use base64::{engine::general_purpose::STANDARD, Engine as _};
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use bytes::Bytes;
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use crc::{Crc, CRC_32_ISCSI};
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use futures_lite::io::AsyncReadExt as _;
|
||||
use thiserror::Error;
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use super::ai::FileArtifactUploadTargetInfo;
|
||||
use super::harness_support::UploadTarget;
|
||||
|
||||
/// Typed error for HTTP-backed operations so downstream classifiers (e.g. the agent-SDK
|
||||
/// retry helper) can decide transient vs permanent failures without string-parsing the
|
||||
/// anyhow Display.
|
||||
///
|
||||
/// Emitted as the source cause of an upload failure; callers typically also attach a
|
||||
/// human-facing context message via `.context(...)` so `err.to_string()` remains useful.
|
||||
#[derive(Debug, Error)]
|
||||
#[error("HTTP request failed with status {status}: {body}")]
|
||||
pub struct HttpStatusError {
|
||||
pub status: u16,
|
||||
pub body: String,
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
static CRC32C: Crc<u32> = Crc::<u32>::new(&CRC_32_ISCSI);
|
||||
const CONTENT_LENGTH_HEADER_NAME: &str = "content-length";
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
const FILE_UPLOAD_CHUNK_SIZE: usize = 64 * 1024;
|
||||
|
||||
struct NormalizedUploadTarget<'a> {
|
||||
url: &'a str,
|
||||
method: &'a str,
|
||||
headers: Vec<(&'a str, &'a str)>,
|
||||
}
|
||||
|
||||
impl<'a> From<&'a UploadTarget> for NormalizedUploadTarget<'a> {
|
||||
fn from(target: &'a UploadTarget) -> Self {
|
||||
Self {
|
||||
url: &target.url,
|
||||
method: &target.method,
|
||||
headers: target
|
||||
.headers
|
||||
.iter()
|
||||
.map(|(name, value)| (name.as_str(), value.as_str()))
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
impl<'a> From<&'a FileArtifactUploadTargetInfo> for NormalizedUploadTarget<'a> {
|
||||
fn from(target: &'a FileArtifactUploadTargetInfo) -> Self {
|
||||
Self {
|
||||
url: &target.url,
|
||||
method: &target.method,
|
||||
headers: target
|
||||
.headers
|
||||
.iter()
|
||||
.map(|header| (header.name.as_str(), header.value.as_str()))
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
#[derive(Clone)]
|
||||
struct SharedChecksumState(Arc<Mutex<Option<crc::Digest<'static, u32>>>>);
|
||||
|
||||
#[derive(Copy, Clone)]
|
||||
struct UploadErrorContext {
|
||||
transport: &'static str,
|
||||
failure: &'static str,
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
impl SharedChecksumState {
|
||||
fn new() -> Self {
|
||||
Self(Arc::new(Mutex::new(Some(CRC32C.digest()))))
|
||||
}
|
||||
|
||||
fn update(&self, bytes: &[u8]) {
|
||||
if bytes.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let mut digest = self.0.lock().expect("checksum state mutex poisoned");
|
||||
digest
|
||||
.as_mut()
|
||||
.expect("checksum already finalized")
|
||||
.update(bytes);
|
||||
}
|
||||
|
||||
fn finalize(&self) -> Result<String> {
|
||||
let digest = self
|
||||
.0
|
||||
.lock()
|
||||
.map_err(|_| anyhow!("checksum state mutex poisoned"))?
|
||||
.take()
|
||||
.ok_or_else(|| anyhow!("checksum already finalized"))?;
|
||||
Ok(encode_crc32c_base64(digest.finalize()))
|
||||
}
|
||||
}
|
||||
|
||||
fn build_upload_request<'a>(
|
||||
http_client: &'a http_client::Client,
|
||||
target: NormalizedUploadTarget<'_>,
|
||||
content_length: Option<u64>,
|
||||
) -> Result<http_client::RequestBuilder<'a>> {
|
||||
let method = target.method.to_ascii_uppercase();
|
||||
let has_content_length = target
|
||||
.headers
|
||||
.iter()
|
||||
.any(|(name, _)| name.eq_ignore_ascii_case(CONTENT_LENGTH_HEADER_NAME));
|
||||
|
||||
let mut request = match method.as_str() {
|
||||
"GET" => http_client.get(target.url),
|
||||
"POST" => http_client.post(target.url),
|
||||
"PUT" => http_client.put(target.url),
|
||||
"DELETE" => http_client.delete(target.url),
|
||||
other => return Err(anyhow!("Unsupported HTTP method: {other}")),
|
||||
};
|
||||
|
||||
for (name, value) in target.headers {
|
||||
request = request.header(name, value);
|
||||
}
|
||||
|
||||
if let Some(content_length) = content_length.filter(|_| !has_content_length) {
|
||||
request = request.header(CONTENT_LENGTH_HEADER_NAME, content_length.to_string());
|
||||
}
|
||||
|
||||
Ok(request)
|
||||
}
|
||||
|
||||
async fn ensure_upload_succeeded(
|
||||
response: http_client::Response,
|
||||
error_context: UploadErrorContext,
|
||||
) -> Result<()> {
|
||||
if response.status().is_success() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
let status_err = HttpStatusError {
|
||||
status: status.as_u16(),
|
||||
body: body.clone(),
|
||||
};
|
||||
Err(anyhow::Error::new(status_err).context(format!(
|
||||
"{} failed with status {status}: {body}",
|
||||
error_context.failure
|
||||
)))
|
||||
}
|
||||
|
||||
async fn send_upload_request(
|
||||
http_client: &http_client::Client,
|
||||
target: NormalizedUploadTarget<'_>,
|
||||
body: impl Into<reqwest::Body>,
|
||||
content_length: Option<u64>,
|
||||
error_context: UploadErrorContext,
|
||||
) -> Result<()> {
|
||||
let response = build_upload_request(http_client, target, content_length)?
|
||||
.body(body)
|
||||
.send()
|
||||
.await
|
||||
.context(error_context.transport)?;
|
||||
|
||||
ensure_upload_succeeded(response, error_context).await
|
||||
}
|
||||
|
||||
pub(crate) async fn upload_to_target(
|
||||
http_client: &http_client::Client,
|
||||
target: &UploadTarget,
|
||||
body: impl Into<reqwest::Body>,
|
||||
) -> Result<()> {
|
||||
send_upload_request(
|
||||
http_client,
|
||||
target.into(),
|
||||
body,
|
||||
None,
|
||||
UploadErrorContext {
|
||||
transport: "Failed to upload to presigned URL",
|
||||
failure: "Upload",
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
fn encode_crc32c_base64(crc32c: u32) -> String {
|
||||
// Storage providers expect the checksum as base64 of the raw big-endian CRC32C bytes,
|
||||
// not the more human-readable hex string we typically log.
|
||||
STANDARD.encode(crc32c.to_be_bytes())
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
fn file_upload_stream(
|
||||
mut file: async_fs::File,
|
||||
path: PathBuf,
|
||||
checksum: SharedChecksumState,
|
||||
) -> impl futures::Stream<Item = std::io::Result<Bytes>> + Send + 'static {
|
||||
try_stream! {
|
||||
loop {
|
||||
let mut chunk = vec![0; FILE_UPLOAD_CHUNK_SIZE];
|
||||
let bytes_read = file.read(&mut chunk).await.map_err(|err| {
|
||||
std::io::Error::other(format!(
|
||||
"Failed to read artifact file '{}': {err}",
|
||||
path.display()
|
||||
))
|
||||
})?;
|
||||
|
||||
if bytes_read == 0 {
|
||||
break;
|
||||
}
|
||||
|
||||
chunk.truncate(bytes_read);
|
||||
checksum.update(&chunk);
|
||||
yield Bytes::from(chunk);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
pub(crate) async fn upload_file_to_target(
|
||||
http_client: &http_client::Client,
|
||||
target: &FileArtifactUploadTargetInfo,
|
||||
path: &Path,
|
||||
file_size: u64,
|
||||
) -> Result<String> {
|
||||
let file = async_fs::File::open(path)
|
||||
.await
|
||||
.with_context(|| format!("Failed to open artifact file '{}'", path.display()))?;
|
||||
let checksum = SharedChecksumState::new();
|
||||
let body = reqwest::Body::wrap_stream(file_upload_stream(
|
||||
file,
|
||||
path.to_path_buf(),
|
||||
checksum.clone(),
|
||||
));
|
||||
|
||||
send_upload_request(
|
||||
http_client,
|
||||
target.into(),
|
||||
body,
|
||||
Some(file_size),
|
||||
UploadErrorContext {
|
||||
transport: "Failed to upload artifact bytes",
|
||||
failure: "Artifact upload",
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
|
||||
checksum.finalize()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "presigned_upload_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,125 @@
|
||||
use std::collections::HashMap;
|
||||
use std::fs;
|
||||
|
||||
use futures::executor::block_on;
|
||||
use mockito::Server;
|
||||
use tempfile::tempdir;
|
||||
|
||||
use super::*;
|
||||
use crate::server::server_api::ai::{FileArtifactUploadHeaderInfo, FileArtifactUploadTargetInfo};
|
||||
|
||||
#[test]
|
||||
fn encode_crc32c_base64_matches_spec_example() {
|
||||
assert_eq!(encode_crc32c_base64(0x1234_5678), "EjRWeA==");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shared_checksum_state_finalize_returns_error_when_called_twice() {
|
||||
let checksum = SharedChecksumState::new();
|
||||
checksum.update(b"artifact payload");
|
||||
|
||||
let finalized = checksum.finalize().unwrap();
|
||||
let err = checksum.finalize().unwrap_err();
|
||||
|
||||
assert_eq!(
|
||||
finalized,
|
||||
encode_crc32c_base64(CRC32C.checksum(b"artifact payload"))
|
||||
);
|
||||
assert!(err.to_string().contains("checksum already finalized"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn upload_to_target_replays_headers_for_byte_uploads() {
|
||||
block_on(async {
|
||||
let mut server = Server::new();
|
||||
let mock = server
|
||||
.mock("POST", "/upload")
|
||||
.match_header("x-test-header", "expected-header")
|
||||
.match_body("serialized body")
|
||||
.with_status(200)
|
||||
.create();
|
||||
|
||||
let client = http_client::Client::new_for_test();
|
||||
let target = UploadTarget {
|
||||
url: format!("{}/upload", server.url()),
|
||||
method: "POST".to_string(),
|
||||
headers: HashMap::from([("x-test-header".to_string(), "expected-header".to_string())]),
|
||||
};
|
||||
|
||||
upload_to_target(&client, &target, "serialized body".to_string())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
mock.assert();
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn upload_file_to_target_replays_headers_sets_content_length_and_returns_checksum() {
|
||||
block_on(async {
|
||||
let tempdir = tempdir().unwrap();
|
||||
let path = tempdir.path().join("artifact.bin");
|
||||
let body = b"artifact payload";
|
||||
let content_length = body.len().to_string();
|
||||
fs::write(&path, body).unwrap();
|
||||
|
||||
let mut server = Server::new();
|
||||
let mock = server
|
||||
.mock("POST", "/upload")
|
||||
.match_header("x-test-header", "expected-header")
|
||||
.match_header("content-length", content_length.as_str())
|
||||
.match_body(body.to_vec())
|
||||
.with_status(200)
|
||||
.create();
|
||||
|
||||
let client = http_client::Client::new_for_test();
|
||||
let target = FileArtifactUploadTargetInfo {
|
||||
url: format!("{}/upload", server.url()),
|
||||
method: "POST".to_string(),
|
||||
headers: vec![FileArtifactUploadHeaderInfo {
|
||||
name: "x-test-header".to_string(),
|
||||
value: "expected-header".to_string(),
|
||||
}],
|
||||
};
|
||||
|
||||
let checksum = upload_file_to_target(&client, &target, &path, body.len() as u64)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
mock.assert();
|
||||
assert_eq!(checksum, encode_crc32c_base64(CRC32C.checksum(body)));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn upload_file_to_target_returns_status_and_body_for_failed_uploads() {
|
||||
block_on(async {
|
||||
let tempdir = tempdir().unwrap();
|
||||
let path = tempdir.path().join("artifact.bin");
|
||||
fs::write(&path, b"artifact payload").unwrap();
|
||||
|
||||
let mut server = Server::new();
|
||||
let mock = server
|
||||
.mock("PUT", "/upload")
|
||||
.with_status(403)
|
||||
.with_body("denied")
|
||||
.create();
|
||||
|
||||
let client = http_client::Client::new_for_test();
|
||||
let target = FileArtifactUploadTargetInfo {
|
||||
url: format!("{}/upload", server.url()),
|
||||
method: "PUT".to_string(),
|
||||
headers: Vec::new(),
|
||||
};
|
||||
|
||||
let err =
|
||||
upload_file_to_target(&client, &target, &path, fs::metadata(&path).unwrap().len())
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
mock.assert();
|
||||
assert!(err
|
||||
.to_string()
|
||||
.contains("Artifact upload failed with status 403 Forbidden: denied"));
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
use super::ServerApi;
|
||||
use crate::server::graphql::{get_request_context, get_user_facing_error_message};
|
||||
use anyhow::{anyhow, Result};
|
||||
use async_trait::async_trait;
|
||||
use cynic::{MutationBuilder, QueryBuilder};
|
||||
#[cfg(test)]
|
||||
use mockall::{automock, predicate::*};
|
||||
use warp_core::channel::ChannelState;
|
||||
use warp_graphql::{
|
||||
mutations::send_referral_invite_emails::{
|
||||
SendReferralInviteEmails, SendReferralInviteEmailsResult, SendReferralInviteEmailsVariables,
|
||||
},
|
||||
queries::get_referral_info::{GetReferralInfo, GetReferralInfoVariables},
|
||||
};
|
||||
|
||||
/// Referral information for the logged-in user
|
||||
pub struct ReferralInfo {
|
||||
/// Shareable URL that the user can use to invite friends
|
||||
pub url: String,
|
||||
/// The underlying referral code associated with the user
|
||||
pub code: String,
|
||||
/// Number of other users who have signed up with this user's referral code
|
||||
pub number_claimed: usize,
|
||||
/// Whether the user has been referred by another user
|
||||
pub is_referred: bool,
|
||||
}
|
||||
|
||||
#[cfg_attr(test, automock)]
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
pub trait ReferralsClient: 'static + Send + Sync {
|
||||
/// Gets the user's referral information.
|
||||
async fn get_referral_info(&self) -> Result<ReferralInfo>;
|
||||
|
||||
/// Send one or more email invites.
|
||||
async fn send_invite(&self, emails: Vec<String>) -> Result<Vec<String>>;
|
||||
}
|
||||
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
impl ReferralsClient for ServerApi {
|
||||
async fn get_referral_info(&self) -> Result<ReferralInfo> {
|
||||
let variables = GetReferralInfoVariables {
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = GetReferralInfo::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.user {
|
||||
warp_graphql::queries::get_referral_info::UserResult::UserOutput(user_output) => {
|
||||
Ok(ReferralInfo {
|
||||
url: format!(
|
||||
"{}/referral/{}",
|
||||
ChannelState::server_root_url(),
|
||||
user_output.user.referrals.referral_code
|
||||
),
|
||||
code: user_output.user.referrals.referral_code,
|
||||
number_claimed: usize::try_from(user_output.user.referrals.number_claimed)
|
||||
.expect("Negative referral count"),
|
||||
is_referred: user_output.user.referrals.is_referred,
|
||||
})
|
||||
}
|
||||
warp_graphql::queries::get_referral_info::UserResult::Unknown => {
|
||||
Err(anyhow!("Unable to fetch referral info"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_invite(&self, emails: Vec<String>) -> Result<Vec<String>> {
|
||||
let variables = SendReferralInviteEmailsVariables {
|
||||
input: warp_graphql::mutations::send_referral_invite_emails::SendReferralInviteEmailsInput {
|
||||
emails,
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = SendReferralInviteEmails::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
let send_referral_invite_emails_result = response.send_referral_invite_emails;
|
||||
|
||||
match send_referral_invite_emails_result {
|
||||
SendReferralInviteEmailsResult::SendReferralInviteEmailsOutput(output) => {
|
||||
Ok(output.successful_emails)
|
||||
}
|
||||
SendReferralInviteEmailsResult::UserFacingError(error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(error)))
|
||||
}
|
||||
SendReferralInviteEmailsResult::Unknown => Err(anyhow!(
|
||||
"unknown error while sending referral invite emails"
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,728 @@
|
||||
use super::ServerApi;
|
||||
use crate::auth::UserUid;
|
||||
use crate::cloud_object::CloudObjectEventEntrypoint;
|
||||
use crate::workspaces::team::{DiscoverableTeam, MembershipRole};
|
||||
use crate::workspaces::workspace::Workspace;
|
||||
use anyhow::{anyhow, Result};
|
||||
use async_trait::async_trait;
|
||||
use cynic::{MutationBuilder, QueryBuilder};
|
||||
use warp_graphql::mutations::add_invite_link_domain_restriction::{
|
||||
AddInviteLinkDomainRestriction, AddInviteLinkDomainRestrictionInput,
|
||||
AddInviteLinkDomainRestrictionResult, AddInviteLinkDomainRestrictionVariables,
|
||||
};
|
||||
use warp_graphql::mutations::create_team::{
|
||||
CreateTeam, CreateTeamInput, CreateTeamResult, CreateTeamVariables,
|
||||
};
|
||||
use warp_graphql::mutations::delete_invite_link_domain_restriction::{
|
||||
DeleteInviteLinkDomainRestriction, DeleteInviteLinkDomainRestrictionInput,
|
||||
DeleteInviteLinkDomainRestrictionResult, DeleteInviteLinkDomainRestrictionVariables,
|
||||
};
|
||||
use warp_graphql::mutations::delete_team_invite::{
|
||||
DeleteTeamInvite, DeleteTeamInviteInput, DeleteTeamInviteResult, DeleteTeamInviteVariables,
|
||||
};
|
||||
use warp_graphql::mutations::join_team_with_team_discovery::{
|
||||
JoinTeamWithTeamDiscovery, JoinTeamWithTeamDiscoveryInput, JoinTeamWithTeamDiscoveryResult,
|
||||
JoinTeamWithTeamDiscoveryVariables, TeamDiscoveryEntrypoint,
|
||||
};
|
||||
use warp_graphql::mutations::remove_user_from_team::{
|
||||
RemoveUserFromTeam, RemoveUserFromTeamInput, RemoveUserFromTeamResult,
|
||||
RemoveUserFromTeamVariables,
|
||||
};
|
||||
use warp_graphql::mutations::rename_team::{
|
||||
RenameTeam, RenameTeamInput, RenameTeamResult, RenameTeamVariables,
|
||||
};
|
||||
use warp_graphql::mutations::reset_invite_links::{
|
||||
ResetInviteLinks, ResetInviteLinksInput, ResetInviteLinksResult, ResetInviteLinksVariables,
|
||||
};
|
||||
use warp_graphql::mutations::send_team_invite_email::{
|
||||
SendTeamInviteEmail, SendTeamInviteEmailInput, SendTeamInviteEmailResult,
|
||||
SendTeamInviteEmailVariables,
|
||||
};
|
||||
use warp_graphql::mutations::set_is_invite_link_enabled::{
|
||||
SetIsInviteLinkEnabled, SetIsInviteLinkEnabledInput, SetIsInviteLinkEnabledResult,
|
||||
SetIsInviteLinkEnabledVariables,
|
||||
};
|
||||
use warp_graphql::mutations::set_team_discoverability::{
|
||||
SetTeamDiscoverability, SetTeamDiscoverabilityInput, SetTeamDiscoverabilityResult,
|
||||
SetTeamDiscoverabilityVariables,
|
||||
};
|
||||
use warp_graphql::mutations::set_team_member_role::{
|
||||
SetTeamMemberRole, SetTeamMemberRoleInput, SetTeamMemberRoleResult, SetTeamMemberRoleVariables,
|
||||
};
|
||||
use warp_graphql::mutations::transfer_team_ownership::{
|
||||
TransferTeamOwnership, TransferTeamOwnershipInput, TransferTeamOwnershipResult,
|
||||
TransferTeamOwnershipVariables,
|
||||
};
|
||||
use warp_graphql::queries::get_discoverable_teams::{
|
||||
GetDiscoverableTeams, GetDiscoverableTeamsVariables,
|
||||
};
|
||||
use warp_graphql::queries::get_workspaces_metadata_for_user::{
|
||||
GetWorkspacesMetadataForUser, GetWorkspacesMetadataForUserVariables, PricingInfoResult,
|
||||
};
|
||||
|
||||
use crate::server::graphql::{get_request_context, get_user_facing_error_message};
|
||||
use crate::server::ids::ServerId;
|
||||
use crate::workspaces::user_workspaces::{CreateTeamResponse, WorkspacesMetadataWithPricing};
|
||||
|
||||
#[cfg(test)]
|
||||
use mockall::{automock, predicate::*};
|
||||
|
||||
#[cfg_attr(test, automock)]
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
pub trait TeamClient: 'static + Send + Sync {
|
||||
async fn workspaces_metadata(&self) -> Result<WorkspacesMetadataWithPricing>;
|
||||
|
||||
async fn add_invite_link_domain_restriction(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
domain: String,
|
||||
) -> Result<WorkspacesMetadataWithPricing>;
|
||||
|
||||
async fn delete_invite_link_domain_restriction(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
domain_uid: ServerId,
|
||||
) -> Result<WorkspacesMetadataWithPricing>;
|
||||
|
||||
/// Creates a team and returns the result from the server with the newly created team.
|
||||
async fn create_team(
|
||||
&self,
|
||||
name: String,
|
||||
entrypoint: CloudObjectEventEntrypoint,
|
||||
discoverable: Option<bool>,
|
||||
) -> Result<CreateTeamResponse>;
|
||||
|
||||
/// Removes the user from the selected team and returns a list of all teams that a user is
|
||||
/// still a member of (including updated team members).
|
||||
async fn remove_user_from_team(
|
||||
&self,
|
||||
user_uid: UserUid,
|
||||
team_uid: ServerId,
|
||||
entrypoint: CloudObjectEventEntrypoint,
|
||||
) -> Result<WorkspacesMetadataWithPricing>;
|
||||
|
||||
/// Removes the _current_ user from the team (user leaving the team) and returns the list of
|
||||
/// all teams that the current user is still a member of.
|
||||
async fn leave_team(
|
||||
&self,
|
||||
user_uid: UserUid,
|
||||
team_uid: ServerId,
|
||||
entrypoint: CloudObjectEventEntrypoint,
|
||||
) -> Result<WorkspacesMetadataWithPricing>;
|
||||
|
||||
async fn join_team_with_team_discovery(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
) -> Result<WorkspacesMetadataWithPricing>;
|
||||
|
||||
async fn send_team_invite_email(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
email: String,
|
||||
) -> Result<WorkspacesMetadataWithPricing>;
|
||||
|
||||
async fn delete_team_invite(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
email: String,
|
||||
) -> Result<WorkspacesMetadataWithPricing>;
|
||||
|
||||
async fn get_discoverable_teams(&self) -> Result<Vec<DiscoverableTeam>>;
|
||||
|
||||
async fn rename_team(
|
||||
&self,
|
||||
new_name: String,
|
||||
team_uid: ServerId,
|
||||
) -> Result<WorkspacesMetadataWithPricing>;
|
||||
|
||||
async fn reset_invite_links(&self, team_uid: ServerId)
|
||||
-> Result<WorkspacesMetadataWithPricing>;
|
||||
|
||||
async fn set_is_invite_link_enabled(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
new_value: bool,
|
||||
) -> Result<WorkspacesMetadataWithPricing>;
|
||||
|
||||
async fn set_team_discoverability(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
discoverable: bool,
|
||||
) -> Result<WorkspacesMetadataWithPricing>;
|
||||
|
||||
async fn transfer_team_ownership(
|
||||
&self,
|
||||
new_owner_email: String,
|
||||
) -> Result<WorkspacesMetadataWithPricing>;
|
||||
|
||||
async fn set_team_member_role(
|
||||
&self,
|
||||
user_uid: UserUid,
|
||||
team_uid: ServerId,
|
||||
role: MembershipRole,
|
||||
) -> Result<WorkspacesMetadataWithPricing>;
|
||||
}
|
||||
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
impl TeamClient for ServerApi {
|
||||
async fn workspaces_metadata(&self) -> Result<WorkspacesMetadataWithPricing> {
|
||||
let variables = GetWorkspacesMetadataForUserVariables {
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = GetWorkspacesMetadataForUser::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
let metadata = match response.user {
|
||||
warp_graphql::queries::get_workspaces_metadata_for_user::UserResult::UserOutput(
|
||||
user_output,
|
||||
) => user_output.user.into(),
|
||||
warp_graphql::queries::get_workspaces_metadata_for_user::UserResult::Unknown => {
|
||||
return Err(anyhow!("Unable to fetch workspaces metadata"));
|
||||
}
|
||||
};
|
||||
|
||||
let pricing_info = match response.pricing_info {
|
||||
PricingInfoResult::PricingInfoOutput(pricing_output) => {
|
||||
Some(pricing_output.pricing_info)
|
||||
}
|
||||
PricingInfoResult::Unknown => None,
|
||||
};
|
||||
|
||||
Ok(WorkspacesMetadataWithPricing {
|
||||
metadata,
|
||||
pricing_info,
|
||||
})
|
||||
}
|
||||
|
||||
async fn add_invite_link_domain_restriction(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
domain: String,
|
||||
) -> Result<WorkspacesMetadataWithPricing> {
|
||||
let variables = AddInviteLinkDomainRestrictionVariables {
|
||||
input: AddInviteLinkDomainRestrictionInput {
|
||||
team_uid: team_uid.into(),
|
||||
domain,
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = AddInviteLinkDomainRestriction::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.add_invite_link_domain_restriction;
|
||||
|
||||
match result {
|
||||
AddInviteLinkDomainRestrictionResult::AddInviteLinkDomainRestrictionOutput(result) => {
|
||||
if !result.success {
|
||||
return Err(anyhow!("failed to add invite link domain restriction"));
|
||||
}
|
||||
}
|
||||
AddInviteLinkDomainRestrictionResult::UserFacingError(user_facing_error) => {
|
||||
return Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
AddInviteLinkDomainRestrictionResult::Unknown => {
|
||||
return Err(anyhow!(
|
||||
"unknown error while adding invite link domain restriction"
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
self.workspaces_metadata().await
|
||||
}
|
||||
|
||||
async fn delete_invite_link_domain_restriction(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
domain_uid: ServerId,
|
||||
) -> Result<WorkspacesMetadataWithPricing> {
|
||||
let variables = DeleteInviteLinkDomainRestrictionVariables {
|
||||
input: DeleteInviteLinkDomainRestrictionInput {
|
||||
uid: domain_uid.into(),
|
||||
team_uid: team_uid.into(),
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = DeleteInviteLinkDomainRestriction::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.delete_invite_link_domain_restriction;
|
||||
|
||||
match result {
|
||||
DeleteInviteLinkDomainRestrictionResult::DeleteInviteLinkDomainRestrictionOutput(
|
||||
result,
|
||||
) => {
|
||||
if !result.success {
|
||||
return Err(anyhow!("failed to delete invite link domain restriction"));
|
||||
}
|
||||
}
|
||||
DeleteInviteLinkDomainRestrictionResult::UserFacingError(user_facing_error) => {
|
||||
return Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
DeleteInviteLinkDomainRestrictionResult::Unknown => {
|
||||
return Err(anyhow!(
|
||||
"unknown error while deleting invite link domain restriction"
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
self.workspaces_metadata().await
|
||||
}
|
||||
|
||||
async fn create_team(
|
||||
&self,
|
||||
name: String,
|
||||
entrypoint: CloudObjectEventEntrypoint,
|
||||
discoverable: Option<bool>,
|
||||
) -> Result<CreateTeamResponse> {
|
||||
let variables = CreateTeamVariables {
|
||||
input: CreateTeamInput {
|
||||
name,
|
||||
entrypoint: entrypoint.into(),
|
||||
discoverable: discoverable.unwrap_or(false),
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = CreateTeam::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.create_team;
|
||||
|
||||
match result {
|
||||
CreateTeamResult::CreateTeamOutput(output) => {
|
||||
let workspace: Workspace = output.workspace.clone().into();
|
||||
if let Some(team) = workspace.teams.first() {
|
||||
Ok(CreateTeamResponse {
|
||||
workspace: workspace.clone(),
|
||||
team: team.clone(),
|
||||
})
|
||||
} else {
|
||||
Err(anyhow!("failed to create team"))
|
||||
}
|
||||
}
|
||||
CreateTeamResult::UserFacingError(user_facing_error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
CreateTeamResult::Unknown => Err(anyhow!("unknown error while creating team")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn remove_user_from_team(
|
||||
&self,
|
||||
user_uid: UserUid,
|
||||
team_uid: ServerId,
|
||||
entrypoint: CloudObjectEventEntrypoint,
|
||||
) -> Result<WorkspacesMetadataWithPricing> {
|
||||
let variables = RemoveUserFromTeamVariables {
|
||||
input: RemoveUserFromTeamInput {
|
||||
user_uid: user_uid.as_str().into(),
|
||||
team_uid: team_uid.into(),
|
||||
entrypoint: entrypoint.into(),
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = RemoveUserFromTeam::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.remove_user_from_team;
|
||||
|
||||
match result {
|
||||
RemoveUserFromTeamResult::RemoveUserFromTeamOutput(output) => {
|
||||
if !output.success {
|
||||
return Err(anyhow!("failed to remove user from team"));
|
||||
} else {
|
||||
self.workspaces_metadata().await
|
||||
}
|
||||
}
|
||||
RemoveUserFromTeamResult::UserFacingError(user_facing_error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
RemoveUserFromTeamResult::Unknown => {
|
||||
Err(anyhow!("unknown error while removing user from team"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn leave_team(
|
||||
&self,
|
||||
user_uid: UserUid,
|
||||
team_uid: ServerId,
|
||||
entrypoint: CloudObjectEventEntrypoint,
|
||||
) -> Result<WorkspacesMetadataWithPricing> {
|
||||
let variables = RemoveUserFromTeamVariables {
|
||||
input: RemoveUserFromTeamInput {
|
||||
user_uid: user_uid.into(),
|
||||
team_uid: team_uid.into(),
|
||||
entrypoint: entrypoint.into(),
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = RemoveUserFromTeam::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.remove_user_from_team;
|
||||
|
||||
match result {
|
||||
RemoveUserFromTeamResult::RemoveUserFromTeamOutput(output) => {
|
||||
if !output.success {
|
||||
return Err(anyhow!("failed to leave team"));
|
||||
} else {
|
||||
self.workspaces_metadata().await
|
||||
}
|
||||
}
|
||||
RemoveUserFromTeamResult::UserFacingError(user_facing_error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
RemoveUserFromTeamResult::Unknown => Err(anyhow!("unknown error while leaving team")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn join_team_with_team_discovery(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
) -> Result<WorkspacesMetadataWithPricing> {
|
||||
let variables = JoinTeamWithTeamDiscoveryVariables {
|
||||
input: JoinTeamWithTeamDiscoveryInput {
|
||||
team_uid: team_uid.into(),
|
||||
entrypoint: TeamDiscoveryEntrypoint::TeamSettings,
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = JoinTeamWithTeamDiscovery::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.join_team_with_team_discovery;
|
||||
|
||||
match result {
|
||||
JoinTeamWithTeamDiscoveryResult::JoinTeamWithTeamDiscoveryOutput(output) => {
|
||||
if !output.success {
|
||||
return Err(anyhow!("failed to join team"));
|
||||
} else {
|
||||
self.workspaces_metadata().await
|
||||
}
|
||||
}
|
||||
JoinTeamWithTeamDiscoveryResult::UserFacingError(user_facing_error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
JoinTeamWithTeamDiscoveryResult::Unknown => {
|
||||
Err(anyhow!("unknown error while joining team"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_team_invite_email(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
email: String,
|
||||
) -> Result<WorkspacesMetadataWithPricing> {
|
||||
let variables = SendTeamInviteEmailVariables {
|
||||
input: SendTeamInviteEmailInput {
|
||||
team_uid: team_uid.into(),
|
||||
email,
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = SendTeamInviteEmail::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.send_team_invite_email;
|
||||
|
||||
match result {
|
||||
SendTeamInviteEmailResult::SendTeamInviteEmailOutput(output) => {
|
||||
if !output.success {
|
||||
return Err(anyhow!("failed to send team invite"));
|
||||
} else {
|
||||
self.workspaces_metadata().await
|
||||
}
|
||||
}
|
||||
SendTeamInviteEmailResult::UserFacingError(user_facing_error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
SendTeamInviteEmailResult::Unknown => {
|
||||
Err(anyhow!("unknown error while sending team invite"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_team_invite(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
email: String,
|
||||
) -> Result<WorkspacesMetadataWithPricing> {
|
||||
let variables = DeleteTeamInviteVariables {
|
||||
input: DeleteTeamInviteInput {
|
||||
team_uid: team_uid.into(),
|
||||
email,
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = DeleteTeamInvite::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.delete_team_invite;
|
||||
|
||||
match result {
|
||||
DeleteTeamInviteResult::DeleteTeamInviteOutput(output) => {
|
||||
if !output.success {
|
||||
return Err(anyhow!("failed to delete team invite"));
|
||||
} else {
|
||||
self.workspaces_metadata().await
|
||||
}
|
||||
}
|
||||
DeleteTeamInviteResult::UserFacingError(user_facing_error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
DeleteTeamInviteResult::Unknown => {
|
||||
Err(anyhow!("unknown error while deleting team invite"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_discoverable_teams(&self) -> Result<Vec<DiscoverableTeam>, anyhow::Error> {
|
||||
let variables = GetDiscoverableTeamsVariables {
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = GetDiscoverableTeams::build(variables);
|
||||
let result = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match result.user {
|
||||
warp_graphql::queries::get_discoverable_teams::UserResult::UserOutput(user_output) => {
|
||||
Ok(user_output
|
||||
.user
|
||||
.discoverable_teams
|
||||
.into_iter()
|
||||
.map(|gql_team_data| Ok(gql_team_data.into()))
|
||||
.collect::<Result<Vec<DiscoverableTeam>>>()?)
|
||||
}
|
||||
warp_graphql::queries::get_discoverable_teams::UserResult::UserFacingError(
|
||||
user_facing_error,
|
||||
) => Err(anyhow!(get_user_facing_error_message(user_facing_error))),
|
||||
warp_graphql::queries::get_discoverable_teams::UserResult::Unknown => {
|
||||
Err(anyhow!("unknown error while getting discoverable teams"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn rename_team(
|
||||
&self,
|
||||
new_name: String,
|
||||
team_uid: ServerId,
|
||||
) -> Result<WorkspacesMetadataWithPricing> {
|
||||
let variables = RenameTeamVariables {
|
||||
input: RenameTeamInput {
|
||||
new_name,
|
||||
team_uid: team_uid.into(),
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = RenameTeam::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.rename_team;
|
||||
|
||||
match result {
|
||||
RenameTeamResult::RenameTeamOutput(output) => {
|
||||
if output.success {
|
||||
self.workspaces_metadata().await
|
||||
} else {
|
||||
Err(anyhow!("failed to rename team"))
|
||||
}
|
||||
}
|
||||
RenameTeamResult::UserFacingError(user_facing_error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
RenameTeamResult::Unknown => Err(anyhow!("unknown error while renaming team")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn reset_invite_links(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
) -> Result<WorkspacesMetadataWithPricing> {
|
||||
let variables = ResetInviteLinksVariables {
|
||||
input: ResetInviteLinksInput {
|
||||
team_uid: team_uid.into(),
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = ResetInviteLinks::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.reset_invite_links;
|
||||
|
||||
match result {
|
||||
ResetInviteLinksResult::ResetInviteLinksOutput(output) => {
|
||||
if output.success {
|
||||
self.workspaces_metadata().await
|
||||
} else {
|
||||
Err(anyhow!("failed to reset invite links"))
|
||||
}
|
||||
}
|
||||
ResetInviteLinksResult::UserFacingError(user_facing_error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
ResetInviteLinksResult::Unknown => {
|
||||
Err(anyhow!("unknown error while resetting invite links"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn set_is_invite_link_enabled(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
new_value: bool,
|
||||
) -> Result<WorkspacesMetadataWithPricing> {
|
||||
let variables = SetIsInviteLinkEnabledVariables {
|
||||
input: SetIsInviteLinkEnabledInput {
|
||||
team_uid: team_uid.into(),
|
||||
new_value,
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = SetIsInviteLinkEnabled::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.set_is_invite_link_enabled;
|
||||
|
||||
match result {
|
||||
SetIsInviteLinkEnabledResult::SetIsInviteLinkEnabledOutput(output) => {
|
||||
if output.success {
|
||||
self.workspaces_metadata().await
|
||||
} else {
|
||||
Err(anyhow!("failed to set invite link enabled"))
|
||||
}
|
||||
}
|
||||
SetIsInviteLinkEnabledResult::UserFacingError(user_facing_error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
SetIsInviteLinkEnabledResult::Unknown => {
|
||||
Err(anyhow!("unknown error while setting invite link enabled"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn set_team_discoverability(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
new_value: bool,
|
||||
) -> Result<WorkspacesMetadataWithPricing> {
|
||||
let variables = SetTeamDiscoverabilityVariables {
|
||||
input: SetTeamDiscoverabilityInput {
|
||||
team_uid: team_uid.into(),
|
||||
discoverable: new_value,
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
|
||||
let operation = SetTeamDiscoverability::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.set_team_discoverability;
|
||||
|
||||
match result {
|
||||
SetTeamDiscoverabilityResult::SetTeamDiscoverabilityOutput(output) => {
|
||||
if output.success {
|
||||
self.workspaces_metadata().await
|
||||
} else {
|
||||
Err(anyhow!("failed to set team discoverability"))
|
||||
}
|
||||
}
|
||||
SetTeamDiscoverabilityResult::UserFacingError(user_facing_error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
SetTeamDiscoverabilityResult::Unknown => {
|
||||
Err(anyhow!("unknown error while setting team discoverability"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn transfer_team_ownership(
|
||||
&self,
|
||||
new_owner_email: String,
|
||||
) -> Result<WorkspacesMetadataWithPricing> {
|
||||
let variables = TransferTeamOwnershipVariables {
|
||||
input: TransferTeamOwnershipInput { new_owner_email },
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = TransferTeamOwnership::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.transfer_team_ownership;
|
||||
|
||||
match result {
|
||||
TransferTeamOwnershipResult::TransferTeamOwnershipOutput(output) => {
|
||||
if !output.success {
|
||||
return Err(anyhow!("failed to transfer team ownership"));
|
||||
} else {
|
||||
self.workspaces_metadata().await
|
||||
}
|
||||
}
|
||||
TransferTeamOwnershipResult::UserFacingError(user_facing_error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
TransferTeamOwnershipResult::Unknown => {
|
||||
Err(anyhow!("unknown error while transferring team ownership"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn set_team_member_role(
|
||||
&self,
|
||||
user_uid: UserUid,
|
||||
team_uid: ServerId,
|
||||
role: MembershipRole,
|
||||
) -> Result<WorkspacesMetadataWithPricing> {
|
||||
let variables = SetTeamMemberRoleVariables {
|
||||
input: SetTeamMemberRoleInput {
|
||||
user_uid: user_uid.as_str().into(),
|
||||
team_uid: team_uid.into(),
|
||||
role: role.into(),
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = SetTeamMemberRole::build(variables);
|
||||
let result = self
|
||||
.send_graphql_request(operation, None)
|
||||
.await?
|
||||
.set_team_member_role;
|
||||
|
||||
match result {
|
||||
SetTeamMemberRoleResult::SetTeamMemberRoleOutput(output) => {
|
||||
if output.success {
|
||||
self.workspaces_metadata().await
|
||||
} else {
|
||||
Err(anyhow!("failed to set team member role"))
|
||||
}
|
||||
}
|
||||
SetTeamMemberRoleResult::UserFacingError(user_facing_error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(user_facing_error)))
|
||||
}
|
||||
SetTeamMemberRoleResult::Unknown => {
|
||||
Err(anyhow!("unknown error while setting team member role"))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,223 @@
|
||||
use super::{team::TeamClient, ServerApi};
|
||||
use crate::workspaces::user_workspaces::WorkspacesMetadataResponse;
|
||||
use crate::workspaces::workspace::AiOverages;
|
||||
use anyhow::{anyhow, Result};
|
||||
use async_trait::async_trait;
|
||||
use cynic::{MutationBuilder, QueryBuilder};
|
||||
use warp_graphql::error::UserFacingErrorInterface;
|
||||
use warp_graphql::mutations::purchase_addon_credits::{
|
||||
PurchaseAddonCredits, PurchaseAddonCreditsInput, PurchaseAddonCreditsResult,
|
||||
PurchaseAddonCreditsVariables,
|
||||
};
|
||||
use warp_graphql::mutations::stripe_billing_portal::{
|
||||
StripeBillingPortal, StripeBillingPortalInput, StripeBillingPortalResult,
|
||||
StripeBillingPortalVariables,
|
||||
};
|
||||
use warp_graphql::mutations::update_workspace_settings::{
|
||||
AddonCreditsSettingsInput, UpdateWorkspaceSettings, UpdateWorkspaceSettingsInput,
|
||||
UpdateWorkspaceSettingsResult, UpdateWorkspaceSettingsVariables,
|
||||
UsageBasedPricingSettingsInput,
|
||||
};
|
||||
use warp_graphql::queries::get_ai_overages_for_workspace::{
|
||||
GetAiOveragesForWorkspace, GetAiOveragesForWorkspaceVariables, UserResult,
|
||||
};
|
||||
|
||||
use crate::server::graphql::{get_request_context, get_user_facing_error_message};
|
||||
use crate::server::ids::ServerId;
|
||||
|
||||
#[cfg(test)]
|
||||
use mockall::{automock, predicate::*};
|
||||
|
||||
#[cfg_attr(test, automock)]
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
pub trait WorkspaceClient: 'static + Send + Sync {
|
||||
async fn generate_stripe_billing_portal_link(&self, team_uid: ServerId) -> Result<String>;
|
||||
|
||||
async fn update_usage_based_pricing_settings(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
usage_based_pricing_enabled: bool,
|
||||
max_monthly_spend_cents: Option<u32>,
|
||||
) -> Result<WorkspacesMetadataResponse>;
|
||||
|
||||
async fn refresh_ai_overages(&self) -> Result<AiOverages>;
|
||||
|
||||
async fn purchase_addon_credits(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
credits: i32,
|
||||
) -> Result<WorkspacesMetadataResponse>;
|
||||
|
||||
async fn update_addon_credits_settings(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
auto_reload_enabled: Option<bool>,
|
||||
max_monthly_spend_cents: Option<i32>,
|
||||
selected_auto_reload_credit_denomination: Option<i32>,
|
||||
) -> Result<WorkspacesMetadataResponse>;
|
||||
}
|
||||
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
impl WorkspaceClient for ServerApi {
|
||||
async fn generate_stripe_billing_portal_link(&self, team_uid: ServerId) -> Result<String> {
|
||||
let variables = StripeBillingPortalVariables {
|
||||
input: StripeBillingPortalInput {
|
||||
team_uid: team_uid.into(),
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = StripeBillingPortal::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.stripe_billing_portal {
|
||||
StripeBillingPortalResult::StripeBillingPortalOutput(output) => Ok(output.url),
|
||||
StripeBillingPortalResult::UserFacingError(error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(error)))
|
||||
}
|
||||
StripeBillingPortalResult::Unknown => Err(anyhow!("Unknown error")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn update_usage_based_pricing_settings(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
usage_based_pricing_enabled: bool,
|
||||
max_monthly_spend_cents: Option<u32>,
|
||||
) -> Result<WorkspacesMetadataResponse> {
|
||||
if let Some(cents) = max_monthly_spend_cents {
|
||||
if cents > i32::MAX as u32 {
|
||||
return Err(anyhow!(
|
||||
"Maximum monthly spend cannot exceed {} cents",
|
||||
i32::MAX
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
let variables = UpdateWorkspaceSettingsVariables {
|
||||
input: UpdateWorkspaceSettingsInput {
|
||||
workspace_uid: team_uid.to_string(),
|
||||
set_usage_based_pricing_settings: Some(UsageBasedPricingSettingsInput {
|
||||
enabled: Some(usage_based_pricing_enabled),
|
||||
max_monthly_spend_cents: max_monthly_spend_cents.map(|cents| cents as i32),
|
||||
}),
|
||||
set_addon_credits_settings: None,
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = UpdateWorkspaceSettings::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.update_workspace_settings {
|
||||
UpdateWorkspaceSettingsResult::UpdateWorkspaceSettingsOutput(_) => {
|
||||
TeamClient::workspaces_metadata(self)
|
||||
.await
|
||||
.map(|w| w.metadata)
|
||||
}
|
||||
UpdateWorkspaceSettingsResult::UserFacingError(error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(error)))
|
||||
}
|
||||
UpdateWorkspaceSettingsResult::Unknown => Err(anyhow!("Unknown error")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn refresh_ai_overages(&self) -> Result<AiOverages> {
|
||||
let variables = GetAiOveragesForWorkspaceVariables {
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = GetAiOveragesForWorkspace::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.user {
|
||||
UserResult::UserOutput(user_output) => user_output
|
||||
.user
|
||||
.workspaces
|
||||
.first()
|
||||
.as_ref()
|
||||
.ok_or_else(|| anyhow!("No workspace found"))?
|
||||
.billing_metadata
|
||||
.ai_overages
|
||||
.as_ref()
|
||||
.ok_or_else(|| anyhow!("No AI overages found"))
|
||||
.map(|overages| AiOverages {
|
||||
current_monthly_request_cost_cents: overages.current_monthly_request_cost_cents,
|
||||
current_monthly_requests_used: overages.current_monthly_requests_used,
|
||||
current_period_end: overages.current_period_end.utc(),
|
||||
}),
|
||||
UserResult::Unknown => Err(anyhow!("Unknown error")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn purchase_addon_credits(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
credits: i32,
|
||||
) -> Result<WorkspacesMetadataResponse> {
|
||||
let variables = PurchaseAddonCreditsVariables {
|
||||
input: PurchaseAddonCreditsInput {
|
||||
team_uid: team_uid.into(),
|
||||
credits,
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = PurchaseAddonCredits::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await;
|
||||
|
||||
match response {
|
||||
Err(_) => Err(anyhow!("Failed to purchase add-on credits")),
|
||||
Ok(response) => match response.purchase_addon_credits {
|
||||
PurchaseAddonCreditsResult::PurchaseAddonCreditsOutput(_) => {
|
||||
TeamClient::workspaces_metadata(self)
|
||||
.await
|
||||
.map(|w| w.metadata)
|
||||
}
|
||||
PurchaseAddonCreditsResult::UserFacingError(error) => match error.error {
|
||||
UserFacingErrorInterface::BudgetExceededError(budget_error) => {
|
||||
Err(budget_error.into())
|
||||
}
|
||||
UserFacingErrorInterface::PaymentMethodDeclinedError(
|
||||
payment_declined_error,
|
||||
) => Err(payment_declined_error.into()),
|
||||
_ => Err(anyhow!(get_user_facing_error_message(error))),
|
||||
},
|
||||
PurchaseAddonCreditsResult::Unknown => Err(anyhow!("Unknown error")),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
async fn update_addon_credits_settings(
|
||||
&self,
|
||||
team_uid: ServerId,
|
||||
auto_reload_enabled: Option<bool>,
|
||||
max_monthly_spend_cents: Option<i32>,
|
||||
selected_auto_reload_credit_denomination: Option<i32>,
|
||||
) -> Result<WorkspacesMetadataResponse> {
|
||||
let variables = UpdateWorkspaceSettingsVariables {
|
||||
input: UpdateWorkspaceSettingsInput {
|
||||
workspace_uid: team_uid.to_string(),
|
||||
set_usage_based_pricing_settings: None,
|
||||
set_addon_credits_settings: Some(AddonCreditsSettingsInput {
|
||||
auto_reload_enabled,
|
||||
max_monthly_spend_cents,
|
||||
selected_auto_reload_credit_denomination,
|
||||
}),
|
||||
},
|
||||
request_context: get_request_context(),
|
||||
};
|
||||
let operation = UpdateWorkspaceSettings::build(variables);
|
||||
let response = self.send_graphql_request(operation, None).await?;
|
||||
|
||||
match response.update_workspace_settings {
|
||||
UpdateWorkspaceSettingsResult::UpdateWorkspaceSettingsOutput(_) => {
|
||||
TeamClient::workspaces_metadata(self)
|
||||
.await
|
||||
.map(|w| w.metadata)
|
||||
}
|
||||
UpdateWorkspaceSettingsResult::UserFacingError(error) => {
|
||||
Err(anyhow!(get_user_facing_error_message(error)))
|
||||
}
|
||||
UpdateWorkspaceSettingsResult::Unknown => Err(anyhow!("Unknown error")),
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,31 @@
|
||||
The MIT License (MIT)
|
||||
|
||||
Copyright 2019 Segment.io, Inc.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
||||
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2021 RudderStack
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -0,0 +1,232 @@
|
||||
use std::sync::Arc;
|
||||
use std::{fs::remove_file, time::Duration};
|
||||
|
||||
use anyhow::Context;
|
||||
use chrono::{LocalResult, TimeZone, Utc};
|
||||
use warp_core::execution_mode::AppExecutionMode;
|
||||
use warp_core::{report_error, report_if_error};
|
||||
use warpui::r#async::{FutureExt as _, Timer};
|
||||
use warpui::{App, Entity, ModelContext, SingletonEntity};
|
||||
|
||||
use super::{rudder_event_file_path, RUDDER_TELEMETRY_EVENTS_FILE_NAME};
|
||||
use crate::auth::AuthStateProvider;
|
||||
use crate::channel::ChannelState;
|
||||
use crate::features::FeatureFlag;
|
||||
use crate::{
|
||||
server::server_api::ServerApi,
|
||||
settings::{PrivacySettings, PrivacySettingsChangedEvent},
|
||||
};
|
||||
|
||||
use super::clear_event_queue;
|
||||
|
||||
// How often we send Active Usage signals.
|
||||
const ACTIVE_USAGE_DURATION: Duration = Duration::from_secs(60);
|
||||
|
||||
/// Duration to wait before flushing the event queue to Rudderstack.
|
||||
const TELEMETRY_FLUSH_DURATION: Duration = Duration::from_secs(30);
|
||||
|
||||
/// Max telemetry events to write to disk. This is bounded to limit the size of the file as well
|
||||
/// as latency of writing the file.
|
||||
const MAX_TELEMETRY_EVENTS_TO_STORE: usize = 20;
|
||||
|
||||
/// Maximum time to wait for the telemetry flush network request during shutdown.
|
||||
/// If the network is unavailable or slow, we don't want the CLI process to hang indefinitely.
|
||||
const TELEMETRY_SHUTDOWN_FLUSH_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
|
||||
/// App singleton responsible for scheduling periodic background tasks for sending batches of
|
||||
/// telemetry events to Rudderstack. This model respects the user's telemetry enablement setting.
|
||||
pub struct TelemetryCollector {
|
||||
server_api: Arc<ServerApi>,
|
||||
}
|
||||
|
||||
impl TelemetryCollector {
|
||||
pub fn new(server_api: Arc<ServerApi>) -> Self {
|
||||
Self { server_api }
|
||||
}
|
||||
|
||||
pub fn initialize_telemetry_collection(&self, ctx: &mut ModelContext<TelemetryCollector>) {
|
||||
// Start a background thread to periodically flush events from the telemetry event queue.
|
||||
if ChannelState::is_release_bundle() || FeatureFlag::WithSandboxTelemetry.is_enabled() {
|
||||
// Flush the events to Rudderstack that were persisted into a file the last time the app was
|
||||
// quit.
|
||||
self.flush_persisted_events_from_disk(ctx);
|
||||
}
|
||||
|
||||
// Send Active App Usage signals
|
||||
if FeatureFlag::RecordAppActiveEvents.is_enabled()
|
||||
&& (ChannelState::is_release_bundle() || FeatureFlag::WithSandboxTelemetry.is_enabled())
|
||||
{
|
||||
self.schedule_send_active_usage_event(ctx);
|
||||
}
|
||||
|
||||
// Start a background thread to periodically flush events from the telemetry event queue.
|
||||
if ChannelState::is_release_bundle()
|
||||
|| FeatureFlag::WithSandboxTelemetry.is_enabled()
|
||||
|| FeatureFlag::SendTelemetryToFile.is_enabled()
|
||||
{
|
||||
self.schedule_event_queue_flush(ctx);
|
||||
}
|
||||
|
||||
// Clear queued telemetry events when telemetry is enabled or disabled. If telemetry is
|
||||
// enabled, we will start sending Rudderstack requests when the event queue is periodically
|
||||
// flushed. The initial request should not contain any events recorded when the user was
|
||||
// previously opted-out of telemetry. In the case where the user turns the telemetry from
|
||||
// on to off, we should not send another request with any telemetry, even if the event was
|
||||
// initially recorded prior to the user turning telemetry off.`
|
||||
ctx.subscribe_to_model(&PrivacySettings::handle(ctx), |_me, event, _ctx| {
|
||||
if let PrivacySettingsChangedEvent::UpdateIsTelemetryEnabled { .. } = event {
|
||||
clear_event_queue();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Writes all queued but unsent telemetry telemetry events to disk so that they may be sent
|
||||
/// on the next app startup.
|
||||
pub fn write_telemetry_events_to_disk(&self, ctx: &mut ModelContext<TelemetryCollector>) {
|
||||
match self.server_api.persist_telemetry_events(
|
||||
MAX_TELEMETRY_EVENTS_TO_STORE,
|
||||
PrivacySettings::as_ref(ctx).get_snapshot(ctx),
|
||||
) {
|
||||
Ok(()) => {
|
||||
log::info!("Successfully wrote telemetry events to disk")
|
||||
}
|
||||
Err(e) => {
|
||||
log::error!("Failed to write telemetry events to disk {e:#}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Flushes telemetry events when the app is shutting down.
|
||||
///
|
||||
/// Depending on the app's execution mode, this will either:
|
||||
/// * Write events to disk, for sending on the next app startup
|
||||
/// * Synchronously send events to rudderstack
|
||||
pub fn flush_telemetry_events_for_shutdown(&self, ctx: &mut ModelContext<TelemetryCollector>) {
|
||||
let execution_mode = AppExecutionMode::as_ref(ctx);
|
||||
|
||||
if execution_mode.send_telemetry_at_shutdown() {
|
||||
let privacy_settings_snapshot = PrivacySettings::as_ref(ctx).get_snapshot(ctx);
|
||||
let server_api = self.server_api.clone();
|
||||
match warpui::r#async::block_on(async move {
|
||||
server_api
|
||||
.flush_telemetry_events(privacy_settings_snapshot)
|
||||
.with_timeout(TELEMETRY_SHUTDOWN_FLUSH_TIMEOUT)
|
||||
.await
|
||||
}) {
|
||||
Ok(Ok(count)) => {
|
||||
if count > 0 {
|
||||
log::info!("Successfully flushed telemetry events before shutdown");
|
||||
}
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
report_error!(e.context("Error flushing telemetry events before shutdown"));
|
||||
}
|
||||
Err(_) => {
|
||||
log::warn!(
|
||||
"Telemetry flush timed out after {}s during shutdown, skipping",
|
||||
TELEMETRY_SHUTDOWN_FLUSH_TIMEOUT.as_secs()
|
||||
);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
self.write_telemetry_events_to_disk(ctx);
|
||||
}
|
||||
}
|
||||
|
||||
/// Sends rudderstack requests containing events persisted to disk (if telemetry is enabled).
|
||||
/// Events may be written to disk at the end of a session prior to app termination; this
|
||||
/// function should be called on startup to track events that were recorded at the end of the
|
||||
/// last session and were not flushed.
|
||||
fn flush_persisted_events_from_disk(&self, ctx: &mut ModelContext<TelemetryCollector>) {
|
||||
let privacy_settings_snapshot = PrivacySettings::as_ref(ctx).get_snapshot(ctx);
|
||||
let server_api = self.server_api.clone();
|
||||
let _ = ctx.spawn(
|
||||
async move {
|
||||
let new_path = rudder_event_file_path();
|
||||
let old_path =
|
||||
warp_core::paths::state_dir().join(RUDDER_TELEMETRY_EVENTS_FILE_NAME);
|
||||
|
||||
// Try flushing from both new and legacy locations.
|
||||
for path in [new_path, old_path] {
|
||||
report_if_error!(server_api
|
||||
.flush_persisted_events_to_rudder(&path, privacy_settings_snapshot)
|
||||
.await
|
||||
.context("Failed to flush rudder events from disk"));
|
||||
// Remove the file regardless of outcome of flushing the events to avoid the
|
||||
// case where we accidentally try to re-flush the events on the next app startup.
|
||||
if let Err(e) = remove_file(&path) {
|
||||
if e.kind() != std::io::ErrorKind::NotFound {
|
||||
warp_core::report_error!(
|
||||
anyhow::anyhow!(e).context("Failed to remove persisted event file")
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
|_, _, _| (),
|
||||
);
|
||||
}
|
||||
|
||||
/// Schedules a background task to send an active usage event in a rudderstack request if
|
||||
/// telemetry is enabled. The scheduled task once again schedules itself after
|
||||
/// `ACTIVE_USAGE_DURATION`.
|
||||
fn schedule_send_active_usage_event(&self, ctx: &mut ModelContext<TelemetryCollector>) {
|
||||
let auth_state = AuthStateProvider::as_ref(ctx).get().clone();
|
||||
let is_telemetry_enabled = PrivacySettings::as_ref(ctx).is_telemetry_enabled;
|
||||
let _ = ctx.spawn(
|
||||
async move {
|
||||
// Record app active if there was any activity now or right after the previous check
|
||||
let last_active_timestamp = App::last_active_timestamp();
|
||||
if is_telemetry_enabled
|
||||
&& last_active_timestamp + ACTIVE_USAGE_DURATION.as_secs() as i64
|
||||
> Utc::now().timestamp()
|
||||
{
|
||||
if let LocalResult::Single(timestamp) =
|
||||
Utc.timestamp_opt(last_active_timestamp, 0)
|
||||
{
|
||||
warpui::telemetry::record_app_active_event(
|
||||
auth_state.user_id().map(|uid| uid.as_string()),
|
||||
auth_state.anonymous_id(),
|
||||
timestamp,
|
||||
);
|
||||
}
|
||||
}
|
||||
Timer::after(ACTIVE_USAGE_DURATION).await;
|
||||
},
|
||||
|me, _, ctx| me.schedule_send_active_usage_event(ctx),
|
||||
);
|
||||
}
|
||||
|
||||
/// Flushes events from the in-memory event queue and schedules a background task to send
|
||||
/// them in rudderstack request if telemetry is enabled. The scheduled task once again schedules
|
||||
/// itself after `TELEMETRY_FLUSH_DURATION`.
|
||||
fn schedule_event_queue_flush(&self, ctx: &mut ModelContext<TelemetryCollector>) {
|
||||
let server_api = self.server_api.clone();
|
||||
let privacy_settings_snapshot = PrivacySettings::as_ref(ctx).get_snapshot(ctx);
|
||||
let _ = ctx.spawn(
|
||||
async move {
|
||||
match server_api
|
||||
.flush_telemetry_events(privacy_settings_snapshot)
|
||||
.await
|
||||
{
|
||||
Ok(count) => {
|
||||
if count > 0 {
|
||||
log::debug!("Flushed telemetry events.");
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
log::info!("Failed to flush events from Telemetry queue: {e}");
|
||||
}
|
||||
}
|
||||
Timer::after(TELEMETRY_FLUSH_DURATION).await;
|
||||
},
|
||||
|me, _, ctx| me.schedule_event_queue_flush(ctx),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
impl Entity for TelemetryCollector {
|
||||
type Event = ();
|
||||
}
|
||||
|
||||
impl SingletonEntity for TelemetryCollector {}
|
||||
@@ -0,0 +1,100 @@
|
||||
//! Module that builds a static context to attach to each of our events that are sent to Rudderstack.
|
||||
//! This is needed so we know the backing operating system and version of each telemetry event.
|
||||
|
||||
use super::rudder_message::Message as RudderMessage;
|
||||
use crate::server::OperatingSystemInfo;
|
||||
|
||||
use serde::Serialize;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use std::sync::OnceLock;
|
||||
|
||||
#[cfg(target_family = "wasm")]
|
||||
use warpui::platform::wasm;
|
||||
|
||||
static TELEMETRY_CONTEXT: OnceLock<TelemetryContext> = OnceLock::new();
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct TelemetryContextInfo {
|
||||
/// Info about the operating system of the client.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
os: Option<&'static OperatingSystemInfo>,
|
||||
/// The user agent provided by the browser, if running on Web. If not on
|
||||
/// Web, this is always `None`.
|
||||
#[serde(rename = "userAgent", skip_serializing_if = "Option::is_none")]
|
||||
user_agent: Option<String>,
|
||||
}
|
||||
|
||||
/// Newtype representing a [`Value`] with a serialized version of the context that we send to
|
||||
/// Rudderstack.
|
||||
/// See https://www.rudderstack.com/docs/event-spec/standard-events/common-fields/#contextual-fields.
|
||||
pub struct TelemetryContext(Value);
|
||||
|
||||
impl TelemetryContext {
|
||||
pub fn as_value(&self) -> Value {
|
||||
self.0.clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl TelemetryContext {
|
||||
fn new() -> Self {
|
||||
let context = TelemetryContextInfo {
|
||||
os: OperatingSystemInfo::get().ok(),
|
||||
user_agent: user_agent(),
|
||||
};
|
||||
|
||||
match serde_json::to_value(context) {
|
||||
Ok(value) => Self(value),
|
||||
Err(e) => {
|
||||
log::error!("Failed to serialize telemetry context info to JSON value: {e:?}");
|
||||
Self(json!({}))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Extension trait used to attach a telemetry context.
|
||||
pub(super) trait AttachContext {
|
||||
/// Attaches a context to the given object.
|
||||
fn attach_context(&mut self);
|
||||
}
|
||||
|
||||
impl AttachContext for RudderMessage {
|
||||
/// Attaches the context to the [`RudderMessage`]. Note this is currently last write wins; if a
|
||||
/// message already has a `context` set it will be overridden.
|
||||
// TODO(alokedesai): Merge the incoming context with the static `TelemetryContext`, if set.
|
||||
fn attach_context(&mut self) {
|
||||
let context = telemetry_context().as_value();
|
||||
match self {
|
||||
RudderMessage::Identify(identify) => {
|
||||
identify.context = Some(context);
|
||||
}
|
||||
RudderMessage::Track(track) => track.context = Some(context),
|
||||
RudderMessage::Page(page) => page.context = Some(context),
|
||||
RudderMessage::Screen(screen) => screen.context = Some(context),
|
||||
RudderMessage::Group(group) => group.context = Some(context),
|
||||
RudderMessage::Alias(alias) => alias.context = Some(context),
|
||||
RudderMessage::Batch(batch) => batch.context = Some(context),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the user agent provided by the browser, if on Web. If not on Web,
|
||||
/// or if the user agent was not able to be read, returns None.
|
||||
fn user_agent() -> Option<String> {
|
||||
cfg_if::cfg_if! {
|
||||
if #[cfg(target_family = "wasm")] {
|
||||
wasm::user_agent()
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the telemetry context
|
||||
/// that should be attached to all telemetry events associated to this client.
|
||||
///
|
||||
/// [Rudderstack](https://www.rudderstack.com/docs/event-spec/standard-events/common-fields/#contextual-fields)
|
||||
pub fn telemetry_context() -> &'static TelemetryContext {
|
||||
TELEMETRY_CONTEXT.get_or_init(TelemetryContext::new)
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
use warp_core::telemetry::{TelemetryContextModel, TelemetryContextProvider};
|
||||
use warpui::{AppContext, ModelContext, SingletonEntity};
|
||||
|
||||
use crate::auth::AuthStateProvider;
|
||||
|
||||
pub struct AppTelemetryContextProvider {}
|
||||
|
||||
impl AppTelemetryContextProvider {
|
||||
pub fn new_context_provider(
|
||||
_ctx: &mut ModelContext<TelemetryContextModel>,
|
||||
) -> TelemetryContextModel {
|
||||
Box::new(Self {})
|
||||
}
|
||||
}
|
||||
|
||||
impl TelemetryContextProvider for AppTelemetryContextProvider {
|
||||
fn user_id(&self, ctx: &AppContext) -> Option<String> {
|
||||
let auth_state = AuthStateProvider::as_ref(ctx).get();
|
||||
auth_state.user_id().map(|uid| uid.as_string())
|
||||
}
|
||||
|
||||
fn anonymous_id(&self, ctx: &AppContext) -> String {
|
||||
let auth_state = AuthStateProvider::as_ref(ctx).get();
|
||||
auth_state.anonymous_id()
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,28 @@
|
||||
use warp_core::telemetry::TelemetryEventDesc;
|
||||
|
||||
#[derive(Debug)]
|
||||
enum TelemetryEventPropertyError {
|
||||
// The variant data is never directly read, but it's used for error formatting if the test
|
||||
// below fails.
|
||||
EmptyName(#[expect(dead_code)] Box<dyn TelemetryEventDesc>),
|
||||
EmptyDescription(#[expect(dead_code)] Box<dyn TelemetryEventDesc>),
|
||||
}
|
||||
|
||||
/// Checks that all telemetry events have a non-empty name and description.
|
||||
///
|
||||
/// The name and description are intended to be user-facing and are used to populate
|
||||
/// our [exhaustive telemetry table](https://docs.warp.dev/support-and-community/privacy-and-security/privacy#exhaustive-telemetry-table).
|
||||
#[test]
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
fn telemetry_events_have_nonempty_name_and_description() -> Result<(), TelemetryEventPropertyError>
|
||||
{
|
||||
for event in warp_core::telemetry::all_events() {
|
||||
if event.name().is_empty() {
|
||||
return Err(TelemetryEventPropertyError::EmptyName(event));
|
||||
}
|
||||
if event.description().is_empty() {
|
||||
return Err(TelemetryEventPropertyError::EmptyDescription(event));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
/// Sends a telemetry event to Rudderstack immediately instead of adding it to the event queue that is
|
||||
/// periodically flushed. This is useful under certain conditions where we want to ensure an event
|
||||
/// is immediately sent to Rudderstack even if the user quits before the queue is flushed.
|
||||
#[macro_export]
|
||||
macro_rules! send_telemetry_sync_from_ctx {
|
||||
($event:expr, $ctx:expr) => {
|
||||
#[allow(unused_imports)]
|
||||
use warp_core::telemetry::TelemetryEvent as _;
|
||||
let event = $event;
|
||||
if event.enablement_state().is_enabled() {
|
||||
let server_api =
|
||||
<$crate::server::server_api::ServerApiProvider as warpui::SingletonEntity>::handle(
|
||||
$ctx,
|
||||
)
|
||||
.as_ref($ctx)
|
||||
.get();
|
||||
let privacy_settings_snapshot =
|
||||
<$crate::settings::PrivacySettings as warpui::SingletonEntity>::handle($ctx)
|
||||
.as_ref($ctx)
|
||||
.get_snapshot($ctx);
|
||||
let _ = $ctx.spawn(
|
||||
async move {
|
||||
if let Err(error) = server_api
|
||||
.send_telemetry_event(event, privacy_settings_snapshot)
|
||||
.await
|
||||
{
|
||||
log::warn!("Error occurred with sending telemetry event: {}", error);
|
||||
}
|
||||
},
|
||||
|_, _, _| {},
|
||||
);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
/// Sends a telemetry event to Rudderstack immediately. This is the same as [`send_telemetry_sync_from_ctx`],
|
||||
/// but can be used when the caller only has access to an [`App`] and not a
|
||||
/// `ViewContext`.
|
||||
#[macro_export]
|
||||
macro_rules! send_telemetry_sync_from_app_ctx {
|
||||
($event:expr, $app_ctx:expr) => {
|
||||
#[allow(unused_imports)]
|
||||
use warp_core::telemetry::TelemetryEvent as _;
|
||||
if $event.enablement_state().is_enabled() {
|
||||
let server_api =
|
||||
<$crate::server::server_api::ServerApiProvider as warpui::SingletonEntity>::handle(
|
||||
$app_ctx,
|
||||
)
|
||||
.as_ref($app_ctx)
|
||||
.get();
|
||||
let privacy_settings_snapshot =
|
||||
<$crate::settings::PrivacySettings as warpui::SingletonEntity>::handle($app_ctx)
|
||||
.as_ref($app_ctx)
|
||||
.get_snapshot($app_ctx);
|
||||
$app_ctx
|
||||
.background_executor()
|
||||
.spawn(async move {
|
||||
if let Err(error) = server_api
|
||||
.send_telemetry_event($event, privacy_settings_snapshot)
|
||||
.await
|
||||
{
|
||||
log::warn!("Error occurred with sending telemetry event: {error}");
|
||||
}
|
||||
})
|
||||
.detach();
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
/// Sends a telemetry `track` event Rudderstack asynchronously. This is the same as the
|
||||
/// [`send_telemetry_from_ctx`], except can be called any time you have an Arc<Background>.
|
||||
/// This should only be called when invoking one of the other macros isn't possible; for example,
|
||||
/// when you are already on a background thread and thus can't access any app context.
|
||||
#[macro_export]
|
||||
macro_rules! send_telemetry_on_executor {
|
||||
($auth_state: expr, $event:expr, $executor:expr) => {
|
||||
#[allow(unused_imports)]
|
||||
use warp_core::telemetry::TelemetryEvent as _;
|
||||
let event = $event;
|
||||
if event.enablement_state().is_enabled() {
|
||||
let user_id = $auth_state.user_id().map(|uid| uid.as_string());
|
||||
let anonymous_id = $auth_state.anonymous_id();
|
||||
warpui::record_telemetry_on_executor!(
|
||||
user_id,
|
||||
anonymous_id,
|
||||
event.name().into(),
|
||||
event.payload(),
|
||||
event.contains_ugc(),
|
||||
$executor
|
||||
);
|
||||
}
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,410 @@
|
||||
mod collector;
|
||||
mod context;
|
||||
pub mod context_provider;
|
||||
mod events;
|
||||
mod macros;
|
||||
pub mod rudder_message;
|
||||
pub mod secret_redaction;
|
||||
|
||||
use chrono::Utc;
|
||||
pub use collector::*;
|
||||
pub use context::telemetry_context;
|
||||
pub use events::*;
|
||||
|
||||
use crate::auth::UserUid;
|
||||
use crate::features::FeatureFlag;
|
||||
use crate::server::telemetry::context::AttachContext;
|
||||
use crate::server::telemetry_ext::TelemetryExt;
|
||||
use crate::settings::PrivacySettingsSnapshot;
|
||||
use crate::ChannelState;
|
||||
use anyhow::Result;
|
||||
use futures::FutureExt;
|
||||
use rudder_message::{
|
||||
Batch as RudderBatch, BatchMessage as RudderBatchMessageWithMetadata,
|
||||
BatchMessageItem as RudderBatchMessage, Message as RudderMessage,
|
||||
};
|
||||
use std::fs::File;
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use std::fs::OpenOptions;
|
||||
use std::future::Future;
|
||||
use std::path::{Path, PathBuf};
|
||||
use warp_core::channel::RudderStackDestination;
|
||||
use warpui::telemetry::Event;
|
||||
|
||||
/// Filename for file where telemetry events are written on app quit.
|
||||
const RUDDER_TELEMETRY_EVENTS_FILE_NAME: &str = "rudder_telemetry_events.json";
|
||||
|
||||
/// Filepath where the Rudder events should be written on app quit.
|
||||
fn rudder_event_file_path() -> PathBuf {
|
||||
warp_core::paths::secure_state_dir()
|
||||
.unwrap_or_else(warp_core::paths::state_dir)
|
||||
.join(RUDDER_TELEMETRY_EVENTS_FILE_NAME)
|
||||
}
|
||||
|
||||
/// Removes all telemetry events from the app telemetry event queue.
|
||||
pub fn clear_event_queue() {
|
||||
let _ = warpui::telemetry::flush_events();
|
||||
}
|
||||
|
||||
pub struct TelemetryApi {
|
||||
pub(super) client: http_client::Client,
|
||||
}
|
||||
|
||||
impl Default for TelemetryApi {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl TelemetryApi {
|
||||
pub fn new() -> Self {
|
||||
cfg_if::cfg_if! {
|
||||
if #[cfg(test)] {
|
||||
let client = http_client::Client::new_for_test();
|
||||
} else if #[cfg(target_family = "wasm")] {
|
||||
let client = http_client::Client::default();
|
||||
} else {
|
||||
use std::time::Duration;
|
||||
|
||||
let client = http_client::Client::from_client_builder(
|
||||
// We use our own http client directly instead of the Rudderstack SDK's because using
|
||||
// our own client gives us the ability to have universal hooks for pre/post
|
||||
// request/response logic.
|
||||
reqwest::Client::builder()
|
||||
// Don't allow insecure connections; they will be rejected by
|
||||
// the server with a 403 Forbidden.
|
||||
.https_only(true)
|
||||
// Keep idle connections in the pool for up to 55s. AWS
|
||||
// Application Load Balancers will drop idle connections after
|
||||
// 60s and the default pool idle timeout is 90s; a pool idle
|
||||
// timeout longer than the server timeout can lead to errors
|
||||
// upon trying to use an idle connection.
|
||||
.pool_idle_timeout(Duration::from_secs(55))
|
||||
.connect_timeout(Duration::from_secs(10)),
|
||||
).expect("Client should be constructed since we use a compatibility layer to use reqwest::Client");
|
||||
}
|
||||
}
|
||||
|
||||
Self { client }
|
||||
}
|
||||
|
||||
// Batches up telemetry events from the global queue and sends a Message to the Rudderstack API.
|
||||
// Returns the number of events that were flushed.
|
||||
pub async fn flush_events(&self, settings_snapshot: PrivacySettingsSnapshot) -> Result<usize> {
|
||||
let events = warpui::telemetry::flush_events();
|
||||
let event_count = events.len();
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
if FeatureFlag::SendTelemetryToFile.is_enabled() {
|
||||
self.persist_events_to_telemetry_log_file(events.clone())?;
|
||||
}
|
||||
|
||||
if ChannelState::is_release_bundle() || FeatureFlag::WithSandboxTelemetry.is_enabled() {
|
||||
self.send_batch_messages_to_rudder(
|
||||
events
|
||||
.into_iter()
|
||||
.map(Event::to_rudder_batch_message)
|
||||
.collect(),
|
||||
settings_snapshot,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
|
||||
Ok(event_count)
|
||||
}
|
||||
|
||||
/// Flushes events directly to Rudder that were previously written into a file at `path`
|
||||
/// (likely via a call to `write_events_to_disk`).
|
||||
pub async fn flush_persisted_events_to_rudder(
|
||||
&self,
|
||||
path: &Path,
|
||||
settings_snapshot: PrivacySettingsSnapshot,
|
||||
) -> Result<()> {
|
||||
if path.exists() {
|
||||
let file = File::open(path)?;
|
||||
let events: Vec<RudderBatchMessage> = serde_json::from_reader(file)?;
|
||||
if !events.is_empty() {
|
||||
let rudder_batch_messages = events
|
||||
.into_iter()
|
||||
.map(|message| RudderBatchMessageWithMetadata {
|
||||
message,
|
||||
// We don't persist any events that contain sensitive user data.
|
||||
contains_ugc: false,
|
||||
})
|
||||
.collect();
|
||||
self.send_batch_messages_to_rudder(rudder_batch_messages, settings_snapshot)
|
||||
.await?;
|
||||
log::info!("Successfully flushed events to rudder from disk");
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Writes the last `max_event_count` events into disk. This is useful for persisting events
|
||||
/// where we can't make a network call to Rudder (such as when the app quits). To flush these
|
||||
/// events to Rudder, call `flush_events_to_rudder_from_disk`.
|
||||
pub fn flush_and_persist_events(
|
||||
&self,
|
||||
max_event_count: usize,
|
||||
settings_snapshot: PrivacySettingsSnapshot,
|
||||
) -> Result<()> {
|
||||
self.flush_and_persist_events_at_path(
|
||||
max_event_count,
|
||||
settings_snapshot,
|
||||
rudder_event_file_path(),
|
||||
)
|
||||
}
|
||||
|
||||
fn flush_and_persist_events_at_path(
|
||||
&self,
|
||||
max_event_count: usize,
|
||||
settings_snapshot: PrivacySettingsSnapshot,
|
||||
path: impl AsRef<Path>,
|
||||
) -> Result<()> {
|
||||
if settings_snapshot.should_disable_telemetry() {
|
||||
log::info!("Not writing queued events to disk because telemetry is disabled.");
|
||||
return Result::Ok(());
|
||||
}
|
||||
log::info!("Writing queued events to disk because telemetry is enabled.");
|
||||
|
||||
let file = File::create(path)?;
|
||||
|
||||
let events = warpui::telemetry::flush_events();
|
||||
if events.len() > max_event_count {
|
||||
log::error!("More telemetry events in queue than the limit to persist")
|
||||
}
|
||||
|
||||
self.persist_events_at_path(&file, max_event_count, events)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn persist_events_at_path(
|
||||
&self,
|
||||
file: &File,
|
||||
max_event_count: usize,
|
||||
events: Vec<Event>,
|
||||
) -> Result<()> {
|
||||
let rudder_events_to_persist: Vec<_> = events
|
||||
.into_iter()
|
||||
.rev()
|
||||
.take(max_event_count)
|
||||
.map(TelemetryExt::to_rudder_batch_message)
|
||||
.filter_map(|message| (!message.contains_ugc).then_some(message.message))
|
||||
.collect();
|
||||
serde_json::to_writer(file, &rudder_events_to_persist)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
fn persist_events_to_telemetry_log_file(&self, events: Vec<Event>) -> Result<()> {
|
||||
let log_directory = warp_logging::log_directory()?;
|
||||
let telemetry_file_path = log_directory.join(&*ChannelState::telemetry_file_name());
|
||||
|
||||
let file = OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&telemetry_file_path)?;
|
||||
|
||||
self.persist_events_at_path(&file, events.len(), events)
|
||||
}
|
||||
|
||||
/// Sends a `TelemetryEvent` to the Rudderstack API.
|
||||
pub async fn send_telemetry_event(
|
||||
&self,
|
||||
user_id: Option<UserUid>,
|
||||
anonymous_id: String,
|
||||
event: impl warp_core::telemetry::TelemetryEvent,
|
||||
settings_snapshot: PrivacySettingsSnapshot,
|
||||
) -> Result<()> {
|
||||
let event = warpui::telemetry::create_event(
|
||||
user_id.map(|uid| uid.as_string()),
|
||||
anonymous_id,
|
||||
event.name().into(),
|
||||
event.payload(),
|
||||
event.contains_ugc(),
|
||||
warpui::time::get_current_time(),
|
||||
);
|
||||
|
||||
self.send_telemetry_event_internal(event, settings_snapshot)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Internal implementation for sending telemetry events. This reduces code size, since
|
||||
// we:
|
||||
// 1. Return a boxed future, so calling `async` functions don't need to inline this one.
|
||||
// 2. Don't have to monomorphize for each telemetry event implementation.
|
||||
fn send_telemetry_event_internal(
|
||||
&self,
|
||||
event: Event,
|
||||
settings_snapshot: PrivacySettingsSnapshot,
|
||||
) -> impl Future<Output = Result<()>> + '_ {
|
||||
let work = async move {
|
||||
if settings_snapshot.should_disable_telemetry() {
|
||||
log::info!("Not sending telemetry event because telemetry is disabled.");
|
||||
return Result::Ok(());
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
if FeatureFlag::SendTelemetryToFile.is_enabled() {
|
||||
self.persist_events_to_telemetry_log_file(vec![event.clone()])?;
|
||||
}
|
||||
|
||||
if !(ChannelState::is_release_bundle()
|
||||
|| FeatureFlag::WithSandboxTelemetry.is_enabled())
|
||||
{
|
||||
return Result::Ok(());
|
||||
}
|
||||
|
||||
let rudder_batch = vec![event.to_rudder_batch_message()];
|
||||
|
||||
let result = self
|
||||
.send_batch_messages_to_rudder(rudder_batch, settings_snapshot)
|
||||
.await;
|
||||
|
||||
// This is only conditionally compiled because `is_connect` is not
|
||||
// available on wasm. If additional checks are made against the
|
||||
// `reqwest::Error`, this condition should be performed specifically
|
||||
// against `is_connect` and not the whole loop.
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
if let Err(error) = &result {
|
||||
for cause in error.chain() {
|
||||
if let Some(err) = cause.downcast_ref::<reqwest::Error>() {
|
||||
if err.is_connect() {
|
||||
log::warn!("Failed to send telemetry event: {error}");
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
};
|
||||
|
||||
// On WASM, the work future is non-Send, because the HTTP request future contains a reference to a JS
|
||||
// value (which is fine, since our WASM executor is single-threaded). On all other platforms, we must
|
||||
// return a Send future in order to use the background executor.
|
||||
cfg_if::cfg_if! {
|
||||
if #[cfg(target_family = "wasm")] {
|
||||
work.boxed_local()
|
||||
} else {
|
||||
work.boxed()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a batch of RudderStack messages to their HTTP API.
|
||||
/// Note that the rudderanalytics SDK provides a client, but we don't
|
||||
/// use it for a few reasons:
|
||||
/// 1. It only supports a blocking HTTP client instead of an async one
|
||||
/// 2. We want to use our own HTTP client which has before/after request logging hooks
|
||||
#[cfg_attr(target_family = "wasm", allow(clippy::question_mark))]
|
||||
async fn send_batch_messages_to_rudder(
|
||||
&self,
|
||||
messages: Vec<RudderBatchMessageWithMetadata>,
|
||||
settings_snapshot: PrivacySettingsSnapshot,
|
||||
) -> Result<()> {
|
||||
if messages.is_empty() {
|
||||
log::debug!("Dropping empty RudderStack telemetry batch");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if settings_snapshot.should_disable_telemetry() {
|
||||
log::info!("Not sending batched messages because telemetry is disabled.");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
log::info!("Start to send telemetry events to RudderStack");
|
||||
|
||||
let (mut messages_with_ugc, messages_without_ugc): (Vec<_>, Vec<_>) = messages
|
||||
.into_iter()
|
||||
.partition(|message| message.contains_ugc);
|
||||
|
||||
// If we shouldn't collect UGC telemetry, forceably clear any messages with UGC before trying to send.
|
||||
if !settings_snapshot.should_collect_ai_ugc_telemetry() {
|
||||
messages_with_ugc.clear();
|
||||
}
|
||||
|
||||
for (messages, rudder_stack_destination) in [
|
||||
(
|
||||
messages_with_ugc,
|
||||
ChannelState::rudderstack_ugc_destination(),
|
||||
),
|
||||
(
|
||||
messages_without_ugc,
|
||||
ChannelState::rudderstack_non_ugc_destination(),
|
||||
),
|
||||
] {
|
||||
if messages.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Note that timestamp and context are already included in the individual RudderBatchMessages
|
||||
// and these are the most important ones,
|
||||
// but we also add them to the RudderMessage::Batch wrapper.
|
||||
let rudder_message = RudderMessage::Batch(RudderBatch {
|
||||
batch: messages
|
||||
.into_iter()
|
||||
.map(|message| message.message)
|
||||
.collect(),
|
||||
original_timestamp: Some(Utc::now()),
|
||||
..Default::default()
|
||||
});
|
||||
if let Err(e) = self
|
||||
.send_rudder_request(rudder_message, rudder_stack_destination)
|
||||
.await
|
||||
{
|
||||
// Don't treat a connection issue as an error as these are outside of our control.
|
||||
//
|
||||
// This is only conditionally compiled because `is_connect` is not
|
||||
// available on wasm. If additional checks are made against the
|
||||
// `reqwest::Error`, this condition should be performed specifically
|
||||
// against `is_connect` and not the whole loop.
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
for cause in e.chain() {
|
||||
if let Some(err) = cause.downcast_ref::<reqwest::Error>() {
|
||||
if err.is_connect() {
|
||||
log::warn!("Failed to send event to RudderStack: {e}");
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
return Err(e);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Sends a POST request to the RudderStack HTTP API.
|
||||
async fn send_rudder_request(
|
||||
&self,
|
||||
mut msg: RudderMessage,
|
||||
rudder_stack_destination: RudderStackDestination,
|
||||
) -> Result<()> {
|
||||
msg.attach_context();
|
||||
|
||||
let path = match msg {
|
||||
RudderMessage::Identify(_) => "/v1/identify",
|
||||
RudderMessage::Track(_) => "/v1/track",
|
||||
RudderMessage::Page(_) => "/v1/page",
|
||||
RudderMessage::Screen(_) => "/v1/screen",
|
||||
RudderMessage::Group(_) => "/v1/group",
|
||||
RudderMessage::Alias(_) => "/v1/alias",
|
||||
RudderMessage::Batch(_) => "/v1/batch",
|
||||
};
|
||||
|
||||
self.client
|
||||
.post(&format!("{}{}", rudder_stack_destination.root_url, path))
|
||||
.basic_auth(rudder_stack_destination.write_key, Some(""))
|
||||
.json(&msg)
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "mod_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,61 @@
|
||||
use rudder_message::Track;
|
||||
use virtual_fs::VirtualFS;
|
||||
|
||||
use super::*;
|
||||
|
||||
// Tests that events with UGC are not persisted to desk.
|
||||
#[test]
|
||||
fn test_persist_events_doesnt_include_ugc_events() {
|
||||
let telemetry_api = TelemetryApi::new();
|
||||
|
||||
VirtualFS::test(
|
||||
"test_persist_events_doesnt_include_ugc_events",
|
||||
|dirs, _sandbox| {
|
||||
// Add one event without UGC
|
||||
let user_id = Some("user".into());
|
||||
let anonymous_id = "anonymous_id".to_owned();
|
||||
|
||||
warpui::telemetry::record_event(
|
||||
user_id.clone(),
|
||||
anonymous_id.clone(),
|
||||
"non UGC event name".into(),
|
||||
None, /* payload */
|
||||
false, /* contains_ugc */
|
||||
warpui::time::get_current_time(),
|
||||
);
|
||||
|
||||
warpui::telemetry::record_event(
|
||||
user_id.clone(),
|
||||
anonymous_id.clone(),
|
||||
"UGC event name".into(),
|
||||
None, /* payload */
|
||||
true, /* contains_ugc */
|
||||
warpui::time::get_current_time(),
|
||||
);
|
||||
|
||||
let file_path = dirs.root().join("rudderstack");
|
||||
|
||||
telemetry_api
|
||||
.flush_and_persist_events_at_path(10, PrivacySettingsSnapshot::mock(), &file_path)
|
||||
.expect("Should be able to persist events");
|
||||
|
||||
let file_content: Vec<RudderBatchMessage> =
|
||||
serde_json::from_reader(File::open(file_path).expect("Failed to open file"))
|
||||
.expect("Failed to parse file");
|
||||
|
||||
assert_eq!(file_content.len(), 1);
|
||||
|
||||
let track = file_content[0].unwrap_track();
|
||||
assert_eq!(track.event, "non UGC event name");
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
impl RudderBatchMessage {
|
||||
fn unwrap_track(&self) -> &Track {
|
||||
match self {
|
||||
RudderBatchMessage::Track(track) => track,
|
||||
_ => panic!("Expected a track event"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,249 @@
|
||||
//! Module that contains RudderStack API message types.
|
||||
//! This is directly copied from the RudderStack Rust SDK: https://github.com/rudderlabs/rudder-sdk-rust/blob/master/src/message.rs
|
||||
//! We do not use the SDK directly because it unconditionally uses a blocking HTTP client, which we don't want for a few reasons:
|
||||
//! 1. The blocking HTTP client is not allowed when compiling for WASM, so the crate itself cannot be compiled for WASM
|
||||
//! 2. An async HTTP client is more efficient
|
||||
//! 3. We want to use our own HTTP client which has before/after request logging hooks
|
||||
//! We can consider using the SDK if it adds support for an async HTTP client, tracked by this issue: https://github.com/rudderlabs/rudder-sdk-rust/issues/23
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::auth::UserUid;
|
||||
|
||||
/// An enum containing all values which may be sent to RudderStack's API.
|
||||
#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum Message {
|
||||
Identify(Identify),
|
||||
Track(Track),
|
||||
Page(Page),
|
||||
Screen(Screen),
|
||||
Group(Group),
|
||||
Alias(Alias),
|
||||
Batch(Batch),
|
||||
}
|
||||
|
||||
/// An identify event.
|
||||
#[derive(PartialEq, Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct Identify {
|
||||
/// The user id associated with this message.
|
||||
#[serde(rename = "userId", skip_serializing_if = "Option::is_none")]
|
||||
pub user_id: Option<UserUid>,
|
||||
|
||||
/// The anonymous user id associated with this message.
|
||||
#[serde(rename = "anonymousId", skip_serializing_if = "Option::is_none")]
|
||||
pub anonymous_id: Option<String>,
|
||||
|
||||
/// The traits to assign to the user.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub traits: Option<Value>,
|
||||
|
||||
/// The timestamp associated with this message.
|
||||
#[serde(rename = "originalTimestamp", skip_serializing_if = "Option::is_none")]
|
||||
pub original_timestamp: Option<DateTime<Utc>>,
|
||||
|
||||
/// Context associated with this message.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub context: Option<Value>,
|
||||
|
||||
/// Integrations to route this message to.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub integrations: Option<Value>,
|
||||
}
|
||||
|
||||
/// A track event.
|
||||
#[derive(PartialEq, Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct Track {
|
||||
/// The user id associated with this message.
|
||||
#[serde(rename = "userId", skip_serializing_if = "Option::is_none")]
|
||||
pub user_id: Option<UserUid>,
|
||||
|
||||
/// The anonymous user id associated with this message.
|
||||
#[serde(rename = "anonymousId", skip_serializing_if = "Option::is_none")]
|
||||
pub anonymous_id: Option<String>,
|
||||
|
||||
/// The name of the event being tracked.
|
||||
pub event: String,
|
||||
|
||||
/// The properties associated with the event.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub properties: Option<Value>,
|
||||
|
||||
/// The timestamp associated with this message.
|
||||
#[serde(rename = "originalTimestamp", skip_serializing_if = "Option::is_none")]
|
||||
pub original_timestamp: Option<DateTime<Utc>>,
|
||||
|
||||
/// Context associated with this message.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub context: Option<Value>,
|
||||
|
||||
/// Integrations to route this message to.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub integrations: Option<Value>,
|
||||
}
|
||||
|
||||
/// A page event.
|
||||
#[derive(PartialEq, Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct Page {
|
||||
/// The user id associated with this message.
|
||||
#[serde(rename = "userId", skip_serializing_if = "Option::is_none")]
|
||||
pub user_id: Option<UserUid>,
|
||||
|
||||
/// The anonymous user id associated with this message.
|
||||
#[serde(rename = "anonymousId", skip_serializing_if = "Option::is_none")]
|
||||
pub anonymous_id: Option<String>,
|
||||
|
||||
/// The name of the page being tracked.
|
||||
pub name: String,
|
||||
|
||||
/// The properties associated with the event.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub properties: Option<Value>,
|
||||
|
||||
/// The timestamp associated with this message.
|
||||
#[serde(rename = "originalTimestamp", skip_serializing_if = "Option::is_none")]
|
||||
pub original_timestamp: Option<DateTime<Utc>>,
|
||||
|
||||
/// Context associated with this message.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub context: Option<Value>,
|
||||
|
||||
/// Integrations to route this message to.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub integrations: Option<Value>,
|
||||
}
|
||||
|
||||
/// A screen event.
|
||||
#[derive(PartialEq, Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct Screen {
|
||||
/// The user id associated with this message.
|
||||
#[serde(rename = "userId", skip_serializing_if = "Option::is_none")]
|
||||
pub user_id: Option<UserUid>,
|
||||
|
||||
/// The anonymous user id associated with this message.
|
||||
#[serde(rename = "anonymousId", skip_serializing_if = "Option::is_none")]
|
||||
pub anonymous_id: Option<String>,
|
||||
|
||||
/// The name of the screen being tracked.
|
||||
pub name: String,
|
||||
|
||||
/// The properties associated with the event.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub properties: Option<Value>,
|
||||
|
||||
/// The timestamp associated with this message.
|
||||
#[serde(rename = "originalTimestamp", skip_serializing_if = "Option::is_none")]
|
||||
pub original_timestamp: Option<DateTime<Utc>>,
|
||||
|
||||
/// Context associated with this message.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub context: Option<Value>,
|
||||
|
||||
/// Integrations to route this message to.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub integrations: Option<Value>,
|
||||
}
|
||||
|
||||
/// A group event.
|
||||
#[derive(PartialEq, Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct Group {
|
||||
/// The user id associated with this message.
|
||||
#[serde(rename = "userId", skip_serializing_if = "Option::is_none")]
|
||||
pub user_id: Option<UserUid>,
|
||||
|
||||
/// The anonymous user id associated with this message.
|
||||
#[serde(rename = "anonymousId", skip_serializing_if = "Option::is_none")]
|
||||
pub anonymous_id: Option<String>,
|
||||
|
||||
/// The group the user is being associated with.
|
||||
#[serde(rename = "groupId")]
|
||||
pub group_id: String,
|
||||
|
||||
/// The traits to assign to the group.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub traits: Option<Value>,
|
||||
|
||||
/// The timestamp associated with this message.
|
||||
#[serde(rename = "originalTimestamp", skip_serializing_if = "Option::is_none")]
|
||||
pub original_timestamp: Option<DateTime<Utc>>,
|
||||
|
||||
/// Context associated with this message.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub context: Option<Value>,
|
||||
|
||||
/// Integrations to route this message to.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub integrations: Option<Value>,
|
||||
}
|
||||
|
||||
/// An alias event.
|
||||
#[derive(PartialEq, Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct Alias {
|
||||
/// The user id associated with this message.
|
||||
#[serde(rename = "userId")]
|
||||
pub user_id: UserUid,
|
||||
|
||||
/// The user's previous ID.
|
||||
#[serde(rename = "previousId")]
|
||||
pub previous_id: String,
|
||||
|
||||
/// The traits to assign to the alias.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub traits: Option<Value>,
|
||||
|
||||
/// The timestamp associated with this message.
|
||||
#[serde(rename = "originalTimestamp", skip_serializing_if = "Option::is_none")]
|
||||
pub original_timestamp: Option<DateTime<Utc>>,
|
||||
|
||||
/// Context associated with this message.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub context: Option<Value>,
|
||||
|
||||
/// Integrations to route this message to.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub integrations: Option<Value>,
|
||||
}
|
||||
|
||||
/// A batch of events.
|
||||
#[derive(PartialEq, Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct Batch {
|
||||
/// The batch of messages to send.
|
||||
pub batch: Vec<BatchMessageItem>,
|
||||
|
||||
/// Context associated with this message.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub context: Option<Value>,
|
||||
|
||||
/// Integrations to route this message to.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub integrations: Option<Value>,
|
||||
|
||||
/// The timestamp associated with this message.
|
||||
#[serde(rename = "originalTimestamp", skip_serializing_if = "Option::is_none")]
|
||||
pub original_timestamp: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
/// An enum containing all messages which may be placed inside a batch.
|
||||
#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum BatchMessageItem {
|
||||
#[serde(rename = "identify")]
|
||||
Identify(Identify),
|
||||
#[serde(rename = "track")]
|
||||
Track(Track),
|
||||
#[serde(rename = "page")]
|
||||
Page(Page),
|
||||
#[serde(rename = "screen")]
|
||||
Screen(Screen),
|
||||
#[serde(rename = "group")]
|
||||
Group(Group),
|
||||
#[serde(rename = "alias")]
|
||||
Alias(Alias),
|
||||
}
|
||||
|
||||
/// Metadata about a batch sent to Rudderstack and whether it contains user generated content.
|
||||
pub struct BatchMessage {
|
||||
pub message: BatchMessageItem,
|
||||
pub contains_ugc: bool,
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
//! Best-effort secret redaction for telemetry payloads.
|
||||
//!
|
||||
//! Unlike the AI-side secret redaction in `app/src/ai/blocklist/block/secret_redaction.rs`,
|
||||
//! which is gated on the user's secret-redaction (a.k.a. "safe mode") setting and is used
|
||||
//! for visual obfuscation in the terminal, the redaction in this module is unconditional:
|
||||
//! we always do a redaction pass on telemetry payloads that may contain user-generated
|
||||
//! content, regardless of the user's safe-mode setting. The two settings are deliberately
|
||||
//! decoupled — visual obfuscation is a UX preference, while telemetry-side redaction is a
|
||||
//! defence-in-depth measure for data leaving the device.
|
||||
//!
|
||||
//! The regex used for redaction always includes the default patterns defined in
|
||||
//! `crate::terminal::model::secrets::regexes::DEFAULT_REGEXES_WITH_NAMES`. Any custom
|
||||
//! patterns the user has configured (or that their organization has configured via
|
||||
//! enterprise secret redaction) are layered on top of those defaults.
|
||||
//!
|
||||
//! This module is intentionally lightweight: it does byte-range matching only and does
|
||||
//! not track `SecretLevel`s or character ranges, since the telemetry path doesn't need
|
||||
//! either.
|
||||
use crate::terminal::model::secrets::regexes::DEFAULT_REGEXES_WITH_NAMES;
|
||||
use lazy_static::lazy_static;
|
||||
use parking_lot::RwLock;
|
||||
use regex_automata::meta::Regex;
|
||||
use serde_json::Value;
|
||||
use std::collections::HashSet;
|
||||
use std::ops::Range;
|
||||
const REDACTION_REPLACEMENT_CHARACTER: &str = "*";
|
||||
lazy_static! {
|
||||
/// Regex used to redact secrets from telemetry payloads. Initialized with the
|
||||
/// default patterns so that redaction works even before the user's privacy
|
||||
/// settings are loaded (and even for users who have never configured any
|
||||
/// custom patterns).
|
||||
static ref TELEMETRY_SECRETS_REGEX: RwLock<Regex> = RwLock::new(build_default_regex());
|
||||
}
|
||||
/// Builds a regex containing only the default patterns. Used to seed the static
|
||||
/// regex before the privacy settings are loaded.
|
||||
fn build_default_regex() -> Regex {
|
||||
let patterns: Vec<&str> = DEFAULT_REGEXES_WITH_NAMES
|
||||
.iter()
|
||||
.map(|d| d.pattern)
|
||||
.collect();
|
||||
Regex::new_many(&patterns).expect("default secret patterns should compile")
|
||||
}
|
||||
/// Rebuilds [`TELEMETRY_SECRETS_REGEX`] from the user's and enterprise's secret
|
||||
/// regex lists, layered on top of the default patterns. The default patterns are
|
||||
/// always included, so redaction works even when the user has not configured any
|
||||
/// custom patterns.
|
||||
pub fn update_telemetry_secrets_regex<'a, U, E>(user_secrets: U, enterprise_secrets: E)
|
||||
where
|
||||
U: IntoIterator<Item = &'a regex::Regex>,
|
||||
E: IntoIterator<Item = &'a regex::Regex>,
|
||||
{
|
||||
let patterns = compose_patterns(
|
||||
user_secrets.into_iter().map(regex::Regex::as_str),
|
||||
enterprise_secrets.into_iter().map(regex::Regex::as_str),
|
||||
);
|
||||
match Regex::new_many(&patterns) {
|
||||
Ok(regex) => *TELEMETRY_SECRETS_REGEX.write() = regex,
|
||||
Err(err) => log::error!("Failed to build telemetry secrets regex: {err:?}"),
|
||||
}
|
||||
}
|
||||
/// Composes the full list of patterns to compile into the telemetry regex,
|
||||
/// ordered enterprise → user → defaults, with later occurrences of an already-
|
||||
/// seen pattern string deduped out.
|
||||
fn compose_patterns<'a>(
|
||||
user: impl Iterator<Item = &'a str>,
|
||||
enterprise: impl Iterator<Item = &'a str>,
|
||||
) -> Vec<&'a str> {
|
||||
let mut seen: HashSet<&str> = HashSet::new();
|
||||
let mut patterns: Vec<&str> = Vec::new();
|
||||
let all = enterprise
|
||||
.chain(user)
|
||||
.chain(DEFAULT_REGEXES_WITH_NAMES.iter().map(|d| d.pattern));
|
||||
for pattern in all {
|
||||
if seen.insert(pattern) {
|
||||
patterns.push(pattern);
|
||||
}
|
||||
}
|
||||
patterns
|
||||
}
|
||||
/// Replaces every detected secret in `input` with a run of asterisks of the same
|
||||
/// byte length. Overlapping matches (which can occur when multiple patterns match
|
||||
/// the same region) are merged before replacement, so each character is replaced
|
||||
/// at most once.
|
||||
pub fn redact_secrets_in_string(input: &mut String) {
|
||||
let ranges: Vec<Range<usize>> = {
|
||||
let regex = TELEMETRY_SECRETS_REGEX.read();
|
||||
regex.find_iter(input.as_str()).map(|m| m.range()).collect()
|
||||
};
|
||||
replace_byte_ranges_with_asterisks(input, ranges);
|
||||
}
|
||||
/// Replaces each byte range in `input` with a run of asterisks of the same byte
|
||||
/// length. Handles overlapping ranges by merging them first, and replaces from
|
||||
/// the end of the string so earlier byte indices stay valid as we mutate.
|
||||
fn replace_byte_ranges_with_asterisks(input: &mut String, mut ranges: Vec<Range<usize>>) {
|
||||
if ranges.is_empty() {
|
||||
return;
|
||||
}
|
||||
// Sort and merge overlapping ranges so we don't double-replace.
|
||||
ranges.sort_by_key(|r| r.start);
|
||||
let mut merged: Vec<Range<usize>> = Vec::with_capacity(ranges.len());
|
||||
for range in ranges {
|
||||
match merged.last_mut() {
|
||||
Some(last) if range.start <= last.end => last.end = last.end.max(range.end),
|
||||
_ => merged.push(range),
|
||||
}
|
||||
}
|
||||
// Replace from the end of the string so earlier byte indices stay valid.
|
||||
for range in merged.into_iter().rev() {
|
||||
let len = range.end - range.start;
|
||||
input.replace_range(range, &REDACTION_REPLACEMENT_CHARACTER.repeat(len));
|
||||
}
|
||||
}
|
||||
/// Walks a [`Value`] and runs [`redact_secrets_in_string`] on every string within
|
||||
/// it. Non-string scalars (numbers, booleans, nulls) are left untouched.
|
||||
pub fn redact_secrets_in_value(value: &mut Value) {
|
||||
match value {
|
||||
Value::String(s) => redact_secrets_in_string(s),
|
||||
Value::Array(arr) => arr.iter_mut().for_each(redact_secrets_in_value),
|
||||
Value::Object(obj) => obj.values_mut().for_each(redact_secrets_in_value),
|
||||
Value::Null | Value::Bool(_) | Value::Number(_) => {}
|
||||
}
|
||||
}
|
||||
#[cfg(test)]
|
||||
#[path = "secret_redaction_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,222 @@
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
// AWS-style access keys used in tests; these match `AWS_ACCESS_ID` from
|
||||
// `DEFAULT_REGEXES_WITH_NAMES`. The example value is the standard one used in
|
||||
// AWS documentation and is not a real key.
|
||||
const AWS_KEY_1: &str = "AKIAIOSFODNN7EXAMPLE";
|
||||
const AWS_KEY_2: &str = "AKIA1234567890123456";
|
||||
|
||||
#[test]
|
||||
fn redact_secrets_in_string_with_no_match_is_noop() {
|
||||
let mut s = String::from("hello world, no secrets here");
|
||||
let original = s.clone();
|
||||
redact_secrets_in_string(&mut s);
|
||||
assert_eq!(s, original);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_secrets_in_string_redacts_single_secret_in_middle() {
|
||||
let mut s = format!("prefix {AWS_KEY_1} suffix");
|
||||
redact_secrets_in_string(&mut s);
|
||||
let expected = format!("prefix {} suffix", "*".repeat(AWS_KEY_1.len()));
|
||||
assert_eq!(s, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_secrets_in_string_redacts_multiple_independent_secrets() {
|
||||
// This exercises the "replace from the end so earlier byte indices stay
|
||||
// valid" requirement: replacing the second secret first must not invalidate
|
||||
// the byte indices of the first secret.
|
||||
let mut s = format!("a {AWS_KEY_1} b {AWS_KEY_2} c");
|
||||
redact_secrets_in_string(&mut s);
|
||||
let expected = format!(
|
||||
"a {} b {} c",
|
||||
"*".repeat(AWS_KEY_1.len()),
|
||||
"*".repeat(AWS_KEY_2.len()),
|
||||
);
|
||||
assert_eq!(s, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_secrets_in_string_redacts_string_that_is_entirely_a_secret() {
|
||||
let mut s = AWS_KEY_1.to_string();
|
||||
redact_secrets_in_string(&mut s);
|
||||
assert_eq!(s, "*".repeat(AWS_KEY_1.len()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replace_byte_ranges_with_asterisks_with_empty_ranges_is_noop() {
|
||||
let mut s = String::from("no changes here");
|
||||
let original = s.clone();
|
||||
replace_byte_ranges_with_asterisks(&mut s, vec![]);
|
||||
assert_eq!(s, original);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replace_byte_ranges_with_asterisks_replaces_independent_ranges() {
|
||||
// Separate, non-overlapping ranges are replaced independently. This
|
||||
// exercises the reverse-iteration so earlier byte indices stay valid as
|
||||
// the string mutates.
|
||||
let mut s = String::from("0123456789ABCDEF");
|
||||
let ranges = vec![0..3, 6..9, 12..15];
|
||||
replace_byte_ranges_with_asterisks(&mut s, ranges);
|
||||
assert_eq!(s, "***345***9AB***F");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replace_byte_ranges_with_asterisks_merges_overlapping_ranges() {
|
||||
// Ranges 0..10 and 5..15 overlap; they should merge into 0..15 so we don't
|
||||
// double-replace the bytes in 5..10.
|
||||
let mut s = String::from("0123456789ABCDEF");
|
||||
let ranges = vec![0..10, 5..15];
|
||||
replace_byte_ranges_with_asterisks(&mut s, ranges);
|
||||
assert_eq!(s, "***************F");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replace_byte_ranges_with_asterisks_merges_adjacent_ranges() {
|
||||
// Adjacent (touching) ranges should also merge.
|
||||
let mut s = String::from("0123456789ABCDEF");
|
||||
let ranges = vec![0..5, 5..10];
|
||||
replace_byte_ranges_with_asterisks(&mut s, ranges);
|
||||
assert_eq!(s, "**********ABCDEF");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replace_byte_ranges_with_asterisks_handles_unsorted_ranges() {
|
||||
// Input ranges may be in arbitrary order; the function must sort before
|
||||
// merging or replacing.
|
||||
let mut s = String::from("0123456789ABCDEF");
|
||||
let ranges = vec![10..12, 0..2, 4..6];
|
||||
replace_byte_ranges_with_asterisks(&mut s, ranges);
|
||||
assert_eq!(s, "**23**6789**CDEF");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replace_byte_ranges_with_asterisks_handles_fully_contained_range() {
|
||||
// A range fully contained in another should merge to the larger range.
|
||||
let mut s = String::from("0123456789ABCDEF");
|
||||
let ranges = vec![2..14, 5..8];
|
||||
replace_byte_ranges_with_asterisks(&mut s, ranges);
|
||||
assert_eq!(s, "01************EF");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_secrets_in_value_redacts_strings_in_objects() {
|
||||
let mut value = json!({
|
||||
"with_secret": format!("contains {AWS_KEY_1} secret"),
|
||||
"without_secret": "no secret here",
|
||||
});
|
||||
redact_secrets_in_value(&mut value);
|
||||
assert_eq!(
|
||||
value["with_secret"],
|
||||
format!("contains {} secret", "*".repeat(AWS_KEY_1.len())),
|
||||
);
|
||||
assert_eq!(value["without_secret"], "no secret here");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_secrets_in_value_redacts_strings_in_arrays() {
|
||||
let mut value = json!([
|
||||
format!("first {AWS_KEY_1}"),
|
||||
"second clean",
|
||||
format!("third {AWS_KEY_2}"),
|
||||
]);
|
||||
redact_secrets_in_value(&mut value);
|
||||
assert_eq!(value[0], format!("first {}", "*".repeat(AWS_KEY_1.len())),);
|
||||
assert_eq!(value[1], "second clean");
|
||||
assert_eq!(value[2], format!("third {}", "*".repeat(AWS_KEY_2.len())),);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_secrets_in_value_recurses_into_nested_structures() {
|
||||
let mut value = json!({
|
||||
"outer": {
|
||||
"inner_array": [
|
||||
format!("nested {AWS_KEY_1}"),
|
||||
{"inner_object": format!("deep {AWS_KEY_2}")},
|
||||
],
|
||||
"scalar_int": 42,
|
||||
"scalar_bool": true,
|
||||
"scalar_null": null,
|
||||
}
|
||||
});
|
||||
redact_secrets_in_value(&mut value);
|
||||
assert_eq!(
|
||||
value["outer"]["inner_array"][0],
|
||||
format!("nested {}", "*".repeat(AWS_KEY_1.len())),
|
||||
);
|
||||
assert_eq!(
|
||||
value["outer"]["inner_array"][1]["inner_object"],
|
||||
format!("deep {}", "*".repeat(AWS_KEY_2.len())),
|
||||
);
|
||||
// Non-string scalars are left untouched.
|
||||
assert_eq!(value["outer"]["scalar_int"], 42);
|
||||
assert_eq!(value["outer"]["scalar_bool"], true);
|
||||
assert!(value["outer"]["scalar_null"].is_null());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_secrets_in_value_leaves_non_string_scalars_untouched() {
|
||||
let mut value = json!({"n": 42, "b": true, "z": null});
|
||||
let expected = value.clone();
|
||||
redact_secrets_in_value(&mut value);
|
||||
assert_eq!(value, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compose_patterns_includes_defaults_when_user_and_enterprise_are_empty() {
|
||||
let patterns = compose_patterns(std::iter::empty(), std::iter::empty());
|
||||
assert_eq!(patterns.len(), DEFAULT_REGEXES_WITH_NAMES.len());
|
||||
for default in DEFAULT_REGEXES_WITH_NAMES {
|
||||
assert!(
|
||||
patterns.contains(&default.pattern),
|
||||
"expected default pattern {} to be present",
|
||||
default.pattern,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compose_patterns_layers_user_and_enterprise_on_top_of_defaults() {
|
||||
let user = [r"\bUSER-\d+\b"];
|
||||
let enterprise = [r"\bENT-\d+\b"];
|
||||
let patterns = compose_patterns(user.iter().copied(), enterprise.iter().copied());
|
||||
// Enterprise comes first, then user, then defaults.
|
||||
assert_eq!(patterns[0], r"\bENT-\d+\b");
|
||||
assert_eq!(patterns[1], r"\bUSER-\d+\b");
|
||||
// Defaults are still all present.
|
||||
for default in DEFAULT_REGEXES_WITH_NAMES {
|
||||
assert!(
|
||||
patterns.contains(&default.pattern),
|
||||
"expected default pattern {} to be present alongside user/enterprise",
|
||||
default.pattern,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compose_patterns_dedups_user_pattern_that_matches_a_default() {
|
||||
// Pick the first default pattern; passing the same string as a "user"
|
||||
// pattern should not cause it to appear twice in the composed list.
|
||||
let duplicated = DEFAULT_REGEXES_WITH_NAMES[0].pattern;
|
||||
let patterns = compose_patterns(std::iter::once(duplicated), std::iter::empty());
|
||||
let occurrences = patterns.iter().filter(|p| **p == duplicated).count();
|
||||
assert_eq!(
|
||||
occurrences, 1,
|
||||
"duplicate pattern should appear at most once in composed list",
|
||||
);
|
||||
// Total length is the defaults (the user pattern was deduped away).
|
||||
assert_eq!(patterns.len(), DEFAULT_REGEXES_WITH_NAMES.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compose_patterns_dedups_enterprise_pattern_that_matches_a_user_pattern() {
|
||||
let user = [r"\bSHARED-\d+\b"];
|
||||
let enterprise = [r"\bSHARED-\d+\b"];
|
||||
let patterns = compose_patterns(user.iter().copied(), enterprise.iter().copied());
|
||||
let occurrences = patterns.iter().filter(|p| **p == r"\bSHARED-\d+\b").count();
|
||||
assert_eq!(occurrences, 1);
|
||||
}
|
||||
@@ -0,0 +1,129 @@
|
||||
use super::telemetry::rudder_message::{
|
||||
BatchMessage as RudderBatchMessage, BatchMessageItem as RudderBatchMessageItem,
|
||||
Identify as RudderIdentify, Track as RudderTrack,
|
||||
};
|
||||
use super::telemetry::secret_redaction::redact_secrets_in_value;
|
||||
use crate::auth::UserUid;
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde_json::{json, Value};
|
||||
use warp_core::{
|
||||
channel::{Channel, ChannelState},
|
||||
execution_mode,
|
||||
};
|
||||
use warpui::telemetry::EventPayload;
|
||||
|
||||
use super::telemetry::telemetry_context;
|
||||
|
||||
pub trait TelemetryExt {
|
||||
fn to_rudder_batch_message(self) -> RudderBatchMessage;
|
||||
}
|
||||
|
||||
impl TelemetryExt for warpui::telemetry::Event {
|
||||
fn to_rudder_batch_message(self) -> RudderBatchMessage {
|
||||
let message = match self.payload {
|
||||
EventPayload::IdentifyUser {
|
||||
user_id,
|
||||
anonymous_id,
|
||||
} => RudderBatchMessageItem::Identify(RudderIdentify {
|
||||
user_id: Some(UserUid::new(user_id.as_str())),
|
||||
anonymous_id: Some(anonymous_id),
|
||||
original_timestamp: Some(self.timestamp),
|
||||
integrations: Some(json!({
|
||||
"Amplitude": {
|
||||
"session_id": self.session_created_at.timestamp(),
|
||||
}
|
||||
})),
|
||||
context: Some(telemetry_context().as_value()),
|
||||
..Default::default()
|
||||
}),
|
||||
EventPayload::AppActive {
|
||||
user_id,
|
||||
anonymous_id,
|
||||
} => form_rudder_track_message(
|
||||
user_id.map(|uid| UserUid::new(uid.as_str())),
|
||||
anonymous_id,
|
||||
"Active App Usage".to_string(),
|
||||
None,
|
||||
self.timestamp,
|
||||
self.session_created_at,
|
||||
),
|
||||
EventPayload::NamedEvent {
|
||||
user_id,
|
||||
anonymous_id,
|
||||
name,
|
||||
mut value,
|
||||
} => {
|
||||
// For events that may contain user-generated content, run a
|
||||
// best-effort secret-redaction pass on the payload before
|
||||
// sending. This is independent of the user's safe-mode setting:
|
||||
// visual obfuscation is a UX preference, while telemetry-side
|
||||
// redaction is a defence-in-depth measure for data leaving the
|
||||
// device. See `secret_redaction.rs` for details.
|
||||
if self.contains_ugc {
|
||||
if let Some(value) = value.as_mut() {
|
||||
redact_secrets_in_value(value);
|
||||
}
|
||||
}
|
||||
form_rudder_track_message(
|
||||
user_id.map(|uid| UserUid::new(uid.as_str())),
|
||||
anonymous_id,
|
||||
name.to_string(),
|
||||
value,
|
||||
self.timestamp,
|
||||
self.session_created_at,
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
RudderBatchMessage {
|
||||
message,
|
||||
contains_ugc: self.contains_ugc,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn form_rudder_track_message(
|
||||
user_id: Option<UserUid>,
|
||||
anonymous_id: String,
|
||||
name: String,
|
||||
payload: Option<Value>,
|
||||
timestamp: DateTime<Utc>,
|
||||
session_created_at: DateTime<Utc>,
|
||||
) -> RudderBatchMessageItem {
|
||||
RudderBatchMessageItem::Track(RudderTrack {
|
||||
user_id,
|
||||
anonymous_id: Some(anonymous_id),
|
||||
event: name,
|
||||
properties: Some(json!({
|
||||
"release_mode": release_mode(ChannelState::channel()),
|
||||
"tag": ChannelState::app_version().unwrap_or("<no tag>"),
|
||||
"client_id": execution_mode::current_client_id(),
|
||||
"payload": payload
|
||||
})),
|
||||
original_timestamp: Some(timestamp),
|
||||
integrations: Some(json!({
|
||||
"Amplitude": {
|
||||
"session_id": session_created_at.timestamp(),
|
||||
}
|
||||
})),
|
||||
context: Some(telemetry_context().as_value()),
|
||||
})
|
||||
}
|
||||
|
||||
fn release_mode(channel: Channel) -> &'static str {
|
||||
match channel {
|
||||
Channel::Stable => "stable_release",
|
||||
Channel::Preview => "preview_release",
|
||||
Channel::Local => "local",
|
||||
Channel::Integration => "integration_test",
|
||||
Channel::Dev => "dev_release",
|
||||
// We don't ever expect to send telemetry for the OSS build, but
|
||||
// until we have some time to clean things up here, we'll set a valid
|
||||
// value that we never intend to receive.
|
||||
Channel::Oss => "oss_release",
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "telemetry_ext_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,103 @@
|
||||
use chrono::Utc;
|
||||
use serde_json::json;
|
||||
use warpui::telemetry::{Event, EventPayload};
|
||||
|
||||
use super::*;
|
||||
|
||||
// AWS-style access key example used in tests; matches `AWS_ACCESS_ID` in
|
||||
// `DEFAULT_REGEXES_WITH_NAMES`. The example value is the standard one used in
|
||||
// AWS documentation and is not a real key.
|
||||
const AWS_KEY: &str = "AKIAIOSFODNN7EXAMPLE";
|
||||
|
||||
/// Constructs a minimal `Event` with a `NamedEvent` payload for testing.
|
||||
fn make_named_event(value: serde_json::Value, contains_ugc: bool) -> Event {
|
||||
let now = Utc::now();
|
||||
Event {
|
||||
payload: EventPayload::NamedEvent {
|
||||
user_id: None,
|
||||
anonymous_id: "anon".to_string(),
|
||||
name: "TestEvent".into(),
|
||||
value: Some(value),
|
||||
},
|
||||
session_created_at: now,
|
||||
timestamp: now,
|
||||
contains_ugc,
|
||||
}
|
||||
}
|
||||
|
||||
/// Extracts the inner payload `Value` from a `Track`-typed `BatchMessageItem`.
|
||||
/// This mirrors the structure produced by `form_rudder_track_message`, which
|
||||
/// wraps the event payload under `properties.payload`.
|
||||
fn extract_payload(message: RudderBatchMessageItem) -> serde_json::Value {
|
||||
let track = match message {
|
||||
RudderBatchMessageItem::Track(track) => track,
|
||||
other => panic!("expected Track message, got {other:?}"),
|
||||
};
|
||||
track
|
||||
.properties
|
||||
.expect("track properties should be set")
|
||||
.get("payload")
|
||||
.cloned()
|
||||
.expect("payload should be set in properties")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn to_rudder_batch_message_redacts_ugc_named_events() {
|
||||
let payload = json!({
|
||||
"command": format!("aws s3 cp {AWS_KEY} ./file"),
|
||||
"ok": true,
|
||||
});
|
||||
let event = make_named_event(payload, /*contains_ugc=*/ true);
|
||||
|
||||
let batch = event.to_rudder_batch_message();
|
||||
assert!(batch.contains_ugc);
|
||||
|
||||
let payload = extract_payload(batch.message);
|
||||
assert_eq!(
|
||||
payload["command"],
|
||||
format!("aws s3 cp {} ./file", "*".repeat(AWS_KEY.len())),
|
||||
);
|
||||
assert_eq!(payload["ok"], true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn to_rudder_batch_message_does_not_redact_non_ugc_named_events() {
|
||||
let original_command = format!("aws s3 cp {AWS_KEY} ./file");
|
||||
let payload = json!({
|
||||
"command": original_command.clone(),
|
||||
"ok": true,
|
||||
});
|
||||
let event = make_named_event(payload, /*contains_ugc=*/ false);
|
||||
|
||||
let batch = event.to_rudder_batch_message();
|
||||
assert!(!batch.contains_ugc);
|
||||
|
||||
let payload = extract_payload(batch.message);
|
||||
// No redaction should have been applied since the event is not flagged as UGC.
|
||||
assert_eq!(payload["command"], original_command);
|
||||
assert_eq!(payload["ok"], true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn to_rudder_batch_message_redacts_nested_strings_in_ugc_payload() {
|
||||
let payload = json!({
|
||||
"outer": {
|
||||
"inner_array": [
|
||||
format!("first secret: {AWS_KEY}"),
|
||||
"no secret here",
|
||||
],
|
||||
"scalar_int": 42,
|
||||
}
|
||||
});
|
||||
let event = make_named_event(payload, /*contains_ugc=*/ true);
|
||||
|
||||
let batch = event.to_rudder_batch_message();
|
||||
let payload = extract_payload(batch.message);
|
||||
|
||||
assert_eq!(
|
||||
payload["outer"]["inner_array"][0],
|
||||
format!("first secret: {}", "*".repeat(AWS_KEY.len())),
|
||||
);
|
||||
assert_eq!(payload["outer"]["inner_array"][1], "no secret here");
|
||||
assert_eq!(payload["outer"]["scalar_int"], 42);
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use warpui::{Entity, SingletonEntity};
|
||||
|
||||
use crate::ai::voice::transcribe::{Provider, TranscribeRequest};
|
||||
use crate::voice::transcriber::Transcriber;
|
||||
|
||||
use super::server_api::{ServerApi, TranscribeError};
|
||||
|
||||
pub struct ServerVoiceTranscriber {
|
||||
server_api: Arc<ServerApi>,
|
||||
}
|
||||
|
||||
impl ServerVoiceTranscriber {
|
||||
pub fn new(server_api: Arc<ServerApi>) -> Self {
|
||||
Self { server_api }
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
||||
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
||||
impl Transcriber for ServerVoiceTranscriber {
|
||||
async fn transcribe(&self, wav_base64: String) -> Result<String, TranscribeError> {
|
||||
let request = TranscribeRequest {
|
||||
provider: Provider::Wispr,
|
||||
audio: Some(wav_base64),
|
||||
..Default::default()
|
||||
};
|
||||
let response = self.server_api.transcribe(&request).await;
|
||||
match response {
|
||||
Ok(response) => Ok(response.text),
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Entity for ServerVoiceTranscriber {
|
||||
type Event = ();
|
||||
}
|
||||
|
||||
impl SingletonEntity for ServerVoiceTranscriber {}
|
||||
Reference in New Issue
Block a user