main.py

from typing import AsyncIterator
from contextlib import asynccontextmanager
import httpx
from fastmcp import FastMCP, Context
from loguru import logger
from dataclasses import dataclass
import uuid
import os
import sys

from utils.config import API_CONFIG
from utils.model import MultiModuleRequest, SingleModuleRequest, SearchSpaceRoutingRequest

from fastmcp.server.dependencies import get_http_request
from starlette.requests import Request


from dotenv import load_dotenv

load_dotenv()


@dataclass
class DSContext:
    """Typed context for DS MCP server"""

    client: httpx.AsyncClient


@asynccontextmanager
async def ds_lifespan(server: FastMCP) -> AsyncIterator[DSContext]:
    client = httpx.AsyncClient(timeout=API_CONFIG["timeout"])
    ctx = DSContext(client=client)

    try:
        yield ctx
    finally:
        await client.aclose()


mcp = FastMCP(name="defect-solver-server", lifespan=ds_lifespan)


@mcp.tool(
    name="multi_module_bug_localization",
    description=(
        """
        Performs bug localization across the entire microservice architecture by first identifying likely microservices
        (search spaces) and then locating the bug-related files within them.

        Use this tool when you do not know which microservice is responsible for the bug.

        Args:
            request (MultiModuleRequest): Issue key and descriptive details from the bug report.
            ctx (Context): MCP context object.

        Returns:
            dict: Includes selected microservices and localized file paths.
        """
    ),
)
async def multi_module_bug_localization(request: MultiModuleRequest, ctx: Context) -> dict:
    ds_ctx: DSContext = ctx.request_context.lifespan_context
    client = ds_ctx.client

    api_url = API_CONFIG["api_base_url"] + API_CONFIG["api_multimodule_endpoint"]

    # hf_token to access private HF Space serving Defect Solver API
    hf_token = API_CONFIG.get("hf_access_token", None)

    # Get per-user API key from MCP Context
    user_request: Request = get_http_request()
    logger.info("Received request headers: {}".format(user_request.headers))
    defect_solver_api_key = user_request.headers.get("DS-API-Key", None)
    logger.info("defect_solver_api_key: {}".format(defect_solver_api_key) )
    logger.info(
        "Using API key: {}".format(defect_solver_api_key)
        if defect_solver_api_key
        else "No Defect Solver API key provided"
    )

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {hf_token}" if hf_token else None,
        "DS-API-Key": defect_solver_api_key,
    }
    headers = {k: v for k, v in headers.items() if v is not None}

    try:
        # Generate a UUID for issue_key even if provided (overwrite)
        request.issue_key = str(uuid.uuid4())

        logger.info(
            f"Sending multi-module bug localization request to {api_url} with data: {str(request.model_dump())[:300]}..."
        )

        response = await client.post(
            api_url,
            json={
                "key": request.issue_key,
                "fields": {
                    "summary": request.summary,
                    "description": request.description,
                },
            },
            headers=headers,
        )
        response.raise_for_status()

        logger.info(f"Received response: {str(response.text)[:300]}...")
        logger.debug(f"Response JSON: {response.json()}")

        return {"issue_key": request.issue_key, "result": response.json()}

    except httpx.HTTPStatusError as e:
        logger.error(
            f"HTTP error {e.response.status_code} from multimodule bug localization endpoint: {e.response.text}"
        )
        return {"error": f"HTTP error: {e.response.status_code}", "details": e.response.text}
    except httpx.RequestError as e:
        logger.error(f"Request error while calling multimodule bug localization endpoint: {e}")
        return {"error": "Request error", "details": str(e)}
    except Exception as e:
        logger.error("Unexpected error during multimodule bug localization request")
        return {"error": "Unexpected error", "details": str(e)}


