load token

wai_service.py

@@ -1,5 +1,5 @@
 # --- wai_service.py (final handler) ---------------------------------
-import base64, sys, os, torch
+import base64, sys, os, torch, tempfile  
 sys.path.append(os.path.join(os.path.dirname(__file__), "code"))
 
 from code.config import set_config
@@ -34,6 +34,17 @@ def _sanitize(cfg):
 
 
 def handler(payload: dict) -> str:
+
+    tmp_token_path = None                           # ➋
+    if "token" in payload:
+        tmp_token_path = tempfile.NamedTemporaryFile(delete=False, dir=".", prefix="TOKEN_", mode="w")
+        tmp_token_path.write(payload["token"])
+        tmp_token_path.close()
+        os.environ["TOKEN_FILE"] = tmp_token_path.name  # optional: let you debug
+
+        # make config.py look for our temp file
+        os.symlink(tmp_token_path.name, "TOKEN")
+        
     # 1️⃣  Build fake argv *only* from recognised keys
     cli_argv = [sys.argv[0]]
     for k in KNOWN_CLI_KEYS & payload.keys():
@@ -45,6 +56,9 @@ def handler(payload: dict) -> str:
         cfg = set_config()           # EasyDict with YAML + CLI-compatible fields
     finally:
         sys.argv = orig_argv
+        if tmp_token_path:        # ➌ clean-up
+            os.remove("TOKEN")            # remove symlink
+            os.unlink(tmp_token_path.name) # remove temp file
 
     _sanitize(cfg)