Skip to content

Commit b7f7034

Browse files
crdantclaude
crdant
and
claude
committed
Add support for per-instance service account tokens
This change enables the SDK to accept and use instance-specific service account tokens, allowing each instance to have its own authentication token separate from the customer-level token. Key changes: - Instance and AsyncInstance now accept optional service_account_token parameter - When API returns service_token on instance creation, it replaces the dynamic_token - User-provided service_account_tokens are stored and used as fallback - Both sync and async implementations updated - Examples updated to display token changes before/after instance creation - Comprehensive test coverage for token replacement behavior The implementation uses a simple token replacement pattern where the instance token replaces the customer token in state, ensuring all subsequent requests use the instance-specific token. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
1 parent b7d0003 commit b7f7034

File tree

4 files changed

+367
-11
lines changed

4 files changed

+367
-11
lines changed

‎examples/basic_example.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,11 +71,23 @@ def main():
7171
)
7272
print(f"✓ Customer created/retrieved - ID: {customer.customer_id}")
7373

74+
# Display service token after customer creation
75+
token_after_customer = client.state_manager.get_dynamic_token()
76+
if token_after_customer:
77+
print(f" Service token: {token_after_customer}")
78+
7479
# Create or get instance
7580
print("\nCreating/getting instance for customer...")
7681
instance = customer.get_or_create_instance()
7782
print(f"✓ Instance created/retrieved - ID: {instance.instance_id}")
7883

84+
# Display service token after instance creation (may have been replaced)
85+
token_after_instance = client.state_manager.get_dynamic_token()
86+
if token_after_instance:
87+
print(f" Service token: {token_after_instance}")
88+
if token_after_customer != token_after_instance:
89+
print(f" ⚠️ Token was replaced by instance-specific token")
90+
7991
# Set instance status
8092
instance.set_status(args.status)
8193
print(f"✓ Instance status set to: {args.status}")
@@ -89,6 +101,14 @@ def main():
89101
print(f"Customer ID: {customer.customer_id}")
90102
print(f"Instance ID: {instance.instance_id}")
91103

104+
# Show final token
105+
final_token = client.state_manager.get_dynamic_token()
106+
print("\nService Token Information:")
107+
if final_token:
108+
print(f" Active service token: {final_token}")
109+
else:
110+
print(" Service token: Not available")
111+
92112

93113
if __name__ == "__main__":
94114
main()

‎examples/metrics_example.py‎

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -72,14 +72,21 @@ async def main():
7272
)
7373
print(f"✓ Customer created/retrieved - ID: {customer.customer_id}")
7474

75+
# Display service token after customer creation
76+
token_after_customer = client.state_manager.get_dynamic_token()
77+
if token_after_customer:
78+
print(f" Service token: {token_after_customer}")
79+
7580
# Get or create the associated instance
7681
instance = await customer.get_or_create_instance()
77-
print(f"Instance ID: {instance.instance_id}")
7882
print(f"✓ Instance created/retrieved - ID: {instance.instance_id}")
7983

80-
# Get or create the associated instance
81-
instance = await customer.get_or_create_instance()
82-
print(f"Instance ID: {instance.instance_id}")
84+
# Display service token after instance creation (may have been replaced)
85+
token_after_instance = client.state_manager.get_dynamic_token()
86+
if token_after_instance:
87+
print(f" Service token: {token_after_instance}")
88+
if token_after_customer != token_after_instance:
89+
print(f" ⚠️ Token was replaced by instance-specific token")
8390

8491
# Set instance status
8592
await instance.set_status(args.status)
@@ -96,9 +103,19 @@ async def main():
96103
instance.send_metric("memory_usage", 0.67),
97104
instance.send_metric("disk_usage", 0.45),
98105
)
99-
print("Metrics sent successfully")
106+
print("✓ Metrics sent successfully")
107+
108+
print("\n🎉 Metrics example completed successfully!")
109+
print(f"Customer ID: {customer.customer_id}")
110+
print(f"Instance ID: {instance.instance_id}")
100111

101-
print(f"Instance ID: {instance.instance_id}")
112+
# Show final token
113+
final_token = client.state_manager.get_dynamic_token()
114+
print("\nService Token Information:")
115+
if final_token:
116+
print(f" Active service token: {final_token}")
117+
else:
118+
print(" Service token: Not available")
102119

103120

104121
if __name__ == "__main__":

‎replicated/resources.py‎

