main.py

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

from utils.config import API_CONFIG
from utils.model import MultiModuleRequest, SingleModuleRequest, SearchSpaceRoutingRequest
from utils.auth import DSContext, authenticate

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


from dotenv import load_dotenv

load_dotenv()

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

    if token:
        client.headers.update({"Authorization": f"Bearer {token}"})

    ctx = DSContext(client=client, access_token=token)

    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()

    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 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)