gobby 0.2.5__py3-none-any.whl → 0.2.6__py3-none-any.whl

gobby/storage/sessions.py

@@ -47,6 +47,7 @@ class Session:
     usage_cache_creation_tokens: int = 0
     usage_cache_read_tokens: int = 0
     usage_total_cost_usd: float = 0.0
+    model: str | None = None  # LLM model used (e.g., "claude-3-5-sonnet-20241022")
     # Terminal context (JSON blob with tty, parent_pid, term_session_id, etc.)
     terminal_context: dict[str, Any] | None = None
     # Global sequence number
@@ -82,6 +83,7 @@ class Session:
             usage_cache_creation_tokens=row["usage_cache_creation_tokens"] or 0,
             usage_cache_read_tokens=row["usage_cache_read_tokens"] or 0,
             usage_total_cost_usd=row["usage_total_cost_usd"] or 0.0,
+            model=row["model"] if "model" in row.keys() else None,
             terminal_context=cls._parse_terminal_context(row["terminal_context"]),
             seq_num=row["seq_num"] if "seq_num" in row.keys() else None,
         )
@@ -773,6 +775,77 @@ class LocalSessionManager:
             logger.error(f"Failed to update session usage {session_id}: {e}")
             return False
 
+    def add_cost(self, session_id: str, cost_usd: float) -> bool:
+        """
+        Add cost to the session's usage_total_cost_usd.
+
+        This is used for internal agent runs that track cost via CostInfo.
+        Unlike update_usage which overwrites, this method adds to the existing cost.
+
+        Args:
+            session_id: Session ID to update.
+            cost_usd: Cost in USD to add.
+
+        Returns:
+            True if update succeeded, False otherwise.
+        """
+        if cost_usd <= 0:
+            return True  # Nothing to add
+
+        query = """
+        UPDATE sessions
+        SET
+            usage_total_cost_usd = COALESCE(usage_total_cost_usd, 0) + ?,
+            updated_at = datetime('now')
+        WHERE id = ?
+        """
+        try:
+            with self.db.transaction():
+                cursor = self.db.execute(query, (cost_usd, session_id))
+                return cursor.rowcount > 0
+        except Exception as e:
+            logger.error(f"Failed to add cost to session {session_id}: {e}")
+            return False
+
+    def get_sessions_since(
+        self, since: datetime, project_id: str | None = None
+    ) -> builtins.list[Session]:
+        """
+        Get sessions created since a given timestamp.
+
+        Used for aggregating usage over a time period (e.g., daily budget tracking).
+
+        Args:
+            since: Datetime to query from (sessions created after this time)
+            project_id: Optional project ID to filter by
+
+        Returns:
+            List of sessions created since the given timestamp
+        """
+        since_str = since.isoformat()
+
+        if project_id:
+            rows = self.db.fetchall(
+                """
+                SELECT * FROM sessions
+                WHERE created_at >= ?
+                AND project_id = ?
+                ORDER BY created_at DESC
+                """,
+                (since_str, project_id),
+            )
+        else:
+            rows = self.db.fetchall(
+                """
+                SELECT * FROM sessions
+                WHERE created_at >= ?
+                ORDER BY created_at DESC
+                """,
+                (since_str,),
+            )
+
+        return [Session.from_row(row) for row in rows]
+
     def update_terminal_pickup_metadata(
         self,
         session_id: str,