Security hardening: server-side sessions, timing-safe tokens

- Server-side session store for workspace/scratch paths — clients
get an opaque session_id, never see filesystem paths
- Server-side cost tracking — client can't bypass limits
- Timing-safe token comparison via hmac.compare_digest
- Path validation ensures workspace is in expected temp directory
- Sanitize error messages (no raw exceptions to client)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

app.py

@@ -6,8 +6,10 @@ streaming question/answer, file upload, and trace endpoints.
 
 from __future__ import annotations
 
+import hmac
 import json
 import os
+import secrets
 import tempfile
 import time
 from collections.abc import Generator
@@ -50,6 +52,41 @@ BASE_PROMPT = (
 hf_api = HfApi(token=HF_TOKEN) if HF_TOKEN else None
 
 
+# ---------------------------------------------------------------------------
+# Server-side session store
+# ---------------------------------------------------------------------------
+
+_sessions: dict[str, dict] = {}
+_TEMP_PREFIX = "/tmp/lh-"
+
+
+def _create_session(workspace_path: str) -> str:
+    session_id = secrets.token_urlsafe(16)
+    scratch_path = tempfile.mkdtemp(prefix="lh-scratch-")
+    _sessions[session_id] = {
+        "workspace": workspace_path,
+        "scratch": scratch_path,
+        "cost": 0.0,
+    }
+    return session_id
+
+
+def _get_session(session_id: str) -> dict | None:
+    session = _sessions.get(session_id)
+    if not session:
+        return None
+    # Validate paths are in expected temp directories
+    if not session["workspace"].startswith(_TEMP_PREFIX):
+        return None
+    return session
+
+
+def _validate_token(token: str) -> bool:
+    if not ACCESS_TOKEN:
+        return True
+    return hmac.compare_digest(token, ACCESS_TOKEN)
+
+
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------
@@ -152,13 +189,11 @@ def format_stats_from_trace(trace: dict) -> str:
 
 def stream_question(
     question: str,
-    workspace_path: str,
-    scratch_path: str,
-    session_cost: float,
+    session_id: str,
     token: str,
 ) -> Generator[str, None, None]:
     """Streaming API — yields JSON event strings."""
-    if ACCESS_TOKEN and token != ACCESS_TOKEN:
+    if not _validate_token(token):
         yield json.dumps({"type": "error", "error": "Invalid access token."})
         return
 
@@ -166,21 +201,20 @@ def stream_question(
         yield json.dumps({"type": "error", "error": "LH_MODEL not set."})
         return
 
-    if session_cost >= MAX_SESSION_COST:
+    session = _get_session(session_id)
+    if not session:
+        yield json.dumps({"type": "error", "error": "Invalid session. Please re-upload your documents."})
+        return
+
+    if session["cost"] >= MAX_SESSION_COST:
         yield json.dumps({
             "type": "error",
-            "error": f"Session cost limit reached (${session_cost:.2f} / ${MAX_SESSION_COST:.2f}).",
+            "error": f"Session cost limit reached (${session['cost']:.2f} / ${MAX_SESSION_COST:.2f}).",
         })
         return
 
-    workspace = Path(workspace_path) if workspace_path else None
-    if not workspace:
-        yield json.dumps({"type": "error", "error": "No documents uploaded."})
-        return
-
-    if not scratch_path:
-        scratch_path = tempfile.mkdtemp(prefix="lh-scratch-")
-    scratch_dir = Path(scratch_path)
+    workspace = Path(session["workspace"])
+    scratch_dir = Path(session["scratch"])
 
     system_prompt = build_system_prompt(base_prompt=BASE_PROMPT, workspace=workspace)
     messages: list[Message] = [
@@ -209,12 +243,16 @@ def stream_question(
                 tool_call_count += 1
                 yield json.dumps({"type": "tool_call", "count": tool_call_count, "name": event.name})
     except Exception as exc:
-        yield json.dumps({"type": "error", "error": str(exc)})
+        yield json.dumps({"type": "error", "error": "An error occurred during processing."})
+        print(f"ERROR in stream_question: {exc}")
         return
 
     trace = agent_run.trace
     trace.wall_time_s = round(time.monotonic() - start, 2)
 
+    # Update server-side session cost
+    session["cost"] += trace.cost or 0
+
     clean_answer, sources = process_citations(trace.answer or "", workspace)
 
     result = {
@@ -234,9 +272,7 @@ def stream_question(
         "sources": sources,
         "stats": format_stats(trace),
         "trace_html": trace_html,
-        "workspace_path": str(workspace),
-        "scratch_path": str(scratch_dir),
-        "session_cost": session_cost + (trace.cost or 0),
+        "session_cost": session["cost"],
     })
 
 
@@ -250,68 +286,61 @@ def build_app() -> gr.Blocks:
         gr.Markdown("# Document Explorer API\n\nThis Space provides the API backend. "
                      "Visit [appsimple.io/explore](https://appsimple.io/explore) for the full interface.")
 
-        # Streaming ask endpoint
+        # Streaming ask endpoint — takes question, session_id, token
         ask_inputs = [
             gr.Textbox(visible=False),  # question
-            gr.Textbox(visible=False),  # workspace_path
-            gr.Textbox(visible=False),  # scratch_path
-            gr.Number(visible=False),   # session_cost
+            gr.Textbox(visible=False),  # session_id
             gr.Textbox(visible=False),  # token
         ]
         ask_output = gr.Textbox(visible=False)
 
-        def api_ask_stream(question, workspace_path, scratch_path, session_cost, token):
-            for event_json in stream_question(
-                question, workspace_path, scratch_path, session_cost, token
-            ):
+        def api_ask_stream(question, session_id, token):
+            for event_json in stream_question(question, session_id, token):
                 yield event_json
 
         ask_btn = gr.Button(visible=False)
         ask_btn.click(api_ask_stream, inputs=ask_inputs, outputs=ask_output, api_name="ask")
 
-        # Upload endpoint — accepts token + JSON list of file paths from /gradio_api/upload
+        # Upload endpoint — accepts files, creates workspace and session
         upload_input = gr.Textbox(visible=False)
         upload_output = gr.Textbox(visible=False)
 
         def api_upload(payload):
-            print(f"[upload] raw payload: {payload!r}")
             try:
                 data = json.loads(payload)
                 token = data.get("token", "")
                 file_paths = data.get("paths", [])
-            except (json.JSONDecodeError, AttributeError) as exc:
-                print(f"[upload] parse error: {exc}")
+            except (json.JSONDecodeError, AttributeError):
                 return json.dumps({"error": "Invalid payload."})
-            print(f"[upload] token={'***' if token else 'empty'}, paths={file_paths}")
-            if ACCESS_TOKEN and token != ACCESS_TOKEN:
+            if not _validate_token(token):
                 return json.dumps({"error": "Invalid access token."})
             if not file_paths:
                 return json.dumps({"error": "No files provided."})
             workspace = save_uploaded_files(file_paths)
-            print(f"[upload] workspace created: {workspace}")
-            return json.dumps({"workspace_path": str(workspace), "file_count": len(file_paths)})
+            session_id = _create_session(str(workspace))
+            return json.dumps({"session_id": session_id, "file_count": len(file_paths)})
 
         upload_btn = gr.Button(visible=False)
         upload_btn.click(api_upload, inputs=upload_input, outputs=upload_output, api_name="upload")
 
-        # Document viewer endpoint
+        # Document viewer endpoint — uses session_id for path lookup
         doc_input = gr.Textbox(visible=False)
-        doc_ws_input = gr.Textbox(visible=False)
+        doc_session_input = gr.Textbox(visible=False)
         doc_output = gr.Textbox(visible=False)
 
-        def api_get_doc(filename, workspace_path):
-            if not workspace_path or not filename:
+        def api_get_doc(filename, session_id):
+            session = _get_session(session_id)
+            if not session or not filename:
                 return json.dumps({"error": "not found"})
             safe_name = Path(filename).name
-            if not safe_name.endswith(".md"):
-                safe_name += ".md"
-            filepath = Path(workspace_path) / safe_name
+            workspace = Path(session["workspace"])
+            filepath = workspace / safe_name
             if not filepath.is_file():
                 return json.dumps({"error": "not found"})
             return json.dumps({"filename": safe_name, "content": filepath.read_text()})
 
         doc_btn = gr.Button(visible=False)
-        doc_btn.click(api_get_doc, inputs=[doc_input, doc_ws_input], outputs=doc_output, api_name="doc")
+        doc_btn.click(api_get_doc, inputs=[doc_input, doc_session_input], outputs=doc_output, api_name="doc")
 
         # Trace list endpoint
         traces_input = gr.Textbox(visible=False)