From 5cd66fbb2c7d03bedc641f7b710d42589fba0808 Mon Sep 17 00:00:00 2001 From: eric_song Date: Fri, 3 Jul 2026 17:03:47 +0800 Subject: [PATCH 1/2] Add LingTai agent event index baseline Signed-off-by: eric_song --- Cargo.lock | 8 +- Cargo.toml | 2 +- bench/Cargo.toml | 4 + bench/agent-interface/README.md | 5 + bench/agent-interface/lingtai/README.md | 180 ++ .../scripts/lingtai_workload_baseline.py | 588 ++++ bench/src/bin/lingtai-index-bench.rs | 1392 +++++++++ crates/nokv-agent/Cargo.toml | 8 +- crates/nokv-agent/src/bin/nokv-agent.rs | 465 +++ crates/nokv-agent/src/event/codec.rs | 63 + crates/nokv-agent/src/event/holt.rs | 859 ++++++ crates/nokv-agent/src/event/ingest.rs | 247 ++ crates/nokv-agent/src/event/key.rs | 276 ++ crates/nokv-agent/src/event/mod.rs | 25 + crates/nokv-agent/src/event/store.rs | 74 + crates/nokv-agent/src/event/types.rs | 263 ++ crates/nokv-agent/src/lib.rs | 2671 +---------------- crates/nokv-agent/src/namespace.rs | 2664 ++++++++++++++++ crates/nokv-agent/tests/event_index.rs | 704 +++++ docs/development/code_contract.md | 1 + docs/development/nokv-agent.md | 150 +- 21 files changed, 7968 insertions(+), 2681 deletions(-) create mode 100644 bench/agent-interface/lingtai/README.md create mode 100755 bench/agent-interface/scripts/lingtai_workload_baseline.py create mode 100644 bench/src/bin/lingtai-index-bench.rs create mode 100644 crates/nokv-agent/src/bin/nokv-agent.rs create mode 100644 crates/nokv-agent/src/event/codec.rs create mode 100644 crates/nokv-agent/src/event/holt.rs create mode 100644 crates/nokv-agent/src/event/ingest.rs create mode 100644 crates/nokv-agent/src/event/key.rs create mode 100644 crates/nokv-agent/src/event/mod.rs create mode 100644 crates/nokv-agent/src/event/store.rs create mode 100644 crates/nokv-agent/src/event/types.rs create mode 100644 crates/nokv-agent/src/namespace.rs create mode 100644 crates/nokv-agent/tests/event_index.rs diff --git a/Cargo.lock b/Cargo.lock index 598d584fb..4b8199e9d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -811,9 +811,9 @@ dependencies = [ [[package]] name = "holt" -version = "0.7.3" +version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4f7180148389fd08f76cdd05ad47640acd1b72b6a5cd2379d1496de96af32e01" +checksum = "4baf6b46ff0e6835daf1c0799fc988608e52af421bda15bf1c08125fe09b5b43" dependencies = [ "crc32fast", "crossbeam-channel", @@ -1395,10 +1395,14 @@ dependencies = [ name = "nokv-agent" version = "0.1.0" dependencies = [ + "holt", "nokv-meta", "nokv-object", "nokv-types", + "serde", "serde_json", + "sha2 0.10.9", + "tempfile", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 67f325b3b..a8dd57158 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -38,7 +38,7 @@ nokv = { path = "crates/nokv", version = "0.1.0" } nokv-python = { path = "crates/nokv-python", version = "0.1.0" } nokv-bench = { path = "bench", version = "0.1.0" } fuser = { version = "0.17.0", default-features = false } -holt = { version = "0.7.3", default-features = false } +holt = { version = "0.8.1", default-features = false } etcd-client = { version = "0.19.0", default-features = false } libc = "0.2" opendal = { version = "0.57.0", default-features = false, features = ["blocking", "executors-tokio", "reqwest-rustls-tls", "services-s3"] } diff --git a/bench/Cargo.toml b/bench/Cargo.toml index 9a3b3701d..078a795a1 100644 --- a/bench/Cargo.toml +++ b/bench/Cargo.toml @@ -36,3 +36,7 @@ path = "src/bin/nokv-bench.rs" [[bin]] name = "yanex-agent-bench" path = "src/bin/yanex-agent-bench.rs" + +[[bin]] +name = "lingtai-index-bench" +path = "src/bin/lingtai-index-bench.rs" diff --git a/bench/agent-interface/README.md b/bench/agent-interface/README.md index f3f9efe3b..6d465cc98 100644 --- a/bench/agent-interface/README.md +++ b/bench/agent-interface/README.md @@ -5,6 +5,11 @@ harness. It compares how different agent-facing surfaces affect correctness, token cost, tool calls, and bytes read over the same fixed Yanex experiment corpus. +The active LingTai event-index benchmark track lives in +`bench/agent-interface/lingtai/`. It measures LingTai `events.jsonl` and +`log.sqlite` exports and is the benchmark path for the Holt-backed +`nokv-agent` derived event index. + This is not a storage-engine throughput benchmark. The product benchmark for metadata, training-read, and checkpoint workloads lives in the root `bench/` crate and is documented in `docs/benchmarks.md`. diff --git a/bench/agent-interface/lingtai/README.md b/bench/agent-interface/lingtai/README.md new file mode 100644 index 000000000..4fec5360d --- /dev/null +++ b/bench/agent-interface/lingtai/README.md @@ -0,0 +1,180 @@ +# LingTai Agent-Index Benchmark Track + +This track measures LingTai's real agent-log query workload. It is separate +from the historical Yanex benchmark corpus and uses LingTai exports as the +active integration target. + +LingTai keeps `logs/events.jsonl` as the authoritative log. `logs/log.sqlite` +is a derived index used by the TUI and session rebuild paths. NoKV's target is +the derived-index role: build a compact, semantic, agent-native index in +`nokv-agent` and compare it against both current and optimized SQLite baselines. + +## Baseline Arms + +| Arm | Meaning | +| --- | --- | +| `sqlite_current_v1` | The exported LingTai `logs/log.sqlite` exactly as produced by LingTai. | +| `sqlite_projected_v1` | SQLite with explicit projected/expression/materialized indexes for the measured query shapes. This is the fairness baseline. | +| `sqlite_rebuilt_from_jsonl_v1` | A SQLite database rebuilt from the same `logs/events.jsonl` input used by `nokv_agent_index_v1`, with the same projected indexes. This is the equality baseline when the exported SQLite sidecar has different coverage from JSONL. | +| `nokv_agent_index_v1` | Holt-backed `nokv-agent` event index built from `logs/events.jsonl`. | + +Do not claim a NoKV win unless the report includes the projected SQLite arm. +Current SQLite can be slow on JSON-field grouping or coverage scans, but many +point queries become fast once SQLite is given the right expression index. + +## Workload Extraction + +Use the read-only extractor to summarize a LingTai export without printing raw +message bodies, tool argument values, or user text: + +```bash +python3 bench/agent-interface/scripts/lingtai_workload_baseline.py \ + --root /path/to/extracted/lingtai-export \ + --projected-sqlite-out /tmp/lingtai-projected.sqlite \ + --json-out /tmp/lingtai-workload.json \ + --markdown-out /tmp/lingtai-workload.md +``` + +For split archives, pass the parts in order: + +```bash +python3 bench/agent-interface/scripts/lingtai_workload_baseline.py \ + --parts /path/to/yuanjiang_lingtai_filtered_logs_20260630.part000 \ + /path/to/yuanjiang_lingtai_filtered_logs_20260630.part001 \ + /path/to/yuanjiang_lingtai_filtered_logs_20260630.part002 \ + /path/to/yuanjiang_lingtai_filtered_logs_20260630.part003 \ + /path/to/yuanjiang_lingtai_filtered_logs_20260630.part004 \ + /path/to/yuanjiang_lingtai_filtered_logs_20260630.part005 \ + /path/to/yuanjiang_lingtai_filtered_logs_20260630.part006 \ + --json-out /tmp/lingtai-workload.json +``` + +The split-archive mode streams `logs/events.jsonl` directly. SQLite timing +requires either `--root` or an explicit `--sqlite /path/to/log.sqlite`. +`--projected-sqlite-out` always writes a copy and leaves the source SQLite file +untouched. + +## SQLite vs NoKV Agent Index + +Use the Rust benchmark when comparing `sqlite_current_v1` with +`nokv_agent_index_v1` in the same process. It ingests JSONL into a Holt-backed +agent index, times the typed queries, and emits safe fingerprints instead of raw +user text or tool arguments: + +```bash +cargo run -p nokv-bench --release --bin lingtai-index-bench -- \ + --events-jsonl /path/to/logs/events.jsonl \ + --sqlite /path/to/logs/log.sqlite \ + --projected-sqlite /tmp/lingtai-projected.sqlite \ + --rebuilt-sqlite /tmp/lingtai-rebuilt.sqlite \ + --agent-index /tmp/lingtai-agent-index \ + --iterations 5 \ + --reset +``` + +## LingTai Adapter Contract + +The LingTai-side replacement should keep `logs/events.jsonl` as the +authoritative log and replace `tui/internal/sqlitelog` as a derived query +sidecar: + +Before querying, run: + +```bash +nokv-agent lingtai ingest \ + --store \ + --events /logs/events.jsonl \ + --source-file logs/events.jsonl \ + --agent-id +``` + +Then map the current `sqlitelog` calls to typed commands: + +| LingTai call | NoKV-agent command | +| --- | --- | +| `QueryEventsIndexCoverage` | `coverage --source-file logs/events.jsonl` | +| `StreamSessionEvents` | `session-rows`; map `ts`, `event_type`, `fields_json` back to `SessionEventRow{TS, Type, FieldsJSON}` | +| `QueryErrorEvents` | `errors` | +| `HasTUIClearCompletionEvent` | `clear-completion --source-offset ` | +| `QueryMoltSessionWindows` | `molt-windows` | +| `QueryRecentMoltTimes` | `recent --type psyche_molt` | +| `QueryRecentRefreshCompleteTimes` | `recent --type refresh_complete` | +| `QueryNotificationBlocks` | `notification-blocks` | +| `QueryNotificationBlockSnapshots` | `notification-block-snapshots` | +| `QueryNotifications` | `notification-events` or `notifications --ref-id/--event-id/--call-id/--channel` when a lifecycle filter is needed | +| `QueryNotificationByID` | `notification-by-id --event-id ` | +| `QueryNotificationBefore` | `notification-before --event-id ` | +| `QueryNotificationAfter` | `notification-after --event-id ` | + +Keep the existing JSONL fallback when ingest, coverage, or row streaming fails, +or when coverage does not reach the tail of `events.jsonl`. + +Do not replace LingTai's live JSONL tail path yet. The current NoKV index is a +derived local index with one writer per index directory; live tail remains the +correct source for fresh rows until LingTai has an explicit background ingest +owner. + +The real-data validation on the filtered 2026-06-30 export found three clear +optimization points even after adding projected SQLite indexes, plus several +SQLite-favored point-query shapes: + +| Query | SQLite current p50 | SQLite projected p50 | JSONL-rebuilt SQLite p50 | `nokv_agent_index_v1` p50 | Notes | +| --- | ---: | ---: | ---: | ---: | --- | +| coverage | 120.339 ms | 109.028 ms | 109.913 ms | 0.001 ms | O(1) coverage key; SQLite still computes over rows. | +| tool/action grouping | 117.273 ms | 1.418 ms | 1.788 ms | 0.028 ms | Materialized projection/facet counts beat expression-indexed SQLite. `nokv-agent` also parses stringified `tool_args`, so facet fingerprints intentionally differ from raw SQLite JSON expressions. | +| tool-name grouping | 109.305 ms | 0.845 ms | 1.063 ms | 0.028 ms | Same materialized facet advantage. | +| TUI clear completion | 0.187 ms | 0.224 ms | 0.225 ms | 0.005 ms | Dedicated newest-first TUI clear index avoids scanning `psyche_molt` / `clear_received` rows. | +| notification before/after | 0.011 / 0.011 ms | 0.010 / 0.010 ms | 0.012 / 0.010 ms | 0.004 / 0.006 ms | Direct neighbor keys are slightly faster than SQLite point traversal. | +| notification by id | 0.008 ms | 0.007 ms | 0.008 ms | 0.014 ms | SQLite wins primary-key lookup. | +| notification events 50 | 0.072 ms | 0.047 ms | 0.049 ms | 0.495 ms | SQLite wins newest-first page materialization. | +| latest notification block/pair | 0.188 / 0.074 ms | 0.024 / 0.020 ms | 0.022 / 0.017 ms | 0.178 / 0.095 ms | Projected SQLite wins this point-query shape. | +| recent molt/refresh times | 0.012 / 0.012 ms | 0.017 / 0.015 ms | 0.016 / 0.015 ms | 0.051 / 0.060 ms | SQLite wins small type/time LIMIT queries. | +| `StreamSessionEvents` row projection | 568.077 ms | 578.335 ms | 628.612 ms | 783.415 ms | JSONL-rebuilt SQLite and `nokv-agent` rows/fingerprints match; SQLite is faster for full row replay. Session index values are event-id pointers to avoid duplicating `fields_json`. | +| full event replay diagnostic | n/a | n/a | n/a | 890.113 ms | Internal diagnostic only; LingTai replacement should use `stream_session_rows`. | + +That export's filtered `events.jsonl` and `log.sqlite` do not have identical +coverage (`events.jsonl` accepted 366127 parseable rows and had 112076 session +events; SQLite had 260642 rows, 86477 session-query rows, and a different +source-offset range), so equality fingerprints are not a merge criterion for +the exported sidecar. The `sqlite_rebuilt_from_jsonl_v1` arm rebuilds SQLite +from the same JSONL source and is the equality check: all public sqlitelog-shaped +queries matched rows and fingerprints against `nokv_agent_index_v1`. + +In the final run, projected-current SQLite index build took 1535.123 ms and the +projected DB was 457809920 bytes; JSONL-rebuilt SQLite took 4433.903 ms and was +391921664 bytes; NoKV ingest took 7173.482 ms and the Holt agent index was +5610024663 bytes. That size is still the main tradeoff for materialized +semantic indexes and should be reduced further before proposing NoKV-agent as a +universal drop-in sidecar. + +## Query Shapes + +The first benchmark task set should cover: + +- index coverage and freshness; +- session event replay; +- latest notification block/pair queries and notification browser navigation; +- recent `psyche_molt` and `refresh_complete` times; +- TUI clear completion checks; +- tool/action top-N and windowed facets; +- slow or failed `bash` triage; +- `tool_call_id` / `tool_trace_id` correlation from call to result; +- notification lifecycle by `ref_id`, event id, or channel; +- file-tool workload summaries such as read offsets, file extensions, and grep + shape. + +These are semantic agent-runtime questions. They should be exposed to NoKV as +typed event-index queries, not as SQL compatibility strings. + +## Evidence Rules + +- JSONL remains the source of truth. A missing or stale index must be detectable + through coverage. +- No benchmark task may depend on hidden answer files or special-case product + logic. +- Reports must include ingest time, index size, warm query latency, and result + equality against SQLite for the comparable event set. +- Do not treat filtered JSONL and SQLite files with mismatched coverage as an + equivalence benchmark. They are still useful for finding optimization points. +- Agent-interface runs must also report token count, turns, tool calls, + correctness, and cost. diff --git a/bench/agent-interface/scripts/lingtai_workload_baseline.py b/bench/agent-interface/scripts/lingtai_workload_baseline.py new file mode 100755 index 000000000..128419228 --- /dev/null +++ b/bench/agent-interface/scripts/lingtai_workload_baseline.py @@ -0,0 +1,588 @@ +#!/usr/bin/env python3 +"""Summarize LingTai agent logs and run SQLite baseline probes. + +The script is deliberately read-only. It emits aggregate counts and timing +statistics without printing raw message bodies, tool arguments, or user text. +""" + +from __future__ import annotations + +import argparse +import codecs +import collections +import contextlib +import io +import json +import math +import shutil +import sqlite3 +import statistics +import tarfile +import time +from pathlib import Path +from typing import BinaryIO, Iterable, Iterator + + +SESSION_EVENT_TYPES = ( + "thinking", + "diary", + "text_input", + "text_output", + "tool_call", + "tool_result", + "llm_call", + "llm_response", + "insight", + "consultation_fire", + "notification_pair_injected", + "apriori_summary_generated", + "apriori_summary_cap_refused", + "apriori_summary_failed", + "apriori_summary_empty", + "apriori_summary_no_summarizer", + "aed_attempt", + "aed_exhausted", + "aed_timeout", +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Analyze LingTai events.jsonl and log.sqlite benchmark baselines." + ) + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument( + "--root", + type=Path, + help="Extracted LingTai export root containing logs/events.jsonl and logs/log.sqlite.", + ) + source.add_argument( + "--parts", + nargs="+", + type=Path, + help="Ordered split archive parts. The archive must contain logs/events.jsonl and logs/log.sqlite.", + ) + parser.add_argument("--sqlite", type=Path, help="Override SQLite database path.") + parser.add_argument( + "--projected-sqlite-out", + type=Path, + help="Copy SQLite here, add projected indexes, and measure sqlite_projected_v1.", + ) + parser.add_argument("--top", type=int, default=20, help="Top-N size for counters.") + parser.add_argument( + "--iterations", + type=int, + default=20, + help="SQLite timing iterations after one warmup.", + ) + parser.add_argument( + "--json-out", + type=Path, + help="Write the full summary JSON to this path instead of stdout.", + ) + parser.add_argument( + "--markdown-out", + type=Path, + help="Optional compact Markdown report path.", + ) + return parser.parse_args() + + +class MultiFileReader(io.RawIOBase): + """Concatenate split archive parts without materializing a full archive.""" + + def __init__(self, paths: Iterable[Path]) -> None: + self._paths = list(paths) + if not self._paths: + raise ValueError("at least one archive part is required") + self._index = 0 + self._current: BinaryIO | None = None + + def readable(self) -> bool: + return True + + def close(self) -> None: + if self._current is not None: + self._current.close() + self._current = None + super().close() + + def readinto(self, buffer: bytearray) -> int: + view = memoryview(buffer) + total = 0 + while total < len(view): + if self._current is None: + if self._index >= len(self._paths): + break + self._current = self._paths[self._index].open("rb") + self._index += 1 + chunk = self._current.read(len(view) - total) + if not chunk: + self._current.close() + self._current = None + continue + view[total : total + len(chunk)] = chunk + total += len(chunk) + return total + + +@contextlib.contextmanager +def open_tar_member(parts: list[Path], member_name: str) -> Iterator[BinaryIO]: + reader = io.BufferedReader(MultiFileReader(parts), buffer_size=1024 * 1024) + with reader: + with tarfile.open(fileobj=reader, mode="r|gz") as archive: + for member in archive: + if member.name == member_name: + extracted = archive.extractfile(member) + if extracted is None: + raise RuntimeError(f"archive member is not readable: {member_name}") + with extracted: + yield extracted + return + raise FileNotFoundError(f"archive member not found: {member_name}") + + +def open_event_lines(args: argparse.Namespace) -> Iterator[str]: + if args.root is not None: + with (args.root / "logs" / "events.jsonl").open( + "r", encoding="utf-8", errors="replace" + ) as file: + yield from file + return + with open_tar_member(args.parts, "logs/events.jsonl") as member: + yield from codecs.iterdecode(member, "utf-8", errors="replace") + + +def sqlite_path(args: argparse.Namespace) -> Path | None: + if args.sqlite is not None: + return args.sqlite + if args.root is not None: + candidate = args.root / "logs" / "log.sqlite" + return candidate if candidate.exists() else None + return None + + +def sorted_counter(counter: collections.Counter, limit: int) -> list[dict[str, object]]: + return [ + {"value": str(value), "count": count} + for value, count in counter.most_common(limit) + ] + + +def bucket_int(value: object, buckets: tuple[tuple[int, str], ...]) -> str: + try: + number = int(value) + except (TypeError, ValueError): + return "missing" if value is None else "non_integer" + for upper, label in buckets: + if number <= upper: + return label + return f">{buckets[-1][0]}" + + +def parse_tool_args(raw: object) -> dict[str, object]: + if isinstance(raw, dict): + return raw + if isinstance(raw, str): + try: + decoded = json.loads(raw) + except json.JSONDecodeError: + return {} + if isinstance(decoded, dict): + return decoded + return {} + + +def command_head(command: object) -> str: + if not isinstance(command, str): + return "missing" + words = command.strip().split() + if not words: + return "empty" + head = words[0] + if "/" in head: + return f"path:{Path(head).name or 'unknown'}" + return head + + +def path_ext(path: object) -> str: + if not isinstance(path, str) or not path: + return "missing" + suffix = Path(path).suffix.lower() + return suffix if suffix else "none" + + +def path_class(path: object) -> str: + if not isinstance(path, str) or not path: + return "missing" + parts = [part for part in path.split("/") if part] + if not parts: + return "/" + if path.startswith("/Users/") or path.startswith("/home/"): + return "home_absolute" + if path.startswith("/tmp/") or path.startswith("/var/folders/"): + return "tmp_absolute" + if path.startswith("/"): + return "other_absolute" + if path.startswith("."): + return "relative_dot" + return "relative" + + +def glob_shape(glob: object) -> str: + if not isinstance(glob, str) or not glob: + return "missing" + if glob == "*": + return "all" + if glob.startswith("*.") and "/" not in glob: + return "extension" + if glob.startswith("*.{") and glob.endswith("}"): + return "extension_set" + if "*" not in glob and "/" not in glob: + return "basename" + return "other" + + +def summarize_events(args: argparse.Namespace) -> dict[str, object]: + event_types: collections.Counter[str] = collections.Counter() + field_keys: collections.Counter[str] = collections.Counter() + tool_names: collections.Counter[str] = collections.Counter() + tool_actions: collections.Counter[tuple[str, str]] = collections.Counter() + tool_arg_key_sets: collections.Counter[str] = collections.Counter() + query_tools: collections.Counter[str] = collections.Counter() + bash_heads: collections.Counter[str] = collections.Counter() + read_exts: collections.Counter[str] = collections.Counter() + read_classes: collections.Counter[str] = collections.Counter() + read_offsets: collections.Counter[str] = collections.Counter() + read_limits: collections.Counter[str] = collections.Counter() + grep_globs: collections.Counter[str] = collections.Counter() + grep_limits: collections.Counter[str] = collections.Counter() + parse_errors = 0 + line_count = 0 + + for line in open_event_lines(args): + if not line.strip(): + continue + line_count += 1 + try: + event = json.loads(line) + except json.JSONDecodeError: + parse_errors += 1 + continue + event_type = str(event.get("type", "missing")) + event_types[event_type] += 1 + for key in event: + field_keys[key] += 1 + if event_type != "tool_call": + continue + tool_name = str(event.get("tool_name") or event.get("tool") or "missing") + tool_names[tool_name] += 1 + tool_args = parse_tool_args(event.get("tool_args")) + action = str(tool_args.get("action", "missing")) + tool_actions[(tool_name, action)] += 1 + tool_arg_key_sets[",".join(sorted(tool_args))] += 1 + if tool_args.get("query"): + query_tools[tool_name] += 1 + if tool_name == "bash": + bash_heads[command_head(tool_args.get("command"))] += 1 + elif tool_name == "read": + read_exts[path_ext(tool_args.get("file_path"))] += 1 + read_classes[path_class(tool_args.get("file_path"))] += 1 + read_offsets[ + "missing" + if tool_args.get("offset") is None + else "zero" if str(tool_args.get("offset")) == "0" else "nonzero" + ] += 1 + read_limits[ + bucket_int( + tool_args.get("limit", tool_args.get("max_chars")), + ((100, "<=100"), (500, "<=500"), (2000, "<=2000"), (10000, "<=10000")), + ) + ] += 1 + elif tool_name == "grep": + grep_globs[glob_shape(tool_args.get("glob"))] += 1 + grep_limits[str(tool_args.get("max_matches", "missing"))] += 1 + + top = args.top + return { + "events_jsonl": { + "lines": line_count, + "parse_errors": parse_errors, + "event_types": sorted_counter(event_types, top), + "field_keys": sorted_counter(field_keys, top), + }, + "tool_calls": { + "count": sum(tool_names.values()), + "tools": sorted_counter(tool_names, top), + "tool_actions": [ + {"tool": tool, "action": action, "count": count} + for (tool, action), count in tool_actions.most_common(top) + ], + "arg_key_sets": sorted_counter(tool_arg_key_sets, top), + "query_tools": sorted_counter(query_tools, top), + "bash_command_heads": sorted_counter(bash_heads, top), + "read_extensions": sorted_counter(read_exts, top), + "read_path_classes": sorted_counter(read_classes, top), + "read_offsets": sorted_counter(read_offsets, top), + "read_limits": sorted_counter(read_limits, top), + "grep_globs": sorted_counter(grep_globs, top), + "grep_limits": sorted_counter(grep_limits, top), + }, + } + + +def percentile(values: list[float], fraction: float) -> float: + if not values: + return 0.0 + index = min(len(values) - 1, max(0, math.ceil(len(values) * fraction) - 1)) + return sorted(values)[index] + + +def timed_query( + conn: sqlite3.Connection, sql: str, params: tuple[object, ...], iterations: int +) -> dict[str, object]: + list(conn.execute(sql, params)) + timings: list[float] = [] + rows = 0 + for _ in range(iterations): + started = time.perf_counter() + result = list(conn.execute(sql, params)) + elapsed_ms = (time.perf_counter() - started) * 1000.0 + timings.append(elapsed_ms) + rows = len(result) + return { + "rows": rows, + "p50_ms": round(statistics.median(timings), 3), + "p95_ms": round(percentile(timings, 0.95), 3), + "avg_ms": round(statistics.mean(timings), 3), + } + + +def sqlite_baseline(db_path: Path, iterations: int) -> dict[str, object]: + conn = sqlite3.connect(db_path) + conn.execute("pragma query_only=on") + placeholders = ",".join("?" for _ in SESSION_EVENT_TYPES) + queries = { + "coverage": ( + "select coalesce(min(source_offset), -1), " + "coalesce(max(source_offset), -1), count(source_offset) from events", + (), + ), + "stream_session_events": ( + f"select ts, type, fields_json from events where type in ({placeholders}) order by id asc", + SESSION_EVENT_TYPES, + ), + "latest_notification_block_10": ( + "select id, ts, fields_json, coalesce(source_file, '') " + "from events where type = 'notification_block_injected' order by id desc limit 10", + (), + ), + "latest_notification_pair_10": ( + "select id, ts, fields_json, coalesce(source_file, '') " + "from events where type = 'notification_pair_injected' order by id desc limit 10", + (), + ), + "recent_psyche_molt_10": ( + "select ts from events where type = 'psyche_molt' order by ts desc limit 10", + (), + ), + "recent_refresh_complete_10": ( + "select ts from events where type = 'refresh_complete' order by ts desc limit 10", + (), + ), + "error_events": ( + "select ts, type, fields_json from events " + "where type in ('aed_attempt', 'aed_exhausted', 'refresh_init_error') " + "order by id desc", + (), + ), + "tool_name_group_json_extract": ( + "select json_extract(fields_json, '$.tool_name') as tool_name, count(*) " + "from events where type = 'tool_call' and json_valid(fields_json) " + "group by tool_name order by count(*) desc limit 30", + (), + ), + "tool_action_group_json_extract": ( + "select json_extract(fields_json, '$.tool_name') as tool_name, " + "json_extract(fields_json, '$.tool_args.action') as action, count(*) " + "from events where type = 'tool_call' and json_valid(fields_json) " + "group by tool_name, action order by count(*) desc limit 50", + (), + ), + } + stats = { + name: timed_query(conn, sql, tuple(params), iterations) + for name, (sql, params) in queries.items() + } + return { + "path": str(db_path), + "bytes": db_path.stat().st_size, + "sqlite_version": conn.execute("select sqlite_version()").fetchone()[0], + "event_rows": conn.execute("select count(*) from events").fetchone()[0], + "invalid_fields_json": conn.execute( + "select count(*) from events where not json_valid(fields_json)" + ).fetchone()[0], + "queries": stats, + } + + +def sqlite_index_sizes(conn: sqlite3.Connection, prefix: str) -> list[dict[str, object]]: + try: + rows = conn.execute( + "select name, sum(pgsize) from dbstat where name like ? group by name order by name", + (f"{prefix}%",), + ).fetchall() + except sqlite3.DatabaseError: + return [] + return [{"name": name, "bytes": bytes_} for name, bytes_ in rows] + + +def projected_sqlite_baseline( + source_db: Path, projected_db: Path, iterations: int +) -> dict[str, object]: + started = time.perf_counter() + projected_db.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source_db, projected_db) + conn = sqlite3.connect(projected_db) + index_sql = ( + "create index if not exists idx_events_source_offset_only " + "on events(source_offset) where source_offset is not null", + "create index if not exists idx_events_type_id_desc " + "on events(type, id desc)", + "create index if not exists idx_events_tool_name_expr " + "on events(json_extract(fields_json, '$.tool_name')) " + "where type = 'tool_call' and json_valid(fields_json)", + "create index if not exists idx_events_tool_action_expr " + "on events(json_extract(fields_json, '$.tool_name'), " + "json_extract(fields_json, '$.tool_args.action')) " + "where type = 'tool_call' and json_valid(fields_json)", + ) + for statement in index_sql: + conn.execute(statement) + conn.commit() + build_ms = (time.perf_counter() - started) * 1000.0 + conn.execute("pragma query_only=on") + placeholders = ",".join("?" for _ in SESSION_EVENT_TYPES) + queries = { + "coverage": ( + "select coalesce(min(source_offset), -1), " + "coalesce(max(source_offset), -1), count(source_offset) from events", + (), + ), + "stream_session_events": ( + f"select ts, type, fields_json from events where type in ({placeholders}) order by id asc", + SESSION_EVENT_TYPES, + ), + "latest_notification_block_10": ( + "select id, ts, fields_json, coalesce(source_file, '') " + "from events where type = 'notification_block_injected' order by id desc limit 10", + (), + ), + "latest_notification_pair_10": ( + "select id, ts, fields_json, coalesce(source_file, '') " + "from events where type = 'notification_pair_injected' order by id desc limit 10", + (), + ), + "tool_name_group_forced_projection": ( + "select json_extract(fields_json, '$.tool_name') as tool_name, count(*) " + "from events indexed by idx_events_tool_name_expr " + "where type = 'tool_call' and json_valid(fields_json) " + "group by tool_name order by count(*) desc limit 30", + (), + ), + "tool_action_group_forced_projection": ( + "select json_extract(fields_json, '$.tool_name') as tool_name, " + "json_extract(fields_json, '$.tool_args.action') as action, count(*) " + "from events indexed by idx_events_tool_action_expr " + "where type = 'tool_call' and json_valid(fields_json) " + "group by tool_name, action order by count(*) desc limit 50", + (), + ), + "latest_read_tool_calls_100": ( + "select id, ts, fields_json from events indexed by idx_events_tool_name_expr " + "where type = 'tool_call' and json_valid(fields_json) " + "and json_extract(fields_json, '$.tool_name') = 'read' " + "order by id desc limit 100", + (), + ), + } + return { + "path": str(projected_db), + "bytes": projected_db.stat().st_size, + "build_ms": round(build_ms, 3), + "projected_index_sizes": sqlite_index_sizes(conn, "idx_events_"), + "queries": { + name: timed_query(conn, sql, tuple(params), iterations) + for name, (sql, params) in queries.items() + }, + } + + +def write_markdown(summary: dict[str, object], path: Path) -> None: + events = summary["events_jsonl"] + tools = summary["tool_calls"] + sqlite = summary.get("sqlite_current_v1") + lines = [ + "# LingTai Workload Baseline", + "", + f"- JSONL lines: {events['lines']}", + f"- JSON parse errors: {events['parse_errors']}", + f"- tool_call rows: {tools['count']}", + "", + "## Top Tools", + "", + "| Tool | Count |", + "| --- | ---: |", + ] + lines.extend( + f"| `{row['value']}` | {row['count']} |" for row in tools["tools"] + ) + lines.extend(["", "## Top Tool Actions", "", "| Tool | Action | Count |", "| --- | --- | ---: |"]) + lines.extend( + f"| `{row['tool']}` | `{row['action']}` | {row['count']} |" + for row in tools["tool_actions"] + ) + if sqlite: + lines.extend(["", "## SQLite Current Baseline", "", "| Query | Rows | p50 ms | p95 ms |", "| --- | ---: | ---: | ---: |"]) + for name, row in sqlite["queries"].items(): + lines.append(f"| `{name}` | {row['rows']} | {row['p50_ms']} | {row['p95_ms']} |") + projected = summary.get("sqlite_projected_v1") + if projected: + lines.extend( + [ + "", + "## SQLite Projected Baseline", + "", + f"- build ms: {projected['build_ms']}", + f"- db bytes: {projected['bytes']}", + "", + "| Query | Rows | p50 ms | p95 ms |", + "| --- | ---: | ---: | ---: |", + ] + ) + for name, row in projected["queries"].items(): + lines.append(f"| `{name}` | {row['rows']} | {row['p50_ms']} | {row['p95_ms']} |") + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + + +def main() -> None: + args = parse_args() + summary = summarize_events(args) + db_path = sqlite_path(args) + if db_path is not None and db_path.exists(): + summary["sqlite_current_v1"] = sqlite_baseline(db_path, args.iterations) + if args.projected_sqlite_out is not None: + summary["sqlite_projected_v1"] = projected_sqlite_baseline( + db_path, args.projected_sqlite_out, args.iterations + ) + if args.markdown_out is not None: + write_markdown(summary, args.markdown_out) + data = json.dumps(summary, indent=2, sort_keys=True) + if args.json_out is not None: + args.json_out.write_text(data + "\n", encoding="utf-8") + else: + print(data) + + +if __name__ == "__main__": + main() diff --git a/bench/src/bin/lingtai-index-bench.rs b/bench/src/bin/lingtai-index-bench.rs new file mode 100644 index 000000000..ae23629de --- /dev/null +++ b/bench/src/bin/lingtai-index-bench.rs @@ -0,0 +1,1392 @@ +use std::collections::{BTreeMap, VecDeque}; +use std::fs::{self, File}; +use std::io::{BufRead, BufReader}; +use std::path::{Path, PathBuf}; +use std::time::Instant; + +use nokv_agent::event::{ + ingest_jsonl_reader, AgentEventStore, ErrorEventsRequest, EventRecord, HoltAgentEventStore, + JsonlIngestOptions, LatestEventsRequest, NotificationEventByIdRequest, + NotificationEventsRequest, NotificationNeighborDirection, NotificationNeighborRequest, + RecentTimesRequest, SessionEventRow, SessionEventsRequest, SessionRowsRequest, ToolFacet, + ToolFacetRequest, TuiClearCompletionRequest, LINGTAI_SESSION_EVENT_TYPES, +}; +use rusqlite::{params_from_iter, Connection}; +use serde::Serialize; +use serde_json::{json, Map, Value}; +use sha2::{Digest, Sha256}; +use walkdir::WalkDir; + +const ERROR_EVENT_TYPES: &[&str] = &["aed_attempt", "aed_exhausted", "refresh_init_error"]; +const MAX_INLINE_FIELDS_JSON_BYTES: usize = 32 * 1024; + +type BenchResult = Result>; + +#[derive(Debug)] +struct Options { + events_jsonl: PathBuf, + sqlite: PathBuf, + projected_sqlite: Option, + rebuilt_sqlite: Option, + agent_index: PathBuf, + source_file: String, + agent_id: String, + iterations: usize, + reset: bool, +} + +#[derive(Clone, Debug, Serialize)] +struct QueryDigest { + rows: usize, + fingerprint: String, +} + +#[derive(Clone, Debug, Serialize)] +struct QueryStats { + rows: usize, + p50_ms: f64, + p95_ms: f64, + avg_ms: f64, + fingerprint: String, +} + +#[derive(Clone, Debug, Serialize)] +struct QueryComparison { + rows_match: bool, + fingerprint_match: bool, +} + +#[derive(Clone, Copy, Debug, Default)] +struct NotificationPivots { + latest: Option, + earliest: Option, +} + +#[derive(Clone, Debug, Serialize)] +struct ProjectedSqliteStats { + path: PathBuf, + bytes: u64, + build_ms: f64, + sqlite_version: String, + projected_index_sizes: Vec, + queries: BTreeMap, +} + +#[derive(Clone, Debug, Serialize)] +struct RebuiltSqliteStats { + path: PathBuf, + bytes: u64, + build_ms: f64, + sqlite_version: String, + event_rows: u64, + projected_index_sizes: Vec, + queries: BTreeMap, +} + +#[derive(Clone, Debug, Serialize)] +struct SqliteIndexSize { + name: String, + bytes: u64, +} + +fn main() { + if let Err(err) = run() { + eprintln!("{err}"); + std::process::exit(2); + } +} + +fn run() -> BenchResult<()> { + let options = parse_args()?; + if options.reset && options.agent_index.exists() { + std::fs::remove_dir_all(&options.agent_index)?; + } + + let sqlite = + Connection::open_with_flags(&options.sqlite, rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY)?; + sqlite.execute_batch("pragma query_only=on")?; + + let ingest_started = Instant::now(); + let store = HoltAgentEventStore::open_file(&options.agent_index)?; + let events_file = File::open(&options.events_jsonl)?; + let file_size = events_file.metadata()?.len(); + let ingest_report = ingest_jsonl_reader( + &store, + JsonlIngestOptions { + agent_id: options.agent_id.clone(), + source_file: options.source_file.clone(), + file_size, + }, + BufReader::new(events_file), + )?; + let ingest_ms = elapsed_ms(ingest_started); + + let sqlite_queries = sqlite_queries(&sqlite, options.iterations)?; + let projected_sqlite = options + .projected_sqlite + .as_ref() + .map(|path| projected_sqlite_baseline(&options.sqlite, path, options.iterations)) + .transpose()?; + let rebuilt_sqlite = options + .rebuilt_sqlite + .as_ref() + .map(|path| rebuilt_sqlite_baseline(&options, path, options.iterations)) + .transpose()?; + let agent_queries = agent_queries(&store, &options, options.iterations)?; + let comparisons = compare_queries(&sqlite_queries, &agent_queries); + let projected_comparisons = projected_sqlite + .as_ref() + .map(|projected| compare_queries(&projected.queries, &agent_queries)); + let rebuilt_comparisons = rebuilt_sqlite + .as_ref() + .map(|rebuilt| compare_queries(&rebuilt.queries, &agent_queries)); + + let output = json!({ + "events_jsonl": { + "path": options.events_jsonl, + "bytes": file_size, + }, + "sqlite_current_v1": { + "path": options.sqlite, + "bytes": options.sqlite.metadata()?.len(), + "sqlite_version": sqlite.query_row("select sqlite_version()", [], |row| row.get::<_, String>(0))?, + "queries": sqlite_queries, + }, + "sqlite_projected_v1": projected_sqlite, + "sqlite_rebuilt_from_jsonl_v1": rebuilt_sqlite, + "nokv_agent_index_v1": { + "path": options.agent_index, + "bytes": dir_size(&options.agent_index)?, + "ingest_ms": round_ms(ingest_ms), + "ingest_report": ingest_report, + "queries": agent_queries, + }, + "comparisons": comparisons, + "projected_comparisons": projected_comparisons, + "rebuilt_comparisons": rebuilt_comparisons, + "notes": [ + "JSONL is the source of truth; nokv_agent_index_v1 is a rebuildable derived index.", + "sqlite_rebuilt_from_jsonl_v1 is emitted only when --rebuilt-sqlite PATH is provided and uses the same events.jsonl input as nokv_agent_index_v1.", + "Fingerprints hash only type, timestamp, source offset, and aggregate labels; raw user text and tool arguments are not printed.", + "sqlite_projected_v1 is emitted only when --projected-sqlite PATH is provided; use it before making optimization claims." + ], + }); + println!("{}", serde_json::to_string_pretty(&output)?); + Ok(()) +} + +fn parse_args() -> BenchResult { + let mut args = std::env::args().skip(1).collect::>(); + if args.iter().any(|arg| arg == "--help" || arg == "-h") { + return Err(usage().into()); + } + let events_jsonl = required_path(&mut args, "--events-jsonl")?; + let sqlite = required_path(&mut args, "--sqlite")?; + let projected_sqlite = option_value(&mut args, "--projected-sqlite")?.map(PathBuf::from); + let rebuilt_sqlite = option_value(&mut args, "--rebuilt-sqlite")?.map(PathBuf::from); + let agent_index = required_path(&mut args, "--agent-index")?; + let source_file = + option_value(&mut args, "--source-file")?.unwrap_or_else(|| "logs/events.jsonl".to_owned()); + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let iterations = option_value(&mut args, "--iterations")? + .map(|value| value.parse::()) + .transpose()? + .unwrap_or(20); + let reset = take_flag(&mut args, "--reset"); + if !args.is_empty() { + return Err(format!("unexpected argument {}", args.front().unwrap()).into()); + } + Ok(Options { + events_jsonl, + sqlite, + projected_sqlite, + rebuilt_sqlite, + agent_index, + source_file, + agent_id, + iterations, + reset, + }) +} + +fn sqlite_queries( + conn: &Connection, + iterations: usize, +) -> BenchResult> { + let notification_pivots = sqlite_notification_pivots(conn)?; + let mut queries = BTreeMap::new(); + queries.insert( + "coverage".to_owned(), + timed(iterations, || sqlite_coverage(conn))?, + ); + queries.insert( + "stream_session_rows".to_owned(), + timed(iterations, || sqlite_session_rows(conn))?, + ); + queries.insert( + "latest_notification_block_10".to_owned(), + timed(iterations, || { + sqlite_latest_by_type(conn, "notification_block_injected", 10) + })?, + ); + queries.insert( + "latest_notification_pair_10".to_owned(), + timed(iterations, || { + sqlite_latest_by_type(conn, "notification_pair_injected", 10) + })?, + ); + queries.insert( + "notification_events_50".to_owned(), + timed(iterations, || sqlite_notification_events(conn, 50))?, + ); + queries.insert( + "notification_by_latest_id".to_owned(), + timed(iterations, || { + sqlite_notification_by_id(conn, notification_pivots.latest) + })?, + ); + queries.insert( + "notification_before_latest".to_owned(), + timed(iterations, || { + sqlite_notification_before(conn, notification_pivots.latest) + })?, + ); + queries.insert( + "notification_after_earliest".to_owned(), + timed(iterations, || { + sqlite_notification_after(conn, notification_pivots.earliest) + })?, + ); + queries.insert( + "recent_psyche_molt_10".to_owned(), + timed(iterations, || sqlite_recent_times(conn, "psyche_molt", 10))?, + ); + queries.insert( + "molt_session_windows".to_owned(), + timed(iterations, || sqlite_molt_session_windows(conn))?, + ); + queries.insert( + "recent_refresh_complete_10".to_owned(), + timed(iterations, || { + sqlite_recent_times(conn, "refresh_complete", 10) + })?, + ); + queries.insert( + "tui_clear_completion_from_zero".to_owned(), + timed(iterations, || sqlite_tui_clear_completion(conn, 0))?, + ); + queries.insert( + "error_events".to_owned(), + timed(iterations, || sqlite_error_events(conn))?, + ); + queries.insert( + "tool_name_group_json_extract".to_owned(), + timed(iterations, || sqlite_tool_name_facets(conn, 30))?, + ); + queries.insert( + "tool_action_group_json_extract".to_owned(), + timed(iterations, || sqlite_tool_action_facets(conn, 50))?, + ); + Ok(queries) +} + +fn projected_sqlite_baseline( + source: &Path, + projected: &Path, + iterations: usize, +) -> BenchResult { + if source == projected { + return Err("--projected-sqlite must not point at --sqlite".into()); + } + if let Some(parent) = projected.parent() { + fs::create_dir_all(parent)?; + } + if projected.exists() { + fs::remove_file(projected)?; + } + let started = Instant::now(); + fs::copy(source, projected)?; + let mut conn = Connection::open(projected)?; + create_projected_sqlite_indexes(&mut conn)?; + let build_ms = elapsed_ms(started); + conn.execute_batch("pragma query_only=on")?; + let queries = projected_sqlite_queries(&conn, iterations)?; + Ok(ProjectedSqliteStats { + path: projected.to_path_buf(), + bytes: projected.metadata()?.len(), + build_ms: round_ms(build_ms), + sqlite_version: conn.query_row("select sqlite_version()", [], |row| row.get(0))?, + projected_index_sizes: sqlite_index_sizes(&conn)?, + queries, + }) +} + +fn rebuilt_sqlite_baseline( + options: &Options, + rebuilt: &Path, + iterations: usize, +) -> BenchResult { + if let Some(parent) = rebuilt.parent() { + fs::create_dir_all(parent)?; + } + if rebuilt.exists() { + fs::remove_file(rebuilt)?; + } + + let started = Instant::now(); + let mut conn = Connection::open(rebuilt)?; + conn.execute_batch( + " + pragma journal_mode=off; + pragma synchronous=off; + create table events ( + id integer primary key autoincrement, + ts real not null, + type text not null, + agent_address text, + agent_name_snapshot text, + fields_json text not null, + source_file text, + source_offset integer, + source_line integer, + source_kind text, + scope text, + run_id text, + inserted_at text not null default (strftime('%Y-%m-%dT%H:%M:%fZ','now')) + ); + ", + )?; + let event_rows = insert_jsonl_sqlite_events(&mut conn, options)?; + create_projected_sqlite_indexes(&mut conn)?; + let build_ms = elapsed_ms(started); + + conn.execute_batch("pragma query_only=on")?; + let queries = projected_sqlite_queries(&conn, iterations)?; + Ok(RebuiltSqliteStats { + path: rebuilt.to_path_buf(), + bytes: rebuilt.metadata()?.len(), + build_ms: round_ms(build_ms), + sqlite_version: conn.query_row("select sqlite_version()", [], |row| row.get(0))?, + event_rows, + projected_index_sizes: sqlite_index_sizes(&conn)?, + queries, + }) +} + +fn insert_jsonl_sqlite_events(conn: &mut Connection, options: &Options) -> BenchResult { + let file = File::open(&options.events_jsonl)?; + let mut reader = BufReader::new(file); + let tx = conn.transaction()?; + let mut statement = tx.prepare( + " + insert into events (ts, type, fields_json, source_file, source_offset, source_line, source_kind) + values (?1, ?2, ?3, ?4, ?5, ?6, 'events_jsonl') + ", + )?; + + let mut offset = 0_u64; + let mut line_no = 0_u64; + let mut inserted = 0_u64; + loop { + let mut line = Vec::new(); + let read = reader.read_until(b'\n', &mut line)?; + if read == 0 { + break; + } + let line_offset = offset; + offset = offset.saturating_add(read as u64); + if line.iter().all(|byte| byte.is_ascii_whitespace()) { + continue; + } + if !line.ends_with(b"\n") { + break; + } + line_no = line_no.saturating_add(1); + let value: Value = match serde_json::from_slice(&line) { + Ok(value) => value, + Err(_) => continue, + }; + let Some(object) = value.as_object() else { + continue; + }; + let Some(event_type) = object.get("type").and_then(Value::as_str) else { + continue; + }; + let ts = object.get("ts").and_then(Value::as_f64).unwrap_or(0.0); + let fields_json = sqlite_fields_json(object)?; + statement.execute(( + ts, + event_type, + fields_json, + options.source_file.as_str(), + line_offset as i64, + line_no as i64, + ))?; + inserted = inserted.saturating_add(1); + } + drop(statement); + tx.commit()?; + Ok(inserted) +} + +fn sqlite_fields_json(object: &Map) -> BenchResult { + let mut fields = object.clone(); + for key in ["type", "ts", "address", "agent_name"] { + fields.remove(key); + } + let value = compact_fields_json(Value::Object(fields)); + Ok(serde_json::to_string(&value)?) +} + +fn compact_fields_json(value: Value) -> Value { + let encoded_len = serde_json::to_vec(&value) + .map(|bytes| bytes.len()) + .unwrap_or(usize::MAX); + if encoded_len <= MAX_INLINE_FIELDS_JSON_BYTES { + return value; + } + let keys = value + .as_object() + .map(|object| { + object + .keys() + .cloned() + .map(Value::String) + .collect::>() + }) + .unwrap_or_default(); + serde_json::json!({ + "_nokv_compacted": true, + "original_json_bytes": encoded_len, + "keys": keys, + }) +} + +fn create_projected_sqlite_indexes(conn: &mut Connection) -> BenchResult<()> { + let session_types = LINGTAI_SESSION_EVENT_TYPES + .iter() + .map(|value| sql_string_literal(value)) + .collect::>() + .join(","); + let tx = conn.transaction()?; + tx.execute_batch(&format!( + " + create index if not exists idx_events_source_offset_only + on events(source_offset) + where source_offset is not null; + create index if not exists idx_events_session_id + on events(id) + where type in ({session_types}); + create index if not exists idx_events_type_id_desc + on events(type, id desc); + create index if not exists idx_events_type_ts_desc + on events(type, ts desc); + create index if not exists idx_events_notification_id_desc + on events(id desc) + where type like '%notification%'; + create index if not exists idx_events_tui_clear_id_desc + on events(id desc) + where type in ('psyche_molt','clear_received') + and source_offset is not null + and json_valid(fields_json) + and json_extract(fields_json, '$.source') = 'tui'; + create index if not exists idx_events_tool_name_expr + on events(coalesce(json_extract(fields_json, '$.tool_name'), '')) + where type = 'tool_call' + and json_valid(fields_json) + and json_extract(fields_json, '$.tool_name') is not null; + create index if not exists idx_events_tool_action_expr + on events( + coalesce(json_extract(fields_json, '$.tool_name'), ''), + coalesce(json_extract(fields_json, '$.tool_args.action'), '') + ) + where type = 'tool_call' + and json_valid(fields_json) + and json_extract(fields_json, '$.tool_name') is not null; + " + ))?; + tx.commit()?; + Ok(()) +} + +fn projected_sqlite_queries( + conn: &Connection, + iterations: usize, +) -> BenchResult> { + let mut queries = sqlite_queries(conn, iterations)?; + queries.insert( + "tool_name_group_json_extract".to_owned(), + timed(iterations, || sqlite_tool_name_facets_indexed(conn, 30))?, + ); + queries.insert( + "tool_action_group_json_extract".to_owned(), + timed(iterations, || sqlite_tool_action_facets_indexed(conn, 50))?, + ); + Ok(queries) +} + +fn sqlite_index_sizes(conn: &Connection) -> BenchResult> { + let mut out = Vec::new(); + let mut statement = match conn.prepare( + "select name, coalesce(sum(pgsize), 0) \ + from dbstat where name like 'idx_events_%' group by name order by name", + ) { + Ok(statement) => statement, + Err(_) => return Ok(out), + }; + let mut rows = statement.query([])?; + while let Some(row) = rows.next()? { + let name: String = row.get(0)?; + let bytes: i64 = row.get(1)?; + out.push(SqliteIndexSize { + name, + bytes: bytes.max(0) as u64, + }); + } + Ok(out) +} + +fn agent_queries( + store: &HoltAgentEventStore, + options: &Options, + iterations: usize, +) -> BenchResult> { + let notification_pivots = agent_notification_pivots(store, options)?; + let mut queries = BTreeMap::new(); + queries.insert( + "coverage".to_owned(), + timed(iterations, || agent_coverage(store, options))?, + ); + queries.insert( + "stream_session_rows".to_owned(), + timed(iterations, || agent_session_rows(store, options))?, + ); + queries.insert( + "stream_session_events_full".to_owned(), + timed(iterations, || agent_session_events(store, options))?, + ); + queries.insert( + "latest_notification_block_10".to_owned(), + timed(iterations, || { + agent_latest_by_type(store, options, "notification_block_injected", 10) + })?, + ); + queries.insert( + "latest_notification_pair_10".to_owned(), + timed(iterations, || { + agent_latest_by_type(store, options, "notification_pair_injected", 10) + })?, + ); + queries.insert( + "notification_events_50".to_owned(), + timed(iterations, || agent_notification_events(store, options, 50))?, + ); + queries.insert( + "notification_by_latest_id".to_owned(), + timed(iterations, || { + agent_notification_by_id(store, options, notification_pivots.latest) + })?, + ); + queries.insert( + "notification_before_latest".to_owned(), + timed(iterations, || { + agent_notification_before(store, options, notification_pivots.latest) + })?, + ); + queries.insert( + "notification_after_earliest".to_owned(), + timed(iterations, || { + agent_notification_after(store, options, notification_pivots.earliest) + })?, + ); + queries.insert( + "recent_psyche_molt_10".to_owned(), + timed(iterations, || { + agent_recent_times(store, options, "psyche_molt", 10) + })?, + ); + queries.insert( + "molt_session_windows".to_owned(), + timed(iterations, || agent_molt_session_windows(store, options))?, + ); + queries.insert( + "recent_refresh_complete_10".to_owned(), + timed(iterations, || { + agent_recent_times(store, options, "refresh_complete", 10) + })?, + ); + queries.insert( + "tui_clear_completion_from_zero".to_owned(), + timed(iterations, || agent_tui_clear_completion(store, options, 0))?, + ); + queries.insert( + "error_events".to_owned(), + timed(iterations, || agent_error_events(store, options))?, + ); + queries.insert( + "tool_name_group_json_extract".to_owned(), + timed(iterations, || agent_tool_name_facets(store, options, 30))?, + ); + queries.insert( + "tool_action_group_json_extract".to_owned(), + timed(iterations, || agent_tool_action_facets(store, options, 50))?, + ); + Ok(queries) +} + +fn sqlite_coverage(conn: &Connection) -> BenchResult { + let row = conn.query_row( + "select coalesce(min(source_offset), -1), coalesce(max(source_offset), -1), count(source_offset) from events", + [], + |row| Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?, row.get::<_, i64>(2)?)), + )?; + let mut digest = SafeDigest::new(); + digest.field("min", row.0); + digest.field("max", row.1); + digest.field("count", row.2); + Ok(QueryDigest { + rows: 1, + fingerprint: digest.finish(), + }) +} + +fn sqlite_session_rows(conn: &Connection) -> BenchResult { + let placeholders = placeholders(LINGTAI_SESSION_EVENT_TYPES.len()); + let sql = format!( + "select ts, type, fields_json from events where type in ({placeholders}) order by id asc" + ); + let mut statement = conn.prepare(&sql)?; + let mut rows = statement.query(params_from_iter(LINGTAI_SESSION_EVENT_TYPES.iter()))?; + let mut count = 0_usize; + let mut digest = SafeDigest::new(); + while let Some(row) = rows.next()? { + let ts: f64 = row.get(0)?; + let event_type: String = row.get(1)?; + let _fields_json: String = row.get(2)?; + digest.event(&event_type, ts, -1); + count += 1; + } + Ok(QueryDigest { + rows: count, + fingerprint: digest.finish(), + }) +} + +fn sqlite_latest_by_type( + conn: &Connection, + event_type: &str, + limit: usize, +) -> BenchResult { + let sql = format!( + "select type, ts, source_offset from events where type = ?1 order by id desc limit {limit}" + ); + sqlite_event_digest(conn, &sql, &[event_type]) +} + +fn sqlite_notification_events(conn: &Connection, limit: usize) -> BenchResult { + let sql = format!( + "select type, ts, source_offset from events \ + where type like '%notification%' order by id desc limit {limit}" + ); + sqlite_event_digest(conn, &sql, &[]) +} + +fn sqlite_notification_pivots(conn: &Connection) -> BenchResult { + let latest = conn.query_row( + "select max(id) from events where type like '%notification%'", + [], + |row| row.get::<_, Option>(0), + )?; + let earliest = conn.query_row( + "select min(id) from events where type like '%notification%'", + [], + |row| row.get::<_, Option>(0), + )?; + Ok(NotificationPivots { + latest: latest.map(|id| id as u64), + earliest: earliest.map(|id| id as u64), + }) +} + +fn sqlite_notification_by_id(conn: &Connection, event_id: Option) -> BenchResult { + let Some(event_id) = event_id else { + return Ok(empty_digest()); + }; + let sql = format!("select type, ts, source_offset from events where id = {event_id}"); + sqlite_event_digest(conn, &sql, &[]) +} + +fn sqlite_notification_before( + conn: &Connection, + pivot_event_id: Option, +) -> BenchResult { + let Some(pivot_event_id) = pivot_event_id else { + return Ok(empty_digest()); + }; + let sql = format!( + "select type, ts, source_offset from events \ + where type like '%notification%' and id < {pivot_event_id} \ + order by id desc limit 1" + ); + sqlite_event_digest(conn, &sql, &[]) +} + +fn sqlite_notification_after( + conn: &Connection, + pivot_event_id: Option, +) -> BenchResult { + let Some(pivot_event_id) = pivot_event_id else { + return Ok(empty_digest()); + }; + let sql = format!( + "select type, ts, source_offset from events \ + where type like '%notification%' and id > {pivot_event_id} \ + order by id asc limit 1" + ); + sqlite_event_digest(conn, &sql, &[]) +} + +fn sqlite_recent_times( + conn: &Connection, + event_type: &str, + limit: usize, +) -> BenchResult { + let sql = + format!("select type, ts, -1 from events where type = ?1 order by ts desc limit {limit}"); + sqlite_event_digest(conn, &sql, &[event_type]) +} + +fn sqlite_molt_session_windows(conn: &Connection) -> BenchResult { + let mut statement = + conn.prepare("select ts from events where type = 'psyche_molt' order by ts desc limit 2")?; + let mut rows = statement.query([])?; + let mut times = Vec::new(); + while let Some(row) = rows.next()? { + times.push(row.get::<_, f64>(0)?); + } + let current_since = times.first().copied(); + let last_since = times.get(1).copied(); + let last_before = last_since.and(current_since); + let mut digest = SafeDigest::new(); + digest.field("ok", 1); + digest.field( + "current_since", + current_since.map(timestamp_micros_i64).unwrap_or(-1), + ); + digest.field( + "last_since", + last_since.map(timestamp_micros_i64).unwrap_or(-1), + ); + digest.field( + "last_before", + last_before.map(timestamp_micros_i64).unwrap_or(-1), + ); + Ok(QueryDigest { + rows: 1, + fingerprint: digest.finish(), + }) +} + +fn sqlite_tui_clear_completion(conn: &Connection, source_offset: u64) -> BenchResult { + let sql = format!( + "select type, ts, source_offset from events \ + where type in ('psyche_molt','clear_received') \ + and source_offset >= {source_offset} \ + and json_valid(fields_json) \ + and json_extract(fields_json, '$.source') = 'tui' \ + order by id desc limit 1" + ); + sqlite_event_digest(conn, &sql, &[]) +} + +fn sqlite_error_events(conn: &Connection) -> BenchResult { + let placeholders = placeholders(ERROR_EVENT_TYPES.len()); + let sql = format!( + "select type, ts, source_offset from events where type in ({placeholders}) order by id desc" + ); + sqlite_event_digest(conn, &sql, ERROR_EVENT_TYPES) +} + +fn sqlite_event_digest(conn: &Connection, sql: &str, params: &[&str]) -> BenchResult { + let mut statement = conn.prepare(sql)?; + let mut rows = statement.query(params_from_iter(params.iter()))?; + let mut count = 0_usize; + let mut digest = SafeDigest::new(); + while let Some(row) = rows.next()? { + let event_type: String = row.get(0)?; + let ts: f64 = row.get(1)?; + let source_offset: Option = row.get(2)?; + digest.event(&event_type, ts, source_offset.unwrap_or(-1)); + count += 1; + } + Ok(QueryDigest { + rows: count, + fingerprint: digest.finish(), + }) +} + +fn sqlite_tool_name_facets(conn: &Connection, limit: usize) -> BenchResult { + let sql = format!( + "select coalesce(json_extract(fields_json, '$.tool_name'), ''), count(*) \ + from events where type = 'tool_call' and json_valid(fields_json) \ + and json_extract(fields_json, '$.tool_name') is not null \ + group by 1 order by count(*) desc, 1 asc limit {limit}" + ); + let mut statement = conn.prepare(&sql)?; + let mut rows = statement.query([])?; + let mut count = 0_usize; + let mut digest = SafeDigest::new(); + while let Some(row) = rows.next()? { + let tool_name: String = row.get(0)?; + let n: i64 = row.get(1)?; + digest.facet(&tool_name, None, n as u64); + count += 1; + } + Ok(QueryDigest { + rows: count, + fingerprint: digest.finish(), + }) +} + +fn sqlite_tool_name_facets_indexed(conn: &Connection, limit: usize) -> BenchResult { + let sql = format!( + "select coalesce(json_extract(fields_json, '$.tool_name'), ''), count(*) \ + from events indexed by idx_events_tool_name_expr \ + where type = 'tool_call' and json_valid(fields_json) \ + and json_extract(fields_json, '$.tool_name') is not null \ + group by 1 order by count(*) desc, 1 asc limit {limit}" + ); + sqlite_facet_digest(conn, &sql, false) +} + +fn sqlite_tool_action_facets(conn: &Connection, limit: usize) -> BenchResult { + let sql = format!( + "select coalesce(json_extract(fields_json, '$.tool_name'), ''), \ + coalesce(json_extract(fields_json, '$.tool_args.action'), ''), count(*) \ + from events where type = 'tool_call' and json_valid(fields_json) \ + and json_extract(fields_json, '$.tool_name') is not null \ + group by 1, 2 order by count(*) desc, 1 asc, 2 asc limit {limit}" + ); + let mut statement = conn.prepare(&sql)?; + let mut rows = statement.query([])?; + let mut count = 0_usize; + let mut digest = SafeDigest::new(); + while let Some(row) = rows.next()? { + let tool_name: String = row.get(0)?; + let action: String = row.get(1)?; + let n: i64 = row.get(2)?; + digest.facet(&tool_name, Some(&action), n as u64); + count += 1; + } + Ok(QueryDigest { + rows: count, + fingerprint: digest.finish(), + }) +} + +fn sqlite_tool_action_facets_indexed(conn: &Connection, limit: usize) -> BenchResult { + let sql = format!( + "select coalesce(json_extract(fields_json, '$.tool_name'), ''), \ + coalesce(json_extract(fields_json, '$.tool_args.action'), ''), count(*) \ + from events indexed by idx_events_tool_action_expr \ + where type = 'tool_call' and json_valid(fields_json) \ + and json_extract(fields_json, '$.tool_name') is not null \ + group by 1, 2 order by count(*) desc, 1 asc, 2 asc limit {limit}" + ); + sqlite_facet_digest(conn, &sql, true) +} + +fn sqlite_facet_digest( + conn: &Connection, + sql: &str, + include_action: bool, +) -> BenchResult { + let mut statement = conn.prepare(sql)?; + let mut rows = statement.query([])?; + let mut count = 0_usize; + let mut digest = SafeDigest::new(); + while let Some(row) = rows.next()? { + let tool_name: String = row.get(0)?; + if include_action { + let action: String = row.get(1)?; + let n: i64 = row.get(2)?; + digest.facet(&tool_name, Some(&action), n as u64); + } else { + let n: i64 = row.get(1)?; + digest.facet(&tool_name, None, n as u64); + } + count += 1; + } + Ok(QueryDigest { + rows: count, + fingerprint: digest.finish(), + }) +} + +fn agent_coverage(store: &HoltAgentEventStore, options: &Options) -> BenchResult { + let coverage = store + .coverage(&options.agent_id, &options.source_file)? + .ok_or("agent coverage is missing")?; + let mut digest = SafeDigest::new(); + digest.field( + "min", + coverage.min_offset.map(|value| value as i64).unwrap_or(-1), + ); + digest.field( + "max", + coverage.max_offset.map(|value| value as i64).unwrap_or(-1), + ); + digest.field("count", coverage.row_count as i64); + Ok(QueryDigest { + rows: 1, + fingerprint: digest.finish(), + }) +} + +fn agent_session_rows(store: &HoltAgentEventStore, options: &Options) -> BenchResult { + let rows = store.stream_session_rows(SessionRowsRequest { + agent_id: options.agent_id.clone(), + limit: None, + })?; + Ok(agent_session_row_digest(rows)) +} + +fn agent_session_events( + store: &HoltAgentEventStore, + options: &Options, +) -> BenchResult { + let records = store.stream_session_events(SessionEventsRequest { + agent_id: options.agent_id.clone(), + event_types: LINGTAI_SESSION_EVENT_TYPES + .iter() + .map(|value| value.to_string()) + .collect(), + limit: None, + })?; + Ok(agent_event_digest(records)) +} + +fn agent_latest_by_type( + store: &HoltAgentEventStore, + options: &Options, + event_type: &str, + limit: usize, +) -> BenchResult { + let records = store.latest_events(LatestEventsRequest { + agent_id: options.agent_id.clone(), + event_type: event_type.to_owned(), + limit, + })?; + Ok(agent_event_digest(records)) +} + +fn agent_notification_events( + store: &HoltAgentEventStore, + options: &Options, + limit: usize, +) -> BenchResult { + let records = store.notification_events(NotificationEventsRequest { + agent_id: options.agent_id.clone(), + limit, + })?; + Ok(agent_event_digest(records)) +} + +fn agent_notification_pivots( + store: &HoltAgentEventStore, + options: &Options, +) -> BenchResult { + let mut events = store.notification_events(NotificationEventsRequest { + agent_id: options.agent_id.clone(), + limit: usize::MAX, + })?; + Ok(NotificationPivots { + latest: events.first().map(|record| record.id), + earliest: events.pop().map(|record| record.id), + }) +} + +fn agent_notification_by_id( + store: &HoltAgentEventStore, + options: &Options, + event_id: Option, +) -> BenchResult { + let Some(event_id) = event_id else { + return Ok(empty_digest()); + }; + let event = store.notification_event_by_id(NotificationEventByIdRequest { + agent_id: options.agent_id.clone(), + event_id, + })?; + Ok(agent_event_digest(event.into_iter().collect())) +} + +fn agent_notification_before( + store: &HoltAgentEventStore, + options: &Options, + pivot_event_id: Option, +) -> BenchResult { + let Some(pivot_event_id) = pivot_event_id else { + return Ok(empty_digest()); + }; + let event = store.notification_neighbor(NotificationNeighborRequest { + agent_id: options.agent_id.clone(), + pivot_event_id, + direction: NotificationNeighborDirection::Before, + })?; + Ok(agent_event_digest(event.into_iter().collect())) +} + +fn agent_notification_after( + store: &HoltAgentEventStore, + options: &Options, + pivot_event_id: Option, +) -> BenchResult { + let Some(pivot_event_id) = pivot_event_id else { + return Ok(empty_digest()); + }; + let event = store.notification_neighbor(NotificationNeighborRequest { + agent_id: options.agent_id.clone(), + pivot_event_id, + direction: NotificationNeighborDirection::After, + })?; + Ok(agent_event_digest(event.into_iter().collect())) +} + +fn agent_recent_times( + store: &HoltAgentEventStore, + options: &Options, + event_type: &str, + limit: usize, +) -> BenchResult { + let times = store.recent_times(RecentTimesRequest { + agent_id: options.agent_id.clone(), + event_type: event_type.to_owned(), + limit, + })?; + let mut digest = SafeDigest::new(); + for item in × { + digest.event(event_type, item.ts, -1); + } + Ok(QueryDigest { + rows: times.len(), + fingerprint: digest.finish(), + }) +} + +fn agent_molt_session_windows( + store: &HoltAgentEventStore, + options: &Options, +) -> BenchResult { + let windows = store.molt_session_windows(&options.agent_id)?; + let mut digest = SafeDigest::new(); + digest.field("ok", i64::from(windows.ok)); + digest.field( + "current_since", + windows + .current_since + .map(timestamp_micros_i64) + .unwrap_or(-1), + ); + digest.field( + "last_since", + windows.last_since.map(timestamp_micros_i64).unwrap_or(-1), + ); + digest.field( + "last_before", + windows.last_before.map(timestamp_micros_i64).unwrap_or(-1), + ); + Ok(QueryDigest { + rows: usize::from(windows.ok), + fingerprint: digest.finish(), + }) +} + +fn agent_tui_clear_completion( + store: &HoltAgentEventStore, + options: &Options, + source_offset: u64, +) -> BenchResult { + let completion = store.tui_clear_completion(TuiClearCompletionRequest { + agent_id: options.agent_id.clone(), + source_offset, + })?; + Ok(agent_event_digest(completion.event.into_iter().collect())) +} + +fn agent_error_events(store: &HoltAgentEventStore, options: &Options) -> BenchResult { + let records = store.error_events(ErrorEventsRequest { + agent_id: options.agent_id.clone(), + event_types: ERROR_EVENT_TYPES + .iter() + .map(|value| value.to_string()) + .collect(), + limit: usize::MAX, + })?; + Ok(agent_event_digest(records)) +} + +fn agent_tool_name_facets( + store: &HoltAgentEventStore, + options: &Options, + limit: usize, +) -> BenchResult { + let action_facets = store.tool_facets(ToolFacetRequest { + agent_id: options.agent_id.clone(), + limit: usize::MAX, + })?; + let mut counts = BTreeMap::::new(); + for facet in action_facets { + *counts.entry(facet.tool_name).or_default() += facet.count; + } + let mut facets = counts + .into_iter() + .map(|(tool_name, count)| ToolFacet { + tool_name, + action: None, + count, + }) + .collect::>(); + sort_facets(&mut facets); + facets.truncate(limit); + Ok(agent_facet_digest(&facets)) +} + +fn agent_tool_action_facets( + store: &HoltAgentEventStore, + options: &Options, + limit: usize, +) -> BenchResult { + let facets = store.tool_facets(ToolFacetRequest { + agent_id: options.agent_id.clone(), + limit, + })?; + Ok(agent_facet_digest(&facets)) +} + +fn agent_event_digest(records: Vec) -> QueryDigest { + let mut digest = SafeDigest::new(); + let rows = records.len(); + for record in records { + digest.event(&record.event_type, record.ts, record.source_offset as i64); + } + QueryDigest { + rows, + fingerprint: digest.finish(), + } +} + +fn agent_session_row_digest(rows: Vec) -> QueryDigest { + let mut digest = SafeDigest::new(); + let count = rows.len(); + for row in rows { + digest.event(&row.event_type, row.ts, -1); + } + QueryDigest { + rows: count, + fingerprint: digest.finish(), + } +} + +fn agent_facet_digest(facets: &[ToolFacet]) -> QueryDigest { + let mut digest = SafeDigest::new(); + for facet in facets { + digest.facet(&facet.tool_name, facet.action.as_deref(), facet.count); + } + QueryDigest { + rows: facets.len(), + fingerprint: digest.finish(), + } +} + +fn empty_digest() -> QueryDigest { + QueryDigest { + rows: 0, + fingerprint: SafeDigest::new().finish(), + } +} + +fn timed(iterations: usize, mut query: F) -> BenchResult +where + F: FnMut() -> BenchResult, +{ + let warmup = query()?; + let mut timings = Vec::with_capacity(iterations.max(1)); + let mut last = warmup.clone(); + for _ in 0..iterations.max(1) { + let started = Instant::now(); + last = query()?; + timings.push(elapsed_ms(started)); + } + timings.sort_by(|left, right| left.total_cmp(right)); + let total = timings.iter().sum::(); + let p50 = timings[timings.len() / 2]; + let p95 = timings[((timings.len() as f64 * 0.95).ceil() as usize).saturating_sub(1)] + .min(*timings.last().unwrap()); + Ok(QueryStats { + rows: last.rows, + p50_ms: round_ms(p50), + p95_ms: round_ms(p95), + avg_ms: round_ms(total / timings.len() as f64), + fingerprint: last.fingerprint, + }) +} + +fn compare_queries( + sqlite: &BTreeMap, + agent: &BTreeMap, +) -> BTreeMap { + let mut out = BTreeMap::new(); + for (name, sqlite_stats) in sqlite { + if let Some(agent_stats) = agent.get(name) { + out.insert( + name.clone(), + QueryComparison { + rows_match: sqlite_stats.rows == agent_stats.rows, + fingerprint_match: sqlite_stats.fingerprint == agent_stats.fingerprint, + }, + ); + } + } + out +} + +struct SafeDigest { + hasher: Sha256, +} + +impl SafeDigest { + fn new() -> Self { + Self { + hasher: Sha256::new(), + } + } + + fn field(&mut self, label: &str, value: i64) { + self.hasher.update(label.as_bytes()); + self.hasher.update(b"="); + self.hasher.update(value.to_string().as_bytes()); + self.hasher.update(b"\n"); + } + + fn event(&mut self, event_type: &str, ts: f64, source_offset: i64) { + self.hasher.update(event_type.as_bytes()); + self.hasher.update(b"|"); + self.hasher + .update(format!("{:.6}", if ts.is_finite() { ts } else { 0.0 }).as_bytes()); + self.hasher.update(b"|"); + self.hasher.update(source_offset.to_string().as_bytes()); + self.hasher.update(b"\n"); + } + + fn facet(&mut self, tool_name: &str, action: Option<&str>, count: u64) { + self.hasher.update(tool_name.as_bytes()); + self.hasher.update(b"|"); + self.hasher.update(action.unwrap_or("").as_bytes()); + self.hasher.update(b"|"); + self.hasher.update(count.to_string().as_bytes()); + self.hasher.update(b"\n"); + } + + fn finish(self) -> String { + let bytes = self.hasher.finalize(); + let mut out = String::with_capacity(bytes.len() * 2); + for byte in bytes { + out.push(nibble(byte >> 4)); + out.push(nibble(byte & 0x0f)); + } + out + } +} + +fn sort_facets(facets: &mut [ToolFacet]) { + facets.sort_by(|left, right| { + right + .count + .cmp(&left.count) + .then_with(|| left.tool_name.cmp(&right.tool_name)) + .then_with(|| left.action.cmp(&right.action)) + }); +} + +fn placeholders(count: usize) -> String { + std::iter::repeat_n("?", count) + .collect::>() + .join(",") +} + +fn sql_string_literal(value: &str) -> String { + format!("'{}'", value.replace('\'', "''")) +} + +fn elapsed_ms(started: Instant) -> f64 { + started.elapsed().as_secs_f64() * 1000.0 +} + +fn round_ms(value: f64) -> f64 { + (value * 1000.0).round() / 1000.0 +} + +fn timestamp_micros_i64(value: f64) -> i64 { + if !value.is_finite() || value <= 0.0 { + return 0; + } + let micros = value * 1_000_000.0; + if micros >= i64::MAX as f64 { + i64::MAX + } else { + micros as i64 + } +} + +fn dir_size(path: &Path) -> BenchResult { + let mut bytes = 0_u64; + if !path.exists() { + return Ok(0); + } + for entry in WalkDir::new(path) { + let entry = entry?; + if entry.file_type().is_file() { + bytes = bytes.saturating_add(entry.metadata()?.len()); + } + } + Ok(bytes) +} + +fn required_path(args: &mut VecDeque, flag: &str) -> BenchResult { + required_value(args, flag).map(PathBuf::from) +} + +fn required_value(args: &mut VecDeque, flag: &str) -> BenchResult { + option_value(args, flag)?.ok_or_else(|| format!("missing required option {flag}").into()) +} + +fn option_value(args: &mut VecDeque, flag: &str) -> BenchResult> { + let Some(index) = args.iter().position(|arg| arg == flag) else { + return Ok(None); + }; + args.remove(index); + args.remove(index) + .map(Some) + .ok_or_else(|| format!("option {flag} requires a value").into()) +} + +fn take_flag(args: &mut VecDeque, flag: &str) -> bool { + let Some(index) = args.iter().position(|arg| arg == flag) else { + return false; + }; + args.remove(index); + true +} + +fn nibble(value: u8) -> char { + match value { + 0..=9 => (b'0' + value) as char, + 10..=15 => (b'a' + (value - 10)) as char, + _ => unreachable!("nibble is four bits"), + } +} + +fn usage() -> &'static str { + "usage: lingtai-index-bench --events-jsonl PATH --sqlite PATH --agent-index PATH [--projected-sqlite PATH] [--rebuilt-sqlite PATH] [--source-file logs/events.jsonl] [--agent-id default] [--iterations 20] [--reset]" +} diff --git a/crates/nokv-agent/Cargo.toml b/crates/nokv-agent/Cargo.toml index 794c94ce5..ed7142bd5 100644 --- a/crates/nokv-agent/Cargo.toml +++ b/crates/nokv-agent/Cargo.toml @@ -7,10 +7,16 @@ license.workspace = true authors.workspace = true repository.workspace = true homepage.workspace = true -description = "NoKV agent tool surface: transport-free dispatch over the metadata namespace verbs." +description = "NoKV agent surfaces: namespace tools and LingTai event indexes." [dependencies] +holt.workspace = true nokv-meta.workspace = true nokv-object.workspace = true nokv-types.workspace = true +serde.workspace = true serde_json.workspace = true +sha2.workspace = true + +[dev-dependencies] +tempfile.workspace = true diff --git a/crates/nokv-agent/src/bin/nokv-agent.rs b/crates/nokv-agent/src/bin/nokv-agent.rs new file mode 100644 index 000000000..8b7997a5d --- /dev/null +++ b/crates/nokv-agent/src/bin/nokv-agent.rs @@ -0,0 +1,465 @@ +use std::collections::VecDeque; +use std::fs::File; +use std::io::BufReader; +use std::path::PathBuf; + +use nokv_agent::event::{ + ingest_jsonl_reader, AgentEventError, AgentEventResult, AgentEventStore, + CompletionAfterRequest, ErrorEventsRequest, HoltAgentEventStore, JsonlIngestOptions, + LatestEventsRequest, NotificationEventByIdRequest, NotificationEventsRequest, + NotificationLifecycleRequest, NotificationNeighborDirection, NotificationNeighborRequest, + RecentTimesRequest, SessionEventsRequest, SessionRowsRequest, ToolFacetRequest, + ToolTraceRequest, TuiClearCompletionRequest, LINGTAI_SESSION_EVENT_TYPES, +}; +use serde_json::json; + +const DEFAULT_ERROR_EVENT_TYPES: &[&str] = &[ + "aed_attempt", + "aed_exhausted", + "aed_over_window_detected", + "aed_timeout", + "aed_transient_retry", + "aed_transient_exhausted", + "daemon_error", + "refresh_failed_permanent", + "refresh_init_error", + "tc_wake_error", + "tool_call_dispatch_failed", + "tool_call_validation_failed", + "daemon_tool_call_validation_failed", +]; + +fn main() { + if let Err(err) = run() { + eprintln!("{err}"); + std::process::exit(2); + } +} + +fn run() -> AgentEventResult<()> { + let mut args = std::env::args().skip(1).collect::>(); + let Some(scope) = args.pop_front() else { + return usage(); + }; + if scope != "lingtai" { + return usage(); + } + let Some(command) = args.pop_front() else { + return usage(); + }; + match command.as_str() { + "ingest" => cmd_ingest(args), + "coverage" => cmd_coverage(args), + "latest" => cmd_latest(args), + "session" => cmd_session(args), + "session-rows" => cmd_session_rows(args), + "recent" => cmd_recent(args), + "molt-windows" => cmd_molt_windows(args), + "errors" => cmd_errors(args), + "completion-after" => cmd_completion_after(args), + "clear-completion" => cmd_clear_completion(args), + "notification-blocks" => cmd_latest_fixed_type(args, "notification_pair_injected"), + "notification-block-snapshots" => { + cmd_latest_fixed_type(args, "notification_block_injected") + } + "notification-events" => cmd_notification_events(args), + "notification-by-id" => cmd_notification_by_id(args), + "notification-before" => { + cmd_notification_neighbor(args, NotificationNeighborDirection::Before) + } + "notification-after" => { + cmd_notification_neighbor(args, NotificationNeighborDirection::After) + } + "notifications" => cmd_notifications(args), + "facets" => cmd_facets(args), + "trace" => cmd_trace(args), + _ => usage(), + } +} + +fn cmd_ingest(mut args: VecDeque) -> AgentEventResult<()> { + let store_path = required_path(&mut args, "--store")?; + let events_path = required_path(&mut args, "--events")?; + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let source_file = option_value(&mut args, "--source-file")? + .unwrap_or_else(|| events_path.display().to_string()); + reject_extra(args)?; + + let store = HoltAgentEventStore::open_file(store_path)?; + let file = File::open(&events_path)?; + let file_size = file.metadata()?.len(); + let report = ingest_jsonl_reader( + &store, + JsonlIngestOptions { + agent_id, + source_file, + file_size, + }, + BufReader::new(file), + )?; + print_json(json!({"report": report})) +} + +fn cmd_coverage(mut args: VecDeque) -> AgentEventResult<()> { + let store_path = required_path(&mut args, "--store")?; + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let source_file = required_value(&mut args, "--source-file")?; + reject_extra(args)?; + + let store = HoltAgentEventStore::open_file(store_path)?; + print_json(json!({ + "coverage": store.coverage(&agent_id, &source_file)? + })) +} + +fn cmd_latest(mut args: VecDeque) -> AgentEventResult<()> { + let store_path = required_path(&mut args, "--store")?; + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let event_type = required_value(&mut args, "--type")?; + let limit = option_value(&mut args, "--limit")? + .map(|value| parse_usize("--limit", &value)) + .transpose()? + .unwrap_or(10); + reject_extra(args)?; + + let store = HoltAgentEventStore::open_file(store_path)?; + print_json(json!({ + "events": store.latest_events(LatestEventsRequest { + agent_id, + event_type, + limit, + })? + })) +} + +fn cmd_latest_fixed_type( + mut args: VecDeque, + event_type: &'static str, +) -> AgentEventResult<()> { + let store_path = required_path(&mut args, "--store")?; + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let limit = option_value(&mut args, "--limit")? + .map(|value| parse_usize("--limit", &value)) + .transpose()? + .unwrap_or(10); + reject_extra(args)?; + + let store = HoltAgentEventStore::open_file(store_path)?; + print_json(json!({ + "events": store.latest_events(LatestEventsRequest { + agent_id, + event_type: event_type.to_owned(), + limit, + })? + })) +} + +fn cmd_session(mut args: VecDeque) -> AgentEventResult<()> { + let store_path = required_path(&mut args, "--store")?; + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let event_types = option_value(&mut args, "--types")? + .map(|value| parse_csv(&value)) + .unwrap_or_else(|| { + LINGTAI_SESSION_EVENT_TYPES + .iter() + .map(|value| value.to_string()) + .collect() + }); + let limit = option_value(&mut args, "--limit")? + .map(|value| parse_usize("--limit", &value)) + .transpose()?; + reject_extra(args)?; + + let store = HoltAgentEventStore::open_file(store_path)?; + print_json(json!({ + "events": store.stream_session_events(SessionEventsRequest { + agent_id, + event_types, + limit, + })? + })) +} + +fn cmd_session_rows(mut args: VecDeque) -> AgentEventResult<()> { + let store_path = required_path(&mut args, "--store")?; + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let limit = option_value(&mut args, "--limit")? + .map(|value| parse_usize("--limit", &value)) + .transpose()?; + reject_extra(args)?; + + let store = HoltAgentEventStore::open_file(store_path)?; + print_json(json!({ + "rows": store.stream_session_rows(SessionRowsRequest { + agent_id, + limit, + })? + })) +} + +fn cmd_recent(mut args: VecDeque) -> AgentEventResult<()> { + let store_path = required_path(&mut args, "--store")?; + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let event_type = required_value(&mut args, "--type")?; + let limit = option_value(&mut args, "--limit")? + .map(|value| parse_usize("--limit", &value)) + .transpose()? + .unwrap_or(10); + reject_extra(args)?; + + let store = HoltAgentEventStore::open_file(store_path)?; + print_json(json!({ + "times": store.recent_times(RecentTimesRequest { + agent_id, + event_type, + limit, + })? + })) +} + +fn cmd_molt_windows(mut args: VecDeque) -> AgentEventResult<()> { + let store_path = required_path(&mut args, "--store")?; + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + reject_extra(args)?; + + let store = HoltAgentEventStore::open_file(store_path)?; + print_json(json!({ + "windows": store.molt_session_windows(&agent_id)? + })) +} + +fn cmd_errors(mut args: VecDeque) -> AgentEventResult<()> { + let store_path = required_path(&mut args, "--store")?; + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let event_types = option_value(&mut args, "--types")? + .map(|value| parse_csv(&value)) + .unwrap_or_else(|| { + DEFAULT_ERROR_EVENT_TYPES + .iter() + .map(|value| value.to_string()) + .collect() + }); + let limit = option_value(&mut args, "--limit")? + .map(|value| parse_usize("--limit", &value)) + .transpose()? + .unwrap_or(50); + reject_extra(args)?; + + let store = HoltAgentEventStore::open_file(store_path)?; + print_json(json!({ + "events": store.error_events(ErrorEventsRequest { + agent_id, + event_types, + limit, + })? + })) +} + +fn cmd_completion_after(mut args: VecDeque) -> AgentEventResult<()> { + let store_path = required_path(&mut args, "--store")?; + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let event_type = + option_value(&mut args, "--type")?.unwrap_or_else(|| "clear_received".to_owned()); + let source_offset = parse_u64( + "--source-offset", + &required_value(&mut args, "--source-offset")?, + )?; + reject_extra(args)?; + + let store = HoltAgentEventStore::open_file(store_path)?; + print_json(json!({ + "completion": store.completion_after(CompletionAfterRequest { + agent_id, + event_type, + source_offset, + })? + })) +} + +fn cmd_clear_completion(mut args: VecDeque) -> AgentEventResult<()> { + let store_path = required_path(&mut args, "--store")?; + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let source_offset = parse_u64( + "--source-offset", + &required_value(&mut args, "--source-offset")?, + )?; + reject_extra(args)?; + + let store = HoltAgentEventStore::open_file(store_path)?; + print_json(json!({ + "completion": store.tui_clear_completion(TuiClearCompletionRequest { + agent_id, + source_offset, + })? + })) +} + +fn cmd_notification_events(mut args: VecDeque) -> AgentEventResult<()> { + let store_path = required_path(&mut args, "--store")?; + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let limit = option_value(&mut args, "--limit")? + .map(|value| parse_usize("--limit", &value)) + .transpose()? + .unwrap_or(50); + reject_extra(args)?; + + let store = HoltAgentEventStore::open_file(store_path)?; + print_json(json!({ + "events": store.notification_events(NotificationEventsRequest { + agent_id, + limit, + })? + })) +} + +fn cmd_notification_by_id(mut args: VecDeque) -> AgentEventResult<()> { + let store_path = required_path(&mut args, "--store")?; + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let event_id = parse_u64("--event-id", &required_value(&mut args, "--event-id")?)?; + reject_extra(args)?; + + let store = HoltAgentEventStore::open_file(store_path)?; + print_json(json!({ + "event": store.notification_event_by_id(NotificationEventByIdRequest { + agent_id, + event_id, + })? + })) +} + +fn cmd_notification_neighbor( + mut args: VecDeque, + direction: NotificationNeighborDirection, +) -> AgentEventResult<()> { + let store_path = required_path(&mut args, "--store")?; + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let pivot_event_id = parse_u64("--event-id", &required_value(&mut args, "--event-id")?)?; + reject_extra(args)?; + + let store = HoltAgentEventStore::open_file(store_path)?; + print_json(json!({ + "event": store.notification_neighbor(NotificationNeighborRequest { + agent_id, + pivot_event_id, + direction, + })? + })) +} + +fn cmd_notifications(mut args: VecDeque) -> AgentEventResult<()> { + let store_path = required_path(&mut args, "--store")?; + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let ref_id = option_value(&mut args, "--ref-id")?; + let event_id = option_value(&mut args, "--event-id")?; + let call_id = option_value(&mut args, "--call-id")?; + let channel = option_value(&mut args, "--channel")?; + let limit = option_value(&mut args, "--limit")? + .map(|value| parse_usize("--limit", &value)) + .transpose()? + .unwrap_or(50); + reject_extra(args)?; + + let store = HoltAgentEventStore::open_file(store_path)?; + print_json(json!({ + "events": store.notification_lifecycle(NotificationLifecycleRequest { + agent_id, + ref_id, + event_id, + call_id, + channel, + limit, + })? + })) +} + +fn cmd_facets(mut args: VecDeque) -> AgentEventResult<()> { + let store_path = required_path(&mut args, "--store")?; + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let limit = option_value(&mut args, "--limit")? + .map(|value| parse_usize("--limit", &value)) + .transpose()? + .unwrap_or(20); + reject_extra(args)?; + + let store = HoltAgentEventStore::open_file(store_path)?; + print_json(json!({ + "facets": store.tool_facets(ToolFacetRequest { agent_id, limit })? + })) +} + +fn cmd_trace(mut args: VecDeque) -> AgentEventResult<()> { + let store_path = required_path(&mut args, "--store")?; + let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let tool_call_id = required_value(&mut args, "--tool-call-id")?; + reject_extra(args)?; + + let store = HoltAgentEventStore::open_file(store_path)?; + print_json(json!({ + "events": store.tool_trace(ToolTraceRequest { + agent_id, + tool_call_id, + })? + })) +} + +fn required_path(args: &mut VecDeque, flag: &str) -> AgentEventResult { + required_value(args, flag).map(PathBuf::from) +} + +fn required_value(args: &mut VecDeque, flag: &str) -> AgentEventResult { + option_value(args, flag)? + .ok_or_else(|| AgentEventError::InvalidArgument(format!("missing required option {flag}"))) +} + +fn option_value(args: &mut VecDeque, flag: &str) -> AgentEventResult> { + let Some(index) = args.iter().position(|arg| arg == flag) else { + return Ok(None); + }; + args.remove(index); + args.remove(index) + .map(Some) + .ok_or_else(|| AgentEventError::InvalidArgument(format!("option {flag} requires a value"))) +} + +fn reject_extra(args: VecDeque) -> AgentEventResult<()> { + if args.is_empty() { + Ok(()) + } else { + Err(AgentEventError::InvalidArgument(format!( + "unexpected argument {}", + args.front().unwrap() + ))) + } +} + +fn parse_usize(flag: &str, value: &str) -> AgentEventResult { + value.parse::().map_err(|_| { + AgentEventError::InvalidArgument(format!("option {flag} must be a positive integer")) + }) +} + +fn parse_u64(flag: &str, value: &str) -> AgentEventResult { + value.parse::().map_err(|_| { + AgentEventError::InvalidArgument(format!("option {flag} must be an unsigned integer")) + }) +} + +fn parse_csv(value: &str) -> Vec { + value + .split(',') + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + .collect() +} + +fn print_json(value: serde_json::Value) -> AgentEventResult<()> { + println!("{}", serde_json::to_string_pretty(&value)?); + Ok(()) +} + +fn usage() -> AgentEventResult<()> { + Err(AgentEventError::InvalidArgument( + "usage: nokv-agent lingtai ingest|coverage|latest|session|session-rows|recent|molt-windows|errors|completion-after|clear-completion|notification-blocks|notification-block-snapshots|notification-events|notification-by-id|notification-before|notification-after|notifications|facets|trace".to_owned(), + )) +} diff --git a/crates/nokv-agent/src/event/codec.rs b/crates/nokv-agent/src/event/codec.rs new file mode 100644 index 000000000..83068642c --- /dev/null +++ b/crates/nokv-agent/src/event/codec.rs @@ -0,0 +1,63 @@ +use super::types::{ + AgentEventError, AgentEventResult, EventRecord, IndexCoverage, SessionEventRow, ToolFacet, +}; + +const EVENT_RECORD_V1: u8 = 1; +const COVERAGE_V1: u8 = 1; +const TOOL_FACET_V1: u8 = 1; +const SESSION_EVENT_ROW_V1: u8 = 1; + +pub fn encode_event(record: &EventRecord) -> AgentEventResult> { + encode_versioned(EVENT_RECORD_V1, record) +} + +pub fn decode_event(bytes: &[u8]) -> AgentEventResult { + decode_versioned(EVENT_RECORD_V1, bytes, "event record") +} + +pub fn encode_coverage(coverage: &IndexCoverage) -> AgentEventResult> { + encode_versioned(COVERAGE_V1, coverage) +} + +pub fn decode_coverage(bytes: &[u8]) -> AgentEventResult { + decode_versioned(COVERAGE_V1, bytes, "coverage") +} + +pub fn encode_tool_facet(facet: &ToolFacet) -> AgentEventResult> { + encode_versioned(TOOL_FACET_V1, facet) +} + +pub fn decode_tool_facet(bytes: &[u8]) -> AgentEventResult { + decode_versioned(TOOL_FACET_V1, bytes, "tool facet") +} + +pub fn encode_session_event_row(row: &SessionEventRow) -> AgentEventResult> { + encode_versioned(SESSION_EVENT_ROW_V1, row) +} + +pub fn decode_session_event_row(bytes: &[u8]) -> AgentEventResult { + decode_versioned(SESSION_EVENT_ROW_V1, bytes, "session event row") +} + +fn encode_versioned(version: u8, value: &T) -> AgentEventResult> { + let mut out = vec![version]; + let mut json = serde_json::to_vec(value)?; + out.append(&mut json); + Ok(out) +} + +fn decode_versioned( + expected: u8, + bytes: &[u8], + label: &str, +) -> AgentEventResult { + let Some((&version, payload)) = bytes.split_first() else { + return Err(AgentEventError::Store(format!("{label} value is empty"))); + }; + if version != expected { + return Err(AgentEventError::Store(format!( + "{label} value version {version} is not supported" + ))); + } + Ok(serde_json::from_slice(payload)?) +} diff --git a/crates/nokv-agent/src/event/holt.rs b/crates/nokv-agent/src/event/holt.rs new file mode 100644 index 000000000..d81b49303 --- /dev/null +++ b/crates/nokv-agent/src/event/holt.rs @@ -0,0 +1,859 @@ +use std::cmp::Reverse; +use std::collections::{BTreeMap, BTreeSet}; +use std::path::Path; + +use holt::{DBAtomicBatch, RangeEntry, Tree, TreeConfig, DB}; + +use super::codec::{ + decode_coverage, decode_event, decode_session_event_row, decode_tool_facet, encode_coverage, + encode_event, encode_tool_facet, +}; +use super::ingest::empty_coverage; +use super::key::{ + coverage_key, event_key, id_from_index_value, id_value, notification_id_key, + notification_id_prefix, notification_key, notification_next_key, notification_prefix, + notification_prev_key, notification_rev_key, notification_rev_prefix, notification_tail_key, + session_key, session_prefix, source_file_hash, source_key, tool_action_facet_key, + tool_action_facet_prefix, tool_name_facet_key, trace_key, trace_prefix, tui_clear_rev_key, + tui_clear_rev_prefix, type_id_key, type_id_prefix, type_ts_key, type_ts_prefix, TREE_COVERAGE, + TREE_EVENTS, TREE_INDEX, +}; +use super::store::AgentEventStore; +use super::types::{ + AgentEventError, AgentEventResult, CompletionAfter, CompletionAfterRequest, ErrorEventsRequest, + EventRecord, EventTime, IndexCoverage, IngestReport, LatestEventsRequest, MoltSessionWindows, + NewEventRecord, NotificationEventByIdRequest, NotificationEventsRequest, + NotificationLifecycleRequest, NotificationNeighborDirection, NotificationNeighborRequest, + RecentTimesRequest, SessionEventRow, SessionEventsRequest, SessionRowsRequest, ToolFacet, + ToolFacetRequest, ToolTraceRequest, TuiClearCompletion, TuiClearCompletionRequest, + LINGTAI_SESSION_EVENT_TYPES, +}; + +pub struct HoltAgentEventStore { + db: DB, + events: Tree, + index: Tree, + coverage: Tree, +} + +impl HoltAgentEventStore { + pub fn open_memory() -> AgentEventResult { + Self::open(TreeConfig::memory()) + } + + pub fn open_file(path: impl AsRef) -> AgentEventResult { + Self::open(TreeConfig::new(path.as_ref())) + } + + pub fn open(config: TreeConfig) -> AgentEventResult { + let db = DB::open(config).map_err(to_store_error)?; + let events = db + .open_or_create_tree(TREE_EVENTS) + .map_err(to_store_error)?; + let index = db.open_or_create_tree(TREE_INDEX).map_err(to_store_error)?; + let coverage = db + .open_or_create_tree(TREE_COVERAGE) + .map_err(to_store_error)?; + Ok(Self { + db, + events, + index, + coverage, + }) + } + + fn event_by_id(&self, agent_id: &str, event_id: u64) -> AgentEventResult> { + let key = event_key(agent_id, event_id); + self.events + .get(&key) + .map_err(to_store_error)? + .map(|bytes| decode_event(&bytes)) + .transpose()? + .map(Some) + // Large file-backed stores can expose a key through range scan + // before the point path observes it; keep secondary-index lookups exact. + .map_or_else(|| self.event_by_id_scan_fallback(&key), Ok) + } + + fn event_by_id_scan_fallback(&self, key: &[u8]) -> AgentEventResult> { + for entry in self.events.scan(key) { + let entry = entry.map_err(to_store_error)?; + let RangeEntry::Key { + key: found, value, .. + } = entry + else { + continue; + }; + if found == key { + return decode_event(&value).map(Some); + } + break; + } + Ok(None) + } + + fn scan_event_ids(&self, prefix: &[u8], limit: usize) -> AgentEventResult> { + let mut out = Vec::new(); + for entry in self.index.scan(prefix) { + let entry = entry.map_err(to_store_error)?; + if let RangeEntry::Key { value, .. } = entry { + if let Some(id) = id_from_index_value(&value) { + out.push(id); + if out.len() == limit { + break; + } + } + } + } + Ok(out) + } + + fn events_by_index_prefix( + &self, + agent_id: &str, + prefix: &[u8], + limit: usize, + ) -> AgentEventResult> { + self.scan_event_ids(prefix, limit)? + .into_iter() + .filter_map(|id| self.event_by_id(agent_id, id).transpose()) + .collect() + } + + fn first_tui_clear_completion( + &self, + agent_id: &str, + source_offset: u64, + ) -> AgentEventResult> { + let prefix = tui_clear_rev_prefix(agent_id); + for entry in self.index.scan(&prefix) { + let entry = entry.map_err(to_store_error)?; + let RangeEntry::Key { value, .. } = entry else { + continue; + }; + let Some(id) = id_from_index_value(&value) else { + continue; + }; + let Some(record) = self.event_by_id(agent_id, id)? else { + continue; + }; + if record.source_offset >= source_offset { + return Ok(Some(record)); + } + } + Ok(None) + } + + fn first_event_after_index_key( + &self, + agent_id: &str, + prefix: &[u8], + start_after: &[u8], + ) -> AgentEventResult> { + for entry in self.index.scan(prefix).start_after(start_after) { + let entry = entry.map_err(to_store_error)?; + let RangeEntry::Key { value, .. } = entry else { + continue; + }; + let Some(id) = id_from_index_value(&value) else { + continue; + }; + if let Some(record) = self.event_by_id(agent_id, id)? { + return Ok(Some(record)); + } + } + Ok(None) + } + + fn facet_by_key(&self, key: &[u8], fallback: ToolFacet) -> AgentEventResult { + self.index + .get(key) + .map_err(to_store_error)? + .map(|bytes| decode_tool_facet(&bytes)) + .transpose() + .map(|facet| facet.unwrap_or(fallback)) + } + + fn id_by_index_key(&self, key: &[u8]) -> AgentEventResult> { + if let Some(id) = self + .index + .get(key) + .map_err(to_store_error)? + .and_then(|bytes| id_from_index_value(&bytes)) + { + return Ok(Some(id)); + } + for entry in self.index.scan(key) { + let entry = entry.map_err(to_store_error)?; + let RangeEntry::Key { + key: found, value, .. + } = entry + else { + continue; + }; + if found == key { + return Ok(id_from_index_value(&value)); + } + break; + } + Ok(None) + } + + fn scan_tool_facets(&self, prefix: &[u8], limit: usize) -> AgentEventResult> { + let mut out = Vec::new(); + for entry in self.index.scan(prefix) { + let entry = entry.map_err(to_store_error)?; + if let RangeEntry::Key { value, .. } = entry { + out.push(decode_tool_facet(&value)?); + } + } + sort_tool_facets(&mut out); + out.truncate(limit); + Ok(out) + } + + fn scan_session_rows( + &self, + agent_id: &str, + prefix: &[u8], + limit: Option, + ) -> AgentEventResult> { + if limit == Some(0) { + return Ok(Vec::new()); + } + let mut out = Vec::new(); + for entry in self.index.scan(prefix) { + let entry = entry.map_err(to_store_error)?; + if let RangeEntry::Key { value, .. } = entry { + if let Some(id) = id_from_index_value(&value) { + if let Some(record) = self.event_by_id(agent_id, id)? { + out.push(session_row_from_event(record)); + } + } else { + out.push(decode_session_event_row(&value)?); + } + if limit.is_some_and(|limit| out.len() == limit) { + break; + } + } + } + Ok(out) + } + + fn materialize_tool_facet_updates( + &self, + materialized: &[(Vec, EventRecord)], + ) -> AgentEventResult, ToolFacet)>> { + let mut increments = BTreeMap::, ToolFacet>::new(); + for (_, record) in materialized { + if record.event_type != "tool_call" { + continue; + } + let Some(tool_name) = record.projection.tool_name.as_deref() else { + continue; + }; + increment_facet( + &mut increments, + tool_name_facet_key(&record.agent_id, tool_name), + tool_name, + None, + ); + increment_facet( + &mut increments, + tool_action_facet_key( + &record.agent_id, + tool_name, + record.projection.tool_action.as_deref(), + ), + tool_name, + record.projection.tool_action.as_deref(), + ); + } + + let mut out = Vec::with_capacity(increments.len()); + for (key, increment) in increments { + let mut facet = self.facet_by_key( + &key, + ToolFacet { + tool_name: increment.tool_name, + action: increment.action, + count: 0, + }, + )?; + facet.count = facet.count.saturating_add(increment.count); + out.push((key, facet)); + } + Ok(out) + } + + fn materialize_notification_neighbor_updates( + &self, + materialized: &[(Vec, EventRecord)], + ) -> AgentEventResult { + let mut by_agent = BTreeMap::>::new(); + for (_, record) in materialized { + if is_notification_event_type(&record.event_type) { + by_agent + .entry(record.agent_id.clone()) + .or_default() + .push(record.id); + } + } + + let mut updates = NotificationNeighborUpdates::default(); + for (agent_id, mut ids) in by_agent { + ids.sort_unstable(); + ids.dedup(); + let mut previous = self.id_by_index_key(¬ification_tail_key(&agent_id))?; + for id in ids { + if let Some(previous_id) = previous { + updates + .next + .push((notification_next_key(&agent_id, previous_id), id_value(id))); + updates + .prev + .push((notification_prev_key(&agent_id, id), id_value(previous_id))); + } + previous = Some(id); + } + if let Some(last_id) = previous { + updates + .tail + .push((notification_tail_key(&agent_id), id_value(last_id))); + } + } + Ok(updates) + } +} + +#[derive(Default)] +struct NotificationNeighborUpdates { + prev: Vec<(Vec, [u8; 8])>, + next: Vec<(Vec, [u8; 8])>, + tail: Vec<(Vec, [u8; 8])>, +} + +impl AgentEventStore for HoltAgentEventStore { + fn ingest_batch( + &self, + records: Vec, + file_size: u64, + ) -> AgentEventResult { + let mut batch_seen = BTreeSet::new(); + let mut accepted_records = Vec::new(); + let mut duplicates = 0_u64; + let mut last_source = None; + + for record in records { + last_source = Some((record.agent_id.clone(), record.source_file.clone())); + let source = source_key(&record.agent_id, &record.source_file, record.source_offset); + if !batch_seen.insert(source.clone()) + || self.index.get(&source).map_err(to_store_error)?.is_some() + { + duplicates = duplicates.saturating_add(1); + continue; + } + accepted_records.push((source, record)); + } + + let mut coverage_updates = BTreeMap::<(String, String), IndexCoverage>::new(); + let mut materialized = Vec::with_capacity(accepted_records.len()); + for (source_key, new) in accepted_records { + let id = event_id_for_source(&new.source_file, new.source_offset); + let record = EventRecord { + id, + agent_id: new.agent_id, + source_file: new.source_file, + source_offset: new.source_offset, + source_line: new.source_line, + ts: new.ts, + event_type: new.event_type, + fields_json: new.fields_json, + projection: new.projection, + }; + let key = (record.agent_id.clone(), record.source_file.clone()); + let coverage = coverage_updates.entry(key.clone()).or_insert_with(|| { + self.coverage(&key.0, &key.1) + .ok() + .flatten() + .unwrap_or_else(|| empty_coverage(&key.0, &key.1, file_size)) + }); + coverage.file_size = coverage.file_size.max(file_size); + coverage.min_offset = + Some(coverage.min_offset.map_or(record.source_offset, |current| { + current.min(record.source_offset) + })); + coverage.max_offset = + Some(coverage.max_offset.map_or(record.source_offset, |current| { + current.max(record.source_offset) + })); + coverage.row_count = coverage.row_count.saturating_add(1); + materialized.push((source_key, record)); + } + let facet_values = self.materialize_tool_facet_updates(&materialized)?; + let notification_neighbor_updates = + self.materialize_notification_neighbor_updates(&materialized)?; + + if !materialized.is_empty() { + self.db + .atomic(|batch| { + for (source, record) in &materialized { + write_record_batch(batch, source, record); + } + for coverage in coverage_updates.values() { + batch.put( + TREE_COVERAGE, + &coverage_key(&coverage.agent_id, &coverage.source_file), + &encode_coverage(coverage).expect("coverage encodes"), + ); + } + for (key, facet) in &facet_values { + batch.put( + TREE_INDEX, + key, + &encode_tool_facet(facet).expect("tool facet encodes"), + ); + } + for (key, id) in ¬ification_neighbor_updates.prev { + batch.put(TREE_INDEX, key, id); + } + for (key, id) in ¬ification_neighbor_updates.next { + batch.put(TREE_INDEX, key, id); + } + for (key, id) in ¬ification_neighbor_updates.tail { + batch.put(TREE_INDEX, key, id); + } + }) + .map_err(to_store_error)?; + } + + let coverage = coverage_updates + .values() + .next() + .cloned() + .or_else(|| { + last_source.as_ref().and_then(|(agent_id, source_file)| { + self.coverage(agent_id, source_file).ok().flatten() + }) + }) + .unwrap_or_default(); + Ok(IngestReport { + accepted: materialized.len() as u64, + duplicates, + parse_errors: 0, + partial_lines: 0, + coverage, + }) + } + + fn coverage( + &self, + agent_id: &str, + source_file: &str, + ) -> AgentEventResult> { + self.coverage + .get(&coverage_key(agent_id, source_file)) + .map_err(to_store_error)? + .map(|bytes| decode_coverage(&bytes)) + .transpose() + } + + fn latest_events(&self, request: LatestEventsRequest) -> AgentEventResult> { + let prefix = type_id_prefix(&request.agent_id, &request.event_type); + self.events_by_index_prefix(&request.agent_id, &prefix, request.limit) + } + + fn stream_session_events( + &self, + request: SessionEventsRequest, + ) -> AgentEventResult> { + let mut ids = Vec::new(); + for event_type in request.event_types { + ids.extend( + self.scan_event_ids(&type_id_prefix(&request.agent_id, &event_type), usize::MAX)?, + ); + } + ids.sort_unstable(); + ids.dedup(); + if let Some(limit) = request.limit { + ids.truncate(limit); + } + ids.into_iter() + .filter_map(|id| self.event_by_id(&request.agent_id, id).transpose()) + .collect() + } + + fn stream_session_rows( + &self, + request: SessionRowsRequest, + ) -> AgentEventResult> { + let prefix = session_prefix(&request.agent_id); + self.scan_session_rows(&request.agent_id, &prefix, request.limit) + } + + fn tool_facets(&self, request: ToolFacetRequest) -> AgentEventResult> { + self.scan_tool_facets(&tool_action_facet_prefix(&request.agent_id), request.limit) + } + + fn tool_trace(&self, request: ToolTraceRequest) -> AgentEventResult> { + let prefix = trace_prefix(&request.agent_id, &request.tool_call_id); + let mut out = self.events_by_index_prefix(&request.agent_id, &prefix, usize::MAX)?; + out.sort_by_key(|record| record.id); + Ok(out) + } + + fn recent_times(&self, request: RecentTimesRequest) -> AgentEventResult> { + let prefix = type_ts_prefix(&request.agent_id, &request.event_type); + let records = self.events_by_index_prefix(&request.agent_id, &prefix, request.limit)?; + Ok(records + .into_iter() + .map(|record| EventTime { + id: record.id, + ts: record.ts, + }) + .collect()) + } + + fn molt_session_windows(&self, agent_id: &str) -> AgentEventResult { + let times = self.recent_times(RecentTimesRequest { + agent_id: agent_id.to_owned(), + event_type: "psyche_molt".to_owned(), + limit: 2, + })?; + let current_since = times.first().map(|event| event.ts); + let last_since = times.get(1).map(|event| event.ts); + let last_before = last_since.and(current_since); + Ok(MoltSessionWindows { + ok: true, + current_since, + last_since, + last_before, + }) + } + + fn error_events(&self, request: ErrorEventsRequest) -> AgentEventResult> { + if request.event_types.is_empty() || request.limit == 0 { + return Ok(Vec::new()); + } + let mut out = Vec::new(); + for event_type in request.event_types { + let prefix = type_id_prefix(&request.agent_id, &event_type); + out.extend(self.events_by_index_prefix(&request.agent_id, &prefix, request.limit)?); + } + out.sort_by_key(|record| Reverse(record.id)); + out.truncate(request.limit); + Ok(out) + } + + fn completion_after( + &self, + request: CompletionAfterRequest, + ) -> AgentEventResult { + let prefix = type_id_prefix(&request.agent_id, &request.event_type); + let mut matches = self + .events_by_index_prefix(&request.agent_id, &prefix, usize::MAX)? + .into_iter() + .filter(|record| record.source_offset > request.source_offset) + .collect::>(); + matches.sort_by_key(|record| (record.source_offset, record.id)); + Ok(CompletionAfter { + found: !matches.is_empty(), + event: matches.into_iter().next(), + }) + } + + fn tui_clear_completion( + &self, + request: TuiClearCompletionRequest, + ) -> AgentEventResult { + if let Some(event) = + self.first_tui_clear_completion(&request.agent_id, request.source_offset)? + { + return Ok(TuiClearCompletion { + found: true, + event: Some(event), + }); + } + + let mut matches = Vec::new(); + for event_type in ["psyche_molt", "clear_received"] { + let prefix = type_id_prefix(&request.agent_id, event_type); + matches.extend( + self.events_by_index_prefix(&request.agent_id, &prefix, usize::MAX)? + .into_iter() + .filter(|record| { + record.source_offset >= request.source_offset + && record + .fields_json + .get("source") + .and_then(|value| value.as_str()) + == Some("tui") + }), + ); + } + matches.sort_by_key(|record| Reverse(record.id)); + Ok(TuiClearCompletion { + found: !matches.is_empty(), + event: matches.into_iter().next(), + }) + } + + fn notification_events( + &self, + request: NotificationEventsRequest, + ) -> AgentEventResult> { + if request.limit == 0 { + return Ok(Vec::new()); + } + let prefix = notification_rev_prefix(&request.agent_id); + self.events_by_index_prefix(&request.agent_id, &prefix, request.limit) + } + + fn notification_event_by_id( + &self, + request: NotificationEventByIdRequest, + ) -> AgentEventResult> { + Ok(self + .event_by_id(&request.agent_id, request.event_id)? + .filter(|record| is_notification_event_type(&record.event_type))) + } + + fn notification_neighbor( + &self, + request: NotificationNeighborRequest, + ) -> AgentEventResult> { + let neighbor_key = match request.direction { + NotificationNeighborDirection::Before => { + if request.pivot_event_id == 0 { + return Ok(None); + } + notification_prev_key(&request.agent_id, request.pivot_event_id) + } + NotificationNeighborDirection::After => { + notification_next_key(&request.agent_id, request.pivot_event_id) + } + }; + if let Some(id) = self.id_by_index_key(&neighbor_key)? { + return self.event_by_id(&request.agent_id, id); + } + + match request.direction { + NotificationNeighborDirection::Before => { + let prefix = notification_rev_prefix(&request.agent_id); + let start_after = notification_rev_key(&request.agent_id, request.pivot_event_id); + self.first_event_after_index_key(&request.agent_id, &prefix, &start_after) + } + NotificationNeighborDirection::After => { + let prefix = notification_id_prefix(&request.agent_id); + let start_after = notification_id_key(&request.agent_id, request.pivot_event_id); + self.first_event_after_index_key(&request.agent_id, &prefix, &start_after) + } + } + } + + fn notification_lifecycle( + &self, + request: NotificationLifecycleRequest, + ) -> AgentEventResult> { + if request.limit == 0 { + return Ok(Vec::new()); + } + let filters = notification_filters(&request); + let Some((field, value)) = filters.first() else { + return Err(AgentEventError::InvalidArgument( + "notification_lifecycle requires ref-id, event-id, call-id, or channel".to_owned(), + )); + }; + let prefix = notification_prefix(&request.agent_id, field, value); + let mut out = self + .events_by_index_prefix(&request.agent_id, &prefix, usize::MAX)? + .into_iter() + .filter(|record| notification_record_matches(record, &request)) + .collect::>(); + out.sort_by_key(|record| record.id); + out.truncate(request.limit); + Ok(out) + } +} + +fn write_record_batch(batch: &mut DBAtomicBatch, source: &[u8], record: &EventRecord) { + let id = id_value(record.id); + batch.put_if_absent(TREE_INDEX, source, &id); + batch.put( + TREE_EVENTS, + &event_key(&record.agent_id, record.id), + &encode_event(record).expect("event encodes"), + ); + batch.put( + TREE_INDEX, + &type_id_key(&record.agent_id, &record.event_type, record.id), + &id, + ); + batch.put( + TREE_INDEX, + &type_ts_key(&record.agent_id, &record.event_type, record.ts, record.id), + &id, + ); + if is_lingtai_session_event_type(&record.event_type) { + batch.put(TREE_INDEX, &session_key(&record.agent_id, record.id), &id); + } + if is_notification_event_type(&record.event_type) { + batch.put( + TREE_INDEX, + ¬ification_id_key(&record.agent_id, record.id), + &id, + ); + batch.put( + TREE_INDEX, + ¬ification_rev_key(&record.agent_id, record.id), + &id, + ); + } + if is_tui_clear_completion_event(record) { + batch.put( + TREE_INDEX, + &tui_clear_rev_key(&record.agent_id, record.id), + &id, + ); + } + if let Some(tool_name) = &record.projection.tool_name { + batch.put( + TREE_INDEX, + &super::key::tool_key( + &record.agent_id, + tool_name, + record.projection.tool_action.as_deref(), + record.id, + ), + &id, + ); + } + for (field, value) in [ + ("ref_id", record.projection.notification_ref_id.as_deref()), + ( + "event_id", + record.projection.notification_event_id.as_deref(), + ), + ("call_id", record.projection.notification_call_id.as_deref()), + ("channel", record.projection.notification_channel.as_deref()), + ] { + if let Some(value) = value { + batch.put( + TREE_INDEX, + ¬ification_key(&record.agent_id, field, value, record.id), + &id, + ); + } + } + if let Some(tool_call_id) = &record.projection.tool_call_id { + batch.put( + TREE_INDEX, + &trace_key(&record.agent_id, tool_call_id, record.id), + &id, + ); + } +} + +fn increment_facet( + increments: &mut BTreeMap, ToolFacet>, + key: Vec, + tool_name: &str, + action: Option<&str>, +) { + let entry = increments.entry(key).or_insert_with(|| ToolFacet { + tool_name: tool_name.to_owned(), + action: action.map(ToOwned::to_owned), + count: 0, + }); + entry.count = entry.count.saturating_add(1); +} + +fn sort_tool_facets(facets: &mut [ToolFacet]) { + facets.sort_by(|left, right| { + right + .count + .cmp(&left.count) + .then_with(|| left.tool_name.cmp(&right.tool_name)) + .then_with(|| left.action.cmp(&right.action)) + }); +} + +fn is_lingtai_session_event_type(event_type: &str) -> bool { + LINGTAI_SESSION_EVENT_TYPES.contains(&event_type) +} + +fn session_row_from_event(record: EventRecord) -> SessionEventRow { + SessionEventRow { + id: record.id, + ts: record.ts, + event_type: record.event_type, + fields_json: record.fields_json, + source_file: record.source_file, + source_offset: record.source_offset, + } +} + +fn is_notification_event_type(event_type: &str) -> bool { + event_type.contains("notification") +} + +fn is_tui_clear_completion_event(record: &EventRecord) -> bool { + matches!(record.event_type.as_str(), "psyche_molt" | "clear_received") + && record + .fields_json + .get("source") + .and_then(|value| value.as_str()) + == Some("tui") +} + +fn notification_filters(request: &NotificationLifecycleRequest) -> Vec<(&'static str, &str)> { + let mut filters = Vec::new(); + if let Some(value) = request.ref_id.as_deref() { + filters.push(("ref_id", value)); + } + if let Some(value) = request.event_id.as_deref() { + filters.push(("event_id", value)); + } + if let Some(value) = request.call_id.as_deref() { + filters.push(("call_id", value)); + } + if let Some(value) = request.channel.as_deref() { + filters.push(("channel", value)); + } + filters +} + +fn notification_record_matches( + record: &EventRecord, + request: &NotificationLifecycleRequest, +) -> bool { + option_matches( + request.ref_id.as_deref(), + record.projection.notification_ref_id.as_deref(), + ) && option_matches( + request.event_id.as_deref(), + record.projection.notification_event_id.as_deref(), + ) && option_matches( + request.call_id.as_deref(), + record.projection.notification_call_id.as_deref(), + ) && option_matches( + request.channel.as_deref(), + record.projection.notification_channel.as_deref(), + ) +} + +fn option_matches(expected: Option<&str>, actual: Option<&str>) -> bool { + expected.is_none_or(|expected| actual == Some(expected)) +} + +fn to_store_error(err: holt::Error) -> AgentEventError { + AgentEventError::Store(err.to_string()) +} + +fn event_id_for_source(source_file: &str, source_offset: u64) -> u64 { + const OFFSET_BITS: u64 = 48; + const OFFSET_MASK: u64 = (1_u64 << OFFSET_BITS) - 1; + let hash = source_file_hash(source_file); + let source_prefix = u64::from_str_radix(&hash[..4], 16).unwrap_or(0); + (source_prefix << OFFSET_BITS) | (source_offset & OFFSET_MASK) +} diff --git a/crates/nokv-agent/src/event/ingest.rs b/crates/nokv-agent/src/event/ingest.rs new file mode 100644 index 000000000..b8116c818 --- /dev/null +++ b/crates/nokv-agent/src/event/ingest.rs @@ -0,0 +1,247 @@ +use std::io::BufRead; + +use serde_json::{Map, Value}; + +const MAX_INLINE_FIELDS_JSON_BYTES: usize = 32 * 1024; +const INGEST_RECORD_BATCH_LIMIT: usize = 512; + +use super::store::AgentEventStore; +use super::types::{ + AgentEventResult, EventProjection, IndexCoverage, IngestReport, NewEventRecord, +}; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct JsonlIngestOptions { + pub agent_id: String, + pub source_file: String, + pub file_size: u64, +} + +pub fn ingest_jsonl_reader( + store: &S, + options: JsonlIngestOptions, + mut reader: R, +) -> AgentEventResult +where + S: AgentEventStore + ?Sized, + R: BufRead, +{ + let mut offset = 0_u64; + let mut line_no = 0_u64; + let mut report = IngestReport { + accepted: 0, + duplicates: 0, + parse_errors: 0, + partial_lines: 0, + coverage: empty_coverage(&options.agent_id, &options.source_file, options.file_size), + }; + let mut records = Vec::new(); + loop { + let mut line = Vec::new(); + let read = reader.read_until(b'\n', &mut line)?; + if read == 0 { + break; + } + let line_offset = offset; + offset = offset.saturating_add(read as u64); + if line.iter().all(|byte| byte.is_ascii_whitespace()) { + continue; + } + if !line.ends_with(b"\n") { + report.partial_lines = report.partial_lines.saturating_add(1); + break; + } + line_no = line_no.saturating_add(1); + match parse_event_line(&options, line_offset, line_no, &line) { + Ok(record) => { + records.push(record); + if records.len() >= INGEST_RECORD_BATCH_LIMIT { + let batch = + store.ingest_batch(std::mem::take(&mut records), options.file_size)?; + merge_report(&mut report, batch); + } + } + Err(_) => report.parse_errors = report.parse_errors.saturating_add(1), + } + } + if !records.is_empty() { + let batch = store.ingest_batch(records, options.file_size)?; + merge_report(&mut report, batch); + } + Ok(report) +} + +fn merge_report(total: &mut IngestReport, batch: IngestReport) { + total.accepted = total.accepted.saturating_add(batch.accepted); + total.duplicates = total.duplicates.saturating_add(batch.duplicates); + total.parse_errors = total.parse_errors.saturating_add(batch.parse_errors); + total.partial_lines = total.partial_lines.saturating_add(batch.partial_lines); + if batch.coverage.has_rows() { + total.coverage = batch.coverage; + } +} + +fn parse_event_line( + options: &JsonlIngestOptions, + source_offset: u64, + source_line: u64, + line: &[u8], +) -> AgentEventResult { + let value: Value = serde_json::from_slice(line)?; + let Some(object) = value.as_object() else { + return Err(super::types::AgentEventError::Json( + "event line is not a JSON object".to_owned(), + )); + }; + let Some(event_type) = object.get("type").and_then(Value::as_str) else { + return Err(super::types::AgentEventError::Json( + "event line is missing string type".to_owned(), + )); + }; + let ts = object.get("ts").and_then(Value::as_f64).unwrap_or(0.0); + Ok(NewEventRecord { + agent_id: options.agent_id.clone(), + source_file: options.source_file.clone(), + source_offset, + source_line, + ts, + event_type: event_type.to_owned(), + fields_json: fields_json(object), + projection: project_event(object), + }) +} + +fn fields_json(object: &Map) -> Value { + let mut fields = object.clone(); + for key in ["type", "ts", "address", "agent_name"] { + fields.remove(key); + } + compact_fields_json(Value::Object(fields)) +} + +fn compact_fields_json(value: Value) -> Value { + let encoded_len = serde_json::to_vec(&value) + .map(|bytes| bytes.len()) + .unwrap_or(usize::MAX); + if encoded_len <= MAX_INLINE_FIELDS_JSON_BYTES { + return value; + } + let keys = value + .as_object() + .map(|object| { + object + .keys() + .cloned() + .map(Value::String) + .collect::>() + }) + .unwrap_or_default(); + serde_json::json!({ + "_nokv_compacted": true, + "original_json_bytes": encoded_len, + "keys": keys, + }) +} + +fn project_event(object: &Map) -> EventProjection { + let parsed_tool_args = object.get("tool_args").and_then(parse_tool_args); + let tool_args = parsed_tool_args.as_ref().and_then(Value::as_object); + EventProjection { + tool_name: object + .get("tool_name") + .or_else(|| object.get("tool")) + .and_then(Value::as_str) + .map(ToOwned::to_owned), + tool_action: tool_args + .and_then(|args| args.get("action")) + .and_then(Value::as_str) + .map(ToOwned::to_owned), + tool_call_id: object + .get("tool_call_id") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + tool_trace_id: object + .get("tool_trace_id") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + api_call_id: object + .get("api_call_id") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + command_head: tool_args + .and_then(|args| args.get("command")) + .and_then(Value::as_str) + .and_then(command_head), + file_extension: tool_args + .and_then(|args| args.get("file_path")) + .and_then(Value::as_str) + .and_then(file_extension), + notification_channel: object + .get("channel") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + notification_ref_id: object + .get("ref_id") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + notification_event_id: object + .get("event_id") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + notification_call_id: object + .get("call_id") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + status: object + .get("status") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + error: object + .get("error") + .or_else(|| object.get("exception_message")) + .or_else(|| object.get("exception")) + .or_else(|| object.get("reason")) + .and_then(Value::as_str) + .map(ToOwned::to_owned), + } +} + +fn parse_tool_args(value: &Value) -> Option { + match value { + Value::Object(_) => Some(value.clone()), + Value::String(text) => serde_json::from_str::(text).ok(), + _ => None, + } +} + +fn command_head(command: &str) -> Option { + command.split_whitespace().next().map(|head| { + if head.contains('/') { + std::path::Path::new(head) + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("unknown") + .to_owned() + } else { + head.to_owned() + } + }) +} + +fn file_extension(path: &str) -> Option { + std::path::Path::new(path) + .extension() + .and_then(|ext| ext.to_str()) + .map(|ext| ext.to_ascii_lowercase()) +} + +pub(crate) fn empty_coverage(agent_id: &str, source_file: &str, file_size: u64) -> IndexCoverage { + IndexCoverage { + agent_id: agent_id.to_owned(), + source_file: source_file.to_owned(), + file_size, + min_offset: None, + max_offset: None, + row_count: 0, + } +} diff --git a/crates/nokv-agent/src/event/key.rs b/crates/nokv-agent/src/event/key.rs new file mode 100644 index 000000000..aa74c4382 --- /dev/null +++ b/crates/nokv-agent/src/event/key.rs @@ -0,0 +1,276 @@ +use sha2::{Digest, Sha256}; + +pub const TREE_EVENTS: &str = "events"; +pub const TREE_INDEX: &str = "index"; +pub const TREE_COVERAGE: &str = "coverage"; + +pub fn source_file_hash(source_file: &str) -> String { + let digest = Sha256::digest(source_file.as_bytes()); + hex_lower(&digest[..16]) +} + +pub fn coverage_key(agent_id: &str, source_file: &str) -> Vec { + format!( + "coverage/{}/{}/", + escape(agent_id), + source_file_hash(source_file) + ) + .into_bytes() +} + +pub fn source_key(agent_id: &str, source_file: &str, offset: u64) -> Vec { + format!( + "source/{}/{}/{offset:020}", + escape(agent_id), + source_file_hash(source_file) + ) + .into_bytes() +} + +pub fn event_key(agent_id: &str, event_id: u64) -> Vec { + format!("event/{}/{event_id:020}", escape(agent_id)).into_bytes() +} + +pub fn event_prefix(agent_id: &str) -> Vec { + format!("event/{}/", escape(agent_id)).into_bytes() +} + +pub fn type_id_key(agent_id: &str, event_type: &str, event_id: u64) -> Vec { + format!( + "type_id/{}/{}/{:020}", + escape(agent_id), + escape(event_type), + u64::MAX - event_id + ) + .into_bytes() +} + +pub fn type_id_prefix(agent_id: &str, event_type: &str) -> Vec { + format!("type_id/{}/{}/", escape(agent_id), escape(event_type)).into_bytes() +} + +pub fn type_ts_key(agent_id: &str, event_type: &str, ts: f64, event_id: u64) -> Vec { + format!( + "type_ts/{}/{}/{:020}/{event_id:020}", + escape(agent_id), + escape(event_type), + u64::MAX - timestamp_micros(ts) + ) + .into_bytes() +} + +pub fn type_ts_prefix(agent_id: &str, event_type: &str) -> Vec { + format!("type_ts/{}/{}/", escape(agent_id), escape(event_type)).into_bytes() +} + +pub fn tool_key(agent_id: &str, tool_name: &str, action: Option<&str>, event_id: u64) -> Vec { + format!( + "tool/{}/{}/{}/{:020}", + escape(agent_id), + escape(tool_name), + escape(action.unwrap_or("")), + u64::MAX - event_id + ) + .into_bytes() +} + +pub fn tool_prefix(agent_id: &str) -> Vec { + format!("tool/{}/", escape(agent_id)).into_bytes() +} + +pub fn tool_name_facet_key(agent_id: &str, tool_name: &str) -> Vec { + format!( + "facet/{}/tool_name/{}/", + escape(agent_id), + escape(tool_name) + ) + .into_bytes() +} + +pub fn tool_name_facet_prefix(agent_id: &str) -> Vec { + format!("facet/{}/tool_name/", escape(agent_id)).into_bytes() +} + +pub fn tool_action_facet_key(agent_id: &str, tool_name: &str, action: Option<&str>) -> Vec { + format!( + "facet/{}/tool_action/{}/{}/", + escape(agent_id), + escape(tool_name), + escape(action.unwrap_or("")) + ) + .into_bytes() +} + +pub fn tool_action_facet_prefix(agent_id: &str) -> Vec { + format!("facet/{}/tool_action/", escape(agent_id)).into_bytes() +} + +pub fn session_key(agent_id: &str, event_id: u64) -> Vec { + format!("session/{}/{event_id:020}", escape(agent_id)).into_bytes() +} + +pub fn session_prefix(agent_id: &str) -> Vec { + format!("session/{}/", escape(agent_id)).into_bytes() +} + +pub fn notification_key(agent_id: &str, field: &str, value: &str, event_id: u64) -> Vec { + format!( + "notification/{}/{}/{}/{event_id:020}", + escape(agent_id), + escape(field), + escape(value) + ) + .into_bytes() +} + +pub fn notification_prefix(agent_id: &str, field: &str, value: &str) -> Vec { + format!( + "notification/{}/{}/{}/", + escape(agent_id), + escape(field), + escape(value) + ) + .into_bytes() +} + +pub fn notification_id_key(agent_id: &str, event_id: u64) -> Vec { + format!("notification_id/{}/{event_id:020}", escape(agent_id)).into_bytes() +} + +pub fn notification_id_prefix(agent_id: &str) -> Vec { + format!("notification_id/{}/", escape(agent_id)).into_bytes() +} + +pub fn notification_rev_key(agent_id: &str, event_id: u64) -> Vec { + format!( + "notification_rev/{}/{:020}", + escape(agent_id), + u64::MAX - event_id + ) + .into_bytes() +} + +pub fn notification_rev_prefix(agent_id: &str) -> Vec { + format!("notification_rev/{}/", escape(agent_id)).into_bytes() +} + +pub fn notification_prev_key(agent_id: &str, event_id: u64) -> Vec { + format!("notification_prev/{}/{event_id:020}", escape(agent_id)).into_bytes() +} + +pub fn notification_next_key(agent_id: &str, event_id: u64) -> Vec { + format!("notification_next/{}/{event_id:020}", escape(agent_id)).into_bytes() +} + +pub fn notification_tail_key(agent_id: &str) -> Vec { + format!("notification_tail/{}/", escape(agent_id)).into_bytes() +} + +pub fn tui_clear_rev_key(agent_id: &str, event_id: u64) -> Vec { + format!( + "tui_clear_rev/{}/{:020}", + escape(agent_id), + u64::MAX - event_id + ) + .into_bytes() +} + +pub fn tui_clear_rev_prefix(agent_id: &str) -> Vec { + format!("tui_clear_rev/{}/", escape(agent_id)).into_bytes() +} + +pub fn trace_key(agent_id: &str, tool_call_id: &str, event_id: u64) -> Vec { + format!( + "trace/{}/{}/{event_id:020}", + escape(agent_id), + escape(tool_call_id) + ) + .into_bytes() +} + +pub fn trace_prefix(agent_id: &str, tool_call_id: &str) -> Vec { + format!("trace/{}/{}/", escape(agent_id), escape(tool_call_id)).into_bytes() +} + +pub fn id_from_index_value(value: &[u8]) -> Option { + let raw: [u8; 8] = value.try_into().ok()?; + Some(u64::from_be_bytes(raw)) +} + +pub fn id_value(event_id: u64) -> [u8; 8] { + event_id.to_be_bytes() +} + +fn timestamp_micros(ts: f64) -> u64 { + if !ts.is_finite() || ts <= 0.0 { + return 0; + } + let micros = ts * 1_000_000.0; + if micros >= u64::MAX as f64 { + u64::MAX + } else { + micros as u64 + } +} + +fn escape(input: &str) -> String { + let mut out = String::with_capacity(input.len()); + for byte in input.bytes() { + match byte { + b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' => { + out.push(byte as char); + } + _ => { + out.push('%'); + out.push(nibble(byte >> 4)); + out.push(nibble(byte & 0x0f)); + } + } + } + out +} + +fn hex_lower(bytes: &[u8]) -> String { + let mut out = String::with_capacity(bytes.len() * 2); + for byte in bytes { + out.push(nibble(byte >> 4)); + out.push(nibble(byte & 0x0f)); + } + out +} + +fn nibble(value: u8) -> char { + match value { + 0..=9 => (b'0' + value) as char, + 10..=15 => (b'a' + (value - 10)) as char, + _ => unreachable!("nibble is four bits"), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn type_id_key_sorts_newest_first() { + let newer = type_id_key("agent", "tool_call", 20); + let older = type_id_key("agent", "tool_call", 10); + assert!(newer < older); + } + + #[test] + fn type_ts_key_sorts_newest_timestamp_first() { + let newer = type_ts_key("agent", "refresh_complete", 20.0, 1); + let older = type_ts_key("agent", "refresh_complete", 10.0, 2); + assert!(newer < older); + } + + #[test] + fn source_file_hash_is_stable_and_short() { + assert_eq!( + source_file_hash("logs/events.jsonl"), + source_file_hash("logs/events.jsonl") + ); + assert_eq!(source_file_hash("logs/events.jsonl").len(), 32); + } +} diff --git a/crates/nokv-agent/src/event/mod.rs b/crates/nokv-agent/src/event/mod.rs new file mode 100644 index 000000000..e9caa76fd --- /dev/null +++ b/crates/nokv-agent/src/event/mod.rs @@ -0,0 +1,25 @@ +//! LingTai event-index surface. +//! +//! This module indexes LingTai `logs/events.jsonl` as a derived view. JSONL is +//! still authoritative; the index is rebuildable and must report coverage +//! clearly enough for LingTai to fall back to JSONL when needed. + +pub mod codec; +pub mod holt; +pub mod ingest; +pub mod key; +pub mod store; +pub mod types; + +pub use holt::HoltAgentEventStore; +pub use ingest::{ingest_jsonl_reader, JsonlIngestOptions}; +pub use store::AgentEventStore; +pub use types::{ + AgentEventError, AgentEventResult, CompletionAfter, CompletionAfterRequest, ErrorEventsRequest, + EventProjection, EventRecord, EventTime, IndexCoverage, IngestReport, LatestEventsRequest, + MoltSessionWindows, NewEventRecord, NotificationEventByIdRequest, NotificationEventsRequest, + NotificationLifecycleRequest, NotificationNeighborDirection, NotificationNeighborRequest, + RecentTimesRequest, SessionEventRow, SessionEventsRequest, SessionRowsRequest, ToolFacet, + ToolFacetRequest, ToolTraceRequest, TuiClearCompletion, TuiClearCompletionRequest, + LINGTAI_SESSION_EVENT_TYPES, +}; diff --git a/crates/nokv-agent/src/event/store.rs b/crates/nokv-agent/src/event/store.rs new file mode 100644 index 000000000..fb8686245 --- /dev/null +++ b/crates/nokv-agent/src/event/store.rs @@ -0,0 +1,74 @@ +use super::types::{ + AgentEventResult, CompletionAfter, CompletionAfterRequest, ErrorEventsRequest, EventRecord, + EventTime, IndexCoverage, IngestReport, LatestEventsRequest, MoltSessionWindows, + NewEventRecord, NotificationEventByIdRequest, NotificationEventsRequest, + NotificationLifecycleRequest, NotificationNeighborRequest, RecentTimesRequest, SessionEventRow, + SessionEventsRequest, SessionRowsRequest, ToolFacet, ToolFacetRequest, ToolTraceRequest, + TuiClearCompletion, TuiClearCompletionRequest, +}; + +pub trait AgentEventStore { + fn ingest_batch( + &self, + records: Vec, + file_size: u64, + ) -> AgentEventResult; + + fn coverage( + &self, + agent_id: &str, + source_file: &str, + ) -> AgentEventResult>; + + fn latest_events(&self, request: LatestEventsRequest) -> AgentEventResult>; + + fn stream_session_events( + &self, + request: SessionEventsRequest, + ) -> AgentEventResult>; + + fn stream_session_rows( + &self, + request: SessionRowsRequest, + ) -> AgentEventResult>; + + fn tool_facets(&self, request: ToolFacetRequest) -> AgentEventResult>; + + fn tool_trace(&self, request: ToolTraceRequest) -> AgentEventResult>; + + fn recent_times(&self, request: RecentTimesRequest) -> AgentEventResult>; + + fn molt_session_windows(&self, agent_id: &str) -> AgentEventResult; + + fn error_events(&self, request: ErrorEventsRequest) -> AgentEventResult>; + + fn completion_after( + &self, + request: CompletionAfterRequest, + ) -> AgentEventResult; + + fn tui_clear_completion( + &self, + request: TuiClearCompletionRequest, + ) -> AgentEventResult; + + fn notification_events( + &self, + request: NotificationEventsRequest, + ) -> AgentEventResult>; + + fn notification_event_by_id( + &self, + request: NotificationEventByIdRequest, + ) -> AgentEventResult>; + + fn notification_neighbor( + &self, + request: NotificationNeighborRequest, + ) -> AgentEventResult>; + + fn notification_lifecycle( + &self, + request: NotificationLifecycleRequest, + ) -> AgentEventResult>; +} diff --git a/crates/nokv-agent/src/event/types.rs b/crates/nokv-agent/src/event/types.rs new file mode 100644 index 000000000..94c423e68 --- /dev/null +++ b/crates/nokv-agent/src/event/types.rs @@ -0,0 +1,263 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +pub type AgentEventResult = Result; + +pub const LINGTAI_SESSION_EVENT_TYPES: &[&str] = &[ + "thinking", + "diary", + "text_input", + "text_output", + "tool_call", + "tool_result", + "llm_call", + "llm_response", + "insight", + "consultation_fire", + "notification_pair_injected", + "apriori_summary_generated", + "apriori_summary_cap_refused", + "apriori_summary_failed", + "apriori_summary_empty", + "apriori_summary_no_summarizer", + "aed_attempt", + "aed_exhausted", + "aed_timeout", +]; + +#[derive(Debug)] +pub enum AgentEventError { + InvalidArgument(String), + Io(String), + Json(String), + Store(String), +} + +impl std::fmt::Display for AgentEventError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::InvalidArgument(msg) => write!(f, "invalid agent event argument: {msg}"), + Self::Io(msg) => write!(f, "agent event io error: {msg}"), + Self::Json(msg) => write!(f, "agent event json error: {msg}"), + Self::Store(msg) => write!(f, "agent event store error: {msg}"), + } + } +} + +impl std::error::Error for AgentEventError {} + +impl From for AgentEventError { + fn from(err: std::io::Error) -> Self { + Self::Io(err.to_string()) + } +} + +impl From for AgentEventError { + fn from(err: serde_json::Error) -> Self { + Self::Json(err.to_string()) + } +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct EventProjection { + pub tool_name: Option, + pub tool_action: Option, + pub tool_call_id: Option, + pub tool_trace_id: Option, + pub api_call_id: Option, + pub command_head: Option, + pub file_extension: Option, + pub notification_channel: Option, + pub notification_ref_id: Option, + pub notification_event_id: Option, + pub notification_call_id: Option, + pub status: Option, + pub error: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct NewEventRecord { + pub agent_id: String, + pub source_file: String, + pub source_offset: u64, + pub source_line: u64, + pub ts: f64, + pub event_type: String, + pub fields_json: Value, + pub projection: EventProjection, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct EventRecord { + pub id: u64, + pub agent_id: String, + pub source_file: String, + pub source_offset: u64, + pub source_line: u64, + pub ts: f64, + pub event_type: String, + pub fields_json: Value, + pub projection: EventProjection, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct IndexCoverage { + pub agent_id: String, + pub source_file: String, + pub file_size: u64, + pub min_offset: Option, + pub max_offset: Option, + pub row_count: u64, +} + +impl IndexCoverage { + pub fn has_rows(&self) -> bool { + self.row_count > 0 && self.min_offset.is_some() && self.max_offset.is_some() + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct IngestReport { + pub accepted: u64, + pub duplicates: u64, + pub parse_errors: u64, + pub partial_lines: u64, + pub coverage: IndexCoverage, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct LatestEventsRequest { + pub agent_id: String, + pub event_type: String, + pub limit: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SessionEventsRequest { + pub agent_id: String, + pub event_types: Vec, + pub limit: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SessionRowsRequest { + pub agent_id: String, + pub limit: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct SessionEventRow { + pub id: u64, + pub ts: f64, + pub event_type: String, + pub fields_json: Value, + pub source_file: String, + pub source_offset: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ToolFacetRequest { + pub agent_id: String, + pub limit: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ToolFacet { + pub tool_name: String, + pub action: Option, + pub count: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ToolTraceRequest { + pub agent_id: String, + pub tool_call_id: String, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RecentTimesRequest { + pub agent_id: String, + pub event_type: String, + pub limit: usize, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct EventTime { + pub id: u64, + pub ts: f64, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct MoltSessionWindows { + pub ok: bool, + pub current_since: Option, + pub last_since: Option, + pub last_before: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ErrorEventsRequest { + pub agent_id: String, + pub event_types: Vec, + pub limit: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CompletionAfterRequest { + pub agent_id: String, + pub event_type: String, + pub source_offset: u64, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct CompletionAfter { + pub found: bool, + pub event: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct TuiClearCompletionRequest { + pub agent_id: String, + pub source_offset: u64, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct TuiClearCompletion { + pub found: bool, + pub event: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct NotificationEventsRequest { + pub agent_id: String, + pub limit: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct NotificationEventByIdRequest { + pub agent_id: String, + pub event_id: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct NotificationNeighborRequest { + pub agent_id: String, + pub pivot_event_id: u64, + pub direction: NotificationNeighborDirection, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum NotificationNeighborDirection { + Before, + After, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct NotificationLifecycleRequest { + pub agent_id: String, + pub ref_id: Option, + pub event_id: Option, + pub call_id: Option, + pub channel: Option, + pub limit: usize, +} diff --git a/crates/nokv-agent/src/lib.rs b/crates/nokv-agent/src/lib.rs index 230e5c94c..bfb1714f2 100644 --- a/crates/nokv-agent/src/lib.rs +++ b/crates/nokv-agent/src/lib.rs @@ -1,2664 +1,13 @@ -//! NoKV agent tool surface. +//! NoKV agent surfaces. //! -//! Transport-free agent tooling: the JSON tool definitions exposed to a model, -//! the dispatcher that maps a tool call onto a namespace verb, argument -//! validation, and result shaping. Depends only on `nokv-meta`, `nokv-object` -//! and `nokv-types` — never on `nokv-client`, `nokv-protocol` or `nokv-control`. +//! The shipped namespace tool surface is transport-free: JSON tool +//! definitions, dispatch, validation, and result shaping over NoKV namespace +//! verbs. LingTai event-index code lives under [`event`] as a separate derived +//! index surface over `logs/events.jsonl`. -use nokv_meta::{ - MetadataStore, NamespaceAggregateGroup, NamespaceAggregateMeasure, NamespaceAggregateOp, - NamespaceAggregateRequest, NamespaceAggregateResult, NamespaceAggregateSort, - NamespaceAggregateValue, NamespaceBodyDescriptor, NamespaceCard, NamespaceCardKind, - NamespaceFacetSummary, NamespaceFacetValue, NamespaceFilterCapability, NamespaceFindField, - NamespaceFindRequest, NamespaceFindResult, NamespaceGrepRequest, NamespaceGrepResult, - NamespaceIndexValue, NamespaceListOptions, NamespaceListPage, NamespacePredicate, - NamespacePredicateOp, NamespacePredicateValue, NamespaceQueryCatalog, NamespaceReadFormat, - NamespaceReadItem, NamespaceReadOptions, NamespaceReadPage, NamespaceRecordCount, - NamespaceRecordType, NamespaceSchema, NamespaceSort, NamespaceSortDirection, - NamespaceSortField, NoKvFs, -}; -use nokv_object::ObjectStore; -use serde_json::{json, Map, Value}; - -/// Error surface for the agent tool layer. -/// -/// The copied dispatcher only ever constructs `Metadata`, `NotFound`, and -/// `InvalidArgument`. `Other` is unused here; it exists so a later step can -/// add `From` in `nokv-client`. -#[derive(Debug)] -pub enum AgentError { - Metadata(nokv_meta::MetadError), - NotFound(String), - InvalidArgument(String), - Other(String), -} - -impl std::fmt::Display for AgentError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - // Display strings MUST stay byte-identical to the old ClientError rendering - // (crates/nokv-client/src/lib.rs:93-119) so model/telemetry output is unchanged: - Self::Metadata(err) => write!(f, "metadata service error: {err}"), - Self::NotFound(path) => write!(f, "path component not found: {path}"), - Self::InvalidArgument(err) => write!(f, "metadata protocol error: {err}"), - Self::Other(msg) => write!(f, "{msg}"), - } - } -} - -impl std::error::Error for AgentError {} - -impl From for AgentError { - fn from(err: nokv_meta::MetadError) -> Self { - Self::Metadata(err) - } -} - -const DEFAULT_AGENT_PAGE_LIMIT: usize = 100; -const MAX_AGENT_PAGE_LIMIT: usize = 100; -const DEFAULT_AGENT_READ_LIMIT: usize = 100; -const MAX_AGENT_READ_LIMIT: usize = 300; -const DEFAULT_AGENT_FIND_LIMIT: usize = 10; -const MAX_AGENT_FIND_LIMIT: usize = 10; -const DEFAULT_AGENT_AGGREGATE_LIMIT: usize = 20; -const MAX_AGENT_AGGREGATE_LIMIT: usize = 100; -const DEFAULT_AGENT_GREP_LIMIT: usize = 100; -const MAX_AGENT_GREP_LIMIT: usize = 300; -/// Mirror of the metadata server's grep pattern cap (nokv-meta -/// MAX_GREP_PATTERNS); enforced there, advertised in the schema here. -const MAX_AGENT_GREP_PATTERNS: usize = 16; - -#[derive(Clone, Debug, PartialEq)] -pub struct AgentToolDefinition { - pub name: &'static str, - pub description: &'static str, - pub parameters: Value, -} - -pub trait AgentNamespace { - fn agent_stat_card(&self, path: &str) -> Result, AgentError>; - - fn agent_list_page( - &self, - path: &str, - options: NamespaceListOptions, - ) -> Result; - - fn agent_find_paths( - &self, - request: NamespaceFindRequest, - ) -> Result; - - fn agent_aggregate_paths( - &self, - request: NamespaceAggregateRequest, - ) -> Result; - - fn agent_grep_paths( - &self, - request: NamespaceGrepRequest, - ) -> Result; - - fn agent_read_page( - &self, - path: &str, - options: NamespaceReadOptions, - ) -> Result; -} - -impl AgentNamespace for NoKvFs -where - M: MetadataStore, - O: ObjectStore, -{ - fn agent_stat_card(&self, path: &str) -> Result, AgentError> { - self.stat_card(path).map_err(AgentError::Metadata) - } - - fn agent_list_page( - &self, - path: &str, - options: NamespaceListOptions, - ) -> Result { - self.list_page(path, options).map_err(AgentError::Metadata) - } - - fn agent_find_paths( - &self, - request: NamespaceFindRequest, - ) -> Result { - self.find_paths(request).map_err(AgentError::Metadata) - } - - fn agent_aggregate_paths( - &self, - request: NamespaceAggregateRequest, - ) -> Result { - self.aggregate_paths(request).map_err(AgentError::Metadata) - } - - fn agent_grep_paths( - &self, - request: NamespaceGrepRequest, - ) -> Result { - self.grep_paths(request).map_err(AgentError::Metadata) - } - - fn agent_read_page( - &self, - path: &str, - options: NamespaceReadOptions, - ) -> Result { - self.read_page(path, options).map_err(AgentError::Metadata) - } -} - -pub fn agent_tool_definitions() -> Vec { - vec![ - AgentToolDefinition { - name: "ls", - description: "List direct children to discover paths. Not recursive; use stat for one path details.", - parameters: json!({ - "type": "object", - "required": ["path"], - "properties": { - "path": {"type": "string"}, - "cursor": {"type": ["string", "null"]}, - "limit": {"type": "integer", "minimum": 1, "maximum": MAX_AGENT_PAGE_LIMIT} - }, - "additionalProperties": false - }), - }, - AgentToolDefinition { - name: "stat", - description: "Inspect a single path compact card: counts, body, schema, sample, catalog, and indexed values. Use read for file body content.", - parameters: json!({ - "type": "object", - "required": ["path"], - "properties": { - "path": {"type": "string"} - }, - "additionalProperties": false - }), - }, - AgentToolDefinition { - name: "catalog", - description: "Discover field ids for find predicates, projections, sort, facets, and aggregate group or measure fields.", - parameters: json!({ - "type": "object", - "required": ["path"], - "properties": { - "path": {"type": "string"}, - "field_prefix": {"type": ["string", "null"]}, - "include_facets": {"type": "boolean"} - }, - "additionalProperties": false - }), - }, - AgentToolDefinition { - name: "read", - description: "Read file body content. Structured mode returns JSON, YAML, or text records; bytes mode returns byte ranges.", - parameters: json!({ - "type": "object", - "required": ["path"], - "properties": { - "path": {"type": "string"}, - "format": {"type": "string", "enum": ["structured", "bytes"]}, - "cursor": {"type": ["string", "null"]}, - "offset": {"type": "integer", "minimum": 0}, - "limit": {"type": "integer", "minimum": 1, "maximum": MAX_AGENT_READ_LIMIT}, - }, - "additionalProperties": false - }), - }, - AgentToolDefinition { - name: "aggregate", - description: "Compute summary rows using catalog field ids: count, sum, avg, min, max, group, filter, and sort.", - parameters: json!({ - "type": "object", - "required": ["path", "measures"], - "properties": { - "path": {"type": "string"}, - "predicates": { - "type": "array", - "items": { - "type": "object", - "required": ["field", "op"], - "properties": { - "field": {"type": "string"}, - "op": { - "type": "string", - "enum": ["eq", "ne", "in", "prefix", "suffix", "contains", "gt", "gte", "lt", "lte", "exists", "not_exists"] - }, - "value": { - "anyOf": [ - {"type": "string"}, - {"type": "integer", "minimum": 0}, - {"type": "number"}, - {"type": "boolean"}, - { - "type": "array", - "items": {"type": ["string", "integer", "number", "boolean"]} - }, - {"type": "null"} - ] - } - }, - "additionalProperties": false - } - }, - "group_by": { - "type": "array", - "items": {"type": "string"} - }, - "measures": { - "type": "array", - "items": { - "type": "object", - "required": ["name", "op"], - "properties": { - "name": {"type": "string"}, - "op": {"type": "string", "enum": ["count", "sum", "avg", "min", "max"]}, - "field": {"type": ["string", "null"]} - }, - "additionalProperties": false - } - }, - "sort": { - "type": "array", - "items": { - "type": "object", - "required": ["field"], - "properties": { - "field": {"type": "string"}, - "direction": {"type": "string", "enum": ["asc", "desc"]} - }, - "additionalProperties": false - } - }, - "limit": {"type": "integer", "minimum": 1, "maximum": MAX_AGENT_AGGREGATE_LIMIT} - }, - "additionalProperties": false - }), - }, - AgentToolDefinition { - name: "find", - description: - "Search paths with catalog field predicates and project fields. Use read for body content and stat for schema or sample.", - parameters: json!({ - "type": "object", - "required": ["path"], - "properties": { - "path": {"type": "string"}, - "predicates": { - "type": "array", - "items": { - "type": "object", - "required": ["field", "op"], - "properties": { - "field": {"type": "string"}, - "op": { - "type": "string", - "enum": ["eq", "ne", "in", "prefix", "suffix", "contains", "gt", "gte", "lt", "lte", "exists", "not_exists"] - }, - "value": { - "anyOf": [ - {"type": "string"}, - {"type": "integer", "minimum": 0}, - {"type": "number"}, - {"type": "boolean"}, - { - "type": "array", - "items": {"type": ["string", "integer", "number", "boolean"]} - }, - {"type": "null"} - ] - } - }, - "additionalProperties": false - } - }, - "sort": { - "type": "array", - "items": { - "type": "object", - "required": ["field"], - "properties": { - "field": {"type": "string"}, - "direction": {"type": "string", "enum": ["asc", "desc"]} - }, - "additionalProperties": false - } - }, - "fields": { - "type": "array", - "items": {"type": "string"} - }, - "facets": { - "type": "array", - "items": {"type": "string"} - }, - "cursor": {"type": ["string", "null"]}, - "limit": {"type": "integer", "minimum": 1, "maximum": MAX_AGENT_FIND_LIMIT} - }, - "additionalProperties": false - }), - }, - AgentToolDefinition { - name: "grep", - description: "Search file bodies for case-insensitive literal substrings and return matching lines with line numbers. patterns adds OR alternatives to pattern (at most 16); glob filters file names. Scope path to one directory or file when known.", - parameters: json!({ - "type": "object", - "required": ["path", "pattern", "recursive"], - "properties": { - "path": {"type": "string"}, - "pattern": {"type": "string"}, - "patterns": {"type": "array", "items": {"type": "string"}, "maxItems": MAX_AGENT_GREP_PATTERNS}, - "glob": {"type": ["string", "null"]}, - "recursive": {"type": "boolean"}, - "cursor": {"type": ["string", "null"]}, - "limit": {"type": "integer", "minimum": 1, "maximum": MAX_AGENT_GREP_LIMIT} - }, - "additionalProperties": false - }), - }, - ] -} - -pub fn execute_agent_tool(namespace: &T, name: &str, args: &Value) -> Result -where - T: AgentNamespace + ?Sized, -{ - match name { - "ls" => execute_ls(namespace, args), - "stat" => execute_stat(namespace, args), - "catalog" => execute_catalog(namespace, args), - "grep" => execute_grep(namespace, args), - "read" => execute_read(namespace, args), - "find" => execute_find(namespace, args), - "aggregate" => execute_aggregate(namespace, args), - other => Err(AgentError::InvalidArgument(format!( - "unknown agent tool {other}" - ))), - } -} - -fn execute_ls(namespace: &T, args: &Value) -> Result -where - T: AgentNamespace + ?Sized, -{ - let path = required_string_arg(args, "path")?; - let page = namespace.agent_list_page( - path, - NamespaceListOptions { - cursor: optional_string_arg(args, "cursor")?, - limit: optional_usize_arg(args, "limit")?.unwrap_or(DEFAULT_AGENT_PAGE_LIMIT), - }, - )?; - Ok(list_page_json(&page)) -} - -fn execute_stat(namespace: &T, args: &Value) -> Result -where - T: AgentNamespace + ?Sized, -{ - let path = required_string_arg(args, "path")?; - let card = namespace - .agent_stat_card(path)? - .ok_or_else(|| AgentError::NotFound(path.to_owned()))?; - Ok(json!({ - "card": card_json(&card), - })) -} - -fn execute_catalog(namespace: &T, args: &Value) -> Result -where - T: AgentNamespace + ?Sized, -{ - let path = required_string_arg(args, "path")?; - let field_prefix = optional_string_arg(args, "field_prefix")?; - let include_facets = optional_bool_arg(args, "include_facets")?.unwrap_or(false); - let card = namespace - .agent_stat_card(path)? - .ok_or_else(|| AgentError::NotFound(path.to_owned()))?; - let catalog = filtered_catalog(&card.catalog, field_prefix.as_deref(), include_facets); - let catalog_empty = catalog_is_empty(&catalog); - let child_catalogs = if catalog_empty && matches!(card.kind, NamespaceCardKind::Directory) { - child_catalogs(namespace, &card.path, include_facets)? - } else { - Vec::new() - }; - Ok(json!({ - "path": card.path, - "catalog_empty": catalog_empty, - "catalog": catalog_json(&catalog), - "child_catalogs": child_catalogs, - })) -} - -fn execute_grep(namespace: &T, args: &Value) -> Result -where - T: AgentNamespace + ?Sized, -{ - let path = required_string_arg(args, "path")?; - let pattern = required_string_arg(args, "pattern")?; - let recursive = optional_bool_arg(args, "recursive")?.ok_or_else(|| { - AgentError::InvalidArgument("grep requires a boolean recursive argument".to_owned()) - })?; - let result = namespace.agent_grep_paths(NamespaceGrepRequest { - path: path.to_owned(), - pattern: pattern.to_owned(), - patterns: optional_string_array_arg(args, "patterns")?.unwrap_or_default(), - recursive, - name_glob: optional_string_arg(args, "glob")?, - cursor: optional_string_arg(args, "cursor")?, - limit: optional_bounded_usize_arg(args, "limit", MAX_AGENT_GREP_LIMIT)? - .unwrap_or(DEFAULT_AGENT_GREP_LIMIT), - max_files: None, - max_bytes: None, - })?; - Ok(grep_result_json(&result)) -} - -fn grep_result_json(result: &NamespaceGrepResult) -> Value { - json!({ - "path": result.path, - "pattern": result.pattern, - "recursive": result.recursive, - "matches": result - .matches - .iter() - .map(|match_| json!({ - "path": match_.path, - "line_number": match_.line_number, - "snippet": match_.snippet, - })) - .collect::>(), - "files_scanned": result.files_scanned, - "next_cursor": result.next_cursor, - "truncated": result.truncated, - }) -} - -fn execute_read(namespace: &T, args: &Value) -> Result -where - T: AgentNamespace + ?Sized, -{ - let path = required_string_arg(args, "path")?; - let options = NamespaceReadOptions { - format: read_format_arg(args)?, - cursor: optional_string_arg(args, "cursor")?, - offset: optional_u64_arg(args, "offset")?.unwrap_or(0), - limit: optional_bounded_usize_arg(args, "limit", MAX_AGENT_READ_LIMIT)? - .unwrap_or(DEFAULT_AGENT_READ_LIMIT), - expected_generation: None, - }; - guard_large_structured_pagination(namespace, path, &options)?; - let page = namespace.agent_read_page(path, options)?; - Ok(read_page_json(&page)) -} - -fn guard_large_structured_pagination( - namespace: &T, - path: &str, - options: &NamespaceReadOptions, -) -> Result<(), AgentError> -where - T: AgentNamespace + ?Sized, -{ - if !matches!(options.format, NamespaceReadFormat::Structured) { - return Ok(()); - } - let Some(card) = namespace.agent_stat_card(path)? else { - return Ok(()); - }; - let Some(record_count) = card.record_count else { - return Ok(()); - }; - if record_count.count <= MAX_AGENT_READ_LIMIT { - return Ok(()); - } - Err(AgentError::InvalidArgument(format!( - "structured pagination for {path} has {} records; use bytes format with offset and limit, grep to locate lines, or stat record_count", - record_count.count - ))) -} - -fn execute_find(namespace: &T, args: &Value) -> Result -where - T: AgentNamespace + ?Sized, -{ - let path = required_string_arg(args, "path")?; - let fields = optional_string_array_arg(args, "fields")?; - reject_unsupported_find_arguments(args)?; - let result = namespace.agent_find_paths(NamespaceFindRequest { - path: path.to_owned(), - predicates: predicates_arg(args)?, - sort: sort_arg(args)?, - include: Vec::new(), - facets: facets_arg(args)?, - cursor: optional_string_arg(args, "cursor")?, - limit: optional_bounded_usize_arg(args, "limit", MAX_AGENT_FIND_LIMIT)? - .unwrap_or(DEFAULT_AGENT_FIND_LIMIT), - })?; - Ok(find_result_json(&result, fields.as_deref())) -} - -fn execute_aggregate(namespace: &T, args: &Value) -> Result -where - T: AgentNamespace + ?Sized, -{ - let path = required_string_arg(args, "path")?; - let predicates = predicates_arg(args)?; - let group_by = group_by_arg(args)?; - let measures = aggregate_measures_arg(args)?; - let sort = aggregate_sort_arg(args)?; - let limit = optional_bounded_usize_arg(args, "limit", MAX_AGENT_AGGREGATE_LIMIT)? - .unwrap_or(DEFAULT_AGENT_AGGREGATE_LIMIT); - let result = namespace.agent_aggregate_paths(NamespaceAggregateRequest { - path: path.to_owned(), - predicates, - group_by, - measures, - sort, - limit, - })?; - Ok(aggregate_result_json(&result)) -} - -fn list_page_json(page: &NamespaceListPage) -> Value { - json!({ - "path": page.path, - "entry_count": page.entry_count, - "entries": page.entries.iter().map(list_entry_json).collect::>(), - "next_cursor": page.next_cursor, - "truncated": page.truncated, - }) -} - -fn find_result_json(result: &NamespaceFindResult, fields: Option<&[String]>) -> Value { - json!({ - "path": result.path, - "match_count": result.match_count, - "matches": result.matches.iter().map(|card| find_match_json(card, fields)).collect::>(), - "facets": result.facets.iter().map(facet_summary_json).collect::>(), - "next_cursor": result.next_cursor, - "truncated": result.truncated, - }) -} - -fn read_page_json(page: &NamespaceReadPage) -> Value { - json!({ - "path": page.path, - "generation": page.generation, - "total_size_bytes": page.total_size_bytes, - "format": read_format_name(&page.format), - "record_type": page.record_type.as_ref().map(record_type_name), - "record_count": page.record_count, - "cursor": page.cursor, - "next_cursor": page.next_cursor, - "truncated": page.truncated, - "items": page.items.iter().map(read_item_json).collect::>(), - "bytes": page.bytes, - }) -} - -fn card_json(card: &NamespaceCard) -> Value { - json!({ - "path": card.path, - "name": card.name, - "kind": card_kind_name(&card.kind), - "size_bytes": card.size_bytes, - "entry_count": card.entry_count, - "record_count": card.record_count.as_ref().map(record_count_json), - "schema": card.schema.as_ref().map(schema_json), - "sample": card.sample, - "body": card.body.as_ref().map(body_json), - "catalog": catalog_json(&card.catalog), - "indexed_values": card.indexed_values.iter().map(index_value_json).collect::>(), - }) -} - -fn list_entry_json(card: &NamespaceCard) -> Value { - json!({ - "path": card.path, - "name": card.name, - "kind": card_kind_name(&card.kind), - "size_bytes": card.size_bytes, - "entry_count": card.entry_count, - }) -} - -fn find_match_json(card: &NamespaceCard, fields: Option<&[String]>) -> Value { - if fields.is_some() { - let mut object = Map::new(); - object.insert("path".to_owned(), json!(card.path)); - object.insert( - "values".to_owned(), - projected_values_json(card, fields.unwrap_or(&[])), - ); - return Value::Object(object); - } - let mut object = Map::new(); - object.insert("path".to_owned(), json!(card.path)); - Value::Object(object) -} - -fn projected_values_json(card: &NamespaceCard, fields: &[String]) -> Value { - let mut output = Map::new(); - for field in fields { - let values = projected_predicate_values(card, field) - .into_iter() - .map(|value| predicate_value_json(&value)) - .collect::>(); - match values.as_slice() { - [] => {} - [value] => { - output.insert(field.clone(), value.clone()); - } - _ => { - output.insert(field.clone(), Value::Array(values)); - } - } - } - Value::Object(output) -} - -fn projected_predicate_values(card: &NamespaceCard, field: &str) -> Vec { - match field { - "path" => vec![NamespacePredicateValue::String(card.path.clone())], - "name" => vec![NamespacePredicateValue::String(card.name.clone())], - "kind" => vec![NamespacePredicateValue::String( - card_kind_name(&card.kind).to_owned(), - )], - "size_bytes" => card - .size_bytes - .map(NamespacePredicateValue::U64) - .into_iter() - .collect(), - "body.content_type" => card - .body - .as_ref() - .map(|body| NamespacePredicateValue::String(body.content_type.clone())) - .into_iter() - .collect(), - "body.producer" => card - .body - .as_ref() - .map(|body| NamespacePredicateValue::String(body.producer.clone())) - .into_iter() - .collect(), - _ => card - .indexed_values - .iter() - .filter(|value| value.field.id == field) - .map(|value| value.value.clone()) - .collect(), - } -} - -fn record_count_json(count: &NamespaceRecordCount) -> Value { - json!(count.count) -} - -fn schema_json(schema: &NamespaceSchema) -> Value { - json!({ - "record_type": record_type_name(&schema.record_type), - "fields": schema.fields, - }) -} - -fn body_json(body: &NamespaceBodyDescriptor) -> Value { - json!({ - "producer": body.producer, - "size": body.size, - "content_type": body.content_type, - }) -} - -fn catalog_json(catalog: &NamespaceQueryCatalog) -> Value { - json!({ - "filterable": grouped_filterable_json(&catalog.filterable), - "sortable": catalog.sortable.iter().map(|field| field.id.clone()).collect::>(), - "facetable": catalog.facetable.iter().map(|field| field.id.clone()).collect::>(), - "facets": catalog.facets.iter().map(facet_summary_json).collect::>(), - }) -} - -fn grouped_filterable_json(capabilities: &[NamespaceFilterCapability]) -> Vec { - let mut groups: Vec<(Vec<&'static str>, Vec)> = Vec::new(); - for capability in capabilities { - let operators = capability - .operators - .iter() - .map(predicate_op_name) - .collect::>(); - match groups.iter_mut().find(|(group, _)| *group == operators) { - Some((_, fields)) => fields.push(capability.field.id.clone()), - None => groups.push((operators, vec![capability.field.id.clone()])), - } - } - groups - .into_iter() - .map(|(operators, fields)| json!({"operators": operators, "fields": fields})) - .collect() -} - -fn facet_summary_json(facet: &NamespaceFacetSummary) -> Value { - json!({ - "field": facet.field.id, - "values": facet.values.iter().map(facet_value_json).collect::>(), - "distinct_count": facet.distinct_count, - "truncated": facet.truncated, - }) -} - -fn facet_value_json(value: &NamespaceFacetValue) -> Value { - json!({ - "value": predicate_value_json(&value.value), - "count": value.count, - }) -} - -fn index_value_json(value: &NamespaceIndexValue) -> Value { - json!({ - "field": value.field.id, - "value": predicate_value_json(&value.value), - }) -} - -fn predicate_value_json(value: &NamespacePredicateValue) -> Value { - match value { - NamespacePredicateValue::String(value) => json!(value), - NamespacePredicateValue::U64(value) => json!(value), - NamespacePredicateValue::F64(value) if value.is_finite() => json!(value), - NamespacePredicateValue::F64(_) => Value::Null, - NamespacePredicateValue::List(values) => { - Value::Array(values.iter().map(predicate_value_json).collect()) - } - } -} - -fn aggregate_result_json(result: &NamespaceAggregateResult) -> Value { - json!({ - "path": result.path, - "input_match_count": result.input_match_count, - "row_count": result.row_count, - "group_count": result.group_count, - "groups": result.groups.iter().map(aggregate_group_json).collect::>(), - "truncated": result.truncated, - }) -} - -fn aggregate_group_json(group: &NamespaceAggregateGroup) -> Value { - let key = group - .key - .iter() - .map(|value| (value.field.id.clone(), predicate_value_json(&value.value))) - .collect::>(); - let values = group - .measures - .iter() - .map(|measure| (measure.name.clone(), aggregate_value_json(&measure.value))) - .collect::>(); - json!({ - "key": key, - "values": values, - }) -} - -fn aggregate_value_json(value: &NamespaceAggregateValue) -> Value { - match value { - NamespaceAggregateValue::U64(value) => json!(value), - NamespaceAggregateValue::F64(value) if value.is_finite() => json!(value), - NamespaceAggregateValue::F64(_) | NamespaceAggregateValue::Null => Value::Null, - } -} - -fn read_item_json(item: &NamespaceReadItem) -> Value { - let value = serde_json::from_str::(&item.value_json) - .unwrap_or_else(|_| Value::String(item.value_json.clone())); - json!({ - "index": item.index, - "value": value, - }) -} - -fn object_args(args: &Value) -> Result<&Map, AgentError> { - args.as_object().ok_or_else(|| { - AgentError::InvalidArgument("agent tool arguments must be a JSON object".to_owned()) - }) -} - -fn required_string_arg<'a>(args: &'a Value, name: &'static str) -> Result<&'a str, AgentError> { - object_args(args)? - .get(name) - .and_then(Value::as_str) - .ok_or_else(|| AgentError::InvalidArgument(format!("missing string argument {name}"))) -} - -fn optional_string_arg(args: &Value, name: &'static str) -> Result, AgentError> { - let Some(value) = object_args(args)?.get(name) else { - return Ok(None); - }; - if value.is_null() { - return Ok(None); - } - value - .as_str() - .map(|value| Some(value.to_owned())) - .ok_or_else(|| AgentError::InvalidArgument(format!("{name} must be a string or null"))) -} - -fn optional_bool_arg(args: &Value, name: &'static str) -> Result, AgentError> { - let Some(value) = object_args(args)?.get(name) else { - return Ok(None); - }; - if value.is_null() { - return Ok(None); - } - value - .as_bool() - .map(Some) - .ok_or_else(|| AgentError::InvalidArgument(format!("{name} must be a boolean or null"))) -} - -fn optional_u64_arg(args: &Value, name: &'static str) -> Result, AgentError> { - let Some(value) = object_args(args)?.get(name) else { - return Ok(None); - }; - if value.is_null() { - return Ok(None); - } - value.as_u64().map(Some).ok_or_else(|| { - AgentError::InvalidArgument(format!("{name} must be a non-negative integer")) - }) -} - -fn optional_usize_arg(args: &Value, name: &'static str) -> Result, AgentError> { - optional_bounded_usize_arg(args, name, MAX_AGENT_PAGE_LIMIT) -} - -fn optional_bounded_usize_arg( - args: &Value, - name: &'static str, - max: usize, -) -> Result, AgentError> { - optional_u64_arg(args, name)? - .map(|value| { - let value = usize::try_from(value).map_err(|_| { - AgentError::InvalidArgument(format!("{name} exceeds platform limit")) - })?; - if value == 0 || value > max { - return Err(AgentError::InvalidArgument(format!( - "{name} must be between 1 and {max}" - ))); - } - Ok(value) - }) - .transpose() -} - -fn read_format_arg(args: &Value) -> Result { - match object_args(args)?.get("format").and_then(Value::as_str) { - None | Some("structured") => Ok(NamespaceReadFormat::Structured), - Some("bytes") => Ok(NamespaceReadFormat::Bytes), - Some(other) => Err(AgentError::InvalidArgument(format!( - "unsupported read format {other}; expected structured or bytes" - ))), - } -} - -fn predicates_arg(args: &Value) -> Result, AgentError> { - let Some(value) = object_args(args)?.get("predicates") else { - return Ok(Vec::new()); - }; - let predicates = value - .as_array() - .ok_or_else(|| AgentError::InvalidArgument("predicates must be an array".to_owned()))?; - predicates.iter().map(predicate_arg).collect() -} - -fn predicate_arg(value: &Value) -> Result { - let object = value - .as_object() - .ok_or_else(|| AgentError::InvalidArgument("predicate must be an object".to_owned()))?; - let field = string_property(object, "field")?; - let op = string_property(object, "op")?; - let op = predicate_op_arg(op)?; - let raw_value = object.get("value").filter(|value| !value.is_null()); - let value = match op { - // Tolerate a stray value on existence checks; agents commonly pass - // {"op": "exists", "value": true} and the value carries no meaning. - NamespacePredicateOp::Exists | NamespacePredicateOp::NotExists => None, - NamespacePredicateOp::In => { - let value = raw_value.ok_or_else(|| { - AgentError::InvalidArgument("predicate op in requires array value".to_owned()) - })?; - if !value.is_array() { - return Err(AgentError::InvalidArgument( - "predicate op in requires array value".to_owned(), - )); - } - Some(predicate_value_arg(value)?) - } - _ => { - let value = raw_value.ok_or_else(|| { - AgentError::InvalidArgument(format!( - "predicate op {} requires value", - predicate_op_name(&op) - )) - })?; - Some(predicate_value_arg(value)?) - } - }; - Ok(NamespacePredicate { - field: NamespaceFindField::new(field), - op, - value, - }) -} - -fn sort_arg(args: &Value) -> Result, AgentError> { - let Some(value) = object_args(args)?.get("sort") else { - return Ok(Vec::new()); - }; - let sort = value - .as_array() - .ok_or_else(|| AgentError::InvalidArgument("sort must be an array".to_owned()))?; - sort.iter().map(sort_item_arg).collect() -} - -fn sort_item_arg(value: &Value) -> Result { - let object = value - .as_object() - .ok_or_else(|| AgentError::InvalidArgument("sort item must be an object".to_owned()))?; - let field = string_property(object, "field")?; - let direction = object - .get("direction") - .and_then(Value::as_str) - .unwrap_or("asc"); - Ok(NamespaceSort { - field: NamespaceSortField::new(field), - direction: match direction { - "asc" => NamespaceSortDirection::Asc, - "desc" => NamespaceSortDirection::Desc, - other => { - return Err(AgentError::InvalidArgument(format!( - "unsupported sort direction {other}" - ))) - } - }, - }) -} - -fn reject_unsupported_find_arguments(args: &Value) -> Result<(), AgentError> { - if object_args(args)?.contains_key("include") { - return Err(AgentError::InvalidArgument( - "unsupported argument include; use stat for schema or sample and read for body content" - .to_owned(), - )); - } - Ok(()) -} - -fn optional_string_array_arg( - args: &Value, - name: &'static str, -) -> Result>, AgentError> { - let Some(value) = object_args(args)?.get(name) else { - return Ok(None); - }; - if value.is_null() { - return Ok(None); - } - let entries = value - .as_array() - .ok_or_else(|| AgentError::InvalidArgument(format!("{name} must be an array")))? - .iter() - .map(|value| { - value.as_str().map(str::to_owned).ok_or_else(|| { - AgentError::InvalidArgument(format!("{name} entries must be strings")) - }) - }) - .collect::, _>>()?; - Ok(Some(entries)) -} - -fn facets_arg(args: &Value) -> Result, AgentError> { - let Some(value) = object_args(args)?.get("facets") else { - return Ok(Vec::new()); - }; - if value.is_null() { - return Ok(Vec::new()); - } - value - .as_array() - .ok_or_else(|| AgentError::InvalidArgument("facets must be an array".to_owned()))? - .iter() - .map(|value| { - value.as_str().map(NamespaceFindField::new).ok_or_else(|| { - AgentError::InvalidArgument("facets entries must be strings".to_owned()) - }) - }) - .collect() -} - -fn group_by_arg(args: &Value) -> Result, AgentError> { - let Some(value) = object_args(args)?.get("group_by") else { - return Ok(Vec::new()); - }; - if value.is_null() { - return Ok(Vec::new()); - } - value - .as_array() - .ok_or_else(|| AgentError::InvalidArgument("group_by must be an array".to_owned()))? - .iter() - .map(|value| { - value.as_str().map(NamespaceFindField::new).ok_or_else(|| { - AgentError::InvalidArgument("group_by entries must be strings".to_owned()) - }) - }) - .collect() -} - -fn aggregate_measures_arg(args: &Value) -> Result, AgentError> { - let value = object_args(args)? - .get("measures") - .ok_or_else(|| AgentError::InvalidArgument("missing array argument measures".to_owned()))?; - let measures = value - .as_array() - .ok_or_else(|| AgentError::InvalidArgument("measures must be an array".to_owned()))?; - if measures.is_empty() { - return Err(AgentError::InvalidArgument( - "measures must contain at least one measure".to_owned(), - )); - } - measures - .iter() - .map(|value| { - let object = value.as_object().ok_or_else(|| { - AgentError::InvalidArgument("measure must be an object".to_owned()) - })?; - let name = string_property(object, "name")?.to_owned(); - if name.is_empty() { - return Err(AgentError::InvalidArgument( - "measure name must not be empty".to_owned(), - )); - } - let op = aggregate_op_arg(string_property(object, "op")?)?; - let field = object - .get("field") - .and_then(|value| (!value.is_null()).then_some(value)) - .map(|value| { - value.as_str().map(str::to_owned).ok_or_else(|| { - AgentError::InvalidArgument( - "measure field must be a string or null".to_owned(), - ) - }) - }) - .transpose()?; - if !matches!(op, NamespaceAggregateOp::Count) && field.is_none() { - return Err(AgentError::InvalidArgument(format!( - "measure {name} with op {} requires field", - aggregate_op_name(&op) - ))); - } - Ok(NamespaceAggregateMeasure { - name, - op, - field: field.map(NamespaceFindField::new), - }) - }) - .collect() -} - -fn aggregate_op_arg(op: &str) -> Result { - match op { - "count" => Ok(NamespaceAggregateOp::Count), - "sum" => Ok(NamespaceAggregateOp::Sum), - "avg" => Ok(NamespaceAggregateOp::Avg), - "min" => Ok(NamespaceAggregateOp::Min), - "max" => Ok(NamespaceAggregateOp::Max), - other => Err(AgentError::InvalidArgument(format!( - "unsupported aggregate op {other}" - ))), - } -} - -fn aggregate_sort_arg(args: &Value) -> Result, AgentError> { - let Some(value) = object_args(args)?.get("sort") else { - return Ok(Vec::new()); - }; - let sort = value - .as_array() - .ok_or_else(|| AgentError::InvalidArgument("sort must be an array".to_owned()))?; - sort.iter() - .map(|value| { - let object = value.as_object().ok_or_else(|| { - AgentError::InvalidArgument("sort item must be an object".to_owned()) - })?; - let field = string_property(object, "field")?.to_owned(); - let direction = object - .get("direction") - .and_then(Value::as_str) - .unwrap_or("asc"); - Ok(NamespaceAggregateSort { - field, - direction: match direction { - "asc" => NamespaceSortDirection::Asc, - "desc" => NamespaceSortDirection::Desc, - other => { - return Err(AgentError::InvalidArgument(format!( - "unsupported sort direction {other}" - ))) - } - }, - }) - }) - .collect() -} - -fn filtered_catalog( - catalog: &NamespaceQueryCatalog, - field_prefix: Option<&str>, - include_facets: bool, -) -> NamespaceQueryCatalog { - NamespaceQueryCatalog { - filterable: catalog - .filterable - .iter() - .filter(|capability| field_matches_prefix(capability.field.id.as_str(), field_prefix)) - .cloned() - .collect(), - sortable: catalog - .sortable - .iter() - .filter(|field| field_matches_prefix(field.id.as_str(), field_prefix)) - .cloned() - .collect(), - facetable: catalog - .facetable - .iter() - .filter(|field| field_matches_prefix(field.id.as_str(), field_prefix)) - .cloned() - .collect(), - facets: if include_facets { - catalog - .facets - .iter() - .filter(|facet| field_matches_prefix(facet.field.id.as_str(), field_prefix)) - .cloned() - .collect() - } else { - Vec::new() - }, - projections: catalog.projections.clone(), - } -} - -fn catalog_is_empty(catalog: &NamespaceQueryCatalog) -> bool { - catalog.filterable.is_empty() && catalog.sortable.is_empty() && catalog.facetable.is_empty() -} - -fn child_catalogs( - namespace: &T, - path: &str, - include_facets: bool, -) -> Result, AgentError> -where - T: AgentNamespace + ?Sized, -{ - let page = namespace.agent_list_page( - path, - NamespaceListOptions { - cursor: None, - limit: 20, - }, - )?; - let mut children = Vec::new(); - for entry in page.entries { - if !matches!(entry.kind, NamespaceCardKind::Directory) { - continue; - } - let Some(card) = namespace.agent_stat_card(&entry.path)? else { - continue; - }; - let catalog = filtered_catalog(&card.catalog, None, include_facets); - if catalog_is_empty(&catalog) { - continue; - } - children.push(json!({ - "path": card.path, - "catalog": catalog_json(&catalog), - })); - if children.len() == 5 { - break; - } - } - Ok(children) -} - -fn field_matches_prefix(field: &str, prefix: Option<&str>) -> bool { - prefix - .map(|prefix| field.starts_with(prefix) || field.contains(prefix)) - .unwrap_or(true) -} - -fn aggregate_op_name(op: &NamespaceAggregateOp) -> &'static str { - match op { - NamespaceAggregateOp::Count => "count", - NamespaceAggregateOp::Sum => "sum", - NamespaceAggregateOp::Avg => "avg", - NamespaceAggregateOp::Min => "min", - NamespaceAggregateOp::Max => "max", - } -} - -fn string_property<'a>( - object: &'a Map, - name: &'static str, -) -> Result<&'a str, AgentError> { - object - .get(name) - .and_then(Value::as_str) - .ok_or_else(|| AgentError::InvalidArgument(format!("missing string property {name}"))) -} - -fn predicate_op_arg(op: &str) -> Result { - match op { - "eq" => Ok(NamespacePredicateOp::Eq), - "ne" | "not_equal" => Ok(NamespacePredicateOp::NotEqual), - "in" => Ok(NamespacePredicateOp::In), - "prefix" => Ok(NamespacePredicateOp::Prefix), - "suffix" => Ok(NamespacePredicateOp::Suffix), - "contains" => Ok(NamespacePredicateOp::Contains), - "gt" | "greater_than" => Ok(NamespacePredicateOp::GreaterThan), - "gte" | "greater_than_or_equal" => Ok(NamespacePredicateOp::GreaterThanOrEqual), - "lt" | "less_than" => Ok(NamespacePredicateOp::LessThan), - "lte" | "less_than_or_equal" => Ok(NamespacePredicateOp::LessThanOrEqual), - "exists" => Ok(NamespacePredicateOp::Exists), - "not_exists" => Ok(NamespacePredicateOp::NotExists), - other => Err(AgentError::InvalidArgument(format!( - "unsupported predicate operator {other}" - ))), - } -} - -fn predicate_value_arg(value: &Value) -> Result { - if let Some(values) = value.as_array() { - return values - .iter() - .map(predicate_scalar_value_arg) - .collect::, _>>() - .map(NamespacePredicateValue::List); - } - predicate_scalar_value_arg(value) -} - -fn predicate_scalar_value_arg(value: &Value) -> Result { - if let Some(value) = value.as_str() { - return Ok(NamespacePredicateValue::String(value.to_owned())); - } - if let Some(value) = value.as_bool() { - return Ok(NamespacePredicateValue::U64(u64::from(value))); - } - if let Some(value) = value.as_u64() { - return Ok(NamespacePredicateValue::U64(value)); - } - if let Some(value) = value.as_f64() { - if value.is_finite() { - return Ok(NamespacePredicateValue::F64(value)); - } - } - Err(AgentError::InvalidArgument( - "predicate value must be a string, boolean, finite number, or array of scalar values" - .to_owned(), - )) -} - -fn card_kind_name(kind: &NamespaceCardKind) -> &'static str { - match kind { - NamespaceCardKind::File => "file", - NamespaceCardKind::Directory => "directory", - NamespaceCardKind::Symlink => "symlink", - NamespaceCardKind::Special => "special", - } -} - -fn record_type_name(record_type: &NamespaceRecordType) -> &'static str { - match record_type { - NamespaceRecordType::DirectoryEntries => "directory_entries", - NamespaceRecordType::JsonArray => "json_array", - NamespaceRecordType::JsonObject => "json_object", - NamespaceRecordType::YamlMapping => "yaml_mapping", - NamespaceRecordType::TextLines => "text_lines", - } -} - -fn read_format_name(format: &NamespaceReadFormat) -> &'static str { - match format { - NamespaceReadFormat::Structured => "structured", - NamespaceReadFormat::Bytes => "bytes", - } -} - -fn predicate_op_name(op: &NamespacePredicateOp) -> &'static str { - match op { - NamespacePredicateOp::Eq => "eq", - NamespacePredicateOp::NotEqual => "ne", - NamespacePredicateOp::In => "in", - NamespacePredicateOp::Prefix => "prefix", - NamespacePredicateOp::Suffix => "suffix", - NamespacePredicateOp::Contains => "contains", - NamespacePredicateOp::GreaterThan => "gt", - NamespacePredicateOp::GreaterThanOrEqual => "gte", - NamespacePredicateOp::LessThan => "lt", - NamespacePredicateOp::LessThanOrEqual => "lte", - NamespacePredicateOp::Exists => "exists", - NamespacePredicateOp::NotExists => "not_exists", - } -} - -#[cfg(test)] -mod tests { - use super::*; - use nokv_meta::{ - NamespaceAggregateOutputMeasure, NamespaceFieldSource, NamespaceFieldSourceKind, - NamespaceFieldValue, NamespaceGrepMatch, NamespaceInclude, NamespaceRecordCount, - RecordCountProvenance, - }; - use nokv_types::InodeId; - use std::cell::RefCell; - use std::collections::BTreeMap; - - struct FakeNamespace { - last_find: RefCell>, - last_aggregate: RefCell>, - last_grep: RefCell>, - last_read: RefCell>, - read_calls: RefCell, - record_count: usize, - find_matches: Vec, - aggregate_result: NamespaceAggregateResult, - stat_cards: BTreeMap, - list_entries: Vec, - } - - impl FakeNamespace { - fn new() -> Self { - Self { - last_find: RefCell::new(None), - last_aggregate: RefCell::new(None), - last_grep: RefCell::new(None), - last_read: RefCell::new(None), - read_calls: RefCell::new(0), - record_count: 1, - find_matches: vec![sample_card("/runs/run-1", 1)], - aggregate_result: sample_aggregate_result(), - stat_cards: BTreeMap::new(), - list_entries: vec![sample_card("/runs/run-1", 1)], - } - } - - fn with_record_count(record_count: usize) -> Self { - Self { - record_count, - find_matches: vec![sample_card("/runs/run-1", record_count)], - list_entries: vec![sample_card("/runs/run-1", record_count)], - ..Self::new() - } - } - - fn with_aggregate_result(aggregate_result: NamespaceAggregateResult) -> Self { - Self { - aggregate_result, - ..Self::new() - } - } - - fn with_root_child_catalog() -> Self { - let mut root = sample_card("/yanex", 1); - root.kind = NamespaceCardKind::Directory; - root.catalog = empty_catalog(); - let mut runs = sample_card("/yanex/runs", 1); - runs.kind = NamespaceCardKind::Directory; - runs.catalog.filterable.push(NamespaceFilterCapability { - field: NamespaceFindField::new("run.script"), - operators: vec![NamespacePredicateOp::Eq], - }); - runs.catalog - .sortable - .push(NamespaceSortField::new("run.script")); - let mut stat_cards = BTreeMap::new(); - stat_cards.insert(root.path.clone(), root); - stat_cards.insert(runs.path.clone(), runs.clone()); - Self { - find_matches: vec![sample_card("/yanex/runs/run-1", 1)], - stat_cards, - list_entries: vec![runs], - ..Self::new() - } - } - } - - impl AgentNamespace for FakeNamespace { - fn agent_stat_card(&self, path: &str) -> Result, AgentError> { - Ok(Some( - self.stat_cards - .get(path) - .cloned() - .unwrap_or_else(|| sample_card(path, self.record_count)), - )) - } - - fn agent_list_page( - &self, - path: &str, - _options: NamespaceListOptions, - ) -> Result { - Ok(NamespaceListPage { - path: path.to_owned(), - evidence: "nokv-native:///runs".to_owned(), - snapshot_id: Some(9), - entry_count: self.list_entries.len(), - entries: self.list_entries.clone(), - next_cursor: None, - truncated: false, - }) - } - - fn agent_find_paths( - &self, - request: NamespaceFindRequest, - ) -> Result { - self.last_find.replace(Some(request)); - let matches = self.find_matches.clone(); - Ok(NamespaceFindResult { - path: "/runs".to_owned(), - evidence: "nokv-native:///runs".to_owned(), - snapshot_id: Some(9), - match_count: matches.len(), - matches, - facets: vec![NamespaceFacetSummary { - field: NamespaceFindField::new("run.script"), - values: vec![NamespaceFacetValue { - value: NamespacePredicateValue::String("train.py".to_owned()), - count: 1, - }], - distinct_count: 1, - truncated: false, - }], - next_cursor: None, - truncated: false, - scanned_entries: 1, - }) - } - - fn agent_aggregate_paths( - &self, - request: NamespaceAggregateRequest, - ) -> Result { - let mut result = self.aggregate_result.clone(); - result.path = request.path.clone(); - result.predicates = request.predicates.clone(); - self.last_aggregate.replace(Some(request)); - Ok(result) - } - - fn agent_grep_paths( - &self, - request: NamespaceGrepRequest, - ) -> Result { - self.last_grep.replace(Some(request.clone())); - let matched = "Checkpoint: best_model_1.0_1.pt"; - let matches = if matched - .to_lowercase() - .contains(&request.pattern.to_lowercase()) - { - vec![NamespaceGrepMatch { - path: format!("{}/artifacts/stdout.txt", request.path), - line_number: 3, - snippet: matched.to_owned(), - evidence: format!( - "nokv-native://{}/artifacts/stdout.txt@generation:1#L3", - request.path - ), - generation: 1, - }] - } else { - Vec::new() - }; - Ok(NamespaceGrepResult { - path: request.path.clone(), - pattern: request.pattern.clone(), - patterns: request.patterns.clone(), - recursive: request.recursive, - evidence: format!("nokv-native://{}", request.path), - snapshot_id: Some(1), - matches, - files_scanned: 1, - bytes_read: matched.len(), - next_cursor: None, - truncated: false, - }) - } - - fn agent_read_page( - &self, - path: &str, - options: NamespaceReadOptions, - ) -> Result { - self.last_read.replace(Some(options)); - *self.read_calls.borrow_mut() += 1; - Ok(NamespaceReadPage { - path: path.to_owned(), - evidence: "nokv-native:///runs/run-1/metadata.json@generation:7".to_owned(), - snapshot_id: Some(9), - generation: 7, - total_size_bytes: 13, - format: NamespaceReadFormat::Structured, - record_type: Some(NamespaceRecordType::JsonObject), - record_count: Some(1), - cursor: None, - next_cursor: None, - truncated: false, - items: vec![NamespaceReadItem { - index: 0, - value_json: r#"{"status":"completed"}"#.to_owned(), - evidence: "nokv-native:///runs/run-1/metadata.json@generation:7#item:0" - .to_owned(), - }], - bytes: None, - }) - } - } - - fn tool_definition<'a>( - tools: &'a [AgentToolDefinition], - name: &str, - ) -> &'a AgentToolDefinition { - tools - .iter() - .find(|tool| tool.name == name) - .unwrap_or_else(|| panic!("{name} tool must be registered")) - } - - fn assert_tool_description_contains( - tools: &[AgentToolDefinition], - name: &str, - expected: &[&str], - ) { - let description = tool_definition(tools, name).description; - let normalized_description = description.to_ascii_lowercase(); - for snippet in expected { - let normalized_snippet = snippet.to_ascii_lowercase(); - assert!( - normalized_description.contains(&normalized_snippet), - "{name} description must contain {snippet:?}; got {description:?}" - ); - } - } - - fn assert_json_lacks_agent_noise(value: &Value) { - // "generation" is intentionally not forbidden: read exposes it so agents - // can feed the conditional read validator (expected_generation). - const FORBIDDEN_KEYS: &[&str] = &[ - "tool", - "bytes_read", - "evidence", - "snapshot_id", - "field_values", - "source_path", - "source_kind", - "key_fields", - "measures", - "sample_matches", - "scope_note", - "predicates", - "scanned_entries", - "digest_uri", - "manifest_id", - "chunk_size", - "block_size", - "inode", - "provenance", - ]; - match value { - Value::Object(object) => { - for forbidden in FORBIDDEN_KEYS { - assert!( - !object.contains_key(*forbidden), - "agent output must not contain noisy key {forbidden:?}: {value}" - ); - } - for child in object.values() { - assert_json_lacks_agent_noise(child); - } - } - Value::Array(values) => { - for child in values { - assert_json_lacks_agent_noise(child); - } - } - _ => {} - } - } - - fn assert_schema_lacks_nested_descriptions(value: &Value) { - match value { - Value::Object(object) => { - assert!( - !object.contains_key("description"), - "tool parameter schema must not contain nested descriptions: {value}" - ); - for child in object.values() { - assert_schema_lacks_nested_descriptions(child); - } - } - Value::Array(values) => { - for child in values { - assert_schema_lacks_nested_descriptions(child); - } - } - _ => {} - } - } - - fn empty_catalog() -> NamespaceQueryCatalog { - NamespaceQueryCatalog { - filterable: Vec::new(), - sortable: Vec::new(), - facetable: Vec::new(), - facets: Vec::new(), - projections: Vec::new(), - } - } - - fn sample_field_value(field: &str, value: NamespacePredicateValue) -> NamespaceFieldValue { - NamespaceFieldValue { - field: NamespaceFindField::new(field), - value, - source: NamespaceFieldSource { - evidence: "nokv-native:///runs/run-1@generation:7".to_owned(), - source_path: "/runs/run-1".to_owned(), - source_kind: NamespaceFieldSourceKind::MaterializedIndex, - }, - } - } - - fn sample_aggregate_result() -> NamespaceAggregateResult { - NamespaceAggregateResult { - path: "/runs".to_owned(), - evidence: "nokv-native:///runs".to_owned(), - snapshot_id: Some(9), - predicates: Vec::new(), - input_match_count: 5, - row_count: 3, - group_count: 2, - groups: vec![NamespaceAggregateGroup { - key: vec![sample_field_value( - "param.lr", - NamespacePredicateValue::String("0.001".to_owned()), - )], - measures: vec![ - NamespaceAggregateOutputMeasure { - name: "run_count".to_owned(), - op: NamespaceAggregateOp::Count, - field: Some(NamespaceFindField::new("metric.val_loss.min")), - value: NamespaceAggregateValue::U64(2), - }, - NamespaceAggregateOutputMeasure { - name: "avg_min_val_loss".to_owned(), - op: NamespaceAggregateOp::Avg, - field: Some(NamespaceFindField::new("metric.val_loss.min")), - value: NamespaceAggregateValue::F64(0.3), - }, - NamespaceAggregateOutputMeasure { - name: "stdout_available".to_owned(), - op: NamespaceAggregateOp::Sum, - field: Some(NamespaceFindField::new("artifact.stdout_available")), - value: NamespaceAggregateValue::U64(1), - }, - ], - evidence: "nokv-native:///runs".to_owned(), - sample_matches: vec![nokv_meta::NamespaceAggregateSample { - path: "/runs/run-1".to_owned(), - evidence: "nokv-native:///runs/run-1@generation:7".to_owned(), - generation: 7, - }], - }], - truncated: false, - scanned_entries: 5, - } - } - - fn sample_card(path: &str, record_count: usize) -> NamespaceCard { - NamespaceCard { - path: path.to_owned(), - name: path.rsplit('/').next().unwrap_or("").to_owned(), - kind: NamespaceCardKind::File, - evidence: format!("nokv-native://{path}@generation:7"), - snapshot_id: Some(9), - inode: InodeId::new(42).unwrap(), - generation: 7, - size_bytes: Some(13), - entry_count: None, - record_count: Some(NamespaceRecordCount { - count: record_count, - provenance: RecordCountProvenance::MaterializedIndex, - }), - schema: Some(NamespaceSchema { - record_type: NamespaceRecordType::JsonObject, - fields: vec!["status".to_owned()], - }), - sample: vec![r#"{"status":"completed"}"#.to_owned()], - body: Some(NamespaceBodyDescriptor { - producer: "unit-test".to_owned(), - digest_uri: "sha256:test".to_owned(), - size: 13, - content_type: "application/json".to_owned(), - manifest_id: "runs/run-1/metadata.json".to_owned(), - generation: 7, - chunk_size: 4096, - block_size: 4096, - }), - catalog: NamespaceQueryCatalog { - filterable: vec![NamespaceFilterCapability { - field: NamespaceFindField::new("run.status"), - operators: vec![NamespacePredicateOp::Eq], - }], - sortable: vec![NamespaceSortField::new("run.status")], - facetable: vec![NamespaceFindField::new("run.status")], - facets: vec![NamespaceFacetSummary { - field: NamespaceFindField::new("run.status"), - values: vec![NamespaceFacetValue { - value: NamespacePredicateValue::String("completed".to_owned()), - count: 1, - }], - distinct_count: 1, - truncated: false, - }], - projections: vec![ - NamespaceInclude::Body, - NamespaceInclude::Schema, - NamespaceInclude::Sample, - NamespaceInclude::Catalog, - ], - }, - indexed_values: vec![NamespaceIndexValue { - field: NamespaceFindField::new("run.status"), - value: NamespacePredicateValue::String("completed".to_owned()), - }], - } - } - - #[test] - fn agent_tool_registry_uses_posix_like_names() { - let names = agent_tool_definitions() - .into_iter() - .map(|tool| tool.name) - .collect::>(); - - assert_eq!( - names, - vec!["ls", "stat", "catalog", "read", "aggregate", "find", "grep"] - ); - } - - #[test] - fn find_tool_schema_excludes_include_and_caps_limit_at_ten() { - let tools = agent_tool_definitions(); - let find = tools - .iter() - .find(|tool| tool.name == "find") - .expect("find tool must be registered"); - - assert!(find.parameters["properties"].get("include").is_none()); - assert_eq!(find.parameters["properties"]["limit"]["maximum"], 10); - } +pub mod event; +mod namespace; - #[test] - fn agent_tool_registry_documents_stable_api_semantics() { - let tools = agent_tool_definitions(); - - assert_tool_description_contains(&tools, "ls", &["direct children", "discover paths"]); - assert_tool_description_contains( - &tools, - "stat", - &[ - "single path", - "compact card", - "body", - "schema", - "sample", - "catalog", - ], - ); - assert_tool_description_contains(&tools, "catalog", &["field ids", "find", "aggregate"]); - assert_tool_description_contains(&tools, "read", &["file body", "json", "yaml", "text"]); - assert_tool_description_contains( - &tools, - "aggregate", - &["catalog field ids", "count", "group", "summary"], - ); - assert_tool_description_contains( - &tools, - "find", - &["search paths", "project fields", "read"], - ); - } - - #[test] - fn agent_tool_parameter_schemas_are_token_compact() { - let tools = agent_tool_definitions(); - let find = tool_definition(&tools, "find"); - let read = tool_definition(&tools, "read"); - - for tool in &tools { - assert_schema_lacks_nested_descriptions(&tool.parameters); - } - assert!(find.parameters["properties"].get("include").is_none()); - assert!(read.parameters["properties"] - .get("expected_generation") - .is_none()); - } - - #[test] - fn stat_tool_returns_compact_card_without_agent_noise() { - let namespace = FakeNamespace::new(); - - let output = execute_agent_tool( - &namespace, - "stat", - &json!({ - "path": "/runs", - }), - ) - .unwrap(); - - assert_json_lacks_agent_noise(&output); - assert_eq!(output["card"]["path"], "/runs"); - assert_eq!( - output["card"]["body"], - json!({ - "producer": "unit-test", - "size": 13, - "content_type": "application/json" - }) - ); - assert_eq!( - output["card"]["catalog"]["filterable"][0]["fields"][0], - "run.status" - ); - assert!(output["card"]["catalog"].get("projections").is_none()); - assert_eq!( - output["card"]["catalog"]["facets"][0], - json!({ - "field": "run.status", - "values": [{"value": "completed", "count": 1}], - "distinct_count": 1, - "truncated": false - }) - ); - assert_eq!(output["card"]["record_count"], 1); - } - - #[test] - fn catalog_tool_returns_compact_filtered_catalog() { - let namespace = FakeNamespace::new(); - - let output = execute_agent_tool( - &namespace, - "catalog", - &json!({ - "path": "/runs", - "field_prefix": "run.", - "include_facets": false - }), - ) - .unwrap(); - - assert_json_lacks_agent_noise(&output); - assert_eq!(output["path"], "/runs"); - assert_eq!( - output["catalog"]["filterable"][0]["fields"], - json!(["run.status"]) - ); - assert_eq!(output["catalog"]["sortable"], json!(["run.status"])); - assert_eq!(output["catalog"]["facetable"], json!(["run.status"])); - assert_eq!(output["catalog"]["facets"], json!([])); - assert!(output.get("card").is_none()); - assert!(output.get("body").is_none()); - } - - #[test] - fn catalog_tool_matches_field_prefix_inside_field_ids() { - let namespace = FakeNamespace::new(); - - let output = execute_agent_tool( - &namespace, - "catalog", - &json!({ - "path": "/runs", - "field_prefix": "status", - "include_facets": false - }), - ) - .unwrap(); - - assert_eq!( - output["catalog"]["filterable"][0]["fields"][0], - "run.status" - ); - assert_eq!(output["catalog"]["sortable"], json!(["run.status"])); - } - - #[test] - fn catalog_tool_suggests_child_catalogs_when_current_scope_is_empty() { - let namespace = FakeNamespace::with_root_child_catalog(); - - let output = execute_agent_tool( - &namespace, - "catalog", - &json!({ - "path": "/yanex", - "field_prefix": "status", - "include_facets": false - }), - ) - .unwrap(); - - assert_eq!(output["catalog_empty"], true); - assert_eq!(output["child_catalogs"][0]["path"], "/yanex/runs"); - let child_filterable = &output["child_catalogs"][0]["catalog"]["filterable"]; - let child_fields = child_filterable - .as_array() - .unwrap() - .iter() - .flat_map(|group| group["fields"].as_array().unwrap().clone()) - .collect::>(); - assert!(child_fields.contains(&json!("run.status"))); - assert!(child_fields.contains(&json!("run.script"))); - assert_json_lacks_agent_noise(&output); - } - - #[test] - fn catalog_tool_does_not_discover_child_catalogs_for_files() { - let namespace = FakeNamespace::new(); - - let output = execute_agent_tool( - &namespace, - "catalog", - &json!({ - "path": "/runs/run-1", - "field_prefix": "missing.", - "include_facets": false - }), - ) - .unwrap(); - - assert_eq!(output["catalog_empty"], true); - assert_eq!(output["child_catalogs"], json!([])); - assert_json_lacks_agent_noise(&output); - } - - #[test] - fn find_tool_translates_catalog_field_ids() { - let namespace = FakeNamespace::new(); - - let output = execute_agent_tool( - &namespace, - "find", - &json!({ - "path": "/runs", - "predicates": [{"field": "run.status", "op": "eq", "value": "completed"}], - "sort": [{"field": "run.status", "direction": "desc"}], - "limit": 5 - }), - ) - .unwrap(); - - let request = namespace.last_find.borrow().clone().unwrap(); - assert_eq!(request.predicates[0].field.id, "run.status"); - assert_eq!( - request.predicates[0].value, - Some(NamespacePredicateValue::String("completed".to_owned())) - ); - assert_eq!(request.sort[0].field.id, "run.status"); - assert_eq!(request.include, Vec::::new()); - assert_eq!(output["match_count"], 1); - assert_eq!(output["matches"][0], json!({"path": "/runs/run-1"})); - assert!(output.get("catalog").is_none() || output["catalog"].is_null()); - assert!(output["matches"][0].get("catalog").is_none()); - assert_json_lacks_agent_noise(&output); - } - - #[test] - fn find_tool_requests_and_returns_filtered_facets() { - let namespace = FakeNamespace::new(); - - let output = execute_agent_tool( - &namespace, - "find", - &json!({ - "path": "/runs", - "predicates": [{"field": "run.status", "op": "eq", "value": "completed"}], - "facets": ["run.script"], - "limit": 1 - }), - ) - .unwrap(); - - let request = namespace.last_find.borrow().clone().unwrap(); - assert_eq!(request.facets, vec![NamespaceFindField::new("run.script")]); - assert_eq!( - output["facets"], - json!([{ - "field": "run.script", - "values": [{"value": "train.py", "count": 1}], - "distinct_count": 1, - "truncated": false - }]) - ); - } - - #[test] - fn find_tool_schema_rejects_body_schema_sample_includes() { - let namespace = FakeNamespace::new(); - - let err = execute_agent_tool( - &namespace, - "find", - &json!({ - "path": "/runs", - "include": ["body", "schema", "sample"], - "limit": 1 - }), - ) - .unwrap_err(); - - assert!( - matches!(err, AgentError::InvalidArgument(ref message) if message.contains("unsupported argument include")), - "unexpected error: {err:?}" - ); - assert!(namespace.last_find.borrow().is_none()); - } - - #[test] - fn find_tool_rejects_catalog_include() { - let namespace = FakeNamespace::new(); - - let err = execute_agent_tool( - &namespace, - "find", - &json!({ - "path": "/runs", - "include": ["catalog"], - "limit": 1 - }), - ) - .unwrap_err(); - - assert!( - matches!(err, AgentError::InvalidArgument(ref message) if message.contains("unsupported argument include")), - "unexpected error: {err:?}" - ); - assert!(namespace.last_find.borrow().is_none()); - } - - #[test] - fn find_tool_defaults_to_limit_ten() { - let namespace = FakeNamespace::new(); - - execute_agent_tool( - &namespace, - "find", - &json!({ - "path": "/runs", - }), - ) - .unwrap(); - - let request = namespace.last_find.borrow().clone().unwrap(); - assert_eq!(request.limit, 10); - } - - #[test] - fn find_tool_rejects_limit_above_ten() { - let namespace = FakeNamespace::new(); - - let err = execute_agent_tool( - &namespace, - "find", - &json!({ - "path": "/runs", - "limit": 11 - }), - ) - .unwrap_err(); - - assert!( - matches!(err, AgentError::InvalidArgument(ref message) if message.contains("limit must be between 1 and 10")), - "unexpected error: {err:?}" - ); - assert!(namespace.last_find.borrow().is_none()); - } - - #[test] - fn find_tool_accepts_boolean_predicates_as_u64_facets() { - let namespace = FakeNamespace::new(); - - execute_agent_tool( - &namespace, - "find", - &json!({ - "path": "/runs", - "predicates": [{"field": "git.has_uncommitted_changes", "op": "eq", "value": true}], - "limit": 5 - }), - ) - .unwrap(); - - let request = namespace.last_find.borrow().clone().unwrap(); - assert_eq!( - request.predicates[0].field.id, - "git.has_uncommitted_changes" - ); - assert_eq!( - request.predicates[0].value, - Some(NamespacePredicateValue::U64(1)) - ); - } - - #[test] - fn find_tool_projects_indexed_values_when_fields_are_requested() { - let namespace = FakeNamespace::new(); - - let output = execute_agent_tool( - &namespace, - "find", - &json!({ - "path": "/runs", - "fields": ["run.id"], - "limit": 5 - }), - ) - .unwrap(); - - assert_eq!(output["matches"][0]["path"], "/runs/run-1"); - assert_eq!(output["matches"][0]["values"], json!({})); - assert!(output["matches"][0].get("indexed_values").is_none()); - assert!(output["matches"][0].get("name").is_none()); - assert!(output.get("catalog").is_none() || output["catalog"].is_null()); - assert!(namespace.last_find.borrow().is_some()); - assert_json_lacks_agent_noise(&output); - } - - #[test] - fn find_tool_projects_requested_fields_without_sources() { - let namespace = FakeNamespace::new(); - - let output = execute_agent_tool( - &namespace, - "find", - &json!({ - "path": "/runs", - "fields": ["run.status"], - "limit": 5 - }), - ) - .unwrap(); - - let match_ = &output["matches"][0]; - assert_eq!(match_["path"], "/runs/run-1"); - assert_eq!(match_["values"], json!({"run.status": "completed"})); - assert!(match_.get("indexed_values").is_none()); - assert!(match_.get("name").is_none()); - assert_json_lacks_agent_noise(&output); - } - - #[test] - fn find_tool_does_not_project_storage_internal_body_fields() { - let namespace = FakeNamespace::new(); - - let output = execute_agent_tool( - &namespace, - "find", - &json!({ - "path": "/runs", - "fields": ["body.manifest_id"], - "limit": 5 - }), - ) - .unwrap(); - - let match_ = &output["matches"][0]; - assert_eq!(match_["path"], "/runs/run-1"); - assert_eq!(match_["values"], json!({})); - assert_json_lacks_agent_noise(&output); - } - - #[test] - fn grep_tool_returns_compact_matches_without_agent_noise() { - let namespace = FakeNamespace::new(); - - let output = execute_agent_tool( - &namespace, - "grep", - &json!({ - "path": "/runs/run-1", - "pattern": "best_model", - "recursive": true, - "limit": 5 - }), - ) - .unwrap(); - - assert_eq!(output["path"], "/runs/run-1"); - assert_eq!(output["matches"][0]["line_number"], 3); - assert_eq!( - output["matches"][0]["snippet"], - "Checkpoint: best_model_1.0_1.pt" - ); - assert!(output.get("tool").is_none()); - assert!(output.get("bytes_read").is_none()); - assert!(output["matches"][0].get("evidence").is_none()); - assert_json_lacks_agent_noise(&output); - } - - #[test] - fn aggregate_tool_groups_and_sorts_indexed_values() { - let namespace = FakeNamespace::with_aggregate_result(sample_aggregate_result()); - - let output = execute_agent_tool( - &namespace, - "aggregate", - &json!({ - "path": "/runs", - "predicates": [{"field": "run.status", "op": "eq", "value": "completed"}], - "group_by": ["param.lr"], - "measures": [ - {"name": "run_count", "op": "count", "field": "metric.val_loss.min"}, - {"name": "avg_min_val_loss", "op": "avg", "field": "metric.val_loss.min"}, - {"name": "stdout_available", "op": "sum", "field": "artifact.stdout_available"} - ], - "sort": [{"field": "avg_min_val_loss", "direction": "asc"}], - "limit": 5 - }), - ) - .unwrap(); - - let request = namespace.last_aggregate.borrow().clone().unwrap(); - assert_eq!(request.predicates[0].field.id, "run.status"); - assert!(namespace.last_find.borrow().is_none()); - assert_eq!(output["input_match_count"], 5); - assert_eq!(output["row_count"], 3); - assert_eq!(output["group_count"], 2); - assert_eq!(output["groups"][0]["key"], json!({"param.lr": "0.001"})); - assert_eq!(output["groups"][0]["values"]["run_count"], 2); - assert_eq!(output["groups"][0]["values"]["avg_min_val_loss"], 0.3); - assert_eq!(output["groups"][0]["values"]["stdout_available"], 1); - assert_json_lacks_agent_noise(&output); - } - - #[test] - fn ls_tool_keeps_entries_lightweight() { - let namespace = FakeNamespace::new(); - - let output = execute_agent_tool( - &namespace, - "ls", - &json!({ - "path": "/runs", - "limit": 1 - }), - ) - .unwrap(); - - assert_eq!(output["entries"][0]["path"], "/runs/run-1"); - assert!(output["entries"][0].get("record_count").is_none()); - assert!(output["entries"][0].get("sample").is_none()); - assert!(output["entries"][0].get("catalog").is_none()); - assert_json_lacks_agent_noise(&output); - } - - #[test] - fn read_tool_returns_body_content_without_evidence_noise() { - let namespace = FakeNamespace::new(); - - let output = execute_agent_tool( - &namespace, - "read", - &json!({ - "path": "/runs/run-1/metadata.json", - "format": "structured", - "limit": 10 - }), - ) - .unwrap(); - - assert_eq!(output["path"], "/runs/run-1/metadata.json"); - assert_eq!(output["items"][0]["index"], 0); - assert_eq!(output["items"][0]["value"]["status"], "completed"); - assert_json_lacks_agent_noise(&output); - } - - #[test] - fn read_tool_rejects_large_structured_pagination() { - let namespace = FakeNamespace::with_record_count(350); - - let err = execute_agent_tool( - &namespace, - "read", - &json!({ - "path": "/index/large.json", - "format": "structured", - "cursor": "100", - "limit": 100 - }), - ) - .unwrap_err(); - - assert!( - matches!(err, AgentError::InvalidArgument(ref message) if message.contains("use bytes format with offset and limit, grep to locate lines, or stat record_count")), - "unexpected error: {err:?}" - ); - assert_eq!(*namespace.read_calls.borrow(), 0); - } - - #[test] - fn read_tool_rejects_large_structured_initial_page() { - let namespace = FakeNamespace::with_record_count(350); - - let err = execute_agent_tool( - &namespace, - "read", - &json!({ - "path": "/index/large.json", - "format": "structured", - "limit": 100 - }), - ) - .unwrap_err(); - - assert!( - matches!(err, AgentError::InvalidArgument(ref message) if message.contains("use bytes format with offset and limit, grep to locate lines, or stat record_count")), - "unexpected error: {err:?}" - ); - assert_eq!(*namespace.read_calls.borrow(), 0); - } - - #[test] - fn read_tool_reads_structured_files_within_guard_threshold() { - let namespace = FakeNamespace::with_record_count(250); - - execute_agent_tool( - &namespace, - "read", - &json!({ - "path": "/index/mid.json", - "format": "structured" - }), - ) - .unwrap(); - - assert_eq!(*namespace.read_calls.borrow(), 1); - } - - #[test] - fn read_tool_structured_guard_allows_exactly_the_read_limit() { - let namespace = FakeNamespace::with_record_count(MAX_AGENT_READ_LIMIT); - - execute_agent_tool( - &namespace, - "read", - &json!({ - "path": "/index/boundary.json", - "format": "structured" - }), - ) - .unwrap(); - - assert_eq!(*namespace.read_calls.borrow(), 1); - } - - #[test] - fn read_tool_returns_generation_for_conditional_reads() { - let namespace = FakeNamespace::new(); - - let output = execute_agent_tool( - &namespace, - "read", - &json!({ - "path": "/runs/run-1/metadata.json", - }), - ) - .unwrap(); - - assert_eq!(output["generation"], 7); - assert_json_lacks_agent_noise(&output); - } - - #[test] - fn read_tool_schema_caps_limit_at_three_hundred() { - let tools = agent_tool_definitions(); - let read = tool_definition(&tools, "read"); - - assert_eq!(read.parameters["properties"]["limit"]["maximum"], 300); - } - - #[test] - fn read_tool_accepts_limit_of_three_hundred() { - let namespace = FakeNamespace::new(); - - execute_agent_tool( - &namespace, - "read", - &json!({ - "path": "/runs/run-1/metadata.json", - "limit": 300 - }), - ) - .unwrap(); - - let options = namespace.last_read.borrow().clone().unwrap(); - assert_eq!(options.limit, 300); - } - - #[test] - fn read_tool_rejects_limit_above_three_hundred() { - let namespace = FakeNamespace::new(); - - let err = execute_agent_tool( - &namespace, - "read", - &json!({ - "path": "/runs/run-1/metadata.json", - "limit": 301 - }), - ) - .unwrap_err(); - - assert!( - matches!(err, AgentError::InvalidArgument(ref message) if message.contains("limit must be between 1 and 300")), - "unexpected error: {err:?}" - ); - assert_eq!(*namespace.read_calls.borrow(), 0); - } - - #[test] - fn ls_tool_keeps_limit_capped_at_one_hundred() { - let namespace = FakeNamespace::new(); - - let err = execute_agent_tool( - &namespace, - "ls", - &json!({ - "path": "/runs", - "limit": 101 - }), - ) - .unwrap_err(); - - assert!( - matches!(err, AgentError::InvalidArgument(ref message) if message.contains("limit must be between 1 and 100")), - "unexpected error: {err:?}" - ); - } - - #[test] - fn grep_tool_schema_caps_limit_and_exposes_patterns_and_glob() { - let tools = agent_tool_definitions(); - let grep = tool_definition(&tools, "grep"); - - assert_eq!(grep.parameters["properties"]["limit"]["maximum"], 300); - assert_eq!( - grep.parameters["properties"]["patterns"], - json!({"type": "array", "items": {"type": "string"}, "maxItems": 16}) - ); - assert_eq!( - grep.parameters["properties"]["glob"], - json!({"type": ["string", "null"]}) - ); - assert_eq!( - grep.parameters["required"], - json!(["path", "pattern", "recursive"]) - ); - } - - #[test] - fn grep_tool_accepts_limit_of_three_hundred() { - let namespace = FakeNamespace::new(); - - execute_agent_tool( - &namespace, - "grep", - &json!({ - "path": "/runs/run-1", - "pattern": "best_model", - "recursive": true, - "limit": 300 - }), - ) - .unwrap(); - - let request = namespace.last_grep.borrow().clone().unwrap(); - assert_eq!(request.limit, 300); - } - - #[test] - fn grep_tool_rejects_limit_above_three_hundred() { - let namespace = FakeNamespace::new(); - - let err = execute_agent_tool( - &namespace, - "grep", - &json!({ - "path": "/runs/run-1", - "pattern": "best_model", - "recursive": true, - "limit": 301 - }), - ) - .unwrap_err(); - - assert!( - matches!(err, AgentError::InvalidArgument(ref message) if message.contains("limit must be between 1 and 300")), - "unexpected error: {err:?}" - ); - assert!(namespace.last_grep.borrow().is_none()); - } - - #[test] - fn grep_tool_rejects_non_string_patterns() { - let namespace = FakeNamespace::new(); - - let err = execute_agent_tool( - &namespace, - "grep", - &json!({ - "path": "/runs/run-1", - "pattern": "best_model", - "recursive": true, - "patterns": [1] - }), - ) - .unwrap_err(); - - assert!( - matches!(err, AgentError::InvalidArgument(ref message) if message.contains("patterns entries must be strings")), - "unexpected error: {err:?}" - ); - assert!(namespace.last_grep.borrow().is_none()); - } - - #[test] - fn grep_tool_rejects_non_string_glob() { - let namespace = FakeNamespace::new(); - - let err = execute_agent_tool( - &namespace, - "grep", - &json!({ - "path": "/runs/run-1", - "pattern": "best_model", - "recursive": true, - "glob": 7 - }), - ) - .unwrap_err(); - - assert!( - matches!(err, AgentError::InvalidArgument(ref message) if message.contains("glob must be a string or null")), - "unexpected error: {err:?}" - ); - assert!(namespace.last_grep.borrow().is_none()); - } - - #[test] - fn grep_tool_forwards_patterns_and_glob() { - let namespace = FakeNamespace::new(); - - execute_agent_tool( - &namespace, - "grep", - &json!({ - "path": "/runs/run-1", - "pattern": "best_model", - "patterns": ["error", "warning"], - "glob": "*.txt", - "recursive": true - }), - ) - .unwrap(); - - let request = namespace.last_grep.borrow().clone().unwrap(); - assert_eq!( - request.patterns, - vec!["error".to_owned(), "warning".to_owned()] - ); - assert_eq!(request.name_glob.as_deref(), Some("*.txt")); - } - - #[test] - fn grep_tool_defaults_patterns_and_glob_to_unset() { - let namespace = FakeNamespace::new(); - - execute_agent_tool( - &namespace, - "grep", - &json!({ - "path": "/runs/run-1", - "pattern": "best_model", - "recursive": true - }), - ) - .unwrap(); - - let request = namespace.last_grep.borrow().clone().unwrap(); - assert_eq!(request.patterns, Vec::::new()); - assert_eq!(request.name_glob, None); - } -} +pub use namespace::{ + agent_tool_definitions, execute_agent_tool, AgentError, AgentNamespace, AgentToolDefinition, +}; diff --git a/crates/nokv-agent/src/namespace.rs b/crates/nokv-agent/src/namespace.rs new file mode 100644 index 000000000..230e5c94c --- /dev/null +++ b/crates/nokv-agent/src/namespace.rs @@ -0,0 +1,2664 @@ +//! NoKV agent tool surface. +//! +//! Transport-free agent tooling: the JSON tool definitions exposed to a model, +//! the dispatcher that maps a tool call onto a namespace verb, argument +//! validation, and result shaping. Depends only on `nokv-meta`, `nokv-object` +//! and `nokv-types` — never on `nokv-client`, `nokv-protocol` or `nokv-control`. + +use nokv_meta::{ + MetadataStore, NamespaceAggregateGroup, NamespaceAggregateMeasure, NamespaceAggregateOp, + NamespaceAggregateRequest, NamespaceAggregateResult, NamespaceAggregateSort, + NamespaceAggregateValue, NamespaceBodyDescriptor, NamespaceCard, NamespaceCardKind, + NamespaceFacetSummary, NamespaceFacetValue, NamespaceFilterCapability, NamespaceFindField, + NamespaceFindRequest, NamespaceFindResult, NamespaceGrepRequest, NamespaceGrepResult, + NamespaceIndexValue, NamespaceListOptions, NamespaceListPage, NamespacePredicate, + NamespacePredicateOp, NamespacePredicateValue, NamespaceQueryCatalog, NamespaceReadFormat, + NamespaceReadItem, NamespaceReadOptions, NamespaceReadPage, NamespaceRecordCount, + NamespaceRecordType, NamespaceSchema, NamespaceSort, NamespaceSortDirection, + NamespaceSortField, NoKvFs, +}; +use nokv_object::ObjectStore; +use serde_json::{json, Map, Value}; + +/// Error surface for the agent tool layer. +/// +/// The copied dispatcher only ever constructs `Metadata`, `NotFound`, and +/// `InvalidArgument`. `Other` is unused here; it exists so a later step can +/// add `From` in `nokv-client`. +#[derive(Debug)] +pub enum AgentError { + Metadata(nokv_meta::MetadError), + NotFound(String), + InvalidArgument(String), + Other(String), +} + +impl std::fmt::Display for AgentError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + // Display strings MUST stay byte-identical to the old ClientError rendering + // (crates/nokv-client/src/lib.rs:93-119) so model/telemetry output is unchanged: + Self::Metadata(err) => write!(f, "metadata service error: {err}"), + Self::NotFound(path) => write!(f, "path component not found: {path}"), + Self::InvalidArgument(err) => write!(f, "metadata protocol error: {err}"), + Self::Other(msg) => write!(f, "{msg}"), + } + } +} + +impl std::error::Error for AgentError {} + +impl From for AgentError { + fn from(err: nokv_meta::MetadError) -> Self { + Self::Metadata(err) + } +} + +const DEFAULT_AGENT_PAGE_LIMIT: usize = 100; +const MAX_AGENT_PAGE_LIMIT: usize = 100; +const DEFAULT_AGENT_READ_LIMIT: usize = 100; +const MAX_AGENT_READ_LIMIT: usize = 300; +const DEFAULT_AGENT_FIND_LIMIT: usize = 10; +const MAX_AGENT_FIND_LIMIT: usize = 10; +const DEFAULT_AGENT_AGGREGATE_LIMIT: usize = 20; +const MAX_AGENT_AGGREGATE_LIMIT: usize = 100; +const DEFAULT_AGENT_GREP_LIMIT: usize = 100; +const MAX_AGENT_GREP_LIMIT: usize = 300; +/// Mirror of the metadata server's grep pattern cap (nokv-meta +/// MAX_GREP_PATTERNS); enforced there, advertised in the schema here. +const MAX_AGENT_GREP_PATTERNS: usize = 16; + +#[derive(Clone, Debug, PartialEq)] +pub struct AgentToolDefinition { + pub name: &'static str, + pub description: &'static str, + pub parameters: Value, +} + +pub trait AgentNamespace { + fn agent_stat_card(&self, path: &str) -> Result, AgentError>; + + fn agent_list_page( + &self, + path: &str, + options: NamespaceListOptions, + ) -> Result; + + fn agent_find_paths( + &self, + request: NamespaceFindRequest, + ) -> Result; + + fn agent_aggregate_paths( + &self, + request: NamespaceAggregateRequest, + ) -> Result; + + fn agent_grep_paths( + &self, + request: NamespaceGrepRequest, + ) -> Result; + + fn agent_read_page( + &self, + path: &str, + options: NamespaceReadOptions, + ) -> Result; +} + +impl AgentNamespace for NoKvFs +where + M: MetadataStore, + O: ObjectStore, +{ + fn agent_stat_card(&self, path: &str) -> Result, AgentError> { + self.stat_card(path).map_err(AgentError::Metadata) + } + + fn agent_list_page( + &self, + path: &str, + options: NamespaceListOptions, + ) -> Result { + self.list_page(path, options).map_err(AgentError::Metadata) + } + + fn agent_find_paths( + &self, + request: NamespaceFindRequest, + ) -> Result { + self.find_paths(request).map_err(AgentError::Metadata) + } + + fn agent_aggregate_paths( + &self, + request: NamespaceAggregateRequest, + ) -> Result { + self.aggregate_paths(request).map_err(AgentError::Metadata) + } + + fn agent_grep_paths( + &self, + request: NamespaceGrepRequest, + ) -> Result { + self.grep_paths(request).map_err(AgentError::Metadata) + } + + fn agent_read_page( + &self, + path: &str, + options: NamespaceReadOptions, + ) -> Result { + self.read_page(path, options).map_err(AgentError::Metadata) + } +} + +pub fn agent_tool_definitions() -> Vec { + vec![ + AgentToolDefinition { + name: "ls", + description: "List direct children to discover paths. Not recursive; use stat for one path details.", + parameters: json!({ + "type": "object", + "required": ["path"], + "properties": { + "path": {"type": "string"}, + "cursor": {"type": ["string", "null"]}, + "limit": {"type": "integer", "minimum": 1, "maximum": MAX_AGENT_PAGE_LIMIT} + }, + "additionalProperties": false + }), + }, + AgentToolDefinition { + name: "stat", + description: "Inspect a single path compact card: counts, body, schema, sample, catalog, and indexed values. Use read for file body content.", + parameters: json!({ + "type": "object", + "required": ["path"], + "properties": { + "path": {"type": "string"} + }, + "additionalProperties": false + }), + }, + AgentToolDefinition { + name: "catalog", + description: "Discover field ids for find predicates, projections, sort, facets, and aggregate group or measure fields.", + parameters: json!({ + "type": "object", + "required": ["path"], + "properties": { + "path": {"type": "string"}, + "field_prefix": {"type": ["string", "null"]}, + "include_facets": {"type": "boolean"} + }, + "additionalProperties": false + }), + }, + AgentToolDefinition { + name: "read", + description: "Read file body content. Structured mode returns JSON, YAML, or text records; bytes mode returns byte ranges.", + parameters: json!({ + "type": "object", + "required": ["path"], + "properties": { + "path": {"type": "string"}, + "format": {"type": "string", "enum": ["structured", "bytes"]}, + "cursor": {"type": ["string", "null"]}, + "offset": {"type": "integer", "minimum": 0}, + "limit": {"type": "integer", "minimum": 1, "maximum": MAX_AGENT_READ_LIMIT}, + }, + "additionalProperties": false + }), + }, + AgentToolDefinition { + name: "aggregate", + description: "Compute summary rows using catalog field ids: count, sum, avg, min, max, group, filter, and sort.", + parameters: json!({ + "type": "object", + "required": ["path", "measures"], + "properties": { + "path": {"type": "string"}, + "predicates": { + "type": "array", + "items": { + "type": "object", + "required": ["field", "op"], + "properties": { + "field": {"type": "string"}, + "op": { + "type": "string", + "enum": ["eq", "ne", "in", "prefix", "suffix", "contains", "gt", "gte", "lt", "lte", "exists", "not_exists"] + }, + "value": { + "anyOf": [ + {"type": "string"}, + {"type": "integer", "minimum": 0}, + {"type": "number"}, + {"type": "boolean"}, + { + "type": "array", + "items": {"type": ["string", "integer", "number", "boolean"]} + }, + {"type": "null"} + ] + } + }, + "additionalProperties": false + } + }, + "group_by": { + "type": "array", + "items": {"type": "string"} + }, + "measures": { + "type": "array", + "items": { + "type": "object", + "required": ["name", "op"], + "properties": { + "name": {"type": "string"}, + "op": {"type": "string", "enum": ["count", "sum", "avg", "min", "max"]}, + "field": {"type": ["string", "null"]} + }, + "additionalProperties": false + } + }, + "sort": { + "type": "array", + "items": { + "type": "object", + "required": ["field"], + "properties": { + "field": {"type": "string"}, + "direction": {"type": "string", "enum": ["asc", "desc"]} + }, + "additionalProperties": false + } + }, + "limit": {"type": "integer", "minimum": 1, "maximum": MAX_AGENT_AGGREGATE_LIMIT} + }, + "additionalProperties": false + }), + }, + AgentToolDefinition { + name: "find", + description: + "Search paths with catalog field predicates and project fields. Use read for body content and stat for schema or sample.", + parameters: json!({ + "type": "object", + "required": ["path"], + "properties": { + "path": {"type": "string"}, + "predicates": { + "type": "array", + "items": { + "type": "object", + "required": ["field", "op"], + "properties": { + "field": {"type": "string"}, + "op": { + "type": "string", + "enum": ["eq", "ne", "in", "prefix", "suffix", "contains", "gt", "gte", "lt", "lte", "exists", "not_exists"] + }, + "value": { + "anyOf": [ + {"type": "string"}, + {"type": "integer", "minimum": 0}, + {"type": "number"}, + {"type": "boolean"}, + { + "type": "array", + "items": {"type": ["string", "integer", "number", "boolean"]} + }, + {"type": "null"} + ] + } + }, + "additionalProperties": false + } + }, + "sort": { + "type": "array", + "items": { + "type": "object", + "required": ["field"], + "properties": { + "field": {"type": "string"}, + "direction": {"type": "string", "enum": ["asc", "desc"]} + }, + "additionalProperties": false + } + }, + "fields": { + "type": "array", + "items": {"type": "string"} + }, + "facets": { + "type": "array", + "items": {"type": "string"} + }, + "cursor": {"type": ["string", "null"]}, + "limit": {"type": "integer", "minimum": 1, "maximum": MAX_AGENT_FIND_LIMIT} + }, + "additionalProperties": false + }), + }, + AgentToolDefinition { + name: "grep", + description: "Search file bodies for case-insensitive literal substrings and return matching lines with line numbers. patterns adds OR alternatives to pattern (at most 16); glob filters file names. Scope path to one directory or file when known.", + parameters: json!({ + "type": "object", + "required": ["path", "pattern", "recursive"], + "properties": { + "path": {"type": "string"}, + "pattern": {"type": "string"}, + "patterns": {"type": "array", "items": {"type": "string"}, "maxItems": MAX_AGENT_GREP_PATTERNS}, + "glob": {"type": ["string", "null"]}, + "recursive": {"type": "boolean"}, + "cursor": {"type": ["string", "null"]}, + "limit": {"type": "integer", "minimum": 1, "maximum": MAX_AGENT_GREP_LIMIT} + }, + "additionalProperties": false + }), + }, + ] +} + +pub fn execute_agent_tool(namespace: &T, name: &str, args: &Value) -> Result +where + T: AgentNamespace + ?Sized, +{ + match name { + "ls" => execute_ls(namespace, args), + "stat" => execute_stat(namespace, args), + "catalog" => execute_catalog(namespace, args), + "grep" => execute_grep(namespace, args), + "read" => execute_read(namespace, args), + "find" => execute_find(namespace, args), + "aggregate" => execute_aggregate(namespace, args), + other => Err(AgentError::InvalidArgument(format!( + "unknown agent tool {other}" + ))), + } +} + +fn execute_ls(namespace: &T, args: &Value) -> Result +where + T: AgentNamespace + ?Sized, +{ + let path = required_string_arg(args, "path")?; + let page = namespace.agent_list_page( + path, + NamespaceListOptions { + cursor: optional_string_arg(args, "cursor")?, + limit: optional_usize_arg(args, "limit")?.unwrap_or(DEFAULT_AGENT_PAGE_LIMIT), + }, + )?; + Ok(list_page_json(&page)) +} + +fn execute_stat(namespace: &T, args: &Value) -> Result +where + T: AgentNamespace + ?Sized, +{ + let path = required_string_arg(args, "path")?; + let card = namespace + .agent_stat_card(path)? + .ok_or_else(|| AgentError::NotFound(path.to_owned()))?; + Ok(json!({ + "card": card_json(&card), + })) +} + +fn execute_catalog(namespace: &T, args: &Value) -> Result +where + T: AgentNamespace + ?Sized, +{ + let path = required_string_arg(args, "path")?; + let field_prefix = optional_string_arg(args, "field_prefix")?; + let include_facets = optional_bool_arg(args, "include_facets")?.unwrap_or(false); + let card = namespace + .agent_stat_card(path)? + .ok_or_else(|| AgentError::NotFound(path.to_owned()))?; + let catalog = filtered_catalog(&card.catalog, field_prefix.as_deref(), include_facets); + let catalog_empty = catalog_is_empty(&catalog); + let child_catalogs = if catalog_empty && matches!(card.kind, NamespaceCardKind::Directory) { + child_catalogs(namespace, &card.path, include_facets)? + } else { + Vec::new() + }; + Ok(json!({ + "path": card.path, + "catalog_empty": catalog_empty, + "catalog": catalog_json(&catalog), + "child_catalogs": child_catalogs, + })) +} + +fn execute_grep(namespace: &T, args: &Value) -> Result +where + T: AgentNamespace + ?Sized, +{ + let path = required_string_arg(args, "path")?; + let pattern = required_string_arg(args, "pattern")?; + let recursive = optional_bool_arg(args, "recursive")?.ok_or_else(|| { + AgentError::InvalidArgument("grep requires a boolean recursive argument".to_owned()) + })?; + let result = namespace.agent_grep_paths(NamespaceGrepRequest { + path: path.to_owned(), + pattern: pattern.to_owned(), + patterns: optional_string_array_arg(args, "patterns")?.unwrap_or_default(), + recursive, + name_glob: optional_string_arg(args, "glob")?, + cursor: optional_string_arg(args, "cursor")?, + limit: optional_bounded_usize_arg(args, "limit", MAX_AGENT_GREP_LIMIT)? + .unwrap_or(DEFAULT_AGENT_GREP_LIMIT), + max_files: None, + max_bytes: None, + })?; + Ok(grep_result_json(&result)) +} + +fn grep_result_json(result: &NamespaceGrepResult) -> Value { + json!({ + "path": result.path, + "pattern": result.pattern, + "recursive": result.recursive, + "matches": result + .matches + .iter() + .map(|match_| json!({ + "path": match_.path, + "line_number": match_.line_number, + "snippet": match_.snippet, + })) + .collect::>(), + "files_scanned": result.files_scanned, + "next_cursor": result.next_cursor, + "truncated": result.truncated, + }) +} + +fn execute_read(namespace: &T, args: &Value) -> Result +where + T: AgentNamespace + ?Sized, +{ + let path = required_string_arg(args, "path")?; + let options = NamespaceReadOptions { + format: read_format_arg(args)?, + cursor: optional_string_arg(args, "cursor")?, + offset: optional_u64_arg(args, "offset")?.unwrap_or(0), + limit: optional_bounded_usize_arg(args, "limit", MAX_AGENT_READ_LIMIT)? + .unwrap_or(DEFAULT_AGENT_READ_LIMIT), + expected_generation: None, + }; + guard_large_structured_pagination(namespace, path, &options)?; + let page = namespace.agent_read_page(path, options)?; + Ok(read_page_json(&page)) +} + +fn guard_large_structured_pagination( + namespace: &T, + path: &str, + options: &NamespaceReadOptions, +) -> Result<(), AgentError> +where + T: AgentNamespace + ?Sized, +{ + if !matches!(options.format, NamespaceReadFormat::Structured) { + return Ok(()); + } + let Some(card) = namespace.agent_stat_card(path)? else { + return Ok(()); + }; + let Some(record_count) = card.record_count else { + return Ok(()); + }; + if record_count.count <= MAX_AGENT_READ_LIMIT { + return Ok(()); + } + Err(AgentError::InvalidArgument(format!( + "structured pagination for {path} has {} records; use bytes format with offset and limit, grep to locate lines, or stat record_count", + record_count.count + ))) +} + +fn execute_find(namespace: &T, args: &Value) -> Result +where + T: AgentNamespace + ?Sized, +{ + let path = required_string_arg(args, "path")?; + let fields = optional_string_array_arg(args, "fields")?; + reject_unsupported_find_arguments(args)?; + let result = namespace.agent_find_paths(NamespaceFindRequest { + path: path.to_owned(), + predicates: predicates_arg(args)?, + sort: sort_arg(args)?, + include: Vec::new(), + facets: facets_arg(args)?, + cursor: optional_string_arg(args, "cursor")?, + limit: optional_bounded_usize_arg(args, "limit", MAX_AGENT_FIND_LIMIT)? + .unwrap_or(DEFAULT_AGENT_FIND_LIMIT), + })?; + Ok(find_result_json(&result, fields.as_deref())) +} + +fn execute_aggregate(namespace: &T, args: &Value) -> Result +where + T: AgentNamespace + ?Sized, +{ + let path = required_string_arg(args, "path")?; + let predicates = predicates_arg(args)?; + let group_by = group_by_arg(args)?; + let measures = aggregate_measures_arg(args)?; + let sort = aggregate_sort_arg(args)?; + let limit = optional_bounded_usize_arg(args, "limit", MAX_AGENT_AGGREGATE_LIMIT)? + .unwrap_or(DEFAULT_AGENT_AGGREGATE_LIMIT); + let result = namespace.agent_aggregate_paths(NamespaceAggregateRequest { + path: path.to_owned(), + predicates, + group_by, + measures, + sort, + limit, + })?; + Ok(aggregate_result_json(&result)) +} + +fn list_page_json(page: &NamespaceListPage) -> Value { + json!({ + "path": page.path, + "entry_count": page.entry_count, + "entries": page.entries.iter().map(list_entry_json).collect::>(), + "next_cursor": page.next_cursor, + "truncated": page.truncated, + }) +} + +fn find_result_json(result: &NamespaceFindResult, fields: Option<&[String]>) -> Value { + json!({ + "path": result.path, + "match_count": result.match_count, + "matches": result.matches.iter().map(|card| find_match_json(card, fields)).collect::>(), + "facets": result.facets.iter().map(facet_summary_json).collect::>(), + "next_cursor": result.next_cursor, + "truncated": result.truncated, + }) +} + +fn read_page_json(page: &NamespaceReadPage) -> Value { + json!({ + "path": page.path, + "generation": page.generation, + "total_size_bytes": page.total_size_bytes, + "format": read_format_name(&page.format), + "record_type": page.record_type.as_ref().map(record_type_name), + "record_count": page.record_count, + "cursor": page.cursor, + "next_cursor": page.next_cursor, + "truncated": page.truncated, + "items": page.items.iter().map(read_item_json).collect::>(), + "bytes": page.bytes, + }) +} + +fn card_json(card: &NamespaceCard) -> Value { + json!({ + "path": card.path, + "name": card.name, + "kind": card_kind_name(&card.kind), + "size_bytes": card.size_bytes, + "entry_count": card.entry_count, + "record_count": card.record_count.as_ref().map(record_count_json), + "schema": card.schema.as_ref().map(schema_json), + "sample": card.sample, + "body": card.body.as_ref().map(body_json), + "catalog": catalog_json(&card.catalog), + "indexed_values": card.indexed_values.iter().map(index_value_json).collect::>(), + }) +} + +fn list_entry_json(card: &NamespaceCard) -> Value { + json!({ + "path": card.path, + "name": card.name, + "kind": card_kind_name(&card.kind), + "size_bytes": card.size_bytes, + "entry_count": card.entry_count, + }) +} + +fn find_match_json(card: &NamespaceCard, fields: Option<&[String]>) -> Value { + if fields.is_some() { + let mut object = Map::new(); + object.insert("path".to_owned(), json!(card.path)); + object.insert( + "values".to_owned(), + projected_values_json(card, fields.unwrap_or(&[])), + ); + return Value::Object(object); + } + let mut object = Map::new(); + object.insert("path".to_owned(), json!(card.path)); + Value::Object(object) +} + +fn projected_values_json(card: &NamespaceCard, fields: &[String]) -> Value { + let mut output = Map::new(); + for field in fields { + let values = projected_predicate_values(card, field) + .into_iter() + .map(|value| predicate_value_json(&value)) + .collect::>(); + match values.as_slice() { + [] => {} + [value] => { + output.insert(field.clone(), value.clone()); + } + _ => { + output.insert(field.clone(), Value::Array(values)); + } + } + } + Value::Object(output) +} + +fn projected_predicate_values(card: &NamespaceCard, field: &str) -> Vec { + match field { + "path" => vec![NamespacePredicateValue::String(card.path.clone())], + "name" => vec![NamespacePredicateValue::String(card.name.clone())], + "kind" => vec![NamespacePredicateValue::String( + card_kind_name(&card.kind).to_owned(), + )], + "size_bytes" => card + .size_bytes + .map(NamespacePredicateValue::U64) + .into_iter() + .collect(), + "body.content_type" => card + .body + .as_ref() + .map(|body| NamespacePredicateValue::String(body.content_type.clone())) + .into_iter() + .collect(), + "body.producer" => card + .body + .as_ref() + .map(|body| NamespacePredicateValue::String(body.producer.clone())) + .into_iter() + .collect(), + _ => card + .indexed_values + .iter() + .filter(|value| value.field.id == field) + .map(|value| value.value.clone()) + .collect(), + } +} + +fn record_count_json(count: &NamespaceRecordCount) -> Value { + json!(count.count) +} + +fn schema_json(schema: &NamespaceSchema) -> Value { + json!({ + "record_type": record_type_name(&schema.record_type), + "fields": schema.fields, + }) +} + +fn body_json(body: &NamespaceBodyDescriptor) -> Value { + json!({ + "producer": body.producer, + "size": body.size, + "content_type": body.content_type, + }) +} + +fn catalog_json(catalog: &NamespaceQueryCatalog) -> Value { + json!({ + "filterable": grouped_filterable_json(&catalog.filterable), + "sortable": catalog.sortable.iter().map(|field| field.id.clone()).collect::>(), + "facetable": catalog.facetable.iter().map(|field| field.id.clone()).collect::>(), + "facets": catalog.facets.iter().map(facet_summary_json).collect::>(), + }) +} + +fn grouped_filterable_json(capabilities: &[NamespaceFilterCapability]) -> Vec { + let mut groups: Vec<(Vec<&'static str>, Vec)> = Vec::new(); + for capability in capabilities { + let operators = capability + .operators + .iter() + .map(predicate_op_name) + .collect::>(); + match groups.iter_mut().find(|(group, _)| *group == operators) { + Some((_, fields)) => fields.push(capability.field.id.clone()), + None => groups.push((operators, vec![capability.field.id.clone()])), + } + } + groups + .into_iter() + .map(|(operators, fields)| json!({"operators": operators, "fields": fields})) + .collect() +} + +fn facet_summary_json(facet: &NamespaceFacetSummary) -> Value { + json!({ + "field": facet.field.id, + "values": facet.values.iter().map(facet_value_json).collect::>(), + "distinct_count": facet.distinct_count, + "truncated": facet.truncated, + }) +} + +fn facet_value_json(value: &NamespaceFacetValue) -> Value { + json!({ + "value": predicate_value_json(&value.value), + "count": value.count, + }) +} + +fn index_value_json(value: &NamespaceIndexValue) -> Value { + json!({ + "field": value.field.id, + "value": predicate_value_json(&value.value), + }) +} + +fn predicate_value_json(value: &NamespacePredicateValue) -> Value { + match value { + NamespacePredicateValue::String(value) => json!(value), + NamespacePredicateValue::U64(value) => json!(value), + NamespacePredicateValue::F64(value) if value.is_finite() => json!(value), + NamespacePredicateValue::F64(_) => Value::Null, + NamespacePredicateValue::List(values) => { + Value::Array(values.iter().map(predicate_value_json).collect()) + } + } +} + +fn aggregate_result_json(result: &NamespaceAggregateResult) -> Value { + json!({ + "path": result.path, + "input_match_count": result.input_match_count, + "row_count": result.row_count, + "group_count": result.group_count, + "groups": result.groups.iter().map(aggregate_group_json).collect::>(), + "truncated": result.truncated, + }) +} + +fn aggregate_group_json(group: &NamespaceAggregateGroup) -> Value { + let key = group + .key + .iter() + .map(|value| (value.field.id.clone(), predicate_value_json(&value.value))) + .collect::>(); + let values = group + .measures + .iter() + .map(|measure| (measure.name.clone(), aggregate_value_json(&measure.value))) + .collect::>(); + json!({ + "key": key, + "values": values, + }) +} + +fn aggregate_value_json(value: &NamespaceAggregateValue) -> Value { + match value { + NamespaceAggregateValue::U64(value) => json!(value), + NamespaceAggregateValue::F64(value) if value.is_finite() => json!(value), + NamespaceAggregateValue::F64(_) | NamespaceAggregateValue::Null => Value::Null, + } +} + +fn read_item_json(item: &NamespaceReadItem) -> Value { + let value = serde_json::from_str::(&item.value_json) + .unwrap_or_else(|_| Value::String(item.value_json.clone())); + json!({ + "index": item.index, + "value": value, + }) +} + +fn object_args(args: &Value) -> Result<&Map, AgentError> { + args.as_object().ok_or_else(|| { + AgentError::InvalidArgument("agent tool arguments must be a JSON object".to_owned()) + }) +} + +fn required_string_arg<'a>(args: &'a Value, name: &'static str) -> Result<&'a str, AgentError> { + object_args(args)? + .get(name) + .and_then(Value::as_str) + .ok_or_else(|| AgentError::InvalidArgument(format!("missing string argument {name}"))) +} + +fn optional_string_arg(args: &Value, name: &'static str) -> Result, AgentError> { + let Some(value) = object_args(args)?.get(name) else { + return Ok(None); + }; + if value.is_null() { + return Ok(None); + } + value + .as_str() + .map(|value| Some(value.to_owned())) + .ok_or_else(|| AgentError::InvalidArgument(format!("{name} must be a string or null"))) +} + +fn optional_bool_arg(args: &Value, name: &'static str) -> Result, AgentError> { + let Some(value) = object_args(args)?.get(name) else { + return Ok(None); + }; + if value.is_null() { + return Ok(None); + } + value + .as_bool() + .map(Some) + .ok_or_else(|| AgentError::InvalidArgument(format!("{name} must be a boolean or null"))) +} + +fn optional_u64_arg(args: &Value, name: &'static str) -> Result, AgentError> { + let Some(value) = object_args(args)?.get(name) else { + return Ok(None); + }; + if value.is_null() { + return Ok(None); + } + value.as_u64().map(Some).ok_or_else(|| { + AgentError::InvalidArgument(format!("{name} must be a non-negative integer")) + }) +} + +fn optional_usize_arg(args: &Value, name: &'static str) -> Result, AgentError> { + optional_bounded_usize_arg(args, name, MAX_AGENT_PAGE_LIMIT) +} + +fn optional_bounded_usize_arg( + args: &Value, + name: &'static str, + max: usize, +) -> Result, AgentError> { + optional_u64_arg(args, name)? + .map(|value| { + let value = usize::try_from(value).map_err(|_| { + AgentError::InvalidArgument(format!("{name} exceeds platform limit")) + })?; + if value == 0 || value > max { + return Err(AgentError::InvalidArgument(format!( + "{name} must be between 1 and {max}" + ))); + } + Ok(value) + }) + .transpose() +} + +fn read_format_arg(args: &Value) -> Result { + match object_args(args)?.get("format").and_then(Value::as_str) { + None | Some("structured") => Ok(NamespaceReadFormat::Structured), + Some("bytes") => Ok(NamespaceReadFormat::Bytes), + Some(other) => Err(AgentError::InvalidArgument(format!( + "unsupported read format {other}; expected structured or bytes" + ))), + } +} + +fn predicates_arg(args: &Value) -> Result, AgentError> { + let Some(value) = object_args(args)?.get("predicates") else { + return Ok(Vec::new()); + }; + let predicates = value + .as_array() + .ok_or_else(|| AgentError::InvalidArgument("predicates must be an array".to_owned()))?; + predicates.iter().map(predicate_arg).collect() +} + +fn predicate_arg(value: &Value) -> Result { + let object = value + .as_object() + .ok_or_else(|| AgentError::InvalidArgument("predicate must be an object".to_owned()))?; + let field = string_property(object, "field")?; + let op = string_property(object, "op")?; + let op = predicate_op_arg(op)?; + let raw_value = object.get("value").filter(|value| !value.is_null()); + let value = match op { + // Tolerate a stray value on existence checks; agents commonly pass + // {"op": "exists", "value": true} and the value carries no meaning. + NamespacePredicateOp::Exists | NamespacePredicateOp::NotExists => None, + NamespacePredicateOp::In => { + let value = raw_value.ok_or_else(|| { + AgentError::InvalidArgument("predicate op in requires array value".to_owned()) + })?; + if !value.is_array() { + return Err(AgentError::InvalidArgument( + "predicate op in requires array value".to_owned(), + )); + } + Some(predicate_value_arg(value)?) + } + _ => { + let value = raw_value.ok_or_else(|| { + AgentError::InvalidArgument(format!( + "predicate op {} requires value", + predicate_op_name(&op) + )) + })?; + Some(predicate_value_arg(value)?) + } + }; + Ok(NamespacePredicate { + field: NamespaceFindField::new(field), + op, + value, + }) +} + +fn sort_arg(args: &Value) -> Result, AgentError> { + let Some(value) = object_args(args)?.get("sort") else { + return Ok(Vec::new()); + }; + let sort = value + .as_array() + .ok_or_else(|| AgentError::InvalidArgument("sort must be an array".to_owned()))?; + sort.iter().map(sort_item_arg).collect() +} + +fn sort_item_arg(value: &Value) -> Result { + let object = value + .as_object() + .ok_or_else(|| AgentError::InvalidArgument("sort item must be an object".to_owned()))?; + let field = string_property(object, "field")?; + let direction = object + .get("direction") + .and_then(Value::as_str) + .unwrap_or("asc"); + Ok(NamespaceSort { + field: NamespaceSortField::new(field), + direction: match direction { + "asc" => NamespaceSortDirection::Asc, + "desc" => NamespaceSortDirection::Desc, + other => { + return Err(AgentError::InvalidArgument(format!( + "unsupported sort direction {other}" + ))) + } + }, + }) +} + +fn reject_unsupported_find_arguments(args: &Value) -> Result<(), AgentError> { + if object_args(args)?.contains_key("include") { + return Err(AgentError::InvalidArgument( + "unsupported argument include; use stat for schema or sample and read for body content" + .to_owned(), + )); + } + Ok(()) +} + +fn optional_string_array_arg( + args: &Value, + name: &'static str, +) -> Result>, AgentError> { + let Some(value) = object_args(args)?.get(name) else { + return Ok(None); + }; + if value.is_null() { + return Ok(None); + } + let entries = value + .as_array() + .ok_or_else(|| AgentError::InvalidArgument(format!("{name} must be an array")))? + .iter() + .map(|value| { + value.as_str().map(str::to_owned).ok_or_else(|| { + AgentError::InvalidArgument(format!("{name} entries must be strings")) + }) + }) + .collect::, _>>()?; + Ok(Some(entries)) +} + +fn facets_arg(args: &Value) -> Result, AgentError> { + let Some(value) = object_args(args)?.get("facets") else { + return Ok(Vec::new()); + }; + if value.is_null() { + return Ok(Vec::new()); + } + value + .as_array() + .ok_or_else(|| AgentError::InvalidArgument("facets must be an array".to_owned()))? + .iter() + .map(|value| { + value.as_str().map(NamespaceFindField::new).ok_or_else(|| { + AgentError::InvalidArgument("facets entries must be strings".to_owned()) + }) + }) + .collect() +} + +fn group_by_arg(args: &Value) -> Result, AgentError> { + let Some(value) = object_args(args)?.get("group_by") else { + return Ok(Vec::new()); + }; + if value.is_null() { + return Ok(Vec::new()); + } + value + .as_array() + .ok_or_else(|| AgentError::InvalidArgument("group_by must be an array".to_owned()))? + .iter() + .map(|value| { + value.as_str().map(NamespaceFindField::new).ok_or_else(|| { + AgentError::InvalidArgument("group_by entries must be strings".to_owned()) + }) + }) + .collect() +} + +fn aggregate_measures_arg(args: &Value) -> Result, AgentError> { + let value = object_args(args)? + .get("measures") + .ok_or_else(|| AgentError::InvalidArgument("missing array argument measures".to_owned()))?; + let measures = value + .as_array() + .ok_or_else(|| AgentError::InvalidArgument("measures must be an array".to_owned()))?; + if measures.is_empty() { + return Err(AgentError::InvalidArgument( + "measures must contain at least one measure".to_owned(), + )); + } + measures + .iter() + .map(|value| { + let object = value.as_object().ok_or_else(|| { + AgentError::InvalidArgument("measure must be an object".to_owned()) + })?; + let name = string_property(object, "name")?.to_owned(); + if name.is_empty() { + return Err(AgentError::InvalidArgument( + "measure name must not be empty".to_owned(), + )); + } + let op = aggregate_op_arg(string_property(object, "op")?)?; + let field = object + .get("field") + .and_then(|value| (!value.is_null()).then_some(value)) + .map(|value| { + value.as_str().map(str::to_owned).ok_or_else(|| { + AgentError::InvalidArgument( + "measure field must be a string or null".to_owned(), + ) + }) + }) + .transpose()?; + if !matches!(op, NamespaceAggregateOp::Count) && field.is_none() { + return Err(AgentError::InvalidArgument(format!( + "measure {name} with op {} requires field", + aggregate_op_name(&op) + ))); + } + Ok(NamespaceAggregateMeasure { + name, + op, + field: field.map(NamespaceFindField::new), + }) + }) + .collect() +} + +fn aggregate_op_arg(op: &str) -> Result { + match op { + "count" => Ok(NamespaceAggregateOp::Count), + "sum" => Ok(NamespaceAggregateOp::Sum), + "avg" => Ok(NamespaceAggregateOp::Avg), + "min" => Ok(NamespaceAggregateOp::Min), + "max" => Ok(NamespaceAggregateOp::Max), + other => Err(AgentError::InvalidArgument(format!( + "unsupported aggregate op {other}" + ))), + } +} + +fn aggregate_sort_arg(args: &Value) -> Result, AgentError> { + let Some(value) = object_args(args)?.get("sort") else { + return Ok(Vec::new()); + }; + let sort = value + .as_array() + .ok_or_else(|| AgentError::InvalidArgument("sort must be an array".to_owned()))?; + sort.iter() + .map(|value| { + let object = value.as_object().ok_or_else(|| { + AgentError::InvalidArgument("sort item must be an object".to_owned()) + })?; + let field = string_property(object, "field")?.to_owned(); + let direction = object + .get("direction") + .and_then(Value::as_str) + .unwrap_or("asc"); + Ok(NamespaceAggregateSort { + field, + direction: match direction { + "asc" => NamespaceSortDirection::Asc, + "desc" => NamespaceSortDirection::Desc, + other => { + return Err(AgentError::InvalidArgument(format!( + "unsupported sort direction {other}" + ))) + } + }, + }) + }) + .collect() +} + +fn filtered_catalog( + catalog: &NamespaceQueryCatalog, + field_prefix: Option<&str>, + include_facets: bool, +) -> NamespaceQueryCatalog { + NamespaceQueryCatalog { + filterable: catalog + .filterable + .iter() + .filter(|capability| field_matches_prefix(capability.field.id.as_str(), field_prefix)) + .cloned() + .collect(), + sortable: catalog + .sortable + .iter() + .filter(|field| field_matches_prefix(field.id.as_str(), field_prefix)) + .cloned() + .collect(), + facetable: catalog + .facetable + .iter() + .filter(|field| field_matches_prefix(field.id.as_str(), field_prefix)) + .cloned() + .collect(), + facets: if include_facets { + catalog + .facets + .iter() + .filter(|facet| field_matches_prefix(facet.field.id.as_str(), field_prefix)) + .cloned() + .collect() + } else { + Vec::new() + }, + projections: catalog.projections.clone(), + } +} + +fn catalog_is_empty(catalog: &NamespaceQueryCatalog) -> bool { + catalog.filterable.is_empty() && catalog.sortable.is_empty() && catalog.facetable.is_empty() +} + +fn child_catalogs( + namespace: &T, + path: &str, + include_facets: bool, +) -> Result, AgentError> +where + T: AgentNamespace + ?Sized, +{ + let page = namespace.agent_list_page( + path, + NamespaceListOptions { + cursor: None, + limit: 20, + }, + )?; + let mut children = Vec::new(); + for entry in page.entries { + if !matches!(entry.kind, NamespaceCardKind::Directory) { + continue; + } + let Some(card) = namespace.agent_stat_card(&entry.path)? else { + continue; + }; + let catalog = filtered_catalog(&card.catalog, None, include_facets); + if catalog_is_empty(&catalog) { + continue; + } + children.push(json!({ + "path": card.path, + "catalog": catalog_json(&catalog), + })); + if children.len() == 5 { + break; + } + } + Ok(children) +} + +fn field_matches_prefix(field: &str, prefix: Option<&str>) -> bool { + prefix + .map(|prefix| field.starts_with(prefix) || field.contains(prefix)) + .unwrap_or(true) +} + +fn aggregate_op_name(op: &NamespaceAggregateOp) -> &'static str { + match op { + NamespaceAggregateOp::Count => "count", + NamespaceAggregateOp::Sum => "sum", + NamespaceAggregateOp::Avg => "avg", + NamespaceAggregateOp::Min => "min", + NamespaceAggregateOp::Max => "max", + } +} + +fn string_property<'a>( + object: &'a Map, + name: &'static str, +) -> Result<&'a str, AgentError> { + object + .get(name) + .and_then(Value::as_str) + .ok_or_else(|| AgentError::InvalidArgument(format!("missing string property {name}"))) +} + +fn predicate_op_arg(op: &str) -> Result { + match op { + "eq" => Ok(NamespacePredicateOp::Eq), + "ne" | "not_equal" => Ok(NamespacePredicateOp::NotEqual), + "in" => Ok(NamespacePredicateOp::In), + "prefix" => Ok(NamespacePredicateOp::Prefix), + "suffix" => Ok(NamespacePredicateOp::Suffix), + "contains" => Ok(NamespacePredicateOp::Contains), + "gt" | "greater_than" => Ok(NamespacePredicateOp::GreaterThan), + "gte" | "greater_than_or_equal" => Ok(NamespacePredicateOp::GreaterThanOrEqual), + "lt" | "less_than" => Ok(NamespacePredicateOp::LessThan), + "lte" | "less_than_or_equal" => Ok(NamespacePredicateOp::LessThanOrEqual), + "exists" => Ok(NamespacePredicateOp::Exists), + "not_exists" => Ok(NamespacePredicateOp::NotExists), + other => Err(AgentError::InvalidArgument(format!( + "unsupported predicate operator {other}" + ))), + } +} + +fn predicate_value_arg(value: &Value) -> Result { + if let Some(values) = value.as_array() { + return values + .iter() + .map(predicate_scalar_value_arg) + .collect::, _>>() + .map(NamespacePredicateValue::List); + } + predicate_scalar_value_arg(value) +} + +fn predicate_scalar_value_arg(value: &Value) -> Result { + if let Some(value) = value.as_str() { + return Ok(NamespacePredicateValue::String(value.to_owned())); + } + if let Some(value) = value.as_bool() { + return Ok(NamespacePredicateValue::U64(u64::from(value))); + } + if let Some(value) = value.as_u64() { + return Ok(NamespacePredicateValue::U64(value)); + } + if let Some(value) = value.as_f64() { + if value.is_finite() { + return Ok(NamespacePredicateValue::F64(value)); + } + } + Err(AgentError::InvalidArgument( + "predicate value must be a string, boolean, finite number, or array of scalar values" + .to_owned(), + )) +} + +fn card_kind_name(kind: &NamespaceCardKind) -> &'static str { + match kind { + NamespaceCardKind::File => "file", + NamespaceCardKind::Directory => "directory", + NamespaceCardKind::Symlink => "symlink", + NamespaceCardKind::Special => "special", + } +} + +fn record_type_name(record_type: &NamespaceRecordType) -> &'static str { + match record_type { + NamespaceRecordType::DirectoryEntries => "directory_entries", + NamespaceRecordType::JsonArray => "json_array", + NamespaceRecordType::JsonObject => "json_object", + NamespaceRecordType::YamlMapping => "yaml_mapping", + NamespaceRecordType::TextLines => "text_lines", + } +} + +fn read_format_name(format: &NamespaceReadFormat) -> &'static str { + match format { + NamespaceReadFormat::Structured => "structured", + NamespaceReadFormat::Bytes => "bytes", + } +} + +fn predicate_op_name(op: &NamespacePredicateOp) -> &'static str { + match op { + NamespacePredicateOp::Eq => "eq", + NamespacePredicateOp::NotEqual => "ne", + NamespacePredicateOp::In => "in", + NamespacePredicateOp::Prefix => "prefix", + NamespacePredicateOp::Suffix => "suffix", + NamespacePredicateOp::Contains => "contains", + NamespacePredicateOp::GreaterThan => "gt", + NamespacePredicateOp::GreaterThanOrEqual => "gte", + NamespacePredicateOp::LessThan => "lt", + NamespacePredicateOp::LessThanOrEqual => "lte", + NamespacePredicateOp::Exists => "exists", + NamespacePredicateOp::NotExists => "not_exists", + } +} + +#[cfg(test)] +mod tests { + use super::*; + use nokv_meta::{ + NamespaceAggregateOutputMeasure, NamespaceFieldSource, NamespaceFieldSourceKind, + NamespaceFieldValue, NamespaceGrepMatch, NamespaceInclude, NamespaceRecordCount, + RecordCountProvenance, + }; + use nokv_types::InodeId; + use std::cell::RefCell; + use std::collections::BTreeMap; + + struct FakeNamespace { + last_find: RefCell>, + last_aggregate: RefCell>, + last_grep: RefCell>, + last_read: RefCell>, + read_calls: RefCell, + record_count: usize, + find_matches: Vec, + aggregate_result: NamespaceAggregateResult, + stat_cards: BTreeMap, + list_entries: Vec, + } + + impl FakeNamespace { + fn new() -> Self { + Self { + last_find: RefCell::new(None), + last_aggregate: RefCell::new(None), + last_grep: RefCell::new(None), + last_read: RefCell::new(None), + read_calls: RefCell::new(0), + record_count: 1, + find_matches: vec![sample_card("/runs/run-1", 1)], + aggregate_result: sample_aggregate_result(), + stat_cards: BTreeMap::new(), + list_entries: vec![sample_card("/runs/run-1", 1)], + } + } + + fn with_record_count(record_count: usize) -> Self { + Self { + record_count, + find_matches: vec![sample_card("/runs/run-1", record_count)], + list_entries: vec![sample_card("/runs/run-1", record_count)], + ..Self::new() + } + } + + fn with_aggregate_result(aggregate_result: NamespaceAggregateResult) -> Self { + Self { + aggregate_result, + ..Self::new() + } + } + + fn with_root_child_catalog() -> Self { + let mut root = sample_card("/yanex", 1); + root.kind = NamespaceCardKind::Directory; + root.catalog = empty_catalog(); + let mut runs = sample_card("/yanex/runs", 1); + runs.kind = NamespaceCardKind::Directory; + runs.catalog.filterable.push(NamespaceFilterCapability { + field: NamespaceFindField::new("run.script"), + operators: vec![NamespacePredicateOp::Eq], + }); + runs.catalog + .sortable + .push(NamespaceSortField::new("run.script")); + let mut stat_cards = BTreeMap::new(); + stat_cards.insert(root.path.clone(), root); + stat_cards.insert(runs.path.clone(), runs.clone()); + Self { + find_matches: vec![sample_card("/yanex/runs/run-1", 1)], + stat_cards, + list_entries: vec![runs], + ..Self::new() + } + } + } + + impl AgentNamespace for FakeNamespace { + fn agent_stat_card(&self, path: &str) -> Result, AgentError> { + Ok(Some( + self.stat_cards + .get(path) + .cloned() + .unwrap_or_else(|| sample_card(path, self.record_count)), + )) + } + + fn agent_list_page( + &self, + path: &str, + _options: NamespaceListOptions, + ) -> Result { + Ok(NamespaceListPage { + path: path.to_owned(), + evidence: "nokv-native:///runs".to_owned(), + snapshot_id: Some(9), + entry_count: self.list_entries.len(), + entries: self.list_entries.clone(), + next_cursor: None, + truncated: false, + }) + } + + fn agent_find_paths( + &self, + request: NamespaceFindRequest, + ) -> Result { + self.last_find.replace(Some(request)); + let matches = self.find_matches.clone(); + Ok(NamespaceFindResult { + path: "/runs".to_owned(), + evidence: "nokv-native:///runs".to_owned(), + snapshot_id: Some(9), + match_count: matches.len(), + matches, + facets: vec![NamespaceFacetSummary { + field: NamespaceFindField::new("run.script"), + values: vec![NamespaceFacetValue { + value: NamespacePredicateValue::String("train.py".to_owned()), + count: 1, + }], + distinct_count: 1, + truncated: false, + }], + next_cursor: None, + truncated: false, + scanned_entries: 1, + }) + } + + fn agent_aggregate_paths( + &self, + request: NamespaceAggregateRequest, + ) -> Result { + let mut result = self.aggregate_result.clone(); + result.path = request.path.clone(); + result.predicates = request.predicates.clone(); + self.last_aggregate.replace(Some(request)); + Ok(result) + } + + fn agent_grep_paths( + &self, + request: NamespaceGrepRequest, + ) -> Result { + self.last_grep.replace(Some(request.clone())); + let matched = "Checkpoint: best_model_1.0_1.pt"; + let matches = if matched + .to_lowercase() + .contains(&request.pattern.to_lowercase()) + { + vec![NamespaceGrepMatch { + path: format!("{}/artifacts/stdout.txt", request.path), + line_number: 3, + snippet: matched.to_owned(), + evidence: format!( + "nokv-native://{}/artifacts/stdout.txt@generation:1#L3", + request.path + ), + generation: 1, + }] + } else { + Vec::new() + }; + Ok(NamespaceGrepResult { + path: request.path.clone(), + pattern: request.pattern.clone(), + patterns: request.patterns.clone(), + recursive: request.recursive, + evidence: format!("nokv-native://{}", request.path), + snapshot_id: Some(1), + matches, + files_scanned: 1, + bytes_read: matched.len(), + next_cursor: None, + truncated: false, + }) + } + + fn agent_read_page( + &self, + path: &str, + options: NamespaceReadOptions, + ) -> Result { + self.last_read.replace(Some(options)); + *self.read_calls.borrow_mut() += 1; + Ok(NamespaceReadPage { + path: path.to_owned(), + evidence: "nokv-native:///runs/run-1/metadata.json@generation:7".to_owned(), + snapshot_id: Some(9), + generation: 7, + total_size_bytes: 13, + format: NamespaceReadFormat::Structured, + record_type: Some(NamespaceRecordType::JsonObject), + record_count: Some(1), + cursor: None, + next_cursor: None, + truncated: false, + items: vec![NamespaceReadItem { + index: 0, + value_json: r#"{"status":"completed"}"#.to_owned(), + evidence: "nokv-native:///runs/run-1/metadata.json@generation:7#item:0" + .to_owned(), + }], + bytes: None, + }) + } + } + + fn tool_definition<'a>( + tools: &'a [AgentToolDefinition], + name: &str, + ) -> &'a AgentToolDefinition { + tools + .iter() + .find(|tool| tool.name == name) + .unwrap_or_else(|| panic!("{name} tool must be registered")) + } + + fn assert_tool_description_contains( + tools: &[AgentToolDefinition], + name: &str, + expected: &[&str], + ) { + let description = tool_definition(tools, name).description; + let normalized_description = description.to_ascii_lowercase(); + for snippet in expected { + let normalized_snippet = snippet.to_ascii_lowercase(); + assert!( + normalized_description.contains(&normalized_snippet), + "{name} description must contain {snippet:?}; got {description:?}" + ); + } + } + + fn assert_json_lacks_agent_noise(value: &Value) { + // "generation" is intentionally not forbidden: read exposes it so agents + // can feed the conditional read validator (expected_generation). + const FORBIDDEN_KEYS: &[&str] = &[ + "tool", + "bytes_read", + "evidence", + "snapshot_id", + "field_values", + "source_path", + "source_kind", + "key_fields", + "measures", + "sample_matches", + "scope_note", + "predicates", + "scanned_entries", + "digest_uri", + "manifest_id", + "chunk_size", + "block_size", + "inode", + "provenance", + ]; + match value { + Value::Object(object) => { + for forbidden in FORBIDDEN_KEYS { + assert!( + !object.contains_key(*forbidden), + "agent output must not contain noisy key {forbidden:?}: {value}" + ); + } + for child in object.values() { + assert_json_lacks_agent_noise(child); + } + } + Value::Array(values) => { + for child in values { + assert_json_lacks_agent_noise(child); + } + } + _ => {} + } + } + + fn assert_schema_lacks_nested_descriptions(value: &Value) { + match value { + Value::Object(object) => { + assert!( + !object.contains_key("description"), + "tool parameter schema must not contain nested descriptions: {value}" + ); + for child in object.values() { + assert_schema_lacks_nested_descriptions(child); + } + } + Value::Array(values) => { + for child in values { + assert_schema_lacks_nested_descriptions(child); + } + } + _ => {} + } + } + + fn empty_catalog() -> NamespaceQueryCatalog { + NamespaceQueryCatalog { + filterable: Vec::new(), + sortable: Vec::new(), + facetable: Vec::new(), + facets: Vec::new(), + projections: Vec::new(), + } + } + + fn sample_field_value(field: &str, value: NamespacePredicateValue) -> NamespaceFieldValue { + NamespaceFieldValue { + field: NamespaceFindField::new(field), + value, + source: NamespaceFieldSource { + evidence: "nokv-native:///runs/run-1@generation:7".to_owned(), + source_path: "/runs/run-1".to_owned(), + source_kind: NamespaceFieldSourceKind::MaterializedIndex, + }, + } + } + + fn sample_aggregate_result() -> NamespaceAggregateResult { + NamespaceAggregateResult { + path: "/runs".to_owned(), + evidence: "nokv-native:///runs".to_owned(), + snapshot_id: Some(9), + predicates: Vec::new(), + input_match_count: 5, + row_count: 3, + group_count: 2, + groups: vec![NamespaceAggregateGroup { + key: vec![sample_field_value( + "param.lr", + NamespacePredicateValue::String("0.001".to_owned()), + )], + measures: vec![ + NamespaceAggregateOutputMeasure { + name: "run_count".to_owned(), + op: NamespaceAggregateOp::Count, + field: Some(NamespaceFindField::new("metric.val_loss.min")), + value: NamespaceAggregateValue::U64(2), + }, + NamespaceAggregateOutputMeasure { + name: "avg_min_val_loss".to_owned(), + op: NamespaceAggregateOp::Avg, + field: Some(NamespaceFindField::new("metric.val_loss.min")), + value: NamespaceAggregateValue::F64(0.3), + }, + NamespaceAggregateOutputMeasure { + name: "stdout_available".to_owned(), + op: NamespaceAggregateOp::Sum, + field: Some(NamespaceFindField::new("artifact.stdout_available")), + value: NamespaceAggregateValue::U64(1), + }, + ], + evidence: "nokv-native:///runs".to_owned(), + sample_matches: vec![nokv_meta::NamespaceAggregateSample { + path: "/runs/run-1".to_owned(), + evidence: "nokv-native:///runs/run-1@generation:7".to_owned(), + generation: 7, + }], + }], + truncated: false, + scanned_entries: 5, + } + } + + fn sample_card(path: &str, record_count: usize) -> NamespaceCard { + NamespaceCard { + path: path.to_owned(), + name: path.rsplit('/').next().unwrap_or("").to_owned(), + kind: NamespaceCardKind::File, + evidence: format!("nokv-native://{path}@generation:7"), + snapshot_id: Some(9), + inode: InodeId::new(42).unwrap(), + generation: 7, + size_bytes: Some(13), + entry_count: None, + record_count: Some(NamespaceRecordCount { + count: record_count, + provenance: RecordCountProvenance::MaterializedIndex, + }), + schema: Some(NamespaceSchema { + record_type: NamespaceRecordType::JsonObject, + fields: vec!["status".to_owned()], + }), + sample: vec![r#"{"status":"completed"}"#.to_owned()], + body: Some(NamespaceBodyDescriptor { + producer: "unit-test".to_owned(), + digest_uri: "sha256:test".to_owned(), + size: 13, + content_type: "application/json".to_owned(), + manifest_id: "runs/run-1/metadata.json".to_owned(), + generation: 7, + chunk_size: 4096, + block_size: 4096, + }), + catalog: NamespaceQueryCatalog { + filterable: vec![NamespaceFilterCapability { + field: NamespaceFindField::new("run.status"), + operators: vec![NamespacePredicateOp::Eq], + }], + sortable: vec![NamespaceSortField::new("run.status")], + facetable: vec![NamespaceFindField::new("run.status")], + facets: vec![NamespaceFacetSummary { + field: NamespaceFindField::new("run.status"), + values: vec![NamespaceFacetValue { + value: NamespacePredicateValue::String("completed".to_owned()), + count: 1, + }], + distinct_count: 1, + truncated: false, + }], + projections: vec![ + NamespaceInclude::Body, + NamespaceInclude::Schema, + NamespaceInclude::Sample, + NamespaceInclude::Catalog, + ], + }, + indexed_values: vec![NamespaceIndexValue { + field: NamespaceFindField::new("run.status"), + value: NamespacePredicateValue::String("completed".to_owned()), + }], + } + } + + #[test] + fn agent_tool_registry_uses_posix_like_names() { + let names = agent_tool_definitions() + .into_iter() + .map(|tool| tool.name) + .collect::>(); + + assert_eq!( + names, + vec!["ls", "stat", "catalog", "read", "aggregate", "find", "grep"] + ); + } + + #[test] + fn find_tool_schema_excludes_include_and_caps_limit_at_ten() { + let tools = agent_tool_definitions(); + let find = tools + .iter() + .find(|tool| tool.name == "find") + .expect("find tool must be registered"); + + assert!(find.parameters["properties"].get("include").is_none()); + assert_eq!(find.parameters["properties"]["limit"]["maximum"], 10); + } + + #[test] + fn agent_tool_registry_documents_stable_api_semantics() { + let tools = agent_tool_definitions(); + + assert_tool_description_contains(&tools, "ls", &["direct children", "discover paths"]); + assert_tool_description_contains( + &tools, + "stat", + &[ + "single path", + "compact card", + "body", + "schema", + "sample", + "catalog", + ], + ); + assert_tool_description_contains(&tools, "catalog", &["field ids", "find", "aggregate"]); + assert_tool_description_contains(&tools, "read", &["file body", "json", "yaml", "text"]); + assert_tool_description_contains( + &tools, + "aggregate", + &["catalog field ids", "count", "group", "summary"], + ); + assert_tool_description_contains( + &tools, + "find", + &["search paths", "project fields", "read"], + ); + } + + #[test] + fn agent_tool_parameter_schemas_are_token_compact() { + let tools = agent_tool_definitions(); + let find = tool_definition(&tools, "find"); + let read = tool_definition(&tools, "read"); + + for tool in &tools { + assert_schema_lacks_nested_descriptions(&tool.parameters); + } + assert!(find.parameters["properties"].get("include").is_none()); + assert!(read.parameters["properties"] + .get("expected_generation") + .is_none()); + } + + #[test] + fn stat_tool_returns_compact_card_without_agent_noise() { + let namespace = FakeNamespace::new(); + + let output = execute_agent_tool( + &namespace, + "stat", + &json!({ + "path": "/runs", + }), + ) + .unwrap(); + + assert_json_lacks_agent_noise(&output); + assert_eq!(output["card"]["path"], "/runs"); + assert_eq!( + output["card"]["body"], + json!({ + "producer": "unit-test", + "size": 13, + "content_type": "application/json" + }) + ); + assert_eq!( + output["card"]["catalog"]["filterable"][0]["fields"][0], + "run.status" + ); + assert!(output["card"]["catalog"].get("projections").is_none()); + assert_eq!( + output["card"]["catalog"]["facets"][0], + json!({ + "field": "run.status", + "values": [{"value": "completed", "count": 1}], + "distinct_count": 1, + "truncated": false + }) + ); + assert_eq!(output["card"]["record_count"], 1); + } + + #[test] + fn catalog_tool_returns_compact_filtered_catalog() { + let namespace = FakeNamespace::new(); + + let output = execute_agent_tool( + &namespace, + "catalog", + &json!({ + "path": "/runs", + "field_prefix": "run.", + "include_facets": false + }), + ) + .unwrap(); + + assert_json_lacks_agent_noise(&output); + assert_eq!(output["path"], "/runs"); + assert_eq!( + output["catalog"]["filterable"][0]["fields"], + json!(["run.status"]) + ); + assert_eq!(output["catalog"]["sortable"], json!(["run.status"])); + assert_eq!(output["catalog"]["facetable"], json!(["run.status"])); + assert_eq!(output["catalog"]["facets"], json!([])); + assert!(output.get("card").is_none()); + assert!(output.get("body").is_none()); + } + + #[test] + fn catalog_tool_matches_field_prefix_inside_field_ids() { + let namespace = FakeNamespace::new(); + + let output = execute_agent_tool( + &namespace, + "catalog", + &json!({ + "path": "/runs", + "field_prefix": "status", + "include_facets": false + }), + ) + .unwrap(); + + assert_eq!( + output["catalog"]["filterable"][0]["fields"][0], + "run.status" + ); + assert_eq!(output["catalog"]["sortable"], json!(["run.status"])); + } + + #[test] + fn catalog_tool_suggests_child_catalogs_when_current_scope_is_empty() { + let namespace = FakeNamespace::with_root_child_catalog(); + + let output = execute_agent_tool( + &namespace, + "catalog", + &json!({ + "path": "/yanex", + "field_prefix": "status", + "include_facets": false + }), + ) + .unwrap(); + + assert_eq!(output["catalog_empty"], true); + assert_eq!(output["child_catalogs"][0]["path"], "/yanex/runs"); + let child_filterable = &output["child_catalogs"][0]["catalog"]["filterable"]; + let child_fields = child_filterable + .as_array() + .unwrap() + .iter() + .flat_map(|group| group["fields"].as_array().unwrap().clone()) + .collect::>(); + assert!(child_fields.contains(&json!("run.status"))); + assert!(child_fields.contains(&json!("run.script"))); + assert_json_lacks_agent_noise(&output); + } + + #[test] + fn catalog_tool_does_not_discover_child_catalogs_for_files() { + let namespace = FakeNamespace::new(); + + let output = execute_agent_tool( + &namespace, + "catalog", + &json!({ + "path": "/runs/run-1", + "field_prefix": "missing.", + "include_facets": false + }), + ) + .unwrap(); + + assert_eq!(output["catalog_empty"], true); + assert_eq!(output["child_catalogs"], json!([])); + assert_json_lacks_agent_noise(&output); + } + + #[test] + fn find_tool_translates_catalog_field_ids() { + let namespace = FakeNamespace::new(); + + let output = execute_agent_tool( + &namespace, + "find", + &json!({ + "path": "/runs", + "predicates": [{"field": "run.status", "op": "eq", "value": "completed"}], + "sort": [{"field": "run.status", "direction": "desc"}], + "limit": 5 + }), + ) + .unwrap(); + + let request = namespace.last_find.borrow().clone().unwrap(); + assert_eq!(request.predicates[0].field.id, "run.status"); + assert_eq!( + request.predicates[0].value, + Some(NamespacePredicateValue::String("completed".to_owned())) + ); + assert_eq!(request.sort[0].field.id, "run.status"); + assert_eq!(request.include, Vec::::new()); + assert_eq!(output["match_count"], 1); + assert_eq!(output["matches"][0], json!({"path": "/runs/run-1"})); + assert!(output.get("catalog").is_none() || output["catalog"].is_null()); + assert!(output["matches"][0].get("catalog").is_none()); + assert_json_lacks_agent_noise(&output); + } + + #[test] + fn find_tool_requests_and_returns_filtered_facets() { + let namespace = FakeNamespace::new(); + + let output = execute_agent_tool( + &namespace, + "find", + &json!({ + "path": "/runs", + "predicates": [{"field": "run.status", "op": "eq", "value": "completed"}], + "facets": ["run.script"], + "limit": 1 + }), + ) + .unwrap(); + + let request = namespace.last_find.borrow().clone().unwrap(); + assert_eq!(request.facets, vec![NamespaceFindField::new("run.script")]); + assert_eq!( + output["facets"], + json!([{ + "field": "run.script", + "values": [{"value": "train.py", "count": 1}], + "distinct_count": 1, + "truncated": false + }]) + ); + } + + #[test] + fn find_tool_schema_rejects_body_schema_sample_includes() { + let namespace = FakeNamespace::new(); + + let err = execute_agent_tool( + &namespace, + "find", + &json!({ + "path": "/runs", + "include": ["body", "schema", "sample"], + "limit": 1 + }), + ) + .unwrap_err(); + + assert!( + matches!(err, AgentError::InvalidArgument(ref message) if message.contains("unsupported argument include")), + "unexpected error: {err:?}" + ); + assert!(namespace.last_find.borrow().is_none()); + } + + #[test] + fn find_tool_rejects_catalog_include() { + let namespace = FakeNamespace::new(); + + let err = execute_agent_tool( + &namespace, + "find", + &json!({ + "path": "/runs", + "include": ["catalog"], + "limit": 1 + }), + ) + .unwrap_err(); + + assert!( + matches!(err, AgentError::InvalidArgument(ref message) if message.contains("unsupported argument include")), + "unexpected error: {err:?}" + ); + assert!(namespace.last_find.borrow().is_none()); + } + + #[test] + fn find_tool_defaults_to_limit_ten() { + let namespace = FakeNamespace::new(); + + execute_agent_tool( + &namespace, + "find", + &json!({ + "path": "/runs", + }), + ) + .unwrap(); + + let request = namespace.last_find.borrow().clone().unwrap(); + assert_eq!(request.limit, 10); + } + + #[test] + fn find_tool_rejects_limit_above_ten() { + let namespace = FakeNamespace::new(); + + let err = execute_agent_tool( + &namespace, + "find", + &json!({ + "path": "/runs", + "limit": 11 + }), + ) + .unwrap_err(); + + assert!( + matches!(err, AgentError::InvalidArgument(ref message) if message.contains("limit must be between 1 and 10")), + "unexpected error: {err:?}" + ); + assert!(namespace.last_find.borrow().is_none()); + } + + #[test] + fn find_tool_accepts_boolean_predicates_as_u64_facets() { + let namespace = FakeNamespace::new(); + + execute_agent_tool( + &namespace, + "find", + &json!({ + "path": "/runs", + "predicates": [{"field": "git.has_uncommitted_changes", "op": "eq", "value": true}], + "limit": 5 + }), + ) + .unwrap(); + + let request = namespace.last_find.borrow().clone().unwrap(); + assert_eq!( + request.predicates[0].field.id, + "git.has_uncommitted_changes" + ); + assert_eq!( + request.predicates[0].value, + Some(NamespacePredicateValue::U64(1)) + ); + } + + #[test] + fn find_tool_projects_indexed_values_when_fields_are_requested() { + let namespace = FakeNamespace::new(); + + let output = execute_agent_tool( + &namespace, + "find", + &json!({ + "path": "/runs", + "fields": ["run.id"], + "limit": 5 + }), + ) + .unwrap(); + + assert_eq!(output["matches"][0]["path"], "/runs/run-1"); + assert_eq!(output["matches"][0]["values"], json!({})); + assert!(output["matches"][0].get("indexed_values").is_none()); + assert!(output["matches"][0].get("name").is_none()); + assert!(output.get("catalog").is_none() || output["catalog"].is_null()); + assert!(namespace.last_find.borrow().is_some()); + assert_json_lacks_agent_noise(&output); + } + + #[test] + fn find_tool_projects_requested_fields_without_sources() { + let namespace = FakeNamespace::new(); + + let output = execute_agent_tool( + &namespace, + "find", + &json!({ + "path": "/runs", + "fields": ["run.status"], + "limit": 5 + }), + ) + .unwrap(); + + let match_ = &output["matches"][0]; + assert_eq!(match_["path"], "/runs/run-1"); + assert_eq!(match_["values"], json!({"run.status": "completed"})); + assert!(match_.get("indexed_values").is_none()); + assert!(match_.get("name").is_none()); + assert_json_lacks_agent_noise(&output); + } + + #[test] + fn find_tool_does_not_project_storage_internal_body_fields() { + let namespace = FakeNamespace::new(); + + let output = execute_agent_tool( + &namespace, + "find", + &json!({ + "path": "/runs", + "fields": ["body.manifest_id"], + "limit": 5 + }), + ) + .unwrap(); + + let match_ = &output["matches"][0]; + assert_eq!(match_["path"], "/runs/run-1"); + assert_eq!(match_["values"], json!({})); + assert_json_lacks_agent_noise(&output); + } + + #[test] + fn grep_tool_returns_compact_matches_without_agent_noise() { + let namespace = FakeNamespace::new(); + + let output = execute_agent_tool( + &namespace, + "grep", + &json!({ + "path": "/runs/run-1", + "pattern": "best_model", + "recursive": true, + "limit": 5 + }), + ) + .unwrap(); + + assert_eq!(output["path"], "/runs/run-1"); + assert_eq!(output["matches"][0]["line_number"], 3); + assert_eq!( + output["matches"][0]["snippet"], + "Checkpoint: best_model_1.0_1.pt" + ); + assert!(output.get("tool").is_none()); + assert!(output.get("bytes_read").is_none()); + assert!(output["matches"][0].get("evidence").is_none()); + assert_json_lacks_agent_noise(&output); + } + + #[test] + fn aggregate_tool_groups_and_sorts_indexed_values() { + let namespace = FakeNamespace::with_aggregate_result(sample_aggregate_result()); + + let output = execute_agent_tool( + &namespace, + "aggregate", + &json!({ + "path": "/runs", + "predicates": [{"field": "run.status", "op": "eq", "value": "completed"}], + "group_by": ["param.lr"], + "measures": [ + {"name": "run_count", "op": "count", "field": "metric.val_loss.min"}, + {"name": "avg_min_val_loss", "op": "avg", "field": "metric.val_loss.min"}, + {"name": "stdout_available", "op": "sum", "field": "artifact.stdout_available"} + ], + "sort": [{"field": "avg_min_val_loss", "direction": "asc"}], + "limit": 5 + }), + ) + .unwrap(); + + let request = namespace.last_aggregate.borrow().clone().unwrap(); + assert_eq!(request.predicates[0].field.id, "run.status"); + assert!(namespace.last_find.borrow().is_none()); + assert_eq!(output["input_match_count"], 5); + assert_eq!(output["row_count"], 3); + assert_eq!(output["group_count"], 2); + assert_eq!(output["groups"][0]["key"], json!({"param.lr": "0.001"})); + assert_eq!(output["groups"][0]["values"]["run_count"], 2); + assert_eq!(output["groups"][0]["values"]["avg_min_val_loss"], 0.3); + assert_eq!(output["groups"][0]["values"]["stdout_available"], 1); + assert_json_lacks_agent_noise(&output); + } + + #[test] + fn ls_tool_keeps_entries_lightweight() { + let namespace = FakeNamespace::new(); + + let output = execute_agent_tool( + &namespace, + "ls", + &json!({ + "path": "/runs", + "limit": 1 + }), + ) + .unwrap(); + + assert_eq!(output["entries"][0]["path"], "/runs/run-1"); + assert!(output["entries"][0].get("record_count").is_none()); + assert!(output["entries"][0].get("sample").is_none()); + assert!(output["entries"][0].get("catalog").is_none()); + assert_json_lacks_agent_noise(&output); + } + + #[test] + fn read_tool_returns_body_content_without_evidence_noise() { + let namespace = FakeNamespace::new(); + + let output = execute_agent_tool( + &namespace, + "read", + &json!({ + "path": "/runs/run-1/metadata.json", + "format": "structured", + "limit": 10 + }), + ) + .unwrap(); + + assert_eq!(output["path"], "/runs/run-1/metadata.json"); + assert_eq!(output["items"][0]["index"], 0); + assert_eq!(output["items"][0]["value"]["status"], "completed"); + assert_json_lacks_agent_noise(&output); + } + + #[test] + fn read_tool_rejects_large_structured_pagination() { + let namespace = FakeNamespace::with_record_count(350); + + let err = execute_agent_tool( + &namespace, + "read", + &json!({ + "path": "/index/large.json", + "format": "structured", + "cursor": "100", + "limit": 100 + }), + ) + .unwrap_err(); + + assert!( + matches!(err, AgentError::InvalidArgument(ref message) if message.contains("use bytes format with offset and limit, grep to locate lines, or stat record_count")), + "unexpected error: {err:?}" + ); + assert_eq!(*namespace.read_calls.borrow(), 0); + } + + #[test] + fn read_tool_rejects_large_structured_initial_page() { + let namespace = FakeNamespace::with_record_count(350); + + let err = execute_agent_tool( + &namespace, + "read", + &json!({ + "path": "/index/large.json", + "format": "structured", + "limit": 100 + }), + ) + .unwrap_err(); + + assert!( + matches!(err, AgentError::InvalidArgument(ref message) if message.contains("use bytes format with offset and limit, grep to locate lines, or stat record_count")), + "unexpected error: {err:?}" + ); + assert_eq!(*namespace.read_calls.borrow(), 0); + } + + #[test] + fn read_tool_reads_structured_files_within_guard_threshold() { + let namespace = FakeNamespace::with_record_count(250); + + execute_agent_tool( + &namespace, + "read", + &json!({ + "path": "/index/mid.json", + "format": "structured" + }), + ) + .unwrap(); + + assert_eq!(*namespace.read_calls.borrow(), 1); + } + + #[test] + fn read_tool_structured_guard_allows_exactly_the_read_limit() { + let namespace = FakeNamespace::with_record_count(MAX_AGENT_READ_LIMIT); + + execute_agent_tool( + &namespace, + "read", + &json!({ + "path": "/index/boundary.json", + "format": "structured" + }), + ) + .unwrap(); + + assert_eq!(*namespace.read_calls.borrow(), 1); + } + + #[test] + fn read_tool_returns_generation_for_conditional_reads() { + let namespace = FakeNamespace::new(); + + let output = execute_agent_tool( + &namespace, + "read", + &json!({ + "path": "/runs/run-1/metadata.json", + }), + ) + .unwrap(); + + assert_eq!(output["generation"], 7); + assert_json_lacks_agent_noise(&output); + } + + #[test] + fn read_tool_schema_caps_limit_at_three_hundred() { + let tools = agent_tool_definitions(); + let read = tool_definition(&tools, "read"); + + assert_eq!(read.parameters["properties"]["limit"]["maximum"], 300); + } + + #[test] + fn read_tool_accepts_limit_of_three_hundred() { + let namespace = FakeNamespace::new(); + + execute_agent_tool( + &namespace, + "read", + &json!({ + "path": "/runs/run-1/metadata.json", + "limit": 300 + }), + ) + .unwrap(); + + let options = namespace.last_read.borrow().clone().unwrap(); + assert_eq!(options.limit, 300); + } + + #[test] + fn read_tool_rejects_limit_above_three_hundred() { + let namespace = FakeNamespace::new(); + + let err = execute_agent_tool( + &namespace, + "read", + &json!({ + "path": "/runs/run-1/metadata.json", + "limit": 301 + }), + ) + .unwrap_err(); + + assert!( + matches!(err, AgentError::InvalidArgument(ref message) if message.contains("limit must be between 1 and 300")), + "unexpected error: {err:?}" + ); + assert_eq!(*namespace.read_calls.borrow(), 0); + } + + #[test] + fn ls_tool_keeps_limit_capped_at_one_hundred() { + let namespace = FakeNamespace::new(); + + let err = execute_agent_tool( + &namespace, + "ls", + &json!({ + "path": "/runs", + "limit": 101 + }), + ) + .unwrap_err(); + + assert!( + matches!(err, AgentError::InvalidArgument(ref message) if message.contains("limit must be between 1 and 100")), + "unexpected error: {err:?}" + ); + } + + #[test] + fn grep_tool_schema_caps_limit_and_exposes_patterns_and_glob() { + let tools = agent_tool_definitions(); + let grep = tool_definition(&tools, "grep"); + + assert_eq!(grep.parameters["properties"]["limit"]["maximum"], 300); + assert_eq!( + grep.parameters["properties"]["patterns"], + json!({"type": "array", "items": {"type": "string"}, "maxItems": 16}) + ); + assert_eq!( + grep.parameters["properties"]["glob"], + json!({"type": ["string", "null"]}) + ); + assert_eq!( + grep.parameters["required"], + json!(["path", "pattern", "recursive"]) + ); + } + + #[test] + fn grep_tool_accepts_limit_of_three_hundred() { + let namespace = FakeNamespace::new(); + + execute_agent_tool( + &namespace, + "grep", + &json!({ + "path": "/runs/run-1", + "pattern": "best_model", + "recursive": true, + "limit": 300 + }), + ) + .unwrap(); + + let request = namespace.last_grep.borrow().clone().unwrap(); + assert_eq!(request.limit, 300); + } + + #[test] + fn grep_tool_rejects_limit_above_three_hundred() { + let namespace = FakeNamespace::new(); + + let err = execute_agent_tool( + &namespace, + "grep", + &json!({ + "path": "/runs/run-1", + "pattern": "best_model", + "recursive": true, + "limit": 301 + }), + ) + .unwrap_err(); + + assert!( + matches!(err, AgentError::InvalidArgument(ref message) if message.contains("limit must be between 1 and 300")), + "unexpected error: {err:?}" + ); + assert!(namespace.last_grep.borrow().is_none()); + } + + #[test] + fn grep_tool_rejects_non_string_patterns() { + let namespace = FakeNamespace::new(); + + let err = execute_agent_tool( + &namespace, + "grep", + &json!({ + "path": "/runs/run-1", + "pattern": "best_model", + "recursive": true, + "patterns": [1] + }), + ) + .unwrap_err(); + + assert!( + matches!(err, AgentError::InvalidArgument(ref message) if message.contains("patterns entries must be strings")), + "unexpected error: {err:?}" + ); + assert!(namespace.last_grep.borrow().is_none()); + } + + #[test] + fn grep_tool_rejects_non_string_glob() { + let namespace = FakeNamespace::new(); + + let err = execute_agent_tool( + &namespace, + "grep", + &json!({ + "path": "/runs/run-1", + "pattern": "best_model", + "recursive": true, + "glob": 7 + }), + ) + .unwrap_err(); + + assert!( + matches!(err, AgentError::InvalidArgument(ref message) if message.contains("glob must be a string or null")), + "unexpected error: {err:?}" + ); + assert!(namespace.last_grep.borrow().is_none()); + } + + #[test] + fn grep_tool_forwards_patterns_and_glob() { + let namespace = FakeNamespace::new(); + + execute_agent_tool( + &namespace, + "grep", + &json!({ + "path": "/runs/run-1", + "pattern": "best_model", + "patterns": ["error", "warning"], + "glob": "*.txt", + "recursive": true + }), + ) + .unwrap(); + + let request = namespace.last_grep.borrow().clone().unwrap(); + assert_eq!( + request.patterns, + vec!["error".to_owned(), "warning".to_owned()] + ); + assert_eq!(request.name_glob.as_deref(), Some("*.txt")); + } + + #[test] + fn grep_tool_defaults_patterns_and_glob_to_unset() { + let namespace = FakeNamespace::new(); + + execute_agent_tool( + &namespace, + "grep", + &json!({ + "path": "/runs/run-1", + "pattern": "best_model", + "recursive": true + }), + ) + .unwrap(); + + let request = namespace.last_grep.borrow().clone().unwrap(); + assert_eq!(request.patterns, Vec::::new()); + assert_eq!(request.name_glob, None); + } +} diff --git a/crates/nokv-agent/tests/event_index.rs b/crates/nokv-agent/tests/event_index.rs new file mode 100644 index 000000000..cdb75fe08 --- /dev/null +++ b/crates/nokv-agent/tests/event_index.rs @@ -0,0 +1,704 @@ +use std::io::Cursor; +use std::process::Command; + +use nokv_agent::event::{ + ingest_jsonl_reader, AgentEventStore, CompletionAfterRequest, ErrorEventsRequest, + HoltAgentEventStore, JsonlIngestOptions, LatestEventsRequest, NotificationLifecycleRequest, + NotificationNeighborDirection, NotificationNeighborRequest, RecentTimesRequest, + SessionEventsRequest, SessionRowsRequest, ToolFacetRequest, ToolTraceRequest, + TuiClearCompletionRequest, +}; + +fn ingest(store: &HoltAgentEventStore, jsonl: &str) -> nokv_agent::event::IngestReport { + ingest_jsonl_reader( + store, + JsonlIngestOptions { + agent_id: "agent-a".to_owned(), + source_file: "logs/events.jsonl".to_owned(), + file_size: jsonl.len() as u64, + }, + Cursor::new(jsonl.as_bytes()), + ) + .unwrap() +} + +#[test] +fn ingest_replay_is_idempotent_and_keeps_coverage_monotonic() { + let store = HoltAgentEventStore::open_memory().unwrap(); + let jsonl = concat!( + r#"{"type":"tool_call","ts":1.0,"tool_name":"read","tool_call_id":"call-1","tool_args":{"file_path":"/tmp/a.md","offset":10,"limit":20}}"#, + "\n", + r#"{"type":"tool_result","ts":2.0,"tool_name":"read","tool_call_id":"call-1","status":"ok"}"#, + "\n", + ); + + let first = ingest(&store, jsonl); + assert_eq!(first.accepted, 2); + assert_eq!(first.duplicates, 0); + assert_eq!(first.coverage.row_count, 2); + assert_eq!(first.coverage.min_offset, Some(0)); + assert!(first.coverage.max_offset.unwrap() > 0); + + let second = ingest(&store, jsonl); + assert_eq!(second.accepted, 0); + assert_eq!(second.duplicates, 2); + + let coverage = store + .coverage("agent-a", "logs/events.jsonl") + .unwrap() + .unwrap(); + assert_eq!(coverage.row_count, 2); + assert_eq!(coverage.file_size, jsonl.len() as u64); +} + +#[test] +fn chunked_ingest_keeps_distinct_type_index_entries() { + let store = HoltAgentEventStore::open_memory().unwrap(); + let mut jsonl = String::new(); + for index in 0..30_000 { + jsonl.push_str( + &serde_json::json!({ + "type": "tool_call", + "ts": index as f64, + "tool_name": "read", + "tool_call_id": format!("call-{index}"), + }) + .to_string(), + ); + jsonl.push('\n'); + } + + let report = ingest(&store, &jsonl); + assert_eq!(report.accepted, 30_000); + + let latest = store + .latest_events(LatestEventsRequest { + agent_id: "agent-a".to_owned(), + event_type: "tool_call".to_owned(), + limit: 3, + }) + .unwrap(); + assert_eq!(latest.len(), 3); + assert_eq!( + latest + .iter() + .map(|record| record.projection.tool_call_id.as_deref()) + .collect::>(), + vec![Some("call-29999"), Some("call-29998"), Some("call-29997")] + ); + assert_eq!( + store + .latest_events(LatestEventsRequest { + agent_id: "agent-a".to_owned(), + event_type: "tool_call".to_owned(), + limit: 40_000, + }) + .unwrap() + .len(), + 30_000 + ); + assert_eq!( + store + .stream_session_rows(SessionRowsRequest { + agent_id: "agent-a".to_owned(), + limit: None, + }) + .unwrap() + .len(), + 30_000 + ); +} + +#[test] +fn event_queries_match_lingtai_sqlite_shapes() { + let store = HoltAgentEventStore::open_memory().unwrap(); + let jsonl = concat!( + r#"{"type":"thinking","ts":0.5,"text":"hidden"}"#, + "\n", + r#"{"type":"tool_call","ts":1.0,"tool_name":"bash","tool_call_id":"call-1","tool_args":{"action":"run","command":"python3 script.py"}}"#, + "\n", + r#"{"type":"tool_result","ts":2.0,"tool_name":"bash","tool_call_id":"call-1","status":"ok"}"#, + "\n", + r#"{"type":"tool_call","ts":3.0,"tool_name":"read","tool_call_id":"call-2","tool_args":{"file_path":"/tmp/a.md","offset":10,"limit":20}}"#, + "\n", + r#"{"type":"notification_block_injected","ts":4.0,"call_id":"notify-call-1","channel":"work"}"#, + "\n", + r#"{"type":"system_notification_published","ts":5.0,"event_id":"event-1","ref_id":"ref-1"}"#, + "\n", + r#"{"type":"notification_event_dismiss","ts":6.0,"event_id":"event-1","ref_id":"ref-1","channel":"work"}"#, + "\n", + r#"{"type":"refresh_complete","ts":7.0}"#, + "\n", + r#"{"type":"psyche_molt","ts":8.0}"#, + "\n", + r#"{"type":"aed_attempt","ts":9.0,"error":"over window"}"#, + "\n", + r#"{"type":"clear_received","ts":10.0,"source":"tui"}"#, + "\n", + ); + ingest(&store, jsonl); + + let latest = store + .latest_events(LatestEventsRequest { + agent_id: "agent-a".to_owned(), + event_type: "tool_call".to_owned(), + limit: 2, + }) + .unwrap(); + assert_eq!(latest.len(), 2); + assert_eq!(latest[0].projection.tool_name.as_deref(), Some("read")); + assert_eq!( + latest[1].projection.command_head.as_deref(), + Some("python3") + ); + + let session = store + .stream_session_events(SessionEventsRequest { + agent_id: "agent-a".to_owned(), + event_types: vec!["thinking".to_owned(), "tool_result".to_owned()], + limit: None, + }) + .unwrap(); + assert_eq!( + session + .iter() + .map(|record| record.event_type.as_str()) + .collect::>(), + vec!["thinking", "tool_result"] + ); + let session_rows = store + .stream_session_rows(SessionRowsRequest { + agent_id: "agent-a".to_owned(), + limit: Some(3), + }) + .unwrap(); + assert_eq!( + session_rows + .iter() + .map(|row| row.event_type.as_str()) + .collect::>(), + vec!["thinking", "tool_call", "tool_result"] + ); + assert_eq!(session_rows[0].fields_json["text"], "hidden"); + assert_eq!(session_rows[1].fields_json["tool_name"], "bash"); + + let facets = store + .tool_facets(ToolFacetRequest { + agent_id: "agent-a".to_owned(), + limit: 4, + }) + .unwrap(); + assert_eq!(facets[0].tool_name, "bash"); + assert_eq!(facets[0].action.as_deref(), Some("run")); + assert_eq!(facets[0].count, 1); + + let trace = store + .tool_trace(ToolTraceRequest { + agent_id: "agent-a".to_owned(), + tool_call_id: "call-1".to_owned(), + }) + .unwrap(); + assert_eq!(trace.len(), 2); + assert_eq!( + trace + .iter() + .map(|record| record.event_type.as_str()) + .collect::>(), + vec!["tool_call", "tool_result"] + ); + + let refresh_times = store + .recent_times(RecentTimesRequest { + agent_id: "agent-a".to_owned(), + event_type: "refresh_complete".to_owned(), + limit: 10, + }) + .unwrap(); + assert_eq!(refresh_times.len(), 1); + assert_eq!(refresh_times[0].ts, 7.0); + + let molt_windows = store.molt_session_windows("agent-a").unwrap(); + assert!(molt_windows.ok); + assert_eq!(molt_windows.current_since, Some(8.0)); + assert_eq!(molt_windows.last_since, None); + assert_eq!(molt_windows.last_before, None); + + let error_events = store + .error_events(ErrorEventsRequest { + agent_id: "agent-a".to_owned(), + event_types: vec!["aed_attempt".to_owned(), "refresh_init_error".to_owned()], + limit: 10, + }) + .unwrap(); + assert_eq!(error_events.len(), 1); + assert_eq!( + error_events[0].projection.error.as_deref(), + Some("over window") + ); + + let completion = store + .completion_after(CompletionAfterRequest { + agent_id: "agent-a".to_owned(), + event_type: "clear_received".to_owned(), + source_offset: 0, + }) + .unwrap(); + assert!(completion.found); + assert_eq!( + completion.event.unwrap().event_type.as_str(), + "clear_received" + ); + let clear_received = store + .latest_events(LatestEventsRequest { + agent_id: "agent-a".to_owned(), + event_type: "clear_received".to_owned(), + limit: 1, + }) + .unwrap() + .pop() + .unwrap(); + let clear_completion = store + .tui_clear_completion(TuiClearCompletionRequest { + agent_id: "agent-a".to_owned(), + source_offset: clear_received.source_offset, + }) + .unwrap(); + assert!(clear_completion.found); + let no_clear_completion = store + .tui_clear_completion(TuiClearCompletionRequest { + agent_id: "agent-a".to_owned(), + source_offset: clear_received.source_offset + 1, + }) + .unwrap(); + assert!(!no_clear_completion.found); + + let notification_by_ref = store + .notification_lifecycle(NotificationLifecycleRequest { + agent_id: "agent-a".to_owned(), + ref_id: Some("ref-1".to_owned()), + limit: 10, + ..NotificationLifecycleRequest::default() + }) + .unwrap(); + assert_eq!( + notification_by_ref + .iter() + .map(|record| record.event_type.as_str()) + .collect::>(), + vec![ + "system_notification_published", + "notification_event_dismiss" + ] + ); + + let notification_by_call = store + .notification_lifecycle(NotificationLifecycleRequest { + agent_id: "agent-a".to_owned(), + call_id: Some("notify-call-1".to_owned()), + limit: 10, + ..NotificationLifecycleRequest::default() + }) + .unwrap(); + assert_eq!(notification_by_call.len(), 1); + assert_eq!( + notification_by_call[0].event_type.as_str(), + "notification_block_injected" + ); + + let notification_events = store + .notification_events(nokv_agent::event::NotificationEventsRequest { + agent_id: "agent-a".to_owned(), + limit: 10, + }) + .unwrap(); + assert_eq!( + notification_events + .iter() + .map(|record| record.event_type.as_str()) + .collect::>(), + vec![ + "notification_event_dismiss", + "system_notification_published", + "notification_block_injected" + ] + ); + let middle_id = notification_events[1].id; + let by_id = store + .notification_event_by_id(nokv_agent::event::NotificationEventByIdRequest { + agent_id: "agent-a".to_owned(), + event_id: middle_id, + }) + .unwrap() + .unwrap(); + assert_eq!(by_id.event_type, "system_notification_published"); + let before = store + .notification_neighbor(NotificationNeighborRequest { + agent_id: "agent-a".to_owned(), + pivot_event_id: middle_id, + direction: NotificationNeighborDirection::Before, + }) + .unwrap() + .unwrap(); + assert_eq!(before.event_type, "notification_block_injected"); + let after = store + .notification_neighbor(NotificationNeighborRequest { + agent_id: "agent-a".to_owned(), + pivot_event_id: middle_id, + direction: NotificationNeighborDirection::After, + }) + .unwrap() + .unwrap(); + assert_eq!(after.event_type, "notification_event_dismiss"); +} + +#[test] +fn partial_trailing_line_is_not_indexed() { + let store = HoltAgentEventStore::open_memory().unwrap(); + let jsonl = concat!( + r#"{"type":"tool_call","ts":1.0,"tool_name":"read","tool_call_id":"call-1"}"#, + "\n", + r#"{"type":"tool_call","ts":2.0,"tool_name":"bash""#, + ); + + let report = ingest(&store, jsonl); + assert_eq!(report.accepted, 1); + assert_eq!(report.partial_lines, 1); + assert_eq!(report.parse_errors, 0); + assert_eq!( + store + .latest_events(LatestEventsRequest { + agent_id: "agent-a".to_owned(), + event_type: "tool_call".to_owned(), + limit: 10, + }) + .unwrap() + .len(), + 1 + ); +} + +#[test] +fn large_fields_json_is_compacted_for_holt_value_limit() { + let store = HoltAgentEventStore::open_memory().unwrap(); + let large_text = "x".repeat(70 * 1024); + let jsonl = format!( + "{}\n", + serde_json::json!({ + "type": "tool_result", + "ts": 1.0, + "tool_name": "bash", + "tool_call_id": "call-large", + "result": large_text, + }) + ); + + let report = ingest(&store, &jsonl); + assert_eq!(report.accepted, 1); + + let events = store + .latest_events(LatestEventsRequest { + agent_id: "agent-a".to_owned(), + event_type: "tool_result".to_owned(), + limit: 1, + }) + .unwrap(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].projection.tool_name.as_deref(), Some("bash")); + assert_eq!(events[0].fields_json["_nokv_compacted"], true); + assert_eq!(events[0].fields_json["keys"][0], "result"); +} + +#[test] +fn file_backed_holt_store_reopens_indexed_events() { + let dir = tempfile::tempdir().unwrap(); + let jsonl = concat!( + r#"{"type":"tool_call","ts":1.0,"tool_name":"read","tool_call_id":"call-1","tool_args":{"file_path":"/tmp/a.md"}}"#, + "\n", + ); + { + let store = HoltAgentEventStore::open_file(dir.path()).unwrap(); + let report = ingest(&store, jsonl); + assert_eq!(report.accepted, 1); + } + { + let store = HoltAgentEventStore::open_file(dir.path()).unwrap(); + let coverage = store + .coverage("agent-a", "logs/events.jsonl") + .unwrap() + .unwrap(); + assert_eq!(coverage.row_count, 1); + let latest = store + .latest_events(LatestEventsRequest { + agent_id: "agent-a".to_owned(), + event_type: "tool_call".to_owned(), + limit: 1, + }) + .unwrap(); + assert_eq!(latest.len(), 1); + assert_eq!(latest[0].projection.tool_name.as_deref(), Some("read")); + } +} + +#[test] +fn nokv_agent_lingtai_cli_round_trips_json() { + let dir = tempfile::tempdir().unwrap(); + let store_dir = dir.path().join("agent-index"); + let events_path = dir.path().join("events.jsonl"); + std::fs::write( + &events_path, + concat!( + r#"{"type":"tool_call","ts":1.0,"tool_name":"read","tool_call_id":"call-1","tool_args":{"file_path":"/tmp/a.md"}}"#, + "\n", + r#"{"type":"notification_pair_injected","ts":2.0,"call_id":"notify-call-1","summary":"note"}"#, + "\n", + r#"{"type":"notification_block_injected","ts":3.0,"call_id":"notify-call-1","summary":"block"}"#, + "\n", + r#"{"type":"clear_received","ts":4.0,"source":"tui"}"#, + "\n", + ), + ) + .unwrap(); + + let bin = env!("CARGO_BIN_EXE_nokv-agent"); + let ingest = Command::new(bin) + .args([ + "lingtai", + "ingest", + "--store", + store_dir.to_str().unwrap(), + "--events", + events_path.to_str().unwrap(), + "--agent-id", + "agent-a", + "--source-file", + "logs/events.jsonl", + ]) + .output() + .unwrap(); + assert!( + ingest.status.success(), + "{}", + String::from_utf8_lossy(&ingest.stderr) + ); + let ingest_json: serde_json::Value = serde_json::from_slice(&ingest.stdout).unwrap(); + assert_eq!(ingest_json["report"]["accepted"], 4); + + let coverage = Command::new(bin) + .args([ + "lingtai", + "coverage", + "--store", + store_dir.to_str().unwrap(), + "--agent-id", + "agent-a", + "--source-file", + "logs/events.jsonl", + ]) + .output() + .unwrap(); + assert!( + coverage.status.success(), + "{}", + String::from_utf8_lossy(&coverage.stderr) + ); + let coverage_json: serde_json::Value = serde_json::from_slice(&coverage.stdout).unwrap(); + assert_eq!(coverage_json["coverage"]["row_count"], 4); + + let latest = Command::new(bin) + .args([ + "lingtai", + "latest", + "--store", + store_dir.to_str().unwrap(), + "--agent-id", + "agent-a", + "--type", + "tool_call", + "--limit", + "1", + ]) + .output() + .unwrap(); + assert!( + latest.status.success(), + "{}", + String::from_utf8_lossy(&latest.stderr) + ); + let latest_json: serde_json::Value = serde_json::from_slice(&latest.stdout).unwrap(); + assert_eq!(latest_json["events"][0]["projection"]["tool_name"], "read"); + + let session_rows = Command::new(bin) + .args([ + "lingtai", + "session-rows", + "--store", + store_dir.to_str().unwrap(), + "--agent-id", + "agent-a", + "--limit", + "1", + ]) + .output() + .unwrap(); + assert!( + session_rows.status.success(), + "{}", + String::from_utf8_lossy(&session_rows.stderr) + ); + let session_rows_json: serde_json::Value = + serde_json::from_slice(&session_rows.stdout).unwrap(); + assert_eq!(session_rows_json["rows"][0]["event_type"], "tool_call"); + assert_eq!( + session_rows_json["rows"][0]["fields_json"]["tool_name"], + "read" + ); + + let recent = Command::new(bin) + .args([ + "lingtai", + "recent", + "--store", + store_dir.to_str().unwrap(), + "--agent-id", + "agent-a", + "--type", + "tool_call", + "--limit", + "1", + ]) + .output() + .unwrap(); + assert!( + recent.status.success(), + "{}", + String::from_utf8_lossy(&recent.stderr) + ); + let recent_json: serde_json::Value = serde_json::from_slice(&recent.stdout).unwrap(); + assert_eq!(recent_json["times"][0]["ts"], 1.0); + + let notification_blocks = Command::new(bin) + .args([ + "lingtai", + "notification-blocks", + "--store", + store_dir.to_str().unwrap(), + "--agent-id", + "agent-a", + "--limit", + "1", + ]) + .output() + .unwrap(); + assert!( + notification_blocks.status.success(), + "{}", + String::from_utf8_lossy(¬ification_blocks.stderr) + ); + let notification_blocks_json: serde_json::Value = + serde_json::from_slice(¬ification_blocks.stdout).unwrap(); + assert_eq!( + notification_blocks_json["events"][0]["event_type"], + "notification_pair_injected" + ); + + let notification_block_snapshots = Command::new(bin) + .args([ + "lingtai", + "notification-block-snapshots", + "--store", + store_dir.to_str().unwrap(), + "--agent-id", + "agent-a", + "--limit", + "1", + ]) + .output() + .unwrap(); + assert!( + notification_block_snapshots.status.success(), + "{}", + String::from_utf8_lossy(¬ification_block_snapshots.stderr) + ); + let notification_block_snapshots_json: serde_json::Value = + serde_json::from_slice(¬ification_block_snapshots.stdout).unwrap(); + assert_eq!( + notification_block_snapshots_json["events"][0]["event_type"], + "notification_block_injected" + ); + + let notification_events = Command::new(bin) + .args([ + "lingtai", + "notification-events", + "--store", + store_dir.to_str().unwrap(), + "--agent-id", + "agent-a", + "--limit", + "1", + ]) + .output() + .unwrap(); + assert!( + notification_events.status.success(), + "{}", + String::from_utf8_lossy(¬ification_events.stderr) + ); + let notification_events_json: serde_json::Value = + serde_json::from_slice(¬ification_events.stdout).unwrap(); + assert_eq!( + notification_events_json["events"][0]["event_type"], + "notification_block_injected" + ); + let notification_id = notification_events_json["events"][0]["id"] + .as_u64() + .unwrap() + .to_string(); + + let notification_by_id = Command::new(bin) + .args([ + "lingtai", + "notification-by-id", + "--store", + store_dir.to_str().unwrap(), + "--agent-id", + "agent-a", + "--event-id", + ¬ification_id, + ]) + .output() + .unwrap(); + assert!( + notification_by_id.status.success(), + "{}", + String::from_utf8_lossy(¬ification_by_id.stderr) + ); + let notification_by_id_json: serde_json::Value = + serde_json::from_slice(¬ification_by_id.stdout).unwrap(); + assert_eq!( + notification_by_id_json["event"]["event_type"], + "notification_block_injected" + ); + + let clear_completion = Command::new(bin) + .args([ + "lingtai", + "clear-completion", + "--store", + store_dir.to_str().unwrap(), + "--agent-id", + "agent-a", + "--source-offset", + "0", + ]) + .output() + .unwrap(); + assert!( + clear_completion.status.success(), + "{}", + String::from_utf8_lossy(&clear_completion.stderr) + ); + let clear_completion_json: serde_json::Value = + serde_json::from_slice(&clear_completion.stdout).unwrap(); + assert_eq!(clear_completion_json["completion"]["found"], true); +} diff --git a/docs/development/code_contract.md b/docs/development/code_contract.md index e1a5239f4..16d174f12 100644 --- a/docs/development/code_contract.md +++ b/docs/development/code_contract.md @@ -22,6 +22,7 @@ paths unless the PR states the operational need and removal condition. | `crates/nokv-meta/` | Metadata schema, Holt-friendly layout, `MetadataCommand`, Holt-backed metastore, in-process metadata service, snapshot retention pins, object-reference GC queue policy, history pruning, and service-level GC workers. | Own provider-specific object behavior, FUSE/kernel cache policy, Python bindings, CSI behavior, or wire-protocol migration behavior. | | `crates/nokv-control/` | Shard map, owner leases, shard epochs, routing metadata, checkpoint/log pointers, and failover coordination. | Own inode/dentry semantics, chunk manifests, object GC policy, Holt internals, data-plane cache placement, FUSE behavior, or provider-specific object-store behavior. | | `crates/nokv-object/` | Object-store boundary, S3-compatible backend, batch object reads, local hot-tier object store, tiered data-fabric helpers, soft placement resolution, and in-memory test object store for file bodies. | Own namespace metadata, import Holt/FUSE/protobuf, implement metadata transactions, or expose filesystem-directory object storage as a product backend. | +| `crates/nokv-agent/` | Transport-free agent surfaces: the NoKV namespace tool dispatcher, model-facing result shaping, LingTai-derived event-index types/query contracts, and the local agent-index CLI. | Depend on `nokv-client`, `nokv-protocol`, or `nokv-control`; own NoKV namespace metadata semantics; treat LingTai JSONL-derived indexes as source of truth; hide benchmark-only behavior in product APIs. | | `crates/nokv-client/` | Path-oriented Rust SDK over the metadata service and object backend. | Own metadata layout, bypass `nokv-meta`, expose object-store internals, implement FUSE/kernel cache semantics, depend on `nokv-fuse`, or define metadata wire formats. | | `crates/nokv-fuse/` | FUSE low-level frontend, inode mapping, kernel-facing attr conversion, range reads, and close-to-open buffered writes through the metadata client/server boundary. | Resolve paths through the path SDK hot path, own metadata layout, import Holt directly, open a production metadata store, or implement object-provider-specific behavior. | | `crates/nokv-python/` | Python SDK and fsspec binding for training workflows over `nokv-client`. | Own metadata layout, bypass `nokv-client`, implement object-provider-specific behavior, import FUSE, or reimplement Rust SDK range planning in Python. | diff --git a/docs/development/nokv-agent.md b/docs/development/nokv-agent.md index 1ef619c67..204b51948 100644 --- a/docs/development/nokv-agent.md +++ b/docs/development/nokv-agent.md @@ -5,15 +5,22 @@ SPDX-License-Identifier: Apache-2.0 # `nokv-agent` — Contributor Handbook -`nokv-agent` is the **agent tool surface** crate: the LLM-facing tool -definitions, the dispatcher that maps a tool call onto a namespace verb, -argument validation, result shaping, and the transport-neutral `AgentError`. -It is deliberately small, deliberately **transport-free**, and deliberately -**read-only** today. - -This handbook is for contributors touching the agent surface. It covers where +`nokv-agent` is the **agent surface** crate. Its current shipped surface is the +LLM-facing NoKV namespace tool layer: tool definitions, the dispatcher that maps +a tool call onto a namespace verb, argument validation, result shaping, and the +transport-neutral `AgentError`. That surface is deliberately **transport-free** +and **read-only**. + +The crate is also the home for LingTai-facing agent event indexes. That index is +a derived local view over LingTai `logs/events.jsonl`, not a NoKV DFS namespace +adapter and not a SQL-compatibility layer. Keep this second surface separate +from the seven namespace verbs so the DFS product path and LingTai event-index +path can evolve independently. + +This handbook is for contributors touching either agent surface. It covers where the crate sits in the workspace, the invariants you must preserve, how to add a -tool, and the roadmap for the verbs we expect to add next. +namespace tool, and how to keep the LingTai event-index work inside its own +boundary. ## 1. Why this crate exists @@ -45,7 +52,8 @@ caller (bench / SDK / MCP) └─ ObjectStore (S3) ← bytes for grep/read, in nokv-object ``` -`nokv-agent` depends on exactly three workspace crates plus `serde_json`: +The namespace tool surface currently depends on exactly three workspace crates +plus `serde_json`: | Dependency | Why it is load-bearing | | --- | --- | @@ -57,6 +65,34 @@ caller (bench / SDK / MCP) It depends on **none** of `nokv-client`, `nokv-protocol`, `nokv-control`. That is the whole point — and it is enforced (see [the cycle rule](#4-invariants-do-not-break-these)). +### LingTai event index — what `nokv-agent` owns + +LingTai's runtime keeps `logs/events.jsonl` as the authoritative log and uses +`logs/log.sqlite` as a derived index for TUI/session queries. NoKV's integration +target is the derived-index layer: + +``` +LingTai logs/events.jsonl + └─ nokv-agent lingtai ingest ← streaming, offset-aware + └─ AgentEventStore ← semantic event-index API + └─ Holt-backed local agent index ← derived view, rebuildable + ├─ coverage by source file + ├─ event type / timestamp indexes + ├─ tool/action facets + └─ tool_call_id / trace correlation +``` + +The event index may open its own Holt store because it is an independent local +agent-index directory. It must not open or mutate a NoKV metadata directory and +must not become a second owner for NoKV namespace state. + +The index is allowed to expose CLI/server entry points such as +`nokv-agent lingtai ingest`, `coverage`, `latest`, `session`, `session-rows`, +`recent`, `molt-windows`, `errors`, `completion-after`, `clear-completion`, +`notification-events`, `notification-by-id`, `notification-before`, +`notification-after`, `notifications`, `facets`, and `trace`. These entry +points return stable JSON for LingTai's Go adapter and the benchmark harness. + ### Remote (RPC) — stays in `nokv-client` ``` @@ -111,6 +147,27 @@ The seven tools are **read-only**: `ls`, `stat`, `catalog`, `read`, `find`, exists (`ls`/`stat`/`catalog`), then query and read only what is needed (`find`/`aggregate`/`read`/`grep`). +The LingTai event-index API is a separate surface. Its initial semantic queries +mirror the current SQLite consumers instead of exposing raw SQL: + +- `coverage(agent, source_file)`; +- `stream_session_events(agent, filter, order = id_asc)`; +- `stream_session_rows(agent, order = id_asc)`, a compact LingTai + `SessionEventRow { ts, type, fields_json }` projection for replacing + SQLite's `StreamSessionEvents` hot path; +- `latest_events(type, limit)`; +- `recent_times(type, limit)`; +- `molt_session_windows(agent)`; +- `error_events(limit)`; +- `completion_after(type = clear_received, source_offset)`; +- `tui_clear_completion(source_offset)`; +- `notification_events(limit)`; +- `notification_event_by_id(event_id)`; +- `notification_neighbor(event_id, before | after)`; +- `notification_lifecycle(ref_id/event_id/call_id/channel)`; +- `tool_facets(window, group_by = tool/action)`; +- `tool_trace(tool_call_id)`. + ## 4. Invariants (do not break these) 1. **Transport-free.** `nokv-agent` must never depend on `nokv-client`, @@ -120,7 +177,8 @@ exists (`ls`/`stat`/`catalog`), then query and read only what is needed 2. **Read-only verb surface.** The `AgentNamespace` trait exposes only the six read verbs. Writes (e.g. `register_namespace_index`) stay as inherent methods on `NoKvFs` in `nokv-meta`, off the trait. Keep the model-facing surface - read-only unless a write contract is explicitly designed (see [the roadmap](#6-roadmap-verbs-we-expect-to-add)). + read-only unless a write contract is explicitly designed (see + [the roadmap](#7-roadmap-verbs-we-expect-to-add)). 3. **Orphan-rule placement.** `impl AgentNamespace for NoKvFs` lives in `nokv-agent` (local trait + foreign type → legal). This is what lets the embedded path work **without** `nokv-meta` gaining a dependency on @@ -137,8 +195,64 @@ exists (`ls`/`stat`/`catalog`), then query and read only what is needed `Display` strings are observed by judges, telemetry, and the model. Treat any change to them as a behavior change, not a refactor — snapshot-test before and after. - -## 5. Adding a new read tool +7. **JSONL remains authoritative for LingTai.** The event index is a derived + acceleration layer. LingTai fallback behavior must remain possible when + coverage is incomplete, a source file is truncated, or the index is absent. +8. **Idempotent event ingest.** The event index keys rows by + `(agent_id, source_file, source_offset)`. Replaying the same JSONL range must + not duplicate events or roll coverage backward. +9. **No raw-log benchmark shortcuts.** Product code may expose semantic event + queries and compact facets, but it must not bake benchmark task answers into + hidden files, tables, or special cases. + +## 5. LingTai event-index design rules + +The first implementation should use responsibility-based files instead of +growing `lib.rs`: + +| File | Contents | +| --- | --- | +| `namespace.rs` | Existing namespace trait, seven tool definitions, dispatcher, parsers, and result JSON. | +| `event/types.rs` | `AgentId`, `SourceFile`, `EventId`, `EventRecord`, projected fields, coverage, and query results. | +| `event/codec.rs` | Durable value encoding with explicit version bytes. | +| `event/key.rs` | Holt key layout for coverage, source-offset de-dupe, event rows, type/time indexes, facets, and traces. | +| `event/store.rs` | `AgentEventStore` trait and batch-ingest/query contracts. | +| `event/holt.rs` | Holt-backed implementation over an independent agent-index directory. | +| `event/ingest.rs` | JSONL streaming parser that records byte offsets and handles partial lines. | +| `src/bin/nokv-agent.rs` | LingTai CLI entry point for ingest and typed queries. | + +Required key families: + +| Family | Purpose | +| --- | --- | +| `coverage/{agent}/{source_file_hash}` | O(1) file coverage: file size, min/max offset, row count. | +| `source/{agent}/{source_file_hash}/{offset}` | Idempotent source-offset de-dupe. | +| `event/{agent}/{event_id}` | Compact event record; event id is derived from source-file hash and source offset. | +| `type_id/{agent}/{type}/{rev_id}` | Latest-by-type scans. | +| `type_ts/{agent}/{type}/{rev_ts}/{event_id}` | Recent timestamp queries. | +| `session/{agent}/{event_id}` | Ordered session-event pointer; values store event ids and replay rows are reconstructed from `event/{agent}/{event_id}`. | +| `notification_id/{agent}/{event_id}` | LingTai notification browser before/after traversal. | +| `notification_rev/{agent}/{rev_id}` | Newest-first `type LIKE '%notification%'` replacement. | +| `notification_prev/{agent}/{event_id}` | Direct older-neighbor pointer for notification browser traversal. | +| `notification_next/{agent}/{event_id}` | Direct newer-neighbor pointer for notification browser traversal. | +| `notification_tail/{agent}` | Append-order tail pointer for notification neighbor materialization. | +| `notification/{agent}/{field}/{value}/{event_id}` | Notification lifecycle lookup by `ref_id`, event id, call id, or channel. | +| `tui_clear_rev/{agent}/{rev_id}` | Newest-first TUI-originated `psyche_molt` / `clear_received` completion checks. | +| `tool/{agent}/{tool_name}/{action}/{rev_id}` | Tool/action history. | +| `trace/{agent}/{tool_call_id}/{phase}` | Tool call/result/reasoning correlation. | +| `facet/{agent}/{facet_name}/{bucket}` | Materialized top-N/facet counts for tool/action grouping. | + +Ingest commits bounded batches so a large LingTai log does not exceed Holt WAL +record limits. Each batch atomically commits event rows, secondary indexes, +materialized facet counts, and coverage. Large `fields_json` values are compacted +into descriptors because the JSONL file remains authoritative; consumers that +need the full payload must read by `source_file` and `source_offset`. + +Tests must cover replay, duplicate offsets inside one batch, chunked ingest, +partial trailing lines, large-field compaction, source-file truncation, and +coverage monotonicity. + +## 6. Adding a new read tool The flow crosses two crates because the verb logic stays in `nokv-meta`: @@ -156,7 +270,7 @@ The flow crosses two crates because the verb logic stays in `nokv-meta`: `nokv-agent` (no engine needed), and confirm the benchmark tool-registry test still matches the arm surface. -## 6. Roadmap (verbs we expect to add) +## 7. Roadmap (verbs we expect to add) The current surface answers **read** questions over a namespace. NoKV already implements the write/stateful semantics below in `nokv-meta`/`nokv-client`; the @@ -186,15 +300,19 @@ contract — argument grammar, idempotency/iflag semantics, evidence URIs, and limits — and an answer to "what does the model see on partial failure." Open a design issue before adding a write verb to the trait. -## 7. Where things live +## 8. Where things live | Concern | Location | | --- | --- | -| Tool definitions, dispatch, validation, `AgentError`, embedded impl | `crates/nokv-agent/src/lib.rs` | +| Namespace tool definitions, dispatch, validation, `AgentError`, embedded impl | `crates/nokv-agent/src/namespace.rs` | +| LingTai event-index types, codecs, keys, ingest, Holt backend | `crates/nokv-agent/src/event/` | +| LingTai agent-index CLI | `crates/nokv-agent/src/bin/nokv-agent.rs` | +| LingTai workload extractor and SQLite baseline probes | `bench/agent-interface/scripts/lingtai_workload_baseline.py` | +| LingTai SQLite vs `nokv-agent` benchmark binary | `bench/src/bin/lingtai-index-bench.rs` | | Public-surface integration test (the seven tools, `AgentError: Error`) | `crates/nokv-agent/tests/public_surface.rs` | +| LingTai event-index integration tests | `crates/nokv-agent/tests/event_index.rs` | | Verb implementations (the six read methods) | `crates/nokv-meta/src/service/agent.rs` | | Remote trait impls + re-export + `From` | `crates/nokv-client/src/agent.rs` | -| Real-world consumer (the agent-interface benchmark) | `bench/src/bin/yanex-agent-bench.rs` | See also the [code contract](code_contract.md) and the [PR review checklist](pr_review_checklist.md). From e09b6e2576ef8625e4e21381bb5a939b7491fb2a Mon Sep 17 00:00:00 2001 From: eric_song Date: Wed, 8 Jul 2026 08:06:48 +0800 Subject: [PATCH 2/2] Complete LingTai SQLite replacement semantics Signed-off-by: eric_song --- bench/agent-interface/lingtai/README.md | 21 +- crates/nokv-agent/src/bin/nokv-agent.rs | 81 ++++--- crates/nokv-agent/src/event/holt.rs | 127 ++++++++-- crates/nokv-agent/src/event/ingest.rs | 27 +-- crates/nokv-agent/src/event/key.rs | 9 + crates/nokv-agent/src/event/mod.rs | 16 +- crates/nokv-agent/src/event/notification.rs | 255 ++++++++++++++++++++ crates/nokv-agent/src/event/types.rs | 8 + crates/nokv-agent/tests/event_index.rs | 130 ++++++++-- docs/development/nokv-agent.md | 25 +- 10 files changed, 585 insertions(+), 114 deletions(-) create mode 100644 crates/nokv-agent/src/event/notification.rs diff --git a/bench/agent-interface/lingtai/README.md b/bench/agent-interface/lingtai/README.md index 4fec5360d..f31694ec3 100644 --- a/bench/agent-interface/lingtai/README.md +++ b/bench/agent-interface/lingtai/README.md @@ -6,7 +6,7 @@ active integration target. LingTai keeps `logs/events.jsonl` as the authoritative log. `logs/log.sqlite` is a derived index used by the TUI and session rebuild paths. NoKV's target is -the derived-index role: build a compact, semantic, agent-native index in +the derived-index role: build a semantic, agent-native index in `nokv-agent` and compare it against both current and optimized SQLite baselines. ## Baseline Arms @@ -99,13 +99,28 @@ Then map the current `sqlitelog` calls to typed commands: | `QueryMoltSessionWindows` | `molt-windows` | | `QueryRecentMoltTimes` | `recent --type psyche_molt` | | `QueryRecentRefreshCompleteTimes` | `recent --type refresh_complete` | -| `QueryNotificationBlocks` | `notification-blocks` | -| `QueryNotificationBlockSnapshots` | `notification-block-snapshots` | +| `QueryNotificationBlocks` | `notification-blocks`; returns parsed `blocks` plus raw `events` | +| `QueryNotificationBlockSnapshots` | `notification-block-snapshots`; returns parsed `snapshots` plus raw `events` | | `QueryNotifications` | `notification-events` or `notifications --ref-id/--event-id/--call-id/--channel` when a lifecycle filter is needed | | `QueryNotificationByID` | `notification-by-id --event-id ` | | `QueryNotificationBefore` | `notification-before --event-id ` | | `QueryNotificationAfter` | `notification-after --event-id ` | +Replacement details that must stay visible in the LingTai adapter: + +- `notification-blocks` parses `notification_pair_injected` fields into the + same summary/meta/source shape used by `QueryNotificationBlocks`. +- `notification-block-snapshots` parses both modern `_meta` envelopes and + legacy `payload`/`meta` rows used by `QueryNotificationBlockSnapshots`. +- `errors` mirrors LingTai's current `QueryErrorEvents` default types: + `aed_attempt`, `aed_exhausted`, and `refresh_init_error`; rows without a + non-empty `error` field are skipped. +- `notification-by-id` follows LingTai's current SQL and returns the row by id + without requiring the type to contain `notification`. +- `molt-windows` returns the indexed window result. LingTai's worker-local + timeout, sidecar freshness cache, and stale-on-error fallback are adapter + responsibilities when the adapter shells out to `nokv-agent`. + Keep the existing JSONL fallback when ingest, coverage, or row streaming fails, or when coverage does not reach the tail of `events.jsonl`. diff --git a/crates/nokv-agent/src/bin/nokv-agent.rs b/crates/nokv-agent/src/bin/nokv-agent.rs index 8b7997a5d..1de790795 100644 --- a/crates/nokv-agent/src/bin/nokv-agent.rs +++ b/crates/nokv-agent/src/bin/nokv-agent.rs @@ -5,29 +5,16 @@ use std::path::PathBuf; use nokv_agent::event::{ ingest_jsonl_reader, AgentEventError, AgentEventResult, AgentEventStore, - CompletionAfterRequest, ErrorEventsRequest, HoltAgentEventStore, JsonlIngestOptions, - LatestEventsRequest, NotificationEventByIdRequest, NotificationEventsRequest, - NotificationLifecycleRequest, NotificationNeighborDirection, NotificationNeighborRequest, + CompletionAfterRequest, ErrorEventsRequest, EventRecord, HoltAgentEventStore, + JsonlIngestOptions, LatestEventsRequest, NotificationBlockSnapshot, + NotificationEventByIdRequest, NotificationEventsRequest, NotificationLifecycleRequest, + NotificationNeighborDirection, NotificationNeighborRequest, NotificationSummaryEntry, RecentTimesRequest, SessionEventsRequest, SessionRowsRequest, ToolFacetRequest, ToolTraceRequest, TuiClearCompletionRequest, LINGTAI_SESSION_EVENT_TYPES, }; use serde_json::json; -const DEFAULT_ERROR_EVENT_TYPES: &[&str] = &[ - "aed_attempt", - "aed_exhausted", - "aed_over_window_detected", - "aed_timeout", - "aed_transient_retry", - "aed_transient_exhausted", - "daemon_error", - "refresh_failed_permanent", - "refresh_init_error", - "tc_wake_error", - "tool_call_dispatch_failed", - "tool_call_validation_failed", - "daemon_tool_call_validation_failed", -]; +const DEFAULT_ERROR_EVENT_TYPES: &[&str] = &["aed_attempt", "aed_exhausted", "refresh_init_error"]; fn main() { if let Err(err) = run() { @@ -58,10 +45,8 @@ fn run() -> AgentEventResult<()> { "errors" => cmd_errors(args), "completion-after" => cmd_completion_after(args), "clear-completion" => cmd_clear_completion(args), - "notification-blocks" => cmd_latest_fixed_type(args, "notification_pair_injected"), - "notification-block-snapshots" => { - cmd_latest_fixed_type(args, "notification_block_injected") - } + "notification-blocks" => cmd_notification_blocks(args), + "notification-block-snapshots" => cmd_notification_block_snapshots(args), "notification-events" => cmd_notification_events(args), "notification-by-id" => cmd_notification_by_id(args), "notification-before" => { @@ -132,26 +117,50 @@ fn cmd_latest(mut args: VecDeque) -> AgentEventResult<()> { })) } -fn cmd_latest_fixed_type( - mut args: VecDeque, +fn cmd_notification_blocks(mut args: VecDeque) -> AgentEventResult<()> { + let events = latest_fixed_type(&mut args, "notification_pair_injected")?; + let blocks = events + .iter() + .cloned() + .map(NotificationSummaryEntry::from_event) + .collect::>(); + print_json(json!({ + "blocks": blocks, + "events": events, + })) +} + +fn cmd_notification_block_snapshots(mut args: VecDeque) -> AgentEventResult<()> { + let events = latest_fixed_type(&mut args, "notification_block_injected")?; + let snapshots = events + .iter() + .cloned() + .map(NotificationBlockSnapshot::from_event) + .collect::>(); + print_json(json!({ + "snapshots": snapshots, + "events": events, + })) +} + +fn latest_fixed_type( + args: &mut VecDeque, event_type: &'static str, -) -> AgentEventResult<()> { - let store_path = required_path(&mut args, "--store")?; - let agent_id = option_value(&mut args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); - let limit = option_value(&mut args, "--limit")? +) -> AgentEventResult> { + let store_path = required_path(args, "--store")?; + let agent_id = option_value(args, "--agent-id")?.unwrap_or_else(|| "default".to_owned()); + let limit = option_value(args, "--limit")? .map(|value| parse_usize("--limit", &value)) .transpose()? .unwrap_or(10); - reject_extra(args)?; + reject_extra(std::mem::take(args))?; let store = HoltAgentEventStore::open_file(store_path)?; - print_json(json!({ - "events": store.latest_events(LatestEventsRequest { - agent_id, - event_type: event_type.to_owned(), - limit, - })? - })) + store.latest_events(LatestEventsRequest { + agent_id, + event_type: event_type.to_owned(), + limit, + }) } fn cmd_session(mut args: VecDeque) -> AgentEventResult<()> { diff --git a/crates/nokv-agent/src/event/holt.rs b/crates/nokv-agent/src/event/holt.rs index d81b49303..37884544b 100644 --- a/crates/nokv-agent/src/event/holt.rs +++ b/crates/nokv-agent/src/event/holt.rs @@ -3,6 +3,7 @@ use std::collections::{BTreeMap, BTreeSet}; use std::path::Path; use holt::{DBAtomicBatch, RangeEntry, Tree, TreeConfig, DB}; +use serde_json::Value; use super::codec::{ decode_coverage, decode_event, decode_session_event_row, decode_tool_facet, encode_coverage, @@ -10,30 +11,36 @@ use super::codec::{ }; use super::ingest::empty_coverage; use super::key::{ - coverage_key, event_key, id_from_index_value, id_value, notification_id_key, - notification_id_prefix, notification_key, notification_next_key, notification_prefix, - notification_prev_key, notification_rev_key, notification_rev_prefix, notification_tail_key, - session_key, session_prefix, source_file_hash, source_key, tool_action_facet_key, - tool_action_facet_prefix, tool_name_facet_key, trace_key, trace_prefix, tui_clear_rev_key, - tui_clear_rev_prefix, type_id_key, type_id_prefix, type_ts_key, type_ts_prefix, TREE_COVERAGE, - TREE_EVENTS, TREE_INDEX, + coverage_key, event_key, fields_json_chunk_key, id_from_index_value, id_value, + notification_id_key, notification_id_prefix, notification_key, notification_next_key, + notification_prefix, notification_prev_key, notification_rev_key, notification_rev_prefix, + notification_tail_key, session_key, session_prefix, source_file_hash, source_key, + tool_action_facet_key, tool_action_facet_prefix, tool_name_facet_key, trace_key, trace_prefix, + tui_clear_rev_key, tui_clear_rev_prefix, type_id_key, type_id_prefix, type_ts_key, + type_ts_prefix, TREE_COVERAGE, TREE_EVENTS, TREE_INDEX, TREE_PAYLOADS, }; use super::store::AgentEventStore; use super::types::{ AgentEventError, AgentEventResult, CompletionAfter, CompletionAfterRequest, ErrorEventsRequest, - EventRecord, EventTime, IndexCoverage, IngestReport, LatestEventsRequest, MoltSessionWindows, - NewEventRecord, NotificationEventByIdRequest, NotificationEventsRequest, - NotificationLifecycleRequest, NotificationNeighborDirection, NotificationNeighborRequest, - RecentTimesRequest, SessionEventRow, SessionEventsRequest, SessionRowsRequest, ToolFacet, - ToolFacetRequest, ToolTraceRequest, TuiClearCompletion, TuiClearCompletionRequest, - LINGTAI_SESSION_EVENT_TYPES, + EventRecord, EventTime, ExternalFieldsJsonRef, IndexCoverage, IngestReport, + LatestEventsRequest, MoltSessionWindows, NewEventRecord, NotificationEventByIdRequest, + NotificationEventsRequest, NotificationLifecycleRequest, NotificationNeighborDirection, + NotificationNeighborRequest, RecentTimesRequest, SessionEventRow, SessionEventsRequest, + SessionRowsRequest, ToolFacet, ToolFacetRequest, ToolTraceRequest, TuiClearCompletion, + TuiClearCompletionRequest, LINGTAI_SESSION_EVENT_TYPES, }; +const MAX_HOLT_VALUE_BYTES: usize = 65_000; +const FIELDS_JSON_CHUNK_BYTES: usize = 48 * 1024; + +type PayloadChunkWrites = Vec<(Vec, Vec)>; + pub struct HoltAgentEventStore { db: DB, events: Tree, index: Tree, coverage: Tree, + payloads: Tree, } impl HoltAgentEventStore { @@ -54,11 +61,15 @@ impl HoltAgentEventStore { let coverage = db .open_or_create_tree(TREE_COVERAGE) .map_err(to_store_error)?; + let payloads = db + .open_or_create_tree(TREE_PAYLOADS) + .map_err(to_store_error)?; Ok(Self { db, events, index, coverage, + payloads, }) } @@ -67,7 +78,7 @@ impl HoltAgentEventStore { self.events .get(&key) .map_err(to_store_error)? - .map(|bytes| decode_event(&bytes)) + .map(|bytes| decode_event(&bytes).and_then(|record| self.resolve_fields_json(record))) .transpose()? .map(Some) // Large file-backed stores can expose a key through range scan @@ -85,13 +96,42 @@ impl HoltAgentEventStore { continue; }; if found == key { - return decode_event(&value).map(Some); + return decode_event(&value) + .and_then(|record| self.resolve_fields_json(record)) + .map(Some); } break; } Ok(None) } + fn resolve_fields_json(&self, mut record: EventRecord) -> AgentEventResult { + let Some(external) = record.external_fields_json.take() else { + return Ok(record); + }; + let mut encoded = Vec::with_capacity(external.encoded_bytes); + for chunk_index in 0..external.chunk_count { + let key = fields_json_chunk_key(&record.agent_id, record.id, chunk_index); + let Some(chunk) = self.payloads.get(&key).map_err(to_store_error)? else { + return Err(AgentEventError::Store(format!( + "missing fields_json chunk {chunk_index} for event {}", + record.id + ))); + }; + encoded.extend_from_slice(&chunk); + } + if encoded.len() != external.encoded_bytes { + return Err(AgentEventError::Store(format!( + "fields_json payload length mismatch for event {}: expected {}, got {}", + record.id, + external.encoded_bytes, + encoded.len() + ))); + } + record.fields_json = serde_json::from_slice(&encoded)?; + Ok(record) + } + fn scan_event_ids(&self, prefix: &[u8], limit: usize) -> AgentEventResult> { let mut out = Vec::new(); for entry in self.index.scan(prefix) { @@ -368,6 +408,7 @@ impl AgentEventStore for HoltAgentEventStore { source_line: new.source_line, ts: new.ts, event_type: new.event_type, + external_fields_json: None, fields_json: new.fields_json, projection: new.projection, }; @@ -538,7 +579,11 @@ impl AgentEventStore for HoltAgentEventStore { let mut out = Vec::new(); for event_type in request.event_types { let prefix = type_id_prefix(&request.agent_id, &event_type); - out.extend(self.events_by_index_prefix(&request.agent_id, &prefix, request.limit)?); + out.extend( + self.events_by_index_prefix(&request.agent_id, &prefix, request.limit)? + .into_iter() + .filter(has_error_text), + ); } out.sort_by_key(|record| Reverse(record.id)); out.truncate(request.limit); @@ -613,9 +658,7 @@ impl AgentEventStore for HoltAgentEventStore { &self, request: NotificationEventByIdRequest, ) -> AgentEventResult> { - Ok(self - .event_by_id(&request.agent_id, request.event_id)? - .filter(|record| is_notification_event_type(&record.event_type))) + self.event_by_id(&request.agent_id, request.event_id) } fn notification_neighbor( @@ -678,11 +721,15 @@ impl AgentEventStore for HoltAgentEventStore { fn write_record_batch(batch: &mut DBAtomicBatch, source: &[u8], record: &EventRecord) { let id = id_value(record.id); + let (encoded_event, fields_json_chunks) = encode_event_for_holt(record); batch.put_if_absent(TREE_INDEX, source, &id); + for (key, chunk) in fields_json_chunks { + batch.put(TREE_PAYLOADS, &key, &chunk); + } batch.put( TREE_EVENTS, &event_key(&record.agent_id, record.id), - &encode_event(record).expect("event encodes"), + &encoded_event, ); batch.put( TREE_INDEX, @@ -754,6 +801,38 @@ fn write_record_batch(batch: &mut DBAtomicBatch, source: &[u8], record: &EventRe } } +fn encode_event_for_holt(record: &EventRecord) -> (Vec, PayloadChunkWrites) { + let inline = encode_event(record).expect("event encodes"); + if inline.len() <= MAX_HOLT_VALUE_BYTES { + return (inline, Vec::new()); + } + + let encoded_fields = serde_json::to_vec(&record.fields_json).expect("fields_json encodes"); + let chunk_count = encoded_fields.chunks(FIELDS_JSON_CHUNK_BYTES).count(); + let mut stored = record.clone(); + stored.fields_json = Value::Null; + stored.external_fields_json = Some(ExternalFieldsJsonRef { + encoded_bytes: encoded_fields.len(), + chunk_count, + }); + let encoded_event = encode_event(&stored).expect("externalized event encodes"); + assert!( + encoded_event.len() <= MAX_HOLT_VALUE_BYTES, + "externalized event still exceeds Holt value limit" + ); + let chunks = encoded_fields + .chunks(FIELDS_JSON_CHUNK_BYTES) + .enumerate() + .map(|(chunk_index, chunk)| { + ( + fields_json_chunk_key(&record.agent_id, record.id, chunk_index), + chunk.to_vec(), + ) + }) + .collect(); + (encoded_event, chunks) +} + fn increment_facet( increments: &mut BTreeMap, ToolFacet>, key: Vec, @@ -797,6 +876,14 @@ fn is_notification_event_type(event_type: &str) -> bool { event_type.contains("notification") } +fn has_error_text(record: &EventRecord) -> bool { + record + .projection + .error + .as_deref() + .is_some_and(|error| !error.trim().is_empty()) +} + fn is_tui_clear_completion_event(record: &EventRecord) -> bool { matches!(record.event_type.as_str(), "psyche_molt" | "clear_received") && record diff --git a/crates/nokv-agent/src/event/ingest.rs b/crates/nokv-agent/src/event/ingest.rs index b8116c818..0c25dd90e 100644 --- a/crates/nokv-agent/src/event/ingest.rs +++ b/crates/nokv-agent/src/event/ingest.rs @@ -2,7 +2,6 @@ use std::io::BufRead; use serde_json::{Map, Value}; -const MAX_INLINE_FIELDS_JSON_BYTES: usize = 32 * 1024; const INGEST_RECORD_BATCH_LIMIT: usize = 512; use super::store::AgentEventStore; @@ -116,31 +115,7 @@ fn fields_json(object: &Map) -> Value { for key in ["type", "ts", "address", "agent_name"] { fields.remove(key); } - compact_fields_json(Value::Object(fields)) -} - -fn compact_fields_json(value: Value) -> Value { - let encoded_len = serde_json::to_vec(&value) - .map(|bytes| bytes.len()) - .unwrap_or(usize::MAX); - if encoded_len <= MAX_INLINE_FIELDS_JSON_BYTES { - return value; - } - let keys = value - .as_object() - .map(|object| { - object - .keys() - .cloned() - .map(Value::String) - .collect::>() - }) - .unwrap_or_default(); - serde_json::json!({ - "_nokv_compacted": true, - "original_json_bytes": encoded_len, - "keys": keys, - }) + Value::Object(fields) } fn project_event(object: &Map) -> EventProjection { diff --git a/crates/nokv-agent/src/event/key.rs b/crates/nokv-agent/src/event/key.rs index aa74c4382..9d2f691e8 100644 --- a/crates/nokv-agent/src/event/key.rs +++ b/crates/nokv-agent/src/event/key.rs @@ -3,6 +3,7 @@ use sha2::{Digest, Sha256}; pub const TREE_EVENTS: &str = "events"; pub const TREE_INDEX: &str = "index"; pub const TREE_COVERAGE: &str = "coverage"; +pub const TREE_PAYLOADS: &str = "payloads"; pub fn source_file_hash(source_file: &str) -> String { let digest = Sha256::digest(source_file.as_bytes()); @@ -35,6 +36,14 @@ pub fn event_prefix(agent_id: &str) -> Vec { format!("event/{}/", escape(agent_id)).into_bytes() } +pub fn fields_json_chunk_key(agent_id: &str, event_id: u64, chunk_index: usize) -> Vec { + format!( + "fields_json/{}/{event_id:020}/{chunk_index:06}", + escape(agent_id) + ) + .into_bytes() +} + pub fn type_id_key(agent_id: &str, event_type: &str, event_id: u64) -> Vec { format!( "type_id/{}/{}/{:020}", diff --git a/crates/nokv-agent/src/event/mod.rs b/crates/nokv-agent/src/event/mod.rs index e9caa76fd..0fd14214f 100644 --- a/crates/nokv-agent/src/event/mod.rs +++ b/crates/nokv-agent/src/event/mod.rs @@ -8,18 +8,22 @@ pub mod codec; pub mod holt; pub mod ingest; pub mod key; +pub mod notification; pub mod store; pub mod types; pub use holt::HoltAgentEventStore; pub use ingest::{ingest_jsonl_reader, JsonlIngestOptions}; +pub use notification::{ + NotificationBlockMeta, NotificationBlockSnapshot, NotificationSummaryEntry, +}; pub use store::AgentEventStore; pub use types::{ AgentEventError, AgentEventResult, CompletionAfter, CompletionAfterRequest, ErrorEventsRequest, - EventProjection, EventRecord, EventTime, IndexCoverage, IngestReport, LatestEventsRequest, - MoltSessionWindows, NewEventRecord, NotificationEventByIdRequest, NotificationEventsRequest, - NotificationLifecycleRequest, NotificationNeighborDirection, NotificationNeighborRequest, - RecentTimesRequest, SessionEventRow, SessionEventsRequest, SessionRowsRequest, ToolFacet, - ToolFacetRequest, ToolTraceRequest, TuiClearCompletion, TuiClearCompletionRequest, - LINGTAI_SESSION_EVENT_TYPES, + EventProjection, EventRecord, EventTime, ExternalFieldsJsonRef, IndexCoverage, IngestReport, + LatestEventsRequest, MoltSessionWindows, NewEventRecord, NotificationEventByIdRequest, + NotificationEventsRequest, NotificationLifecycleRequest, NotificationNeighborDirection, + NotificationNeighborRequest, RecentTimesRequest, SessionEventRow, SessionEventsRequest, + SessionRowsRequest, ToolFacet, ToolFacetRequest, ToolTraceRequest, TuiClearCompletion, + TuiClearCompletionRequest, LINGTAI_SESSION_EVENT_TYPES, }; diff --git a/crates/nokv-agent/src/event/notification.rs b/crates/nokv-agent/src/event/notification.rs new file mode 100644 index 000000000..b40748094 --- /dev/null +++ b/crates/nokv-agent/src/event/notification.rs @@ -0,0 +1,255 @@ +use std::collections::BTreeMap; +use std::path::Path; + +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +use super::types::EventRecord; + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct NotificationBlockMeta { + #[serde(skip_serializing_if = "Option::is_none")] + pub current_time: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub injection_seq: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub context_system_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub context_history_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub context_usage: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct NotificationSummaryEntry { + pub id: u64, + pub ts: f64, + pub source: String, + pub call_id: Option, + pub summary: Option, + pub sources: Vec, + pub meta: Option, + pub raw_event: EventRecord, +} + +impl NotificationSummaryEntry { + pub fn from_event(record: EventRecord) -> Self { + let fields = record.fields_json.as_object(); + let meta = fields + .and_then(|fields| fields.get("meta")) + .and_then(Value::as_object) + .map(parse_block_meta); + Self { + id: record.id, + ts: record.ts, + source: source_base(&record.source_file), + call_id: string_field(fields, "call_id"), + summary: string_field(fields, "summary"), + sources: string_array_field(fields, "sources"), + meta, + raw_event: record, + } + } +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct NotificationBlockSnapshot { + pub id: u64, + pub ts: f64, + pub source: String, + pub mode: Option, + pub call_id: Option, + pub sources: Vec, + pub meta: Option, + pub raw_meta: Option, + pub tool_meta: Option, + pub agent_meta: Option, + pub guidance: Option, + pub notification_guidance: Option, + pub notifications: BTreeMap, + pub raw_event: EventRecord, +} + +impl NotificationBlockSnapshot { + pub fn from_event(record: EventRecord) -> Self { + let fields_json = record.fields_json.clone(); + let fields = fields_json.as_object(); + let mut parsed = Self { + id: record.id, + ts: record.ts, + source: source_base(&record.source_file), + mode: string_field(fields, "mode"), + call_id: string_field(fields, "call_id"), + sources: string_array_field(fields, "sources"), + meta: None, + raw_meta: None, + tool_meta: None, + agent_meta: None, + guidance: None, + notification_guidance: None, + notifications: BTreeMap::new(), + raw_event: record, + }; + + if let Some(meta_env) = fields + .and_then(|fields| fields.get("_meta")) + .and_then(Value::as_object) + { + parse_meta_envelope(meta_env, &mut parsed); + } else if let Some(payload) = fields + .and_then(|fields| fields.get("payload")) + .and_then(Value::as_object) + { + let legacy_meta = fields + .and_then(|fields| fields.get("meta")) + .and_then(Value::as_object); + parse_legacy_payload(payload, legacy_meta, &mut parsed); + } + + if parsed.raw_meta.is_none() { + if let Some(meta) = fields + .and_then(|fields| fields.get("meta")) + .and_then(Value::as_object) + { + parsed.raw_meta = Some(Value::Object(meta.clone())); + } + } + if let Some(raw_meta) = parsed.raw_meta.as_ref().and_then(Value::as_object) { + if parsed.agent_meta.is_none() { + parsed.agent_meta = Some(Value::Object(raw_meta.clone())); + } + parsed.meta = Some(parse_block_meta(raw_meta)); + } + parsed + } +} + +fn parse_meta_envelope(env: &Map, snapshot: &mut NotificationBlockSnapshot) { + if let Some(value) = object_value(env, "tool_meta") { + snapshot.tool_meta = Some(value); + } + if let Some(agent) = object_value(env, "agent_meta") { + snapshot.raw_meta = Some(agent.clone()); + snapshot.agent_meta = Some(agent); + } + if let Some(value) = object_value(env, "guidance") { + snapshot.guidance = Some(value); + } + if let Some(value) = env.get("notification_guidance").and_then(Value::as_str) { + snapshot.notification_guidance = Some(value.to_owned()); + } + if let Some(notifications) = env.get("notifications").and_then(Value::as_object) { + snapshot.notifications = encode_notification_channels(notifications); + } +} + +fn parse_legacy_payload( + payload: &Map, + meta: Option<&Map>, + snapshot: &mut NotificationBlockSnapshot, +) { + if let Some(value) = object_value(payload, "_tool") { + snapshot.tool_meta = Some(value); + } + if let Some(value) = object_value(payload, "_runtime.state") { + snapshot.agent_meta = Some(value); + } + if let Some(value) = object_value(payload, "_runtime.guidance") { + snapshot.guidance = Some(value); + } + if let Some(runtime) = payload.get("_runtime").and_then(Value::as_object) { + if let Some(value) = object_value(runtime, "state") { + snapshot.agent_meta = Some(value); + } + if let Some(value) = object_value(runtime, "guidance") { + snapshot.guidance = Some(value); + } + } + snapshot.notification_guidance = payload + .get("notification_guidance") + .or_else(|| payload.get("_notification_guidance")) + .and_then(Value::as_str) + .map(ToOwned::to_owned); + if let Some(notifications) = payload.get("notifications").and_then(Value::as_object) { + snapshot.notifications = encode_notification_channels(notifications); + } + if let Some(meta) = meta { + snapshot.raw_meta = Some(Value::Object(meta.clone())); + } +} + +fn parse_block_meta(meta: &Map) -> NotificationBlockMeta { + let context = meta.get("context").and_then(Value::as_object); + NotificationBlockMeta { + current_time: meta + .get("current_time") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + injection_seq: meta.get("injection_seq").and_then(number_to_i64), + context_system_tokens: context + .and_then(|context| context.get("system_tokens")) + .and_then(number_to_i64), + context_history_tokens: context + .and_then(|context| context.get("history_tokens")) + .and_then(number_to_i64), + context_usage: context + .and_then(|context| context.get("usage")) + .and_then(Value::as_f64), + } +} + +fn encode_notification_channels(notifications: &Map) -> BTreeMap { + notifications + .iter() + .map(|(channel, value)| { + let encoded = serde_json::to_string_pretty(value).unwrap_or_else(|_| value.to_string()); + (channel.clone(), encoded) + }) + .collect() +} + +fn object_value(object: &Map, field: &str) -> Option { + object + .get(field) + .and_then(Value::as_object) + .map(|value| Value::Object(value.clone())) +} + +fn string_field(fields: Option<&Map>, field: &str) -> Option { + fields + .and_then(|fields| fields.get(field)) + .and_then(Value::as_str) + .map(ToOwned::to_owned) +} + +fn string_array_field(fields: Option<&Map>, field: &str) -> Vec { + fields + .and_then(|fields| fields.get(field)) + .and_then(Value::as_array) + .map(|values| { + values + .iter() + .filter_map(Value::as_str) + .map(ToOwned::to_owned) + .collect() + }) + .unwrap_or_default() +} + +fn number_to_i64(value: &Value) -> Option { + value + .as_i64() + .or_else(|| value.as_u64().and_then(|value| i64::try_from(value).ok())) + .or_else(|| value.as_f64().map(|value| value as i64)) +} + +fn source_base(path: &str) -> String { + if path.is_empty() { + return String::new(); + } + Path::new(path) + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or(path) + .to_owned() +} diff --git a/crates/nokv-agent/src/event/types.rs b/crates/nokv-agent/src/event/types.rs index 94c423e68..9f5556ad2 100644 --- a/crates/nokv-agent/src/event/types.rs +++ b/crates/nokv-agent/src/event/types.rs @@ -96,10 +96,18 @@ pub struct EventRecord { pub source_line: u64, pub ts: f64, pub event_type: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub external_fields_json: Option, pub fields_json: Value, pub projection: EventProjection, } +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ExternalFieldsJsonRef { + pub encoded_bytes: usize, + pub chunk_count: usize, +} + #[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct IndexCoverage { pub agent_id: String, diff --git a/crates/nokv-agent/tests/event_index.rs b/crates/nokv-agent/tests/event_index.rs index cdb75fe08..7c17b88fe 100644 --- a/crates/nokv-agent/tests/event_index.rs +++ b/crates/nokv-agent/tests/event_index.rs @@ -3,10 +3,10 @@ use std::process::Command; use nokv_agent::event::{ ingest_jsonl_reader, AgentEventStore, CompletionAfterRequest, ErrorEventsRequest, - HoltAgentEventStore, JsonlIngestOptions, LatestEventsRequest, NotificationLifecycleRequest, - NotificationNeighborDirection, NotificationNeighborRequest, RecentTimesRequest, - SessionEventsRequest, SessionRowsRequest, ToolFacetRequest, ToolTraceRequest, - TuiClearCompletionRequest, + HoltAgentEventStore, JsonlIngestOptions, LatestEventsRequest, NotificationBlockSnapshot, + NotificationLifecycleRequest, NotificationNeighborDirection, NotificationNeighborRequest, + NotificationSummaryEntry, RecentTimesRequest, SessionEventsRequest, SessionRowsRequest, + ToolFacetRequest, ToolTraceRequest, TuiClearCompletionRequest, }; fn ingest(store: &HoltAgentEventStore, jsonl: &str) -> nokv_agent::event::IngestReport { @@ -133,6 +133,8 @@ fn event_queries_match_lingtai_sqlite_shapes() { "\n", r#"{"type":"aed_attempt","ts":9.0,"error":"over window"}"#, "\n", + r#"{"type":"aed_exhausted","ts":9.5}"#, + "\n", r#"{"type":"clear_received","ts":10.0,"source":"tui"}"#, "\n", ); @@ -331,6 +333,14 @@ fn event_queries_match_lingtai_sqlite_shapes() { .unwrap() .unwrap(); assert_eq!(by_id.event_type, "system_notification_published"); + let non_notification_by_id = store + .notification_event_by_id(nokv_agent::event::NotificationEventByIdRequest { + agent_id: "agent-a".to_owned(), + event_id: latest[0].id, + }) + .unwrap() + .unwrap(); + assert_eq!(non_notification_by_id.event_type, "tool_call"); let before = store .notification_neighbor(NotificationNeighborRequest { agent_id: "agent-a".to_owned(), @@ -378,8 +388,8 @@ fn partial_trailing_line_is_not_indexed() { } #[test] -fn large_fields_json_is_compacted_for_holt_value_limit() { - let store = HoltAgentEventStore::open_memory().unwrap(); +fn large_fields_json_is_preserved_for_sqlite_replacement() { + let dir = tempfile::tempdir().unwrap(); let large_text = "x".repeat(70 * 1024); let jsonl = format!( "{}\n", @@ -392,20 +402,95 @@ fn large_fields_json_is_compacted_for_holt_value_limit() { }) ); - let report = ingest(&store, &jsonl); - assert_eq!(report.accepted, 1); + { + let store = HoltAgentEventStore::open_file(dir.path()).unwrap(); + let report = ingest(&store, &jsonl); + assert_eq!(report.accepted, 1); + } + { + let store = HoltAgentEventStore::open_file(dir.path()).unwrap(); + let events = store + .latest_events(LatestEventsRequest { + agent_id: "agent-a".to_owned(), + event_type: "tool_result".to_owned(), + limit: 1, + }) + .unwrap(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].projection.tool_name.as_deref(), Some("bash")); + assert_eq!(events[0].fields_json["result"], large_text); + assert!(events[0].external_fields_json.is_none()); + } +} + +#[test] +fn parsed_notification_blocks_match_lingtai_struct_fields() { + let store = HoltAgentEventStore::open_memory().unwrap(); + let jsonl = concat!( + r#"{"type":"notification_pair_injected","ts":2.0,"call_id":"notify-call-1","summary":"note","sources":["email","system"],"meta":{"current_time":"2026-07-08T00:00:00Z","injection_seq":7,"context":{"system_tokens":11,"history_tokens":22,"usage":0.5}}}"#, + "\n", + r#"{"type":"notification_block_injected","ts":3.0,"mode":"active_tool_result","call_id":"notify-call-1","sources":["email"],"_meta":{"tool_meta":{"name":"bash"},"agent_meta":{"current_time":"2026-07-08T00:01:00Z","injection_seq":8,"context":{"system_tokens":33,"history_tokens":44,"usage":0.75}},"guidance":{"style":"brief"},"notification_guidance":"check mailbox","notifications":{"email":{"subject":"hello"}}}}"#, + "\n", + r#"{"type":"notification_block_injected","ts":4.0,"mode":"synthetic_notification_pair","call_id":"notify-call-2","payload":{"_tool":{"name":"legacy"},"_runtime":{"state":{"current_time":"legacy-now"},"guidance":{"style":"legacy"}},"_notification_guidance":"legacy guidance","notifications":{"system":{"text":"legacy"}}},"meta":{"current_time":"legacy-meta","context":{"system_tokens":1,"history_tokens":2,"usage":0.25}}}"#, + "\n", + ); + ingest(&store, jsonl); - let events = store + let pair = store .latest_events(LatestEventsRequest { agent_id: "agent-a".to_owned(), - event_type: "tool_result".to_owned(), + event_type: "notification_pair_injected".to_owned(), limit: 1, }) + .unwrap() + .pop() .unwrap(); - assert_eq!(events.len(), 1); - assert_eq!(events[0].projection.tool_name.as_deref(), Some("bash")); - assert_eq!(events[0].fields_json["_nokv_compacted"], true); - assert_eq!(events[0].fields_json["keys"][0], "result"); + let block = NotificationSummaryEntry::from_event(pair); + assert_eq!(block.source, "events.jsonl"); + assert_eq!(block.call_id.as_deref(), Some("notify-call-1")); + assert_eq!(block.summary.as_deref(), Some("note")); + assert_eq!(block.sources, vec!["email", "system"]); + assert_eq!(block.meta.unwrap().context_history_tokens, Some(22)); + + let snapshots = store + .latest_events(LatestEventsRequest { + agent_id: "agent-a".to_owned(), + event_type: "notification_block_injected".to_owned(), + limit: 2, + }) + .unwrap() + .into_iter() + .map(NotificationBlockSnapshot::from_event) + .collect::>(); + assert_eq!( + snapshots[0].notification_guidance.as_deref(), + Some("legacy guidance") + ); + assert_eq!( + snapshots[0].agent_meta.as_ref().unwrap()["current_time"], + "legacy-now" + ); + assert_eq!( + snapshots[0].meta.as_ref().unwrap().current_time.as_deref(), + Some("legacy-meta") + ); + assert_eq!( + snapshots[0].notifications["system"], + "{\n \"text\": \"legacy\"\n}" + ); + assert_eq!(snapshots[1].tool_meta.as_ref().unwrap()["name"], "bash"); + assert_eq!( + snapshots[1].notification_guidance.as_deref(), + Some("check mailbox") + ); + assert_eq!( + snapshots[1].meta.as_ref().unwrap().context_usage, + Some(0.75) + ); + assert_eq!( + snapshots[1].notifications["email"], + "{\n \"subject\": \"hello\"\n}" + ); } #[test] @@ -449,9 +534,9 @@ fn nokv_agent_lingtai_cli_round_trips_json() { concat!( r#"{"type":"tool_call","ts":1.0,"tool_name":"read","tool_call_id":"call-1","tool_args":{"file_path":"/tmp/a.md"}}"#, "\n", - r#"{"type":"notification_pair_injected","ts":2.0,"call_id":"notify-call-1","summary":"note"}"#, + r#"{"type":"notification_pair_injected","ts":2.0,"call_id":"notify-call-1","summary":"note","sources":["email"],"meta":{"current_time":"now","context":{"system_tokens":1,"history_tokens":2,"usage":0.5}}}"#, "\n", - r#"{"type":"notification_block_injected","ts":3.0,"call_id":"notify-call-1","summary":"block"}"#, + r#"{"type":"notification_block_injected","ts":3.0,"mode":"active_tool_result","call_id":"notify-call-1","sources":["email"],"_meta":{"tool_meta":{"name":"read"},"agent_meta":{"current_time":"now","context":{"system_tokens":3,"history_tokens":4,"usage":0.75}},"guidance":{"mode":"brief"},"notification_guidance":"notify","notifications":{"email":{"subject":"hello"}}}}"#, "\n", r#"{"type":"clear_received","ts":4.0,"source":"tui"}"#, "\n", @@ -600,6 +685,11 @@ fn nokv_agent_lingtai_cli_round_trips_json() { notification_blocks_json["events"][0]["event_type"], "notification_pair_injected" ); + assert_eq!(notification_blocks_json["blocks"][0]["summary"], "note"); + assert_eq!( + notification_blocks_json["blocks"][0]["meta"]["context_history_tokens"], + 2 + ); let notification_block_snapshots = Command::new(bin) .args([ @@ -625,6 +715,14 @@ fn nokv_agent_lingtai_cli_round_trips_json() { notification_block_snapshots_json["events"][0]["event_type"], "notification_block_injected" ); + assert_eq!( + notification_block_snapshots_json["snapshots"][0]["tool_meta"]["name"], + "read" + ); + assert_eq!( + notification_block_snapshots_json["snapshots"][0]["notification_guidance"], + "notify" + ); let notification_events = Command::new(bin) .args([ diff --git a/docs/development/nokv-agent.md b/docs/development/nokv-agent.md index 204b51948..a89846013 100644 --- a/docs/development/nokv-agent.md +++ b/docs/development/nokv-agent.md @@ -89,9 +89,10 @@ must not become a second owner for NoKV namespace state. The index is allowed to expose CLI/server entry points such as `nokv-agent lingtai ingest`, `coverage`, `latest`, `session`, `session-rows`, `recent`, `molt-windows`, `errors`, `completion-after`, `clear-completion`, -`notification-events`, `notification-by-id`, `notification-before`, -`notification-after`, `notifications`, `facets`, and `trace`. These entry -points return stable JSON for LingTai's Go adapter and the benchmark harness. +`notification-blocks`, `notification-block-snapshots`, `notification-events`, +`notification-by-id`, `notification-before`, `notification-after`, +`notifications`, `facets`, and `trace`. These entry points return stable JSON +for LingTai's Go adapter and the benchmark harness. ### Remote (RPC) — stays in `nokv-client` @@ -161,6 +162,9 @@ mirror the current SQLite consumers instead of exposing raw SQL: - `error_events(limit)`; - `completion_after(type = clear_received, source_offset)`; - `tui_clear_completion(source_offset)`; +- parsed `notification-blocks` rows for replacing `QueryNotificationBlocks`; +- parsed `notification-block-snapshots` rows for replacing + `QueryNotificationBlockSnapshots`; - `notification_events(limit)`; - `notification_event_by_id(event_id)`; - `notification_neighbor(event_id, before | after)`; @@ -216,6 +220,7 @@ growing `lib.rs`: | `event/types.rs` | `AgentId`, `SourceFile`, `EventId`, `EventRecord`, projected fields, coverage, and query results. | | `event/codec.rs` | Durable value encoding with explicit version bytes. | | `event/key.rs` | Holt key layout for coverage, source-offset de-dupe, event rows, type/time indexes, facets, and traces. | +| `event/notification.rs` | LingTai notification summary/snapshot parsing for modern `_meta` envelopes and legacy `payload`/`meta` rows. | | `event/store.rs` | `AgentEventStore` trait and batch-ingest/query contracts. | | `event/holt.rs` | Holt-backed implementation over an independent agent-index directory. | | `event/ingest.rs` | JSONL streaming parser that records byte offsets and handles partial lines. | @@ -228,6 +233,7 @@ Required key families: | `coverage/{agent}/{source_file_hash}` | O(1) file coverage: file size, min/max offset, row count. | | `source/{agent}/{source_file_hash}/{offset}` | Idempotent source-offset de-dupe. | | `event/{agent}/{event_id}` | Compact event record; event id is derived from source-file hash and source offset. | +| `payloads:fields_json/{agent}/{event_id}/{chunk}` | External `fields_json` chunks for events whose full payload exceeds Holt's single-value limit. Reads transparently rehydrate this into `EventRecord.fields_json`. | | `type_id/{agent}/{type}/{rev_id}` | Latest-by-type scans. | | `type_ts/{agent}/{type}/{rev_ts}/{event_id}` | Recent timestamp queries. | | `session/{agent}/{event_id}` | Ordered session-event pointer; values store event ids and replay rows are reconstructed from `event/{agent}/{event_id}`. | @@ -244,12 +250,17 @@ Required key families: Ingest commits bounded batches so a large LingTai log does not exceed Holt WAL record limits. Each batch atomically commits event rows, secondary indexes, -materialized facet counts, and coverage. Large `fields_json` values are compacted -into descriptors because the JSONL file remains authoritative; consumers that -need the full payload must read by `source_file` and `source_offset`. +materialized facet counts, and coverage. `fields_json` is preserved in the +event record so LingTai's row projections, pretty-fields rendering, and parsed +notification views do not lose payload content relative to SQLite. + +When the LingTai adapter shells out to `nokv-agent`, process-level timeout, +freshness cache, and stale-on-error fallback remain adapter responsibilities. +The Rust event index returns deterministic derived-index state; it does not own +the TUI worker scheduling policy. Tests must cover replay, duplicate offsets inside one batch, chunked ingest, -partial trailing lines, large-field compaction, source-file truncation, and +partial trailing lines, large-field preservation, source-file truncation, and coverage monotonicity. ## 6. Adding a new read tool