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
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 |
|------|-------------|----------|-------------|
| `minimal` | 128 tokens | Simple, straightforward tasks | Lowest cost |
@@ -540,17 +542,17 @@ To help choose the right tool for your needs:
**Examples by scenario:**
```
# Quick style check
"Use o3 to review formatting in utils.py with minimal thinking"
# Quick style check with o3
"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"
# Complex debugging
# Complex debugging, letting claude pick the best model
"Use zen to debug this race condition with max thinking mode"
# Architecture analysis
"Analyze the entire src/ directory architecture with high thinking using zen"
# Architecture analysis with Gemini 2.5 Pro
"Analyze the entire src/ directory architecture with high thinking using pro"
```
## Advanced Features

74
communication_simulator_test.py
View File

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

34
server.py
View File

@@ -310,26 +310,26 @@ final analysis and recommendations."""
remaining_turns = max_turns - current_turn_count - 1
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,
add this JSON block at the very end of your response:
Feel free to ask clarifying questions or suggest areas for deeper exploration naturally within your response.
If something needs clarification or you'd benefit from additional context, simply mention it conversationally.
```json
{{
"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"
}}
```
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:
Good follow-up opportunities:
- "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?"
For optional follow-ups: "Please continue this conversation using the continuation_id from this response if you'd like to explore this further."
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]:
@@ -459,7 +459,7 @@ async def reconstruct_thread_context(arguments: dict[str, Any]) -> dict[str, Any
try:
mcp_activity_logger = logging.getLogger("mcp_activity")
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:
pass

2
simulator_tests/test_basic_conversation.py
View File

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

44
simulator_tests/test_content_validation.py
View File

@@ -27,15 +27,32 @@ class ContentValidationTest(BaseSimulatorTest):
try:
# Check both main server and log monitor for comprehensive logs
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
result_server = subprocess.run(cmd_server, capture_output=True, text=True)
result_monitor = subprocess.run(cmd_monitor, capture_output=True, text=True)
# Combine logs from both containers
combined_logs = result_server.stdout + "\n" + result_monitor.stdout
# Get the internal log files which have more detailed logging
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
except Exception as e:
self.logger.error(f"Failed to get docker logs: {e}")
@@ -140,19 +157,24 @@ DATABASE_CONFIG = {
# Check for proper file 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
deduplication_logs = [
line
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
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
@@ -160,10 +182,10 @@ DATABASE_CONFIG = {
embedding_found = len(embedding_logs) > 0
(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" 📊 Deduplication evidence: {len(deduplication_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" Embedding logs found: {len(embedding_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" Validation file mentioned: {validation_file_mentioned}")
# Log sample evidence for debugging
if self.verbose and embedding_logs:
@@ -179,7 +201,7 @@ DATABASE_CONFIG = {
]
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
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:
"""Test conversation chain and threading functionality"""
try:
self.logger.info("🔗 Test: Conversation chain and threading validation")
self.logger.info("Test: Conversation chain and threading validation")
# Setup test files
self.setup_test_files()
@@ -108,7 +108,7 @@ class TestClass:
conversation_chains = {}
# === 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)
self.logger.info(" Step A1: Chat tool - start new conversation")
@@ -173,7 +173,7 @@ class TestClass:
conversation_chains["A3"] = continuation_id_a3
# === 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)
self.logger.info(" Step B1: Chat tool - start NEW independent conversation")
@@ -215,7 +215,7 @@ class TestClass:
conversation_chains["B2"] = continuation_id_b2
# === 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)
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
# === 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
logs = self.get_recent_server_logs()
@@ -334,7 +334,7 @@ class TestClass:
)
# === VALIDATION RESULTS ===
self.logger.info(" 📊 Thread Relationship Validation:")
self.logger.info(" Thread Relationship Validation:")
relationship_passed = 0
for desc, passed in expected_relationships:
status = "✅" if passed else "❌"
@@ -342,7 +342,7 @@ class TestClass:
if passed:
relationship_passed += 1
self.logger.info(" 📊 History Traversal Validation:")
self.logger.info(" History Traversal Validation:")
traversal_passed = 0
for desc, passed in traversal_validations:
status = "✅" if passed else "❌"
@@ -354,7 +354,7 @@ class TestClass:
total_relationship_checks = len(expected_relationships)
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" 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
traversal_success = True
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
self.logger.info(" 📊 Conversation Chain Structure:")
self.logger.info(" Conversation Chain Structure:")
self.logger.info(
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:
# Check both main server and log monitor for comprehensive logs
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_monitor = subprocess.run(cmd_monitor, capture_output=True, text=True)
# Combine logs from both containers
combined_logs = result_server.stdout + "\n" + result_monitor.stdout
# Get the internal log files which have more detailed logging
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
except Exception as 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"))
# Print comprehensive diagnostics
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" 📊 Conversation logs found: {len(conversation_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" 📊 Cross-tool activity logs: {len(cross_tool_logs)}")
self.logger.info(f" 📊 Auth file mentioned: {auth_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" Tools used: {len(tools_used)} ({', '.join(tools_used)})")
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" File embedding logs found: {len(embedding_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" Auth file mentioned: {auth_file_mentioned}")
self.logger.info(f" Config file mentioned: {config_file_mentioned}")
self.logger.info(f" Improved file mentioned: {improved_file_mentioned}")
if self.verbose:
self.logger.debug(" 📋 Sample tool activity logs:")
@@ -296,9 +313,9 @@ def secure_login(user, pwd):
passed_criteria = sum(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")
return True
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()
# 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 = ""
if monitor_result.returncode == 0:
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:
"""Run all model thinking configuration tests"""
self.logger.info(f"📝 Test: {self.test_description}")
self.logger.info(f" Test: {self.test_description}")
try:
# 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:
"""Test O3 model selection and usage"""
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
self.setup_test_files()
@@ -120,15 +120,15 @@ def multiply(x, y):
logs = self.get_recent_server_logs()
# 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)
openai_http_logs = [
line for line in logs.split("\n") if "HTTP Request: POST https://api.openai.com" in line
# Check for OpenAI model usage logs
openai_model_logs = [
line for line in logs.split("\n") if "Using model:" in line and "openai provider" in line
]
# Check for received responses from OpenAI
openai_response_logs = [line for line in logs.split("\n") if "Received response from openai API" in line]
# Check for successful OpenAI responses
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
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)
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
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
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 response logs: {len(openai_response_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" OpenAI API call logs: {len(openai_api_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" Chat calls to OpenAI: {len(chat_openai_logs)}")
self.logger.info(f" Codereview calls to OpenAI: {len(codereview_openai_logs)}")
# Log sample evidence for debugging
if self.verbose and openai_api_logs:
@@ -164,14 +164,14 @@ def multiply(x, y):
# Success criteria
success_criteria = [
("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),
("Chat tool used OpenAI", chat_calls_to_openai),
("Codereview tool used OpenAI", codereview_calls_to_openai),
]
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:
status = "✅" if passed else "❌"

39
simulator_tests/test_per_tool_deduplication.py
View File

@@ -32,13 +32,30 @@ class PerToolDeduplicationTest(BaseSimulatorTest):
try:
# Check both main server and log monitor for comprehensive logs
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_monitor = subprocess.run(cmd_monitor, capture_output=True, text=True)
# Combine logs from both containers
combined_logs = result_server.stdout + "\n" + result_monitor.stdout
# Get the internal log files which have more detailed logging
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
except Exception as e:
self.logger.error(f"Failed to get docker logs: {e}")
@@ -177,7 +194,7 @@ def subtract(a, b):
embedding_logs = [
line
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
@@ -190,11 +207,11 @@ def subtract(a, b):
new_file_mentioned = any("new_feature.py" in line for line in logs.split("\n"))
# Print diagnostic information
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" 📊 Continuation logs found: {len(continuation_logs)}")
self.logger.info(f" 📊 Dummy file mentioned: {dummy_file_mentioned}")
self.logger.info(f" 📊 New file mentioned: {new_file_mentioned}")
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" Continuation logs found: {len(continuation_logs)}")
self.logger.info(f" Dummy file mentioned: {dummy_file_mentioned}")
self.logger.info(f" New file mentioned: {new_file_mentioned}")
if self.verbose:
self.logger.debug(" 📋 Sample embedding logs:")
@@ -218,9 +235,9 @@ def subtract(a, b):
passed_criteria = sum(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")
return True
else:

2
simulator_tests/test_redis_validation.py
View File

@@ -76,7 +76,7 @@ class RedisValidationTest(BaseSimulatorTest):
return True
else:
# 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_data = {

24
simulator_tests/test_token_allocation_validation.py
View File

@@ -102,7 +102,7 @@ class TokenAllocationValidationTest(BaseSimulatorTest):
def run_test(self) -> bool:
"""Test token allocation and conversation history functionality"""
try:
self.logger.info("🔥 Test: Token allocation and conversation history validation")
self.logger.info(" Test: Token allocation and conversation history validation")
# Setup test files
self.setup_test_files()
@@ -282,7 +282,7 @@ if __name__ == "__main__":
step1_file_tokens = int(match.group(1))
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)
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
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"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
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"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))
# 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 2 - Conversation: {step2_conversation:,}, Remaining: {step2_remaining:,}")
self.logger.info(f" Step 3 - Conversation: {step3_conversation:,}, Remaining: {step3_remaining:,}")
# 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 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)")
@@ -492,7 +492,7 @@ if __name__ == "__main__":
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" File2 (calculator.py) mentioned: {file2_mentioned_step3}")
@@ -504,7 +504,7 @@ if __name__ == "__main__":
passed_criteria = sum(1 for _, passed in criteria if passed)
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:
status = "✅" if passed else "❌"
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()]
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" File embedding logs: {len(file_embedding_logs)}")
self.logger.info(f" Conversation history logs: {len(conversation_logs)}")
# Success criteria: At least 6 out of 8 validation criteria should pass
success = passed_criteria >= 6
# Success criteria: All validation criteria must pass
success = passed_criteria == total_criteria
if success:
self.logger.info(" ✅ Token allocation validation test PASSED")

322
tests/test_claude_continuation.py
View File

@@ -13,7 +13,6 @@ from pydantic import Field
from tests.mock_helpers import create_mock_provider
from tools.base import BaseTool, ToolRequest
from tools.models import ContinuationOffer, ToolOutput
from utils.conversation_memory import MAX_CONVERSATION_TURNS
@@ -59,58 +58,97 @@ class TestClaudeContinuationOffers:
self.tool = ClaudeContinuationTool()
@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"""
mock_client = Mock()
mock_redis.return_value = mock_client
# Test request without continuation_id (new conversation)
request = ContinuationRequest(prompt="Analyze this code")
# Check continuation opportunity
continuation_data = self.tool._check_continuation_opportunity(request)
assert continuation_data is not None
assert continuation_data["remaining_turns"] == MAX_CONVERSATION_TURNS - 1
assert continuation_data["tool_name"] == "test_continuation"
def test_existing_conversation_no_continuation_offer(self):
"""Test that existing threaded conversations don't offer continuation"""
# Test request with continuation_id (existing conversation)
request = ContinuationRequest(
prompt="Continue analysis", continuation_id="12345678-1234-1234-1234-123456789012"
# Mock the model
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
continuation_data = self.tool._check_continuation_opportunity(request)
# Execute tool without continuation_id (new conversation)
arguments = {"prompt": "Analyze this code"}
response = await self.tool.execute(arguments)
assert continuation_data is None
# Parse response
response_data = json.loads(response[0].text)
# Should offer continuation for new conversation
assert response_data["status"] == "continuation_available"
assert "continuation_offer" in response_data
assert response_data["continuation_offer"]["remaining_turns"] == MAX_CONVERSATION_TURNS - 1
@patch("utils.conversation_memory.get_redis_client")
def test_create_continuation_offer_response(self, mock_redis):
"""Test creating continuation offer response"""
@patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
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_redis.return_value = mock_client
request = ContinuationRequest(prompt="Test prompt")
content = "This is the analysis result."
continuation_data = {"remaining_turns": 4, "tool_name": "test_continuation"}
# Mock existing thread context with 2 turns
from utils.conversation_memory import ConversationTurn, ThreadContext
# Create continuation offer response
response = self.tool._create_continuation_offer_response(content, continuation_data, request)
thread_context = ThreadContext(
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)
assert response.status == "continuation_available"
assert response.content == content
assert response.continuation_offer is not None
# Mock the model
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="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
assert isinstance(offer, ContinuationOffer)
assert offer.remaining_turns == 4
assert "continuation_id" in offer.suggested_tool_params
assert "You have 4 more exchange(s) available" in offer.message_to_user
# Execute tool with continuation_id
arguments = {"prompt": "Continue analysis", "continuation_id": "12345678-1234-1234-1234-123456789012"}
response = await self.tool.execute(arguments)
# 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.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
async def test_full_response_flow_with_continuation_offer(self, mock_redis):
"""Test complete response flow that creates continuation offer"""
mock_client = Mock()
@@ -152,26 +190,21 @@ class TestClaudeContinuationOffers:
assert "more exchange(s) available" in offer["message_to_user"]
@patch("utils.conversation_memory.get_redis_client")
async def test_gemini_follow_up_takes_precedence(self, mock_redis):
"""Test that Gemini follow-up questions take precedence over continuation offers"""
@patch.dict("os.environ", {"PYTEST_CURRENT_TEST": ""}, clear=False)
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_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:
mock_provider = create_mock_provider()
mock_provider.get_provider_type.return_value = Mock(value="google")
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.
```json
{
"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"
}
```"""
I'd be happy to examine the error handling patterns in more detail if that would be helpful."""
mock_provider.generate_content.return_value = Mock(
content=content_with_followup,
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
@@ -187,12 +220,13 @@ class TestClaudeContinuationOffers:
# Parse response
response_data = json.loads(response[0].text)
# Should be follow-up, not continuation offer
assert response_data["status"] == "requires_continuation"
assert "follow_up_request" in response_data
assert response_data.get("continuation_offer") is None
# Should always offer continuation
assert response_data["status"] == "continuation_available"
assert "continuation_offer" in response_data
assert response_data["continuation_offer"]["remaining_turns"] == MAX_CONVERSATION_TURNS - 1
@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):
"""Test that threaded conversations still get continuation offers when turns remain"""
mock_client = Mock()
@@ -236,81 +270,60 @@ class TestClaudeContinuationOffers:
assert response_data.get("continuation_offer") is not None
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"""
# 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_redis.return_value = mock_client
request = ContinuationRequest(prompt="Analyze architecture")
content = "Architecture analysis complete."
continuation_data = {"remaining_turns": 3, "tool_name": "test_continuation"}
# Mock existing thread context at max turns
from utils.conversation_memory import ConversationTurn, ThreadContext
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
message = offer.message_to_user
thread_context = ThreadContext(
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
assert "continue this analysis" in message
assert "continuation_id" in message
assert "test_continuation tool call" in message
assert "3 more exchange(s)" in message
# Mock the model
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="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
suggested_params = offer.suggested_tool_params
assert "continuation_id" in suggested_params
assert "prompt" in suggested_params
assert isinstance(suggested_params["continuation_id"], str)
# Execute tool with continuation_id at max turns
arguments = {"prompt": "Final question", "continuation_id": "12345678-1234-1234-1234-123456789012"}
response = await self.tool.execute(arguments)
@patch("utils.conversation_memory.get_redis_client")
def test_continuation_offer_metadata(self, mock_redis):
"""Test that continuation offer includes proper metadata"""
mock_client = Mock()
mock_redis.return_value = mock_client
# Parse response
response_data = json.loads(response[0].text)
request = ContinuationRequest(prompt="Test")
content = "Test content"
continuation_data = {"remaining_turns": 2, "tool_name": "test_continuation"}
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
# Should NOT offer continuation since we're at max turns
assert response_data["status"] == "success"
assert response_data.get("continuation_offer") is None
class TestContinuationIntegration:
@@ -320,7 +333,8 @@ class TestContinuationIntegration:
self.tool = ClaudeContinuationTool()
@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"""
mock_client = Mock()
mock_redis.return_value = mock_client
@@ -336,11 +350,29 @@ class TestContinuationIntegration:
mock_client.get.side_effect = side_effect_get
request = ContinuationRequest(prompt="Initial analysis", files=["/test/file.py"])
content = "Analysis result"
continuation_data = {"remaining_turns": 4, "tool_name": "test_continuation"}
# Mock the model
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 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)
# Parse response
response_data = json.loads(response[0].text)
# Should offer continuation
assert response_data["status"] == "continuation_available"
assert "continuation_offer" in response_data
# Verify thread creation was called (should be called twice: create_thread + add_turn)
assert mock_client.setex.call_count == 2
@@ -359,25 +391,38 @@ class TestContinuationIntegration:
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"] == content
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")
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"""
mock_client = Mock()
mock_redis.return_value = mock_client
# Step 1: Initial request creates continuation offer
request1 = ToolRequest(prompt="Analyze code structure")
continuation_data = {"remaining_turns": 4, "tool_name": "test_continuation"}
response1 = self.tool._create_continuation_offer_response(
"Structure analysis done.", continuation_data, request1
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="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)
# Parse response
response_data = json.loads(response[0].text)
thread_id = response_data["continuation_offer"]["continuation_id"]
# Step 2: Mock the thread context for Claude's follow-up
from utils.conversation_memory import ConversationTurn, ThreadContext
@@ -400,13 +445,24 @@ class TestContinuationIntegration:
mock_client.get.return_value = existing_context.model_dump_json()
# Step 3: Claude uses continuation_id
request2 = ToolRequest(prompt="Now analyze the performance aspects", continuation_id=thread_id)
mock_provider.generate_content.return_value = Mock(
content="Performance analysis done.",
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
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
continuation_data2 = self.tool._check_continuation_opportunity(request2)
assert continuation_data2 is not None
assert continuation_data2["remaining_turns"] == 8 # MAX_CONVERSATION_TURNS(10) - current_turns(1) - 1
assert continuation_data2["tool_name"] == "test_continuation"
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__":

2
tests/test_conversation_history_bug.py
View File

@@ -236,7 +236,7 @@ class TestConversationHistoryBugFix:
# Should include follow-up instructions for new conversation
# (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.add_turn")

99
tests/test_conversation_memory.py
View File

@@ -151,7 +151,6 @@ class TestConversationMemory:
role="assistant",
content="Python is a programming language",
timestamp="2023-01-01T00:01:00Z",
follow_up_question="Would you like examples?",
files=["/home/user/examples/"],
tool_name="chat",
),
@@ -188,11 +187,8 @@ class TestConversationMemory:
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)
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
# Test follow-up attribution
assert "[Gemini's Follow-up: Would you like examples?]" 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
def test_build_conversation_history_empty(self):
"""Test building history with no turns"""
@@ -235,12 +231,11 @@ class TestConversationFlow:
)
mock_client.get.return_value = initial_context.model_dump_json()
# Add assistant response with follow-up
# Add assistant response
success = add_turn(
thread_id,
"assistant",
"Code analysis complete",
follow_up_question="Would you like me to check error handling?",
)
assert success is True
@@ -256,7 +251,6 @@ class TestConversationFlow:
role="assistant",
content="Code analysis complete",
timestamp="2023-01-01T00:00:30Z",
follow_up_question="Would you like me to check error handling?",
)
],
initial_context={"prompt": "Analyze this code"},
@@ -266,9 +260,7 @@ class TestConversationFlow:
success = add_turn(thread_id, "user", "Yes, check error handling")
assert success is True
success = add_turn(
thread_id, "assistant", "Error handling reviewed", follow_up_question="Should I examine the test coverage?"
)
success = add_turn(thread_id, "assistant", "Error handling reviewed")
assert success is True
# REQUEST 3-5: Continue conversation (simulating independent cycles)
@@ -283,14 +275,12 @@ class TestConversationFlow:
role="assistant",
content="Code analysis complete",
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="assistant",
content="Error handling reviewed",
timestamp="2023-01-01T00:02:30Z",
follow_up_question="Should I examine the test coverage?",
),
],
initial_context={"prompt": "Analyze this code"},
@@ -385,18 +375,20 @@ class TestConversationFlow:
# Test early conversation (should allow follow-ups)
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 "Feel free to ask clarifying questions" in early_instructions
# Test mid conversation
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 "Feel free to ask clarifying questions" in mid_instructions
# Test approaching limit (should stop follow-ups)
limit_instructions = get_follow_up_instructions(max_turns - 1, max_turns)
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
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()
# Add Gemini's response with follow-up
# Add Gemini's response
success = add_turn(
thread_id,
"assistant",
"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"],
tool_name="analyze",
)
@@ -514,7 +505,6 @@ class TestConversationFlow:
role="assistant",
content="I've analyzed your codebase structure.",
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"],
tool_name="analyze",
)
@@ -540,7 +530,6 @@ class TestConversationFlow:
role="assistant",
content="I've analyzed your codebase structure.",
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"],
tool_name="analyze",
),
@@ -575,7 +564,6 @@ class TestConversationFlow:
role="assistant",
content="I've analyzed your codebase structure.",
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"],
tool_name="analyze",
),
@@ -604,19 +592,18 @@ class TestConversationFlow:
assert "--- Turn 3 (Gemini using analyze) ---" in history
# 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_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_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_3_files = "Files used in this turn: /project/tests/test_utils.py, /project/coverage.html"
assert turn_1_files in history
assert turn_2_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 "Yes, check the test coverage" 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.)
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
@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")
def test_stateless_request_isolation(self, mock_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()
success = add_turn(
thread_id, "assistant", "Architecture analysis", follow_up_question="Want to explore scalability?"
)
success = add_turn(thread_id, "assistant", "Architecture analysis")
assert success is True
# Process 2: Different "request cycle" accesses same thread
@@ -711,7 +646,6 @@ class TestConversationFlow:
role="assistant",
content="Architecture analysis",
timestamp="2023-01-01T00:00:30Z",
follow_up_question="Want to explore scalability?",
)
],
initial_context={"prompt": "Think about architecture"},
@@ -722,7 +656,6 @@ class TestConversationFlow:
retrieved_context = get_thread(thread_id)
assert retrieved_context is not None
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):
"""Test that build_conversation_history efficiently handles token limits"""
@@ -766,7 +699,7 @@ class TestConversationFlow:
history, tokens = build_conversation_history(context, model_context=None)
# Verify the history was built successfully
assert "=== CONVERSATION HISTORY ===" in history
assert "=== CONVERSATION HISTORY" in history
assert "=== FILES REFERENCED IN THIS CONVERSATION ===" in history
# 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()
@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):
"""Test that a continuation_id from one tool can be used with another tool"""
mock_client = Mock()
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:
mock_provider = create_mock_provider()
mock_provider.get_provider_type.return_value = Mock(value="google")
mock_provider.supports_thinking_mode.return_value = False
# Include follow-up JSON in the content
content_with_followup = """Found potential security issues in authentication logic.
# Simple content without JSON follow-up
content = """Found potential security issues in authentication logic.
```json
{
"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"
}
```"""
I'd be happy to review these security findings in detail if that would be helpful."""
mock_provider.generate_content.return_value = Mock(
content=content_with_followup,
content=content,
usage={"input_tokens": 10, "output_tokens": 20, "total_tokens": 30},
model_name="gemini-2.0-flash-exp",
metadata={"finish_reason": "STOP"},
@@ -126,8 +121,8 @@ class TestCrossToolContinuation:
response = await self.analysis_tool.execute(arguments)
response_data = json.loads(response[0].text)
assert response_data["status"] == "requires_continuation"
continuation_id = response_data["follow_up_request"]["continuation_id"]
assert response_data["status"] == "continuation_available"
continuation_id = response_data["continuation_offer"]["continuation_id"]
# 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
@@ -139,10 +134,9 @@ class TestCrossToolContinuation:
turns=[
ConversationTurn(
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",
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; }"},
@@ -250,6 +244,7 @@ class TestCrossToolContinuation:
@patch("utils.conversation_memory.get_redis_client")
@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):
"""Test that file context is preserved across tool switches"""
mock_client = Mock()

4
tests/test_prompt_regression.py
View File

@@ -109,7 +109,7 @@ class TestPromptRegression:
assert len(result) == 1
output = json.loads(result[0].text)
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"]
@pytest.mark.asyncio
@@ -203,7 +203,7 @@ class TestPromptRegression:
assert len(result) == 1
output = json.loads(result[0].text)
assert output["status"] == "success"
assert "Debug Analysis" in output["content"]
assert "Next Steps:" in output["content"]
assert "Root cause" in output["content"]
@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
mock_get_provider.assert_called_once()
assert mock_get_provider.called
# Verify generate_content was called with thinking_mode
mock_provider.generate_content.assert_called_once()
call_kwargs = mock_provider.generate_content.call_args[1]
@@ -72,7 +72,7 @@ class TestThinkingModes:
response_data = json.loads(result[0].text)
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
@patch("tools.base.BaseTool.get_model_provider")
@@ -96,7 +96,7 @@ class TestThinkingModes:
)
# 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
mock_provider.generate_content.assert_called_once()
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
)
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
@patch("tools.base.BaseTool.get_model_provider")
@@ -127,7 +127,7 @@ class TestThinkingModes:
)
# 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
mock_provider.generate_content.assert_called_once()
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
)
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
@patch("tools.base.BaseTool.get_model_provider")
@@ -159,7 +159,7 @@ class TestThinkingModes:
)
# 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
mock_provider.generate_content.assert_called_once()
call_kwargs = mock_provider.generate_content.call_args[1]
@@ -188,7 +188,7 @@ class TestThinkingModes:
)
# 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
mock_provider.generate_content.assert_called_once()
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
)
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):
"""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
output = json.loads(result[0].text)
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"]
@@ -102,8 +102,8 @@ class TestCodeReviewTool:
)
assert len(result) == 1
assert "Code Review (SECURITY)" in result[0].text
assert "Focus: authentication" in result[0].text
assert "Security issues found" in result[0].text
assert "Claude's Next Steps:" in result[0].text
assert "Security issues found" in result[0].text
@@ -146,7 +146,7 @@ class TestDebugIssueTool:
)
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
@@ -195,8 +195,8 @@ class TestAnalyzeTool:
)
assert len(result) == 1
assert "ARCHITECTURE Analysis" in result[0].text
assert "Analyzed 1 file(s)" in result[0].text
assert "Architecture analysis" in result[0].text
assert "Next Steps:" 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 logging
import os
import re
from abc import ABC, abstractmethod
from typing import Any, Literal, Optional
from mcp.types import TextContent
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 utils import check_token_limit
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 .models import ClarificationRequest, ContinuationOffer, FollowUpRequest, ToolOutput
from .models import ClarificationRequest, ContinuationOffer, ToolOutput
logger = logging.getLogger(__name__)
@@ -363,6 +362,8 @@ class BaseTool(ABC):
if not model_context:
# Manual calculation as fallback
from config import DEFAULT_MODEL
model_name = getattr(self, "_current_model_name", None) or DEFAULT_MODEL
try:
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
model_name = getattr(request, "model", None)
if not model_name:
from config import DEFAULT_MODEL
model_name = DEFAULT_MODEL
# 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:
"""
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
continue the conversation. This method detects such responses and formats them.
This method formats the response and always offers a continuation opportunity
unless max conversation turns have been reached.
Args:
raw_text: The raw text response from the model
request: The original request for context
model_info: Optional dict with model metadata
Returns:
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}")
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 to parse as JSON to check for clarification requests
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
formatted_content = self.format_response(raw_text, request, model_info)
# If we found a follow-up question, prepare the threading response
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
# Always check if we should offer Claude a continuation opportunity
continuation_offer = self._check_continuation_opportunity(request)
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)
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
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},
)
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]:
"""
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_id=thread_id,
message_to_user=(
f"If you'd like to continue this analysis or need further details, "
f"you can use the continuation_id '{thread_id}' in your next {self.name} tool call. "
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}' with any tool and any model. "
f"You have {remaining_turns} more exchange(s) available in this conversation thread."
),
suggested_tool_params={
"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,
)

19
tools/models.py
View File

@@ -7,21 +7,6 @@ from typing import Any, Literal, Optional
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):
"""Offer for Claude to continue conversation when Gemini doesn't ask follow-up"""
@@ -43,15 +28,11 @@ class ToolOutput(BaseModel):
"error",
"requires_clarification",
"requires_file_prompt",
"requires_continuation",
"continuation_available",
] = "success"
content: Optional[str] = Field(None, description="The main content/response from the tool")
content_type: Literal["text", "markdown", "json"] = "text"
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(
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)
content: The actual message content/response
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
tool_name: Which tool generated this turn (for cross-tool tracking)
model_provider: Provider used (e.g., "google", "openai")
@@ -82,7 +81,6 @@ class ConversationTurn(BaseModel):
role: str # "user" or "assistant"
content: str
timestamp: str
follow_up_question: Optional[str] = None
files: Optional[list[str]] = None # Files referenced in this turn
tool_name: Optional[str] = None # Tool used for this turn
model_provider: Optional[str] = None # Model provider (google, openai, etc)
@@ -231,7 +229,6 @@ def add_turn(
thread_id: str,
role: str,
content: str,
follow_up_question: Optional[str] = None,
files: Optional[list[str]] = None,
tool_name: Optional[str] = None,
model_provider: Optional[str] = None,
@@ -249,7 +246,6 @@ def add_turn(
thread_id: UUID of the conversation thread
role: "user" (Claude) or "assistant" (Gemini/O3/etc)
content: The actual message/response content
follow_up_question: Optional follow-up question from assistant
files: Optional list of files referenced in this turn
tool_name: Name of the tool adding this turn (for attribution)
model_provider: Provider used (e.g., "google", "openai")
@@ -287,7 +283,6 @@ def add_turn(
role=role,
content=content,
timestamp=datetime.now(timezone.utc).isoformat(),
follow_up_question=follow_up_question,
files=files, # Preserved for cross-tool file context
tool_name=tool_name, # Track which tool generated this turn
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:,}")
history_parts = [
"=== CONVERSATION HISTORY ===",
"=== CONVERSATION HISTORY (CONTINUATION) ===",
f"Thread: {context.thread_id}",
f"Tool: {context.tool_name}", # Original tool that started the conversation
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
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
turn_content = "\n".join(turn_parts)
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.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