feat: merge develop into hare/develop

This commit is contained in:
2026-09-21 04:30:02 +09:00
77 changed files with 8421 additions and 2482 deletions
+3
View File
@@ -2275,6 +2275,9 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
info!("Paused during response stream");
}
self.timeline.abort_current_block();
// Preserve any UsageEvent already received for billing, but
// do not treat this paused request as a normal completion.
self.timeline.flush_usage();
return Err(EngineError::PauseRequested);
}
cancel = self.cancel_rx.recv() => {
+56 -4
View File
@@ -7,6 +7,7 @@ mod common;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use agen::Item;
use agen::interceptor::{
@@ -18,16 +19,67 @@ use agen::interceptor::{
};
use agen::llm_client::{
ClientError, LlmClient, Request, ResponseStream,
event::{Event, ResponseStatus, StatusEvent},
event::{Event, ResponseStatus, StatusEvent, UsageEvent},
};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use agen::{Engine, EngineError, EngineRunExit, History, RunInterruptionReason};
use async_trait::async_trait;
use common::MockLlmClient;
use futures::StreamExt;
// =============================================================================
// Mutable State Tests
// =============================================================================
#[derive(Clone)]
struct UsageThenPendingClient;
#[async_trait]
impl LlmClient for UsageThenPendingClient {
fn clone_boxed(&self) -> Box<dyn LlmClient> {
Box::new(self.clone())
}
async fn stream(&self, _request: Request) -> Result<ResponseStream, ClientError> {
let usage = Event::Usage(UsageEvent {
input_tokens: Some(90),
output_tokens: Some(2),
total_tokens: Some(92),
cache_read_input_tokens: Some(0),
cache_creation_input_tokens: Some(0),
});
Ok(Box::pin(
futures::stream::once(async move { Ok(usage) })
.chain(futures::stream::pending::<Result<Event, ClientError>>()),
))
}
}
#[tokio::test]
async fn pause_during_response_stream_flushes_observed_usage() {
let observed = Arc::new(Mutex::new(Vec::new()));
let observed_for_callback = Arc::clone(&observed);
let mut engine = Engine::new(UsageThenPendingClient);
engine.on_usage(move |usage| {
observed_for_callback.lock().unwrap().push(usage.clone());
});
let pause = engine.pause_sender();
let mut history = History::new();
let mut engine = engine.lock(&history);
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(20)).await;
pause.send(()).await.unwrap();
});
let exit = tokio::time::timeout(
Duration::from_secs(2),
engine.run(&mut history, "pause after usage"),
)
.await
.expect("run should observe pause");
assert!(matches!(exit, EngineRunExit::Paused));
let observed = observed.lock().unwrap();
assert_eq!(observed.len(), 1);
assert_eq!(observed[0].input_tokens, Some(90));
assert_eq!(observed[0].output_tokens, Some(2));
}
/// Verify that system prompt can be set in Mutable state
#[test]
+21 -15
View File
@@ -943,7 +943,7 @@ mod tests {
profile: Some("builtin:coder".to_string()),
ticket_assignment: None,
initial_submit: Vec::new(),
working_directory: None,
workdir_attachments: Vec::new(),
control_operation_id: None,
};
@@ -964,7 +964,7 @@ mod tests {
assert_eq!(body["display_name"], "Coder one");
assert_eq!(body["profile"], "builtin:coder");
assert_eq!(body["initial_submit"], serde_json::json!([]));
assert_eq!(body["working_directory"], serde_json::Value::Null);
assert_eq!(body["workdir_attachments"], serde_json::json!([]));
}
#[test]
@@ -1014,29 +1014,35 @@ mod tests {
"state": "idle",
"implementation": {"kind": "worker", "display_hint": "Coder"},
"capabilities": {"can_stop": true, "can_spawn_followup": false},
"working_directory": {
"working_directory_id": "wd-1",
"repository_key": "main",
"materializer_kind": "runtime_git_clone",
"status": "active",
"occupied_by": {
"runtime_id": "arcadia",
"worker_id": "worker-opaque-64",
"display_name": "Coder",
"linked_at": "2026-08-12T00:00:00Z"
"workdir_attachments": [{
"alias": "checkout",
"working_directory": {
"working_directory_id": "wd-1",
"repository_key": "main",
"materializer_kind": "runtime_git_clone",
"status": "active",
"occupied_by": {
"runtime_id": "arcadia",
"worker_id": "worker-opaque-64",
"display_name": "Coder",
"linked_at": "2026-08-12T00:00:00Z"
}
}
}
}]
});
let worker: BackendWorkerSummary = serde_json::from_value(payload.clone()).unwrap();
let workdir = worker.working_directory.unwrap();
let attachment = worker.workdir_attachments.into_iter().next().unwrap();
assert_eq!(attachment.alias, "checkout");
let workdir = attachment.working_directory;
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);
stale["workdir_attachments"][0]["working_directory"]["occupied_by"]["runtime_worker_id"] =
serde_json::json!(64);
assert!(serde_json::from_value::<BackendWorkerSummary>(stale).is_err());
}
+1 -1
View File
@@ -312,7 +312,7 @@ impl BackendWorkspaceProductClient {
initial_submit: vec![protocol::Segment::Text {
content: format!("Please handle intake for Ticket {ticket_id}."),
}],
working_directory: None,
workdir_attachments: Vec::new(),
control_operation_id: None,
};
let response: BrowserCreateWorkerResponse =
+83 -42
View File
@@ -572,6 +572,41 @@ pub enum SubscriptionWorkerAvailability {
Unavailable,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
pub struct SubscriptionWorkerWorkdirAttachment {
/// Stable Worker-local routing alias.
pub alias: 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 working_directory_id: SubscriptionWorkdirId,
}
impl SubscriptionWorkerWorkdirAttachment {
pub fn validate(&self) -> Result<(), SubscriptionValidationError> {
validate_identifier(
"workdir_attachment_alias",
&self.alias,
MAX_RESOURCE_ID_BYTES,
)?;
self.working_directory_id.validate()?;
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),
_ => Err(SubscriptionValidationError::InvalidIdentifier {
field: "repository_authority",
}),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
pub struct SubscriptionWorker {
@@ -604,15 +639,8 @@ pub struct SubscriptionWorker {
pub display_name: Option<String>,
#[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>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub workdir_attachments: Vec<SubscriptionWorkerWorkdirAttachment>,
}
impl SubscriptionWorker {
@@ -624,19 +652,20 @@ impl SubscriptionWorker {
if let Some(resource_key) = &self.resource_key {
validate_identifier("resource_key", resource_key, MAX_RESOURCE_ID_BYTES)?;
}
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()?;
let mut aliases = HashSet::new();
let mut workdir_ids = HashSet::new();
for attachment in &self.workdir_attachments {
attachment.validate()?;
if !aliases.insert(attachment.alias.as_str()) {
return Err(SubscriptionValidationError::InvalidIdentifier {
field: "workdir_attachment_alias",
});
}
if !workdir_ids.insert(attachment.working_directory_id.as_str()) {
return Err(SubscriptionValidationError::InvalidIdentifier {
field: "working_directory_id",
});
}
}
Ok(())
}
@@ -654,8 +683,6 @@ pub struct SubscriptionWorkdir {
#[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>,
}
impl SubscriptionWorkdir {
@@ -673,9 +700,6 @@ impl SubscriptionWorkdir {
}
}
validate_identifier("workdir_state", &self.state, MAX_RESOURCE_ID_BYTES)?;
if let Some(worker_id) = &self.primary_worker_id {
worker_id.validate()?;
}
Ok(())
}
}
@@ -688,8 +712,6 @@ 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 {
@@ -697,9 +719,6 @@ impl WorkspaceSubscriptionWorkdir {
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()?;
}
Ok(())
}
}
@@ -903,36 +922,58 @@ mod tests {
workspace_id: Some("workspace-1".to_string()),
display_name: Some(format!("Worker {value}")),
profile: Some("builtin:coder".to_string()),
repository_id: None,
repository_key: None,
working_directory_id: None,
workdir_attachments: Vec::new(),
}
}
#[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
.workdir_attachments
.push(SubscriptionWorkerWorkdirAttachment {
alias: "checkout".to_string(),
repository_id: Some("01890f47-3c22-7cc0-98c4-dc0c0c07398f".to_string()),
repository_key: None,
working_directory_id: SubscriptionWorkdirId::new("workdir-1").unwrap(),
});
runtime_worker.validate().unwrap();
let runtime_json = serde_json::to_value(&runtime_worker).unwrap();
assert_eq!(
runtime_json["repository_id"],
runtime_json["workdir_attachments"][0]["repository_id"],
"01890f47-3c22-7cc0-98c4-dc0c0c07398f"
);
assert!(runtime_json.get("repository_key").is_none());
assert!(
runtime_json["workdir_attachments"][0]
.get("repository_key")
.is_none()
);
let mut workspace_worker = worker("worker-1");
workspace_worker.repository_key = Some("main".to_string());
workspace_worker
.workdir_attachments
.push(SubscriptionWorkerWorkdirAttachment {
alias: "checkout".to_string(),
repository_id: None,
repository_key: Some("main".to_string()),
working_directory_id: SubscriptionWorkdirId::new("workdir-1").unwrap(),
});
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());
assert_eq!(
workspace_json["workdir_attachments"][0]["repository_key"],
"main"
);
assert!(
workspace_json["workdir_attachments"][0]
.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();
+3 -1
View File
@@ -22,7 +22,8 @@ use crate::{
SubscriptionRequestId, SubscriptionResponse, SubscriptionSnapshot,
SubscriptionTerminationCode, SubscriptionWorkdirId, SubscriptionWorker,
SubscriptionWorkerAvailability, SubscriptionWorkerId, SubscriptionWorkerIds,
SubscriptionWorkerProtocolMethod, SubscriptionWorkerState, WorkspaceSubscriptionWorkdir,
SubscriptionWorkerProtocolMethod, SubscriptionWorkerState,
SubscriptionWorkerWorkdirAttachment, WorkspaceSubscriptionWorkdir,
},
};
@@ -114,6 +115,7 @@ pub fn generated_protocol_types() -> String {
push_decl::<SubscriptionWorkerState>(&cfg, &mut output);
push_decl::<SubscriptionWorkerAvailability>(&cfg, &mut output);
push_decl::<EventSubscriptionSelector>(&cfg, &mut output);
push_decl::<SubscriptionWorkerWorkdirAttachment>(&cfg, &mut output);
push_decl::<SubscriptionWorker>(&cfg, &mut output);
push_decl::<WorkspaceSubscriptionWorkdir>(&cfg, &mut output);
push_decl::<SubscriptionSnapshot>(&cfg, &mut output);
+47 -8
View File
@@ -335,6 +335,8 @@ pub struct RepositoryMaterializationContext {
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct WorkingDirectoryRequest {
pub repository: WorkingDirectoryRepository,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub display_name: Option<String>,
#[serde(default)]
pub materializer: MaterializerKind,
#[serde(default, skip_serializing_if = "Option::is_none")]
@@ -344,12 +346,31 @@ pub struct WorkingDirectoryRequest {
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct WorkingDirectoryClaim {
pub struct WorkingDirectoryAttachmentClaim {
/// Stable Worker-local routing key. This is not a Workdir id or display name.
pub alias: workdir::WorkdirAttachmentAlias,
pub working_directory_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub relative_cwd: Option<String>,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct WorkerWorkdirAttachmentsRequest {
pub workdir_attachments: Vec<WorkingDirectoryAttachmentClaim>,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct WorkingDirectoryAttachmentRequest {
pub alias: workdir::WorkdirAttachmentAlias,
pub working_directory: WorkingDirectoryRequest,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct WorkingDirectoryAttachmentStatus {
pub alias: workdir::WorkdirAttachmentAlias,
pub working_directory: WorkingDirectoryStatus,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct WorkingDirectoryStatus {
pub summary: RuntimeWorkingDirectorySummary,
@@ -387,10 +408,10 @@ pub struct CreateWorkerRequest {
pub config_bundle: Option<ConfigBundleRef>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub initial_input: Option<WorkerInput>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub working_directory_request: Option<WorkingDirectoryRequest>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub working_directory: Option<WorkingDirectoryClaim>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub workdir_attachment_requests: Vec<WorkingDirectoryAttachmentRequest>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub workdir_attachments: Vec<WorkingDirectoryAttachmentClaim>,
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
pub worker_observation_enabled: bool,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
@@ -413,8 +434,8 @@ pub struct WorkerSummary {
pub worker_state: Option<WorkerStateSnapshot>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub workspace_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub working_directory: Option<WorkingDirectoryStatus>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub workdir_attachments: Vec<WorkingDirectoryAttachmentStatus>,
pub profile: ProfileSelector,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub display_name: Option<String>,
@@ -726,6 +747,17 @@ pub trait RuntimeApi {
#[body] request: WorkerWorkspaceApiRequest,
) -> Result<WorkerResponse, RuntimeApiError>;
#[post(
"/v1/workers/{worker_id}/workdir-attachments",
status = 200,
error_status = 400
)]
async fn replace_worker_workdir_attachments(
&self,
#[path] worker_id: String,
#[body] request: WorkerWorkdirAttachmentsRequest,
) -> Result<WorkerResponse, RuntimeApiError>;
#[post(
"/v1/workers/{worker_id}/completions",
status = 200,
@@ -1002,6 +1034,13 @@ mod tests {
) -> Result<WorkerResponse, RuntimeApiError> {
Err(test_error(501))
}
async fn replace_worker_workdir_attachments(
&self,
_worker_id: String,
_request: WorkerWorkdirAttachmentsRequest,
) -> Result<WorkerResponse, RuntimeApiError> {
Err(test_error(501))
}
async fn complete_worker_arguments(
&self,
_worker_id: String,
@@ -1072,7 +1111,7 @@ mod tests {
#[test]
fn contract_inventory_is_complete_and_unique() {
let operations = RuntimeApiMetadata::OPERATIONS;
assert_eq!(operations.len(), 15);
assert_eq!(operations.len(), 16);
let mut routes = operations
.iter()
.map(|operation| (format!("{:?}", operation.method), operation.path))
+41 -17
View File
@@ -1583,6 +1583,9 @@ pub struct RuntimeWorkingDirectoryCleanupTarget {
#[serde(deny_unknown_fields)]
pub struct RuntimeWorkingDirectorySummary {
pub working_directory_id: String,
/// Optional human-facing label, never a routing key.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub display_name: Option<String>,
pub repository_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub creation_selector: Option<String>,
@@ -1606,8 +1609,6 @@ pub struct RuntimeWorkingDirectorySummary {
#[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>,
}
@@ -1618,6 +1619,9 @@ pub struct RuntimeWorkingDirectorySummary {
#[serde(deny_unknown_fields)]
pub struct WorkingDirectorySummary {
pub working_directory_id: String,
/// Optional human-facing label, never a routing key.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub display_name: Option<String>,
pub repository_key: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub creation_selector: Option<String>,
@@ -1641,8 +1645,6 @@ pub struct WorkingDirectorySummary {
#[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>,
}
@@ -1713,6 +1715,7 @@ impl From<workdir::workspace::RuntimeWorkingDirectorySummary> for RuntimeWorking
fn from(value: workdir::workspace::RuntimeWorkingDirectorySummary) -> Self {
Self {
working_directory_id: value.working_directory_id,
display_name: value.display_name,
repository_id: value.repository_id,
creation_selector: value.creation_selector,
creation_ref: value.creation_ref,
@@ -1725,7 +1728,6 @@ impl From<workdir::workspace::RuntimeWorkingDirectorySummary> for RuntimeWorking
cleanup_target: value.cleanup_target.map(Into::into),
status: value.status.into(),
cleanliness: value.cleanliness,
primary_worker_id: value.primary_worker_id,
occupied_by: value.occupied_by.map(Into::into),
}
}
@@ -1735,6 +1737,7 @@ impl From<workdir::workspace::WorkingDirectorySummary> for WorkingDirectorySumma
fn from(value: workdir::workspace::WorkingDirectorySummary) -> Self {
Self {
working_directory_id: value.working_directory_id,
display_name: value.display_name,
repository_key: value.repository_key,
creation_selector: value.creation_selector,
creation_ref: value.creation_ref,
@@ -1747,7 +1750,6 @@ impl From<workdir::workspace::WorkingDirectorySummary> for WorkingDirectorySumma
cleanup_target: value.cleanup_target.map(Into::into),
status: value.status.into(),
cleanliness: value.cleanliness,
primary_worker_id: value.primary_worker_id,
occupied_by: value.occupied_by.map(Into::into),
}
}
@@ -1765,6 +1767,9 @@ impl From<workdir::workspace::WorkingDirectorySummary> for WorkingDirectorySumma
pub struct WorkingDirectoryCreateRequest {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub runtime_id: Option<String>,
/// Optional human-facing label, independent of every Worker attachment alias.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub display_name: Option<String>,
pub repository_key: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub selector: Option<String>,
@@ -2445,6 +2450,24 @@ pub struct WorkspaceWorkerDiscoveryPage {
pub next_cursor: Option<String>,
}
/// One Workspace Worker Workdir attachment. `alias` is the stable Worker-local
/// routing key; the nested Workdir id and display name are inventory metadata.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(deny_unknown_fields)]
pub struct WorkerWorkdirAttachmentSummary {
pub alias: String,
pub working_directory: WorkingDirectorySummary,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(deny_unknown_fields)]
pub struct RuntimeWorkerWorkdirAttachmentSummary {
pub alias: String,
pub working_directory: RuntimeWorkingDirectorySummary,
}
/// Workspace-authoritative Worker projection.
///
/// `resource_key` is required here even though Runtime-internal Worker summaries
@@ -2478,8 +2501,8 @@ pub struct WorkerSummary {
pub retention_state: String,
pub implementation: WorkerImplementationSummary,
pub capabilities: WorkerCapabilitySummary,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub working_directory: Option<WorkingDirectorySummary>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub workdir_attachments: Vec<WorkerWorkdirAttachmentSummary>,
#[serde(default)]
pub diagnostics: Vec<Diagnostic>,
}
@@ -2507,9 +2530,8 @@ pub struct WorkerLaunchWorkerSummary {
pub retention_state: String,
pub implementation: WorkerImplementationSummary,
pub capabilities: WorkerCapabilitySummary,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[cfg_attr(feature = "typescript", ts(optional = nullable))]
pub working_directory: Option<RuntimeWorkingDirectorySummary>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub workdir_attachments: Vec<RuntimeWorkerWorkdirAttachmentSummary>,
#[serde(default)]
pub diagnostics: Vec<Diagnostic>,
}
@@ -2563,6 +2585,8 @@ pub struct WorkingDirectoryRepositoryOption {
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(deny_unknown_fields)]
pub struct BrowserWorkerWorkingDirectorySelection {
/// Stable Worker-local routing alias.
pub alias: String,
pub working_directory_id: String,
#[serde(default)]
pub relative_cwd: Option<String>,
@@ -2589,7 +2613,7 @@ pub struct CreateWorkspaceWorkerRequest {
#[serde(default)]
pub initial_submit: Vec<protocol::Segment>,
#[serde(default)]
pub working_directory: Option<BrowserWorkerWorkingDirectorySelection>,
pub workdir_attachments: Vec<BrowserWorkerWorkingDirectorySelection>,
/// Backend idempotency key used only for authenticated Worker-owned spawn/control.
#[serde(default)]
pub control_operation_id: Option<String>,
@@ -3663,6 +3687,7 @@ pub fn worker_launch_api_typescript() -> String {
WorkerWorkspaceSummary::decl(&config),
WorkerImplementationSummary::decl(&config),
WorkerCapabilitySummary::decl(&config),
RuntimeWorkerWorkdirAttachmentSummary::decl(&config),
WorkerLaunchWorkerSummary::decl(&config),
WorkerLaunchRuntimeOption::decl(&config),
WorkerLaunchProfileCandidate::decl(&config),
@@ -4100,7 +4125,7 @@ mod tests {
can_stop: true,
can_spawn_followup: false,
},
working_directory: None,
workdir_attachments: Vec::new(),
diagnostics: Vec::new(),
}
}
@@ -4178,7 +4203,7 @@ mod tests {
profile: None,
ticket_assignment: None,
initial_submit: Vec::new(),
working_directory: None,
workdir_attachments: Vec::new(),
control_operation_id: None,
})
.unwrap();
@@ -4190,7 +4215,7 @@ mod tests {
"profile": null,
"ticket_assignment": null,
"initial_submit": [],
"working_directory": null,
"workdir_attachments": [],
"control_operation_id": null,
})
);
@@ -5097,6 +5122,7 @@ mod tests {
fn workdir_summary_omits_absent_optional_fields_on_the_wire() {
let value = serde_json::to_value(WorkingDirectorySummary {
working_directory_id: "workdir-1".into(),
display_name: None,
repository_key: "main".into(),
creation_selector: None,
creation_ref: None,
@@ -5109,7 +5135,6 @@ mod tests {
cleanup_target: None,
status: WorkingDirectoryStatusKind::Active,
cleanliness: None,
primary_worker_id: None,
occupied_by: None,
})
.expect("serialize Workdir summary");
@@ -5125,7 +5150,6 @@ mod tests {
"observed_at_epoch_seconds",
"cleanup_target",
"cleanliness",
"primary_worker_id",
"occupied_by",
] {
assert!(
+38 -21
View File
@@ -6,7 +6,9 @@ use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use async_trait::async_trait;
use schemars::JsonSchema;
use serde::Deserialize;
use workdir::{CommandHandle, CommandOutputRequest, CommandRequest, WorkdirSessionHandle};
use workdir::{
CommandHandle, CommandOutputRequest, CommandRequest, WorkdirSessionHandle, WorkdirSessionRouter,
};
const DEFAULT_TIMEOUT_SECS: u64 = 120;
const MAX_TIMEOUT_SECS: u64 = 600;
@@ -14,13 +16,16 @@ const INLINE_BYTE_BUDGET: usize = 12 * 1024;
#[derive(Debug, Deserialize, JsonSchema)]
struct BashParams {
/// Worker-local alias of the Workdir attachment to use.
#[serde(default)]
target_workdir: Option<String>,
command: String,
#[serde(default)]
timeout: Option<u64>,
}
pub(crate) struct BashTool {
session: WorkdirSessionHandle,
router: Arc<WorkdirSessionRouter>,
output_dir: PathBuf,
state: Arc<Mutex<BashExecutionState>>,
}
@@ -29,6 +34,7 @@ pub(crate) struct BashTool {
struct ActiveCommand {
call_id: String,
execution_nonce: u64,
session: WorkdirSessionHandle,
handle: CommandHandle,
}
@@ -93,6 +99,12 @@ impl Tool for BashTool {
) -> Result<ToolOutput, ToolError> {
let params: BashParams = serde_json::from_str(input_json)
.map_err(|error| ToolError::InvalidArgument(format!("invalid Bash input: {error}")))?;
let selected = crate::routing::resolve_session(
&self.router,
params.target_workdir.as_deref(),
workdir::WorkdirSessionCapability::Command,
)?;
let session = selected.session;
let timeout_secs = params
.timeout
.unwrap_or(DEFAULT_TIMEOUT_SECS)
@@ -106,14 +118,13 @@ impl Tool for BashTool {
state.next_execution_nonce
};
let mut guard = CommandGuard {
session: self.session.clone(),
session: session.clone(),
state: self.state.clone(),
execution_id: execution_id.clone(),
execution_nonce,
handle: None,
};
let handle = self
.session
let handle = session
.start_command(CommandRequest {
command: params.command,
timeout_secs,
@@ -131,6 +142,7 @@ impl Tool for BashTool {
ActiveCommand {
call_id: call_id.clone(),
execution_nonce,
session: session.clone(),
handle: handle.clone(),
},
);
@@ -139,13 +151,12 @@ impl Tool for BashTool {
};
guard.handle = Some(handle.clone());
if cancel_after_start {
self.session
session
.cancel_command(handle.clone())
.await
.map_err(crate::ToolsError::from)?;
}
let output = self
.session
let output = session
.command_output(CommandOutputRequest {
handle,
cursor: 0,
@@ -222,12 +233,13 @@ impl Tool for BashTool {
.active
.values()
.filter(|active| active.call_id == call_id)
.map(|active| active.handle.clone())
.cloned()
.collect::<Vec<_>>()
};
for handle in handles {
self.session
.cancel_command(handle)
for active in handles {
active
.session
.cancel_command(active.handle)
.await
.map_err(crate::ToolsError::from)?;
}
@@ -242,14 +254,12 @@ impl Tool for BashTool {
let handle = {
let mut state = self.state.lock().unwrap();
state.cancellation_requested.insert(execution_id.clone());
state
.active
.get(&execution_id)
.map(|active| active.handle.clone())
state.active.get(&execution_id).cloned()
};
if let Some(handle) = handle {
self.session
.cancel_command(handle)
if let Some(active) = handle {
active
.session
.cancel_command(active.handle)
.await
.map_err(crate::ToolsError::from)?;
}
@@ -268,13 +278,20 @@ fn truncate_for_summary(command: &str) -> String {
}
pub fn bash_tool(session: WorkdirSessionHandle, output_dir: PathBuf) -> ToolDefinition {
routed_bash_tool(crate::routing::singleton_router(session), output_dir)
}
pub(crate) fn routed_bash_tool(
router: Arc<WorkdirSessionRouter>,
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 inline output, full-output spill, timeout and cancellation are owned by the WorkdirSession provider. This is not a sandbox.")
.description("Execute a shell command in one selected Workdir. Process start, bounded inline output, full-output spill, timeout and cancellation are owned by that 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(),
router: router.clone(),
output_dir: output_dir.clone(),
state: Arc::new(Mutex::new(BashExecutionState::default())),
});
+29 -11
View File
@@ -9,15 +9,18 @@ use serde::Deserialize;
use crate::error::ToolsError;
use crate::tracker::Tracker;
use workdir::{EditRequest, WorkdirPath, WorkdirSessionHandle};
use workdir::{EditRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter};
const DESCRIPTION: &str = "Replace a substring in an existing file. By default \
const DESCRIPTION: &str = "Replace a substring in an existing file in the selected Workdir attachment. By default \
`old_string` must be unique in the file; set `replace_all: true` to replace \
every occurrence. The file must have been read first (via the Read tool) in \
this session. Paths are relative to the bound Workdir.";
every occurrence. The file must have been read first through the same attachment. \
When exactly one Workdir is attached, target_workdir may be omitted.";
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub(crate) struct EditParams {
/// Worker-local alias of the Workdir attachment to use.
#[serde(default)]
pub target_workdir: Option<String>,
/// Logical path relative to the bound Workdir root.
pub file_path: String,
/// String to replace. Must be unique in the file unless `replace_all` is true.
@@ -30,7 +33,7 @@ pub(crate) struct EditParams {
}
pub(crate) struct EditTool {
session: WorkdirSessionHandle,
router: Arc<WorkdirSessionRouter>,
tracker: Tracker,
}
@@ -44,6 +47,14 @@ impl Tool for EditTool {
let params: EditParams = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid Edit input: {e}")))?;
let selected = crate::routing::resolve_session(
&self.router,
params.target_workdir.as_deref(),
workdir::WorkdirSessionCapability::Edit,
)?;
let tracker = self
.tracker
.scoped_attachment(&selected.alias, selected.generation);
let path = WorkdirPath::new(&params.file_path).map_err(ToolsError::from)?;
tracing::debug!(path = %path, replace_all = params.replace_all, "Edit");
@@ -59,9 +70,9 @@ impl Tool for EditTool {
}
let mutation_key = PathBuf::from(path.as_str());
let _mutation_permit = self.tracker.acquire_mutation(&mutation_key, &ctx).await;
let expected_hash = self.tracker.expected_workdir_hash(&path)?;
let result = self
let _mutation_permit = tracker.acquire_mutation(&mutation_key, &ctx).await;
let expected_hash = tracker.expected_workdir_hash(&path)?;
let result = selected
.session
.edit(EditRequest {
path: path.clone(),
@@ -73,7 +84,7 @@ impl Tool for EditTool {
.await
.map_err(ToolsError::from)?;
let replacements = result.replacements;
self.tracker.record_workdir_edit(
tracker.record_workdir_edit(
&path,
result.content_hash,
replacements,
@@ -122,8 +133,15 @@ fn make_preview(text: &str, needle: &str) -> String {
out
}
/// Factory for the `Edit` tool.
/// Factory for the `Edit` tool bound to one compatibility session.
pub fn edit_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefinition {
routed_edit_tool(crate::routing::singleton_router(session), tracker)
}
pub(crate) fn routed_edit_tool(
router: Arc<WorkdirSessionRouter>,
tracker: Tracker,
) -> ToolDefinition {
Arc::new(move || {
let schema = schemars::schema_for!(EditParams);
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
@@ -131,7 +149,7 @@ pub fn edit_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefinit
.description(DESCRIPTION)
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(EditTool {
session: session.clone(),
router: router.clone(),
tracker: tracker.clone(),
});
(meta, tool)
+17 -5
View File
@@ -4,7 +4,7 @@ use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use async_trait::async_trait;
use schemars::JsonSchema;
use serde::Deserialize;
use workdir::{GlobRequest, WorkdirPath, WorkdirSessionHandle};
use workdir::{GlobRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter};
use crate::ToolsError;
@@ -12,6 +12,9 @@ const RESULT_LIMIT: usize = 1000;
#[derive(Debug, Deserialize, JsonSchema)]
struct GlobParams {
/// Worker-local alias of the Workdir attachment to use.
#[serde(default)]
target_workdir: Option<String>,
/// Glob pattern, for example `**/*.rs` or `src/**/test_*.py`.
pattern: String,
/// Logical Workdir-relative directory. Defaults to the Workdir root.
@@ -20,7 +23,7 @@ struct GlobParams {
}
struct GlobTool {
session: WorkdirSessionHandle,
router: Arc<WorkdirSessionRouter>,
}
#[async_trait]
@@ -32,13 +35,18 @@ impl Tool for GlobTool {
) -> Result<ToolOutput, ToolError> {
let params: GlobParams = serde_json::from_str(input_json)
.map_err(|error| ToolError::InvalidArgument(format!("invalid Glob input: {error}")))?;
let selected = crate::routing::resolve_session(
&self.router,
params.target_workdir.as_deref(),
workdir::WorkdirSessionCapability::Glob,
)?;
let path = match params.path {
Some(path) => WorkdirPath::new(&path).map_err(ToolsError::from)?,
None => WorkdirPath::root(),
};
let pattern = params.pattern;
tracing::debug!(%pattern, %path, "Glob");
let result = self
let result = selected
.session
.glob(GlobRequest {
pattern: pattern.clone(),
@@ -75,13 +83,17 @@ impl Tool for GlobTool {
}
pub fn glob_tool(session: WorkdirSessionHandle) -> ToolDefinition {
routed_glob_tool(crate::routing::singleton_router(session))
}
pub(crate) fn routed_glob_tool(router: Arc<WorkdirSessionRouter>) -> ToolDefinition {
Arc::new(move || {
let schema = schemars::schema_for!(GlobParams);
let meta = ToolMeta::new("Glob")
.description("Find files matching a glob pattern inside the bound Workdir. Results are sorted and capped at 1000 entries. Paths are Workdir-relative.")
.description("Find files matching a glob pattern inside the selected Workdir. Results are sorted and capped at 1000 entries. Paths are Workdir-relative.")
.input_schema(serde_json::to_value(schema).expect("Glob schema serialization"));
let tool: Arc<dyn Tool> = Arc::new(GlobTool {
session: session.clone(),
router: router.clone(),
});
(meta, tool)
})
+19 -5
View File
@@ -4,7 +4,9 @@ use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use async_trait::async_trait;
use schemars::JsonSchema;
use serde::Deserialize;
use workdir::{GrepOutputMode, GrepRequest, WorkdirPath, WorkdirSessionHandle};
use workdir::{
GrepOutputMode, GrepRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter,
};
use crate::ToolsError;
@@ -21,6 +23,9 @@ enum OutputMode {
#[derive(Debug, Deserialize, JsonSchema)]
struct GrepParams {
/// Worker-local alias of the Workdir attachment to use.
#[serde(default)]
target_workdir: Option<String>,
pattern: String,
/// Workdir-relative path, or an absolute path covered by readable scope. Defaults to the Workdir root.
#[serde(default)]
@@ -48,7 +53,7 @@ struct GrepParams {
}
struct GrepTool {
session: WorkdirSessionHandle,
router: Arc<WorkdirSessionRouter>,
}
#[async_trait]
@@ -60,6 +65,11 @@ impl Tool for GrepTool {
) -> Result<ToolOutput, ToolError> {
let params: GrepParams = serde_json::from_str(input_json)
.map_err(|error| ToolError::InvalidArgument(format!("invalid Grep input: {error}")))?;
let selected = crate::routing::resolve_session(
&self.router,
params.target_workdir.as_deref(),
workdir::WorkdirSessionCapability::Grep,
)?;
let path = match params.path {
Some(path) => WorkdirPath::new_scoped(&path).map_err(ToolsError::from)?,
None => WorkdirPath::root(),
@@ -74,7 +84,7 @@ impl Tool for GrepTool {
.map(|context| (context, context))
.unwrap_or((params.before.unwrap_or(0), params.after.unwrap_or(0)));
let head_limit = params.head_limit.unwrap_or(DEFAULT_HEAD_LIMIT);
let result = self
let result = selected
.session
.grep(GrepRequest {
pattern: params.pattern,
@@ -126,13 +136,17 @@ impl Tool for GrepTool {
}
pub fn grep_tool(session: WorkdirSessionHandle) -> ToolDefinition {
routed_grep_tool(crate::routing::singleton_router(session))
}
pub(crate) fn routed_grep_tool(router: Arc<WorkdirSessionRouter>) -> ToolDefinition {
Arc::new(move || {
let schema = schemars::schema_for!(GrepParams);
let meta = ToolMeta::new("Grep")
.description("Search a Workdir file or directory with a regex. Content results group lines by file; `>` marks matching lines and unmarked lines are context. Directory traversal executes inside the WorkdirSession provider. Results are bounded and Workdir-relative.")
.description("Search a file or directory in the selected Workdir with a regex. Content results group lines by file; `>` marks matching lines and unmarked lines are context. Directory traversal executes inside the WorkdirSession provider. Results are bounded and Workdir-relative.")
.input_schema(serde_json::to_value(schema).expect("Grep schema serialization"));
let tool: Arc<dyn Tool> = Arc::new(GrepTool {
session: session.clone(),
router: router.clone(),
});
(meta, tool)
})
+83 -8
View File
@@ -17,6 +17,7 @@ mod edit;
mod glob;
mod grep;
mod read;
mod routing;
mod view_image;
mod web;
mod write;
@@ -28,11 +29,11 @@ pub use glob::glob_tool;
pub use grep::grep_tool;
pub use read::read_tool;
pub use tracker::{ChangeStat, Tracker};
pub use view_image::view_image_tool;
pub use view_image::{routed_view_image_tool, view_image_tool};
pub use web::{web_fetch_tool, web_search_tool};
pub use write::write_tool;
/// Build the local filesystem/command tool surface implemented by a WorkdirSession.
/// Build the filesystem/command tool surface for a single Workdir session.
/// Profile/manifest policy may narrow this set further in the Engine.
pub fn core_builtin_tools(
session: workdir::WorkdirSessionHandle,
@@ -42,24 +43,75 @@ pub fn core_builtin_tools(
use workdir::WorkdirSessionCapability;
let capabilities = session.capabilities();
let router = routing::singleton_router(session);
let mut tools = Vec::with_capacity(6);
if capabilities.supports(WorkdirSessionCapability::Read) {
tools.push(read_tool(session.clone(), tracker.clone()));
tools.push(read::routed_read_tool(router.clone(), tracker.clone()));
}
if capabilities.supports(WorkdirSessionCapability::Write) {
tools.push(write_tool(session.clone(), tracker.clone()));
tools.push(write::routed_write_tool(router.clone(), tracker.clone()));
}
if capabilities.supports(WorkdirSessionCapability::Edit) {
tools.push(edit_tool(session.clone(), tracker));
tools.push(edit::routed_edit_tool(router.clone(), tracker));
}
if capabilities.supports(WorkdirSessionCapability::Glob) {
tools.push(glob_tool(session.clone()));
tools.push(glob::routed_glob_tool(router.clone()));
}
if capabilities.supports(WorkdirSessionCapability::Grep) {
tools.push(grep_tool(session.clone()));
tools.push(grep::routed_grep_tool(router.clone()));
}
if capabilities.supports(WorkdirSessionCapability::Command) {
tools.push(bash_tool(session, bash_output_dir));
tools.push(bash::routed_bash_tool(router, bash_output_dir));
}
tools
}
/// Build the stable routed tool surface for a Worker's live attachment set.
///
/// The definitions remain installed while the set is empty or changes at runtime;
/// each invocation resolves `target_workdir` against the router and then lets the
/// selected provider enforce its capabilities.
pub fn routed_builtin_tools(
router: std::sync::Arc<workdir::WorkdirSessionRouter>,
tracker: Tracker,
bash_output_dir: std::path::PathBuf,
) -> Vec<agen::tool::ToolDefinition> {
vec![
read::routed_read_tool(router.clone(), tracker.clone()),
write::routed_write_tool(router.clone(), tracker.clone()),
edit::routed_edit_tool(router.clone(), tracker),
glob::routed_glob_tool(router.clone()),
grep::routed_grep_tool(router.clone()),
bash::routed_bash_tool(router, bash_output_dir),
]
}
pub fn routed_builtin_tools_for_current_capabilities(
router: std::sync::Arc<workdir::WorkdirSessionRouter>,
tracker: Tracker,
bash_output_dir: std::path::PathBuf,
) -> Vec<agen::tool::ToolDefinition> {
use workdir::WorkdirSessionCapability;
let capabilities = router.capabilities();
let mut tools = Vec::with_capacity(6);
if capabilities.supports(WorkdirSessionCapability::Read) {
tools.push(read::routed_read_tool(router.clone(), tracker.clone()));
}
if capabilities.supports(WorkdirSessionCapability::Write) {
tools.push(write::routed_write_tool(router.clone(), tracker.clone()));
}
if capabilities.supports(WorkdirSessionCapability::Edit) {
tools.push(edit::routed_edit_tool(router.clone(), tracker));
}
if capabilities.supports(WorkdirSessionCapability::Glob) {
tools.push(glob::routed_glob_tool(router.clone()));
}
if capabilities.supports(WorkdirSessionCapability::Grep) {
tools.push(grep::routed_grep_tool(router.clone()));
}
if capabilities.supports(WorkdirSessionCapability::Command) {
tools.push(bash::routed_bash_tool(router, bash_output_dir));
}
tools
}
@@ -75,6 +127,29 @@ pub fn read_only_builtin_tools(
core_builtin_tools(session, Tracker::new(), std::path::PathBuf::new())
}
pub fn read_only_routed_builtin_tools(
source: std::sync::Arc<workdir::WorkdirSessionRouter>,
) -> Vec<agen::tool::ToolDefinition> {
let router = std::sync::Arc::new(workdir::WorkdirSessionRouter::new());
for alias in source.aliases() {
let Ok(selected) = source.resolve(Some(alias.as_str())) else {
continue;
};
router
.attach(
alias,
std::sync::Arc::new(workdir::ReadOnlyWorkdirSession::new(selected.session)),
)
.expect("read-only projection preserves unique aliases");
}
let tracker = Tracker::new();
vec![
read::routed_read_tool(router.clone(), tracker),
glob::routed_glob_tool(router.clone()),
grep::routed_grep_tool(router),
]
}
pub fn web_builtin_tools(
web_config: Option<manifest::WebConfig>,
) -> Vec<agen::tool::ToolDefinition> {
+29 -9
View File
@@ -8,18 +8,22 @@ use serde::Deserialize;
use crate::error::ToolsError;
use crate::tracker::Tracker;
use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle};
use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter};
const DESCRIPTION: &str = "Read a text file from the local filesystem. \
const DESCRIPTION: &str = "Read a text file from a Workdir attachment selected by its Worker-local alias. \
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 Workdir-relative unless an absolute path is explicitly readable.";
modify it. Paths are Workdir-relative unless an absolute path is explicitly readable. \
When exactly one Workdir is attached, target_workdir may be omitted.";
const DEFAULT_LIMIT: usize = 2000;
const PROVIDER_BYTE_LIMIT: usize = 256 * 1024;
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub(crate) struct ReadParams {
/// Worker-local alias of the Workdir attachment to use.
#[serde(default)]
pub target_workdir: Option<String>,
/// 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.
@@ -31,7 +35,7 @@ pub(crate) struct ReadParams {
}
pub(crate) struct ReadTool {
session: WorkdirSessionHandle,
router: Arc<WorkdirSessionRouter>,
tracker: Tracker,
}
@@ -44,13 +48,21 @@ impl Tool for ReadTool {
) -> Result<ToolOutput, ToolError> {
let params: ReadParams = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid Read input: {e}")))?;
let selected = crate::routing::resolve_session(
&self.router,
params.target_workdir.as_deref(),
workdir::WorkdirSessionCapability::Read,
)?;
let tracker = self
.tracker
.scoped_attachment(&selected.alias, selected.generation);
let offset = params.offset.unwrap_or(0);
let limit = params.limit.unwrap_or(DEFAULT_LIMIT).max(1);
let path = WorkdirPath::new_scoped(&params.file_path).map_err(ToolsError::from)?;
tracing::debug!(path = %path, offset, limit, "Read");
let result = self
let result = selected
.session
.read(ReadRequest {
path: path.clone(),
@@ -60,7 +72,7 @@ impl Tool for ReadTool {
})
.await
.map_err(ToolsError::from)?;
self.tracker.record_workdir_hash(&path, result.content_hash);
tracker.record_workdir_hash(&path, result.content_hash);
let text = String::from_utf8_lossy(&result.bytes).into_owned();
let rendered = render_provider_read(
@@ -144,8 +156,15 @@ fn render_numbered(text: &str, offset: usize, limit: usize) -> Rendered {
}
}
/// Factory for the `Read` tool.
/// Factory for the `Read` tool bound to one compatibility session.
pub fn read_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefinition {
routed_read_tool(crate::routing::singleton_router(session), tracker)
}
pub(crate) fn routed_read_tool(
router: Arc<WorkdirSessionRouter>,
tracker: Tracker,
) -> ToolDefinition {
Arc::new(move || {
let schema = schemars::schema_for!(ReadParams);
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
@@ -153,7 +172,7 @@ pub fn read_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefinit
.description(DESCRIPTION)
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(ReadTool {
session: session.clone(),
router: router.clone(),
tracker: tracker.clone(),
});
(meta, tool)
@@ -199,6 +218,7 @@ mod tests {
// History recorded
assert!(
tracker
.scoped_attachment(&workdir::WorkdirAttachmentAlias::new("workdir").unwrap(), 0,)
.expected_workdir_hash(&WorkdirPath::new("a.txt").unwrap())
.is_ok()
);
@@ -229,7 +249,7 @@ mod tests {
#[tokio::test]
async fn read_tool_missing_file() {
let (dir, fs, tracker) = setup();
let (_dir, fs, tracker) = setup();
let def = read_tool(fs, tracker);
let (_, tool) = def();
let input = serde_json::json!({
+32
View File
@@ -0,0 +1,32 @@
use std::sync::Arc;
use agen::tool::ToolError;
use workdir::{
ResolvedWorkdirSession, WorkdirAttachmentAlias, WorkdirSessionCapability, WorkdirSessionHandle,
WorkdirSessionRouter,
};
pub(crate) fn singleton_router(session: WorkdirSessionHandle) -> Arc<WorkdirSessionRouter> {
let router = Arc::new(WorkdirSessionRouter::new());
router
.attach(
WorkdirAttachmentAlias::new("workdir").expect("static attachment alias is valid"),
session,
)
.expect("fresh singleton router accepts its session");
router
}
pub(crate) fn resolve_session(
router: &WorkdirSessionRouter,
target_workdir: Option<&str>,
required: WorkdirSessionCapability,
) -> Result<ResolvedWorkdirSession, ToolError> {
let selected = router
.resolve(target_workdir)
.map_err(|error| ToolError::InvalidArgument(error.to_string()))?;
if !selected.session.capabilities().supports(required) {
return Err(crate::ToolsError::from(workdir::WorkdirError::Unsupported(required)).into());
}
Ok(selected)
}
+75 -18
View File
@@ -65,9 +65,15 @@ fn hash_bytes(bytes: &[u8]) -> ContentHash {
hasher.finalize().into()
}
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
struct TrackedPath {
namespace: Option<Arc<str>>,
path: PathBuf,
}
#[derive(Debug, Clone, Default)]
struct FileMutationCoordinator {
locks: Arc<tokio::sync::Mutex<HashMap<PathBuf, Arc<tokio::sync::Mutex<()>>>>>,
locks: Arc<tokio::sync::Mutex<HashMap<TrackedPath, Arc<tokio::sync::Mutex<()>>>>>,
}
pub(crate) struct FileMutationPermit {
@@ -75,8 +81,7 @@ pub(crate) struct FileMutationPermit {
}
impl FileMutationCoordinator {
async fn acquire(&self, path: &Path) -> FileMutationPermit {
let key = file_mutation_key(path);
async fn acquire(&self, key: TrackedPath) -> FileMutationPermit {
let lock = {
let mut locks = self.locks.lock().await;
locks
@@ -127,12 +132,12 @@ pub struct ChangeStat {
#[derive(Debug, Default)]
struct Inner {
/// Hash of each file's last observed contents, keyed by canonical path.
hashes: HashMap<PathBuf, ContentHash>,
/// Hash of each file's last observed contents, keyed by attachment and canonical path.
hashes: HashMap<TrackedPath, ContentHash>,
/// Line count paired with observations that included the file content.
line_counts: HashMap<PathBuf, usize>,
line_counts: HashMap<TrackedPath, usize>,
/// LRU list of touched files. Front = most recently touched.
recency: VecDeque<PathBuf>,
recency: VecDeque<TrackedPath>,
/// Successful Write/Edit mutations attributed to this session's tools.
change_stat: ChangeStat,
}
@@ -146,6 +151,7 @@ struct Inner {
pub struct Tracker {
inner: Arc<Mutex<Inner>>,
mutations: FileMutationCoordinator,
namespace: Option<Arc<str>>,
}
impl Tracker {
@@ -154,6 +160,33 @@ impl Tracker {
Self::default()
}
/// Return a view whose read-before-mutation state is isolated to one
/// attachment incarnation (the Worker-local alias plus router generation),
/// while aggregate recency/change statistics remain shared with the Worker
/// tracker.
pub fn scoped(&self, namespace: impl Into<String>) -> Self {
Self {
inner: self.inner.clone(),
mutations: self.mutations.clone(),
namespace: Some(Arc::from(namespace.into())),
}
}
pub(crate) fn scoped_attachment(
&self,
alias: &workdir::WorkdirAttachmentAlias,
generation: u64,
) -> Self {
self.scoped(format!("{}#{generation}", alias.as_str()))
}
fn key(&self, path: PathBuf) -> TrackedPath {
TrackedPath {
namespace: self.namespace.clone(),
path,
}
}
/// Acquire the per-target-file mutation guard shared by `Write` and `Edit`.
///
/// The guard is keyed by canonical target path where possible so equivalent
@@ -170,7 +203,9 @@ impl Tracker {
call_index = ctx.call_index,
"acquire file mutation guard"
);
self.mutations.acquire(path).await
self.mutations
.acquire(self.key(file_mutation_key(path)))
.await
}
/// Record that `path` has been observed with the given content bytes.
@@ -182,7 +217,7 @@ impl Tracker {
/// Also bumps `path` to the front of the recency list. If the list
/// grows past [`RECENCY_CAPACITY`], the oldest entry is evicted.
pub fn record(&self, path: &Path, bytes: &[u8]) {
let key = canonicalize_or_owned(path);
let key = self.key(canonicalize_or_owned(path));
let hash = hash_bytes(bytes);
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
inner.hashes.insert(key.clone(), hash);
@@ -198,7 +233,7 @@ impl Tracker {
}
pub fn record_workdir_content(&self, path: &workdir::WorkdirPath, content: &[u8]) {
let key = PathBuf::from(path.as_str());
let key = self.key(PathBuf::from(path.as_str()));
let hash = hash_bytes(content);
let line_count = String::from_utf8_lossy(content).lines().count();
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
@@ -212,16 +247,17 @@ impl Tracker {
}
pub fn observed_workdir_line_count(&self, path: &workdir::WorkdirPath) -> Option<usize> {
let key = self.key(PathBuf::from(path.as_str()));
self.inner
.lock()
.unwrap_or_else(|e| e.into_inner())
.line_counts
.get(Path::new(path.as_str()))
.get(&key)
.copied()
}
pub fn record_workdir_hash(&self, path: &workdir::WorkdirPath, hash: workdir::ContentHash) {
let key = PathBuf::from(path.as_str());
let key = self.key(PathBuf::from(path.as_str()));
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
inner.hashes.insert(key.clone(), hash);
inner.recency.retain(|candidate| candidate != &key);
@@ -253,7 +289,7 @@ impl Tracker {
) {
let added = added_lines_per_replacement.saturating_mul(replacements);
let deleted = deleted_lines_per_replacement.saturating_mul(replacements);
let key = PathBuf::from(path.as_str());
let key = self.key(PathBuf::from(path.as_str()));
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
inner.change_stat.added = inner.change_stat.added.saturating_add(added as u64);
inner.change_stat.deleted = inner.change_stat.deleted.saturating_add(deleted as u64);
@@ -279,14 +315,14 @@ impl Tracker {
&self,
path: &workdir::WorkdirPath,
) -> Result<workdir::ContentHash, ToolsError> {
let key = PathBuf::from(path.as_str());
let key = self.key(PathBuf::from(path.as_str()));
self.inner
.lock()
.unwrap_or_else(|error| error.into_inner())
.hashes
.get(&key)
.copied()
.ok_or_else(|| ToolsError::NotRead(key))
.ok_or_else(|| ToolsError::NotRead(key.path.clone()))
}
/// Verify that `path` was previously recorded and its current bytes
@@ -296,7 +332,7 @@ impl Tracker {
/// - If the current content hashes differ from the recorded value,
/// returns [`ToolsError::ExternallyModified`].
pub fn verify(&self, path: &Path, current_bytes: &[u8]) -> Result<(), ToolsError> {
let key = canonicalize_or_owned(path);
let key = self.key(canonicalize_or_owned(path));
let guard = self.inner.lock().unwrap_or_else(|e| e.into_inner());
let recorded = guard
.hashes
@@ -316,13 +352,18 @@ impl Tracker {
/// can pass them as default references to the compaction worker.
pub fn recent_files(&self, n: usize) -> Vec<PathBuf> {
let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
inner.recency.iter().take(n).cloned().collect()
inner
.recency
.iter()
.take(n)
.map(|key| key.path.clone())
.collect()
}
/// Returns true if `path` has a history entry. Test-only.
#[cfg(test)]
pub(crate) fn has(&self, path: &Path) -> bool {
let key = canonicalize_or_owned(path);
let key = self.key(canonicalize_or_owned(path));
self.inner
.lock()
.unwrap_or_else(|e| e.into_inner())
@@ -420,6 +461,22 @@ mod tests {
}
}
#[test]
fn attachment_namespaces_do_not_share_read_state() {
let tracker = Tracker::new();
let left = tracker.scoped("left");
let right = tracker.scoped("right");
let path = workdir::WorkdirPath::new("same.txt").unwrap();
left.record_workdir_content(&path, b"left");
assert!(left.expected_workdir_hash(&path).is_ok());
assert!(matches!(
right.expected_workdir_hash(&path),
Err(ToolsError::NotRead(_))
));
}
#[test]
fn clone_shares_state() {
let dir = TempDir::new().unwrap();
+19 -6
View File
@@ -7,24 +7,28 @@ use agen::tool::{
};
use async_trait::async_trait;
use serde::Deserialize;
use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle};
use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter};
use crate::error::ToolsError;
/// Maximum image body accepted for one model request.
pub const MAX_IMAGE_BYTES: usize = 10 * 1024 * 1024;
const DESCRIPTION: &str = "Attach an image from the bound Workdir to the next model request. \
The path must be logical and Workdir-relative. Supported formats: PNG, JPEG, GIF, and WebP.";
const DESCRIPTION: &str = "Attach an image from the selected Workdir attachment to the next model request. \
The path must be logical and Workdir-relative. Supported formats: PNG, JPEG, GIF, and WebP. \
When exactly one Workdir is attached, target_workdir may be omitted.";
#[derive(Debug, Deserialize, schemars::JsonSchema)]
struct ViewImageParams {
/// Worker-local alias of the Workdir attachment to use.
#[serde(default)]
target_workdir: Option<String>,
/// Logical path relative to the bound Workdir root.
path: String,
}
struct ViewImageTool {
session: WorkdirSessionHandle,
router: Arc<WorkdirSessionRouter>,
}
#[async_trait]
@@ -37,8 +41,13 @@ impl Tool for ViewImageTool {
let input: ViewImageParams = serde_json::from_str(input_json).map_err(|error| {
ToolError::InvalidArgument(format!("invalid ViewImage input: {error}"))
})?;
let selected = crate::routing::resolve_session(
&self.router,
input.target_workdir.as_deref(),
workdir::WorkdirSessionCapability::Read,
)?;
let path = WorkdirPath::new(&input.path).map_err(ToolsError::from)?;
let result = self
let result = selected
.session
.read(ReadRequest {
path: path.clone(),
@@ -75,6 +84,10 @@ impl Tool for ViewImageTool {
}
pub fn view_image_tool(session: WorkdirSessionHandle) -> ToolDefinition {
routed_view_image_tool(crate::routing::singleton_router(session))
}
pub fn routed_view_image_tool(router: Arc<WorkdirSessionRouter>) -> ToolDefinition {
Arc::new(move || {
let schema = schemars::schema_for!(ViewImageParams);
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
@@ -82,7 +95,7 @@ pub fn view_image_tool(session: WorkdirSessionHandle) -> ToolDefinition {
.description(DESCRIPTION)
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(ViewImageTool {
session: session.clone(),
router: router.clone(),
});
(meta, tool)
})
+39 -17
View File
@@ -9,15 +9,20 @@ use serde::Deserialize;
use crate::error::ToolsError;
use crate::tracker::Tracker;
use workdir::{StatRequest, WorkdirError, WorkdirPath, WorkdirSessionHandle, WriteRequest};
use workdir::{
StatRequest, WorkdirError, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter,
WriteRequest,
};
const DESCRIPTION: &str = "Create a new file or overwrite an existing one with \
the given content. Missing parent directories within scope are created \
automatically. Existing files must have been read first (via the Read tool) \
in this session. Paths are relative to the bound Workdir.";
const DESCRIPTION: &str = "Create a new file or overwrite an existing one in the selected Workdir attachment. \
Missing parent directories within scope are created automatically. Existing files must have been read first \
through the same attachment. When exactly one Workdir is attached, target_workdir may be omitted.";
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub(crate) struct WriteParams {
/// Worker-local alias of the Workdir attachment to use.
#[serde(default)]
pub target_workdir: Option<String>,
/// Logical path relative to the bound Workdir root.
pub file_path: String,
/// Full content to write. Overwrites any existing content.
@@ -25,7 +30,7 @@ pub(crate) struct WriteParams {
}
pub(crate) struct WriteTool {
session: WorkdirSessionHandle,
router: Arc<WorkdirSessionRouter>,
tracker: Tracker,
}
@@ -39,19 +44,31 @@ impl Tool for WriteTool {
let params: WriteParams = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid Write input: {e}")))?;
let selected = crate::routing::resolve_session(
&self.router,
params.target_workdir.as_deref(),
workdir::WorkdirSessionCapability::Write,
)?;
let tracker = self
.tracker
.scoped_attachment(&selected.alias, selected.generation);
let path = WorkdirPath::new(&params.file_path).map_err(ToolsError::from)?;
tracing::debug!(path = %path, bytes = params.content.len(), "Write");
let mutation_key = PathBuf::from(path.as_str());
let _mutation_permit = self.tracker.acquire_mutation(&mutation_key, &ctx).await;
let expected_hash = match self.session.stat(StatRequest { path: path.clone() }).await {
Ok(_) => Some(self.tracker.expected_workdir_hash(&path)?),
let _mutation_permit = tracker.acquire_mutation(&mutation_key, &ctx).await;
let expected_hash = match selected
.session
.stat(StatRequest { path: path.clone() })
.await
{
Ok(_) => Some(tracker.expected_workdir_hash(&path)?),
Err(WorkdirError::NotFound(_)) => None,
Err(error) => return Err(ToolsError::from(error).into()),
};
let old_line_count = self.tracker.observed_workdir_line_count(&path).unwrap_or(0);
let outcome = self
let old_line_count = tracker.observed_workdir_line_count(&path).unwrap_or(0);
let outcome = selected
.session
.write(WriteRequest {
path: path.clone(),
@@ -61,10 +78,8 @@ impl Tool for WriteTool {
.await
.map_err(ToolsError::from)?;
self.tracker
.record_change(params.content.lines().count(), old_line_count);
self.tracker
.record_workdir_content(&path, params.content.as_bytes());
tracker.record_change(params.content.lines().count(), old_line_count);
tracker.record_workdir_content(&path, params.content.as_bytes());
let summary = format!(
"{} {} ({} bytes)",
@@ -84,8 +99,15 @@ impl Tool for WriteTool {
}
}
/// Factory for the `Write` tool.
/// Factory for the `Write` tool bound to one compatibility session.
pub fn write_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefinition {
routed_write_tool(crate::routing::singleton_router(session), tracker)
}
pub(crate) fn routed_write_tool(
router: Arc<WorkdirSessionRouter>,
tracker: Tracker,
) -> ToolDefinition {
Arc::new(move || {
let schema = schemars::schema_for!(WriteParams);
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
@@ -93,7 +115,7 @@ pub fn write_tool(session: WorkdirSessionHandle, tracker: Tracker) -> ToolDefini
.description(DESCRIPTION)
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(WriteTool {
session: session.clone(),
router: router.clone(),
tracker: tracker.clone(),
});
(meta, tool)
+297 -2
View File
@@ -14,8 +14,12 @@ use agen::tool::{
use manifest::{Permission, Scope, ScopeConfig, ScopeRule};
use serde_json::json;
use tempfile::TempDir;
use tools::{Tracker, core_builtin_tools, view_image_tool};
use workdir::{LocalWorkdirSession, WorkdirSessionHandle};
use tools::{
Tracker, core_builtin_tools, routed_builtin_tools, routed_view_image_tool, view_image_tool,
};
use workdir::{
LocalWorkdirSession, WorkdirAttachmentAlias, WorkdirSessionHandle, WorkdirSessionRouter,
};
fn scope_with_spill(workspace: &Path, spill: &Path) -> Scope {
let base = Scope::writable(workspace).unwrap();
@@ -66,6 +70,30 @@ fn setup() -> (TempDir, TempDir, Registry) {
(dir, spill, reg)
}
fn setup_routed() -> (
TempDir,
TempDir,
TempDir,
Arc<WorkdirSessionRouter>,
Registry,
) {
let left = TempDir::new().unwrap();
let right = TempDir::new().unwrap();
let spill = TempDir::new().unwrap();
let router = Arc::new(WorkdirSessionRouter::new());
for (alias, dir) in [("left", &left), ("right", &right)] {
let scope = scope_with_spill(dir.path(), spill.path());
let session: WorkdirSessionHandle =
Arc::new(LocalWorkdirSession::new(scope, dir.path().to_path_buf()));
router
.attach(WorkdirAttachmentAlias::new(alias).unwrap(), session)
.unwrap();
}
let definitions =
routed_builtin_tools(router.clone(), Tracker::new(), spill.path().to_path_buf());
(left, right, spill, router, Registry::new(definitions))
}
async fn call(tool: &Arc<dyn Tool>, input: serde_json::Value) -> agen::tool::ToolOutput {
tool.execute(&input.to_string(), Default::default())
.await
@@ -104,6 +132,273 @@ fn meta_has_description_and_schema() {
}
}
#[test]
fn routed_tool_schemas_share_optional_target_workdir() {
let (_left, _right, _spill, router, mut reg) = setup_routed();
reg.entries.push(routed_view_image_tool(router)());
for (meta, _) in &reg.entries {
let target = &meta.input_schema["properties"]["target_workdir"];
assert!(target.is_object(), "{} lacks target_workdir", meta.name);
assert_eq!(target["type"], json!(["string", "null"]));
}
}
#[tokio::test]
async fn routed_tools_require_alias_only_when_attachment_set_is_ambiguous() {
let (_left, _right, _spill, router, reg) = setup_routed();
let error = call_err(&reg.get("Read"), json!({ "file_path": "same.txt" })).await;
let message = error.to_string();
assert!(message.contains("target_workdir_required"), "{message}");
assert!(message.contains("left"), "{message}");
assert!(message.contains("right"), "{message}");
let unknown = call_err(
&reg.get("Glob"),
json!({ "target_workdir": "missing", "pattern": "*" }),
)
.await
.to_string();
assert!(unknown.contains("unknown_target_workdir"), "{unknown}");
assert!(unknown.contains("available_aliases"), "{unknown}");
router.close_all().await.unwrap();
}
#[tokio::test]
async fn empty_routed_surface_returns_no_attachment_error() {
let router = Arc::new(WorkdirSessionRouter::new());
let reg = Registry::new(routed_builtin_tools(
router,
Tracker::new(),
Path::new("unused").to_path_buf(),
));
let message = call_err(&reg.get("Read"), json!({ "file_path": "file.txt" }))
.await
.to_string();
assert!(message.contains("no_workdir_attached"), "{message}");
}
#[tokio::test]
async fn target_alias_routes_each_tool_and_isolates_read_before_edit() {
let (left, right, _spill, router, mut reg) = setup_routed();
std::fs::write(left.path().join("same.txt"), "left value\n").unwrap();
std::fs::write(right.path().join("same.txt"), "right value\n").unwrap();
std::fs::write(left.path().join("only-left.log"), "LEFT-NEEDLE\n").unwrap();
std::fs::write(right.path().join("only-right.log"), "RIGHT-NEEDLE\n").unwrap();
let read = reg.get("Read");
let left_read = call(
&read,
json!({ "target_workdir": "left", "file_path": "same.txt" }),
)
.await;
assert!(left_read.content.unwrap().contains("left value"));
let edit = reg.get("Edit");
let isolated = call_err(
&edit,
json!({
"target_workdir": "right",
"file_path": "same.txt",
"old_string": "right",
"new_string": "edited"
}),
)
.await;
assert!(isolated.to_string().contains("has not been read"));
call(
&read,
json!({ "target_workdir": "right", "file_path": "same.txt" }),
)
.await;
call(
&edit,
json!({
"target_workdir": "right",
"file_path": "same.txt",
"old_string": "right",
"new_string": "edited"
}),
)
.await;
assert_eq!(
std::fs::read_to_string(right.path().join("same.txt")).unwrap(),
"edited value\n"
);
assert_eq!(
std::fs::read_to_string(left.path().join("same.txt")).unwrap(),
"left value\n"
);
let glob = call(
&reg.get("Glob"),
json!({ "target_workdir": "left", "pattern": "*.log" }),
)
.await;
assert!(glob.content.unwrap().contains("only-left.log"));
let grep = call(
&reg.get("Grep"),
json!({
"target_workdir": "right",
"pattern": "RIGHT-NEEDLE",
"output_mode": "content"
}),
)
.await;
assert!(grep.content.unwrap().contains("only-right.log"));
let bash = call(
&reg.get("Bash"),
json!({ "target_workdir": "right", "command": "pwd" }),
)
.await;
assert_eq!(
std::fs::canonicalize(bash.content.unwrap().trim()).unwrap(),
std::fs::canonicalize(right.path()).unwrap()
);
let png = b"\x89PNG\r\n\x1a\nrouted";
std::fs::write(left.path().join("image.png"), png).unwrap();
reg.entries.push(routed_view_image_tool(router.clone())());
let image = call(
&reg.get("ViewImage"),
json!({ "target_workdir": "left", "path": "image.png" }),
)
.await;
let agen::tool::Attachment::Image(image) = &image.attachments[0];
assert_eq!(image.data(), png);
router.close_all().await.unwrap();
}
#[tokio::test]
async fn selected_attachment_capability_is_checked_without_fallback() {
let left = TempDir::new().unwrap();
let right = TempDir::new().unwrap();
let spill = TempDir::new().unwrap();
let router = Arc::new(WorkdirSessionRouter::new());
router
.attach(
WorkdirAttachmentAlias::new("writable").unwrap(),
Arc::new(LocalWorkdirSession::new(
scope_with_spill(left.path(), spill.path()),
left.path().to_path_buf(),
)),
)
.unwrap();
router
.attach(
WorkdirAttachmentAlias::new("readonly").unwrap(),
Arc::new(LocalWorkdirSession::materialized_bound(
workdir::Workdir::new("readonly"),
right.path().to_path_buf(),
right.path().to_path_buf(),
manifest::SharedScope::new(scope_with_spill(right.path(), spill.path())),
workdir::WorkdirSessionCapabilities::READ_ONLY,
)),
)
.unwrap();
let reg = Registry::new(routed_builtin_tools(
router.clone(),
Tracker::new(),
spill.path().to_path_buf(),
));
let message = call_err(
&reg.get("Write"),
json!({
"target_workdir": "readonly",
"file_path": "created.txt",
"content": "must not be routed elsewhere"
}),
)
.await
.to_string();
assert!(message.contains("does not support Write"), "{message}");
assert!(!left.path().join("created.txt").exists());
assert!(!right.path().join("created.txt").exists());
router.close_all().await.unwrap();
}
#[tokio::test]
async fn routed_tools_observe_detach_after_registration() {
let (_left, _right, _spill, router, reg) = setup_routed();
let alias = WorkdirAttachmentAlias::new("left").unwrap();
router.detach(&alias).await.unwrap();
let message = call_err(
&reg.get("Glob"),
json!({ "target_workdir": "left", "pattern": "*" }),
)
.await
.to_string();
assert!(message.contains("unknown_target_workdir"), "{message}");
assert!(message.contains("right"), "{message}");
router.close_all().await.unwrap();
}
#[tokio::test]
async fn detach_and_reattach_same_alias_requires_a_fresh_read() {
let first = TempDir::new().unwrap();
let second = TempDir::new().unwrap();
let spill = TempDir::new().unwrap();
std::fs::write(first.path().join("same.txt"), "same bytes\n").unwrap();
std::fs::write(second.path().join("same.txt"), "same bytes\n").unwrap();
let router = Arc::new(WorkdirSessionRouter::new());
let alias = WorkdirAttachmentAlias::new("checkout").unwrap();
router
.attach(
alias.clone(),
Arc::new(LocalWorkdirSession::new(
scope_with_spill(first.path(), spill.path()),
first.path().to_path_buf(),
)),
)
.unwrap();
let reg = Registry::new(routed_builtin_tools(
router.clone(),
Tracker::new(),
spill.path().to_path_buf(),
));
call(
&reg.get("Read"),
json!({ "target_workdir": "checkout", "file_path": "same.txt" }),
)
.await;
router.detach(&alias).await.unwrap();
router
.attach(
alias.clone(),
Arc::new(LocalWorkdirSession::new(
scope_with_spill(second.path(), spill.path()),
second.path().to_path_buf(),
)),
)
.unwrap();
let error = call_err(
&reg.get("Edit"),
json!({
"target_workdir": "checkout",
"file_path": "same.txt",
"old_string": "same",
"new_string": "changed"
}),
)
.await;
assert!(error.to_string().contains("has not been read"), "{error}");
assert_eq!(
std::fs::read_to_string(second.path().join("same.txt")).unwrap(),
"same bytes\n"
);
router.close_all().await.unwrap();
}
#[tokio::test]
async fn view_image_reads_scoped_bytes_into_durable_tool_detail() {
let dir = TempDir::new().unwrap();
+2 -2
View File
@@ -164,7 +164,7 @@ fn request_from_selection(selection: Selection) -> BackendCreateWorkerRequest {
display_name: selection.display_name,
profile: Some(selection.profile),
initial_submit: Vec::new(),
working_directory: None,
workdir_attachments: Vec::new(),
ticket_assignment: None,
control_operation_id: None,
}
@@ -477,7 +477,7 @@ mod tests {
assert_eq!(request.display_name, "Coder one");
assert_eq!(request.profile.as_deref(), Some("builtin:coder"));
assert!(request.initial_submit.is_empty());
assert!(request.working_directory.is_none());
assert!(request.workdir_attachments.is_empty());
assert!(request.ticket_assignment.is_none());
}
}
+39 -20
View File
@@ -408,10 +408,22 @@ fn short_worker_id(worker: &BackendWorkerSummary) -> String {
}
fn working_directory_text(worker: &BackendWorkerSummary) -> String {
let Some(wd) = worker.working_directory.as_ref() else {
if worker.workdir_attachments.is_empty() {
return "wd:—".to_string();
};
format!("wd:{}・{}", wd.repository_key, wd.working_directory_id)
}
if worker.workdir_attachments.len() > 1 {
return format!("wd:{} attachments", worker.workdir_attachments.len());
}
let attachment = &worker.workdir_attachments[0];
let directory = &attachment.working_directory;
let label = directory
.display_name
.as_deref()
.unwrap_or(&directory.repository_key);
format!(
"wd:{}:{}・{}",
attachment.alias, label, directory.working_directory_id
)
}
#[cfg(test)]
@@ -456,7 +468,7 @@ mod tests {
can_stop: true,
can_spawn_followup: false,
},
working_directory: None,
workdir_attachments: Vec::new(),
diagnostics: Vec::new(),
}
}
@@ -529,22 +541,26 @@ mod tests {
worker.label = "Coder · T-585".to_string();
worker.state = "stopped".to_string();
worker.worker_state = None;
worker.working_directory = Some(
worker.workdir_attachments = vec![
serde_json::from_value(serde_json::json!({
"working_directory_id": "001a06a9f0202000000",
"repository_key": "main",
"materializer_kind": "runtime_git_clone",
"status": "active",
"cleanliness": "clean"
"alias": "checkout",
"working_directory": {
"working_directory_id": "001a06a9f0202000000",
"display_name": "Checkout",
"repository_key": "main",
"materializer_kind": "runtime_git_clone",
"status": "active",
"cleanliness": "clean"
}
}))
.unwrap(),
);
];
let widths = WorkerColumnWidths::from_workers(std::slice::from_ref(&worker));
let text = row_text(&worker, &widths);
assert_eq!(
text,
" W-90 Coder · T-585 [stopped] wd:main・001a06a9f0202000000"
" W-90 Coder · T-585 [stopped] wd:checkout:Checkout・001a06a9f0202000000"
);
assert!(!text.contains("profile:"));
assert!(!text.contains("active clean"));
@@ -570,15 +586,18 @@ mod tests {
long.worker_state = None;
for worker in [&mut short, &mut long] {
worker.working_directory = Some(
worker.workdir_attachments = vec![
serde_json::from_value(serde_json::json!({
"working_directory_id": "workdir-1",
"repository_key": "main",
"materializer_kind": "runtime_git_clone",
"status": "active"
"alias": "checkout",
"working_directory": {
"working_directory_id": "workdir-1",
"repository_key": "main",
"materializer_kind": "runtime_git_clone",
"status": "active"
}
}))
.unwrap(),
);
];
}
let workers = vec![short, long];
@@ -595,8 +614,8 @@ mod tests {
display_column(&second, "[stopped]")
);
assert_eq!(
display_column(&first, "wd:main"),
display_column(&second, "wd:main")
display_column(&first, "wd:checkout:main"),
display_column(&second, "wd:checkout:main")
);
}
+13 -2
View File
@@ -8,6 +8,7 @@
pub mod http;
mod local;
mod operation;
mod router;
mod scope;
pub mod workspace;
@@ -27,10 +28,14 @@ pub use local::{
LocalWorkdirSession, SymlinkInfo, WorkdirSessionResource, direct_symlink, first_symlink,
};
pub use operation::*;
pub use router::{
ResolvedWorkdirSession, RoutedWorkdirSession, WorkdirAttachmentAlias, WorkdirRouteError,
WorkdirRouteErrorCode, WorkdirSessionRouter,
};
pub use scope::{
ReadOnlyWorkdirSession, WorkdirScopeAuthorizationRequest, WorkdirScopeLease,
WorkdirScopeOverlapRequest, WorkdirToolBroker, WorkdirToolScope, WorkdirToolScopePermission,
WorkdirToolScopeRule,
WorkdirScopeLeaseSet, WorkdirScopeOverlapRequest, WorkdirToolBroker, WorkdirToolBrokerRouter,
WorkdirToolScope, WorkdirToolScopePermission, WorkdirToolScopeRule,
};
/// Persistent, opaque identity of one materialized Workdir.
@@ -122,6 +127,12 @@ impl WorkdirSessionCapabilities {
bits: Self::READ | Self::GLOB | Self::GREP,
};
pub const fn union(self, other: Self) -> Self {
Self {
bits: self.bits | other.bits,
}
}
pub const fn supports(self, capability: WorkdirSessionCapability) -> bool {
let bit = match capability {
WorkdirSessionCapability::Read => Self::READ,
+796
View File
@@ -0,0 +1,796 @@
//! Alias-keyed Worker Workdir session routing.
//!
//! The router owns live provider sessions without exposing whether a session is
//! local or transported over HTTP. Attachment aliases are Worker-local stable
//! identities; provider Workdir ids and human display names are deliberately not
//! routing keys.
use std::collections::{BTreeMap, HashSet};
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, RwLock};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use tokio::sync::broadcast;
use crate::{
CommandEvent, CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest,
CommandSnapshot, CommandStatus, EditRequest, EditResult, GlobRequest, GlobResult, GrepRequest,
GrepResult, ListRequest, ListResult, ReadRequest, ReadResult, StatRequest, StatResult, Workdir,
WorkdirError, WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, WorkdirSession,
WorkdirSessionCapabilities, WorkdirSessionHandle, WriteRequest, WriteResult,
};
/// Stable Worker-local routing key for one attached Workdir.
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize)]
#[serde(transparent)]
pub struct WorkdirAttachmentAlias(String);
impl WorkdirAttachmentAlias {
pub fn new(value: impl Into<String>) -> Result<Self, WorkdirError> {
let value = value.into();
let valid = !value.is_empty()
&& value.len() <= 64
&& value.bytes().enumerate().all(|(index, byte)| match byte {
b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9' => true,
b'_' | b'-' | b'.' => index > 0,
_ => false,
});
if !valid {
return Err(WorkdirError::InvalidArgument(
"attachment alias must be 1-64 ASCII letters, digits, '.', '_' or '-', and start with a letter or digit"
.to_string(),
));
}
Ok(Self(value))
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl<'de> Deserialize<'de> for WorkdirAttachmentAlias {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let value = String::deserialize(deserializer)?;
Self::new(value).map_err(serde::de::Error::custom)
}
}
impl std::fmt::Display for WorkdirAttachmentAlias {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str(&self.0)
}
}
/// Stable machine-readable reason why a tool target could not be resolved.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum WorkdirRouteErrorCode {
NoWorkdirAttached,
TargetWorkdirRequired,
UnknownTargetWorkdir,
}
/// Provider-neutral attachment routing failure returned before an operation starts.
#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub struct WorkdirRouteError {
pub code: WorkdirRouteErrorCode,
#[serde(skip_serializing_if = "Option::is_none")]
pub target_workdir: Option<String>,
pub available_aliases: Vec<String>,
}
impl std::fmt::Display for WorkdirRouteError {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let encoded = serde_json::to_string(self).map_err(|_| std::fmt::Error)?;
formatter.write_str(&encoded)
}
}
impl std::error::Error for WorkdirRouteError {}
/// One attachment selected for a single tool invocation.
#[derive(Clone)]
pub struct ResolvedWorkdirSession {
pub alias: WorkdirAttachmentAlias,
pub generation: u64,
pub session: WorkdirSessionHandle,
}
impl std::fmt::Debug for ResolvedWorkdirSession {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("ResolvedWorkdirSession")
.field("alias", &self.alias)
.field("generation", &self.generation)
.field("workdir", self.session.workdir())
.finish_non_exhaustive()
}
}
#[derive(Debug)]
struct RoutedSessionState {
session: WorkdirSessionHandle,
generation: u64,
detached: AtomicBool,
active_operations: AtomicUsize,
active_commands: Mutex<HashSet<CommandHandle>>,
}
impl RoutedSessionState {
fn enter(&self) -> Result<ActiveOperation<'_>, WorkdirError> {
if self.detached.load(Ordering::Acquire) {
return Err(WorkdirError::Unavailable(
"Workdir attachment is detached".to_string(),
));
}
self.active_operations.fetch_add(1, Ordering::AcqRel);
if self.detached.load(Ordering::Acquire) {
self.active_operations.fetch_sub(1, Ordering::AcqRel);
return Err(WorkdirError::Unavailable(
"Workdir attachment is detaching".to_string(),
));
}
Ok(ActiveOperation(self))
}
fn is_busy(&self) -> bool {
if self.active_operations.load(Ordering::Acquire) != 0 {
return true;
}
self.reconcile_terminal_commands(&self.session.command_snapshot());
!self
.active_commands
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.is_empty()
}
async fn is_busy_for_detach(&self) -> bool {
if self.active_operations.load(Ordering::Acquire) != 0 {
return true;
}
self.reconcile_terminal_commands(&self.session.command_snapshot());
let handles = self
.active_commands
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.iter()
.cloned()
.collect::<Vec<_>>();
for handle in handles {
match self.session.command_status(handle.clone()).await {
Ok(status) => self.mark_terminal(&handle, status),
Err(WorkdirError::UnknownCommand(_)) => {
self.active_commands
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.remove(&handle);
}
Err(_) => {}
}
}
self.active_operations.load(Ordering::Acquire) != 0
|| !self
.active_commands
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.is_empty()
}
fn reconcile_terminal_commands(&self, snapshots: &[CommandSnapshot]) {
let terminal = snapshots
.iter()
.filter(|snapshot| snapshot.status != CommandStatus::Running)
.map(|snapshot| snapshot.command_id.as_str())
.collect::<HashSet<_>>();
if terminal.is_empty() {
return;
}
self.active_commands
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.retain(|handle| !terminal.contains(handle.0.as_str()));
}
fn mark_terminal(&self, handle: &CommandHandle, status: CommandStatus) {
if status != CommandStatus::Running {
self.active_commands
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.remove(handle);
}
}
}
struct ActiveOperation<'a>(&'a RoutedSessionState);
impl Drop for ActiveOperation<'_> {
fn drop(&mut self) {
self.0.active_operations.fetch_sub(1, Ordering::AcqRel);
}
}
/// A session handle bound to one alias in a [`WorkdirSessionRouter`].
#[derive(Debug, Clone)]
pub struct RoutedWorkdirSession {
state: Arc<RoutedSessionState>,
}
#[async_trait]
impl WorkdirSession for RoutedWorkdirSession {
fn workdir(&self) -> &Workdir {
self.state.session.workdir()
}
fn capabilities(&self) -> WorkdirSessionCapabilities {
self.state.session.capabilities()
}
async fn authorize_scope_path(
&self,
request: WorkdirScopeAuthorizationRequest,
) -> Result<(), WorkdirError> {
let _active = self.state.enter()?;
self.state.session.authorize_scope_path(request).await
}
async fn scope_rules_overlap(
&self,
request: WorkdirScopeOverlapRequest,
) -> Result<bool, WorkdirError> {
let _active = self.state.enter()?;
self.state.session.scope_rules_overlap(request).await
}
async fn stat(&self, request: StatRequest) -> Result<StatResult, WorkdirError> {
let _active = self.state.enter()?;
self.state.session.stat(request).await
}
async fn read(&self, request: ReadRequest) -> Result<ReadResult, WorkdirError> {
let _active = self.state.enter()?;
self.state.session.read(request).await
}
async fn write(&self, request: WriteRequest) -> Result<WriteResult, WorkdirError> {
let _active = self.state.enter()?;
self.state.session.write(request).await
}
async fn edit(&self, request: EditRequest) -> Result<EditResult, WorkdirError> {
let _active = self.state.enter()?;
self.state.session.edit(request).await
}
async fn list(&self, request: ListRequest) -> Result<ListResult, WorkdirError> {
let _active = self.state.enter()?;
self.state.session.list(request).await
}
async fn glob(&self, request: GlobRequest) -> Result<GlobResult, WorkdirError> {
let _active = self.state.enter()?;
self.state.session.glob(request).await
}
async fn grep(&self, request: GrepRequest) -> Result<GrepResult, WorkdirError> {
let _active = self.state.enter()?;
self.state.session.grep(request).await
}
async fn start_command(&self, request: CommandRequest) -> Result<CommandHandle, WorkdirError> {
let _active = self.state.enter()?;
let handle = self.state.session.start_command(request).await?;
self.state
.active_commands
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(handle.clone());
Ok(handle)
}
async fn command_status(&self, handle: CommandHandle) -> Result<CommandStatus, WorkdirError> {
let _active = self.state.enter()?;
let status = self.state.session.command_status(handle.clone()).await?;
self.state.mark_terminal(&handle, status);
Ok(status)
}
async fn command_output(
&self,
request: CommandOutputRequest,
) -> Result<CommandOutput, WorkdirError> {
let _active = self.state.enter()?;
let handle = request.handle.clone();
let output = self.state.session.command_output(request).await?;
self.state.mark_terminal(&handle, output.status);
Ok(output)
}
async fn cancel_command(&self, handle: CommandHandle) -> Result<(), WorkdirError> {
let _active = self.state.enter()?;
self.state.session.cancel_command(handle.clone()).await?;
self.state
.active_commands
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.remove(&handle);
Ok(())
}
fn subscribe_command_events(&self) -> Option<broadcast::Receiver<CommandEvent>> {
self.state.session.subscribe_command_events()
}
fn command_snapshot(&self) -> Vec<CommandSnapshot> {
let snapshots = self.state.session.command_snapshot();
self.state.reconcile_terminal_commands(&snapshots);
snapshots
}
async fn close(&self) -> Result<(), WorkdirError> {
if self.state.is_busy() {
return Err(WorkdirError::Conflict(
"Workdir attachment has active operations or commands".to_string(),
));
}
self.state.session.close().await
}
}
/// Worker-owned map of stable attachment aliases to provider-neutral sessions.
#[derive(Debug, Default)]
pub struct WorkdirSessionRouter {
sessions: RwLock<BTreeMap<WorkdirAttachmentAlias, Arc<RoutedSessionState>>>,
next_generation: AtomicU64,
}
impl WorkdirSessionRouter {
pub fn new() -> Self {
Self::default()
}
pub fn aliases(&self) -> Vec<WorkdirAttachmentAlias> {
self.sessions
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.keys()
.cloned()
.collect()
}
pub fn capabilities(&self) -> WorkdirSessionCapabilities {
self.sessions
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.values()
.fold(WorkdirSessionCapabilities::EMPTY, |capabilities, state| {
capabilities.union(state.session.capabilities())
})
}
pub fn len(&self) -> usize {
self.sessions
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
/// Resolve one model-visible target without exposing provider identity.
///
/// Omission is accepted only for the unambiguous single-attachment case.
/// Every error includes the current alias snapshot so callers can recover
/// without consulting Workdir ids, display names, or transport metadata.
pub fn resolve(
&self,
target_workdir: Option<&str>,
) -> Result<ResolvedWorkdirSession, WorkdirRouteError> {
let sessions = self
.sessions
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let available_aliases = sessions
.keys()
.map(|alias| alias.as_str().to_owned())
.collect::<Vec<_>>();
let selected = match target_workdir {
Some(target) => sessions
.iter()
.find(|(alias, _)| alias.as_str() == target)
.map(|(alias, state)| (alias.clone(), state.clone()))
.ok_or_else(|| WorkdirRouteError {
code: WorkdirRouteErrorCode::UnknownTargetWorkdir,
target_workdir: Some(target.to_owned()),
available_aliases: available_aliases.clone(),
})?,
None if sessions.is_empty() => {
return Err(WorkdirRouteError {
code: WorkdirRouteErrorCode::NoWorkdirAttached,
target_workdir: None,
available_aliases,
});
}
None if sessions.len() > 1 => {
return Err(WorkdirRouteError {
code: WorkdirRouteErrorCode::TargetWorkdirRequired,
target_workdir: None,
available_aliases,
});
}
None => {
let (alias, state) = sessions
.first_key_value()
.expect("single attachment exists after cardinality checks");
(alias.clone(), state.clone())
}
};
Ok(ResolvedWorkdirSession {
alias: selected.0,
generation: selected.1.generation,
session: Arc::new(RoutedWorkdirSession { state: selected.1 }),
})
}
pub fn attach(
&self,
alias: WorkdirAttachmentAlias,
session: WorkdirSessionHandle,
) -> Result<WorkdirSessionHandle, WorkdirError> {
let mut sessions = self.sessions.write().map_err(|_| {
WorkdirError::Unavailable("Workdir attachment router is poisoned".to_string())
})?;
if sessions.contains_key(&alias) {
return Err(WorkdirError::Conflict(format!(
"Workdir attachment alias `{alias}` already exists"
)));
}
let state = Arc::new(RoutedSessionState {
session,
generation: self.next_generation.fetch_add(1, Ordering::AcqRel),
detached: AtomicBool::new(false),
active_operations: AtomicUsize::new(0),
active_commands: Mutex::new(HashSet::new()),
});
sessions.insert(alias, state.clone());
Ok(Arc::new(RoutedWorkdirSession { state }))
}
pub fn session(&self, alias: &WorkdirAttachmentAlias) -> Option<WorkdirSessionHandle> {
self.sessions
.read()
.ok()?
.get(alias)
.cloned()
.map(|state| Arc::new(RoutedWorkdirSession { state }) as WorkdirSessionHandle)
}
/// Compatibility projection used only while filesystem tool schemas have no
/// `target_workdir`: exactly one attachment is unambiguous; zero or multiple
/// attachments expose no implicit primary.
pub fn only_session(&self) -> Option<WorkdirSessionHandle> {
let sessions = self.sessions.read().ok()?;
if sessions.len() != 1 {
return None;
}
sessions
.values()
.next()
.cloned()
.map(|state| Arc::new(RoutedWorkdirSession { state }) as WorkdirSessionHandle)
}
pub fn can_detach(&self, alias: &WorkdirAttachmentAlias) -> Result<(), WorkdirError> {
let state = self
.sessions
.read()
.map_err(|_| {
WorkdirError::Unavailable("Workdir attachment router is poisoned".to_string())
})?
.get(alias)
.cloned()
.ok_or_else(|| WorkdirError::NotFound(alias.as_str().into()))?;
if state.detached.load(Ordering::Acquire) {
return Err(WorkdirError::Unavailable(
"Workdir attachment is already detaching".to_string(),
));
}
if state.is_busy() {
return Err(WorkdirError::Conflict(format!(
"Workdir attachment `{alias}` has active operations or commands"
)));
}
Ok(())
}
/// Atomically fence new operations before an external detach transaction begins.
pub async fn begin_detach(&self, alias: &WorkdirAttachmentAlias) -> Result<(), WorkdirError> {
let state = self
.sessions
.read()
.map_err(|_| {
WorkdirError::Unavailable("Workdir attachment router is poisoned".to_string())
})?
.get(alias)
.cloned()
.ok_or_else(|| WorkdirError::NotFound(alias.as_str().into()))?;
if state.detached.swap(true, Ordering::AcqRel) {
return Err(WorkdirError::Unavailable(
"Workdir attachment is already detaching".to_string(),
));
}
if state.is_busy_for_detach().await {
state.detached.store(false, Ordering::Release);
return Err(WorkdirError::Conflict(format!(
"Workdir attachment `{alias}` has active operations or commands"
)));
}
Ok(())
}
/// Reopen a route after the external detach transaction failed before commit.
pub fn cancel_detach(&self, alias: &WorkdirAttachmentAlias) {
if let Ok(sessions) = self.sessions.read()
&& let Some(state) = sessions.get(alias)
{
state.detached.store(false, Ordering::Release);
}
}
pub async fn finish_detach(
&self,
alias: &WorkdirAttachmentAlias,
) -> Result<WorkdirSessionHandle, WorkdirError> {
let state = {
let sessions = self.sessions.read().map_err(|_| {
WorkdirError::Unavailable("Workdir attachment router is poisoned".to_string())
})?;
sessions
.get(alias)
.cloned()
.ok_or_else(|| WorkdirError::NotFound(alias.as_str().into()))?
};
if !state.detached.load(Ordering::Acquire) {
return Err(WorkdirError::Conflict(format!(
"Workdir attachment `{alias}` was not prepared for detach"
)));
}
if state.is_busy_for_detach().await {
state.detached.store(false, Ordering::Release);
return Err(WorkdirError::Conflict(format!(
"Workdir attachment `{alias}` has active operations or commands"
)));
}
if let Err(error) = state.session.close().await {
state.detached.store(false, Ordering::Release);
return Err(error);
}
let mut sessions = self.sessions.write().map_err(|_| {
WorkdirError::Unavailable("Workdir attachment router is poisoned".to_string())
})?;
if sessions
.get(alias)
.is_some_and(|current| Arc::ptr_eq(current, &state))
{
sessions.remove(alias);
}
Ok(state.session.clone())
}
pub async fn detach(
&self,
alias: &WorkdirAttachmentAlias,
) -> Result<WorkdirSessionHandle, WorkdirError> {
self.begin_detach(alias).await?;
self.finish_detach(alias).await
}
/// Close every attachment during Worker shutdown or startup rollback.
///
/// Unlike [`Self::detach`], lifecycle teardown is authoritative and delegates
/// cancellation or draining of active commands to each provider session.
pub async fn close_all(&self) -> Result<(), WorkdirError> {
let aliases = self.aliases();
for alias in aliases {
let state = {
let sessions = self.sessions.read().map_err(|_| {
WorkdirError::Unavailable("Workdir attachment router is poisoned".to_string())
})?;
sessions.get(&alias).cloned()
};
let Some(state) = state else {
continue;
};
state.detached.store(true, Ordering::Release);
if let Err(error) = state.session.close().await {
state.detached.store(false, Ordering::Release);
return Err(error);
}
let mut sessions = self.sessions.write().map_err(|_| {
WorkdirError::Unavailable("Workdir attachment router is poisoned".to_string())
})?;
if sessions
.get(&alias)
.is_some_and(|current| Arc::ptr_eq(current, &state))
{
sessions.remove(&alias);
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::LocalWorkdirSession;
use manifest::Scope;
use tempfile::TempDir;
fn session(id: &str, dir: &TempDir) -> WorkdirSessionHandle {
Arc::new(LocalWorkdirSession::materialized_bound(
Workdir::new(id),
dir.path().to_path_buf(),
dir.path().to_path_buf(),
manifest::SharedScope::new(Scope::writable(dir.path()).unwrap()),
WorkdirSessionCapabilities::ALL,
))
}
#[tokio::test]
async fn router_holds_multiple_aliases_without_implicit_primary() {
let left = TempDir::new().unwrap();
let right = TempDir::new().unwrap();
let router = WorkdirSessionRouter::new();
router
.attach(
WorkdirAttachmentAlias::new("left").unwrap(),
session("wd-left", &left),
)
.unwrap();
router
.attach(
WorkdirAttachmentAlias::new("right").unwrap(),
session("wd-right", &right),
)
.unwrap();
assert_eq!(
router
.aliases()
.iter()
.map(WorkdirAttachmentAlias::as_str)
.collect::<Vec<_>>(),
["left", "right"]
);
assert!(router.only_session().is_none());
let missing = router.resolve(None).unwrap_err();
assert_eq!(missing.code, WorkdirRouteErrorCode::TargetWorkdirRequired);
assert_eq!(missing.available_aliases, ["left", "right"]);
assert_eq!(
router
.resolve(Some("right"))
.unwrap()
.session
.workdir()
.id()
.as_str(),
"wd-right"
);
let unknown = router.resolve(Some("display-name")).unwrap_err();
assert_eq!(unknown.code, WorkdirRouteErrorCode::UnknownTargetWorkdir);
assert_eq!(unknown.target_workdir.as_deref(), Some("display-name"));
assert_eq!(unknown.available_aliases, ["left", "right"]);
router.close_all().await.unwrap();
}
#[test]
fn empty_router_returns_structured_no_attachment_error() {
let error = WorkdirSessionRouter::new().resolve(None).unwrap_err();
assert_eq!(error.code, WorkdirRouteErrorCode::NoWorkdirAttached);
assert!(error.available_aliases.is_empty());
assert_eq!(
serde_json::to_value(error).unwrap(),
serde_json::json!({
"code": "no_workdir_attached",
"available_aliases": []
})
);
}
#[tokio::test]
async fn detach_allows_naturally_completed_command_without_final_poll() {
let dir = TempDir::new().unwrap();
let router = WorkdirSessionRouter::new();
let alias = WorkdirAttachmentAlias::new("checkout").unwrap();
let underlying = session("wd", &dir);
let routed = router
.attach(alias.clone(), Arc::clone(&underlying))
.unwrap();
let handle = routed
.start_command(CommandRequest {
command: "printf done".to_string(),
timeout_secs: 5,
output_limit: 1024,
cwd: None,
spill_dir: None,
tool_call_id: None,
})
.await
.unwrap();
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
let completed = underlying.command_snapshot().iter().any(|snapshot| {
snapshot.command_id == handle.0 && snapshot.status != CommandStatus::Running
});
if completed {
break;
}
assert!(tokio::time::Instant::now() < deadline);
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
router.detach(&alias).await.unwrap();
assert!(router.is_empty());
}
#[tokio::test]
async fn lifecycle_close_all_closes_sessions_with_active_commands() {
let dir = TempDir::new().unwrap();
let router = WorkdirSessionRouter::new();
let alias = WorkdirAttachmentAlias::new("checkout").unwrap();
let underlying = session("wd", &dir);
let routed = router.attach(alias, Arc::clone(&underlying)).unwrap();
let handle = routed
.start_command(CommandRequest {
command: "sleep 30".to_string(),
timeout_secs: 60,
output_limit: 1024,
cwd: None,
spill_dir: None,
tool_call_id: None,
})
.await
.unwrap();
router.close_all().await.unwrap();
assert!(router.is_empty());
assert!(matches!(
underlying.command_status(handle).await,
Err(WorkdirError::Unavailable(_))
));
}
#[tokio::test]
async fn detach_rejects_active_command_until_cancelled() {
let dir = TempDir::new().unwrap();
let router = WorkdirSessionRouter::new();
let alias = WorkdirAttachmentAlias::new("checkout").unwrap();
let routed = router.attach(alias.clone(), session("wd", &dir)).unwrap();
let handle = routed
.start_command(CommandRequest {
command: "sleep 30".to_string(),
timeout_secs: 60,
output_limit: 1024,
cwd: None,
spill_dir: None,
tool_call_id: None,
})
.await
.unwrap();
assert!(matches!(
router.detach(&alias).await,
Err(WorkdirError::Conflict(_))
));
routed.cancel_command(handle).await.unwrap();
router.detach(&alias).await.unwrap();
assert!(router.is_empty());
}
}
+262 -1
View File
@@ -1,4 +1,4 @@
use std::collections::{HashMap, HashSet};
use std::collections::{BTreeMap, HashMap, HashSet};
use std::path::Path;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, Weak};
@@ -284,6 +284,204 @@ impl Drop for WorkdirScopeLease {
}
}
/// Alias-keyed parent authority used to attenuate one or more Workdir
/// attachments for an Internal SubWorker.
#[derive(Clone)]
pub struct WorkdirToolBrokerRouter {
sessions: Arc<crate::WorkdirSessionRouter>,
brokers: Arc<Mutex<BTreeMap<crate::WorkdirAttachmentAlias, (u64, WorkdirToolBroker)>>>,
}
impl std::fmt::Debug for WorkdirToolBrokerRouter {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("WorkdirToolBrokerRouter")
.field("aliases", &self.sessions.aliases())
.finish_non_exhaustive()
}
}
impl WorkdirToolBrokerRouter {
pub fn from_single(
alias: crate::WorkdirAttachmentAlias,
session: WorkdirSessionHandle,
) -> Result<Self, WorkdirError> {
let sessions = Arc::new(crate::WorkdirSessionRouter::new());
sessions.attach(alias, session)?;
Ok(Self::new(sessions))
}
pub fn new(sessions: Arc<crate::WorkdirSessionRouter>) -> Self {
Self {
sessions,
brokers: Arc::new(Mutex::new(BTreeMap::new())),
}
}
pub fn sessions(&self) -> Arc<crate::WorkdirSessionRouter> {
self.sessions.clone()
}
pub fn resolve_alias(
&self,
target_workdir: Option<&str>,
) -> Result<crate::WorkdirAttachmentAlias, crate::WorkdirRouteError> {
self.sessions
.resolve(target_workdir)
.map(|selected| selected.alias)
}
fn broker(
&self,
alias: &crate::WorkdirAttachmentAlias,
) -> Result<WorkdirToolBroker, WorkdirError> {
let selected = self
.sessions
.resolve(Some(alias.as_str()))
.map_err(|error| WorkdirError::InvalidArgument(error.to_string()))?;
let mut brokers = self.brokers.lock().map_err(|_| {
WorkdirError::Unavailable("Workdir tool broker router is poisoned".to_string())
})?;
if let Some((generation, broker)) = brokers.get(alias)
&& *generation == selected.generation
{
return Ok(broker.clone());
}
let broker = WorkdirToolBroker::new(selected.session);
brokers.insert(alias.clone(), (selected.generation, broker.clone()));
Ok(broker)
}
pub async fn scope(
&self,
targets: Vec<(crate::WorkdirAttachmentAlias, WorkdirToolScope)>,
) -> Result<WorkdirScopeLeaseSet, WorkdirError> {
if targets.is_empty() {
return Err(WorkdirError::InvalidArgument(
"at least one Workdir attachment scope is required".to_string(),
));
}
let child_sessions = Arc::new(crate::WorkdirSessionRouter::new());
let mut child_brokers = BTreeMap::new();
let mut leases = Vec::with_capacity(targets.len());
for (alias, request) in targets {
if leases.iter().any(
|(existing, _): &(crate::WorkdirAttachmentAlias, WorkdirScopeLease)| {
existing == &alias
},
) {
for (_, lease) in &leases {
let _ = lease.close().await;
}
return Err(WorkdirError::InvalidArgument(format!(
"duplicate Workdir attachment scope `{alias}`"
)));
}
let lease = match self.broker(&alias)?.scope(request).await {
Ok(lease) => lease,
Err(error) => {
for (_, lease) in &leases {
let _ = lease.close().await;
}
return Err(error);
}
};
if let Err(error) = child_sessions.attach(alias.clone(), lease.tool_session()) {
let _ = lease.close().await;
for (_, existing) in &leases {
let _ = existing.close().await;
}
return Err(error);
}
let child_generation = child_sessions
.resolve(Some(alias.as_str()))
.expect("newly attached child route resolves")
.generation;
child_brokers.insert(alias.clone(), (child_generation, lease.broker()));
leases.push((alias, lease));
}
Ok(WorkdirScopeLeaseSet {
broker_router: WorkdirToolBrokerRouter {
sessions: child_sessions,
brokers: Arc::new(Mutex::new(child_brokers)),
},
leases,
})
}
}
pub struct WorkdirScopeLeaseSet {
broker_router: WorkdirToolBrokerRouter,
leases: Vec<(crate::WorkdirAttachmentAlias, WorkdirScopeLease)>,
}
impl std::fmt::Debug for WorkdirScopeLeaseSet {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("WorkdirScopeLeaseSet")
.field(
"aliases",
&self
.leases
.iter()
.map(|(alias, _)| alias)
.collect::<Vec<_>>(),
)
.field("active", &self.is_active())
.finish()
}
}
impl WorkdirScopeLeaseSet {
pub fn from_single(
alias: crate::WorkdirAttachmentAlias,
lease: WorkdirScopeLease,
) -> Result<Self, WorkdirError> {
let sessions = Arc::new(crate::WorkdirSessionRouter::new());
sessions.attach(alias.clone(), lease.tool_session())?;
let generation = sessions
.resolve(Some(alias.as_str()))
.expect("newly attached single route resolves")
.generation;
let brokers = BTreeMap::from([(alias.clone(), (generation, lease.broker()))]);
Ok(Self {
broker_router: WorkdirToolBrokerRouter {
sessions,
brokers: Arc::new(Mutex::new(brokers)),
},
leases: vec![(alias, lease)],
})
}
pub fn broker_router(&self) -> WorkdirToolBrokerRouter {
self.broker_router.clone()
}
pub fn session_router(&self) -> Arc<crate::WorkdirSessionRouter> {
self.broker_router.sessions()
}
pub async fn close(&self) -> Result<(), WorkdirError> {
let mut first_error = None;
for (_, lease) in &self.leases {
if let Err(error) = lease.close().await {
first_error.get_or_insert(error);
}
}
first_error.map_or(Ok(()), Err)
}
pub fn is_active(&self) -> bool {
self.leases.iter().all(|(_, lease)| lease.is_active())
}
pub fn revoke(&self) {
for (_, lease) in &self.leases {
lease.revoke();
}
}
}
#[derive(Debug)]
struct SessionValidity {
active: AtomicBool,
@@ -2303,6 +2501,69 @@ mod tests {
child.close().await.unwrap();
}
#[tokio::test]
async fn broker_router_delegates_multiple_aliases_without_namespace_mixing() {
let left = TempDir::new().unwrap();
let right = TempDir::new().unwrap();
fs::write(left.path().join("same.txt"), "left").unwrap();
fs::write(right.path().join("same.txt"), "right").unwrap();
let sessions = Arc::new(crate::WorkdirSessionRouter::new());
sessions
.attach(
crate::WorkdirAttachmentAlias::new("left").unwrap(),
Arc::new(crate::LocalWorkdirSession::new(
manifest::Scope::writable(left.path()).unwrap(),
left.path().to_path_buf(),
)),
)
.unwrap();
sessions
.attach(
crate::WorkdirAttachmentAlias::new("right").unwrap(),
Arc::new(crate::LocalWorkdirSession::new(
manifest::Scope::writable(right.path()).unwrap(),
right.path().to_path_buf(),
)),
)
.unwrap();
let brokers = WorkdirToolBrokerRouter::new(sessions);
let leases = brokers
.scope(vec![
(
crate::WorkdirAttachmentAlias::new("left").unwrap(),
request("", WorkdirToolScopePermission::Read),
),
(
crate::WorkdirAttachmentAlias::new("right").unwrap(),
request("", WorkdirToolScopePermission::Read),
),
])
.await
.unwrap();
let child = leases.session_router();
let left_content = child
.resolve(Some("left"))
.unwrap()
.session
.read(read("same.txt"))
.await
.unwrap();
let right_content = child
.resolve(Some("right"))
.unwrap()
.session
.read(read("same.txt"))
.await
.unwrap();
assert_eq!(left_content.bytes, b"left");
assert_eq!(right_content.bytes, b"right");
assert!(child.resolve(None).is_err());
leases.close().await.unwrap();
assert!(!leases.is_active());
}
#[tokio::test]
async fn closing_parent_invalidates_scoped_tools() {
let root = TempDir::new().unwrap();
+9 -7
View File
@@ -70,6 +70,9 @@ pub struct RuntimeWorkingDirectoryCleanupTarget {
#[serde(deny_unknown_fields)]
pub struct RuntimeWorkingDirectorySummary {
pub working_directory_id: String,
/// Optional human-facing label. Never used for routing or identity.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub display_name: Option<String>,
pub repository_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub creation_selector: Option<String>,
@@ -92,8 +95,6 @@ pub struct RuntimeWorkingDirectorySummary {
#[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>,
}
@@ -101,6 +102,9 @@ pub struct RuntimeWorkingDirectorySummary {
#[serde(deny_unknown_fields)]
pub struct WorkingDirectorySummary {
pub working_directory_id: String,
/// Optional human-facing label. Never used for attachment routing.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub display_name: Option<String>,
pub repository_key: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub creation_selector: Option<String>,
@@ -123,8 +127,6 @@ pub struct WorkingDirectorySummary {
#[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>,
}
@@ -184,8 +186,6 @@ pub struct WorkingDirectoryCurrentObservation {
#[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>,
}
@@ -214,6 +214,7 @@ mod tests {
fn working_directory_summary_preserves_server_projection_wire_shape() {
let summary = WorkingDirectorySummary {
working_directory_id: "wd-1".to_string(),
display_name: Some("Checkout".to_string()),
repository_key: "repo-1".to_string(),
creation_selector: None,
creation_ref: None,
@@ -226,7 +227,6 @@ mod tests {
cleanup_target: None,
status: WorkingDirectoryStatusKind::Active,
cleanliness: None,
primary_worker_id: None,
occupied_by: None,
};
let value = serde_json::to_value(summary).unwrap();
@@ -240,5 +240,7 @@ mod tests {
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct WorkspaceWorkdirSessionOperationRequest {
/// Worker-local attachment alias selected by the calling tool router.
pub target_workdir: String,
pub operation: crate::http::WorkdirSessionOperation,
}
+24 -9
View File
@@ -161,6 +161,8 @@ pub struct WorkingDirectoryRepositoryAccessRequest {
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct WorkingDirectoryRequest {
pub repository: WorkingDirectoryRepository,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub display_name: Option<String>,
#[serde(default)]
pub materializer: MaterializerKind,
/// Backend-assigned stable Workdir id. Runtimes use this when present so the
@@ -197,12 +199,25 @@ pub struct RepositoryRefObservation {
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct WorkingDirectoryClaim {
pub struct WorkingDirectoryAttachmentClaim {
pub alias: workdir::WorkdirAttachmentAlias,
pub working_directory_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub relative_cwd: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct WorkingDirectoryAttachmentRequest {
pub alias: workdir::WorkdirAttachmentAlias,
pub working_directory: WorkingDirectoryRequest,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct WorkingDirectoryAttachmentStatus {
pub alias: workdir::WorkdirAttachmentAlias,
pub working_directory: WorkingDirectoryStatus,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct WorkingDirectoryStatus {
pub summary: WorkingDirectorySummary,
@@ -247,10 +262,10 @@ pub struct CreateWorkerRequest {
pub config_bundle: Option<ConfigBundleRef>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub initial_input: Option<WorkerInput>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub working_directory_request: Option<WorkingDirectoryRequest>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub working_directory: Option<WorkingDirectoryClaim>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub workdir_attachment_requests: Vec<WorkingDirectoryAttachmentRequest>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub workdir_attachments: Vec<WorkingDirectoryAttachmentClaim>,
/// Backend-only feature enablement. Grants still define local Runtime peers;
/// the Workspace provider reauthorizes its dynamic set per operation.
#[serde(default, skip_serializing_if = "is_false")]
@@ -308,8 +323,8 @@ pub struct WorkerSummary {
pub worker_state: Option<protocol::WorkerStateSnapshot>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub workspace_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub working_directory: Option<WorkingDirectoryStatus>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub workdir_attachments: Vec<WorkingDirectoryAttachmentStatus>,
pub profile: ProfileSelector,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub display_name: Option<String>,
@@ -333,8 +348,8 @@ pub struct WorkerDetail {
pub worker_state: Option<protocol::WorkerStateSnapshot>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub workspace_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub working_directory: Option<WorkingDirectoryStatus>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub workdir_attachments: Vec<WorkingDirectoryAttachmentStatus>,
pub profile: ProfileSelector,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub display_name: Option<String>,
+8 -7
View File
@@ -12,6 +12,7 @@ use crate::observation::WorkerObservationEvent;
use crate::working_directory::{WorkingDirectoryBinding, WorkingDirectoryDiagnostic};
use protocol::{Method, UploadedFileRef};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::fmt;
use std::sync::Arc;
use workdir::WorkdirSessionHandle;
@@ -246,7 +247,7 @@ pub struct WorkerExecutionSpawnRequest {
pub request: crate::catalog::CreateWorkerRequest,
pub workspace_scope: Option<crate::runtime::RuntimeWorkspaceScope>,
pub context: WorkerExecutionContext,
pub working_directory: Option<WorkingDirectoryBinding>,
pub workdir_attachments: BTreeMap<workdir::WorkdirAttachmentAlias, WorkingDirectoryBinding>,
pub config_bundle: Option<ConfigBundle>,
}
@@ -257,8 +258,8 @@ pub struct WorkerExecutionRestoreRequest {
pub request: crate::catalog::CreateWorkerRequest,
pub workspace_scope: Option<crate::runtime::RuntimeWorkspaceScope>,
pub context: WorkerExecutionContext,
pub previous_working_directory: Option<WorkingDirectoryStatus>,
pub working_directory: Option<WorkingDirectoryBinding>,
pub previous_workdir_attachments: Vec<crate::catalog::WorkingDirectoryAttachmentStatus>,
pub workdir_attachments: BTreeMap<workdir::WorkdirAttachmentAlias, WorkingDirectoryBinding>,
pub config_bundle: Option<ConfigBundle>,
}
@@ -284,7 +285,7 @@ pub enum WorkerExecutionSpawnResult {
Connected {
handle: WorkerExecutionHandle,
worker_state: protocol::WorkerStateSnapshot,
working_directory: Option<WorkingDirectoryStatus>,
workdir_attachments: Vec<crate::catalog::WorkingDirectoryAttachmentStatus>,
},
/// A read-only preflight rejected the operation before live work started.
Rejected(WorkerExecutionResult),
@@ -297,7 +298,7 @@ pub enum WorkerExecutionSpawnResult {
result: WorkerExecutionResult,
handle: Option<WorkerExecutionHandle>,
worker_state: Option<protocol::WorkerStateSnapshot>,
working_directory: Option<WorkingDirectoryStatus>,
workdir_attachments: Vec<crate::catalog::WorkingDirectoryAttachmentStatus>,
},
/// Legacy spawn failure. Restore implementations must use `RolledBack` or
/// `ReconciliationRequired` after crossing their live side-effect boundary.
@@ -308,12 +309,12 @@ impl WorkerExecutionSpawnResult {
pub fn connected(
handle: WorkerExecutionHandle,
worker_state: protocol::WorkerStateSnapshot,
working_directory: Option<WorkingDirectoryStatus>,
workdir_attachments: Vec<crate::catalog::WorkingDirectoryAttachmentStatus>,
) -> Self {
Self::Connected {
handle,
worker_state,
working_directory,
workdir_attachments,
}
}
}
+161 -7
View File
@@ -1,6 +1,7 @@
use crate::catalog::{
ConfigBundleRef, CreateWorkerRequest, ProfileSelector, WorkerRestoreIntent, WorkerStatus,
WorkingDirectoryStatus,
WorkingDirectoryAttachmentClaim, WorkingDirectoryAttachmentRequest,
WorkingDirectoryAttachmentStatus, WorkingDirectoryRequest,
};
use crate::config_bundle::ConfigBundle;
use crate::diagnostics::{DiagnosticSeverity, RuntimeDiagnostic};
@@ -442,7 +443,7 @@ pub(crate) struct PersistedWorkerRecord {
pub(crate) status: WorkerStatus,
pub(crate) execution_state: PersistedWorkerExecutionState,
pub(crate) workspace_id: Option<String>,
pub(crate) working_directory: Option<WorkingDirectoryStatus>,
pub(crate) workdir_attachments: Vec<WorkingDirectoryAttachmentStatus>,
}
fn runtime_io_error(operation: &'static str, path: &Path, source: std::io::Error) -> RuntimeError {
@@ -1685,11 +1686,18 @@ struct WorkerIdentityRecord {
status: WorkerStatus,
#[serde(default, skip_serializing_if = "Option::is_none")]
workspace_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
working_directory: Option<WorkingDirectoryStatus>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
workdir_attachments: Vec<WorkingDirectoryAttachmentStatus>,
/// Schema-v8 compatibility input. New records serialize only `workdir_attachments`.
#[serde(
default,
rename = "working_directory",
skip_serializing_if = "Option::is_none"
)]
legacy_working_directory: Option<crate::catalog::WorkingDirectoryStatus>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[derive(Clone, Debug, Serialize)]
#[serde(deny_unknown_fields)]
struct WorkerExecutionRecord {
schema_version: u32,
@@ -1699,6 +1707,72 @@ struct WorkerExecutionRecord {
restore_intent: WorkerRestoreIntent,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct RawWorkerExecutionRecord {
schema_version: u32,
request: serde_json::Value,
#[serde(default)]
binding: Option<PersistedWorkerExecutionBinding>,
restore_intent: WorkerRestoreIntent,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct LegacyWorkingDirectoryClaim {
working_directory_id: String,
#[serde(default)]
relative_cwd: Option<String>,
}
impl<'de> Deserialize<'de> for WorkerExecutionRecord {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let raw = RawWorkerExecutionRecord::deserialize(deserializer)?;
let mut request_value = raw.request;
let request_object = request_value
.as_object_mut()
.ok_or_else(|| serde::de::Error::custom("Worker create request must be an object"))?;
let legacy_request = request_object.remove("working_directory_request");
let legacy_claim = request_object.remove("working_directory");
let mut request: CreateWorkerRequest =
serde_json::from_value(request_value).map_err(serde::de::Error::custom)?;
if request.workdir_attachment_requests.is_empty() && request.workdir_attachments.is_empty()
{
let alias = workdir::WorkdirAttachmentAlias::new("workdir")
.expect("static compatibility alias is valid");
if let Some(value) = legacy_request.filter(|value| !value.is_null()) {
let working_directory: WorkingDirectoryRequest =
serde_json::from_value(value).map_err(serde::de::Error::custom)?;
request
.workdir_attachment_requests
.push(WorkingDirectoryAttachmentRequest {
alias,
working_directory,
});
} else if let Some(value) = legacy_claim.filter(|value| !value.is_null()) {
let claim: LegacyWorkingDirectoryClaim =
serde_json::from_value(value).map_err(serde::de::Error::custom)?;
request
.workdir_attachments
.push(WorkingDirectoryAttachmentClaim {
alias,
working_directory_id: claim.working_directory_id,
relative_cwd: claim.relative_cwd,
});
}
}
Ok(Self {
schema_version: raw.schema_version,
request,
binding: raw.binding,
restore_intent: raw.restore_intent,
})
}
}
impl WorkerIdentityRecord {
fn from_persisted(worker: &PersistedWorkerRecord) -> Self {
Self {
@@ -1712,7 +1786,8 @@ impl WorkerIdentityRecord {
created_at_ms: worker.created_at_ms,
status: worker.status,
workspace_id: worker.workspace_id.clone(),
working_directory: worker.working_directory.clone(),
workdir_attachments: worker.workdir_attachments.clone(),
legacy_working_directory: None,
}
}
@@ -1760,6 +1835,19 @@ impl WorkerIdentityRecord {
self,
execution_state: PersistedWorkerExecutionState,
) -> PersistedWorkerRecord {
let workdir_attachments = if self.workdir_attachments.is_empty() {
self.legacy_working_directory
.map(|working_directory| {
vec![WorkingDirectoryAttachmentStatus {
alias: workdir::WorkdirAttachmentAlias::new("workdir")
.expect("static compatibility alias is valid"),
working_directory,
}]
})
.unwrap_or_default()
} else {
self.workdir_attachments
};
PersistedWorkerRecord {
worker_ref: self.worker_ref,
worker_id: self.worker_id,
@@ -1771,7 +1859,7 @@ impl WorkerIdentityRecord {
status: self.status,
execution_state,
workspace_id: self.workspace_id,
working_directory: self.working_directory,
workdir_attachments,
}
}
}
@@ -2215,6 +2303,72 @@ mod tests {
document
}
#[test]
fn schema_v8_singular_workdir_fields_normalize_to_one_attachment() {
let worker_id = WorkerId::now_v7();
let workdir_status = serde_json::json!({
"summary": {
"working_directory_id": "wd-legacy",
"repository_id": "repo",
"materializer_kind": "runtime_git_clone",
"status": "active"
}
});
let identity: WorkerIdentityRecord = serde_json::from_value(serde_json::json!({
"schema_version": SCHEMA_VERSION,
"worker_ref": { "worker_id": worker_id },
"worker_id": worker_id,
"profile": { "kind": "builtin", "value": "builtin:coder" },
"display_name": null,
"profile_source": {
"id": "archive-1",
"digest": "sha256:archive",
"size_bytes": 0,
"source_graph": {
"source_count": 0,
"total_source_bytes": 0,
"entrypoints": {},
"import_count": 0
}
},
"config_bundle": null,
"status": "stopped",
"working_directory": workdir_status
}))
.unwrap();
let persisted = identity.into_persisted(PersistedWorkerExecutionState::Unavailable);
assert_eq!(persisted.workdir_attachments.len(), 1);
assert_eq!(persisted.workdir_attachments[0].alias.as_str(), "workdir");
assert_eq!(
persisted.workdir_attachments[0]
.working_directory
.summary
.working_directory_id,
"wd-legacy"
);
let mut request = schema_v7_worker_document(worker_id)["request"].clone();
request["working_directory"] = serde_json::json!({
"working_directory_id": "wd-legacy",
"relative_cwd": "crates/yoi"
});
let execution: WorkerExecutionRecord = serde_json::from_value(serde_json::json!({
"schema_version": SCHEMA_VERSION,
"request": request,
"binding": null,
"restore_intent": "explicit"
}))
.unwrap();
assert!(execution.request.workdir_attachment_requests.is_empty());
assert_eq!(execution.request.workdir_attachments.len(), 1);
let attachment = &execution.request.workdir_attachments[0];
assert_eq!(attachment.alias.as_str(), "workdir");
assert_eq!(attachment.working_directory_id, "wd-legacy");
assert_eq!(attachment.relative_cwd.as_deref(), Some("crates/yoi"));
let serialized = serde_json::to_value(execution).unwrap();
assert!(serialized["request"].get("working_directory").is_none());
}
#[test]
fn startup_recovers_staging_created_before_source_backup() {
let temp = tempfile::tempdir().unwrap();
+30 -14
View File
@@ -2282,7 +2282,11 @@ 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") || path.contains("/attachments") {
if path.ends_with("/input")
|| path.ends_with("/restore")
|| path.contains("/attachments")
|| path.contains("/workdir-attachments")
{
return Some("workers:input");
}
if path.ends_with("/stop") || path.ends_with("/cancel") {
@@ -3072,8 +3076,8 @@ mod tests {
digest: bundle.metadata.digest,
}),
initial_input: None,
working_directory_request: None,
working_directory: None,
workdir_attachment_requests: Vec::new(),
workdir_attachments: Vec::new(),
worker_observation_enabled: false,
worker_observation_grants: Vec::new(),
workspace_api: None,
@@ -3319,10 +3323,16 @@ mod tests {
worker_state: protocol::WorkerStateSnapshot {
..protocol::WorkerStatus::Idle.into()
},
working_directory: request
.working_directory
.as_ref()
.map(|binding| binding.status()),
workdir_attachments: request
.workdir_attachments
.iter()
.map(
|(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus {
alias: alias.clone(),
working_directory: binding.status(),
},
)
.collect(),
}
}
@@ -3335,7 +3345,7 @@ mod tests {
worker_state: protocol::WorkerStateSnapshot {
..protocol::WorkerStatus::Idle.into()
},
working_directory: request.previous_working_directory,
workdir_attachments: request.previous_workdir_attachments,
}
}
@@ -3907,10 +3917,16 @@ mod ws_tests {
worker_state: protocol::WorkerStateSnapshot {
..protocol::WorkerStatus::Idle.into()
},
working_directory: request
.working_directory
.as_ref()
.map(|binding| binding.status()),
workdir_attachments: request
.workdir_attachments
.iter()
.map(
|(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus {
alias: alias.clone(),
working_directory: binding.status(),
},
)
.collect(),
}
}
@@ -4020,8 +4036,8 @@ mod ws_tests {
digest: bundle.metadata.digest,
}),
initial_input: None,
working_directory_request: None,
working_directory: None,
workdir_attachment_requests: Vec::new(),
workdir_attachments: Vec::new(),
worker_observation_enabled: false,
worker_observation_grants: Vec::new(),
workspace_api: None,
@@ -423,6 +423,31 @@ impl runtime_api::RuntimeApi for RuntimeManagementApi {
Ok(runtime_api::WorkerResponse { worker })
}
async fn replace_worker_workdir_attachments(
&self,
worker_id: String,
value: runtime_api::WorkerWorkdirAttachmentsRequest,
) -> Result<runtime_api::WorkerResponse, runtime_api::RuntimeApiError> {
let attachments = request(value.workdir_attachments)?;
let worker_ref = worker_ref_for(&self.state.runtime, worker_id).map_err(api_error)?;
let worker = match auth_workspace_scope(&self.state, auth_extension().as_ref())
.map_err(api_error)?
{
Some(scope) => self
.state
.runtime
.replace_worker_workdir_attachments_scoped(&scope, &worker_ref, attachments),
None => self
.state
.runtime
.replace_worker_workdir_attachments(&worker_ref, attachments),
}
.map_err(RuntimeHttpRestError::runtime)
.map_err(api_error)?;
let worker = response(worker)?;
Ok(runtime_api::WorkerResponse { worker })
}
async fn complete_worker_arguments(
&self,
worker_id: String,
+412 -147
View File
@@ -1,7 +1,8 @@
use crate::catalog::{
ConfigBundleRef, CreateWorkerRequest, ProfileSelector, ProfileSourceArchiveSource,
RepositoryRefObservation, RepositoryRefObservationRequest, WorkerDetail, WorkerLifecycleAck,
WorkerRestoreIntent, WorkerStatus, WorkerSummary, WorkingDirectoryRepositoryAccessRequest,
WorkerRestoreIntent, WorkerStatus, WorkerSummary, WorkingDirectoryAttachmentClaim,
WorkingDirectoryAttachmentStatus, WorkingDirectoryRepositoryAccessRequest,
WorkingDirectoryRequest, WorkingDirectoryStatus as CatalogWorkingDirectoryStatus,
WorkspaceApiRef,
};
@@ -43,7 +44,7 @@ use crate::retention::{
use protocol::subscription::{
EventSubscriptionSelector, SubscriptionEventPayload, SubscriptionSnapshot,
SubscriptionValidationError, SubscriptionWorkdirId, SubscriptionWorker, SubscriptionWorkerId,
SubscriptionWorkerState,
SubscriptionWorkerState, SubscriptionWorkerWorkdirAttachment,
};
use protocol::{Event, Method};
use server_api::WorkerRestoreState;
@@ -729,8 +730,9 @@ impl Runtime {
.get(&worker_ref.worker_id)
.is_some_and(|worker| {
worker.belongs_to_workspace(&scope.workspace_id)
&& worker.working_directory.as_ref().is_some_and(|status| {
status.summary.working_directory_id == working_directory_id
&& worker.workdir_attachments.iter().any(|attachment| {
attachment.working_directory.summary.working_directory_id
== working_directory_id
})
});
if !owns_workdir {
@@ -764,7 +766,7 @@ impl Runtime {
let backend = {
let state = self.lock()?;
state.ensure_running()?;
if let Some(worker_id) = state.primary_worker_id_for_workdir(working_directory_id) {
if let Some(worker_id) = state.worker_id_for_workdir(working_directory_id) {
return Err(RuntimeError::InvalidRequest(format!(
"working directory {working_directory_id} is assigned to worker {worker_id}"
)));
@@ -792,12 +794,8 @@ impl Runtime {
fn annotate_working_directory_status(
&self,
mut status: CatalogWorkingDirectoryStatus,
status: CatalogWorkingDirectoryStatus,
) -> Result<CatalogWorkingDirectoryStatus, RuntimeError> {
let state = self.lock()?;
status.summary.primary_worker_id = state
.primary_worker_id_for_workdir(status.summary.working_directory_id.as_str())
.map(|worker_id| worker_id.to_string());
Ok(status)
}
@@ -1023,10 +1021,8 @@ impl Runtime {
return Ok(existing.detail());
}
state.validate_worker_config_boundary(&request)?;
if let Some(working_directory_id) = requested_primary_workdir_id(&request) {
if let Some(owner_worker_id) =
state.primary_worker_id_for_workdir(working_directory_id)
{
for working_directory_id in requested_workdir_ids(&request) {
if let Some(owner_worker_id) = state.worker_id_for_workdir(working_directory_id) {
return Err(RuntimeError::InvalidRequest(format!(
"working directory {working_directory_id} is already assigned to worker {owner_worker_id}"
)));
@@ -1064,7 +1060,7 @@ impl Runtime {
execution_metadata_available: true,
execution_bound: true,
restore_intent: WorkerRestoreIntent::Explicit,
working_directory: None,
workdir_attachments: Vec::new(),
execution_handle: None,
internal_workers: InternalWorkerActivityProjection::default(),
};
@@ -1076,19 +1072,19 @@ impl Runtime {
request,
workspace_scope: scope.cloned(),
context: self.execution_context(worker_ref.clone()),
working_directory: None,
workdir_attachments: BTreeMap::new(),
config_bundle,
};
(backend, worker_ref, spawn_request)
};
let spawn_result = backend.spawn_worker(spawn_request);
let (handle, initial_worker_state, working_directory) = match spawn_result {
let (handle, initial_worker_state, workdir_attachments) = match spawn_result {
WorkerExecutionSpawnResult::Connected {
handle,
worker_state,
working_directory,
} => (handle, worker_state, working_directory),
workdir_attachments,
} => (handle, worker_state, workdir_attachments),
WorkerExecutionSpawnResult::Rejected(result)
| WorkerExecutionSpawnResult::RolledBack(result)
| WorkerExecutionSpawnResult::Errored(result) => {
@@ -1101,8 +1097,30 @@ impl Runtime {
result,
});
}
WorkerExecutionSpawnResult::ReconciliationRequired { result, .. } => {
self.rollback_failed_create(&worker_ref)?;
WorkerExecutionSpawnResult::ReconciliationRequired {
result,
handle,
worker_state,
workdir_attachments,
} => {
let retain_result =
if let (Some(handle), Some(worker_state)) = (handle, worker_state) {
self.retain_restore_execution_evidence(
&worker_ref,
handle,
worker_state,
workdir_attachments,
)
} else {
self.retain_create_reconciliation_pending(&worker_ref, workdir_attachments)
};
if let Err(error) = retain_result {
tracing::error!(
worker_id = %worker_ref.worker_id,
%error,
"failed to persist Worker create reconciliation evidence"
);
}
return Err(RuntimeError::WorkerExecutionRejected {
worker_id: worker_ref.worker_id.clone(),
operation: result.operation,
@@ -1160,7 +1178,7 @@ impl Runtime {
&worker_ref,
handle.clone(),
initial_worker_state.clone(),
working_directory,
workdir_attachments,
dispatch_result,
) {
Ok(detail) => detail,
@@ -1179,7 +1197,7 @@ impl Runtime {
&worker_ref,
handle.clone(),
initial_worker_state,
working_directory,
workdir_attachments,
WorkerExecutionResult::accepted(WorkerExecutionOperation::Spawn),
) {
Ok(detail) => Ok(detail),
@@ -1425,6 +1443,119 @@ impl Runtime {
Ok(state.worker(worker_ref)?.detail())
}
pub fn replace_worker_workdir_attachments_scoped(
&self,
scope: &RuntimeWorkspaceScope,
worker_ref: &WorkerRef,
attachments: Vec<WorkingDirectoryAttachmentClaim>,
) -> Result<WorkerDetail, RuntimeError> {
self.ensure_worker_in_workspace(scope, worker_ref)?;
self.replace_worker_workdir_attachments(worker_ref, attachments)
}
/// Persist the exact current attachment set used by live mutation and future restore.
pub fn replace_worker_workdir_attachments(
&self,
worker_ref: &WorkerRef,
mut attachments: Vec<WorkingDirectoryAttachmentClaim>,
) -> Result<WorkerDetail, RuntimeError> {
let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?;
let _operation_guard = operation_lock
.lock()
.map_err(|_| RuntimeError::StatePoisoned)?;
let backend = {
let state = self.lock()?;
let worker = state.worker(worker_ref)?;
if worker.execution_handle.is_none() || !worker.status.is_active() {
return Err(RuntimeError::WorkerExecutionUnavailable {
worker_id: worker.worker_id,
message: "Workdir attachments can only be changed for a live Worker"
.to_string(),
});
}
let request = worker.request.as_ref().ok_or_else(|| {
RuntimeError::WorkerExecutionUnavailable {
worker_id: worker.worker_id,
message: "persisted Worker restore request is unavailable".to_string(),
}
})?;
for attachment in &mut attachments {
if attachment.relative_cwd.is_none()
&& let Some(previous) = request.workdir_attachments.iter().find(|previous| {
previous.alias == attachment.alias
&& previous.working_directory_id == attachment.working_directory_id
})
{
attachment.relative_cwd = previous.relative_cwd.clone();
}
}
let mut candidate = request.clone();
candidate.workdir_attachment_requests.clear();
candidate.workdir_attachments = attachments.clone();
validate_create_worker_request(&candidate)?;
state.execution_backend.clone().ok_or_else(|| {
RuntimeError::ExecutionBackendUnavailable {
message: "Workdir attachment mutation requires an execution backend"
.to_string(),
}
})?
};
let statuses = attachments
.iter()
.map(|attachment| {
backend
.working_directory(&attachment.working_directory_id)
.map(|working_directory| WorkingDirectoryAttachmentStatus {
alias: attachment.alias.clone(),
working_directory,
})
.map_err(RuntimeError::from)
})
.collect::<Result<Vec<_>, RuntimeError>>()?;
if let Some(inactive) = statuses.iter().find(|attachment| {
attachment.working_directory.summary.status
!= workdir::workspace::WorkingDirectoryStatusKind::Active
}) {
return Err(RuntimeError::InvalidRequest(format!(
"working directory {} is not active ({})",
inactive.working_directory.summary.working_directory_id,
inactive.working_directory.summary.status
)));
}
let mut state = self.lock()?;
for attachment in &attachments {
if let Some(owner_worker_id) =
state.worker_id_for_workdir(&attachment.working_directory_id)
&& owner_worker_id != worker_ref.worker_id
{
return Err(RuntimeError::InvalidRequest(format!(
"working directory {} is already assigned to worker {owner_worker_id}",
attachment.working_directory_id
)));
}
}
let previous = state.worker(worker_ref)?.clone();
{
let worker = state.worker_mut(worker_ref)?;
let request = worker
.request
.as_mut()
.expect("restore request checked above");
request.workdir_attachment_requests.clear();
request.workdir_attachments = attachments;
worker.workdir_attachments = statuses;
}
if let Err(error) = state.persist_worker(&worker_ref.worker_id) {
state.workers.insert(worker_ref.worker_id, previous);
return Err(error);
}
state.publish_worker_upsert(worker_ref.worker_id)?;
Ok(state.worker(worker_ref)?.detail())
}
/// Observe a Worker's live protocol availability or its retained Session.
/// The same lifecycle lock used by stop/restore fences the decision and read.
pub fn worker_session_scoped(
@@ -1570,8 +1701,8 @@ impl Runtime {
request: worker_request,
workspace_scope,
context: self.execution_context(worker_ref.clone()),
previous_working_directory: worker.working_directory.clone(),
working_directory: None,
previous_workdir_attachments: worker.workdir_attachments.clone(),
workdir_attachments: BTreeMap::new(),
config_bundle: None,
};
(backend, request)
@@ -1589,14 +1720,14 @@ impl Runtime {
WorkerExecutionSpawnResult::Connected {
handle,
worker_state,
working_directory,
workdir_attachments,
} => {
let commit = self.commit_restored_worker_execution(
worker_ref,
handle.clone(),
worker_state.clone(),
WorkerStatus::Idle,
working_directory.clone(),
workdir_attachments.clone(),
);
match commit {
Ok(worker) => Ok(RuntimeWorkerRestoreResult::accepted(worker)),
@@ -1625,7 +1756,7 @@ impl Runtime {
worker_ref,
handle,
worker_state,
working_directory,
workdir_attachments,
) {
tracing::error!(
worker_id = %worker_ref.worker_id,
@@ -1663,14 +1794,14 @@ impl Runtime {
result: _result,
handle,
worker_state,
working_directory,
workdir_attachments,
} => {
if let (Some(handle), Some(worker_state)) = (handle, worker_state) {
if let Err(retain_error) = self.retain_restore_execution_evidence(
worker_ref,
handle,
worker_state,
working_directory,
workdir_attachments,
) {
tracing::error!(
worker_id = %worker_ref.worker_id,
@@ -2079,7 +2210,7 @@ impl Runtime {
worker_ref: &WorkerRef,
handle: WorkerExecutionHandle,
initial_worker_state: protocol::WorkerStateSnapshot,
working_directory: Option<CatalogWorkingDirectoryStatus>,
workdir_attachments: Vec<crate::catalog::WorkingDirectoryAttachmentStatus>,
result: WorkerExecutionResult,
) -> Result<WorkerDetail, RuntimeError> {
let mut state = self.lock()?;
@@ -2094,7 +2225,7 @@ impl Runtime {
let _ = worker.apply_worker_state(snapshot);
}
worker.restore_intent = restore_intent_for_status(worker.status);
worker.working_directory = working_directory;
worker.workdir_attachments = workdir_attachments;
worker.detail()
};
state.publish_worker_upsert(worker_ref.worker_id)?;
@@ -2367,9 +2498,13 @@ impl Runtime {
state.ensure_running()?;
state.ensure_worker_ref(worker_ref)?;
let worker = state.worker(worker_ref)?;
if worker.execution_handle.is_some() && worker.status.is_active() {
if worker.status.is_active()
&& (worker.execution_handle.is_some()
|| (!worker.execution_metadata_available
&& !worker.workdir_attachments.is_empty()))
{
return Err(RuntimeError::InvalidRequest(format!(
"worker {} is running and must be stopped before deletion",
"worker {} is active or requires reconciliation and must be stopped before deletion",
worker_ref.worker_id
)));
}
@@ -2398,9 +2533,12 @@ impl Runtime {
state.ensure_running()?;
state.ensure_worker_ref(worker_ref)?;
let worker = state.worker(worker_ref)?;
if worker.execution_handle.is_some() && worker.status.is_active() {
if worker.status.is_active()
&& (worker.execution_handle.is_some()
|| (!worker.execution_metadata_available && !worker.workdir_attachments.is_empty()))
{
return Err(RuntimeError::InvalidRequest(format!(
"worker {} became active before deletion",
"worker {} became active or reconciliation-required before deletion",
worker_ref.worker_id
)));
}
@@ -2662,7 +2800,7 @@ impl Runtime {
handle: WorkerExecutionHandle,
worker_state: protocol::WorkerStateSnapshot,
status: WorkerStatus,
working_directory: Option<CatalogWorkingDirectoryStatus>,
workdir_attachments: Vec<crate::catalog::WorkingDirectoryAttachmentStatus>,
) -> Result<WorkerDetail, RuntimeError> {
let mut state = self.lock()?;
state.ensure_worker_ref(worker_ref)?;
@@ -2679,7 +2817,7 @@ impl Runtime {
candidate.status = status;
let _ = candidate.apply_worker_state(&worker_state);
candidate.restore_intent = restore_intent_for_status(candidate.status);
candidate.working_directory = working_directory;
candidate.workdir_attachments = workdir_attachments;
// Validate the exact subscription projection before persistence. The
// subsequent publish uses the same pure projection while the Worker
@@ -2708,7 +2846,7 @@ impl Runtime {
worker_ref: &WorkerRef,
handle: WorkerExecutionHandle,
worker_state: protocol::WorkerStateSnapshot,
working_directory: Option<CatalogWorkingDirectoryStatus>,
workdir_attachments: Vec<crate::catalog::WorkingDirectoryAttachmentStatus>,
) -> Result<(), RuntimeError> {
let mut state = self.lock()?;
let mut candidate = state.worker(worker_ref)?.clone();
@@ -2718,7 +2856,26 @@ impl Runtime {
candidate.status = WorkerStatus::Idle;
let _ = candidate.apply_worker_state(&worker_state);
candidate.restore_intent = WorkerRestoreIntent::Automatic;
candidate.working_directory = working_directory;
candidate.workdir_attachments = workdir_attachments;
state.workers.insert(worker_ref.worker_id, candidate);
state.publish_worker_upsert(worker_ref.worker_id)?;
state.persist_worker(&worker_ref.worker_id)?;
Ok(())
}
fn retain_create_reconciliation_pending(
&self,
worker_ref: &WorkerRef,
workdir_attachments: Vec<crate::catalog::WorkingDirectoryAttachmentStatus>,
) -> Result<(), RuntimeError> {
let mut state = self.lock()?;
let mut candidate = state.worker(worker_ref)?.clone();
candidate.execution_handle = None;
candidate.execution_metadata_available = false;
candidate.execution_bound = true;
candidate.status = WorkerStatus::Idle;
candidate.restore_intent = WorkerRestoreIntent::Explicit;
candidate.workdir_attachments = workdir_attachments;
state.workers.insert(worker_ref.worker_id, candidate);
state.publish_worker_upsert(worker_ref.worker_id)?;
state.persist_worker(&worker_ref.worker_id)?;
@@ -3078,7 +3235,7 @@ impl RuntimeState {
execution_metadata_available,
execution_bound,
restore_intent,
working_directory: worker.working_directory,
workdir_attachments: worker.workdir_attachments,
execution_handle: None,
internal_workers: InternalWorkerActivityProjection::default(),
},
@@ -3394,18 +3551,25 @@ impl RuntimeState {
) -> Result<SubscriptionWorker, RuntimeError> {
let worker_id = SubscriptionWorkerId::new(worker.worker_id.to_string())
.map_err(subscription_validation_error)?;
let repository_id = worker
.working_directory
.as_ref()
.map(|working_directory| working_directory.summary.repository_id.clone());
let working_directory_id = worker
.working_directory
.as_ref()
.map(|working_directory| {
SubscriptionWorkdirId::new(working_directory.summary.working_directory_id.clone())
.map_err(subscription_validation_error)
let workdir_attachments = worker
.workdir_attachments
.iter()
.map(|attachment| {
Ok(SubscriptionWorkerWorkdirAttachment {
alias: attachment.alias.to_string(),
repository_id: Some(attachment.working_directory.summary.repository_id.clone()),
repository_key: None,
working_directory_id: SubscriptionWorkdirId::new(
attachment
.working_directory
.summary
.working_directory_id
.clone(),
)
.map_err(subscription_validation_error)?,
})
})
.transpose()?;
.collect::<Result<Vec<_>, RuntimeError>>()?;
let profile = match &worker.profile {
ProfileSelector::Builtin(name) | ProfileSelector::Named(name) => Some(name.clone()),
};
@@ -3425,9 +3589,7 @@ impl RuntimeState {
workspace_id: worker.workspace_id.clone(),
display_name: worker.display_name.clone(),
profile,
repository_id,
repository_key: None,
working_directory_id,
workdir_attachments,
})
}
@@ -3525,22 +3687,15 @@ impl RuntimeState {
}
}
fn primary_worker_id_for_workdir(&self, working_directory_id: &str) -> Option<WorkerId> {
fn worker_id_for_workdir(&self, working_directory_id: &str) -> Option<WorkerId> {
self.workers.values().find_map(|worker| {
if worker
.working_directory
.as_ref()
.is_some_and(|binding| binding.summary.working_directory_id == working_directory_id)
|| worker
.request
.as_ref()
.and_then(requested_primary_workdir_id)
== Some(working_directory_id)
{
Some(worker.worker_id)
} else {
None
}
let has_status = worker.workdir_attachments.iter().any(|attachment| {
attachment.working_directory.summary.working_directory_id == working_directory_id
});
let has_request = worker.request.as_ref().is_some_and(|request| {
requested_workdir_ids(request).contains(&working_directory_id)
});
(has_status || has_request).then_some(worker.worker_id)
})
}
@@ -3976,7 +4131,7 @@ struct WorkerRecord {
execution_metadata_available: bool,
execution_bound: bool,
restore_intent: WorkerRestoreIntent,
working_directory: Option<CatalogWorkingDirectoryStatus>,
workdir_attachments: Vec<crate::catalog::WorkingDirectoryAttachmentStatus>,
execution_handle: Option<WorkerExecutionHandle>,
internal_workers: InternalWorkerActivityProjection,
}
@@ -3999,7 +4154,7 @@ impl WorkerRecord {
execution_metadata_available: self.execution_metadata_available,
worker_state: self.worker_state.clone(),
workspace_id: self.workspace_id.clone(),
working_directory: self.working_directory.clone(),
workdir_attachments: self.workdir_attachments.clone(),
profile: self.profile.clone(),
display_name: self.display_name.clone(),
profile_source: self.profile_source.clone(),
@@ -4016,7 +4171,7 @@ impl WorkerRecord {
execution_metadata_available: self.execution_metadata_available,
worker_state: self.worker_state.clone(),
workspace_id: self.workspace_id.clone(),
working_directory: self.working_directory.clone(),
workdir_attachments: self.workdir_attachments.clone(),
profile: self.profile.clone(),
display_name: self.display_name.clone(),
profile_source: self.profile_source.clone(),
@@ -4049,7 +4204,7 @@ impl WorkerRecord {
status: self.status,
execution_state,
workspace_id: self.workspace_id.clone(),
working_directory: self.working_directory.clone(),
workdir_attachments: self.workdir_attachments.clone(),
}
}
}
@@ -4100,25 +4255,28 @@ fn repository_resource_error(error: BackendResourceError) -> RuntimeError {
fn durable_create_worker_request(request: &CreateWorkerRequest) -> CreateWorkerRequest {
let mut durable = request.clone();
if let Some(working_directory) = durable.working_directory_request.as_mut()
&& let Some(materialization) = working_directory.materialization.as_mut()
{
materialization.ssh = None;
for attachment in &mut durable.workdir_attachment_requests {
if let Some(materialization) = attachment.working_directory.materialization.as_mut() {
materialization.ssh = None;
}
}
durable
}
fn requested_primary_workdir_id(request: &CreateWorkerRequest) -> Option<&str> {
fn requested_workdir_ids(request: &CreateWorkerRequest) -> Vec<&str> {
request
.working_directory
.as_ref()
.workdir_attachments
.iter()
.map(|claim| claim.working_directory_id.as_str())
.or_else(|| {
.chain(
request
.working_directory_request
.as_ref()
.and_then(|request| request.backend_workdir_id.as_deref())
})
.workdir_attachment_requests
.iter()
.filter_map(|attachment| {
attachment.working_directory.backend_workdir_id.as_deref()
}),
)
.collect()
}
fn validate_create_worker_request(request: &CreateWorkerRequest) -> Result<(), RuntimeError> {
@@ -4127,6 +4285,42 @@ fn validate_create_worker_request(request: &CreateWorkerRequest) -> Result<(), R
"create_fingerprint must not be empty".to_string(),
));
}
let mut aliases = std::collections::BTreeSet::new();
let mut workdir_ids = std::collections::BTreeSet::new();
for (alias, workdir_id) in request
.workdir_attachments
.iter()
.map(|attachment| {
(
&attachment.alias,
Some(attachment.working_directory_id.as_str()),
)
})
.chain(
request
.workdir_attachment_requests
.iter()
.map(|attachment| {
(
&attachment.alias,
attachment.working_directory.backend_workdir_id.as_deref(),
)
}),
)
{
if !aliases.insert(alias.as_str()) {
return Err(RuntimeError::InvalidRequest(format!(
"duplicate Workdir attachment alias `{alias}`"
)));
}
if let Some(workdir_id) = workdir_id
&& !workdir_ids.insert(workdir_id)
{
return Err(RuntimeError::InvalidRequest(format!(
"working directory `{workdir_id}` cannot be attached more than once"
)));
}
}
match &request.profile_source {
crate::catalog::ProfileSourceArchiveSource::Embedded { archive } => {
archive.verify().map_err(|err| {
@@ -4420,8 +4614,8 @@ mod tests {
use crate::catalog::{
ConfigBundleRef, MaterializerKind, ProfileSelector, RepositoryMaterializationContext,
RepositorySshCredentialCandidate, RepositorySshMaterializationAccess, SensitiveString,
WorkingDirectoryClaim, WorkingDirectoryRepository, WorkingDirectoryRequest,
WorkspaceApiRef,
WorkingDirectoryAttachmentClaim, WorkingDirectoryAttachmentRequest,
WorkingDirectoryRepository, WorkingDirectoryRequest, WorkspaceApiRef,
};
use crate::config_bundle::{
ConfigBundle, ConfigBundleMetadata, ConfigBundleProvenance, ConfigDeclaration,
@@ -4913,8 +5107,8 @@ mod tests {
digest: bundle.metadata.digest,
}),
initial_input: None,
working_directory_request: None,
working_directory: None,
workdir_attachment_requests: Vec::new(),
workdir_attachments: Vec::new(),
worker_observation_enabled: false,
worker_observation_grants: Vec::new(),
workspace_api: None,
@@ -4929,60 +5123,64 @@ mod tests {
#[test]
fn durable_worker_request_omits_repository_credentials() {
let mut request = task_request("worker-secret-redaction");
request.working_directory_request = Some(WorkingDirectoryRequest {
repository: WorkingDirectoryRepository {
id: "repository-1".to_string(),
provider: "git".to_string(),
source: server_api::RepositorySource {
kind: server_api::RepositorySourceKind::Ssh,
uri: "ssh://git@example.test/repo.git".to_string(),
request.workdir_attachment_requests = vec![WorkingDirectoryAttachmentRequest {
alias: workdir::WorkdirAttachmentAlias::new("checkout").unwrap(),
working_directory: WorkingDirectoryRequest {
repository: WorkingDirectoryRepository {
id: "repository-1".to_string(),
provider: "git".to_string(),
source: server_api::RepositorySource {
kind: server_api::RepositorySourceKind::Ssh,
uri: "ssh://git@example.test/repo.git".to_string(),
},
source_revision: 1,
source_fingerprint: "sha256:source".to_string(),
selector: None,
},
source_revision: 1,
source_fingerprint: "sha256:source".to_string(),
selector: None,
},
materializer: MaterializerKind::RuntimeGitClone,
backend_workdir_id: Some("working-directory-1".to_string()),
materialization: Some(RepositoryMaterializationContext {
workspace_id: "workspace-1".to_string(),
runtime_id: "runtime-1".to_string(),
operation_id: "operation-1".to_string(),
config_revision: 1,
config_projection_digest: "sha256:projection".to_string(),
ssh: Some(RepositorySshMaterializationAccess {
credential_candidates: vec![RepositorySshCredentialCandidate {
credential_id: "credential-1".to_string(),
credential_revision: 1,
private_key: SensitiveString::new("private-key-bytes"),
}],
host_trust_id: "host-trust-1".to_string(),
host_trust_revision: 1,
access: server_api::RepositoryAccessMode::ReadOnly,
expires_at_epoch_seconds: u64::MAX,
repository_id: "repository-1".to_string(),
repository_source_fingerprint: "sha256:source".to_string(),
repository_uri: "ssh://git@example.test/repo.git".to_string(),
secret_resource: repository_resource_handle(),
known_hosts_entry: SensitiveString::new("known-hosts-entry"),
display_name: None,
materializer: MaterializerKind::RuntimeGitClone,
backend_workdir_id: Some("working-directory-1".to_string()),
materialization: Some(RepositoryMaterializationContext {
workspace_id: "workspace-1".to_string(),
runtime_id: "runtime-1".to_string(),
operation_id: "operation-1".to_string(),
config_revision: 1,
config_projection_digest: "sha256:projection".to_string(),
ssh: Some(RepositorySshMaterializationAccess {
credential_candidates: vec![RepositorySshCredentialCandidate {
credential_id: "credential-1".to_string(),
credential_revision: 1,
private_key: SensitiveString::new("private-key-bytes"),
}],
host_trust_id: "host-trust-1".to_string(),
host_trust_revision: 1,
access: server_api::RepositoryAccessMode::ReadOnly,
expires_at_epoch_seconds: u64::MAX,
repository_id: "repository-1".to_string(),
repository_source_fingerprint: "sha256:source".to_string(),
repository_uri: "ssh://git@example.test/repo.git".to_string(),
secret_resource: repository_resource_handle(),
known_hosts_entry: SensitiveString::new("known-hosts-entry"),
}),
}),
}),
});
},
}];
let durable = durable_create_worker_request(&request);
assert!(
request
.working_directory_request
.as_ref()
.and_then(|working_directory| working_directory.materialization.as_ref())
.workdir_attachment_requests
.first()
.and_then(|attachment| attachment.working_directory.materialization.as_ref())
.and_then(|materialization| materialization.ssh.as_ref())
.is_some()
);
assert!(
durable
.working_directory_request
.as_ref()
.and_then(|working_directory| working_directory.materialization.as_ref())
.workdir_attachment_requests
.first()
.and_then(|attachment| attachment.working_directory.materialization.as_ref())
.and_then(|materialization| materialization.ssh.as_ref())
.is_none()
);
@@ -5172,6 +5370,7 @@ mod tests {
source_fingerprint: "sha256:source".to_string(),
selector: None,
},
display_name: None,
materializer: MaterializerKind::RuntimeGitClone,
backend_workdir_id: Some("working-directory-1".to_string()),
materialization: Some(RepositoryMaterializationContext {
@@ -5342,6 +5541,7 @@ mod tests {
#[derive(Default)]
struct TestExecutionBackend {
spawn_result: Mutex<Option<WorkerExecutionSpawnResult>>,
dispatch_result: Mutex<Option<WorkerExecutionResult>>,
stop_result: Mutex<Option<WorkerExecutionResult>>,
stop_gate: Mutex<Option<Arc<RestoreGate>>>,
@@ -5447,6 +5647,9 @@ mod tests {
}
fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult {
if let Some(result) = self.spawn_result.lock().unwrap().take() {
return result;
}
self.config_bundles
.lock()
.unwrap()
@@ -5460,10 +5663,16 @@ mod tests {
worker_state: protocol::WorkerStateSnapshot {
..protocol::WorkerStatus::Idle.into()
},
working_directory: request
.working_directory
.as_ref()
.map(|binding| binding.status()),
workdir_attachments: request
.workdir_attachments
.iter()
.map(
|(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus {
alias: alias.clone(),
working_directory: binding.status(),
},
)
.collect(),
}
}
@@ -5502,10 +5711,16 @@ mod tests {
worker_state: protocol::WorkerStateSnapshot {
..protocol::WorkerStatus::Idle.into()
},
working_directory: request
.working_directory
.as_ref()
.map(|binding| binding.status()),
workdir_attachments: request
.workdir_attachments
.iter()
.map(
|(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus {
alias: alias.clone(),
working_directory: binding.status(),
},
)
.collect(),
}
}
@@ -6263,10 +6478,11 @@ mod tests {
let runtime = runtime_with_backend();
let mut request = task_request("idempotent");
request.create_fingerprint = "sha256:input-1".to_string();
request.working_directory = Some(WorkingDirectoryClaim {
request.workdir_attachments = vec![WorkingDirectoryAttachmentClaim {
alias: workdir::WorkdirAttachmentAlias::new("workdir").unwrap(),
working_directory_id: "workdir-idempotent".to_string(),
relative_cwd: None,
});
}];
let first = runtime.create_worker(request.clone()).unwrap();
let workdir_count_after_first = runtime.list_working_directories().unwrap().len();
@@ -6631,6 +6847,49 @@ mod tests {
assert!(runtime.list_workers().unwrap().is_empty());
}
#[test]
fn reconciliation_required_create_retains_durable_worker_evidence() {
let (runtime, backend) = runtime_and_backend();
runtime.store_config_bundle(test_bundle()).unwrap();
*backend.spawn_result.lock().unwrap() =
Some(WorkerExecutionSpawnResult::ReconciliationRequired {
result: WorkerExecutionResult::errored(
WorkerExecutionOperation::Spawn,
"Workdir cleanup could not be proven",
),
handle: None,
worker_state: None,
workdir_attachments: vec![WorkingDirectoryAttachmentStatus {
alias: workdir::WorkdirAttachmentAlias::new("checkout").unwrap(),
working_directory: serde_json::from_value(serde_json::json!({
"summary": {
"working_directory_id": "wd-uncertain",
"repository_id": "repo",
"materializer_kind": "runtime_git_clone",
"status": "active"
}
}))
.unwrap(),
}],
});
let request = task_request("uncertain create cleanup");
let worker_ref = WorkerRef::new(request.worker_id);
let error = runtime.create_worker(request).unwrap_err();
assert!(matches!(
error,
RuntimeError::WorkerExecutionRejected { .. }
));
let retained = runtime.worker_detail(&worker_ref).unwrap();
assert_eq!(retained.status, WorkerStatus::Idle);
assert!(!retained.execution_metadata_available);
assert!(matches!(
runtime.delete_worker(&worker_ref),
Err(RuntimeError::InvalidRequest(_))
));
}
#[test]
fn create_worker_without_execution_backend_is_rejected_and_not_persisted() {
let runtime = Runtime::new_memory();
@@ -6804,10 +7063,16 @@ mod tests {
worker_state: protocol::WorkerStateSnapshot {
..protocol::WorkerStatus::Idle.into()
},
working_directory: request
.working_directory
.as_ref()
.map(|binding| binding.status()),
workdir_attachments: request
.workdir_attachments
.iter()
.map(
|(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus {
alias: alias.clone(),
working_directory: binding.status(),
},
)
.collect(),
}
}
@@ -7667,7 +7932,7 @@ mod tests {
),
handle: None,
worker_state: None,
working_directory: None,
workdir_attachments: Vec::new(),
},
);
+470 -195
View File
@@ -7,7 +7,7 @@
//! socket paths, session paths, manifests, credentials, or handles leave this
//! module.
use std::collections::HashMap;
use std::collections::{BTreeMap, HashMap};
use std::future::Future;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
@@ -17,8 +17,9 @@ use std::time::Duration;
use crate::auth::{BACKEND_RESOURCE_FETCH_PERMISSION, RuntimeIdentityMaterial};
use crate::catalog::{
CreateWorkerRequest, ProfileSourceArchiveSource, RepositoryRefObservation,
RepositoryRefObservationRequest, WorkingDirectoryRepositoryAccessRequest,
WorkingDirectoryRequest, WorkingDirectoryStatus, WorkingDirectoryStatusKind,
RepositoryRefObservationRequest, WorkingDirectoryAttachmentStatus,
WorkingDirectoryRepositoryAccessRequest, WorkingDirectoryRequest, WorkingDirectoryStatus,
WorkingDirectoryStatusKind,
};
use crate::config_bundle::{ConfigBundle, workspace_config_etag};
use crate::execution::{
@@ -44,6 +45,45 @@ use protocol::{Event, Method, Segment, WorkerCommandEnvelope};
static NEXT_INTERNAL_COMMAND_ID: AtomicU64 = AtomicU64::new(1);
fn rollback_materialized_spawn_workdirs(
materializer: &dyn WorkingDirectoryMaterializer,
bindings: &BTreeMap<WorkdirAttachmentAlias, WorkingDirectoryBinding>,
workdir_ids: &[String],
failure_context: &str,
) -> Result<(), WorkerExecutionSpawnResult> {
let mut cleanup_failures = Vec::new();
let mut uncertain_attachments = Vec::new();
for workdir_id in workdir_ids.iter().rev() {
if let Err(error) = materializer.cleanup_working_directory(workdir_id) {
cleanup_failures.push(format!("{workdir_id}: {error}"));
if let Some((alias, binding)) = bindings
.iter()
.find(|(_, binding)| binding.working_directory.id == *workdir_id)
{
uncertain_attachments.push(WorkingDirectoryAttachmentStatus {
alias: alias.clone(),
working_directory: binding.status(),
});
}
}
}
if cleanup_failures.is_empty() {
return Ok(());
}
Err(WorkerExecutionSpawnResult::ReconciliationRequired {
result: WorkerExecutionResult::errored(
WorkerExecutionOperation::Spawn,
format!(
"{failure_context}; Workdir cleanup could not be proven: {}",
cleanup_failures.join("; ")
),
),
handle: None,
worker_state: None,
workdir_attachments: uncertain_attachments,
})
}
fn next_internal_command(
state: &RwLock<protocol::WorkerStateSnapshot>,
) -> Result<WorkerCommandEnvelope, String> {
@@ -68,7 +108,10 @@ use session_store::{FsStore, FsWorkerStore};
use tokio::runtime::Runtime;
#[cfg(feature = "ws-server")]
use tokio::sync::broadcast;
use workdir::{LocalWorkdirSession, Workdir, WorkdirSessionCapabilities, WorkdirSessionHandle};
use workdir::{
LocalWorkdirSession, Workdir, WorkdirAttachmentAlias, WorkdirSessionCapabilities,
WorkdirSessionHandle, WorkdirSessionRouter,
};
#[cfg(test)]
use worker::WorkerController;
@@ -730,6 +773,29 @@ fn runtime_local_workdir_session(
))
}
fn runtime_local_workdir_router(
attachments: &BTreeMap<WorkdirAttachmentAlias, WorkingDirectoryBinding>,
scope: manifest::SharedScope,
) -> Result<Arc<WorkdirSessionRouter>, String> {
let router = Arc::new(WorkdirSessionRouter::new());
for (alias, binding) in attachments {
router
.attach(
alias.clone(),
runtime_local_workdir_session(
&binding.working_directory.id,
binding.root(),
binding.cwd(),
scope.clone(),
binding.command_environment(),
binding.session_resources(),
),
)
.map_err(|error| format!("bind Workdir attachment `{alias}`: {error}"))?;
}
Ok(router)
}
fn bind_workspace_memory_settings(
manifest: &mut manifest::WorkerManifest,
request: &CreateWorkerRequest,
@@ -845,15 +911,14 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
) -> Result<RuntimeWorkerController, String> {
let worker_name = Self::runtime_worker_name(&request);
let profile = Self::runtime_profile(&request);
let has_local_filesystem = request.working_directory.is_some();
let worker_root = request
.working_directory
.as_ref()
let has_local_filesystem = request.workdir_attachments.len() == 1;
let only_binding = (request.workdir_attachments.len() == 1)
.then(|| request.workdir_attachments.values().next())
.flatten();
let worker_root = only_binding
.map(|binding| binding.root().to_path_buf())
.unwrap_or_else(|| self.profile_base_dir.clone());
let filesystem_authority = request
.working_directory
.as_ref()
let filesystem_authority = only_binding
.map(|binding| {
WorkerFilesystemAuthority::local(
binding.root().to_path_buf(),
@@ -965,18 +1030,10 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
})?;
let worker = prepared.worker_mut();
validate_worker_memory_settings(worker.manifest(), &request.request)?;
if let Some(binding) = request.working_directory.as_ref() {
worker.bind_workdir_session(Some(runtime_local_workdir_session(
&binding.working_directory.id,
binding.root(),
binding.cwd(),
worker.scope().clone(),
binding.command_environment(),
binding.session_resources(),
)));
} else {
worker.bind_workdir_session(None);
}
worker.bind_workdir_sessions(runtime_local_workdir_router(
&request.workdir_attachments,
worker.scope().clone(),
)?);
if let (Some(runtime_id), Some(workspace_id)) =
(observation_runtime_id, observation_workspace_id.clone())
&& observation_enabled
@@ -1095,9 +1152,10 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
request: WorkerExecutionRestoreRequest,
) -> Result<RuntimeWorkerController, String> {
let worker_name = Self::runtime_worker_name_for_ref(&request.worker_ref);
let filesystem_authority = request
.working_directory
.as_ref()
let only_binding = (request.workdir_attachments.len() == 1)
.then(|| request.workdir_attachments.values().next())
.flatten();
let filesystem_authority = only_binding
.map(|binding| {
WorkerFilesystemAuthority::local(
binding.root().to_path_buf(),
@@ -1201,18 +1259,10 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
};
validate_worker_memory_settings(worker.manifest(), &request.request)?;
let flow_transition_enabled = worker.manifest().feature.flow.enabled;
if let Some(binding) = request.working_directory.as_ref() {
worker.bind_workdir_session(Some(runtime_local_workdir_session(
&binding.working_directory.id,
binding.root(),
binding.cwd(),
worker.scope().clone(),
binding.command_environment(),
binding.session_resources(),
)));
} else {
worker.bind_workdir_session(None);
}
worker.bind_workdir_sessions(runtime_local_workdir_router(
&request.workdir_attachments,
worker.scope().clone(),
)?);
if let (Some(runtime_id), Some(workspace_id)) =
(observation_runtime_id, observation_workspace_id.clone())
&& observation_enabled
@@ -1738,7 +1788,7 @@ where
handle: WorkerHandle,
shutdown: Arc<tokio::sync::Mutex<Option<worker::ShutdownReceiver>>>,
controller_task: tokio::task::JoinHandle<()>,
working_directory: Option<WorkingDirectoryBinding>,
workdir_attachments: BTreeMap<WorkdirAttachmentAlias, WorkingDirectoryBinding>,
workspace_client: Option<Arc<dyn WorkspaceClient>>,
) -> WorkerExecutionSpawnResult {
#[cfg(feature = "ws-server")]
@@ -1818,8 +1868,13 @@ where
Err(_) => {
let retained_state =
worker_state.read().ok().map(|state| state.clone());
let retained_working_directory =
working_directory.as_ref().map(|binding| binding.status());
let retained_workdir_attachments = workdir_attachments
.iter()
.map(|(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus {
alias: alias.clone(),
working_directory: binding.status(),
})
.collect();
let retained_handle = self
.retain_uncertain_unconnected_controller(
&worker_ref,
@@ -1834,7 +1889,7 @@ where
result,
handle: retained_handle,
worker_state: retained_state,
working_directory: retained_working_directory,
workdir_attachments: retained_workdir_attachments,
}
}
}
@@ -1878,7 +1933,7 @@ where
result,
handle: Some(existing_handle),
worker_state: existing_worker_state,
working_directory: None,
workdir_attachments: Vec::new(),
},
}
} else {
@@ -1904,7 +1959,15 @@ where
WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(worker_ref, self.backend_id()),
worker_state: connected_worker_state,
working_directory: working_directory.map(|binding| binding.status()),
workdir_attachments: workdir_attachments
.into_iter()
.map(
|(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus {
alias,
working_directory: binding.status(),
},
)
.collect(),
}
}
}
@@ -2118,67 +2181,134 @@ where
}
let mut request = request;
let mut rollback_working_directory = None;
let working_directory = match (
request.request.working_directory_request.as_ref(),
request.request.working_directory.as_ref(),
) {
(Some(_), Some(_)) => {
let Some(materializer) = self.working_directory_materializer.as_ref().or_else(|| {
(request.request.workdir_attachment_requests.is_empty()
&& request.request.workdir_attachments.is_empty())
.then_some(())
.and(None)
}) else {
if request.request.workdir_attachment_requests.is_empty()
&& request.request.workdir_attachments.is_empty()
{
// No provider is needed for a Workdir-less Worker.
let factory = self.factory.clone();
let bridge_context = request.context.clone();
let worker_ref = request.worker_ref.clone();
let spawn_result = self
.run_cancellable_on_adapter_runtime(self.spawn_restore_timeout, async move {
factory.spawn_controller(request).await
});
let controller = match spawn_result {
Ok(controller) => controller,
Err(message) => {
return WorkerExecutionSpawnResult::Errored(
WorkerExecutionResult::errored(
WorkerExecutionOperation::Spawn,
message,
),
);
}
};
return self.connect_handle(
WorkerExecutionOperation::Spawn,
worker_ref,
bridge_context,
controller.handle,
controller.shutdown,
controller.controller_task,
BTreeMap::new(),
Some(controller.workspace_client),
);
}
return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Spawn,
"Workdir attachments were requested, but no materializer is configured for this Runtime backend",
));
};
let mut rollback_workdirs: Vec<String> = Vec::new();
for attachment in request.request.workdir_attachment_requests.clone() {
if request.workdir_attachments.contains_key(&attachment.alias) {
let message = format!("duplicate Workdir attachment alias `{}`", attachment.alias);
if let Err(result) = rollback_materialized_spawn_workdirs(
materializer.as_ref(),
&request.workdir_attachments,
&rollback_workdirs,
&message,
) {
return result;
}
return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Spawn,
"worker spawn cannot specify both working_directory_request and working_directory",
message,
));
}
(Some(working_directory_request), None) => {
let Some(materializer) = self.working_directory_materializer.as_ref() else {
match materializer.materialize(&request.worker_ref, &attachment.working_directory) {
Ok(binding) => {
rollback_workdirs.push(binding.working_directory.id.clone());
request
.workdir_attachments
.insert(attachment.alias, binding);
}
Err(error) => {
let message = error.to_string();
if let Err(result) = rollback_materialized_spawn_workdirs(
materializer.as_ref(),
&request.workdir_attachments,
&rollback_workdirs,
&message,
) {
return result;
}
return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Spawn,
"working directory materialization requested, but no materializer is configured for this runtime backend",
message,
));
};
match materializer.materialize(&request.worker_ref, working_directory_request) {
Ok(binding) => {
request.working_directory = Some(binding.clone());
rollback_working_directory = Some(binding.clone());
Some(binding)
}
Err(error) => {
return WorkerExecutionSpawnResult::Rejected(
WorkerExecutionResult::rejected(
WorkerExecutionOperation::Spawn,
error.to_string(),
),
);
}
}
}
(None, Some(working_directory)) => {
let Some(materializer) = self.working_directory_materializer.as_ref() else {
return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Spawn,
"working directory working_directory requested, but no materializer is configured for this runtime backend",
));
};
match materializer.bind_working_directory(
&working_directory.working_directory_id,
working_directory.relative_cwd.as_deref(),
}
for attachment in request.request.workdir_attachments.clone() {
if request.workdir_attachments.contains_key(&attachment.alias) {
let message = format!("duplicate Workdir attachment alias `{}`", attachment.alias);
if let Err(result) = rollback_materialized_spawn_workdirs(
materializer.as_ref(),
&request.workdir_attachments,
&rollback_workdirs,
&message,
) {
Ok(binding) => {
request.working_directory = Some(binding.clone());
Some(binding)
}
Err(error) => {
return WorkerExecutionSpawnResult::Rejected(
WorkerExecutionResult::rejected(
WorkerExecutionOperation::Spawn,
error.to_string(),
),
);
return result;
}
return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Spawn,
message,
));
}
match materializer.bind_working_directory(
&attachment.working_directory_id,
attachment.relative_cwd.as_deref(),
) {
Ok(binding) => {
request
.workdir_attachments
.insert(attachment.alias, binding);
}
Err(error) => {
let message = error.to_string();
if let Err(result) = rollback_materialized_spawn_workdirs(
materializer.as_ref(),
&request.workdir_attachments,
&rollback_workdirs,
&message,
) {
return result;
}
return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Spawn,
message,
));
}
}
(None, None) => None,
};
}
let workdir_attachments = request.workdir_attachments.clone();
let factory = self.factory.clone();
let bridge_context = request.context.clone();
@@ -2191,11 +2321,13 @@ where
let controller = match spawn_result {
Ok(controller) => controller,
Err(message) => {
if let (Some(materializer), Some(binding)) = (
self.working_directory_materializer.as_ref(),
rollback_working_directory.as_ref(),
if let Err(result) = rollback_materialized_spawn_workdirs(
materializer.as_ref(),
&workdir_attachments,
&rollback_workdirs,
&message,
) {
let _ = materializer.cleanup_working_directory(&binding.working_directory.id);
return result;
}
return WorkerExecutionSpawnResult::Errored(WorkerExecutionResult::errored(
WorkerExecutionOperation::Spawn,
@@ -2211,7 +2343,7 @@ where
controller.handle,
controller.shutdown,
controller.controller_task,
working_directory,
workdir_attachments,
Some(controller.workspace_client),
)
}
@@ -2232,46 +2364,56 @@ where
}
drop(workers);
if let Some(previous) = request.previous_working_directory.as_ref() {
if !request.previous_workdir_attachments.is_empty() {
let materializer = self
.working_directory_materializer
.as_ref()
.ok_or_else(|| {
WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
"Persisted Worker Workdir binding cannot be restored by this Runtime",
"Persisted Worker Workdir attachments cannot be restored by this Runtime",
)
})?;
let current = materializer
.working_directory_status(&previous.summary.working_directory_id)
.map_err(|message| {
WorkerExecutionResult::rejected(
for previous in &request.previous_workdir_attachments {
let current = materializer
.working_directory_status(
&previous.working_directory.summary.working_directory_id,
)
.map_err(|message| {
WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
format!(
"Persisted Worker Workdir attachment `{}` is unavailable: {message}",
previous.alias
),
)
})?;
if current.summary.status != WorkingDirectoryStatusKind::Active {
return Err(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
format!("Persisted Worker Workdir is unavailable: {message}"),
)
})?;
if current.summary.status != WorkingDirectoryStatusKind::Active {
return Err(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
"Persisted Worker Workdir is not active",
));
format!(
"Persisted Worker Workdir attachment `{}` is not active",
previous.alias
),
));
}
}
}
if request.previous_working_directory.is_none()
&& request.request.working_directory_request.is_some()
if request.previous_workdir_attachments.is_empty()
&& !request.request.workdir_attachment_requests.is_empty()
{
return Err(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
"Persisted Worker Workdir allocation is unavailable",
"Persisted Worker Workdir allocations are unavailable",
));
}
if request.previous_working_directory.is_none()
&& request.request.working_directory.is_some()
if request.previous_workdir_attachments.is_empty()
&& !request.request.workdir_attachments.is_empty()
&& self.working_directory_materializer.is_none()
{
return Err(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
"Persisted Worker Workdir claim cannot be restored by this Runtime",
"Persisted Worker Workdir claims cannot be restored by this Runtime",
));
}
if let Some(workspace_api) = request.request.workspace_api.as_ref() {
@@ -2321,58 +2463,40 @@ where
}
drop(workers);
let working_directory = match request.previous_working_directory.clone() {
Some(status) => {
let Some(materializer) = self.working_directory_materializer.as_ref() else {
return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
"persisted worker has a working directory binding, but no materializer is configured for this runtime backend",
));
};
let relative_cwd = request
.request
.working_directory
.as_ref()
.and_then(|working_directory| working_directory.relative_cwd.as_deref());
match materializer
.bind_working_directory(&status.summary.working_directory_id, relative_cwd)
{
Ok(binding) => {
request.working_directory = Some(binding.clone());
Some(binding)
}
Err(error) => {
return WorkerExecutionSpawnResult::Rejected(
WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
error.to_string(),
),
);
}
}
}
None if request.request.working_directory_request.is_some() => {
let mut workdir_attachments = BTreeMap::new();
if !request.previous_workdir_attachments.is_empty() {
let Some(materializer) = self.working_directory_materializer.as_ref() else {
return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
"persisted worker requested a working directory, but no persisted working directory binding is available to restore",
"persisted Worker has Workdir attachments, but no materializer is configured for this Runtime backend",
));
}
None if request.request.working_directory.is_some() => {
let Some(materializer) = self.working_directory_materializer.as_ref() else {
return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
"persisted worker has a working directory claim, but no materializer is configured for this runtime backend",
));
};
let working_directory =
request.request.working_directory.as_ref().expect("checked");
};
for status in request.previous_workdir_attachments.clone() {
let relative_cwd = request
.request
.workdir_attachments
.iter()
.find(|claim| claim.alias == status.alias)
.and_then(|claim| claim.relative_cwd.as_deref());
match materializer.bind_working_directory(
&working_directory.working_directory_id,
working_directory.relative_cwd.as_deref(),
&status.working_directory.summary.working_directory_id,
relative_cwd,
) {
Ok(binding) => {
request.working_directory = Some(binding.clone());
Some(binding)
if workdir_attachments
.insert(status.alias.clone(), binding)
.is_some()
{
return WorkerExecutionSpawnResult::Rejected(
WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
format!(
"duplicate persisted Workdir attachment alias `{}`",
status.alias
),
),
);
}
}
Err(error) => {
return WorkerExecutionSpawnResult::Rejected(
@@ -2384,8 +2508,43 @@ where
}
}
}
None => None,
};
} else if !request.request.workdir_attachments.is_empty() {
let Some(materializer) = self.working_directory_materializer.as_ref() else {
return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
"persisted Worker has Workdir claims, but no materializer is configured for this Runtime backend",
));
};
for claim in request.request.workdir_attachments.clone() {
match materializer.bind_working_directory(
&claim.working_directory_id,
claim.relative_cwd.as_deref(),
) {
Ok(binding) => {
if workdir_attachments
.insert(claim.alias.clone(), binding)
.is_some()
{
return WorkerExecutionSpawnResult::Rejected(
WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
format!("duplicate Workdir attachment alias `{}`", claim.alias),
),
);
}
}
Err(error) => {
return WorkerExecutionSpawnResult::Rejected(
WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
error.to_string(),
),
);
}
}
}
}
request.workdir_attachments = workdir_attachments.clone();
let factory = self.factory.clone();
let bridge_context = request.context.clone();
@@ -2405,7 +2564,7 @@ where
),
handle: None,
worker_state: None,
working_directory: None,
workdir_attachments: Vec::new(),
};
}
};
@@ -2417,7 +2576,7 @@ where
controller.handle,
controller.shutdown,
controller.controller_task,
working_directory,
workdir_attachments,
Some(controller.workspace_client),
)
}
@@ -2739,8 +2898,8 @@ mod tests {
use crate::Runtime as EmbeddedRuntime;
use crate::catalog::{
ConfigBundleRef, CreateWorkerRequest, MaterializerKind, ProfileSelector,
RepositorySelector, WorkingDirectoryClaim, WorkingDirectoryRepository,
WorkingDirectoryRequest, WorkspaceApiRef,
RepositorySelector, WorkingDirectoryAttachmentClaim, WorkingDirectoryAttachmentRequest,
WorkingDirectoryRepository, WorkingDirectoryRequest, WorkspaceApiRef,
};
use crate::execution::WorkerExecutionContext;
use crate::identity::WorkerId;
@@ -3279,18 +3438,17 @@ mod tests {
FsStore::new(&self.store_dir).map_err(|err| err.to_string())?,
FsWorkerStore::new(&self.worker_metadata_dir).map_err(|err| err.to_string())?,
);
let filesystem_authority = request
.working_directory
.as_ref()
let only_workdir = (request.workdir_attachments.len() == 1)
.then(|| request.workdir_attachments.values().next())
.flatten();
let filesystem_authority = only_workdir
.map(|binding| {
let cwd = binding.cwd().to_path_buf();
self.observed_cwds.lock().unwrap().push(cwd.clone());
WorkerFilesystemAuthority::local(binding.root().to_path_buf(), cwd)
})
.unwrap_or(WorkerFilesystemAuthority::None);
let scope_root = request
.working_directory
.as_ref()
let scope_root = only_workdir
.map(|binding| binding.root().to_path_buf())
.unwrap_or_else(|| self.cwd.clone());
let workspace_backend_ref = RuntimeWorkspaceBackendRef::from_worker_request(
@@ -3348,7 +3506,7 @@ mod tests {
request: request.request,
workspace_scope: request.workspace_scope,
context: request.context,
working_directory: request.working_directory,
workdir_attachments: request.workdir_attachments,
config_bundle: request.config_bundle,
};
self.spawn_controller(request).await
@@ -3513,8 +3671,8 @@ mod tests {
digest: bundle.metadata.digest,
}),
initial_input: None,
working_directory_request: None,
working_directory: None,
workdir_attachment_requests: Vec::new(),
workdir_attachments: Vec::new(),
worker_observation_enabled: false,
worker_observation_grants: Vec::new(),
workspace_api: None,
@@ -3595,6 +3753,7 @@ mod tests {
source_fingerprint: "sha256:test".to_string(),
selector: Some(RepositorySelector::from("HEAD")),
},
display_name: None,
materializer: MaterializerKind::RuntimeGitClone,
backend_workdir_id: None,
materialization: None,
@@ -3698,7 +3857,7 @@ mod tests {
request: create_request("1"),
workspace_scope: None,
context: test_execution_context(worker_ref),
working_directory: None,
workdir_attachments: BTreeMap::new(),
config_bundle: None,
};
@@ -3830,8 +3989,8 @@ mod tests {
"server-main",
)),
context: test_execution_context(worker_ref),
previous_working_directory: None,
working_directory: None,
previous_workdir_attachments: Vec::new(),
workdir_attachments: BTreeMap::new(),
config_bundle: None,
})
.await
@@ -3920,8 +4079,8 @@ mod tests {
request: create_request("embedded restore"),
workspace_scope: None,
context: test_execution_context(worker_ref),
previous_working_directory: None,
working_directory: None,
previous_workdir_attachments: Vec::new(),
workdir_attachments: BTreeMap::new(),
config_bundle: None,
})
.await
@@ -3988,7 +4147,7 @@ mod tests {
&[
"WorkerBootstrap::new(",
".prepare()",
"worker.bind_workdir_session(",
"worker.bind_workdir_sessions(",
"worker.bind_worker_observation_provider(",
"install_runtime_flow_transition_feature()",
"prepared.start()",
@@ -3998,7 +4157,7 @@ mod tests {
restore,
&[
"Worker::restore_from_worker_metadata_with_context(",
"worker.bind_workdir_session(",
"worker.bind_workdir_sessions(",
"worker.bind_worker_observation_provider(",
"install_runtime_flow_transition_feature()",
"PreparedWorker::new(",
@@ -4354,7 +4513,10 @@ mod tests {
.unwrap();
runtime.store_config_bundle(test_bundle()).unwrap();
let mut request = create_request("chat");
request.working_directory_request = Some(working_directory_request(repo.path()));
request.workdir_attachment_requests = vec![WorkingDirectoryAttachmentRequest {
alias: workdir::WorkdirAttachmentAlias::new("checkout").unwrap(),
working_directory: working_directory_request(repo.path()),
}];
let detail = runtime.create_worker(request).unwrap();
runtime
@@ -4376,7 +4538,7 @@ mod tests {
);
}
assert!(detail.working_directory.is_some());
assert_eq!(detail.workdir_attachments.len(), 1);
let cwds = observed_cwds.lock().unwrap();
assert_eq!(cwds.len(), 1);
let cwd = &cwds[0];
@@ -4387,6 +4549,71 @@ mod tests {
observed_workspace_clients.lock().unwrap().as_slice(),
&[("unavailable".to_string(), None, false)]
);
let workdir_id = detail.workdir_attachments[0]
.working_directory
.summary
.working_directory_id
.clone();
let second = runtime.create_worker(create_request("second")).unwrap();
let conflict = runtime
.replace_worker_workdir_attachments(
&second.worker_ref,
vec![WorkingDirectoryAttachmentClaim {
alias: workdir::WorkdirAttachmentAlias::new("conflict").unwrap(),
working_directory_id: workdir_id.clone(),
relative_cwd: None,
}],
)
.unwrap_err();
assert!(matches!(
conflict,
crate::error::RuntimeError::InvalidRequest(_)
));
let inactive = runtime
.create_working_directory(working_directory_request(repo.path()))
.unwrap();
let inactive_id = inactive.summary.working_directory_id;
runtime.cleanup_working_directory(&inactive_id).unwrap();
let inactive_error = runtime
.replace_worker_workdir_attachments(
&second.worker_ref,
vec![WorkingDirectoryAttachmentClaim {
alias: workdir::WorkdirAttachmentAlias::new("inactive").unwrap(),
working_directory_id: inactive_id,
relative_cwd: None,
}],
)
.unwrap_err();
assert!(matches!(
inactive_error,
crate::error::RuntimeError::InvalidRequest(_)
| crate::error::RuntimeError::WorkingDirectory(_)
));
let detached = runtime
.replace_worker_workdir_attachments(&detail.worker_ref, Vec::new())
.unwrap();
assert!(detached.workdir_attachments.is_empty());
runtime.stop_worker(&detail.worker_ref, None).unwrap();
let restored_without_workdir = runtime.restore_worker(&detail.worker_ref).unwrap();
assert!(restored_without_workdir.workdir_attachments.is_empty());
let alias = workdir::WorkdirAttachmentAlias::new("restored").unwrap();
runtime
.replace_worker_workdir_attachments(
&detail.worker_ref,
vec![WorkingDirectoryAttachmentClaim {
alias: alias.clone(),
working_directory_id: workdir_id,
relative_cwd: None,
}],
)
.unwrap();
runtime.stop_worker(&detail.worker_ref, None).unwrap();
let stopped = runtime.worker_detail(&detail.worker_ref).unwrap();
assert_eq!(stopped.workdir_attachments.len(), 1);
assert_eq!(stopped.workdir_attachments[0].alias, alias);
}
#[test]
@@ -4620,12 +4847,13 @@ mod tests {
.unwrap();
runtime.store_config_bundle(test_bundle()).unwrap();
let mut request = create_request("chat");
request.working_directory_request = Some(working_directory_request(repo.path()));
request.workdir_attachment_requests = vec![WorkingDirectoryAttachmentRequest {
alias: workdir::WorkdirAttachmentAlias::new("checkout").unwrap(),
working_directory: working_directory_request(repo.path()),
}];
let detail = runtime.create_worker(request).unwrap();
let workdir_id = detail
let workdir_id = detail.workdir_attachments[0]
.working_directory
.as_ref()
.unwrap()
.summary
.working_directory_id
.clone();
@@ -4642,7 +4870,6 @@ mod tests {
crate::catalog::WorkingDirectoryStatusKind::Active
);
assert_eq!(status.summary.cleanliness.as_deref(), Some("clean"));
assert_eq!(status.summary.primary_worker_id, None);
}
#[test]
@@ -4663,10 +4890,11 @@ mod tests {
let clone_root = materialized_clone_root(runtime_base.path(), &workdir_id);
assert!(clone_root.join("README.md").exists());
let mut request = create_request("chat");
request.working_directory = Some(WorkingDirectoryClaim {
request.workdir_attachments = vec![WorkingDirectoryAttachmentClaim {
alias: workdir::WorkdirAttachmentAlias::new("workdir").unwrap(),
working_directory_id: workdir_id.clone(),
relative_cwd: None,
});
}];
let error = runtime.create_worker(request).unwrap_err();
@@ -4691,7 +4919,17 @@ mod tests {
.unwrap();
runtime.store_config_bundle(test_bundle()).unwrap();
let mut request = create_request("chat");
request.working_directory_request = Some(working_directory_request(repo.path()));
let second_repo = create_clean_repo();
request.workdir_attachment_requests = vec![
WorkingDirectoryAttachmentRequest {
alias: workdir::WorkdirAttachmentAlias::new("checkout").unwrap(),
working_directory: working_directory_request(repo.path()),
},
WorkingDirectoryAttachmentRequest {
alias: workdir::WorkdirAttachmentAlias::new("docs").unwrap(),
working_directory: working_directory_request(second_repo.path()),
},
];
let error = runtime.create_worker(request).unwrap_err();
@@ -4708,4 +4946,41 @@ mod tests {
assert_eq!(remaining_workdirs, 0);
assert!(!working_directories_root.join(".repository-cache").exists());
}
#[test]
fn second_materialization_failure_rolls_back_first_attachment() {
let runtime_base = tempfile::tempdir().unwrap();
let repo = create_clean_repo();
let missing_repo = runtime_base.path().join("missing-repository");
let backend = WorkerRuntimeExecutionBackend::new(FailingFactory)
.unwrap()
.with_working_directory_materializer(RuntimeGitMaterializer::new(runtime_base.path()));
let runtime =
EmbeddedRuntime::with_execution_backend(RuntimeOptions::default(), Arc::new(backend))
.unwrap();
runtime.store_config_bundle(test_bundle()).unwrap();
let mut request = create_request("chat");
request.workdir_attachment_requests = vec![
WorkingDirectoryAttachmentRequest {
alias: workdir::WorkdirAttachmentAlias::new("checkout").unwrap(),
working_directory: working_directory_request(repo.path()),
},
WorkingDirectoryAttachmentRequest {
alias: workdir::WorkdirAttachmentAlias::new("docs").unwrap(),
working_directory: working_directory_request(&missing_repo),
},
];
let error = runtime.create_worker(request).unwrap_err();
assert!(format!("{error:?}").contains("repository"));
let remaining_workdirs = fs::read_dir(runtime_base.path())
.map(|entries| {
entries
.flatten()
.filter(|entry| !entry.file_name().to_string_lossy().starts_with('.'))
.count()
})
.unwrap_or(0);
assert_eq!(remaining_workdirs, 0);
}
}
@@ -57,6 +57,8 @@ pub struct WorkingDirectoryEvidence {
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct WorkingDirectory {
pub id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub display_name: Option<String>,
pub repository_id: String,
pub materializer_kind: MaterializerKind,
pub evidence: WorkingDirectoryEvidence,
@@ -68,6 +70,7 @@ impl WorkingDirectory {
pub fn status_summary(&self) -> WorkingDirectorySummary {
WorkingDirectorySummary {
working_directory_id: self.id.clone(),
display_name: self.display_name.clone(),
repository_id: self.repository_id.clone(),
creation_selector: self.evidence.requested_selector.clone(),
creation_ref: Some(self.evidence.resolved_commit.clone()),
@@ -80,7 +83,6 @@ impl WorkingDirectory {
cleanup_target: Some(self.cleanup_target.clone()),
status: self.status.clone(),
cleanliness: None,
primary_worker_id: None,
occupied_by: None,
}
}
@@ -717,6 +719,7 @@ impl RuntimeGitMaterializer {
WorkingDirectoryStatus {
summary: WorkingDirectorySummary {
working_directory_id: working_directory_id.to_string(),
display_name: None,
repository_id: "unknown".to_string(),
creation_selector: None,
creation_ref: None,
@@ -733,7 +736,6 @@ impl RuntimeGitMaterializer {
}),
status: WorkingDirectoryStatusKind::Corrupted,
cleanliness: Some("unknown".to_string()),
primary_worker_id: None,
occupied_by: None,
},
}
@@ -1072,6 +1074,7 @@ impl RuntimeGitMaterializer {
let context = request.materialization.as_ref();
let working_directory = WorkingDirectory {
id: working_directory_id.clone(),
display_name: request.display_name.clone(),
repository_id: request.repository.id.clone(),
materializer_kind: MaterializerKind::RuntimeGitClone,
evidence: WorkingDirectoryEvidence {
@@ -1211,6 +1214,7 @@ impl WorkingDirectoryMaterializer for RuntimeGitMaterializer {
let working_request = WorkingDirectoryRequest {
repository: request.repository.clone(),
display_name: None,
materializer: MaterializerKind::RuntimeGitClone,
backend_workdir_id: None,
materialization: request.materialization.clone(),
@@ -3316,6 +3320,7 @@ mod tests {
source_fingerprint: "sha256:test".to_string(),
selector: Some(RepositorySelector::from("HEAD")),
},
display_name: None,
materializer: MaterializerKind::RuntimeGitClone,
backend_workdir_id: None,
materialization: None,
+4 -7
View File
@@ -145,7 +145,7 @@ where
.map_err(WorkerBootstrapError::Worker)?;
if let Some(workdir_session) = self.workdir_session {
worker.bind_workdir_session(Some(workdir_session));
worker.bind_single_workdir_session(Some(workdir_session));
}
Ok(PreparedWorker::new(worker, self.layout, self.transport))
}
@@ -168,7 +168,7 @@ where
.map_err(WorkerBootstrapError::Worker)?;
if let Some(workdir_session) = self.workdir_session {
worker.bind_workdir_session(Some(workdir_session));
worker.bind_single_workdir_session(Some(workdir_session));
}
Ok(PreparedWorker::new(worker, self.layout, self.transport))
}
@@ -222,7 +222,7 @@ where
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
{
let cleanup_session = worker.workdir_session().cloned();
let cleanup_sessions = worker.workdir_sessions();
let controller = match layout {
WorkerBootstrapLayout::Direct {
runtime_base,
@@ -257,10 +257,7 @@ where
controller_task,
}),
Err(source) => {
let cleanup_failed = match cleanup_session {
Some(session) => session.close().await.is_err(),
None => false,
};
let cleanup_failed = cleanup_sessions.close_all().await.is_err();
Err(WorkerBootstrapError::Controller {
source,
cleanup_failed,
+4 -5
View File
@@ -2,9 +2,8 @@
//! Engine callbacks (which run synchronously and cannot themselves
//! perform `async` store writes).
//!
//! Worker drains this buffer in `persist_turn` and writes each metric via
//! `session_metrics::record_metric`, alongside the regular `LlmUsage`
//! entries.
//! Worker drains this buffer at the next successful request accounting
//! boundary, with terminal `persist_turn` as a fallback.
use std::sync::Mutex;
@@ -21,12 +20,12 @@ impl MetricsTracker {
}
}
/// Queue a metric for the next `persist_turn` flush.
/// Queue a metric for the next request or terminal accounting flush.
pub(crate) fn push(&self, metric: Metric) {
self.pending.lock().unwrap().push(metric);
}
/// Drain all queued metrics. Called by Worker after a run completes.
/// Drain all queued metrics at an accounting boundary.
pub(crate) fn drain(&self) -> Vec<Metric> {
std::mem::take(&mut *self.pending.lock().unwrap())
}
+340 -40
View File
@@ -7,17 +7,24 @@
//! - **`on_usage` callback** (sync closure): receives the aggregated final
//! `UsageEvent` for that request after the stream completes.
//!
//! Pairing the two yields one `UsageRecord` per LLM call. Worker drains them
//! in `persist_turn` and writes them as `LogEntry::LlmUsage` entries.
//! Pairing the two yields one `UsageRecord` per LLM call. Worker commits
//! successful requests from the assistant-turn boundary and uses terminal
//! `persist_turn` only for records that have not already crossed that boundary.
//!
//! Multiple LLM calls per Worker run (tool loop) are supported: each call
//! produces its own `(history_len, UsageEvent)` pair, and the records are
//! buffered in chronological order.
use std::collections::VecDeque;
use std::sync::Mutex;
use agen::UsageRecord;
use agen::timeline::event::UsageEvent;
use session_metrics::Metric;
use session_store::{LogEntry, StoreError, segment_log};
use crate::compact::state::CompactState;
use crate::compact::telemetry::correlated_post_request_metric;
/// The metric emitted after the next measured provider request.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -41,11 +48,15 @@ pub(crate) struct PostRequestLink {
pub(crate) metric: PostRequestMetric,
}
/// One drained measurement and its causal metric links.
/// One measured request and its causal metric links.
#[derive(Debug, Clone)]
pub(crate) struct RecordedUsage {
pub(crate) record: UsageRecord,
pub(crate) post_requests: Vec<PostRequestLink>,
/// True only when the provider stream completed normally and reported an
/// input occupancy. Partial/error/cancelled streams remain billable usage
/// records, but cannot satisfy a post-compaction rearm.
pub(crate) completed_with_occupancy: bool,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
@@ -56,29 +67,52 @@ pub(crate) struct UsageSnapshot {
pub(crate) output_tokens: u64,
}
#[derive(Debug)]
struct PendingRequest {
id: u64,
history_len: usize,
usage_recorded: bool,
}
#[derive(Debug)]
struct PendingRecordedUsage {
request_id: u64,
has_input_occupancy: bool,
usage: RecordedUsage,
}
#[derive(Debug, Default)]
struct UsageTrackerState {
next_request_id: u64,
pending_request: Option<PendingRequest>,
pending_correlations: Vec<PostRequestLink>,
pending_records: VecDeque<PendingRecordedUsage>,
}
/// Shared between the pre-request hook, the `on_usage` callback, and Worker.
pub(crate) struct UsageTracker {
/// `history.len()` captured at the most recent `pre_llm_request`.
/// Cleared when paired with an incoming `on_usage` event.
pending_history_len: Mutex<Option<usize>>,
/// Optional causal link consumed by the next measured request.
pending_correlations: Mutex<Vec<PostRequestLink>>,
/// Records accumulated during the current run; drained by Worker.
pending_records: Mutex<Vec<RecordedUsage>>,
state: Mutex<UsageTrackerState>,
}
impl UsageTracker {
pub(crate) fn new() -> Self {
Self {
pending_history_len: Mutex::new(None),
pending_correlations: Mutex::new(Vec::new()),
pending_records: Mutex::new(Vec::new()),
state: Mutex::new(UsageTrackerState::default()),
}
}
/// Called from a `pre_llm_request` hook with the current history length.
/// Replacing an older pending request classifies that request as incomplete;
/// any partial usage already recorded for it remains available for billing.
pub(crate) fn note_request(&self, history_len: usize) {
*self.pending_history_len.lock().unwrap() = Some(history_len);
let mut state = self.state.lock().unwrap();
let id = state.next_request_id;
state.next_request_id = state.next_request_id.wrapping_add(1);
state.pending_request = Some(PendingRequest {
id,
history_len,
usage_recorded: false,
});
}
/// Pair a prune event with the next provider request.
@@ -92,25 +126,31 @@ impl UsageTracker {
}
fn note_post_request(&self, id: String, metric: PostRequestMetric) {
let mut pending = self.pending_correlations.lock().unwrap();
pending.retain(|link| link.metric != metric);
pending.push(PostRequestLink {
let mut state = self.state.lock().unwrap();
state
.pending_correlations
.retain(|link| link.metric != metric);
state.pending_correlations.push(PostRequestLink {
correlation_id: id,
metric,
});
}
/// Called from the `on_usage` callback with the aggregated final
/// UsageEvent. If a `history_len` was previously stashed via
/// `note_request`, builds a `RecordedUsage` and pushes it onto the
/// buffer. If not (e.g. test code that fires Usage outside a request),
/// drops the event.
/// UsageEvent. The record remains associated with the current request until
/// [`Self::request_completed`] proves that the provider stream completed
/// normally. Error/cancel paths still retain the measurement for billing.
pub(crate) fn record_usage(&self, event: &UsageEvent) {
let history_len = match self.pending_history_len.lock().unwrap().take() {
Some(n) => n,
None => return,
let mut state = self.state.lock().unwrap();
let Some(pending) = state.pending_request.as_mut() else {
return;
};
let post_requests = std::mem::take(&mut *self.pending_correlations.lock().unwrap());
if pending.usage_recorded {
return;
}
pending.usage_recorded = true;
let request_id = pending.id;
let history_len = pending.history_len;
// UsageEvent.input_tokens は scheme 層で「占有量(プロンプト全長)」に
// 正規化済みである前提(Anthropic は cache_read + cache_creation を
// 加算して emit する)。
@@ -118,35 +158,166 @@ impl UsageTracker {
let cache_read = event.cache_read_input_tokens.unwrap_or(0);
let cache_write = event.cache_creation_input_tokens.unwrap_or(0);
let output = event.output_tokens.unwrap_or(0);
self.pending_records.lock().unwrap().push(RecordedUsage {
record: UsageRecord {
history_len,
input_total_tokens: input_total,
cache_read_tokens: cache_read,
cache_write_tokens: cache_write,
output_tokens: output,
state.pending_records.push_back(PendingRecordedUsage {
request_id,
has_input_occupancy: event.input_tokens.is_some(),
usage: RecordedUsage {
record: UsageRecord {
history_len,
input_total_tokens: input_total,
cache_read_tokens: cache_read,
cache_write_tokens: cache_write,
output_tokens: output,
},
post_requests: Vec::new(),
completed_with_occupancy: false,
},
post_requests,
});
}
/// Return a clone of the accumulated `UsageRecord`s without clearing them.
/// Mark the current request as normally completed. A post-request causal
/// link is consumed only when that exact request also reported a valid
/// input occupancy; callbacks from partial/error/cancelled streams cannot
/// rearm the compaction guard.
pub(crate) fn request_completed(&self) {
let mut state = self.state.lock().unwrap();
let Some(pending) = state.pending_request.take() else {
return;
};
let Some(recorded_index) = state
.pending_records
.iter()
.rposition(|recorded| recorded.request_id == pending.id)
else {
return;
};
let has_occupancy = state.pending_records[recorded_index].has_input_occupancy;
if !has_occupancy {
return;
}
let post_requests = std::mem::take(&mut state.pending_correlations);
let recorded = &mut state.pending_records[recorded_index].usage;
recorded.completed_with_occupancy = true;
recorded.post_requests = post_requests;
}
/// Return a clone of accumulated measurements without clearing them.
/// Used by request-time circuit breakers that need the same occupancy
/// projection as Worker persistence while the run is still active.
pub(crate) fn records(&self) -> Vec<UsageRecord> {
self.pending_records
self.state
.lock()
.unwrap()
.pending_records
.iter()
.map(|r| r.record.clone())
.map(|r| r.usage.record.clone())
.collect()
}
/// Drain accumulated records. Called by Worker after a run completes,
/// before persisting the turn.
pub(crate) fn drain(&self) -> Vec<RecordedUsage> {
std::mem::take(&mut *self.pending_records.lock().unwrap())
/// Number of measurements that have not crossed the durable commit
/// boundary. Session rewrites must not activate a replacement Segment while
/// this is non-zero, because their history coordinates belong to the current
/// Segment.
pub(crate) fn pending_record_count(&self) -> usize {
self.state.lock().unwrap().pending_records.len()
}
/// Remove the oldest measurement for one durable commit attempt.
pub(crate) fn take_next(&self) -> Option<RecordedUsage> {
self.state
.lock()
.unwrap()
.pending_records
.pop_front()
.map(|recorded| recorded.usage)
}
/// Restore a measurement after its durable commit failed.
pub(crate) fn restore_front(&self, usage: RecordedUsage) {
self.state
.lock()
.unwrap()
.pending_records
.push_front(PendingRecordedUsage {
// Persist ordering is the only remaining identity requirement.
request_id: u64::MAX,
has_input_occupancy: usage.completed_with_occupancy,
usage,
});
}
/// Drain accumulated records at a terminal Worker boundary.
pub(crate) fn drain(&self) -> Vec<RecordedUsage> {
let mut state = self.state.lock().unwrap();
state.pending_request = None;
state
.pending_records
.drain(..)
.map(|recorded| recorded.usage)
.collect()
}
}
/// Commit all currently measured usage records in request order.
///
/// A record is removed from the tracker only after its `LlmUsage` append
/// succeeds. This is the single accounting boundary shared by the in-run
/// interceptor and terminal `persist_turn`, preventing both loss on append
/// failure and duplicate terminal writes. Correlated metrics are handed to the
/// same commit boundary so they stay on the usage record's Segment; the return
/// value is only useful as test/diagnostic evidence and must not be re-emitted.
pub(crate) fn persist_pending_usage(
tracker: &UsageTracker,
usage_history: &Mutex<Vec<UsageRecord>>,
compact_state: Option<&CompactState>,
mut commit: impl FnMut(LogEntry, &mut dyn FnMut(), &[Metric]) -> Result<(), StoreError>,
) -> Result<Vec<Metric>, StoreError> {
let mut metrics = Vec::new();
while let Some(recorded) = tracker.take_next() {
let RecordedUsage {
record,
post_requests,
completed_with_occupancy,
} = recorded;
let post_metrics = post_requests
.iter()
.map(|link| correlated_post_request_metric(link.metric, &link.correlation_id, &record))
.collect::<Vec<_>>();
let usage_entry = LogEntry::LlmUsage {
ts: segment_log::now_millis(),
history_len: record.history_len,
input_total_tokens: record.input_total_tokens,
cache_read_tokens: record.cache_read_tokens,
cache_write_tokens: record.cache_write_tokens,
output_tokens: record.output_tokens,
};
let mut publish_committed_usage = || {
usage_history
.lock()
.expect("usage_history poisoned")
.push(record.clone());
if completed_with_occupancy
&& post_requests
.iter()
.any(|link| link.metric == PostRequestMetric::Compaction)
{
if let Some(state) = compact_state {
state.post_compact_request_committed();
}
}
};
if let Err(error) = commit(usage_entry, &mut publish_committed_usage, &post_metrics) {
tracker.restore_front(RecordedUsage {
record,
post_requests,
completed_with_occupancy,
});
return Err(error);
}
metrics.extend(post_metrics);
}
Ok(metrics)
}
#[cfg(test)]
@@ -230,6 +401,7 @@ mod tests {
tracker.note_correlation_id("prune-id".into());
tracker.note_request(5);
tracker.record_usage(&make_event(100, 10, 2, 20));
tracker.request_completed();
let records = tracker.drain();
assert_eq!(records[0].post_requests.len(), 2);
@@ -248,9 +420,11 @@ mod tests {
tracker.note_correlation_id("abc".into());
tracker.note_request(5);
tracker.record_usage(&make_event(100, 0, 0, 20));
tracker.request_completed();
// Next request without a fresh stash → no correlation_id.
tracker.note_request(10);
tracker.record_usage(&make_event(200, 50, 0, 30));
tracker.request_completed();
let records = tracker.drain();
assert_eq!(records.len(), 2);
@@ -259,4 +433,130 @@ mod tests {
assert_eq!(records[0].post_requests[0].metric, PostRequestMetric::Prune);
assert!(records[1].post_requests.is_empty());
}
#[test]
fn partial_usage_does_not_consume_post_compact_link() {
let tracker = UsageTracker::new();
tracker.note_compaction_correlation_id("compact-id".into());
tracker.note_request(5);
tracker.record_usage(&make_event(90, 0, 0, 2));
// Starting another request proves the prior stream did not complete.
tracker.note_request(6);
tracker.record_usage(&make_event(100, 0, 0, 3));
tracker.request_completed();
let records = tracker.drain();
assert_eq!(records.len(), 2);
assert!(!records[0].completed_with_occupancy);
assert!(records[0].post_requests.is_empty());
assert!(records[1].completed_with_occupancy);
assert_eq!(records[1].post_requests.len(), 1);
assert_eq!(records[1].post_requests[0].correlation_id, "compact-id");
}
#[test]
fn completed_request_without_input_occupancy_does_not_consume_link() {
let tracker = UsageTracker::new();
tracker.note_compaction_correlation_id("compact-id".into());
tracker.note_request(5);
let mut missing_input = make_event(0, 0, 0, 2);
missing_input.input_tokens = None;
tracker.record_usage(&missing_input);
tracker.request_completed();
tracker.note_request(6);
tracker.record_usage(&make_event(100, 0, 0, 3));
tracker.request_completed();
let records = tracker.drain();
assert!(!records[0].completed_with_occupancy);
assert!(records[0].post_requests.is_empty());
assert!(records[1].completed_with_occupancy);
assert_eq!(records[1].post_requests[0].correlation_id, "compact-id");
}
#[test]
fn completed_request_without_usage_does_not_complete_previous_partial_record() {
let tracker = UsageTracker::new();
tracker.note_compaction_correlation_id("compact-id".into());
tracker.note_request(5);
tracker.record_usage(&make_event(90, 0, 0, 2));
tracker.note_request(6);
tracker.request_completed();
let records = tracker.drain();
assert_eq!(records.len(), 1);
assert!(!records[0].completed_with_occupancy);
assert!(records[0].post_requests.is_empty());
}
#[test]
fn failed_commit_can_restore_record_without_reordering() {
let tracker = UsageTracker::new();
tracker.note_request(1);
tracker.record_usage(&make_event(10, 0, 0, 1));
tracker.request_completed();
tracker.note_request(2);
tracker.record_usage(&make_event(20, 0, 0, 2));
tracker.request_completed();
let first = tracker.take_next().expect("first record");
tracker.restore_front(first);
let records = tracker.drain();
assert_eq!(records.len(), 2);
assert_eq!(records[0].record.input_total_tokens, 10);
assert_eq!(records[1].record.input_total_tokens, 20);
}
#[test]
fn compaction_rearms_only_after_usage_append_succeeds() {
use crate::compact::state::{AutomaticCompactDecision, CompactionOutcome};
let tracker = UsageTracker::new();
let state = CompactState::new(None, Some(10), 0);
assert!(matches!(
state.evaluate_request(11),
AutomaticCompactDecision::Start(_)
));
assert!(state.complete_automatic(CompactionOutcome::Succeeded));
tracker.note_compaction_correlation_id("compact-id".into());
tracker.note_request(1);
tracker.record_usage(&make_event(5, 0, 0, 1));
tracker.request_completed();
let usage_history = Mutex::new(Vec::new());
let error = persist_pending_usage(&tracker, &usage_history, Some(&state), |_, _, _| {
Err(StoreError::Io(std::io::Error::other("synthetic failure")))
})
.expect_err("usage append must fail");
assert!(error.to_string().contains("synthetic failure"));
assert!(usage_history.lock().unwrap().is_empty());
assert!(matches!(
state.evaluate_request(11),
AutomaticCompactDecision::Block(_)
));
let mut entries = Vec::new();
let metrics = persist_pending_usage(
&tracker,
&usage_history,
Some(&state),
|entry, after_commit, _post_metrics| {
entries.push(entry);
after_commit();
Ok(())
},
)
.expect("retry should commit");
assert_eq!(entries.len(), 1);
assert_eq!(usage_history.lock().unwrap().len(), 1);
assert_eq!(metrics.len(), 1);
assert!(matches!(
state.evaluate_request(11),
AutomaticCompactDecision::Start(_)
));
}
}
+41 -56
View File
@@ -617,7 +617,7 @@ impl WorkerController {
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
{
let session = worker.workdir_session().cloned();
let sessions = worker.workdir_sessions();
let result = Self::spawn_initialized(
worker,
runtime_base,
@@ -628,10 +628,9 @@ impl WorkerController {
)
.await;
if result.is_err()
&& let Some(session) = session
&& let Err(error) = session.close().await
&& let Err(error) = sessions.close_all().await
{
tracing::warn!(%error, "Workdir session close after controller startup failure failed");
tracing::warn!(%error, "Workdir sessions close after controller startup failure failed");
}
result
}
@@ -1144,7 +1143,7 @@ pub(crate) async fn register_worker_tools<C, St>(
runtime_base: PathBuf,
spawned_registry: Arc<SpawnedWorkerRegistry>,
parent_method_tx: Option<mpsc::WeakSender<Method>>,
inherited_workdir_tool_broker: Option<workdir::WorkdirToolBroker>,
inherited_workdir_tool_broker: Option<workdir::WorkdirToolBrokerRouter>,
) -> std::io::Result<Option<workdir::WorkdirSessionHandle>>
where
C: LlmClient + Clone + 'static,
@@ -1153,26 +1152,13 @@ where
// Worker-immutable snapshots taken before the mutable worker borrow
// below so the worker borrow doesn't conflict with reads on `worker`.
let feature_config = worker.manifest().feature.clone();
let mut workdir_tool_broker = inherited_workdir_tool_broker;
if feature_config.manage_workdir.enabled && worker.workdir_session().is_none() {
let workspace_client = worker.workspace_client_handle();
let broker = workdir::WorkdirToolBroker::new(
crate::feature::builtin::manage_workdir::WorkspaceAttachedWorkdirSession::handle(
workspace_client,
),
);
worker.bind_workdir_session(Some(broker.tool_session()));
workdir_tool_broker = Some(broker);
} else if workdir_tool_broker.is_none()
&& let Some(existing) = worker.workdir_session().cloned()
{
let broker = workdir::WorkdirToolBroker::new(existing);
worker.bind_workdir_session(Some(broker.tool_session()));
workdir_tool_broker = Some(broker);
}
let worker_workdir = workdir_tool_broker
.as_ref()
.map(workdir::WorkdirToolBroker::tool_session);
let workdir_sessions = worker.workdir_sessions();
let restrict_workdir_tools_to_current_capabilities = inherited_workdir_tool_broker.is_some();
let workdir_tool_broker = Some(
inherited_workdir_tool_broker
.unwrap_or_else(|| workdir::WorkdirToolBrokerRouter::new(workdir_sessions.clone())),
);
let workdir_for_view = workdir_sessions.only_session();
let local_filesystem = worker.local_working_directory().cloned();
let local_workspace_root = local_filesystem.as_ref().map(|local| local.root.clone());
let task_feature = worker.task_feature();
@@ -1207,28 +1193,29 @@ where
paste_store,
paste_session_id,
));
// Resolve the existing Worker–Workdir 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.
let (workdir_for_view, tracker) = if let Some(workdir) = worker_workdir {
let tracker = tools::Tracker::new();
// Keep the alias-routed schemas installed for the Worker lifetime. The
// attachment set may be empty or mutate after registration, so target
// resolution and provider capability checks happen per invocation.
let tracker = tools::Tracker::new();
let workdir_tools = if restrict_workdir_tools_to_current_capabilities {
tools::routed_builtin_tools_for_current_capabilities(
workdir_sessions.clone(),
tracker.clone(),
bash_output_dir.clone(),
)
} else {
tools::routed_builtin_tools(
workdir_sessions.clone(),
tracker.clone(),
bash_output_dir.clone(),
)
};
worker.engine_mut().register_tools(workdir_tools);
if feature_config.image.enabled && model_supports_image_attachments(&spawner_manifest.model) {
worker
.engine_mut()
.register_tools(tools::core_builtin_tools(
workdir.clone(),
tracker.clone(),
bash_output_dir.clone(),
));
if feature_config.image.enabled && model_supports_image_attachments(&spawner_manifest.model)
{
worker
.engine_mut()
.register_tool(tools::view_image_tool(workdir.clone()));
}
(Some(workdir), Some(tracker))
} else {
(None, None)
};
.register_tool(tools::routed_view_image_tool(workdir_sessions));
}
if feature_config.web.enabled {
worker
.engine_mut()
@@ -1344,6 +1331,7 @@ where
feature_registry.add_module(
crate::feature::builtin::manage_workdir::ManageWorkdirFeature::with_child_lifecycle(
workspace_client,
worker.workdir_sessions(),
Arc::new(move || {
let child_registry = shutdown_registry.clone();
Box::pin(async move { child_registry.shutdown_internal().await })
@@ -1485,9 +1473,7 @@ where
if let Some((resident_summary, system_prompt_override)) = memory_prompt_contribution {
worker.install_system_prompt_contribution(resident_summary, system_prompt_override);
}
if let Some(tracker) = tracker {
worker.attach_tracker(tracker);
}
worker.attach_tracker(tracker);
Ok(workdir_for_view)
}
@@ -2610,14 +2596,13 @@ async fn controller_loop<C, St>(
}
}
if let Some(session) = worker.workdir_session() {
loop {
match session.close().await {
Ok(()) => break,
Err(error) => {
tracing::warn!(%error, "Workdir session close failed; retrying");
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
}
let sessions = worker.workdir_sessions();
while !sessions.is_empty() {
match sessions.close_all().await {
Ok(()) => break,
Err(error) => {
tracing::warn!(%error, "Workdir session set close failed; retrying");
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
}
}
}
@@ -52,9 +52,9 @@ const DETACH_TOOL: &str = "WorkdirDetach";
const DELETE_TOOL: &str = "WorkdirDelete";
const LIST_DESCRIPTION: &str = "List persistent Workdirs in the current Workspace through Backend Workspace API authority. The result contains safe summaries and diagnostics, never host paths or Runtime connection details.";
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 CREATE_DESCRIPTION: &str = "Materialize a persistent Workdir with an optional human-facing display name on a selected Runtime from a Workspace repository and optional selector. The display name is independent of attachment aliases. This does not change this Worker's attachment; use WorkdirAttach explicitly after creation.";
const ATTACH_DESCRIPTION: &str = "Attach one existing Workdir to this Worker under a stable Worker-local alias. The Backend preserves one active Worker per Workdir while allowing this Worker to own multiple aliases.";
const DETACH_DESCRIPTION: &str = "Detach one Workdir attachment by its Worker-local alias. Active operations, commands, or Internal SubWorkers cause safe rejection rather than implicit authority loss.";
pub(crate) type BeforeWorkdirRelease =
Arc<dyn Fn() -> Pin<Box<dyn Future<Output = std::io::Result<()>> + Send>> + Send + Sync>;
pub(crate) type AfterWorkdirAttach = Arc<dyn Fn() + Send + Sync>;
@@ -64,6 +64,7 @@ const DELETE_DESCRIPTION: &str = "Request removal of one persistent Workdir by i
#[derive(Clone)]
pub struct ManageWorkdirFeature {
client: Arc<dyn WorkspaceClient>,
session_router: Arc<workdir::WorkdirSessionRouter>,
before_workdir_release: Option<BeforeWorkdirRelease>,
after_workdir_attach: Option<AfterWorkdirAttach>,
}
@@ -82,6 +83,7 @@ impl ManageWorkdirFeature {
pub fn new(client: Arc<dyn WorkspaceClient>) -> Self {
Self {
client,
session_router: Arc::new(workdir::WorkdirSessionRouter::new()),
before_workdir_release: None,
after_workdir_attach: None,
}
@@ -89,11 +91,13 @@ impl ManageWorkdirFeature {
pub(crate) fn with_child_lifecycle(
client: Arc<dyn WorkspaceClient>,
session_router: Arc<workdir::WorkdirSessionRouter>,
before_workdir_release: BeforeWorkdirRelease,
after_workdir_attach: AfterWorkdirAttach,
) -> Self {
Self {
client,
session_router,
before_workdir_release: Some(before_workdir_release),
after_workdir_attach: Some(after_workdir_attach),
}
@@ -116,10 +120,12 @@ impl FeatureModule for ManageWorkdirFeature {
}
fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> {
let backend = WorkspaceHttpWorkdirBackend::new(self.client.clone()).with_child_lifecycle(
self.before_workdir_release.clone(),
self.after_workdir_attach.clone(),
);
let backend = WorkspaceHttpWorkdirBackend::new(self.client.clone())
.with_session_router(self.session_router.clone())
.with_child_lifecycle(
self.before_workdir_release.clone(),
self.after_workdir_attach.clone(),
);
for (name, definition) in [
(
LIST_TOOL,
@@ -183,6 +189,8 @@ impl FeatureModule for ManageWorkdirFeature {
#[derive(Clone)]
struct WorkspaceHttpWorkdirBackend {
client: Arc<dyn WorkspaceClient>,
session_router: Arc<workdir::WorkdirSessionRouter>,
attachment_mutation_lock: Arc<tokio::sync::Mutex<()>>,
before_workdir_release: Option<BeforeWorkdirRelease>,
after_workdir_attach: Option<AfterWorkdirAttach>,
}
@@ -205,14 +213,27 @@ impl std::fmt::Debug for WorkspaceHttpWorkdirBackend {
#[derive(Debug)]
pub struct WorkspaceAttachedWorkdirSession {
client: Arc<dyn WorkspaceClient>,
target_workdir: String,
workdir: Workdir,
}
impl WorkspaceAttachedWorkdirSession {
pub fn handle(client: Arc<dyn WorkspaceClient>) -> WorkdirSessionHandle {
pub fn handle(
client: Arc<dyn WorkspaceClient>,
target_workdir: impl Into<String>,
) -> WorkdirSessionHandle {
Self::handle_for_workdir(client, target_workdir, "workspace-attachment")
}
pub fn handle_for_workdir(
client: Arc<dyn WorkspaceClient>,
target_workdir: impl Into<String>,
workdir_id: impl Into<String>,
) -> WorkdirSessionHandle {
Arc::new(Self {
client,
workdir: Workdir::new("workspace-attachment"),
target_workdir: target_workdir.into(),
workdir: Workdir::new(workdir_id),
})
}
@@ -229,13 +250,15 @@ impl WorkspaceAttachedWorkdirSession {
"/api/w/{}/workers/self/workdir-session/operations",
encode_path_segment(workspace_id)
),
serde_json::to_string(&WorkspaceWorkdirSessionOperationRequest { operation }).map_err(
|error| {
WorkdirError::Transport(format!(
"failed to encode Workspace Workdir operation: {error}"
))
},
)?,
serde_json::to_string(&WorkspaceWorkdirSessionOperationRequest {
target_workdir: self.target_workdir.clone(),
operation,
})
.map_err(|error| {
WorkdirError::Transport(format!(
"failed to encode Workspace Workdir operation: {error}"
))
})?,
);
let response = self
.client
@@ -399,11 +422,18 @@ impl WorkspaceHttpWorkdirBackend {
fn new(client: Arc<dyn WorkspaceClient>) -> Self {
Self {
client,
session_router: Arc::new(workdir::WorkdirSessionRouter::new()),
attachment_mutation_lock: Arc::new(tokio::sync::Mutex::new(())),
before_workdir_release: None,
after_workdir_attach: None,
}
}
fn with_session_router(mut self, session_router: Arc<workdir::WorkdirSessionRouter>) -> Self {
self.session_router = session_router;
self
}
fn with_child_lifecycle(
mut self,
before_workdir_release: Option<BeforeWorkdirRelease>,
@@ -449,10 +479,12 @@ impl WorkspaceHttpWorkdirBackend {
.transpose()?;
let repository_key =
validate_identity(&input.repository_key, CREATE_TOOL, "repository_key")?;
let display_name = validate_optional_display_name(input.display_name)?;
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),
display_name,
repository_key: repository_key.to_string(),
selector,
operation_id: Some(operation_id),
@@ -472,32 +504,71 @@ impl WorkspaceHttpWorkdirBackend {
}
fn attach(&self, input: WorkdirAttachInput) -> Result<ToolOutput, ToolError> {
let workdir_id = validate_identity(&input.workdir_id, ATTACH_TOOL, "workdir_id")?;
let response = self.attach_response(workdir_id)?;
workdir_output(format!("Attached to Workdir {workdir_id}"), &response)
let alias = validate_identity(&input.alias, ATTACH_TOOL, "alias")?;
workdir::WorkdirAttachmentAlias::new(alias)
.map_err(|error| ToolError::InvalidArgument(error.to_string()))?;
let working_directory_id = validate_identity(
&input.working_directory_id,
ATTACH_TOOL,
"working_directory_id",
)?;
let alias_key = workdir::WorkdirAttachmentAlias::new(alias)
.map_err(|error| ToolError::InvalidArgument(error.to_string()))?;
if self.session_router.session(&alias_key).is_some() {
return Err(ToolError::InvalidArgument(format!(
"attachment alias `{alias}` already exists"
)));
}
let response = self.attach_response(alias, working_directory_id)?;
if let Err(error) = self.session_router.attach(
alias_key.clone(),
WorkspaceAttachedWorkdirSession::handle_for_workdir(
self.client.clone(),
alias,
working_directory_id,
),
) {
let _ = self.detach(alias);
return Err(ToolError::ExecutionFailed(format!(
"attach Workdir session route: {error}"
)));
}
workdir_output(
format!("Attached Workdir {working_directory_id} as {alias}"),
&response,
)
}
fn attach_response(&self, workdir_id: &str) -> Result<WorkdirAttachmentResponse, ToolError> {
fn attach_response(
&self,
alias: &str,
working_directory_id: &str,
) -> Result<WorkdirAttachmentResponse, ToolError> {
let workspace_id = encode_path_segment(self.workspace_id()?);
self.execute_json::<WorkdirAttachmentResponse>(WorkspaceRequest::json(
WorkspaceRequestMethod::Post,
format!("/api/w/{workspace_id}/workers/self/workdir-attachment"),
format!("/api/w/{workspace_id}/workers/self/workdir-attachments"),
serde_json::to_string(&WorkdirAttachRequest {
workdir_id: workdir_id.to_string(),
alias: alias.to_string(),
working_directory_id: working_directory_id.to_string(),
})
.map_err(decode_error)?,
))
}
fn detach(&self) -> Result<ToolOutput, ToolError> {
fn detach(&self, alias: &str) -> Result<ToolOutput, ToolError> {
let workspace_id = encode_path_segment(self.workspace_id()?);
let alias_path = encode_path_segment(alias);
let response = self.execute_json::<WorkdirAttachmentResponse>(WorkspaceRequest {
method: WorkspaceRequestMethod::Delete,
path: format!("/api/w/{workspace_id}/workers/self/workdir-attachment"),
path: format!("/api/w/{workspace_id}/workers/self/workdir-attachments/{alias_path}"),
body: None,
})?;
workdir_output(
format!("Detached from Workdir {}", response.workdir_id),
format!(
"Detached Workdir {} ({})",
response.working_directory_id, response.alias
),
&response,
)
}
@@ -590,6 +661,7 @@ impl Tool for WorkspaceHttpWorkdirTool {
ctx.call_id.to_string(),
),
WorkdirOperation::Attach => {
let _mutation_guard = self.backend.attachment_mutation_lock.lock().await;
let result = self
.backend
.attach(parse_input::<WorkdirAttachInput>(input_json)?);
@@ -601,15 +673,41 @@ impl Tool for WorkspaceHttpWorkdirTool {
result
}
WorkdirOperation::Detach => {
let _input = parse_input::<WorkdirDetachInput>(input_json)?;
if let Some(before_release) = &self.backend.before_workdir_release {
before_release().await.map_err(|error| {
let _mutation_guard = self.backend.attachment_mutation_lock.lock().await;
let input = parse_input::<WorkdirDetachInput>(input_json)?;
let alias = validate_identity(&input.alias, DETACH_TOOL, "alias")?;
let alias_key = workdir::WorkdirAttachmentAlias::new(alias)
.map_err(|error| ToolError::InvalidArgument(error.to_string()))?;
self.backend
.session_router
.begin_detach(&alias_key)
.await
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?;
if let Some(before_release) = &self.backend.before_workdir_release
&& let Err(error) = before_release().await
{
self.backend.session_router.cancel_detach(&alias_key);
return Err(ToolError::ExecutionFailed(format!(
"stop Internal SubWorkers before Workdir detach: {error}"
)));
}
let result = match self.backend.detach(alias) {
Ok(result) => result,
Err(error) => {
self.backend.session_router.cancel_detach(&alias_key);
return Err(error);
}
};
self.backend
.session_router
.finish_detach(&alias_key)
.await
.map_err(|error| {
ToolError::ExecutionFailed(format!(
"stop Internal SubWorkers before Workdir detach: {error}"
"detach Workdir session route after Backend release: {error}"
))
})?;
}
self.backend.detach()
Ok(result)
}
WorkdirOperation::Delete => self
.backend
@@ -672,6 +770,25 @@ fn validate_delete_reason(reason: &str) -> Result<&str, ToolError> {
Ok(reason)
}
fn validate_optional_display_name(
display_name: Option<String>,
) -> Result<Option<String>, ToolError> {
let Some(display_name) = display_name else {
return Ok(None);
};
let display_name = display_name.trim();
if display_name.is_empty()
|| display_name.len() > 80
|| display_name.chars().any(char::is_control)
{
return Err(ToolError::InvalidArgument(
"WorkdirCreate display_name must be non-empty, contain no control characters, and be at most 80 bytes"
.to_string(),
));
}
Ok(Some(display_name.to_string()))
}
fn validate_optional_selector(selector: Option<String>) -> Result<Option<String>, ToolError> {
let Some(selector) = selector else {
return Ok(None);
@@ -725,6 +842,7 @@ fn create_schema() -> serde_json::Value {
"required": ["repository_key"],
"properties": {
"runtime_id": {"type": ["string", "null"], "minLength": 1},
"display_name": {"type": ["string", "null"], "minLength": 1, "maxLength": 80},
"repository_key": {"type": "string", "minLength": 1},
"selector": {"type": ["string", "null"], "minLength": 1}
}
@@ -735,15 +853,23 @@ fn attach_schema() -> serde_json::Value {
json!({
"type": "object",
"additionalProperties": false,
"required": ["workdir_id"],
"required": ["alias", "working_directory_id"],
"properties": {
"workdir_id": {"type": "string", "minLength": 1}
"alias": {"type": "string", "minLength": 1, "maxLength": 64},
"working_directory_id": {"type": "string", "minLength": 1}
}
})
}
fn detach_schema() -> serde_json::Value {
list_schema()
json!({
"type": "object",
"additionalProperties": false,
"required": ["alias"],
"properties": {
"alias": {"type": "string", "minLength": 1, "maxLength": 64}
}
})
}
fn delete_schema() -> serde_json::Value {
@@ -767,6 +893,8 @@ struct WorkdirListInput {}
struct WorkdirCreateInput {
#[serde(default)]
runtime_id: Option<String>,
#[serde(default)]
display_name: Option<String>,
repository_key: String,
#[serde(default)]
selector: Option<String>,
@@ -775,22 +903,27 @@ struct WorkdirCreateInput {
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct WorkdirAttachInput {
workdir_id: String,
alias: String,
working_directory_id: String,
}
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct WorkdirDetachInput {}
struct WorkdirDetachInput {
alias: String,
}
#[derive(Debug, Serialize)]
struct WorkdirAttachRequest {
workdir_id: String,
alias: String,
working_directory_id: String,
}
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
struct WorkdirAttachmentResponse {
workspace_id: String,
workdir_id: String,
alias: String,
working_directory_id: String,
attached: bool,
}
@@ -1021,7 +1154,10 @@ mod tests {
);
assert!(create["properties"].get("path").is_none());
assert!(create["properties"].get("session_id").is_none());
assert_eq!(attach_schema()["required"], json!(["workdir_id"]));
assert_eq!(
attach_schema()["required"],
json!(["alias", "working_directory_id"])
);
assert!(attach_schema()["properties"].get("session_id").is_none());
assert_eq!(
delete_schema()["required"],
@@ -1057,12 +1193,14 @@ mod tests {
})),
response(json!({
"workspace_id": "workspace/test",
"workdir_id": "wd-created",
"alias": "checkout",
"working_directory_id": "wd-created",
"attached": true
})),
response(json!({
"workspace_id": "workspace/test",
"workdir_id": "wd-created",
"alias": "checkout",
"working_directory_id": "wd-created",
"attached": false
})),
response(json!({
@@ -1089,6 +1227,7 @@ mod tests {
.create(
WorkdirCreateInput {
runtime_id: Some("runtime/one".to_string()),
display_name: Some("Review checkout".to_string()),
repository_key: "main".to_string(),
selector: Some("refs/heads/topic".to_string()),
},
@@ -1106,10 +1245,11 @@ mod tests {
assert_eq!(client.requests().len(), 2);
backend
.attach(WorkdirAttachInput {
workdir_id: "wd-created".to_string(),
alias: "checkout".to_string(),
working_directory_id: "wd-created".to_string(),
})
.unwrap();
backend.detach().unwrap();
backend.detach("checkout").unwrap();
backend
.delete(WorkdirDeleteInput {
working_directory_id: "wd-created".to_string(),
@@ -1132,16 +1272,17 @@ mod tests {
serde_json::from_str(requests[1].body.as_deref().unwrap()).unwrap();
assert_eq!(body["repository_key"], "main");
assert_eq!(body["runtime_id"], "runtime/one");
assert_eq!(body["display_name"], "Review checkout");
assert_eq!(body["operation_id"], "call-create-1");
assert_eq!(body["selector"], "refs/heads/topic");
assert_eq!(
requests[2].path,
"/api/w/workspace%2Ftest/workers/self/workdir-attachment"
"/api/w/workspace%2Ftest/workers/self/workdir-attachments"
);
assert_eq!(requests[2].method, WorkspaceRequestMethod::Post);
assert_eq!(
requests[3].path,
"/api/w/workspace%2Ftest/workers/self/workdir-attachment"
"/api/w/workspace%2Ftest/workers/self/workdir-attachments/checkout"
);
assert_eq!(requests[3].method, WorkspaceRequestMethod::Delete);
assert_eq!(
@@ -1160,7 +1301,7 @@ mod tests {
"operation": "stat",
"result": {"path": "visible.txt", "kind": "file", "size": 8}
}))]));
let session = WorkspaceAttachedWorkdirSession::handle(client.clone());
let session = WorkspaceAttachedWorkdirSession::handle(client.clone(), "workdir");
let result = session
.stat(StatRequest {
path: workdir::WorkdirPath::new("visible.txt").unwrap(),
@@ -1176,6 +1317,7 @@ mod tests {
);
let body: serde_json::Value =
serde_json::from_str(requests[0].body.as_deref().unwrap()).unwrap();
assert_eq!(body["target_workdir"], "workdir");
assert_eq!(body["operation"]["operation"], "stat");
assert!(body.get("expected_session_fence").is_none());
assert!(body.get("runtime_id").is_none());
@@ -1188,7 +1330,7 @@ mod tests {
"operation": "command_start",
"result": "command-1"
}))]));
let session = WorkspaceAttachedWorkdirSession::handle(client.clone());
let session = WorkspaceAttachedWorkdirSession::handle(client.clone(), "workdir");
let handle = session
.start_command(CommandRequest {
@@ -1226,7 +1368,7 @@ mod tests {
workdir::http::WorkdirTransportErrorCode::InvalidRequest,
"Workdir operation request is invalid",
)]));
let session = WorkspaceAttachedWorkdirSession::handle(client);
let session = WorkspaceAttachedWorkdirSession::handle(client, "workdir");
let error = session
.glob(workdir::GlobRequest {
@@ -1246,7 +1388,7 @@ mod tests {
status: 502,
body: "secret token and /host/private/path".to_string(),
}]));
let session = WorkspaceAttachedWorkdirSession::handle(client);
let session = WorkspaceAttachedWorkdirSession::handle(client, "workdir");
let error = session
.stat(StatRequest {
@@ -1273,6 +1415,7 @@ mod tests {
]));
let broker = workdir::WorkdirToolBroker::new(WorkspaceAttachedWorkdirSession::handle(
client.clone(),
"workdir",
));
let scoped = broker
.scope(workdir::WorkdirToolScope {
@@ -1323,6 +1466,7 @@ mod tests {
.create(
WorkdirCreateInput {
runtime_id: None,
display_name: None,
repository_key: "main".to_string(),
selector: None,
},
@@ -1348,6 +1492,7 @@ mod tests {
.create(
WorkdirCreateInput {
runtime_id: Some(" ".to_string()),
display_name: None,
repository_key: "main".to_string(),
selector: None,
},
@@ -1363,9 +1508,17 @@ mod tests {
async fn detach_stops_internal_subworkers_before_backend_release() {
let client = Arc::new(RecordingWorkspaceClient::new(vec![response(json!({
"workspace_id": "workspace/test",
"workdir_id": "wd-attached",
"alias": "checkout",
"working_directory_id": "wd-attached",
"attached": false
}))]));
let router = Arc::new(workdir::WorkdirSessionRouter::new());
router
.attach(
workdir::WorkdirAttachmentAlias::new("checkout").unwrap(),
WorkspaceAttachedWorkdirSession::handle(client.clone(), "workdir"),
)
.unwrap();
let cleanup_calls = Arc::new(AtomicUsize::new(0));
let cleanup_calls_for_guard = cleanup_calls.clone();
let before_release: BeforeWorkdirRelease = Arc::new(move || {
@@ -1377,11 +1530,12 @@ mod tests {
});
let tool = WorkspaceHttpWorkdirTool {
backend: WorkspaceHttpWorkdirBackend::new(client.clone())
.with_session_router(router)
.with_child_lifecycle(Some(before_release), None),
operation: WorkdirOperation::Detach,
};
tool.execute("{}", ToolExecutionContext::default())
tool.execute(r#"{"alias":"checkout"}"#, ToolExecutionContext::default())
.await
.unwrap();
@@ -1389,23 +1543,31 @@ mod tests {
assert_eq!(client.requests().len(), 1);
assert_eq!(
client.requests()[0].path,
"/api/w/workspace%2Ftest/workers/self/workdir-attachment"
"/api/w/workspace%2Ftest/workers/self/workdir-attachments/checkout"
);
}
#[tokio::test]
async fn detach_does_not_release_backend_when_child_cleanup_fails() {
let client = Arc::new(RecordingWorkspaceClient::new(Vec::new()));
let router = Arc::new(workdir::WorkdirSessionRouter::new());
router
.attach(
workdir::WorkdirAttachmentAlias::new("checkout").unwrap(),
WorkspaceAttachedWorkdirSession::handle(client.clone(), "workdir"),
)
.unwrap();
let before_release: BeforeWorkdirRelease =
Arc::new(|| Box::pin(async { Err(std::io::Error::other("child cleanup failed")) }));
let tool = WorkspaceHttpWorkdirTool {
backend: WorkspaceHttpWorkdirBackend::new(client.clone())
.with_session_router(router)
.with_child_lifecycle(Some(before_release), None),
operation: WorkdirOperation::Detach,
};
let error = tool
.execute("{}", ToolExecutionContext::default())
.execute(r#"{"alias":"checkout"}"#, ToolExecutionContext::default())
.await
.unwrap_err();
@@ -1413,11 +1575,77 @@ mod tests {
assert!(client.requests().is_empty());
}
#[tokio::test]
async fn detach_reconciles_naturally_completed_workspace_command() {
let client = Arc::new(RecordingWorkspaceClient::new(vec![
response(json!({
"workspace_id": "workspace/test",
"alias": "checkout",
"working_directory_id": "wd-attached",
"attached": true
})),
response(json!({
"operation": "command_start",
"result": "command-1"
})),
response(json!({
"operation": "command_status",
"result": "completed"
})),
response(json!({
"workspace_id": "workspace/test",
"alias": "checkout",
"working_directory_id": "wd-attached",
"attached": false
})),
]));
let backend = WorkspaceHttpWorkdirBackend::new(client.clone());
backend
.attach(WorkdirAttachInput {
alias: "checkout".to_string(),
working_directory_id: "wd-attached".to_string(),
})
.unwrap();
let alias = workdir::WorkdirAttachmentAlias::new("checkout").unwrap();
backend
.session_router
.session(&alias)
.unwrap()
.start_command(CommandRequest {
command: "printf done".to_string(),
timeout_secs: 5,
output_limit: 1024,
cwd: None,
spill_dir: None,
tool_call_id: None,
})
.await
.unwrap();
let tool = WorkspaceHttpWorkdirTool {
backend,
operation: WorkdirOperation::Detach,
};
tool.execute(r#"{"alias":"checkout"}"#, ToolExecutionContext::default())
.await
.unwrap();
let requests = client.requests();
assert_eq!(requests.len(), 4);
assert!(
requests[2]
.body
.as_deref()
.is_some_and(|body| body.contains("command_status"))
);
}
#[tokio::test]
async fn successful_attach_reopens_internal_subworker_admission() {
let client = Arc::new(RecordingWorkspaceClient::new(vec![response(json!({
"workspace_id": "workspace/test",
"workdir_id": "wd-attached",
"alias": "checkout",
"working_directory_id": "wd-attached",
"attached": true
}))]));
let reopen_calls = Arc::new(AtomicUsize::new(0));
@@ -1432,7 +1660,7 @@ mod tests {
};
tool.execute(
r#"{"workdir_id":"wd-attached"}"#,
r#"{"alias":"checkout","working_directory_id":"wd-attached"}"#,
ToolExecutionContext::default(),
)
.await
+11 -6
View File
@@ -165,7 +165,7 @@ where
})?;
worker.disable_manifest_lifecycle_features();
if let Some(session) = inherited_workdir_session {
worker.bind_workdir_session(Some(session));
worker.bind_single_workdir_session(Some(session));
}
let install_report = worker.install_features(features);
@@ -608,7 +608,7 @@ pub(crate) async fn spawn_internal_worker_session(
})?;
worker.disable_manifest_lifecycle_features();
if let Some(session) = inherited_workdir_session {
worker.bind_workdir_session(Some(session));
worker.bind_single_workdir_session(Some(session));
}
let install_report = worker.install_features(features);
let installed_tools = install_report.installed_tool_names();
@@ -701,7 +701,7 @@ pub(crate) fn prepare_internal_worker_from_spec(
})?;
worker.disable_manifest_lifecycle_features();
if let Some(session) = inherited_workdir_session {
worker.bind_workdir_session(Some(session));
worker.bind_single_workdir_session(Some(session));
}
let install_report = worker.install_features(features);
let installed_tools = install_report.installed_tool_names();
@@ -772,7 +772,7 @@ pub(crate) async fn prepare_internal_worker_session(
visibility: InternalWorkerVisibility,
child_registry: Option<Arc<SpawnedWorkerRegistry>>,
on_turn_end: Option<Arc<dyn Fn(InternalWorkerSessionStatus) + Send + Sync>>,
command_event_broker: Option<workdir::WorkdirToolBroker>,
command_event_broker: Option<workdir::WorkdirToolBrokerRouter>,
) -> Result<InternalWorkerSessionHandle, InternalWorkerSessionError> {
let (event_tx, _event_rx) = broadcast::channel(256);
let sink = worker.sink();
@@ -780,9 +780,14 @@ pub(crate) async fn prepare_internal_worker_session(
let alerter = Alerter::new(event_tx.clone());
let in_flight = InFlightEvents::new(event_tx.clone());
if let Some(broker) = command_event_broker.as_ref() {
wire_workdir_command_events(&broker.tool_session(), &in_flight);
let sessions = broker.sessions();
for alias in sessions.aliases() {
if let Ok(selected) = sessions.resolve(Some(alias.as_str())) {
wire_workdir_command_events(&selected.session, &in_flight);
}
}
} else if let Some(session) = worker.workdir_session() {
wire_workdir_command_events(session, &in_flight);
wire_workdir_command_events(&session, &in_flight);
}
let actor_in_flight = in_flight.clone();
worker.attach_alerter(alerter.clone());
+150 -3
View File
@@ -23,11 +23,12 @@ use agen::interceptor::{
use agen::tool::ToolOutput;
use arc_swap::ArcSwap;
use async_trait::async_trait;
use tracing::info;
use tracing::{info, warn};
use crate::compact::metrics_tracker::MetricsTracker;
use crate::compact::state::{AutomaticCompactDecision, CompactState};
use crate::compact::usage_tracker::UsageTracker;
use session_store::SystemItem;
use crate::compact::usage_tracker::{UsageTracker, persist_pending_usage};
use session_store::{SystemItem, segment_log};
use crate::hook::{
HookEventKind, HookPostToolAction, HookPreRequestAction, HookPreToolAction, HookPromptAction,
@@ -81,6 +82,10 @@ pub(crate) struct WorkerInterceptor {
/// persisted into `usage_history`. Subsequent tool-loop LLM calls must
/// see these records during pre-request safety accounting.
usage_tracker: Option<Arc<UsageTracker>>,
/// Metrics produced synchronously while projecting the request (notably
/// prune.fire/skip). Flushed before the matching request usage so metric
/// correlation order remains stable at the in-run commit boundary.
metrics_tracker: Option<Arc<MetricsTracker>>,
/// Pending-notification buffer drained into `worker.history`
/// via [`Self::pending_history_appends`] just before the next LLM
/// request. The Engine `extend`s these into its persistent history
@@ -152,6 +157,7 @@ impl WorkerInterceptor {
compact_state,
usage_history,
usage_tracker: None,
metrics_tracker: None,
pending_notifies,
pending_attachments,
prompts,
@@ -168,6 +174,11 @@ impl WorkerInterceptor {
self
}
pub(crate) fn with_metrics_tracker(mut self, metrics_tracker: Arc<MetricsTracker>) -> Self {
self.metrics_tracker = Some(metrics_tracker);
self
}
pub(crate) fn with_prompt_workspace_id(mut self, workspace_id: Option<String>) -> Self {
self.prompt_workspace_id = workspace_id;
self
@@ -272,6 +283,56 @@ impl WorkerInterceptor {
}
}
}
fn try_persist_metric(writer: &Arc<dyn SystemItemCommitter>, metric: &session_metrics::Metric) {
let payload = serde_json::to_value(metric).expect("Metric is Serialize");
if writer
.commit_log_entry(session_store::LogEntry::Extension {
ts: segment_log::now_millis(),
domain: session_metrics::DOMAIN.into(),
payload,
})
.is_err()
{
warn!(name = %metric.name, "failed to record session metric; dropping");
}
}
fn persist_completed_request_usage(&self) -> InterceptorResult<()> {
let Some(tracker) = self.usage_tracker.as_ref() else {
return Ok(());
};
tracker.request_completed();
let (Some(writer), Some(usage_history)) =
(self.log_writer.as_ref(), self.usage_history.as_ref())
else {
// Low-level Worker tests may omit the type-erased writer. Their
// terminal persist_turn remains the durable fallback.
return Ok(());
};
if let Some(metrics_tracker) = self.metrics_tracker.as_ref() {
for metric in metrics_tracker.drain() {
Self::try_persist_metric(writer, &metric);
}
}
persist_pending_usage(
tracker,
usage_history,
self.compact_state.as_deref(),
|entry, after_commit, post_metrics| {
writer.commit_usage_boundary(entry, after_commit, post_metrics)
},
)
.map_err(|error| {
InterceptorError::new(
InterceptorErrorCategory::Dependency,
format!("usage persistence failed: {error}"),
)
})?;
Ok(())
}
fn attach_prompt_provenance(&self, items: &mut [SystemItem]) {
let prompts = self.prompts.load();
let projection = prompts.projection();
@@ -596,6 +657,11 @@ impl Interceptor<SessionHistoryMetadata> for WorkerInterceptor {
&self,
context: AssistantTurnEndContext<'_, SessionHistoryMetadata>,
) -> InterceptorResult<TurnEndAction> {
// `on_assistant_turn_end` is reached only after a normal provider
// stream completed and its usage callback was flushed. Commit that
// exact request before tools can extend the same logical run.
self.persist_completed_request_usage()?;
let history = context.history;
let final_text_preview = history
.iter()
@@ -752,6 +818,25 @@ mod tests {
}
}
struct FailingUsageCommitter {
attempts: Arc<AtomicUsize>,
}
impl SystemItemCommitter for FailingUsageCommitter {
fn commit_log_entry(
&self,
entry: session_store::LogEntry,
) -> Result<(), session_store::StoreError> {
if matches!(entry, session_store::LogEntry::LlmUsage { .. }) {
self.attempts.fetch_add(1, Ordering::SeqCst);
return Err(session_store::StoreError::Io(std::io::Error::other(
"synthetic usage failure",
)));
}
Ok(())
}
}
struct AppendingPreRequestHook {
saw_handle: Arc<AtomicBool>,
}
@@ -955,6 +1040,68 @@ mod tests {
);
}
#[tokio::test]
async fn completed_request_usage_failure_stops_before_rearm() {
use crate::compact::state::{AutomaticCompactDecision, CompactionOutcome};
let registry = Arc::new(HookRegistryBuilder::new().build());
let state = Arc::new(CompactState::new(None, Some(10), 0));
assert!(matches!(
state.evaluate_request(11),
AutomaticCompactDecision::Start(_)
));
assert!(state.complete_automatic(CompactionOutcome::Succeeded));
let usage_history = Arc::new(Mutex::new(Vec::new()));
let tracker = Arc::new(UsageTracker::new());
tracker.note_compaction_correlation_id("compact-id".into());
tracker.note_request(1);
tracker.record_usage(&agen::event::UsageEvent {
input_tokens: Some(5),
output_tokens: Some(1),
total_tokens: Some(6),
cache_read_input_tokens: Some(0),
cache_creation_input_tokens: Some(0),
});
let attempts = Arc::new(AtomicUsize::new(0));
let interceptor = WorkerInterceptor::new(
registry,
Some(Arc::clone(&state)),
Some(Arc::clone(&usage_history)),
NotifyBuffer::new(),
Arc::new(Mutex::new(Vec::new())),
test_prompts(),
Some(Arc::new(FailingUsageCommitter {
attempts: Arc::clone(&attempts),
})),
)
.with_usage_tracker(Arc::clone(&tracker));
let error = interceptor
.on_assistant_turn_end(AssistantTurnEndContext {
invocation: Default::default(),
assistant_entries: &[],
history: &[],
tool_calls: &[],
})
.await
.expect_err("usage storage failure must stop the request boundary");
assert_eq!(error.category(), InterceptorErrorCategory::Dependency);
assert!(error.diagnostic().contains("usage persistence failed"));
assert_eq!(attempts.load(Ordering::SeqCst), 1);
assert_eq!(
tracker.records().len(),
1,
"failed record must remain pending"
);
assert!(usage_history.lock().unwrap().is_empty());
assert!(matches!(
state.evaluate_request(11),
AutomaticCompactDecision::Block(_)
));
}
#[tokio::test]
async fn pre_llm_request_counts_in_flight_usage_records() {
let registry = Arc::new(HookRegistryBuilder::new().build());
+8 -4
View File
@@ -27,7 +27,7 @@ use session_store::{
};
use tokio::sync::{Notify, broadcast};
use tracing::warn;
use workdir::WorkdirScopeLease;
use workdir::WorkdirScopeLeaseSet;
use crate::internal_worker::{InternalWorkerSessionHandle, InternalWorkerVisibility};
use crate::runtime::dir::{RuntimeDir, SpawnedWorkerRecord};
@@ -70,7 +70,7 @@ pub(crate) struct SubWorkerStopSummary {
pub(crate) struct InternalSpawnedWorkerRecord {
pub worker_name: String,
pub scope_delegated: Vec<ScopeRule>,
pub workdir_tool_scope: Arc<WorkdirScopeLease>,
pub workdir_tool_scope: Arc<WorkdirScopeLeaseSet>,
#[cfg(test)]
pub installed_tools: Arc<[String]>,
pub session: InternalWorkerSessionHandle,
@@ -89,7 +89,7 @@ impl InternalSpawnedWorkerRecord {
pub(crate) fn new(
worker_name: String,
scope_delegated: Vec<ScopeRule>,
workdir_tool_scope: WorkdirScopeLease,
workdir_tool_scope: WorkdirScopeLeaseSet,
#[cfg(test)] installed_tools: Vec<String>,
session: InternalWorkerSessionHandle,
child_registry: Arc<SpawnedWorkerRegistry>,
@@ -1179,7 +1179,11 @@ mod tests {
InternalSpawnedWorkerRecord::new(
name.into(),
Vec::new(),
delegation,
workdir::WorkdirScopeLeaseSet::from_single(
workdir::WorkdirAttachmentAlias::new("workdir").unwrap(),
delegation,
)
.unwrap(),
Vec::new(),
session,
registry(),
+74 -30
View File
@@ -5,6 +5,7 @@
//! in-process Internal Worker session actor. No Runtime Worker record, OS process, PID, Unix socket,
//! or machine-wide child allocation is created.
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
@@ -22,7 +23,7 @@ use manifest::{
use serde::Deserialize;
use tokio::sync::mpsc;
use workdir::{
WorkdirToolBroker, WorkdirToolScope, WorkdirToolScopePermission, WorkdirToolScopeRule,
WorkdirToolBrokerRouter, WorkdirToolScope, WorkdirToolScopePermission, WorkdirToolScopeRule,
};
use crate::PromptCatalogSource;
@@ -81,6 +82,10 @@ struct ReviewerHandoffInput {
#[derive(Debug, Deserialize, schemars::JsonSchema)]
struct ScopeRuleInput {
/// Worker-local alias of the parent Workdir attachment. May be omitted only
/// when the parent has exactly one attachment.
#[serde(default)]
target_workdir: Option<String>,
/// Logical Workdir-relative target such as `.` or `src`. Absolute host
/// paths and parent traversal are rejected.
target: String,
@@ -310,7 +315,7 @@ pub struct SubWorkerSpawnTool {
/// Directory the spawned SubWorker's tools should use when the LLM did not
/// override it. Defaults to the spawner's cwd.
/// Parent-owned broker for scoped Workdir tool execution.
workdir_tool_broker: Option<WorkdirToolBroker>,
workdir_tool_broker: Option<WorkdirToolBrokerRouter>,
/// Parent-owned in-memory registry shared by the five SubWorker tools.
registry: Arc<SpawnedWorkerRegistry>,
/// Spawner's resolved Manifest. `profile = "inherit"` derives the
@@ -337,7 +342,7 @@ impl SubWorkerSpawnTool {
runtime_base: PathBuf,
bash_output_dir: PathBuf,
workspace_root: PathBuf,
workdir_tool_broker: Option<WorkdirToolBroker>,
workdir_tool_broker: Option<WorkdirToolBrokerRouter>,
registry: Arc<SpawnedWorkerRegistry>,
spawner_manifest: WorkerManifest,
prompt_loader: PromptCatalogSource,
@@ -428,14 +433,19 @@ impl Tool for SubWorkerSpawnTool {
))
})?;
let workdir_tool_broker = require_workdir_tool_broker(self.workdir_tool_broker.as_ref())?;
let tool_scope = workdir_tool_scope(input.cwd.as_deref(), workdir_rules, input.command)?;
let tool_scopes = routed_workdir_tool_scopes(
workdir_tool_broker,
input.cwd.as_deref(),
workdir_rules,
input.command,
)?;
let workdir_scope = workdir_tool_broker
.scope(tool_scope)
.scope(tool_scopes)
.await
.map_err(|error| {
ToolError::InvalidArgument(format!("scope parent-owned Workdir tools: {error}"))
})?;
let child_workdir_tool_broker = workdir_scope.broker();
let child_workdir_tool_broker = workdir_scope.broker_router();
let spawn_selector =
parse_spawn_profile_selector(input.profile.as_deref()).map_err(|msg| {
@@ -537,6 +547,7 @@ impl Tool for SubWorkerSpawnTool {
))
})?;
let child_scope = child.scope().clone();
child.bind_workdir_sessions(workdir_scope.session_router());
let child_registry = SpawnedWorkerRegistry::new_internal(input.name.clone(), child_scope);
register_worker_tools(
&mut child,
@@ -717,22 +728,27 @@ fn logical_workdir_path(value: &str, field: &str) -> Result<FsPath, ToolError> {
})
}
fn parse_workdir_scope(rules: &[ScopeRuleInput]) -> Result<Vec<WorkdirToolScopeRule>, ToolError> {
fn parse_workdir_scope(
rules: &[ScopeRuleInput],
) -> Result<Vec<(Option<String>, WorkdirToolScopeRule)>, ToolError> {
if rules.is_empty() {
return Err(ToolError::InvalidArgument("scope must not be empty".into()));
}
rules
.iter()
.map(|rule| {
Ok(WorkdirToolScopeRule {
target: logical_workdir_path(&rule.target, "scope.target")?,
permission: match rule.permission {
PermissionInput::Read => WorkdirToolScopePermission::Read,
PermissionInput::Write => WorkdirToolScopePermission::Write,
Ok((
rule.target_workdir.clone(),
WorkdirToolScopeRule {
target: logical_workdir_path(&rule.target, "scope.target")?,
permission: match rule.permission {
PermissionInput::Read => WorkdirToolScopePermission::Read,
PermissionInput::Write => WorkdirToolScopePermission::Write,
},
recursive: rule.recursive,
symlink_policy: rule.symlink_policy.into(),
},
recursive: rule.recursive,
symlink_policy: rule.symlink_policy.into(),
})
))
})
.collect()
}
@@ -749,9 +765,28 @@ fn workdir_tool_scope(
})
}
fn routed_workdir_tool_scopes(
broker: &WorkdirToolBrokerRouter,
cwd: Option<&str>,
rules: Vec<(Option<String>, WorkdirToolScopeRule)>,
command: bool,
) -> Result<Vec<(workdir::WorkdirAttachmentAlias, WorkdirToolScope)>, ToolError> {
let mut grouped = BTreeMap::<workdir::WorkdirAttachmentAlias, Vec<WorkdirToolScopeRule>>::new();
for (target_workdir, rule) in rules {
let alias = broker
.resolve_alias(target_workdir.as_deref())
.map_err(|error| ToolError::InvalidArgument(error.to_string()))?;
grouped.entry(alias).or_default().push(rule);
}
grouped
.into_iter()
.map(|(alias, rules)| workdir_tool_scope(cwd, rules, command).map(|scope| (alias, scope)))
.collect()
}
fn require_workdir_tool_broker(
broker: Option<&WorkdirToolBroker>,
) -> Result<&WorkdirToolBroker, ToolError> {
broker: Option<&WorkdirToolBrokerRouter>,
) -> Result<&WorkdirToolBrokerRouter, ToolError> {
broker.ok_or_else(|| {
ToolError::InvalidArgument(
"SubWorkerSpawn requires parent-owned Workdir tools; attach a Workdir before granting filesystem access"
@@ -992,7 +1027,7 @@ pub(crate) fn sub_worker_spawn_tool(
runtime_base: PathBuf,
bash_output_dir: PathBuf,
workspace_root: PathBuf,
workdir_tool_broker: Option<WorkdirToolBroker>,
workdir_tool_broker: Option<WorkdirToolBrokerRouter>,
registry: Arc<SpawnedWorkerRegistry>,
spawner_manifest: WorkerManifest,
prompts: Arc<ArcSwap<PromptCatalog>>,
@@ -1018,7 +1053,7 @@ fn sub_worker_spawn_tool_impl(
runtime_base: PathBuf,
bash_output_dir: PathBuf,
workspace_root: PathBuf,
workdir_tool_broker: Option<WorkdirToolBroker>,
workdir_tool_broker: Option<WorkdirToolBrokerRouter>,
registry: Arc<SpawnedWorkerRegistry>,
spawner_manifest: WorkerManifest,
prompts: Arc<ArcSwap<PromptCatalog>>,
@@ -1096,12 +1131,14 @@ mod tests {
fn workdir_scope_uses_logical_relative_paths() {
let rules = parse_workdir_scope(&[
ScopeRuleInput {
target_workdir: None,
target: ".".to_string(),
permission: PermissionInput::Read,
recursive: true,
symlink_policy: Default::default(),
},
ScopeRuleInput {
target_workdir: Some("checkout".to_string()),
target: "src".to_string(),
permission: PermissionInput::Write,
recursive: false,
@@ -1109,12 +1146,14 @@ mod tests {
},
])
.unwrap();
assert_eq!(rules[0].target.as_str(), "");
assert_eq!(rules[1].target.as_str(), "src");
assert_eq!(rules[0].symlink_policy, SymlinkPolicy::Resolved);
assert_eq!(rules[1].symlink_policy, SymlinkPolicy::Logical);
assert_eq!(rules[0].1.target.as_str(), "");
assert_eq!(rules[1].0.as_deref(), Some("checkout"));
assert_eq!(rules[1].1.target.as_str(), "src");
assert_eq!(rules[0].1.symlink_policy, SymlinkPolicy::Resolved);
assert_eq!(rules[1].1.symlink_policy, SymlinkPolicy::Logical);
for target in ["/host/path", "../escape"] {
let error = parse_workdir_scope(&[ScopeRuleInput {
target_workdir: None,
target: target.to_string(),
permission: PermissionInput::Read,
recursive: true,
@@ -1278,15 +1317,17 @@ enabled = false
let fail_requests = Arc::new(AtomicBool::new(false));
let prompt_loader = PromptCatalogSource::builtins_only();
let (parent_method_tx, mut parent_method_rx) = mpsc::channel(8);
let workdir_tool_broker = workdir::WorkdirToolBroker::new(Arc::new(
workdir::LocalWorkdirSession::materialized_bound(
let workdir_tool_broker = workdir::WorkdirToolBrokerRouter::from_single(
workdir::WorkdirAttachmentAlias::new("workdir").unwrap(),
Arc::new(workdir::LocalWorkdirSession::materialized_bound(
workdir::Workdir::new("test-workdir"),
workspace_root.clone(),
workspace_root.clone(),
spawner_scope.clone(),
workdir::WorkdirSessionCapabilities::ALL,
),
));
)),
)
.unwrap();
let tool = SubWorkerSpawnTool::new(
"parent".into(),
workspace_context,
@@ -1528,9 +1569,11 @@ enabled = false
Arc::new(AvailableWorkspaceClient),
);
let remote_client = Arc::new(StrictRemoteWorkdirWorkspaceClient::default());
let workdir_tool_broker = workdir::WorkdirToolBroker::new(
WorkspaceAttachedWorkdirSession::handle(remote_client.clone()),
);
let workdir_tool_broker = workdir::WorkdirToolBrokerRouter::from_single(
workdir::WorkdirAttachmentAlias::new("workdir").unwrap(),
WorkspaceAttachedWorkdirSession::handle(remote_client.clone(), "workdir"),
)
.unwrap();
let calls = Arc::new(AtomicUsize::new(0));
let (parent_method_tx, _parent_method_rx) = mpsc::channel(8);
let tool = SubWorkerSpawnTool::new(
@@ -1607,6 +1650,7 @@ enabled = false
assert!(properties.contains_key("command"), "schema: {schema}");
let schema_text = serde_json::to_string(&schema).unwrap();
assert!(schema_text.contains("symlink_policy"), "schema: {schema}");
assert!(schema_text.contains("target_workdir"), "schema: {schema}");
assert!(schema_text.contains("logical"), "schema: {schema}");
let required = schema
.get("required")
+289 -88
View File
@@ -45,9 +45,9 @@ use crate::compact::state::{
};
use crate::compact::telemetry::{
CompactAttempt, CompactFailureCategory, CompactMode, CompactSuccessStats,
CompactThresholdPolicy, correlated_post_request_metric, new_compact_metric_correlation_id,
CompactThresholdPolicy, new_compact_metric_correlation_id,
};
use crate::compact::usage_tracker::UsageTracker;
use crate::compact::usage_tracker::{UsageTracker, persist_pending_usage};
use crate::feature::background::{BackgroundTaskRewriteGuard, FeatureBackgroundTaskRegistry};
use crate::feature::builtin::memory::WorkspaceMemoryBackendError;
use crate::feature::builtin::{TaskFeature, WorkerObservationProvider};
@@ -374,7 +374,8 @@ use protocol::{
use tokio::net::UnixStream;
use tokio::sync::broadcast;
use workdir::{
LocalWorkdirSession, ReadOnlyWorkdirSession, WorkdirSessionCapabilities, WorkdirSessionHandle,
LocalWorkdirSession, ReadOnlyWorkdirSession, WorkdirAttachmentAlias,
WorkdirSessionCapabilities, WorkdirSessionHandle, WorkdirSessionRouter,
};
const RESTORE_RECONCILIATION_REACHABILITY_TIMEOUT: Duration = Duration::from_millis(500);
@@ -1946,6 +1947,28 @@ impl PendingSubmissionHandle<session_store::FsStore> {
pub trait SystemItemCommitter: Send + Sync {
fn commit_log_entry(&self, entry: LogEntry) -> Result<(), StoreError>;
/// Commit one authoritative usage record while holding the same append
/// barrier across its in-memory publication, guard rearm, and correlated
/// post-request metrics. Metric appends remain best-effort.
fn commit_usage_boundary(
&self,
entry: LogEntry,
after_commit: &mut dyn FnMut(),
post_metrics: &[session_metrics::Metric],
) -> Result<(), StoreError> {
self.commit_log_entry(entry)?;
after_commit();
for metric in post_metrics {
let payload = serde_json::to_value(metric).expect("Metric is Serialize");
let _ = self.commit_log_entry(LogEntry::Extension {
ts: segment_log::now_millis(),
domain: session_metrics::DOMAIN.into(),
payload,
});
}
Ok(())
}
fn commit_system_item_with_extensions(
&self,
item: SystemItem,
@@ -1977,6 +2000,30 @@ where
fn commit_log_entry(&self, entry: LogEntry) -> Result<(), StoreError> {
self.append_entry(entry)
}
fn commit_usage_boundary(
&self,
entry: LogEntry,
after_commit: &mut dyn FnMut(),
post_metrics: &[session_metrics::Metric],
) -> Result<(), StoreError> {
let _append_guard = self
.state
.append_lock
.lock()
.expect("segment append lock poisoned");
self.append_entry_locked(entry)?;
after_commit();
for metric in post_metrics {
let payload = serde_json::to_value(metric).expect("Metric is Serialize");
let _ = self.append_entry_locked(LogEntry::Extension {
ts: segment_log::now_millis(),
domain: session_metrics::DOMAIN.into(),
payload,
});
}
Ok(())
}
}
struct SessionFlowRuntimeStateCommitter<St: Store + Clone> {
@@ -2128,9 +2175,9 @@ pub struct Worker<C: LlmClient, St: Store> {
/// Explicit local filesystem authority, or `None` for Workers with no
/// local cwd and no filesystem/Bash tool surface.
filesystem_authority: WorkerFilesystemAuthority,
/// Live WorkdirSession provider derived once from the Worker–Workdir binding.
/// Local tools, file views, and compaction workers clone this handle.
workdir_session: Option<WorkdirSessionHandle>,
/// Alias-keyed live Workdir sessions. Provider transport remains hidden
/// behind each handle and no attachment is designated primary.
workdir_sessions: Arc<WorkdirSessionRouter>,
/// Path-free workspace identity/client context injected by Runtime/host.
/// This never grants local filesystem authority.
workspace_context: WorkerWorkspaceContext,
@@ -2154,16 +2201,16 @@ pub struct Worker<C: LlmClient, St: Store> {
/// Shared compaction state (present when threshold is configured).
compact_state: Option<Arc<CompactState>>,
/// Per-LLM-request Usage tracker. Always present after construction.
/// Captures `(history_len, UsageEvent)` pairs during a run; drained
/// in `persist_turn` and persisted as `LogEntry::LlmUsage` entries.
/// Captures `(history_len, UsageEvent)` pairs until their authoritative
/// request-boundary or terminal `LlmUsage` append succeeds.
usage_tracker: Arc<UsageTracker>,
/// Sync-side buffer for `Metric` values queued from inside Engine
/// callbacks (currently the prune observer). Drained in `persist_turn`
/// and written via `session_metrics::record_metric` alongside
/// `LogEntry::LlmUsage`. Always present after construction.
/// callbacks (currently the prune observer). Drained at the same request
/// boundary as usage, with terminal `persist_turn` as the fallback.
metrics_tracker: Arc<crate::compact::metrics_tracker::MetricsTracker>,
/// Cumulative Usage measurement timeline, one entry per LLM call.
/// Restored from session log on `restore`, appended on each persist.
/// Restored from session log on `restore`, appended after each durable
/// request-boundary usage commit.
/// Read by token-accounting APIs (`Worker::total_tokens`, etc.).
///
/// Wrapped in `Arc<Mutex>` so that callbacks injected into the
@@ -2464,7 +2511,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
let prompts = Arc::new(ArcSwap::from(PromptCatalog::builtins_only()?));
DelegationScope::from_config(&manifest.delegation_scope).map_err(WorkerError::Scope)?;
let scope = SharedScope::new(scope);
let workdir_session = workdir_session_from_authority(&filesystem_authority, &scope);
let workdir_sessions = workdir_sessions_from_authority(&filesystem_authority, &scope);
let mut worker = Self {
manifest,
engine: Some(worker),
@@ -2475,7 +2522,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
worker_metadata_segment_cas: None,
segment_state: SegmentState::new(session_id, segment_id, 0),
filesystem_authority,
workdir_session,
workdir_sessions,
workspace_context,
flow_runtime_state: Arc::new(Mutex::new(None)),
flow_feature_enabled: false,
@@ -2652,15 +2699,36 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
self.filesystem_authority.as_local()
}
pub fn workdir_session(&self) -> Option<&WorkdirSessionHandle> {
self.workdir_session.as_ref()
/// Compatibility projection for tool schemas that do not yet carry a
/// target alias. It is available only when the attachment set has exactly
/// one member, so multiple attachments never acquire an implicit primary.
pub fn workdir_session(&self) -> Option<WorkdirSessionHandle> {
self.workdir_sessions.only_session()
}
/// Replace the constructor fallback with the provider binding resolved by
/// the owning Runtime. Runtime calls this before the Worker controller is
/// spawned, so tools only ever observe the Runtime-bound handle.
pub fn bind_workdir_session(&mut self, workdir_session: Option<WorkdirSessionHandle>) {
self.workdir_session = workdir_session;
pub fn workdir_sessions(&self) -> Arc<WorkdirSessionRouter> {
self.workdir_sessions.clone()
}
/// Replace the constructor fallback with the complete alias-keyed provider
/// set resolved by the owning Runtime before controller startup.
pub fn bind_workdir_sessions(&mut self, workdir_sessions: Arc<WorkdirSessionRouter>) {
self.workdir_sessions = workdir_sessions;
}
/// Bind one session for direct/Internal Worker construction. Runtime Worker
/// creation uses [`Self::bind_workdir_sessions`] with the complete set.
pub fn bind_single_workdir_session(&mut self, workdir_session: Option<WorkdirSessionHandle>) {
let router = Arc::new(WorkdirSessionRouter::new());
if let Some(session) = workdir_session {
router
.attach(
WorkdirAttachmentAlias::new("workdir").expect("static alias is valid"),
session,
)
.expect("fresh Workdir router accepts its only attachment");
}
self.workdir_sessions = router;
}
/// Path-free workspace identity, if Runtime/host associated this Worker
@@ -2928,12 +2996,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
location.session_id,
location.segment_id,
);
let read_only_tools = self
.workdir_session
.clone()
.map(|source| Arc::new(ReadOnlyWorkdirSession::new(source)) as WorkdirSessionHandle)
.map(tools::read_only_builtin_tools)
.unwrap_or_default();
let read_only_tools = tools::read_only_routed_builtin_tools(self.workdir_sessions());
let client = self
.engine
.as_ref()
@@ -3309,8 +3372,8 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// Handle to the synchronous `MetricsTracker` buffer.
///
/// Engine callbacks (e.g. the prune observer) clone this `Arc` and
/// `.push(metric)` into it; Worker drains it in `persist_turn` and
/// writes each metric via `session_metrics::record_metric`.
/// `.push(metric)` into it; Worker drains it at the next request accounting
/// boundary, with terminal persistence as fallback.
pub(crate) fn metrics_tracker_handle(
&self,
) -> Arc<crate::compact::metrics_tracker::MetricsTracker> {
@@ -3546,6 +3609,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
self.pending_committed_history.clone(),
)
.with_usage_tracker(self.usage_tracker.clone())
.with_metrics_tracker(self.metrics_tracker.clone())
.with_prompt_workspace_id(
self.workspace_context
.workspace_id()
@@ -4061,7 +4125,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// unresolved placeholder stays in the flattened user message so the LLM
/// still sees the intent.
async fn resolve_file_refs(&self, segments: &[Segment]) -> Vec<SystemItem> {
let Some(workdir) = self.workdir_session.clone() else {
let Some(workdir) = self.workdir_session() else {
for seg in segments {
if let Segment::FileRef { path } = seg {
self.alert(
@@ -4431,6 +4495,15 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
self.handle_worker_result(result, history_before).await
}
fn ensure_no_pending_usage_for_segment_replacement(&self) -> Result<(), WorkerError> {
let pending_records = self.usage_tracker.pending_record_count();
if pending_records == 0 {
Ok(())
} else {
Err(WorkerError::PendingUsagePersistence { pending_records })
}
}
/// Ensure the session exists and the writer's tally still matches
/// the on-disk entry count.
///
@@ -4472,6 +4545,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
if store_count == entries_written {
return Ok(());
}
self.ensure_no_pending_usage_for_segment_replacement()?;
// Auto-fork within the same Session: mint a fresh Segment and
// switch to it. The source segment is left immutable (no terminal
// marker is written back); the fork relationship is recorded
@@ -5008,45 +5082,38 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
self.try_record_metric(&metric);
}
// Persist any LLM Usage measurements collected during this run.
// One LogEntry::LlmUsage per LLM call (the tool loop may have run
// many calls within a single Worker::run). Each is also appended to
// the in-memory `usage_history` so token-accounting APIs see it
// before the next run. Records carrying a `correlation_id` (set
// by an upstream observer such as the prune projection) also get
// a paired `prune.post_request` metric so cache_read/write can be
// joined back to the originating event.
let usage_records = self.usage_tracker.drain();
for recorded in usage_records {
let crate::compact::usage_tracker::RecordedUsage {
record,
post_requests,
} = recorded;
let committed_post_compact_request = post_requests.iter().any(|link| {
link.metric == crate::compact::usage_tracker::PostRequestMetric::Compaction
});
self.commit_entry(LogEntry::LlmUsage {
ts: segment_log::now_millis(),
history_len: record.history_len,
input_total_tokens: record.input_total_tokens,
cache_read_tokens: record.cache_read_tokens,
cache_write_tokens: record.cache_write_tokens,
output_tokens: record.output_tokens,
})?;
if committed_post_compact_request {
if let Some(state) = &self.compact_state {
state.post_compact_request_committed();
}
}
for link in post_requests {
let metric =
correlated_post_request_metric(link.metric, &link.correlation_id, &record);
// Persist any request measurements that did not already cross the
// in-run request boundary (for example partial usage from an errored or
// cancelled stream, or low-level Workers without an attached writer).
// `persist_pending_usage` removes each record only after its authoritative
// LlmUsage append succeeds, so the in-run and terminal paths cannot
// double-write or silently lose a failed append.
let usage_tracker = self.usage_tracker.clone();
let usage_history = self.usage_history.clone();
let compact_state = self.compact_state.clone();
if let Some(writer) = self.log_writer.clone() {
let _ = persist_pending_usage(
&usage_tracker,
&usage_history,
compact_state.as_deref(),
|entry, after_commit, post_metrics| {
writer.commit_usage_boundary(entry, after_commit, post_metrics)
},
)?;
} else {
let post_request_metrics = persist_pending_usage(
&usage_tracker,
&usage_history,
compact_state.as_deref(),
|entry, after_commit, _post_metrics| {
self.commit_entry(entry)?;
after_commit();
Ok(())
},
)?;
for metric in post_request_metrics {
self.try_record_metric(&metric);
}
self.usage_history
.lock()
.expect("usage_history poisoned")
.push(record);
}
let interrupted = matches!(
@@ -5106,6 +5173,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
mut cancel: Option<tokio::sync::watch::Receiver<bool>>,
trigger: CompactionTrigger,
) -> Result<SegmentId, WorkerError> {
self.ensure_no_pending_usage_for_segment_replacement()?;
self.release_pending_compaction_service().await?;
let _rewrite_guard = self
.prepare_session_rewrite(SessionRewriteKind::Compact)
@@ -5380,11 +5448,14 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
// Default references: the N most-recently-touched files in the
// session, surfaced so the compact worker can inspect them and
// decide which (if any) the next session needs.
let default_refs: Vec<PathBuf> = self
.tracker
.as_ref()
.map(|t| t.recent_files(manifest::defaults::COMPACT_DEFAULT_REFERENCE_COUNT))
.unwrap_or_default();
let default_refs: Vec<PathBuf> = if self.workdir_sessions.len() == 1 {
self.tracker
.as_ref()
.map(|t| t.recent_files(manifest::defaults::COMPACT_DEFAULT_REFERENCE_COUNT))
.unwrap_or_default()
} else {
Vec::new()
};
// Input text fed to the compact worker. Includes the default
// references, current TaskStore snapshot, current TaskStore snapshot, and the (pruned) conversation text.
@@ -5428,7 +5499,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
// Build a normal parent-owned Internal Worker over a pinned immutable
// capture. Only SessionExplore and compaction output tools are installed.
let workdir = self.workdir_session.clone();
let workdir = self.workdir_session();
let read_only_workdir = workdir.clone().map(|session| {
Arc::new(ReadOnlyWorkdirSession::new(session)) as workdir::WorkdirSessionHandle
});
@@ -6139,7 +6210,8 @@ where
let worker_metadata_writer = Some(worker_metadata_writer_for_store(&store));
let worker_metadata_segment_cas = Some(worker_metadata_segment_cas_for_store(&store));
let scope = SharedScope::new(common.scope);
let workdir_session = workdir_session_from_authority(&common.filesystem_authority, &scope);
let workdir_sessions =
workdir_sessions_from_authority(&common.filesystem_authority, &scope);
let mut worker = Self {
manifest,
@@ -6151,7 +6223,7 @@ where
worker_metadata_segment_cas,
segment_state: SegmentState::new(session_id, segment_id, 0),
filesystem_authority: common.filesystem_authority,
workdir_session,
workdir_sessions,
workspace_context: common.workspace_context,
flow_runtime_state: Arc::new(Mutex::new(None)),
flow_feature_enabled: false,
@@ -6228,7 +6300,8 @@ where
apply_worker_manifest(&mut engine, &manifest.engine);
engine.set_cache_key(Some(segment_id.to_string()));
let scope = SharedScope::new(common.scope);
let workdir_session = workdir_session_from_authority(&common.filesystem_authority, &scope);
let workdir_sessions =
workdir_sessions_from_authority(&common.filesystem_authority, &scope);
let mut worker = Self {
manifest,
engine: Some(engine),
@@ -6239,7 +6312,7 @@ where
worker_metadata_segment_cas: None,
segment_state: SegmentState::new(session_id, segment_id, 0),
filesystem_authority: common.filesystem_authority,
workdir_session,
workdir_sessions,
workspace_context: common.workspace_context,
flow_runtime_state: Arc::new(Mutex::new(None)),
flow_feature_enabled: false,
@@ -6350,7 +6423,8 @@ where
let worker_metadata_writer = Some(worker_metadata_writer_for_store(&store));
let worker_metadata_segment_cas = Some(worker_metadata_segment_cas_for_store(&store));
let scope = SharedScope::new(common.scope);
let workdir_session = workdir_session_from_authority(&common.filesystem_authority, &scope);
let workdir_sessions =
workdir_sessions_from_authority(&common.filesystem_authority, &scope);
let mut worker = Self {
manifest,
@@ -6362,7 +6436,7 @@ where
worker_metadata_segment_cas,
segment_state: SegmentState::new(session_id, segment_id, 0),
filesystem_authority: common.filesystem_authority,
workdir_session,
workdir_sessions,
workspace_context: common.workspace_context,
flow_runtime_state: Arc::new(Mutex::new(None)),
flow_feature_enabled: false,
@@ -6726,7 +6800,8 @@ where
let worker_metadata_writer = Some(worker_metadata_writer_for_store(&store));
let worker_metadata_segment_cas = Some(worker_metadata_segment_cas_for_store(&store));
let scope = SharedScope::new(common.scope);
let workdir_session = workdir_session_from_authority(&common.filesystem_authority, &scope);
let workdir_sessions =
workdir_sessions_from_authority(&common.filesystem_authority, &scope);
let mut worker = Self {
manifest,
@@ -6738,7 +6813,7 @@ where
worker_metadata_segment_cas,
segment_state: SegmentState::new(session_id, segment_id, state.entries_count),
filesystem_authority: common.filesystem_authority,
workdir_session,
workdir_sessions,
workspace_context: common.workspace_context,
flow_runtime_state: Arc::new(Mutex::new(restored_flow_runtime_state(
&state.extensions,
@@ -7413,7 +7488,9 @@ fn compact_failure_category(error: &WorkerError) -> CompactFailureCategory {
| WorkerError::SegmentActivationIncomplete { .. } => {
CompactFailureCategory::ActiveSegmentCommit
}
WorkerError::Store(_) => CompactFailureCategory::Storage,
WorkerError::Store(_) | WorkerError::PendingUsagePersistence { .. } => {
CompactFailureCategory::Storage
}
WorkerError::InvalidState(_) | WorkerError::Engine(_) => {
CompactFailureCategory::InternalWorker
}
@@ -7477,6 +7554,11 @@ pub enum WorkerError {
#[error(transparent)]
Provider(#[from] crate::model_client::ProviderError),
#[error(
"cannot replace the active Segment while {pending_records} usage record(s) await durable persistence"
)]
PendingUsagePersistence { pending_records: usize },
#[error("active Segment changed before compaction commit")]
CompactActiveSegmentChanged,
@@ -7580,18 +7662,26 @@ pub enum WorkerError {
},
}
fn workdir_session_from_authority(
fn workdir_sessions_from_authority(
authority: &WorkerFilesystemAuthority,
scope: &SharedScope,
) -> Option<WorkdirSessionHandle> {
authority.as_local().map(|local| {
Arc::new(LocalWorkdirSession::materialized(
) -> Arc<WorkdirSessionRouter> {
let router = Arc::new(WorkdirSessionRouter::new());
if let Some(local) = authority.as_local() {
let session = Arc::new(LocalWorkdirSession::materialized(
local.root.clone(),
local.cwd.clone(),
scope.clone(),
WorkdirSessionCapabilities::ALL,
)) as WorkdirSessionHandle
})
)) as WorkdirSessionHandle;
router
.attach(
WorkdirAttachmentAlias::new("workdir").expect("static alias is valid"),
session,
)
.expect("fresh Workdir router accepts its only attachment");
}
router
}
/// Bundle of resources that every high-level Worker constructor needs:
@@ -8534,6 +8624,117 @@ mod build_summary_prompt_tests {
}
}
#[tokio::test]
async fn failed_usage_append_fences_segment_rotation_until_retry_commits_once() {
let dir = tempfile::tempdir().unwrap();
let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap();
let mut worker = Worker::new(
minimal_manifest(),
Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient),
store.clone(),
WorkerWorkspaceContext::no_workspace(),
WorkerFilesystemAuthority::None,
Scope::empty(),
)
.await
.unwrap();
worker.ensure_segment_head().await.unwrap();
let session_id = worker.session_id();
let source_segment = worker.segment_id();
worker
.usage_tracker
.note_compaction_correlation_id("compact-id".into());
worker.usage_tracker.note_request(1);
worker.usage_tracker.record_usage(&agen::event::UsageEvent {
input_tokens: Some(10),
output_tokens: Some(1),
total_tokens: Some(11),
cache_read_input_tokens: Some(0),
cache_creation_input_tokens: Some(0),
});
worker.usage_tracker.request_completed();
let append_error = persist_pending_usage(
&worker.usage_tracker,
&worker.usage_history,
worker.compact_state.as_deref(),
|_, _, _| Err(StoreError::Io(std::io::Error::other("synthetic failure"))),
)
.expect_err("first usage append must fail");
assert!(append_error.to_string().contains("synthetic failure"));
assert_eq!(worker.usage_tracker.pending_record_count(), 1);
// Simulate another writer extending the active Segment. The normal run
// prelude must not auto-fork while the failed record still belongs to
// this Segment.
store
.append(
session_id,
source_segment,
&LogEntry::TurnEnd {
ts: segment_log::now_millis(),
turn_count: 99,
},
)
.unwrap();
let run_error = worker
.run_text("must stop before provider ownership")
.await
.expect_err("head drift must not fork past pending usage");
assert!(matches!(
run_error,
WorkerError::PendingUsagePersistence { pending_records: 1 }
));
assert_eq!(worker.segment_id(), source_segment);
assert_eq!(
store.list_segments(session_id).unwrap(),
vec![source_segment]
);
let error = worker
.compact(0)
.await
.expect_err("pending usage must fence Segment replacement");
assert!(matches!(
error,
WorkerError::PendingUsagePersistence { pending_records: 1 }
));
assert_eq!(worker.segment_id(), source_segment);
assert_eq!(
store.list_segments(session_id).unwrap(),
vec![source_segment]
);
let metrics = persist_pending_usage(
&worker.usage_tracker,
&worker.usage_history,
worker.compact_state.as_deref(),
|entry, after_commit, _| {
store.append(session_id, source_segment, &entry)?;
after_commit();
Ok(())
},
)
.expect("retry should commit on the originating Segment");
assert_eq!(worker.usage_tracker.pending_record_count(), 0);
assert_eq!(worker.usage_history.lock().unwrap().len(), 1);
assert_eq!(metrics.len(), 1, "correlation must be emitted once");
assert_eq!(metrics[0].name, "compact.post_request");
let usage_entries = store
.read_all(session_id, source_segment)
.unwrap()
.into_iter()
.filter(|entry| matches!(entry, LogEntry::LlmUsage { .. }))
.count();
assert_eq!(usage_entries, 1, "retry must not duplicate accounting");
assert_eq!(
store.list_segments(session_id).unwrap(),
vec![source_segment]
);
}
#[tokio::test]
async fn fresh_run_clears_interrupted_budget_before_pre_run_compaction() {
let dir = tempfile::tempdir().unwrap();
+385 -4
View File
@@ -7,15 +7,16 @@
//! (driven by `compact_request_threshold` → `PreRequestAction::Yield`)
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use agen::Engine;
use agen::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent, UsageEvent};
use agen::llm_client::types::Item;
use agen::llm_client::{ClientError, LlmClient, Request};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use async_trait::async_trait;
use futures::Stream;
use futures::{Stream, StreamExt};
use protocol::{Event, Method, RunResult};
use session_store::{
CombinedStore, FsStore, FsWorkerStore, LogEntry, Store, WorkerMetadata, WorkerMetadataStore,
@@ -125,10 +126,13 @@ fn annotated(item: Item) -> session_store::LoggedHistoryEntry {
}
}
type CallObserver = Arc<dyn Fn(usize) + Send + Sync>;
#[derive(Clone)]
struct MockClient {
responses: Arc<Vec<Vec<LlmEvent>>>,
call_count: Arc<AtomicUsize>,
observer: Arc<Mutex<Option<CallObserver>>>,
}
impl MockClient {
@@ -136,8 +140,13 @@ impl MockClient {
Self {
responses: Arc::new(responses),
call_count: Arc::new(AtomicUsize::new(0)),
observer: Arc::new(Mutex::new(None)),
}
}
fn set_observer(&self, observer: impl Fn(usize) + Send + Sync + 'static) {
*self.observer.lock().unwrap() = Some(Arc::new(observer));
}
}
#[async_trait]
@@ -152,6 +161,9 @@ impl LlmClient for MockClient {
) -> Result<Pin<Box<dyn Stream<Item = Result<LlmEvent, ClientError>> + Send>>, ClientError>
{
let count = self.call_count.fetch_add(1, Ordering::SeqCst);
if let Some(observer) = self.observer.lock().unwrap().clone() {
observer(count);
}
if count >= self.responses.len() {
return Err(ClientError::Config("mock client exhausted".into()));
}
@@ -161,6 +173,90 @@ impl LlmClient for MockClient {
}
}
#[derive(Clone)]
struct PauseAfterUsageClient {
responses: Arc<Vec<Vec<LlmEvent>>>,
call_count: Arc<AtomicUsize>,
hanging_call: usize,
}
impl PauseAfterUsageClient {
fn new(responses: Vec<Vec<LlmEvent>>, hanging_call: usize) -> Self {
Self {
responses: Arc::new(responses),
call_count: Arc::new(AtomicUsize::new(0)),
hanging_call,
}
}
}
#[async_trait]
impl LlmClient for PauseAfterUsageClient {
fn clone_boxed(&self) -> Box<dyn LlmClient> {
Box::new(self.clone())
}
async fn stream(
&self,
_request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<LlmEvent, ClientError>> + Send>>, ClientError>
{
let count = self.call_count.fetch_add(1, Ordering::SeqCst);
let events = self
.responses
.get(count)
.cloned()
.ok_or_else(|| ClientError::Config("pause mock client exhausted".into()))?;
let finite = futures::stream::iter(events.into_iter().map(Ok));
if count == self.hanging_call {
Ok(Box::pin(finite.chain(futures::stream::pending())))
} else {
Ok(Box::pin(finite))
}
}
}
#[derive(Clone)]
struct GrowTool;
#[async_trait]
impl Tool for GrowTool {
async fn execute(
&self,
_input: &str,
_ctx: agen::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
Ok(ToolOutput {
summary: "grew context".into(),
content: Some("x".repeat(400_000)),
attachments: Vec::new(),
})
}
}
fn grow_tool_definition() -> ToolDefinition {
Arc::new(|| {
(
ToolMeta::new("grow")
.description("test tool that grows request context")
.input_schema(serde_json::json!({"type": "object"})),
Arc::new(GrowTool) as Arc<dyn Tool>,
)
})
}
fn grow_tool_use_events(call_id: &str, input_tokens: u64) -> Vec<LlmEvent> {
vec![
LlmEvent::tool_use_start(0, call_id, "grow"),
LlmEvent::tool_input_delta(0, "{}"),
LlmEvent::tool_use_stop(0),
LlmEvent::usage(input_tokens, 1),
LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]
}
#[derive(Clone)]
struct BlockingCompactClient {
calls: Arc<AtomicUsize>,
@@ -365,12 +461,13 @@ where
let scope = worker::Scope::writable(&pwd).unwrap();
std::mem::forget(pwd_tmp);
let worker =
let mut engine =
Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client);
engine.register_tool(grow_tool_definition());
let observed_store = store.clone();
let mut worker = Worker::new(
manifest,
worker,
engine,
store,
worker::WorkerWorkspaceContext::local_filesystem(None),
worker::WorkerFilesystemAuthority::local(pwd.clone(), pwd.clone()),
@@ -379,6 +476,7 @@ where
.await
.unwrap();
worker.enable_worker_metadata_write_through().unwrap();
worker.attach_log_writer(Arc::new(worker.log_writer_handle()));
(worker, observed_store)
}
@@ -1057,6 +1155,289 @@ async fn request_threshold_compact_publishes_runtime_progress() {
assert_eq!(post.metric.correlation_id.as_deref(), Some(correlation_id));
}
#[tokio::test]
async fn completed_post_compact_usage_rearms_within_one_tool_loop() {
let client = MockClient::new(vec![
text_events_with_usage("seed", 100_000),
write_summary_tool_use_events("compact-1", "first summary"),
single_text_events("done"),
grow_tool_use_events("grow-1", 1_000),
write_summary_tool_use_events("compact-2", "second summary"),
single_text_events("done"),
text_events_with_usage("finished", 800),
]);
let call_count = Arc::clone(&client.call_count);
let observed_client = client.clone();
let manifest = MID_TURN_MANIFEST_TOML.replace(
"compact_request_threshold = 100",
"compact_threshold = 80000\ncompact_request_threshold = 50000",
);
let (mut worker, store) = make_worker_with_manifest_and_store(&manifest, client).await;
let session_id = worker.session_id();
let durable_at_next_request = Arc::new(AtomicBool::new(false));
let observed_barrier = Arc::clone(&durable_at_next_request);
let observed_store = store.clone();
observed_client.set_observer(move |call_index| {
if call_index != 4 {
return;
}
let usage_is_durable = observed_store
.list_segments(session_id)
.unwrap()
.into_iter()
.flat_map(|segment_id| observed_store.read_all(session_id, segment_id).unwrap())
.any(|entry| {
matches!(
entry,
LogEntry::LlmUsage {
input_total_tokens: 1_000,
..
}
)
});
assert!(
usage_is_durable,
"the first post-compact usage must be durable before the next provider request"
);
observed_barrier.store(true, Ordering::SeqCst);
});
worker.run_text("seed input").await.unwrap();
worker
.run_text("continue through two compactions")
.await
.expect("a successful post-compact request must rearm before the tool-loop threshold");
assert_eq!(
call_count.load(Ordering::SeqCst),
7,
"both compactor cycles and both normal post-compact requests must run"
);
assert!(durable_at_next_request.load(Ordering::SeqCst));
let metrics = session_metrics::read_session_metrics(&store, session_id).unwrap();
let compact_starts = metrics
.iter()
.filter(|record| record.metric.name == "compact.start")
.collect::<Vec<_>>();
assert_eq!(compact_starts.len(), 2);
assert!(compact_starts.iter().any(|record| {
record.metric.dimensions.get("trigger").map(String::as_str) == Some("pre_run")
}));
assert!(compact_starts.iter().any(|record| {
record.metric.dimensions.get("trigger").map(String::as_str) == Some("request_threshold")
}));
let compact_posts = metrics
.iter()
.filter(|record| record.metric.name == "compact.post_request")
.collect::<Vec<_>>();
assert_eq!(compact_posts.len(), 2);
assert_ne!(
compact_posts[0].metric.correlation_id,
compact_posts[1].metric.correlation_id
);
let usage = store
.list_segments(session_id)
.unwrap()
.into_iter()
.flat_map(|segment_id| store.read_all(session_id, segment_id).unwrap())
.filter_map(|entry| match entry {
LogEntry::LlmUsage {
input_total_tokens, ..
} => Some(input_total_tokens),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(usage.iter().filter(|tokens| **tokens == 1_000).count(), 1);
assert_eq!(usage.iter().filter(|tokens| **tokens == 800).count(), 1);
let restored = session_store::restore(&store, session_id, worker.segment_id()).unwrap();
assert_eq!(restored.usage_history.len(), 1);
assert_eq!(restored.usage_history[0].input_total_tokens, 800);
}
#[tokio::test]
async fn post_compact_pause_bills_partial_without_rearming_or_consuming_correlation() {
let client = PauseAfterUsageClient::new(
vec![
text_events_with_usage("seed", 100_000),
write_summary_tool_use_events("manual-compact", "pause test summary"),
single_text_events("done"),
vec![
LlmEvent::text_block_start(0),
LlmEvent::text_delta(0, "partial"),
LlmEvent::usage(90, 2),
],
grow_tool_use_events("grow-after-resume", 80),
write_summary_tool_use_events("automatic-compact", "rearmed summary"),
single_text_events("done"),
text_events_with_usage("finished", 50),
],
3,
);
let call_count = Arc::clone(&client.call_count);
let manifest = MID_TURN_MANIFEST_TOML.replace(
"compact_request_threshold = 100",
"compact_threshold = 80000\ncompact_request_threshold = 50000",
);
let (worker, store) = make_worker_with_manifest_and_store(&manifest, client).await;
let session_id = worker.session_id();
let runtime_tmp = tempfile::tempdir().unwrap();
let bash_output_dir = runtime_tmp.path().join("bash-output");
let (handle, shutdown_receiver) =
WorkerController::spawn(worker, runtime_tmp.path(), &bash_output_dir)
.await
.unwrap();
let mut rx = handle.subscribe();
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"seed history",
))
.await
.unwrap();
loop {
if matches!(
tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("timeout waiting for seed")
.expect("event"),
Event::RunEnd {
result: RunResult::Finished
}
) {
break;
}
}
handle
.send(Method::submit_text(
protocol::new_submission_request_id(),
"pause the first post-compact request",
))
.await
.unwrap();
tokio::time::timeout(std::time::Duration::from_secs(2), async {
while call_count.load(Ordering::SeqCst) < 4 {
tokio::task::yield_now().await;
}
})
.await
.expect("post-compact provider request should start");
handle
.send(Method::Pause {
command: protocol::WorkerCommandEnvelope::new(1),
})
.await
.unwrap();
loop {
let event = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("timeout waiting for pause")
.expect("event");
if matches!(
event,
Event::RunEnd {
result: RunResult::Paused
}
) {
break;
}
}
let paused_usage = store
.list_segments(session_id)
.unwrap()
.into_iter()
.flat_map(|segment_id| store.read_all(session_id, segment_id).unwrap())
.filter_map(|entry| match entry {
LogEntry::LlmUsage {
input_total_tokens, ..
} => Some(input_total_tokens),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(
paused_usage.iter().filter(|tokens| **tokens == 90).count(),
1
);
let paused_metrics = session_metrics::read_session_metrics(&store, session_id).unwrap();
assert!(
paused_metrics
.iter()
.all(|record| record.metric.name != "compact.post_request"),
"partial Pause usage must not consume compact correlation or rearm"
);
handle
.send(Method::Resume {
command: protocol::WorkerCommandEnvelope::new(2),
})
.await
.unwrap();
loop {
if matches!(
tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("timeout waiting for resume completion")
.expect("event"),
Event::RunEnd {
result: RunResult::Finished
}
) {
break;
}
}
assert_eq!(
call_count.load(Ordering::SeqCst),
8,
"the completed resumed request must rearm and permit request-threshold compaction"
);
let metrics = session_metrics::read_session_metrics(&store, session_id).unwrap();
let compact_starts = metrics
.iter()
.filter(|record| record.metric.name == "compact.start")
.collect::<Vec<_>>();
assert_eq!(compact_starts.len(), 2);
let pre_run_correlation = compact_starts
.iter()
.find(|record| {
record.metric.dimensions.get("trigger").map(String::as_str) == Some("pre_run")
})
.and_then(|record| record.metric.correlation_id.as_deref())
.expect("pre-run automatic compact correlation");
let pre_run_posts = metrics
.iter()
.filter(|record| {
record.metric.name == "compact.post_request"
&& record.metric.correlation_id.as_deref() == Some(pre_run_correlation)
})
.collect::<Vec<_>>();
assert_eq!(pre_run_posts.len(), 1);
assert_eq!(
pre_run_posts[0].metric.dimensions["input_total_tokens"],
"80"
);
assert!(metrics.iter().all(|record| {
record.metric.name != "compact.post_request"
|| record.metric.dimensions["input_total_tokens"] != "90"
}));
handle
.send(Method::Shutdown {
command: protocol::WorkerCommandEnvelope::new(3),
})
.await
.unwrap();
tokio::time::timeout(std::time::Duration::from_secs(2), shutdown_receiver)
.await
.expect("controller shutdown timeout")
.expect("shutdown confirmation");
}
#[tokio::test]
async fn compacted_context_above_request_threshold_fails_before_provider_request() {
let client = MockClient::new(vec![
+58 -8
View File
@@ -3,7 +3,9 @@ use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use agen::Engine;
use agen::llm_client::event::{ErrorEvent, Event as LlmEvent, ResponseStatus, StatusEvent};
use agen::llm_client::event::{
ErrorEvent, Event as LlmEvent, ResponseStatus, StatusEvent, UsageEvent,
};
use agen::llm_client::types::Item;
use agen::llm_client::{ClientError, LlmClient, Request};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
@@ -349,7 +351,7 @@ async fn shutdown_closes_bound_workdir_session() {
})
.await
.unwrap();
worker.bind_workdir_session(Some(Arc::clone(&session)));
worker.bind_single_workdir_session(Some(Arc::clone(&session)));
let runtime_base = tempfile::tempdir().unwrap();
let bash_output_dir = runtime_base.path().join("bash-output");
@@ -384,7 +386,7 @@ async fn controller_projects_workdir_command_events_and_snapshot_state() {
worker.scope().clone(),
WorkdirSessionCapabilities::ALL,
));
worker.bind_workdir_session(Some(Arc::clone(&session)));
worker.bind_single_workdir_session(Some(Arc::clone(&session)));
let handle = spawn_controller(worker).await;
let mut events = handle.subscribe();
@@ -495,7 +497,7 @@ async fn controller_refreshes_command_snapshot_after_high_output_provider_lag()
worker.scope().clone(),
WorkdirSessionCapabilities::ALL,
));
worker.bind_workdir_session(Some(Arc::clone(&session)));
worker.bind_single_workdir_session(Some(Arc::clone(&session)));
let handle = spawn_controller(worker).await;
// Local command telemetry uses 8 KiB chunks and a 256-event channel. One
@@ -572,7 +574,7 @@ async fn controller_startup_failure_closes_bound_workdir_session() {
worker.scope().clone(),
WorkdirSessionCapabilities::ALL,
));
worker.bind_workdir_session(Some(Arc::clone(&session)));
worker.bind_single_workdir_session(Some(Arc::clone(&session)));
let runtime_base = tempfile::tempdir().unwrap();
let invalid_runtime_base = runtime_base.path().join("not-a-directory");
std::fs::write(&invalid_runtime_base, "file").unwrap();
@@ -2580,18 +2582,32 @@ async fn drain_until<F: FnMut(&Event) -> bool>(
/// the interrupted turn's history consistency.
#[tokio::test]
async fn pause_then_resume_preserves_notifications_and_history_consistency() {
// Response 1: hang after opening a text block (no stop / completed),
// so the Engine is parked inside the stream read and `cancel_rx`
// races it cleanly on Method::Pause.
// Response 1: report billable partial usage, then hang after opening a text
// block (no stop / completed). Pause must flush that usage without treating
// the provider request as normally completed.
let hang = MockResponse::Hang(vec![
LlmEvent::text_block_start(0),
LlmEvent::text_delta(0, "partial..."),
LlmEvent::Usage(UsageEvent {
input_tokens: Some(90),
output_tokens: Some(2),
total_tokens: Some(92),
cache_read_input_tokens: Some(0),
cache_creation_input_tokens: Some(0),
}),
]);
// Response 2: a clean assistant reply delivered on Resume.
let ok = MockResponse::Complete(vec![
LlmEvent::text_block_start(0),
LlmEvent::text_delta(0, "resumed output"),
LlmEvent::text_block_stop(0, None),
LlmEvent::Usage(UsageEvent {
input_tokens: Some(100),
output_tokens: Some(3),
total_tokens: Some(103),
cache_read_input_tokens: Some(0),
cache_creation_input_tokens: Some(0),
}),
LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
@@ -2645,6 +2661,23 @@ async fn pause_then_resume_preserves_notifications_and_history_consistency() {
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
assert_eq!(handle.shared_state.catalog_status(), WorkerStatus::Paused);
let (paused_entries, _) = handle.sink.subscribe_with_snapshot();
let paused_usage = paused_entries
.iter()
.filter_map(|entry| match entry {
LogEntry::LlmUsage {
input_total_tokens,
output_tokens,
..
} => Some((*input_total_tokens, *output_tokens)),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(
paused_usage,
vec![(90, 2)],
"Pause must durably account for usage already observed on the interrupted request"
);
handle
.send(Method::Notify {
@@ -2729,6 +2762,23 @@ async fn pause_then_resume_preserves_notifications_and_history_consistency() {
})
.unwrap_or_default();
assert_eq!(assistant_text, "resumed output");
let (final_entries, _) = handle.sink.subscribe_with_snapshot();
let final_usage = final_entries
.iter()
.filter_map(|entry| match entry {
LogEntry::LlmUsage {
input_total_tokens,
output_tokens,
..
} => Some((*input_total_tokens, *output_tokens)),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(
final_usage,
vec![(90, 2), (100, 3)],
"the paused request must remain a distinct billing record rather than merging into resume"
);
let has_tool_call = history.iter().any(|i| i.is_tool_call());
assert!(!has_tool_call, "no orphan tool_call in history");
}
+162 -41
View File
@@ -32,8 +32,9 @@ use worker_runtime::catalog::{
ConfigBundleRef, CreateWorkerRequest, ProfileSelector, ProfileSourceArchiveSource,
RepositoryRefObservation, RepositoryRefObservationRequest,
WorkerDetail as EmbeddedWorkerDetail, WorkerStatus as EmbeddedWorkerStatus,
WorkingDirectoryClaim, WorkingDirectoryRepositoryAccessRequest, WorkingDirectoryRequest,
WorkingDirectoryStatus, WorkingDirectorySummary, WorkspaceApiRef,
WorkingDirectoryAttachmentClaim, WorkingDirectoryAttachmentRequest,
WorkingDirectoryAttachmentStatus, WorkingDirectoryRepositoryAccessRequest,
WorkingDirectoryRequest, WorkingDirectoryStatus, WorkspaceApiRef,
};
use worker_runtime::config_bundle::{ConfigBundle, ConfigBundleAvailability, ConfigBundleSummary};
#[cfg(test)]
@@ -270,8 +271,8 @@ pub struct WorkerSummary {
pub retention_state: String,
pub implementation: WorkerImplementationSummary,
pub capabilities: WorkerCapabilitySummary,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub working_directory: Option<WorkingDirectorySummary>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub workdir_attachments: Vec<WorkingDirectoryAttachmentStatus>,
pub diagnostics: Vec<RuntimeDiagnostic>,
}
@@ -336,7 +337,7 @@ impl From<RuntimeSummary> for server_api::RuntimeSummary {
pub(crate) fn workspace_worker_summary(
summary: WorkerSummary,
resource_key: String,
working_directory: Option<server_api::WorkingDirectorySummary>,
workdir_attachments: Vec<server_api::WorkerWorkdirAttachmentSummary>,
) -> server_api::WorkerSummary {
server_api::WorkerSummary {
runtime_id: summary.worker.runtime_id,
@@ -366,7 +367,7 @@ pub(crate) fn workspace_worker_summary(
can_stop: summary.capabilities.can_stop,
can_spawn_followup: summary.capabilities.can_spawn_followup,
},
working_directory,
workdir_attachments,
diagnostics: summary.diagnostics.into_iter().map(Into::into).collect(),
}
}
@@ -483,6 +484,13 @@ pub struct WorkerCreateBinding {
pub create_fingerprint: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct WorkerSpawnWorkingDirectoryAttachmentRequest {
pub alias: workdir::WorkdirAttachmentAlias,
pub working_directory: WorkerSpawnWorkingDirectoryRequest,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct WorkerSpawnRequest {
@@ -495,15 +503,13 @@ pub struct WorkerSpawnRequest {
pub ticket_assignment: Option<WorkerTicketAssignmentRequest>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub initial_submit: Vec<Segment>,
/// Optional safe working-directory creation request. The Workspace server resolves
/// this into a runtime-internal `WorkingDirectoryRequest` from configured
/// repositories before calling a host.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub working_directory_request: Option<WorkerSpawnWorkingDirectoryRequest>,
/// Safe working-directory materialization requests, keyed by stable Worker-local aliases.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub workdir_attachment_requests: Vec<WorkerSpawnWorkingDirectoryAttachmentRequest>,
#[serde(skip, default)]
pub resolved_working_directory_request: Option<WorkingDirectoryRequest>,
pub resolved_workdir_attachment_requests: Vec<WorkingDirectoryAttachmentRequest>,
#[serde(skip, default)]
pub resolved_working_directory: Option<WorkingDirectoryClaim>,
pub resolved_workdir_attachments: Vec<WorkingDirectoryAttachmentClaim>,
#[serde(skip, default)]
pub resolved_config_bundle: Option<ConfigBundle>,
#[serde(skip, default)]
@@ -905,6 +911,24 @@ pub trait WorkspaceWorkerRuntime: Send + Sync {
}
}
fn replace_worker_workdir_attachments(
&self,
worker_id: &str,
_attachments: Vec<WorkingDirectoryAttachmentClaim>,
) -> WorkerWorkspaceApiResult {
WorkerWorkspaceApiResult {
state: WorkerOperationState::Unsupported,
worker: None,
diagnostics: vec![diagnostic(
"worker_workdir_attachments_replace_unsupported",
DiagnosticSeverity::Info,
format!(
"runtime does not support replacing Workdir attachments for worker `{worker_id}`"
),
)],
}
}
fn create_working_directory(
&self,
_request: WorkingDirectoryRequest,
@@ -1489,6 +1513,19 @@ impl RuntimeRegistry {
Ok(runtime.replace_worker_workspace_api(worker_id, workspace_api))
}
pub fn replace_worker_workdir_attachments(
&self,
worker: &RuntimeWorkerRef,
attachments: Vec<WorkingDirectoryAttachmentClaim>,
) -> Result<WorkerWorkspaceApiResult, RuntimeRegistryError> {
let runtime_id = worker.runtime_id.as_str();
let worker_id = worker.worker_id.as_str();
validate_backend_identifier("runtime_id", runtime_id)?;
validate_backend_identifier("worker_id", worker_id)?;
let runtime = self.runtime(runtime_id)?;
Ok(runtime.replace_worker_workdir_attachments(worker_id, attachments))
}
pub fn observe_workspace_prompt_projection(
&self,
projection: worker::WorkspacePromptProjection,
@@ -2176,7 +2213,7 @@ impl EmbeddedWorkerRuntime {
&& self.can_stop_embedded_worker(summary.status),
can_spawn_followup: false,
},
working_directory: summary.working_directory.map(|status| status.summary),
workdir_attachments: summary.workdir_attachments,
diagnostics: embedded_worker_projection_diagnostics(
summary.execution_metadata_available,
),
@@ -2220,7 +2257,7 @@ impl EmbeddedWorkerRuntime {
&& self.can_stop_embedded_worker(detail.status),
can_spawn_followup: false,
},
working_directory: detail.working_directory.map(|status| status.summary),
workdir_attachments: detail.workdir_attachments,
diagnostics: embedded_worker_projection_diagnostics(
detail.execution_metadata_available,
),
@@ -2429,6 +2466,39 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
}
}
fn replace_worker_workdir_attachments(
&self,
worker_id: &str,
attachments: Vec<WorkingDirectoryAttachmentClaim>,
) -> WorkerWorkspaceApiResult {
let Some(worker_ref) = self.worker_ref(worker_id) else {
return WorkerWorkspaceApiResult {
state: WorkerOperationState::Rejected,
worker: None,
diagnostics: vec![diagnostic(
"embedded_worker_id_invalid",
DiagnosticSeverity::Warning,
"Worker id was empty and cannot change Workdir attachments".to_string(),
)],
};
};
match self
.runtime
.replace_worker_workdir_attachments(&worker_ref, attachments)
{
Ok(detail) => WorkerWorkspaceApiResult {
state: WorkerOperationState::Accepted,
worker: Some(self.map_worker_detail(detail)),
diagnostics: Vec::new(),
},
Err(err) => WorkerWorkspaceApiResult {
state: WorkerOperationState::Rejected,
worker: None,
diagnostics: vec![embedded_runtime_diagnostic(&err)],
},
}
}
fn create_working_directory(
&self,
_request: WorkingDirectoryRequest,
@@ -2491,8 +2561,8 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
request: WorkerSpawnRequest,
) -> WorkerSpawnResult {
let mut diagnostics = Vec::new();
if request.resolved_working_directory_request.is_some()
|| request.resolved_working_directory.is_some()
if !request.resolved_workdir_attachment_requests.is_empty()
|| !request.resolved_workdir_attachments.is_empty()
{
diagnostics.push(embedded_workdir_unsupported_diagnostic());
return WorkerSpawnResult {
@@ -2573,8 +2643,8 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
config_bundle,
profile_source,
initial_input: initial_worker_input(&request.initial_submit),
working_directory_request: request.resolved_working_directory_request.clone(),
working_directory: request.resolved_working_directory.clone(),
workdir_attachment_requests: request.resolved_workdir_attachment_requests.clone(),
workdir_attachments: request.resolved_workdir_attachments.clone(),
worker_observation_enabled: request.resolved_worker_observation_enabled,
worker_observation_grants: request.resolved_worker_observation_grants.clone(),
workspace_api: Some(workspace_api),
@@ -3982,7 +4052,7 @@ impl RemoteWorkerRuntime {
&& runtime_worker_can_stop(true, summary.status),
can_spawn_followup: false,
},
working_directory: summary.working_directory.map(|status| status.summary),
workdir_attachments: summary.workdir_attachments,
diagnostics: remote_worker_projection_diagnostics(summary.execution_metadata_available),
}
}
@@ -4024,7 +4094,7 @@ impl RemoteWorkerRuntime {
&& runtime_worker_can_stop(true, detail.status),
can_spawn_followup: false,
},
working_directory: detail.working_directory.map(|status| status.summary),
workdir_attachments: detail.workdir_attachments,
diagnostics: remote_worker_projection_diagnostics(detail.execution_metadata_available),
}
}
@@ -4388,6 +4458,51 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime {
}
}
fn replace_worker_workdir_attachments(
&self,
worker_id: &str,
attachments: Vec<WorkingDirectoryAttachmentClaim>,
) -> WorkerWorkspaceApiResult {
let workdir_attachments = match runtime_contract_convert(attachments) {
Ok(attachments) => attachments,
Err(diagnostic) => {
return WorkerWorkspaceApiResult {
state: WorkerOperationState::Rejected,
worker: None,
diagnostics: vec![diagnostic],
};
}
};
let request = runtime_api::WorkerWorkdirAttachmentsRequest {
workdir_attachments,
};
let worker_id = worker_id.to_string();
match self
.run_runtime_api(
self.request_timeout,
MAX_REMOTE_RUNTIME_RESPONSE_BYTES,
move |client| async move {
client
.replace_worker_workdir_attachments(worker_id, request)
.await
},
)
.and_then(|value| {
runtime_contract_convert::<_, worker_runtime::catalog::WorkerDetail>(value.worker)
}) {
Ok(response) => WorkerWorkspaceApiResult {
state: WorkerOperationState::Accepted,
worker: Some(self.map_worker_detail(response)),
diagnostics: Vec::new(),
},
Err(diagnostic) => WorkerWorkspaceApiResult {
state: WorkerOperationState::Rejected,
worker: None,
diagnostics: vec![diagnostic],
},
}
}
fn create_working_directory(
&self,
request: WorkingDirectoryRequest,
@@ -4567,8 +4682,8 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime {
config_bundle,
profile_source,
initial_input: initial_worker_input(&request.initial_submit),
working_directory_request: request.resolved_working_directory_request.clone(),
working_directory: request.resolved_working_directory.clone(),
workdir_attachment_requests: request.resolved_workdir_attachment_requests.clone(),
workdir_attachments: request.resolved_workdir_attachments.clone(),
worker_observation_enabled: request.resolved_worker_observation_enabled,
worker_observation_grants: request.resolved_worker_observation_grants.clone(),
workspace_api: Some(workspace_api),
@@ -5724,7 +5839,7 @@ pub fn placeholder_worker(host_id: impl Into<String>) -> WorkerSummary {
can_stop: false,
can_spawn_followup: false,
},
working_directory: None,
workdir_attachments: Vec::new(),
diagnostics: vec![diagnostic(
"runtime_capability_unsupported",
DiagnosticSeverity::Info,
@@ -6265,10 +6380,16 @@ mod tests {
worker_state: protocol::WorkerStateSnapshot {
..protocol::WorkerStatus::Idle.into()
},
working_directory: request
.working_directory
.as_ref()
.map(|binding| binding.status()),
workdir_attachments: request
.workdir_attachments
.iter()
.map(|(alias, binding)| {
worker_runtime::catalog::WorkingDirectoryAttachmentStatus {
alias: alias.clone(),
working_directory: binding.status(),
}
})
.collect(),
}
}
@@ -6360,7 +6481,7 @@ mod tests {
can_stop: false,
can_spawn_followup: false,
},
working_directory: None,
workdir_attachments: Vec::new(),
diagnostics: Vec::new(),
}],
observed_prompt_revisions: Arc::new(Mutex::new(Vec::new())),
@@ -6650,9 +6771,9 @@ mod tests {
profile: ProfileSelector::Builtin("builtin:coder".to_string()),
ticket_assignment: None,
initial_submit: Vec::new(),
working_directory_request: None,
resolved_working_directory_request: None,
resolved_working_directory: None,
workdir_attachment_requests: Vec::new(),
resolved_workdir_attachment_requests: Vec::new(),
resolved_workdir_attachments: Vec::new(),
resolved_config_bundle: None,
resolved_worker_observation_enabled: false,
resolved_worker_observation_grants: Vec::new(),
@@ -6890,9 +7011,9 @@ mod tests {
profile: ProfileSelector::Builtin("builtin:coder".to_string()),
ticket_assignment: None,
initial_submit: Vec::new(),
working_directory_request: None,
resolved_working_directory_request: None,
resolved_working_directory: None,
workdir_attachment_requests: Vec::new(),
resolved_workdir_attachment_requests: Vec::new(),
resolved_workdir_attachments: Vec::new(),
resolved_config_bundle: None,
resolved_worker_observation_enabled: false,
resolved_worker_observation_grants: Vec::new(),
@@ -6990,9 +7111,9 @@ mod tests {
profile: ProfileSelector::Builtin("builtin:coder".to_string()),
ticket_assignment: None,
initial_submit: Vec::new(),
working_directory_request: None,
resolved_working_directory_request: None,
resolved_working_directory: None,
workdir_attachment_requests: Vec::new(),
resolved_workdir_attachment_requests: Vec::new(),
resolved_workdir_attachments: Vec::new(),
resolved_config_bundle: None,
resolved_worker_observation_enabled: false,
resolved_worker_observation_grants: Vec::new(),
@@ -7029,9 +7150,9 @@ mod tests {
profile: ProfileSelector::Builtin("builtin:companion".to_string()),
ticket_assignment: None,
initial_submit: Vec::new(),
working_directory_request: None,
resolved_working_directory_request: None,
resolved_working_directory: None,
workdir_attachment_requests: Vec::new(),
resolved_workdir_attachment_requests: Vec::new(),
resolved_workdir_attachments: Vec::new(),
resolved_config_bundle: None,
resolved_worker_observation_enabled: false,
resolved_worker_observation_grants: Vec::new(),
@@ -667,6 +667,7 @@ CREATE TABLE workdir_create_credential_revision_retentions (
CREATE TABLE "workdir_registry" (
workspace_id TEXT NOT NULL,
workdir_id TEXT NOT NULL,
display_name TEXT,
runtime_id TEXT NOT NULL,
repository_id TEXT NOT NULL,
creation_selector TEXT,
@@ -859,10 +860,10 @@ CREATE TABLE "worker_workdir_links" (
runtime_id TEXT NOT NULL,
worker_id TEXT NOT NULL,
workdir_id TEXT NOT NULL,
role TEXT NOT NULL,
alias TEXT NOT NULL,
linked_at TEXT NOT NULL,
unlinked_at TEXT,
PRIMARY KEY (workspace_id, worker_id, workdir_id, role),
PRIMARY KEY (workspace_id, worker_id, workdir_id, alias),
FOREIGN KEY (workspace_id, worker_id)
REFERENCES "worker_registry"(workspace_id, worker_id) ON DELETE CASCADE,
FOREIGN KEY (workspace_id, workdir_id)
@@ -1093,8 +1094,8 @@ CREATE INDEX worker_removal_operations_worker_idx ON worker_removal_operations(w
CREATE UNIQUE INDEX worker_workdir_links_active_workdir_unique
ON worker_workdir_links(workspace_id, workdir_id)
WHERE unlinked_at IS NULL;
CREATE UNIQUE INDEX worker_workdir_links_active_worker_unique
ON worker_workdir_links(workspace_id, worker_id)
CREATE UNIQUE INDEX worker_workdir_links_active_alias_unique
ON worker_workdir_links(workspace_id, worker_id, alias)
WHERE unlinked_at IS NULL;
CREATE INDEX worker_workdir_links_workdir
ON worker_workdir_links(workspace_id, workdir_id);
+132 -5
View File
@@ -2,7 +2,10 @@
//! Runtime receives only resolved dispositions and stable ids; policy authority
//! never comes from prompts, profiles, or model input.
use crate::{Error as StoreError, store::SqliteWorkspaceStore};
use crate::{
Error as StoreError,
store::{SqliteWorkspaceStore, WorkerRemovalCommit, commit_worker_catalog_removal},
};
use chrono::Utc;
use rusqlite::{Connection, OptionalExtension, TransactionBehavior, params};
use serde::{Deserialize, Serialize};
@@ -417,13 +420,19 @@ impl SqliteWorkspaceStore {
operation_id: &str,
fp: &str,
result: &WorkerRetentionExecutionResult,
) -> Result<WorkerRemovalPlan, WorkerRetentionError> {
) -> Result<WorkerRemovalCommit, WorkerRetentionError> {
self.with_conn_mut(|conn|{
let tx=conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
let mut plan=load_plan_op(&tx,operation_id)?.ok_or_else(||StoreError::InvalidInput("operation missing".into()))?;
if plan.workspace_id!=workspace_id{return Err(StoreError::InvalidInput("cross-workspace".into()));}
if plan.input_fingerprint!=fp||result.input_fingerprint!=fp||result.operation_id!=operation_id{return Err(StoreError::InvalidInput(format!("fingerprint:{operation_id}")));}
if plan.state==WorkerRemovalPlanState::Succeeded{tx.commit()?;return Ok(plan);}
if plan.state==WorkerRemovalPlanState::Succeeded{
// Repair projection fences left by the pre-atomic implementation and
// return a replayable change for a broadcast missed after commit.
let catalog=commit_worker_catalog_removal(&tx,workspace_id,&plan.worker)?;
tx.commit()?;
return Ok(WorkerRemovalCommit { plan, catalog });
}
if plan.state != WorkerRemovalPlanState::Executing {
return Err(StoreError::InvalidInput(format!("stale:{}:plan state {} is not committable", plan.plan_id, state_s(plan.state))));
}
@@ -460,12 +469,14 @@ impl SqliteWorkspaceStore {
if plan.metadata_disposition==MetadataDisposition::Tombstone{
tx.execute("INSERT OR IGNORE INTO worker_tombstones(workspace_id,runtime_id,worker_id,display_name,profile,worker_created_at,removed_at,archive_id,policy_id,policy_revision,operation_id) VALUES(?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11)",params![workspace_id,plan.worker.runtime_id,plan.worker.worker_id,worker.display_name,worker.profile,worker.created_at,now,plan.archive_id,plan.policy_id,plan.policy_revision,operation_id])?;
}
tx.execute("UPDATE worker_workdir_links SET unlinked_at=?4 WHERE workspace_id=?1 AND runtime_id=?2 AND worker_id=?3 AND unlinked_at IS NULL",params![workspace_id,plan.worker.runtime_id,plan.worker.worker_id,now])?;
let deleted=tx.execute("DELETE FROM worker_registry WHERE workspace_id=?1 AND runtime_id=?2 AND worker_id=?3 AND updated_at=?4",params![workspace_id,plan.worker.runtime_id,plan.worker.worker_id,plan.worker_revision])?;
tx.execute("UPDATE worker_create_reservations SET state='removed',updated_at=?4 WHERE workspace_id=?1 AND runtime_id=?2 AND worker_id=?3 AND state='created'",params![workspace_id,plan.worker.runtime_id,plan.worker.worker_id,now])?;
if deleted!=1{return Err(StoreError::InvalidInput(format!("stale:{}:removal fence changed",plan.plan_id)));}
let catalog=commit_worker_catalog_removal(&tx,workspace_id,&plan.worker)?;
tx.execute("UPDATE worker_removal_operations SET state='succeeded',failure_category=NULL,updated_at=?1 WHERE operation_id=?2",params![now,operation_id])?;
tx.execute("INSERT OR IGNORE INTO worker_retention_audit_events(event_id,operation_id,workspace_id,event_kind,detail,created_at) VALUES(?1,?2,?3,'worker_removed',?4,?5)",params![stable("wre",operation_id),operation_id,workspace_id,format!("runtime_id={} worker_id={} session={} metadata={} diagnostics={}",plan.worker.runtime_id,plan.worker.worker_id,sess(plan.session_disposition),meta(plan.metadata_disposition),diag(plan.diagnostics_disposition)),now])?;
tx.commit()?;plan.state=WorkerRemovalPlanState::Succeeded;plan.updated_at=now;Ok(plan)
tx.commit()?;plan.state=WorkerRemovalPlanState::Succeeded;plan.updated_at=now;Ok(WorkerRemovalCommit { plan, catalog })
}).map_err(map_error)
}
@@ -948,7 +959,9 @@ fn parse_state(v: &str) -> rusqlite::Result<WorkerRemovalPlanState> {
#[cfg(test)]
mod tests {
use super::*;
use crate::store::{ControlPlaneStore, TicketCoderAssignmentRecord, WorkerRegistryRecord};
use crate::store::{
ControlPlaneStore, TicketCoderAssignmentRecord, WorkerCatalogChange, WorkerRegistryRecord,
};
use worker_runtime::identity::WorkerId;
fn worker_id() -> WorkerId {
WorkerId::from_legacy_u64(1)
@@ -1203,6 +1216,7 @@ mod tests {
assert_eq!(
s.commit_worker_removal("w", &p.operation_id, &p.input_fingerprint, &result)
.unwrap()
.plan
.state,
WorkerRemovalPlanState::Succeeded
);
@@ -1219,6 +1233,7 @@ mod tests {
assert_eq!(
s.commit_worker_removal("w", &p.operation_id, &p.input_fingerprint, &result)
.unwrap()
.plan
.state,
WorkerRemovalPlanState::Succeeded
);
@@ -1351,6 +1366,98 @@ mod tests {
.unwrap();
assert!(s.worker_tombstone("w", &p.worker).unwrap().is_none());
}
#[test]
fn terminal_catalog_failure_rolls_back_removal_projection_and_operation() {
let store = setup();
let purge = WorkerRetentionPolicyUpdate {
policy_id: "purge".into(),
session_disposition: SessionDisposition::Purge,
metadata_disposition: MetadataDisposition::Tombstone,
archive_retention: ArchiveRetention::Forever,
diagnostics_disposition: DiagnosticsDisposition::Purge,
diagnostics_retention_seconds: None,
};
store
.update_worker_retention_policy("w", 1, &purge)
.unwrap();
let plan = store.plan_worker_removal(&req(), &inv()).unwrap();
store
.prepare_worker_removal_execution("w", &plan.plan_id, &plan.input_fingerprint)
.unwrap();
let result = WorkerRetentionExecutionResult {
operation_id: plan.operation_id.clone(),
input_fingerprint: plan.input_fingerprint.clone(),
expected_worker_revision: plan.worker_revision.clone(),
worker_id: worker_id(),
session_disposition: SessionDisposition::Purge,
diagnostics_disposition: DiagnosticsDisposition::Purge,
archive: None,
source_removed: true,
diagnostics_retained: false,
};
store
.with_conn(|conn| {
conn.execute_batch(
"CREATE TRIGGER fail_worker_projection_removal
BEFORE INSERT ON worker_registry_projection_removals
BEGIN SELECT RAISE(ABORT, 'injected projection failure'); END;",
)?;
Ok(())
})
.unwrap();
assert!(
store
.commit_worker_removal("w", &plan.operation_id, &plan.input_fingerprint, &result,)
.is_err()
);
let (registry, removal, state, tombstone): (i64, i64, String, i64) = store
.with_conn(|conn| {
Ok((
conn.query_row(
"SELECT COUNT(*) FROM worker_registry WHERE workspace_id='w' AND runtime_id='r' AND worker_id=?1",
[worker_id().to_string()],
|row| row.get(0),
)?,
conn.query_row(
"SELECT COUNT(*) FROM worker_registry_projection_removals WHERE workspace_id='w' AND runtime_id='r' AND worker_id=?1",
[worker_id().to_string()],
|row| row.get(0),
)?,
conn.query_row(
"SELECT state FROM worker_removal_operations WHERE operation_id=?1",
[&plan.operation_id],
|row| row.get(0),
)?,
conn.query_row(
"SELECT COUNT(*) FROM worker_tombstones WHERE workspace_id='w' AND runtime_id='r' AND worker_id=?1",
[worker_id().to_string()],
|row| row.get(0),
)?,
))
})
.unwrap();
assert_eq!(registry, 1, "catalog deletion must roll back");
assert_eq!(removal, 0, "removal fence must not partially commit");
assert_eq!(state, "executing", "operation must remain recoverable");
assert_eq!(tombstone, 0, "retention metadata must roll back");
store
.with_conn(|conn| {
conn.execute_batch("DROP TRIGGER fail_worker_projection_removal;")?;
Ok(())
})
.unwrap();
let committed = store
.commit_worker_removal("w", &plan.operation_id, &plan.input_fingerprint, &result)
.unwrap();
assert!(matches!(
committed.catalog.changes.as_slice(),
[WorkerCatalogChange::Removed(worker)] if worker == &plan.worker
));
assert_eq!(committed.plan.state, WorkerRemovalPlanState::Succeeded);
}
#[test]
fn commit_requires_executing_state_and_exact_runtime_result() {
let store = setup();
@@ -1597,6 +1704,26 @@ mod tests {
.unwrap()
.is_none()
);
s.with_conn(|conn| {
conn.execute(
"DELETE FROM worker_registry_projection_removals WHERE workspace_id='w' AND runtime_id='r' AND worker_id=?1",
[worker_id().to_string()],
)?;
Ok(())
})
.unwrap();
let repaired = s
.recover_succeeded_worker_removal_catalog("w", &request.worker)
.unwrap();
assert!(matches!(
repaired.changes.as_slice(),
[WorkerCatalogChange::Removed(worker)] if worker == &request.worker
));
assert!(
s.get_worker_registry("w", &request.worker)
.unwrap()
.is_none()
);
let recovered = s
.recover_worker_removal_execution("w", &request.worker)
.unwrap()
@@ -22,7 +22,7 @@ impl WorkerExecutionBackend for TestExecutionBackend {
protocol::WorkerStateSnapshot {
..protocol::WorkerStatus::Idle.into()
},
None,
Vec::new(),
)
}
@@ -64,8 +64,8 @@ fn create_request(name: &str) -> CreateWorkerRequest {
.unwrap(),
},
initial_input: None,
working_directory_request: None,
working_directory: None,
workdir_attachment_requests: Vec::new(),
workdir_attachments: Vec::new(),
worker_observation_enabled: false,
worker_observation_grants: Vec::new(),
workspace_api: None,
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -25,6 +25,7 @@ pub fn request_fingerprint(
repository_id: &str,
selector: Option<&str>,
requested_runtime_id: Option<&str>,
display_name: Option<&str>,
repository_source_fingerprint: &str,
repository_source_revision: u64,
) -> String {
@@ -33,6 +34,7 @@ pub fn request_fingerprint(
Some(repository_id),
selector,
requested_runtime_id,
display_name,
Some(repository_source_fingerprint),
] {
match value {
@@ -569,6 +571,7 @@ mod tests {
"main",
Some("develop"),
None,
None,
"sha256:test",
1,
),
@@ -718,7 +721,7 @@ mod tests {
);
let mut changed_input = record.clone();
changed_input.request_fingerprint =
request_fingerprint("main", Some("main"), None, "sha256:test", 1);
request_fingerprint("main", Some("main"), None, None, "sha256:test", 1);
assert!(
store
.reserve_workdir_create_operation(&changed_input)
+16 -14
View File
@@ -716,7 +716,7 @@ fn load_workdir_record(
workdir_id: &str,
) -> Result<Option<WorkdirRegistryRecord>> {
conn.query_row(
r#"SELECT workspace_id, workdir_id, runtime_id, repository_id,
r#"SELECT workspace_id, workdir_id, display_name, runtime_id, repository_id,
creation_selector, creation_ref, creation_tree,
current_selector, current_ref, current_tree, observed_at_epoch_seconds,
materialization_status, cleanliness, created_at, updated_at
@@ -726,19 +726,20 @@ fn load_workdir_record(
Ok(WorkdirRegistryRecord {
workspace_id: row.get(0)?,
workdir_id: row.get(1)?,
runtime_id: row.get(2)?,
repository_id: row.get(3)?,
creation_selector: row.get(4)?,
creation_ref: row.get(5)?,
creation_tree: row.get(6)?,
current_selector: row.get(7)?,
current_ref: row.get(8)?,
current_tree: row.get(9)?,
observed_at_epoch_seconds: row.get::<_, Option<i64>>(10)?.map(|value| value as u64),
materialization_status: row.get(11)?,
cleanliness: row.get(12)?,
created_at: row.get(13)?,
updated_at: row.get(14)?,
display_name: row.get(2)?,
runtime_id: row.get(3)?,
repository_id: row.get(4)?,
creation_selector: row.get(5)?,
creation_ref: row.get(6)?,
creation_tree: row.get(7)?,
current_selector: row.get(8)?,
current_ref: row.get(9)?,
current_tree: row.get(10)?,
observed_at_epoch_seconds: row.get::<_, Option<i64>>(11)?.map(|value| value as u64),
materialization_status: row.get(12)?,
cleanliness: row.get(13)?,
created_at: row.get(14)?,
updated_at: row.get(15)?,
})
},
)
@@ -934,6 +935,7 @@ mod tests {
let workdir = WorkdirRegistryRecord {
workspace_id: "workspace-a".to_string(),
workdir_id: "workdir-a".to_string(),
display_name: None,
runtime_id: "runtime-a".to_string(),
repository_id: "repository-a".to_string(),
creation_selector: Some("refs/heads/develop".to_string()),
@@ -11,18 +11,12 @@ use tokio::task::JoinHandle;
use crate::runtime_subscription::{
BrokerSubscriptionEvent, RuntimeSubscriptionBroker, RuntimeSubscriptionBrokerError,
};
use crate::store::{
ControlPlaneStore, WorkerRegistryProjectionCommit, WorkerRegistryProjectionRecord,
};
use crate::store::{ControlPlaneStore, WorkerCatalogChange, WorkerRegistryProjectionCommit};
use worker_runtime::identity::RuntimeWorkerRef;
const PROJECTION_EVENT_CAPACITY: usize = 256;
#[derive(Debug, Clone)]
pub enum WorkerProjectionChange {
Upsert(WorkerRegistryProjectionRecord),
Removed(RuntimeWorkerRef),
}
pub type WorkerProjectionChange = WorkerCatalogChange;
#[derive(Debug, Clone)]
pub struct WorkerProjectionEvent {
@@ -144,27 +138,13 @@ impl WorkerProjectionService {
self.publish_commit(commit)
}
pub fn publish_catalog_change(&self, worker: &RuntimeWorkerRef) -> crate::Result<()> {
pub fn refresh(&self, worker: &RuntimeWorkerRef) -> crate::Result<()> {
let commit = self
.store
.publish_worker_registry_catalog_change(self.workspace_id.as_str(), worker)?;
.commit_worker_registry_projection_refresh(self.workspace_id.as_str(), worker)?;
self.publish_commit(commit)
}
pub fn publish_removed(&self, worker: RuntimeWorkerRef) -> crate::Result<()> {
let commit = self
.store
.publish_worker_registry_removal(self.workspace_id.as_str(), &worker)?;
if commit.changed_workers.is_empty() {
return Ok(());
}
let _ = self.events.send(WorkerProjectionEvent {
revision: commit.revision,
changes: vec![WorkerProjectionChange::Removed(worker)],
});
Ok(())
}
async fn apply_broker_event(
&self,
runtime_id: &str,
@@ -236,38 +216,21 @@ impl WorkerProjectionService {
self.publish_commit(commit)
}
fn publish_commit(&self, commit: WorkerRegistryProjectionCommit) -> crate::Result<()> {
if commit.changed_workers.is_empty() {
/// Broadcast an already committed catalog payload. Persistence always happens
/// before this call; a process crash therefore closes subscribers, whose next
/// subscription starts from the durable snapshot. Retrying a removal returns
/// the same durable change and can heal an in-process missed broadcast.
pub(crate) fn publish_commit(
&self,
commit: WorkerRegistryProjectionCommit,
) -> crate::Result<()> {
if commit.changes.is_empty() {
return Ok(());
}
let changes = commit
.changed_workers
.iter()
.filter_map(|worker| {
match self
.store
.worker_registry_projection(self.workspace_id.as_str(), worker)
{
Ok(Some(record)) => Some(WorkerProjectionChange::Upsert(record)),
Ok(None) => None,
Err(error) => {
tracing::warn!(
runtime_id = %worker.runtime_id,
worker_id = %worker.worker_id,
error = %error,
"failed to load committed Worker projection"
);
None
}
}
})
.collect::<Vec<_>>();
if !changes.is_empty() {
let _ = self.events.send(WorkerProjectionEvent {
revision: commit.revision,
changes,
});
}
let _ = self.events.send(WorkerProjectionEvent {
revision: commit.revision,
changes: commit.changes,
});
Ok(())
}
}
@@ -1198,7 +1198,7 @@ mod tests {
);
assert!(
store
!store
.delete_worker_registry(
&workspace_id,
&worker_runtime::identity::RuntimeWorkerRef {
@@ -1207,6 +1207,8 @@ mod tests {
},
)
.expect("remove worker registry")
.changes
.is_empty()
);
let state: String = store
.with_conn(|conn| {
@@ -413,9 +413,7 @@ fn project_registry_worker(
workspace_id: Some(record.registry.workspace_id.clone()),
display_name: Some(record.registry.display_name.clone()),
profile: record.registry.profile.clone(),
repository_id: None,
repository_key: None,
working_directory_id: None,
workdir_attachments: Vec::new(),
}
};
worker.runtime_id = Some(runtime_id);
@@ -440,20 +438,30 @@ fn project_working_directory(
else {
return false;
};
if let Some(link) = links.first() {
let Ok(Some(workdir)) = api
.store
.get_workdir_registry(&api.config.workspace_id, &link.workdir_id)
else {
return false;
};
let Ok(working_directory_id) =
protocol::subscription::SubscriptionWorkdirId::new(workdir.workdir_id)
else {
return false;
};
worker.repository_id = Some(workdir.repository_id);
worker.working_directory_id = Some(working_directory_id);
if !links.is_empty() {
let mut attachments = Vec::with_capacity(links.len());
for link in links {
let Ok(Some(workdir)) = api
.store
.get_workdir_registry(&api.config.workspace_id, &link.workdir_id)
else {
return false;
};
let Ok(working_directory_id) =
protocol::subscription::SubscriptionWorkdirId::new(workdir.workdir_id)
else {
return false;
};
attachments.push(
protocol::subscription::SubscriptionWorkerWorkdirAttachment {
alias: link.alias,
repository_id: Some(workdir.repository_id),
repository_key: None,
working_directory_id,
},
);
}
worker.workdir_attachments = attachments;
} else {
let Ok(has_link_history) = api
.store
@@ -462,26 +470,28 @@ fn project_working_directory(
return false;
};
if has_link_history {
worker.repository_id = None;
worker.repository_key = None;
worker.working_directory_id = None;
worker.workdir_attachments.clear();
}
}
project_repository_key(api, worker)
project_repository_keys(api, worker)
}
fn project_repository_key(api: &WorkspaceApi, worker: &mut SubscriptionWorker) -> bool {
let Some(repository_id) = worker.repository_id.take() else {
worker.repository_key = None;
return true;
};
let Ok(Some(repository)) = api
.store
.get_repository(&api.config.workspace_id, &repository_id)
else {
return false;
};
worker.repository_key = Some(repository.repository_key);
fn project_repository_keys(api: &WorkspaceApi, worker: &mut SubscriptionWorker) -> bool {
for attachment in &mut worker.workdir_attachments {
let Some(repository_id) = attachment.repository_id.take() else {
if attachment.repository_key.is_none() {
return false;
}
continue;
};
let Ok(Some(repository)) = api
.store
.get_repository(&api.config.workspace_id, &repository_id)
else {
return false;
};
attachment.repository_key = Some(repository.repository_key);
}
true
}
+8 -7
View File
@@ -236,6 +236,12 @@ export type SubscriptionWorkerAvailability = "observed" | "unavailable";
export type EventSubscriptionSelector = { "topic": "runtime_workers" } | { "topic": "worker_lifecycle", worker_ids: SubscriptionWorkerIds, } | { "topic": "worker_protocol", worker_id: SubscriptionWorkerId, runtime_id?: string | null, } | { "topic": "workspace_workers" } | { "topic": "workspace_workdirs" };
export type SubscriptionWorkerWorkdirAttachment = {
/**
* Stable Worker-local routing alias.
*/
alias: string, repository_key?: string | null, working_directory_id: SubscriptionWorkdirId, };
export type SubscriptionWorker = { worker_id: SubscriptionWorkerId,
/**
* Set by the Workspace Server when projecting a Runtime-owned Worker to clients.
@@ -264,14 +270,9 @@ worker_state?: WorkerStateSnapshot | null,
/**
* Runtime catalog lifecycle compatibility projection; not foreground-state authority.
*/
state: SubscriptionWorkerState, has_running_internal_workers: boolean, workspace_id?: string | null, display_name?: string | null, profile?: string | null,
/**
* Workspace-facing Repository key. Runtime producers leave this unset and
* Workspace Server projections replace `repository_id` with this field.
*/
repository_key?: string | null, working_directory_id?: SubscriptionWorkdirId | null, };
state: SubscriptionWorkerState, has_running_internal_workers: boolean, workspace_id?: string | null, display_name?: string | null, profile?: string | null, workdir_attachments?: Array<SubscriptionWorkerWorkdirAttachment>, };
export type WorkspaceSubscriptionWorkdir = { working_directory_id: SubscriptionWorkdirId, repository_key: string, state: string, primary_worker_id?: SubscriptionWorkerId | null, };
export type WorkspaceSubscriptionWorkdir = { working_directory_id: SubscriptionWorkdirId, repository_key: string, state: string, };
export type SubscriptionSnapshot = { "topic": "workers", "data": { workers: Array<SubscriptionWorker>, } } | { "topic": "worker_protocol", "data": { worker_id: SubscriptionWorkerId, events: Array<Event>, } } | { "topic": "workspace_workdirs", "data": { workdirs: Array<WorkspaceSubscriptionWorkdir>, } };
+10 -2
View File
@@ -13,9 +13,17 @@ export type WorkingDirectoryCleanupTarget = { kind: string, working_directory_id
export type WorkingDirectoryOccupancy = { runtime_id: string, worker_id: string, display_name: string, linked_at: string, };
export type WorkingDirectorySummary = { working_directory_id: string, repository_key: string, creation_selector?: string | null, creation_ref?: string | null, creation_tree?: string | null, current_selector?: string | null, current_ref?: string | null, current_tree?: string | null, observed_at_epoch_seconds?: number | null, materializer_kind: WorkingDirectoryMaterializerKind, cleanup_target?: WorkingDirectoryCleanupTarget | null, status: WorkingDirectoryStatusKind, cleanliness?: string | null, primary_worker_id?: string | null, occupied_by?: WorkingDirectoryOccupancy | null, };
export type WorkingDirectorySummary = { working_directory_id: string,
/**
* Optional human-facing label, never a routing key.
*/
display_name?: string | null, repository_key: string, creation_selector?: string | null, creation_ref?: string | null, creation_tree?: string | null, current_selector?: string | null, current_ref?: string | null, current_tree?: string | null, observed_at_epoch_seconds?: number | null, materializer_kind: WorkingDirectoryMaterializerKind, cleanup_target?: WorkingDirectoryCleanupTarget | null, status: WorkingDirectoryStatusKind, cleanliness?: string | null, occupied_by?: WorkingDirectoryOccupancy | null, };
export type WorkingDirectoryCreateRequest = { runtime_id?: string | null, repository_key: string, selector?: string | null, operation_id?: string | null, };
export type WorkingDirectoryCreateRequest = { runtime_id?: string | null,
/**
* Optional human-facing label, independent of every Worker attachment alias.
*/
display_name?: string | null, repository_key: string, selector?: string | null, operation_id?: string | null, };
export type WorkingDirectoryListResponse = { workspace_id: string, items: Array<WorkingDirectorySummary>, diagnostics: Array<Diagnostic>, };
@@ -15,11 +15,19 @@ export type WorkingDirectoryCleanupTarget = { kind: string, working_directory_id
export type RuntimeWorkingDirectoryCleanupTarget = { kind: string, working_directory_id: string, repository_id: string, };
export type RuntimeWorkingDirectorySummary = { working_directory_id: string, repository_id: string, creation_selector?: string | null, creation_ref?: string | null, creation_tree?: string | null, current_selector?: string | null, current_ref?: string | null, current_tree?: string | null, observed_at_epoch_seconds?: number | null, materializer_kind: WorkingDirectoryMaterializerKind, cleanup_target?: RuntimeWorkingDirectoryCleanupTarget | null, status: WorkingDirectoryStatusKind, cleanliness?: string | null, primary_worker_id?: string | null, occupied_by?: WorkingDirectoryOccupancy | null, };
export type RuntimeWorkingDirectorySummary = { working_directory_id: string,
/**
* Optional human-facing label, never a routing key.
*/
display_name?: string | null, repository_id: string, creation_selector?: string | null, creation_ref?: string | null, creation_tree?: string | null, current_selector?: string | null, current_ref?: string | null, current_tree?: string | null, observed_at_epoch_seconds?: number | null, materializer_kind: WorkingDirectoryMaterializerKind, cleanup_target?: RuntimeWorkingDirectoryCleanupTarget | null, status: WorkingDirectoryStatusKind, cleanliness?: string | null, occupied_by?: WorkingDirectoryOccupancy | null, };
export type WorkingDirectoryOccupancy = { runtime_id: string, worker_id: string, display_name: string, linked_at: string, };
export type WorkingDirectorySummary = { working_directory_id: string, repository_key: string, creation_selector?: string | null, creation_ref?: string | null, creation_tree?: string | null, current_selector?: string | null, current_ref?: string | null, current_tree?: string | null, observed_at_epoch_seconds?: number | null, materializer_kind: WorkingDirectoryMaterializerKind, cleanup_target?: WorkingDirectoryCleanupTarget | null, status: WorkingDirectoryStatusKind, cleanliness?: string | null, primary_worker_id?: string | null, occupied_by?: WorkingDirectoryOccupancy | null, };
export type WorkingDirectorySummary = { working_directory_id: string,
/**
* Optional human-facing label, never a routing key.
*/
display_name?: string | null, repository_key: string, creation_selector?: string | null, creation_ref?: string | null, creation_tree?: string | null, current_selector?: string | null, current_ref?: string | null, current_tree?: string | null, observed_at_epoch_seconds?: number | null, materializer_kind: WorkingDirectoryMaterializerKind, cleanup_target?: WorkingDirectoryCleanupTarget | null, status: WorkingDirectoryStatusKind, cleanliness?: string | null, occupied_by?: WorkingDirectoryOccupancy | null, };
export type WorkerWorkspaceSummary = { visibility: string, identity: string, workspace_id?: string | null, };
@@ -27,7 +35,9 @@ export type WorkerImplementationSummary = { kind: string, display_hint: string,
export type WorkerCapabilitySummary = { can_stop: boolean, can_spawn_followup: boolean, };
export type WorkerLaunchWorkerSummary = { runtime_id: string, worker_id: string, host_id: string, display_name: string, label: string, profile: string | null, singleton_key: string | null, tags: Array<string>, workspace: WorkerWorkspaceSummary, state: string, last_seen_at: string | null, pinned: boolean, retention_state: string, implementation: WorkerImplementationSummary, capabilities: WorkerCapabilitySummary, working_directory?: RuntimeWorkingDirectorySummary | null, diagnostics: Array<Diagnostic>, };
export type RuntimeWorkerWorkdirAttachmentSummary = { alias: string, working_directory: RuntimeWorkingDirectorySummary, };
export type WorkerLaunchWorkerSummary = { runtime_id: string, worker_id: string, host_id: string, display_name: string, label: string, profile: string | null, singleton_key: string | null, tags: Array<string>, workspace: WorkerWorkspaceSummary, state: string, last_seen_at: string | null, pinned: boolean, retention_state: string, implementation: WorkerImplementationSummary, capabilities: WorkerCapabilitySummary, workdir_attachments?: Array<RuntimeWorkerWorkdirAttachmentSummary>, diagnostics: Array<Diagnostic>, };
export type WorkerLaunchRuntimeOption = { runtime_id: string, display_name: string, built_in: boolean, worker_creation_available: boolean, working_directory_required: boolean, status: string, diagnostics: Array<Diagnostic>, };
@@ -37,11 +47,15 @@ export type WorkingDirectoryRepositoryOption = { repository_key: string, default
export type WorkerLaunchOptionsResponse = { workspace_id: string, runtimes: Array<WorkerLaunchRuntimeOption>, default_profile: string | null, profiles: Array<WorkerLaunchProfileCandidate>, repositories: Array<WorkingDirectoryRepositoryOption>, working_directories: Array<WorkingDirectorySummary>, diagnostics: Array<Diagnostic>, };
export type BrowserWorkerWorkingDirectorySelection = { working_directory_id: string, relative_cwd: string | null, };
export type BrowserWorkerWorkingDirectorySelection = {
/**
* Stable Worker-local routing alias.
*/
alias: string, working_directory_id: string, relative_cwd: string | null, };
export type CreateWorkspaceWorkerTicketAssignmentRequest = { ticket_id: string, operation_id: string, };
export type CreateWorkspaceWorkerRequest = { runtime_id: string, display_name: string, profile: string | null, ticket_assignment: CreateWorkspaceWorkerTicketAssignmentRequest | null, initial_submit: Array<Segment>, working_directory: BrowserWorkerWorkingDirectorySelection | null,
export type CreateWorkspaceWorkerRequest = { runtime_id: string, display_name: string, profile: string | null, ticket_assignment: CreateWorkspaceWorkerTicketAssignmentRequest | null, initial_submit: Array<Segment>, workdir_attachments: Array<BrowserWorkerWorkingDirectorySelection>,
/**
* Backend idempotency key used only for authenticated Worker-owned spawn/control.
*/
@@ -11,6 +11,7 @@ import type {
const SUMMARY_KEYS = new Set([
"working_directory_id",
"display_name",
"repository_key",
"creation_selector",
"creation_ref",
@@ -23,11 +24,11 @@ const SUMMARY_KEYS = new Set([
"cleanup_target",
"status",
"cleanliness",
"primary_worker_id",
"occupied_by",
]);
const CREATE_REQUEST_KEYS = new Set([
"runtime_id",
"display_name",
"repository_key",
"selector",
"operation_id",
@@ -84,6 +85,7 @@ export function validateWorkingDirectoryCreateRequest(
repository_key: stringField(record, "repository_key"),
};
assignOptionalString(request, record, "runtime_id");
assignOptionalString(request, record, "display_name");
assignOptionalString(request, record, "selector");
assignOptionalString(request, record, "operation_id");
return request;
@@ -124,6 +126,7 @@ export function parseWorkingDirectorySummary(
"unknown",
]),
};
assignOptionalString(summary, record, "display_name");
assignOptionalString(summary, record, "creation_selector");
assignOptionalString(summary, record, "creation_ref");
assignOptionalString(summary, record, "creation_tree");
@@ -131,7 +134,6 @@ export function parseWorkingDirectorySummary(
assignOptionalString(summary, record, "current_ref");
assignOptionalString(summary, record, "current_tree");
assignOptionalString(summary, record, "cleanliness");
assignOptionalString(summary, record, "primary_worker_id");
if (record.observed_at_epoch_seconds !== undefined) {
const observedAt = record.observed_at_epoch_seconds;
if (observedAt === null) {
@@ -139,10 +139,11 @@ Deno.test("Worker create request parser requires the complete shared request", (
{ kind: "text", content: "Implement T-565." },
{ kind: "flow", selector: "builtin:coder-review" },
],
working_directory: {
workdir_attachments: [{
alias: "source",
working_directory_id: "workdir-a",
relative_cwd: null,
},
}],
control_operation_id: null,
};
+38 -17
View File
@@ -8,6 +8,7 @@ import type {
DiagnosticSeverity,
RuntimeWorkingDirectoryCleanupTarget,
RuntimeWorkingDirectorySummary,
RuntimeWorkerWorkdirAttachmentSummary,
WorkerCapabilitySummary,
WorkerImplementationSummary,
WorkerLaunchOptionsResponse,
@@ -261,6 +262,7 @@ function runtimeWorkingDirectory(
item,
[
"working_directory_id",
"display_name",
"repository_id",
"creation_selector",
"creation_ref",
@@ -273,7 +275,6 @@ function runtimeWorkingDirectory(
"cleanup_target",
"status",
"cleanliness",
"primary_worker_id",
"occupied_by",
],
label,
@@ -318,6 +319,11 @@ function runtimeWorkingDirectory(
item.working_directory_id,
`${label}.working_directory_id`,
),
display_name: optional(
item.display_name,
`${label}.display_name`,
string,
),
repository_id: string(item.repository_id, `${label}.repository_id`),
creation_selector: optional(
item.creation_selector,
@@ -350,15 +356,25 @@ function runtimeWorkingDirectory(
),
status: status as RuntimeWorkingDirectorySummary["status"],
cleanliness: optional(item.cleanliness, `${label}.cleanliness`, string),
primary_worker_id: optional(
item.primary_worker_id,
`${label}.primary_worker_id`,
string,
),
occupied_by: occupied,
};
}
function runtimeWorkdirAttachment(
value: unknown,
label: string,
): RuntimeWorkerWorkdirAttachmentSummary {
const item = record(value, label);
exact(item, ["alias", "working_directory"], label);
return {
alias: string(item.alias, `${label}.alias`),
working_directory: runtimeWorkingDirectory(
item.working_directory,
`${label}.working_directory`,
),
};
}
function workerSummary(
value: unknown,
label: string,
@@ -382,7 +398,7 @@ function workerSummary(
"retention_state",
"implementation",
"capabilities",
"working_directory",
"workdir_attachments",
"diagnostics",
],
label,
@@ -406,11 +422,13 @@ function workerSummary(
`${label}.implementation`,
),
capabilities: capabilitySummary(item.capabilities, `${label}.capabilities`),
working_directory: optional(
item.working_directory,
`${label}.working_directory`,
runtimeWorkingDirectory,
),
workdir_attachments: item.workdir_attachments === undefined
? undefined
: array(
item.workdir_attachments,
`${label}.workdir_attachments`,
runtimeWorkdirAttachment,
),
diagnostics: array(item.diagnostics, `${label}.diagnostics`, diagnostic),
};
}
@@ -464,8 +482,9 @@ function workingDirectorySelection(
label: string,
): BrowserWorkerWorkingDirectorySelection {
const item = record(value, label);
exact(item, ["working_directory_id", "relative_cwd"], label);
exact(item, ["alias", "working_directory_id", "relative_cwd"], label);
return {
alias: string(item.alias, `${label}.alias`),
working_directory_id: string(
item.working_directory_id,
`${label}.working_directory_id`,
@@ -638,7 +657,7 @@ export function parseCreateWorkspaceWorkerRequest(
"profile",
"ticket_assignment",
"initial_submit",
"working_directory",
"workdir_attachments",
"control_operation_id",
],
"Worker create request",
@@ -651,9 +670,11 @@ export function parseCreateWorkspaceWorkerRequest(
? null
: ticketAssignment(item.ticket_assignment, "ticket_assignment"),
initial_submit: array(item.initial_submit, "initial_submit", segment),
working_directory: item.working_directory === null
? null
: workingDirectorySelection(item.working_directory, "working_directory"),
workdir_attachments: array(
item.workdir_attachments,
"workdir_attachments",
workingDirectorySelection,
),
control_operation_id: nullableString(
item.control_operation_id,
"control_operation_id",
@@ -256,8 +256,10 @@ Deno.test("workspace Worker list lives on the dedicated Workers page", async ()
workersNav.includes("worker.display_name || worker.label") &&
workersNav.includes("worker-status-dot") &&
workersNav.includes("worker-status-spinner") &&
workersNav.includes("worker.repository_key ?? '—'") &&
workersNav.includes("worker.working_directory_id ?? '—'") &&
workersNav.includes("workdirMeta(worker)") &&
workersNav.includes("attachment.alias") &&
workersNav.includes("attachment.repository_key ?? '—'") &&
workersNav.includes("attachment.working_directory_id") &&
!workersNav.includes('aria-disabled="true"'),
"Workers sidebar should link to the Worker list page and show state indicators with repository/workdir metadata",
);
@@ -37,6 +37,14 @@
Math.max(0, workers.length - COLLAPSED_WORKER_COUNT),
);
function workdirMeta(worker: SidebarWorker): string {
const attachments = worker.workdir_attachments;
if (attachments.length === 0) return '—';
return attachments
.map((attachment) => `${attachment.alias}: ${attachment.repository_key ?? '—'}・${attachment.working_directory_id}`)
.join(', ');
}
function workerKey(worker: SidebarWorker): string {
return `${worker.runtime_id}:${worker.worker_id}`;
}
@@ -190,7 +198,7 @@
</span>
<span class="worker-nav-label">{label}</span>
<small class="worker-nav-meta">
{worker.repository_key ?? '—'}・{worker.working_directory_id ?? '—'}
{workdirMeta(worker)}
</small>
</a>
<button
@@ -89,6 +89,11 @@ export type WorkerCapabilities = {
can_spawn_followup: boolean;
};
export type WorkerWorkdirAttachment = {
alias: string;
working_directory: WorkingDirectorySummary;
};
export type Worker = {
runtime_id: string;
worker_id: string;
@@ -107,7 +112,7 @@ export type Worker = {
last_seen_at?: string | null;
implementation: { kind: string; display_hint: string };
capabilities: WorkerCapabilities;
working_directory?: WorkingDirectorySummary | null;
workdir_attachments?: WorkerWorkdirAttachment[];
diagnostics: Diagnostic[];
};
@@ -7,13 +7,14 @@ import type {
} from "./types";
type FetchFn = typeof fetch;
type WorkerActionTarget = Pick<Worker, "runtime_id" | "worker_id" | "state">;
type WorkerLifecycleResponse = {
state: string;
diagnostics?: Diagnostic[];
};
function workerPath(workspaceId: string, worker: Worker): string {
function workerPath(workspaceId: string, worker: WorkerActionTarget): string {
return workspaceApiPath(
workspaceId,
`/runtimes/${encodeURIComponent(worker.runtime_id)}/workers/${
@@ -45,13 +46,13 @@ function diagnosticMessage(
fallback;
}
export function canDeleteSidebarWorker(worker: Worker): boolean {
export function canDeleteSidebarWorker(worker: WorkerActionTarget): boolean {
return worker.state === "stopped";
}
export async function stopSidebarWorker(
workspaceId: string,
worker: Worker,
worker: WorkerActionTarget,
fetchFn: FetchFn = fetch,
): Promise<void> {
const response = await fetchFn(`${workerPath(workspaceId, worker)}/stop`, {
@@ -71,7 +72,7 @@ export async function stopSidebarWorker(
export async function deleteSidebarWorker(
workspaceId: string,
worker: Worker,
worker: WorkerActionTarget,
fetchFn: FetchFn = fetch,
): Promise<void> {
const runtimePath = `/runtimes/${encodeURIComponent(worker.runtime_id)}`;
@@ -1,6 +1,7 @@
import {
buildCreateWorkspaceWorkerRequest,
defaultWorkerLaunchForm,
workerLaunchAttachmentError,
} from "./worker-launch.ts";
import type { WorkerLaunchOptionsResponse } from "./types.ts";
@@ -8,6 +9,262 @@ declare const Deno: {
test(name: string, fn: () => Promise<void> | void): void;
};
Deno.test("defaultWorkerLaunchForm uses the Backend-published defaults and initial Workdir", () => {
const form = defaultWorkerLaunchForm(
options,
emptyForm({ initial_text: "hello" }),
);
assertEquals(form.runtime_id, "remote");
assertEquals(form.display_name, "Worker");
assertEquals(form.profile, "builtin:coder");
assertEquals(form.initial_text, "hello");
assertEquals(form.workdir_attachments, [{
alias: "workdir",
working_directory_id: "wd-1-repo",
relative_cwd: "",
}]);
assertEquals(form.working_directory_repository_key, "repo");
assertEquals(form.working_directory_selector, "HEAD");
});
Deno.test("defaultWorkerLaunchForm preserves an available Ticket role profile", () => {
const reviewerOptions = {
...options,
profiles: [
...options.profiles,
{ id: "builtin:reviewer", label: "Reviewer", description: "review" },
],
};
const form = defaultWorkerLaunchForm(
reviewerOptions,
emptyForm({
display_name: "Review worker",
profile: "builtin:reviewer",
initial_text: "Review the ticket.",
working_directory_repository_key: "repo",
working_directory_selector: "HEAD",
}),
);
assertEquals(form.profile, "builtin:reviewer");
});
Deno.test("defaultWorkerLaunchForm skips occupied Workdirs and leaves an editable attachment", () => {
const form = defaultWorkerLaunchForm(
{
...options,
working_directories: [{
...options.working_directories[0],
occupied_by: {
runtime_id: "embedded",
worker_id: "0198f82e-6d90-7f15-a121-174a02e10e77",
display_name: "Worker 12",
linked_at: "2026-07-24T00:00:00Z",
},
}],
},
emptyForm({ initial_text: "hello" }),
);
assertEquals(form.workdir_attachments, [{
alias: "workdir",
working_directory_id: "",
relative_cwd: "",
}]);
});
Deno.test("defaultWorkerLaunchForm preserves a Ticket repository target", () => {
const form = defaultWorkerLaunchForm(
{
...options,
repositories: [
...options.repositories,
{ repository_key: "ticket-repo", default_selector: "main" },
],
working_directories: [
options.working_directories[0],
{
...options.working_directories[0],
working_directory_id: "ticket-workdir",
repository_key: "ticket-repo",
creation_selector: "work/ticket",
},
],
},
emptyForm({
display_name: "Ticket worker",
profile: "builtin:coder",
initial_text: "Work on a ticket.",
working_directory_repository_key: "ticket-repo",
working_directory_selector: "work/ticket",
}),
);
assertEquals(form.workdir_attachments, [{
alias: "workdir",
working_directory_id: "ticket-workdir",
relative_cwd: "",
}]);
assertEquals(form.working_directory_repository_key, "ticket-repo");
assertEquals(form.working_directory_selector, "work/ticket");
});
Deno.test("defaultWorkerLaunchForm clears attachments for a Workdir-less Runtime", () => {
const form = defaultWorkerLaunchForm(
options,
emptyForm({
runtime_id: "embedded",
workdir_attachments: [{
alias: "checkout",
working_directory_id: "wd-1-repo",
relative_cwd: "src",
}],
}),
);
assertEquals(form.workdir_attachments, []);
});
Deno.test("buildCreateWorkspaceWorkerRequest emits multiple aliased attachments", () => {
const request = buildCreateWorkspaceWorkerRequest(emptyForm({
runtime_id: "remote",
display_name: "Worker",
profile: "builtin:coder",
initial_text: "go",
workdir_attachments: [
{
alias: " checkout ",
working_directory_id: "wd-1-repo",
relative_cwd: "crates/yoi",
},
{
alias: "docs",
working_directory_id: "wd-2-docs",
relative_cwd: " ",
},
],
}));
assertEquals(request, {
runtime_id: "remote",
display_name: "Worker",
profile: "builtin:coder",
ticket_assignment: null,
initial_submit: [{ kind: "text", content: "go" }],
workdir_attachments: [
{
alias: "checkout",
working_directory_id: "wd-1-repo",
relative_cwd: "crates/yoi",
},
{
alias: "docs",
working_directory_id: "wd-2-docs",
relative_cwd: null,
},
],
control_operation_id: null,
});
});
Deno.test("buildCreateWorkspaceWorkerRequest validates attachment aliases and selections", () => {
assertEquals(
workerLaunchAttachmentError([attachment("invalid/alias", "wd-1-repo")]),
"Attachment 1 alias must be 1–64 ASCII letters, digits, dots, underscores, or hyphens, and start with a letter or digit.",
);
assertThrows(
() =>
buildCreateWorkspaceWorkerRequest(emptyForm({
workdir_attachments: [
attachment("checkout", "wd-1-repo"),
attachment("checkout", "wd-2-docs"),
],
})),
"Attachment alias “checkout” is used more than once.",
);
assertThrows(
() =>
buildCreateWorkspaceWorkerRequest(emptyForm({
workdir_attachments: [
attachment("checkout", "wd-1-repo"),
attachment("docs", "wd-1-repo"),
],
})),
"A Workdir cannot be attached more than once to one Worker.",
);
assertThrows(
() =>
buildCreateWorkspaceWorkerRequest(emptyForm({
workdir_attachments: [attachment("checkout", "")],
})),
"Attachment 1 must select a Workdir.",
);
assertThrows(
() =>
buildCreateWorkspaceWorkerRequest(emptyForm({
workdir_attachments: [{
...attachment("checkout", "wd-1-repo"),
relative_cwd: "../outside",
}],
})),
"Attachment 1 relative cwd must stay inside the selected Workdir.",
);
});
Deno.test("buildCreateWorkspaceWorkerRequest sends no initial segments for an empty draft", () => {
const request = buildCreateWorkspaceWorkerRequest(emptyForm({
runtime_id: "embedded",
profile: "builtin:companion",
initial_text: " ",
}));
assertEquals(request.initial_submit, []);
});
Deno.test("buildCreateWorkspaceWorkerRequest emits an empty attachment list for Workdir-less launches", () => {
const request = buildCreateWorkspaceWorkerRequest(emptyForm({
runtime_id: "embedded",
display_name: "Worker",
profile: "builtin:companion",
initial_text: "chat",
}));
assertEquals(request, {
runtime_id: "embedded",
display_name: "Worker",
profile: "builtin:companion",
ticket_assignment: null,
initial_submit: [{ kind: "text", content: "chat" }],
workdir_attachments: [],
control_operation_id: null,
});
});
function emptyForm(
overrides: Partial<Parameters<typeof buildCreateWorkspaceWorkerRequest>[0]> =
{},
): Parameters<typeof buildCreateWorkspaceWorkerRequest>[0] {
return {
runtime_id: "",
display_name: "",
profile: "",
initial_text: "",
workdir_attachments: [],
working_directory_repository_key: "",
working_directory_selector: "",
...overrides,
};
}
function attachment(alias: string, workingDirectoryId: string) {
return {
alias,
working_directory_id: workingDirectoryId,
relative_cwd: "",
};
}
function assertEquals<T>(actual: T, expected: T): void {
const actualJson = JSON.stringify(actual);
const expectedJson = JSON.stringify(expected);
@@ -16,6 +273,18 @@ function assertEquals<T>(actual: T, expected: T): void {
}
}
function assertThrows(fn: () => unknown, expectedMessage: string): void {
try {
fn();
} catch (error) {
if (error instanceof Error && error.message === expectedMessage) return;
throw new Error(
`Expected error ${JSON.stringify(expectedMessage)}, got ${String(error)}`,
);
}
throw new Error(`Expected error ${JSON.stringify(expectedMessage)}`);
}
const options: WorkerLaunchOptionsResponse = {
workspace_id: "workspace",
runtimes: [
@@ -57,7 +326,6 @@ const options: WorkerLaunchOptionsResponse = {
materializer_kind: "runtime_git_clone",
status: "active",
cleanliness: "clean",
primary_worker_id: null,
cleanup_target: {
kind: "runtime_git_clone",
working_directory_id: "wd-1-repo",
@@ -67,179 +335,3 @@ const options: WorkerLaunchOptionsResponse = {
],
diagnostics: [],
};
Deno.test("defaultWorkerLaunchForm uses the Backend-published Workspace default profile", () => {
const form = defaultWorkerLaunchForm(options, {
runtime_id: "",
display_name: "",
profile: "",
initial_text: "hello",
working_directory_id: "",
working_directory_repository_key: "",
working_directory_selector: "",
relative_cwd: "",
});
assertEquals(form.runtime_id, "remote");
assertEquals(form.display_name, "Worker");
assertEquals(form.profile, "builtin:coder");
assertEquals(form.initial_text, "hello");
assertEquals(form.working_directory_id, "wd-1-repo");
assertEquals(form.working_directory_repository_key, "repo");
assertEquals(form.working_directory_selector, "HEAD");
});
Deno.test("defaultWorkerLaunchForm preserves an available Ticket role profile", () => {
const reviewerOptions = {
...options,
profiles: [
...options.profiles,
{ id: "builtin:reviewer", label: "Reviewer", description: "review" },
],
};
const form = defaultWorkerLaunchForm(reviewerOptions, {
runtime_id: "",
display_name: "Review worker",
profile: "builtin:reviewer",
initial_text: "Review the ticket.",
working_directory_id: "",
working_directory_repository_key: "repo",
working_directory_selector: "HEAD",
relative_cwd: "",
});
assertEquals(form.profile, "builtin:reviewer");
});
Deno.test("defaultWorkerLaunchForm skips occupied working directories", () => {
const form = defaultWorkerLaunchForm(
{
...options,
working_directories: [
{
...options.working_directories[0],
occupied_by: {
runtime_id: "embedded",
worker_id: "0198f82e-6d90-7f15-a121-174a02e10e77",
display_name: "Worker 12",
linked_at: "2026-07-24T00:00:00Z",
},
},
],
},
{
runtime_id: "",
display_name: "",
profile: "",
initial_text: "hello",
working_directory_id: "",
working_directory_repository_key: "",
working_directory_selector: "",
relative_cwd: "",
},
);
assertEquals(form.working_directory_id, "");
});
Deno.test("defaultWorkerLaunchForm preserves a Ticket repository target", () => {
const form = defaultWorkerLaunchForm(
{
...options,
repositories: [
...options.repositories,
{
repository_key: "ticket-repo",
default_selector: "main",
},
],
working_directories: [
options.working_directories[0],
{
...options.working_directories[0],
working_directory_id: "ticket-workdir",
repository_key: "ticket-repo",
creation_selector: "work/ticket",
},
],
},
{
runtime_id: "",
display_name: "Ticket worker",
profile: "builtin:coder",
initial_text: "Work on a ticket.",
working_directory_id: "",
working_directory_repository_key: "ticket-repo",
working_directory_selector: "work/ticket",
relative_cwd: "",
},
);
assertEquals(form.working_directory_id, "ticket-workdir");
assertEquals(form.working_directory_repository_key, "ticket-repo");
assertEquals(form.working_directory_selector, "work/ticket");
});
Deno.test("buildCreateWorkspaceWorkerRequest sends working_directory id and relative cwd only", () => {
const request = buildCreateWorkspaceWorkerRequest({
runtime_id: "embedded",
display_name: "Worker",
profile: "builtin:coder",
initial_text: "go",
working_directory_id: "wd-1-repo",
working_directory_repository_key: "repo",
working_directory_selector: "main",
relative_cwd: "crates/yoi",
});
assertEquals(request, {
runtime_id: "embedded",
display_name: "Worker",
profile: "builtin:coder",
ticket_assignment: null,
initial_submit: [{ kind: "text", content: "go" }],
working_directory: {
working_directory_id: "wd-1-repo",
relative_cwd: "crates/yoi",
},
control_operation_id: null,
});
});
Deno.test("buildCreateWorkspaceWorkerRequest sends no initial segments for an empty draft", () => {
const request = buildCreateWorkspaceWorkerRequest({
runtime_id: "embedded",
display_name: "Worker",
profile: "builtin:companion",
initial_text: " ",
working_directory_id: "",
working_directory_repository_key: "",
working_directory_selector: "",
relative_cwd: "",
});
assertEquals(request.initial_submit, []);
});
Deno.test("buildCreateWorkspaceWorkerRequest emits null for embedded no-workdir launches", () => {
const request = buildCreateWorkspaceWorkerRequest({
runtime_id: "embedded",
display_name: "Worker",
profile: "builtin:companion",
initial_text: "chat",
working_directory_id: "",
working_directory_repository_key: "",
working_directory_selector: "",
relative_cwd: "",
});
assertEquals(request, {
runtime_id: "embedded",
display_name: "Worker",
profile: "builtin:companion",
ticket_assignment: null,
initial_submit: [{ kind: "text", content: "chat" }],
working_directory: null,
control_operation_id: null,
});
});
@@ -1,19 +1,29 @@
import type { CreateWorkspaceWorkerRequest } from "$lib/generated/worker-launch-api";
import type {
BrowserWorkerWorkingDirectorySelection,
CreateWorkspaceWorkerRequest,
} from "$lib/generated/worker-launch-api";
import { parseCreateWorkspaceWorkerRequest } from "$lib/workspace/api/workers";
import type { WorkerLaunchOptionsResponse } from "./types";
export type WorkerLaunchAttachmentFormState = {
alias: string;
working_directory_id: string;
relative_cwd: string;
};
export type WorkerLaunchFormState = {
runtime_id: string;
display_name: string;
profile: string;
initial_text: string;
working_directory_id: string;
workdir_attachments: WorkerLaunchAttachmentFormState[];
working_directory_repository_key: string;
working_directory_selector: string;
relative_cwd: string;
};
const WORKDIR_ALIAS_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$/;
export function defaultWorkerLaunchForm(
options: WorkerLaunchOptionsResponse | null,
current: WorkerLaunchFormState,
@@ -31,7 +41,6 @@ export function defaultWorkerLaunchForm(
options?.working_directories.filter((directory) =>
directory.status === "active" &&
directory.cleanliness === "clean" &&
directory.primary_worker_id == null &&
directory.occupied_by == null
) ?? [];
const selectedRuntime = current.runtime_id
@@ -60,6 +69,27 @@ export function defaultWorkerLaunchForm(
repository.repository_key === current.working_directory_repository_key
) ??
options?.repositories[0];
const availableWorkdirIds = new Set(
availableWorkingDirectories.map((directory) =>
directory.working_directory_id
),
);
const currentAttachments = workdirlessRuntime
? []
: current.workdir_attachments.filter((attachment) =>
availableWorkdirIds.has(attachment.working_directory_id)
);
const workdirAttachments = currentAttachments.length > 0
? currentAttachments
: preferredWorkingDirectory
? [{
alias: "workdir",
working_directory_id: preferredWorkingDirectory.working_directory_id,
relative_cwd: "",
}]
: selectedRuntime?.working_directory_required === true
? [{ alias: "workdir", working_directory_id: "", relative_cwd: "" }]
: [];
return {
runtime_id: current.runtime_id || preferredRuntime?.runtime_id || "",
@@ -69,22 +99,68 @@ export function defaultWorkerLaunchForm(
? current.profile
: preferredProfile?.id || "",
initial_text: current.initial_text,
working_directory_id:
!workdirlessRuntime && availableWorkingDirectories.some(
(directory) =>
directory.working_directory_id === current.working_directory_id,
)
? current.working_directory_id
: preferredWorkingDirectory?.working_directory_id || "",
workdir_attachments: workdirAttachments,
working_directory_repository_key:
current.working_directory_repository_key ||
preferredRepository?.repository_key || "",
working_directory_selector: current.working_directory_selector ||
preferredRepository?.default_selector || "HEAD",
relative_cwd: current.relative_cwd,
};
}
export function workerLaunchAttachmentError(
attachments: WorkerLaunchAttachmentFormState[],
): string | null {
const aliases = new Set<string>();
const workdirIds = new Set<string>();
for (const [index, attachment] of attachments.entries()) {
const position = index + 1;
const alias = attachment.alias.trim();
const workdirId = attachment.working_directory_id.trim();
const relativeCwd = attachment.relative_cwd.trim();
if (!WORKDIR_ALIAS_PATTERN.test(alias)) {
return `Attachment ${position} alias must be 1–64 ASCII letters, digits, dots, underscores, or hyphens, and start with a letter or digit.`;
}
if (aliases.has(alias)) {
return `Attachment alias “${alias}” is used more than once.`;
}
aliases.add(alias);
if (!workdirId) {
return `Attachment ${position} must select a Workdir.`;
}
if (workdirIds.has(workdirId)) {
return "A Workdir cannot be attached more than once to one Worker.";
}
workdirIds.add(workdirId);
if (
relativeCwd.startsWith("/") ||
relativeCwd.split("/").some((component) => component === "..")
) {
return `Attachment ${position} relative cwd must stay inside the selected Workdir.`;
}
}
return null;
}
function validatedAttachments(
attachments: WorkerLaunchAttachmentFormState[],
): BrowserWorkerWorkingDirectorySelection[] {
const error = workerLaunchAttachmentError(attachments);
if (error) {
throw new Error(error);
}
return attachments.map((attachment) => ({
alias: attachment.alias.trim(),
working_directory_id: attachment.working_directory_id.trim(),
relative_cwd: attachment.relative_cwd.trim() || null,
}));
}
export function buildCreateWorkspaceWorkerRequest(
form: WorkerLaunchFormState,
): CreateWorkspaceWorkerRequest {
@@ -97,12 +173,7 @@ export function buildCreateWorkspaceWorkerRequest(
initial_submit: initialMessage
? [{ kind: "text", content: form.initial_text }]
: [],
working_directory: form.working_directory_id
? {
working_directory_id: form.working_directory_id,
relative_cwd: form.relative_cwd.trim() || null,
}
: null,
workdir_attachments: validatedAttachments(form.workdir_attachments),
control_operation_id: null,
});
}
@@ -32,7 +32,7 @@ function worker(
workspace_id: 'workspace-test',
display_name: null,
profile: null,
working_directory_id: null,
workdir_attachments: [],
};
}
@@ -1,5 +1,8 @@
import { readable, type Readable } from 'svelte/store';
import type { SubscriptionWorker } from '$lib/generated/protocol';
import type {
SubscriptionWorker,
SubscriptionWorkerWorkdirAttachment,
} from '$lib/generated/protocol';
import { workspaceMultiplexer } from '$lib/workspace/multiplexer';
import {
applyWorkspaceWorkersFrame,
@@ -9,9 +12,8 @@ import { liveWorkerState } from './worker-state';
import { compareWorkersForSidebar } from './workers';
import type { Worker } from './types';
export type SidebarWorker = Worker & {
repository_key: string | null;
working_directory_id: string | null;
export type SidebarWorker = Omit<Worker, 'workdir_attachments'> & {
workdir_attachments: SubscriptionWorkerWorkdirAttachment[];
has_running_internal_workers: boolean;
};
@@ -103,10 +105,8 @@ function projectWorker(worker: SubscriptionWorker): SidebarWorker {
can_stop: worker.availability !== 'unavailable' && worker.state !== 'stopped',
can_spawn_followup: false,
},
repository_key: worker.repository_key ?? null,
working_directory_id: worker.working_directory_id ?? null,
workdir_attachments: worker.workdir_attachments ?? [],
has_running_internal_workers: worker.has_running_internal_workers,
working_directory: null,
diagnostics: [],
};
}
@@ -40,7 +40,7 @@ function worker(overrides: Partial<Worker>): Worker {
can_stop: true,
can_spawn_followup: false,
},
working_directory: null,
workdir_attachments: [],
diagnostics: [],
...overrides,
};
@@ -19,11 +19,15 @@ export function sidebarWorkerActivity(
return 'none';
}
export function canShowWorkerInSidebar(worker: Worker): boolean {
export function canShowWorkerInSidebar(
worker: Pick<Worker, 'implementation'>,
): boolean {
return worker.implementation.kind !== 'backend_worker_registry';
}
export function canOpenWorkerConsole(worker: Worker): boolean {
export function canOpenWorkerConsole(
worker: Pick<Worker, 'implementation'>,
): boolean {
return canShowWorkerInSidebar(worker);
}
@@ -143,19 +143,41 @@
min-width: 0;
max-width: 100%;
}
.worker-inline-select.wd-select {
min-width: 20rem;
.worker-attachment-heading {
display: flex;
align-items: center;
justify-content: space-between;
gap: var(--space-3);
}
.worker-attachment-heading h3,
.new-working-directory-panel h3 {
margin: 0;
color: var(--text-strong);
font-size: var(--font-size-body);
}
.worker-attachment-list {
display: grid;
}
.worker-attachment-row {
display: grid;
grid-template-columns:
minmax(7rem, 0.7fr) minmax(13rem, 1.5fr) minmax(10rem, 1fr) auto;
align-items: end;
gap: var(--space-2);
border-top: 1px solid var(--line);
padding: var(--space-3) 0;
}
.worker-attachment-row:last-child {
border-bottom: 1px solid var(--line);
}
.worker-attachment-remove {
margin-bottom: 1px;
}
.new-working-directory-panel {
display: grid;
gap: var(--space-3);
padding-left: var(--space-3);
}
.new-working-directory-panel h3 {
margin: 0;
color: var(--text-strong);
font-size: var(--font-size-body);
}
.new-working-directory-fields,
.worker-detail-grid {
display: grid;
@@ -182,15 +204,19 @@
flex-direction: column;
white-space: normal;
}
.worker-inline-select,
.worker-inline-select.wd-select {
.worker-inline-select {
width: 100%;
min-width: 0;
}
.new-working-directory-fields,
.worker-detail-grid {
.worker-detail-grid,
.worker-attachment-row {
grid-template-columns: 1fr;
}
.worker-attachment-remove {
justify-self: start;
margin-bottom: 0;
}
}
.worker-submit-error {
display: grid;
@@ -138,7 +138,10 @@
{#each workdirs as workdir}
{@const cleanup = cleanupCandidate(workdir)}
<tr>
<td><code>{workdir.working_directory_id}</code></td>
<td>
<span>{workdir.display_name ?? '—'}</span>
<small><code>{workdir.working_directory_id}</code></small>
</td>
<td>{workdir.repository_key}</td>
<td><code>{currentRevision(workdir)}</code></td>
<td>{workdir.status}</td>
@@ -145,11 +145,14 @@
}
function workerDirectory(worker: Worker): string {
const directory = worker.working_directory;
if (!directory) return '—';
const provider = data.repositories?.items.find((repository) => repository.repository_key === directory.repository_key)
?.provider;
return `${directory.repository_key} · ${formatCurrentWorkdirRevision(directory, provider)}`;
const attachments = worker.workdir_attachments ?? [];
if (attachments.length === 0) return '—';
return attachments.map(({ alias, working_directory: directory }) => {
const provider = data.repositories?.items.find((repository) => repository.repository_key === directory.repository_key)
?.provider;
const label = directory.display_name ?? directory.repository_key;
return `${alias}: ${label} · ${formatCurrentWorkdirRevision(directory, provider)}`;
}).join(', ');
}
</script>
@@ -11,7 +11,12 @@
parseWorkerLaunchOptionsResponse,
} from '$lib/workspace/api/workers';
import { formatCurrentWorkdirRevision } from '$lib/workspace/settings/workdir-revision';
import { buildCreateWorkspaceWorkerRequest, defaultWorkerLaunchForm } from '$lib/workspace/sidebar/worker-launch';
import {
buildCreateWorkspaceWorkerRequest,
defaultWorkerLaunchForm,
workerLaunchAttachmentError,
} from '$lib/workspace/sidebar/worker-launch';
import type { WorkerLaunchAttachmentFormState } from '$lib/workspace/sidebar/worker-launch';
import type {
Diagnostic,
WorkerLaunchOptionsResponse,
@@ -33,7 +38,8 @@
function workdirOptionLabel(directory: WorkingDirectorySummary): string {
const provider = data.repositories?.items.find((repository) => repository.repository_key === directory.repository_key)
?.provider;
return `${directory.repository_key} · ${formatCurrentWorkdirRevision(directory, provider)}`;
const label = directory.display_name ?? directory.repository_key;
return `${label} · ${formatCurrentWorkdirRevision(directory, provider)}`;
}
let { data }: PageProps = $props();
@@ -61,31 +67,37 @@
: '',
);
let initialText = $state(ticketContext?.initialInput ?? '');
let workingDirectoryId = $state('');
let workdirAttachments = $state<WorkerLaunchAttachmentFormState[]>([]);
let workingDirectoryDisplayName = $state('');
let workingDirectoryRepositoryKey = $state(ticketContext?.repositoryKey ?? '');
let workingDirectorySelector = $state(ticketContext?.refSelector ?? 'HEAD');
let relativeCwd = $state('');
let creatingWorkingDirectory = $state(false);
let isNewWorkingDirectorySelected = $derived(workingDirectoryId === NEW_WORKING_DIRECTORY_VALUE);
let newWorkingDirectoryAttachmentIndex = $derived(
workdirAttachments.findIndex((attachment) => attachment.working_directory_id === NEW_WORKING_DIRECTORY_VALUE),
);
let isNewWorkingDirectorySelected = $derived(newWorkingDirectoryAttachmentIndex >= 0);
let selectedRuntime = $derived(options?.runtimes.find((runtime) => runtime.runtime_id === runtimeId));
let selectedRuntimeAllowsNoWorkdir = $derived(selectedRuntime?.working_directory_required === false);
let hasSelectedExistingWorkdir = $derived(Boolean(
workingDirectoryId && !isNewWorkingDirectorySelected && !selectedRuntimeAllowsNoWorkdir,
));
let availableWorkingDirectories = $derived(
selectedRuntimeAllowsNoWorkdir
? []
: (options?.working_directories ?? []).filter((directory) =>
directory.status === 'active' &&
directory.cleanliness === 'clean' &&
directory.primary_worker_id == null &&
directory.occupied_by == null
),
);
let attachmentValidationError = $derived(
isNewWorkingDirectorySelected
? 'Create the new Workdir before starting the Worker.'
: workerLaunchAttachmentError(workdirAttachments),
);
let canStartWorker = $derived(Boolean(
runtimeId &&
profile &&
(hasSelectedExistingWorkdir || (selectedRuntimeAllowsNoWorkdir && !isNewWorkingDirectorySelected)),
(selectedRuntimeAllowsNoWorkdir
? workdirAttachments.length === 0
: workdirAttachments.length > 0 && !attachmentValidationError),
));
function workerApiPath(path: string): string {
@@ -104,12 +116,49 @@
});
$effect(() => {
if (selectedRuntimeAllowsNoWorkdir && workingDirectoryId) {
workingDirectoryId = '';
relativeCwd = '';
if (selectedRuntimeAllowsNoWorkdir && workdirAttachments.length > 0) {
workdirAttachments = [];
} else if (selectedRuntime?.working_directory_required === true && workdirAttachments.length === 0) {
workdirAttachments = [{
alias: 'workdir',
working_directory_id: availableWorkingDirectories[0]?.working_directory_id ?? '',
relative_cwd: '',
}];
}
});
function nextAttachmentAlias(): string {
const aliases = new Set(workdirAttachments.map((attachment) => attachment.alias.trim()));
if (!aliases.has('workdir')) return 'workdir';
let suffix = 2;
while (aliases.has(`workdir${suffix}`)) suffix += 1;
return `workdir${suffix}`;
}
function addAttachment(): void {
const selectedWorkdirs = new Set(
workdirAttachments.map((attachment) => attachment.working_directory_id),
);
const available = availableWorkingDirectories.find((directory) =>
!selectedWorkdirs.has(directory.working_directory_id)
);
workdirAttachments = [...workdirAttachments, {
alias: nextAttachmentAlias(),
working_directory_id: available?.working_directory_id ?? '',
relative_cwd: '',
}];
}
function removeAttachment(index: number): void {
workdirAttachments = workdirAttachments.filter((_, candidateIndex) => candidateIndex !== index);
}
function workdirSelectedByAnotherAttachment(workdirId: string, index: number): boolean {
return workdirAttachments.some((attachment, candidateIndex) =>
candidateIndex !== index && attachment.working_directory_id === workdirId
);
}
async function loadLaunchOptions(signal?: AbortSignal) {
loading = true;
optionsError = null;
@@ -125,19 +174,30 @@
display_name: displayName,
profile,
initial_text: initialText,
working_directory_id: workingDirectoryId,
workdir_attachments: workdirAttachments,
working_directory_repository_key: workingDirectoryRepositoryKey,
working_directory_selector: workingDirectorySelector,
relative_cwd: relativeCwd,
});
runtimeId = form.runtime_id;
displayName = form.display_name;
profile = form.profile;
workingDirectoryId = form.working_directory_id ||
(ticketContext?.repositoryKey ? NEW_WORKING_DIRECTORY_VALUE : '');
const runtimeRequiresWorkdir = payload.runtimes.find((runtime) =>
runtime.runtime_id === form.runtime_id
)?.working_directory_required !== false;
workdirAttachments = form.workdir_attachments;
if (
runtimeRequiresWorkdir &&
ticketContext?.repositoryKey &&
workdirAttachments.every((attachment) => !attachment.working_directory_id)
) {
workdirAttachments = [{
alias: 'workdir',
working_directory_id: NEW_WORKING_DIRECTORY_VALUE,
relative_cwd: '',
}];
}
workingDirectoryRepositoryKey = form.working_directory_repository_key;
workingDirectorySelector = form.working_directory_selector;
relativeCwd = form.relative_cwd;
} catch (err) {
if (err instanceof DOMException && err.name === 'AbortError') {
return;
@@ -151,6 +211,11 @@
}
async function createWorkingDirectory() {
const attachmentIndex = newWorkingDirectoryAttachmentIndex;
if (attachmentIndex < 0) {
submitError = { message: 'select New workdir on an attachment first', diagnostics: [] };
return;
}
if (!runtimeId) {
submitError = { message: 'select a runtime before creating a workdir', diagnostics: [] };
return;
@@ -168,6 +233,7 @@
try {
const request = validateWorkingDirectoryCreateRequest({
runtime_id: runtimeId,
...(workingDirectoryDisplayName.trim() ? { display_name: workingDirectoryDisplayName.trim() } : {}),
repository_key: workingDirectoryRepositoryKey,
...(workingDirectorySelector ? { selector: workingDirectorySelector } : {}),
});
@@ -193,7 +259,11 @@
],
}
: options;
workingDirectoryId = payload.item.working_directory_id;
workdirAttachments = workdirAttachments.map((attachment, index) =>
index === attachmentIndex && attachment.working_directory_id === NEW_WORKING_DIRECTORY_VALUE
? { ...attachment, working_directory_id: payload.item.working_directory_id }
: attachment
);
} catch (err) {
submitError = exceptionDisplayError(err, 'workdir create failed');
} finally {
@@ -206,8 +276,16 @@
submitError = { message: 'workspace id is unavailable', diagnostics: [] };
return;
}
if (isNewWorkingDirectorySelected || (!workingDirectoryId && !selectedRuntimeAllowsNoWorkdir)) {
submitError = { message: 'select or create a workdir before starting a Worker; only embedded Runtime can start without one', diagnostics: [] };
if (selectedRuntimeAllowsNoWorkdir && workdirAttachments.length > 0) {
submitError = { message: 'the selected Runtime must start without Workdir attachments', diagnostics: [] };
return;
}
if (!selectedRuntimeAllowsNoWorkdir && workdirAttachments.length === 0) {
submitError = { message: 'add at least one Workdir attachment before starting a Worker; only embedded Runtime can start without one', diagnostics: [] };
return;
}
if (attachmentValidationError) {
submitError = { message: attachmentValidationError, diagnostics: [] };
return;
}
@@ -222,10 +300,9 @@
display_name: displayName,
profile,
initial_text: initialText,
working_directory_id: workingDirectoryId,
workdir_attachments: workdirAttachments,
working_directory_repository_key: workingDirectoryRepositoryKey,
working_directory_selector: workingDirectorySelector,
relative_cwd: relativeCwd,
})),
});
if (!response.ok) {
@@ -302,21 +379,7 @@
<h2 id="worker-location-heading">Location</h2>
<div class="worker-launch-sentence">
<span>Run at</span>
<select class="worker-inline-select wd-select" bind:value={workingDirectoryId} aria-label="Workdir">
{#if selectedRuntimeAllowsNoWorkdir}
<option value="">No workdir</option>
{:else}
<option value="" disabled>Select workdir</option>
{#each availableWorkingDirectories as directory}
<option value={directory.working_directory_id}>
{workdirOptionLabel(directory)}
</option>
{/each}
<option value={NEW_WORKING_DIRECTORY_VALUE}>New workdir…</option>
{/if}
</select>
<span>in</span>
<span>Run in</span>
<select class="worker-inline-select runtime-select" bind:value={runtimeId} required aria-label="Runtime">
{#if options?.runtimes.length}
{#each options.runtimes as runtime}
@@ -330,16 +393,78 @@
</select>
</div>
{#if !selectedRuntimeAllowsNoWorkdir && !workingDirectoryId && !isNewWorkingDirectorySelected}
<p class="worker-workdir-note">This Runtime requires a selected workdir before starting a Worker.</p>
{:else if selectedRuntimeAllowsNoWorkdir && !workingDirectoryId}
<p class="worker-workdir-note">No filesystem tools or Bash will be available without a workdir.</p>
{#if selectedRuntimeAllowsNoWorkdir}
<p class="worker-workdir-note">No filesystem tools or Bash will be available without a Workdir.</p>
{:else}
<div class="worker-attachment-heading">
<h3>Initial Workdir attachments</h3>
<button type="button" class="secondary-button" onclick={addAttachment}>Add attachment</button>
</div>
<div class="worker-attachment-list">
{#each workdirAttachments as attachment, index}
<div class="worker-attachment-row">
<label>
<span>Alias</span>
<input
bind:value={attachment.alias}
autocomplete="off"
maxlength="64"
placeholder="workdir"
aria-label={`Attachment ${index + 1} alias`}
/>
</label>
<label>
<span>Workdir</span>
<select bind:value={attachment.working_directory_id} aria-label={`Attachment ${index + 1} Workdir`}>
<option value="" disabled>Select Workdir</option>
{#each availableWorkingDirectories as directory}
<option
value={directory.working_directory_id}
disabled={workdirSelectedByAnotherAttachment(directory.working_directory_id, index)}
>
{workdirOptionLabel(directory)}
</option>
{/each}
<option
value={NEW_WORKING_DIRECTORY_VALUE}
disabled={isNewWorkingDirectorySelected && newWorkingDirectoryAttachmentIndex !== index}
>New Workdir…</option>
</select>
</label>
<label>
<span>Relative cwd</span>
<input
bind:value={attachment.relative_cwd}
autocomplete="off"
placeholder="Optional path inside Workdir"
aria-label={`Attachment ${index + 1} relative cwd`}
/>
</label>
<button
type="button"
class="worker-attachment-remove"
onclick={() => removeAttachment(index)}
disabled={creatingWorkingDirectory && newWorkingDirectoryAttachmentIndex === index}
aria-label={`Remove attachment ${index + 1}`}
>Remove</button>
</div>
{/each}
</div>
{#if attachmentValidationError}
<p class="worker-workdir-note error">{attachmentValidationError}</p>
{/if}
{/if}
{#if isNewWorkingDirectorySelected}
<div class="new-working-directory-panel">
<h3>New workdir</h3>
<h3>New Workdir for “{workdirAttachments[newWorkingDirectoryAttachmentIndex]?.alias || `attachment ${newWorkingDirectoryAttachmentIndex + 1}`}”</h3>
<div class="new-working-directory-fields">
<label>
<span>Display name</span>
<input bind:value={workingDirectoryDisplayName} autocomplete="off" placeholder="Optional label" maxlength="80" />
</label>
<label>
<span>Repository</span>
<select bind:value={workingDirectoryRepositoryKey}>
@@ -358,17 +483,10 @@
</label>
</div>
<button type="button" disabled={creatingWorkingDirectory || !runtimeId || !workingDirectoryRepositoryKey} onclick={() => void createWorkingDirectory()}>
{creatingWorkingDirectory ? 'Creating…' : 'Create workdir'}
{creatingWorkingDirectory ? 'Creating…' : 'Create Workdir'}
</button>
</div>
{/if}
{#if hasSelectedExistingWorkdir || isNewWorkingDirectorySelected}
<label class="relative-cwd-field">
<span>Relative cwd inside workdir</span>
<input bind:value={relativeCwd} autocomplete="off" placeholder="Optional path inside workdir" />
</label>
{/if}
</section>
<section class="worker-form-section" aria-labelledby="worker-details-heading">
+8
View File
@@ -10,6 +10,7 @@ import {
const summary = {
working_directory_id: "workdir-1",
display_name: "Checkout",
repository_key: "main",
materializer_kind: "runtime_git_clone",
status: "active",
@@ -27,6 +28,9 @@ Deno.test("Workdir REST validation accepts the generated list and create contrac
items: [summary],
diagnostics: [],
});
if (list.items[0]?.display_name !== "Checkout") {
throw new Error("display name was not preserved");
}
if (list.items[0]?.occupied_by?.runtime_id !== "arcadia") {
throw new Error("occupancy subject was not preserved");
}
@@ -60,6 +64,7 @@ Deno.test("Workdir REST validation rejects stale response JSON", () => {
Deno.test("Workdir REST validation enforces create operation fields", () => {
const request = validateWorkingDirectoryCreateRequest({
runtime_id: "arcadia",
display_name: "Review checkout",
repository_key: "main",
selector: "develop",
operation_id: "operation-1",
@@ -67,6 +72,9 @@ Deno.test("Workdir REST validation enforces create operation fields", () => {
if (request.operation_id !== "operation-1") {
throw new Error("operation id was not preserved");
}
if (request.display_name !== "Review checkout") {
throw new Error("display name was not preserved");
}
for (
const invalid of [