feat: generalized client→Space/FeatureSet mappings + lock-confine (P2/3) (#202)

Signed-off-by: Mohammod Al Amin Ashik <maa.ashik00@gmail.com>

Author: its-mash

apps/desktop/src-tauri/src/commands/oauth.rs

@@ -37,6 +37,7 @@ use tracing::{debug, error, info, warn};
 use url::Url;
 
 use super::gateway::GatewayAppState;
+use crate::state::AppState;
 
 // ============================================================================
 // Deep Link Handling
@@ -900,6 +901,7 @@ pub async fn update_oauth_client(
 #[tauri::command]
 pub async fn delete_oauth_client(
     gateway_state: State<'_, Arc<RwLock<GatewayAppState>>>,
+    app: State<'_, AppState>,
     client_id: String,
 ) -> Result<(), String> {
     let app_state = gateway_state.read().await;
@@ -921,6 +923,23 @@ pub async fn delete_oauth_client(
 
     info!("[OAuth] Deleted client: {}", client_id);
 
+    // Best-effort: remove the auto-mapped clientId id-binding so a deleted
+    // client doesn't leave an orphan "<client_id> → Starter" mapping behind in
+    // the Mapping tab. Only API-key clients have such a binding; for DCR
+    // clients this is a no-op.
+    if let Ok(Some(b)) = app
+        .workspace_binding_repository
+        .find_by_id_key(&client_id)
+        .await
+    {
+        if let Err(e) = app.workspace_binding_repository.delete(&b.id).await {
+            warn!(
+                "[OAuth] failed to remove clientId mapping for {}: {}",
+                client_id, e
+            );
+        }
+    }
+
     // Emit domain event
     state.emit_domain_event(mcpmux_core::DomainEvent::ClientDeleted { client_id });
 
@@ -943,6 +962,7 @@ pub async fn delete_oauth_client(
 pub struct RegisteredApiKeyClient {
     pub client_id: String,
     pub client_name: String,
+    pub locked_space_id: Option<String>,
     pub api_key: String,
     pub key_prefix: String,
 }
@@ -973,12 +993,14 @@ fn generate_api_key() -> (String, String, String) {
     (key_id, plaintext, key_prefix)
 }
 
-/// Register a new pre-approved client authenticated by an API key. The returned
-/// `api_key` is shown once and never stored.
+/// Register a new pre-approved client authenticated by an API key, optionally
+/// locked to a Space. The returned `api_key` is shown once and never stored.
 #[tauri::command]
 pub async fn register_api_key_client(
     gateway_state: State<'_, Arc<RwLock<GatewayAppState>>>,
+    app: State<'_, AppState>,
     name: String,
+    locked_space_id: Option<String>,
 ) -> Result<RegisteredApiKeyClient, String> {
     let app_state = gateway_state.read().await;
     let Some(ref gw_state) = app_state.gateway_state else {
@@ -1024,6 +1046,12 @@ pub async fn register_api_key_client(
         .await
         .map_err(|e| format!("Failed to create client: {}", e))?;
 
+    if let Some(ref space) = locked_space_id {
+        repo.set_locked_space(&client_id, Some(space))
+            .await
+            .map_err(|e| format!("Failed to lock client to space: {}", e))?;
+    }
+
     let (key_id, plaintext, key_prefix) = generate_api_key();
     repo.create_api_key(&key_id, &client_id, &plaintext, &key_prefix, None, None)
         .await
@@ -1034,14 +1062,61 @@ pub async fn register_api_key_client(
         trimmed, client_id
     );
 
+    // Best-effort: auto-create a clientId-keyed mapping → the (locked or
+    // default) Space's Starter, so the client routes sensibly out of the box
+    // and the mapping is visible + editable in the Mapping tab. A failure here
+    // must not undo the registration — without an explicit mapping the resolver
+    // still falls back to the default Starter.
+    if let Err(e) = auto_map_api_key_client(&app, &client_id, locked_space_id.as_deref()).await {
+        warn!(
+            "[OAuth] auto-map for {} failed (non-fatal): {}",
+            client_id, e
+        );
+    }
+
     Ok(RegisteredApiKeyClient {
         client_id,
         client_name: trimmed.to_string(),
+        locked_space_id,
         api_key: plaintext,
         key_prefix,
     })
 }
 
+/// Auto-create a clientId-keyed `id` mapping pointing at the (locked or
+/// default) Space's Starter FeatureSet, so a freshly-registered API-key client
+/// routes somewhere sensible by default and the operator can retarget it from
+/// the Mapping tab.
+async fn auto_map_api_key_client(
+    app: &AppState,
+    client_id: &str,
+    locked_space_id: Option<&str>,
+) -> Result<(), String> {
+    let space_id = match locked_space_id {
+        Some(s) => uuid::Uuid::parse_str(s).map_err(|e| e.to_string())?,
+        None => {
+            app.space_service
+                .get_default()
+                .await
+                .map_err(|e| e.to_string())?
+                .ok_or("no default Space configured")?
+                .id
+        }
+    };
+    let starter = app
+        .feature_set_repository
+        .get_starter_for_space(&space_id.to_string())
+        .await
+        .map_err(|e| e.to_string())?
+        .ok_or("Space has no Starter FeatureSet")?;
+    let binding =
+        mcpmux_core::WorkspaceBinding::new_id(client_id.to_string(), space_id, vec![starter.id]);
+    app.workspace_binding_repository
+        .create(&binding)
+        .await
+        .map_err(|e| e.to_string())
+}
+
 /// Issue an additional API key for an existing client (rotation). Returns the
 /// new key plaintext once.
 #[tauri::command]
@@ -1079,9 +1154,15 @@ pub async fn create_client_api_key(
     .await
     .map_err(|e| format!("Failed to create API key: {}", e))?;
 
+    let locked_space_id = repo
+        .get_locked_space(&client_id)
+        .await
+        .map_err(|e| format!("Failed to read client: {}", e))?;
+
     Ok(RegisteredApiKeyClient {
         client_id,
         client_name: client.client_name,
+        locked_space_id,
         api_key: plaintext,
         key_prefix,
     })

apps/desktop/src-tauri/src/commands/workspace_binding.rs

@@ -8,8 +8,8 @@ use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 use mcpmux_core::{
-    validate_workspace_root as validate_root, DomainEvent, FeatureSet, FeatureSetType, MemberMode,
-    MemberType, ServerFeature, WorkspaceBinding, WorkspaceRootValidation,
+    validate_workspace_root as validate_root, BindingType, DomainEvent, FeatureSet, FeatureSetType,
+    MemberMode, MemberType, ServerFeature, WorkspaceBinding, WorkspaceRootValidation,
 };
 use serde::{Deserialize, Serialize};
 use tauri::State;
@@ -54,6 +54,8 @@ async fn emit_binding_changed(
 pub struct WorkspaceBindingDto {
     pub id: String,
     pub workspace_root: String,
+    /// `path` (folder, normalized) or `id` (arbitrary exact-match key).
+    pub binding_type: String,
     pub space_id: String,
     pub feature_set_ids: Vec<String>,
     pub created_at: String,
@@ -65,6 +67,7 @@ impl From<WorkspaceBinding> for WorkspaceBindingDto {
         Self {
             id: b.id.to_string(),
             workspace_root: b.workspace_root,
+            binding_type: b.binding_type.as_str().to_string(),
             space_id: b.space_id.to_string(),
             feature_set_ids: b.feature_set_ids,
             created_at: b.created_at.to_rfc3339(),
@@ -83,6 +86,10 @@ pub struct WorkspaceBindingInput {
     pub workspace_root: String,
     pub space_id: String,
     pub feature_set_ids: Vec<String>,
+    /// `path` (default — folder, normalized + validated) or `id` (arbitrary
+    /// exact-match key, taken verbatim). Optional for backward compatibility.
+    #[serde(default)]
+    pub binding_type: Option<String>,
 }
 
 fn parse_space_id(input: &WorkspaceBindingInput) -> Result<Uuid, String> {
@@ -229,6 +236,26 @@ fn normalize_and_validate(raw: &str) -> Result<String, String> {
     }
 }
 
+/// Resolve the storage key + type from the input. `path` bindings are
+/// normalized + validated (rejecting relative paths, filesystem roots, …);
+/// `id` bindings take the raw string verbatim (any non-empty label a headless
+/// client sends in `X-Mcpmux-Workspace`, e.g. a client id or machine name).
+fn resolve_key_and_type(input: &WorkspaceBindingInput) -> Result<(String, BindingType), String> {
+    match input.binding_type.as_deref() {
+        Some("id") => {
+            let key = input.workspace_root.trim();
+            if key.is_empty() {
+                return Err("Mapping id cannot be empty".into());
+            }
+            Ok((key.to_string(), BindingType::Id))
+        }
+        _ => Ok((
+            normalize_and_validate(&input.workspace_root)?,
+            BindingType::Path,
+        )),
+    }
+}
+
 /// Create a binding. Path is normalized + validated server-side so the UI
 /// can pass raw input (Windows paths, file:// URIs, trailing slashes).
 #[tauri::command]
@@ -239,23 +266,26 @@ pub async fn create_workspace_binding(
 ) -> Result<WorkspaceBindingDto, String> {
     let space_id = parse_space_id(&input)?;
     let feature_set_ids = validate_fs_list(&input)?;
-    let normalized = normalize_and_validate(&input.workspace_root)?;
+    let (key, binding_type) = resolve_key_and_type(&input)?;
 
-    // Reject a duplicate folder up front with a readable message. The schema
+    // Reject a duplicate key up front with a readable message. The schema
     // already enforces `UNIQUE(workspace_root)`, but that surfaces an opaque
     // SQLite constraint error — this gives the UI something a user can act on.
     let existing = state
         .workspace_binding_repository
         .list()
         .await
         .map_err(|e| e.to_string())?;
-    if existing.iter().any(|b| b.workspace_root == normalized) {
+    if existing.iter().any(|b| b.workspace_root == key) {
         return Err(format!(
-            "A mapping already exists for {normalized}. Edit the existing mapping instead of adding a second one."
+            "A mapping already exists for {key}. Edit the existing mapping instead of adding a second one."
         ));
     }
 
-    let binding = WorkspaceBinding::new_multi(normalized.clone(), space_id, feature_set_ids);
+    let binding = match binding_type {
+        BindingType::Id => WorkspaceBinding::new_id(key.clone(), space_id, feature_set_ids),
+        BindingType::Path => WorkspaceBinding::new_multi(key.clone(), space_id, feature_set_ids),
+    };
 
     state
         .workspace_binding_repository
@@ -292,9 +322,9 @@ pub async fn update_workspace_binding(
     let id_uuid = Uuid::parse_str(&id).map_err(|e| e.to_string())?;
     let space_id = parse_space_id(&input)?;
     let feature_set_ids = validate_fs_list(&input)?;
-    let normalized = normalize_and_validate(&input.workspace_root)?;
+    let (key, binding_type) = resolve_key_and_type(&input)?;
 
-    // If the edit moved the folder onto a path another mapping already owns,
+    // If the edit moved the mapping onto a key another mapping already owns,
     // reject with a readable message rather than tripping the DB UNIQUE
     // constraint. Exclude this binding's own row.
     let all = state
@@ -304,10 +334,10 @@ pub async fn update_workspace_binding(
         .map_err(|e| e.to_string())?;
     if all
         .iter()
-        .any(|b| b.id != id_uuid && b.workspace_root == normalized)
+        .any(|b| b.id != id_uuid && b.workspace_root == key)
     {
         return Err(format!(
-            "Another mapping already uses {normalized}. Pick a different folder."
+            "Another mapping already uses {key}. Pick a different key."
         ));
     }
 
@@ -321,7 +351,8 @@ pub async fn update_workspace_binding(
 
     let updated = WorkspaceBinding {
         id: existing.id,
-        workspace_root: normalized,
+        workspace_root: key,
+        binding_type,
         space_id,
         feature_set_ids,
         created_at: existing.created_at,

apps/desktop/src/features/clients/RegisterApiKeyClientModal.tsx

@@ -10,10 +10,11 @@
  * never display it again — if lost, revoke it and issue a new one.
  */
 
-import { useState } from 'react';
-import { AlertTriangle, Check, Copy, KeyRound, Loader2, ShieldCheck, X } from 'lucide-react';
+import { useEffect, useState } from 'react';
+import { AlertTriangle, Check, Copy, KeyRound, Loader2, Lock, ShieldCheck, X } from 'lucide-react';
 import { Button, Card, CardContent, CardDescription, CardHeader, CardTitle } from '@mcpmux/ui';
 import { registerApiKeyClient, type RegisteredApiKeyClient } from '@/lib/api/gateway';
+import { listSpaces, type Space } from '@/lib/api/spaces';
 
 interface RegisterApiKeyClientModalProps {
   onClose: () => void;
@@ -26,11 +27,23 @@ export function RegisterApiKeyClientModal({
   onRegistered,
 }: RegisterApiKeyClientModalProps) {
   const [name, setName] = useState('');
+  const [lockedSpaceId, setLockedSpaceId] = useState('');
+  const [spaces, setSpaces] = useState<Space[]>([]);
   const [isSubmitting, setIsSubmitting] = useState(false);
   const [error, setError] = useState<string | null>(null);
   const [result, setResult] = useState<RegisteredApiKeyClient | null>(null);
   const [copied, setCopied] = useState(false);
 
+  useEffect(() => {
+    listSpaces()
+      .then(setSpaces)
+      .catch(() => setSpaces([]));
+  }, []);
+
+  const lockedSpaceName = result?.lockedSpaceId
+    ? (spaces.find((s) => s.id === result.lockedSpaceId)?.name ?? 'a Space')
+    : null;
+
   const handleGenerate = async () => {
     const trimmed = name.trim();
     if (!trimmed) {
@@ -40,7 +53,7 @@ export function RegisterApiKeyClientModal({
     setIsSubmitting(true);
     setError(null);
     try {
-      const client = await registerApiKeyClient(trimmed);
+      const client = await registerApiKeyClient(trimmed, lockedSpaceId || null);
       setResult(client);
     } catch (e) {
       setError(e instanceof Error ? e.message : String(e));
@@ -130,6 +143,13 @@ export function RegisterApiKeyClientModal({
                 <code className="block break-all font-mono text-xs text-[rgb(var(--text))]">
                   Authorization: Bearer {result.keyPrefix}…
                 </code>
+                {lockedSpaceName && (
+                  <p className="mt-2 flex items-center gap-1.5 text-xs text-[rgb(var(--muted))]">
+                    <Lock className="h-3.5 w-3.5" />
+                    Locked to <span className="font-medium">{lockedSpaceName}</span> — this key can
+                    only ever reach that Space.
+                  </p>
+                )}
               </div>
 
               <div className="flex justify-end">
@@ -159,6 +179,30 @@ export function RegisterApiKeyClientModal({
                 />
               </div>
 
+              <div>
+                <label htmlFor="api-key-lock-space" className="mb-1.5 block text-sm font-medium">
+                  Lock to a Space <span className="text-[rgb(var(--muted))]">(optional)</span>
+                </label>
+                <select
+                  id="api-key-lock-space"
+                  data-testid="register-api-key-lock-space"
+                  value={lockedSpaceId}
+                  onChange={(e) => setLockedSpaceId(e.target.value)}
+                  className="w-full rounded-xl border border-[rgb(var(--border))] bg-[rgb(var(--surface))] px-3.5 py-2.5 text-sm transition-all focus:border-[rgb(var(--accent))] focus:outline-none focus:ring-2 focus:ring-[rgb(var(--accent))]/40"
+                >
+                  <option value="">No lock — route by mapping (any Space)</option>
+                  {spaces.map((s) => (
+                    <option key={s.id} value={s.id}>
+                      {s.name}
+                    </option>
+                  ))}
+                </select>
+                <p className="mt-1.5 text-xs text-[rgb(var(--muted))]">
+                  Locking confines this client to one Space — a leaked key can never reach the
+                  others. Leave unlocked to route it later from the Workspaces tab.
+                </p>
+              </div>
+
               <div className="flex items-start gap-3 rounded-xl border border-[rgb(var(--border-subtle))] bg-[rgb(var(--surface))] p-3.5">
                 <ShieldCheck className="mt-0.5 h-5 w-5 flex-shrink-0 text-[rgb(var(--accent))]" />
                 <p className="text-xs text-[rgb(var(--muted))]">