use std::{collections::BTreeMap, sync::Arc, time::Duration};
use agent_infra_sdk::{
CallOptions, InfraClient,
workspace_contract::{
CommandNetworkPolicy, CommandSpec, CommandState, CreateCommandRequest, OutputStream,
PutWorkspaceFileRequest, WorkspaceSearchKind, WorkspaceSearchRequest,
},
};
use async_trait::async_trait;
use base64::{Engine as _, engine::general_purpose::STANDARD};
use runtime_ports::{
CommandOutput, CommandRequest, OperationContext, PortFailure, PortFailureKind, WorkspaceEntry,
WorkspaceSession,
};
use runtime_types::WorkspaceId;
use crate::map_sdk_error;
pub struct AgentInfraWorkspaceSession {
client: Arc<InfraClient>,
workspace_id: WorkspaceId,
}
impl AgentInfraWorkspaceSession {
pub fn new(client: Arc<InfraClient>, workspace_id: WorkspaceId) -> Self {
Self {
client,
workspace_id,
}
}
}
#[async_trait]
impl WorkspaceSession for AgentInfraWorkspaceSession {
async fn describe(&self, operation: &OperationContext) -> Result<String, PortFailure> {
let remaining = operation.remaining()?;
let record = self
.client
.workspace()
.workspace_with_options(
self.workspace_id.as_str(),
CallOptions::default().deadline(remaining),
)
.await
.map_err(map_sdk_error)?;
Ok(format!(
"managed workspace {} (backend={}, lifecycle={:?})",
record.id, record.backend, record.lifecycle
))
}
async fn read(&self, operation: &OperationContext, path: &str) -> Result<Vec<u8>, PortFailure> {
let remaining = operation.remaining()?;
let response = self
.client
.workspace()
.read_workspace_file_with_options(
self.workspace_id.as_str(),
path,
CallOptions::default().deadline(remaining),
)
.await
.map_err(map_sdk_error)?;
STANDARD.decode(response.content_base64).map_err(|_| {
PortFailure::new(
PortFailureKind::Protocol,
"WORKSPACE_BASE64_INVALID",
"workspace returned invalid base64 content",
)
})
}
async fn write(
&self,
operation: &OperationContext,
path: &str,
content: &[u8],
) -> Result<(), PortFailure> {
let remaining = operation.remaining()?;
self.client
.workspace()
.put_workspace_file(
self.workspace_id.as_str(),
path,
&PutWorkspaceFileRequest {
content_base64: STANDARD.encode(content),
if_match: None,
},
CallOptions::default()
.deadline(remaining)
.idempotency_key(format!("{}:write:{path}", operation.id)),
)
.await
.map(|_| ())
.map_err(map_sdk_error)
}
async fn list(
&self,
operation: &OperationContext,
path: &str,
) -> Result<Vec<WorkspaceEntry>, PortFailure> {
let remaining = operation.remaining()?;
let response = self
.client
.workspace()
.search_workspace_files_with_options(
self.workspace_id.as_str(),
&WorkspaceSearchRequest {
kind: WorkspaceSearchKind::Walk,
path: path.into(),
pattern: None,
include: None,
max_depth: 1,
cursor: None,
limit: 1000,
},
CallOptions::default().deadline(remaining),
)
.await
.map_err(map_sdk_error)?;
Ok(response
.paths
.into_iter()
.map(|path| WorkspaceEntry {
path,
is_dir: false,
size: None,
})
.collect())
}
async fn search(
&self,
operation: &OperationContext,
path: &str,
query: &str,
limit: usize,
) -> Result<Vec<String>, PortFailure> {
let remaining = operation.remaining()?;
let response = self
.client
.workspace()
.search_workspace_files_with_options(
self.workspace_id.as_str(),
&WorkspaceSearchRequest {
kind: WorkspaceSearchKind::Grep,
path: path.into(),
pattern: Some(query.into()),
include: None,
max_depth: 64,
cursor: None,
limit: limit.min(1000),
},
CallOptions::default().deadline(remaining),
)
.await
.map_err(map_sdk_error)?;
let mut output = response.paths;
output.extend(
response
.matches
.into_iter()
.map(|item| format!("{}:{}:{}", item.file, item.line_number, item.line)),
);
Ok(output)
}
async fn execute(
&self,
operation: &OperationContext,
request: CommandRequest,
) -> Result<CommandOutput, PortFailure> {
let remaining = operation.remaining()?;
let timeout = remaining.min(Duration::from_secs(60));
let command = self
.client
.workspace()
.create_command_with_options(
self.workspace_id.as_str(),
&CreateCommandRequest {
id: None,
command: CommandSpec {
argv: vec!["sh".into(), "-lc".into(), request.command],
cwd: request.cwd,
env: BTreeMap::new(),
shell: false,
timeout_ms: timeout.as_millis() as u64,
memory_bytes: 1024 * 1024 * 1024,
cpu_millis: 1000,
max_processes: 64,
network: CommandNetworkPolicy::Deny,
disk_bytes: 1024 * 1024 * 1024,
max_output_bytes: 256 * 1024,
},
},
CallOptions::default()
.deadline(remaining)
.idempotency_key(operation.id.as_str()),
)
.await
.map_err(map_sdk_error)?;
let command_id = command.id;
loop {
if operation.cancellation.is_cancelled() {
let _ = self
.client
.workspace()
.cancel_command_with_options(
self.workspace_id.as_str(),
&command_id,
CallOptions::default()
.deadline(Duration::from_secs(5))
.idempotency_key(format!(
"{}:cancel-command:{command_id}",
operation.id
)),
)
.await;
return Err(PortFailure::canceled());
}
let remaining = operation.remaining()?;
let current = self
.client
.workspace()
.command_with_options(
self.workspace_id.as_str(),
&command_id,
CallOptions::default().deadline(remaining),
)
.await
.map_err(map_sdk_error)?;
if matches!(
current.state,
CommandState::Succeeded | CommandState::Failed | CommandState::Canceled
) {
let mut stdout = String::new();
let mut stderr = String::new();
let mut after = 0;
let mut truncated = current.output_truncated;
loop {
let remaining = operation.remaining()?;
let page = self
.client
.workspace()
.command_output_with_options(
self.workspace_id.as_str(),
&command_id,
after,
Some(64),
CallOptions::default().deadline(remaining),
)
.await
.map_err(map_sdk_error)?;
for chunk in page.chunks {
match chunk.stream {
OutputStream::Stdout => stdout.push_str(&chunk.content),
OutputStream::Stderr => stderr.push_str(&chunk.content),
}
}
truncated |= page.truncated;
match page.next_cursor {
Some(next) if next > after => after = next,
_ => break,
}
}
return Ok(CommandOutput {
exit_code: current.exit_code.unwrap_or(
if current.state == CommandState::Succeeded {
0
} else {
1
},
),
stdout,
stderr,
truncated,
});
}
tokio::time::sleep(Duration::from_millis(100)).await;
}
}
}