Harden server validation and trace handling

Author: haasonsaas

micro_agent/runtime.py

@@ -10,6 +10,9 @@
 TRACES_DIR = os.getenv("TRACES_DIR", "traces")
 os.makedirs(TRACES_DIR, exist_ok=True)
 
+def _get_traces_dir() -> str:
+    return os.getenv("TRACES_DIR", TRACES_DIR)
+
 class Step(TypedDict):
     tool: str
     args: Dict[str, Any]
@@ -39,8 +42,9 @@ def dump_trace(trace_id: str, question: str, steps: List[Step], answer: str, *,
         rec["usage"] = usage
     if cost_usd is not None:
         rec["cost_usd"] = float(cost_usd)
-    os.makedirs(TRACES_DIR, exist_ok=True)
-    path = os.path.join(TRACES_DIR, f"{trace_id}.jsonl")
+    traces_dir = _get_traces_dir()
+    os.makedirs(traces_dir, exist_ok=True)
+    path = os.path.join(traces_dir, f"{trace_id}.jsonl")
     with open(path, "a", encoding="utf-8") as f:
         f.write(json.dumps(rec, ensure_ascii=False, default=str) + "\n")
     return path

micro_agent/server.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-import os, json
+from fastapi.encoders import jsonable_encoder
+import os, json, re
+from threading import Lock
 from pydantic import BaseModel
 from .costs import estimate_prediction_cost
 from importlib.metadata import version as _pkg_version, PackageNotFoundError
@@ -35,6 +37,8 @@ class AskResponse(BaseModel):
 setup_logging()
 configure_lm()
 _agent = MicroAgent()
+_agent_lock = Lock()
+_serialize = os.getenv("MICRO_AGENT_SERIALIZE", "1").strip().lower() not in {"0", "false", "no", "off"}
 
 @app.get("/healthz")
 def healthz():
@@ -52,16 +56,33 @@ def health():
 
 @app.post("/ask", response_model=AskResponse)
 def ask(req: AskRequest):
-    agent = _agent if req.use_tool_calls is None and req.max_steps == _agent.max_steps else MicroAgent(max_steps=req.max_steps, use_tool_calls=req.use_tool_calls)
-    pred = agent(req.question)
+    question = (req.question or "").strip()
+    if not question:
+        raise HTTPException(status_code=400, detail="question must be a non-empty string")
+    if req.max_steps < 1 or req.max_steps > 20:
+        raise HTTPException(status_code=400, detail="max_steps must be between 1 and 20")
+
+    def _call_agent():
+        agent = _agent if req.use_tool_calls is None and req.max_steps == _agent.max_steps else MicroAgent(max_steps=req.max_steps, use_tool_calls=req.use_tool_calls)
+        return agent(question)
+
+    if _serialize:
+        with _agent_lock:
+            pred = _call_agent()
+    else:
+        pred = _call_agent()
+
     trace_id = new_trace_id()
     usage = getattr(pred, "usage", {}) or {}
-    est = estimate_prediction_cost(req.question, pred.trace, pred.answer, usage)
-    path = dump_trace(trace_id, req.question, pred.trace, pred.answer, usage=usage, cost_usd=est.get("cost_usd"))
-    return AskResponse(answer=pred.answer, trace_id=trace_id, trace_path=path, steps=pred.trace, usage=usage, cost_usd=est.get("cost_usd"))
+    est = estimate_prediction_cost(question, pred.trace, pred.answer, usage)
+    path = dump_trace(trace_id, question, pred.trace, pred.answer, usage=usage, cost_usd=est.get("cost_usd"))
+    steps = jsonable_encoder(pred.trace)
+    return AskResponse(answer=pred.answer, trace_id=trace_id, trace_path=path, steps=steps, usage=usage, cost_usd=est.get("cost_usd"))
 
 @app.get("/trace/{trace_id}")
 def get_trace(trace_id: str):
+    if not re.fullmatch(r"[0-9a-f]{32}", trace_id, flags=re.IGNORECASE):
+        raise HTTPException(status_code=400, detail="Invalid trace id format")
     traces_dir = os.getenv("TRACES_DIR", "traces")
     path = os.path.join(traces_dir, f"{trace_id}.jsonl")
     if not os.path.exists(path):

tests/test_regressions.py

@@ -39,7 +39,7 @@ def test_times_is_math_not_time(monkeypatch):
 
 
 def test_dump_trace_serializes_non_json(tmp_path, monkeypatch):
-    monkeypatch.setattr(runtime, "TRACES_DIR", str(tmp_path))
+    monkeypatch.setenv("TRACES_DIR", str(tmp_path))
     trace_id = runtime.new_trace_id()
     steps = [{"tool": "now", "args": {}, "observation": {"when": datetime.datetime(2020, 1, 1)}}]
     path = runtime.dump_trace(trace_id, "q", steps, "a")

tests/test_server.py

@@ -0,0 +1,30 @@
+import importlib
+
+from fastapi.testclient import TestClient
+
+
+def _client(monkeypatch):
+    monkeypatch.setenv("LLM_PROVIDER", "mock")
+    import micro_agent.server as server
+    importlib.reload(server)
+    return TestClient(server.app)
+
+
+def test_ask_rejects_empty_question(monkeypatch):
+    client = _client(monkeypatch)
+    resp = client.post("/ask", json={"question": "   ", "max_steps": 2})
+    assert resp.status_code == 400
+
+
+def test_ask_max_steps_bounds(monkeypatch):
+    client = _client(monkeypatch)
+    resp = client.post("/ask", json={"question": "hi", "max_steps": 0})
+    assert resp.status_code == 400
+    resp = client.post("/ask", json={"question": "hi", "max_steps": 100})
+    assert resp.status_code == 400
+
+
+def test_trace_id_validation(monkeypatch):
+    client = _client(monkeypatch)
+    resp = client.get("/trace/not-a-uuid")
+    assert resp.status_code == 400