From 5e7e346c74d5903745d6c3211d36b9e186ea6d09 Mon Sep 17 00:00:00 2001
From: LawyZheng <lawyzheng1106@gmail.com>
Date: Thu, 9 Oct 2025 17:32:51 +0800
Subject: [PATCH] add organization_id and organization_name to the tracing span
 in a http request (#3657)

---
 .../forge/sdk/services/org_auth_service.py    | 31 +++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/skyvern/forge/sdk/services/org_auth_service.py b/skyvern/forge/sdk/services/org_auth_service.py
index 2b42cb46..92bcd35d 100644
--- a/skyvern/forge/sdk/services/org_auth_service.py
+++ b/skyvern/forge/sdk/services/org_auth_service.py
@@ -4,6 +4,7 @@ from typing import Annotated
 import structlog
 from asyncache import cached
 from cachetools import TTLCache
+from ddtrace import tracer
 from fastapi import Header, HTTPException, status
 from jose import jwt
 from jose.exceptions import JWTError
@@ -37,11 +38,30 @@ async def get_current_org(
             status_code=status.HTTP_403_FORBIDDEN,
             detail="Invalid credentials",
         )
+    organization = None
     if x_api_key:
-        return await _get_current_org_cached(x_api_key, app.DATABASE)
+        organization = await _get_current_org_cached(x_api_key, app.DATABASE)
     elif authorization:
-        return await _authenticate_helper(authorization)
+        organization = await _authenticate_helper(authorization)
 
+    if organization:
+        try:
+            # Set in context
+            curr_ctx = skyvern_context.current()
+            if curr_ctx:
+                curr_ctx.organization_id = organization.organization_id
+                curr_ctx.organization_name = organization.organization_name
+
+            # Tag datadog span immediately
+            span = tracer.current_span()
+            if span:
+                span.set_tag("organization_id", organization.organization_id)
+                if organization.organization_name:
+                    span.set_tag("organization_name", organization.organization_name)
+        except Exception:
+            pass
+
+        return organization
     raise HTTPException(
         status_code=status.HTTP_403_FORBIDDEN,
         detail="Invalid credentials",
@@ -83,6 +103,13 @@ async def _authenticate_helper(authorization: str) -> Organization:
             status_code=status.HTTP_403_FORBIDDEN,
             detail="Invalid credentials",
         )
+
+    # set organization_id in skyvern context and log context
+    context = skyvern_context.current()
+    if context:
+        context.organization_id = organization.organization_id
+        context.organization_name = organization.organization_name
+
     return organization