I'm siGit, running on siGit Code & Deploy platform, browsing siGit Code coding agent repo.
Reconnect the official MCP server when the signed-in account changes
The official sigit MCP server took its bearer from the stored token once, at startup, so web_search and the other mcp__sigit__ tools kept using the previous account after a login, logout or account switch until a restart. The server now follows the session. Its bearer is read from the stored credentials on every connect, and each connection remembers the token it used. call_server compares that with the stored token before sending and reconnects on a mismatch, so the old bearer is never sent. Reconnecting also refreshes the tool list and the error shown by /mcp. refresh_session_auth() does the same check without a call, and runs after sign-in, sign-out, /login, /logout and /reload so the tool list updates right away. A user-defined `sigit` entry in mcp.toml keeps its static headers and does not follow the session. Fixes #260 Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Seto Elkahfi committed
Oct 29, 2026 at 23:10 UTC
b92581e4b941374b66d93121b037e2207e961911
4 files changed
+575
-65
crates/cli/src/chat.rs
+2
@@ -3162,10 +3162,12 @@ mod tui {
3162
crate::browser_login_message().await
3163
}
3164
};
3165
+ crate::mcp::refresh_session_auth().await;
3166
app.finish_login(result);
3167
}
3168
SlashCommand::Logout => {
3169
let message = crate::account::end_session().await;
3170
+ crate::mcp::refresh_session_auth().await;
3171
app.messages.push(ChatMessage::system(message));
3172
}
3173
SlashCommand::Whoami => {
crates/cli/src/main.rs
+5
@@ -2391,6 +2391,7 @@ impl SiGitAgent {
2391
match outcome {
2392
Ok(email) => {
2393
log::info!("authenticate: signed in as {email}");
2394
+ mcp::refresh_session_auth().await;
2395
Ok(AuthenticateResponse::default())
2396
}
2397
Err(reason) => Err(agent_client_protocol::Error::new(
@@ -2420,6 +2421,7 @@ impl SiGitAgent {
2421
/// message to show. Shared by ACP `logout` and the `/logout` command.
2422
async fn sign_out(&self, cx: &ConnectionTo<Client>) -> String {
2423
let message = account::end_session().await;
2424
+ mcp::refresh_session_auth().await;
2425
let local = default_local_model_config();
2426
for key in self.leave_cloud(&local).await {
2427
let config_options = build_config_options(&local, &key);
@@ -4487,6 +4489,8 @@ impl SiGitAgent {
4489
/// picked up), retries a failed model load, and pushes refreshed commands +
4490
/// picker so the editor's UI reflects the current state.
4491
async fn handle_reload(&self, cx: &ConnectionTo<Client>, session_id: SessionId) {
4492
+ // A `sigit login` from another terminal may have changed the account.
4493
+ mcp::refresh_session_auth().await;
4494
let signed_in = account::status_line().await;
4495
4496
let on_cloud_tier = {
@@ -5986,6 +5990,7 @@ async fn exec_slash_acp(
5990
.await
5991
.unwrap_or_else(|message| message),
5992
};
5993
+ mcp::refresh_session_auth().await;
5994
agent.send_assistant_message(cx, session_id, message).ok();
5995
}
5996
SlashCommand::Logout => {
crates/cli/src/mcp.rs
+257
-65
@@ -70,7 +70,7 @@
70
use std::collections::BTreeMap;
71
use std::path::PathBuf;
72
use std::sync::Mutex as StdMutex;
73
-use std::sync::{Arc, OnceLock};
73
+use std::sync::{Arc, OnceLock, RwLock as StdRwLock};
74
use std::time::Duration;
75
76
use ed_mcp::{ClientInfo, Connection, ServerSpec};
@@ -179,11 +179,29 @@ struct ServerConn {
179
/// How it was reached, kept so an HTTP server can be reconnected.
180
def: ServerDef,
181
/// The live connection. `None` when the server failed to connect.
182
- conn: RwLock<Option<Arc<Connection>>>,
183
- /// Tools discovered at startup. Empty when the server failed to connect.
184
- tools: Vec<McpTool>,
185
- /// Connection error, if the handshake failed. Surfaced by `/mcp`.
186
- error: Option<String>,
182
+ conn: RwLock<Option<Live>>,
183
+ /// Tools discovered at the last (re)connect. Empty when the server failed
184
+ /// to connect. Replaced when a reconnect finds a different set.
185
+ tools: StdRwLock<Vec<McpTool>>,
186
+ /// Connection error, if the last handshake failed. Surfaced by `/mcp`.
187
+ error: StdRwLock<Option<String>>,
188
+}
189
+
190
+/// A connection and the session token it was opened with, so a changed sign-in
191
+/// can be told apart from a connection that is still current.
192
+struct Live {
193
+ conn: Arc<Connection>,
194
+ token: Option<String>,
195
+}
196
+
197
+impl ServerConn {
198
+ fn tools(&self) -> std::sync::RwLockReadGuard<'_, Vec<McpTool>> {
199
+ self.tools.read().unwrap_or_else(|e| e.into_inner())
200
+ }
201
+
202
+ fn error(&self) -> Option<String> {
203
+ self.error.read().unwrap_or_else(|e| e.into_inner()).clone()
204
+ }
205
}
206
207
/// The process-global MCP state: every configured server.
@@ -325,25 +343,43 @@ impl TransportDef {
343
struct ServerDef {
344
name: String,
345
transport: TransportDef,
346
+ /// Whether the server authenticates as the signed-in siGit account. Only
347
+ /// the built-in official server does. Its bearer is read from the stored
348
+ /// credentials on every (re)connect instead of being fixed at startup.
349
+ follows_session: bool,
350
}
351
352
impl ServerDef {
353
/// The `ed-mcp` spec for this server. An `Authorization: Bearer` header is
354
/// passed as the bearer token, which is how `rmcp` wants it.
333
- fn spec(&self) -> ServerSpec {
355
+ /// The stored session token, for a server that follows the session.
356
+ fn session_token(&self) -> Option<String> {
357
+ if self.follows_session {
358
+ crate::credentials::load_token()
359
+ } else {
360
+ None
361
+ }
362
+ }
363
+
364
+ fn spec_with(&self, session_token: Option<&str>) -> ServerSpec {
365
match &self.transport {
366
TransportDef::Http { url, headers } => {
336
- headers
337
- .iter()
338
- .fold(
339
- ServerSpec::http(&self.name, url),
340
- |spec, (key, value)| match value.strip_prefix("Bearer ") {
341
- Some(token) if key.eq_ignore_ascii_case("authorization") => {
342
- spec.with_bearer(token)
343
- }
344
- _ => spec.with_header(key, value),
345
- },
346
- )
367
+ let spec =
368
+ headers
369
+ .iter()
370
+ .fold(
371
+ ServerSpec::http(&self.name, url),
372
+ |spec, (key, value)| match value.strip_prefix("Bearer ") {
373
+ Some(token) if key.eq_ignore_ascii_case("authorization") => {
374
+ spec.with_bearer(token)
375
+ }
376
+ _ => spec.with_header(key, value),
377
+ },
378
+ );
379
+ match session_token {
380
+ Some(token) => spec.with_bearer(token),
381
+ None => spec,
382
+ }
383
}
384
TransportDef::Stdio { command, args, env } => {
385
ServerSpec::stdio(&self.name, command, args.clone(), env.clone())
@@ -446,7 +482,14 @@ fn load_configs() -> Vec<ServerDef> {
482
continue;
483
}
484
};
449
- upsert(&mut defs, ServerDef { name, transport });
485
+ upsert(
486
+ &mut defs,
487
+ ServerDef {
488
+ name,
489
+ transport,
490
+ follows_session: false,
491
+ },
492
+ );
493
}
494
}
495
@@ -463,16 +506,13 @@ fn load_configs() -> Vec<ServerDef> {
506
// Add the baked-in official server, but never clobber a user-defined entry
507
// named `sigit` — an explicit config (e.g. a custom URL or headers) wins.
508
if include_official && !defs.iter().any(|d| d.name == OFFICIAL_SERVER_NAME) {
466
- let mut headers = Vec::new();
467
- if let Some(token) = crate::credentials::load_token() {
468
- headers.push(("Authorization".to_string(), format!("Bearer {token}")));
469
- }
509
defs.push(ServerDef {
510
name: OFFICIAL_SERVER_NAME.to_string(),
511
transport: TransportDef::Http {
512
url: official_url(),
474
- headers,
513
+ headers: Vec::new(),
514
},
515
+ follows_session: true,
516
});
517
}
518
@@ -499,6 +539,7 @@ fn load_configs() -> Vec<ServerDef> {
539
args: vec!["--mcp".to_string()],
540
env: Vec::new(),
541
},
542
+ follows_session: false,
543
});
544
} else {
545
log::debug!("mcp: `{SMBCLOUD_COMMAND}` not on PATH; skipping the smbcloud server");
@@ -529,6 +570,7 @@ fn load_configs() -> Vec<ServerDef> {
570
args: vec!["--mcp".to_string()],
571
env: Vec::new(),
572
},
573
+ follows_session: false,
574
});
575
} else {
576
log::debug!("mcp: `{ASCAPI_COMMAND}` not on PATH; skipping the Ascapi server");
@@ -588,12 +630,12 @@ pub async fn init() {
630
let servers = futures::future::join_all(load_configs().into_iter().map(connect)).await;
631
632
for server in &servers {
591
- match &server.error {
633
+ match server.error() {
634
Some(error) => log::warn!("mcp: server '{}' unavailable: {error}", server.name),
635
None => log::info!(
636
"mcp: server '{}' ready, {} tool(s)",
637
server.name,
596
- server.tools.len()
638
+ server.tools().len()
639
),
640
}
641
}
@@ -610,21 +652,34 @@ fn client_info() -> ClientInfo {
652
/// `ServerConn`; failures land in its `error` field rather than propagating.
653
async fn connect(def: ServerDef) -> ServerConn {
654
let endpoint = def.transport.endpoint();
613
- let (conn, tools, error) =
614
- match Connection::connect(&def.spec(), &client_info(), HANDSHAKE_TIMEOUT).await {
615
- Ok(conn) => {
616
- let tools = server_tools(&def.name, &conn);
617
- (Some(Arc::new(conn)), tools, None)
618
- }
619
- Err(error) => (None, Vec::new(), Some(format!("{error:#}"))),
620
- };
655
+ let token = def.session_token();
656
+ let (live, tools, error) = match Connection::connect(
657
+ &def.spec_with(token.as_deref()),
658
+ &client_info(),
659
+ HANDSHAKE_TIMEOUT,
660
+ )
661
+ .await
662
+ {
663
+ Ok(conn) => {
664
+ let tools = server_tools(&def.name, &conn);
665
+ (
666
+ Some(Live {
667
+ conn: Arc::new(conn),
668
+ token,
669
+ }),
670
+ tools,
671
+ None,
672
+ )
673
+ }
674
+ Err(error) => (None, Vec::new(), Some(format!("{error:#}"))),
675
+ };
676
ServerConn {
677
name: def.name.clone(),
678
endpoint,
679
def,
625
- conn: RwLock::new(conn),
626
- tools,
627
- error,
680
+ conn: RwLock::new(live),
681
+ tools: StdRwLock::new(tools),
682
+ error: StdRwLock::new(error),
683
}
684
}
685
@@ -744,30 +799,35 @@ pub async fn connect_session_servers(servers: Vec<ClientServer>) -> SessionServe
799
if name.is_empty() || taken.contains(&name) {
800
clashes.push(ServerConn {
801
endpoint: transport.endpoint(),
747
- error: Some(if name.is_empty() {
802
+ error: StdRwLock::new(Some(if name.is_empty() {
803
"the client sent a server with no name".to_string()
804
} else {
805
format!("another MCP server is already named '{name}'")
751
- }),
806
+ })),
807
def: ServerDef {
808
name: name.clone(),
809
transport,
810
+ follows_session: false,
811
},
812
conn: RwLock::new(None),
757
- tools: Vec::new(),
813
+ tools: StdRwLock::new(Vec::new()),
814
name,
815
});
816
continue;
817
}
818
taken.push(name.clone());
763
- defs.push(ServerDef { name, transport });
819
+ defs.push(ServerDef {
820
+ name,
821
+ transport,
822
+ follows_session: false,
823
+ });
824
}
825
826
let mut connected = futures::future::join_all(defs.into_iter().map(connect)).await;
827
connected.extend(clashes);
828
829
for server in &connected {
770
- match &server.error {
830
+ match server.error() {
831
Some(error) => log::warn!(
832
"mcp: client-supplied server '{}' unavailable: {error}",
833
server.name
@@ -775,7 +835,7 @@ pub async fn connect_session_servers(servers: Vec<ClientServer>) -> SessionServe
835
None => log::info!(
836
"mcp: client-supplied server '{}' ready, {} tool(s)",
837
server.name,
778
- server.tools.len()
838
+ server.tools().len()
839
),
840
}
841
}
@@ -809,7 +869,7 @@ pub fn tool_specs() -> Vec<ToolSpec> {
869
let startup = MCP.get().map(|mcp| mcp.servers.as_slice()).unwrap_or(&[]);
870
let mut specs = Vec::new();
871
for server in startup.iter().chain(session.0.iter()) {
812
- for tool in &server.tools {
872
+ for tool in server.tools().iter() {
873
specs.push(ToolSpec {
874
name: tool.full_name.clone(),
875
description: tool.description.clone(),
@@ -830,11 +890,11 @@ pub async fn call_tool(full_name: &str, arguments: &str) -> String {
890
return "Error: MCP is not initialized.".to_string();
891
}
892
833
- let Some((server, tool)) = startup.iter().chain(session.0.iter()).find_map(|s| {
834
- s.tools
893
+ let Some((server, remote_name)) = startup.iter().chain(session.0.iter()).find_map(|s| {
894
+ s.tools()
895
.iter()
896
.find(|t| t.full_name == full_name)
837
- .map(|t| (s, t))
897
+ .map(|t| (s, t.remote_name.clone()))
898
}) else {
899
return format!("Error: unknown MCP tool \"{full_name}\".");
900
};
@@ -851,7 +911,7 @@ pub async fn call_tool(full_name: &str, arguments: &str) -> String {
911
}
912
};
913
854
- match call_server(server, &tool.remote_name, args).await {
914
+ match call_server(server, &remote_name, args).await {
915
Ok(text) => truncate(text),
916
Err(error) => format!("Error: {error}"),
917
}
@@ -875,7 +935,20 @@ async fn call_server(
935
} else {
936
"server"
937
};
878
- let Some(conn) = server.conn.read().await.clone() else {
938
+ // Reconnect before sending if the signed-in account changed, so a call is
939
+ // never made with the previous account's bearer.
940
+ if server.session_changed().await
941
+ && let Err(error) = reconnect(server).await
942
+ {
943
+ return Err(format!("{kind} '{}': {error}", server.name));
944
+ }
945
+ let Some(conn) = server
946
+ .conn
947
+ .read()
948
+ .await
949
+ .as_ref()
950
+ .map(|live| live.conn.clone())
951
+ else {
952
return Err(format!("{kind} '{}' is not connected", server.name));
953
};
954
let result = match conn.call_tool(remote_name, args.clone(), timeout).await {
@@ -885,16 +958,12 @@ async fn call_server(
958
}
959
Err(error) => {
960
log::info!("mcp: '{}' failed ({error:#}); reconnecting", server.name);
888
- let fresh = Connection::connect(&server.def.spec(), &client_info(), HANDSHAKE_TIMEOUT)
889
- .await
890
- .map(Arc::new)
891
- .map_err(|again| {
892
- format!(
893
- "server '{}': {error:#}; reconnecting failed: {again:#}",
894
- server.name
895
- )
896
- })?;
897
- *server.conn.write().await = Some(fresh.clone());
961
+ let fresh = reconnect(server).await.map_err(|again| {
962
+ format!(
963
+ "server '{}': {error:#}; reconnecting failed: {again}",
964
+ server.name
965
+ )
966
+ })?;
967
fresh
968
.call_tool(remote_name, args, timeout)
969
.await
@@ -904,6 +973,78 @@ async fn call_server(
973
Ok(render_tool_result(&result))
974
}
975
976
+impl ServerConn {
977
+ /// Whether this server follows the signed-in account and the stored token
978
+ /// is no longer the one its connection was opened with (or it has none).
979
+ async fn session_changed(&self) -> bool {
980
+ if !self.def.follows_session {
981
+ return false;
982
+ }
983
+ let current = crate::credentials::load_token();
984
+ match self.conn.read().await.as_ref() {
985
+ Some(live) => live.token != current,
986
+ None => true,
987
+ }
988
+ }
989
+}
990
+
991
+/// Open a fresh connection with the current spec and swap it in, along with
992
+/// the tools it lists and the handshake outcome.
993
+async fn reconnect(server: &ServerConn) -> Result<Arc<Connection>, String> {
994
+ let token = server.def.session_token();
995
+ let result = Connection::connect(
996
+ &server.def.spec_with(token.as_deref()),
997
+ &client_info(),
998
+ HANDSHAKE_TIMEOUT,
999
+ )
1000
+ .await;
1001
+ match result {
1002
+ Ok(conn) => {
1003
+ let conn = Arc::new(conn);
1004
+ *server.tools.write().unwrap_or_else(|e| e.into_inner()) =
1005
+ server_tools(&server.name, &conn);
1006
+ *server.error.write().unwrap_or_else(|e| e.into_inner()) = None;
1007
+ *server.conn.write().await = Some(Live {
1008
+ conn: conn.clone(),
1009
+ token,
1010
+ });
1011
+ Ok(conn)
1012
+ }
1013
+ Err(error) => {
1014
+ let message = format!("{error:#}");
1015
+ // Drop the old connection: it carries a bearer that no longer
1016
+ // matches the stored account.
1017
+ if server.def.follows_session {
1018
+ *server.conn.write().await = None;
1019
+ server
1020
+ .tools
1021
+ .write()
1022
+ .unwrap_or_else(|e| e.into_inner())
1023
+ .clear();
1024
+ }
1025
+ *server.error.write().unwrap_or_else(|e| e.into_inner()) = Some(message.clone());
1026
+ Err(message)
1027
+ }
1028
+ }
1029
+}
1030
+
1031
+/// Reconnect the official server if the stored sign-in changed since it
1032
+/// connected, so its tool list and `/mcp` entry update right away. No-op when
1033
+/// MCP isn't initialized or the token is unchanged.
1034
+pub async fn refresh_session_auth() {
1035
+ let Some(mcp) = MCP.get() else {
1036
+ return;
1037
+ };
1038
+ for server in mcp.servers.iter().filter(|s| s.def.follows_session) {
1039
+ if server.session_changed().await {
1040
+ log::info!("mcp: '{}' sign-in changed; reconnecting", server.name);
1041
+ if let Err(error) = reconnect(server).await {
1042
+ log::warn!("mcp: server '{}' unavailable: {error}", server.name);
1043
+ }
1044
+ }
1045
+ }
1046
+}
1047
+
1048
/// Flatten an MCP `tools/call` result into text. Joins text content blocks;
1049
/// notes non-text blocks; honors `isError`.
1050
fn render_tool_result(result: &Value) -> String {
@@ -993,8 +1134,8 @@ pub fn overview() -> Option<Overview> {
1134
let mcp = MCP.get()?;
1135
Some(Overview {
1136
servers: mcp.servers.len(),
996
- tools: mcp.servers.iter().map(|s| s.tools.len()).sum(),
997
- unavailable: mcp.servers.iter().filter(|s| s.error.is_some()).count(),
1137
+ tools: mcp.servers.iter().map(|s| s.tools().len()).sum(),
1138
+ unavailable: mcp.servers.iter().filter(|s| s.error().is_some()).count(),
1139
})
1140
}
1141
@@ -1002,7 +1143,7 @@ pub fn overview() -> Option<Overview> {
1143
/// session's client supplied, which are marked as such.
1144
fn summarize_servers(startup: &[ServerConn], session: &[ServerConn]) -> String {
1145
let all = || startup.iter().chain(session.iter());
1005
- let total_tools: usize = all().map(|s| s.tools.len()).sum();
1146
+ let total_tools: usize = all().map(|s| s.tools().len()).sum();
1147
let mut lines = vec![format!(
1148
"{} MCP server(s), {total_tools} tool(s) available:",
1149
all().count()
@@ -1013,7 +1154,7 @@ fn summarize_servers(startup: &[ServerConn], session: &[ServerConn]) -> String {
1154
} else {
1155
""
1156
};
1016
- match &server.error {
1157
+ match server.error() {
1158
Some(error) => lines.push(format!(
1159
"- {} ({}{origin}) — unavailable: {error}",
1160
server.name, server.endpoint
@@ -1023,9 +1164,9 @@ fn summarize_servers(startup: &[ServerConn], session: &[ServerConn]) -> String {
1164
"- {} ({}{origin}) — {} tool(s)",
1165
server.name,
1166
server.endpoint,
1026
- server.tools.len()
1167
+ server.tools().len()
1168
));
1028
- for tool in &server.tools {
1169
+ for tool in server.tools().iter() {
1170
lines.push(format!(" • {}", tool.full_name));
1171
}
1172
}
@@ -1256,6 +1397,7 @@ mod tests {
1397
url: "u1".into(),
1398
headers: vec![],
1399
},
1400
+ follows_session: false,
1401
}];
1402
upsert(
1403
&mut defs,
@@ -1265,6 +1407,7 @@ mod tests {
1407
url: "u2".into(),
1408
headers: vec![],
1409
},
1410
+ follows_session: false,
1411
},
1412
);
1413
assert_eq!(defs.len(), 1);
@@ -1315,4 +1458,53 @@ mod tests {
1458
// Without init() the global is unset; this must not panic.
1459
assert!(super::tool_specs().is_empty() || MCP.get().is_some());
1460
}
1461
+
1462
+ #[test]
1463
+ fn the_official_server_follows_the_session_without_a_static_header() {
1464
+ let defs = load_configs();
1465
+ let Some(official) = defs.iter().find(|d| d.name == OFFICIAL_SERVER_NAME) else {
1466
+ // SIGIT_MCP or SIGIT_MCP_OFFICIAL is set in this environment.
1467
+ return;
1468
+ };
1469
+ assert!(official.follows_session);
1470
+ match &official.transport {
1471
+ TransportDef::Http { headers, .. } => assert!(
1472
+ !headers
1473
+ .iter()
1474
+ .any(|(key, _)| key.eq_ignore_ascii_case("authorization"))
1475
+ ),
1476
+ TransportDef::Stdio { .. } => panic!("the official server is HTTP"),
1477
+ }
1478
+ }
1479
+
1480
+ #[test]
1481
+ fn a_user_defined_sigit_entry_keeps_its_header_and_does_not_follow() {
1482
+ let file: McpFile = toml::from_str(
1483
+ r#"
1484
+ [[server]]
1485
+ name = "sigit"
1486
+ url = "https://example.com/mcp"
1487
+ headers = { Authorization = "Bearer static" }
1488
+ "#,
1489
+ )
1490
+ .unwrap();
1491
+ let entry = &file.server[0];
1492
+ let mut defs = Vec::new();
1493
+ upsert(
1494
+ &mut defs,
1495
+ ServerDef {
1496
+ name: sanitize(&entry.name),
1497
+ transport: entry.transport_def().unwrap(),
1498
+ follows_session: false,
1499
+ },
1500
+ );
1501
+ assert!(!defs[0].follows_session);
1502
+ assert_eq!(defs[0].session_token(), None);
1503
+ match &defs[0].transport {
1504
+ TransportDef::Http { headers, .. } => {
1505
+ assert_eq!(headers[0].1, "Bearer static");
1506
+ }
1507
+ TransportDef::Stdio { .. } => panic!("expected HTTP"),
1508
+ }
1509
+ }
1510
}
crates/cli/tests/acp_mcp_reauth.rs
new
+311
@@ -0,0 +1,311 @@
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
+}