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

workspace: remove worker credential refresh flow

Browse Source
This commit is contained in:
Hare
2026-08-03 17:08:23 +09:00
parent ddadc830ac
commit 0ffaa6c741
17 changed files with 184 additions and 976 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+1 -15
View File
@@ -26,10 +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,
WorkspaceClientError,
};
use crate::worker::{SystemItemCommitter, Worker, WorkerError, WorkerRunResult};
use protocol::{
AlertLevel, AlertSource, ErrorCode, Event, Method, RewindTargetId, RunResult, Segment,
TurnResult, WorkerStatus,
@@ -43,7 +40,6 @@ use protocol::{
pub struct WorkerHandle {
method_tx: mpsc::Sender<Method>,
event_tx: broadcast::Sender<Event>,
workspace_client: Arc<dyn WorkspaceClient>,
pub shared_state: Arc<WorkerSharedState>,
pub runtime_dir: Arc<RuntimeDir>,
pub alerter: Alerter,
@@ -119,14 +115,6 @@ impl WorkerHandle {
pub fn alert(&self, level: AlertLevel, source: AlertSource, message: String) {
self.alerter.alert(level, source, message);
}
/// Replace the Runtime-issued Workspace access token used by this live Worker.
pub fn replace_workspace_access_token(
&self,
access_token: String,
) -> Result<(), WorkspaceClientError> {
self.workspace_client.replace_access_token(access_token)
}
}
async fn set_controller_status(
@@ -248,7 +236,6 @@ impl WorkerController {
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
let (method_tx, method_rx) = mpsc::channel::<Method>(32);
let (event_tx, _) = broadcast::channel::<Event>(256);
let workspace_client = worker.workspace_client_handle();
let alerter = Alerter::new(event_tx.clone());
let in_flight = InFlightEvents::new(event_tx.clone());
worker.attach_in_flight_events(in_flight.clone());
@@ -367,7 +354,6 @@ impl WorkerController {
let handle = WorkerHandle {
method_tx,
event_tx: event_tx.clone(),
workspace_client,
shared_state: shared_state.clone(),
runtime_dir: runtime_dir.clone(),
alerter: alerter.clone(),
crates/worker/src/feature/builtin/memory.rs
+1
View File
@@ -348,6 +348,7 @@ mod tests {
Arc::new(crate::worker::RuntimeWorkspaceHttpClient::new(
"workspace",
"http://backend",
"test-runtime",
"test-worker",
))
}
crates/worker/src/feature/builtin/objective.rs
+1
View File
@@ -627,6 +627,7 @@ mod tests {
crate::worker::RuntimeWorkspaceHttpClient::new(
"workspace",
"http://backend",
"test-runtime",
"test-worker",
),
)));
crates/worker/src/feature/builtin/session_explore.rs
+1
View File
@@ -662,6 +662,7 @@ mod tests {
Arc::new(crate::worker::RuntimeWorkspaceHttpClient::new(
"test-workspace",
format!("http://{addr}"),
"test-runtime",
"test-worker",
)),
rx,
crates/worker/src/feature/builtin/ticket.rs
+8 -1
View File
@@ -1372,6 +1372,7 @@ provider = "github"
crate::worker::RuntimeWorkspaceHttpClient::new(
"workspace-a",
"not-a-url",
"test-runtime",
"test-worker",
),
));
@@ -1407,6 +1408,7 @@ provider = "github"
let client = Arc::new(crate::worker::RuntimeWorkspaceHttpClient::new(
"workspace-a",
format!("http://{address}"),
"test-runtime",
"worker-a",
));
let backend = WorkspaceHttpTicketBackend::new(client);
@@ -1448,7 +1450,12 @@ provider = "github"
});
let backend = WorkspaceHttpTicketBackend::new(Arc::new(
crate::worker::RuntimeWorkspaceHttpClient::new("workspace-a", base_url, "test-worker"),
crate::worker::RuntimeWorkspaceHttpClient::new(
"workspace-a",
base_url,
"test-runtime",
"test-worker",
),
));
let created = backend.create(NewTicket::new("HTTP ticket")).unwrap();
crates/worker/src/skill.rs
+8 -3
View File
@@ -193,11 +193,15 @@ 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 runtime_header = None;
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-runtime-id: ") {
runtime_header = Some(value.trim().to_string());
}
if let Some(value) = line.strip_prefix("x-yoi-worker-id: ") {
worker_header = Some(value.trim().to_string());
}
@@ -208,8 +212,9 @@ mod tests {
break;
}
}
assert_eq!(runtime_header.as_deref(), Some("runtime-test"));
assert_eq!(worker_header.as_deref(), Some("test-worker"));
assert_eq!(authorization.as_deref(), Some("Bearer test-credential"));
assert_eq!(authorization, None);
let body = serde_json::json!({
"authority": "workspace-backend-skills-v0",
"entries": [{
@@ -234,9 +239,9 @@ mod tests {
let client = crate::worker::RuntimeWorkspaceHttpClient::new(
"ws-1",
format!("http://{addr}"),
"runtime-test",
"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");
crates/worker/src/worker.rs
+45 -173
View File
@@ -216,13 +216,6 @@ pub trait WorkspaceClient: std::fmt::Debug + Send + Sync {
fn is_available(&self) -> bool;
fn execute(&self, request: WorkspaceRequest)
-> Result<WorkspaceResponse, WorkspaceClientError>;
/// Replace the Runtime-issued Workspace access token for this live client.
fn replace_access_token(&self, _access_token: String) -> Result<(), WorkspaceClientError> {
Err(WorkspaceClientError::Unavailable(
"Workspace client does not support access token replacement".to_string(),
))
}
}
/// HTTP forwarding client created by Runtime for one concrete Worker execution.
@@ -233,8 +226,8 @@ pub trait WorkspaceClient: std::fmt::Debug + Send + Sync {
pub struct RuntimeWorkspaceHttpClient {
workspace_id: String,
base_url: String,
runtime_id: String,
worker_id: String,
access_token: Mutex<Option<String>>,
}
impl std::fmt::Debug for RuntimeWorkspaceHttpClient {
@@ -243,15 +236,8 @@ impl std::fmt::Debug for RuntimeWorkspaceHttpClient {
.debug_struct("RuntimeWorkspaceHttpClient")
.field("workspace_id", &self.workspace_id)
.field("base_url", &self.base_url)
.field("runtime_id", &self.runtime_id)
.field("worker_id", &self.worker_id)
.field(
"access_token",
&self
.access_token
.lock()
.ok()
.and_then(|token| token.as_ref().map(|_| "[redacted]")),
)
.finish()
}
}
@@ -260,20 +246,16 @@ impl RuntimeWorkspaceHttpClient {
pub fn new(
workspace_id: impl Into<String>,
base_url: impl Into<String>,
runtime_id: impl Into<String>,
worker_id: impl Into<String>,
) -> Self {
Self {
workspace_id: workspace_id.into(),
base_url: base_url.into().trim_end_matches('/').to_string(),
runtime_id: runtime_id.into(),
worker_id: worker_id.into(),
access_token: Mutex::new(None),
}
}
pub fn with_access_token(self, access_token: Option<String>) -> Self {
*self.access_token.lock().expect("new credential mutex") = access_token;
self
}
}
impl WorkspaceClient for RuntimeWorkspaceHttpClient {
@@ -294,105 +276,26 @@ impl WorkspaceClient for RuntimeWorkspaceHttpClient {
request: WorkspaceRequest,
) -> Result<WorkspaceResponse, WorkspaceClientError> {
let base_url = self.base_url.clone();
let runtime_id = self.runtime_id.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() {
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,
)
execute_runtime_workspace_http(&base_url, &runtime_id, &worker_id, request)
})
.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, &runtime_id, &worker_id, request)
}
Ok(result.0)
}
fn replace_access_token(&self, access_token: String) -> Result<(), WorkspaceClientError> {
*self.access_token.lock().map_err(|_| {
WorkspaceClientError::Request("workspace credential lock poisoned".to_string())
})? = Some(access_token);
Ok(())
}
}
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,
runtime_id: &str,
worker_id: &str,
access_token: Option<&str>,
request: WorkspaceRequest,
) -> Result<WorkspaceResponse, WorkspaceClientError> {
if !request.path.starts_with('/') || request.path.starts_with("//") {
@@ -409,10 +312,8 @@ fn execute_runtime_workspace_http(
let client = reqwest::blocking::Client::new();
let mut request_builder = client
.request(method, url)
.header("x-yoi-runtime-id", runtime_id)
.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")
@@ -6132,6 +6033,7 @@ mod build_summary_prompt_tests {
Arc::new(RuntimeWorkspaceHttpClient::new(
"test-memory",
format!("http://{addr}"),
"test-runtime",
"test-worker",
)),
)
@@ -6268,6 +6170,7 @@ mod build_summary_prompt_tests {
Arc::new(RuntimeWorkspaceHttpClient::new(
"ws-skill",
format!("http://{addr}"),
"test-runtime",
"test-worker",
)),
),
@@ -6311,88 +6214,57 @@ mod build_summary_prompt_tests {
}
#[test]
fn runtime_workspace_client_refreshes_expired_credential_and_retries() {
fn runtime_workspace_client_sends_runtime_worker_identity_without_bearer() {
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;
}
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();
assert!(first_line.contains("/api/w/workspace-a/tickets/search"));
let mut runtime_id = String::new();
let mut worker_id = String::new();
let mut authorization = String::new();
loop {
let mut line = String::new();
reader.read_line(&mut line).unwrap();
if let Some(value) = line.strip_prefix("x-yoi-runtime-id: ") {
runtime_id = value.trim().to_string();
}
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();
}
if let Some(value) = line.strip_prefix("x-yoi-worker-id: ") {
worker_id = value.trim().to_string();
}
if let Some(value) = line.strip_prefix("authorization: ") {
authorization = value.trim().to_string();
}
if line == "\r\n" || line.is_empty() {
break;
}
}
assert_eq!(runtime_id, "runtime-a");
assert_eq!(worker_id, "worker-a");
assert!(authorization.is_empty());
stream
.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\n{}")
.unwrap();
});
let client = RuntimeWorkspaceHttpClient::new(
"workspace-refresh",
"workspace-a",
format!("http://{address}"),
"worker-refresh",
)
.with_access_token(Some("expired-token".to_string()));
"runtime-a",
"worker-a",
);
let response = client
.execute(WorkspaceRequest::get(
"/api/w/workspace-refresh/tickets/search",
))
.execute(WorkspaceRequest::get("/api/w/workspace-a/tickets/search"))
.unwrap();
assert_eq!(response.status, 200);
server.join().unwrap();
}
#[test]
fn runtime_workspace_client_can_install_missing_access_token() {
let client =
RuntimeWorkspaceHttpClient::new("workspace-a", "https://workspace.example", "worker-a");
client
.replace_access_token("replacement-token".to_string())
.unwrap();
assert_eq!(
client.access_token.lock().unwrap().as_deref(),
Some("replacement-token")
);
}
fn minimal_manifest() -> WorkerManifest {
let toml_str = r#"
[worker]