fix: Logging fix for python 3.13 (#1745)

* t Logging fix for python 3.13

* Flake8 fixes

* Added unit tests

* Reverting uv change

* Flake8 fixes

* Minor fix

Author: gavin-aguiar

workers/proxy_worker/dispatcher.py

@@ -40,6 +40,14 @@
 # Library worker import reloaded in init and reload request
 _library_worker = None
 
+# Thread-local invocation ID registry for efficient lookup
+_thread_invocation_registry: typing.Dict[int, str] = {}
+_registry_lock = threading.Lock()
+
+# Global current invocation tracker (as a fallback)
+_current_invocation_id: Optional[str] = None
+_current_invocation_lock = threading.Lock()
+
 
 class ContextEnabledTask(asyncio.Task):
     AZURE_INVOCATION_ID = '__azure_function_invocation_id__'
@@ -61,16 +69,63 @@ def set_azure_invocation_id(self, invocation_id: str) -> None:
 _invocation_id_local = threading.local()
 
 
+def set_thread_invocation_id(thread_id: int, invocation_id: str) -> None:
+    """Set the invocation ID for a specific thread"""
+    with _registry_lock:
+        _thread_invocation_registry[thread_id] = invocation_id
+
+
+def clear_thread_invocation_id(thread_id: int) -> None:
+    """Clear the invocation ID for a specific thread"""
+    with _registry_lock:
+        _thread_invocation_registry.pop(thread_id, None)
+
+
+def get_thread_invocation_id(thread_id: int) -> Optional[str]:
+    """Get the invocation ID for a specific thread"""
+    with _registry_lock:
+        return _thread_invocation_registry.get(thread_id)
+
+
+def set_current_invocation_id(invocation_id: str) -> None:
+    """Set the global current invocation ID"""
+    global _current_invocation_id
+    with _current_invocation_lock:
+        _current_invocation_id = invocation_id
+
+
+def get_global_current_invocation_id() -> Optional[str]:
+    """Get the global current invocation ID"""
+    with _current_invocation_lock:
+        return _current_invocation_id
+
+
 def get_current_invocation_id() -> Optional[Any]:
-    loop = asyncio._get_running_loop()
-    if loop is not None:
-        current_task = asyncio.current_task(loop)
-        if current_task is not None:
-            task_invocation_id = getattr(current_task,
-                                         ContextEnabledTask.AZURE_INVOCATION_ID,
-                                         None)
-            if task_invocation_id is not None:
-                return task_invocation_id
+    # Check global current invocation first (most up-to-date)
+    global_invocation_id = get_global_current_invocation_id()
+    if global_invocation_id is not None:
+        return global_invocation_id
+
+    # Check asyncio task context
+    try:
+        loop = asyncio._get_running_loop()
+        if loop is not None:
+            current_task = asyncio.current_task(loop)
+            if current_task is not None:
+                task_invocation_id = getattr(current_task,
+                                             ContextEnabledTask.AZURE_INVOCATION_ID,
+                                             None)
+                if task_invocation_id is not None:
+                    return task_invocation_id
+    except RuntimeError:
+        # No event loop running
+        pass
+
+    # Check the thread-local invocation ID registry
+    current_thread_id = threading.get_ident()
+    thread_invocation_id = get_thread_invocation_id(current_thread_id)
+    if thread_invocation_id is not None:
+        return thread_invocation_id
 
     return getattr(_invocation_id_local, 'invocation_id', None)
 
@@ -516,14 +571,32 @@ async def _handle__invocation_request(self, request):
             'invocation_id: %s, worker_id: %s',
             self.request_id, function_id, invocation_id, self.worker_id)
 
-        invocation_request = WorkerRequest(name="FunctionInvocationRequest",
-                                           request=request,
-                                           properties={
-                                               "threadpool": self._sync_call_tp})
-        invocation_response = await (
-            _library_worker.invocation_request(  # type: ignore[union-attr]
-                invocation_request))
+        # Set the global current invocation ID first (for all threads to access)
+        set_current_invocation_id(invocation_id)
 
-        return protos.StreamingMessage(
-            request_id=self.request_id,
-            invocation_response=invocation_response)
+        # Set the current `invocation_id` to the current task so
+        # that our logging handler can find it.
+        current_task = asyncio.current_task()
+        if current_task is not None and isinstance(current_task, ContextEnabledTask):
+            current_task.set_azure_invocation_id(invocation_id)
+
+        # Register the invocation ID for the current thread
+        current_thread_id = threading.get_ident()
+        set_thread_invocation_id(current_thread_id, invocation_id)
+
+        try:
+            invocation_request = WorkerRequest(name="FunctionInvocationRequest",
+                                               request=request,
+                                               properties={
+                                                   "threadpool": self._sync_call_tp})
+            invocation_response = await (
+                _library_worker.invocation_request(  # type: ignore[union-attr]
+                    invocation_request))
+
+            return protos.StreamingMessage(
+                request_id=self.request_id,
+                invocation_response=invocation_response)
+        except Exception:
+            # Clear thread registry on exception to prevent stale IDs
+            clear_thread_invocation_id(current_thread_id)
+            raise