From 4f401e81b61b8e26ea16ad3066956398f975b8ee Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Wed, 2 Apr 2025 17:03:39 -0400
Subject: [PATCH] support extra_headers in SkyvernAgent (#2080)

---
 skyvern/agent/agent.py  | 4 +++-
 skyvern/agent/client.py | 9 ++++++++-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/skyvern/agent/agent.py b/skyvern/agent/agent.py
index 369186b1..3a107135 100644
--- a/skyvern/agent/agent.py
+++ b/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")
diff --git a/skyvern/agent/client.py b/skyvern/agent/client.py
index c55cbcee..393ce2de 100644
--- a/skyvern/agent/client.py
+++ b/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())