chore: simplify agent executor implementation (#1730)

Signed-off-by: Radek Ježek <[email protected]>

Author: jezekra1

apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/base.py

@@ -111,7 +111,7 @@ def from_agent_card(cls, agent: a2a.types.AgentCard) -> typing.Self | None:
         return None
 
 
-ExtensionSpecT = typing.TypeVar("ExtensionSpecT", bound=BaseExtensionSpec)
+ExtensionSpecT = typing.TypeVar("ExtensionSpecT", bound=BaseExtensionSpec[typing.Any])
 
 
 class BaseExtensionServer(abc.ABC, typing.Generic[ExtensionSpecT, MetadataFromClientT]):

apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/ui/citation.py

@@ -4,7 +4,6 @@
 from __future__ import annotations
 
 from types import NoneType
-from typing import Any
 
 import pydantic
 from a2a.types import DataPart, FilePart, Part, TextPart
@@ -59,11 +58,7 @@ class CitationExtensionSpec(NoParamsBaseExtensionSpec):
 
 
 class CitationExtensionServer(BaseExtensionServer[CitationExtensionSpec, NoneType]):
-    def citation_metadata(
-        self,
-        *,
-        citations: list[Citation],
-    ) -> Metadata[str, Any]:
+    def citation_metadata(self, *, citations: list[Citation]) -> Metadata:
         return Metadata({self.spec.URI: CitationMetadata(citations=citations).model_dump(mode="json")})
 
     def message(

apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/ui/error.py

@@ -144,7 +144,7 @@ def context(self) -> JsonDict:
             )
             return {}
 
-    def error_metadata(self, error: BaseException) -> Metadata[str, Any]:
+    def error_metadata(self, error: BaseException) -> Metadata:
         """
         Create metadata for an error.
 

apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/ui/trajectory.py

@@ -4,7 +4,6 @@
 from __future__ import annotations
 
 from types import NoneType
-from typing import Any
 
 import pydantic
 from a2a.types import DataPart, FilePart, Part, TextPart
@@ -49,7 +48,7 @@ class TrajectoryExtensionSpec(NoParamsBaseExtensionSpec):
 class TrajectoryExtensionServer(BaseExtensionServer[TrajectoryExtensionSpec, NoneType]):
     def trajectory_metadata(
         self, *, title: str | None = None, content: str | None = None, group_id: str | None = None
-    ) -> Metadata[str, Any]:
+    ) -> Metadata:
         return Metadata(
             {self.spec.URI: Trajectory(title=title, content=content, group_id=group_id).model_dump(mode="json")}
         )

apps/agentstack-sdk-py/src/agentstack_sdk/a2a/types.py

@@ -1,8 +1,7 @@
 # Copyright 2025 © BeeAI a Series of LF Projects, LLC
 # SPDX-License-Identifier: Apache-2.0
-import typing
 import uuid
-from typing import Generic, Literal, TypeAlias, Union
+from typing import TYPE_CHECKING, Literal, TypeAlias
 
 from a2a.types import (
     Artifact,
@@ -20,13 +19,20 @@
     TextPart,
 )
 from pydantic import Field, model_validator
-from typing_extensions import TypeAliasType
 
-K = typing.TypeVar("K")
-V = typing.TypeVar("V")
+if TYPE_CHECKING:
+    JsonValue: TypeAlias = list["JsonValue"] | dict[str, "JsonValue"] | str | bool | int | float | None
+    JsonDict: TypeAlias = dict[str, JsonValue]
+else:
+    from typing import Union
 
+    from typing_extensions import TypeAliasType
 
-class Metadata(dict[K, V], Generic[K, V]): ...
+    JsonValue = TypeAliasType("JsonValue", "Union[dict[str, JsonValue], list[JsonValue], str, int, float, bool, None]")  # noqa: UP007
+    JsonDict = TypeAliasType("JsonDict", "dict[str, JsonValue]")
+
+
+class Metadata(dict[str, JsonValue]): ...
 
 
 RunYield: TypeAlias = (
@@ -43,15 +49,15 @@ class Metadata(dict[K, V], Generic[K, V]): ...
     | TaskStatusUpdateEvent
     | TaskArtifactUpdateEvent
     | str
-    | dict
+    | JsonDict
     | Exception
 )
 RunYieldResume: TypeAlias = Message | None
 
 
 class AgentArtifact(Artifact):
     artifact_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
-    parts: list[Part | TextPart | FilePart | DataPart]  # pyright: ignore [reportIncompatibleVariableOverride]
+    parts: list[Part | TextPart | FilePart | DataPart]
 
     @model_validator(mode="after")
     def text_message_validate(self):
@@ -61,7 +67,7 @@ def text_message_validate(self):
 
 class ArtifactChunk(Artifact):
     last_chunk: bool = False
-    parts: list[Part | TextPart | FilePart | DataPart]  # pyright: ignore [reportIncompatibleVariableOverride]
+    parts: list[Part | TextPart | FilePart | DataPart]
 
     @model_validator(mode="after")
     def text_message_validate(self):
@@ -105,9 +111,3 @@ def text_message_validate(self):
 
 class AuthRequired(InputRequired):
     state: Literal[TaskState.auth_required] = TaskState.auth_required  # pyright: ignore [reportIncompatibleVariableOverride]
-
-
-JsonDict = TypeAliasType(
-    "JsonDict",
-    "Union[dict[str, JsonDict], list[JsonDict], str, int, float, bool, None]",  # pyright: ignore[reportDeprecated]  # noqa: UP007
-)