diff --git a/application/rebalance_service.py b/application/rebalance_service.py index aee3fc9..bcd5fc6 100644 --- a/application/rebalance_service.py +++ b/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 = [ { diff --git a/application/state_persistence.py b/application/state_persistence.py index 255fd09..1e636f3 100644 --- a/application/state_persistence.py +++ b/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, diff --git a/application/strategy_run_persistence.py b/application/strategy_run_persistence.py index 0143f74..8066d69 100644 --- a/application/strategy_run_persistence.py +++ b/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, diff --git a/tests/test_rebalance_service.py b/tests/test_rebalance_service.py index 7dced56..18b3182 100644 --- a/tests/test_rebalance_service.py +++ b/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")] diff --git a/tests/test_strategy_run_claim.py b/tests/test_strategy_run_claim.py new file mode 100644 index 0000000..845496b --- /dev/null +++ b/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