Adds support for per-instance service account tokens

examples/basic_example.py

@@ -71,11 +71,23 @@ def main():
         )
         print(f"✓ Customer created/retrieved - ID: {customer.customer_id}")
 
+        # Display service token after customer creation
+        token_after_customer = client.state_manager.get_dynamic_token()
+        if token_after_customer:
+            print(f"  Service token: {token_after_customer}")
+
         # Create or get instance
         print("\nCreating/getting instance for customer...")
         instance = customer.get_or_create_instance()
         print(f"✓ Instance created/retrieved - ID: {instance.instance_id}")
 
+        # Display service token after instance creation (may have been replaced)
+        token_after_instance = client.state_manager.get_dynamic_token()
+        if token_after_instance:
+            print(f"  Service token: {token_after_instance}")
+            if token_after_customer != token_after_instance:
+                print("  ⚠️  Token was replaced by instance-specific token")
+
         # Set instance status
         instance.set_status(args.status)
         print(f"✓ Instance status set to: {args.status}")
@@ -89,6 +101,14 @@ def main():
         print(f"Customer ID: {customer.customer_id}")
         print(f"Instance ID: {instance.instance_id}")
 
+        # Show final token
+        final_token = client.state_manager.get_dynamic_token()
+        print("\nService Token Information:")
+        if final_token:
+            print(f"  Active service token: {final_token}")
+        else:
+            print("  Service token: Not available")
+
 
 if __name__ == "__main__":
     main()

examples/metrics_example.py

@@ -72,14 +72,21 @@ async def main():
         )
         print(f"✓ Customer created/retrieved - ID: {customer.customer_id}")
 
+        # Display service token after customer creation
+        token_after_customer = client.state_manager.get_dynamic_token()
+        if token_after_customer:
+            print(f"  Service token: {token_after_customer}")
+
         # Get or create the associated instance
         instance = await customer.get_or_create_instance()
-        print(f"Instance ID: {instance.instance_id}")
         print(f"✓ Instance created/retrieved - ID: {instance.instance_id}")
 
-        # Get or create the associated instance
-        instance = await customer.get_or_create_instance()
-        print(f"Instance ID: {instance.instance_id}")
+        # Display service token after instance creation (may have been replaced)
+        token_after_instance = client.state_manager.get_dynamic_token()
+        if token_after_instance:
+            print(f"  Service token: {token_after_instance}")
+            if token_after_customer != token_after_instance:
+                print("  ⚠️  Token was replaced by instance-specific token")
 
         # Set instance status
         await instance.set_status(args.status)
@@ -96,9 +103,19 @@ async def main():
             instance.send_metric("memory_usage", 0.67),
             instance.send_metric("disk_usage", 0.45),
         )
-        print("Metrics sent successfully")
+        print("✓ Metrics sent successfully")
+
+        print("\n🎉 Metrics example completed successfully!")
+        print(f"Customer ID: {customer.customer_id}")
+        print(f"Instance ID: {instance.instance_id}")
 
-    print(f"Instance ID: {instance.instance_id}")
+        # Show final token
+        final_token = client.state_manager.get_dynamic_token()
+        print("\nService Token Information:")
+        if final_token:
+            print(f"  Active service token: {final_token}")
+        else:
+            print("  Service token: Not available")
 
 
 if __name__ == "__main__":

replicated/resources.py

