vamp-cli 0.1.0__py3-none-any.whl

vamp_cli/persistence/emitter.py

@@ -0,0 +1,371 @@
+"""
+Translate LangChain `astream_events(version='v2')` frames into events
+appended to an EventLog.
+
+Covers:
+    assistant_message   from on_chat_model_start/stream/end
+    tool_result         from on_tool_start/end/error
+
+Design rules:
+    - The emitter is a passive translator. It never mutates LangGraph state,
+      never raises past its own boundary, and never blocks the display path.
+    - All appends go through EventLog.append, which is already lock-serialized.
+    - Streaming chunks are accumulated locally using AIMessageChunk.__add__,
+      which correctly merges content, tool_call_chunks, usage, and
+      additional_kwargs (incl. reasoning_content for reasoning models).
+"""
+
+from __future__ import annotations
+
+import sys
+import time
+from typing import Any
+
+from langchain_core.messages import (
+    AIMessage,
+    AIMessageChunk,
+    BaseMessage,
+    ToolMessage,
+)
+
+from vamp_cli.persistence.event_log import EventLog
+
+
+# helpers
+def _content_to_str(content: Any) -> str:
+    """Normalize message content to plain text."""
+    if content is None:
+        return ""
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        parts: list[str] = []
+        for block in content:
+            if isinstance(block, str):
+                parts.append(block)
+            elif isinstance(block, dict):
+                # Anthropic / OpenAI content blocks
+                if block.get("type") in ("text", "output_text") and "text" in block:
+                    parts.append(str(block["text"]))
+                elif block.get("type") == "thinking" and "thinking" in block:
+                    parts.append(str(block["thinking"]))
+                # other block types (image, tool_use, etc) skipped
+        return "".join(parts)
+    return str(content)
+
+
+def _extract_reasoning(msg: Any) -> str:
+    """Pull reasoning text out of whatever field the provider used."""
+    if msg is None:
+        return ""
+    direct = getattr(msg, "reasoning_content", None)
+    if isinstance(direct, str) and direct:
+        return direct
+    extra = getattr(msg, "additional_kwargs", None) or {}
+    for key in ("reasoning_content", "reasoning", "thinking"):
+        val = extra.get(key)
+        if isinstance(val, str) and val:
+            return val
+    return ""
+
+
+def _extract_tool_calls(msg: Any) -> list[dict]:
+    """
+    Return normalized tool_calls as plain dicts.
+
+    LangChain normalizes them into {name, args, id, type} on AIMessage.
+    Chunks may instead expose `tool_call_chunks`, which merge into the
+    same shape as a best-effort.
+    """
+    if msg is None:
+        return []
+
+    calls = getattr(msg, "tool_calls", None)
+    if calls:
+        out: list[dict] = []
+        for tc in calls:
+            if isinstance(tc, dict):
+                out.append({
+                    "id":   tc.get("id"),
+                    "name": tc.get("name"),
+                    "args": tc.get("args") or {},
+                })
+            else:
+                out.append({
+                    "id":   getattr(tc, "id", None),
+                    "name": getattr(tc, "name", None),
+                    "args": getattr(tc, "args", {}) or {},
+                })
+        return out
+
+    chunks = getattr(msg, "tool_call_chunks", None) or []
+    merged: dict[int, dict] = {}
+    for i, c in enumerate(chunks):
+        idx = c.get("index", i) if isinstance(c, dict) else i
+        bucket = merged.setdefault(idx, {"id": None, "name": None, "args_str": ""})
+        if isinstance(c, dict):
+            bucket["id"] = bucket["id"] or c.get("id")
+            bucket["name"] = bucket["name"] or c.get("name")
+            if c.get("args"):
+                bucket["args_str"] += str(c["args"])
+    out2: list[dict] = []
+    for idx in sorted(merged):
+        b = merged[idx]
+        args: Any = {}
+        if b["args_str"]:
+            try:
+                import json
+                args = json.loads(b["args_str"])
+            except Exception:
+                args = {"_raw": b["args_str"]}
+        out2.append({"id": b["id"], "name": b["name"], "args": args})
+    return out2
+
+
+def _unwrap_chat_output(output: Any) -> BaseMessage | None:
+    """
+    on_chat_model_end.data.output is either an AIMessage/AIMessageChunk
+    or an LLMResult wrapping one. Unwrap both.
+    """
+    if output is None:
+        return None
+    if hasattr(output, "generations"):
+        gens = getattr(output, "generations") or []
+        if gens and gens[0]:
+            return getattr(gens[0][0], "message", None)
+        return None
+    if isinstance(output, (AIMessage, AIMessageChunk)):
+        return output
+    return output if isinstance(output, BaseMessage) else None
+
+
+def _truncate(text: str, limit: int) -> tuple[str, bool]:
+    if len(text) <= limit:
+        return text, False
+    # keep head and tail so both the context and the error message survive
+    head = limit * 2 // 3
+    tail = limit - head
+    return f"{text[:head]}\n...[{len(text) - limit} bytes elided]...\n{text[-tail:]}", True
+
+
+# EventEmitter
+class EventEmitter:
+    """
+    Consumes astream_events frames and appends persistence events.
+
+    Usage:
+        emitter = EventEmitter(log)
+        async for event in graph.astream_events(payload, config=cfg, version="v2"):
+            await emitter.consume(event)
+                display logic unchanged
+    """
+
+    def __init__(self, log: EventLog, *, tool_output_limit: int = 16_000) -> None:
+        self.log = log
+        self._tool_output_limit = tool_output_limit
+
+        # run_id -> in-flight chat state
+        self._chat: dict[str, dict[str, Any]] = {}
+
+        # run_id -> in-flight tool state
+        self._tools: dict[str, dict[str, Any]] = {}
+
+        # tool_call_id -> originating node ("coder" | "planner" | "tool")
+        # Populated when an assistant_message carrying tool_calls is emitted,
+        # consumed when the matching tool_result is emitted.
+        self._tool_call_to_node: dict[str, str | None] = {}
+        
+        # last input_tokens reported by any model call in this turn
+        self.last_input_tokens: int = 0
+
+    # public entry point 
+    async def consume(self, event: dict) -> None:
+        """Route one astream_events frame. Never raises."""
+        kind = event.get("event", "")
+        handler = _HANDLERS.get(kind)
+        if handler is None:
+            return
+        try:
+            await handler(self, event)
+        except Exception as e:
+            # Persistence must never break the run. Report and move on.
+            print(f"[persistence] failed to emit '{kind}': {e!r}", file=sys.stderr)
+
+    # chat model 
+    async def _on_chat_start(self, event: dict) -> None:
+        run_id = event["run_id"]
+        self._chat[run_id] = {
+            "chunk": None,          # accumulated AIMessageChunk
+            "meta": event.get("metadata") or {},
+            "started_at": time.time(),
+        }
+
+    async def _on_chat_stream(self, event: dict) -> None:
+        run_id = event.get("run_id")
+        buf = self._chat.get(run_id) # type: ignore
+        if buf is None:
+            return
+        chunk = (event.get("data") or {}).get("chunk")
+        if chunk is None:
+            return
+        try:
+            buf["chunk"] = chunk if buf["chunk"] is None else buf["chunk"] + chunk
+        except Exception:
+            # merge failure is rare and non-fatal - keep what we have
+            pass
+
+    async def _on_chat_end(self, event: dict) -> None:
+        run_id = event.get("run_id")
+        buf = self._chat.pop(run_id, None) #type: ignore
+        if buf is None:
+            return
+
+        accumulated: AIMessageChunk | None = buf["chunk"]
+        final_msg = _unwrap_chat_output((event.get("data") or {}).get("output"))
+
+        # Prefer the accumulated stream (it has reasoning_content that the
+        # final message often drops). Supplement with final_msg when needed.
+        content = _content_to_str(getattr(accumulated, "content", "")) if accumulated else ""
+        reasoning = _extract_reasoning(accumulated) if accumulated else ""
+        tool_calls = _extract_tool_calls(accumulated) if accumulated else []
+
+        if final_msg is not None:
+            if not content:
+                content = _content_to_str(getattr(final_msg, "content", ""))
+            if not reasoning:
+                reasoning = _extract_reasoning(final_msg)
+            if not tool_calls:
+                tool_calls = _extract_tool_calls(final_msg)
+
+        usage = None
+        if final_msg is not None:
+            um = getattr(final_msg, "usage_metadata", None)
+            if um:
+                usage = dict(um)
+        
+        if usage and isinstance(usage.get("input_tokens"), int):
+            self.last_input_tokens = usage["input_tokens"]
+
+        meta = buf.get("meta") or {}
+        node = meta.get("langgraph_node")
+
+        # Remember which node each tool call came from so tool_result
+        # events can carry the same attribution.
+        for tc in tool_calls:
+            tid = tc.get("id")
+            if tid:
+                self._tool_call_to_node[tid] = node
+
+        self.log.emit("assistant_message", {
+            "node": node,
+            "content": content,
+            "reasoning_content": reasoning or None,
+            "tool_calls": tool_calls,
+            "usage": usage,
+            "model": meta.get("ls_model_name"),
+            "duration_ms": int((time.time() - buf["started_at"]) * 1000),
+        })
+
+    # tools ---------------------------------------------------------
+    async def _on_tool_start(self, event: dict) -> None:
+        run_id = event.get("run_id")
+        data = event.get("data") or {}
+        meta = event.get("metadata") or {}
+
+        # LangChain names the tool in `event["name"]`, the args in data["input"].
+        self._tools[run_id] = { # type: ignore
+            "tool_name": event.get("name"),
+            "tool_call_id": data.get("tool_call_id"),
+            "input": data.get("input"),
+            "node_hint": meta.get("langgraph_node"),
+            "started_at": time.time(),
+        }
+
+    async def _on_tool_end(self, event: dict) -> None:
+        run_id = event.get("run_id")
+        meta = self._tools.pop(run_id, {}) # type: ignore
+        output = (event.get("data") or {}).get("output")
+
+        tool_call_id = meta.get("tool_call_id")
+        tool_name = meta.get("tool_name")
+        status = "success"
+        content = ""
+
+        if isinstance(output, ToolMessage):
+            tool_call_id = getattr(output, "tool_call_id", None) or tool_call_id
+            tool_name = getattr(output, "name", None) or tool_name
+            content = _content_to_str(getattr(output, "content", ""))
+            status = getattr(output, "status", None) or "success"
+        else:
+            content = _content_to_str(output)
+
+        text, truncated = _truncate(content, self._tool_output_limit)
+
+        # Prefer the model-node when the assistant emitted the call. 
+        # Fall back to langgraph_node from the tool event,
+        # which will usually be "execute_tools" and thus less useful.
+        node = self._tool_call_to_node.get(tool_call_id) or meta.get("node_hint")
+
+        self.log.emit("tool_result", {
+            "tool_call_id": tool_call_id,
+            "tool_name": tool_name,
+            "node": node,
+            "status": status,
+            "content": text,
+            "truncated": truncated,
+            "content_len": len(content),
+            "duration_ms": int((time.time() - meta.get("started_at", time.time())) * 1000),
+        })
+
+    async def _on_tool_error(self, event: dict) -> None:
+        run_id = event.get("run_id")
+        meta = self._tools.pop(run_id, {}) # type: ignore
+        err = (event.get("data") or {}).get("error")
+
+        tool_call_id = meta.get("tool_call_id")
+        node = self._tool_call_to_node.get(tool_call_id) or meta.get("node_hint")
+
+        self.log.emit("tool_result", {
+            "tool_call_id": tool_call_id,
+            "tool_name":    meta.get("tool_name"),
+            "node":         node,
+            "status":       "error",
+            "content":      f"{type(err).__name__}: {err}" if err else "unknown error",
+            "truncated":    False,
+            "content_len":  0,
+            "duration_ms":  int((time.time() - meta.get("started_at", time.time())) * 1000),
+        })
+        
+    def emit_tool_result(
+        self,
+        *,
+        tool_call_id: str,
+        tool_name:    str,
+        node:         str | None,
+        status:       str,
+        content:      str,
+        duration_ms:  int,
+    ) -> None:
+        """Public hook for callers that bypass LangChain's callback stream."""
+        text, truncated = _truncate(content or "", self._tool_output_limit)
+        self.log.emit("tool_result", {
+            "tool_call_id": tool_call_id,
+            "tool_name":    tool_name,
+            "node":         node,
+            "status":       status,
+            "content":      text,
+            "truncated":    truncated,
+            "content_len":  len(content or ""),
+            "duration_ms":  duration_ms,
+        })
+
+
+# dispatch table - populated after the class to keep the definition clean
+_HANDLERS = {
+    "on_chat_model_start":  EventEmitter._on_chat_start,
+    "on_chat_model_stream": EventEmitter._on_chat_stream,
+    "on_chat_model_end":    EventEmitter._on_chat_end,
+    "on_tool_start":        EventEmitter._on_tool_start,
+    "on_tool_end":          EventEmitter._on_tool_end,
+    "on_tool_error":        EventEmitter._on_tool_error,
+}

