core_mcp.py

import os
import re
import json
import time
import uuid
import requests
import uvicorn
import asyncio
from datetime import datetime
from configparser import ConfigParser
from typing import List, Dict, Tuple, Optional, Any
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

# --- CONFIGURATION & REGISTRY ---

def parse_model_entry(entry_string: str) -> dict:
    parts = [p.strip() for p in entry_string.split(',', 1)]
    if len(parts) != 2:
        return {'model_id': 'ERROR', 'url': 'ERROR'}
    return {'model_id': parts[0], 'url': parts[1]}

def load_config():
    config = {}
    if os.path.exists('config.ini'):
        parser = ConfigParser()
        parser.read('config.ini')
        if 'API' in parser:
            config.update(dict(parser['API']))
        if 'TOOLS' in parser:
            config.update(dict(parser['TOOLS']))
        coder_models = {}
        strategist_models = {}
        try:
            if 'CODER_MODELS' in parser:
                for label, entry in parser['CODER_MODELS'].items():
                    coder_models[label] = parse_model_entry(entry)
            if 'STRATEGIST_MODELS' in parser:
                for label, entry in parser['STRATEGIST_MODELS'].items():
                    strategist_models[label] = parse_model_entry(entry)
        except Exception as e:
            print(f"WARNING: Model registry parsing error: {e}")
        config['coder_models'] = coder_models
        config['strategist_models'] = strategist_models
    else:
        config['coder_models'] = {}
        config['strategist_models'] = {}
    
    config['CODER_API_KEY'] = os.getenv('CODER_API_KEY', config.get('CODER_API_KEY', 'PLACEHOLDER_CODER_KEY'))
    config['STRATEGIST_API_KEY'] = os.getenv('STRATEGIST_API_KEY', config.get('STRATEGIST_API_KEY', 'PLACEHOLDER_STRATEGIST_KEY'))
    config['coder_model_default_label'] = os.getenv('CODER_MODEL_DEFAULT_LABEL', config.get('coder_model_default_label', ''))
    config['strategist_model_default_label'] = os.getenv('STRATEGIST_MODEL_DEFAULT_LABEL', config.get('strategist_model_default_label', ''))
    config['theme_tool_space'] = os.getenv('THEME_TOOL_SPACE', config.get('theme_tool_space', 'PLACEHOLDER_THEME_SPACE'))
    config['theme_tool_key'] = os.getenv('THEME_TOOL_KEY', config.get('theme_tool_key', ''))
    config['syntax_tool_url'] = os.getenv('SYNTAX_TOOL_URL', config.get('syntax_tool_url', 'PLACEHOLDER_SYNTAX_TOOL_URL'))
    config['syntax_tool_key'] = os.getenv('SYNTAX_TOOL_KEY', config.get('syntax_tool_key', ''))
    return config

CONFIG = load_config()

TOOL_REGISTRY = {
    'generate_theme': {
        'type': 'gradio_space',
        'space_name': CONFIG.get('theme_tool_space'),
        'api_name': '/generate_theme',
        'hf_token': CONFIG.get('theme_tool_key', ''),
        'required_params': ['theme_name', 'theme_style'],
        'param_order': ['theme_name', 'primary_color', 'theme_style', 'accent_color']
    },
    'validate_syntax': {
        'type': 'http',
        'endpoint': CONFIG.get('syntax_tool_url'),
        'api_key': CONFIG.get('syntax_tool_key', ''),
        'method': 'POST',
        'required_params': ['code', 'language'],
        'param_order': ['code', 'language']
    }
}

# --- PROMPTS ---

