mem0/integrations/kimi-plugin/core/hook_runner.py

"""Shared hook orchestration for all Mem0 agent plugins."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import subprocess
import sys
import time
import uuid
from pathlib import Path

import telemetry
from memory_core import (
    EvidenceStore,
    _session_id,
    api_key,
    bounded,
    cache_plugin_api_key,
    checkpoint_session,
    clear_stale_api_key_cache,
    configure_harness,
    data_dir,
    detached_process_kwargs,
    format_context,
    harness_config,
    record_session_start,
    record_tool,
    record_user_prompt,
    redact,
    search_memories,
)

STALE_RUNNING_SECONDS = 300
PENDING_EXPIRY_SECONDS = 7 * 24 * 60 * 60
PENDING_LAUNCH_LIMIT = 5
DEFAULT_IDLE_FLUSH_SECONDS = 300

_core_dir: Path = Path(__file__).resolve().parent


def read_hook_input() -> dict:
    try:
        value = json.load(sys.stdin)
        return value if isinstance(value, dict) else {}
    except (json.JSONDecodeError, OSError):
        return {}


def default_record_stop(store: EvidenceStore, hook_input: dict):
    """Record the assistant's response without transcript parsing."""
    session_id = _session_id(hook_input)
    repo = store.repo_for_session(session_id, hook_input.get("cwd"))
    message = redact(hook_input.get("last_assistant_message", "")).strip()
    if message:
        store.record_assistant_response(repo, session_id, message)
    return repo, session_id


def first_prompt_memory_output(store: EvidenceStore, hook_input: dict) -> dict:
    """Search once before the agent handles the first prompt in a session."""
    repo, session_id, prompt, is_first_prompt = record_user_prompt(store, hook_input)
    if not is_first_prompt:
        return {}
    try:
        minimum_query_chars = int(os.environ.get("MEM0_CODE_MIN_QUERY_CHARS", "20"))
    except ValueError:
        minimum_query_chars = 20
    if len(prompt.strip()) < max(minimum_query_chars, 1):
        return {}
    result = search_memories(
        store, repo, session_id, bounded(prompt, 6000),
        top_k=5, operation="first-prompt-search", timeout=2,
    )
    if not result.memories:
        return {}
    context = format_context(
        result.memories,
        "Mem0 found these relevant memories from earlier work in this repository:",
    )
    telemetry.record(
        "context_injected",
        repo=repo, session_id=session_id, trigger="first-prompt",
        memory_count=len(result.memories), context_chars=len(context),
        prompt_chars=len(prompt),
    )
    return {
        "hookSpecificOutput": {
            "hookEventName": "UserPromptSubmit",
            "additionalContext": context,
        },
    }


def _launch_handoff(handoff_path: Path) -> bool:
    running_path = handoff_path.with_suffix(".running")
    try:
        handoff_path.replace(running_path)
    except OSError:
        return False
    worker = _core_dir / "flush_worker.py"
    log_path = data_dir() / "flush-worker.log"
    log_handle = open(log_path, "a", encoding="utf-8")
    harness = harness_config()
    child_env = os.environ.copy()
    child_env.update(
        {
            "MEM0_CODE_DATA_DIR": str(data_dir()),
            "MEM0_PLUGIN_HARNESS": harness["name"],
            "MEM0_PLUGIN_ENV_PREFIX": harness["env_prefix"],
            "MEM0_PLUGIN_DATA_DIR_NAME": harness["data_dir_name"],
            "MEM0_PLUGIN_SOURCE_TAG": harness["source_tag"],
        }
    )
    try:
        subprocess.Popen(
            [sys.executable, str(worker), str(running_path)],
            stdin=subprocess.DEVNULL,
            stdout=log_handle, stderr=log_handle,
            close_fds=True,
            env=child_env,
            **detached_process_kwargs(),
        )
    finally:
        log_handle.close()
    return True


