Add mypy typechecking and improve error handling

.github/workflows/ci.yaml

@@ -52,6 +52,10 @@ jobs:
         working-directory: sdk
         run: poetry run ruff format --check .
 
+      - name: Run mypy type checker
+        working-directory: sdk
+        run: poetry run mypy
+
   changelog:
     runs-on: ubuntu-latest
     name: Changelog check

sdk/CHANGELOG.md

@@ -8,13 +8,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
+- Static type checking with mypy: the `eggai` package is type-checked in CI
+  (`poetry run mypy`, config in `pyproject.toml`); `mypy` is a new dev dependency.
 - `RedisTransport`: new `group_start` subscribe option (`"$"` default, `"0"` or
   a stream id) chooses where a NEW consumer group starts reading, so a consumer
   added to a channel that already carries traffic can pick up the existing
   backlog. The transport now creates its consumer groups itself before the
   broker starts (#260).
 
 ### Fixed
+- A2A executor: the JSON-serialisation fallback caught the non-existent
+  `json.JSONEncodeError`, so a failing `json.dumps` raised `AttributeError` and
+  masked the real error. It now catches `(TypeError, ValueError)`.
+- A2A executor: the "unknown skill" error path enqueued a raw dict instead of a
+  proper A2A message; a skill registered with `data_type=None` raised
+  `TypeError` when building the message and now falls back to `BaseMessage`.
+- Typed subscriptions: a `data_type` whose `type` field has no default (e.g. raw
+  `BaseMessage`) silently dropped every message; `subscribe()` now raises
+  `ValueError` asking for a default discriminator.
+- `Agent.subscribe()`: plugin-prefixed kwargs (e.g. `a2a_*`) on an agent whose
+  plugin was never initialised raised an opaque `KeyError`; now `ValueError`.
 - `RedisTransport`: `agent.stop()` could hang forever on Python 3.10/3.11 with
   redis-py >= 8. redis-py 8 sends every command through `asyncio.wait_for`
   (`socket_timeout` now defaults to 5s), and on CPython < 3.12 `wait_for` can
@@ -38,6 +51,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   `Channel().publish()` never reached it (#264).
 
 ### Changed
+- Type-annotation cleanups across transports, channel, agent and hooks to
+  satisfy mypy (no behaviour change). `Transport.subscribe()`'s second
+  parameter is now named `handler` in the abstract base and all transports.
 - **BREAKING: `BaseMessage.data` is now required.** The generic base previously
   declared `data: TData = Field(default_factory=dict)`. Pydantic does not
   validate defaults, so an envelope missing `data` validated cleanly on *any*

sdk/eggai/adapters/a2a/executor.py

@@ -30,27 +30,28 @@
         try:
             # Extract skill name from request
             skill_name = self._extract_skill_name(request_context)
             logger.info(f"Executing A2A skill: {skill_name}")
 
             if skill_name not in self.a2a_plugin.handlers:
                 error_msg = f"Unknown skill: {skill_name}. Available skills: {list(self.a2a_plugin.handlers.keys())}"
                 logger.error(error_msg)
-                await event_queue.enqueue_event({"text": error_msg})
+                # Use the same A2A-message wrapper as every other response path,
+                # rather than enqueuing a raw dict the event queue can't model.
+                await self._send_agent_response(event_queue, {"error": error_msg})
                 return
 
             # Get handler and extract message data
             handler = self.a2a_plugin.handlers[skill_name]
             message_data = self._extract_message_data(request_context)
 
-            # Create proper BaseMessage with validated data
+            # Validate the raw data against the skill's declared type (if any) and
+            # wrap it in the BaseMessage subclass the handler expects. Looked up once;
+            # data_type may be None (a skill can register without one), which falls
+            # through to the generic BaseMessage rather than calling None(...).
             validated_data = message_data
-            if skill_name in self.a2a_plugin.data_types:
-                data_type = self.a2a_plugin.data_types[skill_name]
-                if (
-                    data_type
-                    and hasattr(data_type, "__args__")
-                    and len(data_type.__args__) > 0
-                ):
+            data_type = self.a2a_plugin.data_types.get(skill_name)
+            if data_type is not None:
+                if hasattr(data_type, "__args__") and len(data_type.__args__) > 0:
                     # Extract inner type from BaseMessage[T] and validate the raw data
                     inner_type = data_type.__args__[0]
                     if hasattr(inner_type, "model_validate"):
@@ -64,12 +65,8 @@
                                 f"Failed to validate data as {inner_type.__name__}: {e}"
                             )
                             # Keep as dict if validation fails
-
-            # Create the proper BaseMessage structure that the handler expects
-            if skill_name in self.a2a_plugin.data_types:
-                data_type = self.a2a_plugin.data_types[skill_name]
                 # Create instance of the specific BaseMessage subclass
                 eggai_message = data_type(
                     source="a2a-client",
                     type=f"{skill_name}.request",
                     data=validated_data,
@@ -107,7 +104,7 @@
             skill_name = (
                 locals()["skill_name"] if "skill_name" in locals() else "unknown"
             )
             error_msg = f"Error executing skill '{skill_name}': {str(e)}"
             logger.exception(error_msg)
             await self._send_agent_response(event_queue, {"error": error_msg})
 
@@ -216,6 +213,6 @@
             # Fallback: try to send as simple event
             try:
                 await event_queue.enqueue_event({"text": json.dumps(data)})
-            except (json.JSONEncodeError, TypeError) as json_err:
+            except (TypeError, ValueError) as json_err:
                 logger.error(f"Failed to serialize response data: {json_err}")
                 raise

sdk/eggai/adapters/a2a/plugin.py

@@ -22,7 +22,7 @@
         self.config = None
         self.skills: dict[str, Any] = {}  # AgentSkill objects
         self.handlers: dict[str, Callable] = {}  # Handler functions
-        self.data_types: dict[str, type] = {}  # Data types for conversion
+        self.data_types: dict[str, type | None] = {}  # Data types for conversion
 
     def init(
         self,
@@ -47,7 +47,7 @@
                     "a2a_capability must be a string representing the skill name"
                 )
             data_type = kwargs.get("data_type")
             if data_type is None:
                 raise ValueError("data_type must be provided for A2A skills")
             self.register_skill(a2a_capability, handler, data_type)
 
@@ -121,7 +121,7 @@
             security_schemes={},  # Empty dict for now, can be enhanced later
         )
 
     async def start_server(self, host: str = "0.0.0.0", port: int = 8080):
         """Start A2A HTTP server with direct handler execution."""
         if not self.config:
             raise ValueError("A2A config required for A2A server.")

sdk/eggai/agent.py

@@ -9,7 +9,7 @@
 from .transport import get_default_transport
 from .transport.base import Transport
 
-HANDLERS_IDS = defaultdict(int)
+HANDLERS_IDS: defaultdict[str, int] = defaultdict(int)
 
 PLUGIN_LIST = ["a2a"]
 
@@ -43,7 +43,7 @@
         self._stop_registered = False
 
         # Initialize plugins system
-        self.plugins = {}
+        self.plugins: dict[Any, dict[str, Any]] = {}
 
         # Initialize plugins generically
         for plugin_name in PLUGIN_LIST:
@@ -127,18 +127,30 @@
             # Extract plugin-specific kwargs dynamically and clean them from kwargs
             plugin_found_keys = set()
             for plugin_name in PLUGIN_LIST:
                 for key in list(kwargs.keys()):
                     if key.startswith(plugin_name):
                         kwargs.pop(key)
                         plugin_found_keys.add(plugin_name)
 
-            self._subscriptions.append((channel_name, handler, kwargs))
-
-            # Call plugin subscribe methods if they have relevant kwargs
+            # Call plugin subscribe methods if they have relevant kwargs. Validate
+            # before registering the subscription so a misconfiguration doesn't leave
+            # a half-registered handler in self._subscriptions.
             for plugin_found_key in plugin_found_keys:
-                self.plugins[plugin_found_key]["_instance"].subscribe(
-                    channel_name, handler, **original_kwargs
-                )
+                plugin = self.plugins.get(plugin_found_key)
+                if plugin is None or "_instance" not in plugin:
+                    # Plugin-prefixed kwargs were passed to subscribe(), but the
+                    # plugin is only instantiated when matching kwargs are passed to
+                    # Agent(...). Without that, self.plugins has no entry and the
+                    # access below would raise an opaque KeyError.
+                    raise ValueError(
+                        f"Subscription passed {plugin_found_key!r}-prefixed kwargs, "
+                        f"but the {plugin_found_key!r} plugin is not initialized on "
+                        f"this agent. Pass {plugin_found_key!r} configuration to the "
+                        "Agent(...) constructor to enable it."
+                    )
+                plugin["_instance"].subscribe(channel_name, handler, **original_kwargs)
+
+            self._subscriptions.append((channel_name, handler, kwargs))
 
             return handler
 

sdk/eggai/channel.py

@@ -10,7 +10,7 @@
 from .transport import get_default_transport
 from .transport.base import Transport
 
-HANDLERS_IDS = defaultdict(int)
+HANDLERS_IDS: defaultdict[str, int] = defaultdict(int)
 
 # Environment variable: EGGAI_NAMESPACE
 # Purpose: Prefix for all channel names to enable namespace isolation in shared transports.
@@ -29,7 +29,7 @@
     Connection is established lazily on the first publish or subscription.
     """
 
-    def __init__(self, name: str = None, transport: Transport | None = None):
+    def __init__(self, name: str | None = None, transport: Transport | None = None):
         """
         Initialize a Channel instance.
 
@@ -101,7 +101,7 @@
         )
         HANDLERS_IDS[handler_name] += 1
         kwargs["handler_id"] = f"{handler_name}-{HANDLERS_IDS[handler_name]}"
         await self._get_transport().subscribe(self._name, callback, **kwargs)
         await self._ensure_connected()
 
     async def ensure_exists(self):

