use crate::{
config::ProjectConfig,
error::AoError,
error::Result,
prompt_builder,
scm::{
AutomatedComment, CheckRun, CiStatus, CreateIssueInput, Issue, IssueFilters, IssueUpdate,
MergeMethod, MergeReadiness, PrState, PrSummary, PullRequest, Review, ReviewComment,
ReviewDecision, ScmWebhookEvent, ScmWebhookRequest, ScmWebhookVerificationResult,
},
scm_transitions::ScmObservation,
types::{ActivityState, CostEstimate, Session, WorkspaceCreateConfig},
};
use async_trait::async_trait;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
#[async_trait]
pub trait Runtime: Send + Sync {
async fn create(
&self,
session_id: &str,
cwd: &Path,
launch_command: &str,
env: &[(String, String)],
) -> Result<String>;
async fn send_message(&self, handle: &str, msg: &str) -> Result<()>;
async fn is_alive(&self, handle: &str) -> Result<bool>;
async fn destroy(&self, handle: &str) -> Result<()>;
}
#[async_trait]
pub trait Workspace: Send + Sync {
async fn create(&self, cfg: &WorkspaceCreateConfig) -> Result<PathBuf>;
async fn destroy(&self, workspace_path: &Path) -> Result<()>;
async fn exists(&self, workspace_path: &Path) -> Result<bool> {
Ok(workspace_path.exists())
}
}
#[async_trait]
pub trait Agent: Send + Sync {
fn launch_command(&self, session: &Session) -> String;
fn environment(&self, session: &Session) -> Vec<(String, String)>;
fn initial_prompt(&self, session: &Session) -> String;
fn system_prompt(&self) -> Option<String> {
None
}
async fn detect_activity(&self, session: &Session) -> Result<ActivityState> {
if let Some(ref ws) = session.workspace_path {
if let Some(state) = crate::activity_log::detect_activity_from_log(ws) {
return Ok(state);
}
}
Ok(ActivityState::Ready)
}
async fn cost_estimate(&self, session: &Session) -> Result<Option<CostEstimate>> {
let Some(ref ws) = session.workspace_path else {
return Ok(None);
};
let ws = ws.clone();
let estimate = tokio::task::spawn_blocking(move || crate::cost_log::parse_usage_jsonl(&ws))
.await
.unwrap_or_else(|e| {
tracing::warn!("cost_estimate task failed: {e}");
None
});
Ok(estimate)
}
}
#[async_trait]
pub trait Scm: Send + Sync {
fn name(&self) -> &str;
async fn detect_pr(&self, session: &Session) -> Result<Option<PullRequest>>;
async fn pr_state(&self, pr: &PullRequest) -> Result<PrState>;
async fn ci_checks(&self, pr: &PullRequest) -> Result<Vec<CheckRun>>;
async fn ci_status(&self, pr: &PullRequest) -> Result<CiStatus>;
async fn reviews(&self, pr: &PullRequest) -> Result<Vec<Review>>;
async fn review_decision(&self, pr: &PullRequest) -> Result<ReviewDecision>;
async fn pending_comments(&self, pr: &PullRequest) -> Result<Vec<ReviewComment>>;
async fn mergeability(&self, pr: &PullRequest) -> Result<MergeReadiness>;
async fn merge(&self, pr: &PullRequest, method: Option<MergeMethod>) -> Result<()>;
async fn verify_webhook(
&self,
_request: &ScmWebhookRequest,
_project: &ProjectConfig,
) -> Result<ScmWebhookVerificationResult> {
Ok(ScmWebhookVerificationResult {
ok: false,
reason: Some("scm plugin does not support webhook verification".into()),
..Default::default()
})
}
async fn parse_webhook(
&self,
_request: &ScmWebhookRequest,
_project: &ProjectConfig,
) -> Result<Option<ScmWebhookEvent>> {
Ok(None)
}
async fn resolve_pr(&self, _reference: &str, _project: &ProjectConfig) -> Result<PullRequest> {
Err(AoError::Scm(
"scm plugin does not support PR resolution".into(),
))
}
async fn assign_pr_to_current_user(&self, _pr: &PullRequest) -> Result<()> {
Err(AoError::Scm(
"scm plugin does not support PR assignment".into(),
))
}
async fn checkout_pr(&self, _pr: &PullRequest, _workspace_path: &Path) -> Result<bool> {
Err(AoError::Scm(
"scm plugin does not support PR checkout".into(),
))
}
async fn pr_summary(&self, _pr: &PullRequest) -> Result<PrSummary> {
Err(AoError::Scm(
"scm plugin does not support PR summary".into(),
))
}
async fn close_pr(&self, _pr: &PullRequest) -> Result<()> {
Err(AoError::Scm(
"scm plugin does not support closing PRs".into(),
))
}
async fn automated_comments(&self, _pr: &PullRequest) -> Result<Vec<AutomatedComment>> {
Ok(Vec::new())
}
async fn enrich_prs_batch(
&self,
_prs: &[PullRequest],
) -> Result<HashMap<String, ScmObservation>> {
Ok(HashMap::new())
}
}
#[async_trait]
pub trait Tracker: Send + Sync {
fn name(&self) -> &str;
async fn get_issue(&self, identifier: &str) -> Result<Issue>;
async fn is_completed(&self, identifier: &str) -> Result<bool>;
fn issue_url(&self, identifier: &str) -> String;
fn branch_name(&self, identifier: &str) -> String;
async fn comment_issue(&self, _identifier: &str, _body: &str) -> Result<()> {
Err(AoError::Other(
"tracker does not support commenting".to_string(),
))
}
async fn assign_to_me(&self, _identifier: &str) -> Result<()> {
Err(AoError::Other(
"tracker does not support assignment".to_string(),
))
}
fn generate_prompt(&self, issue: &Issue) -> String {
prompt_builder::format_issue_context(issue)
}
async fn list_issues(&self, _filters: &IssueFilters) -> Result<Vec<Issue>> {
Err(AoError::Other(
"tracker does not support listing issues".to_string(),
))
}
async fn update_issue(&self, _identifier: &str, _update: &IssueUpdate) -> Result<()> {
Err(AoError::Other(
"tracker does not support updating issues".to_string(),
))
}
async fn create_issue(&self, _input: &CreateIssueInput) -> Result<Issue> {
Err(AoError::Other(
"tracker does not support creating issues".to_string(),
))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::activity_log::{
append_activity_entry, ActivityLogEntry, ACTIVITY_INPUT_STALENESS_SECS,
};
use crate::cost_log::usage_log_path;
use crate::types::{now_ms, SessionId, SessionStatus};
use std::io::Write;
use std::time::{SystemTime, UNIX_EPOCH};
struct StubAgent;
#[async_trait]
impl Agent for StubAgent {
fn launch_command(&self, _session: &Session) -> String {
String::new()
}
fn environment(&self, _session: &Session) -> Vec<(String, String)> {
Vec::new()
}
fn initial_prompt(&self, _session: &Session) -> String {
String::new()
}
}
fn unique_workspace(label: &str) -> PathBuf {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
let p = std::env::temp_dir().join(format!("ao-rs-trait-default-{label}-{nanos}"));
std::fs::create_dir_all(&p).unwrap();
p
}
fn session_with(workspace: Option<PathBuf>) -> Session {
Session {
id: SessionId("trait-default".into()),
project_id: "demo".into(),
status: SessionStatus::Working,
agent: "stub".into(),
agent_config: None,
branch: "feat".into(),
task: "t".into(),
workspace_path: workspace,
runtime_handle: None,
runtime: "tmux".into(),
activity: None,
created_at: now_ms(),
cost: None,
issue_id: None,
issue_url: None,
claimed_pr_number: None,
claimed_pr_url: None,
initial_prompt_override: None,
spawned_by: None,
last_merge_conflict_dispatched: None,
last_review_backlog_fingerprint: None,
}
}
#[tokio::test]
async fn detect_activity_default_no_workspace_returns_ready() {
let agent = StubAgent;
let session = session_with(None);
assert_eq!(
agent.detect_activity(&session).await.unwrap(),
ActivityState::Ready
);
}
#[tokio::test]
async fn detect_activity_default_no_log_returns_ready() {
let agent = StubAgent;
let ws = unique_workspace("no-log");
let session = session_with(Some(ws));
assert_eq!(
agent.detect_activity(&session).await.unwrap(),
ActivityState::Ready
);
}
#[tokio::test]
async fn detect_activity_default_surfaces_exited_from_log() {
let agent = StubAgent;
let ws = unique_workspace("exited");
append_activity_entry(
&ws,
&ActivityLogEntry {
ts: now_ms().to_string(),
state: ActivityState::Exited,
source: "terminal".into(),
trigger: None,
},
)
.unwrap();
let session = session_with(Some(ws));
assert_eq!(
agent.detect_activity(&session).await.unwrap(),
ActivityState::Exited
);
}
#[tokio::test]
async fn detect_activity_default_surfaces_fresh_waiting_input() {
let agent = StubAgent;
let ws = unique_workspace("waiting");
append_activity_entry(
&ws,
&ActivityLogEntry {
ts: now_ms().to_string(),
state: ActivityState::WaitingInput,
source: "terminal".into(),
trigger: Some("approve?".into()),
},
)
.unwrap();
let session = session_with(Some(ws));
assert_eq!(
agent.detect_activity(&session).await.unwrap(),
ActivityState::WaitingInput
);
}
#[tokio::test]
async fn detect_activity_default_stale_waiting_falls_back_to_ready() {
let agent = StubAgent;
let ws = unique_workspace("stale-waiting");
let stale_ms = now_ms().saturating_sub((ACTIVITY_INPUT_STALENESS_SECS + 60) * 1000);
append_activity_entry(
&ws,
&ActivityLogEntry {
ts: stale_ms.to_string(),
state: ActivityState::WaitingInput,
source: "terminal".into(),
trigger: None,
},
)
.unwrap();
let session = session_with(Some(ws));
assert_eq!(
agent.detect_activity(&session).await.unwrap(),
ActivityState::Ready
);
}
#[tokio::test]
async fn cost_estimate_default_no_workspace_returns_none() {
let agent = StubAgent;
let session = session_with(None);
assert!(agent.cost_estimate(&session).await.unwrap().is_none());
}
#[tokio::test]
async fn cost_estimate_default_no_log_returns_none() {
let agent = StubAgent;
let ws = unique_workspace("cost-missing");
let session = session_with(Some(ws));
assert!(agent.cost_estimate(&session).await.unwrap().is_none());
}
#[tokio::test]
async fn cost_estimate_default_reads_usage_log() {
let agent = StubAgent;
let ws = unique_workspace("cost-present");
let path = usage_log_path(&ws);
std::fs::create_dir_all(path.parent().unwrap()).unwrap();
let mut f = std::fs::File::create(&path).unwrap();
writeln!(
f,
r#"{{"input_tokens":100,"output_tokens":50,"cost_usd":0.5}}"#
)
.unwrap();
writeln!(
f,
r#"{{"input_tokens":200,"output_tokens":75,"cost_usd":0.25}}"#
)
.unwrap();
let session = session_with(Some(ws));
let cost = agent.cost_estimate(&session).await.unwrap().expect("some");
assert_eq!(cost.input_tokens, 300);
assert_eq!(cost.output_tokens, 125);
assert!((cost.cost_usd.unwrap() - 0.75).abs() < 1e-9);
}
}