refactor(state): pack hanging signatures/tuples in messages/compression/usage (AST-identical)

2026-09-02 16:38:21 -07:00
parent 5c3acca66b
commit 0d7acf67b5
3 changed files with 57 additions and 164 deletions
@@ -151,23 +151,12 @@ class SessionCompressionMixin:
        )

    def publish_compression_child(
-        self,
-        *,
-        parent_session_id: str,
-        child_session_id: str,
-        source: str,
-        messages: List[Dict[str, Any]],
-        model: str = None,
-        model_config: Dict[str, Any] = None,
-        system_prompt: str = None,
-        cwd: str = None,
-        profile_name: str = None,
-        compression_lock_holder: str = None,
-        require_compression_lease: bool = True,
-        require_lease_refresh: bool = False,
-        lease_ttl_seconds: float = 300.0,
-        watermark: Optional[int] = None,
-        watermark_ceiling: Optional[int] = None,
+        self, *, parent_session_id: str, child_session_id: str, source: str,
+        messages: List[Dict[str, Any]], model: str = None, model_config: Dict[str, Any] = None,
+        system_prompt: str = None, cwd: str = None, profile_name: str = None,
+        compression_lock_holder: str = None, require_compression_lease: bool = True,
+        require_lease_refresh: bool = False, lease_ttl_seconds: float = 300.0,
+        watermark: Optional[int] = None, watermark_ceiling: Optional[int] = None,
    ) -> None:
        """Atomically close a parent and publish its durable compression child.

@@ -575,16 +564,9 @@ class SessionCompressionMixin:
        return bool(self._execute_write(_do, patience_s=patience_s))

    def acquire_session_turn_lease(
-        self,
-        session_id: str,
-        holder: str,
-        *,
-        ttl_seconds: float = 300.0,
-        wait_seconds: float = 1800.0,
-        poll_interval_seconds: float = 1.0,
-        on_wait=None,
-        wait_notice_interval_seconds: float = 15.0,
-        should_abort=None,
+        self, session_id: str, holder: str, *, ttl_seconds: float = 300.0,
+        wait_seconds: float = 1800.0, poll_interval_seconds: float = 1.0, on_wait=None,
+        wait_notice_interval_seconds: float = 15.0, should_abort=None,
        acquire_patience_s: float = 0.5,
    ) -> bool:
        """Wait for a cross-process turn lease without holding a SQLite lock.
