feat(sandbox): support dynamic WS limits, multipart upload and migrate conflict…

feat(sandbox): support dynamic WS limits, multipart upload and migrate conflict detection in rust agent (#7206)

projects/agent-sandbox-proxy/src/auth.rs

@@ -5,10 +5,27 @@ use std::sync::{LazyLock, OnceLock};
 use std::time::Duration;
 use tracing::{debug, error};
 
+#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
+pub struct WsLimits {
+    pub max_message_bytes: usize,
+    pub max_frame_bytes: usize,
+}
+
+impl Default for WsLimits {
+    fn default() -> Self {
+        Self {
+            max_message_bytes: 64 * 1024 * 1024,
+            max_frame_bytes: 16 * 1024 * 1024,
+        }
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct SandboxAddress {
     pub sandbox_url: Option<String>,
     pub agent_token: Option<String>,
+    #[serde(default)]
+    pub ws_limits: WsLimits,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]

projects/agent-sandbox-proxy/src/main.rs

@@ -15,9 +15,6 @@ mod relay;
 use auth::resolve_sandbox_address;
 use relay::handle_relay;
 
-const MAX_WS_MESSAGE_SIZE: usize = 16 * 1024 * 1024;
-const MAX_WS_FRAME_SIZE: usize = 4 * 1024 * 1024;
-
 #[derive(Deserialize)]
 struct WsQuery {
     ticket: Option<String>,
@@ -124,8 +121,9 @@ async fn fs_handler(ws: WebSocketUpgrade, Query(query): Query<WsQuery>) -> impl
             info!(
                 "[Auth] Ticket verified & address resolved successfully. Upgrading to WebSocket (FS)..."
             );
-            ws.max_message_size(MAX_WS_MESSAGE_SIZE)
-                .max_frame_size(MAX_WS_FRAME_SIZE)
+            let ws_limits = address.ws_limits;
+            ws.max_message_size(ws_limits.max_message_bytes)
+                .max_frame_size(ws_limits.max_frame_bytes)
                 .on_upgrade(move |socket| handle_relay(socket, address, claims, false))
                 .into_response()
         }
@@ -139,8 +137,9 @@ async fn terminal_handler(ws: WebSocketUpgrade, Query(query): Query<WsQuery>) ->
             info!(
                 "[Auth] Ticket verified & address resolved successfully. Upgrading to WebSocket (TERMINAL)..."
             );
-            ws.max_message_size(MAX_WS_MESSAGE_SIZE)
-                .max_frame_size(MAX_WS_FRAME_SIZE)
+            let ws_limits = address.ws_limits;
+            ws.max_message_size(ws_limits.max_message_bytes)
+                .max_frame_size(ws_limits.max_frame_bytes)
                 .on_upgrade(move |socket| handle_relay(socket, address, claims, true))
                 .into_response()
         }

projects/agent-sandbox-proxy/src/relay.rs

@@ -7,15 +7,13 @@ use tokio_tungstenite::{
     tungstenite::{
         Error as WsError,
         client::IntoClientRequest,
+        error::CapacityError,
         protocol::{Message as WsMsg, WebSocketConfig},
     },
 };
 use tracing::{debug, error, info};
 
-use crate::auth::{SandboxAddress, get_http_client, get_proxy_secret};
-
-const MAX_WS_MESSAGE_SIZE: usize = 16 * 1024 * 1024;
-const MAX_WS_FRAME_SIZE: usize = 4 * 1024 * 1024;
+use crate::auth::{SandboxAddress, WsLimits, get_http_client, get_proxy_secret};
 
 type UpstreamWsStream =
     tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
@@ -29,11 +27,13 @@ enum UpstreamControl {
 const UPSTREAM_CONNECT_MAX_ATTEMPTS: u8 = 10;
 const UPSTREAM_CONNECT_RETRY_DELAY: Duration = Duration::from_millis(300);
 const LOOPBACK_REWRITE_HOST_ENV: &str = "AGENT_SANDBOX_PROXY_REWRITE_HOST";
+const WS_CLOSE_MESSAGE_TOO_BIG_CODE: u16 = 1009;
+const WS_CLOSE_INTERNAL_ERROR_CODE: u16 = 1011;
 
-fn upstream_ws_config() -> WebSocketConfig {
+fn upstream_ws_config(ws_limits: WsLimits) -> WebSocketConfig {
     WebSocketConfig::default()
-        .max_message_size(Some(MAX_WS_MESSAGE_SIZE))
-        .max_frame_size(Some(MAX_WS_FRAME_SIZE))
+        .max_message_size(Some(ws_limits.max_message_bytes))
+        .max_frame_size(Some(ws_limits.max_frame_bytes))
 }
 
 fn is_loopback_host(host: &str) -> bool {
@@ -177,7 +177,10 @@ mod tests {
 }
 
 /// 连接沙盒内的 IDE Agent，允许 agent 冷启动时出现短暂端口不可用。
-async fn connect_upstream_with_retry(target_url: String) -> Result<UpstreamWsStream, String> {
+async fn connect_upstream_with_retry(
+    target_url: String,
+    ws_limits: WsLimits,
+) -> Result<UpstreamWsStream, String> {
     let mut attempts = 0;
     let safe_target_url = redact_sensitive_query(&target_url);
 
@@ -189,7 +192,7 @@ async fn connect_upstream_with_retry(target_url: String) -> Result<UpstreamWsStr
             Err(err) => return Err(format!("Failed to build WebSocket request: {}", err)),
         };
 
-        match connect_async_with_config(request, Some(upstream_ws_config()), false).await {
+        match connect_async_with_config(request, Some(upstream_ws_config(ws_limits)), false).await {
             Ok((ws, _)) => return Ok(ws),
             Err(err) => {
                 let err_str = err.to_string();
@@ -230,6 +233,7 @@ pub async fn handle_relay(
         return;
     };
     let permission_to_forward = claims.permission.as_str();
+    let ws_limits = address.ws_limits;
 
     let target_url = match address.sandbox_url.as_deref().filter(|url| !url.is_empty()) {
         Some(url) => {
@@ -262,7 +266,7 @@ pub async fn handle_relay(
         redact_sensitive_query(&target_url)
     );
 
-    let connect_fut = connect_upstream_with_retry(target_url);
+    let connect_fut = connect_upstream_with_retry(target_url, ws_limits);
     tokio::pin!(connect_fut);
 
     let mut buffer: Vec<AxumMsg> = Vec::new();
@@ -415,7 +419,7 @@ pub async fn handle_relay(
                         } else {
                             error!("[WSProxy] Error sending Ping frame to Upstream Devbox: {}", err);
                         }
-                        send_client_close(&client_to_upstream_close_tx, 1011, "Sandbox agent connection lost");
+                        send_client_close(&client_to_upstream_close_tx, WS_CLOSE_INTERNAL_ERROR_CODE, "Sandbox agent connection lost");
                         break;
                     }
                 }
@@ -443,7 +447,8 @@ pub async fn handle_relay(
                                         } else {
                                             error!("[WSProxy] Error forwarding client message to upstream: {}", err);
                                         }
-                                        send_client_close(&client_to_upstream_close_tx, 1011, "Sandbox agent connection lost");
+                                        let (code, reason) = upstream_error_client_close(&err);
+                                        send_client_close(&client_to_upstream_close_tx, code, reason);
                                         break;
                                     },
                                     Err(_) => debug!("[WSProxy] Dropping unsupported client WebSocket message."),
@@ -495,11 +500,8 @@ pub async fn handle_relay(
                     } else {
                         error!("[WSProxy] Upstream stream read error: {}", err);
                     }
-                    send_client_close(
-                        &upstream_to_client_close_tx,
-                        1011,
-                        "Sandbox agent connection lost",
-                    );
+                    let (code, reason) = upstream_error_client_close(&err);
+                    send_client_close(&upstream_to_client_close_tx, code, reason);
                     break;
                 }
             }
@@ -671,6 +673,21 @@ fn is_upstream_closed_error(err: &WsError) -> bool {
         )
 }
 
+fn is_ws_message_too_big_error(err: &WsError) -> bool {
+    matches!(err, WsError::Capacity(CapacityError::MessageTooLong { .. }))
+}
+
+fn upstream_error_client_close(err: &WsError) -> (u16, &'static str) {
+    if is_ws_message_too_big_error(err) {
+        (WS_CLOSE_MESSAGE_TOO_BIG_CODE, "Sandbox message too large")
+    } else {
+        (
+            WS_CLOSE_INTERNAL_ERROR_CODE,
+            "Sandbox agent connection lost",
+        )
+    }
+}
+
 async fn close_client_ws(client_sink: &mut ClientWsSink, code: u16, reason: &str) {
     let _ = client_sink.send(client_close_message(code, reason)).await;
 }

projects/fastgpt-ide-agent/Cargo.lock

@@ -116,6 +116,7 @@ dependencies = [
  "portable-pty",
  "serde",
  "serde_json",
+ "sha2",
  "tempfile",
  "tokio",
  "tokio-tungstenite",
@@ -605,6 +606,17 @@ dependencies = [
 ]
 
 [[package]]
+name = "sha2"
+version = "0.10.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283"
+dependencies = [
+ "cfg-if",
+ "cpufeatures",
+ "digest",
+]
+
+[[package]]
 name = "shared_library"
 version = "0.1.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"

projects/fastgpt-ide-agent/Cargo.toml

@@ -15,6 +15,7 @@ futures-util = "0.3.32"
 path-security = "0.2.0"
 notify = "9.0.0-rc.4"
 notify-debouncer-full = "0.8.0-rc.2"
+sha2 = "0.10.9"
 
 [dev-dependencies]
 tempfile = "3.27.0"

projects/fastgpt-ide-agent/src/connection.rs

@@ -3,10 +3,19 @@ use std::sync::Arc;
 use tokio::net::TcpStream;
 
 use crate::fs::{FsPermission, handle_fs_session};
-const MAX_WS_MESSAGE_SIZE: usize = 16 * 1024 * 1024;
-const MAX_WS_FRAME_SIZE: usize = 4 * 1024 * 1024;
 use crate::terminal::handle_terminal_session;
 
+const WS_MAX_MESSAGE_BYTES_ENV: &str = "FASTGPT_IDE_WS_MAX_MESSAGE_BYTES";
+const WS_MAX_FRAME_BYTES_ENV: &str = "FASTGPT_IDE_WS_MAX_FRAME_BYTES";
+
+fn read_positive_usize_env(key: &str) -> Option<usize> {
+    std::env::var(key)
+        .ok()?
+        .parse::<usize>()
+        .ok()
+        .filter(|value| *value > 0)
+}
+
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 enum Channel {
     Fs,
@@ -14,9 +23,16 @@ enum Channel {
 }
 
 fn ws_config() -> tokio_tungstenite::tungstenite::protocol::WebSocketConfig {
-    tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default()
-        .max_message_size(Some(MAX_WS_MESSAGE_SIZE))
-        .max_frame_size(Some(MAX_WS_FRAME_SIZE))
+    let config = tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default();
+    let config = match read_positive_usize_env(WS_MAX_MESSAGE_BYTES_ENV) {
+        Some(max_message_bytes) => config.max_message_size(Some(max_message_bytes)),
+        None => config,
+    };
+
+    match read_positive_usize_env(WS_MAX_FRAME_BYTES_ENV) {
+        Some(max_frame_bytes) => config.max_frame_size(Some(max_frame_bytes)),
+        None => config,
+    }
 }
 
 fn extract_token(

projects/fastgpt-ide-agent/src/protocol.rs

-use serde::{Deserialize, Serialize};
+use serde::{Deserialize, Deserializer, Serialize, Serializer};
+
+/// JSON-RPC 错误码的语义枚举，序列化时仍保持协议要求的数字格式。
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum JsonRpcErrorCode {
+    MethodNotFound,
+    InternalError,
+    FileConflict,
+    PermissionDenied,
+    FileTooLarge,
+    Unknown(i32),
+}
+
+impl JsonRpcErrorCode {
+    pub const fn as_i32(self) -> i32 {
+        match self {
+            Self::MethodNotFound => -32601,
+            Self::InternalError => -32603,
+            Self::FileConflict => -32001,
+            Self::PermissionDenied => -32003,
+            Self::FileTooLarge => -32004,
+            Self::Unknown(code) => code,
+        }
+    }
+
+    pub const fn from_i32(code: i32) -> Self {
+        match code {
+            -32601 => Self::MethodNotFound,
+            -32603 => Self::InternalError,
+            -32001 => Self::FileConflict,
+            -32003 => Self::PermissionDenied,
+            -32004 => Self::FileTooLarge,
+            _ => Self::Unknown(code),
+        }
+    }
+}
+
+impl Serialize for JsonRpcErrorCode {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        serializer.serialize_i32(self.as_i32())
+    }
+}
+
+impl<'de> Deserialize<'de> for JsonRpcErrorCode {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let code = i32::deserialize(deserializer)?;
+        Ok(Self::from_i32(code))
+    }
+}
 
 #[derive(Deserialize, Serialize, Debug)]
 pub struct JsonRpcRequest {
@@ -20,7 +74,7 @@ pub struct JsonRpcResponse {
 
 #[derive(Serialize, Deserialize, Debug)]
 pub struct JsonRpcError {
-    pub code: i32,
+    pub code: JsonRpcErrorCode,
     pub message: String,
 }
 
@@ -30,3 +84,28 @@ pub struct JsonRpcNotification {
     pub method: String,
     pub params: serde_json::Value,
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    #[test]
+    fn json_rpc_error_code_uses_numeric_wire_format() {
+        let error = JsonRpcError {
+            code: JsonRpcErrorCode::FileTooLarge,
+            message: "File is too large".to_string(),
+        };
+
+        let serialized = serde_json::to_value(error).unwrap();
+        assert_eq!(serialized["code"], json!(-32004));
+
+        let known_error: JsonRpcError =
+            serde_json::from_value(json!({ "code": -32001, "message": "conflict" })).unwrap();
+        assert_eq!(known_error.code, JsonRpcErrorCode::FileConflict);
+
+        let unknown_error: JsonRpcError =
+            serde_json::from_value(json!({ "code": -32099, "message": "custom" })).unwrap();
+        assert_eq!(unknown_error.code, JsonRpcErrorCode::Unknown(-32099));
+    }
+}