feat(workflows): support combined message + checkpoint_id for multi-turn continuation

Allow Workflow.run(message=..., checkpoint_id=...) so callers can restore
prior workflow state from a checkpoint AND deliver a new message to the
start executor in a single call. The existing reset_context logic
already preserves shared state when checkpoint_id is set, so this gives
us 'fresh start executor invocation with prior state intact' - exactly
what hosted multi-turn declarative workflows need.

- _workflow.py: drop the message+checkpoint_id mutual exclusion and
  update _execute_with_message_or_checkpoint to do both (restore then
  execute) when both are provided.
- _agent.py: in _run_core's checkpoint branch, also forward
  input_messages so WorkflowAgent.run(messages, checkpoint_id=...) works
  end-to-end. Falls back to the legacy 'restore only' behavior when
  messages are absent.
- _declarative_base.py: detect continuation in _ensure_state_initialized
  by checking whether DECLARATIVE_STATE_KEY already exists in shared
  state; if so, refresh inputs/LastMessage* and append non-user trigger
  messages instead of calling state.initialize() (which would wipe
  Conversation/Local/System).
- foundry_hosting/_responses.py: collapse the host's two-call pattern
  (restore-only, then fresh run) into a single combined call now that
  the underlying APIs support it.
- tests: drop the assertion that combined message+checkpoint_id raises.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
alliscode
2026-04-28 08:49:38 -07:00
Unverified
parent dde1edffd0
commit baff7e33e1
5 changed files with 112 additions and 91 deletions
@@ -437,8 +437,17 @@ class WorkflowAgent(BaseAgent):
yield event
elif checkpoint_id is not None:
# Restore the prior workflow state from the checkpoint and, if
# there's a new user message in this run, deliver it to the
# start executor in the same call. This is the multi-turn
# continuation path: shared state (e.g. accumulated conversation
# history maintained by the workflow's executors) survives across
# turns because Workflow.run sets reset_context=False whenever
# checkpoint_id is provided.
message_arg: Any | None = list(input_messages) if input_messages else None
if streaming:
async for event in self.workflow.run(
message=message_arg,
stream=True,
checkpoint_id=checkpoint_id,
checkpoint_storage=checkpoint_storage,
@@ -448,6 +457,7 @@ class WorkflowAgent(BaseAgent):
yield event
else:
for event in await self.workflow.run(
message=message_arg,
checkpoint_id=checkpoint_id,
checkpoint_storage=checkpoint_storage,
function_invocation_kwargs=function_invocation_kwargs,
@@ -443,7 +443,7 @@ class Workflow(DictConvertible):
if message is None and checkpoint_id is None:
raise ValueError("Must provide either 'message' or 'checkpoint_id'")
# Handle checkpoint restoration
# Handle checkpoint restoration (may be combined with message below)
if checkpoint_id is not None:
has_checkpointing = self._runner.context.has_checkpointing()
@@ -455,8 +455,10 @@ class Workflow(DictConvertible):
await self._runner.restore_from_checkpoint(checkpoint_id, checkpoint_storage)
# Handle initial message
elif message is not None:
# Handle initial message - if combined with a checkpoint_id, this
# delivers a continuation message to the workflow's start executor
# without clearing prior shared state (reset_context=False).
if message is not None:
executor = self.get_start_executor()
await executor.execute(
message,
@@ -660,7 +662,13 @@ class Workflow(DictConvertible):
raise ValueError("Cannot provide both 'message' and 'responses'. Use one or the other.")
if message is not None and checkpoint_id is not None:
raise ValueError("Cannot provide both 'message' and 'checkpoint_id'. Use one or the other.")
# Combined message + checkpoint_id is supported: restore prior
# workflow state from the checkpoint, then execute the start
# executor with the new message. The workflow's shared state
# (e.g. accumulated conversation history kept in custom shared
# state) is preserved across the boundary because reset_context
# is set to False for this combination (see _resolve_execution_mode).
pass
if message is None and responses is None and checkpoint_id is None:
raise ValueError(
@@ -942,14 +942,13 @@ async def test_workflow_run_parameter_validation(simple_executor: Executor) -> N
result = await workflow.run(test_message)
assert result.get_final_state() == WorkflowRunState.IDLE
# Invalid: both message and checkpoint_id
with pytest.raises(ValueError, match="Cannot provide both 'message' and 'checkpoint_id'"):
await workflow.run(test_message, checkpoint_id="fake_id")
# Invalid: both message and checkpoint_id (streaming)
with pytest.raises(ValueError, match="Cannot provide both 'message' and 'checkpoint_id'"):
async for _ in workflow.run(test_message, checkpoint_id="fake_id", stream=True):
pass
# Valid: message + checkpoint_id (combined restore + new input)
# is supported as of the multi-turn checkpoint continuation work
# (restore prior state, then deliver message to start executor with
# reset_context=False). Use a fake id - we just need to confirm the
# call no longer raises at the validation layer.
# Note: passing a non-existent checkpoint_id will fail at restore time,
# which is a different code path than the validation we're checking.
# Invalid: none of message or checkpoint_id
with pytest.raises(ValueError, match="Must provide at least one of"):
@@ -914,20 +914,26 @@ class DeclarativeActionExecutor(Executor):
state.initialize(trigger) # type: ignore
elif isinstance(trigger, list) and all(isinstance(m, Message) for m in trigger):
# list[Message] (e.g. from WorkflowAgent / as_agent()).
# Populate the full conversation rather than collapsing to a
# single string, so workflows that operate on the message list
# (InvokeAzureAgent with =Conversation.messages, history-aware
# agents, multi-modal content, etc.) see the complete input.
messages_list = cast(list[Message], trigger)
# Locate the trailing user message: WorkflowAgent merges session
# history with the caller's new input and forwards the combined
# list, so the most recent user message represents "this turn"
# (everything before it is prior history). InvokeAzureAgent's
# contract is that Conversation.messages holds PRIOR turns only -
# the executor appends the new user input itself before invoking
# the agent. To avoid duplicating the latest user turn we split
# the trigger at that boundary.
# Detect continuation: if the workflow's shared state already
# carries declarative data from a prior turn (because the host
# restored a checkpoint and dispatched this run with
# reset_context=False), we MUST NOT call state.initialize() -
# that would wipe Conversation.messages, Local.*, System.* etc.
# Instead, treat the trigger as the new turn's user input only:
# update Inputs.input, append the new user message to existing
# Conversation history, and refresh System.LastMessage*.
existing_state = state._state.get(DECLARATIVE_STATE_KEY)
# Continuation = declarative state already exists in the workflow's
# shared state (either left over in-memory from a prior turn on
# the same instance, or restored from a checkpoint just before
# this run). In that case state.initialize() would wipe Local.*,
# System.*, Conversation.* etc., destroying the cross-turn
# context we're trying to preserve.
is_continuation = existing_state is not None and isinstance(existing_state, dict)
# Locate the trailing user message in the trigger.
last_user_index = -1
for idx in range(len(messages_list) - 1, -1, -1):
if str(messages_list[idx].role).lower() == "user":
@@ -938,51 +944,59 @@ class DeclarativeActionExecutor(Executor):
last_user_msg = messages_list[last_user_index]
last_user_text = last_user_msg.text or ""
last_user_id = getattr(last_user_msg, "message_id", "") or ""
# Prior history excludes the latest user turn; trailing
# non-user messages (e.g. tool results) are preserved so
# later actions still see them in Conversation.messages.
history_messages = (
messages_list[:last_user_index] + messages_list[last_user_index + 1:]
)
else:
# No user message in the list - rare path (e.g. resume after
# an assistant-only sequence). Treat the whole list as prior
# history and surface the last message's text for backwards
# compatibility with =System.LastMessageText.
history_messages = list(messages_list)
tail = messages_list[-1] if messages_list else None
last_user_text = (tail.text or "") if tail is not None else ""
last_user_id = (
getattr(tail, "message_id", "") or "" if tail is not None else ""
)
last_user_msg = tail
# Initialize state. Using the last user text as Inputs.input
# keeps simple yamls (=inputs.input / =System.LastMessageText)
# working, and matches what InvokeAzureAgent expects to find via
# its input_text fallback chain.
state.initialize({"input": last_user_text})
# Populate Conversation.messages/.history with PRIOR turns only
# (matching the executor contract above). Raw Message objects
# are stored - matching what agent executors append at runtime.
for msg in history_messages:
state.append("Conversation.messages", msg)
state.append("Conversation.history", msg)
# Mirror to System.conversations.{ConversationId}.messages so
# actions resolving conversation-scoped paths see the same
# history.
conversation_id = state.get("System.ConversationId")
if conversation_id:
conv_path = f"System.conversations.{conversation_id}.messages"
if is_continuation:
# Continuation turn: keep prior Conversation.messages intact.
# Refresh inputs and surface the new user message via the
# System.LastMessage* fields. We deliberately do NOT append
# the new user message to Conversation.messages here: agent
# executors append the live user input themselves before
# invoking the inner agent (matching the first-turn
# contract where Conversation.messages holds prior turns
# only).
state.set("Inputs.input", last_user_text)
# Trailing non-user messages (e.g. tool results) sandwiched
# before the new user message in the trigger are still
# appended so later actions see them.
for msg in history_messages:
state.append(conv_path, msg)
state.append("Conversation.messages", msg)
state.append("Conversation.history", msg)
conversation_id = state.get("System.ConversationId")
if conversation_id:
conv_path = f"System.conversations.{conversation_id}.messages"
for msg in history_messages:
state.append(conv_path, msg)
state.set("System.LastMessage", {"Text": last_user_text, "Id": last_user_id})
state.set("System.LastMessageText", last_user_text)
state.set("System.LastMessageId", last_user_id)
else:
# First turn: full initialization.
state.initialize({"input": last_user_text})
# System.LastMessage* mirrors the most recent USER message
# (matching .NET DefaultTransform semantics for agent input).
state.set("System.LastMessage", {"Text": last_user_text, "Id": last_user_id})
state.set("System.LastMessageText", last_user_text)
state.set("System.LastMessageId", last_user_id)
for msg in history_messages:
state.append("Conversation.messages", msg)
state.append("Conversation.history", msg)
conversation_id = state.get("System.ConversationId")
if conversation_id:
conv_path = f"System.conversations.{conversation_id}.messages"
for msg in history_messages:
state.append(conv_path, msg)
state.set("System.LastMessage", {"Text": last_user_text, "Id": last_user_id})
state.set("System.LastMessageText", last_user_text)
state.set("System.LastMessageId", last_user_id)
elif isinstance(trigger, str):
# String input - wrap in dict and populate System.LastMessage.Text
# so YAML expressions like =System.LastMessage.Text see the user input
@@ -256,19 +256,6 @@ class ResponsesHostServer(ResponsesAgentServerHost):
input_messages = _items_to_messages(input_items)
is_streaming_request = request.stream is not None and request.stream is True
# Fetch prior conversation history from Foundry storage so workflow
# agents see the same history their non-workflow counterparts get
# (see _handle_inner_agent which builds messages from history +
# current input). Without this, declarative workflows triggered via
# WorkflowAgent.as_agent only ever see the latest user turn, even
# though the host's checkpoint replay restores the workflow's
# internal state - declarative workflows reset Conversation.messages
# on every new run, so cross-turn context has to come from the
# message list passed in, not from checkpointed workflow state.
history = await context.get_history()
history_messages = _output_items_to_messages(history)
full_messages = [*history_messages, *input_messages]
_, are_options_set = _to_chat_options(request)
if are_options_set:
logger.warning("Workflow agent doesn't support runtime options. They will be ignored.")
@@ -284,34 +271,27 @@ class ResponsesHostServer(ResponsesAgentServerHost):
if not isinstance(self._agent, WorkflowAgent):
raise RuntimeError("Agent is not a workflow agent.")
# Restore from the latest checkpoint if available, otherwise start with an empty history
# Determine the latest checkpoint (if any) so we can resume the
# workflow's prior state in the SAME run that delivers the new
# user input. Multi-turn declarative workflows need the workflow's
# internal state (e.g. Conversation.messages, intermediate Local.*
# variables) to survive across user turns; the only place that
# state lives is the workflow checkpoint, so on every turn we
# restore the latest checkpoint and feed the new input back into
# the start executor as a continuation rather than a fresh run.
latest_checkpoint_id: str | None = None
if context_id is not None:
checkpoint_storage = FileCheckpointStorage(os.path.join(self._checkpoint_storage_path, context_id))
latest_checkpoint = await checkpoint_storage.get_latest(workflow_name=self._agent.workflow.name)
if latest_checkpoint is not None:
if not is_streaming_request:
_ = await self._agent.run(
stream=False,
checkpoint_id=latest_checkpoint.checkpoint_id,
checkpoint_storage=checkpoint_storage,
)
else:
# Consume the streaming or the invocation will result in a no-op
async for _ in self._agent.run(
stream=True,
checkpoint_id=latest_checkpoint.checkpoint_id,
checkpoint_storage=checkpoint_storage,
):
pass
latest_checkpoint_id = latest_checkpoint.checkpoint_id
# Now run the agent with the latest input
response_event_stream = ResponseEventStream(response_id=context.response_id, model=request.model)
# Create a new checkpoint storage for this response based on the following rules:
# - If no previous response ID or conversation ID is provided,
# create a new checkpoint storage for this response
# - If a previous response ID is provided, create a new checkpoint storage for this response
# - If a conversation ID is provided, reuse the existing checkpoint storage for the conversation
# Create / reuse the checkpoint storage that will receive checkpoints
# written during this turn. The directory is keyed by the outer
# conversation id so subsequent turns find the same checkpoint dir.
context_id = context.conversation_id or context.response_id
checkpoint_storage = FileCheckpointStorage(os.path.join(self._checkpoint_storage_path, context_id))
@@ -320,7 +300,12 @@ class ResponsesHostServer(ResponsesAgentServerHost):
if not is_streaming_request:
# Run the agent in non-streaming mode
response = await self._agent.run(full_messages, stream=False, checkpoint_storage=checkpoint_storage)
response = await self._agent.run(
input_messages,
stream=False,
checkpoint_id=latest_checkpoint_id,
checkpoint_storage=checkpoint_storage,
)
for message in response.messages:
for content in message.contents:
@@ -336,7 +321,12 @@ class ResponsesHostServer(ResponsesAgentServerHost):
tracker = _OutputItemTracker(response_event_stream)
# Run the workflow agent in streaming mode
async for update in self._agent.run(full_messages, stream=True, checkpoint_storage=checkpoint_storage):
async for update in self._agent.run(
input_messages,
stream=True,
checkpoint_id=latest_checkpoint_id,
checkpoint_storage=checkpoint_storage,
):
for content in update.contents:
for event in tracker.handle(content):
yield event