Lines changed: 54 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,14 +27,26 @@ def __init__(
2727
self.channel = channel
2828
self._data = kwargs
2929

30-
def get_or_create_instance(self) -> Union["Instance", "AsyncInstance"]:
30+
def get_or_create_instance(
31+
self, service_account_token: Optional[str] = None
32+
) -> Union["Instance", "AsyncInstance"]:
3133
"""Get or create an instance for this customer."""
3234
if hasattr(self._client, "_get_or_create_instance_async"):
3335
# type: ignore[arg-type]
34-
return AsyncInstance(self._client, self.customer_id, self.instance_id)
36+
return AsyncInstance(
37+
self._client,
38+
self.customer_id,
39+
self.instance_id,
40+
service_account_token=service_account_token,
41+
)
3542
else:
3643
# type: ignore[arg-type]
37-
return Instance(self._client, self.customer_id, self.instance_id)
44+
return Instance(
45+
self._client,
46+
self.customer_id,
47+
self.instance_id,
48+
service_account_token=service_account_token,
49+
)
3850

3951
def __getattr__(self, name: str) -> Any:
4052
"""Access additional customer data."""
@@ -45,10 +57,17 @@ class AsyncCustomer(Customer):
4557
"""Async version of Customer."""
4658

4759
# type: ignore[override]
48-
async def get_or_create_instance(self) -> "AsyncInstance":
60+
async def get_or_create_instance(
61+
self, service_account_token: Optional[str] = None
62+
) -> "AsyncInstance":
4963
"""Get or create an instance for this customer."""
5064
# type: ignore[arg-type]
51-
return AsyncInstance(self._client, self.customer_id, self.instance_id)
65+
return AsyncInstance(
66+
self._client,
67+
self.customer_id,
68+
self.instance_id,
69+
service_account_token=service_account_token,
70+
)
5271

5372

5473
class Instance:
@@ -59,11 +78,13 @@ def __init__(
5978
client: "ReplicatedClient",
6079
customer_id: str,
6180
instance_id: Optional[str] = None,
81+
service_account_token: Optional[str] = None,
6282
**kwargs: Any,
6383
) -> None:
6484
self._client = client
6585
self.customer_id = customer_id
6686
self.instance_id = instance_id
87+
self._service_account_token = service_account_token
6788
self._machine_id = client._machine_id
6889
self._data = kwargs
6990
self._status = "ready"
@@ -112,12 +133,18 @@ def set_version(self, version: str) -> None:
112133
def _ensure_instance(self) -> None:
113134
"""Ensure the instance ID is generated and cached."""
114135
if self.instance_id:
136+
# If we have an instance ID but a service token was provided, replace dynamic token
137+
if self._service_account_token:
138+
self._client.state_manager.set_dynamic_token(self._service_account_token)
115139
return
116140

117141
# Check if instance ID is cached
118142
cached_instance_id = self._client.state_manager.get_instance_id()
119143
if cached_instance_id:
120144
self.instance_id = cached_instance_id
145+
# If we have a service token provided, replace dynamic token
146+
if self._service_account_token:
147+
self._client.state_manager.set_dynamic_token(self._service_account_token)
121148
return
122149

123150
# Create new instance
@@ -135,6 +162,13 @@ def _ensure_instance(self) -> None:
135162
self.instance_id = response["instance_id"]
136163
self._client.state_manager.set_instance_id(self.instance_id)
137164

165+
# If API returns a service_token, replace the dynamic token with it
166+
if "service_token" in response:
167+
self._client.state_manager.set_dynamic_token(response["service_token"])
168+
# Otherwise, if user provided a service token, use that
169+
elif self._service_account_token:
170+
self._client.state_manager.set_dynamic_token(self._service_account_token)
171+
138172
def _report_instance(self) -> None:
139173
"""Send instance telemetry to vandoor."""
140174
if not self.instance_id:
@@ -190,11 +224,13 @@ def __init__(
190224
client: "AsyncReplicatedClient",
191225
customer_id: str,
192226
instance_id: Optional[str] = None,
227+
service_account_token: Optional[str] = None,
193228
**kwargs: Any,
194229
) -> None:
195230
self._client = client
196231
self.customer_id = customer_id
197232
self.instance_id = instance_id
233+
self._service_account_token = service_account_token
198234
self._machine_id = client._machine_id
199235
self._data = kwargs
200236
self._status = "ready"
@@ -243,12 +279,18 @@ async def set_version(self, version: str) -> None:
243279
async def _ensure_instance(self) -> None:
244280
"""Ensure the instance ID is generated and cached."""
245281
if self.instance_id:
282+
# If we have an instance ID but a service token was provided, replace dynamic token
283+
if self._service_account_token:
284+
self._client.state_manager.set_dynamic_token(self._service_account_token)
246285
return
247286

248287
# Check if instance ID is cached
249288
cached_instance_id = self._client.state_manager.get_instance_id()
250289
if cached_instance_id:
251290
self.instance_id = cached_instance_id
291+
# If we have a service token provided, replace dynamic token
292+
if self._service_account_token:
293+
self._client.state_manager.set_dynamic_token(self._service_account_token)
252294
return
253295

254296
# Create new instance
@@ -266,6 +308,13 @@ async def _ensure_instance(self) -> None:
266308
self.instance_id = response["instance_id"]
267309
self._client.state_manager.set_instance_id(self.instance_id)
268310

311+
# If API returns a service_token, replace the dynamic token with it
312+
if "service_token" in response:
313+
self._client.state_manager.set_dynamic_token(response["service_token"])
314+
# Otherwise, if user provided a service token, use that
315+
elif self._service_account_token:
316+
self._client.state_manager.set_dynamic_token(self._service_account_token)
317+
269318
async def _report_instance(self) -> None:
270319
"""Send instance telemetry to vandoor."""
271320
if not self.instance_id:

0 commit comments

Comments
 (0)