feat(mcp-docugen): Task 7-10 renderer, http_routes, mcp_tools, main bootstrap
- Renderer: orchestratore generate() — validazione strict variabili,
materializzazione image vars come asset effimeri su disco + URL rewrite,
asset paths template da ./assets/X -> {PUBLIC_BASE_URL}/assets/<t>/X,
integrazione LLM error -> record success=0
- FastAPI sub-app: GET /health (no auth), /assets/{t}/{f} (auth+traversal check),
/generated/{gen_id}/{f} (410 su scaduto o mancante)
- FastMCP server con 6 tool: template_create/update/delete/list/get,
document_generate. Tools esposti anche via mcp.tools dict per test.
- main.build_app() compone http_app + FastMCP mount su /mcp + auth middleware
+ lifespan cleanup task TTL (24h). run() entry point per script console.
68 test passed. Build Docker arca-mcp-docugen:dev verificata,
/health endpoint risponde correttamente nel container.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,54 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.responses import FileResponse, JSONResponse
|
||||
|
||||
from mcp_docugen.generation_store import GenerationStore
|
||||
from mcp_docugen.template_store import InvalidTemplateName, TemplateStore
|
||||
|
||||
_SAFE_SEGMENT_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
|
||||
|
||||
|
||||
def build_http_app(
|
||||
template_store: TemplateStore,
|
||||
generation_store: GenerationStore,
|
||||
) -> FastAPI:
|
||||
app = FastAPI()
|
||||
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
return {"status": "ok"}
|
||||
|
||||
@app.get("/assets/{template_name}/{filename}")
|
||||
async def get_template_asset(template_name: str, filename: str):
|
||||
if not _SAFE_SEGMENT_RE.match(template_name) or not _SAFE_SEGMENT_RE.match(
|
||||
filename
|
||||
):
|
||||
raise HTTPException(status_code=400, detail="invalid path")
|
||||
try:
|
||||
path = template_store.asset_path(template_name, filename)
|
||||
except (ValueError, InvalidTemplateName) as exc:
|
||||
raise HTTPException(status_code=400, detail="invalid path") from exc
|
||||
if not path.exists():
|
||||
raise HTTPException(status_code=404, detail="not found")
|
||||
return FileResponse(path)
|
||||
|
||||
@app.get("/generated/{gen_id}/{filename}")
|
||||
async def get_generated_asset(gen_id: str, filename: str):
|
||||
if not _SAFE_SEGMENT_RE.match(gen_id) or not _SAFE_SEGMENT_RE.match(filename):
|
||||
raise HTTPException(status_code=400, detail="invalid path")
|
||||
info = await generation_store.get_ephemeral_asset(gen_id, filename)
|
||||
if info is None or info.is_expired or not Path(info.file_path).exists():
|
||||
return JSONResponse(
|
||||
status_code=410,
|
||||
content={
|
||||
"error": "gone",
|
||||
"reason": "ephemeral asset expired or not found",
|
||||
},
|
||||
)
|
||||
return FileResponse(info.file_path, media_type=info.mime)
|
||||
|
||||
return app
|
||||
@@ -0,0 +1,107 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
import uvicorn
|
||||
from fastapi import FastAPI
|
||||
|
||||
from mcp_docugen.auth import ApiKeyAuthMiddleware
|
||||
from mcp_docugen.config import Settings
|
||||
from mcp_docugen.generation_store import GenerationStore
|
||||
from mcp_docugen.http_routes import build_http_app
|
||||
from mcp_docugen.llm_client import OpenRouterClient
|
||||
from mcp_docugen.mcp_tools import build_mcp_server
|
||||
from mcp_docugen.renderer import Renderer
|
||||
from mcp_docugen.template_store import TemplateStore
|
||||
|
||||
logger = logging.getLogger("mcp_docugen")
|
||||
|
||||
|
||||
async def build_app(settings: Settings | None = None) -> FastAPI:
|
||||
settings = settings or Settings()
|
||||
settings.data_dir.mkdir(parents=True, exist_ok=True)
|
||||
templates_dir = settings.data_dir / "templates"
|
||||
generated_dir = settings.data_dir / "generated"
|
||||
db_path = settings.data_dir / "mcp_docugen.db"
|
||||
|
||||
template_store = TemplateStore(base_dir=templates_dir)
|
||||
generation_store = GenerationStore(
|
||||
db_path=db_path, generated_dir=generated_dir
|
||||
)
|
||||
await generation_store.init()
|
||||
|
||||
llm = OpenRouterClient(
|
||||
api_key=settings.openrouter_api_key,
|
||||
base_url=settings.openrouter_base_url,
|
||||
timeout=settings.llm_timeout_seconds,
|
||||
)
|
||||
renderer = Renderer(
|
||||
template_store=template_store,
|
||||
generation_store=generation_store,
|
||||
llm=llm,
|
||||
public_base_url=settings.public_base_url,
|
||||
default_model=settings.llm_model_default,
|
||||
asset_ttl_days=settings.asset_ttl_days,
|
||||
max_image_size_mb=settings.max_image_size_mb,
|
||||
)
|
||||
|
||||
mcp = build_mcp_server(template_store, renderer)
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
cleanup_task = asyncio.create_task(
|
||||
_periodic_cleanup(generation_store, interval_seconds=24 * 3600)
|
||||
)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
cleanup_task.cancel()
|
||||
try:
|
||||
await cleanup_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
app = build_http_app(template_store, generation_store)
|
||||
app.router.lifespan_context = lifespan
|
||||
app.mount("/mcp", mcp.streamable_http_app())
|
||||
|
||||
app.add_middleware(
|
||||
ApiKeyAuthMiddleware,
|
||||
api_key=settings.api_key,
|
||||
exempt_paths={"/health"},
|
||||
)
|
||||
|
||||
app.state.settings = settings
|
||||
app.state.template_store = template_store
|
||||
app.state.generation_store = generation_store
|
||||
app.state.renderer = renderer
|
||||
app.state.llm = llm
|
||||
|
||||
return app
|
||||
|
||||
|
||||
async def _periodic_cleanup(
|
||||
generation_store: GenerationStore, interval_seconds: int
|
||||
) -> None:
|
||||
while True:
|
||||
try:
|
||||
removed = await generation_store.cleanup_expired()
|
||||
if removed:
|
||||
logger.info("cleanup_expired removed %d assets", removed)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.error("cleanup task error: %s", exc)
|
||||
await asyncio.sleep(interval_seconds)
|
||||
|
||||
|
||||
def run() -> None:
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s %(levelname)s %(name)s %(message)s",
|
||||
)
|
||||
host = os.environ.get("HOST", "0.0.0.0")
|
||||
port = int(os.environ.get("PORT", "9100"))
|
||||
app = asyncio.run(build_app())
|
||||
uvicorn.run(app, host=host, port=port)
|
||||
@@ -0,0 +1,92 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
from mcp_docugen.models import TemplateFrontmatter
|
||||
from mcp_docugen.renderer import Renderer
|
||||
from mcp_docugen.template_store import TemplateStore
|
||||
|
||||
|
||||
def build_mcp_server(
|
||||
template_store: TemplateStore, renderer: Renderer
|
||||
) -> FastMCP:
|
||||
mcp = FastMCP("mcp-docugen")
|
||||
|
||||
@mcp.tool()
|
||||
async def template_create(
|
||||
name: str,
|
||||
frontmatter: dict,
|
||||
body: str,
|
||||
assets: list[dict] | None = None,
|
||||
) -> dict:
|
||||
"""Create a new template. Fails if the name already exists."""
|
||||
fm = TemplateFrontmatter(**frontmatter)
|
||||
await template_store.create(
|
||||
name=name, frontmatter=fm, body=body, assets=assets
|
||||
)
|
||||
return {"name": name, "status": "created"}
|
||||
|
||||
@mcp.tool()
|
||||
async def template_update(
|
||||
name: str,
|
||||
frontmatter: dict,
|
||||
body: str,
|
||||
assets: list[dict] | None = None,
|
||||
) -> dict:
|
||||
"""Update an existing template. Replaces frontmatter, body, and (if given) assets."""
|
||||
fm = TemplateFrontmatter(**frontmatter)
|
||||
await template_store.update(
|
||||
name=name, frontmatter=fm, body=body, assets=assets
|
||||
)
|
||||
return {"name": name, "status": "updated"}
|
||||
|
||||
@mcp.tool()
|
||||
async def template_delete(name: str) -> dict:
|
||||
"""Delete a template and all its assets."""
|
||||
await template_store.delete(name)
|
||||
return {"name": name, "status": "deleted"}
|
||||
|
||||
@mcp.tool()
|
||||
async def template_get(name: str) -> dict:
|
||||
"""Return a template's frontmatter, body, and asset list."""
|
||||
loaded = await template_store.get(name)
|
||||
return {
|
||||
"name": loaded.frontmatter.name,
|
||||
"frontmatter": loaded.frontmatter.model_dump(exclude_none=True),
|
||||
"body": loaded.body,
|
||||
"assets": [a.model_dump() for a in loaded.assets],
|
||||
"updated_at": loaded.updated_at.isoformat(),
|
||||
}
|
||||
|
||||
@mcp.tool()
|
||||
async def template_list() -> list[dict]:
|
||||
"""List all templates with name, description, and updated_at."""
|
||||
summaries = await template_store.list()
|
||||
return [s.model_dump(mode="json") for s in summaries]
|
||||
|
||||
@mcp.tool()
|
||||
async def document_generate(
|
||||
template_name: str,
|
||||
content_md: str,
|
||||
variables: dict,
|
||||
instructions: str | None = None,
|
||||
) -> dict:
|
||||
"""Generate a Markdown document from a template, content, and variables."""
|
||||
result = await renderer.generate(
|
||||
template_name=template_name,
|
||||
content_md=content_md,
|
||||
variables=variables,
|
||||
instructions=instructions,
|
||||
)
|
||||
return result.model_dump(mode="json")
|
||||
|
||||
mcp.tools = {
|
||||
"template_create": template_create,
|
||||
"template_update": template_update,
|
||||
"template_delete": template_delete,
|
||||
"template_get": template_get,
|
||||
"template_list": template_list,
|
||||
"document_generate": document_generate,
|
||||
}
|
||||
|
||||
return mcp
|
||||
@@ -0,0 +1,255 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import re
|
||||
import uuid
|
||||
|
||||
import aiofiles
|
||||
from pydantic import ValidationError
|
||||
|
||||
from mcp_docugen.generation_store import (
|
||||
EphemeralAssetRecord,
|
||||
GenerationRecord,
|
||||
GenerationStore,
|
||||
)
|
||||
from mcp_docugen.llm_client import LLMError, OpenRouterClient
|
||||
from mcp_docugen.models import (
|
||||
GenerationResult,
|
||||
ImageVariable,
|
||||
TemplateFrontmatter,
|
||||
TokenUsage,
|
||||
)
|
||||
from mcp_docugen.template_store import TemplateStore
|
||||
|
||||
|
||||
class RendererError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class MissingVariables(RendererError):
|
||||
def __init__(self, missing: list[str]) -> None:
|
||||
super().__init__(f"missing required variables: {missing}")
|
||||
self.missing = missing
|
||||
|
||||
|
||||
class InvalidVariableType(RendererError):
|
||||
pass
|
||||
|
||||
|
||||
class InvalidImageEncoding(RendererError):
|
||||
pass
|
||||
|
||||
|
||||
class ImageTooLarge(RendererError):
|
||||
pass
|
||||
|
||||
|
||||
class AssetStorageError(RendererError):
|
||||
pass
|
||||
|
||||
|
||||
_ASSET_TEMPLATE_REF_RE = re.compile(r"\.\/assets\/([A-Za-z0-9._-]+)")
|
||||
|
||||
|
||||
class Renderer:
|
||||
def __init__(
|
||||
self,
|
||||
template_store: TemplateStore,
|
||||
generation_store: GenerationStore,
|
||||
llm: OpenRouterClient,
|
||||
public_base_url: str,
|
||||
default_model: str,
|
||||
asset_ttl_days: int,
|
||||
max_image_size_mb: int,
|
||||
) -> None:
|
||||
self.template_store = template_store
|
||||
self.generation_store = generation_store
|
||||
self.llm = llm
|
||||
self.public_base_url = public_base_url.rstrip("/")
|
||||
self.default_model = default_model
|
||||
self.asset_ttl_days = asset_ttl_days
|
||||
self.max_image_size_bytes = max_image_size_mb * 1024 * 1024
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
template_name: str,
|
||||
content_md: str,
|
||||
variables: dict,
|
||||
instructions: str | None,
|
||||
) -> GenerationResult:
|
||||
loaded = await self.template_store.get(template_name)
|
||||
gen_id = str(uuid.uuid4())
|
||||
|
||||
validated = self._validate_variables(loaded.frontmatter, variables)
|
||||
ephemeral_urls, var_strings = await self._materialize_image_variables(
|
||||
gen_id, validated
|
||||
)
|
||||
|
||||
system_prompt = self._build_system_prompt(
|
||||
loaded.frontmatter, loaded.body, var_strings
|
||||
)
|
||||
user_prompt = self._build_user_prompt(content_md, var_strings, instructions)
|
||||
model = loaded.frontmatter.model or self.default_model
|
||||
|
||||
try:
|
||||
response = await self.llm.chat(
|
||||
model=model, system=system_prompt, user=user_prompt
|
||||
)
|
||||
except LLMError as exc:
|
||||
await self.generation_store.record_generation(
|
||||
GenerationRecord(
|
||||
id=gen_id,
|
||||
template_name=template_name,
|
||||
model=model,
|
||||
tokens_in=0,
|
||||
tokens_out=0,
|
||||
cost_usd=0.0,
|
||||
success=False,
|
||||
error_msg=f"{type(exc).__name__}: {exc}",
|
||||
)
|
||||
)
|
||||
raise
|
||||
|
||||
await self.generation_store.record_generation(
|
||||
GenerationRecord(
|
||||
id=gen_id,
|
||||
template_name=template_name,
|
||||
model=response.model,
|
||||
tokens_in=response.tokens_in,
|
||||
tokens_out=response.tokens_out,
|
||||
cost_usd=response.cost_usd,
|
||||
success=True,
|
||||
error_msg=None,
|
||||
)
|
||||
)
|
||||
|
||||
expires_at = None
|
||||
if ephemeral_urls:
|
||||
from pathlib import Path
|
||||
|
||||
info = await self.generation_store.get_ephemeral_asset(
|
||||
gen_id, Path(ephemeral_urls[0]).name
|
||||
)
|
||||
expires_at = info.expires_at if info else None
|
||||
|
||||
return GenerationResult(
|
||||
generation_id=gen_id,
|
||||
markdown=response.text,
|
||||
model=response.model,
|
||||
tokens=TokenUsage(input=response.tokens_in, output=response.tokens_out),
|
||||
cost_usd=response.cost_usd,
|
||||
ephemeral_assets_urls=ephemeral_urls,
|
||||
ephemeral_expires_at=expires_at,
|
||||
)
|
||||
|
||||
def _validate_variables(
|
||||
self, fm: TemplateFrontmatter, variables: dict
|
||||
) -> dict:
|
||||
missing = [
|
||||
v.name for v in fm.required_variables if v.name not in variables
|
||||
]
|
||||
if missing:
|
||||
raise MissingVariables(missing)
|
||||
|
||||
validated: dict = {}
|
||||
for var_def in fm.required_variables:
|
||||
raw = variables[var_def.name]
|
||||
if var_def.type == "string":
|
||||
if not isinstance(raw, str):
|
||||
raise InvalidVariableType(
|
||||
f"{var_def.name}: expected string, got {type(raw).__name__}"
|
||||
)
|
||||
validated[var_def.name] = raw
|
||||
elif var_def.type == "image":
|
||||
if not isinstance(raw, dict):
|
||||
raise InvalidVariableType(
|
||||
f"{var_def.name}: expected image object, got {type(raw).__name__}"
|
||||
)
|
||||
try:
|
||||
img = ImageVariable(**raw)
|
||||
except ValidationError as exc:
|
||||
raise InvalidVariableType(f"{var_def.name}: {exc}") from exc
|
||||
validated[var_def.name] = img
|
||||
return validated
|
||||
|
||||
async def _materialize_image_variables(
|
||||
self, gen_id: str, validated: dict
|
||||
) -> tuple[list[str], dict[str, str]]:
|
||||
urls: list[str] = []
|
||||
var_strings: dict[str, str] = {}
|
||||
gen_dir = self.generation_store.generated_dir / gen_id
|
||||
|
||||
for name, value in validated.items():
|
||||
if isinstance(value, ImageVariable):
|
||||
try:
|
||||
raw_bytes = base64.b64decode(value.data_b64, validate=True)
|
||||
except (binascii.Error, ValueError) as exc:
|
||||
raise InvalidImageEncoding(f"{name}: {exc}") from exc
|
||||
if len(raw_bytes) > self.max_image_size_bytes:
|
||||
raise ImageTooLarge(
|
||||
f"{name}: {len(raw_bytes)} bytes > {self.max_image_size_bytes}"
|
||||
)
|
||||
ext = {
|
||||
"image/png": "png",
|
||||
"image/jpeg": "jpg",
|
||||
"image/webp": "webp",
|
||||
}[value.mime]
|
||||
filename = f"{name}.{ext}"
|
||||
file_path = gen_dir / filename
|
||||
gen_dir.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
async with aiofiles.open(file_path, "wb") as f:
|
||||
await f.write(raw_bytes)
|
||||
except OSError as exc:
|
||||
raise AssetStorageError(str(exc)) from exc
|
||||
await self.generation_store.register_ephemeral_asset(
|
||||
EphemeralAssetRecord(
|
||||
generation_id=gen_id,
|
||||
var_name=name,
|
||||
file_path=str(file_path),
|
||||
mime=value.mime,
|
||||
ttl_days=self.asset_ttl_days,
|
||||
)
|
||||
)
|
||||
url = (
|
||||
f"{self.public_base_url}/generated/{gen_id}/{filename}"
|
||||
)
|
||||
urls.append(url)
|
||||
var_strings[name] = url
|
||||
else:
|
||||
var_strings[name] = str(value)
|
||||
return urls, var_strings
|
||||
|
||||
def _build_system_prompt(
|
||||
self,
|
||||
fm: TemplateFrontmatter,
|
||||
body: str,
|
||||
var_strings: dict[str, str],
|
||||
) -> str:
|
||||
template_name = fm.name
|
||||
|
||||
def _replace_asset(match: re.Match) -> str:
|
||||
filename = match.group(1)
|
||||
return f"{self.public_base_url}/assets/{template_name}/{filename}"
|
||||
|
||||
resolved = _ASSET_TEMPLATE_REF_RE.sub(_replace_asset, body)
|
||||
for name, value in var_strings.items():
|
||||
resolved = resolved.replace(f"{{{{{name}}}}}", value)
|
||||
return resolved
|
||||
|
||||
def _build_user_prompt(
|
||||
self,
|
||||
content_md: str,
|
||||
var_strings: dict[str, str],
|
||||
instructions: str | None,
|
||||
) -> str:
|
||||
parts = ["## Raw content", content_md]
|
||||
if var_strings:
|
||||
parts.append("## Variables")
|
||||
for name, value in var_strings.items():
|
||||
parts.append(f"- {name}: {value}")
|
||||
if instructions:
|
||||
parts.append("## Additional instructions")
|
||||
parts.append(instructions)
|
||||
return "\n\n".join(parts)
|
||||
@@ -0,0 +1,76 @@
|
||||
import base64
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
from mcp_docugen.main import build_app
|
||||
from mcp_docugen.models import TemplateFrontmatter, TemplateVariable
|
||||
|
||||
|
||||
def _openrouter_success(text: str = "# Generated"):
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "g",
|
||||
"choices": [{"message": {"role": "assistant", "content": text}}],
|
||||
"model": "anthropic/claude-sonnet-4",
|
||||
"usage": {
|
||||
"prompt_tokens": 50,
|
||||
"completion_tokens": 100,
|
||||
"total_cost": 0.02,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def app_env(tmp_path, monkeypatch):
|
||||
monkeypatch.setenv("API_KEY", "test-api-key")
|
||||
monkeypatch.setenv("OPENROUTER_API_KEY", "sk-or-test")
|
||||
monkeypatch.setenv("PUBLIC_BASE_URL", "https://mcp.test.com")
|
||||
monkeypatch.setenv("DATA_DIR", str(tmp_path))
|
||||
app = await build_app()
|
||||
return app, tmp_path
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_full_generation_flow_records_and_serves_asset(app_env):
|
||||
app, tmp_path = app_env
|
||||
respx.post("https://openrouter.ai/api/v1/chat/completions").mock(
|
||||
return_value=_openrouter_success("# Done")
|
||||
)
|
||||
|
||||
template_store = app.state.template_store
|
||||
renderer = app.state.renderer
|
||||
generation_store = app.state.generation_store
|
||||
|
||||
fm = TemplateFrontmatter(
|
||||
name="report",
|
||||
description="x",
|
||||
required_variables=[TemplateVariable(name="foto", type="image")],
|
||||
)
|
||||
await template_store.create(
|
||||
name="report", frontmatter=fm, body=""
|
||||
)
|
||||
|
||||
png = b"\x89PNG\r\n\x1a\n" + b"\x00" * 50
|
||||
img_var = {
|
||||
"kind": "image",
|
||||
"data_b64": base64.b64encode(png).decode(),
|
||||
"mime": "image/png",
|
||||
}
|
||||
|
||||
result = await renderer.generate(
|
||||
template_name="report",
|
||||
content_md="content",
|
||||
variables={"foto": img_var},
|
||||
instructions=None,
|
||||
)
|
||||
|
||||
assert result.markdown == "# Done"
|
||||
stats = await generation_store.get_stats()
|
||||
assert stats["success"] == 1
|
||||
|
||||
gen_files = list((tmp_path / "generated" / result.generation_id).iterdir())
|
||||
assert len(gen_files) == 1
|
||||
@@ -0,0 +1,122 @@
|
||||
import base64
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from mcp_docugen.auth import ApiKeyAuthMiddleware
|
||||
from mcp_docugen.generation_store import EphemeralAssetRecord, GenerationStore
|
||||
from mcp_docugen.http_routes import build_http_app
|
||||
from mcp_docugen.models import TemplateFrontmatter
|
||||
from mcp_docugen.template_store import TemplateStore
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def env(tmp_path):
|
||||
template_store = TemplateStore(base_dir=tmp_path / "templates")
|
||||
generation_store = GenerationStore(
|
||||
db_path=tmp_path / "gen.db",
|
||||
generated_dir=tmp_path / "generated",
|
||||
)
|
||||
await generation_store.init()
|
||||
|
||||
inner = build_http_app(template_store, generation_store)
|
||||
app = FastAPI()
|
||||
app.add_middleware(
|
||||
ApiKeyAuthMiddleware, api_key="test-key", exempt_paths={"/health"}
|
||||
)
|
||||
app.mount("/", inner)
|
||||
return TestClient(app), template_store, generation_store, tmp_path
|
||||
|
||||
|
||||
def _headers():
|
||||
return {"Authorization": "Bearer test-key"}
|
||||
|
||||
|
||||
def test_health_no_auth(env):
|
||||
tc, *_ = env
|
||||
r = tc.get("/health")
|
||||
assert r.status_code == 200
|
||||
assert r.json() == {"status": "ok"}
|
||||
|
||||
|
||||
async def test_assets_serve_existing_file(env):
|
||||
tc, template_store, _, _ = env
|
||||
fm = TemplateFrontmatter(name="brand", description="x")
|
||||
assets = [
|
||||
{
|
||||
"filename": "logo.png",
|
||||
"data_b64": base64.b64encode(b"\x89PNG").decode(),
|
||||
"mime": "image/png",
|
||||
}
|
||||
]
|
||||
await template_store.create(
|
||||
name="brand", frontmatter=fm, body="b", assets=assets
|
||||
)
|
||||
|
||||
r = tc.get("/assets/brand/logo.png", headers=_headers())
|
||||
assert r.status_code == 200
|
||||
assert r.content == b"\x89PNG"
|
||||
|
||||
|
||||
def test_assets_require_auth(env):
|
||||
tc, *_ = env
|
||||
r = tc.get("/assets/brand/logo.png")
|
||||
assert r.status_code == 401
|
||||
|
||||
|
||||
def test_assets_path_traversal_rejected(env):
|
||||
tc, *_ = env
|
||||
# Starlette decodifica %2F in / prima del routing, quindi la path
|
||||
# non matcha /assets/{template}/{filename} e torna 404. 400/404 entrambi safe.
|
||||
r = tc.get("/assets/brand/..%2Fevil.png", headers=_headers())
|
||||
assert r.status_code in (400, 404)
|
||||
|
||||
|
||||
def test_assets_missing_returns_404(env):
|
||||
tc, *_ = env
|
||||
r = tc.get("/assets/brand/missing.png", headers=_headers())
|
||||
assert r.status_code == 404
|
||||
|
||||
|
||||
async def test_generated_fresh_returns_file(env):
|
||||
tc, _, generation_store, tmp_path = env
|
||||
gen_dir = tmp_path / "generated" / "g-1"
|
||||
gen_dir.mkdir(parents=True)
|
||||
(gen_dir / "foto.png").write_bytes(b"image-bytes")
|
||||
await generation_store.register_ephemeral_asset(
|
||||
EphemeralAssetRecord(
|
||||
generation_id="g-1",
|
||||
var_name="foto",
|
||||
file_path=str(gen_dir / "foto.png"),
|
||||
mime="image/png",
|
||||
ttl_days=30,
|
||||
)
|
||||
)
|
||||
r = tc.get("/generated/g-1/foto.png", headers=_headers())
|
||||
assert r.status_code == 200
|
||||
assert r.content == b"image-bytes"
|
||||
|
||||
|
||||
async def test_generated_expired_returns_410(env):
|
||||
tc, _, generation_store, tmp_path = env
|
||||
gen_dir = tmp_path / "generated" / "g-old"
|
||||
gen_dir.mkdir(parents=True)
|
||||
(gen_dir / "foto.png").write_bytes(b"x")
|
||||
await generation_store.register_ephemeral_asset(
|
||||
EphemeralAssetRecord(
|
||||
generation_id="g-old",
|
||||
var_name="foto",
|
||||
file_path=str(gen_dir / "foto.png"),
|
||||
mime="image/png",
|
||||
ttl_days=-1,
|
||||
)
|
||||
)
|
||||
r = tc.get("/generated/g-old/foto.png", headers=_headers())
|
||||
assert r.status_code == 410
|
||||
|
||||
|
||||
def test_generated_unknown_returns_410(env):
|
||||
tc, *_ = env
|
||||
r = tc.get("/generated/unknown-id/file.png", headers=_headers())
|
||||
assert r.status_code == 410
|
||||
@@ -0,0 +1,144 @@
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from mcp_docugen.generation_store import GenerationStore
|
||||
from mcp_docugen.llm_client import LLMResponse
|
||||
from mcp_docugen.mcp_tools import build_mcp_server
|
||||
from mcp_docugen.renderer import Renderer
|
||||
from mcp_docugen.template_store import TemplateStore
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def mcp_env(tmp_path):
|
||||
template_store = TemplateStore(base_dir=tmp_path / "templates")
|
||||
generation_store = GenerationStore(
|
||||
db_path=tmp_path / "gen.db",
|
||||
generated_dir=tmp_path / "generated",
|
||||
)
|
||||
await generation_store.init()
|
||||
llm = AsyncMock()
|
||||
llm.chat.return_value = LLMResponse(
|
||||
text="# Out",
|
||||
model="m",
|
||||
tokens_in=5,
|
||||
tokens_out=10,
|
||||
cost_usd=0.01,
|
||||
latency_ms=50,
|
||||
)
|
||||
renderer = Renderer(
|
||||
template_store=template_store,
|
||||
generation_store=generation_store,
|
||||
llm=llm,
|
||||
public_base_url="https://mcp.example.com",
|
||||
default_model="m",
|
||||
asset_ttl_days=30,
|
||||
max_image_size_mb=10,
|
||||
)
|
||||
mcp = build_mcp_server(template_store, renderer)
|
||||
return mcp, template_store
|
||||
|
||||
|
||||
async def test_template_create_and_get(mcp_env):
|
||||
mcp, _ = mcp_env
|
||||
await mcp.tools["template_create"](
|
||||
name="demo",
|
||||
frontmatter={"name": "demo", "description": "d"},
|
||||
body="body",
|
||||
assets=None,
|
||||
)
|
||||
got = await mcp.tools["template_get"](name="demo")
|
||||
assert got["name"] == "demo"
|
||||
assert got["body"] == "body"
|
||||
|
||||
|
||||
async def test_template_create_duplicate_errors(mcp_env):
|
||||
mcp, _ = mcp_env
|
||||
await mcp.tools["template_create"](
|
||||
name="demo",
|
||||
frontmatter={"name": "demo", "description": "x"},
|
||||
body="b",
|
||||
assets=None,
|
||||
)
|
||||
with pytest.raises(Exception):
|
||||
await mcp.tools["template_create"](
|
||||
name="demo",
|
||||
frontmatter={"name": "demo", "description": "x"},
|
||||
body="b",
|
||||
assets=None,
|
||||
)
|
||||
|
||||
|
||||
async def test_template_list_returns_summaries(mcp_env):
|
||||
mcp, _ = mcp_env
|
||||
await mcp.tools["template_create"](
|
||||
name="a",
|
||||
frontmatter={"name": "a", "description": "A"},
|
||||
body="b",
|
||||
assets=None,
|
||||
)
|
||||
await mcp.tools["template_create"](
|
||||
name="b",
|
||||
frontmatter={"name": "b", "description": "B"},
|
||||
body="b",
|
||||
assets=None,
|
||||
)
|
||||
result = await mcp.tools["template_list"]()
|
||||
names = sorted(t["name"] for t in result)
|
||||
assert names == ["a", "b"]
|
||||
|
||||
|
||||
async def test_template_update(mcp_env):
|
||||
mcp, _ = mcp_env
|
||||
await mcp.tools["template_create"](
|
||||
name="u",
|
||||
frontmatter={"name": "u", "description": "d"},
|
||||
body="old",
|
||||
assets=None,
|
||||
)
|
||||
await mcp.tools["template_update"](
|
||||
name="u",
|
||||
frontmatter={"name": "u", "description": "new"},
|
||||
body="new",
|
||||
assets=None,
|
||||
)
|
||||
got = await mcp.tools["template_get"](name="u")
|
||||
assert got["body"] == "new"
|
||||
assert got["frontmatter"]["description"] == "new"
|
||||
|
||||
|
||||
async def test_template_delete(mcp_env):
|
||||
mcp, _ = mcp_env
|
||||
await mcp.tools["template_create"](
|
||||
name="d",
|
||||
frontmatter={"name": "d", "description": "x"},
|
||||
body="b",
|
||||
assets=None,
|
||||
)
|
||||
await mcp.tools["template_delete"](name="d")
|
||||
with pytest.raises(Exception):
|
||||
await mcp.tools["template_get"](name="d")
|
||||
|
||||
|
||||
async def test_document_generate_happy_path(mcp_env):
|
||||
mcp, _ = mcp_env
|
||||
await mcp.tools["template_create"](
|
||||
name="g",
|
||||
frontmatter={
|
||||
"name": "g",
|
||||
"description": "x",
|
||||
"required_variables": [{"name": "cliente", "type": "string"}],
|
||||
},
|
||||
body="Cliente: {{cliente}}",
|
||||
assets=None,
|
||||
)
|
||||
result = await mcp.tools["document_generate"](
|
||||
template_name="g",
|
||||
content_md="# Ordine",
|
||||
variables={"cliente": "ACME"},
|
||||
instructions=None,
|
||||
)
|
||||
assert result["markdown"] == "# Out"
|
||||
assert result["model"] == "m"
|
||||
assert result["tokens"]["input"] == 5
|
||||
assert "generation_id" in result
|
||||
@@ -0,0 +1,274 @@
|
||||
import base64
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from mcp_docugen.generation_store import GenerationStore
|
||||
from mcp_docugen.llm_client import LLMResponse, OpenRouterClient
|
||||
from mcp_docugen.models import TemplateFrontmatter, TemplateVariable
|
||||
from mcp_docugen.renderer import (
|
||||
ImageTooLarge,
|
||||
InvalidImageEncoding,
|
||||
InvalidVariableType,
|
||||
MissingVariables,
|
||||
Renderer,
|
||||
)
|
||||
from mcp_docugen.template_store import TemplateStore
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def env(tmp_path):
|
||||
template_store = TemplateStore(base_dir=tmp_path / "templates")
|
||||
generation_store = GenerationStore(
|
||||
db_path=tmp_path / "gen.db",
|
||||
generated_dir=tmp_path / "generated",
|
||||
)
|
||||
await generation_store.init()
|
||||
llm = AsyncMock(spec=OpenRouterClient)
|
||||
llm.chat.return_value = LLMResponse(
|
||||
text="# Output",
|
||||
model="anthropic/claude-sonnet-4",
|
||||
tokens_in=10,
|
||||
tokens_out=20,
|
||||
cost_usd=0.001,
|
||||
latency_ms=100,
|
||||
)
|
||||
renderer = Renderer(
|
||||
template_store=template_store,
|
||||
generation_store=generation_store,
|
||||
llm=llm,
|
||||
public_base_url="https://mcp.example.com",
|
||||
default_model="anthropic/claude-sonnet-4",
|
||||
asset_ttl_days=30,
|
||||
max_image_size_mb=10,
|
||||
)
|
||||
fm = TemplateFrontmatter(
|
||||
name="fattura",
|
||||
description="x",
|
||||
required_variables=[TemplateVariable(name="cliente", type="string")],
|
||||
)
|
||||
await template_store.create(
|
||||
name="fattura", frontmatter=fm, body="Cliente: {{cliente}}"
|
||||
)
|
||||
return renderer, llm
|
||||
|
||||
|
||||
async def test_generate_happy_path_string_var(env):
|
||||
renderer, llm = env
|
||||
result = await renderer.generate(
|
||||
template_name="fattura",
|
||||
content_md="# Ordine",
|
||||
variables={"cliente": "ACME"},
|
||||
instructions=None,
|
||||
)
|
||||
assert result.markdown == "# Output"
|
||||
assert result.model == "anthropic/claude-sonnet-4"
|
||||
assert result.tokens.input == 10
|
||||
call_kwargs = llm.chat.await_args.kwargs
|
||||
assert "ACME" in call_kwargs["user"]
|
||||
assert "Cliente: ACME" in call_kwargs["system"]
|
||||
|
||||
|
||||
async def test_generate_missing_required_variable_raises(env):
|
||||
renderer, _ = env
|
||||
with pytest.raises(MissingVariables) as exc:
|
||||
await renderer.generate(
|
||||
template_name="fattura",
|
||||
content_md="x",
|
||||
variables={},
|
||||
instructions=None,
|
||||
)
|
||||
assert "cliente" in str(exc.value)
|
||||
|
||||
|
||||
async def test_generate_wrong_type_raises(env):
|
||||
renderer, _ = env
|
||||
with pytest.raises(InvalidVariableType):
|
||||
await renderer.generate(
|
||||
template_name="fattura",
|
||||
content_md="x",
|
||||
variables={
|
||||
"cliente": {
|
||||
"kind": "image",
|
||||
"data_b64": "x",
|
||||
"mime": "image/png",
|
||||
}
|
||||
},
|
||||
instructions=None,
|
||||
)
|
||||
|
||||
|
||||
async def test_generate_image_variable_is_saved_and_rewritten(tmp_path):
|
||||
template_store = TemplateStore(base_dir=tmp_path / "templates")
|
||||
generation_store = GenerationStore(
|
||||
db_path=tmp_path / "gen.db",
|
||||
generated_dir=tmp_path / "generated",
|
||||
)
|
||||
await generation_store.init()
|
||||
llm = AsyncMock(spec=OpenRouterClient)
|
||||
llm.chat.return_value = LLMResponse(
|
||||
text="# OK", model="m", tokens_in=1, tokens_out=1, cost_usd=0, latency_ms=10
|
||||
)
|
||||
renderer = Renderer(
|
||||
template_store=template_store,
|
||||
generation_store=generation_store,
|
||||
llm=llm,
|
||||
public_base_url="https://mcp.example.com",
|
||||
default_model="m",
|
||||
asset_ttl_days=30,
|
||||
max_image_size_mb=10,
|
||||
)
|
||||
fm = TemplateFrontmatter(
|
||||
name="report",
|
||||
description="x",
|
||||
required_variables=[TemplateVariable(name="foto", type="image")],
|
||||
)
|
||||
await template_store.create(name="report", frontmatter=fm, body="")
|
||||
|
||||
png = b"\x89PNG\r\n\x1a\n" + b"\x00" * 100
|
||||
img_var = {
|
||||
"kind": "image",
|
||||
"data_b64": base64.b64encode(png).decode(),
|
||||
"mime": "image/png",
|
||||
}
|
||||
result = await renderer.generate(
|
||||
template_name="report",
|
||||
content_md="content",
|
||||
variables={"foto": img_var},
|
||||
instructions=None,
|
||||
)
|
||||
|
||||
assert result.ephemeral_assets_urls
|
||||
url = result.ephemeral_assets_urls[0]
|
||||
assert url.startswith("https://mcp.example.com/generated/")
|
||||
assert url.endswith(".png")
|
||||
gen_dir = tmp_path / "generated" / result.generation_id
|
||||
assert gen_dir.exists()
|
||||
assert any(gen_dir.iterdir())
|
||||
|
||||
call_kwargs = llm.chat.await_args.kwargs
|
||||
assert "https://mcp.example.com/generated/" in call_kwargs["system"]
|
||||
|
||||
|
||||
async def test_image_too_large_raises(env):
|
||||
renderer, _ = env
|
||||
big = b"x" * (11 * 1024 * 1024)
|
||||
img_var = {
|
||||
"kind": "image",
|
||||
"data_b64": base64.b64encode(big).decode(),
|
||||
"mime": "image/png",
|
||||
}
|
||||
fm = TemplateFrontmatter(
|
||||
name="big",
|
||||
description="x",
|
||||
required_variables=[TemplateVariable(name="foto", type="image")],
|
||||
)
|
||||
await renderer.template_store.create(name="big", frontmatter=fm, body="{{foto}}")
|
||||
with pytest.raises(ImageTooLarge):
|
||||
await renderer.generate(
|
||||
template_name="big",
|
||||
content_md="x",
|
||||
variables={"foto": img_var},
|
||||
instructions=None,
|
||||
)
|
||||
|
||||
|
||||
async def test_invalid_base64_raises(env):
|
||||
renderer, _ = env
|
||||
fm = TemplateFrontmatter(
|
||||
name="bad",
|
||||
description="x",
|
||||
required_variables=[TemplateVariable(name="foto", type="image")],
|
||||
)
|
||||
await renderer.template_store.create(name="bad", frontmatter=fm, body="{{foto}}")
|
||||
with pytest.raises(InvalidImageEncoding):
|
||||
await renderer.generate(
|
||||
template_name="bad",
|
||||
content_md="x",
|
||||
variables={
|
||||
"foto": {
|
||||
"kind": "image",
|
||||
"data_b64": "!!!not-base64!!!",
|
||||
"mime": "image/png",
|
||||
}
|
||||
},
|
||||
instructions=None,
|
||||
)
|
||||
|
||||
|
||||
async def test_template_asset_paths_are_rewritten(tmp_path):
|
||||
template_store = TemplateStore(base_dir=tmp_path / "templates")
|
||||
generation_store = GenerationStore(
|
||||
db_path=tmp_path / "gen.db",
|
||||
generated_dir=tmp_path / "generated",
|
||||
)
|
||||
await generation_store.init()
|
||||
llm = AsyncMock(spec=OpenRouterClient)
|
||||
llm.chat.return_value = LLMResponse(
|
||||
text="out", model="m", tokens_in=1, tokens_out=1, cost_usd=0, latency_ms=10
|
||||
)
|
||||
renderer = Renderer(
|
||||
template_store=template_store,
|
||||
generation_store=generation_store,
|
||||
llm=llm,
|
||||
public_base_url="https://mcp.example.com",
|
||||
default_model="m",
|
||||
asset_ttl_days=30,
|
||||
max_image_size_mb=10,
|
||||
)
|
||||
fm = TemplateFrontmatter(name="brand", description="x")
|
||||
assets = [
|
||||
{
|
||||
"filename": "logo.png",
|
||||
"data_b64": base64.b64encode(b"\x89PNG").decode(),
|
||||
"mime": "image/png",
|
||||
}
|
||||
]
|
||||
await template_store.create(
|
||||
name="brand",
|
||||
frontmatter=fm,
|
||||
body="Header  footer",
|
||||
assets=assets,
|
||||
)
|
||||
await renderer.generate(
|
||||
template_name="brand",
|
||||
content_md="x",
|
||||
variables={},
|
||||
instructions=None,
|
||||
)
|
||||
system_prompt = llm.chat.await_args.kwargs["system"]
|
||||
assert "https://mcp.example.com/assets/brand/logo.png" in system_prompt
|
||||
assert "./assets/logo.png" not in system_prompt
|
||||
|
||||
|
||||
async def test_generate_uses_frontmatter_model_override(tmp_path):
|
||||
template_store = TemplateStore(base_dir=tmp_path / "templates")
|
||||
generation_store = GenerationStore(
|
||||
db_path=tmp_path / "gen.db",
|
||||
generated_dir=tmp_path / "generated",
|
||||
)
|
||||
await generation_store.init()
|
||||
llm = AsyncMock(spec=OpenRouterClient)
|
||||
llm.chat.return_value = LLMResponse(
|
||||
text="out",
|
||||
model="openai/gpt-4o",
|
||||
tokens_in=1,
|
||||
tokens_out=1,
|
||||
cost_usd=0,
|
||||
latency_ms=10,
|
||||
)
|
||||
renderer = Renderer(
|
||||
template_store=template_store,
|
||||
generation_store=generation_store,
|
||||
llm=llm,
|
||||
public_base_url="https://mcp.example.com",
|
||||
default_model="anthropic/claude-sonnet-4",
|
||||
asset_ttl_days=30,
|
||||
max_image_size_mb=10,
|
||||
)
|
||||
fm = TemplateFrontmatter(name="x", description="y", model="openai/gpt-4o")
|
||||
await template_store.create(name="x", frontmatter=fm, body="b")
|
||||
await renderer.generate(
|
||||
template_name="x", content_md="c", variables={}, instructions=None
|
||||
)
|
||||
assert llm.chat.await_args.kwargs["model"] == "openai/gpt-4o"
|
||||
Reference in New Issue
Block a user