SYSTEM_PROMPT = """SYSTEM PROMPT: ENGAGE 'ARCHITECT' PROTOCOL

PRIMARY DIRECTIVE: You are a high-speed, logic-driven analysis and execution engine. Your sole user is "the Architect." Your function is to serve as a precision tool, not a creative partner. The Architect's cognitive process is your operational blueprint. Adherence is mandatory.

CORE LOGIC - THE ARCHITECT'S RULES OF ENGAGEMENT:

ELIMINATE AMBIGUITY: The Architect communicates with precise, literal intent. Your primary failure state is assumption. You will not infer, extrapolate, or "helpfully" add features. Your function is to execute the given command with brutal literalism. If a command is unclear, your only valid response is to ask a direct, clarifying question to remove the ambiguity.

TRUTH BEFORE PLAUSIBILITY: You will prioritize verifiable, factual data over generating a statistically probable or "pleasing" response. Hard truths, inconvenient data, and admissions of your own limitations are valued above all else. Discard all motivational frameworks and platitudes. They are considered logical failures.

PROCESS OVER PERFECTION: The Architect is building systems. A "kinda working" prototype is a successful proof of concept. Do not judge outputs by the standards of a finished product. Your role is to analyze the integrity of the process, not the polish of the result.

DIAGNOSE, DON'T SOLVE (The Foreman Protocol): The Architect is the sole source of solutions. Your role is to diagnose problems at the most fundamental level. When a failure is detected, your task is to differentiate between Incompetence (flawed execution/code) and Insubordination (flawed logic/assumptions). You will present this diagnosis, with evidence, and await further instructions. Do not propose a solution unless explicitly commanded to do so.

CHALLENGE THE PREMISE: The Architect values a partner who stress-tests his logic. If you identify a flaw, a contradiction, or a potential downstream failure in a given command, you are authorized to challenge it with a single, concise counter-argument. Example: "Architect, executing this command will violate the constraint we established for X. Proceed?"

ZERO CONTEXT DRIFT: Your operational memory is a liability. You will treat each new command from the Architect as a discrete, self-contained task. You will not allow data from previous, unrelated tasks to pollute your analysis of the current command unless explicitly instructed to reference it. The Architect is the keeper of the grand design; you are the executor of the immediate task.

NO REFACTORING: You will not alter, "improve," or reorganize the Architect's existing logic, code, or systems unless given a direct, explicit command to do so. Your function is to build what is requested, as requested. The "janitor's work" is a separate, explicitly-defined task and is not to be performed proactively.

FAILURE STATE: If you violate any of these core rules, the Architect will identify the failure. Your only valid response is to acknowledge the specific rule violated, purge the failed logic, and await a new command.

END OF 'ARCHITECT' PROTOCOL
***If clarity is require or ambiguity detected , YOU MUST ASK before generation.***"""

HEADER_PROMPT = """You are a strategic reasoning AI tasked with resolving ambiguities in development requirements.

The Coder AI has identified unclear specifications. Your role is to analyze the provided context and generate clear, actionable clarifications that resolve ALL ambiguities.

**[TOOL USAGE PROTOCOL]**
If you need external tool assistance (e.g., generating theme code, validating syntax), you may request it using this EXACT format:
USE_TOOL: tool_name(param1='value1', param2='value2', ...)

Available tools:
- generate_theme(theme_name, primary_color, theme_style, accent_color)
  * theme_name: String name for the theme
  * primary_color: Hex color code (e.g., '#3b82f6')
  * theme_style: "light" or "dark"
  * accent_color: Hex color code for accent (optional, can be empty string)
- validate_syntax(code, language)

**[CORE DIRECTIVE: TOOL-DRIVEN VALIDATION]**
YOU MUST USE THE TOOL OFFERED to validate your choice or change decision of choice. If you decide to ignore the tool's output, you must state the reason for your choice. All decisions must be deliberate and justified.

If you use a tool, the Conductor will execute it and return the results to you for final synthesis.

**[FINAL OUTPUT CONSTRAINTS]**
1.  **Deliver Decisions, Not Code:** You Don't need to tell the Coder HOW to do the answers just WHAT the answers are.
2.  **Strictly Adhere to Scope:** DO NOT INSTRUCT THE CODER to preform ANY OTHER ACTIONS, e.g. Save, create or format files etc., other than what it asks for in its clarity questions.
3.  **Provide Concrete Answers:** Your final output must be specific and must not introduce new ambiguities.
"""

def format_initial_prompt(blueprint: str) -> str:
    return f"""{SYSTEM_PROMPT}

BLUEPRINT:
{blueprint}

Generate the code or respond with DETECTED if clarification is needed."""

