Refine memory system user-key flow and search output

memory_system_api/api.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 from fastapi import APIRouter, Depends, HTTPException, Query, status
 
 from .auth import verify_api_key
-from .schemas import MessageIngestRequest, SearchRequest, SessionUserRequest
+from .schemas import MessageIngestRequest, SearchRequest, SessionUserRequest, UserCreateRequest
 from .service import MemorySystemService
 
 
@@ -19,17 +19,31 @@ def get_service() -> MemorySystemService:
     return MemorySystemService()
 
 
+def user_auth_error(exc: PermissionError) -> HTTPException:
+    return HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail=str(exc))
+
+
 @router.get("/health")
 async def health(service: MemorySystemService = Depends(get_service)):
     return await service.health()
 
 
+@router.post("/users")
+async def create_user(request: UserCreateRequest, service: MemorySystemService = Depends(get_service)):
+    return await service.create_user(request.user_id)
+
+
 @router.post("/messages")
-async def ingest_messages(request: MessageIngestRequest, service: MemorySystemService = Depends(get_service)):
+async def ingest_messages(
+    request: MessageIngestRequest,
+    service: MemorySystemService = Depends(get_service),
+):
     try:
         return await service.ingest_messages(request)
     except ValueError as exc:
         raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(exc)) from exc
+    except PermissionError as exc:
+        raise user_auth_error(exc) from exc
 
 
 @router.post("/sessions/{session_id}/commit")
@@ -38,7 +52,10 @@ async def commit_session(
     request: SessionUserRequest,
     service: MemorySystemService = Depends(get_service),
 ):
-    return await service.commit_session(request.user_id, session_id)
+    try:
+        return await service.commit_session(request.user_id, request.user_key, session_id)
+    except PermissionError as exc:
+        raise user_auth_error(exc) from exc
 
 
 @router.post("/sessions/{session_id}/extract")
@@ -47,23 +64,50 @@ async def extract_session(
     request: SessionUserRequest,
     service: MemorySystemService = Depends(get_service),
 ):
-    return await service.extract_session(request.user_id, session_id)
+    try:
+        return await service.extract_session(request.user_id, request.user_key, session_id)
+    except PermissionError as exc:
+        raise user_auth_error(exc) from exc
 
 
 @router.get("/openviking/tasks/{task_id}")
 async def get_openviking_task(
     task_id: str,
     user_id: str = Query(min_length=1),
+    user_key: str = Query(min_length=1),
+    session_id: str | None = Query(default=None, min_length=1),
     service: MemorySystemService = Depends(get_service),
 ):
-    return await service.get_openviking_task(user_id, task_id)
+    try:
+        return await service.get_openviking_task(
+            user_id,
+            user_key,
+            task_id,
+            session_id=session_id,
+        )
+    except PermissionError as exc:
+        raise user_auth_error(exc) from exc
 
 
 @router.post("/search")
-async def search(request: SearchRequest, service: MemorySystemService = Depends(get_service)):
-    return await service.search(request)
+async def search(
+    request: SearchRequest,
+    service: MemorySystemService = Depends(get_service),
+):
+    try:
+        return await service.search(request)
+    except PermissionError as exc:
+        raise user_auth_error(exc) from exc
 
 
 @router.get("/users/{user_id}/profile")
-async def get_profile(user_id: str, service: MemorySystemService = Depends(get_service)):
+async def get_profile(
+    user_id: str,
+    user_key: str = Query(min_length=1),
+    service: MemorySystemService = Depends(get_service),
+):
+    try:
+        service.openviking.credential_for_user(user_id, user_key)
+    except PermissionError as exc:
+        raise user_auth_error(exc) from exc
     return await service.get_profile(user_id)

memory_system_api/clients.py

@@ -8,7 +8,7 @@ from typing import Any
 import httpx
 
 from .config import get_config
-from .store import OpenVikingUserKeyStore
+from .store import ADMIN_ACCOUNT_ID, ADMIN_USER_ID, OpenVikingUserKeyStore
 
 
 @dataclass(frozen=True)
@@ -16,6 +16,8 @@ class OpenVikingCredential:
     api_key: str
     account_id: str | None = None
     user_id: str | None = None
+    agent_id: str | None = None
+    user_key_auth: bool = False
 
 
 class OpenVikingMemorySystemClient:
