fix(messages): stamp agent messages with workflow.now() for monotonic ordering

examples/tutorials/10_async/00_base/010_multiturn/tests/test_agent.py

@@ -145,6 +145,7 @@ async def test_send_event_and_stream(self, client: AsyncAgentex, agent_id: str):
         task = task_response.result
         assert task is not None
 
+        await asyncio.sleep(1)  # wait for state to be initialized
         # Check initial state
         states = await client.states.list(agent_id=agent_id, task_id=task.id)
         assert len(states) == 1

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

@@ -1,7 +1,7 @@
 # ruff: noqa: I001
 # Import order matters - AsyncTracer must come after client import to avoid circular imports
 from __future__ import annotations
-from datetime import timedelta
+from datetime import datetime, timedelta
 
 from temporalio.common import RetryPolicy
 
@@ -22,7 +22,7 @@
 from agentex.lib.core.tracing.tracer import AsyncTracer
 from agentex.types.task_message import TaskMessage, TaskMessageContent
 from agentex.lib.utils.logging import make_logger
-from agentex.lib.utils.temporal import in_temporal_workflow
+from agentex.lib.utils.temporal import in_temporal_workflow, workflow_now_if_in_workflow
 
 logger = make_logger(__name__)
 
@@ -66,6 +66,7 @@ async def create(
         start_to_close_timeout: timedelta = timedelta(seconds=5),
         heartbeat_timeout: timedelta = timedelta(seconds=5),
         retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+        created_at: datetime | None = None,
     ) -> TaskMessage:
         """
         Create a new message for a task.
@@ -82,12 +83,17 @@ async def create(
         Returns:
             TaskMessageEntity: The created message.
         """
