Added a shutdown custom command that cannot be called by the agent

README.md

@@ -136,6 +136,10 @@ configuration file.
 
 All tools return typed Python objects backed by `audithub-sdk` models.
 
+The optional authenticated `POST /shutdown` operational route is deliberately not an MCP tool
+and never appears in tool discovery. It is enabled only for local HTTP mode when
+`--shutdown-secret-file PATH` is supplied; see [the configuration guide](./docs/configuration.md#local-http-shutdown).
+
 The objects returned by `get_organizations`, `get_projects`, and `get_versions`
 only exposes items belonging to organizations/projects on the configured
 allowlists.

docs/architecture.md

@@ -15,10 +15,13 @@ Every server must satisfy four requirements:
 1. **Explicitly scoped tool surface**: every tool must be narrowly scoped and
    documented. Mutation tools must have clear authorization and bounded effects.
 2. **Least-privilege HTTP surface**: expose only the HTTP methods and paths a
-   server needs. Admin endpoints must not be exposed.
-3. **Credential isolation**: secrets come from configuration files or
-   environment variables, read once at startup. Never include credentials in
-   tool output, error messages, or logs.
+   server needs. AuditHub admin endpoints must not be exposed. Local operational
+   endpoints must be opt-in, independently authenticated, and excluded from the
+   MCP tool surface.
+3. **Credential isolation**: credentials come from configuration files or
+   environment variables and are read once at startup. Generated operational
+   secrets are written only to explicitly configured protected files. Never
+   include credentials or secrets in tool output, error messages, or logs.
 4. **ID allowlisting**: access is restricted to explicitly configured IDs, such
    as organization and project IDs.
    Reject disallowed IDs on the client side before making any network request.

docs/configuration.md

@@ -25,6 +25,36 @@ CLI flags can also override the allowlist and feature toggles after the base set
 Use CLI overrides when you want a temporary change without editing the config file or exporting
 new environment variables.
 
+### Local HTTP Shutdown
+
+Local Streamable HTTP mode can optionally expose an authenticated operational endpoint that
+gracefully exits the server process:
+
+```bash
+audithub-mcp \
+  --local-http-server \
+  --local-http-port 8000 \
+  --shutdown-secret-file /shared/audithub-mcp-shutdown-token
+```
+
+When `--shutdown-secret-file PATH` is supplied, the server generates a fresh random secret at
+startup and atomically writes it to `PATH` with owner-only (`0600`) permissions. The parent
+directory must already exist. Send `POST /shutdown` with the file contents as a bearer token:
+
+```http
+POST /shutdown
+Authorization: Bearer <generated-secret>
+```
+
+An authenticated request returns HTTP 202 and then sends `SIGTERM` to the server process so the
+HTTP runtime can perform graceful shutdown. The route is not an MCP tool and is not included in
+MCP discovery. Without `--shutdown-secret-file`, it returns HTTP 404. The option is rejected
+outside local HTTP mode.
+
+In multi-container deployments, place the output file on a volume mounted only into the MCP
+server and its authorized controller. Do not mount that volume into the agent container. The
+orchestrator remains responsible for its container restart policy after the server exits.
+
 ## Configuration Reference
 
 The JSON and YAML config file uses the same keys listed below. Environment variables use the same

src/audithub_mcp/server.py

@@ -14,7 +14,11 @@
 import base64
 import inspect
 import json
+import os
+import secrets
+import signal
 import sys
+import tempfile
 import time
 import urllib.request
 from collections.abc import Awaitable, Callable, Mapping, Sequence
@@ -56,6 +60,9 @@
 from mcp.server.transport_security import TransportSecuritySettings
 from mcp.shared.exceptions import McpError
 from pydantic import Field, TypeAdapter, ValidationError
+from starlette.background import BackgroundTask
+from starlette.requests import Request
+from starlette.responses import JSONResponse, Response
 
 from audithub_mcp import config as server_config
 from audithub_mcp import paql, parse_fio_logs, vanguard
@@ -128,6 +135,8 @@
 mcp = FastMCP("audithub-mcp")
 mcp._mcp_server.experimental.enable_tasks()
 
+_shutdown_secret: str | None = None
+
 _VANGUARD_CUSTOM_DETECTOR_DOCS_RESOURCE_URI = "docs://vanguard/custom-detectors"
 _VANGUARD_CUSTOM_DETECTOR_DOCS = {
     "Custom Detector Definition": "https://docs.audithub.dev/vanguard/custom-detectors/custom-detector-definition.md",
@@ -368,16 +377,112 @@ def _build_arg_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--local-http-server",
         action="store_true",
-        help="List all registered MCP tools and their schemas, then exit.",
+        help="Run a local Streamable HTTP server instead of using stdio.",
     )
     parser.add_argument(
         "--local-http-port",
         type=int,
-        help="List all registered MCP tools and their schemas, then exit.",
+        help="Port for the local Streamable HTTP server.",
+    )
+    parser.add_argument(
+        "--shutdown-secret-file",
+        metavar="PATH",
+        help=(
+            "Generate a shutdown secret at PATH and enable POST /shutdown; "
+            "valid only in local HTTP mode."
+        ),
     )
     return parser
 
 
+def _write_shutdown_secret(path: Path, secret: str) -> None:
+    """Atomically write a process shutdown secret with owner-only permissions."""
+    if not path.parent.is_dir():
+        raise RuntimeError(f"Shutdown secret directory does not exist: {path.parent}")
+
+    file_descriptor, temporary_path_text = tempfile.mkstemp(
+        dir=path.parent,
+        prefix=f".{path.name}.",
+    )
+    temporary_path = Path(temporary_path_text)
+    try:
+        os.fchmod(file_descriptor, 0o600)
+        output_file = os.fdopen(file_descriptor, "w", encoding="utf-8")
+        file_descriptor = -1
+        with output_file:
+            output_file.write(f"{secret}\n")
+            output_file.flush()
+            os.fsync(output_file.fileno())
+        temporary_path.replace(path)
+    except Exception:
+        if file_descriptor >= 0:
+            os.close(file_descriptor)
+        temporary_path.unlink(missing_ok=True)
+        raise
+
+
+def _terminate_server() -> None:
+    """Ask the current process to follow its normal SIGTERM shutdown path."""
+    os.kill(os.getpid(), signal.SIGTERM)
+
+
+async def _shutdown_server(request: Request) -> Response:
+    """Authenticate and schedule graceful shutdown of the HTTP server process."""
+    if _shutdown_secret is None:
+        return JSONResponse({"detail": "Not found"}, status_code=404)
+
+    authorization = request.headers.get("authorization", "")
+    scheme, separator, supplied_secret = authorization.partition(" ")
+    if (
+        separator == ""
+        or scheme.casefold() != "bearer"
+        or not secrets.compare_digest(supplied_secret, _shutdown_secret)
+    ):
+        return JSONResponse({"detail": "Unauthorized"}, status_code=401)
+
+    return JSONResponse(
+        {"status": "shutting_down"},
+        status_code=202,
+        background=BackgroundTask(_terminate_server),
+    )
+
+
+mcp.custom_route("/shutdown", methods=["POST"], include_in_schema=False)(_shutdown_server)
+
+
+def _run_local_http_server(
+    *,
+    port: int,
+    shutdown_secret_file: str | None,
+    parser: argparse.ArgumentParser,
+) -> None:
+    """Configure and run the local HTTP transport and optional shutdown endpoint."""
+    global _shutdown_secret
+
+    mcp.settings.host = "0.0.0.0"
+    mcp.settings.port = port
+    transport_security = mcp.settings.transport_security or TransportSecuritySettings()
+    transport_security.allowed_hosts = [
+        "127.0.0.1:*",
+        "localhost:*",
+        "host.docker.internal:*",
+    ]
+    mcp.settings.transport_security = transport_security
+
+    if shutdown_secret_file is not None:
+        shutdown_secret = secrets.token_urlsafe(32)
+        try:
+            _write_shutdown_secret(Path(shutdown_secret_file), shutdown_secret)
+        except (OSError, RuntimeError) as exc:
+            parser.error(str(exc))
+        _shutdown_secret = shutdown_secret
+
+    try:
+        mcp.run(transport="streamable-http")
+    finally:
+        _shutdown_secret = None
+
+
 def _apply_cli_args_to_config(
     config: AuditHubServerConfig, args: argparse.Namespace
 ) -> AuditHubServerConfig:
@@ -1931,8 +2036,12 @@ async def _handler(request: mcp_types.ListToolsRequest) -> mcp_types.ServerResul
 
 def main() -> None:
     """Entry point for the ``audithub-mcp`` console script."""
+    global _allowed_org_ids, _allowed_project_ids, _context
+
     parser = _build_arg_parser()
     args, _ = parser.parse_known_args()
+    if args.shutdown_secret_file is not None and not args.local_http_server:
+        parser.error("--shutdown-secret-file requires --local-http-server")
     try:
         if args.list_tools:
             _set_task_runs_enabled(True)
@@ -1945,7 +2054,6 @@ def main() -> None:
     except RuntimeError as exc:
         sys.exit(str(exc))
 
-    global _allowed_org_ids, _allowed_project_ids, _context
     _allowed_org_ids = settings.allowed_org_ids
     _allowed_project_ids = settings.allowed_project_ids
     _set_task_runs_enabled(settings.task_runs_enabled)
@@ -1955,16 +2063,11 @@ def main() -> None:
     if args.local_http_server:
         if not args.local_http_port:
             parser.error("Port is required when server is run in http server mode")
-        mcp.settings.host = "0.0.0.0"
-        mcp.settings.port = args.local_http_port
-        transport_security = mcp.settings.transport_security or TransportSecuritySettings()
-        transport_security.allowed_hosts = [
-            "127.0.0.1:*",
-            "localhost:*",
-            "host.docker.internal:*",
-        ]
-        mcp.settings.transport_security = transport_security
-        mcp.run(transport="streamable-http")
+        _run_local_http_server(
+            port=args.local_http_port,
+            shutdown_secret_file=args.shutdown_secret_file,
+            parser=parser,
+        )
     else:
         mcp.run()
 

tests/test_server.py

@@ -22,6 +22,7 @@
 from mcp import types as mcp_types
 from mcp.server.fastmcp.exceptions import ResourceError  # noqa: E402
 from pydantic import ValidationError
+from starlette.requests import Request
 
 import audithub_mcp.server as server  # noqa: E402
 import audithub_mcp.vanguard as vanguard  # noqa: E402
@@ -579,6 +580,7 @@ def test_main_lists_tools_and_exits(self) -> None:
             self.assertIn('"name": "create_version_from_url"', rendered)
             self.assertIn('"name": "upload_custom_detector"', rendered)
             self.assertIn('"name": "parse_findings_from_task_log"', rendered)
+            self.assertNotIn('"name": "shutdown"', rendered)
             self.assertIsNone(server._context)
         finally:
             server._set_task_runs_enabled(False)
@@ -620,6 +622,101 @@ def test_main_can_disable_env_overrides_for_config_file(self) -> None:
             server._set_edit_custom_detectors_enabled(False)
 
 
+class TestShutdownEndpoint(unittest.IsolatedAsyncioTestCase):
+    def tearDown(self) -> None:
+        server._shutdown_secret = None
+
+    @staticmethod
+    def _request(authorization: str | None = None) -> Request:
+        headers = []
+        if authorization is not None:
+            headers.append((b"authorization", authorization.encode("ascii")))
+        return Request(
+            {
+                "type": "http",
+                "method": "POST",
+                "path": "/shutdown",
+                "headers": headers,
+            }
+        )
+
+    async def test_disabled_endpoint_returns_not_found(self) -> None:
+        server._shutdown_secret = None
+
+        response = await server._shutdown_server(self._request("Bearer unused"))
+
+        self.assertEqual(response.status_code, 404)
+        self.assertIsNone(response.background)
+
+    async def test_missing_or_invalid_secret_is_rejected(self) -> None:
+        server._shutdown_secret = "correct-secret"
+        for authorization in (None, "Basic correct-secret", "Bearer wrong-secret"):
+            with self.subTest(authorization=authorization):
+                response = await server._shutdown_server(self._request(authorization))
+                self.assertEqual(response.status_code, 401)
+                self.assertIsNone(response.background)
+
+    async def test_valid_secret_schedules_process_termination(self) -> None:
+        server._shutdown_secret = "correct-secret"
+        with patch.object(server, "_terminate_server") as terminate_server:
+            response = await server._shutdown_server(self._request("Bearer correct-secret"))
+            self.assertEqual(response.status_code, 202)
+            self.assertEqual(json.loads(response.body), {"status": "shutting_down"})
+            self.assertIsNotNone(response.background)
+            assert response.background is not None
+            await response.background()
+        terminate_server.assert_called_once_with()
+
+    def test_shutdown_secret_is_written_atomically_with_owner_only_permissions(self) -> None:
+        with TemporaryDirectory() as tmpdir:
+            secret_path = Path(tmpdir) / "shutdown-token"
+            secret_path.write_text("stale-secret\n", encoding="utf-8")
+            secret_path.chmod(0o644)
+
+            server._write_shutdown_secret(secret_path, "fresh-secret")
+
+            self.assertEqual(secret_path.read_text(encoding="utf-8"), "fresh-secret\n")
+            self.assertEqual(secret_path.stat().st_mode & 0o777, 0o600)
+
+    def test_main_rejects_shutdown_secret_file_outside_http_mode(self) -> None:
+        argv = ["audithub-mcp", "--shutdown-secret-file", "/tmp/shutdown-token"]
+        with (
+            patch.dict(os.environ, _FULL_ENV, clear=True),
+            patch.object(sys, "argv", argv),
+            self.assertRaises(SystemExit) as cm,
+        ):
+            server.main()
+        self.assertEqual(cm.exception.code, 2)
+
+    def test_main_generates_secret_and_enables_http_shutdown(self) -> None:
+        allow_env = {
+            "AH_ALLOWED_ORG_IDS": "1",
+            "AH_ALLOWED_PROJECT_IDS": "10",
+            **_FULL_ENV,
+        }
+        with TemporaryDirectory() as tmpdir:
+            secret_path = Path(tmpdir) / "shutdown-token"
+            argv = [
+                "audithub-mcp",
+                "--local-http-server",
+                "--local-http-port",
+                "8080",
+                "--shutdown-secret-file",
+                str(secret_path),
+            ]
+            with (
+                patch.dict(os.environ, allow_env, clear=True),
+                patch.object(sys, "argv", argv),
+                patch.object(server.secrets, "token_urlsafe", return_value="generated-secret"),
+                patch.object(server.mcp, "run") as run_server,
+            ):
+                server.main()
+
+            self.assertEqual(secret_path.read_text(encoding="utf-8"), "generated-secret\n")
+            run_server.assert_called_once_with(transport="streamable-http")
+            self.assertIsNone(server._shutdown_secret)
+
+
 class TestRunTool(unittest.IsolatedAsyncioTestCase):
     async def test_returns_value(self) -> None:
         self.assertEqual(await server._run_tool(lambda: {"ok": True}), {"ok": True})