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

Make direct-provider agent runs durable

Browse Source
This commit is contained in:
rkw6086
2026-08-14 22:02:15 -05:00
parent f4a04d0240
commit b079f036fa
50 changed files with 9473 additions and 3189 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/agent/api.rs
+8 -14
View File
@@ -21,10 +21,10 @@ use galaxy_core::features::FeatureFlag;
use galaxy_core::user_preferences::GetUserPreferences;
use galaxyui::{AppContext, EntityId, SingletonEntity as _};
use mcp::TemplatableMCPServerInfo;
pub use r#impl::generate_multi_agent_output;
pub(crate) use r#impl::prepare_direct_provider_params;
use serde::Serialize;
use super::{AIAgentAction, AIAgentInput, MCPContext, MCPServer, RequestMetadata, Suggestions};
use super::{AIAgentInput, MCPContext, MCPServer, RequestMetadata, Suggestions};
use crate::ai::agent::conversation::AIConversationId;
use crate::ai::ambient_agents::AmbientAgentTaskId;
use crate::ai::blocklist::{BlocklistAIPermissions, RequestInput, SessionContext};
@@ -112,7 +112,7 @@ pub struct RequestParams {
/// locally so ACP-provided Galaxy tools can be pinned to the exact pane.
pub terminal_view_id: Option<EntityId>,
pub input: Vec<AIAgentInput>,
/// Normalized results consumed directly by Rig-selected models.
/// Normalized action results appended to direct-provider run history.
pub tool_results: Vec<ToolResult>,
pub conversation_token: Option<ServerConversationToken>,
pub forked_from_conversation_token: Option<ServerConversationToken>,
@@ -151,8 +151,8 @@ pub struct RequestParams {
pub research_agent_enabled: bool,
pub orchestration_enabled: bool,
pub supported_tools_override: Option<Vec<warp_multi_agent_api::ToolType>>,
/// The root task ID for the conversation — needed for direct Bedrock streaming
/// since optimistic tasks don't appear in the proto task_context.
/// The root task ID used to anchor direct-provider projection when optimistic tasks are not
/// present in the proto task context.
pub root_task_id: Option<String>,
/// The conversation ID of the parent agent that spawned this child agent, if any.
pub parent_agent_id: Option<String>,
@@ -167,9 +167,8 @@ pub struct RequestParams {
/// Kept separately so `recall_tool_history` can search archived results even after
/// they've been summarized away from live history.
pub tool_result_archive: Vec<crate::ai::provider::types::ConversationMessage>,
/// Populated by direct-provider paths after building the message list.
/// Contains the full messages sent (old history + new input) so the controller
/// can store them back into the conversation for the next request cycle.
/// Populated while preparing a direct-provider run with the durable transcript that the
/// controller persists for restoration and future turns.
pub messages_sent:
std::sync::Arc<std::sync::Mutex<Vec<crate::ai::provider::types::ConversationMessage>>>,
/// Global rules (name, content) from the local CloudModel (AIFact/AIMemory).
@@ -177,15 +176,10 @@ pub struct RequestParams {
pub global_rules: Vec<(String, String)>,
}
/// Provider/runtime events consumed by the local conversation controller.
///
/// The legacy response envelope remains at the UI boundary while the local Rig runtime emits
/// executable tool proposals directly as Galaxy domain actions. This avoids translating Rig tool
/// calls into protobuf only to immediately translate them back before execution.
/// Response event projected into the local conversation controller.
#[derive(Debug)]
pub enum StreamEvent {
Response(warp_multi_agent_api::ResponseEvent),
ToolProposed(AIAgentAction),
}
pub type Event = Result<StreamEvent, Arc<AIApiError>>;
app/src/ai/agent/api/convert_conversation.rs
+2
View File
@@ -73,6 +73,7 @@ pub fn convert_conversation_data_to_ai_conversation(
let agent_conversation_data = match restoration_mode {
RestorationMode::Fork => AgentConversationData {
agent_backend: Default::default(),
active_provider_run_json: None,
server_conversation_token: None,
conversation_usage_metadata: usage_metadata,
reverted_action_ids: None,
@@ -96,6 +97,7 @@ pub fn convert_conversation_data_to_ai_conversation(
},
RestorationMode::Continue => AgentConversationData {
agent_backend: Default::default(),
active_provider_run_json: None,
server_conversation_token: Some(
metadata.server_conversation_token.as_str().to_string(),
),
app/src/ai/agent/api/impl.rs
+22 -262
View File
@@ -1,270 +1,10 @@
use std::collections::HashMap;
use std::sync::Arc;
use futures_util::StreamExt;
use galaxy_core::features::FeatureFlag;
use warp_multi_agent_api as api;
use super::convert_to::convert_input;
use super::{ConvertToAPITypeError, RequestParams, ResponseStream, StreamEvent};
use crate::ai::agent::{redaction, AIAgentInput};
use crate::ai::openai::translator as openai_translator;
use crate::ai::provider::ProviderConfig;
use crate::server::server_api::AIApiError;
use super::RequestParams;
use crate::ai::agent::redaction;
use crate::terminal::model::session::SessionType;
pub async fn generate_multi_agent_output(
provider_config: ProviderConfig,
mut params: RequestParams,
cancellation_rx: futures::channel::oneshot::Receiver<()>,
) -> Result<ResponseStream, ConvertToAPITypeError> {
let supported_tools_override = params.supported_tools_override.take();
let mut supported_tools = supported_tools_override
.clone()
.unwrap_or_else(|| get_supported_tools(&params));
remove_orchestration_tools_if_disabled(&mut supported_tools, params.orchestration_enabled);
let mut supported_cli_agent_tools =
supported_tools_override.unwrap_or_else(|| get_supported_cli_agent_tools(&params));
remove_orchestration_tools_if_disabled(
&mut supported_cli_agent_tools,
params.orchestration_enabled,
);
if params.should_redact_secrets {
redaction::redact_inputs(&mut params.input);
}
match &provider_config {
ProviderConfig::OpenAI(config) if config.use_rig => {
return Ok(crate::ai::runtime::rig_openai_response_stream(
config.clone(),
params,
supported_tools,
supported_cli_agent_tools,
cancellation_rx,
));
}
ProviderConfig::Bedrock(config) if config.use_rig => {
return match crate::ai::runtime::rig_bedrock_response_stream(
config.clone(),
params,
supported_tools,
supported_cli_agent_tools,
cancellation_rx,
)
.await
{
Ok(stream) => Ok(stream),
Err(error) => {
log::error!("[rig/bedrock] Runtime error: {error}");
let error = Arc::new(crate::server::server_api::AIApiError::Stream {
stream_type: "rig_bedrock",
source: error,
});
let (sender, receiver) = async_channel::unbounded();
let _ = sender.send(Err(error)).await;
Ok(Box::pin(receiver))
}
};
}
ProviderConfig::OpenAI(_) | ProviderConfig::Bedrock(_) | ProviderConfig::None => {}
}
let mut logging_metadata = HashMap::new();
if let Some(ref metadata) = params.metadata {
logging_metadata.insert(
"is_autodetected_user_query".to_owned(),
prost_types::Value {
kind: Some(prost_types::value::Kind::BoolValue(
metadata.is_autodetected_user_query,
)),
},
);
logging_metadata.insert(
"entrypoint".to_owned(),
prost_types::Value {
kind: Some(prost_types::value::Kind::StringValue(
metadata.entrypoint.entrypoint(),
)),
},
);
logging_metadata.insert(
"is_auto_resume_after_error".to_owned(),
prost_types::Value {
kind: Some(prost_types::value::Kind::BoolValue(
metadata.is_auto_resume_after_error,
)),
},
);
}
let emit_user_query_message = !params
.input
.iter()
.any(|input| matches!(input, AIAgentInput::CommandCompletionAssessment { .. }));
let mut request = api::Request {
task_context: Some(api::request::TaskContext {
tasks: params.tasks,
}),
input: Some(convert_input(params.input)?),
settings: Some(api::request::Settings {
model_config: Some(api::request::settings::ModelConfig {
base: params.model.clone().into(),
cli_agent: params.cli_agent_model.clone().into(),
computer_use_agent: params.computer_use_model.clone().into(),
base_model_context_window_limit: params.context_window_limit.unwrap_or(0),
..Default::default()
}),
rules_enabled: params.is_memory_enabled,
warp_drive_context_enabled: params.warp_drive_context_enabled,
web_context_retrieval_enabled: true,
supports_parallel_tool_calls: true,
use_anthropic_text_editor_tools: false,
planning_enabled: params.planning_enabled,
supports_create_files: true,
supported_tools: supported_tools.into_iter().map(Into::into).collect(),
supports_long_running_commands: true,
should_preserve_file_content_in_history: true,
supports_todos_ui: true,
supports_linked_code_blocks: FeatureFlag::LinkedCodeBlocks.is_enabled(),
supports_started_child_task_message: true,
// Galaxy's direct providers only receive tools with local schemas and
// executors. Hosted-only suggestion/orchestration capability bits must
// remain false so models do not plan around unavailable Warp services.
supports_suggest_prompt: false,
supports_read_image_files: FeatureFlag::ReadImageFiles.is_enabled(),
supports_reasoning_message: true,
api_keys: params.api_keys,
autonomy_level: params.autonomy_level.into(),
isolation_level: params.isolation_level.into(),
web_search_enabled: params.web_search_enabled,
supported_cli_agent_tools: supported_cli_agent_tools
.into_iter()
.map(Into::into)
.collect(),
supports_v4a_file_diffs: FeatureFlag::V4AFileDiffs.is_enabled(),
supports_summarization_via_message_replacement:
FeatureFlag::SummarizationViaMessageReplacement.is_enabled(),
supports_bundled_skills: FeatureFlag::BundledSkills.is_enabled(),
supports_research_agent: params.research_agent_enabled,
supports_orchestration_v2: false,
supports_background_computer_use: FeatureFlag::BackgroundComputerUse.is_enabled()
&& computer_use::background_supported(),
custom_model_providers: params.custom_model_providers,
custom_model_routers: params.custom_model_routers,
}),
metadata: Some(api::request::Metadata {
logging: logging_metadata,
conversation_id: params
.conversation_token
.as_ref()
.map(|token| token.as_str().to_string())
.unwrap_or_default(),
ambient_agent_task_id: params
.ambient_agent_task_id
.map(|id| id.to_string())
.unwrap_or_default(),
forked_from_conversation_id: if params.conversation_token.is_none() {
// We only include this param on our initial request to the server
// (when the forked conversation has not been assigned a new id yet).
params
.forked_from_conversation_token
.map(|token| token.as_str().to_string())
.unwrap_or_default()
} else {
String::new()
},
parent_agent_id: params.parent_agent_id.unwrap_or_default(),
agent_name: params.agent_name.unwrap_or_default(),
}),
existing_suggestions: params
.existing_suggestions
.map(|suggestions| suggestions.into()),
mcp_context: params.mcp_context.map(Into::into),
};
match provider_config {
ProviderConfig::OpenAI(config) => {
let translator_request = openai_translator::TranslatorRequest {
config,
model_id: params.model.as_str().to_string(),
root_task_id: params.root_task_id.clone(),
message_history: params.message_history.clone(),
tool_result_archive: params.tool_result_archive.clone(),
progressive_summary: params.progressive_summary.clone(),
messages_sent: params.messages_sent.clone(),
global_rules: params.global_rules.clone(),
emit_user_query_message,
};
match openai_translator::execute(translator_request, &mut request).await {
Ok(stream) => {
let output_stream = stream
.map(|event| event.map(StreamEvent::Response))
.take_until(cancellation_rx);
Ok(Box::pin(output_stream))
}
Err(e) => {
log::error!("[openai] Translator error: {e}");
let err = Arc::new(
crate::server::server_api::AIApiError::Stream {
stream_type: "openai_chat_completions",
source: anyhow::anyhow!("{e}"),
}
.into_quota_limit_if_provider_budget_exhausted(),
);
let (tx, rx) = async_channel::unbounded();
let _ = tx.send(Err(err)).await;
Ok(Box::pin(rx))
}
}
}
ProviderConfig::Bedrock(config) => {
let translator_request = crate::ai::bedrock::translator::TranslatorRequest {
config,
model_id: params.model.as_str().to_string(),
root_task_id: params.root_task_id.clone(),
bedrock_message_history: params.message_history.clone(),
bedrock_tool_result_archive: params.tool_result_archive.clone(),
bedrock_progressive_summary: params.progressive_summary.clone(),
bedrock_messages_sent: params.messages_sent.clone(),
global_rules: params.global_rules.clone(),
emit_user_query_message,
};
match crate::ai::bedrock::translator::execute(translator_request, &mut request).await {
Ok(stream) => {
let output_stream = stream
.map(|event| event.map(StreamEvent::Response))
.take_until(cancellation_rx);
Ok(Box::pin(output_stream))
}
Err(e) => {
log::error!("[bedrock] Translator error: {e}");
let err = Arc::new(crate::server::server_api::AIApiError::Stream {
stream_type: "bedrock",
source: anyhow::anyhow!("{e}"),
});
let (tx, rx) = async_channel::unbounded();
let _ = tx.send(Err(err)).await;
Ok(Box::pin(rx))
}
}
}
ProviderConfig::None => {
// No provider configured — do not fall back to Warp's cloud API.
let err = Arc::new(crate::server::server_api::AIApiError::Stream {
stream_type: "none",
source: anyhow::anyhow!(
"No AI runtime configured. Enable an agent runtime or model provider in settings."
),
});
let (tx, rx) = async_channel::unbounded();
let _ = tx.send(Err(err)).await;
Ok(Box::pin(rx))
}
}
}
fn remove_orchestration_tools_if_disabled(
supported_tools: &mut Vec<api::ToolType>,
orchestration_enabled: bool,
@@ -280,6 +20,26 @@ fn remove_orchestration_tools_if_disabled(
});
}
pub(crate) fn prepare_direct_provider_params(
params: &mut RequestParams,
) -> (Vec<api::ToolType>, Vec<api::ToolType>) {
let supported_tools_override = params.supported_tools_override.take();
let mut supported_tools = supported_tools_override
.clone()
.unwrap_or_else(|| get_supported_tools(params));
remove_orchestration_tools_if_disabled(&mut supported_tools, params.orchestration_enabled);
let mut supported_cli_agent_tools =
supported_tools_override.unwrap_or_else(|| get_supported_cli_agent_tools(params));
remove_orchestration_tools_if_disabled(
&mut supported_cli_agent_tools,
params.orchestration_enabled,
);
if params.should_redact_secrets {
redaction::redact_inputs(&mut params.input);
}
(supported_tools, supported_cli_agent_tools)
}
fn get_supported_tools(params: &RequestParams) -> Vec<api::ToolType> {
let mut supported_tools = vec![
api::ToolType::Grep,
app/src/ai/agent/conversation.rs
+99
View File
@@ -229,6 +229,9 @@ pub struct AIConversation {
/// Runtime responsible for executing this conversation.
agent_backend: AgentBackend,
/// Opaque, versioned snapshot of the active direct-provider run.
active_provider_run_json: Option<String>,
/// The server-generated unique "token" for this conversation.
///
/// This must be roundtripped to the server when sending follow-ups within a given conversation.
@@ -380,6 +383,7 @@ impl AIConversation {
has_opened_code_review: false,
conversation_usage_metadata: ConversationUsageMetadata::default(),
agent_backend,
active_provider_run_json: None,
server_conversation_token: None,
task_id: None,
forked_from_server_conversation_token: None,
@@ -539,6 +543,7 @@ impl AIConversation {
let (
agent_backend,
active_provider_run_json,
server_conversation_token,
forked_from_server_conversation_token,
conversation_usage_metadata,
@@ -589,6 +594,7 @@ impl AIConversation {
};
(
data.agent_backend,
data.active_provider_run_json,
server_conversation_token,
forked_from_server_conversation_token,
conversation_usage_metadata,
@@ -611,6 +617,7 @@ impl AIConversation {
AgentBackend::default(),
None,
None,
None,
ConversationUsageMetadata::default(),
HashSet::new(),
Vec::new(),
@@ -663,6 +670,7 @@ impl AIConversation {
has_opened_code_review: false,
conversation_usage_metadata,
agent_backend,
active_provider_run_json,
server_conversation_token,
task_id: run_id.as_deref().and_then(|id| id.parse().ok()),
forked_from_server_conversation_token,
@@ -705,6 +713,14 @@ impl AIConversation {
&self.agent_backend
}
pub(crate) fn active_provider_run_json(&self) -> Option<&str> {
self.active_provider_run_json.as_deref()
}
pub(crate) fn set_active_provider_run_json(&mut self, snapshot: Option<String>) {
self.active_provider_run_json = snapshot;
}
/// Updates the backend of a conversation that has not produced agent output.
///
/// Provider failures without output are safe to retry through a newly enabled runtime. Once
@@ -2109,6 +2125,81 @@ impl AIConversation {
Ok(())
}
/// Reopens an exact restored exchange for continued provider projection.
///
/// This only restores the process-local stream association; it never adds input or provider
/// history, so the persisted provider run remains the sole continuation source of truth.
pub(crate) fn provider_projection_target(
&self,
response_stream_id: &ResponseStreamId,
) -> Option<(TaskId, AIAgentExchangeId)> {
let mut exchanges = self
.added_exchanges_by_response
.get(response_stream_id)?
.iter();
let target = exchanges.next()?;
exchanges
.next()
.is_none()
.then(|| (target.task_id.clone(), target.exchange_id))
}
pub(crate) fn rebind_provider_projection(
&mut self,
task_id: &TaskId,
exchange_id: AIAgentExchangeId,
response_stream_id: ResponseStreamId,
terminal_surface_id: EntityId,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) -> Result<(), UpdateConversationError> {
let Some(task) = self.task_store.get(task_id) else {
return Err(UpdateConversationError::TaskNotFound);
};
if !task.exchanges().any(|exchange| exchange.id == exchange_id) {
return if self.exchange_with_id(exchange_id).is_some() {
Err(UpdateConversationError::ExchangeTaskMismatch)
} else {
Err(UpdateConversationError::ExchangeNotFound)
};
}
if self
.added_exchanges_by_response
.contains_key(&response_stream_id)
{
return Err(UpdateConversationError::ResponseStreamAlreadyBound);
}
let exchange = self.get_exchange_to_update(exchange_id)?;
let previous_status = std::mem::replace(
&mut exchange.output_status,
AIAgentOutputStatus::Streaming { output: None },
);
let output = match previous_status {
AIAgentOutputStatus::Streaming { output } => output,
AIAgentOutputStatus::Finished { finished_output } => match finished_output {
FinishedAIAgentOutput::Cancelled { output, .. }
| FinishedAIAgentOutput::Error { output, .. } => output,
FinishedAIAgentOutput::Success { output } => Some(output),
},
};
exchange.output_status = AIAgentOutputStatus::Streaming { output };
exchange.finish_time = None;
self.added_exchanges_by_response.insert(
response_stream_id,
Vec1::new(AddedExchange {
task_id: task_id.clone(),
exchange_id,
}),
);
ctx.emit(BlocklistAIHistoryEvent::UpdatedStreamingExchange {
exchange_id,
terminal_surface_id,
conversation_id: self.id,
is_hidden: self.hidden_exchanges.contains(&exchange_id),
});
Ok(())
}
pub fn append_reassigned_exchange(
&mut self,
response_stream_id: &ResponseStreamId,
@@ -2255,6 +2346,9 @@ impl AIConversation {
action: AIAgentAction,
ctx: &mut ModelContext<BlocklistAIHistoryModel>,
) -> Result<(), UpdateConversationError> {
if self.contains_action(&action.id) {
return Ok(());
}
let added_exchanges = self
.added_exchanges_by_response
.get(stream_id)
@@ -3886,6 +3980,7 @@ impl AIConversation {
.collect(),
conversation_data: AgentConversationData {
agent_backend: self.agent_backend.clone(),
active_provider_run_json: self.active_provider_run_json.clone(),
server_conversation_token: self
.server_conversation_token
.clone()
@@ -4763,6 +4858,10 @@ fn cleanup_conversation_search_temp_dir(
pub enum UpdateConversationError {
#[error("Exchange not found.")]
ExchangeNotFound,
#[error("Exchange does not belong to the persisted task.")]
ExchangeTaskMismatch,
#[error("Response stream is already bound to an exchange.")]
ResponseStreamAlreadyBound,
#[error("Could not update task: {0:?}")]
UpdateTask(#[from] UpdateTaskError),
#[error("Could not update upgrade optimistic task for server task: {0:?}")]
app/src/ai/agent/conversation_tests.rs
+13
View File
@@ -193,6 +193,19 @@ fn restored_conversation_defaults_autoexecute_override_when_not_persisted() {
);
}
#[test]
fn restored_conversation_retains_active_provider_run_json() {
let snapshot = r#"{"version":1,"run":{"state":"awaiting_model"}}"#;
let conversation_data = AgentConversationData {
active_provider_run_json: Some(snapshot.to_string()),
..Default::default()
};
let conversation = restored_conversation(Some(conversation_data));
assert_eq!(conversation.active_provider_run_json(), Some(snapshot));
}
#[test]
fn restored_conversation_uses_persisted_last_event_sequence() {
let conversation_data: AgentConversationData =