Add Firstrade session check persistence (#12)

Author: Pigbibi

.env.example

@@ -24,8 +24,16 @@ GLOBAL_TELEGRAM_CHAT_ID=
 
 # Runtime safety controls.
 FIRSTRADE_COOKIE_DIR=.runtime/firstrade-cookies
+FIRSTRADE_REUSE_SESSION=false
+FIRSTRADE_SESSION_CACHE_TTL_SECONDS=21600
+FIRSTRADE_PERSIST_SESSION_CACHE=false
+FIRSTRADE_GCS_STATE_BUCKET=
+FIRSTRADE_STATE_PREFIX=firstrade-platform
+FIRSTRADE_PERSIST_ACCOUNT_SNAPSHOT=false
 FIRSTRADE_ENABLE_LIVE_TRADING=false
 FIRSTRADE_RUN_SMOKE_ON_HTTP=false
+FIRSTRADE_RUN_SESSION_CHECK_ON_HTTP=false
+FIRSTRADE_SESSION_CHECK_INCLUDE_POSITIONS=false
 FIRSTRADE_RUN_STRATEGY_ON_HTTP=false
 FIRSTRADE_LIVE_ORDER_ACK=false
 FIRSTRADE_MAX_ORDER_NOTIONAL_USD=25

.github/workflows/sync-cloud-run-env.yml

@@ -50,6 +50,12 @@ jobs:
       FIRSTRADE_SMOKE_SYMBOL: ${{ vars.FIRSTRADE_SMOKE_SYMBOL }}
       FIRSTRADE_FEATURE_SNAPSHOT_PATH: ${{ vars.FIRSTRADE_FEATURE_SNAPSHOT_PATH }}
       FIRSTRADE_FEATURE_SNAPSHOT_MANIFEST_PATH: ${{ vars.FIRSTRADE_FEATURE_SNAPSHOT_MANIFEST_PATH }}
+      FIRSTRADE_GCS_STATE_BUCKET: ${{ vars.FIRSTRADE_GCS_STATE_BUCKET }}
+      FIRSTRADE_PERSIST_ACCOUNT_SNAPSHOT: ${{ vars.FIRSTRADE_PERSIST_ACCOUNT_SNAPSHOT }}
+      FIRSTRADE_PERSIST_SESSION_CACHE: ${{ vars.FIRSTRADE_PERSIST_SESSION_CACHE }}
+      FIRSTRADE_RUN_SESSION_CHECK_ON_HTTP: ${{ vars.FIRSTRADE_RUN_SESSION_CHECK_ON_HTTP }}
+      FIRSTRADE_SESSION_CHECK_INCLUDE_POSITIONS: ${{ vars.FIRSTRADE_SESSION_CHECK_INCLUDE_POSITIONS }}
+      FIRSTRADE_STATE_PREFIX: ${{ vars.FIRSTRADE_STATE_PREFIX }}
       FIRSTRADE_STRATEGY_CONFIG_PATH: ${{ vars.FIRSTRADE_STRATEGY_CONFIG_PATH }}
       FIRSTRADE_STRATEGY_PLUGIN_MOUNTS_JSON: ${{ vars.FIRSTRADE_STRATEGY_PLUGIN_MOUNTS_JSON }}
       FIRSTRADE_TECH_RUNTIME_EXECUTION_WINDOW_TRADING_DAYS: ${{ vars.FIRSTRADE_TECH_RUNTIME_EXECUTION_WINDOW_TRADING_DAYS }}
@@ -394,8 +400,14 @@ jobs:
           add_optional_env FIRSTRADE_DRY_RUN_ONLY
           add_optional_env FIRSTRADE_REUSE_SESSION
           add_optional_env FIRSTRADE_SESSION_CACHE_TTL_SECONDS
+          add_optional_env FIRSTRADE_PERSIST_SESSION_CACHE
+          add_optional_env FIRSTRADE_GCS_STATE_BUCKET
+          add_optional_env FIRSTRADE_STATE_PREFIX
+          add_optional_env FIRSTRADE_PERSIST_ACCOUNT_SNAPSHOT
           add_optional_env FIRSTRADE_ENABLE_LIVE_TRADING
           add_optional_env FIRSTRADE_RUN_SMOKE_ON_HTTP
+          add_optional_env FIRSTRADE_RUN_SESSION_CHECK_ON_HTTP
+          add_optional_env FIRSTRADE_SESSION_CHECK_INCLUDE_POSITIONS
           add_optional_env FIRSTRADE_RUN_STRATEGY_ON_HTTP
           add_optional_env FIRSTRADE_LIVE_ORDER_ACK
           add_optional_env FIRSTRADE_MAX_ORDER_NOTIONAL_USD

README.md

@@ -73,6 +73,10 @@ commit credentials.
 | `FIRSTRADE_DRY_RUN_ONLY` | Optional | Defaults to `true` for platform runtime |
 | `FIRSTRADE_REUSE_SESSION` | Optional | Reuse cached Firstrade session headers inside the same warm runtime instance before logging in again. Defaults to `false` |
 | `FIRSTRADE_SESSION_CACHE_TTL_SECONDS` | Optional | Max age for local session header reuse when `FIRSTRADE_REUSE_SESSION=true`. Defaults to `21600` |
+| `FIRSTRADE_PERSIST_SESSION_CACHE` | Optional | Persist Firstrade session headers to the configured GCS state bucket when `FIRSTRADE_REUSE_SESSION=true`. Defaults to `false` |
+| `FIRSTRADE_GCS_STATE_BUCKET` | Optional | GCS bucket for runtime state JSON, including persisted session cache and account funds snapshots |
+| `FIRSTRADE_STATE_PREFIX` | Optional | Object prefix within `FIRSTRADE_GCS_STATE_BUCKET`, default `firstrade-platform` |
+| `FIRSTRADE_PERSIST_ACCOUNT_SNAPSHOT` | Optional | Persist compact masked account funds snapshots from `/session-check`. Defaults to `false` |
 | `ACCOUNT_PREFIX` | Optional | Alert/log prefix, default `FIRSTRADE` |
 | `ACCOUNT_REGION` | Optional | Runtime account scope, default `US` |
 | `NOTIFY_LANG` | Optional | Notification language, `en` or `zh` |
@@ -81,6 +85,8 @@ commit credentials.
 | `FIRSTRADE_COOKIE_DIR` | Optional | Cookie cache directory, default `.runtime/firstrade-cookies` |
 | `FIRSTRADE_ENABLE_LIVE_TRADING` | Optional | Must be `true` before any live order can be submitted |
 | `FIRSTRADE_RUN_SMOKE_ON_HTTP` | Optional | Must be `true` before `/smoke` performs a real login/quote |
+| `FIRSTRADE_RUN_SESSION_CHECK_ON_HTTP` | Optional | Must be `true` before `/session-check` performs a read-only login/session/account-state check |
+| `FIRSTRADE_SESSION_CHECK_INCLUDE_POSITIONS` | Optional | Include compact symbol/quantity/market-value positions in `/session-check` funds snapshots. Defaults to `false` |
 | `FIRSTRADE_RUN_STRATEGY_ON_HTTP` | Optional | Must be `true` before `/run` performs strategy evaluation and order routing |
 | `FIRSTRADE_LIVE_ORDER_ACK` | Optional | Must be `true` before `/run` can submit live orders |
 | `FIRSTRADE_MAX_ORDER_NOTIONAL_USD` | Optional | Single-order cap for strategy-generated orders, default `25` |
@@ -175,11 +181,21 @@ The strategy execution service uses whole-share limit orders for generated
 strategy orders. If the notional cap is below the current price of a target
 symbol, that order is skipped instead of being enlarged.
 
-`FIRSTRADE_REUSE_SESSION=true` reduces repeated login attempts while the same
-Cloud Run instance stays warm. It stores the current session headers only in the
-container-local cookie directory and tries that session before calling Firstrade
-login again. A cold start, new revision, expired session, or broker-side
-invalidation still falls back to a fresh login.
+`FIRSTRADE_REUSE_SESSION=true` reduces repeated login attempts by trying cached
+session headers before calling Firstrade login again. By default this cache is
+container-local. When `FIRSTRADE_PERSIST_SESSION_CACHE=true` and
+`FIRSTRADE_GCS_STATE_BUCKET` is set, the same cache is also written to GCS so a
+cold start can try the last known session first. Expired sessions, new broker
+sessions from another device, or broker-side invalidation still fall back to a
+fresh login.
+
+`/session-check` is a read-only route for session keepalive experiments and
+account-state persistence. It connects to Firstrade, selects the account, reads
+balances, optionally reads positions, and returns a compact masked funds
+snapshot. With `FIRSTRADE_PERSIST_ACCOUNT_SNAPSHOT=true`, it writes the snapshot
+to `accounts/<masked-account>/funds/latest.json` plus a timestamped history path
+under the configured GCS prefix. Raw account IDs and login secrets are not
+included in the snapshot.
 
 ## Cloud Run Shape
 
@@ -190,6 +206,8 @@ invalidation still falls back to a fresh login.
 - `/probe` health metadata only
 - `/profiles` shared US equity strategy matrix
 - `/smoke` login + quote only when `FIRSTRADE_RUN_SMOKE_ON_HTTP=true`
+- `/session-check` read-only session/account-state check only when
+  `FIRSTRADE_RUN_SESSION_CHECK_ON_HTTP=true`
 - `/run` strategy evaluation + guarded order routing only when
   `FIRSTRADE_RUN_STRATEGY_ON_HTTP=true`
 

application/firstrade_client.py

@@ -13,6 +13,8 @@
 from time import time
 from typing import Any, Callable
 
+from application.state_persistence import GcsStateStore
+
 
 class FirstradePlatformError(RuntimeError):
     """Base error for platform integration failures."""
@@ -42,6 +44,9 @@ class FirstradeCredentials:
     cookie_dir: str = ".runtime/firstrade-cookies"
     reuse_session: bool = False
     session_cache_ttl_seconds: int = 21_600
+    persist_session_cache: bool = False
+    gcs_state_bucket: str = ""
+    gcs_state_prefix: str = "firstrade-platform"
     debug: bool = False
 
     @classmethod
@@ -63,6 +68,13 @@ def from_env(cls, env: Callable[[str, str | None], str | None] = os.getenv) -> "
                 env("FIRSTRADE_SESSION_CACHE_TTL_SECONDS", "21600"),
                 default=21_600,
             ),
