From 4adadee82a84c5a8a9c861f22c99bbf9609778e2 Mon Sep 17 00:00:00 2001
From: Marc Kelechava <marc@skyvern.com>
Date: Mon, 5 Jan 2026 17:43:05 -0800
Subject: [PATCH] New Billing - per action credit consumption (#4398)

---
 skyvern/forge/sdk/db/agent_db.py | 20 ++++++++++++++++++++
 skyvern/forge/sdk/db/id.py       |  6 ++++++
 2 files changed, 26 insertions(+)

diff --git a/skyvern/forge/sdk/db/agent_db.py b/skyvern/forge/sdk/db/agent_db.py
index 34a417f8..a1660d4f 100644
--- a/skyvern/forge/sdk/db/agent_db.py
+++ b/skyvern/forge/sdk/db/agent_db.py
@@ -560,6 +560,26 @@ class AgentDB(BaseAlchemyDB):
             LOG.error("UnexpectedError", exc_info=True)
             raise
 
+    async def get_action_count_for_step(self, step_id: str, task_id: str, organization_id: str) -> int:
+        """Get count of actions for a step. Uses composite index for efficiency."""
+        try:
+            async with self.Session() as session:
+                query = (
+                    select(func.count())
+                    .select_from(ActionModel)
+                    .where(ActionModel.organization_id == organization_id)
+                    .where(ActionModel.task_id == task_id)
+                    .where(ActionModel.step_id == step_id)
+                )
+                result = await session.scalar(query)
+                return result or 0
+        except SQLAlchemyError:
+            LOG.error("SQLAlchemyError", exc_info=True)
+            raise
+        except Exception:
+            LOG.error("UnexpectedError", exc_info=True)
+            raise
+
     async def get_first_step(self, task_id: str, organization_id: str | None = None) -> Step | None:
         try:
             async with self.Session() as session:
diff --git a/skyvern/forge/sdk/db/id.py b/skyvern/forge/sdk/db/id.py
index 1af0bafd..98c8f917 100644
--- a/skyvern/forge/sdk/db/id.py
+++ b/skyvern/forge/sdk/db/id.py
@@ -68,6 +68,7 @@ WORKFLOW_RUN_BLOCK_PREFIX = "wrb"
 WORKFLOW_RUN_PREFIX = "wr"
 WORKFLOW_SCRIPT_PREFIX = "ws"
 WORKFLOW_TEMPLATE_PREFIX = "wt"
+ORGANIZATION_BILLING_PREFIX = "ob"
 
 
 def generate_workflow_id() -> str:
@@ -260,6 +261,11 @@ def generate_script_block_id() -> str:
     return f"{SCRIPT_BLOCK_PREFIX}_{int_id}"
 
 
+def generate_billing_id() -> str:
+    int_id = generate_id()
+    return f"{ORGANIZATION_BILLING_PREFIX}_{int_id}"
+
+
 ############# Helper functions below ##############
 def generate_id() -> int:
     """