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

server: address Ticket orchestration review

Browse Source
This commit is contained in:
Hare
2026-07-31 22:15:55 +09:00
parent 005f6cb498
commit e9a6269d9f
6 changed files with 2081 additions and 361 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/ticket/src/lib.rs
+98 -3
View File
@@ -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<String>,
event_attributes: BTreeMap<String, String>,
mutation_hook: Option<Arc<SqliteTicketMutationHook>>,
}
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<String, String>) -> Self {
self.event_attributes = attributes;
self
}
pub fn with_mutation_hook(mut self, hook: Arc<SqliteTicketMutationHook>) -> 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();
crates/ticket/src/tool.rs
+148 -79
View File
@@ -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<String>,
/// Optional thread author for the create event.
#[serde(default)]
author: Option<String>,
/// Optional assignee frontmatter value.
#[serde(default)]
assignee: Option<String>,
@@ -376,9 +387,6 @@ struct TicketCreateParams {
/// Optional state frontmatter value. Defaults to `planning`.
#[serde(default)]
state: Option<TicketWorkflowStateParam>,
/// Optional queued_by frontmatter value.
#[serde(default)]
queued_by: Option<String>,
/// Optional queued_at frontmatter value.
#[serde(default)]
queued_at: Option<String>,
@@ -412,9 +420,6 @@ struct TicketEditItemParams {
/// Optional target repository/ref update.
#[serde(default)]
target: Option<crate::TicketTargetEdit>,
/// Optional thread author for the audited item_edit event.
#[serde(default)]
author: Option<String>,
}
#[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<String>,
}
#[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<String>,
}
#[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<String>,
/// Reason attached to the state_changed event. Defaults to `planning_ready`.
#[serde(default)]
reason: Option<String>,
@@ -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<String>,
}
#[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<String>,
}
#[derive(Debug, Deserialize, schemars::JsonSchema)]
@@ -673,9 +652,6 @@ struct TicketRelationRecordParams {
/// Optional bounded rationale/note.
#[serde(default)]
note: Option<String>,
/// Optional record author.
#[serde(default)]
author: Option<String>,
}
#[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<AcceptedOrchestrationPlanParams>,
/// Optional record author.
#[serde(default)]
author: Option<String>,
}
#[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<ToolOutput, ToolError> {
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<ToolOutput, ToolError> {
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<ToolOutput, ToolError> {
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<ToolOutput, ToolError> {
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<TicketToolBackend>) -> Vec<ToolDefinition
tool_definition::<TicketListTool>("TicketList", backend.clone()),
tool_definition::<TicketShowTool>("TicketShow", backend.clone()),
tool_definition::<TicketCommentTool>("TicketComment", backend.clone()),
tool_definition::<TicketPlanTool>("TicketPlan", backend.clone()),
tool_definition::<TicketDecisionTool>("TicketDecision", backend.clone()),
tool_definition::<TicketImplementationReportTool>(
"TicketImplementationReport",
backend.clone(),
),
tool_definition::<TicketReviewTool>("TicketReview", backend.clone()),
tool_definition::<TicketIntakeReadyTool>("TicketIntakeReady", backend.clone()),
tool_definition::<TicketQueueTool>("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)
crates/worker/src/worker.rs
+153 -13
View File
@@ -226,7 +226,7 @@ pub struct RuntimeWorkspaceHttpClient {
workspace_id: String,
base_url: String,
worker_id: String,
access_token: Option<String>,
access_token: Mutex<Option<String>>,
}
impl std::fmt::Debug for RuntimeWorkspaceHttpClient {
@@ -238,7 +238,11 @@ impl std::fmt::Debug for RuntimeWorkspaceHttpClient {
.field("worker_id", &self.worker_id)
.field(
"access_token",
&self.access_token.as_ref().map(|_| "[redacted]"),
&self
.access_token
.lock()
.ok()
.and_then(|token| token.as_ref().map(|_| "[redacted]")),
)
.finish()
}
@@ -254,12 +258,12 @@ impl RuntimeWorkspaceHttpClient {
workspace_id: workspace_id.into(),
base_url: base_url.into().trim_end_matches('/').to_string(),
worker_id: worker_id.into(),
access_token: None,
access_token: Mutex::new(None),
}
}
pub fn with_access_token(mut self, access_token: Option<String>) -> Self {
self.access_token = access_token;
pub fn with_access_token(self, access_token: Option<String>) -> Self {
*self.access_token.lock().expect("new credential mutex") = access_token;
self
}
}
@@ -283,25 +287,93 @@ impl WorkspaceClient for RuntimeWorkspaceHttpClient {
) -> Result<WorkspaceResponse, WorkspaceClientError> {
let base_url = self.base_url.clone();
let worker_id = self.worker_id.clone();
let access_token = self.access_token.clone();
if tokio::runtime::Handle::try_current().is_ok() {
return std::thread::spawn(move || {
execute_runtime_workspace_http(
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.as_deref(),
request,
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);
}
execute_runtime_workspace_http(&base_url, &worker_id, access_token.as_deref(), request)
Ok(result.0)
}
}
fn execute_runtime_workspace_http_with_refresh(
base_url: &str,
worker_id: &str,
access_token: Option<String>,
request: WorkspaceRequest,
) -> Result<(WorkspaceResponse, Option<String>), 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,
@@ -6185,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]
crates/workspace-server/src/hosts.rs
+15
View File
@@ -308,6 +308,13 @@ pub struct WorkerSpawnWorkingDirectoryRequest {
pub selector: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct WorkerTicketAssignmentRequest {
pub ticket_id: String,
pub operation_id: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct WorkerSpawnRequest {
@@ -317,6 +324,8 @@ pub struct WorkerSpawnRequest {
pub acceptance: WorkerSpawnAcceptanceRequirement,
pub profile: ProfileSelector,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub ticket_assignment: Option<WorkerTicketAssignmentRequest>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub initial_input: Option<EmbeddedWorkerInput>,
/// Optional safe working-directory creation request. The Workspace server resolves
/// this into a runtime-internal `WorkingDirectoryRequest` from configured
@@ -429,6 +438,8 @@ pub struct WorkerStopResult {
pub struct WorkerLifecycleRequest {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reason: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub ticket_assignment: Option<WorkerTicketAssignmentRequest>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
@@ -4177,6 +4188,7 @@ 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,
@@ -4304,6 +4316,7 @@ 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,
@@ -4401,6 +4414,7 @@ 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,
@@ -4434,6 +4448,7 @@ 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,
crates/workspace-server/src/server.rs
+888 -225
View File
File diff suppressed because it is too large Load Diff
crates/workspace-server/src/store.rs
+779 -41
View File
File diff suppressed because it is too large Load Diff