Harden direct provider lifecycle handling
This commit is contained in:
@@ -132,8 +132,11 @@ context_size = 128000
|
|||||||
|
|
||||||
Key invariants:
|
Key invariants:
|
||||||
- Every direct-provider `AgentRuntime::start_turn` performs exactly one model call; only `ProviderRun` may schedule another turn or retry
|
- Every direct-provider `AgentRuntime::start_turn` performs exactly one model call; only `ProviderRun` may schedule another turn or retry
|
||||||
|
- Direct-provider model calls allow 120 seconds for stream startup and 300 seconds between stream events; either timeout is a recoverable transport failure that enters the existing bounded retry lifecycle with the same work identity
|
||||||
|
- Direct-provider remote telemetry records requested, started, retry-scheduled, and finished model-turn phases with explicit `llm_finished` state; root `provider_run_finished` records distinguish clean completion from failure or cancellation and mark the response stream terminal
|
||||||
- `use_rig` and provider selection may choose request/transport details but must never choose lifecycle ownership
|
- `use_rig` and provider selection may choose request/transport details but must never choose lifecycle ownership
|
||||||
- Direct-provider output may be projected through `ResponseStream`, but provider progress must not depend on response-stream result draining or `AfterStreamFinished`
|
- Direct-provider output may be projected through `ResponseStream`, but provider progress must not depend on response-stream result draining or `AfterStreamFinished`
|
||||||
|
- A clean direct-provider `ProviderRunOutcome::Completed` explicitly finalizes the conversation as `Success` after terminal output projection, even if earlier turns added tool actions; child-completion waits rely on that status
|
||||||
- Provider actions and results must correlate by `(conversation_id, run_id, epoch, call_id)`; stale or duplicate callbacks must not advance a run
|
- Provider actions and results must correlate by `(conversation_id, run_id, epoch, call_id)`; stale or duplicate callbacks must not advance a run
|
||||||
- Active provider runs must checkpoint before external work, persist without credentials, normalize unsafe restored states, and reconcile command state before continuing
|
- Active provider runs must checkpoint before external work, persist without credentials, normalize unsafe restored states, and reconcile command state before continuing
|
||||||
- Known tools are in `KNOWN_TOOLS` in `response_translator.rs`; definitions are built by `tool_definition_for_name()` in `convert_request.rs`
|
- Known tools are in `KNOWN_TOOLS` in `response_translator.rs`; definitions are built by `tool_definition_for_name()` in `convert_request.rs`
|
||||||
|
|||||||
@@ -78,9 +78,9 @@ use crate::ai::provider::types::{ContentPart, ConversationMessage, MessageConten
|
|||||||
use crate::ai::remote_logging::{self, RemoteLogLevel, RemoteLogRecord};
|
use crate::ai::remote_logging::{self, RemoteLogLevel, RemoteLogRecord};
|
||||||
use crate::ai::runtime::{
|
use crate::ai::runtime::{
|
||||||
prepare_provider_run, provider_runtime_for_request, PreparedProviderRun, ProviderActionContext,
|
prepare_provider_run, provider_runtime_for_request, PreparedProviderRun, ProviderActionContext,
|
||||||
ProviderRunBlock, ProviderRunCoordinator, ProviderRunProfile, ProviderRunResponseProjector,
|
ProviderRunBlock, ProviderRunCoordinator, ProviderRunProfile, ProviderRunProjection,
|
||||||
ProviderToolExecutionRef, ProviderToolLifecycleOutcome, RuntimeResponseConfig,
|
ProviderRunResponseProjector, ProviderToolExecutionRef, ProviderToolLifecycleOutcome,
|
||||||
BASE_PROVIDER_PROFILE, CLI_MONITOR_PROVIDER_PROFILE,
|
RuntimeResponseConfig, BASE_PROVIDER_PROFILE, CLI_MONITOR_PROVIDER_PROFILE,
|
||||||
};
|
};
|
||||||
use crate::ai::AIRequestUsageModel;
|
use crate::ai::AIRequestUsageModel;
|
||||||
use crate::cloud_object::model::persistence::CloudModel;
|
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 {
|
enum ProviderDriveMessage {
|
||||||
Response(warp_multi_agent_api::ResponseEvent),
|
Response(warp_multi_agent_api::ResponseEvent),
|
||||||
|
Lifecycle(ProviderLlmLifecycle),
|
||||||
Checkpoint {
|
Checkpoint {
|
||||||
checkpoint: ActiveProviderRunCheckpoint,
|
checkpoint: ActiveProviderRunCheckpoint,
|
||||||
acknowledgement: oneshot::Sender<Result<(), String>>,
|
acknowledgement: oneshot::Sender<Result<(), String>>,
|
||||||
@@ -5075,6 +5323,14 @@ impl BlocklistAIController {
|
|||||||
.drive_until_blocked_with_checkpoint(
|
.drive_until_blocked_with_checkpoint(
|
||||||
turn_control,
|
turn_control,
|
||||||
|projection| {
|
|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)? {
|
for event in run.projector.project(projection)? {
|
||||||
projection_sender
|
projection_sender
|
||||||
.try_send(ProviderDriveMessage::Response(event))
|
.try_send(ProviderDriveMessage::Response(event))
|
||||||
@@ -5140,6 +5396,19 @@ impl BlocklistAIController {
|
|||||||
ctx,
|
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 {
|
ProviderDriveMessage::Checkpoint {
|
||||||
checkpoint,
|
checkpoint,
|
||||||
acknowledgement,
|
acknowledgement,
|
||||||
@@ -5894,6 +6163,28 @@ impl BlocklistAIController {
|
|||||||
self.drive_active_provider_run(conversation_id, ctx);
|
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(
|
fn finish_active_provider_run(
|
||||||
&mut self,
|
&mut self,
|
||||||
conversation_id: AIConversationId,
|
conversation_id: AIConversationId,
|
||||||
@@ -5935,32 +6226,49 @@ impl BlocklistAIController {
|
|||||||
ctx,
|
ctx,
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
if matches!(outcome, ProviderRunOutcome::Cancelled { .. }) {
|
#[cfg(not(target_family = "wasm"))]
|
||||||
let cancellation_reason =
|
remote_logging::log_model_event(
|
||||||
self.active_provider_runs[&conversation_id].cancellation_reason;
|
ctx,
|
||||||
if let Some(reason) = cancellation_reason {
|
provider_run_terminal_remote_log_record(
|
||||||
let status = match reason.conversation_outcome() {
|
conversation_id,
|
||||||
CancellationOutcome::KeepInProgress => ConversationStatus::InProgress,
|
&stream_id,
|
||||||
CancellationOutcome::Succeeded => ConversationStatus::Success,
|
run.coordinator.run(),
|
||||||
CancellationOutcome::Cancelled => ConversationStatus::Cancelled,
|
&outcome,
|
||||||
CancellationOutcome::FinalizedExternally => {
|
),
|
||||||
self.cleanup_active_provider_run(
|
);
|
||||||
|
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,
|
conversation_id,
|
||||||
&stream_id,
|
status,
|
||||||
&response_stream,
|
|
||||||
ctx,
|
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);
|
self.cleanup_active_provider_run(conversation_id, &stream_id, &response_stream, ctx);
|
||||||
|
|||||||
@@ -313,6 +313,8 @@ impl ResponseStream {
|
|||||||
"ask_user_question_enabled": params.ask_user_question_enabled,
|
"ask_user_question_enabled": params.ask_user_question_enabled,
|
||||||
"orchestration_enabled": params.orchestration_enabled,
|
"orchestration_enabled": params.orchestration_enabled,
|
||||||
"is_remote_session": params.session_context.is_remote(),
|
"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!({})),
|
"identifiers": serde_json::to_value(ai_identifiers).unwrap_or_else(|_| serde_json::json!({})),
|
||||||
}),
|
}),
|
||||||
},
|
},
|
||||||
@@ -372,6 +374,14 @@ impl ResponseStream {
|
|||||||
"reason".to_string(),
|
"reason".to_string(),
|
||||||
serde_json::json!(stream_finished_reason_name(&finished_event.reason)),
|
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(
|
context.insert(
|
||||||
"elapsed_ms".to_string(),
|
"elapsed_ms".to_string(),
|
||||||
serde_json::json!(self.time_to_latest_event.num_milliseconds()),
|
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()),
|
serde_json::json!(self.time_to_latest_event.num_milliseconds()),
|
||||||
);
|
);
|
||||||
context.insert("recovery".to_string(), serde_json::json!(recovery));
|
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(
|
context.insert(
|
||||||
"error".to_string(),
|
"error".to_string(),
|
||||||
serde_json::json!(remote_logging::sanitize_error(error)),
|
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"))]
|
#[cfg(not(target_family = "wasm"))]
|
||||||
fn token_usage_context(
|
fn token_usage_context(
|
||||||
token_usage: &[response_event::stream_finished::TokenUsage],
|
token_usage: &[response_event::stream_finished::TokenUsage],
|
||||||
|
|||||||
@@ -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]
|
#[test]
|
||||||
fn raw_interactive_ssh_is_treated_as_remote_for_acp() {
|
fn raw_interactive_ssh_is_treated_as_remote_for_acp() {
|
||||||
for command in [
|
for command in [
|
||||||
|
|||||||
@@ -4,10 +4,11 @@ use std::sync::{Arc, Mutex};
|
|||||||
use ai::agent::action::{AskUserQuestionItem, AskUserQuestionType};
|
use ai::agent::action::{AskUserQuestionItem, AskUserQuestionType};
|
||||||
use chrono::Local;
|
use chrono::Local;
|
||||||
use galaxy_agent_core::{
|
use galaxy_agent_core::{
|
||||||
CompletedModelTurn, ContentPart, ConversationMessage, ExternalWorkId, MessageContent,
|
AgentError, AgentErrorKind, CompletedModelTurn, ContentPart, ConversationMessage,
|
||||||
MessageRole, PermissionKind, PermissionRequest, ProviderRun, ProviderRunId, ProviderRunLimits,
|
ExternalWorkId, MessageContent, MessageRole, PermissionKind, PermissionRequest, ProviderRun,
|
||||||
ProviderRunState, ProviderRunStep, RunEpoch, RuntimeCapabilities, StopReason, ToolCall,
|
ProviderRunFailure, ProviderRunFailureKind, ProviderRunId, ProviderRunLimits,
|
||||||
TurnRequest, Usage,
|
ProviderRunOutcome, ProviderRunState, ProviderRunStep, RunEpoch, RuntimeCapabilities,
|
||||||
|
StopReason, ToolCall, TurnRequest, Usage,
|
||||||
};
|
};
|
||||||
use galaxy_core::command::ExitCode;
|
use galaxy_core::command::ExitCode;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
@@ -26,9 +27,11 @@ use crate::ai::agent::{
|
|||||||
use crate::ai::ambient_agents::AmbientAgentTaskId;
|
use crate::ai::ambient_agents::AmbientAgentTaskId;
|
||||||
use crate::ai::blocklist::{
|
use crate::ai::blocklist::{
|
||||||
BlocklistAIHistoryEvent, BlocklistAIHistoryModel, PendingAttachment, PendingFile, RequestInput,
|
BlocklistAIHistoryEvent, BlocklistAIHistoryModel, PendingAttachment, PendingFile, RequestInput,
|
||||||
ResponseStream, ResponseStreamId,
|
ResponseStream, ResponseStreamId, StartAgentExecutor,
|
||||||
};
|
};
|
||||||
use crate::ai::llms::LLMId;
|
use crate::ai::llms::LLMId;
|
||||||
|
use crate::ai::remote_logging::RemoteLogLevel;
|
||||||
|
use crate::ai::runtime::ProviderRunProjection;
|
||||||
use crate::persistence::model::{AcpConversationData, AgentBackend};
|
use crate::persistence::model::{AcpConversationData, AgentBackend};
|
||||||
use crate::terminal::model::block::{BlockId, BlockState};
|
use crate::terminal::model::block::{BlockId, BlockState};
|
||||||
use crate::test_util::settings::initialize_history_persistence_for_tests;
|
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()
|
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 {
|
fn ask_user_question_action(action_id: &str) -> AIAgentAction {
|
||||||
AIAgentAction {
|
AIAgentAction {
|
||||||
id: AIAgentActionId::from(action_id.to_string()),
|
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
|
/// When an agent command exits the shell, the conversation must be finalized as
|
||||||
/// `Error` (not `Cancelled`), and a subsequent `ManuallyCancelled` (as fired by
|
/// `Error` (not `Cancelled`), and a subsequent `ManuallyCancelled` (as fired by
|
||||||
/// the pane-close path) must not overwrite that failure.
|
/// the pane-close path) must not overwrite that failure.
|
||||||
|
|||||||
@@ -88,7 +88,9 @@ impl ProviderRunResponseProjector {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
ProviderRunProjection::ModelEvent { event, .. } => self.translator.translate(event),
|
ProviderRunProjection::ModelEvent { event, .. } => self.translator.translate(event),
|
||||||
ProviderRunProjection::ModelRetry { .. }
|
ProviderRunProjection::ModelTurnRequested { .. }
|
||||||
|
| ProviderRunProjection::ModelTurnFinished { .. }
|
||||||
|
| ProviderRunProjection::ModelRetry { .. }
|
||||||
| ProviderRunProjection::ToolBatchReady { .. } => Ok(Vec::new()),
|
| ProviderRunProjection::ToolBatchReady { .. } => Ok(Vec::new()),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -54,8 +54,12 @@ fn restored_provider_projection_skips_stream_initialization() {
|
|||||||
assert!(projector
|
assert!(projector
|
||||||
.project(ProviderRunProjection::ModelTurnStarted {
|
.project(ProviderRunProjection::ModelTurnStarted {
|
||||||
work_id: work_id.clone(),
|
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(),
|
runtime_request_id: "request".to_owned(),
|
||||||
retry_attempt: 0,
|
retry_attempt: 0,
|
||||||
|
elapsed_ms: 1,
|
||||||
})
|
})
|
||||||
.unwrap()
|
.unwrap()
|
||||||
.is_empty());
|
.is_empty());
|
||||||
|
|||||||
@@ -8,8 +8,9 @@ pub(crate) use event_translator::{
|
|||||||
ProviderRunResponseProjector, RuntimeResponseConfig, RuntimeResponseTranslator,
|
ProviderRunResponseProjector, RuntimeResponseConfig, RuntimeResponseTranslator,
|
||||||
};
|
};
|
||||||
pub(crate) use provider_run_coordinator::{
|
pub(crate) use provider_run_coordinator::{
|
||||||
ProviderRunBlock, ProviderRunCoordinator, ProviderRunProfile, ProviderToolExecutionRef,
|
ProviderRunBlock, ProviderRunCoordinator, ProviderRunProfile, ProviderRunProjection,
|
||||||
ProviderToolLifecycleOutcome, BASE_PROVIDER_PROFILE, CLI_MONITOR_PROVIDER_PROFILE,
|
ProviderToolExecutionRef, ProviderToolLifecycleOutcome, BASE_PROVIDER_PROFILE,
|
||||||
|
CLI_MONITOR_PROVIDER_PROFILE,
|
||||||
};
|
};
|
||||||
pub(crate) use rig::{
|
pub(crate) use rig::{
|
||||||
prepare_provider_run, provider_runtime_for_request, PreparedProviderRun, ProviderActionContext,
|
prepare_provider_run, provider_runtime_for_request, PreparedProviderRun, ProviderActionContext,
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ use std::collections::{BTreeMap, BTreeSet};
|
|||||||
use std::error::Error;
|
use std::error::Error;
|
||||||
use std::fmt;
|
use std::fmt;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
use futures::future::BoxFuture;
|
use futures::future::BoxFuture;
|
||||||
use futures::StreamExt;
|
use futures::StreamExt;
|
||||||
@@ -12,18 +13,43 @@ use galaxy_agent_core::{
|
|||||||
ProviderRunOutcome, ProviderRunPhase, ProviderRunProtocolError, ProviderRunState,
|
ProviderRunOutcome, ProviderRunPhase, ProviderRunProtocolError, ProviderRunState,
|
||||||
ProviderRunStep, RunEpoch, RuntimeKind, StopReason, ToolEvent, TurnControl, TurnRequest, Usage,
|
ProviderRunStep, RunEpoch, RuntimeKind, StopReason, ToolEvent, TurnControl, TurnRequest, Usage,
|
||||||
};
|
};
|
||||||
|
use instant::Instant;
|
||||||
|
use warpui::r#async::FutureExt as _;
|
||||||
|
|
||||||
use crate::ai::agent::conversation::AIConversationId;
|
use crate::ai::agent::conversation::AIConversationId;
|
||||||
|
|
||||||
pub(crate) const BASE_PROVIDER_PROFILE: &str = "base";
|
pub(crate) const BASE_PROVIDER_PROFILE: &str = "base";
|
||||||
pub(crate) const CLI_MONITOR_PROVIDER_PROFILE: &str = "cli-monitor";
|
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)]
|
#[derive(Clone, Debug, PartialEq)]
|
||||||
pub(crate) enum ProviderRunProjection {
|
pub(crate) enum ProviderRunProjection {
|
||||||
|
ModelTurnRequested {
|
||||||
|
work_id: ExternalWorkId,
|
||||||
|
profile: ProviderRequestProfile,
|
||||||
|
runtime_id: String,
|
||||||
|
model_id: String,
|
||||||
|
retry_attempt: u32,
|
||||||
|
},
|
||||||
ModelTurnStarted {
|
ModelTurnStarted {
|
||||||
work_id: ExternalWorkId,
|
work_id: ExternalWorkId,
|
||||||
|
profile: ProviderRequestProfile,
|
||||||
|
runtime_id: String,
|
||||||
|
model_id: String,
|
||||||
runtime_request_id: String,
|
runtime_request_id: String,
|
||||||
retry_attempt: u32,
|
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 {
|
ModelEvent {
|
||||||
work_id: ExternalWorkId,
|
work_id: ExternalWorkId,
|
||||||
@@ -31,7 +57,11 @@ pub(crate) enum ProviderRunProjection {
|
|||||||
},
|
},
|
||||||
ModelRetry {
|
ModelRetry {
|
||||||
work_id: ExternalWorkId,
|
work_id: ExternalWorkId,
|
||||||
|
profile: ProviderRequestProfile,
|
||||||
|
runtime_id: String,
|
||||||
|
model_id: String,
|
||||||
retry_attempt: u32,
|
retry_attempt: u32,
|
||||||
|
elapsed_ms: u64,
|
||||||
error: AgentError,
|
error: AgentError,
|
||||||
},
|
},
|
||||||
ToolBatchReady {
|
ToolBatchReady {
|
||||||
@@ -133,6 +163,8 @@ impl ProviderRunProfile {
|
|||||||
pub(crate) struct ProviderRunCoordinator {
|
pub(crate) struct ProviderRunCoordinator {
|
||||||
run: ProviderRun,
|
run: ProviderRun,
|
||||||
profiles: BTreeMap<String, ProviderRunProfile>,
|
profiles: BTreeMap<String, ProviderRunProfile>,
|
||||||
|
model_start_timeout: Duration,
|
||||||
|
model_event_idle_timeout: Duration,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ProviderRunCoordinator {
|
impl ProviderRunCoordinator {
|
||||||
@@ -171,7 +203,18 @@ impl ProviderRunCoordinator {
|
|||||||
for (profile, config) in &profiles {
|
for (profile, config) in &profiles {
|
||||||
validate_profile_runtime(profile, config.runtime.as_ref())?;
|
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 {
|
pub(crate) fn run(&self) -> &ProviderRun {
|
||||||
@@ -395,22 +438,67 @@ impl ProviderRunCoordinator {
|
|||||||
.iter()
|
.iter()
|
||||||
.map(|tool| tool.name.clone())
|
.map(|tool| tool.name.clone())
|
||||||
.collect::<BTreeSet<_>>();
|
.collect::<BTreeSet<_>>();
|
||||||
let request = request_for_model_call(profile.request, &call);
|
let runtime_id = profile.runtime.descriptor().id.clone();
|
||||||
let stream = match profile.runtime.start_turn(request, control).await {
|
let model_id = profile.request.model.as_str().to_string();
|
||||||
Ok(stream) => stream,
|
if !self.project_or_fail(
|
||||||
Err(error) => {
|
ProviderRunProjection::ModelTurnRequested {
|
||||||
self.handle_model_failure(&call.work_id, error, project)?;
|
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(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
futures::pin_mut!(stream);
|
futures::pin_mut!(stream);
|
||||||
let mut buffer = ModelTurnBuffer::default();
|
let mut buffer = ModelTurnBuffer::default();
|
||||||
|
|
||||||
while let Some(event) = stream.next().await {
|
loop {
|
||||||
let event = match event {
|
let event = match stream
|
||||||
Ok(event) => event,
|
.next()
|
||||||
Err(error) => {
|
.with_timeout(self.model_event_idle_timeout)
|
||||||
self.handle_model_failure(&call.work_id, error, project)?;
|
.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(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -418,7 +506,9 @@ impl ProviderRunCoordinator {
|
|||||||
AgentEvent::TurnStarted { runtime_request_id } => {
|
AgentEvent::TurnStarted { runtime_request_id } => {
|
||||||
if buffer.started {
|
if buffer.started {
|
||||||
self.handle_model_failure(
|
self.handle_model_failure(
|
||||||
&call.work_id,
|
&call,
|
||||||
|
&profile,
|
||||||
|
started_at,
|
||||||
protocol_error("provider emitted more than one TurnStarted event"),
|
protocol_error("provider emitted more than one TurnStarted event"),
|
||||||
project,
|
project,
|
||||||
)?;
|
)?;
|
||||||
@@ -426,7 +516,9 @@ impl ProviderRunCoordinator {
|
|||||||
}
|
}
|
||||||
if runtime_request_id.is_empty() {
|
if runtime_request_id.is_empty() {
|
||||||
self.handle_model_failure(
|
self.handle_model_failure(
|
||||||
&call.work_id,
|
&call,
|
||||||
|
&profile,
|
||||||
|
started_at,
|
||||||
protocol_error("provider emitted an empty runtime request ID"),
|
protocol_error("provider emitted an empty runtime request ID"),
|
||||||
project,
|
project,
|
||||||
)?;
|
)?;
|
||||||
@@ -436,8 +528,12 @@ impl ProviderRunCoordinator {
|
|||||||
if !self.project_or_fail(
|
if !self.project_or_fail(
|
||||||
ProviderRunProjection::ModelTurnStarted {
|
ProviderRunProjection::ModelTurnStarted {
|
||||||
work_id: call.work_id.clone(),
|
work_id: call.work_id.clone(),
|
||||||
|
profile: call.profile.clone(),
|
||||||
|
runtime_id: runtime_id.clone(),
|
||||||
|
model_id: model_id.clone(),
|
||||||
runtime_request_id,
|
runtime_request_id,
|
||||||
retry_attempt: call.retry_attempt,
|
retry_attempt: call.retry_attempt,
|
||||||
|
elapsed_ms: elapsed_millis(started_at),
|
||||||
},
|
},
|
||||||
project,
|
project,
|
||||||
)? {
|
)? {
|
||||||
@@ -445,7 +541,7 @@ impl ProviderRunCoordinator {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
AgentEvent::TextDelta { text } => {
|
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(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
buffer.text.push_str(&text);
|
buffer.text.push_str(&text);
|
||||||
@@ -460,7 +556,7 @@ impl ProviderRunCoordinator {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
AgentEvent::ReasoningDelta { text } => {
|
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(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
buffer.reasoning.push_str(&text);
|
buffer.reasoning.push_str(&text);
|
||||||
@@ -475,7 +571,7 @@ impl ProviderRunCoordinator {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
AgentEvent::ReasoningCompleted { text, signature } => {
|
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(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
if !text.is_empty() {
|
if !text.is_empty() {
|
||||||
@@ -495,13 +591,13 @@ impl ProviderRunCoordinator {
|
|||||||
AgentEvent::Tool {
|
AgentEvent::Tool {
|
||||||
event: ToolEvent::Proposed { call: tool_call },
|
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(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
buffer.tool_calls.push(tool_call);
|
buffer.tool_calls.push(tool_call);
|
||||||
}
|
}
|
||||||
AgentEvent::UsageUpdated { usage } => {
|
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(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
buffer.usage.clone_from(&usage);
|
buffer.usage.clone_from(&usage);
|
||||||
@@ -519,7 +615,22 @@ impl ProviderRunCoordinator {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
AgentEvent::TurnStopped { reason } => {
|
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(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
if reason == StopReason::Cancelled {
|
if reason == StopReason::Cancelled {
|
||||||
@@ -547,7 +658,9 @@ impl ProviderRunCoordinator {
|
|||||||
| AgentEvent::UserInputAccepted { .. }
|
| AgentEvent::UserInputAccepted { .. }
|
||||||
| AgentEvent::RuntimeNotice { .. } => {
|
| AgentEvent::RuntimeNotice { .. } => {
|
||||||
self.handle_model_failure(
|
self.handle_model_failure(
|
||||||
&call.work_id,
|
&call,
|
||||||
|
&profile,
|
||||||
|
started_at,
|
||||||
protocol_error(
|
protocol_error(
|
||||||
"direct-provider transport emitted a non-model lifecycle event",
|
"direct-provider transport emitted a non-model lifecycle event",
|
||||||
),
|
),
|
||||||
@@ -559,7 +672,9 @@ impl ProviderRunCoordinator {
|
|||||||
}
|
}
|
||||||
|
|
||||||
self.handle_model_failure(
|
self.handle_model_failure(
|
||||||
&call.work_id,
|
&call,
|
||||||
|
&profile,
|
||||||
|
started_at,
|
||||||
protocol_error("provider stream ended before TurnStopped"),
|
protocol_error("provider stream ended before TurnStopped"),
|
||||||
project,
|
project,
|
||||||
)?;
|
)?;
|
||||||
@@ -568,7 +683,9 @@ impl ProviderRunCoordinator {
|
|||||||
|
|
||||||
fn ensure_model_started<F>(
|
fn ensure_model_started<F>(
|
||||||
&mut self,
|
&mut self,
|
||||||
work_id: &ExternalWorkId,
|
call: &ProviderModelCall,
|
||||||
|
profile: &ProviderRunProfile,
|
||||||
|
started_at: Instant,
|
||||||
buffer: &ModelTurnBuffer,
|
buffer: &ModelTurnBuffer,
|
||||||
project: &mut F,
|
project: &mut F,
|
||||||
) -> Result<bool, ProviderRunCoordinatorError>
|
) -> Result<bool, ProviderRunCoordinatorError>
|
||||||
@@ -579,7 +696,9 @@ impl ProviderRunCoordinator {
|
|||||||
return Ok(true);
|
return Ok(true);
|
||||||
}
|
}
|
||||||
self.handle_model_failure(
|
self.handle_model_failure(
|
||||||
work_id,
|
call,
|
||||||
|
profile,
|
||||||
|
started_at,
|
||||||
protocol_error("provider emitted model output before TurnStarted"),
|
protocol_error("provider emitted model output before TurnStarted"),
|
||||||
project,
|
project,
|
||||||
)?;
|
)?;
|
||||||
@@ -588,14 +707,18 @@ impl ProviderRunCoordinator {
|
|||||||
|
|
||||||
fn handle_model_failure<F>(
|
fn handle_model_failure<F>(
|
||||||
&mut self,
|
&mut self,
|
||||||
work_id: &ExternalWorkId,
|
call: &ProviderModelCall,
|
||||||
|
profile: &ProviderRunProfile,
|
||||||
|
started_at: Instant,
|
||||||
error: AgentError,
|
error: AgentError,
|
||||||
project: &mut F,
|
project: &mut F,
|
||||||
) -> Result<(), ProviderRunCoordinatorError>
|
) -> Result<(), ProviderRunCoordinatorError>
|
||||||
where
|
where
|
||||||
F: FnMut(ProviderRunProjection) -> Result<(), String>,
|
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 {
|
if disposition == ModelFailureDisposition::RetryScheduled {
|
||||||
let retry_attempt = match self.run.state() {
|
let retry_attempt = match self.run.state() {
|
||||||
ProviderRunState::AwaitingModel { call } => call.retry_attempt,
|
ProviderRunState::AwaitingModel { call } => call.retry_attempt,
|
||||||
@@ -616,8 +739,12 @@ impl ProviderRunCoordinator {
|
|||||||
};
|
};
|
||||||
self.project_or_fail(
|
self.project_or_fail(
|
||||||
ProviderRunProjection::ModelRetry {
|
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,
|
retry_attempt,
|
||||||
|
elapsed_ms: elapsed_millis(started_at),
|
||||||
error,
|
error,
|
||||||
},
|
},
|
||||||
project,
|
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(
|
fn validate_profile_runtime(
|
||||||
profile: &str,
|
profile: &str,
|
||||||
runtime: &dyn AgentRuntime,
|
runtime: &dyn AgentRuntime,
|
||||||
|
|||||||
@@ -1,5 +1,7 @@
|
|||||||
use std::collections::VecDeque;
|
use std::collections::VecDeque;
|
||||||
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||||
use std::sync::Mutex;
|
use std::sync::Mutex;
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use galaxy_agent_core::{
|
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 {
|
fn request() -> TurnRequest {
|
||||||
let mut request = TurnRequest::new(
|
let mut request = TurnRequest::new(
|
||||||
"test-model",
|
"test-model",
|
||||||
@@ -690,7 +751,9 @@ async fn recoverable_start_failure_retries_the_same_work_identity() {
|
|||||||
retry_attempt,
|
retry_attempt,
|
||||||
..
|
..
|
||||||
} => Some((work_id.clone(), *retry_attempt)),
|
} => Some((work_id.clone(), *retry_attempt)),
|
||||||
ProviderRunProjection::ModelTurnStarted { .. }
|
ProviderRunProjection::ModelTurnRequested { .. }
|
||||||
|
| ProviderRunProjection::ModelTurnStarted { .. }
|
||||||
|
| ProviderRunProjection::ModelTurnFinished { .. }
|
||||||
| ProviderRunProjection::ModelEvent { .. }
|
| ProviderRunProjection::ModelEvent { .. }
|
||||||
| ProviderRunProjection::ToolBatchReady { .. } => None,
|
| ProviderRunProjection::ToolBatchReady { .. } => None,
|
||||||
})
|
})
|
||||||
@@ -703,7 +766,9 @@ async fn recoverable_start_failure_retries_the_same_work_identity() {
|
|||||||
retry_attempt: 1,
|
retry_attempt: 1,
|
||||||
..
|
..
|
||||||
} => Some(work_id.clone()),
|
} => Some(work_id.clone()),
|
||||||
ProviderRunProjection::ModelTurnStarted { .. }
|
ProviderRunProjection::ModelTurnRequested { .. }
|
||||||
|
| ProviderRunProjection::ModelTurnStarted { .. }
|
||||||
|
| ProviderRunProjection::ModelTurnFinished { .. }
|
||||||
| ProviderRunProjection::ModelEvent { .. }
|
| ProviderRunProjection::ModelEvent { .. }
|
||||||
| ProviderRunProjection::ModelRetry { .. }
|
| ProviderRunProjection::ModelRetry { .. }
|
||||||
| ProviderRunProjection::ToolBatchReady { .. } => None,
|
| ProviderRunProjection::ToolBatchReady { .. } => None,
|
||||||
@@ -713,6 +778,119 @@ async fn recoverable_start_failure_retries_the_same_work_identity() {
|
|||||||
assert_eq!(retry.1, 1);
|
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]
|
#[tokio::test]
|
||||||
async fn tool_projection_failure_synthesizes_a_result_and_fails_the_run() {
|
async fn tool_projection_failure_synthesizes_a_result_and_fails_the_run() {
|
||||||
let runtime = Arc::new(ScriptedRuntime::new(vec![tool_turn()]));
|
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 { .. } => {
|
ProviderRunProjection::ToolBatchReady { .. } => {
|
||||||
Err("task projection disappeared".to_string())
|
Err("task projection disappeared".to_string())
|
||||||
}
|
}
|
||||||
ProviderRunProjection::ModelTurnStarted { .. }
|
ProviderRunProjection::ModelTurnRequested { .. }
|
||||||
|
| ProviderRunProjection::ModelTurnStarted { .. }
|
||||||
|
| ProviderRunProjection::ModelTurnFinished { .. }
|
||||||
| ProviderRunProjection::ModelEvent { .. }
|
| ProviderRunProjection::ModelEvent { .. }
|
||||||
| ProviderRunProjection::ModelRetry { .. } => Ok(()),
|
| ProviderRunProjection::ModelRetry { .. } => Ok(()),
|
||||||
})
|
})
|
||||||
|
|||||||
Reference in New Issue
Block a user