agent-memory-git 0.1.0__py3-none-any.whl

agent_memory/context.py

@@ -0,0 +1,73 @@
+"""Where a session is running: folder, host and git identity."""
+
+import os
+import re
+import socket
+import subprocess
+
+_GIT_ENV = {**os.environ, "GIT_OPTIONAL_LOCKS": "0", "GIT_TERMINAL_PROMPT": "0"}
+
+
+def _git(cwd, *args):
+    try:
+        r = subprocess.run(["git", "-C", cwd, *args], capture_output=True, text=True,
+                           timeout=3, env=_GIT_ENV)
+    except (OSError, subprocess.SubprocessError):
+        return ""
+    return r.stdout.strip() if r.returncode == 0 else ""
+
+
+def normalize_remote(url):
+    """Reduce a git remote URL to host/owner/repo, dropping credentials."""
+    url = (url or "").strip()
+    if not url:
+        return ""
+    m = re.match(r"^[\w.-]+@([^:/]+):(.+)$", url)
+    if not m:
+        m = re.match(r"^[a-z][a-z0-9+.-]*://(?:[^@/]+@)?([^/:]+)(?::\d+)?/(.+)$", url, re.I)
+    if not m:
+        return url.lower()
+    host, path = m.groups()
+    path = path.strip("/")
+    if path.endswith(".git"):
+        path = path[:-4]
+    return f"{host}/{path}".lower()
+
+
+def hostname():
+    h = socket.gethostname()
+    return h[:-6] if h.endswith(".local") else h
+
+
+def detect(cwd=None):
+    cwd = os.path.realpath(os.path.expanduser(cwd or os.getcwd()))
+    ctx = {"host": hostname(), "cwd": cwd}
+    if not os.path.isdir(cwd):
+        return ctx
+    root = _git(cwd, "rev-parse", "--show-toplevel")
+    if not root:
+        return ctx
+    ctx["git_root"] = os.path.realpath(root)
+    remote = _git(cwd, "remote", "get-url", "origin")
+    if not remote:
+        names = _git(cwd, "remote").split()
+        remote = _git(cwd, "remote", "get-url", names[0]) if names else ""
+    if remote:
+        ctx["git_remote"] = normalize_remote(remote)
+    branch = _git(cwd, "rev-parse", "--abbrev-ref", "HEAD")
+    if branch and branch != "HEAD":
+        ctx["git_branch"] = branch
+    commit = _git(cwd, "rev-parse", "--short", "HEAD")
+    if commit:
+        ctx["git_commit"] = commit
+    return ctx
+
+
+def describe(ctx):
+    parts = [f"cwd={ctx.get('cwd', '?')}"]
+    if ctx.get("git_remote"):
+        parts.append(f"remote={ctx['git_remote']}")
+    elif ctx.get("git_root"):
+        parts.append(f"git_root={ctx['git_root']}")
+    parts.append(f"host={ctx.get('host', '?')}")
+    return ", ".join(parts)

agent_memory/embed_worker.py

