Log elapsed time and token rate when the response arrives.

Add a hard timeout (e.g., 300s) and return a friendly 504.
Add a rule to drop system prompts above a cap unless explicitly allowed.

app.py

@@ -116,6 +116,9 @@ CONTEXT_SIZE = int(os.getenv("CONTEXT_SIZE", "2048"))
 
 PROMPT_MARGIN_TOKENS = int(os.getenv("PROMPT_MARGIN_TOKENS", "256"))
 CHARS_PER_TOKEN_EST = float(os.getenv("CHARS_PER_TOKEN_EST", "4.0"))
+SYSTEM_PROMPT_MAX_TOKENS = int(os.getenv("SYSTEM_PROMPT_MAX_TOKENS", "512"))
+ALLOW_LONG_SYSTEM_PROMPT = os.getenv("ALLOW_LONG_SYSTEM_PROMPT", "0") == "1"
+HARD_REQUEST_TIMEOUT = int(os.getenv("HARD_REQUEST_TIMEOUT", "300"))
 
 
 def _estimate_tokens(text: str) -> int:
@@ -157,7 +160,10 @@ def _compact_messages(messages: list[dict], max_tokens: int) -> list[dict]:
     system_cap = min(1024, max(256, prompt_budget // 3))
     for msg in compacted:
         if msg.get("role") == "system" and "content" in msg:
-            msg["content"] = _truncate_text_to_tokens(str(msg["content"]), system_cap)
+            if not ALLOW_LONG_SYSTEM_PROMPT and _estimate_tokens(str(msg["content"])) > SYSTEM_PROMPT_MAX_TOKENS:
+                msg["content"] = ""
+            else:
+                msg["content"] = _truncate_text_to_tokens(str(msg["content"]), system_cap)
 
     def total_tokens(msgs: list[dict]) -> int:
         return sum(_estimate_tokens(str(m.get("content", ""))) for m in msgs)
@@ -979,6 +985,16 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request)
                 f"original_tokens≈{original_tokens} budget≈{prompt_budget}"
             )
 
+        # Drop system prompts above cap unless allowed
+        if not ALLOW_LONG_SYSTEM_PROMPT:
+            for msg in request.messages:
+                if msg.get("role") == "system" and _estimate_tokens(str(msg.get("content", ""))) > SYSTEM_PROMPT_MAX_TOKENS:
+                    logger.warning(
+                        f"request_id={request_id} system_prompt_dropped "
+                        f"tokens≈{_estimate_tokens(str(msg.get('content', '')))} cap≈{SYSTEM_PROMPT_MAX_TOKENS}"
+                    )
+                    break
+
         compacted_messages = _compact_messages(request.messages, request.max_tokens)
 
         compacted_tokens = _estimate_messages_tokens(compacted_messages)
@@ -1000,18 +1016,21 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request)
             "max_tokens": request.max_tokens,
             "temperature": request.temperature,
         }
-        async with http_session.post(
-            f"{cached_model.url}/v1/chat/completions",
-            json=payload
-        ) as response:
-            if response.status >= 400:
-                error_text = await response.text()
-                logger.error(
-                    f"request_id={request_id} llama-server {response.status} "
-                    f"error_body={error_text[:1000]}"
-                )
-                response.raise_for_status()
-            result = await response.json()
+        async def _do_request():
+            async with http_session.post(
+                f"{cached_model.url}/v1/chat/completions",
+                json=payload
+            ) as response:
+                if response.status >= 400:
+                    error_text = await response.text()
+                    logger.error(
+                        f"request_id={request_id} llama-server {response.status} "
+                        f"error_body={error_text[:1000]}"
+                    )
+                    response.raise_for_status()
+                return await response.json()
+
+        result = await asyncio.wait_for(_do_request(), timeout=HARD_REQUEST_TIMEOUT)
 
         # Update metrics
         request_latency = time.time() - request_start