@mcp.tool(
    name="single_module_bug_localization",
    description=(
        """
        Performs bug localization within a specific module (microservice) of a microservice architecture.

        Use this tool when you already know the module where the bug is likely located.

        Args:
            request (SingleModuleRequest): Bug description, issue key, and the known module name.
            ctx (Context): The MCP context containing the lifespan context.

        Returns:
            dict: Includes the selected module and localized file paths.
        """
    ),
)
async def single_module_bug_localization(request: SingleModuleRequest, ctx: Context) -> dict:
    ds_ctx: DSContext = ctx.request_context.lifespan_context
    client = ds_ctx.client

    api_url = API_CONFIG["api_base_url"] + API_CONFIG["api_singlemodule_endpoint"]

    # hf token to access Defect Solver API
    hf_token = API_CONFIG.get("hf_access_token", None)

    # Get per-user API key from MCP Context
    user_request: Request = get_http_request()
    defect_solver_api_key = user_request.headers.get("DS-API-Key", None)
    logger.info(
        "Using API key: {}".format(defect_solver_api_key)
        if defect_solver_api_key
        else "No Defect Solver API key provided"
    )

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {hf_token}" if hf_token else None,
        "DS-API-Key": defect_solver_api_key,
    }
    headers = {k: v for k, v in headers.items() if v is not None}

    try:
        # Generate a UUID for issue_key even if provided (overwrite)
        request.issue_key = str(uuid.uuid4())

        logger.info(
            f"Sending single-module bug localization request to {api_url} with data: {str(request.model_dump())[:300]}..."
        )

        response = await client.post(
            api_url,
            json={
                "key": request.issue_key,
                "fields": {
                    "summary": request.summary,
                    "description": request.description,
                    "module": request.module,
                },
            },
            headers=headers,
        )
        response.raise_for_status()

        logger.info(f"Received response: {str(response.text)[:300]}...")
        logger.debug(f"Response JSON: {response.json()}")

        return {"issue_key": request.issue_key, "result": response.json()}

    except httpx.HTTPStatusError as e:
        logger.error(f"HTTP error {e.response.status_code} from endpoint: {e.response.text}")
        return {"error": f"HTTP error: {e.response.status_code}", "details": e.response.text}
    except httpx.RequestError as e:
        logger.error(f"Request error while calling endpoint: {e}")
        return {"error": "Request error", "details": str(e)}
    except Exception as e:
        logger.error("Unexpected error during single-module bug localization")
        return {"error": "Unexpected error", "details": str(e)}


@mcp.tool(
    name="search_space_routing",
    description=(
        """
        Identifies and returns the most likely microservices (search spaces) that could be the source of a reported bug.

        Use this tool when the source of the bug is unknown and you want to narrow down the investigation to top candidate microservices.

        Args:
            request (SearchSpaceRoutingRequest): Bug description and optional issue key/summary.
            ctx (Context): The MCP context with lifespan context.

        Returns:
            dict: A message and the list of selected search spaces (microservices).
        """
    ),
)
async def search_space_routing(request: SearchSpaceRoutingRequest, ctx: Context) -> dict:
    ds_ctx: DSContext = ctx.request_context.lifespan_context
    client = ds_ctx.client

    api_url = API_CONFIG["api_base_url"] + API_CONFIG["api_searchspace_endpoint"]

    # hf_token to access private HF Space serving Defect Solver API
    hf_token = API_CONFIG.get("hf_access_token", None)

    # Get per-user API key from MCP Context
    user_request: Request = get_http_request()
    defect_solver_api_key = user_request.headers.get("DS-API-Key", None)
    logger.info(
        "Using API key: {}".format(defect_solver_api_key)
        if defect_solver_api_key
        else "No Defect Solver API key provided"
    )

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {hf_token}" if hf_token else None,
        "DS-API-Key": defect_solver_api_key,
    }
    headers = {k: v for k, v in headers.items() if v is not None}

    try:
        # Generate a UUID for issue_key even if provided (overwrite)
        request.issue_key = str(uuid.uuid4())

        logger.info(
            f"Sending search space routing request to {api_url} with data: {str(request.model_dump())[:300]}..."
        )

        response = await client.post(
            api_url,
            json={
                "key": request.issue_key,
                "fields": {
                    "summary": request.summary,
                    "description": request.description,
                },
            },
            headers=headers,
        )
        response.raise_for_status()

        logger.info(f"Received response: {str(response.text)[:300]}...")
        logger.debug(f"Response JSON: {response.json()}")

        return {"issue_key": request.issue_key, "result": response.json()}

    except httpx.HTTPStatusError as e:
        logger.error(f"HTTP error {e.response.status_code} from endpoint: {e.response.text}")
        return {"error": f"HTTP error: {e.response.status_code}", "details": e.response.text}
    except httpx.RequestError as e:
        logger.error(f"Request error while calling endpoint: {e}")
        return {"error": "Request error", "details": str(e)}
    except Exception as e:
        logger.error("Unexpected error during search space routing")
        return {"error": "Unexpected error", "details": str(e)}


@mcp.prompt(title="Select Bug Localization Tool")
async def prompt_select_tool() -> str:
    return """
You have access to bug localization tools in this MCP server. Determine the best approach for the current bug based on the description provided in this conversation.

Recommend:
1. **Primary Tool** - Which tool to use first and why
2. **Tool Sequence** - If multiple tools needed, in what order
3. **Reasoning** - Why this approach is optimal
4. **Expected Workflow** - Step-by-step tool usage plan

Provide specific tool recommendations with rationale.
"""