@@ -27,14 +27,26 @@ def __init__(
         self.channel = channel
         self._data = kwargs
 
-    def get_or_create_instance(self) -> Union["Instance", "AsyncInstance"]:
+    def get_or_create_instance(
+        self, service_account_token: Optional[str] = None
+    ) -> Union["Instance", "AsyncInstance"]:
         """Get or create an instance for this customer."""
         if hasattr(self._client, "_get_or_create_instance_async"):
             # type: ignore[arg-type]
-            return AsyncInstance(self._client, self.customer_id, self.instance_id)
+            return AsyncInstance(
+                self._client,
+                self.customer_id,
+                self.instance_id,
+                service_account_token=service_account_token,
+            )
         else:
             # type: ignore[arg-type]
-            return Instance(self._client, self.customer_id, self.instance_id)
+            return Instance(
+                self._client,
+                self.customer_id,
+                self.instance_id,
+                service_account_token=service_account_token,
+            )
 
     def __getattr__(self, name: str) -> Any:
         """Access additional customer data."""
@@ -45,10 +57,17 @@ class AsyncCustomer(Customer):
     """Async version of Customer."""
 
     # type: ignore[override]
-    async def get_or_create_instance(self) -> "AsyncInstance":
+    async def get_or_create_instance(
+        self, service_account_token: Optional[str] = None
+    ) -> "AsyncInstance":
         """Get or create an instance for this customer."""
         # type: ignore[arg-type]
-        return AsyncInstance(self._client, self.customer_id, self.instance_id)
+        return AsyncInstance(
+            self._client,
+            self.customer_id,
+            self.instance_id,
+            service_account_token=service_account_token,
+        )
 
 
 class Instance:
@@ -59,11 +78,13 @@ def __init__(
         client: "ReplicatedClient",
         customer_id: str,
         instance_id: Optional[str] = None,
+        service_account_token: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
         self._client = client
         self.customer_id = customer_id
         self.instance_id = instance_id
+        self._service_account_token = service_account_token
         self._machine_id = client._machine_id
         self._data = kwargs
         self._status = "ready"
@@ -112,12 +133,23 @@ def set_version(self, version: str) -> None:
     def _ensure_instance(self) -> None:
         """Ensure the instance ID is generated and cached."""
         if self.instance_id:
+            # If we have an instance ID but a service token was provided,
+            # replace dynamic token
+            if self._service_account_token:
+                self._client.state_manager.set_dynamic_token(
+                    self._service_account_token
+                )
             return
 
         # Check if instance ID is cached
         cached_instance_id = self._client.state_manager.get_instance_id()
         if cached_instance_id:
             self.instance_id = cached_instance_id
+            # If we have a service token provided, replace dynamic token
+            if self._service_account_token:
+                self._client.state_manager.set_dynamic_token(
+                    self._service_account_token
+                )
             return
 
         # Create new instance
@@ -135,6 +167,13 @@ def _ensure_instance(self) -> None:
         self.instance_id = response["instance_id"]
         self._client.state_manager.set_instance_id(self.instance_id)
 
+        # If API returns a service_token, replace the dynamic token with it
+        if "service_token" in response:
+            self._client.state_manager.set_dynamic_token(response["service_token"])
+        # Otherwise, if user provided a service token, use that
+        elif self._service_account_token:
+            self._client.state_manager.set_dynamic_token(self._service_account_token)
+
     def _report_instance(self) -> None:
         """Send instance telemetry to vandoor."""
         if not self.instance_id:
@@ -190,11 +229,13 @@ def __init__(
         client: "AsyncReplicatedClient",
         customer_id: str,
         instance_id: Optional[str] = None,
+        service_account_token: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
         self._client = client
         self.customer_id = customer_id
         self.instance_id = instance_id
+        self._service_account_token = service_account_token
         self._machine_id = client._machine_id
         self._data = kwargs
         self._status = "ready"
@@ -243,12 +284,23 @@ async def set_version(self, version: str) -> None:
     async def _ensure_instance(self) -> None:
         """Ensure the instance ID is generated and cached."""
         if self.instance_id:
+            # If we have an instance ID but a service token was provided,
+            # replace dynamic token
+            if self._service_account_token:
+                self._client.state_manager.set_dynamic_token(
+                    self._service_account_token
+                )
             return
 
         # Check if instance ID is cached
         cached_instance_id = self._client.state_manager.get_instance_id()
         if cached_instance_id:
             self.instance_id = cached_instance_id
+            # If we have a service token provided, replace dynamic token
+            if self._service_account_token:
+                self._client.state_manager.set_dynamic_token(
+                    self._service_account_token
+                )
             return
 
         # Create new instance
@@ -266,6 +318,13 @@ async def _ensure_instance(self) -> None:
         self.instance_id = response["instance_id"]
         self._client.state_manager.set_instance_id(self.instance_id)
 
+        # If API returns a service_token, replace the dynamic token with it
+        if "service_token" in response:
+            self._client.state_manager.set_dynamic_token(response["service_token"])
+        # Otherwise, if user provided a service token, use that
+        elif self._service_account_token:
+            self._client.state_manager.set_dynamic_token(self._service_account_token)
+
     async def _report_instance(self) -> None:
         """Send instance telemetry to vandoor."""
         if not self.instance_id: