Update replit_identity_token_manager.py

src/replit/ai/modelfarm/replit_identity_token_manager.py

@@ -24,15 +24,15 @@ def __init__(self, token_timeout: int = 300):
         self.last_update: Optional[float] = None
         self.token: Optional[str] = None
         self.__set_token_type()
-        if self.token_type != "L402":
-            self.__update_token()
+        self.__update_token()
 
     def __set_token_type(self):
         """Sets the type of token to be used. Bearer if in replit, L402 if not."""
-        if self.__in_deployment():
-            self.token_type = "Bearer"
-        else:
+        if self.__using_L402():
             self.token_type = "L402"
+        else :
+            self.token_type = "Bearer"
+
 
     def get_token(self) -> Optional[str]:
         """Returns the token, updates if the current token has expired.
@@ -58,14 +58,16 @@ def get_new_token(self) -> str:
         Returns:
           str: The most recent token.
         """
-        try:
+        try: 
             if self.__in_deployment():
                 return self.get_deployment_token()
-            return self.get_interactive_token()
-        # If the environment variables are not set, they're not in Replit, so use L402
-        except MissingEnvironmentVariable:
-            print("Using L402 token")
-            return self.get_l402_token()
+            elif self.__using_L402():
+                return self.get_l402_token()
+            else:
+                return self.get_interactive_token()
+        # Fallback, try to let them use L402
+        except:
+            return self.get_402_token()
 
     def get_deployment_token(self) -> str:
         """Fetches deployment token from hostindpid1.
@@ -108,7 +110,7 @@ def get_l402_token(self) -> str:
         try:
             l402_token = self.get_env_var("REPLIT_L402_TOKEN")
             l402_preimage = self.get_env_var("REPLIT_L402_PREIMAGE")
-            return "L402 " + l402_token + ":" + l402_preimage
+            return l402_token + ":" + l402_preimage
         except MissingEnvironmentVariable:
             print("Generating new L402 token")
             res = requests.get("https://matador-replit.kody.repl.co")
@@ -142,6 +144,14 @@ def __in_deployment(self) -> bool:
         """
         return "REPLIT_DEPLOYMENT" in os.environ
 
+    def __using_L402(self) -> bool:
+         """Determines if using L402.
+
+        Returns:
+          bool: True if using L402, False otherwise.
+        """
+        return "REPLIT_L402_TOKEN" in os.environ and "REPLIT_L402_PREIMAGE" in os.environ
+
 
 def split_L402(l402: str) -> tuple[str, str]:
     """Splits L402 token into token and invoice.