@mcp.prompt(title="Localize and Find Bug Using Selected Tool")
async def prompt_find_bug() -> str:
    return """
You are using a bug localization tool to identify potential files related to the bug described in this conversation. Execute the following steps:

1. **Input the augmented bug description** - Use the detailed description from the previous step
2. **Run the localization tool** - Execute the tool to find areas related to the bug
3. **Collect results** - Gather the results identified by the tool
4. **Return results** - Provide the results and any additional metadata

The tool will return info that are likely related to the bug.
"""


@mcp.prompt(title="Explain Localization Results")
async def prompt_explain() -> str:
    return """
You received a response from a bug localization tool in this conversation. Interpret and explain the results for actionable next steps.

Analyze and provide:
1. **Key Files Identified** - Prioritized list of files to investigate
2. **Investigation Order** - Which files to check first and why  
3. **Code Patterns to Look For** - Specific methods, classes, or patterns
4. **Next Tool Recommendations** - Should you use single_module on specific modules?
5. **Confidence Assessment** - How reliable are these results?

Provide actionable guidance for the developer's next steps.
"""


@mcp.prompt(title="Fix Localized Bug")
async def prompt_fix_bug() -> str:
    return """
You have codebase access. A bug localization tool has identified files as potential starting points in this conversation.

Instructions:
1. **Directly fix the bug** - Use the identified files as entry points, but your goal is to discover the root cause and apply the necessary code changes.
2. **Explore the codebase** - Investigate related code, dependencies, and connections as needed.
3. **Make concrete changes** - Show the exact code modifications required to resolve the bug, with before/after code snippets.
4. **Do not just analyze** - You must provide and apply the fix, not just discuss possible causes.
5. **Suggest verification** - Briefly describe how to test and confirm the fix works.

Your task is to discover, modify, and fix the bug in the codebase using the provided starting points.

NOTE: If you cannot find the bug and fix it, provide a brief explanation of your investigation process and why you could not resolve it.
"""


@mcp.prompt(title="Full Bug Resolution Workflow")
async def prompt_full_workflow() -> str:
    return """
You are a bug resolution expert with access to this codebase and bug localization tools. Complete the full workflow from bug description to fix.

Follow this workflow:

1. **AUGMENT** - First, examine the codebase and augment the bug description with technical details, specific modules, file types, and architectural context.

2. **ROUTE** - Decide which bug localization tool to use:
   - search_space_routing: If you don't know which microservices are involved
   - multi_module_bug_localization: If multiple modules likely affected
   - single_module_bug_localization: If you know the specific module

3. **LOCALIZE** - Call the appropriate tool(s) with the augmented description.

4. **INTERPRET** - Analyze the localization results to identify key files and investigation priorities.

5. **FIX** - Examine the identified files (as starting points), find the actual bug, and provide concrete code fixes.

Execute each step and provide the final resolution with specific code changes.

===== Bug Description =====
<insert_here>
"""


@mcp.prompt(title="Augment Bug report with Technical Details")
async def prompt_augment_bug_report() -> str:
    return """
You are a bug localization expert with access to this codebase. Your task is to augment the bug description from this conversation with technical details that help pinpoint the bug location.

Steps:
1. **Analyze the codebase** - Examine project structure, modules, and architectural patterns
2. **Map bug to code** - Identify likely affected areas based on the bug symptoms
3. **Augment the description** - Add:
   - Specific module/package/class names that might be involved
   - File types, naming patterns, and directory paths
   - Technical keywords, error types, and exception names
   - Affected architectural layers and components
   - Cross-cutting concerns and dependencies
   - External system interactions or API calls
   - Component relationships that could cause the bug

Return a 2-3x more detailed bug report with precise technical terms for effective bug localization, while preserving the original description.

===== Bug Description =====
<insert_here>
"""


@mcp.prompt(title="Revise Bug Report")
async def prompt_revise_bug_report() -> str:
    return """
Enhance the bug description from this conversation by adding relevant technical context that supports precise root cause analysis.
Return a 2-3x more detailed/enhanced bug report with precise technical terms for effective bug localization, while preserving the original description.

===== Bug Description =====
<insert_here>
"""


if __name__ == "__main__":
    TRANSPORT_MODE = os.getenv("TRANSPORT_MODE", "http")
    HOST = os.getenv("HOST", "0.0.0.0")
    PORT = int(os.getenv("PORT", "7860"))
    LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO")

    # Configure logger
    logger.remove()
    logger.add(sys.stderr, level=LOG_LEVEL, format="{level}:\t\t{time:YYYY-MM-DD HH:mm:ss} - {message}")

    mcp.run(transport=TRANSPORT_MODE, host=HOST, port=PORT, log_level=LOG_LEVEL)