Python: Refactor ag-ui to clean up some patterns (#2363)

* Refactor ag-ui to clean up some patterns

* Mypy fixes

* Fix imports, typing, tests, logging.

* Fix test import error

* Fix imports again

* Fix thread handling

python/packages/ag-ui/tests/test_state_manager.py

@@ -0,0 +1,51 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from ag_ui.core import CustomEvent, EventType
+from agent_framework import ChatMessage, TextContent
+
+from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+from agent_framework_ag_ui._orchestration._state_manager import StateManager
+
+
+def test_state_manager_initializes_defaults_and_snapshot() -> None:
+    state_manager = StateManager(
+        state_schema={"items": {"type": "array"}, "metadata": {"type": "object"}},
+        predict_state_config=None,
+        require_confirmation=True,
+    )
+    current_state = state_manager.initialize({"metadata": {"a": 1}})
+    bridge = AgentFrameworkEventBridge(run_id="run", thread_id="thread", current_state=current_state)
+
+    snapshot_event = state_manager.initial_snapshot_event(bridge)
+    assert snapshot_event is not None
+    assert snapshot_event.snapshot["items"] == []
+    assert snapshot_event.snapshot["metadata"] == {"a": 1}
+
+
+def test_state_manager_predict_state_event_shape() -> None:
+    state_manager = StateManager(
+        state_schema=None,
+        predict_state_config={"doc": {"tool": "write_document_local", "tool_argument": "document"}},
+        require_confirmation=True,
+    )
+    predict_event = state_manager.predict_state_event()
+    assert isinstance(predict_event, CustomEvent)
+    assert predict_event.type == EventType.CUSTOM
+    assert predict_event.name == "PredictState"
+    assert predict_event.value[0]["state_key"] == "doc"
+
+
+def test_state_context_only_when_new_user_turn() -> None:
+    state_manager = StateManager(
+        state_schema={"items": {"type": "array"}},
+        predict_state_config=None,
+        require_confirmation=True,
+    )
+    state_manager.initialize({"items": [1]})
+
+    assert state_manager.state_context_message(is_new_user_turn=False, conversation_has_tool_calls=False) is None
+
+    message = state_manager.state_context_message(is_new_user_turn=True, conversation_has_tool_calls=False)
+    assert isinstance(message, ChatMessage)
+    assert isinstance(message.contents[0], TextContent)
+    assert "Current state of the application" in message.contents[0].text