Shutdown

README.md

@@ -97,6 +97,7 @@ Tuning the server's behavior is done via `config.py`.
 | `MODEL_PRESETS_PATH` | Path to the `model_presets.toml` file. | `ROOT_DIR / "model_presets.toml"` |
 | `RES_PRESETS_PATH` | Path to the `resolution_presets.toml` file. | `ROOT_DIR / "resolution_presets.toml"` |
 | `TAG_DATABASE_PATH` | Path to the Danbooru `tags.csv` file. | (Path to extension folder) |
+| `TAG_SEARCH_LIMIT` | Number of results returned by `search_tags` (direct or similar). | `20` |
 
 ### 🧠 Model & Token Tuning (Optimized for Gemma 4)
 | Parameter | Description | Default |

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
@@ -17,6 +19,7 @@ SD_URL = "http://127.0.0.1:7860"
 MODEL_PRESETS_PATH = str(ROOT_DIR / "model_presets.toml")
 RES_PRESETS_PATH = str(ROOT_DIR / "resolution_presets.toml")
 TAG_DATABASE_PATH = "/home/matt/stable-diffusion-webui/extensions/a1111-sd-webui-tagcomplete/tags/danbooru.csv"
+TAG_SEARCH_LIMIT = 20
 
 # --- Model Specific Token Tuning (Tuned for Gemma 4) ---
 # Patch size is typically (clip.vision.patch_size * n_merge)

main.py

@@ -3,6 +3,7 @@ import logging
 import uvicorn
 import signal
 import os
+import sys
 from starlette.applications import Starlette
 from starlette.routing import Route
 from starlette.responses import Response, StreamingResponse, FileResponse
@@ -56,6 +57,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 +121,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 +130,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")
@@ -140,6 +146,5 @@ if __name__ == "__main__":
     except KeyboardInterrupt:
         pass
     finally:
-        # Final fallback to ensure the process actually dies
-        # if the event loop is still hanging on a connection
-        os._exit(0)
+        logger.info("Server process exiting.")
+        sys.exit(0)

mcp_logic.py

@@ -42,23 +42,13 @@ class MCPServer:
                 )
             
             try:
-                result_data = await tool.handler(args)
+                content = await tool.handler(args)
             except ToolError as e:
                 # Convert tool errors into a text response so the LLM can understand and potentially fix the input
-                result_data = {"text": str(e)}
+                content = [{"type": "text", "text": str(e)}]
             except Exception as e:
                 # Catch-all for unexpected crashes to prevent server death
-                result_data = {"text": f"Unexpected internal error: {str(e)}"}
-            
-            # Handle results that are already lists of content (for multimodal/multi-part responses)
-            if isinstance(result_data, list):
-                content = result_data
-            elif isinstance(result_data, dict) and "text" in result_data:
-                content = [{"type": "text", "text": result_data["text"]}]
-            elif isinstance(result_data, dict) and result_data.get("type") == "image":
-                content = [result_data]
-            else:
-                content = [result_data]
+                content = [{"type": "text", "text": f"Unexpected internal error: {str(e)}"}]
             
             return {"content": content}
 

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)

tools/get_model_info.py

@@ -3,7 +3,7 @@ from typing import Any, Dict, List
 import config
 from tools.utils import ToolError, load_toml
 
-async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
+async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
     """
     Returns information about available models or detailed guides for a specific model.
     """
@@ -20,10 +20,13 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
                 "name": name,
                 "description": data.get("description", "No description provided.")
             })
-        return {
-            "text": "Available models:\n\n" + "\n".join([f"- {m['name']}: {m['description']}" for m in catalog]) + 
-                    "\n\nTo get a detailed prompting guide and available resolutions for a specific model, call this tool again with the 'model_name' argument."
-        }
+        return [
+            {
+                "type": "text", 
+                "text": "Available models:\n\n" + "\n".join([f"- {m['name']}: {m['description']}" for m in catalog]) + 
+                        "\n\nTo get a detailed prompting guide and available resolutions for a specific model, call this tool again with the 'model_name' argument."
+            }
+        ]
     
     if model_name not in models:
         # Return the full catalog if the specific model isn't found
