diff --git a/python/packages/core/agent_framework/_harness/_loop.py b/python/packages/core/agent_framework/_harness/_loop.py index 62200518ee..7b14765868 100644 --- a/python/packages/core/agent_framework/_harness/_loop.py +++ b/python/packages/core/agent_framework/_harness/_loop.py @@ -34,7 +34,9 @@ from pydantic import BaseModel, Field from typing_extensions import Self -from .._agents import _LOOP_ITERATION_TOKEN_KEY # pyright: ignore[reportPrivateUsage] -- shared loop-turn marker, see _agents.py +from .._agents import ( + _LOOP_ITERATION_TOKEN_KEY, # pyright: ignore[reportPrivateUsage] -- shared loop-turn marker, see _agents.py +) from .._feature_stage import ExperimentalFeature, experimental from .._middleware import AgentContext, AgentMiddleware, MiddlewareTermination from .._sessions import SessionContext @@ -490,8 +492,7 @@ async def _fire_turn_scoped_after_providers( if response is None or run_after is None: return if not any( - getattr(provider, "after_run_once_per_turn", False) - for provider in getattr(agent, "context_providers", []) + getattr(provider, "after_run_once_per_turn", False) for provider in getattr(agent, "context_providers", []) ): return session_context = SessionContext( diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index 09bd592e92..353a1a2efb 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -713,11 +713,7 @@ def _process_request_info_event( Note: Text requests use the function-call envelope so callers can reply with a matching function result. """ - if ( - isinstance(event.data, Content) - and event.data.user_input_request - and event.data.type != "text" - ): + if isinstance(event.data, Content) and event.data.user_input_request and event.data.type != "text": # Preserve specialized requests that callers already understand how to present. return event.data diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 3de9460c86..62f225b04a 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -130,7 +130,7 @@ class CheckpointStorage(Protocol): """Protocol for checkpoint storage backends.""" async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: - """Save a checkpoint and return its ID. + """Create a copy of the given checkpoint and store it, returning its ID. Args: checkpoint: The WorkflowCheckpoint object to save. @@ -147,7 +147,7 @@ async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint: checkpoint_id: The unique ID of the checkpoint to load. Returns: - The WorkflowCheckpoint object corresponding to the given ID. + A copy of the WorkflowCheckpoint object corresponding to the given ID. Raises: WorkflowCheckpointException: If no checkpoint with the given ID exists. @@ -161,7 +161,7 @@ async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoi workflow_name: The name of the workflow to list checkpoints for. Returns: - A list of WorkflowCheckpoint objects for the specified workflow name. + A list of copies of WorkflowCheckpoint objects for the specified workflow name. """ ... @@ -183,7 +183,8 @@ async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: workflow_name: The name of the workflow to get the latest checkpoint for. Returns: - The latest WorkflowCheckpoint object for the specified workflow name, or None if no checkpoints exist. + A copy of the latest WorkflowCheckpoint object for the specified workflow name, + or None if no checkpoints exist. """ ... @@ -207,7 +208,7 @@ def __init__(self) -> None: self._checkpoints: dict[CheckpointID, WorkflowCheckpoint] = {} async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: - """Save a checkpoint and return its ID.""" + """Create a copy of the given checkpoint and store it, returning its ID.""" self._checkpoints[checkpoint.checkpoint_id] = copy.deepcopy(checkpoint) logger.debug(f"Saved checkpoint {checkpoint.checkpoint_id} to memory") return checkpoint.checkpoint_id @@ -217,12 +218,12 @@ async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint: checkpoint = self._checkpoints.get(checkpoint_id) if checkpoint: logger.debug(f"Loaded checkpoint {checkpoint_id} from memory") - return checkpoint + return copy.deepcopy(checkpoint) raise WorkflowCheckpointException(f"No checkpoint found with ID {checkpoint_id}") async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoint]: """List checkpoint objects for a given workflow name.""" - return [cp for cp in self._checkpoints.values() if cp.workflow_name == workflow_name] + return [copy.deepcopy(cp) for cp in self._checkpoints.values() if cp.workflow_name == workflow_name] async def delete(self, checkpoint_id: CheckpointID) -> bool: """Delete a checkpoint by ID.""" @@ -239,7 +240,7 @@ async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: return None latest_checkpoint = max(checkpoints, key=lambda cp: datetime.fromisoformat(cp.timestamp)) logger.debug(f"Latest checkpoint for workflow {workflow_name} is {latest_checkpoint.checkpoint_id}") - return latest_checkpoint + return copy.deepcopy(latest_checkpoint) async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID]: """List checkpoint IDs. If workflow_id is provided, filter by that workflow.""" diff --git a/python/packages/core/agent_framework/_workflows/_state.py b/python/packages/core/agent_framework/_workflows/_state.py index 093cfea8b6..64759da72c 100644 --- a/python/packages/core/agent_framework/_workflows/_state.py +++ b/python/packages/core/agent_framework/_workflows/_state.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import copy from typing import Any @@ -31,7 +32,8 @@ def set(self, key: str, value: Any) -> None: """Set a value in the pending state buffer. The value will be visible to subsequent `get()` calls but won't be - committed to the actual state until `commit()` is called. + committed to the actual state until `commit()` is called. A deep copy is + stored so later caller mutations do not change pending state. Note: When multiple executors run concurrently within the same superstep, @@ -40,7 +42,7 @@ def set(self, key: str, value: Any) -> None: the .NET behavior and the superstep execution model where all executors in a superstep see the same committed state at the start. """ - self._pending[key] = value + self._pending[key] = copy.deepcopy(value) def get(self, key: str, default: Any = None) -> Any: """Get a value from state, checking pending first then committed. @@ -50,14 +52,17 @@ def get(self, key: str, default: Any = None) -> Any: default: Value to return if key is not found. Defaults to None. Returns: - The value if found, otherwise the default value. + A deep copy of the value if found, otherwise the default value. Mutate + the returned value and pass it to :meth:`set` to update workflow state. """ if key in self._pending: value = self._pending[key] if value is _DeleteSentinel: return default - return value - return self._committed.get(key, default) + return copy.deepcopy(value) + if key in self._committed: + return copy.deepcopy(self._committed[key]) + return default def has(self, key: str) -> bool: """Check if a key exists in pending or committed state.""" @@ -104,18 +109,20 @@ def discard(self) -> None: self._pending.clear() def export_state(self) -> dict[str, Any]: - """Export a serialized copy of the committed state. + """Export a deepcopy of the committed state. - Note: Does not include pending changes. + Note: + Does not include pending changes. Values must support :func:`copy.deepcopy`. """ - return dict(self._committed) + return copy.deepcopy(self._committed) def import_state(self, state: dict[str, Any]) -> None: """Import state from a serialized dictionary. - Merges into committed state. Does not affect pending changes. + Merges a deepcopy into committed state. Does not affect pending changes. + Values must support :func:`copy.deepcopy`. """ - self._committed.update(state) + self._committed.update(copy.deepcopy(state)) class _DeleteSentinelType: diff --git a/python/packages/core/tests/workflow/test_checkpoint_storage_conformance.py b/python/packages/core/tests/workflow/test_checkpoint_storage_conformance.py new file mode 100644 index 0000000000..e1ce83d81f --- /dev/null +++ b/python/packages/core/tests/workflow/test_checkpoint_storage_conformance.py @@ -0,0 +1,154 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Conformance tests for the CheckpointStorage ownership contract.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, cast + +import pytest + +from agent_framework import ( + CheckpointStorage, + FileCheckpointStorage, + InMemoryCheckpointStorage, + WorkflowCheckpoint, + WorkflowCheckpointException, + WorkflowEvent, +) +from agent_framework._workflows._runner_context import WorkflowMessage + + +@pytest.fixture(params=["memory", "file"]) +def conformance_storage(request: pytest.FixtureRequest, tmp_path: Path) -> CheckpointStorage: + """Yield each in-tree checkpoint storage backend.""" + if request.param == "memory": + return InMemoryCheckpointStorage() + return FileCheckpointStorage(tmp_path / "conformance") + + +def _conformance_checkpoint(workflow_name: str = "conformance-workflow") -> WorkflowCheckpoint: + return WorkflowCheckpoint( + workflow_name=workflow_name, + graph_signature_hash="conformance-hash", + state={ + "shared": {"counter": 0, "history": ["initial"]}, + "_executor_state": {"executor1": {"visits": ["first"]}}, + }, + messages={ + "executor1": [ + WorkflowMessage(data={"text": "hello", "tags": ["initial"]}, source_id="src", target_id="tgt") + ] + }, + pending_request_info_events={ + "req1": WorkflowEvent.request_info( + request_id="req1", + source_executor_id="executor1", + request_data={"payload": ["initial"]}, + response_type=str, + ) + }, + metadata={"tags": ["initial"]}, + ) + + +def _mutate(checkpoint: WorkflowCheckpoint) -> None: + checkpoint.state["shared"]["counter"] = 999 + checkpoint.state["shared"]["history"].append("mutated") + checkpoint.state["_executor_state"]["executor1"]["visits"].append("mutated") + cast(dict[str, Any], checkpoint.messages["executor1"][0].data)["tags"].append("mutated") + cast(dict[str, Any], checkpoint.pending_request_info_events["req1"].data)["payload"].append("mutated") + checkpoint.metadata["tags"].append("mutated") + + +def _assert_pristine(checkpoint: WorkflowCheckpoint) -> None: + assert checkpoint.state["shared"] == {"counter": 0, "history": ["initial"]} + assert checkpoint.state["_executor_state"]["executor1"]["visits"] == ["first"] + assert cast(dict[str, Any], checkpoint.messages["executor1"][0].data)["tags"] == ["initial"] + assert cast(dict[str, Any], checkpoint.pending_request_info_events["req1"].data)["payload"] == ["initial"] + assert checkpoint.metadata["tags"] == ["initial"] + + +async def test_load_returns_caller_owned_copy(conformance_storage: CheckpointStorage) -> None: + checkpoint = _conformance_checkpoint() + await conformance_storage.save(checkpoint) + + loaded = await conformance_storage.load(checkpoint.checkpoint_id) + _mutate(loaded) + + _assert_pristine(await conformance_storage.load(checkpoint.checkpoint_id)) + + +async def test_repeated_loads_are_independent(conformance_storage: CheckpointStorage) -> None: + checkpoint = _conformance_checkpoint() + await conformance_storage.save(checkpoint) + + first = await conformance_storage.load(checkpoint.checkpoint_id) + second = await conformance_storage.load(checkpoint.checkpoint_id) + assert first is not second + + _mutate(first) + _assert_pristine(second) + + +async def test_get_latest_returns_caller_owned_copy(conformance_storage: CheckpointStorage) -> None: + checkpoint = _conformance_checkpoint() + await conformance_storage.save(checkpoint) + + latest = await conformance_storage.get_latest(workflow_name=checkpoint.workflow_name) + assert latest is not None + _mutate(latest) + + reloaded = await conformance_storage.get_latest(workflow_name=checkpoint.workflow_name) + assert reloaded is not None + _assert_pristine(reloaded) + + +async def test_list_checkpoints_returns_caller_owned_copies(conformance_storage: CheckpointStorage) -> None: + checkpoint = _conformance_checkpoint() + await conformance_storage.save(checkpoint) + + listed = await conformance_storage.list_checkpoints(workflow_name=checkpoint.workflow_name) + assert len(listed) == 1 + _mutate(listed[0]) + + relisted = await conformance_storage.list_checkpoints(workflow_name=checkpoint.workflow_name) + assert len(relisted) == 1 + _assert_pristine(relisted[0]) + + +async def test_save_snapshots_at_call_time(conformance_storage: CheckpointStorage) -> None: + checkpoint = _conformance_checkpoint() + await conformance_storage.save(checkpoint) + + _mutate(checkpoint) + + _assert_pristine(await conformance_storage.load(checkpoint.checkpoint_id)) + + +async def test_list_checkpoints_filters_by_workflow_name(conformance_storage: CheckpointStorage) -> None: + first = _conformance_checkpoint("workflow-a") + second = _conformance_checkpoint("workflow-b") + await conformance_storage.save(first) + await conformance_storage.save(second) + + listed_a = await conformance_storage.list_checkpoints(workflow_name="workflow-a") + listed_b = await conformance_storage.list_checkpoints(workflow_name="workflow-b") + + assert {checkpoint.checkpoint_id for checkpoint in listed_a} == {first.checkpoint_id} + assert {checkpoint.checkpoint_id for checkpoint in listed_b} == {second.checkpoint_id} + + +async def test_get_latest_returns_none_when_empty(conformance_storage: CheckpointStorage) -> None: + assert await conformance_storage.get_latest(workflow_name="missing-workflow") is None + + +async def test_load_missing_id_raises(conformance_storage: CheckpointStorage) -> None: + with pytest.raises(WorkflowCheckpointException): + await conformance_storage.load("does-not-exist") + + +async def test_save_returns_id(conformance_storage: CheckpointStorage) -> None: + checkpoint = _conformance_checkpoint() + assert await conformance_storage.save(checkpoint) == checkpoint.checkpoint_id diff --git a/python/packages/core/tests/workflow/test_runner.py b/python/packages/core/tests/workflow/test_runner.py index 3690cdaa1f..2299c38243 100644 --- a/python/packages/core/tests/workflow/test_runner.py +++ b/python/packages/core/tests/workflow/test_runner.py @@ -542,7 +542,7 @@ async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: # Establish some state to capture. executor.count = 7 - state.set("shared_key", "shared_value") + state.set("shared_key", {"history": ["shared_value"]}) state.commit() checkpoint = await runner.build_checkpoint() @@ -550,13 +550,19 @@ async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: # Mutate after capture; restoring must roll back to the captured snapshot. executor.count = 999 - state.set("shared_key", "mutated") + shared_state = state.get("shared_key") + shared_state["history"].append("mutated") + state.set("shared_key", shared_state) state.commit() await runner.restore_checkpoint(checkpoint) assert executor.count == 7 - assert state.get("shared_key") == "shared_value" + assert state.get("shared_key") == {"history": ["shared_value"]} + restored_state = state.get("shared_key") + restored_state["history"].append("restored-mutation") + state.set("shared_key", restored_state) + assert checkpoint.state["shared_key"] == {"history": ["shared_value"]} assert runner._previous_checkpoint_id == checkpoint.checkpoint_id # pyright: ignore[reportPrivateUsage] diff --git a/python/packages/core/tests/workflow/test_state.py b/python/packages/core/tests/workflow/test_state.py index 7781eb4141..db9dd17ee5 100644 --- a/python/packages/core/tests/workflow/test_state.py +++ b/python/packages/core/tests/workflow/test_state.py @@ -15,6 +15,45 @@ def test_set_and_get(self) -> None: state.set("key", "value") assert state.get("key") == "value" + def test_set_does_not_alias_caller_value(self) -> None: + state = State() + value = {"history": ["step-1"]} + + state.set("key", value) + value["history"].append("step-2") + + assert state.get("key") == {"history": ["step-1"]} + + def test_get_does_not_expose_pending_value(self) -> None: + state = State() + state.set("key", {"history": ["step-1"]}) + + value = state.get("key") + value["history"].append("step-2") + + assert state.get("key") == {"history": ["step-1"]} + + def test_get_does_not_expose_committed_value(self) -> None: + state = State() + state.set("key", {"history": ["step-1"]}) + state.commit() + + value = state.get("key") + value["history"].append("step-2") + + assert state.get("key") == {"history": ["step-1"]} + + def test_get_mutate_set_updates_state(self) -> None: + state = State() + state.set("key", {"history": ["step-1"]}) + state.commit() + + value = state.get("key") + value["history"].append("step-2") + state.set("key", value) + + assert state.get("key") == {"history": ["step-1", "step-2"]} + def test_get_with_default(self) -> None: state = State() assert state.get("missing") is None @@ -301,3 +340,55 @@ def test_import_does_not_affect_pending(self) -> None: # Pending is still there assert state.get("pending_key") == "pending_value" assert "pending_key" in state._pending # pyright: ignore[reportPrivateUsage] + + def test_export_isolates_nested_mutable_values(self) -> None: + state = State() + state.set("history", ["step-1"]) + state.set("settings", {"enabled": True}) + state.commit() + + exported = state.export_state() + history = state.get("history") + history.append("step-2") + state.set("history", history) + settings = state.get("settings") + settings["enabled"] = False + state.set("settings", settings) + state.commit() + + assert exported == {"history": ["step-1"], "settings": {"enabled": True}} + + def test_exported_dict_does_not_mutate_state(self) -> None: + state = State() + state.set("key", "value") + state.commit() + + exported = state.export_state() + exported["added"] = True + del exported["key"] + + assert state.get("key") == "value" + assert state.has("added") is False + + def test_import_does_not_alias_caller_state(self) -> None: + state = State() + incoming = {"history": ["step-1"]} + + state.import_state(incoming) + incoming["history"].append("step-2") + + assert state.get("history") == ["step-1"] + + def test_export_import_roundtrip_isolates_snapshot(self) -> None: + source = State() + source.set("history", ["step-1"]) + source.commit() + snapshot = source.export_state() + + restored = State() + restored.import_state(snapshot) + history = restored.get("history") + history.append("step-2") + restored.set("history", history) + + assert snapshot == {"history": ["step-1"]}