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']
    }
}

ARCHITECT_HEADER = """# ============================================================================
# WARNING: AI-GENERATED CODE
# ----------------------------------------------------------------------------
# This code was generated by the Project: Conductor AI system.
# It is provided for demonstration and validation purposes only. The user 
# assumes all responsibility for testing, security, and implementation.
# ============================================================================
"""

ARCHITECT_FOOTER = """
# ============================================================================
# ARCHITECT'S NOTE ON CODE VARIABILITY:
# The quality of this output is a function of a complex system. Key variables include:
#   - The LLM Agent: The specific capabilities of the selected Coder/Strategist models.
#   - The Directives: The clarity and detail of the Blueprint and Intent.
#   - The Personas: The logic protocols governing the AI's behavior.
# Always treat this code as a sophisticated first draft, not a finished product.
# ============================================================================
"""

# --- 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}\n\nBLUEPRINT:\n{blueprint}\n\nGenerate 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}\n\nORIGINAL BLUEPRINT:\n{blueprint}\n\nORIGINAL INTENT:\n{intent}\n\nCODER AI'S QUESTIONS:\n{coder_questions}\n\nProvide 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}\n\nORIGINAL BLUEPRINT:\n{blueprint}\n\nORIGINAL INTENT:\n{intent}\n\nCODER AI'S QUESTIONS:\n{coder_questions}\n\nTOOL EXECUTION RESULTS:\n{tool_results}\n\nUsing the tool results above, provide the final, complete clarification for the Coder AI."

def format_resolution_prompt(blueprint: str, clarification: str) -> str:
    return f"{SYSTEM_PROMPT}\n\nBLUEPRINT:\n{blueprint}\n\nCLARIFICATION PROVIDED:\n{clarification}\n\nGenerate the complete code based on these clarified requirements."

def format_strategist_audit_prompt(blueprint: str, intent: str, generated_code: str) -> str:
    return f"{HEADER_PROMPT}\n\nORIGINAL BLUEPRINT:\n{blueprint}\n\nORIGINAL INTENT:\n{intent}\n\nGENERATED CODE:\n{generated_code}\n\n**[AUDIT DIRECTIVE]**..." # Assuming audit prompt is correct

# --- 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):
        params[param_match.group(1)] = param_match.group(3)
    return {'tool_name': tool_name, 'params': params}

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

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 apply_architect_signature(code: str) -> str:
    return f"{ARCHITECT_HEADER}\n{code}\n{ARCHITECT_FOOTER}"

# --- OPTION A FIX: Adaptive response parser ---
def parse_gradio_response(raw_response: Any) -> Any:
    if isinstance(raw_response, dict):
        return raw_response
    if isinstance(raw_response, str):
        if not raw_response.strip():
            return {"raw_response": "", "parse_note": "Empty response from Gradio Space"}
        try:
            return json.loads(raw_response)
        except json.JSONDecodeError:
            return {"raw_response": raw_response, "parse_note": "Non-JSON string response"}
    return {"raw_response": str(raw_response), "parse_note": f"Unexpected type: {type(raw_response).__name__}"}

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, 'status': 'FAILURE', 'error_message': f"Tool '{tool_name}' not found."}
    
    tool_config = TOOL_REGISTRY[tool_name]
    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, 'status': 'FAILURE', 'error_message': f"Unknown tool type: {tool_config['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]:
    base_result = {'transaction_id': transaction_id, 'timestamp': timestamp, 'tool_name': tool_name, 'tool_type': 'gradio_space', 'endpoint': tool_config['space_name'], 'request_payload': params}
    try:
        from gradio_client import Client
        
        # --- DEFINITIVE FIX IS HERE ---
        # 1. Define space_name from the config first.
        space_name = tool_config['space_name']
        
        # 2. Get the token.
        hf_token = tool_config.get('hf_token')
        
        # 3. Conditionally create the client ONLY if a token exists.
        if hf_token:
            client = Client(space_name, hf_token=hf_token)
        else:
            client = Client(space_name)
        # --- END OF FIX ---

        param_order = tool_config.get('param_order', list(params.keys()))
        param_values = [params.get(p, '') for p in param_order]
        api_name = tool_config.get('api_name', '/predict')
        
        raw_result = client.predict(*param_values, api_name=api_name)
        parsed_result = parse_gradio_response(raw_result)
        
        summary = f"✅ Executed {tool_name}"
        if isinstance(parsed_result, dict) and 'parse_note' in parsed_result:
            summary += f" (Note: {parsed_result['parse_note']})"
        
        base_result.update({'response_payload': parsed_result, 'status': 'SUCCESS', 'summary': summary})
        return base_result
    except Exception as e:
        base_result.update({'response_payload': None, 'status': 'FAILURE', 'error_message': str(e), 'summary': f"❌ Execution failed: {str(e)[:50]}..."})
        return base_result

