feat: observe retained worker sessions without restore

This commit is contained in:
2026-09-18 04:27:38 +09:00
parent ae7f0f0446
commit 7a61a26670
26 changed files with 834 additions and 124 deletions
+1
View File
@@ -27,6 +27,7 @@ session-store = { workspace = true }
ticket = { workspace = true }
serde = { workspace = true, features = ["derive"] }
pulldown-cmark = { version = "0.13.3", default-features = false }
async-trait.workspace = true
agen.workspace = true
[dev-dependencies]
+25 -111
View File
@@ -3,9 +3,8 @@ use std::io;
use std::time::Duration;
use client::{
BackendRuntimeListTarget, BackendWorkerRestoreResponse, BackendWorkerRestoreState,
BackendWorkerSummary, list_backend_stopped_workers, list_backend_workers,
restore_backend_worker,
BackendRuntimeListTarget, BackendWorkerSummary, ServerApiClient, WorkerSessionAvailability,
list_backend_stopped_workers, list_backend_workers,
};
use crossterm::event::{self, Event as TermEvent, KeyCode, KeyEventKind, KeyModifiers};
use ratatui::Frame;
@@ -81,62 +80,39 @@ pub(crate) async fn run(
}
WorkerPickerResult::Selected(selected) => selected,
};
let worker = if selected.state == "stopped" {
let restore_target = target
.runtime_target(selected.runtime_id.clone(), selected.worker_id.clone())
let mut attach_target = target
.runtime_target(selected.runtime_id.clone(), selected.worker_id.clone())
.map_err(|error| io::Error::other(error.to_string()))?;
if selected.state == "stopped" {
let api = ServerApiClient::builder(&target.base_url)
.map_err(|error| io::Error::other(error.to_string()))?
.build()
.map_err(|error| io::Error::other(error.to_string()))?;
let restore = restore_backend_worker(&restore_target)
let response = api
.worker_session(
attach_target.workspace_id.clone(),
selected.runtime_id.clone(),
selected.worker_id.clone(),
)
.await
.map_err(|error| {
io::Error::other(format!(
"failed to restore Backend worker {}/{}: {error}",
"failed to observe retained Worker Session {}/{}: {error}",
selected.runtime_id, selected.worker_id
))
})?;
restored_worker(restore).map_err(|error| {
io::Error::other(format!(
"failed to restore Backend worker {}/{}: {error}",
selected.runtime_id, selected.worker_id
))
})?
} else {
selected
};
let attach_target = target
.runtime_target(worker.runtime_id, worker.worker_id)
.map_err(|error| io::Error::other(error.to_string()))?;
attach_target.initial_snapshot = match response.observation {
WorkerSessionAvailability::RetainedSnapshot { snapshot, .. } => Some(snapshot),
WorkerSessionAvailability::LiveProtocol => None,
WorkerSessionAvailability::Unavailable { message, .. } => {
return Err(io::Error::other(message).into());
}
};
}
return console::run_backend_runtime(attach_target).await;
}
}
fn restored_worker(response: BackendWorkerRestoreResponse) -> Result<BackendWorkerSummary, String> {
if response.result.state != BackendWorkerRestoreState::Accepted {
let diagnostics = response
.result
.diagnostics
.iter()
.map(|diagnostic| format!("{}: {}", diagnostic.code, diagnostic.message))
.collect::<Vec<_>>()
.join("; ");
let state = match response.result.state {
BackendWorkerRestoreState::Accepted => unreachable!(),
BackendWorkerRestoreState::Rejected => "rejected",
BackendWorkerRestoreState::RolledBack => "rolled back",
BackendWorkerRestoreState::ReconciliationRequired => "reconciliation required",
};
return Err(if diagnostics.is_empty() {
format!("restore was {state} without a diagnostic")
} else {
format!("restore was {state}: {diagnostics}")
});
}
response
.result
.worker
.ok_or_else(|| "restore was accepted without a Worker snapshot".to_string())
}
fn dedup_workers(workers: &mut Vec<BackendWorkerSummary>) {
let mut seen = std::collections::HashSet::new();
workers.retain(|worker| seen.insert((worker.runtime_id.clone(), worker.worker_id.clone())));
@@ -437,8 +413,7 @@ fn working_directory_text(worker: &BackendWorkerSummary) -> String {
mod tests {
use super::*;
use client::{
BackendDiagnostic, BackendDiagnosticSeverity, BackendWorkerCapabilitySummary,
BackendWorkerImplementationSummary, BackendWorkerRestoreResult,
BackendWorkerCapabilitySummary, BackendWorkerImplementationSummary,
BackendWorkerWorkspaceSummary,
};
@@ -494,67 +469,6 @@ mod tests {
text_width(&text[..byte_offset])
}
fn restore_response(
state: BackendWorkerRestoreState,
worker: Option<BackendWorkerSummary>,
diagnostics: Vec<BackendDiagnostic>,
) -> BackendWorkerRestoreResponse {
BackendWorkerRestoreResponse {
workspace_id: "workspace-a".to_string(),
runtime_id: "runtime-a".to_string(),
worker_id: "worker-a".to_string(),
result: BackendWorkerRestoreResult {
state,
worker,
diagnostics,
},
}
}
#[test]
fn rejected_restore_surfaces_diagnostic_instead_of_attaching_selected_worker() {
let error = restored_worker(restore_response(
BackendWorkerRestoreState::Rejected,
None,
vec![BackendDiagnostic {
code: "working_directory_not_found".to_string(),
severity: BackendDiagnosticSeverity::Error,
message: "working directory was not found".to_string(),
}],
))
.expect_err("rejected restore must not produce a Worker to attach");
assert_eq!(
error,
"restore was rejected: working_directory_not_found: working directory was not found"
);
}
#[test]
fn accepted_restore_requires_returned_worker_snapshot() {
let error = restored_worker(restore_response(
BackendWorkerRestoreState::Accepted,
None,
Vec::new(),
))
.expect_err("accepted restore without a Worker must not attach the stale selection");
assert_eq!(error, "restore was accepted without a Worker snapshot");
}
#[test]
fn accepted_restore_returns_authoritative_worker_snapshot() {
let worker = worker("runtime-a", "worker-a", Some("builtin:companion"));
let restored = restored_worker(restore_response(
BackendWorkerRestoreState::Accepted,
Some(worker.clone()),
Vec::new(),
))
.expect("accepted restore should return its Worker snapshot");
assert_eq!(restored, worker);
}
#[test]
fn worker_row_orders_and_simplifies_columns() {
let mut worker = worker("runtime-a", "worker-b", Some("builtin:coder"));
+65 -4
View File
@@ -457,20 +457,81 @@ async fn run_standalone_host(
pub(crate) async fn run_backend_runtime(
target: BackendRuntimeTarget,
) -> Result<(), Box<dyn std::error::Error>> {
if let Some(snapshot) = target.initial_snapshot.clone() {
let client = Client::new(RetainedSnapshotSocket::new(snapshot)?);
return run_backend_console(target, client, false).await;
}
let client = connect_backend_runtime(target.clone()).await?;
run_backend_console(target, client, true).await
}
async fn run_backend_console<T: Socket>(
target: BackendRuntimeTarget,
client: Client<T>,
connected: bool,
) -> Result<(), Box<dyn std::error::Error>> {
let worker_label = target.display_label();
let attachment_target = target.clone();
let client = connect_backend_runtime(target).await?;
let mut terminal = enter_fullscreen()?;
let workspace_root = std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("."));
let mut app = App::new_with_persistent_input_history(worker_label, &workspace_root);
app.connected = true;
let mut connection = ConsoleConnection::with_backend_target(client, attachment_target);
app.connected = connected;
let mut connection = ConsoleConnection::with_backend_target(client, target);
let result = run_loop(&mut terminal, &mut app, &mut connection).await;
let _ = leave_fullscreen(&mut terminal);
result
}
struct RetainedSnapshotSocket {
event: Option<String>,
}
impl RetainedSnapshotSocket {
fn new(snapshot: protocol::SessionSnapshot) -> Result<Self, serde_json::Error> {
Ok(Self {
event: Some(serde_json::to_string(&Event::Snapshot {
session: snapshot,
greeting: protocol::Greeting {
worker_name: "retained worker".to_string(),
cwd: String::new(),
provider: "retained session".to_string(),
model: String::new(),
scope_summary: "read-only retained session".to_string(),
tools: Vec::new(),
context_window: 0,
context_tokens: 0,
},
state: protocol::WorkerStateSnapshot {
last_command_id: 0,
state: protocol::WorkerState::Idle,
},
in_flight: protocol::InFlightSnapshot::default(),
internal_workers: Vec::new(),
})?),
})
}
}
#[async_trait::async_trait]
impl Socket for RetainedSnapshotSocket {
type Error = std::convert::Infallible;
async fn send(&mut self, _message: String) -> Result<(), Self::Error> {
Ok(())
}
async fn next(&mut self) -> Result<Option<String>, Self::Error> {
if self.event.is_some() {
return Ok(self.event.take());
}
std::future::pending::<Result<Option<String>, Self::Error>>().await
}
fn try_next(&mut self) -> Result<Option<String>, Self::Error> {
Ok(self.event.take())
}
}
fn enter_fullscreen() -> Result<ConsoleTerminal, Box<dyn std::error::Error>> {
let mut stdout = io::stdout();
// Enable button-event tracking so the transcript can own drag selection;