add functionality to cache task_run (#1755)

skyvern/forge/sdk/db/client.py

@@ -2672,3 +2672,30 @@ class AgentDB:
             await session.commit()
             await session.refresh(task_run)
             return TaskRun.model_validate(task_run)
+
+    async def cache_task_run(self, run_id: str, organization_id: str | None = None) -> TaskRun:
+        async with self.Session() as session:
+            task_run = await session.scalars(
+                select(TaskRunModel).filter_by(organization_id=organization_id).filter_by(run_id=run_id)
+            ).first()
+            if task_run:
+                task_run.cached = True
+                await session.commit()
+                await session.refresh(task_run)
+                return TaskRun.model_validate(task_run)
+            raise NotFoundError(f"TaskRun {run_id} not found")
+
+    async def get_cached_task_run(
+        self, task_run_type: TaskRunType, url_hash: str | None = None, organization_id: str | None = None
+    ) -> TaskRun | None:
+        async with self.Session() as session:
+            query = select(TaskRunModel)
+            if task_run_type:
+                query = query.filter_by(task_run_type=task_run_type)
+            if url_hash:
+                query = query.filter_by(url_hash=url_hash)
+            if organization_id:
+                query = query.filter_by(organization_id=organization_id)
+            query = query.filter_by(cached=True).order_by(TaskRunModel.created_at.desc())
+            task_run = await session.scalars(query).first()
+            return TaskRun.model_validate(task_run) if task_run else None

skyvern/forge/sdk/db/models.py

@@ -614,7 +614,10 @@ class PersistentBrowserSessionModel(Base):
 
 class TaskRunModel(Base):
     __tablename__ = "task_runs"
-    __table_args__ = (Index("task_run_org_url_index", "organization_id", "url_hash", "cached"),)
+    __table_args__ = (
+        Index("task_run_org_url_index", "organization_id", "url_hash", "cached"),
+        Index("task_run_org_run_id_index", "organization_id", "run_id"),
+    )
 
     task_run_id = Column(String, primary_key=True, default=generate_task_run_id)
     organization_id = Column(String, nullable=False)