From 35f848d27f9fcae4d1a02f9e38ecf37ae900eb2d Mon Sep 17 00:00:00 2001
From: jnjpng <jin@letta.com>
Date: Tue, 9 Sep 2025 16:29:38 -0700
Subject: [PATCH] fix: handle long running send message integration tests that
 emit LettaPing

---
 .github/workflows/reusable-test-workflow.yml |  6 ++++--
 tests/integration_test_send_message.py       | 21 ++++++++++++++++++++
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/reusable-test-workflow.yml b/.github/workflows/reusable-test-workflow.yml
index dd06c26a..2dcfb785 100644
--- a/.github/workflows/reusable-test-workflow.yml
+++ b/.github/workflows/reusable-test-workflow.yml
@@ -369,8 +369,10 @@ jobs:
 
           # Feature flags (shared across all test types)
           LETTA_ENABLE_BATCH_JOB_POLLING: true
-          LETTA_GEMINI_FORCE_MINIMUM_THINKING_BUDGET: true
-          LETTA_GEMINI_MAX_RETRIES: 10
+
+          # Gemini flags
+          GEMINI_FORCE_MINIMUM_THINKING_BUDGET: true
+          GEMINI_MAX_RETRIES: 10
 
           # Pinecone flags
           LETTA_PINECONE_API_KEY: ${{ secrets.LETTA_PINECONE_API_KEY }}
diff --git a/tests/integration_test_send_message.py b/tests/integration_test_send_message.py
index dd0b6008..da634a46 100644
--- a/tests/integration_test_send_message.py
+++ b/tests/integration_test_send_message.py
@@ -34,6 +34,7 @@ from letta.helpers.reasoning_helper import is_reasoning_completely_disabled
 from letta.llm_api.openai_client import is_openai_reasoning_model
 from letta.log import get_logger
 from letta.schemas.agent import AgentState
+from letta.schemas.letta_ping import LettaPing
 from letta.schemas.llm_config import LLMConfig
 
 logger = get_logger(__name__)
@@ -210,6 +211,10 @@ def assert_greeting_with_assistant_message_response(
     Asserts that the messages list follows the expected sequence:
     ReasoningMessage -> AssistantMessage.
     """
+    # Filter out LettaPing messages which are keep-alive messages for SSE streams
+    messages = [
+        msg for msg in messages if not (isinstance(msg, LettaPing) or (hasattr(msg, "message_type") and msg.message_type == "ping"))
+    ]
     expected_message_count = 4 if streaming else 3 if from_db else 2
     assert len(messages) == expected_message_count
 
@@ -256,6 +261,10 @@ def assert_greeting_no_reasoning_response(
     Asserts that the messages list follows the expected sequence without reasoning:
     AssistantMessage (no ReasoningMessage when put_inner_thoughts_in_kwargs is False).
     """
+    # Filter out LettaPing messages which are keep-alive messages for SSE streams
+    messages = [
+        msg for msg in messages if not (isinstance(msg, LettaPing) or (hasattr(msg, "message_type") and msg.message_type == "ping"))
+    ]
     expected_message_count = 3 if streaming else 2 if from_db else 1
     assert len(messages) == expected_message_count
 
@@ -294,6 +303,10 @@ def assert_greeting_without_assistant_message_response(
     Asserts that the messages list follows the expected sequence:
     ReasoningMessage -> ToolCallMessage -> ToolReturnMessage.
     """
+    # Filter out LettaPing messages which are keep-alive messages for SSE streams
+    messages = [
+        msg for msg in messages if not (isinstance(msg, LettaPing) or (hasattr(msg, "message_type") and msg.message_type == "ping"))
+    ]
     expected_message_count = 5 if streaming else 4 if from_db else 3
     assert len(messages) == expected_message_count
 
@@ -345,6 +358,10 @@ def assert_tool_call_response(
     ReasoningMessage -> ToolCallMessage -> ToolReturnMessage ->
     ReasoningMessage -> AssistantMessage.
     """
+    # Filter out LettaPing messages which are keep-alive messages for SSE streams
+    messages = [
+        msg for msg in messages if not (isinstance(msg, LettaPing) or (hasattr(msg, "message_type") and msg.message_type == "ping"))
+    ]
     expected_message_count = 7 if streaming or from_db else 5
     assert len(messages) == expected_message_count, messages
 
@@ -503,6 +520,10 @@ def assert_image_input_response(
     Asserts that the messages list follows the expected sequence:
     ReasoningMessage -> AssistantMessage.
     """
+    # Filter out LettaPing messages which are keep-alive messages for SSE streams
+    messages = [
+        msg for msg in messages if not (isinstance(msg, LettaPing) or (hasattr(msg, "message_type") and msg.message_type == "ping"))
+    ]
     expected_message_count = 4 if streaming else 3 if from_db else 2
     assert len(messages) == expected_message_count