wehub-resource-sync
203 lines
6.8 KiB
Python
203 lines
6.8 KiB
Python
"""Session persistence for ``ohmo``."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import hashlib
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from uuid import uuid4
|
|
|
|
from openharness.api.usage import UsageSnapshot
|
|
from openharness.engine.messages import ConversationMessage, sanitize_conversation_messages
|
|
from openharness.services.session_backend import SessionBackend
|
|
from openharness.services.session_storage import (
|
|
_persistable_tool_metadata,
|
|
_sanitize_snapshot_payload,
|
|
)
|
|
from openharness.utils.fs import atomic_write_text
|
|
|
|
from ohmo.workspace import get_sessions_dir
|
|
|
|
|
|
def get_session_dir(workspace: str | Path | None = None) -> Path:
|
|
"""Return the ohmo sessions directory."""
|
|
session_dir = get_sessions_dir(workspace)
|
|
session_dir.mkdir(parents=True, exist_ok=True)
|
|
return session_dir
|
|
|
|
|
|
def _session_key_token(session_key: str) -> str:
|
|
return hashlib.sha1(session_key.encode("utf-8")).hexdigest()[:12]
|
|
|
|
|
|
def _session_key_latest_path(workspace: str | Path | None, session_key: str) -> Path:
|
|
session_dir = get_session_dir(workspace)
|
|
token = _session_key_token(session_key)
|
|
return session_dir / f"latest-{token}.json"
|
|
|
|
|
|
def save_session_snapshot(
|
|
*,
|
|
cwd: str | Path,
|
|
workspace: str | Path | None = None,
|
|
model: str,
|
|
system_prompt: str,
|
|
messages: list[ConversationMessage],
|
|
usage: UsageSnapshot,
|
|
session_id: str | None = None,
|
|
session_key: str | None = None,
|
|
tool_metadata: dict[str, object] | None = None,
|
|
) -> Path:
|
|
"""Persist the latest ohmo session snapshot."""
|
|
session_dir = get_session_dir(workspace)
|
|
sid = session_id or uuid4().hex[:12]
|
|
now = time.time()
|
|
messages = sanitize_conversation_messages(messages)
|
|
summary = ""
|
|
for msg in messages:
|
|
if msg.role == "user" and msg.text.strip():
|
|
summary = msg.text.strip()[:80]
|
|
break
|
|
|
|
payload = {
|
|
"app": "ohmo",
|
|
"session_id": sid,
|
|
"session_key": session_key,
|
|
"cwd": str(Path(cwd).resolve()),
|
|
"model": model,
|
|
"system_prompt": system_prompt,
|
|
"messages": [message.model_dump(mode="json") for message in messages],
|
|
"usage": usage.model_dump(),
|
|
"tool_metadata": _persistable_tool_metadata(tool_metadata),
|
|
"created_at": now,
|
|
"summary": summary,
|
|
"message_count": len(messages),
|
|
}
|
|
data = json.dumps(payload, indent=2) + "\n"
|
|
latest_path = session_dir / "latest.json"
|
|
atomic_write_text(latest_path, data)
|
|
if session_key:
|
|
atomic_write_text(_session_key_latest_path(workspace, session_key), data)
|
|
session_path = session_dir / f"session-{sid}.json"
|
|
atomic_write_text(session_path, data)
|
|
return latest_path
|
|
|
|
|
|
def load_latest(workspace: str | Path | None = None) -> dict[str, Any] | None:
|
|
path = get_session_dir(workspace) / "latest.json"
|
|
if not path.exists():
|
|
return None
|
|
return _sanitize_snapshot_payload(json.loads(path.read_text(encoding="utf-8")))
|
|
|
|
|
|
def load_latest_for_session_key(workspace: str | Path | None, session_key: str) -> dict[str, Any] | None:
|
|
path = _session_key_latest_path(workspace, session_key)
|
|
if path.exists():
|
|
return _sanitize_snapshot_payload(json.loads(path.read_text(encoding="utf-8")))
|
|
return None
|
|
|
|
|
|
def list_snapshots(workspace: str | Path | None = None, limit: int = 20) -> list[dict[str, Any]]:
|
|
session_dir = get_session_dir(workspace)
|
|
sessions: list[dict[str, Any]] = []
|
|
for path in sorted(session_dir.glob("session-*.json"), key=lambda p: p.stat().st_mtime, reverse=True):
|
|
try:
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
except (json.JSONDecodeError, OSError):
|
|
continue
|
|
sessions.append(
|
|
{
|
|
"session_id": data.get("session_id", path.stem.replace("session-", "")),
|
|
"summary": data.get("summary", ""),
|
|
"message_count": data.get("message_count", len(data.get("messages", []))),
|
|
"model": data.get("model", ""),
|
|
"created_at": data.get("created_at", path.stat().st_mtime),
|
|
}
|
|
)
|
|
if len(sessions) >= limit:
|
|
break
|
|
return sessions
|
|
|
|
|
|
def load_by_id(workspace: str | Path | None, session_id: str) -> dict[str, Any] | None:
|
|
path = get_session_dir(workspace) / f"session-{session_id}.json"
|
|
if path.exists():
|
|
return _sanitize_snapshot_payload(json.loads(path.read_text(encoding="utf-8")))
|
|
latest = load_latest(workspace)
|
|
if latest and (latest.get("session_id") == session_id or session_id == "latest"):
|
|
return latest
|
|
return None
|
|
|
|
|
|
def export_session_markdown(
|
|
*,
|
|
cwd: str | Path,
|
|
workspace: str | Path | None = None,
|
|
messages: list[ConversationMessage],
|
|
) -> Path:
|
|
path = get_session_dir(workspace) / "transcript.md"
|
|
parts = ["# ohmo Session Transcript"]
|
|
for message in messages:
|
|
parts.append(f"\n## {message.role.capitalize()}\n")
|
|
text = message.text.strip()
|
|
if text:
|
|
parts.append(text)
|
|
atomic_write_text(path, "\n".join(parts).strip() + "\n")
|
|
return path
|
|
|
|
|
|
class OhmoSessionBackend(SessionBackend):
|
|
"""Session backend rooted in ``.ohmo/sessions``."""
|
|
|
|
def __init__(self, workspace: str | Path | None = None) -> None:
|
|
self._workspace = workspace
|
|
|
|
def get_session_dir(self, cwd: str | Path) -> Path:
|
|
return get_session_dir(self._workspace)
|
|
|
|
def save_snapshot(
|
|
self,
|
|
*,
|
|
cwd: str | Path,
|
|
model: str,
|
|
system_prompt: str,
|
|
messages: list[ConversationMessage],
|
|
usage: UsageSnapshot,
|
|
session_id: str | None = None,
|
|
session_key: str | None = None,
|
|
tool_metadata: dict[str, object] | None = None,
|
|
) -> Path:
|
|
return save_session_snapshot(
|
|
cwd=cwd,
|
|
workspace=self._workspace,
|
|
model=model,
|
|
system_prompt=system_prompt,
|
|
messages=messages,
|
|
usage=usage,
|
|
session_id=session_id,
|
|
session_key=session_key,
|
|
tool_metadata=tool_metadata,
|
|
)
|
|
|
|
def load_latest(self, cwd: str | Path) -> dict[str, Any] | None:
|
|
return load_latest(self._workspace)
|
|
|
|
def list_snapshots(self, cwd: str | Path, limit: int = 20) -> list[dict[str, Any]]:
|
|
return list_snapshots(self._workspace, limit=limit)
|
|
|
|
def load_by_id(self, cwd: str | Path, session_id: str) -> dict[str, Any] | None:
|
|
return load_by_id(self._workspace, session_id)
|
|
|
|
def load_latest_for_session_key(self, session_key: str) -> dict[str, Any] | None:
|
|
return load_latest_for_session_key(self._workspace, session_key)
|
|
|
|
def export_markdown(
|
|
self,
|
|
*,
|
|
cwd: str | Path,
|
|
messages: list[ConversationMessage],
|
|
) -> Path:
|
|
return export_session_markdown(cwd=cwd, workspace=self._workspace, messages=messages)
|