| from __future__ import annotations |
| import json, time |
| from datetime import datetime |
| from pathlib import Path |
| from typing import Any, Callable |
|
|
| class UserStateStore: |
| """Server-side persistence for workspace/settings/favorites. |
| |
| Keeps file persistence concerns out of the FastAPI assembly module. |
| """ |
|
|
| def __init__(self, root: Path, validator: Callable[[str], bool], max_favorites: int = 300): |
| self.root = Path(root) |
| self.validator = validator |
| self.max_favorites = int(max_favorites) |
|
|
| def _require_user(self, user_id: str) -> str: |
| uid = str(user_id or "").strip() |
| if not self.validator(uid): |
| raise ValueError("invalid user_id") |
| return uid |
|
|
| def workspace_state_file(self, user_id: str) -> Path: |
| return self.root / f"{self._require_user(user_id)}.workspace.json" |
|
|
| def sanitize_workspace_sessions(self, items: Any) -> list[dict[str, Any]]: |
| out=[] |
| for x in items if isinstance(items,list) else []: |
| if not isinstance(x,dict): continue |
| thread_id=str(x.get("thread_id") or "")[:200] |
| if not thread_id: continue |
| messages=[] |
| for m in (x.get("messages") if isinstance(x.get("messages"),list) else [])[-80:]: |
| if not isinstance(m,dict): continue |
| messages.append({ |
| "role":"assistant" if m.get("role")=="assistant" else "user", |
| "text":str(m.get("text") or "")[:50000], |
| "meta":str(m.get("meta") or "")[:1000], |
| "time":int(m.get("time") or int(time.time()*1000)), |
| }) |
| out.append({ |
| "thread_id":thread_id, |
| "title":str(x.get("title") or "")[:500], |
| "messages":messages, |
| "created_at":int(x.get("created_at") or int(time.time()*1000)), |
| "updated_at":int(x.get("updated_at") or int(time.time()*1000)), |
| }) |
| out.sort(key=lambda x:int(x.get("updated_at") or 0), reverse=True) |
| return out[:100] |
|
|
| def sanitize_workspace_settings(self, value: Any) -> dict[str, Any]: |
| if not isinstance(value,dict): return {} |
| allowed={"fontSize","enterToSend","refreshSeconds","accent","defaultView","compact","reduceMotion","longitudeRange","timeDisplay","exportFormat","answerDetail","historyLimit","confirmDownload","messageActionTrigger","sidebarToolsCollapsed"} |
| return {str(k):v for k,v in value.items() if k in allowed} |
|
|
| def read_workspace(self, user_id: str) -> dict[str, Any]: |
| path=self.workspace_state_file(user_id) |
| if not path.exists(): return {"sessions":[],"settings":{},"updated_at":""} |
| try: raw=json.loads(path.read_text(encoding="utf-8")) |
| except Exception: return {"sessions":[],"settings":{},"updated_at":""} |
| return { |
| "sessions":self.sanitize_workspace_sessions(raw.get("sessions",[])), |
| "settings":self.sanitize_workspace_settings(raw.get("settings",{})), |
| "updated_at":str(raw.get("updated_at") or ""), |
| } |
|
|
| def write_workspace(self, user_id: str, sessions: Any, settings: Any) -> dict[str, Any]: |
| clean_sessions=self.sanitize_workspace_sessions(sessions) |
| clean_settings=self.sanitize_workspace_settings(settings) |
| payload={ |
| "schema":"squid-server-workspace", |
| "schema_version":1, |
| "updated_at":datetime.now().astimezone().isoformat(timespec="seconds"), |
| "sessions":clean_sessions, |
| "settings":clean_settings, |
| } |
| path=self.workspace_state_file(user_id) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp=path.with_suffix(path.suffix+".tmp") |
| tmp.write_text(json.dumps(payload,ensure_ascii=False,indent=2),encoding="utf-8") |
| tmp.replace(path) |
| return payload |
|
|
| def favorite_state_file(self, user_id: str) -> Path: |
| return self.root / f"{self._require_user(user_id)}.favorites.json" |
|
|
| def sanitize_favorite(self, item: Any) -> dict[str, Any] | None: |
| if not isinstance(item, dict): return None |
| kind=str(item.get("kind") or "answer") |
| if kind not in {"dataset","file","answer"}: kind="answer" |
| def cut(key: str, n: int) -> str: return str(item.get(key) or "")[:n] |
| return { |
| "kind":kind, |
| "id":cut("id",240), |
| "key":cut("key",800), |
| "title":cut("title",400), |
| "content":cut("content",60000), |
| "prompt":cut("prompt",10000), |
| "repository":cut("repository",500), |
| "path":cut("path",2000), |
| "created_at":int(item.get("created_at") or int(time.time()*1000)), |
| } |
|
|
| def read_favorites(self, user_id: str) -> list[dict[str, Any]]: |
| path=self.favorite_state_file(user_id) |
| if not path.exists(): return [] |
| try: raw=json.loads(path.read_text(encoding="utf-8")) |
| except Exception: return [] |
| items=raw.get("favorites",[]) if isinstance(raw,dict) else raw |
| out=[] |
| for item in items if isinstance(items,list) else []: |
| clean=self.sanitize_favorite(item) |
| if clean: out.append(clean) |
| return out[:self.max_favorites] |
|
|
| def write_favorites(self, user_id: str, favorites: list[dict[str, Any]]) -> list[dict[str, Any]]: |
| cleaned=[]; seen=set() |
| for item in favorites: |
| clean=self.sanitize_favorite(item) |
| if not clean: continue |
| ident=clean.get("id") or f"{clean.get('kind')}:{clean.get('key')}:{clean.get('title')}" |
| if ident in seen: continue |
| seen.add(ident); cleaned.append(clean) |
| if len(cleaned)>=self.max_favorites: break |
| path=self.favorite_state_file(user_id) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp=path.with_suffix(path.suffix+".tmp") |
| payload={ |
| "schema":"squid-server-favorites", |
| "schema_version":1, |
| "updated_at":datetime.now().astimezone().isoformat(timespec="seconds"), |
| "favorites":cleaned, |
| } |
| tmp.write_text(json.dumps(payload,ensure_ascii=False,indent=2),encoding="utf-8") |
| tmp.replace(path) |
| return cleaned |
|
|
| def storage_mode(self) -> str: |
| try: |
| return "persistent" if str(self.root.resolve()).startswith("/data/") else "server_session" |
| except Exception: |
| return "server_session" |
|
|
