315 lines
14 KiB
Rust
315 lines
14 KiB
Rust
//! Stream-based API for spawning and monitoring ambient agents.
|
|
#![cfg_attr(target_family = "wasm", expect(dead_code))]
|
|
|
|
use std::str::FromStr;
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
use anyhow::anyhow;
|
|
use futures::{select, FutureExt, Stream, StreamExt};
|
|
use session_sharing_protocol::common::SessionId;
|
|
|
|
use super::{AmbientAgentTask, AmbientAgentTaskId, AmbientAgentTaskState};
|
|
use crate::server::retry_strategies::with_bounded_retry;
|
|
use crate::server::server_api::ai::{
|
|
AIClient, RunFollowupRequest, SpawnAgentRequest, TaskStatusMessage,
|
|
};
|
|
use crate::terminal::shared_session;
|
|
|
|
/// How long to poll for the agent to be ready.
|
|
/// This should be long enough that the shared session will be joinable.
|
|
pub const TASK_STATUS_POLLING_DURATION: Duration = Duration::from_secs(80);
|
|
|
|
#[cfg(not(test))]
|
|
const TASK_STATUS_POLL_INTERVAL: Duration = Duration::from_secs(3);
|
|
#[cfg(test)]
|
|
const TASK_STATUS_POLL_INTERVAL: Duration = Duration::from_millis(1);
|
|
|
|
/// Maximum number of consecutive polls that may observe a non-working, non-cancelled task
|
|
/// state for a follow-up before we give up and surface a failure. Bounds the worst-case
|
|
/// wait when the server is wedged and never transitions the task off its prior terminal
|
|
/// state. At the production `TASK_STATUS_POLL_INTERVAL` of 3s, 10 skipped observations
|
|
/// is ~30s — comfortably longer than the dispatcher's `ProcessingInterval` plus typical
|
|
/// worker claim latency.
|
|
const MAX_STALE_POLLS_BEFORE_FAILURE: usize = 10;
|
|
|
|
/// Information about a session join link for an ambient agent task.
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub struct SessionJoinInfo {
|
|
pub session_id: Option<SessionId>,
|
|
pub session_link: String,
|
|
}
|
|
|
|
impl SessionJoinInfo {
|
|
pub fn from_task(task: &AmbientAgentTask) -> Option<Self> {
|
|
let run_execution = task.active_run_execution();
|
|
// The cloud-mode pane joins on `session_id`; a standalone `session_link` isn't
|
|
// actionable without it.
|
|
let session_id_str = run_execution.session_id?;
|
|
let session_id = SessionId::from_str(session_id_str).ok()?;
|
|
|
|
// Prefer the server-provided `session_link`; fall back to constructing one from
|
|
// `session_id`. `active_run_execution()` already filters out empty links.
|
|
let session_link = run_execution
|
|
.session_link
|
|
.map(String::from)
|
|
.unwrap_or_else(|| shared_session::join_link(&session_id));
|
|
Some(Self {
|
|
session_id: Some(session_id),
|
|
session_link,
|
|
})
|
|
}
|
|
}
|
|
|
|
/// Lifecycle events during ambient agent startup.
|
|
#[derive(Debug)]
|
|
pub enum AmbientAgentEvent {
|
|
/// The task was successfully spawned with the given task ID and run ID.
|
|
TaskSpawned {
|
|
task_id: AmbientAgentTaskId,
|
|
run_id: String,
|
|
},
|
|
/// The task state changed.
|
|
StateChanged {
|
|
state: AmbientAgentTaskState,
|
|
status_message: Option<TaskStatusMessage>,
|
|
},
|
|
/// Session started and join information became available.
|
|
SessionStarted { session_join_info: SessionJoinInfo },
|
|
/// Timed out waiting for the agent session to be ready.
|
|
TimedOut,
|
|
/// Cloud agent capacity limit has been reached. This does not block
|
|
/// the task from eventually starting.
|
|
AtCapacity,
|
|
}
|
|
|
|
enum RunPollMode {
|
|
InitialRun,
|
|
Followup {
|
|
previous_session_id: Option<SessionId>,
|
|
},
|
|
}
|
|
|
|
/// Spawns an ambient agent task and monitors its state.
|
|
///
|
|
/// The stream completes when:
|
|
/// - The task completes (either successfully or with a failure)
|
|
/// - The task's shared session is ready to join
|
|
/// - The timeout expires (if provided)
|
|
/// - An error occurs
|
|
///
|
|
/// If `timeout` is `None`, there is no timeout.
|
|
pub fn spawn_task(
|
|
request: SpawnAgentRequest,
|
|
ai_client: Arc<dyn AIClient>,
|
|
timeout: Option<Duration>,
|
|
) -> impl Stream<Item = Result<AmbientAgentEvent, anyhow::Error>> {
|
|
// We can't use try_stream! because of the select! macro invocation.
|
|
// See https://github.com/tokio-rs/async-stream/issues/63.
|
|
async_stream::stream! {
|
|
// First, spawn the ambient agent task.
|
|
let (task_id, run_id, at_capacity) = match ai_client.spawn_agent(request).await {
|
|
Ok(response) => (response.task_id, response.run_id, response.at_capacity),
|
|
Err(err) => {
|
|
yield Err(err);
|
|
return;
|
|
},
|
|
};
|
|
|
|
yield Ok(AmbientAgentEvent::TaskSpawned { task_id, run_id });
|
|
|
|
// Emit AtCapacity event if the server indicates capacity limit reached.
|
|
if at_capacity {
|
|
yield Ok(AmbientAgentEvent::AtCapacity);
|
|
}
|
|
|
|
let mut stream = Box::pin(poll_run_until_joinable_session(
|
|
task_id,
|
|
ai_client,
|
|
RunPollMode::InitialRun,
|
|
timeout,
|
|
));
|
|
while let Some(event) = stream.next().await {
|
|
yield event;
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn submit_run_followup(
|
|
message: String,
|
|
run_id: AmbientAgentTaskId,
|
|
previous_session_id: Option<SessionId>,
|
|
ai_client: Arc<dyn AIClient>,
|
|
timeout: Option<Duration>,
|
|
) -> impl Stream<Item = Result<AmbientAgentEvent, anyhow::Error>> {
|
|
async_stream::stream! {
|
|
let request = RunFollowupRequest { message };
|
|
if let Err(err) = ai_client.submit_run_followup(&run_id, request).await {
|
|
yield Err(err);
|
|
return;
|
|
}
|
|
|
|
let mut stream = Box::pin(poll_run_until_joinable_session(
|
|
run_id,
|
|
ai_client,
|
|
RunPollMode::Followup {
|
|
previous_session_id,
|
|
},
|
|
timeout,
|
|
));
|
|
while let Some(event) = stream.next().await {
|
|
yield event;
|
|
}
|
|
}
|
|
}
|
|
|
|
fn poll_run_until_joinable_session(
|
|
run_id: AmbientAgentTaskId,
|
|
ai_client: Arc<dyn AIClient>,
|
|
mode: RunPollMode,
|
|
timeout: Option<Duration>,
|
|
) -> impl Stream<Item = Result<AmbientAgentEvent, anyhow::Error>> {
|
|
async_stream::stream! {
|
|
// Poll for the task until it completes OR has session join info.
|
|
// We use a timeout to ensure we don't wait indefinitely for session info.
|
|
// If no timeout is provided, we use a future that never completes.
|
|
let mut timeout_timer = FutureExt::fuse(match timeout {
|
|
Some(d) => galaxyui::r#async::Timer::after(d),
|
|
None => galaxyui::r#async::Timer::never(),
|
|
});
|
|
let mut last_state = None;
|
|
// For follow-ups, the server does NOT synchronously transition the task off its
|
|
// prior terminal state when `submit_run_followup` returns — the transition happens
|
|
// later, asynchronously, when the dispatcher loop picks up the newly enqueued
|
|
// execution and the worker claims it. If we treated the first observed state as
|
|
// the follow-up's outcome we'd misreport the prior run's `status_message` as a
|
|
// failure and end the stream, leaving the model permanently stuck in `Failed`
|
|
// even though the new run is actually about to start.
|
|
//
|
|
// To avoid that, gate event emission for follow-ups on having observed at least
|
|
// one working state. Initial spawns don't need this — they start from a fresh
|
|
// task whose first observation reflects the spawn itself.
|
|
let mut seen_working_state = matches!(&mode, RunPollMode::InitialRun);
|
|
let mut skipped_stale_polls: usize = 0;
|
|
loop {
|
|
let mut poll_timer = FutureExt::fuse(galaxyui::r#async::Timer::after(TASK_STATUS_POLL_INTERVAL));
|
|
|
|
select! {
|
|
_ = timeout_timer => {
|
|
yield Ok(AmbientAgentEvent::TimedOut);
|
|
return;
|
|
}
|
|
_ = poll_timer => {
|
|
// Wrap the status poll in with_bounded_retry so transient
|
|
// HTTP errors (429, 5xx) are retried with exponential
|
|
// backoff instead of immediately killing the CLI.
|
|
let poll_result = {
|
|
let client = ai_client.clone();
|
|
with_bounded_retry(
|
|
&format!("poll agent {run_id}"),
|
|
|| {
|
|
let client = client.clone();
|
|
async move { client.get_ambient_agent_task(&run_id).await }
|
|
},
|
|
)
|
|
.await
|
|
};
|
|
match poll_result {
|
|
Ok(task) => {
|
|
// Log every non-InProgress observation BEFORE the skip check
|
|
// so we retain visibility into stalls where the server is
|
|
// wedged on the prior terminal state.
|
|
if task.state != AmbientAgentTaskState::InProgress {
|
|
log::info!("Agent {run_id} state: {:?}", task.state);
|
|
}
|
|
|
|
if task.state.is_working() {
|
|
seen_working_state = true;
|
|
} else if !seen_working_state
|
|
&& task.state != AmbientAgentTaskState::Cancelled
|
|
&& skipped_stale_polls < MAX_STALE_POLLS_BEFORE_FAILURE
|
|
{
|
|
// Likely the prior run's residual terminal state; the
|
|
// server hasn't transitioned the task yet. Skip without
|
|
// emitting events or ending the stream.
|
|
//
|
|
// Carve-outs:
|
|
// - `Cancelled` always falls through: it's only reached
|
|
// via explicit user/admin/server cancellation and the
|
|
// server will never transition out of it on its own.
|
|
// - After `MAX_STALE_POLLS_BEFORE_FAILURE` skipped polls
|
|
// we give up so a wedged server can't keep the stream
|
|
// alive indefinitely.
|
|
skipped_stale_polls += 1;
|
|
continue;
|
|
}
|
|
|
|
if last_state.as_ref() != Some(&task.state) {
|
|
last_state = Some(task.state.clone());
|
|
yield Ok(AmbientAgentEvent::StateChanged {
|
|
state: task.state.clone(),
|
|
status_message: task.status_message.clone(),
|
|
});
|
|
}
|
|
|
|
if task.state.is_terminal() {
|
|
if matches!(&mode, RunPollMode::Followup { .. }) {
|
|
let exhausted_stale_skips = !seen_working_state
|
|
&& skipped_stale_polls >= MAX_STALE_POLLS_BEFORE_FAILURE;
|
|
let message = if exhausted_stale_skips {
|
|
"Cloud follow-up did not start in time".to_string()
|
|
} else {
|
|
task.status_message
|
|
.as_ref()
|
|
.map(|msg| msg.message.clone())
|
|
.unwrap_or_else(|| {
|
|
if task.state.is_failure_like() {
|
|
"Cloud agent failed".to_string()
|
|
} else {
|
|
"Cloud follow-up finished before a new session became available".to_string()
|
|
}
|
|
})
|
|
};
|
|
yield Err(anyhow!(message));
|
|
}
|
|
return;
|
|
}
|
|
|
|
if task.state == AmbientAgentTaskState::InProgress {
|
|
if let Some(session_join_info) = SessionJoinInfo::from_task(&task) {
|
|
let has_new_session = match &mode {
|
|
RunPollMode::InitialRun
|
|
| RunPollMode::Followup {
|
|
previous_session_id: None,
|
|
} => true,
|
|
RunPollMode::Followup {
|
|
previous_session_id: Some(previous_session_id),
|
|
} => session_join_info
|
|
.session_id
|
|
.as_ref()
|
|
.is_some_and(|session_id| session_id != previous_session_id),
|
|
};
|
|
if has_new_session {
|
|
yield Ok(AmbientAgentEvent::SessionStarted {
|
|
session_join_info,
|
|
});
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Err(err) => {
|
|
yield Err(err);
|
|
return;
|
|
},
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[path = "spawn_tests.rs"]
|
|
mod tests;
|