@@ -201,14 +201,9 @@ class SessionMessagesMixin:
        return value if isinstance(value, str) else json.dumps(value)

    def _check_transcript_write_guards(
-        self,
-        conn,
-        session_id: str,
-        compression_lock_holder: Optional[str],
-        turn_lease_holder: Optional[str] = None,
-        turn_lease_ttl_seconds: float = 300.0,
-        reject_active_turn_lease: bool = False,
-        reject_active_compression_lock: bool = False,
+        self, conn, session_id: str, compression_lock_holder: Optional[str],
+        turn_lease_holder: Optional[str] = None, turn_lease_ttl_seconds: float = 300.0,
+        reject_active_turn_lease: bool = False, reject_active_compression_lock: bool = False,
        allow_closed_compression_parent: bool = False,
    ) -> None:
        """Transcript-write admission checks, run INSIDE the write txn.
@@ -317,30 +312,15 @@ class SessionMessagesMixin:
        )

    def append_message(
-        self,
-        session_id: str,
-        role: str,
-        content: str = None,
-        tool_name: str = None,
-        tool_calls: Any = None,
-        tool_call_id: str = None,
-        token_count: int = None,
-        finish_reason: str = None,
-        reasoning: str = None,
-        reasoning_content: str = None,
-        reasoning_details: Any = None,
-        codex_reasoning_items: Any = None,
-        codex_message_items: Any = None,
-        platform_message_id: str = None,
-        observed: bool = False,
-        effect_disposition: Optional[str] = None,
-        _compressed_summary: bool = False,
-        timestamp: Any = None,
-        api_content: Optional[str] = None,
-        display_kind: Optional[str] = None,
-        display_metadata: Optional[Dict[str, Any]] = None,
-        compression_lock_holder: Optional[str] = None,
-        turn_lease_holder: Optional[str] = None,
+        self, session_id: str, role: str, content: str = None, tool_name: str = None,
+        tool_calls: Any = None, tool_call_id: str = None, token_count: int = None,
+        finish_reason: str = None, reasoning: str = None, reasoning_content: str = None,
+        reasoning_details: Any = None, codex_reasoning_items: Any = None,
+        codex_message_items: Any = None, platform_message_id: str = None, observed: bool = False,
+        effect_disposition: Optional[str] = None, _compressed_summary: bool = False,
+        timestamp: Any = None, api_content: Optional[str] = None,
+        display_kind: Optional[str] = None, display_metadata: Optional[Dict[str, Any]] = None,
+        compression_lock_holder: Optional[str] = None, turn_lease_holder: Optional[str] = None,
        turn_lease_ttl_seconds: float = 300.0,
    ) -> int:
        """Append one message; returns the row id. Bumps ``message_count`` (and
@@ -394,13 +374,9 @@ class SessionMessagesMixin:
        return self._execute_write(_do, patience_s=self._TRANSCRIPT_WRITE_PATIENCE_S)

    def append_messages_batch(
-        self,
-        session_id: str,
-        messages: List[Dict[str, Any]],
-        compression_lock_holder: Optional[str] = None,
-        turn_lease_holder: Optional[str] = None,
-        chunk_rows: Optional[int] = None,
-        turn_lease_ttl_seconds: float = 300.0,
+        self, session_id: str, messages: List[Dict[str, Any]],
+        compression_lock_holder: Optional[str] = None, turn_lease_holder: Optional[str] = None,
+        chunk_rows: Optional[int] = None, turn_lease_ttl_seconds: float = 300.0,
    ) -> int:
        """Append *messages* (``_insert_message_rows`` dict shape) in ONE write txn.

@@ -638,12 +614,8 @@ class SessionMessagesMixin:
        return inserted, tool_calls_total

    def replace_messages(
-        self,
-        session_id: str,
-        messages: List[Dict[str, Any]],
-        active_only: bool = False,
-        archive_dropped: bool = False,
-        reject_active_turn_lease: bool = False,
+        self, session_id: str, messages: List[Dict[str, Any]], active_only: bool = False,
+        archive_dropped: bool = False, reject_active_turn_lease: bool = False,
    ) -> None:
        """Atomically replace the stored messages for a session (/retry, /undo, /compress).

@@ -733,13 +705,9 @@ class SessionMessagesMixin:
            )

    def archive_and_compact(
-        self,
-        session_id: str,
-        compacted_messages: List[Dict[str, Any]],
-        model_config_patch: Optional[Dict[str, Any]] = None,
-        watermark: Optional[int] = None,
-        lock_holder: Optional[str] = None,
-        tail_count: int = 0,
+        self, session_id: str, compacted_messages: List[Dict[str, Any]],
+        model_config_patch: Optional[Dict[str, Any]] = None, watermark: Optional[int] = None,
+        lock_holder: Optional[str] = None, tail_count: int = 0,
    ) -> int:
        """Non-destructive in-place compaction under ONE durable session id.

@@ -940,13 +908,8 @@ class SessionMessagesMixin:
        return _DISPLAY_ACTIVE_CLAUSE if include_compacted else " AND active = 1"

    def get_messages(
-        self,
-        session_id: str,
-        include_inactive: bool = False,
-        include_compacted: bool = False,
-        limit: Optional[int] = None,
-        offset: int = 0,
-        latest: bool = False,
+        self, session_id: str, include_inactive: bool = False, include_compacted: bool = False,
+        limit: Optional[int] = None, offset: int = 0, latest: bool = False,
        after_id: Optional[int] = None,
    ) -> List[Dict[str, Any]]:
        """Load messages for a session in insertion order (AUTOINCREMENT id, never
@@ -1117,12 +1080,8 @@ class SessionMessagesMixin:
            ).fetchall()

    def get_messages_as_conversation(
-        self,
-        session_id: str,
-        include_ancestors: bool = False,
-        include_inactive: bool = False,
-        repair_alternation: bool = False,
-        include_row_ids: bool = False,
+        self, session_id: str, include_ancestors: bool = False, include_inactive: bool = False,
+        repair_alternation: bool = False, include_row_ids: bool = False,
        include_compacted: bool = False,
    ) -> List[Dict[str, Any]]:
        """Load messages in OpenAI conversation format (gateway history restore).
@@ -1176,14 +1135,8 @@ class SessionMessagesMixin:
        return False, exact_clone_key

    def _rows_to_conversation(
-        self,
-        rows,
-        *,
-        session_id: str,
-        include_ancestors: bool,
-        repair_alternation: bool,
-        include_row_ids: bool = False,
-        include_summary_markers: bool = False,
+        self, rows, *, session_id: str, include_ancestors: bool, repair_alternation: bool,
+        include_row_ids: bool = False, include_summary_markers: bool = False,
    ) -> List[Dict[str, Any]]:
        """Decode fetched message rows (ordered by id, pre-filtered) into OpenAI format.