@@ -33,11 +36,14 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
                 "name": name,
                 "description": data.get("description", "No description provided.")
             })
-        return {
-            "text": f"Model '{model_name}' not found.\n\nAvailable models:\n\n" + 
-                    "\n".join([f"- {m['name']}: {m['description']}" for m in catalog]) + 
-                    "\n\nTo get a detailed prompting guide and available resolutions for a specific model, call this tool again with the 'model_name' argument."
-        }
+        return [
+            {
+                "type": "text", 
+                "text": f"Model '{model_name}' not found.\n\nAvailable models:\n\n" + 
+                        "\n".join([f"- {m['name']}: {m['description']}" for m in catalog]) + 
+                        "\n\nTo get a detailed prompting guide and available resolutions for a specific model, call this tool again with the 'model_name' argument."
+            }
+        ]
     
     model_data = models[model_name]
     res_set_name = model_data.get("Resolution Set")
@@ -58,13 +64,16 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
     preset_text = "\n".join([f"- {k}: {v}" for k, v in presets_to_show.items()])
     res_text = ", ".join(res_options)
     
-    return {
-        "text": (
-            f"Model: {model_name}\n"
-            f"Description: {description}\n\n"
-            f"--- Prompting Guide ---\n{guide}\n\n"
-            f"--- Active Presets ---\n{preset_text}\n\n"
-            f"--- Available Resolution Presets ---\n{res_text}\n\n"
-            f"Use 'resolution_preset' in generate_image to choose one of these."
-        )
-    }
+    return [
+        {
+            "type": "text", 
+            "text": (
+                f"Model: {model_name}\n"
+                f"Description: {description}\n\n"
+                f"--- Prompting Guide ---\n{guide}\n\n"
+                f"--- Active Presets ---\n{preset_text}\n\n"
+                f"--- Available Resolution Presets ---\n{res_text}\n\n"
+                f"Use 'resolution_preset' in generate_image to choose one of these."
+            )
+        }
+    ]

tools/get_text_context.py

@@ -35,7 +35,7 @@ async def handle(args: Dict[str, Any]):
     # Case 1: No pattern and no lines provided -> Dump whole file
     if not pattern and not lines_str:
         output = [f"{i+1}: {line}" for i, line in enumerate(all_lines)]
-        return {"text": "".join(output)}
+        return [{"type": "text", "text": "".join(output)}]
 
     # Determine target line numbers (1-based)
     target_lines: Set[int] = set()
@@ -61,7 +61,7 @@ async def handle(args: Dict[str, Any]):
             raise ToolError(f"Error parsing line indices: {e}")
 
     if not target_lines:
-        return {"text": "No matching lines found."}
+        return [{"type": "text", "text": "No matching lines found."}]
 
     # Expand targets with context
     lines_to_show: Set[int] = set()
@@ -83,4 +83,4 @@ async def handle(args: Dict[str, Any]):
         output.append(f"{ln}: {all_lines[ln-1]}")
         last_line = ln
 
-    return {"text": "".join(output)}
+    return [{"type": "text", "text": "".join(output)}]

tools/list_directory.py

@@ -43,7 +43,7 @@ async def handle(args: Dict[str, Any]):
     paged_items = all_items[start_idx:end_idx]
 
     if not paged_items:
-        return {"text": f"No items found on page {page}."}
+        return [{"type": "text", "text": f"No items found on page {page}."}]
 
     output = []
     for item in paged_items:
@@ -61,4 +61,4 @@ async def handle(args: Dict[str, Any]):
 
     header = f"Listing of {path_str}\nPage {effective_page} of {total_pages} ({total_items} total items). Showing {start_idx+1}-{min(end_idx, total_items)}."
     
