use std::path::Path;
use std::sync::Arc;
use newt_core::Caveats;
use newt_inference::{ChatRequest, InferenceBackend};
use crate::emission::{normalize_emission, Emission};
use crate::error::{CoderError, Result};
use crate::prompt::{build_prompt, build_reprompt, CoderPrompt};
pub struct Coder {
backend: Arc<dyn InferenceBackend>,
}
#[derive(Debug, Clone)]
pub struct CoderRun {
pub emission_shape: String,
pub model_id: String,
pub files_written: Vec<String>,
pub raw_reply: String,
pub first_emission: String,
}
impl Coder {
pub fn new(backend: Arc<dyn InferenceBackend>) -> Self {
Self { backend }
}
pub async fn run(&self, workspace: &Path, task: &str, caveats: &Caveats) -> Result<CoderRun> {
let prompt = build_prompt(workspace, task)?;
check_fs_read(caveats, &prompt)?;
tracing::info!(
files_included = prompt.included_files.len(),
user_chars = prompt.user.len(),
"newt-coder prompt built"
);
let mut calls_used: u64 = 0;
check_call_budget(caveats, calls_used)?;
check_net(caveats, self.backend.as_ref())?;
let req = ChatRequest::new().system(prompt.system).user(prompt.user);
let reply = self
.backend
.complete(req)
.await
.map_err(|e| CoderError::Inference(e.to_string()))?;
calls_used += 1;
let raw = reply.content.clone();
let model_id = reply.model_id.clone();
let emission = normalize_emission(&raw)?;
let shape_label = emission.shape_label().to_string();
match self.apply(&emission, workspace, caveats) {
Ok(files_written) => {
tracing::info!(
emission_shape = %shape_label,
files_written = files_written.len(),
"newt-coder run complete"
);
Ok(CoderRun {
emission_shape: shape_label,
model_id,
files_written,
first_emission: raw.clone(),
raw_reply: raw,
})
}
Err(first_err)
if matches!(emission, Emission::UnifiedDiff(_))
|| matches!(first_err, CoderError::LooksLikeDiff { .. }) =>
{
tracing::warn!(
error = %first_err,
"newt-coder: diff-shaped emission did not apply, re-prompting for whole files"
);
self.reprompt_whole_files(workspace, task, raw, first_err, calls_used, caveats)
.await
}
Err(other) => Err(other),
}
}
async fn reprompt_whole_files(
&self,
workspace: &Path,
task: &str,
first_raw: String,
original_err: CoderError,
calls_used: u64,
caveats: &Caveats,
) -> Result<CoderRun> {
if !caveats.max_calls.permits_one_more(calls_used) {
tracing::warn!(
calls_used,
"newt-coder: re-prompt skipped, max_calls budget exhausted"
);
return Err(original_err);
}
let prompt = match build_reprompt(workspace, task) {
Ok(p) => p,
Err(e) => {
tracing::warn!(error = %e, "newt-coder: re-prompt build failed");
return Err(original_err);
}
};
if let Err(e) = check_fs_read(caveats, &prompt) {
tracing::warn!(error = %e, "newt-coder: re-prompt fs_read denied");
return Err(original_err);
}
let req = ChatRequest::new().system(prompt.system).user(prompt.user);
let reply = match self.backend.complete(req).await {
Ok(r) => r,
Err(e) => {
tracing::warn!(error = %e, "newt-coder: re-prompt inference failed");
return Err(original_err);
}
};
let retry_raw = reply.content.clone();
let model_id = reply.model_id.clone();
let emission = match normalize_emission(&retry_raw) {
Ok(em @ Emission::WholeFiles(_)) => em,
Ok(other) => {
tracing::warn!(
emission_shape = %other.shape_label(),
"newt-coder: re-prompt did not return whole files"
);
return Err(original_err);
}
Err(e) => {
tracing::warn!(error = %e, "newt-coder: re-prompt emission malformed");
return Err(original_err);
}
};
let shape_label = emission.shape_label().to_string();
match self.apply(&emission, workspace, caveats) {
Ok(files_written) => {
tracing::info!(
emission_shape = %shape_label,
files_written = files_written.len(),
"newt-coder: re-prompt whole-file fallback applied"
);
Ok(CoderRun {
emission_shape: shape_label,
model_id,
files_written,
first_emission: first_raw.clone(),
raw_reply: format!(
"[diff-apply failed, re-prompted for whole files]\n\
--- first reply ---\n{first_raw}\n\
--- retry reply ---\n{retry_raw}"
),
})
}
Err(e) => {
tracing::warn!(
error = %e,
"newt-coder: re-prompt whole-file apply failed"
);
Err(original_err)
}
}
}
fn apply(
&self,
emission: &Emission,
workspace: &Path,
caveats: &Caveats,
) -> Result<Vec<String>> {
match emission {
Emission::WholeFiles(files) => {
for (path, contents) in files {
reject_bad_shape(path, contents)?;
}
for path in files.keys() {
if !caveats.permits_fs_write(path) {
return Err(CoderError::CapabilityDenied {
kind: "fs_write",
target: path.clone(),
});
}
}
let pairs: Vec<(String, String)> =
files.iter().map(|(k, v)| (k.clone(), v.clone())).collect();
let written = newt_tools::apply_whole_files(workspace, pairs)
.map_err(|e| CoderError::FileWrite(e.to_string()))?;
Ok(written)
}
Emission::UnifiedDiff(diff) => {
if !matches!(caveats.fs_write, newt_core::Scope::All) {
return Err(CoderError::CapabilityDenied {
kind: "fs_write",
target: "<unified_diff: paths not enumerable>".to_string(),
});
}
newt_tools::apply_patch(workspace, diff)
.map_err(|e| CoderError::FileWrite(e.to_string()))?;
Ok(Vec::new())
}
Emission::Prose(prose) => {
tracing::warn!(
prose_len = prose.len(),
"newt-coder: prose-only emission, no edits"
);
Ok(Vec::new())
}
}
}
}
fn check_call_budget(caveats: &Caveats, used_so_far: u64) -> Result<()> {
if caveats.max_calls.permits_one_more(used_so_far) {
Ok(())
} else {
Err(CoderError::CapabilityDenied {
kind: "max_calls",
target: format!("turn #{}", used_so_far + 1),
})
}
}
fn check_net(caveats: &Caveats, backend: &dyn InferenceBackend) -> Result<()> {
let endpoint = match backend.endpoint() {
Some(e) => e,
None => return Ok(()),
};
let host = host_from_endpoint(endpoint);
if caveats.permits_net(host) {
Ok(())
} else {
Err(CoderError::CapabilityDenied {
kind: "net",
target: host.to_string(),
})
}
}
fn check_fs_read(caveats: &Caveats, prompt: &CoderPrompt) -> Result<()> {
for path in &prompt.included_files {
let s = path.to_string_lossy();
if !caveats.permits_fs_read(&s) {
return Err(CoderError::CapabilityDenied {
kind: "fs_read",
target: s.into_owned(),
});
}
}
Ok(())
}
fn host_from_endpoint(endpoint: &str) -> &str {
let after_scheme = endpoint
.find("://")
.map(|i| &endpoint[i + 3..])
.unwrap_or(endpoint);
let end = after_scheme
.find(['/', ':', '?'])
.unwrap_or(after_scheme.len());
&after_scheme[..end]
}
fn reject_bad_shape(path: &str, contents: &str) -> Result<()> {
let first_non_blank = contents.lines().find(|l| !l.trim().is_empty());
match first_non_blank {
None => Err(CoderError::EmptyEmission {
path: path.to_string(),
}),
Some(first) => {
let trimmed = first.trim_start();
if trimmed.starts_with("--- ")
|| trimmed.starts_with("+++ ")
|| trimmed.starts_with("@@")
{
return Err(CoderError::LooksLikeDiff {
path: path.to_string(),
});
}
if trimmed.starts_with("FILE:") {
return Err(CoderError::LeakedMarker {
path: path.to_string(),
});
}
Ok(())
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::BTreeMap;
use std::fs;
use tempfile::TempDir;
fn coder_with_no_backend_used() -> Coder {
struct Stub;
#[async_trait::async_trait]
impl InferenceBackend for Stub {
fn name(&self) -> &str {
"stub"
}
fn model_id(&self) -> &str {
"stub-model"
}
fn supports_tier(&self, _t: newt_core::router::Tier) -> bool {
false
}
async fn complete(
&self,
_req: ChatRequest,
) -> anyhow::Result<newt_inference::ChatReply> {
unreachable!("apply tests do not call the backend")
}
}
Coder::new(Arc::new(Stub))
}
#[test]
fn apply_whole_files_writes_to_workspace() {
let tmp = TempDir::new().unwrap();
let coder = coder_with_no_backend_used();
let mut files = BTreeMap::new();
files.insert("src/lib.rs".to_string(), "pub fn hello() {}\n".to_string());
let written = coder
.apply(&Emission::WholeFiles(files), tmp.path(), &Caveats::top())
.unwrap();
assert_eq!(written, vec!["src/lib.rs".to_string()]);
let content = fs::read_to_string(tmp.path().join("src/lib.rs")).unwrap();
assert_eq!(content, "pub fn hello() {}\n");
}
#[test]
fn apply_prose_writes_nothing() {
let tmp = TempDir::new().unwrap();
let coder = coder_with_no_backend_used();
let written = coder
.apply(
&Emission::Prose("I've updated it.".to_string()),
tmp.path(),
&Caveats::top(),
)
.unwrap();
assert!(written.is_empty());
}
#[test]
fn apply_unified_diff_returns_empty_files_written() {
let tmp = TempDir::new().unwrap();
fs::write(tmp.path().join("a.txt"), "old\n").unwrap();
let diff = "\
--- a/a.txt
+++ b/a.txt
@@ -1 +1 @@
-old
+new
";
let coder = coder_with_no_backend_used();
let written = coder
.apply(
&Emission::UnifiedDiff(diff.to_string()),
tmp.path(),
&Caveats::top(),
)
.unwrap();
assert!(written.is_empty(), "diff path returns empty files_written");
let content = fs::read_to_string(tmp.path().join("a.txt")).unwrap();
assert_eq!(content, "new\n");
}
#[test]
fn apply_bad_diff_surfaces_filewrite_error() {
let tmp = TempDir::new().unwrap();
let coder = coder_with_no_backend_used();
let bad = Emission::UnifiedDiff("not a real diff".to_string());
let err = coder.apply(&bad, tmp.path(), &Caveats::top()).unwrap_err();
assert!(matches!(err, CoderError::FileWrite(_)));
}
fn whole_files(path: &str, contents: &str) -> Emission {
let mut m = BTreeMap::new();
m.insert(path.to_string(), contents.to_string());
Emission::WholeFiles(m)
}
#[test]
fn apply_whole_files_accepts_line_one_change() {
let tmp = TempDir::new().unwrap();
let coder = coder_with_no_backend_used();
fs::create_dir_all(tmp.path().join("src")).unwrap();
fs::write(
tmp.path().join("src/lib.rs"),
"pub fn hello(name: &str) -> String {\n format!(\"hi {name}\")\n}\n",
)
.unwrap();
let new_body = "pub fn greet(name: &str) -> String {\n format!(\"hi {name}\")\n}\n";
let written = coder
.apply(
&whole_files("src/lib.rs", new_body),
tmp.path(),
&Caveats::top(),
)
.unwrap();
assert_eq!(written, vec!["src/lib.rs".to_string()]);
assert_eq!(
fs::read_to_string(tmp.path().join("src/lib.rs")).unwrap(),
new_body
);
}
#[test]
fn apply_whole_files_rejects_diff_shaped_contents() {
let tmp = TempDir::new().unwrap();
let coder = coder_with_no_backend_used();
fs::write(tmp.path().join("a.txt"), "old\n").unwrap();
let diff = "--- a/a.txt\n+++ b/a.txt\n@@ -1 +1 @@\n-old\n+new\n";
let err = coder
.apply(&whole_files("a.txt", diff), tmp.path(), &Caveats::top())
.unwrap_err();
assert!(matches!(err, CoderError::LooksLikeDiff { ref path } if path == "a.txt"));
assert_eq!(
fs::read_to_string(tmp.path().join("a.txt")).unwrap(),
"old\n"
);
}
#[test]
fn apply_whole_files_rejects_hunk_only_contents() {
let tmp = TempDir::new().unwrap();
let coder = coder_with_no_backend_used();
let hunk = "@@ -1,2 +1,2 @@\n-old\n+new\n";
let err = coder
.apply(&whole_files("a.txt", hunk), tmp.path(), &Caveats::top())
.unwrap_err();
assert!(matches!(err, CoderError::LooksLikeDiff { .. }));
}
#[test]
fn apply_whole_files_rejects_empty_contents() {
let tmp = TempDir::new().unwrap();
let coder = coder_with_no_backend_used();
let err = coder
.apply(&whole_files("a.txt", ""), tmp.path(), &Caveats::top())
.unwrap_err();
assert!(matches!(err, CoderError::EmptyEmission { ref path } if path == "a.txt"));
}
#[test]
fn apply_whole_files_rejects_whitespace_only_contents() {
let tmp = TempDir::new().unwrap();
let coder = coder_with_no_backend_used();
let err = coder
.apply(
&whole_files("a.txt", " \n\t\n"),
tmp.path(),
&Caveats::top(),
)
.unwrap_err();
assert!(matches!(err, CoderError::EmptyEmission { .. }));
}
#[test]
fn apply_whole_files_rejects_leaked_file_marker() {
let tmp = TempDir::new().unwrap();
let coder = coder_with_no_backend_used();
let body = "FILE: src/lib.rs\npub fn add(a: i32, b: i32) -> i32 { a + b }\n";
let err = coder
.apply(
&whole_files("src/lib.rs", body),
tmp.path(),
&Caveats::top(),
)
.unwrap_err();
assert!(matches!(err, CoderError::LeakedMarker { ref path } if path == "src/lib.rs"));
}
#[test]
fn reject_bad_shape_messages_start_with_file_write_failed() {
for err in [
super::reject_bad_shape("p", "").unwrap_err(),
super::reject_bad_shape("p", "--- a/p\n").unwrap_err(),
super::reject_bad_shape("p", "FILE: p\n").unwrap_err(),
] {
assert!(
err.to_string().starts_with("file write failed:"),
"message did not start with prefix: {err}"
);
}
}
#[test]
fn apply_whole_files_denies_path_outside_fs_write_scope() {
let tmp = TempDir::new().unwrap();
let coder = coder_with_no_backend_used();
let caveats = Caveats {
fs_write: newt_core::Scope::only(["allowed.rs".to_string()]),
..Caveats::top()
};
let allowed = coder
.apply(
&whole_files("allowed.rs", "fn ok() {}\n"),
tmp.path(),
&caveats,
)
.expect("permitted write must succeed");
assert_eq!(allowed, vec!["allowed.rs".to_string()]);
let err = coder
.apply(
&whole_files("forbidden.rs", "fn evil() {}\n"),
tmp.path(),
&caveats,
)
.unwrap_err();
match err {
CoderError::CapabilityDenied { kind, target } => {
assert_eq!(kind, "fs_write");
assert_eq!(target, "forbidden.rs");
}
other => panic!("expected CapabilityDenied, got {other:?}"),
}
assert!(!tmp.path().join("forbidden.rs").exists());
}
#[test]
fn apply_whole_files_denies_atomically_on_partial_scope() {
let tmp = TempDir::new().unwrap();
let coder = coder_with_no_backend_used();
let caveats = Caveats {
fs_write: newt_core::Scope::only(["a.rs".to_string()]),
..Caveats::top()
};
let mut files = BTreeMap::new();
files.insert("a.rs".to_string(), "fn a() {}\n".to_string());
files.insert("b.rs".to_string(), "fn b() {}\n".to_string());
let err = coder
.apply(&Emission::WholeFiles(files), tmp.path(), &caveats)
.unwrap_err();
assert!(matches!(
err,
CoderError::CapabilityDenied {
kind: "fs_write",
..
}
));
assert!(!tmp.path().join("a.rs").exists());
assert!(!tmp.path().join("b.rs").exists());
}
#[test]
fn apply_unified_diff_denied_under_bounded_fs_write() {
let tmp = TempDir::new().unwrap();
let coder = coder_with_no_backend_used();
let caveats = Caveats {
fs_write: newt_core::Scope::only(["whatever.rs".to_string()]),
..Caveats::top()
};
let diff = Emission::UnifiedDiff(
"--- a/whatever.rs\n+++ b/whatever.rs\n@@ -1 +1 @@\n-x\n+y\n".to_string(),
);
let err = coder.apply(&diff, tmp.path(), &caveats).unwrap_err();
assert!(matches!(
err,
CoderError::CapabilityDenied {
kind: "fs_write",
..
}
));
}
#[test]
fn host_from_endpoint_strips_scheme_and_path() {
assert_eq!(
super::host_from_endpoint("http://localhost:11434/api/chat"),
"localhost"
);
assert_eq!(
super::host_from_endpoint("https://allowed.example.com/v1/chat"),
"allowed.example.com"
);
assert_eq!(
super::host_from_endpoint("bare.host.local"),
"bare.host.local"
);
assert_eq!(super::host_from_endpoint("http://h:8080"), "h");
assert_eq!(super::host_from_endpoint("https://only.host/"), "only.host");
}
#[test]
fn check_call_budget_passes_under_unlimited() {
super::check_call_budget(&Caveats::top(), 0).unwrap();
super::check_call_budget(&Caveats::top(), 999_999).unwrap();
}
#[test]
fn check_call_budget_passes_within_bound() {
let caveats = Caveats {
max_calls: newt_core::CountBound::AtMost(3),
..Caveats::top()
};
super::check_call_budget(&caveats, 0).unwrap();
super::check_call_budget(&caveats, 2).unwrap();
}
#[test]
fn check_call_budget_denies_at_bound() {
let caveats = Caveats {
max_calls: newt_core::CountBound::AtMost(2),
..Caveats::top()
};
let err = super::check_call_budget(&caveats, 2).unwrap_err();
match err {
CoderError::CapabilityDenied { kind, target } => {
assert_eq!(kind, "max_calls");
assert!(target.contains("#3"));
}
other => panic!("expected CapabilityDenied, got {other:?}"),
}
}
}