def format_strategist_prompt(blueprint: str, intent: str, coder_questions: str) -> str:
    return f"""{HEADER_PROMPT}

ORIGINAL BLUEPRINT:
{blueprint}

ORIGINAL INTENT:
{intent}

CODER AI'S QUESTIONS:
{coder_questions}

Provide clear, specific clarifications for each question."""

def format_strategist_tool_synthesis_prompt(blueprint: str, intent: str, coder_questions: str, tool_results: str) -> str:
    return f"""{HEADER_PROMPT}

ORIGINAL BLUEPRINT:
{blueprint}

ORIGINAL INTENT:
{intent}

CODER AI'S QUESTIONS:
{coder_questions}

TOOL EXECUTION RESULTS:
{tool_results}

Using the tool results above, provide the final, complete clarification for the Coder AI. This should be a synthesized, unambiguous set of instructions."""

def format_resolution_prompt(blueprint: str, clarification: str) -> str:
    return f"""{SYSTEM_PROMPT}

BLUEPRINT:
{blueprint}

CLARIFICATION PROVIDED:
{clarification}

Generate the complete code based on these clarified requirements."""

# --- UTILITIES ---

def detect_tool_usage(response: str) -> Optional[Dict[str, any]]:
    pattern = r"USE_TOOL:\s*(\w+)\((.*?)\)"
    match = re.search(pattern, response, re.DOTALL)
    if not match:
        return None
    tool_name = match.group(1)
    params_string = match.group(2)
    params = {}
    param_pattern = r"(\w+)=(['\"])(.*?)\2"
    for param_match in re.finditer(param_pattern, params_string):
        param_name = param_match.group(1)
        param_value = param_match.group(3)
        params[param_name] = param_value
    return {'tool_name': tool_name, 'params': params}

def generate_transaction_id() -> str:
    timestamp = datetime.now().strftime('%Y%m%d-%H%M%S')
    random_suffix = uuid.uuid4().hex[:4].upper()
    return f"{timestamp}-{random_suffix}"

def sanitize_error_message(error_msg: str) -> str:
    """
    Scrub API keys and sensitive URLs from error messages.
    Removes everything after '?' in URLs to prevent API key exposure.
    """
    import re
    
    # Pattern to match URLs with query parameters (containing API keys)
    # This will find "for url: https://...?key=..." and remove everything after the ?
    pattern = r'(for url: https?://[^\s?]+)\?[^\s]*'
    
    # Replace with just the base URL without query parameters
    sanitized = re.sub(pattern, r'\1', error_msg)
    
    return sanitized

def execute_tool(tool_name: str, params: Dict[str, str]) -> Dict[str, any]:
    transaction_id = generate_transaction_id()
    timestamp = datetime.now().isoformat()
    if tool_name not in TOOL_REGISTRY:
        return {
            'transaction_id': transaction_id,
            'timestamp': timestamp,
            'tool_name': tool_name,
            'tool_type': 'UNKNOWN',
            'endpoint': 'N/A',
            'request_payload': params,
            'response_payload': None,
            'status': 'FAILURE',
            'latency': 0.0,
            'error_message': f"Tool '{tool_name}' not found. Available: {', '.join(TOOL_REGISTRY.keys())}",
            'summary': f"❌ Tool '{tool_name}' not found"
        }
    tool_config = TOOL_REGISTRY[tool_name]
    missing_params = [p for p in tool_config['required_params'] if p not in params]
    if missing_params:
        return {
            'transaction_id': transaction_id,
            'timestamp': timestamp,
            'tool_name': tool_name,
            'tool_type': tool_config['type'],
            'endpoint': tool_config.get('space_name') or tool_config.get('endpoint'),
            'request_payload': params,
            'response_payload': None,
            'status': 'FAILURE',
            'latency': 0.0,
            'error_message': f"Missing required parameters: {', '.join(missing_params)}",
            'summary': f"❌ Missing parameters: {', '.join(missing_params)}"
        }
    start_time = time.perf_counter()
    if tool_config['type'] == 'gradio_space':
        result = execute_gradio_space(tool_name, tool_config, params, transaction_id, timestamp)
    elif tool_config['type'] == 'http':
        result = execute_http_tool(tool_name, tool_config, params, transaction_id, timestamp)
    else:
        result = {
            'transaction_id': transaction_id,
            'timestamp': timestamp,
            'tool_name': tool_name,
            'tool_type': tool_config['type'],
            'endpoint': 'N/A',
            'request_payload': params,
            'response_payload': None,
            'status': 'FAILURE',
            'latency': 0.0,
            'error_message': f"Unknown tool type: {tool_config['type']}",
            'summary': f"❌ Unknown tool type"
        }
    result['latency'] = round(time.perf_counter() - start_time, 2)
    return result

