from __future__ import annotations import io import json from datetime import datetime from pathlib import Path from typing import Any, cast import pytest from openai.types.responses import ResponseCustomToolCall, ResponseFunctionToolCall from openai.types.responses.response_output_message import ResponseOutputMessage from openai.types.responses.response_reasoning_item import ResponseReasoningItem import agents.sandbox.capabilities.memory as memory_module import agents.sandbox.memory.manager as memory_manager_module import agents.sandbox.memory.phase_one as phase_one_module from agents import ( Agent, ReasoningItem, RunConfig, Runner, ShellTool, SQLiteSession, TResponseInputItem, ) from agents.exceptions import UserError from agents.items import ( CompactionItem, MessageOutputItem, ToolApprovalItem, TResponseOutputItem, ) from agents.result import RunResultStreaming from agents.run import _sandbox_memory_input from agents.run_context import RunContextWrapper from agents.sandbox import ( Manifest, MemoryGenerateConfig, MemoryLayoutConfig, MemoryReadConfig, SandboxAgent, SandboxRunConfig, ) from agents.sandbox.capabilities import Memory from agents.sandbox.memory.manager import ( _rollout_file_name_for_rollout_id, get_or_create_memory_generation_manager, ) from agents.sandbox.memory.phase_one import render_phase_one_prompt from agents.sandbox.memory.prompts import ( render_memory_consolidation_prompt, render_rollout_extraction_prompt, ) from agents.sandbox.memory.rollouts import ( RolloutTerminalMetadata, build_rollout_payload, build_rollout_payload_from_result, dump_rollout_json, ) from agents.sandbox.memory.storage import ( PhaseTwoInputSelection, PhaseTwoSelectionItem, SandboxMemoryStorage, _updated_at_sort_key, ) from agents.sandbox.runtime import _stream_memory_input_override from agents.sandbox.sandboxes.unix_local import UnixLocalSandboxClient from tests.fake_model import FakeModel from tests.test_responses import get_final_output_message, get_text_message from tests.utils.hitl import make_shell_call class _DeleteTrackingUnixLocalSandboxClient(UnixLocalSandboxClient): def __init__(self) -> None: super().__init__() self.deleted_roots: list[Path] = [] async def delete(self, session: Any) -> Any: self.deleted_roots.append(Path(session.state.manifest.root)) return await super().delete(session) def _phase_one_message( *, slug: str = "task_memory", summary: str = "# Task summary\n", raw_memory: str = "raw memory entry\n", ) -> Any: return get_final_output_message( json.dumps( { "rollout_slug": slug, "rollout_summary": summary, "raw_memory": raw_memory, } ) ) def test_rollout_file_name_for_rollout_id_uses_file_safe_id_directly() -> None: assert _rollout_file_name_for_rollout_id("chat-session.2026_04") == "chat-session.2026_04.jsonl" def test_rollout_file_name_for_rollout_id_rejects_path_like_ids() -> None: with pytest.raises(ValueError, match="file-safe ID"): _rollout_file_name_for_rollout_id("../chat-session") def test_rollout_file_name_for_rollout_id_rejects_empty_ids() -> None: with pytest.raises(ValueError, match="file-safe ID"): _rollout_file_name_for_rollout_id(" ") def _patch_update_call(call_id: str, path: str, text: str) -> Any: diff = "@@\n" + "".join(f"+{line}\n" for line in text.splitlines()) return ResponseCustomToolCall( type="custom_tool_call", name="apply_patch", call_id=call_id, input=json.dumps({"type": "update_file", "path": path, "diff": diff}), ) def _memory_config( *, max_raw_memories_for_consolidation: int = 256, extra_prompt: str | None = None, layout: MemoryLayoutConfig | None = None, read: MemoryReadConfig | None = None, phase_one_model: FakeModel | None = None, phase_two_model: FakeModel | None = None, ) -> Memory: return Memory( layout=layout or MemoryLayoutConfig(), read=read, generate=MemoryGenerateConfig( max_raw_memories_for_consolidation=max_raw_memories_for_consolidation, extra_prompt=extra_prompt, phase_one_model=phase_one_model or FakeModel(initial_output=[_phase_one_message()]), phase_two_model=phase_two_model or FakeModel( initial_output=[ _patch_update_call("memory-md", "memories/MEMORY.md", "memory entry"), _patch_update_call( "memory-summary", "memories/memory_summary.md", "summary entry" ), ] ), ), ) def _run_config_for_session(session: Any) -> RunConfig: return RunConfig(sandbox=SandboxRunConfig(session=session)) def _extract_user_text(fake_model: FakeModel) -> str: assert fake_model.first_turn_args is not None return _extract_user_text_from_turn_args(fake_model.first_turn_args) def _extract_user_text_from_turn_args(turn_args: dict[str, Any]) -> str: input_items = turn_args["input"] assert isinstance(input_items, list) first_item = cast(dict[str, Any], input_items[0]) content = first_item["content"] if isinstance(content, str): return content first_content = cast(dict[str, Any], content[0]) return cast(str, first_content["text"]) def _empty_phase_two_selection() -> PhaseTwoInputSelection: return PhaseTwoInputSelection(selected=[], retained_rollout_ids=set(), removed=[]) def _raw_memory_record( *, rollout_id: str, updated_at: str, rollout_summary_file: str, raw_memory: str, ) -> str: return ( f"rollout_id: {rollout_id}\n" f"updated_at: {updated_at}\n" f"rollout_path: sessions/{rollout_id}.jsonl\n" f"rollout_summary_file: {rollout_summary_file}\n" "terminal_state: completed\n\n" f"{raw_memory.rstrip()}\n" ) async def _cleanup_session( client: UnixLocalSandboxClient, session: Any, *, close: bool = True, ) -> None: try: if close: await session.aclose() finally: await client.delete(session) def test_build_rollout_payload_filters_developer_and_noisy_items() -> None: agent = Agent(name="test") assistant_message = cast(ResponseOutputMessage, get_text_message("assistant")) reasoning_item = ReasoningItem( agent=agent, raw_item=ResponseReasoningItem(id="rs_1", summary=[], type="reasoning"), ) compaction_item = CompactionItem( agent=agent, raw_item=cast( TResponseInputItem, { "type": "compaction", "summary": "compact", "encrypted_content": "encrypted", }, ), ) message_item = MessageOutputItem( agent=agent, raw_item=assistant_message, ) payload = build_rollout_payload( input=[ {"role": "developer", "content": "debug"}, {"role": "system", "content": "system"}, {"role": "user", "content": "hello"}, cast(TResponseInputItem, {"type": "reasoning", "summary": []}), cast( TResponseInputItem, { "type": "compaction", "summary": "compact", "encrypted_content": "encrypted", }, ), ], new_items=[reasoning_item, compaction_item, message_item], final_output="done", interruptions=[], terminal_metadata=RolloutTerminalMetadata( terminal_state="completed", has_final_output=True, ), ) updated_at = cast(str, payload.pop("updated_at")) assert datetime.fromisoformat(updated_at) assert list(payload) == ["input", "generated_items", "terminal_metadata", "final_output"] assert payload["input"] == [ {"role": "user", "content": "hello"}, ] assert payload["generated_items"] == [ assistant_message.model_dump(exclude_unset=True), ] assert payload["final_output"] == "done" def test_build_rollout_payload_serializes_model_interruptions_as_dicts() -> None: agent = Agent(name="test") raw = ResponseFunctionToolCall( id="fc_1", call_id="call_1", name="get_weather", arguments='{"city":"Paris"}', type="function_call", ) approval = ToolApprovalItem(agent=agent, raw_item=raw) payload = build_rollout_payload( input="hello", new_items=[], final_output=None, interruptions=[approval], terminal_metadata=RolloutTerminalMetadata(terminal_state="interrupted"), ) interruption = payload["interruptions"][0] assert isinstance(interruption, dict) assert interruption == raw.model_dump(exclude_unset=True) assert interruption["type"] == "function_call" assert interruption["call_id"] == "call_1" assert interruption["name"] == "get_weather" assert interruption["arguments"] == '{"city":"Paris"}' def test_render_phase_one_prompt_truncates_large_rollout_contents() -> None: payload = { "input": [{"role": "user", "content": f"start{'a' * 700_000}middle{'z' * 700_000}end"}], "generated_items": [], "terminal_metadata": {"terminal_state": "completed", "has_final_output": False}, } prompt = render_phase_one_prompt(rollout_contents=dump_rollout_json(payload)) assert "start" in prompt assert "end" in prompt assert "middle" not in prompt assert "tokens truncated" in prompt assert "rollout content omitted" in prompt assert "Do not assume the rendered rollout below is complete" in prompt def test_sandbox_memory_input_preserves_empty_session_delta() -> None: assert ( _sandbox_memory_input( memory_input_items_for_persistence=[], original_user_input=[{"content": "old turn", "role": "user"}], original_input=[{"content": "old turn", "role": "user"}], ) == [] ) def test_sandbox_memory_input_uses_saved_session_delta_after_persistence() -> None: assert _sandbox_memory_input( memory_input_items_for_persistence=[{"content": "current turn", "role": "user"}], original_user_input=[{"content": "old turn", "role": "user"}], original_input=[{"content": "old turn", "role": "user"}], ) == [{"content": "current turn", "role": "user"}] def test_streaming_memory_payload_preserves_empty_input_override() -> None: agent = Agent(name="test") result = RunResultStreaming( input=[{"content": "old turn", "role": "user"}], new_items=[], raw_responses=[], final_output="done", input_guardrail_results=[], output_guardrail_results=[], tool_input_guardrail_results=[], tool_output_guardrail_results=[], context_wrapper=RunContextWrapper(context=None), current_agent=agent, current_turn=0, max_turns=1, _current_agent_output_schema=None, trace=None, is_complete=True, ) assert result._original_input_for_persistence is None result._original_input_for_persistence = [] assert _stream_memory_input_override(result) == [] payload = build_rollout_payload_from_result( result, input_override=_stream_memory_input_override(result), ) assert payload["input"] == [] @pytest.mark.parametrize( ("conversation_id", "previous_response_id", "auto_previous_response_id"), [ ("conversation-123", None, False), (None, "resp_123", False), (None, None, True), ], ) def test_streaming_memory_payload_uses_result_input_for_server_managed_conversation( conversation_id: str | None, previous_response_id: str | None, auto_previous_response_id: bool, ) -> None: agent = Agent(name="test") result = RunResultStreaming( input=[{"content": "current turn", "role": "user"}], new_items=[], raw_responses=[], final_output="done", input_guardrail_results=[], output_guardrail_results=[], tool_input_guardrail_results=[], tool_output_guardrail_results=[], context_wrapper=RunContextWrapper(context=None), current_agent=agent, current_turn=0, max_turns=1, _current_agent_output_schema=None, trace=None, is_complete=True, ) result._conversation_id = conversation_id result._previous_response_id = previous_response_id result._auto_previous_response_id = auto_previous_response_id result._original_input_for_persistence = [] assert _stream_memory_input_override(result) is None payload = build_rollout_payload_from_result( result, input_override=_stream_memory_input_override(result), ) assert payload["input"] == [{"content": "current turn", "role": "user"}] def test_render_memory_prompts_omit_extra_prompt_section_by_default() -> None: rollout_prompt = render_rollout_extraction_prompt() consolidation_prompt = render_memory_consolidation_prompt( memory_root="memory", selection=_empty_phase_two_selection(), ) assert "{{ extra_prompt_section }}" not in rollout_prompt assert "{{ extra_prompt_section }}" not in consolidation_prompt assert "DEVELOPER-SPECIFIC EXTRA GUIDANCE" not in rollout_prompt assert "DEVELOPER-SPECIFIC EXTRA GUIDANCE" not in consolidation_prompt def test_render_memory_prompts_include_extra_prompt_section() -> None: rollout_prompt = render_rollout_extraction_prompt(extra_prompt="Focus on user preferences.") consolidation_prompt = render_memory_consolidation_prompt( memory_root="memory", selection=_empty_phase_two_selection(), extra_prompt="Focus on user preferences.", ) assert "DEVELOPER-SPECIFIC EXTRA GUIDANCE" in rollout_prompt assert "Focus on user preferences." in rollout_prompt assert "DEVELOPER-SPECIFIC EXTRA GUIDANCE" in consolidation_prompt assert "Focus on user preferences." in consolidation_prompt def test_render_memory_consolidation_prompt_lists_removed_rollouts() -> None: selection = PhaseTwoInputSelection( selected=[], retained_rollout_ids=set(), removed=[ PhaseTwoSelectionItem( rollout_id="old-rollout", updated_at="", rollout_path="sessions/old-rollout.jsonl", rollout_summary_file="memories/rollout_summaries/old.md", terminal_state="completed", ) ], ) prompt = render_memory_consolidation_prompt(memory_root="memory", selection=selection) assert "- removed from the last successful Phase 2 run: 1" in prompt assert "rollout_id=old-rollout" in prompt assert "updated_at=unknown" in prompt def test_updated_at_sort_key_places_unknown_timestamps_last() -> None: assert _updated_at_sort_key("updated_at: 2025-03-01T00:00:00Z\n") > _updated_at_sort_key( "updated_at: unknown\n" ) assert _updated_at_sort_key("updated_at: unknown\n") == _updated_at_sort_key("updated_at:\n") assert _updated_at_sort_key("updated_at: unknown\n") == _updated_at_sort_key("no metadata\n") @pytest.mark.asyncio async def test_phase_two_selection_tracks_added_retained_and_removed_rollouts() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) try: storage = SandboxMemoryStorage(session=session, layout=MemoryLayoutConfig()) await storage.ensure_layout() old_item = PhaseTwoSelectionItem( rollout_id="old-rollout", updated_at="2025-03-01T00:00:00Z", rollout_path="sessions/old-rollout.jsonl", rollout_summary_file="rollout_summaries/old-rollout.md", terminal_state="completed", ) await storage.write_text( storage.raw_memories_dir / "old-rollout.md", _raw_memory_record( rollout_id=old_item.rollout_id, updated_at=old_item.updated_at, rollout_summary_file=old_item.rollout_summary_file, raw_memory="old raw", ), ) await storage.write_text( storage.raw_memories_dir / "new-rollout.md", _raw_memory_record( rollout_id="new-rollout", updated_at="2025-03-02T00:00:00Z", rollout_summary_file="rollout_summaries/new-rollout.md", raw_memory="new raw", ), ) await storage.write_phase_two_selection(selected_items=[old_item]) selection = await storage.build_phase_two_input_selection( max_raw_memories_for_consolidation=1 ) assert [item.rollout_id for item in selection.selected] == ["new-rollout"] assert selection.retained_rollout_ids == set() assert [item.rollout_id for item in selection.removed] == ["old-rollout"] finally: await _cleanup_session(client, session) @pytest.mark.asyncio async def test_runner_memory_generation_sanitizes_and_truncates_phase_one_prompt( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(phase_one_module, "_PHASE_ONE_ROLLOUT_TOKEN_LIMIT", 1000) client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) phase_one_model = FakeModel(initial_output=[_phase_one_message()]) memory = _memory_config(phase_one_model=phase_one_model) agent = SandboxAgent( name="worker", model=FakeModel( initial_output=[ ResponseReasoningItem(id="rs_1", summary=[], type="reasoning"), cast( TResponseOutputItem, { "id": "compaction_1", "type": "compaction", "summary": "compacted-so-far", "encrypted_content": "encrypted", }, ), get_text_message("done"), ] ), instructions="Worker.", capabilities=[memory], ) closed = False try: result = await Runner.run( agent, [ {"role": "developer", "content": "developer debug"}, {"role": "system", "content": "system note"}, {"role": "user", "content": f"start{'a' * 20_000}middle{'z' * 20_000}end"}, cast(TResponseInputItem, {"type": "reasoning", "summary": []}), cast( TResponseInputItem, { "type": "compaction", "summary": "input-compact", "encrypted_content": "encrypted", }, ), ], run_config=_run_config_for_session(session), ) assert result.final_output == "done" assert phase_one_model.first_turn_args is None await session.aclose() closed = True prompt = _extract_user_text(phase_one_model) assert "developer debug" not in prompt assert "system note" not in prompt assert "reasoning" not in prompt assert "encrypted_content" not in prompt assert "input-compact" not in prompt assert "compacted-so-far" not in prompt assert "start" in prompt assert "middle" not in prompt assert "end" in prompt assert "tokens truncated" in prompt assert "rollout content omitted" in prompt finally: await _cleanup_session(client, session, close=not closed) @pytest.mark.asyncio async def test_sandbox_agent_without_memory_capability_skips_memory_generation() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) agent = SandboxAgent( name="worker", model=FakeModel(initial_output=[get_final_output_message("done")]), instructions="Worker.", ) try: result = await Runner.run( agent, "hello", run_config=_run_config_for_session(session), ) root = Path(session.state.manifest.root) assert result.final_output == "done" assert not (root / "sessions").exists() assert not (root / "memories").exists() finally: await _cleanup_session(client, session) @pytest.mark.asyncio async def test_memory_capability_returns_none_without_memory_summary() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) capability = Memory(generate=None) try: async with session: capability.bind(session) assert await capability.instructions(session.state.manifest) is None await session.mkdir("memories", parents=True) await session.write( Path("memories/memory_summary.md"), io.BytesIO(b""), ) assert await capability.instructions(session.state.manifest) is None finally: await client.delete(session) @pytest.mark.parametrize( ("memories_dir", "match"), [ ("/memory", "memories_dir must be relative"), ("../memory", "memories_dir must not escape root"), ("", "memories_dir must be non-empty"), (".", "memories_dir must be non-empty"), ], ) def test_memory_capability_rejects_invalid_memories_dir( memories_dir: str, match: str, ) -> None: with pytest.raises(ValueError, match=match): Memory(layout=MemoryLayoutConfig(memories_dir=memories_dir), generate=None) @pytest.mark.parametrize( ("sessions_dir", "match"), [ ("/sessions", "sessions_dir must be relative"), ("../sessions", "sessions_dir must not escape root"), ("", "sessions_dir must be non-empty"), (".", "sessions_dir must be non-empty"), ], ) def test_memory_capability_rejects_invalid_sessions_dir( sessions_dir: str, match: str, ) -> None: with pytest.raises(ValueError, match=match): Memory(layout=MemoryLayoutConfig(sessions_dir=sessions_dir), generate=None) def test_memory_capability_requires_read_or_generate() -> None: with pytest.raises(ValueError, match="Memory requires at least one of `read` or `generate`"): Memory(read=None, generate=None) @pytest.mark.asyncio async def test_memory_capability_instructions_requires_bound_session() -> None: capability = Memory(generate=None) with pytest.raises(ValueError, match="Memory capability is not bound to a SandboxSession"): await capability.instructions(Manifest()) def test_memory_generate_config_rejects_non_positive_recent_rollout_limit() -> None: with pytest.raises( ValueError, match=("MemoryGenerateConfig.max_raw_memories_for_consolidation must be greater than 0"), ): MemoryGenerateConfig(max_raw_memories_for_consolidation=0) def test_memory_layout_config_defaults_match_codex_names() -> None: config = MemoryLayoutConfig() assert config.memories_dir == "memories" assert config.sessions_dir == "sessions" def test_memory_generate_config_accepts_renamed_limit_field() -> None: config = MemoryGenerateConfig(max_raw_memories_for_consolidation=123) assert config.max_raw_memories_for_consolidation == 123 def test_memory_generate_config_rejects_too_many_raw_memories() -> None: with pytest.raises( ValueError, match=( "MemoryGenerateConfig.max_raw_memories_for_consolidation " "must be less than or equal to 4096" ), ): MemoryGenerateConfig(max_raw_memories_for_consolidation=4097) @pytest.mark.asyncio async def test_memory_capability_injects_truncated_memory_summary( monkeypatch: pytest.MonkeyPatch, ) -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) capability = Memory(generate=None) try: async with session: monkeypatch.setattr(memory_module, "_MEMORY_SUMMARY_MAX_TOKENS", 1) await session.mkdir("memories", parents=True) await session.write( Path("memories/memory_summary.md"), io.BytesIO(b"abcdefg"), ) capability.bind(session) instructions = await capability.instructions(session.state.manifest) assert instructions is not None assert ( "memories/memory_summary.md (already provided below; do NOT open again)" in instructions ) assert "MEMORY_SUMMARY BEGINS" in instructions assert "tokens truncated" in instructions finally: await client.delete(session) @pytest.mark.asyncio async def test_memory_capability_live_update_instructions() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) capability = Memory(generate=None) try: async with session: await session.mkdir("memories", parents=True) await session.write( Path("memories/memory_summary.md"), io.BytesIO(b"summary entry"), ) capability.bind(session) instructions = await capability.instructions(session.state.manifest) assert instructions is not None assert "Memory is writable." in instructions assert "memories/MEMORY.md" in instructions assert "same turn" in instructions assert "Never update memories." not in instructions finally: await client.delete(session) @pytest.mark.asyncio async def test_sandbox_memory_writes_rollouts_and_memory_files() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) phase_one_model = FakeModel(initial_output=[_phase_one_message()]) phase_two_model = FakeModel( initial_output=[ _patch_update_call("memory-md", "memories/MEMORY.md", "memory entry"), _patch_update_call("memory-summary", "memories/memory_summary.md", "summary entry"), ] ) phase_two_model.set_next_output([get_final_output_message("consolidated")]) memory = _memory_config( extra_prompt="Track durable user preferences.", phase_one_model=phase_one_model, phase_two_model=phase_two_model, ) agent = SandboxAgent( name="worker", model=FakeModel(initial_output=[get_final_output_message("done")]), instructions="Worker.", capabilities=[memory], ) closed = False try: result = await Runner.run( agent, "hello", run_config=_run_config_for_session(session), ) root = Path(session.state.manifest.root) rollouts = sorted((root / "sessions").glob("*.jsonl")) assert result.final_output == "done" assert len(rollouts) == 1 assert phase_one_model.first_turn_args is None await session.aclose() closed = True raw_memories = sorted((root / "memories" / "raw_memories").glob("*.md")) rollout_summaries = sorted((root / "memories" / "rollout_summaries").glob("*.md")) assert len(raw_memories) == 1 assert len(rollout_summaries) == 1 assert (root / "memories" / "MEMORY.md").read_text() == "memory entry\n" assert (root / "memories" / "memory_summary.md").read_text() == "summary entry\n" assert "rollout_id: " in (root / "memories" / "raw_memories.md").read_text() assert "updated_at: " in (root / "memories" / "raw_memories.md").read_text() assert "rollout_path: sessions/" in (root / "memories" / "raw_memories.md").read_text() assert ( "rollout_summary_file: rollout_summaries/" in (root / "memories" / "raw_memories.md").read_text() ) assert "terminal_state: completed" in (root / "memories" / "raw_memories.md").read_text() assert "session_id: " in rollout_summaries[0].read_text() assert "updated_at: " in rollout_summaries[0].read_text() assert "rollout_path: sessions/" in rollout_summaries[0].read_text() assert "terminal_state: completed" in rollout_summaries[0].read_text() assert '"terminal_state":"completed"' in _extract_user_text(phase_one_model) assert phase_one_model.first_turn_args is not None assert ( "DEVELOPER-SPECIFIC EXTRA GUIDANCE" in phase_one_model.first_turn_args["system_instructions"] ) assert ( "Track durable user preferences." in phase_one_model.first_turn_args["system_instructions"] ) assert phase_two_model.first_turn_args is not None assert "DEVELOPER-SPECIFIC EXTRA GUIDANCE" in _extract_user_text(phase_two_model) assert "Track durable user preferences." in _extract_user_text(phase_two_model) finally: await _cleanup_session(client, session, close=not closed) @pytest.mark.asyncio async def test_sandbox_memory_uses_custom_layout() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) phase_two_model = FakeModel( initial_output=[ _patch_update_call("memory-md", "agent_memory/MEMORY.md", "memory entry"), _patch_update_call("memory-summary", "agent_memory/memory_summary.md", "summary entry"), ] ) phase_two_model.set_next_output([get_final_output_message("consolidated")]) memory = Memory( layout=MemoryLayoutConfig(memories_dir="agent_memory", sessions_dir="agent_sessions"), read=None, generate=MemoryGenerateConfig( phase_one_model=FakeModel(initial_output=[_phase_one_message()]), phase_two_model=phase_two_model, ), ) agent = SandboxAgent( name="worker", model=FakeModel(initial_output=[get_final_output_message("done")]), instructions="Worker.", capabilities=[memory], ) closed = False try: await Runner.run( agent, "hello", run_config=_run_config_for_session(session), ) root = Path(session.state.manifest.root) assert len(list((root / "agent_sessions").glob("*.jsonl"))) == 1 await session.aclose() closed = True assert (root / "agent_memory" / "MEMORY.md").read_text() == "memory entry\n" assert (root / "agent_memory" / "memory_summary.md").read_text() == "summary entry\n" finally: await _cleanup_session(client, session, close=not closed) @pytest.mark.asyncio async def test_sandbox_memory_supports_multiple_generating_layouts_in_one_session() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) phase_two_model_a = FakeModel( initial_output=[ _patch_update_call("a-memory", "agent_a_memory/MEMORY.md", "agent a entry"), _patch_update_call( "a-summary", "agent_a_memory/memory_summary.md", "agent a summary", ), ] ) phase_two_model_a.set_next_output([get_final_output_message("agent a consolidated")]) phase_two_model_b = FakeModel( initial_output=[ _patch_update_call("b-memory", "agent_b_memory/MEMORY.md", "agent b entry"), _patch_update_call( "b-summary", "agent_b_memory/memory_summary.md", "agent b summary", ), ] ) phase_two_model_b.set_next_output([get_final_output_message("agent b consolidated")]) memory_a = _memory_config( layout=MemoryLayoutConfig(memories_dir="agent_a_memory", sessions_dir="agent_a_sessions"), phase_one_model=FakeModel(initial_output=[_phase_one_message(raw_memory="agent a raw\n")]), phase_two_model=phase_two_model_a, ) memory_b = _memory_config( layout=MemoryLayoutConfig(memories_dir="agent_b_memory", sessions_dir="agent_b_sessions"), phase_one_model=FakeModel(initial_output=[_phase_one_message(raw_memory="agent b raw\n")]), phase_two_model=phase_two_model_b, ) agent_a = SandboxAgent( name="agent-a", model=FakeModel(initial_output=[get_final_output_message("a done")]), instructions="Agent A.", capabilities=[memory_a], ) agent_b = SandboxAgent( name="agent-b", model=FakeModel(initial_output=[get_final_output_message("b done")]), instructions="Agent B.", capabilities=[memory_b], ) closed = False try: await Runner.run(agent_a, "first", run_config=_run_config_for_session(session)) await Runner.run(agent_b, "second", run_config=_run_config_for_session(session)) root = Path(session.state.manifest.root) assert len(list((root / "agent_a_sessions").glob("*.jsonl"))) == 1 assert len(list((root / "agent_b_sessions").glob("*.jsonl"))) == 1 await session.aclose() closed = True assert (root / "agent_a_memory" / "MEMORY.md").read_text() == "agent a entry\n" assert (root / "agent_b_memory" / "MEMORY.md").read_text() == "agent b entry\n" finally: await _cleanup_session(client, session, close=not closed) @pytest.mark.asyncio async def test_sandbox_memory_rejects_different_generate_configs_for_same_layout() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) memory = _memory_config() different_memory = _memory_config( phase_one_model=FakeModel(initial_output=[_phase_one_message(raw_memory="different\n")]) ) try: get_or_create_memory_generation_manager(session=session, memory=memory) with pytest.raises(UserError, match="different Memory generation config"): get_or_create_memory_generation_manager(session=session, memory=different_memory) finally: await _cleanup_session(client, session) @pytest.mark.asyncio async def test_sandbox_memory_rollout_payload_uses_validated_rollout_id() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) memory = _memory_config() try: manager = get_or_create_memory_generation_manager(session=session, memory=memory) await manager.enqueue_rollout_payload( { "updated_at": "2026-04-15T00:00:00+00:00", "rollout_id": "payload-id", "input": [], "generated_items": [], "terminal_metadata": {"terminal_state": "completed", "has_final_output": False}, }, rollout_id="canonical-id", ) root = Path(session.state.manifest.root) rollout_path = root / "sessions" / "canonical-id.jsonl" payload = json.loads(rollout_path.read_text()) assert payload["rollout_id"] == "canonical-id" finally: await client.delete(session) @pytest.mark.asyncio async def test_sandbox_memory_rejects_different_sessions_dirs_for_same_memories_dir() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) first_memory = _memory_config( layout=MemoryLayoutConfig(memories_dir="shared_memory", sessions_dir="sessions_a") ) second_memory = _memory_config( layout=MemoryLayoutConfig(memories_dir="shared_memory", sessions_dir="sessions_b") ) try: get_or_create_memory_generation_manager(session=session, memory=first_memory) with pytest.raises(UserError, match="already has a Memory generation capability"): get_or_create_memory_generation_manager(session=session, memory=second_memory) finally: await _cleanup_session(client, session) @pytest.mark.asyncio async def test_sandbox_memory_rejects_shared_sessions_dir_for_different_memories_dirs() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) first_memory = _memory_config( layout=MemoryLayoutConfig(memories_dir="memory_a", sessions_dir="shared_sessions") ) second_memory = _memory_config( layout=MemoryLayoutConfig(memories_dir="memory_b", sessions_dir="shared_sessions") ) try: get_or_create_memory_generation_manager(session=session, memory=first_memory) with pytest.raises(UserError, match="sessions_dir='shared_sessions'"): get_or_create_memory_generation_manager(session=session, memory=second_memory) finally: await _cleanup_session(client, session) @pytest.mark.asyncio async def test_sandbox_memory_groups_segments_by_sdk_session_until_close() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) phase_one_model = FakeModel(initial_output=[_phase_one_message(raw_memory="joined raw\n")]) phase_two_model = FakeModel( initial_output=[ _patch_update_call("memory-md", "memories/MEMORY.md", "joined entry"), _patch_update_call("memory-summary", "memories/memory_summary.md", "joined summary"), ] ) phase_two_model.set_next_output([get_final_output_message("joined")]) memory = _memory_config( phase_one_model=phase_one_model, phase_two_model=phase_two_model, ) first_agent = SandboxAgent( name="first-worker", model=FakeModel(initial_output=[get_final_output_message("first done")]), instructions="Worker.", capabilities=[memory], ) second_agent = SandboxAgent( name="second-worker", model=FakeModel(initial_output=[get_final_output_message("second done")]), instructions="Worker.", capabilities=[memory], ) closed = False try: chat_session = SQLiteSession("chat-session") run_config = _run_config_for_session(session) first = await Runner.run( first_agent, "first", session=chat_session, run_config=run_config, ) second = await Runner.run( second_agent, "second", session=chat_session, run_config=run_config, ) root = Path(session.state.manifest.root) rollouts = sorted((root / "sessions").glob("*.jsonl")) assert first.final_output == "first done" assert second.final_output == "second done" assert len(rollouts) == 1 assert rollouts[0].name == "chat-session.jsonl" assert len(rollouts[0].read_text().splitlines()) == 2 segments = [json.loads(line) for line in rollouts[0].read_text().splitlines()] assert list(segments[0])[:4] == [ "updated_at", "rollout_id", "input", "generated_items", ] assert segments[0]["input"] == [{"content": "first", "role": "user"}] assert segments[1]["input"] == [{"content": "second", "role": "user"}] assert phase_one_model.first_turn_args is None await session.aclose() closed = True prompt = _extract_user_text(phase_one_model) assert "first" in prompt assert "second" in prompt assert '"segment_count":2' in prompt raw_memory_files = list((root / "memories" / "raw_memories").glob("*.md")) assert len(raw_memory_files) == 1 assert f"updated_at: {segments[-1]['updated_at']}\n" in raw_memory_files[0].read_text() assert (root / "memories" / "MEMORY.md").read_text() == "joined entry\n" finally: await _cleanup_session(client, session, close=not closed) @pytest.mark.asyncio async def test_sandbox_memory_fallback_does_not_mutate_run_config() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) agent_model = FakeModel() agent_model.add_multiple_turn_outputs( [ [get_final_output_message("first done")], [get_final_output_message("second done")], ] ) agent = SandboxAgent( name="worker", model=agent_model, instructions="Worker.", capabilities=[_memory_config()], ) try: run_config = _run_config_for_session(session) await Runner.run( agent, "first", session=SQLiteSession("first-chat"), run_config=run_config, ) await Runner.run( agent, "second", session=SQLiteSession("second-chat"), run_config=run_config, ) root = Path(session.state.manifest.root) rollouts = sorted(path.name for path in (root / "sessions").glob("*.jsonl")) assert rollouts == ["first-chat.jsonl", "second-chat.jsonl"] finally: await _cleanup_session(client, session) @pytest.mark.asyncio async def test_sandbox_memory_uses_conversation_id_when_sdk_session_is_absent() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) agent = SandboxAgent( name="worker", model=FakeModel(initial_output=[get_final_output_message("done")]), instructions="Worker.", capabilities=[_memory_config()], ) try: result = await Runner.run( agent, "remember this conversation", conversation_id="conversation-123", run_config=_run_config_for_session(session), ) root = Path(session.state.manifest.root) rollouts = sorted((root / "sessions").glob("*.jsonl")) assert result.final_output == "done" assert len(rollouts) == 1 assert rollouts[0].name == "conversation-123.jsonl" finally: await _cleanup_session(client, session) @pytest.mark.asyncio async def test_sandbox_memory_uses_group_id_when_sdk_session_is_absent() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) agent_model = FakeModel() agent_model.add_multiple_turn_outputs( [ [get_final_output_message("first done")], [get_final_output_message("second done")], ] ) agent = SandboxAgent( name="worker", model=agent_model, instructions="Worker.", capabilities=[_memory_config()], ) try: run_config = RunConfig( sandbox=SandboxRunConfig(session=session), group_id="trace-thread-123", ) first = await Runner.run(agent, "first", run_config=run_config) second = await Runner.run(agent, "second", run_config=run_config) root = Path(session.state.manifest.root) rollouts = sorted((root / "sessions").glob("*.jsonl")) assert first.final_output == "first done" assert second.final_output == "second done" assert len(rollouts) == 1 assert rollouts[0].name == "trace-thread-123.jsonl" assert len(rollouts[0].read_text().splitlines()) == 2 finally: await _cleanup_session(client, session) @pytest.mark.asyncio async def test_sandbox_memory_uses_per_run_conversation_when_no_conversation_id() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) agent_model = FakeModel() agent_model.add_multiple_turn_outputs( [ [get_final_output_message("first done")], [get_final_output_message("second done")], ] ) agent = SandboxAgent( name="worker", model=agent_model, instructions="Worker.", capabilities=[_memory_config()], ) try: run_config = _run_config_for_session(session) first = await Runner.run(agent, "first", run_config=run_config) second = await Runner.run(agent, "second", run_config=run_config) root = Path(session.state.manifest.root) rollouts = sorted(path.name for path in (root / "sessions").glob("*.jsonl")) assert first.final_output == "first done" assert second.final_output == "second done" assert len(rollouts) == 2 assert all(name.startswith("run-") and name.endswith(".jsonl") for name in rollouts) finally: await _cleanup_session(client, session) @pytest.mark.asyncio async def test_sandbox_memory_caps_phase_two_selection_and_surfaces_removed_rollouts() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) phase_one_model = FakeModel() phase_one_model.add_multiple_turn_outputs( [ [_phase_one_message(slug="first", raw_memory="first raw\n")], [_phase_one_message(slug="second", raw_memory="second raw\n")], ] ) phase_two_model = FakeModel( initial_output=[ _patch_update_call("memory-md", "memories/MEMORY.md", "first entry"), _patch_update_call("memory-summary", "memories/memory_summary.md", "first summary"), ] ) phase_two_model.set_next_output([get_final_output_message("consolidated")]) memory = _memory_config( max_raw_memories_for_consolidation=1, phase_one_model=phase_one_model, phase_two_model=phase_two_model, ) agent_model = FakeModel() agent_model.add_multiple_turn_outputs( [ [get_final_output_message("first done")], [get_final_output_message("second done")], ] ) agent = SandboxAgent( name="worker", model=agent_model, instructions="Worker.", capabilities=[memory], ) closed = False try: root = Path(session.state.manifest.root) await Runner.run( agent, "first", run_config=RunConfig( sandbox=SandboxRunConfig(session=session), group_id="first-chat", ), ) await Runner.run( agent, "second", run_config=RunConfig( sandbox=SandboxRunConfig(session=session), group_id="second-chat", ), ) assert len(list((root / "sessions").glob("*.jsonl"))) == 2 await session.aclose() closed = True selection_payload = json.loads((root / "memories" / "phase_two_selection.json").read_text()) selected_rollout_ids = [ cast(str, item["rollout_id"]) for item in selection_payload["selected"] ] assert len(selected_rollout_ids) == 1 merged_raw_memories = (root / "memories" / "raw_memories.md").read_text() assert "second raw" in merged_raw_memories assert "first raw" not in merged_raw_memories assert phase_two_model.first_turn_args is not None prompt = _extract_user_text_from_turn_args(phase_two_model.first_turn_args) assert "newly added since the last successful Phase 2 run: 1" in prompt assert f"rollout_id={selected_rollout_ids[0]}" in prompt finally: await _cleanup_session(client, session, close=not closed) @pytest.mark.asyncio async def test_sandbox_memory_runs_phase_one_and_phase_two_on_session_close() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) phase_one_model = FakeModel(initial_output=[_phase_one_message()]) phase_two_model = FakeModel( initial_output=[ _patch_update_call("memory-md", "memories/MEMORY.md", "shutdown entry"), _patch_update_call("memory-summary", "memories/memory_summary.md", "shutdown summary"), ] ) phase_two_model.set_next_output([get_final_output_message("shutdown")]) memory = _memory_config( phase_one_model=phase_one_model, phase_two_model=phase_two_model, ) agent = SandboxAgent( name="worker", model=FakeModel(initial_output=[get_final_output_message("done")]), instructions="Worker.", capabilities=[memory], ) root = Path(session.state.manifest.root) try: await Runner.run(agent, "hello", run_config=_run_config_for_session(session)) manager = get_or_create_memory_generation_manager(session=session, memory=memory) await manager._queue.join() assert (root / "memories" / "MEMORY.md").read_text() == "" await session.aclose() assert (root / "memories" / "MEMORY.md").read_text() == "shutdown entry\n" assert (root / "memories" / "memory_summary.md").read_text() == "shutdown summary\n" finally: await client.delete(session) @pytest.mark.asyncio async def test_sandbox_memory_unregisters_manager_on_session_close() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) memory = _memory_config() try: manager = get_or_create_memory_generation_manager(session=session, memory=memory) managers_by_layout = memory_manager_module._MEMORY_GENERATION_MANAGERS.get(session) assert managers_by_layout is not None assert manager in managers_by_layout.values() await session.aclose() assert memory_manager_module._MEMORY_GENERATION_MANAGERS.get(session) is None finally: await client.delete(session) @pytest.mark.asyncio async def test_sandbox_memory_enqueue_failure_still_cleans_up_owned_session( monkeypatch: pytest.MonkeyPatch, ) -> None: async def _raise_write_rollout(*args: Any, **kwargs: Any) -> Path: _ = args, kwargs raise RuntimeError("write_rollout failed") monkeypatch.setattr(memory_manager_module, "write_rollout", _raise_write_rollout) client = _DeleteTrackingUnixLocalSandboxClient() agent = SandboxAgent( name="worker", model=FakeModel(initial_output=[get_final_output_message("done")]), instructions="Worker.", capabilities=[_memory_config()], ) result = await Runner.run( agent, "hello", run_config=RunConfig(sandbox=SandboxRunConfig(client=client)), ) assert result.final_output == "done" assert len(client.deleted_roots) == 1 assert not client.deleted_roots[0].exists() @pytest.mark.asyncio async def test_sandbox_memory_marks_interrupted_runs_in_phase_one_prompt() -> None: client = UnixLocalSandboxClient() session = await client.create(manifest=Manifest()) phase_one_model = FakeModel(initial_output=[_phase_one_message()]) phase_two_model = FakeModel( initial_output=[ _patch_update_call("memory-md", "memories/MEMORY.md", "interrupted entry"), _patch_update_call( "memory-summary", "memories/memory_summary.md", "interrupted summary" ), ] ) phase_two_model.set_next_output([get_final_output_message("done")]) memory = _memory_config( phase_one_model=phase_one_model, phase_two_model=phase_two_model, ) agent = SandboxAgent( name="worker", model=FakeModel(initial_output=[make_shell_call("approval-call")]), instructions="Worker.", tools=[ShellTool(executor=lambda _request: "ok", needs_approval=True)], capabilities=[memory], ) closed = False try: result = await Runner.run( agent, "interrupt me", run_config=_run_config_for_session(session), ) assert result.interruptions await session.aclose() closed = True assert '"terminal_state":"interrupted"' in _extract_user_text(phase_one_model) finally: await _cleanup_session(client, session, close=not closed)