@@ -1486,13 +1439,8 @@ class SessionMessagesMixin:
        return handoff if preserve_compaction_handoff else None

    def rewind_to_message(
-        self,
-        session_id: str,
-        target_message_id: int,
-        *,
-        preserve_compaction_handoff: bool = False,
-        expected_active_ids: Optional[List[int]] = None,
-        expected_target_content: Any = None,
+        self, session_id: str, target_message_id: int, *, preserve_compaction_handoff: bool = False,
+        expected_active_ids: Optional[List[int]] = None, expected_target_content: Any = None,
    ) -> Dict[str, Any]:
        """Soft-delete (``active=0``) every message with id >= *target_message_id*.

@@ -279,24 +279,13 @@ class SessionUsageMixin:
            pass  # never fatal at interpreter shutdown

    def update_token_counts(
-        self,
-        session_id: str,
-        input_tokens: int = 0,
-        output_tokens: int = 0,
-        model: str = None,
-        cache_read_tokens: int = 0,
-        cache_write_tokens: int = 0,
-        reasoning_tokens: int = 0,
-        estimated_cost_usd: Optional[float] = None,
-        actual_cost_usd: Optional[float] = None,
-        cost_status: Optional[str] = None,
-        cost_source: Optional[str] = None,
-        pricing_version: Optional[str] = None,
-        billing_provider: Optional[str] = None,
-        billing_base_url: Optional[str] = None,
-        billing_mode: Optional[str] = None,
-        api_call_count: int = 0,
-        absolute: bool = False,
+        self, session_id: str, input_tokens: int = 0, output_tokens: int = 0, model: str = None,
+        cache_read_tokens: int = 0, cache_write_tokens: int = 0, reasoning_tokens: int = 0,
+        estimated_cost_usd: Optional[float] = None, actual_cost_usd: Optional[float] = None,
+        cost_status: Optional[str] = None, cost_source: Optional[str] = None,
+        pricing_version: Optional[str] = None, billing_provider: Optional[str] = None,
+        billing_base_url: Optional[str] = None, billing_mode: Optional[str] = None,
+        api_call_count: int = 0, absolute: bool = False,
    ) -> None:
        """Update token counters and backfill model if unset. *absolute*=False
        increments (per-API-call deltas, CLI path); *absolute*=True sets directly
@@ -315,10 +304,8 @@ class SessionUsageMixin:
            estimated_cost_usd, actual_cost_usd, actual_cost_usd, cost_status, cost_source, pricing_version,
            billing_provider if has_accounted_usage else None,
            billing_base_url if has_accounted_usage else None,
-            billing_mode if has_accounted_usage else None,
-            model if has_accounted_usage else None,
-            api_call_count,
-            session_id,
+            billing_mode if has_accounted_usage else None, model if has_accounted_usage else None,
+            api_call_count, session_id,
        )
        # Per-model attribution: the sessions row keeps one (model, provider) pair, so a
        # mid-session /model switch would attribute every token to the initial model.
@@ -365,25 +352,11 @@ class SessionUsageMixin:
        self._execute_write(_do)

    def _record_model_usage(
-        self,
-        conn,
-        session_id: str,
-        *,
-        model: Optional[str],
-        billing_provider: Optional[str],
-        billing_base_url: Optional[str],
-        billing_mode: Optional[str],
-        input_tokens: int,
-        output_tokens: int,
-        cache_read_tokens: int,
-        cache_write_tokens: int,
-        reasoning_tokens: int,
-        estimated_cost_usd: Optional[float],
-        actual_cost_usd: Optional[float],
-        cost_status: Optional[str],
-        cost_source: Optional[str],
-        api_call_count: int,
-        task: str = "",
+        self, conn, session_id: str, *, model: Optional[str], billing_provider: Optional[str],
+        billing_base_url: Optional[str], billing_mode: Optional[str], input_tokens: int,
+        output_tokens: int, cache_read_tokens: int, cache_write_tokens: int, reasoning_tokens: int,
+        estimated_cost_usd: Optional[float], actual_cost_usd: Optional[float],
+        cost_status: Optional[str], cost_source: Optional[str], api_call_count: int, task: str = "",
    ) -> None:
        """Accumulate a per-API-call usage delta into session_model_usage, inside the
        caller's write txn after the ``sessions`` UPDATE. A missing model/provider falls
@@ -414,20 +387,11 @@ class SessionUsageMixin:
        )

    def record_auxiliary_usage(
-        self,
-        session_id: str,
-        task: str,
-        *,
-        model: Optional[str] = None,
-        billing_provider: Optional[str] = None,
-        billing_base_url: Optional[str] = None,
-        input_tokens: int = 0,
-        output_tokens: int = 0,
-        cache_read_tokens: int = 0,
-        cache_write_tokens: int = 0,
-        reasoning_tokens: int = 0,
-        estimated_cost_usd: Optional[float] = None,
-        api_call_count: int = 1,
+        self, session_id: str, task: str, *, model: Optional[str] = None,
+        billing_provider: Optional[str] = None, billing_base_url: Optional[str] = None,
+        input_tokens: int = 0, output_tokens: int = 0, cache_read_tokens: int = 0,
+        cache_write_tokens: int = 0, reasoning_tokens: int = 0,
+        estimated_cost_usd: Optional[float] = None, api_call_count: int = 1,
    ) -> None:
        """Record an auxiliary LLM call's usage (vision, compression, title generation,
        ...) as a per-(model, provider, task) delta in ``session_model_usage`` WITHOUT
@@ -447,8 +411,7 @@ class SessionUsageMixin:
                cache_read_tokens=cache_read_tokens or 0, cache_write_tokens=cache_write_tokens or 0,
                reasoning_tokens=reasoning_tokens or 0, estimated_cost_usd=estimated_cost_usd,
                actual_cost_usd=None, cost_status=None, cost_source=None,
-                api_call_count=1 if api_call_count is None else int(api_call_count),
-                task=task,
+                api_call_count=1 if api_call_count is None else int(api_call_count), task=task,
            )
        self._execute_write(_do)