@@ -37,7 +37,9 @@ pub use crate::commands::provider_readiness::{ProviderWaitOptions, sandbox_provi
3737
3838use crate :: color:: Colorize ;
3939use crate :: policy_update:: build_policy_update_plan;
40- use crate :: tls:: { TlsOptions , grpc_client} ;
40+ use crate :: tls:: {
41+ TlsOptions , grpc_client, owner_replica, retry_unrouted, routed_to, with_owner_hint,
42+ } ;
4143use futures:: StreamExt ;
4244use miette:: { IntoDiagnostic , Result , WrapErr , miette} ;
4345use openshell_bootstrap:: {
@@ -71,6 +73,7 @@ use std::io::{ErrorKind, IsTerminal, Read, Write};
7173use std:: path:: { Path , PathBuf } ;
7274use std:: process:: Command ;
7375use std:: time:: { Duration , Instant } ;
76+ use tonic:: metadata:: AsciiMetadataValue ;
7477use tonic:: { Code , Status } ;
7578
7679const PROVISIONAL_CONTAINER_EXIT_RECONCILIATION_TIMEOUT : Duration = Duration :: from_secs ( 5 ) ;
@@ -1990,15 +1993,17 @@ pub async fn sandbox_exec_grpc(
19901993 let mut client = grpc_client ( server, tls) . await ?;
19911994
19921995 // Resolve sandbox name to id.
1993- let sandbox = client
1994- . get_sandbox ( GetSandboxRequest {
1996+ let response = client
1997+ . get_sandbox ( with_owner_hint ( GetSandboxRequest {
19951998 name : name. to_string ( ) ,
19961999 workspace_scope : Some ( openshell_core:: proto:: workspace_selector (
19972000 workspace. to_string ( ) ,
19982001 ) ) ,
1999- } )
2002+ } ) )
20002003 . await
2001- . into_diagnostic ( ) ?
2004+ . into_diagnostic ( ) ?;
2005+ let replica = owner_replica ( & response) ;
2006+ let sandbox = response
20022007 . into_inner ( )
20032008 . sandbox
20042009 . ok_or_else ( || miette:: miette!( "sandbox not found" ) ) ?;
@@ -2086,16 +2091,26 @@ pub async fn sandbox_exec_grpc(
20862091 stdin_is_terminal,
20872092 std:: mem:: take ( & mut request. stdin ) ,
20882093 stdin_rest,
2094+ replica. as_ref ( ) ,
20892095 )
20902096 . await ;
20912097 }
20922098
20932099 // Make the streaming gRPC call.
2094- let mut stream = client
2095- . exec_sandbox ( request)
2096- . await
2097- . into_diagnostic ( ) ?
2098- . into_inner ( ) ;
2100+ let retry = replica. is_some ( ) . then ( || request. clone ( ) ) ;
2101+ let mut stream = match (
2102+ client
2103+ . exec_sandbox ( routed_to ( request, replica. as_ref ( ) ) )
2104+ . await ,
2105+ retry,
2106+ ) {
2107+ ( Err ( status) , Some ( request) ) if retry_unrouted ( replica. as_ref ( ) , & status) => {
2108+ client. exec_sandbox ( request) . await
2109+ }
2110+ ( response, _) => response,
2111+ }
2112+ . into_diagnostic ( ) ?
2113+ . into_inner ( ) ;
20992114
21002115 // Stream output to terminal in real-time.
21012116 let mut exit_code = 0i32 ;
@@ -2148,7 +2163,7 @@ pub async fn service_forward_tcp(
21482163 let ( bind_addr, bind_port) = parse_tcp_forward_spec ( local, target_port) ?;
21492164 let mut client = grpc_client ( server, tls) . await ?;
21502165
2151- fetch_ready_sandbox_for_forward ( & mut client, name, workspace) . await ?;
2166+ let mut replica = fetch_ready_sandbox_for_forward ( & mut client, name, workspace) . await ?;
21522167
21532168 let listener = tokio:: net:: TcpListener :: bind ( ( bind_addr. as_str ( ) , bind_port) )
21542169 . await
@@ -2179,7 +2194,7 @@ pub async fn service_forward_tcp(
21792194 }
21802195
21812196 _ = health_check. tick( ) => {
2182- fetch_ready_sandbox_for_forward( & mut client, name, workspace) . await ?;
2197+ replica = fetch_ready_sandbox_for_forward( & mut client, name, workspace) . await ?;
21832198 }
21842199
21852200 accepted = listener. accept( ) => {
@@ -2193,6 +2208,7 @@ pub async fn service_forward_tcp(
21932208 let target_host = target_host. to_string( ) ;
21942209 let service_id = format!( "service-forward:{name}:{target_host}:{target_port}" ) ;
21952210 let fatal_tx = fatal_tx. clone( ) ;
2211+ let replica = replica. clone( ) ;
21962212 tokio:: spawn( async move {
21972213 let token = match create_forward_session_token(
21982214 & mut client,
@@ -2217,6 +2233,7 @@ pub async fn service_forward_tcp(
22172233 target_port,
22182234 service_id,
22192235 token. clone( ) ,
2236+ replica. as_ref( ) ,
22202237 )
22212238 . await
22222239 {
@@ -2251,18 +2268,19 @@ async fn create_forward_session_token(
22512268 Ok ( response. into_inner ( ) . token )
22522269}
22532270
2271+ /// Confirm the sandbox is still ready, returning the replica that owns it.
22542272async fn fetch_ready_sandbox_for_forward (
22552273 client : & mut crate :: tls:: GrpcClient ,
22562274 name : & str ,
22572275 workspace : & str ,
2258- ) -> Result < Sandbox > {
2276+ ) -> Result < Option < AsciiMetadataValue > > {
22592277 let response = match client
2260- . get_sandbox ( GetSandboxRequest {
2278+ . get_sandbox ( with_owner_hint ( GetSandboxRequest {
22612279 name : name. to_string ( ) ,
22622280 workspace_scope : Some ( openshell_core:: proto:: workspace_selector (
22632281 workspace. to_string ( ) ,
22642282 ) ) ,
2265- } )
2283+ } ) )
22662284 . await
22672285 {
22682286 Ok ( response) => response,
@@ -2274,6 +2292,7 @@ async fn fetch_ready_sandbox_for_forward(
22742292 Err ( status) => return Err ( status) . into_diagnostic ( ) ,
22752293 } ;
22762294
2295+ let replica = owner_replica ( & response) ;
22772296 let sandbox = response
22782297 . into_inner ( )
22792298 . sandbox
@@ -2287,7 +2306,7 @@ async fn fetch_ready_sandbox_for_forward(
22872306 ) ) ;
22882307 }
22892308
2290- Ok ( sandbox )
2309+ Ok ( replica )
22912310}
22922311
22932312#[ derive( Debug ) ]
@@ -2353,11 +2372,9 @@ async fn forward_one_tcp_connection(
23532372 target_port : u16 ,
23542373 service_id : String ,
23552374 authorization_token : String ,
2375+ replica : Option < & AsciiMetadataValue > ,
23562376) -> std:: result:: Result < ( ) , ForwardTcpConnectionError > {
2357- use tokio_stream:: wrappers:: ReceiverStream ;
2358-
2359- let ( tx, rx) = tokio:: sync:: mpsc:: channel :: < TcpForwardFrame > ( 16 ) ;
2360- tx. send ( TcpForwardFrame {
2377+ let init = TcpForwardFrame {
23612378 payload : Some ( openshell_core:: proto:: tcp_forward_frame:: Payload :: Init (
23622379 TcpForwardInit {
23632380 sandbox : sandbox_name,
@@ -2370,11 +2387,10 @@ async fn forward_one_tcp_connection(
23702387 authorization_token,
23712388 } ,
23722389 ) ) ,
2373- } )
2374- . await
2375- . map_err ( |_| ForwardTcpConnectionError :: transient ( "failed to initialize forward stream" ) ) ?;
2390+ } ;
2391+ let ( tx, response) = crate :: ssh:: open_forward_tcp ( client, init, replica) . await ;
23762392
2377- let response = match client . forward_tcp ( ReceiverStream :: new ( rx ) ) . await {
2393+ let response = match response {
23782394 Ok ( response) => response. into_inner ( ) ,
23792395 Err ( status) => {
23802396 let err = ForwardTcpConnectionError :: from_status ( status) ;
@@ -2493,6 +2509,7 @@ async fn sandbox_exec_streaming_grpc(
24932509 stdin_is_terminal : bool ,
24942510 stdin_prefix : Vec < u8 > ,
24952511 stdin_rest : Option < tokio:: sync:: mpsc:: Receiver < std:: io:: Result < Vec < u8 > > > > ,
2512+ replica : Option < & AsciiMetadataValue > ,
24962513) -> Result < i32 > {
24972514 #[ cfg( unix) ]
24982515 use openshell_core:: proto:: ExecSandboxWindowResize ;
@@ -2505,36 +2522,44 @@ async fn sandbox_exec_streaming_grpc(
25052522 ( 0 , 0 )
25062523 } ;
25072524
2508- let ( input_tx, input_rx) = tokio:: sync:: mpsc:: channel :: < ExecSandboxInput > ( 64 ) ;
2509-
2510- // Send the start message with exec metadata.
2511- input_tx
2512- . send ( ExecSandboxInput {
2513- payload : Some ( exec_sandbox_input:: Payload :: Start ( ExecSandboxRequest {
2514- request_id : String :: new ( ) ,
2515- sandbox : sandbox. object_name ( ) . to_string ( ) ,
2516- workspace_scope : Some ( openshell_core:: proto:: workspace_selector (
2517- ( sandbox. object_workspace ( ) ) . to_string ( ) ,
2518- ) ) ,
2519- command : command. to_vec ( ) ,
2520- workdir : workdir. unwrap_or_default ( ) . to_string ( ) ,
2521- environment : environment. clone ( ) ,
2522- no_login_shell,
2523- execution_timeout : proto_execution_timeout ( timeout_seconds) ?,
2524- stdin : Vec :: new ( ) ,
2525- tty,
2526- cols,
2527- rows,
2528- } ) ) ,
2529- } )
2530- . await
2531- . into_diagnostic ( ) ?;
2525+ // The start message carries the exec metadata.
2526+ let start = ExecSandboxInput {
2527+ payload : Some ( exec_sandbox_input:: Payload :: Start ( ExecSandboxRequest {
2528+ request_id : String :: new ( ) ,
2529+ sandbox : sandbox. object_name ( ) . to_string ( ) ,
2530+ workspace_scope : Some ( openshell_core:: proto:: workspace_selector (
2531+ ( sandbox. object_workspace ( ) ) . to_string ( ) ,
2532+ ) ) ,
2533+ command : command. to_vec ( ) ,
2534+ workdir : workdir. unwrap_or_default ( ) . to_string ( ) ,
2535+ environment : environment. clone ( ) ,
2536+ no_login_shell,
2537+ execution_timeout : proto_execution_timeout ( timeout_seconds) ?,
2538+ stdin : Vec :: new ( ) ,
2539+ tty,
2540+ cols,
2541+ rows,
2542+ } ) ) ,
2543+ } ;
25322544
2533- let mut stream = client
2534- . exec_sandbox_interactive ( ReceiverStream :: new ( input_rx) )
2545+ let ( mut input_tx, input_rx) = tokio:: sync:: mpsc:: channel :: < ExecSandboxInput > ( 64 ) ;
2546+ input_tx. send ( start. clone ( ) ) . await . into_diagnostic ( ) ?;
2547+ let mut stream = match client
2548+ . exec_sandbox_interactive ( routed_to ( ReceiverStream :: new ( input_rx) , replica) )
25352549 . await
2536- . into_diagnostic ( ) ?
2537- . into_inner ( ) ;
2550+ {
2551+ Err ( status) if retry_unrouted ( replica, & status) => {
2552+ let ( tx, rx) = tokio:: sync:: mpsc:: channel :: < ExecSandboxInput > ( 64 ) ;
2553+ tx. send ( start) . await . into_diagnostic ( ) ?;
2554+ input_tx = tx;
2555+ client
2556+ . exec_sandbox_interactive ( ReceiverStream :: new ( rx) )
2557+ . await
2558+ }
2559+ response => response,
2560+ }
2561+ . into_diagnostic ( ) ?
2562+ . into_inner ( ) ;
25382563
25392564 // Raw mode is only appropriate for an interactive terminal, not a pipe.
25402565 let raw_guard = if tty && stdin_is_terminal {
0 commit comments