Resolve task return types

Author: declan-scale

src/agentex/lib/adk/_modules/tasks.py

@@ -16,6 +16,8 @@
 )
 from agentex.lib.core.tracing.tracer import AsyncTracer
 from agentex.types.task import Task
+from agentex.types.task_retrieve_response import TaskRetrieveResponse
+from agentex.types.task_retrieve_by_name_response import TaskRetrieveByNameResponse
 from agentex.lib.utils.logging import make_logger
 from agentex.lib.utils.temporal import in_temporal_workflow
 
@@ -53,7 +55,7 @@ async def get(
         start_to_close_timeout: timedelta = timedelta(seconds=5),
         heartbeat_timeout: timedelta = timedelta(seconds=5),
         retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
-    ) -> Task:
+    ) -> TaskRetrieveResponse | TaskRetrieveByNameResponse:
         """
         Get a task by ID or name.
         Args:

src/agentex/lib/core/services/adk/tasks.py

@@ -6,6 +6,8 @@
 from agentex.lib.utils.logging import make_logger
 from agentex.lib.utils.temporal import heartbeat_if_in_workflow
 from agentex.lib.core.tracing.tracer import AsyncTracer
+from agentex.types.task_retrieve_response import TaskRetrieveResponse
+from agentex.types.task_retrieve_by_name_response import TaskRetrieveByNameResponse
 
 logger = make_logger(__name__)
 
@@ -25,7 +27,7 @@ async def get_task(
         task_name: str | None = None,
         trace_id: str | None = None,
         parent_span_id: str | None = None,
-    ) -> Task:
+    ) -> TaskRetrieveResponse | TaskRetrieveByNameResponse:
         trace = self._tracer.trace(trace_id)
         async with trace.span(
             parent_id=parent_span_id,

src/agentex/lib/core/temporal/activities/adk/tasks_activities.py

@@ -8,6 +8,8 @@
 from agentex.lib.types.tracing import BaseModelWithTraceParams
 from agentex.lib.utils.logging import make_logger
 from agentex.lib.core.services.adk.tasks import TasksService
+from agentex.types.task_retrieve_response import TaskRetrieveResponse
+from agentex.types.task_retrieve_by_name_response import TaskRetrieveByNameResponse
 
 logger = make_logger(__name__)
 
@@ -32,7 +34,7 @@ def __init__(self, tasks_service: TasksService):
         self._tasks_service = tasks_service
 
     @activity.defn(name=TasksActivityName.GET_TASK)
-    async def get_task(self, params: GetTaskParams) -> Task | None:
+    async def get_task(self, params: GetTaskParams) -> TaskRetrieveResponse | TaskRetrieveByNameResponse:
         return await self._tasks_service.get_task(
             task_id=params.task_id,
             task_name=params.task_name,