@@ -0,0 +1,90 @@
+"""Embed texts with a local ONNX embedding model. Runs in its own environment.
+
+`mem ui` starts this with `uv run --with onnxruntime --with tokenizers`, so
+neither package has to be installed where the hooks run. It reads a JSON job
+on stdin and writes an .npy matrix of L2-normalized vectors:
+
+    {"model_dir": ..., "texts": [...], "out": "vectors.npy", "max_tokens": 512}
+
+With `--serve MODEL_DIR` it stays up and answers one {"text": ...} per line
+with {"vector": [...]}, for search queries.
+
+The model directory holds model_quantized.onnx (or model.onnx) and a
+Hugging Face tokenizer.json. Pooling follows the model family: last token for
+decoder models such as Qwen3-Embedding, attention-masked mean otherwise.
+"""
+
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import onnxruntime as ort
+from tokenizers import Tokenizer
+
+
+def load(model_dir):
+    model_dir = Path(model_dir)
+    onnx = next(p for p in (model_dir / "model_quantized.onnx", model_dir / "model.onnx") if p.exists())
+    opts = ort.SessionOptions()
+    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    session = ort.InferenceSession(str(onnx), opts, providers=["CPUExecutionProvider"])
+    tok = Tokenizer.from_file(str(model_dir / "tokenizer.json"))
+    config = json.loads((model_dir / "config.json").read_text())
+    return session, tok, config
+
+
+def embed_one(session, tok, config, text, max_tokens):
+    ids = tok.encode(text).ids[: max_tokens - 1]
+    decoder = config.get("model_type", "").startswith("qwen")
+    eos = config.get("eos_token_id")
+    if decoder and isinstance(eos, int) and (not ids or ids[-1] != eos):
+        ids.append(eos)  # last-token pooling reads the end-of-text position
+    n = len(ids)
+    feed = {}
+    for inp in session.get_inputs():
+        name = inp.name
+        if name == "input_ids":
+            feed[name] = np.array([ids], dtype=np.int64)
+        elif name == "attention_mask":
+            feed[name] = np.ones((1, n), dtype=np.int64)
+        elif name == "position_ids":
+            feed[name] = np.arange(n, dtype=np.int64)[None, :]
+        elif name == "token_type_ids":
+            feed[name] = np.zeros((1, n), dtype=np.int64)
+        elif name.startswith("past_key_values"):
+            shape = [1 if d == "batch_size" else 0 if isinstance(d, str) else d for d in inp.shape]
+            feed[name] = np.zeros(shape, dtype=np.float16 if "float16" in inp.type else np.float32)
+        else:
+            raise SystemExit(f"unsupported model input {name}")
+    hidden = session.run(None, feed)[0][0]  # (tokens, dim)
+    vec = hidden[-1] if decoder else hidden.mean(axis=0)
+    return vec / (np.linalg.norm(vec) or 1.0)
+
+
+def serve(model_dir):
+    """One JSON request per line, {"text": ...} -> {"vector": [...]}, until stdin closes."""
+    session, tok, config = load(model_dir)
+    for line in sys.stdin:
+        text = json.loads(line)["text"]
+        vec = embed_one(session, tok, config, text, 256)
+        sys.stdout.write(json.dumps({"vector": [round(float(v), 6) for v in vec]}) + "\n")
+        sys.stdout.flush()
+
+
+def main():
+    if len(sys.argv) == 3 and sys.argv[1] == "--serve":
+        return serve(sys.argv[2])
+    job = json.load(sys.stdin)
+    session, tok, config = load(job["model_dir"])
+    texts = job["texts"]
+    out = []
+    for i, text in enumerate(texts, 1):
+        out.append(embed_one(session, tok, config, text, job.get("max_tokens", 512)))
+        if i % 10 == 0 or i == len(texts):
+            print(f"embedded {i}/{len(texts)}", file=sys.stderr, flush=True)
+    np.save(job["out"], np.array(out, dtype=np.float32))
+
+
+if __name__ == "__main__":
+    main()

agent_memory/frontmatter.py