+            persist_session_cache=(env("FIRSTRADE_PERSIST_SESSION_CACHE", "false") or "")
+            .strip()
+            .lower()
+            == "true",
+            gcs_state_bucket=(env("FIRSTRADE_GCS_STATE_BUCKET", "") or "").strip(),
+            gcs_state_prefix=env("FIRSTRADE_STATE_PREFIX", "firstrade-platform")
+            or "firstrade-platform",
             debug=(env("FIRSTRADE_DEBUG", "false") or "").lower() == "true",
         )
 
@@ -195,6 +207,7 @@ def __init__(
         order_factory: Callable[[Any], Any] | None = None,
         quote_factory: Callable[[Any, str, str], Any] | None = None,
         ohlc_factory: Callable[[Any, str, str], Any] | None = None,
+        session_cache_store: GcsStateStore | None = None,
     ) -> None:
         self.credentials = credentials
         self.live_trading_enabled = live_trading_enabled
@@ -203,6 +216,7 @@ def __init__(
         self._order_factory = order_factory
         self._quote_factory = quote_factory
         self._ohlc_factory = ohlc_factory
+        self._session_cache_store = session_cache_store
         self.session: Any | None = None
         self.account_data: Any | None = None
         self.session_reused = False
@@ -261,19 +275,31 @@ def _load_session_cache(self, cookie_dir: Path) -> dict[str, Any] | None:
         try:
             payload = json.loads(path.read_text())
         except (OSError, json.JSONDecodeError):
+            payload = None
+        if self._is_valid_session_cache_payload(payload):
+            return payload
+        store = self._session_state_store()
+        if store is None:
             return None
-        if not isinstance(payload, dict):
+        try:
+            persisted_payload = store.read_json(self._session_state_key())
+        except Exception:
             return None
+        if self._is_valid_session_cache_payload(persisted_payload):
+            return persisted_payload
+        return None
+
+    def _is_valid_session_cache_payload(self, payload: Any) -> bool:
+        if not isinstance(payload, dict):
+            return False
         try:
             saved_at = float(payload.get("saved_at") or 0.0)
         except (TypeError, ValueError):
-            return None
+            return False
         ttl = max(1, int(self.credentials.session_cache_ttl_seconds or 1))
         if saved_at <= 0.0 or (time() - saved_at) > ttl:
-            return None
-        if not payload.get("ftat") or not payload.get("sid"):
-            return None
-        return payload
+            return False
+        return bool(payload.get("ftat") and payload.get("sid"))
 
     def _try_cached_session(
         self,
@@ -288,10 +314,18 @@ def _try_cached_session(
         try:
             from firstrade import urls
 
-            session.session.headers.update(urls.session_headers())
-            session.session.headers["access-token"] = urls.access_token()
-            session.session.headers["ftat"] = str(payload["ftat"])
-            session.session.headers["sid"] = str(payload["sid"])
+            if hasattr(session, "build_session_from_tokens"):
+                session.build_session_from_tokens(payload)
+            else:
+                session.session.headers.update(urls.session_headers())
+                session.session.headers["access-token"] = str(
+                    payload.get("access-token") or urls.access_token()
+                )
+                session.session.headers["ftat"] = str(payload["ftat"])
+                session.session.headers["sid"] = str(payload["sid"])
+                cookies = payload.get("cookies")
+                if isinstance(cookies, dict) and hasattr(session.session, "cookies"):
+                    session.session.cookies.update(cookies)
             account_data = account_data_factory(session)
         except Exception:
             try:
@@ -302,24 +336,55 @@ def _try_cached_session(
         self.session = session
         self.account_data = account_data
         self.session_reused = True
+        self._save_session_cache(cookie_dir)
         return True
 
     def _save_session_cache(self, cookie_dir: Path) -> None:
         if not self.credentials.reuse_session or self.session is None:
             return
-        headers = getattr(getattr(self.session, "session", None), "headers", {}) or {}
+        session_obj = getattr(self.session, "session", None)
+        headers = getattr(session_obj, "headers", {}) or {}
+        cookies = {}
+        if hasattr(session_obj, "cookies"):
+            try:
+                cookies = session_obj.cookies.get_dict()
+            except Exception:
+                cookies = {}
         payload = {
+            "access-token": headers.get("access-token"),
             "ftat": headers.get("ftat"),
             "sid": headers.get("sid"),
+            "cookies": cookies,
             "saved_at": time(),
         }
         if not payload["ftat"] or not payload["sid"]:
             return
         try:
             self._session_cache_path(cookie_dir).write_text(json.dumps(payload), encoding="utf-8")
         except OSError:
+            pass
+        store = self._session_state_store()
+        if store is None:
+            return
+        try:
+            store.write_json(self._session_state_key(), payload)
+        except Exception:
             return
 
+    def _session_state_store(self) -> GcsStateStore | None:
+        if self._session_cache_store is not None:
+            return self._session_cache_store
+        if not self.credentials.persist_session_cache or not self.credentials.gcs_state_bucket:
+            return None
+        return GcsStateStore(
+            bucket=self.credentials.gcs_state_bucket,
+            prefix=self.credentials.gcs_state_prefix,
+        )
+
+    def _session_state_key(self) -> str:
+        safe_username = "".join(ch for ch in self.credentials.username if ch.isalnum() or ch in ("-", "_"))
+        return f"sessions/{safe_username or 'unknown'}/latest.json"
+
     def require_connected(self) -> tuple[Any, Any]:
         if self.session is None or self.account_data is None:
             raise FirstradePlatformError("Firstrade client is not connected.")