From a79875754169c93d8f591c5020277a214775728d Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Mon, 13 Jan 2025 19:28:14 -0800
Subject: [PATCH] add helper to generate webhook headers (#1543)

---
 skyvern/forge/agent.py                | 13 ++-----------
 skyvern/forge/sdk/core/security.py    | 10 ++++++++++
 skyvern/forge/sdk/workflow/service.py | 11 ++---------
 3 files changed, 14 insertions(+), 20 deletions(-)

diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py
index b7df7687..117ab520 100644
--- a/skyvern/forge/agent.py
+++ b/skyvern/forge/agent.py
@@ -47,7 +47,7 @@ from skyvern.forge.prompts import prompt_engine
 from skyvern.forge.sdk.api.files import get_path_for_workflow_download_directory, list_files_in_directory, rename_file
 from skyvern.forge.sdk.artifact.models import ArtifactType
 from skyvern.forge.sdk.core import skyvern_context
-from skyvern.forge.sdk.core.security import generate_skyvern_signature
+from skyvern.forge.sdk.core.security import generate_skyvern_webhook_headers
 from skyvern.forge.sdk.db.enums import TaskType
 from skyvern.forge.sdk.log_artifacts import save_step_logs, save_task_logs
 from skyvern.forge.sdk.models import Step, StepStatus
@@ -1603,17 +1603,8 @@ class ForgeAgent:
         task_response = await self.build_task_response(task=task, last_step=last_step)
 
         # send task_response to the webhook callback url
-        timestamp = str(int(datetime.utcnow().timestamp()))
         payload = task_response.model_dump_json(exclude={"request"})
-        signature = generate_skyvern_signature(
-            payload=payload,
-            api_key=api_key,
-        )
-        headers = {
-            "x-skyvern-timestamp": timestamp,
-            "x-skyvern-signature": signature,
-            "Content-Type": "application/json",
-        }
+        headers = generate_skyvern_webhook_headers(payload=payload, api_key=api_key)
         LOG.info(
             "Sending task response to webhook callback url",
             task_id=task.task_id,
diff --git a/skyvern/forge/sdk/core/security.py b/skyvern/forge/sdk/core/security.py
index d197cfcb..48a17450 100644
--- a/skyvern/forge/sdk/core/security.py
+++ b/skyvern/forge/sdk/core/security.py
@@ -41,3 +41,13 @@ def generate_skyvern_signature(
     """
     hash_obj = hmac.new(api_key.encode("utf-8"), msg=payload.encode("utf-8"), digestmod=hashlib.sha256)
     return hash_obj.hexdigest()
+
+
+def generate_skyvern_webhook_headers(payload: str, api_key: str) -> dict[str, str]:
+    signature = generate_skyvern_signature(payload=payload, api_key=api_key)
+    timestamp = str(int(datetime.utcnow().timestamp()))
+    return {
+        "x-skyvern-timestamp": timestamp,
+        "x-skyvern-signature": signature,
+        "Content-Type": "application/json",
+    }
diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py
index c4a88d36..35fd7d90 100644
--- a/skyvern/forge/sdk/workflow/service.py
+++ b/skyvern/forge/sdk/workflow/service.py
@@ -1,6 +1,5 @@
 import asyncio
 import json
-from datetime import datetime
 from typing import Any
 
 import httpx
@@ -19,7 +18,7 @@ from skyvern.exceptions import (
 from skyvern.forge import app
 from skyvern.forge.sdk.artifact.models import ArtifactType
 from skyvern.forge.sdk.core import skyvern_context
-from skyvern.forge.sdk.core.security import generate_skyvern_signature
+from skyvern.forge.sdk.core.security import generate_skyvern_webhook_headers
 from skyvern.forge.sdk.core.skyvern_context import SkyvernContext
 from skyvern.forge.sdk.db.enums import TaskType
 from skyvern.forge.sdk.models import Step, StepStatus
@@ -974,17 +973,11 @@ class WorkflowService:
             return
 
         # send webhook to the webhook callback url
-        timestamp = str(int(datetime.utcnow().timestamp()))
         payload = workflow_run_status_response.model_dump_json()
-        signature = generate_skyvern_signature(
+        headers = generate_skyvern_webhook_headers(
             payload=payload,
             api_key=api_key,
         )
-        headers = {
-            "x-skyvern-timestamp": timestamp,
-            "x-skyvern-signature": signature,
-            "Content-Type": "application/json",
-        }
         LOG.info(
             "Sending webhook run status to webhook callback url",
             workflow_id=workflow_id,