public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first pass of merging in warp (doesn't build)

Browse Source
This commit is contained in:
Ryan Ward
2026-07-01 16:08:58 -05:00
parent 2f64909469
commit 4770ac06b5
3662 changed files with 414574 additions and 89772 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/agent_sdk/mod.rs
+466 -176
View File
@@ -7,6 +7,40 @@ use std::path::Path;
use std::pin::Pin;
use std::sync::Arc;
use ai::api_keys::{ApiKeyManager, AwsCredentialsRefreshStrategy};
use anyhow::Context;
pub(crate) use driver::harness::{task_env_vars, validate_cli_installed, ClaudeHarness};
pub use driver::AgentDriver;
use driver::AgentDriverError;
use telemetry::CliTelemetryEvent;
use tracing::Instrument as _;
use warp_cli::agent::{
AgentCommand, AgentProfileCommand, Harness, OutputFormat, Prompt, RunAgentArgs,
};
use warp_cli::api_key::ApiKeyCommand;
use warp_cli::artifact::ArtifactCommand;
use warp_cli::environment::{EnvironmentCommand, ImageCommand};
use warp_cli::federate::FederateCommand;
use warp_cli::harness_support::{HarnessSupportCommand, ReportArtifactCommand, TaskStatus};
use warp_cli::integration::IntegrationCommand;
use warp_cli::mcp::MCPCommand;
use warp_cli::memory_store::{MemoryCommand, MemoryStoreCommand};
use warp_cli::model::ModelCommand;
use warp_cli::provider::ProviderCommand;
use warp_cli::schedule::ScheduleSubcommand;
use warp_cli::secret::SecretCommand;
use warp_cli::share::ShareRequest;
use warp_cli::task::{MessageCommand, TaskCommand};
use warp_cli::{CliCommand, GlobalOptions, OZ_HARNESS_ENV};
use galaxy_core::features::FeatureFlag;
use warp_graphql::object_permissions::OwnerType;
use warp_isolation_platform::IsolationPlatformError;
#[cfg(not(target_family = "wasm"))]
use warp_logging::log_file_path;
use warp_managed_secrets::ManagedSecretManager;
use warpui::platform::TerminationMode;
use warpui::{AppContext, ModelSpawner, SingletonEntity};
use crate::ai::agent::api::convert_conversation::{
convert_conversation_data_to_ai_conversation, RestorationMode,
};
@@ -15,69 +49,35 @@ use crate::ai::agent::conversation::AIConversationId;
use crate::ai::agent_sdk::driver::harness::{harness_kind, HarnessKind};
use crate::ai::agent_sdk::driver::{AgentDriverOptions, AgentRunPrompt, Task};
use crate::ai::agent_sdk::mcp_config::build_mcp_servers_from_specs;
use crate::ai::agent_sdk::setup_observability::{
OzRunTimelineEvent, SetupClientEventReporter, SetupStep,
};
use crate::ai::ambient_agents::task::HarnessConfig;
use crate::ai::ambient_agents::AmbientAgentTaskId;
use crate::ai::attachment_utils::attachments_download_dir;
#[cfg(not(target_family = "wasm"))]
use crate::ai::aws_credentials::refresh_aws_credentials;
use crate::ai::cloud_environments::CloudAmbientAgentEnvironment;
use crate::ai::llms::LLMId;
use crate::auth::auth_manager::{AuthManager, AuthManagerEvent};
use crate::cloud_object::model::persistence::CloudModel;
use crate::server::server_api::ai::AIClient;
use crate::workflows::workflow::Workflow;
use ai::api_keys::{ApiKeyManager, AwsCredentialsRefreshStrategy};
use anyhow::Context;
use galaxy_cli::{
agent::{AgentCommand, AgentProfileCommand, OutputFormat},
artifact::ArtifactCommand,
environment::{EnvironmentCommand, ImageCommand},
federate::FederateCommand,
harness_support::{HarnessSupportCommand, ReportArtifactCommand, TaskStatus},
integration::IntegrationCommand,
mcp::MCPCommand,
model::ModelCommand,
provider::ProviderCommand,
schedule::ScheduleSubcommand,
secret::SecretCommand,
share::ShareRequest,
task::{MessageCommand, TaskCommand},
CliCommand, GlobalOptions,
};
use galaxy_core::features::FeatureFlag;
use galaxy_isolation_platform::IsolationPlatformError;
#[cfg(not(target_family = "wasm"))]
use galaxy_logging::log_file_path;
use galaxy_managed_secrets::ManagedSecretManager;
use galaxyui::ModelSpawner;
use galaxyui::{platform::TerminationMode, AppContext, SingletonEntity};
use crate::{
ai::ambient_agents::{task::HarnessConfig, AmbientAgentTaskId},
ai::cloud_environments::CloudAmbientAgentEnvironment,
auth::AuthStateProvider,
send_telemetry_sync_from_app_ctx,
server::{
ids::{ServerId, SyncId},
server_api::{ai::AgentConfigSnapshot, ServerApiProvider},
},
terminal::view::ConversationRestorationInNewPaneType,
};
use driver::AgentDriverError;
use galaxy_graphql::object_permissions::OwnerType;
use crate::ai::attachment_utils::attachments_download_dir;
use crate::ai::skills::{
clone_repo_for_skill, resolve_skill_spec, ResolveSkillError, ResolvedSkill,
};
pub(crate) use driver::harness::{
task_env_vars, validate_cli_installed, ClaudeHarness, ThirdPartyHarness,
};
pub use driver::AgentDriver;
use galaxy_cli::agent::{Harness, Prompt, RunAgentArgs};
use galaxy_cli::OZ_HARNESS_ENV;
use telemetry::CliTelemetryEvent;
use crate::auth::auth_manager::{AuthManager, AuthManagerEvent};
use crate::auth::AuthStateProvider;
use crate::cloud_object::model::persistence::CloudModel;
use crate::cloud_object::CloudObjectLookup as _;
use crate::send_telemetry_sync_from_app_ctx;
use crate::server::ids::{ServerId, SyncId};
use crate::server::server_api::ai::{AIClient, AgentConfigSnapshot, GitCredential};
use crate::server::server_api::ServerApiProvider;
use crate::terminal::view::ConversationRestorationInNewPaneType;
use crate::workflows::workflow::Workflow;
mod admin;
mod agent_config;
mod agent_management;
mod ambient;
mod api_key;
mod artifact;
pub(crate) mod artifact_upload;
mod common;
@@ -92,6 +92,7 @@ mod integration;
mod integration_output;
mod mcp;
mod mcp_config;
mod memory_store;
mod model;
mod oauth_flow;
pub mod output;
@@ -100,6 +101,7 @@ mod provider;
pub(crate) mod retry;
mod schedule;
mod secret;
pub(crate) mod setup_observability;
mod telemetry;
#[cfg(test)]
mod test_support;
@@ -120,6 +122,7 @@ fn maybe_warn_team_api_key(ctx: &AppContext) {
}
/// Run a Warp CLI command.
#[tracing::instrument(name = "agent_sdk::run", skip_all, err, fields(tags.cloud_agent = true))]
pub fn run(
ctx: &mut AppContext,
command: CliCommand,
@@ -148,6 +151,10 @@ fn dispatch_command(
CliCommand::MCP(mcp_cmd) => mcp::run(ctx, global_options, mcp_cmd),
CliCommand::Run(task_cmd) => run_task(ctx, global_options, task_cmd),
CliCommand::Model(model_cmd) => model::run(ctx, global_options, model_cmd),
CliCommand::MemoryStore(memory_store_cmd) => {
memory_store::run(ctx, global_options, memory_store_cmd)
}
CliCommand::Memory(memory_cmd) => memory_store::run_memory(ctx, global_options, memory_cmd),
CliCommand::Login => admin::login(ctx),
CliCommand::Logout => admin::logout(ctx),
CliCommand::Whoami => admin::whoami(ctx, global_options.output_format),
@@ -198,6 +205,12 @@ fn dispatch_command(
}
artifact::run(ctx, global_options, artifact_cmd)
}
CliCommand::ApiKey(api_key_cmd) => {
if !FeatureFlag::APIKeyManagement.is_enabled() {
return Err(anyhow::anyhow!("invalid value 'api-key'"));
}
api_key::run(ctx, global_options, api_key_cmd)
}
}
}
@@ -305,10 +318,28 @@ fn run_agent(
"--claude-auth-secret is only valid with --harness claude."
));
}
if args.runner.is_some() && !FeatureFlag::CloudRunners.is_enabled() {
return Err(anyhow::anyhow!("unexpected argument '--runner' found"));
}
ambient::run_ambient_agent(ctx, args)
}
AgentCommand::Profile(sub) => profiles::run(ctx, global_options, sub),
AgentCommand::List(args) => agent_config::list_agents(ctx, args),
AgentCommand::List(args) => {
agent_management::list_agents(ctx, global_options.output_format, args)
}
AgentCommand::Get(args) => {
agent_management::get_agent(ctx, global_options.output_format, args)
}
AgentCommand::Create(args) => {
agent_management::create_agent(ctx, global_options.output_format, args)
}
AgentCommand::Update(args) => {
agent_management::update_agent(ctx, global_options.output_format, args)
}
AgentCommand::Delete(args) => {
agent_management::delete_agent(ctx, global_options.output_format, args)
}
AgentCommand::Skills(args) => agent_config::list_skills(ctx, args),
}
}
@@ -337,21 +368,44 @@ fn build_merged_config_and_task(
// Merge precedence: file < CLI < skill
let file_merged = config_file::merge_with_precedence(loaded_file.as_ref(), Default::default());
// Runner support is gated. The `run` command has no `--runner` flag, but a
// config file can still set `runner_id`, so reject it when the flag is off.
if file_merged.runner_id.is_some() && !FeatureFlag::CloudRunners.is_enabled() {
return Err(anyhow::anyhow!(
"`runner_id` is set in the config file but runner support is not enabled"
));
}
// Skill provides base_prompt and optionally name
let (skill_name, runtime_base_prompt) = match resolved_skill {
Some(skill) => (Some(skill.name.clone()), Some(skill.instructions.clone())),
None => (None, None),
};
// When a non-Oz harness is active, --model targets the harness rather than the Oz model.
let harness_model_id = if args.harness != Harness::Oz {
args.model.model.clone()
} else {
None
};
let harness_override = (args.harness != Harness::Oz).then_some(HarnessConfig {
harness_type: args.harness,
model_id: harness_model_id,
reasoning_level: None,
});
let oz_model = if args.harness == Harness::Oz {
args.model.model.clone().or(file_merged.model_id)
} else {
None
};
let mut merged_config = AgentConfigSnapshot {
// CLI name > skill name > file name
name: args.name.clone().or(skill_name).or(file_merged.name),
environment_id: args.environment.clone().or(file_merged.environment_id),
model_id: args.model.model.clone().or(file_merged.model_id),
runner_id: file_merged.runner_id,
model_id: oz_model,
// Skill base_prompt takes precedence over file base_prompt
base_prompt: runtime_base_prompt.clone().or(file_merged.base_prompt),
mcp_servers: config_file::merge_mcp_servers(file_merged.mcp_servers, cli_mcp_servers),
@@ -374,6 +428,7 @@ fn build_merged_config_and_task(
let model_override: Option<LLMId> = merged_config
.model_id
.as_deref()
.filter(|_| args.harness == Harness::Oz)
.map(|model_id| common::validate_agent_mode_base_model_id(model_id, ctx))
.transpose()?;
@@ -420,15 +475,25 @@ fn build_server_side_task(
None => Vec::new(),
};
let model_override: Option<LLMId> = args
.model
.model
.as_deref()
.map(|model_id| common::validate_agent_mode_base_model_id(model_id, ctx))
.transpose()?;
let harness_model_id = if args.harness != Harness::Oz {
args.model.model.clone()
} else {
None
};
let model_override: Option<LLMId> = if args.harness == Harness::Oz {
args.model
.model
.as_deref()
.map(|model_id| common::validate_agent_mode_base_model_id(model_id, ctx))
.transpose()?
} else {
None
};
let harness_override = (args.harness != Harness::Oz).then_some(HarnessConfig {
harness_type: args.harness,
model_id: harness_model_id,
reasoning_level: None,
});
let skill_name = resolved_skill.as_ref().map(|s| s.name.clone());
@@ -439,6 +504,7 @@ fn build_server_side_task(
let config = AgentConfigSnapshot {
name: args.name.clone().or(skill_name),
environment_id: environment.clone(),
runner_id: None,
model_id: model_id_string,
base_prompt: None,
mcp_servers: cli_mcp_servers,
@@ -466,6 +532,24 @@ fn build_server_side_task(
Ok((config, task))
}
fn reconcile_task_harness(
task_id: &str,
selected_harness: &mut Harness,
task_harness: Harness,
) -> Result<HarnessKind, AgentDriverError> {
if *selected_harness == Harness::Oz {
*selected_harness = task_harness;
} else if task_harness != *selected_harness {
return Err(AgentDriverError::TaskHarnessMismatch {
task_id: task_id.to_string(),
expected: task_harness.to_string(),
got: selected_harness.to_string(),
});
}
harness_kind(*selected_harness)
}
/// Resolve a `Prompt` to a plain string.
fn resolve_prompt(prompt: &Prompt, ctx: &AppContext) -> Result<String, AgentDriverError> {
match prompt {
@@ -515,14 +599,7 @@ fn run_task(
}
}
}
TaskCommand::Message(message_cmd) => {
if !FeatureFlag::OrchestrationV2.is_enabled() {
return Err(anyhow::anyhow!(
"The 'message' subcommand is not available in this build"
));
}
ambient::run_message(ctx, global_options, message_cmd)
}
TaskCommand::Message(message_cmd) => ambient::run_message(ctx, global_options, message_cmd),
}
}
@@ -538,53 +615,80 @@ impl galaxyui::Entity for AgentDriverRunner {
impl galaxyui::SingletonEntity for AgentDriverRunner {}
impl AgentDriverRunner {
#[tracing::instrument(skip_all, err, fields(
tags.cloud_agent = true,
args.sandboxed = args.sandboxed,
args.computer_use = args.computer_use.computer_use,
args.no_computer_use = args.computer_use.no_computer_use
))]
async fn setup_and_run_driver(
foreground: ModelSpawner<Self>,
args: RunAgentArgs,
server_api: Arc<dyn AIClient>,
output_format: OutputFormat,
) -> Result<(), AgentDriverError> {
// Extract the task ID as early as possible for best-effort setup observability.
// Local CLI-created runs may not have a task yet, so those setup events explicitly no-op.
let mut task_id: Option<AmbientAgentTaskId> =
args.task_id.as_deref().and_then(|s| s.parse().ok());
Self::set_ambient_agent_task_id(&foreground, task_id).await?;
let background = foreground.spawn(|_, ctx| ctx.background_executor()).await?;
let setup_events = match task_id {
Some(task_id) => SetupClientEventReporter::new(task_id, server_api.clone(), background),
None => SetupClientEventReporter::noop(server_api.clone(), background),
};
setup_events
.post_timeline_event(OzRunTimelineEvent::WorkerContainerReady)
.await;
// Ensure we've synced team state before starting the driver.
Self::refresh_team_metadata(&foreground).await?;
setup_events
.record_result(
SetupStep::TeamMetadataRefresh,
Self::refresh_team_metadata(&foreground),
)
.await?;
// Wait for Galaxy Drive to sync before building the task config, since
// prompt resolution (SavedPrompt -> workflow lookup) and environment
// resolution (CloudAmbientAgentEnvironment lookup) depend on it.
if foreground
.spawn(|_, ctx| common::refresh_warp_drive(ctx))
.await?
.await
.is_err()
{
return Err(AgentDriverError::WarpDriveSyncFailed);
}
// Extract the task ID if available, so that if there are setup errors and we have
// a server-provided task ID, we can report them. If we create a task for a local CLI
// run, its ID will be stored in the inner future.
let mut task_id: Option<AmbientAgentTaskId> =
args.task_id.as_deref().and_then(|s| s.parse().ok());
setup_events
.record_result(SetupStep::WarpDriveSync, async {
if foreground
.spawn(|_, ctx| common::refresh_warp_drive(ctx))
.await?
.await
.is_err()
{
return Err(AgentDriverError::WarpDriveSyncFailed);
}
Ok(())
})
.await?;
// Set up and run the driver, reporting any errors back to the server.
let result: Result<(), AgentDriverError> = async {
// Pull relevant variables out of args before moving it into the closure.
let share_requests = args.share.share.clone();
let bedrock_inference_role = args.bedrock_inference_role.clone();
let bedrock_role_region = args.bedrock_role_region.clone();
let has_task_id = args.task_id.is_some();
let args_harness = args.harness;
// `--conversation` path (user-invoked local resume): validate before any task side
// effects so mismatches fail fast. The `--task-id` path derives its conversation id
// from the server-side task metadata inside `build_driver_options_and_task`. Both
// can currently be passed together (the worker server-side appends `--conversation`
// alongside `--task-id` for Slack/Linear followups); when both are set, the explicit
// `--conversation` value wins via the merge below.
if let Some(conversation_id) = args.conversation.as_deref() {
common::fetch_and_validate_conversation_harness(
server_api.clone(),
conversation_id,
args_harness,
)
.await?;
if !has_task_id {
if let Some(conversation_id) = args.conversation.as_deref() {
common::fetch_and_validate_conversation_harness(
server_api.clone(),
conversation_id,
args_harness,
)
.await?;
}
}
let resume_conversation_id = args.conversation.clone();
@@ -593,7 +697,8 @@ impl AgentDriverRunner {
// the fetched `AmbientAgentTask` (set by the server when linking the task to an
// existing conversation, e.g. via `run-cloud --conversation`).
let (mut driver_options, task, task_conversation_id) =
Self::build_driver_options_and_task(&foreground, args, &server_api).await?;
Self::build_driver_options_and_task(&foreground, args, &server_api, &setup_events)
.await?;
// Update the effective task ID so errors are reported correctly.
// This only matters if we created a task ID locally.
@@ -609,6 +714,16 @@ impl AgentDriverRunner {
#[cfg(not(target_family = "wasm"))]
if let Some(role_arn) = bedrock_inference_role {
// clap's `requires` constraint enforces this at parse time, so a missing
// region here means a caller is constructing `RunAgentArgs` directly
// without the flag. Fail loudly so callers don't silently fall back to a
// hard-coded STS region.
let role_region = bedrock_role_region.ok_or_else(|| {
AgentDriverError::AwsBedrockCredentialsFailed(
"--bedrock-role-region is required when --bedrock-inference-role is set"
.to_string(),
)
})?;
// Set the OIDC strategy on the UI thread and kick off the refresh; the
// returned future resolves when credentials are committed to the model.
let refresh_future = foreground
@@ -619,6 +734,7 @@ impl AgentDriverRunner {
AwsCredentialsRefreshStrategy::OidcManaged {
task_id: bedrock_task_id,
role_arn,
region: role_region,
},
);
refresh_aws_credentials(manager, ctx)
@@ -654,12 +770,16 @@ impl AgentDriverRunner {
// Pull conversation information, if we have it
if let Some(conversation_id) = resume_conversation_id {
driver_options.resume = Self::load_conversation_information(
&foreground,
conversation_id,
&task.harness,
)
.await?;
driver_options.resume = setup_events
.record_result(
SetupStep::ConversationResumeLoading,
Self::load_conversation_information(
&foreground,
conversation_id,
&task.harness,
),
)
.await?;
}
// Run the driver
@@ -701,6 +821,114 @@ impl AgentDriverRunner {
.map_err(|_| AgentDriverError::TeamMetadataRefreshTimeout)
}
async fn set_ambient_agent_task_id(
foreground: &ModelSpawner<Self>,
task_id: Option<AmbientAgentTaskId>,
) -> Result<(), AgentDriverError> {
foreground
.spawn(move |_, ctx| {
ServerApiProvider::handle(ctx)
.as_ref(ctx)
.get()
.set_ambient_agent_task_id(task_id);
})
.await?;
Ok(())
}
async fn fetch_task_git_credentials(
task_id_str: String,
ai_client: Arc<dyn AIClient>,
) -> anyhow::Result<Vec<GitCredential>> {
let workload_token = warp_isolation_platform::issue_workload_token(Some(
std::time::Duration::from_secs(5 * 60),
))
.await?
.token;
ai_client
.get_task_git_credentials(task_id_str, workload_token)
.await
}
async fn bootstrap_git_credentials_for_task(
foreground: &ModelSpawner<Self>,
task_id_str: &str,
args: &RunAgentArgs,
) -> Result<(), AgentDriverError> {
if warp_isolation_platform::detect().is_none() && args.configure_git_credentials_with_github
{
foreground
.spawn(|_, _| {
command::blocking::Command::new("gh")
.args(["auth", "setup-git"])
.spawn()
.map_err(|err| {
AgentDriverError::ConfigBuildFailed(anyhow::anyhow!(
"gh auth setup-git failed: {err:?}"
))
})
})
.await?
.map(|_| ())?;
return Ok(());
}
if !FeatureFlag::GitCredentialRefresh.is_enabled() {
return Ok(());
}
if task_id_str.parse::<AmbientAgentTaskId>().is_err() {
log::debug!(
"Skipping git credentials bootstrap: could not parse task ID '{task_id_str}'"
);
return Ok(());
}
let (ai_client, task_id_str) = foreground
.spawn({
let task_id_str = task_id_str.to_string();
move |_, ctx| {
let ai_client = ServerApiProvider::handle(ctx)
.as_ref(ctx)
.get_ai_client()
.clone();
(ai_client, task_id_str)
}
})
.await?;
let credentials = match Self::fetch_task_git_credentials(task_id_str, ai_client).await {
Ok(credentials) => credentials,
Err(err)
if err
.downcast_ref::<IsolationPlatformError>()
.is_some_and(|err| {
matches!(err, IsolationPlatformError::NoIsolationPlatformDetected)
}) =>
{
log::debug!("Skipping git credentials bootstrap: {err}");
return Ok(());
}
Err(err) => {
return Err(AgentDriverError::SkillResolutionFailed(format!(
"Failed to fetch git credentials before skill resolution: {err:#}"
)));
}
};
if credentials.is_empty() {
log::debug!("No git credentials returned before skill resolution");
return Ok(());
}
driver::git_credentials::configure_git_credentials(&credentials).map_err(|err| {
AgentDriverError::SkillResolutionFailed(format!(
"Failed to write git credentials before skill resolution: {err:#}"
))
})?;
log::info!("Git credentials configured before task setup");
Ok(())
}
/// Resolve the skill spec from args, if one was provided.
///
/// In sandboxed mode with a fully-qualified spec (org + repo), the repo is
@@ -710,6 +938,7 @@ impl AgentDriverRunner {
foreground: &ModelSpawner<Self>,
args: &RunAgentArgs,
working_dir: &Path,
setup_events: &SetupClientEventReporter,
) -> Result<Option<ResolvedSkill>, AgentDriverError> {
if !FeatureFlag::OzPlatformSkills.is_enabled() {
return Ok(None);
@@ -724,11 +953,17 @@ impl AgentDriverRunner {
let org = skill_spec.org.as_ref().expect("org checked above");
let repo_name = skill_spec.repo.as_ref().expect("repo checked above");
log::info!("Cloning {org}/{repo_name} for skill resolution in sandboxed mode");
clone_repo_for_skill(org, repo_name, working_dir)
.await
.map_err(|err| {
AgentDriverError::SkillResolutionFailed(format_skill_resolution_error(err))
})?;
setup_events
.record_result(SetupStep::SkillRepoClone, async {
clone_repo_for_skill(org, repo_name, working_dir)
.await
.map_err(|err| {
AgentDriverError::SkillResolutionFailed(format_skill_resolution_error(
err,
))
})
})
.await?;
}
let working_dir_buf = working_dir.to_path_buf();
@@ -756,6 +991,7 @@ impl AgentDriverRunner {
foreground: &ModelSpawner<Self>,
args: RunAgentArgs,
server_api: &Arc<dyn AIClient>,
setup_events: &SetupClientEventReporter,
) -> Result<(AgentDriverOptions, Task, Option<String>), AgentDriverError> {
// Get the working directory
let working_dir = match args.cwd.as_ref() {
@@ -765,8 +1001,12 @@ impl AgentDriverRunner {
}
.map_err(AgentDriverError::ConfigBuildFailed)?;
if let Some(task_id_str) = args.task_id.as_ref() {
Self::bootstrap_git_credentials_for_task(foreground, task_id_str, &args).await?;
}
// Resolve the skill, if we have one
let resolved_skill = Self::resolve_skill(foreground, &args, &working_dir).await?;
let resolved_skill =
Self::resolve_skill(foreground, &args, &working_dir, setup_events).await?;
// Extract variables we want to use later before moving args into the closure
let task_id_str = args.task_id.clone();
@@ -784,6 +1024,10 @@ impl AgentDriverRunner {
let should_share = (args.share.is_shared() || args.task_id.is_some())
&& FeatureFlag::AgentSharedSessions.is_enabled();
let third_party_harness_model_config = merged_config
.harness
.as_ref()
.and_then(|h| h.model_config());
let driver_options = driver::AgentDriverOptions {
working_dir: working_dir.clone(),
task_id,
@@ -795,6 +1039,7 @@ impl AgentDriverRunner {
cloud_providers: Vec::new(),
environment: None,
selected_harness: args.harness,
third_party_harness_model_config,
snapshot_disabled: args.snapshot.no_snapshot.then_some(true),
snapshot_upload_timeout: args
.snapshot
@@ -804,6 +1049,9 @@ impl AgentDriverRunner {
.snapshot
.snapshot_script_timeout
.map(|duration| duration.into()),
skip_initial_turn: args.skip_initial_turn,
strict_mcp_startup: args.strict_mcp_startup,
mcp_startup_timeout: args.mcp_startup_timeout.map(|duration| duration.into()),
};
Ok((merged_config, task, driver_options))
@@ -817,13 +1065,17 @@ impl AgentDriverRunner {
// The existing-task branch also surfaces the task's `conversation_id` (if any) so
// the caller can wire up resume without a separate `--conversation` arg.
let task_conversation_id = if let Some(task_id_str) = task_id_str {
Self::fetch_secrets_and_attachments(
foreground,
task_id_str,
&mut driver_options,
&mut task,
)
.await?
setup_events
.record_result(
SetupStep::TaskDataFetch,
Self::fetch_secrets_and_attachments(
foreground,
task_id_str,
&mut driver_options,
&mut task,
),
)
.await?
} else {
// Extract the prompt text that we'll pass up to the server when we create the task.
let prompt_for_task_creation = match &prompt {
@@ -847,9 +1099,13 @@ impl AgentDriverRunner {
.await?;
None
};
// Resolve environment and cloud providers.
Self::resolve_environment(foreground, environment_id, &mut driver_options).await?;
setup_events
.record_result(
SetupStep::EnvironmentResolution,
Self::resolve_environment(foreground, environment_id, &mut driver_options),
)
.await?;
Ok((driver_options, task, task_conversation_id))
}
@@ -889,15 +1145,8 @@ impl AgentDriverRunner {
}
};
foreground
.spawn(move |_, ctx| {
// Set the task ID on the ServerApi so it's sent with all subsequent requests.
ServerApiProvider::handle(ctx)
.as_ref(ctx)
.get()
.set_ambient_agent_task_id(task_id);
})
.await?;
// Set the task ID on the ServerApi so it's sent with all subsequent requests.
Self::set_ambient_agent_task_id(foreground, task_id).await?;
driver_options.task_id = task_id;
Ok(())
@@ -939,6 +1188,9 @@ impl AgentDriverRunner {
None
}
};
// Set the task ID on the ServerApi before any task-scoped server calls below, so failures
// during setup can still be reported with cloud-agent context.
Self::set_ambient_agent_task_id(foreground, parsed_task_id).await?;
// Fetch secrets, task metadata, regular attachments, and handoff snapshot
// attachments in parallel. The handoff snapshot fetch is independent of the
@@ -976,19 +1228,18 @@ impl AgentDriverRunner {
)
.await
};
let (secrets_result, attachments_result, task_metadata_result, handoff_snapshot_result) =
futures::future::join4(
task_secrets,
driver::attachments::fetch_and_download_attachments(
ai_client.clone(),
server_api.clone(),
task_id_str.clone(),
attachments_download_dir.clone(),
),
task_metadata,
handoff_snapshot,
)
.await;
let (secrets_result, attachments_result, task_metadata_result, handoff_snapshot_result) = futures::join!(
task_secrets,
driver::attachments::fetch_and_download_attachments(
ai_client.clone(),
server_api.clone(),
task_id_str.clone(),
attachments_download_dir.clone(),
),
task_metadata,
handoff_snapshot,
);
// Extract attachments_dir from successful result, log errors
let mut attachments_dir = match attachments_result {
@@ -1010,6 +1261,7 @@ impl AgentDriverRunner {
log::warn!("Failed to fetch handoff snapshot attachments: {e:#}");
}
}
let secrets = match secrets_result {
Ok(secrets) => secrets,
Err(err) => {
@@ -1028,55 +1280,52 @@ impl AgentDriverRunner {
}
}
};
let (parent_run_id, task_conversation_id, task_harness) = match task_metadata_result {
Ok(Some(task_metadata)) => {
// The task's harness is stored on the snapshot; if absent, it's the default Oz.
let task_harness = task_metadata
.agent_config_snapshot
.as_ref()
.and_then(|c| c.harness.as_ref())
.map(|h| h.harness_type)
.unwrap_or(Harness::Oz);
(
task_metadata.parent_run_id,
task_metadata.conversation_id,
Some(task_harness),
)
}
Ok(None) => (None, None, None),
Err(err) => {
log::warn!("Failed to fetch task metadata: {err:#}");
(None, None, None)
}
};
let (parent_run_id, task_conversation_id, task_harness, task_harness_model_config) =
match task_metadata_result {
Ok(Some(task_metadata)) => {
// The task's harness is stored on the snapshot; if absent, it's the default Oz.
let task_harness_config = task_metadata
.agent_config_snapshot
.as_ref()
.and_then(|c| c.harness.as_ref());
let task_harness = task_harness_config
.map(|h| h.harness_type)
.unwrap_or(Harness::Oz);
let task_harness_model_config =
task_harness_config.and_then(|h| h.model_config());
(
task_metadata.parent_run_id,
task_metadata.conversation_id,
Some(task_harness),
task_harness_model_config,
)
}
Ok(None) => (None, None, None, None),
Err(err) => {
log::warn!("Failed to fetch task metadata: {err:#}");
(None, None, None, None)
}
};
// Validate the requested `--harness` against the task's harness setting. This avoids the
// extra conversation-metadata roundtrip that would otherwise be needed downstream when the
// task is linked to an existing conversation, since task harness and conversation harness
// always match (the task spawned the conversation).
if let Some(task_harness) = task_harness {
if task_harness != driver_options.selected_harness {
return Err(AgentDriverError::TaskHarnessMismatch {
task_id: task_id_str,
expected: task_harness.to_string(),
got: driver_options.selected_harness.to_string(),
});
}
task.harness = reconcile_task_harness(
&task_id_str,
&mut driver_options.selected_harness,
task_harness,
)?;
}
// Set the task ID on the ServerApi so it's sent with all subsequent requests.
foreground
.spawn(move |_, ctx| {
ServerApiProvider::handle(ctx)
.as_ref(ctx)
.get()
.set_ambient_agent_task_id(parsed_task_id);
})
.await?;
driver_options.task_id = parsed_task_id;
driver_options.parent_run_id = parent_run_id;
driver_options.secrets = secrets;
// CLI flags continue to take precedence so users can still override per-invocation.
if driver_options.third_party_harness_model_config.is_none() {
driver_options.third_party_harness_model_config = task_harness_model_config;
}
// Update the task prompt to include the downloaded attachments dir
if let AgentRunPrompt::ServerSide {
@@ -1102,6 +1351,7 @@ impl AgentDriverRunner {
/// wraps the returned payload (if any) in [`driver::ResumeOptions::ThirdParty`]; each harness
/// owns its server call and error mapping. Returns `None` if a third-party harness has no
/// resume payload to surface.
#[tracing::instrument(skip_all, err, fields(tags.cloud_agent = true, conversation_id = conversation_id))]
async fn load_conversation_information(
foreground: &ModelSpawner<Self>,
conversation_id: String,
@@ -1133,7 +1383,6 @@ impl AgentDriverRunner {
"Failed to convert conversation data to AIConversation".into(),
)
})?;
Ok(Some(driver::ResumeOptions::Oz(Box::new(
ConversationRestorationInNewPaneType::Historical {
conversation,
@@ -1164,6 +1413,7 @@ impl AgentDriverRunner {
}
/// Resolve the environment and store into `driver_options`.
#[tracing::instrument(skip_all, err, fields(tags.cloud_agent = true, ?environment_id))]
async fn resolve_environment(
foreground: &ModelSpawner<Self>,
environment_id: Option<String>,
@@ -1207,6 +1457,7 @@ impl AgentDriverRunner {
}
/// Create the AgentDriver and start running the task.
#[tracing::instrument(skip_all, fields(tags.cloud_agent = true))]
fn create_and_run_driver(
ctx: &mut AppContext,
driver_options: driver::AgentDriverOptions,
@@ -1227,7 +1478,9 @@ impl AgentDriverRunner {
if let Some(share_requests) = share_requests {
driver.add_share_requests(share_requests, ctx);
}
let agent_future = driver.run(task, ctx);
let span =
tracing::info_span!("AgentDriver::run", tags.cloud_agent = true, ?task.model, ?task.harness);
let agent_future = driver.run(task, ctx).instrument(span);
ctx.spawn(agent_future, |_, result, ctx| match result {
Ok(()) => {
@@ -1251,6 +1504,11 @@ fn command_requires_auth(command: &CliCommand) -> bool {
AgentProfileCommand::List => true,
},
AgentCommand::List(_) => true,
AgentCommand::Get(_) => true,
AgentCommand::Create(_) => true,
AgentCommand::Update(_) => true,
AgentCommand::Delete(_) => true,
AgentCommand::Skills(_) => true,
},
CliCommand::Environment(environment_cmd) => match environment_cmd {
EnvironmentCommand::List => true,
@@ -1272,6 +1530,8 @@ fn command_requires_auth(command: &CliCommand) -> bool {
CliCommand::Model(model_cmd) => match model_cmd {
ModelCommand::List => true,
},
CliCommand::MemoryStore(_) => true,
CliCommand::Memory(_) => true,
CliCommand::Login => false,
CliCommand::Logout => false,
CliCommand::Whoami => true,
@@ -1282,6 +1542,7 @@ fn command_requires_auth(command: &CliCommand) -> bool {
CliCommand::Federate(_) => true,
CliCommand::HarnessSupport(_) => true,
CliCommand::Artifact(_) => true,
CliCommand::ApiKey(_) => true,
}
}
@@ -1365,6 +1626,8 @@ fn report_fatal_error(err: anyhow::Error, ctx: &mut AppContext) {
let _ = write!(&mut message, "\n=> {cause}");
}
tracing::event!(tracing::Level::ERROR, tags.cloud_agent = true, message);
#[cfg(not(target_family = "wasm"))]
{
if let Ok(path) = log_file_path() {
@@ -1389,6 +1652,7 @@ fn resolve_orchestration_harness_label() -> &'static str {
Some(Harness::Claude) => "claude",
Some(Harness::OpenCode) => "opencode",
Some(Harness::Gemini) => "gemini",
Some(Harness::Codex) => "codex",
Some(Harness::Unknown) | None => "unknown",
}
}
@@ -1408,6 +1672,11 @@ fn command_to_telemetry_event(command: &CliCommand) -> CliTelemetryEvent {
AgentProfileCommand::List => CliTelemetryEvent::AgentProfileList,
},
CliCommand::Agent(AgentCommand::List(_)) => CliTelemetryEvent::AgentList,
CliCommand::Agent(AgentCommand::Get(_)) => CliTelemetryEvent::AgentGet,
CliCommand::Agent(AgentCommand::Create(_)) => CliTelemetryEvent::AgentCreate,
CliCommand::Agent(AgentCommand::Update(_)) => CliTelemetryEvent::AgentUpdate,
CliCommand::Agent(AgentCommand::Delete(_)) => CliTelemetryEvent::AgentDelete,
CliCommand::Agent(AgentCommand::Skills(_)) => CliTelemetryEvent::AgentSkills,
CliCommand::Environment(EnvironmentCommand::List) => CliTelemetryEvent::EnvironmentList,
CliCommand::Environment(EnvironmentCommand::Create { .. }) => {
CliTelemetryEvent::EnvironmentCreate
@@ -1452,6 +1721,19 @@ fn command_to_telemetry_event(command: &CliCommand) -> CliTelemetryEvent {
},
},
CliCommand::Model(ModelCommand::List) => CliTelemetryEvent::ModelList,
CliCommand::MemoryStore(memory_store_cmd) => match memory_store_cmd {
MemoryStoreCommand::List => CliTelemetryEvent::MemoryStoreList,
MemoryStoreCommand::Get(_) => CliTelemetryEvent::MemoryStoreGetStore,
MemoryStoreCommand::Update(_) => CliTelemetryEvent::MemoryStoreUpdateStore,
MemoryStoreCommand::ListStoreAgents(_) => CliTelemetryEvent::MemoryStoreListStoreAgents,
},
CliCommand::Memory(memory_cmd) => match memory_cmd {
MemoryCommand::List(_) => CliTelemetryEvent::MemoryStoreListMemories,
MemoryCommand::Create(_) => CliTelemetryEvent::MemoryStoreCreateMemory,
MemoryCommand::Update(_) => CliTelemetryEvent::MemoryStoreUpdateMemory,
MemoryCommand::Delete(_) => CliTelemetryEvent::MemoryStoreDeleteMemory,
MemoryCommand::Versions(_) => CliTelemetryEvent::MemoryStoreListVersions,
},
CliCommand::Login => CliTelemetryEvent::Login,
CliCommand::Logout => CliTelemetryEvent::Logout,
CliCommand::Whoami => CliTelemetryEvent::Whoami,
@@ -1496,12 +1778,20 @@ fn command_to_telemetry_event(command: &CliCommand) -> CliTelemetryEvent {
success: finish_args.status == TaskStatus::Success,
}
}
HarnessSupportCommand::ReportShutdown(_) => {
CliTelemetryEvent::HarnessSupportReportShutdown
}
},
CliCommand::Artifact(artifact_cmd) => match artifact_cmd {
ArtifactCommand::Upload(_) => CliTelemetryEvent::ArtifactUpload,
ArtifactCommand::Get(_) => CliTelemetryEvent::ArtifactGet,
ArtifactCommand::Download(_) => CliTelemetryEvent::ArtifactDownload,
},
CliCommand::ApiKey(api_key_cmd) => match api_key_cmd {
ApiKeyCommand::List(_) => CliTelemetryEvent::ApiKeyList,
ApiKeyCommand::Create(_) => CliTelemetryEvent::ApiKeyCreate,
ApiKeyCommand::Expire(_) => CliTelemetryEvent::ApiKeyExpire,
},
}
}