Skip to content

Conversation

@adtyavrdhn
Copy link
Contributor

@adtyavrdhn adtyavrdhn commented Nov 7, 2025

Closes #3171

Adds a UUID run_id at agent start, thread it through GraphAgentState, and assign it to every ModelRequest/ModelResponse as they’re created so all execution paths share one identifier.

Keep run_id optional on message classes for backward compatibility.

Cover the change with tests:
(1) assert captured messages and result.all_messages() all carry the same non-empty run_id and
(2) confirm ModelMessagesTypeAdapter preserves the field through serialization.

@adtyavrdhn adtyavrdhn changed the title Adding run_id to ModelRequest and ModelResponse, propagating it via G… Adding run_id to ModelRequest and ModelResponse Nov 7, 2025
@adtyavrdhn
Copy link
Contributor Author

@DouweM A lot of tests will need to be modified to accommodate run_id in the snapshots.

ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
response: _messages.ModelResponse,
) -> CallToolsNode[DepsT, NodeRunEndT]:
response.run_id = response.run_id or ctx.state.run_id
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We may also need to do this in StreamedRunResult._marked_completed to get this to work right when using agent.run_stream:

if message is not None:
self._all_messages.append(message)

You'll want to test if it currently works, and if not add it.

We also definitely need to set the run_id here:

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah I made it a part of RunContext and passed it, added for the others two as well

@DouweM
Copy link
Collaborator

DouweM commented Nov 7, 2025

@DouweM A lot of tests will need to be modified to accommodate run_id in the snapshots.

@adtyavrdhn Yep, we'll need run_id=IsStr() everywhere.

@adtyavrdhn
Copy link
Contributor Author

This ended up being a much bigger diff than what I was expecting

def serialize_run_context(cls, ctx: RunContext[Any]) -> dict[str, Any]:
"""Serialize the run context to a `dict[str, Any]`."""
return {
'run_id': ctx.run_id,
Copy link
Contributor Author

@adtyavrdhn adtyavrdhn Nov 8, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added it as a part of serialization for temporal and I added a small test for the round trip.

@adtyavrdhn adtyavrdhn requested a review from DouweM November 8, 2025 13:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add run_id field to ModelMessage for reliable run grouping

2 participants