New Billing - per action credit consumption (#4398)

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:

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:
     """