diff --git a/crates/ticket/src/lib.rs b/crates/ticket/src/lib.rs index 999be609..88f1811c 100644 --- a/crates/ticket/src/lib.rs +++ b/crates/ticket/src/lib.rs @@ -9,6 +9,7 @@ use std::fmt; use std::fs::{self, File, OpenOptions}; use std::io::{self, Write}; use std::path::{Component, Path, PathBuf}; +use std::sync::Arc; use chrono::Utc; use fs4::fs_std::FileExt; @@ -2274,11 +2275,40 @@ impl LocalTicketBackend { } } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SqliteTicketMutationEvent { + pub workspace_id: String, + pub ticket_id: String, + pub event_index: i64, + pub event_kind: TicketEventKind, +} + +pub type SqliteTicketMutationHook = + dyn Fn(&Connection, &SqliteTicketMutationEvent) -> Result<()> + Send + Sync; + +#[derive(Clone)] pub struct SqliteTicketBackend { db_path: PathBuf, workspace_id: String, record_language: Option, + event_attributes: BTreeMap, + mutation_hook: Option>, +} + +impl fmt::Debug for SqliteTicketBackend { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("SqliteTicketBackend") + .field("db_path", &self.db_path) + .field("workspace_id", &self.workspace_id) + .field("record_language", &self.record_language) + .field("event_attributes", &self.event_attributes) + .field( + "mutation_hook", + &self.mutation_hook.as_ref().map(|_| "configured"), + ) + .finish() + } } impl SqliteTicketBackend { @@ -2287,9 +2317,21 @@ impl SqliteTicketBackend { db_path: db_path.into(), workspace_id: workspace_id.into(), record_language: None, + event_attributes: BTreeMap::new(), + mutation_hook: None, } } + pub fn with_event_attributes(mut self, attributes: BTreeMap) -> Self { + self.event_attributes = attributes; + self + } + + pub fn with_mutation_hook(mut self, hook: Arc) -> Self { + self.mutation_hook = Some(hook); + self + } + pub fn with_record_language(mut self, language: Option<&str>) -> Self { self.record_language = language.and_then(normalized_record_language); self @@ -2506,10 +2548,25 @@ CREATE TABLE IF NOT EXISTS typed_ticket_artifacts ( conn.execute("INSERT INTO typed_ticket_event_references (workspace_id, ticket_id, event_index, ordinal, kind, target) VALUES (?1, ?2, ?3, ?4, ?5, ?6)", params![self.workspace_id, ticket_id, next_index, ordinal as i64, reference.kind, reference.target]).map_err(sqlite_err)?; } - for (key, value) in &event.attributes { + let mut attributes = event.attributes.clone(); + for (key, value) in &self.event_attributes { + attributes.insert(key.clone(), value.clone()); + } + for (key, value) in &attributes { conn.execute("INSERT INTO typed_ticket_event_attributes (workspace_id, ticket_id, event_index, key, value) VALUES (?1, ?2, ?3, ?4, ?5)", params![self.workspace_id, ticket_id, next_index, key, value]).map_err(sqlite_err)?; } + if let Some(hook) = &self.mutation_hook { + hook( + conn, + &SqliteTicketMutationEvent { + workspace_id: self.workspace_id.clone(), + ticket_id: ticket_id.to_string(), + event_index: next_index, + event_kind: event.kind.clone(), + }, + )?; + } Ok(()) } @@ -2718,6 +2775,9 @@ CREATE TABLE IF NOT EXISTS typed_ticket_artifacts ( for row in rows { let (index, kind, author, at, status, from, to, reason, state_field, heading, body) = row.map_err(sqlite_err)?; + let mut attributes = self.load_event_attributes(conn, ticket_id, index)?; + attributes.insert("event_id".to_string(), format!("{ticket_id}:{index}")); + attributes.insert("event_sequence".to_string(), index.to_string()); events.push(TicketEvent { kind: TicketEventKind::from(kind.as_str()), author, @@ -2730,7 +2790,7 @@ CREATE TABLE IF NOT EXISTS typed_ticket_artifacts ( heading, body: MarkdownText::new(body), references: self.load_event_references(conn, ticket_id, index)?, - attributes: self.load_event_attributes(conn, ticket_id, index)?, + attributes, }); } Ok(events) @@ -6393,6 +6453,41 @@ state: planning assert_partial_body_replacement_semantics(&backend); } + #[test] + fn sqlite_mutation_hook_failure_rolls_back_ticket_event() { + let tmp = TempDir::new().unwrap(); + let db_path = tmp.path().join("workspace.db"); + let backend = SqliteTicketBackend::new(&db_path, "workspace-test"); + let created = backend.create(NewTicket::new("Atomic mutation")).unwrap(); + let before = backend + .show(TicketIdOrSlug::Id(created.id.clone())) + .unwrap() + .events + .len(); + let failing = backend.clone().with_mutation_hook(Arc::new(|_, event| { + Err(TicketError::Conflict(format!( + "reject outbox for {}:{}", + event.ticket_id, event.event_index + ))) + })); + assert!( + failing + .add_event( + TicketIdOrSlug::Id(created.id.clone()), + NewTicketEvent::new(TicketEventKind::Comment, "must roll back"), + ) + .is_err() + ); + let after = backend.show(TicketIdOrSlug::Id(created.id)).unwrap(); + assert_eq!(after.events.len(), before); + assert!( + after + .events + .iter() + .all(|event| event.body.as_str() != "must roll back") + ); + } + #[test] fn sqlite_backend_persists_core_ticket_operations() { let tmp = TempDir::new().unwrap(); diff --git a/crates/ticket/src/tool.rs b/crates/ticket/src/tool.rs index fb88153e..0095f9b0 100644 --- a/crates/ticket/src/tool.rs +++ b/crates/ticket/src/tool.rs @@ -34,12 +34,15 @@ const MAX_BODY_MAX_BYTES: usize = 64 * 1024; const DEFAULT_DIAGNOSTIC_LIMIT: usize = 100; const MAX_DIAGNOSTIC_LIMIT: usize = 500; -pub const TICKET_BASE_TOOL_NAMES: [&str; 12] = [ +pub const TICKET_BASE_TOOL_NAMES: [&str; 15] = [ "TicketCreate", "TicketEditItem", "TicketList", "TicketShow", "TicketComment", + "TicketPlan", + "TicketDecision", + "TicketImplementationReport", "TicketReview", "TicketIntakeReady", "TicketQueue", @@ -66,12 +69,15 @@ pub const TICKET_ORCHESTRATION_TOOL_NAMES: [&str; 4] = [ pub const TICKET_ORCHESTRATION_READ_ONLY_TOOL_NAMES: [&str; 2] = ["TicketRelationQuery", "TicketOrchestrationPlanQuery"]; -pub const TICKET_TOOL_NAMES: [&str; 16] = [ +pub const TICKET_TOOL_NAMES: [&str; 19] = [ "TicketCreate", "TicketEditItem", "TicketList", "TicketShow", "TicketComment", + "TicketPlan", + "TicketDecision", + "TicketImplementationReport", "TicketReview", "TicketIntakeReady", "TicketQueue", @@ -94,10 +100,13 @@ pub const TICKET_READ_ONLY_TOOL_NAMES: [&str; 6] = [ "TicketOrchestrationPlanQuery", ]; -pub const TICKET_MUTATING_TOOL_NAMES: [&str; 10] = [ +pub const TICKET_MUTATING_TOOL_NAMES: [&str; 13] = [ "TicketCreate", "TicketEditItem", "TicketComment", + "TicketPlan", + "TicketDecision", + "TicketImplementationReport", "TicketReview", "TicketIntakeReady", "TicketQueue", @@ -120,9 +129,11 @@ routing, closing, planning, or implementation decisions."; const SHOW_DESCRIPTION: &str = "Show one Ticket by id or exact query through the configured \ typed Ticket backend. Output includes bounded Markdown body, recent thread events, resolution, and \ artifact metadata."; -const COMMENT_DESCRIPTION: &str = "Append a typed Ticket thread event. `role` must be `comment`, \ -`plan`, `decision`, or `implementation_report`; `body` is Markdown. Writes stay inside the \ -configured Ticket backend root."; +const COMMENT_DESCRIPTION: &str = "Append a typed Ticket comment event. `body` is Markdown."; +const PLAN_DESCRIPTION: &str = "Append a typed Ticket plan event. `body` is Markdown."; +const DECISION_DESCRIPTION: &str = "Append a typed Ticket decision event. `body` is Markdown."; +const IMPLEMENTATION_REPORT_DESCRIPTION: &str = + "Append a typed Ticket implementation_report event. `body` is Markdown."; const REVIEW_DESCRIPTION: &str = "Append a Ticket review event. `result` must be `approve` or \ `request_changes`; `body` is Markdown. Writes stay inside the configured Ticket backend root."; const INTAKE_READY_DESCRIPTION: &str = "Mark an existing Ticket planning lane ready through the typed \ @@ -161,6 +172,9 @@ fn base_tool_description(name: &str) -> &'static str { "TicketList" => LIST_DESCRIPTION, "TicketShow" => SHOW_DESCRIPTION, "TicketComment" => COMMENT_DESCRIPTION, + "TicketPlan" => PLAN_DESCRIPTION, + "TicketDecision" => DECISION_DESCRIPTION, + "TicketImplementationReport" => IMPLEMENTATION_REPORT_DESCRIPTION, "TicketReview" => REVIEW_DESCRIPTION, "TicketIntakeReady" => INTAKE_READY_DESCRIPTION, "TicketQueue" => QUEUE_DESCRIPTION, @@ -361,9 +375,6 @@ struct TicketCreateParams { /// Markdown body for item.md. If omitted, a small default body is used. #[serde(default)] body: Option, - /// Optional thread author for the create event. - #[serde(default)] - author: Option, /// Optional assignee frontmatter value. #[serde(default)] assignee: Option, @@ -376,9 +387,6 @@ struct TicketCreateParams { /// Optional state frontmatter value. Defaults to `planning`. #[serde(default)] state: Option, - /// Optional queued_by frontmatter value. - #[serde(default)] - queued_by: Option, /// Optional queued_at frontmatter value. #[serde(default)] queued_at: Option, @@ -412,9 +420,6 @@ struct TicketEditItemParams { /// Optional target repository/ref update. #[serde(default)] target: Option, - /// Optional thread author for the audited item_edit event. - #[serde(default)] - author: Option, } #[derive(Debug, Clone, Copy, Deserialize, schemars::JsonSchema)] @@ -542,25 +547,11 @@ struct TicketShowParams { } #[derive(Debug, Deserialize, schemars::JsonSchema)] -#[serde(rename_all = "snake_case")] -enum TicketCommentRoleParam { - Comment, - Plan, - Decision, - ImplementationReport, -} - -#[derive(Debug, Deserialize, schemars::JsonSchema)] -struct TicketCommentParams { +struct TicketThreadEventParams { /// Ticket id. ticket: String, - /// Thread event role: `comment`, `plan`, `decision`, or `implementation_report`. - role: TicketCommentRoleParam, /// Markdown event body. body: String, - /// Optional thread author. - #[serde(default)] - author: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] @@ -578,9 +569,6 @@ struct TicketReviewParams { result: TicketReviewResultParam, /// Markdown review body. body: String, - /// Optional thread author. - #[serde(default)] - author: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] @@ -589,9 +577,6 @@ struct TicketIntakeReadyParams { ticket: String, /// Concise bounded intake summary to append as a typed intake_summary event. intake_summary: String, - /// Optional author for both intake_summary and state_changed events. - #[serde(default)] - author: Option, /// Reason attached to the state_changed event. Defaults to `planning_ready`. #[serde(default)] reason: Option, @@ -604,9 +589,6 @@ struct TicketIntakeReadyParams { struct TicketQueueParams { /// Ticket id. ticket: String, - /// Optional queued_by frontmatter value. Defaults to the backend/user default. - #[serde(default)] - queued_by: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] @@ -621,9 +603,6 @@ struct TicketWorkflowStateParams { reason: String, /// Markdown body for the typed state_changed event. body: String, - /// Optional thread author. - #[serde(default)] - author: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] @@ -673,9 +652,6 @@ struct TicketRelationRecordParams { /// Optional bounded rationale/note. #[serde(default)] note: Option, - /// Optional record author. - #[serde(default)] - author: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] @@ -757,9 +733,6 @@ struct TicketOrchestrationPlanRecordParams { /// Accepted plan fields. Required for accepted_plan and invalid for other kinds. #[serde(default)] accepted_plan: Option, - /// Optional record author. - #[serde(default)] - author: Option, } #[derive(Debug, Deserialize, schemars::JsonSchema)] @@ -851,6 +824,21 @@ struct TicketCommentTool { backend: TicketToolBackend, } +#[derive(Clone)] +struct TicketPlanTool { + backend: TicketToolBackend, +} + +#[derive(Clone)] +struct TicketDecisionTool { + backend: TicketToolBackend, +} + +#[derive(Clone)] +struct TicketImplementationReportTool { + backend: TicketToolBackend, +} + #[derive(Clone)] struct TicketReviewTool { backend: TicketToolBackend, @@ -918,12 +906,12 @@ impl Tool for TicketCreateTool { if let Some(body) = params.body { input.body = MarkdownText::new(body); } - input.author = params.author; + input.author = None; input.assignee = params.assignee; input.readiness = params.readiness; input.risk_flags = params.risk_flags; input.workflow_state = params.state.map(TicketWorkflowStateParam::into_state); - input.queued_by = params.queued_by; + input.queued_by = None; input.queued_at = params.queued_at; input.repository_id = params.repository_id; input.ref_selector = params.ref_selector; @@ -971,7 +959,7 @@ impl Tool for TicketEditItemTool { body: params.body.map(MarkdownText::new), body_replacement, target: params.target, - author: params.author, + author: None, }; let ticket = self .backend @@ -1066,6 +1054,26 @@ impl Tool for TicketShowTool { } } +fn execute_ticket_thread_event( + backend: &TicketToolBackend, + tool_name: &str, + kind: TicketEventKind, + input_json: &str, +) -> Result { + let params: TicketThreadEventParams = parse_input(tool_name, input_json)?; + let role = kind.as_str().to_string(); + backend + .add_event( + TicketIdOrSlug::Query(params.ticket.clone()), + NewTicketEvent::new(kind, params.body), + ) + .map_err(|error| backend_error(tool_name, error))?; + Ok(json_output( + format!("Appended {role} event to ticket {}", params.ticket), + json!({ "ticket": params.ticket, "event": role, "ok": true }), + )) +} + #[async_trait] impl Tool for TicketCommentTool { async fn execute( @@ -1073,26 +1081,42 @@ impl Tool for TicketCommentTool { input_json: &str, _ctx: llm_engine::tool::ToolExecutionContext, ) -> Result { - let params: TicketCommentParams = parse_input("TicketComment", input_json)?; - let kind = match params.role { - TicketCommentRoleParam::Comment => TicketEventKind::Comment, - TicketCommentRoleParam::Plan => TicketEventKind::Plan, - TicketCommentRoleParam::Decision => TicketEventKind::Decision, - TicketCommentRoleParam::ImplementationReport => TicketEventKind::ImplementationReport, - }; - let role = kind.as_str().to_string(); - let mut event = NewTicketEvent::new(kind, params.body); - event.author = params.author; - self.backend - .add_event(TicketIdOrSlug::Query(params.ticket.clone()), event) - .map_err(|error| backend_error("TicketComment", error))?; - Ok(json_output( - format!("Appended {role} event to ticket {}", params.ticket), - json!({ "ticket": params.ticket, "event": role, "ok": true }), - )) + execute_ticket_thread_event( + &self.backend, + "TicketComment", + TicketEventKind::Comment, + input_json, + ) } } +macro_rules! impl_ticket_thread_event_tool { + ($tool:ty, $name:literal, $kind:expr) => { + #[async_trait] + impl Tool for $tool { + async fn execute( + &self, + input_json: &str, + _ctx: llm_engine::tool::ToolExecutionContext, + ) -> Result { + execute_ticket_thread_event(&self.backend, $name, $kind, input_json) + } + } + }; +} + +impl_ticket_thread_event_tool!(TicketPlanTool, "TicketPlan", TicketEventKind::Plan); +impl_ticket_thread_event_tool!( + TicketDecisionTool, + "TicketDecision", + TicketEventKind::Decision +); +impl_ticket_thread_event_tool!( + TicketImplementationReportTool, + "TicketImplementationReport", + TicketEventKind::ImplementationReport +); + #[async_trait] impl Tool for TicketReviewTool { async fn execute( @@ -1108,7 +1132,7 @@ impl Tool for TicketReviewTool { let result_str = result.as_str().to_string(); let review = TicketReview { result, - author: params.author, + author: None, body: MarkdownText::new(params.body), }; self.backend @@ -1138,14 +1162,14 @@ impl Tool for TicketIntakeReadyTool { .default_intake_ready_state_change_body(from.as_str()) }); let mut summary = TicketIntakeSummary::new(params.intake_summary); - summary.author = params.author.clone(); + summary.author = None; let mut change = TicketStateChange::new( from.as_str(), TicketWorkflowState::Ready.as_str(), reason, body, ); - change.author = params.author; + change.author = None; self.backend .mark_intake_ready( TicketIdOrSlug::Query(params.ticket.clone()), @@ -1168,7 +1192,7 @@ impl Tool for TicketQueueTool { _ctx: llm_engine::tool::ToolExecutionContext, ) -> Result { let params: TicketQueueParams = parse_input("TicketQueue", input_json)?; - let queued_by = params.queued_by.unwrap_or_else(default_author); + let queued_by = default_author(); self.backend .queue_ready(TicketIdOrSlug::Query(params.ticket.clone()), &queued_by) .map_err(|error| backend_error("TicketQueue", error))?; @@ -1196,7 +1220,7 @@ impl Tool for TicketWorkflowStateTool { } let mut change = TicketStateChange::new(from.as_str(), to.as_str(), params.reason, params.body); - change.author = params.author; + change.author = None; self.backend .set_workflow_state(TicketIdOrSlug::Query(params.ticket.clone()), change) .map_err(|error| backend_error("TicketWorkflowState", error))?; @@ -1251,7 +1275,7 @@ impl Tool for TicketRelationRecordTool { kind: params.kind.into_kind(), target: params.target.clone(), note: params.note, - author: params.author, + author: None, }; let output = self .backend @@ -1325,7 +1349,7 @@ impl Tool for TicketOrchestrationPlanRecordTool { related_ticket: params.related_ticket, note: params.note, accepted_plan, - author: params.author, + author: None, }; let output = self .backend @@ -1703,7 +1727,9 @@ fn input_schema(name: &str) -> Value { "TicketEditItem" => serde_json::to_value(schemars::schema_for!(TicketEditItemParams)), "TicketList" => serde_json::to_value(schemars::schema_for!(TicketListParams)), "TicketShow" => serde_json::to_value(schemars::schema_for!(TicketShowParams)), - "TicketComment" => serde_json::to_value(schemars::schema_for!(TicketCommentParams)), + "TicketComment" | "TicketPlan" | "TicketDecision" | "TicketImplementationReport" => { + serde_json::to_value(schemars::schema_for!(TicketThreadEventParams)) + } "TicketReview" => serde_json::to_value(schemars::schema_for!(TicketReviewParams)), "TicketIntakeReady" => serde_json::to_value(schemars::schema_for!(TicketIntakeReadyParams)), "TicketQueue" => serde_json::to_value(schemars::schema_for!(TicketQueueParams)), @@ -1747,6 +1773,9 @@ impl_from_backend!(TicketEditItemTool); impl_from_backend!(TicketListTool); impl_from_backend!(TicketShowTool); impl_from_backend!(TicketCommentTool); +impl_from_backend!(TicketPlanTool); +impl_from_backend!(TicketDecisionTool); +impl_from_backend!(TicketImplementationReportTool); impl_from_backend!(TicketReviewTool); impl_from_backend!(TicketIntakeReadyTool); impl_from_backend!(TicketQueueTool); @@ -1768,6 +1797,12 @@ pub fn ticket_tools(backend: impl Into) -> Vec("TicketList", backend.clone()), tool_definition::("TicketShow", backend.clone()), tool_definition::("TicketComment", backend.clone()), + tool_definition::("TicketPlan", backend.clone()), + tool_definition::("TicketDecision", backend.clone()), + tool_definition::( + "TicketImplementationReport", + backend.clone(), + ), tool_definition::("TicketReview", backend.clone()), tool_definition::("TicketIntakeReady", backend.clone()), tool_definition::("TicketQueue", backend.clone()), @@ -1841,6 +1876,9 @@ mod tests { "TicketCreate", "TicketEditItem", "TicketComment", + "TicketPlan", + "TicketDecision", + "TicketImplementationReport", "TicketReview", "TicketIntakeReady", "TicketQueue", @@ -2338,16 +2376,15 @@ mod tests { let temp = TempDir::new().unwrap(); let backend = backend(&temp); let created = backend.create(NewTicket::new("Flow Tool")).unwrap(); - let comment = tool_by_name(backend.clone(), "TicketComment"); + let report = tool_by_name(backend.clone(), "TicketImplementationReport"); let review = tool_by_name(backend.clone(), "TicketReview"); let close = tool_by_name(backend.clone(), "TicketClose"); let doctor = tool_by_name(backend.clone(), "TicketDoctor"); - comment + report .execute( &json!({ "ticket": created.id.clone(), - "role": "implementation_report", "body": "Implemented." }) .to_string(), @@ -2807,6 +2844,38 @@ mod tests { assert!(edit_schema.contains("old_string")); assert!(edit_schema.contains("new_string")); assert!(edit_schema.contains("replace_all")); + for name in [ + "TicketCreate", + "TicketEditItem", + "TicketComment", + "TicketPlan", + "TicketDecision", + "TicketImplementationReport", + "TicketReview", + "TicketIntakeReady", + "TicketQueue", + "TicketRelationRecord", + "TicketOrchestrationPlanRecord", + ] { + let schema = tools + .iter() + .map(|definition| definition().0) + .find(|meta| meta.name == name) + .unwrap() + .input_schema; + let properties = schema["properties"].as_object().unwrap(); + assert!(!properties.contains_key("author"), "{name} exposes author"); + assert!( + !properties.contains_key("queued_by"), + "{name} exposes queued_by" + ); + if matches!( + name, + "TicketComment" | "TicketPlan" | "TicketDecision" | "TicketImplementationReport" + ) { + assert!(!properties.contains_key("role"), "{name} exposes role"); + } + } let names = tools .into_iter() .map(|definition| definition().0) diff --git a/crates/worker-runtime/src/catalog.rs b/crates/worker-runtime/src/catalog.rs index b412e0c8..4e7fb3ca 100644 --- a/crates/worker-runtime/src/catalog.rs +++ b/crates/worker-runtime/src/catalog.rs @@ -172,10 +172,29 @@ pub struct WorkingDirectoryStatus { pub summary: WorkingDirectorySummary, } -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct WorkspaceApiRef { pub workspace_id: String, pub base_url: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub runtime_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub access_token: Option, +} + +impl std::fmt::Debug for WorkspaceApiRef { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WorkspaceApiRef") + .field("workspace_id", &self.workspace_id) + .field("base_url", &self.base_url) + .field("runtime_id", &self.runtime_id) + .field( + "access_token", + &self.access_token.as_ref().map(|_| "[redacted]"), + ) + .finish() + } } /// Canonical Runtime Worker creation request. @@ -189,6 +208,10 @@ pub struct WorkspaceApiRef { /// summarized without exposing raw host paths. #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct CreateWorkerRequest { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub idempotency_key: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub idempotency_fingerprint: Option, pub profile: ProfileSelector, #[serde(default, skip_serializing_if = "Option::is_none")] pub display_name: Option, diff --git a/crates/worker-runtime/src/http_server.rs b/crates/worker-runtime/src/http_server.rs index 1565bf90..4f0f4c95 100644 --- a/crates/worker-runtime/src/http_server.rs +++ b/crates/worker-runtime/src/http_server.rs @@ -1153,6 +1153,8 @@ mod tests { request.workspace_api = Some(WorkspaceApiRef { workspace_id: workspace_id.to_string(), base_url: format!("https://workspace.example/{workspace_id}"), + runtime_id: None, + access_token: None, }); request } @@ -1410,6 +1412,8 @@ mod tests { let profile = ProfileSelector::Builtin("builtin:coder".to_string()); let bundle = test_bundle(profile.clone()); CreateWorkerRequest { + idempotency_key: None, + idempotency_fingerprint: None, profile, display_name: None, profile_source: crate::catalog::ProfileSourceArchiveSource::Http { @@ -1810,6 +1814,8 @@ mod ws_tests { fn ws_create_request() -> CreateWorkerRequest { let bundle = ws_test_bundle(ProfileSelector::Builtin("builtin:companion".to_string())); CreateWorkerRequest { + idempotency_key: None, + idempotency_fingerprint: None, profile: ProfileSelector::Builtin("builtin:companion".to_string()), display_name: None, profile_source: crate::catalog::ProfileSourceArchiveSource::Http { diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index 8bf620cb..678de3f5 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -354,6 +354,11 @@ impl Runtime { request: CreateWorkerRequest, scope: Option<&RuntimeWorkspaceScope>, ) -> Result { + if request.idempotency_key.is_some() != request.idempotency_fingerprint.is_some() { + return Err(RuntimeError::InvalidRequest( + "idempotency_key and idempotency_fingerprint must be provided together".to_string(), + )); + } let (backend, worker_ref, spawn_request) = { let mut state = self.lock()?; state.ensure_running()?; @@ -365,6 +370,20 @@ impl Runtime { if let Some(scope) = scope { state.ensure_workspace_owner(scope, true)?; }; + if let Some(idempotency_key) = request.idempotency_key.as_deref() { + let workspace_id = scope.map(|scope| scope.workspace_id.as_str()); + if let Some(existing) = state.workers.values().find(|record| { + record.workspace_id.as_deref() == workspace_id + && record.request.idempotency_key.as_deref() == Some(idempotency_key) + }) { + if existing.request.idempotency_fingerprint != request.idempotency_fingerprint { + return Err(RuntimeError::InvalidRequest(format!( + "worker creation idempotency key {idempotency_key} was already used with different input" + ))); + } + 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) = @@ -2108,7 +2127,9 @@ fn input_protocol_event(input: &WorkerInput) -> protocol::Event { #[cfg(test)] mod tests { use super::*; - use crate::catalog::{ConfigBundleRef, ProfileSelector, WorkspaceApiRef}; + use crate::catalog::{ + ConfigBundleRef, ProfileSelector, WorkingDirectoryClaim, WorkspaceApiRef, + }; use crate::config_bundle::{ ConfigBundle, ConfigBundleMetadata, ConfigBundleProvenance, ConfigDeclaration, ConfigDeclarationKind, ConfigProfileDescriptor, @@ -2126,6 +2147,8 @@ mod tests { let profile = ProfileSelector::Builtin("builtin:coder".to_string()); let bundle = test_bundle_for_profile(profile.clone()); CreateWorkerRequest { + idempotency_key: None, + idempotency_fingerprint: None, profile, display_name: None, profile_source: crate::catalog::ProfileSourceArchiveSource::Http { @@ -2161,6 +2184,8 @@ mod tests { request.workspace_api = Some(WorkspaceApiRef { workspace_id: workspace_id.to_string(), base_url: format!("https://workspace.example/{workspace_id}"), + runtime_id: None, + access_token: None, }); request } @@ -2612,6 +2637,33 @@ mod tests { )); } + #[test] + fn create_worker_idempotency_reuses_worker_and_rejects_different_input() { + let runtime = runtime_with_backend(); + let mut request = task_request("idempotent"); + request.idempotency_key = Some("operation-1".to_string()); + request.idempotency_fingerprint = Some("sha256:input-1".to_string()); + request.working_directory = Some(WorkingDirectoryClaim { + 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(); + let replayed = runtime.create_worker(request.clone()).unwrap(); + assert_eq!(replayed.worker_ref, first.worker_ref); + assert_eq!(runtime.list_workers().unwrap().len(), 1); + assert_eq!( + runtime.list_working_directories().unwrap().len(), + workdir_count_after_first + ); + + request.idempotency_fingerprint = Some("sha256:different".to_string()); + let error = runtime.create_worker(request).unwrap_err(); + assert!(matches!(error, RuntimeError::InvalidRequest(_))); + assert_eq!(runtime.list_workers().unwrap().len(), 1); + } + #[test] fn create_worker_rejects_system_initial_input_without_persisting_worker() { let runtime = runtime_with_backend(); diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 113e40b8..79f5fe27 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -23,6 +23,7 @@ use crate::execution::{ WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionRunState, WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult, }; +use crate::identity::WorkerRef; use crate::interaction::{WorkerInput, WorkerInputKind}; use crate::resource::{BackendResourceClient, ProfileSourceArchiveCache}; use crate::working_directory::{ @@ -40,8 +41,8 @@ use tokio::sync::broadcast; #[cfg(feature = "ws-server")] use worker::ipc::protocol_session::{live_log_entry_event, subscribe_worker_protocol_session}; use worker::{ - PromptLoader, Worker, WorkerController, WorkerError, WorkerFilesystemAuthority, WorkerHandle, - WorkerWorkspaceContext, WorkspaceClient, WorkspaceId, + PromptLoader, RuntimeWorkspaceHttpClient, Worker, WorkerController, WorkerError, + WorkerFilesystemAuthority, WorkerHandle, WorkerWorkspaceContext, WorkspaceId, }; const DEFAULT_BACKEND_ID: &str = "worker-crate"; @@ -259,6 +260,7 @@ enum RuntimeWorkspaceBackendRef { Http { workspace_id: String, base_url: String, + access_token: Option, }, } @@ -268,20 +270,29 @@ impl RuntimeWorkspaceBackendRef { return Self::Http { workspace_id: api.workspace_id.clone(), base_url: api.base_url.clone(), + access_token: api.access_token.clone(), }; } Self::None } - fn worker_context(&self) -> WorkerWorkspaceContext { + fn worker_context(&self, worker_ref: &WorkerRef) -> WorkerWorkspaceContext { match self { Self::None => WorkerWorkspaceContext::no_workspace(), Self::Http { workspace_id, base_url, + access_token, } => WorkerWorkspaceContext::with_client( WorkspaceId::new(workspace_id.clone()).ok(), - WorkspaceClient::http(workspace_id.clone(), base_url.clone()), + Arc::new( + RuntimeWorkspaceHttpClient::new( + workspace_id.clone(), + base_url.clone(), + worker_ref.worker_id.to_string(), + ) + .with_access_token(access_token.clone()), + ), ), } } @@ -368,7 +379,7 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { .unwrap_or(WorkerFilesystemAuthority::None); let workspace_backend_ref = RuntimeWorkspaceBackendRef::from_worker_request(&request.request); - let workspace_context = workspace_backend_ref.worker_context(); + let workspace_context = workspace_backend_ref.worker_context(&request.worker_ref); let selector = profile.as_ref(); let archive = self .resolve_profile_source_archive(&request.request.profile_source) @@ -442,7 +453,7 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { .unwrap_or(WorkerFilesystemAuthority::None); let workspace_backend_ref = RuntimeWorkspaceBackendRef::from_worker_request(&request.request); - let workspace_context = workspace_backend_ref.worker_context(); + let workspace_context = workspace_backend_ref.worker_context(&request.worker_ref); let (manifest, loader) = Self::restore_fallback_manifest(&worker_name)?; let store_dir = self.store_dir()?; @@ -1276,7 +1287,7 @@ mod tests { store_dir: PathBuf, worker_metadata_dir: PathBuf, observed_cwds: Arc>>, - observed_workspace_clients: Arc>>, + observed_workspace_clients: Arc, bool)>>>, } #[async_trait] @@ -1325,11 +1336,13 @@ mod tests { .unwrap_or_else(|| self.cwd.clone()); let workspace_backend_ref = RuntimeWorkspaceBackendRef::from_worker_request(&request.request); - let workspace_context = workspace_backend_ref.worker_context(); - self.observed_workspace_clients - .lock() - .unwrap() - .push(workspace_context.client().clone()); + let workspace_context = workspace_backend_ref.worker_context(&request.worker_ref); + let workspace_client = workspace_context.client(); + self.observed_workspace_clients.lock().unwrap().push(( + workspace_client.kind().to_string(), + workspace_client.workspace_id().map(str::to_string), + workspace_client.is_available(), + )); let scope = Scope::writable(&scope_root).map_err(|err| err.to_string())?; let worker = Worker::new( manifest, @@ -1438,6 +1451,8 @@ mod tests { fn create_request(_name: &str) -> CreateWorkerRequest { let bundle = test_bundle(); CreateWorkerRequest { + idempotency_key: None, + idempotency_fingerprint: None, profile: ProfileSelector::Builtin("builtin:companion".to_string()), display_name: None, profile_source: crate::catalog::ProfileSourceArchiveSource::Embedded { @@ -1673,6 +1688,8 @@ mod tests { request.workspace_api = Some(crate::catalog::WorkspaceApiRef { workspace_id: "ws-test".to_string(), base_url: "http://127.0.0.1:3999".to_string(), + runtime_id: None, + access_token: None, }); let detail = runtime.create_worker(request).unwrap(); @@ -1704,7 +1721,11 @@ mod tests { assert!(observed_cwds.lock().unwrap().is_empty()); assert_eq!( observed_workspace_clients.lock().unwrap().as_slice(), - &[WorkspaceClient::http("ws-test", "http://127.0.0.1:3999")] + &[( + "runtime-http-proxy".to_string(), + Some("ws-test".to_string()), + true, + )] ); let names = captured_tool_names(&client, 0); for forbidden in core_filesystem_tool_names() { @@ -1781,9 +1802,7 @@ mod tests { assert!(cwd.join("README.md").exists()); assert_eq!( observed_workspace_clients.lock().unwrap().as_slice(), - &[WorkspaceClient::Unavailable { - reason: "no workspace configured".to_string() - }] + &[("unavailable".to_string(), None, false)] ); } diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index b8c35fcd..48a43235 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -26,7 +26,7 @@ use crate::shutdown_after_idle::{ use crate::spawn::comm_tools::{read_worker_output_tool, send_to_worker_tool, stop_worker_tool}; use crate::spawn::registry::SpawnedWorkerRegistry; use crate::spawn::tool::spawn_worker_tool; -use crate::worker::{SystemItemCommitter, Worker, WorkerError, WorkerRunResult, WorkspaceClient}; +use crate::worker::{SystemItemCommitter, Worker, WorkerError, WorkerRunResult}; use protocol::{ AlertLevel, AlertSource, ErrorCode, Event, Method, RewindTargetId, RunResult, Segment, TurnResult, WorkerStatus, @@ -627,21 +627,16 @@ where // Ticket tools are typed operations over the current workspace Ticket backend. // Workspace access must be authority-bound to the Backend Workspace API; the // Worker must not fall back to a local `.yoi/tickets` store. - let ticket_backend = match worker.workspace_client() { - WorkspaceClient::Http { - workspace_id, - base_url, - } => crate::feature::builtin::ticket::TicketFeatureBackend::WorkspaceHttp { - workspace_id: workspace_id.clone(), - base_url: base_url.clone(), - }, - _ => { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "ticket tools require Backend Workspace API authority", - )); - } - }; + let workspace_client = worker.workspace_client_handle(); + if !workspace_client.is_available() || workspace_client.workspace_id().is_none() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "ticket tools require Backend Workspace API authority", + )); + } + let ticket_backend = crate::feature::builtin::ticket::TicketFeatureBackend::WorkspaceClient( + workspace_client, + ); feature_registry.add_module( crate::feature::builtin::ticket::ticket_tools_feature_with_backend( ticket_backend, @@ -668,21 +663,16 @@ where } { - let workspace_client = worker.workspace_client().clone(); + let workspace_client = worker.workspace_client_handle(); let engine = worker.engine_mut(); // Objective tools expose read-only project Objective context through the // Backend Workspace API. Workers must not guess local `.yoi/objectives` // paths or read Objective files directly. if feature_config.objective.enabled { - if let WorkspaceClient::Http { - workspace_id, - base_url, - } = &workspace_client - { + if workspace_client.is_available() && workspace_client.workspace_id().is_some() { for definition in crate::feature::builtin::objective::workspace_http_objective_tools( - workspace_id.clone(), - base_url.clone(), + workspace_client.clone(), ) { engine.register_tool(definition); } @@ -705,20 +695,14 @@ where "[feature.memory].enabled = true requires a [memory] configuration section", ) })?; - if let WorkspaceClient::Http { - workspace_id, - base_url, - } = workspace_client - { + if workspace_client.is_available() && workspace_client.workspace_id().is_some() { let definitions = if feature_config.memory.staging { crate::feature::builtin::memory::workspace_http_memory_consolidation_tools( - workspace_id, - base_url, + workspace_client.clone(), ) } else { crate::feature::builtin::memory::workspace_http_memory_tools( - workspace_id, - base_url, + workspace_client.clone(), ) }; for definition in definitions { diff --git a/crates/worker/src/feature/builtin/memory.rs b/crates/worker/src/feature/builtin/memory.rs index 4e034fa9..97fc334a 100644 --- a/crates/worker/src/feature/builtin/memory.rs +++ b/crates/worker/src/feature/builtin/memory.rs @@ -20,27 +20,25 @@ use schemars::JsonSchema; use serde::de::DeserializeOwned; use serde_json::json; -use crate::worker::WorkspaceClient; +use crate::worker::{ + WorkspaceClient, WorkspaceClientError, WorkspaceRequest, WorkspaceRequestMethod, +}; #[derive(Clone, Debug)] pub struct WorkspaceHttpMemoryBackend { - workspace_id: String, - base_url: String, + client: Arc, } impl WorkspaceHttpMemoryBackend { - pub fn new(workspace_id: impl Into, base_url: impl Into) -> Self { - Self { - workspace_id: workspace_id.into(), - base_url: base_url.into(), - } + pub fn new(client: Arc) -> Self { + Self { client } } pub async fn execute_operation( &self, operation: MemoryBackendOperation, ) -> Result { - execute_http_memory_backend(&self.workspace_id, &self.base_url, operation).await + execute_memory_backend(self.client.as_ref(), operation).await } async fn execute(&self, operation: MemoryBackendOperation) -> Result { @@ -59,7 +57,7 @@ pub enum WorkspaceMemoryBackendError { #[error("workspace memory backend is unavailable: {reason}")] Unavailable { reason: String }, #[error("workspace memory backend request failed: {0}")] - Request(#[from] reqwest::Error), + Request(#[from] WorkspaceClientError), #[error("workspace memory backend returned HTTP {status}: {body}")] Http { status: reqwest::StatusCode, @@ -71,73 +69,49 @@ pub enum WorkspaceMemoryBackendError { Backend(String), } -impl WorkspaceClient { +impl dyn WorkspaceClient + '_ { pub async fn execute_memory_backend_operation( &self, operation: MemoryBackendOperation, ) -> Result { - match self { - WorkspaceClient::Http { - workspace_id, - base_url, - } => execute_http_memory_backend(workspace_id, base_url, operation).await, - WorkspaceClient::Available { kind } => Err(WorkspaceMemoryBackendError::Unavailable { - reason: format!( - "workspace client kind `{kind}` does not expose the Backend Workspace API" - ), - }), - WorkspaceClient::Unavailable { reason } => { - Err(WorkspaceMemoryBackendError::Unavailable { - reason: reason.clone(), - }) - } - } + execute_memory_backend(self, operation).await } pub async fn request_memory_staging_consolidation( &self, operation: MemoryConsolidateStagingOperation, ) -> Result { - match self { - WorkspaceClient::Http { - workspace_id, - base_url, - } => execute_http_memory_consolidation(workspace_id, base_url, operation).await, - WorkspaceClient::Available { kind } => Err(WorkspaceMemoryBackendError::Unavailable { - reason: format!( - "workspace client kind `{kind}` does not expose the Backend Workspace API" - ), - }), - WorkspaceClient::Unavailable { reason } => { - Err(WorkspaceMemoryBackendError::Unavailable { - reason: reason.clone(), - }) - } - } + execute_memory_consolidation(self, operation).await } } -async fn execute_http_memory_backend( - workspace_id: &str, - base_url: &str, +async fn execute_memory_backend( + client: &dyn WorkspaceClient, operation: MemoryBackendOperation, ) -> Result { - let url = format!( - "{}/api/w/{}/memory/backend", - base_url.trim_end_matches('/'), - workspace_id - ); - let response = reqwest::Client::new() - .post(url) - .json(&operation) - .send() - .await?; - let status = response.status(); - let body = response.text().await?; - if !status.is_success() { - return Err(WorkspaceMemoryBackendError::Http { status, body }); + let workspace_id = + client + .workspace_id() + .ok_or_else(|| WorkspaceMemoryBackendError::Unavailable { + reason: format!( + "workspace client kind `{}` has no workspace id", + client.kind() + ), + })?; + let response = client.execute(WorkspaceRequest::json( + WorkspaceRequestMethod::Post, + format!("/api/w/{workspace_id}/memory/backend"), + serde_json::to_string(&operation)?, + ))?; + let status = reqwest::StatusCode::from_u16(response.status) + .unwrap_or(reqwest::StatusCode::INTERNAL_SERVER_ERROR); + if !response.is_success() { + return Err(WorkspaceMemoryBackendError::Http { + status, + body: response.body, + }); } - match serde_json::from_str::(&body)? { + match serde_json::from_str::(&response.body)? { MemoryBackendHttpResponse::Ok { result } => Ok(result), MemoryBackendHttpResponse::Error { message } => { Err(WorkspaceMemoryBackendError::Backend(message)) @@ -145,34 +119,37 @@ async fn execute_http_memory_backend( } } -async fn execute_http_memory_consolidation( - workspace_id: &str, - base_url: &str, +async fn execute_memory_consolidation( + client: &dyn WorkspaceClient, operation: MemoryConsolidateStagingOperation, ) -> Result { - let url = format!( - "{}/api/w/{}/memory/consolidation", - base_url.trim_end_matches('/'), - workspace_id - ); - let response = reqwest::Client::new() - .post(url) - .json(&operation) - .send() - .await?; - let status = response.status(); - let body = response.text().await?; - if !status.is_success() { - return Err(WorkspaceMemoryBackendError::Http { status, body }); + let workspace_id = + client + .workspace_id() + .ok_or_else(|| WorkspaceMemoryBackendError::Unavailable { + reason: format!( + "workspace client kind `{}` has no workspace id", + client.kind() + ), + })?; + let response = client.execute(WorkspaceRequest::json( + WorkspaceRequestMethod::Post, + format!("/api/w/{workspace_id}/memory/consolidation"), + serde_json::to_string(&operation)?, + ))?; + let status = reqwest::StatusCode::from_u16(response.status) + .unwrap_or(reqwest::StatusCode::INTERNAL_SERVER_ERROR); + if !response.is_success() { + return Err(WorkspaceMemoryBackendError::Http { + status, + body: response.body, + }); } - serde_json::from_str::(&body).map_err(Into::into) + serde_json::from_str::(&response.body).map_err(Into::into) } -pub fn workspace_http_memory_tools( - workspace_id: impl Into, - base_url: impl Into, -) -> Vec { - let backend = WorkspaceHttpMemoryBackend::new(workspace_id, base_url); +pub fn workspace_http_memory_tools(client: Arc) -> Vec { + let backend = WorkspaceHttpMemoryBackend::new(client); vec![ memory_tool( "MemoryReadDocument", @@ -215,13 +192,10 @@ pub fn workspace_http_memory_tools( } pub fn workspace_http_memory_consolidation_tools( - workspace_id: impl Into, - base_url: impl Into, + client: Arc, ) -> Vec { - let workspace_id = workspace_id.into(); - let base_url = base_url.into(); - let mut tools = workspace_http_memory_tools(workspace_id.clone(), base_url.clone()); - let backend = WorkspaceHttpMemoryBackend::new(workspace_id, base_url); + let mut tools = workspace_http_memory_tools(client.clone()); + let backend = WorkspaceHttpMemoryBackend::new(client); tools.extend([ memory_tool( "MemoryStagingList", @@ -370,6 +344,14 @@ mod tests { use super::*; use llm_engine::tool::ToolDefinition; + fn test_client() -> Arc { + Arc::new(crate::worker::RuntimeWorkspaceHttpClient::new( + "workspace", + "http://backend", + "test-worker", + )) + } + fn tool_names(definitions: Vec) -> Vec { let mut names = definitions .into_iter() @@ -390,10 +372,7 @@ mod tests { #[test] fn normal_workspace_memory_tools_do_not_include_staging_tools() { - let names = tool_names(workspace_http_memory_tools( - "workspace".to_string(), - "http://backend".to_string(), - )); + let names = tool_names(workspace_http_memory_tools(test_client())); assert!(names.contains(&"MemoryQuery".to_string())); assert!(names.contains(&"MemoryReadDocument".to_string())); @@ -410,7 +389,7 @@ mod tests { #[test] fn document_update_schema_is_edit_like_and_staging_close_has_no_legacy_kinds() { let update_schema = tool_meta( - workspace_http_memory_tools("workspace".to_string(), "http://backend".to_string()), + workspace_http_memory_tools(test_client()), "MemoryUpdateDocument", ); assert_eq!( @@ -423,10 +402,7 @@ mod tests { assert!(update_schema["properties"].get("body_md").is_none()); let close_schema_text = tool_meta( - workspace_http_memory_consolidation_tools( - "workspace".to_string(), - "http://backend".to_string(), - ), + workspace_http_memory_consolidation_tools(test_client()), "MemoryStagingClose", ) .to_string(); @@ -440,10 +416,7 @@ mod tests { #[test] fn consolidation_workspace_memory_tools_include_staging_tools() { - let names = tool_names(workspace_http_memory_consolidation_tools( - "workspace".to_string(), - "http://backend".to_string(), - )); + let names = tool_names(workspace_http_memory_consolidation_tools(test_client())); assert!(names.contains(&"MemoryQuery".to_string())); assert!(names.contains(&"MemoryReadDocument".to_string())); diff --git a/crates/worker/src/feature/builtin/objective.rs b/crates/worker/src/feature/builtin/objective.rs index 9c1cbc9e..49f46515 100644 --- a/crates/worker/src/feature/builtin/objective.rs +++ b/crates/worker/src/feature/builtin/objective.rs @@ -14,26 +14,27 @@ use llm_engine::tool::{ use serde::{Deserialize, Serialize}; use serde_json::json; +use crate::worker::{WorkspaceClient, WorkspaceRequest, WorkspaceRequestMethod}; + #[derive(Clone, Debug)] pub struct WorkspaceHttpObjectiveBackend { - workspace_id: String, - base_url: String, + client: Arc, } impl WorkspaceHttpObjectiveBackend { - pub fn new(workspace_id: impl Into, base_url: impl Into) -> Self { - Self { - workspace_id: workspace_id.into(), - base_url: base_url.into().trim_end_matches('/').to_string(), - } + pub fn new(client: Arc) -> Self { + Self { client } } async fn list(&self, input: ObjectiveListInput) -> Result { - let mut url = format!("{}/api/w/{}/objectives", self.base_url, self.workspace_id); + let mut url = format!( + "/api/w/{}/objectives", + self.client.workspace_id().unwrap_or_default() + ); if let Some(limit) = input.limit { url.push_str(&format!("?limit={}", limit.min(1000))); } - let response = get_json::(&url) + let response = get_json::(self.client.as_ref(), &url) .await .map_err(backend_error)?; let count = response.items.len(); @@ -46,7 +47,7 @@ impl WorkspaceHttpObjectiveBackend { async fn show(&self, input: ObjectiveShowInput) -> Result { let id = validate_id(&input.id, "ObjectiveShow")?; let url = self.objective_url(id); - let response = get_json::(&url) + let response = get_json::(self.client.as_ref(), &url) .await .map_err(backend_error)?; Ok(objective_output( @@ -61,11 +62,18 @@ impl WorkspaceHttpObjectiveBackend { "ObjectiveCreate requires non-empty title".to_string(), )); } - let url = format!("{}/api/w/{}/objectives", self.base_url, self.workspace_id); - let response = - send_json::(reqwest::Method::POST, &url, &input) - .await - .map_err(backend_error)?; + let url = format!( + "/api/w/{}/objectives", + self.client.workspace_id().unwrap_or_default() + ); + let response = send_json::( + self.client.as_ref(), + reqwest::Method::POST, + &url, + &input, + ) + .await + .map_err(backend_error)?; Ok(objective_output( format!("Created objective {}", response.id), response, @@ -86,10 +94,14 @@ impl WorkspaceHttpObjectiveBackend { new_string: input.new_string, replace_all: input.replace_all, }; - let response = - send_json::(reqwest::Method::PATCH, &url, &body) - .await - .map_err(backend_error)?; + let response = send_json::( + self.client.as_ref(), + reqwest::Method::PATCH, + &url, + &body, + ) + .await + .map_err(backend_error)?; Ok(objective_output( format!("Edited objective {}", response.id), response, @@ -105,6 +117,7 @@ impl WorkspaceHttpObjectiveBackend { } let url = format!("{}/state", self.objective_url(id)); let response = send_json::( + self.client.as_ref(), reqwest::Method::POST, &url, &ObjectiveSetStateRequest { state: input.state }, @@ -122,6 +135,7 @@ impl WorkspaceHttpObjectiveBackend { let ticket_id = validate_id(&input.ticket_id, "ObjectiveLinkTicket")?; let url = format!("{}/ticket-links", self.objective_url(id)); let response = send_json::( + self.client.as_ref(), reqwest::Method::POST, &url, &ObjectiveLinkTicketRequest { @@ -143,7 +157,7 @@ impl WorkspaceHttpObjectiveBackend { let id = validate_id(&input.id, "ObjectiveUnlinkTicket")?; let ticket_id = validate_id(&input.ticket_id, "ObjectiveUnlinkTicket")?; let url = format!("{}/ticket-links/{}", self.objective_url(id), ticket_id); - let response = delete_json::(&url) + let response = delete_json::(self.client.as_ref(), &url) .await .map_err(backend_error)?; Ok(objective_output( @@ -153,17 +167,15 @@ impl WorkspaceHttpObjectiveBackend { } fn objective_url(&self, id: &str) -> String { - format!( - "{}/api/w/{}/objectives/{}", - self.base_url, self.workspace_id, id - ) + let workspace_id = self.client.workspace_id().unwrap_or_default(); + format!("/api/w/{workspace_id}/objectives/{id}") } } #[derive(Debug, thiserror::Error)] pub enum WorkspaceObjectiveBackendError { #[error("workspace objective backend request failed: {0}")] - Request(#[from] reqwest::Error), + Request(#[from] crate::worker::WorkspaceClientError), #[error("workspace objective backend returned HTTP {status}: {body}")] Http { status: reqwest::StatusCode, @@ -182,41 +194,55 @@ fn backend_error(error: WorkspaceObjectiveBackendError) -> ToolError { } async fn get_json Deserialize<'de>>( - url: &str, + client: &dyn WorkspaceClient, + path: &str, ) -> Result { - let response = reqwest::Client::new().get(url).send().await?; - decode_response(response).await + decode_response(client.execute(WorkspaceRequest::get(path))?) } async fn send_json Deserialize<'de>>( + client: &dyn WorkspaceClient, method: reqwest::Method, - url: &str, + path: &str, body: &B, ) -> Result { - let response = reqwest::Client::new() - .request(method, url) - .json(body) - .send() - .await?; - decode_response(response).await + let method = match method { + reqwest::Method::POST => WorkspaceRequestMethod::Post, + reqwest::Method::PUT => WorkspaceRequestMethod::Put, + reqwest::Method::PATCH => WorkspaceRequestMethod::Patch, + reqwest::Method::DELETE => WorkspaceRequestMethod::Delete, + _ => WorkspaceRequestMethod::Get, + }; + decode_response(client.execute(WorkspaceRequest::json( + method, + path, + serde_json::to_string(body)?, + ))?) } async fn delete_json Deserialize<'de>>( - url: &str, + client: &dyn WorkspaceClient, + path: &str, ) -> Result { - let response = reqwest::Client::new().delete(url).send().await?; - decode_response(response).await + decode_response(client.execute(WorkspaceRequest { + method: WorkspaceRequestMethod::Delete, + path: path.to_string(), + body: None, + })?) } -async fn decode_response Deserialize<'de>>( - response: reqwest::Response, +fn decode_response Deserialize<'de>>( + response: crate::worker::WorkspaceResponse, ) -> Result { - let status = response.status(); - let body = response.text().await?; - if !status.is_success() { - return Err(WorkspaceObjectiveBackendError::Http { status, body }); + let status = reqwest::StatusCode::from_u16(response.status) + .unwrap_or(reqwest::StatusCode::INTERNAL_SERVER_ERROR); + if !response.is_success() { + return Err(WorkspaceObjectiveBackendError::Http { + status, + body: response.body, + }); } - serde_json::from_str(&body).map_err(Into::into) + serde_json::from_str(&response.body).map_err(Into::into) } fn objective_output(summary: String, response: ObjectiveDetail) -> Result { @@ -236,11 +262,8 @@ fn validate_id<'a>(id: &'a str, tool_name: &str) -> Result<&'a str, ToolError> { Ok(id) } -pub fn workspace_http_objective_tools( - workspace_id: impl Into, - base_url: impl Into, -) -> Vec { - let backend = WorkspaceHttpObjectiveBackend::new(workspace_id, base_url); +pub fn workspace_http_objective_tools(client: Arc) -> Vec { + let backend = WorkspaceHttpObjectiveBackend::new(client); vec![ objective_tool( "ObjectiveList", @@ -600,10 +623,13 @@ mod tests { #[test] fn workspace_http_objective_tools_include_objective_crud_tools() { - let names = tool_names(workspace_http_objective_tools( - "workspace".to_string(), - "http://backend".to_string(), - )); + let names = tool_names(workspace_http_objective_tools(Arc::new( + crate::worker::RuntimeWorkspaceHttpClient::new( + "workspace", + "http://backend", + "test-worker", + ), + ))); assert_eq!( names, diff --git a/crates/worker/src/feature/builtin/session_explore.rs b/crates/worker/src/feature/builtin/session_explore.rs index 5585a332..b71d48b7 100644 --- a/crates/worker/src/feature/builtin/session_explore.rs +++ b/crates/worker/src/feature/builtin/session_explore.rs @@ -29,7 +29,7 @@ const FINISH_EXTRACTION_DESCRIPTION: &str = "Finish the extract worker run after #[derive(Clone)] pub(crate) struct SessionExploreState { view: Arc, - workspace_client: WorkspaceClient, + workspace_client: Arc, source: SourceRef, extract_run_id: String, staged: Arc>>, @@ -39,7 +39,7 @@ pub(crate) struct SessionExploreState { impl SessionExploreState { pub(crate) fn new( view: SessionReferenceView, - workspace_client: WorkspaceClient, + workspace_client: Arc, source: SourceRef, ) -> Self { Self { @@ -615,7 +615,7 @@ mod tests { fn stub_memory_backend_response( body: &'static str, - ) -> (WorkspaceClient, mpsc::Receiver) { + ) -> (Arc, mpsc::Receiver) { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let addr = listener.local_addr().unwrap(); let (tx, rx) = mpsc::channel(); @@ -659,7 +659,11 @@ mod tests { stream.write_all(response.as_bytes()).unwrap(); }); ( - WorkspaceClient::http("test-workspace", format!("http://{addr}")), + Arc::new(crate::worker::RuntimeWorkspaceHttpClient::new( + "test-workspace", + format!("http://{addr}"), + "test-worker", + )), rx, ) } @@ -668,7 +672,7 @@ mod tests { fn descriptor_declares_session_explore_tools() { let state = SessionExploreState::new( SessionReferenceView::new("segment-1", vec![Item::user_message("remember this")]), - WorkspaceClient::available("test-backend"), + crate::worker::marker_workspace_client(None, "test-backend"), SourceRef { segment_id: "segment-1".to_string(), range: [0, 0], diff --git a/crates/worker/src/feature/builtin/ticket.rs b/crates/worker/src/feature/builtin/ticket.rs index 3b1f7e2f..b9e36510 100644 --- a/crates/worker/src/feature/builtin/ticket.rs +++ b/crates/worker/src/feature/builtin/ticket.rs @@ -4,7 +4,10 @@ //! module only resolves the local backend root, declares the built-in feature, //! and contributes those tools through the normal feature registry path. -use std::path::{Path, PathBuf}; +use std::{ + path::{Path, PathBuf}, + sync::Arc, +}; use ticket::{ LocalTicketBackend, MarkdownText, NewOrchestrationPlanRecord, NewTicket, NewTicketEvent, @@ -22,6 +25,7 @@ use crate::feature::{ FeatureInstructionContribution, FeatureInstructionDeclaration, FeatureInstructionId, FeatureModule, ToolContribution, ToolDeclaration, }; +use crate::worker::{WorkspaceClient, WorkspaceRequest, WorkspaceRequestMethod}; const FEATURE_ID: &str = "ticket"; const FEATURE_NAME: &str = "Ticket tools"; @@ -183,13 +187,8 @@ const ORCHESTRATION_CONTROL_ADDITIONAL_TOOL_NAMES: &[&str] = &[ #[derive(Clone, Debug)] pub enum TicketFeatureBackend { - Local { - root: PathBuf, - }, - WorkspaceHttp { - workspace_id: String, - base_url: String, - }, + Local { root: PathBuf }, + WorkspaceClient(Arc), } impl From for TicketFeatureBackend { @@ -274,7 +273,7 @@ impl TicketFeature { pub fn backend_root(&self) -> Option<&Path> { match &self.backend { TicketFeatureBackend::Local { root } => Some(root), - TicketFeatureBackend::WorkspaceHttp { .. } => None, + TicketFeatureBackend::WorkspaceClient(_) => None, } } @@ -321,15 +320,9 @@ impl TicketFeature { .into(), ) } - TicketFeatureBackend::WorkspaceHttp { - workspace_id, - base_url, - } => Some( - TicketToolBackend::new(WorkspaceHttpTicketBackend::new( - workspace_id.clone(), - base_url.clone(), - )) - .with_record_language(self.record_language.as_deref()), + TicketFeatureBackend::WorkspaceClient(client) => Some( + TicketToolBackend::new(WorkspaceHttpTicketBackend::new(client.clone())) + .with_record_language(self.record_language.as_deref()), ), } } @@ -386,22 +379,18 @@ impl FeatureModule for TicketFeature { #[derive(Clone, Debug)] struct WorkspaceHttpTicketBackend { - workspace_id: String, - base_url: String, + client: Arc, } impl WorkspaceHttpTicketBackend { - fn new(workspace_id: String, base_url: String) -> Self { - Self { - workspace_id, - base_url: base_url.trim_end_matches('/').to_string(), - } + fn new(client: Arc) -> Self { + Self { client } } fn endpoint(&self) -> String { format!( - "{}/api/w/{}/tickets/backend", - self.base_url, self.workspace_id + "/api/w/{}/tickets/backend", + self.client.workspace_id().unwrap_or_default() ) } @@ -409,44 +398,44 @@ impl WorkspaceHttpTicketBackend { &self, operation: TicketBackendOperation, ) -> TicketResult { + let client = self.client.clone(); let endpoint = self.endpoint(); if tokio::runtime::Handle::try_current().is_ok() { - return std::thread::spawn(move || Self::invoke_http(endpoint, operation)) + return std::thread::spawn(move || Self::invoke_client(client, endpoint, operation)) .join() .map_err(|_| { TicketError::Conflict("ticket backend request thread panicked".to_string()) })?; } - Self::invoke_http(endpoint, operation) + Self::invoke_client(client, endpoint, operation) } - fn invoke_http( + fn invoke_client( + client: Arc, endpoint: String, operation: TicketBackendOperation, ) -> TicketResult { let body = serde_json::to_string(&operation).map_err(|error| { TicketError::Conflict(format!("serialize ticket operation: {error}")) })?; - let response = reqwest::blocking::Client::new() - .post(endpoint) - .header(reqwest::header::CONTENT_TYPE, "application/json") - .body(body) - .send() + let response = client + .execute(WorkspaceRequest::json( + WorkspaceRequestMethod::Post, + endpoint, + body, + )) .map_err(|error| { TicketError::Conflict(format!("ticket backend request failed: {error}")) })?; - let status = response.status(); - let text = response.text().map_err(|error| { - TicketError::Conflict(format!("ticket backend response failed: {error}")) - })?; - if !status.is_success() { + if !response.is_success() { return Err(TicketError::Conflict(format!( - "ticket backend returned HTTP {status}: {text}" + "ticket backend returned HTTP {}: {}", + response.status, response.body ))); } - match serde_json::from_str::(&text).map_err(|error| { - TicketError::Conflict(format!("decode ticket backend response: {error}")) - })? { + match serde_json::from_str::(&response.body).map_err( + |error| TicketError::Conflict(format!("decode ticket backend response: {error}")), + )? { TicketBackendHttpResponse::Ok { result } => Ok(result), TicketBackendHttpResponse::Error { message } => Err(TicketError::Conflict(message)), } @@ -1126,8 +1115,13 @@ provider = "github" #[tokio::test(flavor = "multi_thread")] async fn workspace_http_backend_invoke_is_safe_inside_async_context() { - let backend = - WorkspaceHttpTicketBackend::new("workspace-a".to_string(), "not-a-url".to_string()); + let backend = WorkspaceHttpTicketBackend::new(Arc::new( + crate::worker::RuntimeWorkspaceHttpClient::new( + "workspace-a", + "not-a-url", + "test-worker", + ), + )); let error = backend .invoke(TicketBackendOperation::DefaultIntakeReadyStateChangeBody { @@ -1167,7 +1161,9 @@ provider = "github" .unwrap(); }); - let backend = WorkspaceHttpTicketBackend::new("workspace-a".to_string(), base_url); + let backend = WorkspaceHttpTicketBackend::new(Arc::new( + crate::worker::RuntimeWorkspaceHttpClient::new("workspace-a", base_url, "test-worker"), + )); let created = backend.create(NewTicket::new("HTTP ticket")).unwrap(); server.join().unwrap(); diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index e792863e..66e74468 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -40,6 +40,9 @@ pub use runtime::dir::RuntimeDir; pub use segment_log_sink::SegmentLogSink; pub use shared_state::WorkerSharedState; pub use worker::{ - LocalWorkingDirectory, Worker, WorkerError, WorkerFilesystemAuthority, WorkerRunResult, - WorkerWorkspaceContext, WorkspaceClient, WorkspaceId, WorkspaceIdError, apply_worker_manifest, + LocalWorkingDirectory, RuntimeWorkspaceHttpClient, Worker, WorkerError, + WorkerFilesystemAuthority, WorkerRunResult, WorkerWorkspaceContext, WorkspaceClient, + WorkspaceClientError, WorkspaceId, WorkspaceIdError, WorkspaceRequest, WorkspaceRequestMethod, + WorkspaceResponse, apply_worker_manifest, marker_workspace_client, + unavailable_workspace_client, }; diff --git a/crates/worker/src/skill.rs b/crates/worker/src/skill.rs index 2c430786..2b1a1691 100644 --- a/crates/worker/src/skill.rs +++ b/crates/worker/src/skill.rs @@ -128,7 +128,7 @@ pub enum SkillClientError { #[error("workspace client kind `{0}` does not expose direct Skill HTTP operations")] UnsupportedClient(String), #[error("Skill request failed: {0}")] - Request(#[from] reqwest::Error), + Request(#[from] crate::worker::WorkspaceClientError), #[error("Skill API response JSON is invalid: {0}")] Json(#[from] serde_json::Error), #[error("Skill API returned HTTP {status}: {body}")] @@ -140,7 +140,7 @@ pub enum SkillClientError { InvalidBaseUrl(String), } -impl WorkspaceClient { +impl dyn WorkspaceClient + '_ { pub fn list_skills(&self) -> Result { self.get_skill_json("skills") } @@ -157,29 +157,21 @@ impl WorkspaceClient { &self, path: &str, ) -> Result { - let Self::Http { - workspace_id, - base_url, - } = self - else { - return match self { - Self::Available { kind } => Err(SkillClientError::UnsupportedClient(kind.clone())), - Self::Unavailable { reason } => Err(SkillClientError::Unavailable(reason.clone())), - Self::Http { .. } => unreachable!(), - }; - }; - if base_url.trim().is_empty() { - return Err(SkillClientError::InvalidBaseUrl(base_url.clone())); + let workspace_id = self + .workspace_id() + .ok_or_else(|| SkillClientError::UnsupportedClient(self.kind().to_string()))?; + let response = self.execute(crate::worker::WorkspaceRequest::get(format!( + "/api/w/{workspace_id}/{path}" + )))?; + let status = reqwest::StatusCode::from_u16(response.status) + .unwrap_or(reqwest::StatusCode::INTERNAL_SERVER_ERROR); + if !response.is_success() { + return Err(SkillClientError::Http { + status, + body: response.body, + }); } - let base = base_url.trim_end_matches('/'); - let url = format!("{base}/api/w/{workspace_id}/{path}"); - let response = reqwest::blocking::Client::new().get(url).send()?; - let status = response.status(); - let body = response.text()?; - if !status.is_success() { - return Err(SkillClientError::Http { status, body }); - } - Ok(serde_json::from_str(&body)?) + Ok(serde_json::from_str(&response.body)?) } } @@ -201,13 +193,23 @@ mod tests { let mut request_line = String::new(); reader.read_line(&mut request_line).unwrap(); assert!(request_line.starts_with("GET /api/w/ws-1/skills HTTP/1.1")); + let mut worker_header = None; + let mut authorization = None; loop { let mut line = String::new(); reader.read_line(&mut line).unwrap(); + if let Some(value) = line.strip_prefix("x-yoi-worker-id: ") { + worker_header = Some(value.trim().to_string()); + } + if let Some(value) = line.strip_prefix("authorization: ") { + authorization = Some(value.trim().to_string()); + } if line == "\r\n" || line.is_empty() { break; } } + assert_eq!(worker_header.as_deref(), Some("test-worker")); + assert_eq!(authorization.as_deref(), Some("Bearer test-credential")); let body = serde_json::json!({ "authority": "workspace-backend-skills-v0", "entries": [{ @@ -229,8 +231,13 @@ mod tests { .unwrap(); }); - let client = WorkspaceClient::http("ws-1", format!("http://{addr}")); - let catalog = client.list_skills().unwrap(); + let client = crate::worker::RuntimeWorkspaceHttpClient::new( + "ws-1", + format!("http://{addr}"), + "test-worker", + ) + .with_access_token(Some("test-credential".to_string())); + let catalog = (&client as &dyn WorkspaceClient).list_skills().unwrap(); assert_eq!(catalog.entries[0].name, "triage-errors"); assert_eq!(catalog.entries[0].provenance.id, "workspace:triage-errors"); handle.join().unwrap(); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 445b8512..82c10245 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -143,77 +143,368 @@ pub enum WorkspaceIdError { Empty, } -/// Narrow path-free workspace API handle injected by Runtime/host code. -/// -/// This is deliberately not a filesystem authority surface. A Worker may have a -/// workspace client without local filesystem authority, or neither. Local -/// path-backed implementations are represented only as a capability marker here; -/// the actual paths remain under [`WorkerFilesystemAuthority::Local`] or in host -/// adapter code. +/// One authority-bound operation sent through the Runtime-supplied Workspace client. #[derive(Debug, Clone, PartialEq, Eq)] -pub enum WorkspaceClient { - /// Runtime/host supplied an HTTP workspace API endpoint. - Http { - workspace_id: String, - base_url: String, - }, - /// Runtime/host supplied a workspace API handle. The string is an opaque - /// diagnostic/backend kind, not an endpoint, path, or secret-bearing value. - Available { kind: String }, - /// Workspace-aware operations must fail closed or stay disabled. - Unavailable { reason: String }, +pub struct WorkspaceRequest { + pub method: WorkspaceRequestMethod, + pub path: String, + pub body: Option, } -impl WorkspaceClient { - pub fn available(kind: impl Into) -> Self { - Self::Available { kind: kind.into() } +impl WorkspaceRequest { + pub fn get(path: impl Into) -> Self { + Self { + method: WorkspaceRequestMethod::Get, + path: path.into(), + body: None, + } } - pub fn http(workspace_id: impl Into, base_url: impl Into) -> Self { - Self::Http { + pub fn json( + method: WorkspaceRequestMethod, + path: impl Into, + body: impl Into, + ) -> Self { + Self { + method, + path: path.into(), + body: Some(body.into()), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WorkspaceRequestMethod { + Get, + Post, + Put, + Patch, + Delete, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WorkspaceResponse { + pub status: u16, + pub body: String, +} + +impl WorkspaceResponse { + pub fn is_success(&self) -> bool { + (200..300).contains(&self.status) + } +} + +#[derive(Debug, thiserror::Error)] +pub enum WorkspaceClientError { + #[error("workspace client is unavailable: {0}")] + Unavailable(String), + #[error("workspace request path must start with '/': {0}")] + InvalidPath(String), + #[error("workspace request failed: {0}")] + Request(String), +} + +/// Path-free Workspace operation authority injected by Runtime/host code. +/// +/// Workers receive this trait object rather than a Backend URL. The concrete +/// implementation is responsible for binding Runtime/Worker identity and +/// forwarding operations to the Workspace authority. +pub trait WorkspaceClient: std::fmt::Debug + Send + Sync { + fn workspace_id(&self) -> Option<&str>; + fn kind(&self) -> &str; + fn is_available(&self) -> bool; + fn execute(&self, request: WorkspaceRequest) + -> Result; +} + +/// HTTP forwarding client created by Runtime for one concrete Worker execution. +/// +/// The upstream endpoint and source headers are private implementation details; +/// model-visible tools can only submit [`WorkspaceRequest`] values through the +/// [`WorkspaceClient`] trait. +pub struct RuntimeWorkspaceHttpClient { + workspace_id: String, + base_url: String, + worker_id: String, + access_token: Mutex>, +} + +impl std::fmt::Debug for RuntimeWorkspaceHttpClient { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("RuntimeWorkspaceHttpClient") + .field("workspace_id", &self.workspace_id) + .field("base_url", &self.base_url) + .field("worker_id", &self.worker_id) + .field( + "access_token", + &self + .access_token + .lock() + .ok() + .and_then(|token| token.as_ref().map(|_| "[redacted]")), + ) + .finish() + } +} + +impl RuntimeWorkspaceHttpClient { + pub fn new( + workspace_id: impl Into, + base_url: impl Into, + worker_id: impl Into, + ) -> Self { + Self { workspace_id: workspace_id.into(), - base_url: base_url.into(), + base_url: base_url.into().trim_end_matches('/').to_string(), + worker_id: worker_id.into(), + access_token: Mutex::new(None), } } - pub fn unavailable(reason: impl Into) -> Self { - Self::Unavailable { - reason: reason.into(), + pub fn with_access_token(self, access_token: Option) -> Self { + *self.access_token.lock().expect("new credential mutex") = access_token; + self + } +} + +impl WorkspaceClient for RuntimeWorkspaceHttpClient { + fn workspace_id(&self) -> Option<&str> { + Some(&self.workspace_id) + } + + fn kind(&self) -> &str { + "runtime-http-proxy" + } + + fn is_available(&self) -> bool { + true + } + + fn execute( + &self, + request: WorkspaceRequest, + ) -> Result { + let base_url = self.base_url.clone(); + let worker_id = self.worker_id.clone(); + let access_token = self + .access_token + .lock() + .map_err(|_| { + WorkspaceClientError::Request("workspace credential lock poisoned".to_string()) + })? + .clone(); + let request_copy = request.clone(); + let result = if tokio::runtime::Handle::try_current().is_ok() { + std::thread::spawn(move || { + execute_runtime_workspace_http_with_refresh( + &base_url, + &worker_id, + access_token, + request_copy, + ) + }) + .join() + .map_err(|_| { + WorkspaceClientError::Request("workspace request thread panicked".to_string()) + })? + } else { + execute_runtime_workspace_http_with_refresh( + &base_url, + &worker_id, + access_token, + request, + ) + }?; + if let Some(new_token) = result.1 { + *self.access_token.lock().map_err(|_| { + WorkspaceClientError::Request("workspace credential lock poisoned".to_string()) + })? = Some(new_token); } + Ok(result.0) + } +} + +fn execute_runtime_workspace_http_with_refresh( + base_url: &str, + worker_id: &str, + access_token: Option, + request: WorkspaceRequest, +) -> Result<(WorkspaceResponse, Option), WorkspaceClientError> { + let response = execute_runtime_workspace_http( + base_url, + worker_id, + access_token.as_deref(), + request.clone(), + )?; + if response.status != 401 { + return Ok((response, None)); + } + let Some(expired_token) = access_token else { + return Ok((response, None)); + }; + let workspace_id = request + .path + .strip_prefix("/api/w/") + .and_then(|path| path.split('/').next()) + .ok_or_else(|| WorkspaceClientError::InvalidPath(request.path.clone()))?; + let refresh_url = format!("{base_url}/api/w/{workspace_id}/worker-credentials/refresh"); + let refresh = reqwest::blocking::Client::new() + .post(refresh_url) + .bearer_auth(expired_token) + .header("x-yoi-worker-id", worker_id) + .send() + .map_err(|error| WorkspaceClientError::Request(error.to_string()))?; + if !refresh.status().is_success() { + return Ok((response, None)); + } + let body: serde_json::Value = refresh + .json() + .map_err(|error| WorkspaceClientError::Request(error.to_string()))?; + let new_token = body + .get("access_token") + .and_then(serde_json::Value::as_str) + .ok_or_else(|| { + WorkspaceClientError::Request( + "Workspace credential refresh response omitted access_token".to_string(), + ) + })? + .to_string(); + let retried = execute_runtime_workspace_http(base_url, worker_id, Some(&new_token), request)?; + Ok((retried, Some(new_token))) +} + +fn execute_runtime_workspace_http( + base_url: &str, + worker_id: &str, + access_token: Option<&str>, + request: WorkspaceRequest, +) -> Result { + if !request.path.starts_with('/') || request.path.starts_with("//") { + return Err(WorkspaceClientError::InvalidPath(request.path)); + } + let url = format!("{base_url}{}", request.path); + let method = match request.method { + WorkspaceRequestMethod::Get => reqwest::Method::GET, + WorkspaceRequestMethod::Post => reqwest::Method::POST, + WorkspaceRequestMethod::Put => reqwest::Method::PUT, + WorkspaceRequestMethod::Patch => reqwest::Method::PATCH, + WorkspaceRequestMethod::Delete => reqwest::Method::DELETE, + }; + let client = reqwest::blocking::Client::new(); + let mut request_builder = client + .request(method, url) + .header("x-yoi-worker-id", worker_id); + if let Some(access_token) = access_token { + request_builder = request_builder.bearer_auth(access_token); + } + if let Some(body) = request.body { + request_builder = request_builder + .header(reqwest::header::CONTENT_TYPE, "application/json") + .body(body); + } + let response = request_builder + .send() + .map_err(|error| WorkspaceClientError::Request(error.to_string()))?; + let status = response.status().as_u16(); + let body = response + .text() + .map_err(|error| WorkspaceClientError::Request(error.to_string()))?; + Ok(WorkspaceResponse { status, body }) +} + +#[derive(Debug)] +struct MarkerWorkspaceClient { + workspace_id: Option, + kind: String, + available: bool, + reason: String, +} + +impl WorkspaceClient for MarkerWorkspaceClient { + fn workspace_id(&self) -> Option<&str> { + self.workspace_id.as_deref() } - pub fn local_filesystem() -> Self { - Self::available("local-filesystem") + fn kind(&self) -> &str { + &self.kind } - pub fn is_available(&self) -> bool { - matches!(self, Self::Available { .. } | Self::Http { .. }) + fn is_available(&self) -> bool { + self.available } + + fn execute( + &self, + _request: WorkspaceRequest, + ) -> Result { + Err(WorkspaceClientError::Unavailable(self.reason.clone())) + } +} + +pub fn unavailable_workspace_client( + workspace_id: Option<&WorkspaceId>, + reason: impl Into, +) -> Arc { + Arc::new(MarkerWorkspaceClient { + workspace_id: workspace_id.map(|id| id.as_str().to_string()), + kind: "unavailable".to_string(), + available: false, + reason: reason.into(), + }) +} + +pub fn marker_workspace_client( + workspace_id: Option<&WorkspaceId>, + kind: impl Into, +) -> Arc { + let kind = kind.into(); + Arc::new(MarkerWorkspaceClient { + workspace_id: workspace_id.map(|id| id.as_str().to_string()), + reason: format!("workspace client kind `{kind}` does not expose Workspace operations"), + kind, + available: true, + }) } /// Workspace context supplied to a Worker separately from filesystem authority. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Clone)] pub struct WorkerWorkspaceContext { workspace_id: Option, - client: WorkspaceClient, + client: Arc, +} + +impl std::fmt::Debug for WorkerWorkspaceContext { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WorkerWorkspaceContext") + .field("workspace_id", &self.workspace_id) + .field("client_kind", &self.client.kind()) + .field("client_available", &self.client.is_available()) + .finish() + } } impl WorkerWorkspaceContext { pub fn no_workspace() -> Self { Self { workspace_id: None, - client: WorkspaceClient::unavailable("no workspace configured"), + client: unavailable_workspace_client(None, "no workspace configured"), } } pub fn unavailable(workspace_id: Option, reason: impl Into) -> Self { + let client = unavailable_workspace_client(workspace_id.as_ref(), reason); Self { workspace_id, - client: WorkspaceClient::unavailable(reason), + client, } } - pub fn with_client(workspace_id: Option, client: WorkspaceClient) -> Self { + pub fn with_client( + workspace_id: Option, + client: Arc, + ) -> Self { Self { workspace_id, client, @@ -221,15 +512,23 @@ impl WorkerWorkspaceContext { } pub fn local_filesystem(workspace_id: Option) -> Self { - Self::with_client(workspace_id, WorkspaceClient::local_filesystem()) + let client = marker_workspace_client(workspace_id.as_ref(), "local-filesystem"); + Self { + workspace_id, + client, + } } pub fn workspace_id(&self) -> Option<&WorkspaceId> { self.workspace_id.as_ref() } - pub fn client(&self) -> &WorkspaceClient { - &self.client + pub fn client(&self) -> &dyn WorkspaceClient { + self.client.as_ref() + } + + pub fn client_handle(&self) -> Arc { + self.client.clone() } } @@ -926,10 +1225,14 @@ impl Worker { /// Narrow workspace client/availability handle injected by Runtime/host. /// This never grants local filesystem authority. - pub fn workspace_client(&self) -> &WorkspaceClient { + pub fn workspace_client(&self) -> &dyn WorkspaceClient { self.workspace_context.client() } + pub fn workspace_client_handle(&self) -> Arc { + self.workspace_context.client_handle() + } + async fn resident_summary_from_workspace_authority( &self, ) -> Result, WorkerError> { @@ -3197,7 +3500,7 @@ impl Worker { items_to_extract, ); let session_explore_state = - SessionExploreState::new(session_view, self.workspace_client().clone(), source); + SessionExploreState::new(session_view, self.workspace_client_handle(), source); let input_text = render_extract_input(session_explore_state.view()); let mut internal_tools = Vec::new(); let mut internal_hook_builder = HookRegistryBuilder::new(); @@ -3464,7 +3767,7 @@ impl WorkerAuditBase { async fn emit( &self, - workspace_client: &WorkspaceClient, + workspace_client: &dyn WorkspaceClient, event_tx: Option<&broadcast::Sender>, status: memory::audit::WorkerLifecycleStatus, reason: impl Into, @@ -4936,7 +5239,7 @@ mod spawned_context_tests { false, WorkerWorkspaceContext::with_client( Some(workspace_id.clone()), - WorkspaceClient::available("test-api"), + marker_workspace_client(Some(&workspace_id), "test-api"), ), WorkerFilesystemAuthority::None, manifest.scope.clone(), @@ -5773,7 +6076,11 @@ mod build_summary_prompt_tests { }); WorkerWorkspaceContext::with_client( Some(WorkspaceId::new("test-memory").unwrap()), - WorkspaceClient::http("test-memory", format!("http://{addr}")), + Arc::new(RuntimeWorkspaceHttpClient::new( + "test-memory", + format!("http://{addr}"), + "test-worker", + )), ) } @@ -5905,7 +6212,11 @@ mod build_summary_prompt_tests { store, WorkerWorkspaceContext::with_client( Some(WorkspaceId::new("ws-skill").unwrap()), - WorkspaceClient::http("ws-skill", format!("http://{addr}")), + Arc::new(RuntimeWorkspaceHttpClient::new( + "ws-skill", + format!("http://{addr}"), + "test-worker", + )), ), authority, scope, @@ -5946,6 +6257,74 @@ mod build_summary_prompt_tests { })); } + #[test] + fn runtime_workspace_client_refreshes_expired_credential_and_retries() { + use std::io::{BufRead, BufReader, Write}; + use std::net::TcpListener; + + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let server = std::thread::spawn(move || { + for step in 0..3 { + let (mut stream, _) = listener.accept().unwrap(); + let mut reader = BufReader::new(stream.try_clone().unwrap()); + let mut first_line = String::new(); + reader.read_line(&mut first_line).unwrap(); + let mut authorization = String::new(); + loop { + let mut line = String::new(); + reader.read_line(&mut line).unwrap(); + if let Some(value) = line.strip_prefix("authorization: ") { + authorization = value.trim().to_string(); + } + if line == "\r\n" || line.is_empty() { + break; + } + } + match step { + 0 => { + assert_eq!(authorization, "Bearer expired-token"); + stream + .write_all(b"HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\n\r\n") + .unwrap(); + } + 1 => { + assert!(first_line.contains("/worker-credentials/refresh")); + assert_eq!(authorization, "Bearer expired-token"); + let body = + r#"{"access_token":"fresh-token","expires_at":"2099-01-01T00:00:00Z"}"#; + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", + body.len(), + body + ) + .unwrap(); + } + _ => { + assert_eq!(authorization, "Bearer fresh-token"); + stream + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\n{}") + .unwrap(); + } + } + } + }); + let client = RuntimeWorkspaceHttpClient::new( + "workspace-refresh", + format!("http://{address}"), + "worker-refresh", + ) + .with_access_token(Some("expired-token".to_string())); + let response = client + .execute(WorkspaceRequest::get( + "/api/w/workspace-refresh/tickets/backend", + )) + .unwrap(); + assert_eq!(response.status, 200); + server.join().unwrap(); + } + fn minimal_manifest() -> WorkerManifest { let toml_str = r#" [worker] diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index d860992b..e44f828a 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -308,6 +308,27 @@ pub struct WorkerSpawnWorkingDirectoryRequest { pub selector: Option, } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct WorkerTicketAssignmentRequest { + pub ticket_id: String, + pub operation_id: String, +} + +pub(crate) fn worker_spawn_idempotency( + request: &WorkerSpawnRequest, +) -> Result, String> { + let Some(assignment) = request.ticket_assignment.as_ref() else { + return Ok(None); + }; + let encoded = serde_json::to_vec(request) + .map_err(|error| format!("serialize Worker spawn idempotency input: {error}"))?; + Ok(Some(( + assignment.operation_id.clone(), + format!("sha256:{}", digest_hex(&encoded, 64)), + ))) +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] #[serde(deny_unknown_fields)] pub struct WorkerSpawnRequest { @@ -317,6 +338,8 @@ pub struct WorkerSpawnRequest { pub acceptance: WorkerSpawnAcceptanceRequirement, pub profile: ProfileSelector, #[serde(default, skip_serializing_if = "Option::is_none")] + pub ticket_assignment: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] pub initial_input: Option, /// Optional safe working-directory creation request. The Workspace server resolves /// this into a runtime-internal `WorkingDirectoryRequest` from configured @@ -329,6 +352,8 @@ pub struct WorkerSpawnRequest { pub resolved_working_directory: Option, #[serde(skip, default)] pub resolved_config_bundle: Option, + #[serde(skip, default)] + pub resolved_workspace_api: Option, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] @@ -427,6 +452,8 @@ pub struct WorkerStopResult { pub struct WorkerLifecycleRequest { #[serde(default, skip_serializing_if = "Option::is_none")] pub reason: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ticket_assignment: Option, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] @@ -1695,7 +1722,14 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime { }; } }; + let (idempotency_key, idempotency_fingerprint) = worker_spawn_idempotency(&request) + .expect("WorkerSpawnRequest serialization is infallible") + .map_or((None, None), |(key, fingerprint)| { + (Some(key), Some(fingerprint)) + }); let create_request = CreateWorkerRequest { + idempotency_key, + idempotency_fingerprint, profile, display_name: request.requested_worker_name.clone(), config_bundle: None, @@ -1703,13 +1737,16 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime { initial_input: request.initial_input.clone(), working_directory_request: request.resolved_working_directory_request.clone(), working_directory: request.resolved_working_directory.clone(), - workspace_api: self - .backend_base_url - .as_ref() - .map(|base_url| WorkspaceApiRef { - workspace_id: self.workspace_id.clone(), - base_url: base_url.clone(), - }), + workspace_api: request.resolved_workspace_api.clone().or_else(|| { + self.backend_base_url + .as_ref() + .map(|base_url| WorkspaceApiRef { + workspace_id: self.workspace_id.clone(), + base_url: base_url.clone(), + runtime_id: Some(self.runtime_id.clone()), + access_token: None, + }) + }), }; match self.runtime.create_worker(create_request) { Ok(detail) => WorkerSpawnResult { @@ -2669,7 +2706,14 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime { }; } }; + let (idempotency_key, idempotency_fingerprint) = worker_spawn_idempotency(&request) + .expect("WorkerSpawnRequest serialization is infallible") + .map_or((None, None), |(key, fingerprint)| { + (Some(key), Some(fingerprint)) + }); let create = CreateWorkerRequest { + idempotency_key, + idempotency_fingerprint, profile, display_name: request.requested_worker_name.clone(), config_bundle: None, @@ -2677,9 +2721,13 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime { initial_input: request.initial_input.clone(), working_directory_request: request.resolved_working_directory_request.clone(), working_directory: request.resolved_working_directory.clone(), - workspace_api: Some(WorkspaceApiRef { - workspace_id: self.workspace_id.clone(), - base_url: self.backend_base_url.clone(), + workspace_api: request.resolved_workspace_api.clone().or_else(|| { + Some(WorkspaceApiRef { + workspace_id: self.workspace_id.clone(), + base_url: self.backend_base_url.clone(), + runtime_id: Some(self.runtime_id.clone()), + access_token: None, + }) }), }; match self.post_json::<_, RuntimeHttpWorkerResponse>("/v1/workers", &create) { @@ -3121,8 +3169,11 @@ fn embedded_profile_path(profile: &ProfileSelector) -> Result { fn embedded_profile_label(profile: &ProfileSelector) -> Option { Some(match profile { ProfileSelector::Builtin(name) | ProfileSelector::Named(name) => { - if name.strip_prefix("builtin:").unwrap_or(name) == MEMORY_CONSOLIDATION_PROFILE { + let builtin_name = name.strip_prefix("builtin:").unwrap_or(name); + if builtin_name == MEMORY_CONSOLIDATION_PROFILE { MEMORY_CONSOLIDATION_PROFILE.to_string() + } else if builtin_name == WORKSPACE_ORCHESTRATOR_PROFILE { + WORKSPACE_ORCHESTRATOR_PROFILE.to_string() } else { safe_display_hint(name) } @@ -3132,6 +3183,8 @@ fn embedded_profile_label(profile: &ProfileSelector) -> Option { const MEMORY_CONSOLIDATION_PROFILE: &str = "memory-consolidation"; const MEMORY_CONSOLIDATION_SINGLETON_KEY: &str = "workspace-memory-consolidation"; +const WORKSPACE_ORCHESTRATOR_PROFILE: &str = "orchestrator"; +pub(crate) const WORKSPACE_ORCHESTRATOR_SINGLETON_KEY: &str = "workspace-orchestrator"; struct WorkerDisplayMetadata { display_name: String, @@ -3160,6 +3213,20 @@ fn worker_display_metadata( tags, }; } + if profile_label == Some(WORKSPACE_ORCHESTRATOR_PROFILE) { + let mut tags = vec!["orchestrator".to_string(), "singleton".to_string()]; + if internal { + tags.insert(0, "internal".to_string()); + } + return WorkerDisplayMetadata { + display_name: requested_display_name + .filter(|value| !value.trim().is_empty()) + .map(safe_display_hint) + .unwrap_or_else(|| "Workspace Orchestrator".to_string()), + singleton_key: Some(WORKSPACE_ORCHESTRATOR_SINGLETON_KEY.to_string()), + tags, + }; + } let display_name = requested_display_name .filter(|value| !value.trim().is_empty()) .map(safe_display_hint) @@ -4149,11 +4216,13 @@ mod tests { expected_segments: 0, }, profile: ProfileSelector::Builtin("builtin:coder".to_string()), + ticket_assignment: None, initial_input: None, working_directory_request: None, resolved_working_directory_request: None, resolved_working_directory: None, resolved_config_bundle: None, + resolved_workspace_api: None, } } @@ -4275,11 +4344,13 @@ mod tests { expected_segments: 0, }, profile: ProfileSelector::Builtin("builtin:coder".to_string()), + ticket_assignment: None, initial_input: None, working_directory_request: None, resolved_working_directory_request: None, resolved_working_directory: None, resolved_config_bundle: None, + resolved_workspace_api: None, }, ) .unwrap(); @@ -4371,11 +4442,13 @@ mod tests { expected_segments: 0, }, profile: ProfileSelector::Builtin("builtin:coder".to_string()), + ticket_assignment: None, initial_input: None, working_directory_request: None, resolved_working_directory_request: None, resolved_working_directory: None, resolved_config_bundle: None, + resolved_workspace_api: None, }, ) .unwrap(); @@ -4403,11 +4476,13 @@ mod tests { requested_worker_name: None, acceptance: WorkerSpawnAcceptanceRequirement::SocketReady, profile: ProfileSelector::Builtin("builtin:companion".to_string()), + ticket_assignment: None, initial_input: None, working_directory_request: None, resolved_working_directory_request: None, resolved_working_directory: None, resolved_config_bundle: None, + resolved_workspace_api: None, }, ) .unwrap(); diff --git a/crates/workspace-server/src/lib.rs b/crates/workspace-server/src/lib.rs index 0747f467..6ced71e4 100644 --- a/crates/workspace-server/src/lib.rs +++ b/crates/workspace-server/src/lib.rs @@ -85,6 +85,10 @@ pub enum Error { UnknownRepository(String), #[error("workspace id does not match this Workspace backend")] WorkspaceIdMismatch, + #[error("Ticket assignment conflict: {0}")] + TicketAssignmentConflict(String), + #[error("Worker Workspace authentication failed: {0}")] + WorkerWorkspaceAuthentication(String), #[error("workspace identity error: {0}")] WorkspaceIdentity(String), #[error("store error: {0}")] diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 15804702..bd0c7d37 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -1,7 +1,7 @@ -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::path::{Component, Path, PathBuf}; use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade}; use axum::extract::{Path as AxumPath, Query, State}; @@ -17,6 +17,7 @@ use memory::backend::{ MemoryConsolidationOutput, }; use protocol::stream::{decode_method, encode_event}; +use rusqlite::OptionalExtension; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use ticket::{ @@ -91,8 +92,9 @@ use crate::resource_broker::BackendResourceBroker; use crate::skills; use crate::store::{ AccountRecord, ApiTokenRecord, AuthChallengeRecord, BrowserSessionRecord, ControlPlaneStore, - DeviceLoginFlowRecord, PasskeyCredentialRecord, RepositoryRecord, UserRecord, - WorkdirRegistryRecord, WorkerRegistryRecord, WorkerWorkdirLinkRecord, WorkspaceRecord, + DeviceLoginFlowRecord, PasskeyCredentialRecord, RepositoryRecord, TicketWorkerAssignmentRecord, + UserRecord, WorkdirRegistryRecord, WorkerRegistryRecord, WorkerWorkdirLinkRecord, + WorkerWorkspaceCredentialRecord, WorkspaceRecord, }; use crate::{Error, Result}; use worker_runtime::catalog::{ @@ -488,6 +490,10 @@ pub fn build_router(api: WorkspaceApi) -> Router { ) .route("/api/tickets", get(list_tickets)) .route("/api/w/{workspace_id}/tickets", get(scoped_list_tickets)) + .route( + "/api/w/{workspace_id}/worker-credentials/refresh", + post(scoped_refresh_worker_workspace_credential), + ) .route( "/api/w/{workspace_id}/tickets/backend", post(scoped_ticket_backend_operation), @@ -520,6 +526,16 @@ pub fn build_router(api: WorkspaceApi) -> Router { "/api/w/{workspace_id}/tickets/{id}", get(scoped_get_ticket).patch(scoped_edit_ticket_item), ) + .route( + "/api/w/{workspace_id}/tickets/{id}/assignment", + get(scoped_get_ticket_worker_assignment) + .put(scoped_set_ticket_worker_assignment) + .delete(scoped_clear_ticket_worker_assignment), + ) + .route( + "/api/w/{workspace_id}/tickets/{id}/assignment/reassign", + post(scoped_reassign_ticket_worker_assignment), + ) .route( "/api/w/{workspace_id}/tickets/{id}/state", post(scoped_transition_ticket_state), @@ -780,7 +796,22 @@ pub async fn serve( listener: TcpListener, ) -> Result<()> { let api = WorkspaceApi::new(config, store).await?; - axum::serve(listener, build_router(api)).await?; + let dispatcher_api = api.clone(); + let dispatcher_workspace_id = dispatcher_api.config.workspace_id.clone(); + let dispatcher = tokio::spawn(async move { + loop { + let api = dispatcher_api.clone(); + let workspace_id = dispatcher_workspace_id.clone(); + let _ = tokio::task::spawn_blocking(move || { + dispatch_pending_ticket_notifications(&api, &workspace_id) + }) + .await; + tokio::time::sleep(std::time::Duration::from_secs(2)).await; + } + }); + let result = axum::serve(listener, build_router(api)).await; + dispatcher.abort(); + result?; Ok(()) } @@ -1530,6 +1561,241 @@ async fn scoped_get_ticket( get_ticket(State(api), AxumPath(path.id)).await } +#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)] +struct TicketWorkerAssignmentResponse { + workspace_id: String, + ticket_id: String, + assignment: Option, + worker: Option, +} + +#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)] +struct TicketWorkerAssignmentMutationResponse { + workspace_id: String, + ticket_id: String, + assignment: Option, + previous_assignment_id: Option, +} + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct SetTicketWorkerAssignmentRequest { + operation_id: String, + runtime_id: String, + worker_id: String, + expected_assignment_id: Option, + assigned_by: Option, +} + +#[derive(Debug, Default, Deserialize)] +#[serde(deny_unknown_fields)] +struct ClearTicketWorkerAssignmentQuery { + operation_id: Option, + expected_assignment_id: Option, + actor: Option, +} + +async fn scoped_get_ticket_worker_assignment( + State(api): State, + AxumPath(path): AxumPath, +) -> ApiResult> { + validate_workspace_scope(&api, &path.workspace_id)?; + let ticket = api.authority.ticket(&path.id)?; + let assignment = api + .store + .get_current_ticket_worker_assignment(&path.workspace_id, &ticket.id)?; + let worker = assignment.as_ref().and_then(|assignment| { + api.runtime + .worker(&assignment.runtime_id, &assignment.worker_id) + .ok() + }); + Ok(Json(TicketWorkerAssignmentResponse { + workspace_id: path.workspace_id, + ticket_id: ticket.id, + assignment, + worker, + })) +} + +async fn scoped_set_ticket_worker_assignment( + State(api): State, + AxumPath(path): AxumPath, + Json(request): Json, +) -> ApiResult> { + set_ticket_worker_assignment(api, path, request, false).await +} + +async fn scoped_reassign_ticket_worker_assignment( + State(api): State, + AxumPath(path): AxumPath, + Json(request): Json, +) -> ApiResult> { + set_ticket_worker_assignment(api, path, request, true).await +} + +async fn set_ticket_worker_assignment( + api: WorkspaceApi, + path: ScopedRecordPath, + request: SetTicketWorkerAssignmentRequest, + allow_reassign: bool, +) -> ApiResult> { + validate_workspace_scope(&api, &path.workspace_id)?; + let ticket = api.authority.ticket(&path.id)?; + let operation_id = require_ticket_assignment_value("operation_id", request.operation_id)?; + let runtime_id = require_ticket_assignment_value("runtime_id", request.runtime_id)?; + let worker_id = require_ticket_assignment_value("worker_id", request.worker_id)?; + let expected_assignment_id = request + .expected_assignment_id + .map(|value| require_ticket_assignment_value("expected_assignment_id", value)) + .transpose()?; + let assigned_by = request + .assigned_by + .map(|value| require_ticket_assignment_value("assigned_by", value)) + .transpose()? + .unwrap_or_else(|| "workspace-api".to_string()); + let worker = api + .runtime + .worker(&runtime_id, &worker_id) + .map_err(|err| err.into_error())?; + let assigned_at = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); + let record = TicketWorkerAssignmentRecord { + workspace_id: path.workspace_id.clone(), + ticket_id: ticket.id.clone(), + assignment_id: new_id("tasg"), + runtime_id: worker.runtime_id, + worker_id: worker.worker_id, + assigned_by, + assigned_at, + }; + let update = api.store.set_current_ticket_worker_assignment( + &record, + expected_assignment_id.as_deref(), + &new_id("tasev"), + &operation_id, + allow_reassign, + )?; + Ok(Json(TicketWorkerAssignmentMutationResponse { + workspace_id: path.workspace_id, + ticket_id: ticket.id, + assignment: Some(update.current), + previous_assignment_id: update.previous.map(|assignment| assignment.assignment_id), + })) +} + +async fn scoped_clear_ticket_worker_assignment( + State(api): State, + AxumPath(path): AxumPath, + Query(query): Query, +) -> ApiResult> { + validate_workspace_scope(&api, &path.workspace_id)?; + let ticket = api.authority.ticket(&path.id)?; + let operation_id = query + .operation_id + .map(|value| require_ticket_assignment_value("operation_id", value)) + .transpose()? + .ok_or_else(|| { + Error::TicketAssignmentConflict("unassign requires operation_id".to_string()) + })?; + let expected_assignment_id = query + .expected_assignment_id + .map(|value| require_ticket_assignment_value("expected_assignment_id", value)) + .transpose()?; + let actor = query + .actor + .map(|value| require_ticket_assignment_value("actor", value)) + .transpose()? + .unwrap_or_else(|| "workspace-api".to_string()); + let previous = api.store.clear_current_ticket_worker_assignment( + &path.workspace_id, + &ticket.id, + expected_assignment_id.as_deref(), + &operation_id, + &new_id("tasev"), + &actor, + &Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + )?; + Ok(Json(TicketWorkerAssignmentMutationResponse { + workspace_id: path.workspace_id, + ticket_id: ticket.id, + assignment: None, + previous_assignment_id: previous.map(|assignment| assignment.assignment_id), + })) +} + +fn assign_ticket_worker_from_lifecycle( + api: &WorkspaceApi, + assignment: &crate::hosts::WorkerTicketAssignmentRequest, + runtime_id: &str, + worker_id: &str, +) -> Result { + let ticket = api.authority.ticket(&assignment.ticket_id)?; + let assigned_at = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); + let record = TicketWorkerAssignmentRecord { + workspace_id: api.config.workspace_id.clone(), + ticket_id: ticket.id, + assignment_id: new_id("tasg"), + runtime_id: runtime_id.to_string(), + worker_id: worker_id.to_string(), + assigned_by: "worker-lifecycle".to_string(), + assigned_at, + }; + Ok(api + .store + .set_current_ticket_worker_assignment( + &record, + None, + &new_id("tasev"), + &assignment.operation_id, + false, + )? + .current) +} + +fn existing_lifecycle_assignment_worker( + api: &WorkspaceApi, + assignment: &crate::hosts::WorkerTicketAssignmentRequest, + runtime_id: &str, +) -> Result> { + let Some(operation) = api + .store + .get_ticket_assignment_operation(&api.config.workspace_id, &assignment.operation_id)? + else { + return Ok(None); + }; + if operation.action != "assign" + || operation.ticket_id != assignment.ticket_id + || operation.runtime_id.as_deref() != Some(runtime_id) + { + return Err(Error::TicketAssignmentConflict(format!( + "assignment operation {} was already used with different lifecycle input", + assignment.operation_id + ))); + } + let Some(worker_id) = operation.worker_id else { + return Ok(None); + }; + let worker = api + .runtime + .worker(runtime_id, &worker_id) + .map_err(|error| error.into_error())?; + if operation.assignment_id.is_none() && worker.state == "stopped" { + return Ok(None); + } + Ok(Some(worker)) +} + +fn require_ticket_assignment_value(field: &str, value: String) -> Result { + let value = value.trim(); + if value.is_empty() { + return Err(Error::RuntimeOperationFailed { + runtime_id: "workspace-server".to_string(), + code: "invalid_ticket_assignment".to_string(), + message: format!("{field} must not be empty"), + }); + } + Ok(value.to_string()) +} + #[derive(Debug, Deserialize)] #[serde(deny_unknown_fields)] struct BrowserEditTicketRequest { @@ -1744,21 +2010,130 @@ async fn scoped_close_ticket( browser_ticket_detail(&api, &path.id) } +#[derive(Debug, Serialize)] +struct WorkerWorkspaceCredentialRefreshResponse { + access_token: String, + expires_at: String, +} + +async fn scoped_refresh_worker_workspace_credential( + State(api): State, + AxumPath(path): AxumPath, + headers: HeaderMap, +) -> ApiResult> { + validate_workspace_scope(&api, &path.workspace_id)?; + let token = headers + .get(axum::http::header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.strip_prefix("Bearer ")) + .filter(|value| !value.trim().is_empty()) + .ok_or_else(|| { + Error::WorkerWorkspaceAuthentication("missing expired credential".to_string()) + })?; + let worker_id = headers + .get("x-yoi-worker-id") + .and_then(|value| value.to_str().ok()) + .filter(|value| !value.trim().is_empty()) + .ok_or_else(|| { + Error::WorkerWorkspaceAuthentication("missing Runtime-bound Worker id".to_string()) + })?; + let new_token = mint_secret("wac"); + let expires_at = + (Utc::now() + chrono::Duration::hours(1)).to_rfc3339_opts(SecondsFormat::Secs, true); + let credential = api + .store + .refresh_worker_workspace_credential( + token, + &path.workspace_id, + worker_id, + &new_token, + &expires_at, + )? + .ok_or_else(|| { + Error::WorkerWorkspaceAuthentication("credential cannot be refreshed".to_string()) + })?; + api.runtime + .worker(&credential.runtime_id, worker_id) + .map_err(|_| { + Error::WorkerWorkspaceAuthentication("credential Worker no longer exists".to_string()) + })?; + Ok(Json(WorkerWorkspaceCredentialRefreshResponse { + access_token: new_token, + expires_at, + })) +} + async fn scoped_ticket_backend_operation( State(api): State, AxumPath(path): AxumPath, - Json(operation): Json, + headers: HeaderMap, + Json(mut operation): Json, ) -> ApiResult> { validate_workspace_scope(&api, &path.workspace_id)?; let config = ticket::config::TicketConfig::load_workspace(&api.config.workspace_root) .map_err(|error| Error::Config(format!("load Ticket workspace settings: {error}")))?; - let backend = SqliteTicketBackend::new( + let mut backend = SqliteTicketBackend::new( api.config.database_path.clone(), api.config.workspace_id.clone(), ) .with_record_language(config.ticket_record_language()); + let operation_kind = ticket_mutation_operation_kind(&operation); + let is_mutation = operation_kind != "read"; + let target = ticket_mutation_target(&operation).cloned(); + let read_target = ticket_read_target(&operation).cloned(); + let has_worker_credential = headers.contains_key(axum::http::header::AUTHORIZATION); + let source = if is_mutation || (read_target.is_some() && has_worker_credential) { + Some(authenticate_worker_mutation_source( + &api, + &path.workspace_id, + &headers, + )?) + } else { + None + }; + let before = target.as_ref().and_then(|id| backend.show(id.clone()).ok()); + if let Some(source) = source.as_ref() { + bind_worker_ticket_operation_source(source, &mut operation); + let source_context = + worker_ticket_source_context(&api, &path.workspace_id, source, before.as_ref()); + backend = backend + .with_event_attributes(source_context.attributes(operation_kind)) + .with_mutation_hook(build_ticket_notification_hook( + &api, + source_context, + operation_kind, + before + .as_ref() + .map(|ticket| ticket.meta.workflow_state.as_str().to_string()) + .unwrap_or_else(|| ticket_operation_initial_state(&operation)), + )); + } let response = match execute_ticket_backend_operation(&backend, operation) { - Ok(result) => TicketBackendHttpResponse::Ok { result }, + Ok(result) => { + if let (Some(source), Some(read_target)) = (source.as_ref(), read_target.as_ref()) { + if let Ok(ticket) = backend.show(read_target.clone()) { + if let Some(event_index) = ticket.events.last().and_then(|event| { + event + .attributes + .get("event_sequence") + .and_then(|value| value.parse::().ok()) + }) { + api.store.upsert_ticket_notification_cursor( + &path.workspace_id, + &ticket.meta.id, + &source.runtime_id, + &source.worker_id, + event_index, + &Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + )?; + } + } + } + if is_mutation && source.is_some() { + dispatch_pending_ticket_notifications(&api, &path.workspace_id); + } + TicketBackendHttpResponse::Ok { result } + } Err(error) => TicketBackendHttpResponse::Error { message: error.to_string(), }, @@ -1766,6 +2141,441 @@ async fn scoped_ticket_backend_operation( Ok(Json(response)) } +#[derive(Debug, Clone, PartialEq, Eq)] +struct WorkerMutationSource { + runtime_id: String, + worker_id: String, +} + +fn ticket_mutation_target(operation: &TicketBackendOperation) -> Option<&TicketIdOrSlug> { + match operation { + TicketBackendOperation::EditItem { id, .. } + | TicketBackendOperation::AddEvent { id, .. } + | TicketBackendOperation::AddStateChanged { id, .. } + | TicketBackendOperation::AddIntakeSummary { id, .. } + | TicketBackendOperation::SetStateField { id, .. } + | TicketBackendOperation::SetWorkflowState { id, .. } + | TicketBackendOperation::MarkIntakeReady { id, .. } + | TicketBackendOperation::QueueReady { id, .. } + | TicketBackendOperation::Review { id, .. } + | TicketBackendOperation::Close { id, .. } + | TicketBackendOperation::AddTicketRelation { id, .. } + | TicketBackendOperation::AddOrchestrationPlanRecord { id, .. } => Some(id), + _ => None, + } +} + +fn bind_worker_ticket_operation_source( + source: &WorkerMutationSource, + operation: &mut TicketBackendOperation, +) { + let author = format!("worker:{}/{}", source.runtime_id, source.worker_id); + match operation { + TicketBackendOperation::Create { input } => input.author = Some(author), + TicketBackendOperation::EditItem { edit, .. } => edit.author = Some(author), + TicketBackendOperation::AddEvent { event, .. } => event.author = Some(author), + TicketBackendOperation::AddStateChanged { change, .. } + | TicketBackendOperation::SetStateField { change, .. } + | TicketBackendOperation::SetWorkflowState { change, .. } => change.author = Some(author), + TicketBackendOperation::AddIntakeSummary { summary, .. } => summary.author = Some(author), + TicketBackendOperation::MarkIntakeReady { + summary, change, .. + } => { + summary.author = Some(author.clone()); + change.author = Some(author); + } + TicketBackendOperation::QueueReady { queued_by, .. } => *queued_by = author, + TicketBackendOperation::Review { review, .. } => review.author = Some(author), + TicketBackendOperation::AddTicketRelation { relation, .. } => { + relation.author = Some(author) + } + TicketBackendOperation::AddOrchestrationPlanRecord { record, .. } => { + record.author = Some(author) + } + _ => {} + } +} + +fn ticket_read_target(operation: &TicketBackendOperation) -> Option<&TicketIdOrSlug> { + match operation { + TicketBackendOperation::Show { id } => Some(id), + _ => None, + } +} + +fn ticket_mutation_operation_kind(operation: &TicketBackendOperation) -> &'static str { + match operation { + TicketBackendOperation::Create { .. } => "create", + TicketBackendOperation::EditItem { .. } => "edit_item", + TicketBackendOperation::AddEvent { .. } => "add_event", + TicketBackendOperation::AddStateChanged { .. } => "add_state_changed", + TicketBackendOperation::AddIntakeSummary { .. } => "add_intake_summary", + TicketBackendOperation::SetStateField { .. } => "set_state_field", + TicketBackendOperation::SetWorkflowState { .. } => "set_workflow_state", + TicketBackendOperation::MarkIntakeReady { .. } => "mark_intake_ready", + TicketBackendOperation::QueueReady { .. } => "queue_ready", + TicketBackendOperation::Review { .. } => "review", + TicketBackendOperation::Close { .. } => "close", + TicketBackendOperation::AddTicketRelation { .. } => "add_relation", + TicketBackendOperation::AddOrchestrationPlanRecord { .. } => "add_plan_record", + _ => "read", + } +} + +fn ticket_operation_initial_state(operation: &TicketBackendOperation) -> String { + match operation { + TicketBackendOperation::Create { input } => input + .workflow_state + .as_ref() + .map(|state| state.as_str().to_string()) + .unwrap_or_else(|| TicketWorkflowState::Planning.as_str().to_string()), + _ => TicketWorkflowState::Planning.as_str().to_string(), + } +} + +#[derive(Debug, Clone)] +struct WorkerTicketSourceContext { + workspace_id: String, + runtime_id: String, + worker_id: String, + actor_role: String, + assignment_id: Option, + orchestrator: Option<(String, String)>, +} + +impl WorkerTicketSourceContext { + fn attributes(&self, operation_kind: &str) -> BTreeMap { + let mut attributes = BTreeMap::from([ + ("source_runtime_id".to_string(), self.runtime_id.clone()), + ("source_worker_id".to_string(), self.worker_id.clone()), + ("source_actor_role".to_string(), self.actor_role.clone()), + ( + "source_operation_kind".to_string(), + operation_kind.to_string(), + ), + ]); + if let Some(assignment_id) = &self.assignment_id { + attributes.insert("source_assignment_id".to_string(), assignment_id.clone()); + } + attributes + } +} + +fn worker_source_actor_role(is_current_assignment: bool, is_orchestrator: bool) -> &'static str { + if is_current_assignment { + "coder" + } else if is_orchestrator { + "orchestrator" + } else { + "worker" + } +} + +fn worker_ticket_source_context( + api: &WorkspaceApi, + workspace_id: &str, + source: &WorkerMutationSource, + ticket: Option<&ticket::Ticket>, +) -> WorkerTicketSourceContext { + let assignment = ticket.and_then(|ticket| { + api.store + .get_current_ticket_worker_assignment(workspace_id, &ticket.meta.id) + .ok() + .flatten() + }); + let orchestrator = find_workspace_orchestrator(api); + let is_current_assignment = assignment.as_ref().is_some_and(|assignment| { + assignment.runtime_id == source.runtime_id && assignment.worker_id == source.worker_id + }); + let is_orchestrator = orchestrator.as_ref().is_some_and(|worker| { + worker.runtime_id == source.runtime_id && worker.worker_id == source.worker_id + }); + let actor_role = worker_source_actor_role(is_current_assignment, is_orchestrator); + WorkerTicketSourceContext { + workspace_id: workspace_id.to_string(), + runtime_id: source.runtime_id.clone(), + worker_id: source.worker_id.clone(), + actor_role: actor_role.to_string(), + assignment_id: assignment.and_then(|assignment| { + (assignment.runtime_id == source.runtime_id && assignment.worker_id == source.worker_id) + .then_some(assignment.assignment_id) + }), + orchestrator: orchestrator.map(|worker| (worker.runtime_id, worker.worker_id)), + } +} + +fn build_ticket_notification_hook( + _api: &WorkspaceApi, + source: WorkerTicketSourceContext, + operation_kind: &'static str, + previous_state: String, +) -> Arc { + let invoked = AtomicBool::new(false); + let notification_id = new_id("tnfy"); + Arc::new(move |conn, event| { + if invoked.swap(true, Ordering::SeqCst) { + return Ok(()); + } + let current_state: String = conn + .query_row( + "SELECT workflow_state FROM typed_tickets WHERE workspace_id = ?1 AND ticket_id = ?2", + rusqlite::params![source.workspace_id, event.ticket_id], + |row| row.get(0), + ) + .map_err(|error| ticket::TicketError::Conflict(format!("read committed Ticket state for outbox: {error}")))?; + conn.execute( + r#"INSERT INTO ticket_notification_outbox ( + notification_id, workspace_id, ticket_id, event_sequence, + source_runtime_id, source_worker_id, previous_state, current_state, created_at, + event_kind, source_operation_kind, source_actor_role, source_assignment_id + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)"#, + rusqlite::params![ + notification_id, + source.workspace_id, + event.ticket_id, + event.event_index, + source.runtime_id, + source.worker_id, + previous_state, + current_state, + Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + event.event_kind.as_str(), + operation_kind, + source.actor_role, + source.assignment_id, + ], + ) + .map_err(|error| { + ticket::TicketError::Conflict(format!("insert Ticket notification outbox: {error}")) + })?; + + let assigned: Option<(String, String)> = conn + .query_row( + r#"SELECT runtime_id, worker_id FROM ticket_current_worker_assignments + WHERE workspace_id = ?1 AND ticket_id = ?2"#, + rusqlite::params![source.workspace_id, event.ticket_id], + |row| Ok((row.get(0)?, row.get(1)?)), + ) + .optional() + .map_err(|error| { + ticket::TicketError::Conflict(format!( + "resolve assigned notification recipient: {error}" + )) + })?; + let mut recipients = Vec::new(); + if let Some((runtime_id, worker_id)) = assigned { + if runtime_id != source.runtime_id || worker_id != source.worker_id { + recipients.push((runtime_id, worker_id, "assigned")); + } + } + if (matches!(previous_state.as_str(), "queued" | "inprogress") + || matches!(current_state.as_str(), "queued" | "inprogress")) + && let Some((runtime_id, worker_id)) = &source.orchestrator + && (*runtime_id != source.runtime_id || *worker_id != source.worker_id) + { + recipients.push((runtime_id.clone(), worker_id.clone(), "orchestrator")); + } + recipients.sort(); + recipients.dedup_by(|left, right| left.0 == right.0 && left.1 == right.1); + for (runtime_id, worker_id, recipient_kind) in recipients { + conn.execute( + r#"INSERT OR IGNORE INTO ticket_notification_deliveries ( + notification_id, recipient_runtime_id, recipient_worker_id, recipient_kind, attempts + ) VALUES (?1, ?2, ?3, ?4, 0)"#, + rusqlite::params![notification_id, runtime_id, worker_id, recipient_kind], + ) + .map_err(|error| ticket::TicketError::Conflict(format!("insert Ticket notification delivery: {error}")))?; + } + Ok(()) + }) +} + +fn authenticate_worker_mutation_source( + api: &WorkspaceApi, + workspace_id: &str, + headers: &HeaderMap, +) -> Result { + let token = headers + .get(axum::http::header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.strip_prefix("Bearer ")) + .filter(|value| !value.trim().is_empty()) + .ok_or_else(|| { + Error::WorkerWorkspaceAuthentication("missing Runtime Workspace credential".to_string()) + })?; + let worker_id = headers + .get("x-yoi-worker-id") + .and_then(|value| value.to_str().ok()) + .filter(|value| !value.trim().is_empty()) + .ok_or_else(|| { + Error::WorkerWorkspaceAuthentication("missing Runtime-bound Worker id".to_string()) + })?; + let credential = api + .store + .authenticate_worker_workspace_credential(token, workspace_id, worker_id)? + .ok_or_else(|| { + Error::WorkerWorkspaceAuthentication("invalid Runtime Workspace credential".to_string()) + })?; + api.runtime + .worker(&credential.runtime_id, worker_id) + .map_err(|_| { + Error::WorkerWorkspaceAuthentication( + "credential does not identify a current Runtime Worker".to_string(), + ) + })?; + Ok(WorkerMutationSource { + runtime_id: credential.runtime_id, + worker_id: worker_id.to_string(), + }) +} + +fn dispatch_pending_ticket_notifications(api: &WorkspaceApi, workspace_id: &str) { + let Ok(deliveries) = api + .store + .list_pending_ticket_notification_deliveries(workspace_id, 100) + else { + return; + }; + for delivery in deliveries { + let Some((current_runtime_id, current_worker_id)) = + current_ticket_notification_recipient(api, &delivery) + else { + continue; + }; + if current_runtime_id == delivery.source_runtime_id + && current_worker_id == delivery.source_worker_id + { + let _ = api.store.mark_ticket_notification_delivered( + &delivery.notification_id, + &delivery.recipient_runtime_id, + &delivery.recipient_worker_id, + &Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + ); + continue; + } + if current_runtime_id != delivery.recipient_runtime_id + || current_worker_id != delivery.recipient_worker_id + { + let _ = api.store.reroute_ticket_notification_delivery( + &delivery.notification_id, + &delivery.recipient_runtime_id, + &delivery.recipient_worker_id, + ¤t_runtime_id, + ¤t_worker_id, + ); + continue; + } + if api + .store + .get_ticket_notification_cursor( + &delivery.workspace_id, + &delivery.ticket_id, + ¤t_runtime_id, + ¤t_worker_id, + ) + .ok() + .flatten() + .is_some_and(|cursor| cursor >= delivery.event_sequence) + { + let _ = api.store.mark_ticket_notification_delivered( + &delivery.notification_id, + &delivery.recipient_runtime_id, + &delivery.recipient_worker_id, + &Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + ); + continue; + } + let result = api.runtime.send_input( + &delivery.recipient_runtime_id, + &delivery.recipient_worker_id, + WorkerInputRequest { + kind: WorkerInputKind::System, + content: format!( + "Ticket notification: workspace_id={} ticket_id={} event_sequence={} event_kind={} source_operation_kind={} source_runtime_id={} source_worker_id={}. Reread the Ticket before acting.", + delivery.workspace_id, + delivery.ticket_id, + delivery.event_sequence, + delivery.event_kind, + delivery.source_operation_kind, + delivery.source_runtime_id, + delivery.source_worker_id + ), + segments: None, + }, + ); + match result { + Ok(result) if result.state == WorkerOperationState::Accepted => { + let _ = api.store.mark_ticket_notification_delivered( + &delivery.notification_id, + &delivery.recipient_runtime_id, + &delivery.recipient_worker_id, + &Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + ); + } + Ok(result) => { + let _ = api.store.mark_ticket_notification_failed( + &delivery.notification_id, + &delivery.recipient_runtime_id, + &delivery.recipient_worker_id, + &format!("Runtime rejected notification: {:?}", result.diagnostics), + ); + } + Err(error) => { + let _ = api.store.mark_ticket_notification_failed( + &delivery.notification_id, + &delivery.recipient_runtime_id, + &delivery.recipient_worker_id, + &error.into_error().to_string(), + ); + } + } + } +} + +fn find_workspace_orchestrator(api: &WorkspaceApi) -> Option { + let is_orchestrator = |worker: &WorkerSummary| { + worker.singleton_key.as_deref() == Some(crate::hosts::WORKSPACE_ORCHESTRATOR_SINGLETON_KEY) + }; + if let Some(worker) = api + .runtime + .list_workers(1000) + .items + .into_iter() + .find(is_orchestrator) + { + return Some(worker); + } + for runtime in api.runtime.list_runtimes(1000).items { + if let Ok(stopped) = api + .runtime + .list_stopped_workers_for_runtime(&runtime.runtime_id, 1000) + { + if let Some(worker) = stopped.items.into_iter().find(is_orchestrator) { + return Some(worker); + } + } + } + None +} + +fn current_ticket_notification_recipient( + api: &WorkspaceApi, + delivery: &crate::store::TicketNotificationDeliveryRecord, +) -> Option<(String, String)> { + match delivery.recipient_kind.as_str() { + "assigned" => api + .store + .get_current_ticket_worker_assignment(&delivery.workspace_id, &delivery.ticket_id) + .ok() + .flatten() + .map(|assignment| (assignment.runtime_id, assignment.worker_id)), + "orchestrator" => { + find_workspace_orchestrator(api).map(|worker| (worker.runtime_id, worker.worker_id)) + } + _ => None, + } +} + #[derive(Debug, Clone, Serialize)] struct MemoryDocumentResponse { body_md: String, @@ -1898,11 +2708,13 @@ fn start_memory_staging_consolidation( expected_segments: 1, }, profile: profile_selector, + ticket_assignment: None, initial_input: Some(input), working_directory_request: None, resolved_working_directory_request: None, resolved_working_directory: None, resolved_config_bundle, + resolved_workspace_api: None, }, ) .map_err(|err| err.into_error())?; @@ -2959,6 +3771,7 @@ fn cleanup_runtime_worker_for_execution( candidate.runtime_worker_id.as_str(), WorkerLifecycleRequest { reason: Some("cleanup worker before deletion".to_string()), + ticket_assignment: None, }, ) { Ok(result) if result.state == WorkerOperationState::Accepted => {} @@ -2980,7 +3793,15 @@ fn cleanup_runtime_worker_for_execution( .runtime .delete_worker(runtime_id, candidate.runtime_worker_id.as_str()) { - Ok(result) if result.deleted && result.state == WorkerOperationState::Accepted => Ok(()), + Ok(result) if result.deleted && result.state == WorkerOperationState::Accepted => { + api.store.revoke_worker_workspace_credentials( + &api.config.workspace_id, + runtime_id, + candidate.runtime_worker_id.as_str(), + &Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + )?; + Ok(()) + } Ok(result) => Err(ApiError::with_diagnostics( Error::RuntimeOperationFailed { runtime_id: runtime_id.to_string(), @@ -3142,12 +3963,92 @@ async fn scoped_get_runtime_worker( get_runtime_worker(State(api), AxumPath((path.runtime_id, path.worker_id))).await } +#[derive(Debug, Default, Deserialize)] +struct RestoreTicketAssignmentQuery { + ticket_id: Option, + assignment_operation_id: Option, +} + async fn scoped_restore_runtime_worker( State(api): State, AxumPath(path): AxumPath, + Query(query): Query, ) -> ApiResult> { validate_workspace_scope(&api, &path.workspace_id)?; - restore_runtime_worker(State(api), AxumPath((path.runtime_id, path.worker_id))).await + let workspace_id = path.workspace_id.clone(); + let runtime_id = path.runtime_id.clone(); + let worker_id = path.worker_id.clone(); + let assignment_request = match ( + query.ticket_id.clone(), + query.assignment_operation_id.clone(), + ) { + (Some(ticket_id), Some(operation_id)) => { + Some(crate::hosts::WorkerTicketAssignmentRequest { + ticket_id, + operation_id, + }) + } + (None, None) => None, + _ => { + return Err(Error::TicketAssignmentConflict( + "restore assignment requires both ticket_id and assignment_operation_id" + .to_string(), + ) + .into()); + } + }; + if let Some(assignment) = assignment_request.as_ref() { + let fingerprint = format!( + "sha256:{}", + Sha256::digest(format!( + "restore\0{}\0{}\0{}", + assignment.ticket_id, runtime_id, worker_id + )) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::() + ); + api.store.reserve_ticket_assignment_operation( + &workspace_id, + &assignment.operation_id, + &assignment.ticket_id, + &runtime_id, + Some(&worker_id), + &fingerprint, + &Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + )?; + if let Some(worker) = existing_lifecycle_assignment_worker(&api, assignment, &runtime_id)? { + if worker.worker_id != worker_id { + return Err(Error::TicketAssignmentConflict(format!( + "assignment operation {} belongs to worker {}, not {}", + assignment.operation_id, worker.worker_id, worker_id + )) + .into()); + } + assign_ticket_worker_from_lifecycle(&api, assignment, &runtime_id, &worker_id)?; + dispatch_pending_ticket_notifications(&api, &workspace_id); + return Ok(Json(WorkerRestoreResponse { + workspace_id, + runtime_id, + worker_id, + result: crate::hosts::WorkerRestoreResult { + state: WorkerOperationState::Accepted, + worker: Some(worker), + diagnostics: Vec::new(), + }, + })); + } + } + let response = restore_runtime_worker( + State(api.clone()), + AxumPath((runtime_id.clone(), worker_id.clone())), + ) + .await?; + if let Some(assignment) = assignment_request.as_ref() { + assign_ticket_worker_from_lifecycle(&api, assignment, &runtime_id, &worker_id)?; + } + dispatch_pending_ticket_notifications(&api, &workspace_id); + Ok(response) } async fn scoped_pin_runtime_worker( @@ -4520,11 +5421,13 @@ async fn create_workspace_worker( expected_segments: if initial_input.is_some() { 1 } else { 0 }, }, profile: profile_selector, + ticket_assignment: None, initial_input, working_directory_request: None, resolved_working_directory_request: None, resolved_working_directory, resolved_config_bundle, + resolved_workspace_api: None, }, ) .map_err(|err| err.into_error())?; @@ -4817,6 +5720,19 @@ async fn create_runtime_worker( AxumPath(runtime_id): AxumPath, Json(mut request): Json, ) -> ApiResult> { + if let Some(assignment) = request.ticket_assignment.as_ref() { + if let Some(worker) = existing_lifecycle_assignment_worker(&api, assignment, &runtime_id)? { + assign_ticket_worker_from_lifecycle(&api, assignment, &runtime_id, &worker.worker_id)?; + dispatch_pending_ticket_notifications(&api, api.workspace_id()); + return Ok(Json(WorkerSpawnResult { + state: WorkerOperationState::Accepted, + worker: Some(worker), + acceptance_evidence: Vec::new(), + diagnostics: Vec::new(), + })); + } + } + let lifecycle_assignment = request.ticket_assignment.clone(); reject_workdir_for_embedded_runtime( &runtime_id, request.working_directory_request.is_some() || request.resolved_working_directory.is_some(), @@ -4846,6 +5762,41 @@ async fn create_runtime_worker( .map(|claim| claim.working_directory_id.clone()) }; let requested_worker_name = request.requested_worker_name.clone(); + if let Some(base_url) = api.config.backend_base_url.clone() { + let credential = WorkerWorkspaceCredentialRecord { + credential_id: new_id("wac"), + token: mint_secret("wac"), + workspace_id: api.config.workspace_id.clone(), + runtime_id: runtime_id.clone(), + worker_id: None, + created_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + expires_at: (Utc::now() + chrono::Duration::hours(1)) + .to_rfc3339_opts(SecondsFormat::Secs, true), + revoked_at: None, + }; + api.store.upsert_worker_workspace_credential(&credential)?; + request.resolved_workspace_api = Some(worker_runtime::catalog::WorkspaceApiRef { + workspace_id: api.config.workspace_id.clone(), + base_url, + runtime_id: Some(runtime_id.clone()), + access_token: Some(credential.token), + }); + } + let spawn_idempotency = + crate::hosts::worker_spawn_idempotency(&request).map_err(Error::Config)?; + if let (Some(assignment), Some((_, fingerprint))) = + (lifecycle_assignment.as_ref(), spawn_idempotency.as_ref()) + { + api.store.reserve_ticket_assignment_operation( + &api.config.workspace_id, + &assignment.operation_id, + &assignment.ticket_id, + &runtime_id, + None, + fingerprint, + &Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + )?; + } let result = api .runtime .spawn_worker(&runtime_id, request) @@ -4863,6 +5814,14 @@ async fn create_runtime_worker( worker.profile.clone(), WorkerRegistryDisplayNamePolicy::UseProvided, )?; + if let Some(assignment) = lifecycle_assignment.as_ref() { + api.store.bind_ticket_assignment_operation_worker( + &api.config.workspace_id, + &assignment.operation_id, + &worker.worker_id, + )?; + assign_ticket_worker_from_lifecycle(&api, assignment, &runtime_id, &worker.worker_id)?; + } if worker.working_directory.is_none() { if let Some(workdir_id) = prepared_workdir_id.as_deref() { if api @@ -6952,6 +7911,8 @@ impl ApiError { impl IntoResponse for ApiError { fn into_response(self) -> Response { let status = match &self.error { + Error::TicketAssignmentConflict(_) => StatusCode::CONFLICT, + Error::WorkerWorkspaceAuthentication(_) => StatusCode::UNAUTHORIZED, Error::InvalidRuntimeIdentifier { .. } => StatusCode::BAD_REQUEST, Error::Ticket(ticket::TicketError::NotFound(_)) => StatusCode::NOT_FOUND, Error::Ticket( @@ -7879,11 +8840,13 @@ mod tests { expected_segments: 0, }, profile: ProfileSelector::Builtin(MEMORY_CONSOLIDATION_PROFILE.to_string()), + ticket_assignment: None, initial_input: None, working_directory_request: None, resolved_working_directory_request: None, resolved_working_directory: None, resolved_config_bundle, + resolved_workspace_api: None, }, ) .unwrap(); @@ -7962,65 +8925,690 @@ mod tests { } } + #[test] + fn worker_source_actor_roles_use_canonical_vocabulary() { + assert_eq!(worker_source_actor_role(true, false), "coder"); + assert_eq!(worker_source_actor_role(false, true), "orchestrator"); + assert_eq!(worker_source_actor_role(false, false), "worker"); + assert_eq!(worker_source_actor_role(true, true), "coder"); + } + #[tokio::test] - async fn ticket_browser_endpoints_mutate_typed_backend_and_return_thread() { + async fn ticket_assignment_endpoints_read_and_clear_current_assignment() { let dir = tempfile::tempdir().unwrap(); let api = test_api(dir.path()).await; - let Json(created) = scoped_ticket_backend_operation( - State(api.clone()), - AxumPath(ScopedWorkspacePath { - workspace_id: TEST_WORKSPACE_ID.to_string(), - }), - Json(TicketBackendOperation::Create { - input: ticket::NewTicket::new("Browser Ticket API"), - }), - ) - .await - .unwrap(); - let ticket_id = match created { - TicketBackendHttpResponse::Ok { - result: ticket::TicketBackendOperationResult::TicketRef(ticket_ref), - } => ticket_ref.id, - other => panic!("unexpected create response: {other:?}"), + let created = browser_ticket_backend(&api) + .unwrap() + .create(ticket::NewTicket::new("Assigned Ticket")) + .unwrap(); + let ticket_id = created.id; + let assignment = TicketWorkerAssignmentRecord { + workspace_id: TEST_WORKSPACE_ID.to_string(), + ticket_id: ticket_id.clone(), + assignment_id: "assignment-api-1".to_string(), + runtime_id: "embedded".to_string(), + worker_id: "42".to_string(), + assigned_by: "test-user".to_string(), + assigned_at: TEST_CREATED_AT.to_string(), }; + api.store + .set_current_ticket_worker_assignment( + &assignment, + None, + "event-api-1", + "operation-api-1", + false, + ) + .unwrap(); let path = || ScopedRecordPath { workspace_id: TEST_WORKSPACE_ID.to_string(), id: ticket_id.clone(), }; - let Json(related) = scoped_ticket_backend_operation( + + let Json(read) = scoped_get_ticket_worker_assignment(State(api.clone()), AxumPath(path())) + .await + .unwrap(); + assert_eq!(read.assignment, Some(assignment)); + + let stale = scoped_clear_ticket_worker_assignment( State(api.clone()), - AxumPath(ScopedWorkspacePath { - workspace_id: TEST_WORKSPACE_ID.to_string(), + AxumPath(path()), + Query(ClearTicketWorkerAssignmentQuery { + operation_id: Some("clear-stale".to_string()), + expected_assignment_id: Some("stale-assignment".to_string()), + actor: Some("test-user".to_string()), }), - Json(TicketBackendOperation::Create { - input: ticket::NewTicket::new("Related Browser Ticket"), + ) + .await + .unwrap_err() + .into_response(); + assert_eq!(stale.status(), StatusCode::CONFLICT); + + let Json(cleared) = scoped_clear_ticket_worker_assignment( + State(api.clone()), + AxumPath(path()), + Query(ClearTicketWorkerAssignmentQuery { + operation_id: Some("clear-current".to_string()), + expected_assignment_id: Some("assignment-api-1".to_string()), + actor: Some("test-user".to_string()), }), ) .await .unwrap(); - let related_ticket_id = match related { - TicketBackendHttpResponse::Ok { - result: ticket::TicketBackendOperationResult::TicketRef(ticket_ref), - } => ticket_ref.id, - other => panic!("unexpected related create response: {other:?}"), + assert_eq!( + cleared.previous_assignment_id.as_deref(), + Some("assignment-api-1") + ); + assert_eq!(cleared.assignment, None); + } + + #[tokio::test] + async fn authenticated_worker_ticket_mutation_routes_durable_assignment_notification() { + let dir = tempfile::tempdir().unwrap(); + let api = test_api(dir.path()).await; + let spawn = |name: &str| WorkerSpawnRequest { + requested_worker_name: Some(name.to_string()), + intent: WorkerSpawnIntent::TicketRole { + ticket_id: name.to_string(), + role: TicketWorkerRole::Coder, + }, + acceptance: WorkerSpawnAcceptanceRequirement::RunAccepted { + expected_segments: 0, + }, + profile: ProfileSelector::Builtin("builtin:coder".to_string()), + ticket_assignment: None, + initial_input: None, + working_directory_request: None, + resolved_working_directory_request: None, + resolved_working_directory: None, + resolved_config_bundle: None, + resolved_workspace_api: None, }; + let source_worker = api + .runtime + .spawn_worker(EMBEDDED_WORKER_RUNTIME_ID, spawn("source-worker")) + .unwrap() + .worker + .unwrap(); + let recipient_worker = api + .runtime + .spawn_worker(EMBEDDED_WORKER_RUNTIME_ID, spawn("recipient-worker")) + .unwrap() + .worker + .unwrap(); + let backend = browser_ticket_backend(&api).unwrap(); + let ticket_ref = backend + .create(ticket::NewTicket::new("Notify assigned Worker")) + .unwrap(); + api.store + .set_current_ticket_worker_assignment( + &TicketWorkerAssignmentRecord { + workspace_id: TEST_WORKSPACE_ID.to_string(), + ticket_id: ticket_ref.id.clone(), + assignment_id: "notify-assignment".to_string(), + runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + worker_id: recipient_worker.worker_id.clone(), + assigned_by: "test-user".to_string(), + assigned_at: TEST_CREATED_AT.to_string(), + }, + None, + "notify-assignment-event", + "notify-assignment-operation", + false, + ) + .unwrap(); + api.store + .upsert_worker_workspace_credential(&WorkerWorkspaceCredentialRecord { + credential_id: "source-credential".to_string(), + token: "source-secret".to_string(), + workspace_id: TEST_WORKSPACE_ID.to_string(), + runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + worker_id: Some(source_worker.worker_id.clone()), + created_at: TEST_CREATED_AT.to_string(), + expires_at: "2099-01-01T00:00:00Z".to_string(), + revoked_at: None, + }) + .unwrap(); + let mut headers = HeaderMap::new(); + headers.insert( + axum::http::header::AUTHORIZATION, + axum::http::HeaderValue::from_static("Bearer source-secret"), + ); + headers.insert( + "x-yoi-worker-id", + axum::http::HeaderValue::from_str(&source_worker.worker_id).unwrap(), + ); + let Json(response) = scoped_ticket_backend_operation( + State(api.clone()), + AxumPath(ScopedWorkspacePath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + }), + headers.clone(), + Json(TicketBackendOperation::AddEvent { + id: ticket_ref.id.clone().into(), + event: NewTicketEvent::new(TicketEventKind::Comment, "implementation update"), + }), + ) + .await + .unwrap(); + assert!( + matches!(response, TicketBackendHttpResponse::Ok { .. }), + "unexpected response: {response:?}" + ); + let committed = backend.show(ticket_ref.id.clone().into()).unwrap(); + let committed_event = committed.events.last().unwrap(); + assert_eq!( + committed_event + .attributes + .get("source_runtime_id") + .map(String::as_str), + Some(EMBEDDED_WORKER_RUNTIME_ID) + ); + assert_eq!( + committed_event + .attributes + .get("source_worker_id") + .map(String::as_str), + Some(source_worker.worker_id.as_str()) + ); + assert_eq!( + committed_event + .attributes + .get("source_operation_kind") + .map(String::as_str), + Some("add_event") + ); + assert!(committed_event.attributes.contains_key("event_id")); + assert!(committed_event.attributes.contains_key("event_sequence")); + let event_sequence = committed_event + .attributes + .get("event_sequence") + .unwrap() + .parse::() + .unwrap(); let _ = scoped_ticket_backend_operation( State(api.clone()), AxumPath(ScopedWorkspacePath { workspace_id: TEST_WORKSPACE_ID.to_string(), }), - Json(TicketBackendOperation::AddTicketRelation { - id: ticket_id.clone().into(), - relation: ticket::NewTicketRelation { + headers.clone(), + Json(TicketBackendOperation::Show { + id: ticket_ref.id.clone().into(), + }), + ) + .await + .unwrap(); + assert_eq!( + api.store + .get_ticket_notification_cursor( + TEST_WORKSPACE_ID, + &ticket_ref.id, + EMBEDDED_WORKER_RUNTIME_ID, + &source_worker.worker_id, + ) + .unwrap(), + Some(event_sequence) + ); + assert!( + api.store + .list_pending_ticket_notification_deliveries(TEST_WORKSPACE_ID, 10) + .unwrap() + .is_empty(), + "accepted Runtime system input must complete the outbox delivery" + ); + + let Json(stale_report) = scoped_ticket_backend_operation( + State(api.clone()), + AxumPath(ScopedWorkspacePath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + }), + headers.clone(), + Json(TicketBackendOperation::AddEvent { + id: ticket_ref.id.clone().into(), + event: NewTicketEvent::new( + TicketEventKind::ImplementationReport, + "non-assigned Worker report", + ), + }), + ) + .await + .unwrap(); + assert!(matches!(stale_report, TicketBackendHttpResponse::Ok { .. })); + let non_assigned_report = backend.show(ticket_ref.id.clone().into()).unwrap(); + assert!( + !non_assigned_report + .events + .last() + .unwrap() + .attributes + .contains_key("source_assignment_id") + ); + + api.store + .set_current_ticket_worker_assignment( + &TicketWorkerAssignmentRecord { + workspace_id: TEST_WORKSPACE_ID.to_string(), + ticket_id: ticket_ref.id.clone(), + assignment_id: "source-assignment".to_string(), + runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + worker_id: source_worker.worker_id.clone(), + assigned_by: "test-user".to_string(), + assigned_at: TEST_CREATED_AT.to_string(), + }, + Some("notify-assignment"), + "source-assignment-event", + "source-assignment-operation", + true, + ) + .unwrap(); + let _ = scoped_ticket_backend_operation( + State(api.clone()), + AxumPath(ScopedWorkspacePath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + }), + headers.clone(), + Json(TicketBackendOperation::AddEvent { + id: ticket_ref.id.clone().into(), + event: NewTicketEvent::new( + TicketEventKind::ImplementationReport, + "current assignment report", + ), + }), + ) + .await + .unwrap(); + let reported = backend.show(ticket_ref.id.clone().into()).unwrap(); + let report_event = reported.events.last().unwrap(); + assert_eq!( + report_event + .attributes + .get("source_assignment_id") + .map(String::as_str), + Some("source-assignment") + ); + assert_eq!( + report_event + .attributes + .get("source_actor_role") + .map(String::as_str), + Some("coder") + ); + + let unauthorized = scoped_ticket_backend_operation( + State(api.clone()), + AxumPath(ScopedWorkspacePath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + }), + HeaderMap::new(), + Json(TicketBackendOperation::AddEvent { + id: ticket_ref.id.clone().into(), + event: NewTicketEvent::new(TicketEventKind::Comment, "spoofed update"), + }), + ) + .await + .unwrap_err() + .into_response(); + assert_eq!(unauthorized.status(), StatusCode::UNAUTHORIZED); + } + + #[tokio::test] + async fn queued_ticket_mutation_targets_current_orchestrator() { + let dir = tempfile::tempdir().unwrap(); + let api = test_api(dir.path()).await; + let source = api + .runtime + .spawn_worker( + EMBEDDED_WORKER_RUNTIME_ID, + WorkerSpawnRequest { + requested_worker_name: Some("orchestrator-source".to_string()), + intent: WorkerSpawnIntent::TicketRole { + ticket_id: "source-ticket".to_string(), + role: TicketWorkerRole::Coder, + }, + acceptance: WorkerSpawnAcceptanceRequirement::RunAccepted { + expected_segments: 0, + }, + profile: ProfileSelector::Builtin("builtin:coder".to_string()), + ticket_assignment: None, + initial_input: None, + working_directory_request: None, + resolved_working_directory_request: None, + resolved_working_directory: None, + resolved_config_bundle: None, + resolved_workspace_api: None, + }, + ) + .unwrap() + .worker + .unwrap(); + let orchestrator = api + .runtime + .spawn_worker( + EMBEDDED_WORKER_RUNTIME_ID, + WorkerSpawnRequest { + requested_worker_name: Some("workspace-orchestrator".to_string()), + intent: WorkerSpawnIntent::WorkspaceOrchestrator, + acceptance: WorkerSpawnAcceptanceRequirement::RunAccepted { + expected_segments: 0, + }, + profile: ProfileSelector::Builtin("builtin:orchestrator".to_string()), + ticket_assignment: None, + initial_input: None, + working_directory_request: None, + resolved_working_directory_request: None, + resolved_working_directory: None, + resolved_config_bundle: None, + resolved_workspace_api: None, + }, + ) + .unwrap() + .worker + .unwrap(); + api.runtime + .stop_worker( + EMBEDDED_WORKER_RUNTIME_ID, + &orchestrator.worker_id, + WorkerLifecycleRequest { + reason: Some("test pending delivery".to_string()), + ticket_assignment: None, + }, + ) + .unwrap(); + api.store + .upsert_worker_workspace_credential(&WorkerWorkspaceCredentialRecord { + credential_id: "orchestrator-source-credential".to_string(), + token: "orchestrator-source-secret".to_string(), + workspace_id: TEST_WORKSPACE_ID.to_string(), + runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + worker_id: Some(source.worker_id.clone()), + created_at: TEST_CREATED_AT.to_string(), + expires_at: "2099-01-01T00:00:00Z".to_string(), + revoked_at: None, + }) + .unwrap(); + let backend = browser_ticket_backend(&api).unwrap(); + let mut input = ticket::NewTicket::new("Queued notification"); + input.workflow_state = Some(TicketWorkflowState::Queued); + let ticket_ref = backend.create(input).unwrap(); + let mut headers = HeaderMap::new(); + headers.insert( + axum::http::header::AUTHORIZATION, + axum::http::HeaderValue::from_static("Bearer orchestrator-source-secret"), + ); + headers.insert( + "x-yoi-worker-id", + axum::http::HeaderValue::from_str(&source.worker_id).unwrap(), + ); + let _ = scoped_ticket_backend_operation( + State(api.clone()), + AxumPath(ScopedWorkspacePath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + }), + headers, + Json(TicketBackendOperation::AddEvent { + id: ticket_ref.id.clone().into(), + event: NewTicketEvent::new(TicketEventKind::Comment, "queued update"), + }), + ) + .await + .unwrap(); + assert_eq!( + api.store + .count_ticket_notification_deliveries_for_recipient( + TEST_WORKSPACE_ID, + &ticket_ref.id, + EMBEDDED_WORKER_RUNTIME_ID, + &orchestrator.worker_id, + ) + .unwrap(), + 1 + ); + } + + #[tokio::test] + async fn worker_spawn_and_restore_assignment_operations_are_idempotent() { + let dir = tempfile::tempdir().unwrap(); + let api = test_api(dir.path()).await; + let backend = browser_ticket_backend(&api).unwrap(); + let first_ticket = backend + .create(ticket::NewTicket::new("Spawn assignment")) + .unwrap(); + let request = WorkerSpawnRequest { + requested_worker_name: Some("assigned-spawn".to_string()), + intent: WorkerSpawnIntent::TicketRole { + ticket_id: first_ticket.id.clone(), + role: TicketWorkerRole::Coder, + }, + acceptance: WorkerSpawnAcceptanceRequirement::RunAccepted { + expected_segments: 0, + }, + profile: ProfileSelector::Builtin("builtin:coder".to_string()), + ticket_assignment: Some(crate::hosts::WorkerTicketAssignmentRequest { + ticket_id: first_ticket.id.clone(), + operation_id: "spawn-assignment-operation".to_string(), + }), + initial_input: None, + working_directory_request: None, + resolved_working_directory_request: None, + resolved_working_directory: None, + resolved_config_bundle: None, + resolved_workspace_api: None, + }; + let Json(first) = scoped_create_runtime_worker( + State(api.clone()), + AxumPath(ScopedRuntimePath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + }), + Json(request.clone()), + ) + .await + .unwrap(); + let first_worker = first.worker.unwrap(); + let Json(projected) = scoped_get_ticket_worker_assignment( + State(api.clone()), + AxumPath(ScopedRecordPath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + id: first_ticket.id.clone(), + }), + ) + .await + .unwrap(); + assert_eq!( + projected + .worker + .as_ref() + .map(|worker| worker.worker_id.as_str()), + Some(first_worker.worker_id.as_str()) + ); + let Json(retried) = scoped_create_runtime_worker( + State(api.clone()), + AxumPath(ScopedRuntimePath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + }), + Json(request.clone()), + ) + .await + .unwrap(); + assert_eq!(retried.worker.unwrap().worker_id, first_worker.worker_id); + assert_eq!( + api.store + .list_ticket_worker_assignment_events(TEST_WORKSPACE_ID, &first_ticket.id, 10,) + .unwrap() + .len(), + 1 + ); + + let current = api + .store + .get_current_ticket_worker_assignment(TEST_WORKSPACE_ID, &first_ticket.id) + .unwrap() + .unwrap(); + api.store + .clear_current_ticket_worker_assignment( + TEST_WORKSPACE_ID, + &first_ticket.id, + Some(¤t.assignment_id), + "spawn-unassign-operation", + "spawn-unassign-event", + "test-user", + TEST_CREATED_AT, + ) + .unwrap(); + api.runtime + .stop_worker( + EMBEDDED_WORKER_RUNTIME_ID, + &first_worker.worker_id, + WorkerLifecycleRequest { + reason: Some("restore assignment test".to_string()), + ticket_assignment: None, + }, + ) + .unwrap(); + let second_ticket = backend + .create(ticket::NewTicket::new("Restore assignment")) + .unwrap(); + let _ = scoped_restore_runtime_worker( + State(api.clone()), + AxumPath(ScopedRuntimeWorkerPath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + worker_id: first_worker.worker_id.clone(), + }), + Query(RestoreTicketAssignmentQuery { + ticket_id: Some(second_ticket.id.clone()), + assignment_operation_id: Some("restore-assignment-operation".to_string()), + }), + ) + .await + .unwrap(); + let Json(retried_restore) = scoped_restore_runtime_worker( + State(api.clone()), + AxumPath(ScopedRuntimeWorkerPath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + worker_id: first_worker.worker_id.clone(), + }), + Query(RestoreTicketAssignmentQuery { + ticket_id: Some(second_ticket.id.clone()), + assignment_operation_id: Some("restore-assignment-operation".to_string()), + }), + ) + .await + .unwrap(); + assert_eq!(retried_restore.worker_id, first_worker.worker_id); + assert_eq!(retried_restore.result.state, WorkerOperationState::Accepted); + let restored_assignment = api + .store + .get_current_ticket_worker_assignment(TEST_WORKSPACE_ID, &second_ticket.id) + .unwrap() + .unwrap(); + assert_eq!(restored_assignment.worker_id, first_worker.worker_id); + + api.store + .clear_current_ticket_worker_assignment( + TEST_WORKSPACE_ID, + &second_ticket.id, + Some(&restored_assignment.assignment_id), + "restore-clear-operation", + "restore-clear-event", + "test", + TEST_CREATED_AT, + ) + .unwrap(); + let pending_request = WorkerSpawnRequest { + ticket_assignment: Some(crate::hosts::WorkerTicketAssignmentRequest { + ticket_id: second_ticket.id.clone(), + operation_id: "pending-spawn-operation".to_string(), + }), + ..request + }; + let (_, pending_fingerprint) = crate::hosts::worker_spawn_idempotency(&pending_request) + .unwrap() + .unwrap(); + api.store + .reserve_ticket_assignment_operation( + TEST_WORKSPACE_ID, + "pending-spawn-operation", + &second_ticket.id, + EMBEDDED_WORKER_RUNTIME_ID, + None, + &pending_fingerprint, + TEST_CREATED_AT, + ) + .unwrap(); + let spawned_before_backend_failure = api + .runtime + .spawn_worker(EMBEDDED_WORKER_RUNTIME_ID, pending_request.clone()) + .unwrap() + .worker + .unwrap(); + assert!( + api.store + .get_ticket_assignment_operation(TEST_WORKSPACE_ID, "pending-spawn-operation") + .unwrap() + .is_some_and(|operation| operation.worker_id.is_none()) + ); + let worker_count_before_retry = api.runtime.list_workers(100).items.len(); + let Json(reconciled) = scoped_create_runtime_worker( + State(api.clone()), + AxumPath(ScopedRuntimePath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + }), + Json(pending_request), + ) + .await + .unwrap(); + assert_eq!( + reconciled.worker.unwrap().worker_id, + spawned_before_backend_failure.worker_id + ); + assert_eq!( + api.runtime.list_workers(100).items.len(), + worker_count_before_retry, + "retrying a reserved lifecycle operation must not spawn another Worker" + ); + assert!( + api.store + .get_ticket_assignment_operation(TEST_WORKSPACE_ID, "pending-spawn-operation") + .unwrap() + .and_then(|operation| operation.assignment_id) + .is_some() + ); + } + + #[tokio::test] + async fn ticket_browser_endpoints_mutate_typed_backend_and_return_thread() { + let dir = tempfile::tempdir().unwrap(); + let api = test_api(dir.path()).await; + let ticket_ref = browser_ticket_backend(&api) + .unwrap() + .create(ticket::NewTicket::new("Browser Ticket API")) + .unwrap(); + let ticket_id = ticket_ref.id; + let path = || ScopedRecordPath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + id: ticket_id.clone(), + }; + let related_ticket_id = browser_ticket_backend(&api) + .unwrap() + .create(ticket::NewTicket::new("Related Browser Ticket")) + .unwrap() + .id; + browser_ticket_backend(&api) + .unwrap() + .add_ticket_relation( + ticket_id.clone().into(), + ticket::NewTicketRelation { kind: ticket::TicketRelationKind::Related, target: related_ticket_id.clone(), note: Some("Browser relation".to_string()), author: Some("browser-user".to_string()), }, - }), - ) - .await - .unwrap(); + ) + .unwrap(); let Json(edited) = scoped_edit_ticket_item( State(api.clone()), @@ -8134,24 +9722,10 @@ mod tests { .unwrap(); let api = test_api(dir.path()).await; - let Json(response) = scoped_ticket_backend_operation( - State(api.clone()), - AxumPath(ScopedWorkspacePath { - workspace_id: TEST_WORKSPACE_ID.to_string(), - }), - Json(TicketBackendOperation::Create { - input: ticket::NewTicket::new("Endpoint configured root"), - }), - ) - .await - .unwrap_or_else(|error| panic!("ticket backend operation failed: {}", error.error)); - - let ticket_ref = match response { - TicketBackendHttpResponse::Ok { - result: ticket::TicketBackendOperationResult::TicketRef(ticket_ref), - } => ticket_ref, - other => panic!("unexpected ticket backend response: {other:?}"), - }; + let ticket_ref = browser_ticket_backend(&api) + .unwrap() + .create(ticket::NewTicket::new("Endpoint configured root")) + .unwrap(); assert!(api.config.database_path.is_file()); assert!( !dir.path() @@ -8215,9 +9789,10 @@ mod tests { } async fn test_api(workspace_root: impl Into) -> WorkspaceApi { - let store = SqliteWorkspaceStore::in_memory().unwrap(); + let config = test_server_config(workspace_root); + let store = SqliteWorkspaceStore::open(config.database_path.clone()).unwrap(); WorkspaceApi::new_with_execution_backend( - test_server_config(workspace_root), + config, Arc::new(store), Arc::new(DeterministicExecutionBackend::default()), ) @@ -8610,6 +10185,8 @@ mod tests { fn runtime_create_request() -> worker_runtime::catalog::CreateWorkerRequest { let bundle = runtime_test_bundle(); worker_runtime::catalog::CreateWorkerRequest { + idempotency_key: None, + idempotency_fingerprint: None, profile: worker_runtime::catalog::ProfileSelector::Builtin( "builtin:companion".to_string(), ), @@ -9731,11 +11308,13 @@ mod tests { expected_segments: 0, }, profile: ProfileSelector::Builtin("builtin:coder".to_string()), + ticket_assignment: None, initial_input: None, working_directory_request: None, resolved_working_directory_request: None, resolved_working_directory: None, resolved_config_bundle: None, + resolved_workspace_api: None, }, ) .expect("spawn worker"); diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index 39513eb8..459c9123 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -92,6 +92,31 @@ const MIGRATIONS: &[Migration] = &[ name: "separate workdir creation evidence from current revision observation", apply: add_workdir_revision_observations, }, + Migration { + version: 16, + name: "ticket worker current assignment authority", + apply: create_ticket_worker_assignment_tables, + }, + Migration { + version: 17, + name: "worker workspace credentials and Ticket notification outbox", + apply: create_ticket_notification_tables, + }, + Migration { + version: 18, + name: "bidirectional idempotent Ticket Worker assignments", + apply: strengthen_ticket_worker_assignments, + }, + Migration { + version: 19, + name: "atomic Ticket notification identity credentials and cursors", + apply: strengthen_ticket_notifications, + }, + Migration { + version: 20, + name: "reconcile Workdir revision and crash safe Worker lifecycle reservations", + apply: strengthen_ticket_assignment_lifecycle_reservations, + }, ]; struct Migration { @@ -237,6 +262,72 @@ pub struct WorkerRegistryRecord { pub updated_at: String, } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct TicketWorkerAssignmentRecord { + pub workspace_id: String, + pub ticket_id: String, + pub assignment_id: String, + pub runtime_id: String, + pub worker_id: String, + pub assigned_by: String, + pub assigned_at: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct TicketWorkerAssignmentEventRecord { + pub workspace_id: String, + pub ticket_id: String, + pub event_id: String, + pub action: String, + pub assignment_id: Option, + pub previous_assignment_id: Option, + pub actor: String, + pub created_at: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct TicketWorkerAssignmentUpdate { + pub current: TicketWorkerAssignmentRecord, + pub previous: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct WorkerWorkspaceCredentialRecord { + pub credential_id: String, + pub token: String, + pub workspace_id: String, + pub runtime_id: String, + pub worker_id: Option, + pub created_at: String, + pub expires_at: String, + pub revoked_at: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct TicketNotificationRecipient { + pub runtime_id: String, + pub worker_id: String, + pub recipient_kind: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct TicketNotificationDeliveryRecord { + pub notification_id: String, + pub workspace_id: String, + pub ticket_id: String, + pub event_sequence: i64, + pub event_kind: String, + pub source_operation_kind: String, + pub source_actor_role: String, + pub source_assignment_id: Option, + pub source_runtime_id: String, + pub source_worker_id: String, + pub recipient_runtime_id: String, + pub recipient_worker_id: String, + pub recipient_kind: String, + pub attempts: i64, +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct WorkdirRegistryRecord { pub workspace_id: String, @@ -490,6 +581,146 @@ pub trait ControlPlaneStore: Send + Sync { runtime_worker_id: u64, ) -> Result; + fn get_ticket_assignment_operation( + &self, + workspace_id: &str, + operation_id: &str, + ) -> Result>; + fn reserve_ticket_assignment_operation( + &self, + workspace_id: &str, + operation_id: &str, + ticket_id: &str, + runtime_id: &str, + worker_id: Option<&str>, + request_fingerprint: &str, + created_at: &str, + ) -> Result<()>; + fn bind_ticket_assignment_operation_worker( + &self, + workspace_id: &str, + operation_id: &str, + worker_id: &str, + ) -> Result<()>; + fn get_current_ticket_worker_assignment( + &self, + workspace_id: &str, + ticket_id: &str, + ) -> Result>; + fn set_current_ticket_worker_assignment( + &self, + record: &TicketWorkerAssignmentRecord, + expected_assignment_id: Option<&str>, + event_id: &str, + operation_id: &str, + allow_reassign: bool, + ) -> Result; + fn clear_current_ticket_worker_assignment( + &self, + workspace_id: &str, + ticket_id: &str, + expected_assignment_id: Option<&str>, + operation_id: &str, + event_id: &str, + actor: &str, + created_at: &str, + ) -> Result>; + fn list_ticket_worker_assignment_events( + &self, + workspace_id: &str, + ticket_id: &str, + limit: usize, + ) -> Result>; + + fn upsert_worker_workspace_credential( + &self, + record: &WorkerWorkspaceCredentialRecord, + ) -> Result<()>; + fn authenticate_worker_workspace_credential( + &self, + token: &str, + workspace_id: &str, + worker_id: &str, + ) -> Result>; + fn refresh_worker_workspace_credential( + &self, + token: &str, + workspace_id: &str, + worker_id: &str, + new_token: &str, + new_expires_at: &str, + ) -> Result>; + fn revoke_worker_workspace_credentials( + &self, + workspace_id: &str, + runtime_id: &str, + worker_id: &str, + revoked_at: &str, + ) -> Result<()>; + fn enqueue_ticket_notification( + &self, + notification_id: &str, + workspace_id: &str, + ticket_id: &str, + event_sequence: i64, + source_runtime_id: &str, + source_worker_id: &str, + previous_state: &str, + current_state: &str, + created_at: &str, + recipients: &[TicketNotificationRecipient], + ) -> Result<()>; + fn list_pending_ticket_notification_deliveries( + &self, + workspace_id: &str, + limit: usize, + ) -> Result>; + fn count_ticket_notification_deliveries_for_recipient( + &self, + workspace_id: &str, + ticket_id: &str, + runtime_id: &str, + worker_id: &str, + ) -> Result; + fn mark_ticket_notification_delivered( + &self, + notification_id: &str, + recipient_runtime_id: &str, + recipient_worker_id: &str, + delivered_at: &str, + ) -> Result<()>; + fn mark_ticket_notification_failed( + &self, + notification_id: &str, + recipient_runtime_id: &str, + recipient_worker_id: &str, + error: &str, + ) -> Result<()>; + fn reroute_ticket_notification_delivery( + &self, + notification_id: &str, + old_runtime_id: &str, + old_worker_id: &str, + new_runtime_id: &str, + new_worker_id: &str, + ) -> Result<()>; + fn upsert_ticket_notification_cursor( + &self, + workspace_id: &str, + ticket_id: &str, + runtime_id: &str, + worker_id: &str, + event_index: i64, + updated_at: &str, + ) -> Result<()>; + fn get_ticket_notification_cursor( + &self, + workspace_id: &str, + ticket_id: &str, + runtime_id: &str, + worker_id: &str, + ) -> Result>; + fn upsert_workdir_registry(&self, record: &WorkdirRegistryRecord) -> Result<()>; fn get_workdir_registry( &self, @@ -1573,6 +1804,790 @@ impl ControlPlaneStore for SqliteWorkspaceStore { }) } + fn get_ticket_assignment_operation( + &self, + workspace_id: &str, + operation_id: &str, + ) -> Result> { + self.with_conn(|conn| read_assignment_operation(conn, workspace_id, operation_id)) + } + + fn reserve_ticket_assignment_operation( + &self, + workspace_id: &str, + operation_id: &str, + ticket_id: &str, + runtime_id: &str, + worker_id: Option<&str>, + request_fingerprint: &str, + created_at: &str, + ) -> Result<()> { + self.with_conn(|conn| { + let inserted = conn.execute( + r#"INSERT OR IGNORE INTO ticket_assignment_operations ( + workspace_id, operation_id, action, ticket_id, runtime_id, worker_id, + assignment_id, expected_assignment_id, created_at, request_fingerprint + ) VALUES (?1, ?2, 'assign', ?3, ?4, ?5, NULL, NULL, ?6, ?7)"#, + params![ + workspace_id, + operation_id, + ticket_id, + runtime_id, + worker_id, + created_at, + request_fingerprint, + ], + )?; + if inserted > 0 { + return Ok(()); + } + let existing = read_assignment_operation(conn, workspace_id, operation_id)? + .ok_or_else(|| { + Error::TicketAssignmentConflict(format!( + "assignment operation {operation_id} could not be reserved" + )) + })?; + if existing.action == "assign" + && existing.ticket_id == ticket_id + && existing.runtime_id.as_deref() == Some(runtime_id) + && (worker_id.is_none() || existing.worker_id.as_deref() == worker_id) + && existing.expected_assignment_id.is_none() + && existing.request_fingerprint.as_deref() == Some(request_fingerprint) + { + Ok(()) + } else { + Err(Error::TicketAssignmentConflict(format!( + "assignment operation {operation_id} was already used with different input" + ))) + } + }) + } + + fn bind_ticket_assignment_operation_worker( + &self, + workspace_id: &str, + operation_id: &str, + worker_id: &str, + ) -> Result<()> { + self.with_conn(|conn| { + let updated = conn.execute( + r#"UPDATE ticket_assignment_operations + SET worker_id = ?3 + WHERE workspace_id = ?1 AND operation_id = ?2 + AND assignment_id IS NULL AND (worker_id IS NULL OR worker_id = ?3)"#, + params![workspace_id, operation_id, worker_id], + )?; + if updated == 1 { + return Ok(()); + } + Err(Error::TicketAssignmentConflict(format!( + "assignment operation {operation_id} cannot bind Worker {worker_id}" + ))) + }) + } + + fn get_current_ticket_worker_assignment( + &self, + workspace_id: &str, + ticket_id: &str, + ) -> Result> { + self.with_conn(|conn| { + conn.query_row( + current_ticket_worker_assignment_select_sql().as_str(), + params![workspace_id, ticket_id], + read_ticket_worker_assignment_record, + ) + .optional() + .map_err(Error::from) + }) + } + + fn set_current_ticket_worker_assignment( + &self, + record: &TicketWorkerAssignmentRecord, + expected_assignment_id: Option<&str>, + event_id: &str, + operation_id: &str, + allow_reassign: bool, + ) -> Result { + self.with_conn(|conn| { + let tx = conn.unchecked_transaction()?; + let mut reserved_operation = false; + if let Some(existing) = + read_assignment_operation(&tx, &record.workspace_id, operation_id)? + { + if existing.action != if allow_reassign { "reassign" } else { "assign" } + || existing.ticket_id != record.ticket_id + || existing.runtime_id.as_deref() != Some(record.runtime_id.as_str()) + || existing.worker_id.as_deref() != Some(record.worker_id.as_str()) + || existing.expected_assignment_id.as_deref() != expected_assignment_id + { + return Err(Error::TicketAssignmentConflict(format!( + "assignment operation {operation_id} was already used with different input" + ))); + } + if let Some(assignment_id) = existing.assignment_id { + let current = tx.query_row( + r#"SELECT workspace_id, ticket_id, assignment_id, runtime_id, worker_id, + assigned_by, assigned_at + FROM ticket_worker_assignments + WHERE workspace_id = ?1 AND ticket_id = ?2 AND assignment_id = ?3"#, + params![record.workspace_id, record.ticket_id, assignment_id], + read_ticket_worker_assignment_record, + )?; + let previous = if existing.action == "reassign" { + existing + .expected_assignment_id + .as_deref() + .map(|previous_assignment_id| { + tx.query_row( + r#"SELECT workspace_id, ticket_id, assignment_id, runtime_id, worker_id, + assigned_by, assigned_at + FROM ticket_worker_assignments + WHERE workspace_id = ?1 AND ticket_id = ?2 AND assignment_id = ?3"#, + params![ + record.workspace_id, + record.ticket_id, + previous_assignment_id, + ], + read_ticket_worker_assignment_record, + ) + }) + .transpose()? + } else { + None + }; + tx.commit()?; + return Ok(TicketWorkerAssignmentUpdate { current, previous }); + } + reserved_operation = true; + } + let previous = tx + .query_row( + current_ticket_worker_assignment_select_sql().as_str(), + params![record.workspace_id, record.ticket_id], + read_ticket_worker_assignment_record, + ) + .optional()?; + if previous.is_some() && !allow_reassign { + return Err(Error::TicketAssignmentConflict(format!( + "Ticket {} is already assigned; use the explicit reassign operation", + record.ticket_id + ))); + } + if allow_reassign { + let expected_assignment_id = expected_assignment_id.ok_or_else(|| { + Error::TicketAssignmentConflict( + "reassign requires expected_assignment_id".to_string(), + ) + })?; + require_expected_ticket_assignment( + record.ticket_id.as_str(), + previous.as_ref(), + Some(expected_assignment_id), + )?; + } else if expected_assignment_id.is_some() { + return Err(Error::TicketAssignmentConflict( + "assign does not accept expected_assignment_id".to_string(), + )); + } + tx.execute( + r#"INSERT INTO ticket_worker_assignments ( + workspace_id, ticket_id, assignment_id, runtime_id, worker_id, + assigned_by, assigned_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)"#, + params![ + record.workspace_id, + record.ticket_id, + record.assignment_id, + record.runtime_id, + record.worker_id, + record.assigned_by, + record.assigned_at, + ], + )?; + let current_write = if allow_reassign { + tx.execute( + r#"UPDATE ticket_current_worker_assignments + SET assignment_id = ?3, runtime_id = ?4, worker_id = ?5, updated_at = ?6 + WHERE workspace_id = ?1 AND ticket_id = ?2"#, + params![ + record.workspace_id, + record.ticket_id, + record.assignment_id, + record.runtime_id, + record.worker_id, + record.assigned_at, + ], + ) + } else { + tx.execute( + r#"INSERT INTO ticket_current_worker_assignments ( + workspace_id, ticket_id, assignment_id, runtime_id, worker_id, updated_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6)"#, + params![ + record.workspace_id, + record.ticket_id, + record.assignment_id, + record.runtime_id, + record.worker_id, + record.assigned_at, + ], + ) + }; + if let Err(error) = current_write { + return Err(map_assignment_constraint( + error, + &record.ticket_id, + &record.worker_id, + )); + } + tx.execute( + r#"INSERT INTO ticket_worker_assignment_events ( + workspace_id, ticket_id, event_id, action, assignment_id, + previous_assignment_id, actor, created_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)"#, + params![ + record.workspace_id, + record.ticket_id, + event_id, + if previous.is_some() { + "reassigned" + } else { + "assigned" + }, + record.assignment_id, + previous + .as_ref() + .map(|assignment| assignment.assignment_id.as_str()), + record.assigned_by, + record.assigned_at, + ], + )?; + if reserved_operation { + let updated = tx.execute( + r#"UPDATE ticket_assignment_operations + SET assignment_id = ?3 + WHERE workspace_id = ?1 AND operation_id = ?2 AND assignment_id IS NULL"#, + params![record.workspace_id, operation_id, record.assignment_id], + )?; + if updated != 1 { + return Err(Error::TicketAssignmentConflict(format!( + "assignment operation {operation_id} reservation was not current" + ))); + } + } else { + tx.execute( + r#"INSERT INTO ticket_assignment_operations ( + workspace_id, operation_id, action, ticket_id, runtime_id, worker_id, + assignment_id, expected_assignment_id, created_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)"#, + params![ + record.workspace_id, + operation_id, + if allow_reassign { "reassign" } else { "assign" }, + record.ticket_id, + record.runtime_id, + record.worker_id, + record.assignment_id, + expected_assignment_id, + record.assigned_at, + ], + )?; + } + tx.commit()?; + Ok(TicketWorkerAssignmentUpdate { + current: record.clone(), + previous, + }) + }) + } + + fn clear_current_ticket_worker_assignment( + &self, + workspace_id: &str, + ticket_id: &str, + expected_assignment_id: Option<&str>, + operation_id: &str, + event_id: &str, + actor: &str, + created_at: &str, + ) -> Result> { + self.with_conn(|conn| { + let tx = conn.unchecked_transaction()?; + if let Some(existing) = read_assignment_operation(&tx, workspace_id, operation_id)? { + if existing.action != "unassign" + || existing.ticket_id != ticket_id + || existing.expected_assignment_id.as_deref() != expected_assignment_id + { + return Err(Error::TicketAssignmentConflict(format!( + "assignment operation {operation_id} was already used with different input" + ))); + } + let assignment = existing + .assignment_id + .map(|assignment_id| { + tx.query_row( + r#"SELECT workspace_id, ticket_id, assignment_id, runtime_id, worker_id, + assigned_by, assigned_at + FROM ticket_worker_assignments + WHERE workspace_id = ?1 AND ticket_id = ?2 AND assignment_id = ?3"#, + params![workspace_id, ticket_id, assignment_id], + read_ticket_worker_assignment_record, + ) + }) + .transpose()?; + tx.commit()?; + return Ok(assignment); + } + let previous = tx + .query_row( + current_ticket_worker_assignment_select_sql().as_str(), + params![workspace_id, ticket_id], + read_ticket_worker_assignment_record, + ) + .optional()?; + require_expected_ticket_assignment(ticket_id, previous.as_ref(), expected_assignment_id)?; + let Some(previous) = previous else { + tx.commit()?; + return Ok(None); + }; + tx.execute( + "DELETE FROM ticket_current_worker_assignments WHERE workspace_id = ?1 AND ticket_id = ?2", + params![workspace_id, ticket_id], + )?; + tx.execute( + r#"INSERT INTO ticket_worker_assignment_events ( + workspace_id, ticket_id, event_id, action, assignment_id, + previous_assignment_id, actor, created_at + ) VALUES (?1, ?2, ?3, 'unassigned', NULL, ?4, ?5, ?6)"#, + params![ + workspace_id, + ticket_id, + event_id, + previous.assignment_id, + actor, + created_at, + ], + )?; + tx.execute( + r#"INSERT INTO ticket_assignment_operations ( + workspace_id, operation_id, action, ticket_id, runtime_id, worker_id, + assignment_id, expected_assignment_id, created_at + ) VALUES (?1, ?2, 'unassign', ?3, ?4, ?5, ?6, ?7, ?8)"#, + params![ + workspace_id, + operation_id, + ticket_id, + previous.runtime_id, + previous.worker_id, + previous.assignment_id, + expected_assignment_id, + created_at, + ], + )?; + tx.commit()?; + Ok(Some(previous)) + }) + } + + fn list_ticket_worker_assignment_events( + &self, + workspace_id: &str, + ticket_id: &str, + limit: usize, + ) -> Result> { + self.with_conn(|conn| { + let mut stmt = conn.prepare( + r#"SELECT workspace_id, ticket_id, event_id, action, assignment_id, + previous_assignment_id, actor, created_at + FROM ticket_worker_assignment_events + WHERE workspace_id = ?1 AND ticket_id = ?2 + ORDER BY created_at DESC, event_id DESC + LIMIT ?3"#, + )?; + let rows = stmt.query_map( + params![workspace_id, ticket_id, limit as i64], + read_ticket_worker_assignment_event_record, + )?; + rows.collect::, _>>() + .map_err(Error::from) + }) + } + + fn upsert_worker_workspace_credential( + &self, + record: &WorkerWorkspaceCredentialRecord, + ) -> Result<()> { + self.with_conn(|conn| { + conn.execute( + r#"INSERT INTO worker_workspace_credentials ( + credential_id, token, workspace_id, runtime_id, worker_id, created_at, + expires_at, revoked_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) + ON CONFLICT(credential_id) DO UPDATE SET + token = excluded.token, + workspace_id = excluded.workspace_id, + runtime_id = excluded.runtime_id, + worker_id = excluded.worker_id, + created_at = excluded.created_at, + expires_at = excluded.expires_at, + revoked_at = excluded.revoked_at"#, + params![ + record.credential_id, + record.token, + record.workspace_id, + record.runtime_id, + record.worker_id, + record.created_at, + record.expires_at, + record.revoked_at, + ], + )?; + Ok(()) + }) + } + + fn authenticate_worker_workspace_credential( + &self, + token: &str, + workspace_id: &str, + worker_id: &str, + ) -> Result> { + self.with_conn(|conn| { + let tx = conn.unchecked_transaction()?; + let record = tx + .query_row( + r#"SELECT credential_id, token, workspace_id, runtime_id, worker_id, created_at, + expires_at, revoked_at + FROM worker_workspace_credentials + WHERE token = ?1 AND workspace_id = ?2 + AND revoked_at IS NULL AND datetime(expires_at) > datetime('now')"#, + params![token, workspace_id], + |row| { + Ok(WorkerWorkspaceCredentialRecord { + credential_id: row.get(0)?, + token: row.get(1)?, + workspace_id: row.get(2)?, + runtime_id: row.get(3)?, + worker_id: row.get(4)?, + created_at: row.get(5)?, + expires_at: row.get(6)?, + revoked_at: row.get(7)?, + }) + }, + ) + .optional()?; + let Some(mut record) = record else { + tx.commit()?; + return Ok(None); + }; + if record.worker_id.as_deref().is_some_and(|bound| bound != worker_id) { + tx.commit()?; + return Ok(None); + } + if record.worker_id.is_none() { + tx.execute( + "UPDATE worker_workspace_credentials SET worker_id = ?1 WHERE credential_id = ?2 AND worker_id IS NULL", + params![worker_id, record.credential_id], + )?; + record.worker_id = Some(worker_id.to_string()); + } + tx.commit()?; + Ok(Some(record)) + }) + } + + fn refresh_worker_workspace_credential( + &self, + token: &str, + workspace_id: &str, + worker_id: &str, + new_token: &str, + new_expires_at: &str, + ) -> Result> { + self.with_conn(|conn| { + let tx = conn.unchecked_transaction()?; + let record = tx + .query_row( + r#"SELECT credential_id, runtime_id, created_at FROM worker_workspace_credentials + WHERE token = ?1 AND workspace_id = ?2 AND worker_id = ?3 AND revoked_at IS NULL"#, + params![token, workspace_id, worker_id], + |row| Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?, row.get::<_, String>(2)?)), + ) + .optional()?; + let Some((credential_id, runtime_id, created_at)) = record else { + tx.commit()?; + return Ok(None); + }; + tx.execute( + "UPDATE worker_workspace_credentials SET token = ?1, expires_at = ?2 WHERE credential_id = ?3", + params![new_token, new_expires_at, credential_id], + )?; + tx.commit()?; + Ok(Some(WorkerWorkspaceCredentialRecord { + credential_id, + token: new_token.to_string(), + workspace_id: workspace_id.to_string(), + runtime_id, + worker_id: Some(worker_id.to_string()), + created_at, + expires_at: new_expires_at.to_string(), + revoked_at: None, + })) + }) + } + + fn revoke_worker_workspace_credentials( + &self, + workspace_id: &str, + runtime_id: &str, + worker_id: &str, + revoked_at: &str, + ) -> Result<()> { + self.with_conn(|conn| { + conn.execute( + r#"UPDATE worker_workspace_credentials SET revoked_at = ?4 + WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3 AND revoked_at IS NULL"#, + params![workspace_id, runtime_id, worker_id, revoked_at], + )?; + Ok(()) + }) + } + + fn enqueue_ticket_notification( + &self, + notification_id: &str, + workspace_id: &str, + ticket_id: &str, + event_sequence: i64, + source_runtime_id: &str, + source_worker_id: &str, + previous_state: &str, + current_state: &str, + created_at: &str, + recipients: &[TicketNotificationRecipient], + ) -> Result<()> { + self.with_conn(|conn| { + let tx = conn.unchecked_transaction()?; + tx.execute( + r#"INSERT INTO ticket_notification_outbox ( + notification_id, workspace_id, ticket_id, event_sequence, + source_runtime_id, source_worker_id, previous_state, current_state, created_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)"#, + params![ + notification_id, + workspace_id, + ticket_id, + event_sequence, + source_runtime_id, + source_worker_id, + previous_state, + current_state, + created_at, + ], + )?; + for recipient in recipients { + tx.execute( + r#"INSERT OR IGNORE INTO ticket_notification_deliveries ( + notification_id, recipient_runtime_id, recipient_worker_id, + recipient_kind, attempts + ) VALUES (?1, ?2, ?3, ?4, 0)"#, + params![ + notification_id, + recipient.runtime_id, + recipient.worker_id, + recipient.recipient_kind, + ], + )?; + } + tx.commit()?; + Ok(()) + }) + } + + fn list_pending_ticket_notification_deliveries( + &self, + workspace_id: &str, + limit: usize, + ) -> Result> { + self.with_conn(|conn| { + let mut stmt = conn.prepare( + r#"SELECT o.notification_id, o.workspace_id, o.ticket_id, o.event_sequence, + o.event_kind, o.source_operation_kind, o.source_actor_role, + o.source_assignment_id, o.source_runtime_id, o.source_worker_id, + d.recipient_runtime_id, d.recipient_worker_id, d.recipient_kind, d.attempts + FROM ticket_notification_deliveries AS d + JOIN ticket_notification_outbox AS o ON o.notification_id = d.notification_id + WHERE o.workspace_id = ?1 AND d.delivered_at IS NULL + ORDER BY o.created_at ASC, o.notification_id ASC + LIMIT ?2"#, + )?; + let rows = stmt.query_map(params![workspace_id, limit as i64], |row| { + Ok(TicketNotificationDeliveryRecord { + notification_id: row.get(0)?, + workspace_id: row.get(1)?, + ticket_id: row.get(2)?, + event_sequence: row.get(3)?, + event_kind: row.get(4)?, + source_operation_kind: row.get(5)?, + source_actor_role: row.get(6)?, + source_assignment_id: row.get(7)?, + source_runtime_id: row.get(8)?, + source_worker_id: row.get(9)?, + recipient_runtime_id: row.get(10)?, + recipient_worker_id: row.get(11)?, + recipient_kind: row.get(12)?, + attempts: row.get(13)?, + }) + })?; + rows.collect::, _>>() + .map_err(Error::from) + }) + } + + fn count_ticket_notification_deliveries_for_recipient( + &self, + workspace_id: &str, + ticket_id: &str, + runtime_id: &str, + worker_id: &str, + ) -> Result { + self.with_conn(|conn| { + let count = conn.query_row( + r#"SELECT COUNT(*) + FROM ticket_notification_deliveries AS d + JOIN ticket_notification_outbox AS o ON o.notification_id = d.notification_id + WHERE o.workspace_id = ?1 AND o.ticket_id = ?2 + AND d.recipient_runtime_id = ?3 AND d.recipient_worker_id = ?4"#, + params![workspace_id, ticket_id, runtime_id, worker_id], + |row| row.get::<_, i64>(0), + )?; + Ok(count as usize) + }) + } + + fn mark_ticket_notification_delivered( + &self, + notification_id: &str, + recipient_runtime_id: &str, + recipient_worker_id: &str, + delivered_at: &str, + ) -> Result<()> { + self.with_conn(|conn| { + conn.execute( + r#"UPDATE ticket_notification_deliveries + SET delivered_at = ?4, last_error = NULL, attempts = attempts + 1 + WHERE notification_id = ?1 AND recipient_runtime_id = ?2 AND recipient_worker_id = ?3"#, + params![notification_id, recipient_runtime_id, recipient_worker_id, delivered_at], + )?; + Ok(()) + }) + } + + fn mark_ticket_notification_failed( + &self, + notification_id: &str, + recipient_runtime_id: &str, + recipient_worker_id: &str, + error: &str, + ) -> Result<()> { + self.with_conn(|conn| { + conn.execute( + r#"UPDATE ticket_notification_deliveries + SET last_error = ?4, attempts = attempts + 1 + WHERE notification_id = ?1 AND recipient_runtime_id = ?2 AND recipient_worker_id = ?3"#, + params![notification_id, recipient_runtime_id, recipient_worker_id, error], + )?; + Ok(()) + }) + } + + fn reroute_ticket_notification_delivery( + &self, + notification_id: &str, + old_runtime_id: &str, + old_worker_id: &str, + new_runtime_id: &str, + new_worker_id: &str, + ) -> Result<()> { + self.with_conn(|conn| { + let tx = conn.unchecked_transaction()?; + let recipient_kind: Option = tx + .query_row( + r#"SELECT recipient_kind FROM ticket_notification_deliveries + WHERE notification_id = ?1 AND recipient_runtime_id = ?2 AND recipient_worker_id = ?3"#, + params![notification_id, old_runtime_id, old_worker_id], + |row| row.get(0), + ) + .optional()?; + if let Some(recipient_kind) = recipient_kind { + tx.execute( + r#"INSERT OR IGNORE INTO ticket_notification_deliveries ( + notification_id, recipient_runtime_id, recipient_worker_id, recipient_kind, attempts + ) VALUES (?1, ?2, ?3, ?4, 0)"#, + params![notification_id, new_runtime_id, new_worker_id, recipient_kind], + )?; + tx.execute( + r#"DELETE FROM ticket_notification_deliveries + WHERE notification_id = ?1 AND recipient_runtime_id = ?2 AND recipient_worker_id = ?3"#, + params![notification_id, old_runtime_id, old_worker_id], + )?; + } + tx.commit()?; + Ok(()) + }) + } + + fn upsert_ticket_notification_cursor( + &self, + workspace_id: &str, + ticket_id: &str, + runtime_id: &str, + worker_id: &str, + event_index: i64, + updated_at: &str, + ) -> Result<()> { + self.with_conn(|conn| { + conn.execute( + r#"INSERT INTO ticket_notification_cursors ( + workspace_id, ticket_id, runtime_id, worker_id, last_event_index, updated_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6) + ON CONFLICT(workspace_id, ticket_id, runtime_id, worker_id) DO UPDATE SET + last_event_index = MAX(last_event_index, excluded.last_event_index), + updated_at = excluded.updated_at"#, + params![ + workspace_id, + ticket_id, + runtime_id, + worker_id, + event_index, + updated_at + ], + )?; + Ok(()) + }) + } + + fn get_ticket_notification_cursor( + &self, + workspace_id: &str, + ticket_id: &str, + runtime_id: &str, + worker_id: &str, + ) -> Result> { + self.with_conn(|conn| { + conn.query_row( + r#"SELECT last_event_index FROM ticket_notification_cursors + WHERE workspace_id = ?1 AND ticket_id = ?2 AND runtime_id = ?3 AND worker_id = ?4"#, + params![workspace_id, ticket_id, runtime_id, worker_id], + |row| row.get(0), + ) + .optional() + .map_err(Error::from) + }) + } + fn upsert_workdir_registry(&self, record: &WorkdirRegistryRecord) -> Result<()> { self.with_conn(|conn| { conn.execute( @@ -2008,6 +3023,111 @@ fn read_worker_registry_record(row: &rusqlite::Row<'_>) -> rusqlite::Result String { + "SELECT a.workspace_id, a.ticket_id, a.assignment_id, a.runtime_id, a.worker_id, \ + a.assigned_by, a.assigned_at \ + FROM ticket_current_worker_assignments AS current \ + JOIN ticket_worker_assignments AS a \ + ON a.workspace_id = current.workspace_id \ + AND a.ticket_id = current.ticket_id \ + AND a.assignment_id = current.assignment_id \ + WHERE current.workspace_id = ?1 AND current.ticket_id = ?2" + .to_owned() +} + +fn read_ticket_worker_assignment_record( + row: &rusqlite::Row<'_>, +) -> rusqlite::Result { + Ok(TicketWorkerAssignmentRecord { + workspace_id: row.get(0)?, + ticket_id: row.get(1)?, + assignment_id: row.get(2)?, + runtime_id: row.get(3)?, + worker_id: row.get(4)?, + assigned_by: row.get(5)?, + assigned_at: row.get(6)?, + }) +} + +fn read_ticket_worker_assignment_event_record( + row: &rusqlite::Row<'_>, +) -> rusqlite::Result { + Ok(TicketWorkerAssignmentEventRecord { + workspace_id: row.get(0)?, + ticket_id: row.get(1)?, + event_id: row.get(2)?, + action: row.get(3)?, + assignment_id: row.get(4)?, + previous_assignment_id: row.get(5)?, + actor: row.get(6)?, + created_at: row.get(7)?, + }) +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TicketAssignmentOperationRecord { + pub action: String, + pub ticket_id: String, + pub runtime_id: Option, + pub worker_id: Option, + pub assignment_id: Option, + pub expected_assignment_id: Option, + pub request_fingerprint: Option, +} + +fn read_assignment_operation( + conn: &Connection, + workspace_id: &str, + operation_id: &str, +) -> Result> { + conn.query_row( + r#"SELECT action, ticket_id, runtime_id, worker_id, assignment_id, expected_assignment_id, + request_fingerprint + FROM ticket_assignment_operations + WHERE workspace_id = ?1 AND operation_id = ?2"#, + params![workspace_id, operation_id], + |row| { + Ok(TicketAssignmentOperationRecord { + action: row.get(0)?, + ticket_id: row.get(1)?, + runtime_id: row.get(2)?, + worker_id: row.get(3)?, + assignment_id: row.get(4)?, + expected_assignment_id: row.get(5)?, + request_fingerprint: row.get(6)?, + }) + }, + ) + .optional() + .map_err(Error::from) +} + +fn map_assignment_constraint(error: rusqlite::Error, ticket_id: &str, worker_id: &str) -> Error { + if matches!(error, rusqlite::Error::SqliteFailure(_, _)) { + Error::TicketAssignmentConflict(format!( + "Ticket {ticket_id} or Worker {worker_id} already has a current assignment" + )) + } else { + Error::Sqlite(error) + } +} + +fn require_expected_ticket_assignment( + ticket_id: &str, + current: Option<&TicketWorkerAssignmentRecord>, + expected_assignment_id: Option<&str>, +) -> Result<()> { + let Some(expected_assignment_id) = expected_assignment_id else { + return Ok(()); + }; + if current.map(|assignment| assignment.assignment_id.as_str()) == Some(expected_assignment_id) { + return Ok(()); + } + Err(Error::TicketAssignmentConflict(format!( + "Ticket {ticket_id} is no longer assigned to {expected_assignment_id}" + ))) +} + fn workdir_registry_select_sql(where_clause: &str) -> String { format!( "SELECT workspace_id, workdir_id, runtime_id, repository_id, \ @@ -2191,17 +3311,205 @@ DROP TABLE IF EXISTS tickets; } fn add_workdir_revision_observations(conn: &Connection) -> Result<()> { + if column_exists(conn, "workdir_registry", "selector")? + && !column_exists(conn, "workdir_registry", "creation_selector")? + { + conn.execute_batch( + "ALTER TABLE workdir_registry RENAME COLUMN selector TO creation_selector;", + )?; + } + if column_exists(conn, "workdir_registry", "resolved_commit")? + && !column_exists(conn, "workdir_registry", "creation_ref")? + { + conn.execute_batch( + "ALTER TABLE workdir_registry RENAME COLUMN resolved_commit TO creation_ref;", + )?; + } + if !column_exists(conn, "workdir_registry", "current_selector")? { + conn.execute_batch("ALTER TABLE workdir_registry ADD COLUMN current_selector TEXT;")?; + } + if !column_exists(conn, "workdir_registry", "current_ref")? { + conn.execute_batch("ALTER TABLE workdir_registry ADD COLUMN current_ref TEXT;")?; + } + Ok(()) +} + +fn create_ticket_worker_assignment_tables(conn: &Connection) -> Result<()> { conn.execute_batch( r#" -ALTER TABLE workdir_registry RENAME COLUMN selector TO creation_selector; -ALTER TABLE workdir_registry RENAME COLUMN resolved_commit TO creation_ref; -ALTER TABLE workdir_registry ADD COLUMN current_selector TEXT; -ALTER TABLE workdir_registry ADD COLUMN current_ref TEXT; +CREATE TABLE IF NOT EXISTS ticket_worker_assignments ( + workspace_id TEXT NOT NULL REFERENCES workspaces(workspace_id) ON DELETE CASCADE, + ticket_id TEXT NOT NULL, + assignment_id TEXT NOT NULL, + runtime_id TEXT NOT NULL, + worker_id TEXT NOT NULL, + assigned_by TEXT NOT NULL, + assigned_at TEXT NOT NULL, + PRIMARY KEY (workspace_id, assignment_id), + UNIQUE (workspace_id, ticket_id, assignment_id) +); + +CREATE TABLE IF NOT EXISTS ticket_current_worker_assignments ( + workspace_id TEXT NOT NULL, + ticket_id TEXT NOT NULL, + assignment_id TEXT NOT NULL, + updated_at TEXT NOT NULL, + PRIMARY KEY (workspace_id, ticket_id), + FOREIGN KEY (workspace_id, ticket_id, assignment_id) + REFERENCES ticket_worker_assignments(workspace_id, ticket_id, assignment_id) + ON DELETE CASCADE +); + +CREATE TABLE IF NOT EXISTS ticket_worker_assignment_events ( + workspace_id TEXT NOT NULL REFERENCES workspaces(workspace_id) ON DELETE CASCADE, + ticket_id TEXT NOT NULL, + event_id TEXT NOT NULL, + action TEXT NOT NULL CHECK (action IN ('assigned', 'reassigned', 'unassigned')), + assignment_id TEXT, + previous_assignment_id TEXT, + actor TEXT NOT NULL, + created_at TEXT NOT NULL, + PRIMARY KEY (workspace_id, event_id) +); + +CREATE INDEX IF NOT EXISTS idx_ticket_assignments_worker + ON ticket_worker_assignments(workspace_id, runtime_id, worker_id, assigned_at DESC); +CREATE INDEX IF NOT EXISTS idx_ticket_assignment_events_ticket + ON ticket_worker_assignment_events(workspace_id, ticket_id, created_at DESC); "#, )?; Ok(()) } +fn create_ticket_notification_tables(conn: &Connection) -> Result<()> { + conn.execute_batch( + r#" +CREATE TABLE IF NOT EXISTS worker_workspace_credentials ( + credential_id TEXT PRIMARY KEY, + token TEXT NOT NULL UNIQUE, + workspace_id TEXT NOT NULL REFERENCES workspaces(workspace_id) ON DELETE CASCADE, + runtime_id TEXT NOT NULL, + worker_id TEXT, + created_at TEXT NOT NULL +); + +CREATE TABLE IF NOT EXISTS ticket_notification_outbox ( + notification_id TEXT PRIMARY KEY, + workspace_id TEXT NOT NULL REFERENCES workspaces(workspace_id) ON DELETE CASCADE, + ticket_id TEXT NOT NULL, + event_sequence INTEGER NOT NULL, + source_runtime_id TEXT NOT NULL, + source_worker_id TEXT NOT NULL, + previous_state TEXT NOT NULL, + current_state TEXT NOT NULL, + created_at TEXT NOT NULL +); + +CREATE TABLE IF NOT EXISTS ticket_notification_deliveries ( + notification_id TEXT NOT NULL REFERENCES ticket_notification_outbox(notification_id) ON DELETE CASCADE, + recipient_runtime_id TEXT NOT NULL, + recipient_worker_id TEXT NOT NULL, + recipient_kind TEXT NOT NULL CHECK (recipient_kind IN ('assigned', 'orchestrator')), + attempts INTEGER NOT NULL DEFAULT 0, + delivered_at TEXT, + last_error TEXT, + PRIMARY KEY (notification_id, recipient_runtime_id, recipient_worker_id) +); + +CREATE INDEX IF NOT EXISTS idx_ticket_notification_pending + ON ticket_notification_deliveries(delivered_at, attempts); +"#, + )?; + Ok(()) +} + +fn strengthen_ticket_worker_assignments(conn: &Connection) -> Result<()> { + conn.execute_batch( + r#" +ALTER TABLE ticket_current_worker_assignments RENAME TO ticket_current_worker_assignments_v16; + +CREATE TABLE ticket_current_worker_assignments ( + workspace_id TEXT NOT NULL, + ticket_id TEXT NOT NULL, + assignment_id TEXT NOT NULL, + runtime_id TEXT NOT NULL, + worker_id TEXT NOT NULL, + updated_at TEXT NOT NULL, + PRIMARY KEY (workspace_id, ticket_id), + UNIQUE (workspace_id, runtime_id, worker_id), + FOREIGN KEY (workspace_id, ticket_id, assignment_id) + REFERENCES ticket_worker_assignments(workspace_id, ticket_id, assignment_id) + ON DELETE CASCADE +); + +INSERT INTO ticket_current_worker_assignments ( + workspace_id, ticket_id, assignment_id, runtime_id, worker_id, updated_at +) +SELECT current.workspace_id, current.ticket_id, current.assignment_id, + assignment.runtime_id, assignment.worker_id, current.updated_at +FROM ticket_current_worker_assignments_v16 AS current +JOIN ticket_worker_assignments AS assignment + ON assignment.workspace_id = current.workspace_id + AND assignment.ticket_id = current.ticket_id + AND assignment.assignment_id = current.assignment_id; + +DROP TABLE ticket_current_worker_assignments_v16; + +CREATE TABLE ticket_assignment_operations ( + workspace_id TEXT NOT NULL REFERENCES workspaces(workspace_id) ON DELETE CASCADE, + operation_id TEXT NOT NULL, + action TEXT NOT NULL CHECK (action IN ('assign', 'reassign', 'unassign')), + ticket_id TEXT NOT NULL, + runtime_id TEXT, + worker_id TEXT, + assignment_id TEXT, + expected_assignment_id TEXT, + created_at TEXT NOT NULL, + PRIMARY KEY (workspace_id, operation_id) +); +"#, + )?; + Ok(()) +} + +fn strengthen_ticket_notifications(conn: &Connection) -> Result<()> { + conn.execute_batch( + r#" +ALTER TABLE worker_workspace_credentials ADD COLUMN expires_at TEXT; +ALTER TABLE worker_workspace_credentials ADD COLUMN revoked_at TEXT; +ALTER TABLE ticket_notification_outbox ADD COLUMN event_kind TEXT NOT NULL DEFAULT 'comment'; +ALTER TABLE ticket_notification_outbox ADD COLUMN source_operation_kind TEXT NOT NULL DEFAULT 'unknown'; +ALTER TABLE ticket_notification_outbox ADD COLUMN source_actor_role TEXT NOT NULL DEFAULT 'worker'; +ALTER TABLE ticket_notification_outbox ADD COLUMN source_assignment_id TEXT; + +CREATE TABLE ticket_notification_cursors ( + workspace_id TEXT NOT NULL REFERENCES workspaces(workspace_id) ON DELETE CASCADE, + ticket_id TEXT NOT NULL, + runtime_id TEXT NOT NULL, + worker_id TEXT NOT NULL, + last_event_index INTEGER NOT NULL, + updated_at TEXT NOT NULL, + PRIMARY KEY (workspace_id, ticket_id, runtime_id, worker_id) +); +"#, + )?; + Ok(()) +} + +fn strengthen_ticket_assignment_lifecycle_reservations(conn: &Connection) -> Result<()> { + if !column_exists(conn, "ticket_assignment_operations", "request_fingerprint")? { + conn.execute_batch( + "ALTER TABLE ticket_assignment_operations ADD COLUMN request_fingerprint TEXT;", + )?; + } + // The Workdir and Ticket branches both used schema version 15 before + // integration. A database that ran the Ticket branch through v19 still + // needs the Workdir revision projection while already having the lifecycle + // reservation column. Reconcile both shapes at the combined v20 boundary. + add_workdir_revision_observations(conn)?; + Ok(()) +} + fn create_objective_event_tables(conn: &Connection) -> Result<()> { conn.execute_batch( r#" @@ -2884,7 +4192,7 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT); let db = dir.path().join("control-plane.sqlite"); let store = SqliteWorkspaceStore::open(&db).unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 15); + assert_eq!(store.schema_version().await.unwrap(), 20); let record = WorkspaceRecord { workspace_id: "local-dev".to_string(), @@ -2897,13 +4205,412 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT); store.upsert_workspace(&record).await.unwrap(); let reopened = SqliteWorkspaceStore::open(&db).unwrap(); - assert_eq!(reopened.schema_version().await.unwrap(), 15); + assert_eq!(reopened.schema_version().await.unwrap(), 20); assert_eq!( reopened.get_workspace("local-dev").await.unwrap(), Some(record) ); } + #[tokio::test] + async fn ticket_worker_assignment_replaces_current_and_preserves_audit_history() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("server.db"); + let store = SqliteWorkspaceStore::open(&db).unwrap(); + store + .upsert_workspace(&WorkspaceRecord { + workspace_id: "workspace-a".to_string(), + owner_account_id: None, + display_name: "Workspace A".to_string(), + state: "active".to_string(), + created_at: "2026-07-31T00:00:00Z".to_string(), + updated_at: "2026-07-31T00:00:00Z".to_string(), + }) + .await + .unwrap(); + + let first = TicketWorkerAssignmentRecord { + workspace_id: "workspace-a".to_string(), + ticket_id: "ticket-1".to_string(), + assignment_id: "assignment-1".to_string(), + runtime_id: "runtime-1".to_string(), + worker_id: "worker-1".to_string(), + assigned_by: "user-1".to_string(), + assigned_at: "2026-07-31T00:00:01Z".to_string(), + }; + let created = store + .set_current_ticket_worker_assignment(&first, None, "event-1", "operation-1", false) + .unwrap(); + assert_eq!(created.current, first); + assert_eq!(created.previous, None); + let retried = store + .set_current_ticket_worker_assignment( + &TicketWorkerAssignmentRecord { + assignment_id: "ignored-retry-assignment".to_string(), + ..first.clone() + }, + None, + "ignored-retry-event", + "operation-1", + false, + ) + .unwrap(); + assert_eq!(retried.current, first); + assert_eq!( + store + .list_ticket_worker_assignment_events("workspace-a", "ticket-1", 10) + .unwrap() + .len(), + 1, + "idempotent retry must not append another assignment event" + ); + let implicit_reassign = store + .set_current_ticket_worker_assignment( + &TicketWorkerAssignmentRecord { + assignment_id: "implicit-reassign".to_string(), + worker_id: "worker-other".to_string(), + ..first.clone() + }, + None, + "implicit-event", + "implicit-operation", + false, + ) + .unwrap_err(); + assert!(matches!( + implicit_reassign, + Error::TicketAssignmentConflict(_) + )); + let worker_conflict = store + .set_current_ticket_worker_assignment( + &TicketWorkerAssignmentRecord { + ticket_id: "ticket-2".to_string(), + assignment_id: "worker-conflict".to_string(), + ..first.clone() + }, + None, + "worker-conflict-event", + "worker-conflict-operation", + false, + ) + .unwrap_err(); + assert!(matches!( + worker_conflict, + Error::TicketAssignmentConflict(_) + )); + + let second = TicketWorkerAssignmentRecord { + assignment_id: "assignment-2".to_string(), + runtime_id: "runtime-2".to_string(), + worker_id: "worker-2".to_string(), + assigned_by: "user-2".to_string(), + assigned_at: "2026-07-31T00:00:02Z".to_string(), + ..first.clone() + }; + let replaced = store + .set_current_ticket_worker_assignment( + &second, + Some("assignment-1"), + "event-2", + "operation-2", + true, + ) + .unwrap(); + assert_eq!(replaced.current, second); + assert_eq!(replaced.previous, Some(first.clone())); + let replayed_reassignment = store + .set_current_ticket_worker_assignment( + &second, + Some("assignment-1"), + "ignored-reassign-event", + "operation-2", + true, + ) + .unwrap(); + assert_eq!(replayed_reassignment, replaced); + assert_eq!( + store + .get_current_ticket_worker_assignment("workspace-a", "ticket-1") + .unwrap(), + Some(second.clone()) + ); + + let stale = store + .clear_current_ticket_worker_assignment( + "workspace-a", + "ticket-1", + Some("assignment-1"), + "unassign-operation-stale", + "event-stale", + "user-1", + "2026-07-31T00:00:03Z", + ) + .unwrap_err(); + assert!(matches!(stale, Error::TicketAssignmentConflict(_))); + + let cleared = store + .clear_current_ticket_worker_assignment( + "workspace-a", + "ticket-1", + Some("assignment-2"), + "unassign-operation-2", + "event-3", + "user-2", + "2026-07-31T00:00:03Z", + ) + .unwrap(); + assert_eq!(cleared, Some(second.clone())); + let retried_clear = store + .clear_current_ticket_worker_assignment( + "workspace-a", + "ticket-1", + Some("assignment-2"), + "unassign-operation-2", + "ignored-clear-event", + "user-2", + "2026-07-31T00:00:04Z", + ) + .unwrap(); + assert_eq!(retried_clear, Some(second)); + store + .reserve_ticket_assignment_operation( + "workspace-a", + "reserved-operation", + "ticket-3", + "runtime-3", + None, + "sha256:reserved", + "2026-07-31T00:00:05Z", + ) + .unwrap(); + drop(store); + let store = SqliteWorkspaceStore::open(&db).unwrap(); + let pending = store + .get_ticket_assignment_operation("workspace-a", "reserved-operation") + .unwrap() + .unwrap(); + assert_eq!(pending.worker_id, None); + assert_eq!( + pending.request_fingerprint.as_deref(), + Some("sha256:reserved") + ); + store + .bind_ticket_assignment_operation_worker( + "workspace-a", + "reserved-operation", + "worker-3", + ) + .unwrap(); + let reserved_assignment = TicketWorkerAssignmentRecord { + workspace_id: "workspace-a".to_string(), + ticket_id: "ticket-3".to_string(), + assignment_id: "assignment-3".to_string(), + runtime_id: "runtime-3".to_string(), + worker_id: "worker-3".to_string(), + assigned_by: "runtime".to_string(), + assigned_at: "2026-07-31T00:00:06Z".to_string(), + }; + let completed_reservation = store + .set_current_ticket_worker_assignment( + &reserved_assignment, + None, + "reserved-event", + "reserved-operation", + false, + ) + .unwrap(); + assert_eq!(completed_reservation.current, reserved_assignment); + assert_eq!( + store + .get_ticket_assignment_operation("workspace-a", "reserved-operation") + .unwrap() + .and_then(|operation| operation.assignment_id), + Some("assignment-3".to_string()) + ); + assert_eq!( + store + .get_current_ticket_worker_assignment("workspace-a", "ticket-1") + .unwrap(), + None + ); + + let events = store + .list_ticket_worker_assignment_events("workspace-a", "ticket-1", 10) + .unwrap(); + assert_eq!( + events + .iter() + .map(|event| event.action.as_str()) + .collect::>(), + vec!["unassigned", "reassigned", "assigned"] + ); + assert_eq!(events[1].assignment_id.as_deref(), Some("assignment-2")); + assert_eq!( + events[1].previous_assignment_id.as_deref(), + Some("assignment-1") + ); + } + + #[tokio::test] + async fn worker_credential_binds_once_and_notification_outbox_is_durable() { + let dir = tempfile::tempdir().unwrap(); + let store = SqliteWorkspaceStore::open(dir.path().join("server.db")).unwrap(); + store + .upsert_workspace(&WorkspaceRecord { + workspace_id: "workspace-a".to_string(), + owner_account_id: None, + display_name: "Workspace A".to_string(), + state: "active".to_string(), + created_at: "2026-07-31T00:00:00Z".to_string(), + updated_at: "2026-07-31T00:00:00Z".to_string(), + }) + .await + .unwrap(); + store + .upsert_worker_workspace_credential(&WorkerWorkspaceCredentialRecord { + credential_id: "credential-1".to_string(), + token: "secret-token".to_string(), + workspace_id: "workspace-a".to_string(), + runtime_id: "runtime-1".to_string(), + worker_id: None, + created_at: "2026-07-31T00:00:01Z".to_string(), + expires_at: "2099-01-01T00:00:00Z".to_string(), + revoked_at: None, + }) + .unwrap(); + let bound = store + .authenticate_worker_workspace_credential("secret-token", "workspace-a", "worker-1") + .unwrap() + .unwrap(); + assert_eq!(bound.worker_id.as_deref(), Some("worker-1")); + let refreshed = store + .refresh_worker_workspace_credential( + "secret-token", + "workspace-a", + "worker-1", + "refreshed-token", + "2099-02-01T00:00:00Z", + ) + .unwrap() + .unwrap(); + assert_eq!(refreshed.token, "refreshed-token"); + assert!(store + .authenticate_worker_workspace_credential( + "secret-token", + "workspace-a", + "worker-1", + ) + .unwrap() + .is_none()); + assert!( + store + .authenticate_worker_workspace_credential( + "refreshed-token", + "workspace-a", + "worker-1", + ) + .unwrap() + .is_some() + ); + store + .revoke_worker_workspace_credentials( + "workspace-a", + "runtime-1", + "worker-1", + "2026-08-01T00:00:00Z", + ) + .unwrap(); + assert!( + store + .authenticate_worker_workspace_credential( + "refreshed-token", + "workspace-a", + "worker-1", + ) + .unwrap() + .is_none() + ); + assert!( + store + .authenticate_worker_workspace_credential( + "secret-token", + "workspace-a", + "worker-2", + ) + .unwrap() + .is_none() + ); + + store + .enqueue_ticket_notification( + "notification-1", + "workspace-a", + "ticket-1", + 4, + "runtime-1", + "worker-1", + "queued", + "inprogress", + "2026-07-31T00:00:02Z", + &[TicketNotificationRecipient { + runtime_id: "runtime-1".to_string(), + worker_id: "worker-2".to_string(), + recipient_kind: "assigned".to_string(), + }], + ) + .unwrap(); + let pending = store + .list_pending_ticket_notification_deliveries("workspace-a", 10) + .unwrap(); + assert_eq!(pending.len(), 1); + assert_eq!(pending[0].event_sequence, 4); + store + .reroute_ticket_notification_delivery( + "notification-1", + "runtime-1", + "worker-2", + "runtime-2", + "worker-3", + ) + .unwrap(); + assert_eq!( + store + .count_ticket_notification_deliveries_for_recipient( + "workspace-a", + "ticket-1", + "runtime-1", + "worker-2", + ) + .unwrap(), + 0 + ); + assert_eq!( + store + .count_ticket_notification_deliveries_for_recipient( + "workspace-a", + "ticket-1", + "runtime-2", + "worker-3", + ) + .unwrap(), + 1 + ); + store + .mark_ticket_notification_delivered( + "notification-1", + "runtime-2", + "worker-3", + "2026-07-31T00:00:03Z", + ) + .unwrap(); + assert!( + store + .list_pending_ticket_notification_deliveries("workspace-a", 10) + .unwrap() + .is_empty() + ); + } + #[test] fn fresh_schema_matches_workspace_db_v0_boundaries() { let conn = Connection::open_in_memory().unwrap(); @@ -2923,6 +4630,9 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT); "artifacts", "audit_events", "worker_registry", + "ticket_worker_assignments", + "ticket_current_worker_assignments", + "ticket_worker_assignment_events", "workdir_registry", "worker_workdir_links", "accounts", @@ -3103,7 +4813,7 @@ CREATE TABLE ticket_worker_links (ticket_id TEXT, worker_ref_key TEXT); .unwrap(); let store = SqliteWorkspaceStore::from_connection(conn).unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 15); + assert_eq!(store.schema_version().await.unwrap(), 20); store .with_conn(|conn| { @@ -3259,10 +4969,37 @@ INSERT INTO workdir_registry ( ); } + #[test] + fn combined_v20_migration_reconciles_ticket_branch_schema() { + let conn = Connection::open_in_memory().unwrap(); + conn.execute_batch( + r#" +CREATE TABLE workdir_registry ( + selector TEXT, + resolved_commit TEXT +); +CREATE TABLE ticket_assignment_operations ( + request_fingerprint TEXT +); +"#, + ) + .unwrap(); + + strengthen_ticket_assignment_lifecycle_reservations(&conn).unwrap(); + + assert!(column_exists(&conn, "workdir_registry", "creation_selector").unwrap()); + assert!(column_exists(&conn, "workdir_registry", "creation_ref").unwrap()); + assert!(column_exists(&conn, "workdir_registry", "current_selector").unwrap()); + assert!(column_exists(&conn, "workdir_registry", "current_ref").unwrap()); + assert!( + column_exists(&conn, "ticket_assignment_operations", "request_fingerprint").unwrap() + ); + } + #[tokio::test] async fn repository_records_round_trip() { let store = SqliteWorkspaceStore::in_memory().unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 15); + assert_eq!(store.schema_version().await.unwrap(), 20); let workspace = WorkspaceRecord { workspace_id: "local-dev".to_string(), owner_account_id: None, @@ -3300,7 +5037,7 @@ INSERT INTO workdir_registry ( #[tokio::test] async fn memory_authority_records_round_trip_and_close_staging() { let store = SqliteWorkspaceStore::in_memory().unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 15); + assert_eq!(store.schema_version().await.unwrap(), 20); let workspace = WorkspaceRecord { workspace_id: "local-dev".to_string(), owner_account_id: None, @@ -3478,7 +5215,7 @@ INSERT INTO workdir_registry ( #[tokio::test] async fn account_and_login_records_round_trip() { let store = SqliteWorkspaceStore::in_memory().unwrap(); - assert_eq!(store.schema_version().await.unwrap(), 15); + assert_eq!(store.schema_version().await.unwrap(), 20); let now = "2026-07-22T00:00:00Z".to_string(); let account = AccountRecord { account_id: "acct-user-alice".to_string(),