vamp_cli/persistence/event_log.py

@@ -0,0 +1,303 @@
+"""
+Append-only JSONL event log for a single agent session.
+
+File layout (per workspace):
+    <user-home-dir>/.vamp/sessions/<sessionId>.jsonl   # append-only event log
+    <user-home-dir>/.vamp/sessions/<sessionId>.head    # sidecar: current leaf uuid
+
+Concurrency model:
+    - One background writer task drains an asyncio.Queue in FIFO order.
+    - `emit()` is fire-and-forget: enqueue and return. Display path is free.
+    - `append()` enqueues and awaits the write. Use for durability-critical
+      events (user_message, interrupt_request, interrupt_response, ...).
+    - `parentUuid` chaining happens *inside* the writer task, so it is
+      race-free even under bursts of concurrent emit() calls.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import subprocess
+import time
+import uuid
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Callable, Iterable, Iterator
+
+
+
+# uuid7 - time-sortable, canonical-string form
+def uuid7() -> str:
+    ms = int(time.time() * 1000) & ((1 << 48) - 1)
+    rand_a = int.from_bytes(os.urandom(2), "big") & 0x0FFF
+    rand_b = int.from_bytes(os.urandom(8), "big") & ((1 << 62) - 1)
+    value = (ms << 80) | (0x7 << 76) | (rand_a << 64) | (0b10 << 62) | rand_b
+    return str(uuid.UUID(int=value))
+
+
+# git branch
+def _read_git_branch(cwd: str) -> str | None:
+    try:
+        result = subprocess.run(
+            ["git", "rev-parse", "--abbrev-ref", "HEAD"],
+            cwd=cwd, capture_output=True, text=True, timeout=1,
+        )
+        if result.returncode == 0:
+            branch = result.stdout.strip()
+            return branch or None
+    except Exception:
+        pass
+    return None
+
+
+# redaction hook
+Redactor = Callable[[str, dict[str, Any]], dict[str, Any]]
+
+
+def identity_redactor(event_type: str, data: dict[str, Any]) -> dict[str, Any]:
+    return data
+
+
+# fsync policy
+DEFAULT_FSYNC_TYPES = frozenset({
+    "session_start",
+    "session_end",
+    "user_message",
+    "mode_switch",
+    "interrupt_request",
+    "interrupt_response",
+})
+
+
+# Queue item: (event_type, data, done_future_or_None)
+# _QueueItem type
+_QueueItem = tuple[str, dict[str, Any], "asyncio.Future[str] | None", "str | None"]
+
+
+# EventLog
+class EventLog:
+    """Async append-only writer for one session, backed by a writer task."""
+
+    def __init__(
+        self,
+        *,
+        session_id: str,
+        path: Path,
+        head_path: Path,
+        cwd: str,
+        version: str,
+        git_branch: str | None,
+        redactor: Redactor,
+        fsync_types: frozenset[str],
+    ) -> None:
+        self.session_id = session_id
+        self.path = path
+        self.head_path = head_path
+        self.cwd = cwd
+        self.version = version
+        self.git_branch = git_branch
+        self._redactor = redactor
+        self._fsync_types = fsync_types
+
+        self._file = None
+        self._last_uuid: str | None = None
+        self._queue: asyncio.Queue[_QueueItem | None] = asyncio.Queue()
+        self._writer_task: asyncio.Task[None] | None = None
+        self._closed = False
+
+    # construction 
+    @classmethod
+    async def open(
+        cls,
+        session_id: str,
+        *,
+        workspace_root: str | os.PathLike[str],
+        version: str,
+        redactor: Redactor | None = None,
+        fsync_on: Iterable[str] | None = None,
+        cwd: str | None = None,
+    ) -> "EventLog":
+        sessions_dir = Path.home() / ".vamp" / "sessions"
+        sessions_dir.mkdir(parents=True, exist_ok=True)
+
+        path = sessions_dir / f"{session_id}.jsonl"
+        head_path = sessions_dir / f"{session_id}.head"
+        cwd = cwd or str(Path.cwd())
+
+        log = cls(
+            session_id=session_id,
+            path=path,
+            head_path=head_path,
+            cwd=cwd,
+            version=version,
+            git_branch=_read_git_branch(cwd),
+            redactor=redactor or identity_redactor,
+            fsync_types=(
+                frozenset(fsync_on) if fsync_on is not None
+                else DEFAULT_FSYNC_TYPES
+            ),
+        )
+
+        if head_path.exists():
+            log._last_uuid = head_path.read_text(encoding="utf-8").strip() or None
+        elif path.exists():
+            log._last_uuid = _scan_leaf_uuid(path)
+
+        log._file = open(path, "a", encoding="utf-8", buffering=1)
+        log._writer_task = asyncio.create_task(log._writer_loop())
+        return log
+
+    
+    @property
+    def head_uuid(self) -> str | None:
+        return self._last_uuid
+
+    def refresh_branch(self) -> None:
+        self.git_branch = _read_git_branch(self.cwd)
+
+    # writer task 
+    async def _writer_loop(self) -> None:
+        while True:
+            item = await self._queue.get()
+            if item is None:
+                self._queue.task_done()
+                return
+            event_type, data, done, parent_uuid = item
+            try:
+                event_uuid = self._write_event(event_type, data, parent_uuid)
+                if done is not None and not done.done():
+                    done.set_result(event_uuid)
+            except Exception as e:
+                if done is not None and not done.done():
+                    done.set_exception(e)
+            finally:
+                self._queue.task_done()
+
+    def _write_event(
+        self, 
+        event_type: str, 
+        data: dict[str, Any],
+        parent_uuid: str | None = None,
+    ) -> str:
+        """Synchronous - runs only inside the writer task."""
+        payload = self._redactor(event_type, dict(data))
+
+        event_uuid = uuid7()
+        envelope = {
+            "type": event_type,
+            "uuid": event_uuid,
+            "parentUuid":  parent_uuid if parent_uuid is not None else self._last_uuid,
+            "sessionId": self.session_id,
+            "timestamp": _iso_now(),
+            "cwd": self.cwd,
+            "version": self.version,
+            "gitBranch": self.git_branch,
+            "data": payload,
+        }
+
+        line = json.dumps(envelope, ensure_ascii=False, separators=(",", ":"))
+        self._file.write(line + "\n") # type: ignore 
+        self._file.flush() # type: ignore
+        if event_type in self._fsync_types:
+            os.fsync(self._file.fileno()) # type: ignore
+
+        self._last_uuid = event_uuid
+        self.head_path.write_text(event_uuid, encoding="utf-8")
+        return event_uuid
+
+    # public API 
+    def emit(self, event_type: str, data: dict[str, Any] | None = None, *, parent_uuid=None) -> None:
+        """Fire-and-forget. Non-blocking. FIFO ordering guaranteed."""
+        if self._closed:
+            raise RuntimeError("EventLog is closed")
+        self._queue.put_nowait((event_type, dict(data or {}), None, parent_uuid))
+
+    async def append(
+        self,
+        event_type: str,
+        data: dict[str, Any] | None = None,
+        *,
+        parent_uuid=None,
+    ) -> str:
+        """Enqueue and wait for the durable write to complete."""
+        if self._closed:
+            raise RuntimeError("EventLog is closed")
+        done: asyncio.Future[str] = asyncio.get_running_loop().create_future()
+        await self._queue.put((event_type, dict(data or {}), done, parent_uuid))
+        return await done
+
+    async def flush(self) -> None:
+        """Wait until every queued write has been processed."""
+        await self._queue.join()
+
+    async def close(self, reason: str = "user_exit") -> None:
+        if self._closed:
+            return
+        try:
+            await self.append("session_end", {"reason": reason})
+            await self.flush()
+        finally:
+            self._closed = True
+            await self._queue.put(None)          # sentinel
+            if self._writer_task is not None:
+                try:
+                    await self._writer_task
+                except Exception:
+                    pass
+            if self._file is not None:
+                self._file.flush()
+                try:
+                    os.fsync(self._file.fileno())
+                except Exception:
+                    pass
+                self._file.close()
+                self._file = None
+
+    # branch control (fork; no writes, no truncation)
+    async def set_head(self, event_uuid: str) -> None:
+        """Repoint the head. Flushes first so pending writes are durable."""
+        await self.flush()
+        self._last_uuid = event_uuid
+        self.head_path.write_text(event_uuid, encoding="utf-8")
+
+
+
+# helpers
+def _iso_now() -> str:
+    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
+
+
+def _scan_leaf_uuid(path: Path) -> str | None:
+    all_uuids: list[str] = []
+    parents: set[str] = set()
+    with open(path, "r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                event = json.loads(line)
+            except json.JSONDecodeError:
+                continue
+            u = event.get("uuid")
+            p = event.get("parentUuid")
+            if u:
+                all_uuids.append(u)
+            if p:
+                parents.add(p)
+    leaves = [u for u in all_uuids if u not in parents]
+    return max(leaves) if leaves else None
+
+
+def iter_events(path: Path) -> Iterator[dict[str, Any]]:
+    with open(path, "r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                yield json.loads(line)
+            except json.JSONDecodeError:
+                continue