use std::borrow::Cow;
use std::ffi::OsStr;
#[cfg(not(test))]
use std::io::IsTerminal as _;
use std::path::{Path, PathBuf};
use std::process::ExitCode;
use std::{env, io};
use ag_harness::{
Harness, ModelConfiguration, ModelConfigurationError, ModelProvider, OutputSchema, Repository,
Session, SessionInfo, Tool, TurnOutcome,
};
use clap::builder::{PossibleValuesParser, TypedValueParser};
use clap::{Args, Parser, Subcommand};
use serde_json::{Map, Value};
use thiserror::Error;
use tokio::io::{AsyncBufRead, AsyncBufReadExt as _, AsyncWrite, AsyncWriteExt as _, BufReader};
const READ_ONLY_SYSTEM_PROMPT: &str = concat!(
"You are operating in a read-only repository harness. The read tool supports file, list, ",
"search, diff, and show actions. For change review, call diff first, then use search, file, ",
"list, or show for evidence. Call the tool immediately and use its result before answering. ",
"Never narrate, promise, or defer a future tool call. Never claim that you created, ",
"modified, deleted, or executed files or commands because filesystem mutation and command ",
"execution are unavailable. If asked to perform an unsupported action, state that it is ",
"unsupported."
);
const READ_WRITE_SYSTEM_PROMPT: &str = concat!(
"You are operating in a repository harness with read and write tools. The read tool supports ",
"file, list, search, diff, and show actions. For change review, call diff first. When a user ",
"asks about repository contents, call read immediately and use its result before answering. ",
"When a user asks to create or modify a file, call the write tool ",
"immediately in the same response. Never narrate, promise, or defer a future tool call. Only ",
"claim that a file was created or modified after the write tool succeeds. File deletion and ",
"command execution are unavailable."
);
#[derive(Debug, Parser)]
#[command(
name = "ag-harness",
version,
about = "Chats with models through a repository harness",
after_help = provider_help()
)]
struct Cli {
#[arg(long, global = true, value_name = "FILE")]
database: Option<PathBuf>,
#[arg(long, global = true, value_name = "FILE")]
git_executable: Option<PathBuf>,
#[command(subcommand)]
command: Command,
}
#[derive(Debug, Subcommand)]
enum Command {
Run(RunArgs),
Resume(ResumeArgs),
}
#[derive(Debug, Args)]
#[command(after_help = provider_help())]
struct RunArgs {
model: String,
#[arg(value_parser = parse_prompt)]
prompt: Option<String>,
#[arg(long, value_name = "URL")]
base_url: Option<String>,
#[arg(long)]
allow_write: bool,
#[arg(
long,
default_value_t = ModelProvider::Muse,
value_parser = model_provider_parser()
)]
provider: ModelProvider,
#[arg(long, value_name = "DIR", default_value = ".")]
read_dir: PathBuf,
#[arg(long, value_name = "ID")]
session: Option<String>,
}
#[derive(Debug, Args)]
struct ResumeArgs {
session: String,
#[arg(value_parser = parse_prompt)]
prompt: Option<String>,
#[arg(long, value_name = "URL")]
base_url: Option<String>,
#[arg(long)]
allow_write: bool,
#[arg(long, value_name = "DIR", default_value = ".")]
read_dir: PathBuf,
}
fn model_provider_parser() -> impl TypedValueParser<Value = ModelProvider> {
PossibleValuesParser::new(
ModelProvider::all()
.iter()
.map(|provider| provider.as_str()),
)
.try_map(|provider| provider.parse::<ModelProvider>())
}
fn provider_help() -> String {
let mut help =
String::from("Supported models (other endpoint-supported model IDs also work):\n");
for provider in ModelProvider::all() {
help.push_str(" ");
help.push_str(provider.as_str());
help.push_str(": ");
help.push_str(&provider.known_models().join(", "));
help.push('\n');
}
help.push_str("\nCredentials:\n");
for provider in ModelProvider::all() {
help.push_str(" ");
help.push_str(provider.as_str());
help.push_str(": ");
help.push_str(provider.api_key_environment());
if provider.default_base_url().is_some() {
help.push_str(" (");
help.push_str(provider.base_url_environment());
help.push_str(" optional)");
} else {
help.push_str(", ");
help.push_str(provider.base_url_environment());
}
help.push('\n');
}
help.pop();
help
}
fn parse_prompt(prompt: &str) -> Result<String, String> {
if prompt.trim().is_empty() {
return Err("prompt must contain a non-whitespace character".to_string());
}
Ok(prompt.to_string())
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum ChatMode {
Interactive,
NonInteractive,
OneShot,
}
impl ChatMode {
fn detect(cli: &Cli, stdin_is_terminal: bool, stdout_is_terminal: bool) -> Self {
if stdin_is_terminal && stdout_is_terminal {
return Self::Interactive;
}
let initial_prompt = match &cli.command {
Command::Run(args) => &args.prompt,
Command::Resume(args) => &args.prompt,
};
if stdin_is_terminal && initial_prompt.is_some() {
return Self::OneShot;
}
Self::NonInteractive
}
}
#[cfg(not(test))]
#[tokio::main]
async fn main() -> ExitCode {
let cli = Cli::parse();
let stdin_is_terminal = io::stdin().is_terminal();
let stdout_is_terminal = io::stdout().is_terminal();
let mode = ChatMode::detect(&cli, stdin_is_terminal, stdout_is_terminal);
let input = BufReader::new(tokio::io::stdin());
let output = tokio::io::stdout();
report_exit(
execute(cli, |name| env::var(name), input, output, mode).await,
io::stderr().lock(),
)
}
fn report_exit(result: Result<(), CliError>, mut error_output: impl io::Write) -> ExitCode {
match result {
Ok(()) => ExitCode::SUCCESS,
Err(error) => {
let error = error.to_string();
let error = single_line_terminal_text(&error);
let _ = writeln!(error_output, "{error}");
ExitCode::FAILURE
}
}
}
async fn execute<Input, Output>(
cli: Cli,
mut environment: impl FnMut(&str) -> Result<String, env::VarError>,
input: Input,
output: Output,
mode: ChatMode,
) -> Result<(), CliError>
where
Input: AsyncBufRead + Unpin,
Output: AsyncWrite + Unpin,
{
let database = database_path(cli.database, &mut environment)?;
let git_executable = cli.git_executable;
match cli.command {
Command::Run(args) => {
let session_id = args
.session
.unwrap_or_else(|| uuid::Uuid::new_v4().simple().to_string());
let client = model_client(
args.provider,
&args.model,
args.base_url.as_deref(),
&mut environment,
)?;
let repository = repository_or_default(args.read_dir, git_executable)?;
let (harness, system_prompt) =
configured_harness(client, database, repository, args.allow_write);
let mut session = harness
.session(&session_id, chat_schema()?)
.system_prompt(system_prompt)
.create()
.await?;
let mut output = output;
announce_session(&mut output, &session_id).await?;
run_chat(&mut session, &args.model, args.prompt, input, output, mode).await
}
Command::Resume(args) => {
let info = SessionInfo::load(&database, &args.session).await?;
let (provider, model) = stored_model_identity(&info)?;
let client =
model_client(provider, &model, args.base_url.as_deref(), &mut environment)?;
let repository = repository_or_default(args.read_dir, git_executable)?;
let (harness, _) = configured_harness(client, database, repository, args.allow_write);
let mut session = harness.resume(&args.session).await?;
let mut output = output;
announce_session(&mut output, &args.session).await?;
run_chat(&mut session, &model, args.prompt, input, output, mode).await
}
}
}
async fn announce_session(
output: &mut (impl AsyncWrite + Unpin),
session_id: &str,
) -> Result<(), io::Error> {
let session_id = single_line_terminal_text(session_id);
output
.write_all(format!("session: {session_id}\n").as_bytes())
.await?;
output.flush().await
}
fn database_path(
explicit: Option<PathBuf>,
environment: &mut impl FnMut(&str) -> Result<String, env::VarError>,
) -> Result<PathBuf, CliError> {
if let Some(path) = explicit {
return Ok(path);
}
if let Ok(root) = environment("AG_HARNESS_ROOT")
&& !root.trim().is_empty()
{
return Ok(PathBuf::from(root).join("db").join("harness.db"));
}
let home = environment("HOME").map_err(|_| CliError::DatabaseLocation)?;
if home.trim().is_empty() {
return Err(CliError::DatabaseLocation);
}
Ok(PathBuf::from(home).join(".ag-harness/db/harness.db"))
}
fn repository_or_default(root: PathBuf, explicit: Option<PathBuf>) -> Result<Repository, CliError> {
if let Some(explicit) = explicit {
return Repository::new(root, explicit).map_err(CliError::from);
}
let path = env::var_os("PATH");
repository_from_path(&root, path.as_deref())
}
fn repository_from_path(root: &Path, path: Option<&OsStr>) -> Result<Repository, CliError> {
let executable_name = format!("git{}", env::consts::EXE_SUFFIX);
let candidates = path
.iter()
.flat_map(env::split_paths)
.filter(|directory| directory.is_absolute())
.map(|directory| directory.join(&executable_name));
for candidate in candidates {
match Repository::new(root, candidate) {
Ok(repository) => return Ok(repository),
Err(
error @ (ag_harness::RepositoryError::Root { .. }
| ag_harness::RepositoryError::RootIsGitAdministrative { .. }),
) => return Err(error.into()),
Err(_) => {}
}
}
Err(CliError::GitExecutableNotFound)
}
fn model_client(
provider: ModelProvider,
model: &str,
base_url: Option<&str>,
environment: &mut impl FnMut(&str) -> Result<String, env::VarError>,
) -> Result<ag_harness::ModelClient, CliError> {
let mut configuration = ModelConfiguration::new(provider, model);
if let Some(base_url) = base_url {
configuration = configuration.base_url(base_url);
}
configuration
.client_from_environment(environment)
.map_err(CliError::from)
}
fn configured_harness(
client: ag_harness::ModelClient,
database: PathBuf,
repository: Repository,
allow_write: bool,
) -> (Harness, &'static str) {
let mut harness = Harness::new(client)
.database(database)
.repository(repository)
.allow(Tool::Read);
if allow_write {
harness = harness.allow(Tool::Write);
(harness, READ_WRITE_SYSTEM_PROMPT)
} else {
(harness, READ_ONLY_SYSTEM_PROMPT)
}
}
fn stored_model_identity(info: &SessionInfo) -> Result<(ModelProvider, String), CliError> {
stored_model_identity_parts(info.provider(), info.model())
}
fn stored_model_identity_parts(
provider: Option<&str>,
model: Option<&str>,
) -> Result<(ModelProvider, String), CliError> {
let model = model.ok_or(CliError::MissingModelIdentity)?.to_string();
let provider = match provider {
Some("meta") => ModelProvider::Muse,
Some("moonshot_ai") => ModelProvider::Kimi,
Some("alibaba_cloud") => ModelProvider::Qwen,
_ => return Err(CliError::MissingModelIdentity),
};
Ok((provider, model))
}
async fn run_chat<Input, Output>(
session: &mut Session<'_>,
requested_model: &str,
initial_prompt: Option<String>,
mut input: Input,
mut output: Output,
mode: ChatMode,
) -> Result<(), CliError>
where
Input: AsyncBufRead + Unpin,
Output: AsyncWrite + Unpin,
{
if mode == ChatMode::Interactive {
let requested_model = single_line_terminal_text(requested_model);
output
.write_all(format!("Chat with {requested_model}. Ctrl-D to exit.\n").as_bytes())
.await?;
}
let mut pending_prompt = initial_prompt;
let mut turn_failed = false;
loop {
let Some(prompt) = read_prompt(&mut pending_prompt, &mut input, &mut output, mode).await?
else {
break;
};
if prompt.trim().is_empty() {
continue;
}
match session.send(prompt).await {
Ok(outcome) => write_outcome(&mut output, requested_model, &outcome).await?,
Err(error) if mode == ChatMode::Interactive => {
write_turn_error(&mut output, &error).await?;
}
Err(error) if mode == ChatMode::OneShot => return Err(error.into()),
Err(error) => {
write_turn_error(&mut output, &error).await?;
turn_failed = true;
}
}
if mode == ChatMode::OneShot {
break;
}
}
if turn_failed {
Err(CliError::ChatTurnsFailed)
} else {
Ok(())
}
}
async fn read_prompt<Input, Output>(
pending_prompt: &mut Option<String>,
input: &mut Input,
output: &mut Output,
mode: ChatMode,
) -> Result<Option<String>, io::Error>
where
Input: AsyncBufRead + Unpin,
Output: AsyncWrite + Unpin,
{
if let Some(prompt) = pending_prompt.take() {
return Ok(Some(prompt));
}
if mode == ChatMode::Interactive {
output.write_all(b">>> ").await?;
output.flush().await?;
}
let mut prompt = String::new();
if input.read_line(&mut prompt).await? == 0 {
return Ok(None);
}
trim_line_ending(&mut prompt);
Ok(Some(prompt))
}
async fn write_outcome(
output: &mut (impl AsyncWrite + Unpin),
requested_model: &str,
outcome: &TurnOutcome,
) -> Result<(), CliError> {
let message = outcome
.output()
.get("message")
.and_then(serde_json::Value::as_str)
.ok_or(CliError::MissingMessage)?;
output.write_all(assistant_text(message).as_bytes()).await?;
output.write_all(b"---\n").await?;
output
.write_all(format!("turn: {}\n", format_duration(outcome.report().duration())).as_bytes())
.await?;
output
.write_all(format!("model calls: {}\n", outcome.report().model_requests().len()).as_bytes())
.await?;
for (index, request) in outcome.report().model_requests().iter().enumerate() {
let response_type = request.response_type();
let completion = request.completion();
let model = completion
.and_then(|metadata| metadata.response_model())
.unwrap_or(requested_model);
let finish_reason =
completion.map_or("unavailable", ag_harness::CompletionMetadata::finish_reason);
let model = single_line_terminal_text(model);
let finish_reason = single_line_terminal_text(finish_reason);
let usage = completion
.and_then(|metadata| metadata.usage())
.map_or_else(|| "tokens unavailable".to_string(), format_usage);
output
.write_all(
format!(
" {}. {response_type}; {model}; {finish_reason}; {}; {usage}\n",
index + 1,
format_duration(request.duration()),
)
.as_bytes(),
)
.await?;
}
if outcome.report().tool_calls().is_empty() {
output.write_all(b"tools: none\n").await?;
} else {
output.write_all(b"tools:\n").await?;
for activity in outcome.report().tool_calls() {
output
.write_all(format!(" {activity}\n").as_bytes())
.await?;
}
}
output.flush().await?;
Ok(())
}
async fn write_turn_error(
output: &mut (impl AsyncWrite + Unpin),
error: &(impl std::fmt::Display + ?Sized),
) -> Result<(), io::Error> {
let error = error.to_string();
let error = single_line_terminal_text(&error);
output
.write_all(format!("error: {error}\n").as_bytes())
.await?;
output.flush().await
}
fn format_usage(usage: &ag_harness::CompletionUsage) -> String {
let input = usage
.input_tokens()
.map_or_else(|| "?".to_string(), |tokens| tokens.to_string());
let output = usage
.output_tokens()
.map_or_else(|| "?".to_string(), |tokens| tokens.to_string());
let total = usage
.total_tokens()
.map_or_else(|| "?".to_string(), |tokens| tokens.to_string());
format!("tokens {input} in, {output} out, {total} total")
}
fn format_duration(duration: std::time::Duration) -> String {
if duration.as_millis() == 0 {
"<1 ms".to_string()
} else {
format!("{} ms", duration.as_millis())
}
}
fn trim_line_ending(line: &mut String) {
if line.ends_with('\n') {
line.pop();
if line.ends_with('\r') {
line.pop();
}
}
}
fn assistant_text(text: &str) -> String {
let text = terminal_text(text);
let mut framed = String::new();
for (index, line) in text.split('\n').enumerate() {
framed.push_str(if index == 0 {
"assistant> "
} else {
" "
});
framed.push_str(line);
framed.push('\n');
}
framed
}
fn terminal_text(text: &str) -> Cow<'_, str> {
if text.chars().all(is_terminal_safe) {
return Cow::Borrowed(text);
}
Cow::Owned(
text.chars()
.map(|character| {
if is_terminal_safe(character) {
character
} else {
'\u{fffd}'
}
})
.collect(),
)
}
fn single_line_terminal_text(text: &str) -> Cow<'_, str> {
if text.chars().all(|character| !character.is_control()) {
return Cow::Borrowed(text);
}
Cow::Owned(
text.chars()
.map(|character| {
if character.is_control() {
'\u{fffd}'
} else {
character
}
})
.collect(),
)
}
fn is_terminal_safe(character: char) -> bool {
!character.is_control() || matches!(character, '\n' | '\t')
}
fn chat_schema() -> Result<OutputSchema, CliError> {
let message = Value::Object(Map::from_iter([(
"type".to_string(),
Value::String("string".to_string()),
)]));
let properties = Value::Object(Map::from_iter([("message".to_string(), message)]));
let schema = Value::Object(Map::from_iter([
("type".to_string(), Value::String("object".to_string())),
("properties".to_string(), properties),
(
"required".to_string(),
Value::Array(vec![Value::String("message".to_string())]),
),
("additionalProperties".to_string(), Value::Bool(false)),
]));
OutputSchema::new(schema).map_err(CliError::from)
}
#[derive(Debug, Error)]
enum CliError {
#[error("--base-url or {name} is required")]
BaseUrlRequired { name: &'static str },
#[error("one or more chat turns failed")]
ChatTurnsFailed,
#[error("--database, AG_HARNESS_ROOT, or HOME is required for durable session storage")]
DatabaseLocation,
#[error("No valid Git executable was found in PATH; pass --git-executable <FILE>")]
GitExecutableNotFound,
#[error("model output did not contain a message")]
MissingMessage,
#[error("stored session does not identify a supported built-in model")]
MissingModelIdentity,
#[error(transparent)]
Io(#[from] io::Error),
#[error(transparent)]
ModelConfiguration(ModelConfigurationError),
#[error(transparent)]
OutputSchema(#[from] ag_harness::OutputSchemaError),
#[error(transparent)]
Repository(#[from] ag_harness::RepositoryError),
#[error(transparent)]
Session(#[from] ag_harness::SessionError),
#[error(transparent)]
Turn(#[from] ag_harness::TurnError),
}
impl From<ModelConfigurationError> for CliError {
fn from(error: ModelConfigurationError) -> Self {
match error {
ModelConfigurationError::BaseUrl { name } => Self::BaseUrlRequired { name },
error => Self::ModelConfiguration(error),
}
}
}
#[cfg(test)]
mod tests {
use std::ffi::OsString;
use std::path::Path;
use std::sync::atomic::{AtomicUsize, Ordering};
use async_trait::async_trait;
use serde_json::json;
use wiremock::matchers::{body_string_contains, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
use super::*;
struct FixedModel(Value);
#[async_trait]
impl ag_harness::Model for FixedModel {
async fn complete(
&self,
_request: ag_harness::ModelRequest,
) -> Result<ag_harness::ModelCompletion, ag_harness::ModelError> {
Ok(ag_harness::ModelCompletion::from_response(
ag_harness::ModelResponse::Output(self.0.clone()),
))
}
}
struct FailOnceModel {
requests: AtomicUsize,
}
#[async_trait]
impl ag_harness::Model for FailOnceModel {
async fn complete(
&self,
_request: ag_harness::ModelRequest,
) -> Result<ag_harness::ModelCompletion, ag_harness::ModelError> {
if self.requests.fetch_add(1, Ordering::SeqCst) == 0 {
return Err(ag_harness::ModelError::InvalidResponse);
}
Ok(ag_harness::ModelCompletion::from_response(
ag_harness::ModelResponse::Output(json!({"message": "recovered"})),
))
}
}
fn provider_response(message: &str) -> ResponseTemplate {
ResponseTemplate::new(200).set_body_json(json!({
"choices": [{
"finish_reason": "stop",
"message": {"content": json!({"message": message}).to_string()}
}]
}))
}
fn run_arguments(command: Command) -> Option<RunArgs> {
match command {
Command::Run(args) => Some(args),
Command::Resume(_) => None,
}
}
fn resume_arguments(command: Command) -> Option<ResumeArgs> {
match command {
Command::Resume(args) => Some(args),
Command::Run(_) => None,
}
}
fn with_repository_controlled_git(mut cli: Cli) -> Result<Cli, io::Error> {
let git_executable = std::env::current_exe()?;
let repository_root = git_executable
.parent()
.map(Path::to_path_buf)
.ok_or_else(|| io::Error::other("test executable should have a parent"))?;
cli.git_executable = Some(git_executable);
match &mut cli.command {
Command::Run(arguments) => arguments.read_dir = repository_root,
Command::Resume(arguments) => arguments.read_dir = repository_root,
}
Ok(cli)
}
fn parse_cli<I, T>(arguments: I) -> Result<Cli, clap::Error>
where
I: IntoIterator<Item = T>,
T: Into<OsString>,
{
let mut arguments = arguments.into_iter().map(Into::into).collect::<Vec<_>>();
arguments.splice(
1..1,
[
OsString::from("--git-executable"),
test_git_executable().into_os_string(),
],
);
Cli::try_parse_from(arguments)
}
fn test_git_executable() -> PathBuf {
let executable_name = format!("git{}", std::env::consts::EXE_SUFFIX);
let path = std::env::var_os("PATH");
assert!(path.is_some(), "test PATH should be configured");
let executables = path
.iter()
.flat_map(|path| std::env::split_paths(path))
.filter(|directory| directory.is_absolute())
.map(|directory| directory.join(&executable_name))
.filter_map(|candidate| candidate.canonicalize().ok())
.filter(|path| path.is_file())
.collect::<Vec<_>>();
assert!(
!executables.is_empty(),
"trusted Git executable should be available on PATH"
);
executables[0].clone()
}
#[test]
fn cli_accepts_chat_with_or_without_an_initial_prompt() {
let without_prompt =
parse_cli(["ag-harness", "run", "muse-custom"]).expect("chat arguments should parse");
let with_prompt = parse_cli([
"ag-harness",
"run",
"muse-custom",
"Summarize this change",
"--provider",
"qwen",
"--base-url",
"https://models.example/v1",
"--read-dir",
"repo",
"--allow-write",
])
.expect("an initial prompt should parse");
let blank_prompt = parse_cli(["ag-harness", "run", "muse-custom", " "])
.expect_err("a blank initial prompt should be rejected");
let unknown_provider =
parse_cli(["ag-harness", "run", "muse-custom", "--provider", "unknown"])
.expect_err("an unknown provider should be rejected");
let without_prompt = run_arguments(without_prompt.command)
.expect("run command should contain run arguments");
assert_eq!(without_prompt.prompt, None);
assert!(!without_prompt.allow_write);
assert_eq!(without_prompt.provider, ModelProvider::Muse);
assert_eq!(without_prompt.read_dir, PathBuf::from("."));
let with_prompt =
run_arguments(with_prompt.command).expect("run command should contain run arguments");
assert_eq!(with_prompt.model, "muse-custom");
assert_eq!(with_prompt.prompt.as_deref(), Some("Summarize this change"));
assert_eq!(with_prompt.provider, ModelProvider::Qwen);
assert_eq!(
with_prompt.base_url.as_deref(),
Some("https://models.example/v1")
);
assert_eq!(with_prompt.read_dir, PathBuf::from("repo"));
assert!(with_prompt.allow_write);
assert!(
blank_prompt
.to_string()
.contains("prompt must contain a non-whitespace character")
);
assert!(
unknown_provider
.to_string()
.contains("invalid value 'unknown'")
);
}
#[test]
fn cli_accepts_every_catalog_provider() {
let providers = ModelProvider::all()
.iter()
.map(|provider| {
parse_cli([
"ag-harness",
"run",
"model-id",
"--provider",
provider.as_str(),
])
.expect("catalog provider should parse")
})
.collect::<Vec<_>>();
for (cli, expected) in providers.into_iter().zip(ModelProvider::all()) {
let args = run_arguments(cli.command).expect("run command should contain arguments");
assert_eq!(args.provider, *expected);
}
}
#[test]
fn cli_accepts_resume_and_database_override() {
let cli = parse_cli([
"ag-harness",
"--database",
"state.db",
"resume",
"session-a",
"continue",
"--allow-write",
])
.expect("resume arguments should parse");
assert_eq!(cli.database, Some(PathBuf::from("state.db")));
let args = resume_arguments(cli.command).expect("resume command should contain arguments");
assert_eq!(args.session, "session-a");
assert_eq!(args.prompt.as_deref(), Some("continue"));
assert!(args.allow_write);
let resume_probe = parse_cli(["ag-harness", "resume", "session-b"])
.expect("resume arguments should parse");
let run_probe =
parse_cli(["ag-harness", "run", "model"]).expect("run arguments should parse");
assert!(run_arguments(resume_probe.command).is_none());
assert!(resume_arguments(run_probe.command).is_none());
}
#[test]
fn database_path_uses_explicit_root_or_home_and_rejects_missing_storage_root() {
let explicit_environment_calls = std::cell::Cell::new(0);
let mut explicit_environment = |_: &str| {
explicit_environment_calls.set(explicit_environment_calls.get() + 1);
Err(env::VarError::NotPresent)
};
let explicit = database_path(
Some(PathBuf::from("explicit.db")),
&mut explicit_environment,
);
let rooted_variables =
std::collections::HashMap::from([("AG_HARNESS_ROOT", "/state/harness")]);
let rooted = database_path(None, &mut |name| {
rooted_variables
.get(name)
.map(ToString::to_string)
.ok_or(env::VarError::NotPresent)
});
let home_variables = std::collections::HashMap::from([("HOME", "/home/user")]);
let home = database_path(None, &mut |name| {
home_variables
.get(name)
.map(ToString::to_string)
.ok_or(env::VarError::NotPresent)
});
let missing = database_path(None, &mut |_| Err(env::VarError::NotPresent));
let empty_home = database_path(None, &mut |name| match name {
"HOME" => Ok(String::new()),
_ => Err(env::VarError::NotPresent),
});
assert_eq!(
explicit.expect("explicit database should resolve"),
PathBuf::from("explicit.db")
);
assert_eq!(explicit_environment_calls.get(), 0);
assert!(explicit_environment("unused").is_err());
assert_eq!(explicit_environment_calls.get(), 1);
assert_eq!(
rooted.expect("rooted database should resolve"),
PathBuf::from("/state/harness/db/harness.db")
);
assert_eq!(
home.expect("home database should resolve"),
PathBuf::from("/home/user/.ag-harness/db/harness.db")
);
assert!(matches!(missing, Err(CliError::DatabaseLocation)));
assert!(matches!(empty_home, Err(CliError::DatabaseLocation)));
}
#[test]
fn chat_mode_accounts_for_both_terminal_streams_and_initial_prompt() {
let with_prompt =
parse_cli(["ag-harness", "run", "muse", "hello"]).expect("chat arguments should parse");
let without_prompt =
parse_cli(["ag-harness", "run", "muse"]).expect("chat arguments should parse");
let resume = parse_cli(["ag-harness", "resume", "session-a", "continue"])
.expect("resume arguments should parse");
assert_eq!(
ChatMode::detect(&with_prompt, true, true),
ChatMode::Interactive
);
assert_eq!(
ChatMode::detect(&with_prompt, true, false),
ChatMode::OneShot
);
assert_eq!(
ChatMode::detect(&with_prompt, false, false),
ChatMode::NonInteractive
);
assert_eq!(
ChatMode::detect(&without_prompt, true, false),
ChatMode::NonInteractive
);
assert_eq!(ChatMode::detect(&resume, true, false), ChatMode::OneShot);
}
#[test]
fn stored_model_identity_maps_supported_providers_and_rejects_incomplete_identity() {
let muse = stored_model_identity_parts(Some("meta"), Some("muse-model"));
let kimi = stored_model_identity_parts(Some("moonshot_ai"), Some("kimi-model"));
let qwen = stored_model_identity_parts(Some("alibaba_cloud"), Some("qwen-model"));
let unknown = stored_model_identity_parts(Some("unknown"), Some("model"));
let missing_model = stored_model_identity_parts(Some("meta"), None);
assert!(matches!(muse, Ok((ModelProvider::Muse, model)) if model == "muse-model"));
assert!(matches!(kimi, Ok((ModelProvider::Kimi, model)) if model == "kimi-model"));
assert!(matches!(qwen, Ok((ModelProvider::Qwen, model)) if model == "qwen-model"));
assert!(matches!(unknown, Err(CliError::MissingModelIdentity)));
assert!(matches!(missing_model, Err(CliError::MissingModelIdentity)));
}
#[test]
fn exit_reporting_sanitizes_errors_and_preserves_success() {
let mut success_output = Vec::new();
let mut error_output = Vec::new();
let error = CliError::Turn(ag_harness::TurnError::Model(
ag_harness::ModelError::IncompleteResponse {
reason: "stop\u{1b}]52;c;Y2xpcGJvYXJk\u{7}".to_string(),
},
));
let success = report_exit(Ok(()), &mut success_output);
let failure = report_exit(Err(error), &mut error_output);
assert_eq!(success, ExitCode::SUCCESS);
assert_eq!(failure, ExitCode::FAILURE);
assert_eq!(success_output, [] as [u8; 0]);
assert!(!error_output.contains(&0x1b));
assert!(!error_output.contains(&0x07));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn execute_advertises_repository_reads_by_default() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/chat/completions"))
.and(body_string_contains(r#""name":"read""#))
.and(body_string_contains(r#""content":"Hello","role":"user""#))
.respond_with(provider_response("hello"))
.expect(1)
.mount(&server)
.await;
let storage = tempfile::tempdir().expect("temporary storage should exist");
let database = storage.path().join("harness.db");
let cli = parse_cli([
"ag-harness",
"run",
"muse-test",
"Hello",
"--base-url",
&server.uri(),
"--database",
&database.to_string_lossy(),
])
.expect("chat arguments should parse");
let input = BufReader::new(&b""[..]);
let mut output = Vec::new();
execute(
cli,
|_| Ok("test-key".to_string()),
input,
&mut output,
ChatMode::OneShot,
)
.await
.expect("chat with default repository reads should succeed");
assert!(
String::from_utf8(output)
.expect("chat output should be UTF-8")
.contains("assistant> hello\n---\n")
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn execute_resumes_a_saved_session_with_its_model_identity_and_history() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/chat/completions"))
.and(body_string_contains(r#""content":"second","role":"user""#))
.and(body_string_contains(r#""content":"first","role":"user""#))
.respond_with(provider_response("second answer"))
.with_priority(1)
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/chat/completions"))
.and(body_string_contains(r#""content":"first","role":"user""#))
.respond_with(provider_response("first answer"))
.with_priority(2)
.expect(1)
.mount(&server)
.await;
let storage = tempfile::tempdir().expect("temporary storage should exist");
let database = storage.path().join("harness.db");
let run = parse_cli([
"ag-harness",
"--database",
&database.to_string_lossy(),
"run",
"muse-test",
"first",
"--session",
"session-a",
"--base-url",
&server.uri(),
])
.expect("run arguments should parse");
let resume = parse_cli([
"ag-harness",
"--database",
&database.to_string_lossy(),
"resume",
"session-a",
"second",
"--base-url",
&server.uri(),
])
.expect("resume arguments should parse");
let invalid_resume = with_repository_controlled_git(
parse_cli([
"ag-harness",
"--database",
&database.to_string_lossy(),
"resume",
"session-a",
"third",
"--base-url",
&server.uri(),
])
.expect("invalid resume fixture should parse"),
)
.expect("repository-controlled Git fixture should resolve");
let mut first_output = Vec::new();
let mut second_output = Vec::new();
execute(
run,
|_| Ok("test-key".to_string()),
BufReader::new(&b""[..]),
&mut first_output,
ChatMode::OneShot,
)
.await
.expect("first process should create the session");
execute(
resume,
|_| Ok("test-key".to_string()),
BufReader::new(&b""[..]),
&mut second_output,
ChatMode::OneShot,
)
.await
.expect("second process should resume the session");
let invalid_error = execute(
invalid_resume,
|_| Ok("test-key".to_string()),
BufReader::new(&b""[..]),
Vec::new(),
ChatMode::OneShot,
)
.await
.expect_err("repository-controlled Git should reject resume");
let first_output = String::from_utf8(first_output).expect("output should be UTF-8");
let second_output = String::from_utf8(second_output).expect("output should be UTF-8");
assert!(first_output.contains("session: session-a\n"));
assert!(first_output.contains("assistant> first answer\n"));
assert!(second_output.contains("session: session-a\n"));
assert!(second_output.contains("assistant> second answer\n"));
assert!(matches!(
invalid_error,
CliError::Repository(ag_harness::RepositoryError::GitExecutableInsideRepository { .. })
));
}
#[tokio::test]
async fn execute_reports_missing_provider_credentials_before_creating_a_session() {
let cli = parse_cli([
"ag-harness",
"--database",
"unused.db",
"run",
"muse-test",
"Hello",
])
.expect("run arguments should parse");
let mut output = Vec::new();
let error = execute(
cli,
|_| Err(env::VarError::NotPresent),
BufReader::new(&b""[..]),
&mut output,
ChatMode::OneShot,
)
.await
.expect_err("missing credentials should fail");
assert!(matches!(
error,
CliError::ModelConfiguration(ModelConfigurationError::ApiKey { .. })
));
assert_eq!(output, [] as [u8; 0]);
}
#[test]
fn cli_accepts_the_default_git_executable() {
let arguments = [
"ag-harness",
"--database",
"unused.db",
"run",
"muse-test",
"Hello",
];
let cli = Cli::try_parse_from(arguments)
.expect("missing Git executable override should use the default");
assert_eq!(cli.git_executable, None);
}
#[test]
#[cfg(unix)]
fn git_executable_default_skips_a_non_executable_file() {
let storage = tempfile::tempdir().expect("temporary storage should exist");
let executable_name = format!("git{}", env::consts::EXE_SUFFIX);
let inert = storage.path().join(executable_name);
std::fs::write(&inert, "not executable").expect("inert Git fixture should be written");
let trusted_git = test_git_executable();
let trusted_directory = trusted_git
.parent()
.expect("trusted Git executable should have a parent");
let path = env::join_paths([storage.path(), trusted_directory])
.expect("test PATH should be valid");
let root = env::current_dir().expect("current directory should resolve");
let expected = Repository::new(&root, trusted_git)
.expect("trusted Git executable should configure the repository");
let actual = repository_from_path(&root, Some(path.as_os_str()));
assert_eq!(
actual.expect("non-executable Git should be skipped"),
expected
);
}
#[test]
fn git_executable_default_uses_the_process_path() {
let root = env::current_dir().expect("current directory should resolve");
let expected = Repository::new(&root, test_git_executable())
.expect("trusted Git executable should configure the repository");
let actual = repository_or_default(root, None)
.expect("test PATH should contain a trusted Git executable");
assert_eq!(actual, expected);
}
#[test]
fn git_executable_default_requires_git_on_path() {
let root = env::current_dir().expect("current directory should resolve");
let error = repository_from_path(&root, None)
.expect_err("missing PATH should not produce a Git executable");
assert!(matches!(error, CliError::GitExecutableNotFound));
}
#[test]
fn git_executable_default_preserves_repository_root_errors() {
let storage = tempfile::tempdir().expect("temporary storage should exist");
let missing_root = storage.path().join("missing-root");
let path = env::var_os("PATH").expect("test PATH should be configured");
let error = repository_from_path(&missing_root, Some(path.as_os_str()))
.expect_err("missing repository root should fail");
assert!(matches!(
error,
CliError::Repository(ag_harness::RepositoryError::Root { .. })
));
}
#[tokio::test]
async fn execute_rejects_repository_controlled_git_for_new_sessions() {
let cli = with_repository_controlled_git(
parse_cli([
"ag-harness",
"--database",
"unused.db",
"run",
"muse-test",
"Hello",
])
.expect("run arguments should parse"),
)
.expect("repository-controlled Git fixture should resolve");
let error = execute(
cli,
|_| Ok("test-key".to_string()),
BufReader::new(&b""[..]),
Vec::new(),
ChatMode::OneShot,
)
.await
.expect_err("repository-controlled Git should reject a new session");
assert!(matches!(
error,
CliError::Repository(ag_harness::RepositoryError::GitExecutableInsideRepository { .. })
));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn execute_advertises_writes_only_when_explicitly_enabled() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/chat/completions"))
.and(body_string_contains(READ_WRITE_SYSTEM_PROMPT))
.and(body_string_contains(r#""name":"read""#))
.and(body_string_contains(r#""name":"write""#))
.respond_with(provider_response("ready"))
.expect(1)
.mount(&server)
.await;
let storage = tempfile::tempdir().expect("temporary storage should exist");
let database = storage.path().join("harness.db");
let cli = parse_cli([
"ag-harness",
"run",
"muse-test",
"Hello",
"--base-url",
&server.uri(),
"--allow-write",
"--database",
&database.to_string_lossy(),
])
.expect("write-enabled chat arguments should parse");
let input = BufReader::new(&b""[..]);
let mut output = Vec::new();
execute(
cli,
|_| Ok("test-key".to_string()),
input,
&mut output,
ChatMode::OneShot,
)
.await
.expect("write-enabled chat should succeed");
assert!(
String::from_utf8(output)
.expect("chat output should be UTF-8")
.contains("assistant> ready\n---\n")
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn execute_advertises_read_only_with_an_explicit_directory() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/chat/completions"))
.and(body_string_contains(r#""name":"read""#))
.and(body_string_contains(r#""content":"Hello","role":"user""#))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"choices": [{
"finish_reason": "tool_calls",
"message": {
"content": null,
"tool_calls": [{
"id": "call-read",
"type": "function",
"function": {
"name": "read",
"arguments": r#"{"path":"input.txt"}"#
}
}]
}
}]
})))
.with_priority(2)
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/chat/completions"))
.and(body_string_contains(r#""tool_call_id":"call-read""#))
.respond_with(provider_response("hello"))
.with_priority(1)
.expect(1)
.mount(&server)
.await;
let repository = tempfile::TempDir::new().expect("temporary repository should exist");
let database = repository.path().join("harness.db");
std::fs::write(repository.path().join("input.txt"), "contents")
.expect("read fixture should be written");
let cli = parse_cli([
"ag-harness",
"run",
"muse-test",
"Hello",
"--base-url",
&server.uri(),
"--read-dir",
&repository.path().to_string_lossy(),
"--database",
&database.to_string_lossy(),
])
.expect("chat arguments should parse");
let input = BufReader::new(&b""[..]);
let mut output = Vec::new();
execute(
cli,
|_| Ok("test-key".to_string()),
input,
&mut output,
ChatMode::OneShot,
)
.await
.expect("chat with explicit read access should succeed");
let output = String::from_utf8(output).expect("chat output should be UTF-8");
assert!(output.contains("assistant> hello\n---\n"));
assert!(output.contains("tools:\n read input.txt (lines 1-1;"));
}
#[test]
fn cli_configuration_errors_preserve_cli_specific_guidance() {
let base_url = ModelConfigurationError::BaseUrl {
name: "KIMI_BASE_URL",
};
let api_key = ModelConfigurationError::ApiKey {
name: "MODEL_API_KEY",
};
let base_url = CliError::from(base_url);
let api_key = CliError::from(api_key);
assert_eq!(
base_url.to_string(),
"--base-url or KIMI_BASE_URL is required"
);
assert_eq!(api_key.to_string(), "MODEL_API_KEY is unavailable");
}
#[test]
fn chat_schema_requires_one_message_string() {
let schema = chat_schema().expect("chat schema should compile");
assert_eq!(schema.value()["required"], json!(["message"]));
assert_eq!(schema.value()["additionalProperties"], json!(false));
}
#[test]
fn line_endings_are_trimmed_without_changing_prompt_content() {
let mut unix = "hello\n".to_string();
let mut windows = "hello\r\n".to_string();
let mut unchanged = "hello".to_string();
trim_line_ending(&mut unix);
trim_line_ending(&mut windows);
trim_line_ending(&mut unchanged);
assert_eq!(unix, "hello");
assert_eq!(windows, "hello");
assert_eq!(unchanged, "hello");
}
#[test]
fn durations_have_compact_terminal_formatting() {
let short = format_duration(std::time::Duration::ZERO);
let measured = format_duration(std::time::Duration::from_millis(12));
assert_eq!(short, "<1 ms");
assert_eq!(measured, "12 ms");
}
#[tokio::test]
async fn interactive_chat_prints_prompts_and_handles_blank_input() {
let directory = tempfile::tempdir().expect("temporary directory should be created");
let repository = Repository::new(env!("CARGO_MANIFEST_DIR"), test_git_executable())
.expect("repository fixture should be valid");
let harness = Harness::new(FixedModel(json!({"message": "hello"})))
.database(directory.path().join("harness.db"))
.repository(repository)
.allow(Tool::Read);
let mut session = harness
.session(
"session-a",
chat_schema().expect("chat schema should compile"),
)
.create()
.await
.expect("session should be created");
let input = BufReader::new(&b"\nquestion\n"[..]);
let mut output = Vec::new();
run_chat(
&mut session,
"test-model",
None,
input,
&mut output,
ChatMode::Interactive,
)
.await
.expect("interactive chat should finish at EOF");
let output = String::from_utf8(output).expect("chat output should be UTF-8");
assert!(
output.starts_with("Chat with test-model. Ctrl-D to exit.\n>>> >>> assistant> hello\n")
);
assert!(output.contains("output; test-model; unavailable;"));
assert!(output.contains("tokens unavailable"));
assert!(output.ends_with("tools: none\n>>> "));
}
#[tokio::test]
async fn interactive_chat_continues_after_a_failed_turn() {
let directory = tempfile::tempdir().expect("temporary directory should be created");
let harness = Harness::new(FailOnceModel {
requests: AtomicUsize::new(0),
})
.database(directory.path().join("harness.db"));
let mut session = harness
.session(
"session-a",
chat_schema().expect("chat schema should compile"),
)
.create()
.await
.expect("session should be created");
let input = BufReader::new(&b"first\nretry\n"[..]);
let mut output = Vec::new();
run_chat(
&mut session,
"test-model",
None,
input,
&mut output,
ChatMode::Interactive,
)
.await
.expect("interactive chat should recover and finish at EOF");
let output = String::from_utf8(output).expect("chat output should be UTF-8");
assert!(output.contains("error: model returned no response content\n"));
assert!(output.contains(">>> assistant> recovered\n---\n"));
assert!(output.ends_with("tools: none\n>>> "));
}
#[tokio::test]
async fn noninteractive_chat_reports_a_failure_before_retrying() {
let directory = tempfile::tempdir().expect("temporary directory should be created");
let harness = Harness::new(FailOnceModel {
requests: AtomicUsize::new(0),
})
.database(directory.path().join("harness.db"));
let mut session = harness
.session(
"session-a",
chat_schema().expect("chat schema should compile"),
)
.create()
.await
.expect("session should be created");
let input = BufReader::new(&b"first\nretry\n"[..]);
let mut output = Vec::new();
let error = run_chat(
&mut session,
"test-model",
None,
input,
&mut output,
ChatMode::NonInteractive,
)
.await
.expect_err("a recovered chat should retain its failed exit status");
assert!(matches!(error, CliError::ChatTurnsFailed));
let output = String::from_utf8(output).expect("chat output should be UTF-8");
assert!(output.starts_with("error: model returned no response content\n"));
assert!(output.contains("assistant> recovered\n---\n"));
}
#[tokio::test]
async fn noninteractive_chat_returns_the_last_failure_at_eof() {
let directory = tempfile::tempdir().expect("temporary directory should be created");
let harness = Harness::new(FailOnceModel {
requests: AtomicUsize::new(0),
})
.database(directory.path().join("harness.db"));
let mut session = harness
.session(
"session-a",
chat_schema().expect("chat schema should compile"),
)
.create()
.await
.expect("session should be created");
let input = BufReader::new(&b"first\n"[..]);
let mut output = Vec::new();
let error = run_chat(
&mut session,
"test-model",
None,
input,
&mut output,
ChatMode::NonInteractive,
)
.await
.expect_err("the final failed turn should be returned at EOF");
assert!(matches!(error, CliError::ChatTurnsFailed));
assert_eq!(
String::from_utf8(output).expect("chat output should be UTF-8"),
"error: model returned no response content\n"
);
}
#[tokio::test]
async fn chat_rejects_model_output_that_violates_schema() {
let directory = tempfile::tempdir().expect("temporary directory should be created");
let harness = Harness::new(FixedModel(json!({"unexpected": true})))
.database(directory.path().join("harness.db"));
let mut session = harness
.session(
"session-a",
chat_schema().expect("chat schema should compile"),
)
.create()
.await
.expect("session should be created");
let input = BufReader::new(&b""[..]);
let mut output = Vec::new();
let error = run_chat(
&mut session,
"test-model",
Some("question".to_string()),
input,
&mut output,
ChatMode::OneShot,
)
.await
.expect_err("schema-invalid output should fail");
assert!(matches!(
error,
CliError::Session(ag_harness::SessionError::Turn(
ag_harness::TurnError::Model(
ag_harness::ModelError::SchemaViolation { path, .. }
)
)) if path == "$"
));
assert_eq!(output, [] as [u8; 0]);
}
#[tokio::test]
async fn one_shot_chat_does_not_read_follow_up_terminal_input() {
let directory = tempfile::tempdir().expect("temporary directory should be created");
let harness = Harness::new(FixedModel(json!({"message": "hello"})))
.database(directory.path().join("harness.db"));
let mut session = harness
.session(
"session-a",
chat_schema().expect("chat schema should compile"),
)
.create()
.await
.expect("session should be created");
let input = BufReader::new(&b"unexpected follow-up\n"[..]);
let mut output = Vec::new();
run_chat(
&mut session,
"test-model",
Some("question".to_string()),
input,
&mut output,
ChatMode::OneShot,
)
.await
.expect("one-shot chat should finish after the initial prompt");
let output = String::from_utf8(output).expect("chat output should be UTF-8");
assert_eq!(output.matches("assistant> hello\n---\n").count(), 1);
assert!(!output.contains("Chat with"));
assert!(!output.contains(">>>"));
}
#[tokio::test]
async fn one_shot_chat_returns_turn_failures() {
let directory = tempfile::tempdir().expect("temporary directory should be created");
let harness = Harness::new(FailOnceModel {
requests: AtomicUsize::new(0),
})
.database(directory.path().join("harness.db"));
let mut session = harness
.session(
"session-a",
chat_schema().expect("chat schema should compile"),
)
.create()
.await
.expect("session should be created");
let input = BufReader::new(&b""[..]);
let mut output = Vec::new();
let error = run_chat(
&mut session,
"test-model",
Some("question".to_string()),
input,
&mut output,
ChatMode::OneShot,
)
.await
.expect_err("one-shot chat should return its failed turn");
assert!(matches!(error, CliError::Session(_)));
assert_eq!(output, [] as [u8; 0]);
}
#[test]
fn terminal_text_replaces_control_sequences_and_preserves_safe_whitespace() {
let text = "before\n\t\u{1b}]52;c;Y2xpcGJvYXJk\u{7}after\r";
let sanitized = terminal_text(text);
assert_eq!(
sanitized,
"before\n\t\u{fffd}]52;c;Y2xpcGJvYXJk\u{fffd}after\u{fffd}"
);
assert!(
sanitized
.chars()
.all(|character| !character.is_control() || matches!(character, '\n' | '\t'))
);
}
#[test]
fn assistant_text_indents_continuation_lines_and_sanitizes_them() {
let text = "answer\n---\nturn: forged\u{1b}";
let framed = assistant_text(text);
assert_eq!(
framed,
"assistant> answer\n ---\n turn: forged\u{fffd}\n"
);
}
#[test]
fn single_line_terminal_text_replaces_all_control_characters() {
let text = "model\nname\t\u{1b}";
let sanitized = single_line_terminal_text(text);
assert_eq!(sanitized, "model\u{fffd}name\u{fffd}\u{fffd}");
assert!(sanitized.chars().all(|character| !character.is_control()));
}
#[test]
fn usage_format_marks_missing_counts() {
let usage = ag_harness::CompletionUsage::new(None, None, None, Some(4), None, None);
let formatted = format_usage(&usage);
assert_eq!(formatted, "tokens ? in, 4 out, ? total");
}
}