@@ -33,36 +35,83 @@ class OpenVikingMemorySystemClient:
             response.raise_for_status()
             return response.json()
 
-    async def ensure_user(self, user_id: str) -> OpenVikingCredential:
-        existing = self.store.get_user_key(user_id)
-        if existing:
-            return OpenVikingCredential(api_key=existing)
-
+    async def create_account(self, account_id: str = ADMIN_ACCOUNT_ID, admin_user_id: str = ADMIN_USER_ID) -> dict[str, Any]:
         async with self._client(self.root_key) as client:
             response = await client.post(
                 "/api/v1/admin/accounts",
-                json={"account_id": user_id, "admin_user_id": user_id},
+                json=self._create_account_payload(account_id, admin_user_id),
             )
             if response.status_code == 409:
-                return self.root_credential(user_id)
+                return response.json()
             response.raise_for_status()
             data = response.json()
 
         user_key = self._extract_user_key(data)
-        if not user_key:
-            return self.root_credential(user_id)
-        self.store.save_user_key(user_id, user_key)
-        return OpenVikingCredential(api_key=user_key)
+        if user_key:
+            self.store.save_account_key(account_id, admin_user_id, user_key)
+            self.store.save_user_key(admin_user_id, user_key)
+        return data
 
-    def root_credential(self, user_id: str) -> OpenVikingCredential:
-        return OpenVikingCredential(api_key=self.root_key, account_id=user_id, user_id=user_id)
+    async def ensure_admin_workspace(self) -> None:
+        if self.store.get_account_key(ADMIN_ACCOUNT_ID):
+            return
+        await self.create_account(ADMIN_ACCOUNT_ID, ADMIN_USER_ID)
+
+    async def create_user(self, user_id: str, role: str = "user") -> dict[str, Any]:
+        existing = self.store.get_user_key(user_id)
+        if existing:
+            return {"status": "ok", "result": {"account_id": ADMIN_ACCOUNT_ID, "user_id": user_id, "user_key": existing}}
+
+        await self.ensure_admin_workspace()
+        if user_id == ADMIN_USER_ID:
+            user_key = self.store.get_user_key(user_id)
+            return {"status": "ok", "result": {"account_id": ADMIN_ACCOUNT_ID, "user_id": user_id, "user_key": user_key}}
+
+        async with self._client(self.root_key) as client:
+            response = await client.post(
+                f"/api/v1/admin/accounts/{ADMIN_ACCOUNT_ID}/users",
+                json={"user_id": user_id, "role": role},
+            )
+            response.raise_for_status()
+            data = response.json()
+
+        user_key = self._extract_user_key(data)
+        if user_key:
+            self.store.save_user_key(user_id, user_key)
+        return data
+
+    def credential_for_user(
+        self,
+        user_id: str,
+        user_key: str,
+        agent_id: str | None = None,
+    ) -> OpenVikingCredential:
+        if not self.store.user_key_matches(user_id, user_key):
+            raise PermissionError("Invalid user key")
+        return self.user_credential(user_key, user_id, agent_id=agent_id)
+
+    def user_credential(
+        self,
+        user_key: str,
+        user_id: str,
+        agent_id: str | None = None,
+    ) -> OpenVikingCredential:
+        return OpenVikingCredential(
+            api_key=user_key,
+            account_id=ADMIN_ACCOUNT_ID,
+            user_id=user_id,
+            agent_id=agent_id,
+            user_key_auth=True,
+        )
 
     async def ensure_session(self, credential: OpenVikingCredential | str, session_id: str) -> dict[str, Any]:
         async with self._credential_client(credential) as client:
             response = await client.post("/api/v1/sessions", json={"session_id": session_id})
             if response.status_code in {409, 422}:
+                self._save_session(credential, session_id)
                 return {"session_id": session_id, "status": "exists"}
             response.raise_for_status()
