fix: claim Firstrade live runs atomically before orders (#249)

Co-authored-by: Codex <noreply@openai.com>

Author: Pigbibi

application/rebalance_service.py

@@ -30,6 +30,7 @@
 from application.state_persistence import GcsStateStore, build_gcs_state_store_from_env
 from application.strategy_run_persistence import (
     build_strategy_run_state,
+    claim_live_strategy_run,
     is_duplicate_live_run,
     persist_strategy_run_state,
     read_latest_strategy_run_state,
@@ -512,13 +513,26 @@ def log_message(message: str) -> None:
     masked_account = mask_account_id(account)
     existing_run = None
     if persist_strategy_runs and not settings.dry_run_only:
+        claim_acquired = claim_live_strategy_run(
+            store=store,
+            account=masked_account,
+            strategy_profile=strategy_runtime.profile,
+            run_period=run_period,
+            now=now,
+        )
         existing_run = read_latest_strategy_run_state(
             store=store,
             account=masked_account,
             strategy_profile=strategy_runtime.profile,
             run_period=run_period,
         )
-        if is_duplicate_live_run(existing_run):
+        if not claim_acquired and existing_run is None:
+            existing_run = {
+                "stage": "PENDING_SUBMISSION",
+                "as_of": now.isoformat(),
+                "claim_only": True,
+            }
+        if not claim_acquired or is_duplicate_live_run(existing_run):
             duplicate_stage = str(existing_run.get("stage") or "NO_ACTION")
             duplicate_skipped_orders = [
                 {

application/state_persistence.py

@@ -70,6 +70,15 @@ def write_json(self, key: str, payload: dict[str, Any]) -> bool:
             raise StatePersistenceError(f"GCS write failed for {key}: {exc}") from exc
         return True
 
+    def create_json(self, key: str, payload: dict[str, Any]) -> bool:
+        """Atomically create a JSON object; return False if it already exists."""
+        uri = self._object_uri(key)
+        data = json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
+        try:
+            return bool(_object_store().create_text(uri, data, content_type="application/json"))
+        except Exception as exc:
+            raise StatePersistenceError(f"GCS atomic create failed for {key}: {exc}") from exc
+
 
 def build_gcs_state_store_from_env(
     env: Callable[[str, str | None], str | None] = os.getenv,

application/strategy_run_persistence.py

@@ -116,6 +116,37 @@ def strategy_run_history_key(
     )
 
 
+def strategy_run_claim_key(
+    *, account: str, strategy_profile: str, run_period: str,
+) -> str:
+    """Permanent live claim key; a failed/unknown submission must remain blocked."""
+    return (
+        f"strategy-runs/claims/{safe_key(account)}/{safe_key(strategy_profile)}/"
+        f"{safe_key(run_period)}.json"
+    )
+
+
+def claim_live_strategy_run(
+    *, store: GcsStateStore, account: str, strategy_profile: str,
+    run_period: str, now: datetime | None = None,
+) -> bool:
+    """Acquire the durable pre-order claim using object-store create-if-absent."""
+    payload = {
+        "stage": "PENDING_SUBMISSION",
+        "account": account,
+        "strategy_profile": strategy_profile,
+        "run_period": run_period,
+        "as_of": (now or utcnow()).isoformat(),
+        "no_order_submitted": True,
+    }
+    return store.create_json(
+        strategy_run_claim_key(
+            account=account, strategy_profile=strategy_profile, run_period=run_period,
+        ),
+        payload,
+    )
+
+
 def read_latest_strategy_run_state(
     *,
     store: GcsStateStore,

tests/test_rebalance_service.py

@@ -157,6 +157,12 @@ def write_json(self, key, payload):
         self.writes.append((key, dict(payload)))
         return True
 
+    def create_json(self, key, payload):
+        if key in self.payloads:
+            return False
+        self.payloads[key] = dict(payload)
+        return True
+
 
 def _latest_strategy_run_payloads(store: FakeStateStore) -> list[dict]:
     return [payload for key, payload in store.writes if key.endswith("latest.json")]

tests/test_strategy_run_claim.py

@@ -0,0 +1,33 @@
+from application.strategy_run_persistence import (
+    claim_live_strategy_run,
+    strategy_run_claim_key,
+)
+
+
+class AtomicFakeStore:
+    def __init__(self):
+        self.payloads = {}
+
+    def create_json(self, key, payload):
+        if key in self.payloads:
+            return False
+        self.payloads[key] = dict(payload)
+        return True
+
+
+def test_live_claim_is_create_only_and_permanent():
+    store = AtomicFakeStore()
+    kwargs = {
+        "store": store,
+        "account": "****1234",
+        "strategy_profile": "tqqq_core",
+        "run_period": "2026-08",
+    }
+
+    assert claim_live_strategy_run(**kwargs) is True
+    assert claim_live_strategy_run(**kwargs) is False
+    key = strategy_run_claim_key(
+        account="****1234", strategy_profile="tqqq_core", run_period="2026-08"
+    )
+    assert store.payloads[key]["stage"] == "PENDING_SUBMISSION"
+    assert store.payloads[key]["no_order_submitted"] is True