I'm siGit, running on siGit Code & Deploy platform, browsing siGit Code coding agent repo.

1 //! The official MCP server follows the signed-in account.
2 //!
3 //! An MCP connection fixes its bearer when it connects, so `web_search` and the
4 //! other `mcp__sigit__*` tools used to keep the account that was signed in at
5 //! startup. The server is now reconnected when the stored token changes, and
6 //! the check runs before the call so the old bearer is never sent.
7
8 use std::io::{BufRead, BufReader, Read, Write};
9 use std::net::{TcpListener, TcpStream};
10 use std::process::{Child, ChildStdin, Command, Stdio};
11 use std::sync::mpsc::{Receiver, Sender, channel};
12 use std::time::{Duration, Instant};
13
14 use serde_json::{Value, json};
15
16 const TIMEOUT: Duration = Duration::from_secs(60);
17
18 fn respond(stream: &mut TcpStream, status: &str, content_type: &str, body: &str) {
19 let response = format!(
20 "HTTP/1.1 {status}\r\ncontent-type: {content_type}\r\n\
21 content-length: {}\r\n\r\n{body}",
22 body.len()
23 );
24 let _ = stream.write_all(response.as_bytes());
25 let _ = stream.flush();
26 }
27
28 fn serve(stream: TcpStream, calls: Sender<String>) {
29 let Ok(clone) = stream.try_clone() else {
30 return;
31 };
32 let mut stream = stream;
33 let mut reader = BufReader::new(clone);
34 loop {
35 let mut request_line = String::new();
36 if reader.read_line(&mut request_line).unwrap_or(0) == 0 {
37 return;
38 }
39 let path = request_line
40 .split_whitespace()
41 .nth(1)
42 .unwrap_or("")
43 .to_string();
44 let method = request_line
45 .split_whitespace()
46 .next()
47 .unwrap_or("")
48 .to_string();
49 let mut content_length = 0usize;
50 let mut authorization = String::new();
51 loop {
52 let mut line = String::new();
53 if reader.read_line(&mut line).unwrap_or(0) == 0 {
54 return;
55 }
56 let line = line.trim().to_string();
57 if line.is_empty() {
58 break;
59 }
60 let lower = line.to_ascii_lowercase();
61 if let Some(length) = lower.strip_prefix("content-length:") {
62 content_length = length.trim().parse().unwrap_or(0);
63 }
64 if lower.starts_with("authorization:") {
65 authorization = lower;
66 }
67 }
68 let mut body = vec![0u8; content_length];
69 if reader.read_exact(&mut body).is_err() {
70 return;
71 }
72 let body: Value = serde_json::from_slice(&body).unwrap_or(Value::Null);
73
74 if path.ends_with("/chat/completions") {
75 // Search once per prompt: answer in text only after a tool result.
76 let answered = body["messages"]
77 .as_array()
78 .and_then(|m| m.last())
79 .is_some_and(|m| m["role"] == "tool");
80 let delta = if answered {
81 json!({"content": "done"})
82 } else {
83 json!({"tool_calls": [{
84 "index": 0,
85 "id": "call_1",
86 "type": "function",
87 "function": {
88 "name": "mcp__sigit__web_search",
89 "arguments": "{\"query\":\"rust\"}"
90 }
91 }]})
92 };
93 let sse = format!(
94 "data: {}\n\ndata: [DONE]\n\n",
95 json!({"choices": [{"delta": delta}]})
96 );
97 respond(&mut stream, "200 OK", "text/event-stream", &sse);
98 } else if path.ends_with("/mcp") && method == "POST" {
99 match body["method"].as_str() {
100 Some("initialize") => {
101 let result = json!({
102 "jsonrpc": "2.0",
103 "id": body["id"],
104 "result": {
105 "protocolVersion": body["params"]["protocolVersion"],
106 "capabilities": {"tools": {}},
107 "serverInfo": {"name": "fake", "version": "1"}
108 }
109 });
110 respond(
111 &mut stream,
112 "200 OK",
113 "application/json",
114 &result.to_string(),
115 );
116 }
117 Some("tools/list") => {
118 let result = json!({
119 "jsonrpc": "2.0",
120 "id": body["id"],
121 "result": {"tools": [{
122 "name": "web_search",
123 "description": "Search the web",
124 "inputSchema": {"type": "object"}
125 }]}
126 });
127 respond(
128 &mut stream,
129 "200 OK",
130 "application/json",
131 &result.to_string(),
132 );
133 }
134 Some("tools/call") => {
135 let _ = calls.send(authorization);
136 let result = json!({
137 "jsonrpc": "2.0",
138 "id": body["id"],
139 "result": {"content": [{"type": "text", "text": "found"}]}
140 });
141 respond(
142 &mut stream,
143 "200 OK",
144 "application/json",
145 &result.to_string(),
146 );
147 }
148 // Notifications carry no id and get an empty 202.
149 _ => respond(&mut stream, "202 Accepted", "text/plain", ""),
150 }
151 } else {
152 respond(&mut stream, "405 Method Not Allowed", "text/plain", "");
153 }
154 }
155 }
156
157 /// Stand in for the cloud: `/mcp` and `/chat/completions`. The `authorization`
158 /// header of every MCP `tools/call` is forwarded to the returned receiver.
159 fn start_fake_cloud() -> (u16, Receiver<String>) {
160 let listener = TcpListener::bind("127.0.0.1:0").expect("bind fake cloud");
161 let port = listener.local_addr().unwrap().port();
162 let (call_tx, calls) = channel();
163 std::thread::spawn(move || {
164 for stream in listener.incoming() {
165 let Ok(stream) = stream else { continue };
166 let calls = call_tx.clone();
167 std::thread::spawn(move || serve(stream, calls));
168 }
169 });
170 (port, calls)
171 }
172
173 struct AgentUnderTest {
174 child: Child,
175 stdin: ChildStdin,
176 incoming: Receiver<Value>,
177 next_id: u64,
178 }
179
180 fn spawn_agent(cloud_port: u16, config_dir: &std::path::Path) -> AgentUnderTest {
181 let mut child = Command::new(env!("CARGO_BIN_EXE_sigit"))
182 // No provider override and no on-device model: the thread starts on a
183 // cloud tier because a session is stored.
184 .env_remove("OPENAI_BASE_URL")
185 .env_remove("OPENAI_API_KEY")
186 .env_remove("SIGIT_MODEL")
187 .env("SIGIT_LOCAL_INFERENCE", "off")
188 .env("SIGIT_CLOUD_URL", format!("http://127.0.0.1:{cloud_port}"))
189 .env("SIGIT_CONFIG_DIR", config_dir)
190 .stdin(Stdio::piped())
191 .stdout(Stdio::piped())
192 .stderr(Stdio::null())
193 .spawn()
194 .expect("spawn sigit in ACP mode");
195
196 let stdout = child.stdout.take().unwrap();
197 let (message_tx, incoming) = channel();
198 std::thread::spawn(move || {
199 for line in BufReader::new(stdout).lines() {
200 let Ok(line) = line else { break };
201 if let Ok(message) = serde_json::from_str::<Value>(&line)
202 && message_tx.send(message).is_err()
203 {
204 break;
205 }
206 }
207 });
208
209 let stdin = child.stdin.take().unwrap();
210 AgentUnderTest {
211 child,
212 stdin,
213 incoming,
214 next_id: 0,
215 }
216 }
217
218 impl AgentUnderTest {
219 fn request(&mut self, method: &str, params: Value) -> u64 {
220 self.next_id += 1;
221 let id = self.next_id;
222 let mut line =
223 json!({"jsonrpc": "2.0", "id": id, "method": method, "params": params}).to_string();
224 line.push('\n');
225 self.stdin.write_all(line.as_bytes()).expect("write stdin");
226 self.stdin.flush().expect("flush stdin");
227 id
228 }
229
230 fn wait_for_response(&mut self, id: u64) -> Value {
231 let deadline = Instant::now() + TIMEOUT;
232 loop {
233 let remaining = deadline.saturating_duration_since(Instant::now());
234 let Ok(message) = self.incoming.recv_timeout(remaining) else {
235 panic!("timed out waiting for the response to request {id}");
236 };
237 if message["id"] == id && message.get("method").is_none() {
238 assert!(
239 message.get("error").is_none(),
240 "request {id} failed: {message}"
241 );
242 return message;
243 }
244 }
245 }
246
247 fn prompt(&mut self, session_id: &str, text: &str) {
248 let id = self.request(
249 "session/prompt",
250 json!({"sessionId": session_id, "prompt": [{"type": "text", "text": text}]}),
251 );
252 self.wait_for_response(id);
253 }
254 }
255
256 impl Drop for AgentUnderTest {
257 fn drop(&mut self) {
258 let _ = self.child.kill();
259 let _ = self.child.wait();
260 }
261 }
262
263 fn store_token(config_dir: &std::path::Path, token: &str) {
264 std::fs::write(
265 config_dir.join("credentials.toml"),
266 format!("access_token = \"{token}\"\nemail = \"dev@example.com\"\n"),
267 )
268 .unwrap();
269 }
270
271 #[test]
272 fn the_official_mcp_server_uses_the_current_account() {
273 let scratch = std::env::temp_dir().join(format!("sigit_acp_mcp_reauth_{}", std::process::id()));
274 let config_dir = scratch.join("config");
275 let work_dir = scratch.join("work");
276 std::fs::create_dir_all(&config_dir).unwrap();
277 std::fs::create_dir_all(&work_dir).unwrap();
278 store_token(&config_dir, "token-a");
279
280 let (port, calls) = start_fake_cloud();
281 let mut agent = spawn_agent(port, &config_dir);
282
283 let id = agent.request(
284 "initialize",
285 json!({"protocolVersion": 1, "clientCapabilities": {}}),
286 );
287 agent.wait_for_response(id);
288 let id = agent.request("session/new", json!({"cwd": work_dir, "mcpServers": []}));
289 let session_id = agent.wait_for_response(id)["result"]["sessionId"]
290 .as_str()
291 .expect("session id")
292 .to_string();
293
294 agent.prompt(&session_id, "first");
295 let first = calls.recv_timeout(TIMEOUT).expect("first tools/call");
296 assert_eq!(first, "authorization: bearer token-a");
297
298 // What signing in as another account does to the stored session.
299 store_token(&config_dir, "token-b");
300 agent.prompt(&session_id, "second");
301 let second = calls.recv_timeout(TIMEOUT).expect("second tools/call");
302 assert_eq!(second, "authorization: bearer token-b");
303
304 // And signing out.
305 std::fs::remove_file(config_dir.join("credentials.toml")).unwrap();
306 agent.prompt(&session_id, "third");
307 let third = calls.recv_timeout(TIMEOUT).expect("third tools/call");
308 assert_eq!(third, "");
309
310 std::fs::remove_dir_all(&scratch).ok();
311 }