add new tool: to_captcha

api/services/task_service.py

@@ -89,6 +89,93 @@ class TaskService:
     async def unsubscribe_run_stream(self, run_id: str, queue) -> None:
         await self._store.unsubscribe(run_id, queue)
 
+    async def notify_captcha(
+            self,
+            task_id: str,
+            kind: str | None,
+            reason: str | None,
+            view_url: str | None,
+            timeout_seconds: int,
+    ):
+        rec = await self._store.set_captcha_awaiting(
+            task_id=task_id,
+            kind=kind,
+            reason=reason,
+            view_url=view_url,
+            timeout_seconds=timeout_seconds,
+        )
+        if rec is not None:
+            await self._store.publish(
+                task_id,
+                self._event(task_id, "captcha_required", {
+                    "captcha_kind": rec.captcha_kind,
+                    "reason": rec.captcha_reason,
+                    "browser_view_url": rec.captcha_view_url,
+                    "deadline": rec.captcha_deadline,
+                    "timeout_seconds": timeout_seconds,
+                }),
+            )
+        return rec
+
+    async def mark_captcha_solved(self, task_id: str, detector: str | None = None):
+        rec = await self._store.set_captcha_solved(task_id)
+        if rec is not None:
+            await self._store.publish(
+                task_id,
+                self._event(task_id, "captcha_solved", {
+                    "detector": detector or "unknown",
+                    "solved_at": rec.captcha_solved_at,
+                }),
+            )
+        return rec
+
+    async def extend_captcha(self, task_id: str, extra_seconds: int):
+        rec = await self._store.set_captcha_extended(task_id=task_id, extra_seconds=extra_seconds)
+        if rec is not None:
+            await self._store.publish(
+                task_id,
+                self._event(task_id, "captcha_extended", {
+                    "extra_seconds": extra_seconds,
+                    "deadline": rec.captcha_deadline,
+                }),
+            )
+        return rec
+
+    async def prompt_captcha_timeout(self, task_id: str):
+        rec = await self._store.set_captcha_timeout_prompt(task_id)
+        if rec is not None:
+            await self._store.publish(
+                task_id,
+                self._event(task_id, "captcha_timeout_prompt", {
+                    "captcha_kind": rec.captcha_kind,
+                    "browser_view_url": rec.captcha_view_url,
+                    "deadline": rec.captcha_deadline,
+                    "actions": ["extend", "abort"],
+                }),
+            )
+        return rec
+
+    async def abort_captcha(self, task_id: str, reason: str | None = None):
+        rec = await self._store.set_captcha_aborted(task_id=task_id, reason=reason)
+        if rec is not None:
+            await self._store.publish(
+                task_id,
+                self._event(task_id, "captcha_aborted", {
+                    "reason": rec.captcha_reason,
+                }),
+            )
+        return rec
+
+    async def wait_captcha(self, task_id: str, timeout: float):
+        rec = await self._store.get(task_id)
+        if rec is None:
+            return None
+        try:
+            await asyncio.wait_for(rec.captcha_event.wait(), timeout=timeout)
+        except asyncio.TimeoutError:
+            pass
+        return await self._store.get(task_id)
+
     async def close(self) -> None:
         if not self._background_tasks:
             return
@@ -127,7 +214,12 @@ class TaskService:
                     rpc_timeout = min(rpc_timeout, self._rpc_timeout_cap)
 
                 raw = await asyncio.wait_for(
-                    self._rpc_client.run(task=rec.task, timeout_sec=rpc_timeout, rpc_url=runtime.get("rpc_url")),
+                    self._rpc_client.run(
+                        task=rec.task,
+                        timeout_sec=rpc_timeout,
+                        rpc_url=runtime.get("rpc_url"),
+                        task_id=task_id,
+                    ),
                     timeout=float(rec.timeout) + 5,
                 )
                 raw = self._with_runtime_metadata(raw, runtime)