Run agent shell commands as independent background jobs
This commit is contained in:
@@ -49,6 +49,7 @@ message HostScopedRequest {
|
||||
GitGenerateCommitMessageRequest git_generate_commit_message = 18;
|
||||
GitGetCommittedBranchFilesRequest git_get_committed_branch_files = 19;
|
||||
RipgrepSearchRequest ripgrep_search = 21;
|
||||
BackgroundJobRequest background_job = 22;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -125,6 +126,7 @@ message ServerMessage {
|
||||
GitHubRepositoryInfoPush github_repository_info_push = 36;
|
||||
RipgrepSearchResponse ripgrep_search_response = 37;
|
||||
RemoteAgentContextSnapshot remote_agent_context_snapshot = 33;
|
||||
BackgroundJobResponse background_job_response = 38;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -177,6 +179,42 @@ message Abort {
|
||||
message InitializeResponse {
|
||||
string server_version = 1;
|
||||
string host_id = 2;
|
||||
bool supports_background_jobs = 3;
|
||||
}
|
||||
|
||||
// Agent commands run independently of the requesting connection. Repeated starts with the
|
||||
// same conversation/action identity return the existing job rather than running it again.
|
||||
message BackgroundJobRequest {
|
||||
enum Operation {
|
||||
START = 0;
|
||||
STATUS = 1;
|
||||
CANCEL = 2;
|
||||
}
|
||||
Operation operation = 1;
|
||||
string job_id = 2;
|
||||
string conversation_id = 3;
|
||||
string action_id = 4;
|
||||
string command = 5;
|
||||
string working_directory = 6;
|
||||
map<string, string> environment = 7;
|
||||
uint64 session_id = 8;
|
||||
uint64 log_offset = 9;
|
||||
// Distinguish an absent shell snapshot from a known, explicitly empty environment.
|
||||
bool has_environment = 10;
|
||||
}
|
||||
|
||||
message BackgroundJobResponse {
|
||||
string job_json = 1;
|
||||
bytes log_chunk = 2;
|
||||
uint64 next_log_offset = 3;
|
||||
uint64 total_log_bytes = 4;
|
||||
bool logs_available = 5;
|
||||
string error = 6;
|
||||
// Bounded completion evidence survives unlinking the daemon's temporary log file.
|
||||
bytes completion_log_tail = 7;
|
||||
uint64 completion_log_bytes = 8;
|
||||
bool completion_log_tail_available = 9;
|
||||
bool job_not_found = 10;
|
||||
}
|
||||
|
||||
// Sent by the client after the SSH session has been bootstrapped.
|
||||
|
||||
@@ -201,6 +201,7 @@ pub struct RemoteServerClient {
|
||||
/// `send_request` after inserting into `pending_requests` to avoid hanging
|
||||
/// on a dead connection.
|
||||
disconnected: Arc<AtomicBool>,
|
||||
background_jobs_supported: AtomicBool,
|
||||
|
||||
/// Dedicated channel for `RequestFailed` telemetry events. Separate from
|
||||
/// the lifecycle `event_tx` so that holding this sender does not keep the
|
||||
@@ -313,6 +314,7 @@ impl RemoteServerClient {
|
||||
outbound_tx,
|
||||
pending_requests,
|
||||
disconnected,
|
||||
background_jobs_supported: AtomicBool::new(false),
|
||||
failure_tx,
|
||||
host_response_tx,
|
||||
},
|
||||
@@ -336,6 +338,10 @@ impl RemoteServerClient {
|
||||
self.disconnected.load(Ordering::Acquire)
|
||||
}
|
||||
|
||||
pub fn supports_background_jobs(&self) -> bool {
|
||||
self.background_jobs_supported.load(Ordering::Acquire)
|
||||
}
|
||||
|
||||
/// Sends an `Initialize` request and awaits the `InitializeResponse`.
|
||||
pub async fn initialize(
|
||||
&self,
|
||||
@@ -357,7 +363,11 @@ impl RemoteServerClient {
|
||||
let response = self.send_request_internal(request_id, msg).await?;
|
||||
|
||||
match response.message {
|
||||
Some(server_message::Message::InitializeResponse(resp)) => Ok(resp),
|
||||
Some(server_message::Message::InitializeResponse(resp)) => {
|
||||
self.background_jobs_supported
|
||||
.store(resp.supports_background_jobs, Ordering::Release);
|
||||
Ok(resp)
|
||||
}
|
||||
other => {
|
||||
safe_error!(
|
||||
safe: ("Remote server unexpected response for Initialize"),
|
||||
|
||||
@@ -214,8 +214,10 @@ async fn initialize_round_trip() {
|
||||
server_message::Message::InitializeResponse(InitializeResponse {
|
||||
server_version: "test-0.1.0".to_string(),
|
||||
host_id: "test-host-id".to_string(),
|
||||
supports_background_jobs: true,
|
||||
})
|
||||
});
|
||||
assert!(!client.supports_background_jobs());
|
||||
|
||||
let resp = client
|
||||
.initialize(
|
||||
@@ -229,6 +231,8 @@ async fn initialize_round_trip() {
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(resp.supports_background_jobs);
|
||||
assert!(client.supports_background_jobs());
|
||||
assert_eq!(resp.server_version, "test-0.1.0");
|
||||
assert_eq!(resp.host_id, "test-host-id");
|
||||
}
|
||||
@@ -243,6 +247,7 @@ async fn initialize_sends_empty_auth_token_when_none() {
|
||||
server_message::Message::InitializeResponse(InitializeResponse {
|
||||
server_version: "test-0.1.0".to_string(),
|
||||
host_id: "test-host-id".to_string(),
|
||||
supports_background_jobs: true,
|
||||
})
|
||||
});
|
||||
|
||||
@@ -270,6 +275,7 @@ async fn initialize_sends_auth_token_when_provided() {
|
||||
server_message::Message::InitializeResponse(InitializeResponse {
|
||||
server_version: "test-0.1.0".to_string(),
|
||||
host_id: "test-host-id".to_string(),
|
||||
supports_background_jobs: true,
|
||||
})
|
||||
});
|
||||
|
||||
@@ -381,6 +387,7 @@ async fn is_disconnected_starts_false() {
|
||||
server_message::Message::InitializeResponse(InitializeResponse {
|
||||
server_version: "test-0.1.0".to_string(),
|
||||
host_id: "test-host-id".to_string(),
|
||||
supports_background_jobs: true,
|
||||
})
|
||||
});
|
||||
|
||||
@@ -427,6 +434,7 @@ async fn concurrent_in_flight_requests() {
|
||||
server_message::Message::InitializeResponse(InitializeResponse {
|
||||
server_version: "test-0.1.0".to_string(),
|
||||
host_id: "test-host-id".to_string(),
|
||||
supports_background_jobs: true,
|
||||
})
|
||||
});
|
||||
let client = std::sync::Arc::new(client);
|
||||
@@ -472,6 +480,7 @@ async fn mock_server_with_error_handling(
|
||||
InitializeResponse {
|
||||
server_version: "test-0.1.0".to_string(),
|
||||
host_id: "test-host-id".to_string(),
|
||||
supports_background_jobs: true,
|
||||
},
|
||||
)),
|
||||
};
|
||||
|
||||
@@ -167,6 +167,7 @@ fn every_host_scoped_request_has_a_response_disposition() {
|
||||
M::GitGenerateCommitMessage(_) => "manager::generate_commit_message",
|
||||
M::GitGetCommittedBranchFiles(_) => "manager::get_committed_branch_files",
|
||||
M::RipgrepSearch(_) => "manager::start_ripgrep_search",
|
||||
M::BackgroundJob(_) => "manager::HostRequestHandle::background_job",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -888,6 +888,23 @@ pub struct HostRequestHandle {
|
||||
}
|
||||
|
||||
impl HostRequestHandle {
|
||||
pub async fn background_job(
|
||||
&self,
|
||||
request: crate::proto::BackgroundJobRequest,
|
||||
) -> Result<crate::proto::BackgroundJobResponse, HostRequestError> {
|
||||
let message = self
|
||||
.send(crate::proto::host_scoped_request::Message::BackgroundJob(
|
||||
request,
|
||||
))
|
||||
.await?;
|
||||
match message.message {
|
||||
Some(crate::proto::server_message::Message::BackgroundJobResponse(response)) => {
|
||||
Ok(response)
|
||||
}
|
||||
_ => Err(HostRequestError::UnexpectedResponse),
|
||||
}
|
||||
}
|
||||
|
||||
/// Sends a host-scoped request and awaits the raw `ServerMessage`.
|
||||
///
|
||||
/// Bounces to the main thread to call `send_host_request`, then
|
||||
|
||||
@@ -7,6 +7,50 @@ use crate::proto::{
|
||||
RemoteAgentContextSnapshot, RemoteContextFileProto, RemoteSkillProto, ServerMessage,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn background_job_requests_preserve_identity_and_byte_log_cursors() {
|
||||
let request = ClientMessage::host_scoped(
|
||||
"background-status".to_owned(),
|
||||
crate::proto::host_scoped_request::Message::BackgroundJob(
|
||||
crate::proto::BackgroundJobRequest {
|
||||
operation: crate::proto::background_job_request::Operation::Status as i32,
|
||||
job_id: "2ce65749-03d4-4482-910f-b9224575e9bd".to_owned(),
|
||||
conversation_id: "43131875-abf9-4f98-b20a-b607d1b99125".to_owned(),
|
||||
action_id: "job-action".to_owned(),
|
||||
log_offset: 128 * 1024 + 3,
|
||||
has_environment: true,
|
||||
..Default::default()
|
||||
},
|
||||
),
|
||||
);
|
||||
let mut buffer = Vec::new();
|
||||
write_message(&mut buffer, &request).await.unwrap();
|
||||
let decoded: ClientMessage = read_message(&mut buffer.as_slice()).await.unwrap();
|
||||
assert_eq!(decoded, request);
|
||||
|
||||
let response = ServerMessage {
|
||||
request_id: request.request_id,
|
||||
message: Some(server_message::Message::BackgroundJobResponse(
|
||||
crate::proto::BackgroundJobResponse {
|
||||
job_json: "{}".to_owned(),
|
||||
log_chunk: vec![0xf0, 0x9f, 0x92, 0xa1, 0, 0xff],
|
||||
next_log_offset: 128 * 1024 + 9,
|
||||
total_log_bytes: 1024 * 1024,
|
||||
logs_available: true,
|
||||
error: String::new(),
|
||||
completion_log_tail: "retained completion output".as_bytes().to_vec(),
|
||||
completion_log_bytes: 1024 * 1024 + 7,
|
||||
completion_log_tail_available: true,
|
||||
job_not_found: false,
|
||||
},
|
||||
)),
|
||||
};
|
||||
let mut buffer = Vec::new();
|
||||
write_message(&mut buffer, &response).await.unwrap();
|
||||
let decoded: ServerMessage = read_message(&mut buffer.as_slice()).await.unwrap();
|
||||
assert_eq!(decoded, response);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn round_trip_client_message() {
|
||||
let msg = ClientMessage::session_scoped(
|
||||
@@ -108,6 +152,7 @@ async fn round_trip_server_message() {
|
||||
InitializeResponse {
|
||||
server_version: "0.1.0".to_string(),
|
||||
host_id: "test-host".to_string(),
|
||||
supports_background_jobs: true,
|
||||
},
|
||||
)),
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user