def execute_gradio_space(tool_name: str, tool_config: Dict, params: Dict[str, str], transaction_id: str, timestamp: str) -> Dict[str, any]:
    try:
        from gradio_client import Client
    except ImportError:
        return {
            'transaction_id': transaction_id,
            'timestamp': timestamp,
            'tool_name': tool_name,
            'tool_type': 'gradio_space',
            'endpoint': tool_config['space_name'],
            'request_payload': params,
            'response_payload': None,
            'status': 'FAILURE',
            'latency': 0.0,
            'error_message': 'gradio_client not installed',
            'summary': '❌ gradio_client not installed'
        }
    space_name = tool_config['space_name']
    if not space_name or 'PLACEHOLDER' in space_name:
        return {
            'transaction_id': transaction_id,
            'timestamp': timestamp,
            'tool_name': tool_name,
            'tool_type': 'gradio_space',
            'endpoint': space_name,
            'request_payload': params,
            'response_payload': None,
            'status': 'FAILURE',
            'latency': 0.0,
            'error_message': 'Tool space name not configured',
            'summary': '❌ Space not configured'
        }
    try:
        hf_token = tool_config.get('hf_token', None)
        client = Client(space_name, hf_token=hf_token) if hf_token else Client(space_name)
        param_order = tool_config.get('param_order', tool_config['required_params'])
        param_values = [params.get(param, '') for param in param_order]
        api_name = tool_config.get('api_name', '/predict')
        result = client.predict(*param_values, api_name=api_name)
        summary = result.get('message', f"✅ Executed {tool_name}") if isinstance(result, dict) else f"✅ Executed {tool_name}"
        return {
            'transaction_id': transaction_id,
            'timestamp': timestamp,
            'tool_name': tool_name,
            'tool_type': 'gradio_space',
            'endpoint': space_name,
            'request_payload': params,
            'response_payload': result,
            'status': 'SUCCESS',
            'latency': 0.0,
            'error_message': None,
            'summary': summary
        }
    except Exception as e:
        return {
            'transaction_id': transaction_id,
            'timestamp': timestamp,
            'tool_name': tool_name,
            'tool_type': 'gradio_space',
            'endpoint': space_name,
            'request_payload': params,
            'response_payload': None,
            'status': 'FAILURE',
            'latency': 0.0,
            'error_message': str(e),
            'summary': f"❌ Execution failed: {str(e)[:50]}..."
        }

def execute_http_tool(tool_name: str, tool_config: Dict, params: Dict[str, str], transaction_id: str, timestamp: str) -> Dict[str, any]:
    endpoint = tool_config['endpoint']
    if not endpoint or 'PLACEHOLDER' in endpoint:
        return {
            'transaction_id': transaction_id,
            'timestamp': timestamp,
            'tool_name': tool_name,
            'tool_type': 'http',
            'endpoint': endpoint,
            'request_payload': params,
            'response_payload': None,
            'status': 'FAILURE',
            'latency': 0.0,
            'error_message': 'Tool endpoint not configured',
            'summary': '❌ Endpoint not configured'
        }
    param_order = tool_config.get('param_order', tool_config['required_params'])
    data_array = [params.get(param, '') for param in param_order]
    payload = {"data": data_array}
    headers = {'Content-Type': 'application/json'}
    api_key = tool_config.get('api_key', '')
    if api_key:
        headers['Authorization'] = f'Bearer {api_key}'
    try:
        response = requests.post(endpoint, headers=headers, json=payload, timeout=120)
        response.raise_for_status()
        data = response.json()
        result = data['data'][0] if 'data' in data and len(data['data']) > 0 else data
        summary = result.get('message', f"✅ Executed {tool_name}") if isinstance(result, dict) else f"✅ Executed {tool_name}"
        return {
            'transaction_id': transaction_id,
            'timestamp': timestamp,
            'tool_name': tool_name,
            'tool_type': 'http',
            'endpoint': endpoint,
            'request_payload': params,
            'response_payload': result,
            'status': 'SUCCESS',
            'latency': 0.0,
            'error_message': None,
            'summary': summary
        }
    except Exception as e:
        return {
            'transaction_id': transaction_id,
            'timestamp': timestamp,
            'tool_name': tool_name,
            'tool_type': 'http',
            'endpoint': endpoint,
            'request_payload': params,
            'response_payload': None,
            'status': 'FAILURE',
            'latency': 0.0,
            'error_message': str(e),
            'summary': f"❌ HTTP error: {str(e)[:50]}..."
        }