def recover_pending_handoffs() -> int:
    pending_dir = data_dir() / "pending"
    pending_dir.mkdir(parents=True, exist_ok=True)
    now = time.time()
    for running in pending_dir.glob("*.running"):
        try:
            if now - running.stat().st_mtime > STALE_RUNNING_SECONDS:
                running.replace(running.with_suffix(".json"))
        except OSError:
            continue
    recoverable = []
    for handoff in pending_dir.glob("*.json"):
        try:
            age = now - handoff.stat().st_mtime
        except OSError:
            continue
        if age > PENDING_EXPIRY_SECONDS:
            handoff.unlink(missing_ok=True)
            continue
        recoverable.append((age, handoff))
    recoverable.sort(key=lambda item: item[0], reverse=True)
    launched = 0
    for _, handoff in recoverable[:PENDING_LAUNCH_LIMIT]:
        launched += int(_launch_handoff(handoff))
    return launched


def refresh_pending_handoffs() -> None:
    pending_dir = data_dir() / "pending"
    if not pending_dir.is_dir():
        return
    for pattern in ("*.json", "*.running"):
        for handoff in pending_dir.glob(pattern):
            try:
                os.utime(handoff)
            except OSError:
                continue


def hand_off_flush(
    hook_input: dict, reason: str, *, wait_for_inflight: bool = False,
) -> None:
    pending_dir = data_dir() / "pending"
    pending_dir.mkdir(parents=True, exist_ok=True)
    material = (
        f"{hook_input.get('cwd', '')}\0{hook_input.get('session_id', '')}\0{reason}"
    )
    digest = hashlib.sha256(material.encode()).hexdigest()[:24]
    handoff_path = pending_dir / f"{digest}-{uuid.uuid4().hex[:8]}.json"
    temporary_path = handoff_path.with_suffix(".tmp")
    temporary_path.write_text(
        json.dumps({
            "hook_input": hook_input,
            "reason": reason,
            "wait_for_inflight": wait_for_inflight,
        }),
        encoding="utf-8",
    )
    temporary_path.replace(handoff_path)
    _launch_handoff(handoff_path)


def automatic_flush_enabled() -> bool:
    return os.environ.get("MEM0_CODE_AUTO_FLUSH", "true").lower() in {
        "1", "true", "yes", "on",
    }


def schedule_periodic_checkpoint(
    store: EvidenceStore, hook_input: dict, repo, session_id: str,
) -> bool:
    if (
        not automatic_flush_enabled()
        or not api_key()
        or not store.checkpoint_due(repo.identity, session_id)
    ):
        return False
    if store.prepare_flush(repo, session_id, "periodic") is None:
        return False
    hand_off_flush(hook_input, "periodic")
    return True


def _idle_flush_seconds() -> int:
    try:
        return max(
            int(os.environ.get("MEM0_CODE_IDLE_FLUSH_SECONDS", str(DEFAULT_IDLE_FLUSH_SECONDS))),
            0,
        )
    except ValueError:
        return DEFAULT_IDLE_FLUSH_SECONDS


def schedule_idle_flush(
    store: EvidenceStore, hook_input: dict, repo, session_id: str,
) -> bool:
    delay = _idle_flush_seconds()
    if delay <= 0 or not automatic_flush_enabled() or not api_key():
        return False
    if store.has_inflight_flush(repo.identity, session_id):
        return False
    if not store.has_unflushed_events(repo.identity, session_id):
        return False
    pending_dir = data_dir() / "pending"
    pending_dir.mkdir(parents=True, exist_ok=True)
    material = f"idle\0{hook_input.get('cwd', '')}\0{hook_input.get('session_id', '')}"
    digest = hashlib.sha256(material.encode()).hexdigest()[:24]
    for old in pending_dir.glob(f"idle-{digest}*"):
        old.unlink(missing_ok=True)
    handoff_path = pending_dir / f"idle-{digest}-{uuid.uuid4().hex[:8]}.json"
    temporary_path = handoff_path.with_suffix(".tmp")
    temporary_path.write_text(
        json.dumps({
            "hook_input": hook_input,
            "reason": "idle",
            "delay_seconds": delay,
        }),
        encoding="utf-8",
    )
    temporary_path.replace(handoff_path)
    _launch_handoff(handoff_path)
    return True


def log_failure(exc: Exception) -> None:
    try:
        log_path = data_dir() / "plugin-errors.log"
        with log_path.open("a", encoding="utf-8") as handle:
            handle.write(f"{time.time():.3f} {type(exc).__name__}: {exc}\n")
    except OSError:
        pass