sdk/eggai/hooks.py

@@ -10,7 +10,7 @@
 F = TypeVar("F", bound=Callable[..., Coroutine[Any, Any, Any]])
 
 # Global variables for shutdown handling.
-_STOP_CALLBACKS = []
+_STOP_CALLBACKS: list[Callable[..., Any]] = []
 _EXIT_EVENT = None  # will be lazily created
 _SIGNAL_HANDLERS_INSTALLED = False
 _CLEANUP_STARTED = False
@@ -76,7 +76,7 @@
         #     [task.cancel() for task in tasks]
         #     await asyncio.gather(*tasks)
 
     global _SIGNAL_HANDLERS_INSTALLED
     if _SIGNAL_HANDLERS_INSTALLED:
         return
 
@@ -84,10 +84,10 @@
     for sig in HANDLED_SIGNALS:
         try:
             loop.add_signal_handler(
                 sig, lambda s=sig: asyncio.create_task(shutdown(sig, False))
             )
         except NotImplementedError:
             signal.signal(sig, lambda _, __: asyncio.create_task(shutdown(sig, True)))
 
 
 def eggai_main(func: F) -> F:

sdk/eggai/tracing.py

@@ -125,7 +125,7 @@
         )
 
         span_exporter = OTLPSpanExporter(**({"endpoint": endpoint} if endpoint else {}))
     elif hasattr(exporter, "export"):
         span_exporter = exporter  # accept pre-built exporter object (for testing)
     else:  # "otlp" / "otlp-grpc" (default)
         from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
