11# ruff: noqa: I001
22# Import order matters - AsyncTracer must come after client import to avoid circular imports
33from __future__ import annotations
4- from datetime import timedelta
4+ from datetime import datetime , timedelta
55
66from temporalio .common import RetryPolicy
77
2222from agentex .lib .core .tracing .tracer import AsyncTracer
2323from agentex .types .task_message import TaskMessage , TaskMessageContent
2424from agentex .lib .utils .logging import make_logger
25- from agentex .lib .utils .temporal import in_temporal_workflow
25+ from agentex .lib .utils .temporal import in_temporal_workflow , workflow_now_if_in_workflow
2626
2727logger = make_logger (__name__ )
2828
@@ -66,6 +66,7 @@ async def create(
6666 start_to_close_timeout : timedelta = timedelta (seconds = 5 ),
6767 heartbeat_timeout : timedelta = timedelta (seconds = 5 ),
6868 retry_policy : RetryPolicy = DEFAULT_RETRY_POLICY ,
69+ created_at : datetime | None = None ,
6970 ) -> TaskMessage :
7071 """
7172 Create a new message for a task.
@@ -82,12 +83,17 @@ async def create(
8283 Returns:
8384 TaskMessageEntity: The created message.
8485 """
86+ # Default created_at to workflow.now() so two awaited adk.messages.create
87+ # calls from the same workflow are guaranteed monotonic at the server.
88+ if created_at is None :
89+ created_at = workflow_now_if_in_workflow ()
8590 params = CreateMessageParams (
8691 trace_id = trace_id ,
8792 parent_span_id = parent_span_id ,
8893 task_id = task_id ,
8994 content = content ,
9095 emit_updates = emit_updates ,
96+ created_at = created_at ,
9197 )
9298 if in_temporal_workflow ():
9399 return await ActivityHelpers .execute_activity (
@@ -103,6 +109,7 @@ async def create(
103109 task_id = task_id ,
104110 content = content ,
105111 emit_updates = emit_updates ,
112+ created_at = created_at ,
106113 )
107114
108115 async def update (
@@ -163,6 +170,7 @@ async def create_batch(
163170 start_to_close_timeout : timedelta = timedelta (seconds = 5 ),
164171 heartbeat_timeout : timedelta = timedelta (seconds = 5 ),
165172 retry_policy : RetryPolicy = DEFAULT_RETRY_POLICY ,
173+ created_at : datetime | None = None ,
166174 ) -> list [TaskMessage ]:
167175 """
168176 Create a batch of messages for a task.
@@ -177,12 +185,15 @@ async def create_batch(
177185 Returns:
178186 List[TaskMessageEntity]: The created messages.
179187 """
188+ if created_at is None :
189+ created_at = workflow_now_if_in_workflow ()
180190 params = CreateMessagesBatchParams (
181191 task_id = task_id ,
182192 contents = contents ,
183193 emit_updates = emit_updates ,
184194 trace_id = trace_id ,
185195 parent_span_id = parent_span_id ,
196+ created_at = created_at ,
186197 )
187198 if in_temporal_workflow ():
188199 return await ActivityHelpers .execute_activity (
@@ -198,6 +209,7 @@ async def create_batch(
198209 task_id = task_id ,
199210 contents = contents ,
200211 emit_updates = emit_updates ,
212+ created_at = created_at ,
201213 )
202214
203215 async def update_batch (
0 commit comments