From c939513ff75f9a4b26bcf582c7e570087ec40a27 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Tue, 9 Dec 2025 15:23:47 -0800
Subject: [PATCH] temporal_v2: job cost reporting (#4240)

---
 ...4da0203e_add_columns_to_task_runs_table.py | 37 +++++++++++++++++++
 skyvern/forge/sdk/db/client.py                | 34 +++++++++++++++++
 skyvern/forge/sdk/db/models.py                |  6 +++
 skyvern/forge/sdk/schemas/runs.py             |  5 +++
 4 files changed, 82 insertions(+)
 create mode 100644 alembic/versions/2025_12_09_2308-67784da0203e_add_columns_to_task_runs_table.py

diff --git a/alembic/versions/2025_12_09_2308-67784da0203e_add_columns_to_task_runs_table.py b/alembic/versions/2025_12_09_2308-67784da0203e_add_columns_to_task_runs_table.py
new file mode 100644
index 00000000..97edd769
--- /dev/null
+++ b/alembic/versions/2025_12_09_2308-67784da0203e_add_columns_to_task_runs_table.py
@@ -0,0 +1,37 @@
+"""add columns to task_runs table
+
+Revision ID: 67784da0203e
+Revises: 1faa2a5869cd
+Create Date: 2025-12-09 23:08:29.686747+00:00
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "67784da0203e"
+down_revision: Union[str, None] = "1faa2a5869cd"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("task_runs", sa.Column("instance_type", sa.String(), nullable=True))
+    op.add_column("task_runs", sa.Column("vcpu_millicores", sa.Integer(), nullable=True))
+    op.add_column("task_runs", sa.Column("duration_ms", sa.BigInteger(), nullable=True))
+    op.add_column("task_runs", sa.Column("compute_cost", sa.Numeric(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("task_runs", "compute_cost")
+    op.drop_column("task_runs", "duration_ms")
+    op.drop_column("task_runs", "vcpu_millicores")
+    op.drop_column("task_runs", "instance_type")
+    # ### end Alembic commands ###
diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py
index 4b76a644..b3c59a7e 100644
--- a/skyvern/forge/sdk/db/client.py
+++ b/skyvern/forge/sdk/db/client.py
@@ -4561,6 +4561,40 @@ class AgentDB:
                 task_run.url_hash = url_hash
             await session.commit()
 
+    async def update_job_run_compute_cost(
+        self,
+        organization_id: str,
+        run_id: str,
+        instance_type: str | None = None,
+        vcpu_millicores: int | None = None,
+        duration_ms: int | None = None,
+        compute_cost: float | None = None,
+    ) -> None:
+        """Update compute cost metrics for a job run."""
+        async with self.Session() as session:
+            task_run = (
+                await session.scalars(
+                    select(TaskRunModel).filter_by(run_id=run_id).filter_by(organization_id=organization_id)
+                )
+            ).first()
+            if not task_run:
+                LOG.warning(
+                    "TaskRun not found for compute cost update",
+                    run_id=run_id,
+                    organization_id=organization_id,
+                )
+                return
+
+            if instance_type is not None:
+                task_run.instance_type = instance_type
+            if vcpu_millicores is not None:
+                task_run.vcpu_millicores = vcpu_millicores
+            if duration_ms is not None:
+                task_run.duration_ms = duration_ms
+            if compute_cost is not None:
+                task_run.compute_cost = compute_cost
+            await session.commit()
+
     async def create_credential(
         self,
         organization_id: str,
diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py
index 83159759..83c29106 100644
--- a/skyvern/forge/sdk/db/models.py
+++ b/skyvern/forge/sdk/db/models.py
@@ -3,6 +3,7 @@ import datetime
 import sqlalchemy
 from sqlalchemy import (
     JSON,
+    BigInteger,
     Boolean,
     Column,
     DateTime,
@@ -867,6 +868,11 @@ class TaskRunModel(Base):
     url = Column(String, nullable=True)
     url_hash = Column(String, nullable=True)
     cached = Column(Boolean, nullable=False, default=False)
+    # Compute cost tracking fields
+    instance_type = Column(String, nullable=True)
+    vcpu_millicores = Column(Integer, nullable=True)
+    duration_ms = Column(BigInteger, nullable=True)
+    compute_cost = Column(Numeric, nullable=True)
     created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
     modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
 
diff --git a/skyvern/forge/sdk/schemas/runs.py b/skyvern/forge/sdk/schemas/runs.py
index 409baa78..dab90a49 100644
--- a/skyvern/forge/sdk/schemas/runs.py
+++ b/skyvern/forge/sdk/schemas/runs.py
@@ -15,5 +15,10 @@ class Run(BaseModel):
     title: str | None = None
     url: str | None = None
     cached: bool = False
+    # Compute cost tracking fields
+    instance_type: str | None = None
+    vcpu_millicores: int | None = None
+    duration_ms: int | None = None
+    compute_cost: float | None = None
     created_at: datetime
     modified_at: datetime