dara-core 1.17.6__py3-none-any.whl → 1.18.0__py3-none-any.whl

dara/core/internal/registry_lookup.py

@@ -15,7 +15,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from typing import Callable, Coroutine, Dict, Literal
+from collections.abc import Coroutine
+from typing import Callable, Dict, Literal, Union
 
 from dara.core.internal.registry import Registry, RegistryType
 from dara.core.internal.utils import async_dedupe
@@ -28,6 +29,7 @@ RegistryLookupKey = Literal[
     RegistryType.STATIC_KWARGS,
     RegistryType.UPLOAD_RESOLVER,
     RegistryType.BACKEND_STORE,
+    RegistryType.DOWNLOAD_CODE,
 ]
 CustomRegistryLookup = Dict[RegistryLookupKey, Callable[[str], Coroutine]]
 
@@ -37,7 +39,9 @@ class RegistryLookup:
     Manages registry Lookup.
     """
 
-    def __init__(self, handlers: CustomRegistryLookup = {}):
+    def __init__(self, handlers: Union[CustomRegistryLookup, None] = None):
+        if handlers is None:
+            handlers = {}
         self.handlers = handlers
 
     @async_dedupe
@@ -62,4 +66,4 @@ class RegistryLookup:
                 return entry
             raise ValueError(
                 f'Could not find uid {uid} in {registry.name} registry, did you register it before the app was initialized?'
-            ).with_traceback(e.__traceback__)
+            ) from e

dara/core/internal/routing.py

@@ -18,9 +18,10 @@ limitations under the License.
 import inspect
 import json
 import os
+from collections.abc import Mapping
 from functools import wraps
 from importlib.metadata import version
-from typing import Any, Callable, Dict, List, Mapping, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 import anyio
 import pandas
@@ -57,6 +58,7 @@ from dara.core.internal.registries import (
     component_registry,
     data_variable_registry,
     derived_variable_registry,
+    download_code_registry,
     latest_value_registry,
     static_kwargs_registry,
     template_registry,
@@ -86,7 +88,7 @@ def error_decorator(handler: Callable[..., Any]):
                 if isinstance(err, HTTPException):
                     raise err
                 dev_logger.error('Unhandled error', error=err)
-                raise HTTPException(status_code=500, detail=str(err))
+                raise HTTPException(status_code=500, detail=str(err)) from err
 
         return _async_inner_func
 
@@ -99,7 +101,7 @@ def error_decorator(handler: Callable[..., Any]):
             if isinstance(err, HTTPException):
                 raise err
             dev_logger.error('Unhandled error', error=err)
-            raise HTTPException(status_code=500, detail=str(err))
+            raise HTTPException(status_code=500, detail=str(err)) from err
 
     return _inner_func
 
@@ -113,7 +115,7 @@ def create_router(config: Configuration):
     core_api_router = APIRouter()
 
     @core_api_router.get('/actions', dependencies=[Depends(verify_session)])
-    async def get_actions():   # pylint: disable=unused-variable
+    async def get_actions():
         return action_def_registry.get_all().items()
 
     class ActionRequestBody(BaseModel):
@@ -133,7 +135,7 @@ def create_router(config: Configuration):
         """Execution id, unique to this request"""
 
     @core_api_router.post('/action/{uid}', dependencies=[Depends(verify_session)])
-    async def get_action(uid: str, body: ActionRequestBody):  # pylint: disable=unused-variable
+    async def get_action(uid: str, body: ActionRequestBody):
         store: CacheStore = utils_registry.get('Store')
         task_mgr: TaskManager = utils_registry.get('TaskManager')
         registry_mgr: RegistryLookup = utils_registry.get('RegistryLookup')
@@ -150,7 +152,14 @@ def create_router(config: Configuration):
 
         # Execute the action - kick off a background task to run the handler
         response = await action_def.execute_action(
-            action_def, body.input, values, static_kwargs, body.execution_id, body.ws_channel, store, task_mgr
+            action_def,
+            body.input,
+            values,
+            static_kwargs,
+            body.execution_id,
+            body.ws_channel,
+            store,
+            task_mgr,
         )
 
         if isinstance(response, BaseTask):
@@ -159,15 +168,22 @@ def create_router(config: Configuration):
 
         return {'execution_id': response}
 
-    @core_api_router.get('/download')
-    async def get_download(code: str):   # pylint: disable=unused-variable
+    @core_api_router.get('/download')  # explicitly unauthenticated
+    async def get_download(code: str):
         store: CacheStore = utils_registry.get('Store')
 
         try:
             data_entry = await store.get(DownloadRegistryEntry, key=code)
 
+            # If not found directly in the store, use the override registry
+            # to check if we can get the download entry from there
             if data_entry is None:
-                raise ValueError('Invalid or expired download code')
+                registry_mgr: RegistryLookup = utils_registry.get('RegistryLookup')
+                # NOTE: This will throw a Value/KeyError if the code is not found so no need to rethrow
+                data_entry = await registry_mgr.get(download_code_registry, code)
+                # We managed to find one from the lookup,
+                # remove it from the registry immediately because it's one time use
+                download_code_registry.remove(code)
 
             async_file, cleanup = await data_entry.download(data_entry)
 
@@ -188,11 +204,11 @@ def create_router(config: Configuration):
                 background=BackgroundTask(cleanup),
             )
 
-        except KeyError:
-            raise ValueError('Invalid or expired download code')
+        except (KeyError, ValueError) as e:
+            raise ValueError('Invalid or expired download code') from e
 
     @core_api_router.get('/config', dependencies=[Depends(verify_session)])
-    async def get_config():  # pylint: disable=unused-variable
+    async def get_config():
         return {
             **config.model_dump(
                 include={
@@ -209,13 +225,13 @@ def create_router(config: Configuration):
         }
 
     @core_api_router.get('/auth-config')
-    async def get_auth_config():  # pylint: disable=unused-variable
+    async def get_auth_config():
         return {
             'auth_components': config.auth_config.component_config.model_dump(),
         }
 
     @core_api_router.get('/components', dependencies=[Depends(verify_session)])
-    async def get_components(name: Optional[str] = None):  # pylint: disable=unused-variable
+    async def get_components(name: Optional[str] = None):
         """
         If name is passed, will try to register the component
 
