message syncing WIP

Author: nemozak1

src/schema/schema.py

@@ -243,6 +243,11 @@ class UserInput(BaseModel):
         default=None,
         examples=["847c6285-8fc9-4560-a83f-4e6285809254"],
     )
+    correlation_id: str | None = Field(
+        description="Frontend-generated correlation ID for reliable message matching.",
+        default=None,
+        examples=["550e8400-e29b-41d4-a716-446655440000"],
+    )
     tool_call_approval: ToolCallApproval = Field(
         description="Whether this input is a tool call approval.",
         default=False,

src/service/routers/chat.py

@@ -335,4 +335,94 @@ async def stream_generator():
         ):
             yield stream_manager.to_sse_format(stream_event)
 
-    return StreamingResponse(stream_generator(), media_type="text/event-stream")
+    return StreamingResponse(stream_generator(), media_type="text/event-stream")
+
+
+@router.get("/threads/{thread_id}/messages", dependencies=[Depends(session_cookie)])
+async def get_thread_messages(
+    thread_id: str,
+    stream_manager: StreamManager = Depends(get_stream_manager)
+) -> dict[str, Any]:
+    """
+    Get the authoritative message history for a thread.
+    
+    This endpoint serves as the single source of truth for thread state,
+    enabling frontend recovery/reconciliation when real-time sync fails.
+    
+    Use cases:
+    - Initial load/page refresh
+    - After reconnection
+    - When state seems inconsistent
+    
+    Args:
+        thread_id: The conversation thread ID
+        stream_manager: Injected stream manager dependency
+    
+    Returns:
+        Dictionary containing thread_id and list of messages with their metadata
+    """
+    logger.info(f"Fetching message history for thread: {thread_id}")
+    
+    # Get the graph from the session
+    agent = stream_manager.opey_session.graph
+    
+    # Build config for the thread
+    config = stream_manager.opey_session.build_config({
+        'configurable': {
+            'thread_id': thread_id,
+        }
+    })
+    
+    try:
+        # Get current state
+        current_state = await agent.aget_state(config)
+        
+        if not current_state or not current_state.values:
+            # Thread doesn't exist yet or has no messages
+            return {
+                "thread_id": thread_id,
+                "messages": [],
+                "message_count": 0
+            }
+        
+        # Get messages from state
+        messages = current_state.values.get('messages', [])
+        
+        # Convert LangChain messages to our ChatMessage format
+        formatted_messages = []
+        for msg in messages:
+            try:
+                chat_msg = ChatMessage.from_langchain(msg)
+                formatted_messages.append({
+                    "id": getattr(msg, 'id', None),
+                    "type": chat_msg.type,
+                    "content": chat_msg.content,
+                    "timestamp": getattr(msg, 'timestamp', None),
+                    "tool_calls": chat_msg.tool_calls if hasattr(chat_msg, 'tool_calls') else [],
+                    "tool_call_id": chat_msg.tool_call_id if hasattr(chat_msg, 'tool_call_id') else None,
+                    "tool_status": chat_msg.tool_status if hasattr(chat_msg, 'tool_status') else None,
+                })
+            except Exception as e:
+                logger.warning(f"Failed to convert message to ChatMessage: {e}")
+                # Include a minimal representation if conversion fails
+                formatted_messages.append({
+                    "id": getattr(msg, 'id', None),
+                    "type": msg.__class__.__name__,
+                    "content": str(msg.content) if hasattr(msg, 'content') else "",
+                    "error": f"Conversion failed: {str(e)}"
+                })
+        
+        logger.info(f"Retrieved {len(formatted_messages)} messages for thread {thread_id}")
+        
+        return {
+            "thread_id": thread_id,
+            "messages": formatted_messages,
+            "message_count": len(formatted_messages)
+        }
+        
+    except Exception as e:
+        logger.error(f"Error retrieving thread messages for {thread_id}: {e}", exc_info=True)
+        raise HTTPException(
+            status_code=500,
+            detail=f"Failed to retrieve thread messages: {str(e)}"
+        )

src/service/streaming/events.py

@@ -142,12 +142,22 @@ class UserMessageConfirmEvent(BaseStreamEvent):
     """Event fired when a user message is confirmed with its backend ID"""
     type: Literal["user_message_confirmed"] = "user_message_confirmed"
     message_id: str = Field(description="Backend-assigned message ID")
+    correlation_id: str = Field(description="Frontend-generated correlation ID for reliable matching")
     content: str = Field(description="The user's message content")
 
     def to_sse_data(self) -> str:
         return f"data: {self.model_dump_json()}\n\n"
 
 
