Scaling Phase 1: multi-tenant onboarding + replica-safe coordination
Decouples "many users" + "many server replicas" from "many machines" so the platform is tenant-isolated and horizontally safe on the current single node. - Per-signup workspaces (cm-auth): a new hosted-identity sign-in provisions and owns its own workspace instead of joining the first. Config-gated by auth.per_signup_workspace (default off); concurrent first-logins serialized by a per-subject advisory lock so no duplicate workspaces. - Terminal tickets in Postgres (migration 0016, hashed, single-use): any replica can redeem a ticket minted by another. Drops the in-process ticket map. - Container registry in Postgres (migration 0017, agent_containers): Terminal and Sandbox managers resolve an agent's container through a shared registry, so a 2nd replica reuses it instead of spawning a duplicate. node_id recorded as 'local' (Phase 2 hook). Boot reconcile removes only true orphans, so terminals now survive a redeploy (tmux sessions resume). - Per-workspace quotas (cm-api/quota.rs): plan-tier caps on agents + live containers, enforced at agent create + terminal spin-up (reconnects allowed), returned as HTTP 402. New GET /api/quota surfaces usage vs limits. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
This commit is contained in:
Omar Sobh
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
f8f2b65e1f
commit
e9ce368ec1
@@ -5,9 +5,7 @@
|
||||
//! ticket, then opens the WS with `?ticket=`. The WS itself is routed straight
|
||||
//! to this server by the edge (Traefik), bypassing the HTTP-only Next proxy.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, Instant};
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
|
||||
use axum::extract::{Path, Query, State};
|
||||
@@ -19,56 +17,67 @@ use cm_domain::{AgentId, WorkspaceId};
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
use sqlx::{PgPool, Row};
|
||||
use tokio::io::AsyncWriteExt;
|
||||
|
||||
use crate::routes::claws::workspace_agent;
|
||||
use crate::{ApiError, AppState, Authed};
|
||||
|
||||
const TICKET_TTL: Duration = Duration::from_secs(30);
|
||||
/// Single-use WS tickets live in Postgres (hashed), so any server replica can
|
||||
/// redeem one — not just the instance that minted it. A ticket is deleted on
|
||||
/// redeem; expired rows are swept opportunistically. TTL is 30s (in the SQL).
|
||||
fn hash_ticket(token: &str) -> String {
|
||||
use sha2::{Digest, Sha256};
|
||||
let mut h = Sha256::new();
|
||||
h.update(token.as_bytes());
|
||||
format!("{:x}", h.finalize())
|
||||
}
|
||||
|
||||
/// In-memory single-use WS tickets (token → agent + greeting label + expiry).
|
||||
/// Tiny + ephemeral; a ticket is consumed on redeem, expired ones swept lazily.
|
||||
struct Ticket {
|
||||
/// Mint a short-lived single-use ticket; returns the raw token.
|
||||
async fn issue_ticket(
|
||||
pool: &PgPool,
|
||||
agent_id: AgentId,
|
||||
workspace_id: WorkspaceId,
|
||||
/// Display name to greet the user with in the terminal MOTD.
|
||||
label: String,
|
||||
expires: Instant,
|
||||
label: &str,
|
||||
) -> Result<String, sqlx::Error> {
|
||||
let token = uuid::Uuid::new_v4().simple().to_string();
|
||||
let _ = sqlx::query("DELETE FROM terminal_ws_tickets WHERE expires_at < now()")
|
||||
.execute(pool)
|
||||
.await;
|
||||
sqlx::query(
|
||||
"INSERT INTO terminal_ws_tickets (token_hash, agent_id, workspace_id, label, expires_at)
|
||||
VALUES ($1, $2, $3, $4, now() + interval '30 seconds')",
|
||||
)
|
||||
.bind(hash_ticket(&token))
|
||||
.bind(agent_id.as_uuid())
|
||||
.bind(workspace_id.as_uuid())
|
||||
.bind(label)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(token)
|
||||
}
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct TerminalTickets {
|
||||
inner: Arc<Mutex<HashMap<String, Ticket>>>,
|
||||
}
|
||||
|
||||
impl TerminalTickets {
|
||||
pub fn issue(&self, agent_id: AgentId, workspace_id: WorkspaceId, label: String) -> String {
|
||||
let token = uuid::Uuid::new_v4().simple().to_string();
|
||||
let mut g = self.inner.lock().expect("tickets lock");
|
||||
let now = Instant::now();
|
||||
g.retain(|_, t| t.expires > now);
|
||||
g.insert(
|
||||
token.clone(),
|
||||
Ticket {
|
||||
agent_id,
|
||||
workspace_id,
|
||||
label,
|
||||
expires: now + TICKET_TTL,
|
||||
},
|
||||
);
|
||||
token
|
||||
}
|
||||
|
||||
/// Consume a ticket; returns its (workspace, greeting label) iff valid for this agent.
|
||||
pub fn redeem(&self, token: &str, agent_id: AgentId) -> Option<(WorkspaceId, String)> {
|
||||
let now = Instant::now();
|
||||
let mut g = self.inner.lock().expect("tickets lock");
|
||||
g.retain(|_, t| t.expires > now);
|
||||
match g.remove(token) {
|
||||
Some(t) if t.agent_id == agent_id && t.expires > now => Some((t.workspace_id, t.label)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
/// Consume a ticket; returns (workspace, greeting label) iff valid for this agent.
|
||||
async fn redeem_ticket(
|
||||
pool: &PgPool,
|
||||
token: &str,
|
||||
agent_id: AgentId,
|
||||
) -> Option<(WorkspaceId, String)> {
|
||||
let row = sqlx::query(
|
||||
"DELETE FROM terminal_ws_tickets
|
||||
WHERE token_hash = $1 AND agent_id = $2 AND expires_at > now()
|
||||
RETURNING workspace_id, label",
|
||||
)
|
||||
.bind(hash_ticket(token))
|
||||
.bind(agent_id.as_uuid())
|
||||
.fetch_optional(pool)
|
||||
.await
|
||||
.ok()
|
||||
.flatten()?;
|
||||
Some((
|
||||
WorkspaceId::from(row.get::<uuid::Uuid, _>("workspace_id")),
|
||||
row.get::<String, _>("label"),
|
||||
))
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
@@ -84,6 +93,15 @@ pub async fn ticket(
|
||||
) -> Result<Json<TicketResponse>, ApiError> {
|
||||
// Tenant isolation: a foreign agent looks non-existent.
|
||||
workspace_agent(&state, &user, agent_id).await?;
|
||||
// Quota: gate only when this would spin up a NEW container (a reconnect to an
|
||||
// already-running terminal is always allowed).
|
||||
let existing = cm_db::repo::agent_containers::get(&state.pool, agent_id, "terminal")
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
if existing.is_none() {
|
||||
crate::quota::enforce_new_container(&state, user.workspace_id).await?;
|
||||
}
|
||||
// Greet by display name, falling back to the email local-part.
|
||||
let label = match cm_db::repo::users::get(&state.pool, user.user_id).await {
|
||||
Ok(u) if !u.display_name.trim().is_empty() => u.display_name.trim().to_string(),
|
||||
@@ -96,9 +114,7 @@ pub async fn ticket(
|
||||
.to_string(),
|
||||
Err(_) => "there".to_string(),
|
||||
};
|
||||
let ticket = state
|
||||
.terminal_tickets
|
||||
.issue(agent_id, user.workspace_id, label);
|
||||
let ticket = issue_ticket(&state.pool, agent_id, user.workspace_id, &label).await?;
|
||||
cm_db::repo::audit::append(
|
||||
&state.pool,
|
||||
user.workspace_id,
|
||||
@@ -200,7 +216,7 @@ pub async fn ws(
|
||||
Query(q): Query<WsQuery>,
|
||||
upgrade: WebSocketUpgrade,
|
||||
) -> Response {
|
||||
let (workspace_id, label) = match state.terminal_tickets.redeem(&q.ticket, agent_id) {
|
||||
let (workspace_id, label) = match redeem_ticket(&state.pool, &q.ticket, agent_id).await {
|
||||
Some(v) => v,
|
||||
None => return (StatusCode::UNAUTHORIZED, "invalid or expired ticket").into_response(),
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user