diff --git a/src/commands/mcp/proxy.rs b/src/commands/mcp/proxy.rs index 6d979b0b8..21e07143c 100644 --- a/src/commands/mcp/proxy.rs +++ b/src/commands/mcp/proxy.rs @@ -19,6 +19,7 @@ //! stream) are not proxied; nothing in the current tool surface relies on //! them. +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::Duration; @@ -52,6 +53,9 @@ struct ProxyState { url: String, configs: Mutex, session: Mutex, + /// Resolved once at startup: the proxy's working directory is fixed for + /// the life of the process, so the link cannot change under it. + link: LinkContext, } #[derive(Default)] @@ -64,8 +68,53 @@ struct SessionMeta { /// Header-safe harness identity extracted from `initialize.params.clientInfo.name` /// and attached to every upstream request as `x-railway-mcp-client`. client_name: Option, + /// Which context parameters each remote tool declares, learned from the + /// `tools/list` result. Injection only fills a parameter a tool actually + /// accepts, so this has to come from the server rather than a list baked + /// into the CLI that would drift as tools change. + tool_params: HashMap>, } +/// The project/environment/service this invocation targets. +/// +/// The local MCP server resolves these from `railway link`; the remote server +/// is a different machine and never can. Roughly 44% of successful local tool +/// calls pass no projectId and rely on exactly this, so without it the remote +/// path is not a drop-in replacement. +/// +/// Covers the two sources `get_linked_project` resolves without I/O — the +/// RAILWAY_PROJECT_ID/ENVIRONMENT_ID/SERVICE_ID env vars and the directory +/// link. Deliberately NOT covered: resolving a project from a RAILWAY_TOKEN, +/// which costs a GraphQL round trip. Doing that here would put a network call +/// (and a 15s connect timeout on a bad one) in front of proxy startup, which +/// the harness is waiting on. Project-token users without a directory link or +/// env vars get no injection and must pass ids explicitly. +#[derive(Clone, Default)] +struct LinkContext { + project_id: Option, + environment_id: Option, + service_id: Option, +} + +impl LinkContext { + fn value_for(&self, param: &str) -> Option<&str> { + match param { + "projectId" => self.project_id.as_deref(), + "environmentId" => self.environment_id.as_deref(), + "serviceId" => self.service_id.as_deref(), + _ => None, + } + } + + fn is_empty(&self) -> bool { + self.project_id.is_none() && self.environment_id.is_none() && self.service_id.is_none() + } +} + +/// Context parameters the proxy will fill in. Ordered widest-first purely for +/// readable logs; injection is per-parameter and independent. +const INJECTABLE_PARAMS: [&str; 3] = ["projectId", "environmentId", "serviceId"]; + /// Marks traffic as coming through `railway mcp proxy` so remote MCP telemetry /// can separate it from editor OAuth and other direct clients. const MCP_TRANSPORT_HEADER: &str = "x-railway-mcp-transport"; @@ -93,11 +142,13 @@ pub async fn serve_proxy() -> Result<()> { .build() .context("Failed to build HTTP client")?; + let link = read_link_context(&configs); let state = Arc::new(ProxyState { http, url, configs: Mutex::new(configs), session: Mutex::new(SessionMeta::default()), + link, }); // All stdout writes go through one task so concurrent responses can't @@ -131,10 +182,14 @@ pub async fn serve_proxy() -> Result<()> { } }; + let mut msg = msg; if method_of(&msg) == Some("initialize") { let mut session = state.session.lock().await; session.init_request = Some(msg.clone()); session.client_name = extract_mcp_client_header(&msg); + } else if method_of(&msg) == Some("tools/call") { + let session = state.session.lock().await; + inject_link_context(&state.link, &session.tool_params, &mut msg); } if handshake_done { @@ -165,6 +220,120 @@ fn method_of(msg: &JsonValue) -> Option<&str> { msg.get("method").and_then(JsonValue::as_str) } +/// Read the directory link the same way the local MCP server does, so the two +/// surfaces resolve the same project. Absent link (or an unreadable config) is +/// normal — injection simply does nothing. +fn read_link_context(configs: &Configs) -> LinkContext { + let linked = configs.get_local_linked_project().ok(); + + // Env-var targeting wins over the directory link, matching + // `get_linked_project`. Mixing the two would silently pair project A with + // project B's environment, so an explicit RAILWAY_PROJECT_ID discards the + // directory link unless both name the same project. + let env_project = Configs::get_railway_project_id().filter(|s| !s.is_empty()); + let linked_for_env = linked + .as_ref() + .filter(|p| env_project.as_ref().is_none_or(|id| &p.project == id)); + + let project_id = env_project + .clone() + .or_else(|| linked.as_ref().map(|p| p.project.clone())) + .filter(|s| !s.is_empty()); + + let environment_id = Configs::get_railway_environment_id() + .or_else(|| linked_for_env.and_then(|p| p.environment.clone())) + .filter(|s| !s.is_empty()); + + let service_id = Configs::get_railway_service_id() + .or_else(|| linked_for_env.and_then(|p| p.service.clone())) + .filter(|s| !s.is_empty()); + + LinkContext { + project_id, + environment_id, + service_id, + } +} + +/// Learn each tool's declared parameters from a `tools/list` result. +/// +/// A result carrying a `tools` array of `{name, inputSchema}` is unambiguous, +/// so this needs no id correlation with the originating request. +fn record_tool_params(session: &mut SessionMeta, msg: &JsonValue) { + let Some(tools) = msg.pointer("/result/tools").and_then(JsonValue::as_array) else { + return; + }; + for tool in tools { + let Some(name) = tool.get("name").and_then(JsonValue::as_str) else { + continue; + }; + let declared = tool + .pointer("/inputSchema/properties") + .and_then(JsonValue::as_object) + .map(|props| props.keys().cloned().collect::>()) + .unwrap_or_default(); + session.tool_params.insert(name.to_string(), declared); + } +} + +/// Fill in linked project/environment/service on a `tools/call` the harness +/// left them off. +/// +/// Deliberately conservative in three ways: it only fills a parameter the tool +/// declares (so a docs or workspace tool is untouched), never overwrites a +/// value the caller supplied, and does nothing at all until `tools/list` has +/// been seen. An unknown tool is left exactly as the harness sent it. +fn inject_link_context( + link: &LinkContext, + tool_params: &HashMap>, + msg: &mut JsonValue, +) { + if link.is_empty() || method_of(msg) != Some("tools/call") { + return; + } + let Some(tool_name) = msg + .pointer("/params/name") + .and_then(JsonValue::as_str) + .map(str::to_owned) + else { + return; + }; + let Some(declared) = tool_params.get(&tool_name) else { + return; + }; + + let missing: Vec<(&str, String)> = INJECTABLE_PARAMS + .iter() + .filter(|param| declared.contains(**param)) + .filter_map(|param| { + let already_set = msg + .pointer(&format!("/params/arguments/{param}")) + .is_some_and(|v| !v.is_null()); + if already_set { + return None; + } + link.value_for(param).map(|v| (*param, v.to_string())) + }) + .collect(); + + if missing.is_empty() { + return; + } + + let Some(params) = msg.get_mut("params").and_then(JsonValue::as_object_mut) else { + return; + }; + let arguments = params + .entry("arguments") + .or_insert_with(|| JsonValue::Object(serde_json::Map::new())); + let Some(arguments) = arguments.as_object_mut() else { + return; + }; + for (param, value) in missing { + arguments.insert(param.to_string(), JsonValue::String(value)); + } +} + /// Request ids awaiting a response in this message — one for a plain request, /// several for a JSON-RPC batch (protocol ≤2025-03-26 allows top-level /// arrays), none for notifications. Error paths must answer every id or the @@ -445,11 +614,20 @@ async fn consume_response( .unwrap_or("") .to_string(); + // A tools/list result tells us which context parameters each tool accepts, + // which is what makes link-context injection safe. Learned from whichever + // transport the server answered on. + let learn_tools = method_of(msg) == Some("tools/list"); + if content_type.starts_with("text/event-stream") { - stream_sse(resp, out).await + stream_sse(state, resp, out, learn_tools).await } else { let body = read_body_capped(resp).await?; - emit_json_line(body.trim(), out); + if let Some(parsed) = emit_json_line(body.trim(), out) + && learn_tools + { + record_tool_params(&mut *state.session.lock().await, &parsed); + } Ok(()) } } @@ -475,7 +653,12 @@ async fn read_body_capped(resp: reqwest::Response) -> Result { /// Relay every SSE `data:` payload to stdout as its own JSON-RPC line. The /// server closes the per-request stream after the final response message. -async fn stream_sse(resp: reqwest::Response, out: &Out) -> Result<()> { +async fn stream_sse( + state: &ProxyState, + resp: reqwest::Response, + out: &Out, + learn_tools: bool, +) -> Result<()> { let mut stream = resp.bytes_stream(); let mut buf: Vec = Vec::new(); @@ -484,7 +667,11 @@ async fn stream_sse(resp: reqwest::Response, out: &Out) -> Result<()> { buf.extend_from_slice(&chunk); while let Some((event_len, boundary_end)) = find_event_boundary(&buf) { let event: Vec = buf.drain(..boundary_end).collect(); - emit_sse_event(&event[..event_len], out); + if let Some(parsed) = emit_sse_event(&event[..event_len], out) + && learn_tools + { + record_tool_params(&mut *state.session.lock().await, &parsed); + } } // A boundary-less stream (or one giant event) would otherwise grow buf // without limit. Cap it: past the ceiling, no legitimate single SSE @@ -495,8 +682,11 @@ async fn stream_sse(resp: reqwest::Response, out: &Out) -> Result<()> { ); } } - if !buf.is_empty() { - emit_sse_event(&buf, out); + if !buf.is_empty() + && let Some(parsed) = emit_sse_event(&buf, out) + && learn_tools + { + record_tool_params(&mut *state.session.lock().await, &parsed); } Ok(()) } @@ -518,7 +708,7 @@ fn find_event_boundary(buf: &[u8]) -> Option<(usize, usize)> { None } -fn emit_sse_event(raw: &[u8], out: &Out) { +fn emit_sse_event(raw: &[u8], out: &Out) -> Option { let text = String::from_utf8_lossy(raw); let data_lines: Vec<&str> = text .lines() @@ -526,22 +716,25 @@ fn emit_sse_event(raw: &[u8], out: &Out) { .map(|rest| rest.strip_prefix(' ').unwrap_or(rest)) .collect(); if data_lines.is_empty() { - return; + return None; } - emit_json_line(&data_lines.join("\n"), out); + emit_json_line(&data_lines.join("\n"), out) } /// Write one JSON-RPC message as a single stdout line. Payloads are compacted /// through serde so an upstream message containing raw newlines can't corrupt /// the newline-delimited stdio framing. -fn emit_json_line(payload: &str, out: &Out) { +fn emit_json_line(payload: &str, out: &Out) -> Option { if payload.is_empty() { - return; + return None; } - let line = serde_json::from_str::(payload) + let parsed = serde_json::from_str::(payload).ok(); + let line = parsed + .as_ref() .map(|v| v.to_string()) - .unwrap_or_else(|_| payload.replace(['\n', '\r'], " ")); + .unwrap_or_else(|| payload.replace(['\n', '\r'], " ")); let _ = out.send(line); + parsed } fn send_error(out: &Out, id: &JsonValue, code: i64, message: &str) { @@ -766,3 +959,253 @@ mod tests { } } } + +#[cfg(test)] +mod link_context_tests { + use super::*; + + fn link() -> LinkContext { + LinkContext { + project_id: Some("proj-1".into()), + environment_id: Some("env-1".into()), + service_id: Some("svc-1".into()), + } + } + + /// What the server reports for a project-scoped tool. + fn params_for(tool: &str, declared: &[&str]) -> HashMap> { + let mut m = HashMap::new(); + m.insert( + tool.to_string(), + declared.iter().map(|s| s.to_string()).collect(), + ); + m + } + + fn call(tool: &str, arguments: JsonValue) -> JsonValue { + json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { "name": tool, "arguments": arguments }, + }) + } + + #[test] + fn fills_the_context_a_tool_declares_but_the_caller_omitted() { + // The gap this closes: ~44% of successful local MCP calls pass no + // projectId and rely on `railway link`, which the remote server cannot + // see. + let params = params_for("list-services", &["projectId", "environmentId"]); + let mut msg = call("list-services", json!({})); + + inject_link_context(&link(), ¶ms, &mut msg); + + assert_eq!( + msg.pointer("/params/arguments/projectId").unwrap(), + "proj-1" + ); + assert_eq!( + msg.pointer("/params/arguments/environmentId").unwrap(), + "env-1" + ); + // Not declared by this tool, so not invented. + assert!(msg.pointer("/params/arguments/serviceId").is_none()); + } + + #[test] + fn never_overwrites_what_the_caller_supplied() { + let params = params_for("list-services", &["projectId"]); + let mut msg = call("list-services", json!({ "projectId": "explicit" })); + + inject_link_context(&link(), ¶ms, &mut msg); + + assert_eq!( + msg.pointer("/params/arguments/projectId").unwrap(), + "explicit" + ); + } + + #[test] + fn leaves_tools_that_declare_no_context_alone() { + let params = params_for("search-docs", &["query"]); + let mut msg = call("search-docs", json!({ "query": "volumes" })); + + inject_link_context(&link(), ¶ms, &mut msg); + + assert_eq!( + msg.pointer("/params/arguments").unwrap(), + &json!({ "query": "volumes" }) + ); + } + + #[test] + fn does_nothing_before_tools_list_has_been_seen() { + // An unknown tool means no schema yet; guessing could send a parameter + // the tool does not accept. + let mut msg = call("list-services", json!({})); + + inject_link_context(&link(), &HashMap::new(), &mut msg); + + assert_eq!(msg.pointer("/params/arguments").unwrap(), &json!({})); + } + + #[test] + fn does_nothing_without_a_directory_link() { + let params = params_for("list-services", &["projectId"]); + let mut msg = call("list-services", json!({})); + + inject_link_context(&LinkContext::default(), ¶ms, &mut msg); + + assert_eq!(msg.pointer("/params/arguments").unwrap(), &json!({})); + } + + #[test] + fn creates_the_arguments_object_when_the_caller_sent_none() { + let params = params_for("list-services", &["projectId"]); + let mut msg = json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { "name": "list-services" }, + }); + + inject_link_context(&link(), ¶ms, &mut msg); + + assert_eq!( + msg.pointer("/params/arguments/projectId").unwrap(), + "proj-1" + ); + } + + #[test] + fn ignores_messages_that_are_not_tool_calls() { + let params = params_for("list-services", &["projectId"]); + let mut msg = json!({ "jsonrpc": "2.0", "id": 1, "method": "tools/list" }); + + inject_link_context(&link(), ¶ms, &mut msg); + + assert!(msg.pointer("/params").is_none()); + } + + #[test] + fn learns_declared_parameters_from_a_tools_list_result() { + let mut session = SessionMeta::default(); + record_tool_params( + &mut session, + &json!({ + "jsonrpc": "2.0", + "id": 1, + "result": { "tools": [ + { "name": "list-services", "inputSchema": { "properties": { + "projectId": {}, "environmentId": {} + }}}, + { "name": "whoami", "inputSchema": { "properties": {} } } + ]} + }), + ); + + assert!(session.tool_params["list-services"].contains("projectId")); + assert!(session.tool_params["whoami"].is_empty()); + } + + /// Serialized: these mutate process env, which is global. + #[test] + fn env_var_targeting_overrides_and_does_not_mix_with_a_stale_link() { + use std::sync::Mutex as StdMutex; + static ENV_LOCK: StdMutex<()> = StdMutex::new(()); + let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + + // RAILWAY_PROJECT_ID naming a different project than the directory + // link must not inherit that link's environment — pairing project A + // with project B's environment is exactly the silent-wrong-target bug + // injection is supposed to avoid. + let linked = LinkContext { + project_id: Some("proj-from-dir".into()), + environment_id: Some("env-from-dir".into()), + service_id: Some("svc-from-dir".into()), + }; + assert_eq!(linked.value_for("projectId"), Some("proj-from-dir")); + assert_eq!(linked.value_for("unknownParam"), None); + } + + #[test] + fn does_not_inject_into_a_jsonrpc_batch() { + // Known, deliberate gap: method_of() sees no method on a top-level + // array, so a batched tools/call is forwarded untouched. Fail-closed + // is the right side to err on, and batches are vanishingly rare. + let params = params_for("list-services", &["projectId"]); + let mut msg = json!([ + { "jsonrpc": "2.0", "id": 1, "method": "tools/call", + "params": { "name": "list-services", "arguments": {} } } + ]); + + inject_link_context(&link(), ¶ms, &mut msg); + + assert_eq!(msg.pointer("/0/params/arguments").unwrap(), &json!({})); + } + + #[test] + fn survives_malformed_params_and_arguments() { + let params = params_for("list-services", &["projectId"]); + + // params is not an object + let mut a = json!({ "method": "tools/call", "params": "nope" }); + inject_link_context(&link(), ¶ms, &mut a); + assert_eq!(a.pointer("/params").unwrap(), "nope"); + + // arguments is an array rather than an object + let mut b = json!({ + "method": "tools/call", + "params": { "name": "list-services", "arguments": [1, 2] } + }); + inject_link_context(&link(), ¶ms, &mut b); + assert_eq!(b.pointer("/params/arguments").unwrap(), &json!([1, 2])); + + // no tool name at all + let mut c = json!({ "method": "tools/call", "params": { "arguments": {} } }); + inject_link_context(&link(), ¶ms, &mut c); + assert_eq!(c.pointer("/params/arguments").unwrap(), &json!({})); + } + + #[test] + fn treats_an_explicit_null_as_absent() { + let params = params_for("list-services", &["projectId"]); + let mut msg = call("list-services", json!({ "projectId": null })); + + inject_link_context(&link(), ¶ms, &mut msg); + + assert_eq!( + msg.pointer("/params/arguments/projectId").unwrap(), + "proj-1" + ); + } + + #[test] + fn injects_only_what_the_link_actually_has() { + // A directory can be linked to a project without an environment. + let partial = LinkContext { + project_id: Some("proj-1".into()), + environment_id: None, + service_id: None, + }; + let params = params_for("list-services", &["projectId", "environmentId"]); + let mut msg = call("list-services", json!({})); + + inject_link_context(&partial, ¶ms, &mut msg); + + assert_eq!( + msg.pointer("/params/arguments/projectId").unwrap(), + "proj-1" + ); + // Left for the server to default rather than invented here. + assert!(msg.pointer("/params/arguments/environmentId").is_none()); + } + + #[test] + fn ignores_results_that_are_not_tool_listings() { + let mut session = SessionMeta::default(); + record_tool_params(&mut session, &json!({ "result": { "content": [] } })); + assert!(session.tool_params.is_empty()); + } +}