@@ -0,0 +1,146 @@
+"""Minimal YAML-subset frontmatter parser and serializer.
+
+Memory files only need scalars, quoted strings, flow lists, block lists and
+one level of nested mappings, so this avoids a PyYAML dependency.
+"""
+
+import json
+import re
+
+FENCE = "---"
+
+
+def split(text):
+    """Return (meta, body). A file without frontmatter yields ({}, text)."""
+    lines = text.split("\n")
+    if not lines or lines[0].strip() != FENCE:
+        return {}, text
+    for i in range(1, len(lines)):
+        if lines[i].strip() == FENCE:
+            meta = parse("\n".join(lines[1:i]))
+            body = "\n".join(lines[i + 1:]).lstrip("\n")
+            return meta, body
+    return {}, text
+
+
+def join(meta, body):
+    return f"{FENCE}\n{dump(meta)}{FENCE}\n{body.rstrip()}\n"
+
+
+def parse(block):
+    lines = [line.rstrip() for line in block.split("\n")
+             if line.strip() and not line.lstrip().startswith("#")]
+    if not lines:
+        return {}
+    value, _ = _parse_map(lines, 0, _indent(lines[0]))
+    return value
+
+
+def _indent(line):
+    return len(line) - len(line.lstrip(" "))
+
+
+def _parse_map(lines, i, indent):
+    out = {}
+    while i < len(lines):
+        line = lines[i]
+        ind = _indent(line)
+        if ind < indent:
+            break
+        stripped = line.strip()
+        if ind > indent or stripped.startswith("- "):
+            i += 1
+            continue
+        key, sep, rest = stripped.partition(":")
+        i += 1
+        if not sep:
+            continue
+        key, rest = key.strip(), rest.strip()
+        if rest:
+            out[key] = _scalar(rest)
+        elif i < len(lines) and _indent(lines[i]) > indent:
+            child = _indent(lines[i])
+            if lines[i].strip().startswith("- "):
+                out[key], i = _parse_list(lines, i, child)
+            else:
+                out[key], i = _parse_map(lines, i, child)
+        else:
+            out[key] = ""
+    return out, i
+
+
+def _parse_list(lines, i, indent):
+    out = []
+    while i < len(lines) and _indent(lines[i]) >= indent:
+        stripped = lines[i].strip()
+        if _indent(lines[i]) == indent and stripped.startswith("- "):
+            out.append(_scalar(stripped[2:].strip()))
+        i += 1
+    return out, i
+
+
+def _scalar(s):
+    if s.startswith('"'):
+        try:
+            return json.loads(s)
+        except ValueError:
+            return s.strip('"')
+    if s.startswith("'") and s.endswith("'") and len(s) >= 2:
+        return s[1:-1].replace("''", "'")
+    if s.startswith("[") and s.endswith("]"):
+        try:
+            value = json.loads(s)
+            if isinstance(value, list):
+                return value
+        except ValueError:
+            pass
+        inner = s[1:-1].strip()
+        return [_scalar(p.strip()) for p in inner.split(",") if p.strip()] if inner else []
+    low = s.lower()
+    if low == "true":
+        return True
+    if low == "false":
+        return False
+    if low in ("null", "~"):
+        return None
+    return s
+
+
+_RESERVED = {"true", "false", "null", "yes", "no", "on", "off", "~"}
+
+
+def _needs_quote(s):
+    if s == "" or s != s.strip() or "\n" in s:
+        return True
+    if s.lower() in _RESERVED:
+        return True
+    if s[0] in "-?:,[]{}#&*!|>'\"%@`":
+        return True
+    if ": " in s or " #" in s or s.endswith(":"):
+        return True
+    return bool(re.fullmatch(r"[-+]?\d+(\.\d+)?", s))
+
+
+def _emit(value):
+    if value is None:
+        return "null"
+    if isinstance(value, bool):
+        return "true" if value else "false"
+    if isinstance(value, (int, float)):
+        return str(value)
+    if isinstance(value, (list, tuple)):
+        return "[" + ", ".join(_emit(v) for v in value) + "]"
+    s = str(value)
+    return json.dumps(s, ensure_ascii=False) if _needs_quote(s) else s
+
+
+def dump(meta):
+    out = []
+    for key, value in meta.items():
+        if isinstance(value, dict):
+            out.append(f"{key}:")
+            for k, v in value.items():
+                out.append(f"  {k}: {_emit(v)}")
+        else:
+            out.append(f"{key}: {_emit(value)}")
+    return "\n".join(out) + "\n"

agent_memory/gitsync.py

