diff --git a/config.py b/config.py index 54c7624..fb20c7c 100644 --- a/config.py +++ b/config.py @@ -1,5 +1,6 @@ import logging from pathlib import Path +from contextvars import ContextVar # --- Project Root --- # Get the directory where config.py is located @@ -8,6 +9,7 @@ ROOT_DIR = Path(__file__).parent.resolve() # --- Server & Logs --- HOST = "127.0.0.1" PORT = 8000 +request_host = ContextVar("request_host", default=f"{HOST}:{PORT}") LOG_LEVEL = "WARNING" # Options: "DEBUG", "INFO", "WARNING", "ERROR" LOG_FILE = str(ROOT_DIR / "debug.log") USER_AGENT = "Mozilla/5.0 (X11; Linux x86_64; rv:151.0) Gecko/20100101 Firefox/151.0" # Stealth User-Agent to bypass Wikipedia's bot detection diff --git a/main.py b/main.py index b753470..2e44d29 100644 --- a/main.py +++ b/main.py @@ -56,6 +56,10 @@ async def messages_endpoint(request): if not sid or sid not in mcp_logic.sessions: return Response("Session not found", status_code=404) + # Capture the host header to ensure image links work across the network + host_header = request.headers.get("host", f"{config.HOST}:{config.PORT}") + config.request_host.set(host_header) + try: body = await request.json() except Exception as e: @@ -116,7 +120,8 @@ app.add_middleware( ) if __name__ == "__main__": - config = uvicorn.Config( + # Use a separate variable for uvicorn config to avoid shadowing the config module + uvicorn_config = uvicorn.Config( app=app, host=config.HOST, port=config.PORT, @@ -124,7 +129,7 @@ if __name__ == "__main__": timeout_graceful_shutdown=2 # Reduced from 5 to 2 seconds ) - server = uvicorn.Server(config) + server = uvicorn.Server(uvicorn_config) def signal_handler(sig, frame): logger.info("Shutdown signal received") diff --git a/tools/generate_image.py b/tools/generate_image.py index 6508370..9898dc1 100644 --- a/tools/generate_image.py +++ b/tools/generate_image.py @@ -106,7 +106,7 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]: res_cfg = load_toml(config.RES_PRESETS_PATH) if model_name not in models_cfg: - raise ToolError(f"Model '{model_name}' not found in presets. Available: {', '.join(models_cfg.keys())}") + raise ToolError(f"Model '{model_name}' not found in presets. Please call get_model_info to see available models and their correct names.") model_preset = models_cfg[model_name] @@ -187,9 +187,9 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]: raise ToolError(f"Could not extract file path from SD URL: {sd_img_url}") # Construct proxy URL through our MCP server - # We strip leading slash from raw_path to avoid double slashes in the proxy URL if we want, - # but the current endpoint handles it. - proxy_url = f"http://localhost:{config.PORT}/file{raw_path}" + # Use the captured request host to ensure links work across the network + host = config.request_host.get() + proxy_url = f"http://{host}/file{raw_path}" # Download the image and convert to base64 for the AI's vision img_resp = await asyncio.to_thread(requests.get, sd_img_url, timeout=30)