Files
wehub-resource-sync 25576b0be6
Checks (magui2.0) / python-lint (push) Failing after 1s
Checks (magui2.0) / python-test (push) Failing after 0s
Checks (magui2.0) / frontend-lint (push) Failing after 0s
CodeQL Advanced / Analyze (actions) (push) Failing after 1s
Checks (magui2.0) / python-format (push) Failing after 1s
CodeQL Advanced / Analyze (python) (push) Failing after 0s
Checks (magui2.0) / python-pyright (push) Failing after 1s
Checks (magui2.0) / frontend-format (push) Failing after 1s
Checks (magui2.0) / frontend-typecheck (push) Failing after 0s
Checks (magui2.0) / frontend-test (push) Failing after 2s
CodeQL Advanced / Analyze (javascript-typescript) (push) Failing after 1s
chore: import upstream snapshot with attribution
2026-07-13 12:24:32 +08:00

315 lines
11 KiB
Python

"""Conversation history is preserved across run_stream calls."""
from __future__ import annotations
import json
from pathlib import Path
from unittest.mock import MagicMock
import pytest
from magentic_ui.magentic_ui_config import ApprovalPolicy
from magentic_ui.sandbox._null import NullSandbox
from magentic_ui.teams.omniagent._omni_agent import _PATH_QUOTING_NUDGE, OmniAgent
def _make_completion(text: str, total_tokens: int = 50) -> MagicMock:
response = MagicMock()
response.choices = [MagicMock()]
response.choices[0].message.content = text
response.usage = MagicMock(
prompt_tokens=total_tokens - 5,
completion_tokens=5,
total_tokens=total_tokens,
)
return response
def _mock_llm_client(responses: list[str]) -> MagicMock:
from ._stream_mock import install_stream_mock
client = MagicMock()
install_stream_mock(client, [_make_completion(t) for t in responses])
return client
async def _build_agent(
tmp_path: Path,
client: MagicMock,
state_dir: Path | None = None,
) -> tuple[OmniAgent, NullSandbox]:
sandbox = NullSandbox(workspace=tmp_path)
await sandbox.__aenter__()
agent = OmniAgent(
client=client,
model="test-model",
host_workspace=tmp_path,
sandbox=sandbox,
approval_policy=ApprovalPolicy.AUTO_APPROVE,
state_dir=state_dir,
)
return agent, sandbox
@pytest.mark.asyncio
async def test_second_run_stream_sees_first_turn_history(tmp_path: Path) -> None:
"""After a final answer in turn 1, turn 2's LLM call must include
both the prior task and the prior answer."""
client = _mock_llm_client(
[
"<answer>first answer</answer>",
"<answer>second answer</answer>",
]
)
agent, sandbox = await _build_agent(tmp_path, client)
try:
async for _ in agent.run_stream("first task"):
pass
async for _ in agent.run_stream("second task"):
pass
finally:
await sandbox.__aexit__(None, None, None)
calls = client.chat.completions.stream.call_args_list
assert len(calls) == 2
# The first user message of a session is prepended with a one-shot
# bash path-quoting nudge; subsequent user messages are unchanged.
first_task_msg = f"{_PATH_QUOTING_NUDGE}\nfirst task"
# Turn 1: just system + first task
turn1_msgs = calls[0].kwargs["messages"]
turn1_roles = [m["role"] for m in turn1_msgs]
assert turn1_roles == ["system", "user"]
assert turn1_msgs[1]["content"] == first_task_msg
# Turn 2: system + first task + first answer + second task
turn2_msgs = calls[1].kwargs["messages"]
turn2_roles = [m["role"] for m in turn2_msgs]
assert turn2_roles == ["system", "user", "assistant", "user"]
assert turn2_msgs[1]["content"] == first_task_msg
assert turn2_msgs[2]["content"] == "<answer>first answer</answer>"
assert turn2_msgs[3]["content"] == "second task"
# Same system prompt across turns (built once)
assert turn1_msgs[0]["content"] == turn2_msgs[0]["content"]
@pytest.mark.asyncio
async def test_chat_attribute_reused_across_runs(tmp_path: Path) -> None:
"""The OmniResponses instance is the same object across run_stream calls."""
client = _mock_llm_client(
[
"<answer>a</answer>",
"<answer>b</answer>",
]
)
agent, sandbox = await _build_agent(tmp_path, client)
try:
assert agent._chat is None
async for _ in agent.run_stream("first"):
pass
chat_after_first = agent._chat
assert chat_after_first is not None
async for _ in agent.run_stream("second"):
pass
assert agent._chat is chat_after_first
finally:
await sandbox.__aexit__(None, None, None)
@pytest.mark.asyncio
async def test_state_persisted_to_disk_after_turn(tmp_path: Path) -> None:
"""A finished turn writes ``omni_state.json`` containing the conversation."""
state_dir = tmp_path / "session"
state_dir.mkdir()
workspace = tmp_path / "run"
workspace.mkdir()
client = _mock_llm_client(["<answer>done</answer>"])
agent, sandbox = await _build_agent(workspace, client, state_dir=state_dir)
try:
async for _ in agent.run_stream("first task"):
pass
finally:
await sandbox.__aexit__(None, None, None)
state_path = state_dir / "omni_state.json"
assert state_path.exists()
payload = json.loads(state_path.read_text())
assert isinstance(payload["messages"], list)
roles = [m["role"] for m in payload["messages"]]
assert roles == ["system", "user", "assistant"]
assert payload["messages"][1]["content"].endswith("first task")
assert payload["messages"][2]["content"] == "<answer>done</answer>"
assert payload["compaction_count"] == 0
@pytest.mark.asyncio
async def test_second_agent_resumes_from_disk(tmp_path: Path) -> None:
"""A new OmniAgent pointed at the same state_dir sees prior history."""
state_dir = tmp_path / "session"
state_dir.mkdir()
workspace = tmp_path / "run"
workspace.mkdir()
client_a = _mock_llm_client(["<answer>first answer</answer>"])
agent_a, sandbox_a = await _build_agent(workspace, client_a, state_dir=state_dir)
try:
async for _ in agent_a.run_stream("first task"):
pass
finally:
await sandbox_a.__aexit__(None, None, None)
client_b = _mock_llm_client(["<answer>second answer</answer>"])
agent_b, sandbox_b = await _build_agent(workspace, client_b, state_dir=state_dir)
try:
async for _ in agent_b.run_stream("second task"):
pass
finally:
await sandbox_b.__aexit__(None, None, None)
calls = client_b.chat.completions.stream.call_args_list
assert len(calls) == 1
# Turn 1 nudged the original task; turn 2 in a new agent must keep that
# nudged first message verbatim and append a plain follow-up.
nudged_first = f"{_PATH_QUOTING_NUDGE}\nfirst task"
resumed_msgs = calls[0].kwargs["messages"]
roles = [m["role"] for m in resumed_msgs]
assert roles == ["system", "user", "assistant", "user"]
assert resumed_msgs[1]["content"] == nudged_first
assert resumed_msgs[2]["content"] == "<answer>first answer</answer>"
assert resumed_msgs[3]["content"] == "second task"
@pytest.mark.asyncio
async def test_resume_preserves_delegate_cua_tool_response(tmp_path: Path) -> None:
"""A persisted ``<tool_response>`` from delegate_cua survives resume as plain text."""
state_dir = tmp_path / "session"
state_dir.mkdir()
workspace = tmp_path / "run"
workspace.mkdir()
nudged_first = f"{_PATH_QUOTING_NUDGE}\nbrowse to example.com"
seeded = {
"messages": [
{"role": "system", "content": "ignored — replaced by current prompt"},
{"role": "user", "content": nudged_first},
{
"role": "assistant",
"content": "<tool_call>delegate_cua...</tool_call>",
},
{
"role": "user",
"content": "<tool_response>Fara opened example.com and found the headline</tool_response>",
},
{"role": "assistant", "content": "<answer>headline found</answer>"},
],
"total_tokens": 0,
"prev_handoff": None,
"compaction_count": 0,
"viewport": {
"current_file": None,
"current_line": 0,
"scroll_count": 0,
},
}
(state_dir / "omni_state.json").write_text(json.dumps(seeded))
client = _mock_llm_client(["<answer>followed up</answer>"])
agent, sandbox = await _build_agent(workspace, client, state_dir=state_dir)
try:
async for _ in agent.run_stream("what else did you see?"):
pass
finally:
await sandbox.__aexit__(None, None, None)
messages = client.chat.completions.stream.call_args_list[0].kwargs["messages"]
contents = [m["content"] for m in messages]
# Fresh system prompt at index 0 — should not equal the seeded placeholder.
assert contents[0] != "ignored — replaced by current prompt"
# Delegate-cua tool_call and tool_response survive round-trip as plain text.
assert any("delegate_cua" in c for c in contents)
assert any("<tool_response>Fara opened example.com" in c for c in contents)
# Follow-up appears without the nudge prefix on the resumed turn.
assert contents[-1] == "what else did you see?"
@pytest.mark.asyncio
async def test_resume_restores_viewport(tmp_path: Path) -> None:
"""Persisted viewport values are restored on the next run_stream."""
state_dir = tmp_path / "session"
state_dir.mkdir()
workspace = tmp_path / "run"
workspace.mkdir()
seeded = {
"messages": [
{"role": "system", "content": "old"},
{"role": "user", "content": "task"},
{"role": "assistant", "content": "<answer>done</answer>"},
],
"total_tokens": 0,
"prev_handoff": None,
"compaction_count": 0,
"viewport": {
"current_file": "/workspace/notes.md",
"current_line": 42,
"scroll_count": 3,
},
}
(state_dir / "omni_state.json").write_text(json.dumps(seeded))
client = _mock_llm_client(["<answer>noop</answer>"])
agent, sandbox = await _build_agent(workspace, client, state_dir=state_dir)
try:
# Restoration happens during run_stream's lazy-init, so viewport is
# empty until the first turn opens. After that, prior values are
# available to tool calls.
async for _ in agent.run_stream("follow-up"):
pass
assert agent._viewport.current_file == "/workspace/notes.md"
assert agent._viewport.current_line == 42
assert agent._viewport.scroll_count == 3
finally:
await sandbox.__aexit__(None, None, None)
@pytest.mark.asyncio
async def test_resume_with_corrupted_state_starts_fresh(tmp_path: Path) -> None:
"""A garbage state file leaves the agent in a normal fresh-start path."""
state_dir = tmp_path / "session"
state_dir.mkdir()
(state_dir / "omni_state.json").write_text("not json {{{")
workspace = tmp_path / "run"
workspace.mkdir()
client = _mock_llm_client(["<answer>ok</answer>"])
agent, sandbox = await _build_agent(workspace, client, state_dir=state_dir)
try:
async for _ in agent.run_stream("hello"):
pass
finally:
await sandbox.__aexit__(None, None, None)
# Fresh start: the first user message should carry the nudge prefix,
# exactly like a never-resumed agent.
msgs = client.chat.completions.stream.call_args_list[0].kwargs["messages"]
assert msgs[1]["content"] == f"{_PATH_QUOTING_NUDGE}\nhello"
@pytest.mark.asyncio
async def test_no_state_dir_disables_persistence(tmp_path: Path) -> None:
"""With ``state_dir=None`` no file is written and behavior is unchanged."""
client = _mock_llm_client(["<answer>ok</answer>"])
agent, sandbox = await _build_agent(tmp_path, client, state_dir=None)
try:
async for _ in agent.run_stream("hello"):
pass
assert agent._state_path is None
finally:
await sandbox.__aexit__(None, None, None)
# No omni_state.json anywhere under tmp_path
assert list(tmp_path.rglob("omni_state.json")) == []