fix: preserve authenticated submit source

crates/protocol/src/lib.rs

@@ -32,6 +32,13 @@ fn is_false(value: &bool) -> bool {
 // Method (Client → Worker via Unix Socket)
 // ---------------------------------------------------------------------------
 
+/// Trusted Server → Runtime transport header carrying the authenticated
+/// browser Account identity for one Worker protocol connection.
+///
+/// Runtime accepts this only after its normal HTTP authentication succeeds;
+/// serialized [`Method`] payloads cannot set authenticated source identity.
+pub const AUTHENTICATED_ACCOUNT_ID_HEADER: &str = "x-yoi-authenticated-account-id";
+
 /// Trusted source identity attached by an authenticated transport boundary.
 ///
 /// Public clients cannot select this value directly. Runtime/Backend adapters

crates/worker-runtime/src/http_server.rs

@@ -33,7 +33,7 @@ use axum::extract::rejection::{JsonRejection, QueryRejection};
 #[cfg(feature = "ws-server")]
 use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade};
 use axum::extract::{DefaultBodyLimit, Extension, Path, Query, State};
-use axum::http::{Method, Request, StatusCode, header};
+use axum::http::{HeaderMap, Method, Request, StatusCode, header};
 use axum::middleware::{self, Next};
 use axum::response::{IntoResponse, Response};
 use axum::routing::{delete, get, post};