@@ -0,0 +1,119 @@
+"""Git operations on the data repository, serialized by a file lock."""
+
+import fcntl
+import os
+import subprocess
+import time
+from pathlib import Path
+
+from .config import mem_command
+from .util import log
+
+_ENV = {**os.environ, "GIT_TERMINAL_PROMPT": "0"}
+
+
+class Lock:
+    """Exclusive lock shared by every hook and CLI process on this machine."""
+
+    def __init__(self, data_dir, timeout=10.0):
+        self.path = Path(data_dir) / ".cache" / "lock"
+        self.timeout = timeout
+        self.fd = None
+
+    def __enter__(self):
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        self.fd = os.open(self.path, os.O_RDWR | os.O_CREAT, 0o644)
+        deadline = time.monotonic() + self.timeout
+        while True:
+            try:
+                fcntl.flock(self.fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                return True
+            except BlockingIOError:
+                if time.monotonic() >= deadline:
+                    os.close(self.fd)
+                    self.fd = None
+                    return False
+                time.sleep(0.05)
+
+    def __exit__(self, *exc):
+        if self.fd is not None:
+            fcntl.flock(self.fd, fcntl.LOCK_UN)
+            os.close(self.fd)
+            self.fd = None
+
+
+def git(data_dir, *args, timeout=30):
+    return subprocess.run(["git", "-C", str(data_dir), *args], capture_output=True,
+                          text=True, timeout=timeout, env=_ENV)
+
+
+def has_remote(data_dir):
+    # A system gitconfig can declare remote.origin.* without a URL, so look for URLs.
+    r = git(data_dir, "config", "--get-regexp", r"^remote\..*\.url$", timeout=5)
+    return r.returncode == 0 and bool(r.stdout.strip())
+
+
+def commit(data_dir, message):
+    """Stage everything under memories/ and commit. Returns True if a commit was made."""
+    git(data_dir, "add", "-A", "--", "memories")
+    if git(data_dir, "diff", "--cached", "--quiet").returncode == 0:
+        return False
+    r = git(data_dir, "-c", "commit.gpgsign=false", "commit", "-q", "--no-verify", "-m", message)
+    if r.returncode != 0:
+        log(data_dir, f"commit failed: {r.stderr.strip()}")
+        return False
+    return True
+
+
+def _rebasing(data_dir):
+    git_dir = Path(data_dir) / ".git"
+    return (git_dir / "rebase-merge").exists() or (git_dir / "rebase-apply").exists()
+
+
+def sync(data_dir, push=True):
+    """Fetch, rebase onto upstream and optionally push.
+
+    Only the local rebase holds the lock; network round trips do not, so hooks
+    that commit are never stuck behind a slow fetch or push.
+    """
+    if not has_remote(data_dir):
+        return True
+    r = git(data_dir, "fetch", "-q", timeout=60)
+    if r.returncode != 0:
+        log(data_dir, f"fetch failed: {r.stderr.strip()}")
+        return False
+    with Lock(data_dir, timeout=30) as locked:
+        if not locked:
+            log(data_dir, "sync skipped, lock busy")
+            return False
+        behind = git(data_dir, "rev-list", "--count", "HEAD..@{u}")
+        if behind.returncode == 0 and behind.stdout.strip() != "0":
+            r = git(data_dir, "-c", "commit.gpgsign=false", "rebase", "-q", "--autostash", "@{u}")
+            if r.returncode != 0:
+                log(data_dir, f"rebase failed, aborted: {r.stderr.strip()}")
+                if _rebasing(data_dir):
+                    git(data_dir, "rebase", "--abort")
+                return False
+    if push:
+        ahead = git(data_dir, "rev-list", "--count", "@{u}..HEAD")
+        if ahead.returncode != 0 or ahead.stdout.strip() != "0":
+            r = git(data_dir, "push", "-q", timeout=60)
+            if r.returncode != 0:
+                log(data_dir, f"push failed: {r.stderr.strip()}")
+                return False
+    return True
+
+
+def spawn_sync(data_dir, push):
+    """Run `mem sync` detached so hooks return immediately."""
+    args = mem_command() + ["sync", "--quiet"]
+    if not push:
+        args.append("--pull-only")
+    log_path = Path(data_dir) / ".cache" / "sync.out"
+    try:
+        with open(log_path, "a") as out:
+            subprocess.Popen(args, stdin=subprocess.DEVNULL, stdout=out, stderr=out,
+                             start_new_session=True,
+                             env={**os.environ, "AGENT_MEMORY_DATA": str(data_dir)})
+    except OSError as e:
+        log(data_dir, f"spawn sync failed: {e}")

agent_memory/hooks.py

@@ -0,0 +1,95 @@
+"""Hook entry points shared by Claude Code and Codex.
+
+Both tools send a JSON payload on stdin with at least `session_id` and `cwd`,
+and accept `hookSpecificOutput.additionalContext` from SessionStart hooks.
+A hook must never break the host session, so every failure is logged and
+swallowed.
+"""
+
+import json
+import secrets
+import sys
+import time
+from pathlib import Path
+
+from . import capture, config, gitsync, select, store, workdir
+from .context import detect
+from .util import log, log_exc
+
+PULL_INTERVAL = 60
+
+
+def _payload():
+    try:
+        data = sys.stdin.read()
+        return json.loads(data) if data.strip() else {}
+    except (OSError, ValueError):
+        return {}
+
+
+def _session_id(payload):
+    return str(payload.get("session_id") or f"nosession-{secrets.token_hex(4)}")
+
+
+def session_start(payload, tool, data_dir):
+    sid = _session_id(payload)
+    ctx = detect(payload.get("cwd"))
+    sdir = capture.session_dir(data_dir, tool, sid)
+    sdir.mkdir(parents=True, exist_ok=True)
+    capture.write_session(sdir, {
+        "tool": tool,
+        "session_id": sid,
+        "session_title": payload.get("session_title"),
+        "model": payload.get("model"),
+        "transcript_path": payload.get("transcript_path"),
+        "ctx": ctx,
+    })
+
+    recovered = capture.recover_orphans(data_dir, keep=sdir)
+    stamp = Path(data_dir) / ".cache" / "last-pull"
+    if gitsync.has_remote(data_dir):
+        last = stamp.stat().st_mtime if stamp.exists() else 0
+        if recovered or time.time() - last > PULL_INTERVAL:
+            stamp.parent.mkdir(parents=True, exist_ok=True)
+            stamp.touch()
+            gitsync.spawn_sync(data_dir, push=recovered)
+
+    paths = {
+        "memories": str(store.memories_dir(data_dir)),
+        "staging": str(sdir),
+        "mem": str(config.BIN),
+    }
+    return select.render(store.load_all(data_dir), ctx, paths,
+                         throwaway=workdir.is_throwaway(ctx.get("cwd")))
+
+
+def flush(payload, tool, data_dir, final):
+    sdir = capture.session_dir(data_dir, tool, _session_id(payload))
+    if not sdir.is_dir():
+        return
+    if payload.get("transcript_path"):
+        capture.write_session(sdir, {"transcript_path": payload["transcript_path"]})
+    committed = capture.flush(data_dir, sdir, final=final, lock_timeout=1.0 if final else 5.0)
+    if committed and gitsync.has_remote(data_dir):
+        gitsync.spawn_sync(data_dir, push=True)
+
+
+def main(event, tool):
+    payload = _payload()
+    data_dir = config.data_dir()
+    if not data_dir or not (data_dir / ".git").exists():
+        return 0
+    try:
+        if event == "session-start":
+            text = session_start(payload, tool, data_dir)
+            print(json.dumps({"hookSpecificOutput": {
+                "hookEventName": "SessionStart",
+                "additionalContext": text,
+            }}, ensure_ascii=False))
+        elif event in ("stop", "session-end"):
+            flush(payload, tool, data_dir, final=event == "session-end")
+        else:
+            log(data_dir, f"unknown hook event {event!r}")
+    except Exception:
+        log_exc(data_dir, f"hook {event} ({tool})")
+    return 0