@@ -236,7 +252,7 @@ def create_router(config: Configuration):
         ws_channel: str
 
     @core_api_router.post('/components/{component}', dependencies=[Depends(verify_session)])
-    async def get_component(component: str, body: ComponentRequestBody):  # pylint: disable=unused-variable
+    async def get_component(component: str, body: ComponentRequestBody):
         CURRENT_COMPONENT_ID.set(body.uid)
         WS_CHANNEL.set(body.ws_channel)
         store: CacheStore = utils_registry.get('Store')
@@ -265,7 +281,7 @@ def create_router(config: Configuration):
         raise HTTPException(status_code=400, detail='Requesting this type of component is not supported')
 
     @core_api_router.get('/derived-variable/{uid}/latest', dependencies=[Depends(verify_session)])
-    async def get_latest_derived_variable(uid: str):  # pylint: disable=unused-variable
+    async def get_latest_derived_variable(uid: str):
         try:
             store: CacheStore = utils_registry.get('Store')
             latest_value_entry = latest_value_registry.get(uid)
@@ -289,9 +305,7 @@ def create_router(config: Configuration):
             return latest_value
 
         except KeyError as err:
-            raise ValueError(f'Could not find latest value for derived variable with uid: {uid}').with_traceback(
-                err.__traceback__
-            )
+            raise ValueError(f'Could not find latest value for derived variable with uid: {uid}') from err
 
     class DataVariableRequestBody(BaseModel):
         filters: Optional[FilterQuery] = None
@@ -306,7 +320,7 @@ def create_router(config: Configuration):
         limit: Optional[int] = None,
         order_by: Optional[str] = None,
         index: Optional[str] = None,
-    ):   # pylint: disable=unused-variable
+    ):
         try:
             store: CacheStore = utils_registry.get('Store')
             task_mgr: TaskManager = utils_registry.get('TaskManager')
@@ -318,11 +332,15 @@ def create_router(config: Configuration):
 
             if data_variable_entry.type == 'derived':
                 if body.cache_key is None:
-                    raise HTTPException(status_code=400, detail='Cache key is required for derived data variables')
+                    raise HTTPException(
+                        status_code=400,
+                        detail='Cache key is required for derived data variables',
+                    )
 
                 if body.ws_channel is None:
                     raise HTTPException(
-                        status_code=400, detail='Websocket channel is required for derived data variables'
+                        status_code=400,
+                        detail='Websocket channel is required for derived data variables',
                     )
 
                 derived_variable_entry = await registry_mgr.get(derived_variable_registry, uid)