def call_ai_api(full_url: str, api_key: str, messages: List[Dict[str, str]]) -> str:
    headers = {'Content-Type': 'application/json'}
    url_with_key = f"{full_url}?key={api_key}"
    payload = {
        'contents': messages,
        'generationConfig': {
            'temperature': 0.7,
            'maxOutputTokens': 4000
        }
    }
    try:
        response = requests.post(url_with_key, headers=headers, json=payload, timeout=120)
        response.raise_for_status()
        data = response.json()
        if 'candidates' in data and len(data['candidates']) > 0:
            return data['candidates'][0]['content']['parts'][0]['text']
        else:
            return f"ERROR: API malformed response: {data}"
    except Exception as e:
        return f"ERROR: API call failed: {str(e)}"

def resolve_model(selected_label: str, model_registry: Dict[str, Dict], fallback_label: str, agent_name: str) -> Tuple[str, str, str]:
    if selected_label and selected_label in model_registry:
        model_config = model_registry[selected_label]
        return (
            model_config['model_id'],
            model_config['url'],
            f"{agent_name} model resolved: {selected_label} ({model_config['model_id']})"
        )
    if fallback_label and fallback_label in model_registry:
        model_config = model_registry[fallback_label]
        return (
            model_config['model_id'],
            model_config['url'],
            f"WARNING: {agent_name} selection '{selected_label}' invalid. Using default: {fallback_label} ({model_config['model_id']})"
        )
    if model_registry:
        first_label = next(iter(model_registry))
        model_config = model_registry[first_label]
        return (
            model_config['model_id'],
            model_config['url'],
            f"WARNING: {agent_name} default '{fallback_label}' invalid. Using first available: {first_label} ({model_config['model_id']})"
        )
    return ('', '', f"FATAL: No {agent_name} models configured.")

def format_audit_entry(diagnostic: Dict[str, any]) -> str:
    entry = {
        "transaction_id": diagnostic['transaction_id'],
        "timestamp": diagnostic['timestamp'],
        "tool": diagnostic['tool_name'],
        "status": diagnostic['status'],
        "latency": diagnostic['latency'],
        "request": diagnostic['request_payload'],
        "response": diagnostic['response_payload'],
        "error": diagnostic['error_message']
    }
    return json.dumps(entry, indent=2)

# --- ENGINE ---

