use std::borrow::Cow;
#[cfg(not(test))]
use std::io::IsTerminal as _;
use std::path::PathBuf;
use std::process::ExitCode;
use std::{env, io};
use ag_harness::{
ChatSession, Harness, KimiConfig, ModelClient, MuseConfig, OutputSchema, QwenConfig, Tool,
TurnOutcome,
};
use clap::{Args, Parser, Subcommand, ValueEnum};
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. When a user asks about repository ",
"contents, call the read tool immediately in the same response 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. When a user asks about ",
"repository contents, call the read tool immediately in the same response 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. To create an empty file, pass a write patch containing ",
"only `--- /dev/null` and `+++ b/<path>` headers with no hunk."
);
const KIMI_API_KEY_ENV: &str = "KIMI_API_KEY";
const KIMI_BASE_URL_ENV: &str = "KIMI_BASE_URL";
const MUSE_DEFAULT_BASE_URL: &str = "https://api.meta.ai/v1";
const MUSE_API_KEY_ENV: &str = "MODEL_API_KEY";
const MUSE_BASE_URL_ENV: &str = "MODEL_API_BASE_URL";
const QWEN_API_KEY_ENV: &str = "DASHSCOPE_API_KEY";
const QWEN_BASE_URL_ENV: &str = "DASHSCOPE_BASE_URL";
const PROVIDER_HELP: &str = "\
Supported models (other endpoint-supported model IDs also work):
muse: muse-spark-1.2, muse-spark-1.2-contributor
kimi: kimi-k2.6
qwen: qwen-plus
Credentials:
muse: MODEL_API_KEY (MODEL_API_BASE_URL optional)
kimi: KIMI_API_KEY, KIMI_BASE_URL
qwen: DASHSCOPE_API_KEY, DASHSCOPE_BASE_URL";
#[derive(Debug, Parser)]
#[command(
name = "ag-harness",
version,
about = "Chats with models through a repository harness",
after_help = PROVIDER_HELP
)]
struct Cli {
#[command(subcommand)]
command: Command,
}
#[derive(Debug, Subcommand)]
enum Command {
Run(RunArgs),
}
#[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, value_enum, default_value_t = Provider::Muse)]
provider: Provider,
#[arg(long, value_name = "DIR", default_value = ".")]
read_dir: PathBuf,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)]
enum Provider {
Muse,
Kimi,
Qwen,
}
impl Provider {
fn api_key_environment(self) -> &'static str {
match self {
Self::Muse => MUSE_API_KEY_ENV,
Self::Kimi => KIMI_API_KEY_ENV,
Self::Qwen => QWEN_API_KEY_ENV,
}
}
fn base_url_environment(self) -> &'static str {
match self {
Self::Muse => MUSE_BASE_URL_ENV,
Self::Kimi => KIMI_BASE_URL_ENV,
Self::Qwen => QWEN_BASE_URL_ENV,
}
}
fn default_base_url(self) -> Option<&'static str> {
match self {
Self::Muse => Some(MUSE_DEFAULT_BASE_URL),
Self::Kimi | Self::Qwen => None,
}
}
fn model_client(
self,
configuration: ModelConfiguration,
) -> Result<ModelClient, ag_harness::ModelMetadataError> {
let ModelConfiguration {
api_key,
base_url,
model,
} = configuration;
match self {
Self::Muse => ModelClient::muse(MuseConfig {
api_key,
base_url,
model,
}),
Self::Kimi => ModelClient::kimi(KimiConfig {
api_key,
base_url,
model,
}),
Self::Qwen => ModelClient::qwen(QwenConfig {
api_key,
base_url,
model,
}),
}
}
}
struct ModelConfiguration {
api_key: String,
base_url: String,
model: String,
}
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 Command::Run(args) = &cli.command;
if stdin_is_terminal && args.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,
environment: impl FnMut(&str) -> Result<String, env::VarError>,
input: Input,
output: Output,
mode: ChatMode,
) -> Result<(), CliError>
where
Input: AsyncBufRead + Unpin,
Output: AsyncWrite + Unpin,
{
let Command::Run(args) = cli.command;
let client = model_client(&args, environment)?;
let mut harness = Harness::new(client)
.repository(args.read_dir.clone())
.allow(Tool::Read);
let system_prompt = if args.allow_write {
harness = harness.allow(Tool::Write);
READ_WRITE_SYSTEM_PROMPT
} else {
READ_ONLY_SYSTEM_PROMPT
};
let mut session = harness
.chat(chat_schema()?)
.with_system_prompt(system_prompt);
run_chat(&mut session, &args.model, args.prompt, input, output, mode).await
}
async fn run_chat<Input, Output>(
session: &mut ChatSession<'_>,
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 prompt = if let Some(prompt) = pending_prompt.take() {
prompt
} else {
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 {
break;
}
trim_line_ending(&mut prompt);
prompt
};
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 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: &ag_harness::TurnError,
) -> 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)
}
fn model_client(
args: &RunArgs,
environment: impl FnMut(&str) -> Result<String, env::VarError>,
) -> Result<ModelClient, CliError> {
let configuration = model_config(args, environment)?;
Ok(args.provider.model_client(configuration)?)
}
fn model_config(
args: &RunArgs,
mut environment: impl FnMut(&str) -> Result<String, env::VarError>,
) -> Result<ModelConfiguration, CliError> {
let api_key_environment = args.provider.api_key_environment();
let api_key = environment(api_key_environment).map_err(|_| CliError::ApiKeyUnavailable {
name: api_key_environment,
})?;
let base_url_environment = args.provider.base_url_environment();
let base_url = if let Some(base_url) = &args.base_url {
base_url.clone()
} else {
match environment(base_url_environment) {
Ok(base_url) => base_url,
Err(env::VarError::NotPresent) => {
args.provider.default_base_url().map(str::to_string).ok_or(
CliError::BaseUrlRequired {
name: base_url_environment,
},
)?
}
Err(source) => {
return Err(CliError::Environment {
name: base_url_environment,
source,
});
}
}
};
Ok(ModelConfiguration {
api_key,
base_url,
model: args.model.clone(),
})
}
#[derive(Debug, Error)]
enum CliError {
#[error("{name} is unavailable")]
ApiKeyUnavailable { name: &'static str },
#[error("--base-url or {name} is required")]
BaseUrlRequired { name: &'static str },
#[error("one or more chat turns failed")]
ChatTurnsFailed,
#[error("{name} is unavailable: {source}")]
Environment {
name: &'static str,
source: env::VarError,
},
#[error("model output did not contain a message")]
MissingMessage,
#[error(transparent)]
Io(#[from] io::Error),
#[error(transparent)]
ModelConfiguration(#[from] ag_harness::ModelMetadataError),
#[error(transparent)]
OutputSchema(#[from] ag_harness::OutputSchemaError),
#[error(transparent)]
Turn(#[from] ag_harness::TurnError),
}
#[cfg(test)]
mod tests {
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::ModelResponse, ag_harness::ModelError> {
Ok(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::ModelResponse, ag_harness::ModelError> {
if self.requests.fetch_add(1, Ordering::SeqCst) == 0 {
return Err(ag_harness::ModelError::InvalidResponse);
}
Ok(ag_harness::ModelResponse::Output(
json!({"message": "recovered"}),
))
}
}
fn run_args() -> RunArgs {
RunArgs {
allow_write: false,
base_url: None,
model: "muse-model".to_string(),
prompt: None,
provider: Provider::Muse,
read_dir: PathBuf::from("."),
}
}
fn missing_environment(_: &str) -> Result<String, env::VarError> {
Err(env::VarError::NotPresent)
}
fn provider_response(message: &str) -> ResponseTemplate {
ResponseTemplate::new(200).set_body_json(json!({
"choices": [{
"finish_reason": "stop",
"message": {"content": json!({"message": message}).to_string()}
}]
}))
}
#[test]
fn cli_accepts_chat_with_or_without_an_initial_prompt() {
let without_prompt = Cli::try_parse_from(["ag-harness", "run", "muse-custom"])
.expect("chat arguments should parse");
let with_prompt = Cli::try_parse_from([
"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 = Cli::try_parse_from(["ag-harness", "run", "muse-custom", " "])
.expect_err("a blank initial prompt should be rejected");
let unknown_provider =
Cli::try_parse_from(["ag-harness", "run", "muse-custom", "--provider", "unknown"])
.expect_err("an unknown provider should be rejected");
let Command::Run(without_prompt) = without_prompt.command;
assert_eq!(without_prompt.prompt, None);
assert!(!without_prompt.allow_write);
assert_eq!(without_prompt.provider, Provider::Muse);
assert_eq!(without_prompt.read_dir, PathBuf::from("."));
let Command::Run(with_prompt) = with_prompt.command;
assert_eq!(with_prompt.model, "muse-custom");
assert_eq!(with_prompt.prompt.as_deref(), Some("Summarize this change"));
assert_eq!(with_prompt.provider, Provider::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 chat_mode_accounts_for_both_terminal_streams_and_initial_prompt() {
let with_prompt = Cli::try_parse_from(["ag-harness", "run", "muse", "hello"])
.expect("chat arguments should parse");
let without_prompt = Cli::try_parse_from(["ag-harness", "run", "muse"])
.expect("chat 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
);
}
#[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 cli = Cli::try_parse_from([
"ag-harness",
"run",
"muse-test",
"Hello",
"--base-url",
&server.uri(),
])
.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")
.starts_with("assistant> hello\n---\n")
);
}
#[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 cli = Cli::try_parse_from([
"ag-harness",
"run",
"muse-test",
"Hello",
"--base-url",
&server.uri(),
"--allow-write",
])
.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")
.starts_with("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");
std::fs::write(repository.path().join("input.txt"), "contents")
.expect("read fixture should be written");
let cli = Cli::try_parse_from([
"ag-harness",
"run",
"muse-test",
"Hello",
"--base-url",
&server.uri(),
"--read-dir",
&repository.path().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.starts_with("assistant> hello\n---\n"));
assert!(output.contains("tools:\n read input.txt (lines 1-1;"));
}
#[test]
fn model_configuration_uses_environment_defaults() {
let args = run_args();
let config = model_config(&args, |name| match name {
MUSE_API_KEY_ENV => Ok("test-key".to_string()),
_ => Err(env::VarError::NotPresent),
})
.expect("default model configuration should be valid");
assert_eq!(config.api_key, "test-key");
assert_eq!(config.base_url, MUSE_DEFAULT_BASE_URL);
assert_eq!(config.model, "muse-model");
}
#[test]
fn model_configuration_uses_provider_environment() {
let providers = [
(Provider::Muse, MUSE_API_KEY_ENV, MUSE_BASE_URL_ENV),
(Provider::Kimi, KIMI_API_KEY_ENV, KIMI_BASE_URL_ENV),
(Provider::Qwen, QWEN_API_KEY_ENV, QWEN_BASE_URL_ENV),
];
for (provider, api_key_environment, base_url_environment) in providers {
let mut args = run_args();
args.provider = provider;
let mut requested_environment = Vec::new();
let config = model_config(&args, |name| {
requested_environment.push(name.to_string());
if name == api_key_environment {
Ok("provider-key".to_string())
} else {
assert_eq!(name, base_url_environment);
Ok("https://provider.example/v1".to_string())
}
})
.expect("provider environment should produce a valid configuration");
assert_eq!(
requested_environment,
[
api_key_environment.to_string(),
base_url_environment.to_string()
]
);
assert_eq!(config.api_key, "provider-key");
assert_eq!(config.base_url, "https://provider.example/v1");
}
}
#[test]
fn model_clients_select_every_supported_provider() {
let providers = [
(Provider::Muse, "meta"),
(Provider::Kimi, "moonshot_ai"),
(Provider::Qwen, "alibaba_cloud"),
];
let identities = providers.map(|(provider, expected_provider)| {
let mut args = run_args();
args.base_url = Some("https://models.example/v1".to_string());
args.provider = provider;
let client = model_client(&args, |_| Ok("test-key".to_string()))
.expect("supported provider configuration should be valid");
(client.metadata().provider().to_string(), expected_provider)
});
for (provider, expected_provider) in identities {
assert_eq!(provider, expected_provider);
}
}
#[test]
fn model_configuration_uses_environment_base_url() {
let args = run_args();
let config = model_config(&args, |name| {
if name == MUSE_BASE_URL_ENV {
Ok("https://environment.example/v1".to_string())
} else {
Ok("test-key".to_string())
}
})
.expect("environment model configuration should be valid");
assert_eq!(config.api_key, "test-key");
assert_eq!(config.base_url, "https://environment.example/v1");
}
#[test]
fn model_configuration_prefers_base_url_flag() {
let mut args = run_args();
args.base_url = Some("https://cli.example/v1".to_string());
let config = model_config(&args, |_| Ok("test-key".to_string()))
.expect("CLI overrides should produce valid configuration");
assert_eq!(config.base_url, "https://cli.example/v1");
}
#[test]
fn non_muse_configuration_requires_a_base_url() {
let mut args = run_args();
args.provider = Provider::Kimi;
let error = model_config(&args, |name| {
if name == KIMI_API_KEY_ENV {
Ok("test-key".to_string())
} else {
Err(env::VarError::NotPresent)
}
})
.err()
.expect("Kimi without an endpoint should be rejected");
assert_eq!(error.to_string(), "--base-url or KIMI_BASE_URL is required");
}
#[test]
fn model_configuration_reports_missing_api_key() {
let args = run_args();
let error = model_config(&args, missing_environment)
.err()
.expect("a missing API key should be rejected");
assert_eq!(error.to_string(), "MODEL_API_KEY is unavailable");
}
#[test]
fn api_key_errors_redact_non_unicode_values_from_output() {
let args = run_args();
let secret = "visible-secret-material";
let error = model_config(&args, |_| {
Err(env::VarError::NotUnicode(std::ffi::OsString::from(secret)))
})
.err()
.expect("a non-Unicode API key should be rejected");
let mut error_output = Vec::new();
let exit = report_exit(Err(error), &mut error_output);
assert_eq!(exit, ExitCode::FAILURE);
let error_output = String::from_utf8(error_output).expect("error output should be UTF-8");
assert_eq!(error_output, "MODEL_API_KEY is unavailable\n");
assert!(!error_output.contains(secret));
}
#[test]
fn model_configuration_reports_non_unicode_optional_environment() {
let args = run_args();
let error = model_config(&args, |name| {
if name == MUSE_BASE_URL_ENV {
Err(env::VarError::NotUnicode("invalid".into()))
} else {
Ok("test-key".to_string())
}
})
.err()
.expect("non-Unicode configuration should be rejected");
assert!(
error
.to_string()
.starts_with("MODEL_API_BASE_URL 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 harness = Harness::new(FixedModel(json!({"message": "hello"})))
.repository(".")
.allow(Tool::Read);
let mut session = harness.chat(chat_schema().expect("chat schema should compile"));
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 harness = Harness::new(FailOnceModel {
requests: AtomicUsize::new(0),
});
let mut session = harness.chat(chat_schema().expect("chat schema should compile"));
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 harness = Harness::new(FailOnceModel {
requests: AtomicUsize::new(0),
});
let mut session = harness.chat(chat_schema().expect("chat schema should compile"));
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 harness = Harness::new(FailOnceModel {
requests: AtomicUsize::new(0),
});
let mut session = harness.chat(chat_schema().expect("chat schema should compile"));
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_without_a_message() {
let harness = Harness::new(FixedModel(json!({"unexpected": true})));
let mut session = harness.chat(chat_schema().expect("chat schema should compile"));
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::NonInteractive,
)
.await
.expect_err("missing message output should fail");
assert!(matches!(error, CliError::MissingMessage));
}
#[tokio::test]
async fn one_shot_chat_does_not_read_follow_up_terminal_input() {
let harness = Harness::new(FixedModel(json!({"message": "hello"})));
let mut session = harness.chat(chat_schema().expect("chat schema should compile"));
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 harness = Harness::new(FailOnceModel {
requests: AtomicUsize::new(0),
});
let mut session = harness.chat(chat_schema().expect("chat schema should compile"));
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::Turn(_)));
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");
}
}