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

Harden direct provider lifecycle handling

Browse Source
This commit is contained in:
rkw6086
2026-08-16 09:11:22 -05:00
parent 04bc7f7055
commit ae3a8c7b40
10 changed files with 1003 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/blocklist/controller.rs
+334 -26
View File
@@ -78,9 +78,9 @@ use crate::ai::provider::types::{ContentPart, ConversationMessage, MessageConten
use crate::ai::remote_logging::{self, RemoteLogLevel, RemoteLogRecord};
use crate::ai::runtime::{
prepare_provider_run, provider_runtime_for_request, PreparedProviderRun, ProviderActionContext,
ProviderRunBlock, ProviderRunCoordinator, ProviderRunProfile, ProviderRunResponseProjector,
ProviderToolExecutionRef, ProviderToolLifecycleOutcome, RuntimeResponseConfig,
BASE_PROVIDER_PROFILE, CLI_MONITOR_PROVIDER_PROFILE,
ProviderRunBlock, ProviderRunCoordinator, ProviderRunProfile, ProviderRunProjection,
ProviderRunResponseProjector, ProviderToolExecutionRef, ProviderToolLifecycleOutcome,
RuntimeResponseConfig, BASE_PROVIDER_PROFILE, CLI_MONITOR_PROVIDER_PROFILE,
};
use crate::ai::AIRequestUsageModel;
use crate::cloud_object::model::persistence::CloudModel;
@@ -1287,8 +1287,256 @@ fn provider_boundary_intent(
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum ProviderLlmLifecyclePhase {
Requested,
Started,
RetryScheduled,
Finished,
}
impl ProviderLlmLifecyclePhase {
fn event(self) -> &'static str {
match self {
Self::Requested => "provider_model_turn_requested",
Self::Started => "provider_model_turn_started",
Self::RetryScheduled => "provider_model_turn_retry_scheduled",
Self::Finished => "provider_model_turn_finished",
}
}
fn message(self) -> &'static str {
match self {
Self::Requested => "Provider model turn requested",
Self::Started => "Provider model turn started",
Self::RetryScheduled => "Provider model turn retry scheduled",
Self::Finished => "Provider model turn finished",
}
}
fn llm_finished(self) -> bool {
matches!(self, Self::Finished)
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
struct ProviderLlmLifecycle {
phase: ProviderLlmLifecyclePhase,
work_id: ExternalWorkId,
profile: String,
runtime_id: String,
model_id: String,
runtime_request_id: Option<String>,
retry_attempt: u32,
elapsed_ms: Option<u64>,
stop_reason: Option<String>,
tool_call_count: Option<usize>,
error_kind: Option<String>,
error_recoverable: Option<bool>,
error: Option<String>,
}
fn provider_llm_lifecycle(projection: &ProviderRunProjection) -> Option<ProviderLlmLifecycle> {
let lifecycle = match projection {
ProviderRunProjection::ModelTurnRequested {
work_id,
profile,
runtime_id,
model_id,
retry_attempt,
} => ProviderLlmLifecycle {
phase: ProviderLlmLifecyclePhase::Requested,
work_id: work_id.clone(),
profile: profile.as_str().to_string(),
runtime_id: runtime_id.clone(),
model_id: model_id.clone(),
runtime_request_id: None,
retry_attempt: *retry_attempt,
elapsed_ms: None,
stop_reason: None,
tool_call_count: None,
error_kind: None,
error_recoverable: None,
error: None,
},
ProviderRunProjection::ModelTurnStarted {
work_id,
profile,
runtime_id,
model_id,
runtime_request_id,
retry_attempt,
elapsed_ms,
} => ProviderLlmLifecycle {
phase: ProviderLlmLifecyclePhase::Started,
work_id: work_id.clone(),
profile: profile.as_str().to_string(),
runtime_id: runtime_id.clone(),
model_id: model_id.clone(),
runtime_request_id: Some(runtime_request_id.clone()),
retry_attempt: *retry_attempt,
elapsed_ms: Some(*elapsed_ms),
stop_reason: None,
tool_call_count: None,
error_kind: None,
error_recoverable: None,
error: None,
},
ProviderRunProjection::ModelTurnFinished {
work_id,
profile,
runtime_id,
model_id,
stop_reason,
retry_attempt,
elapsed_ms,
tool_call_count,
} => ProviderLlmLifecycle {
phase: ProviderLlmLifecyclePhase::Finished,
work_id: work_id.clone(),
profile: profile.as_str().to_string(),
runtime_id: runtime_id.clone(),
model_id: model_id.clone(),
runtime_request_id: None,
retry_attempt: *retry_attempt,
elapsed_ms: Some(*elapsed_ms),
stop_reason: Some(format!("{stop_reason:?}")),
tool_call_count: Some(*tool_call_count),
error_kind: None,
error_recoverable: None,
error: None,
},
ProviderRunProjection::ModelRetry {
work_id,
profile,
runtime_id,
model_id,
retry_attempt,
elapsed_ms,
error,
} => ProviderLlmLifecycle {
phase: ProviderLlmLifecyclePhase::RetryScheduled,
work_id: work_id.clone(),
profile: profile.as_str().to_string(),
runtime_id: runtime_id.clone(),
model_id: model_id.clone(),
runtime_request_id: None,
retry_attempt: *retry_attempt,
elapsed_ms: Some(*elapsed_ms),
stop_reason: None,
tool_call_count: None,
error_kind: Some(format!("{:?}", error.kind)),
error_recoverable: Some(error.recoverable),
error: Some(error.message.clone()),
},
ProviderRunProjection::ModelEvent { .. } | ProviderRunProjection::ToolBatchReady { .. } => {
return None
}
};
Some(lifecycle)
}
#[cfg(not(target_family = "wasm"))]
fn provider_llm_lifecycle_remote_log_record(
conversation_id: AIConversationId,
stream_id: &ResponseStreamId,
lifecycle: &ProviderLlmLifecycle,
) -> RemoteLogRecord {
let level = if lifecycle.phase == ProviderLlmLifecyclePhase::RetryScheduled {
RemoteLogLevel::Warn
} else {
RemoteLogLevel::Info
};
RemoteLogRecord {
level,
message: lifecycle.phase.message().to_string(),
context: serde_json::json!({
"event": lifecycle.phase.event(),
"conversation_id": conversation_id.to_string(),
"stream_id": stream_id.as_str(),
"provider_run_id": lifecycle.work_id.run_id.as_str(),
"provider_epoch": lifecycle.work_id.epoch.get(),
"provider_work_id": format!(
"{}:{}",
lifecycle.work_id.run_id.as_str(),
lifecycle.work_id.epoch.get()
),
"profile": lifecycle.profile,
"runtime_id": lifecycle.runtime_id,
"model_id": lifecycle.model_id,
"runtime_request_id": lifecycle.runtime_request_id,
"retry_attempt": lifecycle.retry_attempt,
"elapsed_ms": lifecycle.elapsed_ms,
"stop_reason": lifecycle.stop_reason,
"tool_call_count": lifecycle.tool_call_count,
"error_kind": lifecycle.error_kind,
"error_recoverable": lifecycle.error_recoverable,
"error": lifecycle.error.as_deref().map(remote_logging::sanitize_error),
"llm_finished": lifecycle.phase.llm_finished(),
"response_stream_terminal": false,
}),
}
}
#[cfg(not(target_family = "wasm"))]
fn provider_run_terminal_remote_log_record(
conversation_id: AIConversationId,
stream_id: &ResponseStreamId,
run: &ProviderRun,
outcome: &ProviderRunOutcome,
) -> RemoteLogRecord {
let (level, outcome_name, llm_finished, stop_reason, failure_kind, error) = match outcome {
ProviderRunOutcome::Completed(completion) => (
RemoteLogLevel::Info,
"completed",
true,
Some(format!("{:?}", completion.stop_reason)),
None,
None,
),
ProviderRunOutcome::Failed(failure) => (
RemoteLogLevel::Error,
"failed",
false,
None,
Some(format!("{:?}", failure.kind)),
Some(remote_logging::sanitize_error(&failure.message)),
),
ProviderRunOutcome::Cancelled { reason } => (
RemoteLogLevel::Info,
"cancelled",
false,
None,
None,
Some(remote_logging::sanitize_error(reason)),
),
};
RemoteLogRecord {
level,
message: "Provider run finished".to_string(),
context: serde_json::json!({
"event": "provider_run_finished",
"conversation_id": conversation_id.to_string(),
"stream_id": stream_id.as_str(),
"provider_run_id": run.id().as_str(),
"provider_epoch": run.epoch().get(),
"profile": run.profile().as_str(),
"model_turn_count": run.model_turns(),
"model_retry_count": run.model_retries(),
"outcome": outcome_name,
"stop_reason": stop_reason,
"failure_kind": failure_kind,
"error": error,
"llm_finished": llm_finished,
"provider_run_finished": true,
"response_stream_terminal": true,
}),
}
}
enum ProviderDriveMessage {
Response(warp_multi_agent_api::ResponseEvent),
Lifecycle(ProviderLlmLifecycle),
Checkpoint {
checkpoint: ActiveProviderRunCheckpoint,
acknowledgement: oneshot::Sender<Result<(), String>>,
@@ -5075,6 +5323,14 @@ impl BlocklistAIController {
.drive_until_blocked_with_checkpoint(
turn_control,
|projection| {
if let Some(lifecycle) = provider_llm_lifecycle(&projection) {
projection_sender
.try_send(ProviderDriveMessage::Lifecycle(lifecycle))
.map_err(|_| {
"provider lifecycle projection receiver was closed"
.to_string()
})?;
}
for event in run.projector.project(projection)? {
projection_sender
.try_send(ProviderDriveMessage::Response(event))
@@ -5140,6 +5396,19 @@ impl BlocklistAIController {
ctx,
);
}
ProviderDriveMessage::Lifecycle(lifecycle) => {
#[cfg(not(target_family = "wasm"))]
remote_logging::log_model_event(
ctx,
provider_llm_lifecycle_remote_log_record(
conversation_id,
stream_id,
&lifecycle,
),
);
#[cfg(target_family = "wasm")]
let _ = lifecycle;
}
ProviderDriveMessage::Checkpoint {
checkpoint,
acknowledgement,
@@ -5894,6 +6163,28 @@ impl BlocklistAIController {
self.drive_active_provider_run(conversation_id, ctx);
}
fn finalize_completed_provider_conversation(
&self,
conversation_id: AIConversationId,
ctx: &mut ModelContext<Self>,
) {
let history_model = BlocklistAIHistoryModel::handle(ctx);
let should_finalize = history_model
.as_ref(ctx)
.conversation_status(&conversation_id)
.is_some_and(|status| status != &ConversationStatus::Success);
if should_finalize {
history_model.update(ctx, |history_model, ctx| {
history_model.update_conversation_status(
self.terminal_surface_id,
conversation_id,
ConversationStatus::Success,
ctx,
);
});
}
}
fn finish_active_provider_run(
&mut self,
conversation_id: AIConversationId,
@@ -5935,32 +6226,49 @@ impl BlocklistAIController {
ctx,
);
}
if matches!(outcome, ProviderRunOutcome::Cancelled { .. }) {
let cancellation_reason =
self.active_provider_runs[&conversation_id].cancellation_reason;
if let Some(reason) = cancellation_reason {
let status = match reason.conversation_outcome() {
CancellationOutcome::KeepInProgress => ConversationStatus::InProgress,
CancellationOutcome::Succeeded => ConversationStatus::Success,
CancellationOutcome::Cancelled => ConversationStatus::Cancelled,
CancellationOutcome::FinalizedExternally => {
self.cleanup_active_provider_run(
#[cfg(not(target_family = "wasm"))]
remote_logging::log_model_event(
ctx,
provider_run_terminal_remote_log_record(
conversation_id,
&stream_id,
run.coordinator.run(),
&outcome,
),
);
match outcome {
ProviderRunOutcome::Completed(_) => {
self.finalize_completed_provider_conversation(conversation_id, ctx);
}
// Failed outcomes are finalized by the projected InternalError event.
ProviderRunOutcome::Failed(_) => {}
ProviderRunOutcome::Cancelled { .. } => {
let cancellation_reason =
self.active_provider_runs[&conversation_id].cancellation_reason;
if let Some(reason) = cancellation_reason {
let status = match reason.conversation_outcome() {
CancellationOutcome::KeepInProgress => ConversationStatus::InProgress,
CancellationOutcome::Succeeded => ConversationStatus::Success,
CancellationOutcome::Cancelled => ConversationStatus::Cancelled,
CancellationOutcome::FinalizedExternally => {
self.cleanup_active_provider_run(
conversation_id,
&stream_id,
&response_stream,
ctx,
);
return;
}
};
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| {
history_model.update_conversation_status(
self.terminal_surface_id,
conversation_id,
&stream_id,
&response_stream,
status,
ctx,
);
return;
}
};
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| {
history_model.update_conversation_status(
self.terminal_surface_id,
conversation_id,
status,
ctx,
);
});
});
}
}
}
self.cleanup_active_provider_run(conversation_id, &stream_id, &response_stream, ctx);
app/src/ai/blocklist/controller/response_stream.rs
+24
View File
@@ -313,6 +313,8 @@ impl ResponseStream {
"ask_user_question_enabled": params.ask_user_question_enabled,
"orchestration_enabled": params.orchestration_enabled,
"is_remote_session": params.session_context.is_remote(),
"llm_finished": false,
"response_stream_terminal": false,
"identifiers": serde_json::to_value(ai_identifiers).unwrap_or_else(|_| serde_json::json!({})),
}),
},
@@ -372,6 +374,14 @@ impl ResponseStream {
"reason".to_string(),
serde_json::json!(stream_finished_reason_name(&finished_event.reason)),
);
context.insert(
"llm_finished".to_string(),
serde_json::json!(stream_finished_llm_finished(&finished_event.reason)),
);
context.insert(
"response_stream_terminal".to_string(),
serde_json::json!(true),
);
context.insert(
"elapsed_ms".to_string(),
serde_json::json!(self.time_to_latest_event.num_milliseconds()),
@@ -417,6 +427,11 @@ impl ResponseStream {
serde_json::json!(self.time_to_latest_event.num_milliseconds()),
);
context.insert("recovery".to_string(), serde_json::json!(recovery));
context.insert("llm_finished".to_string(), serde_json::json!(false));
context.insert(
"response_stream_terminal".to_string(),
serde_json::json!(true),
);
context.insert(
"error".to_string(),
serde_json::json!(remote_logging::sanitize_error(error)),
@@ -1115,6 +1130,15 @@ fn stream_finished_reason_name(
}
}
#[cfg(not(target_family = "wasm"))]
fn stream_finished_llm_finished(reason: &Option<response_event::stream_finished::Reason>) -> bool {
matches!(
reason,
None | Some(response_event::stream_finished::Reason::Done(_))
| Some(response_event::stream_finished::Reason::MaxTokenLimit(_))
)
}
#[cfg(not(target_family = "wasm"))]
fn token_usage_context(
token_usage: &[response_event::stream_finished::TokenUsage],
app/src/ai/blocklist/controller/response_stream_tests.rs
+21 -1
View File
@@ -1,4 +1,24 @@
use super::is_interactive_remote_command;
use warp_multi_agent_api::response_event::stream_finished;
use super::{is_interactive_remote_command, stream_finished_llm_finished};
#[test]
fn response_finish_reason_reports_whether_the_llm_completed() {
assert!(stream_finished_llm_finished(&None));
assert!(stream_finished_llm_finished(&Some(
stream_finished::Reason::Done(stream_finished::Done {})
)));
assert!(stream_finished_llm_finished(&Some(
stream_finished::Reason::MaxTokenLimit(stream_finished::ReachedMaxTokenLimit {})
)));
assert!(!stream_finished_llm_finished(&Some(
stream_finished::Reason::Other(stream_finished::Other {})
)));
assert!(!stream_finished_llm_finished(&Some(
stream_finished::Reason::LlmUnavailable(stream_finished::LlmUnavailable {})
)));
}
#[test]
fn raw_interactive_ssh_is_treated_as_remote_for_acp() {
for command in [
app/src/ai/blocklist/controller_tests.rs
+259 -5
View File
@@ -4,10 +4,11 @@ use std::sync::{Arc, Mutex};
use ai::agent::action::{AskUserQuestionItem, AskUserQuestionType};
use chrono::Local;
use galaxy_agent_core::{
CompletedModelTurn, ContentPart, ConversationMessage, ExternalWorkId, MessageContent,
MessageRole, PermissionKind, PermissionRequest, ProviderRun, ProviderRunId, ProviderRunLimits,
ProviderRunState, ProviderRunStep, RunEpoch, RuntimeCapabilities, StopReason, ToolCall,
TurnRequest, Usage,
AgentError, AgentErrorKind, CompletedModelTurn, ContentPart, ConversationMessage,
ExternalWorkId, MessageContent, MessageRole, PermissionKind, PermissionRequest, ProviderRun,
ProviderRunFailure, ProviderRunFailureKind, ProviderRunId, ProviderRunLimits,
ProviderRunOutcome, ProviderRunState, ProviderRunStep, RunEpoch, RuntimeCapabilities,
StopReason, ToolCall, TurnRequest, Usage,
};
use galaxy_core::command::ExitCode;
use uuid::Uuid;
@@ -26,9 +27,11 @@ use crate::ai::agent::{
use crate::ai::ambient_agents::AmbientAgentTaskId;
use crate::ai::blocklist::{
BlocklistAIHistoryEvent, BlocklistAIHistoryModel, PendingAttachment, PendingFile, RequestInput,
ResponseStream, ResponseStreamId,
ResponseStream, ResponseStreamId, StartAgentExecutor,
};
use crate::ai::llms::LLMId;
use crate::ai::remote_logging::RemoteLogLevel;
use crate::ai::runtime::ProviderRunProjection;
use crate::persistence::model::{AcpConversationData, AgentBackend};
use crate::terminal::model::block::{BlockId, BlockState};
use crate::test_util::settings::initialize_history_persistence_for_tests;
@@ -38,6 +41,128 @@ fn new_ambient_agent_task_id() -> AmbientAgentTaskId {
Uuid::new_v4().to_string().parse().unwrap()
}
#[test]
fn provider_lifecycle_logs_expose_llm_completion_and_sanitize_errors() {
let conversation_id = AIConversationId::new();
let stream_id = ResponseStreamId::new_for_test();
let work_id = ExternalWorkId {
run_id: ProviderRunId::new("run-1"),
epoch: RunEpoch::new(3),
};
let mut retry_error = AgentError::new(AgentErrorKind::Transport, "sk-secret connection failed");
retry_error.recoverable = true;
let projections = [
ProviderRunProjection::ModelTurnRequested {
work_id: work_id.clone(),
profile: "base".into(),
runtime_id: "rig:openai".to_owned(),
model_id: "test-model".to_owned(),
retry_attempt: 0,
},
ProviderRunProjection::ModelTurnStarted {
work_id: work_id.clone(),
profile: "base".into(),
runtime_id: "rig:openai".to_owned(),
model_id: "test-model".to_owned(),
runtime_request_id: "request-1".to_owned(),
retry_attempt: 0,
elapsed_ms: 12,
},
ProviderRunProjection::ModelRetry {
work_id: work_id.clone(),
profile: "base".into(),
runtime_id: "rig:openai".to_owned(),
model_id: "test-model".to_owned(),
retry_attempt: 1,
elapsed_ms: 120_000,
error: retry_error,
},
ProviderRunProjection::ModelTurnFinished {
work_id,
profile: "base".into(),
runtime_id: "rig:openai".to_owned(),
model_id: "test-model".to_owned(),
stop_reason: StopReason::Completed,
retry_attempt: 1,
elapsed_ms: 140,
tool_call_count: 2,
},
];
let records = projections
.iter()
.map(|projection| {
let lifecycle = super::provider_llm_lifecycle(projection).unwrap();
super::provider_llm_lifecycle_remote_log_record(conversation_id, &stream_id, &lifecycle)
})
.collect::<Vec<_>>();
assert_eq!(
records
.iter()
.map(|record| record.context["event"].as_str().unwrap())
.collect::<Vec<_>>(),
[
"provider_model_turn_requested",
"provider_model_turn_started",
"provider_model_turn_retry_scheduled",
"provider_model_turn_finished",
]
);
assert_eq!(records[0].context["llm_finished"], false);
assert_eq!(records[1].context["llm_finished"], false);
assert_eq!(records[2].context["llm_finished"], false);
assert_eq!(records[3].context["llm_finished"], true);
assert_eq!(records[2].level, RemoteLogLevel::Warn);
assert_eq!(records[2].context["error"], "[redacted] connection failed");
assert_eq!(records[3].context["provider_run_id"], "run-1");
assert_eq!(records[3].context["provider_epoch"], 3);
assert_eq!(records[3].context["profile"], "base");
assert_eq!(records[3].context["runtime_id"], "rig:openai");
assert_eq!(records[3].context["model_id"], "test-model");
assert_eq!(records[3].context["stop_reason"], "Completed");
assert_eq!(records[3].context["tool_call_count"], 2);
}
#[test]
fn provider_terminal_logs_distinguish_clean_completion_from_failure() {
let conversation_id = AIConversationId::new();
let stream_id = ResponseStreamId::new_for_test();
let run = ProviderRun::new(
"run-1",
Vec::new(),
crate::ai::runtime::BASE_PROVIDER_PROFILE,
ProviderRunLimits::default(),
);
let completed = super::provider_run_terminal_remote_log_record(
conversation_id,
&stream_id,
&run,
&ProviderRunOutcome::Completed(galaxy_agent_core::ProviderRunCompletion {
stop_reason: StopReason::Completed,
}),
);
let failed = super::provider_run_terminal_remote_log_record(
conversation_id,
&stream_id,
&run,
&ProviderRunOutcome::Failed(ProviderRunFailure {
kind: ProviderRunFailureKind::RetryLimitExceeded,
message: "sk-secret timeout".to_owned(),
source: None,
}),
);
assert_eq!(completed.context["llm_finished"], true);
assert_eq!(completed.context["provider_run_finished"], true);
assert_eq!(completed.context["response_stream_terminal"], true);
assert_eq!(completed.context["outcome"], "completed");
assert_eq!(failed.level, RemoteLogLevel::Error);
assert_eq!(failed.context["llm_finished"], false);
assert_eq!(failed.context["outcome"], "failed");
assert_eq!(failed.context["failure_kind"], "RetryLimitExceeded");
assert_eq!(failed.context["error"], "[redacted] timeout");
}
fn ask_user_question_action(action_id: &str) -> AIAgentAction {
AIAgentAction {
id: AIAgentActionId::from(action_id.to_string()),
@@ -1849,6 +1974,135 @@ fn mock_response_stream_updates_history_through_controller() {
});
}
#[test]
fn completed_provider_run_with_prior_action_resolves_child_completion_wait() {
App::test((), |mut app| async move {
initialize_app_for_terminal_view(&mut app);
let terminal = add_window_with_terminal(&mut app, None);
let start_agent_executor = app.add_model(StartAgentExecutor::new);
let terminal_surface_id = terminal.read(&app, |terminal, _| terminal.id());
let stream_id = ResponseStreamId::new_for_test();
let history_model = BlocklistAIHistoryModel::handle(&app);
let (parent_conversation_id, child_conversation_id, child_task_id) =
history_model.update(&mut app, |history_model, ctx| {
let parent_conversation_id = history_model.start_new_conversation(
terminal_surface_id,
false,
false,
false,
ctx,
);
let child_conversation_id = history_model.start_new_child_conversation(
terminal_surface_id,
"child".to_owned(),
parent_conversation_id,
None,
ctx,
);
let child_task_id = history_model
.conversation(&child_conversation_id)
.expect("child conversation should exist")
.get_root_task_id()
.clone();
history_model
.update_conversation_for_new_request_input(
RequestInput {
conversation_id: child_conversation_id,
input_messages: HashMap::from([(child_task_id.clone(), vec![])]),
working_directory: None,
model_id: LLMId::from("test-model"),
coding_model_id: LLMId::from("test-coding-model"),
cli_agent_model_id: LLMId::from("test-cli-agent-model"),
computer_use_model_id: LLMId::from("test-computer-use-model"),
shared_session_response_initiator: None,
request_start_ts: Local::now(),
supported_tools_override: None,
},
stream_id.clone(),
terminal_surface_id,
ctx,
)
.expect("child request should be recorded");
history_model.initialize_output_for_response_stream(
&stream_id,
child_conversation_id,
terminal_surface_id,
response_event::StreamInit {
request_id: "provider-request".to_owned(),
conversation_id: "provider-conversation".to_owned(),
run_id: "provider-run".to_owned(),
},
ctx,
);
(parent_conversation_id, child_conversation_id, child_task_id)
});
let dispatch = start_agent_executor.update(&mut app, |executor, ctx| {
executor.reattach(
AIAgentActionId::from("run-agents-action".to_owned()),
"child".to_owned(),
parent_conversation_id,
child_conversation_id,
None,
ctx,
)
});
history_model.update(&mut app, |history_model, ctx| {
history_model
.apply_domain_tool_proposal(
&stream_id,
child_conversation_id,
terminal_surface_id,
AIAgentAction {
id: AIAgentActionId::from("prior-tool-call".to_owned()),
task_id: child_task_id,
action: AIAgentActionType::FileGlob {
patterns: vec!["*.rs".to_owned()],
path: None,
},
requires_result: true,
tool_name: Some("file_glob".to_owned()),
},
ctx,
)
.expect("provider tool proposal should attach to child history");
history_model.mark_response_stream_completed_successfully(
&stream_id,
child_conversation_id,
terminal_surface_id,
ctx,
);
});
history_model.read(&app, |history_model, _| {
let child = history_model
.conversation(&child_conversation_id)
.expect("child conversation should exist");
assert_eq!(child.count_all_actions(), 1);
assert_eq!(child.status(), &ConversationStatus::InProgress);
});
assert!(matches!(
dispatch.receiver.try_recv(),
Err(async_channel::TryRecvError::Empty)
));
terminal.update(&mut app, |terminal, ctx| {
terminal.ai_controller().update(ctx, |controller, ctx| {
controller.finalize_completed_provider_conversation(child_conversation_id, ctx);
});
});
history_model.read(&app, |history_model, _| {
assert_eq!(
history_model
.conversation(&child_conversation_id)
.map(|conversation| conversation.status()),
Some(&ConversationStatus::Success)
);
});
assert!(dispatch.receiver.try_recv().is_ok());
});
}
/// When an agent command exits the shell, the conversation must be finalized as
/// `Error` (not `Cancelled`), and a subsequent `ManuallyCancelled` (as fired by
/// the pane-close path) must not overwrite that failure.
app/src/ai/runtime/event_translator.rs
+3 -1
View File
@@ -88,7 +88,9 @@ impl ProviderRunResponseProjector {
})
}
ProviderRunProjection::ModelEvent { event, .. } => self.translator.translate(event),
ProviderRunProjection::ModelRetry { .. }
ProviderRunProjection::ModelTurnRequested { .. }
| ProviderRunProjection::ModelTurnFinished { .. }
| ProviderRunProjection::ModelRetry { .. }
| ProviderRunProjection::ToolBatchReady { .. } => Ok(Vec::new()),
}
}
app/src/ai/runtime/event_translator_tests.rs
+4
View File
@@ -54,8 +54,12 @@ fn restored_provider_projection_skips_stream_initialization() {
assert!(projector
.project(ProviderRunProjection::ModelTurnStarted {
work_id: work_id.clone(),
profile: galaxy_agent_core::ProviderRequestProfile::new("base"),
runtime_id: "runtime".to_owned(),
model_id: "model".to_owned(),
runtime_request_id: "request".to_owned(),
retry_attempt: 0,
elapsed_ms: 1,
})
.unwrap()
.is_empty());
app/src/ai/runtime/mod.rs
+3 -2
View File
@@ -8,8 +8,9 @@ pub(crate) use event_translator::{
ProviderRunResponseProjector, RuntimeResponseConfig, RuntimeResponseTranslator,
};
pub(crate) use provider_run_coordinator::{
ProviderRunBlock, ProviderRunCoordinator, ProviderRunProfile, ProviderToolExecutionRef,
ProviderToolLifecycleOutcome, BASE_PROVIDER_PROFILE, CLI_MONITOR_PROVIDER_PROFILE,
ProviderRunBlock, ProviderRunCoordinator, ProviderRunProfile, ProviderRunProjection,
ProviderToolExecutionRef, ProviderToolLifecycleOutcome, BASE_PROVIDER_PROFILE,
CLI_MONITOR_PROVIDER_PROFILE,
};
pub(crate) use rig::{
prepare_provider_run, provider_runtime_for_request, PreparedProviderRun, ProviderActionContext,
app/src/ai/runtime/provider_run_coordinator.rs
+169 -26
View File
@@ -2,6 +2,7 @@ use std::collections::{BTreeMap, BTreeSet};
use std::error::Error;
use std::fmt;
use std::sync::Arc;
use std::time::Duration;
use futures::future::BoxFuture;
use futures::StreamExt;
@@ -12,18 +13,43 @@ use galaxy_agent_core::{
ProviderRunOutcome, ProviderRunPhase, ProviderRunProtocolError, ProviderRunState,
ProviderRunStep, RunEpoch, RuntimeKind, StopReason, ToolEvent, TurnControl, TurnRequest, Usage,
};
use instant::Instant;
use warpui::r#async::FutureExt as _;
use crate::ai::agent::conversation::AIConversationId;
pub(crate) const BASE_PROVIDER_PROFILE: &str = "base";
pub(crate) const CLI_MONITOR_PROVIDER_PROFILE: &str = "cli-monitor";
const PROVIDER_MODEL_START_TIMEOUT: Duration = Duration::from_secs(120);
const PROVIDER_MODEL_EVENT_IDLE_TIMEOUT: Duration = Duration::from_secs(300);
#[derive(Clone, Debug, PartialEq)]
pub(crate) enum ProviderRunProjection {
ModelTurnRequested {
work_id: ExternalWorkId,
profile: ProviderRequestProfile,
runtime_id: String,
model_id: String,
retry_attempt: u32,
},
ModelTurnStarted {
work_id: ExternalWorkId,
profile: ProviderRequestProfile,
runtime_id: String,
model_id: String,
runtime_request_id: String,
retry_attempt: u32,
elapsed_ms: u64,
},
ModelTurnFinished {
work_id: ExternalWorkId,
profile: ProviderRequestProfile,
runtime_id: String,
model_id: String,
stop_reason: StopReason,
retry_attempt: u32,
elapsed_ms: u64,
tool_call_count: usize,
},
ModelEvent {
work_id: ExternalWorkId,
@@ -31,7 +57,11 @@ pub(crate) enum ProviderRunProjection {
},
ModelRetry {
work_id: ExternalWorkId,
profile: ProviderRequestProfile,
runtime_id: String,
model_id: String,
retry_attempt: u32,
elapsed_ms: u64,
error: AgentError,
},
ToolBatchReady {
@@ -133,6 +163,8 @@ impl ProviderRunProfile {
pub(crate) struct ProviderRunCoordinator {
run: ProviderRun,
profiles: BTreeMap<String, ProviderRunProfile>,
model_start_timeout: Duration,
model_event_idle_timeout: Duration,
}
impl ProviderRunCoordinator {
@@ -171,7 +203,18 @@ impl ProviderRunCoordinator {
for (profile, config) in &profiles {
validate_profile_runtime(profile, config.runtime.as_ref())?;
}
Ok(Self { run, profiles })
Ok(Self {
run,
profiles,
model_start_timeout: PROVIDER_MODEL_START_TIMEOUT,
model_event_idle_timeout: PROVIDER_MODEL_EVENT_IDLE_TIMEOUT,
})
}
#[cfg(test)]
fn set_model_timeouts(&mut self, start: Duration, event_idle: Duration) {
self.model_start_timeout = start;
self.model_event_idle_timeout = event_idle;
}
pub(crate) fn run(&self) -> &ProviderRun {
@@ -395,22 +438,67 @@ impl ProviderRunCoordinator {
.iter()
.map(|tool| tool.name.clone())
.collect::<BTreeSet<_>>();
let request = request_for_model_call(profile.request, &call);
let stream = match profile.runtime.start_turn(request, control).await {
Ok(stream) => stream,
Err(error) => {
self.handle_model_failure(&call.work_id, error, project)?;
let runtime_id = profile.runtime.descriptor().id.clone();
let model_id = profile.request.model.as_str().to_string();
if !self.project_or_fail(
ProviderRunProjection::ModelTurnRequested {
work_id: call.work_id.clone(),
profile: call.profile.clone(),
runtime_id: runtime_id.clone(),
model_id: model_id.clone(),
retry_attempt: call.retry_attempt,
},
project,
)? {
return Ok(());
}
let request = request_for_model_call(profile.request.clone(), &call);
let started_at = Instant::now();
let stream = match profile
.runtime
.start_turn(request, control)
.with_timeout(self.model_start_timeout)
.await
{
Ok(Ok(stream)) => stream,
Ok(Err(error)) => {
self.handle_model_failure(&call, &profile, started_at, error, project)?;
return Ok(());
}
Err(_) => {
self.handle_model_failure(
&call,
&profile,
started_at,
provider_timeout_error("start", self.model_start_timeout),
project,
)?;
return Ok(());
}
};
futures::pin_mut!(stream);
let mut buffer = ModelTurnBuffer::default();
while let Some(event) = stream.next().await {
let event = match event {
Ok(event) => event,
Err(error) => {
self.handle_model_failure(&call.work_id, error, project)?;
loop {
let event = match stream
.next()
.with_timeout(self.model_event_idle_timeout)
.await
{
Ok(Some(Ok(event))) => event,
Ok(Some(Err(error))) => {
self.handle_model_failure(&call, &profile, started_at, error, project)?;
return Ok(());
}
Ok(None) => break,
Err(_) => {
self.handle_model_failure(
&call,
&profile,
started_at,
provider_timeout_error("event", self.model_event_idle_timeout),
project,
)?;
return Ok(());
}
};
@@ -418,7 +506,9 @@ impl ProviderRunCoordinator {
AgentEvent::TurnStarted { runtime_request_id } => {
if buffer.started {
self.handle_model_failure(
&call.work_id,
&call,
&profile,
started_at,
protocol_error("provider emitted more than one TurnStarted event"),
project,
)?;
@@ -426,7 +516,9 @@ impl ProviderRunCoordinator {
}
if runtime_request_id.is_empty() {
self.handle_model_failure(
&call.work_id,
&call,
&profile,
started_at,
protocol_error("provider emitted an empty runtime request ID"),
project,
)?;
@@ -436,8 +528,12 @@ impl ProviderRunCoordinator {
if !self.project_or_fail(
ProviderRunProjection::ModelTurnStarted {
work_id: call.work_id.clone(),
profile: call.profile.clone(),
runtime_id: runtime_id.clone(),
model_id: model_id.clone(),
runtime_request_id,
retry_attempt: call.retry_attempt,
elapsed_ms: elapsed_millis(started_at),
},
project,
)? {
@@ -445,7 +541,7 @@ impl ProviderRunCoordinator {
}
}
AgentEvent::TextDelta { text } => {
if !self.ensure_model_started(&call.work_id, &buffer, project)? {
if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? {
return Ok(());
}
buffer.text.push_str(&text);
@@ -460,7 +556,7 @@ impl ProviderRunCoordinator {
}
}
AgentEvent::ReasoningDelta { text } => {
if !self.ensure_model_started(&call.work_id, &buffer, project)? {
if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? {
return Ok(());
}
buffer.reasoning.push_str(&text);
@@ -475,7 +571,7 @@ impl ProviderRunCoordinator {
}
}
AgentEvent::ReasoningCompleted { text, signature } => {
if !self.ensure_model_started(&call.work_id, &buffer, project)? {
if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? {
return Ok(());
}
if !text.is_empty() {
@@ -495,13 +591,13 @@ impl ProviderRunCoordinator {
AgentEvent::Tool {
event: ToolEvent::Proposed { call: tool_call },
} => {
if !self.ensure_model_started(&call.work_id, &buffer, project)? {
if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? {
return Ok(());
}
buffer.tool_calls.push(tool_call);
}
AgentEvent::UsageUpdated { usage } => {
if !self.ensure_model_started(&call.work_id, &buffer, project)? {
if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? {
return Ok(());
}
buffer.usage.clone_from(&usage);
@@ -519,7 +615,22 @@ impl ProviderRunCoordinator {
}
}
AgentEvent::TurnStopped { reason } => {
if !self.ensure_model_started(&call.work_id, &buffer, project)? {
if !self.ensure_model_started(&call, &profile, started_at, &buffer, project)? {
return Ok(());
}
if !self.project_or_fail(
ProviderRunProjection::ModelTurnFinished {
work_id: call.work_id.clone(),
profile: call.profile.clone(),
runtime_id: runtime_id.clone(),
model_id: model_id.clone(),
stop_reason: reason.clone(),
retry_attempt: call.retry_attempt,
elapsed_ms: elapsed_millis(started_at),
tool_call_count: buffer.tool_calls.len(),
},
project,
)? {
return Ok(());
}
if reason == StopReason::Cancelled {
@@ -547,7 +658,9 @@ impl ProviderRunCoordinator {
| AgentEvent::UserInputAccepted { .. }
| AgentEvent::RuntimeNotice { .. } => {
self.handle_model_failure(
&call.work_id,
&call,
&profile,
started_at,
protocol_error(
"direct-provider transport emitted a non-model lifecycle event",
),
@@ -559,7 +672,9 @@ impl ProviderRunCoordinator {
}
self.handle_model_failure(
&call.work_id,
&call,
&profile,
started_at,
protocol_error("provider stream ended before TurnStopped"),
project,
)?;
@@ -568,7 +683,9 @@ impl ProviderRunCoordinator {
fn ensure_model_started<F>(
&mut self,
work_id: &ExternalWorkId,
call: &ProviderModelCall,
profile: &ProviderRunProfile,
started_at: Instant,
buffer: &ModelTurnBuffer,
project: &mut F,
) -> Result<bool, ProviderRunCoordinatorError>
@@ -579,7 +696,9 @@ impl ProviderRunCoordinator {
return Ok(true);
}
self.handle_model_failure(
work_id,
call,
profile,
started_at,
protocol_error("provider emitted model output before TurnStarted"),
project,
)?;
@@ -588,14 +707,18 @@ impl ProviderRunCoordinator {
fn handle_model_failure<F>(
&mut self,
work_id: &ExternalWorkId,
call: &ProviderModelCall,
profile: &ProviderRunProfile,
started_at: Instant,
error: AgentError,
project: &mut F,
) -> Result<(), ProviderRunCoordinatorError>
where
F: FnMut(ProviderRunProjection) -> Result<(), String>,
{
let disposition = self.run.register_model_failure(work_id, error.clone())?;
let disposition = self
.run
.register_model_failure(&call.work_id, error.clone())?;
if disposition == ModelFailureDisposition::RetryScheduled {
let retry_attempt = match self.run.state() {
ProviderRunState::AwaitingModel { call } => call.retry_attempt,
@@ -616,8 +739,12 @@ impl ProviderRunCoordinator {
};
self.project_or_fail(
ProviderRunProjection::ModelRetry {
work_id: work_id.clone(),
work_id: call.work_id.clone(),
profile: call.profile.clone(),
runtime_id: profile.runtime.descriptor().id.clone(),
model_id: profile.request.model.as_str().to_string(),
retry_attempt,
elapsed_ms: elapsed_millis(started_at),
error,
},
project,
@@ -683,6 +810,22 @@ impl ModelTurnBuffer {
}
}
fn elapsed_millis(started_at: Instant) -> u64 {
u64::try_from(started_at.elapsed().as_millis()).unwrap_or(u64::MAX)
}
fn provider_timeout_error(stage: &str, timeout: Duration) -> AgentError {
let mut error = AgentError::new(
AgentErrorKind::Transport,
format!(
"provider model {stage} timed out after {} seconds",
timeout.as_secs()
),
);
error.recoverable = true;
error
}
fn validate_profile_runtime(
profile: &str,
runtime: &dyn AgentRuntime,
app/src/ai/runtime/provider_run_coordinator_tests.rs
+183 -3
View File
@@ -1,5 +1,7 @@
use std::collections::VecDeque;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Mutex;
use std::time::Duration;
use async_trait::async_trait;
use galaxy_agent_core::{
@@ -66,6 +68,65 @@ impl AgentRuntime for ScriptedRuntime {
}
}
#[derive(Clone, Copy)]
enum FirstAttemptStall {
Start,
Event,
}
struct StallingRuntime {
descriptor: RuntimeDescriptor,
first_attempt_stall: FirstAttemptStall,
attempts: AtomicUsize,
requests: Mutex<Vec<TurnRequest>>,
}
impl StallingRuntime {
fn new(first_attempt_stall: FirstAttemptStall) -> Self {
Self {
descriptor: RuntimeDescriptor {
id: "stalling".to_string(),
display_name: "Stalling provider".to_string(),
kind: RuntimeKind::Provider,
capabilities: RuntimeCapabilities::provider(),
},
first_attempt_stall,
attempts: AtomicUsize::new(0),
requests: Mutex::new(Vec::new()),
}
}
fn requests(&self) -> Vec<TurnRequest> {
self.requests.lock().unwrap().clone()
}
}
#[async_trait]
impl AgentRuntime for StallingRuntime {
fn descriptor(&self) -> &RuntimeDescriptor {
&self.descriptor
}
async fn start_turn(
&self,
request: TurnRequest,
_control: TurnControl,
) -> Result<AgentEventStream, AgentError> {
self.requests.lock().unwrap().push(request);
let attempt = self.attempts.fetch_add(1, Ordering::Relaxed);
if attempt == 0 {
match self.first_attempt_stall {
FirstAttemptStall::Start => return futures::future::pending().await,
FirstAttemptStall::Event => {
let started = futures::stream::iter(vec![started("request-stalled")]);
return Ok(Box::pin(started.chain(futures::stream::pending())));
}
}
}
Ok(Box::pin(futures::stream::iter(answer_turn().unwrap())))
}
}
fn request() -> TurnRequest {
let mut request = TurnRequest::new(
"test-model",
@@ -690,7 +751,9 @@ async fn recoverable_start_failure_retries_the_same_work_identity() {
retry_attempt,
..
} => Some((work_id.clone(), *retry_attempt)),
ProviderRunProjection::ModelTurnStarted { .. }
ProviderRunProjection::ModelTurnRequested { .. }
| ProviderRunProjection::ModelTurnStarted { .. }
| ProviderRunProjection::ModelTurnFinished { .. }
| ProviderRunProjection::ModelEvent { .. }
| ProviderRunProjection::ToolBatchReady { .. } => None,
})
@@ -703,7 +766,9 @@ async fn recoverable_start_failure_retries_the_same_work_identity() {
retry_attempt: 1,
..
} => Some(work_id.clone()),
ProviderRunProjection::ModelTurnStarted { .. }
ProviderRunProjection::ModelTurnRequested { .. }
| ProviderRunProjection::ModelTurnStarted { .. }
| ProviderRunProjection::ModelTurnFinished { .. }
| ProviderRunProjection::ModelEvent { .. }
| ProviderRunProjection::ModelRetry { .. }
| ProviderRunProjection::ToolBatchReady { .. } => None,
@@ -713,6 +778,119 @@ async fn recoverable_start_failure_retries_the_same_work_identity() {
assert_eq!(retry.1, 1);
}
fn assert_single_retry_lifecycle(
projections: &[ProviderRunProjection],
expected_timeout_stage: &str,
expected_initial_start: bool,
) {
let mut work_ids = Vec::new();
let mut phases = Vec::new();
let mut retry_error = None;
for projection in projections {
match projection {
ProviderRunProjection::ModelTurnRequested {
work_id,
retry_attempt,
..
} => {
work_ids.push(work_id.clone());
phases.push(format!("requested:{retry_attempt}"));
}
ProviderRunProjection::ModelTurnStarted {
work_id,
retry_attempt,
..
} => {
work_ids.push(work_id.clone());
phases.push(format!("started:{retry_attempt}"));
}
ProviderRunProjection::ModelRetry {
work_id,
retry_attempt,
error,
..
} => {
work_ids.push(work_id.clone());
phases.push(format!("retry:{retry_attempt}"));
retry_error = Some(error);
}
ProviderRunProjection::ModelTurnFinished {
work_id,
retry_attempt,
..
} => {
work_ids.push(work_id.clone());
phases.push(format!("finished:{retry_attempt}"));
}
ProviderRunProjection::ModelEvent { .. }
| ProviderRunProjection::ToolBatchReady { .. } => {}
}
}
let expected = if expected_initial_start {
vec![
"requested:0",
"started:0",
"retry:1",
"requested:1",
"started:1",
"finished:1",
]
} else {
vec![
"requested:0",
"retry:1",
"requested:1",
"started:1",
"finished:1",
]
};
assert_eq!(phases, expected);
assert!(work_ids.windows(2).all(|ids| ids[0] == ids[1]));
let retry_error = retry_error.expect("timeout retry error");
assert_eq!(retry_error.kind, AgentErrorKind::Transport);
assert!(retry_error.recoverable);
assert!(retry_error.message.contains(expected_timeout_stage));
}
#[tokio::test]
async fn model_start_timeout_retries_the_same_work_identity() {
let runtime = Arc::new(StallingRuntime::new(FirstAttemptStall::Start));
let mut coordinator = coordinator(runtime.clone());
coordinator.set_model_timeouts(Duration::from_millis(10), Duration::from_secs(1));
let mut projections = Vec::new();
let (_sender, control) = turn_control();
let block = coordinator
.drive_until_blocked(control, collect_projection(&mut projections))
.await
.unwrap();
assert!(matches!(block, ProviderRunBlock::AwaitingDriver { .. }));
assert_eq!(runtime.requests().len(), 2);
assert_eq!(coordinator.run().model_retries(), 1);
assert_single_retry_lifecycle(&projections, "start timed out", false);
}
#[tokio::test]
async fn model_event_idle_timeout_retries_the_same_work_identity() {
let runtime = Arc::new(StallingRuntime::new(FirstAttemptStall::Event));
let mut coordinator = coordinator(runtime.clone());
coordinator.set_model_timeouts(Duration::from_secs(1), Duration::from_millis(10));
let mut projections = Vec::new();
let (_sender, control) = turn_control();
let block = coordinator
.drive_until_blocked(control, collect_projection(&mut projections))
.await
.unwrap();
assert!(matches!(block, ProviderRunBlock::AwaitingDriver { .. }));
assert_eq!(runtime.requests().len(), 2);
assert_eq!(coordinator.run().model_retries(), 1);
assert_single_retry_lifecycle(&projections, "event timed out", true);
}
#[tokio::test]
async fn tool_projection_failure_synthesizes_a_result_and_fails_the_run() {
let runtime = Arc::new(ScriptedRuntime::new(vec![tool_turn()]));
@@ -724,7 +902,9 @@ async fn tool_projection_failure_synthesizes_a_result_and_fails_the_run() {
ProviderRunProjection::ToolBatchReady { .. } => {
Err("task projection disappeared".to_string())
}
ProviderRunProjection::ModelTurnStarted { .. }
ProviderRunProjection::ModelTurnRequested { .. }
| ProviderRunProjection::ModelTurnStarted { .. }
| ProviderRunProjection::ModelTurnFinished { .. }
| ProviderRunProjection::ModelEvent { .. }
| ProviderRunProjection::ModelRetry { .. } => Ok(()),
})