1134 lines
38 KiB
Rust
1134 lines
38 KiB
Rust
use makepad_studio_hub::{HubConfig, MountConfig, StudioHub};
|
|
use makepad_studio_protocol::hub_protocol::{
|
|
AiMessageRole, AiMountState, ClientToHub, HubToClient,
|
|
};
|
|
use std::fs;
|
|
use std::io::{Read, Write};
|
|
use std::net::TcpListener;
|
|
use std::sync::{Mutex, OnceLock};
|
|
use std::thread;
|
|
use std::time::{Duration, Instant};
|
|
|
|
fn ai_env_lock() -> &'static Mutex<()> {
|
|
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
|
|
LOCK.get_or_init(|| Mutex::new(()))
|
|
}
|
|
|
|
struct EnvGuard {
|
|
key: &'static str,
|
|
previous: Option<String>,
|
|
}
|
|
|
|
impl EnvGuard {
|
|
fn set(key: &'static str, value: String) -> Self {
|
|
let previous = std::env::var(key).ok();
|
|
std::env::set_var(key, value);
|
|
Self { key, previous }
|
|
}
|
|
}
|
|
|
|
impl Drop for EnvGuard {
|
|
fn drop(&mut self) {
|
|
if let Some(previous) = self.previous.take() {
|
|
std::env::set_var(self.key, previous);
|
|
} else {
|
|
std::env::remove_var(self.key);
|
|
}
|
|
}
|
|
}
|
|
|
|
fn read_http_request(stream: &mut std::net::TcpStream) -> String {
|
|
stream
|
|
.set_read_timeout(Some(Duration::from_secs(2)))
|
|
.expect("set read timeout");
|
|
let mut request = Vec::new();
|
|
let mut buf = [0u8; 4096];
|
|
let mut header_end = None;
|
|
let mut content_length = 0usize;
|
|
|
|
loop {
|
|
let read = stream.read(&mut buf).expect("read request bytes");
|
|
if read == 0 {
|
|
break;
|
|
}
|
|
request.extend_from_slice(&buf[..read]);
|
|
if header_end.is_none() {
|
|
if let Some(pos) = request.windows(4).position(|window| window == b"\r\n\r\n") {
|
|
let end = pos + 4;
|
|
header_end = Some(end);
|
|
let headers = String::from_utf8_lossy(&request[..end]);
|
|
for line in headers.lines() {
|
|
if let Some((name, value)) = line.split_once(':') {
|
|
if name.eq_ignore_ascii_case("content-length") {
|
|
content_length = value.trim().parse().unwrap_or(0);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
if let Some(end) = header_end {
|
|
if request.len() >= end + content_length {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
String::from_utf8_lossy(&request).to_string()
|
|
}
|
|
|
|
fn write_chunked_sse(stream: &mut std::net::TcpStream, chunks: &[&str]) {
|
|
let headers = concat!(
|
|
"HTTP/1.1 200 OK\r\n",
|
|
"Content-Type: text/event-stream\r\n",
|
|
"Transfer-Encoding: chunked\r\n",
|
|
"Connection: close\r\n\r\n"
|
|
);
|
|
stream
|
|
.write_all(headers.as_bytes())
|
|
.expect("write sse headers");
|
|
stream.flush().expect("flush sse headers");
|
|
for chunk in chunks {
|
|
let bytes = chunk.as_bytes();
|
|
let prefix = format!("{:X}\r\n", bytes.len());
|
|
stream
|
|
.write_all(prefix.as_bytes())
|
|
.expect("write sse chunk prefix");
|
|
stream.write_all(bytes).expect("write sse chunk body");
|
|
stream.write_all(b"\r\n").expect("write sse chunk trailer");
|
|
stream.flush().expect("flush sse chunk");
|
|
thread::sleep(Duration::from_millis(60));
|
|
}
|
|
stream
|
|
.write_all(b"0\r\n\r\n")
|
|
.expect("write sse terminator");
|
|
stream.flush().expect("flush sse terminator");
|
|
}
|
|
|
|
fn write_chunked_sse_and_hold_open(
|
|
mut stream: std::net::TcpStream,
|
|
chunks: &[&str],
|
|
hold_open_for: Duration,
|
|
) {
|
|
write_chunked_sse(&mut stream, chunks);
|
|
thread::sleep(hold_open_for);
|
|
}
|
|
|
|
fn wait_for_ai_state(
|
|
connection: &makepad_studio_hub::HubConnection,
|
|
mount: &str,
|
|
timeout: Duration,
|
|
predicate: impl Fn(&AiMountState) -> bool,
|
|
) -> AiMountState {
|
|
let deadline = Instant::now() + timeout;
|
|
while Instant::now() < deadline {
|
|
let Some(msg) = connection.recv_timeout(Duration::from_millis(100)) else {
|
|
continue;
|
|
};
|
|
if let HubToClient::AiMountState { mount: got, state } = msg {
|
|
if got == mount && predicate(&state) {
|
|
return state;
|
|
}
|
|
}
|
|
}
|
|
panic!("timed out waiting for AiMountState for mount {}", mount);
|
|
}
|
|
|
|
fn wait_for_message(
|
|
connection: &makepad_studio_hub::HubConnection,
|
|
timeout: Duration,
|
|
predicate: impl Fn(&HubToClient) -> bool,
|
|
) -> HubToClient {
|
|
let deadline = Instant::now() + timeout;
|
|
while Instant::now() < deadline {
|
|
let Some(msg) = connection.recv_timeout(Duration::from_millis(100)) else {
|
|
continue;
|
|
};
|
|
if predicate(&msg) {
|
|
return msg;
|
|
}
|
|
}
|
|
panic!("timed out waiting for HubToClient message");
|
|
}
|
|
|
|
#[test]
|
|
fn ai_manager_round_trips_prompt_through_local_backend() {
|
|
let _env_lock = ai_env_lock()
|
|
.lock()
|
|
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
|
|
|
let listener = TcpListener::bind("127.0.0.1:0").expect("bind local ai server");
|
|
let addr = listener.local_addr().expect("local addr");
|
|
let _base_url = EnvGuard::set(
|
|
"MAKEPAD_STUDIO_AI_BASE_URL",
|
|
format!("http://{}/v1/chat/completions", addr),
|
|
);
|
|
|
|
let server = thread::spawn(move || {
|
|
let (mut stream, _) = listener.accept().expect("accept ai request");
|
|
let request_text = read_http_request(&mut stream);
|
|
assert!(request_text.contains("POST /v1/chat/completions"));
|
|
assert!(request_text.contains("\"stream\":true"));
|
|
write_chunked_sse(
|
|
&mut stream,
|
|
&[
|
|
"data: {\"choices\":[{\"delta\":{\"content\":\"assistant reply\"}}]}\n\n",
|
|
"data: [DONE]\n\n",
|
|
],
|
|
);
|
|
});
|
|
|
|
let root = std::env::current_dir().expect("current_dir");
|
|
let config = HubConfig {
|
|
mounts: vec![MountConfig {
|
|
name: "repo".to_string(),
|
|
path: root,
|
|
}],
|
|
enable_in_process_gateway: false,
|
|
..Default::default()
|
|
};
|
|
let mut connection = StudioHub::start_in_process(config).expect("start backend");
|
|
|
|
let _ = connection.send(ClientToHub::AiGetState {
|
|
mount: "repo".to_string(),
|
|
});
|
|
let initial = wait_for_ai_state(&connection, "repo", Duration::from_secs(2), |state| {
|
|
state.active_agent.is_some()
|
|
});
|
|
let agent_id = initial.active_agent_id.expect("default ai agent");
|
|
|
|
let _ = connection.send(ClientToHub::AiSetBackend {
|
|
mount: "repo".to_string(),
|
|
backend_id: "openai_localhost".to_string(),
|
|
});
|
|
let _ = wait_for_ai_state(&connection, "repo", Duration::from_secs(2), |state| {
|
|
state.active_backend_id.as_deref() == Some("openai_localhost")
|
|
});
|
|
|
|
let _ = connection.send(ClientToHub::AiSendPrompt {
|
|
mount: "repo".to_string(),
|
|
agent_id,
|
|
text: "hello from test".to_string(),
|
|
});
|
|
|
|
let _pending = wait_for_ai_state(&connection, "repo", Duration::from_secs(2), |state| {
|
|
state
|
|
.active_agent
|
|
.as_ref()
|
|
.map(|agent| agent.pending)
|
|
.unwrap_or(false)
|
|
});
|
|
|
|
let done = wait_for_ai_state(&connection, "repo", Duration::from_secs(5), |state| {
|
|
state
|
|
.active_agent
|
|
.as_ref()
|
|
.map(|agent| {
|
|
!agent.pending
|
|
&& agent
|
|
.messages
|
|
.iter()
|
|
.any(|message| message.text == "assistant reply")
|
|
})
|
|
.unwrap_or(false)
|
|
});
|
|
|
|
let messages = &done.active_agent.expect("active agent").messages;
|
|
assert!(messages
|
|
.iter()
|
|
.any(|message| message.text == "hello from test"));
|
|
assert!(messages
|
|
.iter()
|
|
.any(|message| message.text == "assistant reply"));
|
|
|
|
server.join().expect("join ai server");
|
|
}
|
|
|
|
#[test]
|
|
fn ai_manager_persists_chats_per_mount_and_loads_them_on_restart() {
|
|
let _env_lock = ai_env_lock()
|
|
.lock()
|
|
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
|
|
|
let listener = TcpListener::bind("127.0.0.1:0").expect("bind local ai server");
|
|
let addr = listener.local_addr().expect("local addr");
|
|
let _base_url = EnvGuard::set(
|
|
"MAKEPAD_STUDIO_AI_BASE_URL",
|
|
format!("http://{}/v1/chat/completions", addr),
|
|
);
|
|
|
|
let repo = makepad_studio_hub::test_support::tempdir().expect("tempdir");
|
|
let config = HubConfig {
|
|
mounts: vec![MountConfig {
|
|
name: "repo".to_string(),
|
|
path: repo.path().to_path_buf(),
|
|
}],
|
|
enable_in_process_gateway: false,
|
|
..Default::default()
|
|
};
|
|
|
|
let server = thread::spawn(move || {
|
|
let (mut stream, _) = listener.accept().expect("accept ai request");
|
|
let request_text = read_http_request(&mut stream);
|
|
assert!(request_text.contains("\"stream\":true"));
|
|
write_chunked_sse(
|
|
&mut stream,
|
|
&[
|
|
"data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"thinking now\"}}]}\n\n",
|
|
"data: {\"choices\":[{\"delta\":{\"content\":\"assistant reply\"}}]}\n\n",
|
|
"data: [DONE]\n\n",
|
|
],
|
|
);
|
|
});
|
|
|
|
{
|
|
let mut connection = StudioHub::start_in_process(config.clone()).expect("start backend");
|
|
|
|
let _ = connection.send(ClientToHub::AiGetState {
|
|
mount: "repo".to_string(),
|
|
});
|
|
let initial = wait_for_ai_state(&connection, "repo", Duration::from_secs(2), |state| {
|
|
state.active_agent.is_some()
|
|
});
|
|
let agent_id = initial.active_agent_id.expect("default ai agent");
|
|
|
|
let _ = connection.send(ClientToHub::AiSendPrompt {
|
|
mount: "repo".to_string(),
|
|
agent_id,
|
|
text: "persist this chat".to_string(),
|
|
});
|
|
|
|
let done = wait_for_ai_state(&connection, "repo", Duration::from_secs(5), |state| {
|
|
state
|
|
.active_agent
|
|
.as_ref()
|
|
.map(|agent| {
|
|
!agent.pending
|
|
&& agent
|
|
.messages
|
|
.iter()
|
|
.any(|message| message.text.contains("assistant reply"))
|
|
})
|
|
.unwrap_or(false)
|
|
});
|
|
assert!(done
|
|
.active_agent
|
|
.as_ref()
|
|
.unwrap()
|
|
.messages
|
|
.iter()
|
|
.any(|message| message.text.contains("thinking now")));
|
|
|
|
let _ = connection.send(ClientToHub::AiCreateAgent {
|
|
mount: "repo".to_string(),
|
|
title: Some("Second chat".to_string()),
|
|
});
|
|
let two_chats = wait_for_ai_state(&connection, "repo", Duration::from_secs(2), |state| {
|
|
state.agents.len() == 2
|
|
});
|
|
assert_eq!(two_chats.agents.len(), 2);
|
|
}
|
|
|
|
server.join().expect("join ai server");
|
|
|
|
let chats_dir = repo.path().join(".makepad/ai_chats");
|
|
let chat_files = fs::read_dir(&chats_dir)
|
|
.expect("read chats dir")
|
|
.filter_map(Result::ok)
|
|
.filter(|entry| entry.path().extension().and_then(|ext| ext.to_str()) == Some("json"))
|
|
.count();
|
|
assert_eq!(chat_files, 2);
|
|
|
|
let mut connection = StudioHub::start_in_process(config).expect("restart backend");
|
|
let _ = connection.send(ClientToHub::AiGetState {
|
|
mount: "repo".to_string(),
|
|
});
|
|
let restored = wait_for_ai_state(&connection, "repo", Duration::from_secs(2), |state| {
|
|
state.agents.len() == 2
|
|
&& state.agents.iter().any(|agent| agent.title == "persist this chat")
|
|
&& state.agents.iter().any(|agent| agent.title == "Second chat")
|
|
&& state
|
|
.active_agent
|
|
.as_ref()
|
|
.map(|agent| agent.title == "Second chat")
|
|
.unwrap_or(false)
|
|
});
|
|
assert_eq!(restored.agents.len(), 2);
|
|
assert_eq!(restored.active_agent.as_ref().unwrap().title, "Second chat");
|
|
|
|
let first_agent_id = restored
|
|
.agents
|
|
.iter()
|
|
.find(|agent| agent.title == "persist this chat")
|
|
.expect("first restored agent")
|
|
.agent_id;
|
|
let _ = connection.send(ClientToHub::AiSelectAgent {
|
|
mount: "repo".to_string(),
|
|
agent_id: first_agent_id,
|
|
});
|
|
let restored_first_chat =
|
|
wait_for_ai_state(&connection, "repo", Duration::from_secs(2), |state| {
|
|
state
|
|
.active_agent
|
|
.as_ref()
|
|
.map(|agent| {
|
|
agent.agent_id == first_agent_id
|
|
&& agent
|
|
.messages
|
|
.iter()
|
|
.any(|message| message.text.contains("persist this chat"))
|
|
&& agent
|
|
.messages
|
|
.iter()
|
|
.any(|message| message.text.contains("assistant reply"))
|
|
})
|
|
.unwrap_or(false)
|
|
});
|
|
assert_eq!(
|
|
restored_first_chat.active_agent.as_ref().unwrap().title,
|
|
"persist this chat"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn ai_manager_executes_tool_calls_inside_the_hub() {
|
|
let _env_lock = ai_env_lock()
|
|
.lock()
|
|
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
|
|
|
let listener = TcpListener::bind("127.0.0.1:0").expect("bind local ai server");
|
|
let addr = listener.local_addr().expect("local addr");
|
|
let _base_url = EnvGuard::set(
|
|
"MAKEPAD_STUDIO_AI_BASE_URL",
|
|
format!("http://{}/v1/chat/completions", addr),
|
|
);
|
|
|
|
let server = thread::spawn(move || {
|
|
let (mut stream1, _) = listener.accept().expect("accept first ai request");
|
|
let request1 = read_http_request(&mut stream1);
|
|
assert!(request1.contains("\"tools\""));
|
|
assert!(request1.contains("\"read_file\""));
|
|
assert!(request1.contains("\"stream\":true"));
|
|
write_chunked_sse(
|
|
&mut stream1,
|
|
&[
|
|
concat!(
|
|
"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,",
|
|
"\"id\":\"call_1\",\"type\":\"function\",\"function\":{",
|
|
"\"name\":\"read_file\",",
|
|
"\"arguments\":\"{\\\"path\\\":\\\"Cargo.toml\\\",\\\"offset\\\":1,\\\"limit\\\":2}\"",
|
|
"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n"
|
|
),
|
|
"data: [DONE]\n\n",
|
|
],
|
|
);
|
|
|
|
let (mut stream2, _) = listener.accept().expect("accept second ai request");
|
|
let request2 = read_http_request(&mut stream2);
|
|
assert!(request2.contains("\"role\":\"tool\""));
|
|
assert!(request2.contains("\"tool_call_id\":\"call_1\""));
|
|
assert!(request2.contains("Cargo.toml"));
|
|
write_chunked_sse(
|
|
&mut stream2,
|
|
&[
|
|
"data: {\"choices\":[{\"delta\":{\"content\":\"finished after tool call\"}}]}\n\n",
|
|
"data: [DONE]\n\n",
|
|
],
|
|
);
|
|
});
|
|
|
|
let root = std::env::current_dir().expect("current_dir");
|
|
let config = HubConfig {
|
|
mounts: vec![MountConfig {
|
|
name: "repo".to_string(),
|
|
path: root,
|
|
}],
|
|
enable_in_process_gateway: false,
|
|
..Default::default()
|
|
};
|
|
let mut connection = StudioHub::start_in_process(config).expect("start backend");
|
|
|
|
let _ = connection.send(ClientToHub::AiGetState {
|
|
mount: "repo".to_string(),
|
|
});
|
|
let initial = wait_for_ai_state(&connection, "repo", Duration::from_secs(2), |state| {
|
|
state.active_agent.is_some()
|
|
});
|
|
let agent_id = initial.active_agent_id.expect("default ai agent");
|
|
|
|
let _ = connection.send(ClientToHub::AiSendPrompt {
|
|
mount: "repo".to_string(),
|
|
agent_id,
|
|
text: "inspect Cargo.toml".to_string(),
|
|
});
|
|
|
|
let done = wait_for_ai_state(&connection, "repo", Duration::from_secs(5), |state| {
|
|
state
|
|
.active_agent
|
|
.as_ref()
|
|
.map(|agent| {
|
|
!agent.pending
|
|
&& agent
|
|
.messages
|
|
.iter()
|
|
.any(|message| message.text.contains("finished after tool call"))
|
|
})
|
|
.unwrap_or(false)
|
|
});
|
|
|
|
let messages = &done.active_agent.expect("active agent").messages;
|
|
assert!(messages
|
|
.iter()
|
|
.any(|message| message.text.contains("read_file")));
|
|
assert!(messages
|
|
.iter()
|
|
.any(|message| message.text.contains("finished after tool call")));
|
|
|
|
server.join().expect("join ai server");
|
|
}
|
|
|
|
#[test]
|
|
fn ai_manager_open_editor_tool_opens_file_in_primary_ui() {
|
|
let _env_lock = ai_env_lock()
|
|
.lock()
|
|
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
|
|
|
let listener = TcpListener::bind("127.0.0.1:0").expect("bind local ai server");
|
|
let addr = listener.local_addr().expect("local addr");
|
|
let _base_url = EnvGuard::set(
|
|
"MAKEPAD_STUDIO_AI_BASE_URL",
|
|
format!("http://{}/v1/chat/completions", addr),
|
|
);
|
|
|
|
let repo = makepad_studio_hub::test_support::tempdir().expect("tempdir");
|
|
fs::create_dir_all(repo.path().join("src")).expect("create src");
|
|
fs::write(repo.path().join("src/lib.rs"), "pub fn opened_by_ai() {}\n").expect("write file");
|
|
|
|
let server = thread::spawn(move || {
|
|
let (mut stream1, _) = listener.accept().expect("accept first ai request");
|
|
let request1 = read_http_request(&mut stream1);
|
|
assert!(request1.contains("\"open_editor\""));
|
|
write_chunked_sse(
|
|
&mut stream1,
|
|
&[
|
|
concat!(
|
|
"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,",
|
|
"\"id\":\"call_open_editor\",\"type\":\"function\",\"function\":{",
|
|
"\"name\":\"open_editor\",",
|
|
"\"arguments\":\"{\\\"path\\\":\\\"src/lib.rs\\\"}\"",
|
|
"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n"
|
|
),
|
|
"data: [DONE]\n\n",
|
|
],
|
|
);
|
|
|
|
let (mut stream2, _) = listener.accept().expect("accept second ai request");
|
|
let request2 = read_http_request(&mut stream2);
|
|
assert!(request2.contains("\"tool_call_id\":\"call_open_editor\""));
|
|
assert!(request2.contains("Opened repo/src/lib.rs in Studio editor."));
|
|
write_chunked_sse(
|
|
&mut stream2,
|
|
&[
|
|
"data: {\"choices\":[{\"delta\":{\"content\":\"editor opened\"}}]}\n\n",
|
|
"data: [DONE]\n\n",
|
|
],
|
|
);
|
|
});
|
|
|
|
let config = HubConfig {
|
|
mounts: vec![MountConfig {
|
|
name: "repo".to_string(),
|
|
path: repo.path().to_path_buf(),
|
|
}],
|
|
enable_in_process_gateway: false,
|
|
..Default::default()
|
|
};
|
|
let mut connection = StudioHub::start_in_process(config).expect("start backend");
|
|
|
|
let _ = connection.send(ClientToHub::ObserveMount {
|
|
mount: "repo".to_string(),
|
|
primary: Some(true),
|
|
});
|
|
let _ = connection.send(ClientToHub::AiGetState {
|
|
mount: "repo".to_string(),
|
|
});
|
|
let initial = wait_for_ai_state(&connection, "repo", Duration::from_secs(2), |state| {
|
|
state.active_agent.is_some()
|
|
});
|
|
let agent_id = initial.active_agent_id.expect("default ai agent");
|
|
|
|
let _ = connection.send(ClientToHub::AiSendPrompt {
|
|
mount: "repo".to_string(),
|
|
agent_id,
|
|
text: "open src/lib.rs".to_string(),
|
|
});
|
|
|
|
let opened = wait_for_message(&connection, Duration::from_secs(5), |msg| {
|
|
matches!(
|
|
msg,
|
|
HubToClient::TextFileOpened { path, content, .. }
|
|
if path == "repo/src/lib.rs" && content == "pub fn opened_by_ai() {}\n"
|
|
)
|
|
});
|
|
match opened {
|
|
HubToClient::TextFileOpened { path, content, .. } => {
|
|
assert_eq!(path, "repo/src/lib.rs");
|
|
assert_eq!(content, "pub fn opened_by_ai() {}\n");
|
|
}
|
|
_ => unreachable!(),
|
|
}
|
|
|
|
let done = wait_for_ai_state(&connection, "repo", Duration::from_secs(5), |state| {
|
|
state
|
|
.active_agent
|
|
.as_ref()
|
|
.map(|agent| {
|
|
!agent.pending
|
|
&& agent
|
|
.messages
|
|
.iter()
|
|
.any(|message| message.text.contains("editor opened"))
|
|
})
|
|
.unwrap_or(false)
|
|
});
|
|
|
|
assert!(done
|
|
.active_agent
|
|
.as_ref()
|
|
.unwrap()
|
|
.messages
|
|
.iter()
|
|
.any(|message| message.text.contains("open_editor")));
|
|
|
|
server.join().expect("join ai server");
|
|
}
|
|
|
|
#[test]
|
|
fn ai_manager_open_editor_tool_forwards_jump_location_to_primary_ui() {
|
|
let _env_lock = ai_env_lock()
|
|
.lock()
|
|
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
|
|
|
let listener = TcpListener::bind("127.0.0.1:0").expect("bind local ai server");
|
|
let addr = listener.local_addr().expect("local addr");
|
|
let _base_url = EnvGuard::set(
|
|
"MAKEPAD_STUDIO_AI_BASE_URL",
|
|
format!("http://{}/v1/chat/completions", addr),
|
|
);
|
|
|
|
let repo = makepad_studio_hub::test_support::tempdir().expect("tempdir");
|
|
fs::create_dir_all(repo.path().join("src")).expect("create src");
|
|
fs::write(repo.path().join("src/lib.rs"), "pub fn jumped_by_ai() {}\n").expect("write file");
|
|
|
|
let server = thread::spawn(move || {
|
|
let (mut stream1, _) = listener.accept().expect("accept first ai request");
|
|
let request1 = read_http_request(&mut stream1);
|
|
assert!(request1.contains("\"open_editor\""));
|
|
write_chunked_sse(
|
|
&mut stream1,
|
|
&[
|
|
concat!(
|
|
"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,",
|
|
"\"id\":\"call_open_editor\",\"type\":\"function\",\"function\":{",
|
|
"\"name\":\"open_editor\",",
|
|
"\"arguments\":\"{\\\"path\\\":\\\"src/lib.rs\\\",\\\"line\\\":1,\\\"column\\\":8}\"",
|
|
"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n"
|
|
),
|
|
"data: [DONE]\n\n",
|
|
],
|
|
);
|
|
|
|
let (mut stream2, _) = listener.accept().expect("accept second ai request");
|
|
let request2 = read_http_request(&mut stream2);
|
|
assert!(request2.contains("\"tool_call_id\":\"call_open_editor\""));
|
|
assert!(request2.contains("Opened repo/src/lib.rs at 1:8 in Studio editor."));
|
|
write_chunked_sse(
|
|
&mut stream2,
|
|
&[
|
|
"data: {\"choices\":[{\"delta\":{\"content\":\"editor jumped\"}}]}\n\n",
|
|
"data: [DONE]\n\n",
|
|
],
|
|
);
|
|
});
|
|
|
|
let config = HubConfig {
|
|
mounts: vec![MountConfig {
|
|
name: "repo".to_string(),
|
|
path: repo.path().to_path_buf(),
|
|
}],
|
|
enable_in_process_gateway: false,
|
|
..Default::default()
|
|
};
|
|
let mut connection = StudioHub::start_in_process(config).expect("start backend");
|
|
|
|
let _ = connection.send(ClientToHub::ObserveMount {
|
|
mount: "repo".to_string(),
|
|
primary: Some(true),
|
|
});
|
|
let initial = wait_for_ai_state(&connection, "repo", Duration::from_secs(2), |state| {
|
|
state.active_agent.is_some()
|
|
});
|
|
let agent_id = initial.active_agent_id.expect("default ai agent");
|
|
|
|
let _ = connection.send(ClientToHub::AiSendPrompt {
|
|
mount: "repo".to_string(),
|
|
agent_id,
|
|
text: "open src/lib.rs at line 1 column 8".to_string(),
|
|
});
|
|
|
|
let opened = wait_for_message(&connection, Duration::from_secs(5), |msg| {
|
|
matches!(
|
|
msg,
|
|
HubToClient::TextFileOpened {
|
|
path,
|
|
content,
|
|
line,
|
|
column,
|
|
..
|
|
} if path == "repo/src/lib.rs"
|
|
&& content == "pub fn jumped_by_ai() {}\n"
|
|
&& *line == Some(1)
|
|
&& *column == Some(8)
|
|
)
|
|
});
|
|
match opened {
|
|
HubToClient::TextFileOpened {
|
|
path,
|
|
content,
|
|
line,
|
|
column,
|
|
..
|
|
} => {
|
|
assert_eq!(path, "repo/src/lib.rs");
|
|
assert_eq!(content, "pub fn jumped_by_ai() {}\n");
|
|
assert_eq!(line, Some(1));
|
|
assert_eq!(column, Some(8));
|
|
}
|
|
_ => unreachable!(),
|
|
}
|
|
|
|
let done = wait_for_ai_state(&connection, "repo", Duration::from_secs(5), |state| {
|
|
state
|
|
.active_agent
|
|
.as_ref()
|
|
.map(|agent| {
|
|
!agent.pending
|
|
&& agent
|
|
.messages
|
|
.iter()
|
|
.any(|message| message.text.contains("editor jumped"))
|
|
})
|
|
.unwrap_or(false)
|
|
});
|
|
assert!(done
|
|
.active_agent
|
|
.as_ref()
|
|
.unwrap()
|
|
.messages
|
|
.iter()
|
|
.any(|message| message.text.contains("open_editor")));
|
|
|
|
server.join().expect("join ai server");
|
|
}
|
|
|
|
#[test]
|
|
fn ai_manager_observe_filesystem_tool_reports_recent_changes() {
|
|
let _env_lock = ai_env_lock()
|
|
.lock()
|
|
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
|
|
|
let listener = TcpListener::bind("127.0.0.1:0").expect("bind local ai server");
|
|
let addr = listener.local_addr().expect("local addr");
|
|
let _base_url = EnvGuard::set(
|
|
"MAKEPAD_STUDIO_AI_BASE_URL",
|
|
format!("http://{}/v1/chat/completions", addr),
|
|
);
|
|
|
|
let repo = makepad_studio_hub::test_support::tempdir().expect("tempdir");
|
|
fs::create_dir_all(repo.path().join("src")).expect("create src");
|
|
fs::write(repo.path().join("src/lib.rs"), "pub fn before() {}\n").expect("write file");
|
|
|
|
let server = thread::spawn(move || {
|
|
let (mut stream1, _) = listener.accept().expect("accept first ai request");
|
|
let request1 = read_http_request(&mut stream1);
|
|
assert!(request1.contains("\"observe_filesystem\""));
|
|
write_chunked_sse(
|
|
&mut stream1,
|
|
&[
|
|
concat!(
|
|
"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,",
|
|
"\"id\":\"call_observe_fs\",\"type\":\"function\",\"function\":{",
|
|
"\"name\":\"observe_filesystem\",",
|
|
"\"arguments\":\"{\\\"path\\\":\\\"src\\\",\\\"since_secs\\\":300}\"",
|
|
"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n"
|
|
),
|
|
"data: [DONE]\n\n",
|
|
],
|
|
);
|
|
|
|
let (mut stream2, _) = listener.accept().expect("accept second ai request");
|
|
let request2 = read_http_request(&mut stream2);
|
|
assert!(request2.contains("\"tool_call_id\":\"call_observe_fs\""));
|
|
assert!(request2.contains("src/lib.rs"));
|
|
write_chunked_sse(
|
|
&mut stream2,
|
|
&[
|
|
"data: {\"choices\":[{\"delta\":{\"content\":\"observed changes\"}}]}\n\n",
|
|
"data: [DONE]\n\n",
|
|
],
|
|
);
|
|
});
|
|
|
|
let config = HubConfig {
|
|
mounts: vec![MountConfig {
|
|
name: "repo".to_string(),
|
|
path: repo.path().to_path_buf(),
|
|
}],
|
|
enable_in_process_gateway: false,
|
|
..Default::default()
|
|
};
|
|
let mut connection = StudioHub::start_in_process(config).expect("start backend");
|
|
|
|
let _ = connection.send(ClientToHub::AiGetState {
|
|
mount: "repo".to_string(),
|
|
});
|
|
let initial = wait_for_ai_state(&connection, "repo", Duration::from_secs(2), |state| {
|
|
state.active_agent.is_some()
|
|
});
|
|
let agent_id = initial.active_agent_id.expect("default ai agent");
|
|
|
|
fs::write(repo.path().join("src/lib.rs"), "pub fn after() {}\n").expect("update file");
|
|
let _ = wait_for_message(&connection, Duration::from_secs(6), |msg| {
|
|
matches!(
|
|
msg,
|
|
HubToClient::FileChanged { path }
|
|
if path == "repo/src/lib.rs" || path == "repo"
|
|
)
|
|
});
|
|
|
|
let _ = connection.send(ClientToHub::AiSendPrompt {
|
|
mount: "repo".to_string(),
|
|
agent_id,
|
|
text: "what changed under src recently?".to_string(),
|
|
});
|
|
|
|
let done = wait_for_ai_state(&connection, "repo", Duration::from_secs(5), |state| {
|
|
state
|
|
.active_agent
|
|
.as_ref()
|
|
.map(|agent| {
|
|
!agent.pending
|
|
&& agent
|
|
.messages
|
|
.iter()
|
|
.any(|message| message.text.contains("observed changes"))
|
|
})
|
|
.unwrap_or(false)
|
|
});
|
|
assert!(done
|
|
.active_agent
|
|
.as_ref()
|
|
.unwrap()
|
|
.messages
|
|
.iter()
|
|
.any(|message| message.text.contains("observe_filesystem")));
|
|
|
|
server.join().expect("join ai server");
|
|
}
|
|
|
|
#[test]
|
|
fn ai_manager_streams_thinking_before_completion() {
|
|
let _env_lock = ai_env_lock()
|
|
.lock()
|
|
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
|
|
|
let listener = TcpListener::bind("127.0.0.1:0").expect("bind local ai server");
|
|
let addr = listener.local_addr().expect("local addr");
|
|
let _base_url = EnvGuard::set(
|
|
"MAKEPAD_STUDIO_AI_BASE_URL",
|
|
format!("http://{}/v1/chat/completions", addr),
|
|
);
|
|
|
|
let server = thread::spawn(move || {
|
|
let (mut stream, _) = listener.accept().expect("accept ai request");
|
|
let request_text = read_http_request(&mut stream);
|
|
assert!(request_text.contains("\"stream\":true"));
|
|
write_chunked_sse(
|
|
&mut stream,
|
|
&[
|
|
"data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"thinking now\"}}]}\n\n",
|
|
"data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\n",
|
|
"data: [DONE]\n\n",
|
|
],
|
|
);
|
|
});
|
|
|
|
let root = std::env::current_dir().expect("current_dir");
|
|
let config = HubConfig {
|
|
mounts: vec![MountConfig {
|
|
name: "repo".to_string(),
|
|
path: root,
|
|
}],
|
|
enable_in_process_gateway: false,
|
|
..Default::default()
|
|
};
|
|
let mut connection = StudioHub::start_in_process(config).expect("start backend");
|
|
|
|
let _ = connection.send(ClientToHub::AiGetState {
|
|
mount: "repo".to_string(),
|
|
});
|
|
let initial = wait_for_ai_state(&connection, "repo", Duration::from_secs(2), |state| {
|
|
state.active_agent.is_some()
|
|
});
|
|
let agent_id = initial.active_agent_id.expect("default ai agent");
|
|
|
|
let _ = connection.send(ClientToHub::AiSendPrompt {
|
|
mount: "repo".to_string(),
|
|
agent_id,
|
|
text: "stream it".to_string(),
|
|
});
|
|
|
|
let thinking = wait_for_ai_state(&connection, "repo", Duration::from_secs(5), |state| {
|
|
state
|
|
.active_agent
|
|
.as_ref()
|
|
.map(|agent| {
|
|
agent.pending
|
|
&& agent.messages.iter().any(|message| {
|
|
matches!(message.role, AiMessageRole::Thinking)
|
|
&& message.text.contains("thinking now")
|
|
})
|
|
})
|
|
.unwrap_or(false)
|
|
});
|
|
assert!(thinking
|
|
.active_agent
|
|
.as_ref()
|
|
.unwrap()
|
|
.messages
|
|
.iter()
|
|
.any(|message| matches!(message.role, AiMessageRole::Thinking)));
|
|
|
|
let done = wait_for_ai_state(&connection, "repo", Duration::from_secs(5), |state| {
|
|
state
|
|
.active_agent
|
|
.as_ref()
|
|
.map(|agent| {
|
|
!agent.pending
|
|
&& agent.messages.iter().any(|message| {
|
|
matches!(message.role, AiMessageRole::Assistant)
|
|
&& message.text.contains("hello")
|
|
})
|
|
})
|
|
.unwrap_or(false)
|
|
});
|
|
assert!(done
|
|
.active_agent
|
|
.as_ref()
|
|
.unwrap()
|
|
.messages
|
|
.iter()
|
|
.any(|message| matches!(message.role, AiMessageRole::Assistant)));
|
|
|
|
server.join().expect("join ai server");
|
|
}
|
|
|
|
#[test]
|
|
fn ai_manager_preserves_streamed_thinking_whitespace() {
|
|
let _env_lock = ai_env_lock()
|
|
.lock()
|
|
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
|
|
|
let listener = TcpListener::bind("127.0.0.1:0").expect("bind local ai server");
|
|
let addr = listener.local_addr().expect("local addr");
|
|
let _base_url = EnvGuard::set(
|
|
"MAKEPAD_STUDIO_AI_BASE_URL",
|
|
format!("http://{}/v1/chat/completions", addr),
|
|
);
|
|
|
|
let server = thread::spawn(move || {
|
|
let (mut stream, _) = listener.accept().expect("accept ai request");
|
|
let request_text = read_http_request(&mut stream);
|
|
assert!(request_text.contains("\"stream\":true"));
|
|
write_chunked_sse(
|
|
&mut stream,
|
|
&[
|
|
"data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"The\"}}]}\n\n",
|
|
"data: {\"choices\":[{\"delta\":{\"reasoning_content\":\" user\"}}]}\n\n",
|
|
"data: {\"choices\":[{\"delta\":{\"reasoning_content\":\" says hi\"}}]}\n\n",
|
|
"data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n",
|
|
"data: [DONE]\n\n",
|
|
],
|
|
);
|
|
});
|
|
|
|
let root = std::env::current_dir().expect("current_dir");
|
|
let config = HubConfig {
|
|
mounts: vec![MountConfig {
|
|
name: "repo".to_string(),
|
|
path: root,
|
|
}],
|
|
enable_in_process_gateway: false,
|
|
..Default::default()
|
|
};
|
|
let mut connection = StudioHub::start_in_process(config).expect("start backend");
|
|
|
|
let _ = connection.send(ClientToHub::AiGetState {
|
|
mount: "repo".to_string(),
|
|
});
|
|
let initial = wait_for_ai_state(&connection, "repo", Duration::from_secs(2), |state| {
|
|
state.active_agent.is_some()
|
|
});
|
|
let agent_id = initial.active_agent_id.expect("default ai agent");
|
|
|
|
let _ = connection.send(ClientToHub::AiSendPrompt {
|
|
mount: "repo".to_string(),
|
|
agent_id,
|
|
text: "say hi".to_string(),
|
|
});
|
|
|
|
let done = wait_for_ai_state(&connection, "repo", Duration::from_secs(5), |state| {
|
|
state
|
|
.active_agent
|
|
.as_ref()
|
|
.map(|agent| {
|
|
!agent.pending
|
|
&& agent.messages.iter().any(|message| {
|
|
matches!(message.role, AiMessageRole::Thinking)
|
|
&& message.text.contains("The user says hi")
|
|
})
|
|
})
|
|
.unwrap_or(false)
|
|
});
|
|
|
|
assert!(done
|
|
.active_agent
|
|
.as_ref()
|
|
.unwrap()
|
|
.messages
|
|
.iter()
|
|
.any(|message| {
|
|
matches!(message.role, AiMessageRole::Thinking)
|
|
&& message.text.contains("The user says hi")
|
|
}));
|
|
|
|
server.join().expect("join ai server");
|
|
}
|
|
|
|
#[test]
|
|
fn ai_manager_accepts_second_prompt_after_done_before_socket_close() {
|
|
let _env_lock = ai_env_lock()
|
|
.lock()
|
|
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
|
|
|
let listener = TcpListener::bind("127.0.0.1:0").expect("bind local ai server");
|
|
let addr = listener.local_addr().expect("local addr");
|
|
let _base_url = EnvGuard::set(
|
|
"MAKEPAD_STUDIO_AI_BASE_URL",
|
|
format!("http://{}/v1/chat/completions", addr),
|
|
);
|
|
|
|
let server = thread::spawn(move || {
|
|
let (stream1, _) = listener.accept().expect("accept first ai request");
|
|
let mut stream1_reader = stream1.try_clone().expect("clone first stream");
|
|
let request1 = read_http_request(&mut stream1_reader);
|
|
assert!(request1.contains("\"stream\":true"));
|
|
let first_handler = thread::spawn(move || {
|
|
write_chunked_sse_and_hold_open(
|
|
stream1,
|
|
&[
|
|
"data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"hi reasoning\"}}]}\n\n",
|
|
"data: {\"choices\":[{\"delta\":{\"content\":\"Hi!\"}}]}\n\n",
|
|
"data: [DONE]\n\n",
|
|
],
|
|
Duration::from_secs(2),
|
|
);
|
|
});
|
|
|
|
let (mut stream2, _) = listener.accept().expect("accept second ai request");
|
|
let request2 = read_http_request(&mut stream2);
|
|
assert!(request2.contains("write a poem"));
|
|
write_chunked_sse(
|
|
&mut stream2,
|
|
&[
|
|
"data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"poem reasoning\"}}]}\n\n",
|
|
"data: {\"choices\":[{\"delta\":{\"content\":\"Roses are red\"}}]}\n\n",
|
|
"data: [DONE]\n\n",
|
|
],
|
|
);
|
|
|
|
first_handler.join().expect("join first handler");
|
|
});
|
|
|
|
let root = std::env::current_dir().expect("current_dir");
|
|
let config = HubConfig {
|
|
mounts: vec![MountConfig {
|
|
name: "repo".to_string(),
|
|
path: root,
|
|
}],
|
|
enable_in_process_gateway: false,
|
|
..Default::default()
|
|
};
|
|
let mut connection = StudioHub::start_in_process(config).expect("start backend");
|
|
|
|
let _ = connection.send(ClientToHub::AiGetState {
|
|
mount: "repo".to_string(),
|
|
});
|
|
let initial = wait_for_ai_state(&connection, "repo", Duration::from_secs(2), |state| {
|
|
state.active_agent.is_some()
|
|
});
|
|
let agent_id = initial.active_agent_id.expect("default ai agent");
|
|
|
|
let _ = connection.send(ClientToHub::AiSendPrompt {
|
|
mount: "repo".to_string(),
|
|
agent_id,
|
|
text: "say hi".to_string(),
|
|
});
|
|
|
|
let first_done = wait_for_ai_state(&connection, "repo", Duration::from_secs(5), |state| {
|
|
state
|
|
.active_agent
|
|
.as_ref()
|
|
.map(|agent| {
|
|
!agent.pending
|
|
&& agent.messages.iter().any(|message| {
|
|
matches!(message.role, AiMessageRole::Assistant)
|
|
&& message.text.contains("Hi!")
|
|
})
|
|
})
|
|
.unwrap_or(false)
|
|
});
|
|
assert!(!first_done.active_agent.as_ref().unwrap().pending);
|
|
|
|
let _ = connection.send(ClientToHub::AiSendPrompt {
|
|
mount: "repo".to_string(),
|
|
agent_id,
|
|
text: "write a poem".to_string(),
|
|
});
|
|
|
|
let second_done = wait_for_ai_state(&connection, "repo", Duration::from_secs(5), |state| {
|
|
state
|
|
.active_agent
|
|
.as_ref()
|
|
.map(|agent| {
|
|
!agent.pending
|
|
&& agent.messages.iter().any(|message| {
|
|
matches!(message.role, AiMessageRole::Assistant)
|
|
&& message.text.contains("Roses are red")
|
|
})
|
|
})
|
|
.unwrap_or(false)
|
|
});
|
|
|
|
assert!(second_done
|
|
.active_agent
|
|
.as_ref()
|
|
.unwrap()
|
|
.messages
|
|
.iter()
|
|
.any(|message| {
|
|
matches!(message.role, AiMessageRole::Assistant)
|
|
&& message.text.contains("Roses are red")
|
|
}));
|
|
|
|
server.join().expect("join ai server");
|
|
}
|