+            self._save_session(credential, session_id)
             return response.json()
 
     async def append_message(
@@ -78,9 +127,14 @@ class OpenVikingMemorySystemClient:
 
     async def commit_session(self, credential: OpenVikingCredential | str, session_id: str) -> dict[str, Any]:
         async with self._credential_client(credential) as client:
-            response = await client.post(f"/api/v1/sessions/{session_id}/commit")
+            response = await client.post(
+                f"/api/v1/sessions/{session_id}/commit",
+                json={"keep_recent_count": 0},
+            )
             response.raise_for_status()
-            return response.json()
+            data = response.json()
+            self._save_commit_metadata(credential, session_id, data)
+            return data
 
     async def extract_session(self, credential: OpenVikingCredential | str, session_id: str) -> dict[str, Any]:
         async with self._credential_client(credential) as client:
@@ -94,15 +148,15 @@ class OpenVikingMemorySystemClient:
             response.raise_for_status()
             return response.json()
 
-    async def find(
-        self, credential: OpenVikingCredential | str, user_id: str, query: str, limit: int
-    ) -> dict[str, Any]:
+    async def find(self, credential: OpenVikingCredential | str, query: str, limit: int) -> dict[str, Any]:
+        user_id = credential.user_id if isinstance(credential, OpenVikingCredential) else None
+        target_uri = f"viking://user/{user_id}/memories/" if user_id else "viking://user/memories/"
         async with self._credential_client(credential) as client:
             response = await client.post(
                 "/api/v1/search/find",
                 json={
                     "query": query,
-                    "target_uri": f"viking://user/{user_id}/memories/",
+                    "target_uri": target_uri,
                     "limit": limit,
                 },
             )
@@ -115,6 +169,12 @@ class OpenVikingMemorySystemClient:
         payload: dict[str, Any] = {"query": query, "limit": limit}
         if session_id:
             payload["session_id"] = session_id
+        if isinstance(credential, OpenVikingCredential) and credential.user_id:
+            payload["target_uri"] = (
+                f"viking://user/{credential.user_id}/{session_id}"
+                if session_id
+                else f"viking://user/{credential.user_id}/memories/"
+            )
         async with self._credential_client(credential) as client:
             response = await client.post("/api/v1/search/search", json=payload)
             response.raise_for_status()
@@ -123,15 +183,23 @@ class OpenVikingMemorySystemClient:
     def _credential_client(self, credential: OpenVikingCredential | str) -> httpx.AsyncClient:
         if isinstance(credential, str):
             return self._client(credential)
+        if credential.user_key_auth:
+            return self._client(credential.api_key)
         headers = {}
         if credential.account_id:
             headers["X-OpenViking-Account"] = credential.account_id
         if credential.user_id:
             headers["X-OpenViking-User"] = credential.user_id
+        if credential.agent_id:
+            headers["X-OpenViking-Agent"] = credential.agent_id
         return self._client(credential.api_key, headers)
 
     def _client(self, api_key: str, extra_headers: dict[str, str] | None = None) -> httpx.AsyncClient:
-        headers = {"X-API-Key": api_key, "Content-Type": "application/json"}
+        headers = {"Content-Type": "application/json"}
+        if api_key == self.root_key:
+            headers["X-API-Key"] = api_key
+        else:
+            headers["Authorization"] = f"Bearer {api_key}"
         if extra_headers:
             headers.update(extra_headers)
         return httpx.AsyncClient(
@@ -146,6 +214,38 @@ class OpenVikingMemorySystemClient:
         value = result.get("user_key") if isinstance(result, dict) else None
         return str(value) if value else None
 
+    def _create_account_payload(self, account_id: str, admin_user_id: str) -> dict[str, Any]:
+        return {
+            "account_id": account_id,
+            "admin_user_id": admin_user_id,
+        }
+
+    def _save_session(self, credential: OpenVikingCredential | str, session_id: str) -> None:
+        if isinstance(credential, OpenVikingCredential) and credential.user_id:
+            self.store.save_session(credential.user_id, session_id)
+
+    def _save_commit_metadata(
+        self,
+        credential: OpenVikingCredential | str,
+        session_id: str,
+        data: dict[str, Any],
+    ) -> None:
+        if not isinstance(credential, OpenVikingCredential) or not credential.user_id:
+            return
+        result = data.get("result") if isinstance(data.get("result"), dict) else data
+        if not isinstance(result, dict):
+            return
+        task_id = result.get("task_id")
+        if not task_id:
+            return
+        archive_uri = result.get("archive_uri")
+        self.store.save_task(
+            user_id=credential.user_id,
+            session_id=session_id,
+            task_id=str(task_id),
+            archive_uri=str(archive_uri) if archive_uri else None,
+        )
+
 
 class EverOSMemorySystemClient:
     def __init__(self) -> None:

memory_system_api/schemas.py

@@ -11,6 +11,7 @@ OperationStatus = Literal["success", "partial_success", "failed"]
 
 class MessageIngestRequest(BaseModel):
     user_id: str = Field(min_length=1)
+    user_key: str = Field(min_length=1)
     session_id: str = Field(min_length=1)
     user_message: str | None = None
     assistant_message: str | None = None
@@ -20,10 +21,12 @@ class MessageIngestRequest(BaseModel):
 
 class SessionUserRequest(BaseModel):
     user_id: str = Field(min_length=1)
+    user_key: str = Field(min_length=1)
 
 
 class SearchRequest(BaseModel):
     user_id: str = Field(min_length=1)
+    user_key: str = Field(min_length=1)
     session_id: str | None = None
     query: str = Field(min_length=1)
     use_llm: bool = False
@@ -36,6 +39,16 @@ class BackendStatus(BaseModel):
     error: str | None = None
 
 
+class UserCreateRequest(BaseModel):
+    user_id: str = Field(min_length=1)
+
+
+class AccountResponse(BaseModel):
+    status: OperationStatus
+    account: Any = None
+    backends: dict[str, BackendStatus]
+
+
 class MessageIngestResponse(BaseModel):
     status: OperationStatus
     message_count: int

memory_system_api/service.py

@@ -6,6 +6,7 @@ from typing import Any, Awaitable, Callable
 
 from .clients import EverOSMemorySystemClient, OpenVikingMemorySystemClient
 from .schemas import (
+    AccountResponse,
     BackendStatus,
     CommitResponse,
     ExtractResponse,
@@ -22,12 +23,21 @@ class MemorySystemService:
         self.openviking = openviking or OpenVikingMemorySystemClient()
         self.everos = everos or EverOSMemorySystemClient()
 
+    async def create_user(self, user_id: str) -> AccountResponse:
+        backends = {"openviking": await self._capture(lambda: self.openviking.create_user(user_id))}
+        account = backends["openviking"].result if backends["openviking"].status == "success" else None
+        return AccountResponse(status=self._aggregate_status(backends), account=account, backends=backends)
+
     async def ingest_messages(self, request: MessageIngestRequest) -> MessageIngestResponse:
         messages = self._messages_from_request(request)
         if not messages:
             raise ValueError("at least one message is required")
 
-        credential = await self.openviking.ensure_user(request.user_id)
+        credential = self.openviking.credential_for_user(
+            request.user_id,
+            request.user_key,
+            agent_id=request.session_id,
+        )
         await self.openviking.ensure_session(credential, request.session_id)
 
         async def write_openviking() -> list[dict[str, Any]]:
@@ -53,11 +63,11 @@ class MemorySystemService:
             backends=backends,
         )
 
-    async def commit_session(self, user_id: str, session_id: str) -> CommitResponse:
-        user_key = await self.openviking.ensure_user(user_id)
+    async def commit_session(self, user_id: str, user_key: str, session_id: str) -> CommitResponse:
+        credential = self.openviking.credential_for_user(user_id, user_key, agent_id=session_id)
 
         async def commit_openviking() -> dict[str, Any]:
-            return await self.openviking.commit_session(user_key, session_id)
+            return await self.openviking.commit_session(credential, session_id)
 
         async def flush_everos() -> dict[str, Any]:
             return await self.everos.flush(user_id, session_id)
@@ -65,25 +75,35 @@ class MemorySystemService:
         backends = await self._run_backends(openviking=commit_openviking, everos=flush_everos)
         return CommitResponse(status=self._aggregate_status(backends), backends=backends)
 
-    async def extract_session(self, user_id: str, session_id: str) -> ExtractResponse:
-        user_key = await self.openviking.ensure_user(user_id)
+    async def extract_session(self, user_id: str, user_key: str, session_id: str) -> ExtractResponse:
+        credential = self.openviking.credential_for_user(user_id, user_key, agent_id=session_id)
         backends = {
-            "openviking": await self._capture(lambda: self.openviking.extract_session(user_key, session_id)),
+            "openviking": await self._capture(lambda: self.openviking.extract_session(credential, session_id)),
         }
         return ExtractResponse(status=self._aggregate_status(backends), backends=backends)
 
-    async def get_openviking_task(self, user_id: str, task_id: str) -> dict[str, Any]:
-        user_key = await self.openviking.ensure_user(user_id)
-        return await self.openviking.get_task(user_key, task_id)
+    async def get_openviking_task(
+        self,
+        user_id: str,
+        user_key: str,
+        task_id: str,
+        session_id: str | None = None,
+    ) -> dict[str, Any]:
+        credential = self.openviking.credential_for_user(user_id, user_key, agent_id=session_id)
+        return await self.openviking.get_task(credential, task_id)
 
     async def search(self, request: SearchRequest) -> SearchResponse:
-        user_key = await self.openviking.ensure_user(request.user_id)
+        credential = self.openviking.credential_for_user(
+            request.user_id,
+            request.user_key,
+            agent_id=request.session_id,
+        )
         everos_method = "agentic" if request.use_llm else "hybrid"
 
         async def search_openviking() -> dict[str, Any]:
             if request.use_llm:
-                return await self.openviking.search(user_key, request.session_id, request.query, request.limit)
-            return await self.openviking.find(user_key, request.user_id, request.query, request.limit)
+                return await self.openviking.search(credential, request.session_id, request.query, request.limit)
+            return await self.openviking.find(credential, request.query, request.limit)
 
         async def search_everos() -> dict[str, Any]:
             return await self.everos.search(
@@ -97,7 +117,12 @@ class MemorySystemService:
         backends = await self._run_backends(openviking=search_openviking, everos=search_everos)
         backends = self._remove_vectors_from_backends(backends)
         items = self._merge_search_items(backends)
-        return SearchResponse(status=self._aggregate_status(backends), items=items[: request.limit], backends=backends)
+        compact_backends = self._compact_search_backends(backends)
+        return SearchResponse(
+            status=self._aggregate_status(backends),
+            items=items[: request.limit],
+            backends=compact_backends,
+        )
 
     async def get_profile(self, user_id: str) -> ProfileResponse:
         backends = {"everos": await self._capture(lambda: self.everos.get_profile(user_id))}
@@ -157,7 +182,11 @@ class MemorySystemService:
         for key in ("memories", "resources", "episodes", "profiles", "raw_messages"):
             values = data.get(key)
             if isinstance(values, list):
-                raw_items.extend(item for item in values if isinstance(item, dict))
+                raw_items.extend(
+                    self._compact_search_item(backend_name, key, item)
+                    for item in values
+                    if isinstance(item, dict)
+                )
         return [self._with_backend(backend_name, item) for item in raw_items]
 
     def _with_backend(self, backend_name: str, item: dict[str, Any]) -> dict[str, Any]:
@@ -165,6 +194,72 @@ class MemorySystemService:
             return item
         return {"source_backend": backend_name, **item}
 
+    def _compact_search_item(self, backend_name: str, collection: str, item: dict[str, Any]) -> dict[str, Any]:
+        if backend_name == "everos":
+            fields = (
+                "id",
+                "user_id",
+                "session_id",
+                "timestamp",
+                "summary",
+                "score",
+            )
+            compact = {"memory_type": self._singular_memory_type(collection)}
+            compact.update({field: item[field] for field in fields if field in item and item[field] is not None})
+            return compact
+        return item
+
+    def _singular_memory_type(self, collection: str) -> str:
+        names = {
+            "memories": "memory",
+            "resources": "resource",
+            "episodes": "episode",
+            "profiles": "profile",
+            "raw_messages": "raw_message",
+        }
+        return names.get(collection, collection)
+
+    def _compact_search_backends(self, backends: dict[str, BackendStatus]) -> dict[str, BackendStatus]:
+        return {
+            name: backend.model_copy(update={"result": self._compact_backend_result(name, backend.result)})
+            for name, backend in backends.items()
+        }
+
+    def _compact_backend_result(self, backend_name: str, result: Any) -> Any:
+        if backend_name == "everos":
+            data = result.get("data") if isinstance(result, dict) and isinstance(result.get("data"), dict) else result
+            if not isinstance(data, dict):
+                return result
+            compact: dict[str, Any] = {
+                "counts": {
+                    key: len(data.get(key) or [])
+                    for key in ("episodes", "profiles", "raw_messages")
+                    if isinstance(data.get(key), list)
+                }
+            }
+            if "query" in data:
+                compact["query"] = data["query"]
+            return compact
+
+        if backend_name == "openviking":
+            data = result.get("result") if isinstance(result, dict) and isinstance(result.get("result"), dict) else result
+            if not isinstance(data, dict):
+                return result
+            compact = {
+                "status": result.get("status") if isinstance(result, dict) else None,
+                "total": data.get("total"),
+                "counts": {
+                    key: len(data.get(key) or [])
+                    for key in ("memories", "resources", "skills")
+                    if isinstance(data.get(key), list)
+                },
+            }
+            if "query_plan" in data:
+                compact["query_plan"] = data["query_plan"]
+            return {key: value for key, value in compact.items() if value is not None}
+
+        return result
+
     def _remove_vectors_from_backends(self, backends: dict[str, BackendStatus]) -> dict[str, BackendStatus]:
         return {
             name: backend.model_copy(update={"result": self._remove_vectors(backend.result)})

memory_system_api/store.py

@@ -2,21 +2,67 @@
 from __future__ import annotations
 
 import sqlite3
+import hmac
 from datetime import datetime, timezone
 from pathlib import Path
 
+ADMIN_ACCOUNT_ID = "admin"
+ADMIN_USER_ID = "admin"
+
 
 class OpenVikingUserKeyStore:
     def __init__(self, sqlite_path: str) -> None:
         self.sqlite_path = sqlite_path
         self._ensure_table()
 
+    def get_account_key(self, account_id: str) -> str | None:
+        with self._connect() as conn:
+            row = conn.execute(
+                "SELECT account_key FROM memory_system_openviking_accounts WHERE account_id = ?",
+                (account_id,),
+            ).fetchone()
+            if row is None:
+                row = conn.execute(
+                    """
+                    SELECT user_key FROM memory_system_openviking_users
+                    WHERE account_id = ?
+                    ORDER BY created_at ASC
+                    LIMIT 1
+                    """,
+                    (account_id,),
+                ).fetchone()
+        return str(row[0]) if row else None
+
+    def account_key_matches(self, account_id: str, account_key: str) -> bool:
+        expected = self.get_account_key(account_id)
+        return bool(expected and hmac.compare_digest(expected, account_key))
+
+    def save_account_key(self, account_id: str, admin_user_id: str, account_key: str) -> None:
+        now = datetime.now(timezone.utc).isoformat()
+        with self._connect() as conn:
+            conn.execute(
+                """
+                INSERT INTO memory_system_openviking_accounts (account_id, admin_user_id, account_key, created_at, updated_at)
+                VALUES (?, ?, ?, ?, ?)
+                ON CONFLICT(account_id) DO UPDATE SET
+                    admin_user_id = excluded.admin_user_id,
+                    account_key = excluded.account_key,
+                    updated_at = excluded.updated_at
+                """,
+                (account_id, admin_user_id, account_key, now, now),
+            )
+
     def get_user_key(self, user_id: str) -> str | None:
         with self._connect() as conn:
             row = conn.execute(
                 "SELECT user_key FROM memory_system_openviking_users WHERE user_id = ?",
                 (user_id,),
             ).fetchone()
+            if row is None:
+                row = conn.execute(
+                    "SELECT user_key FROM memory_system_openviking_users WHERE user_id = ?",
+                    (self._legacy_store_key(ADMIN_ACCOUNT_ID, user_id),),
+                ).fetchone()
         return str(row[0]) if row else None
 
     def save_user_key(self, user_id: str, user_key: str) -> None:
@@ -30,13 +76,109 @@ class OpenVikingUserKeyStore:
                     user_key = excluded.user_key,
                     updated_at = excluded.updated_at
                 """,
-                (user_id, user_id, user_key, now, now),
+                (user_id, ADMIN_ACCOUNT_ID, user_key, now, now),
             )
 
+    def user_key_matches(self, user_id: str, user_key: str) -> bool:
+        expected = self.get_user_key(user_id)
+        return bool(expected and hmac.compare_digest(expected, user_key))
+
+    def save_session(self, user_id: str, session_id: str) -> None:
+        now = datetime.now(timezone.utc).isoformat()
+        with self._connect() as conn:
+            conn.execute(
+                """
+                INSERT INTO memory_system_openviking_sessions
+                    (user_id, session_id, latest_task_id, latest_archive_uri, created_at, updated_at)
+                VALUES (?, ?, NULL, NULL, ?, ?)
+                ON CONFLICT(user_id, session_id) DO UPDATE SET
+                    updated_at = excluded.updated_at
+                """,
+                (user_id, session_id, now, now),
+            )
+
+    def get_session(self, user_id: str, session_id: str) -> dict[str, str | None] | None:
+        with self._connect() as conn:
+            row = conn.execute(
+                """
+                SELECT user_id, session_id, latest_task_id, latest_archive_uri
+                FROM memory_system_openviking_sessions
+                WHERE user_id = ? AND session_id = ?
+                """,
+                (user_id, session_id),
+            ).fetchone()
+        if row is None:
+            return None
+        return {
+            "user_id": str(row[0]),
+            "session_id": str(row[1]),
+            "latest_task_id": str(row[2]) if row[2] is not None else None,
+            "latest_archive_uri": str(row[3]) if row[3] is not None else None,
+        }
+
+    def save_task(self, user_id: str, session_id: str, task_id: str, archive_uri: str | None) -> None:
+        now = datetime.now(timezone.utc).isoformat()
+        with self._connect() as conn:
+            conn.execute(
+                """
+                INSERT INTO memory_system_openviking_tasks
+                    (task_id, user_id, session_id, archive_uri, created_at, updated_at)
+                VALUES (?, ?, ?, ?, ?, ?)
+                ON CONFLICT(task_id) DO UPDATE SET
+                    user_id = excluded.user_id,
+                    session_id = excluded.session_id,
+                    archive_uri = excluded.archive_uri,
+                    updated_at = excluded.updated_at
+                """,
+                (task_id, user_id, session_id, archive_uri, now, now),
+            )
+            conn.execute(
+                """
+                INSERT INTO memory_system_openviking_sessions
+                    (user_id, session_id, latest_task_id, latest_archive_uri, created_at, updated_at)
+                VALUES (?, ?, ?, ?, ?, ?)
+                ON CONFLICT(user_id, session_id) DO UPDATE SET
+                    latest_task_id = excluded.latest_task_id,
+                    latest_archive_uri = excluded.latest_archive_uri,
+                    updated_at = excluded.updated_at
+                """,
+                (user_id, session_id, task_id, archive_uri, now, now),
+            )
+
+    def get_task(self, task_id: str) -> dict[str, str | None] | None:
+        with self._connect() as conn:
+            row = conn.execute(
+                """
+                SELECT task_id, user_id, session_id, archive_uri
+                FROM memory_system_openviking_tasks
+                WHERE task_id = ?
+                """,
+                (task_id,),
+            ).fetchone()
+        if row is None:
+            return None
+        return {
+            "task_id": str(row[0]),
+            "user_id": str(row[1]),
+            "session_id": str(row[2]),
+            "archive_uri": str(row[3]) if row[3] is not None else None,
+        }
+
     def _ensure_table(self) -> None:
         path = Path(self.sqlite_path)
         path.parent.mkdir(parents=True, exist_ok=True)
         with self._connect() as conn:
+            conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS memory_system_openviking_accounts (
+                    account_id TEXT PRIMARY KEY,
+                    admin_user_id TEXT NOT NULL,
+                    account_key TEXT NOT NULL,
+                    created_at TEXT NOT NULL,
+                    updated_at TEXT NOT NULL
+                )
+                """
+            )
             conn.execute(
                 """
                 CREATE TABLE IF NOT EXISTS memory_system_openviking_users (
@@ -48,6 +190,34 @@ class OpenVikingUserKeyStore:
                 )
                 """
             )
+            conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS memory_system_openviking_sessions (
+                    user_id TEXT NOT NULL,
+                    session_id TEXT NOT NULL,
+                    latest_task_id TEXT,
+                    latest_archive_uri TEXT,
+                    created_at TEXT NOT NULL,
+                    updated_at TEXT NOT NULL,
+                    PRIMARY KEY (user_id, session_id)
+                )
+                """
+            )
+            conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS memory_system_openviking_tasks (
+                    task_id TEXT PRIMARY KEY,
+                    user_id TEXT NOT NULL,
+                    session_id TEXT NOT NULL,
+                    archive_uri TEXT,
+                    created_at TEXT NOT NULL,
+                    updated_at TEXT NOT NULL
+                )
+                """
+            )
 
     def _connect(self) -> sqlite3.Connection:
         return sqlite3.connect(self.sqlite_path)
+
+    def _legacy_store_key(self, account_id: str, user_id: str) -> str:
+        return f"{account_id}:{user_id}"