mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
fix: restore owner-scoped gateway startup (#1625)
* fix: restore owner-scoped gateway startup * fix: split gateway owner and sender scope * fix: keep multi-user gateway sender identity * test: cover gateway sender scope regression * test: harden e2e startup teardown race * fix: align gateway owner scope across auth modes
This commit is contained in:
+1
-7
@@ -312,13 +312,7 @@ impl AppBuilder {
|
||||
.create_provider(&self.config.llm.nearai.base_url, self.session.clone());
|
||||
|
||||
// Register memory tools if database is available
|
||||
let workspace_user_id = self
|
||||
.config
|
||||
.channels
|
||||
.gateway
|
||||
.as_ref()
|
||||
.map(|gw| gw.user_id.as_str())
|
||||
.unwrap_or("default");
|
||||
let workspace_user_id = self.config.owner_id.as_str();
|
||||
let workspace = if let Some(ref db) = self.db {
|
||||
let emb_cache_config = EmbeddingCacheConfig {
|
||||
max_entries: self.config.embeddings.cache_size,
|
||||
|
||||
+22
-3
@@ -98,7 +98,8 @@ impl GatewayChannel {
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
default_user_id: config.user_id.clone(),
|
||||
owner_id: config.user_id.clone(),
|
||||
default_sender_id: config.user_id.clone(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(ws::WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
@@ -121,6 +122,22 @@ impl GatewayChannel {
|
||||
}
|
||||
}
|
||||
|
||||
/// Rebind the single-user auth identity to the durable owner scope while
|
||||
/// preserving the configured gateway sender/routing identity.
|
||||
pub fn with_owner_scope(mut self, owner_id: impl Into<String>) -> Self {
|
||||
let owner_id = owner_id.into();
|
||||
let single_user_token = if self.config.user_tokens.is_none() {
|
||||
self.auth.first_token().map(ToOwned::to_owned)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if let Some(token) = single_user_token {
|
||||
self.auth = MultiAuthState::single(token, owner_id.clone());
|
||||
}
|
||||
self.rebuild_state(|s| s.owner_id = owner_id);
|
||||
self
|
||||
}
|
||||
|
||||
/// Create a gateway channel with a pre-built multi-user auth state.
|
||||
pub fn new_multi_auth(config: GatewayConfig, auth: MultiAuthState) -> Self {
|
||||
let state = Arc::new(GatewayState {
|
||||
@@ -137,7 +154,8 @@ impl GatewayChannel {
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
default_user_id: config.user_id.clone(),
|
||||
owner_id: config.user_id.clone(),
|
||||
default_sender_id: config.user_id.clone(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(ws::WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
@@ -177,7 +195,8 @@ impl GatewayChannel {
|
||||
job_manager: self.state.job_manager.clone(),
|
||||
prompt_queue: self.state.prompt_queue.clone(),
|
||||
scheduler: self.state.scheduler.clone(),
|
||||
default_user_id: self.state.default_user_id.clone(),
|
||||
owner_id: self.state.owner_id.clone(),
|
||||
default_sender_id: self.state.default_sender_id.clone(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: self.state.ws_tracker.clone(),
|
||||
llm_provider: self.state.llm_provider.clone(),
|
||||
|
||||
+17
-15
@@ -345,8 +345,10 @@ pub struct GatewayState {
|
||||
pub job_manager: Option<Arc<ContainerJobManager>>,
|
||||
/// Prompt queue for Claude Code follow-up prompts.
|
||||
pub prompt_queue: Option<PromptQueue>,
|
||||
/// Default user ID (fallback for non-request contexts like heartbeat/routines).
|
||||
pub default_user_id: String,
|
||||
/// Durable owner scope for persistence and unauthenticated callback flows.
|
||||
pub owner_id: String,
|
||||
/// Default sender/routing identity for gateway-originated messages.
|
||||
pub default_sender_id: String,
|
||||
/// Shutdown signal sender.
|
||||
pub shutdown_tx: tokio::sync::RwLock<Option<oneshot::Sender<()>>>,
|
||||
/// WebSocket connection tracker.
|
||||
@@ -775,7 +777,7 @@ async fn oauth_callback_handler(
|
||||
error = %error,
|
||||
"OAuth callback received with malformed state"
|
||||
);
|
||||
clear_auth_mode(&state, &state.default_user_id).await;
|
||||
clear_auth_mode(&state, &state.owner_id).await;
|
||||
return oauth_error_page("IronClaw");
|
||||
}
|
||||
};
|
||||
@@ -1136,7 +1138,7 @@ async fn slack_relay_oauth_callback_handler(
|
||||
let state_key = format!("relay:{}:oauth_state", DEFAULT_RELAY_NAME);
|
||||
let stored_state = match ext_mgr
|
||||
.secrets()
|
||||
.get_decrypted(&state.default_user_id, &state_key)
|
||||
.get_decrypted(&state.owner_id, &state_key)
|
||||
.await
|
||||
{
|
||||
Ok(secret) => secret.expose().to_string(),
|
||||
@@ -1160,10 +1162,7 @@ async fn slack_relay_oauth_callback_handler(
|
||||
}
|
||||
|
||||
// Delete the nonce (one-time use)
|
||||
let _ = ext_mgr
|
||||
.secrets()
|
||||
.delete(&state.default_user_id, &state_key)
|
||||
.await;
|
||||
let _ = ext_mgr.secrets().delete(&state.owner_id, &state_key).await;
|
||||
|
||||
let result: Result<(), String> = async {
|
||||
let store = state.store.as_ref().ok_or_else(|| {
|
||||
@@ -1174,16 +1173,12 @@ async fn slack_relay_oauth_callback_handler(
|
||||
// Store team_id in settings
|
||||
let team_id_key = format!("relay:{}:team_id", DEFAULT_RELAY_NAME);
|
||||
let _ = store
|
||||
.set_setting(
|
||||
&state.default_user_id,
|
||||
&team_id_key,
|
||||
&serde_json::json!(team_id),
|
||||
)
|
||||
.set_setting(&state.owner_id, &team_id_key, &serde_json::json!(team_id))
|
||||
.await;
|
||||
|
||||
// Activate the relay channel
|
||||
ext_mgr
|
||||
.activate_stored_relay(DEFAULT_RELAY_NAME, &state.default_user_id)
|
||||
.activate_stored_relay(DEFAULT_RELAY_NAME, &state.owner_id)
|
||||
.await
|
||||
.map_err(|e| format!("Failed to activate relay channel: {}", e))?;
|
||||
|
||||
@@ -1303,6 +1298,9 @@ async fn chat_send_handler(
|
||||
}
|
||||
|
||||
let mut msg = IncomingMessage::new("gateway", &user.user_id, &req.content);
|
||||
if state.owner_id != state.default_sender_id && user.user_id == state.owner_id {
|
||||
msg = msg.with_sender_id(&state.default_sender_id);
|
||||
}
|
||||
// Prefer timezone from JSON body, fall back to X-Timezone header
|
||||
let tz = req
|
||||
.timezone
|
||||
@@ -1404,6 +1402,9 @@ async fn chat_approval_handler(
|
||||
})?;
|
||||
|
||||
let mut msg = IncomingMessage::new("gateway", &user.user_id, content);
|
||||
if state.owner_id != state.default_sender_id && user.user_id == state.owner_id {
|
||||
msg = msg.with_sender_id(&state.default_sender_id);
|
||||
}
|
||||
|
||||
if let Some(ref thread_id) = req.thread_id {
|
||||
msg = msg.with_thread(thread_id);
|
||||
@@ -2976,7 +2977,8 @@ mod tests {
|
||||
store: None,
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
default_user_id: "test".to_string(),
|
||||
owner_id: "test".to_string(),
|
||||
default_sender_id: "test".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: None,
|
||||
llm_provider: None,
|
||||
|
||||
@@ -76,7 +76,8 @@ impl TestGatewayBuilder {
|
||||
store: None,
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
default_user_id: self.user_id,
|
||||
owner_id: self.user_id.clone(),
|
||||
default_sender_id: self.user_id,
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: self.llm_provider,
|
||||
|
||||
@@ -16,6 +16,7 @@ use axum::routing::{delete, get, post};
|
||||
use tower::ServiceExt;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::channels::web::GatewayChannel;
|
||||
use crate::channels::web::auth::{
|
||||
AuthenticatedUser, MultiAuthState, UserIdentity, auth_middleware,
|
||||
};
|
||||
@@ -23,6 +24,7 @@ use crate::channels::web::server::{
|
||||
ActiveConfigSnapshot, GatewayState, PerUserRateLimiter, PromptQueue, RateLimiter, WorkspacePool,
|
||||
};
|
||||
use crate::channels::web::sse::SseManager;
|
||||
use crate::config::GatewayConfig;
|
||||
|
||||
// ── Helpers ────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -64,7 +66,8 @@ fn build_state(
|
||||
store,
|
||||
job_manager: None,
|
||||
prompt_queue,
|
||||
default_user_id: "test".to_string(),
|
||||
owner_id: "test".to_string(),
|
||||
default_sender_id: "test".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: None,
|
||||
llm_provider: None,
|
||||
@@ -82,6 +85,40 @@ fn build_state(
|
||||
})
|
||||
}
|
||||
|
||||
fn gateway_config() -> GatewayConfig {
|
||||
GatewayConfig {
|
||||
host: "127.0.0.1".to_string(),
|
||||
port: 3000,
|
||||
auth_token: Some("gateway-auth".to_string()),
|
||||
user_id: "gateway-sender".to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
memory_layers: Vec::new(),
|
||||
user_tokens: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn with_owner_scope_updates_gateway_owner_scope_in_multi_user_mode() {
|
||||
let mut gateway = GatewayChannel::new(gateway_config());
|
||||
gateway.auth = two_user_auth();
|
||||
gateway.config.user_tokens = Some(HashMap::new());
|
||||
let gateway = gateway.with_owner_scope("owner-scope");
|
||||
|
||||
assert_eq!(gateway.state.owner_id, "owner-scope");
|
||||
assert_eq!(gateway.state.default_sender_id, "gateway-sender");
|
||||
|
||||
let alice = gateway
|
||||
.auth
|
||||
.authenticate("tok-alice")
|
||||
.expect("alice token should remain valid");
|
||||
let bob = gateway
|
||||
.auth
|
||||
.authenticate("tok-bob")
|
||||
.expect("bob token should remain valid");
|
||||
assert_eq!(alice.user_id, "alice");
|
||||
assert_eq!(bob.user_id, "bob");
|
||||
}
|
||||
|
||||
/// Create a libSQL-backed test database in a temporary directory.
|
||||
///
|
||||
/// Returns the database and a `TempDir` guard — the database file is
|
||||
|
||||
@@ -520,7 +520,8 @@ mod tests {
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
default_user_id: "test".to_string(),
|
||||
owner_id: "test".to_string(),
|
||||
default_sender_id: "test".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
|
||||
+5
-7
@@ -312,13 +312,11 @@ impl Config {
|
||||
let tunnel = TunnelConfig::resolve(settings)?;
|
||||
let channels = ChannelsConfig::resolve(settings, &owner_id)?;
|
||||
|
||||
// Resolve workspace config using the gateway user_id for default layers.
|
||||
let workspace_user_id = channels
|
||||
.gateway
|
||||
.as_ref()
|
||||
.map(|gw| gw.user_id.as_str())
|
||||
.unwrap_or("default");
|
||||
let workspace = WorkspaceConfig::resolve(workspace_user_id)?;
|
||||
// Resolve the startup workspace against the durable owner scope. The
|
||||
// gateway may expose a distinct sender identity, but the base runtime
|
||||
// workspace stays owner-scoped and per-user gateway workspaces are
|
||||
// handled separately by WorkspacePool.
|
||||
let workspace = WorkspaceConfig::resolve(&owner_id)?;
|
||||
|
||||
Ok(Self {
|
||||
owner_id: owner_id.clone(),
|
||||
|
||||
@@ -611,6 +611,7 @@ async fn async_main() -> anyhow::Result<()> {
|
||||
} else {
|
||||
GatewayChannel::new(gw_config.clone())
|
||||
};
|
||||
gw = gw.with_owner_scope(config.owner_id.clone());
|
||||
gw = gw.with_llm_provider(Arc::clone(&components.llm));
|
||||
if let Some(ref ws) = components.workspace {
|
||||
gw = gw.with_workspace(Arc::clone(ws));
|
||||
|
||||
+57
-22
@@ -112,6 +112,30 @@ def _reserve_loopback_sockets(count: int) -> list[socket.socket]:
|
||||
sock.close()
|
||||
raise
|
||||
|
||||
async def _stop_process(
|
||||
proc: asyncio.subprocess.Process, *, sig: int | None = None, timeout: float
|
||||
) -> None:
|
||||
"""Signal a subprocess and wait briefly without masking exit races."""
|
||||
if proc.returncode is not None:
|
||||
return
|
||||
|
||||
try:
|
||||
if sig is None:
|
||||
proc.kill()
|
||||
else:
|
||||
proc.send_signal(sig)
|
||||
except ProcessLookupError:
|
||||
try:
|
||||
await asyncio.wait_for(proc.wait(), timeout=timeout)
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
return
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(proc.wait(), timeout=timeout)
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
|
||||
|
||||
def _forward_coverage_env(env: dict[str, str]) -> None:
|
||||
"""Forward cargo-llvm-cov env vars into child processes when present."""
|
||||
@@ -281,35 +305,39 @@ async def ironclaw_server(
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
env=env,
|
||||
)
|
||||
startup_kill_attempted = False
|
||||
base_url = f"http://127.0.0.1:{gateway_port}"
|
||||
try:
|
||||
await wait_for_ready(f"{base_url}/api/health", timeout=60)
|
||||
yield base_url
|
||||
except TimeoutError:
|
||||
# Dump stderr so CI logs show why the server failed to start
|
||||
if proc.returncode is None:
|
||||
startup_kill_attempted = True
|
||||
await _stop_process(proc, timeout=2)
|
||||
returncode = proc.returncode
|
||||
stderr_bytes = b""
|
||||
if proc.stderr:
|
||||
try:
|
||||
stderr_bytes = await asyncio.wait_for(proc.stderr.read(8192), timeout=2)
|
||||
except (asyncio.TimeoutError, Exception):
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
stderr_text = stderr_bytes.decode("utf-8", errors="replace")
|
||||
proc.kill()
|
||||
pytest.fail(
|
||||
f"ironclaw server failed to start on port {gateway_port} "
|
||||
f"(returncode={returncode}).\nstderr:\n{stderr_text}"
|
||||
)
|
||||
finally:
|
||||
if proc.returncode is None:
|
||||
if startup_kill_attempted:
|
||||
await _stop_process(proc, timeout=2)
|
||||
else:
|
||||
# Use SIGINT (not SIGTERM) so tokio's ctrl_c handler triggers a
|
||||
# graceful shutdown. This lets the LLVM coverage runtime run its
|
||||
# atexit handler and flush .profraw files for cargo-llvm-cov.
|
||||
proc.send_signal(signal.SIGINT)
|
||||
try:
|
||||
await asyncio.wait_for(proc.wait(), timeout=10)
|
||||
except asyncio.TimeoutError:
|
||||
proc.kill()
|
||||
await _stop_process(proc, sig=signal.SIGINT, timeout=10)
|
||||
if proc.returncode is None:
|
||||
await _stop_process(proc, timeout=2)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
@@ -376,6 +404,7 @@ async def hosted_oauth_refresh_server(
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
env=env,
|
||||
)
|
||||
startup_kill_attempted = False
|
||||
base_url = f"http://127.0.0.1:{gateway_port}"
|
||||
try:
|
||||
await wait_for_ready(f"{base_url}/api/health", timeout=60)
|
||||
@@ -386,27 +415,29 @@ async def hosted_oauth_refresh_server(
|
||||
"mock_llm_url": mock_llm_server,
|
||||
}
|
||||
except TimeoutError:
|
||||
if proc.returncode is None:
|
||||
startup_kill_attempted = True
|
||||
await _stop_process(proc, timeout=2)
|
||||
returncode = proc.returncode
|
||||
stderr_bytes = b""
|
||||
if proc.stderr:
|
||||
try:
|
||||
stderr_bytes = await asyncio.wait_for(proc.stderr.read(8192), timeout=2)
|
||||
except (asyncio.TimeoutError, Exception):
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
stderr_text = stderr_bytes.decode("utf-8", errors="replace")
|
||||
if proc.returncode is None:
|
||||
proc.kill()
|
||||
pytest.fail(
|
||||
f"hosted oauth refresh server failed to start on port {gateway_port} "
|
||||
f"(returncode={returncode}).\nstderr:\n{stderr_text}"
|
||||
)
|
||||
finally:
|
||||
if proc.returncode is None:
|
||||
proc.send_signal(signal.SIGINT)
|
||||
try:
|
||||
await asyncio.wait_for(proc.wait(), timeout=10)
|
||||
except asyncio.TimeoutError:
|
||||
proc.kill()
|
||||
if startup_kill_attempted:
|
||||
await _stop_process(proc, timeout=2)
|
||||
else:
|
||||
await _stop_process(proc, sig=signal.SIGINT, timeout=10)
|
||||
if proc.returncode is None:
|
||||
await _stop_process(proc, timeout=2)
|
||||
finally:
|
||||
for sock in reserved:
|
||||
if sock.fileno() != -1:
|
||||
@@ -475,6 +506,7 @@ async def http_channel_server_without_secret(
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
env=env,
|
||||
)
|
||||
startup_kill_attempted = False
|
||||
gateway_url = f"http://127.0.0.1:{gateway_port}"
|
||||
http_base_url = f"http://127.0.0.1:{http_port}"
|
||||
try:
|
||||
@@ -483,15 +515,17 @@ async def http_channel_server_without_secret(
|
||||
yield http_base_url
|
||||
except TimeoutError:
|
||||
# Dump stderr so CI logs show why the server failed to start
|
||||
if proc.returncode is None:
|
||||
startup_kill_attempted = True
|
||||
await _stop_process(proc, timeout=2)
|
||||
returncode = proc.returncode
|
||||
stderr_bytes = b""
|
||||
if proc.stderr:
|
||||
try:
|
||||
stderr_bytes = await asyncio.wait_for(proc.stderr.read(8192), timeout=2)
|
||||
except (asyncio.TimeoutError, Exception):
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
stderr_text = stderr_bytes.decode("utf-8", errors="replace")
|
||||
proc.kill()
|
||||
pytest.fail(
|
||||
f"ironclaw server without webhook secret failed to start on ports "
|
||||
f"gateway={gateway_port}, http={http_port} "
|
||||
@@ -499,14 +533,15 @@ async def http_channel_server_without_secret(
|
||||
)
|
||||
finally:
|
||||
if proc.returncode is None:
|
||||
if startup_kill_attempted:
|
||||
await _stop_process(proc, timeout=2)
|
||||
else:
|
||||
# Use SIGINT (not SIGTERM) so tokio's ctrl_c handler triggers a
|
||||
# graceful shutdown. This lets the LLVM coverage runtime run its
|
||||
# atexit handler and flush .profraw files for cargo-llvm-cov.
|
||||
proc.send_signal(signal.SIGINT)
|
||||
try:
|
||||
await asyncio.wait_for(proc.wait(), timeout=10)
|
||||
except asyncio.TimeoutError:
|
||||
proc.kill()
|
||||
await _stop_process(proc, sig=signal.SIGINT, timeout=10)
|
||||
if proc.returncode is None:
|
||||
await _stop_process(proc, timeout=2)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
|
||||
@@ -19,10 +19,13 @@ use axum::middleware;
|
||||
use axum::routing::{get, post};
|
||||
use tower::ServiceExt;
|
||||
|
||||
use ironclaw::channels::IncomingMessage;
|
||||
use ironclaw::channels::web::auth::{
|
||||
AuthenticatedUser, MultiAuthState, UserIdentity, auth_middleware,
|
||||
};
|
||||
use ironclaw::channels::web::server::{GatewayState, PerUserRateLimiter, RateLimiter};
|
||||
use ironclaw::channels::web::server::{
|
||||
GatewayState, PerUserRateLimiter, RateLimiter, start_server,
|
||||
};
|
||||
use ironclaw::channels::web::sse::SseManager;
|
||||
use ironclaw::channels::web::test_helpers::TestGatewayBuilder;
|
||||
use ironclaw::channels::web::ws::WsConnectionTracker;
|
||||
@@ -37,6 +40,9 @@ const ALICE_TOKEN: &str = "tok-alice-secret";
|
||||
const BOB_TOKEN: &str = "tok-bob-secret";
|
||||
const ALICE_USER_ID: &str = "alice";
|
||||
const BOB_USER_ID: &str = "bob";
|
||||
const OWNER_TOKEN: &str = "tok-owner-secret";
|
||||
const OWNER_SCOPE_ID: &str = "owner-scope";
|
||||
const GATEWAY_SENDER_ID: &str = "gateway-sender";
|
||||
|
||||
/// Build a MultiAuthState with two users.
|
||||
fn two_user_auth() -> MultiAuthState {
|
||||
@@ -537,7 +543,8 @@ fn gateway_state_has_multi_tenant_fields() {
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
default_user_id: "fallback".to_string(), // Multi-tenant: renamed from user_id
|
||||
owner_id: "fallback".to_string(),
|
||||
default_sender_id: "fallback".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
@@ -553,7 +560,8 @@ fn gateway_state_has_multi_tenant_fields() {
|
||||
active_config: Default::default(),
|
||||
};
|
||||
|
||||
assert_eq!(state.default_user_id, "fallback");
|
||||
assert_eq!(state.owner_id, "fallback");
|
||||
assert_eq!(state.default_sender_id, "fallback");
|
||||
assert!(state.workspace_pool.is_none());
|
||||
}
|
||||
|
||||
@@ -572,6 +580,69 @@ async fn start_multi_user_server() -> (SocketAddr, Arc<GatewayState>) {
|
||||
.expect("Failed to start multi-user test server")
|
||||
}
|
||||
|
||||
async fn start_owner_scoped_sender_server() -> (
|
||||
SocketAddr,
|
||||
Arc<GatewayState>,
|
||||
tokio::sync::mpsc::Receiver<IncomingMessage>,
|
||||
) {
|
||||
let (agent_tx, agent_rx) = tokio::sync::mpsc::channel(64);
|
||||
|
||||
let mut tokens = HashMap::new();
|
||||
tokens.insert(
|
||||
OWNER_TOKEN.to_string(),
|
||||
UserIdentity {
|
||||
user_id: OWNER_SCOPE_ID.to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
},
|
||||
);
|
||||
tokens.insert(
|
||||
BOB_TOKEN.to_string(),
|
||||
UserIdentity {
|
||||
user_id: BOB_USER_ID.to_string(),
|
||||
workspace_read_scopes: Vec::new(),
|
||||
},
|
||||
);
|
||||
|
||||
let state = Arc::new(GatewayState {
|
||||
msg_tx: tokio::sync::RwLock::new(Some(agent_tx)),
|
||||
sse: Arc::new(SseManager::new()),
|
||||
workspace: None,
|
||||
workspace_pool: None,
|
||||
session_manager: None,
|
||||
log_broadcaster: None,
|
||||
log_level_handle: None,
|
||||
extension_manager: None,
|
||||
tool_registry: None,
|
||||
store: None,
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
owner_id: OWNER_SCOPE_ID.to_string(),
|
||||
default_sender_id: GATEWAY_SENDER_ID.to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
skill_registry: None,
|
||||
skill_catalog: None,
|
||||
chat_rate_limiter: PerUserRateLimiter::new(30, 60),
|
||||
oauth_rate_limiter: RateLimiter::new(10, 60),
|
||||
webhook_rate_limiter: RateLimiter::new(10, 60),
|
||||
registry_entries: Vec::new(),
|
||||
cost_guard: None,
|
||||
routine_engine: Arc::new(tokio::sync::RwLock::new(None)),
|
||||
startup_time: std::time::Instant::now(),
|
||||
active_config: Default::default(),
|
||||
});
|
||||
|
||||
let auth = MultiAuthState::multi(tokens);
|
||||
let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
|
||||
let bound = start_server(addr, state.clone(), auth)
|
||||
.await
|
||||
.expect("Failed to start owner-scoped sender test server");
|
||||
|
||||
(bound, state, agent_rx)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn full_server_alice_can_access_protected_endpoint() {
|
||||
let (addr, _state) = start_multi_user_server().await;
|
||||
@@ -677,6 +748,49 @@ async fn full_server_chat_send_accepted_for_alice() {
|
||||
assert_eq!(msg.channel, "gateway");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn full_server_chat_send_rewrites_sender_only_for_owner_scope_rebind() {
|
||||
let (addr, _state, mut agent_rx) = start_owner_scoped_sender_server().await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
let owner_resp = client
|
||||
.post(format!("http://{}/api/chat/send", addr))
|
||||
.header("Authorization", format!("Bearer {}", OWNER_TOKEN))
|
||||
.header("Content-Type", "application/json")
|
||||
.body(r#"{"content":"hello from owner"}"#)
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(owner_resp.status(), 202);
|
||||
|
||||
let owner_msg = tokio::time::timeout(Duration::from_secs(2), agent_rx.recv())
|
||||
.await
|
||||
.expect("Timed out waiting for owner message")
|
||||
.expect("Agent channel closed");
|
||||
assert_eq!(owner_msg.user_id, OWNER_SCOPE_ID);
|
||||
assert_eq!(owner_msg.sender_id, GATEWAY_SENDER_ID);
|
||||
assert_eq!(owner_msg.content, "hello from owner");
|
||||
|
||||
let other_resp = client
|
||||
.post(format!("http://{}/api/chat/send", addr))
|
||||
.header("Authorization", format!("Bearer {}", BOB_TOKEN))
|
||||
.header("Content-Type", "application/json")
|
||||
.body(r#"{"content":"hello from bob"}"#)
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(other_resp.status(), 202);
|
||||
|
||||
let other_msg = tokio::time::timeout(Duration::from_secs(2), agent_rx.recv())
|
||||
.await
|
||||
.expect("Timed out waiting for non-owner message")
|
||||
.expect("Agent channel closed");
|
||||
assert_eq!(other_msg.user_id, BOB_USER_ID);
|
||||
assert_eq!(other_msg.sender_id, BOB_USER_ID);
|
||||
assert_eq!(other_msg.content, "hello from bob");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn full_server_chat_send_rejected_without_auth() {
|
||||
let (addr, _state) = start_multi_user_server().await;
|
||||
@@ -888,7 +1002,8 @@ async fn start_multi_user_server_with_db() -> (
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
default_user_id: ALICE_USER_ID.to_string(),
|
||||
owner_id: ALICE_USER_ID.to_string(),
|
||||
default_sender_id: ALICE_USER_ID.to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
|
||||
@@ -203,7 +203,8 @@ async fn start_test_server_with_provider(
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
default_user_id: "test-user".to_string(),
|
||||
owner_id: "test-user".to_string(),
|
||||
default_sender_id: "test-user".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: Some(llm_provider),
|
||||
@@ -701,7 +702,8 @@ async fn test_no_llm_provider_returns_503() {
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
default_user_id: "test-user".to_string(),
|
||||
owner_id: "test-user".to_string(),
|
||||
default_sender_id: "test-user".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None, // No LLM!
|
||||
|
||||
@@ -226,7 +226,8 @@ impl GatewayWorkflowHarness {
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: Some(scheduler_slot.clone()),
|
||||
default_user_id: user_id.clone(),
|
||||
owner_id: user_id.clone(),
|
||||
default_sender_id: user_id.clone(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: Some(Arc::clone(&components.llm)),
|
||||
|
||||
@@ -51,7 +51,8 @@ async fn start_test_server() -> (
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: None,
|
||||
default_user_id: "test-user".to_string(),
|
||||
owner_id: "test-user".to_string(),
|
||||
default_sender_id: "test-user".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: None,
|
||||
|
||||
Reference in New Issue
Block a user