def execute_http_tool(tool_name: str, tool_config: Dict, params: Dict[str, str], transaction_id: str, timestamp: str) -> Dict[str, any]:
    base_result = {'transaction_id': transaction_id, 'timestamp': timestamp, 'tool_name': tool_name, 'tool_type': 'http', 'endpoint': tool_config['endpoint'], 'request_payload': params}
    try:
        param_order = tool_config.get('param_order', list(params.keys()))
        data_array = [params.get(p, '') for p in param_order]
        payload = {"data": data_array}
        headers = {'Content-Type': 'application/json'}
        if tool_config.get('api_key'): headers['Authorization'] = f"Bearer {tool_config['api_key']}"
        
        response = requests.post(tool_config['endpoint'], headers=headers, json=payload, timeout=120)
        response.raise_for_status()
        data = response.json()
        result = data['data'][0] if 'data' in data and data['data'] else data
        
        base_result.update({'response_payload': result, 'status': 'SUCCESS', 'summary': f"✅ Executed {tool_name}"})
        return base_result
    except Exception as e:
        base_result.update({'response_payload': None, 'status': 'FAILURE', 'error_message': str(e), 'summary': f"❌ HTTP error: {str(e)[:50]}..."})
        return base_result

def call_ai_api(full_url: str, api_key: str, messages: List[Dict[str, str]]) -> str:
    headers = {'Content-Type': 'application/json'}
    # We construct the URL with the key here for the request...
    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() # This is what usually triggers the error containing the URL
        data = response.json()
        if 'candidates' in data and len(data['candidates']) > 0:
            # Added a check to prevent key errors on malformed responses
            if 'parts' in data['candidates'][0]['content'] and len(data['candidates'][0]['content']['parts']) > 0:
                 return data['candidates'][0]['content']['parts'][0]['text']
            else:
                return f"ERROR: API malformed response, missing 'parts': {data}"
        else:
            return f"ERROR: API malformed response, missing 'candidates': {data}"
    except Exception as e:
        # SECURITY FIX: Sanitize the exception string before returning it
        sanitized_error = sanitize_error_message(str(e))
        return f"ERROR: API call failed: {sanitized_error}"

def resolve_model(selected_label: str, model_registry: Dict, fallback_label: str, agent_name: str) -> Tuple[str, str, str]:
    model_config = model_registry.get(selected_label) or model_registry.get(fallback_label)
    if model_config:
        return model_config['model_id'], model_config['url'], f"{agent_name} model resolved: {selected_label or fallback_label}"
    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: Using first available {agent_name} model: {first_label}"
    return '', '', f"FATAL: No {agent_name} models configured."

# --- ENGINE ---
class ConductorEngine:
    def __init__(self, coder_model_label: str, strategist_model_label: str, enable_audit: bool = False):
        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 = 6
        self.enable_audit: bool = enable_audit
        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 = apply_architect_signature(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.")
        final_code = apply_architect_signature(coder_response)
        return self._compile_results(final_code)

    def execute_workflow_streaming(self, blueprint: str, intent: str):
        """STREAMING version with Architect's Signature"""
        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})")
        if self.enable_audit:
            self.log_system("AUDIT MODE: ENABLED")
        yield self._current_state()
        
        initial_prompt = format_initial_prompt(blueprint)
        self.log_coder("user", initial_prompt)
        yield self._current_state()
        
        coder_response = ""
        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:
                if not self.enable_audit:
                    self.log_system("MISSION COMPLETE: Code generated.")
                    final_state = self._current_state()
                    final_state["final_code"] = apply_architect_signature(coder_response)
                    yield final_state
                    return

                self.log_system("AUDIT PHASE INITIATED.")
                yield self._current_state()
                
                audit_prompt = format_strategist_audit_prompt(blueprint, intent, coder_response)
                self.log_strategist("user", audit_prompt)
                yield self._current_state()
                
                audit_response = call_ai_api(self.strategist_url, CONFIG['STRATEGIST_API_KEY'], self.strategist_history)
                self.log_strategist("assistant", audit_response)
                yield self._current_state()
                
                if "AUDIT_PASSED" in audit_response:
                    self.log_system("AUDIT PASSED. MISSION COMPLETE.")
                    final_state = self._current_state()
                    final_state["final_code"] = apply_architect_signature(coder_response)
                    yield final_state
                    return
                
                self.log_system("AUDIT FAILED. Generating revisions.")
                yield self._current_state()
                
                revision_prompt = format_resolution_prompt(blueprint, f"AUDIT CHANGES REQUIRED:\n{audit_response}")
                self.log_coder("user", revision_prompt)
                yield self._current_state()
                continue
            
            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.")
        final_state = self._current_state()
        final_state["final_code"] = apply_architect_signature(coder_response)
        yield final_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
    enable_audit: Optional[bool] = False

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, req.enable_audit)
    result = engine.execute_workflow(req.blueprint, req.intent)
    return result

@app.post("/execute_stream")
async def execute_conductor_stream(req: ExecutionRequest):
    """STREAMING endpoint with audit support"""
    
    async def generate_updates():
        engine = ConductorEngine(req.coder_model_label, req.strategist_model_label, req.enable_audit)
        
        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)