@@ -1019,6 +1038,18 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request)
         cached_model.total_latency += request_latency
         metrics.record_request(current_model, request_latency)
 
+        # Log elapsed time and token rate (if usage available)
+        usage = result.get("usage") if isinstance(result, dict) else None
+        if usage and usage.get("completion_tokens"):
+            completion_tokens = usage.get("completion_tokens", 0)
+            tok_per_sec = completion_tokens / max(request_latency, 1e-6)
+            logger.info(
+                f"request_id={request_id} done "
+                f"time={request_latency:.2f}s tokens={completion_tokens} tok/s={tok_per_sec:.1f}"
+            )
+        else:
+            logger.info(f"request_id={request_id} done time={request_latency:.2f}s")
+
         return result
     except aiohttp.ClientResponseError as e:
         logger.exception(f"request_id={request_id} llama-server error")
@@ -1026,6 +1057,9 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request)
     except aiohttp.ClientError as e:
         logger.exception(f"request_id={request_id} llama-server error")
         raise HTTPException(status_code=500, detail=f"llama-server error: {str(e)}")
+    except asyncio.TimeoutError:
+        logger.error(f"request_id={request_id} timeout after {HARD_REQUEST_TIMEOUT}s")
+        raise HTTPException(status_code=504, detail="Upstream model timed out. Please retry.")
     except Exception:
         logger.exception(f"request_id={request_id} chat_completions error")
         raise
@@ -1150,6 +1184,16 @@ Always cite sources when using information from the search results."""
                 f"original_tokens≈{original_tokens} budget≈{prompt_budget}"
             )
 
+        # Drop system prompts above cap unless allowed
+        if not ALLOW_LONG_SYSTEM_PROMPT:
+            for msg in augmented_messages:
+                if msg.get("role") == "system" and _estimate_tokens(str(msg.get("content", ""))) > SYSTEM_PROMPT_MAX_TOKENS:
+                    logger.warning(
+                        f"request_id={request_id} system_prompt_dropped "
+                        f"tokens≈{_estimate_tokens(str(msg.get('content', '')))} cap≈{SYSTEM_PROMPT_MAX_TOKENS}"
+                    )
+                    break
+
         augmented_messages = _compact_messages(augmented_messages, request.max_tokens)
 
         compacted_tokens = _estimate_messages_tokens(augmented_messages)
@@ -1170,17 +1214,20 @@ Always cite sources when using information from the search results."""
         if not cached_model:
             raise HTTPException(status_code=500, detail="Current model not loaded")
 
-        # Forward to llama-server with augmented context
-        async with http_session.post(
-            f"{cached_model.url}/v1/chat/completions",
-            json={
-                "messages": augmented_messages,
-                "max_tokens": request.max_tokens,
-                "temperature": request.temperature,
-            }
-        ) as response:
-            response.raise_for_status()
-            result = await response.json()
+        async def _do_request():
+            # Forward to llama-server with augmented context
+            async with http_session.post(
+                f"{cached_model.url}/v1/chat/completions",
+                json={
+                    "messages": augmented_messages,
+                    "max_tokens": request.max_tokens,
+                    "temperature": request.temperature,
+                }
+            ) as response:
+                response.raise_for_status()
+                return await response.json()
+
+        result = await asyncio.wait_for(_do_request(), timeout=HARD_REQUEST_TIMEOUT)
 
         # Add metadata about search results
         result["web_search"] = {
@@ -1195,6 +1242,9 @@ Always cite sources when using information from the search results."""
     except aiohttp.ClientError as e:
         logger.exception(f"request_id={request_id} llama-server error")
         raise HTTPException(status_code=500, detail=f"llama-server error: {str(e)}")
+    except asyncio.TimeoutError:
+        logger.error(f"request_id={request_id} timeout after {HARD_REQUEST_TIMEOUT}s")
+        raise HTTPException(status_code=504, detail="Upstream model timed out. Please retry.")
     except HTTPException:
         raise
     except Exception as e: