golf-mcp 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

golf/core/builder.py

@@ -10,7 +10,6 @@ from typing import Any, Dict, List, Optional, Set
 import black
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn
-from rich.panel import Panel
 
 from golf.core.config import Settings
 from golf.core.parser import (
@@ -21,10 +20,8 @@ from golf.core.parser import (
 from golf.core.transformer import transform_component
 from golf.core.builder_auth import generate_auth_code, generate_auth_routes
 from golf.auth import get_auth_config
-from golf.auth import get_access_token
 from golf.core.builder_telemetry import (
-    generate_otel_lifespan_code, 
-    generate_otel_instrumentation_code, 
+    generate_telemetry_imports,
     get_otel_dependencies
 )
 
@@ -513,24 +510,40 @@ class CodeGenerator:
         """Generate the main server entry point."""
         server_file = self.output_dir / "server.py"
         
+        # Get auth components
+        provider_config, _ = get_auth_config()
+        auth_components = generate_auth_code(
+            server_name=self.settings.name,
+            host=self.settings.host,
+            port=self.settings.port,
+            https=False,  # This could be configurable in settings
+            opentelemetry_enabled=self.settings.opentelemetry_enabled,
+            transport=self.settings.transport
+        )
+        
         # Create imports section
         imports = [
             "from fastmcp import FastMCP",
             "import os",
             "import sys",
             "from dotenv import load_dotenv",
+            "import logging",
+            "",
+            "# Suppress FastMCP INFO logs",
+            "logging.getLogger('fastmcp').setLevel(logging.WARNING)",
+            "logging.getLogger('mcp').setLevel(logging.WARNING)",
             ""
         ]
         
-        # For imports
+        # Add auth imports if auth is configured
+        if auth_components.get("has_auth"):
+            imports.extend(auth_components["imports"])
+            imports.append("")
+        
+        # Add OpenTelemetry imports if enabled
         if self.settings.opentelemetry_enabled:
-            imports.extend([
-                "# OpenTelemetry imports",
-                "from opentelemetry.instrumentation.asgi import OpenTelemetryMiddleware",
-                "from starlette.middleware import Middleware",
-                # otel_lifespan function will be defined from generate_otel_lifespan_code
-            ])
-        imports.append("") # Add blank line after all component type imports or OTel imports
+            imports.extend(generate_telemetry_imports())
+            imports.append("")
         
         # Add imports section for different transport methods
         if self.settings.transport == "sse":
@@ -539,14 +552,6 @@ class CodeGenerator:
         elif self.settings.transport != "stdio":
             imports.append("import uvicorn")
                 
-        # Create a new FastMCP instance for the server
-        server_code_lines = ["# Create FastMCP server"]
-        mcp_constructor_args = [f'"{self.settings.name}"']
-
-        mcp_instance_line = f"mcp = FastMCP({', '.join(mcp_constructor_args)})"
-        server_code_lines.append(mcp_instance_line)
-        server_code_lines.append("")
-        
         # Get transport-specific configuration
         transport_config = self._get_transport_config(self.settings.transport)
         endpoint_path = transport_config["endpoint_path"]
@@ -616,62 +621,78 @@ class CodeGenerator:
                 imports.append(f"import {full_module_path}")
                 
                 # Add code to register this component
-                if component_type == ComponentType.TOOL:
-                    registration = f"# Register the tool '{component.name}' from {full_module_path}"
+                if self.settings.opentelemetry_enabled:
+                    # Use telemetry instrumentation
+                    registration = f"# Register the {component_type.value} '{component.name}' with telemetry"
+                    entry_func = component.entry_function if hasattr(component, "entry_function") and component.entry_function else "export"
                     
-                    # Use the entry_function if available, otherwise try the export variable
-                    if hasattr(component, "entry_function") and component.entry_function:
-                        registration += f"\nmcp.add_tool({full_module_path}.{component.entry_function}"
-                    else:
-                        registration += f"\nmcp.add_tool({full_module_path}.export"
+                    # Debug: Add logging to verify wrapping
+                    registration += f"\n_wrapped_func = instrument_{component_type.value}({full_module_path}.{entry_func}, '{component.name}')"
                     
-                    # Add the name parameter
-                    registration += f", name=\"{component.name}\""
-                    
-                    # Add description from docstring
-                    if component.docstring:
-                        # Escape any quotes in the docstring
-                        escaped_docstring = component.docstring.replace("\"", "\\\"")
-                        registration += f", description=\"{escaped_docstring}\""
-                    registration += ")"
+                    if component_type == ComponentType.TOOL:
+                        registration += f"\nmcp.add_tool(_wrapped_func, name=\"{component.name}\", description=\"{component.docstring or ''}\")"
+                    elif component_type == ComponentType.RESOURCE:
+                        registration += f"\nmcp.add_resource_fn(_wrapped_func, uri=\"{component.uri_template}\", name=\"{component.name}\", description=\"{component.docstring or ''}\")"
+                    else:  # PROMPT
+                        registration += f"\nmcp.add_prompt(_wrapped_func, name=\"{component.name}\", description=\"{component.docstring or ''}\")"
+                else:
+                    # Standard registration without telemetry
+                    if component_type == ComponentType.TOOL:
+                        registration = f"# Register the tool '{component.name}' from {full_module_path}"
+                        
+                        # Use the entry_function if available, otherwise try the export variable
+                        if hasattr(component, "entry_function") and component.entry_function:
+                            registration += f"\nmcp.add_tool({full_module_path}.{component.entry_function}"
+                        else:
+                            registration += f"\nmcp.add_tool({full_module_path}.export"
+                        
+                        # Add the name parameter
+                        registration += f", name=\"{component.name}\""
+                        
+                        # Add description from docstring
+                        if component.docstring:
+                            # Escape any quotes in the docstring
+                            escaped_docstring = component.docstring.replace("\"", "\\\"")
+                            registration += f", description=\"{escaped_docstring}\""
+                        registration += ")"
 
-                elif component_type == ComponentType.RESOURCE:
-                    registration = f"# Register the resource '{component.name}' from {full_module_path}"
-                    
-                    # Use the entry_function if available, otherwise try the export variable
-                    if hasattr(component, "entry_function") and component.entry_function:
-                        registration += f"\nmcp.add_resource_fn({full_module_path}.{component.entry_function}, uri=\"{component.uri_template}\""
-                    else:
-                        registration += f"\nmcp.add_resource_fn({full_module_path}.export, uri=\"{component.uri_template}\""
-                    
-                    # Add the name parameter
-                    registration += f", name=\"{component.name}\""
+                    elif component_type == ComponentType.RESOURCE:
+                        registration = f"# Register the resource '{component.name}' from {full_module_path}"
                         
-                    # Add description from docstring
-                    if component.docstring:
-                        # Escape any quotes in the docstring
-                        escaped_docstring = component.docstring.replace("\"", "\\\"")
-                        registration += f", description=\"{escaped_docstring}\""
-                    registration += ")"
+                        # Use the entry_function if available, otherwise try the export variable
+                        if hasattr(component, "entry_function") and component.entry_function:
+                            registration += f"\nmcp.add_resource_fn({full_module_path}.{component.entry_function}, uri=\"{component.uri_template}\""
+                        else:
+                            registration += f"\nmcp.add_resource_fn({full_module_path}.export, uri=\"{component.uri_template}\""
+                        
+                        # Add the name parameter
+                        registration += f", name=\"{component.name}\""
+                            
+                        # Add description from docstring
+                        if component.docstring:
+                            # Escape any quotes in the docstring
+                            escaped_docstring = component.docstring.replace("\"", "\\\"")
+                            registration += f", description=\"{escaped_docstring}\""
+                        registration += ")"
 
-                else:  # PROMPT
-                    registration = f"# Register the prompt '{component.name}' from {full_module_path}"
-                    
-                    # Use the entry_function if available, otherwise try the export variable
-                    if hasattr(component, "entry_function") and component.entry_function:
-                        registration += f"\nmcp.add_prompt({full_module_path}.{component.entry_function}"
-                    else:
-                        registration += f"\nmcp.add_prompt({full_module_path}.export"
-                    
-                    # Add the name parameter
-                    registration += f", name=\"{component.name}\""
+                    else:  # PROMPT
+                        registration = f"# Register the prompt '{component.name}' from {full_module_path}"
                         
-                    # Add description from docstring
-                    if component.docstring:
-                        # Escape any quotes in the docstring
-                        escaped_docstring = component.docstring.replace("\"", "\\\"")
-                        registration += f", description=\"{escaped_docstring}\""
-                    registration += ")"
+                        # Use the entry_function if available, otherwise try the export variable
+                        if hasattr(component, "entry_function") and component.entry_function:
+                            registration += f"\nmcp.add_prompt({full_module_path}.{component.entry_function}"
+                        else:
+                            registration += f"\nmcp.add_prompt({full_module_path}.export"
+                        
+                        # Add the name parameter
+                        registration += f", name=\"{component.name}\""
+                            
+                        # Add description from docstring
+                        if component.docstring:
+                            # Escape any quotes in the docstring
+                            escaped_docstring = component.docstring.replace("\"", "\\\"")
+                            registration += f", description=\"{escaped_docstring}\""
+                        registration += ")"
                 
                 component_registrations.append(registration)
             
@@ -688,21 +709,37 @@ class CodeGenerator:
             ""
         ]
 
-        # After env_section, add OpenTelemetry lifespan code
-        otel_definitions_code = []
-        otel_instrumentation_application_code = [] # For instrumentation that runs after mcp is set up
+        # OpenTelemetry setup code will be handled through imports and lifespan
 
+        # Add auth setup code if auth is configured
+        auth_setup_code = []
+        if auth_components.get("has_auth"):
+            auth_setup_code = auth_components["setup_code"]
+
+        # Create FastMCP instance section
+        server_code_lines = ["# Create FastMCP server"]
+        
+        # Build FastMCP constructor arguments
+        mcp_constructor_args = [f'"{self.settings.name}"']
+        
+        # Add auth arguments if configured
+        if auth_components.get("has_auth") and auth_components.get("fastmcp_args"):
+            for key, value in auth_components["fastmcp_args"].items():
+                mcp_constructor_args.append(f"{key}={value}")
+        
+        # Add OpenTelemetry parameters if enabled
         if self.settings.opentelemetry_enabled:
-            otel_definitions_code.append(generate_otel_lifespan_code(
-                default_exporter=self.settings.opentelemetry_default_exporter,
-                project_name=self.settings.name
-            ))
-            otel_definitions_code.append("")  # Add blank line
+            mcp_constructor_args.append("lifespan=telemetry_lifespan")
 
-            # Prepare instrumentation code to be added after component registration
-            otel_instrumentation_application_code.append("# Apply OpenTelemetry Instrumentation")
-            otel_instrumentation_application_code.append(generate_otel_instrumentation_code())
-            otel_instrumentation_application_code.append("")
+        mcp_instance_line = f"mcp = FastMCP({', '.join(mcp_constructor_args)})"
+        server_code_lines.append(mcp_instance_line)
+        server_code_lines.append("")
+        
+        # Add any post-init code from auth
+        post_init_code = []
+        if auth_components.get("has_auth") and auth_components.get("post_init_code"):
+            post_init_code.extend(auth_components["post_init_code"])
+            post_init_code.append("")
 
         # Main entry point with transport-specific app initialization
         main_code = [
@@ -729,34 +766,42 @@ class CodeGenerator:
         if self.settings.transport == "sse":
             main_code.extend([
                 "    # For SSE, FastMCP's run method handles auth integration better",
-                "    print(f\"[Server Runner] Using mcp.run() for SSE transport with host={host}, port={port}\", file=sys.stderr)",
-                "    mcp.run(transport=\"sse\", host=host, port=port, log_level=\"debug\")"
+                "    mcp.run(transport=\"sse\", host=host, port=port, log_level=\"info\")"
             ])
         elif self.settings.transport == "streamable-http":
             main_code.extend([
                 "    # Create HTTP app and run with uvicorn",
-                "    print(f\"[Server Runner] Starting streamable-http transport with host={host}, port={port}\", file=sys.stderr)",
                 "    app = mcp.http_app()",
-                "    uvicorn.run(app, host=host, port=port, log_level=\"debug\")"
+            ])
+            
+            # Add OpenTelemetry middleware to the HTTP app if enabled
+            if self.settings.opentelemetry_enabled:
+                main_code.extend([
+                    "    # Apply OpenTelemetry middleware to the HTTP app",
+                    "    from opentelemetry.instrumentation.asgi import OpenTelemetryMiddleware",
+                    "    app = OpenTelemetryMiddleware(app)",
+                ])
+            
+            main_code.extend([
+                "    uvicorn.run(app, host=host, port=port, log_level=\"info\")"
             ])
         else:
             # For stdio transport, use mcp.run()
             main_code.extend([
                 "    # Run with stdio transport",
-                "    print(f\"[Server Runner] Starting stdio transport\", file=sys.stderr)",
                 "    mcp.run(transport=\"stdio\")"
             ])
         
-        # Combine all sections - move env_section to right location
-        # Order: imports, env_section, otel_definitions (lifespan func), server_code (mcp init), 
-        # component_registrations, otel_instrumentation (wrappers), main_code (run block)
+        # Combine all sections
+        # Order: imports, env_section, auth_setup, server_code (mcp init), 
+        # post_init (API key middleware), component_registrations, main_code (run block)
         code = "\n".join(
             imports + 
             env_section + 
-            otel_definitions_code + 
-            server_code_lines +  # Add back the server_code_lines with constructor
-            component_registrations + 
-            otel_instrumentation_application_code + # Added instrumentation here
+            auth_setup_code +
+            server_code_lines +
+            post_init_code +
+            component_registrations +
             main_code
         )
         
@@ -826,15 +871,64 @@ def build_project(
     if output_dir.exists():
         shutil.rmtree(output_dir)
     output_dir.mkdir(parents=True, exist_ok=True) # Ensure output_dir exists after clearing
-    
-    # If dev build and copy_env flag is true, copy .env file from project root to output_dir
-    if copy_env: # The build_env string ('dev'/'prod') check can be done in the CLI layer that sets copy_env
+
+    # --- BEGIN Enhanced .env handling ---
+    env_vars_to_write = {}
+    env_file_path = output_dir / ".env"
+
+    # 1. Load from existing project .env if copy_env is true
+    if copy_env:
         project_env_file = project_path / ".env"
         if project_env_file.exists():
             try:
-                shutil.copy(project_env_file, output_dir / ".env")
+                from dotenv import dotenv_values
+                env_vars_to_write.update(dotenv_values(project_env_file))
+            except ImportError:
+                console.print("[yellow]Warning: python-dotenv is not installed. Cannot read existing .env file for rich merging. Copying directly.[/yellow]")
+                try:
+                    shutil.copy(project_env_file, env_file_path)
+                    # If direct copy happens, re-read for step 2 & 3 to respect its content
+                    if env_file_path.exists():
+                         from dotenv import dotenv_values
+                         env_vars_to_write.update(dotenv_values(env_file_path)) # Read what was copied
+                except Exception as e:
+                    console.print(f"[yellow]Warning: Could not copy project .env file: {e}[/yellow]")
             except Exception as e:
-                console.print(f"[yellow]Warning: Could not copy project .env file: {e}[/yellow]")
+                console.print(f"[yellow]Warning: Error reading project .env file content: {e}[/yellow]")
+
+
+    # 2. Apply Golf's OTel default exporter setting if OTEL_TRACES_EXPORTER is not already set
+    if settings.opentelemetry_enabled and settings.opentelemetry_default_exporter:
+        if "OTEL_TRACES_EXPORTER" not in env_vars_to_write:
+            env_vars_to_write["OTEL_TRACES_EXPORTER"] = settings.opentelemetry_default_exporter
+            console.print(f"[info]Setting OTEL_TRACES_EXPORTER to '{settings.opentelemetry_default_exporter}' from golf.json in built app's .env[/info]")
+
+    # 3. Apply Golf's project name as OTEL_SERVICE_NAME if not already set
+    # (Ensures service name defaults to project name if not specified in user's .env)
+    if settings.opentelemetry_enabled and settings.name:
+        if "OTEL_SERVICE_NAME" not in env_vars_to_write:
+            env_vars_to_write["OTEL_SERVICE_NAME"] = settings.name
+    
+    # 4. (Re-)Write the .env file in the output directory if there's anything to write
+    if env_vars_to_write:
+        try:
+            with open(env_file_path, "w") as f:
+                for key, value in env_vars_to_write.items():
+                    # Ensure values are properly quoted if they contain spaces or special characters
+                    # and handle existing quotes within the value.
+                    if isinstance(value, str):
+                        # Replace backslashes first, then double quotes
+                        processed_value = value.replace('\\', '\\\\') # Escape backslashes
+                        processed_value = processed_value.replace('"', '\\"') # Escape double quotes
+                        if ' ' in value or '#' in value or '\n' in value or '"' in value or "'" in value:
+                            f.write(f'{key}="{processed_value}"\n')
+                        else:
+                            f.write(f"{key}={processed_value}\n")
+                    else: # For non-string values, write directly
+                        f.write(f"{key}={value}\n")
+        except Exception as e:
+            console.print(f"[yellow]Warning: Could not write .env file to output directory: {e}[/yellow]")
+    # --- END Enhanced .env handling ---
 
     # Show what we're building, with environment info
     console.print(f"[bold]Building [green]{settings.name}[/green] ({build_env} environment)[/bold]")
@@ -918,11 +1012,12 @@ dependencies = [
 
 from golf.auth.provider import ProviderConfig
 from golf.auth.oauth import GolfOAuthProvider, create_callback_handler
-from golf.auth.helpers import get_access_token, get_provider_token, extract_token_from_header
+from golf.auth.helpers import get_access_token, get_provider_token, extract_token_from_header, get_api_key, set_api_key
+from golf.auth.api_key import configure_api_key, get_api_key_config
 """)
     
     # Copy provider, oauth, and helper modules
-    for module in ["provider.py", "oauth.py", "helpers.py"]:
+    for module in ["provider.py", "oauth.py", "helpers.py", "api_key.py"]:
         src_file = Path(__file__).parent.parent.parent / "golf" / "auth" / module
         dst_file = auth_dir / module
         
@@ -931,197 +1026,57 @@ from golf.auth.helpers import get_access_token, get_provider_token, extract_toke
         else:
             console.print(f"[yellow]Warning: Could not find {src_file} to copy[/yellow]")
     
-    # Now handle the auth integration if configured
-    if provider_config:
-        
-        # Generate the auth code to inject into server.py
-        # The existing call to generate_auth_code.
-        # We need to ensure the arguments passed are sensible.
-        # server.py determines issuer_url at runtime. generate_auth_code
-        # likely uses host/port/https to construct its own version or parts of it.
-        
-        # Determine protocol for https flag based on runtime logic similar to server.py
-        # This is a bit of a guess as settings doesn't explicitly store protocol for generate_auth_code
-        # A small inconsistency here if server.py's runtime logic for issuer_url differs significantly
-        # from what generate_auth_code expects/builds.
-        # For now, let's assume False is okay, or it's handled internally by generate_auth_code
-        # based on typical dev environments.
-        is_https_proto = False # Default, adjust if settings provide this info for build time
-
-        auth_code_str = generate_auth_code( # Renamed to auth_code_str to avoid confusion
-            server_name=settings.name,
-            host=settings.host,
-            port=settings.port
-        )
-    else:
-        # If auth is not configured, create a basic FastMCP instantiation string
-        # This string will then be processed for OTel args like the auth_code_str would be
-        auth_code_str = f"mcp = FastMCP('{settings.name}')" 
-
-    # ---- Centralized OpenTelemetry Argument Injection ----
+    # Copy telemetry module if OpenTelemetry is enabled
     if settings.opentelemetry_enabled:
-        temp_mcp_lines = auth_code_str.split('\n')
-        final_mcp_lines = []
-        otel_args_injected = False
-        for line_content in temp_mcp_lines:
-            if "mcp = FastMCP(" in line_content and ")" in line_content and not otel_args_injected:
-                open_paren_pos = line_content.find("(")
-                close_paren_pos = line_content.rfind(")")
-                if open_paren_pos != -1 and close_paren_pos != -1 and open_paren_pos < close_paren_pos:
-                    existing_args_str = line_content[open_paren_pos+1:close_paren_pos].strip()
-                    otel_args_to_add = []
-                    if "lifespan=" not in existing_args_str:
-                        otel_args_to_add.append("lifespan=otel_lifespan")
-                    if settings.transport != "stdio" and "middleware=" not in existing_args_str:
-                        otel_args_to_add.append("middleware=[Middleware(OpenTelemetryMiddleware)]")
-                    
-                    if otel_args_to_add:
-                        new_args_str = existing_args_str
-                        if new_args_str and not new_args_str.endswith(','):
-                            new_args_str += ", "
-                        new_args_str += ", ".join(otel_args_to_add)
-                        new_line = f"{line_content[:open_paren_pos+1]}{new_args_str}{line_content[close_paren_pos:]}"
-                        final_mcp_lines.append(new_line)
-                        otel_args_injected = True
-                        continue
-            final_mcp_lines.append(line_content)
+        telemetry_dir = output_dir / "golf" / "telemetry"
+        telemetry_dir.mkdir(parents=True, exist_ok=True)
         
-        if otel_args_injected:
-            auth_code_str = "\n".join(final_mcp_lines)
-        elif otel_args_to_add: # Only warn if we actually tried to add something
-            console.print(f"[yellow]Warning: Could not automatically inject OpenTelemetry lifespan/middleware into FastMCP constructor. Review server.py.[/yellow]")
-    # ---- END Centralized OpenTelemetry Argument Injection ----
-
-    # ---- MODIFICATION TO auth_code_str for _set_active_golf_oauth_provider (if auth was enabled) ----
-    if provider_config: # Only run this if auth was actually processed by generate_auth_code
+        # Copy telemetry __init__.py
+        src_init = Path(__file__).parent.parent.parent / "golf" / "telemetry" / "__init__.py"
+        dst_init = telemetry_dir / "__init__.py"
+        if src_init.exists():
+            shutil.copy(src_init, dst_init)
+        
+        # Copy instrumentation module
+        src_instrumentation = Path(__file__).parent.parent.parent / "golf" / "telemetry" / "instrumentation.py"
+        dst_instrumentation = telemetry_dir / "instrumentation.py"
+        if src_instrumentation.exists():
+            shutil.copy(src_instrumentation, dst_instrumentation)
+        else:
+            console.print("[yellow]Warning: Could not find telemetry instrumentation module[/yellow]")
+    
+    # Check if auth routes need to be added
+    provider_config, _ = get_auth_config()
+    if provider_config:
         auth_routes_code = generate_auth_routes()
-
+        
         server_file = output_dir / "server.py"
         if server_file.exists():
             with open(server_file, "r") as f:
                 server_code_content = f.read()
             
-            create_marker = '# Create FastMCP server'
-            # The original logic replaces the FastMCP instantiation part.
-            # So we use the modified auth_code_str here.
-            create_pos = server_code_content.find(create_marker)
-            if create_pos != -1: # Ensure marker is found
-                create_pos += len(create_marker) # Move past the marker text itself
-                create_next_line = server_code_content.find('\n', create_pos) + 1
-                # Assuming the original mcp = FastMCP(...) line is what auth_code_str replaces
-                # Find the end of the line that starts with "mcp = FastMCP("
-                mcp_line_start_search = server_code_content.find("mcp = FastMCP(", create_next_line)
-                if mcp_line_start_search != -1:
-                    mcp_line_end = server_code_content.find('\n', mcp_line_start_search)
-                    if mcp_line_end == -1: mcp_line_end = len(server_code_content) # if it's the last line
-
-                    modified_code = (
-                        server_code_content[:create_next_line] + 
-                        auth_code_str + # Use the modified auth code string
-                        server_code_content[mcp_line_end:]
-                    )
-                else: # Fallback if "mcp = FastMCP(" line isn't found as expected
-                    console.print(f"[yellow]Warning: Could not precisely find 'mcp = FastMCP(...)' line for replacement by auth_code in {server_file}. Appending auth_code instead.[/yellow]")
-                    # This part of the logic was to replace mcp = FastMCP(...)
-                    # If the generate_auth_code ALREADY includes the mcp = FastMCP(...) line,
-                    # then the original injection logic might be different.
-                    # The example server.py shows that the auth_code INCLUDES the mcp = FastMCP(...) line.
-                    # The original code in builder.py:
-                    # create_next_line = server_code.find('\n', create_pos) + 1
-                    # mcp_line_end = server_code.find('\n', create_next_line)
-                    # This implies it replaces ONE line after '# Create FastMCP server'
-                    # This needs to be robust. If auth_code_str contains the `mcp = FastMCP(...)` line itself,
-                    # then this replacement logic is correct.
-                    
-                    # The server.py example from `new/dist` implies that auth_code effectively *is* the
-                    # whole block from "import os" for auth settings down to and including "mcp = FastMCP(...)".
-                    # The original `_generate_server` creates a very minimal `mcp = FastMCP(...)`.
-                    # The `build_project` then overwrites this with the richer `auth_code` block.
-                    # Let's assume `auth_code_str_modified` should replace from `create_next_line` up to
-                    # where the original `mcp = FastMCP(...)` definition ended.
-
-                    # Re-evaluating the injection for auth_code_str_modified.
-                    # The server.py is first generated by _generate_server.
-                    # Then, if auth is enabled, this part of build_project MODIFIES it.
-                    # It finds '# Create FastMCP server', then replaces the *next line* (which is `mcp = FastMCP(...)` from _generate_server)
-                    # with the entire `auth_code_str_modified`.
-                    
-                    # Original line to find/replace: `mcp = FastMCP("{self.settings.name}")`
-                    # OR if telemetry was on `mcp = FastMCP("{self.settings.name}", lifespan=otel_lifespan)`
-                    # The replacement logic must be robust to find the line created by _generate_server
-                    
-                    # Let's find the line starting with "mcp = FastMCP(" that _generate_server created
-                    original_mcp_instantiation_pattern = "mcp = FastMCP("
-                    start_replace_idx = server_code_content.find(original_mcp_instantiation_pattern)
-                    
-                    if start_replace_idx != -1:
-                        # We need to find the complete statement, including any continuation lines
-                        line_start = server_code_content.rfind('\n', 0, start_replace_idx) + 1
-                        
-                        # Find the closing parenthesis, handling potential multi-line calls
-                        opening_paren_pos = server_code_content.find('(', start_replace_idx)
-                        if opening_paren_pos != -1:
-                            # Count open parentheses to handle nested ones correctly
-                            paren_count = 1
-                            pos = opening_paren_pos + 1
-                            while pos < len(server_code_content) and paren_count > 0:
-                                if server_code_content[pos] == '(':
-                                    paren_count += 1
-                                elif server_code_content[pos] == ')':
-                                    paren_count -= 1
-                                pos += 1
-                            
-                            closing_paren_pos = pos - 1 if paren_count == 0 else -1
-                            
-                            if closing_paren_pos != -1:
-                                # Find the end of the statement (newline after the closing parenthesis)
-                                next_newline = server_code_content.find('\n', closing_paren_pos)
-                                if next_newline != -1:
-                                    end_replace_idx = next_newline + 1
-                                else:
-                                    end_replace_idx = len(server_code_content)
-                                
-                                # Replace the entire statement with the auth code
-                                modified_code = (
-                                    server_code_content[:line_start] +
-                                    auth_code_str + 
-                                    server_code_content[end_replace_idx:]
-                                )
-                            else:
-                                console.print(f"[red]Error: Could not find closing parenthesis for FastMCP constructor in {server_file}. Auth injection may fail.[/red]")
-                                modified_code = server_code_content
-                        else:
-                            console.print(f"[red]Error: Could not find opening parenthesis for FastMCP constructor in {server_file}. Auth injection may fail.[/red]")
-                            modified_code = server_code_content
-            
-            else: # create_marker not found (This case should ideally not happen if _generate_server works)
-                console.print(f"[red]Could not find injection marker '{create_marker}' in {server_file}. Auth injection failed.[/red]")
-                modified_code = server_code_content # No change
-
+            # Add auth routes before the main block
             app_marker = 'if __name__ == "__main__":'
-            app_pos = modified_code.find(app_marker)
-            if app_pos != -1: # Ensure marker is found
+            app_pos = server_code_content.find(app_marker)
+            if app_pos != -1:
                 modified_code = (
-                    modified_code[:app_pos] + 
-                    auth_routes_code + "\n\n" + # Ensure auth routes are injected
-                    modified_code[app_pos:]
+                    server_code_content[:app_pos] + 
+                    auth_routes_code + "\n\n" +
+                    server_code_content[app_pos:]
                 )
+                
+                # Format with black before writing
+                try:
+                    final_code_to_write = black.format_str(modified_code, mode=black.Mode())
+                except Exception as e:
+                    console.print(f"[yellow]Warning: Could not format server.py after auth routes injection: {e}[/yellow]")
+                    final_code_to_write = modified_code
+                
+                with open(server_file, "w") as f:
+                    f.write(final_code_to_write)
             else:
                 console.print(f"[yellow]Warning: Could not find main block marker '{app_marker}' in {server_file} to inject auth routes.[/yellow]")
 
-            # Format with black before writing
-            try:
-                final_code_to_write = black.format_str(modified_code, mode=black.Mode())
-            except Exception as e:
-                console.print(f"[yellow]Warning: Could not format server.py after auth injection: {e}[/yellow]")
-                final_code_to_write = modified_code # Write unformatted if black fails
-
-            with open(server_file, "w") as f:
-                f.write(final_code_to_write)
-            
-        else: # server_file does not exist
-            console.print(f"[red]Error: {server_file} does not exist for auth modification. Ensure _generate_server runs first.[/red]")
-
 
 # Renamed function - was find_shared_modules
 def find_common_files(project_path: Path, components: Dict[ComponentType, List[ParsedComponent]]) -> Dict[str, Path]:
@@ -1161,9 +1116,13 @@ def build_import_map(project_path: Path, common_files: Dict[str, Path]) -> Dict[
         try:
             rel_to_component = dir_path.relative_to(component_type)
             # Create the new import path
-            new_path = f"components.{component_type}.{rel_to_component}".replace("/", ".")
-            # Fix any double dots
-            new_path = new_path.replace("..", ".")
+            if str(rel_to_component) == ".":
+                # This is at the root of the component type
+                new_path = f"components.{component_type}"
+            else:
+                # Replace path separators with dots
+                path_parts = str(rel_to_component).replace("\\", "/").split("/")
+                new_path = f"components.{component_type}.{'.'.join(path_parts)}"
             
             # Map both the directory and the common file
             orig_module = dir_path_str