+class ThreadSyncEvent(BaseStreamEvent):
+    """Event fired to sync thread_id with the frontend"""
+    type: Literal["thread_sync"] = "thread_sync"
+    thread_id: str = Field(description="Thread ID assigned/confirmed by backend")
+
+    def to_sse_data(self) -> str:
+        return f"data: {self.model_dump_json()}\n\n"
+
+
 class StreamEndEvent(BaseStreamEvent):
     """Event fired when the stream ends"""
     type: Literal["stream_end"] = "stream_end"
@@ -169,6 +179,7 @@ def to_sse_data(self) -> str:
     ApprovalRequestEvent,
     BatchApprovalRequestEvent,
     UserMessageConfirmEvent,
+    ThreadSyncEvent,
     StreamEndEvent
 ]
 
@@ -409,16 +420,34 @@ def batch_approval_request(
         return event
 
     @staticmethod
-    def user_message_confirmed(message_id: str, content: str) -> UserMessageConfirmEvent:
+    def user_message_confirmed(message_id: str, correlation_id: str, content: str) -> UserMessageConfirmEvent:
         """
         Create a user message confirmation event.
         This is sent after the backend accepts a user message and assigns it an ID.
         """
-        event = UserMessageConfirmEvent(message_id=message_id, content=content)
+        event = UserMessageConfirmEvent(
+            message_id=message_id,
+            correlation_id=correlation_id,
+            content=content
+        )
         StreamEventFactory._log_event(
             event,
             "USER_MESSAGE_CONFIRMED",
-            {"message_id": message_id, "content_length": len(content)}
+            {"message_id": message_id, "correlation_id": correlation_id[:8], "content_length": len(content)}
+        )
+        return event
+
+    @staticmethod
+    def thread_sync(thread_id: str) -> ThreadSyncEvent:
+        """
+        Create a thread sync event.
+        This is sent at the start of a stream to sync the thread_id with the frontend.
+        """
+        event = ThreadSyncEvent(thread_id=thread_id)
+        StreamEventFactory._log_event(
+            event,
+            "THREAD_SYNC",
+            {"thread_id": thread_id}
         )
         return event
 

src/service/streaming/processors.py

@@ -67,8 +67,13 @@ async def process(self, event: LangGraphStreamEvent) -> AsyncGenerator[StreamEve
                     # (Compare content to ensure we're syncing the right message)
                     if message_id and content and content == self.stream_input.message:
                         self.user_message_confirmed = True
+                        
+                        # Get correlation_id from stream_input (frontend provides this)
+                        correlation_id = getattr(self.stream_input, 'correlation_id', None) or message_id
+                        
                         yield StreamEventFactory.user_message_confirmed(
                             message_id=message_id,
+                            correlation_id=correlation_id,
                             content=content
                         )
                         return

src/service/streaming/stream_manager.py

@@ -46,6 +46,10 @@ async def stream_response(
             "message_length": len(stream_input.message) if stream_input.message else 0
         })
 
+        # Emit thread_sync event first to ensure frontend has the correct thread_id
+        # This is critical when the backend generates the thread_id (when not provided)
+        yield StreamEventFactory.thread_sync(thread_id)
+
         orchestrator = orchestrator_repository.get_or_create(thread_id, stream_input)
 
         # Track if generator is being closed to avoid yielding in finally block

test/test_correlation_id_implementation.py

@@ -0,0 +1,177 @@
+"""
+Test suite for correlation_id and thread_sync implementation.
+
+Tests verify:
+1. Frontend generates a correlation_id when sending a message
+2. Backend echoes it back in user_message_confirmed event
+3. Backend sends thread_sync event at stream start
+4. Frontend can retrieve authoritative message history via GET endpoint
+"""
+
+import pytest
+import uuid
+import json
+
+from schema.schema import StreamInput
+from service.streaming.events import StreamEventFactory
+
+
+def generate_correlation_id() -> str:
+    """Simulate frontend generating a correlation ID"""
+    return str(uuid.uuid4())
+
+
+def test_stream_input_accepts_correlation_id():
+    """Test that StreamInput accepts and stores correlation_id"""
+    correlation_id = generate_correlation_id()
+    
+    stream_input = StreamInput(
+        message="Test message",
+        thread_id="test-thread-123",
+        correlation_id=correlation_id,
+        stream_tokens=True
+    )
+    
+    assert stream_input.message == "Test message"
+    assert stream_input.thread_id == "test-thread-123"
+    assert stream_input.correlation_id == correlation_id
+    assert stream_input.stream_tokens is True
+
+
+def test_user_message_confirmed_event_includes_correlation_id():
+    """Test that UserMessageConfirmEvent includes correlation_id"""
+    correlation_id = generate_correlation_id()
+    backend_id = str(uuid.uuid4())
+    
+    event = StreamEventFactory.user_message_confirmed(
+        message_id=backend_id,
+        correlation_id=correlation_id,
+        content="Test message content"
+    )
+    
+    assert event.message_id == backend_id
+    assert event.correlation_id == correlation_id
+    assert event.content == "Test message content"
+    assert event.type == "user_message_confirmed"
+
+
+def test_thread_sync_event_creation():
+    """Test that ThreadSyncEvent is created correctly"""
+    thread_id = "test-thread-456"
+    
+    sync_event = StreamEventFactory.thread_sync(thread_id=thread_id)
+    
+    assert sync_event.thread_id == thread_id
+    assert sync_event.type == "thread_sync"
+
+
+def test_user_message_confirmed_sse_serialization():
+    """Test SSE serialization of user_message_confirmed event"""
+    correlation_id = generate_correlation_id()
+    backend_id = str(uuid.uuid4())
+    
+    event = StreamEventFactory.user_message_confirmed(
+        message_id=backend_id,
+        correlation_id=correlation_id,
+        content="Test content"
+    )
+    
+    sse_data = event.to_sse_data()
+    
+    assert sse_data.startswith("data: ")
+    assert sse_data.endswith("\n\n")
+    
+    # Parse the JSON data
+    json_str = sse_data.replace("data: ", "").strip()
+    parsed = json.loads(json_str)
+    
+    assert parsed["type"] == "user_message_confirmed"
+    assert parsed["message_id"] == backend_id
+    assert parsed["correlation_id"] == correlation_id
+    assert parsed["content"] == "Test content"
+
+
+def test_thread_sync_sse_serialization():
+    """Test SSE serialization of thread_sync event"""
+    thread_id = "test-thread-789"
+    
+    sync_event = StreamEventFactory.thread_sync(thread_id=thread_id)
+    sse_data = sync_event.to_sse_data()
+    
+    assert sse_data.startswith("data: ")
+    assert sse_data.endswith("\n\n")
+    
+    # Parse the JSON data
+    json_str = sse_data.replace("data: ", "").strip()
+    parsed = json.loads(json_str)
+    
+    assert parsed["type"] == "thread_sync"
+    assert parsed["thread_id"] == thread_id
+
+
+def test_correlation_id_matching_flow():
+    """
+    Test the complete correlation ID flow for reliable message matching
+    """
+    # 1. Frontend generates correlation_id
+    frontend_correlation_id = generate_correlation_id()
+    
+    # 2. Frontend sends message with correlation_id
+    stream_input = StreamInput(
+        message="Hello, how can I help?",
+        correlation_id=frontend_correlation_id,
+        stream_tokens=True
+    )
+    
+    assert stream_input.correlation_id == frontend_correlation_id
+    
+    # 3. Backend assigns backend message_id
+    backend_message_id = str(uuid.uuid4())
+    
+    # 4. Backend emits user_message_confirmed with both IDs
+    confirm_event = StreamEventFactory.user_message_confirmed(
+        message_id=backend_message_id,
+        correlation_id=frontend_correlation_id,
+        content="Hello, how can I help?"
+    )
+    
+    # 5. Verify frontend can match by correlation_id
+    assert confirm_event.correlation_id == frontend_correlation_id
+    assert confirm_event.message_id == backend_message_id
+    
+    # Frontend would find message with correlation_id and update to backend_id
+    # This eliminates the fragile content-based matching
+
+
+def test_thread_sync_with_provided_thread_id():
+    """Test thread sync when frontend provides thread_id"""
+    frontend_thread_id = "existing-thread-789"
+    
+    # Backend uses provided thread_id
+    sync_event = StreamEventFactory.thread_sync(thread_id=frontend_thread_id)
+    
+    assert sync_event.thread_id == frontend_thread_id
+
+
+def test_thread_sync_with_generated_thread_id():
+    """Test thread sync when backend generates new thread_id"""
+    # Frontend doesn't provide thread_id (new conversation)
+    # Backend generates one
+    backend_thread_id = str(uuid.uuid4())
+    
+    sync_event = StreamEventFactory.thread_sync(thread_id=backend_thread_id)
+    
+    assert sync_event.thread_id == backend_thread_id
+    # Frontend would store this as the authoritative thread_id
+
+
+def test_correlation_id_optional_for_backward_compatibility():
+    """Test that correlation_id is optional for backward compatibility"""
+    # Should work without correlation_id
+    stream_input = StreamInput(
+        message="Test message",
+        stream_tokens=True
+    )
+    
+    assert stream_input.message == "Test message"
+    assert stream_input.correlation_id is None  # Optional field