+        # Default created_at to workflow.now() so two awaited adk.messages.create
+        # calls from the same workflow are guaranteed monotonic at the server.
+        if created_at is None:
+            created_at = workflow_now_if_in_workflow()
         params = CreateMessageParams(
             trace_id=trace_id,
             parent_span_id=parent_span_id,
             task_id=task_id,
             content=content,
             emit_updates=emit_updates,
+            created_at=created_at,
         )
         if in_temporal_workflow():
             return await ActivityHelpers.execute_activity(
@@ -103,6 +109,7 @@ async def create(
                 task_id=task_id,
                 content=content,
                 emit_updates=emit_updates,
+                created_at=created_at,
             )
 
     async def update(
@@ -163,6 +170,7 @@ async def create_batch(
         start_to_close_timeout: timedelta = timedelta(seconds=5),
         heartbeat_timeout: timedelta = timedelta(seconds=5),
         retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+        created_at: datetime | None = None,
     ) -> list[TaskMessage]:
         """
         Create a batch of messages for a task.
@@ -177,12 +185,15 @@ async def create_batch(
         Returns:
             List[TaskMessageEntity]: The created messages.
         """
+        if created_at is None:
+            created_at = workflow_now_if_in_workflow()
         params = CreateMessagesBatchParams(
             task_id=task_id,
             contents=contents,
             emit_updates=emit_updates,
             trace_id=trace_id,
             parent_span_id=parent_span_id,
+            created_at=created_at,
         )
         if in_temporal_workflow():
             return await ActivityHelpers.execute_activity(
@@ -198,6 +209,7 @@ async def create_batch(
                 task_id=task_id,
                 contents=contents,
                 emit_updates=emit_updates,
+                created_at=created_at,
             )
 
     async def update_batch(

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

@@ -1,6 +1,7 @@
 # ruff: noqa: I001
 # Import order matters - AsyncTracer must come after client import to avoid circular imports
 from __future__ import annotations
+from datetime import datetime
 from temporalio.common import RetryPolicy
 
 from agentex import AsyncAgentex  # noqa: F401
@@ -52,6 +53,7 @@ def streaming_task_message_context(
         task_id: str,
         initial_content: TaskMessageContent,
         streaming_mode: StreamingMode = "coalesced",
+        created_at: datetime | None = None,
     ) -> StreamingTaskMessageContext:
         """
         Create a streaming context for managing TaskMessage lifecycle.
@@ -83,4 +85,5 @@ def streaming_task_message_context(
             task_id=task_id,
             initial_content=initial_content,
             streaming_mode=streaming_mode,
+            created_at=created_at,
         )

src/agentex/lib/adk/providers/_modules/litellm.py

@@ -6,7 +6,7 @@
 from temporalio.common import RetryPolicy
 
 from agentex.lib.utils.logging import make_logger
-from agentex.lib.utils.temporal import in_temporal_workflow
+from agentex.lib.utils.temporal import in_temporal_workflow, workflow_now_if_in_workflow
 from agentex.types.task_message import TaskMessage
 from agentex.lib.types.llm_messages import LLMConfig, Completion
 from agentex.lib.core.tracing.tracer import AsyncTracer
@@ -88,9 +88,7 @@ async def chat_completion(
             Completion: An OpenAI compatible Completion object
         """
         if in_temporal_workflow():
-            params = ChatCompletionParams(
-                trace_id=trace_id, parent_span_id=parent_span_id, llm_config=llm_config
-            )
+            params = ChatCompletionParams(trace_id=trace_id, parent_span_id=parent_span_id, llm_config=llm_config)
             return await ActivityHelpers.execute_activity(
                 activity_name=LiteLLMActivityName.CHAT_COMPLETION,
                 request=params,
@@ -138,6 +136,7 @@ async def chat_completion_auto_send(
                 parent_span_id=parent_span_id,
                 task_id=task_id,
                 llm_config=llm_config,
+                created_at=workflow_now_if_in_workflow(),
             )
             return await ActivityHelpers.execute_activity(
                 activity_name=LiteLLMActivityName.CHAT_COMPLETION_AUTO_SEND,
@@ -222,6 +221,7 @@ async def chat_completion_stream_auto_send(
                 parent_span_id=parent_span_id,
                 task_id=task_id,
                 llm_config=llm_config,
+                created_at=workflow_now_if_in_workflow(),
             )
             return await ActivityHelpers.execute_activity(
                 activity_name=LiteLLMActivityName.CHAT_COMPLETION_STREAM_AUTO_SEND,

src/agentex/lib/adk/providers/_modules/openai.py

@@ -20,7 +20,7 @@
     from typing_extensions import deprecated
 
 from agentex.lib.utils.logging import make_logger
-from agentex.lib.utils.temporal import in_temporal_workflow
+from agentex.lib.utils.temporal import in_temporal_workflow, workflow_now_if_in_workflow
 from agentex.lib.core.tracing.tracer import AsyncTracer
 from agentex.lib.types.agent_results import (
     SerializableRunResult,
@@ -265,6 +265,7 @@ async def run_agent_auto_send(
                 output_guardrails=output_guardrails,  # type: ignore[arg-type]
                 max_turns=max_turns,
                 previous_response_id=previous_response_id,
+                created_at=workflow_now_if_in_workflow(),
             )
             return await ActivityHelpers.execute_activity(
                 activity_name=OpenAIActivityName.RUN_AGENT_AUTO_SEND,
@@ -479,6 +480,7 @@ async def run_agent_streamed_auto_send(
                 input_guardrails=input_guardrails,
                 output_guardrails=output_guardrails,
                 max_turns=max_turns,
+                created_at=workflow_now_if_in_workflow(),
             )
             return await ActivityHelpers.execute_activity(
                 activity_name=OpenAIActivityName.RUN_AGENT_STREAMED_AUTO_SEND,
@@ -509,4 +511,4 @@ async def run_agent_streamed_auto_send(
                 output_guardrails=output_guardrails,
                 max_turns=max_turns,
                 previous_response_id=previous_response_id,
-            )
+            )

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

@@ -2,8 +2,10 @@
 
 import asyncio
 from typing import Any, Coroutine
+from datetime import datetime
 
 from agentex import AsyncAgentex
+from agentex._types import omit
 from agentex.lib.utils.logging import make_logger
 from agentex.lib.utils.temporal import heartbeat_if_in_workflow
 from agentex.types.task_message import TaskMessage, TaskMessageContent
@@ -32,6 +34,7 @@ async def create_message(
         emit_updates: bool = True,
         trace_id: str | None = None,
         parent_span_id: str | None = None,
+        created_at: datetime | None = None,
     ) -> TaskMessage:
         trace = self._tracer.trace(trace_id)
         async with trace.span(
@@ -43,6 +46,7 @@ async def create_message(
             task_message = await self._agentex_client.messages.create(
                 task_id=task_id,
                 content=content.model_dump(),
+                created_at=created_at if created_at is not None else omit,
             )
             if emit_updates:
                 await self._emit_updates([task_message])
@@ -85,6 +89,7 @@ async def create_messages_batch(
         emit_updates: bool = True,
         trace_id: str | None = None,
         parent_span_id: str | None = None,
+        created_at: datetime | None = None,
     ) -> list[TaskMessage]:
         trace = self._tracer.trace(trace_id)
         async with trace.span(
@@ -96,6 +101,7 @@ async def create_messages_batch(
             task_messages = await self._agentex_client.messages.batch.create(
                 task_id=task_id,
                 contents=[content.model_dump() for content in contents],
+                created_at=created_at if created_at is not None else omit,
             )
             if emit_updates:
                 await self._emit_updates(task_messages)
@@ -119,10 +125,7 @@ async def update_messages_batch(
             heartbeat_if_in_workflow("update messages batch")
             task_messages = await self._agentex_client.messages.batch.update(
                 task_id=task_id,
-                updates={
-                    message_id: content.model_dump()
-                    for message_id, content in updates.items()
-                },
+                updates={message_id: content.model_dump() for message_id, content in updates.items()},
             )
             if span:
                 span.output = [task_message.model_dump() for task_message in task_messages]

src/agentex/lib/core/services/adk/providers/litellm.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from datetime import datetime
 from collections.abc import AsyncGenerator
 
 from agentex import AsyncAgentex
@@ -63,6 +64,7 @@ async def chat_completion_auto_send(
         llm_config: LLMConfig,
         trace_id: str | None = None,
         parent_span_id: str | None = None,
+        created_at: datetime | None = None,
     ) -> TaskMessage | None:
         """
         Chat completion with automatic TaskMessage creation. This does not stream the completion. To stream use chat_completion_stream_auto_send.
@@ -98,13 +100,10 @@ async def chat_completion_auto_send(
                     content="",
                     format="markdown",
                 ),
+                created_at=created_at,
             ) as streaming_context:
                 completion = await self.llm_gateway.acompletion(**llm_config.model_dump())
-                if (
-                    completion.choices
-                    and len(completion.choices) > 0
-                    and completion.choices[0].message
-                ):
+                if completion.choices and len(completion.choices) > 0 and completion.choices[0].message:
                     final_content = TextContent(
                         author="agent",
                         content=completion.choices[0].message.content or "",
@@ -159,9 +158,7 @@ async def chat_completion_stream(
         ) as span:
             # Direct streaming outside temporal - yield each chunk as it comes
             chunks: list[Completion] = []
-            async for chunk in self.llm_gateway.acompletion_stream(
-                **llm_config.model_dump()
-            ):
+            async for chunk in self.llm_gateway.acompletion_stream(**llm_config.model_dump()):
                 chunks.append(chunk)
                 yield chunk
             if span:
@@ -173,6 +170,7 @@ async def chat_completion_stream_auto_send(
         llm_config: LLMConfig,
         trace_id: str | None = None,
         parent_span_id: str | None = None,
+        created_at: datetime | None = None,
     ) -> TaskMessage | None:
         """
         Stream chat completion with automatic TaskMessage creation and streaming.
@@ -206,18 +204,13 @@ async def chat_completion_stream_auto_send(
                     content="",
                     format="markdown",
                 ),
+                created_at=created_at,
             ) as streaming_context:
                 # Get the streaming response
                 chunks = []
-                async for response in self.llm_gateway.acompletion_stream(
-                    **llm_config.model_dump()
-                ):
+                async for response in self.llm_gateway.acompletion_stream(**llm_config.model_dump()):
                     heartbeat_if_in_workflow("chat completion streaming")
-                    if (
-                        response.choices
-                        and len(response.choices) > 0
-                        and response.choices[0].delta
-                    ):
+                    if response.choices and len(response.choices) > 0 and response.choices[0].delta:
                         delta = response.choices[0].delta.content
                         if delta:
                             # Stream the chunk via the context manager
@@ -235,11 +228,7 @@ async def chat_completion_stream_auto_send(
 
                 # Update the final message content
                 complete_message = concat_completion_chunks(chunks)
-                if (
-                    complete_message
-                    and complete_message.choices
-                    and complete_message.choices[0].message
-                ):
+                if complete_message and complete_message.choices and complete_message.choices[0].message:
                     final_content = TextContent(
                         author="agent",
                         content=complete_message.choices[0].message.content or "",