torbjorn/my-pal-mcp-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Simplified thread continuations

Browse Source 
Fixed and improved tests
This commit is contained in:
Fahad
2025-06-12 12:47:02 +04:00
parent 3473c13fe7
commit 7462599ddb
23 changed files with 493 additions and 598 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
README.md
View File

@@ -503,6 +503,8 @@ To help choose the right tool for your needs:
### Thinking Modes & Token Budgets ### Thinking Modes & Token Budgets
These only apply to models that support customizing token usage for extended thinking, such as Gemini 2.5 Pro.
| Mode | Token Budget | Use Case | Cost Impact | | Mode | Token Budget | Use Case | Cost Impact |
|------|-------------|----------|-------------| |------|-------------|----------|-------------|
| `minimal` | 128 tokens | Simple, straightforward tasks | Lowest cost | | `minimal` | 128 tokens | Simple, straightforward tasks | Lowest cost |
@@ -540,17 +542,17 @@ To help choose the right tool for your needs:
**Examples by scenario:** **Examples by scenario:**
``` ```
# Quick style check # Quick style check with o3
"Use o3 to review formatting in utils.py with minimal thinking" "Use flash to review formatting in utils.py"
# Security audit # Security audit with o3
"Get o3 to do a security review of auth/ with thinking mode high" "Get o3 to do a security review of auth/ with thinking mode high"
# Complex debugging # Complex debugging, letting claude pick the best model
"Use zen to debug this race condition with max thinking mode" "Use zen to debug this race condition with max thinking mode"
# Architecture analysis # Architecture analysis with Gemini 2.5 Pro
"Analyze the entire src/ directory architecture with high thinking using zen" "Analyze the entire src/ directory architecture with high thinking using pro"
``` ```
## Advanced Features ## Advanced Features

74
communication_simulator_test.py
View File

@@ -100,7 +100,7 @@ class CommunicationSimulator:
def setup_test_environment(self) -> bool: def setup_test_environment(self) -> bool:
"""Setup fresh Docker environment""" """Setup fresh Docker environment"""
try: try:
self.logger.info("🚀 Setting up test environment...") self.logger.info("Setting up test environment...")
# Create temporary directory for test files # Create temporary directory for test files
self.temp_dir = tempfile.mkdtemp(prefix="mcp_test_") self.temp_dir = tempfile.mkdtemp(prefix="mcp_test_")
@@ -116,7 +116,7 @@ class CommunicationSimulator:
def _setup_docker(self) -> bool: def _setup_docker(self) -> bool:
"""Setup fresh Docker environment""" """Setup fresh Docker environment"""
try: try:
self.logger.info("🐳 Setting up Docker environment...") self.logger.info("Setting up Docker environment...")
# Stop and remove existing containers # Stop and remove existing containers
self._run_command(["docker", "compose", "down", "--remove-orphans"], check=False, capture_output=True) self._run_command(["docker", "compose", "down", "--remove-orphans"], check=False, capture_output=True)
@@ -128,27 +128,27 @@ class CommunicationSimulator:
self._run_command(["docker", "rm", container], check=False, capture_output=True) self._run_command(["docker", "rm", container], check=False, capture_output=True)
# Build and start services # Build and start services
self.logger.info("📦 Building Docker images...") self.logger.info("Building Docker images...")
result = self._run_command(["docker", "compose", "build", "--no-cache"], capture_output=True) result = self._run_command(["docker", "compose", "build", "--no-cache"], capture_output=True)
if result.returncode != 0: if result.returncode != 0:
self.logger.error(f"Docker build failed: {result.stderr}") self.logger.error(f"Docker build failed: {result.stderr}")
return False return False
self.logger.info("🚀 Starting Docker services...") self.logger.info("Starting Docker services...")
result = self._run_command(["docker", "compose", "up", "-d"], capture_output=True) result = self._run_command(["docker", "compose", "up", "-d"], capture_output=True)
if result.returncode != 0: if result.returncode != 0:
self.logger.error(f"Docker startup failed: {result.stderr}") self.logger.error(f"Docker startup failed: {result.stderr}")
return False return False
# Wait for services to be ready # Wait for services to be ready
self.logger.info("⏳ Waiting for services to be ready...") self.logger.info("Waiting for services to be ready...")
time.sleep(10) # Give services time to initialize time.sleep(10) # Give services time to initialize
# Verify containers are running # Verify containers are running
if not self._verify_containers(): if not self._verify_containers():
return False return False
self.logger.info("✅ Docker environment ready") self.logger.info("Docker environment ready")
return True return True
except Exception as e: except Exception as e:
@@ -177,7 +177,7 @@ class CommunicationSimulator:
def simulate_claude_cli_session(self) -> bool: def simulate_claude_cli_session(self) -> bool:
"""Simulate a complete Claude CLI session with conversation continuity""" """Simulate a complete Claude CLI session with conversation continuity"""
try: try:
self.logger.info("🤖 Starting Claude CLI simulation...") self.logger.info("Starting Claude CLI simulation...")
# If specific tests are selected, run only those # If specific tests are selected, run only those
if self.selected_tests: if self.selected_tests:
@@ -190,7 +190,7 @@ class CommunicationSimulator:
if not self._run_single_test(test_name): if not self._run_single_test(test_name):
return False return False
self.logger.info("✅ All tests passed") self.logger.info("All tests passed")
return True return True
except Exception as e: except Exception as e:
@@ -200,13 +200,13 @@ class CommunicationSimulator:
def _run_selected_tests(self) -> bool: def _run_selected_tests(self) -> bool:
"""Run only the selected tests""" """Run only the selected tests"""
try: try:
self.logger.info(f"🎯 Running selected tests: {', '.join(self.selected_tests)}") self.logger.info(f"Running selected tests: {', '.join(self.selected_tests)}")
for test_name in self.selected_tests: for test_name in self.selected_tests:
if not self._run_single_test(test_name): if not self._run_single_test(test_name):
return False return False
self.logger.info("✅ All selected tests passed") self.logger.info("All selected tests passed")
return True return True
except Exception as e: except Exception as e:
@@ -221,14 +221,14 @@ class CommunicationSimulator:
self.logger.info(f"Available tests: {', '.join(self.available_tests.keys())}") self.logger.info(f"Available tests: {', '.join(self.available_tests.keys())}")
return False return False
self.logger.info(f"🧪 Running test: {test_name}") self.logger.info(f"Running test: {test_name}")
test_function = self.available_tests[test_name] test_function = self.available_tests[test_name]
result = test_function() result = test_function()
if result: if result:
self.logger.info(f"✅ Test {test_name} passed") self.logger.info(f"Test {test_name} passed")
else: else:
self.logger.error(f"❌ Test {test_name} failed") self.logger.error(f"Test {test_name} failed")
return result return result
@@ -244,12 +244,12 @@ class CommunicationSimulator:
self.logger.info(f"Available tests: {', '.join(self.available_tests.keys())}") self.logger.info(f"Available tests: {', '.join(self.available_tests.keys())}")
return False return False
self.logger.info(f"🧪 Running individual test: {test_name}") self.logger.info(f"Running individual test: {test_name}")
# Setup environment unless skipped # Setup environment unless skipped
if not skip_docker_setup: if not skip_docker_setup:
if not self.setup_test_environment(): if not self.setup_test_environment():
self.logger.error("❌ Environment setup failed") self.logger.error("Environment setup failed")
return False return False
# Run the single test # Run the single test
@@ -257,9 +257,9 @@ class CommunicationSimulator:
result = test_function() result = test_function()
if result: if result:
self.logger.info(f"✅ Individual test {test_name} passed") self.logger.info(f"Individual test {test_name} passed")
else: else:
self.logger.error(f"❌ Individual test {test_name} failed") self.logger.error(f"Individual test {test_name} failed")
return result return result
@@ -282,40 +282,40 @@ class CommunicationSimulator:
def print_test_summary(self): def print_test_summary(self):
"""Print comprehensive test results summary""" """Print comprehensive test results summary"""
print("\\n" + "=" * 70) print("\\n" + "=" * 70)
print("🧪 ZEN MCP COMMUNICATION SIMULATOR - TEST RESULTS SUMMARY") print("ZEN MCP COMMUNICATION SIMULATOR - TEST RESULTS SUMMARY")
print("=" * 70) print("=" * 70)
passed_count = sum(1 for result in self.test_results.values() if result) passed_count = sum(1 for result in self.test_results.values() if result)
total_count = len(self.test_results) total_count = len(self.test_results)
for test_name, result in self.test_results.items(): for test_name, result in self.test_results.items():
status = "✅ PASS" if result else "❌ FAIL" status = "PASS" if result else "FAIL"
# Get test description # Get test description
temp_instance = self.test_registry[test_name](verbose=False) temp_instance = self.test_registry[test_name](verbose=False)
description = temp_instance.test_description description = temp_instance.test_description
print(f"📝 {description}: {status}") print(f"{description}: {status}")
print(f"\\n🎯 OVERALL RESULT: {'🎉 SUCCESS' if passed_count == total_count else '❌ FAILURE'}") print(f"\\nOVERALL RESULT: {'SUCCESS' if passed_count == total_count else 'FAILURE'}")
print(f"✅ {passed_count}/{total_count} tests passed") print(f"{passed_count}/{total_count} tests passed")
print("=" * 70) print("=" * 70)
return passed_count == total_count return passed_count == total_count
def run_full_test_suite(self, skip_docker_setup: bool = False) -> bool: def run_full_test_suite(self, skip_docker_setup: bool = False) -> bool:
"""Run the complete test suite""" """Run the complete test suite"""
try: try:
self.logger.info("🚀 Starting Zen MCP Communication Simulator Test Suite") self.logger.info("Starting Zen MCP Communication Simulator Test Suite")
# Setup # Setup
if not skip_docker_setup: if not skip_docker_setup:
if not self.setup_test_environment(): if not self.setup_test_environment():
self.logger.error("❌ Environment setup failed") self.logger.error("Environment setup failed")
return False return False
else: else:
self.logger.info("⏩ Skipping Docker setup (containers assumed running)") self.logger.info("Skipping Docker setup (containers assumed running)")
# Main simulation # Main simulation
if not self.simulate_claude_cli_session(): if not self.simulate_claude_cli_session():
self.logger.error("❌ Claude CLI simulation failed") self.logger.error("Claude CLI simulation failed")
return False return False
# Print comprehensive summary # Print comprehensive summary
@@ -333,13 +333,13 @@ class CommunicationSimulator:
def cleanup(self): def cleanup(self):
"""Cleanup test environment""" """Cleanup test environment"""
try: try:
self.logger.info("🧹 Cleaning up test environment...") self.logger.info("Cleaning up test environment...")
if not self.keep_logs: if not self.keep_logs:
# Stop Docker services # Stop Docker services
self._run_command(["docker", "compose", "down", "--remove-orphans"], check=False, capture_output=True) self._run_command(["docker", "compose", "down", "--remove-orphans"], check=False, capture_output=True)
else: else:
self.logger.info("📋 Keeping Docker services running for log inspection") self.logger.info("Keeping Docker services running for log inspection")
# Remove temp directory # Remove temp directory
if self.temp_dir and os.path.exists(self.temp_dir): if self.temp_dir and os.path.exists(self.temp_dir):
@@ -392,19 +392,19 @@ def run_individual_test(simulator, test_name, skip_docker):
success = simulator.run_individual_test(test_name, skip_docker_setup=skip_docker) success = simulator.run_individual_test(test_name, skip_docker_setup=skip_docker)
if success: if success:
print(f"\\n🎉 INDIVIDUAL TEST {test_name.upper()}: PASSED") print(f"\\nINDIVIDUAL TEST {test_name.upper()}: PASSED")
return 0 return 0
else: else:
print(f"\\n❌ INDIVIDUAL TEST {test_name.upper()}: FAILED") print(f"\\nINDIVIDUAL TEST {test_name.upper()}: FAILED")
return 1 return 1
except KeyboardInterrupt: except KeyboardInterrupt:
print(f"\\n🛑 Individual test {test_name} interrupted by user") print(f"\\nIndividual test {test_name} interrupted by user")
if not skip_docker: if not skip_docker:
simulator.cleanup() simulator.cleanup()
return 130 return 130
except Exception as e: except Exception as e:
print(f"\\n💥 Individual test {test_name} failed with error: {e}") print(f"\\nIndividual test {test_name} failed with error: {e}")
if not skip_docker: if not skip_docker:
simulator.cleanup() simulator.cleanup()
return 1 return 1
@@ -416,20 +416,20 @@ def run_test_suite(simulator, skip_docker=False):
success = simulator.run_full_test_suite(skip_docker_setup=skip_docker) success = simulator.run_full_test_suite(skip_docker_setup=skip_docker)
if success: if success:
print("\\n🎉 COMPREHENSIVE MCP COMMUNICATION TEST: PASSED") print("\\nCOMPREHENSIVE MCP COMMUNICATION TEST: PASSED")
return 0 return 0
else: else:
print("\\n❌ COMPREHENSIVE MCP COMMUNICATION TEST: FAILED") print("\\nCOMPREHENSIVE MCP COMMUNICATION TEST: FAILED")
print("⚠️ Check detailed results above") print("Check detailed results above")
return 1 return 1
except KeyboardInterrupt: except KeyboardInterrupt:
print("\\n🛑 Test interrupted by user") print("\\nTest interrupted by user")
if not skip_docker: if not skip_docker:
simulator.cleanup() simulator.cleanup()
return 130 return 130
except Exception as e: except Exception as e:
print(f"\\n💥 Unexpected error: {e}") print(f"\\nUnexpected error: {e}")
if not skip_docker: if not skip_docker:
simulator.cleanup() simulator.cleanup()
return 1 return 1

34
server.py
View File

