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_endpoint.py

@@ -3,7 +3,8 @@
 """Tests for FastAPI endpoint creation (_endpoint.py)."""
 
 import json
-from typing import Any
+import sys
+from pathlib import Path
 
 from agent_framework import ChatAgent, TextContent
 from agent_framework._types import ChatResponseUpdate
@@ -13,22 +14,20 @@ from fastapi.testclient import TestClient
 from agent_framework_ag_ui._agent import AgentFrameworkAgent
 from agent_framework_ag_ui._endpoint import add_agent_framework_fastapi_endpoint
 
+sys.path.insert(0, str(Path(__file__).parent))
+from test_helpers_ag_ui import StreamingChatClientStub, stream_from_updates
 
-class MockChatClient:
-    """Mock chat client for testing."""
 
-    def __init__(self, response_text: str = "Test response"):
-        self.response_text = response_text
-
-    async def get_streaming_response(self, messages: list[Any], chat_options: Any, **kwargs: Any):
-        """Mock streaming response."""
-        yield ChatResponseUpdate(contents=[TextContent(text=self.response_text)])
+def build_chat_client(response_text: str = "Test response") -> StreamingChatClientStub:
+    """Create a typed chat client stub for endpoint tests."""
+    updates = [ChatResponseUpdate(contents=[TextContent(text=response_text)])]
+    return StreamingChatClientStub(stream_from_updates(updates))
 
 
 async def test_add_endpoint_with_agent_protocol():
     """Test adding endpoint with raw AgentProtocol."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/test-agent")
 
@@ -42,7 +41,7 @@ async def test_add_endpoint_with_agent_protocol():
 async def test_add_endpoint_with_wrapped_agent():
     """Test adding endpoint with pre-wrapped AgentFrameworkAgent."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
     wrapped_agent = AgentFrameworkAgent(agent=agent, name="wrapped")
 
     add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/wrapped-agent")
@@ -57,7 +56,7 @@ async def test_add_endpoint_with_wrapped_agent():
 async def test_endpoint_with_state_schema():
     """Test endpoint with state_schema parameter."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
     state_schema = {"document": {"type": "string"}}
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/stateful", state_schema=state_schema)
@@ -70,10 +69,37 @@ async def test_endpoint_with_state_schema():
     assert response.status_code == 200
 
 
+async def test_endpoint_with_default_state_seed():
+    """Test endpoint seeds default state when client omits it."""
+    app = FastAPI()
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    state_schema = {"proverbs": {"type": "array"}}
+    default_state = {"proverbs": ["Keep the original."]}
+
+    add_agent_framework_fastapi_endpoint(
+        app,
+        agent,
+        path="/default-state",
+        state_schema=state_schema,
+        default_state=default_state,
+    )
+
+    client = TestClient(app)
+    response = client.post("/default-state", json={"messages": [{"role": "user", "content": "Hello"}]})
+
+    assert response.status_code == 200
+
+    content = response.content.decode("utf-8")
+    lines = [line for line in content.split("\n") if line.startswith("data: ")]
+    snapshots = [json.loads(line[6:]) for line in lines if json.loads(line[6:]).get("type") == "STATE_SNAPSHOT"]
+    assert snapshots, "Expected a STATE_SNAPSHOT event"
+    assert snapshots[0]["snapshot"]["proverbs"] == default_state["proverbs"]
+
+
 async def test_endpoint_with_predict_state_config():
     """Test endpoint with predict_state_config parameter."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
     predict_config = {"document": {"tool": "write_doc", "tool_argument": "content"}}
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/predictive", predict_state_config=predict_config)
@@ -87,7 +113,7 @@ async def test_endpoint_with_predict_state_config():
 async def test_endpoint_request_logging():
     """Test that endpoint logs request details."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/logged")
 
@@ -107,7 +133,7 @@ async def test_endpoint_request_logging():
 async def test_endpoint_event_streaming():
     """Test that endpoint streams events correctly."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient("Streamed response"))
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client("Streamed response"))
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/stream")
 
@@ -141,14 +167,14 @@ async def test_endpoint_event_streaming():
 async def test_endpoint_error_handling():
     """Test endpoint error handling during request parsing."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/failing")
 
     client = TestClient(app)
 
     # Send invalid JSON to trigger parsing error before streaming
-    response = client.post("/failing", data="invalid json", headers={"content-type": "application/json"})
+    response = client.post("/failing", data=b"invalid json", headers={"content-type": "application/json"})  # type: ignore
 
     # The exception handler catches it and returns JSON error
     assert response.status_code == 200
@@ -160,8 +186,8 @@ async def test_endpoint_error_handling():
 async def test_endpoint_multiple_paths():
     """Test adding multiple endpoints with different paths."""
     app = FastAPI()
-    agent1 = ChatAgent(name="agent1", instructions="First agent", chat_client=MockChatClient("Response 1"))
-    agent2 = ChatAgent(name="agent2", instructions="Second agent", chat_client=MockChatClient("Response 2"))
+    agent1 = ChatAgent(name="agent1", instructions="First agent", chat_client=build_chat_client("Response 1"))
+    agent2 = ChatAgent(name="agent2", instructions="Second agent", chat_client=build_chat_client("Response 2"))
 
     add_agent_framework_fastapi_endpoint(app, agent1, path="/agent1")
     add_agent_framework_fastapi_endpoint(app, agent2, path="/agent2")
@@ -178,7 +204,7 @@ async def test_endpoint_multiple_paths():
 async def test_endpoint_default_path():
     """Test endpoint with default path."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent)
 
@@ -191,7 +217,7 @@ async def test_endpoint_default_path():
 async def test_endpoint_response_headers():
     """Test that endpoint sets correct response headers."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/headers")
 
@@ -207,7 +233,7 @@ async def test_endpoint_response_headers():
 async def test_endpoint_empty_messages():
     """Test endpoint with empty messages list."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/empty")
 
@@ -220,7 +246,7 @@ async def test_endpoint_empty_messages():
 async def test_endpoint_complex_input():
     """Test endpoint with complex input data."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/complex")