@@ -1184,10 +1184,12 @@ async fn worker_protocol_ws(
     auth: Option<Extension<RuntimeAuthContext>>,
     Path(worker_id): Path<String>,
     Query(query): Query<RuntimeWorkerEventsWsQuery>,
+    headers: HeaderMap,
     ws: WebSocketUpgrade,
 ) -> Result<Response, RuntimeHttpRestError> {
     let worker_ref = worker_ref_for(&state.runtime, worker_id)?;
     let scope = auth_workspace_scope(&state, auth.as_ref())?;
+    let input_source = authenticated_protocol_input_source(&headers)?;
     match scope.as_ref() {
         Some(scope) => state
             .runtime
@@ -1198,22 +1200,60 @@ async fn worker_protocol_ws(
     .map_err(RuntimeHttpRestError::runtime)?;
     Ok(ws
         .on_upgrade(move |socket| {
-            worker_protocol_ws_session(state.runtime, scope, worker_ref, query, socket)
+            worker_protocol_ws_session(
+                state.runtime,
+                scope,
+                worker_ref,
+                query,
+                input_source,
+                socket,
+            )
         })
         .into_response())
 }
 
 #[cfg(feature = "ws-server")]
-fn authorize_runtime_protocol_method(method: protocol::Method) -> protocol::Method {
+fn authenticated_protocol_input_source(
+    headers: &HeaderMap,
+) -> Result<Option<protocol::AuthenticatedInputSource>, RuntimeHttpRestError> {
+    let Some(value) = headers.get(protocol::AUTHENTICATED_ACCOUNT_ID_HEADER) else {
+        return Ok(None);
+    };
+    let account_id = value.to_str().map_err(|_| {
+        RuntimeHttpRestError::new(
+            StatusCode::BAD_REQUEST,
+            "authenticated_input_source_invalid",
+            "authenticated Worker input source is invalid",
+        )
+    })?;
+    if account_id.trim().is_empty() || account_id.len() > 128 {
+        return Err(RuntimeHttpRestError::new(
+            StatusCode::BAD_REQUEST,
+            "authenticated_input_source_invalid",
+            "authenticated Worker input source is invalid",
+        ));
+    }
+    Ok(Some(protocol::AuthenticatedInputSource::Account {
+        account_id: account_id.to_owned(),
+    }))
+}
+
+#[cfg(feature = "ws-server")]
+fn authorize_runtime_protocol_method(
+    method: protocol::Method,
+    transport_source: Option<&protocol::AuthenticatedInputSource>,
+) -> protocol::Method {
     match method {
         protocol::Method::SubmitTracked {
             submission_request_id,
             input,
             ..
         } => protocol::Method::SubmitTracked {
-            source: protocol::AuthenticatedInputSource::Backend {
+            source: transport_source.cloned().unwrap_or_else(|| {
-                operation_id: submission_request_id.clone(),
+                protocol::AuthenticatedInputSource::Backend {
-            },
+                    operation_id: submission_request_id.clone(),
+                }
+            }),
             submission_request_id,
             input,
         },
@@ -1223,9 +1263,11 @@ fn authorize_runtime_protocol_method(method: protocol::Method) -> protocol::Meth
             auto_run,
             ..
         } => protocol::Method::NotifyTracked {
-            source: protocol::AuthenticatedInputSource::Backend {
+            source: transport_source.cloned().unwrap_or_else(|| {
-                operation_id: notification_request_id.clone(),
+                protocol::AuthenticatedInputSource::Backend {
-            },
+                    operation_id: notification_request_id.clone(),
+                }
+            }),
             notification_request_id,
             message,
             auto_run,
@@ -1240,6 +1282,7 @@ async fn worker_protocol_ws_session(
     scope: Option<RuntimeWorkspaceScope>,
     worker_ref: WorkerRef,
     query: RuntimeWorkerEventsWsQuery,
+    input_source: Option<protocol::AuthenticatedInputSource>,
     mut socket: WebSocket,
 ) {
     let mut cursor = match query.cursor.as_deref() {
@@ -1322,7 +1365,8 @@ async fn worker_protocol_ws_session(
                 match inbound {
                     Some(Ok(WsMessage::Text(text))) => match decode_method(&text) {
                         Ok(method) => {
-                            let method = authorize_runtime_protocol_method(method);
+                            let method =
+                                authorize_runtime_protocol_method(method, input_source.as_ref());
                             let result = match scope.as_ref() {
                                 Some(scope) => {
                                     runtime.send_protocol_method_scoped(scope, &worker_ref, method)
@@ -2135,7 +2179,7 @@ mod tests {
             }
         ));
         assert!(matches!(
-            authorize_runtime_protocol_method(decoded),
+            authorize_runtime_protocol_method(decoded, None),
             protocol::Method::SubmitTracked {
                 source: protocol::AuthenticatedInputSource::Backend { operation_id },
                 ..
@@ -2143,6 +2187,36 @@ mod tests {
         ));
     }
 
+    #[test]
+    fn runtime_protocol_uses_transport_authenticated_account_source() {
+        let mut headers = HeaderMap::new();
+        headers.insert(
+            protocol::AUTHENTICATED_ACCOUNT_ID_HEADER,
+            "account-1".parse().unwrap(),
+        );
+        let source = authenticated_protocol_input_source(&headers)
+            .unwrap()
+            .expect("account source header must resolve");
+        let wire = serde_json::to_string(&protocol::Method::NotifyTracked {
+            notification_request_id: "notification-1".into(),
+            message: "hello".into(),
+            auto_run: true,
+            source: protocol::AuthenticatedInputSource::Account {
+                account_id: "forged".into(),
+            },
+        })
+        .unwrap();
+        let decoded: protocol::Method = serde_json::from_str(&wire).unwrap();
+
+        assert!(matches!(
+            authorize_runtime_protocol_method(decoded, Some(&source)),
+            protocol::Method::NotifyTracked {
+                source: protocol::AuthenticatedInputSource::Account { account_id },
+                ..
+            } if account_id == "account-1"
+        ));
+    }
+
     #[test]
     fn attachment_routes_require_worker_input_permission() {
         assert_eq!(

crates/worker/src/worker.rs

@@ -1290,10 +1290,31 @@ where
         pending: &PendingSubmission,
     ) -> Result<(), PendingSubmissionError> {
         let session_id = self.writer.state.location().session_id;
+        let mut pinned = Vec::new();
         for reference in submission_uploaded_file_refs(&pending.input) {
-            self.writer
+            if pinned.iter().any(|existing: &protocol::UploadedFileRef| {
-                .store
+                existing.artifact_id == reference.artifact_id
-                .pin_uploaded_file(session_id, reference, &pending.submission_id)?;
+            }) {
+                continue;
+            }
+            if let Err(pin_error) =
+                self.writer
+                    .store
+                    .pin_uploaded_file(session_id, reference, &pending.submission_id)
+            {
+                let mut rollback_error = None;
+                for acquired in pinned.iter().rev() {
+                    if let Err(error) = self.writer.store.release_uploaded_file_pin(
+                        session_id,
+                        &acquired.artifact_id,
+                        &pending.submission_id,
+                    ) {
+                        rollback_error.get_or_insert(error);
+                    }
+                }
+                return Err(rollback_error.unwrap_or(pin_error).into());
+            }
+            pinned.push(reference.clone());
         }
         Ok(())
     }
@@ -1303,14 +1324,28 @@ where
         pending: &PendingSubmission,
     ) -> Result<(), PendingSubmissionError> {
         let session_id = self.writer.state.location().session_id;
+        let mut released = Vec::new();
+        let mut first_error = None;
         for reference in submission_uploaded_file_refs(&pending.input) {
-            self.writer.store.release_uploaded_file_pin(
+            if released
+                .iter()
+                .any(|artifact_id: &String| artifact_id == &reference.artifact_id)
+            {
+                continue;
+            }
+            if let Err(error) = self.writer.store.release_uploaded_file_pin(
                 session_id,
                 &reference.artifact_id,
                 &pending.submission_id,
-            )?;
+            ) {
+                first_error.get_or_insert(error);
+            }
+            released.push(reference.artifact_id.clone());
+        }
+        match first_error {
+            Some(error) => Err(error.into()),
+            None => Ok(()),
         }
-        Ok(())
     }
 
     #[cfg(test)]
@@ -9668,6 +9703,76 @@ mod build_summary_prompt_tests {
         assert_eq!(state.pending[1].provenance, account_b);
     }
 
+    #[test]
+    fn rejected_submission_rolls_back_uploaded_file_pins_acquired_before_conflict() {
+        let temp = tempfile::tempdir().unwrap();
+        let handle = PendingSubmissionHandle::for_test(temp.path());
+        let session_id = handle.writer.state.session_id();
+        let limits = session_store::UploadedFileLimits {
+            max_file_bytes: 1024,
+            max_session_bytes: 2048,
+        };
+        let first = handle
+            .writer
+            .store
+            .write_uploaded_file(session_id, "first.txt", "text/plain", b"first", limits)
+            .unwrap();
+        let second = handle
+            .writer
+            .store
+            .write_uploaded_file(session_id, "second.txt", "text/plain", b"second", limits)
+            .unwrap();
+        handle
+            .writer
+            .store
+            .pin_uploaded_file(session_id, &second, "other-submission")
+            .unwrap();
+
+        assert!(
+            handle
+                .accept(
+                    "request-partial-pin".into(),
+                    vec![
+                        Segment::UploadedFile {
+                            file: first.clone(),
+                        },
+                        Segment::UploadedFile {
+                            file: second.clone(),
+                        },
+                    ],
+                    false,
+                )
+                .is_err()
+        );
+        assert!(handle.snapshot().submissions.is_empty());
+        assert!(
+            handle
+                .writer
+                .store
+                .delete_uploaded_file(session_id, &first.artifact_id)
+                .unwrap()
+        );
+        assert!(matches!(
+            handle
+                .writer
+                .store
+                .delete_uploaded_file(session_id, &second.artifact_id),
+            Err(StoreError::ArtifactAlreadyCommitted)
+        ));
+        handle
+            .writer
+            .store
+            .release_uploaded_file_pin(session_id, &second.artifact_id, "other-submission")
+            .unwrap();
+        assert!(
+            handle
+                .writer
+                .store
+                .delete_uploaded_file(session_id, &second.artifact_id)
+                .unwrap()
+        );
+    }
+
     #[test]
     fn queued_submission_pins_uploaded_file_until_cancelled() {
         let temp = tempfile::tempdir().unwrap();

crates/workspace-server/src/server.rs

@@ -9251,7 +9251,7 @@ async fn scoped_capture_worker_observation_session(
         return Err(ApiError::from(Error::UnknownWorker { worker: target }));
     }
 
-    let mut connection = connect_workspace_worker_protocol(&api, &target).await?;
+    let mut connection = connect_workspace_worker_protocol(&api, &target, None).await?;
     let event = tokio::time::timeout(std::time::Duration::from_secs(10), connection.events.recv())
         .await
         .map_err(|_| {
@@ -13955,6 +13955,7 @@ pub(crate) struct WorkspaceWorkerProtocolConnection {
 pub(crate) async fn connect_workspace_worker_protocol(
     api: &WorkspaceApi,
     worker: &RuntimeWorkerRef,
+    input_source: Option<&protocol::AuthenticatedInputSource>,
 ) -> Result<WorkspaceWorkerProtocolConnection> {
     let source = match api.observation_proxy.source(worker) {
         Ok(source) => source,
@@ -13971,15 +13972,39 @@ pub(crate) async fn connect_workspace_worker_protocol(
         }
     };
     match source {
-        RuntimeObservationSource::RemoteWs(config) => connect_remote_worker_protocol(config).await,
+        RuntimeObservationSource::RemoteWs(config) => {
+            connect_remote_worker_protocol(config, input_source).await
+        }
         RuntimeObservationSource::Embedded(source) => {
             connect_embedded_worker_protocol(source).await
         }
     }
 }
 
+fn insert_authenticated_input_source_header(
+    headers: &mut HeaderMap,
+    input_source: Option<&protocol::AuthenticatedInputSource>,
+) -> Result<()> {
+    let Some(input_source) = input_source else {
+        return Ok(());
+    };
+    let protocol::AuthenticatedInputSource::Account { account_id } = input_source else {
+        return Err(Error::Config(
+            "remote Worker protocol transport supports only Account input source".into(),
+        ));
+    };
+    headers.insert(
+        protocol::AUTHENTICATED_ACCOUNT_ID_HEADER,
+        account_id.parse().map_err(|error| {
+            Error::Config(format!("invalid authenticated Account identity: {error}"))
+        })?,
+    );
+    Ok(())
+}
+
 async fn connect_remote_worker_protocol(
     config: RuntimeObservationSourceConfig,
+    input_source: Option<&protocol::AuthenticatedInputSource>,
 ) -> Result<WorkspaceWorkerProtocolConnection> {
     let mut request = config
         .endpoint
@@ -13994,6 +14019,7 @@ async fn connect_remote_worker_protocol(
             })?,
         );
     }
+    insert_authenticated_input_source_header(request.headers_mut(), input_source)?;
     let (socket, _) =
         connect_async(request)
             .await
@@ -14117,6 +14143,16 @@ async fn remote_worker_protocol_ws_session(
             }
         }
     }
+    if let Err(error) =
+        insert_authenticated_input_source_header(request.headers_mut(), Some(&input_source))
+    {
+        let mut socket = socket;
+        let event = protocol_error_event(format!(
+            "failed to build authenticated Account identity header: {error}"
+        ));
+        let _ = send_protocol_event(&mut socket, &event).await;
+        return;
+    }
 
     let (upstream, _) = match connect_async(request).await {
         Ok(connection) => connection,
@@ -16678,6 +16714,25 @@ mod tests {
         assert!(authorize_browser_worker_method(method, &source).is_err());
     }
 
+    #[test]
+    fn remote_worker_protocol_header_preserves_authenticated_account_source() {
+        let mut headers = HeaderMap::new();
+        insert_authenticated_input_source_header(
+            &mut headers,
+            Some(&protocol::AuthenticatedInputSource::Account {
+                account_id: "account-1".into(),
+            }),
+        )
+        .unwrap();
+
+        assert_eq!(
+            headers
+                .get(protocol::AUTHENTICATED_ACCOUNT_ID_HEADER)
+                .unwrap(),
+            "account-1"
+        );
+    }
+
     #[test]
     fn merge_request_http_paths_observe_refs_through_runtime_provider_authority() {
         let source = include_str!("server.rs");

crates/workspace-server/src/workspace_subscription.rs

@@ -91,7 +91,13 @@ pub(crate) async fn serve_workspace_subscription(
                                 runtime_id: Some(runtime_id),
                             } => {
                                 let worker = RuntimeWorkerRef::new(&runtime_id, worker_id.as_str());
-                                match connect_workspace_worker_protocol(&api, &worker).await {
+                                match connect_workspace_worker_protocol(
+                                    &api,
+                                    &worker,
+                                    Some(&input_source),
+                                )
+                                .await
+                                {
                                     Ok(connection) => {
                                         let methods = connection.methods.clone();
                                         let task = tokio::spawn(run_worker_protocol(