|
| 1 | +import uuid |
1 | 2 | import pytest |
2 | 3 |
|
3 | 4 | from agno.models.openai.chat import OpenAIChat |
@@ -142,3 +143,49 @@ def test_team_session_state_on_run(route_team): |
142 | 143 | "test_key": "test_value", |
143 | 144 | "something_else": "other_value", |
144 | 145 | }, "Merging session state should work" |
| 146 | + |
| 147 | + |
| 148 | +def test_session_state_db_precedence(route_team, team_storage): |
| 149 | + """Test that DB session_state takes precedence over agent's in-memory session_state when switching sessions.""" |
| 150 | + # Set up two sessions with different session_state |
| 151 | + session_id_1 = "session_db" |
| 152 | + session_id_2 = "session_mem" |
| 153 | + |
| 154 | + # Simulate a session in storage with a specific session_state |
| 155 | + db_state = {"db_key": "db_value", "shared_key": "db"} |
| 156 | + team_storage.upsert( |
| 157 | + session=type("TeamSession", (), { |
| 158 | + "session_id": session_id_1, |
| 159 | + "session_data": { |
| 160 | + "session_state": db_state.copy(), |
| 161 | + "session_name": "db_session" |
| 162 | + }, |
| 163 | + "team_session_id": str(uuid.uuid4()), |
| 164 | + "team_id": str(uuid.uuid4()), |
| 165 | + "user_id": None, |
| 166 | + "team_data": None, |
| 167 | + "extra_data": None, |
| 168 | + })() |
| 169 | + ) |
| 170 | + |
| 171 | + # Set agent's in-memory session_state to something different |
| 172 | + route_team.session_state = {"mem_key": "mem_value", "shared_key": "mem"} |
| 173 | + route_team.session_id = session_id_2 |
| 174 | + route_team.session_name = "mem_session" |
| 175 | + |
| 176 | + # Run with the in-memory session (should use in-memory state) |
| 177 | + route_team.run("Test in-memory", session_id=session_id_2) |
| 178 | + assert route_team.session_state == {"current_session_id": session_id_2, "mem_key": "mem_value", "shared_key": "mem"} |
| 179 | + assert route_team.session_name == "mem_session" |
| 180 | + |
| 181 | + # Now switch to the DB session (should load and take precedence) |
| 182 | + route_team.run("Test DB", session_id=session_id_1) |
| 183 | + # The session_state should now match the DB's, not the in-memory one |
| 184 | + expected_state = {"current_session_id": session_id_1, **db_state} |
| 185 | + assert route_team.session_state == expected_state |
| 186 | + assert route_team.session_name == "db_session" |
| 187 | + |
| 188 | + # If we switch back to the in-memory session, it should restore the old state |
| 189 | + route_team.run("Test in-memory again", session_id=session_id_2) |
| 190 | + assert route_team.session_state == {"current_session_id": session_id_2, "mem_key": "mem_value", "shared_key": "mem"} |
| 191 | + assert route_team.session_name == "mem_session" |
0 commit comments