feat: add session-owned uploaded file attachments
This commit is contained in:
@@ -8,7 +8,7 @@ use crate::interaction::WorkerInput;
|
||||
#[cfg(feature = "ws-server")]
|
||||
use crate::observation::WorkerObservationEvent;
|
||||
use crate::working_directory::{WorkingDirectoryBinding, WorkingDirectoryDiagnostic};
|
||||
use protocol::Method;
|
||||
use protocol::{Method, UploadedFileRef};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
@@ -33,6 +33,8 @@ pub enum WorkerExecutionOperation {
|
||||
Spawn,
|
||||
Restore,
|
||||
Input,
|
||||
UploadFile,
|
||||
DeleteUploadedFile,
|
||||
ProtocolMethod,
|
||||
Stop,
|
||||
Cancel,
|
||||
@@ -385,6 +387,30 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static {
|
||||
input: WorkerInput,
|
||||
) -> WorkerExecutionResult;
|
||||
|
||||
fn upload_file(
|
||||
&self,
|
||||
_handle: &WorkerExecutionHandle,
|
||||
_file_name: &str,
|
||||
_media_type: &str,
|
||||
_content: &[u8],
|
||||
) -> Result<UploadedFileRef, WorkerExecutionResult> {
|
||||
Err(WorkerExecutionResult::unsupported(
|
||||
WorkerExecutionOperation::UploadFile,
|
||||
"execution backend does not support file upload",
|
||||
))
|
||||
}
|
||||
|
||||
fn delete_uploaded_file(
|
||||
&self,
|
||||
_handle: &WorkerExecutionHandle,
|
||||
_artifact_id: &str,
|
||||
) -> WorkerExecutionResult {
|
||||
WorkerExecutionResult::unsupported(
|
||||
WorkerExecutionOperation::DeleteUploadedFile,
|
||||
"execution backend does not support uploaded-file deletion",
|
||||
)
|
||||
}
|
||||
|
||||
fn dispatch_method(
|
||||
&self,
|
||||
_handle: &WorkerExecutionHandle,
|
||||
@@ -514,6 +540,25 @@ impl WorkerExecutionBackendRef {
|
||||
self.backend.dispatch_input(handle, input)
|
||||
}
|
||||
|
||||
pub(crate) fn upload_file(
|
||||
&self,
|
||||
handle: &WorkerExecutionHandle,
|
||||
file_name: &str,
|
||||
media_type: &str,
|
||||
content: &[u8],
|
||||
) -> Result<UploadedFileRef, WorkerExecutionResult> {
|
||||
self.backend
|
||||
.upload_file(handle, file_name, media_type, content)
|
||||
}
|
||||
|
||||
pub(crate) fn delete_uploaded_file(
|
||||
&self,
|
||||
handle: &WorkerExecutionHandle,
|
||||
artifact_id: &str,
|
||||
) -> WorkerExecutionResult {
|
||||
self.backend.delete_uploaded_file(handle, artifact_id)
|
||||
}
|
||||
|
||||
pub(crate) fn dispatch_method(
|
||||
&self,
|
||||
handle: &WorkerExecutionHandle,
|
||||
|
||||
@@ -32,7 +32,7 @@ use axum::body::{Body, Bytes};
|
||||
use axum::extract::rejection::{JsonRejection, QueryRejection};
|
||||
#[cfg(feature = "ws-server")]
|
||||
use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade};
|
||||
use axum::extract::{Extension, Path, Query, State};
|
||||
use axum::extract::{DefaultBodyLimit, Extension, Path, Query, State};
|
||||
use axum::http::{Method, Request, StatusCode, header};
|
||||
use axum::middleware::{self, Next};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
@@ -238,6 +238,14 @@ fn runtime_http_router_with_optional_auth(
|
||||
post(execute_worker_retention),
|
||||
)
|
||||
.route("/v1/workers/{worker_id}/input", post(send_worker_input))
|
||||
.route(
|
||||
"/v1/workers/{worker_id}/attachments",
|
||||
post(upload_worker_file).layer(DefaultBodyLimit::max(MAX_WORKER_FILE_UPLOAD_BYTES)),
|
||||
)
|
||||
.route(
|
||||
"/v1/workers/{worker_id}/attachments/{artifact_id}",
|
||||
delete(delete_worker_uploaded_file),
|
||||
)
|
||||
.route("/v1/workers/{worker_id}/restore", post(restore_worker))
|
||||
.route(
|
||||
"/v1/workers/{worker_id}/workspace-api",
|
||||
@@ -263,6 +271,9 @@ fn runtime_http_router_with_optional_auth(
|
||||
.layer(middleware::from_fn_with_state(state, require_runtime_auth))
|
||||
}
|
||||
|
||||
pub const MAX_WORKER_FILE_UPLOAD_BYTES: usize =
|
||||
session_store::DEFAULT_MAX_UPLOADED_FILE_BYTES as usize;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct RuntimeHttpState {
|
||||
runtime: Runtime,
|
||||
@@ -375,6 +386,22 @@ pub struct RuntimeHttpWorkerInputResponse {
|
||||
pub ack: WorkerInteractionAck,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
pub struct RuntimeHttpUploadFileQuery {
|
||||
pub file_name: String,
|
||||
pub media_type: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct RuntimeHttpUploadedFileResponse {
|
||||
pub file: protocol::UploadedFileRef,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct RuntimeHttpUploadedFileDeleteResponse {
|
||||
pub deleted: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct RuntimeHttpWorkerCompletionsRequest {
|
||||
pub kind: protocol::CompletionKind,
|
||||
@@ -1420,6 +1447,55 @@ async fn worker_completions(
|
||||
}))
|
||||
}
|
||||
|
||||
async fn upload_worker_file(
|
||||
State(state): State<RuntimeHttpState>,
|
||||
auth: Option<Extension<RuntimeAuthContext>>,
|
||||
Path(worker_id): Path<String>,
|
||||
Query(query): Query<RuntimeHttpUploadFileQuery>,
|
||||
body: Bytes,
|
||||
) -> RestResult<RuntimeHttpUploadedFileResponse> {
|
||||
let worker_ref = worker_ref_for(&state.runtime, worker_id)?;
|
||||
let file = match auth_workspace_scope(&state, auth.as_ref())? {
|
||||
Some(scope) => state.runtime.upload_worker_file_scoped(
|
||||
&scope,
|
||||
&worker_ref,
|
||||
&query.file_name,
|
||||
&query.media_type,
|
||||
&body,
|
||||
),
|
||||
None => state.runtime.upload_worker_file(
|
||||
&worker_ref,
|
||||
&query.file_name,
|
||||
&query.media_type,
|
||||
&body,
|
||||
),
|
||||
}
|
||||
.map_err(RuntimeHttpRestError::runtime)?;
|
||||
Ok(Json(RuntimeHttpUploadedFileResponse { file }))
|
||||
}
|
||||
|
||||
async fn delete_worker_uploaded_file(
|
||||
State(state): State<RuntimeHttpState>,
|
||||
auth: Option<Extension<RuntimeAuthContext>>,
|
||||
Path((worker_id, artifact_id)): Path<(String, String)>,
|
||||
) -> RestResult<RuntimeHttpUploadedFileDeleteResponse> {
|
||||
let worker_ref = worker_ref_for(&state.runtime, worker_id)?;
|
||||
match auth_workspace_scope(&state, auth.as_ref())? {
|
||||
Some(scope) => {
|
||||
state
|
||||
.runtime
|
||||
.delete_worker_uploaded_file_scoped(&scope, &worker_ref, &artifact_id)
|
||||
}
|
||||
None => state
|
||||
.runtime
|
||||
.delete_worker_uploaded_file(&worker_ref, &artifact_id),
|
||||
}
|
||||
.map_err(RuntimeHttpRestError::runtime)?;
|
||||
Ok(Json(RuntimeHttpUploadedFileDeleteResponse {
|
||||
deleted: true,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn stop_worker(
|
||||
State(state): State<RuntimeHttpState>,
|
||||
auth: Option<Extension<RuntimeAuthContext>>,
|
||||
@@ -1621,7 +1697,7 @@ fn required_runtime_permission(method: &Method, path: &str) -> Option<&'static s
|
||||
if path.ends_with("/workspace-api") {
|
||||
return Some("workers:create");
|
||||
}
|
||||
if path.ends_with("/input") || path.ends_with("/restore") {
|
||||
if path.ends_with("/input") || path.ends_with("/restore") || path.contains("/attachments") {
|
||||
return Some("workers:input");
|
||||
}
|
||||
if path.ends_with("/stop") || path.ends_with("/cancel") {
|
||||
@@ -1882,6 +1958,21 @@ mod tests {
|
||||
WorkdirPath, WorkdirSessionCapabilities,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn attachment_routes_require_worker_input_permission() {
|
||||
assert_eq!(
|
||||
required_runtime_permission(&Method::POST, "/v1/workers/7/attachments"),
|
||||
Some("workers:input")
|
||||
);
|
||||
assert_eq!(
|
||||
required_runtime_permission(
|
||||
&Method::DELETE,
|
||||
"/v1/workers/7/attachments/019ca7c8-57b6-7f05-8edf-524147aba7b3"
|
||||
),
|
||||
Some("workers:input")
|
||||
);
|
||||
}
|
||||
|
||||
fn test_bundle(profile: ProfileSelector) -> ConfigBundle {
|
||||
ConfigBundle {
|
||||
metadata: ConfigBundleMetadata {
|
||||
|
||||
@@ -1205,6 +1205,103 @@ impl Runtime {
|
||||
})
|
||||
}
|
||||
|
||||
/// Store a client-local file in the owning Worker session before input submit.
|
||||
pub fn upload_worker_file_scoped(
|
||||
&self,
|
||||
scope: &RuntimeWorkspaceScope,
|
||||
worker_ref: &WorkerRef,
|
||||
file_name: &str,
|
||||
media_type: &str,
|
||||
content: &[u8],
|
||||
) -> Result<protocol::UploadedFileRef, RuntimeError> {
|
||||
self.ensure_worker_in_workspace(scope, worker_ref)?;
|
||||
self.upload_worker_file(worker_ref, file_name, media_type, content)
|
||||
}
|
||||
|
||||
pub fn upload_worker_file(
|
||||
&self,
|
||||
worker_ref: &WorkerRef,
|
||||
file_name: &str,
|
||||
media_type: &str,
|
||||
content: &[u8],
|
||||
) -> Result<protocol::UploadedFileRef, RuntimeError> {
|
||||
let (backend, handle) = {
|
||||
let state = self.lock()?;
|
||||
state.ensure_running()?;
|
||||
state.ensure_worker_ref(worker_ref)?;
|
||||
let worker = state.worker(worker_ref)?;
|
||||
match (
|
||||
state.execution_backend.clone(),
|
||||
worker.execution_handle.clone(),
|
||||
) {
|
||||
(Some(backend), Some(handle)) => (backend, handle),
|
||||
_ => {
|
||||
return Err(RuntimeError::WorkerExecutionUnavailable {
|
||||
worker_id: worker_ref.worker_id.clone(),
|
||||
message: "worker has no live execution handle".to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
};
|
||||
backend
|
||||
.upload_file(&handle, file_name, media_type, content)
|
||||
.map_err(|result| RuntimeError::WorkerExecutionRejected {
|
||||
worker_id: worker_ref.worker_id.clone(),
|
||||
operation: result.operation,
|
||||
outcome: result.outcome,
|
||||
message: result.message_or_default(),
|
||||
result,
|
||||
})
|
||||
}
|
||||
|
||||
/// Delete an unsubmitted uploaded file from the owning Worker session.
|
||||
pub fn delete_worker_uploaded_file_scoped(
|
||||
&self,
|
||||
scope: &RuntimeWorkspaceScope,
|
||||
worker_ref: &WorkerRef,
|
||||
artifact_id: &str,
|
||||
) -> Result<(), RuntimeError> {
|
||||
self.ensure_worker_in_workspace(scope, worker_ref)?;
|
||||
self.delete_worker_uploaded_file(worker_ref, artifact_id)
|
||||
}
|
||||
|
||||
pub fn delete_worker_uploaded_file(
|
||||
&self,
|
||||
worker_ref: &WorkerRef,
|
||||
artifact_id: &str,
|
||||
) -> Result<(), RuntimeError> {
|
||||
let (backend, handle) = {
|
||||
let state = self.lock()?;
|
||||
state.ensure_running()?;
|
||||
state.ensure_worker_ref(worker_ref)?;
|
||||
let worker = state.worker(worker_ref)?;
|
||||
match (
|
||||
state.execution_backend.clone(),
|
||||
worker.execution_handle.clone(),
|
||||
) {
|
||||
(Some(backend), Some(handle)) => (backend, handle),
|
||||
_ => {
|
||||
return Err(RuntimeError::WorkerExecutionUnavailable {
|
||||
worker_id: worker_ref.worker_id.clone(),
|
||||
message: "worker has no live execution handle".to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
};
|
||||
let result = backend.delete_uploaded_file(&handle, artifact_id);
|
||||
if result.is_accepted() {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(RuntimeError::WorkerExecutionRejected {
|
||||
worker_id: worker_ref.worker_id.clone(),
|
||||
operation: result.operation,
|
||||
outcome: result.outcome,
|
||||
message: result.message_or_default(),
|
||||
result,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Return live completion entries through a workspace-scoped Runtime authorization context.
|
||||
pub fn worker_completions_scoped(
|
||||
&self,
|
||||
@@ -1387,6 +1484,8 @@ impl Runtime {
|
||||
WorkerExecutionOperation::Spawn
|
||||
| WorkerExecutionOperation::Restore
|
||||
| WorkerExecutionOperation::Input
|
||||
| WorkerExecutionOperation::UploadFile
|
||||
| WorkerExecutionOperation::DeleteUploadedFile
|
||||
| WorkerExecutionOperation::ProtocolMethod => return Ok(()),
|
||||
};
|
||||
if result.is_accepted() {
|
||||
|
||||
@@ -2025,6 +2025,51 @@ where
|
||||
result
|
||||
}
|
||||
|
||||
fn upload_file(
|
||||
&self,
|
||||
handle: &WorkerExecutionHandle,
|
||||
file_name: &str,
|
||||
media_type: &str,
|
||||
content: &[u8],
|
||||
) -> Result<protocol::UploadedFileRef, WorkerExecutionResult> {
|
||||
let (worker, _, _) = self.get_execution(handle).map_err(|mut result| {
|
||||
result.operation = WorkerExecutionOperation::UploadFile;
|
||||
result
|
||||
})?;
|
||||
worker
|
||||
.upload_file(file_name, media_type, content)
|
||||
.map_err(|error| {
|
||||
WorkerExecutionResult::rejected(
|
||||
WorkerExecutionOperation::UploadFile,
|
||||
format!("uploaded_file_rejected: {error}"),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn delete_uploaded_file(
|
||||
&self,
|
||||
handle: &WorkerExecutionHandle,
|
||||
artifact_id: &str,
|
||||
) -> WorkerExecutionResult {
|
||||
let (worker, _, _) = match self.get_execution(handle) {
|
||||
Ok(execution) => execution,
|
||||
Err(mut result) => {
|
||||
result.operation = WorkerExecutionOperation::DeleteUploadedFile;
|
||||
return result;
|
||||
}
|
||||
};
|
||||
match worker.delete_uploaded_file(artifact_id) {
|
||||
Ok(_) => WorkerExecutionResult::accepted(
|
||||
WorkerExecutionOperation::DeleteUploadedFile,
|
||||
WorkerExecutionRunState::Idle,
|
||||
),
|
||||
Err(error) => WorkerExecutionResult::rejected(
|
||||
WorkerExecutionOperation::DeleteUploadedFile,
|
||||
format!("uploaded_file_delete_rejected: {error}"),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch_method(
|
||||
&self,
|
||||
handle: &WorkerExecutionHandle,
|
||||
|
||||
Reference in New Issue
Block a user