From 6f34f272aa434826459b76f315f9f33c70d969f4 Mon Sep 17 00:00:00 2001
From: Niklas van Schrick <mc.taucher2003@gmail.com>
Date: Sat, 23 May 2026 00:28:05 +0200
Subject: [PATCH] Implement health service and proper shutdown

The health service is accessable without authentication

The server didn't handle SIGINT or SIGTERM correctly because
it relied on rescuing KeyboardInterrupt which isn't reliable
in containers.
---
 main.py                             | 20 +++++++++++++++-----
 pyproject.toml                      |  1 +
 src/interceptor/auth_interceptor.py |  4 ++++
 uv.lock                             | 15 +++++++++++++++
 4 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/main.py b/main.py
index f1eb33a..cc024a4 100644
--- a/main.py
+++ b/main.py
@@ -1,5 +1,6 @@
 import concurrent.futures
 import os
+import signal
 import sys
 import time
 
@@ -21,6 +22,7 @@
 import tucana.generated.velorum.info_pb2_grpc as info_pb2_grpc
 import tucana.generated.velorum.info_pb2 as info_pb2
 from grpc_reflection.v1alpha import reflection
+from grpc_health.v1 import health, health_pb2, health_pb2_grpc
 
 from src.endpoint.info.info_endpoint import InfoService
 from src.endpoint.generation.generate_endpoint import GenerateService
@@ -34,6 +36,11 @@
     generate_pb2_grpc.add_GenerateServiceServicer_to_server(GenerateService(), server)
     info_pb2_grpc.add_InfoServiceServicer_to_server(InfoService(), server)
 
+    health_servicer = health.HealthServicer()
+    health_servicer.set('liveness', health_pb2.HealthCheckResponse.SERVING)
+    health_servicer.set('readiness', health_pb2.HealthCheckResponse.SERVING)
+    health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server)
+
     port = "0.0.0.0:50051"
     server.add_insecure_port(port)
     print(f"Velorum GenerateService läuft auf {port}...")
@@ -47,9 +54,12 @@
 
     server.start()
 
-    try:
-        while True:
-            time.sleep(86400)
-    except KeyboardInterrupt:
+    def shutdown(signum, frame):
         print("\nServer wird sauber heruntergefahren...")
-        server.stop(0)
+        health_servicer.enter_graceful_shutdown()
+        server.stop(5).wait()
+
+    signal.signal(signal.SIGTERM, shutdown)
+    signal.signal(signal.SIGINT, shutdown)
+
+    server.wait_for_termination()
diff --git a/pyproject.toml b/pyproject.toml
index bfdf249..3b025d2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -7,6 +7,7 @@ requires-python = ">=3.10,<3.13"
 dependencies = [
     "apscheduler>=3.11.2",
     "grpcio>=1.80.0",
+    "grpcio-health-checking>=1.80.0",
     "grpcio-reflection>=1.80.0",
     "instructor>=1.15.1",
     "PyJWT>=2.4.0",
diff --git a/src/interceptor/auth_interceptor.py b/src/interceptor/auth_interceptor.py
index 0bba06f..02660c2 100644
--- a/src/interceptor/auth_interceptor.py
+++ b/src/interceptor/auth_interceptor.py
@@ -8,6 +8,10 @@
 class AuthInterceptor(grpc.ServerInterceptor):
 
     def intercept_service(self, continuation, handler_call_details):
+        method = handler_call_details.method or ''
+        if method.startswith('/grpc.health.v1.Health/'):
+            return continuation(handler_call_details)
+
         metadata = dict(handler_call_details.invocation_metadata or [])
         auth_token = metadata.get('authorization', '')
 
diff --git a/uv.lock b/uv.lock
index be11fe7..60be949 100644
--- a/uv.lock
+++ b/uv.lock
@@ -501,6 +501,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/2f/2c/296f6138caca1f4b92a31ace4ae1b87dab692fc16a7a3417af3bb3c805bf/grpcio-1.80.0-cp312-cp312-win_amd64.whl", hash = "sha256:9fe648599c0e37594c4809d81a9e77bd138cc82eb8baa71b6a86af65426723ff", size = 4880944, upload-time = "2026-03-30T08:47:37.831Z" },
 ]
 
+[[package]]
+name = "grpcio-health-checking"
+version = "1.80.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "grpcio" },
+    { name = "protobuf" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/d1/a2/aa3cc47f19c03f8e5287b987317059753141a3af8f66b96d5a64b3be10b8/grpcio_health_checking-1.80.0.tar.gz", hash = "sha256:2cc5f08bc8b816b8655ab6f59c71450063ba20766d31e21a493e912e3560c8b1", size = 17117, upload-time = "2026-03-30T08:54:41.899Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/5e/d1/d97eb30386feff6ac2a662620e2ed68be352e9a182d62e06213db694906a/grpcio_health_checking-1.80.0-py3-none-any.whl", hash = "sha256:d804d4549cbb71e90ca2c7bf0c501060135dfd220aca8e2c54f96d3e79e210e5", size = 19125, upload-time = "2026-03-30T08:53:44.835Z" },
+]
+
 [[package]]
 name = "grpcio-reflection"
 version = "1.80.0"
@@ -2126,6 +2139,7 @@ source = { virtual = "." }
 dependencies = [
     { name = "apscheduler" },
     { name = "grpcio" },
+    { name = "grpcio-health-checking" },
     { name = "grpcio-reflection" },
     { name = "instructor" },
     { name = "litellm" },
@@ -2140,6 +2154,7 @@ dependencies = [
 requires-dist = [
     { name = "apscheduler", specifier = ">=3.11.2" },
     { name = "grpcio", specifier = ">=1.80.0" },
+    { name = "grpcio-health-checking", specifier = ">=1.80.0" },
     { name = "grpcio-reflection", specifier = ">=1.80.0" },
     { name = "instructor", specifier = ">=1.15.1" },
     { name = "litellm", specifier = "==1.83.7" },