| 1 | import copy |
| 2 | import json |
| 3 | from pathlib import Path |
| 4 | import sys |
| 5 | from types import SimpleNamespace |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | |
| 10 | ROOT = Path(__file__).resolve().parents[3] |
| 11 | if str(ROOT) not in sys.path: |
| 12 | sys.path.insert(0, str(ROOT)) |
| 13 | |
| 14 | from helpers import state_monitor_integration |
| 15 | from helpers.persist_chat import _collect_response_ids |
| 16 | from plugins._chat_branching.api import branch_chat |
| 17 | |
| 18 | |
| 19 | @pytest.mark.asyncio |
| 20 | async def test_branch_rebuilds_provider_and_context_state_from_trimmed_history( |
| 21 | monkeypatch, |
| 22 | ): |
| 23 | history = json.dumps( |
| 24 | { |
| 25 | "_cls": "History", |
| 26 | "counter": 2, |
| 27 | "bulks": [], |
| 28 | "topics": [], |
| 29 | "current": { |
| 30 | "summary": "", |
| 31 | "messages": [ |
| 32 | { |
| 33 | "id": "kept-message", |
| 34 | "content": "before", |
| 35 | "metadata": { |
| 36 | "responses": { |
| 37 | "response_id": "resp_kept", |
| 38 | "previous_response_id": "resp_previous", |
| 39 | "output_items": [{"type": "message"}], |
| 40 | } |
| 41 | }, |
| 42 | }, |
| 43 | {"id": "removed-message", "content": "after"}, |
| 44 | ], |
| 45 | }, |
| 46 | } |
| 47 | ) |
| 48 | serialized = { |
| 49 | "id": "source-chat", |
| 50 | "name": "Source chat", |
| 51 | "log": { |
| 52 | "logs": [ |
| 53 | {"no": 4, "id": "kept-message"}, |
| 54 | {"no": 5, "id": "removed-message"}, |
| 55 | ] |
| 56 | }, |
| 57 | "agents": [ |
| 58 | { |
| 59 | "history": history, |
| 60 | "data": { |
| 61 | "responses_state": { |
| 62 | "response_id": "resp_current", |
| 63 | "response_ids": ["resp_kept", "resp_current"], |
| 64 | }, |
| 65 | "ctx_window": {"text": "source context"}, |
| 66 | }, |
| 67 | } |
| 68 | for _ in range(2) |
| 69 | ], |
| 70 | } |
| 71 | |
| 72 | branched = [] |
| 73 | |
| 74 | monkeypatch.setattr( |
| 75 | branch_chat.AgentContext, |
| 76 | "get", |
| 77 | lambda context_id: object() if context_id == "source-chat" else None, |
| 78 | ) |
| 79 | monkeypatch.setattr( |
| 80 | branch_chat, |
| 81 | "_serialize_context", |
| 82 | lambda _context: copy.deepcopy(serialized), |
| 83 | ) |
| 84 | |
| 85 | def deserialize(data): |
| 86 | branched.append(copy.deepcopy(data)) |
| 87 | return SimpleNamespace(id="branch-chat") |
| 88 | |
| 89 | monkeypatch.setattr(branch_chat, "_deserialize_context", deserialize) |
| 90 | monkeypatch.setattr(branch_chat, "save_tmp_chat", lambda _context: None) |
| 91 | monkeypatch.setattr( |
| 92 | state_monitor_integration, |
| 93 | "mark_dirty_all", |
| 94 | lambda **_kwargs: None, |
| 95 | ) |
| 96 | |
| 97 | result = await branch_chat.BranchChat.process( |
| 98 | None, |
| 99 | {"context": "source-chat", "log_no": 4}, |
| 100 | None, |
| 101 | ) |
| 102 | |
| 103 | assert result["ctxid"] == "branch-chat" |
| 104 | assert len(branched) == 1 |
| 105 | assert _collect_response_ids(branched[0]) == [] |
| 106 | for agent_data in branched[0]["agents"]: |
| 107 | assert "ctx_window" not in agent_data["data"] |
| 108 | assert "responses_state" not in agent_data["data"] |
| 109 | trimmed = json.loads(agent_data["history"]) |
| 110 | messages = trimmed["current"]["messages"] |
| 111 | assert [message["id"] for message in messages] == ["kept-message"] |
| 112 | responses = messages[0]["metadata"]["responses"] |
| 113 | assert "response_id" not in responses |
| 114 | assert "previous_response_id" not in responses |
| 115 | assert responses["output_items"] == [{"type": "message"}] |
| 116 | |
| 117 | assert serialized["agents"][0]["data"]["responses_state"]["response_id"] == ( |
| 118 | "resp_current" |
| 119 | ) |
| 120 | original_message = json.loads(serialized["agents"][0]["history"])["current"][ |
| 121 | "messages" |
| 122 | ][0] |
| 123 | assert original_message["metadata"]["responses"]["response_id"] == "resp_kept" |