From 473ba573a39aef82135a9eafcd0dbe89ba0e3ae9 Mon Sep 17 00:00:00 2001
From: dadachann <185672915+dadachann@users.noreply.github.com>
Date: Thu, 30 Jul 2026 21:37:06 +0000
Subject: [PATCH] fix(cloud): retain launch replay records through clock skew

---
 src/langbot/pkg/cloud/launch.py             | 10 ++++++----
 tests/unit_tests/cloud/test_space_launch.py | 16 ++++++++++++++++
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/langbot/pkg/cloud/launch.py b/src/langbot/pkg/cloud/launch.py
index 3725d03d8..26fcf32ae 100644
--- a/src/langbot/pkg/cloud/launch.py
+++ b/src/langbot/pkg/cloud/launch.py
@@ -7,6 +7,7 @@ import datetime
 import hashlib
 import heapq
 import json
+import math
 import os
 import time
 import typing
@@ -124,7 +125,7 @@ class SpaceLaunchService:
         *,
         expected_workspace_uuid: str | None = None,
     ) -> dict[str, str]:
-        claims = self._verify_assertion(assertion)
+        claims, clock_skew_seconds = self._verify_assertion(assertion)
         payload = claims.get('payload')
         if not isinstance(payload, dict):
             raise SpaceLaunchError('Launch assertion payload must be a JSON object')
@@ -139,14 +140,15 @@ class SpaceLaunchService:
             raise SpaceLaunchError('Launch assertion return path is invalid')
         if expected_workspace_uuid is not None and workspace_uuid != expected_workspace_uuid:
             raise SpaceLaunchError('Launch assertion targets another Workspace')
-        await self._consume_jti(_required_string(claims, 'jti'), _required_int(claims, 'exp', minimum=1))
+        replay_retention_expires_at = _required_int(claims, 'exp', minimum=1) + math.ceil(clock_skew_seconds)
+        await self._consume_jti(_required_string(claims, 'jti'), replay_retention_expires_at)
         return {
             'account_uuid': account_uuid,
             'workspace_uuid': workspace_uuid,
             'return_path': return_path,
         }
 
-    def _verify_assertion(self, token: str) -> dict[str, typing.Any]:
+    def _verify_assertion(self, token: str) -> tuple[dict[str, typing.Any], float]:
         if not getattr(getattr(self.ap, 'deployment', None), 'multi_workspace_enabled', False):
             raise SpaceLaunchError('Space direct launch requires verified Cloud mode')
         public_key, key_id, clock_skew_seconds = self._trust_config()
@@ -197,7 +199,7 @@ class SpaceLaunchService:
             raise SpaceLaunchError('Launch assertion is expired')
         if expires_at <= max(issued_at, not_before):
             raise SpaceLaunchError('Launch assertion expiry must follow issue time')
-        return claims
+        return claims, clock_skew_seconds
 
     def _trust_config(self) -> tuple[Ed25519PublicKey, str, float]:
         data = getattr(getattr(self.ap, 'instance_config', None), 'data', {}) or {}
diff --git a/tests/unit_tests/cloud/test_space_launch.py b/tests/unit_tests/cloud/test_space_launch.py
index 07d0f80de..51aa0edb8 100644
--- a/tests/unit_tests/cloud/test_space_launch.py
+++ b/tests/unit_tests/cloud/test_space_launch.py
@@ -91,6 +91,22 @@ async def test_consumes_valid_workspace_launch_assertion_once():
         await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
 
 
+async def test_consumed_assertion_remains_blocked_through_clock_skew_window():
+    private_key = Ed25519PrivateKey.generate()
+    now = int(time.time())
+    service = _service(private_key, now=now)
+    claims = _claims(now=now)
+    claims['iat'] = now - 10
+    claims['nbf'] = now - 10
+    claims['exp'] = now - 1
+    token = _sign(private_key, claims)
+
+    await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
+
+    with pytest.raises(SpaceLaunchError, match='already been consumed'):
+        await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
+
+
 async def test_replay_cache_does_not_scan_all_live_assertions(monkeypatch):
     private_key = Ed25519PrivateKey.generate()
     now = int(time.time())