diff --git a/.repository-projection.json b/.repository-projection.json index 5d3dff09b..7e6fb8001 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "deixic-code", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "0dda4b314d332ad6f95f189d10b066d1a080ee34", + "sourceSha": "373dd517beb6babba8f831ede8501edef7cd0520", "destinationRepository": "dx-corp/code", - "priorProjectedBase": "8991ebcc446033952c75e0858f46245ab8b0580a", + "priorProjectedBase": "d30b1b01dee74cc2a51d546423e51e80ed66f31f", "definitionDigest": "82936441c776e3e8edb5d215a75007ec9714a233f489d460075d79d5ef5ba32f", "toolDigest": "c244d99199a7ae3eb8ff644a99462163c23b0bb6a83ef50af01efbdca0b81d04", - "contentDigest": "93b1fddc5de9161c81eb617fc50c86e6f9bb877a46f1a1206a46d8732b6b39f8", + "contentDigest": "67278d7f71d18a40a13eb757faff25c35ae2c1e05c6f963b74ac7d6da00953b9", "publicationEligible": true } diff --git a/packages/local-host-rs/src/hosted_runner/tests.rs b/packages/local-host-rs/src/hosted_runner/tests.rs index f1ca95681..82b7ec294 100644 --- a/packages/local-host-rs/src/hosted_runner/tests.rs +++ b/packages/local-host-rs/src/hosted_runner/tests.rs @@ -16,6 +16,10 @@ use crate::hosted_runner::rendezvous_protocol::RendezvousMode; #[cfg(unix)] mod deferred_rejection; +#[cfg(unix)] +mod replay_ownership; +#[cfg(unix)] +use replay_ownership::governed_response_for_ack_test; mod env_aliases; mod transport_flush; @@ -11281,105 +11285,6 @@ async fn queued_response_restarts_pending_and_consumes_once_after_child_exit() { .shutdown(); } -#[cfg(unix)] -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn correlated_protocol_rejection_rolls_back_ownership_and_allows_retry() { - let workspace = tempdir().expect("workspace"); - let fixtures = tempdir().expect("fixtures"); - let log_path = fixtures.path().join("rejected-responses.log"); - let script = create_reject_then_accept_script(fixtures.path(), &log_path, None); - let supervisor = connected_supervisor_for_script(&script).await; - let executor = Arc::new(AgentSupervisorHostedRunnerMessageExecutor::new(Arc::clone( - &supervisor, - ))); - let handle = start_hosted_runner_with_message_executor( - test_config(workspace.path().to_path_buf()), - executor.clone(), - ) - .await - .expect("hosted runner"); - let client = reqwest::Client::new(); - let (capability, subscription_id) = - attach_thread_controller(&client, &handle.base_url(), "conn_rejection").await; - let headers = response_headers( - "conn_rejection", - &subscription_id, - &capability, - "rejected-key", - ); - let response = ToAgentMessage::ToolResponse { - call_id: "retry-call".to_string(), - tool_execution_id: Some("retry-execution".to_string()), - approved: true, - result: None, - }; - - let error = match handle_message( - handle.shared.clone(), - "sess_test", - headers.clone(), - response.clone(), - ) - .await - { - Err(error) => error, - Ok(_) => panic!("correlated protocol rejection must not return success"), - }; - assert_eq!(error.code, HostedRunnerErrorCode::RuntimeFailed); - assert!(error.message.contains("not awaiting a decision")); - { - let state = handle - .shared - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - assert!( - !state - .pending_response_idempotency - .contains_key("rejected-key") - ); - } - assert!( - !executor - .queued_responses - .lock() - .expect("queued responses") - .contains_key("rejected-key") - ); - assert!( - !load_executor_response_ledger(workspace.path(), "sess_test") - .expect("response ledger") - .iter() - .any(|(key, _)| key == "rejected-key") - ); - - handle_message( - handle.shared.clone(), - "sess_test", - headers.clone(), - response.clone(), - ) - .await - .expect("corrected retry dispatches"); - let replay = handle_message(handle.shared.clone(), "sess_test", headers, response) - .await - .expect("accepted retry replays"); - let ResponseBody::Json { body, .. } = replay else { - panic!("accepted retry replay must return JSON"); - }; - assert_eq!(body["replayed"], true); - assert_eq!( - std::fs::read_to_string(&log_path) - .expect("rejected response log") - .lines() - .count(), - 2, - "one rejected dispatch and one corrected dispatch are expected" - ); - handle.shutdown().await; - supervisor.lock().expect("supervisor").shutdown(); -} - #[cfg(unix)] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn rejected_response_releases_ownership_when_ledger_cleanup_fails() { @@ -11585,234 +11490,6 @@ async fn unmatched_governed_ack_does_not_finalize_or_redispatch() { supervisor.lock().expect("supervisor").shutdown(); } -#[cfg(unix)] -async fn assert_unique_protocol_request_owner_across_restart( - message: ToAgentMessage, - message_type: &str, - request_id: &str, -) { - let workspace = tempdir().expect("workspace"); - let fixtures = tempdir().expect("fixtures"); - let first_log = fixtures.path().join(format!("{message_type}-first.log")); - let first_script = create_delayed_identity_ack_script( - fixtures.path(), - &format!("{message_type}-first.sh"), - &first_log, - message_type, - request_id, - ); - let first_supervisor = connected_supervisor_for_script(&first_script).await; - let first_executor = Arc::new(AgentSupervisorHostedRunnerMessageExecutor::new(Arc::clone( - &first_supervisor, - ))); - let first = start_hosted_runner_with_message_executor( - test_config(workspace.path().to_path_buf()), - first_executor, - ) - .await - .expect("first hosted runner"); - let client = reqwest::Client::new(); - let (capability, subscription_id) = - attach_thread_controller(&client, &first.base_url(), "conn_identity_first").await; - let owner_headers = response_headers( - "conn_identity_first", - &subscription_id, - &capability, - "identity-owner-key", - ); - let competing_headers = response_headers( - "conn_identity_first", - &subscription_id, - &capability, - "identity-competing-key", - ); - - handle_message( - first.shared.clone(), - "sess_test", - owner_headers.clone(), - message.clone(), - ) - .await - .expect("owner response queues"); - let conflict = match handle_message( - first.shared.clone(), - "sess_test", - competing_headers.clone(), - message.clone(), - ) - .await - { - Err(error) => error, - Ok(_) => panic!("second key must not own the same protocol request"), - }; - assert_eq!(conflict.code, HostedRunnerErrorCode::IdempotencyConflict); - tokio::time::sleep(Duration::from_millis(300)).await; - let replay = handle_message( - first.shared.clone(), - "sess_test", - owner_headers, - message.clone(), - ) - .await - .expect("owner key replays after delayed acknowledgement"); - let ResponseBody::Json { body, .. } = replay else { - panic!("owner replay must return JSON"); - }; - assert_eq!(body["replayed"], true); - assert_eq!( - std::fs::read_to_string(&first_log) - .expect("first child log") - .lines() - .count(), - 1 - ); - first.shutdown().await; - first_supervisor - .lock() - .expect("first supervisor") - .shutdown(); - - let second_log = fixtures.path().join(format!("{message_type}-second.log")); - let second_script = create_delayed_identity_ack_script( - fixtures.path(), - &format!("{message_type}-second.sh"), - &second_log, - message_type, - request_id, - ); - let second_supervisor = connected_supervisor_for_script(&second_script).await; - let second_executor = Arc::new(AgentSupervisorHostedRunnerMessageExecutor::new(Arc::clone( - &second_supervisor, - ))); - let second = start_hosted_runner_with_message_executor( - test_config(workspace.path().to_path_buf()), - second_executor, - ) - .await - .expect("restarted hosted runner"); - let second_connection_id = if matches!( - &message, - ToAgentMessage::GovernedClientToolResult(tool_wire::GovernedClientToolResult { .. }) - ) { - "conn_identity_first" - } else { - "conn_identity_second" - }; - let (capability, subscription_id) = - attach_thread_controller(&client, &second.base_url(), second_connection_id).await; - let replay = handle_message( - second.shared.clone(), - "sess_test", - response_headers( - second_connection_id, - &subscription_id, - &capability, - "identity-owner-key", - ), - message.clone(), - ) - .await - .expect("durable owner key replays"); - let ResponseBody::Json { body, .. } = replay else { - panic!("durable owner replay must return JSON"); - }; - assert_eq!(body["replayed"], true); - let conflict = match handle_message( - second.shared.clone(), - "sess_test", - response_headers( - second_connection_id, - &subscription_id, - &capability, - "identity-competing-key", - ), - message, - ) - .await - { - Err(error) => error, - Ok(_) => panic!("request ownership must survive restart"), - }; - assert_eq!(conflict.code, HostedRunnerErrorCode::IdempotencyConflict); - assert!( - !second_log.exists(), - "restart must not redispatch either key" - ); - second.shutdown().await; - second_supervisor - .lock() - .expect("second supervisor") - .shutdown(); -} - -#[cfg(unix)] -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn delayed_tool_response_has_one_idempotency_owner_across_restart() { - assert_unique_protocol_request_owner_across_restart( - ToAgentMessage::ToolResponse { - call_id: "unique-tool-call".to_string(), - tool_execution_id: Some("unique-tool-execution".to_string()), - approved: true, - result: None, - }, - "tool_response", - "unique-tool-call", - ) - .await; -} - -#[cfg(unix)] -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn delayed_server_request_response_has_one_idempotency_owner_across_restart() { - assert_unique_protocol_request_owner_across_restart( - ToAgentMessage::ServerRequestResponse { - request_id: "unique-server-request".to_string(), - request_type: ServerRequestType::UserInput, - approved: None, - result: None, - content: Some(Vec::new()), - is_error: Some(false), - decision_action: None, - reason: Some("answer".to_string()), - }, - "server_request_response", - "unique-server-request", - ) - .await; -} - -#[cfg(unix)] -fn governed_response_for_ack_test() -> ToAgentMessage { - ToAgentMessage::GovernedClientToolResult(tool_wire::GovernedClientToolResult { - process_tool_cost_micros: None, - call_id: "unique-governed-call".to_string(), - content: Vec::new(), - is_error: false, - tool_execution_id: "unique-governed-execution".to_string(), - client_instance_id: "conn_identity_first".to_string(), - grant_id: "grant-1".to_string(), - grant_version: 1, - grant_hash: "hash".to_string(), - turn_digest: "turn-digest".to_string(), - definition_digest: "definition-digest".to_string(), - args_digest: "args-digest".to_string(), - owner_lease_epoch: 1, - idempotency_key: "identity-owner-key".to_string(), - }) -} - -#[cfg(unix)] -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn delayed_governed_result_has_one_idempotency_owner_across_restart() { - assert_unique_protocol_request_owner_across_restart( - governed_response_for_ack_test(), - "governed_client_tool_result", - "unique-governed-call", - ) - .await; -} - #[cfg(unix)] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn delayed_ack_event_pump_finalizes_before_restart_without_redispatch() { diff --git a/packages/local-host-rs/src/hosted_runner/tests/replay_ownership.rs b/packages/local-host-rs/src/hosted_runner/tests/replay_ownership.rs new file mode 100644 index 000000000..ab904597b --- /dev/null +++ b/packages/local-host-rs/src/hosted_runner/tests/replay_ownership.rs @@ -0,0 +1,455 @@ +use super::*; + +// Reuse the fixture's bounded condition wait, rather than assuming elapsed +// time means the child and event pump have persisted consumed ownership. +async fn await_acknowledged_owner( + handle: &HostedRunnerHandle, + workspace: &Path, + key: &str, + request_id: &str, +) { + let journal_path = std::fs::read_dir(workspace.join(".maestro/hosted-runner/threads")) + .expect("thread journal directory") + .map(|entry| entry.expect("thread journal entry").path()) + .find(|path| { + path.extension() + .is_some_and(|extension| extension == "json") + }) + .expect("thread journal JSON path"); + wait_for_condition(|| { + let completed = { + let state = handle + .shared + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.response_idempotency_keys.contains(key) + && !state.pending_response_idempotency.contains_key(key) + && state + .response_request_owners + .get(request_id) + .map(String::as_str) + == Some(key) + }; + let ledger_consumed = load_executor_response_ledger(workspace, "sess_test") + .expect("executor response ledger") + .iter() + .any(|(owner, consumed)| owner == key && *consumed); + let journal: serde_json::Value = + serde_json::from_slice(&std::fs::read(&journal_path).expect("durable thread journal")) + .expect("thread journal JSON"); + completed + && journal["thread_id"] == "sess_test" + && ledger_consumed + && journal["response_idempotency_keys"] + .as_array() + .is_some_and(|keys| keys.contains(&json!(key))) + && journal["pending_response_idempotency"] + .as_object() + .is_some_and(|pending| !pending.contains_key(key)) + && journal["response_request_owners"][request_id].as_str() == Some(key) + }) + .await; +} + +fn gated_identity_ack_script( + directory: &Path, + name: &str, + log_path: &Path, + message_type: &str, + request_id: &str, +) -> (PathBuf, PathBuf) { + let script_path = directory.join(name); + let release = directory.join(format!("{name}.release")); + std::fs::write( + &script_path, + format!( + r#"#!/bin/sh +printf '{{"type":"ready","model":"test","provider":"test"}}\n' +while IFS= read -r line; do + case "$line" in + *'"type":"{message_type}"'*) + printf '%s\n' "$line" >> "{}" + while [ ! -f '{}' ]; do sleep 0.01; done + printf '{{"type":"response_accepted","request_id":"{request_id}"}}\n' + ;; + esac +done +"#, + log_path.display(), + release.display(), + ), + ) + .expect("gated identity acknowledgement script"); + let mut permissions = std::fs::metadata(&script_path) + .expect("script metadata") + .permissions(); + permissions.set_mode(0o755); + std::fs::set_permissions(&script_path, permissions).expect("script permissions"); + (script_path, release) +} + +#[cfg(unix)] +async fn assert_unique_protocol_request_owner_across_restart( + message: ToAgentMessage, + message_type: &str, + request_id: &str, +) { + let workspace = tempdir().expect("workspace"); + let fixtures = tempdir().expect("fixtures"); + let first_log = fixtures.path().join(format!("{message_type}-first.log")); + let (first_script, release_ack) = gated_identity_ack_script( + fixtures.path(), + &format!("{message_type}-first.sh"), + &first_log, + message_type, + request_id, + ); + let first_supervisor = connected_supervisor_for_script(&first_script).await; + let first_executor = Arc::new(AgentSupervisorHostedRunnerMessageExecutor::new(Arc::clone( + &first_supervisor, + ))); + let first = start_hosted_runner_with_message_executor( + test_config(workspace.path().to_path_buf()), + first_executor, + ) + .await + .expect("first hosted runner"); + let client = reqwest::Client::new(); + let (capability, subscription_id) = + attach_thread_controller(&client, &first.base_url(), "conn_identity_first").await; + let owner_headers = response_headers( + "conn_identity_first", + &subscription_id, + &capability, + "identity-owner-key", + ); + let competing_headers = response_headers( + "conn_identity_first", + &subscription_id, + &capability, + "identity-competing-key", + ); + + handle_message( + first.shared.clone(), + "sess_test", + owner_headers.clone(), + message.clone(), + ) + .await + .expect("owner response queues"); + let conflict = match handle_message( + first.shared.clone(), + "sess_test", + competing_headers.clone(), + message.clone(), + ) + .await + { + Err(error) => error, + Ok(_) => panic!("second key must not own the same protocol request"), + }; + assert_eq!(conflict.code, HostedRunnerErrorCode::IdempotencyConflict); + // Until the native child acknowledges, a retry reconciles the same owner; + // it is not yet a replay of completed ownership. + let pending = handle_message( + first.shared.clone(), + "sess_test", + owner_headers.clone(), + message.clone(), + ) + .await + .expect("same owner reconciles while acknowledgement is withheld"); + let ResponseBody::Json { body, .. } = pending else { + panic!("pending owner retry must return JSON"); + }; + assert_eq!(body["replayed"], false); + assert_eq!( + std::fs::read_to_string(&first_log) + .expect("pending child log") + .lines() + .count(), + 1, + "pending reconciliation must not redispatch the response" + ); + std::fs::write(&release_ack, b"accept").expect("release native acknowledgement"); + await_acknowledged_owner(&first, workspace.path(), "identity-owner-key", request_id).await; + let replay = handle_message( + first.shared.clone(), + "sess_test", + owner_headers, + message.clone(), + ) + .await + .expect("owner key replays after delayed acknowledgement"); + let ResponseBody::Json { body, .. } = replay else { + panic!("owner replay must return JSON"); + }; + assert_eq!(body["replayed"], true); + assert_eq!( + std::fs::read_to_string(&first_log) + .expect("first child log") + .lines() + .count(), + 1 + ); + first.shutdown().await; + first_supervisor + .lock() + .expect("first supervisor") + .shutdown(); + + let second_log = fixtures.path().join(format!("{message_type}-second.log")); + let second_script = create_delayed_identity_ack_script( + fixtures.path(), + &format!("{message_type}-second.sh"), + &second_log, + message_type, + request_id, + ); + let second_supervisor = connected_supervisor_for_script(&second_script).await; + let second_executor = Arc::new(AgentSupervisorHostedRunnerMessageExecutor::new(Arc::clone( + &second_supervisor, + ))); + let second = start_hosted_runner_with_message_executor( + test_config(workspace.path().to_path_buf()), + second_executor, + ) + .await + .expect("restarted hosted runner"); + let second_connection_id = if matches!( + &message, + ToAgentMessage::GovernedClientToolResult(tool_wire::GovernedClientToolResult { .. }) + ) { + "conn_identity_first" + } else { + "conn_identity_second" + }; + let (capability, subscription_id) = + attach_thread_controller(&client, &second.base_url(), second_connection_id).await; + let replay = handle_message( + second.shared.clone(), + "sess_test", + response_headers( + second_connection_id, + &subscription_id, + &capability, + "identity-owner-key", + ), + message.clone(), + ) + .await + .expect("durable owner key replays"); + let ResponseBody::Json { body, .. } = replay else { + panic!("durable owner replay must return JSON"); + }; + assert_eq!(body["replayed"], true); + let conflict = match handle_message( + second.shared.clone(), + "sess_test", + response_headers( + second_connection_id, + &subscription_id, + &capability, + "identity-competing-key", + ), + message, + ) + .await + { + Err(error) => error, + Ok(_) => panic!("request ownership must survive restart"), + }; + assert_eq!(conflict.code, HostedRunnerErrorCode::IdempotencyConflict); + assert!( + !second_log.exists(), + "restart must not redispatch either key" + ); + second.shutdown().await; + second_supervisor + .lock() + .expect("second supervisor") + .shutdown(); +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn delayed_tool_response_has_one_idempotency_owner_across_restart() { + assert_unique_protocol_request_owner_across_restart( + ToAgentMessage::ToolResponse { + call_id: "unique-tool-call".to_string(), + tool_execution_id: Some("unique-tool-execution".to_string()), + approved: true, + result: None, + }, + "tool_response", + "unique-tool-call", + ) + .await; +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn delayed_server_request_response_has_one_idempotency_owner_across_restart() { + assert_unique_protocol_request_owner_across_restart( + ToAgentMessage::ServerRequestResponse { + request_id: "unique-server-request".to_string(), + request_type: ServerRequestType::UserInput, + approved: None, + result: None, + content: Some(Vec::new()), + is_error: Some(false), + decision_action: None, + reason: Some("answer".to_string()), + }, + "server_request_response", + "unique-server-request", + ) + .await; +} + +#[cfg(unix)] +pub(super) fn governed_response_for_ack_test() -> ToAgentMessage { + ToAgentMessage::GovernedClientToolResult(tool_wire::GovernedClientToolResult { + process_tool_cost_micros: None, + call_id: "unique-governed-call".to_string(), + content: Vec::new(), + is_error: false, + tool_execution_id: "unique-governed-execution".to_string(), + client_instance_id: "conn_identity_first".to_string(), + grant_id: "grant-1".to_string(), + grant_version: 1, + grant_hash: "hash".to_string(), + turn_digest: "turn-digest".to_string(), + definition_digest: "definition-digest".to_string(), + args_digest: "args-digest".to_string(), + owner_lease_epoch: 1, + idempotency_key: "identity-owner-key".to_string(), + }) +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn delayed_governed_result_has_one_idempotency_owner_across_restart() { + assert_unique_protocol_request_owner_across_restart( + governed_response_for_ack_test(), + "governed_client_tool_result", + "unique-governed-call", + ) + .await; +} + +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn correlated_protocol_rejection_rolls_back_ownership_and_allows_retry() { + let workspace = tempdir().expect("workspace"); + let fixtures = tempdir().expect("fixtures"); + let log_path = fixtures.path().join("rejected-responses.log"); + let script = create_reject_then_accept_script(fixtures.path(), &log_path, None); + let release_ack = fixtures.path().join("release-corrected-ack"); + let scripted = std::fs::read_to_string(&script).expect("rejection script"); + std::fs::write( + &script, + scripted.replace( + " else\n", + &format!( + " else\n while [ ! -f '{}' ]; do sleep 0.01; done\n", + release_ack.display() + ), + ), + ) + .expect("gate corrected native acknowledgement"); + let supervisor = connected_supervisor_for_script(&script).await; + let executor = Arc::new(AgentSupervisorHostedRunnerMessageExecutor::new(Arc::clone( + &supervisor, + ))); + let handle = start_hosted_runner_with_message_executor( + test_config(workspace.path().to_path_buf()), + executor.clone(), + ) + .await + .expect("hosted runner"); + let client = reqwest::Client::new(); + let (capability, subscription_id) = + attach_thread_controller(&client, &handle.base_url(), "conn_rejection").await; + let headers = response_headers( + "conn_rejection", + &subscription_id, + &capability, + "rejected-key", + ); + let response = ToAgentMessage::ToolResponse { + call_id: "retry-call".to_string(), + tool_execution_id: Some("retry-execution".to_string()), + approved: true, + result: None, + }; + + let error = match handle_message( + handle.shared.clone(), + "sess_test", + headers.clone(), + response.clone(), + ) + .await + { + Err(error) => error, + Ok(_) => panic!("correlated protocol rejection must not return success"), + }; + assert_eq!(error.code, HostedRunnerErrorCode::RuntimeFailed); + assert!(error.message.contains("not awaiting a decision")); + { + let state = handle + .shared + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + assert!( + !state + .pending_response_idempotency + .contains_key("rejected-key") + ); + } + assert!( + !executor + .queued_responses + .lock() + .expect("queued responses") + .contains_key("rejected-key") + ); + assert!( + !load_executor_response_ledger(workspace.path(), "sess_test") + .expect("response ledger") + .iter() + .any(|(key, _)| key == "rejected-key") + ); + + handle_message( + handle.shared.clone(), + "sess_test", + headers.clone(), + response.clone(), + ) + .await + .expect("corrected retry dispatches"); + std::fs::write(&release_ack, b"accept").expect("release corrected acknowledgement"); + await_acknowledged_owner(&handle, workspace.path(), "rejected-key", "retry-call").await; + let replay = handle_message(handle.shared.clone(), "sess_test", headers, response) + .await + .expect("accepted retry replays"); + let ResponseBody::Json { body, .. } = replay else { + panic!("accepted retry replay must return JSON"); + }; + assert_eq!(body["replayed"], true); + assert_eq!( + std::fs::read_to_string(&log_path) + .expect("rejected response log") + .lines() + .count(), + 2, + "one rejected dispatch and one corrected dispatch are expected" + ); + handle.shutdown().await; + supervisor.lock().expect("supervisor").shutdown(); +} diff --git a/packages/local-host-rs/src/tools/bash/shield.rs b/packages/local-host-rs/src/tools/bash/shield.rs index b5e391705..0fe506a75 100644 --- a/packages/local-host-rs/src/tools/bash/shield.rs +++ b/packages/local-host-rs/src/tools/bash/shield.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; use std::ops::ControlFlow; use std::path::{Path, PathBuf}; -use std::time::{Duration, Instant}; +use std::time::Duration; use anyhow::{Context, Result, bail}; use tokio::io::AsyncReadExt; @@ -15,6 +15,7 @@ use crate::agent::credential_store::{publication_secret_kind, redact_credentials const MAX_OUTPUT: u64 = 8 * 1024 * 1024; const MAX_COMMITS: usize = 256; const MAX_FINDINGS: usize = 20; +const PARSER_CPU_BUDGET: Duration = Duration::from_millis(50); const HELP: &str = "Maestro Shield could not safely scan this Git operation. Run staging separately, then use a single literal git commit or git push command."; struct Publication { @@ -23,10 +24,83 @@ struct Publication { args: Vec, } +// Parsing is synchronous: thread CPU time excludes descheduling without +// charging other threads or weakening the original 50 ms computation limit. +#[cfg(unix)] +fn thread_cpu_time() -> Result { + let mut time = std::mem::MaybeUninit::::uninit(); + // SAFETY: clock_gettime initializes the pointed-to timespec on success. + if unsafe { libc::clock_gettime(libc::CLOCK_THREAD_CPUTIME_ID, time.as_mut_ptr()) } != 0 { + return Err(std::io::Error::last_os_error().into()); + } + // SAFETY: the successful call above initialized time. + let time = unsafe { time.assume_init() }; + let seconds = u64::try_from(time.tv_sec).context("invalid thread CPU clock seconds")?; + let nanos = u32::try_from(time.tv_nsec).context("invalid thread CPU clock nanoseconds")?; + if nanos >= 1_000_000_000 { + bail!("invalid thread CPU clock nanoseconds"); + } + Ok(Duration::new(seconds, nanos)) +} + +#[cfg(windows)] +fn thread_cpu_time() -> Result { + use windows_sys::Win32::Foundation::FILETIME; + use windows_sys::Win32::System::Threading::{GetCurrentThread, GetThreadTimes}; + + let mut creation = FILETIME::default(); + let mut exit = FILETIME::default(); + let mut kernel = FILETIME::default(); + let mut user = FILETIME::default(); + // SAFETY: the current-thread pseudo-handle is valid for GetThreadTimes; + // all four output pointers refer to initialized writable FILETIME values. + if unsafe { + GetThreadTimes( + GetCurrentThread(), + &mut creation, + &mut exit, + &mut kernel, + &mut user, + ) + } == 0 + { + return Err(std::io::Error::last_os_error().into()); + } + let ticks = + |time: FILETIME| (u64::from(time.dwHighDateTime) << 32) | u64::from(time.dwLowDateTime); + let total = ticks(kernel) + .checked_add(ticks(user)) + .context("thread CPU clock overflow")?; + Ok(Duration::new( + total / 10_000_000, + ((total % 10_000_000) * 100) as u32, + )) +} + +#[cfg(not(any(unix, windows)))] +fn thread_cpu_time() -> Result { + bail!("thread CPU clock is unavailable"); +} + /// Detect literal publication commands with the same Bash grammar used by /// command approval. A compound publication must be split so staging or a /// directory change cannot invalidate the scan before the command executes. fn publication(command: &str, cwd: &Path) -> Result> { + let bytes = command.as_bytes(); + publication_with_reader( + command, + cwd, + |offset, _| bytes.get(offset..).unwrap_or_default(), + thread_cpu_time, + ) +} + +fn publication_with_reader<'a>( + command: &'a str, + cwd: &Path, + mut read: impl FnMut(usize, tree_sitter::Point) -> &'a [u8], + mut cpu_time: impl FnMut() -> Result, +) -> Result> { if command.len() > 64 * 1024 { bail!("{HELP}"); } @@ -34,23 +108,37 @@ fn publication(command: &str, cwd: &Path) -> Result> { parser .set_language(&tree_sitter_bash::LANGUAGE.into()) .context(HELP)?; - let started_at = Instant::now(); - let timeout = Duration::from_millis(50); - let mut cancel_after_timeout = |_: &tree_sitter::ParseState| { - if started_at.elapsed() >= timeout { - ControlFlow::Break(()) - } else { - ControlFlow::Continue(()) + let started_at = cpu_time().context(HELP)?; + let mut budget_error = None; + let mut check_budget = || -> Result<()> { + let elapsed = cpu_time()? + .checked_sub(started_at) + .context("thread CPU clock moved backwards")?; + if elapsed >= PARSER_CPU_BUDGET { + bail!("Shield parser CPU budget exhausted"); } + Ok(()) }; - let bytes = command.as_bytes(); - let tree = parser - .parse_with_options( - &mut |offset, _| bytes.get(offset..).unwrap_or_default(), + let tree = { + let mut cancel_after_budget = |_: &tree_sitter::ParseState| match check_budget() { + Ok(()) => ControlFlow::Continue(()), + Err(error) => { + budget_error = Some(error); + ControlFlow::Break(()) + } + }; + parser.parse_with_options( + &mut read, None, - Some(ParseOptions::new().progress_callback(&mut cancel_after_timeout)), + Some(ParseOptions::new().progress_callback(&mut cancel_after_budget)), ) - .context(HELP)?; + }; + if let Some(error) = budget_error { + return Err(error.context(HELP)); + } + let tree = tree.context(HELP)?; + // Short parses may finish without invoking the progress callback. + check_budget().context(HELP)?; let root = tree.root_node(); let mut stack = vec![root]; let mut candidate = None; @@ -687,6 +775,121 @@ mod tests { assert_eq!(parsed.cwd, Path::new("root/a directory")); } + #[test] + fn parser_pause_does_not_reject_background_lifecycle_command() { + let command = "printf '%s\\n' \"$$\" > '/tmp/background-parent.pid'; sleep 30 & child=$!; printf '%s\\n' \"$child\" > '/tmp/background-child.pid'; wait \"$child\""; + let mut paused = false; + let parsed = publication_with_reader( + command, + Path::new("."), + |offset, _| { + if !paused { + paused = true; + std::thread::sleep(Duration::from_millis(60)); + } + command.as_bytes().get(offset..).unwrap_or_default() + }, + thread_cpu_time, + ); + assert!(parsed.unwrap().is_none()); + } + + #[test] + fn parser_cpu_budget_exhaustion_fails_closed() { + // The longer command invokes progress; the literal push also checks + // the final clock even when parsing completes without a callback. + for command in [ + "printf '%s\\n' \"$$\" > '/tmp/parent.pid'; sleep 30 & child=$!; wait \"$child\"", + "git push origin HEAD", + ] { + let mut started = false; + let error = publication_with_reader( + command, + Path::new("."), + |offset, _| command.as_bytes().get(offset..).unwrap_or_default(), + || { + if started { + Ok(Duration::from_millis(50)) + } else { + started = true; + Ok(Duration::ZERO) + } + }, + ) + .err() + .expect("exhausted parser budget must refuse the command"); + assert_eq!(error.to_string(), HELP); + assert!(format!("{error:#}").contains("parser CPU budget exhausted")); + } + } + + #[test] + fn parser_cpu_clock_failures_do_not_approve_commands() { + for command in [ + "echo x; echo y; echo z; echo w; echo v", + "git push origin HEAD", + ] { + for fail_at in [0, 1] { + let mut reads = 0; + let error = publication_with_reader( + command, + Path::new("."), + |offset, _| command.as_bytes().get(offset..).unwrap_or_default(), + || { + let fail = reads >= fail_at; + reads += 1; + if fail { + bail!("fixture clock unavailable"); + } + Ok(Duration::from_millis(100)) + }, + ) + .err() + .expect("failed parser clock must refuse the command"); + assert_eq!(error.to_string(), HELP); + assert!(format!("{error:#}").contains("fixture clock unavailable")); + } + let mut reads = 0; + let error = publication_with_reader( + command, + Path::new("."), + |offset, _| command.as_bytes().get(offset..).unwrap_or_default(), + || { + reads += 1; + Ok(if reads == 1 { + Duration::from_millis(100) + } else { + Duration::ZERO + }) + }, + ) + .err() + .expect("backward parser clock must refuse the command"); + assert!(format!("{error:#}").contains("clock moved backwards")); + } + } + + #[test] + fn malformed_parser_work_exhausts_cpu_budget() { + let command = "a[0]=(".repeat(10_000); + let error = publication(&command, Path::new(".")) + .err() + .expect("expensive malformed syntax must exhaust the parser budget"); + assert_eq!(error.to_string(), HELP); + assert!(format!("{error:#}").contains("parser CPU budget exhausted")); + } + + #[test] + fn parser_keeps_maximum_size_literal_commit_and_input_limit() { + let message = "x".repeat(65_520); + let command = format!("git commit -m '{message}'"); + assert_eq!(command.len(), 65_536); + let parsed = publication(&command, Path::new(".")).unwrap().unwrap(); + assert_eq!(parsed.operation, "commit"); + assert_eq!(parsed.args, ["-m", &message]); + assert!(publication(&"x".repeat(65_537), Path::new(".")).is_err()); + } + #[cfg(unix)] #[tokio::test] async fn inspection_keeps_the_original_sandbox_working_directory() {