support extra_headers in SkyvernAgent (#2080)

skyvern/agent/agent.py

@@ -30,7 +30,9 @@ class SkyvernAgent:
         cdp_url: str | None = None,
         browser_path: str | None = None,
         browser_type: str | None = None,
+        extra_headers: dict[str, str] | None = None,
     ) -> None:
+        self.extra_headers = extra_headers
         self.client: SkyvernClient | None = None
         if base_url is None and api_key is None:
             # TODO: run at the root wherever the code is initiated
@@ -69,7 +71,7 @@ class SkyvernAgent:
             self.client = SkyvernClient(
                 base_url=base_url,
                 api_key=api_key,
-                extra_headers={"X-User-Agent": "skyvern-mcp"},
+                extra_headers=self.extra_headers,
             )
         else:
             raise ValueError("base_url and api_key must be both provided")

skyvern/agent/client.py

@@ -13,7 +13,12 @@ class SkyvernClient:
         self.base_url = base_url
         self.api_key = api_key
         self.client = AsyncSkyvern(base_url=base_url, api_key=api_key)
-        self.extra_headers = extra_headers
+        self.extra_headers = extra_headers or {}
+        self.user_agent = None
+        if "X-User-Agent" in self.extra_headers:
+            self.user_agent = self.extra_headers["X-User-Agent"]
+        elif "x-user-agent" in self.extra_headers:
+            self.user_agent = self.extra_headers["x-user-agent"]
 
     async def run_task(
         self,
@@ -43,6 +48,7 @@ class SkyvernClient:
             max_steps=max_steps,
             browser_session_id=browser_session_id,
             publish_workflow=publish_workflow,
+            user_agent=self.user_agent,
         )
         return TaskRunResponse.model_validate(task_run_obj.dict())
 
@@ -66,6 +72,7 @@ class SkyvernClient:
             totp_url=totp_url,
             browser_session_id=browser_session_id,
             template=template,
+            user_agent=self.user_agent,
         )
         return WorkflowRunResponse.model_validate(workflow_run_obj.dict())