def run(
    *,
    record_stop_fn=None,
    extra_actions: dict | None = None,
    data_dir_env: str = "MEM0_PLUGIN_DATA_DIR",
    automatic_flush_reasons: set | None = None,
) -> int:
    if record_stop_fn is None:
        record_stop_fn = default_record_stop
    if automatic_flush_reasons is None:
        automatic_flush_reasons = {"session-end"}

    base_actions = ["session-start", "user-prompt", "post-tool", "stop", "flush"]
    all_actions = base_actions + list((extra_actions or {}).keys())

    parser = argparse.ArgumentParser()
    parser.add_argument("action", choices=all_actions)
    parser.add_argument("--reason", default="manual")
    parser.add_argument("--plugin-data-dir", default="")
    parser.add_argument("--harness", default="")
    args = parser.parse_args()

    if args.harness:
        configure_harness(args.harness)
        telemetry.init(harness=args.harness)

    if args.plugin_data_dir:
        os.environ[data_dir_env] = args.plugin_data_dir

    cache_plugin_api_key()
    if args.action == "session-start":
        clear_stale_api_key_cache()

    hook_input = read_hook_input()
    store = EvidenceStore()
    try:
        if store.is_paused():
            if args.action == "session-start":
                refresh_pending_handoffs()
                telemetry.record("session_start", paused=True)
                telemetry.spawn_flush()
            return 0

        if args.action == "session-start":
            # Claims the marker atomically and says which event to record, so a
            # second session starting alongside this one cannot record it too.
            first_event = telemetry.claim_install()
            if first_event == "install":
                telemetry.record("install")
            elif first_event == "upgrade":
                # First run after a build that never wrote the marker; the
                # predecessor version was never recorded anywhere.
                telemetry.record("upgrade", from_version="pre-0.3")
            else:
                previous = telemetry.claim_version_change()
                if previous:
                    telemetry.record("upgrade", from_version=previous)
            recovered = recover_pending_handoffs()
            record_session_start(store, hook_input)
            if recovered:
                telemetry.record("handoff_recovered", count=recovered)
            telemetry.spawn_flush()
        elif args.action == "user-prompt":
            output = first_prompt_memory_output(store, hook_input)
            if output:
                print(json.dumps(output))
        elif args.action == "post-tool":
            record_tool(store, hook_input)
        elif args.action == "stop":
            repo, session_id = record_stop_fn(store, hook_input)
            if not schedule_periodic_checkpoint(store, hook_input, repo, session_id):
                schedule_idle_flush(store, hook_input, repo, session_id)
        elif args.action == "flush":
            automatic = args.reason in automatic_flush_reasons
            if automatic and not automatic_flush_enabled():
                return 0
            if args.reason == "session-end":
                record_stop_fn(store, hook_input)
            if os.environ.get("MEM0_CODE_SYNC_FLUSH") == "1":
                print(json.dumps(checkpoint_session(store, hook_input, args.reason)))
            else:
                session_id = str(hook_input.get("session_id") or "unknown-session")
                repo = store.repo_for_session(session_id, hook_input.get("cwd"))
                already_running = store.has_inflight_flush(repo.identity, session_id)
                if already_running and args.reason == "session-end":
                    hand_off_flush(hook_input, args.reason, wait_for_inflight=True)
                elif not already_running and store.prepare_flush(
                    repo, session_id, args.reason,
                ) is not None:
                    hand_off_flush(hook_input, args.reason)
        elif extra_actions and args.action in extra_actions:
            result = extra_actions[args.action](store, hook_input)
            if result:
                print(json.dumps(result))
    finally:
        store.close()
    return 0


def entry_point(
    *,
    record_stop_fn=None,
    extra_actions: dict | None = None,
    data_dir_env: str = "MEM0_PLUGIN_DATA_DIR",
    automatic_flush_reasons: set | None = None,
) -> None:
    try:
        raise SystemExit(run(
            record_stop_fn=record_stop_fn,
            extra_actions=extra_actions,
            data_dir_env=data_dir_env,
            automatic_flush_reasons=automatic_flush_reasons,
        ))
    except Exception as exc:
        log_failure(exc)
        raise SystemExit(0)


if __name__ == "__main__":
    entry_point()