huskies: merge 1208 story Ops/LLM sessions can reach gateway-mode + cross-project MCP (the biggest shell-fallback cause)
This commit is contained in:
@@ -298,7 +298,21 @@ pub async fn gateway_mcp_post_handler(
|
||||
handle_gateway_tool(tool_name, &rpc.params, &state, rpc.id.clone()).await,
|
||||
)
|
||||
} else {
|
||||
proxy_and_respond(&state, &bytes, rpc.id).await
|
||||
// Story 1208 AC 1: an explicit `project` argument on any
|
||||
// proxied tool call targets that project directly, without
|
||||
// requiring a prior `switch_project`.
|
||||
let explicit_project = rpc
|
||||
.params
|
||||
.get("arguments")
|
||||
.and_then(|a| a.get("project"))
|
||||
.and_then(|v| v.as_str())
|
||||
.filter(|p| !p.is_empty());
|
||||
match explicit_project {
|
||||
Some(project) => {
|
||||
proxy_and_respond_for_project(&state, project, &bytes, rpc.id).await
|
||||
}
|
||||
None => proxy_and_respond(&state, &bytes, rpc.id).await,
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => proxy_and_respond(&state, &bytes, rpc.id).await,
|
||||
@@ -323,6 +337,36 @@ async fn proxy_and_respond(state: &GatewayState, bytes: &[u8], id: Option<Value>
|
||||
}
|
||||
}
|
||||
|
||||
/// Proxy a request to an explicitly named project (story 1208 AC 1) rather
|
||||
/// than whatever project is currently active, so a single ops/LLM session
|
||||
/// can address any registered project per-call without a prior
|
||||
/// `switch_project`. Returns a JSON-RPC error when `project` is not a
|
||||
/// registered project name.
|
||||
async fn proxy_and_respond_for_project(
|
||||
state: &GatewayState,
|
||||
project: &str,
|
||||
bytes: &[u8],
|
||||
id: Option<Value>,
|
||||
) -> Response {
|
||||
{
|
||||
let projects = state.projects.read().await;
|
||||
if let Err(e) = gateway::config::validate_project_exists(&projects, project) {
|
||||
return to_json_response(JsonRpcResponse::error(id, -32602, e));
|
||||
}
|
||||
}
|
||||
match state.proxy_mcp_for_project(project, bytes).await {
|
||||
Ok(resp_body) => Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.header("Content-Type", "application/json")
|
||||
.body(Body::from(resp_body)),
|
||||
Err(e) => to_json_response(JsonRpcResponse::error(
|
||||
id,
|
||||
-32603,
|
||||
format!("proxy error: {e}"),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Stream an MCP tool call to the active sled as SSE, re-emitting each `data:`
|
||||
/// event from the sled to the originating gateway client without buffering.
|
||||
///
|
||||
@@ -441,12 +485,35 @@ async fn handle_tools_list(
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
|
||||
inject_project_arg_schema(&mut tools);
|
||||
|
||||
let mut all_tools = gateway_tool_definitions();
|
||||
all_tools.append(&mut tools);
|
||||
|
||||
Ok(JsonRpcResponse::success(id, json!({ "tools": all_tools })))
|
||||
}
|
||||
|
||||
/// Advertise the optional per-call `project` argument (story 1208 AC 1) on
|
||||
/// every proxied tool's `inputSchema.properties`, so MCP clients that
|
||||
/// validate call arguments against the declared schema don't reject it.
|
||||
/// Leaves any tool without an object `inputSchema.properties` untouched.
|
||||
fn inject_project_arg_schema(tools: &mut [Value]) {
|
||||
for tool in tools.iter_mut() {
|
||||
if let Some(props) = tool
|
||||
.get_mut("inputSchema")
|
||||
.and_then(|s| s.get_mut("properties"))
|
||||
.and_then(|p| p.as_object_mut())
|
||||
{
|
||||
props.entry("project".to_string()).or_insert_with(|| {
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "Optional: name of a registered gateway project to target instead of the currently active one (see list_projects). Omit to use the active project."
|
||||
})
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Gateway tool dispatch ────────────────────────────────────────────────────
|
||||
|
||||
/// Dispatch a gateway-specific tool call.
|
||||
@@ -1335,4 +1402,79 @@ mod tests {
|
||||
fn fleet_identity_is_in_gateway_tools() {
|
||||
assert!(GATEWAY_TOOLS.contains(&"fleet_identity"));
|
||||
}
|
||||
|
||||
// ── explicit per-call `project` targeting (story 1208 AC 1) ─────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_for_unknown_project_returns_invalid_params_error() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let state = make_test_state(dir.path());
|
||||
let resp =
|
||||
proxy_and_respond_for_project(&state, "nonexistent", b"{}", Some(json!(1))).await;
|
||||
let bytes = resp.into_body().into_bytes().await.unwrap();
|
||||
let parsed: Value = serde_json::from_slice(&bytes).unwrap();
|
||||
assert_eq!(parsed["error"]["code"], -32602);
|
||||
assert!(
|
||||
parsed["error"]["message"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("unknown project"),
|
||||
"got: {parsed}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_for_known_project_without_live_connection_returns_proxy_error() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let state = make_test_state(dir.path());
|
||||
let resp =
|
||||
proxy_and_respond_for_project(&state, "test-project", b"{}", Some(json!(1))).await;
|
||||
let bytes = resp.into_body().into_bytes().await.unwrap();
|
||||
let parsed: Value = serde_json::from_slice(&bytes).unwrap();
|
||||
assert_eq!(parsed["error"]["code"], -32603);
|
||||
assert!(
|
||||
parsed["error"]["message"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("no live WS uplink"),
|
||||
"got: {parsed}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inject_project_arg_schema_adds_property_to_object_schema() {
|
||||
let mut tools = vec![json!({
|
||||
"name": "create_story",
|
||||
"inputSchema": {"type": "object", "properties": {"name": {"type": "string"}}}
|
||||
})];
|
||||
inject_project_arg_schema(&mut tools);
|
||||
assert_eq!(
|
||||
tools[0]["inputSchema"]["properties"]["project"]["type"],
|
||||
"string"
|
||||
);
|
||||
assert_eq!(
|
||||
tools[0]["inputSchema"]["properties"]["name"]["type"],
|
||||
"string"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inject_project_arg_schema_does_not_overwrite_existing_project_property() {
|
||||
let mut tools = vec![json!({
|
||||
"name": "weird_tool",
|
||||
"inputSchema": {"type": "object", "properties": {"project": {"type": "integer"}}}
|
||||
})];
|
||||
inject_project_arg_schema(&mut tools);
|
||||
assert_eq!(
|
||||
tools[0]["inputSchema"]["properties"]["project"]["type"],
|
||||
"integer"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inject_project_arg_schema_skips_tool_without_properties() {
|
||||
let mut tools = vec![json!({"name": "no_schema"})];
|
||||
inject_project_arg_schema(&mut tools);
|
||||
assert_eq!(tools[0], json!({"name": "no_schema"}));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -354,6 +354,26 @@ impl GatewayState {
|
||||
self.sled_connections.write().await.remove(project_name);
|
||||
}
|
||||
|
||||
/// Look up the live sled connection for an explicit project name,
|
||||
/// returning a clone if one exists and has a recent heartbeat.
|
||||
///
|
||||
/// Returns `None` when no sled has connected for this project or when its
|
||||
/// heartbeat is overdue. Generalizes [`GatewayState::active_sled_connection`]
|
||||
/// to any registered project rather than only the active one (story 1208).
|
||||
pub async fn sled_connection_for(&self, project_name: &str) -> Option<SledConnection> {
|
||||
let conn = self
|
||||
.sled_connections
|
||||
.read()
|
||||
.await
|
||||
.get(project_name)
|
||||
.cloned()?;
|
||||
if conn.is_alive(HEARTBEAT_MAX_AGE_MS) {
|
||||
Some(conn)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Look up the live sled connection for the active project, returning a
|
||||
/// clone if one exists and has a recent heartbeat.
|
||||
///
|
||||
@@ -361,12 +381,7 @@ impl GatewayState {
|
||||
/// heartbeat is overdue.
|
||||
pub async fn active_sled_connection(&self) -> Option<SledConnection> {
|
||||
let name = self.active_project.read().await.clone();
|
||||
let conn = self.sled_connections.read().await.get(&name).cloned()?;
|
||||
if conn.is_alive(HEARTBEAT_MAX_AGE_MS) {
|
||||
Some(conn)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
self.sled_connection_for(&name).await
|
||||
}
|
||||
|
||||
/// Proxy an MCP request to the active project over its live sled-uplink
|
||||
@@ -389,6 +404,31 @@ impl GatewayState {
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Proxy an MCP request to an explicitly named project over its live
|
||||
/// sled-uplink WebSocket, without reading or mutating `active_project`
|
||||
/// (story 1208 AC 1).
|
||||
///
|
||||
/// Lets a single ops/LLM session address any registered project by name
|
||||
/// on a per-call basis instead of first mutating the shared
|
||||
/// `active_project` via `switch_project` — which avoids the race where a
|
||||
/// second concurrent caller's tool calls get proxied to the wrong
|
||||
/// project mid-session.
|
||||
///
|
||||
/// Returns the raw response body bytes ready to be relayed to the caller.
|
||||
pub async fn proxy_mcp_for_project(
|
||||
&self,
|
||||
project: &str,
|
||||
bytes: &[u8],
|
||||
) -> Result<Vec<u8>, String> {
|
||||
match self.sled_connection_for(project).await {
|
||||
Some(conn) => proxy_mcp_via_ws(&conn, bytes).await,
|
||||
None => Err(format!(
|
||||
"sled '{project}' has no live WS uplink connection; \
|
||||
ensure the sled is running and connected to this gateway"
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Public API ──────────────────────────────────────────────────────────────
|
||||
@@ -989,6 +1029,45 @@ mod tests {
|
||||
assert!(!state.sled_connections.read().await.contains_key("myproj"));
|
||||
}
|
||||
|
||||
// ── explicit per-project proxy targeting (story 1208) ──────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn sled_connection_for_unregistered_project_returns_none() {
|
||||
let config = make_config(&[("alpha", "http://a:3001")]);
|
||||
let state = GatewayState::new(config, PathBuf::new(), 3000).unwrap();
|
||||
assert!(state.sled_connection_for("nonexistent").await.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sled_connection_for_finds_connection_by_explicit_name_not_active() {
|
||||
let config = make_config(&[("alpha", "http://a:3001"), ("beta", "http://b:3002")]);
|
||||
let state = GatewayState::new(config, PathBuf::new(), 3000).unwrap();
|
||||
// "alpha" is active (first project), but the live connection is for "beta".
|
||||
assert_eq!(*state.active_project.read().await, "alpha");
|
||||
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let conn = SledConnection {
|
||||
tx,
|
||||
last_heartbeat_ms: Arc::new(AtomicI64::new(chrono::Utc::now().timestamp_millis())),
|
||||
in_flight: Arc::new(TokioMutex::new(HashMap::new())),
|
||||
};
|
||||
state
|
||||
.register_sled_connection("beta".to_string(), conn)
|
||||
.await;
|
||||
|
||||
assert!(state.sled_connection_for("beta").await.is_some());
|
||||
assert!(state.active_sled_connection().await.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_mcp_for_project_without_live_connection_fails() {
|
||||
let config = make_config(&[("alpha", "http://a:3001")]);
|
||||
let state = GatewayState::new(config, PathBuf::new(), 3000).unwrap();
|
||||
let result = state.proxy_mcp_for_project("alpha", b"{}").await;
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("no live WS uplink"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn auth_token_in_project_entry_populates_sled_tokens_map() {
|
||||
let mut projects = BTreeMap::new();
|
||||
|
||||
Reference in New Issue
Block a user