use serde_json::Value;
use std::{
fs,
io::{Read, Write},
net::{TcpListener, TcpStream},
path::PathBuf,
sync::{
Arc, Mutex,
atomic::{AtomicBool, AtomicUsize, Ordering},
},
thread,
time::{Duration, Instant, SystemTime, UNIX_EPOCH},
};
use ziyuan_app::{
Session,
text_project::{NewProject, Project, Provenance, TextProjects},
translation::{ReferencePreview, ReferenceRequest, ReferenceSelection, ReferenceSubmission},
};
use ziyuan_core::{Capability, CourtEvent, Department, FileChange, Seal, TaskId, TranslationKind};
use ziyuan_governance::{Court, TaskStatus};
use ziyuan_llm::{OpenAiCompatible, Retry};
use ziyuan_runtime::{AgentBackend, CancelToken, Limits, Scheduler};
use ziyuan_storage::{EventLog, EventStore};
const B1: &str = "c0001-b0001";
const B3: &str = "c0003-b0001";
const SOURCES: [&str; 3] = [
"ALPHA_SOURCE_UNSELECTED",
"BETA_SOURCE_SELECTED",
"GAMMA_SOURCE_UNSELECTED",
];
const TARGETS: [&str; 3] = [
"ALPHA_TARGET_UNSELECTED",
"BETA_TARGET_SELECTED",
"GAMMA_TARGET_UNSELECTED",
];
struct Fixture {
root: PathBuf,
project: Project,
}
impl Fixture {
fn new() -> Self {
static NEXT: AtomicUsize = AtomicUsize::new(0);
let root = std::env::temp_dir().join(format!(
"ziyuan-reference-http-{}-{}-{}",
std::process::id(),
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos(),
NEXT.fetch_add(1, Ordering::Relaxed)
));
fs::create_dir(&root).unwrap();
Session::found(EventLog::new(&root)).unwrap();
let texts = TextProjects::new(&root);
let project = texts
.create(NewProject {
title: "HTTP references".into(),
source_language: "en".into(),
target_language: "zh-CN".into(),
provenance: Provenance::default(),
})
.unwrap();
let files = [
SOURCES.join("\n\n"),
"Neighbor context without reference prose.".into(),
"Actual translation target.".into(),
];
let paths: Vec<_> = files
.iter()
.enumerate()
.map(|(i, body)| {
let path = root.join(format!("chapter-{i}.txt"));
fs::write(&path, body).unwrap();
path
})
.collect();
let project = texts.import(&project.id, &paths).unwrap();
let fixture = Self { root, project };
fixture.draft(B1, true);
fixture
}
fn session(&self) -> Session<EventLog> {
Session::open(EventLog::new(&self.root)).unwrap()
}
fn draft_path(&self) -> PathBuf {
self.root
.join(format!("translations/{}/{B1}/zh.json", self.project.id))
}
fn draft(&self, block: &str, finalize: bool) {
let mut session = self.session();
let task = session
.issue_translation(&self.root, &self.project.id, block)
.unwrap()
.id;
let worker = session
.court()
.route(&session.court().task(task).unwrap().edict)
.unwrap()
.worker;
let mut candidate = self
.project
.chapters
.iter()
.flat_map(|c| &c.blocks)
.find(|b| b.block_id == block)
.unwrap()
.clone();
for (i, paragraph) in candidate.paragraphs.iter_mut().enumerate() {
paragraph.text = if block == B1 {
TARGETS[i].into()
} else {
"Existing target draft.".into()
};
}
let events = vec![
CourtEvent::TaskAssigned {
task,
worker,
model: None,
},
CourtEvent::ChangesProposed {
task,
changes: vec![FileChange {
path: format!("translations/{}/{block}/zh.json", self.project.id),
before: None,
after: String::from_utf8(candidate.serialized().unwrap()).unwrap(),
}],
},
CourtEvent::TaskCompleted {
task,
worker,
summary: "fixture done".into(),
},
];
let mut log = EventLog::new(&self.root);
let history = log.load().unwrap();
let mut court = Court::replay(&history).unwrap();
for event in &events {
court.apply(event).unwrap();
}
log.append(history.len(), &events).unwrap();
assert!(
session
.apply_changes(task, &self.root, vec![])
.unwrap()
.is_empty()
);
if finalize {
session
.finalize_translation(&self.root, &self.project.id, block)
.unwrap();
}
}
fn selected(&self, kind: TranslationKind) -> (ReferencePreview, ReferenceSubmission) {
let session = self.session();
let pairs = session.memory_pairs(&self.root, &self.project.id).unwrap();
assert_eq!(pairs.pairs.len(), 3);
let beta = pairs
.pairs
.iter()
.find(|p| p.source.contains(SOURCES[1]))
.unwrap();
let request = ReferenceRequest {
selections: vec![ReferenceSelection {
identity: beta.identity.clone(),
digests: beta.digests.clone(),
origin: beta.origin.clone(),
}],
authorized_projects: Default::default(),
};
let preview = session
.preview_references(&self.root, &self.project.id, B3, kind, &request)
.unwrap();
let submission = ReferenceSubmission {
request,
fingerprint: preview.fingerprint.clone(),
submission: "loopback-submission".into(),
};
(preview, submission)
}
fn issue(&self, kind: TranslationKind, submission: &ReferenceSubmission) -> TaskId {
self.session()
.issue_with_references(&self.root, &self.project.id, B3, kind, submission)
.unwrap()
.id
}
}
impl Drop for Fixture {
fn drop(&mut self) {
let _ = fs::remove_dir_all(&self.root);
}
}
struct Server {
url: String,
requests: Arc<Mutex<Vec<Value>>>,
stop: Arc<AtomicBool>,
worker: Option<thread::JoinHandle<()>>,
}
impl Server {
fn new(callback: impl FnOnce() + Send + 'static) -> Self {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
listener.set_nonblocking(true).unwrap();
let url = format!("http://{}/v1", listener.local_addr().unwrap());
let requests = Arc::new(Mutex::new(Vec::new()));
let stop = Arc::new(AtomicBool::new(false));
let (captured, stopped) = (requests.clone(), stop.clone());
let worker = thread::spawn(move || {
let deadline = Instant::now() + Duration::from_secs(20);
let mut callback = Some(callback);
while !stopped.load(Ordering::SeqCst) && Instant::now() < deadline {
match listener.accept() {
Ok((mut stream, _)) => {
let request = read_request(&mut stream);
captured.lock().unwrap().push(request);
if let Some(callback) = callback.take() {
callback();
}
let body = "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"结论:离线执行测试。\\n依据:未调用工具。\\n未办之事:未修改译稿。\"},\"finish_reason\":null}]}\n\ndata: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n";
write!(stream, "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).unwrap();
}
Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
thread::sleep(Duration::from_millis(5))
}
Err(e) => panic!("loopback accept: {e}"),
}
}
});
Self {
url,
requests,
stop,
worker: Some(worker),
}
}
fn scheduler(&self, fixture: &Fixture) -> Scheduler {
let model =
OpenAiCompatible::new(&self.url, None, "loopback-reference-test").with_retry(Retry {
attempts: 0,
base: Duration::ZERO,
max: Duration::ZERO,
});
Scheduler::new(
Arc::new(AgentBackend::new(Box::new(model), fixture.root.clone(), 2)),
Limits {
max_concurrent: 1,
timeout: Duration::from_secs(8),
},
)
}
fn captured(&self) -> Vec<Value> {
self.requests.lock().unwrap().clone()
}
}
impl Drop for Server {
fn drop(&mut self) {
self.stop.store(true, Ordering::SeqCst);
if let Some(worker) = self.worker.take() {
worker.join().expect("loopback server");
}
}
}
fn read_request(stream: &mut TcpStream) -> Value {
stream.set_nonblocking(false).unwrap();
stream
.set_read_timeout(Some(Duration::from_secs(3)))
.unwrap();
stream
.set_write_timeout(Some(Duration::from_secs(3)))
.unwrap();
let mut bytes = Vec::new();
let mut chunk = [0; 4096];
let (header_end, length) = loop {
let n = stream.read(&mut chunk).unwrap();
assert!(n > 0, "HTTP closed before headers");
bytes.extend_from_slice(&chunk[..n]);
assert!(bytes.len() < 2_000_000);
if let Some(end) = bytes.windows(4).position(|w| w == b"\r\n\r\n") {
let header = std::str::from_utf8(&bytes[..end]).unwrap();
assert!(
header.starts_with("POST /v1/chat/completions HTTP/1.1"),
"{header}"
);
let length = header
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().unwrap())
})
.expect("content length");
break (end + 4, length);
}
};
while bytes.len() < header_end + length {
let n = stream.read(&mut chunk).unwrap();
assert!(n > 0, "HTTP closed before body");
bytes.extend_from_slice(&chunk[..n]);
}
serde_json::from_slice(&bytes[header_end..header_end + length]).unwrap()
}
fn assert_writer_tools(request: &Value) {
let mut names: Vec<_> = request["tools"]
.as_array()
.unwrap()
.iter()
.map(|t| t["function"]["name"].as_str().unwrap())
.collect();
names.sort_unstable();
assert_eq!(
names,
["list_files", "read_file", "search_records", "write_file"]
);
}
fn assert_failed(fixture: &Fixture, task: TaskId) {
let session = fixture.session();
assert!(
matches!(
&session.court().task(task).unwrap().status,
TaskStatus::Failed { reason, .. } if !reason.is_empty()
),
"{:?}",
session.court().task(task).unwrap().status
);
let view = session.reference_task(&fixture.root, task).unwrap();
assert!(!view.current);
assert!(view.reason.is_some());
assert!(!session.history().unwrap().iter().any(|event| matches!(
event,
CourtEvent::ModelUsage { task: id, .. } | CourtEvent::ToolInvoked { task: id, .. } if *id == task
)));
assert!(
session
.history()
.unwrap()
.iter()
.any(|e| matches!(e, CourtEvent::TaskFailed { task: id, .. } if *id == task))
);
}
#[tokio::test]
async fn selected_beta_reaches_real_http_for_translate_and_revise() {
for kind in [TranslationKind::Translate, TranslationKind::Revise] {
let fixture = Fixture::new();
if kind == TranslationKind::Revise {
fixture.draft(B3, false);
}
let (preview, submission) = fixture.selected(kind);
let task = fixture.issue(kind, &submission);
let mut session = fixture.session();
assert_eq!(
session.court().task(task).unwrap().edict.seal.capabilities,
[Capability::ReadFile, Capability::WriteFile].into()
);
let server = Server::new(|| {});
session
.run(
&[task],
&server.scheduler(&fixture),
&CancelToken::new(),
|_| {},
)
.await
.unwrap();
assert!(
matches!(
session.court().task(task).unwrap().status,
TaskStatus::Completed { .. }
),
"{:?}",
session.court().task(task).unwrap().status
);
let requests = server.captured();
assert_eq!(requests.len(), 2);
for request in &requests {
let messages = request["messages"].as_array().unwrap();
assert_eq!(
messages
.iter()
.filter_map(|m| m["content"].as_str())
.map(|text| text.matches(&preview.prompt).count())
.sum::<usize>(),
1,
"the exact confirmed payload must occur once, without rewriting or duplication"
);
assert_eq!(
messages
.iter()
.filter(|m| m["role"] == "user"
&& m["content"]
.as_str()
.is_some_and(|s| s.contains(&preview.prompt)))
.count(),
1
);
assert!(preview.prompt.contains(SOURCES[1]) && preview.prompt.contains(TARGETS[1]));
let wire = serde_json::to_string(request).unwrap();
for excluded in [SOURCES[0], SOURCES[2], TARGETS[0], TARGETS[2]] {
assert!(!preview.prompt.contains(excluded));
assert!(
!wire.contains(excluded),
"unselected prose leaked: {excluded}"
);
}
assert_writer_tools(request);
}
}
}
#[tokio::test]
async fn no_reference_control_uses_same_real_writer_tools() {
let fixture = Fixture::new();
let mut session = fixture.session();
let task = session
.issue_translation(&fixture.root, &fixture.project.id, B3)
.unwrap()
.id;
assert!(session.court().task(task).unwrap().references.is_none());
assert_eq!(
session.court().task(task).unwrap().edict.seal.capabilities,
[Capability::ReadFile, Capability::WriteFile].into()
);
let server = Server::new(|| {});
session
.run(
&[task],
&server.scheduler(&fixture),
&CancelToken::new(),
|_| {},
)
.await
.unwrap();
assert!(matches!(
session.court().task(task).unwrap().status,
TaskStatus::Completed { .. }
));
let requests = server.captured();
assert_eq!(requests.len(), 2, "one no-tools-used send-back is expected");
assert_writer_tools(&requests[0]);
let wire = serde_json::to_string(&requests).unwrap();
for prose in SOURCES.into_iter().chain(TARGETS) {
assert!(!wire.contains(prose));
}
}
#[tokio::test]
async fn stale_sources_fail_without_http_and_retry_recovers_original_task() {
for reopen in [false, true] {
let fixture = Fixture::new();
let (preview, submission) = fixture.selected(TranslationKind::Translate);
let task = fixture.issue(TranslationKind::Translate, &submission);
if reopen {
fixture
.session()
.reopen_translation(&fixture.root, &fixture.project.id, B1, "source withdrawn")
.unwrap();
} else {
let path = fixture.draft_path();
fs::write(
&path,
fs::read_to_string(&path)
.unwrap()
.replace(TARGETS[1], "Edited selected translation"),
)
.unwrap();
}
let mut session = fixture.session();
let before = session.event_count();
assert_eq!(
session
.issue_with_references(
&fixture.root,
&fixture.project.id,
B3,
TranslationKind::Translate,
&submission
)
.unwrap()
.id,
task
);
assert_eq!(session.event_count(), before);
let view = session.reference_task(&fixture.root, task).unwrap();
assert!(!view.current);
assert_eq!(view.snapshot, preview);
let server = Server::new(|| {});
session
.run(
&[task],
&server.scheduler(&fixture),
&CancelToken::new(),
|_| {},
)
.await
.unwrap();
assert!(server.captured().is_empty());
assert_failed(&fixture, task);
}
}
#[tokio::test]
async fn queued_reference_is_rechecked_after_worker_permit() {
let fixture = Fixture::new();
let (_, submission) = fixture.selected(TranslationKind::Translate);
let mut session = fixture.session();
let first = session
.issue(
Department::Works,
"Unrelated first task; reply briefly.".into(),
Seal {
capabilities: [Capability::ReadFile].into(),
scope: vec![".".into()],
mcp: Default::default(),
},
)
.unwrap()
.id;
let bound = fixture.issue(TranslationKind::Translate, &submission);
let root = fixture.root.clone();
let path = fixture.draft_path();
let changed = Arc::new(AtomicBool::new(false));
let did_change = changed.clone();
let server = Server::new(move || {
let queued = Session::open(EventLog::new(&root)).unwrap();
assert!(matches!(
queued.court().task(bound).unwrap().status,
TaskStatus::Assigned(_)
));
fs::write(
&path,
fs::read_to_string(&path)
.unwrap()
.replace(TARGETS[1], "Changed while queued"),
)
.unwrap();
did_change.store(true, Ordering::SeqCst);
});
session
.run(
&[first, bound],
&server.scheduler(&fixture),
&CancelToken::new(),
|_| {},
)
.await
.unwrap();
assert!(changed.load(Ordering::SeqCst));
let requests = server.captured();
assert_eq!(
requests.len(),
2,
"only the first task and its no-tools-used send-back may call HTTP"
);
for request in &requests {
assert!(request["messages"].as_array().unwrap().iter().any(|m| {
m["role"] == "user"
&& m["content"]
.as_str()
.is_some_and(|s| s.contains("Unrelated first task; reply briefly."))
}));
assert!(!serde_json::to_string(request).unwrap().contains(SOURCES[1]));
}
assert!(matches!(
session.court().task(first).unwrap().status,
TaskStatus::Completed { .. }
));
assert_failed(&fixture, bound);
}