Add async hook support to hook infrastructure

awalker4 · claude · awalker4 · commit 674653952293 · 2026-02-09T16:34:05.000-05:00
- Add AsyncSDKInitHook, AsyncBeforeRequestHook, AsyncAfterSuccessHook, and AsyncAfterErrorHook interfaces
- Update SDKHooks to store and execute both sync and async hooks
- Add async executor methods that run async hooks first, then sync hooks for backward compatibility
- Update basesdk.py do_request_async to call async hook methods
- Enables proper async/await patterns in hooks without blocking the event loop

This is the first step toward fixing ENG-792 where blocking calls in the split PDF hook prevent concurrent async requests.

Co-Authored-By: Claude Sonnet 4.5 &lt;noreply@anthropic.com&gt;
diff --git a/src/unstructured_client/_hooks/sdkhooks.py b/src/unstructured_client/_hooks/sdkhooks.py
@@ -1,6 +1,7 @@
 """Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT."""
 
 import httpx
+import inspect
 from .types import (
     SDKInitHook,
     BeforeRequestContext,
@@ -9,11 +10,15 @@
     AfterSuccessHook,
     AfterErrorContext,
     AfterErrorHook,
+    AsyncSDKInitHook,
+    AsyncBeforeRequestHook,
+    AsyncAfterSuccessHook,
+    AsyncAfterErrorHook,
     Hooks,
 )
 from .registration import init_hooks
-from typing import List, Optional, Tuple
-from unstructured_client.httpclient import HttpClient
+from typing import List, Optional, Tuple, Union
+from unstructured_client.httpclient import AsyncHttpClient, HttpClient
 
 
 class SDKHooks(Hooks):
@@ -22,6 +27,10 @@ def __init__(self) -> None:
         self.before_request_hooks: List[BeforeRequestHook] = []
         self.after_success_hooks: List[AfterSuccessHook] = []
         self.after_error_hooks: List[AfterErrorHook] = []
+        self.async_sdk_init_hooks: List[AsyncSDKInitHook] = []
+        self.async_before_request_hooks: List[AsyncBeforeRequestHook] = []
+        self.async_after_success_hooks: List[AsyncAfterSuccessHook] = []
+        self.async_after_error_hooks: List[AsyncAfterErrorHook] = []
         init_hooks(self)
 
     def register_sdk_init_hook(self, hook: SDKInitHook) -> None:
@@ -36,6 +45,18 @@ def register_after_success_hook(self, hook: AfterSuccessHook) -> None:
     def register_after_error_hook(self, hook: AfterErrorHook) -> None:
         self.after_error_hooks.append(hook)
 
+    def register_async_sdk_init_hook(self, hook: AsyncSDKInitHook) -> None:
+        self.async_sdk_init_hooks.append(hook)
+
+    def register_async_before_request_hook(self, hook: AsyncBeforeRequestHook) -> None:
+        self.async_before_request_hooks.append(hook)
+
+    def register_async_after_success_hook(self, hook: AsyncAfterSuccessHook) -> None:
+        self.async_after_success_hooks.append(hook)
+
+    def register_async_after_error_hook(self, hook: AsyncAfterErrorHook) -> None:
+        self.async_after_error_hooks.append(hook)
+
     def sdk_init(self, base_url: str, client: HttpClient) -> Tuple[str, HttpClient]:
         for hook in self.sdk_init_hooks:
             base_url, client = hook.sdk_init(base_url, client)
