Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

feat: add retained Worker session observation

Browse Source
This commit is contained in:
Hare
2026-09-18 05:53:53 +09:00
parent 195aeb69be 9d8f294d75
commit 9566ad3482
30 changed files with 1277 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/tui/Cargo.toml
+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]
crates/tui/src/backend_worker_picker.rs
+75 -109
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, WorkerSessionAvailability,
list_backend_stopped_workers, list_backend_workers, observe_backend_worker_session,
};
use crossterm::event::{self, Event as TermEvent, KeyCode, KeyEventKind, KeyModifiers};
use ratatui::Frame;
@@ -81,62 +80,22 @@ 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())
.map_err(|error| io::Error::other(error.to_string()))?;
let restore = restore_backend_worker(&restore_target)
.await
.map_err(|error| {
io::Error::other(format!(
"failed to restore Backend worker {}/{}: {error}",
selected.runtime_id, selected.worker_id
))
})?;
restored_worker(restore).map_err(|error| {
let mut attach_target = target
.runtime_target(selected.runtime_id.clone(), selected.worker_id.clone())
.map_err(|error| io::Error::other(error.to_string()))?;
let response = observe_backend_worker_session(&attach_target)
.await
.map_err(|error| {
io::Error::other(format!(
"failed to restore Backend worker {}/{}: {error}",
"failed to observe Worker Session {}/{}: {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()))?;
})?;
apply_worker_session_observation(&mut attach_target, response.observation);
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())));
@@ -147,6 +106,28 @@ enum WorkerPickerResult {
SwitchWorkspace,
}
fn apply_worker_session_observation(
target: &mut client::BackendRuntimeTarget,
observation: WorkerSessionAvailability,
) {
match observation {
WorkerSessionAvailability::LiveProtocol => {}
WorkerSessionAvailability::RetainedSnapshot { snapshot, .. } => {
target.initial_snapshot = Some(snapshot);
target.initial_notice = Some("read-only retained Session snapshot".to_string());
}
WorkerSessionAvailability::Unavailable { message, .. } => {
let bounded_message: String = message.chars().take(512).collect();
target.initial_snapshot = Some(protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
});
target.initial_notice =
Some(format!("retained Session unavailable: {bounded_message}"));
}
}
}
fn pick_worker(
target: BackendRuntimeListTarget,
mut workers: Vec<BackendWorkerSummary>,
@@ -437,8 +418,7 @@ fn working_directory_text(worker: &BackendWorkerSummary) -> String {
mod tests {
use super::*;
use client::{
BackendDiagnostic, BackendDiagnosticSeverity, BackendWorkerCapabilitySummary,
BackendWorkerImplementationSummary, BackendWorkerRestoreResult,
BackendWorkerCapabilitySummary, BackendWorkerImplementationSummary,
BackendWorkerWorkspaceSummary,
};
@@ -494,65 +474,51 @@ 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"
fn authoritative_observation_selects_live_retained_and_unavailable_detail_modes() {
let mut live = client::BackendRuntimeTarget::new(
"http://127.0.0.1:3000",
"workspace-a",
"runtime-a",
"worker-a",
);
}
apply_worker_session_observation(&mut live, WorkerSessionAvailability::LiveProtocol);
assert!(live.initial_snapshot.is_none());
#[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");
let snapshot = protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
};
let mut retained = live.clone();
apply_worker_session_observation(
&mut retained,
WorkerSessionAvailability::RetainedSnapshot {
identity: client::RetainedSessionIdentity {
session_id: "session-a".to_string(),
segment_id: "segment-a".to_string(),
entry_count: 0,
},
snapshot: snapshot.clone(),
},
);
assert_eq!(retained.initial_snapshot, Some(snapshot));
assert!(retained.initial_notice.unwrap().contains("read-only"));
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);
let mut unavailable = live;
apply_worker_session_observation(
&mut unavailable,
WorkerSessionAvailability::Unavailable {
reason: client::WorkerSessionUnavailableReason::RetentionMissing,
message: "retention expired".to_string(),
},
);
assert!(unavailable.initial_snapshot.is_some());
assert!(
unavailable
.initial_notice
.unwrap()
.contains("retention expired")
);
}
#[test]
crates/tui/src/console/mod.rs
+69 -5
View File
@@ -458,19 +458,83 @@ async fn run_standalone_host(
pub(crate) async fn run_backend_runtime(
target: BackendRuntimeTarget,
) -> Result<(), Box<dyn std::error::Error>> {
let worker_label = target.display_label();
let attachment_target = target.clone();
let client = connect_backend_runtime(target).await?;
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 = match target.initial_notice.as_deref() {
Some(notice) => format!("{} — {notice}", target.display_label()),
None => target.display_label(),
};
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;