-    return {"text": f"{header}\n\n" + "\n".join(output)}
+    return [{"type": "text", "text": f"{header}\n\n" + "\n".join(output)}]

tools/read_image.py

@@ -25,6 +25,6 @@ async def handle(args: Dict[str, Any]):
     
     try:
         data = base64.b64encode(path.read_bytes()).decode("utf-8")
-        return {"type": "image", "data": data, "mimeType": mime_type}
+        return [{"type": "image", "data": data, "mimeType": mime_type}]
     except Exception as e:
         raise ToolError(f"Error reading image: {str(e)}")

tools/read_metadata.py

@@ -26,11 +26,11 @@ async def handle(args: Dict[str, Any]):
             # Pillow parses PNG text chunks into the .info dictionary
             metadata = img.info
             if not metadata:
-                return {"text": "No metadata found in this PNG."}
+                return [{"type": "text", "text": "No metadata found in this PNG."}]
             
             # Format as a simple key: value list
             output = "\n".join([f"{k}: {v}" for k, v in metadata.items()])
-            return {"text": f"PNG Metadata:\n{output}"}
+            return [{"type": "text", "text": f"PNG Metadata:\n{output}"}]
     except ToolError:
         raise
     except Exception as e:

tools/search_tags.py

@@ -41,7 +41,7 @@ def _load_tags() -> List[Dict[str, Any]]:
     
     return tags
 
-async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
+async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
     """
     Searches the Danbooru tag database for tags matching a query. 
     Returns the most popular tags including alias matches.
@@ -74,9 +74,9 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
     # Sort by count descending
     matches.sort(key=lambda x: x["count"], reverse=True)
 
-    # Format top 20 results
+    # Format top results
     results = []
-    for m in matches[:20]:
+    for m in matches[:config.TAG_SEARCH_LIMIT]:
         count_fmt = format_count(str(m["count"]))
         type_sfx = get_type_suffix(m["type"])
         
@@ -90,10 +90,10 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
     if not results:
         # Attempt to find similar tags using difflib
         all_names_lower = [t["name_lower"] for t in _TAG_CACHE]
-        suggestions_lower = difflib.get_close_matches(query, all_names_lower, n=10, cutoff=0.5)
+        suggestions_lower = difflib.get_close_matches(query, all_names_lower, n=config.TAG_SEARCH_LIMIT, cutoff=0.5)
         
         if not suggestions_lower:
-            return {"text": f"No tags found matching '{query}'."}
+            return [{"type": "text", "text": f"No tags found matching '{query}'."}]
         
         # Map lowercased suggestions back to original tag objects
         suggestions = []
@@ -105,6 +105,6 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
                 type_sfx = get_type_suffix(tag["type"])
                 suggestions.append(f"{tag['name']} ({count_fmt}){type_sfx}")
         
-        return {"text": f"No exact matches for '{query}'. Did you mean:\n" + "\n".join(suggestions)}
+        return [{"type": "text", "text": f"No exact matches for '{query}'. Did you mean:\n" + "\n".join(suggestions)}]
 
-    return {"text": "Top matches:\n" + "\n".join(results)}
+    return [{"type": "text", "text": "Top matches:\n" + "\n".join(results)}]

tools/wikipedia.py

@@ -3,7 +3,7 @@ import urllib.parse
 import json
 import re
 import asyncio
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 from .utils import ToolError
 import config
 
@@ -130,7 +130,7 @@ async def _fetch_toc(title: str) -> Optional[str]:
         
     return "\n".join(lines)
 
-async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
+async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
     """
     Main handler for the browse_wikipedia tool.
     Supports modes: summary, toc, section, and search.
@@ -170,4 +170,4 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
     else:
         raise ToolError(f"Invalid mode '{mode}'. Supported modes: summary, toc, section, search")
             
-    return {"text": result}
+    return [{"type": "text", "text": result}]