105 Commits
Author SHA1 Message Date
Hare 6b1b8a8846 fix: migrate legacy repository references 2026-09-03 11:09:55 +09:00
Hare 9fb1b90856 Merge remote-tracking branch 'refs/remotes/origin/hare/develop' into develop 2026-09-03 09:47:53 +09:00
Hare 30d4023475 fix: preserve blocked TUI submissions 2026-09-03 08:57:28 +09:00
Hare 70432f3d12 fix: retain attachments through TUI retry 2026-09-03 08:44:28 +09:00
Hare 7ee6c307fc fix: require exact attachment acceptance 2026-09-03 08:38:15 +09:00
Hare 36cfbbe6d2 fix: preserve attachment replay evidence 2026-09-03 08:21:11 +09:00
Hare a2e1a3d939 fix: preserve attachment lifecycle boundaries 2026-09-03 08:10:35 +09:00
Hare 9dc8d9a77a fix: fence attachment grants and lifecycle cleanup 2026-09-03 07:22:07 +09:00
Hare 3f6bb65eb1 test: align fork checkpoint regression 2026-09-03 05:55:48 +09:00
Hare 8a70f3cb26 fix: harden uploaded attachment retention and replay 2026-09-03 05:50:47 +09:00
Hare 6c5b8315a3 chore: remove generated test bindings 2026-09-03 05:07:14 +09:00
Hare 690ed0f121 fix: fence uploaded attachment lifecycle 2026-09-03 05:07:03 +09:00
Hare fd60c2b8be fix: render TUI attachment notices with supported alert level 2026-09-03 05:03:49 +09:00
Hare e27b4feb25 feat: upload client-local files from Web and TUI 2026-09-03 04:59:35 +09:00
Hare 09a33e7283 feat: add session-owned uploaded file attachments 2026-09-03 04:58:11 +09:00
Hare d87441448e test: normalize formatted workdir TypeScript 2026-09-03 03:18:09 +09:00
Hare f783f10f6e Merge branch 'work/companion' into hare/develop 2026-09-03 03:06:10 +09:00
Hare 70bdb2d723 test: align client workdir fixture with repository key 2026-09-03 03:03:32 +09:00
Hare 4c1ef04378 fix: show workspace list throughout global sidebar 2026-09-03 02:49:58 +09:00
Hare a595af133c feat: merge durable Workdir updates from develop
# Conflicts:
#	docs/README.md
#	docs/design/durable-operations.md
#	web/workspace/deno.json
2026-09-03 02:47:50 +09:00
Hare f74f3cd133 fix: reject conflicting repository import intent 2026-09-03 02:44:24 +09:00
Hare bcd4848458 Merge branch 'work/companion' into hare/develop 2026-09-03 02:39:54 +09:00
Hare fc05bf9711 fix: restore cursor after inline terminal menus 2026-09-03 02:34:27 +09:00
Hare 63ad590262 fix: remove implicit repository startup selection 2026-09-02 20:00:12 +09:00
Hare 0fd1193b6b fix: harden repository key projections 2026-09-02 19:53:18 +09:00
Hare d996822957 feat: list accessible workspaces in global sidebar 2026-09-02 19:22:45 +09:00
Hare 96349721cb feat: project repository keys across public clients 2026-09-02 19:03:55 +09:00
Hare 8344921b65 feat: add Composer input history navigation 2026-09-02 18:33:43 +09:00
Hare c97b3b7b77 feat: introduce workspace-scoped repository keys 2026-09-02 16:43:28 +09:00
Hare e00e675ed1 feat: add Worker actions to workspace sidebar 2026-09-02 16:24:29 +09:00
Hare 538da1f2b2 fix: validate workspace profile responses 2026-09-02 15:24:33 +09:00
Hare bad37ddc7d feat: centralize workspace profile DTOs 2026-09-02 15:24:12 +09:00
Hare 175eda9f29 fix: authenticate remote worker removal through workspace router 2026-09-02 14:43:49 +09:00
Hare 14c806d38f docs: define durable operation boundaries 2026-09-02 14:43:49 +09:00
Hare 925100fb82 fix: stabilize configuration editor completion 2026-09-02 14:18:51 +09:00
Hare b29b003ea3 feat: define orchestrator cleanup ownership 2026-09-02 02:26:35 +09:00
Hare faa727965b fix: fence orphan recovery claims 2026-09-02 01:59:32 +09:00
Hare d2ffbf2c40 fix: guard Workdir removal recovery ownership 2026-09-02 01:40:48 +09:00
Hare 5418fad7d7 fix: serialize Workdir removal attempts 2026-09-02 01:00:24 +09:00
Hare 510795f1c5 fix: fence Workdir removal retries 2026-09-02 00:25:21 +09:00
Hare 9e0d499987 Merge remote-tracking branch 'origin/develop' into work/T-578-durable-workdir-removal 2026-09-02 00:01:26 +09:00
Hare e96fde0632 feat: add durable Workdir removal authority 2026-09-02 00:01:16 +09:00
Hare eea79dead4 fix: bound Web UX review evidence and cleanup 2026-09-01 23:44:24 +09:00
Hare c4a3f4ba1e fix: harden Web UX capture profiles 2026-09-01 23:00:28 +09:00
Hare 4a4a01b730 feat: add repeatable Web UX inspection workbench 2026-09-01 22:48:26 +09:00
Hare a664e72488 feat: merge Composer paste handling from develop 2026-09-01 22:06:01 +09:00
Hare 21317123a4 fix: deduplicate CodeMirror completion state 2026-09-01 22:03:25 +09:00
Hare 409245cb52 fix: fence disabled Composer mutations 2026-09-01 21:00:26 +09:00
Hare 1aeb6fdb35 fix: enforce Composer read-only state 2026-09-01 20:48:16 +09:00
Hare 816fa96e07 fix: restore workers with remote workdir attachments 2026-09-01 20:41:37 +09:00
Hare 8fbe4218c6 fix: delete complete Composer selections 2026-09-01 20:34:23 +09:00
Hare 00c8df0fc9 fix: preserve exact short paste text 2026-09-01 20:19:53 +09:00
Hare c52c7ead19 feat: add atomic paste chips to Web Composer 2026-09-01 19:50:26 +09:00
Hare d1e8a827c2 fix: preserve configuration completion context 2026-09-01 18:59:16 +09:00
Hare 12d96fb03d fix: preserve short paste line endings 2026-09-01 18:49:44 +09:00
Hare 171a191873 test: run shared composer paste contract 2026-09-01 18:27:21 +09:00
Hare 47dabd8793 fix: keep short composer pastes as text 2026-09-01 18:23:23 +09:00
Hare 2bb661f1cf fix: complete paste artifact storage contract 2026-09-01 17:51:24 +09:00
Hare 04e296a4ef feat: store large paste inputs as artifacts 2026-09-01 17:26:58 +09:00
Hare 4bba227af5 fix: remove legacy resume flag 2026-09-01 16:10:59 +09:00
Hare 5cc78d63c6 fix: merge Companion repository response validation 2026-09-01 13:21:13 +09:00
Hare 323f5dc09c fix: validate repository settings wire response 2026-09-01 13:13:01 +09:00
Hare fb97edfe95 fix: validate workspace schema before owner migration 2026-09-01 12:41:14 +09:00
Hare 070f62ef12 fix: harden workspace owner migration preflight 2026-09-01 12:22:29 +09:00
Hare 981749aa3d feat: require account owners for workspaces 2026-09-01 12:10:52 +09:00
Hare e01b46b30a chore: shard Cargo build cache by workspace 2026-09-01 09:33:47 +09:00
Hare a1b659c45d chore: merge develop into hare/develop
# Conflicts:
#	crates/client/src/lib.rs
#	web/workspace/deno.json
2026-09-01 09:15:10 +09:00
Hare 37a012ef92 chore: configure shared Cargo build caching 2026-09-01 09:04:00 +09:00
Hare 4927e8a843 chore: merge current develop into workspace catalog DTO source
# Conflicts:
#	crates/workspace-api/src/lib.rs
#	crates/workspace-server/src/server.rs
#	web/workspace/deno.json
#	web/workspace/src/lib/workspace/sidebar/types.ts
2026-09-01 08:17:25 +09:00
Hare 025d6ddb47 refactor: make workspace API own workdir projections 2026-09-01 03:16:31 +09:00
Hare 01a4dfd5d3 fix: remove unused workspace API dependency 2026-09-01 02:52:11 +09:00
Hare 1d7158a0bf chore: merge current develop into T-563 source 2026-09-01 02:39:08 +09:00
Hare 9de2afbfc6 fix: validate ticket repository summaries 2026-09-01 02:37:44 +09:00
Hare 9013754a3a refactor: name the workspace catalog list response 2026-09-01 02:24:11 +09:00
Hare 21eea0b104 fix: consume shared workdir create contract 2026-09-01 02:14:57 +09:00
Hare 15e8d7365c test: run workspace DTO boundary regressions 2026-09-01 02:05:07 +09:00
Hare 88e3bf7065 fix: preserve optional workdir wire fields 2026-09-01 02:00:19 +09:00
Hare 2765138bf3 chore: merge current develop into Companion DTO work
# Conflicts:
#	web/workspace/deno.json
2026-09-01 01:59:39 +09:00
Hare 6b20ceac46 fix: validate workspace repository browser payloads 2026-09-01 01:49:38 +09:00
Hare d748274905 refactor: centralize workspace repository response DTOs 2026-09-01 01:49:25 +09:00
Hare e1578217d5 fix: preserve workdir status display contract 2026-09-01 01:45:32 +09:00
Hare 3481682cb4 feat: validate shared workdir REST contracts in web 2026-09-01 01:42:17 +09:00
Hare 996b7f2468 refactor: centralize workdir REST DTOs 2026-09-01 01:42:05 +09:00
Hare 879993b9b1 fix: accept zero-length Companion transcripts 2026-09-01 01:41:16 +09:00
Hare 6604154e3f test: cover repository access loader failures 2026-09-01 01:39:55 +09:00
Hare a9ad42a970 fix: align Companion REST DTO contracts 2026-09-01 01:29:04 +09:00
Hare 8b3d1302c6 fix: share repository access API contracts 2026-09-01 01:22:27 +09:00
Hare ac9269d6ce fix: omit local bash spill path from remote commands 2026-08-31 23:14:33 +09:00
Hare 8ffb716817 fix: preserve workdir operation error classifications 2026-08-31 23:14:23 +09:00
Hare 23f671fa48 fix: remove redundant settings overview 2026-08-31 20:09:29 +09:00
Hare d7cdcde443 fix: preserve standalone scope conflict details 2026-08-31 19:28:09 +09:00
Hare 95a81faf63 chore: merge hare/develop into companion work 2026-08-31 18:45:32 +09:00
Hare bb8bb6d099 feat: restore interactive standalone profile selection 2026-08-31 18:43:46 +09:00
Hare 310801a29b feat: add repository settings registration 2026-08-31 18:20:17 +09:00
Hare 7d09b20445 feat: add workspace repository registration API 2026-08-31 18:12:08 +09:00
Hare 456a06f194 fix: propagate browser actor to workspace routes 2026-08-31 16:55:45 +09:00
Hare 62eaefb1fa feat: spill long bash output to worker temp storage 2026-08-31 16:52:27 +09:00
Hare 10264b4019 Merge branch 'work/companion' into hare/develop 2026-08-31 15:51:36 +09:00
Hare ab9765d91d feat: refine workspace navigation chrome 2026-08-31 15:50:49 +09:00
Hare e7f4c6864f fix: make WorkerId the standalone primary identity 2026-08-31 14:37:46 +09:00
Hare bde1dea2a5 fix: require confirmation before Ctrl-X shutdown 2026-08-31 13:07:39 +09:00
Hare 13a021c480 feat: add TUI run status spinner 2026-08-31 12:46:47 +09:00
Hare a7f09fad98 fix: route standalone protocol through client transports 2026-08-31 12:11:35 +09:00
Hare 10eaf4a5fb refactor: name worker event channels by role 2026-08-31 11:21:14 +09:00
Hare bcada300e3 fix: mark standalone TUI connection ready 2026-08-31 09:44:13 +09:00
256 changed files with 26947 additions and 4360 deletions
+3
View File
@@ -0,0 +1,3 @@
[build]
build-dir = "{cargo-cache-home}/build/{workspace-path-hash}"
incremental = false
Generated
+26 -2
View File
@@ -637,6 +637,7 @@ checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9"
name = "client"
version = "0.1.0"
dependencies = [
"async-trait",
"chrono",
"futures",
"protocol",
@@ -649,7 +650,6 @@ dependencies = [
"tokio",
"tokio-tungstenite 0.29.0",
"uuid",
"workdir",
"workspace-api",
]
@@ -3507,6 +3507,7 @@ dependencies = [
"schemars",
"serde",
"serde_json",
"sha2 0.11.0",
"tokio",
"ts-rs",
"uuid",
@@ -4401,14 +4402,19 @@ dependencies = [
"agen",
"async-trait",
"base64 0.22.1",
"fs4",
"futures",
"protocol",
"serde",
"serde_json",
"sha2 0.11.0",
"tempfile",
"thiserror 2.0.18",
"tokio",
"tracing",
"unicode-normalization",
"unicode-properties",
"unicode-security",
"uuid",
]
@@ -4622,6 +4628,7 @@ version = "0.1.0"
dependencies = [
"agen",
"async-trait",
"client",
"fs4",
"futures",
"manifest",
@@ -5322,6 +5329,7 @@ name = "tui"
version = "0.1.0"
dependencies = [
"agen",
"async-trait",
"base64 0.22.1",
"client",
"crossterm 0.28.1",
@@ -5432,6 +5440,22 @@ version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7df058c713841ad818f1dc5d3fd88063241cc61f49f5fbea4b951e8cf5a8d71d"
[[package]]
name = "unicode-script"
version = "0.5.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "383ad40bb927465ec0ce7720e033cb4ca06912855fc35db31b5755d0de75b1ee"
[[package]]
name = "unicode-security"
version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2e4ddba1535dd35ed8b61c52166b7155d7f4e4b8847cec6f48e71dc66d8b5e50"
dependencies = [
"unicode-normalization",
"unicode-script",
]
[[package]]
name = "unicode-segmentation"
version = "1.13.2"
@@ -6591,6 +6615,7 @@ dependencies = [
"tempfile",
"thiserror 2.0.18",
"tokio",
"workspace-api",
]
[[package]]
@@ -6684,7 +6709,6 @@ dependencies = [
"serde",
"serde_json",
"ts-rs",
"workdir",
]
[[package]]
+1 -1
View File
@@ -5,6 +5,7 @@ edition.workspace = true
license.workspace = true
[dependencies]
async-trait.workspace = true
chrono = { version = "0.4", default-features = false, features = ["clock"] }
protocol = { workspace = true }
ticket = { workspace = true }
@@ -17,7 +18,6 @@ tokio = { workspace = true, features = ["rt", "macros", "net", "io-util", "sync"
tokio-tungstenite = { workspace = true }
uuid = { workspace = true }
workspace-api.workspace = true
workdir = { workspace = true }
[dev-dependencies]
tempfile = { workspace = true }
+71
View File
@@ -192,6 +192,32 @@ impl BackendApiClient {
format!("Bearer {}", self.access_token.0)
}
pub async fn require_success(
&self,
response: reqwest::Response,
) -> Result<reqwest::Response, BackendApiClientError> {
let status = response.status();
match status {
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => {
self.check_status(status)?;
}
status if !status.is_success() => {
let detail = response
.bytes()
.await
.ok()
.and_then(|body| backend_error_detail(&body));
return Err(BackendApiClientError::BackendResponse {
origin: self.origin.clone(),
status: status.as_u16(),
detail,
});
}
_ => {}
}
Ok(response)
}
pub fn check_status(&self, status: StatusCode) -> Result<(), BackendApiClientError> {
match status {
StatusCode::UNAUTHORIZED => Err(BackendApiClientError::Unauthorized {
@@ -235,6 +261,18 @@ fn redirect_policy(origin: BackendOrigin) -> redirect::Policy {
})
}
#[derive(Deserialize)]
struct BackendErrorBody {
message: String,
}
fn backend_error_detail(body: &[u8]) -> Option<String> {
serde_json::from_slice::<BackendErrorBody>(body)
.ok()
.map(|body| body.message)
.filter(|message| !message.trim().is_empty())
}
#[derive(Debug)]
pub enum BackendApiClientError {
InvalidBackendOrigin(String),
@@ -266,6 +304,11 @@ pub enum BackendApiClientError {
origin: BackendOrigin,
status: u16,
},
BackendResponse {
origin: BackendOrigin,
status: u16,
detail: Option<String>,
},
Io {
path: PathBuf,
source: std::io::Error,
@@ -312,6 +355,17 @@ impl fmt::Display for BackendApiClientError {
Self::BackendStatus { origin, status } => {
write!(f, "Backend {origin} returned HTTP {status}")
}
Self::BackendResponse {
origin,
status,
detail,
} => {
write!(f, "Backend {origin} returned HTTP {status}")?;
if let Some(detail) = detail {
write!(f, ": {detail}")?;
}
Ok(())
}
Self::Io { path, source } => {
write!(f, "failed to access {}: {source}", path.display())
}
@@ -584,6 +638,23 @@ mod tests {
);
}
#[test]
fn backend_error_detail_preserves_public_server_message() {
let detail = backend_error_detail(
br#"{"error":"Bad Request","message":"working_directory_runtime_mismatch: Working directory is owned by a different Runtime","diagnostics":[{"code":"working_directory_runtime_mismatch"}]}"#,
);
let error = BackendApiClientError::BackendResponse {
origin: BackendOrigin::parse("http://127.0.0.1:8787").unwrap(),
status: 400,
detail,
};
assert_eq!(
error.to_string(),
"Backend http://127.0.0.1:8787 returned HTTP 400: working_directory_runtime_mismatch: Working directory is owned by a different Runtime"
);
}
#[test]
fn backend_origin_rejects_unsafe_authority_changes() {
for invalid in [
+149 -178
View File
@@ -1,17 +1,11 @@
use crate::{BackendApiClient, BackendApiClientError};
use futures::{SinkExt, StreamExt};
use protocol::stream::{decode_event, encode_method};
use protocol::{ErrorCode, Event, Method};
use crate::transport::websocket::{Socket as WebSocket, SocketError as WebSocketError};
use crate::{BackendApiClient, BackendApiClientError, Client};
use reqwest::Method as HttpMethod;
use std::collections::VecDeque;
use serde::Deserialize;
use std::fmt;
use tokio::sync::mpsc;
use tokio_tungstenite::connect_async;
use tokio_tungstenite::tungstenite::Message as TungsteniteMessage;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
pub use workdir::workspace::WorkingDirectorySummary as BackendWorkingDirectorySummary;
pub use workspace_api::{
Diagnostic as BackendDiagnostic, DiagnosticSeverity as BackendDiagnosticSeverity,
ListResponse as BackendRuntimeListResponse, RuntimeSummary as BackendRuntimeSummary,
@@ -20,6 +14,11 @@ pub use workspace_api::{
WorkerRestoreResponse as BackendWorkerRestoreResponse,
WorkerRestoreResult as BackendWorkerRestoreResult, WorkerSummary as BackendWorkerSummary,
WorkerWorkspaceSummary as BackendWorkerWorkspaceSummary,
WorkingDirectoryCreateRequest as BackendWorkingDirectoryCreateRequest,
WorkingDirectoryCreateResponse as BackendWorkingDirectoryCreateResponse,
WorkingDirectoryDetailResponse as BackendWorkingDirectoryDetailResponse,
WorkingDirectoryListResponse as BackendWorkingDirectoryListResponse,
WorkingDirectorySummary as BackendWorkingDirectorySummary,
};
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -53,6 +52,123 @@ impl BackendRuntimeTarget {
pub fn display_label(&self) -> String {
format!("{}:{}", self.runtime_id, self.worker_id)
}
pub async fn upload_file(
&self,
file_name: &str,
media_type: &str,
content: Vec<u8>,
) -> Result<protocol::UploadedFileRef, BackendRuntimeClientError> {
self.upload_file_with_id(
&uuid::Uuid::now_v7().to_string(),
file_name,
media_type,
content,
)
.await
}
pub async fn upload_file_with_id(
&self,
upload_id: &str,
file_name: &str,
media_type: &str,
content: Vec<u8>,
) -> Result<protocol::UploadedFileRef, BackendRuntimeClientError> {
let api = BackendApiClient::from_stored_token(&self.base_url)?;
let worker_path = format!(
"/api/w/{}/runtimes/{}/workers/{}",
path_segment_encode(&self.workspace_id),
path_segment_encode(&self.runtime_id),
path_segment_encode(&self.worker_id),
);
let grant_path = format!(
"{worker_path}/attachment-upload-grants?file_name={}&media_type={}&upload_id={}",
path_segment_encode(file_name),
path_segment_encode(media_type),
path_segment_encode(&upload_id),
);
let grant_response = api
.request(HttpMethod::POST, &grant_path)?
.send()
.await
.map_err(BackendRuntimeClientError::Http)?;
api.check_status(grant_response.status())?;
let grant = grant_response
.json::<AttachmentUploadGrantResponse>()
.await
.map_err(BackendRuntimeClientError::Http)?;
let upload_path = format!(
"{worker_path}/attachment-uploads/{}",
path_segment_encode(&grant.upload_id),
);
let response = api
.request(HttpMethod::PUT, &upload_path)?
.body(content)
.send()
.await
.map_err(BackendRuntimeClientError::Http)?;
api.check_status(response.status())?;
response
.json::<UploadedFileResponse>()
.await
.map(|response| response.file)
.map_err(BackendRuntimeClientError::Http)
}
pub async fn cancel_file_upload(
&self,
upload_id: &str,
) -> Result<(), BackendRuntimeClientError> {
let api = BackendApiClient::from_stored_token(&self.base_url)?;
let path = format!(
"/api/w/{}/runtimes/{}/workers/{}/attachment-uploads/{}",
path_segment_encode(&self.workspace_id),
path_segment_encode(&self.runtime_id),
path_segment_encode(&self.worker_id),
path_segment_encode(upload_id),
);
let response = api
.request(HttpMethod::DELETE, &path)?
.send()
.await
.map_err(BackendRuntimeClientError::Http)?;
api.check_status(response.status())?;
Ok(())
}
pub async fn delete_uploaded_file(
&self,
artifact_id: &str,
) -> Result<(), BackendRuntimeClientError> {
let api = BackendApiClient::from_stored_token(&self.base_url)?;
let path = format!(
"/api/w/{}/runtimes/{}/workers/{}/attachments/{}",
path_segment_encode(&self.workspace_id),
path_segment_encode(&self.runtime_id),
path_segment_encode(&self.worker_id),
path_segment_encode(artifact_id),
);
let response = api
.request(HttpMethod::DELETE, &path)?
.send()
.await
.map_err(BackendRuntimeClientError::Http)?;
api.check_status(response.status())?;
Ok(())
}
}
#[derive(Deserialize)]
struct AttachmentUploadGrantResponse {
upload_id: String,
#[allow(dead_code)]
expires_at_ms: u64,
}
#[derive(Deserialize)]
struct UploadedFileResponse {
file: protocol::UploadedFileRef,
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -106,20 +222,12 @@ impl BackendRuntimeListTarget {
}
}
#[derive(Debug)]
pub struct BackendRuntimeClient {
target: BackendRuntimeTarget,
command_tx: mpsc::UnboundedSender<Method>,
events: mpsc::UnboundedReceiver<Event>,
diagnostics: VecDeque<Event>,
_protocol_task: tokio::task::JoinHandle<()>,
}
#[derive(Debug)]
pub enum BackendRuntimeClientError {
InvalidTarget(String),
Api(BackendApiClientError),
Http(reqwest::Error),
Protocol(String),
}
impl fmt::Display for BackendRuntimeClientError {
@@ -128,6 +236,7 @@ impl fmt::Display for BackendRuntimeClientError {
Self::InvalidTarget(message) => f.write_str(message),
Self::Api(error) => write!(f, "{error}"),
Self::Http(error) => write!(f, "{error}"),
Self::Protocol(message) => f.write_str(message),
}
}
}
@@ -278,155 +387,26 @@ pub async fn restore_backend_worker(
.json(&serde_json::json!({}))
.send()
.await?;
api.check_status(response.status())?;
let response = api.require_success(response).await?;
Ok(response.json::<BackendWorkerRestoreResponse>().await?)
}
impl BackendRuntimeClient {
pub async fn connect(target: BackendRuntimeTarget) -> Result<Self, BackendRuntimeClientError> {
validate_target(&target)?;
let api = BackendApiClient::from_stored_token(&target.base_url)?;
let (event_tx, rx) = mpsc::unbounded_channel();
let (command_tx, command_rx) = mpsc::unbounded_channel();
let protocol_target = target.clone();
let protocol_event_tx = event_tx.clone();
let protocol_task = tokio::spawn(async move {
run_worker_protocol_transport(protocol_target, api, command_rx, protocol_event_tx)
.await;
});
Ok(Self {
target,
command_tx,
events: rx,
diagnostics: VecDeque::new(),
_protocol_task: protocol_task,
})
}
pub fn try_next_event(&mut self) -> Option<Event> {
if let Some(event) = self.diagnostics.pop_front() {
return Some(event);
}
self.events.try_recv().ok()
}
pub async fn next_event(&mut self) -> Option<Event> {
if let Some(event) = self.diagnostics.pop_front() {
return Some(event);
}
self.events.recv().await
}
pub async fn send(&mut self, method: &Method) -> Result<(), BackendRuntimeClientError> {
self.command_tx.send(method.clone()).map_err(|_| {
BackendRuntimeClientError::InvalidTarget(format!(
"Backend protocol command stream is closed for {}",
self.target.display_label()
))
})?;
Ok(())
}
}
impl Drop for BackendRuntimeClient {
fn drop(&mut self) {
self._protocol_task.abort();
}
}
async fn run_worker_protocol_transport(
pub async fn connect_backend_runtime(
target: BackendRuntimeTarget,
api: BackendApiClient,
mut commands: mpsc::UnboundedReceiver<Method>,
tx: mpsc::UnboundedSender<Event>,
) {
let request = match protocol_ws_request(&target, &api) {
Ok(request) => request,
Err(error) => {
let _ = tx.send(diagnostic_event(format!(
"Backend protocol request could not be constructed for {}: {error}",
target.display_label()
)));
return;
}
};
match connect_async(request).await {
Ok((ws, _)) => {
let (mut sink, mut stream) = ws.split();
loop {
tokio::select! {
maybe_method = commands.recv() => {
let Some(method) = maybe_method else {
break;
};
match encode_method(&method) {
Ok(text) => {
if let Err(error) = sink.send(TungsteniteMessage::Text(text.into())).await {
let _ = tx.send(diagnostic_event(format!(
"Backend protocol command send failed for {}: {error}",
target.display_label()
)));
break;
}
}
Err(error) => {
let _ = tx.send(diagnostic_event(format!(
"Backend protocol command could not serialize method for {}: {error}",
target.display_label()
)));
}
}
}
frame = stream.next() => {
match frame {
Some(Ok(TungsteniteMessage::Text(text))) => {
match decode_event(&text) {
Ok(event) => {
let _ = tx.send(event);
}
Err(error) => {
let _ = tx.send(diagnostic_event(format!(
"Backend protocol response was not valid Event JSON for {}: {error}",
target.display_label()
)));
}
}
}
Some(Ok(TungsteniteMessage::Close(_))) | None => {
let _ = tx.send(diagnostic_event(format!(
"Backend protocol command stream closed for {}",
target.display_label()
)));
break;
}
Some(Ok(TungsteniteMessage::Ping(_)))
| Some(Ok(TungsteniteMessage::Pong(_)))
| Some(Ok(TungsteniteMessage::Binary(_)))
| Some(Ok(TungsteniteMessage::Frame(_))) => {}
Some(Err(error)) => {
let _ = tx.send(diagnostic_event(format!(
"Backend protocol WebSocket error for {}: {error}",
target.display_label()
)));
break;
}
}
}
}
}
}
Err(error) => {
let message = protocol_connect_error_message(&target, &api, &error);
let _ = tx.send(diagnostic_event(message));
while commands.recv().await.is_some() {
let _ = tx.send(diagnostic_event(format!(
"Backend protocol command was not sent because command stream is unavailable for {}",
target.display_label()
)));
}
}
) -> Result<Client<WebSocket>, BackendRuntimeClientError> {
validate_target(&target)?;
let api = BackendApiClient::from_stored_token(&target.base_url)?;
let request = protocol_ws_request(&target, &api).map_err(|error| {
BackendRuntimeClientError::Protocol(format!(
"Backend protocol request could not be constructed for {}: {error}",
target.display_label()
))
})?;
match WebSocket::connect(request).await {
Ok(socket) => Ok(Client::new(socket)),
Err(WebSocketError::WebSocket(error)) => Err(BackendRuntimeClientError::Protocol(
protocol_connect_error_message(&target, &api, &error),
)),
}
}
@@ -453,13 +433,6 @@ fn protocol_connect_error_message(
)
}
fn diagnostic_event(message: impl Into<String>) -> Event {
Event::Error {
code: ErrorCode::Internal,
message: message.into(),
}
}
fn validate_target(target: &BackendRuntimeTarget) -> Result<(), BackendRuntimeClientError> {
if target.base_url.trim().is_empty() {
return Err(BackendRuntimeClientError::InvalidTarget(
@@ -657,7 +630,7 @@ mod tests {
"capabilities": {"can_stop": true, "can_spawn_followup": false},
"working_directory": {
"working_directory_id": "wd-1",
"repository_id": "main",
"repository_key": "main",
"materializer_kind": "local_git_worktree",
"status": "active",
"occupied_by": {
@@ -670,13 +643,11 @@ mod tests {
});
let worker: BackendWorkerSummary = serde_json::from_value(payload.clone()).unwrap();
let occupied_by = worker
.working_directory
.unwrap()
.occupied_by
.expect("occupied Workdir");
assert_eq!(occupied_by.worker.runtime_id, "arcadia");
assert_eq!(occupied_by.worker.worker_id, "worker-opaque-64");
let workdir = worker.working_directory.unwrap();
assert_eq!(workdir.repository_key, "main");
let occupied_by = workdir.occupied_by.expect("occupied Workdir");
assert_eq!(occupied_by.runtime_id, "arcadia");
assert_eq!(occupied_by.worker_id, "worker-opaque-64");
let mut stale = payload;
stale["working_directory"]["occupied_by"]["runtime_worker_id"] = serde_json::json!(64);
+8 -35
View File
@@ -2,19 +2,16 @@ use crate::{BackendApiClient, BackendApiClientError};
use reqwest::Method;
use serde::{Deserialize, Serialize};
use std::fmt;
use workspace_api::{RepositoryObservedStatus, RepositorySource};
use workspace_api::{
WorkspaceCatalogListResponse, WorkspaceCreateResponse, WorkspaceRepositoryRecord,
WorkspaceSummary,
};
const DEFAULT_WORKSPACE_LIMIT: usize = 200;
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
pub struct BackendWorkspace {
pub workspace_id: String,
pub owner_account_id: Option<String>,
pub display_name: String,
pub state: String,
pub created_at: String,
pub updated_at: String,
}
pub type BackendWorkspace = WorkspaceSummary;
pub type CreateBackendWorkspaceResponse = WorkspaceCreateResponse;
pub type CreateBackendWorkspaceRepositoryRecord = WorkspaceRepositoryRecord;
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
@@ -32,30 +29,6 @@ pub struct CreateBackendWorkspaceRepository {
pub default_ref: Option<String>,
}
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
pub struct CreateBackendWorkspaceResponse {
pub workspace: BackendWorkspace,
pub repository: CreateBackendWorkspaceRepositoryRecord,
pub config_revision: u64,
pub request_fingerprint: String,
pub replayed: bool,
}
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
pub struct CreateBackendWorkspaceRepositoryRecord {
pub workspace_id: String,
pub repository_id: String,
pub name: String,
pub kind: String,
pub provider: Option<String>,
pub source: RepositorySource,
pub default_ref: Option<String>,
pub source_revision: u64,
pub source_fingerprint: String,
pub observed_status: RepositoryObservedStatus,
pub observed_at: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BackendWorkspaceCatalogTarget {
pub base_url: String,
@@ -118,7 +91,7 @@ async fn list_backend_workspaces_with_client(
.send()
.await?;
client.check_status(response.status())?;
Ok(response.json::<Vec<BackendWorkspace>>().await?)
Ok(response.json::<WorkspaceCatalogListResponse>().await?.0)
}
pub async fn create_backend_workspace(
+137
View File
@@ -0,0 +1,137 @@
use std::error::Error;
use std::fmt;
use protocol::stream::{decode_event, encode_method};
use protocol::{Event, Method};
use crate::transport::Socket;
/// Typed Worker protocol client over an injected message transport.
pub struct Client<T> {
socket: T,
}
#[derive(Debug)]
pub enum ClientError<E> {
Transport(E),
Protocol(serde_json::Error),
}
impl<T> Client<T> {
pub fn new(socket: T) -> Self {
Self { socket }
}
pub fn into_inner(self) -> T {
self.socket
}
}
impl<T: Socket> Client<T> {
pub async fn send(&mut self, method: &Method) -> Result<(), ClientError<T::Error>> {
let message = encode_method(method).map_err(ClientError::Protocol)?;
self.socket
.send(message)
.await
.map_err(ClientError::Transport)
}
pub async fn next_event(&mut self) -> Result<Option<Event>, ClientError<T::Error>> {
self.socket
.next()
.await
.map_err(ClientError::Transport)?
.map(|message| decode_event(&message).map_err(ClientError::Protocol))
.transpose()
}
pub fn try_next_event(&mut self) -> Result<Option<Event>, ClientError<T::Error>> {
self.socket
.try_next()
.map_err(ClientError::Transport)?
.map(|message| decode_event(&message).map_err(ClientError::Protocol))
.transpose()
}
}
impl<E: fmt::Display> fmt::Display for ClientError<E> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Transport(error) => write!(formatter, "Worker transport error: {error}"),
Self::Protocol(error) => write!(formatter, "Worker protocol error: {error}"),
}
}
}
impl<E: Error + 'static> Error for ClientError<E> {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
Self::Transport(error) => Some(error),
Self::Protocol(error) => Some(error),
}
}
}
#[cfg(test)]
mod tests {
use std::collections::VecDeque;
use std::convert::Infallible;
use async_trait::async_trait;
use protocol::stream::{decode_method, encode_event};
use protocol::{Event, Method, WorkerStatus};
use super::Client;
use crate::transport::Socket;
#[derive(Default)]
struct TestSocket {
sent: Vec<String>,
incoming: VecDeque<String>,
}
#[async_trait]
impl Socket for TestSocket {
type Error = Infallible;
async fn send(&mut self, message: String) -> Result<(), Self::Error> {
self.sent.push(message);
Ok(())
}
async fn next(&mut self) -> Result<Option<String>, Self::Error> {
Ok(self.incoming.pop_front())
}
fn try_next(&mut self) -> Result<Option<String>, Self::Error> {
Ok(self.incoming.pop_front())
}
}
#[tokio::test]
async fn encodes_methods_and_decodes_events_above_transport() {
let mut socket = TestSocket::default();
socket.incoming.push_back(
encode_event(&Event::Status {
status: WorkerStatus::Idle,
})
.expect("encode event"),
);
let mut client = Client::new(socket);
client
.send(&Method::run_text("hello"))
.await
.expect("send method");
assert!(matches!(
decode_method(&client.socket.sent[0]),
Ok(Method::Run { .. })
));
assert!(matches!(
client.next_event().await,
Ok(Some(Event::Status {
status: WorkerStatus::Idle
}))
));
}
}
+15 -9
View File
@@ -7,8 +7,9 @@ pub mod backend_api;
mod backend_auth;
pub mod backend_runtime;
pub mod backend_workspace;
mod client;
pub mod target;
mod worker_client;
pub mod transport;
mod workspace_product;
pub use backend_api::{
@@ -20,23 +21,28 @@ pub use backend_auth::{
poll_device_login, start_device_login, wait_for_device_login,
};
pub use backend_runtime::{
BackendDiagnostic, BackendDiagnosticSeverity, BackendRuntimeClient, BackendRuntimeClientError,
BackendDiagnostic, BackendDiagnosticSeverity, BackendRuntimeClientError,
BackendRuntimeListResponse, BackendRuntimeListTarget, BackendRuntimeSummary,
BackendRuntimeTarget, BackendWorkerCapabilitySummary, BackendWorkerImplementationSummary,
BackendWorkerRestoreResponse, BackendWorkerRestoreResult, BackendWorkerSummary,
BackendWorkerWorkspaceSummary, BackendWorkingDirectorySummary, list_backend_stopped_workers,
list_backend_workers, restore_backend_worker,
BackendWorkerWorkspaceSummary, BackendWorkingDirectorySummary, connect_backend_runtime,
list_backend_stopped_workers, list_backend_workers, restore_backend_worker,
};
pub use backend_workspace::{
BackendWorkspace, BackendWorkspaceCatalogTarget, BackendWorkspaceClientError,
CreateBackendWorkspaceRepository, CreateBackendWorkspaceRequest,
CreateBackendWorkspaceResponse, create_backend_workspace, list_backend_workspaces,
};
pub use client::{Client, ClientError};
pub use target::{
BackendTarget, Dashboard, ResolvedTarget, StandaloneSessionListIntent,
StandaloneSessionResumeIntent, StandaloneTarget, Target, TargetError, TargetKind,
WorkerConnection, WorkerConnectionSelector, WorkerList, WorkerListRequest, WorkerSpawn,
BackendTarget, Dashboard, ResolvedTarget, StandaloneTarget, StandaloneWorkerListIntent,
StandaloneWorkerResumeIntent, Target, TargetError, TargetKind, WorkerConnection,
WorkerConnectionSelector, WorkerList, WorkerListRequest, WorkerSpawn,
};
pub use workspace_api::{
CompanionCancelRequest, CompanionLifecycleState, CompanionMessageDisposition,
CompanionMessageRequest, CompanionMessageResponse, CompanionStatusResponse,
CompanionTranscriptItem, CompanionTranscriptProjection, CompanionTranscriptRole,
CompanionTransportSummary, ObjectiveDetail, ObjectiveSummary,
};
pub use worker_client::WorkerClient;
pub use workspace_api::{ObjectiveDetail, ObjectiveSummary};
pub use workspace_product::BackendWorkspaceProductClient;
+24 -24
View File
@@ -105,16 +105,16 @@ pub struct WorkerSpawn {
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StandaloneSessionListIntent {
pub struct StandaloneWorkerListIntent {
pub state_dir: PathBuf,
pub cwd: PathBuf,
pub include_all: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StandaloneSessionResumeIntent {
pub struct StandaloneWorkerResumeIntent {
pub state_dir: PathBuf,
pub session_id: String,
pub worker_id: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -175,22 +175,22 @@ pub trait Target: fmt::Debug + Send + Sync {
Err(TargetError::unsupported("Worker spawn", self.kind()))
}
fn standalone_session_list(
fn standalone_worker_list(
&self,
_include_all: bool,
) -> Result<StandaloneSessionListIntent, TargetError> {
) -> Result<StandaloneWorkerListIntent, TargetError> {
Err(TargetError::unsupported(
"standalone session listing",
"standalone Worker listing",
self.kind(),
))
}
fn standalone_session_resume(
fn standalone_worker_resume(
&self,
_session_id: String,
) -> Result<StandaloneSessionResumeIntent, TargetError> {
_worker_id: String,
) -> Result<StandaloneWorkerResumeIntent, TargetError> {
Err(TargetError::unsupported(
"standalone session restore",
"standalone Worker restore",
self.kind(),
))
}
@@ -243,26 +243,26 @@ impl Target for StandaloneTarget {
})
}
fn standalone_session_list(
fn standalone_worker_list(
&self,
include_all: bool,
) -> Result<StandaloneSessionListIntent, TargetError> {
) -> Result<StandaloneWorkerListIntent, TargetError> {
let cwd = std::env::current_dir()
.map_err(|error| TargetError::invalid(self.kind(), error.to_string()))?;
Ok(StandaloneSessionListIntent {
Ok(StandaloneWorkerListIntent {
state_dir: self.state_dir.clone(),
cwd,
include_all,
})
}
fn standalone_session_resume(
fn standalone_worker_resume(
&self,
session_id: String,
) -> Result<StandaloneSessionResumeIntent, TargetError> {
Ok(StandaloneSessionResumeIntent {
worker_id: String,
) -> Result<StandaloneWorkerResumeIntent, TargetError> {
Ok(StandaloneWorkerResumeIntent {
state_dir: self.state_dir.clone(),
session_id,
worker_id,
})
}
}
@@ -437,17 +437,17 @@ mod tests {
}
#[test]
fn standalone_target_builds_explicit_session_intents() {
let target = StandaloneTarget::new("/tmp/yoi-client-sessions");
let list = target.standalone_session_list(true).unwrap();
assert_eq!(list.state_dir, PathBuf::from("/tmp/yoi-client-sessions"));
fn standalone_target_builds_explicit_worker_intents() {
let target = StandaloneTarget::new("/tmp/yoi-client-workers");
let list = target.standalone_worker_list(true).unwrap();
assert_eq!(list.state_dir, PathBuf::from("/tmp/yoi-client-workers"));
assert!(list.include_all);
assert!(list.cwd.is_absolute());
let resume = target
.standalone_session_resume("019d1234-0000-7000-8000-000000000000".to_string())
.standalone_worker_resume("019d1234-0000-7000-8000-000000000000".to_string())
.unwrap();
assert_eq!(resume.state_dir, list.state_dir);
assert_eq!(resume.session_id, "019d1234-0000-7000-8000-000000000000");
assert_eq!(resume.worker_id, "019d1234-0000-7000-8000-000000000000");
}
}
+115
View File
@@ -0,0 +1,115 @@
use async_trait::async_trait;
use thiserror::Error;
use tokio::sync::mpsc;
use super::Socket as SocketContract;
const CHANNEL_CAPACITY: usize = 256;
pub struct Socket {
outgoing: mpsc::Sender<String>,
incoming: mpsc::Receiver<String>,
}
/// Host-side endpoint paired with an in-process client transport.
pub struct Peer {
incoming: mpsc::Receiver<String>,
outgoing: mpsc::Sender<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
pub enum SocketError {
#[error("in-process Worker protocol transport closed")]
Closed,
}
impl Socket {
pub fn pair() -> (Self, Peer) {
let (client_tx, peer_rx) = mpsc::channel(CHANNEL_CAPACITY);
let (peer_tx, client_rx) = mpsc::channel(CHANNEL_CAPACITY);
(
Self {
outgoing: client_tx,
incoming: client_rx,
},
Peer {
incoming: peer_rx,
outgoing: peer_tx,
},
)
}
}
#[async_trait]
impl SocketContract for Socket {
type Error = SocketError;
async fn send(&mut self, message: String) -> Result<(), Self::Error> {
self.outgoing
.send(message)
.await
.map_err(|_| SocketError::Closed)
}
async fn next(&mut self) -> Result<Option<String>, Self::Error> {
Ok(self.incoming.recv().await)
}
fn try_next(&mut self) -> Result<Option<String>, Self::Error> {
match self.incoming.try_recv() {
Ok(message) => Ok(Some(message)),
Err(mpsc::error::TryRecvError::Empty | mpsc::error::TryRecvError::Disconnected) => {
Ok(None)
}
}
}
}
impl Peer {
pub async fn next(&mut self) -> Option<String> {
self.incoming.recv().await
}
pub async fn send(&self, message: String) -> Result<(), String> {
self.outgoing.send(message).await.map_err(|error| error.0)
}
}
#[cfg(test)]
mod tests {
use protocol::stream::{decode_method, encode_event};
use protocol::{Event, Method, WorkerStatus};
use super::Socket;
use crate::Client;
#[tokio::test]
async fn pair_carries_typed_protocol_through_generic_client() {
let (socket, mut peer) = Socket::pair();
let mut client = Client::new(socket);
client
.send(&Method::run_text("hello"))
.await
.expect("send method");
assert!(matches!(
peer.next().await.as_deref().map(decode_method),
Some(Ok(Method::Run { .. }))
));
peer.send(
encode_event(&Event::Status {
status: WorkerStatus::Idle,
})
.expect("encode event"),
)
.await
.expect("send event");
assert!(matches!(
client.next_event().await,
Ok(Some(Event::Status {
status: WorkerStatus::Idle
}))
));
}
}
+22
View File
@@ -0,0 +1,22 @@
use std::error::Error;
use async_trait::async_trait;
pub mod in_process;
pub mod unix_socket;
pub mod websocket;
/// Message-oriented transport for one Worker protocol connection.
///
/// Implementations own physical framing. `client::Client` owns the typed
/// Method/Event protocol encoding layered on top of these UTF-8 messages.
#[async_trait]
pub trait Socket {
type Error: Error + Send + Sync + 'static;
async fn send(&mut self, message: String) -> Result<(), Self::Error>;
async fn next(&mut self) -> Result<Option<String>, Self::Error>;
fn try_next(&mut self) -> Result<Option<String>, Self::Error>;
}
+172
View File
@@ -0,0 +1,172 @@
use std::io;
use std::path::Path;
use async_trait::async_trait;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::UnixStream;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use super::Socket as SocketContract;
pub struct Socket {
writer: tokio::io::WriteHalf<UnixStream>,
messages: mpsc::Receiver<io::Result<String>>,
reader_task: JoinHandle<()>,
}
impl Socket {
pub async fn connect(path: &Path) -> io::Result<Self> {
let stream = UnixStream::connect(path).await?;
let (reader, writer) = tokio::io::split(stream);
let (message_tx, messages) = mpsc::channel(256);
let reader_task = tokio::spawn(async move {
let mut lines = BufReader::new(reader).lines();
loop {
match lines.next_line().await {
Ok(Some(message)) if message.trim().is_empty() => {}
Ok(Some(message)) => {
if message_tx.send(Ok(message)).await.is_err() {
return;
}
}
Ok(None) => return,
Err(error) => {
let _ = message_tx.send(Err(error)).await;
return;
}
}
}
});
Ok(Self {
writer,
messages,
reader_task,
})
}
}
#[async_trait]
impl SocketContract for Socket {
type Error = io::Error;
async fn send(&mut self, message: String) -> Result<(), Self::Error> {
self.writer.write_all(message.as_bytes()).await?;
self.writer.write_all(b"\n").await?;
self.writer.flush().await
}
async fn next(&mut self) -> Result<Option<String>, Self::Error> {
match self.messages.recv().await {
Some(message) => message.map(Some),
None => Ok(None),
}
}
fn try_next(&mut self) -> Result<Option<String>, Self::Error> {
match self.messages.try_recv() {
Ok(message) => message.map(Some),
Err(mpsc::error::TryRecvError::Empty | mpsc::error::TryRecvError::Disconnected) => {
Ok(None)
}
}
}
}
impl Drop for Socket {
fn drop(&mut self) {
self.reader_task.abort();
}
}
#[cfg(test)]
mod tests {
use std::io::ErrorKind;
use std::time::Duration;
use protocol::stream::{decode_method, encode_event};
use protocol::{Event, Method, WorkerStatus};
use tempfile::tempdir;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::UnixListener;
use super::*;
use crate::Client;
async fn assert_peer_closed(stream: &mut UnixStream, reason: &str) {
let mut buf = [0_u8; 1];
match tokio::time::timeout(Duration::from_secs(1), stream.read(&mut buf))
.await
.expect(reason)
{
Ok(0) => {}
Err(error) if error.kind() == ErrorKind::ConnectionReset => {}
Ok(n) => panic!("server should observe peer close, read {n} byte(s)"),
Err(error) => panic!("server read failed unexpectedly: {error}"),
}
}
#[tokio::test]
async fn client_receives_events_over_unix_socket() {
let socket_dir = tempdir().unwrap();
let socket_path = socket_dir.path().join("events.sock");
let listener = UnixListener::bind(&socket_path).unwrap();
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
let event = encode_event(&Event::Status {
status: WorkerStatus::Idle,
})
.unwrap();
stream.write_all(event.as_bytes()).await.unwrap();
stream.write_all(b"\n").await.unwrap();
});
let mut client = Client::new(Socket::connect(&socket_path).await.unwrap());
let event = tokio::time::timeout(Duration::from_secs(1), client.next_event())
.await
.expect("client should receive event while alive")
.expect("transport should succeed");
assert!(matches!(
event,
Some(Event::Status {
status: WorkerStatus::Idle
})
));
server.await.unwrap();
}
#[tokio::test]
async fn client_sends_methods_over_unix_socket() {
let socket_dir = tempdir().unwrap();
let socket_path = socket_dir.path().join("send.sock");
let listener = UnixListener::bind(&socket_path).unwrap();
let server = tokio::spawn(async move {
let (reader, _) = listener.accept().await.unwrap();
BufReader::new(reader).lines().next_line().await.unwrap()
});
let mut client = Client::new(Socket::connect(&socket_path).await.unwrap());
client
.send(&Method::run_text("hello"))
.await
.expect("send method");
let received = server.await.unwrap().expect("method message");
assert!(matches!(decode_method(&received), Ok(Method::Run { .. })));
}
#[tokio::test]
async fn dropping_socket_closes_server_connection() {
let socket_dir = tempdir().unwrap();
let socket_path = socket_dir.path().join("drop.sock");
let listener = UnixListener::bind(&socket_path).unwrap();
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
assert_peer_closed(&mut stream, "dropped socket should close promptly").await;
});
let socket = Socket::connect(&socket_path).await.unwrap();
drop(socket);
server.await.unwrap();
}
}
+140
View File
@@ -0,0 +1,140 @@
use async_trait::async_trait;
use futures::{SinkExt, StreamExt};
use thiserror::Error;
use tokio::net::TcpStream;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use tokio_tungstenite::tungstenite::http::Request;
use tokio_tungstenite::tungstenite::{self, Message};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
use super::Socket as SocketContract;
type Writer = futures::stream::SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>;
pub struct Socket {
writer: Writer,
messages: mpsc::Receiver<Result<String, SocketError>>,
reader_task: JoinHandle<()>,
}
#[derive(Debug, Error)]
pub enum SocketError {
#[error("WebSocket transport failed: {0}")]
WebSocket(#[from] tungstenite::Error),
}
impl Socket {
pub async fn connect(request: Request<()>) -> Result<Self, SocketError> {
let (stream, _) = connect_async(request).await?;
let (writer, mut reader) = stream.split();
let (message_tx, messages) = mpsc::channel(256);
let reader_task = tokio::spawn(async move {
loop {
match reader.next().await {
Some(Ok(Message::Text(message))) => {
if message_tx.send(Ok(message.to_string())).await.is_err() {
return;
}
}
Some(Ok(Message::Close(_))) | None => return,
Some(Ok(
Message::Binary(_)
| Message::Ping(_)
| Message::Pong(_)
| Message::Frame(_),
)) => {}
Some(Err(error)) => {
let _ = message_tx.send(Err(SocketError::WebSocket(error))).await;
return;
}
}
}
});
Ok(Self {
writer,
messages,
reader_task,
})
}
}
#[async_trait]
impl SocketContract for Socket {
type Error = SocketError;
async fn send(&mut self, message: String) -> Result<(), Self::Error> {
self.writer.send(Message::Text(message.into())).await?;
Ok(())
}
async fn next(&mut self) -> Result<Option<String>, Self::Error> {
match self.messages.recv().await {
Some(message) => message.map(Some),
None => Ok(None),
}
}
fn try_next(&mut self) -> Result<Option<String>, Self::Error> {
match self.messages.try_recv() {
Ok(message) => message.map(Some),
Err(mpsc::error::TryRecvError::Empty | mpsc::error::TryRecvError::Disconnected) => {
Ok(None)
}
}
}
}
impl Drop for Socket {
fn drop(&mut self) {
self.reader_task.abort();
}
}
#[cfg(test)]
mod tests {
use futures::{SinkExt, StreamExt};
use protocol::stream::{decode_method, encode_event};
use protocol::{Event, Method, WorkerStatus};
use tokio::net::TcpListener;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use super::*;
use crate::Client;
#[tokio::test]
async fn carries_typed_protocol_through_generic_client() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut socket = accept_async(stream).await.unwrap();
let message = socket.next().await.unwrap().unwrap();
assert!(matches!(
message,
Message::Text(ref text)
if matches!(decode_method(text), Ok(Method::Run { .. }))
));
let event = encode_event(&Event::Status {
status: WorkerStatus::Idle,
})
.unwrap();
socket.send(Message::Text(event.into())).await.unwrap();
});
let request = format!("ws://{address}").into_client_request().unwrap();
let mut client = Client::new(Socket::connect(request).await.unwrap());
client
.send(&Method::run_text("hello"))
.await
.expect("send method");
assert!(matches!(
client.next_event().await,
Ok(Some(Event::Status {
status: WorkerStatus::Idle
}))
));
server.await.unwrap();
}
}
-186
View File
@@ -1,186 +0,0 @@
use std::io;
use std::path::Path;
use protocol::stream::{JsonLineReader, JsonLineWriter};
use protocol::{Event, Method};
use tokio::net::UnixStream;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
pub struct WorkerClient {
writer: JsonLineWriter<tokio::io::WriteHalf<UnixStream>>,
event_rx: mpsc::Receiver<Event>,
reader_task: JoinHandle<()>,
}
impl WorkerClient {
pub async fn connect(path: &Path) -> Result<Self, io::Error> {
let stream = UnixStream::connect(path).await?;
let (reader, writer) = tokio::io::split(stream);
let writer = JsonLineWriter::new(writer);
let (event_tx, event_rx) = mpsc::channel::<Event>(256);
let reader_task = tokio::spawn(async move {
let mut reader = JsonLineReader::new(reader);
while let Ok(Some(event)) = reader.next::<Event>().await {
if event_tx.send(event).await.is_err() {
break;
}
}
});
Ok(Self {
writer,
event_rx,
reader_task,
})
}
pub async fn send(&mut self, method: &Method) -> Result<(), io::Error> {
self.writer.write(method).await
}
pub fn try_next_event(&mut self) -> Option<Event> {
self.event_rx.try_recv().ok()
}
pub async fn next_event(&mut self) -> Option<Event> {
self.event_rx.recv().await
}
}
impl Drop for WorkerClient {
fn drop(&mut self) {
self.reader_task.abort();
}
}
#[cfg(test)]
mod tests {
use std::io::ErrorKind;
use std::time::Duration;
use protocol::{Segment, WorkerStatus};
use tempfile::tempdir;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::UnixListener;
use super::*;
async fn assert_peer_closed(stream: &mut UnixStream, reason: &str) {
let mut buf = [0_u8; 1];
match tokio::time::timeout(Duration::from_secs(1), stream.read(&mut buf))
.await
.expect(reason)
{
Ok(0) => {}
Err(error) if error.kind() == ErrorKind::ConnectionReset => {}
Ok(n) => panic!("server should observe peer close, read {n} byte(s)"),
Err(error) => panic!("server read failed unexpectedly: {error}"),
}
}
#[tokio::test]
async fn receives_events_while_client_is_alive() {
let socket_dir = tempdir().unwrap();
let socket_path = socket_dir.path().join("events.sock");
let listener = UnixListener::bind(&socket_path).unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut writer = JsonLineWriter::new(stream);
writer
.write(&Event::Status {
status: WorkerStatus::Idle,
})
.await
.unwrap();
});
let mut client = WorkerClient::connect(&socket_path).await.unwrap();
let event = tokio::time::timeout(Duration::from_secs(1), client.next_event())
.await
.expect("client should receive event while alive");
assert!(matches!(
event,
Some(Event::Status {
status: WorkerStatus::Idle
})
));
server.await.unwrap();
}
#[tokio::test]
async fn send_writes_methods_while_client_is_alive() {
let socket_dir = tempdir().unwrap();
let socket_path = socket_dir.path().join("send.sock");
let listener = UnixListener::bind(&socket_path).unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut reader = JsonLineReader::new(stream);
reader.next::<Method>().await.unwrap()
});
let mut client = WorkerClient::connect(&socket_path).await.unwrap();
let method = Method::Run {
input: vec![Segment::text("hello")],
};
client.send(&method).await.unwrap();
let received = tokio::time::timeout(Duration::from_secs(1), server)
.await
.expect("server should receive method while client is alive")
.unwrap();
match received {
Some(Method::Run { input }) => assert_eq!(input, vec![Segment::text("hello")]),
other => panic!("expected Run method, got {other:?}"),
}
}
#[tokio::test]
async fn dropping_repeated_clients_closes_server_connections() {
let socket_dir = tempdir().unwrap();
let socket_path = socket_dir.path().join("drop.sock");
let listener = UnixListener::bind(&socket_path).unwrap();
let server = tokio::spawn(async move {
for _ in 0..16 {
let (mut stream, _) = listener.accept().await.unwrap();
assert_peer_closed(
&mut stream,
"dropped client should close its socket promptly",
)
.await;
}
});
for _ in 0..16 {
let client = WorkerClient::connect(&socket_path).await.unwrap();
drop(client);
}
server.await.unwrap();
}
#[tokio::test]
async fn dropping_client_aborts_blocked_reader_task() {
let socket_dir = tempdir().unwrap();
let socket_path = socket_dir.path().join("blocked-reader.sock");
let listener = UnixListener::bind(&socket_path).unwrap();
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
stream.write_all(b"{\"event\"").await.unwrap();
assert_peer_closed(
&mut stream,
"aborting the blocked client reader should close the socket",
)
.await;
});
let client = WorkerClient::connect(&socket_path).await.unwrap();
tokio::task::yield_now().await;
drop(client);
server.await.unwrap();
}
}
+28 -14
View File
@@ -101,20 +101,24 @@ pub fn complete_current(
let utf8_byte_offset = utf16_to_utf8_offset(&source, utf16_offset)?;
let result = session_environment(snapshot.clone())
.complete_config(&entrypoint, &source, utf8_byte_offset, explicit)
.map_err(|error| JsValue::from_str(&format!("{error:?}")))?
.map(|result| WasmCompletionResult {
from: result.from,
items: result
.items
.into_iter()
.map(|item| WasmCompletionItem {
label: item.label,
kind: format!("{:?}", item.kind).to_lowercase(),
detail: item.detail,
priority: item.priority,
})
.collect(),
});
.map_err(|error| JsValue::from_str(&format!("{error:?}")))?;
let result = result
.map(|result| {
Ok::<WasmCompletionResult, JsValue>(WasmCompletionResult {
from: utf8_to_utf16_offset(&source, result.from)?,
items: result
.items
.into_iter()
.map(|item| WasmCompletionItem {
label: item.label,
kind: format!("{:?}", item.kind).to_lowercase(),
detail: item.detail,
priority: item.priority,
})
.collect(),
})
})
.transpose()?;
encode(result)
})
}
@@ -177,6 +181,16 @@ fn utf16_to_utf8_offset(source: &str, utf16_offset: usize) -> Result<usize, JsVa
}
}
fn utf8_to_utf16_offset(source: &str, utf8_offset: usize) -> Result<usize, JsValue> {
if utf8_offset > source.len() {
return Err(JsValue::from_str("UTF-8 offset is outside the source"));
}
if !source.is_char_boundary(utf8_offset) {
return Err(JsValue::from_str("UTF-8 offset splits a character"));
}
Ok(source[..utf8_offset].encode_utf16().count())
}
fn decode<T: serde::de::DeserializeOwned>(value: JsValue) -> Result<T, JsValue> {
from_value(value).map_err(|error| JsValue::from_str(&error.to_string()))
}
+28
View File
@@ -1203,6 +1203,9 @@ impl SnapshotEnvironment {
{
let mut member_source = format!("{WORKSPACE_CONFIG_SCHEMA_GLOBAL}.");
member_source.push_str(&context.schema_path.join("."));
if !context.schema_path.is_empty() && context.from == utf8_byte_offset {
member_source.push('.');
}
let mut completion = LanguageService::new(self).complete(
entrypoint.as_str(),
&member_source,
@@ -1961,6 +1964,31 @@ mod tests {
.iter()
.any(|item| item.label == "default_profile")
);
let blank_nested_source = "{ profile = { } } as WorkspaceConfigSchema";
let blank_nested_cursor = blank_nested_source.find("{ }").unwrap() + 2;
let blank_nested = environment
.complete_config(
&path("main.dcdl"),
blank_nested_source,
blank_nested_cursor,
true,
)
.unwrap()
.unwrap();
assert_eq!(blank_nested.from, blank_nested_cursor);
assert!(
blank_nested
.items
.iter()
.any(|item| item.label == "default_profile")
);
assert!(
!blank_nested
.items
.iter()
.any(|item| item.label == "profile")
);
}
#[test]
+2
View File
@@ -177,6 +177,8 @@ mod tests {
fn logical_paths_reject_absolute_parent_and_backslash_forms() {
assert!(FsPath::new("src/lib.rs").is_ok());
assert!(FsPath::new("/tmp/file").is_err());
assert!(FsPath::new_scoped("/tmp/file").is_ok());
assert!(FsPath::new_scoped("/tmp/../secret").is_err());
assert!(FsPath::new("../file").is_err());
assert!(FsPath::new("src\\lib.rs").is_err());
}
+22 -2
View File
@@ -5,7 +5,8 @@ use serde::{Deserialize, Serialize};
use crate::FsError;
/// Logical path relative to the bound Workdir root.
/// Scope-checked filesystem path. Relative paths resolve below the bound
/// Workdir root; absolute paths require an explicit matching scope rule.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize)]
#[serde(transparent)]
pub struct FsPath(String);
@@ -16,11 +17,30 @@ impl<'de> Deserialize<'de> for FsPath {
D: serde::Deserializer<'de>,
{
let value = String::deserialize(deserializer)?;
Self::new(&value).map_err(serde::de::Error::custom)
Self::new_scoped(&value).map_err(serde::de::Error::custom)
}
}
impl FsPath {
/// Construct a path for a scope-checked operation that may target an
/// explicitly granted absolute path outside the provider root.
pub fn new_scoped(value: impl Into<String>) -> Result<Self, FsError> {
let value = value.into();
if !Path::new(&value).is_absolute() {
return Self::new(value);
}
if value.contains('\\') {
return Err(FsError::InvalidPath(value));
}
if Path::new(&value)
.components()
.any(|component| component == Component::ParentDir)
{
return Err(FsError::InvalidPath(value));
}
Ok(Self(value))
}
pub fn root() -> Self {
Self(String::new())
}
+22
View File
@@ -1438,6 +1438,28 @@ mod tests {
assert!(resolved.manifest.feature.workspace_worker_discovery.enabled);
}
#[test]
fn builtin_orchestrator_keeps_cleanup_tool_providers_enabled() {
let tmp = TempDir::new().unwrap();
let resolved = ProfileResolver::new()
.with_workspace_base(tmp.path())
.resolve(
&ProfileSelector::source_named(ProfileRegistrySource::Builtin, "orchestrator"),
ProfileResolveOptions::with_worker_name("orchestrator-worker"),
)
.unwrap();
let feature = resolved.manifest.feature;
assert!(feature.worker.enabled);
assert!(!feature.worker.direct_spawn);
assert!(feature.manage_workdir.enabled);
assert!(feature.merge_request.show);
assert!(feature.merge_request.readiness_check);
assert!(feature.merge_request.complete);
assert!(!feature.merge_request.open);
assert!(!feature.merge_request.review);
}
#[test]
fn profile_resolution_requires_runtime_worker_name() {
let tmp = TempDir::new().unwrap();
+2 -1
View File
@@ -14,6 +14,7 @@ json-schema = ["dep:schemars"]
schemars = { workspace = true, optional = true }
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
sha2.workspace = true
tokio = { workspace = true, features = ["io-util"], optional = true }
ts-rs = { version = "12.0.1", optional = true }
uuid = { workspace = true, features = ["serde"] }
uuid = { workspace = true, features = ["serde", "v7"] }
+132
View File
@@ -0,0 +1,132 @@
use std::{fmt, str::FromStr};
use serde::{Deserialize, Deserializer, Serialize, Serializer, de};
use sha2::{Digest, Sha256};
use uuid::{Uuid, Version};
/// Stable Worker identity independent of its current Runtime placement or
/// conversation Session.
///
/// Workspace authority allocates this ID for managed Workers. A standalone
/// Worker store allocates it locally when no Workspace authority is present.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct WorkerId(Uuid);
impl WorkerId {
pub fn now_v7() -> Self {
Self(Uuid::now_v7())
}
/// Converts a legacy Runtime-local numeric id into a syntactically valid
/// migration-only UUIDv7 value. New Worker allocation must use `now_v7`.
pub fn from_legacy_u64(value: u64) -> Self {
let mut bytes = [0_u8; 16];
bytes[8..].copy_from_slice(&value.to_be_bytes());
bytes[6] = 0x70;
bytes[8] = (bytes[8] & 0x3f) | 0x80;
Self(Uuid::from_bytes(bytes))
}
pub fn from_legacy_binding(workspace_id: &str, runtime_id: &str, value: u64) -> Self {
let mut hasher = Sha256::new();
hasher.update(b"yoi.workspace-worker-id.v1\0");
hasher.update(workspace_id.as_bytes());
hasher.update([0]);
hasher.update(runtime_id.as_bytes());
hasher.update([0]);
hasher.update(value.to_be_bytes());
let digest = hasher.finalize();
let mut bytes = [0_u8; 16];
bytes.copy_from_slice(&digest[..16]);
// Migrated ids sort before normally allocated UUIDv7 values while retaining
// deterministic collision-resistant payload bits.
bytes[..6].fill(0);
bytes[6] = (bytes[6] & 0x0f) | 0x70;
bytes[8] = (bytes[8] & 0x3f) | 0x80;
Self(Uuid::from_bytes(bytes))
}
pub fn parse(value: &str) -> Option<Self> {
let value = Uuid::parse_str(value).ok()?;
(value.get_version() == Some(Version::SortRand)).then_some(Self(value))
}
pub const fn as_uuid(self) -> Uuid {
self.0
}
#[must_use]
pub fn short(self) -> String {
let simple = self.0.simple().to_string();
simple[simple.len() - 12..].to_string()
}
}
impl fmt::Display for WorkerId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt(formatter)
}
}
impl FromStr for WorkerId {
type Err = WorkerIdParseError;
fn from_str(value: &str) -> Result<Self, Self::Err> {
Self::parse(value).ok_or(WorkerIdParseError)
}
}
impl Serialize for WorkerId {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_str(&self.to_string())
}
}
impl<'de> Deserialize<'de> for WorkerId {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let value = String::deserialize(deserializer)?;
Self::parse(&value).ok_or_else(|| de::Error::custom("Worker id must be a UUIDv7"))
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct WorkerIdParseError;
impl fmt::Display for WorkerIdParseError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("Worker id must be a UUIDv7")
}
}
impl std::error::Error for WorkerIdParseError {}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn worker_id_accepts_only_uuid_v7() {
let worker_id = WorkerId::now_v7();
assert_eq!(WorkerId::parse(&worker_id.to_string()), Some(worker_id));
assert!(WorkerId::parse("30").is_none());
assert!(WorkerId::parse(&Uuid::nil().to_string()).is_none());
}
#[test]
fn legacy_worker_id_mapping_is_stable() {
assert_eq!(
WorkerId::from_legacy_binding("workspace", "runtime", 42),
WorkerId::from_legacy_binding("workspace", "runtime", 42)
);
assert_ne!(
WorkerId::from_legacy_binding("workspace", "runtime", 42),
WorkerId::from_legacy_binding("workspace", "runtime", 43)
);
}
}
+189
View File
@@ -1,3 +1,4 @@
pub mod identity;
#[cfg(feature = "stream")]
pub mod stream;
pub mod subscription;
@@ -8,6 +9,8 @@ use std::path::PathBuf;
use serde::{Deserialize, Serialize};
pub use identity::{WorkerId, WorkerIdParseError};
fn default_true() -> bool {
true
}
@@ -190,6 +193,106 @@ impl WorkerEvent {
/// variants — emits an alert and inserts a `[unknown input segment]`
/// placeholder into the LLM context so neither user nor LLM is blind to
/// the dropped intent.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum PasteArtifactMediaType {
TextPlainUtf8,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum PasteArtifactAvailability {
Available,
Unavailable,
IntegrityFailed,
}
impl PasteArtifactMediaType {
pub fn as_str(self) -> &'static str {
match self {
Self::TextPlainUtf8 => "text/plain; charset=utf-8",
}
}
}
impl PasteArtifactAvailability {
pub fn as_str(self) -> &'static str {
match self {
Self::Available => "available",
Self::Unavailable => "unavailable",
Self::IntegrityFailed => "integrity_failed",
}
}
}
/// Session-owned reference to a large pasted-input artifact.
///
/// The reference contains only bounded integrity and provenance metadata. The
/// artifact body remains in session storage and is available to the model only
/// through the scoped paste-artifact tools installed by Worker.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))]
pub struct PasteArtifactRef {
pub artifact_id: String,
pub created_at_ms: u64,
pub media_type: PasteArtifactMediaType,
/// Availability observed when this immutable reference was committed.
/// Reads revalidate storage and integrity rather than trusting this field.
pub availability: PasteArtifactAvailability,
pub byte_len: u64,
pub char_count: u64,
pub line_count: u64,
pub sha256: String,
pub source_entry_id: String,
}
/// Availability recorded for an uploaded client-local file.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum UploadedFileAvailability {
Available,
Unavailable,
IntegrityFailed,
}
impl UploadedFileAvailability {
pub fn as_str(self) -> &'static str {
match self {
Self::Available => "available",
Self::Unavailable => "unavailable",
Self::IntegrityFailed => "integrity_failed",
}
}
}
/// Session-owned immutable reference to a client-local uploaded file.
///
/// Upload transports return an unbound reference. Worker fills
/// `source_entry_id` immediately before the containing user input is committed;
/// committed Session Log and public snapshot records therefore always retain
/// the durable source-entry identity without storing the file body.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))]
pub struct UploadedFileRef {
pub artifact_id: String,
pub file_name: String,
pub media_type: String,
pub created_at_ms: u64,
pub availability: UploadedFileAvailability,
pub byte_len: u64,
pub sha256: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub source_entry_id: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))]
@@ -207,6 +310,14 @@ pub enum Segment {
lines: u32,
content: String,
},
/// Internal reference produced when Worker stores a large `Paste` before
/// committing input. Clients may receive this in history/event projections;
/// the body is intentionally absent.
PasteArtifact { artifact: PasteArtifactRef },
/// Client-local file uploaded into the owning Worker session before submit.
/// The Session Log stores only this immutable reference, never file bytes or
/// the client's local path.
UploadedFile { file: UploadedFileRef },
/// `@<path>` file-system reference. Worker resolves readable files to
/// `[File: <path>]` attachments and readable normal directories to shallow
/// `[Dir: <path>]` listings; the flattened user text keeps the literal
@@ -247,6 +358,35 @@ impl Segment {
match seg {
Segment::Text { content } => out.push_str(content),
Segment::Paste { content, .. } => out.push_str(content),
Segment::PasteArtifact { artifact } => {
use std::fmt::Write as _;
let _ = write!(
out,
"[Large paste stored as artifact {}: {} bytes, {} chars, {} lines, {}, {}, created at {} ms, sha256 {}; use SearchInputArtifact and ReadInputArtifact to inspect it]",
artifact.artifact_id,
artifact.byte_len,
artifact.char_count,
artifact.line_count,
artifact.media_type.as_str(),
artifact.availability.as_str(),
artifact.created_at_ms,
artifact.sha256
);
}
Segment::UploadedFile { file } => {
use std::fmt::Write as _;
let _ = write!(
out,
"[Attached file {} stored as input artifact {}: {} bytes, {}, {}, created at {} ms, sha256 {}; use SearchInputArtifact and ReadInputArtifact for supported text content]",
file.file_name,
file.artifact_id,
file.byte_len,
file.media_type,
file.availability.as_str(),
file.created_at_ms,
file.sha256
);
}
Segment::FileRef { path } => {
out.push('@');
out.push_str(path);
@@ -1199,6 +1339,55 @@ mod tests {
}
}
#[test]
fn paste_artifact_segment_roundtrips_without_body() {
let artifact = PasteArtifactRef {
artifact_id: "019ca7c8-57b6-7f05-8edf-524147aba7b2".to_string(),
created_at_ms: 1_700_000_000_000,
media_type: PasteArtifactMediaType::TextPlainUtf8,
availability: PasteArtifactAvailability::Available,
byte_len: 65_536,
char_count: 65_530,
line_count: 200,
sha256: "a".repeat(64),
source_entry_id: "entry-1".to_string(),
};
let segment = Segment::PasteArtifact {
artifact: artifact.clone(),
};
let json = serde_json::to_string(&segment).unwrap();
assert!(!json.contains("pasted body"));
assert_eq!(serde_json::from_str::<Segment>(&json).unwrap(), segment);
let projected = Segment::flatten_to_text(&[segment]);
assert!(projected.contains(&artifact.artifact_id));
assert!(projected.contains("SearchInputArtifact"));
assert!(projected.contains("ReadInputArtifact"));
assert!(!projected.contains("pasted body"));
}
#[test]
fn uploaded_file_segment_roundtrips_without_path_or_body() {
let file = UploadedFileRef {
artifact_id: "019ca7c8-57b6-7f05-8edf-524147aba7b3".to_string(),
file_name: "notes.md".to_string(),
media_type: "text/markdown".to_string(),
created_at_ms: 1_700_000_000_001,
availability: UploadedFileAvailability::Available,
byte_len: 128,
sha256: "b".repeat(64),
source_entry_id: Some("entry-2".to_string()),
};
let segment = Segment::UploadedFile { file: file.clone() };
let json = serde_json::to_string(&segment).unwrap();
assert!(!json.contains("/home/user/private"));
assert!(!json.contains("file body"));
assert_eq!(serde_json::from_str::<Segment>(&json).unwrap(), segment);
let projected = Segment::flatten_to_text(&[segment]);
assert!(projected.contains("notes.md"));
assert!(projected.contains(&file.artifact_id));
assert!(projected.contains("ReadInputArtifact"));
}
#[test]
fn method_run_flow_segment_roundtrip() {
let method = Method::Run {
+106 -4
View File
@@ -170,6 +170,23 @@ fn validate_identifier(
Ok(())
}
fn validate_repository_key(value: &str) -> Result<(), SubscriptionValidationError> {
let bytes = value.as_bytes();
if bytes.is_empty()
|| bytes.len() > 64
|| bytes.first() == Some(&b'-')
|| bytes.last() == Some(&b'-')
|| !bytes
.iter()
.all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || *byte == b'-')
{
return Err(SubscriptionValidationError::InvalidIdentifier {
field: "repository_key",
});
}
Ok(())
}
fn validate_rejection_message(message: &str) -> Result<(), SubscriptionValidationError> {
if message.is_empty() {
return Err(SubscriptionValidationError::EmptyRejectionMessage);
@@ -567,7 +584,12 @@ pub struct SubscriptionWorker {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub profile: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[cfg_attr(feature = "typescript", ts(skip))]
pub repository_id: Option<String>,
/// Workspace-facing Repository key. Runtime producers leave this unset and
/// Workspace Server projections replace `repository_id` with this field.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub repository_key: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub working_directory_id: Option<SubscriptionWorkdirId>,
}
@@ -584,6 +606,14 @@ impl SubscriptionWorker {
if let Some(repository_id) = &self.repository_id {
validate_identifier("repository_id", repository_id, MAX_RESOURCE_ID_BYTES)?;
}
if let Some(repository_key) = &self.repository_key {
validate_repository_key(repository_key)?;
}
if self.repository_id.is_some() && self.repository_key.is_some() {
return Err(SubscriptionValidationError::InvalidIdentifier {
field: "repository_authority",
});
}
if let Some(working_directory_id) = &self.working_directory_id {
working_directory_id.validate()?;
}
@@ -595,7 +625,13 @@ impl SubscriptionWorker {
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
pub struct SubscriptionWorkdir {
pub working_directory_id: SubscriptionWorkdirId,
pub repository_id: String,
/// Runtime-internal Repository id. Workspace-facing TypeScript contracts
/// omit this field and require `repository_key` from the Server projection.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[cfg_attr(feature = "typescript", ts(skip))]
pub repository_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub repository_key: Option<String>,
pub state: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub primary_worker_id: Option<SubscriptionWorkerId>,
@@ -604,7 +640,41 @@ pub struct SubscriptionWorkdir {
impl SubscriptionWorkdir {
pub fn validate(&self) -> Result<(), SubscriptionValidationError> {
self.working_directory_id.validate()?;
validate_identifier("repository_id", &self.repository_id, MAX_RESOURCE_ID_BYTES)?;
match (&self.repository_id, &self.repository_key) {
(Some(repository_id), None) => {
validate_identifier("repository_id", repository_id, MAX_RESOURCE_ID_BYTES)?;
}
(None, Some(repository_key)) => validate_repository_key(repository_key)?,
_ => {
return Err(SubscriptionValidationError::InvalidIdentifier {
field: "repository_authority",
});
}
}
validate_identifier("workdir_state", &self.state, MAX_RESOURCE_ID_BYTES)?;
if let Some(worker_id) = &self.primary_worker_id {
worker_id.validate()?;
}
Ok(())
}
}
/// Workspace-facing Workdir summary. Backend-generated Repository UUIDs never
/// enter this DTO; Workspace Server must resolve the required Repository key.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
pub struct WorkspaceSubscriptionWorkdir {
pub working_directory_id: SubscriptionWorkdirId,
pub repository_key: String,
pub state: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub primary_worker_id: Option<SubscriptionWorkerId>,
}
impl WorkspaceSubscriptionWorkdir {
pub fn validate(&self) -> Result<(), SubscriptionValidationError> {
self.working_directory_id.validate()?;
validate_repository_key(&self.repository_key)?;
validate_identifier("workdir_state", &self.state, MAX_RESOURCE_ID_BYTES)?;
if let Some(worker_id) = &self.primary_worker_id {
worker_id.validate()?;
@@ -625,7 +695,7 @@ pub enum SubscriptionSnapshot {
events: Vec<WorkerProtocolEvent>,
},
WorkspaceWorkdirs {
workdirs: Vec<SubscriptionWorkdir>,
workdirs: Vec<WorkspaceSubscriptionWorkdir>,
},
}
@@ -693,7 +763,7 @@ pub enum SubscriptionEventPayload {
event: WorkerProtocolEvent,
},
WorkdirUpserted {
workdir: SubscriptionWorkdir,
workdir: WorkspaceSubscriptionWorkdir,
},
WorkdirRemoved {
working_directory_id: SubscriptionWorkdirId,
@@ -811,10 +881,42 @@ mod tests {
display_name: Some(format!("Worker {value}")),
profile: Some("builtin:coder".to_string()),
repository_id: None,
repository_key: None,
working_directory_id: None,
}
}
#[test]
fn runtime_and_workspace_repository_identity_projections_do_not_alias() {
let mut runtime_worker = worker("worker-1");
runtime_worker.repository_id = Some("01890f47-3c22-7cc0-98c4-dc0c0c07398f".to_string());
runtime_worker.validate().unwrap();
let runtime_json = serde_json::to_value(&runtime_worker).unwrap();
assert_eq!(
runtime_json["repository_id"],
"01890f47-3c22-7cc0-98c4-dc0c0c07398f"
);
assert!(runtime_json.get("repository_key").is_none());
let mut workspace_worker = worker("worker-1");
workspace_worker.repository_key = Some("main".to_string());
workspace_worker.validate().unwrap();
let workspace_json = serde_json::to_value(&workspace_worker).unwrap();
assert_eq!(workspace_json["repository_key"], "main");
assert!(workspace_json.get("repository_id").is_none());
let workspace_workdir = WorkspaceSubscriptionWorkdir {
working_directory_id: SubscriptionWorkdirId::new("workdir-1").unwrap(),
repository_key: "main".to_string(),
state: "active".to_string(),
primary_worker_id: Some(worker_id("worker-1")),
};
workspace_workdir.validate().unwrap();
let workdir_json = serde_json::to_value(&workspace_workdir).unwrap();
assert_eq!(workdir_json["repository_key"], "main");
assert!(workdir_json.get("repository_id").is_none());
}
#[test]
fn subscribe_frame_has_stable_versioned_json_shape() {
let frame = SubscriptionFrame::new(SubscriptionFramePayload::Request(
+22 -7
View File
@@ -7,17 +7,19 @@ use crate::{
CommandStreamSlice, CompactionLifecycle, CompactionLifecycleState, CompletionEntry,
CompletionKind, ErrorCode, Event, Greeting, InFlightBlock, InFlightSnapshot,
InFlightToolCallState, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnapshot,
InvokeKind, MemoryWorkerEvent, Method, Permission, RewindSummary, RewindTarget, RewindTargetId,
RunResult, ScopeRule, Segment, SessionContentPart, SessionEntryProvenance, SessionMessageRole,
InvokeKind, MemoryWorkerEvent, Method, PasteArtifactAvailability, PasteArtifactMediaType,
PasteArtifactRef, Permission, RewindSummary, RewindTarget, RewindTargetId, RunResult,
ScopeRule, Segment, SessionContentPart, SessionEntryProvenance, SessionMessageRole,
SessionSnapshot, SessionSnapshotEntry, SessionSnapshotEntryData, SessionToolAttachment,
ToolResultDisposition, TurnResult, WorkerEvent, WorkerStatus,
ToolResultDisposition, TurnResult, UploadedFileAvailability, UploadedFileRef, WorkerEvent,
WorkerStatus,
subscription::{
EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame,
SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest,
SubscriptionRequestId, SubscriptionResponse, SubscriptionSnapshot,
SubscriptionTerminationCode, SubscriptionWorkdir, SubscriptionWorkdirId,
SubscriptionWorker, SubscriptionWorkerId, SubscriptionWorkerIds,
SubscriptionWorkerProtocolMethod, SubscriptionWorkerState,
SubscriptionTerminationCode, SubscriptionWorkdirId, SubscriptionWorker,
SubscriptionWorkerId, SubscriptionWorkerIds, SubscriptionWorkerProtocolMethod,
SubscriptionWorkerState, WorkspaceSubscriptionWorkdir,
},
};
@@ -58,6 +60,8 @@ pub fn generated_protocol_types() -> String {
push_decl::<CommandEvent>(&cfg, &mut output);
push_decl::<CompactionLifecycleState>(&cfg, &mut output);
push_decl::<CompactionLifecycle>(&cfg, &mut output);
push_decl::<UploadedFileAvailability>(&cfg, &mut output);
push_decl::<UploadedFileRef>(&cfg, &mut output);
push_decl::<ScopeRule>(&cfg, &mut output);
push_decl::<CompletionEntry>(&cfg, &mut output);
push_decl::<RewindTargetId>(&cfg, &mut output);
@@ -78,6 +82,9 @@ pub fn generated_protocol_types() -> String {
push_decl::<Greeting>(&cfg, &mut output);
push_decl::<Alert>(&cfg, &mut output);
push_decl::<MemoryWorkerEvent>(&cfg, &mut output);
push_decl::<PasteArtifactMediaType>(&cfg, &mut output);
push_decl::<PasteArtifactAvailability>(&cfg, &mut output);
push_decl::<PasteArtifactRef>(&cfg, &mut output);
push_decl::<Segment>(&cfg, &mut output);
push_decl::<WorkerEvent>(&cfg, &mut output);
push_decl::<SubscriptionRequestId>(&cfg, &mut output);
@@ -88,7 +95,7 @@ pub fn generated_protocol_types() -> String {
push_decl::<SubscriptionWorkerState>(&cfg, &mut output);
push_decl::<EventSubscriptionSelector>(&cfg, &mut output);
push_decl::<SubscriptionWorker>(&cfg, &mut output);
push_decl::<SubscriptionWorkdir>(&cfg, &mut output);
push_decl::<WorkspaceSubscriptionWorkdir>(&cfg, &mut output);
push_decl::<SubscriptionSnapshot>(&cfg, &mut output);
push_decl::<SubscriptionEventPayload>(&cfg, &mut output);
push_decl::<SubscriptionRejectionCode>(&cfg, &mut output);
@@ -132,6 +139,14 @@ fn export_decl(decl: &str) -> String {
mod tests {
use super::*;
#[test]
fn workspace_typescript_omits_runtime_repository_ids() {
let generated = generated_protocol_types();
assert!(!generated.contains("repository_id?:"), "{generated}");
assert!(!generated.contains("repository_id:"), "{generated}");
assert!(generated.contains("repository_key"), "{generated}");
}
#[test]
fn generated_protocol_types_are_current() {
let expected = generated_protocol_types();
+5
View File
@@ -8,12 +8,17 @@ license.workspace = true
[dependencies]
base64.workspace = true
agen = { workspace = true }
fs4.workspace = true
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
sha2.workspace = true
uuid = { workspace = true, features = ["v7", "serde"] }
thiserror = { workspace = true }
protocol = { workspace = true }
tracing.workspace = true
unicode-normalization = "0.1.25"
unicode-properties = { version = "0.1.4", features = ["general-category"] }
unicode-security = "0.1.2"
[dev-dependencies]
async-trait = { workspace = true }
+639 -1
View File
@@ -16,9 +16,18 @@
//! enumerable by the picker.
use crate::event_trace::TraceEntry;
use crate::paste_artifact::{read_from_dir, write_to_dir};
use crate::segment_log::LogEntry;
use crate::store::{Store, StoreError};
use crate::{SegmentId, SessionId};
use crate::uploaded_file::{
bind_uploaded_file, clear_uploaded_file_binding, copy_committed_uploaded_files,
delete_uncommitted_uploaded_files, delete_uploaded_file, list_uploaded_file_refs,
read_uploaded_file, read_uploaded_file_by_id, write_uploaded_file,
};
use crate::{
PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits, UploadedFileUploadContext,
};
use protocol::{PasteArtifactRef, UploadedFileRef};
use std::fs;
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
@@ -109,6 +118,50 @@ impl FsStore {
.join(format!("{segment_id}.trace.jsonl"))
}
fn paste_artifact_dir(&self, session_id: SessionId) -> PathBuf {
self.session_dir(session_id).join("artifacts").join("paste")
}
fn uploaded_file_is_referenced(
&self,
session_id: SessionId,
artifact_id: &str,
) -> Result<bool, StoreError> {
fn segments_contain(segments: &[protocol::Segment], artifact_id: &str) -> bool {
segments.iter().any(|segment| {
matches!(
segment,
protocol::Segment::UploadedFile { file }
if file.artifact_id == artifact_id
)
})
}
for segment_id in self.list_segments(session_id)? {
for entry in self.read_all(session_id, segment_id)? {
let referenced = match entry {
LogEntry::AnnotatedUserInput { segments, .. } => {
segments_contain(&segments, artifact_id)
}
LogEntry::InputSegmentsCheckpoint { user_segments, .. } => user_segments
.iter()
.any(|segments| segments_contain(segments, artifact_id)),
_ => false,
};
if referenced {
return Ok(true);
}
}
}
Ok(false)
}
#[cfg(test)]
fn paste_artifact_path(&self, session_id: SessionId, artifact_id: &str) -> PathBuf {
self.paste_artifact_dir(session_id)
.join(format!("{artifact_id}.json"))
}
fn append_line(&self, path: &Path, line: &str) -> Result<(), StoreError> {
let _guard = self
.append_lock
@@ -350,6 +403,171 @@ impl Store for FsStore {
Ok(complete.lines().filter(|l| !l.trim().is_empty()).count())
}
fn write_paste_artifact(
&self,
session_id: SessionId,
source_entry_id: &str,
content: &str,
limits: PasteArtifactLimits,
) -> Result<PasteArtifactRef, StoreError> {
let _guard = self
.append_lock
.lock()
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
write_to_dir(
&self.paste_artifact_dir(session_id),
source_entry_id,
content,
limits,
)
}
fn read_paste_artifact(
&self,
session_id: SessionId,
artifact_id: &str,
) -> Result<(PasteArtifactRef, String), StoreError> {
read_from_dir(&self.paste_artifact_dir(session_id), artifact_id)
}
fn write_uploaded_file(
&self,
session_id: SessionId,
file_name: &str,
media_type: &str,
content: &[u8],
limits: UploadedFileLimits,
) -> Result<UploadedFileRef, StoreError> {
let _guard = self
.append_lock
.lock()
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
write_uploaded_file(
&self.paste_artifact_dir(session_id),
file_name,
media_type,
content,
None,
limits,
)
}
fn write_uploaded_file_with_context(
&self,
session_id: SessionId,
file_name: &str,
media_type: &str,
content: &[u8],
context: &UploadedFileUploadContext,
limits: UploadedFileLimits,
) -> Result<UploadedFileRef, StoreError> {
let _guard = self
.append_lock
.lock()
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
write_uploaded_file(
&self.paste_artifact_dir(session_id),
file_name,
media_type,
content,
Some(context),
limits,
)
}
fn read_uploaded_file(
&self,
session_id: SessionId,
reference: &UploadedFileRef,
) -> Result<Vec<u8>, StoreError> {
read_uploaded_file(&self.paste_artifact_dir(session_id), reference)
}
fn read_uploaded_file_by_id(
&self,
session_id: SessionId,
artifact_id: &str,
) -> Result<(UploadedFileRef, Vec<u8>), StoreError> {
read_uploaded_file_by_id(&self.paste_artifact_dir(session_id), artifact_id)
}
fn bind_uploaded_file(
&self,
session_id: SessionId,
reference: &UploadedFileRef,
source_entry_id: &str,
) -> Result<UploadedFileRef, StoreError> {
let _guard = self
.append_lock
.lock()
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
let dir = self.paste_artifact_dir(session_id);
match bind_uploaded_file(&dir, reference, source_entry_id) {
Err(StoreError::ArtifactAlreadyCommitted) => {
let (stored, _) = read_uploaded_file_by_id(&dir, &reference.artifact_id)?;
let previous_source = stored
.source_entry_id
.ok_or(StoreError::ArtifactIntegrityMismatch)?;
if self.uploaded_file_is_referenced(session_id, &reference.artifact_id)? {
return Err(StoreError::ArtifactAlreadyCommitted);
}
clear_uploaded_file_binding(&dir, &reference.artifact_id, &previous_source)?;
bind_uploaded_file(&dir, reference, source_entry_id)
}
result => result,
}
}
fn delete_uploaded_file(
&self,
session_id: SessionId,
artifact_id: &str,
) -> Result<bool, StoreError> {
let _guard = self
.append_lock
.lock()
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
delete_uploaded_file(&self.paste_artifact_dir(session_id), artifact_id)
}
fn delete_uncommitted_uploaded_files(&self, session_id: SessionId) -> Result<u64, StoreError> {
let _guard = self
.append_lock
.lock()
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
let dir = self.paste_artifact_dir(session_id);
let mut removed = delete_uncommitted_uploaded_files(&dir)?;
for reference in list_uploaded_file_refs(&dir)? {
let Some(source_entry_id) = reference.source_entry_id.as_deref() else {
continue;
};
if !self.uploaded_file_is_referenced(session_id, &reference.artifact_id)? {
clear_uploaded_file_binding(&dir, &reference.artifact_id, source_entry_id)?;
if delete_uploaded_file(&dir, &reference.artifact_id)? {
removed = removed
.checked_add(1)
.ok_or(StoreError::ArtifactQuotaExceeded)?;
}
}
}
Ok(removed)
}
fn copy_committed_uploaded_files(
&self,
source_session_id: SessionId,
target_session_id: SessionId,
) -> Result<u64, StoreError> {
let _guard = self
.append_lock
.lock()
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
copy_committed_uploaded_files(
&self.paste_artifact_dir(source_session_id),
&self.paste_artifact_dir(target_session_id),
)
}
fn append_trace(
&self,
session_id: SessionId,
@@ -398,4 +616,424 @@ mod tests {
store.create_segment(session_id, segment_id, &[]).unwrap();
assert!(store.session_modified_at(session_id).unwrap().is_some());
}
#[test]
fn paste_artifacts_are_atomic_integrity_checked_and_session_scoped() {
let tmp = tempfile::TempDir::new().unwrap();
let store = FsStore::new(tmp.path()).unwrap();
let owner = new_session_id();
let other = new_session_id();
let content = "αβγ\nsecond line\n";
let reference = store
.write_paste_artifact(owner, "entry-1", content, PasteArtifactLimits::default())
.unwrap();
assert_eq!(reference.byte_len, content.len() as u64);
assert!(reference.created_at_ms > 0);
assert_eq!(
reference.media_type,
protocol::PasteArtifactMediaType::TextPlainUtf8
);
assert_eq!(
reference.availability,
protocol::PasteArtifactAvailability::Available
);
assert_eq!(reference.char_count, content.chars().count() as u64);
assert_eq!(reference.source_entry_id, "entry-1");
assert_eq!(
store
.read_paste_artifact(owner, &reference.artifact_id)
.unwrap()
.1,
content
);
assert!(matches!(
store.read_paste_artifact(other, &reference.artifact_id),
Err(StoreError::PasteArtifactNotFound(_))
));
assert!(
self::fs::read_dir(store.paste_artifact_dir(owner))
.unwrap()
.all(|entry| !entry
.unwrap()
.file_name()
.to_string_lossy()
.ends_with(".tmp"))
);
let very_large = "z".repeat(1024 * 1024);
let very_large_ref = store
.write_paste_artifact(
owner,
"entry-2",
&very_large,
PasteArtifactLimits::default(),
)
.unwrap();
assert_eq!(
store
.read_paste_artifact(owner, &very_large_ref.artifact_id)
.unwrap()
.1,
very_large
);
}
#[test]
fn concurrent_paste_writes_atomically_enforce_aggregate_caps() {
let tmp = tempfile::TempDir::new().unwrap();
let session_id = new_session_id();
let barrier = std::sync::Arc::new(std::sync::Barrier::new(3));
let limits = PasteArtifactLimits {
max_artifact_bytes: 4,
max_session_bytes: 8,
max_session_artifacts: 1,
};
let mut handles = Vec::new();
for entry_id in ["entry-1", "entry-2"] {
let root = tmp.path().to_path_buf();
let barrier = barrier.clone();
handles.push(std::thread::spawn(move || {
let store = FsStore::new(root).unwrap();
barrier.wait();
store.write_paste_artifact(session_id, entry_id, "1234", limits)
}));
}
barrier.wait();
let results = handles
.into_iter()
.map(|handle| handle.join().unwrap())
.collect::<Vec<_>>();
assert_eq!(results.iter().filter(|result| result.is_ok()).count(), 1);
assert_eq!(
results
.iter()
.filter(|result| matches!(result, Err(StoreError::PasteArtifactLimit(_))))
.count(),
1
);
assert_eq!(
std::fs::read_dir(
FsStore::new(tmp.path())
.unwrap()
.paste_artifact_dir(session_id)
)
.unwrap()
.filter_map(Result::ok)
.filter(
|entry| entry.path().extension().and_then(|value| value.to_str()) == Some("json")
)
.count(),
1
);
}
#[test]
fn uploaded_file_persists_trusted_upload_context_without_projecting_it() {
let tmp = tempfile::TempDir::new().unwrap();
let store = FsStore::new(tmp.path()).unwrap();
let session_id = new_session_id();
let context = UploadedFileUploadContext {
upload_id: "upload-1".into(),
principal_id: "account-1".into(),
workspace_id: "workspace-1".into(),
runtime_id: "runtime-1".into(),
worker_id: "worker-1".into(),
};
let reference = store
.write_uploaded_file_with_context(
session_id,
"notes.txt",
"text/plain",
b"hello",
&context,
UploadedFileLimits::default(),
)
.unwrap();
let raw = fs::read_to_string(
store
.paste_artifact_dir(session_id)
.join(format!("{}.file.json", reference.artifact_id)),
)
.unwrap();
assert!(raw.contains("account-1"));
assert!(raw.contains("workspace-1"));
assert!(raw.contains("runtime-1"));
assert!(raw.contains("worker-1"));
assert!(
!serde_json::to_string(&reference)
.unwrap()
.contains("account-1")
);
let replay = store
.write_uploaded_file_with_context(
session_id,
"notes.txt",
"text/plain",
b"hello",
&context,
UploadedFileLimits::default(),
)
.unwrap();
assert_eq!(replay.artifact_id, reference.artifact_id);
assert!(matches!(
store.write_uploaded_file_with_context(
session_id,
"renamed.txt",
"text/plain",
b"hello",
&context,
UploadedFileLimits::default(),
),
Err(StoreError::InvalidUploadedFileName)
));
}
#[test]
fn uploaded_file_exact_replay_succeeds_at_session_count_limit() {
let tmp = tempfile::TempDir::new().unwrap();
let store = FsStore::new(tmp.path()).unwrap();
let session_id = new_session_id();
let limits = UploadedFileLimits {
max_file_bytes: 1,
max_session_bytes: crate::DEFAULT_MAX_SESSION_UPLOADED_FILES,
};
let mut first = None;
for index in 0..crate::DEFAULT_MAX_SESSION_UPLOADED_FILES {
let reference = store
.write_uploaded_file(
session_id,
&format!("file-{index}.txt"),
"text/plain",
b"x",
limits,
)
.unwrap();
first.get_or_insert(reference);
}
let replay = store
.write_uploaded_file(session_id, "file-0.txt", "text/plain", b"x", limits)
.unwrap();
assert_eq!(replay.artifact_id, first.unwrap().artifact_id);
assert!(matches!(
store.write_uploaded_file(session_id, "overflow.txt", "text/plain", b"x", limits),
Err(StoreError::ArtifactQuotaExceeded)
));
}
#[test]
fn uploaded_files_are_session_scoped_integrity_checked_and_removable() {
let tmp = tempfile::TempDir::new().unwrap();
let store = FsStore::new(tmp.path()).unwrap();
let owner = new_session_id();
let other = new_session_id();
let limits = UploadedFileLimits {
max_file_bytes: 16,
max_session_bytes: 16,
};
let reference = store
.write_uploaded_file(owner, "notes.txt", "text/plain", b"hello", limits)
.unwrap();
assert_eq!(reference.file_name, "notes.txt");
assert_eq!(reference.media_type, "text/plain");
assert_eq!(reference.byte_len, 5);
assert_eq!(reference.source_entry_id, None);
assert_eq!(
store.read_uploaded_file(owner, &reference).unwrap(),
b"hello"
);
assert!(store.read_uploaded_file(other, &reference).is_err());
let mut forged = reference.clone();
forged.file_name = "other.txt".to_string();
assert!(matches!(
store.read_uploaded_file(owner, &forged),
Err(StoreError::ArtifactIntegrityMismatch)
));
assert!(
store
.delete_uploaded_file(owner, &reference.artifact_id)
.unwrap()
);
assert!(
!store
.delete_uploaded_file(owner, &reference.artifact_id)
.unwrap()
);
assert!(store.read_uploaded_file(owner, &reference).is_err());
}
#[test]
fn uploaded_file_validation_and_shared_quota_fail_closed() {
let tmp = tempfile::TempDir::new().unwrap();
let store = FsStore::new(tmp.path()).unwrap();
let session_id = new_session_id();
let limits = UploadedFileLimits {
max_file_bytes: 8,
max_session_bytes: 8,
};
assert!(matches!(
store.write_uploaded_file(session_id, "../secret", "text/plain", b"x", limits),
Err(StoreError::InvalidUploadedFileName)
));
assert!(matches!(
store.write_uploaded_file(session_id, "notes.txt", "not a type", b"x", limits),
Err(StoreError::InvalidUploadedFileMediaType)
));
assert!(matches!(
store.write_uploaded_file(
session_id,
"safe\u{202e}txt.exe",
"text/plain",
b"x",
limits
),
Err(StoreError::InvalidUploadedFileName)
));
assert!(matches!(
store.write_uploaded_file(session_id, "image.png", "image/png", b"not a png", limits),
Err(StoreError::ArtifactIntegrityMismatch)
));
let pending = store
.write_uploaded_file(session_id, "Readme.txt", "text/plain", b"x", limits)
.unwrap();
let replay = store
.write_uploaded_file(session_id, "Readme.txt", "text/plain", b"x", limits)
.unwrap();
assert_eq!(replay.artifact_id, pending.artifact_id);
assert!(matches!(
store.write_uploaded_file(session_id, "README.txt", "text/plain", b"changed", limits),
Err(StoreError::InvalidUploadedFileName)
));
assert!(matches!(
store.write_uploaded_file(session_id, "README.txt", "text/plain", b"y", limits),
Err(StoreError::InvalidUploadedFileName)
));
store
.bind_uploaded_file(session_id, &pending, "entry-from-failed-submit")
.unwrap();
let bound = store
.bind_uploaded_file(session_id, &pending, "entry-upload")
.unwrap();
store
.create_segment(
session_id,
new_segment_id(),
&[LogEntry::InputSegmentsCheckpoint {
ts: 1,
user_segments: vec![vec![protocol::Segment::UploadedFile {
file: bound.clone(),
}]],
}],
)
.unwrap();
let other = store
.write_uploaded_file(session_id, "other.txt", "text/plain", b"z", limits)
.unwrap();
let stale = store
.write_uploaded_file(session_id, "stale.txt", "text/plain", b"s", limits)
.unwrap();
store
.bind_uploaded_file(session_id, &stale, "entry-never-committed")
.unwrap();
assert_eq!(
store.delete_uncommitted_uploaded_files(session_id).unwrap(),
2
);
assert!(store.read_uploaded_file(session_id, &other).is_err());
assert!(store.read_uploaded_file(session_id, &stale).is_err());
assert_eq!(store.read_uploaded_file(session_id, &bound).unwrap(), b"x");
let fork_session_id = new_session_id();
assert_eq!(
store
.copy_committed_uploaded_files(session_id, fork_session_id)
.unwrap(),
1
);
assert_eq!(
store.read_uploaded_file(fork_session_id, &bound).unwrap(),
b"x"
);
store
.write_paste_artifact(
session_id,
"entry-1",
"1234",
PasteArtifactLimits {
max_artifact_bytes: 8,
max_session_bytes: 8,
max_session_artifacts: 4,
},
)
.unwrap();
assert!(matches!(
store.write_uploaded_file(session_id, "notes.txt", "text/plain", b"56789", limits),
Err(StoreError::ArtifactQuotaExceeded)
));
}
#[test]
fn uploaded_file_names_reject_format_mixed_script_and_confusable_forms() {
let tmp = tempfile::TempDir::new().unwrap();
let store = FsStore::new(tmp.path()).unwrap();
let session_id = new_session_id();
let limits = UploadedFileLimits::default();
for file_name in [
"safe\u{00ad}name.txt",
"safe\u{061c}name.txt",
"safe\u{180e}name.txt",
"safe\u{e0001}name.txt",
"p\u{0430}ypal.txt",
"report.\u{03c1}df",
"\u{0440}\u{0430}\u{0443}\u{0440}\u{0430}\u{04cf}.txt",
"\u{ff26}\u{ff49}\u{ff4c}\u{ff45}.txt",
"re\u{0301}sume\u{0301}.txt",
] {
assert!(matches!(
store.write_uploaded_file(session_id, file_name, "text/plain", b"safe", limits),
Err(StoreError::InvalidUploadedFileName)
));
}
for file_name in ["notes.txt", "résumé.txt", "日本語.txt", "📎.txt"] {
store
.write_uploaded_file(session_id, file_name, "text/plain", b"safe", limits)
.unwrap();
}
}
#[test]
fn paste_artifact_limits_and_corruption_fail_closed() {
let tmp = tempfile::TempDir::new().unwrap();
let store = FsStore::new(tmp.path()).unwrap();
let session_id = new_session_id();
let limits = PasteArtifactLimits {
max_artifact_bytes: 5,
max_session_bytes: 8,
max_session_artifacts: 2,
};
let first = store
.write_paste_artifact(session_id, "entry-1", "1234", limits)
.unwrap();
assert!(matches!(
store.write_paste_artifact(session_id, "entry-2", "56789", limits),
Err(StoreError::PasteArtifactLimit(_))
));
assert!(matches!(
store.write_paste_artifact(session_id, "entry-2", "5678", limits),
Ok(_)
));
std::fs::write(
store.paste_artifact_path(session_id, &first.artifact_id),
b"{}",
)
.unwrap();
assert!(matches!(
store.read_paste_artifact(session_id, &first.artifact_id),
Err(StoreError::Serde(_)) | Err(StoreError::PasteArtifactIntegrity(_))
));
}
}
+9
View File
@@ -27,6 +27,7 @@
//! system_prompt: None,
//! config: &config,
//! history: Vec::new(),
//! user_segments: Vec::new(),
//! })?;
//! ```
@@ -35,11 +36,13 @@ pub mod fs_store;
pub mod history;
mod legacy_session_log;
pub mod logged_item;
mod paste_artifact;
pub mod public_snapshot;
pub mod segment;
pub mod segment_log;
pub mod store;
pub mod system_item;
pub mod uploaded_file;
pub mod worker_metadata;
pub mod worker_session_store;
@@ -53,6 +56,7 @@ pub use history::{
LoggedWorkerSubject,
};
pub use logged_item::{LoggedContentPart, LoggedItem, LoggedRole, from_logged, to_logged};
pub use paste_artifact::PasteArtifactLimits;
pub use segment::{
SegmentStartState, append_entry, append_system_item, classify_logged_history_entry,
create_compacted_segment, create_segment, create_segment_with_ids, ensure_head_or_fork, fork,
@@ -64,6 +68,11 @@ pub use store::{Store, StoreError};
pub use system_item::{
PromptRenderProvenance, SystemItem, SystemReminder, SystemReminderSource, render_worker_event,
};
pub use uploaded_file::{
DEFAULT_MAX_FILES_PER_SUBMISSION, DEFAULT_MAX_SESSION_ARTIFACT_BYTES,
DEFAULT_MAX_SESSION_UPLOADED_FILES, DEFAULT_MAX_UPLOADED_FILE_BYTES, UploadedFileLimits,
UploadedFileUploadContext,
};
pub use worker_metadata::{
CombinedStore, FsWorkerStore, WorkerActiveSegmentRef, WorkerAggregateStore, WorkerMetadata,
WorkerMetadataStore, WorkerPeer, WorkerReclaimedChild, WorkerSpawnedChild,
+205
View File
@@ -0,0 +1,205 @@
//! Session-owned storage for large pasted-input artifacts.
use std::fs;
use std::io::Write as _;
use std::path::Path;
use std::time::{SystemTime, UNIX_EPOCH};
use fs4::fs_std::FileExt;
use protocol::{PasteArtifactAvailability, PasteArtifactMediaType, PasteArtifactRef};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use crate::StoreError;
/// Bounded storage policy applied before a large paste becomes durable input.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PasteArtifactLimits {
pub max_artifact_bytes: u64,
pub max_session_bytes: u64,
pub max_session_artifacts: u64,
}
impl Default for PasteArtifactLimits {
fn default() -> Self {
Self {
max_artifact_bytes: 8 * 1024 * 1024,
max_session_bytes: 64 * 1024 * 1024,
max_session_artifacts: 1_024,
}
}
}
/// Integrity-bearing on-disk record. The body and metadata are committed in one
/// atomic file replacement so readers never observe a half-written artifact.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct StoredPasteArtifact {
pub reference: PasteArtifactRef,
pub content: String,
}
pub(crate) fn stored_paste_usage(artifact_dir: &Path) -> Result<(u64, u64), StoreError> {
if !artifact_dir.exists() {
return Ok((0, 0));
}
let mut aggregate = 0_u64;
let mut artifact_count = 0_u64;
for entry in fs::read_dir(artifact_dir)? {
let path = entry?.path();
let Some(name) = path.file_name().and_then(|value| value.to_str()) else {
continue;
};
if !name.ends_with(".json") || name.ends_with(".file.json") {
continue;
}
let stored: StoredPasteArtifact = serde_json::from_slice(&fs::read(&path)?)?;
verify(&stored, &stored.reference.artifact_id)?;
artifact_count = artifact_count.checked_add(1).ok_or_else(|| {
StoreError::PasteArtifactLimit("session artifact count overflow".to_string())
})?;
aggregate = aggregate
.checked_add(stored.reference.byte_len)
.ok_or_else(|| {
StoreError::PasteArtifactLimit("session aggregate size overflow".to_string())
})?;
}
Ok((aggregate, artifact_count))
}
pub(crate) fn write_to_dir(
artifact_dir: &Path,
source_entry_id: &str,
content: &str,
limits: PasteArtifactLimits,
) -> Result<PasteArtifactRef, StoreError> {
let byte_len = content.len() as u64;
if byte_len > limits.max_artifact_bytes {
return Err(StoreError::PasteArtifactLimit(format!(
"artifact has {byte_len} bytes; maximum is {}",
limits.max_artifact_bytes
)));
}
fs::create_dir_all(artifact_dir)?;
let aggregate_lock = fs::OpenOptions::new()
.create(true)
.read(true)
.write(true)
.open(artifact_dir.join(".aggregate.lock"))?;
FileExt::lock_exclusive(&aggregate_lock)?;
let (paste_bytes, artifact_count) = stored_paste_usage(artifact_dir)?;
let (uploaded_bytes, uploaded_count) =
crate::uploaded_file::stored_uploaded_file_usage(artifact_dir)?;
let aggregate = paste_bytes.checked_add(uploaded_bytes).ok_or_else(|| {
StoreError::PasteArtifactLimit("session aggregate size overflow".to_string())
})?;
let artifact_count = artifact_count.checked_add(uploaded_count).ok_or_else(|| {
StoreError::PasteArtifactLimit("session artifact count overflow".to_string())
})?;
let projected = aggregate.checked_add(byte_len).ok_or_else(|| {
StoreError::PasteArtifactLimit("session aggregate size overflow".to_string())
})?;
if projected > limits.max_session_bytes {
return Err(StoreError::PasteArtifactLimit(format!(
"session artifacts would use {projected} bytes; maximum is {}",
limits.max_session_bytes
)));
}
if artifact_count >= limits.max_session_artifacts {
return Err(StoreError::PasteArtifactLimit(format!(
"session already has {artifact_count} artifacts; maximum is {}",
limits.max_session_artifacts
)));
}
let artifact_id = uuid::Uuid::now_v7().to_string();
let created_at_ms = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_err(|error| StoreError::PasteArtifactIntegrity(error.to_string()))?
.as_millis() as u64;
let reference = PasteArtifactRef {
artifact_id: artifact_id.clone(),
created_at_ms,
media_type: PasteArtifactMediaType::TextPlainUtf8,
availability: PasteArtifactAvailability::Available,
byte_len,
char_count: content.chars().count() as u64,
line_count: line_count(content),
sha256: sha256_hex(content),
source_entry_id: source_entry_id.to_string(),
};
let bytes = serde_json::to_vec(&StoredPasteArtifact {
reference: reference.clone(),
content: content.to_string(),
})?;
let target = artifact_dir.join(format!("{artifact_id}.json"));
let temporary = artifact_dir.join(format!(".{artifact_id}.tmp"));
let mut file = fs::OpenOptions::new()
.create_new(true)
.write(true)
.open(&temporary)?;
if let Err(error) = file.write_all(&bytes).and_then(|_| file.sync_all()) {
let _ = fs::remove_file(&temporary);
return Err(error.into());
}
if let Err(error) = fs::rename(&temporary, &target) {
let _ = fs::remove_file(&temporary);
return Err(error.into());
}
if let Ok(directory) = fs::File::open(artifact_dir) {
directory.sync_all()?;
}
Ok(reference)
}
pub(crate) fn read_from_dir(
artifact_dir: &Path,
artifact_id: &str,
) -> Result<(PasteArtifactRef, String), StoreError> {
let parsed = uuid::Uuid::parse_str(artifact_id)
.map_err(|_| StoreError::PasteArtifactNotFound(artifact_id.to_string()))?;
if parsed.to_string() != artifact_id {
return Err(StoreError::PasteArtifactNotFound(artifact_id.to_string()));
}
let path = artifact_dir.join(format!("{artifact_id}.json"));
let bytes = match fs::read(path) {
Ok(bytes) => bytes,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
return Err(StoreError::PasteArtifactNotFound(artifact_id.to_string()));
}
Err(error) => return Err(error.into()),
};
let stored: StoredPasteArtifact = serde_json::from_slice(&bytes)?;
verify(&stored, artifact_id)?;
Ok((stored.reference, stored.content))
}
fn verify(stored: &StoredPasteArtifact, artifact_id: &str) -> Result<(), StoreError> {
let actual_digest = sha256_hex(&stored.content);
if stored.reference.artifact_id != artifact_id
|| stored.reference.created_at_ms == 0
|| stored.reference.media_type != PasteArtifactMediaType::TextPlainUtf8
|| stored.reference.availability != PasteArtifactAvailability::Available
|| stored.reference.byte_len != stored.content.len() as u64
|| stored.reference.char_count != stored.content.chars().count() as u64
|| stored.reference.line_count != line_count(&stored.content)
|| stored.reference.sha256 != actual_digest
{
return Err(StoreError::PasteArtifactIntegrity(artifact_id.to_string()));
}
Ok(())
}
fn sha256_hex(content: &str) -> String {
Sha256::digest(content.as_bytes())
.iter()
.map(|byte| format!("{byte:02x}"))
.collect()
}
fn line_count(content: &str) -> u64 {
if content.is_empty() {
0
} else {
content.lines().count().max(1) as u64
}
}
@@ -41,6 +41,24 @@ pub fn project_session_snapshot(session_id: SessionId, log: &[LogEntry]) -> Sess
entries.clear();
extend_history(&mut entries, history, None, *ts);
}
LogEntry::InputSegmentsCheckpoint { user_segments, .. } => {
let mut segments = user_segments.iter();
for entry in &mut entries {
let is_user = matches!(
&entry.data,
SessionSnapshotEntryData::UserInput { .. }
| SessionSnapshotEntryData::Message {
role: SessionMessageRole::User,
..
}
);
if is_user && let Some(checkpoint) = segments.next() {
entry.data = SessionSnapshotEntryData::UserInput {
segments: checkpoint.clone(),
};
}
}
}
LogEntry::AnnotatedUserInput {
ts,
segments,
@@ -357,6 +375,63 @@ mod tests {
assert!(json.contains("visible"));
}
#[test]
fn compacted_checkpoint_restores_uploaded_file_segments() {
let session_id = crate::new_session_id();
let user_entry_id = LoggedSessionHistoryEntryId::new();
let file = protocol::UploadedFileRef {
artifact_id: "019ca7c8-57b6-7f05-8edf-524147aba7b3".into(),
file_name: "notes.md".into(),
media_type: "text/markdown".into(),
created_at_ms: 7,
availability: protocol::UploadedFileAvailability::Available,
byte_len: 12,
sha256: "a".repeat(64),
source_entry_id: Some(user_entry_id.0.clone()),
};
let segment = Segment::UploadedFile { file };
let log = vec![
LogEntry::AnnotatedSegmentStart {
ts: 10,
session_id,
system_prompt: None,
config: RequestConfig::default(),
history: vec![LoggedHistoryEntry {
item: LoggedItem::Message {
role: LoggedRole::User,
content: vec![LoggedContentPart::Text {
text: "[Attached file: notes.md]".into(),
}],
},
metadata: LoggedSessionHistoryMetadata {
entry_id: user_entry_id,
origin: LoggedSessionHistoryOrigin::HumanInput {
account_id: "account-1".into(),
},
derivation: None,
},
}],
forked_from: None,
compacted_from: Some(crate::SegmentOrigin {
segment_id: crate::new_segment_id(),
at_turn_index: 1,
}),
},
LogEntry::InputSegmentsCheckpoint {
ts: 10,
user_segments: vec![vec![segment.clone()]],
},
];
let snapshot = project_current_session_snapshot(&log);
assert_eq!(
snapshot.entries[0].data,
SessionSnapshotEntryData::UserInput {
segments: vec![segment]
}
);
}
#[test]
fn annotated_user_input_attaches_segments_to_first_user_role_entry_for_any_origin() {
let session_id = crate::new_session_id();
+69 -45
View File
@@ -17,6 +17,33 @@ pub struct SegmentStartState<'a> {
pub system_prompt: Option<&'a str>,
pub config: &'a RequestConfig,
pub history: Vec<LoggedHistoryEntry>,
pub user_segments: Vec<Vec<Segment>>,
}
fn seed_entries(
ts: u64,
session_id: SessionId,
state: SegmentStartState<'_>,
forked_from: Option<SegmentOrigin>,
compacted_from: Option<SegmentOrigin>,
) -> Vec<LogEntry> {
let entry = LogEntry::AnnotatedSegmentStart {
ts,
session_id,
system_prompt: state.system_prompt.map(String::from),
config: state.config.clone(),
history: state.history,
forked_from,
compacted_from,
};
let mut entries = vec![entry];
if !state.user_segments.is_empty() {
entries.push(LogEntry::InputSegmentsCheckpoint {
ts,
user_segments: state.user_segments,
});
}
entries
}
/// Create a new session + initial segment, writing the initial
@@ -42,16 +69,8 @@ pub fn create_segment_with_ids(
segment_id: SegmentId,
state: SegmentStartState<'_>,
) -> Result<(), StoreError> {
let entry = LogEntry::AnnotatedSegmentStart {
ts: segment_log::now_millis(),
session_id,
system_prompt: state.system_prompt.map(String::from),
config: state.config.clone(),
history: state.history.to_vec(),
forked_from: None,
compacted_from: None,
};
store.append(session_id, segment_id, &entry)
let entries = seed_entries(segment_log::now_millis(), session_id, state, None, None);
store.create_segment(session_id, segment_id, &entries)
}
/// Create a compacted segment from an existing one. Inherits the source's
@@ -68,19 +87,17 @@ pub fn create_compacted_segment(
source_turn_count: usize,
) -> Result<SegmentId, StoreError> {
let segment_id = crate::new_segment_id();
let entry = LogEntry::AnnotatedSegmentStart {
ts: segment_log::now_millis(),
session_id: source_session_id,
system_prompt: state.system_prompt.map(String::from),
config: state.config.clone(),
history: state.history.to_vec(),
forked_from: None,
compacted_from: Some(SegmentOrigin {
let entries = seed_entries(
segment_log::now_millis(),
source_session_id,
state,
None,
Some(SegmentOrigin {
segment_id: source_segment_id,
at_turn_index: source_turn_count,
}),
};
store.append(source_session_id, segment_id, &entry)?;
);
store.create_segment(source_session_id, segment_id, &entries)?;
Ok(segment_id)
}
@@ -152,21 +169,19 @@ pub fn ensure_head_or_fork(
}
let source_segment_id = *segment_id;
let fork_id = crate::new_segment_id();
let entry = LogEntry::AnnotatedSegmentStart {
ts: segment_log::now_millis(),
let entries = seed_entries(
segment_log::now_millis(),
session_id,
system_prompt: state.system_prompt.map(String::from),
config: state.config.clone(),
history: state.history.to_vec(),
forked_from: Some(SegmentOrigin {
state,
Some(SegmentOrigin {
segment_id: source_segment_id,
at_turn_index,
}),
compacted_from: None,
};
store.create_segment(session_id, fork_id, &[entry])?;
None,
);
store.create_segment(session_id, fork_id, &entries)?;
*segment_id = fork_id;
*entries_written = 1;
*entries_written = entries.len();
Ok(())
}
@@ -420,20 +435,14 @@ pub fn save_config_changed(
/// [`fork_at`] or [`ensure_head_or_fork`] instead.
pub fn fork(
store: &impl Store,
source_session_id: SessionId,
state: SegmentStartState<'_>,
) -> Result<(SessionId, SegmentId), StoreError> {
let session_id = crate::new_session_id();
let fork_id = crate::new_segment_id();
let entry = LogEntry::AnnotatedSegmentStart {
ts: segment_log::now_millis(),
session_id,
system_prompt: state.system_prompt.map(String::from),
config: state.config.clone(),
history: state.history.to_vec(),
forked_from: None,
compacted_from: None,
};
store.create_segment(session_id, fork_id, &[entry])?;
let entries = seed_entries(segment_log::now_millis(), session_id, state, None, None);
store.create_segment(session_id, fork_id, &entries)?;
store.copy_committed_uploaded_files(source_session_id, session_id)?;
Ok((session_id, fork_id))
}
@@ -460,11 +469,18 @@ pub fn fork_at(
) -> Result<SegmentId, StoreError> {
let entries = store.read_all(source_session_id, source_id)?;
let cut = if at_turn_index == 0 {
// Branch directly after the SegmentStart (or whatever opens the
// segment), before any turn completes.
// Branch from the seeded state before any new turn completes. A typed
// input checkpoint immediately following SegmentStart is part of that
// seed and must stay atomic with its annotated history.
entries
.iter()
.position(|e| !matches!(e, LogEntry::AnnotatedSegmentStart { .. }))
.position(|entry| {
!matches!(
entry,
LogEntry::AnnotatedSegmentStart { .. }
| LogEntry::InputSegmentsCheckpoint { .. }
)
})
.unwrap_or(entries.len())
} else {
entries
@@ -476,8 +492,9 @@ pub fn fork_at(
let state = segment_log::collect_state(&entries[..cut]);
let fork_id = crate::new_segment_id();
let ts = segment_log::now_millis();
let entry = LogEntry::AnnotatedSegmentStart {
ts: segment_log::now_millis(),
ts,
session_id: source_session_id,
system_prompt: state.system_prompt,
config: state.config,
@@ -488,7 +505,14 @@ pub fn fork_at(
}),
compacted_from: None,
};
store.create_segment(source_session_id, fork_id, &[entry])?;
let mut fork_entries = vec![entry];
if !state.user_segments.is_empty() {
fork_entries.push(LogEntry::InputSegmentsCheckpoint {
ts,
user_segments: state.user_segments,
});
}
store.create_segment(source_session_id, fork_id, &fork_entries)?;
Ok(fork_id)
}
+11
View File
@@ -63,6 +63,14 @@ pub enum LogEntry {
compacted_from: Option<SegmentOrigin>,
},
/// Typed user-segment projection accompanying a compacted or forked
/// SegmentStart history snapshot. This keeps attachment identity and
/// metadata aligned with retained user entries without embedding bodies.
InputSegmentsCheckpoint {
ts: u64,
user_segments: Vec<Vec<Segment>>,
},
/// IDLE → active marker. Records the start of a new self-driving
/// cycle (Invoke range). The range extends implicitly until the
/// next `Invoke` entry; this entry carries the trigger only — the
@@ -273,6 +281,9 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState {
.map(|entry| Item::from(entry.item))
.collect();
}
LogEntry::InputSegmentsCheckpoint { user_segments, .. } => {
state.user_segments = user_segments.clone();
}
LogEntry::Invoke { .. } => {
// A terminal run record below clears or refines this. If the
// log ends first, restore must treat the turn as interrupted.
+131 -1
View File
@@ -13,7 +13,10 @@
use crate::event_trace::TraceEntry;
use crate::segment_log::LogEntry;
use crate::{SegmentId, SessionId};
use crate::{
PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits, UploadedFileUploadContext,
};
use protocol::{PasteArtifactRef, UploadedFileRef};
/// Errors from the persistence store.
#[derive(Debug, thiserror::Error)]
@@ -29,6 +32,42 @@ pub enum StoreError {
#[error("log corrupted at line {line}: {message}")]
Corrupt { line: usize, message: String },
#[error("paste artifact storage is unavailable")]
PasteArtifactUnsupported,
#[error("paste artifact not found: {0}")]
PasteArtifactNotFound(String),
#[error("paste artifact integrity check failed: {0}")]
PasteArtifactIntegrity(String),
#[error("paste artifact size limit exceeded: {0}")]
PasteArtifactLimit(String),
#[error("uploaded file is too large")]
ArtifactTooLarge,
#[error("session artifact aggregate quota exceeded")]
ArtifactQuotaExceeded,
#[error("uploaded file reference integrity check failed")]
ArtifactIntegrityMismatch,
#[error("uploaded file name is invalid")]
InvalidUploadedFileName,
#[error("uploaded file media type is invalid")]
InvalidUploadedFileMediaType,
#[error("uploaded file is already committed to session history")]
ArtifactAlreadyCommitted,
#[error("artifact id is invalid")]
InvalidArtifactId,
#[error("artifact timestamp is invalid")]
InvalidTimestamp,
}
/// Sync persistence backend for segment logs.
@@ -117,6 +156,97 @@ pub trait Store: Send + Sync {
segment_id: SegmentId,
) -> Result<usize, StoreError>;
/// Store a large paste before its reference is committed to history.
fn write_paste_artifact(
&self,
_session_id: SessionId,
_source_entry_id: &str,
_content: &str,
_limits: PasteArtifactLimits,
) -> Result<PasteArtifactRef, StoreError> {
Err(StoreError::PasteArtifactUnsupported)
}
/// Read and verify one artifact owned by `session_id`.
fn read_paste_artifact(
&self,
_session_id: SessionId,
_artifact_id: &str,
) -> Result<(PasteArtifactRef, String), StoreError> {
Err(StoreError::PasteArtifactUnsupported)
}
/// Persist a client-local file before a submission references it.
fn write_uploaded_file(
&self,
_session_id: SessionId,
_file_name: &str,
_media_type: &str,
_content: &[u8],
_limits: UploadedFileLimits,
) -> Result<UploadedFileRef, StoreError> {
Err(StoreError::PasteArtifactUnsupported)
}
fn write_uploaded_file_with_context(
&self,
session_id: SessionId,
file_name: &str,
media_type: &str,
content: &[u8],
_context: &UploadedFileUploadContext,
limits: UploadedFileLimits,
) -> Result<UploadedFileRef, StoreError> {
self.write_uploaded_file(session_id, file_name, media_type, content, limits)
}
/// Read and integrity-check an uploaded file owned by `session_id`.
fn read_uploaded_file(
&self,
_session_id: SessionId,
_reference: &UploadedFileRef,
) -> Result<Vec<u8>, StoreError> {
Err(StoreError::PasteArtifactUnsupported)
}
fn read_uploaded_file_by_id(
&self,
_session_id: SessionId,
_artifact_id: &str,
) -> Result<(UploadedFileRef, Vec<u8>), StoreError> {
Err(StoreError::PasteArtifactUnsupported)
}
fn bind_uploaded_file(
&self,
_session_id: SessionId,
_reference: &UploadedFileRef,
_source_entry_id: &str,
) -> Result<UploadedFileRef, StoreError> {
Err(StoreError::PasteArtifactUnsupported)
}
/// Delete an uncommitted uploaded file owned by `session_id`.
fn delete_uploaded_file(
&self,
_session_id: SessionId,
_artifact_id: &str,
) -> Result<bool, StoreError> {
Err(StoreError::PasteArtifactUnsupported)
}
fn delete_uncommitted_uploaded_files(&self, _session_id: SessionId) -> Result<u64, StoreError> {
Ok(0)
}
fn copy_committed_uploaded_files(
&self,
_source_session_id: SessionId,
_target_session_id: SessionId,
) -> Result<u64, StoreError> {
Ok(0)
}
/// Append a trace entry to the debug event trace file.
fn append_trace(
&self,
+534
View File
@@ -0,0 +1,534 @@
use std::{
fs,
path::Path,
time::{SystemTime, UNIX_EPOCH},
};
use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
use fs4::fs_std::FileExt;
use protocol::{UploadedFileAvailability, UploadedFileRef};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use unicode_normalization::UnicodeNormalization;
use unicode_properties::general_category::{GeneralCategory, UnicodeGeneralCategory};
use unicode_security::{confusable_detection::skeleton, mixed_script::MixedScript};
use uuid::Uuid;
use crate::StoreError;
type Result<T> = std::result::Result<T, StoreError>;
pub const DEFAULT_MAX_UPLOADED_FILE_BYTES: u64 = 10 * 1024 * 1024;
pub const DEFAULT_MAX_SESSION_ARTIFACT_BYTES: u64 = 32 * 1024 * 1024;
pub const DEFAULT_MAX_FILES_PER_SUBMISSION: usize = 8;
pub const DEFAULT_MAX_SESSION_UPLOADED_FILES: u64 = 256;
const MAX_FILE_NAME_CHARS: usize = 255;
const MAX_MEDIA_TYPE_BYTES: usize = 127;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct UploadedFileLimits {
pub max_file_bytes: u64,
pub max_session_bytes: u64,
}
impl Default for UploadedFileLimits {
fn default() -> Self {
Self {
max_file_bytes: DEFAULT_MAX_UPLOADED_FILE_BYTES,
max_session_bytes: DEFAULT_MAX_SESSION_ARTIFACT_BYTES,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UploadedFileUploadContext {
pub upload_id: String,
pub principal_id: String,
pub workspace_id: String,
pub runtime_id: String,
pub worker_id: String,
}
#[derive(Debug, Serialize, Deserialize)]
struct StoredUploadedFile {
file_name: String,
media_type: String,
created_at_ms: u64,
byte_len: u64,
sha256: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
source_entry_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
upload_context: Option<UploadedFileUploadContext>,
content_base64: String,
}
pub(crate) fn validate_file_name(file_name: &str) -> Result<()> {
let normalized: String = file_name.nfkc().collect();
let has_unsafe_component = file_name
.split('.')
.filter(|part| !part.is_empty())
.any(|part| {
let confusable_skeleton: String = skeleton(part).collect();
let ascii_confusable = part.chars().any(|ch| !ch.is_ascii())
&& confusable_skeleton.is_ascii()
&& !confusable_skeleton.eq_ignore_ascii_case(part);
!part.is_single_script() || ascii_confusable
});
if file_name.is_empty()
|| file_name.chars().count() > MAX_FILE_NAME_CHARS
|| file_name == "."
|| file_name == ".."
|| normalized != file_name
|| has_unsafe_component
|| file_name.chars().any(|ch| {
ch.is_control()
|| ch.general_category() == GeneralCategory::Format
|| matches!(ch, '/' | '\\')
})
{
return Err(StoreError::InvalidUploadedFileName);
}
Ok(())
}
pub(crate) fn validate_media_type(media_type: &str) -> Result<()> {
let valid = !media_type.is_empty()
&& media_type.len() <= MAX_MEDIA_TYPE_BYTES
&& media_type.is_ascii()
&& !media_type
.bytes()
.any(|byte| byte.is_ascii_control() || byte == b' ')
&& media_type.split_once('/').is_some_and(|(kind, subtype)| {
!kind.is_empty()
&& !subtype.is_empty()
&& kind.bytes().chain(subtype.bytes()).all(|byte| {
byte.is_ascii_alphanumeric()
|| matches!(
byte,
b'!' | b'#' | b'$' | b'&' | b'^' | b'_' | b'.' | b'+' | b'-'
)
})
});
let allowed = media_type.starts_with("text/")
|| matches!(
media_type,
"application/json"
| "application/pdf"
| "image/png"
| "image/jpeg"
| "image/gif"
| "image/webp"
);
if !valid || !allowed {
return Err(StoreError::InvalidUploadedFileMediaType);
}
Ok(())
}
fn normalized_file_name(file_name: &str) -> String {
file_name.nfkc().flat_map(char::to_lowercase).collect()
}
fn validate_content(media_type: &str, content: &[u8]) -> Result<()> {
if content.is_empty() {
return Err(StoreError::InvalidUploadedFileMediaType);
}
let matches_declared_type = if media_type.starts_with("text/") {
std::str::from_utf8(content).is_ok()
} else {
match media_type {
"application/json" => serde_json::from_slice::<serde_json::Value>(content).is_ok(),
"application/pdf" => content.starts_with(b"%PDF-"),
"image/png" => content.starts_with(b"\x89PNG\r\n\x1a\n"),
"image/jpeg" => content.starts_with(&[0xff, 0xd8, 0xff]),
"image/gif" => content.starts_with(b"GIF87a") || content.starts_with(b"GIF89a"),
"image/webp" => {
content.len() >= 12 && content.starts_with(b"RIFF") && &content[8..12] == b"WEBP"
}
_ => false,
}
};
if !matches_declared_type {
return Err(StoreError::ArtifactIntegrityMismatch);
}
Ok(())
}
fn record_path(dir: &Path, artifact_id: &str) -> Result<std::path::PathBuf> {
let id = Uuid::parse_str(artifact_id).map_err(|_| StoreError::InvalidArtifactId)?;
Ok(dir.join(format!("{id}.file.json")))
}
fn now_ms() -> Result<u64> {
let value = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_err(|_| StoreError::InvalidTimestamp)?
.as_millis();
u64::try_from(value).map_err(|_| StoreError::InvalidTimestamp)
}
fn digest(bytes: &[u8]) -> String {
Sha256::digest(bytes)
.iter()
.map(|byte| format!("{byte:02x}"))
.collect()
}
pub(crate) fn stored_uploaded_file_usage(dir: &Path) -> Result<(u64, u64)> {
if !dir.exists() {
return Ok((0, 0));
}
let mut bytes = 0_u64;
let mut count = 0_u64;
for entry in fs::read_dir(dir)? {
let entry = entry?;
let path = entry.path();
if !entry.file_type()?.is_file()
|| !path
.file_name()
.and_then(|name| name.to_str())
.is_some_and(|name| name.ends_with(".file.json"))
{
continue;
}
let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
bytes = bytes
.checked_add(stored.byte_len)
.ok_or(StoreError::ArtifactQuotaExceeded)?;
count = count
.checked_add(1)
.ok_or(StoreError::ArtifactQuotaExceeded)?;
}
Ok((bytes, count))
}
pub(crate) fn write_uploaded_file(
dir: &Path,
file_name: &str,
media_type: &str,
content: &[u8],
context: Option<&UploadedFileUploadContext>,
limits: UploadedFileLimits,
) -> Result<UploadedFileRef> {
validate_file_name(file_name)?;
validate_media_type(media_type)?;
validate_content(media_type, content)?;
let byte_len = u64::try_from(content.len()).map_err(|_| StoreError::ArtifactTooLarge)?;
let sha256 = digest(content);
if byte_len > limits.max_file_bytes {
return Err(StoreError::ArtifactTooLarge);
}
fs::create_dir_all(dir)?;
let aggregate_lock = fs::OpenOptions::new()
.create(true)
.read(true)
.write(true)
.open(dir.join(".aggregate.lock"))?;
FileExt::lock_exclusive(&aggregate_lock)?;
let (paste_bytes, _) = crate::paste_artifact::stored_paste_usage(dir)?;
let (file_bytes, file_count) = stored_uploaded_file_usage(dir)?;
let normalized_name = normalized_file_name(file_name);
for entry in fs::read_dir(dir)? {
let path = entry?.path();
if !path
.file_name()
.and_then(|name| name.to_str())
.is_some_and(|name| name.ends_with(".file.json"))
{
continue;
}
let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
let same_context = context.is_some() && stored.upload_context.as_ref() == context;
let same_uncommitted_name = stored.source_entry_id.is_none()
&& normalized_file_name(&stored.file_name) == normalized_name;
if same_context || same_uncommitted_name {
if stored.file_name == file_name
&& stored.media_type == media_type
&& stored.byte_len == byte_len
&& stored.sha256 == sha256
&& stored.upload_context.as_ref() == context
{
let artifact_id = path
.file_name()
.and_then(|name| name.to_str())
.and_then(|name| name.strip_suffix(".file.json"))
.ok_or(StoreError::InvalidArtifactId)?
.to_string();
return Ok(UploadedFileRef {
artifact_id,
file_name: stored.file_name,
media_type: stored.media_type,
created_at_ms: stored.created_at_ms,
availability: UploadedFileAvailability::Available,
byte_len: stored.byte_len,
sha256: stored.sha256,
source_entry_id: None,
});
}
return Err(StoreError::InvalidUploadedFileName);
}
}
if file_count >= DEFAULT_MAX_SESSION_UPLOADED_FILES {
return Err(StoreError::ArtifactQuotaExceeded);
}
if paste_bytes
.checked_add(file_bytes)
.and_then(|total| total.checked_add(byte_len))
.is_none_or(|total| total > limits.max_session_bytes)
{
return Err(StoreError::ArtifactQuotaExceeded);
}
let artifact_id = Uuid::now_v7().to_string();
let created_at_ms = now_ms()?;
let stored = StoredUploadedFile {
file_name: file_name.to_owned(),
media_type: media_type.to_owned(),
created_at_ms,
byte_len,
sha256: sha256.clone(),
source_entry_id: None,
upload_context: context.cloned(),
content_base64: BASE64.encode(content),
};
let path = record_path(dir, &artifact_id)?;
let temp = dir.join(format!(".{artifact_id}.file.tmp"));
fs::write(&temp, serde_json::to_vec(&stored)?)?;
fs::rename(&temp, &path)?;
Ok(UploadedFileRef {
artifact_id,
file_name: file_name.to_owned(),
media_type: media_type.to_owned(),
created_at_ms,
availability: UploadedFileAvailability::Available,
byte_len,
sha256,
source_entry_id: None,
})
}
pub(crate) fn read_uploaded_file_by_id(
dir: &Path,
artifact_id: &str,
) -> Result<(UploadedFileRef, Vec<u8>)> {
let stored: StoredUploadedFile =
serde_json::from_slice(&fs::read(record_path(dir, artifact_id)?)?)?;
let content = BASE64
.decode(&stored.content_base64)
.map_err(|_| StoreError::ArtifactIntegrityMismatch)?;
if u64::try_from(content.len()).ok() != Some(stored.byte_len)
|| digest(&content) != stored.sha256
{
return Err(StoreError::ArtifactIntegrityMismatch);
}
let reference = UploadedFileRef {
artifact_id: artifact_id.to_owned(),
file_name: stored.file_name,
media_type: stored.media_type,
created_at_ms: stored.created_at_ms,
availability: UploadedFileAvailability::Available,
byte_len: stored.byte_len,
sha256: stored.sha256,
source_entry_id: stored.source_entry_id,
};
Ok((reference, content))
}
pub(crate) fn read_uploaded_file(dir: &Path, reference: &UploadedFileRef) -> Result<Vec<u8>> {
let (stored_reference, content) = read_uploaded_file_by_id(dir, &reference.artifact_id)?;
if stored_reference.file_name != reference.file_name
|| stored_reference.media_type != reference.media_type
|| stored_reference.created_at_ms != reference.created_at_ms
|| stored_reference.byte_len != reference.byte_len
|| stored_reference.sha256 != reference.sha256
|| stored_reference.source_entry_id != reference.source_entry_id
{
return Err(StoreError::ArtifactIntegrityMismatch);
}
Ok(content)
}
pub(crate) fn clear_uploaded_file_binding(
dir: &Path,
artifact_id: &str,
expected_source_entry_id: &str,
) -> Result<()> {
fs::create_dir_all(dir)?;
let aggregate_lock = fs::OpenOptions::new()
.create(true)
.read(true)
.write(true)
.open(dir.join(".aggregate.lock"))?;
FileExt::lock_exclusive(&aggregate_lock)?;
let path = record_path(dir, artifact_id)?;
let mut stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
if stored.source_entry_id.as_deref() != Some(expected_source_entry_id) {
return Err(StoreError::ArtifactIntegrityMismatch);
}
stored.source_entry_id = None;
let temp = dir.join(format!(".{artifact_id}.file.unbind.tmp"));
fs::write(&temp, serde_json::to_vec(&stored)?)?;
fs::rename(temp, path)?;
Ok(())
}
pub(crate) fn bind_uploaded_file(
dir: &Path,
reference: &UploadedFileRef,
source_entry_id: &str,
) -> Result<UploadedFileRef> {
if source_entry_id.is_empty() || reference.source_entry_id.is_some() {
return Err(StoreError::ArtifactIntegrityMismatch);
}
let aggregate_lock = fs::OpenOptions::new()
.create(true)
.read(true)
.write(true)
.open(dir.join(".aggregate.lock"))?;
FileExt::lock_exclusive(&aggregate_lock)?;
let (stored_reference, _) = read_uploaded_file_by_id(dir, &reference.artifact_id)?;
if stored_reference.file_name != reference.file_name
|| stored_reference.media_type != reference.media_type
|| stored_reference.created_at_ms != reference.created_at_ms
|| stored_reference.byte_len != reference.byte_len
|| stored_reference.sha256 != reference.sha256
{
return Err(StoreError::ArtifactIntegrityMismatch);
}
let path = record_path(dir, &reference.artifact_id)?;
let mut stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
if stored.source_entry_id.is_some() {
return Err(StoreError::ArtifactAlreadyCommitted);
}
stored.source_entry_id = Some(source_entry_id.to_owned());
let temp = dir.join(format!(".{}.file.bind.tmp", reference.artifact_id));
fs::write(&temp, serde_json::to_vec(&stored)?)?;
fs::rename(&temp, path)?;
let mut bound = reference.clone();
bound.source_entry_id = Some(source_entry_id.to_owned());
Ok(bound)
}
pub(crate) fn list_uploaded_file_refs(dir: &Path) -> Result<Vec<UploadedFileRef>> {
if !dir.exists() {
return Ok(Vec::new());
}
let mut refs = Vec::new();
for entry in fs::read_dir(dir)? {
let path = entry?.path();
let Some(artifact_id) = path
.file_name()
.and_then(|name| name.to_str())
.and_then(|name| name.strip_suffix(".file.json"))
else {
continue;
};
refs.push(read_uploaded_file_by_id(dir, artifact_id)?.0);
}
Ok(refs)
}
pub(crate) fn copy_committed_uploaded_files(source_dir: &Path, target_dir: &Path) -> Result<u64> {
if !source_dir.exists() {
return Ok(0);
}
fs::create_dir_all(target_dir)?;
let target_lock = fs::OpenOptions::new()
.create(true)
.read(true)
.write(true)
.open(target_dir.join(".aggregate.lock"))?;
FileExt::lock_exclusive(&target_lock)?;
let mut copied = 0_u64;
for entry in fs::read_dir(source_dir)? {
let entry = entry?;
let path = entry.path();
let Some(name) = path.file_name().and_then(|name| name.to_str()) else {
continue;
};
if !name.ends_with(".file.json") {
continue;
}
let bytes = fs::read(&path)?;
let stored: StoredUploadedFile = serde_json::from_slice(&bytes)?;
if stored.source_entry_id.is_none() {
continue;
}
let target = target_dir.join(name);
if target.exists() {
let existing: StoredUploadedFile = serde_json::from_slice(&fs::read(&target)?)?;
if existing.sha256 != stored.sha256
|| existing.file_name != stored.file_name
|| existing.source_entry_id != stored.source_entry_id
{
return Err(StoreError::ArtifactIntegrityMismatch);
}
continue;
}
let temp = target_dir.join(format!(".{name}.copy.tmp"));
fs::write(&temp, &bytes)?;
fs::rename(temp, target)?;
copied = copied
.checked_add(1)
.ok_or(StoreError::ArtifactQuotaExceeded)?;
}
Ok(copied)
}
pub(crate) fn delete_uncommitted_uploaded_files(dir: &Path) -> Result<u64> {
fs::create_dir_all(dir)?;
let aggregate_lock = fs::OpenOptions::new()
.create(true)
.read(true)
.write(true)
.open(dir.join(".aggregate.lock"))?;
FileExt::lock_exclusive(&aggregate_lock)?;
let mut removed = 0_u64;
for entry in fs::read_dir(dir)? {
let entry = entry?;
let path = entry.path();
if !path
.file_name()
.and_then(|name| name.to_str())
.is_some_and(|name| name.ends_with(".file.json"))
{
continue;
}
let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
if stored.source_entry_id.is_none() {
fs::remove_file(path)?;
removed = removed
.checked_add(1)
.ok_or(StoreError::ArtifactQuotaExceeded)?;
}
}
Ok(removed)
}
pub(crate) fn delete_uploaded_file(dir: &Path, artifact_id: &str) -> Result<bool> {
fs::create_dir_all(dir)?;
let aggregate_lock = fs::OpenOptions::new()
.create(true)
.read(true)
.write(true)
.open(dir.join(".aggregate.lock"))?;
FileExt::lock_exclusive(&aggregate_lock)?;
let path = record_path(dir, artifact_id)?;
let stored = match fs::read(&path) {
Ok(bytes) => serde_json::from_slice::<StoredUploadedFile>(&bytes)?,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
Err(error) => return Err(error.into()),
};
if stored.source_entry_id.is_some() {
return Err(StoreError::ArtifactAlreadyCommitted);
}
match fs::remove_file(path) {
Ok(()) => Ok(true),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false),
Err(error) => Err(error.into()),
}
}
@@ -608,6 +608,24 @@ where
) -> Result<usize, crate::StoreError> {
self.session_store.read_entry_count(session_id, segment_id)
}
fn write_paste_artifact(
&self,
session_id: SessionId,
source_entry_id: &str,
content: &str,
limits: crate::PasteArtifactLimits,
) -> Result<protocol::PasteArtifactRef, crate::StoreError> {
self.session_store
.write_paste_artifact(session_id, source_entry_id, content, limits)
}
fn read_paste_artifact(
&self,
session_id: SessionId,
artifact_id: &str,
) -> Result<(protocol::PasteArtifactRef, String), crate::StoreError> {
self.session_store
.read_paste_artifact(session_id, artifact_id)
}
fn append_trace(
&self,
session_id: SessionId,
@@ -10,9 +10,11 @@
//! every later operation must use that same ID.
use crate::event_trace::TraceEntry;
use crate::paste_artifact::{read_from_dir, write_to_dir};
use crate::segment_log::LogEntry;
use crate::store::{Store, StoreError};
use crate::{SegmentId, SessionId};
use crate::{PasteArtifactLimits, SegmentId, SessionId};
use protocol::PasteArtifactRef;
use serde::{Deserialize, Serialize};
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write};
@@ -25,6 +27,7 @@ const PREVIOUS_SESSION_SCHEMA_VERSION: u32 = 2;
const LEGACY_SESSION_SCHEMA_VERSION: u32 = 1;
const SESSION_FILE: &str = "session.json";
const SEGMENTS_DIR: &str = "segments";
const PASTE_ARTIFACTS_DIR: &str = "artifacts/paste";
#[derive(Clone)]
pub struct WorkerSessionStore {
@@ -317,6 +320,35 @@ impl Store for WorkerSessionStore {
.count())
}
fn write_paste_artifact(
&self,
session_id: SessionId,
source_entry_id: &str,
content: &str,
limits: PasteArtifactLimits,
) -> Result<PasteArtifactRef, StoreError> {
self.ensure_session(session_id, true)?;
let _guard = self
.append_lock
.lock()
.map_err(|_| std::io::Error::other("Worker Session append lock was poisoned"))?;
write_to_dir(
&self.root.join(PASTE_ARTIFACTS_DIR),
source_entry_id,
content,
limits,
)
}
fn read_paste_artifact(
&self,
session_id: SessionId,
artifact_id: &str,
) -> Result<(PasteArtifactRef, String), StoreError> {
self.ensure_session(session_id, false)?;
read_from_dir(&self.root.join(PASTE_ARTIFACTS_DIR), artifact_id)
}
fn append_trace(
&self,
session_id: SessionId,
@@ -601,6 +633,45 @@ mod tests {
assert_eq!(store.list_sessions().unwrap(), vec![session_id]);
}
#[test]
fn worker_session_store_keeps_paste_artifacts_inside_retention_root() {
let root = tempfile::tempdir().unwrap();
let store = WorkerSessionStore::new(root.path().join("session")).unwrap();
let session_id = new_session_id();
store
.create_segment(session_id, new_segment_id(), &[])
.unwrap();
let content = "large paste body\n終端\n";
let reference = store
.write_paste_artifact(
session_id,
"entry-1",
content,
PasteArtifactLimits::default(),
)
.unwrap();
assert!(
root.path()
.join(format!(
"session/{PASTE_ARTIFACTS_DIR}/{}.json",
reference.artifact_id
))
.is_file()
);
assert_eq!(
store
.read_paste_artifact(session_id, &reference.artifact_id)
.unwrap()
.1,
content
);
assert!(matches!(
store.read_paste_artifact(new_session_id(), &reference.artifact_id),
Err(StoreError::Corrupt { .. })
));
}
#[test]
fn schema_v1_logs_are_rewritten_and_promoted_to_v3() {
let root = tempfile::tempdir().unwrap();
+118 -7
View File
@@ -10,6 +10,7 @@ use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use agen::{Engine, History};
use async_trait::async_trait;
use common::MockLlmClient;
use protocol::{Segment, SessionSnapshotEntryData, UploadedFileAvailability, UploadedFileRef};
use session_store::{FsStore, LogEntry, SegmentStartState, Store, collect_state};
// =============================================================================
@@ -236,6 +237,7 @@ async fn session_run_logs_entries() {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: annotated(&worker.history()),
user_segments: Vec::new(),
},
)
.unwrap();
@@ -284,6 +286,7 @@ async fn session_restore_round_trip() {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: annotated(&worker.history()),
user_segments: Vec::new(),
},
)
.unwrap();
@@ -323,6 +326,7 @@ async fn session_run_with_tool_call() {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: annotated(&worker.history()),
user_segments: Vec::new(),
},
)
.unwrap();
@@ -358,6 +362,7 @@ async fn session_resume_after_pause() {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: annotated(&worker.history()),
user_segments: Vec::new(),
},
)
.unwrap();
@@ -397,6 +402,7 @@ async fn session_fork_creates_new_session() {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: annotated(&worker.history()),
user_segments: Vec::new(),
},
)
.unwrap();
@@ -404,28 +410,38 @@ async fn session_fork_creates_new_session() {
let (worker, _) = run_and_persist(worker, &store, sid, segid, "Hello").await;
let original_history_len = worker.history().len();
let source_user_segments = session_store::restore(&store, sid, segid)
.unwrap()
.user_segments;
let (fork_sid, fork_segid) = session_store::fork(
&store,
sid,
SegmentStartState {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: annotated(&worker.history()),
user_segments: source_user_segments.clone(),
},
)
.unwrap();
assert_ne!(fork_sid, sid, "`fork` mints a fresh Session");
// Fork should have a SegmentStart with the current history
// Fork should have an annotated seed and typed input checkpoint.
let fork_entries = store.read_all(fork_sid, fork_segid).unwrap();
assert_eq!(fork_entries.len(), 1);
assert_eq!(fork_entries.len(), 2);
assert!(matches!(
&fork_entries[0],
LogEntry::AnnotatedSegmentStart { .. }
));
assert!(matches!(
&fork_entries[1],
LogEntry::InputSegmentsCheckpoint { .. }
));
let fork_state = collect_state(&fork_entries);
assert_eq!(fork_state.session_id, Some(fork_sid));
assert_eq!(fork_state.history.len(), original_history_len);
assert_eq!(fork_state.user_segments, source_user_segments);
assert_eq!(fork_state.system_prompt.as_deref(), Some("System prompt"));
}
@@ -441,6 +457,7 @@ async fn session_fork_at_truncates_within_session() {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: annotated(&worker.history()),
user_segments: Vec::new(),
},
)
.unwrap();
@@ -454,7 +471,11 @@ async fn session_fork_at_truncates_within_session() {
let fork_segid = session_store::fork_at(&store, sid, segid, worker.turn_count()).unwrap();
let fork_entries = store.read_all(sid, fork_segid).unwrap();
assert_eq!(fork_entries.len(), 1); // Just the new SegmentStart
assert_eq!(fork_entries.len(), 2);
assert!(matches!(
&fork_entries[1],
LogEntry::InputSegmentsCheckpoint { .. }
));
let fork_state = collect_state(&fork_entries);
assert_eq!(fork_state.session_id, Some(sid), "fork_at inherits Session");
@@ -466,6 +487,7 @@ async fn session_fork_at_truncates_within_session() {
.position(|e| matches!(e, LogEntry::TurnEnd { turn_count, .. } if *turn_count == worker.turn_count()))
.expect("source segment has the matching TurnEnd");
let source_state_at_fork = collect_state(&all_entries[..=turn_end_pos]);
assert_eq!(fork_state.user_segments, source_state_at_fork.user_segments);
assert_eq!(fork_state.history.len(), source_state_at_fork.history.len());
assert_eq!(
fork_state.annotated_history, source_state_at_fork.annotated_history,
@@ -491,6 +513,84 @@ async fn session_fork_at_truncates_within_session() {
assert!(segs.contains(&fork_segid));
}
#[test]
fn rewound_fork_preserves_uploaded_file_segments_in_snapshot() {
let (_dir, store) = make_store();
let config = RequestConfig::default();
let (sid, segid) = session_store::create_segment(
&store,
SegmentStartState {
system_prompt: Some("System prompt"),
config: &config,
history: Vec::new(),
user_segments: Vec::new(),
},
)
.unwrap();
let uploaded = UploadedFileRef {
artifact_id: "uploaded-file-1".into(),
file_name: "notes.txt".into(),
media_type: "text/plain".into(),
created_at_ms: 123,
availability: UploadedFileAvailability::Available,
byte_len: 5,
sha256: "a".repeat(64),
source_entry_id: Some("entry-1".into()),
};
let segments = vec![Segment::UploadedFile {
file: uploaded.clone(),
}];
session_store::save_user_input(
&store,
sid,
segid,
segments.clone(),
annotated(&[Item::user_message(Segment::flatten_to_text(&segments))]),
)
.unwrap();
session_store::save_turn_end(&store, sid, segid, 1).unwrap();
let fork_segid = session_store::fork_at(&store, sid, segid, 1).unwrap();
let fork_entries = store.read_all(sid, fork_segid).unwrap();
let snapshot = session_store::public_snapshot::project_session_snapshot(sid, &fork_entries);
assert!(fork_entries.iter().any(|entry| matches!(
entry,
LogEntry::InputSegmentsCheckpoint { user_segments, .. }
if user_segments == &vec![segments.clone()]
)));
assert!(snapshot.entries.iter().any(|entry| matches!(
&entry.data,
SessionSnapshotEntryData::UserInput { segments: restored }
if restored == &segments
)));
let fork_state = collect_state(&fork_entries);
let (copied_session_id, copied_segment_id) = session_store::fork(
&store,
sid,
SegmentStartState {
system_prompt: fork_state.system_prompt.as_deref(),
config: &fork_state.config,
history: fork_state.annotated_history.clone(),
user_segments: fork_state.user_segments.clone(),
},
)
.unwrap();
let copied_entries = store
.read_all(copied_session_id, copied_segment_id)
.unwrap();
let copied_snapshot = session_store::public_snapshot::project_session_snapshot(
copied_session_id,
&copied_entries,
);
assert!(copied_snapshot.entries.iter().any(|entry| matches!(
&entry.data,
SessionSnapshotEntryData::UserInput { segments: restored }
if restored == &segments
)));
}
#[tokio::test]
async fn session_config_changed_logged() {
let (_dir, store) = make_store();
@@ -503,6 +603,7 @@ async fn session_config_changed_logged() {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: annotated(&worker.history()),
user_segments: Vec::new(),
},
)
.unwrap();
@@ -536,6 +637,7 @@ async fn session_auto_forks_on_conflict() {
system_prompt: worker_a.get_system_prompt(),
config: worker_a.request_config(),
history: annotated(&worker_a.history()),
user_segments: Vec::new(),
},
)
.unwrap();
@@ -564,6 +666,7 @@ async fn session_auto_forks_on_conflict() {
system_prompt: worker_a.get_system_prompt(),
config: worker_a.request_config(),
history: annotated(&worker_a.history()),
user_segments: Vec::new(),
},
)
.unwrap();
@@ -623,6 +726,7 @@ async fn nested_past_fork_leaves_ancestors_immutable() {
system_prompt: worker.get_system_prompt(),
config: worker.request_config(),
history: annotated(&worker.history()),
user_segments: Vec::new(),
},
)
.unwrap();
@@ -653,12 +757,19 @@ async fn nested_past_fork_leaves_ancestors_immutable() {
let fork1_entries = store.read_all(sid, fork1).unwrap();
assert_eq!(
fork1_entries.len(),
1,
"fork1 is just its SegmentStart seed"
2,
"fork1 stores its SegmentStart and typed input checkpoint"
);
// fork2's lineage points at fork1, not the root.
match &store.read_all(sid, fork2).unwrap()[0] {
// fork2's lineage points at fork1, not the root, and the typed seed remains
// intact across the nested turn-zero fork.
let fork2_entries = store.read_all(sid, fork2).unwrap();
assert_eq!(fork2_entries.len(), 2);
assert_eq!(
collect_state(&fork2_entries).user_segments,
collect_state(&fork1_entries).user_segments
);
match &fork2_entries[0] {
LogEntry::AnnotatedSegmentStart {
forked_from: Some(origin),
..
+1
View File
@@ -7,6 +7,7 @@ license.workspace = true
[dependencies]
agen.workspace = true
client.workspace = true
fs4.workspace = true
manifest.workspace = true
protocol.workspace = true
+227 -93
View File
@@ -2,26 +2,36 @@ use std::path::PathBuf;
use std::time::Duration;
use agen::llm_client::client::LlmClient;
use protocol::{Event, Method};
use client::Client;
use client::transport::in_process::{Peer as InProcessPeer, Socket as InProcessSocket};
use manifest::ScopeRule;
use protocol::stream::{decode_method, encode_event};
use protocol::{Event, Method, WorkerId};
use session_store::{
CombinedStore, FsStore, FsWorkerStore, WorkerActiveSegmentRef, WorkerMetadataStore,
};
use thiserror::Error;
use tokio::sync::broadcast;
use worker::bootstrap::{WorkerBootstrap, WorkerBootstrapError, WorkerBootstrapLayout};
use worker::bootstrap::{
WorkerBootstrap, WorkerBootstrapError, WorkerBootstrapLayout, bash_output_dir_for_worker_id,
};
use worker::controller::WorkerControllerTransport;
use worker::ipc::protocol_session::{
WorkerProtocolSessionStreams, dispatch_worker_protocol_method, live_log_entry_event,
subscribe_worker_protocol_session,
};
use worker::runtime::worker_allocation::ScopeLockError;
use worker::{BootstrappedWorker, WorkerError, WorkerFilesystemAuthority, WorkerWorkspaceContext};
use crate::launch::ResolvedStandaloneLaunch;
use crate::store::{
StaleLeasePolicy, StandaloneSessionId, StandaloneSessionLease, StandaloneSessionRecord,
StandaloneSessionStore, StandaloneShutdownReason, StandaloneStoreError,
StaleLeasePolicy, StandaloneShutdownReason, StandaloneStoreError, StandaloneWorkerLease,
StandaloneWorkerRecord, StandaloneWorkerStore,
};
const DEFAULT_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10);
type StandaloneBackingStore = CombinedStore<FsStore, FsWorkerStore>;
/// One client-owned top-level Worker and its standalone session authority.
/// One client-owned top-level Worker and its standalone Worker authority.
///
/// The host deliberately exposes the existing typed Worker protocol rather than owning an
/// HTTP/WebSocket server or creating Runtime/Workspace/Ticket/Workdir domain records.
@@ -29,22 +39,32 @@ pub struct StandaloneHost {
handle: worker::WorkerHandle,
shutdown: Option<worker::controller::ShutdownReceiver>,
shutdown_timeout: Duration,
store: StandaloneSessionStore,
store: StandaloneWorkerStore,
worker_store: FsWorkerStore,
record: StandaloneSessionRecord,
lease: Option<StandaloneSessionLease>,
record: StandaloneWorkerRecord,
lease: Option<StandaloneWorkerLease>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum StandaloneStartupError {
#[error("the standalone state store could not be opened or validated")]
StateStore,
#[error("the standalone session is already active")]
SessionActive,
#[error("the standalone session lease cannot be observed safely; recovery is rejected")]
#[error("the standalone Worker is already active")]
WorkerActive,
#[error("the standalone Worker lease cannot be observed safely; recovery is rejected")]
LeaseLivenessUnknown,
#[error("the standalone session working directory is unavailable or changed")]
#[error("the standalone Worker working directory is unavailable or changed")]
WorkingDirectoryUnavailable,
#[error(
"requested scope `{}` conflicts with worker allocation `{competitor}` rule `{}`",
requested_rule.target.display(),
competitor_rule.target.display()
)]
ScopeConflict {
competitor: String,
requested_rule: ScopeRule,
competitor_rule: ScopeRule,
},
#[error("the resolved Worker configuration or persisted history is invalid")]
WorkerConfiguration,
#[error("the configured model provider is unavailable")]
@@ -55,19 +75,13 @@ pub enum StandaloneStartupError {
Controller,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
pub enum StandaloneRequestError {
#[error("the standalone Worker is no longer accepting requests")]
WorkerUnavailable,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
pub enum StandaloneShutdownError {
#[error("the standalone Worker did not stop before the shutdown deadline")]
DeadlineExceeded,
#[error("the standalone Worker shutdown confirmation was lost")]
ConfirmationLost,
#[error("the standalone session final state could not be committed")]
#[error("the standalone Worker final state could not be committed")]
StateStore,
}
@@ -87,22 +101,24 @@ impl StandaloneHost {
}
async fn start_with_optional_model_client(
mut launch: ResolvedStandaloneLaunch,
launch: ResolvedStandaloneLaunch,
model_client: Option<Box<dyn LlmClient>>,
) -> Result<Self, StandaloneStartupError> {
let store = StandaloneSessionStore::open(&launch.state_dir)
.map_err(classify_store_startup_error)?;
let store =
StandaloneWorkerStore::open(&launch.state_dir).map_err(classify_store_startup_error)?;
let allocation = store
.allocate(&launch.cwd, StaleLeasePolicy::Reject)
.map_err(classify_store_startup_error)?;
let id = allocation.id();
let worker_id = allocation.worker_id();
// The standalone session ID is the local identity. A unique internal Worker name avoids
// process-global allocation collisions without creating a Runtime/Workspace Worker ID.
launch.profile.manifest.worker.name = format!("standalone-{id}");
// WorkerId is the stable identity. The current Worker store remains
// name-keyed, so keep its derived storage key separate from the
// user-facing profile name.
let manifest = launch.profile.manifest.clone();
let worker_name = manifest.worker.name.clone();
let (backing_store, worker_store) = match backing_store(&store, id) {
let storage_key = format!("standalone-{worker_id}");
let mut bootstrap_manifest = manifest.clone();
bootstrap_manifest.worker.name = storage_key.clone();
let (backing_store, worker_store) = match backing_store(&store, worker_id) {
Ok(stores) => stores,
Err(error) => {
let _ = store.abandon_allocation(allocation);
@@ -112,15 +128,19 @@ impl StandaloneHost {
let filesystem_authority =
WorkerFilesystemAuthority::local(launch.cwd.clone(), launch.cwd.clone());
let workspace_context = WorkerWorkspaceContext::local_filesystem(None);
let runtime_base = store.runtime_dir(id);
let runtime_base = store.runtime_dir(worker_id);
let bash_output_dir = bash_output_dir_for_worker_id(worker_id);
let mut bootstrap = WorkerBootstrap::new(
manifest.clone(),
bootstrap_manifest,
backing_store,
launch.prompt_catalog,
workspace_context,
filesystem_authority,
WorkerBootstrapLayout::Direct { runtime_base },
WorkerBootstrapLayout::Direct {
runtime_base,
bash_output_dir,
},
WorkerControllerTransport::InProcess,
);
if let Some(model_client) = model_client {
@@ -133,7 +153,7 @@ impl StandaloneHost {
return Err(classify_startup_error(error));
}
};
let active = match active_pointer(&worker_store, &worker_name) {
let active = match active_pointer(&worker_store, &storage_key) {
Ok(active) => active,
Err(error) => {
stop_started_worker(started).await;
@@ -141,16 +161,20 @@ impl StandaloneHost {
return Err(error);
}
};
let record =
match store.commit_created(&allocation, manifest, active.session_id, active.segment_id)
{
Ok(record) => record,
Err(_) => {
stop_started_worker(started).await;
let _ = store.abandon_allocation(allocation);
return Err(StandaloneStartupError::StateStore);
}
};
let record = match store.commit_created(
&allocation,
manifest,
storage_key,
active.session_id,
active.segment_id,
) {
Ok(record) => record,
Err(_) => {
stop_started_worker(started).await;
let _ = store.abandon_allocation(allocation);
return Err(StandaloneStartupError::StateStore);
}
};
Ok(Self::from_started(
started,
store,
@@ -162,50 +186,47 @@ impl StandaloneHost {
pub async fn restore(
state_dir: PathBuf,
session_id: StandaloneSessionId,
worker_id: WorkerId,
) -> Result<Self, StandaloneStartupError> {
Self::restore_with_optional_model_client(state_dir, session_id, None).await
Self::restore_with_optional_model_client(state_dir, worker_id, None).await
}
pub async fn restore_with_model_client<C>(
state_dir: PathBuf,
session_id: StandaloneSessionId,
worker_id: WorkerId,
model_client: C,
) -> Result<Self, StandaloneStartupError>
where
C: LlmClient + 'static,
{
Self::restore_with_optional_model_client(
state_dir,
session_id,
Some(Box::new(model_client)),
)
.await
Self::restore_with_optional_model_client(state_dir, worker_id, Some(Box::new(model_client)))
.await
}
async fn restore_with_optional_model_client(
state_dir: PathBuf,
session_id: StandaloneSessionId,
worker_id: WorkerId,
model_client: Option<Box<dyn LlmClient>>,
) -> Result<Self, StandaloneStartupError> {
let store =
StandaloneSessionStore::open(state_dir).map_err(classify_store_startup_error)?;
let store = StandaloneWorkerStore::open(state_dir).map_err(classify_store_startup_error)?;
let record = store
.load(session_id)
.load(worker_id)
.map_err(classify_store_startup_error)?;
record.cwd.verify().map_err(classify_store_startup_error)?;
let lease = store
.acquire_lease(session_id, StaleLeasePolicy::Recover)
.acquire_lease(worker_id, StaleLeasePolicy::Recover)
.map_err(classify_store_startup_error)?;
let (backing_store, worker_store) = backing_store(&store, session_id)?;
let worker_name = record.worker_name.clone();
let manifest = record.manifest.clone();
let (backing_store, worker_store) = backing_store(&store, worker_id)?;
let storage_key = record.storage_key.clone();
let mut manifest = record.manifest.clone();
manifest.worker.name = storage_key.clone();
let filesystem_authority = WorkerFilesystemAuthority::local(
record.cwd.canonical_path.clone(),
record.cwd.canonical_path.clone(),
);
let workspace_context = WorkerWorkspaceContext::local_filesystem(None);
let runtime_base = store.runtime_dir(session_id);
let runtime_base = store.runtime_dir(worker_id);
let bash_output_dir = bash_output_dir_for_worker_id(worker_id);
let mut bootstrap = WorkerBootstrap::new(
manifest,
@@ -213,18 +234,21 @@ impl StandaloneHost {
worker::PromptCatalogSource::builtins_only(),
workspace_context,
filesystem_authority,
WorkerBootstrapLayout::Direct { runtime_base },
WorkerBootstrapLayout::Direct {
runtime_base,
bash_output_dir,
},
WorkerControllerTransport::InProcess,
);
if let Some(model_client) = model_client {
bootstrap = bootstrap.with_model_client(model_client);
}
let prepared = bootstrap
.prepare_restored(&worker_name)
.prepare_restored(&storage_key)
.await
.map_err(classify_startup_error)?;
let started = prepared.start().await.map_err(classify_startup_error)?;
let active = match active_pointer(&worker_store, &worker_name) {
let active = match active_pointer(&worker_store, &storage_key) {
Ok(active) => active,
Err(error) => {
stop_started_worker(started).await;
@@ -251,10 +275,10 @@ impl StandaloneHost {
fn from_started(
started: BootstrappedWorker,
store: StandaloneSessionStore,
store: StandaloneWorkerStore,
worker_store: FsWorkerStore,
record: StandaloneSessionRecord,
lease: StandaloneSessionLease,
record: StandaloneWorkerRecord,
lease: StandaloneWorkerLease,
) -> Self {
Self {
handle: started.handle,
@@ -268,28 +292,24 @@ impl StandaloneHost {
}
#[must_use]
pub fn session_id(&self) -> StandaloneSessionId {
self.record.session_id
pub fn worker_id(&self) -> WorkerId {
self.record.worker_id
}
#[must_use]
pub fn record(&self) -> &StandaloneSessionRecord {
pub fn record(&self) -> &StandaloneWorkerRecord {
&self.record
}
pub async fn send(&self, method: Method) -> Result<(), StandaloneRequestError> {
self.handle
.send(method)
.await
.map_err(|_| StandaloneRequestError::WorkerUnavailable)
}
pub fn subscribe(&self) -> broadcast::Receiver<Event> {
self.handle.subscribe()
}
pub fn snapshot(&self) -> Event {
self.handle.snapshot_event()
/// Open one complete client-side Worker protocol session.
///
/// Working events, committed session entries, alert snapshots, and the
/// initial history snapshot are merged behind the client boundary.
pub fn connect(&self) -> Client<InProcessSocket> {
let streams = subscribe_worker_protocol_session(&self.handle);
let (socket, peer) = InProcessSocket::pair();
tokio::spawn(run_protocol_session(self.handle.clone(), streams, peer));
Client::new(socket)
}
pub fn with_shutdown_timeout(mut self, shutdown_timeout: Duration) -> Self {
@@ -314,7 +334,7 @@ impl StandaloneHost {
return Err(StandaloneShutdownError::DeadlineExceeded);
}
}
let active = match active_pointer(&self.worker_store, &self.record.worker_name) {
let active = match active_pointer(&self.worker_store, &self.record.storage_key) {
Ok(active) => active,
Err(_) => {
self.retain_lease();
@@ -349,13 +369,118 @@ impl StandaloneHost {
}
}
async fn run_protocol_session(
handle: worker::WorkerHandle,
streams: WorkerProtocolSessionStreams,
mut peer: InProcessPeer,
) {
let WorkerProtocolSessionStreams {
snapshot_event,
mut log_entries,
alert_snapshot,
mut events,
} = streams;
if !send_protocol_snapshot(&peer, alert_snapshot, snapshot_event).await {
return;
}
loop {
tokio::select! {
message = peer.next() => {
let Some(message) = message else {
return;
};
let Ok(method) = decode_method(&message) else {
return;
};
if let Some(event) = dispatch_worker_protocol_method(&handle, method).await
&& !send_protocol_event(&peer, event).await
{
return;
}
}
event = events.recv() => {
match event {
Ok(event) => {
if !send_protocol_event(&peer, event).await {
return;
}
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
let replacement = subscribe_worker_protocol_session(&handle);
let WorkerProtocolSessionStreams {
snapshot_event,
log_entries: replacement_log_entries,
alert_snapshot,
events: replacement_events,
} = replacement;
log_entries = replacement_log_entries;
events = replacement_events;
if !send_protocol_snapshot(&peer, alert_snapshot, snapshot_event).await {
return;
}
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => return,
}
}
entry = log_entries.recv() => {
match entry {
Ok(entry) => {
if let Some(event) = live_log_entry_event(entry)
&& !send_protocol_event(&peer, event).await
{
return;
}
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
let replacement = subscribe_worker_protocol_session(&handle);
let WorkerProtocolSessionStreams {
snapshot_event,
log_entries: replacement_log_entries,
alert_snapshot,
events: replacement_events,
} = replacement;
log_entries = replacement_log_entries;
events = replacement_events;
if !send_protocol_snapshot(&peer, alert_snapshot, snapshot_event).await {
return;
}
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => return,
}
}
}
}
}
async fn send_protocol_snapshot(
peer: &InProcessPeer,
alert_snapshot: Vec<protocol::Alert>,
snapshot_event: Event,
) -> bool {
for alert in alert_snapshot {
if !send_protocol_event(peer, Event::Alert(alert)).await {
return false;
}
}
send_protocol_event(peer, snapshot_event).await
}
async fn send_protocol_event(peer: &InProcessPeer, event: Event) -> bool {
let Ok(message) = encode_event(&event) else {
return false;
};
peer.send(message).await.is_ok()
}
fn backing_store(
store: &StandaloneSessionStore,
id: StandaloneSessionId,
store: &StandaloneWorkerStore,
worker_id: WorkerId,
) -> Result<(StandaloneBackingStore, FsWorkerStore), StandaloneStartupError> {
let session_store =
FsStore::new(store.session_log_dir(id)).map_err(|_| StandaloneStartupError::StateStore)?;
let worker_store = FsWorkerStore::new(store.worker_metadata_dir(id))
let session_store = FsStore::new(store.sessions_dir(worker_id))
.map_err(|_| StandaloneStartupError::StateStore)?;
let worker_store = FsWorkerStore::new(store.worker_metadata_dir(worker_id))
.map_err(|_| StandaloneStartupError::StateStore)?;
Ok((
CombinedStore::new(session_store, worker_store.clone()),
@@ -365,10 +490,10 @@ fn backing_store(
fn active_pointer(
worker_store: &FsWorkerStore,
worker_name: &str,
storage_key: &str,
) -> Result<WorkerActiveSegmentRef, StandaloneStartupError> {
worker_store
.read_by_name(worker_name)
.read_by_name(storage_key)
.map_err(|_| StandaloneStartupError::StateStore)?
.and_then(|metadata| metadata.active)
.ok_or(StandaloneStartupError::StateStore)
@@ -381,7 +506,7 @@ async fn stop_started_worker(started: BootstrappedWorker) {
fn classify_store_startup_error(error: StandaloneStoreError) -> StandaloneStartupError {
match error {
StandaloneStoreError::SessionLeased(_) => StandaloneStartupError::SessionActive,
StandaloneStoreError::WorkerLeased(_) => StandaloneStartupError::WorkerActive,
StandaloneStoreError::LeaseLivenessUnknown(_) => {
StandaloneStartupError::LeaseLivenessUnknown
}
@@ -396,6 +521,15 @@ fn classify_store_startup_error(error: StandaloneStoreError) -> StandaloneStartu
fn classify_startup_error(error: WorkerBootstrapError) -> StandaloneStartupError {
match error {
WorkerBootstrapError::Worker(WorkerError::ScopeLock(ScopeLockError::WriteConflict {
competitor,
rule,
competitor_rule,
})) => StandaloneStartupError::ScopeConflict {
competitor,
requested_rule: rule,
competitor_rule,
},
WorkerBootstrapError::Worker(WorkerError::Provider(_)) => {
StandaloneStartupError::ModelProvider
}
+4 -6
View File
@@ -8,12 +8,10 @@ pub mod host;
pub mod launch;
pub mod store;
pub use host::{
StandaloneHost, StandaloneRequestError, StandaloneShutdownError, StandaloneStartupError,
};
pub use host::{StandaloneHost, StandaloneShutdownError, StandaloneStartupError};
pub use launch::{ResolvedStandaloneLaunch, StandaloneLaunchConfig, StandaloneLaunchError};
pub use protocol::WorkerId;
pub use store::{
StaleLeasePolicy, StandaloneCwdIdentity, StandaloneListScope, StandaloneSessionId,
StandaloneSessionRecord, StandaloneSessionStatus, StandaloneSessionStore,
StandaloneShutdownReason, StandaloneStoreError,
StaleLeasePolicy, StandaloneCwdIdentity, StandaloneListScope, StandaloneShutdownReason,
StandaloneStoreError, StandaloneWorkerRecord, StandaloneWorkerStatus, StandaloneWorkerStore,
};
+106 -143
View File
@@ -1,12 +1,11 @@
use std::fmt;
use std::fs::{self, File, OpenOptions};
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::str::FromStr;
use std::time::{SystemTime, UNIX_EPOCH};
use fs4::fs_std::FileExt;
use manifest::WorkerManifest;
use protocol::WorkerId;
use serde::{Deserialize, Serialize};
use session_store::{SegmentId, SessionId};
use thiserror::Error;
@@ -16,47 +15,10 @@ const RECORD_FILE: &str = "record.json";
const COMMIT_MARKER: &str = "commit.pending";
const LEASE_FILE: &str = "lease.json";
const LEASE_LOCK_FILE: &str = "lease.lock";
const SESSION_DIR: &str = "session";
const SESSIONS_DIR: &str = "sessions";
const WORKER_DIR: &str = "worker";
const SCHEMA_VERSION: u32 = 1;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct StandaloneSessionId(Uuid);
impl StandaloneSessionId {
#[must_use]
pub fn new() -> Self {
Self(Uuid::now_v7())
}
#[must_use]
pub fn short(self) -> String {
let simple = self.0.simple().to_string();
simple[simple.len() - 12..].to_string()
}
}
impl Default for StandaloneSessionId {
fn default() -> Self {
Self::new()
}
}
impl fmt::Display for StandaloneSessionId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt(formatter)
}
}
impl FromStr for StandaloneSessionId {
type Err = uuid::Error;
fn from_str(value: &str) -> Result<Self, Self::Err> {
Uuid::parse_str(value).map(Self)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct StandaloneCwdIdentity {
pub canonical_path: PathBuf,
@@ -100,7 +62,7 @@ impl StandaloneCwdIdentity {
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StandaloneSessionStatus {
pub enum StandaloneWorkerStatus {
Active,
Stopped,
}
@@ -115,17 +77,20 @@ pub enum StandaloneShutdownReason {
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StandaloneSessionRecord {
pub struct StandaloneWorkerRecord {
pub schema_version: u32,
pub revision: u64,
pub session_id: StandaloneSessionId,
pub worker_id: WorkerId,
/// User-facing Worker name resolved from the profile.
pub worker_name: String,
/// Internal key used by the current name-keyed Worker store.
pub storage_key: String,
pub cwd: StandaloneCwdIdentity,
pub manifest: WorkerManifest,
pub active_session_id: SessionId,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub active_segment_id: Option<SegmentId>,
pub status: StandaloneSessionStatus,
pub status: StandaloneWorkerStatus,
pub created_at_unix_ms: u64,
pub updated_at_unix_ms: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
@@ -145,11 +110,11 @@ pub enum StaleLeasePolicy {
}
#[derive(Debug, Clone)]
pub struct StandaloneSessionStore {
pub struct StandaloneWorkerStore {
root: PathBuf,
}
impl StandaloneSessionStore {
impl StandaloneWorkerStore {
pub fn open(root: impl Into<PathBuf>) -> Result<Self, StandaloneStoreError> {
let root = root.into();
fs::create_dir_all(&root).map_err(StandaloneStoreError::Io)?;
@@ -171,35 +136,41 @@ impl StandaloneSessionStore {
&self,
cwd: impl AsRef<Path>,
policy: StaleLeasePolicy,
) -> Result<StandaloneSessionAllocation, StandaloneStoreError> {
let id = StandaloneSessionId::new();
) -> Result<StandaloneWorkerAllocation, StandaloneStoreError> {
let worker_id = WorkerId::now_v7();
let cwd = StandaloneCwdIdentity::capture(cwd)?;
let dir = self.session_dir(id);
let dir = self.worker_dir(worker_id);
fs::create_dir(&dir).map_err(StandaloneStoreError::Io)?;
fs::create_dir(dir.join(SESSION_DIR)).map_err(StandaloneStoreError::Io)?;
fs::create_dir(dir.join(SESSIONS_DIR)).map_err(StandaloneStoreError::Io)?;
fs::create_dir(dir.join(WORKER_DIR)).map_err(StandaloneStoreError::Io)?;
let lease = self.acquire_lease(id, policy)?;
Ok(StandaloneSessionAllocation { id, cwd, lease })
let lease = self.acquire_lease(worker_id, policy)?;
Ok(StandaloneWorkerAllocation {
worker_id,
cwd,
lease,
})
}
pub fn commit_created(
&self,
allocation: &StandaloneSessionAllocation,
allocation: &StandaloneWorkerAllocation,
manifest: WorkerManifest,
storage_key: String,
active_session_id: SessionId,
active_segment_id: Option<SegmentId>,
) -> Result<StandaloneSessionRecord, StandaloneStoreError> {
) -> Result<StandaloneWorkerRecord, StandaloneStoreError> {
let now = now_unix_ms()?;
let record = StandaloneSessionRecord {
let record = StandaloneWorkerRecord {
schema_version: SCHEMA_VERSION,
revision: 1,
session_id: allocation.id,
worker_id: allocation.worker_id,
worker_name: manifest.worker.name.clone(),
storage_key,
cwd: allocation.cwd.clone(),
manifest,
active_session_id,
active_segment_id,
status: StandaloneSessionStatus::Active,
status: StandaloneWorkerStatus::Active,
created_at_unix_ms: now,
updated_at_unix_ms: now,
shutdown_reason: None,
@@ -208,22 +179,19 @@ impl StandaloneSessionStore {
Ok(record)
}
pub fn load(
&self,
id: StandaloneSessionId,
) -> Result<StandaloneSessionRecord, StandaloneStoreError> {
let dir = self.session_dir(id);
pub fn load(&self, id: WorkerId) -> Result<StandaloneWorkerRecord, StandaloneStoreError> {
let dir = self.worker_dir(id);
if dir.join(COMMIT_MARKER).exists() {
return Err(StandaloneStoreError::IncompleteCommit(id));
}
let bytes = fs::read(dir.join(RECORD_FILE)).map_err(|error| {
if error.kind() == io::ErrorKind::NotFound {
StandaloneStoreError::SessionNotFound(id)
StandaloneStoreError::WorkerNotFound(id)
} else {
StandaloneStoreError::Io(error)
}
})?;
let record: StandaloneSessionRecord = serde_json::from_slice(&bytes)
let record: StandaloneWorkerRecord = serde_json::from_slice(&bytes)
.map_err(|source| StandaloneStoreError::CorruptRecord { id, source })?;
if record.schema_version > SCHEMA_VERSION {
return Err(StandaloneStoreError::NewerSchema {
@@ -232,7 +200,7 @@ impl StandaloneSessionStore {
supported: SCHEMA_VERSION,
});
}
if record.schema_version != SCHEMA_VERSION || record.session_id != id {
if record.schema_version != SCHEMA_VERSION || record.worker_id != id {
return Err(StandaloneStoreError::InvalidRecord(id));
}
Ok(record)
@@ -243,7 +211,7 @@ impl StandaloneSessionStore {
cwd: impl AsRef<Path>,
scope: StandaloneListScope,
limit: usize,
) -> Result<Vec<StandaloneSessionRecord>, StandaloneStoreError> {
) -> Result<Vec<StandaloneWorkerRecord>, StandaloneStoreError> {
let current_cwd = (scope == StandaloneListScope::CurrentCwd)
.then(|| StandaloneCwdIdentity::capture(cwd))
.transpose()?;
@@ -269,12 +237,7 @@ impl StandaloneSessionStore {
right
.updated_at_unix_ms
.cmp(&left.updated_at_unix_ms)
.then_with(|| {
right
.session_id
.to_string()
.cmp(&left.session_id.to_string())
})
.then_with(|| right.worker_id.to_string().cmp(&left.worker_id.to_string()))
});
records.truncate(limit);
Ok(records)
@@ -282,10 +245,10 @@ impl StandaloneSessionStore {
pub fn acquire_lease(
&self,
id: StandaloneSessionId,
id: WorkerId,
policy: StaleLeasePolicy,
) -> Result<StandaloneSessionLease, StandaloneStoreError> {
let dir = self.session_dir(id);
) -> Result<StandaloneWorkerLease, StandaloneStoreError> {
let dir = self.worker_dir(id);
let path = dir.join(LEASE_FILE);
let _guard = LeaseMutationGuard::acquire(&dir)?;
let lease = LeaseRecord::current()?;
@@ -296,7 +259,7 @@ impl StandaloneSessionStore {
file.write_all(b"\n").map_err(StandaloneStoreError::Io)?;
file.sync_all().map_err(StandaloneStoreError::Io)?;
sync_directory(&dir)?;
return Ok(StandaloneSessionLease {
return Ok(StandaloneWorkerLease {
path,
lease_id: lease.lease_id,
released: false,
@@ -306,7 +269,7 @@ impl StandaloneSessionStore {
let existing = read_lease(&path, id)?;
match existing.liveness() {
LeaseLiveness::Live => {
return Err(StandaloneStoreError::SessionLeased(id));
return Err(StandaloneStoreError::WorkerLeased(id));
}
LeaseLiveness::Unknown => {
return Err(StandaloneStoreError::LeaseLivenessUnknown(id));
@@ -326,16 +289,16 @@ impl StandaloneSessionStore {
pub fn update_active_pointer(
&self,
record: &StandaloneSessionRecord,
record: &StandaloneWorkerRecord,
active_session_id: SessionId,
active_segment_id: Option<SegmentId>,
) -> Result<StandaloneSessionRecord, StandaloneStoreError> {
) -> Result<StandaloneWorkerRecord, StandaloneStoreError> {
let mut next = record.clone();
next.revision = next.revision.saturating_add(1);
next.updated_at_unix_ms = now_unix_ms()?;
next.active_session_id = active_session_id;
next.active_segment_id = active_segment_id;
next.status = StandaloneSessionStatus::Active;
next.status = StandaloneWorkerStatus::Active;
next.shutdown_reason = None;
self.commit_record(Some(record.revision), &next)?;
Ok(next)
@@ -343,73 +306,73 @@ impl StandaloneSessionStore {
pub fn mark_stopped(
&self,
record: &StandaloneSessionRecord,
record: &StandaloneWorkerRecord,
active_session_id: SessionId,
active_segment_id: Option<SegmentId>,
reason: StandaloneShutdownReason,
) -> Result<StandaloneSessionRecord, StandaloneStoreError> {
) -> Result<StandaloneWorkerRecord, StandaloneStoreError> {
let mut next = record.clone();
next.revision = next.revision.saturating_add(1);
next.updated_at_unix_ms = now_unix_ms()?;
next.active_session_id = active_session_id;
next.active_segment_id = active_segment_id;
next.status = StandaloneSessionStatus::Stopped;
next.status = StandaloneWorkerStatus::Stopped;
next.shutdown_reason = Some(reason);
self.commit_record(Some(record.revision), &next)?;
Ok(next)
}
pub fn delete(&self, id: StandaloneSessionId) -> Result<(), StandaloneStoreError> {
pub fn delete(&self, id: WorkerId) -> Result<(), StandaloneStoreError> {
let record = self.load(id)?;
if record.status != StandaloneSessionStatus::Stopped {
if record.status != StandaloneWorkerStatus::Stopped {
return Err(StandaloneStoreError::DeleteActive(id));
}
let session_dir = self.session_dir(id);
let _guard = LeaseMutationGuard::acquire(&session_dir)?;
let lease_path = session_dir.join(LEASE_FILE);
let worker_dir = self.worker_dir(id);
let _guard = LeaseMutationGuard::acquire(&worker_dir)?;
let lease_path = worker_dir.join(LEASE_FILE);
if lease_path.exists() {
let lease = read_lease(&lease_path, id)?;
return Err(match lease.liveness() {
LeaseLiveness::Live => StandaloneStoreError::SessionLeased(id),
LeaseLiveness::Live => StandaloneStoreError::WorkerLeased(id),
LeaseLiveness::Stale => StandaloneStoreError::StaleLease(id),
LeaseLiveness::Unknown => StandaloneStoreError::LeaseLivenessUnknown(id),
});
}
fs::remove_dir_all(self.session_dir(id)).map_err(StandaloneStoreError::Io)?;
fs::remove_dir_all(self.worker_dir(id)).map_err(StandaloneStoreError::Io)?;
sync_directory(&self.root)
}
#[must_use]
pub fn session_log_dir(&self, id: StandaloneSessionId) -> PathBuf {
self.session_dir(id).join(SESSION_DIR)
pub fn sessions_dir(&self, id: WorkerId) -> PathBuf {
self.worker_dir(id).join(SESSIONS_DIR)
}
#[must_use]
pub fn worker_metadata_dir(&self, id: StandaloneSessionId) -> PathBuf {
self.session_dir(id).join(WORKER_DIR)
pub fn worker_metadata_dir(&self, id: WorkerId) -> PathBuf {
self.worker_dir(id).join(WORKER_DIR)
}
#[must_use]
pub(crate) fn runtime_dir(&self, id: StandaloneSessionId) -> PathBuf {
self.session_dir(id).join("runtime")
pub(crate) fn runtime_dir(&self, id: WorkerId) -> PathBuf {
self.worker_dir(id).join("runtime")
}
pub(crate) fn abandon_allocation(
&self,
allocation: StandaloneSessionAllocation,
allocation: StandaloneWorkerAllocation,
) -> Result<(), StandaloneStoreError> {
let id = allocation.id;
let worker_id = allocation.worker_id;
allocation.lease.release()?;
fs::remove_dir_all(self.session_dir(id)).map_err(StandaloneStoreError::Io)?;
fs::remove_dir_all(self.worker_dir(worker_id)).map_err(StandaloneStoreError::Io)?;
sync_directory(&self.root)
}
fn commit_record(
&self,
expected_revision: Option<u64>,
next: &StandaloneSessionRecord,
next: &StandaloneWorkerRecord,
) -> Result<(), StandaloneStoreError> {
let dir = self.session_dir(next.session_id);
let dir = self.worker_dir(next.worker_id);
let marker = dir.join(COMMIT_MARKER);
let mut marker_file = OpenOptions::new()
.write(true)
@@ -417,7 +380,7 @@ impl StandaloneSessionStore {
.open(&marker)
.map_err(|error| {
if error.kind() == io::ErrorKind::AlreadyExists {
StandaloneStoreError::IncompleteCommit(next.session_id)
StandaloneStoreError::IncompleteCommit(next.worker_id)
} else {
StandaloneStoreError::Io(error)
}
@@ -427,11 +390,11 @@ impl StandaloneSessionStore {
sync_directory(&dir)?;
if let Some(expected) = expected_revision {
let current = self.load_record_while_committing(next.session_id)?;
let current = self.load_record_while_committing(next.worker_id)?;
if current.revision != expected {
let _ = fs::remove_file(&marker);
return Err(StandaloneStoreError::RevisionConflict {
id: next.session_id,
id: next.worker_id,
expected,
found: current.revision,
});
@@ -461,30 +424,30 @@ impl StandaloneSessionStore {
fn load_record_while_committing(
&self,
id: StandaloneSessionId,
) -> Result<StandaloneSessionRecord, StandaloneStoreError> {
id: WorkerId,
) -> Result<StandaloneWorkerRecord, StandaloneStoreError> {
let bytes =
fs::read(self.session_dir(id).join(RECORD_FILE)).map_err(StandaloneStoreError::Io)?;
fs::read(self.worker_dir(id).join(RECORD_FILE)).map_err(StandaloneStoreError::Io)?;
serde_json::from_slice(&bytes)
.map_err(|source| StandaloneStoreError::CorruptRecord { id, source })
}
fn session_dir(&self, id: StandaloneSessionId) -> PathBuf {
fn worker_dir(&self, id: WorkerId) -> PathBuf {
self.root.join(id.to_string())
}
}
#[derive(Debug)]
pub struct StandaloneSessionAllocation {
id: StandaloneSessionId,
pub struct StandaloneWorkerAllocation {
worker_id: WorkerId,
cwd: StandaloneCwdIdentity,
lease: StandaloneSessionLease,
lease: StandaloneWorkerLease,
}
impl StandaloneSessionAllocation {
impl StandaloneWorkerAllocation {
#[must_use]
pub fn id(&self) -> StandaloneSessionId {
self.id
pub fn worker_id(&self) -> WorkerId {
self.worker_id
}
#[must_use]
@@ -492,19 +455,19 @@ impl StandaloneSessionAllocation {
&self.cwd
}
pub fn into_lease(self) -> StandaloneSessionLease {
pub fn into_lease(self) -> StandaloneWorkerLease {
self.lease
}
}
#[derive(Debug)]
pub struct StandaloneSessionLease {
pub struct StandaloneWorkerLease {
path: PathBuf,
lease_id: Uuid,
released: bool,
}
impl StandaloneSessionLease {
impl StandaloneWorkerLease {
pub fn release(mut self) -> Result<(), StandaloneStoreError> {
self.release_inner()
}
@@ -534,7 +497,7 @@ impl StandaloneSessionLease {
}
}
impl Drop for StandaloneSessionLease {
impl Drop for StandaloneWorkerLease {
fn drop(&mut self) {
let _ = self.release_inner();
}
@@ -624,7 +587,7 @@ fn classify_lease_liveness(
}
}
fn read_lease(path: &Path, id: StandaloneSessionId) -> Result<LeaseRecord, StandaloneStoreError> {
fn read_lease(path: &Path, id: WorkerId) -> Result<LeaseRecord, StandaloneStoreError> {
let bytes = fs::read(path).map_err(StandaloneStoreError::Io)?;
serde_json::from_slice(&bytes)
.map_err(|source| StandaloneStoreError::CorruptLease { id, source })
@@ -692,47 +655,47 @@ pub enum StandaloneStoreError {
CwdUnavailable(#[source] io::Error),
#[error("standalone cwd is not a directory")]
CwdNotDirectory,
#[error("standalone cwd identity no longer matches the persisted session")]
#[error("standalone cwd identity no longer matches the persisted Worker")]
CwdIdentityMismatch,
#[error("standalone session {0} was not found")]
SessionNotFound(StandaloneSessionId),
#[error("standalone session {0} has an incomplete metadata commit")]
IncompleteCommit(StandaloneSessionId),
#[error("standalone session {0} has invalid metadata")]
InvalidRecord(StandaloneSessionId),
#[error("standalone session {id} metadata is corrupt")]
#[error("standalone Worker {0} was not found")]
WorkerNotFound(WorkerId),
#[error("standalone Worker {0} has an incomplete metadata commit")]
IncompleteCommit(WorkerId),
#[error("standalone Worker {0} has invalid metadata")]
InvalidRecord(WorkerId),
#[error("standalone Worker {id} metadata is corrupt")]
CorruptRecord {
id: StandaloneSessionId,
id: WorkerId,
#[source]
source: serde_json::Error,
},
#[error("standalone session {id} lease is corrupt")]
#[error("standalone Worker {id} lease is corrupt")]
CorruptLease {
id: StandaloneSessionId,
id: WorkerId,
#[source]
source: serde_json::Error,
},
#[error("standalone session {id} uses schema {found}, newer than supported schema {supported}")]
#[error("standalone Worker {id} uses schema {found}, newer than supported schema {supported}")]
NewerSchema {
id: StandaloneSessionId,
id: WorkerId,
found: u32,
supported: u32,
},
#[error("standalone session {0} is already active")]
SessionLeased(StandaloneSessionId),
#[error("standalone session {0} lease liveness cannot be proven; recovery is rejected")]
LeaseLivenessUnknown(StandaloneSessionId),
#[error("standalone session {0} has a stale lease; explicit recovery is required")]
StaleLease(StandaloneSessionId),
#[error("standalone session lease ownership changed")]
#[error("standalone Worker {0} is already active")]
WorkerLeased(WorkerId),
#[error("standalone Worker {0} lease liveness cannot be proven; recovery is rejected")]
LeaseLivenessUnknown(WorkerId),
#[error("standalone Worker {0} has a stale lease; explicit recovery is required")]
StaleLease(WorkerId),
#[error("standalone Worker lease ownership changed")]
LeaseOwnershipLost,
#[error("standalone session {0} must be stopped before deletion")]
DeleteActive(StandaloneSessionId),
#[error("standalone Worker {0} must be stopped before deletion")]
DeleteActive(WorkerId),
#[error(
"standalone session {id} metadata revision changed (expected {expected}, found {found})"
"standalone Worker {id} metadata revision changed (expected {expected}, found {found})"
)]
RevisionConflict {
id: StandaloneSessionId,
id: WorkerId,
expected: u64,
found: u64,
},
+157 -61
View File
@@ -8,11 +8,13 @@ use agen::llm_client::error::ClientError;
use agen::llm_client::event::{Event as LlmEvent, StopReason};
use agen::llm_client::types::Request;
use async_trait::async_trait;
use client::Client;
use client::transport::in_process::Socket as InProcessSocket;
use futures::{Stream, stream};
use protocol::{Event, Method};
use standalone::{
StaleLeasePolicy, StandaloneHost, StandaloneLaunchConfig, StandaloneListScope,
StandaloneSessionStatus, StandaloneSessionStore, StandaloneStartupError, StandaloneStoreError,
StandaloneStartupError, StandaloneStoreError, StandaloneWorkerStatus, StandaloneWorkerStore,
};
use uuid::Uuid;
@@ -88,17 +90,35 @@ async fn in_process_host_runs_text_and_read_tool_then_shuts_down() {
let host = StandaloneHost::start_with_model_client(launch, client)
.await
.expect("start in-process host");
let mut events = host.subscribe();
assert_eq!(host.record().worker_name, worker_name);
assert_eq!(host.record().manifest.worker.name, worker_name);
assert_eq!(
host.record().storage_key,
format!("standalone-{}", host.worker_id())
);
let mut protocol_client = host.connect();
host.send(Method::run_text("read the probe"))
protocol_client
.send(&Method::run_text("read the probe"))
.await
.expect("submit input");
tokio::time::timeout(Duration::from_secs(30), async {
let mut saw_user_message = false;
let mut saw_text = false;
let mut saw_tool_result = false;
loop {
match events.recv().await.expect("worker event") {
match protocol_client
.next_event()
.await
.expect("protocol event")
.expect("worker event")
{
Event::UserMessage { segments }
if format!("{segments:?}").contains("read the probe") =>
{
saw_user_message = true;
}
Event::TextDelta { text } if text.contains("standalone response") => {
saw_text = true;
}
@@ -106,6 +126,10 @@ async fn in_process_host_runs_text_and_read_tool_then_shuts_down() {
saw_tool_result = true;
}
Event::RunEnd { .. } => {
assert!(
saw_user_message,
"stream must expose the committed user message"
);
assert!(saw_text, "stream must expose the model text delta");
assert!(saw_tool_result, "stream must expose the tool result");
break;
@@ -140,6 +164,65 @@ async fn in_process_host_runs_text_and_read_tool_then_shuts_down() {
host.shutdown().await.expect("graceful shutdown");
}
#[tokio::test]
async fn startup_preserves_occupied_scope_conflict_details() {
let temp = tempfile::tempdir().expect("tempdir");
let cwd = temp.path().join("project");
std::fs::create_dir(&cwd).expect("create project");
let first_launch = StandaloneLaunchConfig::new(
&cwd,
temp.path().join("first-state"),
manifest::ProfileSelector::Default,
"first-worker",
)
.resolve()
.expect("resolve first launch");
let first_host =
StandaloneHost::start_with_model_client(first_launch, ScriptedClient::new(Vec::new()))
.await
.expect("start first host");
let competitor = first_host.record().storage_key.clone();
let second_launch = StandaloneLaunchConfig::new(
&cwd,
temp.path().join("second-state"),
manifest::ProfileSelector::Default,
"second-worker",
)
.resolve()
.expect("resolve second launch");
let error =
StandaloneHost::start_with_model_client(second_launch, ScriptedClient::new(Vec::new()))
.await
.err()
.expect("occupied scope rejected");
first_host.shutdown().await.expect("shutdown first host");
let canonical_cwd = cwd.canonicalize().expect("canonical cwd");
match &error {
StandaloneStartupError::ScopeConflict {
competitor: actual_competitor,
requested_rule,
competitor_rule,
} => {
assert_eq!(actual_competitor, &competitor);
assert_eq!(requested_rule.target, canonical_cwd);
assert_eq!(competitor_rule.target, canonical_cwd);
}
other => panic!("expected scope conflict, got {other:?}"),
}
assert_eq!(
error.to_string(),
format!(
"requested scope `{}` conflicts with worker allocation `{competitor}` rule `{}`",
canonical_cwd.display(),
canonical_cwd.display()
)
);
}
#[tokio::test]
async fn state_store_failure_is_redacted_and_starts_no_controller() {
let temp = tempfile::tempdir().expect("tempdir");
@@ -220,7 +303,7 @@ type TestResult = Result<(), Box<dyn std::error::Error>>;
async fn standalone_restore_preserves_history_tasks_notifications_and_cwd_scope() -> TestResult {
let temp = tempfile::tempdir()?;
let cwd = temp.path().join("project");
let state_dir = temp.path().join("client").join("standalone-sessions");
let state_dir = temp.path().join("client").join("standalone-workers");
std::fs::create_dir_all(&cwd)?;
let launch = StandaloneLaunchConfig::new(
&cwd,
@@ -250,23 +333,26 @@ async fn standalone_restore_preserves_history_tasks_notifications_and_cwd_scope(
],
]);
let host = StandaloneHost::start_with_model_client(launch, first_client).await?;
let session_id = host.session_id();
let mut events = host.subscribe();
host.send(Method::run_text("first request")).await?;
wait_for_run_end(&mut events).await?;
host.send(Method::Notify {
message: "persisted notification".to_string(),
auto_run: true,
})
.await?;
wait_for_run_end(&mut events).await?;
let worker_id = host.worker_id();
let mut protocol_client = host.connect();
protocol_client
.send(&Method::run_text("first request"))
.await?;
wait_for_run_end(&mut protocol_client).await?;
protocol_client
.send(&Method::Notify {
message: "persisted notification".to_string(),
auto_run: true,
})
.await?;
wait_for_run_end(&mut protocol_client).await?;
host.shutdown().await?;
let store = StandaloneSessionStore::open(&state_dir)?;
let store = StandaloneWorkerStore::open(&state_dir)?;
let current = store.list(&cwd, StandaloneListScope::CurrentCwd, 100)?;
assert_eq!(current.len(), 1);
assert_eq!(current[0].session_id, session_id);
assert_eq!(current[0].status, StandaloneSessionStatus::Stopped);
assert_eq!(current[0].worker_id, worker_id);
assert_eq!(current[0].status, StandaloneWorkerStatus::Stopped);
let other_cwd = temp.path().join("other");
std::fs::create_dir(&other_cwd)?;
assert!(
@@ -286,18 +372,31 @@ async fn standalone_restore_preserves_history_tasks_notifications_and_cwd_scope(
]]);
let second_inspection = second_client.clone();
let host =
StandaloneHost::restore_with_model_client(state_dir.clone(), session_id, second_client)
StandaloneHost::restore_with_model_client(state_dir.clone(), worker_id, second_client)
.await?;
let snapshot = format!("{:?}", host.snapshot());
assert_eq!(
host.record().worker_name,
"display-name-is-not-session-identity"
);
assert_eq!(host.record().storage_key, format!("standalone-{worker_id}"));
let mut protocol_client = host.connect();
let snapshot = format!(
"{:?}",
protocol_client
.next_event()
.await
.expect("restored protocol stream")
.expect("restored snapshot")
);
assert!(snapshot.contains("first request"), "{snapshot}");
assert!(snapshot.contains("first answer"), "{snapshot}");
assert!(snapshot.contains("persisted task"), "{snapshot}");
assert!(snapshot.contains("persisted notification"), "{snapshot}");
let mut events = host.subscribe();
host.send(Method::run_text("continue after restore"))
protocol_client
.send(&Method::run_text("continue after restore"))
.await?;
wait_for_run_end(&mut events).await?;
wait_for_run_end(&mut protocol_client).await?;
let request = second_inspection
.requests()
.into_iter()
@@ -309,11 +408,11 @@ async fn standalone_restore_preserves_history_tasks_notifications_and_cwd_scope(
assert!(projected.contains("persisted task"), "{projected}");
host.shutdown().await?;
store.delete(session_id)?;
store.delete(worker_id)?;
assert!(cwd.exists(), "deleting session state must not mutate cwd");
assert!(matches!(
store.load(session_id),
Err(StandaloneStoreError::SessionNotFound(_))
store.load(worker_id),
Err(StandaloneStoreError::WorkerNotFound(_))
));
Ok(())
}
@@ -334,28 +433,25 @@ async fn standalone_restore_rejects_concurrent_lease_and_missing_cwd() -> TestRe
.resolve()?;
let host =
StandaloneHost::start_with_model_client(launch, ScriptedClient::new(Vec::new())).await?;
let session_id = host.session_id();
let store = StandaloneSessionStore::open(&state_dir)?;
let worker_id = host.worker_id();
let store = StandaloneWorkerStore::open(&state_dir)?;
assert!(matches!(
store.acquire_lease(session_id, StaleLeasePolicy::Recover),
Err(StandaloneStoreError::SessionLeased(id)) if id == session_id
store.acquire_lease(worker_id, StaleLeasePolicy::Recover),
Err(StandaloneStoreError::WorkerLeased(id)) if id == worker_id
));
let restore = StandaloneHost::restore_with_model_client(
state_dir.clone(),
session_id,
worker_id,
ScriptedClient::new(Vec::new()),
)
.await;
assert!(matches!(
restore,
Err(StandaloneStartupError::SessionActive)
));
assert!(matches!(restore, Err(StandaloneStartupError::WorkerActive)));
host.shutdown().await?;
std::fs::rename(&cwd, &moved)?;
let restore = StandaloneHost::restore_with_model_client(
state_dir,
session_id,
worker_id,
ScriptedClient::new(Vec::new()),
)
.await;
@@ -392,11 +488,11 @@ async fn standalone_restore_recovers_only_a_proven_stale_lease() -> TestResult {
});
let host =
StandaloneHost::start_with_model_client(launch, ScriptedClient::new(Vec::new())).await?;
let session_id = host.session_id();
let worker_id = host.worker_id();
host.shutdown().await?;
let store = StandaloneSessionStore::open(&state_dir)?;
let store = StandaloneWorkerStore::open(&state_dir)?;
assert!(matches!(
store.load(session_id)?.manifest.profile,
store.load(worker_id)?.manifest.profile,
Some(manifest::ProfileManifestSnapshot {
source: manifest::ProfileSource::Registry {
source: manifest::ProfileRegistrySource::User,
@@ -405,9 +501,9 @@ async fn standalone_restore_recovers_only_a_proven_stale_lease() -> TestResult {
..
})
));
let session_dir = state_dir.join(session_id.to_string());
let worker_dir = state_dir.join(worker_id.to_string());
std::fs::write(
session_dir.join("lease.json"),
worker_dir.join("lease.json"),
serde_json::to_vec(&serde_json::json!({
"lease_id": uuid::Uuid::now_v7(),
"pid": u32::MAX,
@@ -418,7 +514,7 @@ async fn standalone_restore_recovers_only_a_proven_stale_lease() -> TestResult {
let host = StandaloneHost::restore_with_model_client(
state_dir,
session_id,
worker_id,
ScriptedClient::new(Vec::new()),
)
.await?;
@@ -439,11 +535,11 @@ async fn standalone_restore_rejects_lease_with_missing_start_marker() -> TestRes
.resolve()?;
let host =
StandaloneHost::start_with_model_client(launch, ScriptedClient::new(Vec::new())).await?;
let session_id = host.session_id();
let worker_id = host.worker_id();
host.shutdown().await?;
let session_dir = state_dir.join(session_id.to_string());
let worker_dir = state_dir.join(worker_id.to_string());
std::fs::write(
session_dir.join("lease.json"),
worker_dir.join("lease.json"),
serde_json::to_vec(&serde_json::json!({
"lease_id": uuid::Uuid::now_v7(),
"pid": std::process::id(),
@@ -451,14 +547,14 @@ async fn standalone_restore_rejects_lease_with_missing_start_marker() -> TestRes
}))?,
)?;
let store = StandaloneSessionStore::open(&state_dir)?;
let store = StandaloneWorkerStore::open(&state_dir)?;
assert!(matches!(
store.acquire_lease(session_id, StaleLeasePolicy::Recover),
Err(StandaloneStoreError::LeaseLivenessUnknown(id)) if id == session_id
store.acquire_lease(worker_id, StaleLeasePolicy::Recover),
Err(StandaloneStoreError::LeaseLivenessUnknown(id)) if id == worker_id
));
let restore = StandaloneHost::restore_with_model_client(
state_dir,
session_id,
worker_id,
ScriptedClient::new(Vec::new()),
)
.await;
@@ -482,31 +578,31 @@ async fn standalone_metadata_fails_closed_on_incomplete_or_newer_records() -> Te
.resolve()?;
let host =
StandaloneHost::start_with_model_client(launch, ScriptedClient::new(Vec::new())).await?;
let session_id = host.session_id();
let worker_id = host.worker_id();
host.shutdown().await?;
let store = StandaloneSessionStore::open(&state_dir)?;
let session_dir = state_dir.join(session_id.to_string());
std::fs::write(session_dir.join("commit.pending"), b"interrupted\n")?;
let store = StandaloneWorkerStore::open(&state_dir)?;
let worker_dir = state_dir.join(worker_id.to_string());
std::fs::write(worker_dir.join("commit.pending"), b"interrupted\n")?;
assert!(matches!(
store.load(session_id),
Err(StandaloneStoreError::IncompleteCommit(id)) if id == session_id
store.load(worker_id),
Err(StandaloneStoreError::IncompleteCommit(id)) if id == worker_id
));
std::fs::remove_file(session_dir.join("commit.pending"))?;
let record_path = session_dir.join("record.json");
std::fs::remove_file(worker_dir.join("commit.pending"))?;
let record_path = worker_dir.join("record.json");
let mut record: serde_json::Value = serde_json::from_slice(&std::fs::read(&record_path)?)?;
record["schema_version"] = serde_json::json!(u32::MAX);
std::fs::write(&record_path, serde_json::to_vec_pretty(&record)?)?;
assert!(matches!(
store.load(session_id),
Err(StandaloneStoreError::NewerSchema { id, .. }) if id == session_id
store.load(worker_id),
Err(StandaloneStoreError::NewerSchema { id, .. }) if id == worker_id
));
Ok(())
}
async fn wait_for_run_end(events: &mut tokio::sync::broadcast::Receiver<Event>) -> TestResult {
async fn wait_for_run_end(client: &mut Client<InProcessSocket>) -> TestResult {
tokio::time::timeout(Duration::from_secs(10), async {
loop {
if matches!(events.recv().await, Ok(Event::RunEnd { .. })) {
if matches!(client.next_event().await, Ok(Some(Event::RunEnd { .. }))) {
break;
}
}
+3
View File
@@ -489,6 +489,7 @@ pub struct NewTicket {
pub workflow_state: Option<TicketWorkflowState>,
pub queued_by: Option<String>,
pub queued_at: Option<String>,
#[serde(rename = "repository_key")]
pub repository_id: Option<String>,
pub ref_selector: Option<String>,
}
@@ -519,6 +520,7 @@ impl NewTicket {
#[serde(tag = "action", rename_all = "snake_case")]
pub enum TicketTargetEdit {
Set {
#[serde(rename = "repository_key")]
repository_id: String,
ref_selector: Option<String>,
},
@@ -1610,6 +1612,7 @@ pub struct TicketMeta {
pub workflow_state_explicit: bool,
pub queued_by: Option<String>,
pub queued_at: Option<String>,
#[serde(rename = "repository_key")]
pub repository_id: Option<String>,
pub ref_selector: Option<String>,
pub raw: BTreeMap<String, String>,
+7 -7
View File
@@ -402,8 +402,8 @@ struct TicketCreateParams {
queued_at: Option<String>,
/// Optional target Workspace repository id.
#[serde(default)]
repository_id: Option<String>,
/// Optional target Git ref selector. Requires `repository_id`.
repository_key: Option<String>,
/// Optional target Git ref selector. Requires `repository_key`.
#[serde(default)]
ref_selector: Option<String>,
}
@@ -944,7 +944,7 @@ impl Tool for TicketCreateTool {
input.workflow_state = params.state.map(TicketWorkflowStateParam::into_state);
input.queued_by = None;
input.queued_at = params.queued_at;
input.repository_id = params.repository_id;
input.repository_id = params.repository_key;
input.ref_selector = params.ref_selector;
let created = self
@@ -1173,7 +1173,7 @@ impl Tool for TicketMarkReadyTool {
json!({
"ticket": ticket.meta.id,
"state": ticket.meta.workflow_state.as_str(),
"repository_id": ticket.meta.repository_id,
"repository_key": ticket.meta.repository_id,
"ref_selector": ticket.meta.ref_selector,
"ok": true
}),
@@ -1206,7 +1206,7 @@ impl Tool for TicketIntakeReadyTool {
json!({
"ticket": ticket.meta.id,
"state": ticket.meta.workflow_state.as_str(),
"repository_id": ticket.meta.repository_id,
"repository_key": ticket.meta.repository_id,
"ref_selector": ticket.meta.ref_selector,
"ok": true
}),
@@ -1940,11 +1940,11 @@ mod tests {
fn resolve_target(
&self,
_workspace_id: &str,
repository_id: Option<&str>,
repository_key: Option<&str>,
ref_selector: Option<&str>,
) -> crate::Result<crate::ResolvedTicketTarget> {
Ok(crate::ResolvedTicketTarget {
repository_id: repository_id.unwrap_or("main").to_owned(),
repository_id: repository_key.unwrap_or("main").to_owned(),
ref_selector: ref_selector.unwrap_or("develop").to_owned(),
})
}
+134 -6
View File
@@ -21,6 +21,7 @@ struct BashParams {
pub(crate) struct BashTool {
session: WorkdirSessionHandle,
output_dir: PathBuf,
state: Arc<Mutex<BashExecutionState>>,
}
@@ -117,6 +118,7 @@ impl Tool for BashTool {
command: params.command,
timeout_secs,
output_limit: INLINE_BYTE_BUDGET,
spill_dir: Some(self.output_dir.clone()),
tool_call_id: Some(call_id.clone()),
})
.await
@@ -183,10 +185,15 @@ impl Tool for BashTool {
let content = if output.content.is_empty() {
None
} else if output.truncated {
Some(format!(
"[showing bounded WorkdirSession command output; additional output was truncated]\n{}",
output.content
))
let notice = match output.output_path {
Some(path) => format!(
"[showing bounded WorkdirSession command output; full output saved to {}]",
path.display()
),
None => "[showing bounded WorkdirSession command output; additional output was truncated]"
.to_owned(),
};
Some(format!("{notice}\n{}", output.content))
} else {
Some(output.content)
};
@@ -259,16 +266,137 @@ fn truncate_for_summary(command: &str) -> String {
summary
}
pub fn bash_tool(session: WorkdirSessionHandle, _output_dir: PathBuf) -> ToolDefinition {
pub fn bash_tool(session: WorkdirSessionHandle, output_dir: PathBuf) -> ToolDefinition {
Arc::new(move || {
let schema = schemars::schema_for!(BashParams);
let meta = ToolMeta::new("Bash")
.description("Execute a shell command in the bound Workdir. Process start, bounded output, timeout and cancellation are owned by the WorkdirSession provider. This is not a sandbox.")
.description("Execute a shell command in the bound Workdir. Process start, bounded inline output, full-output spill, timeout and cancellation are owned by the WorkdirSession provider. This is not a sandbox.")
.input_schema(serde_json::to_value(schema).expect("Bash schema serialization"));
let tool: Arc<dyn Tool> = Arc::new(BashTool {
session: session.clone(),
output_dir: output_dir.clone(),
state: Arc::new(Mutex::new(BashExecutionState::default())),
});
(meta, tool)
})
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use manifest::{Permission, Scope, ScopeConfig, ScopeRule};
use tempfile::TempDir;
use workdir::{LocalWorkdirSession, WorkdirSessionHandle};
use super::bash_tool;
use crate::{grep::grep_tool, read::read_tool, tracker::Tracker};
fn session_with_output_scope(root: &TempDir, output: &TempDir) -> WorkdirSessionHandle {
let scope = Scope::from_config(&ScopeConfig {
allow: vec![
ScopeRule {
target: root.path().to_path_buf(),
permission: Permission::Write,
recursive: true,
},
ScopeRule {
target: output.path().to_path_buf(),
permission: Permission::Read,
recursive: true,
},
],
deny: Vec::new(),
})
.unwrap();
Arc::new(LocalWorkdirSession::new(scope, root.path().to_path_buf()))
}
#[tokio::test]
async fn long_output_is_spilled_and_available_to_read_and_grep() {
let root = TempDir::new().unwrap();
let output = TempDir::new().unwrap();
let session = session_with_output_scope(&root, &output);
let (_, bash) = bash_tool(session.clone(), output.path().to_path_buf())();
let command = "i=0; while [ $i -lt 2000 ]; do printf 'line-%04d\\n' \"$i\"; i=$((i+1)); done; printf 'FINAL-NEEDLE\\n'";
let result = bash
.execute(
&serde_json::json!({ "command": command }).to_string(),
Default::default(),
)
.await
.unwrap();
let rendered = result.content.expect("bounded Bash output");
let artifact = std::fs::read_dir(output.path())
.unwrap()
.next()
.expect("artifact entry")
.unwrap()
.path();
assert!(rendered.contains("full output saved to"));
assert!(rendered.contains(&artifact.display().to_string()));
let retained = std::fs::read_to_string(&artifact).unwrap();
assert!(retained.starts_with("line-0000\n"));
assert!(retained.ends_with("FINAL-NEEDLE\n"));
assert_eq!(retained.lines().count(), 2001);
let (_, read) = read_tool(session.clone(), Tracker::new())();
let read_result = read
.execute(
&serde_json::json!({
"file_path": artifact,
"offset": 2000,
"limit": 1,
})
.to_string(),
Default::default(),
)
.await
.unwrap();
assert!(
read_result
.content
.expect("Read content")
.contains("FINAL-NEEDLE")
);
let (_, grep) = grep_tool(session)();
let grep_result = grep
.execute(
&serde_json::json!({
"pattern": "FINAL-NEEDLE",
"path": artifact,
"output_mode": "content",
})
.to_string(),
Default::default(),
)
.await
.unwrap();
let grep_content = grep_result.content.expect("Grep content");
assert!(
grep_content.contains("FINAL-NEEDLE"),
"unexpected Grep content: {grep_content:?}"
);
}
#[tokio::test]
async fn short_output_does_not_leave_a_spill_artifact() {
let root = TempDir::new().unwrap();
let output = TempDir::new().unwrap();
let session = session_with_output_scope(&root, &output);
let (_, bash) = bash_tool(session, output.path().to_path_buf())();
let result = bash
.execute(
&serde_json::json!({ "command": "printf short" }).to_string(),
Default::default(),
)
.await
.unwrap();
assert_eq!(result.content.as_deref(), Some("short"));
assert_eq!(std::fs::read_dir(output.path()).unwrap().count(), 0);
}
}
+1
View File
@@ -42,6 +42,7 @@ impl From<ToolsError> for ToolError {
workdir::WorkdirError::NotFound(_)
| workdir::WorkdirError::Io { .. }
| workdir::WorkdirError::Unavailable(_)
| workdir::WorkdirError::OperationFailed
| workdir::WorkdirError::Transport(_),
) => ToolError::ExecutionFailed(err.to_string()),
ToolsError::FileSystem(_)
+2 -2
View File
@@ -22,7 +22,7 @@ enum OutputMode {
#[derive(Debug, Deserialize, JsonSchema)]
struct GrepParams {
pattern: String,
/// Logical Workdir-relative file or directory to search. Defaults to the Workdir root.
/// Workdir-relative path, or an absolute path covered by readable scope. Defaults to the Workdir root.
#[serde(default)]
path: Option<String>,
#[serde(default)]
@@ -61,7 +61,7 @@ impl Tool for GrepTool {
let params: GrepParams = serde_json::from_str(input_json)
.map_err(|error| ToolError::InvalidArgument(format!("invalid Grep input: {error}")))?;
let path = match params.path {
Some(path) => WorkdirPath::new(&path).map_err(ToolsError::from)?,
Some(path) => WorkdirPath::new_scoped(&path).map_err(ToolsError::from)?,
None => WorkdirPath::root(),
};
let mode = match params.output_mode.unwrap_or_default() {
+3 -3
View File
@@ -13,14 +13,14 @@ use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle};
const DESCRIPTION: &str = "Read a text file from the local filesystem. \
Supports offset/limit for large files. Returns line-numbered output (1-based). \
Directories cannot be read. The file must be read before Write or Edit can \
modify it. Paths are relative to the bound Workdir.";
modify it. Paths are Workdir-relative unless an absolute path is explicitly readable.";
const DEFAULT_LIMIT: usize = 2000;
const PROVIDER_BYTE_LIMIT: usize = 256 * 1024;
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub(crate) struct ReadParams {
/// Logical path relative to the bound Workdir root.
/// Workdir-relative path, or an absolute path covered by readable scope.
pub file_path: String,
/// 0-based line offset from the start. Defaults to 0.
#[serde(default)]
@@ -47,7 +47,7 @@ impl Tool for ReadTool {
let offset = params.offset.unwrap_or(0);
let limit = params.limit.unwrap_or(DEFAULT_LIMIT).max(1);
let path = WorkdirPath::new(&params.file_path).map_err(ToolsError::from)?;
let path = WorkdirPath::new_scoped(&params.file_path).map_err(ToolsError::from)?;
tracing::debug!(path = %path, offset, limit, "Read");
let result = self
+8 -5
View File
@@ -224,20 +224,23 @@ async fn very_long_single_line() {
}
#[tokio::test]
async fn absolute_path_is_rejected() {
let (dir, _spill, reg) = setup();
async fn absolute_path_requires_matching_read_scope() {
let (_dir, _spill, reg) = setup();
let outside = tempfile::tempdir().unwrap();
let outside_file = outside.path().join("outside.txt");
std::fs::write(&outside_file, "secret").unwrap();
let read = reg.get("Read");
let err = read
.execute(
&json!({ "file_path": dir.path().join("outside.txt") }).to_string(),
&json!({ "file_path": outside_file }).to_string(),
Default::default(),
)
.await
.unwrap_err();
let msg = format!("{err}");
assert!(
msg.contains("invalid logical filesystem path"),
"absolute path was not rejected as invalid: {msg}"
msg.contains("outside allowed scope"),
"absolute path escaped readable scope: {msg}"
);
}
+10 -3
View File
@@ -394,14 +394,21 @@ async fn bash_inherits_workdir_cwd() {
}
#[tokio::test]
async fn bash_provider_output_does_not_expose_internal_paths() {
async fn bash_provider_output_exposes_readable_retained_path() {
let (_dir, spill, reg) = setup();
let bash = reg.get("Bash");
let out = call(&bash, json!({ "command": "printf 'x%.0s' {1..20480}" })).await;
let body = out.content.unwrap();
assert!(body.contains("bounded WorkdirSession command output"));
assert!(!body.contains(spill.path().to_str().unwrap()));
assert_eq!(std::fs::read_dir(spill.path()).unwrap().count(), 0);
assert!(body.contains("full output saved to"));
assert!(body.contains(spill.path().to_str().unwrap()));
let artifact = std::fs::read_dir(spill.path())
.unwrap()
.next()
.expect("retained output")
.unwrap()
.path();
assert_eq!(std::fs::metadata(artifact).unwrap().len(), 20_480);
}
#[tokio::test]
+1
View File
@@ -30,4 +30,5 @@ pulldown-cmark = { version = "0.13.3", default-features = false }
agen.workspace = true
[dev-dependencies]
async-trait.workspace = true
tempfile = { workspace = true }
+51 -5
View File
@@ -249,6 +249,9 @@ pub struct App {
pub running: bool,
/// True while the Worker is in `WorkerStatus::Paused`.
pub paused: bool,
/// Local observation time for the current run. Used only for live UI
/// elapsed time and spinner animation; it is not persisted in history.
pub run_started_at: Option<Instant>,
pub run_requests: usize,
/// Sum of `input_tokens - cache_read_input_tokens` across the
/// current turn's LLM requests — i.e. the net tokens this turn
@@ -281,6 +284,9 @@ pub struct App {
/// records the instant; a second press within the timeout exits the
/// TUI (the Worker itself stays alive).
pub quit_confirm: Option<std::time::Instant>,
/// Independent 2-tap guard for `Ctrl-X` when the Worker is idle or
/// stopped. A second press within the timeout shuts down the Worker.
pub shutdown_confirm: Option<std::time::Instant>,
/// Full display history in render order.
pub blocks: Vec<Block>,
/// Turn/protocol errors retained when a real `SegmentStart` replaces the
@@ -352,6 +358,7 @@ impl App {
worker_status: WorkerStatus::Idle,
running: false,
paused: false,
run_started_at: None,
run_requests: 0,
run_upload_tokens: 0,
run_output_tokens: 0,
@@ -369,6 +376,7 @@ impl App {
command_completion_selected: None,
quit: false,
quit_confirm: None,
shutdown_confirm: None,
blocks: Vec::new(),
run_error_messages: Vec::new(),
internal_workers: Vec::new(),
@@ -553,11 +561,18 @@ impl App {
}
pub fn set_worker_status(&mut self, status: WorkerStatus) {
let was_running = self.running;
self.worker_status = status;
self.running = status == WorkerStatus::Running;
self.paused = status == WorkerStatus::Paused;
if self.running {
if !was_running {
self.run_started_at = Some(Instant::now());
}
self.quit_confirm = None;
self.shutdown_confirm = None;
} else {
self.run_started_at = None;
}
}
@@ -763,6 +778,20 @@ impl App {
Some(self.method_for_run(segments))
}
pub fn restore_unsent_run(&mut self, method: &Method) {
let Method::Run { input } = method else {
return;
};
self.pending_submit_rollback = None;
if self.input.is_empty() {
self.input.replace_with_segments(input);
self.completion = None;
} else {
self.queued_inputs
.push_front(QueuedInput::new(input.clone()));
}
}
fn method_for_run(&mut self, segments: Vec<Segment>) -> Method {
// TurnHeader / UserMessage blocks are pushed only after the Worker
// emits `Event::UserMessage` from a committed `LogEntry::AnnotatedUserInput`.
@@ -913,6 +942,10 @@ impl App {
Some(self.method_for_run(queued.segments))
}
pub fn clear_actionbar_notice(&mut self) {
self.actionbar_notice = None;
}
pub fn push_error(&mut self, message: impl Into<String>) {
self.blocks.push(Block::Alert {
level: AlertLevel::Error,
@@ -1121,11 +1154,13 @@ impl App {
self.latest_llm_wait_event = None;
self.assistant_streaming = false;
}
// UI consumers of Invoke / LlmCall semantics are out of scope
// for `tickets/invoke-turn-llmcall-semantics.md`; events flow
// through to subscribers but the TUI currently derives its
// turn header from `UserMessage` / `SystemItem` arrivals.
Event::InvokeStart { .. } | Event::LlmCallStart { .. } | Event::LlmCallEnd { .. } => {
Event::InvokeStart { .. } => {
self.set_worker_status(WorkerStatus::Running);
}
// UI consumers of per-attempt LlmCall semantics remain out of scope;
// the run-level status starts at InvokeStart and TurnStart counts each
// LLM request within that run.
Event::LlmCallStart { .. } | Event::LlmCallEnd { .. } => {
self.latest_llm_wait_event = None;
}
Event::LlmRetry {
@@ -3377,6 +3412,17 @@ mod completion_flow_tests {
}
}
#[test]
fn running_status_starts_and_stops_live_run_clock() {
let mut app = App::new("test".into());
app.set_worker_status(WorkerStatus::Running);
assert!(app.run_started_at.is_some());
app.set_worker_status(WorkerStatus::Idle);
assert!(app.run_started_at.is_none());
}
#[test]
fn running_submit_is_queued_locally_and_clears_composer() {
let mut app = App::new("test".into());
+28 -48
View File
@@ -7,15 +7,15 @@ use client::{
list_backend_workers, restore_backend_worker,
};
use crossterm::event::{self, Event as TermEvent, KeyCode, KeyEventKind, KeyModifiers};
use ratatui::backend::CrosstermBackend;
use ratatui::Frame;
use ratatui::layout::{Constraint, Layout};
use ratatui::style::{Color, Modifier, Style};
use ratatui::text::{Line, Span};
use ratatui::widgets::Paragraph;
use ratatui::{Frame, Terminal, TerminalOptions, Viewport};
use crate::backend_workspace_picker::select_backend_workspace;
use crate::console;
use crate::inline_terminal::with_inline_terminal;
const MAX_ROWS: usize = 10;
const VIEWPORT_LINES: u16 = MAX_ROWS as u16 + 4;
@@ -127,31 +127,32 @@ fn pick_worker(
workers.truncate(MAX_ROWS);
let mut state = BackendWorkerPickerState::new(target, workers);
let mut terminal = make_inline_terminal()?;
loop {
terminal.draw(|frame| draw(frame, &state))?;
match poll_event()? {
None => continue,
Some(Action::Up) => state.previous(),
Some(Action::Down) => state.next(),
Some(Action::Submit) => {
close_viewport(&mut terminal)?;
return Ok(WorkerPickerResult::Selected(
state.selected_worker().clone(),
));
with_inline_terminal(
VIEWPORT_LINES,
|terminal| -> Result<_, Box<dyn std::error::Error>> {
loop {
terminal.draw(|frame| draw(frame, &state))?;
match poll_event()? {
None => continue,
Some(Action::Up) => state.previous(),
Some(Action::Down) => state.next(),
Some(Action::Submit) => {
return Ok(WorkerPickerResult::Selected(
state.selected_worker().clone(),
));
}
Some(Action::SwitchWorkspace) => {
return Ok(WorkerPickerResult::SwitchWorkspace);
}
Some(Action::Cancel) => {
return Err(Box::new(io::Error::other(
"Backend worker picker cancelled",
)));
}
}
}
Some(Action::SwitchWorkspace) => {
close_viewport(&mut terminal)?;
return Ok(WorkerPickerResult::SwitchWorkspace);
}
Some(Action::Cancel) => {
close_viewport(&mut terminal)?;
return Err(Box::new(io::Error::other(
"Backend worker picker cancelled",
)));
}
}
}
},
)
}
struct BackendWorkerPickerState {
@@ -184,27 +185,6 @@ impl BackendWorkerPickerState {
}
}
fn make_inline_terminal() -> io::Result<Terminal<CrosstermBackend<io::Stdout>>> {
let backend = CrosstermBackend::new(io::stdout());
Terminal::with_options(
backend,
TerminalOptions {
viewport: Viewport::Inline(VIEWPORT_LINES),
},
)
}
fn close_viewport(terminal: &mut Terminal<CrosstermBackend<io::Stdout>>) -> io::Result<()> {
let area = terminal.get_frame().area();
let last_row = area.bottom().saturating_sub(1);
terminal.set_cursor_position((0, last_row))?;
use std::io::Write;
let mut out = io::stdout();
out.write_all(b"\r\n")?;
out.flush()?;
Ok(())
}
enum Action {
Up,
Down,
@@ -370,7 +350,7 @@ fn working_directory_text(worker: &BackendWorkerSummary) -> String {
let cleanliness = wd.cleanliness.as_deref().unwrap_or("unknown");
format!(
"wd:{}:{} {} {}",
wd.repository_id, wd.working_directory_id, wd.status, cleanliness
wd.repository_key, wd.working_directory_id, wd.status, cleanliness
)
}
File diff suppressed because it is too large Load Diff
+109
View File
@@ -0,0 +1,109 @@
use std::io::{self, Stdout, Write};
use ratatui::Terminal;
use ratatui::backend::CrosstermBackend;
use ratatui::{TerminalOptions, Viewport};
pub(crate) type InlineTerminal = Terminal<CrosstermBackend<Stdout>>;
struct InlineTerminalGuard {
terminal: InlineTerminal,
closed: bool,
}
impl InlineTerminalGuard {
fn open(height: u16) -> io::Result<Self> {
let terminal = Terminal::with_options(
CrosstermBackend::new(io::stdout()),
TerminalOptions {
viewport: Viewport::Inline(height),
},
)?;
Ok(Self {
terminal,
closed: false,
})
}
fn close(&mut self) -> io::Result<()> {
if self.closed {
return Ok(());
}
self.closed = true;
let area = self.terminal.get_frame().area();
let last_row = area.bottom().saturating_sub(1);
let cursor_result = self.terminal.set_cursor_position((0, last_row));
let output_result = write_viewport_terminator(&mut io::stdout());
cursor_result?;
output_result
}
}
impl Drop for InlineTerminalGuard {
fn drop(&mut self) {
let _ = self.close();
}
}
pub(crate) fn with_inline_terminal<T, E>(
height: u16,
run: impl FnOnce(&mut InlineTerminal) -> Result<T, E>,
) -> Result<T, E>
where
E: From<io::Error>,
{
let mut guard = InlineTerminalGuard::open(height).map_err(E::from)?;
let result = run(&mut guard.terminal);
let close_result = guard.close();
match result {
Ok(value) => {
close_result.map_err(E::from)?;
Ok(value)
}
Err(error) => Err(error),
}
}
fn write_viewport_terminator(output: &mut impl Write) -> io::Result<()> {
output.write_all(b"\r\n")?;
output.flush()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn viewport_terminator_moves_following_output_to_a_fresh_line() {
let mut output = Vec::new();
write_viewport_terminator(&mut output).unwrap();
assert_eq!(output, b"\r\n");
}
#[test]
fn inline_viewport_construction_is_owned_by_this_module() {
fn assert_shared_owner(path: &std::path::Path) {
for entry in std::fs::read_dir(path).unwrap() {
let path = entry.unwrap().path();
if path.is_dir() {
assert_shared_owner(&path);
} else if path.extension().and_then(|value| value.to_str()) == Some("rs")
&& path.file_name().and_then(|value| value.to_str())
!= Some("inline_terminal.rs")
{
let source = std::fs::read_to_string(&path).unwrap();
assert!(
!source.contains("Viewport::Inline"),
"{} constructs an inline viewport outside its shared owner",
path.display()
);
}
}
}
assert_shared_owner(&std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("src"));
}
}
+388 -74
View File
@@ -15,6 +15,64 @@ use ratatui::style::{Color, Style};
use ratatui::text::{Line, Span};
use unicode_width::UnicodeWidthChar;
pub const MAX_PLAIN_TEXT_PASTE_CHARS: usize = 50;
pub const MAX_PLAIN_TEXT_PASTE_LOGICAL_LINES: usize = 3;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PasteMeasurement {
pub chars: usize,
pub logical_lines: usize,
}
impl PasteMeasurement {
pub fn presentation(self) -> PastePresentation {
if self.chars <= MAX_PLAIN_TEXT_PASTE_CHARS
&& self.logical_lines <= MAX_PLAIN_TEXT_PASTE_LOGICAL_LINES
{
PastePresentation::Text
} else {
PastePresentation::Chip
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PastePresentation {
Text,
Chip,
}
pub fn measure_paste(content: &str) -> PasteMeasurement {
PasteMeasurement {
chars: content.chars().count(),
logical_lines: logical_line_count(content),
}
}
/// Empty content has zero logical lines. Otherwise LF, lone CR, and CRLF each
/// advance one line; a CRLF pair is one break rather than two.
pub fn logical_line_count(content: &str) -> usize {
if content.is_empty() {
return 0;
}
let mut lines = 1;
let mut chars = content.chars().peekable();
while let Some(ch) = chars.next() {
match ch {
'\r' => {
if chars.peek() == Some(&'\n') {
chars.next();
}
lines += 1;
}
'\n' => lines += 1,
_ => {}
}
}
lines
}
#[derive(Debug, Clone)]
pub struct PasteRef {
pub id: u32,
@@ -61,6 +119,7 @@ impl FlowRefAtom {
pub enum Atom {
Char(char),
Paste(PasteRef),
PasteArtifact(protocol::PasteArtifactRef),
FileRef(FileRefAtom),
FlowRef(FlowRefAtom),
}
@@ -72,6 +131,18 @@ impl Atom {
match self {
Atom::Char(_) => None,
Atom::Paste(p) => Some((Style::default().fg(Color::Magenta), p.label())),
Atom::PasteArtifact(artifact) => Some((
Style::default().fg(Color::Magenta),
format!(
"[Paste artifact {} | {} chars, {} lines, {}, {}, created {} ms]",
artifact.artifact_id,
artifact.char_count,
artifact.line_count,
artifact.media_type.as_str(),
artifact.availability.as_str(),
artifact.created_at_ms
),
)),
Atom::FileRef(r) => Some((Style::default().fg(Color::Cyan), r.label())),
Atom::FlowRef(r) => Some((Style::default().fg(Color::Yellow), r.label())),
}
@@ -102,7 +173,9 @@ enum WordKind {
fn atom_class(atom: &Atom) -> AtomClass {
match atom {
Atom::Char(c) => char_class(*c),
Atom::Paste(_) | Atom::FileRef(_) | Atom::FlowRef(_) => AtomClass::Chip,
Atom::Paste(_) | Atom::PasteArtifact(_) | Atom::FileRef(_) | Atom::FlowRef(_) => {
AtomClass::Chip
}
}
}
@@ -190,6 +263,16 @@ impl InputBuffer {
content: content.clone(),
}));
}
protocol::Segment::PasteArtifact { artifact } => {
self.atoms.push(Atom::PasteArtifact(artifact.clone()));
}
protocol::Segment::UploadedFile { file } => {
self.atoms.extend(
format!("[Attached file: {}]", file.file_name)
.chars()
.map(Atom::Char),
);
}
protocol::Segment::FileRef { path } => {
self.atoms
.push(Atom::FileRef(FileRefAtom { path: path.clone() }));
@@ -225,6 +308,13 @@ impl InputBuffer {
match atom {
Atom::Char(c) => text.push(*c),
Atom::Paste(paste) => text.push_str(&paste.content),
Atom::PasteArtifact(artifact) => {
text.push_str(&protocol::Segment::flatten_to_text(&[
protocol::Segment::PasteArtifact {
artifact: artifact.clone(),
},
]))
}
Atom::FileRef(file) => text.push_str(&file.path),
Atom::FlowRef(flow) => text.push_str(&flow.selector),
}
@@ -237,16 +327,20 @@ impl InputBuffer {
}
pub fn insert_paste(&mut self, content: String) {
let measurement = measure_paste(&content);
if measurement.presentation() == PastePresentation::Text {
self.insert_str(&content);
return;
}
let id = self.next_paste_id;
self.next_paste_id = self.next_paste_id.wrapping_add(1);
let chars = content.chars().count();
let lines = content.lines().count().max(1);
self.atoms.insert(
self.cursor,
Atom::Paste(PasteRef {
id,
chars,
lines,
chars: measurement.chars,
lines: measurement.logical_lines,
content,
}),
);
@@ -395,80 +489,78 @@ impl InputBuffer {
self.cursor = 0;
}
pub fn move_home(&mut self) {
while self.cursor > 0 {
if matches!(self.atoms[self.cursor - 1], Atom::Char('\n')) {
break;
}
self.cursor -= 1;
fn logical_line_ranges(&self) -> Vec<(usize, usize)> {
let mut ranges = Vec::new();
let mut start = 0;
let mut index = 0;
while index < self.atoms.len() {
let break_len = match self.atoms[index] {
Atom::Char('\r') => {
if matches!(self.atoms.get(index + 1), Some(Atom::Char('\n'))) {
2
} else {
1
}
}
Atom::Char('\n') => 1,
_ => {
index += 1;
continue;
}
};
ranges.push((start, index));
index += break_len;
start = index;
}
ranges.push((start, self.atoms.len()));
ranges
}
fn logical_line_and_col(&self) -> (Vec<(usize, usize)>, usize, usize) {
let ranges = self.logical_line_ranges();
for (line, &(start, end)) in ranges.iter().enumerate() {
if self.cursor <= end {
return (ranges, line, self.cursor.saturating_sub(start));
}
if let Some(&(next_start, _)) = ranges.get(line + 1)
&& self.cursor < next_start
{
return (ranges, line + 1, 0);
}
}
let line = ranges.len().saturating_sub(1);
let col = self.cursor.saturating_sub(ranges[line].0);
(ranges, line, col)
}
pub fn move_home(&mut self) {
let (ranges, line, _) = self.logical_line_and_col();
self.cursor = ranges[line].0;
}
pub fn move_end(&mut self) {
while self.cursor < self.atoms.len() {
if matches!(self.atoms[self.cursor], Atom::Char('\n')) {
break;
}
self.cursor += 1;
}
let (ranges, line, _) = self.logical_line_and_col();
self.cursor = ranges[line].1;
}
/// Move one logical line up, preserving column (atom count from
/// current line start). No-op if already on the first line.
pub fn move_up(&mut self) {
let (line_start, col) = self.line_start_and_col();
if line_start == 0 {
let (ranges, line, col) = self.logical_line_and_col();
if line == 0 {
return;
}
// `atoms[line_start - 1]` is the '\n' that opens the current
// line; find the previous line's start.
let prev_end = line_start - 1;
let mut prev_start = 0;
for i in (0..prev_end).rev() {
if matches!(self.atoms[i], Atom::Char('\n')) {
prev_start = i + 1;
break;
}
}
let prev_len = prev_end - prev_start;
self.cursor = prev_start + col.min(prev_len);
let (start, end) = ranges[line - 1];
self.cursor = start + col.min(end - start);
}
/// Move one logical line down, preserving column.
pub fn move_down(&mut self) {
let (line_start, col) = self.line_start_and_col();
// End of current line.
let mut cur_end = self.atoms.len();
for i in line_start..self.atoms.len() {
if matches!(self.atoms[i], Atom::Char('\n')) {
cur_end = i;
break;
}
}
if cur_end == self.atoms.len() {
return; // no next line
}
let next_start = cur_end + 1;
let mut next_end = self.atoms.len();
for i in next_start..self.atoms.len() {
if matches!(self.atoms[i], Atom::Char('\n')) {
next_end = i;
break;
}
}
let next_len = next_end - next_start;
self.cursor = next_start + col.min(next_len);
}
fn line_start_and_col(&self) -> (usize, usize) {
let mut start = 0;
for i in (0..self.cursor).rev() {
if matches!(self.atoms[i], Atom::Char('\n')) {
start = i + 1;
break;
}
}
(start, self.cursor - start)
let (ranges, line, col) = self.logical_line_and_col();
let Some(&(start, end)) = ranges.get(line + 1) else {
return;
};
self.cursor = start + col.min(end - start);
}
/// Build the typed `Vec<Segment>` sent over the protocol. Adjacent
@@ -497,6 +589,12 @@ impl InputBuffer {
content: p.content.clone(),
});
}
Atom::PasteArtifact(artifact) => {
flush_text(&mut buf, &mut out);
out.push(protocol::Segment::PasteArtifact {
artifact: artifact.clone(),
});
}
Atom::FileRef(r) => {
flush_text(&mut buf, &mut out);
out.push(protocol::Segment::FileRef {
@@ -535,6 +633,7 @@ impl InputBuffer {
let mut cursor_row: u16 = 0;
let mut cursor_col: u16 = 0;
let mut cursor_set = false;
let mut previous_was_cr = false;
// Record cursor once, at the point right before `atom` would be
// placed — accounting for a wrap that the atom itself will cause.
@@ -558,7 +657,7 @@ impl InputBuffer {
for (i, atom) in self.atoms.iter().enumerate() {
if !cursor_set && i == self.cursor {
let leading = match atom {
Atom::Char('\n') => 0,
Atom::Char('\n' | '\r') => 0,
Atom::Char(c) => UnicodeWidthChar::width(*c).unwrap_or(0),
other => other
.chip()
@@ -573,6 +672,21 @@ impl InputBuffer {
}
match atom {
Atom::Char('\r') => {
flush_pending(
&mut pending,
&mut pending_width,
pending_style,
&mut rows,
&mut row_width,
);
rows.push(Vec::new());
row_width = 0;
previous_was_cr = true;
}
Atom::Char('\n') if previous_was_cr => {
previous_was_cr = false;
}
Atom::Char('\n') => {
flush_pending(
&mut pending,
@@ -583,8 +697,10 @@ impl InputBuffer {
);
rows.push(Vec::new());
row_width = 0;
previous_was_cr = false;
}
Atom::Char(c) => {
previous_was_cr = false;
let cw = UnicodeWidthChar::width(*c).unwrap_or(0);
if pending_style != text_style && !pending.is_empty() {
flush_pending(
@@ -608,6 +724,7 @@ impl InputBuffer {
);
}
other => {
previous_was_cr = false;
let (chip_style, label) = other.chip().expect("non-char atom has a chip");
if pending_style != chip_style && !pending.is_empty() {
flush_pending(
@@ -848,6 +965,161 @@ mod render_viewport_tests {
}
}
#[cfg(test)]
mod paste_policy_tests {
use super::*;
use protocol::Segment;
use serde::Deserialize;
#[derive(Debug, Deserialize)]
struct Fixture {
max_plain_text_chars: usize,
max_plain_text_logical_lines: usize,
cases: Vec<FixtureCase>,
}
#[derive(Debug, Deserialize)]
struct FixtureCase {
name: String,
parts: Vec<FixturePart>,
char_count: usize,
logical_line_count: usize,
presentation: FixturePresentation,
}
#[derive(Debug, Deserialize)]
struct FixturePart {
value: String,
repeat: usize,
}
#[derive(Debug, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
enum FixturePresentation {
Text,
Chip,
}
fn fixture() -> Fixture {
serde_json::from_str(include_str!(
"../../../tests/fixtures/composer-paste-policy.json"
))
.expect("shared composer paste policy fixture must be valid")
}
fn fixture_content(case: &FixtureCase) -> String {
case.parts
.iter()
.map(|part| part.value.repeat(part.repeat))
.collect()
}
#[test]
fn tui_follows_shared_paste_presentation_contract() {
let fixture = fixture();
assert_eq!(fixture.max_plain_text_chars, MAX_PLAIN_TEXT_PASTE_CHARS);
assert_eq!(
fixture.max_plain_text_logical_lines,
MAX_PLAIN_TEXT_PASTE_LOGICAL_LINES
);
for case in fixture.cases {
let content = fixture_content(&case);
let measurement = measure_paste(&content);
let expected_presentation = match case.presentation {
FixturePresentation::Text => PastePresentation::Text,
FixturePresentation::Chip => PastePresentation::Chip,
};
assert_eq!(measurement.chars, case.char_count, "{} chars", case.name);
assert_eq!(
measurement.logical_lines, case.logical_line_count,
"{} logical lines",
case.name
);
assert_eq!(
measurement.presentation(),
expected_presentation,
"{} presentation",
case.name
);
}
}
#[test]
fn short_paste_is_editable_text_at_the_cursor() {
let mut buffer = InputBuffer::new();
buffer.insert_str("ac");
buffer.move_left();
buffer.insert_paste("b".to_owned());
assert_eq!(buffer.plain_text(), "abc");
assert!(
buffer
.atoms
.iter()
.all(|atom| matches!(atom, Atom::Char(_)))
);
assert_eq!(
buffer.submit_segments(),
vec![Segment::text("abc".to_owned())]
);
}
#[test]
fn short_multiline_paste_preserves_original_line_endings_as_text() {
let content = "ab\r\ncd\ref";
let mut buffer = InputBuffer::new();
buffer.insert_paste(content.to_owned());
assert_eq!(buffer.plain_text(), content);
assert!(
buffer
.atoms
.iter()
.all(|atom| matches!(atom, Atom::Char(_)))
);
assert_eq!(
buffer.submit_segments(),
vec![Segment::text(content.to_owned())]
);
let rendered: Vec<String> = buffer
.render(80)
.lines
.iter()
.map(|line| {
line.spans
.iter()
.map(|span| span.content.as_ref())
.collect()
})
.collect();
assert_eq!(rendered, vec!["ab", "cd", "ef"]);
buffer.move_up();
assert_eq!(buffer.cursor, 6);
buffer.move_up();
assert_eq!(buffer.cursor, 2);
buffer.move_down();
assert_eq!(buffer.cursor, 6);
buffer.move_home();
assert_eq!(buffer.cursor, 4);
buffer.move_end();
assert_eq!(buffer.cursor, 6);
}
#[test]
fn empty_paste_is_a_noop() {
let mut buffer = InputBuffer::new();
buffer.insert_str("unchanged");
let paste_id = buffer.next_paste_id;
buffer.insert_paste(String::new());
assert_eq!(buffer.plain_text(), "unchanged");
assert_eq!(buffer.next_paste_id, paste_id);
}
}
#[cfg(test)]
mod submit_segments_tests {
use super::*;
@@ -873,7 +1145,8 @@ mod submit_segments_tests {
for c in "see ".chars() {
buf.insert_char(c);
}
buf.insert_paste("line1\nline2".into());
let pasted = "line1\nline2\nline3\nline4";
buf.insert_paste(pasted.into());
for c in " end".chars() {
buf.insert_char(c);
}
@@ -890,9 +1163,9 @@ mod submit_segments_tests {
content,
..
} => {
assert_eq!(content, "line1\nline2");
assert_eq!(*chars, "line1\nline2".chars().count() as u32);
assert_eq!(*lines, 2);
assert_eq!(content, pasted);
assert_eq!(*chars, pasted.chars().count() as u32);
assert_eq!(*lines, 4);
}
other => panic!("expected Paste, got {other:?}"),
}
@@ -902,6 +1175,45 @@ mod submit_segments_tests {
}
}
#[test]
fn restored_direct_paste_remains_a_typed_segment_without_reclassification() {
let original = Segment::Paste {
id: 7,
chars: 1,
lines: 1,
content: "x".to_owned(),
};
let mut buf = InputBuffer::new();
buf.replace_with_segments(std::slice::from_ref(&original));
assert_eq!(buf.submit_segments(), vec![original]);
}
#[test]
fn restored_paste_artifact_remains_a_typed_segment() {
let artifact = protocol::PasteArtifactRef {
artifact_id: "019ca7c8-57b6-7f05-8edf-524147aba7b2".to_string(),
created_at_ms: 1_700_000_000_000,
media_type: protocol::PasteArtifactMediaType::TextPlainUtf8,
availability: protocol::PasteArtifactAvailability::Available,
byte_len: 65_536,
char_count: 65_530,
line_count: 200,
sha256: "a".repeat(64),
source_entry_id: "entry-1".to_string(),
};
let original = Segment::PasteArtifact {
artifact: artifact.clone(),
};
let mut buf = InputBuffer::new();
buf.replace_with_segments(std::slice::from_ref(&original));
assert_eq!(
buf.submit_segments(),
vec![Segment::PasteArtifact { artifact }]
);
}
#[test]
fn empty_buffer_yields_empty_segments() {
let buf = InputBuffer::new();
@@ -911,7 +1223,7 @@ mod submit_segments_tests {
#[test]
fn leading_paste_does_not_emit_empty_text() {
let mut buf = InputBuffer::new();
buf.insert_paste("X".into());
buf.insert_paste("X".repeat(MAX_PLAIN_TEXT_PASTE_CHARS + 1));
let segs = buf.submit_segments();
assert_eq!(segs.len(), 1);
assert!(matches!(segs[0], Segment::Paste { .. }));
@@ -1011,7 +1323,7 @@ mod completion_prefix_tests {
#[test]
fn trigger_after_chip_atom() {
let mut buf = InputBuffer::new();
buf.insert_paste("X".into());
buf.insert_paste("X".repeat(MAX_PLAIN_TEXT_PASTE_CHARS + 1));
for c in "@sr".chars() {
buf.insert_char(c);
}
@@ -1120,7 +1432,7 @@ mod word_motion_tests {
for c in "foo ".chars() {
buf.insert_char(c);
}
buf.insert_paste("anything".into());
buf.insert_paste("anything".repeat(MAX_PLAIN_TEXT_PASTE_CHARS + 1));
for c in " bar".chars() {
buf.insert_char(c);
}
@@ -1219,7 +1531,9 @@ mod word_motion_tests {
for a in &buf.atoms {
match a {
Atom::Char(c) => out.push(*c),
Atom::Paste(_) | Atom::FileRef(_) | Atom::FlowRef(_) => out.push_str("<P>"),
Atom::Paste(_) | Atom::PasteArtifact(_) | Atom::FileRef(_) | Atom::FlowRef(_) => {
out.push_str("<P>")
}
}
}
out
@@ -1277,7 +1591,7 @@ mod word_motion_tests {
for c in "foo ".chars() {
buf.insert_char(c);
}
buf.insert_paste("anything".into());
buf.insert_paste("anything".repeat(MAX_PLAIN_TEXT_PASTE_CHARS + 1));
for c in " bar".chars() {
buf.insert_char(c);
}
+5 -34
View File
@@ -1,17 +1,17 @@
use std::io::{self, Stdout, Write};
use std::process::ExitCode;
use std::time::Duration;
use crossterm::event::{self, Event, KeyCode, KeyEventKind, KeyModifiers};
use crossterm::terminal::{disable_raw_mode, enable_raw_mode};
use ratatui::backend::CrosstermBackend;
use ratatui::Frame;
use ratatui::layout::{Constraint, Layout};
use ratatui::style::{Color, Modifier, Style};
use ratatui::text::{Line, Span};
use ratatui::widgets::Paragraph;
use ratatui::{Frame, Terminal, TerminalOptions, Viewport};
use secrets::{SecretStore, SecretValue};
use crate::inline_terminal::{InlineTerminal, with_inline_terminal};
#[derive(Debug, Clone, PartialEq, Eq)]
enum Mode {
Normal,
@@ -235,7 +235,6 @@ pub async fn launch() -> ExitCode {
}
type UiResult<T> = Result<T, Box<dyn std::error::Error>>;
type InlineTerminal = Terminal<CrosstermBackend<Stdout>>;
const MAX_ROWS: usize = 10;
const VIEWPORT_LINES: u16 = MAX_ROWS as u16 + 5;
@@ -270,37 +269,9 @@ impl Drop for RawModeGuard {
fn run(store: SecretStore) -> UiResult<()> {
enable_raw_mode()?;
let guard = RawModeGuard::new();
let mut terminal = make_inline_terminal()?;
let result = run_loop(&mut terminal, store);
let close_result = close_viewport(&mut terminal);
drop(terminal);
let result = with_inline_terminal(VIEWPORT_LINES, |terminal| run_loop(terminal, store));
guard.restore();
result?;
close_result?;
Ok(())
}
fn make_inline_terminal() -> io::Result<InlineTerminal> {
let backend = CrosstermBackend::new(io::stdout());
Terminal::with_options(
backend,
TerminalOptions {
viewport: Viewport::Inline(VIEWPORT_LINES),
},
)
}
/// Park the cursor at the very bottom of the inline viewport and emit one
/// newline before dropping the terminal. This matches the resume picker and
/// keeps the shell prompt (or a later inline viewport) from drawing over rows.
fn close_viewport(terminal: &mut InlineTerminal) -> io::Result<()> {
let area = terminal.get_frame().area();
let last_row = area.bottom().saturating_sub(1);
terminal.set_cursor_position((0, last_row))?;
let mut out = io::stdout();
out.write_all(b"\r\n")?;
out.flush()?;
Ok(())
result
}
fn run_loop(terminal: &mut InlineTerminal, store: SecretStore) -> UiResult<()> {
+19 -13
View File
@@ -10,12 +10,14 @@ mod composer_keys;
mod console;
#[cfg(feature = "e2e-test")]
mod e2e_observer;
mod inline_terminal;
mod input;
pub mod keys;
mod markdown;
mod scroll;
pub mod setup_model;
mod standalone_picker;
mod standalone_spawn;
mod task;
mod text_selection;
mod tool;
@@ -46,8 +48,8 @@ pub enum LaunchMode {
worker_name: Option<String>,
profile: Option<String>,
},
/// Restore one client-owned standalone session. The current cwd is the default scope;
/// `include_all` opts into all standalone sessions under the same client data root.
/// Restore one client-owned standalone Worker. The current cwd is the default scope;
/// `include_all` opts into all standalone Workers under the same client data root.
StandaloneResume { include_all: bool },
/// List Backend Workers and attach to the selected Worker.
Workers {
@@ -136,17 +138,21 @@ pub async fn launch(options: LaunchOptions) -> ExitCode {
LaunchMode::Spawn {
worker_name,
profile,
} => match target.spawn_worker() {
Ok(spawn) => {
console::run_standalone(
workspace_root.clone(),
spawn.state_dir,
worker_name,
profile,
)
.await
}
Err(e) => Err(Box::new(e) as Box<dyn std::error::Error>),
} => match standalone_spawn::select(&workspace_root, worker_name, profile) {
Ok(Some(selection)) => match target.spawn_worker() {
Ok(spawn) => {
console::run_standalone(
workspace_root.clone(),
spawn.state_dir,
Some(selection.worker_name),
Some(selection.profile),
)
.await
}
Err(error) => Err(Box::new(error) as Box<dyn std::error::Error>),
},
Ok(None) => Ok(()),
Err(error) => Err(Box::new(error) as Box<dyn std::error::Error>),
},
LaunchMode::StandaloneResume { include_all } => {
match standalone_picker::pick(target.as_ref(), include_all) {
+51 -54
View File
@@ -1,44 +1,43 @@
use std::io;
use std::time::Duration;
use client::{StandaloneSessionListIntent, StandaloneSessionResumeIntent, Target};
use client::{StandaloneWorkerListIntent, StandaloneWorkerResumeIntent, Target};
use crossterm::event::{self, Event as TermEvent, KeyCode, KeyEventKind, KeyModifiers};
use ratatui::Terminal;
use ratatui::backend::CrosstermBackend;
use ratatui::layout::{Constraint, Layout};
use ratatui::prelude::{Color, Line, Modifier, Span, Style};
use ratatui::widgets::Paragraph;
use ratatui::{TerminalOptions, Viewport};
use standalone::{StandaloneListScope, StandaloneSessionRecord, StandaloneSessionStore};
use standalone::{StandaloneListScope, StandaloneWorkerRecord, StandaloneWorkerStore};
use thiserror::Error;
use crate::inline_terminal::with_inline_terminal;
const LIMIT: usize = 100;
pub(crate) fn pick(
target: &dyn Target,
include_all: bool,
) -> Result<Option<StandaloneSessionResumeIntent>, StandalonePickerError> {
) -> Result<Option<StandaloneWorkerResumeIntent>, StandalonePickerError> {
let intent = target
.standalone_session_list(include_all)
.standalone_worker_list(include_all)
.map_err(StandalonePickerError::Target)?;
let records = load_records(&intent)?;
if records.is_empty() {
return Err(StandalonePickerError::NoSessions { include_all });
return Err(StandalonePickerError::NoWorkers { include_all });
}
let selected = run_picker(records)?;
selected
.map(|record| {
target
.standalone_session_resume(record.session_id.to_string())
.standalone_worker_resume(record.worker_id.to_string())
.map_err(StandalonePickerError::Target)
})
.transpose()
}
fn load_records(
intent: &StandaloneSessionListIntent,
) -> Result<Vec<StandaloneSessionRecord>, StandalonePickerError> {
let store = StandaloneSessionStore::open(&intent.state_dir)
intent: &StandaloneWorkerListIntent,
) -> Result<Vec<StandaloneWorkerRecord>, StandalonePickerError> {
let store = StandaloneWorkerStore::open(&intent.state_dir)
.map_err(StandalonePickerError::StateStore)?;
store
.list(
@@ -54,54 +53,49 @@ fn load_records(
}
fn run_picker(
records: Vec<StandaloneSessionRecord>,
) -> Result<Option<StandaloneSessionRecord>, StandalonePickerError> {
records: Vec<StandaloneWorkerRecord>,
) -> Result<Option<StandaloneWorkerRecord>, StandalonePickerError> {
let height = u16::try_from(records.len().saturating_add(3).min(20)).unwrap_or(20);
let mut terminal = Terminal::with_options(
CrosstermBackend::new(io::stdout()),
TerminalOptions {
viewport: Viewport::Inline(height),
},
)
.map_err(StandalonePickerError::Io)?;
let mut selected = 0usize;
loop {
terminal
.draw(|frame| draw(frame, &records, selected))
.map_err(StandalonePickerError::Io)?;
if !event::poll(Duration::from_millis(100)).map_err(StandalonePickerError::Io)? {
continue;
}
let TermEvent::Key(key) = event::read().map_err(StandalonePickerError::Io)? else {
continue;
};
if key.kind == KeyEventKind::Release {
continue;
}
let ctrl = key.modifiers.contains(KeyModifiers::CONTROL);
match key.code {
KeyCode::Up | KeyCode::Char('k') if !ctrl => {
selected = selected.saturating_sub(1);
with_inline_terminal(height, |terminal| {
let mut selected = 0usize;
loop {
terminal
.draw(|frame| draw(frame, &records, selected))
.map_err(StandalonePickerError::Io)?;
if !event::poll(Duration::from_millis(100)).map_err(StandalonePickerError::Io)? {
continue;
}
KeyCode::Down | KeyCode::Char('j') if !ctrl => {
selected = (selected + 1).min(records.len() - 1);
let TermEvent::Key(key) = event::read().map_err(StandalonePickerError::Io)? else {
continue;
};
if key.kind == KeyEventKind::Release {
continue;
}
let ctrl = key.modifiers.contains(KeyModifiers::CONTROL);
match key.code {
KeyCode::Up | KeyCode::Char('k') if !ctrl => {
selected = selected.saturating_sub(1);
}
KeyCode::Down | KeyCode::Char('j') if !ctrl => {
selected = (selected + 1).min(records.len() - 1);
}
KeyCode::Enter => return Ok(Some(records[selected].clone())),
KeyCode::Esc => return Ok(None),
KeyCode::Char('c') if ctrl => return Ok(None),
_ => {}
}
KeyCode::Enter => return Ok(Some(records[selected].clone())),
KeyCode::Esc => return Ok(None),
KeyCode::Char('c') if ctrl => return Ok(None),
_ => {}
}
}
})
}
fn draw(frame: &mut ratatui::Frame<'_>, records: &[StandaloneSessionRecord], selected: usize) {
fn draw(frame: &mut ratatui::Frame<'_>, records: &[StandaloneWorkerRecord], selected: usize) {
let mut constraints = vec![Constraint::Length(1)];
constraints.extend(records.iter().map(|_| Constraint::Length(1)));
constraints.push(Constraint::Length(1));
let rows = Layout::vertical(constraints).split(frame.area());
frame.render_widget(
Paragraph::new(Line::from(Span::styled(
"resume standalone session",
"resume standalone Worker",
Style::default().add_modifier(Modifier::BOLD),
))),
rows[0],
@@ -120,7 +114,10 @@ fn draw(frame: &mut ratatui::Frame<'_>, records: &[StandaloneSessionRecord], sel
frame.render_widget(
Paragraph::new(Line::from(vec![
Span::raw(marker),
Span::styled(record.session_id.short(), style),
Span::styled(
format!("{} ({})", record.worker_name, record.worker_id.short()),
style,
),
Span::raw(format!(
" [{:?}] updated:{} {}",
record.status, record.updated_at_unix_ms, cwd
@@ -139,14 +136,14 @@ fn draw(frame: &mut ratatui::Frame<'_>, records: &[StandaloneSessionRecord], sel
pub(crate) enum StandalonePickerError {
#[error("standalone target error: {0}")]
Target(#[source] client::TargetError),
#[error("standalone session state is unavailable: {0}")]
#[error("standalone Worker state is unavailable: {0}")]
StateStore(#[source] standalone::StandaloneStoreError),
#[error(
"no standalone sessions found for this cwd; use `yoi --local --resume --all` to include all cwd identities"
"no standalone Workers found for this cwd; use `yoi --local resume --all` to include all cwd identities"
)]
NoSessions { include_all: bool },
#[error("standalone session picker I/O failed: {0}")]
Io(#[source] io::Error),
NoWorkers { include_all: bool },
#[error("standalone Worker picker I/O failed: {0}")]
Io(#[from] io::Error),
}
#[cfg(test)]
+488
View File
@@ -0,0 +1,488 @@
use std::io;
use std::path::Path;
use std::time::Duration;
use crossterm::event::{self, Event, KeyCode, KeyEvent, KeyEventKind, KeyModifiers};
use manifest::ProfileDiscovery;
use ratatui::layout::{Constraint, Direction, Layout};
use ratatui::style::{Color, Modifier, Style};
use ratatui::text::{Line, Span};
use ratatui::widgets::Paragraph;
use thiserror::Error;
use crate::inline_terminal::{InlineTerminal, with_inline_terminal};
const VIEWPORT_HEIGHT: u16 = 6;
const FALLBACK_WORKER_NAME: &str = "worker";
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct StandaloneSpawnSelection {
pub worker_name: String,
pub profile: String,
}
#[derive(Debug, Error)]
pub(crate) enum StandaloneSpawnError {
#[error("profile discovery failed: {0}")]
ProfileDiscovery(#[from] manifest::ProfileError),
#[error("no profiles are available")]
NoProfiles,
#[error("standalone spawn picker terminal error: {0}")]
Terminal(#[from] io::Error),
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct ProfileChoice {
selector: String,
label: String,
is_default: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StatusKind {
Info,
Progress,
Error,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SpawnAction {
None,
Submit,
Cancel,
}
struct SpawnForm {
worker_name: String,
cursor: usize,
profile_choices: Vec<ProfileChoice>,
selected_profile: usize,
status: Option<(String, StatusKind)>,
}
impl SpawnForm {
fn new(
worker_name: Option<String>,
default_worker_name: String,
profile_choices: Vec<ProfileChoice>,
) -> Self {
let worker_name = worker_name.unwrap_or(default_worker_name);
let cursor = worker_name.chars().count();
let selected_profile = profile_choices
.iter()
.position(|choice| choice.is_default)
.unwrap_or(0);
Self {
worker_name,
cursor,
profile_choices,
selected_profile,
status: None,
}
}
fn selected_profile(&self) -> &ProfileChoice {
&self.profile_choices[self.selected_profile]
}
fn apply_key(&mut self, key: KeyEvent) -> SpawnAction {
if key.kind == KeyEventKind::Release {
return SpawnAction::None;
}
if key.modifiers.contains(KeyModifiers::CONTROL) {
match key.code {
KeyCode::Char('c') | KeyCode::Char('u') => return SpawnAction::Cancel,
_ => return SpawnAction::None,
}
}
self.status = None;
match key.code {
KeyCode::Esc => SpawnAction::Cancel,
KeyCode::Enter => {
if self.worker_name.trim().is_empty() {
self.status =
Some(("worker name cannot be empty".to_owned(), StatusKind::Error));
SpawnAction::None
} else {
SpawnAction::Submit
}
}
KeyCode::Tab | KeyCode::Down => {
self.selected_profile = (self.selected_profile + 1) % self.profile_choices.len();
SpawnAction::None
}
KeyCode::BackTab | KeyCode::Up => {
self.selected_profile = if self.selected_profile == 0 {
self.profile_choices.len() - 1
} else {
self.selected_profile - 1
};
SpawnAction::None
}
KeyCode::Left => {
self.cursor = self.cursor.saturating_sub(1);
SpawnAction::None
}
KeyCode::Right => {
self.cursor = (self.cursor + 1).min(self.worker_name.chars().count());
SpawnAction::None
}
KeyCode::Home => {
self.cursor = 0;
SpawnAction::None
}
KeyCode::End => {
self.cursor = self.worker_name.chars().count();
SpawnAction::None
}
KeyCode::Backspace => {
if self.cursor > 0 {
let idx = byte_index(&self.worker_name, self.cursor - 1);
self.worker_name.remove(idx);
self.cursor -= 1;
}
SpawnAction::None
}
KeyCode::Delete => {
if self.cursor < self.worker_name.chars().count() {
let idx = byte_index(&self.worker_name, self.cursor);
self.worker_name.remove(idx);
}
SpawnAction::None
}
KeyCode::Char(ch) if is_safe_worker_char(ch) => {
let idx = byte_index(&self.worker_name, self.cursor);
self.worker_name.insert(idx, ch);
self.cursor += 1;
SpawnAction::None
}
_ => SpawnAction::None,
}
}
}
pub(crate) fn select(
workspace_root: &Path,
worker_name: Option<String>,
profile: Option<String>,
) -> Result<Option<StandaloneSpawnSelection>, StandaloneSpawnError> {
let default_worker_name = default_worker_name(workspace_root);
if let Some(profile) = profile {
return Ok(Some(StandaloneSpawnSelection {
worker_name: worker_name.unwrap_or(default_worker_name),
profile,
}));
}
let registry = ProfileDiscovery::user_settings().discover()?;
let choices = profile_choices(&registry);
if choices.is_empty() {
return Err(StandaloneSpawnError::NoProfiles);
}
with_inline_terminal(VIEWPORT_HEIGHT, |terminal| {
run_picker(
terminal,
SpawnForm::new(worker_name, default_worker_name, choices),
)
})
}
fn run_picker(
terminal: &mut InlineTerminal,
mut form: SpawnForm,
) -> Result<Option<StandaloneSpawnSelection>, StandaloneSpawnError> {
loop {
terminal.draw(|frame| draw_form(frame, &form))?;
if !event::poll(Duration::from_millis(100))? {
continue;
}
let Event::Key(key) = event::read()? else {
continue;
};
match form.apply_key(key) {
SpawnAction::None => {}
SpawnAction::Cancel => {
form.status = Some(("cancelled".to_owned(), StatusKind::Info));
terminal.draw(|frame| draw_form(frame, &form))?;
return Ok(None);
}
SpawnAction::Submit => {
let selection = StandaloneSpawnSelection {
worker_name: form.worker_name.trim().to_owned(),
profile: form.selected_profile().selector.clone(),
};
form.status = Some(("starting worker...".to_owned(), StatusKind::Progress));
terminal.draw(|frame| draw_form(frame, &form))?;
return Ok(Some(selection));
}
}
}
}
fn profile_choices(registry: &manifest::ProfileRegistry) -> Vec<ProfileChoice> {
registry
.entries()
.iter()
.map(|entry| {
let selector = entry.qualified_name();
let default_marker = if entry.is_default { " (default)" } else { "" };
let mut label = format!("{selector}{default_marker}");
if let Some(description) = &entry.description {
label.push_str("");
label.push_str(description);
}
ProfileChoice {
selector,
label,
is_default: entry.is_default,
}
})
.collect()
}
fn draw_form(frame: &mut ratatui::Frame<'_>, form: &SpawnForm) {
let chunks = Layout::default()
.direction(Direction::Vertical)
.constraints([
Constraint::Length(1),
Constraint::Length(1),
Constraint::Length(1),
Constraint::Length(1),
Constraint::Length(1),
Constraint::Min(0),
])
.split(frame.area());
frame.render_widget(
Paragraph::new(Line::from(vec![
Span::raw(" "),
Span::styled(
"spawn worker",
Style::default().add_modifier(Modifier::BOLD),
),
])),
chunks[0],
);
frame.render_widget(
Paragraph::new(Line::from(vec![
Span::raw(" "),
Span::styled("name: ", Style::default().fg(Color::DarkGray)),
Span::styled(
&form.worker_name,
Style::default()
.fg(Color::Cyan)
.add_modifier(Modifier::BOLD),
),
])),
chunks[1],
);
frame.render_widget(
Paragraph::new(Line::from(vec![
Span::raw(" "),
Span::styled("profile: ", Style::default().fg(Color::DarkGray)),
Span::styled(
&form.selected_profile().label,
Style::default().fg(Color::Green),
),
Span::styled(
" (tab/down to change)",
Style::default().fg(Color::DarkGray),
),
])),
chunks[2],
);
frame.render_widget(
Paragraph::new(Line::from(Span::styled(
" enter spawn · left/right edit · esc cancel",
Style::default().fg(Color::DarkGray),
))),
chunks[3],
);
let (message, color) = form
.status
.as_ref()
.map(|(message, kind)| {
let color = match kind {
StatusKind::Info => Color::DarkGray,
StatusKind::Progress => Color::Yellow,
StatusKind::Error => Color::Red,
};
(message.as_str(), color)
})
.unwrap_or(("", Color::Reset));
frame.render_widget(
Paragraph::new(Line::from(vec![
Span::raw(" "),
Span::styled(message, Style::default().fg(color)),
])),
chunks[4],
);
let prefix_width = " name: ".chars().count() as u16;
let x = chunks[1]
.x
.saturating_add(prefix_width)
.saturating_add(form.cursor as u16)
.min(chunks[1].right().saturating_sub(1));
frame.set_cursor_position((x, chunks[1].y));
}
fn default_worker_name(workspace_root: &Path) -> String {
workspace_root
.file_name()
.and_then(|name| name.to_str())
.map(sanitise_default_name)
.filter(|name| !name.is_empty())
.unwrap_or_else(|| FALLBACK_WORKER_NAME.to_owned())
}
fn sanitise_default_name(name: &str) -> String {
name.chars()
.map(|ch| if is_safe_worker_char(ch) { ch } else { '-' })
.collect()
}
fn is_safe_worker_char(ch: char) -> bool {
ch.is_ascii_alphanumeric() || matches!(ch, '-' | '_' | '.')
}
fn byte_index(input: &str, char_index: usize) -> usize {
input
.char_indices()
.nth(char_index)
.map_or(input.len(), |(idx, _)| idx)
}
#[cfg(test)]
mod tests {
use crossterm::event::{KeyEvent, KeyModifiers};
use super::*;
fn choices() -> Vec<ProfileChoice> {
vec![
ProfileChoice {
selector: "builtin:default".to_owned(),
label: "builtin:default (default) — Default".to_owned(),
is_default: true,
},
ProfileChoice {
selector: "builtin:coder".to_owned(),
label: "builtin:coder — Coder".to_owned(),
is_default: false,
},
]
}
#[test]
fn default_form_preserves_old_spawn_layout_defaults() {
let form = SpawnForm::new(None, "yoi".to_owned(), choices());
assert_eq!(form.worker_name, "yoi");
assert_eq!(form.selected_profile().selector, "builtin:default");
}
#[test]
fn tab_and_arrows_cycle_profiles() {
let mut form = SpawnForm::new(None, "yoi".to_owned(), choices());
assert_eq!(
form.apply_key(KeyEvent::new(KeyCode::Tab, KeyModifiers::NONE)),
SpawnAction::None
);
assert_eq!(form.selected_profile().selector, "builtin:coder");
form.apply_key(KeyEvent::new(KeyCode::Down, KeyModifiers::NONE));
assert_eq!(form.selected_profile().selector, "builtin:default");
form.apply_key(KeyEvent::new(KeyCode::Up, KeyModifiers::NONE));
assert_eq!(form.selected_profile().selector, "builtin:coder");
}
#[test]
fn name_input_uses_old_safe_character_policy() {
let mut form = SpawnForm::new(Some("worker".to_owned()), "yoi".to_owned(), choices());
form.apply_key(KeyEvent::new(KeyCode::Char('-'), KeyModifiers::NONE));
form.apply_key(KeyEvent::new(KeyCode::Char('1'), KeyModifiers::NONE));
form.apply_key(KeyEvent::new(KeyCode::Char('/'), KeyModifiers::NONE));
assert_eq!(form.worker_name, "worker-1");
}
#[test]
fn enter_rejects_empty_name_and_escape_cancels() {
let mut form = SpawnForm::new(Some(String::new()), "yoi".to_owned(), choices());
assert_eq!(
form.apply_key(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)),
SpawnAction::None
);
assert_eq!(
form.status.as_ref().map(|(message, _)| message.as_str()),
Some("worker name cannot be empty")
);
assert_eq!(
form.apply_key(KeyEvent::new(KeyCode::Esc, KeyModifiers::NONE)),
SpawnAction::Cancel
);
}
#[test]
fn renderer_preserves_legacy_inline_spawn_form() {
let backend = ratatui::backend::TestBackend::new(100, VIEWPORT_HEIGHT);
let mut terminal = ratatui::Terminal::new(backend).unwrap();
let form = SpawnForm::new(None, "yoi".to_owned(), choices());
terminal.draw(|frame| draw_form(frame, &form)).unwrap();
let buffer = terminal.backend().buffer();
let rendered = buffer
.content
.chunks(buffer.area.width as usize)
.map(|row| row.iter().map(|cell| cell.symbol()).collect::<String>())
.collect::<Vec<_>>()
.join("\n");
assert!(rendered.contains("spawn worker"));
assert!(rendered.contains("name: yoi"));
assert!(rendered.contains("profile: builtin:default (default) — Default"));
assert!(rendered.contains("enter spawn · left/right edit · esc cancel"));
}
#[test]
fn builtin_discovery_produces_a_default_profile_choice() {
let registry = ProfileDiscovery::with_sources(None, None)
.discover()
.unwrap();
let choices = profile_choices(&registry);
let default = choices.iter().find(|choice| choice.is_default).unwrap();
assert_eq!(default.selector, "builtin:default");
assert!(default.label.contains("(default)"));
}
#[test]
fn default_worker_name_comes_from_sanitised_directory_basename() {
assert_eq!(
default_worker_name(Path::new("/home/hare/Project/yoi")),
"yoi"
);
assert_eq!(
default_worker_name(Path::new("/home/hare/Project/my project")),
"my-project"
);
assert_eq!(default_worker_name(Path::new("/")), "worker");
}
#[test]
fn explicit_profile_bypasses_discovery_and_uses_directory_name() {
let selection = select(
Path::new("/home/hare/Project/yoi"),
None,
Some("builtin:coder".to_owned()),
)
.unwrap()
.unwrap();
assert_eq!(selection.worker_name, "yoi");
assert_eq!(selection.profile, "builtin:coder");
}
}
+153 -45
View File
@@ -36,6 +36,9 @@ use crate::task::{TaskCounts, TaskEntry, TaskStatus, TaskStore};
use crate::text_selection::{HistoryViewport, SelectionRow};
use crate::view_mode::Mode;
const RUN_SPINNER_FRAMES: [&str; 8] = ["", "", "", "", "", "", "", ""];
const RUN_SPINNER_FRAME_MS: u128 = 80;
pub fn draw(frame: &mut Frame, app: &mut App) {
let area = frame.area();
// Input content starts after the prompt (`> ` or `: `), so the width
@@ -57,19 +60,27 @@ pub fn draw(frame: &mut Frame, app: &mut App) {
let tabs = app.worker_view_tabs();
let show_tabs = tabs.len() > 1;
let mini_view_h = task_mini_view_height(&app.selected_worker_view().task_store, show_tabs);
// One blank row separates the history tail from the mini-view so
// the latest message doesn't visually crash into the task summary.
// Folds away with the mini-view when there are no tasks.
let mini_view_gap = if mini_view_h > 0 { 1 } else { 0 };
let run_status_h = u16::from(app.running);
let run_status_gap = run_status_h;
// One blank row separates the history tail from the run/task mini-view so
// the latest message doesn't visually crash into operational status.
// Folds away when neither run status nor tasks are visible.
let mini_view_gap = if mini_view_h > 0 || run_status_h > 0 {
1
} else {
0
};
let chunks = Layout::vertical([
Constraint::Min(0), // history view
Constraint::Length(mini_view_gap), // gap above mini-view
Constraint::Length(mini_view_h), // task mini-view (0 when empty)
Constraint::Length(1), // separator
Constraint::Length(1), // status
Constraint::Length(input_height), // input area
Constraint::Length(1), // actionbar
Constraint::Min(0), // history view
Constraint::Length(mini_view_gap), // gap above run/task mini-view
Constraint::Length(run_status_h), // active run status
Constraint::Length(run_status_gap), // gap below active run status
Constraint::Length(mini_view_h), // task mini-view (0 when empty)
Constraint::Length(1), // separator
Constraint::Length(1), // status
Constraint::Length(input_height), // input area
Constraint::Length(1), // actionbar
])
.split(area);
@@ -82,24 +93,27 @@ pub fn draw(frame: &mut Frame, app: &mut App) {
} else {
draw_history(frame, app, chunks[0]);
}
if run_status_h > 0 {
draw_run_status(frame, app, chunks[2]);
}
if mini_view_h > 0 {
draw_task_mini_view(
frame,
&app.selected_worker_view().task_store,
&tabs,
chunks[2],
chunks[4],
);
}
draw_separator(frame, chunks[3]);
draw_separator(frame, chunks[5]);
// Status/composer/control surfaces remain parent-owned. View selection changes
// only transcript/task presentation and never implies SubWorker control.
draw_status(frame, app, chunks[4]);
draw_input(frame, app, &input_render, chunks[5]);
draw_actionbar(frame, app, chunks[6]);
draw_status(frame, app, chunks[6]);
draw_input(frame, app, &input_render, chunks[7]);
draw_actionbar(frame, app, chunks[8]);
if app.is_command_mode() {
draw_command_popup(frame, app, chunks[5]);
draw_command_popup(frame, app, chunks[7]);
} else if let Some(state) = app.completion.as_ref().filter(|c| c.is_active()) {
draw_completion_popup(frame, state, chunks[5]);
draw_completion_popup(frame, state, chunks[7]);
}
}
@@ -120,6 +134,65 @@ fn task_mini_view_height(store: &TaskStore, show_tabs: bool) -> u16 {
(active_shown as u16).saturating_add(1)
}
fn draw_run_status(frame: &mut Frame, app: &App, area: Rect) {
frame.render_widget(Paragraph::new(run_status_line(app, Instant::now())), area);
}
fn run_status_line(app: &App, now: Instant) -> Line<'static> {
let elapsed = app
.run_started_at
.and_then(|started_at| now.checked_duration_since(started_at))
.unwrap_or_default();
let spinner_index =
((elapsed.as_millis() / RUN_SPINNER_FRAME_MS) as usize) % RUN_SPINNER_FRAMES.len();
let request_label = if app.run_requests == 1 {
"1 req".to_owned()
} else {
format!("{} reqs", app.run_requests)
};
Line::from(vec![
Span::styled(
RUN_SPINNER_FRAMES[spinner_index],
Style::default()
.fg(Color::Cyan)
.add_modifier(Modifier::BOLD),
),
Span::raw(" "),
Span::styled(
fmt_run_elapsed(elapsed.as_secs()),
Style::default().fg(Color::Gray),
),
Span::styled("", Style::default().fg(Color::DarkGray)),
Span::styled(request_label, Style::default().fg(Color::Gray)),
Span::styled(" | ", Style::default().fg(Color::DarkGray)),
Span::styled("", Style::default().fg(Color::Green)),
Span::styled(
fmt_tokens(app.run_upload_tokens),
Style::default().fg(Color::Green),
),
Span::styled("/", Style::default().fg(Color::DarkGray)),
Span::styled("", Style::default().fg(Color::Yellow)),
Span::styled(
fmt_tokens(app.run_output_tokens),
Style::default().fg(Color::Yellow),
),
])
}
fn fmt_run_elapsed(secs: u64) -> String {
let hours = secs / 3600;
let minutes = (secs % 3600) / 60;
let seconds = secs % 60;
if hours > 0 {
format!("{hours}h {minutes}m {seconds:02}s")
} else if minutes > 0 {
format!("{minutes}m {seconds:02}s")
} else {
format!("{seconds}s")
}
}
fn draw_task_mini_view(frame: &mut Frame, store: &TaskStore, tabs: &[WorkerViewTab], area: Rect) {
if area.height == 0 || area.width == 0 {
return;
@@ -1223,6 +1296,28 @@ fn chip_span_for(seg: &Segment, fallback: Style) -> (Style, String) {
Style::default().fg(Color::Magenta),
format!("[Clipboard #{id} | {chars} chars, {line_count} lines]"),
),
Segment::PasteArtifact { artifact } => (
Style::default().fg(Color::Magenta),
format!(
"[Paste artifact {} | {} chars, {} lines, {}, {}, created {} ms]",
artifact.artifact_id,
artifact.char_count,
artifact.line_count,
artifact.media_type.as_str(),
artifact.availability.as_str(),
artifact.created_at_ms
),
),
Segment::UploadedFile { file } => (
Style::default().fg(Color::Cyan),
format!(
"[Attached {} | {} bytes, {}, {}]",
file.file_name,
file.byte_len,
file.media_type,
file.availability.as_str()
),
),
Segment::FileRef { path } => (Style::default().fg(Color::Cyan), format!("@{path}")),
Segment::Flow { selector } => (
Style::default().fg(Color::Yellow),
@@ -1241,6 +1336,22 @@ fn segment_display_text(seg: &Segment) -> String {
Segment::Paste {
id, chars, lines, ..
} => format!("[Clipboard #{id} | {chars} chars, {lines} lines]"),
Segment::PasteArtifact { artifact } => format!(
"[Paste artifact {} | {} chars, {} lines, {}, {}, created {} ms]",
artifact.artifact_id,
artifact.char_count,
artifact.line_count,
artifact.media_type.as_str(),
artifact.availability.as_str(),
artifact.created_at_ms
),
Segment::UploadedFile { file } => format!(
"[Attached {} | {} bytes, {}, {}]",
file.file_name,
file.byte_len,
file.media_type,
file.availability.as_str()
),
Segment::FileRef { path } => format!("@{path}"),
Segment::Flow { selector } => format!("[Flow: {selector}]"),
Segment::Unknown => "[unknown segment]".to_owned(),
@@ -1726,32 +1837,7 @@ fn draw_status(frame: &mut Frame, app: &App, area: Rect) {
),
];
if app.running {
let status = if let Some(wait_event) = &app.latest_llm_wait_event {
format!(
"request: {} | ↑{}/↓{} | {wait_event}",
app.run_requests,
fmt_tokens(app.run_upload_tokens),
fmt_tokens(app.run_output_tokens),
)
} else if let Some(tool) = &app.current_tool {
format!(
"request: {} | ↑{}/↓{} | tool: {tool}",
app.run_requests,
fmt_tokens(app.run_upload_tokens),
fmt_tokens(app.run_output_tokens),
)
} else {
format!(
"request: {} | ↑{}/↓{}",
app.run_requests,
fmt_tokens(app.run_upload_tokens),
fmt_tokens(app.run_output_tokens),
)
};
spans.push(Span::raw(" | "));
spans.push(Span::styled(status, Style::default().fg(Color::Yellow)));
} else if app.paused {
if app.paused {
spans.push(Span::raw(" | "));
spans.push(Span::styled(
"paused",
@@ -1763,7 +1849,7 @@ fn draw_status(frame: &mut Frame, app: &App, area: Rect) {
" — Enter to resume, Ctrl-X to cancel, type to start new turn",
Style::default().fg(Color::DarkGray),
));
} else {
} else if !app.running {
spans.push(Span::styled(" idle", Style::default().fg(Color::DarkGray)));
}
@@ -2053,6 +2139,28 @@ mod tests {
use protocol::WorkerStatus;
use std::time::{Duration, Instant};
#[test]
fn run_status_line_matches_console_metrics_and_spinner_frame() {
let now = Instant::now();
let mut app = App::new("worker".into());
app.run_started_at = now.checked_sub(Duration::from_millis(160));
app.run_requests = 1;
app.run_upload_tokens = 1_200;
app.run_output_tokens = 45;
assert_eq!(
line_text(&run_status_line(&app, now)),
"⣟ 0s ・ 1 req | ↑1.2k/↓45"
);
}
#[test]
fn run_elapsed_uses_console_style_units() {
assert_eq!(fmt_run_elapsed(9), "9s");
assert_eq!(fmt_run_elapsed(65), "1m 05s");
assert_eq!(fmt_run_elapsed(3_726), "1h 2m 06s");
}
#[test]
fn task_summary_right_aligns_worker_tabs_and_highlights_selection() {
let tabs = vec![
+1
View File
@@ -18,6 +18,7 @@ sha2.workspace = true
tempfile.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["process", "rt", "sync", "time"] }
workspace-api = { workspace = true }
[dev-dependencies]
serde_json.workspace = true
+5
View File
@@ -743,6 +743,7 @@ mod tests {
command: command.into(),
timeout_secs: 5,
output_limit: 1024,
spill_dir: None,
tool_call_id: Some(tool_call_id.into()),
})
.await
@@ -770,6 +771,7 @@ mod tests {
command: "printf ready; sleep 0.2; printf done".into(),
timeout_secs: 5,
output_limit: 1024,
spill_dir: None,
tool_call_id: Some("tool-delegated".into()),
})
.await
@@ -858,6 +860,7 @@ mod tests {
command: "printf denied".into(),
timeout_secs: 5,
output_limit: 1024,
spill_dir: None,
tool_call_id: Some("read-only-command".into()),
})
.await,
@@ -993,6 +996,7 @@ mod tests {
command: "printf revoked".into(),
timeout_secs: 5,
output_limit: 1024,
spill_dir: None,
tool_call_id: Some("revoked-child-command".into()),
})
.await,
@@ -1171,6 +1175,7 @@ mod tests {
command: "printf closed".into(),
timeout_secs: 5,
output_limit: 1024,
spill_dir: None,
tool_call_id: Some("closed-parent-command".into()),
})
.await,
+149 -17
View File
@@ -102,8 +102,18 @@ pub enum WorkdirTransportErrorCode {
Conflict,
Unsupported,
InvalidRequest,
Denied,
OutOfScope,
SymlinkOutOfScope,
BrokenSymlink,
SymlinkTargetIsDirectory,
ReadOnly,
IsDirectory,
SymlinkDirectoryNotTraversed,
UnknownCommand,
Unavailable,
Io,
Transport,
Internal,
}
@@ -114,8 +124,18 @@ impl WorkdirTransportErrorCode {
Self::Conflict => "conflict",
Self::Unsupported => "unsupported",
Self::InvalidRequest => "invalid_request",
Self::Denied => "denied",
Self::OutOfScope => "out_of_scope",
Self::SymlinkOutOfScope => "symlink_out_of_scope",
Self::BrokenSymlink => "broken_symlink",
Self::SymlinkTargetIsDirectory => "symlink_target_is_directory",
Self::ReadOnly => "read_only",
Self::IsDirectory => "is_directory",
Self::SymlinkDirectoryNotTraversed => "symlink_directory_not_traversed",
Self::UnknownCommand => "unknown_command",
Self::Unavailable => "unavailable",
Self::Io => "io",
Self::Transport => "transport",
Self::Internal => "internal",
}
}
@@ -125,9 +145,16 @@ impl WorkdirTransportErrorCode {
match self {
Self::NotFound | Self::UnknownCommand => 404,
Self::Conflict => 409,
Self::Unsupported | Self::InvalidRequest => 400,
Self::Denied | Self::OutOfScope | Self::SymlinkOutOfScope | Self::ReadOnly => 403,
Self::Unsupported
| Self::InvalidRequest
| Self::BrokenSymlink
| Self::SymlinkTargetIsDirectory
| Self::IsDirectory
| Self::SymlinkDirectoryNotTraversed => 400,
Self::Unavailable => 503,
Self::Internal => 500,
Self::Io | Self::Internal => 500,
Self::Transport => 502,
}
}
}
@@ -160,8 +187,41 @@ impl WorkdirTransportError {
WorkdirError::Unavailable(_) | WorkdirError::SessionClosed => {
(Code::Unavailable, "Workdir session is unavailable")
}
WorkdirError::Denied(_) => (Code::InvalidRequest, "Workdir operation was denied"),
WorkdirError::Transport(_) => (Code::Internal, "Workdir transport failed"),
WorkdirError::Denied(_) => (Code::Denied, "Workdir operation was denied"),
WorkdirError::OutOfScope(_) => (Code::OutOfScope, "Workdir path is out of scope"),
WorkdirError::SymlinkOutOfScope { .. } => (
Code::SymlinkOutOfScope,
"Workdir symlink target is out of scope",
),
WorkdirError::BrokenSymlink { .. } => {
(Code::BrokenSymlink, "Workdir symlink target does not exist")
}
WorkdirError::SymlinkTargetIsDirectory { .. } => (
Code::SymlinkTargetIsDirectory,
"Workdir symlink target is a directory",
),
WorkdirError::ReadOnly(_) => (Code::ReadOnly, "Workdir path is read-only"),
WorkdirError::IsDirectory(_) => (Code::IsDirectory, "Workdir path is a directory"),
WorkdirError::SymlinkDirectoryNotTraversed { .. } => (
Code::SymlinkDirectoryNotTraversed,
"Workdir symlink directory was not traversed",
),
WorkdirError::Io { source, .. } => match source.kind() {
std::io::ErrorKind::NotFound => (Code::NotFound, "Workdir path was not found"),
std::io::ErrorKind::PermissionDenied => {
(Code::Denied, "Workdir operation was denied")
}
std::io::ErrorKind::AlreadyExists => {
(Code::Conflict, "Workdir resource already exists")
}
std::io::ErrorKind::InvalidInput | std::io::ErrorKind::InvalidData => {
(Code::InvalidRequest, "Workdir operation request is invalid")
}
std::io::ErrorKind::TimedOut => (Code::Unavailable, "Workdir operation timed out"),
_ => (Code::Io, "Workdir I/O operation failed"),
},
WorkdirError::OperationFailed => (Code::Internal, "Workdir operation failed"),
WorkdirError::Transport(_) => (Code::Transport, "Workdir transport failed"),
WorkdirError::InvalidPath(_)
| WorkdirError::RelativePath(_)
| WorkdirError::InvalidGlob(_)
@@ -169,14 +229,6 @@ impl WorkdirTransportError {
| WorkdirError::InvalidArgument(_) => {
(Code::InvalidRequest, "Workdir operation request is invalid")
}
WorkdirError::OutOfScope(_)
| WorkdirError::SymlinkOutOfScope { .. }
| WorkdirError::BrokenSymlink { .. }
| WorkdirError::SymlinkTargetIsDirectory { .. }
| WorkdirError::ReadOnly(_)
| WorkdirError::IsDirectory(_)
| WorkdirError::SymlinkDirectoryNotTraversed { .. }
| WorkdirError::Io { .. } => (Code::Internal, "Workdir operation failed"),
};
Self {
code,
@@ -190,10 +242,38 @@ impl WorkdirTransportError {
Code::NotFound => WorkdirError::NotFound("<remote>".into()),
Code::Conflict => WorkdirError::Conflict(self.message),
Code::Unsupported => WorkdirError::UnsupportedOperation(self.message),
Code::UnknownCommand => WorkdirError::UnknownCommand("<remote>".to_string()),
Code::InvalidRequest => WorkdirError::InvalidArgument(self.message),
Code::Denied => WorkdirError::Denied(self.message),
Code::OutOfScope => WorkdirError::OutOfScope("<remote>".into()),
Code::SymlinkOutOfScope => WorkdirError::SymlinkOutOfScope {
path: "<remote>".into(),
target: "<remote-target>".into(),
required_permission: "requested",
},
Code::BrokenSymlink => WorkdirError::BrokenSymlink {
path: "<remote>".into(),
link: "<remote-link>".into(),
target: "<remote-target>".into(),
},
Code::SymlinkTargetIsDirectory => WorkdirError::SymlinkTargetIsDirectory {
path: "<remote>".into(),
target: "<remote-target>".into(),
},
Code::ReadOnly => WorkdirError::ReadOnly("<remote>".into()),
Code::IsDirectory => WorkdirError::IsDirectory("<remote>".into()),
Code::SymlinkDirectoryNotTraversed => WorkdirError::SymlinkDirectoryNotTraversed {
tool: "remote operation",
path: "<remote>".into(),
target: "<remote-target>".into(),
},
Code::UnknownCommand => WorkdirError::UnknownCommand("<remote>".to_string()),
Code::Unavailable => WorkdirError::Unavailable(self.message),
Code::Internal => WorkdirError::Transport(self.message),
Code::Io => WorkdirError::Io {
path: "<remote>".into(),
source: std::io::Error::other(self.message),
},
Code::Transport => WorkdirError::Transport(self.message),
Code::Internal => WorkdirError::OperationFailed,
}
}
}
@@ -584,8 +664,42 @@ mod tests {
"modified externally",
),
(WorkdirTransportErrorCode::Unsupported, 400, "unsupported"),
(WorkdirTransportErrorCode::Denied, 403, "denied"),
(
WorkdirTransportErrorCode::OutOfScope,
403,
"outside allowed scope",
),
(
WorkdirTransportErrorCode::SymlinkOutOfScope,
403,
"outside allowed requested scope",
),
(
WorkdirTransportErrorCode::BrokenSymlink,
400,
"broken symlink",
),
(
WorkdirTransportErrorCode::SymlinkTargetIsDirectory,
400,
"symlink to a directory",
),
(WorkdirTransportErrorCode::ReadOnly, 403, "read-only"),
(WorkdirTransportErrorCode::IsDirectory, 400, "expected file"),
(
WorkdirTransportErrorCode::SymlinkDirectoryNotTraversed,
400,
"does not follow symlink directories",
),
(WorkdirTransportErrorCode::Unavailable, 503, "unavailable"),
(WorkdirTransportErrorCode::Internal, 500, "transport failed"),
(WorkdirTransportErrorCode::Io, 500, "I/O error"),
(
WorkdirTransportErrorCode::Transport,
502,
"transport failed",
),
(WorkdirTransportErrorCode::Internal, 500, "operation failed"),
] {
let transport = WorkdirTransportError {
code,
@@ -620,7 +734,8 @@ mod tests {
let transport = WorkdirTransportError::from_workdir_error(&WorkdirError::Transport(
"Workspace API request timed out".to_string(),
));
assert_eq!(transport.code, WorkdirTransportErrorCode::Internal);
assert_eq!(transport.code, WorkdirTransportErrorCode::Transport);
assert_eq!(transport.code.http_status(), 502);
assert_eq!(transport.message, "Workdir transport failed");
assert!(matches!(
transport.into_workdir_error(),
@@ -635,8 +750,25 @@ mod tests {
source: std::io::Error::new(std::io::ErrorKind::PermissionDenied, "host detail"),
};
let transport = WorkdirTransportError::from_workdir_error(&error);
assert_eq!(transport.code, WorkdirTransportErrorCode::Internal);
assert_eq!(transport.code, WorkdirTransportErrorCode::Denied);
assert!(!transport.message.contains("/secret"));
assert!(!transport.message.contains("host detail"));
assert!(matches!(
transport.into_workdir_error(),
WorkdirError::Denied(_)
));
let error = WorkdirError::Io {
path: "/secret/runtime/root/file".into(),
source: std::io::Error::other("host detail"),
};
let transport = WorkdirTransportError::from_workdir_error(&error);
assert_eq!(transport.code, WorkdirTransportErrorCode::Io);
assert!(!transport.message.contains("/secret"));
assert!(!transport.message.contains("host detail"));
assert!(matches!(
transport.into_workdir_error(),
WorkdirError::Io { .. }
));
}
}
+3
View File
@@ -234,6 +234,9 @@ pub enum WorkdirError {
#[error("Workdir session is unavailable: {0}")]
Unavailable(String),
#[error("Workdir operation failed")]
OperationFailed,
#[error("Workdir transport failed: {0}")]
Transport(String),
+231 -3
View File
@@ -10,9 +10,7 @@
use std::collections::{BTreeMap, HashMap};
use std::fmt::Debug;
#[cfg(test)]
use std::io::Write as _;
use std::io::{Read as _, Seek as _, SeekFrom};
use std::io::{Read as _, Seek as _, SeekFrom, Write as _};
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
@@ -691,6 +689,11 @@ impl WorkdirSession for LocalWorkdirSession {
async fn start_command(&self, request: CommandRequest) -> Result<CommandHandle, WorkdirError> {
self.ensure_capability(WorkdirSessionCapability::Command)?;
self.ensure_open()?;
if let Some(spill_dir) = request.spill_dir.as_deref()
&& !self.inner.scope.snapshot().is_readable(spill_dir)
{
return Err(WorkdirError::OutOfScope(spill_dir.to_path_buf()));
}
let id = self.inner.next_command_id.fetch_add(1, Ordering::Relaxed);
let handle = CommandHandle(format!("command-{id}"));
let cwd = self.inner.cwd.clone();
@@ -776,6 +779,7 @@ impl WorkdirSession for LocalWorkdirSession {
content: String::new(),
next_cursor: None,
truncated: false,
output_path: None,
});
}
drop(commands);
@@ -792,6 +796,7 @@ impl WorkdirSession for LocalWorkdirSession {
content: String::new(),
next_cursor: None,
truncated: false,
output_path: None,
});
}
break commands
@@ -901,6 +906,7 @@ fn command_output_page(output: &CommandOutput, cursor: usize, limit: usize) -> C
content,
next_cursor: (end < total_chars).then_some(end),
truncated: output.truncated || end < total_chars,
output_path: output.output_path.clone(),
}
}
@@ -1059,6 +1065,22 @@ async fn run_command(
let (content, truncated) =
read_command_output_files(&stdout_path, &stderr_path, request.output_limit.max(1))?;
let output_path = match (truncated, request.spill_dir) {
(true, Some(spill_dir)) => {
let stdout_path = stdout_path.to_path_buf();
let stderr_path = stderr_path.to_path_buf();
Some(
tokio::task::spawn_blocking(move || {
persist_command_output(&stdout_path, &stderr_path, &spill_dir)
})
.await
.map_err(|error| {
WorkdirError::Unavailable(format!("Bash output spill task failed: {error}"))
})??,
)
}
_ => None,
};
Ok(CommandOutput {
status,
exit_code,
@@ -1066,6 +1088,7 @@ async fn run_command(
content,
next_cursor: None,
truncated,
output_path,
})
}
@@ -1154,6 +1177,59 @@ fn stable_utf8_prefix_len(bytes: &[u8]) -> usize {
inspected
}
fn persist_command_output(
stdout_path: &Path,
stderr_path: &Path,
spill_dir: &Path,
) -> Result<PathBuf, WorkdirError> {
std::fs::create_dir_all(spill_dir).map_err(|error| WorkdirError::io(spill_dir, error))?;
let mut artifact = tempfile::Builder::new()
.prefix("bash-")
.suffix(".log")
.tempfile_in(spill_dir)
.map_err(|error| WorkdirError::io(spill_dir, error))?;
let artifact_path = artifact.path().to_path_buf();
let mut stdout =
std::fs::File::open(stdout_path).map_err(|error| WorkdirError::io(stdout_path, error))?;
let stdout_len = stdout
.metadata()
.map_err(|error| WorkdirError::io(stdout_path, error))?
.len();
std::io::copy(&mut stdout, &mut artifact)
.map_err(|error| WorkdirError::io(&artifact_path, error))?;
let mut stderr =
std::fs::File::open(stderr_path).map_err(|error| WorkdirError::io(stderr_path, error))?;
let stderr_len = stderr
.metadata()
.map_err(|error| WorkdirError::io(stderr_path, error))?
.len();
if stdout_len > 0 && stderr_len > 0 {
stdout
.seek(SeekFrom::End(-1))
.map_err(|error| WorkdirError::io(stdout_path, error))?;
let mut last = [0_u8; 1];
stdout
.read_exact(&mut last)
.map_err(|error| WorkdirError::io(stdout_path, error))?;
if last[0] != b'\n' {
artifact
.write_all(b"\n")
.map_err(|error| WorkdirError::io(&artifact_path, error))?;
}
}
std::io::copy(&mut stderr, &mut artifact)
.map_err(|error| WorkdirError::io(&artifact_path, error))?;
artifact
.flush()
.map_err(|error| WorkdirError::io(&artifact_path, error))?;
artifact
.keep()
.map(|(_, path)| path)
.map_err(|error| WorkdirError::io(&artifact_path, error.error))
}
fn read_command_output_files(
stdout_path: &Path,
stderr_path: &Path,
@@ -1440,6 +1516,7 @@ mod tests {
command: "sleep 30".to_owned(),
timeout_secs: 60,
output_limit: 1024,
spill_dir: None,
tool_call_id: None,
},
)
@@ -1966,6 +2043,7 @@ mod tests {
command: "pwd && printf provider-command".into(),
timeout_secs: 5,
output_limit: 4096,
spill_dir: None,
tool_call_id: None,
},
)
@@ -1991,6 +2069,151 @@ mod tests {
);
}
#[tokio::test]
async fn explicitly_scoped_absolute_artifact_can_be_read_and_grepped() {
let dir = TempDir::new().unwrap();
let spill = TempDir::new().unwrap();
let artifact = spill.path().join("bash-output.log");
std::fs::write(&artifact, "first\nFINAL-NEEDLE\nlast\n").unwrap();
let scope = Scope::from_config(&ScopeConfig {
allow: vec![
ScopeRule {
target: dir.path().to_path_buf(),
permission: Permission::Write,
recursive: true,
},
ScopeRule {
target: spill.path().to_path_buf(),
permission: Permission::Read,
recursive: true,
},
],
deny: Vec::new(),
})
.unwrap();
let workdir = LocalWorkdirSession::new(scope, dir.path().to_path_buf());
let artifact_path = WorkdirPath::new_scoped(artifact.to_string_lossy()).unwrap();
let read = WorkdirSession::read(
&workdir,
ReadRequest {
path: artifact_path.clone(),
offset: 1,
limit: 1,
max_bytes: 1024,
},
)
.await
.unwrap();
assert_eq!(String::from_utf8(read.bytes).unwrap(), "FINAL-NEEDLE\n");
let grep = WorkdirSession::grep(
&workdir,
GrepRequest {
pattern: "FINAL-NEEDLE".into(),
path: artifact_path,
glob: None,
file_type: None,
case_insensitive: false,
before_context: 0,
after_context: 0,
multiline: false,
output_mode: crate::GrepOutputMode::Content,
limit: 10,
offset: 0,
},
)
.await
.unwrap();
assert_eq!(grep.match_count, 1);
assert!(grep.output.contains("FINAL-NEEDLE"));
}
#[tokio::test]
async fn command_rejects_spill_directory_without_read_scope() {
let dir = TempDir::new().unwrap();
let spill = TempDir::new().unwrap();
let workdir = make_fs(&dir);
let error = WorkdirSession::start_command(
&workdir,
CommandRequest {
command: "printf hidden".into(),
timeout_secs: 5,
output_limit: 1,
spill_dir: Some(spill.path().to_path_buf()),
tool_call_id: None,
},
)
.await
.unwrap_err();
assert!(matches!(error, WorkdirError::OutOfScope(path) if path == spill.path()));
}
#[tokio::test]
async fn truncated_command_output_is_retained_in_the_requested_spill_directory() {
let dir = TempDir::new().unwrap();
let spill = TempDir::new().unwrap();
let scope = Scope::from_config(&ScopeConfig {
allow: vec![
ScopeRule {
target: dir.path().to_path_buf(),
permission: Permission::Write,
recursive: true,
},
ScopeRule {
target: spill.path().to_path_buf(),
permission: Permission::Read,
recursive: true,
},
],
deny: Vec::new(),
})
.unwrap();
let workdir = LocalWorkdirSession::new(scope, dir.path().to_path_buf());
let handle = WorkdirSession::start_command(
&workdir,
CommandRequest {
command: "i=0; while [ $i -lt 200 ]; do printf 'line-%03d\\n' \"$i\"; i=$((i+1)); done; printf 'FINAL-NEEDLE\\n'".into(),
timeout_secs: 5,
output_limit: 64,
spill_dir: Some(spill.path().to_path_buf()),
tool_call_id: None,
},
)
.await
.unwrap();
let output = WorkdirSession::command_output(
&workdir,
CommandOutputRequest {
handle,
cursor: 0,
limit: 4096,
wait: true,
},
)
.await
.unwrap();
assert!(output.truncated);
let output_path = output.output_path.expect("retained output path");
assert_eq!(output_path.parent(), Some(spill.path()));
let retained = std::fs::read_to_string(&output_path).unwrap();
assert!(retained.starts_with("line-000\n"));
assert!(retained.ends_with("FINAL-NEEDLE\n"));
assert_eq!(retained.lines().count(), 201);
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt as _;
assert_eq!(
std::fs::metadata(output_path).unwrap().permissions().mode() & 0o777,
0o600
);
}
}
#[tokio::test]
async fn completed_command_output_can_be_read_in_bounded_unicode_pages() {
let dir = TempDir::new().unwrap();
@@ -2001,6 +2224,7 @@ mod tests {
command: "printf 'aéz'".into(),
timeout_secs: 5,
output_limit: 1024,
spill_dir: None,
tool_call_id: None,
},
)
@@ -2120,6 +2344,7 @@ mod tests {
content: "done".into(),
next_cursor: None,
truncated: false,
output_path: None,
})
});
workdir.inner.commands.lock().await.insert(
@@ -2224,6 +2449,7 @@ mod tests {
command: "printf ready; printf warning >&2; sleep 0.2; printf done".into(),
timeout_secs: 5,
output_limit: 1024,
spill_dir: None,
tool_call_id: Some("tool-7".into()),
},
)
@@ -2327,6 +2553,7 @@ mod tests {
command: "sleep 30".into(),
timeout_secs: 1,
output_limit: 1024,
spill_dir: None,
tool_call_id: None,
},
)
@@ -2396,6 +2623,7 @@ mod tests {
command: "sleep 30".into(),
timeout_secs: 60,
output_limit: 1024,
spill_dir: None,
tool_call_id: None,
},
)
+7
View File
@@ -1,3 +1,5 @@
use std::path::PathBuf;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
@@ -9,6 +11,9 @@ pub struct CommandRequest {
pub command: String,
pub timeout_secs: u64,
pub output_limit: usize,
/// Provider-local directory where complete output is retained when the
/// inline result exceeds `output_limit`.
pub spill_dir: Option<PathBuf>,
/// Optional caller-owned correlation id. Bash supplies its tool-call id so
/// user-facing command telemetry can update the corresponding Console row
/// without exposing provider/session handles.
@@ -96,4 +101,6 @@ pub struct CommandOutput {
pub content: String,
pub next_cursor: Option<usize>,
pub truncated: bool,
/// Complete output retained by the provider when `truncated` is true.
pub output_path: Option<PathBuf>,
}
+12 -251
View File
@@ -6,7 +6,12 @@
//! [`crate::http`].
use serde::{Deserialize, Serialize};
use std::fmt;
pub use workspace_api::{
RuntimeWorkingDirectoryCleanupTarget, RuntimeWorkingDirectorySummary,
WorkingDirectoryCleanupTarget, WorkingDirectoryMaterializerKind as MaterializerKind,
WorkingDirectoryOccupancy, WorkingDirectoryStatusKind, WorkingDirectorySummary,
};
/// Stable Workspace identity for a Worker hosted by a Runtime.
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
@@ -26,83 +31,6 @@ impl RuntimeWorkerRef {
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MaterializerKind {
#[default]
RuntimeGitCache,
/// Legacy persisted value from the pre-cache local `git worktree` materializer.
LocalGitWorktree,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum WorkingDirectoryStatusKind {
Active,
CleanupPending,
Corrupted,
NotFound,
Unknown,
}
impl WorkingDirectoryStatusKind {
pub const fn as_str(&self) -> &'static str {
match self {
Self::Active => "active",
Self::CleanupPending => "cleanup_pending",
Self::Corrupted => "corrupted",
Self::NotFound => "not_found",
Self::Unknown => "unknown",
}
}
}
impl fmt::Display for WorkingDirectoryStatusKind {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(self.as_str())
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct WorkingDirectoryCleanupTarget {
pub kind: String,
pub working_directory_id: String,
pub repository_id: String,
}
/// Durable Workspace occupancy projection for one Workdir.
#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub struct WorkingDirectoryOccupancy {
#[serde(flatten)]
pub worker: RuntimeWorkerRef,
pub display_name: String,
pub linked_at: String,
}
impl<'de> Deserialize<'de> for WorkingDirectoryOccupancy {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct Wire {
runtime_id: String,
worker_id: String,
display_name: String,
linked_at: String,
}
let wire = Wire::deserialize(deserializer)?;
Ok(Self {
worker: RuntimeWorkerRef::new(wire.runtime_id, wire.worker_id),
display_name: wire.display_name,
linked_at: wire.linked_at,
})
}
}
/// Immutable materialization provenance retained by Workspace inventory.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
@@ -139,100 +67,6 @@ pub struct WorkingDirectoryCurrentObservation {
pub occupied_by: Option<WorkingDirectoryOccupancy>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct WorkingDirectorySummary {
pub working_directory_id: String,
pub repository_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub creation_selector: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub creation_ref: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub creation_tree: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub current_selector: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub current_ref: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub current_tree: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub observed_at_epoch_seconds: Option<u64>,
pub materializer_kind: MaterializerKind,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cleanup_target: Option<WorkingDirectoryCleanupTarget>,
pub status: WorkingDirectoryStatusKind,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cleanliness: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub primary_worker_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub occupied_by: Option<WorkingDirectoryOccupancy>,
}
impl WorkingDirectorySummary {
/// Workspace-managed inventory rows carry explicit cleanup authority.
pub fn is_workspace_managed(&self) -> bool {
self.cleanup_target.is_some()
}
pub fn provenance(&self) -> WorkingDirectoryProvenance {
WorkingDirectoryProvenance {
creation_selector: self.creation_selector.clone(),
creation_ref: self.creation_ref.clone(),
creation_tree: self.creation_tree.clone(),
materializer_kind: self.materializer_kind.clone(),
cleanup_target: self.cleanup_target.clone(),
}
}
pub fn current_observation(&self) -> WorkingDirectoryCurrentObservation {
WorkingDirectoryCurrentObservation {
current_selector: self.current_selector.clone(),
current_ref: self.current_ref.clone(),
current_tree: self.current_tree.clone(),
observed_at_epoch_seconds: self.observed_at_epoch_seconds,
status: self.status.clone(),
cleanliness: self.cleanliness.clone(),
primary_worker_id: self.primary_worker_id.clone(),
occupied_by: self.occupied_by.clone(),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum WorkingDirectoryDiagnosticSeverity {
Info,
Warning,
Error,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct WorkingDirectoryDiagnostic {
pub code: String,
pub severity: WorkingDirectoryDiagnosticSeverity,
pub message: String,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct WorkingDirectoryListResponse {
pub workspace_id: String,
pub items: Vec<WorkingDirectorySummary>,
pub diagnostics: Vec<WorkingDirectoryDiagnostic>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct WorkingDirectoryDetailResponse {
pub workspace_id: String,
pub runtime_id: String,
pub item: WorkingDirectorySummary,
pub diagnostics: Vec<WorkingDirectoryDiagnostic>,
}
#[cfg(test)]
mod tests {
use super::*;
@@ -255,88 +89,15 @@ mod tests {
}
#[test]
fn occupied_and_free_list_response_round_trips() {
let response = WorkingDirectoryListResponse {
workspace_id: "workspace".to_string(),
items: vec![
WorkingDirectorySummary {
working_directory_id: "occupied".to_string(),
repository_id: "repo".to_string(),
creation_selector: Some("develop".to_string()),
creation_ref: Some("abc123".to_string()),
creation_tree: Some("tree123".to_string()),
current_selector: Some("work/ticket".to_string()),
current_ref: Some("def456".to_string()),
current_tree: Some("tree456".to_string()),
observed_at_epoch_seconds: Some(1_777_777_777),
materializer_kind: MaterializerKind::LocalGitWorktree,
cleanup_target: Some(WorkingDirectoryCleanupTarget {
kind: "git_worktree".to_string(),
working_directory_id: "occupied".to_string(),
repository_id: "repo".to_string(),
}),
status: WorkingDirectoryStatusKind::Active,
cleanliness: Some("clean".to_string()),
primary_worker_id: None,
occupied_by: Some(WorkingDirectoryOccupancy {
worker: RuntimeWorkerRef::new("arcadia", "worker-opaque-64"),
display_name: "Coder".to_string(),
linked_at: "2026-08-12T00:00:00Z".to_string(),
}),
},
WorkingDirectorySummary {
working_directory_id: "free".to_string(),
repository_id: "repo".to_string(),
creation_selector: None,
creation_ref: None,
creation_tree: None,
current_selector: None,
current_ref: Some("987fed".to_string()),
current_tree: None,
observed_at_epoch_seconds: None,
materializer_kind: MaterializerKind::LocalGitWorktree,
cleanup_target: None,
status: WorkingDirectoryStatusKind::Active,
cleanliness: Some("unknown".to_string()),
primary_worker_id: None,
occupied_by: None,
},
],
diagnostics: vec![WorkingDirectoryDiagnostic {
code: "observed".to_string(),
severity: WorkingDirectoryDiagnosticSeverity::Info,
message: "inventory observed".to_string(),
}],
};
let encoded = serde_json::to_value(&response).unwrap();
fn workspace_workdir_projection_reexports_workspace_api_authority() {
assert_eq!(
encoded["items"][0]["occupied_by"]["worker_id"],
"worker-opaque-64"
std::any::TypeId::of::<WorkingDirectorySummary>(),
std::any::TypeId::of::<workspace_api::WorkingDirectorySummary>()
);
assert!(
encoded["items"][0]["occupied_by"]
.get("runtime_worker_id")
.is_none()
assert_eq!(
std::any::TypeId::of::<WorkingDirectoryOccupancy>(),
std::any::TypeId::of::<workspace_api::WorkingDirectoryOccupancy>()
);
assert!(encoded["items"][1].get("occupied_by").is_none());
let mut stale = encoded.clone();
stale["items"][0]["occupied_by"]["runtime_worker_id"] = serde_json::json!(64);
assert!(serde_json::from_value::<WorkingDirectoryListResponse>(stale).is_err());
let decoded: WorkingDirectoryListResponse = serde_json::from_value(encoded).unwrap();
assert_eq!(decoded, response);
let detail = WorkingDirectoryDetailResponse {
workspace_id: decoded.workspace_id.clone(),
runtime_id: "arcadia".to_string(),
item: decoded.items[0].clone(),
diagnostics: decoded.diagnostics.clone(),
};
let encoded = serde_json::to_value(&detail).unwrap();
let decoded: WorkingDirectoryDetailResponse = serde_json::from_value(encoded).unwrap();
assert_eq!(decoded, detail);
}
}
+2 -2
View File
@@ -92,9 +92,9 @@ pub struct WorkingDirectoryRepository {
}
pub use workdir::workspace::{
MaterializerKind, WorkingDirectoryCleanupTarget, WorkingDirectoryCurrentObservation,
MaterializerKind, RuntimeWorkingDirectoryCleanupTarget as WorkingDirectoryCleanupTarget,
RuntimeWorkingDirectorySummary as WorkingDirectorySummary, WorkingDirectoryCurrentObservation,
WorkingDirectoryOccupancy, WorkingDirectoryProvenance, WorkingDirectoryStatusKind,
WorkingDirectorySummary,
};
#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
+48 -1
View File
@@ -8,7 +8,7 @@ use crate::interaction::WorkerInput;
#[cfg(feature = "ws-server")]
use crate::observation::WorkerObservationEvent;
use crate::working_directory::{WorkingDirectoryBinding, WorkingDirectoryDiagnostic};
use protocol::Method;
use protocol::{Method, UploadedFileRef};
use serde::{Deserialize, Serialize};
use std::fmt;
use std::sync::Arc;
@@ -33,6 +33,8 @@ pub enum WorkerExecutionOperation {
Spawn,
Restore,
Input,
UploadFile,
DeleteUploadedFile,
ProtocolMethod,
Stop,
Cancel,
@@ -385,6 +387,31 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static {
input: WorkerInput,
) -> WorkerExecutionResult;
fn upload_file(
&self,
_handle: &WorkerExecutionHandle,
_file_name: &str,
_media_type: &str,
_content: &[u8],
_context: Option<&session_store::UploadedFileUploadContext>,
) -> Result<UploadedFileRef, WorkerExecutionResult> {
Err(WorkerExecutionResult::unsupported(
WorkerExecutionOperation::UploadFile,
"execution backend does not support file upload",
))
}
fn delete_uploaded_file(
&self,
_handle: &WorkerExecutionHandle,
_artifact_id: &str,
) -> WorkerExecutionResult {
WorkerExecutionResult::unsupported(
WorkerExecutionOperation::DeleteUploadedFile,
"execution backend does not support uploaded-file deletion",
)
}
fn dispatch_method(
&self,
_handle: &WorkerExecutionHandle,
@@ -514,6 +541,26 @@ impl WorkerExecutionBackendRef {
self.backend.dispatch_input(handle, input)
}
pub(crate) fn upload_file(
&self,
handle: &WorkerExecutionHandle,
file_name: &str,
media_type: &str,
content: &[u8],
context: Option<&session_store::UploadedFileUploadContext>,
) -> Result<UploadedFileRef, WorkerExecutionResult> {
self.backend
.upload_file(handle, file_name, media_type, content, context)
}
pub(crate) fn delete_uploaded_file(
&self,
handle: &WorkerExecutionHandle,
artifact_id: &str,
) -> WorkerExecutionResult {
self.backend.delete_uploaded_file(handle, artifact_id)
}
pub(crate) fn dispatch_method(
&self,
handle: &WorkerExecutionHandle,
+164 -2
View File
@@ -32,7 +32,7 @@ use axum::body::{Body, Bytes};
use axum::extract::rejection::{JsonRejection, QueryRejection};
#[cfg(feature = "ws-server")]
use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade};
use axum::extract::{Extension, Path, Query, State};
use axum::extract::{DefaultBodyLimit, Extension, Path, Query, State};
use axum::http::{Method, Request, StatusCode, header};
use axum::middleware::{self, Next};
use axum::response::{IntoResponse, Response};
@@ -238,6 +238,14 @@ fn runtime_http_router_with_optional_auth(
post(execute_worker_retention),
)
.route("/v1/workers/{worker_id}/input", post(send_worker_input))
.route(
"/v1/workers/{worker_id}/attachments",
post(upload_worker_file).layer(DefaultBodyLimit::max(MAX_WORKER_FILE_UPLOAD_BYTES)),
)
.route(
"/v1/workers/{worker_id}/attachments/{artifact_id}",
delete(delete_worker_uploaded_file),
)
.route("/v1/workers/{worker_id}/restore", post(restore_worker))
.route(
"/v1/workers/{worker_id}/workspace-api",
@@ -263,6 +271,9 @@ fn runtime_http_router_with_optional_auth(
.layer(middleware::from_fn_with_state(state, require_runtime_auth))
}
pub const MAX_WORKER_FILE_UPLOAD_BYTES: usize =
session_store::DEFAULT_MAX_UPLOADED_FILE_BYTES as usize;
#[derive(Clone)]
struct RuntimeHttpState {
runtime: Runtime,
@@ -375,6 +386,32 @@ pub struct RuntimeHttpWorkerInputResponse {
pub ack: WorkerInteractionAck,
}
#[derive(Clone, Debug, Deserialize)]
pub struct RuntimeHttpUploadFileQuery {
pub file_name: String,
pub media_type: String,
#[serde(default)]
pub upload_id: Option<String>,
#[serde(default)]
pub principal_id: Option<String>,
#[serde(default)]
pub workspace_id: Option<String>,
#[serde(default)]
pub runtime_id: Option<String>,
#[serde(default)]
pub owner_worker_id: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct RuntimeHttpUploadedFileResponse {
pub file: protocol::UploadedFileRef,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct RuntimeHttpUploadedFileDeleteResponse {
pub deleted: bool,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct RuntimeHttpWorkerCompletionsRequest {
pub kind: protocol::CompletionKind,
@@ -703,6 +740,7 @@ async fn run_workdir_session_operation(
content: String::new(),
next_cursor: Some(cursor),
truncated: false,
output_path: None,
},
};
WorkdirSessionOperationResult::CommandOutput(output)
@@ -1419,6 +1457,115 @@ async fn worker_completions(
}))
}
async fn upload_worker_file(
State(state): State<RuntimeHttpState>,
auth: Option<Extension<RuntimeAuthContext>>,
Path(worker_id): Path<String>,
Query(query): Query<RuntimeHttpUploadFileQuery>,
body: Bytes,
) -> RestResult<RuntimeHttpUploadedFileResponse> {
let worker_ref = worker_ref_for(&state.runtime, worker_id)?;
let context = match (
query.upload_id,
query.principal_id,
query.workspace_id,
query.runtime_id,
query.owner_worker_id,
) {
(None, None, None, None, None) => None,
(
Some(upload_id),
Some(principal_id),
Some(workspace_id),
Some(runtime_id),
Some(owner_worker_id),
) => {
if owner_worker_id != worker_ref.worker_id.to_string() {
return Err(RuntimeHttpRestError::new(
StatusCode::FORBIDDEN,
"uploaded_file_owner_mismatch",
"uploaded file context does not match the target Worker",
));
}
Some(session_store::UploadedFileUploadContext {
upload_id,
principal_id,
workspace_id,
runtime_id,
worker_id: owner_worker_id,
})
}
_ => {
return Err(RuntimeHttpRestError::new(
StatusCode::BAD_REQUEST,
"uploaded_file_context_incomplete",
"uploaded file context fields must be provided together",
));
}
};
let file = match auth_workspace_scope(&state, auth.as_ref())? {
Some(scope) => {
if context
.as_ref()
.is_some_and(|context| context.workspace_id != scope.workspace_id)
{
return Err(RuntimeHttpRestError::new(
StatusCode::FORBIDDEN,
"uploaded_file_workspace_mismatch",
"uploaded file context does not match the authenticated Workspace",
));
}
match context.as_ref() {
Some(context) => state.runtime.upload_worker_file_with_context_scoped(
&scope,
&worker_ref,
&query.file_name,
&query.media_type,
&body,
context,
),
None => state.runtime.upload_worker_file_scoped(
&scope,
&worker_ref,
&query.file_name,
&query.media_type,
&body,
),
}
}
None => state.runtime.upload_worker_file(
&worker_ref,
&query.file_name,
&query.media_type,
&body,
),
}
.map_err(RuntimeHttpRestError::runtime)?;
Ok(Json(RuntimeHttpUploadedFileResponse { file }))
}
async fn delete_worker_uploaded_file(
State(state): State<RuntimeHttpState>,
auth: Option<Extension<RuntimeAuthContext>>,
Path((worker_id, artifact_id)): Path<(String, String)>,
) -> RestResult<RuntimeHttpUploadedFileDeleteResponse> {
let worker_ref = worker_ref_for(&state.runtime, worker_id)?;
match auth_workspace_scope(&state, auth.as_ref())? {
Some(scope) => {
state
.runtime
.delete_worker_uploaded_file_scoped(&scope, &worker_ref, &artifact_id)
}
None => state
.runtime
.delete_worker_uploaded_file(&worker_ref, &artifact_id),
}
.map_err(RuntimeHttpRestError::runtime)?;
Ok(Json(RuntimeHttpUploadedFileDeleteResponse {
deleted: true,
}))
}
async fn stop_worker(
State(state): State<RuntimeHttpState>,
auth: Option<Extension<RuntimeAuthContext>>,
@@ -1620,7 +1767,7 @@ fn required_runtime_permission(method: &Method, path: &str) -> Option<&'static s
if path.ends_with("/workspace-api") {
return Some("workers:create");
}
if path.ends_with("/input") || path.ends_with("/restore") {
if path.ends_with("/input") || path.ends_with("/restore") || path.contains("/attachments") {
return Some("workers:input");
}
if path.ends_with("/stop") || path.ends_with("/cancel") {
@@ -1881,6 +2028,21 @@ mod tests {
WorkdirPath, WorkdirSessionCapabilities,
};
#[test]
fn attachment_routes_require_worker_input_permission() {
assert_eq!(
required_runtime_permission(&Method::POST, "/v1/workers/7/attachments"),
Some("workers:input")
);
assert_eq!(
required_runtime_permission(
&Method::DELETE,
"/v1/workers/7/attachments/019ca7c8-57b6-7f05-8edf-524147aba7b3"
),
Some("workers:input")
);
}
fn test_bundle(profile: ProfileSelector) -> ConfigBundle {
ConfigBundle {
metadata: ConfigBundleMetadata {
+3 -107
View File
@@ -1,105 +1,9 @@
use serde::{Deserialize, Deserializer, Serialize, Serializer, de};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::{fmt, str::FromStr};
use uuid::{Uuid, Version};
pub use protocol::{WorkerId, WorkerIdParseError};
pub use workdir::workspace::RuntimeWorkerRef;
/// Stable Workspace-owned Worker identity.
///
/// Runtime placement is deliberately not part of this value. New identities are
/// allocated by Workspace authority before a Runtime create request is sent.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct WorkerId(Uuid);
impl WorkerId {
pub fn now_v7() -> Self {
Self(Uuid::now_v7())
}
/// Converts a legacy Runtime-local numeric id into a syntactically valid
/// migration-only UUIDv7 value. New Worker allocation must use `now_v7`.
pub fn from_legacy_u64(value: u64) -> Self {
let mut bytes = [0_u8; 16];
bytes[8..].copy_from_slice(&value.to_be_bytes());
bytes[6] = 0x70;
bytes[8] = (bytes[8] & 0x3f) | 0x80;
Self(Uuid::from_bytes(bytes))
}
pub fn from_legacy_binding(workspace_id: &str, runtime_id: &str, value: u64) -> Self {
let mut hasher = Sha256::new();
hasher.update(b"yoi.workspace-worker-id.v1\0");
hasher.update(workspace_id.as_bytes());
hasher.update([0]);
hasher.update(runtime_id.as_bytes());
hasher.update([0]);
hasher.update(value.to_be_bytes());
let digest = hasher.finalize();
let mut bytes = [0_u8; 16];
bytes.copy_from_slice(&digest[..16]);
// Migrated ids sort before normally allocated UUIDv7 values while retaining
// deterministic collision-resistant payload bits.
bytes[..6].fill(0);
bytes[6] = (bytes[6] & 0x0f) | 0x70;
bytes[8] = (bytes[8] & 0x3f) | 0x80;
Self(Uuid::from_bytes(bytes))
}
pub fn parse(value: &str) -> Option<Self> {
let value = Uuid::parse_str(value).ok()?;
(value.get_version() == Some(Version::SortRand)).then_some(Self(value))
}
pub const fn as_uuid(self) -> Uuid {
self.0
}
}
impl fmt::Display for WorkerId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt(formatter)
}
}
impl FromStr for WorkerId {
type Err = WorkerIdParseError;
fn from_str(value: &str) -> Result<Self, Self::Err> {
Self::parse(value).ok_or(WorkerIdParseError)
}
}
impl Serialize for WorkerId {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_str(&self.to_string())
}
}
impl<'de> Deserialize<'de> for WorkerId {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let value = String::deserialize(deserializer)?;
Self::parse(&value).ok_or_else(|| de::Error::custom("Worker id must be a UUIDv7"))
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct WorkerIdParseError;
impl fmt::Display for WorkerIdParseError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("Worker id must be a UUIDv7")
}
}
impl std::error::Error for WorkerIdParseError {}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct LegacyWorkerIdentityMapping {
pub workspace_id: String,
@@ -140,7 +44,7 @@ pub fn legacy_worker_identity_mapping_digest(mappings: &[LegacyWorkerIdentityMap
}
/// Runtime-local authority reference for Worker operations. The contained id is
/// nevertheless the Workspace-owned stable identity; the Runtime does not mint it.
/// nevertheless the stable Worker identity; the Runtime does not mint it.
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct WorkerRef {
pub worker_id: WorkerId,
@@ -164,14 +68,6 @@ impl TryFrom<&RuntimeWorkerRef> for WorkerRef {
mod tests {
use super::*;
#[test]
fn worker_id_accepts_only_uuid_v7() {
let worker_id = WorkerId::now_v7();
assert_eq!(WorkerId::parse(&worker_id.to_string()), Some(worker_id));
assert!(WorkerId::parse("30").is_none());
assert!(WorkerId::parse(&Uuid::nil().to_string()).is_none());
}
#[test]
fn runtime_worker_ref_preserves_stable_worker_identity() {
let worker_id = WorkerId::now_v7();
+1
View File
@@ -33,3 +33,4 @@ pub mod working_directory;
pub use fs_store::{FsRuntimeStore, FsRuntimeStoreOptions};
pub use management::RuntimeOptions;
pub use runtime::{Runtime, RuntimeWorkspaceScope};
pub use session_store::UploadedFileUploadContext;
+135
View File
@@ -1205,6 +1205,138 @@ impl Runtime {
})
}
/// Store a client-local file in the owning Worker session before input submit.
pub fn upload_worker_file_scoped(
&self,
scope: &RuntimeWorkspaceScope,
worker_ref: &WorkerRef,
file_name: &str,
media_type: &str,
content: &[u8],
) -> Result<protocol::UploadedFileRef, RuntimeError> {
self.ensure_worker_in_workspace(scope, worker_ref)?;
self.upload_worker_file(worker_ref, file_name, media_type, content)
}
pub fn upload_worker_file(
&self,
worker_ref: &WorkerRef,
file_name: &str,
media_type: &str,
content: &[u8],
) -> Result<protocol::UploadedFileRef, RuntimeError> {
self.upload_worker_file_inner(worker_ref, file_name, media_type, content, None)
}
pub fn upload_worker_file_with_context_scoped(
&self,
scope: &RuntimeWorkspaceScope,
worker_ref: &WorkerRef,
file_name: &str,
media_type: &str,
content: &[u8],
context: &session_store::UploadedFileUploadContext,
) -> Result<protocol::UploadedFileRef, RuntimeError> {
self.ensure_worker_in_workspace(scope, worker_ref)?;
self.upload_worker_file_inner(worker_ref, file_name, media_type, content, Some(context))
}
pub fn upload_worker_file_with_context(
&self,
worker_ref: &WorkerRef,
file_name: &str,
media_type: &str,
content: &[u8],
context: &session_store::UploadedFileUploadContext,
) -> Result<protocol::UploadedFileRef, RuntimeError> {
self.upload_worker_file_inner(worker_ref, file_name, media_type, content, Some(context))
}
fn upload_worker_file_inner(
&self,
worker_ref: &WorkerRef,
file_name: &str,
media_type: &str,
content: &[u8],
context: Option<&session_store::UploadedFileUploadContext>,
) -> Result<protocol::UploadedFileRef, RuntimeError> {
let (backend, handle) = {
let state = self.lock()?;
state.ensure_running()?;
state.ensure_worker_ref(worker_ref)?;
let worker = state.worker(worker_ref)?;
match (
state.execution_backend.clone(),
worker.execution_handle.clone(),
) {
(Some(backend), Some(handle)) => (backend, handle),
_ => {
return Err(RuntimeError::WorkerExecutionUnavailable {
worker_id: worker_ref.worker_id.clone(),
message: "worker has no live execution handle".to_string(),
});
}
}
};
backend
.upload_file(&handle, file_name, media_type, content, context)
.map_err(|result| RuntimeError::WorkerExecutionRejected {
worker_id: worker_ref.worker_id.clone(),
operation: result.operation,
outcome: result.outcome,
message: result.message_or_default(),
result,
})
}
/// Delete an unsubmitted uploaded file from the owning Worker session.
pub fn delete_worker_uploaded_file_scoped(
&self,
scope: &RuntimeWorkspaceScope,
worker_ref: &WorkerRef,
artifact_id: &str,
) -> Result<(), RuntimeError> {
self.ensure_worker_in_workspace(scope, worker_ref)?;
self.delete_worker_uploaded_file(worker_ref, artifact_id)
}
pub fn delete_worker_uploaded_file(
&self,
worker_ref: &WorkerRef,
artifact_id: &str,
) -> Result<(), RuntimeError> {
let (backend, handle) = {
let state = self.lock()?;
state.ensure_running()?;
state.ensure_worker_ref(worker_ref)?;
let worker = state.worker(worker_ref)?;
match (
state.execution_backend.clone(),
worker.execution_handle.clone(),
) {
(Some(backend), Some(handle)) => (backend, handle),
_ => {
return Err(RuntimeError::WorkerExecutionUnavailable {
worker_id: worker_ref.worker_id.clone(),
message: "worker has no live execution handle".to_string(),
});
}
}
};
let result = backend.delete_uploaded_file(&handle, artifact_id);
if result.is_accepted() {
Ok(())
} else {
Err(RuntimeError::WorkerExecutionRejected {
worker_id: worker_ref.worker_id.clone(),
operation: result.operation,
outcome: result.outcome,
message: result.message_or_default(),
result,
})
}
}
/// Return live completion entries through a workspace-scoped Runtime authorization context.
pub fn worker_completions_scoped(
&self,
@@ -1387,6 +1519,8 @@ impl Runtime {
WorkerExecutionOperation::Spawn
| WorkerExecutionOperation::Restore
| WorkerExecutionOperation::Input
| WorkerExecutionOperation::UploadFile
| WorkerExecutionOperation::DeleteUploadedFile
| WorkerExecutionOperation::ProtocolMethod => return Ok(()),
};
if result.is_accepted() {
@@ -2402,6 +2536,7 @@ impl RuntimeState {
display_name: worker.request.display_name.clone(),
profile,
repository_id,
repository_key: None,
working_directory_id,
})
}
+71 -5
View File
@@ -60,6 +60,7 @@ use worker::{
Worker, WorkerBootstrap, WorkerBootstrapError, WorkerBootstrapLayout,
WorkerControllerTransport, WorkerError, WorkerFilesystemAuthority, WorkerHandle,
WorkerSharedState, WorkerWorkspaceContext, WorkspaceClient, WorkspaceId,
bash_output_dir_for_worker_id,
};
const DEFAULT_BACKEND_ID: &str = "worker-crate";
@@ -886,6 +887,7 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
let run_dir = worker_aggregate_dir
.join("runs")
.join(request.run_generation.to_string());
let bash_output_dir = bash_output_dir_for_worker_id(&request.worker_ref.worker_id);
let mut prepared = WorkerBootstrap::new(
manifest,
store,
@@ -894,6 +896,7 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
filesystem_authority,
WorkerBootstrapLayout::RuntimeManagedRun {
run_dir: run_dir.clone(),
bash_output_dir,
},
self.controller_transport,
)
@@ -1131,10 +1134,12 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
let run_dir = worker_aggregate_dir
.join("runs")
.join(request.run_generation.to_string());
let bash_output_dir = bash_output_dir_for_worker_id(&request.worker_ref.worker_id);
let started = PreparedWorker::new(
worker,
WorkerBootstrapLayout::RuntimeManagedRun {
run_dir: run_dir.clone(),
bash_output_dir,
},
self.controller_transport,
)
@@ -2020,6 +2025,56 @@ where
result
}
fn upload_file(
&self,
handle: &WorkerExecutionHandle,
file_name: &str,
media_type: &str,
content: &[u8],
context: Option<&session_store::UploadedFileUploadContext>,
) -> Result<protocol::UploadedFileRef, WorkerExecutionResult> {
let (worker, _, _) = self.get_execution(handle).map_err(|mut result| {
result.operation = WorkerExecutionOperation::UploadFile;
result
})?;
let uploaded = match context {
Some(context) => {
worker.upload_file_with_context(file_name, media_type, content, context)
}
None => worker.upload_file(file_name, media_type, content),
};
uploaded.map_err(|error| {
WorkerExecutionResult::rejected(
WorkerExecutionOperation::UploadFile,
format!("uploaded_file_rejected: {error}"),
)
})
}
fn delete_uploaded_file(
&self,
handle: &WorkerExecutionHandle,
artifact_id: &str,
) -> WorkerExecutionResult {
let (worker, _, _) = match self.get_execution(handle) {
Ok(execution) => execution,
Err(mut result) => {
result.operation = WorkerExecutionOperation::DeleteUploadedFile;
return result;
}
};
match worker.delete_uploaded_file(artifact_id) {
Ok(_) => WorkerExecutionResult::accepted(
WorkerExecutionOperation::DeleteUploadedFile,
WorkerExecutionRunState::Idle,
),
Err(error) => WorkerExecutionResult::rejected(
WorkerExecutionOperation::DeleteUploadedFile,
format!("uploaded_file_delete_rejected: {error}"),
),
}
}
fn dispatch_method(
&self,
handle: &WorkerExecutionHandle,
@@ -2110,6 +2165,7 @@ where
"execution handle does not reference a live Worker",
);
};
let artifact_cleanup = execution.handle.clone();
let shutdown = execution.shutdown.clone();
let result = self.send_method(
WorkerExecutionOperation::Stop,
@@ -2129,7 +2185,13 @@ where
}
Ok(())
}) {
Ok(()) => result,
Ok(()) => match artifact_cleanup.delete_uncommitted_uploaded_files() {
Ok(_) => result,
Err(error) => WorkerExecutionResult::errored(
WorkerExecutionOperation::Stop,
format!("uploaded_file_cleanup_failed: {error}"),
),
},
Err(message) => WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message),
}
}
@@ -2552,10 +2614,14 @@ mod tests {
)
.await
.map_err(|err| err.to_string())?;
let (handle, shutdown_rx) =
WorkerController::spawn_runtime_managed(worker, &self.runtime_base)
.await
.map_err(|err| err.to_string())?;
let bash_output_dir = self.runtime_base.join("bash-output");
let (handle, shutdown_rx) = WorkerController::spawn_runtime_managed(
worker,
&self.runtime_base,
&bash_output_dir,
)
.await
.map_err(|err| err.to_string())?;
Ok(RuntimeWorkerController {
handle,
shutdown: Arc::new(tokio::sync::Mutex::new(Some(shutdown_rx))),
+72 -10
View File
@@ -134,6 +134,8 @@ pub trait EmbeddedWorkerMutationDispatcher: Send + Sync {
enum RuntimeWorkerMutationTransport {
Remote {
base_url: String,
request_source_signer: RuntimeRequestSourceSigner,
request_source_audience: String,
},
Embedded {
dispatcher: Arc<dyn EmbeddedWorkerMutationDispatcher>,
@@ -157,10 +159,12 @@ impl RuntimeWorkerMutationForwarder {
) -> Self {
Self {
authority: RuntimeWorkerMutationSourceAuthority::remote(identity),
scope,
scope: scope.clone(),
source_worker_id: source_worker_id.into(),
transport: RuntimeWorkerMutationTransport::Remote {
base_url: base_url.into().trim_end_matches('/').to_string(),
request_source_signer: RuntimeRequestSourceSigner::from_identity(identity),
request_source_audience: scope.server_id,
},
}
}
@@ -197,11 +201,18 @@ impl RuntimeWorkerMutationForwarder {
)?;
match (&self.transport, proof) {
(
RuntimeWorkerMutationTransport::Remote { base_url },
RuntimeWorkerMutationTransport::Remote {
base_url,
request_source_signer,
request_source_audience,
},
RuntimeOwnedWorkerMutationProof::Remote(token),
) => execute_remote_worker_remove_http(RemoteWorkerRemoveHttpRequest {
base_url: base_url.clone(),
workspace_id: self.scope.workspace_id.clone(),
source_worker_id: self.source_worker_id.clone(),
request_source_signer: request_source_signer.clone(),
request_source_audience: request_source_audience.clone(),
token,
target_runtime_id: target_runtime_id.to_string(),
target_worker_id: target_worker_id.to_string(),
@@ -224,6 +235,9 @@ impl RuntimeWorkerMutationForwarder {
struct RemoteWorkerRemoveHttpRequest {
base_url: String,
workspace_id: String,
source_worker_id: String,
request_source_signer: RuntimeRequestSourceSigner,
request_source_audience: String,
token: String,
target_runtime_id: String,
target_worker_id: String,
@@ -256,23 +270,35 @@ fn execute_remote_worker_remove_http(
fn execute_remote_worker_remove_http_blocking(
request: RemoteWorkerRemoveHttpRequest,
) -> Result<WorkspaceResponse, RuntimeWorkerMutationForwardError> {
let url = format!(
"{}/api/w/{}/workers/remove",
request.base_url, request.workspace_id
);
let body = serde_json::json!({
let path = format!("/api/w/{}/workers/remove", request.workspace_id);
let url = format!("{}{}", request.base_url, path);
let body = serde_json::to_string(&serde_json::json!({
"target_runtime_id": request.target_runtime_id,
"target_worker_id": request.target_worker_id,
"reason": request.reason,
});
}))
.map_err(|error| RuntimeWorkerMutationForwardError::Transport(error.to_string()))?;
let request_source_proof = request.request_source_signer.issue(
&request.request_source_audience,
&request.workspace_id,
Some(&request.source_worker_id),
WORKSPACE_REQUEST_PERMISSION,
"POST",
&path,
body.as_bytes(),
i64::try_from(unix_now_seconds()).unwrap_or(i64::MAX),
30,
)?;
let client = reqwest::blocking::Client::new();
let response = client
.post(url)
.header(RUNTIME_REQUEST_SOURCE_PROOF_HEADER, request_source_proof)
.header(
crate::auth::WORKER_MUTATION_SOURCE_PROOF_HEADER,
request.token,
)
.json(&body)
.header(reqwest::header::CONTENT_TYPE, "application/json")
.body(body)
.send()
.map_err(|error| RuntimeWorkerMutationForwardError::Transport(error.to_string()))?;
let status = response.status().as_u16();
@@ -697,7 +723,8 @@ mod tests {
use super::*;
use crate::auth::{
WorkerMutationSourceExpectation, decode_runtime_request_source_claims,
decode_worker_mutation_source_claims, verify_worker_mutation_source_proof,
decode_worker_mutation_source_claims, request_body_digest,
verify_worker_mutation_source_proof,
};
#[test]
@@ -1119,6 +1146,41 @@ mod tests {
assert!(request.contains("\"target_worker_id\":\"worker-target\""));
assert!(!request.contains("expected_worker_revision"));
assert!(request.contains("\"reason\":\"retire obsolete Worker\""));
let request_source_token = request
.lines()
.find_map(|line| {
line.split_once(':').and_then(|(name, value)| {
name.eq_ignore_ascii_case(RUNTIME_REQUEST_SOURCE_PROOF_HEADER)
.then(|| value.trim())
})
})
.expect("runtime request source proof header");
let request_source_claims =
decode_runtime_request_source_claims(request_source_token).unwrap();
assert_eq!(request_source_claims.iss, "runtime-a");
assert_eq!(request_source_claims.aud, "server-a");
assert_eq!(request_source_claims.workspace_id, "workspace-a");
assert_eq!(
request_source_claims.worker_id.as_deref(),
Some("worker-source")
);
assert_eq!(
request_source_claims.permission,
WORKSPACE_REQUEST_PERMISSION
);
assert_eq!(request_source_claims.method, "POST");
assert_eq!(
request_source_claims.path,
"/api/w/workspace-a/workers/remove"
);
let request_body = request
.split_once("\r\n\r\n")
.map(|(_, body)| body)
.expect("WorkerRemove request body");
assert_eq!(
request_source_claims.body_digest,
request_body_digest(request_body.as_bytes())
);
let token = request
.lines()
.find_map(|line| {
+3 -1
View File
@@ -47,7 +47,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let worker = worker::Worker::from_manifest_toml(&toml, store).await?;
let runtime_tmp = tempfile::tempdir()?;
let (handle, _shutdown_rx) = WorkerController::spawn(worker, runtime_tmp.path()).await?;
let bash_output_dir = runtime_tmp.path().join("bash-output");
let (handle, _shutdown_rx) =
WorkerController::spawn(worker, runtime_tmp.path(), &bash_output_dir).await?;
// Check initial status via shared state
println!("[shared_state] {}", handle.shared_state.status_json());
+62 -7
View File
@@ -17,9 +17,28 @@ use manifest::WorkerManifest;
#[derive(Debug, Clone)]
pub enum WorkerBootstrapLayout {
/// A direct Worker rooted below the supplied runtime base directory.
Direct { runtime_base: PathBuf },
Direct {
runtime_base: PathBuf,
bash_output_dir: PathBuf,
},
/// A runtime-managed Worker with an exact persisted run directory.
RuntimeManagedRun { run_dir: PathBuf },
RuntimeManagedRun {
run_dir: PathBuf,
bash_output_dir: PathBuf,
},
}
/// Return the temporary Bash spill directory owned by a stable Worker identity.
///
/// The directory deliberately lives outside session/run-generation storage so a
/// restarted controller for the same Worker keeps the same readable artifact
/// boundary.
pub fn bash_output_dir_for_worker_id(worker_id: impl std::fmt::Display) -> PathBuf {
std::env::temp_dir()
.join("yoi")
.join("workers")
.join(worker_id.to_string())
.join("bash-output")
}
/// Construction and controller inputs that are stable for one Worker launch.
@@ -204,12 +223,29 @@ where
{
let cleanup_session = worker.workdir_session().cloned();
let controller = match layout {
WorkerBootstrapLayout::Direct { runtime_base } => {
WorkerController::spawn_with_transport(worker, &runtime_base, transport).await
WorkerBootstrapLayout::Direct {
runtime_base,
bash_output_dir,
} => {
WorkerController::spawn_with_transport(
worker,
&runtime_base,
&bash_output_dir,
transport,
)
.await
}
WorkerBootstrapLayout::RuntimeManagedRun { run_dir } => {
WorkerController::spawn_runtime_managed_run_with_transport(worker, &run_dir, transport)
.await
WorkerBootstrapLayout::RuntimeManagedRun {
run_dir,
bash_output_dir,
} => {
WorkerController::spawn_runtime_managed_run_with_transport(
worker,
&run_dir,
&bash_output_dir,
transport,
)
.await
}
};
@@ -227,3 +263,22 @@ where
}
}
}
#[cfg(test)]
mod tests {
use super::bash_output_dir_for_worker_id;
#[test]
fn bash_output_directory_is_stable_per_worker_below_system_temp() {
let path = bash_output_dir_for_worker_id("019c1234-worker");
assert_eq!(
path,
std::env::temp_dir()
.join("yoi")
.join("workers")
.join("019c1234-worker")
.join("bash-output")
);
}
}
+213 -113
View File
@@ -31,7 +31,8 @@ use protocol::{
AlertLevel, AlertSource, CommandEvent as ProtocolCommandEvent,
CommandSnapshot as ProtocolCommandSnapshot, CommandStatus as ProtocolCommandStatus,
CommandStream as ProtocolCommandStream, CommandStreamSlice as ProtocolCommandStreamSlice,
ErrorCode, Event, Method, RewindTargetId, RunResult, Segment, TurnResult, WorkerStatus,
ErrorCode, Event, Method, RewindTargetId, RunResult, Segment, TurnResult, UploadedFileRef,
WorkerStatus,
};
use workdir::{
CommandEvent as WorkdirCommandEvent, CommandSnapshot as WorkdirCommandSnapshot,
@@ -45,16 +46,18 @@ use workdir::{
#[derive(Clone)]
pub struct WorkerHandle {
method_tx: mpsc::Sender<Method>,
event_tx: broadcast::Sender<Event>,
working_event_tx: broadcast::Sender<Event>,
pub shared_state: Arc<WorkerSharedState>,
pub runtime_dir: Arc<RuntimeDir>,
pub alerter: Alerter,
pub in_flight: InFlightEvents,
/// Segment-log mirror + broadcast handle. The IPC server snapshots
/// Segment-log mirror + session-entry channel. The IPC server snapshots
/// it on every new connection (Event::Snapshot) and forwards
/// subsequent commits (Event::Entry) on the receiver.
pub sink: SegmentLogSink,
spawned_registry: Arc<SpawnedWorkerRegistry>,
artifact_store: Arc<dyn Store>,
session_id: session_store::SessionId,
}
impl WorkerHandle {
@@ -62,8 +65,53 @@ impl WorkerHandle {
self.method_tx.send(method).await
}
pub fn upload_file(
&self,
file_name: &str,
media_type: &str,
content: &[u8],
) -> Result<UploadedFileRef, session_store::StoreError> {
self.artifact_store.write_uploaded_file(
self.session_id,
file_name,
media_type,
content,
session_store::UploadedFileLimits::default(),
)
}
pub fn upload_file_with_context(
&self,
file_name: &str,
media_type: &str,
content: &[u8],
context: &session_store::UploadedFileUploadContext,
) -> Result<UploadedFileRef, session_store::StoreError> {
self.artifact_store.write_uploaded_file_with_context(
self.session_id,
file_name,
media_type,
content,
context,
session_store::UploadedFileLimits::default(),
)
}
pub fn delete_uploaded_file(
&self,
artifact_id: &str,
) -> Result<bool, session_store::StoreError> {
self.artifact_store
.delete_uploaded_file(self.session_id, artifact_id)
}
pub fn delete_uncommitted_uploaded_files(&self) -> Result<u64, session_store::StoreError> {
self.artifact_store
.delete_uncommitted_uploaded_files(self.session_id)
}
pub fn subscribe(&self) -> broadcast::Receiver<Event> {
self.event_tx.subscribe()
self.working_event_tx.subscribe()
}
pub fn committed_entries(&self) -> Vec<LogEntry> {
@@ -117,7 +165,7 @@ impl WorkerHandle {
/// Broadcast an event to all listeners (including socket clients).
pub fn send_event(&self, event: Event) -> Result<usize, broadcast::error::SendError<Event>> {
self.event_tx.send(event)
self.working_event_tx.send(event)
}
/// Emit a user-facing alert. Thin wrapper over `Alerter::alert`.
@@ -129,19 +177,19 @@ impl WorkerHandle {
async fn set_controller_status(
shared_state: &Arc<WorkerSharedState>,
runtime_dir: &RuntimeDir,
event_tx: &broadcast::Sender<Event>,
working_event_tx: &broadcast::Sender<Event>,
status: WorkerStatus,
) {
shared_state.set_status(status);
let _ = runtime_dir.write_status(shared_state).await;
let _ = event_tx.send(Event::Status { status });
let _ = working_event_tx.send(Event::Status { status });
}
async fn finish_controller_run<C, St>(
worker: &mut Worker<C, St>,
shared_state: &Arc<WorkerSharedState>,
runtime_dir: &RuntimeDir,
event_tx: &broadcast::Sender<Event>,
working_event_tx: &broadcast::Sender<Event>,
new_status: WorkerStatus,
) where
C: LlmClient + Clone + 'static,
@@ -157,7 +205,7 @@ async fn finish_controller_run<C, St>(
// the terminal run boundary so reconnect snapshots cannot append stale
// partial text/tool arguments after newer entries.
worker.clear_in_flight_events();
set_controller_status(shared_state, runtime_dir, event_tx, new_status).await;
set_controller_status(shared_state, runtime_dir, working_event_tx, new_status).await;
worker.spawn_post_run_memory_jobs();
}
@@ -222,6 +270,7 @@ impl WorkerController {
pub async fn spawn<C, St>(
worker: Worker<C, St>,
runtime_base: &Path,
bash_output_dir: &Path,
) -> Result<(WorkerHandle, ShutdownReceiver), std::io::Error>
where
C: LlmClient + Clone + 'static,
@@ -230,6 +279,7 @@ impl WorkerController {
Self::spawn_inner(
worker,
runtime_base,
bash_output_dir,
false,
None,
WorkerControllerTransport::UnixSocket,
@@ -242,24 +292,9 @@ impl WorkerController {
pub async fn spawn_with_transport<C, St>(
worker: Worker<C, St>,
runtime_base: &Path,
bash_output_dir: &Path,
transport: WorkerControllerTransport,
) -> Result<(WorkerHandle, ShutdownReceiver), std::io::Error>
where
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
{
Self::spawn_inner(worker, runtime_base, false, None, transport).await
}
/// Spawn a Worker owned by `worker-runtime`.
///
/// The controller still uses an ephemeral directory for Unix sockets and
/// tool spill artifacts, but does not write legacy pid/status/manifest
/// liveness projections.
pub async fn spawn_runtime_managed<C, St>(
worker: Worker<C, St>,
runtime_base: &Path,
) -> Result<(WorkerHandle, ShutdownReceiver), std::io::Error>
where
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
@@ -267,6 +302,33 @@ impl WorkerController {
Self::spawn_inner(
worker,
runtime_base,
bash_output_dir,
false,
None,
transport,
)
.await
}
/// Spawn a Worker owned by `worker-runtime`.
///
/// The controller uses an ephemeral directory for Unix sockets while tool
/// spill artifacts use the separately supplied Worker-owned temporary path.
/// Runtime-managed Workers do not write legacy pid/status/manifest liveness
/// projections.
pub async fn spawn_runtime_managed<C, St>(
worker: Worker<C, St>,
runtime_base: &Path,
bash_output_dir: &Path,
) -> Result<(WorkerHandle, ShutdownReceiver), std::io::Error>
where
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
{
Self::spawn_inner(
worker,
runtime_base,
bash_output_dir,
true,
None,
WorkerControllerTransport::UnixSocket,
@@ -278,6 +340,7 @@ impl WorkerController {
pub async fn spawn_runtime_managed_run<C, St>(
worker: Worker<C, St>,
run_dir: &Path,
bash_output_dir: &Path,
) -> Result<(WorkerHandle, ShutdownReceiver), std::io::Error>
where
C: LlmClient + Clone + 'static,
@@ -286,6 +349,7 @@ impl WorkerController {
Self::spawn_runtime_managed_run_with_transport(
worker,
run_dir,
bash_output_dir,
WorkerControllerTransport::UnixSocket,
)
.await
@@ -296,6 +360,7 @@ impl WorkerController {
pub async fn spawn_runtime_managed_run_with_transport<C, St>(
worker: Worker<C, St>,
run_dir: &Path,
bash_output_dir: &Path,
transport: WorkerControllerTransport,
) -> Result<(WorkerHandle, ShutdownReceiver), std::io::Error>
where
@@ -305,12 +370,21 @@ impl WorkerController {
let parent = run_dir
.parent()
.ok_or_else(|| std::io::Error::other("run path has no parent"))?;
Self::spawn_inner(worker, parent, true, Some(run_dir), transport).await
Self::spawn_inner(
worker,
parent,
bash_output_dir,
true,
Some(run_dir),
transport,
)
.await
}
async fn spawn_inner<C, St>(
worker: Worker<C, St>,
runtime_base: &Path,
bash_output_dir: &Path,
runtime_managed: bool,
runtime_run: Option<&Path>,
transport: WorkerControllerTransport,
@@ -323,6 +397,7 @@ impl WorkerController {
let result = Self::spawn_initialized(
worker,
runtime_base,
bash_output_dir,
runtime_managed,
runtime_run,
transport,
@@ -340,6 +415,7 @@ impl WorkerController {
async fn spawn_initialized<C, St>(
mut worker: Worker<C, St>,
runtime_base: &Path,
bash_output_dir: &Path,
runtime_managed: bool,
runtime_run: Option<&Path>,
transport: WorkerControllerTransport,
@@ -353,9 +429,9 @@ impl WorkerController {
// bash-output scope) ===
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
let (method_tx, method_rx) = mpsc::channel::<Method>(32);
let (event_tx, _) = broadcast::channel::<Event>(256);
let alerter = Alerter::new(event_tx.clone());
let in_flight = InFlightEvents::new(event_tx.clone());
let (working_event_tx, _) = broadcast::channel::<Event>(256);
let alerter = Alerter::new(working_event_tx.clone());
let in_flight = InFlightEvents::new(working_event_tx.clone());
worker.attach_in_flight_events(in_flight.clone());
// Runtime directory is created before tool registration because it owns
@@ -395,13 +471,13 @@ impl WorkerController {
// Also hand the raw broadcast sender so Worker-internal operations
// can emit typed lifecycle `Event`s (currently: compact progress).
worker.attach_internal_worker_registry(spawned_registry.clone());
worker.attach_event_tx(event_tx.clone());
worker.attach_working_event_tx(working_event_tx.clone());
// Bash spills long outputs to a per-worker subdir under the runtime
// dir. Push a recursive `allow(Read)` for that path into the
// Worker's runtime scope so the agent can `Read` saved files
// without polluting the workspace.
let bash_output_dir = runtime_dir.path().join("bash-output");
// Bash spill artifacts are owned by the stable Worker identity rather
// than a controller session/run generation. Push a recursive
// `allow(Read)` for the exact tool output path into the Worker's shared
// runtime scope so the Workdir session and system prompt stay aligned.
let bash_output_dir = bash_output_dir.to_path_buf();
std::fs::create_dir_all(&bash_output_dir).map_err(|e| {
std::io::Error::other(format!(
"create bash output dir {}: {e}",
@@ -430,7 +506,7 @@ impl WorkerController {
worker.wire_history_persistence();
// === 2. Engine event bridge wiring ===
wire_event_bridges_on_engine(&mut worker, &event_tx, &alerter, &in_flight);
wire_event_bridges_on_engine(&mut worker, &working_event_tx, &alerter, &in_flight);
// === 3. Tool registration (builtin / memory / spawn-orchestration) ===
let fs_for_view = register_worker_tools(
@@ -475,15 +551,19 @@ impl WorkerController {
runtime_dir.write_manifest(&manifest_toml).await?;
runtime_dir.write_status(&shared_state).await?;
let artifact_store: Arc<dyn Store> = Arc::new(worker.store().clone());
let session_id = worker.session_id();
let handle = WorkerHandle {
method_tx,
event_tx: event_tx.clone(),
working_event_tx: working_event_tx.clone(),
shared_state: shared_state.clone(),
runtime_dir: runtime_dir.clone(),
alerter: alerter.clone(),
in_flight: in_flight.clone(),
sink: worker.sink(),
spawned_registry: spawned_registry.clone(),
artifact_store,
session_id,
};
let socket_server = match transport {
@@ -502,7 +582,7 @@ impl WorkerController {
tokio::spawn(controller_loop(
worker,
method_rx,
event_tx,
working_event_tx,
shared_state,
runtime_dir,
cancel_tx,
@@ -640,7 +720,7 @@ fn protocol_command_status(status: WorkdirCommandStatus) -> ProtocolCommandStatu
}
/// Wire the per-event broadcast bridges on the Worker's Engine. Each callback
/// re-publishes a worker-level signal as a `protocol::Event` on `event_tx`
/// re-publishes a worker-level signal as a `protocol::Event` on `working_event_tx`
/// so subscribers (TUI, socket clients) get a single typed stream.
///
/// `Worker::wire_history_persistence` is called separately to wire the
@@ -649,7 +729,7 @@ fn protocol_command_status(status: WorkdirCommandStatus) -> ProtocolCommandStatu
/// / `AnnotatedToolResult` commit through the sync writer.
pub(crate) fn wire_event_bridges_on_engine<C, St>(
worker: &mut Worker<C, St>,
event_tx: &broadcast::Sender<Event>,
working_event_tx: &broadcast::Sender<Event>,
alerter: &Alerter,
in_flight: &InFlightEvents,
) where
@@ -659,12 +739,12 @@ pub(crate) fn wire_event_bridges_on_engine<C, St>(
let ai_activity = worker.ai_activity_counter();
let worker = worker.engine_mut();
let tx = event_tx.clone();
let tx = working_event_tx.clone();
worker.on_turn_start(move |turn| {
let _ = tx.send(Event::TurnStart { turn });
});
let tx = event_tx.clone();
let tx = working_event_tx.clone();
worker.on_turn_end(move |turn| {
let _ = tx.send(Event::TurnEnd {
turn,
@@ -672,17 +752,17 @@ pub(crate) fn wire_event_bridges_on_engine<C, St>(
});
});
let tx = event_tx.clone();
let tx = working_event_tx.clone();
worker.on_llm_call_start(move |llm_call| {
let _ = tx.send(Event::LlmCallStart { llm_call });
});
let tx = event_tx.clone();
let tx = working_event_tx.clone();
worker.on_llm_call_end(move |llm_call| {
let _ = tx.send(Event::LlmCallEnd { llm_call });
});
let tx = event_tx.clone();
let tx = working_event_tx.clone();
worker.on_llm_retry(move |llm_call, notice| {
let _ = tx.send(Event::LlmRetry {
llm_call,
@@ -695,7 +775,7 @@ pub(crate) fn wire_event_bridges_on_engine<C, St>(
});
});
let tx = event_tx.clone();
let tx = working_event_tx.clone();
worker.on_llm_continuation(move |llm_call, attempt, max_attempts, reason| {
let _ = tx.send(Event::LlmContinuation {
llm_call,
@@ -768,7 +848,7 @@ pub(crate) fn wire_event_bridges_on_engine<C, St>(
});
});
let tx = event_tx.clone();
let tx = working_event_tx.clone();
let activity = ai_activity.clone();
worker.on_tool_result(move |result| {
activity.fetch_add(1, Ordering::SeqCst);
@@ -793,7 +873,7 @@ pub(crate) fn wire_event_bridges_on_engine<C, St>(
});
});
let tx = event_tx.clone();
let tx = working_event_tx.clone();
worker.on_usage(move |event| {
let _ = tx.send(Event::Usage {
input_tokens: event.input_tokens,
@@ -802,7 +882,7 @@ pub(crate) fn wire_event_bridges_on_engine<C, St>(
});
});
let tx = event_tx.clone();
let tx = working_event_tx.clone();
worker.on_error(move |event| {
let _ = tx.send(Event::Error {
code: ErrorCode::ProviderError,
@@ -870,6 +950,20 @@ where
crate::spawn::tool::ParentNotificationTarget::Buffer(worker.notify_buffer_handle())
});
let prompts = worker.prompts().clone();
let paste_store = worker.store().clone();
let paste_session_id = worker.session_id();
worker
.engine_mut()
.register_tool(crate::paste_artifact_tool::search_input_artifact_tool(
paste_store.clone(),
paste_session_id,
));
worker
.engine_mut()
.register_tool(crate::paste_artifact_tool::read_input_artifact_tool(
paste_store,
paste_session_id,
));
// Resolve the existing WorkerWorkdir binding into the domain provider.
// Tools only consume the provider handle; they do not own its root, cwd,
// scope, or lifecycle. No-workdir Workers expose no local tools.
@@ -880,7 +974,7 @@ where
.register_tools(tools::core_builtin_tools(
workdir.clone(),
tracker.clone(),
bash_output_dir,
bash_output_dir.clone(),
));
if feature_config.image.enabled && model_supports_image_attachments(&spawner_manifest.model)
{
@@ -1103,6 +1197,7 @@ where
spawner_workspace_context,
parent_notifications,
runtime_base.clone(),
bash_output_dir.clone(),
spawner_workspace_root,
source_workdir_session,
spawned_registry.clone(),
@@ -1156,7 +1251,7 @@ where
async fn controller_loop<C, St>(
mut worker: Worker<C, St>,
mut method_rx: mpsc::Receiver<Method>,
event_tx: broadcast::Sender<Event>,
working_event_tx: broadcast::Sender<Event>,
shared_state: Arc<WorkerSharedState>,
runtime_dir: Arc<RuntimeDir>,
cancel_tx: mpsc::Sender<()>,
@@ -1213,7 +1308,7 @@ async fn controller_loop<C, St>(
set_controller_status(
&shared_state,
&runtime_dir,
&event_tx,
&working_event_tx,
WorkerStatus::Running,
)
.await;
@@ -1230,7 +1325,7 @@ async fn controller_loop<C, St>(
},
),
&mut method_rx,
&event_tx,
&working_event_tx,
&cancel_tx,
&pause_tx,
&shared_state,
@@ -1255,7 +1350,7 @@ async fn controller_loop<C, St>(
},
),
&mut method_rx,
&event_tx,
&working_event_tx,
&cancel_tx,
&pause_tx,
&shared_state,
@@ -1273,7 +1368,7 @@ async fn controller_loop<C, St>(
drive_turn(
worker.run_for_notification(kind),
&mut method_rx,
&event_tx,
&working_event_tx,
&cancel_tx,
&pause_tx,
&shared_state,
@@ -1291,7 +1386,7 @@ async fn controller_loop<C, St>(
drive_turn(
worker.resume(),
&mut method_rx,
&event_tx,
&working_event_tx,
&cancel_tx,
&pause_tx,
&shared_state,
@@ -1315,16 +1410,16 @@ async fn controller_loop<C, St>(
&mut worker,
&shared_state,
&runtime_dir,
&event_tx,
&working_event_tx,
new_status,
)
.await;
if shutdown {
let _ = event_tx.send(Event::Shutdown);
let _ = working_event_tx.send(Event::Shutdown);
break;
}
if take_shutdown_request_after_status(&shutdown_after_idle, new_status) {
let _ = event_tx.send(Event::Shutdown);
let _ = working_event_tx.send(Event::Shutdown);
break;
}
continue;
@@ -1342,7 +1437,7 @@ async fn controller_loop<C, St>(
// already rejects `Run` while a turn is live, so
// this branch is only reachable across a race window
// around status flips.
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn".into(),
});
@@ -1393,7 +1488,7 @@ async fn controller_loop<C, St>(
Method::Resume => {
if shared_state.get_status() != WorkerStatus::Paused {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::NotPaused,
message: "Worker is not paused".into(),
});
@@ -1409,20 +1504,20 @@ async fn controller_loop<C, St>(
set_controller_status(
&shared_state,
&runtime_dir,
&event_tx,
&working_event_tx,
WorkerStatus::Idle,
)
.await;
}
Err(error) => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: worker_error_code(&error),
message: error.to_string(),
});
}
},
WorkerStatus::Idle | WorkerStatus::Stopped => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::NotRunning,
message: "Worker is not running".into(),
});
@@ -1439,7 +1534,7 @@ async fn controller_loop<C, St>(
// Worker is Idle (Running turns go through `drive_turn`,
// not this outer match), so there is nothing to pause.
if shared_state.get_status() != WorkerStatus::Paused {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::NotRunning,
message: "Worker is not running".into(),
});
@@ -1449,21 +1544,21 @@ async fn controller_loop<C, St>(
Method::Compact => match shared_state.get_status() {
WorkerStatus::Idle => {
if let Err(error) = worker.manual_compact().await {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: worker_error_code(&error),
message: error.to_string(),
});
}
}
WorkerStatus::Paused => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: "Cannot compact while the Worker is paused; resume or start a fresh turn first"
.into(),
});
}
WorkerStatus::Running | WorkerStatus::Stopped => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message:
"Worker is already executing a turn; compact can only run while idle"
@@ -1474,10 +1569,10 @@ async fn controller_loop<C, St>(
Method::ListRewindTargets => match shared_state.get_status() {
WorkerStatus::Idle | WorkerStatus::Paused => {
emit_rewind_targets(&worker, &event_tx)
emit_rewind_targets(&worker, &working_event_tx)
}
WorkerStatus::Running | WorkerStatus::Stopped => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn; rewind can only run while idle or paused"
.into(),
@@ -1490,23 +1585,28 @@ async fn controller_loop<C, St>(
expected_head_entries,
} => match shared_state.get_status() {
WorkerStatus::Idle => {
if apply_rewind(&mut worker, &event_tx, target, expected_head_entries) {
if apply_rewind(
&mut worker,
&working_event_tx,
target,
expected_head_entries,
) {
worker.clear_in_flight_events();
shared_state.set_status(WorkerStatus::Idle);
let _ = event_tx.send(Event::Status {
let _ = working_event_tx.send(Event::Status {
status: WorkerStatus::Idle,
});
}
}
WorkerStatus::Paused => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: "Cannot apply rewind while the Worker is paused; resume or wait for idle first"
.into(),
});
}
WorkerStatus::Running | WorkerStatus::Stopped => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn; rewind can only run while idle or paused"
.into(),
@@ -1515,24 +1615,24 @@ async fn controller_loop<C, St>(
},
Method::Shutdown => {
let _ = event_tx.send(Event::Shutdown);
let _ = working_event_tx.send(Event::Shutdown);
break;
}
Method::ListWorkers => match discovery.list_visible().await {
Ok(workers) => match serde_json::to_value(workers) {
Ok(workers) => {
let _ = event_tx.send(Event::WorkersListed { workers });
let _ = working_event_tx.send(Event::WorkersListed { workers });
}
Err(error) => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::Internal,
message: format!("serialize visible workers: {error}"),
});
}
},
Err(error) => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: error.to_string(),
});
@@ -1542,17 +1642,17 @@ async fn controller_loop<C, St>(
Method::RestoreWorker { name } => match discovery.restore(&name).await {
Ok(result) => match serde_json::to_value(result) {
Ok(result) => {
let _ = event_tx.send(Event::WorkerRestored { result });
let _ = working_event_tx.send(Event::WorkerRestored { result });
}
Err(error) => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::Internal,
message: format!("serialize worker restore result: {error}"),
});
}
},
Err(error) => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: error.to_string(),
});
@@ -1562,17 +1662,17 @@ async fn controller_loop<C, St>(
Method::RegisterPeer { name } => match discovery.register_peer(&name) {
Ok(result) => match serde_json::to_value(result) {
Ok(result) => {
let _ = event_tx.send(Event::PeerRegistered { result });
let _ = working_event_tx.send(Event::PeerRegistered { result });
}
Err(error) => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::Internal,
message: format!("serialize peer registration result: {error}"),
});
}
},
Err(error) => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: error.to_string(),
});
@@ -1691,7 +1791,7 @@ async fn handle_inbound_worker_event(
async fn drive_turn<F>(
worker_future: F,
method_rx: &mut mpsc::Receiver<Method>,
event_tx: &broadcast::Sender<Event>,
working_event_tx: &broadcast::Sender<Event>,
cancel_tx: &mpsc::Sender<()>,
pause_tx: &mpsc::Sender<()>,
shared_state: &Arc<WorkerSharedState>,
@@ -1727,7 +1827,7 @@ where
set_controller_status(
shared_state,
runtime_dir,
event_tx,
working_event_tx,
WorkerStatus::Running,
)
.await;
@@ -1745,11 +1845,11 @@ where
WorkerRunResult::LimitReached => (WorkerStatus::Idle, RunResult::LimitReached),
WorkerRunResult::RolledBack => (WorkerStatus::Idle, RunResult::RolledBack),
WorkerRunResult::Interrupted { .. } if pause_requested => {
let _ = event_tx.send(Event::RunEnd { result: RunResult::Paused });
let _ = working_event_tx.send(Event::RunEnd { result: RunResult::Paused });
return (WorkerStatus::Paused, shutdown_requested);
}
WorkerRunResult::Interrupted { code, message } => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code,
message: message.clone(),
});
@@ -1765,7 +1865,7 @@ where
return (WorkerStatus::Idle, shutdown_requested);
}
};
let _ = event_tx.send(Event::RunEnd { result: run_result });
let _ = working_event_tx.send(Event::RunEnd { result: run_result });
if parent_originated && matches!(run_result, RunResult::Finished) {
crate::ipc::event::fire_and_forget(
parent_socket.cloned(),
@@ -1782,13 +1882,13 @@ where
// intentionally skip `WorkerEvent::Errored` upward:
// that channel is reserved for worker runtime
// failures, not deliberate interruptions.
let _ = event_tx.send(Event::RunEnd { result: RunResult::Paused });
let _ = working_event_tx.send(Event::RunEnd { result: RunResult::Paused });
(WorkerStatus::Paused, shutdown_requested)
}
Err(e) => {
let code = worker_error_code(&e);
let message = e.to_string();
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code,
message: message.clone(),
});
@@ -1819,13 +1919,13 @@ where
let _ = cancel_tx.try_send(());
}
Some(Method::Run { .. } | Method::RunTracked { .. } | Method::Resume) => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn".into(),
});
}
Some(Method::Compact | Method::ListRewindTargets | Method::RewindTo { .. }) => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn; rewind/compact can only run while idle or paused"
.into(),
@@ -1839,7 +1939,7 @@ where
}
Some(Method::ListCompletions { .. }) => {}
Some(Method::ListWorkers | Method::RestoreWorker { .. } | Method::RegisterPeer { .. }) => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Worker discovery/control requests are only handled while the Worker is idle or paused"
.into(),
@@ -1872,20 +1972,20 @@ where
}
}
fn emit_rewind_targets<C, St>(worker: &Worker<C, St>, event_tx: &broadcast::Sender<Event>)
fn emit_rewind_targets<C, St>(worker: &Worker<C, St>, working_event_tx: &broadcast::Sender<Event>)
where
C: LlmClient + 'static,
St: Store,
{
match worker.list_rewind_targets() {
Ok((head_entries, targets)) => {
let _ = event_tx.send(Event::RewindTargets {
let _ = working_event_tx.send(Event::RewindTargets {
head_entries,
targets,
});
}
Err(err) => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::Internal,
message: err.to_string(),
});
@@ -1895,7 +1995,7 @@ where
fn apply_rewind<C, St>(
worker: &mut Worker<C, St>,
event_tx: &broadcast::Sender<Event>,
working_event_tx: &broadcast::Sender<Event>,
target: RewindTargetId,
expected_head_entries: usize,
) -> bool
@@ -1907,7 +2007,7 @@ where
Ok(applied) => {
let session =
session_store::public_snapshot::project_current_session_snapshot(&applied.entries);
let _ = event_tx.send(Event::RewindApplied {
let _ = working_event_tx.send(Event::RewindApplied {
session,
input: applied.input,
summary: applied.summary,
@@ -1915,7 +2015,7 @@ where
true
}
Err(err) => {
let _ = event_tx.send(Event::Error {
let _ = working_event_tx.send(Event::Error {
code: ErrorCode::InvalidRequest,
message: err.to_string(),
});
@@ -2049,7 +2149,7 @@ mod tests {
// would observe channel-closed and confuse the select! arm.
_method_tx: mpsc::Sender<Method>,
method_rx: mpsc::Receiver<Method>,
event_tx: broadcast::Sender<Event>,
working_event_tx: broadcast::Sender<Event>,
cancel_tx: mpsc::Sender<()>,
_cancel_rx: mpsc::Receiver<()>,
pause_tx: mpsc::Sender<()>,
@@ -2070,7 +2170,7 @@ mod tests {
.expect("runtime dir create"),
);
let (method_tx, method_rx) = mpsc::channel::<Method>(16);
let (event_tx, _) = broadcast::channel::<Event>(16);
let (working_event_tx, _) = broadcast::channel::<Event>(16);
let (cancel_tx, cancel_rx) = mpsc::channel::<()>(1);
let (pause_tx, pause_rx) = mpsc::channel::<()>(1);
let shared_state = Arc::new(WorkerSharedState::new(
@@ -2095,7 +2195,7 @@ mod tests {
DriveTurnEnv {
_method_tx: method_tx,
method_rx,
event_tx,
working_event_tx,
cancel_tx,
_cancel_rx: cancel_rx,
pause_tx,
@@ -2157,7 +2257,7 @@ mod tests {
let (status, shutdown) = drive_turn(
worker_future,
&mut env.method_rx,
&env.event_tx,
&env.working_event_tx,
&env.cancel_tx,
&env.pause_tx,
&env.shared_state,
@@ -2200,7 +2300,7 @@ mod tests {
let (status, shutdown) = drive_turn(
worker_future,
&mut env.method_rx,
&env.event_tx,
&env.working_event_tx,
&env.cancel_tx,
&env.pause_tx,
&env.shared_state,
@@ -2230,7 +2330,7 @@ mod tests {
let (status, _) = drive_turn(
worker_future,
&mut env.method_rx,
&env.event_tx,
&env.working_event_tx,
&env.cancel_tx,
&env.pause_tx,
&env.shared_state,
@@ -2268,7 +2368,7 @@ mod tests {
let (status, _) = drive_turn(
worker_future,
&mut env.method_rx,
&env.event_tx,
&env.working_event_tx,
&env.cancel_tx,
&env.pause_tx,
&env.shared_state,
@@ -2312,7 +2412,7 @@ mod tests {
let (status, _) = drive_turn(
worker_future,
&mut env.method_rx,
&env.event_tx,
&env.working_event_tx,
&env.cancel_tx,
&env.pause_tx,
&env.shared_state,
@@ -2354,7 +2454,7 @@ mod tests {
let (status, shutdown) = drive_turn(
worker_future,
&mut env.method_rx,
&env.event_tx,
&env.working_event_tx,
&env.cancel_tx,
&env.pause_tx,
&env.shared_state,
@@ -2393,7 +2493,7 @@ mod tests {
let (status, shutdown) = drive_turn(
worker_future,
&mut env.method_rx,
&env.event_tx,
&env.working_event_tx,
&env.cancel_tx,
&env.pause_tx,
&env.shared_state,
@@ -2430,7 +2530,7 @@ mod tests {
let (status, shutdown) = drive_turn(
worker_future,
&mut env.method_rx,
&env.event_tx,
&env.working_event_tx,
&env.cancel_tx,
&env.pause_tx,
&env.shared_state,
@@ -2453,7 +2553,7 @@ mod tests {
#[tokio::test]
async fn compact_method_is_rejected_while_running() {
let mut env = make_env().await;
let mut events = env.event_tx.subscribe();
let mut events = env.working_event_tx.subscribe();
env._method_tx
.send(Method::Compact)
.await
@@ -2466,7 +2566,7 @@ mod tests {
let (status, shutdown) = drive_turn(
worker_future,
&mut env.method_rx,
&env.event_tx,
&env.working_event_tx,
&env.cancel_tx,
&env.pause_tx,
&env.shared_state,
+2
View File
@@ -634,10 +634,12 @@ async fn run_cli_inner(cli: Cli) -> ExitCode {
return ExitCode::FAILURE;
}
};
let bash_output_dir = crate::bash_output_dir_for_worker_id(&worker_name);
let started = match start_worker_controller(
worker,
WorkerBootstrapLayout::Direct {
runtime_base: runtime_base.clone(),
bash_output_dir,
},
WorkerControllerTransport::UnixSocket,
)
@@ -12,11 +12,7 @@ use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::json;
use workdir::http::{WorkdirSessionOperation, WorkdirSessionOperationResult};
use workdir::workspace::{
WorkingDirectoryDetailResponse as WorkdirDetailResponse,
WorkingDirectoryListResponse as WorkdirListResponse, WorkspaceWorkdirSessionFence,
WorkspaceWorkdirSessionOperationRequest,
};
use workdir::workspace::{WorkspaceWorkdirSessionFence, WorkspaceWorkdirSessionOperationRequest};
use workdir::{
CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandStatus, EditRequest,
EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult,
@@ -24,6 +20,14 @@ use workdir::{
WorkdirSessionCapabilities, WorkdirSessionHandle, WriteRequest, WriteResult,
};
use workspace_api::{
WorkingDirectoryCreateRequest as WorkdirCreateRequest,
WorkingDirectoryCreateResponse as WorkdirCreateResponse,
WorkingDirectoryListResponse as WorkdirListResponse,
WorkingDirectoryRemovalRequest as WorkdirRemovalRequest,
WorkingDirectoryRemovalResponse as WorkdirRemovalResponse,
};
use crate::feature::{
FeatureDescriptor, FeatureInstallContext, FeatureInstallError, FeatureModule, ToolContribution,
ToolDeclaration,
@@ -48,7 +52,7 @@ const LIST_DESCRIPTION: &str = "List persistent Workdirs in the current Workspac
const CREATE_DESCRIPTION: &str = "Materialize a persistent Workdir on a selected Runtime from a Workspace repository and optional selector. This does not change this Worker's attachment; use WorkdirAttach explicitly after creation.";
const ATTACH_DESCRIPTION: &str = "Attach this Worker to one existing Workdir. The Backend enforces one active Workdir per Worker and one active Worker per Workdir, then opens an ephemeral operation session.";
const DETACH_DESCRIPTION: &str = "Detach this Worker from its active Workdir and release Workdir occupancy. Any ephemeral operation session is closed.";
const DELETE_DESCRIPTION: &str = "Delete one persistent Workdir by id through Backend Workspace API authority. Occupied, blocked, or dirty Workdirs requiring confirmation are rejected.";
const DELETE_DESCRIPTION: &str = "Request removal of one persistent Workdir by id through durable Backend Workspace authority. The input includes only the Workdir id and a bounded reason. The result reports removed, retained, or attention_required without exposing operation-table or provider internals.";
#[derive(Clone, Debug)]
pub struct ManageWorkdirFeature {
@@ -339,7 +343,13 @@ impl WorkdirSession for WorkspaceAttachedWorkdirSession {
}
}
async fn start_command(&self, request: CommandRequest) -> Result<CommandHandle, WorkdirError> {
async fn start_command(
&self,
mut request: CommandRequest,
) -> Result<CommandHandle, WorkdirError> {
// `spill_dir` belongs to the Worker host running BashTool. This remote
// WorkdirSession provider cannot safely resolve or write that host path.
request.spill_dir = None;
match self.operate(WorkdirSessionOperation::CommandStart(request))? {
WorkdirSessionOperationResult::CommandStart(result) => Ok(result),
_ => Err(Self::mismatch("command_start")),
@@ -413,16 +423,17 @@ impl WorkspaceHttpWorkdirBackend {
.as_deref()
.map(|value| validate_identity(value, CREATE_TOOL, "runtime_id"))
.transpose()?;
let repository_id = validate_identity(&input.repository_id, CREATE_TOOL, "repository_id")?;
let repository_key =
validate_identity(&input.repository_key, CREATE_TOOL, "repository_key")?;
let selector = validate_optional_selector(input.selector)?;
let workspace_id = encode_path_segment(self.workspace_id()?);
let request = WorkdirCreateRequest {
runtime_id: runtime_id.map(str::to_string),
repository_id: repository_id.to_string(),
repository_key: repository_key.to_string(),
selector,
operation_id,
operation_id: Some(operation_id),
};
let response = self.execute_json::<WorkdirDetailResponse>(WorkspaceRequest::json(
let response = self.execute_json::<WorkdirCreateResponse>(WorkspaceRequest::json(
WorkspaceRequestMethod::Post,
format!("/api/w/{workspace_id}/working-directories"),
serde_json::to_string(&request).map_err(decode_error)?,
@@ -475,12 +486,21 @@ impl WorkspaceHttpWorkdirBackend {
)?;
let workspace_id = encode_path_segment(self.workspace_id()?);
let workdir_path = encode_path_segment(workdir_id);
let response = self.execute_json::<WorkdirDetailResponse>(WorkspaceRequest {
method: WorkspaceRequestMethod::Delete,
path: format!("/api/w/{workspace_id}/working-directories/{workdir_path}"),
body: None,
})?;
workdir_output(format!("Deleted Workdir {workdir_id}"), &response)
let response = self.execute_json::<WorkdirRemovalResponse>(WorkspaceRequest::json(
WorkspaceRequestMethod::Delete,
format!("/api/w/{workspace_id}/working-directories/{workdir_path}"),
serde_json::to_string(&WorkdirRemovalRequest {
reason: validate_delete_reason(&input.reason)?.to_string(),
})
.map_err(decode_error)?,
))?;
workdir_output(
format!(
"Workdir {workdir_id} removal disposition: {:?}",
response.disposition
),
&response,
)
}
fn execute_json<T: for<'de> Deserialize<'de>>(
@@ -602,6 +622,17 @@ fn validate_identity<'a>(
Ok(value)
}
fn validate_delete_reason(reason: &str) -> Result<&str, ToolError> {
let reason = reason.trim();
if reason.is_empty() || reason.len() > 500 || reason.chars().any(char::is_control) {
return Err(ToolError::InvalidArgument(
"WorkdirDelete reason must be non-empty, contain no control characters, and be at most 500 bytes"
.to_string(),
));
}
Ok(reason)
}
fn validate_optional_selector(selector: Option<String>) -> Result<Option<String>, ToolError> {
let Some(selector) = selector else {
return Ok(None);
@@ -652,10 +683,10 @@ fn create_schema() -> serde_json::Value {
json!({
"type": "object",
"additionalProperties": false,
"required": ["repository_id"],
"required": ["repository_key"],
"properties": {
"runtime_id": {"type": ["string", "null"], "minLength": 1},
"repository_id": {"type": "string", "minLength": 1},
"repository_key": {"type": "string", "minLength": 1},
"selector": {"type": ["string", "null"], "minLength": 1}
}
})
@@ -680,9 +711,10 @@ fn delete_schema() -> serde_json::Value {
json!({
"type": "object",
"additionalProperties": false,
"required": ["working_directory_id"],
"required": ["working_directory_id", "reason"],
"properties": {
"working_directory_id": {"type": "string", "minLength": 1}
"working_directory_id": {"type": "string", "minLength": 1},
"reason": {"type": "string", "minLength": 1, "maxLength": 500}
}
})
}
@@ -696,21 +728,11 @@ struct WorkdirListInput {}
struct WorkdirCreateInput {
#[serde(default)]
runtime_id: Option<String>,
repository_id: String,
repository_key: String,
#[serde(default)]
selector: Option<String>,
}
#[derive(Debug, Serialize)]
struct WorkdirCreateRequest {
#[serde(skip_serializing_if = "Option::is_none")]
runtime_id: Option<String>,
repository_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
selector: Option<String>,
operation_id: String,
}
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct WorkdirAttachInput {
@@ -737,6 +759,7 @@ struct WorkdirAttachmentResponse {
#[serde(deny_unknown_fields)]
struct WorkdirDeleteInput {
working_directory_id: String,
reason: String,
}
#[cfg(test)]
@@ -818,14 +841,14 @@ mod tests {
fn workdir_json(id: &str) -> serde_json::Value {
json!({
"working_directory_id": id,
"repository_id": "main",
"repository_key": "main",
"creation_selector": "refs/heads/main",
"creation_ref": "0123456789abcdef",
"materializer_kind": "local_git_worktree",
"cleanup_target": {
"kind": "git_worktree",
"working_directory_id": id,
"repository_id": "main"
"repository_key": "main"
},
"status": "active",
"cleanliness": "clean",
@@ -951,7 +974,7 @@ mod tests {
#[test]
fn schemas_expose_identities_without_paths_or_session_handles() {
let create = create_schema();
assert_eq!(create["required"], json!(["repository_id"]));
assert_eq!(create["required"], json!(["repository_key"]));
assert_eq!(
create["properties"]["runtime_id"]["type"],
json!(["string", "null"])
@@ -960,7 +983,10 @@ mod tests {
assert!(create["properties"].get("session_id").is_none());
assert_eq!(attach_schema()["required"], json!(["workdir_id"]));
assert!(attach_schema()["properties"].get("session_id").is_none());
assert_eq!(delete_schema()["required"], json!(["working_directory_id"]));
assert_eq!(
delete_schema()["required"],
json!(["working_directory_id", "reason"])
);
}
#[test]
@@ -1000,15 +1026,9 @@ mod tests {
"attached": false
})),
response(json!({
"workspace_id": "workspace/test",
"runtime_id": "runtime/one",
"item": {
"working_directory_id": "wd-created",
"repository_id": "main",
"materializer_kind": "local_git_worktree",
"status": "not_found"
},
"diagnostics": []
"working_directory_id": "wd-created",
"disposition": "removed",
"retryable": false
})),
]));
let backend = WorkspaceHttpWorkdirBackend::new(client.clone());
@@ -1029,7 +1049,7 @@ mod tests {
.create(
WorkdirCreateInput {
runtime_id: Some("runtime/one".to_string()),
repository_id: "main".to_string(),
repository_key: "main".to_string(),
selector: Some("refs/heads/topic".to_string()),
},
"call-create-1".to_string(),
@@ -1053,6 +1073,7 @@ mod tests {
backend
.delete(WorkdirDeleteInput {
working_directory_id: "wd-created".to_string(),
reason: "remove stale Workdir".to_string(),
})
.unwrap();
@@ -1069,7 +1090,7 @@ mod tests {
assert_eq!(requests[1].method, WorkspaceRequestMethod::Post);
let body: serde_json::Value =
serde_json::from_str(requests[1].body.as_deref().unwrap()).unwrap();
assert_eq!(body["repository_id"], "main");
assert_eq!(body["repository_key"], "main");
assert_eq!(body["runtime_id"], "runtime/one");
assert_eq!(body["operation_id"], "call-create-1");
assert_eq!(body["selector"], "refs/heads/topic");
@@ -1088,6 +1109,9 @@ mod tests {
"/api/w/workspace%2Ftest/working-directories/wd-created"
);
assert_eq!(requests[4].method, WorkspaceRequestMethod::Delete);
let body: serde_json::Value =
serde_json::from_str(requests[4].body.as_deref().unwrap()).unwrap();
assert_eq!(body, json!({"reason": "remove stale Workdir"}));
}
#[tokio::test]
@@ -1118,6 +1142,42 @@ mod tests {
assert!(body.get("session_id").is_none());
}
#[tokio::test]
async fn attached_session_does_not_send_worker_local_bash_spill_path() {
let client = Arc::new(RecordingWorkspaceClient::new(vec![response(json!({
"operation": "command_start",
"result": "command-1"
}))]));
let session = WorkspaceAttachedWorkdirSession::handle(client.clone());
let handle = session
.start_command(CommandRequest {
command: "true".to_string(),
timeout_secs: 120,
output_limit: 1024,
spill_dir: Some("/worker-local/bash-output".into()),
tool_call_id: Some("call-1".to_string()),
})
.await
.unwrap();
assert_eq!(handle, CommandHandle("command-1".to_string()));
let requests = client.requests();
let body: serde_json::Value =
serde_json::from_str(requests[0].body.as_deref().unwrap()).unwrap();
assert_eq!(body["operation"]["operation"], "command_start");
assert_eq!(body["operation"]["request"]["command"], "true");
assert!(body["operation"]["request"]["spill_dir"].is_null());
assert_eq!(body["operation"]["request"]["tool_call_id"], "call-1");
assert!(
!requests[0]
.body
.as_deref()
.unwrap()
.contains("/worker-local/bash-output")
);
}
#[tokio::test]
async fn delegated_attached_session_carries_captured_fence_on_operations() {
let client = Arc::new(RecordingWorkspaceClient::new(vec![
@@ -1321,7 +1381,7 @@ mod tests {
.create(
WorkdirCreateInput {
runtime_id: None,
repository_id: "main".to_string(),
repository_key: "main".to_string(),
selector: None,
},
"call-default".to_string(),
@@ -1346,7 +1406,7 @@ mod tests {
.create(
WorkdirCreateInput {
runtime_id: Some(" ".to_string()),
repository_id: "main".to_string(),
repository_key: "main".to_string(),
selector: None,
},
"call-invalid".to_string(),
@@ -56,7 +56,7 @@ struct TicketInput {
#[derive(Debug, Deserialize, JsonSchema)]
struct OpenMergeRequestInput {
ticket: String,
repository_id: String,
repository_key: String,
selector_from: String,
selector_to: String,
#[serde(default)]
@@ -178,7 +178,7 @@ impl Tool for MergeRequestTool {
WorkspaceRequestMethod::Post,
format!("/api/w/{ws}/tickets/{}/merge-request", v.ticket),
Some(
json!({"repository_id":v.repository_id,"selector_from":v.selector_from,"selector_to":v.selector_to,"summary":v.summary}),
json!({"repository_key":v.repository_key,"selector_from":v.selector_from,"selector_to":v.selector_to,"summary":v.summary}),
),
)
}
+35 -33
View File
@@ -15,7 +15,7 @@ pub struct InFlightBlockId(u64);
#[derive(Debug, Clone)]
pub struct InFlightEvents {
inner: Arc<Mutex<InFlightInner>>,
event_tx: broadcast::Sender<Event>,
working_event_tx: broadcast::Sender<Event>,
}
#[derive(Debug)]
@@ -47,14 +47,14 @@ enum TrackedBlock {
}
impl InFlightEvents {
pub(crate) fn new(event_tx: broadcast::Sender<Event>) -> Self {
pub(crate) fn new(working_event_tx: broadcast::Sender<Event>) -> Self {
Self {
inner: Arc::new(Mutex::new(InFlightInner {
next_block_id: 1,
blocks: Vec::new(),
commands: Vec::new(),
})),
event_tx,
working_event_tx,
}
}
@@ -84,7 +84,7 @@ impl InFlightEvents {
current.push_str(&text);
*finished = false;
}
let _ = self.event_tx.send(Event::TextDelta { text });
let _ = self.working_event_tx.send(Event::TextDelta { text });
}
pub(crate) fn text_done(&self, block_id: InFlightBlockId, text: String) {
@@ -100,7 +100,7 @@ impl InFlightEvents {
}
*finished = true;
}
let _ = self.event_tx.send(Event::TextDone { text });
let _ = self.working_event_tx.send(Event::TextDone { text });
}
pub(crate) fn thinking_start(&self) -> InFlightBlockId {
@@ -111,7 +111,7 @@ impl InFlightEvents {
text: String::new(),
finished: false,
});
let _ = self.event_tx.send(Event::ThinkingStart);
let _ = self.working_event_tx.send(Event::ThinkingStart);
block_id
}
@@ -126,7 +126,7 @@ impl InFlightEvents {
current.push_str(&text);
*finished = false;
}
let _ = self.event_tx.send(Event::ThinkingDelta { text });
let _ = self.working_event_tx.send(Event::ThinkingDelta { text });
}
pub(crate) fn thinking_done(&self, block_id: InFlightBlockId, text: String) {
@@ -142,7 +142,7 @@ impl InFlightEvents {
}
*finished = true;
}
let _ = self.event_tx.send(Event::ThinkingDone { text });
let _ = self.working_event_tx.send(Event::ThinkingDone { text });
}
pub(crate) fn tool_call_start(&self, id: String, name: String) -> InFlightBlockId {
@@ -155,7 +155,9 @@ impl InFlightEvents {
args: String::new(),
state: InFlightToolCallState::Pending,
});
let _ = self.event_tx.send(Event::ToolCallStart { id, name });
let _ = self
.working_event_tx
.send(Event::ToolCallStart { id, name });
block_id
}
@@ -171,7 +173,7 @@ impl InFlightEvents {
*state = InFlightToolCallState::StreamingArgs;
}
let _ = self
.event_tx
.working_event_tx
.send(Event::ToolCallArgsDelta { id, json: delta });
}
@@ -191,7 +193,7 @@ impl InFlightEvents {
}
*state = InFlightToolCallState::Done;
}
let _ = self.event_tx.send(Event::ToolCallDone {
let _ = self.working_event_tx.send(Event::ToolCallDone {
id,
name,
arguments: args,
@@ -210,7 +212,7 @@ impl InFlightEvents {
pub(crate) fn publish_command_event(&self, event: CommandEvent) {
self.lock().apply_command_event(&event);
let _ = self.event_tx.send(Event::Command { event });
let _ = self.working_event_tx.send(Event::Command { event });
}
pub(crate) fn replace_command_snapshot(&self, commands: Vec<CommandSnapshot>) {
@@ -492,13 +494,13 @@ mod tests {
#[test]
fn snapshot_boundary_does_not_duplicate_or_gap_delta_sent_after_subscribe() {
let (event_tx, _) = broadcast::channel(16);
let in_flight = InFlightEvents::new(event_tx.clone());
let (working_event_tx, _) = broadcast::channel(16);
let in_flight = InFlightEvents::new(working_event_tx.clone());
let block_id = in_flight.start_text_block();
in_flight.text_delta(block_id, "hel".into());
let guard = in_flight.snapshot_guard();
let mut rx = event_tx.subscribe();
let mut rx = working_event_tx.subscribe();
let snapshot = snapshot_from_guard(&guard);
drop(guard);
@@ -526,9 +528,9 @@ mod tests {
use crate::segment_log_sink::SegmentLogSink;
use session_store::{LogEntry, LoggedRole};
let (event_tx, _) = broadcast::channel(16);
let (working_event_tx, _) = broadcast::channel(16);
let sink = SegmentLogSink::new();
let in_flight = InFlightEvents::new(event_tx);
let in_flight = InFlightEvents::new(working_event_tx);
let block_id = in_flight.start_text_block();
in_flight.text_delta(block_id, "done".into());
in_flight.text_done(block_id, "done".into());
@@ -580,9 +582,9 @@ mod tests {
use crate::segment_log_sink::SegmentLogSink;
use session_store::{LogEntry, LoggedRole};
let (event_tx, _) = broadcast::channel(16);
let (working_event_tx, _) = broadcast::channel(16);
let sink = SegmentLogSink::new();
let in_flight = InFlightEvents::new(event_tx);
let in_flight = InFlightEvents::new(working_event_tx);
let block_id = in_flight.start_text_block();
in_flight.text_delta(block_id, "done".into());
in_flight.text_done(block_id, "done".into());
@@ -615,8 +617,8 @@ mod tests {
#[test]
fn committed_item_clears_matching_in_flight_block() {
let (event_tx, _) = broadcast::channel(16);
let in_flight = InFlightEvents::new(event_tx);
let (working_event_tx, _) = broadcast::channel(16);
let in_flight = InFlightEvents::new(working_event_tx);
let block_id = in_flight.start_text_block();
in_flight.text_delta(block_id, "done".into());
in_flight.clear_for_committed_item_then(
@@ -635,8 +637,8 @@ mod tests {
#[test]
fn committed_reasoning_summary_clears_matching_in_flight_thinking_blocks() {
let (event_tx, _) = broadcast::channel(16);
let in_flight = InFlightEvents::new(event_tx);
let (working_event_tx, _) = broadcast::channel(16);
let in_flight = InFlightEvents::new(working_event_tx);
let first = in_flight.thinking_start();
in_flight.thinking_delta(first, "summary A".into());
in_flight.thinking_done(first, "".into());
@@ -660,8 +662,8 @@ mod tests {
#[test]
fn committed_encrypted_only_reasoning_clears_empty_finished_thinking_block() {
let (event_tx, _) = broadcast::channel(16);
let in_flight = InFlightEvents::new(event_tx);
let (working_event_tx, _) = broadcast::channel(16);
let in_flight = InFlightEvents::new(working_event_tx);
let first = in_flight.thinking_start();
in_flight.thinking_done(first, "".into());
let second = in_flight.thinking_start();
@@ -689,9 +691,9 @@ mod tests {
#[test]
fn command_events_are_bounded_and_recoverable_from_snapshot() {
let (event_tx, _) = broadcast::channel(16);
let mut rx = event_tx.subscribe();
let in_flight = InFlightEvents::new(event_tx);
let (working_event_tx, _) = broadcast::channel(16);
let mut rx = working_event_tx.subscribe();
let in_flight = InFlightEvents::new(working_event_tx);
in_flight.publish_command_event(CommandEvent::Started {
command_id: "command-1".into(),
tool_call_id: Some("tool-1".into()),
@@ -740,9 +742,9 @@ mod tests {
#[test]
fn clear_discards_uncommitted_blocks_without_protocol_event() {
let (event_tx, _) = broadcast::channel(16);
let mut rx = event_tx.subscribe();
let in_flight = InFlightEvents::new(event_tx);
let (working_event_tx, _) = broadcast::channel(16);
let mut rx = working_event_tx.subscribe();
let in_flight = InFlightEvents::new(working_event_tx);
let text = in_flight.start_text_block();
in_flight.text_delta(text, "stale".into());
let tool = in_flight.tool_call_start("call-1".into(), "Bash".into());
@@ -770,8 +772,8 @@ mod tests {
#[test]
fn snapshot_omits_empty_finished_thinking_blocks() {
let (event_tx, _) = broadcast::channel(16);
let in_flight = InFlightEvents::new(event_tx);
let (working_event_tx, _) = broadcast::channel(16);
let in_flight = InFlightEvents::new(working_event_tx);
let empty_finished = in_flight.thinking_start();
in_flight.thinking_done(empty_finished, "".into());
let empty_running = in_flight.thinking_start();
+1 -1
View File
@@ -752,7 +752,7 @@ pub(crate) async fn prepare_internal_worker_session(
}
let actor_in_flight = in_flight.clone();
worker.attach_alerter(alerter.clone());
worker.attach_event_tx(event_tx.clone());
worker.attach_working_event_tx(event_tx.clone());
worker.attach_in_flight_events(in_flight.clone());
wire_event_bridges_on_engine(&mut worker, &event_tx, &alerter, &in_flight);
+5 -5
View File
@@ -28,15 +28,15 @@ pub struct Alerter {
}
struct Inner {
event_tx: broadcast::Sender<Event>,
working_event_tx: broadcast::Sender<Event>,
buffer: Mutex<VecDeque<Alert>>,
}
impl Alerter {
pub fn new(event_tx: broadcast::Sender<Event>) -> Self {
pub fn new(working_event_tx: broadcast::Sender<Event>) -> Self {
Self {
inner: Arc::new(Inner {
event_tx,
working_event_tx,
buffer: Mutex::new(VecDeque::with_capacity(MAX_BUFFERED_ALERTS)),
}),
}
@@ -66,7 +66,7 @@ impl Alerter {
buf.pop_front();
}
buf.push_back(alert.clone());
let _ = self.inner.event_tx.send(Event::Alert(alert));
let _ = self.inner.working_event_tx.send(Event::Alert(alert));
}
}
@@ -81,7 +81,7 @@ impl Alerter {
.buffer
.lock()
.expect("alerter buffer mutex poisoned");
let rx = self.inner.event_tx.subscribe();
let rx = self.inner.working_event_tx.subscribe();
let snapshot: Vec<Alert> = buf.iter().cloned().collect();
(snapshot, rx)
}
+2 -1
View File
@@ -9,6 +9,7 @@ pub mod hook;
pub(crate) mod in_flight;
pub mod ipc;
pub mod model_client;
mod paste_artifact_tool;
pub mod prompt;
pub mod runtime;
pub mod runtime_command;
@@ -27,7 +28,7 @@ mod worker;
pub use bootstrap::{
BootstrappedWorker, PreparedWorker, WorkerBootstrap, WorkerBootstrapError,
WorkerBootstrapLayout, start_worker_controller,
WorkerBootstrapLayout, bash_output_dir_for_worker_id, start_worker_controller,
};
pub use compact::token_counter::{EstimateSource, SplitPoint, TokenEstimate};
pub use controller::{ShutdownReceiver, WorkerController, WorkerControllerTransport, WorkerHandle};
+443
View File
@@ -0,0 +1,443 @@
//! Bounded model-facing access to session-owned large paste artifacts.
use std::sync::Arc;
use agen::tool::{Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput};
use async_trait::async_trait;
use protocol::{PasteArtifactAvailability, PasteArtifactMediaType, PasteArtifactRef};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use session_store::{SessionId, Store, StoreError};
const MAX_QUERY_BYTES: usize = 256;
const DEFAULT_SEARCH_RESULTS: usize = 20;
const MAX_SEARCH_RESULTS: usize = 100;
const MAX_SNIPPET_CHARS: usize = 300;
const DEFAULT_READ_BYTES: usize = 8 * 1024;
const MAX_READ_BYTES: usize = 16 * 1024;
const SEARCH_DESCRIPTION: &str = "Search one large pasted-input artifact owned by the current Worker. Returns bounded matching line snippets; never returns the whole artifact.";
const READ_DESCRIPTION: &str = "Read a bounded UTF-8 byte range from one large pasted-input artifact owned by the current Worker. Use next_offset for repeated calls instead of requesting the whole artifact.";
#[derive(Clone)]
struct ArtifactAccess<St: Store + Clone> {
store: St,
session_id: SessionId,
}
#[derive(Debug, Deserialize, JsonSchema)]
struct SearchInputArtifactInput {
/// Opaque artifact id from a large-paste history reference.
artifact_id: String,
/// Literal case-sensitive text to find.
query: String,
/// Maximum matching lines to return (1..=100).
max_results: Option<usize>,
}
#[derive(Debug, Serialize)]
struct SearchInputArtifactOutput {
artifact_id: String,
matches: Vec<SearchMatch>,
truncated: bool,
}
#[derive(Debug, Serialize)]
struct SearchMatch {
line: u64,
byte_offset: u64,
snippet: String,
}
struct SearchInputArtifactTool<St: Store + Clone> {
access: ArtifactAccess<St>,
}
#[async_trait]
impl<St> Tool for SearchInputArtifactTool<St>
where
St: Store + Clone + Send + Sync + 'static,
{
async fn execute(
&self,
input_json: &str,
_context: ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
let input: SearchInputArtifactInput =
serde_json::from_str(input_json).map_err(|error| {
ToolError::InvalidArgument(format!("invalid SearchInputArtifact input: {error}"))
})?;
if input.query.is_empty() || input.query.len() > MAX_QUERY_BYTES {
return Err(ToolError::InvalidArgument(
"query must contain 1..=256 UTF-8 bytes".to_string(),
));
}
let max_results = input
.max_results
.unwrap_or(DEFAULT_SEARCH_RESULTS)
.clamp(1, MAX_SEARCH_RESULTS);
let (_, content) =
read_artifact_text(&self.access, &input.artifact_id).map_err(tool_store_error)?;
let mut matches = Vec::new();
let mut truncated = false;
let mut byte_offset = 0_u64;
for (index, raw_line) in content.split_inclusive('\n').enumerate() {
let line = raw_line.strip_suffix('\n').unwrap_or(raw_line);
if line.contains(&input.query) {
if matches.len() == max_results {
truncated = true;
break;
}
matches.push(SearchMatch {
line: index as u64 + 1,
byte_offset,
snippet: truncate_chars(line, MAX_SNIPPET_CHARS),
});
}
byte_offset += raw_line.len() as u64;
}
json_output(
format!("Found {} matching pasted-input line(s).", matches.len()),
&SearchInputArtifactOutput {
artifact_id: input.artifact_id,
matches,
truncated,
},
)
}
}
#[derive(Debug, Deserialize, JsonSchema)]
struct ReadInputArtifactInput {
/// Opaque artifact id from a large-paste history reference.
artifact_id: String,
/// UTF-8 byte offset to start reading. Defaults to 0 and must be a character boundary.
offset: Option<u64>,
/// Maximum UTF-8 bytes to return (4..=16384). Defaults to 8192.
max_bytes: Option<usize>,
}
#[derive(Debug, Serialize)]
struct ReadInputArtifactOutput {
artifact_id: String,
offset: u64,
content: String,
next_offset: Option<u64>,
truncated: bool,
}
struct ReadInputArtifactTool<St: Store + Clone> {
access: ArtifactAccess<St>,
}
#[async_trait]
impl<St> Tool for ReadInputArtifactTool<St>
where
St: Store + Clone + Send + Sync + 'static,
{
async fn execute(
&self,
input_json: &str,
_context: ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
let input: ReadInputArtifactInput = serde_json::from_str(input_json).map_err(|error| {
ToolError::InvalidArgument(format!("invalid ReadInputArtifact input: {error}"))
})?;
let offset = input.offset.unwrap_or(0);
let max_bytes = input
.max_bytes
.unwrap_or(DEFAULT_READ_BYTES)
.clamp(4, MAX_READ_BYTES);
let (_, content) =
read_artifact_text(&self.access, &input.artifact_id).map_err(tool_store_error)?;
let offset = usize::try_from(offset).map_err(|_| {
ToolError::InvalidArgument("offset exceeds the artifact size".to_string())
})?;
if offset > content.len() || !content.is_char_boundary(offset) {
return Err(ToolError::InvalidArgument(
"offset must be a UTF-8 character boundary within the artifact".to_string(),
));
}
let mut end = offset.saturating_add(max_bytes).min(content.len());
while end > offset && !content.is_char_boundary(end) {
end -= 1;
}
let output = content[offset..end].to_string();
let next_offset = (end < content.len()).then_some(end as u64);
let truncated = next_offset.is_some();
json_output(
format!("Read {} pasted-input byte(s).", output.len()),
&ReadInputArtifactOutput {
artifact_id: input.artifact_id,
offset: offset as u64,
content: output,
next_offset,
truncated,
},
)
}
}
pub(crate) fn search_input_artifact_tool<St>(store: St, session_id: SessionId) -> ToolDefinition
where
St: Store + Clone + Send + Sync + 'static,
{
Arc::new(move || {
let schema = serde_json::to_value(schemars::schema_for!(SearchInputArtifactInput))
.unwrap_or_else(|_| serde_json::json!({}));
let meta = ToolMeta::new("SearchInputArtifact")
.description(SEARCH_DESCRIPTION)
.input_schema(schema);
let tool: Arc<dyn Tool> = Arc::new(SearchInputArtifactTool {
access: ArtifactAccess {
store: store.clone(),
session_id,
},
});
(meta, tool)
})
}
pub(crate) fn read_input_artifact_tool<St>(store: St, session_id: SessionId) -> ToolDefinition
where
St: Store + Clone + Send + Sync + 'static,
{
Arc::new(move || {
let schema = serde_json::to_value(schemars::schema_for!(ReadInputArtifactInput))
.unwrap_or_else(|_| serde_json::json!({}));
let meta = ToolMeta::new("ReadInputArtifact")
.description(READ_DESCRIPTION)
.input_schema(schema);
let tool: Arc<dyn Tool> = Arc::new(ReadInputArtifactTool {
access: ArtifactAccess {
store: store.clone(),
session_id,
},
});
(meta, tool)
})
}
fn json_output(summary: String, value: &impl Serialize) -> Result<ToolOutput, ToolError> {
let content = serde_json::to_string(value)
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?;
Ok(ToolOutput {
summary,
content: Some(content),
attachments: Vec::new(),
})
}
fn read_artifact_text<St: Store + Clone>(
access: &ArtifactAccess<St>,
artifact_id: &str,
) -> Result<(PasteArtifactRef, String), StoreError> {
match access
.store
.read_paste_artifact(access.session_id, artifact_id)
{
Ok(result) => Ok(result),
Err(paste_error) => {
let (file, bytes) = match access
.store
.read_uploaded_file_by_id(access.session_id, artifact_id)
{
Ok(result) => result,
Err(_) => return Err(paste_error),
};
let content =
String::from_utf8(bytes).map_err(|_| StoreError::ArtifactIntegrityMismatch)?;
let char_count =
u64::try_from(content.chars().count()).map_err(|_| StoreError::ArtifactTooLarge)?;
let line_count =
u64::try_from(content.lines().count()).map_err(|_| StoreError::ArtifactTooLarge)?;
Ok((
PasteArtifactRef {
artifact_id: file.artifact_id,
created_at_ms: file.created_at_ms,
media_type: PasteArtifactMediaType::TextPlainUtf8,
availability: PasteArtifactAvailability::Available,
byte_len: file.byte_len,
char_count,
line_count,
sha256: file.sha256,
source_entry_id: file.source_entry_id.unwrap_or_default(),
},
content,
))
}
}
}
fn tool_store_error(error: StoreError) -> ToolError {
let message = match error {
StoreError::PasteArtifactNotFound(_) => "paste artifact not found",
StoreError::PasteArtifactIntegrity(_) | StoreError::Corrupt { .. } => {
"paste artifact failed its integrity check"
}
StoreError::PasteArtifactUnsupported => "paste artifact storage is unavailable",
_ => "paste artifact is unavailable",
};
ToolError::ExecutionFailed(message.to_string())
}
fn truncate_chars(value: &str, limit: usize) -> String {
let mut chars = value.chars();
let truncated = chars.by_ref().take(limit).collect::<String>();
if chars.next().is_some() {
format!("{truncated}")
} else {
truncated
}
}
#[cfg(test)]
mod tests {
use agen::tool::ToolExecutionContext;
use session_store::{FsStore, PasteArtifactLimits, Store, new_session_id};
use super::*;
#[tokio::test]
async fn read_input_artifact_reads_uploaded_text_but_rejects_binary_content() {
let temp = tempfile::TempDir::new().unwrap();
let store = FsStore::new(temp.path()).unwrap();
let owner = new_session_id();
let text = store
.write_uploaded_file(
owner,
"notes.md",
"text/markdown",
b"alpha\nbeta",
session_store::UploadedFileLimits::default(),
)
.unwrap();
let read = ReadInputArtifactTool {
access: ArtifactAccess {
store: store.clone(),
session_id: owner,
},
};
let output = read
.execute(
&serde_json::json!({
"artifact_id": text.artifact_id,
"offset": 0,
"max_bytes": 64
})
.to_string(),
ToolExecutionContext::default(),
)
.await
.unwrap();
let output: serde_json::Value =
serde_json::from_str(output.content.as_deref().unwrap()).unwrap();
assert_eq!(output["content"], "alpha\nbeta");
let binary = store
.write_uploaded_file(
owner,
"image.png",
"image/png",
b"\x89PNG\r\n\x1a\nbody",
session_store::UploadedFileLimits::default(),
)
.unwrap();
let error = read
.execute(
&serde_json::json!({
"artifact_id": binary.artifact_id,
"offset": 0,
"max_bytes": 64
})
.to_string(),
ToolExecutionContext::default(),
)
.await
.unwrap_err();
assert!(error.to_string().contains("unavailable"));
}
#[tokio::test]
async fn search_and_read_are_bounded_and_owner_scoped() {
let temp = tempfile::TempDir::new().unwrap();
let store = FsStore::new(temp.path()).unwrap();
let owner = new_session_id();
let other = new_session_id();
let content = (0..700)
.map(|index| format!("line {index}: needle {}", "x".repeat(80)))
.collect::<Vec<_>>()
.join("\n");
let artifact = store
.write_paste_artifact(owner, "entry-1", &content, PasteArtifactLimits::default())
.unwrap();
let search = SearchInputArtifactTool {
access: ArtifactAccess {
store: store.clone(),
session_id: owner,
},
};
let searched = search
.execute(
&serde_json::json!({
"artifact_id": artifact.artifact_id,
"query": "needle",
"max_results": 3
})
.to_string(),
ToolExecutionContext::default(),
)
.await
.unwrap();
let searched: serde_json::Value =
serde_json::from_str(searched.content.as_deref().unwrap()).unwrap();
assert_eq!(searched["matches"].as_array().unwrap().len(), 3);
assert_eq!(searched["truncated"], true);
let read = ReadInputArtifactTool {
access: ArtifactAccess {
store: store.clone(),
session_id: owner,
},
};
let read_output = read
.execute(
&serde_json::json!({
"artifact_id": artifact.artifact_id,
"offset": 2,
"max_bytes": 999999
})
.to_string(),
ToolExecutionContext::default(),
)
.await
.unwrap();
let read_output: serde_json::Value =
serde_json::from_str(read_output.content.as_deref().unwrap()).unwrap();
assert!(read_output["content"].as_str().unwrap().len() <= MAX_READ_BYTES);
assert_eq!(read_output["truncated"], true);
assert!(read_output["next_offset"].as_u64().is_some());
let foreign = ReadInputArtifactTool {
access: ArtifactAccess {
store,
session_id: other,
},
};
let error = foreign
.execute(
&serde_json::json!({
"artifact_id": artifact.artifact_id,
"offset": 0,
"max_bytes": 1
})
.to_string(),
ToolExecutionContext::default(),
)
.await
.unwrap_err();
assert!(error.to_string().contains("paste artifact not found"));
}
}
+18
View File
@@ -920,4 +920,22 @@ mod tests {
);
}
}
#[test]
fn builtin_orchestrator_cleanup_policy_renders_with_common_includes() {
let rendered = PromptCatalog::builtins_only()
.unwrap()
.render_name("role.orchestrator", Value::UNDEFINED)
.unwrap();
assert!(rendered.contains("This policy governs naming only"));
assert!(rendered.contains("Coder cleanup is a separate post-completion decision"));
assert!(rendered.contains("perform one cleanup pass before ending the orchestration turn"));
assert!(rendered.contains("Never predeclare `delete_on_completion`"));
assert!(rendered.contains("call `WorkerStop`"));
assert!(rendered.contains("call `WorkerRemove`"));
assert!(rendered.contains("only then call `WorkdirDelete`"));
assert!(rendered.contains("`CurrentAssignment` means unassign and reread"));
assert!(!rendered.contains("{% include"));
}
}
+9 -9
View File
@@ -51,7 +51,7 @@ struct SinkInner {
/// survives session swaps so existing subscribers keep their
/// receiver — they observe the swap as a freshly broadcast
/// `LogEntry::AnnotatedSegmentStart` and reset their view accordingly.
broadcast_tx: broadcast::Sender<LogEntry>,
session_entry_tx: broadcast::Sender<LogEntry>,
}
impl SegmentLogSink {
@@ -59,11 +59,11 @@ impl SegmentLogSink {
/// has been written (deferred SegmentStart) or as a placeholder in
/// tests.
pub fn new() -> Self {
let (broadcast_tx, _) = broadcast::channel(BROADCAST_CAPACITY);
let (session_entry_tx, _) = broadcast::channel(BROADCAST_CAPACITY);
Self {
inner: Arc::new(SinkInner {
mirror: StdMutex::new(Vec::new()),
broadcast_tx,
session_entry_tx,
}),
}
}
@@ -72,11 +72,11 @@ impl SegmentLogSink {
/// Used by restore / fork-at-restore code paths that materialise
/// the existing log before the sink starts taking new commits.
pub fn with_initial(entries: Vec<LogEntry>) -> Self {
let (broadcast_tx, _) = broadcast::channel(BROADCAST_CAPACITY);
let (session_entry_tx, _) = broadcast::channel(BROADCAST_CAPACITY);
Self {
inner: Arc::new(SinkInner {
mirror: StdMutex::new(entries),
broadcast_tx,
session_entry_tx,
}),
}
}
@@ -111,7 +111,7 @@ impl SegmentLogSink {
// SendError means there are zero subscribers; harmless. The
// mirror lock is held across `send` so subscribers cannot
// observe an inconsistent (snapshot, receiver) pair.
let _ = self.inner.broadcast_tx.send(entry);
let _ = self.inner.session_entry_tx.send(entry);
}
}
@@ -144,7 +144,7 @@ impl SegmentLogSink {
.expect("session log mirror mutex poisoned");
mirror.clear();
mirror.push(initial.clone());
let _ = self.inner.broadcast_tx.send(initial);
let _ = self.inner.session_entry_tx.send(initial);
}
/// Atomically swap the mirror to the supplied replacement-session prefix
@@ -161,7 +161,7 @@ impl SegmentLogSink {
.expect("session log mirror mutex poisoned");
*mirror = entries;
if let Some(initial) = first {
let _ = self.inner.broadcast_tx.send(initial);
let _ = self.inner.session_entry_tx.send(initial);
}
}
@@ -199,7 +199,7 @@ impl SegmentLogSink {
.lock()
.expect("session log mirror mutex poisoned");
let snapshot = mirror.clone();
let rx = self.inner.broadcast_tx.subscribe();
let rx = self.inner.session_entry_tx.subscribe();
(snapshot, rx)
}
+16
View File
@@ -62,6 +62,7 @@ pub(crate) fn metadata(
}
}
#[cfg(test)]
pub(crate) fn history_entry(
item: Item,
origin: WorkerHistoryProvenance,
@@ -69,6 +70,21 @@ pub(crate) fn history_entry(
HistoryEntry::new(item, metadata(origin, None))
}
pub(crate) fn history_entry_with_id(
item: Item,
entry_id: SessionHistoryEntryId,
origin: WorkerHistoryProvenance,
) -> HistoryEntry<SessionHistoryMetadata> {
HistoryEntry::new(
item,
SessionHistoryMetadata {
entry_id,
origin,
derivation: None,
},
)
}
pub(crate) fn to_logged_history_entry(
entry: &HistoryEntry<SessionHistoryMetadata>,
) -> LoggedHistoryEntry {
+2 -2
View File
@@ -414,10 +414,10 @@ impl SpawnedWorkerRegistry {
pub(crate) fn attach_parent_protocol(
&self,
event_tx: broadcast::Sender<Event>,
working_event_tx: broadcast::Sender<Event>,
parent_session_id: String,
) {
*self.parent_protocol.lock().unwrap() = Some((event_tx, parent_session_id));
*self.parent_protocol.lock().unwrap() = Some((working_event_tx, parent_session_id));
for record in self.internal_records.lock().unwrap().clone() {
self.start_protocol_forwarding(record);
}
+50 -8
View File
@@ -22,7 +22,7 @@ use manifest::{
use serde::Deserialize;
use tokio::sync::mpsc;
use workdir::{
WorkdirDelegationPermission, WorkdirDelegationRequest, WorkdirDelegationRule,
WorkdirDelegationPermission, WorkdirDelegationRequest, WorkdirDelegationRule, WorkdirPath,
WorkdirSessionHandle,
};
@@ -258,9 +258,10 @@ pub struct SubWorkerSpawnTool {
spawner_name: String,
workspace_context: crate::worker::WorkerWorkspaceContext,
parent_notifications: ParentNotificationTarget,
/// Runtime-owned root used only for bounded Internal Worker tool artifacts such as Bash spill
/// output. It is not an Internal Worker identity or catalog location.
/// Runtime-owned root used for Internal Worker controller state.
runtime_base: PathBuf,
/// Parent Worker-owned temporary root used for bounded Bash spill output.
bash_output_dir: PathBuf,
/// Inherited runtime workspace root for Profile/project/Ticket/workflow/
/// memory context. SubWorkerSpawn `cwd` must not affect this value.
workspace_root: PathBuf,
@@ -292,6 +293,7 @@ impl SubWorkerSpawnTool {
workspace_context: crate::worker::WorkerWorkspaceContext,
parent_notifications: ParentNotificationTarget,
runtime_base: PathBuf,
bash_output_dir: PathBuf,
workspace_root: PathBuf,
source_workdir_session: Option<WorkdirSessionHandle>,
registry: Arc<SpawnedWorkerRegistry>,
@@ -304,6 +306,7 @@ impl SubWorkerSpawnTool {
workspace_context,
parent_notifications,
runtime_base,
bash_output_dir,
workspace_root,
source_workdir_session,
registry,
@@ -367,7 +370,22 @@ impl Tool for SubWorkerSpawnTool {
.reserve_internal_name(input.name.clone())
.map_err(|error| ToolError::InvalidArgument(error.to_string()))?;
let workdir_rules = parse_workdir_scope(&input.scope)?;
let mut workdir_rules = parse_workdir_scope(&input.scope)?;
let child_bash_output_dir = self.bash_output_dir.join("sub-workers").join(&input.name);
tokio::fs::create_dir_all(&child_bash_output_dir)
.await
.map_err(|error| {
ToolError::ExecutionFailed(format!(
"create Internal Worker Bash output directory {}: {error}",
child_bash_output_dir.display()
))
})?;
workdir_rules.push(WorkdirDelegationRule {
target: WorkdirPath::new_scoped(child_bash_output_dir.to_string_lossy())
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?,
permission: WorkdirDelegationPermission::Read,
recursive: true,
});
let source_workdir_session =
require_active_workdir_session(self.source_workdir_session.as_ref())?;
let delegation_request = workdir_delegation_request(input.cwd.as_deref(), workdir_rules)?;
@@ -464,14 +482,22 @@ impl Tool for SubWorkerSpawnTool {
.await
.map_err(|error| ToolError::ExecutionFailed(format!("build Internal Worker: {error}")))?;
child.bind_workdir_session(Some(workdir_delegation.scoped_session.clone()));
child
.add_scope_rules([ScopeRule {
target: child_bash_output_dir.clone(),
permission: manifest::Permission::Read,
recursive: true,
}])
.map_err(|error| {
ToolError::ExecutionFailed(format!(
"grant Internal Worker Bash output scope: {error}"
))
})?;
let child_scope = child.scope().clone();
let child_registry = SpawnedWorkerRegistry::new_internal(input.name.clone(), child_scope);
register_worker_tools(
&mut child,
self.runtime_base
.join("internal-workers")
.join(&input.name)
.join("bash-output"),
child_bash_output_dir,
self.runtime_base.clone(),
child_registry.clone(),
None,
@@ -883,6 +909,7 @@ pub(crate) fn sub_worker_spawn_tool(
workspace_context: crate::worker::WorkerWorkspaceContext,
parent_notifications: ParentNotificationTarget,
runtime_base: PathBuf,
bash_output_dir: PathBuf,
workspace_root: PathBuf,
source_workdir_session: Option<WorkdirSessionHandle>,
registry: Arc<SpawnedWorkerRegistry>,
@@ -894,6 +921,7 @@ pub(crate) fn sub_worker_spawn_tool(
workspace_context,
parent_notifications,
runtime_base,
bash_output_dir,
workspace_root,
source_workdir_session,
registry,
@@ -907,6 +935,7 @@ fn sub_worker_spawn_tool_impl(
workspace_context: crate::worker::WorkerWorkspaceContext,
parent_notifications: ParentNotificationTarget,
runtime_base: PathBuf,
bash_output_dir: PathBuf,
workspace_root: PathBuf,
source_workdir_session: Option<WorkdirSessionHandle>,
registry: Arc<SpawnedWorkerRegistry>,
@@ -938,6 +967,7 @@ fn sub_worker_spawn_tool_impl(
workspace_context.clone(),
parent_notifications.clone(),
runtime_base.clone(),
bash_output_dir.clone(),
workspace_root.clone(),
source_workdir_session.clone(),
registry.clone(),
@@ -1082,12 +1112,17 @@ extract_threshold = 4000
async fn reviewer_profile_write_scope_exposes_command_tools_and_notifies_parent_controller() {
let runtime = TempDir::new().unwrap();
let workspace_root = runtime.path().join("project");
let bash_output_dir = runtime.path().join("bash-output");
let available_profiles = write_project_profile_registry(
&workspace_root,
Some("reviewer"),
&[("reviewer", "reviewer.toml", INTERNAL_REVIEWER_PROFILE)],
);
let mut manifest = parent_manifest(&workspace_root, None);
manifest
.scope
.allow
.push(abs_rule(&bash_output_dir, Permission::Read));
manifest.delegation_scope = ScopeConfig {
allow: vec![abs_rule(&workspace_root, Permission::Write)],
deny: Vec::new(),
@@ -1118,6 +1153,7 @@ extract_threshold = 4000
workspace_context,
ParentNotificationTarget::Controller(parent_method_tx.downgrade()),
runtime.path().to_path_buf(),
bash_output_dir.clone(),
workspace_root.clone(),
Some(source_workdir_session),
registry.clone(),
@@ -1167,6 +1203,12 @@ extract_threshold = 4000
.await
.expect("spawn project reviewer as Internal Worker");
assert!(output.summary.contains("internal worker `reviewer-child`"));
assert!(
bash_output_dir
.join("sub-workers")
.join("reviewer-child")
.is_dir()
);
assert!(spawner_scope.snapshot().is_writable(&workspace_root));
let record = registry
.get_internal("reviewer-child")
+358 -50
View File
@@ -15,8 +15,8 @@ use agen::{
};
use arc_swap::ArcSwap;
use session_store::{
LogEntry, PromptRenderProvenance, SegmentId, SessionExtension, SessionId, Store, StoreError,
SystemItem, segment_log,
LogEntry, PasteArtifactLimits, PromptRenderProvenance, SegmentId, SessionExtension, SessionId,
Store, StoreError, SystemItem, segment_log,
};
use session_store::{
WorkerActiveSegmentRef, WorkerMetadata, WorkerMetadataStore, WorkerReclaimedChild,
@@ -25,10 +25,12 @@ use session_store::{
use tracing::{info, warn};
use crate::segment_log_sink::SegmentLogSink;
#[cfg(test)]
use crate::session_history::history_entry;
use crate::session_history::{
SessionHistoryDerivation, SessionHistoryMetadata, WorkerHistoryProvenance, history_entry,
metadata as new_history_metadata, restore_history_entries, to_logged_history_entry,
worker_subject,
SessionHistoryDerivation, SessionHistoryEntryId, SessionHistoryMetadata,
WorkerHistoryProvenance, history_entry_with_id, metadata as new_history_metadata,
restore_history_entries, to_logged_history_entry, worker_subject,
};
use manifest::{
@@ -58,6 +60,7 @@ use crate::internal_worker::{
};
const COMPACTION_EXTENSION_DOMAIN: &str = "yoi.compaction";
const LARGE_PASTE_INLINE_MAX_BYTES: usize = 32 * 1024;
const WORKER_ORCHESTRATION_INSTRUCTION_ID: &str = "worker.orchestration";
const WORKER_ORCHESTRATION_PROMPT_REF: &str = "common.worker_orchestration";
@@ -1149,7 +1152,7 @@ pub struct Worker<C: LlmClient, St: Store> {
/// etc.). Attached by the Controller alongside `alerter`. Unlike
/// notifications, events sent here are NOT replayed to clients that
/// connect after the fact — they are fire-and-forget broadcasts.
event_tx: Option<broadcast::Sender<Event>>,
working_event_tx: Option<broadcast::Sender<Event>>,
/// Parent-owned projection/control boundary for observable Internal service Workers.
/// Service Workers are never exposed through the model-facing SubWorker control surface.
internal_worker_registry: Option<Arc<crate::spawn::registry::SpawnedWorkerRegistry>>,
@@ -1304,7 +1307,7 @@ impl<C: LlmClient + Clone + 'static, St: Store + Clone + 'static> Worker<C, St>
system_prompt_template: None,
feature_instructions: self.feature_instructions.clone(),
alerter: self.alerter.clone(),
event_tx: self.event_tx.clone(),
working_event_tx: self.working_event_tx.clone(),
internal_worker_registry: self.internal_worker_registry.clone(),
in_flight: self.in_flight.clone(),
ai_activity_counter: self.ai_activity_counter.clone(),
@@ -1497,7 +1500,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
system_prompt_template: None,
feature_instructions: Vec::new(),
alerter: None,
event_tx: None,
working_event_tx: None,
internal_worker_registry: None,
in_flight: None,
ai_activity_counter: Arc::new(AtomicUsize::new(0)),
@@ -2190,13 +2193,13 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// The Controller wires this alongside [`attach_alerter`] so that
/// Worker-internal operations (currently: compaction) can surface
/// progress to connected clients.
pub fn attach_event_tx(&mut self, event_tx: broadcast::Sender<Event>) {
pub fn attach_working_event_tx(&mut self, working_event_tx: broadcast::Sender<Event>) {
let session_id = self.session_id().to_string();
let registry = self.internal_worker_registry.get_or_insert_with(
crate::spawn::registry::SpawnedWorkerRegistry::new_for_internal_services,
);
registry.attach_parent_protocol(event_tx.clone(), session_id);
self.event_tx = Some(event_tx);
registry.attach_parent_protocol(working_event_tx.clone(), session_id);
self.working_event_tx = Some(working_event_tx);
}
pub(crate) fn attach_internal_worker_registry(
@@ -2240,10 +2243,10 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
}
/// Broadcast a typed `Event` to connected clients. No-op when no
/// `event_tx` is attached (tests / direct `Worker::new` usage) or when
/// `working_event_tx` is attached (tests / direct `Worker::new` usage) or when
/// no clients are currently subscribed.
fn send_event(&self, event: Event) {
if let Some(tx) = self.event_tx.as_ref() {
if let Some(tx) = self.working_event_tx.as_ref() {
let _ = tx.send(event);
}
}
@@ -2797,7 +2800,15 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
St: Clone + 'static,
F: FnOnce(),
{
let (input, pending_flow_state, flow_projection) = self.prepare_flow_input(input)?;
let (mut input, pending_flow_state, flow_projection) = self.prepare_flow_input(input)?;
let projected_entry_ids = if flow_projection.is_some() {
(0..input.len())
.map(|_| SessionHistoryEntryId::new())
.collect::<Vec<_>>()
} else {
vec![SessionHistoryEntryId::new()]
};
self.materialize_large_pastes(&mut input, &projected_entry_ids, flow_projection.is_some())?;
if let Some(state) = pending_flow_state.as_ref() {
let payload = serde_json::to_value(state).map_err(|error| {
WorkerError::FlowInput(format!("serialize Flow runtime state: {error}"))
@@ -2830,7 +2841,8 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
trigger: protocol::InvokeKind::UserSend,
})?;
let projected_input = self.projected_input_history(&input, flow_projection.as_ref());
let projected_input =
self.projected_input_history(&input, flow_projection.as_ref(), &projected_entry_ids);
// Persist original typed segments together with the exact ordered
// model-visible item+origin projection before any entry becomes live.
@@ -3064,17 +3076,86 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
Ok(())
}
fn materialize_large_pastes(
&self,
input: &mut [Segment],
projected_entry_ids: &[SessionHistoryEntryId],
one_entry_per_segment: bool,
) -> Result<(), WorkerError> {
let uploaded_file_count = input
.iter()
.filter(|segment| matches!(segment, Segment::UploadedFile { .. }))
.count();
if uploaded_file_count > session_store::DEFAULT_MAX_FILES_PER_SUBMISSION {
return Err(WorkerError::Store(StoreError::ArtifactQuotaExceeded));
}
for (index, segment) in input.iter_mut().enumerate() {
if let Segment::UploadedFile { file } = segment {
if file.source_entry_id.is_some()
|| file.availability != protocol::UploadedFileAvailability::Available
{
return Err(WorkerError::Store(StoreError::ArtifactIntegrityMismatch));
}
let entry_index = if one_entry_per_segment { index } else { 0 };
let source_entry_id = projected_entry_ids
.get(entry_index)
.expect("projected input id exists for every uploaded file")
.0
.clone();
*file = self
.store
.bind_uploaded_file(self.session_id(), file, &source_entry_id)
.map_err(WorkerError::Store)?;
continue;
}
if let Segment::PasteArtifact { artifact } = segment {
let (stored, _) = self
.store
.read_paste_artifact(self.session_id(), &artifact.artifact_id)?;
if &stored != artifact {
return Err(WorkerError::Store(StoreError::PasteArtifactIntegrity(
artifact.artifact_id.clone(),
)));
}
continue;
}
let Segment::Paste { content, .. } = segment else {
continue;
};
if content.len() <= LARGE_PASTE_INLINE_MAX_BYTES {
continue;
}
let entry_index = if one_entry_per_segment { index } else { 0 };
let source_entry_id = projected_entry_ids
.get(entry_index)
.expect("projected input id exists for every paste")
.0
.as_str();
let artifact = self.store.write_paste_artifact(
self.session_id(),
source_entry_id,
content,
PasteArtifactLimits::default(),
)?;
*segment = Segment::PasteArtifact { artifact };
}
Ok(())
}
fn projected_input_history(
&self,
input: &[Segment],
flow_projection: Option<&PreparedFlowProjection>,
entry_ids: &[SessionHistoryEntryId],
) -> Vec<HistoryEntry<SessionHistoryMetadata>> {
if let Some(flow) = flow_projection {
return input
.iter()
.map(|segment| match segment {
Segment::Flow { .. } => history_entry(
.zip(entry_ids)
.map(|(segment, entry_id)| match segment {
Segment::Flow { .. } => history_entry_with_id(
Item::user_message(flow.instructions.clone()),
entry_id.clone(),
WorkerHistoryProvenance::FlowInstruction {
selector: flow.selector.clone(),
definition_id: flow.definition_id.clone(),
@@ -3083,8 +3164,9 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
state_id: flow.state_id.clone(),
},
),
other => history_entry(
other => history_entry_with_id(
Item::user_message(Segment::flatten_to_text(std::slice::from_ref(other))),
entry_id.clone(),
// Current public submit transport does not carry a
// trusted account/Worker subject envelope. Fail closed
// instead of promoting role=user to HumanInput.
@@ -3094,8 +3176,12 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
.collect();
}
vec![history_entry(
vec![history_entry_with_id(
Item::user_message(Segment::flatten_to_text(input)),
entry_ids
.first()
.expect("projected Worker input always has one entry id")
.clone(),
WorkerHistoryProvenance::LegacyUnknown,
)]
}
@@ -3245,7 +3331,13 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
}),
compacted_from: None,
};
let mut initial_entries = vec![entry.clone()];
let mut initial_entries = vec![
entry.clone(),
LogEntry::InputSegmentsCheckpoint {
ts: segment_log::now_millis(),
user_segments: self.user_segments.clone(),
},
];
if let Some(checkpoint) =
active_run_checkpoint_entry(w.active_run_turn_count(), w.turn_count())
{
@@ -4209,6 +4301,13 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
})
.collect::<Vec<_>>();
let retained_user_segments = self
.user_segments
.iter()
.skip(self.user_segments.len().saturating_sub(retained_user_msgs))
.cloned()
.collect::<Vec<_>>();
// Build the SegmentStart entry for the new compacted segment.
// Inherits the source Segment's session_id so the compacted
// lineage stays grouped under the same Session. Atomically
@@ -4234,7 +4333,13 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
at_turn_index: source_turn_count,
}),
};
let mut initial_entries = vec![entry.clone()];
let mut initial_entries = vec![
entry.clone(),
LogEntry::InputSegmentsCheckpoint {
ts: segment_log::now_millis(),
user_segments: retained_user_segments.clone(),
},
];
if let Some(checkpoint) =
active_run_checkpoint_entry(w.active_run_turn_count(), source_turn_count)
{
@@ -4286,10 +4391,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
// segments; only the user_messages surviving in retained_items
// keep them. They are always the trailing K entries of
// `self.user_segments` because submissions are appended in order.
let drop_n = self.user_segments.len().saturating_sub(retained_user_msgs);
if drop_n > 0 {
self.user_segments.drain(..drop_n);
}
self.user_segments = retained_user_segments;
self.session.replace_history(compacted_history_entries);
// Compaction-introduced system messages are part of the new
@@ -4407,7 +4509,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
.with_memory_settings(&memory_cfg)
.emit(
self.workspace_client(),
self.event_tx.as_ref(),
self.working_event_tx.as_ref(),
memory::audit::WorkerLifecycleStatus::Skipped,
"extract_threshold_disabled",
None,
@@ -4438,7 +4540,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
.with_memory_settings(&memory_cfg)
.emit(
self.workspace_client(),
self.event_tx.as_ref(),
self.working_event_tx.as_ref(),
memory::audit::WorkerLifecycleStatus::Skipped,
"extract_already_in_flight",
None,
@@ -4503,7 +4605,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
Some(model_audit_from_manifest(model)),
)
.with_memory_settings(memory_cfg);
let event_tx = self.event_tx.as_ref();
let working_event_tx = self.working_event_tx.as_ref();
let pointer_snapshot = self
.extract_pointer
@@ -4519,7 +4621,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
if tokens_since < threshold {
audit.emit(
self.workspace_client(),
event_tx,
working_event_tx,
memory::audit::WorkerLifecycleStatus::Skipped,
format!(
"token_threshold_not_reached tokens_since={tokens_since} threshold={threshold}"
@@ -4536,7 +4638,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
audit
.emit(
self.workspace_client(),
event_tx,
working_event_tx,
memory::audit::WorkerLifecycleStatus::Skipped,
"no_new_history_items",
None,
@@ -4564,7 +4666,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
audit
.emit(
self.workspace_client(),
event_tx,
working_event_tx,
memory::audit::WorkerLifecycleStatus::Skipped,
"empty_segment_log",
None,
@@ -4583,7 +4685,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
audit
.emit(
self.workspace_client(),
event_tx,
working_event_tx,
memory::audit::WorkerLifecycleStatus::Skipped,
"no_new_segment_entries",
None,
@@ -4613,7 +4715,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
audit
.emit(
self.workspace_client(),
event_tx,
working_event_tx,
memory::audit::WorkerLifecycleStatus::Started,
format!(
"token_threshold_reached tokens_since={tokens_since} threshold={threshold}"
@@ -4637,7 +4739,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
audit
.emit(
self.workspace_client(),
event_tx,
working_event_tx,
memory::audit::WorkerLifecycleStatus::Failed,
format!("client_build_failed: {err}"),
None,
@@ -4659,7 +4761,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
audit
.emit(
self.workspace_client(),
event_tx,
working_event_tx,
memory::audit::WorkerLifecycleStatus::Failed,
format!("prompt_render_failed: {err}"),
None,
@@ -4737,7 +4839,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
audit
.emit(
self.workspace_client(),
event_tx,
working_event_tx,
memory::audit::WorkerLifecycleStatus::Cancelled,
"worker_cancelled: internal Worker run rolled back before AI output",
usage,
@@ -4760,7 +4862,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
audit
.emit(
self.workspace_client(),
event_tx,
working_event_tx,
lifecycle_status_for_worker_error(&err.source),
format!("worker_failed: {}", err.source),
usage,
@@ -4812,7 +4914,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
audit
.emit(
self.workspace_client(),
event_tx,
working_event_tx,
memory::audit::WorkerLifecycleStatus::Completed,
reason,
usage,
@@ -4847,7 +4949,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
.with_memory_settings(&memory_cfg)
.emit(
self.workspace_client(),
self.event_tx.as_ref(),
self.working_event_tx.as_ref(),
memory::audit::WorkerLifecycleStatus::Skipped,
"consolidation_threshold_disabled",
None,
@@ -4889,7 +4991,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
.with_memory_settings(&memory_cfg)
.emit(
self.workspace_client(),
self.event_tx.as_ref(),
self.working_event_tx.as_ref(),
memory::audit::WorkerLifecycleStatus::Skipped,
"consolidation_backend_operation_failed",
None,
@@ -4942,18 +5044,18 @@ fn model_audit_from_manifest(model: &manifest::ModelManifest) -> memory::audit::
}
fn emit_memory_worker_event(
event_tx: Option<&broadcast::Sender<Event>>,
working_event_tx: Option<&broadcast::Sender<Event>>,
run_id: uuid::Uuid,
worker: memory::audit::AuditWorker,
status: memory::audit::WorkerLifecycleStatus,
trigger: memory::audit::AuditTrigger,
reason: &str,
) {
let Some(event_tx) = event_tx else {
let Some(working_event_tx) = working_event_tx else {
return;
};
let message = format!("memory {} {}: {reason}", worker.label(), status.label());
let _ = event_tx.send(Event::MemoryWorker(protocol::MemoryWorkerEvent {
let _ = working_event_tx.send(Event::MemoryWorker(protocol::MemoryWorkerEvent {
worker: worker.label().to_string(),
status: status.label().to_string(),
run_id: run_id.to_string(),
@@ -5003,7 +5105,7 @@ impl WorkerAuditBase {
async fn emit(
&self,
workspace_client: &dyn WorkspaceClient,
event_tx: Option<&broadcast::Sender<Event>>,
working_event_tx: Option<&broadcast::Sender<Event>>,
status: memory::audit::WorkerLifecycleStatus,
reason: impl Into<String>,
usage: Option<memory::audit::UsageAudit>,
@@ -5034,7 +5136,7 @@ impl WorkerAuditBase {
.await;
if should_emit_memory_worker_event(self.worker, status, &reason) {
emit_memory_worker_event(
event_tx,
working_event_tx,
self.run_id,
self.worker,
status,
@@ -5218,7 +5320,7 @@ where
system_prompt_template: common.system_prompt_template,
feature_instructions: common.feature_instructions,
alerter: None,
event_tx: None,
working_event_tx: None,
internal_worker_registry: None,
in_flight: None,
ai_activity_counter: Arc::new(AtomicUsize::new(0)),
@@ -5302,7 +5404,7 @@ where
system_prompt_template: common.system_prompt_template,
feature_instructions: common.feature_instructions,
alerter: None,
event_tx: None,
working_event_tx: None,
internal_worker_registry: None,
in_flight: None,
ai_activity_counter: Arc::new(AtomicUsize::new(0)),
@@ -5421,7 +5523,7 @@ where
system_prompt_template: common.system_prompt_template,
feature_instructions: common.feature_instructions,
alerter: None,
event_tx: None,
working_event_tx: None,
internal_worker_registry: None,
in_flight: None,
ai_activity_counter: Arc::new(AtomicUsize::new(0)),
@@ -5797,7 +5899,7 @@ where
system_prompt_template: None,
feature_instructions: common.feature_instructions,
alerter: None,
event_tx: None,
working_event_tx: None,
internal_worker_registry: None,
in_flight: None,
ai_activity_counter: Arc::new(AtomicUsize::new(0)),
@@ -6405,6 +6507,16 @@ fn preview_segments(segments: &[Segment]) -> String {
match segment {
Segment::Text { content } => preview.push_str(content.trim()),
Segment::Paste { content, .. } => preview.push_str(content.trim()),
Segment::PasteArtifact { artifact } => {
preview.push_str("[Large paste artifact: ");
preview.push_str(&artifact.artifact_id);
preview.push(']');
}
Segment::UploadedFile { file } => {
preview.push_str("[Attached file: ");
preview.push_str(&file.file_name);
preview.push(']');
}
Segment::FileRef { path } => {
preview.push('@');
preview.push_str(path);
@@ -7807,6 +7919,7 @@ mod build_summary_prompt_tests {
fork_entries.as_slice(),
[
LogEntry::AnnotatedSegmentStart { .. },
LogEntry::InputSegmentsCheckpoint { .. },
LogEntry::ActiveRunCheckpoint {
active_turn_count: 3,
total_turn_count: 7,
@@ -7909,7 +8022,9 @@ mod build_summary_prompt_tests {
FLOW_RUNTIME_EXTENSION_DOMAIN,
serde_json::to_value(&state).unwrap(),
);
let projected = worker.projected_input_history(&segments, projection.as_ref());
let projected_ids = vec![SessionHistoryEntryId::new(), SessionHistoryEntryId::new()];
let projected =
worker.projected_input_history(&segments, projection.as_ref(), &projected_ids);
worker
.commit_entry(LogEntry::AnnotatedUserInput {
ts: segment_log::now_millis(),
@@ -7981,6 +8096,199 @@ mod build_summary_prompt_tests {
);
}
#[tokio::test]
async fn uploaded_file_is_verified_and_bound_to_projected_entry_before_commit() {
let (_dir, worker) = rewind_test_worker().await;
let reference = worker
.store
.write_uploaded_file(
worker.session_id(),
"notes.md",
"text/markdown",
b"# private body",
session_store::UploadedFileLimits::default(),
)
.unwrap();
let entry_id = SessionHistoryEntryId::new();
let mut input = vec![Segment::UploadedFile { file: reference }];
worker
.materialize_large_pastes(&mut input, std::slice::from_ref(&entry_id), false)
.unwrap();
let file = match &input[0] {
Segment::UploadedFile { file } => file,
other => panic!("expected uploaded file, got {other:?}"),
};
assert_eq!(file.source_entry_id.as_deref(), Some(entry_id.0.as_str()));
assert_eq!(
worker
.store
.read_uploaded_file(worker.session_id(), file)
.unwrap(),
b"# private body"
);
assert!(matches!(
worker
.store
.delete_uploaded_file(worker.session_id(), &file.artifact_id),
Err(StoreError::ArtifactAlreadyCommitted)
));
let projected = worker.projected_input_history(&input, None, &[entry_id]);
let text = projected[0].item.as_text().unwrap();
assert!(text.contains("notes.md"));
assert!(text.contains(&file.artifact_id));
assert!(!text.contains("private body"));
let mut forged = input.clone();
let Segment::UploadedFile { file } = &mut forged[0] else {
unreachable!();
};
file.source_entry_id = None;
file.sha256 = "0".repeat(64);
assert!(matches!(
worker.materialize_large_pastes(&mut forged, &[SessionHistoryEntryId::new()], false),
Err(WorkerError::Store(StoreError::ArtifactIntegrityMismatch))
));
}
#[tokio::test]
async fn large_paste_is_stored_before_compact_history_is_committed() {
let (_dir, worker) = rewind_test_worker().await;
let exact = "x".repeat(LARGE_PASTE_INLINE_MAX_BYTES);
let exact_ids = vec![SessionHistoryEntryId::new()];
let mut exact_input = vec![Segment::Paste {
id: 1,
chars: exact.len() as u32,
lines: 1,
content: exact.clone(),
}];
worker
.materialize_large_pastes(&mut exact_input, &exact_ids, false)
.unwrap();
assert!(matches!(&exact_input[0], Segment::Paste { content, .. } if content == &exact));
let mut empty_input = vec![Segment::Paste {
id: 0,
chars: 0,
lines: 0,
content: String::new(),
}];
worker
.materialize_large_pastes(&mut empty_input, &exact_ids, false)
.unwrap();
assert!(matches!(
&empty_input[0],
Segment::Paste { content, .. } if content.is_empty()
));
let body = format!("{}\n終端\n", "".repeat(12_000));
let entry_id = SessionHistoryEntryId::new();
let mut input = vec![Segment::Paste {
id: 2,
chars: body.chars().count() as u32,
lines: 3,
content: body.clone(),
}];
worker
.materialize_large_pastes(&mut input, std::slice::from_ref(&entry_id), false)
.unwrap();
let artifact = match &input[0] {
Segment::PasteArtifact { artifact } => artifact.clone(),
other => panic!("expected stored paste reference, got {other:?}"),
};
assert_eq!(artifact.source_entry_id, entry_id.0);
assert_eq!(artifact.byte_len, body.len() as u64);
assert_eq!(artifact.char_count, body.chars().count() as u64);
assert_eq!(
worker
.store
.read_paste_artifact(worker.session_id(), &artifact.artifact_id)
.unwrap()
.1,
body
);
worker
.materialize_large_pastes(&mut input, &[SessionHistoryEntryId::new()], false)
.unwrap();
assert!(matches!(
&input[0],
Segment::PasteArtifact { artifact: retained }
if retained.source_entry_id == artifact.source_entry_id
));
let history = worker.projected_input_history(&input, None, &[entry_id]);
assert!(!history[0].item.as_text().unwrap().contains("終端"));
append_test_entry(
&worker,
LogEntry::Invoke {
ts: segment_log::now_millis(),
trigger: protocol::InvokeKind::UserSend,
},
);
worker
.commit_entry(LogEntry::AnnotatedUserInput {
ts: segment_log::now_millis(),
segments: input.clone(),
extensions: Vec::new(),
history: history.iter().map(to_logged_history_entry).collect(),
})
.unwrap();
let location = worker.segment_state.location();
let entries = worker
.store
.read_all(location.session_id, location.segment_id)
.unwrap();
let persisted = serde_json::to_string(&entries).unwrap();
assert!(!persisted.contains("終端"));
let state = session_store::collect_state(&entries);
assert!(matches!(
&state.user_segments[0][0],
Segment::PasteArtifact { artifact: restored }
if restored.artifact_id == artifact.artifact_id
));
}
#[tokio::test]
async fn large_paste_storage_failure_commits_no_input() {
let temp = tempfile::TempDir::new().unwrap();
let store = session_store::FsStore::new(temp.path()).unwrap();
let mut worker = Worker::new(
minimal_manifest(),
Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient),
store.clone(),
WorkerWorkspaceContext::unavailable(None, "test unavailable"),
WorkerFilesystemAuthority::None,
Scope::empty(),
)
.await
.unwrap();
worker.ensure_segment_head().unwrap();
std::fs::write(
temp.path()
.join(worker.session_id().to_string())
.join("artifacts"),
"block artifact directory creation",
)
.unwrap();
let result = worker
.run(vec![Segment::Paste {
id: 1,
chars: (LARGE_PASTE_INLINE_MAX_BYTES + 1) as u32,
lines: 1,
content: "x".repeat(LARGE_PASTE_INLINE_MAX_BYTES + 1),
}])
.await;
assert!(matches!(result, Err(WorkerError::Store(StoreError::Io(_)))));
let location = worker.segment_state.location();
let entries = store
.read_all(location.session_id, location.segment_id)
.unwrap();
assert!(
!entries
.iter()
.any(|entry| matches!(entry, LogEntry::AnnotatedUserInput { .. }))
);
}
async fn rewind_test_worker() -> (
tempfile::TempDir,
Worker<NoopClient, session_store::FsStore>,
+6 -5
View File
@@ -382,7 +382,7 @@ async fn compact_emits_session_start_carrying_summary_and_task_snapshot() {
let mut worker = make_worker(client).await;
let (tx, _rx_keep) = broadcast::channel::<Event>(64);
worker.attach_event_tx(tx);
worker.attach_working_event_tx(tx);
worker.run_text("first").await.unwrap();
let session_id = worker.session_id();
@@ -429,7 +429,7 @@ async fn pre_run_compact_success_broadcasts_start_and_done() {
let mut worker = make_worker(client).await;
let (tx, mut rx) = broadcast::channel::<Event>(64);
worker.attach_event_tx(tx);
worker.attach_working_event_tx(tx);
worker.run_text("first").await.unwrap();
// Drain run events so only compact events remain in `rx`.
@@ -539,7 +539,7 @@ async fn mid_turn_compact_success_broadcasts_start_and_done() {
let mut worker = make_worker_with_manifest(MID_TURN_MANIFEST_TOML, client).await;
let (tx, mut rx) = broadcast::channel::<Event>(64);
worker.attach_event_tx(tx);
worker.attach_working_event_tx(tx);
// First run populates usage_history above the request threshold.
worker.run_text("first").await.unwrap();
@@ -718,7 +718,7 @@ async fn pre_run_compact_failure_broadcasts_start_and_failed() {
let mut worker = make_worker(client).await;
let (tx, mut rx) = broadcast::channel::<Event>(64);
worker.attach_event_tx(tx);
worker.attach_working_event_tx(tx);
worker.run_text("first").await.unwrap();
let _ = drain(&mut rx);
@@ -861,7 +861,8 @@ async fn controller_compact_method_emits_start_and_done() {
]);
let worker = make_worker_with_manifest(POST_RUN_MANIFEST_TOML, client).await;
let runtime_tmp = tempfile::tempdir().unwrap();
let (handle, _shutdown) = WorkerController::spawn(worker, runtime_tmp.path())
let bash_output_dir = runtime_tmp.path().join("bash-output");
let (handle, _shutdown) = WorkerController::spawn(worker, runtime_tmp.path(), &bash_output_dir)
.await
.unwrap();
let mut rx = handle.subscribe();
+95 -9
View File
@@ -9,6 +9,7 @@ use agen::llm_client::{ClientError, LlmClient, Request};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use async_trait::async_trait;
use futures::{Stream, StreamExt};
use manifest::{ProfileRegistrySource, ProfileResolveOptions, ProfileResolver, ProfileSelector};
use session_store::{CombinedStore, FsWorkerStore};
use session_store::{FsStore, LogEntry};
use workdir::{
@@ -248,6 +249,14 @@ async fn make_worker_with_pwd_manifest_and_workspace_context(
workspace_context: WorkerWorkspaceContext,
) -> (Worker<MockClient, TestStore>, std::path::PathBuf) {
let manifest = WorkerManifest::from_toml(manifest_toml).unwrap();
make_worker_with_manifest_and_workspace_context(client, manifest, workspace_context).await
}
async fn make_worker_with_manifest_and_workspace_context(
client: MockClient,
manifest: WorkerManifest,
workspace_context: WorkerWorkspaceContext,
) -> (Worker<MockClient, TestStore>, std::path::PathBuf) {
let store_tmp = tempfile::tempdir().unwrap();
let store = CombinedStore::new(
FsStore::new(store_tmp.path()).unwrap(),
@@ -276,12 +285,38 @@ async fn spawn_controller(worker: Worker<MockClient, TestStore>) -> WorkerHandle
let tmp = tempfile::tempdir().unwrap();
let runtime_base = tmp.path().to_owned();
std::mem::forget(tmp);
let (handle, _shutdown_rx) = WorkerController::spawn(worker, &runtime_base)
let bash_output_dir = runtime_base.join("bash-output");
let (handle, _shutdown_rx) = WorkerController::spawn(worker, &runtime_base, &bash_output_dir)
.await
.unwrap();
handle
}
#[tokio::test]
async fn controller_grants_read_scope_for_exact_bash_output_directory() {
let worker = make_worker(MockClient::new(simple_text_events())).await;
let shared_scope = worker.scope().clone();
let runtime_base = tempfile::tempdir().unwrap();
let worker_tmp = tempfile::tempdir().unwrap();
let bash_output_dir = worker_tmp.path().join("worker-1").join("bash-output");
let (handle, shutdown_rx) =
WorkerController::spawn(worker, runtime_base.path(), &bash_output_dir)
.await
.unwrap();
assert!(bash_output_dir.is_dir());
assert!(shared_scope.snapshot().allow_rules().iter().any(|rule| {
rule.target == bash_output_dir
&& rule.permission == manifest::Permission::Read
&& rule.recursive
}));
assert!(!handle.runtime_dir.path().join("bash-output").exists());
handle.send(Method::Shutdown).await.unwrap();
shutdown_rx.await.unwrap();
}
#[tokio::test]
async fn shutdown_closes_bound_workdir_session() {
let (mut worker, pwd) = make_worker_with_pwd(MockClient::new(simple_text_events())).await;
@@ -297,6 +332,7 @@ async fn shutdown_closes_bound_workdir_session() {
command: "sleep 30".to_owned(),
timeout_secs: 60,
output_limit: 1024,
spill_dir: None,
tool_call_id: None,
})
.await
@@ -304,9 +340,11 @@ async fn shutdown_closes_bound_workdir_session() {
worker.bind_workdir_session(Some(Arc::clone(&session)));
let runtime_base = tempfile::tempdir().unwrap();
let (handle, shutdown_rx) = WorkerController::spawn(worker, runtime_base.path())
.await
.unwrap();
let bash_output_dir = runtime_base.path().join("bash-output");
let (handle, shutdown_rx) =
WorkerController::spawn(worker, runtime_base.path(), &bash_output_dir)
.await
.unwrap();
handle.send(Method::Shutdown).await.unwrap();
tokio::time::timeout(std::time::Duration::from_secs(5), shutdown_rx)
.await
@@ -338,6 +376,7 @@ async fn controller_projects_workdir_command_events_and_snapshot_state() {
command: "printf ready; sleep 0.3; printf done".to_owned(),
timeout_secs: 5,
output_limit: 1024,
spill_dir: None,
tool_call_id: Some("tool-command-1".into()),
})
.await
@@ -445,6 +484,7 @@ async fn controller_refreshes_command_snapshot_after_high_output_provider_lag()
.to_owned(),
timeout_secs: 10,
output_limit: 1024,
spill_dir: None,
tool_call_id: Some("tool-high-output".into()),
})
.await
@@ -508,8 +548,9 @@ async fn controller_startup_failure_closes_bound_workdir_session() {
let invalid_runtime_base = runtime_base.path().join("not-a-directory");
std::fs::write(&invalid_runtime_base, "file").unwrap();
let bash_output_dir = runtime_base.path().join("bash-output");
assert!(
WorkerController::spawn(worker, &invalid_runtime_base)
WorkerController::spawn(worker, &invalid_runtime_base, &bash_output_dir)
.await
.is_err()
);
@@ -519,6 +560,7 @@ async fn controller_startup_failure_closes_bound_workdir_session() {
command: "printf unreachable".to_owned(),
timeout_secs: 5,
output_limit: 1024,
spill_dir: None,
tool_call_id: None,
})
.await,
@@ -580,7 +622,19 @@ async fn feature_flags_default_to_core_tool_surface_only() {
let request = wait_for_captured_request(&client_for_assert).await;
let names = request_tool_names(&request);
assert_eq!(names, vec!["Bash", "Edit", "Glob", "Grep", "Read", "Write"]);
assert_eq!(
names,
vec![
"Bash",
"Edit",
"Glob",
"Grep",
"Read",
"ReadInputArtifact",
"SearchInputArtifact",
"Write",
]
);
assert!(!names.iter().any(|name| name == "TaskCreate"));
assert!(!names.iter().any(|name| name == "WebSearch"));
assert!(!names.iter().any(|name| name == "SubWorkerSpawn"));
@@ -738,6 +792,37 @@ permission = "write"
}
}
#[tokio::test]
async fn builtin_orchestrator_exposes_worker_remove_and_workdir_delete() {
let workspace = tempfile::tempdir().unwrap();
let resolved = ProfileResolver::new()
.with_workspace_base(workspace.path())
.resolve(
&ProfileSelector::source_named(ProfileRegistrySource::Builtin, "orchestrator"),
ProfileResolveOptions::with_worker_name("orchestrator-worker"),
)
.unwrap();
let workspace_context =
WorkerWorkspaceContext::with_client(None, Arc::new(NoopWorkspaceClient));
let client = MockClient::new(simple_text_events());
let client_for_assert = client.clone();
let (worker, _pwd) = make_worker_with_manifest_and_workspace_context(
client,
resolved.manifest,
workspace_context,
)
.await;
let handle = spawn_controller(worker).await;
handle.send(Method::run_text("Hello")).await.unwrap();
wait_for_status(&handle, WorkerStatus::Idle).await;
let request = wait_for_captured_request(&client_for_assert).await;
let installed = request_tool_names(&request);
assert!(installed.iter().any(|name| name == "WorkerRemove"));
assert!(installed.iter().any(|name| name == "WorkdirDelete"));
}
#[tokio::test]
async fn worker_and_sub_worker_features_install_one_canonical_control_surface() {
let manifest = r#"
@@ -863,7 +948,8 @@ permission = "write"
let client = MockClient::new(simple_text_events());
let worker = make_worker_with_pwd_and_manifest(client, manifest).await.0;
let tmp = tempfile::tempdir().unwrap();
let result = WorkerController::spawn(worker, tmp.path()).await;
let bash_output_dir = tmp.path().join("bash-output");
let result = WorkerController::spawn(worker, tmp.path(), &bash_output_dir).await;
assert!(
result.is_ok(),
"feature exposure must not imply delegation authority"
@@ -1515,7 +1601,7 @@ async fn notify_while_idle_auto_starts_turn_and_injects_system_message() {
// after that.
wait_for_status(&handle, WorkerStatus::Idle).await;
// The live echo arrives via the sink's `Event::SystemItem` lane,
// not on the `event_tx` broadcast that `handle.subscribe()` taps.
// not on the `working_event_tx` broadcast that `handle.subscribe()` taps.
// Verify the notification landed on the sink mirror instead.
let (entries, _) = handle.sink.subscribe_with_snapshot();
let saw_notify_in_mirror = entries.iter().any(|e| {
@@ -1899,7 +1985,7 @@ async fn socket_worker_event_turn_ended_while_idle_auto_starts_turn() {
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
// The SystemItem and TurnEnd events arrive through independent
// broadcast lanes (sink fan-out vs `event_tx`), so their relative
// broadcast lanes (sink fan-out vs `working_event_tx`), so their relative
// order on the wire is non-deterministic. Keep reading until both
// are observed (or the deadline trips), rather than breaking on
// the first TurnEnd.
+16 -1
View File
@@ -12,7 +12,22 @@ typescript = ["dep:ts-rs"]
[dependencies]
serde = { workspace = true, features = ["derive"] }
ts-rs = { version = "12.0.1", optional = true }
workdir.workspace = true
[[example]]
name = "generate_typescript"
required-features = ["typescript"]
[dev-dependencies]
serde_json.workspace = true
[[example]]
name = "generate_workdir_api_types"
required-features = ["typescript"]
[[example]]
name = "generate_companion_api_types"
required-features = ["typescript"]
[[example]]
name = "generate_repository_access_types"
required-features = ["typescript"]
@@ -0,0 +1,3 @@
fn main() {
print!("{}", workspace_api::companion_api_typescript());
}
@@ -0,0 +1,3 @@
fn main() {
print!("{}", workspace_api::repository_access_api_typescript());
}
@@ -0,0 +1,3 @@
fn main() {
print!("{}", workspace_api::catalog_typescript());
}
@@ -0,0 +1,3 @@
fn main() {
print!("{}", workspace_api::workdir_api_typescript());
}
File diff suppressed because it is too large Load Diff

Some files were not shown because too many files have changed in this diff Show More