feat: async memorize via Temporal + task status endpoint (#23)

Author: SparkZou

app/main.py

@@ -1,24 +1,33 @@
 """memU Server - FastAPI application entry point."""
 
+import asyncio
 import json
 import logging
+import re
 import uuid
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from pathlib import Path
-from typing import Any
+from typing import Any, cast
 
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import JSONResponse
+from temporalio.client import Client
+from temporalio.service import RPCError, RPCStatusCode
 
 from app.schemas.memory import (
     CategoryObject,
     ClearMemoriesRequest,
     ClearMemoriesResponse,
     ListCategoriesRequest,
     ListCategoriesResponse,
+    MemorizeRequest,
+    MemorizeResponse,
+    TaskStatusResponse,
 )
 from app.services.memu import create_memory_service
+from app.workers.memorize_workflow import MemorizeWorkflow
+from app.workers.worker import TASK_QUEUE
 from config.settings import Settings
 
 logger = logging.getLogger(__name__)
@@ -38,14 +47,37 @@
 storage_dir = Path(settings.STORAGE_PATH)
 
 
+async def _get_temporal_client(app: FastAPI) -> Client:
+    """Return the cached Temporal client, connecting lazily on first call."""
+    # Treat any non-None value as the cached client to support mocking/DI.
+    client = getattr(app.state, "temporal", None)
+    if client is not None:
+        return cast(Client, client)
+    # Create the lock lazily on app.state so it's bound to the running event loop
+    # (module-level asyncio.Lock() can raise RuntimeError in Python 3.13+).
+    lock: asyncio.Lock = getattr(app.state, "_temporal_lock", None) or asyncio.Lock()
+    app.state._temporal_lock = lock
+    async with lock:
+        # Double-check after acquiring the lock
+        client = getattr(app.state, "temporal", None)
+        if client is not None:
+            return cast(Client, client)
+        client = await Client.connect(
+            settings.temporal_url,
+            namespace=settings.TEMPORAL_NAMESPACE,
+        )
+        app.state.temporal = client
+        logger.info("Connected to Temporal at %s", settings.temporal_url)
+        return client
+
+
 @asynccontextmanager
 async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
-    """Initialise MemoryService on startup (defers DB connection until the app runs)."""
+    """Initialise MemoryService on startup. Temporal connects lazily on first use."""
     try:
         storage_dir.mkdir(parents=True, exist_ok=True)
         _app.state.service = create_memory_service(settings)
     except Exception as exc:
-        # Log full traceback for operators and wrap in a clearer startup error
         msg = "Failed to initialize MemoryService during application startup"
         logger.exception(msg)
         raise RuntimeError(msg) from exc
@@ -56,27 +88,121 @@ async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
 
 
 @app.post("/memorize")
-async def memorize(request: Request, payload: dict[str, Any]):
+async def memorize(request: Request, body: MemorizeRequest):
+    """Submit an async memorization task via Temporal workflow."""
+    file_path: Path | None = None
+    workflow_started = False
     try:
-        service = request.app.state.service
-        file_path = storage_dir / f"conversation-{uuid.uuid4().hex}.json"
-        with file_path.open("w", encoding="utf-8") as f:
-            json.dump(payload, f, ensure_ascii=False)
+        # 1. Save conversation to local storage (offload sync I/O to threadpool)
+        task_id = uuid.uuid4().hex
+        file_path = storage_dir / f"conversation-{task_id}.json"
+        data = json.dumps(body.conversation, ensure_ascii=False)
+        await asyncio.to_thread(file_path.write_text, data, "utf-8")
 
-        result = await service.memorize(resource_url=str(file_path), modality="conversation")
-        return JSONResponse(content={"status": "success", "result": result})
+        # 2. Build workflow spec
+        # Pass the filename only; the worker reconstructs the full path
+        # from its own STORAGE_PATH, so it works across containers/hosts.
+        spec = {
+            "task_id": task_id,
+            "resource_url": file_path.name,
+            "user_id": body.user_id,
+            "agent_id": body.agent_id,
+            "override_config": body.override_config,
+        }
+
+        # 3. Start Temporal workflow
+        temporal = await _get_temporal_client(request.app)
+        workflow_id = f"memorize-{task_id}"
+
+        await temporal.start_workflow(
+            MemorizeWorkflow.run,
+            spec,
+            id=workflow_id,
+            task_queue=TASK_QUEUE,
+        )
+        workflow_started = True
+
+        logger.info("Memorize workflow started: %s", workflow_id)
+
+        result = MemorizeResponse(
+            task_id=workflow_id,
+            status="PENDING",
+            message=f"Memorization task submitted for user {body.user_id}",
+        )
+        return JSONResponse(content={"status": "success", "result": result.model_dump()})
+    except Exception as exc:
+        # Only clean up the conversation file if the workflow has NOT started,
+        # because a running workflow still needs its input file.
+        if not workflow_started and file_path is not None and file_path.exists():
+            try:
+                file_path.unlink(missing_ok=True)
+            except Exception:
+                logger.warning(
+                    "Failed to clean up conversation file %s during error handling",
+                    file_path,
+                    exc_info=True,
+                )
+        logger.exception("Failed to submit memorize task")
+        raise HTTPException(status_code=500, detail="Failed to submit memorization task") from exc
+
+
+# Regex for valid memorize workflow IDs: memorize-<32 hex chars>
+_MEMORIZE_WORKFLOW_ID_RE = re.compile(r"^memorize-[0-9a-f]{32}$")
+
+
+@app.get("/memorize/status/{task_id}")
+async def get_memorize_status(request: Request, task_id: str):
+    """Get the status of a memorization task."""
+    if not _MEMORIZE_WORKFLOW_ID_RE.match(task_id):
+        raise HTTPException(
+            status_code=422,
+            detail="task_id must match the format 'memorize-<uuid4hex>' (e.g. memorize-abc123def456...)",
+        )
+    try:
+        temporal = await _get_temporal_client(request.app)
+        handle = temporal.get_workflow_handle(task_id)
+
+        describe = await handle.describe()
+        status = describe.status.name if describe.status else "UNKNOWN"
+
+        detail = None
+        if status == "COMPLETED":
+            result = await handle.result()
+            if isinstance(result, dict):
+                detail = result.get("status", "SUCCESS")
+            elif result is not None:
+                detail = str(result)
+            else:
+                detail = "SUCCESS"
+        elif status == "FAILED":
+            detail = "Task execution failed"
+
+        task_status = TaskStatusResponse(
+            task_id=task_id,
+            status=status,
+            detail=detail,
+        )
+        return JSONResponse(content={"status": "success", "result": task_status.model_dump()})
+    except RPCError as exc:
+        if exc.status == RPCStatusCode.NOT_FOUND:
+            raise HTTPException(status_code=404, detail=f"Task {task_id} not found") from exc
+        logger.exception("Temporal RPC error for task %s", task_id)
+        raise HTTPException(status_code=500, detail="Internal server error") from exc
     except Exception as exc:
-        logger.exception("Memorize request failed")
+        logger.exception("Failed to get task status for %s", task_id)
         raise HTTPException(status_code=500, detail="Internal server error") from exc
 
 
 @app.post("/retrieve")
 async def retrieve(request: Request, payload: dict[str, Any]):
     if "query" not in payload:
         raise HTTPException(status_code=400, detail="Missing 'query' in request body")
+    query = payload["query"]
+    if not isinstance(query, str) or not query.strip():
+        raise HTTPException(status_code=400, detail="'query' must be a non-empty string")
     try:
         service = request.app.state.service
-        result = await service.retrieve([payload["query"]])
+        result = await service.retrieve([query.strip()])
         return JSONResponse(content={"status": "success", "result": result})
     except Exception as exc:
         logger.exception("Retrieve request failed")

app/schemas/memory.py

@@ -1,8 +1,47 @@
 """Request/response schemas for memory endpoints."""
 
-from pydantic import BaseModel, Field, model_validator
+from pydantic import BaseModel, Field, field_validator, model_validator
 
-# ── Clear ──
+
+# ── Memorize (async) ──
+class MemorizeRequest(BaseModel):
+    """Request to memorize a conversation."""
+
+    conversation: dict | list = Field(..., description="Conversation data to memorize")
+    user_id: str = Field(..., min_length=1, description="User ID (non-empty)")
+    agent_id: str = Field(default="", description="Agent ID")
+    override_config: dict | None = Field(default=None, description="Override MemU config")
+
+    @field_validator("user_id", mode="before")
+    @classmethod
+    def strip_user_id(cls, v: str) -> str:
+        """Strip whitespace and reject blank user_id."""
+        if isinstance(v, str):
+            return v.strip()
+        return v
+
+
+class MemorizeResponse(BaseModel):
+    """Response after submitting an async memorize task."""
+
+    task_id: str = Field(..., description="Task ID for tracking (Temporal workflow ID)")
+    status: str = Field(default="PENDING", description="Initial task status")
+    message: str = Field(default="Memorization task submitted", description="Response message")
+
+
+# ── Task Status ──
+class TaskStatusResponse(BaseModel):
+    """Response for task status query."""
+
+    task_id: str = Field(..., description="Task ID")
+    status: str = Field(
+        ...,
+        description=(
+            "Task status from Temporal: RUNNING, COMPLETED, FAILED, UNKNOWN, CANCELED, TERMINATED. "
+            "PENDING is returned only by the initial POST /memorize response before Temporal picks up the task."
+        ),
+    )
+    detail: str | None = Field(default=None, description="Status detail or error message")
 
 
 class ClearMemoriesRequest(BaseModel):
@@ -11,6 +50,15 @@ class ClearMemoriesRequest(BaseModel):
     user_id: str | None = Field(default=None, description="User ID")
     agent_id: str | None = Field(default=None, description="Agent ID")
 
+    @field_validator("user_id", "agent_id", mode="before")
+    @classmethod
+    def strip_whitespace(cls, v: str | None) -> str | None:
+        """Strip whitespace; treat blank strings as None."""
+        if isinstance(v, str):
+            v = v.strip()
+            return v if v else None
+        return v
+
     @model_validator(mode="after")
     def check_user_or_agent(self) -> "ClearMemoriesRequest":
         if self.user_id is None and self.agent_id is None:
@@ -33,9 +81,17 @@ class ClearMemoriesResponse(BaseModel):
 class ListCategoriesRequest(BaseModel):
     """Request to list memory categories."""
 
-    user_id: str = Field(..., description="User ID")
+    user_id: str = Field(..., min_length=1, description="User ID (non-empty)")
     agent_id: str | None = Field(default=None, description="Agent ID")
 
+    @field_validator("user_id", mode="before")
+    @classmethod
+    def strip_user_id(cls, v: str) -> str:
+        """Strip whitespace and reject blank user_id."""
+        if isinstance(v, str):
+            return v.strip()
+        return v
+
 
 class CategoryObject(BaseModel):
     """A single memory category."""

app/workers/memorize_activity.py

@@ -3,6 +3,7 @@
 import json
 import logging
 from datetime import UTC, datetime
+from pathlib import Path
 from typing import Any
 
 from temporalio import activity
@@ -60,6 +61,19 @@ async def task_memorize(spec: dict) -> dict[str, Any]:
             msg = f"override_config must be a dict, got {type(override_config).__name__}"
             raise ApplicationError(msg, non_retryable=True)
 
+        # Validate and resolve resource_url BEFORE building the service so
+        # invalid specs fail fast without opening DB connections or other resources.
+        raw_url = spec["resource_url"]
+        candidate = Path(raw_url)
+        # Reject absolute paths, path traversal, and any directory components.
+        # candidate.name != raw_url catches inputs like "subdir/file.json".
+        if candidate.is_absolute() or ".." in candidate.parts or candidate.name != raw_url:
+            raise ApplicationError(
+                "Invalid resource_url: must be a bare filename without path separators",
+                non_retryable=True,
+            )
+        resource_url = str(Path(settings.STORAGE_PATH).resolve() / candidate.name)
+
         # Build MemoryService with optional config override
         if override_config:
             service = create_memory_service(
@@ -71,7 +85,7 @@ async def task_memorize(spec: dict) -> dict[str, Any]:
 
         # Execute memorization
         result = await service.memorize(
-            resource_url=spec["resource_url"],
+            resource_url=resource_url,
             modality="conversation",
             user={
                 "user_id": spec["user_id"],