add task_runs (#1752)

skyvern/forge/sdk/db/client.py

@@ -29,6 +29,7 @@ from skyvern.forge.sdk.db.models import (
     StepModel,
     TaskGenerationModel,
     TaskModel,
+    TaskRunModel,
     TOTPCodeModel,
     WorkflowModel,
     WorkflowParameterModel,
@@ -62,6 +63,7 @@ from skyvern.forge.sdk.schemas.observers import ObserverTask, ObserverTaskStatus
 from skyvern.forge.sdk.schemas.organizations import Organization, OrganizationAuthToken
 from skyvern.forge.sdk.schemas.persistent_browser_sessions import PersistentBrowserSession
 from skyvern.forge.sdk.schemas.task_generations import TaskGeneration
+from skyvern.forge.sdk.schemas.task_runs import TaskRun, TaskRunType
 from skyvern.forge.sdk.schemas.tasks import OrderBy, ProxyLocation, SortDirection, Task, TaskStatus
 from skyvern.forge.sdk.schemas.totp_codes import TOTPCode
 from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunBlock
@@ -2647,3 +2649,26 @@ class AgentDB:
         except Exception:
             LOG.error("UnexpectedError", exc_info=True)
             raise
+
+    async def create_task_run(
+        self,
+        task_run_type: TaskRunType,
+        organization_id: str,
+        run_id: str,
+        title: str | None = None,
+        url: str | None = None,
+        url_hash: str | None = None,
+    ) -> TaskRun:
+        async with self.Session() as session:
+            task_run = TaskRunModel(
+                task_run_type=task_run_type,
+                organization_id=organization_id,
+                run_id=run_id,
+                title=title,
+                url=url,
+                url_hash=url_hash,
+            )
+            session.add(task_run)
+            await session.commit()
+            await session.refresh(task_run)
+            return TaskRun.model_validate(task_run)

skyvern/forge/sdk/db/id.py

@@ -43,6 +43,7 @@ PERSISTENT_BROWSER_SESSION_ID = "pbs"
 STEP_PREFIX = "stp"
 TASK_GENERATION_PREFIX = "tg"
 TASK_PREFIX = "tsk"
+TASK_RUN_PREFIX = "tr"
 TOTP_CODE_PREFIX = "totp"
 USER_PREFIX = "u"
 WORKFLOW_PARAMETER_PREFIX = "wp"
@@ -167,6 +168,11 @@ def generate_persistent_browser_session_id() -> str:
     return f"{PERSISTENT_BROWSER_SESSION_ID}_{int_id}"
 
 
+def generate_task_run_id() -> str:
+    int_id = generate_id()
+    return f"{TASK_RUN_PREFIX}_{int_id}"
+
+
 def generate_id() -> int:
     """
     generate a 64-bit int ID

skyvern/forge/sdk/db/models.py

@@ -36,6 +36,7 @@ from skyvern.forge.sdk.db.id import (
     generate_step_id,
     generate_task_generation_id,
     generate_task_id,
+    generate_task_run_id,
     generate_totp_code_id,
     generate_workflow_id,
     generate_workflow_parameter_id,
@@ -609,3 +610,19 @@ class PersistentBrowserSessionModel(Base):
     created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
     modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
     deleted_at = Column(DateTime, nullable=True)
+
+
+class TaskRunModel(Base):
+    __tablename__ = "task_runs"
+    __table_args__ = (Index("task_run_org_url_index", "organization_id", "url_hash", "cached"),)
+
+    task_run_id = Column(String, primary_key=True, default=generate_task_run_id)
+    organization_id = Column(String, nullable=False)
+    task_run_type = Column(String, nullable=False)
+    run_id = Column(String, nullable=False)
+    title = Column(String, nullable=True)
+    url = Column(String, nullable=True)
+    url_hash = Column(String, nullable=True)
+    cached = Column(Boolean, nullable=False, default=False)
+    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
+    modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)