@@ -74,3 +95,75 @@ def after_error(
                 raise result
             response, error = result
         return response, error
+
+    async def sdk_init_async(
+        self, base_url: str, client: AsyncHttpClient
+    ) -> Tuple[str, AsyncHttpClient]:
+        # Run async hooks first
+        for hook in self.async_sdk_init_hooks:
+            base_url, client = await hook.sdk_init_async(base_url, client)
+        # Then run sync hooks for backward compatibility (if hook has sync version)
+        for hook in self.sdk_init_hooks:
+            # Cast the client temporarily - sync hooks may not touch async client
+            base_url, _ = hook.sdk_init(base_url, client)  # type: ignore
+        return base_url, client
+
+    async def before_request_async(
+        self, hook_ctx: BeforeRequestContext, request: httpx.Request
+    ) -> httpx.Request:
+        # Run async hooks first
+        for hook in self.async_before_request_hooks:
+            out = await hook.before_request_async(hook_ctx, request)
+            if isinstance(out, Exception):
+                raise out
+            request = out
+
+        # Then run sync hooks for backward compatibility
+        for hook in self.before_request_hooks:
+            out = hook.before_request(hook_ctx, request)
+            if isinstance(out, Exception):
+                raise out
+            request = out
+
+        return request
+
+    async def after_success_async(
+        self, hook_ctx: AfterSuccessContext, response: httpx.Response
+    ) -> httpx.Response:
+        # Run async hooks first
+        for hook in self.async_after_success_hooks:
+            out = await hook.after_success_async(hook_ctx, response)
+            if isinstance(out, Exception):
+                raise out
+            response = out
+
+        # Then run sync hooks for backward compatibility
+        for hook in self.after_success_hooks:
+            out = hook.after_success(hook_ctx, response)
+            if isinstance(out, Exception):
+                raise out
+            response = out
+
+        return response
+
+    async def after_error_async(
+        self,
+        hook_ctx: AfterErrorContext,
+        response: Optional[httpx.Response],
+        error: Optional[Exception],
+    ) -> Tuple[Optional[httpx.Response], Optional[Exception]]:
+        # Run async hooks first
+        for hook in self.async_after_error_hooks:
+            result = await hook.after_error_async(hook_ctx, response, error)
+            if isinstance(result, Exception):
+                raise result
+            response, error = result
+
+        # Then run sync hooks for backward compatibility
+        for hook in self.after_error_hooks:
+            result = hook.after_error(hook_ctx, response, error)
+            if isinstance(result, Exception):
+                raise result
+            response, error = result
+
+        return response, error
diff --git a/src/unstructured_client/_hooks/types.py b/src/unstructured_client/_hooks/types.py
@@ -3,7 +3,7 @@
 from abc import ABC, abstractmethod
 import httpx
 from typing import Any, Callable, List, Optional, Tuple, Union
-from unstructured_client.httpclient import HttpClient
+from unstructured_client.httpclient import AsyncHttpClient, HttpClient
 from unstructured_client.sdkconfiguration import SDKConfiguration
 
 
@@ -95,6 +95,41 @@ def after_error(
         pass
 
 
+class AsyncSDKInitHook(ABC):
+    @abstractmethod
+    async def sdk_init_async(
+        self, base_url: str, client: AsyncHttpClient
+    ) -> Tuple[str, AsyncHttpClient]:
+        pass
+
+
+class AsyncBeforeRequestHook(ABC):
+    @abstractmethod
+    async def before_request_async(
+        self, hook_ctx: BeforeRequestContext, request: httpx.Request
+    ) -> Union[httpx.Request, Exception]:
+        pass
+
+
+class AsyncAfterSuccessHook(ABC):
+    @abstractmethod
+    async def after_success_async(
+        self, hook_ctx: AfterSuccessContext, response: httpx.Response
+    ) -> Union[httpx.Response, Exception]:
+        pass
+
+
+class AsyncAfterErrorHook(ABC):
+    @abstractmethod
+    async def after_error_async(
+        self,
+        hook_ctx: AfterErrorContext,
+        response: Optional[httpx.Response],
+        error: Optional[Exception],
+    ) -> Union[Tuple[Optional[httpx.Response], Optional[Exception]], Exception]:
+        pass
+
+
 class Hooks(ABC):
     @abstractmethod
     def register_sdk_init_hook(self, hook: SDKInitHook):
@@ -111,3 +146,19 @@ def register_after_success_hook(self, hook: AfterSuccessHook):
     @abstractmethod
     def register_after_error_hook(self, hook: AfterErrorHook):
         pass
+
+    @abstractmethod
+    def register_async_sdk_init_hook(self, hook: AsyncSDKInitHook):
+        pass
+
+    @abstractmethod
+    def register_async_before_request_hook(self, hook: AsyncBeforeRequestHook):
+        pass
+
+    @abstractmethod
+    def register_async_after_success_hook(self, hook: AsyncAfterSuccessHook):
+        pass
+
+    @abstractmethod
+    def register_async_after_error_hook(self, hook: AsyncAfterErrorHook):
+        pass
diff --git a/src/unstructured_client/basesdk.py b/src/unstructured_client/basesdk.py
@@ -300,7 +300,7 @@ async def do_request_async(
         async def do():
             http_res = None
             try:
-                req = hooks.before_request(BeforeRequestContext(hook_ctx), request)
+                req = await hooks.before_request_async(BeforeRequestContext(hook_ctx), request)
                 logger.debug(
                     "Request:\nMethod: %s\nURL: %s\nHeaders: %s\nBody: %s",
                     req.method,
@@ -314,7 +314,7 @@ async def do():
 
                 http_res = await client.send(req, stream=stream)
             except Exception as e:
-                _, e = hooks.after_error(AfterErrorContext(hook_ctx), None, e)
+                _, e = await hooks.after_error_async(AfterErrorContext(hook_ctx), None, e)
                 if e is not None:
                     logger.debug("Request Exception", exc_info=True)
                     raise e
@@ -332,7 +332,7 @@ async def do():
             )
 
             if utils.match_status_codes(error_status_codes, http_res.status_code):
-                result, err = hooks.after_error(
+                result, err = await hooks.after_error_async(
                     AfterErrorContext(hook_ctx), http_res, None
                 )
                 if err is not None:
@@ -354,6 +354,6 @@ async def do():
             http_res = await do()
 
         if not utils.match_status_codes(error_status_codes, http_res.status_code):
-            http_res = hooks.after_success(AfterSuccessContext(hook_ctx), http_res)
+            http_res = await hooks.after_success_async(AfterSuccessContext(hook_ctx), http_res)
 
         return http_res