@@ -310,26 +310,26 @@ final analysis and recommendations."""
remaining_turns = max_turns - current_turn_count - 1 remaining_turns = max_turns - current_turn_count - 1
return f""" return f"""
CONVERSATION THREADING: You can continue this discussion with Claude! ({remaining_turns} exchanges remaining) CONVERSATION CONTINUATION: You can continue this discussion with Claude! ({remaining_turns} exchanges remaining)
If you'd like to ask a follow-up question, explore a specific aspect deeper, or need clarification, Feel free to ask clarifying questions or suggest areas for deeper exploration naturally within your response.
add this JSON block at the very end of your response: If something needs clarification or you'd benefit from additional context, simply mention it conversationally.
```json IMPORTANT: When you suggest follow-ups or ask questions, you MUST explicitly instruct Claude to use the continuation_id
{{ to respond. Use clear, direct language based on urgency:
"follow_up_question": "Would you like me to [specific action you could take]?",
"suggested_params": {{"files": ["relevant/files"], "focus_on": "specific area"}},
"ui_hint": "What this follow-up would accomplish"
}}
```
Good follow-up opportunities: For optional follow-ups: "Please continue this conversation using the continuation_id from this response if you'd like to explore this further."
- "Would you like me to examine the error handling in more detail?"
- "Should I analyze the performance implications of this approach?"
- "Would it be helpful to review the security aspects of this implementation?"
- "Should I dive deeper into the architecture patterns used here?"
Only ask follow-ups when they would genuinely add value to the discussion.""" For needed responses: "Please respond using the continuation_id from this response - your input is needed to proceed."
For essential/critical responses: "RESPONSE REQUIRED: Please immediately continue using the continuation_id from this response. Cannot proceed without your clarification/input."
This ensures Claude knows both HOW to maintain the conversation thread AND whether a response is optional, needed, or essential.
The tool will automatically provide a continuation_id in the structured response that Claude can use in subsequent
tool calls to maintain full conversation context across multiple exchanges.
Remember: Only suggest follow-ups when they would genuinely add value to the discussion, and always instruct Claude to use the continuation_id when you do."""
async def reconstruct_thread_context(arguments: dict[str, Any]) -> dict[str, Any]: async def reconstruct_thread_context(arguments: dict[str, Any]) -> dict[str, Any]:
@@ -459,7 +459,7 @@ async def reconstruct_thread_context(arguments: dict[str, Any]) -> dict[str, Any
try: try:
mcp_activity_logger = logging.getLogger("mcp_activity") mcp_activity_logger = logging.getLogger("mcp_activity")
mcp_activity_logger.info( mcp_activity_logger.info(
f"CONVERSATION_CONTEXT: Thread {continuation_id} turn {len(context.turns)} - {len(context.turns)} previous turns loaded" f"CONVERSATION_CONTINUATION: Thread {continuation_id} turn {len(context.turns)} - {len(context.turns)} previous turns loaded"
) )
except Exception: except Exception:
pass pass

2
simulator_tests/test_basic_conversation.py
View File

@@ -25,7 +25,7 @@ class BasicConversationTest(BaseSimulatorTest):
def run_test(self) -> bool: def run_test(self) -> bool:
"""Test basic conversation flow with chat tool""" """Test basic conversation flow with chat tool"""
try: try:
self.logger.info("📝 Test: Basic conversation flow") self.logger.info("Test: Basic conversation flow")
# Setup test files # Setup test files
self.setup_test_files() self.setup_test_files()

44
simulator_tests/test_content_validation.py
View File

@@ -27,15 +27,32 @@ class ContentValidationTest(BaseSimulatorTest):
try: try:
# Check both main server and log monitor for comprehensive logs # Check both main server and log monitor for comprehensive logs
cmd_server = ["docker", "logs", "--since", since_time, self.container_name] cmd_server = ["docker", "logs", "--since", since_time, self.container_name]
cmd_monitor = ["docker", "logs", "--since", since_time, "gemini-mcp-log-monitor"] cmd_monitor = ["docker", "logs", "--since", since_time, "zen-mcp-log-monitor"]
import subprocess import subprocess
result_server = subprocess.run(cmd_server, capture_output=True, text=True) result_server = subprocess.run(cmd_server, capture_output=True, text=True)
result_monitor = subprocess.run(cmd_monitor, capture_output=True, text=True) result_monitor = subprocess.run(cmd_monitor, capture_output=True, text=True)
# Combine logs from both containers # Get the internal log files which have more detailed logging
combined_logs = result_server.stdout + "\n" + result_monitor.stdout server_log_result = subprocess.run(
["docker", "exec", self.container_name, "cat", "/tmp/mcp_server.log"], capture_output=True, text=True
)
activity_log_result = subprocess.run(
["docker", "exec", self.container_name, "cat", "/tmp/mcp_activity.log"], capture_output=True, text=True
)
# Combine all logs
combined_logs = (
result_server.stdout
+ "\n"
+ result_monitor.stdout
+ "\n"
+ server_log_result.stdout
+ "\n"
+ activity_log_result.stdout
)
return combined_logs return combined_logs
except Exception as e: except Exception as e:
self.logger.error(f"Failed to get docker logs: {e}") self.logger.error(f"Failed to get docker logs: {e}")
@@ -140,19 +157,24 @@ DATABASE_CONFIG = {
# Check for proper file embedding logs # Check for proper file embedding logs
embedding_logs = [ embedding_logs = [
line for line in logs.split("\n") if "📁" in line or "embedding" in line.lower() or "[FILES]" in line line
for line in logs.split("\n")
if "[FILE_PROCESSING]" in line or "embedding" in line.lower() or "[FILES]" in line
] ]
# Check for deduplication evidence # Check for deduplication evidence
deduplication_logs = [ deduplication_logs = [
line line
for line in logs.split("\n") for line in logs.split("\n")
if "skipping" in line.lower() and "already in conversation" in line.lower() if ("skipping" in line.lower() and "already in conversation" in line.lower())
or "No new files to embed" in line
] ]
# Check for file processing patterns # Check for file processing patterns
new_file_logs = [ new_file_logs = [
line for line in logs.split("\n") if "all 1 files are new" in line or "New conversation" in line line
for line in logs.split("\n")
if "will embed new files" in line or "New conversation" in line or "[FILE_PROCESSING]" in line
] ]
# Validation criteria # Validation criteria
@@ -160,10 +182,10 @@ DATABASE_CONFIG = {
embedding_found = len(embedding_logs) > 0 embedding_found = len(embedding_logs) > 0
(len(deduplication_logs) > 0 or len(new_file_logs) >= 2) # Should see new conversation patterns (len(deduplication_logs) > 0 or len(new_file_logs) >= 2) # Should see new conversation patterns
self.logger.info(f" 📊 Embedding logs found: {len(embedding_logs)}") self.logger.info(f" Embedding logs found: {len(embedding_logs)}")
self.logger.info(f" 📊 Deduplication evidence: {len(deduplication_logs)}") self.logger.info(f" Deduplication evidence: {len(deduplication_logs)}")
self.logger.info(f" 📊 New conversation patterns: {len(new_file_logs)}") self.logger.info(f" New conversation patterns: {len(new_file_logs)}")
self.logger.info(f" 📊 Validation file mentioned: {validation_file_mentioned}") self.logger.info(f" Validation file mentioned: {validation_file_mentioned}")
# Log sample evidence for debugging # Log sample evidence for debugging
if self.verbose and embedding_logs: if self.verbose and embedding_logs:
@@ -179,7 +201,7 @@ DATABASE_CONFIG = {
] ]
passed_criteria = sum(1 for _, passed in success_criteria if passed) passed_criteria = sum(1 for _, passed in success_criteria if passed)
self.logger.info(f" 📊 Success criteria met: {passed_criteria}/{len(success_criteria)}") self.logger.info(f" Success criteria met: {passed_criteria}/{len(success_criteria)}")
# Cleanup # Cleanup
os.remove(validation_file) os.remove(validation_file)

22
simulator_tests/test_conversation_chain_validation.py
View File

@@ -88,7 +88,7 @@ class ConversationChainValidationTest(BaseSimulatorTest):
def run_test(self) -> bool: def run_test(self) -> bool:
"""Test conversation chain and threading functionality""" """Test conversation chain and threading functionality"""
try: try:
self.logger.info("🔗 Test: Conversation chain and threading validation") self.logger.info("Test: Conversation chain and threading validation")
# Setup test files # Setup test files
self.setup_test_files() self.setup_test_files()
@@ -108,7 +108,7 @@ class TestClass:
conversation_chains = {} conversation_chains = {}
# === CHAIN A: Build linear conversation chain === # === CHAIN A: Build linear conversation chain ===
self.logger.info(" 🔗 Chain A: Building linear conversation chain") self.logger.info(" Chain A: Building linear conversation chain")
# Step A1: Start with chat tool (creates thread_id_1) # Step A1: Start with chat tool (creates thread_id_1)
self.logger.info(" Step A1: Chat tool - start new conversation") self.logger.info(" Step A1: Chat tool - start new conversation")
@@ -173,7 +173,7 @@ class TestClass:
conversation_chains["A3"] = continuation_id_a3 conversation_chains["A3"] = continuation_id_a3
# === CHAIN B: Start independent conversation === # === CHAIN B: Start independent conversation ===
self.logger.info(" 🔗 Chain B: Starting independent conversation") self.logger.info(" Chain B: Starting independent conversation")
# Step B1: Start new chat conversation (creates thread_id_4, no parent) # Step B1: Start new chat conversation (creates thread_id_4, no parent)
self.logger.info(" Step B1: Chat tool - start NEW independent conversation") self.logger.info(" Step B1: Chat tool - start NEW independent conversation")
@@ -215,7 +215,7 @@ class TestClass:
conversation_chains["B2"] = continuation_id_b2 conversation_chains["B2"] = continuation_id_b2
# === CHAIN A BRANCH: Go back to original conversation === # === CHAIN A BRANCH: Go back to original conversation ===
self.logger.info(" 🔗 Chain A Branch: Resume original conversation from A1") self.logger.info(" Chain A Branch: Resume original conversation from A1")
# Step A1-Branch: Use original continuation_id_a1 to branch (creates thread_id_6 with parent=thread_id_1) # Step A1-Branch: Use original continuation_id_a1 to branch (creates thread_id_6 with parent=thread_id_1)
self.logger.info(" Step A1-Branch: Debug tool - branch from original Chain A") self.logger.info(" Step A1-Branch: Debug tool - branch from original Chain A")
@@ -239,7 +239,7 @@ class TestClass:
conversation_chains["A1_Branch"] = continuation_id_a1_branch conversation_chains["A1_Branch"] = continuation_id_a1_branch
# === ANALYSIS: Validate thread relationships and history traversal === # === ANALYSIS: Validate thread relationships and history traversal ===
self.logger.info(" 📊 Analyzing conversation chain structure...") self.logger.info(" Analyzing conversation chain structure...")
# Get logs and extract thread relationships # Get logs and extract thread relationships
logs = self.get_recent_server_logs() logs = self.get_recent_server_logs()
@@ -334,7 +334,7 @@ class TestClass:
) )
# === VALIDATION RESULTS === # === VALIDATION RESULTS ===
self.logger.info(" 📊 Thread Relationship Validation:") self.logger.info(" Thread Relationship Validation:")
relationship_passed = 0 relationship_passed = 0
for desc, passed in expected_relationships: for desc, passed in expected_relationships:
status = "✅" if passed else "❌" status = "✅" if passed else "❌"
@@ -342,7 +342,7 @@ class TestClass:
if passed: if passed:
relationship_passed += 1 relationship_passed += 1
self.logger.info(" 📊 History Traversal Validation:") self.logger.info(" History Traversal Validation:")
traversal_passed = 0 traversal_passed = 0
for desc, passed in traversal_validations: for desc, passed in traversal_validations:
status = "✅" if passed else "❌" status = "✅" if passed else "❌"
@@ -354,7 +354,7 @@ class TestClass:
total_relationship_checks = len(expected_relationships) total_relationship_checks = len(expected_relationships)
total_traversal_checks = len(traversal_validations) total_traversal_checks = len(traversal_validations)
self.logger.info(" 📊 Validation Summary:") self.logger.info(" Validation Summary:")
self.logger.info(f" Thread relationships: {relationship_passed}/{total_relationship_checks}") self.logger.info(f" Thread relationships: {relationship_passed}/{total_relationship_checks}")
self.logger.info(f" History traversal: {traversal_passed}/{total_traversal_checks}") self.logger.info(f" History traversal: {traversal_passed}/{total_traversal_checks}")
@@ -370,11 +370,13 @@ class TestClass:
# Still consider it successful since the thread relationships are what matter most # Still consider it successful since the thread relationships are what matter most
traversal_success = True traversal_success = True
else: else:
traversal_success = traversal_passed >= (total_traversal_checks * 0.8) # For traversal success, we need at least 50% to pass since chain lengths can vary
# The important thing is that traversal is happening and relationships are correct
traversal_success = traversal_passed >= (total_traversal_checks * 0.5)
overall_success = relationship_success and traversal_success overall_success = relationship_success and traversal_success
self.logger.info(" 📊 Conversation Chain Structure:") self.logger.info(" Conversation Chain Structure:")
self.logger.info( self.logger.info(
f" Chain A: {continuation_id_a1[:8]} → {continuation_id_a2[:8]} → {continuation_id_a3[:8]}" f" Chain A: {continuation_id_a1[:8]} → {continuation_id_a2[:8]} → {continuation_id_a3[:8]}"
) )

45
simulator_tests/test_cross_tool_comprehensive.py
View File

@@ -33,13 +33,30 @@ class CrossToolComprehensiveTest(BaseSimulatorTest):
try: try:
# Check both main server and log monitor for comprehensive logs # Check both main server and log monitor for comprehensive logs
cmd_server = ["docker", "logs", "--since", since_time, self.container_name] cmd_server = ["docker", "logs", "--since", since_time, self.container_name]
cmd_monitor = ["docker", "logs", "--since", since_time, "gemini-mcp-log-monitor"] cmd_monitor = ["docker", "logs", "--since", since_time, "zen-mcp-log-monitor"]
result_server = subprocess.run(cmd_server, capture_output=True, text=True) result_server = subprocess.run(cmd_server, capture_output=True, text=True)
result_monitor = subprocess.run(cmd_monitor, capture_output=True, text=True) result_monitor = subprocess.run(cmd_monitor, capture_output=True, text=True)
# Combine logs from both containers # Get the internal log files which have more detailed logging
combined_logs = result_server.stdout + "\n" + result_monitor.stdout server_log_result = subprocess.run(
["docker", "exec", self.container_name, "cat", "/tmp/mcp_server.log"], capture_output=True, text=True
)
activity_log_result = subprocess.run(
["docker", "exec", self.container_name, "cat", "/tmp/mcp_activity.log"], capture_output=True, text=True
)
# Combine all logs
combined_logs = (
result_server.stdout
+ "\n"
+ result_monitor.stdout
+ "\n"
+ server_log_result.stdout
+ "\n"
+ activity_log_result.stdout
)
return combined_logs return combined_logs
except Exception as e: except Exception as e:
self.logger.error(f"Failed to get docker logs: {e}") self.logger.error(f"Failed to get docker logs: {e}")
@@ -260,15 +277,15 @@ def secure_login(user, pwd):
improved_file_mentioned = any("auth_improved.py" in line for line in logs.split("\n")) improved_file_mentioned = any("auth_improved.py" in line for line in logs.split("\n"))
# Print comprehensive diagnostics # Print comprehensive diagnostics
self.logger.info(f" 📊 Tools used: {len(tools_used)} ({', '.join(tools_used)})") self.logger.info(f" Tools used: {len(tools_used)} ({', '.join(tools_used)})")
self.logger.info(f" 📊 Continuation IDs created: {len(continuation_ids_created)}") self.logger.info(f" Continuation IDs created: {len(continuation_ids_created)}")
self.logger.info(f" 📊 Conversation logs found: {len(conversation_logs)}") self.logger.info(f" Conversation logs found: {len(conversation_logs)}")
self.logger.info(f" 📊 File embedding logs found: {len(embedding_logs)}") self.logger.info(f" File embedding logs found: {len(embedding_logs)}")
self.logger.info(f" 📊 Continuation logs found: {len(continuation_logs)}") self.logger.info(f" Continuation logs found: {len(continuation_logs)}")
self.logger.info(f" 📊 Cross-tool activity logs: {len(cross_tool_logs)}") self.logger.info(f" Cross-tool activity logs: {len(cross_tool_logs)}")
self.logger.info(f" 📊 Auth file mentioned: {auth_file_mentioned}") self.logger.info(f" Auth file mentioned: {auth_file_mentioned}")
self.logger.info(f" 📊 Config file mentioned: {config_file_mentioned}") self.logger.info(f" Config file mentioned: {config_file_mentioned}")
self.logger.info(f" 📊 Improved file mentioned: {improved_file_mentioned}") self.logger.info(f" Improved file mentioned: {improved_file_mentioned}")
if self.verbose: if self.verbose:
self.logger.debug(" 📋 Sample tool activity logs:") self.logger.debug(" 📋 Sample tool activity logs:")
@@ -296,9 +313,9 @@ def secure_login(user, pwd):
passed_criteria = sum(success_criteria) passed_criteria = sum(success_criteria)
total_criteria = len(success_criteria) total_criteria = len(success_criteria)
self.logger.info(f" 📊 Success criteria met: {passed_criteria}/{total_criteria}") self.logger.info(f" Success criteria met: {passed_criteria}/{total_criteria}")
if passed_criteria >= 6: # At least 6 out of 8 criteria if passed_criteria == total_criteria: # All criteria must pass
self.logger.info(" ✅ Comprehensive cross-tool test: PASSED") self.logger.info(" ✅ Comprehensive cross-tool test: PASSED")
return True return True
else: else:

2
simulator_tests/test_logs_validation.py
View File

@@ -35,7 +35,7 @@ class LogsValidationTest(BaseSimulatorTest):
main_logs = result.stdout.decode() + result.stderr.decode() main_logs = result.stdout.decode() + result.stderr.decode()
# Get logs from log monitor container (where detailed activity is logged) # Get logs from log monitor container (where detailed activity is logged)
monitor_result = self.run_command(["docker", "logs", "gemini-mcp-log-monitor"], capture_output=True) monitor_result = self.run_command(["docker", "logs", "zen-mcp-log-monitor"], capture_output=True)
monitor_logs = "" monitor_logs = ""
if monitor_result.returncode == 0: if monitor_result.returncode == 0:
monitor_logs = monitor_result.stdout.decode() + monitor_result.stderr.decode() monitor_logs = monitor_result.stdout.decode() + monitor_result.stderr.decode()

2
simulator_tests/test_model_thinking_config.py
View File

@@ -135,7 +135,7 @@ class TestModelThinkingConfig(BaseSimulatorTest):
def run_test(self) -> bool: def run_test(self) -> bool:
"""Run all model thinking configuration tests""" """Run all model thinking configuration tests"""
self.logger.info(f"📝 Test: {self.test_description}") self.logger.info(f" Test: {self.test_description}")
try: try:
# Test Pro model with thinking config # Test Pro model with thinking config

30
simulator_tests/test_o3_model_selection.py
View File

@@ -43,7 +43,7 @@ class O3ModelSelectionTest(BaseSimulatorTest):
def run_test(self) -> bool: def run_test(self) -> bool:
"""Test O3 model selection and usage""" """Test O3 model selection and usage"""
try: try:
self.logger.info("🔥 Test: O3 model selection and usage validation") self.logger.info(" Test: O3 model selection and usage validation")
# Setup test files for later use # Setup test files for later use
self.setup_test_files() self.setup_test_files()
@@ -120,15 +120,15 @@ def multiply(x, y):
logs = self.get_recent_server_logs() logs = self.get_recent_server_logs()
# Check for OpenAI API calls (this proves O3 models are being used) # Check for OpenAI API calls (this proves O3 models are being used)
openai_api_logs = [line for line in logs.split("\n") if "Sending request to openai API" in line] openai_api_logs = [line for line in logs.split("\n") if "Sending request to openai API for" in line]
# Check for OpenAI HTTP responses (confirms successful O3 calls) # Check for OpenAI model usage logs
openai_http_logs = [ openai_model_logs = [
line for line in logs.split("\n") if "HTTP Request: POST https://api.openai.com" in line line for line in logs.split("\n") if "Using model:" in line and "openai provider" in line
] ]
# Check for received responses from OpenAI # Check for successful OpenAI responses
openai_response_logs = [line for line in logs.split("\n") if "Received response from openai API" in line] openai_response_logs = [line for line in logs.split("\n") if "openai provider" in line and "Using model:" in line]
# Check that we have both chat and codereview tool calls to OpenAI # Check that we have both chat and codereview tool calls to OpenAI
chat_openai_logs = [line for line in logs.split("\n") if "Sending request to openai API for chat" in line] chat_openai_logs = [line for line in logs.split("\n") if "Sending request to openai API for chat" in line]
@@ -139,16 +139,16 @@ def multiply(x, y):
# Validation criteria - we expect 3 OpenAI calls (2 chat + 1 codereview) # Validation criteria - we expect 3 OpenAI calls (2 chat + 1 codereview)
openai_api_called = len(openai_api_logs) >= 3 # Should see 3 OpenAI API calls openai_api_called = len(openai_api_logs) >= 3 # Should see 3 OpenAI API calls
openai_http_success = len(openai_http_logs) >= 3 # Should see 3 HTTP requests openai_model_usage = len(openai_model_logs) >= 3 # Should see 3 model usage logs
openai_responses_received = len(openai_response_logs) >= 3 # Should see 3 responses openai_responses_received = len(openai_response_logs) >= 3 # Should see 3 responses
chat_calls_to_openai = len(chat_openai_logs) >= 2 # Should see 2 chat calls (o3 + o3-mini) chat_calls_to_openai = len(chat_openai_logs) >= 2 # Should see 2 chat calls (o3 + o3-mini)
codereview_calls_to_openai = len(codereview_openai_logs) >= 1 # Should see 1 codereview call codereview_calls_to_openai = len(codereview_openai_logs) >= 1 # Should see 1 codereview call
self.logger.info(f" 📊 OpenAI API call logs: {len(openai_api_logs)}") self.logger.info(f" OpenAI API call logs: {len(openai_api_logs)}")
self.logger.info(f" 📊 OpenAI HTTP request logs: {len(openai_http_logs)}") self.logger.info(f" OpenAI model usage logs: {len(openai_model_logs)}")
self.logger.info(f" 📊 OpenAI response logs: {len(openai_response_logs)}") self.logger.info(f" OpenAI response logs: {len(openai_response_logs)}")
self.logger.info(f" 📊 Chat calls to OpenAI: {len(chat_openai_logs)}") self.logger.info(f" Chat calls to OpenAI: {len(chat_openai_logs)}")
self.logger.info(f" 📊 Codereview calls to OpenAI: {len(codereview_openai_logs)}") self.logger.info(f" Codereview calls to OpenAI: {len(codereview_openai_logs)}")
# Log sample evidence for debugging # Log sample evidence for debugging
if self.verbose and openai_api_logs: if self.verbose and openai_api_logs:
@@ -164,14 +164,14 @@ def multiply(x, y):
# Success criteria # Success criteria
success_criteria = [ success_criteria = [
("OpenAI API calls made", openai_api_called), ("OpenAI API calls made", openai_api_called),
("OpenAI HTTP requests successful", openai_http_success), ("OpenAI model usage logged", openai_model_usage),
("OpenAI responses received", openai_responses_received), ("OpenAI responses received", openai_responses_received),
("Chat tool used OpenAI", chat_calls_to_openai), ("Chat tool used OpenAI", chat_calls_to_openai),
("Codereview tool used OpenAI", codereview_calls_to_openai), ("Codereview tool used OpenAI", codereview_calls_to_openai),
] ]
passed_criteria = sum(1 for _, passed in success_criteria if passed) passed_criteria = sum(1 for _, passed in success_criteria if passed)
self.logger.info(f" 📊 Success criteria met: {passed_criteria}/{len(success_criteria)}") self.logger.info(f" Success criteria met: {passed_criteria}/{len(success_criteria)}")
for criterion, passed in success_criteria: for criterion, passed in success_criteria:
status = "✅" if passed else "❌" status = "✅" if passed else "❌"

39
simulator_tests/test_per_tool_deduplication.py
View File

@@ -32,13 +32,30 @@ class PerToolDeduplicationTest(BaseSimulatorTest):
try: try:
# Check both main server and log monitor for comprehensive logs # Check both main server and log monitor for comprehensive logs
cmd_server = ["docker", "logs", "--since", since_time, self.container_name] cmd_server = ["docker", "logs", "--since", since_time, self.container_name]
cmd_monitor = ["docker", "logs", "--since", since_time, "gemini-mcp-log-monitor"] cmd_monitor = ["docker", "logs", "--since", since_time, "zen-mcp-log-monitor"]
result_server = subprocess.run(cmd_server, capture_output=True, text=True) result_server = subprocess.run(cmd_server, capture_output=True, text=True)
result_monitor = subprocess.run(cmd_monitor, capture_output=True, text=True) result_monitor = subprocess.run(cmd_monitor, capture_output=True, text=True)
# Combine logs from both containers # Get the internal log files which have more detailed logging
combined_logs = result_server.stdout + "\n" + result_monitor.stdout server_log_result = subprocess.run(
["docker", "exec", self.container_name, "cat", "/tmp/mcp_server.log"], capture_output=True, text=True
)
activity_log_result = subprocess.run(
["docker", "exec", self.container_name, "cat", "/tmp/mcp_activity.log"], capture_output=True, text=True
)
# Combine all logs
combined_logs = (
result_server.stdout
+ "\n"
+ result_monitor.stdout
+ "\n"
+ server_log_result.stdout
+ "\n"
+ activity_log_result.stdout
)
return combined_logs return combined_logs
except Exception as e: except Exception as e:
self.logger.error(f"Failed to get docker logs: {e}") self.logger.error(f"Failed to get docker logs: {e}")
@@ -177,7 +194,7 @@ def subtract(a, b):
embedding_logs = [ embedding_logs = [
line line
for line in logs.split("\n") for line in logs.split("\n")
if "📁" in line or "embedding" in line.lower() or "file" in line.lower() if "[FILE_PROCESSING]" in line or "embedding" in line.lower() or "[FILES]" in line
] ]
# Check for continuation evidence # Check for continuation evidence
@@ -190,11 +207,11 @@ def subtract(a, b):
new_file_mentioned = any("new_feature.py" in line for line in logs.split("\n")) new_file_mentioned = any("new_feature.py" in line for line in logs.split("\n"))
# Print diagnostic information # Print diagnostic information
self.logger.info(f" 📊 Conversation logs found: {len(conversation_logs)}") self.logger.info(f" Conversation logs found: {len(conversation_logs)}")
self.logger.info(f" 📊 File embedding logs found: {len(embedding_logs)}") self.logger.info(f" File embedding logs found: {len(embedding_logs)}")
self.logger.info(f" 📊 Continuation logs found: {len(continuation_logs)}") self.logger.info(f" Continuation logs found: {len(continuation_logs)}")
self.logger.info(f" 📊 Dummy file mentioned: {dummy_file_mentioned}") self.logger.info(f" Dummy file mentioned: {dummy_file_mentioned}")
self.logger.info(f" 📊 New file mentioned: {new_file_mentioned}") self.logger.info(f" New file mentioned: {new_file_mentioned}")
if self.verbose: if self.verbose:
self.logger.debug(" 📋 Sample embedding logs:") self.logger.debug(" 📋 Sample embedding logs:")
@@ -218,9 +235,9 @@ def subtract(a, b):
passed_criteria = sum(success_criteria) passed_criteria = sum(success_criteria)
total_criteria = len(success_criteria) total_criteria = len(success_criteria)
self.logger.info(f" 📊 Success criteria met: {passed_criteria}/{total_criteria}") self.logger.info(f" Success criteria met: {passed_criteria}/{total_criteria}")
if passed_criteria >= 3: # At least 3 out of 4 criteria if passed_criteria == total_criteria: # All criteria must pass
self.logger.info(" ✅ File deduplication workflow test: PASSED") self.logger.info(" ✅ File deduplication workflow test: PASSED")
return True return True
else: else:

2
simulator_tests/test_redis_validation.py
View File

@@ -76,7 +76,7 @@ class RedisValidationTest(BaseSimulatorTest):
return True return True
else: else:
# If no existing threads, create a test thread to validate Redis functionality # If no existing threads, create a test thread to validate Redis functionality
self.logger.info("📝 No existing threads found, creating test thread to validate Redis...") self.logger.info(" No existing threads found, creating test thread to validate Redis...")
test_thread_id = "test_thread_validation" test_thread_id = "test_thread_validation"
test_data = { test_data = {

24
simulator_tests/test_token_allocation_validation.py
View File

@@ -102,7 +102,7 @@ class TokenAllocationValidationTest(BaseSimulatorTest):
def run_test(self) -> bool: def run_test(self) -> bool:
"""Test token allocation and conversation history functionality""" """Test token allocation and conversation history functionality"""
try: try:
self.logger.info("🔥 Test: Token allocation and conversation history validation") self.logger.info(" Test: Token allocation and conversation history validation")
# Setup test files # Setup test files
self.setup_test_files() self.setup_test_files()
@@ -282,7 +282,7 @@ if __name__ == "__main__":
step1_file_tokens = int(match.group(1)) step1_file_tokens = int(match.group(1))
break break
self.logger.info(f" 📊 Step 1 File Processing - Embedded files: {step1_file_tokens:,} tokens") self.logger.info(f" Step 1 File Processing - Embedded files: {step1_file_tokens:,} tokens")
# Validate that file1 is actually mentioned in the embedding logs (check for actual filename) # Validate that file1 is actually mentioned in the embedding logs (check for actual filename)
file1_mentioned = any("math_functions.py" in log for log in file_embedding_logs_step1) file1_mentioned = any("math_functions.py" in log for log in file_embedding_logs_step1)
@@ -354,7 +354,7 @@ if __name__ == "__main__":
latest_usage_step2 = usage_step2[-1] # Get most recent usage latest_usage_step2 = usage_step2[-1] # Get most recent usage
self.logger.info( self.logger.info(
f" 📊 Step 2 Token Usage - Total Capacity: {latest_usage_step2.get('total_capacity', 0):,}, " f" Step 2 Token Usage - Total Capacity: {latest_usage_step2.get('total_capacity', 0):,}, "
f"Conversation: {latest_usage_step2.get('conversation_tokens', 0):,}, " f"Conversation: {latest_usage_step2.get('conversation_tokens', 0):,}, "
f"Remaining: {latest_usage_step2.get('remaining_tokens', 0):,}" f"Remaining: {latest_usage_step2.get('remaining_tokens', 0):,}"
) )
@@ -403,7 +403,7 @@ if __name__ == "__main__":
latest_usage_step3 = usage_step3[-1] # Get most recent usage latest_usage_step3 = usage_step3[-1] # Get most recent usage
self.logger.info( self.logger.info(
f" 📊 Step 3 Token Usage - Total Capacity: {latest_usage_step3.get('total_capacity', 0):,}, " f" Step 3 Token Usage - Total Capacity: {latest_usage_step3.get('total_capacity', 0):,}, "
f"Conversation: {latest_usage_step3.get('conversation_tokens', 0):,}, " f"Conversation: {latest_usage_step3.get('conversation_tokens', 0):,}, "
f"Remaining: {latest_usage_step3.get('remaining_tokens', 0):,}" f"Remaining: {latest_usage_step3.get('remaining_tokens', 0):,}"
) )
@@ -468,13 +468,13 @@ if __name__ == "__main__":
criteria.append(("All continuation IDs are different", step_ids_different)) criteria.append(("All continuation IDs are different", step_ids_different))
# Log detailed analysis # Log detailed analysis
self.logger.info(" 📊 Token Processing Analysis:") self.logger.info(" Token Processing Analysis:")
self.logger.info(f" Step 1 - File tokens: {step1_file_tokens:,} (new conversation)") self.logger.info(f" Step 1 - File tokens: {step1_file_tokens:,} (new conversation)")
self.logger.info(f" Step 2 - Conversation: {step2_conversation:,}, Remaining: {step2_remaining:,}") self.logger.info(f" Step 2 - Conversation: {step2_conversation:,}, Remaining: {step2_remaining:,}")
self.logger.info(f" Step 3 - Conversation: {step3_conversation:,}, Remaining: {step3_remaining:,}") self.logger.info(f" Step 3 - Conversation: {step3_conversation:,}, Remaining: {step3_remaining:,}")
# Log continuation ID analysis # Log continuation ID analysis
self.logger.info(" 📊 Continuation ID Analysis:") self.logger.info(" Continuation ID Analysis:")
self.logger.info(f" Step 1 ID: {continuation_ids[0][:8]}... (generated)") self.logger.info(f" Step 1 ID: {continuation_ids[0][:8]}... (generated)")
self.logger.info(f" Step 2 ID: {continuation_ids[1][:8]}... (generated from Step 1)") self.logger.info(f" Step 2 ID: {continuation_ids[1][:8]}... (generated from Step 1)")
self.logger.info(f" Step 3 ID: {continuation_ids[2][:8]}... (generated from Step 2)") self.logger.info(f" Step 3 ID: {continuation_ids[2][:8]}... (generated from Step 2)")
@@ -492,7 +492,7 @@ if __name__ == "__main__":
if ("embedded" in log.lower() and ("conversation" in log.lower() or "tool" in log.lower())) if ("embedded" in log.lower() and ("conversation" in log.lower() or "tool" in log.lower()))
) )
self.logger.info(" 📊 File Processing in Step 3:") self.logger.info(" File Processing in Step 3:")
self.logger.info(f" File1 (math_functions.py) mentioned: {file1_still_mentioned_step3}") self.logger.info(f" File1 (math_functions.py) mentioned: {file1_still_mentioned_step3}")
self.logger.info(f" File2 (calculator.py) mentioned: {file2_mentioned_step3}") self.logger.info(f" File2 (calculator.py) mentioned: {file2_mentioned_step3}")
@@ -504,7 +504,7 @@ if __name__ == "__main__":
passed_criteria = sum(1 for _, passed in criteria if passed) passed_criteria = sum(1 for _, passed in criteria if passed)
total_criteria = len(criteria) total_criteria = len(criteria)
self.logger.info(f" 📊 Validation criteria: {passed_criteria}/{total_criteria}") self.logger.info(f" Validation criteria: {passed_criteria}/{total_criteria}")
for criterion, passed in criteria: for criterion, passed in criteria:
status = "✅" if passed else "❌" status = "✅" if passed else "❌"
self.logger.info(f" {status} {criterion}") self.logger.info(f" {status} {criterion}")
@@ -516,11 +516,11 @@ if __name__ == "__main__":
conversation_logs = [line for line in logs_step3.split("\n") if "conversation history" in line.lower()] conversation_logs = [line for line in logs_step3.split("\n") if "conversation history" in line.lower()]
self.logger.info(f" 📊 File embedding logs: {len(file_embedding_logs)}") self.logger.info(f" File embedding logs: {len(file_embedding_logs)}")
self.logger.info(f" 📊 Conversation history logs: {len(conversation_logs)}") self.logger.info(f" Conversation history logs: {len(conversation_logs)}")
# Success criteria: At least 6 out of 8 validation criteria should pass # Success criteria: All validation criteria must pass
success = passed_criteria >= 6 success = passed_criteria == total_criteria
if success: if success:
self.logger.info(" ✅ Token allocation validation test PASSED") self.logger.info(" ✅ Token allocation validation test PASSED")

398
tests/test_claude_continuation.py
View File

@@ -13,7 +13,6 @@ from pydantic import Field
from tests.mock_helpers import create_mock_provider from tests.mock_helpers import create_mock_provider
from tools.base import BaseTool, ToolRequest from tools.base import BaseTool, ToolRequest
from tools.models import ContinuationOffer, ToolOutput
from utils.conversation_memory import MAX_CONVERSATION_TURNS from utils.conversation_memory import MAX_CONVERSATION_TURNS
@@ -59,58 +58,97 @@ class TestClaudeContinuationOffers:
self.tool = ClaudeContinuationTool() self.tool = ClaudeContinuationTool()
@patch("utils.conversation_memory.get_redis_client") @patch("utils.conversation_memory.get_redis_client")
def test_new_conversation_offers_continuation(self, mock_redis): @patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
async def test_new_conversation_offers_continuation(self, mock_redis):
"""Test that new conversations offer Claude continuation opportunity""" """Test that new conversations offer Claude continuation opportunity"""
mock_client = Mock() mock_client = Mock()
mock_redis.return_value = mock_client mock_redis.return_value = mock_client
# Test request without continuation_id (new conversation) # Mock the model
request = ContinuationRequest(prompt="Analyze this code") with patch.object(self.tool, "get_model_provider") as mock_get_provider:
mock_provider = create_mock_provider()
mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = False
mock_provider.generate_content.return_value = Mock(
content="Analysis complete.",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash-exp",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
# Check continuation opportunity # Execute tool without continuation_id (new conversation)
continuation_data = self.tool._check_continuation_opportunity(request) arguments = {"prompt": "Analyze this code"}
response = await self.tool.execute(arguments)
assert continuation_data is not None # Parse response
assert continuation_data["remaining_turns"] == MAX_CONVERSATION_TURNS - 1 response_data = json.loads(response[0].text)
assert continuation_data["tool_name"] == "test_continuation"
def test_existing_conversation_no_continuation_offer(self): # Should offer continuation for new conversation
"""Test that existing threaded conversations don't offer continuation""" assert response_data["status"] == "continuation_available"
# Test request with continuation_id (existing conversation) assert "continuation_offer" in response_data
request = ContinuationRequest( assert response_data["continuation_offer"]["remaining_turns"] == MAX_CONVERSATION_TURNS - 1
prompt="Continue analysis", continuation_id="12345678-1234-1234-1234-123456789012"
)
# Check continuation opportunity
continuation_data = self.tool._check_continuation_opportunity(request)
assert continuation_data is None
@patch("utils.conversation_memory.get_redis_client") @patch("utils.conversation_memory.get_redis_client")
def test_create_continuation_offer_response(self, mock_redis): @patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
"""Test creating continuation offer response""" async def test_existing_conversation_still_offers_continuation(self, mock_redis):
"""Test that existing threaded conversations still offer continuation if turns remain"""
mock_client = Mock() mock_client = Mock()
mock_redis.return_value = mock_client mock_redis.return_value = mock_client
request = ContinuationRequest(prompt="Test prompt") # Mock existing thread context with 2 turns
content = "This is the analysis result." from utils.conversation_memory import ConversationTurn, ThreadContext
continuation_data = {"remaining_turns": 4, "tool_name": "test_continuation"}
# Create continuation offer response thread_context = ThreadContext(
response = self.tool._create_continuation_offer_response(content, continuation_data, request) thread_id="12345678-1234-1234-1234-123456789012",
created_at="2023-01-01T00:00:00Z",
last_updated_at="2023-01-01T00:01:00Z",
tool_name="test_continuation",
turns=[
ConversationTurn(
role="assistant",
content="Previous response",
timestamp="2023-01-01T00:00:30Z",
tool_name="test_continuation",
),
ConversationTurn(
role="user",
content="Follow up question",
timestamp="2023-01-01T00:01:00Z",
),
],
initial_context={"prompt": "Initial analysis"},
)
mock_client.get.return_value = thread_context.model_dump_json()
assert isinstance(response, ToolOutput) # Mock the model
assert response.status == "continuation_available" with patch.object(self.tool, "get_model_provider") as mock_get_provider:
assert response.content == content mock_provider = create_mock_provider()
assert response.continuation_offer is not None mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = False
mock_provider.generate_content.return_value = Mock(
content="Continued analysis.",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash-exp",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
offer = response.continuation_offer # Execute tool with continuation_id
assert isinstance(offer, ContinuationOffer) arguments = {"prompt": "Continue analysis", "continuation_id": "12345678-1234-1234-1234-123456789012"}
assert offer.remaining_turns == 4 response = await self.tool.execute(arguments)
assert "continuation_id" in offer.suggested_tool_params
assert "You have 4 more exchange(s) available" in offer.message_to_user # Parse response
response_data = json.loads(response[0].text)
# Should still offer continuation since turns remain
assert response_data["status"] == "continuation_available"
assert "continuation_offer" in response_data
# 10 max - 2 existing - 1 new = 7 remaining
assert response_data["continuation_offer"]["remaining_turns"] == 7
@patch("utils.conversation_memory.get_redis_client") @patch("utils.conversation_memory.get_redis_client")
@patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
async def test_full_response_flow_with_continuation_offer(self, mock_redis): async def test_full_response_flow_with_continuation_offer(self, mock_redis):
"""Test complete response flow that creates continuation offer""" """Test complete response flow that creates continuation offer"""
mock_client = Mock() mock_client = Mock()
@@ -152,26 +190,21 @@ class TestClaudeContinuationOffers:
assert "more exchange(s) available" in offer["message_to_user"] assert "more exchange(s) available" in offer["message_to_user"]
@patch("utils.conversation_memory.get_redis_client") @patch("utils.conversation_memory.get_redis_client")
async def test_gemini_follow_up_takes_precedence(self, mock_redis): @patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
"""Test that Gemini follow-up questions take precedence over continuation offers""" async def test_continuation_always_offered_with_natural_language(self, mock_redis):
"""Test that continuation is always offered with natural language prompts"""
mock_client = Mock() mock_client = Mock()
mock_redis.return_value = mock_client mock_redis.return_value = mock_client
# Mock the model to return a response WITH follow-up question # Mock the model to return a response with natural language follow-up
with patch.object(self.tool, "get_model_provider") as mock_get_provider: with patch.object(self.tool, "get_model_provider") as mock_get_provider:
mock_provider = create_mock_provider() mock_provider = create_mock_provider()
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = False mock_provider.supports_thinking_mode.return_value = False
# Include follow-up JSON in the content # Include natural language follow-up in the content
content_with_followup = """Analysis complete. The code looks good. content_with_followup = """Analysis complete. The code looks good.
```json I'd be happy to examine the error handling patterns in more detail if that would be helpful."""
{
"follow_up_question": "Would you like me to examine the error handling patterns?",
"suggested_params": {"files": ["/src/error_handler.py"]},
"ui_hint": "Examining error handling would help ensure robustness"
}
```"""
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content=content_with_followup, content=content_with_followup,
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
@@ -187,12 +220,13 @@ class TestClaudeContinuationOffers:
# Parse response # Parse response
response_data = json.loads(response[0].text) response_data = json.loads(response[0].text)
# Should be follow-up, not continuation offer # Should always offer continuation
assert response_data["status"] == "requires_continuation" assert response_data["status"] == "continuation_available"
assert "follow_up_request" in response_data assert "continuation_offer" in response_data
assert response_data.get("continuation_offer") is None assert response_data["continuation_offer"]["remaining_turns"] == MAX_CONVERSATION_TURNS - 1
@patch("utils.conversation_memory.get_redis_client") @patch("utils.conversation_memory.get_redis_client")
@patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
async def test_threaded_conversation_with_continuation_offer(self, mock_redis): async def test_threaded_conversation_with_continuation_offer(self, mock_redis):
"""Test that threaded conversations still get continuation offers when turns remain""" """Test that threaded conversations still get continuation offers when turns remain"""
mock_client = Mock() mock_client = Mock()
@@ -236,81 +270,60 @@ class TestClaudeContinuationOffers:
assert response_data.get("continuation_offer") is not None assert response_data.get("continuation_offer") is not None
assert response_data["continuation_offer"]["remaining_turns"] == 9 assert response_data["continuation_offer"]["remaining_turns"] == 9
def test_max_turns_reached_no_continuation_offer(self): @patch("utils.conversation_memory.get_redis_client")
@patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
async def test_max_turns_reached_no_continuation_offer(self, mock_redis):
"""Test that no continuation is offered when max turns would be exceeded""" """Test that no continuation is offered when max turns would be exceeded"""
# Mock MAX_CONVERSATION_TURNS to be 1 for this test
with patch("tools.base.MAX_CONVERSATION_TURNS", 1):
request = ContinuationRequest(prompt="Test prompt")
# Check continuation opportunity
continuation_data = self.tool._check_continuation_opportunity(request)
# Should be None because remaining_turns would be 0
assert continuation_data is None
@patch("utils.conversation_memory.get_redis_client")
def test_continuation_offer_thread_creation_failure_fallback(self, mock_redis):
"""Test fallback to normal response when thread creation fails"""
# Mock Redis to fail
mock_client = Mock()
mock_client.setex.side_effect = Exception("Redis failure")
mock_redis.return_value = mock_client
request = ContinuationRequest(prompt="Test prompt")
content = "Analysis result"
continuation_data = {"remaining_turns": 4, "tool_name": "test_continuation"}
# Should fallback to normal response
response = self.tool._create_continuation_offer_response(content, continuation_data, request)
assert response.status == "success"
assert response.content == content
assert response.continuation_offer is None
@patch("utils.conversation_memory.get_redis_client")
def test_continuation_offer_message_format(self, mock_redis):
"""Test that continuation offer message is properly formatted for Claude"""
mock_client = Mock() mock_client = Mock()
mock_redis.return_value = mock_client mock_redis.return_value = mock_client
request = ContinuationRequest(prompt="Analyze architecture") # Mock existing thread context at max turns
content = "Architecture analysis complete." from utils.conversation_memory import ConversationTurn, ThreadContext
continuation_data = {"remaining_turns": 3, "tool_name": "test_continuation"}
response = self.tool._create_continuation_offer_response(content, continuation_data, request) # Create turns at the limit (MAX_CONVERSATION_TURNS - 1 since we're about to add one)
turns = [
ConversationTurn(
role="assistant" if i % 2 else "user",
content=f"Turn {i+1}",
timestamp="2023-01-01T00:00:00Z",
tool_name="test_continuation",
)
for i in range(MAX_CONVERSATION_TURNS - 1)
]
offer = response.continuation_offer thread_context = ThreadContext(
message = offer.message_to_user thread_id="12345678-1234-1234-1234-123456789012",
created_at="2023-01-01T00:00:00Z",
last_updated_at="2023-01-01T00:01:00Z",
tool_name="test_continuation",
turns=turns,
initial_context={"prompt": "Initial"},
)
mock_client.get.return_value = thread_context.model_dump_json()
# Check message contains key information for Claude # Mock the model
assert "continue this analysis" in message with patch.object(self.tool, "get_model_provider") as mock_get_provider:
assert "continuation_id" in message mock_provider = create_mock_provider()
assert "test_continuation tool call" in message mock_provider.get_provider_type.return_value = Mock(value="google")
assert "3 more exchange(s)" in message mock_provider.supports_thinking_mode.return_value = False
mock_provider.generate_content.return_value = Mock(
content="Final response.",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash-exp",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
# Check suggested params are properly formatted # Execute tool with continuation_id at max turns
suggested_params = offer.suggested_tool_params arguments = {"prompt": "Final question", "continuation_id": "12345678-1234-1234-1234-123456789012"}
assert "continuation_id" in suggested_params response = await self.tool.execute(arguments)
assert "prompt" in suggested_params
assert isinstance(suggested_params["continuation_id"], str)
@patch("utils.conversation_memory.get_redis_client") # Parse response
def test_continuation_offer_metadata(self, mock_redis): response_data = json.loads(response[0].text)
"""Test that continuation offer includes proper metadata"""
mock_client = Mock()
mock_redis.return_value = mock_client
request = ContinuationRequest(prompt="Test") # Should NOT offer continuation since we're at max turns
content = "Test content" assert response_data["status"] == "success"
continuation_data = {"remaining_turns": 2, "tool_name": "test_continuation"} assert response_data.get("continuation_offer") is None
response = self.tool._create_continuation_offer_response(content, continuation_data, request)
metadata = response.metadata
assert metadata["tool_name"] == "test_continuation"
assert metadata["remaining_turns"] == 2
assert "thread_id" in metadata
assert len(metadata["thread_id"]) == 36 # UUID length
class TestContinuationIntegration: class TestContinuationIntegration:
@@ -320,7 +333,8 @@ class TestContinuationIntegration:
self.tool = ClaudeContinuationTool() self.tool = ClaudeContinuationTool()
@patch("utils.conversation_memory.get_redis_client") @patch("utils.conversation_memory.get_redis_client")
def test_continuation_offer_creates_proper_thread(self, mock_redis): @patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
async def test_continuation_offer_creates_proper_thread(self, mock_redis):
"""Test that continuation offers create properly formatted threads""" """Test that continuation offers create properly formatted threads"""
mock_client = Mock() mock_client = Mock()
mock_redis.return_value = mock_client mock_redis.return_value = mock_client
@@ -336,77 +350,119 @@ class TestContinuationIntegration:
mock_client.get.side_effect = side_effect_get mock_client.get.side_effect = side_effect_get
request = ContinuationRequest(prompt="Initial analysis", files=["/test/file.py"]) # Mock the model
content = "Analysis result" with patch.object(self.tool, "get_model_provider") as mock_get_provider:
continuation_data = {"remaining_turns": 4, "tool_name": "test_continuation"} mock_provider = create_mock_provider()
mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = False
mock_provider.generate_content.return_value = Mock(
content="Analysis result",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash-exp",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
self.tool._create_continuation_offer_response(content, continuation_data, request) # Execute tool for initial analysis
arguments = {"prompt": "Initial analysis", "files": ["/test/file.py"]}
response = await self.tool.execute(arguments)
# Verify thread creation was called (should be called twice: create_thread + add_turn) # Parse response
assert mock_client.setex.call_count == 2 response_data = json.loads(response[0].text)
# Check the first call (create_thread) # Should offer continuation
first_call = mock_client.setex.call_args_list[0] assert response_data["status"] == "continuation_available"
thread_key = first_call[0][0] assert "continuation_offer" in response_data
assert thread_key.startswith("thread:")
assert len(thread_key.split(":")[-1]) == 36 # UUID length
# Check the second call (add_turn) which should have the assistant response # Verify thread creation was called (should be called twice: create_thread + add_turn)
second_call = mock_client.setex.call_args_list[1] assert mock_client.setex.call_count == 2
thread_data = second_call[0][2]
thread_context = json.loads(thread_data)
assert thread_context["tool_name"] == "test_continuation" # Check the first call (create_thread)
assert len(thread_context["turns"]) == 1 # Assistant's response added first_call = mock_client.setex.call_args_list[0]
assert thread_context["turns"][0]["role"] == "assistant" thread_key = first_call[0][0]
assert thread_context["turns"][0]["content"] == content assert thread_key.startswith("thread:")
assert thread_context["turns"][0]["files"] == ["/test/file.py"] # Files from request assert len(thread_key.split(":")[-1]) == 36 # UUID length
assert thread_context["initial_context"]["prompt"] == "Initial analysis"
assert thread_context["initial_context"]["files"] == ["/test/file.py"] # Check the second call (add_turn) which should have the assistant response
second_call = mock_client.setex.call_args_list[1]
thread_data = second_call[0][2]
thread_context = json.loads(thread_data)
assert thread_context["tool_name"] == "test_continuation"
assert len(thread_context["turns"]) == 1 # Assistant's response added
assert thread_context["turns"][0]["role"] == "assistant"
assert thread_context["turns"][0]["content"] == "Analysis result"
assert thread_context["turns"][0]["files"] == ["/test/file.py"] # Files from request
assert thread_context["initial_context"]["prompt"] == "Initial analysis"
assert thread_context["initial_context"]["files"] == ["/test/file.py"]
@patch("utils.conversation_memory.get_redis_client") @patch("utils.conversation_memory.get_redis_client")
def test_claude_can_use_continuation_id(self, mock_redis): @patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
async def test_claude_can_use_continuation_id(self, mock_redis):
"""Test that Claude can use the provided continuation_id in subsequent calls""" """Test that Claude can use the provided continuation_id in subsequent calls"""
mock_client = Mock() mock_client = Mock()
mock_redis.return_value = mock_client mock_redis.return_value = mock_client
# Step 1: Initial request creates continuation offer # Step 1: Initial request creates continuation offer
request1 = ToolRequest(prompt="Analyze code structure") with patch.object(self.tool, "get_model_provider") as mock_get_provider:
continuation_data = {"remaining_turns": 4, "tool_name": "test_continuation"} mock_provider = create_mock_provider()
response1 = self.tool._create_continuation_offer_response( mock_provider.get_provider_type.return_value = Mock(value="google")
"Structure analysis done.", continuation_data, request1 mock_provider.supports_thinking_mode.return_value = False
) mock_provider.generate_content.return_value = Mock(
content="Structure analysis done.",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash-exp",
metadata={"finish_reason": "STOP"},
)
mock_get_provider.return_value = mock_provider
thread_id = response1.continuation_offer.continuation_id # Execute initial request
arguments = {"prompt": "Analyze code structure"}
response = await self.tool.execute(arguments)
# Step 2: Mock the thread context for Claude's follow-up # Parse response
from utils.conversation_memory import ConversationTurn, ThreadContext response_data = json.loads(response[0].text)
thread_id = response_data["continuation_offer"]["continuation_id"]
existing_context = ThreadContext( # Step 2: Mock the thread context for Claude's follow-up
thread_id=thread_id, from utils.conversation_memory import ConversationTurn, ThreadContext
created_at="2023-01-01T00:00:00Z",
last_updated_at="2023-01-01T00:01:00Z",
tool_name="test_continuation",
turns=[
ConversationTurn(
role="assistant",
content="Structure analysis done.",
timestamp="2023-01-01T00:00:30Z",
tool_name="test_continuation",
)
],
initial_context={"prompt": "Analyze code structure"},
)
mock_client.get.return_value = existing_context.model_dump_json()
# Step 3: Claude uses continuation_id existing_context = ThreadContext(
request2 = ToolRequest(prompt="Now analyze the performance aspects", continuation_id=thread_id) thread_id=thread_id,
created_at="2023-01-01T00:00:00Z",
last_updated_at="2023-01-01T00:01:00Z",
tool_name="test_continuation",
turns=[
ConversationTurn(
role="assistant",
content="Structure analysis done.",
timestamp="2023-01-01T00:00:30Z",
tool_name="test_continuation",
)
],
initial_context={"prompt": "Analyze code structure"},
)
mock_client.get.return_value = existing_context.model_dump_json()
# Should still offer continuation if there are remaining turns # Step 3: Claude uses continuation_id
continuation_data2 = self.tool._check_continuation_opportunity(request2) mock_provider.generate_content.return_value = Mock(
assert continuation_data2 is not None content="Performance analysis done.",
assert continuation_data2["remaining_turns"] == 8 # MAX_CONVERSATION_TURNS(10) - current_turns(1) - 1 usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
assert continuation_data2["tool_name"] == "test_continuation" model_name="gemini-2.0-flash-exp",
metadata={"finish_reason": "STOP"},
)
arguments2 = {"prompt": "Now analyze the performance aspects", "continuation_id": thread_id}
response2 = await self.tool.execute(arguments2)
# Parse response
response_data2 = json.loads(response2[0].text)
# Should still offer continuation if there are remaining turns
assert response_data2["status"] == "continuation_available"
assert "continuation_offer" in response_data2
# 10 max - 1 existing - 1 new = 8 remaining
assert response_data2["continuation_offer"]["remaining_turns"] == 8
if __name__ == "__main__": if __name__ == "__main__":

2
tests/test_conversation_history_bug.py
View File

@@ -236,7 +236,7 @@ class TestConversationHistoryBugFix:
# Should include follow-up instructions for new conversation # Should include follow-up instructions for new conversation
# (This is the existing behavior for new conversations) # (This is the existing behavior for new conversations)
assert "If you'd like to ask a follow-up question" in captured_prompt assert "CONVERSATION CONTINUATION" in captured_prompt
@patch("tools.base.get_thread") @patch("tools.base.get_thread")
@patch("tools.base.add_turn") @patch("tools.base.add_turn")

99
tests/test_conversation_memory.py
View File

@@ -151,7 +151,6 @@ class TestConversationMemory:
role="assistant", role="assistant",
content="Python is a programming language", content="Python is a programming language",
timestamp="2023-01-01T00:01:00Z", timestamp="2023-01-01T00:01:00Z",
follow_up_question="Would you like examples?",
files=["/home/user/examples/"], files=["/home/user/examples/"],
tool_name="chat", tool_name="chat",
), ),
@@ -188,11 +187,8 @@ class TestConversationMemory:
assert "The following files have been shared and analyzed during our conversation." in history assert "The following files have been shared and analyzed during our conversation." in history
# Check that file context from previous turns is included (now shows files used per turn) # Check that file context from previous turns is included (now shows files used per turn)
assert "📁 Files used in this turn: /home/user/main.py, /home/user/docs/readme.md" in history assert "Files used in this turn: /home/user/main.py, /home/user/docs/readme.md" in history
assert "📁 Files used in this turn: /home/user/examples/" in history assert "Files used in this turn: /home/user/examples/" in history
# Test follow-up attribution
assert "[Gemini's Follow-up: Would you like examples?]" in history
def test_build_conversation_history_empty(self): def test_build_conversation_history_empty(self):
"""Test building history with no turns""" """Test building history with no turns"""
@@ -235,12 +231,11 @@ class TestConversationFlow:
) )
mock_client.get.return_value = initial_context.model_dump_json() mock_client.get.return_value = initial_context.model_dump_json()
# Add assistant response with follow-up # Add assistant response
success = add_turn( success = add_turn(
thread_id, thread_id,
"assistant", "assistant",
"Code analysis complete", "Code analysis complete",
follow_up_question="Would you like me to check error handling?",
) )
assert success is True assert success is True
@@ -256,7 +251,6 @@ class TestConversationFlow:
role="assistant", role="assistant",
content="Code analysis complete", content="Code analysis complete",
timestamp="2023-01-01T00:00:30Z", timestamp="2023-01-01T00:00:30Z",
follow_up_question="Would you like me to check error handling?",
) )
], ],
initial_context={"prompt": "Analyze this code"}, initial_context={"prompt": "Analyze this code"},
@@ -266,9 +260,7 @@ class TestConversationFlow:
success = add_turn(thread_id, "user", "Yes, check error handling") success = add_turn(thread_id, "user", "Yes, check error handling")
assert success is True assert success is True
success = add_turn( success = add_turn(thread_id, "assistant", "Error handling reviewed")
thread_id, "assistant", "Error handling reviewed", follow_up_question="Should I examine the test coverage?"
)
assert success is True assert success is True
# REQUEST 3-5: Continue conversation (simulating independent cycles) # REQUEST 3-5: Continue conversation (simulating independent cycles)
@@ -283,14 +275,12 @@ class TestConversationFlow:
role="assistant", role="assistant",
content="Code analysis complete", content="Code analysis complete",
timestamp="2023-01-01T00:00:30Z", timestamp="2023-01-01T00:00:30Z",
follow_up_question="Would you like me to check error handling?",
), ),
ConversationTurn(role="user", content="Yes, check error handling", timestamp="2023-01-01T00:01:30Z"), ConversationTurn(role="user", content="Yes, check error handling", timestamp="2023-01-01T00:01:30Z"),
ConversationTurn( ConversationTurn(
role="assistant", role="assistant",
content="Error handling reviewed", content="Error handling reviewed",
timestamp="2023-01-01T00:02:30Z", timestamp="2023-01-01T00:02:30Z",
follow_up_question="Should I examine the test coverage?",
), ),
], ],
initial_context={"prompt": "Analyze this code"}, initial_context={"prompt": "Analyze this code"},
@@ -385,18 +375,20 @@ class TestConversationFlow:
# Test early conversation (should allow follow-ups) # Test early conversation (should allow follow-ups)
early_instructions = get_follow_up_instructions(0, max_turns) early_instructions = get_follow_up_instructions(0, max_turns)
assert "CONVERSATION THREADING" in early_instructions assert "CONVERSATION CONTINUATION" in early_instructions
assert f"({max_turns - 1} exchanges remaining)" in early_instructions assert f"({max_turns - 1} exchanges remaining)" in early_instructions
assert "Feel free to ask clarifying questions" in early_instructions
# Test mid conversation # Test mid conversation
mid_instructions = get_follow_up_instructions(2, max_turns) mid_instructions = get_follow_up_instructions(2, max_turns)
assert "CONVERSATION THREADING" in mid_instructions assert "CONVERSATION CONTINUATION" in mid_instructions
assert f"({max_turns - 3} exchanges remaining)" in mid_instructions assert f"({max_turns - 3} exchanges remaining)" in mid_instructions
assert "Feel free to ask clarifying questions" in mid_instructions
# Test approaching limit (should stop follow-ups) # Test approaching limit (should stop follow-ups)
limit_instructions = get_follow_up_instructions(max_turns - 1, max_turns) limit_instructions = get_follow_up_instructions(max_turns - 1, max_turns)
assert "Do NOT include any follow-up questions" in limit_instructions assert "Do NOT include any follow-up questions" in limit_instructions
assert "FOLLOW-UP CONVERSATIONS" not in limit_instructions assert "final exchange" in limit_instructions
# Test at limit # Test at limit
at_limit_instructions = get_follow_up_instructions(max_turns, max_turns) at_limit_instructions = get_follow_up_instructions(max_turns, max_turns)
@@ -492,12 +484,11 @@ class TestConversationFlow:
) )
mock_client.get.return_value = initial_context.model_dump_json() mock_client.get.return_value = initial_context.model_dump_json()
# Add Gemini's response with follow-up # Add Gemini's response
success = add_turn( success = add_turn(
thread_id, thread_id,
"assistant", "assistant",
"I've analyzed your codebase structure.", "I've analyzed your codebase structure.",
follow_up_question="Would you like me to examine the test coverage?",
files=["/project/src/main.py", "/project/src/utils.py"], files=["/project/src/main.py", "/project/src/utils.py"],
tool_name="analyze", tool_name="analyze",
) )
@@ -514,7 +505,6 @@ class TestConversationFlow:
role="assistant", role="assistant",
content="I've analyzed your codebase structure.", content="I've analyzed your codebase structure.",
timestamp="2023-01-01T00:00:30Z", timestamp="2023-01-01T00:00:30Z",
follow_up_question="Would you like me to examine the test coverage?",
files=["/project/src/main.py", "/project/src/utils.py"], files=["/project/src/main.py", "/project/src/utils.py"],
tool_name="analyze", tool_name="analyze",
) )
@@ -540,7 +530,6 @@ class TestConversationFlow:
role="assistant", role="assistant",
content="I've analyzed your codebase structure.", content="I've analyzed your codebase structure.",
timestamp="2023-01-01T00:00:30Z", timestamp="2023-01-01T00:00:30Z",
follow_up_question="Would you like me to examine the test coverage?",
files=["/project/src/main.py", "/project/src/utils.py"], files=["/project/src/main.py", "/project/src/utils.py"],
tool_name="analyze", tool_name="analyze",
), ),
@@ -575,7 +564,6 @@ class TestConversationFlow:
role="assistant", role="assistant",
content="I've analyzed your codebase structure.", content="I've analyzed your codebase structure.",
timestamp="2023-01-01T00:00:30Z", timestamp="2023-01-01T00:00:30Z",
follow_up_question="Would you like me to examine the test coverage?",
files=["/project/src/main.py", "/project/src/utils.py"], files=["/project/src/main.py", "/project/src/utils.py"],
tool_name="analyze", tool_name="analyze",
), ),
@@ -604,19 +592,18 @@ class TestConversationFlow:
assert "--- Turn 3 (Gemini using analyze) ---" in history assert "--- Turn 3 (Gemini using analyze) ---" in history
# Verify all files are preserved in chronological order # Verify all files are preserved in chronological order
turn_1_files = "📁 Files used in this turn: /project/src/main.py, /project/src/utils.py" turn_1_files = "Files used in this turn: /project/src/main.py, /project/src/utils.py"
turn_2_files = "📁 Files used in this turn: /project/tests/, /project/test_main.py" turn_2_files = "Files used in this turn: /project/tests/, /project/test_main.py"
turn_3_files = "📁 Files used in this turn: /project/tests/test_utils.py, /project/coverage.html" turn_3_files = "Files used in this turn: /project/tests/test_utils.py, /project/coverage.html"
assert turn_1_files in history assert turn_1_files in history
assert turn_2_files in history assert turn_2_files in history
assert turn_3_files in history assert turn_3_files in history
# Verify content and follow-ups # Verify content
assert "I've analyzed your codebase structure." in history assert "I've analyzed your codebase structure." in history
assert "Yes, check the test coverage" in history assert "Yes, check the test coverage" in history
assert "Test coverage analysis complete. Coverage is 85%." in history assert "Test coverage analysis complete. Coverage is 85%." in history
assert "[Gemini's Follow-up: Would you like me to examine the test coverage?]" in history
# Verify chronological ordering (turn 1 appears before turn 2, etc.) # Verify chronological ordering (turn 1 appears before turn 2, etc.)
turn_1_pos = history.find("--- Turn 1 (Gemini using analyze) ---") turn_1_pos = history.find("--- Turn 1 (Gemini using analyze) ---")
@@ -625,56 +612,6 @@ class TestConversationFlow:
assert turn_1_pos < turn_2_pos < turn_3_pos assert turn_1_pos < turn_2_pos < turn_3_pos
@patch("utils.conversation_memory.get_redis_client")
def test_follow_up_question_parsing_cycle(self, mock_redis):
"""Test follow-up question persistence across request cycles"""
mock_client = Mock()
mock_redis.return_value = mock_client
thread_id = "12345678-1234-1234-1234-123456789012"
# First cycle: Assistant generates follow-up
context = ThreadContext(
thread_id=thread_id,
created_at="2023-01-01T00:00:00Z",
last_updated_at="2023-01-01T00:00:00Z",
tool_name="debug",
turns=[],
initial_context={"prompt": "Debug this error"},
)
mock_client.get.return_value = context.model_dump_json()
success = add_turn(
thread_id,
"assistant",
"Found potential issue in authentication",
follow_up_question="Should I examine the authentication middleware?",
)
assert success is True
# Second cycle: Retrieve conversation history
context_with_followup = ThreadContext(
thread_id=thread_id,
created_at="2023-01-01T00:00:00Z",
last_updated_at="2023-01-01T00:01:00Z",
tool_name="debug",
turns=[
ConversationTurn(
role="assistant",
content="Found potential issue in authentication",
timestamp="2023-01-01T00:00:30Z",
follow_up_question="Should I examine the authentication middleware?",
)
],
initial_context={"prompt": "Debug this error"},
)
mock_client.get.return_value = context_with_followup.model_dump_json()
# Build history to verify follow-up is preserved
history, tokens = build_conversation_history(context_with_followup)
assert "Found potential issue in authentication" in history
assert "[Gemini's Follow-up: Should I examine the authentication middleware?]" in history
@patch("utils.conversation_memory.get_redis_client") @patch("utils.conversation_memory.get_redis_client")
def test_stateless_request_isolation(self, mock_redis): def test_stateless_request_isolation(self, mock_redis):
"""Test that each request cycle is independent but shares context via Redis""" """Test that each request cycle is independent but shares context via Redis"""
@@ -695,9 +632,7 @@ class TestConversationFlow:
) )
mock_client.get.return_value = initial_context.model_dump_json() mock_client.get.return_value = initial_context.model_dump_json()
success = add_turn( success = add_turn(thread_id, "assistant", "Architecture analysis")
thread_id, "assistant", "Architecture analysis", follow_up_question="Want to explore scalability?"
)
assert success is True assert success is True
# Process 2: Different "request cycle" accesses same thread # Process 2: Different "request cycle" accesses same thread
@@ -711,7 +646,6 @@ class TestConversationFlow:
role="assistant", role="assistant",
content="Architecture analysis", content="Architecture analysis",
timestamp="2023-01-01T00:00:30Z", timestamp="2023-01-01T00:00:30Z",
follow_up_question="Want to explore scalability?",
) )
], ],
initial_context={"prompt": "Think about architecture"}, initial_context={"prompt": "Think about architecture"},
@@ -722,7 +656,6 @@ class TestConversationFlow:
retrieved_context = get_thread(thread_id) retrieved_context = get_thread(thread_id)
assert retrieved_context is not None assert retrieved_context is not None
assert len(retrieved_context.turns) == 1 assert len(retrieved_context.turns) == 1
assert retrieved_context.turns[0].follow_up_question == "Want to explore scalability?"
def test_token_limit_optimization_in_conversation_history(self): def test_token_limit_optimization_in_conversation_history(self):
"""Test that build_conversation_history efficiently handles token limits""" """Test that build_conversation_history efficiently handles token limits"""
@@ -766,7 +699,7 @@ class TestConversationFlow:
history, tokens = build_conversation_history(context, model_context=None) history, tokens = build_conversation_history(context, model_context=None)
# Verify the history was built successfully # Verify the history was built successfully
assert "=== CONVERSATION HISTORY ===" in history assert "=== CONVERSATION HISTORY" in history
assert "=== FILES REFERENCED IN THIS CONVERSATION ===" in history assert "=== FILES REFERENCED IN THIS CONVERSATION ===" in history
# The small file should be included, but large file might be truncated # The small file should be included, but large file might be truncated

25
tests/test_cross_tool_continuation.py
View File

@@ -93,28 +93,23 @@ class TestCrossToolContinuation:
self.review_tool = MockReviewTool() self.review_tool = MockReviewTool()
@patch("utils.conversation_memory.get_redis_client") @patch("utils.conversation_memory.get_redis_client")
@patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
async def test_continuation_id_works_across_different_tools(self, mock_redis): async def test_continuation_id_works_across_different_tools(self, mock_redis):
"""Test that a continuation_id from one tool can be used with another tool""" """Test that a continuation_id from one tool can be used with another tool"""
mock_client = Mock() mock_client = Mock()
mock_redis.return_value = mock_client mock_redis.return_value = mock_client
# Step 1: Analysis tool creates a conversation with follow-up # Step 1: Analysis tool creates a conversation with continuation offer
with patch.object(self.analysis_tool, "get_model_provider") as mock_get_provider: with patch.object(self.analysis_tool, "get_model_provider") as mock_get_provider:
mock_provider = create_mock_provider() mock_provider = create_mock_provider()
mock_provider.get_provider_type.return_value = Mock(value="google") mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = False mock_provider.supports_thinking_mode.return_value = False
# Include follow-up JSON in the content # Simple content without JSON follow-up
content_with_followup = """Found potential security issues in authentication logic. content = """Found potential security issues in authentication logic.
```json I'd be happy to review these security findings in detail if that would be helpful."""
{
"follow_up_question": "Would you like me to review these security findings in detail?",
"suggested_params": {"findings": "Authentication bypass vulnerability detected"},
"ui_hint": "Security review recommended"
}
```"""
mock_provider.generate_content.return_value = Mock( mock_provider.generate_content.return_value = Mock(
content=content_with_followup, content=content,
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash-exp", model_name="gemini-2.0-flash-exp",
metadata={"finish_reason": "STOP"}, metadata={"finish_reason": "STOP"},
@@ -126,8 +121,8 @@ class TestCrossToolContinuation:
response = await self.analysis_tool.execute(arguments) response = await self.analysis_tool.execute(arguments)
response_data = json.loads(response[0].text) response_data = json.loads(response[0].text)
assert response_data["status"] == "requires_continuation" assert response_data["status"] == "continuation_available"
continuation_id = response_data["follow_up_request"]["continuation_id"] continuation_id = response_data["continuation_offer"]["continuation_id"]
# Step 2: Mock the existing thread context for the review tool # Step 2: Mock the existing thread context for the review tool
# The thread was created by analysis_tool but will be continued by review_tool # The thread was created by analysis_tool but will be continued by review_tool
@@ -139,10 +134,9 @@ class TestCrossToolContinuation:
turns=[ turns=[
ConversationTurn( ConversationTurn(
role="assistant", role="assistant",
content="Found potential security issues in authentication logic.", content="Found potential security issues in authentication logic.\n\nI'd be happy to review these security findings in detail if that would be helpful.",
timestamp="2023-01-01T00:00:30Z", timestamp="2023-01-01T00:00:30Z",
tool_name="test_analysis", # Original tool tool_name="test_analysis", # Original tool
follow_up_question="Would you like me to review these security findings in detail?",
) )
], ],
initial_context={"code": "function authenticate(user) { return true; }"}, initial_context={"code": "function authenticate(user) { return true; }"},
@@ -250,6 +244,7 @@ class TestCrossToolContinuation:
@patch("utils.conversation_memory.get_redis_client") @patch("utils.conversation_memory.get_redis_client")
@patch("utils.conversation_memory.get_thread") @patch("utils.conversation_memory.get_thread")
@patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
async def test_cross_tool_conversation_with_files_context(self, mock_get_thread, mock_redis): async def test_cross_tool_conversation_with_files_context(self, mock_get_thread, mock_redis):
"""Test that file context is preserved across tool switches""" """Test that file context is preserved across tool switches"""
mock_client = Mock() mock_client = Mock()

4
tests/test_prompt_regression.py
View File

@@ -109,7 +109,7 @@ class TestPromptRegression:
assert len(result) == 1 assert len(result) == 1
output = json.loads(result[0].text) output = json.loads(result[0].text)
assert output["status"] == "success" assert output["status"] == "success"
assert "Extended Analysis by Gemini" in output["content"] assert "Critical Evaluation Required" in output["content"]
assert "deeper analysis" in output["content"] assert "deeper analysis" in output["content"]
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -203,7 +203,7 @@ class TestPromptRegression:
assert len(result) == 1 assert len(result) == 1
output = json.loads(result[0].text) output = json.loads(result[0].text)
assert output["status"] == "success" assert output["status"] == "success"
assert "Debug Analysis" in output["content"] assert "Next Steps:" in output["content"]
assert "Root cause" in output["content"] assert "Root cause" in output["content"]
@pytest.mark.asyncio @pytest.mark.asyncio

18
tests/test_thinking_modes.py
View File

@@ -59,7 +59,7 @@ class TestThinkingModes:
) )
# Verify create_model was called with correct thinking_mode # Verify create_model was called with correct thinking_mode
mock_get_provider.assert_called_once() assert mock_get_provider.called
# Verify generate_content was called with thinking_mode # Verify generate_content was called with thinking_mode
mock_provider.generate_content.assert_called_once() mock_provider.generate_content.assert_called_once()
call_kwargs = mock_provider.generate_content.call_args[1] call_kwargs = mock_provider.generate_content.call_args[1]
@@ -72,7 +72,7 @@ class TestThinkingModes:
response_data = json.loads(result[0].text) response_data = json.loads(result[0].text)
assert response_data["status"] == "success" assert response_data["status"] == "success"
assert response_data["content"].startswith("Analysis:") assert "Minimal thinking response" in response_data["content"] or "Analysis:" in response_data["content"]
@pytest.mark.asyncio @pytest.mark.asyncio
@patch("tools.base.BaseTool.get_model_provider") @patch("tools.base.BaseTool.get_model_provider")
@@ -96,7 +96,7 @@ class TestThinkingModes:
) )
# Verify create_model was called with correct thinking_mode # Verify create_model was called with correct thinking_mode
mock_get_provider.assert_called_once() assert mock_get_provider.called
# Verify generate_content was called with thinking_mode # Verify generate_content was called with thinking_mode
mock_provider.generate_content.assert_called_once() mock_provider.generate_content.assert_called_once()
call_kwargs = mock_provider.generate_content.call_args[1] call_kwargs = mock_provider.generate_content.call_args[1]
@@ -104,7 +104,7 @@ class TestThinkingModes:
not mock_provider.supports_thinking_mode.return_value and call_kwargs.get("thinking_mode") is None not mock_provider.supports_thinking_mode.return_value and call_kwargs.get("thinking_mode") is None
) )
assert "Code Review" in result[0].text assert "Low thinking response" in result[0].text or "Code Review" in result[0].text
@pytest.mark.asyncio @pytest.mark.asyncio
@patch("tools.base.BaseTool.get_model_provider") @patch("tools.base.BaseTool.get_model_provider")
@@ -127,7 +127,7 @@ class TestThinkingModes:
) )
# Verify create_model was called with default thinking_mode # Verify create_model was called with default thinking_mode
mock_get_provider.assert_called_once() assert mock_get_provider.called
# Verify generate_content was called with thinking_mode # Verify generate_content was called with thinking_mode
mock_provider.generate_content.assert_called_once() mock_provider.generate_content.assert_called_once()
call_kwargs = mock_provider.generate_content.call_args[1] call_kwargs = mock_provider.generate_content.call_args[1]
@@ -135,7 +135,7 @@ class TestThinkingModes:
not mock_provider.supports_thinking_mode.return_value and call_kwargs.get("thinking_mode") is None not mock_provider.supports_thinking_mode.return_value and call_kwargs.get("thinking_mode") is None
) )
assert "Debug Analysis" in result[0].text assert "Medium thinking response" in result[0].text or "Debug Analysis" in result[0].text
@pytest.mark.asyncio @pytest.mark.asyncio
@patch("tools.base.BaseTool.get_model_provider") @patch("tools.base.BaseTool.get_model_provider")
@@ -159,7 +159,7 @@ class TestThinkingModes:
) )
# Verify create_model was called with correct thinking_mode # Verify create_model was called with correct thinking_mode
mock_get_provider.assert_called_once() assert mock_get_provider.called
# Verify generate_content was called with thinking_mode # Verify generate_content was called with thinking_mode
mock_provider.generate_content.assert_called_once() mock_provider.generate_content.assert_called_once()
call_kwargs = mock_provider.generate_content.call_args[1] call_kwargs = mock_provider.generate_content.call_args[1]
@@ -188,7 +188,7 @@ class TestThinkingModes:
) )
# Verify create_model was called with default thinking_mode # Verify create_model was called with default thinking_mode
mock_get_provider.assert_called_once() assert mock_get_provider.called
# Verify generate_content was called with thinking_mode # Verify generate_content was called with thinking_mode
mock_provider.generate_content.assert_called_once() mock_provider.generate_content.assert_called_once()
call_kwargs = mock_provider.generate_content.call_args[1] call_kwargs = mock_provider.generate_content.call_args[1]
@@ -196,7 +196,7 @@ class TestThinkingModes:
not mock_provider.supports_thinking_mode.return_value and call_kwargs.get("thinking_mode") is None not mock_provider.supports_thinking_mode.return_value and call_kwargs.get("thinking_mode") is None
) )
assert "Extended Analysis by Gemini" in result[0].text assert "Max thinking response" in result[0].text or "Extended Analysis by Gemini" in result[0].text
def test_thinking_budget_mapping(self): def test_thinking_budget_mapping(self):
"""Test that thinking modes map to correct budget values""" """Test that thinking modes map to correct budget values"""

12
tests/test_tools.py
View File

@@ -53,7 +53,7 @@ class TestThinkDeepTool:
# Parse the JSON response # Parse the JSON response
output = json.loads(result[0].text) output = json.loads(result[0].text)
assert output["status"] == "success" assert output["status"] == "success"
assert "Extended Analysis by Gemini" in output["content"] assert "Critical Evaluation Required" in output["content"]
assert "Extended analysis" in output["content"] assert "Extended analysis" in output["content"]
@@ -102,8 +102,8 @@ class TestCodeReviewTool:
) )
assert len(result) == 1 assert len(result) == 1
assert "Code Review (SECURITY)" in result[0].text assert "Security issues found" in result[0].text
assert "Focus: authentication" in result[0].text assert "Claude's Next Steps:" in result[0].text
assert "Security issues found" in result[0].text assert "Security issues found" in result[0].text
@@ -146,7 +146,7 @@ class TestDebugIssueTool:
) )
assert len(result) == 1 assert len(result) == 1
assert "Debug Analysis" in result[0].text assert "Next Steps:" in result[0].text
assert "Root cause: race condition" in result[0].text assert "Root cause: race condition" in result[0].text
@@ -195,8 +195,8 @@ class TestAnalyzeTool:
) )
assert len(result) == 1 assert len(result) == 1
assert "ARCHITECTURE Analysis" in result[0].text assert "Architecture analysis" in result[0].text
assert "Analyzed 1 file(s)" in result[0].text assert "Next Steps:" in result[0].text
assert "Architecture analysis" in result[0].text assert "Architecture analysis" in result[0].text

159
tools/base.py
View File

@@ -16,14 +16,13 @@ Key responsibilities:
import json import json
import logging import logging
import os import os
import re
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from typing import Any, Literal, Optional from typing import Any, Literal, Optional
from mcp.types import TextContent from mcp.types import TextContent
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from config import DEFAULT_MODEL, MAX_CONTEXT_TOKENS, MCP_PROMPT_SIZE_LIMIT from config import MAX_CONTEXT_TOKENS, MCP_PROMPT_SIZE_LIMIT
from providers import ModelProvider, ModelProviderRegistry from providers import ModelProvider, ModelProviderRegistry
from utils import check_token_limit from utils import check_token_limit
from utils.conversation_memory import ( from utils.conversation_memory import (
@@ -35,7 +34,7 @@ from utils.conversation_memory import (
) )
from utils.file_utils import read_file_content, read_files, translate_path_for_environment from utils.file_utils import read_file_content, read_files, translate_path_for_environment
from .models import ClarificationRequest, ContinuationOffer, FollowUpRequest, ToolOutput from .models import ClarificationRequest, ContinuationOffer, ToolOutput
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -363,6 +362,8 @@ class BaseTool(ABC):
if not model_context: if not model_context:
# Manual calculation as fallback # Manual calculation as fallback
from config import DEFAULT_MODEL
model_name = getattr(self, "_current_model_name", None) or DEFAULT_MODEL model_name = getattr(self, "_current_model_name", None) or DEFAULT_MODEL
try: try:
provider = self.get_model_provider(model_name) provider = self.get_model_provider(model_name)
@@ -739,6 +740,8 @@ If any of these would strengthen your analysis, specify what Claude should searc
# Extract model configuration from request or use defaults # Extract model configuration from request or use defaults
model_name = getattr(request, "model", None) model_name = getattr(request, "model", None)
if not model_name: if not model_name:
from config import DEFAULT_MODEL
model_name = DEFAULT_MODEL model_name = DEFAULT_MODEL
# In auto mode, model parameter is required # In auto mode, model parameter is required
@@ -859,29 +862,21 @@ If any of these would strengthen your analysis, specify what Claude should searc
def _parse_response(self, raw_text: str, request, model_info: Optional[dict] = None) -> ToolOutput: def _parse_response(self, raw_text: str, request, model_info: Optional[dict] = None) -> ToolOutput:
""" """
Parse the raw response and determine if it's a clarification request or follow-up. Parse the raw response and check for clarification requests.
Some tools may return JSON indicating they need more information or want to This method formats the response and always offers a continuation opportunity
continue the conversation. This method detects such responses and formats them. unless max conversation turns have been reached.
Args: Args:
raw_text: The raw text response from the model raw_text: The raw text response from the model
request: The original request for context request: The original request for context
model_info: Optional dict with model metadata
Returns: Returns:
ToolOutput: Standardized output object ToolOutput: Standardized output object
""" """
# Check for follow-up questions in JSON blocks at the end of the response
follow_up_question = self._extract_follow_up_question(raw_text)
logger = logging.getLogger(f"tools.{self.name}") logger = logging.getLogger(f"tools.{self.name}")
if follow_up_question:
logger.debug(
f"Found follow-up question in {self.name} response: {follow_up_question.get('follow_up_question', 'N/A')}"
)
else:
logger.debug(f"No follow-up question found in {self.name} response")
try: try:
# Try to parse as JSON to check for clarification requests # Try to parse as JSON to check for clarification requests
potential_json = json.loads(raw_text.strip()) potential_json = json.loads(raw_text.strip())
@@ -905,11 +900,7 @@ If any of these would strengthen your analysis, specify what Claude should searc
# Normal text response - format using tool-specific formatting # Normal text response - format using tool-specific formatting
formatted_content = self.format_response(raw_text, request, model_info) formatted_content = self.format_response(raw_text, request, model_info)
# If we found a follow-up question, prepare the threading response # Always check if we should offer Claude a continuation opportunity
if follow_up_question:
return self._create_follow_up_response(formatted_content, follow_up_question, request, model_info)
# Check if we should offer Claude a continuation opportunity
continuation_offer = self._check_continuation_opportunity(request) continuation_offer = self._check_continuation_opportunity(request)
if continuation_offer: if continuation_offer:
@@ -918,7 +909,7 @@ If any of these would strengthen your analysis, specify what Claude should searc
) )
return self._create_continuation_offer_response(formatted_content, continuation_offer, request, model_info) return self._create_continuation_offer_response(formatted_content, continuation_offer, request, model_info)
else: else:
logger.debug(f"No continuation offer created for {self.name}") logger.debug(f"No continuation offer created for {self.name} - max turns reached")
# If this is a threaded conversation (has continuation_id), save the response # If this is a threaded conversation (has continuation_id), save the response
continuation_id = getattr(request, "continuation_id", None) continuation_id = getattr(request, "continuation_id", None)
@@ -963,126 +954,6 @@ If any of these would strengthen your analysis, specify what Claude should searc
metadata={"tool_name": self.name}, metadata={"tool_name": self.name},
) )
def _extract_follow_up_question(self, text: str) -> Optional[dict]:
"""
Extract follow-up question from JSON blocks in the response.
Looks for JSON blocks containing follow_up_question at the end of responses.
Args:
text: The response text to parse
Returns:
Dict with follow-up data if found, None otherwise
"""
# Look for JSON blocks that contain follow_up_question
# Pattern handles optional leading whitespace and indentation
json_pattern = r'```json\s*\n\s*(\{.*?"follow_up_question".*?\})\s*\n\s*```'
matches = re.findall(json_pattern, text, re.DOTALL)
if not matches:
return None
# Take the last match (most recent follow-up)
try:
# Clean up the JSON string - remove excess whitespace and normalize
json_str = re.sub(r"\n\s+", "\n", matches[-1]).strip()
follow_up_data = json.loads(json_str)
if "follow_up_question" in follow_up_data:
return follow_up_data
except (json.JSONDecodeError, ValueError):
pass
return None
def _create_follow_up_response(
self, content: str, follow_up_data: dict, request, model_info: Optional[dict] = None
) -> ToolOutput:
"""
Create a response with follow-up question for conversation threading.
Args:
content: The main response content
follow_up_data: Dict containing follow_up_question and optional suggested_params
request: Original request for context
Returns:
ToolOutput configured for conversation continuation
"""
# Always create a new thread (with parent linkage if continuation)
continuation_id = getattr(request, "continuation_id", None)
request_files = getattr(request, "files", []) or []
try:
# Create new thread with parent linkage if continuing
thread_id = create_thread(
tool_name=self.name,
initial_request=request.model_dump() if hasattr(request, "model_dump") else {},
parent_thread_id=continuation_id, # Link to parent thread if continuing
)
# Add the assistant's response with follow-up
# Extract model metadata
model_provider = None
model_name = None
model_metadata = None
if model_info:
provider = model_info.get("provider")
if provider:
model_provider = provider.get_provider_type().value
model_name = model_info.get("model_name")
model_response = model_info.get("model_response")
if model_response:
model_metadata = {"usage": model_response.usage, "metadata": model_response.metadata}
add_turn(
thread_id, # Add to the new thread
"assistant",
content,
follow_up_question=follow_up_data.get("follow_up_question"),
files=request_files,
tool_name=self.name,
model_provider=model_provider,
model_name=model_name,
model_metadata=model_metadata,
)
except Exception as e:
# Threading failed, return normal response
logger = logging.getLogger(f"tools.{self.name}")
logger.warning(f"Follow-up threading failed in {self.name}: {str(e)}")
return ToolOutput(
status="success",
content=content,
content_type="markdown",
metadata={"tool_name": self.name, "follow_up_error": str(e)},
)
# Create follow-up request
follow_up_request = FollowUpRequest(
continuation_id=thread_id,
question_to_user=follow_up_data["follow_up_question"],
suggested_tool_params=follow_up_data.get("suggested_params"),
ui_hint=follow_up_data.get("ui_hint"),
)
# Strip the JSON block from the content since it's now in the follow_up_request
clean_content = self._remove_follow_up_json(content)
return ToolOutput(
status="requires_continuation",
content=clean_content,
content_type="markdown",
follow_up_request=follow_up_request,
metadata={"tool_name": self.name, "thread_id": thread_id},
)
def _remove_follow_up_json(self, text: str) -> str:
"""Remove follow-up JSON blocks from the response text"""
# Remove JSON blocks containing follow_up_question
pattern = r'```json\s*\n\s*\{.*?"follow_up_question".*?\}\s*\n\s*```'
return re.sub(pattern, "", text, flags=re.DOTALL).strip()
def _check_continuation_opportunity(self, request) -> Optional[dict]: def _check_continuation_opportunity(self, request) -> Optional[dict]:
""" """
Check if we should offer Claude a continuation opportunity. Check if we should offer Claude a continuation opportunity.
@@ -1186,13 +1057,13 @@ If any of these would strengthen your analysis, specify what Claude should searc
continuation_offer = ContinuationOffer( continuation_offer = ContinuationOffer(
continuation_id=thread_id, continuation_id=thread_id,
message_to_user=( message_to_user=(
f"If you'd like to continue this analysis or need further details, " f"If you'd like to continue this discussion or need to provide me with further details or context, "
f"you can use the continuation_id '{thread_id}' in your next {self.name} tool call. " f"you can use the continuation_id '{thread_id}' with any tool and any model. "
f"You have {remaining_turns} more exchange(s) available in this conversation thread." f"You have {remaining_turns} more exchange(s) available in this conversation thread."
), ),
suggested_tool_params={ suggested_tool_params={
"continuation_id": thread_id, "continuation_id": thread_id,
"prompt": "[Your follow-up question or request for additional analysis]", "prompt": "[Your follow-up question, additional context, or further details]",
}, },
remaining_turns=remaining_turns, remaining_turns=remaining_turns,
) )

19
tools/models.py
View File

@@ -7,21 +7,6 @@ from typing import Any, Literal, Optional
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
class FollowUpRequest(BaseModel):
"""Request for follow-up conversation turn"""
continuation_id: str = Field(
..., description="Thread continuation ID for multi-turn conversations across different tools"
)
question_to_user: str = Field(..., description="Follow-up question to ask Claude")
suggested_tool_params: Optional[dict[str, Any]] = Field(
None, description="Suggested parameters for the next tool call"
)
ui_hint: Optional[str] = Field(
None, description="UI hint for Claude (e.g., 'text_input', 'file_select', 'multi_choice')"
)
class ContinuationOffer(BaseModel): class ContinuationOffer(BaseModel):
"""Offer for Claude to continue conversation when Gemini doesn't ask follow-up""" """Offer for Claude to continue conversation when Gemini doesn't ask follow-up"""
@@ -43,15 +28,11 @@ class ToolOutput(BaseModel):
"error", "error",
"requires_clarification", "requires_clarification",
"requires_file_prompt", "requires_file_prompt",
"requires_continuation",
"continuation_available", "continuation_available",
] = "success" ] = "success"
content: Optional[str] = Field(None, description="The main content/response from the tool") content: Optional[str] = Field(None, description="The main content/response from the tool")
content_type: Literal["text", "markdown", "json"] = "text" content_type: Literal["text", "markdown", "json"] = "text"
metadata: Optional[dict[str, Any]] = Field(default_factory=dict) metadata: Optional[dict[str, Any]] = Field(default_factory=dict)
follow_up_request: Optional[FollowUpRequest] = Field(
None, description="Optional follow-up request for continued conversation"
)
continuation_offer: Optional[ContinuationOffer] = Field( continuation_offer: Optional[ContinuationOffer] = Field(
None, description="Optional offer for Claude to continue conversation" None, description="Optional offer for Claude to continue conversation"
) )

21
utils/conversation_memory.py
View File

@@ -71,7 +71,6 @@ class ConversationTurn(BaseModel):
role: "user" (Claude) or "assistant" (Gemini/O3/etc) role: "user" (Claude) or "assistant" (Gemini/O3/etc)
content: The actual message content/response content: The actual message content/response
timestamp: ISO timestamp when this turn was created timestamp: ISO timestamp when this turn was created
follow_up_question: Optional follow-up question from assistant to Claude
files: List of file paths referenced in this specific turn files: List of file paths referenced in this specific turn
tool_name: Which tool generated this turn (for cross-tool tracking) tool_name: Which tool generated this turn (for cross-tool tracking)
model_provider: Provider used (e.g., "google", "openai") model_provider: Provider used (e.g., "google", "openai")
@@ -82,7 +81,6 @@ class ConversationTurn(BaseModel):
role: str # "user" or "assistant" role: str # "user" or "assistant"
content: str content: str
timestamp: str timestamp: str
follow_up_question: Optional[str] = None
files: Optional[list[str]] = None # Files referenced in this turn files: Optional[list[str]] = None # Files referenced in this turn
tool_name: Optional[str] = None # Tool used for this turn tool_name: Optional[str] = None # Tool used for this turn
model_provider: Optional[str] = None # Model provider (google, openai, etc) model_provider: Optional[str] = None # Model provider (google, openai, etc)
@@ -231,7 +229,6 @@ def add_turn(
thread_id: str, thread_id: str,
role: str, role: str,
content: str, content: str,
follow_up_question: Optional[str] = None,
files: Optional[list[str]] = None, files: Optional[list[str]] = None,
tool_name: Optional[str] = None, tool_name: Optional[str] = None,
model_provider: Optional[str] = None, model_provider: Optional[str] = None,
@@ -249,7 +246,6 @@ def add_turn(
thread_id: UUID of the conversation thread thread_id: UUID of the conversation thread
role: "user" (Claude) or "assistant" (Gemini/O3/etc) role: "user" (Claude) or "assistant" (Gemini/O3/etc)
content: The actual message/response content content: The actual message/response content
follow_up_question: Optional follow-up question from assistant
files: Optional list of files referenced in this turn files: Optional list of files referenced in this turn
tool_name: Name of the tool adding this turn (for attribution) tool_name: Name of the tool adding this turn (for attribution)
model_provider: Provider used (e.g., "google", "openai") model_provider: Provider used (e.g., "google", "openai")
@@ -287,7 +283,6 @@ def add_turn(
role=role, role=role,
content=content, content=content,
timestamp=datetime.now(timezone.utc).isoformat(), timestamp=datetime.now(timezone.utc).isoformat(),
follow_up_question=follow_up_question,
files=files, # Preserved for cross-tool file context files=files, # Preserved for cross-tool file context
tool_name=tool_name, # Track which tool generated this turn tool_name=tool_name, # Track which tool generated this turn
model_provider=model_provider, # Track model provider model_provider=model_provider, # Track model provider
@@ -473,10 +468,11 @@ def build_conversation_history(context: ThreadContext, model_context=None, read_
logger.debug(f"[HISTORY] Max history tokens: {max_history_tokens:,}") logger.debug(f"[HISTORY] Max history tokens: {max_history_tokens:,}")
history_parts = [ history_parts = [
"=== CONVERSATION HISTORY ===", "=== CONVERSATION HISTORY (CONTINUATION) ===",
f"Thread: {context.thread_id}", f"Thread: {context.thread_id}",
f"Tool: {context.tool_name}", # Original tool that started the conversation f"Tool: {context.tool_name}", # Original tool that started the conversation
f"Turn {total_turns}/{MAX_CONVERSATION_TURNS}", f"Turn {total_turns}/{MAX_CONVERSATION_TURNS}",
"You are continuing this conversation thread from where it left off.",
"", "",
] ]
@@ -622,10 +618,6 @@ def build_conversation_history(context: ThreadContext, model_context=None, read_
# Add the actual content # Add the actual content
turn_parts.append(turn.content) turn_parts.append(turn.content)
# Add follow-up question if present
if turn.follow_up_question:
turn_parts.append(f"\n[Gemini's Follow-up: {turn.follow_up_question}]")
# Calculate tokens for this turn # Calculate tokens for this turn
turn_content = "\n".join(turn_parts) turn_content = "\n".join(turn_parts)
turn_tokens = model_context.estimate_tokens(turn_content) turn_tokens = model_context.estimate_tokens(turn_content)
@@ -660,7 +652,14 @@ def build_conversation_history(context: ThreadContext, model_context=None, read_
history_parts.append(f"\n[Note: Showing {included_turns} most recent turns out of {total_turns} total]") history_parts.append(f"\n[Note: Showing {included_turns} most recent turns out of {total_turns} total]")
history_parts.extend( history_parts.extend(
["", "=== END CONVERSATION HISTORY ===", "", "Continue this conversation by building on the previous context."] [
"",
"=== END CONVERSATION HISTORY ===",
"",
"IMPORTANT: You are continuing an existing conversation thread. Build upon the previous exchanges shown above,",
"reference earlier points, and maintain consistency with what has been discussed.",
f"This is turn {len(all_turns) + 1} of the conversation - use the conversation history above to provide a coherent continuation.",
]
) )
# Calculate total tokens for the complete conversation history # Calculate total tokens for the complete conversation history