@@ -351,7 +369,10 @@ def create_router(config: Configuration):
             dev_logger.debug(
                 f'DataVariable {data_variable_entry.uid[:3]}..{data_variable_entry.uid[-3:]}',
                 'return value',
-                {'value': data.describe() if isinstance(data, pandas.DataFrame) else None, 'uid': uid},  # type: ignore
+                {
+                    'value': data.describe() if isinstance(data, pandas.DataFrame) else None,
+                    'uid': uid,
+                },  # type: ignore
             )
 
             if data is None:
@@ -360,10 +381,11 @@ def create_router(config: Configuration):
             # Explicitly convert to JSON to avoid implicit serialization;
             # return as records as that makes more sense in a JSON structure
             return Response(
-                content=df_to_json(data) if isinstance(data, pandas.DataFrame) else data, media_type='application/json'
-            )   # type: ignore
+                content=df_to_json(data) if isinstance(data, pandas.DataFrame) else data,
+                media_type='application/json',
+            )  # type: ignore
         except ValueError as e:
-            raise HTTPException(status_code=400, detail=str(e))
+            raise HTTPException(status_code=400, detail=str(e)) from e
 
     class DataVariableCountRequestBody(BaseModel):
         cache_key: Optional[str] = None
@@ -383,12 +405,13 @@ def create_router(config: Configuration):
 
             if body is None or body.cache_key is None:
                 raise HTTPException(
-                    status_code=400, detail="Cache key is required when requesting DerivedDataVariable's count"
+                    status_code=400,
+                    detail="Cache key is required when requesting DerivedDataVariable's count",
                 )
 
             return await variable_def.get_total_count(variable_def, store, body.cache_key, body.filters)
         except ValueError as e:
-            raise HTTPException(status_code=400, detail=str(e))
+            raise HTTPException(status_code=400, detail=str(e)) from e
 
     @core_api_router.get('/data-variable/{uid}/schema', dependencies=[Depends(verify_session)])
     async def get_data_variable_schema(uid: str, cache_key: Optional[str] = None):
@@ -402,7 +425,8 @@ def create_router(config: Configuration):
 
             if cache_key is None:
                 raise HTTPException(
-                    status_code=400, detail='Cache key is required when requesting DerivedDataVariable schema'
+                    status_code=400,
+                    detail='Cache key is required when requesting DerivedDataVariable schema',
                 )
 
             # Use the other registry for derived variables
@@ -411,7 +435,7 @@ def create_router(config: Configuration):
             content = json.dumps(jsonable_encoder(data)) if isinstance(data, dict) else data
             return Response(content=content, media_type='application/json')
         except ValueError as e:
