changes

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