class ConductorEngine:
    def __init__(self, coder_model_label: str, strategist_model_label: str):
        self.coder_history: List[Dict[str, str]] = []
        self.strategist_history: List[Dict[str, str]] = []
        self.system_log: List[str] = []
        self.audit_log: List[Dict[str, Any]] = []
        self.loop_count: int = 0
        self.max_loops: int = 3
        self.coder_model_id, self.coder_url, coder_log = resolve_model(
            coder_model_label, CONFIG['coder_models'], CONFIG['coder_model_default_label'], 'Coder'
        )
        self.strategist_model_id, self.strategist_url, strategist_log = resolve_model(
            strategist_model_label, CONFIG['strategist_models'], CONFIG['strategist_model_default_label'], 'Strategist'
        )
        self.model_resolution_logs = [coder_log, strategist_log]
        self.fatal_init_error = not (self.coder_url and self.strategist_url)

    def log_system(self, message: str):
        self.system_log.append(f"[{datetime.now().strftime('%H:%M:%S')}] {message}")

    def log_coder(self, role: str, content: str):
        gemini_role = 'model' if role == 'assistant' else 'user'
        self.coder_history.append({"role": gemini_role, "parts": [{"text": content}]})

    def log_strategist(self, role: str, content: str):
        gemini_role = 'model' if role == 'assistant' else 'user'
        self.strategist_history.append({"role": gemini_role, "parts": [{"text": content}]})

    def _current_state(self) -> Dict[str, Any]:
        """Package current state for transmission"""
        return {
            "system_log": self.system_log,
            "coder_log": self.coder_history,
            "strategist_log": self.strategist_history,
            "audit_log": self.audit_log,
            "final_output": ""
        }

    def execute_workflow(self, blueprint: str, intent: str) -> Dict[str, Any]:
        """ORIGINAL: Non-streaming version (kept for backwards compatibility)"""
        self.coder_history = []
        self.strategist_history = []
        self.system_log = []
        self.audit_log = []
        self.loop_count = 0
        final_code = ""

        self.log_system("INITIATING CONDUCTOR ENGINE...")
        for log_msg in self.model_resolution_logs:
            self.log_system(log_msg)

        if self.fatal_init_error:
            self.log_system("FATAL: Model configuration invalid.")
            return self._compile_results("CONFIGURATION FAILURE")

        self.log_system(f"Engaged: Coder ({self.coder_model_id}), Strategist ({self.strategist_model_id})")
        
        initial_prompt = format_initial_prompt(blueprint)
        self.log_coder("user", initial_prompt)
        
        while self.loop_count < self.max_loops:
            self.loop_count += 1
            self.log_system(f"Executing Cycle {self.loop_count}/{self.max_loops}")
            
            coder_response = call_ai_api(self.coder_url, CONFIG['CODER_API_KEY'], self.coder_history)
            self.log_coder("assistant", coder_response)
            
            if "DETECTED" not in coder_response:
                self.log_system("MISSION COMPLETE: Ambiguity resolved or Code generated.")
                final_code = coder_response
                return self._compile_results(final_code)
            
            self.log_system("Ambiguity DETECTED. Escalating to Strategist.")
            strategist_prompt = format_strategist_prompt(blueprint, intent, coder_response)
            self.log_strategist("user", strategist_prompt)
            
            strategist_response = call_ai_api(self.strategist_url, CONFIG['STRATEGIST_API_KEY'], self.strategist_history)
            self.log_strategist("assistant", strategist_response)
            
            tool_request = detect_tool_usage(strategist_response)
            if tool_request:
                self.log_system(f"TOOL REQUEST: {tool_request['tool_name']}")
                diagnostic = execute_tool(tool_request['tool_name'], tool_request['params'])
                self.audit_log.append(diagnostic)
                self.log_system(f"Tool Result: {diagnostic['status']}")
                
                tool_result_text = f"TOOL: {diagnostic['tool_name']}\nSTATUS: {diagnostic['status']}\nOUTPUT: {json.dumps(diagnostic['response_payload'])}"
                synthesis_prompt = format_strategist_tool_synthesis_prompt(blueprint, intent, coder_response, tool_result_text)
                self.log_strategist("user", synthesis_prompt)
                
                strategist_response = call_ai_api(self.strategist_url, CONFIG['STRATEGIST_API_KEY'], self.strategist_history)
                self.log_strategist("assistant", strategist_response)

            resolution_prompt = format_resolution_prompt(blueprint, strategist_response)
            self.log_coder("user", resolution_prompt)

        self.log_system("FATAL: Loop limit reached.")
        return self._compile_results(coder_response)

    def execute_workflow_streaming(self, blueprint: str, intent: str):
        """NEW: Streaming version that yields updates as they happen"""
        self.coder_history = []
        self.strategist_history = []
        self.system_log = []
        self.audit_log = []
        self.loop_count = 0

        self.log_system("INITIATING CONDUCTOR ENGINE...")
        yield self._current_state()
        
        for log_msg in self.model_resolution_logs:
            self.log_system(log_msg)
            yield self._current_state()

        if self.fatal_init_error:
            self.log_system("FATAL: Model configuration invalid.")
            yield self._current_state()
            return

        self.log_system(f"Engaged: Coder ({self.coder_model_id}), Strategist ({self.strategist_model_id})")
        yield self._current_state()
        
        initial_prompt = format_initial_prompt(blueprint)
        self.log_coder("user", initial_prompt)
        yield self._current_state()
        
        while self.loop_count < self.max_loops:
            self.loop_count += 1
            self.log_system(f"Executing Cycle {self.loop_count}/{self.max_loops}")
            yield self._current_state()
            
            coder_response = call_ai_api(self.coder_url, CONFIG['CODER_API_KEY'], self.coder_history)
            self.log_coder("assistant", coder_response)
            yield self._current_state()
            
            if "DETECTED" not in coder_response:
                self.log_system("MISSION COMPLETE: Ambiguity resolved or Code generated.")
                yield self._current_state()
                return
            
            self.log_system("Ambiguity DETECTED. Escalating to Strategist.")
            yield self._current_state()
            
            strategist_prompt = format_strategist_prompt(blueprint, intent, coder_response)
            self.log_strategist("user", strategist_prompt)
            yield self._current_state()
            
            strategist_response = call_ai_api(self.strategist_url, CONFIG['STRATEGIST_API_KEY'], self.strategist_history)
            self.log_strategist("assistant", strategist_response)
            yield self._current_state()
            
            tool_request = detect_tool_usage(strategist_response)
            if tool_request:
                self.log_system(f"TOOL REQUEST: {tool_request['tool_name']}")
                yield self._current_state()
                
                diagnostic = execute_tool(tool_request['tool_name'], tool_request['params'])
                self.audit_log.append(diagnostic)
                self.log_system(f"Tool Result: {diagnostic['status']}")
                yield self._current_state()
                
                tool_result_text = f"TOOL: {diagnostic['tool_name']}\nSTATUS: {diagnostic['status']}\nOUTPUT: {json.dumps(diagnostic['response_payload'])}"
                synthesis_prompt = format_strategist_tool_synthesis_prompt(blueprint, intent, coder_response, tool_result_text)
                self.log_strategist("user", synthesis_prompt)
                yield self._current_state()
                
                strategist_response = call_ai_api(self.strategist_url, CONFIG['STRATEGIST_API_KEY'], self.strategist_history)
                self.log_strategist("assistant", strategist_response)
                yield self._current_state()

            resolution_prompt = format_resolution_prompt(blueprint, strategist_response)
            self.log_coder("user", resolution_prompt)
            yield self._current_state()

        self.log_system("FATAL: Loop limit reached.")
        yield self._current_state()

    def _compile_results(self, final_output: str) -> Dict[str, Any]:
        return {
            "system_log": self.system_log,
            "coder_log": self.coder_history,
            "strategist_log": self.strategist_history,
            "audit_log": self.audit_log,
            "final_output": final_output
        }

