fix(assistant): improve middleware dispatch and inject kwargs in middleware (#1456)

Author: WilliamBergamin

slack_bolt/app/app.py

@@ -22,7 +22,6 @@
 
 from slack_bolt.context.assistant.thread_context_store.store import AssistantThreadContextStore
 
-from slack_bolt.context.assistant.assistant_utilities import AssistantUtilities
 from slack_bolt.error import BoltError, BoltUnhandledRequestError
 from slack_bolt.lazy_listener.thread_runner import ThreadLazyListenerRunner
 from slack_bolt.listener.builtins import TokenRevocationListeners
@@ -70,6 +69,7 @@
     IgnoringSelfEvents,
     CustomMiddleware,
     AttachingFunctionToken,
+    AttachingAgentKwargs,
 )
 from slack_bolt.middleware.assistant import Assistant
 from slack_bolt.middleware.message_listener_matches import MessageListenerMatches
@@ -83,10 +83,6 @@
 from slack_bolt.oauth.internals import select_consistent_installation_store
 from slack_bolt.oauth.oauth_settings import OAuthSettings
 from slack_bolt.request import BoltRequest
-from slack_bolt.request.payload_utils import (
-    is_assistant_event,
-    to_event,
-)
 from slack_bolt.response import BoltResponse
 from slack_bolt.util.utils import (
     create_web_client,
@@ -137,6 +133,7 @@ def __init__(
         listener_executor: Optional[Executor] = None,
         # for AI Agents & Assistants
         assistant_thread_context_store: Optional[AssistantThreadContextStore] = None,
+        attaching_agent_kwargs_enabled: bool = True,
     ):
         """Bolt App that provides functionalities to register middleware/listeners.
 
@@ -357,6 +354,7 @@ def message_hello(message, say):
             listener_executor = ThreadPoolExecutor(max_workers=5)
 
         self._assistant_thread_context_store = assistant_thread_context_store
+        self._attaching_agent_kwargs_enabled = attaching_agent_kwargs_enabled
 
         self._process_before_response = process_before_response
         self._listener_runner = ThreadListenerRunner(
@@ -841,10 +839,13 @@ def ask_for_introduction(event, say):
             middleware: A list of lister middleware functions.
                 Only when all the middleware call `next()` method, the listener function can be invoked.
         """
+        middleware = list(middleware) if middleware else []
 
         def __call__(*args, **kwargs):
             functions = self._to_listener_functions(kwargs) if kwargs else list(args)
             primary_matcher = builtin_matchers.event(event, base_logger=self._base_logger)
+            if self._attaching_agent_kwargs_enabled:
+                middleware.insert(0, AttachingAgentKwargs(self._assistant_thread_context_store))
             return self._register_listener(list(functions), primary_matcher, matchers, middleware, True)
 
         return __call__
@@ -902,6 +903,8 @@ def __call__(*args, **kwargs):
             primary_matcher = builtin_matchers.message_event(
                 keyword=keyword, constraints=constraints, base_logger=self._base_logger
             )
+            if self._attaching_agent_kwargs_enabled:
+                middleware.insert(0, AttachingAgentKwargs(self._assistant_thread_context_store))
             middleware.insert(0, MessageListenerMatches(keyword))
             return self._register_listener(list(functions), primary_matcher, matchers, middleware, True)
 
@@ -1398,20 +1401,6 @@ def _init_context(self, req: BoltRequest):
         # It is intended for apps that start lazy listeners from their custom global middleware.
         req.context["listener_runner"] = self.listener_runner
 
-        # For AI Agents & Assistants
-        if is_assistant_event(req.body):
-            assistant = AssistantUtilities(
-                payload=to_event(req.body),  # type: ignore[arg-type]
-                context=req.context,
-                thread_context_store=self._assistant_thread_context_store,
-            )
-            req.context["say"] = assistant.say
-            req.context["set_status"] = assistant.set_status
-            req.context["set_title"] = assistant.set_title
-            req.context["set_suggested_prompts"] = assistant.set_suggested_prompts
-            req.context["get_thread_context"] = assistant.get_thread_context
-            req.context["save_thread_context"] = assistant.save_thread_context
-
     @staticmethod
     def _to_listener_functions(
         kwargs: dict,

slack_bolt/app/async_app.py

@@ -8,7 +8,6 @@
 from aiohttp import web
 
 from slack_bolt.app.async_server import AsyncSlackAppServer
-from slack_bolt.context.assistant.async_assistant_utilities import AsyncAssistantUtilities
 from slack_bolt.context.assistant.thread_context_store.async_store import (
     AsyncAssistantThreadContextStore,
 )
@@ -30,7 +29,6 @@
     AsyncMessageListenerMatches,
 )
 from slack_bolt.oauth.async_internals import select_consistent_installation_store
-from slack_bolt.request.payload_utils import is_assistant_event, to_event
 from slack_bolt.util.utils import get_name_for_callable, is_callable_coroutine
 from slack_bolt.workflows.step.async_step import (
     AsyncWorkflowStep,
@@ -88,6 +86,7 @@
     AsyncIgnoringSelfEvents,
     AsyncUrlVerification,
     AsyncAttachingFunctionToken,
+    AsyncAttachingAgentKwargs,
 )
 from slack_bolt.middleware.async_custom_middleware import (
     AsyncMiddleware,
@@ -143,6 +142,7 @@ def __init__(
         verification_token: Optional[str] = None,
         # for AI Agents & Assistants
         assistant_thread_context_store: Optional[AsyncAssistantThreadContextStore] = None,
+        attaching_agent_kwargs_enabled: bool = True,
     ):
         """Bolt App that provides functionalities to register middleware/listeners.
 
@@ -363,6 +363,7 @@ async def message_hello(message, say):  # async function
         self._async_listeners: List[AsyncListener] = []
 
         self._assistant_thread_context_store = assistant_thread_context_store
+        self._attaching_agent_kwargs_enabled = attaching_agent_kwargs_enabled
 
         self._process_before_response = process_before_response
         self._async_listener_runner = AsyncioListenerRunner(
@@ -866,10 +867,13 @@ async def ask_for_introduction(event, say):
             middleware: A list of lister middleware functions.
                 Only when all the middleware call `next()` method, the listener function can be invoked.
         """
+        middleware = list(middleware) if middleware else []
 
         def __call__(*args, **kwargs):
             functions = self._to_listener_functions(kwargs) if kwargs else list(args)
             primary_matcher = builtin_matchers.event(event, True, base_logger=self._base_logger)
+            if self._attaching_agent_kwargs_enabled:
+                middleware.insert(0, AsyncAttachingAgentKwargs(self._assistant_thread_context_store))
             return self._register_listener(list(functions), primary_matcher, matchers, middleware, True)
 
         return __call__
@@ -930,6 +934,8 @@ def __call__(*args, **kwargs):
                 asyncio=True,
                 base_logger=self._base_logger,
             )
+            if self._attaching_agent_kwargs_enabled:
+                middleware.insert(0, AsyncAttachingAgentKwargs(self._assistant_thread_context_store))
             middleware.insert(0, AsyncMessageListenerMatches(keyword))
             return self._register_listener(list(functions), primary_matcher, matchers, middleware, True)
 
@@ -1431,20 +1437,6 @@ def _init_context(self, req: AsyncBoltRequest):
         # It is intended for apps that start lazy listeners from their custom global middleware.
         req.context["listener_runner"] = self.listener_runner
 
-        # For AI Agents & Assistants
-        if is_assistant_event(req.body):
-            assistant = AsyncAssistantUtilities(
-                payload=to_event(req.body),  # type: ignore[arg-type]
-                context=req.context,
-                thread_context_store=self._assistant_thread_context_store,
-            )
-            req.context["say"] = assistant.say
-            req.context["set_status"] = assistant.set_status
-            req.context["set_title"] = assistant.set_title
-            req.context["set_suggested_prompts"] = assistant.set_suggested_prompts
-            req.context["get_thread_context"] = assistant.get_thread_context
-            req.context["save_thread_context"] = assistant.save_thread_context
-
     @staticmethod
     def _to_listener_functions(
         kwargs: dict,

slack_bolt/context/async_context.py

@@ -110,7 +110,7 @@ async def handle_button_clicks(ack, say):
             Callable `say()` function
         """
         if "say" not in self:
-            self["say"] = AsyncSay(client=self.client, channel=self.channel_id, thread_ts=self.thread_ts)
+            self["say"] = AsyncSay(client=self.client, channel=self.channel_id)
         return self["say"]
 
     @property

slack_bolt/context/context.py

@@ -111,7 +111,7 @@ def handle_button_clicks(ack, say):
             Callable `say()` function
         """
         if "say" not in self:
-            self["say"] = Say(client=self.client, channel=self.channel_id, thread_ts=self.thread_ts)
+            self["say"] = Say(client=self.client, channel=self.channel_id)
         return self["say"]
 
     @property

slack_bolt/middleware/__init__.py

@@ -17,6 +17,7 @@
 from .ssl_check import SslCheck
 from .url_verification import UrlVerification
 from .attaching_function_token import AttachingFunctionToken
+from .attaching_agent_kwargs import AttachingAgentKwargs
 
 builtin_middleware_classes = [
     SslCheck,
@@ -41,5 +42,6 @@
     "SslCheck",
     "UrlVerification",
     "AttachingFunctionToken",
+    "AttachingAgentKwargs",
     "builtin_middleware_classes",
 ]

slack_bolt/middleware/assistant/assistant.py

@@ -7,6 +7,7 @@
 from slack_bolt.context.assistant.thread_context_store.store import AssistantThreadContextStore
 from slack_bolt.listener_matcher.builtins import build_listener_matcher
 
+from slack_bolt.middleware.attaching_agent_kwargs import AttachingAgentKwargs
 from slack_bolt.request.request import BoltRequest
 from slack_bolt.response.response import BoltResponse
 from slack_bolt.listener_matcher import CustomListenerMatcher
@@ -236,6 +237,15 @@ def process(  # type: ignore[return]
             if listeners is not None:
                 for listener in listeners:
                     if listener.matches(req=req, resp=resp):
+                        middleware_resp, next_was_not_called = listener.run_middleware(req=req, resp=resp)
+                        if next_was_not_called:
+                            if middleware_resp is not None:
+                                return middleware_resp
+                            # The listener middleware didn't call next().
+                            # Skip this listener and try the next one.
+                            continue
+                        if middleware_resp is not None:
+                            resp = middleware_resp
                         return listener_runner.run(
                             request=req,
                             response=resp,
@@ -262,6 +272,7 @@ def build_listener(
             return listener_or_functions
         elif isinstance(listener_or_functions, list):
             middleware = middleware if middleware else []
+            middleware.insert(0, AttachingAgentKwargs(self.thread_context_store))
             functions = listener_or_functions
             ack_function = functions.pop(0)
 

slack_bolt/middleware/assistant/async_assistant.py

@@ -8,6 +8,7 @@
 
 from slack_bolt.listener.asyncio_runner import AsyncioListenerRunner
 from slack_bolt.listener_matcher.builtins import build_listener_matcher
+from slack_bolt.middleware.attaching_agent_kwargs.async_attaching_agent_kwargs import AsyncAttachingAgentKwargs
 from slack_bolt.request.async_request import AsyncBoltRequest
 from slack_bolt.response import BoltResponse
 from slack_bolt.error import BoltError
@@ -265,6 +266,15 @@ async def async_process(  # type: ignore[return]
             if listeners is not None:
                 for listener in listeners:
                     if listener is not None and await listener.async_matches(req=req, resp=resp):
+                        middleware_resp, next_was_not_called = await listener.run_async_middleware(req=req, resp=resp)
+                        if next_was_not_called:
+                            if middleware_resp is not None:
+                                return middleware_resp
+                            # The listener middleware didn't call next().
+                            # Skip this listener and try the next one.
+                            continue
+                        if middleware_resp is not None:
+                            resp = middleware_resp
                         return await listener_runner.run(
                             request=req,
                             response=resp,
@@ -291,6 +301,7 @@ def build_listener(
             return listener_or_functions
         elif isinstance(listener_or_functions, list):
             middleware = middleware if middleware else []
+            middleware.insert(0, AsyncAttachingAgentKwargs(self.thread_context_store))
             functions = listener_or_functions
             ack_function = functions.pop(0)
 

slack_bolt/middleware/async_builtins.py

@@ -10,6 +10,7 @@
     AsyncMessageListenerMatches,
 )
 from .attaching_function_token.async_attaching_function_token import AsyncAttachingFunctionToken
+from .attaching_agent_kwargs.async_attaching_agent_kwargs import AsyncAttachingAgentKwargs
 
 __all__ = [
     "AsyncIgnoringSelfEvents",
@@ -18,4 +19,5 @@
     "AsyncUrlVerification",
     "AsyncMessageListenerMatches",
     "AsyncAttachingFunctionToken",
+    "AsyncAttachingAgentKwargs",
 ]

slack_bolt/middleware/attaching_agent_kwargs/__init__.py

@@ -0,0 +1,5 @@
+from .attaching_agent_kwargs import AttachingAgentKwargs
+
+__all__ = [
+    "AttachingAgentKwargs",
+]

slack_bolt/middleware/attaching_agent_kwargs/async_attaching_agent_kwargs.py

@@ -0,0 +1,39 @@
+from typing import Optional, Callable, Awaitable
+
+from slack_bolt.context.assistant.async_assistant_utilities import AsyncAssistantUtilities
+from slack_bolt.context.assistant.thread_context_store.async_store import AsyncAssistantThreadContextStore
+from slack_bolt.middleware.async_middleware import AsyncMiddleware
+from slack_bolt.request.async_request import AsyncBoltRequest
+from slack_bolt.request.payload_utils import is_assistant_event, to_event
+from slack_bolt.response import BoltResponse
+
+
+class AsyncAttachingAgentKwargs(AsyncMiddleware):
+
+    thread_context_store: Optional[AsyncAssistantThreadContextStore]
+
+    def __init__(self, thread_context_store: Optional[AsyncAssistantThreadContextStore] = None):
+        self.thread_context_store = thread_context_store
+
+    async def async_process(
+        self,
+        *,
+        req: AsyncBoltRequest,
+        resp: BoltResponse,
+        next: Callable[[], Awaitable[BoltResponse]],
+    ) -> Optional[BoltResponse]:
+        event = to_event(req.body)
+        if event is not None:
+            if is_assistant_event(req.body):
+                assistant = AsyncAssistantUtilities(
+                    payload=event,
+                    context=req.context,
+                    thread_context_store=self.thread_context_store,
+                )
+                req.context["say"] = assistant.say
+                req.context["set_status"] = assistant.set_status
+                req.context["set_title"] = assistant.set_title
+                req.context["set_suggested_prompts"] = assistant.set_suggested_prompts
+                req.context["get_thread_context"] = assistant.get_thread_context
+                req.context["save_thread_context"] = assistant.save_thread_context
+        return await next()