-            raise HTTPException(status_code=400, detail=str(e))
+            raise HTTPException(status_code=400, detail=str(e)) from e
 
     @core_api_router.post('/data/upload', dependencies=[Depends(verify_session)])
     async def upload_data(
@@ -427,7 +451,10 @@ def create_router(config: Configuration):
         registry_mgr: RegistryLookup = utils_registry.get('RegistryLookup')
 
         if data_uid is None and resolver_id is None:
-            raise HTTPException(400, 'Neither resolver_id or data_uid specified, at least one of them is required')
+            raise HTTPException(
+                400,
+                'Neither resolver_id or data_uid specified, at least one of them is required',
+            )
 
         try:
             # If resolver id is provided, run the custom
@@ -440,16 +467,16 @@ def create_router(config: Configuration):
 
             return {'status': 'SUCCESS'}
         except Exception as e:
-            raise HTTPException(status_code=400, detail=str(e))
+            raise HTTPException(status_code=400, detail=str(e)) from e
 
     class DerivedStateRequestBody(BaseModel):
         values: NormalizedPayload[List[Any]]
-        force: bool
+        force_key: Optional[str] = None
         ws_channel: str
         is_data_variable: Optional[bool] = False
 
     @core_api_router.post('/derived-variable/{uid}', dependencies=[Depends(verify_session)])
-    async def get_derived_variable(uid: str, body: DerivedStateRequestBody):  # pylint: disable=unused-variable
+    async def get_derived_variable(uid: str, body: DerivedStateRequestBody):
         task_mgr: TaskManager = utils_registry.get('TaskManager')
         store: CacheStore = utils_registry.get('Store')
         registry_mgr: RegistryLookup = utils_registry.get('RegistryLookup')
@@ -457,7 +484,7 @@ def create_router(config: Configuration):
 
         values = denormalize(body.values.data, body.values.lookup)
 
-        result = await variable_def.get_value(variable_def, store, task_mgr, values, body.force)
+        result = await variable_def.get_value(variable_def, store, task_mgr, values, body.force_key)
 
         response: Any = result
 
@@ -466,7 +493,10 @@ def create_router(config: Configuration):
         if isinstance(result['value'], BaseTask):
             # Kick off the task
             await task_mgr.run_task(result['value'], body.ws_channel)
-            response = {'task_id': result['value'].task_id, 'cache_key': result['cache_key']}
+            response = {
+                'task_id': result['value'].task_id,
+                'cache_key': result['cache_key'],
+            }
 
         dev_logger.debug(
             f'DerivedVariable {variable_def.uid[:3]}..{variable_def.uid[-3:]}',
@@ -512,7 +542,7 @@ def create_router(config: Configuration):
                 tg.start_soon(_write, store_uid, value)
 
     @core_api_router.get('/tasks/{task_id}', dependencies=[Depends(verify_session)])
-    async def get_task_result(task_id: str):   # pylint: disable=unused-variable
+    async def get_task_result(task_id: str):
         try:
             task_mgr: TaskManager = utils_registry.get('TaskManager')
             res = await task_mgr.get_result(task_id)
@@ -528,25 +558,28 @@ def create_router(config: Configuration):
                 return Response(df_to_json(res))
 
             return res
-        except Exception as e:
-            raise ValueError(f'The result for task id {task_id} could not be found').with_traceback(e.__traceback__)
+        except Exception as err:
+            raise ValueError(f'The result for task id {task_id} could not be found') from err
 
     @core_api_router.delete('/tasks/{task_id}', dependencies=[Depends(verify_session)])
-    async def cancel_task(task_id: str):   # pylint: disable=unused-variable
+    async def cancel_task(task_id: str):
         try:
             task_mgr: TaskManager = utils_registry.get('TaskManager')
             return await task_mgr.cancel_task(task_id)
         except TaskManagerError as e:
-            dev_logger.error(f'The task id {task_id} could not be found, it may have already been cancelled', e)
+            dev_logger.error(
+                f'The task id {task_id} could not be found, it may have already been cancelled',
+                e,
+            )
 
     @core_api_router.get('/template/{template}', dependencies=[Depends(verify_session)])
-    async def get_template(template: str):  # pylint: disable=unused-variable
+    async def get_template(template: str):
         try:
             selected_template = template_registry.get(template)
             normalized_template, lookup = normalize(jsonable_encoder(selected_template))
             return {'data': normalized_template, 'lookup': lookup}
-        except KeyError:
-            raise HTTPException(status_code=404, detail=f'Template: {template}, not found in registry')
+        except KeyError as err:
+            raise HTTPException(status_code=404, detail=f'Template: {template}, not found in registry') from err
         except Exception as e:
             dev_logger.error('Something went wrong while trying to get the template', e)
 

dara/core/internal/scheduler.py

@@ -20,7 +20,7 @@ from datetime import datetime
 from multiprocessing import get_context
 from multiprocessing.process import BaseProcess
 from pickle import PicklingError
-from typing import Any, List, Optional, Union
+from typing import Any, List, Optional, Union, cast
 
 from croniter import croniter
 from pydantic import BaseModel, field_validator
@@ -56,25 +56,28 @@ class ScheduledJob(BaseModel):
             job_process = ctx.Process(target=self._refresh_timer, args=(func, args), daemon=True)
             job_process.start()
             return job_process
-        except PicklingError:
+        except PicklingError as err:
             raise PicklingError(
                 """
             Unable to pickle scheduled function. Please ensure that the function you are trying
             to schedule is not in the same file as the ConfigurationBuilder is defined and that
             the function is not a lambda.
             """
-            )
+            ) from err
 
     def _refresh_timer(self, func, args):
         while self.continue_running and not (self.run_once and not self.first_execution):
-            interval = self.interval
+            interval: int
             # If there's more than one interval to wait, i.e. this is a weekday process
-            if type(interval) == list:
+            if isinstance(self.interval, list):
                 # Wait the first interval if this is the first execution of the job
                 interval = self.interval[0] if self.first_execution else self.interval[1]
+            else:
+                interval = self.interval
+
             self.first_execution = False
             # Wait the interval and then run the job
-            time.sleep(interval)
+            time.sleep(cast(int, interval))
             func(*args)
 
 
@@ -176,7 +179,7 @@ class ScheduledJobFactory(BaseModel):
 
     @field_validator('weekday', mode='before')
     @classmethod
-    def validate_weekday(cls, weekday: Any) -> datetime:  # pylint: disable=E0213
+    def validate_weekday(cls, weekday: Any) -> datetime:
         if isinstance(weekday, datetime):
             return weekday
         if isinstance(weekday, str):
@@ -283,7 +286,9 @@ class Scheduler:
     def _weekday(self, weekday: int):
         # The job must execute on a weekly interval
         return ScheduledJobFactory(
-            interval=self.interval * 604800, run_once=self._run_once, weekday=str(weekday)  # type: ignore
+            interval=self.interval * 604800,
+            run_once=self._run_once,
+            weekday=str(weekday),  # type: ignore
         )
 
     def monday(self) -> ScheduledJobFactory:

dara/core/internal/settings.py

@@ -74,7 +74,7 @@ def generate_env_file(filename='.env'):
         f.write(env_content)
 
 
-@lru_cache()
+@lru_cache
 def get_settings():
     """
     Get a cached instance of the settings, loading values from the .env if present.
@@ -83,7 +83,7 @@ def get_settings():
 
     # Test purposes - if DARA_TEST_FLAG is set then override env with .env.test
     if os.environ.get('DARA_TEST_FLAG', None) is not None:
-        return Settings(**dotenv_values('.env.test'))
+        return Settings(**dotenv_values('.env.test'))  # type: ignore
 
     env_error = False
 

dara/core/internal/tasks.py

@@ -15,9 +15,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import contextlib
 import inspect
 import math
-from typing import Any, Awaitable, Callable, Dict, List, Optional, Union, overload
+from collections.abc import Awaitable
+from typing import Any, Callable, Dict, List, Optional, Union, overload
 
 from anyio import (
     CancelScope,
@@ -137,14 +139,12 @@ class Task(BaseTask):
 
             async def on_progress(progress: float, msg: str):
                 if send_stream is not None:
-                    try:
+                    with contextlib.suppress(ClosedResourceError):
                         await send_stream.send(TaskProgressUpdate(task_id=self.task_id, progress=progress, message=msg))
-                    except ClosedResourceError:
-                        pass
 
             async def on_result(result: Any):
                 if send_stream is not None:
-                    try:
+                    with contextlib.suppress(ClosedResourceError):
                         await send_stream.send(
                             TaskResult(
                                 task_id=self.task_id,
@@ -153,19 +153,15 @@ class Task(BaseTask):
                                 reg_entry=self.reg_entry,
                             )
                         )
-                    except ClosedResourceError:
-                        pass
 
             async def on_error(exc: BaseException):
                 if send_stream is not None:
-                    try:
+                    with contextlib.suppress(ClosedResourceError):
                         await send_stream.send(
                             TaskError(
                                 task_id=self.task_id, error=exc, cache_key=self.cache_key, reg_entry=self.reg_entry
                             )
                         )
-                    except ClosedResourceError:
-                        pass
 
             with pool.on_progress(self.task_id, on_progress):
                 pool_task_def = pool.submit(self.task_id, self._func_name, args=tuple(self._args), kwargs=self._kwargs)
@@ -362,12 +358,10 @@ class TaskManager:
         self.store = store
 
     @overload
-    async def run_task(self, task: PendingTask, ws_channel: Optional[str] = None) -> Any:
-        ...
+    async def run_task(self, task: PendingTask, ws_channel: Optional[str] = None) -> Any: ...
 
     @overload
-    async def run_task(self, task: BaseTask, ws_channel: Optional[str] = None) -> PendingTask:
-        ...
+    async def run_task(self, task: BaseTask, ws_channel: Optional[str] = None) -> PendingTask: ...
 
     async def run_task(self, task: BaseTask, ws_channel: Optional[str] = None):
         """

dara/core/internal/utils.py

@@ -20,6 +20,7 @@ from __future__ import annotations
 import asyncio
 import inspect
 import os
+from collections.abc import Awaitable, Coroutine, Sequence
 from functools import wraps
 from importlib import import_module
 from importlib.util import find_spec
@@ -27,13 +28,10 @@ from types import ModuleType
 from typing import (
     TYPE_CHECKING,
     Any,
-    Awaitable,
     Callable,
-    Coroutine,
     Dict,
     Literal,
     Optional,
-    Sequence,
     Tuple,
     Union,
 )
@@ -79,7 +77,7 @@ def get_cache_scope(cache_type: Optional[CacheType]) -> CacheScope:
     return 'global'
 
 
-async def run_user_handler(handler: Callable, args: Sequence = [], kwargs: dict = {}):
+async def run_user_handler(handler: Callable, args: Union[Sequence, None] = None, kwargs: Union[dict, None] = None):
     """
     Run a user-defined handler function. Runs sync functions in a threadpool.
     Handles SystemExits cleanly.
@@ -88,6 +86,10 @@ async def run_user_handler(handler: Callable, args: Sequence = [], kwargs: dict
     :param args: list of arguments to pass to the function
     :param kwargs: dict of kwargs to past to the function
     """
+    if args is None:
+        args = []
+    if kwargs is None:
+        kwargs = {}
     with handle_system_exit('User defined function quit unexpectedly'):
         if inspect.iscoroutinefunction(handler):
             return await handler(*args, **kwargs)
@@ -164,14 +166,14 @@ def enforce_sso(conf: ConfigurationBuilder):
     Raises if SSO is not used
     """
     try:
-        from dara.enterprise import SSOAuthConfig
+        from dara.enterprise import SSOAuthConfig  # pyright: ignore[reportMissingImports]
 
         if conf.auth_config is None or not isinstance(conf.auth_config, SSOAuthConfig):
             raise ValueError('Config does not have SSO auth enabled. Please update your application to configure SSO.')
-    except ImportError:
+    except ImportError as err:
         raise ValueError(
             'SSO is not enabled. Please install the dara_enterprise package and configure SSO to use this feature.'
-        )
+        ) from err
 
 
 def async_dedupe(fn: Callable[..., Awaitable]):
@@ -228,8 +230,7 @@ def resolve_exception_group(error: Any):
 
     :param error: The error to resolve
     """
-    if isinstance(error, ExceptionGroup):
-        if len(error.exceptions) == 1:
-            return resolve_exception_group(error.exceptions[0])
+    if isinstance(error, ExceptionGroup) and len(error.exceptions) == 1:
+        return resolve_exception_group(error.exceptions[0])
 
     return error

dara/core/internal/websocket.py

@@ -204,28 +204,27 @@ class WebSocketHandler:
             message_id = message.channel
 
             # If the message has a channel ID, it's a response to a previous message
-            if message_id:
-                if message_id in self.pending_responses:
-                    event, existing_messages = self.pending_responses[message_id]
-
-                    # If the response is chunked then collect the messages in pending responses
-                    if message.chunk_count is not None:
-                        if existing_messages is not None and isinstance(existing_messages, list):
-                            existing_messages.append(message.message)
-                        else:
-                            existing_messages = [message.message]
-                            self.pending_responses[message_id] = (
-                                event,
-                                existing_messages,
-                            )
+            if message_id and message_id in self.pending_responses:
+                event, existing_messages = self.pending_responses[message_id]
 
-                        # If all chunks have been received, set the event to notify the waiting coroutine
-                        if len(existing_messages) == message.chunk_count:
-                            event.set()
+                # If the response is chunked then collect the messages in pending responses
+                if message.chunk_count is not None:
+                    if existing_messages is not None and isinstance(existing_messages, list):
+                        existing_messages.append(message.message)
                     else:
-                        # Store the response and set the event to notify the waiting coroutine
-                        self.pending_responses[message_id] = (event, message.message)
+                        existing_messages = [message.message]
+                        self.pending_responses[message_id] = (
+                            event,
+                            existing_messages,
+                        )
+
+                    # If all chunks have been received, set the event to notify the waiting coroutine
+                    if len(existing_messages) == message.chunk_count:
                         event.set()
+                else:
+                    # Store the response and set the event to notify the waiting coroutine
+                    self.pending_responses[message_id] = (event, message.message)
+                    event.set()
 
             return None