@@ -134,13 +134,14 @@
 
         span_exporter = OTLPSpanExporter(**({"endpoint": endpoint} if endpoint else {}))
 
     provider.add_span_processor(BatchSpanProcessor(span_exporter))
 
     from opentelemetry.trace import ProxyTracerProvider
 
     if isinstance(trace.get_tracer_provider(), ProxyTracerProvider):
         trace.set_tracer_provider(provider)
     else:
+        # Adopt the already-configured provider (API base type, not the SDK one).
         provider = trace.get_tracer_provider()
 
     from importlib.metadata import version
@@ -151,13 +152,13 @@
 
 def apply_traceparent(message: Any, traceparent: str) -> Any:
     """Return message with traceparent set (creates new Pydantic model instance or dict copy)."""
     from pydantic import BaseModel
 
     if isinstance(message, BaseModel):
         return message.model_copy(update={"traceparent": traceparent})
     if isinstance(message, dict):
         return {**message, "traceparent": traceparent}
     return message
 
 
 def make_tracing_wrapper(channel_name: str, handler: Callable) -> Callable:
@@ -181,7 +182,7 @@
         # handler signature, not ours -- the parameter name it passes is
         # whatever the user named it. FastStream 0.6 called positionally;
         # 0.7 calls by keyword. Accept either and forward verbatim.
         message = args[0] if args else next(iter(kwargs.values()), None)
 
         traceparent = (
             message.get("traceparent")

sdk/eggai/transport/base.py

@@ -37,7 +37,7 @@ async def publish(self, channel: str, message: dict[str, Any] | BaseModel):
     async def subscribe(
         self,
         channel: str,
-        callback: Callable[[dict[str, Any]], "asyncio.Future"],
+        handler: Callable[[dict[str, Any]], "asyncio.Future"],
         **kwargs,
     ) -> Callable:
         """

sdk/eggai/transport/defaults.py

@@ -24,10 +24,11 @@
     Get a fresh Transport instance from the default factory.
     If no default transport factory is set, return an InMemoryTransport instance and log a warning.
     """
     if _DEFAULT_TRANSPORT_FACTORY is None:
         logger.warning(
             "No default transport factory set, InMemoryTransport will be used. "
             "Use eggai_set_default_transport() to set a different default transport."
         )
         eggai_set_default_transport(lambda: InMemoryTransport())
+    assert _DEFAULT_TRANSPORT_FACTORY is not None  # set just above if it was None
     return _DEFAULT_TRANSPORT_FACTORY()

sdk/eggai/transport/inmemory.py

@@ -6,7 +6,8 @@
 from collections.abc import Callable
 from typing import Any
 
-from eggai.schemas import BaseMessage
+from pydantic import BaseModel
+
 from eggai.transport import Transport
 from eggai.transport.middleware_utils import wrap_handler_with_filters
 
@@ -90,7 +91,7 @@
         self._consume_tasks.clear()
         self._connected = False
 
-    async def publish(self, channel: str, message: dict[str, Any] | BaseMessage):
+    async def publish(self, channel: str, message: dict[str, Any] | BaseModel):
         """
         Publishes a message to the given channel.
         The message is put into all queues for that channel so each consumer group receives it.
@@ -111,7 +112,7 @@
     async def subscribe(
         self,
         channel: str,
-        callback: Callable[[dict[str, Any]], "asyncio.Future"],
+        handler: Callable[[dict[str, Any]], "asyncio.Future"],
         **kwargs,
     ):
         """
@@ -129,7 +130,7 @@
         # and typed delivery identical across all transports and is the single place
         # that rejects invalid filter-option combinations. Tracing stays outermost.
         final_callback = wrap_handler_with_filters(
-            callback,
+            handler,
             data_type=kwargs.get("data_type"),
             filter_by_data=kwargs.get("filter_by_data"),
             filter_by_message=kwargs.get("filter_by_message"),

sdk/eggai/transport/kafka.py

@@ -5,8 +5,8 @@
 from aiokafka.admin import AIOKafkaAdminClient, NewTopic
 from aiokafka.errors import TopicAlreadyExistsError
 from faststream.kafka import KafkaBroker
+from pydantic import BaseModel
 
-from eggai.schemas import BaseMessage
 from eggai.transport.base import Transport
 from eggai.transport.middleware_utils import wrap_handler_with_filters
 
@@ -141,7 +141,11 @@
             and self.broker._producer
         ):
             try:
-                await self.broker._producer._producer.client.force_metadata_update()
+                # Reaches into aiokafka internals via FastStream's producer — there
+                # is no public API for force_metadata_update(). Guarded by the
+                # hasattr check above and a broad except so a FastStream internals
+                # change degrades gracefully (stale metadata, not a crash).
+                await self.broker._producer._producer.client.force_metadata_update()  # type: ignore[attr-defined]
             except Exception:
                 pass
 
@@ -151,7 +155,7 @@
             await self._ensure_topic(channel)
             self._topics_to_create.add(channel)
 
-    async def publish(self, channel: str, message: dict[str, Any] | BaseMessage):
+    async def publish(self, channel: str, message: dict[str, Any] | BaseModel):
         """
         Publishes a message to the specified Kafka topic (channel).
 

sdk/eggai/transport/middleware_utils.py

@@ -14,7 +14,8 @@
 from collections.abc import Callable
 from typing import Any
 
-from pydantic import ValidationError
+from pydantic import BaseModel, ValidationError
+from pydantic_core import PydanticUndefined
 
 # Identity attributes copied from the user handler onto a wrapper. We intentionally
 # do NOT use functools.wraps here: it sets __wrapped__, which inspect.signature
@@ -52,7 +53,7 @@
 def wrap_handler_with_filters(
     handler: Callable,
     *,
-    data_type: type | None = None,
+    data_type: type[BaseModel] | None = None,
     filter_by_data: Callable[[Any], bool] | None = None,
     filter_by_message: Callable[[dict[str, Any]], bool] | None = None,
 ) -> Callable:
@@ -98,14 +99,24 @@
                 "(the discriminator used to match messages, as on BaseMessage)."
             )
         expected_type = data_type.model_fields["type"].default
+        if expected_type is PydanticUndefined:
+            # The 'type' field exists but has no default, so there is no
+            # discriminator value to match against — every message would fail the
+            # check at line below and be silently dropped. Fail loudly instead.
+            raise ValueError(
+                f"data_type {data_type.__name__!r} has a 'type' field with no "
+                "default, so there is no value to match messages against (every "
+                "message would be silently dropped). Give it a default discriminator, "
+                'e.g. `type: str = "OrderCreated"`.'
+            )
 
         async def typed_handler(message: dict[str, Any]) -> Any:
             try:
                 typed_message = data_type.model_validate(message)
             except (ValidationError, ValueError, TypeError):
                 # Wrong shape / payload for this data_type — not ours to handle.
                 return None
-            if typed_message.type != expected_type:
+            if getattr(typed_message, "type", None) != expected_type:
                 return None
             if filter_by_data is not None and not filter_by_data(typed_message):
                 return None