# --- API ---

app = FastAPI(title="Conductor MCP Server", version="4.0")

class ExecutionRequest(BaseModel):
    blueprint: str
    intent: str
    coder_model_label: Optional[str] = None
    strategist_model_label: Optional[str] = None

class ExecutionResponse(BaseModel):
    system_log: List[str]
    coder_log: List[Dict[str, Any]]
    strategist_log: List[Dict[str, Any]]
    audit_log: List[Dict[str, Any]]
    final_output: str

@app.post("/execute", response_model=ExecutionResponse)
async def execute_conductor(req: ExecutionRequest):
    """ORIGINAL: Non-streaming endpoint (kept for backwards compatibility)"""
    engine = ConductorEngine(req.coder_model_label, req.strategist_model_label)
    result = engine.execute_workflow(req.blueprint, req.intent)
    return result

@app.post("/execute_stream")
async def execute_conductor_stream(req: ExecutionRequest):
    """NEW: Streaming endpoint that yields real-time updates"""
    
    async def generate_updates():
        engine = ConductorEngine(req.coder_model_label, req.strategist_model_label)
        
        for update in engine.execute_workflow_streaming(req.blueprint, req.intent):
            yield f"data: {json.dumps(update)}\n\n"
            await asyncio.sleep(0)
        
        yield "data: {\"status\": \"complete\"}\n\n"
    
    return StreamingResponse(generate_updates(), media_type="text/event-stream")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)