diff --git a/CLAUDE.md b/CLAUDE.md index 9a09edf..d2596ca 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -12,7 +12,8 @@ Keep the repository focused on these public event-service responsibilities. - `identity` validates the deployment bearer token and returns the stable standalone owner identifier. - `usage` owns event contracts plus backend-independent validation and ordering. -- `search` owns shared search contracts and safe projection extraction. +- `search` owns shared search contracts and the literal text matching both + storage backends share. - `storage` is the `dyn StorageBackend` compatibility boundary and contains the conditionally compiled PostgreSQL/Elasticsearch and SQLite/FTS5 backends. - `db` owns PostgreSQL-only pooling, migrations, and Diesel models. diff --git a/Cargo.lock b/Cargo.lock index 1753509..e75658f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10,6 +10,7 @@ dependencies = [ "axum-extra", "base64 0.21.7", "chrono", + "chrono-tz", "diesel", "diesel_migrations", "hex", @@ -226,6 +227,16 @@ dependencies = [ "windows-link", ] +[[package]] +name = "chrono-tz" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6139a8597ed92cf816dfb33f5dd6cf0bb93a6adc938f11039f371bc5bcd26c3" +dependencies = [ + "chrono", + "phf", +] + [[package]] name = "core-foundation-sys" version = "0.8.7" @@ -982,6 +993,24 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "phf" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "913273894cec178f401a31ec4b656318d95473527be05c0752cc41cdc32be8b7" +dependencies = [ + "phf_shared", +] + +[[package]] +name = "phf_shared" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06005508882fb681fd97892ecff4b7fd0fee13ef1aa569f8695dae7ab9099981" +dependencies = [ + "siphasher", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -1427,6 +1456,12 @@ dependencies = [ "libc", ] +[[package]] +name = "siphasher" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33f4fe9184a62d842c9ef383018f3306d8ba224fd9d836f56d7288308847c256" + [[package]] name = "slab" version = "0.4.12" diff --git a/README.md b/README.md index 506716c..dfd00e2 100644 --- a/README.md +++ b/README.md @@ -2,37 +2,45 @@ `abyss-backend` is the open-source, self-hostable Agent event store for Abyss. It accepts normalized Agent events and exposes APIs for event history, -summaries, session timelines, attachments, and full-text session search. +summaries, session streams, attachments, and session search. This repository is intentionally focused on standalone event storage, queries, -and optional full-text search. +and search. ## Capabilities - Idempotent Agent event ingestion and raw event queries. -- Token-usage summaries and ordered session timelines. +- Token-usage summaries and derived session streams. - Validated image attachment storage and authorized download. -- Full-text session search through Elasticsearch or embedded SQLite FTS5. +- Session search through Elasticsearch or an embedded SQLite trigram index, + and, with SQLite, search within one session. - Health and storage-readiness probes. - Embedded migrations for both supported storage profiles. -A session is titled by the `session_title` the collector reports for a call in -which the Agent names its own session (Claude Code); otherwise by its first -nonempty user input, collapsed to one line and truncated to 80 characters with -an ellipsis. An Agent title is never replaced by a default title. Requests whose -`purpose` marks them as naming calls are not user input. Codex names tasks from -separate threads that do not say which session they name; those threads are -marked `title_generation` from their `thread_source` and are not linked to a -task. Naming calls retain their own events and usage. Existing sessions receive -default titles through a one-time backfill. +Every session view uses the same title: the label of the session's earliest +request, ordered by observation time and event primary key, even when that +request has no label. A request's label is the last `user` entry with text in +the request's new part, as `{"event_pk", "entry_index", "text", "length"}`: +`text` is the entry's text cut to its first 50 UTF-16 code units, unformatted, +and `length` is the whole text's length in UTF-16 code units. A session whose +earliest request has no label has a `null` title, and later delivery of an +earlier request corrects it. A segment label is the label of the segment's first +request; it does not name the session. + +`20261014150000_initial_schema` creates the owner, device, session, event, +attachment, diagnostic, and session title tables, and in PostgreSQL the search +outbox. `20261030130000_event_conversation_and_stream_facts` adds the event +`conversation` column and the stream facts table, and in SQLite the search text +tables. Neither migration moves or rebuilds stored data: databases created by +earlier releases are not upgraded and must be recreated. Exactly one storage profile is compiled into a binary: - `postgres-es` uses PostgreSQL as the source of truth and optionally projects search documents to Elasticsearch. It is the default profile. -- `sqlite-fts` stores both authoritative data and its transactional FTS5 index - in one local SQLite file through Diesel, with FTS5 operations expressed as - SQL. It does not compile PostgreSQL, reqwest, or Elasticsearch worker code +- `sqlite-fts` stores both authoritative data and its search index, written in + the same transactions, in one local SQLite file through Diesel, with FTS5 + operations expressed as SQL. It does not compile PostgreSQL, reqwest, or Elasticsearch worker code into the binary. ## Authentication @@ -59,8 +67,8 @@ cargo run --locked --package abyss-backend \ --no-default-features --features sqlite-fts ``` -The backend creates the parent directory, database, schema, and FTS5 index on -startup. A `sqlite://` prefix is also accepted, for example +The backend creates the parent directory, database, schema, and search index +on startup. A `sqlite://` prefix is also accepted, for example `sqlite://./data/abyss.sqlite`. ## Run with PostgreSQL @@ -96,11 +104,26 @@ curl \ | `GET` | `/readyz` | Selected storage readiness. | | `POST` | `/v1/agent-usage/events` | Ingest events and correlated diagnostic captures. | | `GET` | `/v1/agent-usage/events` | Query raw events. | +| `GET` | `/v1/agent-usage/events/{event_pk}` | Read one event. | +| `GET` | `/v1/agent-usage/events/{event_pk}/conversation` | Read one event's conversation, all of it or the asked entries. | +| `GET` | `/v1/agent-usage/events/{event_pk}/location` | Resolve an event to its current session, segment, and exchange. | | `GET` | `/v1/agent-usage/attachments/{id}` | Download stored image content. | | `GET` | `/v1/agent-usage/summary` | Aggregate event and token usage. | +| `GET` | `/v1/agent-usage/sessions` | List sessions with their titles, counts, and token totals. | +| `GET` | `/v1/agent-usage/sessions/{session_pk}/header` | Session title, counts, and token totals. | +| `GET` | `/v1/agent-usage/sessions/{session_pk}/streams` | Read the session’s derived segments. | +| `GET` | `/v1/agent-usage/sessions/{session_pk}/turns` | One node per turn: its executors' main-line and side segments. | +| `GET` | `/v1/agent-usage/sessions/{session_pk}/segments/{head}` | Read one segment’s exchanges. | | `POST` | `/v1/agent-usage/comparison` | Compare token totals and daily activity across dimensions. | -| `GET` | `/v1/agent-usage/sessions/{id}` | Read an ordered session timeline. | -| `GET` | `/v1/agent-usage/search` | Search sessions through Elasticsearch or SQLite FTS5. | +| `GET` | `/v1/agent-usage/search` | Search sessions through Elasticsearch or SQLite. | +| `GET` | `/v1/agent-usage/sessions/{session_pk}/search` | Every hit of a query in one session (SQLite). | + +A turn is one user input and all the work it causes; the backend reads only +the `turn_id` an event declares. Turn counts (session list, summary `turns`, +header `turn_count`) are, per session, the number of distinct `turn_id` values +when the session declares any, else the number of requests whose persisted +`starts_turn` fact is true; a summary row applies the same rule to the events +of its group. Every `/v1/agent-usage/*` endpoint requires the bearer token. Health and readiness endpoints are unauthenticated for orchestrator probes. @@ -119,6 +142,84 @@ builder for both storage backends. A small `QueryFragment` extension supplies th bound `VALUES` CTE for time buckets, which Diesel does not provide as a built-in query source. Bucket values remain bound and result types remain checked. +### Session structure and event links + +The header, streams, turns, segment details, in-session search, and event location +share one immutable derived structure for an owner/session/revision. The +process retains up to 128 idle revisions with LRU eviction. Concurrent reads +of the same revision share the load and derivation; a cache hit reads the +session revision without loading stream facts again. Event bodies are not +stored in this cache. + +Ingestion advances `agent_sessions.structure_revision` within its existing +session update and event transaction. Structure readers load the revision and +facts from the same database snapshot (a read transaction in SQLite, read-only +repeatable read in PostgreSQL). On a new revision or after eviction/restart, +the first read derives the full structure again. This is a shared query cache, +not an incremental derivation at ingest time. + +A request's base is the exchange that emitted the last identity in its +history, when exactly one other exchange did and the model output after that +identity was all said by it. Identities come from the conversation entries: +an entry's `item_id` (`item:`), a `tool_call` entry's `call_id`, and for an +`assistant` or `reasoning` entry with neither and non-blank text, a digest of +that exact text (`text::`). Text identities let a reply that +the protocol gives no id, such as a Chat Completions answer or an Anthropic text +block without thinking, link the next turn; rewritten text links nothing, and +text emitted by more than one exchange is ambiguous. + +Stable links retain `event_pk`. The location endpoint returns +`{"session_pk", "event_pk", "segment", "turn", "segment_kind", "exchange"}` using the current structure, +and returns 404 for an event outside the authenticated owner's data. A later +fork can change an event's segment, so links resolve the event when opened. +Segment IDs remain their first exchange IDs; a segment is not a turn number. + +`GET /events?locate=true` includes a `location` object on each returned event; +the list projection also preserves it. Session-search matches and the single +event endpoint include the same object: `{segment, turn, segment_kind, +exchange}`. `turn` is the one-based number of the turn of the event's exchange, +in the order the session's turns first appear, or null when the exchange belongs +to no turn; it is for display only. A response page looks up each session structure +once, sharing the cache with the detail view. The dashboard uses these values +in its segment column, search-hit labels, and inspector properties without +sending a separate location request for every row. + +### Search + +In the SQLite profile, search reads the pieces of each event's +`conversation.entries`: the `text` of every entry, whatever its +role, and its `name`, the tool name of a tool call. Entries are never joined, +an empty value is not a piece, and nothing else is searched. The query is one +literal substring, matched after folding every character to its lowercase +form when that is exactly one character. A hit is a piece of an event whose +folded text holds the folded query, so each event that replays history hits +again. A hit is new when its event is a response, or when its entry index is +at least the request's `new_from`; otherwise it is in the inherited context. + +`GET /v1/agent-usage/search?q=…` also takes `from`, `to`, `agent_name`, +`llm_provider`, `llm_model`, `event_type`, `page` (from 1), and `page_size` +(at most 50). Sessions are ordered by their latest matching event, newest +first. Each item carries `match_count`, its matching events, +`new_match_count`, those with a new hit, and up to three `matches`: the events +with a new hit, earliest first, then the earliest events whose hits are all +inherited. A match says whether it is `new` and carries up to two `fragments` +of distinct matching texts, new hits first, each about 240 characters around +the first occurrence, `…` on a cut side, as `{"text", "highlighted"}` +segments. + +`GET /v1/agent-usage/sessions/{session_pk}/search?q=…` returns every hit in +one session as `hits`, each `{"event_pk", "entry_index", "field", "new", +"text_id", "segment"}` with `field` `text` or `name` and `segment` the head of +the segment holding the event's exchange, as `/streams` gives it, or null. Hits +are in observation order, then by event, entry, and `name` before `text`. +`texts` lists each named text once as `{"id", "ranges"}`, the non-overlapping +occurrences as half-open ranges of the original text's UTF-16 code units. At +most 100,000 hits are returned; `truncated` says whether more were left out. + +The PostgreSQL profile searches through Elasticsearch, requires every +whitespace-separated term to match, counts every match as new, and does not +support search within a session. + ## Configuration | Variable | Required | Default | Description | @@ -143,7 +244,7 @@ query source. Bucket values remain bound and result types remain checked. Elasticsearch settings exist only in `postgres-es` builds. Username and password must be provided together. Search remains disabled in that profile when no URL is configured; event storage and queries continue to work. The -`sqlite-fts` profile always provides search through the local FTS5 index. +`sqlite-fts` profile always provides search through its local index. ## Native releases @@ -289,4 +390,4 @@ make k8s-render The PostgreSQL case starts an ephemeral container and an Elasticsearch contract double; the SQLite case uses only a temporary local database file. Both verify schema creation, authentication, event ingestion, attachment retrieval, raw -diagnostic retention, summary, timeline, and search behavior. +diagnostic retention, summary, streams, and search behavior. diff --git a/crates/abyss-backend/Cargo.toml b/crates/abyss-backend/Cargo.toml index 7e254c6..c038e9f 100644 --- a/crates/abyss-backend/Cargo.toml +++ b/crates/abyss-backend/Cargo.toml @@ -34,6 +34,7 @@ axum = "0.8" axum-extra = { version = "0.12", default-features = false, features = ["query"] } base64 = "0.21" chrono = { version = "0.4", features = ["serde"] } +chrono-tz = "0.10" diesel = { version = "2.3", default-features = false, features = ["32-column-tables", "r2d2", "with-deprecated"], optional = true } diesel_migrations = { version = "2.2", default-features = false, optional = true } hex = "0.4" diff --git a/crates/abyss-backend/migrations-sqlite/20261014150000_initial_schema/down.sql b/crates/abyss-backend/migrations-sqlite/20261014150000_initial_schema/down.sql index fbf1f4a..6ca5012 100644 --- a/crates/abyss-backend/migrations-sqlite/20261014150000_initial_schema/down.sql +++ b/crates/abyss-backend/migrations-sqlite/20261014150000_initial_schema/down.sql @@ -1,10 +1,8 @@ DROP TABLE session_titles; -DROP TABLE usage_events_fts; DROP TABLE agent_diagnostic_capture_events; DROP TABLE agent_diagnostic_captures; DROP TABLE llm_usage_event_attachments; DROP TABLE llm_usage_events; -DROP TABLE agent_turns; DROP TABLE agent_sessions; DROP TABLE devices; DROP TABLE app_users; diff --git a/crates/abyss-backend/migrations-sqlite/20261014150000_initial_schema/up.sql b/crates/abyss-backend/migrations-sqlite/20261014150000_initial_schema/up.sql index 483d88e..bddea07 100644 --- a/crates/abyss-backend/migrations-sqlite/20261014150000_initial_schema/up.sql +++ b/crates/abyss-backend/migrations-sqlite/20261014150000_initial_schema/up.sql @@ -41,48 +41,26 @@ CREATE TABLE agent_sessions ( metadata TEXT NOT NULL DEFAULT '{}' CHECK (json_valid(metadata)), created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, + -- Advances on event ingestion so derived stream caches share one revision. + structure_revision INTEGER NOT NULL DEFAULT 0, UNIQUE (user_id, agent_name, session_id) ) STRICT; CREATE INDEX agent_sessions_user_time_idx ON agent_sessions (user_id, started_at DESC); --- turn_key is the turn the Agent named for its requests, unique within the --- session; it is NULL for the turn opened by requests naming none. turn_index --- is the turn's position in the session by start time. -CREATE TABLE agent_turns ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL REFERENCES app_users(id) ON DELETE CASCADE, - session_pk TEXT NOT NULL REFERENCES agent_sessions(id) ON DELETE CASCADE, - turn_index INTEGER NOT NULL CHECK (turn_index >= 1), - turn_key TEXT, - started_at INTEGER NOT NULL, - ended_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL, - UNIQUE (session_pk, turn_index), - UNIQUE (session_pk, turn_key) -) STRICT; - -CREATE INDEX agent_turns_user_session_idx - ON agent_turns (user_id, session_pk, turn_index); - CREATE TABLE llm_usage_events ( id TEXT PRIMARY KEY, user_id TEXT NOT NULL REFERENCES app_users(id) ON DELETE CASCADE, device_context_id TEXT NOT NULL REFERENCES devices(id), session_pk TEXT NOT NULL REFERENCES agent_sessions(id) ON DELETE CASCADE, - turn_pk TEXT NOT NULL REFERENCES agent_turns(id) ON DELETE CASCADE, event_id TEXT NOT NULL UNIQUE, agent_name TEXT NOT NULL, agent_version TEXT, session_id TEXT NOT NULL, - turn_index INTEGER NOT NULL CHECK (turn_index >= 1), llm_provider TEXT NOT NULL, llm_model TEXT NOT NULL, event_type TEXT NOT NULL CHECK (event_type IN ('request', 'response')), - text TEXT, - text_sha256 BLOB, input_tokens INTEGER NOT NULL DEFAULT 0 CHECK (input_tokens >= 0), output_tokens INTEGER NOT NULL DEFAULT 0 CHECK (output_tokens >= 0), cache_read_tokens INTEGER NOT NULL DEFAULT 0 CHECK (cache_read_tokens >= 0), @@ -103,8 +81,8 @@ CREATE TABLE llm_usage_events ( CREATE INDEX llm_usage_events_user_time_idx ON llm_usage_events (user_id, observed_at DESC, id DESC); -CREATE INDEX llm_usage_events_session_turn_idx - ON llm_usage_events (session_pk, turn_index, observed_at, id); +CREATE INDEX llm_usage_events_session_order_idx + ON llm_usage_events (session_pk, observed_at, event_type, id); CREATE INDEX llm_usage_events_agent_model_time_idx ON llm_usage_events (agent_name, llm_provider, llm_model, observed_at DESC); @@ -157,22 +135,9 @@ CREATE TABLE agent_diagnostic_capture_events ( CREATE INDEX agent_diagnostic_capture_events_event_idx ON agent_diagnostic_capture_events (event_pk); -CREATE VIRTUAL TABLE usage_events_fts USING fts5( - event_pk UNINDEXED, - user_id UNINDEXED, - session_pk UNINDEXED, - session_id, - content, - tool_names, - tool_content, - tokenize = 'unicode61' -); - - --- Session names read back from the Agents' own title-generation calls, kept --- apart from agent_sessions so session rows keep their shape. purpose marks --- sessions that exist only to generate a title, such as Codex naming threads, --- and parent_session_pk the task session such a thread was matched to. +-- The earliest request is the title anchor, including requests without a label. +-- API titles read this event's stream facts rather than the title column. +-- purpose and parent_session_pk are separate session metadata. CREATE TABLE session_titles ( session_pk TEXT PRIMARY KEY REFERENCES agent_sessions(id) ON DELETE CASCADE, user_id TEXT NOT NULL REFERENCES app_users(id) ON DELETE CASCADE, @@ -184,16 +149,17 @@ CREATE TABLE session_titles ( updated_at INTEGER NOT NULL ) STRICT; --- Event rows hold whole request bodies in text and metadata, so reading a --- token column from the row walks through megabytes of overflow pages. These --- indexes carry every column the usage summary and the session list read, --- letting both answer from the index alone. +-- Event rows hold each event's conversation, tens of kilobytes, so a leaf +-- page holds about one row and reading any column from the row reads a page +-- per event. These indexes carry every column the usage summary and the +-- session list read, letting both answer from the index alone; the summary +-- joins stream facts by id. CREATE INDEX llm_usage_events_summary_cover_idx ON llm_usage_events ( user_id, observed_at, agent_name, llm_provider, llm_model, event_type, - device_context_id, session_pk, turn_pk, session_id, + device_context_id, session_pk, session_id, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, - reasoning_tokens, total_tokens + reasoning_tokens, total_tokens, id ); CREATE INDEX llm_usage_events_session_activity_idx ON llm_usage_events (session_pk, observed_at, total_tokens); diff --git a/crates/abyss-backend/migrations-sqlite/20261030130000_event_conversation_and_stream_facts/down.sql b/crates/abyss-backend/migrations-sqlite/20261030130000_event_conversation_and_stream_facts/down.sql new file mode 100644 index 0000000..0444924 --- /dev/null +++ b/crates/abyss-backend/migrations-sqlite/20261030130000_event_conversation_and_stream_facts/down.sql @@ -0,0 +1,8 @@ +DROP INDEX search_refs_text; +DROP TABLE search_refs; +DROP TABLE search_texts_trigram; +DROP TABLE search_texts; + +DROP TABLE llm_usage_event_stream_facts; + +ALTER TABLE llm_usage_events DROP COLUMN conversation; diff --git a/crates/abyss-backend/migrations-sqlite/20261030130000_event_conversation_and_stream_facts/up.sql b/crates/abyss-backend/migrations-sqlite/20261030130000_event_conversation_and_stream_facts/up.sql new file mode 100644 index 0000000..42a809b --- /dev/null +++ b/crates/abyss-backend/migrations-sqlite/20261030130000_event_conversation_and_stream_facts/up.sql @@ -0,0 +1,81 @@ +-- The AgentEvent's conversation, stored as received; NULL when the event +-- carries none. +ALTER TABLE llm_usage_events ADD COLUMN conversation TEXT CHECK (json_valid(conversation)); + +-- Stream facts read at ingest from one event's own conversation entries, +-- kept beside the event so a session's streams read them without walking +-- the overflow pages of the conversation. A request has new_from, +-- continues_context, starts_turn, last_identity, previous_identity, +-- trailing, result_ids, and its label; a response has emitted_ids and +-- output_skeleton; the other side's columns stay NULL. +-- continues_context: the request holds model output (its new part begins +-- after the first entry) or names the response it continues in +-- previous_response_id, as a request sending only what is new does. In a +-- session without declared turns, a turn-opening request opens a turn only +-- when it continues a context, or when no turn began before it. +-- previous_identity: the identity before the request's last identity, +-- repeats of it skipped; it tells apart the responses that emitted the same +-- last identity, such as an answer given word for word in every round. NULL +-- when nothing before the last identity has one. +-- The label is the last user entry with text in the request's new part: its +-- index, its first 50 UTF-16 code units, and its whole length in them. +-- trailing is a JSON array of roles, id lists JSON arrays of strings, and +-- output_skeleton a JSON array of [role, key] pairs, the key being the call +-- id, else the item id. +-- subagent_id, parent_subagent_id, and turn_id are the ids the client +-- declared on the event, kept as sent and NULL when absent. turn_id names +-- the user input the event works for; it decides the event's turn and +-- where a sub-agent's segment shows. +CREATE TABLE llm_usage_event_stream_facts ( + event_pk TEXT PRIMARY KEY REFERENCES llm_usage_events(id) ON DELETE CASCADE, + new_from INTEGER, + continues_context INTEGER, + starts_turn INTEGER, + last_identity TEXT, + previous_identity TEXT, + trailing TEXT CHECK (json_valid(trailing)), + result_ids TEXT CHECK (json_valid(result_ids)), + emitted_ids TEXT CHECK (json_valid(emitted_ids)), + output_skeleton TEXT CHECK (json_valid(output_skeleton)), + label_index INTEGER, + label_text TEXT, + label_length INTEGER, + subagent_id TEXT, + parent_subagent_id TEXT, + turn_id TEXT +) STRICT; + +-- Search reads the text and the name of every conversation entry. Each +-- distinct piece of text is stored once per owner, and every piece of an +-- event refers to its text; replayed history therefore adds references, not +-- text. A trigram index over the folded texts, without positions, narrows +-- the candidates the backend verifies. + +CREATE TABLE search_texts ( + id INTEGER PRIMARY KEY, + user_id TEXT NOT NULL, + -- SHA-256 of the text's UTF-8 bytes. + sha256 BLOB NOT NULL, + text TEXT NOT NULL, + UNIQUE (user_id, sha256) +) STRICT; + +-- Contentless and without positions; it indexes the folded text of each +-- row of search_texts under that row's id. +CREATE VIRTUAL TABLE search_texts_trigram USING fts5( + folded, + content = '', + detail = none, + tokenize = 'trigram case_sensitive 1' +); + +-- One row per nonempty `text` or `name` of an entry of an event. +CREATE TABLE search_refs ( + event_pk TEXT NOT NULL, + entry_index INTEGER NOT NULL, + field TEXT NOT NULL CHECK (field IN ('text', 'name')), + text_id INTEGER NOT NULL, + PRIMARY KEY (event_pk, entry_index, field) +) STRICT, WITHOUT ROWID; + +CREATE INDEX search_refs_text ON search_refs (text_id); diff --git a/crates/abyss-backend/migrations/20261014150000_initial_schema/down.sql b/crates/abyss-backend/migrations/20261014150000_initial_schema/down.sql index 06486fa..bf103b6 100644 --- a/crates/abyss-backend/migrations/20261014150000_initial_schema/down.sql +++ b/crates/abyss-backend/migrations/20261014150000_initial_schema/down.sql @@ -6,7 +6,6 @@ DROP TABLE IF EXISTS agent_diagnostic_capture_events; DROP TABLE IF EXISTS agent_diagnostic_captures; DROP TABLE IF EXISTS llm_usage_event_attachments; DROP TABLE IF EXISTS llm_usage_events; -DROP TABLE IF EXISTS agent_turns; DROP TABLE IF EXISTS agent_sessions; DROP TABLE IF EXISTS devices; DROP TABLE IF EXISTS app_users; diff --git a/crates/abyss-backend/migrations/20261014150000_initial_schema/up.sql b/crates/abyss-backend/migrations/20261014150000_initial_schema/up.sql index 8faca4a..dd5e7dc 100644 --- a/crates/abyss-backend/migrations/20261014150000_initial_schema/up.sql +++ b/crates/abyss-backend/migrations/20261014150000_initial_schema/up.sql @@ -47,32 +47,14 @@ CREATE TABLE agent_sessions ( metadata jsonb NOT NULL DEFAULT '{}'::jsonb, created_at timestamptz NOT NULL DEFAULT now(), updated_at timestamptz NOT NULL DEFAULT now(), + -- Advances on event ingestion so derived stream caches share one revision. + structure_revision bigint NOT NULL DEFAULT 0, UNIQUE (user_id, agent_name, session_id) ); CREATE INDEX agent_sessions_user_time_idx ON agent_sessions (user_id, started_at DESC); --- turn_key is the turn the Agent named for its requests, unique within the --- session; it is NULL for the turn opened by requests naming none. turn_index --- is the turn's position in the session by start time. -CREATE TABLE agent_turns ( - id uuid PRIMARY KEY, - user_id uuid NOT NULL REFERENCES app_users(id) ON DELETE CASCADE, - session_pk uuid NOT NULL REFERENCES agent_sessions(id) ON DELETE CASCADE, - turn_index integer NOT NULL CHECK (turn_index >= 1), - turn_key text, - started_at timestamptz NOT NULL, - ended_at timestamptz, - created_at timestamptz NOT NULL DEFAULT now(), - updated_at timestamptz NOT NULL DEFAULT now(), - UNIQUE (session_pk, turn_index), - UNIQUE (session_pk, turn_key) -); - -CREATE INDEX agent_turns_user_session_idx - ON agent_turns (user_id, session_pk, turn_index); - -- Usage events are immutable source records. event_id provides global ingest -- idempotency, while denormalized labels keep filters and aggregation direct. CREATE TABLE llm_usage_events ( @@ -80,17 +62,13 @@ CREATE TABLE llm_usage_events ( user_id uuid NOT NULL REFERENCES app_users(id) ON DELETE CASCADE, device_context_id uuid NOT NULL REFERENCES devices(id), session_pk uuid NOT NULL REFERENCES agent_sessions(id) ON DELETE CASCADE, - turn_pk uuid NOT NULL REFERENCES agent_turns(id) ON DELETE CASCADE, event_id text NOT NULL UNIQUE, agent_name text NOT NULL, agent_version text, session_id text NOT NULL, - turn_index integer NOT NULL CHECK (turn_index >= 1), llm_provider text NOT NULL, llm_model text NOT NULL, event_type text NOT NULL CHECK (event_type IN ('request', 'response')), - text text, - text_sha256 bytea, input_tokens bigint NOT NULL DEFAULT 0 CHECK (input_tokens >= 0), output_tokens bigint NOT NULL DEFAULT 0 CHECK (output_tokens >= 0), cache_read_tokens bigint NOT NULL DEFAULT 0 CHECK (cache_read_tokens >= 0), @@ -114,8 +92,8 @@ CREATE INDEX llm_usage_events_user_time_idx CREATE INDEX llm_usage_events_observed_at_idx ON llm_usage_events (observed_at DESC); -CREATE INDEX llm_usage_events_session_turn_idx - ON llm_usage_events (session_pk, turn_index, observed_at); +CREATE INDEX llm_usage_events_session_order_idx + ON llm_usage_events (session_pk, observed_at, event_type, id); CREATE INDEX llm_usage_events_agent_model_time_idx ON llm_usage_events (agent_name, llm_provider, llm_model, observed_at DESC); @@ -196,11 +174,10 @@ CREATE INDEX search_outbox_pending_idx ON search_outbox (available_at, id) WHERE processed_at IS NULL AND dead_lettered_at IS NULL; - --- Session names read back from the Agents' own title-generation calls, kept --- apart from agent_sessions so session rows keep their shape. purpose marks --- sessions that exist only to generate a title, such as Codex naming threads, --- and parent_session_pk the task session such a thread was matched to. +-- Session titles register the earliest request by observation time and event +-- primary key, including requests without a label. Readers join that event's +-- stream facts for its label rather than the title column. purpose and +-- parent_session_pk are separate session metadata. CREATE TABLE session_titles ( session_pk uuid PRIMARY KEY REFERENCES agent_sessions(id) ON DELETE CASCADE, user_id uuid NOT NULL REFERENCES app_users(id) ON DELETE CASCADE, diff --git a/crates/abyss-backend/migrations/20261030130000_event_conversation_and_stream_facts/down.sql b/crates/abyss-backend/migrations/20261030130000_event_conversation_and_stream_facts/down.sql new file mode 100644 index 0000000..b719941 --- /dev/null +++ b/crates/abyss-backend/migrations/20261030130000_event_conversation_and_stream_facts/down.sql @@ -0,0 +1,3 @@ +DROP TABLE llm_usage_event_stream_facts; + +ALTER TABLE llm_usage_events DROP COLUMN conversation; diff --git a/crates/abyss-backend/migrations/20261030130000_event_conversation_and_stream_facts/up.sql b/crates/abyss-backend/migrations/20261030130000_event_conversation_and_stream_facts/up.sql new file mode 100644 index 0000000..9760fdb --- /dev/null +++ b/crates/abyss-backend/migrations/20261030130000_event_conversation_and_stream_facts/up.sql @@ -0,0 +1,46 @@ +-- The AgentEvent's conversation, stored as received; NULL when the event +-- carries none. +ALTER TABLE llm_usage_events ADD COLUMN conversation jsonb; + +-- Stream facts read at ingest from one event's own conversation entries, +-- kept beside the event so a session's streams read them without touching +-- the conversation. A request has new_from, continues_context, starts_turn, +-- last_identity, previous_identity, trailing, result_ids, and its label; a +-- response has emitted_ids and output_skeleton; the other side's columns +-- stay NULL. +-- continues_context: the request holds model output (its new part begins +-- after the first entry) or names the response it continues in +-- previous_response_id, as a request sending only what is new does. In a +-- session without declared turns, a turn-opening request opens a turn only +-- when it continues a context, or when no turn began before it. +-- previous_identity: the identity before the request's last identity, +-- repeats of it skipped; it tells apart the responses that emitted the same +-- last identity, such as an answer given word for word in every round. NULL +-- when nothing before the last identity has one. +-- The label is the last user entry with text in the request's new part: its +-- index, its first 50 UTF-16 code units, and its whole length in them. +-- trailing is a JSON array of roles, id lists JSON arrays of strings, and +-- output_skeleton a JSON array of [role, key] pairs, the key being the call +-- id, else the item id. +-- subagent_id, parent_subagent_id, and turn_id are the ids the client +-- declared on the event, kept as sent and NULL when absent. turn_id names +-- the user input the event works for; it decides the event's turn and +-- where a sub-agent's segment shows. +CREATE TABLE llm_usage_event_stream_facts ( + event_pk uuid PRIMARY KEY REFERENCES llm_usage_events(id) ON DELETE CASCADE, + new_from integer, + continues_context boolean, + starts_turn boolean, + last_identity text, + previous_identity text, + "trailing" jsonb, + result_ids jsonb, + emitted_ids jsonb, + output_skeleton jsonb, + label_index integer, + label_text text, + label_length integer, + subagent_id text, + parent_subagent_id text, + turn_id text +); diff --git a/crates/abyss-backend/src/api.rs b/crates/abyss-backend/src/api.rs index 738e7ed..069d3cf 100644 --- a/crates/abyss-backend/src/api.rs +++ b/crates/abyss-backend/src/api.rs @@ -1,7 +1,7 @@ //! HTTP boundary for standalone Agent event ingestion and queries. //! //! The health endpoints are public, while every event, attachment, summary, -//! timeline, and search endpoint authenticates the deployment bearer token. +//! session, and search endpoint authenticates the deployment bearer token. //! Storage implementation details remain behind the configured backend trait. use axum::{ @@ -22,15 +22,19 @@ use uuid::Uuid; use crate::{ error::AppError, identity::IdentityAuthenticator, - search::{SessionSearchQuery, SessionSearchResponse}, + search::{ + InSessionSearchQuery, InSessionSearchResponse, SessionSearchQuery, SessionSearchResponse, + }, storage::StorageBackend, usage::{ IngestEventsRequest, IngestEventsResponse, RawEventListResponse, RawEventsQuery, - SessionOutlineResponse, SessionTimelineResponse, SummaryFields, SummaryQuery, - TokenUsageSummaryResponse, UsageEventResponse, + SummaryFields, SummaryQuery, TokenUsageSummaryResponse, UsageEventResponse, comparison::{ComparisonQuery, ComparisonResponse}, + preview::{ConversationPreview, PreviewQuery, Window}, + segments::{SegmentQuery, SegmentResponse}, + session_header::SessionHeaderResponse, session_list::{SessionListQuery, SessionListResponse}, - trajectory::{self, SessionTrajectoryResponse}, + streams::{EventLocationResponse, SessionStreamsResponse, SessionTurnsResponse}, }, }; @@ -49,7 +53,7 @@ pub struct AppState { pub default_page_size: i64, /// Deployment-wide bearer-token validator. pub identity: IdentityAuthenticator, - /// Compile-time-selected event storage and full-text search implementation. + /// Compile-time-selected event storage and search implementation. pub storage: std::sync::Arc, } @@ -69,6 +73,10 @@ pub fn router(state: AppState) -> Router { .layer(DefaultBodyLimit::max(MAX_INGEST_REQUEST_BODY_BYTES)), ) .route("/v1/agent-usage/events/{event_pk}", get(usage_event)) + .route( + "/v1/agent-usage/events/{event_pk}/location", + get(event_location), + ) .route( "/v1/agent-usage/attachments/{attachment_id}", get(image_attachment), @@ -78,16 +86,28 @@ pub fn router(state: AppState) -> Router { .route("/v1/agent-usage/search", get(session_search)) .route("/v1/agent-usage/sessions", get(list_sessions)) .route( - "/v1/agent-usage/sessions/{session_pk}", - get(session_timeline), + "/v1/agent-usage/sessions/{session_pk}/header", + get(session_header), ) .route( - "/v1/agent-usage/sessions/{session_pk}/outline", - get(session_outline), + "/v1/agent-usage/sessions/{session_pk}/streams", + get(session_streams), ) .route( - "/v1/agent-usage/sessions/{session_pk}/trajectory", - get(session_trajectory), + "/v1/agent-usage/sessions/{session_pk}/turns", + get(session_turns), + ) + .route( + "/v1/agent-usage/sessions/{session_pk}/segments/{head}", + get(session_segment), + ) + .route( + "/v1/agent-usage/sessions/{session_pk}/search", + get(search_in_session), + ) + .route( + "/v1/agent-usage/events/{event_pk}/conversation", + get(event_conversation), ) .layer(TraceLayer::new_for_http()) .with_state(state) @@ -179,35 +199,102 @@ async fn list_sessions( Ok(Json(state.storage.list_sessions(user_id, query).await?)) } -async fn session_timeline( +/// Segments of one session, without content. +async fn session_streams( State(state): State, headers: HeaderMap, Path(session_pk): Path, -) -> Result, AppError> { +) -> Result, AppError> { let user_id = state.identity.authenticate(&headers)?; - let response = state.storage.session_timeline(user_id, session_pk).await?; - Ok(Json(response)) + Ok(Json( + state.storage.session_streams(user_id, session_pk).await?, + )) } -async fn session_outline( +/// Turns of one session, each with its executors' segments, without content. +async fn session_turns( State(state): State, headers: HeaderMap, Path(session_pk): Path, -) -> Result, AppError> { +) -> Result, AppError> { let user_id = state.identity.authenticate(&headers)?; Ok(Json( - state.storage.session_outline(user_id, session_pk).await?, + state.storage.session_turns(user_id, session_pk).await?, )) } -async fn session_trajectory( +/// Every hit of a query in one session, each with the head of the segment +/// holding its event. +async fn search_in_session( State(state): State, headers: HeaderMap, Path(session_pk): Path, -) -> Result, AppError> { + Query(query): Query, +) -> Result, AppError> { + let user_id = state.identity.authenticate(&headers)?; + Ok(Json( + state + .storage + .search_in_session(user_id, session_pk, query) + .await?, + )) +} + +/// Agents, time span, counts, and token totals of one session. +async fn session_header( + State(state): State, + headers: HeaderMap, + Path(session_pk): Path, +) -> Result, AppError> { + let user_id = state.identity.authenticate(&headers)?; + Ok(Json( + state.storage.session_header(user_id, session_pk).await?, + )) +} + +/// A segment's exchanges from exchange `head` to the segment end, with a +/// preview of each request and response. +async fn session_segment( + State(state): State, + headers: HeaderMap, + Path((session_pk, head)): Path<(Uuid, String)>, + Query(query): Query, +) -> Result, AppError> { + let user_id = state.identity.authenticate(&headers)?; + let query = query.validate()?; + Ok(Json( + state + .storage + .session_segment(user_id, session_pk, head, query) + .await?, + )) +} + +/// The conversation of one event, all of it or the asked entries. +async fn event_conversation( + State(state): State, + headers: HeaderMap, + Path(event_pk): Path, + Query(query): Query, +) -> Result, AppError> { + let user_id = state.identity.authenticate(&headers)?; + let window = Window::parse(&query)?; + Ok(Json( + state + .storage + .event_conversation(user_id, event_pk, window) + .await?, + )) +} + +/// Resolves an event anchor against the current session structure. +async fn event_location( + State(state): State, + headers: HeaderMap, + Path(event_pk): Path, +) -> Result, AppError> { let user_id = state.identity.authenticate(&headers)?; - let timeline = state.storage.session_timeline(user_id, session_pk).await?; - Ok(Json(trajectory::build(timeline))) + Ok(Json(state.storage.event_location(user_id, event_pk).await?)) } async fn usage_event( diff --git a/crates/abyss-backend/src/db/models.rs b/crates/abyss-backend/src/db/models.rs index 74cb971..45c45ee 100644 --- a/crates/abyss-backend/src/db/models.rs +++ b/crates/abyss-backend/src/db/models.rs @@ -1,7 +1,7 @@ //! Diesel row and insert models owned by the Agent event service. //! -//! Query models mirror complete table rows because Diesel uses them for typed -//! selection. Insert models make server-assigned fields explicit and prevent +//! Query models select the fields each consumer needs through Diesel. +//! Insert models make server-assigned fields explicit and prevent //! HTTP payloads from being written directly without normalization. use chrono::{DateTime, Utc}; @@ -10,8 +10,8 @@ use serde_json::Value; use uuid::Uuid; use super::schema::{ - agent_diagnostic_capture_events, agent_diagnostic_captures, agent_sessions, agent_turns, - devices, llm_usage_event_attachments, llm_usage_events, search_outbox, + agent_diagnostic_capture_events, agent_diagnostic_captures, agent_sessions, devices, + llm_usage_event_attachments, llm_usage_event_stream_facts, llm_usage_events, search_outbox, }; #[derive(Clone, Debug, Queryable, Selectable)] @@ -66,14 +66,10 @@ pub struct NewDevice { #[derive(Clone, Debug, Queryable, Selectable)] #[diesel(table_name = agent_sessions)] #[diesel(check_for_backend(diesel::pg::Pg))] -/// Persisted Agent session containing many normalized turns. +/// Persisted Agent session containing request and response events. pub struct AgentSession { /// Server-generated session primary key. pub id: Uuid, - /// Owning user identifier. - pub user_id: Uuid, - /// Most recently associated device-context key. - pub device_context_id: Uuid, /// Agent name as reported by the collector. pub agent_name: String, /// Most recently observed Agent version. @@ -84,12 +80,6 @@ pub struct AgentSession { pub started_at: DateTime, /// Latest event observation in the session. pub ended_at: Option>, - /// Stable allowlisted session metadata. - pub metadata: Value, - /// Row creation time. - pub created_at: DateTime, - /// Last session metadata update time. - pub updated_at: DateTime, } #[derive(Debug, Insertable)] @@ -118,57 +108,8 @@ pub struct NewAgentSession { pub created_at: DateTime, /// Initial metadata update time. pub updated_at: DateTime, -} - -#[derive(Clone, Debug, Queryable, Selectable)] -#[diesel(table_name = agent_turns)] -#[diesel(check_for_backend(diesel::pg::Pg))] -/// Persisted logical turn within one Agent session. -pub struct AgentTurn { - /// Server-generated turn primary key. - pub id: Uuid, - /// Owning user identifier. - pub user_id: Uuid, - /// Parent session primary key. - pub session_pk: Uuid, - /// Position of the turn in its session by start time. - pub turn_index: i32, - /// Turn the Agent named for its requests; `None` for the turn opened by - /// requests naming none. - pub turn_key: Option, - /// Earliest event observation in the turn. - pub started_at: DateTime, - /// Latest event observation in the turn. - pub ended_at: Option>, - /// Row creation time. - pub created_at: DateTime, - /// Last turn-boundary update time. - pub updated_at: DateTime, -} - -#[derive(Debug, Insertable)] -#[diesel(table_name = agent_turns)] -/// Values used to create a logical Agent turn. -pub struct NewAgentTurn { - /// Server-generated turn primary key. - pub id: Uuid, - /// Owning user identifier. - pub user_id: Uuid, - /// Parent session primary key. - pub session_pk: Uuid, - /// Position of the turn in its session by start time. - pub turn_index: i32, - /// Turn the Agent named for its requests; `None` for the turn opened by - /// requests naming none. - pub turn_key: Option, - /// Initial earliest observation time. - pub started_at: DateTime, - /// Initial latest observation time. - pub ended_at: Option>, - /// Row creation time. - pub created_at: DateTime, - /// Initial boundary update time. - pub updated_at: DateTime, + /// Changes whenever ingestion adds an event to the session. + pub structure_revision: i64, } #[derive(Clone, Debug, Queryable, Selectable)] @@ -184,8 +125,6 @@ pub struct UsageEvent { pub device_context_id: Uuid, /// Parent session primary key. pub session_pk: Uuid, - /// Parent normalized-turn primary key. - pub turn_pk: Uuid, /// Collector-generated idempotency key. pub event_id: String, /// Agent name as reported by the collector. @@ -194,18 +133,12 @@ pub struct UsageEvent { pub agent_version: Option, /// Agent-native session identifier copied for direct filtering. pub session_id: String, - /// Normalized turn number copied for direct filtering. - pub turn_index: i32, /// LLM provider as reported by the collector. pub llm_provider: String, /// Provider-specific model identifier. pub llm_model: String, /// Canonical request or response value. pub event_type: String, - /// Optional prompt or response text. - pub text: Option, - /// Binary SHA-256 digest of text when present. - pub text_sha256: Option>, /// Normalized input-token count. pub input_tokens: i64, /// Normalized output-token count. @@ -230,6 +163,8 @@ pub struct UsageEvent { pub response_id: Option, /// Response whose conversation the exchange's request continues. pub previous_response_id: Option, + /// Conversation of the event, as the collector reported it. + pub conversation: Option, } #[derive(Queryable, Selectable)] @@ -305,8 +240,6 @@ pub struct NewUsageEvent { pub device_context_id: Uuid, /// Parent session primary key. pub session_pk: Uuid, - /// Parent normalized-turn primary key. - pub turn_pk: Uuid, /// Collector-generated idempotency key. pub event_id: String, /// Agent name as reported by the collector. @@ -315,18 +248,12 @@ pub struct NewUsageEvent { pub agent_version: Option, /// Agent-native session identifier. pub session_id: String, - /// Backend-normalized turn number. - pub turn_index: i32, /// LLM provider as reported by the collector. pub llm_provider: String, /// Provider-specific model identifier. pub llm_model: String, /// Canonical request or response value. pub event_type: String, - /// Optional prompt or response text. - pub text: Option, - /// Optional binary SHA-256 of text. - pub text_sha256: Option>, /// Normalized input-token count. pub input_tokens: i64, /// Normalized output-token count. @@ -351,6 +278,93 @@ pub struct NewUsageEvent { pub response_id: Option, /// Response whose conversation the exchange's request continues. pub previous_response_id: Option, + /// Conversation of the event, as the collector reported it. + pub conversation: Option, +} + +#[derive(Queryable)] +/// The columns of one session event that its streams read; the facts are +/// NULL for an event stored before they were read. +pub struct StreamEventRow { + /// Event primary key. + pub id: Uuid, + /// Collector-generated idempotency key. + pub event_id: String, + /// Canonical request or response value. + pub event_type: String, + /// Collector observation time. + pub observed_at: DateTime, + /// Provider identity of the response the exchange produced. + pub response_id: Option, + /// Response whose conversation the exchange's request continues. + pub previous_response_id: Option, + /// Request: index where the new part begins. + pub new_from: Option, + /// Request: whether the new part opens a turn. + pub starts_turn: Option, + /// Request: last identity of the last run. + pub last_identity: Option, + /// Request: the identity before the last one. + pub previous_identity: Option, + /// Request: roles of the model-side entries after it. + pub trailing: Option, + /// Request without model output: call ids of its tool results. + pub result_ids: Option, + /// Response: identities of its entries. + pub emitted_ids: Option, + /// Response: skeleton of its entries. + pub output_skeleton: Option, + /// Request: index of its label entry. + pub label_index: Option, + /// Request: start of its label entry's text. + pub label_text: Option, + /// Request: length of its label entry's text in UTF-16 code units. + pub label_length: Option, + /// Sub-agent id the client declared for the event. + pub subagent_id: Option, + /// Id of the sub-agent the client declared as this one's parent. + pub parent_subagent_id: Option, + /// Turn the client declared the event works for. + pub turn_id: Option, +} + +#[derive(Debug, Default, Insertable)] +#[diesel(table_name = llm_usage_event_stream_facts)] +/// Stream facts of a new event; the other side's columns stay NULL. +pub struct NewStreamFacts { + /// Event the facts were read from. + pub event_pk: Uuid, + /// Request: index where the new part begins. + pub new_from: Option, + /// Request: whether the new part opens a turn. + pub starts_turn: Option, + /// Request: last identity of the last run. + pub last_identity: Option, + /// Request: the identity before the last one. + pub previous_identity: Option, + /// Request: roles of the model-side entries after it. + pub trailing: Option, + /// Request without model output: call ids of its tool results. + pub result_ids: Option, + /// Response: identities of its entries. + pub emitted_ids: Option, + /// Response: skeleton of its entries. + pub output_skeleton: Option, + /// Request: index of its label entry. + pub label_index: Option, + /// Request: start of its label entry's text. + pub label_text: Option, + /// Request: length of its label entry's text in UTF-16 code units. + pub label_length: Option, + /// Sub-agent id the client declared for the event. + pub subagent_id: Option, + /// Id of the sub-agent the client declared as this one's parent. + pub parent_subagent_id: Option, + /// Turn the client declared the event works for. + pub turn_id: Option, + /// Request: whether it holds model output or names the response it + /// continues. + pub continues_context: Option, } #[derive(Insertable)] diff --git a/crates/abyss-backend/src/db/schema.rs b/crates/abyss-backend/src/db/schema.rs index 5cbd236..c050bf9 100644 --- a/crates/abyss-backend/src/db/schema.rs +++ b/crates/abyss-backend/src/db/schema.rs @@ -28,7 +28,7 @@ diesel::table! { } diesel::table! { - use diesel::sql_types::{Jsonb, Nullable, Text, Timestamptz, Uuid}; + use diesel::sql_types::{BigInt, Jsonb, Nullable, Text, Timestamptz, Uuid}; agent_sessions (id) { id -> Uuid, @@ -42,44 +42,25 @@ diesel::table! { metadata -> Jsonb, created_at -> Timestamptz, updated_at -> Timestamptz, + structure_revision -> BigInt, } } diesel::table! { - use diesel::sql_types::{Integer, Nullable, Text, Timestamptz, Uuid}; - - agent_turns (id) { - id -> Uuid, - user_id -> Uuid, - session_pk -> Uuid, - turn_index -> Integer, - turn_key -> Nullable, - started_at -> Timestamptz, - ended_at -> Nullable, - created_at -> Timestamptz, - updated_at -> Timestamptz, - } -} - -diesel::table! { - use diesel::sql_types::{BigInt, Bytea, Integer, Jsonb, Nullable, Text, Timestamptz, Uuid}; + use diesel::sql_types::{BigInt, Integer, Jsonb, Nullable, Text, Timestamptz, Uuid}; llm_usage_events (id) { id -> Uuid, user_id -> Uuid, device_context_id -> Uuid, session_pk -> Uuid, - turn_pk -> Uuid, event_id -> Text, agent_name -> Text, agent_version -> Nullable, session_id -> Text, - turn_index -> Integer, llm_provider -> Text, llm_model -> Text, event_type -> Text, - text -> Nullable, - text_sha256 -> Nullable, input_tokens -> BigInt, output_tokens -> BigInt, cache_read_tokens -> BigInt, @@ -92,6 +73,30 @@ diesel::table! { stream_id -> Nullable, response_id -> Nullable, previous_response_id -> Nullable, + conversation -> Nullable, + } +} + +diesel::table! { + use diesel::sql_types::{Bool, Integer, Jsonb, Nullable, Text, Uuid}; + + llm_usage_event_stream_facts (event_pk) { + event_pk -> Uuid, + new_from -> Nullable, + continues_context -> Nullable, + starts_turn -> Nullable, + last_identity -> Nullable, + previous_identity -> Nullable, + trailing -> Nullable, + result_ids -> Nullable, + emitted_ids -> Nullable, + output_skeleton -> Nullable, + label_index -> Nullable, + label_text -> Nullable, + label_length -> Nullable, + subagent_id -> Nullable, + parent_subagent_id -> Nullable, + turn_id -> Nullable, } } @@ -173,15 +178,13 @@ diesel::table! { diesel::joinable!(agent_sessions -> app_users (user_id)); diesel::joinable!(agent_sessions -> devices (device_context_id)); -diesel::joinable!(agent_turns -> agent_sessions (session_pk)); -diesel::joinable!(agent_turns -> app_users (user_id)); diesel::joinable!(devices -> app_users (user_id)); diesel::joinable!(llm_usage_events -> agent_sessions (session_pk)); -diesel::joinable!(llm_usage_events -> agent_turns (turn_pk)); diesel::joinable!(llm_usage_events -> app_users (user_id)); diesel::joinable!(llm_usage_events -> devices (device_context_id)); diesel::joinable!(llm_usage_event_attachments -> app_users (user_id)); diesel::joinable!(llm_usage_event_attachments -> llm_usage_events (event_pk)); +diesel::joinable!(llm_usage_event_stream_facts -> llm_usage_events (event_pk)); diesel::joinable!(agent_diagnostic_capture_events -> agent_diagnostic_captures (capture_pk)); diesel::joinable!(agent_diagnostic_capture_events -> llm_usage_events (event_pk)); diesel::joinable!(agent_diagnostic_captures -> agent_sessions (session_pk)); @@ -193,11 +196,11 @@ diesel::allow_tables_to_appear_in_same_query!( agent_diagnostic_capture_events, agent_diagnostic_captures, agent_sessions, - agent_turns, app_users, devices, llm_usage_events, llm_usage_event_attachments, + llm_usage_event_stream_facts, search_outbox, session_titles, ); diff --git a/crates/abyss-backend/src/search/matching/fold.rs b/crates/abyss-backend/src/search/matching/fold.rs new file mode 100644 index 0000000..7332e7a --- /dev/null +++ b/crates/abyss-backend/src/search/matching/fold.rs @@ -0,0 +1,75 @@ +//! Per-character case folding applied alike to texts and queries. + +/// One character's lowercase form when that is exactly one character, else +/// the character itself. +#[must_use] +pub fn fold_char(character: char) -> char { + if character.is_ascii() { + return character.to_ascii_lowercase(); + } + let mut lowercase = character.to_lowercase(); + match (lowercase.next(), lowercase.next()) { + (Some(folded), None) => folded, + _ => character, + } +} + +/// Folds every character of `text`. The result has as many characters as +/// `text`, so a character position in it is the same position in `text`. +#[must_use] +pub fn fold(text: &str) -> String { + if text.is_ascii() { + return text.to_ascii_lowercase(); + } + text.chars().map(fold_char).collect() +} + +#[cfg(test)] +mod tests { + use super::fold; + + #[test] + fn folds_case_one_character_at_a_time() { + assert_eq!(fold("ÄBC"), fold("äbc")); + assert_eq!(fold("ÄBC"), "äbc"); + assert_eq!( + fold("ΣΊΣΥΦΟΣ"), + "σίσυφοσ", + "Greek folds without final-sigma context" + ); + assert_eq!(fold("Cargo TEST"), "cargo test"); + assert_eq!(fold("搜索 STATUS.md"), "搜索 status.md"); + } + + #[test] + fn keeps_characters_whose_lowercase_is_not_one_character() { + // İ lowercases to two characters, i and a combining dot. + assert_eq!(fold("İstanbul"), "İstanbul"); + assert_eq!(fold("ß"), "ß"); + assert_ne!(fold("Straße"), fold("STRASSE"), "ß does not expand to ss"); + assert_eq!( + fold("ẞ"), + "ß", + "capital sharp s lowercases to one character" + ); + } + + #[test] + fn leaves_width_and_compatibility_forms_alone() { + assert_eq!( + fold("ＡＢＣ"), + "ａｂｃ", + "full-width letters fold within their width" + ); + assert_ne!(fold("ＡＢＣ"), fold("abc")); + assert_ne!(fold("fi"), fold("fi")); + assert_eq!(fold("😀 Ü"), "😀 ü"); + } + + #[test] + fn keeps_the_character_count() { + for text in ["İİ", "ΣΑΣ", "Kelvin K", "Dž", "😀a", ""] { + assert_eq!(fold(text).chars().count(), text.chars().count(), "{text}"); + } + } +} diff --git a/crates/abyss-backend/src/search/matching/fragment.rs b/crates/abyss-backend/src/search/matching/fragment.rs new file mode 100644 index 0000000..9695f5d --- /dev/null +++ b/crates/abyss-backend/src/search/matching/fragment.rs @@ -0,0 +1,194 @@ +//! A highlighted window of a matching text around its first occurrence. + +use std::ops::Range; + +use crate::search::{SearchFragment, SearchFragmentSegment}; + +use super::pattern::character_offsets; + +/// Characters a fragment keeps, counted in code points, unless the +/// occurrence alone is longer. +pub const FRAGMENT_CHARACTERS: usize = 240; + +/// Marks the side where a fragment cuts its text. +const ELLIPSIS: &str = "…"; + +/// The window of `text` centered on the first of `occurrences`, which are +/// sorted, non-overlapping character ranges, with every occurrence in the +/// window highlighted. An occurrence the window cuts is highlighted where +/// it is visible. A cut side starts or ends with `…`. None without an +/// occurrence. +#[must_use] +pub fn fragment(text: &str, occurrences: &[Range]) -> Option { + let first = occurrences.first()?; + let characters = text.chars().count(); + let window = window(first, characters); + let visible = occurrences + .iter() + .filter(|occurrence| occurrence.start < window.end && occurrence.end > window.start) + .map(|occurrence| occurrence.start.max(window.start)..occurrence.end.min(window.end)) + .collect::>(); + let boundaries = std::iter::once(window.start) + .chain(visible.iter().flat_map(|range| [range.start, range.end])) + .chain(std::iter::once(window.end)); + let bytes = character_offsets(text, boundaries, char::len_utf8); + let (window_start, rest) = bytes.split_first()?; + let (window_end, highlights) = rest.split_last()?; + + let mut segments = Vec::new(); + let mut plain_start = *window_start; + let (highlights, _) = highlights.as_chunks::<2>(); + for &[start, end] in highlights { + push_segment(&mut segments, text.get(plain_start..start)?, false); + push_segment(&mut segments, text.get(start..end)?, true); + plain_start = end; + } + push_segment(&mut segments, text.get(plain_start..*window_end)?, false); + if window.start > 0 { + match segments.first_mut() { + Some(segment) if !segment.highlighted => segment.text.insert_str(0, ELLIPSIS), + _ => segments.insert(0, plain(ELLIPSIS)), + } + } + if window.end < characters { + match segments.last_mut() { + Some(segment) if !segment.highlighted => segment.text.push_str(ELLIPSIS), + _ => segments.push(plain(ELLIPSIS)), + } + } + Some(SearchFragment { segments }) +} + +/// About [`FRAGMENT_CHARACTERS`] characters with `occurrence` in the middle, +/// shifted inward where the text ends sooner on one side. +fn window(occurrence: &Range, characters: usize) -> Range { + let length = occurrence.end.saturating_sub(occurrence.start); + let width = FRAGMENT_CHARACTERS.max(length); + let before = width + .saturating_sub(length) + .checked_div(2) + .unwrap_or_default(); + let start = occurrence.start.saturating_sub(before); + let end = characters.min(start.saturating_add(width)); + end.saturating_sub(width).min(start)..end +} + +fn push_segment(segments: &mut Vec, text: &str, highlighted: bool) { + if !text.is_empty() { + segments.push(SearchFragmentSegment { + text: text.to_owned(), + highlighted, + }); + } +} + +fn plain(text: &str) -> SearchFragmentSegment { + SearchFragmentSegment { + text: text.to_owned(), + highlighted: false, + } +} + +#[cfg(test)] +mod tests { + use super::{FRAGMENT_CHARACTERS, fragment}; + use crate::search::{SearchFragment, matching::pattern::Pattern}; + + fn render(fragment: &SearchFragment) -> String { + fragment + .segments + .iter() + .map(|segment| { + if segment.highlighted { + format!("[{}]", segment.text) + } else { + segment.text.clone() + } + }) + .collect() + } + + fn fragment_of(text: &str, query: &str) -> SearchFragment { + fragment(text, &Pattern::new(query).occurrences(text)).expect("the text matches") + } + + #[test] + fn a_short_text_is_kept_whole_with_every_occurrence_highlighted() { + assert_eq!( + render(&fragment_of("Find the Needle, then the needle.", "needle")), + "Find the [Needle], then the [needle]." + ); + assert_eq!(render(&fragment_of("needle", "needle")), "[needle]"); + assert_eq!(fragment("no match", &[]), None); + } + + #[test] + fn a_long_text_is_cut_around_the_first_occurrence() { + let text = format!("{}needle{}", "a".repeat(500), "b".repeat(500)); + let rendered = render(&fragment_of(&text, "needle")); + // (240 - 6) / 2 characters on each side. + let half = 117; + assert_eq!( + rendered, + format!("…{}[needle]{}…", "a".repeat(half), "b".repeat(half)) + ); + let shown = rendered + .chars() + .filter(|&character| character != '…') + .count(); + assert_eq!( + shown, + FRAGMENT_CHARACTERS + 2, + "two brackets around the highlight" + ); + } + + #[test] + fn the_window_shifts_inward_at_either_end() { + let start = format!("needle{}", "b".repeat(500)); + assert_eq!( + render(&fragment_of(&start, "needle")), + format!("[needle]{}…", "b".repeat(FRAGMENT_CHARACTERS - 6)) + ); + let end = format!("{}needle", "会".repeat(500)); + assert_eq!( + render(&fragment_of(&end, "needle")), + format!("…{}[needle]", "会".repeat(FRAGMENT_CHARACTERS - 6)), + "the window counts code points" + ); + } + + #[test] + fn later_occurrences_in_the_window_are_highlighted_and_cut_ones_partly() { + let text = format!( + "{}needle{}needle{}", + "x".repeat(200), + "b".repeat(114), + "c".repeat(100) + ); + // The window is [83, 323): it ends inside the second occurrence, + // which starts at 320. + let rendered = render(&fragment_of(&text, "needle")); + assert!(rendered.starts_with('…'), "{rendered}"); + assert!(rendered.contains("[needle]"), "{rendered}"); + assert!(rendered.ends_with("[nee]…"), "{rendered}"); + } + + #[test] + fn an_occurrence_longer_than_the_window_is_shown_whole() { + let query = "q".repeat(FRAGMENT_CHARACTERS + 10); + let text = format!("before {query} after"); + assert_eq!(render(&fragment_of(&text, &query)), format!("…[{query}]…")); + } + + #[test] + fn highlights_follow_the_original_case_and_astral_characters() { + assert_eq!( + render(&fragment_of( + "😀 NIMBUS.example.org 😀", + "nimbus.EXAMPLE.org" + )), + "😀 [NIMBUS.example.org] 😀" + ); + } +} diff --git a/crates/abyss-backend/src/search/matching/in_session.rs b/crates/abyss-backend/src/search/matching/in_session.rs new file mode 100644 index 0000000..56b8f56 --- /dev/null +++ b/crates/abyss-backend/src/search/matching/in_session.rs @@ -0,0 +1,137 @@ +//! The bounded response of a search within one session. + +use std::collections::{BTreeSet, HashMap}; + +use crate::search::{InSessionHit, InSessionSearchResponse, MatchedText}; + +/// Hits a search within one session returns at most. +pub const MAX_IN_SESSION_HITS: usize = 100_000; + +impl InSessionSearchResponse { + /// Keeps the first [`MAX_IN_SESSION_HITS`] of `hits`, which are in + /// response order and hold at most one more, with the occurrence ranges + /// of each text they name, ordered by text id. + #[must_use] + pub fn from_hits( + query: String, + mut hits: Vec, + mut ranges: HashMap>, + ) -> Self { + let truncated = hits.len() > MAX_IN_SESSION_HITS; + hits.truncate(MAX_IN_SESSION_HITS); + let texts = hits + .iter() + .map(|hit| hit.text_id) + .collect::>() + .into_iter() + .filter_map(|id| { + Some(MatchedText { + id, + ranges: ranges.remove(&id)?, + }) + }) + .collect(); + Self { + query, + texts, + hits, + truncated, + } + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use serde_json::json; + use uuid::Uuid; + + use super::MAX_IN_SESSION_HITS; + use crate::search::{InSessionHit, InSessionSearchResponse, MatchedText, PieceField}; + + fn hit(text_id: i64) -> InSessionHit { + InSessionHit { + event_pk: Uuid::from_u128(1), + entry_index: 0, + field: PieceField::Text, + new: true, + text_id, + segment: Some("e1".to_owned()), + } + } + + #[test] + fn names_only_the_texts_of_the_returned_hits() { + let response = InSessionSearchResponse::from_hits( + "nimbus".to_owned(), + vec![hit(12), hit(3), hit(12)], + HashMap::from([ + (3, vec![[0, 6]]), + (12, vec![[3, 9], [40, 46]]), + (99, vec![[0, 1]]), + ]), + ); + assert!(!response.truncated); + assert_eq!( + response.texts, + [ + MatchedText { + id: 3, + ranges: vec![[0, 6]] + }, + MatchedText { + id: 12, + ranges: vec![[3, 9], [40, 46]] + } + ] + ); + assert_eq!( + serde_json::to_value(&response).expect("serializes"), + json!({ + "query": "nimbus", + "texts": [{"id": 3_i64, "ranges": [[0_i64, 6_i64]]}, {"id": 12_i64, "ranges": [[3_i64, 9_i64], [40_i64, 46_i64]]}], + "hits": [ + {"event_pk": Uuid::from_u128(1), "entry_index": 0_i64, "field": "text", "new": true, "text_id": 12_i64, "segment": "e1"}, + {"event_pk": Uuid::from_u128(1), "entry_index": 0_i64, "field": "text", "new": true, "text_id": 3_i64, "segment": "e1"}, + {"event_pk": Uuid::from_u128(1), "entry_index": 0_i64, "field": "text", "new": true, "text_id": 12_i64, "segment": "e1"} + ], + "truncated": false + }) + ); + } + + #[test] + fn caps_the_hits_and_reports_truncation() { + let at_cap = InSessionSearchResponse::from_hits( + "x".to_owned(), + vec![hit(1); MAX_IN_SESSION_HITS], + HashMap::from([(1, vec![[0, 1]])]), + ); + assert_eq!( + (at_cap.hits.len(), at_cap.truncated), + (MAX_IN_SESSION_HITS, false) + ); + + let mut over = vec![hit(1); MAX_IN_SESSION_HITS]; + over.push(hit(2)); + let over_cap = InSessionSearchResponse::from_hits( + "x".to_owned(), + over, + HashMap::from([(1, vec![[0, 1]]), (2, vec![[0, 1]])]), + ); + assert_eq!( + (over_cap.hits.len(), over_cap.truncated), + (MAX_IN_SESSION_HITS, true) + ); + assert_eq!( + over_cap + .texts + .iter() + .map(|text| text.id) + .collect::>(), + [1], + "a text only the dropped hit names is left out" + ); + } +} diff --git a/crates/abyss-backend/src/search/matching/mod.rs b/crates/abyss-backend/src/search/matching/mod.rs new file mode 100644 index 0000000..1b9e184 --- /dev/null +++ b/crates/abyss-backend/src/search/matching/mod.rs @@ -0,0 +1,15 @@ +//! Literal substring matching every storage backend shares: folding, the +//! pieces of an event, candidate expressions, occurrences, fragments, and +//! the order and bounds of results. +//! +//! The reference model is to fold every piece of every event and look for +//! the folded query in it. Backends may narrow the texts they verify, but a +//! hit is always a piece whose folded text holds the folded query, verified +//! here, so every backend returns the reference model's hits. + +pub mod fold; +pub mod fragment; +pub mod in_session; +pub mod pattern; +pub mod pieces; +pub mod ranking; diff --git a/crates/abyss-backend/src/search/matching/pattern.rs b/crates/abyss-backend/src/search/matching/pattern.rs new file mode 100644 index 0000000..2848380 --- /dev/null +++ b/crates/abyss-backend/src/search/matching/pattern.rs @@ -0,0 +1,230 @@ +//! A folded query: the trigram expression that narrows its candidate texts, +//! and its occurrences in a text. + +use std::{collections::HashSet, ops::Range}; + +use super::fold::fold; + +/// Characters in one trigram of the candidate index. +const TRIGRAM_CHARACTERS: usize = 3; + +/// A nonempty query folded for matching. +pub struct Pattern { + folded: String, + characters: usize, +} + +impl Pattern { + /// Folds a validated, nonempty query. + #[must_use] + pub fn new(query: &str) -> Self { + let folded = fold(query); + let characters = folded.chars().count(); + Self { folded, characters } + } + + /// An FTS5 expression over a trigram index of folded texts that every + /// text holding the query matches: each distinct trigram of the folded + /// query as its own quoted string, `"` written as `""`, joined with + /// `AND`. None when the query is shorter than a trigram; every text is + /// then a candidate. + /// + /// A trigram holding NUL is left out, because an FTS5 string ends there; + /// leaving a trigram out only widens the candidates. + #[must_use] + pub fn trigram_expression(&self) -> Option { + let characters = self.folded.chars().collect::>(); + let mut seen = HashSet::new(); + let terms = characters + .windows(TRIGRAM_CHARACTERS) + .filter(|trigram| !trigram.contains(&'\0')) + .map(|trigram| trigram.iter().collect::()) + .filter(|trigram| seen.insert(trigram.clone())) + .map(|trigram| format!("\"{}\"", trigram.replace('"', "\"\""))) + .collect::>(); + (!terms.is_empty()).then(|| terms.join(" AND ")) + } + + /// Whether the folded `text` holds the folded query. + #[must_use] + pub fn matches(&self, text: &str) -> bool { + fold(text).contains(self.folded.as_str()) + } + + /// Non-overlapping occurrences of the query in `text`, left to right, as + /// character ranges. Folding keeps the character count, so the ranges + /// found in the folded text are the same ranges of `text`. + #[must_use] + pub fn occurrences(&self, text: &str) -> Vec> { + let folded = fold(text); + let mut ranges = Vec::new(); + let mut scanned_bytes = 0_usize; + let mut scanned_characters = 0_usize; + for (start, _) in folded.match_indices(self.folded.as_str()) { + scanned_characters = scanned_characters.saturating_add( + folded + .get(scanned_bytes..start) + .map_or(0, |skipped| skipped.chars().count()), + ); + scanned_bytes = start; + ranges.push(scanned_characters..scanned_characters.saturating_add(self.characters)); + } + ranges + } +} + +/// Character ranges of `text`, sorted and non-overlapping, as half-open +/// ranges of its UTF-16 code units. +#[must_use] +pub fn utf16_ranges(text: &str, ranges: &[Range]) -> Vec<[usize; 2]> { + let offsets = character_offsets( + text, + ranges.iter().flat_map(|range| [range.start, range.end]), + char::len_utf16, + ); + let (pairs, _) = offsets.as_chunks::<2>(); + pairs.to_vec() +} + +/// The offset of each of `positions`, character positions in non-decreasing +/// order, counted in the units `width` gives each character of `text`. +pub fn character_offsets>( + text: &str, + positions: P, + width: fn(char) -> usize, +) -> Vec { + let mut characters = text.chars(); + let mut position = 0_usize; + let mut offset = 0_usize; + positions + .into_iter() + .map(|target| { + while position < target { + let Some(character) = characters.next() else { + break; + }; + offset = offset.saturating_add(width(character)); + position = position.saturating_add(1); + } + offset + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::{Pattern, character_offsets, utf16_ranges}; + + #[test] + fn trigram_expression_quotes_each_distinct_trigram() { + assert_eq!( + Pattern::new("Cargo").trigram_expression().as_deref(), + Some(r#""car" AND "arg" AND "rgo""#) + ); + assert_eq!( + Pattern::new("aaaa").trigram_expression().as_deref(), + Some(r#""aaa""#), + "repeated trigrams appear once" + ); + assert_eq!( + Pattern::new(r#"a"b"#).trigram_expression().as_deref(), + Some(r#""a""b""#), + "a double quote is written twice inside the string" + ); + assert_eq!( + Pattern::new("开始新任务").trigram_expression().as_deref(), + Some(r#""开始新" AND "始新任" AND "新任务""#) + ); + assert_eq!( + Pattern::new("a b").trigram_expression().as_deref(), + Some(r#""a b""#), + "spaces are characters of the query" + ); + } + + #[test] + fn short_queries_have_no_trigram_and_scan_every_text() { + assert_eq!(Pattern::new("搜").trigram_expression(), None); + assert_eq!(Pattern::new("搜索").trigram_expression(), None); + assert_eq!(Pattern::new("a\0").trigram_expression(), None); + assert_eq!( + Pattern::new("ab\0cd").trigram_expression(), + None, + "every trigram holds NUL" + ); + assert_eq!( + Pattern::new("abc\0d").trigram_expression().as_deref(), + Some(r#""abc""#) + ); + } + + #[test] + fn matches_after_folding_both_sides() { + let pattern = Pattern::new("ÄBC"); + assert!(pattern.matches("xxäbcxx")); + assert!(pattern.matches("ÄBC")); + assert!(!pattern.matches("abc")); + assert!(Pattern::new("straße").matches("STRAßE")); + assert!(!Pattern::new("strasse").matches("STRAßE")); + assert!(Pattern::new("cargo test").matches("run `CARGO TEST` now")); + assert!(!Pattern::new("cargo test").matches("cargo test")); + } + + #[test] + fn occurrences_do_not_overlap() { + assert_eq!(Pattern::new("aa").occurrences("aaaaa"), [0..2, 2..4]); + let once = Pattern::new("aba").occurrences("ababa"); + assert_eq!((once.len(), once.first()), (1, Some(&(0..3)))); + assert_eq!( + Pattern::new("x").occurrences("abc"), + Vec::>::new() + ); + // The Kelvin sign folds to an ASCII k, so byte positions shift while + // character positions do not. + assert_eq!(Pattern::new("k").occurrences("\u{212a}ok"), [0..1, 2..3]); + assert_eq!( + Pattern::new("搜索").occurrences("先搜索，再搜索"), + [1..3, 5..7] + ); + } + + #[test] + fn utf16_ranges_count_astral_characters_twice() { + let text = "😀 nimbus 😀NIMBUS"; + let ranges = Pattern::new("nimbus").occurrences(text); + assert_eq!(ranges, [2..8, 10..16]); + let utf16 = utf16_ranges(text, &ranges); + assert_eq!(utf16, [[3, 9], [12, 18]]); + let units = text.encode_utf16().collect::>(); + for [start, end] in utf16 { + assert_eq!( + String::from_utf16(&units[start..end]) + .expect("ranges fall on character boundaries") + .to_lowercase(), + "nimbus" + ); + } + assert_eq!( + utf16_ranges("İx", &Pattern::new("x").occurrences("İx")), + [[1, 2]] + ); + } + + #[test] + fn character_offsets_follow_the_unit_width() { + let text = "a会😀b"; + assert_eq!( + character_offsets(text, [0, 1, 2, 3, 4], char::len_utf8), + [0, 1, 4, 8, 9] + ); + assert_eq!( + character_offsets(text, [0, 1, 2, 3, 4], char::len_utf16), + [0, 1, 2, 4, 5] + ); + assert_eq!( + character_offsets(text, [9], char::len_utf8), + [9], + "past the end stops at it" + ); + } +} diff --git a/crates/abyss-backend/src/search/matching/pieces.rs b/crates/abyss-backend/src/search/matching/pieces.rs new file mode 100644 index 0000000..dc06f47 --- /dev/null +++ b/crates/abyss-backend/src/search/matching/pieces.rs @@ -0,0 +1,127 @@ +//! The searchable pieces of one event: the `text` and the `name` of each of +//! its conversation entries. + +use serde_json::Value; + +use crate::search::PieceField; + +impl PieceField { + /// The stored and serialized name of the field. + #[must_use] + pub const fn as_str(self) -> &'static str { + match self { + Self::Name => "name", + Self::Text => "text", + } + } + + /// Reads a stored field name. + #[must_use] + pub fn parse(value: &str) -> Option { + match value { + "name" => Some(Self::Name), + "text" => Some(Self::Text), + _ => None, + } + } +} + +/// One nonempty string value of one entry. +#[derive(Debug, PartialEq, Eq)] +pub struct Piece<'a> { + /// Position of the entry in `conversation.entries`. + pub entry_index: usize, + /// Which value of the entry it is. + pub field: PieceField, + /// The value as the event holds it. + pub text: &'a str, +} + +/// The pieces of the event's `conversation.entries`, in entry order and +/// `name` before `text` within an entry. Nothing outside the entries is read, +/// and metadata never is, so headers are never searched; an empty value is +/// not a piece, and entries are never joined. +#[must_use] +pub fn pieces(conversation: Option<&Value>) -> Vec> { + conversation + .and_then(|conversation| conversation.get("entries")) + .and_then(Value::as_array) + .into_iter() + .flatten() + .enumerate() + .flat_map(|(entry_index, entry)| { + [PieceField::Name, PieceField::Text] + .into_iter() + .filter_map(move |field| { + let text = entry.get(field.as_str())?.as_str()?; + (!text.is_empty()).then_some(Piece { + entry_index, + field, + text, + }) + }) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::{Piece, PieceField, pieces}; + + #[test] + fn reads_the_text_and_name_of_every_entry_and_nothing_else() { + let conversation = json!({ + "entries": [ + {"role": "system", "text": "You are an agent"}, + {"role": "reasoning", "text": ""}, + {"role": "tool_call", "text": "{\"command\":\"cargo test\"}", "call_id": "c1", "name": "Bash"}, + {"role": "tool_result", "text": "all tests passed", "call_id": "c1"}, + {"role": "unknown"}, + {"role": "user", "text": 3_i64, "name": ""} + ] + }); + + assert_eq!( + pieces(Some(&conversation)), + [ + Piece { + entry_index: 0, + field: PieceField::Text, + text: "You are an agent" + }, + Piece { + entry_index: 2, + field: PieceField::Name, + text: "Bash" + }, + Piece { + entry_index: 2, + field: PieceField::Text, + text: "{\"command\":\"cargo test\"}" + }, + Piece { + entry_index: 3, + field: PieceField::Text, + text: "all tests passed" + }, + ] + ); + assert_eq!(pieces(None), []); + assert_eq!(pieces(Some(&json!({"entries": "none"}))), []); + } + + #[test] + fn field_names_round_trip_and_order_name_first() { + for field in [PieceField::Name, PieceField::Text] { + assert_eq!(PieceField::parse(field.as_str()), Some(field)); + assert_eq!( + serde_json::to_value(field).expect("serializes"), + json!(field.as_str()) + ); + } + assert_eq!(PieceField::parse("role"), None); + assert!(PieceField::Name < PieceField::Text); + } +} diff --git a/crates/abyss-backend/src/search/matching/ranking.rs b/crates/abyss-backend/src/search/matching/ranking.rs new file mode 100644 index 0000000..559369b --- /dev/null +++ b/crates/abyss-backend/src/search/matching/ranking.rs @@ -0,0 +1,628 @@ +//! Which hits are new, how matching sessions and their events are ordered, +//! and which events and texts a page of session search results shows. + +use std::collections::{BTreeSet, HashMap, HashSet}; + +use chrono::{DateTime, Utc}; +use uuid::Uuid; + +use crate::{ + search::{ + PieceField, SearchFragment, SessionSearchMatch, SessionSearchResponse, SessionSearchResult, + ValidatedSearchQuery, + }, + usage::UsageEventType, +}; + +/// Matching events a session result shows. +pub const MATCHES_PER_SESSION: usize = 3; +/// Fragments a matching event shows. +pub const FRAGMENTS_PER_MATCH: usize = 2; + +/// Whether a hit is outside the event's inherited context: every hit of a +/// response, and the hits of a request from its `new_from` on, or from its +/// first entry without one. +#[must_use] +pub fn is_new(event_type: UsageEventType, new_from: Option, entry_index: u32) -> bool { + event_type == UsageEventType::Response || entry_index >= new_from.unwrap_or_default() +} + +/// One hit of a session search, with the event facts ranking reads. +#[derive(Clone, Debug)] +pub struct SearchHit { + /// Session of the event. + pub session_pk: Uuid, + /// Agent-native identifier of that session. + pub session_id: String, + /// Event holding the piece. + pub event_pk: Uuid, + /// Observation time of the event. + pub observed_at: DateTime, + /// Position of the piece's entry. + pub entry_index: u32, + /// Which value of the entry the piece is. + pub field: PieceField, + /// Text of the piece. + pub text_id: i64, + /// Whether the hit is new, by [`is_new`]. + pub new: bool, +} + +/// One page of matching sessions, before the event store hydrates it. +#[derive(Debug, PartialEq, Eq)] +pub struct RankedPage { + /// Sessions with at least one matching event. + pub total_sessions: u64, + /// The page's sessions, most recently matching first. + pub sessions: Vec, +} + +/// One matching session of a page. +#[derive(Debug, PartialEq, Eq)] +pub struct RankedSession { + /// Session primary key. + pub session_pk: Uuid, + /// Matching events. + pub match_count: u64, + /// Matching events with a new hit. + pub new_match_count: u64, + /// The matching events shown, at most [`MATCHES_PER_SESSION`]. + pub matches: Vec, +} + +/// One matching event shown for a session. +#[derive(Debug, PartialEq, Eq)] +pub struct RankedMatch { + /// The event. + pub event_pk: Uuid, + /// Observation time of the event. + pub observed_at: DateTime, + /// Whether the event has a new hit. + pub new: bool, + /// Texts its fragments come from, at most [`FRAGMENTS_PER_MATCH`]. + pub fragment_texts: Vec, +} + +/// Event columns a shown match reports. +pub struct MatchedEvent { + /// Request or response side. + pub event_type: String, + /// LLM provider as reported. + pub llm_provider: String, + /// Model as reported. + pub llm_model: String, +} + +/// Session and device columns a session result reports. +pub struct SessionDetails { + /// Agent-native session identifier. + pub session_id: String, + /// Agent name as reported by the collector. + pub agent_name: String, + /// Most recently observed Agent version. + pub agent_version: Option, + /// Authoritative device host name. + pub host_name: String, + /// Authoritative device platform. + pub platform: String, + /// Earliest event observation in the session. + pub started_at: DateTime, + /// Latest event observation in the session. + pub ended_at: Option>, +} + +/// The hits of one event. +struct EventHits { + session_pk: Uuid, + session_id: String, + observed_at: DateTime, + hits: Vec, +} + +#[derive(Clone, Copy)] +struct PieceHit { + new: bool, + entry_index: u32, + field: PieceField, + text_id: i64, +} + +impl EventHits { + fn has_new(&self) -> bool { + self.hits.iter().any(|hit| hit.new) + } + + /// Distinct texts of the hits: new hits first, then those in the + /// inherited context, each by entry and `name` before `text`. + fn fragment_texts(&self) -> Vec { + let mut hits = self.hits.clone(); + hits.sort_by_key(|hit| (!hit.new, hit.entry_index, hit.field)); + let mut texts = Vec::with_capacity(FRAGMENTS_PER_MATCH); + for hit in hits { + if texts.len() == FRAGMENTS_PER_MATCH { + break; + } + if !texts.contains(&hit.text_id) { + texts.push(hit.text_id); + } + } + texts + } +} + +/// Groups hits into matching events and sessions and keeps one page. +/// +/// Sessions order by their latest matching event, newest first, then by +/// session id, then by session key. A session shows its matching events with a new hit, earliest +/// first, then fills up with its earliest events whose hits are all in the +/// inherited context; events at the same time order by key. +#[must_use] +pub fn rank_sessions(hits: Vec, offset: usize, page_size: usize) -> RankedPage { + let mut events = HashMap::::new(); + for hit in hits { + let SearchHit { + session_pk, + session_id, + event_pk, + observed_at, + entry_index, + field, + text_id, + new, + } = hit; + events + .entry(event_pk) + .or_insert_with(|| EventHits { + session_pk, + session_id, + observed_at, + hits: Vec::new(), + }) + .hits + .push(PieceHit { + new, + entry_index, + field, + text_id, + }); + } + let mut sessions = HashMap::>::new(); + for (event_pk, event) in events { + sessions + .entry(event.session_pk) + .or_default() + .push((event_pk, event)); + } + let mut sessions = sessions + .into_iter() + .map(|(session_pk, events)| { + let latest = events.iter().map(|(_, event)| event.observed_at).max(); + let session_id = events + .first() + .map(|(_, event)| event.session_id.clone()) + .unwrap_or_default(); + ((latest, session_id, session_pk), events) + }) + .collect::>(); + sessions.sort_by( + |((left_latest, left_id, left_pk), _), ((right_latest, right_id, right_pk), _)| { + right_latest + .cmp(left_latest) + .then_with(|| left_id.cmp(right_id)) + .then_with(|| left_pk.cmp(right_pk)) + }, + ); + RankedPage { + total_sessions: count(sessions.len()), + sessions: sessions + .into_iter() + .skip(offset) + .take(page_size) + .map(|((_, _, session_pk), events)| rank_session(session_pk, events)) + .collect(), + } +} + +fn rank_session(session_pk: Uuid, mut events: Vec<(Uuid, EventHits)>) -> RankedSession { + events.sort_by_key(|(event_pk, event)| (event.observed_at, *event_pk)); + let new_match_count = events.iter().filter(|(_, event)| event.has_new()).count(); + let matches = events + .iter() + .filter(|(_, event)| event.has_new()) + .chain(events.iter().filter(|(_, event)| !event.has_new())) + .take(MATCHES_PER_SESSION) + .map(|(event_pk, event)| RankedMatch { + event_pk: *event_pk, + observed_at: event.observed_at, + new: event.has_new(), + fragment_texts: event.fragment_texts(), + }) + .collect(); + RankedSession { + session_pk, + match_count: count(events.len()), + new_match_count: count(new_match_count), + matches, + } +} + +fn count(value: usize) -> u64 { + u64::try_from(value).unwrap_or(u64::MAX) +} + +impl RankedPage { + /// Session keys of the page, in order. + #[must_use] + pub fn session_pks(&self) -> Vec { + self.sessions + .iter() + .map(|session| session.session_pk) + .collect() + } + + /// Every shown event. + #[must_use] + pub fn event_pks(&self) -> Vec { + self.matches().map(|matched| matched.event_pk).collect() + } + + /// Every text a shown fragment comes from, once. + #[must_use] + pub fn fragment_texts(&self) -> Vec { + self.matches() + .flat_map(|matched| matched.fragment_texts.iter().copied()) + .collect::>() + .into_iter() + .collect() + } + + fn matches(&self) -> impl Iterator { + self.sessions + .iter() + .flat_map(|session| session.matches.iter()) + } + + /// Hydrates the page with its sessions' details, its events' columns, + /// and one fragment per text. A session or event the store no longer + /// holds is left out. A session lists the providers and models of the + /// matches it shows. + #[must_use] + pub fn into_response( + self, + query: ValidatedSearchQuery, + mut details: HashMap, + events: &HashMap, + fragments: &HashMap, + ) -> SessionSearchResponse { + let items = self + .sessions + .into_iter() + .filter_map(|ranked| { + let details = details.remove(&ranked.session_pk)?; + let matches = ranked + .matches + .into_iter() + .filter_map(|matched| { + let event = events.get(&matched.event_pk)?; + Some(SessionSearchMatch { + event_pk: matched.event_pk, + location: None, + event_type: event.event_type.clone(), + llm_provider: event.llm_provider.clone(), + llm_model: event.llm_model.clone(), + observed_at: matched.observed_at, + new: matched.new, + fragments: matched + .fragment_texts + .iter() + .filter_map(|text_id| fragments.get(text_id).cloned()) + .collect(), + }) + }) + .collect::>(); + Some(SessionSearchResult { + session_pk: ranked.session_pk, + title: None, + session_id: details.session_id, + agent_name: details.agent_name, + agent_version: details.agent_version, + host_name: details.host_name, + platform: details.platform, + started_at: details.started_at, + ended_at: details.ended_at, + providers: sorted_distinct(matches.iter().map(|m| &m.llm_provider)), + models: sorted_distinct(matches.iter().map(|m| &m.llm_model)), + match_count: ranked.match_count, + new_match_count: ranked.new_match_count, + matches, + }) + }) + .collect(); + let consumed = u64::from(query.offset).saturating_add(u64::from(query.page_size)); + SessionSearchResponse { + query: query.text, + total_sessions: self.total_sessions, + page: query.page, + page_size: query.page_size, + has_more: consumed < self.total_sessions, + items, + } + } +} + +fn sorted_distinct<'a>(values: impl Iterator) -> Vec { + let mut values = values + .cloned() + .collect::>() + .into_iter() + .collect::>(); + values.sort(); + values +} + +#[cfg(test)] +mod tests { + use chrono::{DateTime, TimeZone as _, Utc}; + use uuid::Uuid; + + use std::collections::HashMap; + + use super::{ + MatchedEvent, RankedMatch, RankedPage, SearchHit, SessionDetails, is_new, rank_sessions, + }; + use crate::{ + search::{ + PieceField, SearchFragment, SearchFragmentSegment, SessionSearchQuery, + ValidatedSearchQuery, + }, + usage::UsageEventType, + }; + + fn at(second: u32) -> DateTime { + Utc.with_ymd_and_hms(2026, 10, 18, 0, 0, second) + .single() + .expect("valid time") + } + + /// A session key standing for a session name of at most 16 bytes. + fn pk(session: &str) -> Uuid { + Uuid::from_u128( + session + .bytes() + .fold(0_u128, |key, byte| (key << 8_u32) | u128::from(byte)), + ) + } + + fn names(page: &RankedPage, sessions: &[&str]) -> Vec { + page.session_pks() + .iter() + .map(|session_pk| { + sessions + .iter() + .find(|session| pk(session) == *session_pk) + .map(|session| (*session).to_owned()) + .unwrap_or_default() + }) + .collect() + } + + fn hit( + session: &str, + event: u128, + second: u32, + entry: u32, + text_id: i64, + new: bool, + ) -> SearchHit { + SearchHit { + session_pk: pk(session), + session_id: session.to_owned(), + event_pk: Uuid::from_u128(event), + observed_at: at(second), + entry_index: entry, + field: PieceField::Text, + text_id, + new, + } + } + + #[test] + fn hits_are_new_in_responses_and_from_a_requests_new_from() { + assert!(is_new(UsageEventType::Response, None, 0)); + assert!( + is_new(UsageEventType::Response, Some(9), 0), + "a response has no context" + ); + assert!(!is_new(UsageEventType::Request, Some(4), 3)); + assert!(is_new(UsageEventType::Request, Some(4), 4)); + assert!(is_new(UsageEventType::Request, Some(0), 0)); + assert!( + is_new(UsageEventType::Request, None, 0), + "a request without new_from is all new" + ); + } + + #[test] + fn sessions_order_by_their_latest_matching_event_then_by_id() { + let hits = vec![ + hit("old", 1, 1, 0, 1, true), + hit("tie-b", 2, 5, 0, 1, true), + hit("tie-a", 3, 5, 0, 1, false), + hit("old", 4, 2, 0, 1, true), + hit("new", 5, 9, 0, 1, false), + hit("new", 6, 0, 0, 1, true), + ]; + let all = ["old", "tie-a", "tie-b", "new"]; + let page = rank_sessions(hits.clone(), 0, 10); + assert_eq!(page.total_sessions, 4); + assert_eq!(names(&page, &all), ["new", "tie-a", "tie-b", "old"]); + + let second = rank_sessions(hits.clone(), 1, 2); + assert_eq!(second.total_sessions, 4); + assert_eq!(names(&second, &all), ["tie-a", "tie-b"]); + assert_eq!(rank_sessions(hits, 4, 2).sessions, []); + } + + #[test] + fn sessions_sharing_an_id_stay_apart_and_order_by_key() { + let mut later = hit("b", 1, 5, 0, 1, true); + later.session_id = "same".to_owned(); + let mut earlier = hit("a", 2, 5, 0, 1, true); + earlier.session_id = "same".to_owned(); + let page = rank_sessions(vec![later, earlier], 0, 10); + assert_eq!(page.total_sessions, 2); + assert_eq!(names(&page, &["a", "b"]), ["a", "b"]); + } + + #[test] + fn a_session_shows_new_matches_first_then_its_earliest_context_matches() { + let hits = vec![ + hit("s", 1, 1, 0, 10, false), + hit("s", 2, 2, 5, 11, true), + hit("s", 2, 2, 0, 10, false), + hit("s", 3, 3, 0, 10, false), + hit("s", 4, 4, 1, 12, true), + hit("s", 5, 4, 0, 10, false), + ]; + let page = rank_sessions(hits, 0, 10); + let session = &page.sessions[0]; + assert_eq!((session.match_count, session.new_match_count), (5, 2)); + let shown = session + .matches + .iter() + .map(|matched| (matched.event_pk.as_u128(), matched.new)) + .collect::>(); + assert_eq!(shown, [(2, true), (4, true), (1, false)]); + + let all_context = rank_sessions( + vec![ + hit("c", 9, 3, 0, 1, false), + hit("c", 8, 3, 0, 1, false), + hit("c", 7, 1, 0, 1, false), + hit("c", 6, 2, 0, 1, false), + ], + 0, + 1, + ); + let shown = all_context.sessions[0] + .matches + .iter() + .map(|matched| matched.event_pk.as_u128()) + .collect::>(); + assert_eq!(shown, [7, 6, 8], "equal times order by event key"); + } + + #[test] + fn fragments_come_from_distinct_texts_new_hits_first() { + let mut name = hit("s", 1, 1, 3, 30, true); + name.field = PieceField::Name; + let hits = vec![ + hit("s", 1, 1, 0, 10, false), + hit("s", 1, 1, 3, 31, true), + name, + hit("s", 1, 1, 4, 30, true), + hit("s", 2, 2, 0, 10, false), + hit("s", 2, 2, 1, 10, false), + hit("s", 2, 2, 2, 20, false), + ]; + let page = rank_sessions(hits, 0, 10); + let texts = page.sessions[0] + .matches + .iter() + .map(|matched: &RankedMatch| matched.fragment_texts.clone()) + .collect::>(); + assert_eq!(texts, [vec![30, 31], vec![10, 20]]); + assert_eq!(page.fragment_texts(), [10, 20, 30, 31]); + assert_eq!(page.event_pks(), [Uuid::from_u128(1), Uuid::from_u128(2)]); + } + + fn details(session_id: &str) -> SessionDetails { + SessionDetails { + session_id: session_id.to_owned(), + agent_name: "codex".to_owned(), + agent_version: None, + host_name: "alice-mbp".to_owned(), + platform: "macos".to_owned(), + started_at: at(0), + ended_at: Some(at(9)), + } + } + + fn query(page: u32, page_size: u32) -> ValidatedSearchQuery { + SessionSearchQuery { + q: "needle".to_owned(), + from: None, + to: None, + agent_name: Vec::new(), + llm_provider: Vec::new(), + llm_model: Vec::new(), + event_type: None, + page: Some(page), + page_size: Some(page_size), + } + .validate() + .expect("query is valid") + } + + #[test] + fn a_page_is_hydrated_from_the_store_and_drops_what_it_lacks() { + let hits = vec![ + hit("kept", 1, 5, 0, 10, true), + hit("kept", 2, 6, 0, 11, false), + hit("gone", 3, 7, 0, 10, true), + hit("later", 4, 1, 0, 10, true), + ]; + let page = rank_sessions(hits, 0, 2); + assert_eq!(names(&page, &["gone", "kept"]), ["gone", "kept"]); + let fragment = |text: &str| SearchFragment { + segments: vec![SearchFragmentSegment { + text: text.to_owned(), + highlighted: true, + }], + }; + let events = HashMap::from([( + uuid::Uuid::from_u128(1), + MatchedEvent { + event_type: "request".to_owned(), + llm_provider: "openai".to_owned(), + llm_model: "gpt-5".to_owned(), + }, + )]); + let fragments = HashMap::from([(10, fragment("needle"))]); + let response = page.into_response( + query(1, 2), + HashMap::from([(pk("kept"), details("kept"))]), + &events, + &fragments, + ); + + assert_eq!(response.total_sessions, 3); + assert!(response.has_more, "a third session follows the page"); + assert_eq!( + response.items.len(), + 1, + "a session without details is left out" + ); + let kept = &response.items[0]; + assert_eq!((kept.match_count, kept.new_match_count), (2, 1)); + assert_eq!( + kept.matches + .iter() + .map(|matched| ( + matched.event_pk.as_u128(), + matched.new, + matched.fragments.len() + )) + .collect::>(), + [(1, true, 1)], + "an event the store lacks is left out" + ); + assert_eq!(kept.matches[0].llm_model, "gpt-5"); + assert_eq!( + (kept.session_id.as_str(), kept.models.as_slice()), + ("kept", ["gpt-5".to_owned()].as_slice()), + "a session lists the models of the matches it shows" + ); + assert_eq!(kept.matches[0].fragments[0], fragment("needle")); + } +} diff --git a/crates/abyss-backend/src/search/mod.rs b/crates/abyss-backend/src/search/mod.rs index fa04921..6e9ad54 100644 --- a/crates/abyss-backend/src/search/mod.rs +++ b/crates/abyss-backend/src/search/mod.rs @@ -1,21 +1,30 @@ -//! Backend-independent full-text search contracts and safe source projection. +//! Backend-independent search contracts and the literal text matching both +//! searches share. //! -//! Concrete search engines live inside their selected storage backend. This -//! module defines only the query, response, highlight, and bounded projection -//! types shared by those implementations and the HTTP layer. +//! A search reads the pieces of each event's conversation entries, each +//! entry's `text` and its `name`, and matches the query as one literal +//! substring after per-character case folding. Storage backends find the +//! candidate texts and expand matching texts into events; folding, matching, +//! occurrence ranges, fragments, and result ordering live in [`matching`], so +//! every backend returns the same hits. -pub mod projection; +#[cfg_attr( + all(feature = "postgres-es", not(test)), + expect( + dead_code, + reason = "the PostgreSQL profile keeps its Elasticsearch search until its replacement is decided" + ) +)] +pub mod matching; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use uuid::Uuid; -use crate::error::AppError; - -/// Sentinel inserted before text matched by either search implementation. -pub const HIGHLIGHT_START: &str = "[[[ABYSS_SEARCH_HIGHLIGHT_START]]]"; -/// Sentinel inserted after text matched by either search implementation. -pub const HIGHLIGHT_END: &str = "[[[ABYSS_SEARCH_HIGHLIGHT_END]]]"; +use crate::{ + error::AppError, + usage::streams::{EntryLabel, EventLocation}, +}; const DEFAULT_PAGE_SIZE: u32 = 20; const MAX_PAGE_SIZE: u32 = 50; @@ -24,9 +33,9 @@ const MAX_FILTER_CHARACTERS: usize = 256; const MAX_RESULT_WINDOW: u32 = 10_000; #[derive(Debug, Deserialize)] -/// Query-string contract for session full-text search. +/// Query-string contract for session search. pub struct SessionSearchQuery { - /// Required full-text query. + /// Required query, matched as one literal substring. pub q: String, /// Inclusive lower observation-time bound. pub from: Option>, @@ -51,7 +60,7 @@ pub struct SessionSearchQuery { impl SessionSearchQuery { pub(crate) fn validate(self) -> Result { - let text = normalized_required(&self.q, "q", MAX_QUERY_CHARACTERS)?; + let text = normalized_query(&self.q)?; if let (Some(from), Some(to)) = (self.from, self.to) && from >= to { @@ -109,7 +118,7 @@ impl SessionSearchQuery { /// Normalized query safe to translate into a backend-specific search request. pub struct ValidatedSearchQuery { - /// Trimmed full-text query. + /// Trimmed query. pub text: String, /// Inclusive lower observation-time bound. pub from: Option>, @@ -134,9 +143,9 @@ pub struct ValidatedSearchQuery { #[derive(Serialize)] /// Paginated session search response. pub struct SessionSearchResponse { - /// Normalized full-text query. + /// Normalized query. pub query: String, - /// Distinct matching-session count reported by the selected backend. + /// Sessions with at least one matching event. pub total_sessions: u64, /// One-based current page. pub page: u32, @@ -149,12 +158,14 @@ pub struct SessionSearchResponse { } #[derive(Serialize)] -/// One hydrated session with its strongest matching events. +/// One hydrated session with up to three of its matching events. pub struct SessionSearchResult { /// Backend session primary key. pub session_pk: Uuid, /// Agent-native session identifier. pub session_id: String, + /// The same earliest-request label returned by the session list and header. + pub title: Option, /// Agent name as reported by the collector. pub agent_name: String, /// Most recently observed Agent version. @@ -171,9 +182,12 @@ pub struct SessionSearchResult { pub providers: Vec, /// Sorted models represented by the returned matching events. pub models: Vec, - /// Total matching events in the session. + /// Matching events in the session. pub match_count: u64, - /// Bounded strongest matching events with fragments. + /// Matching events with at least one new hit. + pub new_match_count: u64, + /// Matching events with fragments: those with a new hit first, then + /// those whose hits are all in the inherited context. pub matches: Vec, } @@ -182,10 +196,9 @@ pub struct SessionSearchResult { pub struct SessionSearchMatch { /// Backend event primary key. pub event_pk: Uuid, - /// Backend turn primary key. - pub turn_pk: Uuid, - /// Normalized turn number. - pub turn_index: i32, + /// Current position in this session's stream view, when available. + #[serde(skip_serializing_if = "Option::is_none")] + pub location: Option, /// Request or response side. pub event_type: String, /// LLM provider as reported by the collector. @@ -194,62 +207,98 @@ pub struct SessionSearchMatch { pub llm_model: String, /// Collector observation time. pub observed_at: DateTime, - /// Safely segmented highlighted snippets. + /// Whether the event has a hit outside its inherited context. + pub new: bool, + /// Highlighted fragments of the event's matching pieces. pub fragments: Vec, } -#[derive(Serialize)] +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] /// One search highlight fragment split into plain and matching text. pub struct SearchFragment { /// Ordered text segments suitable for structured UI rendering. pub segments: Vec, } -impl SearchFragment { - pub(crate) fn parse(fragment: &str) -> Self { - let mut segments = Vec::new(); - let mut remainder = fragment; - while let Some(start) = remainder.find(HIGHLIGHT_START) { - let (plain, tagged) = remainder.split_at(start); - push_fragment_segment(&mut segments, plain, false); - let highlighted = tagged - .strip_prefix(HIGHLIGHT_START) - .expect("the marker position came from find"); - // Treat malformed or truncated marker pairs as plain text. The API - // never emits raw HTML and therefore does not trust ES fragments. - let Some(end) = highlighted.find(HIGHLIGHT_END) else { - push_fragment_segment(&mut segments, tagged, false); - remainder = ""; - break; - }; - let (matched, after_match) = highlighted.split_at(end); - push_fragment_segment(&mut segments, matched, true); - remainder = after_match - .strip_prefix(HIGHLIGHT_END) - .expect("the marker position came from find"); - } - push_fragment_segment(&mut segments, remainder, false); - Self { segments } - } -} - -#[derive(Serialize)] +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] /// Plain or highlighted portion of a search fragment. pub struct SearchFragmentSegment { - /// Fragment text with internal marker tokens removed by parsing. + /// Fragment text. pub text: String, - /// Whether this segment matched the full-text query. + /// Whether this segment is an occurrence of the query. pub highlighted: bool, } -fn push_fragment_segment(segments: &mut Vec, text: &str, highlighted: bool) { - if text.is_empty() { - return; +#[derive(Debug, Deserialize)] +/// Query-string contract for search within one session. +pub struct InSessionSearchQuery { + /// Required query, matched as one literal substring. + pub q: String, +} + +impl InSessionSearchQuery { + /// The trimmed query. + pub(crate) fn validate(&self) -> Result { + normalized_query(&self.q) } - segments.push(SearchFragmentSegment { - text: text.to_owned(), - highlighted, - }); +} + +#[derive(Debug, Serialize)] +/// Every hit of the query in one session, with the occurrences in each +/// matching text. +pub struct InSessionSearchResponse { + /// Normalized query. + pub query: String, + /// The texts the hits name, by id, with their occurrences. + pub texts: Vec, + /// Hits in observation order, then by event, entry, and field, `name` + /// before `text`. + pub hits: Vec, + /// Whether hits beyond the returned ones were left out. + pub truncated: bool, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +/// One matching text and the occurrences of the query in it. +pub struct MatchedText { + /// Text id the hits name. + pub id: i64, + /// Non-overlapping occurrences, left to right, as half-open ranges of the + /// text's UTF-16 code units. + pub ranges: Vec<[usize; 2]>, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +/// One piece of one event whose text holds the query. +pub struct InSessionHit { + /// Event the piece belongs to. + pub event_pk: Uuid, + /// Position of the entry in the event's `conversation.entries`. + pub entry_index: u32, + /// Which value of the entry the piece is. + pub field: PieceField, + /// Whether the hit is outside the event's inherited context. + pub new: bool, + /// Text of the piece, in [`InSessionSearchResponse::texts`]. + pub text_id: i64, + /// Head of the segment holding the event's exchange, as the session's + /// streams give it; null when the event is in no exchange. + pub segment: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize)] +#[serde(rename_all = "lowercase")] +/// Which value of an entry a piece is. Within an entry, `name` orders +/// before `text`. +pub enum PieceField { + /// The entry's `name`, the tool name of a tool call. + Name, + /// The entry's `text`. + Text, +} + +fn normalized_query(value: &str) -> Result { + normalized_required(value, "q", MAX_QUERY_CHARACTERS) } fn normalized_required(value: &str, field: &str, maximum: usize) -> Result { @@ -284,7 +333,7 @@ fn normalized_list(values: &[String], field: &str) -> Result, AppErr #[cfg(test)] mod tests { - use super::{HIGHLIGHT_END, HIGHLIGHT_START, SearchFragment, SessionSearchQuery}; + use super::{InSessionSearchQuery, SessionSearchQuery}; #[test] fn rejects_empty_and_invalid_pagination_queries() { @@ -327,14 +376,22 @@ mod tests { } #[test] - fn parses_highlights_without_rendering_html() { - let fragment = SearchFragment::parse(&format!( - "before {HIGHLIGHT_START}timeout{HIGHLIGHT_END} after" - )); - - assert_eq!(fragment.segments.len(), 3); - assert_eq!(fragment.segments[1].text, "timeout"); - assert!(fragment.segments[1].highlighted); + fn a_query_is_trimmed_and_bounded_by_characters() { + let trimmed = query(" cargo test \n").validate().expect("query is valid"); + assert_eq!( + trimmed.text, "cargo test", + "inner spaces stay part of the query" + ); + assert!(query(&"会".repeat(256)).validate().is_ok()); + assert!(query(&"会".repeat(257)).validate().is_err()); + + let in_session = |q: &str| InSessionSearchQuery { q: q.to_owned() }.validate(); + assert_eq!( + in_session(" STATUS.md ").expect("query is valid"), + "STATUS.md" + ); + assert!(in_session(" ").is_err()); + assert!(in_session(&"a".repeat(257)).is_err()); } fn query(q: &str) -> SessionSearchQuery { diff --git a/crates/abyss-backend/src/search/projection.rs b/crates/abyss-backend/src/search/projection.rs deleted file mode 100644 index c4e4733..0000000 --- a/crates/abyss-backend/src/search/projection.rs +++ /dev/null @@ -1,97 +0,0 @@ -//! Bounded, allowlisted extraction shared by every full-text implementation. - -use serde_json::Value; - -use super::{HIGHLIGHT_END, HIGHLIGHT_START}; - -pub const MAX_CONTENT_CHARACTERS: usize = 1_000_000; -const MAX_METADATA_FIELD_CHARACTERS: usize = 100_000; -const MAX_EXTRACTED_VALUES: usize = 64; - -/// Searchable values derived from one source usage event. -pub struct SearchProjection { - pub content: Option, - pub tool_names: Vec, - pub tool_content: Vec, -} - -impl SearchProjection { - pub fn from_source(text: Option, metadata: &Value) -> Self { - let extracted = SearchableMetadata::extract(metadata); - Self { - content: text - .map(|value| sanitize_search_text(&value)) - .map(|value| truncate_characters(value, MAX_CONTENT_CHARACTERS)), - tool_names: extracted.tool_names, - tool_content: extracted.tool_content, - } - } -} - -struct SearchableMetadata { - tool_names: Vec, - tool_content: Vec, -} - -impl SearchableMetadata { - fn extract(metadata: &Value) -> Self { - let mut extracted = Self { - tool_names: Vec::new(), - tool_content: Vec::new(), - }; - - // Ignore every other top-level key so headers, credentials, and image - // data never enter a derived full-text index. - let Some(segments) = metadata.get("content_segments").and_then(Value::as_array) else { - return extracted; - }; - for segment in segments.iter().take(MAX_EXTRACTED_VALUES) { - let Some(object) = segment.as_object() else { - continue; - }; - match object.get("type").and_then(Value::as_str) { - Some("tool_call") => { - if let Some(name) = object.get("name").and_then(Value::as_str) { - push_bounded(&mut extracted.tool_names, name); - } - if let Some(input) = object.get("input").and_then(Value::as_str) { - push_bounded(&mut extracted.tool_content, input); - } - } - Some("tool_result") => { - if let Some(output) = object.get("output").and_then(Value::as_str) { - push_bounded(&mut extracted.tool_content, output); - } - } - _ => {} - } - } - extracted - } -} - -fn push_bounded(values: &mut Vec, value: &str) { - let value = sanitize_search_text(value); - let value = value.trim(); - if value.is_empty() || values.len() >= MAX_EXTRACTED_VALUES { - return; - } - let value = truncate_characters(value.to_owned(), MAX_METADATA_FIELD_CHARACTERS); - if !values.contains(&value) { - values.push(value); - } -} - -pub fn sanitize_search_text(value: &str) -> String { - value - .replace(HIGHLIGHT_START, "") - .replace(HIGHLIGHT_END, "") -} - -fn truncate_characters(mut value: String, maximum: usize) -> String { - let Some((byte_index, _character)) = value.char_indices().nth(maximum) else { - return value; - }; - value.truncate(byte_index); - value -} diff --git a/crates/abyss-backend/src/storage/event_locations.rs b/crates/abyss-backend/src/storage/event_locations.rs new file mode 100644 index 0000000..3aaa758 --- /dev/null +++ b/crates/abyss-backend/src/storage/event_locations.rs @@ -0,0 +1,51 @@ +//! Attaches derived positions to a response page using one structure lookup +//! per session. The storage caller owns the read snapshot and owner scope. + +use std::{collections::HashMap, sync::Arc}; + +use uuid::Uuid; + +use crate::{ + error::AppError, + search::SessionSearchResponse, + usage::{UsageEventResponse, streams::SessionStructure}, +}; + +pub(super) fn locate_events( + events: &mut [UsageEventResponse], + mut load: F, +) -> Result<(), AppError> +where + F: FnMut(Uuid) -> Result, AppError>, +{ + let mut structures = HashMap::new(); + for event in events { + let structure = match structures.entry(event.session_pk) { + std::collections::hash_map::Entry::Occupied(entry) => entry.into_mut(), + std::collections::hash_map::Entry::Vacant(entry) => { + entry.insert(load(event.session_pk)?) + } + }; + event.location = structure.event_position(event.id); + } + Ok(()) +} + +pub(super) fn locate_search_matches( + response: &mut SessionSearchResponse, + mut load: F, +) -> Result<(), AppError> +where + F: FnMut(Uuid) -> Result, AppError>, +{ + for session in &mut response.items { + if session.matches.is_empty() { + continue; + } + let structure = load(session.session_pk)?; + for hit in &mut session.matches { + hit.location = structure.event_position(hit.event_pk); + } + } + Ok(()) +} diff --git a/crates/abyss-backend/src/storage/mod.rs b/crates/abyss-backend/src/storage/mod.rs index 61f7dd3..5d1d8a7 100644 --- a/crates/abyss-backend/src/storage/mod.rs +++ b/crates/abyss-backend/src/storage/mod.rs @@ -1,4 +1,4 @@ -//! Compile-time-selected persistence and full-text search compatibility boundary. +//! Compile-time-selected persistence and search compatibility boundary. //! //! HTTP handlers depend only on [`StorageBackend`]. Exactly one concrete //! implementation is compiled into a binary, while dynamic dispatch keeps @@ -6,11 +6,12 @@ //! of the API layer. mod comparison_buckets; - +mod event_locations; #[cfg(feature = "postgres-es")] mod postgres; #[cfg(feature = "sqlite-fts")] mod sqlite; +mod stream_cache; use std::{future::Future, pin::Pin, sync::Arc}; @@ -19,14 +20,19 @@ use uuid::Uuid; use crate::{ config::Config, error::AppError, - search::{SessionSearchQuery, SessionSearchResponse}, + search::{ + InSessionSearchQuery, InSessionSearchResponse, SessionSearchQuery, SessionSearchResponse, + }, usage::{ - IngestEventsRequest, IngestEventsResponse, RawEventsQuery, RawEventsResponse, - SessionOutlineResponse, SessionTimelineResponse, SummaryQuery, SummaryResponse, - UsageEventResponse, + IngestEventsRequest, IngestEventsResponse, RawEventsQuery, RawEventsResponse, SummaryQuery, + SummaryResponse, UsageEventResponse, attachments::StoredImageAttachment, comparison::{ComparisonQuery, ComparisonResponse}, + preview::{ConversationPreview, Window}, + segments::{SegmentResponse, ValidatedSegmentQuery}, + session_header::SessionHeaderResponse, session_list::{SessionListResponse, ValidatedSessionListQuery}, + streams::{EventLocationResponse, SessionStreamsResponse, SessionTurnsResponse}, }, }; @@ -61,6 +67,13 @@ pub trait StorageBackend: Send + Sync { event_pk: Uuid, ) -> BoxedFuture<'_, Result>; + /// Resolves a stable event identifier to its current session segment and exchange. + fn event_location( + &self, + user_id: Uuid, + event_pk: Uuid, + ) -> BoxedFuture<'_, Result>; + /// Aggregates owner-scoped event and token usage. fn usage_summary( &self, @@ -76,19 +89,48 @@ pub trait StorageBackend: Send + Sync { query: ComparisonQuery, ) -> BoxedFuture<'_, Result>; - /// Loads one owner-scoped session timeline. - fn session_timeline( + /// The header of an owner-scoped session; not found when the owner has + /// no event in it. + fn session_header( + &self, + user_id: Uuid, + session_pk: Uuid, + ) -> BoxedFuture<'_, Result>; + + /// Derives the segments of an owner-scoped session; not found when the + /// owner has no event in it. + fn session_streams( + &self, + user_id: Uuid, + session_pk: Uuid, + ) -> BoxedFuture<'_, Result>; + + /// Derives the turn view of an owner-scoped session; not found when the + /// owner has no event in it. + fn session_turns( &self, user_id: Uuid, session_pk: Uuid, - ) -> BoxedFuture<'_, Result>; + ) -> BoxedFuture<'_, Result>; - /// Loads one owner-scoped session timeline without event content bodies. - fn session_outline( + /// The exchanges of the segment holding exchange `head` in + /// an owner-scoped session, from `head` to the segment end; not found when the session + /// has no such exchange. + fn session_segment( &self, user_id: Uuid, session_pk: Uuid, - ) -> BoxedFuture<'_, Result>; + head: String, + query: ValidatedSegmentQuery, + ) -> BoxedFuture<'_, Result>; + + /// The conversation of one owner-scoped event through `window`. + fn event_conversation( + &self, + user_id: Uuid, + event_pk: Uuid, + window: Window, + ) -> BoxedFuture<'_, Result>; /// Lists one page of owner-scoped sessions, most recently active first. fn list_sessions( @@ -104,13 +146,22 @@ pub trait StorageBackend: Send + Sync { attachment_id: Uuid, ) -> BoxedFuture<'_, Result>; - /// Executes one owner-scoped full-text session search. + /// Executes one owner-scoped session search. fn session_search( &self, user_id: Uuid, query: SessionSearchQuery, ) -> BoxedFuture<'_, Result>; + /// Finds every hit of a query in one owner-scoped session, with the + /// segment of each; not found when the owner has no such session. + fn search_in_session( + &self, + user_id: Uuid, + session_pk: Uuid, + query: InSessionSearchQuery, + ) -> BoxedFuture<'_, Result>; + /// Stops backend-owned background work and releases lifecycle resources. fn shutdown(&self) -> BoxedFuture<'_, ()>; } diff --git a/crates/abyss-backend/src/storage/postgres.rs b/crates/abyss-backend/src/storage/postgres.rs index 745cabd..7c7b493 100644 --- a/crates/abyss-backend/src/storage/postgres.rs +++ b/crates/abyss-backend/src/storage/postgres.rs @@ -5,8 +5,9 @@ mod repository; mod search; mod session_list; mod session_titles; +mod stream_facts; -use std::time::Duration; +use std::{sync::Arc, time::Duration}; use diesel::PgConnection; use tokio::{sync::Mutex, task::JoinHandle}; @@ -16,15 +17,24 @@ use crate::{ config::Config, db::{self, DbPool}, error::AppError, - search::{SessionSearchQuery, SessionSearchResponse}, - storage::{BoxedFuture, StorageBackend}, + search::{ + InSessionSearchQuery, InSessionSearchResponse, SessionSearchQuery, SessionSearchResponse, + }, + storage::{ + BoxedFuture, StorageBackend, + event_locations::{locate_events, locate_search_matches}, + stream_cache::StreamCache, + }, usage::{ - IngestEventsRequest, IngestEventsResponse, RawEventsQuery, RawEventsResponse, - SessionOutlineResponse, SessionTimelineResponse, SummaryQuery, SummaryResponse, - UsageEventResponse, + IngestEventsRequest, IngestEventsResponse, RawEventsQuery, RawEventsResponse, SummaryQuery, + SummaryResponse, UsageEventResponse, attachments::StoredImageAttachment, comparison::{ComparisonQuery, ComparisonResponse}, + preview::{ConversationPreview, Window}, + segments::{SegmentResponse, ValidatedSegmentQuery}, + session_header::SessionHeaderResponse, session_list::{SessionListResponse, ValidatedSessionListQuery}, + streams::{EventLocationResponse, SessionStreamsResponse, SessionTurnsResponse}, }, }; @@ -38,6 +48,7 @@ struct SearchWorker { /// Existing PostgreSQL and optional Elasticsearch storage stack. pub struct PostgresEsBackend { pool: DbPool, + structures: Arc, search: Option, search_worker: Option, } @@ -67,6 +78,7 @@ impl PostgresEsBackend { Ok(Self { pool, + structures: Arc::new(StreamCache::default()), search, search_worker, }) @@ -112,9 +124,29 @@ impl StorageBackend for PostgresEsBackend { query: RawEventsQuery, default_page_size: i64, ) -> BoxedFuture<'_, Result> { + let structures = Arc::clone(&self.structures); Box::pin(async move { self.run_db(move |connection| { - repository::raw_events(connection, &query, user_id, default_page_size) + if !query.locate { + return repository::raw_events(connection, &query, user_id, default_page_size); + } + connection + .build_transaction() + .read_only() + .repeatable_read() + .run(|connection| { + let mut response = + repository::raw_events(connection, &query, user_id, default_page_size)?; + locate_events(&mut response.events, |session_pk| { + repository::session_structure( + connection, + user_id, + session_pk, + &structures, + ) + })?; + Ok(response) + }) }) .await }) @@ -125,9 +157,47 @@ impl StorageBackend for PostgresEsBackend { user_id: Uuid, event_pk: Uuid, ) -> BoxedFuture<'_, Result> { + let structures = Arc::clone(&self.structures); Box::pin(async move { - self.run_db(move |connection| repository::usage_event(connection, user_id, event_pk)) - .await + self.run_db(move |connection| { + connection + .build_transaction() + .read_only() + .repeatable_read() + .run(|connection| { + let mut event = repository::usage_event(connection, user_id, event_pk)?; + locate_events(std::slice::from_mut(&mut event), |session_pk| { + repository::session_structure( + connection, + user_id, + session_pk, + &structures, + ) + })?; + Ok(event) + }) + }) + .await + }) + } + + fn event_location( + &self, + user_id: Uuid, + event_pk: Uuid, + ) -> BoxedFuture<'_, Result> { + let structures = Arc::clone(&self.structures); + Box::pin(async move { + self.run_db(move |connection| { + connection + .build_transaction() + .read_only() + .repeatable_read() + .run(|connection| { + repository::event_location(connection, user_id, event_pk, &structures) + }) + }) + .await }) } @@ -156,27 +226,104 @@ impl StorageBackend for PostgresEsBackend { }) } - fn session_timeline( + fn session_header( + &self, + user_id: Uuid, + session_pk: Uuid, + ) -> BoxedFuture<'_, Result> { + let structures = Arc::clone(&self.structures); + Box::pin(async move { + self.run_db(move |connection| { + connection + .build_transaction() + .read_only() + .repeatable_read() + .run(|connection| { + repository::session_header(connection, user_id, session_pk, &structures) + }) + }) + .await + }) + } + + fn session_streams( + &self, + user_id: Uuid, + session_pk: Uuid, + ) -> BoxedFuture<'_, Result> { + let structures = Arc::clone(&self.structures); + Box::pin(async move { + self.run_db(move |connection| { + connection + .build_transaction() + .read_only() + .repeatable_read() + .run(|connection| { + repository::session_streams(connection, user_id, session_pk, &structures) + }) + }) + .await + }) + } + + fn session_turns( &self, user_id: Uuid, session_pk: Uuid, - ) -> BoxedFuture<'_, Result> { + ) -> BoxedFuture<'_, Result> { + let structures = Arc::clone(&self.structures); Box::pin(async move { self.run_db(move |connection| { - repository::session_timeline(connection, user_id, session_pk) + connection + .build_transaction() + .read_only() + .repeatable_read() + .run(|connection| { + repository::session_turns(connection, user_id, session_pk, &structures) + }) }) .await }) } - fn session_outline( + fn session_segment( &self, user_id: Uuid, session_pk: Uuid, - ) -> BoxedFuture<'_, Result> { + head: String, + query: ValidatedSegmentQuery, + ) -> BoxedFuture<'_, Result> { + let structures = Arc::clone(&self.structures); + Box::pin(async move { + self.run_db(move |connection| { + connection + .build_transaction() + .read_only() + .repeatable_read() + .run(|connection| { + repository::session_segment( + connection, + user_id, + session_pk, + &head, + &query, + &structures, + ) + }) + }) + .await + }) + } + + fn event_conversation( + &self, + user_id: Uuid, + event_pk: Uuid, + window: Window, + ) -> BoxedFuture<'_, Result> { Box::pin(async move { self.run_db(move |connection| { - repository::session_outline(connection, user_id, session_pk) + repository::event_conversation(connection, user_id, event_pk, &window) }) .await }) @@ -211,18 +358,53 @@ impl StorageBackend for PostgresEsBackend { user_id: Uuid, query: SessionSearchQuery, ) -> BoxedFuture<'_, Result> { + let structures = Arc::clone(&self.structures); Box::pin(async move { let search = self.search.clone().ok_or_else(|| { AppError::unavailable("session search is not configured".to_owned()) })?; let execution = search.search(user_id, query).await?; let session_ids = execution.session_ids(); - let details = self - .run_db(move |connection| { - SearchOutboxRepository::session_details(connection, user_id, &session_ids) - }) - .await?; - Ok(execution.hydrate(details)) + self.run_db(move |connection| { + connection + .build_transaction() + .read_only() + .repeatable_read() + .run(|connection| { + let details = SearchOutboxRepository::session_details( + connection, + user_id, + &session_ids, + )?; + let mut response = execution.hydrate(details); + locate_search_matches(&mut response, |session_pk| { + repository::session_structure( + connection, + user_id, + session_pk, + &structures, + ) + })?; + Ok(response) + }) + }) + .await + }) + } + + /// Elasticsearch holds each event's text as one document without entry + /// positions, so this profile cannot list the hits within a session. + fn search_in_session( + &self, + _user_id: Uuid, + _session_pk: Uuid, + query: InSessionSearchQuery, + ) -> BoxedFuture<'_, Result> { + Box::pin(async move { + query.validate()?; + Err(AppError::unavailable( + "search within a session is not supported by this storage backend".to_owned(), + )) }) } diff --git a/crates/abyss-backend/src/storage/postgres/repository.rs b/crates/abyss-backend/src/storage/postgres/repository.rs index 5c24343..440118c 100644 --- a/crates/abyss-backend/src/storage/postgres/repository.rs +++ b/crates/abyss-backend/src/storage/postgres/repository.rs @@ -2,57 +2,70 @@ //! //! This module is the transactional boundary for the event hierarchy. Ingest //! validates the complete request before opening a transaction, upserts device, -//! session, and turn aggregates, inserts immutable events idempotently, and +//! session aggregates, inserts immutable events idempotently, and //! enqueues search projection in the same commit. Query functions always apply //! the authenticated owner before returning conversation or attachment data. -use std::collections::{HashMap, HashSet}; +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, +}; use chrono::{DateTime, Utc}; +use chrono_tz::Tz; use diesel::{ - Connection, ExpressionMethods, IntoSql, JoinOnDsl, PgConnection, QueryDsl, QueryableByName, - RunQueryDsl, SelectableHelper, define_sql_function, - dsl::{count_star, max, sql}, + Connection, ExpressionMethods, IntoSql, NullableExpressionMethods, PgConnection, QueryDsl, + QueryableByName, RunQueryDsl, define_sql_function, result::OptionalExtension, sql_query, sql_types::{ - Array, BigInt, Bool, Bytea, Integer, Jsonb, Nullable, SingleValue, Text, Timestamptz, - Uuid as SqlUuid, + Array, BigInt, Bool, Jsonb, Nullable, SingleValue, Text, Timestamptz, Uuid as SqlUuid, }, upsert::excluded, }; +use serde::de::DeserializeOwned; +use serde_json::Value; use uuid::Uuid; use crate::{ db::{ models::{ - AgentSession, AgentTurn, Device, NewAgentDiagnosticCapture, - NewAgentDiagnosticCaptureEvent, NewAgentSession, NewAgentTurn, NewDevice, - NewSearchOutboxTask, NewUsageEvent, NewUsageEventAttachment, UsageEvent, - UsageEventAttachment, + Device, NewAgentDiagnosticCapture, NewAgentDiagnosticCaptureEvent, NewAgentSession, + NewDevice, NewSearchOutboxTask, NewStreamFacts, NewUsageEvent, NewUsageEventAttachment, + StreamEventRow, UsageEvent, UsageEventAttachment, }, schema::{ - agent_diagnostic_capture_events, agent_diagnostic_captures, agent_sessions, - agent_turns, devices, llm_usage_event_attachments, llm_usage_events, search_outbox, + agent_diagnostic_capture_events, agent_diagnostic_captures, agent_sessions, devices, + llm_usage_event_attachments, llm_usage_event_stream_facts, llm_usage_events, + search_outbox, }, }, error::AppError, + storage::stream_cache::StreamCache, usage::{ - IngestEventsRequest, IngestEventsResponse, IngestUsageEvent, OUTLINE_PREVIEW_CHARS, - RawEventsQuery, RawEventsResponse, SessionInfo, SessionOutlineResponse, - SessionTimelineResponse, SummaryQuery, SummaryResponse, SummaryRow, TurnTimeline, - UsageEventOutline, UsageEventResponse, + IngestEventsRequest, IngestEventsResponse, IngestUsageEvent, RawEventsQuery, + RawEventsResponse, SummaryQuery, SummaryResponse, SummaryRow, UsageEventResponse, + UsageEventType, attachments::{ ImageMediaType, StoredImageAttachment, UsageEventAttachmentResponse, ValidatedImageAttachment, }, diagnostics::IngestDiagnosticCapture, - event_order::UsageEventTimelineOrder, persistence::{ escape_like_pattern, exact_values_filter, non_empty, non_negative_count, - normalize_limit, normalize_text, parse_group_by, sha256_bytes, unix_epoch, - validate_batch, validate_event_type, + normalize_limit, normalize_text, parse_group_by, unix_epoch, validate_batch, + validate_event_type, + }, + preview::{ConversationPreview, Window, preview}, + segments::{ + SegmentResponse, TokenCounts, ValidatedSegmentQuery, run_event_pks, segment_response, + }, + session_header::{HeaderEvent, SessionHeaderResponse, session_header as header_of}, + streams::{ + EventFacts, EventLocationResponse, Label, RequestFacts, ResponseFacts, + SessionStreamsResponse, SessionStructure, SessionTurnsResponse, StreamEvent, }, + time_zone::parse_time_zone, }, }; @@ -62,53 +75,90 @@ define_sql_function!(fn coalesce(value: Nullable, fallback: define_sql_function!(fn least(left: ST, right: ST) -> ST); define_sql_function!(fn greatest(left: ST, right: ST) -> ST); +/// Usage totals grouped in SQL. Turns are counted per session within each +/// group, then summed: the distinct declared `turn_id` values when the +/// session declares any turn, else the requests that open a turn. const SUMMARY_AGGREGATE_SQL: &str = "\ + WITH first_openers AS ( \ + SELECT id FROM ( \ + SELECT fe.id, ROW_NUMBER() OVER ( \ + PARTITION BY fe.session_pk ORDER BY fe.observed_at, fe.id \ + ) AS n \ + FROM llm_usage_events fe \ + JOIN llm_usage_event_stream_facts ff ON ff.event_pk = fe.id \ + WHERE ff.starts_turn \ + ) firsts WHERE n = 1 \ + ), \ + by_session AS ( \ + SELECT \ + CASE WHEN $1 THEN to_char(e.observed_at AT TIME ZONE $30, 'YYYY-MM-DD') END AS day, \ + CASE WHEN $2 THEN e.user_id END AS user_id, \ + CASE WHEN $2 THEN NULLIF(btrim(u.name), '') END AS user_name, \ + CASE WHEN $2 THEN u.email END AS user_email, \ + CASE WHEN $3 THEN d.host_name END AS host_name, \ + CASE WHEN $3 THEN d.platform END AS platform, \ + CASE WHEN $3 THEN d.os_version END AS os_version, \ + CASE WHEN $4 THEN e.agent_name END AS agent_name, \ + CASE WHEN $5 THEN e.llm_provider END AS llm_provider, \ + CASE WHEN $6 THEN e.llm_model END AS llm_model, \ + CASE WHEN $7 THEN e.event_type END AS event_type, \ + e.session_pk AS session_pk, \ + EXISTS ( \ + SELECT 1 FROM llm_usage_events de \ + JOIN llm_usage_event_stream_facts df ON df.event_pk = de.id \ + WHERE de.session_pk = e.session_pk AND df.turn_id IS NOT NULL \ + ) AS declares_turns, \ + COUNT(DISTINCT f.turn_id)::bigint AS declared_turns, \ + COUNT(*) FILTER (WHERE f.starts_turn \ + AND (f.continues_context OR e.id IN (SELECT id FROM first_openers)))::bigint AS opened_turns, \ + COUNT(*) FILTER (WHERE e.event_type = 'request')::bigint AS requests, \ + COUNT(*) FILTER (WHERE e.event_type = 'response')::bigint AS responses, \ + COALESCE(SUM(e.input_tokens), 0)::bigint AS input_tokens, \ + COALESCE(SUM(e.output_tokens), 0)::bigint AS output_tokens, \ + COALESCE(SUM(e.cache_read_tokens), 0)::bigint AS cache_read_tokens, \ + COALESCE(SUM(e.cache_write_tokens), 0)::bigint AS cache_write_tokens, \ + COALESCE(SUM(e.reasoning_tokens), 0)::bigint AS reasoning_tokens, \ + COALESCE(SUM(e.total_tokens), 0)::bigint AS total_tokens \ + FROM llm_usage_events e \ + LEFT JOIN llm_usage_event_stream_facts f ON f.event_pk = e.id \ + LEFT JOIN app_users u ON u.id = e.user_id \ + LEFT JOIN devices d ON d.id = e.device_context_id \ + WHERE e.user_id = $8 \ + AND ($9 = false OR e.observed_at >= $10) \ + AND ($11 = false OR e.observed_at < $12) \ + AND ( \ + $13 = false \ + OR ($15 = true AND e.user_id = $14) \ + OR ($15 = false AND ( \ + u.email ILIKE $16 ESCAPE '\\' \ + OR COALESCE(u.name, '') ILIKE $16 ESCAPE '\\' \ + )) \ + ) \ + AND ($17 = false OR e.agent_name = ANY($18)) \ + AND ($19 = false OR e.session_id = ANY($20)) \ + AND ($21 = false OR e.session_pk = $22) \ + AND ($23 = false OR e.llm_provider = ANY($24)) \ + AND ($25 = false OR e.llm_model = ANY($26)) \ + AND ($27 = false OR e.event_type = $28) \ + GROUP BY 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 \ + ) \ SELECT \ - CASE WHEN $1 THEN to_char(e.observed_at AT TIME ZONE 'UTC', 'YYYY-MM-DD') END AS day, \ - CASE WHEN $2 THEN e.user_id END AS user_id, \ - CASE WHEN $2 THEN NULLIF(btrim(u.name), '') END AS user_name, \ - CASE WHEN $2 THEN u.email END AS user_email, \ - CASE WHEN $3 THEN d.host_name END AS host_name, \ - CASE WHEN $3 THEN d.platform END AS platform, \ - CASE WHEN $3 THEN d.os_version END AS os_version, \ - CASE WHEN $4 THEN e.agent_name END AS agent_name, \ - CASE WHEN $5 THEN e.llm_provider END AS llm_provider, \ - CASE WHEN $6 THEN e.llm_model END AS llm_model, \ - CASE WHEN $7 THEN e.event_type END AS event_type, \ - COUNT(DISTINCT e.session_pk)::bigint AS sessions, \ - COUNT(DISTINCT e.turn_pk)::bigint AS turns, \ - COUNT(*) FILTER (WHERE e.event_type = 'request')::bigint AS requests, \ - COUNT(*) FILTER (WHERE e.event_type = 'response')::bigint AS responses, \ - COALESCE(SUM(e.input_tokens), 0)::bigint AS input_tokens, \ - COALESCE(SUM(e.output_tokens), 0)::bigint AS output_tokens, \ - COALESCE(SUM(e.cache_read_tokens), 0)::bigint AS cache_read_tokens, \ - COALESCE(SUM(e.cache_write_tokens), 0)::bigint AS cache_write_tokens, \ - COALESCE(SUM(e.reasoning_tokens), 0)::bigint AS reasoning_tokens, \ - COALESCE(SUM(e.total_tokens), 0)::bigint AS total_tokens \ - FROM llm_usage_events e \ - LEFT JOIN app_users u ON u.id = e.user_id \ - LEFT JOIN devices d ON d.id = e.device_context_id \ - WHERE e.user_id = $8 \ - AND ($9 = false OR e.observed_at >= $10) \ - AND ($11 = false OR e.observed_at < $12) \ - AND ( \ - $13 = false \ - OR ($15 = true AND e.user_id = $14) \ - OR ($15 = false AND ( \ - u.email ILIKE $16 ESCAPE '\\' \ - OR COALESCE(u.name, '') ILIKE $16 ESCAPE '\\' \ - )) \ - ) \ - AND ($17 = false OR e.agent_name = ANY($18)) \ - AND ($19 = false OR e.session_id = ANY($20)) \ - AND ($21 = false OR e.session_pk = $22) \ - AND ($23 = false OR e.turn_index = $24) \ - AND ($25 = false OR e.llm_provider = ANY($26)) \ - AND ($27 = false OR e.llm_model = ANY($28)) \ - AND ($29 = false OR e.event_type = $30) \ + day, user_id, user_name, user_email, host_name, platform, os_version, \ + agent_name, llm_provider, llm_model, event_type, \ + COUNT(*)::bigint AS sessions, \ + COALESCE(SUM(CASE WHEN declares_turns THEN declared_turns ELSE opened_turns END), 0)::bigint AS turns, \ + COALESCE(SUM(requests), 0)::bigint AS requests, \ + COALESCE(SUM(responses), 0)::bigint AS responses, \ + COALESCE(SUM(input_tokens), 0)::bigint AS input_tokens, \ + COALESCE(SUM(output_tokens), 0)::bigint AS output_tokens, \ + COALESCE(SUM(cache_read_tokens), 0)::bigint AS cache_read_tokens, \ + COALESCE(SUM(cache_write_tokens), 0)::bigint AS cache_write_tokens, \ + COALESCE(SUM(reasoning_tokens), 0)::bigint AS reasoning_tokens, \ + COALESCE(SUM(total_tokens), 0)::bigint AS total_tokens \ + FROM by_session g \ GROUP BY 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11 \ ORDER BY total_tokens DESC, agent_name ASC NULLS LAST, user_email ASC NULLS LAST, day ASC NULLS LAST \ - LIMIT $31"; + LIMIT $29"; /// Validates and atomically ingests one batch for the authenticated owner. /// @@ -122,7 +172,7 @@ pub fn ingest_events( ) -> Result { let validated_attachments = validate_batch(request, max_batch_size)?; - // Process the batch atomically so the device/session/turn hierarchy cannot + // Process the batch atomically so the device/session hierarchy cannot // be partially created without its corresponding usage events. let (accepted, duplicates, accepted_diagnostic_captures, duplicate_diagnostic_captures) = connection.transaction(|transaction| { @@ -240,6 +290,7 @@ pub fn usage_summary( default_limit: i64, ) -> Result { let group_by = parse_group_by(query.group_by.as_deref()); + let zone = parse_time_zone(query.tz.as_deref())?; let limit = normalize_limit(query.limit, default_limit, 100_000); let rows = load_filtered_summary_rows( connection, @@ -251,7 +302,6 @@ pub fn usage_summary( agent_names: &query.agent_name, session_ids: &query.session_id, session_pk: None, - turn_index: None, llm_providers: &query.llm_provider, llm_models: &query.llm_model, event_type: query.event_type.as_deref(), @@ -259,6 +309,7 @@ pub fn usage_summary( offset: 0, }, &group_by, + zone, )?; Ok(SummaryResponse { @@ -274,6 +325,7 @@ fn load_filtered_summary_rows( connection: &mut PgConnection, filters: &EventFilters<'_>, group_by: &[String], + zone: Tz, ) -> Result, AppError> { let dimensions = SummaryDimensions::from_group_by(group_by); let sql_filters = SummarySqlFilters::from_event_filters(filters)?; @@ -292,8 +344,6 @@ fn load_filtered_summary_rows( let has_session_id_filter = !session_ids.is_empty(); let session_pk = sql_filters.session_pk.unwrap_or_else(Uuid::nil); let has_session_pk_filter = session_pk != Uuid::nil(); - let turn_index = sql_filters.turn_index.unwrap_or_default(); - let has_turn_index_filter = turn_index > 0_i32; let providers = sql_filters.providers; let has_provider_filter = !providers.is_empty(); let models = sql_filters.models; @@ -327,8 +377,6 @@ fn load_filtered_summary_rows( .bind::, _>(session_ids) .bind::(has_session_pk_filter) .bind::(session_pk) - .bind::(has_turn_index_filter) - .bind::(turn_index) .bind::(has_provider_filter) .bind::, _>(providers) .bind::(has_model_filter) @@ -336,200 +384,12 @@ fn load_filtered_summary_rows( .bind::(has_event_type_filter) .bind::(&event_type) .bind::(filters.limit) + .bind::(zone.name()) .load::(connection) .map(|rows| rows.into_iter().map(SummaryRow::from).collect()) .map_err(AppError::from) } -/// Loads one owner-scoped session with normalized turns and ordered events. -pub fn session_timeline( - connection: &mut PgConnection, - user_id: Uuid, - session_pk: Uuid, -) -> Result { - let (session, turns) = load_session_turns(connection, user_id, session_pk)?; - - let mut events = llm_usage_events::table - .filter(llm_usage_events::session_pk.eq(session.id)) - .order(( - llm_usage_events::turn_index.asc(), - llm_usage_events::observed_at.asc(), - )) - .load::(connection)?; - UsageEventTimelineOrder::sort(&mut events); - - let device_contexts = load_device_contexts_for_events(connection, &events)?; - let mut attachments_by_event = load_attachments_for_events(connection, &events)?; - let mut events_by_turn: HashMap> = HashMap::new(); - for event in events { - let event_pk = event.id; - let device_context_id = event.device_context_id; - events_by_turn - .entry(event.turn_pk) - .or_default() - .push(event_response( - event, - device_contexts.get(&device_context_id), - attachments_by_event.remove(&event_pk).unwrap_or_default(), - )); - } - let mut timeline = assemble_timeline(session, turns, events_by_turn); - session_titles::fill(connection, &mut timeline.session)?; - Ok(timeline) -} - -type OutlineRow = ( - Uuid, - String, - Uuid, - i32, - String, - String, - String, - i64, - i64, - i64, - i64, - i64, - i64, - DateTime, - i64, - Option, -); - -/// Loads one owner-scoped session timeline without event content bodies. -pub fn session_outline( - connection: &mut PgConnection, - user_id: Uuid, - session_pk: Uuid, -) -> Result { - let (session, turns) = load_session_turns(connection, user_id, session_pk)?; - // Only fixed-size columns and SQL-computed summaries are selected, so the - // outline stays small no matter how large event text or metadata grows. - let rows = llm_usage_events::table - .filter(llm_usage_events::session_pk.eq(session.id)) - .select(( - llm_usage_events::id, - llm_usage_events::event_id, - llm_usage_events::turn_pk, - llm_usage_events::turn_index, - llm_usage_events::llm_provider, - llm_usage_events::llm_model, - llm_usage_events::event_type, - llm_usage_events::input_tokens, - llm_usage_events::output_tokens, - llm_usage_events::cache_read_tokens, - llm_usage_events::cache_write_tokens, - llm_usage_events::reasoning_tokens, - llm_usage_events::total_tokens, - llm_usage_events::observed_at, - sql::("coalesce(char_length(text), 0)::bigint"), - sql::>(&format!( - "CASE WHEN event_type = 'request' THEN right(text, {OUTLINE_PREVIEW_CHARS}) \ - ELSE left(text, {OUTLINE_PREVIEW_CHARS}) END" - )), - )) - .load::(connection)?; - let event_ids = rows.iter().map(|row| row.0).collect::>(); - let attachment_counts = llm_usage_event_attachments::table - .filter(llm_usage_event_attachments::event_pk.eq_any(event_ids)) - .group_by(llm_usage_event_attachments::event_pk) - .select((llm_usage_event_attachments::event_pk, count_star())) - .load::<(Uuid, i64)>(connection)? - .into_iter() - .collect::>(); - let mut events = rows - .into_iter() - .map(|row| UsageEventOutline { - id: row.0, - event_id: row.1, - turn_pk: row.2, - turn_index: row.3, - llm_provider: row.4, - llm_model: row.5, - event_type: row.6, - input_tokens: row.7, - output_tokens: row.8, - cache_read_tokens: row.9, - cache_write_tokens: row.10, - reasoning_tokens: row.11, - total_tokens: row.12, - observed_at: row.13, - text_chars: row.14, - preview: row.15, - attachment_count: attachment_counts.get(&row.0).copied().unwrap_or(0), - }) - .collect::>(); - UsageEventTimelineOrder::sort(&mut events); - let mut events_by_turn: HashMap> = HashMap::new(); - for event in events { - events_by_turn.entry(event.turn_pk).or_default().push(event); - } - let mut timeline = assemble_timeline(session, turns, events_by_turn); - session_titles::fill(connection, &mut timeline.session)?; - Ok(timeline) -} - -fn load_session_turns( - connection: &mut PgConnection, - user_id: Uuid, - session_pk: Uuid, -) -> Result<(AgentSession, Vec), AppError> { - // Timeline views include prompt/response text, so they are always limited - // to the owner. - let session = agent_sessions::table - .filter(agent_sessions::id.eq(session_pk)) - .filter(agent_sessions::user_id.eq(user_id)) - .first::(connection) - .optional()? - .ok_or_else(|| AppError::not_found(format!("session not found: {session_pk}")))?; - touch_session_metadata(&session); - - let turns = agent_turns::table - .filter(agent_turns::session_pk.eq(session.id)) - .order(agent_turns::turn_index.asc()) - .load::(connection)?; - Ok((session, turns)) -} - -fn assemble_timeline( - session: AgentSession, - turns: Vec, - mut events_by_turn: HashMap>, -) -> SessionTimelineResponse { - let turn_timelines = turns - .into_iter() - .map(|turn| { - touch_turn_metadata(&turn); - TurnTimeline { - turn_pk: turn.id, - turn_index: turn.turn_index, - started_at: turn.started_at, - ended_at: turn.ended_at, - events: events_by_turn.remove(&turn.id).unwrap_or_default(), - } - }) - .collect(); - - SessionTimelineResponse { - session: SessionInfo { - session_pk: session.id, - user_id: session.user_id, - device_context_id: session.device_context_id, - agent_name: session.agent_name, - agent_version: session.agent_version, - session_id: session.session_id, - started_at: session.started_at, - ended_at: session.ended_at, - metadata: session.metadata, - title: None, - purpose: None, - parent_session_pk: None, - }, - turns: turn_timelines, - } -} - /// Returns one newest-first page of raw owner-scoped events. pub fn raw_events( connection: &mut PgConnection, @@ -553,7 +413,6 @@ pub fn raw_events( agent_names: &query.agent_name, session_ids: &query.session_id, session_pk: query.session_pk, - turn_index: query.turn_index, llm_providers: &query.llm_provider, llm_models: &query.llm_model, event_type: query.event_type.as_deref(), @@ -659,12 +518,349 @@ pub fn image_attachment( }) } +/// Derives the streams of one session's events from their stream facts, +/// without reading any request body. +pub fn session_streams( + connection: &mut PgConnection, + user_id: Uuid, + session_pk: Uuid, + cache: &StreamCache, +) -> Result { + let structure = session_structure(connection, user_id, session_pk, cache)?; + Ok(structure.session_streams(session_pk)) +} + +/// Derives the turn view of one session from its stream facts, without +/// reading any request body. +pub fn session_turns( + connection: &mut PgConnection, + user_id: Uuid, + session_pk: Uuid, + cache: &StreamCache, +) -> Result { + let structure = session_structure(connection, user_id, session_pk, cache)?; + Ok(structure.session_turns(session_pk)) +} + +/// Uses the caller's read transaction so the revision and facts are one snapshot. +/// A cache hit reads only the session revision; it does not load the event facts. +pub fn session_structure( + connection: &mut PgConnection, + user_id: Uuid, + session_pk: Uuid, + cache: &StreamCache, +) -> Result, AppError> { + let revision = agent_sessions::table + .filter(agent_sessions::user_id.eq(user_id)) + .filter(agent_sessions::id.eq(session_pk)) + .select(agent_sessions::structure_revision) + .first::(connection) + .optional()? + .ok_or_else(|| AppError::not_found(format!("session not found: {session_pk}")))?; + cache.get_or_load(user_id, session_pk, revision, || { + let events = stream_events(connection, session_pk)?; + if events.is_empty() { + return Err(AppError::not_found(format!( + "session not found: {session_pk}" + ))); + } + Ok(events) + }) +} + +/// Resolves an authorized event without reading its body or attachments. +pub fn event_location( + connection: &mut PgConnection, + user_id: Uuid, + event_pk: Uuid, + cache: &StreamCache, +) -> Result { + let session_pk = llm_usage_events::table + .filter(llm_usage_events::user_id.eq(user_id)) + .filter(llm_usage_events::id.eq(event_pk)) + .select(llm_usage_events::session_pk) + .first::(connection) + .optional()? + .ok_or_else(|| AppError::not_found(format!("event not found: {event_pk}")))?; + session_structure(connection, user_id, session_pk, cache)? + .event_location(session_pk, event_pk) + .ok_or_else(|| AppError::not_found(format!("event not found: {event_pk}"))) +} + +/// Columns of one session event the header reads. +type HeaderColumns = ( + String, + String, + Option, + String, + String, + DateTime, + i64, + i64, + i64, + i64, + i64, + i64, +); + +/// The header of an owner-scoped session. +pub fn session_header( + connection: &mut PgConnection, + user_id: Uuid, + session_pk: Uuid, + cache: &StreamCache, +) -> Result { + let owned = agent_sessions::table + .filter(agent_sessions::user_id.eq(user_id)) + .filter(agent_sessions::id.eq(session_pk)) + .select(agent_sessions::id) + .first::(connection) + .optional()?; + if owned.is_none() { + return Err(AppError::not_found(format!( + "session not found: {session_pk}" + ))); + } + // The owner is checked on the session row, so events are read by session + // alone, the same way as the SQLite header. + let events = llm_usage_events::table + .filter(llm_usage_events::session_pk.eq(session_pk)) + .select(( + llm_usage_events::session_id, + llm_usage_events::agent_name, + llm_usage_events::agent_version, + llm_usage_events::llm_provider, + llm_usage_events::llm_model, + llm_usage_events::observed_at, + llm_usage_events::input_tokens, + llm_usage_events::output_tokens, + llm_usage_events::cache_read_tokens, + llm_usage_events::cache_write_tokens, + llm_usage_events::reasoning_tokens, + llm_usage_events::total_tokens, + )) + .load::(connection)? + .into_iter() + .map(header_event) + .collect::>(); + if events.is_empty() { + return Err(AppError::not_found(format!( + "session not found: {session_pk}" + ))); + } + let summary = session_structure(connection, user_id, session_pk, cache)?.summary(); + let title = session_titles::title(connection, session_pk)?; + Ok(header_of(session_pk, events, &summary, title)) +} + +fn header_event( + ( + session_id, + agent_name, + agent_version, + provider, + model, + observed_at, + input, + output, + cache_read, + cache_write, + reasoning, + total, + ): HeaderColumns, +) -> HeaderEvent { + HeaderEvent { + session_id, + agent_name, + agent_version, + provider, + model, + observed_at, + tokens: TokenCounts { + input, + output, + cache_read, + cache_write, + reasoning, + total, + }, + } +} + +/// The exchanges of the segment holding exchange `head`, from it to the +/// segment end, with a preview of each event. +pub fn session_segment( + connection: &mut PgConnection, + user_id: Uuid, + session_pk: Uuid, + head: &str, + query: &ValidatedSegmentQuery, + cache: &StreamCache, +) -> Result { + let structure = session_structure(connection, user_id, session_pk, cache)?; + let run = structure.segment_run(head).ok_or_else(|| { + AppError::not_found(format!( + "exchange not found in session {session_pk}: {head}" + )) + })?; + let events = usage_events(connection, user_id, &run_event_pks(&run))?; + Ok(segment_response(session_pk, run, events, &query.window)) +} + +/// The conversation of one owner-scoped event through `window`. +pub fn event_conversation( + connection: &mut PgConnection, + user_id: Uuid, + event_pk: Uuid, + window: &Window, +) -> Result { + usage_event(connection, user_id, event_pk).map(|event| preview(event, window)) +} + +/// Owner-scoped events by primary key, with their device labels and +/// attachments. +fn usage_events( + connection: &mut PgConnection, + user_id: Uuid, + pks: &[Uuid], +) -> Result, AppError> { + let events = llm_usage_events::table + .filter(llm_usage_events::id.eq_any(pks)) + .filter(llm_usage_events::user_id.eq(user_id)) + .load::(connection)?; + let devices = load_device_contexts_for_events(connection, &events)?; + let mut attachments = load_attachments_for_events(connection, &events)?; + Ok(events + .into_iter() + .map(|event| { + let (event_pk, device_id) = (event.id, event.device_context_id); + let attachments = attachments.remove(&event_pk).unwrap_or_default(); + ( + event_pk, + event_response(event, devices.get(&device_id), attachments), + ) + }) + .collect()) +} + +/// The events of one session in session order, with their stream facts and +/// without any request body. The caller has checked the session's owner on +/// its row, so events are read by session alone. +fn stream_events( + connection: &mut PgConnection, + session_pk: Uuid, +) -> Result, AppError> { + llm_usage_events::table + .left_join(llm_usage_event_stream_facts::table) + .filter(llm_usage_events::session_pk.eq(session_pk)) + .order(( + llm_usage_events::observed_at.asc(), + llm_usage_events::event_type.asc(), + llm_usage_events::id.asc(), + )) + .select(( + llm_usage_events::id, + llm_usage_events::event_id, + llm_usage_events::event_type, + llm_usage_events::observed_at, + llm_usage_events::response_id, + llm_usage_events::previous_response_id, + llm_usage_event_stream_facts::new_from.nullable(), + llm_usage_event_stream_facts::starts_turn.nullable(), + llm_usage_event_stream_facts::last_identity.nullable(), + llm_usage_event_stream_facts::previous_identity.nullable(), + llm_usage_event_stream_facts::trailing.nullable(), + llm_usage_event_stream_facts::result_ids.nullable(), + llm_usage_event_stream_facts::emitted_ids.nullable(), + llm_usage_event_stream_facts::output_skeleton.nullable(), + llm_usage_event_stream_facts::label_index.nullable(), + llm_usage_event_stream_facts::label_text.nullable(), + llm_usage_event_stream_facts::label_length.nullable(), + llm_usage_event_stream_facts::subagent_id.nullable(), + llm_usage_event_stream_facts::parent_subagent_id.nullable(), + llm_usage_event_stream_facts::turn_id.nullable(), + )) + .load::(connection)? + .into_iter() + .map(stream_event) + .collect() +} + +fn stream_event(row: StreamEventRow) -> Result { + let facts = if row.event_type == UsageEventType::Request.as_str() { + EventFacts::Request(RequestFacts { + new_from: row.new_from.unwrap_or_default(), + starts_turn: row.starts_turn.unwrap_or_default(), + last_identity: row.last_identity, + previous_identity: row.previous_identity, + trailing: json_list(row.trailing, "trailing roles")?, + result_ids: json_list(row.result_ids, "result ids")?, + label: Label::from_columns(row.label_index, row.label_text, row.label_length), + }) + } else { + EventFacts::Response(ResponseFacts { + emitted_ids: json_list(row.emitted_ids, "emitted ids")?, + output_skeleton: json_list(row.output_skeleton, "output skeleton")?, + }) + }; + Ok(StreamEvent { + pk: row.id, + event_id: row.event_id, + observed_at: row.observed_at, + response_id: row.response_id, + previous_response_id: row.previous_response_id, + subagent_id: row.subagent_id, + parent_subagent_id: row.parent_subagent_id, + turn_id: row.turn_id, + facts, + }) +} + +/// A stored JSON list; absent reads as empty. +fn json_list( + value: Option, + field: &str, +) -> Result { + value.map_or_else( + || Ok(T::default()), + |value| { + serde_json::from_value(value) + .map_err(|error| AppError::internal(format!("stored {field} are invalid: {error}"))) + }, + ) +} + +/// Stream fact columns of one ingested event, with the ids it declares. +fn stream_facts(event_pk: Uuid, event: &IngestUsageEvent) -> Result { + Ok(NewStreamFacts { + subagent_id: event.subagent_id.clone(), + parent_subagent_id: event.parent_subagent_id.clone(), + turn_id: event.turn_id.clone(), + ..super::stream_facts::read( + event_pk, + event.event_type.as_str(), + event.conversation.as_ref(), + event.previous_response_id.is_some(), + )? + }) +} + fn ingest_one_event( connection: &mut PgConnection, user_id: Uuid, event: &IngestUsageEvent, attachments: Vec, ) -> Result { + let exists = llm_usage_events::table + .filter(llm_usage_events::event_id.eq(&event.event_id)) + .select(llm_usage_events::id) + .first::(connection) + .optional()? + .is_some(); + if exists { + return Ok(false); + } + let observed_at = event.observed_at; let now = Utc::now(); let agent_name = event.agent.name.clone(); @@ -673,7 +869,7 @@ fn ingest_one_event( let model = event.llm.model.clone(); let device = upsert_device(connection, user_id, event, observed_at, now)?; - let session = upsert_session( + let session_pk = upsert_session( connection, UpsertSessionInput { device: &device, @@ -685,17 +881,7 @@ fn ingest_one_event( now, }, )?; - let (turn_pk, turn_index) = place_in_turn( - connection, - user_id, - session.id, - event.turn_id.as_deref(), - event.response_id.as_deref(), - observed_at, - now, - )?; let tokens = &event.token_usage; - let text = event.text.clone(); // `event_id` is the idempotency key supplied by the collector. Duplicate // observations are accepted as no-ops instead of failing the whole batch. @@ -705,18 +891,14 @@ fn ingest_one_event( id: usage_event_id, user_id, device_context_id: device.id, - session_pk: session.id, - turn_pk, + session_pk, event_id: event.event_id.clone(), agent_name: agent_name.clone(), agent_version, session_id: event.session_id.clone(), - turn_index, llm_provider: provider, llm_model: model, event_type: event.event_type.as_str().to_owned(), - text_sha256: text.as_deref().map(sha256_bytes), - text, input_tokens: tokens.input_tokens.unwrap_or_default(), output_tokens: tokens.output_tokens.unwrap_or_default(), cache_read_tokens: tokens.cache_read_tokens.unwrap_or_default(), @@ -729,6 +911,7 @@ fn ingest_one_event( stream_id: event.stream_id.clone(), response_id: event.response_id.clone(), previous_response_id: event.previous_response_id.clone(), + conversation: event.conversation.clone(), }) .on_conflict(llm_usage_events::event_id) .do_nothing() @@ -743,6 +926,9 @@ fn ingest_one_event( // The enclosing transaction commits the event and its projection task // together, providing at-least-once delivery to Elasticsearch. if inserted == 1 { + diesel::insert_into(llm_usage_event_stream_facts::table) + .values(stream_facts(usage_event_id, event)?) + .execute(connection)?; diesel::insert_into(search_outbox::table) .values(NewSearchOutboxTask { event_pk: usage_event_id, @@ -754,11 +940,9 @@ fn ingest_one_event( connection, &session_titles::CapturedEvent { user_id, - session_pk: session.id, + session_pk, event_pk: usage_event_id, event_type: event.event_type.as_str(), - metadata: &event.metadata, - naming_thread: event.naming_thread, observed_at, }, )?; @@ -834,10 +1018,14 @@ fn upsert_device( .map_err(AppError::from) } +#[expect( + clippy::arithmetic_side_effects, + reason = "Diesel builds an atomic SQL increment, not a Rust integer addition." +)] fn upsert_session( connection: &mut PgConnection, input: UpsertSessionInput<'_>, -) -> Result { +) -> Result { let session_id = input.event.session_id.clone(); diesel::insert_into(agent_sessions::table) .values(NewAgentSession { @@ -852,6 +1040,7 @@ fn upsert_session( metadata: serde_json::Value::Object(serde_json::Map::new()), created_at: input.now, updated_at: input.now, + structure_revision: 1, }) .on_conflict(( agent_sessions::user_id, @@ -874,202 +1063,16 @@ fn upsert_session( excluded(agent_sessions::ended_at), )), agent_sessions::updated_at.eq(input.now), + agent_sessions::structure_revision.eq(agent_sessions::structure_revision + 1_i64), )) - .get_result::(connection) - .map_err(AppError::from) -} - -/// Places an event in a turn of its session and returns the turn's id and -/// position. -/// -/// Explicit turn ids take priority; otherwise each response id gets its own -/// fallback turn. Events with neither id join the turn in progress. -/// Turns are numbered by when they started, so a turn that arrives late -/// moves the turns after it. -fn place_in_turn( - connection: &mut PgConnection, - user_id: Uuid, - session_pk: Uuid, - turn_id: Option<&str>, - response_id: Option<&str>, - observed_at: DateTime, - now: DateTime, -) -> Result<(Uuid, i32), AppError> { - let found = match (turn_id, response_id) { - (Some(key), _) => named_turn(connection, session_pk, key)?, - (None, Some(key)) => response_turn(connection, session_pk, key)?, - (None, None) => turn_in_progress(connection, session_pk, observed_at)?, - }; - let (turn_pk, turn_index, reordered) = if let Some(turn) = found { - diesel::update(agent_turns::table.find(turn.id)) - .set(( - agent_turns::started_at.eq(turn.started_at.min(observed_at)), - agent_turns::ended_at - .eq(Some(turn.ended_at.unwrap_or(observed_at).max(observed_at))), - agent_turns::updated_at.eq(now), - )) - .execute(connection)?; - (turn.id, turn.turn_index, observed_at < turn.started_at) - } else { - let last = agent_turns::table - .filter(agent_turns::session_pk.eq(session_pk)) - .select(max(agent_turns::turn_index)) - .get_result::>(connection)?; - let id = Uuid::now_v7(); - let turn_index = last.unwrap_or(0_i32).saturating_add(1_i32); - diesel::insert_into(agent_turns::table) - .values(NewAgentTurn { - id, - user_id, - session_pk, - turn_index, - turn_key: turn_id.map(str::to_owned), - started_at: observed_at, - ended_at: Some(observed_at), - created_at: now, - updated_at: now, - }) - .execute(connection)?; - (id, turn_index, last.is_some()) - }; - let turn_index = if reordered && number_turns(connection, session_pk)? { - agent_turns::table - .find(turn_pk) - .select(agent_turns::turn_index) - .get_result::(connection)? - } else { - turn_index - }; - Ok((turn_pk, turn_index)) -} - -/// The turn named `key` in a session. The first named turn adopts the turn -/// that events with neither turn nor response id opened at the start of the -/// session. A response fallback turn keeps its separate identity. -fn named_turn( - connection: &mut PgConnection, - session_pk: Uuid, - key: &str, -) -> Result, AppError> { - let named = agent_turns::table - .filter(agent_turns::session_pk.eq(session_pk)) - .filter(agent_turns::turn_key.eq(key)) - .select(AgentTurn::as_select()) - .first::(connection) - .optional()?; - if named.is_some() { - return Ok(named); - } - let mut turns = agent_turns::table - .filter(agent_turns::session_pk.eq(session_pk)) - .select(AgentTurn::as_select()) - .limit(2_i64) - .load::(connection)?; - let opening = match turns.as_slice() { - [turn] if turn.turn_key.is_none() => turns.pop(), - _ => None, - }; - if let Some(turn) = &opening { - let has_response = llm_usage_events::table - .filter(llm_usage_events::turn_pk.eq(turn.id)) - .filter(llm_usage_events::response_id.is_not_null()) - .select(llm_usage_events::id) - .first::(connection) - .optional()? - .is_some(); - if has_response { - return Ok(None); - } - diesel::update(agent_turns::table.find(turn.id)) - .set(agent_turns::turn_key.eq(key)) - .execute(connection)?; - } - Ok(opening) -} - -/// Reuses the fallback turn for a model response within this session. -/// Explicit turn keys and response ids occupy separate identity domains. -fn response_turn( - connection: &mut PgConnection, - session_pk: Uuid, - response_id: &str, -) -> Result, AppError> { - agent_turns::table - .inner_join(llm_usage_events::table.on(llm_usage_events::turn_pk.eq(agent_turns::id))) - .filter(llm_usage_events::session_pk.eq(session_pk)) - .filter(llm_usage_events::response_id.eq(response_id)) - .filter(agent_turns::turn_key.is_null()) - .select(AgentTurn::as_select()) - .first::(connection) - .optional() - .map_err(AppError::from) -} - -/// The turn in progress when an event naming no turn was observed: the latest -/// turn started by then, else the first turn of the session. -fn turn_in_progress( - connection: &mut PgConnection, - session_pk: Uuid, - observed_at: DateTime, -) -> Result, AppError> { - let started = agent_turns::table - .filter(agent_turns::session_pk.eq(session_pk)) - .filter(agent_turns::started_at.le(observed_at)) - .order_by((agent_turns::started_at.desc(), agent_turns::id.desc())) - .select(AgentTurn::as_select()) - .first::(connection) - .optional()?; - if started.is_some() { - return Ok(started); - } - agent_turns::table - .filter(agent_turns::session_pk.eq(session_pk)) - .order_by((agent_turns::started_at.asc(), agent_turns::id.asc())) - .select(AgentTurn::as_select()) - .first::(connection) - .optional() + .returning(agent_sessions::id) + .get_result::(connection) .map_err(AppError::from) } -/// Numbers the turns of a session, and their events, by when the turns -/// started. Writes nothing and returns `false` while the numbers already -/// follow that order. -fn number_turns(connection: &mut PgConnection, session_pk: Uuid) -> Result { - let turns = agent_turns::table - .filter(agent_turns::session_pk.eq(session_pk)) - .order_by((agent_turns::started_at.asc(), agent_turns::id.asc())) - .select((agent_turns::id, agent_turns::turn_index)) - .load::<(Uuid, i32)>(connection)?; - let moved = turns - .iter() - .zip(1_i32..) - .filter(|((_, index), position)| index != position) - .map(|((id, _), position)| (*id, position)) - .collect::>(); - if moved.is_empty() { - return Ok(false); - } - // Moved turns first step above every current number, so no single-row - // update collides with another turn's number or goes below one. - let above = turns.iter().map(|(_, index)| *index).max().unwrap_or(0_i32); - for (id, position) in &moved { - diesel::update(agent_turns::table.find(*id)) - .set(agent_turns::turn_index.eq(above.saturating_add(*position))) - .execute(connection)?; - } - for (id, position) in moved { - diesel::update(agent_turns::table.find(id)) - .set(agent_turns::turn_index.eq(position)) - .execute(connection)?; - diesel::update(llm_usage_events::table.filter(llm_usage_events::turn_pk.eq(id))) - .set(llm_usage_events::turn_index.eq(position)) - .execute(connection)?; - } - Ok(true) -} - -/// Loads matching events newest first. Without content, text and metadata are -/// not read: they hold whole request bodies, and an event list never shows them. +/// Loads matching events newest first. Without content, the conversation and +/// metadata are not read: the conversation holds whole request bodies, and an +/// event list shows neither. fn load_filtered_events( connection: &mut PgConnection, filters: &EventFilters<'_>, @@ -1102,9 +1105,6 @@ fn load_filtered_events( if let Some(session_pk) = filters.session_pk { query = query.filter(llm_usage_events::session_pk.eq(session_pk)); } - if let Some(turn_index) = filters.turn_index { - query = query.filter(llm_usage_events::turn_index.eq(turn_index)); - } let providers = exact_values_filter(filters.llm_providers); if !providers.is_empty() { query = query.filter(llm_usage_events::llm_provider.eq_any(providers)); @@ -1118,7 +1118,11 @@ fn load_filtered_events( } let query = query - .order(llm_usage_events::observed_at.desc()) + .order(( + llm_usage_events::observed_at.desc(), + llm_usage_events::event_type.asc(), + llm_usage_events::id.desc(), + )) .limit(filters.limit) .offset(filters.offset); if include_content { @@ -1130,17 +1134,13 @@ fn load_filtered_events( llm_usage_events::user_id, llm_usage_events::device_context_id, llm_usage_events::session_pk, - llm_usage_events::turn_pk, llm_usage_events::event_id, llm_usage_events::agent_name, llm_usage_events::agent_version, llm_usage_events::session_id, - llm_usage_events::turn_index, llm_usage_events::llm_provider, llm_usage_events::llm_model, llm_usage_events::event_type, - None::.into_sql::>(), - None::>.into_sql::>(), llm_usage_events::input_tokens, llm_usage_events::output_tokens, llm_usage_events::cache_read_tokens, @@ -1153,6 +1153,7 @@ fn load_filtered_events( llm_usage_events::stream_id, llm_usage_events::response_id, llm_usage_events::previous_response_id, + None::.into_sql::>(), )) .load::(connection) .map_err(AppError::from) @@ -1234,16 +1235,13 @@ fn event_response( host_name: device.map(|value| value.host_name.clone()), platform: device.map(|value| value.platform.clone()), session_pk: event.session_pk, - turn_pk: event.turn_pk, + location: None, agent_name: event.agent_name, agent_version: event.agent_version, session_id: event.session_id, - turn_index: event.turn_index, llm_provider: event.llm_provider, llm_model: event.llm_model, event_type: event.event_type, - text: event.text, - text_sha256: event.text_sha256.map(hex::encode), input_tokens: event.input_tokens, output_tokens: event.output_tokens, cache_read_tokens: event.cache_read_tokens, @@ -1254,6 +1252,7 @@ fn event_response( stream_id: event.stream_id, response_id: event.response_id, previous_response_id: event.previous_response_id, + conversation: event.conversation, metadata: event.metadata, attachments, } @@ -1272,19 +1271,6 @@ const fn touch_device_metadata(device: &Device) { )); } -const fn touch_session_metadata(session: &AgentSession) { - std::hint::black_box((session.created_at, session.updated_at, &session.metadata)); -} - -const fn touch_turn_metadata(turn: &AgentTurn) { - std::hint::black_box(( - turn.user_id, - turn.session_pk, - turn.created_at, - turn.updated_at, - )); -} - #[cfg(test)] fn latest_optional( left: Option>, @@ -1306,7 +1292,6 @@ struct EventFilters<'a> { agent_names: &'a [String], session_ids: &'a [String], session_pk: Option, - turn_index: Option, llm_providers: &'a [String], llm_models: &'a [String], event_type: Option<&'a str>, @@ -1371,7 +1356,6 @@ struct SummarySqlFilters { agent_names: Vec, session_ids: Vec, session_pk: Option, - turn_index: Option, providers: Vec, models: Vec, event_type: Option, @@ -1398,7 +1382,6 @@ impl SummarySqlFilters { agent_names, session_ids, session_pk: filters.session_pk, - turn_index: filters.turn_index, providers, models, event_type, @@ -1507,7 +1490,7 @@ mod tests { error::AppError, usage::{ AgentPayload, DevicePayload, IngestEventsRequest, IngestUsageEvent, LlmPayload, - SummaryRow, TokenUsagePayload, UsageEventType, + SummaryRow, TokenUsagePayload, UsageEventType, streams::RequestFacts, }, }; @@ -1567,7 +1550,12 @@ mod tests { let key = summary_key(event, device_contexts, user_contexts, group_by); let accumulator = accumulators.entry(key).or_default(); accumulator.sessions.insert(event.session_pk); - accumulator.turns.insert(event.turn_pk); + // Ingest stores the same fact for every request it reads. + if event.event_type == "request" + && RequestFacts::read(event.conversation.as_ref()).starts_turn + { + accumulator.turns.insert(event.id); + } if event.event_type == "request" { accumulator.requests = accumulator.requests.saturating_add(1); } @@ -1719,7 +1707,6 @@ mod tests { agent_names: &[" claude-desktop ".to_owned(), "codex".to_owned()], session_ids: &[" session-1 ".to_owned(), "session-2".to_owned()], session_pk: None, - turn_index: None, llm_providers: &["OpenAI".to_owned(), " ".to_owned()], llm_models: &[ " gpt-5.5 ".to_owned(), @@ -1745,7 +1732,6 @@ mod tests { assert_eq!(sql_filters.agent_names, ["claude-desktop", "codex"]); assert_eq!(sql_filters.session_ids, ["session-1", "session-2"]); assert_eq!(sql_filters.session_pk, None); - assert_eq!(sql_filters.turn_index, None); assert_eq!(sql_filters.providers, ["OpenAI"]); assert_eq!(sql_filters.models, ["gpt-5.5", "o3"]); assert_eq!(sql_filters.event_type.as_deref(), Some("response")); @@ -1812,7 +1798,7 @@ mod tests { #[test] fn aggregate_summary_counts_sessions_turns_events_and_tokens() { let device_context_id = fixed_uuid(1_u128); - let events = vec![ + let mut events = vec![ usage_event( "evt-summary-request", "request", @@ -1830,6 +1816,11 @@ mod tests { 15_i64, ), ]; + // The request asks in words, so it opens a turn. + events[0].conversation = Some(json!({ + "entries": [{"role": "user", "text": "ask"}], + "completeness": "complete" + })); let device_contexts = HashMap::from([(device_context_id, device_context(device_context_id))]); let group_by = vec![ @@ -1863,7 +1854,7 @@ mod tests { ); assert_eq!( row.turns, 1_usize, - "summary row should count distinct turns" + "summary row should count turn-opening requests" ); assert_eq!( row.requests, 1_i64, @@ -1983,17 +1974,18 @@ mod tests { version: Some("1.0.0".to_owned()), }, session_id: "session-001".to_owned(), - turn_id: None, - naming_thread: false, stream_id: None, response_id: None, previous_response_id: None, + subagent_id: None, + parent_subagent_id: None, + turn_id: None, llm: LlmPayload { provider: "OpenAI".to_owned(), model: "gpt-5.5".to_owned(), }, event_type: UsageEventType::Request, - text: Some("hello".to_owned()), + conversation: None, token_usage: TokenUsagePayload { input_tokens: Some(5_i64), output_tokens: Some(7_i64), @@ -2020,17 +2012,13 @@ mod tests { user_id: fixed_uuid(100_u128), device_context_id: fixed_uuid(1_u128), session_pk: fixed_uuid(2_u128), - turn_pk: fixed_uuid(3_u128), event_id: event_id.to_owned(), agent_name: "codex".to_owned(), agent_version: Some("1.0.0".to_owned()), session_id: "session-001".to_owned(), - turn_index: 1_i32, llm_provider: "openai".to_owned(), llm_model: "gpt-5.5".to_owned(), event_type: event_type.to_owned(), - text: None, - text_sha256: None, input_tokens, output_tokens, cache_read_tokens: 0_i64, @@ -2043,6 +2031,7 @@ mod tests { stream_id: None, response_id: None, previous_response_id: None, + conversation: None, } } diff --git a/crates/abyss-backend/src/storage/postgres/search/document.rs b/crates/abyss-backend/src/storage/postgres/search/document.rs index 1120039..5f8ab8f 100644 --- a/crates/abyss-backend/src/storage/postgres/search/document.rs +++ b/crates/abyss-backend/src/storage/postgres/search/document.rs @@ -1,20 +1,19 @@ -//! Search document projection and bounded extraction from usage-event metadata. +//! Search document projection and bounded extraction from usage-event +//! conversations. //! -//! Only explicitly allowlisted metadata enters the index. Extraction limits -//! depth, value count, and character length to bound CPU, allocation, request -//! size, and accidental indexing of unrelated collector metadata. PostgreSQL +//! Only conversation entries enter the index; metadata never does. Extraction +//! limits depth, value count, and character length to bound CPU, allocation, +//! and request size. PostgreSQL //! retains the complete source event and remains the recovery source. use serde::Serialize; use uuid::Uuid; -use crate::{ - db::models::UsageEvent, - search::projection::{SearchProjection, sanitize_search_text}, -}; +use super::projection::{SearchProjection, sanitize_search_text}; +use crate::db::models::UsageEvent; #[cfg(test)] -use crate::search::{HIGHLIGHT_END, HIGHLIGHT_START, projection::MAX_CONTENT_CHARACTERS}; +use super::projection::{HIGHLIGHT_END, HIGHLIGHT_START, MAX_CONTENT_CHARACTERS}; /// Full-text search document for one immutable usage event. #[derive(Serialize)] @@ -27,10 +26,6 @@ pub struct SearchDocument { pub session_pk: Uuid, /// Agent-native session identifier searchable as text and keyword. pub session_id: String, - /// Parent turn primary key. - pub turn_pk: Uuid, - /// Backend-normalized turn number. - pub turn_index: i32, /// Agent name as reported by the collector, used for exact filtering. pub agent_name: String, /// LLM provider as reported by the collector, used for exact filtering. @@ -44,7 +39,7 @@ pub struct SearchDocument { /// Bounded prompt or response text. #[serde(skip_serializing_if = "Option::is_none")] pub content: Option, - /// Bounded names extracted from tool-call metadata. + /// Bounded names extracted from tool-call entries. #[serde(skip_serializing_if = "Vec::is_empty")] pub tool_names: Vec, /// Bounded raw tool inputs and outputs. @@ -56,14 +51,12 @@ impl SearchDocument { /// Projects one source event into the strict Elasticsearch mapping. #[must_use] pub fn from_event(event: UsageEvent) -> Self { - let projection = SearchProjection::from_source(event.text, &event.metadata); + let projection = SearchProjection::from_source(event.conversation.as_ref()); Self { event_pk: event.id, user_id: event.user_id, session_pk: event.session_pk, session_id: sanitize_search_text(&event.session_id), - turn_pk: event.turn_pk, - turn_index: event.turn_index, agent_name: event.agent_name, llm_provider: event.llm_provider, llm_model: event.llm_model, @@ -86,46 +79,39 @@ mod tests { use super::{HIGHLIGHT_END, HIGHLIGHT_START, MAX_CONTENT_CHARACTERS, SearchDocument}; + /// `entries` as an event's complete conversation. + fn conversation(entries: &serde_json::Value) -> serde_json::Value { + json!({"entries": entries, "completeness": "complete"}) + } + #[test] - fn extracts_only_searchable_tool_metadata() { + fn indexes_every_entry_text_and_the_called_tool_names_only() { let mut event = usage_event(); - event.metadata = json!({ - "authorization": "Bearer secret", - "content_segments": [ - { - "type": "tool_call", - "name": "Bash", - "input": "{\"command\":\"cargo test\",\"nested\":{\"file_path\":\"src/lib.rs\"}}" - }, - { - "type": "tool_result", - "output": "all tests passed" - }, - { - "type": "image", - "content_base64": "secret-image" - } - ] - }); + event.conversation = Some(conversation(&json!([ + {"role": "user", "text": "run the tests"}, + {"role": "tool_call", "text": "{\"command\":\"cargo test\"}", "call_id": "c1", "name": "Bash"}, + {"role": "tool_result", "text": "all tests passed", "call_id": "c1"} + ]))); + event.metadata["headers"] = json!([{"name": "authorization", "value": "Bearer secret"}]); let document = SearchDocument::from_event(event); assert_eq!(document.tool_names, ["Bash"]); - // Tool inputs are searchable as written; no field is pulled out by name. assert_eq!( - document.tool_content[0], - "{\"command\":\"cargo test\",\"nested\":{\"file_path\":\"src/lib.rs\"}}" + document.content.as_deref(), + Some("run the tests\n{\"command\":\"cargo test\"}\nall tests passed") ); - assert_eq!(document.tool_content.len(), 2); + assert_eq!(document.tool_content, Vec::::new()); let serialized = serde_json::to_string(&document).expect("document should serialize"); assert!(!serialized.contains("Bearer secret")); - assert!(!serialized.contains("secret-image")); } #[test] fn truncates_oversized_content_on_a_utf8_boundary() { let mut event = usage_event(); - event.text = Some("会".repeat(MAX_CONTENT_CHARACTERS + 1)); + event.conversation = Some(conversation(&json!([ + {"role": "user", "text": "会".repeat(MAX_CONTENT_CHARACTERS + 1)} + ]))); let document = SearchDocument::from_event(event); @@ -143,26 +129,15 @@ mod tests { fn strips_highlight_markers_from_every_searchable_field() { let mut event = usage_event(); event.session_id = format!("session-{HIGHLIGHT_START}one{HIGHLIGHT_END}"); - event.text = Some(format!( - "before {HIGHLIGHT_START}match{HIGHLIGHT_END} after" - )); - event.metadata = json!({ - "content_segments": [ - { - "type": "tool_call", - "name": format!("{HIGHLIGHT_START}Bash{HIGHLIGHT_END}"), - "input": serde_json::to_string(&json!({ - "command": format!("cargo {HIGHLIGHT_START}test{HIGHLIGHT_END}"), - "file_path": format!("src/{HIGHLIGHT_END}lib.rs") - })) - .expect("tool input should serialize") - }, - { - "type": "tool_result", - "output": format!("{HIGHLIGHT_START}passed{HIGHLIGHT_END}") - } - ] - }); + event.conversation = Some(conversation(&json!([ + {"role": "user", "text": format!("before {HIGHLIGHT_START}match{HIGHLIGHT_END} after")}, + { + "role": "tool_call", + "text": format!("cargo {HIGHLIGHT_START}test{HIGHLIGHT_END}"), + "call_id": "c1", + "name": format!("{HIGHLIGHT_START}Bash{HIGHLIGHT_END}") + } + ]))); let serialized = serde_json::to_string(&SearchDocument::from_event(event)) .expect("document should serialize"); @@ -177,17 +152,13 @@ mod tests { user_id: Uuid::now_v7(), device_context_id: Uuid::now_v7(), session_pk: Uuid::now_v7(), - turn_pk: Uuid::now_v7(), event_id: "event-1".to_owned(), agent_name: "codex".to_owned(), agent_version: None, session_id: "session-1".to_owned(), - turn_index: 1, llm_provider: "openai".to_owned(), llm_model: "gpt-5".to_owned(), event_type: "request".to_owned(), - text: Some("hello".to_owned()), - text_sha256: None, input_tokens: 1, output_tokens: 0, cache_read_tokens: 0, @@ -200,6 +171,7 @@ mod tests { stream_id: None, response_id: None, previous_response_id: None, + conversation: None, } } } diff --git a/crates/abyss-backend/src/storage/postgres/search/elasticsearch.rs b/crates/abyss-backend/src/storage/postgres/search/elasticsearch.rs index 8757c4b..1d02340 100644 --- a/crates/abyss-backend/src/storage/postgres/search/elasticsearch.rs +++ b/crates/abyss-backend/src/storage/postgres/search/elasticsearch.rs @@ -13,13 +13,12 @@ use serde::Deserialize; use serde_json::{Value, json}; use uuid::Uuid; -use crate::{ - config::SearchConfig, - error::AppError, - search::{HIGHLIGHT_END, HIGHLIGHT_START, ValidatedSearchQuery}, -}; +use crate::{config::SearchConfig, error::AppError, search::ValidatedSearchQuery}; -use super::document::SearchDocument; +use super::{ + document::SearchDocument, + projection::{HIGHLIGHT_END, HIGHLIGHT_START}, +}; /// Fixed name of the derived usage-event index. pub const SEARCH_INDEX: &str = "abyss_usage_events"; @@ -240,10 +239,6 @@ pub struct SessionMatches { pub struct MatchedEvent { /// Backend event primary key. pub event_pk: Uuid, - /// Backend turn primary key. - pub turn_pk: Uuid, - /// Normalized turn number. - pub turn_index: i32, /// Request or response side. pub event_type: String, /// LLM provider as reported by the collector. @@ -291,8 +286,6 @@ fn index_definition() -> Value { "user_id": {"type": "keyword"}, "session_pk": {"type": "keyword"}, "session_id": {"type": "text", "fields": {"keyword": {"type": "keyword"}}}, - "turn_pk": {"type": "keyword"}, - "turn_index": {"type": "integer"}, "agent_name": {"type": "keyword"}, "llm_provider": {"type": "keyword"}, "llm_model": {"type": "keyword"}, @@ -357,8 +350,6 @@ fn search_request(user_id: Uuid, query: &ValidatedSearchQuery) -> Value { "size": 3, "_source": [ "event_pk", - "turn_pk", - "turn_index", "event_type", "llm_provider", "llm_model", @@ -541,8 +532,6 @@ struct RawMatchedHit { #[derive(Deserialize)] struct RawMatchedSource { event_pk: Uuid, - turn_pk: Uuid, - turn_index: i32, event_type: String, llm_provider: String, llm_model: String, @@ -573,8 +562,6 @@ impl TryFrom for SearchMatchPage { .iter() .map(|event| MatchedEvent { event_pk: event.source.event_pk, - turn_pk: event.source.turn_pk, - turn_index: event.source.turn_index, event_type: event.source.event_type.clone(), llm_provider: event.source.llm_provider.clone(), llm_model: event.source.llm_model.clone(), @@ -673,7 +660,6 @@ mod tests { fn parses_collapsed_hits_and_highlights() { let session_pk = Uuid::now_v7(); let event_pk = Uuid::now_v7(); - let turn_pk = Uuid::now_v7(); let response = serde_json::from_value::(json!({ "aggregations": {"session_count": {"value": 1_u64}}, "hits": {"hits": [{ @@ -683,8 +669,6 @@ mod tests { "hits": [{ "_source": { "event_pk": event_pk, - "turn_pk": turn_pk, - "turn_index": 4_i32, "event_type": "response", "llm_provider": "openai", "llm_model": "gpt-5", diff --git a/crates/abyss-backend/src/storage/postgres/search/mod.rs b/crates/abyss-backend/src/storage/postgres/search/mod.rs index 1c2f4f3..67658fd 100644 --- a/crates/abyss-backend/src/storage/postgres/search/mod.rs +++ b/crates/abyss-backend/src/storage/postgres/search/mod.rs @@ -8,6 +8,7 @@ mod document; mod elasticsearch; mod outbox; +mod projection; mod worker; use std::collections::{HashMap, HashSet}; @@ -18,14 +19,15 @@ use crate::{ config::SearchConfig, error::AppError, search::{ - SearchFragment, SessionSearchMatch, SessionSearchQuery, SessionSearchResponse, - SessionSearchResult, ValidatedSearchQuery, + SessionSearchMatch, SessionSearchQuery, SessionSearchResponse, SessionSearchResult, + ValidatedSearchQuery, }, }; use self::{ elasticsearch::{ElasticsearchClient, SearchMatchPage}, outbox::SearchSessionDetails, + projection::parse_fragment, }; pub(super) use self::{outbox::SearchOutboxRepository, worker::SearchIndexer}; @@ -85,7 +87,9 @@ impl SearchExecution { /// Combines search matches with authoritative session/device details. /// /// Stale index entries whose session no longer exists are omitted instead - /// of returning partially authorized or incomplete data. + /// of returning partially authorized or incomplete data. Elasticsearch + /// indexes an event's text as one document and cannot tell a new hit from + /// one in the inherited context, so every match counts as new. #[must_use] pub(super) fn hydrate( self, @@ -116,6 +120,7 @@ impl SearchExecution { Some(SessionSearchResult { session_pk: details.session.id, session_id: details.session.session_id, + title: details.title, agent_name: details.session.agent_name, agent_version: details.session.agent_version, host_name: details.device.host_name, @@ -125,21 +130,22 @@ impl SearchExecution { providers, models, match_count: matches.match_count, + new_match_count: matches.match_count, matches: matches .events .into_iter() .map(|event| SessionSearchMatch { event_pk: event.event_pk, - turn_pk: event.turn_pk, - turn_index: event.turn_index, + location: None, event_type: event.event_type, llm_provider: event.llm_provider, llm_model: event.llm_model, observed_at: event.observed_at, + new: true, fragments: event .fragments - .into_iter() - .map(|fragment| SearchFragment::parse(&fragment)) + .iter() + .map(|fragment| parse_fragment(fragment)) .collect(), }) .collect(), diff --git a/crates/abyss-backend/src/storage/postgres/search/outbox.rs b/crates/abyss-backend/src/storage/postgres/search/outbox.rs index f9187a7..827f83c 100644 --- a/crates/abyss-backend/src/storage/postgres/search/outbox.rs +++ b/crates/abyss-backend/src/storage/postgres/search/outbox.rs @@ -24,6 +24,7 @@ use crate::{ schema::{agent_sessions, devices, llm_usage_events, search_outbox}, }, error::AppError, + usage::streams::EntryLabel, }; use super::{document::SearchDocument, elasticsearch::BulkOperation}; @@ -59,6 +60,8 @@ pub struct SearchSessionDetails { pub session: AgentSession, /// Device row associated with the session. pub device: Device, + /// Label of the session's earliest request, read from its existing facts. + pub title: Option, } /// Synchronous persistence boundary for search projection state. @@ -265,9 +268,24 @@ impl SearchOutboxRepository { .filter(agent_sessions::id.eq_any(session_ids)) .select((AgentSession::as_select(), Device::as_select())) .load::<(AgentSession, Device)>(connection)?; + let authorized_ids = rows + .iter() + .map(|(session, _)| session.id) + .collect::>(); + let mut titles = super::super::session_titles::titles(connection, &authorized_ids)?; Ok(rows .into_iter() - .map(|(session, device)| (session.id, SearchSessionDetails { session, device })) + .map(|(session, device)| { + let session_pk = session.id; + ( + session_pk, + SearchSessionDetails { + session, + device, + title: titles.remove(&session_pk), + }, + ) + }) .collect()) } } diff --git a/crates/abyss-backend/src/storage/postgres/search/projection.rs b/crates/abyss-backend/src/storage/postgres/search/projection.rs new file mode 100644 index 0000000..c167cb1 --- /dev/null +++ b/crates/abyss-backend/src/storage/postgres/search/projection.rs @@ -0,0 +1,136 @@ +//! Bounded Elasticsearch projection of an event's conversation entries, and +//! the marker pair Elasticsearch places around highlighted text. + +use serde_json::Value; + +use crate::search::{SearchFragment, SearchFragmentSegment}; + +/// Sentinel Elasticsearch inserts before highlighted text. +pub const HIGHLIGHT_START: &str = "[[[ABYSS_SEARCH_HIGHLIGHT_START]]]"; +/// Sentinel Elasticsearch inserts after highlighted text. +pub const HIGHLIGHT_END: &str = "[[[ABYSS_SEARCH_HIGHLIGHT_END]]]"; + +pub const MAX_CONTENT_CHARACTERS: usize = 1_000_000; +const MAX_METADATA_FIELD_CHARACTERS: usize = 100_000; +const MAX_EXTRACTED_VALUES: usize = 64; + +/// Searchable values derived from one source usage event. +pub struct SearchProjection { + pub content: Option, + pub tool_names: Vec, + pub tool_content: Vec, +} + +impl SearchProjection { + /// Reads the event's conversation entries: the text of every entry, of + /// every role, in order, and the names of the tools it calls. Metadata is + /// never read, so headers, credentials, and image data never enter a + /// derived full-text index. + pub fn from_source(conversation: Option<&Value>) -> Self { + let entries = conversation + .and_then(|conversation| conversation.get("entries")) + .and_then(Value::as_array) + .map(Vec::as_slice) + .unwrap_or_default(); + let mut tool_names = Vec::new(); + let mut texts = Vec::new(); + for entry in entries { + if entry.get("role").and_then(Value::as_str) == Some("tool_call") + && let Some(name) = entry.get("name").and_then(Value::as_str) + { + push_bounded(&mut tool_names, name); + } + if let Some(text) = entry.get("text").and_then(Value::as_str) + && !text.is_empty() + { + texts.push(text); + } + } + Self { + content: (!texts.is_empty()) + .then(|| sanitize_search_text(&texts.join("\n"))) + .map(|value| truncate_characters(value, MAX_CONTENT_CHARACTERS)), + tool_names, + tool_content: Vec::new(), + } + } +} + +/// Splits an Elasticsearch fragment at its highlight markers. +pub fn parse_fragment(fragment: &str) -> SearchFragment { + let mut segments = Vec::new(); + let mut remainder = fragment; + while let Some(start) = remainder.find(HIGHLIGHT_START) { + let (plain, tagged) = remainder.split_at(start); + push_fragment_segment(&mut segments, plain, false); + let highlighted = tagged + .strip_prefix(HIGHLIGHT_START) + .expect("the marker position came from find"); + // Treat malformed or truncated marker pairs as plain text. The API + // never emits raw HTML and therefore does not trust ES fragments. + let Some(end) = highlighted.find(HIGHLIGHT_END) else { + push_fragment_segment(&mut segments, tagged, false); + remainder = ""; + break; + }; + let (matched, after_match) = highlighted.split_at(end); + push_fragment_segment(&mut segments, matched, true); + remainder = after_match + .strip_prefix(HIGHLIGHT_END) + .expect("the marker position came from find"); + } + push_fragment_segment(&mut segments, remainder, false); + SearchFragment { segments } +} + +fn push_fragment_segment(segments: &mut Vec, text: &str, highlighted: bool) { + if text.is_empty() { + return; + } + segments.push(SearchFragmentSegment { + text: text.to_owned(), + highlighted, + }); +} + +fn push_bounded(values: &mut Vec, value: &str) { + let value = sanitize_search_text(value); + let value = value.trim(); + if value.is_empty() || values.len() >= MAX_EXTRACTED_VALUES { + return; + } + let value = truncate_characters(value.to_owned(), MAX_METADATA_FIELD_CHARACTERS); + if !values.contains(&value) { + values.push(value); + } +} + +pub fn sanitize_search_text(value: &str) -> String { + value + .replace(HIGHLIGHT_START, "") + .replace(HIGHLIGHT_END, "") +} + +fn truncate_characters(mut value: String, maximum: usize) -> String { + let Some((byte_index, _character)) = value.char_indices().nth(maximum) else { + return value; + }; + value.truncate(byte_index); + value +} + +#[cfg(test)] +mod tests { + use super::{HIGHLIGHT_END, HIGHLIGHT_START, parse_fragment}; + + #[test] + fn parses_highlights_without_rendering_html() { + let fragment = parse_fragment(&format!( + "before {HIGHLIGHT_START}timeout{HIGHLIGHT_END} after" + )); + + assert_eq!(fragment.segments.len(), 3); + assert_eq!(fragment.segments[1].text, "timeout"); + assert!(fragment.segments[1].highlighted); + } +} diff --git a/crates/abyss-backend/src/storage/postgres/session_list.rs b/crates/abyss-backend/src/storage/postgres/session_list.rs index 112ef4d..6d5aa23 100644 --- a/crates/abyss-backend/src/storage/postgres/session_list.rs +++ b/crates/abyss-backend/src/storage/postgres/session_list.rs @@ -3,13 +3,16 @@ use chrono::{DateTime, Utc}; use diesel::{ PgConnection, QueryableByName, RunQueryDsl, sql_query, - sql_types::{Array, BigInt, Bool, Nullable, Text, Timestamptz, Uuid as SqlUuid}, + sql_types::{Array, BigInt, Bool, Integer, Nullable, Text, Timestamptz, Uuid as SqlUuid}, }; use uuid::Uuid; use crate::{ error::AppError, - usage::session_list::{SessionListItem, SessionListResponse, ValidatedSessionListQuery}, + usage::{ + session_list::{SessionListItem, SessionListResponse, ValidatedSessionListQuery}, + streams::EntryLabel, + }, }; /// Per-session activity with the filters that apply before paging. @@ -29,14 +32,27 @@ const ACTIVITY_CTE: &str = "\ s.started_at AS started_at, \ s.ended_at AS ended_at, \ COALESCE(MAX(e.observed_at), s.started_at) AS last_activity_at, \ - (SELECT COUNT(*) FROM agent_turns t WHERE t.session_pk = s.id)::bigint AS turns, \ + (SELECT CASE WHEN COUNT(tf.turn_id) > 0 THEN COUNT(DISTINCT tf.turn_id) \ + ELSE COUNT(*) FILTER (WHERE tf.starts_turn AND tf.continues_context) \ + + COALESCE((SELECT CASE WHEN ff.continues_context THEN 0 ELSE 1 END \ + FROM llm_usage_events fe \ + JOIN llm_usage_event_stream_facts ff ON ff.event_pk = fe.id \ + WHERE fe.session_pk = s.id AND ff.starts_turn \ + ORDER BY fe.observed_at, fe.id LIMIT 1), 0) END \ + FROM llm_usage_events te \ + JOIN llm_usage_event_stream_facts tf ON tf.event_pk = te.id \ + WHERE te.session_pk = s.id)::bigint AS turns, \ COUNT(e.id)::bigint AS events, \ COALESCE(SUM(e.total_tokens), 0)::bigint AS total_tokens, \ - st.title AS title, \ + st.title_event_pk AS title_event_pk, \ + title_facts.label_index AS title_entry_index, \ + title_facts.label_text AS title_text, \ + title_facts.label_length AS title_length, \ st.purpose AS purpose \ FROM agent_sessions s \ JOIN devices d ON d.id = s.device_context_id \ LEFT JOIN session_titles st ON st.session_pk = s.id \ + LEFT JOIN llm_usage_event_stream_facts title_facts ON title_facts.event_pk = st.title_event_pk \ LEFT JOIN llm_usage_events e ON e.session_pk = s.id \ WHERE s.user_id = $1 \ AND ($2 = false OR s.agent_name = ANY($3)) \ @@ -46,7 +62,8 @@ const ACTIVITY_CTE: &str = "\ AND ($10 = false OR EXISTS ( \ SELECT 1 FROM llm_usage_events m \ WHERE m.session_pk = s.id AND m.llm_model = ANY($11))) \ - GROUP BY s.id, d.host_name, d.platform, st.title, st.purpose \ + GROUP BY s.id, d.host_name, d.platform, st.title_event_pk, \ + title_facts.label_index, title_facts.label_text, title_facts.label_length, st.purpose \ ), \ filtered AS ( \ SELECT * FROM activity \ @@ -80,8 +97,14 @@ struct SessionListRow { events: i64, #[diesel(sql_type = BigInt)] total_tokens: i64, + #[diesel(sql_type = Nullable)] + title_event_pk: Option, + #[diesel(sql_type = Nullable)] + title_entry_index: Option, #[diesel(sql_type = Nullable)] - title: Option, + title_text: Option, + #[diesel(sql_type = Nullable)] + title_length: Option, #[diesel(sql_type = Nullable)] purpose: Option, } @@ -152,7 +175,12 @@ pub(super) fn list_sessions( turns: row.turns, events: row.events, total_tokens: row.total_tokens, - title: row.title, + title: EntryLabel::from_columns( + row.title_event_pk, + row.title_entry_index, + row.title_text, + row.title_length, + ), purpose: row.purpose, }) .collect(); diff --git a/crates/abyss-backend/src/storage/postgres/session_titles.rs b/crates/abyss-backend/src/storage/postgres/session_titles.rs index f9686f0..af82ba6 100644 --- a/crates/abyss-backend/src/storage/postgres/session_titles.rs +++ b/crates/abyss-backend/src/storage/postgres/session_titles.rs @@ -1,157 +1,97 @@ -//! PostgreSQL persistence of session names reported by the collector. +//! The earliest request anchors a session's title to its stored stream label. -use chrono::{DateTime, Utc}; -use diesel::{ExpressionMethods, OptionalExtension, PgConnection, QueryDsl, RunQueryDsl}; -use serde_json::Value; -use uuid::Uuid; +use std::collections::HashMap; -use crate::{ - db::schema::{llm_usage_events, session_titles}, - error::AppError, - usage::{ - SessionInfo, - session_title::{TITLE_GENERATION_PURPOSE, TitleContribution, default_title}, - }, +use chrono::{DateTime, Utc}; +use diesel::{ + PgConnection, QueryableByName, RunQueryDsl, sql_query, + sql_types::{Array, Integer, Nullable, Text, Timestamptz, Uuid as SqlUuid}, }; +use uuid::Uuid; -/// Stored title, purpose, and matched task session of one session. -type StoredTitle = (Option, Option, Option); +use crate::{error::AppError, usage::streams::EntryLabel}; -/// One stored event considered for title capture. +/// One stored event considered for the session's first-request anchor. pub(super) struct CapturedEvent<'a> { pub(super) user_id: Uuid, pub(super) session_pk: Uuid, pub(super) event_pk: Uuid, pub(super) event_type: &'a str, - pub(super) metadata: &'a Value, - pub(super) naming_thread: bool, pub(super) observed_at: DateTime, } -/// Records what a newly stored event says about its session's name. -/// -/// An Agent title replaces a default title and an older Agent title. A default -/// title keeps the earliest user input, even after late delivery, and never -/// replaces an Agent title. +/// Registers the earliest request even when that request has no label. +/// A request delivered later can replace the anchor if its order is earlier. pub(super) fn capture( connection: &mut PgConnection, event: &CapturedEvent<'_>, -) -> Result<(), AppError> { - let session = event.session_pk; - let agent_title = - match TitleContribution::of(event.event_type, event.metadata, event.naming_thread) { - None => return Ok(()), - Some(TitleContribution::TitleThread) => { - return mark_title_thread(connection, session, event.user_id); - } - Some(TitleContribution::Agent(title)) => Some(title), - Some(TitleContribution::UserInput) => None, - }; - let stored = session_titles::table - .find(session) - .select(( - session_titles::title_event_pk, - session_titles::title_observed_at, - )) - .first::<(Option, Option>)>(connection) - .optional()? - .and_then(|(event_pk, observed_at)| event_pk.zip(observed_at)); - let agent = agent_title.is_some(); - let replaces = match stored { - None => true, - Some((_, stored_at)) if agent && stored_at < event.observed_at => true, - Some((_, stored_at)) if !agent && stored_at <= event.observed_at => false, - // Otherwise only a stored default title gives way. - Some((event_pk, _)) => !is_agent_title(connection, event_pk)?, - }; - if !replaces { - return Ok(()); +) -> Result { + if event.event_type != "request" { + return Ok(false); } - // Read the user input only once it could name the session. - let Some(title) = agent_title - .map(str::to_owned) - .or_else(|| default_title(event.metadata)) - else { - return Ok(()); - }; - let now = Utc::now(); - diesel::insert_into(session_titles::table) - .values(( - session_titles::session_pk.eq(session), - session_titles::user_id.eq(event.user_id), - session_titles::title.eq(&title), - session_titles::title_event_pk.eq(event.event_pk), - session_titles::title_observed_at.eq(event.observed_at), - session_titles::updated_at.eq(now), - )) - .on_conflict(session_titles::session_pk) - .do_update() - .set(( - session_titles::title.eq(&title), - session_titles::title_event_pk.eq(event.event_pk), - session_titles::title_observed_at.eq(event.observed_at), - session_titles::updated_at.eq(now), - )) - .execute(connection)?; - Ok(()) + let changed = sql_query( + "INSERT INTO session_titles \ + (session_pk, user_id, title, title_event_pk, title_observed_at, updated_at) \ + VALUES ($1, $2, NULL, $3, $4, CURRENT_TIMESTAMP) \ + ON CONFLICT (session_pk) DO UPDATE SET \ + title = NULL, title_event_pk = EXCLUDED.title_event_pk, \ + title_observed_at = EXCLUDED.title_observed_at, updated_at = EXCLUDED.updated_at \ + WHERE session_titles.title_event_pk IS NULL \ + OR session_titles.title_observed_at IS NULL \ + OR (EXCLUDED.title_observed_at, EXCLUDED.title_event_pk) \ + < (session_titles.title_observed_at, session_titles.title_event_pk)", + ) + .bind::(event.session_pk) + .bind::(event.user_id) + .bind::(event.event_pk) + .bind::(event.observed_at) + .execute(connection)?; + Ok(changed == 1) } -/// Agent titles are read from responses and default titles from requests, so -/// the side of the event that gave the stored title tells them apart. -fn is_agent_title(connection: &mut PgConnection, event_pk: Uuid) -> Result { - Ok(llm_usage_events::table - .find(event_pk) - .select(llm_usage_events::event_type) - .first::(connection) - .optional()? - .is_some_and(|event_type| event_type == "response")) +#[derive(QueryableByName)] +struct StoredLabel { + #[diesel(sql_type = SqlUuid)] + session_pk: Uuid, + #[diesel(sql_type = Nullable)] + event_pk: Option, + #[diesel(sql_type = Nullable)] + entry_index: Option, + #[diesel(sql_type = Nullable)] + text: Option, + #[diesel(sql_type = Nullable)] + length: Option, } -/// Marks a session as a thread that names another session. -fn mark_title_thread( +/// Reads the anchor request's existing label for both the header and picker. +pub(super) fn title( connection: &mut PgConnection, - session: Uuid, - user_id: Uuid, -) -> Result<(), AppError> { - let now = Utc::now(); - diesel::insert_into(session_titles::table) - .values(( - session_titles::session_pk.eq(session), - session_titles::user_id.eq(user_id), - session_titles::purpose.eq(TITLE_GENERATION_PURPOSE), - session_titles::updated_at.eq(now), - )) - .on_conflict(session_titles::session_pk) - .do_update() - .set(( - session_titles::purpose.eq(TITLE_GENERATION_PURPOSE), - session_titles::updated_at.eq(now), - )) - .execute(connection)?; - Ok(()) + session_pk: Uuid, +) -> Result, AppError> { + Ok(titles(connection, &[session_pk])?.remove(&session_pk)) } -/// Adds the stored title, purpose, and matched task session to a session. -pub(super) fn fill( +/// Reads labels for an already authorized page of sessions in one query. +pub(super) fn titles( connection: &mut PgConnection, - session: &mut SessionInfo, -) -> Result<(), AppError> { - (session.title, session.purpose, session.parent_session_pk) = - load(connection, session.session_pk)?; - Ok(()) -} - -/// Returns the stored title and purpose of one session, and for a naming -/// thread the task session it was matched to. -fn load(connection: &mut PgConnection, session_pk: Uuid) -> Result { - Ok(session_titles::table - .find(session_pk) - .select(( - session_titles::title, - session_titles::purpose, - session_titles::parent_session_pk, - )) - .first::<(Option, Option, Option)>(connection) - .optional()? - .unwrap_or((None, None, None))) + session_pks: &[Uuid], +) -> Result, AppError> { + if session_pks.is_empty() { + return Ok(HashMap::new()); + } + Ok(sql_query( + "SELECT st.session_pk, st.title_event_pk AS event_pk, f.label_index AS entry_index, \ + f.label_text AS text, f.label_length AS length \ + FROM session_titles st \ + LEFT JOIN llm_usage_event_stream_facts f ON f.event_pk = st.title_event_pk \ + WHERE st.session_pk = ANY($1)", + ) + .bind::, _>(session_pks) + .load::(connection)? + .into_iter() + .filter_map(|row| { + EntryLabel::from_columns(row.event_pk, row.entry_index, row.text, row.length) + .map(|label| (row.session_pk, label)) + }) + .collect()) } diff --git a/crates/abyss-backend/src/storage/postgres/stream_facts.rs b/crates/abyss-backend/src/storage/postgres/stream_facts.rs new file mode 100644 index 0000000..067e1bb --- /dev/null +++ b/crates/abyss-backend/src/storage/postgres/stream_facts.rs @@ -0,0 +1,54 @@ +//! Stream facts of PostgreSQL events read at ingest from their conversations. + +use serde::Serialize; +use serde_json::Value; +use uuid::Uuid; + +use crate::{ + db::models::NewStreamFacts, + error::AppError, + usage::{ + UsageEventType, + streams::{RequestFacts, ResponseFacts}, + }, +}; + +fn json_value(value: &T, field: &str) -> Result { + serde_json::to_value(value) + .map_err(|error| AppError::internal(format!("serialize {field}: {error}"))) +} + +/// Fact columns read from an event's conversation, by its type; the other +/// side's columns and the ids the client declared stay NULL. +pub(super) fn read( + event_pk: Uuid, + event_type: &str, + conversation: Option<&Value>, + names_previous: bool, +) -> Result { + Ok(if event_type == UsageEventType::Request.as_str() { + let facts = RequestFacts::read(conversation); + NewStreamFacts { + event_pk, + new_from: Some(facts.new_from), + continues_context: Some(facts.new_from > 0_i32 || names_previous), + starts_turn: Some(facts.starts_turn), + last_identity: facts.last_identity, + previous_identity: facts.previous_identity, + trailing: Some(json_value(&facts.trailing, "trailing roles")?), + result_ids: Some(json_value(&facts.result_ids, "result ids")?), + label_index: facts.label.as_ref().map(|label| label.index), + label_length: facts.label.as_ref().map(|label| label.length), + label_text: facts.label.map(|label| label.text), + ..NewStreamFacts::default() + } + } else { + let facts = ResponseFacts::read(conversation); + NewStreamFacts { + event_pk, + emitted_ids: Some(json_value(&facts.emitted_ids, "emitted ids")?), + output_skeleton: Some(json_value(&facts.output_skeleton, "output skeleton")?), + ..NewStreamFacts::default() + } + }) +} diff --git a/crates/abyss-backend/src/storage/sqlite/day_offsets.rs b/crates/abyss-backend/src/storage/sqlite/day_offsets.rs new file mode 100644 index 0000000..b0a8f24 --- /dev/null +++ b/crates/abyss-backend/src/storage/sqlite/day_offsets.rs @@ -0,0 +1,170 @@ +//! UTC offsets that let SQLite group events by a time zone's calendar days. +//! +//! SQLite has no time zone database, so the offsets a zone applies across the +//! summarized events are worked out here and handed to the query as JSON. + +use chrono::{DateTime, Offset, TimeZone, Utc}; +use chrono_tz::Tz; + +use crate::error::AppError; + +const MICROS_PER_SECOND: i64 = 1_000_000; +/// Seconds in one day; no zone changes its offset more than once a day. +const DAY_SECONDS: i64 = 86_400; + +/// One UTC offset of a zone and the instant it takes effect. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct UtcOffsetChange { + /// Unix second from which the offset applies. + start_second: i64, + /// Local time minus UTC, in seconds. + offset_seconds: i32, +} + +/// The offsets `zone` applies to events observed from `first` to `last` +/// (inclusive, in microseconds), as a JSON list of `[start, offset]` +/// microsecond pairs ordered by start. The first pair also covers every +/// earlier instant; `None` means no events, which UTC's single offset covers. +pub(super) fn day_offsets_json(zone: Tz, events: Option<(i64, i64)>) -> Result { + let changes = match events { + Some((first, last)) if zone != Tz::UTC => utc_offset_changes( + zone, + first.div_euclid(MICROS_PER_SECOND), + last.div_euclid(MICROS_PER_SECOND).saturating_add(1), + ), + _ => Vec::new(), + }; + let mut offsets = changes + .iter() + .map(|change| { + [ + change.start_second.saturating_mul(MICROS_PER_SECOND), + i64::from(change.offset_seconds).saturating_mul(MICROS_PER_SECOND), + ] + }) + .collect::>(); + match offsets.first_mut() { + Some(first) => first[0] = i64::MIN, + None => offsets.push([i64::MIN, 0_i64]), + } + serde_json::to_string(&offsets) + .map_err(|error| AppError::internal(format!("serialize day offsets: {error}"))) +} + +/// The offsets `zone` applies from `from_second` up to `to_second`, in order. +/// +/// The range is stepped a day at a time and each change found is narrowed to +/// the exact second it takes effect. +fn utc_offset_changes(zone: Tz, from_second: i64, to_second: i64) -> Vec { + let mut current = offset_at(zone, from_second); + let mut changes = vec![UtcOffsetChange { + start_second: from_second, + offset_seconds: current, + }]; + let mut cursor = from_second; + while cursor < to_second { + let next = cursor.saturating_add(DAY_SECONDS).min(to_second); + if offset_at(zone, next) == current { + cursor = next; + continue; + } + // `low` keeps the current offset and `high` does not. + let (mut low, mut high) = (cursor, next); + while high.saturating_sub(low) > 1_i64 { + let middle = low.midpoint(high); + if offset_at(zone, middle) == current { + low = middle; + } else { + high = middle; + } + } + current = offset_at(zone, high); + changes.push(UtcOffsetChange { + start_second: high, + offset_seconds: current, + }); + cursor = high; + } + changes +} + +fn offset_at(zone: Tz, unix_second: i64) -> i32 { + DateTime::::from_timestamp(unix_second, 0).map_or(0_i32, |instant| { + zone.offset_from_utc_datetime(&instant.naive_utc()) + .fix() + .local_minus_utc() + }) +} + +#[cfg(test)] +mod tests { + use chrono::DateTime; + use chrono_tz::Tz; + + use super::{UtcOffsetChange, day_offsets_json, utc_offset_changes}; + + fn second(value: &str) -> i64 { + DateTime::parse_from_rfc3339(value) + .expect("timestamp should parse") + .timestamp() + } + + #[test] + fn zone_without_changes_keeps_one_offset() { + assert_eq!( + utc_offset_changes( + Tz::Asia__Shanghai, + second("2026-01-01T00:00:00Z"), + second("2026-12-31T00:00:00Z"), + ), + [UtcOffsetChange { + start_second: second("2026-01-01T00:00:00Z"), + offset_seconds: 8 * 3600, + }] + ); + } + + #[test] + fn daylight_saving_changes_start_at_their_exact_second() { + // Los Angeles moves to PDT at 2026-03-08 02:00 PST and back to PST at + // 2026-11-01 02:00 PDT. + assert_eq!( + utc_offset_changes( + Tz::America__Los_Angeles, + second("2026-02-01T00:00:00Z"), + second("2026-12-01T00:00:00Z"), + ), + [ + UtcOffsetChange { + start_second: second("2026-02-01T00:00:00Z"), + offset_seconds: -8 * 3600, + }, + UtcOffsetChange { + start_second: second("2026-03-08T10:00:00Z"), + offset_seconds: -7 * 3600, + }, + UtcOffsetChange { + start_second: second("2026-11-01T09:00:00Z"), + offset_seconds: -8 * 3600, + }, + ] + ); + } + + #[test] + fn offsets_cover_every_instant_and_default_to_utc() { + let minimum = i64::MIN; + assert_eq!( + day_offsets_json(Tz::UTC, Some((0, 1))).expect("utc offsets"), + format!("[[{minimum},0]]") + ); + assert_eq!( + day_offsets_json(Tz::Asia__Shanghai, None).expect("no events"), + format!("[[{minimum},0]]") + ); + assert_eq!( + day_offsets_json(Tz::Asia__Shanghai, Some((0, 1))).expect("shanghai offsets"), + format!("[[{minimum},28800000000]]") + ); + } +} diff --git a/crates/abyss-backend/src/storage/sqlite/mod.rs b/crates/abyss-backend/src/storage/sqlite/mod.rs index 70fd552..d1ac7d4 100644 --- a/crates/abyss-backend/src/storage/sqlite/mod.rs +++ b/crates/abyss-backend/src/storage/sqlite/mod.rs @@ -1,7 +1,9 @@ -//! Self-contained SQLite source-of-truth with transactional FTS5 projection. +//! Self-contained SQLite source-of-truth with its search index written in +//! the same transactions. mod comparison; mod connection; +mod day_offsets; mod migrations; mod models; mod repository; @@ -10,29 +12,41 @@ mod search; mod session_list; mod session_titles; -use diesel::{QueryDsl, RunQueryDsl, SqliteConnection, dsl::count_star}; +use std::sync::Arc; + +use diesel::{Connection, QueryDsl, RunQueryDsl, SqliteConnection, dsl::count_star}; use uuid::Uuid; use crate::{ config::Config, error::AppError, - search::{SessionSearchQuery, SessionSearchResponse}, - storage::{BoxedFuture, StorageBackend}, + search::{ + InSessionSearchQuery, InSessionSearchResponse, SessionSearchQuery, SessionSearchResponse, + }, + storage::{ + BoxedFuture, StorageBackend, + event_locations::{locate_events, locate_search_matches}, + stream_cache::StreamCache, + }, usage::{ - IngestEventsRequest, IngestEventsResponse, RawEventsQuery, RawEventsResponse, - SessionOutlineResponse, SessionTimelineResponse, SummaryQuery, SummaryResponse, - UsageEventResponse, + IngestEventsRequest, IngestEventsResponse, RawEventsQuery, RawEventsResponse, SummaryQuery, + SummaryResponse, UsageEventResponse, attachments::StoredImageAttachment, comparison::{ComparisonQuery, ComparisonResponse}, + preview::{ConversationPreview, Window}, + segments::{SegmentResponse, ValidatedSegmentQuery}, + session_header::SessionHeaderResponse, session_list::{SessionListResponse, ValidatedSessionListQuery}, + streams::{EventLocationResponse, SessionStreamsResponse, SessionTurnsResponse}, }, }; use self::connection::SqlitePool; -/// SQLite event store with an FTS5 projection committed beside each event. +/// SQLite event store with a search index committed beside each event. pub(super) struct SqliteFtsBackend { pool: SqlitePool, + structures: Arc, } impl SqliteFtsBackend { @@ -45,7 +59,10 @@ impl SqliteFtsBackend { migrations::run(&mut database)?; } } - Ok(Self { pool }) + Ok(Self { + pool, + structures: Arc::new(StreamCache::default()), + }) } async fn run_db(&self, task_fn: F) -> Result @@ -96,9 +113,20 @@ impl StorageBackend for SqliteFtsBackend { query: RawEventsQuery, default_page_size: i64, ) -> BoxedFuture<'_, Result> { + let structures = Arc::clone(&self.structures); Box::pin(async move { self.run_db(move |connection| { - repository::raw_events(connection, &query, user_id, default_page_size) + if !query.locate { + return repository::raw_events(connection, &query, user_id, default_page_size); + } + connection.transaction(|connection| { + let mut response = + repository::raw_events(connection, &query, user_id, default_page_size)?; + locate_events(&mut response.events, |session_pk| { + repository::session_structure(connection, user_id, session_pk, &structures) + })?; + Ok(response) + }) }) .await }) @@ -109,9 +137,34 @@ impl StorageBackend for SqliteFtsBackend { user_id: Uuid, event_pk: Uuid, ) -> BoxedFuture<'_, Result> { + let structures = Arc::clone(&self.structures); Box::pin(async move { - self.run_db(move |connection| repository::usage_event(connection, user_id, event_pk)) - .await + self.run_db(move |connection| { + connection.transaction(|connection| { + let mut event = repository::usage_event(connection, user_id, event_pk)?; + locate_events(std::slice::from_mut(&mut event), |session_pk| { + repository::session_structure(connection, user_id, session_pk, &structures) + })?; + Ok(event) + }) + }) + .await + }) + } + + fn event_location( + &self, + user_id: Uuid, + event_pk: Uuid, + ) -> BoxedFuture<'_, Result> { + let structures = Arc::clone(&self.structures); + Box::pin(async move { + self.run_db(move |connection| { + connection.transaction(|connection| { + repository::event_location(connection, user_id, event_pk, &structures) + }) + }) + .await }) } @@ -140,27 +193,88 @@ impl StorageBackend for SqliteFtsBackend { }) } - fn session_timeline( + fn session_header( + &self, + user_id: Uuid, + session_pk: Uuid, + ) -> BoxedFuture<'_, Result> { + let structures = Arc::clone(&self.structures); + Box::pin(async move { + self.run_db(move |connection| { + connection.transaction(|connection| { + repository::session_header(connection, user_id, session_pk, &structures) + }) + }) + .await + }) + } + + fn session_streams( + &self, + user_id: Uuid, + session_pk: Uuid, + ) -> BoxedFuture<'_, Result> { + let structures = Arc::clone(&self.structures); + Box::pin(async move { + self.run_db(move |connection| { + connection.transaction(|connection| { + repository::session_streams(connection, user_id, session_pk, &structures) + }) + }) + .await + }) + } + + fn session_turns( &self, user_id: Uuid, session_pk: Uuid, - ) -> BoxedFuture<'_, Result> { + ) -> BoxedFuture<'_, Result> { + let structures = Arc::clone(&self.structures); Box::pin(async move { self.run_db(move |connection| { - repository::session_timeline(connection, user_id, session_pk) + connection.transaction(|connection| { + repository::session_turns(connection, user_id, session_pk, &structures) + }) }) .await }) } - fn session_outline( + fn session_segment( &self, user_id: Uuid, session_pk: Uuid, - ) -> BoxedFuture<'_, Result> { + head: String, + query: ValidatedSegmentQuery, + ) -> BoxedFuture<'_, Result> { + let structures = Arc::clone(&self.structures); Box::pin(async move { self.run_db(move |connection| { - repository::session_outline(connection, user_id, session_pk) + connection.transaction(|connection| { + repository::session_segment( + connection, + user_id, + session_pk, + &head, + &query, + &structures, + ) + }) + }) + .await + }) + } + + fn event_conversation( + &self, + user_id: Uuid, + event_pk: Uuid, + window: Window, + ) -> BoxedFuture<'_, Result> { + Box::pin(async move { + self.run_db(move |connection| { + repository::event_conversation(connection, user_id, event_pk, &window) }) .await }) @@ -195,9 +309,33 @@ impl StorageBackend for SqliteFtsBackend { user_id: Uuid, query: SessionSearchQuery, ) -> BoxedFuture<'_, Result> { + let structures = Arc::clone(&self.structures); Box::pin(async move { - self.run_db(move |connection| search::session_search(connection, user_id, query)) - .await + self.run_db(move |connection| { + connection.transaction(|connection| { + let mut response = search::session_search(connection, user_id, query)?; + locate_search_matches(&mut response, |session_pk| { + repository::session_structure(connection, user_id, session_pk, &structures) + })?; + Ok(response) + }) + }) + .await + }) + } + + fn search_in_session( + &self, + user_id: Uuid, + session_pk: Uuid, + query: InSessionSearchQuery, + ) -> BoxedFuture<'_, Result> { + let structures = Arc::clone(&self.structures); + Box::pin(async move { + self.run_db(move |connection| { + search::search_in_session(connection, user_id, session_pk, &query, &structures) + }) + .await }) } diff --git a/crates/abyss-backend/src/storage/sqlite/models.rs b/crates/abyss-backend/src/storage/sqlite/models.rs index 0bc692f..34d533a 100644 --- a/crates/abyss-backend/src/storage/sqlite/models.rs +++ b/crates/abyss-backend/src/storage/sqlite/models.rs @@ -2,14 +2,22 @@ use chrono::{DateTime, Utc}; use diesel::{Insertable, Queryable, Selectable}; +use serde::{Serialize, de::DeserializeOwned}; use serde_json::Value; use uuid::Uuid; -use crate::{error::AppError, usage::event_order::TimelineEvent}; +use crate::{ + error::AppError, + usage::{ + UsageEventType, + streams::{EventFacts, Label, RequestFacts, ResponseFacts, StreamEvent}, + }, +}; use super::schema::{ - agent_diagnostic_capture_events, agent_diagnostic_captures, agent_sessions, agent_turns, - devices, llm_usage_event_attachments, llm_usage_events, + agent_diagnostic_capture_events, agent_diagnostic_captures, agent_sessions, devices, + llm_usage_event_attachments, llm_usage_event_stream_facts, llm_usage_events, search_refs, + search_texts, }; #[derive(Queryable, Selectable)] @@ -46,16 +54,9 @@ pub(super) struct NewDevice { #[diesel(check_for_backend(diesel::sqlite::Sqlite))] pub(super) struct SessionRow { pub(super) id: String, - pub(super) user_id: String, - pub(super) device_context_id: String, - pub(super) agent_name: String, pub(super) agent_version: Option, - pub(super) session_id: String, pub(super) started_at: i64, pub(super) ended_at: Option, - pub(super) metadata: String, - pub(super) created_at: i64, - pub(super) updated_at: i64, } #[derive(Insertable)] @@ -72,35 +73,8 @@ pub(super) struct NewSession { pub(super) metadata: String, pub(super) created_at: i64, pub(super) updated_at: i64, -} - -#[derive(Queryable, Selectable)] -#[diesel(table_name = agent_turns)] -#[diesel(check_for_backend(diesel::sqlite::Sqlite))] -pub(super) struct TurnRow { - pub(super) id: String, - pub(super) user_id: String, - pub(super) session_pk: String, - pub(super) turn_index: i32, - pub(super) turn_key: Option, - pub(super) started_at: i64, - pub(super) ended_at: Option, - pub(super) created_at: i64, - pub(super) updated_at: i64, -} - -#[derive(Insertable)] -#[diesel(table_name = agent_turns)] -pub(super) struct NewTurn { - pub(super) id: String, - pub(super) user_id: String, - pub(super) session_pk: String, - pub(super) turn_index: i32, - pub(super) turn_key: Option, - pub(super) started_at: i64, - pub(super) ended_at: Option, - pub(super) created_at: i64, - pub(super) updated_at: i64, + /// Changes whenever ingestion adds an event to the session. + pub(super) structure_revision: i64, } #[derive(Clone, Queryable, Selectable)] @@ -111,17 +85,13 @@ pub(super) struct EventRow { pub(super) user_id: String, pub(super) device_context_id: String, pub(super) session_pk: String, - pub(super) turn_pk: String, pub(super) event_id: String, pub(super) agent_name: String, pub(super) agent_version: Option, pub(super) session_id: String, - pub(super) turn_index: i32, pub(super) llm_provider: String, pub(super) llm_model: String, pub(super) event_type: String, - pub(super) text: Option, - pub(super) text_sha256: Option>, pub(super) input_tokens: i64, pub(super) output_tokens: i64, pub(super) cache_read_tokens: i64, @@ -134,6 +104,7 @@ pub(super) struct EventRow { pub(super) stream_id: Option, pub(super) response_id: Option, pub(super) previous_response_id: Option, + pub(super) conversation: Option, } #[derive(Insertable)] @@ -143,17 +114,13 @@ pub(super) struct NewEvent { pub(super) user_id: String, pub(super) device_context_id: String, pub(super) session_pk: String, - pub(super) turn_pk: String, pub(super) event_id: String, pub(super) agent_name: String, pub(super) agent_version: Option, pub(super) session_id: String, - pub(super) turn_index: i32, pub(super) llm_provider: String, pub(super) llm_model: String, pub(super) event_type: String, - pub(super) text: Option, - pub(super) text_sha256: Option>, pub(super) input_tokens: i64, pub(super) output_tokens: i64, pub(super) cache_read_tokens: i64, @@ -166,6 +133,156 @@ pub(super) struct NewEvent { pub(super) stream_id: Option, pub(super) response_id: Option, pub(super) previous_response_id: Option, + pub(super) conversation: Option, +} + +/// Stream facts of a new event; the other side's columns stay NULL. +#[derive(Default, Insertable)] +#[diesel(table_name = llm_usage_event_stream_facts)] +pub(super) struct NewStreamFacts { + pub(super) event_pk: String, + pub(super) new_from: Option, + pub(super) starts_turn: Option, + pub(super) last_identity: Option, + pub(super) trailing: Option, + pub(super) result_ids: Option, + pub(super) emitted_ids: Option, + pub(super) output_skeleton: Option, + pub(super) label_index: Option, + pub(super) label_text: Option, + pub(super) label_length: Option, + pub(super) subagent_id: Option, + pub(super) parent_subagent_id: Option, + pub(super) turn_id: Option, + pub(super) continues_context: Option, + pub(super) previous_identity: Option, +} + +impl NewStreamFacts { + /// Columns of the facts read from an event's conversation, by its type, + /// and whether the event names the response it continues; without the + /// ids the client declared. + pub(super) fn read( + event_pk: String, + event_type: &str, + conversation: Option<&Value>, + names_previous: bool, + ) -> Result { + if event_type == UsageEventType::Request.as_str() { + Self::request(event_pk, &RequestFacts::read(conversation), names_previous) + } else { + Self::response(event_pk, &ResponseFacts::read(conversation)) + } + } + + /// Columns of a request's facts. + pub(super) fn request( + event_pk: String, + facts: &RequestFacts, + names_previous: bool, + ) -> Result { + Ok(Self { + event_pk, + new_from: Some(facts.new_from), + continues_context: Some(facts.new_from > 0_i32 || names_previous), + starts_turn: Some(facts.starts_turn), + last_identity: facts.last_identity.clone(), + previous_identity: facts.previous_identity.clone(), + trailing: Some(json_text(&facts.trailing, "trailing roles")?), + result_ids: Some(json_text(&facts.result_ids, "result ids")?), + label_index: facts.label.as_ref().map(|label| label.index), + label_text: facts.label.as_ref().map(|label| label.text.clone()), + label_length: facts.label.as_ref().map(|label| label.length), + ..Self::default() + }) + } + + /// Columns of a response's facts. + pub(super) fn response(event_pk: String, facts: &ResponseFacts) -> Result { + Ok(Self { + event_pk, + emitted_ids: Some(json_text(&facts.emitted_ids, "emitted ids")?), + output_skeleton: Some(json_text(&facts.output_skeleton, "output skeleton")?), + ..Self::default() + }) + } +} + +/// The columns of one session event that its streams read; the facts are +/// NULL for an event stored before they were read. +#[derive(Queryable)] +pub(super) struct StreamEventRow { + pub(super) id: String, + pub(super) event_id: String, + pub(super) event_type: String, + pub(super) observed_at: i64, + pub(super) response_id: Option, + pub(super) previous_response_id: Option, + pub(super) new_from: Option, + pub(super) starts_turn: Option, + pub(super) last_identity: Option, + pub(super) previous_identity: Option, + pub(super) trailing: Option, + pub(super) result_ids: Option, + pub(super) emitted_ids: Option, + pub(super) output_skeleton: Option, + pub(super) label_index: Option, + pub(super) label_text: Option, + pub(super) label_length: Option, + pub(super) subagent_id: Option, + pub(super) parent_subagent_id: Option, + pub(super) turn_id: Option, +} + +impl StreamEventRow { + pub(super) fn into_stream_event(self) -> Result { + let facts = if self.event_type == "request" { + EventFacts::Request(RequestFacts { + new_from: self.new_from.unwrap_or_default(), + starts_turn: self.starts_turn.unwrap_or_default(), + last_identity: self.last_identity, + previous_identity: self.previous_identity, + trailing: json_list(self.trailing.as_deref(), "trailing roles")?, + result_ids: json_list(self.result_ids.as_deref(), "result ids")?, + label: Label::from_columns(self.label_index, self.label_text, self.label_length), + }) + } else { + EventFacts::Response(ResponseFacts { + emitted_ids: json_list(self.emitted_ids.as_deref(), "emitted ids")?, + output_skeleton: json_list(self.output_skeleton.as_deref(), "output skeleton")?, + }) + }; + Ok(StreamEvent { + pk: parse_uuid(&self.id, "event id")?, + event_id: self.event_id, + observed_at: timestamp_from_micros(self.observed_at)?, + response_id: self.response_id, + previous_response_id: self.previous_response_id, + subagent_id: self.subagent_id, + parent_subagent_id: self.parent_subagent_id, + turn_id: self.turn_id, + facts, + }) + } +} + +/// A stored JSON list; absent reads as empty. +fn json_list( + text: Option<&str>, + field: &str, +) -> Result { + text.map_or_else( + || Ok(T::default()), + |text| { + serde_json::from_str(text) + .map_err(|error| AppError::internal(format!("stored {field} are invalid: {error}"))) + }, + ) +} + +fn json_text(value: &T, field: &str) -> Result { + serde_json::to_string(value) + .map_err(|error| AppError::internal(format!("serialize {field}: {error}"))) } #[derive(Queryable, Selectable)] @@ -219,48 +336,43 @@ pub(super) struct NewDiagnosticCaptureEvent { pub(super) event_pk: String, } -pub(super) struct DeviceRecord { - pub(super) id: Uuid, - pub(super) host_name: String, - pub(super) platform: String, +/// A distinct text of one owner. +#[derive(Insertable)] +#[diesel(table_name = search_texts)] +pub(super) struct NewSearchText<'a> { + pub(super) user_id: &'a str, + pub(super) sha256: &'a [u8], + pub(super) text: &'a str, } -pub(super) struct SessionRecord { - pub(super) id: Uuid, - pub(super) user_id: Uuid, - pub(super) device_context_id: Uuid, - pub(super) agent_name: String, - pub(super) agent_version: Option, - pub(super) session_id: String, - pub(super) started_at: DateTime, - pub(super) ended_at: Option>, - pub(super) metadata: Value, +/// One piece of one event, referring to its text. +#[derive(Insertable)] +#[diesel(table_name = search_refs)] +pub(super) struct NewSearchRef<'a> { + pub(super) event_pk: &'a str, + pub(super) entry_index: i32, + pub(super) field: &'static str, + pub(super) text_id: i64, } -pub(super) struct TurnRecord { +pub(super) struct DeviceRecord { pub(super) id: Uuid, - pub(super) turn_index: i32, - pub(super) started_at: DateTime, - pub(super) ended_at: Option>, + pub(super) host_name: String, + pub(super) platform: String, } -#[derive(Clone)] pub(super) struct EventRecord { pub(super) id: Uuid, pub(super) user_id: Uuid, pub(super) device_context_id: Uuid, pub(super) session_pk: Uuid, - pub(super) turn_pk: Uuid, pub(super) event_id: String, pub(super) agent_name: String, pub(super) agent_version: Option, pub(super) session_id: String, - pub(super) turn_index: i32, pub(super) llm_provider: String, pub(super) llm_model: String, pub(super) event_type: String, - pub(super) text: Option, - pub(super) text_sha256: Option>, pub(super) input_tokens: i64, pub(super) output_tokens: i64, pub(super) cache_read_tokens: i64, @@ -272,24 +384,7 @@ pub(super) struct EventRecord { pub(super) stream_id: Option, pub(super) response_id: Option, pub(super) previous_response_id: Option, -} - -impl TimelineEvent for EventRecord { - fn turn_index(&self) -> i32 { - self.turn_index - } - - fn observed_at(&self) -> DateTime { - self.observed_at - } - - fn event_type(&self) -> &str { - &self.event_type - } - - fn event_id(&self) -> &str { - &self.event_id - } + pub(super) conversation: Option, } pub(super) struct AttachmentRecord { @@ -318,40 +413,6 @@ impl DeviceRow { } } -impl SessionRow { - pub(super) fn into_record(self) -> Result { - let _ = (self.created_at, self.updated_at); - Ok(SessionRecord { - id: parse_uuid(&self.id, "session id")?, - user_id: parse_uuid(&self.user_id, "session user id")?, - device_context_id: parse_uuid(&self.device_context_id, "session device id")?, - agent_name: self.agent_name, - agent_version: self.agent_version, - session_id: self.session_id, - started_at: timestamp_from_micros(self.started_at)?, - ended_at: self.ended_at.map(timestamp_from_micros).transpose()?, - metadata: parse_json(&self.metadata, "session metadata")?, - }) - } -} - -impl TurnRow { - pub(super) fn into_record(self) -> Result { - let _ = ( - &self.user_id, - &self.session_pk, - self.created_at, - self.updated_at, - ); - Ok(TurnRecord { - id: parse_uuid(&self.id, "turn id")?, - turn_index: self.turn_index, - started_at: timestamp_from_micros(self.started_at)?, - ended_at: self.ended_at.map(timestamp_from_micros).transpose()?, - }) - } -} - impl EventRow { pub(super) fn into_record(self) -> Result { let _ = self.created_at; @@ -360,17 +421,13 @@ impl EventRow { user_id: parse_uuid(&self.user_id, "event user id")?, device_context_id: parse_uuid(&self.device_context_id, "event device id")?, session_pk: parse_uuid(&self.session_pk, "event session id")?, - turn_pk: parse_uuid(&self.turn_pk, "event turn id")?, event_id: self.event_id, agent_name: self.agent_name, agent_version: self.agent_version, session_id: self.session_id, - turn_index: self.turn_index, llm_provider: self.llm_provider, llm_model: self.llm_model, event_type: self.event_type, - text: self.text, - text_sha256: self.text_sha256, input_tokens: self.input_tokens, output_tokens: self.output_tokens, cache_read_tokens: self.cache_read_tokens, @@ -382,6 +439,11 @@ impl EventRow { stream_id: self.stream_id, response_id: self.response_id, previous_response_id: self.previous_response_id, + conversation: self + .conversation + .as_deref() + .map(|conversation| parse_json(conversation, "event conversation")) + .transpose()?, }) } } diff --git a/crates/abyss-backend/src/storage/sqlite/repository.rs b/crates/abyss-backend/src/storage/sqlite/repository.rs index 1000bcd..c9280e3 100644 --- a/crates/abyss-backend/src/storage/sqlite/repository.rs +++ b/crates/abyss-backend/src/storage/sqlite/repository.rs @@ -1,49 +1,62 @@ //! Diesel ORM implementation of SQLite usage ingest and relational queries. -use std::collections::{HashMap, HashSet}; +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, +}; use chrono::{DateTime, Utc}; +use chrono_tz::Tz; use diesel::{ - ExpressionMethods, IntoSql, JoinOnDsl, QueryDsl, QueryableByName, RunQueryDsl, + ExpressionMethods, IntoSql, NullableExpressionMethods, QueryDsl, QueryableByName, RunQueryDsl, SelectableHelper, SqliteConnection, - dsl::{count_star, max, sql}, + dsl::{max, min}, result::OptionalExtension, sql_query, - sql_types::{BigInt, Binary, Bool, Nullable, Text}, + sql_types::{BigInt, Bool, Nullable, Text}, }; use uuid::Uuid; use crate::{ error::AppError, - search::projection::SearchProjection, + storage::stream_cache::StreamCache, usage::{ - IngestEventsRequest, IngestEventsResponse, IngestUsageEvent, OUTLINE_PREVIEW_CHARS, - RawEventsQuery, RawEventsResponse, SessionInfo, SessionOutlineResponse, - SessionTimelineResponse, SummaryQuery, SummaryResponse, SummaryRow, TurnTimeline, - UsageEventOutline, UsageEventResponse, + IngestEventsRequest, IngestEventsResponse, IngestUsageEvent, RawEventsQuery, + RawEventsResponse, SummaryQuery, SummaryResponse, SummaryRow, UsageEventResponse, attachments::{ ImageMediaType, StoredImageAttachment, UsageEventAttachmentResponse, ValidatedImageAttachment, }, diagnostics::IngestDiagnosticCapture, - event_order::UsageEventTimelineOrder, persistence::{ exact_values_filter, non_empty, normalize_limit, normalize_text, parse_group_by, - sha256_bytes, validate_batch, validate_event_type, + validate_batch, validate_event_type, + }, + preview::{ConversationPreview, Window, preview}, + segments::{ + SegmentResponse, TokenCounts, ValidatedSegmentQuery, run_event_pks, segment_response, }, + session_header::{HeaderEvent, SessionHeaderResponse, session_header as header_of}, + streams::{ + EventLocationResponse, SessionStreamsResponse, SessionStructure, SessionTurnsResponse, + StreamEvent, + }, + time_zone::parse_time_zone, }, }; use super::{ + day_offsets::day_offsets_json, models::{ AttachmentRow, DeviceRecord, DeviceRow, EventRecord, EventRow, NewAttachment, NewDevice, - NewDiagnosticCapture, NewDiagnosticCaptureEvent, NewEvent, NewSession, NewTurn, - SessionRecord, SessionRow, TurnRecord, TurnRow, parse_uuid, timestamp_from_micros, + NewDiagnosticCapture, NewDiagnosticCaptureEvent, NewEvent, NewSession, NewStreamFacts, + SessionRow, StreamEventRow, parse_uuid, timestamp_from_micros, }, schema::{ - agent_diagnostic_capture_events, agent_diagnostic_captures, agent_sessions, agent_turns, - app_users, devices, llm_usage_event_attachments, llm_usage_events, + agent_diagnostic_capture_events, agent_diagnostic_captures, agent_sessions, app_users, + devices, llm_usage_event_attachments, llm_usage_event_stream_facts, llm_usage_events, }, + search::index_event, session_titles, }; @@ -54,7 +67,7 @@ pub(super) fn ingest_events( max_batch_size: usize, ) -> Result { let validated_attachments = validate_batch(request, max_batch_size)?; - // Reserve the single SQLite writer before allocating turn indexes. + // Reserve the single SQLite writer before updating session metadata. let (accepted, duplicates, accepted_captures, duplicate_captures) = connection .immediate_transaction(|transaction| { let mut accepted = 0_usize; @@ -123,19 +136,14 @@ fn ingest_one_event( observed_at, now, )?; - let (turn_pk, turn_index) = place_in_turn( - connection, - user_id, - session_pk, - event.turn_id.as_deref(), - event.response_id.as_deref(), - observed_at, - now, - )?; let tokens = &event.token_usage; - let text = event.text.clone(); let event_pk = Uuid::now_v7(); let metadata = serialize_json(&event.metadata, "event metadata")?; + let conversation = event + .conversation + .as_ref() + .map(|conversation| serialize_json(conversation, "event conversation")) + .transpose()?; let inserted = diesel::insert_into(llm_usage_events::table) .values(NewEvent { @@ -143,17 +151,13 @@ fn ingest_one_event( user_id: user_id.to_string(), device_context_id: device_id.to_string(), session_pk: session_pk.to_string(), - turn_pk: turn_pk.to_string(), event_id, agent_name: agent_name.clone(), agent_version, session_id: event.session_id.clone(), - turn_index, llm_provider: provider, llm_model: model, event_type: event.event_type.as_str().to_owned(), - text: text.clone(), - text_sha256: text.as_deref().map(sha256_bytes), input_tokens: tokens.input_tokens.unwrap_or_default(), output_tokens: tokens.output_tokens.unwrap_or_default(), cache_read_tokens: tokens.cache_read_tokens.unwrap_or_default(), @@ -166,6 +170,7 @@ fn ingest_one_event( stream_id: event.stream_id.clone(), response_id: event.response_id.clone(), previous_response_id: event.previous_response_id.clone(), + conversation, }) .on_conflict(llm_usage_events::event_id) .do_nothing() @@ -174,8 +179,14 @@ fn ingest_one_event( return Ok(false); } + insert_stream_facts(connection, event_pk, event)?; insert_attachments(connection, user_id, event_pk, attachments, now)?; - insert_search_projection(connection, event_pk, user_id, session_pk, event)?; + index_event( + connection, + &user_id.to_string(), + &event_pk.to_string(), + event.conversation.as_ref(), + )?; session_titles::capture( connection, &session_titles::CapturedEvent { @@ -183,8 +194,6 @@ fn ingest_one_event( session_pk, event_pk, event_type: event.event_type.as_str(), - metadata: &event.metadata, - naming_thread: event.naming_thread, observed_at, }, )?; @@ -273,6 +282,10 @@ fn upsert_device( clippy::too_many_arguments, reason = "The session upsert receives one explicit value for every evolving boundary." )] +#[expect( + clippy::arithmetic_side_effects, + reason = "Diesel builds an atomic SQL increment, not a Rust integer addition." +)] fn upsert_session( connection: &mut SqliteConnection, user_id: Uuid, @@ -304,6 +317,7 @@ fn upsert_session( existing.ended_at.unwrap_or(observed_at).max(observed_at), )), agent_sessions::updated_at.eq(now), + agent_sessions::structure_revision.eq(agent_sessions::structure_revision + 1_i64), )) .execute(connection)?; return Ok(id); @@ -323,224 +337,307 @@ fn upsert_session( metadata: "{}".to_owned(), created_at: now, updated_at: now, + structure_revision: 1, }) .execute(connection)?; Ok(id) } -/// Places an event in a turn of its session and returns the turn's id and -/// position. -/// -/// Explicit turn ids take priority; otherwise each response id gets its own -/// fallback turn. Events with neither id join the turn in progress. -/// Turns are numbered by when they started, so a turn that arrives late -/// moves the turns after it. -fn place_in_turn( +/// Derives the streams of one session's events from their stream facts, +/// without reading any request body. +pub(super) fn session_streams( connection: &mut SqliteConnection, user_id: Uuid, session_pk: Uuid, - turn_id: Option<&str>, - response_id: Option<&str>, - observed_at: i64, - now: i64, -) -> Result<(Uuid, i32), AppError> { - let session = session_pk.to_string(); - let found = match (turn_id, response_id) { - (Some(key), _) => named_turn(connection, &session, key)?, - (None, Some(key)) => response_turn(connection, &session, key)?, - (None, None) => turn_in_progress(connection, &session, observed_at)?, - }; - let (turn_pk, turn_index, reordered) = if let Some(turn) = found { - diesel::update(agent_turns::table.find(&turn.id)) - .set(( - agent_turns::started_at.eq(turn.started_at.min(observed_at)), - agent_turns::ended_at - .eq(Some(turn.ended_at.unwrap_or(observed_at).max(observed_at))), - agent_turns::updated_at.eq(now), - )) - .execute(connection)?; - (turn.id, turn.turn_index, observed_at < turn.started_at) - } else { - let last = agent_turns::table - .filter(agent_turns::session_pk.eq(&session)) - .select(max(agent_turns::turn_index)) - .first::>(connection)?; - let id = Uuid::now_v7().to_string(); - let turn_index = last.unwrap_or(0_i32).saturating_add(1_i32); - diesel::insert_into(agent_turns::table) - .values(NewTurn { - id: id.clone(), - user_id: user_id.to_string(), - session_pk: session.clone(), - turn_index, - turn_key: turn_id.map(str::to_owned), - started_at: observed_at, - ended_at: Some(observed_at), - created_at: now, - updated_at: now, - }) - .execute(connection)?; - (id, turn_index, last.is_some()) - }; - let turn_index = if reordered && number_turns(connection, &session)? { - agent_turns::table - .find(&turn_pk) - .select(agent_turns::turn_index) - .first::(connection)? - } else { - turn_index - }; - Ok((parse_uuid(&turn_pk, "turn id")?, turn_index)) + cache: &StreamCache, +) -> Result { + let structure = session_structure(connection, user_id, session_pk, cache)?; + Ok(structure.session_streams(session_pk)) } -/// The turn named `key` in a session. The first named turn adopts the turn -/// that events with neither turn nor response id opened at the start of the -/// session. A response fallback turn keeps its separate identity. -fn named_turn( +/// Derives the turn view of one session from its stream facts, without +/// reading any request body. +pub(super) fn session_turns( connection: &mut SqliteConnection, - session: &str, - key: &str, -) -> Result, AppError> { - let named = agent_turns::table - .filter(agent_turns::session_pk.eq(session)) - .filter(agent_turns::turn_key.eq(key)) - .select(TurnRow::as_select()) - .first::(connection) - .optional()?; - if named.is_some() { - return Ok(named); - } - let mut turns = agent_turns::table - .filter(agent_turns::session_pk.eq(session)) - .select(TurnRow::as_select()) - .limit(2_i64) - .load::(connection)?; - let opening = match turns.as_slice() { - [turn] if turn.turn_key.is_none() => turns.pop(), - _ => None, - }; - if let Some(turn) = &opening { - let has_response = llm_usage_events::table - .filter(llm_usage_events::turn_pk.eq(&turn.id)) - .filter(llm_usage_events::response_id.is_not_null()) - .select(llm_usage_events::id) - .first::(connection) - .optional()? - .is_some(); - if has_response { - return Ok(None); + user_id: Uuid, + session_pk: Uuid, + cache: &StreamCache, +) -> Result { + let structure = session_structure(connection, user_id, session_pk, cache)?; + Ok(structure.session_turns(session_pk)) +} + +/// Uses the caller's read transaction so the revision and facts are one snapshot. +/// A cache hit reads only the session revision; it does not load the event facts. +pub(super) fn session_structure( + connection: &mut SqliteConnection, + user_id: Uuid, + session_pk: Uuid, + cache: &StreamCache, +) -> Result, AppError> { + let revision = agent_sessions::table + .filter(agent_sessions::user_id.eq(user_id.to_string())) + .filter(agent_sessions::id.eq(session_pk.to_string())) + .select(agent_sessions::structure_revision) + .first::(connection) + .optional()? + .ok_or_else(|| AppError::not_found(format!("session not found: {session_pk}")))?; + cache.get_or_load(user_id, session_pk, revision, || { + let events = stream_events(connection, session_pk)?; + if events.is_empty() { + return Err(AppError::not_found(format!( + "session not found: {session_pk}" + ))); } - diesel::update(agent_turns::table.find(&turn.id)) - .set(agent_turns::turn_key.eq(key)) - .execute(connection)?; - } - Ok(opening) + Ok(events) + }) } -/// Reuses the fallback turn for a model response within this session. -/// Explicit turn keys and response ids occupy separate identity domains. -fn response_turn( +/// Resolves an authorized event without reading its body or attachments. +pub(super) fn event_location( connection: &mut SqliteConnection, - session: &str, - response_id: &str, -) -> Result, AppError> { - agent_turns::table - .inner_join(llm_usage_events::table.on(llm_usage_events::turn_pk.eq(agent_turns::id))) - .filter(llm_usage_events::session_pk.eq(session)) - .filter(llm_usage_events::response_id.eq(response_id)) - .filter(agent_turns::turn_key.is_null()) - .select(TurnRow::as_select()) - .first::(connection) - .optional() - .map_err(AppError::from) + user_id: Uuid, + event_pk: Uuid, + cache: &StreamCache, +) -> Result { + let session_pk = llm_usage_events::table + .filter(llm_usage_events::user_id.eq(user_id.to_string())) + .filter(llm_usage_events::id.eq(event_pk.to_string())) + .select(llm_usage_events::session_pk) + .first::(connection) + .optional()? + .ok_or_else(|| AppError::not_found(format!("event not found: {event_pk}")))?; + let session_pk = parse_uuid(&session_pk, "session primary key")?; + session_structure(connection, user_id, session_pk, cache)? + .event_location(session_pk, event_pk) + .ok_or_else(|| AppError::not_found(format!("event not found: {event_pk}"))) } -/// The turn in progress when an event naming no turn was observed: the latest -/// turn started by then, else the first turn of the session. -fn turn_in_progress( +/// Columns of one session event the header reads. +type HeaderColumns = ( + String, + String, + Option, + String, + String, + i64, + i64, + i64, + i64, + i64, + i64, + i64, +); + +/// The header of an owner-scoped session. +pub(super) fn session_header( connection: &mut SqliteConnection, - session: &str, - observed_at: i64, -) -> Result, AppError> { - let started = agent_turns::table - .filter(agent_turns::session_pk.eq(session)) - .filter(agent_turns::started_at.le(observed_at)) - .order_by((agent_turns::started_at.desc(), agent_turns::id.desc())) - .select(TurnRow::as_select()) - .first::(connection) + user_id: Uuid, + session_pk: Uuid, + cache: &StreamCache, +) -> Result { + let owned = agent_sessions::table + .filter(agent_sessions::user_id.eq(user_id.to_string())) + .filter(agent_sessions::id.eq(session_pk.to_string())) + .select(agent_sessions::id) + .first::(connection) .optional()?; - if started.is_some() { - return Ok(started); + if owned.is_none() { + return Err(AppError::not_found(format!( + "session not found: {session_pk}" + ))); + } + // The owner is checked on the session row. Filtering events by the owner + // as well lets SQLite pick the owner's time index and read every event + // the owner has; the session alone uses the session index. + let events = llm_usage_events::table + .filter(llm_usage_events::session_pk.eq(session_pk.to_string())) + .select(( + llm_usage_events::session_id, + llm_usage_events::agent_name, + llm_usage_events::agent_version, + llm_usage_events::llm_provider, + llm_usage_events::llm_model, + llm_usage_events::observed_at, + llm_usage_events::input_tokens, + llm_usage_events::output_tokens, + llm_usage_events::cache_read_tokens, + llm_usage_events::cache_write_tokens, + llm_usage_events::reasoning_tokens, + llm_usage_events::total_tokens, + )) + .load::(connection)? + .into_iter() + .map(header_event) + .collect::, AppError>>()?; + if events.is_empty() { + return Err(AppError::not_found(format!( + "session not found: {session_pk}" + ))); } - agent_turns::table - .filter(agent_turns::session_pk.eq(session)) - .order_by((agent_turns::started_at.asc(), agent_turns::id.asc())) - .select(TurnRow::as_select()) - .first::(connection) - .optional() - .map_err(AppError::from) + let summary = session_structure(connection, user_id, session_pk, cache)?.summary(); + let title = session_titles::title(connection, session_pk)?; + Ok(header_of(session_pk, events, &summary, title)) } -/// Numbers the turns of a session, and their events, by when the turns -/// started. Writes nothing and returns `false` while the numbers already -/// follow that order. -fn number_turns(connection: &mut SqliteConnection, session: &str) -> Result { - let turns = agent_turns::table - .filter(agent_turns::session_pk.eq(session)) - .order_by((agent_turns::started_at.asc(), agent_turns::id.asc())) - .select((agent_turns::id, agent_turns::turn_index)) - .load::<(String, i32)>(connection)?; - let moved = turns - .iter() - .zip(1_i32..) - .filter(|((_, index), position)| index != position) - .collect::>(); - if moved.is_empty() { - return Ok(false); - } - // Moved turns first step above every current number, so no single-row - // update collides with another turn's number or goes below one. - let above = turns.iter().map(|(_, index)| *index).max().unwrap_or(0_i32); - for ((id, _), position) in &moved { - diesel::update(agent_turns::table.find(id)) - .set(agent_turns::turn_index.eq(above.saturating_add(*position))) - .execute(connection)?; - } - for ((id, _), position) in moved { - diesel::update(agent_turns::table.find(id)) - .set(agent_turns::turn_index.eq(position)) - .execute(connection)?; - diesel::update(llm_usage_events::table.filter(llm_usage_events::turn_pk.eq(id))) - .set(llm_usage_events::turn_index.eq(position)) - .execute(connection)?; - } - Ok(true) +fn header_event( + ( + session_id, + agent_name, + agent_version, + provider, + model, + observed_at, + input, + output, + cache_read, + cache_write, + reasoning, + total, + ): HeaderColumns, +) -> Result { + Ok(HeaderEvent { + session_id, + agent_name, + agent_version, + provider, + model, + observed_at: timestamp_from_micros(observed_at)?, + tokens: TokenCounts { + input, + output, + cache_read, + cache_write, + reasoning, + total, + }, + }) +} + +/// The exchanges of the segment holding exchange `head`, from it to the +/// segment end, with a preview of each event. +pub(super) fn session_segment( + connection: &mut SqliteConnection, + user_id: Uuid, + session_pk: Uuid, + head: &str, + query: &ValidatedSegmentQuery, + cache: &StreamCache, +) -> Result { + let structure = session_structure(connection, user_id, session_pk, cache)?; + let run = structure.segment_run(head).ok_or_else(|| { + AppError::not_found(format!( + "exchange not found in session {session_pk}: {head}" + )) + })?; + let events = usage_events(connection, user_id, &run_event_pks(&run))?; + Ok(segment_response(session_pk, run, events, &query.window)) } -/// Writes the SQLite FTS5 projection. Virtual tables are intentionally outside -/// Diesel's relational schema and use a fully bound raw statement. -fn insert_search_projection( +/// The conversation of one owner-scoped event through `window`. +pub(super) fn event_conversation( connection: &mut SqliteConnection, + user_id: Uuid, event_pk: Uuid, + window: &Window, +) -> Result { + usage_event(connection, user_id, event_pk).map(|event| preview(event, window)) +} + +/// Owner-scoped events by primary key, with their device labels and +/// attachments. +fn usage_events( + connection: &mut SqliteConnection, user_id: Uuid, + pks: &[Uuid], +) -> Result, AppError> { + let events = llm_usage_events::table + .filter(llm_usage_events::id.eq_any(pks.iter().map(Uuid::to_string))) + .filter(llm_usage_events::user_id.eq(user_id.to_string())) + .select(EventRow::as_select()) + .load::(connection)? + .into_iter() + .map(EventRow::into_record) + .collect::, _>>()?; + let devices = load_device_contexts(connection, &events)?; + let mut attachments = load_attachments(connection, &events)?; + Ok(events + .into_iter() + .map(|event| { + let (event_pk, device_id) = (event.id, event.device_context_id); + let attachments = attachments.remove(&event_pk).unwrap_or_default(); + ( + event_pk, + event_response(event, devices.get(&device_id), attachments), + ) + }) + .collect()) +} + +/// The events of one session in session order, with their stream facts and +/// without any request body. The caller has checked the session's owner on +/// its row, so events are read by session alone and use the session index. +pub(super) fn stream_events( + connection: &mut SqliteConnection, session_pk: Uuid, +) -> Result, AppError> { + llm_usage_events::table + .left_join(llm_usage_event_stream_facts::table) + .filter(llm_usage_events::session_pk.eq(session_pk.to_string())) + .order_by(( + llm_usage_events::observed_at.asc(), + llm_usage_events::event_type.asc(), + llm_usage_events::id.asc(), + )) + .select(( + llm_usage_events::id, + llm_usage_events::event_id, + llm_usage_events::event_type, + llm_usage_events::observed_at, + llm_usage_events::response_id, + llm_usage_events::previous_response_id, + llm_usage_event_stream_facts::new_from.nullable(), + llm_usage_event_stream_facts::starts_turn.nullable(), + llm_usage_event_stream_facts::last_identity.nullable(), + llm_usage_event_stream_facts::previous_identity.nullable(), + llm_usage_event_stream_facts::trailing.nullable(), + llm_usage_event_stream_facts::result_ids.nullable(), + llm_usage_event_stream_facts::emitted_ids.nullable(), + llm_usage_event_stream_facts::output_skeleton.nullable(), + llm_usage_event_stream_facts::label_index.nullable(), + llm_usage_event_stream_facts::label_text.nullable(), + llm_usage_event_stream_facts::label_length.nullable(), + llm_usage_event_stream_facts::subagent_id.nullable(), + llm_usage_event_stream_facts::parent_subagent_id.nullable(), + llm_usage_event_stream_facts::turn_id.nullable(), + )) + .load::(connection)? + .into_iter() + .map(StreamEventRow::into_stream_event) + .collect() +} + +/// Stores the stream facts read from one new event's own entries. +fn insert_stream_facts( + connection: &mut SqliteConnection, + event_pk: Uuid, event: &IngestUsageEvent, ) -> Result<(), AppError> { - let projection = SearchProjection::from_source(event.text.clone(), &event.metadata); - sql_query( - "INSERT INTO usage_events_fts ( - event_pk, user_id, session_pk, session_id, content, tool_names, - tool_content - ) VALUES (?, ?, ?, ?, ?, ?, ?)", - ) - .bind::(event_pk.to_string()) - .bind::(user_id.to_string()) - .bind::(session_pk.to_string()) - .bind::(&event.session_id) - .bind::, _>(projection.content) - .bind::(projection.tool_names.join("\n")) - .bind::(projection.tool_content.join("\n")) - .execute(connection)?; + let facts = NewStreamFacts::read( + event_pk.to_string(), + event.event_type.as_str(), + event.conversation.as_ref(), + event.previous_response_id.is_some(), + )?; + let facts = NewStreamFacts { + subagent_id: event.subagent_id.clone(), + parent_subagent_id: event.parent_subagent_id.clone(), + turn_id: event.turn_id.clone(), + ..facts + }; + diesel::insert_into(llm_usage_event_stream_facts::table) + .values(facts) + .execute(connection)?; Ok(()) } @@ -631,7 +728,6 @@ pub(super) fn raw_events( agent_names: &query.agent_name, session_ids: &query.session_id, session_pk: query.session_pk, - turn_index: query.turn_index, llm_providers: &query.llm_provider, llm_models: &query.llm_model, event_type: query.event_type.as_deref(), @@ -676,7 +772,6 @@ struct EventFilters<'a> { agent_names: &'a [String], session_ids: &'a [String], session_pk: Option, - turn_index: Option, llm_providers: &'a [String], llm_models: &'a [String], event_type: Option<&'a str>, @@ -684,8 +779,9 @@ struct EventFilters<'a> { offset: i64, } -/// Loads matching events newest first. Without content, text and metadata are -/// not read: they hold whole request bodies, and an event list never shows them. +/// Loads matching events newest first. Without content, the conversation and +/// metadata are not read: the conversation holds whole request bodies, and an +/// event list shows neither. fn load_filtered_events( connection: &mut SqliteConnection, filters: &EventFilters<'_>, @@ -711,9 +807,6 @@ fn load_filtered_events( if let Some(session_pk) = filters.session_pk { query = query.filter(llm_usage_events::session_pk.eq(session_pk.to_string())); } - if let Some(turn_index) = filters.turn_index { - query = query.filter(llm_usage_events::turn_index.eq(turn_index)); - } let providers = exact_values_filter(filters.llm_providers); if !providers.is_empty() { query = query.filter(llm_usage_events::llm_provider.eq_any(providers)); @@ -728,6 +821,7 @@ fn load_filtered_events( let query = query .order_by(( llm_usage_events::observed_at.desc(), + llm_usage_events::event_type.asc(), llm_usage_events::id.desc(), )) .limit(filters.limit) @@ -743,17 +837,13 @@ fn load_filtered_events( llm_usage_events::user_id, llm_usage_events::device_context_id, llm_usage_events::session_pk, - llm_usage_events::turn_pk, llm_usage_events::event_id, llm_usage_events::agent_name, llm_usage_events::agent_version, llm_usage_events::session_id, - llm_usage_events::turn_index, llm_usage_events::llm_provider, llm_usage_events::llm_model, llm_usage_events::event_type, - None::.into_sql::>(), - None::>.into_sql::>(), llm_usage_events::input_tokens, llm_usage_events::output_tokens, llm_usage_events::cache_read_tokens, @@ -766,243 +856,13 @@ fn load_filtered_events( llm_usage_events::stream_id, llm_usage_events::response_id, llm_usage_events::previous_response_id, + None::.into_sql::>(), )) .load::(connection)? }; rows.into_iter().map(EventRow::into_record).collect() } -pub(super) fn session_timeline( - connection: &mut SqliteConnection, - user_id: Uuid, - session_pk: Uuid, -) -> Result { - let (session, turns) = load_session_turns(connection, user_id, session_pk)?; - let mut events = llm_usage_events::table - .filter(llm_usage_events::session_pk.eq(session.id.to_string())) - .order_by(( - llm_usage_events::turn_index.asc(), - llm_usage_events::observed_at.asc(), - llm_usage_events::id.asc(), - )) - .select(EventRow::as_select()) - .load::(connection)? - .into_iter() - .map(EventRow::into_record) - .collect::, _>>()?; - UsageEventTimelineOrder::sort(&mut events); - let devices = load_device_contexts(connection, &events)?; - let mut attachments = load_attachments(connection, &events)?; - let mut events_by_turn: HashMap> = HashMap::new(); - for event in events { - let event_pk = event.id; - let device_id = event.device_context_id; - events_by_turn - .entry(event.turn_pk) - .or_default() - .push(event_response( - event, - devices.get(&device_id), - attachments.remove(&event_pk).unwrap_or_default(), - )); - } - let mut timeline = assemble_timeline(session, turns, events_by_turn); - session_titles::fill(connection, &mut timeline.session)?; - Ok(timeline) -} - -type OutlineRow = ( - String, - String, - String, - i32, - String, - String, - String, - i64, - i64, - i64, - i64, - i64, - i64, - i64, - i64, - Option, - bool, -); - -pub(super) fn session_outline( - connection: &mut SqliteConnection, - user_id: Uuid, - session_pk: Uuid, -) -> Result { - let (session, turns) = load_session_turns(connection, user_id, session_pk)?; - // Only fixed-size columns and SQL-computed summaries are selected, so the - // outline stays small no matter how large event text or metadata grows. - let rows = llm_usage_events::table - .filter(llm_usage_events::session_pk.eq(session.id.to_string())) - .select(( - llm_usage_events::id, - llm_usage_events::event_id, - llm_usage_events::turn_pk, - llm_usage_events::turn_index, - llm_usage_events::llm_provider, - llm_usage_events::llm_model, - llm_usage_events::event_type, - llm_usage_events::input_tokens, - llm_usage_events::output_tokens, - llm_usage_events::cache_read_tokens, - llm_usage_events::cache_write_tokens, - llm_usage_events::reasoning_tokens, - llm_usage_events::total_tokens, - llm_usage_events::observed_at, - sql::("coalesce(length(text), 0)"), - sql::>(&format!( - "CASE WHEN event_type = 'request' THEN substr(text, -{OUTLINE_PREVIEW_CHARS}) \ - ELSE substr(text, 1, {OUTLINE_PREVIEW_CHARS}) END" - )), - // SQLite text functions stop at the first NUL character, so the - // length and preview above are wrong for text containing NUL. - sql::("coalesce(instr(CAST(text AS BLOB), x'00') > 0, 0)"), - )) - .load::(connection)?; - let nul_text_ids = rows - .iter() - .filter(|row| row.16) - .map(|row| row.0.clone()) - .collect::>(); - let nul_texts = if nul_text_ids.is_empty() { - HashMap::new() - } else { - llm_usage_events::table - .filter(llm_usage_events::id.eq_any(nul_text_ids)) - .select((llm_usage_events::id, llm_usage_events::text)) - .load::<(String, Option)>(connection)? - .into_iter() - .filter_map(|(id, text)| text.map(|text| (id, text))) - .collect::>() - }; - let event_ids = rows.iter().map(|row| row.0.clone()).collect::>(); - let attachment_counts = llm_usage_event_attachments::table - .filter(llm_usage_event_attachments::event_pk.eq_any(event_ids)) - .group_by(llm_usage_event_attachments::event_pk) - .select((llm_usage_event_attachments::event_pk, count_star())) - .load::<(String, i64)>(connection)? - .into_iter() - .collect::>(); - let mut events = rows - .into_iter() - .map(|mut row| { - let attachment_count = attachment_counts.get(&row.0).copied().unwrap_or(0); - if let Some(text) = nul_texts.get(&row.0) { - (row.14, row.15) = text_outline(text, &row.6)?; - } - Ok(UsageEventOutline { - id: parse_uuid(&row.0, "event id")?, - event_id: row.1, - turn_pk: parse_uuid(&row.2, "event turn id")?, - turn_index: row.3, - llm_provider: row.4, - llm_model: row.5, - event_type: row.6, - input_tokens: row.7, - output_tokens: row.8, - cache_read_tokens: row.9, - cache_write_tokens: row.10, - reasoning_tokens: row.11, - total_tokens: row.12, - observed_at: timestamp_from_micros(row.13)?, - text_chars: row.14, - preview: row.15, - attachment_count, - }) - }) - .collect::, AppError>>()?; - UsageEventTimelineOrder::sort(&mut events); - let mut events_by_turn: HashMap> = HashMap::new(); - for event in events { - events_by_turn.entry(event.turn_pk).or_default().push(event); - } - let mut timeline = assemble_timeline(session, turns, events_by_turn); - session_titles::fill(connection, &mut timeline.session)?; - Ok(timeline) -} - -/// Computes the outline length and preview in Rust for text SQLite cannot measure. -fn text_outline(text: &str, event_type: &str) -> Result<(i64, Option), AppError> { - let chars = text.chars().count(); - let preview_chars = usize::try_from(OUTLINE_PREVIEW_CHARS) - .map_err(|error| AppError::internal(format!("invalid outline preview size: {error}")))?; - let preview = if event_type == "request" { - text.chars() - .skip(chars.saturating_sub(preview_chars)) - .collect() - } else { - text.chars().take(preview_chars).collect() - }; - let chars = i64::try_from(chars) - .map_err(|error| AppError::internal(format!("invalid event text length: {error}")))?; - Ok((chars, Some(preview))) -} - -fn load_session_turns( - connection: &mut SqliteConnection, - user_id: Uuid, - session_pk: Uuid, -) -> Result<(SessionRecord, Vec), AppError> { - let session = agent_sessions::table - .filter(agent_sessions::id.eq(session_pk.to_string())) - .filter(agent_sessions::user_id.eq(user_id.to_string())) - .select(SessionRow::as_select()) - .first::(connection) - .optional()? - .ok_or_else(|| AppError::not_found(format!("session not found: {session_pk}")))? - .into_record()?; - let turns = agent_turns::table - .filter(agent_turns::session_pk.eq(session.id.to_string())) - .order_by(agent_turns::turn_index.asc()) - .select(TurnRow::as_select()) - .load::(connection)? - .into_iter() - .map(TurnRow::into_record) - .collect::, _>>()?; - Ok((session, turns)) -} - -fn assemble_timeline( - session: SessionRecord, - turns: Vec, - mut events_by_turn: HashMap>, -) -> SessionTimelineResponse { - let turn_timelines = turns - .into_iter() - .map(|turn| TurnTimeline { - turn_pk: turn.id, - turn_index: turn.turn_index, - started_at: turn.started_at, - ended_at: turn.ended_at, - events: events_by_turn.remove(&turn.id).unwrap_or_default(), - }) - .collect(); - SessionTimelineResponse { - session: SessionInfo { - session_pk: session.id, - user_id: session.user_id, - device_context_id: session.device_context_id, - agent_name: session.agent_name, - agent_version: session.agent_version, - session_id: session.session_id, - started_at: session.started_at, - ended_at: session.ended_at, - metadata: session.metadata, - title: None, - purpose: None, - parent_session_pk: None, - }, - turns: turn_timelines, - } -} - pub(super) fn usage_event( connection: &mut SqliteConnection, user_id: Uuid, @@ -1067,6 +927,7 @@ pub(super) fn usage_summary( ) -> Result { let group_by = parse_group_by(query.group_by.as_deref()); let dimensions = SummaryDimensions::from_group_by(&group_by); + let zone = parse_time_zone(query.tz.as_deref())?; let owner = load_owner(connection, user_id)?; if !owner_matches_filter(user_id, &owner, query.user_filter.as_deref()) { return Ok(empty_summary(query, group_by)); @@ -1088,6 +949,13 @@ pub(super) fn usage_summary( .map(validate_event_type) .transpose()?; let limit = normalize_limit(query.limit, default_limit, 100_000); + let from = query.from.map_or(i64::MIN, timestamp_to_micros); + let to = query.to.map_or(i64::MAX, timestamp_to_micros); + let day_offsets = if dimensions.enabled(SummaryDimension::Day) { + day_offsets(connection, user_id, zone, from, to)? + } else { + "[]".to_owned() + }; // Bind order mirrors the numbered placeholders in the query. let rows = sql_query(SUMMARY_AGGREGATE_SQL) @@ -1098,8 +966,8 @@ pub(super) fn usage_summary( .bind::(dimensions.enabled(SummaryDimension::Model)) .bind::(dimensions.enabled(SummaryDimension::EventType)) .bind::(user_id.to_string()) - .bind::(query.from.map_or(i64::MIN, timestamp_to_micros)) - .bind::(query.to.map_or(i64::MAX, timestamp_to_micros)) + .bind::(from) + .bind::(to) .bind::(has_agents) .bind::(&agent_names) .bind::(has_sessions) @@ -1111,6 +979,7 @@ pub(super) fn usage_summary( .bind::(event_type.is_some()) .bind::(event_type.unwrap_or_default()) .bind::(limit) + .bind::(&day_offsets) .load::(connection)?; let include_user = dimensions.enabled(SummaryDimension::User); let rows = rows @@ -1126,48 +995,122 @@ pub(super) fn usage_summary( }) } +/// Day offsets for `zone` across the owner's events observed in `[from, to)`; +/// see [`day_offsets_json`]. +fn day_offsets( + connection: &mut SqliteConnection, + user_id: Uuid, + zone: Tz, + from: i64, + to: i64, +) -> Result { + if zone == Tz::UTC { + return day_offsets_json(zone, None); + } + let (first, last) = llm_usage_events::table + .filter(llm_usage_events::user_id.eq(user_id.to_string())) + .filter(llm_usage_events::observed_at.ge(from)) + .filter(llm_usage_events::observed_at.lt(to)) + .select(( + min(llm_usage_events::observed_at), + max(llm_usage_events::observed_at), + )) + .first::<(Option, Option)>(connection)?; + day_offsets_json(zone, first.zip(last)) +} + /// Usage totals grouped in SQL, reading only the columns a summary reports. /// -/// Event text and metadata hold whole request bodies, so loading rows to sum -/// them in memory read gigabytes for a week of usage. The owner is the single +/// Each event row holds the event's conversation, so reading event rows +/// reads a page per event; the summary cover index carries every event column +/// used here, `id` included for the stream facts join. The owner is the single /// standalone user, so the user dimension is filled in afterwards. /// +/// Turns are counted per session within each group, then summed: the +/// distinct declared `turn_id` values when the session declares any turn, +/// else the requests that open a turn. +/// /// Bind order: `?1`..`?6` day, device, Agent, provider, model, and event-type /// dimensions; `?7` owner; `?8`/`?9` observation bounds in microseconds, kept /// unconditional so the owner-time index bounds the scan; `?10`..`?17` flags /// and JSON value lists for the Agent, session, provider, and model filters, /// each matching any listed value; `?18`/`?19` event-type filter; `?20` row -/// limit. +/// limit; `?21` the day offsets from [`day_offsets`], read only for the day +/// dimension. const SUMMARY_AGGREGATE_SQL: &str = "\ + WITH day_offsets(start_us, offset_us) AS MATERIALIZED ( \ + SELECT json_extract(value, '$[0]'), json_extract(value, '$[1]') FROM json_each(?21) \ + ), \ + first_openers AS ( \ + SELECT id FROM ( \ + SELECT fe.id, ROW_NUMBER() OVER ( \ + PARTITION BY fe.session_pk ORDER BY fe.observed_at, fe.id \ + ) AS n \ + FROM llm_usage_events fe \ + JOIN llm_usage_event_stream_facts ff ON ff.event_pk = fe.id \ + WHERE ff.starts_turn = 1 \ + ) firsts WHERE n = 1 \ + ), \ + by_session AS ( \ + SELECT \ + CASE WHEN ?1 THEN strftime( \ + '%Y-%m-%d', \ + (e.observed_at + ( \ + SELECT o.offset_us FROM day_offsets o WHERE o.start_us <= e.observed_at \ + ORDER BY o.start_us DESC LIMIT 1 \ + )) / 1000000, \ + 'unixepoch' \ + ) END AS day, \ + CASE WHEN ?2 THEN d.host_name END AS host_name, \ + CASE WHEN ?2 THEN d.platform END AS platform, \ + CASE WHEN ?2 THEN d.os_version END AS os_version, \ + CASE WHEN ?3 THEN e.agent_name END AS agent_name, \ + CASE WHEN ?4 THEN e.llm_provider END AS llm_provider, \ + CASE WHEN ?5 THEN e.llm_model END AS llm_model, \ + CASE WHEN ?6 THEN e.event_type END AS event_type, \ + e.session_pk AS session_pk, \ + EXISTS ( \ + SELECT 1 FROM llm_usage_events de \ + JOIN llm_usage_event_stream_facts df ON df.event_pk = de.id \ + WHERE de.session_pk = e.session_pk AND df.turn_id IS NOT NULL \ + ) AS declares_turns, \ + COUNT(DISTINCT f.turn_id) AS declared_turns, \ + COUNT(*) FILTER (WHERE f.starts_turn = 1 \ + AND (f.continues_context = 1 OR e.id IN first_openers)) AS opened_turns, \ + COUNT(*) FILTER (WHERE e.event_type = 'request') AS requests, \ + COUNT(*) FILTER (WHERE e.event_type = 'response') AS responses, \ + COALESCE(SUM(e.input_tokens), 0) AS input_tokens, \ + COALESCE(SUM(e.output_tokens), 0) AS output_tokens, \ + COALESCE(SUM(e.cache_read_tokens), 0) AS cache_read_tokens, \ + COALESCE(SUM(e.cache_write_tokens), 0) AS cache_write_tokens, \ + COALESCE(SUM(e.reasoning_tokens), 0) AS reasoning_tokens, \ + COALESCE(SUM(e.total_tokens), 0) AS total_tokens \ + FROM llm_usage_events e \ + LEFT JOIN llm_usage_event_stream_facts f ON f.event_pk = e.id \ + LEFT JOIN devices d ON ?2 AND d.id = e.device_context_id \ + WHERE e.user_id = ?7 \ + AND e.observed_at >= ?8 \ + AND e.observed_at < ?9 \ + AND (?10 = 0 OR e.agent_name IN (SELECT value FROM json_each(?11))) \ + AND (?12 = 0 OR e.session_id IN (SELECT value FROM json_each(?13))) \ + AND (?14 = 0 OR e.llm_provider IN (SELECT value FROM json_each(?15))) \ + AND (?16 = 0 OR e.llm_model IN (SELECT value FROM json_each(?17))) \ + AND (?18 = 0 OR e.event_type = ?19) \ + GROUP BY 1, 2, 3, 4, 5, 6, 7, 8, 9 \ + ) \ SELECT \ - CASE WHEN ?1 THEN strftime('%Y-%m-%d', e.observed_at / 1000000, 'unixepoch') END AS day, \ - CASE WHEN ?2 THEN d.host_name END AS host_name, \ - CASE WHEN ?2 THEN d.platform END AS platform, \ - CASE WHEN ?2 THEN d.os_version END AS os_version, \ - CASE WHEN ?3 THEN e.agent_name END AS agent_name, \ - CASE WHEN ?4 THEN e.llm_provider END AS llm_provider, \ - CASE WHEN ?5 THEN e.llm_model END AS llm_model, \ - CASE WHEN ?6 THEN e.event_type END AS event_type, \ - COUNT(DISTINCT e.session_pk) AS sessions, \ - COUNT(DISTINCT e.turn_pk) AS turns, \ - COUNT(*) FILTER (WHERE e.event_type = 'request') AS requests, \ - COUNT(*) FILTER (WHERE e.event_type = 'response') AS responses, \ - COALESCE(SUM(e.input_tokens), 0) AS input_tokens, \ - COALESCE(SUM(e.output_tokens), 0) AS output_tokens, \ - COALESCE(SUM(e.cache_read_tokens), 0) AS cache_read_tokens, \ - COALESCE(SUM(e.cache_write_tokens), 0) AS cache_write_tokens, \ - COALESCE(SUM(e.reasoning_tokens), 0) AS reasoning_tokens, \ - COALESCE(SUM(e.total_tokens), 0) AS total_tokens \ - FROM llm_usage_events e \ - LEFT JOIN devices d ON ?2 AND d.id = e.device_context_id \ - WHERE e.user_id = ?7 \ - AND e.observed_at >= ?8 \ - AND e.observed_at < ?9 \ - AND (?10 = 0 OR e.agent_name IN (SELECT value FROM json_each(?11))) \ - AND (?12 = 0 OR e.session_id IN (SELECT value FROM json_each(?13))) \ - AND (?14 = 0 OR e.llm_provider IN (SELECT value FROM json_each(?15))) \ - AND (?16 = 0 OR e.llm_model IN (SELECT value FROM json_each(?17))) \ - AND (?18 = 0 OR e.event_type = ?19) \ + day, host_name, platform, os_version, agent_name, llm_provider, llm_model, event_type, \ + COUNT(*) AS sessions, \ + COALESCE(SUM(CASE WHEN declares_turns THEN declared_turns ELSE opened_turns END), 0) AS turns, \ + COALESCE(SUM(requests), 0) AS requests, \ + COALESCE(SUM(responses), 0) AS responses, \ + COALESCE(SUM(input_tokens), 0) AS input_tokens, \ + COALESCE(SUM(output_tokens), 0) AS output_tokens, \ + COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens, \ + COALESCE(SUM(cache_write_tokens), 0) AS cache_write_tokens, \ + COALESCE(SUM(reasoning_tokens), 0) AS reasoning_tokens, \ + COALESCE(SUM(total_tokens), 0) AS total_tokens \ + FROM by_session g \ GROUP BY 1, 2, 3, 4, 5, 6, 7, 8 \ ORDER BY total_tokens DESC, agent_name ASC, day ASC \ LIMIT ?20"; @@ -1406,16 +1349,13 @@ fn event_response( host_name: device.map(|value| value.host_name.clone()), platform: device.map(|value| value.platform.clone()), session_pk: event.session_pk, - turn_pk: event.turn_pk, + location: None, agent_name: event.agent_name, agent_version: event.agent_version, session_id: event.session_id, - turn_index: event.turn_index, llm_provider: event.llm_provider, llm_model: event.llm_model, event_type: event.event_type, - text: event.text, - text_sha256: event.text_sha256.map(hex::encode), input_tokens: event.input_tokens, output_tokens: event.output_tokens, cache_read_tokens: event.cache_read_tokens, @@ -1426,6 +1366,7 @@ fn event_response( stream_id: event.stream_id, response_id: event.response_id, previous_response_id: event.previous_response_id, + conversation: event.conversation, metadata: event.metadata, attachments, } @@ -1442,6 +1383,7 @@ pub(super) const fn timestamp_to_micros(value: DateTime) -> i64 { #[cfg(test)] mod tests { + use crate::storage::stream_cache::StreamCache; use diesel::{ Connection, ExpressionMethods, QueryDsl, RunQueryDsl, SqliteConnection, dsl::count_star, sql_query, @@ -1451,16 +1393,21 @@ mod tests { use crate::{ search::SessionSearchQuery, - usage::{IngestEventsRequest, RawEventsQuery, SummaryQuery}, + usage::{ + IngestEventsRequest, RawEventsQuery, SummaryQuery, + preview::{PreviewQuery, Window}, + segments::SegmentQuery, + }, }; use super::{ super::{ migrations, - schema::{agent_turns, llm_usage_events}, + schema::{agent_sessions, llm_usage_event_stream_facts, llm_usage_events}, search, }, - ingest_events, raw_events, session_outline, usage_summary, + event_conversation, event_location, ingest_events, raw_events, session_header, + session_segment, session_streams, usage_summary, }; const OWNER_ID: Uuid = Uuid::from_u128(1); @@ -1491,11 +1438,11 @@ mod tests { assert_eq!(duplicate.accepted, 0); assert_eq!(duplicate.duplicates, 1); - let projection_count = sql_query("SELECT count(*) AS count FROM usage_events_fts") + let reference_count = sql_query("SELECT count(*) AS count FROM search_refs") .get_result::(&mut database) - .expect("FTS projection should be queryable") + .expect("search references should be queryable") .count; - assert_eq!(projection_count, 1); + assert_eq!(reference_count, 1, "one piece, indexed once"); let results = search::session_search( &mut database, OWNER_ID, @@ -1521,12 +1468,12 @@ mod tests { .iter() .flat_map(|fragment| &fragment.segments) .any(|segment| segment.highlighted), - "FTS result should contain a structured highlight" + "a search result should contain a structured highlight" ); } #[test] - fn validation_failure_leaves_relational_and_fts_tables_empty() { + fn validation_failure_leaves_relational_and_search_tables_empty() { let mut database = database(); let request = request_with_text("invalid event"); @@ -1538,38 +1485,18 @@ mod tests { .expect("relational table should be queryable"); assert_eq!(relational_count, 0); - let fts_count = sql_query("SELECT count(*) AS count FROM usage_events_fts") - .get_result::(&mut database) - .expect("FTS table should be queryable") - .count; - assert_eq!(fts_count, 0); - } - - #[test] - fn outline_measures_text_containing_nul_characters() { - let mut database = database(); - let text = format!("ab\0{}", "c".repeat(300)); - ingest_events(&mut database, &request_with_text(&text), OWNER_ID, 100) - .expect("ingest should succeed"); - let session_pk = llm_usage_events::table - .select(llm_usage_events::session_pk) - .first::(&mut database) - .expect("event should exist"); - let session_pk = Uuid::parse_str(&session_pk).expect("session pk should be a UUID"); - - let outline = - session_outline(&mut database, OWNER_ID, session_pk).expect("outline should load"); - - let event = &outline.turns[0].events[0]; - assert_eq!(event.text_chars, 303); - assert_eq!( - event.preview.as_deref(), - Some(format!("ab\0{}", "c".repeat(197)).as_str()) - ); + let search_count = sql_query( + "SELECT (SELECT count(*) FROM search_texts) + (SELECT count(*) FROM search_refs) \ + AS count", + ) + .get_result::(&mut database) + .expect("search tables should be queryable") + .count; + assert_eq!(search_count, 0); } #[test] - fn raw_body_archive_survives_ingest_retries_and_trajectory_projection() { + fn opaque_metadata_survives_ingest_retries_and_preview_projection() { use base64::{Engine as _, engine::general_purpose::STANDARD}; use sha2::Digest as _; let original = @@ -1577,36 +1504,37 @@ mod tests { let mut database = database(); let mut request = request_with_text("hello"); request.events[0].event_type = crate::usage::UsageEventType::Request; - request.events[0].metadata["raw_body"] = json!({"transport":"http", "parts":[{ + request.events[0].metadata["archive"] = json!({"transport":"http", "parts":[{ "content_base64": STANDARD.encode(original), "byte_size": original.len(), "truncated": false, "sha256": hex::encode(sha2::Sha256::digest(original)) }]}); let expected = request.events[0].metadata.clone(); ingest_events(&mut database, &request, OWNER_ID, 100).unwrap(); ingest_events(&mut database, &request, OWNER_ID, 100).unwrap(); - let session_pk = llm_usage_events::table - .select(llm_usage_events::session_pk) + let event_pk = llm_usage_events::table + .select(llm_usage_events::id) .first::(&mut database) .unwrap(); - let session_pk = Uuid::parse_str(&session_pk).unwrap(); - let timeline = super::session_timeline(&mut database, OWNER_ID, session_pk).unwrap(); - let stored = &timeline.turns[0].events[0].metadata; + let event_pk = Uuid::parse_str(&event_pk).unwrap(); + let event = super::usage_event(&mut database, OWNER_ID, event_pk).unwrap(); + let stored = &event.metadata; assert_eq!(*stored, expected); + assert_eq!(event.conversation, request.events[0].conversation); assert_eq!( STANDARD .decode( - stored["raw_body"]["parts"][0]["content_base64"] + stored["archive"]["parts"][0]["content_base64"] .as_str() .unwrap() ) .unwrap(), original ); - let trajectory = crate::usage::trajectory::build(timeline); - assert_eq!(trajectory.items[0].text.as_deref(), Some("hello")); - let after = super::session_timeline(&mut database, OWNER_ID, session_pk).unwrap(); - assert_eq!(after.turns[0].events.len(), 1); - assert_eq!(after.turns[0].events[0].metadata, expected); + let preview = + event_conversation(&mut database, OWNER_ID, event_pk, &Window::default()).unwrap(); + assert_eq!(preview.chunks[0].entries[0]["text"].as_str(), Some("hello")); + let after = super::usage_event(&mut database, OWNER_ID, event_pk).unwrap(); + assert_eq!(after.metadata, expected); } #[test] @@ -1619,10 +1547,6 @@ mod tests { event.llm.model = " Model ".to_owned(); event.device.platform = "macOS".to_owned(); ingest_events(&mut database, &request, OWNER_ID, 100).expect("ingest should succeed"); - let mut empty = request_with_text(""); - empty.events[0].event_id = "sqlite-empty-text".to_owned(); - ingest_events(&mut database, &empty, OWNER_ID, 100).expect("ingest should succeed"); - let query: RawEventsQuery = serde_json::from_value(json!({})).expect("query should deserialize"); let events = raw_events(&mut database, &query, OWNER_ID, 50) @@ -1636,12 +1560,595 @@ mod tests { assert_eq!(stored.llm_provider, "Gateway_Host"); assert_eq!(stored.llm_model, " Model "); assert_eq!(stored.platform.as_deref(), Some("macOS")); - assert_eq!(stored.text.as_deref(), Some(" hello \n")); - let empty = events + assert_eq!( + stored.conversation.as_ref().expect("conversation")["entries"][0]["text"], + json!(" hello \n") + ); + } + + #[test] + fn ingest_stores_the_stream_facts_of_each_new_event_beside_it() { + type Facts = ( + String, + Option, + Option, + Option, + Option, + Option, + ); + let mut database = database(); + let event = |id: &str, event_type: &str, entries: serde_json::Value| { + json!({ + "event_id": id, + "observed_at": "2026-10-19T00:00:00Z", + "device": {"host_name": "local", "platform": "macos"}, + "agent": {"name": "codex"}, + "session_id": "sqlite-session", + "llm": {"provider": "openai", "model": "gpt-5"}, + "event_type": event_type, + "conversation": {"entries": entries, "completeness": "complete"} + }) + }; + let request: IngestEventsRequest = serde_json::from_value(json!({"events": [ + event("ask", "request", json!([ + {"role": "user", "text": "fix it"}, + {"role": "tool_call", "text": "{}", "call_id": "a", "item_id": "fc_1"}, + {"role": "tool_result", "text": "ok", "call_id": "a"} + ])), + event("answer", "response", json!([ + {"role": "assistant", "text": "done", "item_id": "msg_1"} + ])) + ]})) + .expect("test ingest request should deserialize"); + for _ in 0_u8..2 { + ingest_events(&mut database, &request, OWNER_ID, 100).expect("ingest should succeed"); + } + + let facts = llm_usage_event_stream_facts::table + .inner_join(llm_usage_events::table) + .order(llm_usage_events::event_id.desc()) + .select(( + llm_usage_events::event_id, + llm_usage_event_stream_facts::new_from, + llm_usage_event_stream_facts::starts_turn, + llm_usage_event_stream_facts::last_identity, + llm_usage_event_stream_facts::trailing, + llm_usage_event_stream_facts::emitted_ids, + )) + .load::(&mut database) + .expect("stream facts should load"); + assert_eq!( + facts, + [ + ( + "ask".to_owned(), + Some(2_i32), + Some(false), + Some("a".to_owned()), + Some("[]".to_owned()), + None + ), + ( + "answer".to_owned(), + None, + None, + None, + None, + Some(r#"["item:msg_1"]"#.to_owned()) + ), + ], + "one row per event, written once, the other side's columns NULL" + ); + } + + /// A database holding session `streams-session`: exchange r1 asks and + /// calls c1, exchange r2 sends its result back and gets an answer. + fn streams_database() -> (SqliteConnection, Uuid) { + let mut database = database(); + let event = |id: &str, second: u8, event_type: &str, entries: serde_json::Value| { + let response = if second < 3 { "r1" } else { "r2" }; + json!({ + "event_id": id, + "observed_at": format!("2026-10-20T00:00:0{second}Z"), + "device": {"host_name": "local", "platform": "macos"}, + "agent": {"name": "codex"}, + "session_id": "streams-session", + "llm": {"provider": "openai", "model": "gpt-5"}, + "event_type": event_type, + "response_id": response, + "conversation": {"entries": entries, "completeness": "complete"} + }) + }; + let request: IngestEventsRequest = serde_json::from_value(json!({"events": [ + event("a", 1, "request", json!([{"role": "user", "text": "fix the bug"}])), + event("ab", 2, "response", json!([ + {"role": "tool_call", "text": "{}", "call_id": "c1"} + ])), + event("abc", 3, "request", json!([ + {"role": "user", "text": "fix the bug"}, + {"role": "tool_call", "text": "{}", "call_id": "c1"}, + {"role": "tool_result", "text": "ok", "call_id": "c1"} + ])), + event("abcd", 4, "response", json!([{"role": "assistant", "text": "fixed"}])) + ]})) + .expect("test ingest request should deserialize"); + ingest_events(&mut database, &request, OWNER_ID, 100).expect("ingest should succeed"); + let session_pk = agent_sessions::table + .select(agent_sessions::id) + .first::(&mut database) + .expect("the session is stored"); + let session_pk = Uuid::parse_str(&session_pk).expect("a UUID"); + (database, session_pk) + } + + #[test] + fn session_streams_link_stored_exchanges() { + let (mut database, session_pk) = streams_database(); + let streams = session_streams(&mut database, OWNER_ID, session_pk, &StreamCache::default()) + .expect("the session's streams derive"); + let segments = streams + .segments + .iter() + .map(|segment| (segment.head.as_str(), segment.exchange_count)) + .collect::>(); + assert_eq!(segments, [("r1", 2)], "r2 continues r1 by its call"); + let label = streams.segments[0] + .label + .as_ref() + .expect("r1 asks in words"); + assert_eq!( + (label.entry_index, label.text.as_str(), label.length), + (0_i32, "fix the bug", 11_i32) + ); + assert!( + session_streams( + &mut database, + OWNER_ID, + Uuid::nil(), + &StreamCache::default() + ) + .is_err(), + "a session without events is not found" + ); + } + + #[test] + fn declared_sub_agent_ids_are_stored_and_group_the_sub_agent_segment() { + let mut database = database(); + let event = |id: &str, second: u8, event_type: &str, declared: serde_json::Value| { + let mut event = json!({ + "event_id": id, + "observed_at": format!("2026-10-20T00:00:0{second}Z"), + "device": {"host_name": "local", "platform": "macos"}, + "agent": {"name": "claude-code"}, + "session_id": "declared-session", + "llm": {"provider": "anthropic", "model": "claude"}, + "event_type": event_type, + "response_id": &id[..1], + "conversation": { + "entries": [{"role": "user", "text": id}], + "completeness": "complete" + } + }); + for (key, value) in declared.as_object().into_iter().flatten() { + event[key] = value.clone(); + } + event + }; + let main = json!({"turn_id": "p1"}); + let sub = json!({"subagent_id": "a1", "turn_id": "p1"}); + let request: IngestEventsRequest = serde_json::from_value(json!({"events": [ + event("m-req", 1, "request", main.clone()), + event("m-res", 5, "response", main), + event("s-req", 2, "request", sub.clone()), + event("s-res", 3, "response", sub), + // Without a request, the response's declared ids stand for the + // exchange; with no turn it attaches by time. + event("o-res", 4, "response", json!({"subagent_id": "a2", "parent_subagent_id": null})), + ]})) + .expect("test ingest request should deserialize"); + ingest_events(&mut database, &request, OWNER_ID, 100).expect("ingest should succeed"); + + let stored = llm_usage_event_stream_facts::table + .inner_join(llm_usage_events::table) + .filter(llm_usage_events::event_id.eq("s-res")) + .select(( + llm_usage_event_stream_facts::subagent_id, + llm_usage_event_stream_facts::parent_subagent_id, + llm_usage_event_stream_facts::turn_id, + )) + .first::<(Option, Option, Option)>(&mut database) + .expect("the response's facts"); + assert_eq!(stored, (Some("a1".to_owned()), None, Some("p1".to_owned()))); + + let session_pk = agent_sessions::table + .select(agent_sessions::id) + .first::(&mut database) + .expect("the session is stored"); + let session_pk = Uuid::parse_str(&session_pk).expect("a UUID"); + let streams = session_streams(&mut database, OWNER_ID, session_pk, &StreamCache::default()) + .expect("the session's streams derive"); + let attachments = streams + .segments .iter() - .find(|event| event.event_id == "sqlite-empty-text") - .expect("the empty-text event should be stored"); - assert_eq!(empty.text.as_deref(), Some("")); + .map(|segment| { + ( + segment.head.as_str(), + segment.subagent_id.as_deref(), + segment.attached_to.as_deref(), + segment.attached_by, + ) + }) + .collect::>(); + assert_eq!( + attachments, + [ + ("m", None, None, None), + // Sub-agent segments are never attached. + ("s", Some("a1"), None, None), + ("o", Some("a2"), None, None), + ] + ); + } + + /// One exchange's events: `(name, session, agent, second, declared ids, + /// request entries, response entries)`. + type TurnExchange = ( + &'static str, + &'static str, + &'static str, + u8, + serde_json::Value, + serde_json::Value, + serde_json::Value, + ); + + fn turn_events(exchanges: &[TurnExchange]) -> IngestEventsRequest { + let mut events = Vec::new(); + for (name, session, agent, second, declared, request, response) in exchanges { + for (side, entries, at) in [("request", request, "0"), ("response", response, "5")] { + let mut event = json!({ + "event_id": format!("{name}-{side}"), + "observed_at": format!("2026-10-22T00:00:{second:02}.{at}Z"), + "device": {"host_name": "local", "platform": "macos"}, + "agent": {"name": agent}, + "session_id": session, + "llm": {"provider": "anthropic", "model": "claude"}, + "event_type": side, + "response_id": name, + "conversation": {"entries": entries, "completeness": "complete"} + }); + for (key, value) in declared.as_object().into_iter().flatten() { + event[key] = value.clone(); + } + events.push(event); + } + } + serde_json::from_value(json!({ "events": events })).expect("a valid ingest request") + } + + /// A Claude Code session `declared` with two declared turns, and a Codex + /// session `plain` without declarations whose three requests open turns. + fn turn_database() -> SqliteConnection { + let mut database = database(); + let user = |text: &str| json!({"role": "user", "text": text}); + let call = |id: &str| json!({"role": "tool_call", "text": "{}", "call_id": id}); + let result = |id: &str| json!({"role": "tool_result", "text": "ok", "call_id": id}); + let answer = |text: &str| json!({"role": "assistant", "text": text}); + let said = |text: &str| json!([answer(text)]); + let (cc, codex) = ("claude-code", "codex"); + let none = json!({}); + let request = turn_events(&[ + // An auxiliary request before the first prompt belongs to no turn. + ( + "x", + "declared", + cc, + 1, + none.clone(), + json!([user("quota")]), + said("ok"), + ), + ( + "m", + "declared", + cc, + 2, + json!({"turn_id": "p1"}), + json!([user("fix")]), + json!([call("c1")]), + ), + ( + "s", + "declared", + cc, + 3, + json!({"subagent_id": "a1", "turn_id": "p1"}), + json!([user("explore")]), + said("found"), + ), + // A title request declares nothing and joins p1. + ( + "t", + "declared", + cc, + 4, + none.clone(), + json!([user("name it")]), + said("Fix"), + ), + ( + "m2", + "declared", + cc, + 5, + json!({"turn_id": "p1"}), + json!([user("fix"), call("c1"), result("c1")]), + said("fixed"), + ), + ( + "i", + "declared", + cc, + 6, + json!({"turn_id": "p2"}), + json!([user("next")]), + said("ok"), + ), + ( + "u1", + "plain", + codex, + 7, + none.clone(), + json!([user("one")]), + said("1"), + ), + ( + "u2", + "plain", + codex, + 8, + none.clone(), + // A plain chat carries its history: each later question + // follows an answer, so each opens a turn. + json!([user("one"), answer("1"), user("two")]), + said("2"), + ), + ( + "u3", + "plain", + codex, + 9, + none, + json!([ + user("one"), + answer("1"), + user("two"), + answer("2"), + user("three") + ]), + said("3"), + ), + ]); + ingest_events(&mut database, &request, OWNER_ID, 100).expect("ingest should succeed"); + database + } + + #[test] + fn turns_follow_declarations_or_turn_opening_requests_everywhere() { + let mut database = turn_database(); + let stored = llm_usage_event_stream_facts::table + .inner_join(llm_usage_events::table) + .filter(llm_usage_events::event_id.eq("i-response")) + .select(llm_usage_event_stream_facts::turn_id) + .first::>(&mut database) + .expect("the response's facts"); + assert_eq!(stored.as_deref(), Some("p2")); + + let query = crate::usage::session_list::SessionListQuery::default() + .validate() + .expect("valid"); + let listed = super::super::session_list::list_sessions(&mut database, OWNER_ID, &query) + .expect("sessions list") + .items + .into_iter() + .map(|item| (item.session_id, (item.session_pk, item.turns))) + .collect::>(); + let cache = StreamCache::default(); + for (session, turns) in [("declared", 2_i64), ("plain", 3_i64)] { + let (session_pk, listed_turns) = listed[session]; + assert_eq!(listed_turns, turns, "{session}: listed turns"); + let header = + session_header(&mut database, OWNER_ID, session_pk, &cache).expect("the header"); + assert_eq!( + i64::try_from(header.turn_count).expect("a small count"), + turns, + "{session}: the header counts as the list does" + ); + } + + let (declared_pk, _) = listed["declared"]; + let streams = + session_streams(&mut database, OWNER_ID, declared_pk, &cache).expect("streams"); + let segments = streams + .segments + .iter() + .map(|segment| { + ( + segment.head.as_str(), + segment.turn_id.as_deref(), + segment.turn, + ) + }) + .collect::>(); + assert_eq!( + segments, + [ + ("x", None, None), + ("m", Some("p1"), Some(1)), + ("s", Some("p1"), Some(1)), + ("t", Some("p1"), Some(1)), + ("i", Some("p2"), Some(2)), + ] + ); + let event_pk = |database: &mut SqliteConnection, event_id: &str| { + let pk = llm_usage_events::table + .filter(llm_usage_events::event_id.eq(event_id)) + .select(llm_usage_events::id) + .first::(database) + .expect("the event is stored"); + Uuid::parse_str(&pk).expect("a UUID") + }; + for (event_id, turn) in [ + ("x-request", None), + ("t-response", Some(1)), + ("i-request", Some(2)), + ] { + let pk = event_pk(&mut database, event_id); + let location = event_location(&mut database, OWNER_ID, pk, &cache).expect("located"); + assert_eq!(location.location.turn, turn, "{event_id}"); + } + + let rows = summary(&mut database, json!({"group_by": "agent"})); + let turns = rows + .as_array() + .expect("rows") + .iter() + .map(|row| { + ( + row["agent_name"].clone(), + row["turns"].clone(), + row["sessions"].clone(), + ) + }) + .collect::>(); + assert!( + turns.contains(&(json!("claude-code"), json!(2_i64), json!(1_i64))), + "{turns:?}" + ); + assert!( + turns.contains(&(json!("codex"), json!(3_i64), json!(1_i64))), + "{turns:?}" + ); + let all = summary(&mut database, json!({"group_by": "day"})); + assert_eq!( + (all[0]["turns"].clone(), all[0]["sessions"].clone()), + (json!(5_i64), json!(2_i64)) + ); + } + + #[test] + fn segments_and_conversations_preview_stored_events() { + let (mut database, session_pk) = streams_database(); + let query = SegmentQuery { + entries: Some("0".to_owned()), + text_preview: Some(3), + } + .validate() + .expect("a valid query"); + let segment = session_segment( + &mut database, + OWNER_ID, + session_pk, + "r1", + &query, + &StreamCache::default(), + ) + .expect("the whole segment"); + assert_eq!(segment.segment, "r1"); + let r2 = &segment.exchanges[1]; + assert_eq!( + (r2.exchange.base.as_deref(), r2.model.as_deref()), + (Some("r1"), Some("gpt-5")) + ); + let request = r2.request.as_ref().expect("a request preview"); + assert_eq!( + ( + request.entries_total, + request.omitted_entries, + request.chunks.len() + ), + (3_usize, 2_usize, 1_usize) + ); + assert_eq!( + request.chunks[0].entries[0], + json!({"role": "user", "text": "fix", "text_length": 11_usize}) + ); + assert!( + session_segment( + &mut database, + OWNER_ID, + session_pk, + "zz", + &query, + &StreamCache::default() + ) + .is_err() + ); + + let window = Window::parse(&PreviewQuery { + entries: Some("-1".to_owned()), + text_preview: None, + }) + .expect("a valid window"); + let conversation = event_conversation(&mut database, OWNER_ID, request.event_pk, &window) + .expect("the event's conversation"); + assert_eq!( + ( + conversation.chunks[0].from, + conversation.chunks[0].entries[0]["role"].clone() + ), + (2, json!("tool_result")) + ); + } + + #[test] + fn the_session_header_sums_its_events() { + let (mut database, session_pk) = streams_database(); + let header = session_header(&mut database, OWNER_ID, session_pk, &StreamCache::default()) + .expect("the header"); + assert_eq!( + ( + header.event_count, + header.exchange_count, + header.turn_count, + header.segment_count + ), + (4, 2, 1, 1) + ); + assert_eq!( + (header.agents[0].name.as_str(), header.models.as_slice()), + ("codex", ["gpt-5".to_owned()].as_slice()) + ); + assert_eq!( + header.end.signed_duration_since(header.start), + chrono::TimeDelta::seconds(3) + ); + assert_eq!( + header.title.as_ref().map(|label| label.text.as_str()), + Some("fix the bug"), + "the title is the first request label" + ); + assert!( + session_header( + &mut database, + OWNER_ID, + Uuid::nil(), + &StreamCache::default() + ) + .is_err() + ); + assert!( + session_header( + &mut database, + Uuid::from_u128(2), + session_pk, + &StreamCache::default() + ) + .is_err(), + "another owner cannot read the session" + ); } #[test] @@ -1662,17 +2169,16 @@ mod tests { assert_eq!(listed.id, complete.id); assert_eq!(listed.total_tokens, complete.total_tokens); assert_eq!(listed.host_name.as_deref(), Some("alice-mbp")); - assert_eq!(listed.text, None); assert_eq!(listed.metadata, json!({})); + assert!(listed.conversation.is_none()); assert!(listed.attachments.is_empty()); - assert_eq!(complete.text.as_deref(), Some("body")); - assert_eq!(complete.metadata["raw_body"], json!("large")); + assert_eq!(complete.metadata["archive"], json!("large")); } } fn summary_seeded() -> SqliteConnection { let mut database = database(); - let event = |id: &str, agent: &str, session: &str, turn: i32, observed_at: &str| { + let event = |id: &str, agent: &str, session: &str, observed_at: &str| { let (provider, model) = if agent == "codex" { ("openai", "gpt-5") } else { @@ -1684,25 +2190,32 @@ mod tests { "device": {"host_name": "alice-mbp", "platform": "macos", "os_version": "15.1"}, "agent": {"name": agent}, "session_id": session, - "turn_id": format!("{session}-{turn}"), "llm": {"provider": provider, "model": model}, "event_type": if id.ends_with("req") { "request" } else { "response" }, - "text": "body", "token_usage": if id.ends_with("req") { json!({"input_tokens": 10_i64, "output_tokens": 0_i64, "cache_read_tokens": 4_i64, "cache_write_tokens": 0_i64, "reasoning_tokens": 0_i64, "total_tokens": 14_i64}) } else { json!({"input_tokens": 0_i64, "output_tokens": 5_i64, "cache_read_tokens": 0_i64, "cache_write_tokens": 0_i64, "reasoning_tokens": 0_i64, "total_tokens": 5_i64}) }, - "metadata": {"raw_body": "large"} + // A request asking in words after an answer opens a turn; a + // response opens none. + "conversation": id.ends_with("req").then(|| json!({ + "entries": [ + {"role": "assistant", "text": "earlier"}, + {"role": "user", "text": "ask"} + ], + "completeness": "complete" + })), + "metadata": {"archive": "large"} }) }; let request: IngestEventsRequest = serde_json::from_value(json!({ "events": [ - event("a-req", "codex", "alpha", 1_i32, "2026-09-01T23:00:00Z"), - event("a-res", "codex", "alpha", 1_i32, "2026-09-01T23:30:00Z"), - event("b-req", "claude-desktop", "beta", 1_i32, "2026-09-02T01:00:00Z"), - event("c-res", "claude-code", "gamma", 1_i32, "2026-09-02T02:00:00Z"), - event("a2-req", "codex", "alpha", 2_i32, "2026-09-03T00:00:00Z"), + event("a-req", "codex", "alpha", "2026-09-01T23:00:00Z"), + event("a-res", "codex", "alpha", "2026-09-01T23:30:00Z"), + event("b-req", "claude-desktop", "beta", "2026-09-02T01:00:00Z"), + event("c-res", "claude-code", "gamma", "2026-09-02T02:00:00Z"), + event("a2-req", "codex", "alpha", "2026-09-03T00:00:00Z"), ] })) .expect("request should deserialize"); @@ -1717,6 +2230,42 @@ mod tests { serde_json::to_value(response.rows).expect("rows should serialize") } + #[test] + fn summary_reads_events_only_from_covering_indexes() { + #[derive(diesel::QueryableByName)] + struct PlanStep { + #[diesel(sql_type = diesel::sql_types::Text)] + detail: String, + } + let mut database = summary_seeded(); + + let plan = diesel::sql_query(format!( + "EXPLAIN QUERY PLAN {}", + super::SUMMARY_AGGREGATE_SQL + )) + .load::(&mut database) + .expect("summary should plan"); + + // Each event row holds its conversation, so reading rows costs a page + // per event; `e` and the turn check's `de` are both events. + let event_reads = plan + .iter() + .map(|step| step.detail.as_str()) + .filter(|detail| { + let mut words = detail.split_whitespace(); + matches!(words.next(), Some("SEARCH" | "SCAN")) + && matches!(words.next(), Some("e" | "de")) + }) + .collect::>(); + assert!( + !event_reads.is_empty() + && event_reads + .iter() + .all(|detail| detail.contains("USING COVERING INDEX")), + "summary should not read event rows: {event_reads:?}" + ); + } + #[test] fn summary_groups_in_sql_by_reported_agent_and_distinct_counts() { let mut database = summary_seeded(); @@ -1746,7 +2295,7 @@ mod tests { { "agent_name": "claude-code", "llm_provider": "anthropic", "llm_model": "claude-x", - "sessions": 1_i64, "turns": 1_i64, "requests": 0_i64, "responses": 1_i64, + "sessions": 1_i64, "turns": 0_i64, "requests": 0_i64, "responses": 1_i64, "input_tokens": 0_i64, "output_tokens": 5_i64, "cache_read_tokens": 0_i64, "cache_write_tokens": 0_i64, "reasoning_tokens": 0_i64, "total_tokens": 5_i64 } @@ -1837,195 +2386,108 @@ mod tests { ); } - #[test] - fn summary_fills_owner_and_device_dimensions() { - let mut database = summary_seeded(); - - let rows = summary(&mut database, json!({"group_by": "user,device"})); - - assert_eq!(rows.as_array().map(Vec::len), Some(1)); - assert_eq!(rows[0]["user_id"], json!(OWNER_ID)); - assert!(rows[0]["user_email"].is_string()); - assert_eq!(rows[0]["host_name"], json!("alice-mbp")); - assert_eq!(rows[0]["platform"], json!("macos")); - assert_eq!(rows[0]["os_version"], json!("15.1")); - assert_eq!(rows[0]["sessions"], json!(3_i64)); - assert!(rows[0].get("agent_name").is_none()); - } - - /// One request observed `second` seconds into a fixed minute, working for - /// `turn` when it names one. - fn turn_event(id: &str, agent: &str, second: u32, turn: Option<&str>) -> serde_json::Value { - json!({ - "event_id": id, - "observed_at": format!("2026-10-08T00:00:{second:02}Z"), - "device": {"host_name": "local", "platform": "macos"}, - "agent": {"name": agent}, - "session_id": "thread-session", - "turn_id": turn, - "llm": {"provider": "Anthropic", "model": "claude-x"}, - "event_type": "request", - "token_usage": {"input_tokens": 1_i64, "output_tokens": 0_i64, "cache_read_tokens": 0_i64, "cache_write_tokens": 0_i64, "reasoning_tokens": 0_i64, "total_tokens": 1_i64}, - "metadata": {} - }) - } - - fn ingest_turn_events(database: &mut SqliteConnection, events: Vec) { - let request: IngestEventsRequest = - serde_json::from_value(json!({"events": serde_json::Value::Array(events)})) - .expect("request should deserialize"); - ingest_events(database, &request, OWNER_ID, 100).expect("ingest should succeed"); - } - - fn turn_of(database: &mut SqliteConnection, event_id: &str) -> i32 { - llm_usage_events::table - .filter(llm_usage_events::event_id.eq(event_id)) - .select(llm_usage_events::turn_index) - .first::(database) - .expect("event should be stored") - } - - /// Stored turns in number order: key, number, start and end second. - fn turns(database: &mut SqliteConnection) -> Vec<(Option, i32, i64, Option)> { - agent_turns::table - .order_by(agent_turns::turn_index.asc()) - .select(( - agent_turns::turn_key, - agent_turns::turn_index, - agent_turns::started_at, - agent_turns::ended_at, - )) - .load::<(Option, i32, i64, Option)>(database) - .expect("turns should load") - } - - fn second(value: i64) -> i64 { - super::timestamp_to_micros( - chrono::DateTime::parse_from_rfc3339(&format!("2026-10-08T00:00:{value:02}Z")) - .expect("timestamp should parse") - .into(), - ) + /// Token totals of each `day` row, in day order. + fn day_totals(rows: &serde_json::Value) -> Vec<(serde_json::Value, serde_json::Value)> { + let mut days = rows + .as_array() + .expect("rows should be an array") + .iter() + .map(|row| (row["day"].clone(), row["total_tokens"].clone())) + .collect::>(); + days.sort_by_key(|day| day.0.to_string()); + days } #[test] - fn events_naming_a_turn_share_it_and_others_join_the_turn_in_progress() { - let mut database = database(); - ingest_turn_events( + fn summary_groups_days_in_the_requested_time_zone() { + let mut database = summary_seeded(); + + let utc = summary(&mut database, json!({"group_by": "day"})); + let shanghai = summary( &mut database, - vec![ - turn_event("main-1", "codex", 10, Some("p1")), - turn_event("subagent-1", "codex", 12, Some("p1")), - turn_event("recap", "codex", 20, None), - turn_event("main-2", "codex", 30, Some("p2")), - turn_event("subagent-2", "codex", 40, Some("p2")), - ], + json!({"group_by": "day", "tz": "Asia/Shanghai"}), ); - // Delivered late: a request naming no turn observed during turn 2, and - // a background agent of turn 1 still working during turn 2. - ingest_turn_events( - &mut database, + + assert_eq!( + day_totals(&utc), vec![ - turn_event("late-title", "codex", 32, None), - turn_event("late-subagent", "codex", 45, Some("p1")), - ], + (json!("2026-09-01"), json!(19_i64)), + (json!("2026-09-02"), json!(19_i64)), + (json!("2026-09-03"), json!(14_i64)), + ] ); - - for (event, turn) in [ - ("main-1", 1_i32), - ("subagent-1", 1_i32), - ("recap", 1_i32), - ("late-subagent", 1_i32), - ("main-2", 2_i32), - ("subagent-2", 2_i32), - ("late-title", 2_i32), - ] { - assert_eq!(turn_of(&mut database, event), turn, "{event}"); - } assert_eq!( - turns(&mut database), + day_totals(&shanghai), vec![ - (Some("p1".to_owned()), 1_i32, second(10), Some(second(45))), - (Some("p2".to_owned()), 2_i32, second(30), Some(second(40))), + (json!("2026-09-02"), json!(38_i64)), + (json!("2026-09-03"), json!(14_i64)), ] ); + let query: SummaryQuery = serde_json::from_value(json!({"group_by": "day", "tz": "UTC+8"})) + .expect("query should deserialize"); + assert!(usage_summary(&mut database, &query, OWNER_ID, 500).is_err()); } #[test] - fn the_first_named_turn_adopts_the_turn_opened_without_one() { + fn summary_days_follow_daylight_saving_changes() { let mut database = database(); - ingest_turn_events( - &mut database, - vec![turn_event("probe", "claude-code", 5, None)], - ); - ingest_turn_events( + let event = |id: &str, observed_at: &str| { + json!({ + "event_id": id, + "observed_at": observed_at, + "device": {"host_name": "alice-mbp", "platform": "macos", "os_version": "15.1"}, + "agent": {"name": "Codex"}, + "session_id": "dst", + "llm": {"provider": "OpenAI", "model": "gpt-5"}, + "event_type": "request", + "token_usage": {"input_tokens": 1_i64, "total_tokens": 1_i64}, + "metadata": {"turn_id": id} + }) + }; + // Los Angeles moves from PST (-8) to PDT (-7) at 2026-03-08T10:00Z. + let request: IngestEventsRequest = serde_json::from_value(json!({ + "events": [ + event("before", "2026-03-08T07:30:00Z"), // 03-07 23:30 PST + event("late", "2026-03-09T06:30:00Z"), // 03-08 23:30 PDT + event("after", "2026-03-09T07:30:00Z"), // 03-09 00:30 PDT + ] + })) + .expect("request should deserialize"); + ingest_events(&mut database, &request, OWNER_ID, 100).expect("ingest should succeed"); + + let rows = summary( &mut database, - vec![ - turn_event("main-1", "claude-code", 10, Some("p1")), - turn_event("main-2", "claude-code", 30, Some("p2")), - // Earlier than every turn: joins the first one. - turn_event("early", "claude-code", 1, None), - ], + json!({"group_by": "day", "tz": "America/Los_Angeles"}), ); - assert_eq!(turn_of(&mut database, "probe"), 1_i32); - assert_eq!(turn_of(&mut database, "main-1"), 1_i32); - assert_eq!(turn_of(&mut database, "early"), 1_i32); - assert_eq!(turn_of(&mut database, "main-2"), 2_i32); assert_eq!( - turns(&mut database), + day_totals(&rows), vec![ - (Some("p1".to_owned()), 1_i32, second(1), Some(second(10))), - (Some("p2".to_owned()), 2_i32, second(30), Some(second(30))), + (json!("2026-03-07"), json!(1_i64)), + (json!("2026-03-08"), json!(1_i64)), + (json!("2026-03-09"), json!(1_i64)), ] ); } #[test] - fn turns_are_numbered_by_start_time_whatever_order_they_arrive_in() { - let mut database = database(); - ingest_turn_events( - &mut database, - vec![ - turn_event("b", "codex", 30, Some("b")), - turn_event("b-next", "codex", 31, Some("b")), - turn_event("c", "codex", 50, Some("c")), - ], - ); - assert_eq!(turn_of(&mut database, "c"), 2_i32); + fn summary_fills_owner_and_device_dimensions() { + let mut database = summary_seeded(); - // A turn delivered late started first. - ingest_turn_events(&mut database, vec![turn_event("a", "codex", 10, Some("a"))]); - for (event, turn) in [("a", 1_i32), ("b", 2_i32), ("b-next", 2_i32), ("c", 3_i32)] { - assert_eq!(turn_of(&mut database, event), turn, "{event}"); - } + let rows = summary(&mut database, json!({"group_by": "user,device"})); - // A late event of a known turn moves its start before every other turn. - ingest_turn_events( - &mut database, - vec![turn_event("c-early", "codex", 5, Some("c"))], - ); - for (event, turn) in [ - ("c-early", 1_i32), - ("c", 1_i32), - ("a", 2_i32), - ("b", 3_i32), - ("b-next", 3_i32), - ] { - assert_eq!(turn_of(&mut database, event), turn, "{event}"); - } - assert_eq!( - turns(&mut database) - .into_iter() - .map(|(key, index, ..)| (key, index)) - .collect::>(), - vec![ - (Some("c".to_owned()), 1_i32), - (Some("a".to_owned()), 2_i32), - (Some("b".to_owned()), 3_i32), - ] - ); + assert_eq!(rows.as_array().map(Vec::len), Some(1)); + assert_eq!(rows[0]["user_id"], json!(OWNER_ID)); + assert!(rows[0]["user_email"].is_string()); + assert_eq!(rows[0]["host_name"], json!("alice-mbp")); + assert_eq!(rows[0]["platform"], json!("macos")); + assert_eq!(rows[0]["os_version"], json!("15.1")); + assert_eq!(rows[0]["sessions"], json!(3_i64)); + assert!(rows[0].get("agent_name").is_none()); } + /// One request observed `second` seconds into a fixed minute, working for + /// `turn` when it names one. fn request_with_text(text: &str) -> IngestEventsRequest { serde_json::from_value(json!({ "events": [{ @@ -2036,9 +2498,11 @@ mod tests { "session_id": "sqlite-session", "llm": {"provider": "openai", "model": "gpt-5"}, "event_type": "response", - "text": text, "token_usage": {"input_tokens": 0_i64, "output_tokens": 3_i64, "cache_read_tokens": 0_i64, "cache_write_tokens": 0_i64, "reasoning_tokens": 0_i64, "total_tokens": 3_i64}, - "metadata": {} + "conversation": { + "entries": [{"role": "assistant", "text": text}], + "completeness": "complete" + } }] })) .expect("test ingest request should deserialize") diff --git a/crates/abyss-backend/src/storage/sqlite/schema.rs b/crates/abyss-backend/src/storage/sqlite/schema.rs index 45eaec8..88f3b46 100644 --- a/crates/abyss-backend/src/storage/sqlite/schema.rs +++ b/crates/abyss-backend/src/storage/sqlite/schema.rs @@ -1,4 +1,5 @@ -//! Diesel schema for SQLite relational tables; the FTS5 virtual table stays raw SQL. +//! Diesel schema for SQLite relational tables; the FTS5 trigram index is a +//! virtual table outside it. diesel::table! { use diesel::sql_types::{BigInt, Nullable, Text}; @@ -42,44 +43,25 @@ diesel::table! { metadata -> Text, created_at -> BigInt, updated_at -> BigInt, + structure_revision -> BigInt, } } diesel::table! { use diesel::sql_types::{BigInt, Integer, Nullable, Text}; - agent_turns (id) { - id -> Text, - user_id -> Text, - session_pk -> Text, - turn_index -> Integer, - turn_key -> Nullable, - started_at -> BigInt, - ended_at -> Nullable, - created_at -> BigInt, - updated_at -> BigInt, - } -} - -diesel::table! { - use diesel::sql_types::{BigInt, Binary, Integer, Nullable, Text}; - llm_usage_events (id) { id -> Text, user_id -> Text, device_context_id -> Text, session_pk -> Text, - turn_pk -> Text, event_id -> Text, agent_name -> Text, agent_version -> Nullable, session_id -> Text, - turn_index -> Integer, llm_provider -> Text, llm_model -> Text, event_type -> Text, - text -> Nullable, - text_sha256 -> Nullable, input_tokens -> BigInt, output_tokens -> BigInt, cache_read_tokens -> BigInt, @@ -92,6 +74,30 @@ diesel::table! { stream_id -> Nullable, response_id -> Nullable, previous_response_id -> Nullable, + conversation -> Nullable, + } +} + +diesel::table! { + use diesel::sql_types::{Bool, Integer, Nullable, Text}; + + llm_usage_event_stream_facts (event_pk) { + event_pk -> Text, + new_from -> Nullable, + continues_context -> Nullable, + starts_turn -> Nullable, + last_identity -> Nullable, + previous_identity -> Nullable, + trailing -> Nullable, + result_ids -> Nullable, + emitted_ids -> Nullable, + output_skeleton -> Nullable, + label_index -> Nullable, + label_text -> Nullable, + label_length -> Nullable, + subagent_id -> Nullable, + parent_subagent_id -> Nullable, + turn_id -> Nullable, } } @@ -152,17 +158,37 @@ diesel::table! { } } +diesel::table! { + use diesel::sql_types::{BigInt, Binary, Text}; + + search_texts (id) { + id -> BigInt, + user_id -> Text, + sha256 -> Binary, + text -> Text, + } +} + +diesel::table! { + use diesel::sql_types::{BigInt, Integer, Text}; + + search_refs (event_pk, entry_index, field) { + event_pk -> Text, + entry_index -> Integer, + field -> Text, + text_id -> BigInt, + } +} + diesel::joinable!(agent_sessions -> app_users (user_id)); diesel::joinable!(agent_sessions -> devices (device_context_id)); -diesel::joinable!(agent_turns -> agent_sessions (session_pk)); -diesel::joinable!(agent_turns -> app_users (user_id)); diesel::joinable!(devices -> app_users (user_id)); diesel::joinable!(llm_usage_events -> agent_sessions (session_pk)); -diesel::joinable!(llm_usage_events -> agent_turns (turn_pk)); diesel::joinable!(llm_usage_events -> app_users (user_id)); diesel::joinable!(llm_usage_events -> devices (device_context_id)); diesel::joinable!(llm_usage_event_attachments -> app_users (user_id)); diesel::joinable!(llm_usage_event_attachments -> llm_usage_events (event_pk)); +diesel::joinable!(llm_usage_event_stream_facts -> llm_usage_events (event_pk)); diesel::joinable!(agent_diagnostic_capture_events -> agent_diagnostic_captures (capture_pk)); diesel::joinable!(agent_diagnostic_capture_events -> llm_usage_events (event_pk)); diesel::joinable!(agent_diagnostic_captures -> agent_sessions (session_pk)); @@ -174,10 +200,12 @@ diesel::allow_tables_to_appear_in_same_query!( agent_diagnostic_capture_events, agent_diagnostic_captures, agent_sessions, - agent_turns, app_users, devices, llm_usage_events, llm_usage_event_attachments, + llm_usage_event_stream_facts, + search_refs, + search_texts, session_titles, ); diff --git a/crates/abyss-backend/src/storage/sqlite/search.rs b/crates/abyss-backend/src/storage/sqlite/search.rs deleted file mode 100644 index 2bc5c85..0000000 --- a/crates/abyss-backend/src/storage/sqlite/search.rs +++ /dev/null @@ -1,744 +0,0 @@ -//! SQLite FTS5 raw queries with Diesel ORM hydration of relational session data. - -use std::collections::{HashMap, HashSet}; - -use chrono::{DateTime, Utc}; -use diesel::{ - ExpressionMethods, QueryDsl, QueryableByName, RunQueryDsl, SqliteConnection, - query_builder::{BoxedSqlQuery, SqlQuery}, - sql_query, - sql_types::{BigInt, Integer, Nullable, Text}, - sqlite::Sqlite, -}; -use uuid::Uuid; - -use crate::{ - error::AppError, - search::{ - HIGHLIGHT_END, HIGHLIGHT_START, SearchFragment, SessionSearchMatch, SessionSearchQuery, - SessionSearchResponse, SessionSearchResult, ValidatedSearchQuery, - }, -}; - -use super::{ - models::{parse_uuid, timestamp_from_micros}, - repository::timestamp_to_micros, - schema::{agent_sessions, devices}, -}; - -pub(super) fn session_search( - connection: &mut SqliteConnection, - user_id: Uuid, - query: SessionSearchQuery, -) -> Result { - let query = query.validate()?; - let hits = count_hits(connection, &query)?; - search(connection, user_id, query, hits) -} - -/// Runs a validated search; `hits` only chooses how observation times are read. -fn search( - connection: &mut SqliteConnection, - user_id: Uuid, - query: ValidatedSearchQuery, - hits: i64, -) -> Result { - let source = MatchSource::new(&query, user_id, hits); - let (page, total_sessions) = load_session_page(connection, &query, &source)?; - let session_ids = page - .iter() - .map(|(session_pk, _)| *session_pk) - .collect::>(); - let mut matches = load_matches(connection, &session_ids, &source)?; - let mut details = load_session_details(connection, user_id, &session_ids)?; - - let items = page - .into_iter() - .filter_map(|(session_pk, match_count)| { - let details = details.remove(&session_pk)?; - let session_matches = matches.remove(&session_pk).unwrap_or_default(); - let mut providers = session_matches - .iter() - .map(|matched| matched.llm_provider.clone()) - .collect::>() - .into_iter() - .collect::>(); - providers.sort(); - let mut models = session_matches - .iter() - .map(|matched| matched.llm_model.clone()) - .collect::>() - .into_iter() - .collect::>(); - models.sort(); - Some(SessionSearchResult { - session_pk, - session_id: details.session_id, - agent_name: details.agent_name, - agent_version: details.agent_version, - host_name: details.host_name, - platform: details.platform, - started_at: details.started_at, - ended_at: details.ended_at, - providers, - models, - match_count, - matches: session_matches - .into_iter() - .map(MatchedEvent::into_response) - .collect(), - }) - }) - .collect(); - let consumed = u64::from(query.offset).saturating_add(u64::from(query.page_size)); - Ok(SessionSearchResponse { - query: query.text, - total_sessions, - page: query.page, - page_size: query.page_size, - has_more: consumed < total_sessions, - items, - }) -} - -/// Matches the strongest events rank first: the lowest bm25, then the latest -/// observation, then the event key, so results are deterministic. -const MATCHES_PER_SESSION: i64 = 3; - -#[derive(Clone)] -enum FtsBind { - Text(String), - BigInt(i64), -} - -/// Above this many FTS hits, observation times come from the owner-time -/// covering index for all of the owner's events at once; below it, each hit -/// reads its own event row, which is cheaper than building that index. -const ROW_LOOKUP_MAX_HITS: i64 = 256; - -/// Where a page's session list is inserted into the shared expressions. -const SESSION_FILTER: &str = "/* session filter */"; - -/// Common table expressions that list every matching event of the owner with -/// its rank, observation time, and position and count within its session. -/// -/// Every step avoids reading wide rows: `hits` takes only rowids and bm25 from -/// the FTS index; event and session keys come from FTS5's content table rather -/// than through the virtual table, which would load each matching document. -/// Observation times sit after the event text in its row, so for many hits -/// they come from the owner-time covering index instead, and the time window -/// is applied after the join so SQLite keeps an index on that list. Event rows -/// are otherwise joined only for filters whose columns precede the text. -struct MatchSource { - sql: String, - binds: Vec, -} - -impl MatchSource { - fn new(query: &ValidatedSearchQuery, user_id: Uuid, hits: i64) -> Self { - let mut binds = vec![FtsBind::Text(fts_query(&query.text))]; - let read_rows = hits <= ROW_LOOKUP_MAX_HITS; - let (events_cte, joins, time_column) = if read_rows { - ( - "", - "INNER JOIN llm_usage_events e ON e.id = c.c0", - "e.observed_at", - ) - } else { - binds.push(FtsBind::Text(user_id.to_string())); - ( - ", - events AS MATERIALIZED ( - SELECT id, observed_at FROM llm_usage_events WHERE user_id = ? - )", - "INNER JOIN events ev ON ev.id = c.c0", - "ev.observed_at", - ) - }; - let mut conditions = String::from("1 = 1"); - if read_rows { - conditions.push_str(" AND e.user_id = ?"); - binds.push(FtsBind::Text(user_id.to_string())); - } - if let Some(from) = query.from { - conditions.push_str(" AND "); - conditions.push_str(time_column); - conditions.push_str(" >= ?"); - binds.push(FtsBind::BigInt(timestamp_to_micros(from))); - } - if let Some(to) = query.to { - conditions.push_str(" AND "); - conditions.push_str(time_column); - conditions.push_str(" < ?"); - binds.push(FtsBind::BigInt(timestamp_to_micros(to))); - } - let mut filters = String::new(); - push_any_filter(&mut filters, &mut binds, "e.agent_name", &query.agent_names); - push_any_filter( - &mut filters, - &mut binds, - "e.llm_provider", - &query.llm_providers, - ); - push_any_filter(&mut filters, &mut binds, "e.llm_model", &query.llm_models); - push_exact_filter( - &mut filters, - &mut binds, - "e.event_type", - query.event_type.as_deref(), - ); - let event_join = if read_rows || filters.is_empty() { - "" - } else { - "INNER JOIN llm_usage_events e ON e.id = c.c0" - }; - let sql = format!( - "WITH hits AS MATERIALIZED ( - SELECT rowid AS fts_rowid, bm25(usage_events_fts) AS rank - FROM usage_events_fts - WHERE usage_events_fts MATCH ? - ){events_cte}, - matched AS MATERIALIZED ( - SELECT c.c0 AS event_pk, c.c2 AS session_pk, h.fts_rowid, h.rank, - {time_column} AS observed_at - FROM hits h - INNER JOIN usage_events_fts_content c ON c.id = h.fts_rowid - {joins} - {event_join} - WHERE {conditions}{filters} {SESSION_FILTER} - ), - ranked AS ( - SELECT matched.*, - ROW_NUMBER() OVER ( - PARTITION BY session_pk - ORDER BY rank ASC, observed_at DESC, event_pk ASC - ) AS position, - COUNT(*) OVER (PARTITION BY session_pk) AS match_count - FROM matched - )" - ); - Self { sql, binds } - } - - /// The expressions with matches limited to `session_ids`, or all sessions. - fn sql_for(&self, session_ids: &[Uuid]) -> (String, Vec) { - let mut binds = self.binds.clone(); - let session_filter = if session_ids.is_empty() { - String::new() - } else { - binds.extend( - session_ids - .iter() - .map(|session_pk| FtsBind::Text(session_pk.to_string())), - ); - format!("AND c.c2 IN ({})", vec!["?"; session_ids.len()].join(", ")) - }; - (self.sql.replace(SESSION_FILTER, &session_filter), binds) - } -} - -fn push_exact_filter( - sql: &mut String, - binds: &mut Vec, - column: &str, - value: Option<&str>, -) { - if let Some(value) = value { - sql.push_str(" AND "); - sql.push_str(column); - sql.push_str(" = ?"); - binds.push(FtsBind::Text(value.to_owned())); - } -} - -/// Matches any of `values`; an empty list adds no filter. -fn push_any_filter(sql: &mut String, binds: &mut Vec, column: &str, values: &[String]) { - if values.is_empty() { - return; - } - sql.push_str(" AND "); - sql.push_str(column); - sql.push_str(" IN ("); - sql.push_str(&vec!["?"; values.len()].join(", ")); - sql.push(')'); - binds.extend(values.iter().cloned().map(FtsBind::Text)); -} - -fn bound_query(sql: String, binds: Vec) -> BoxedSqlQuery<'static, Sqlite, SqlQuery> { - let mut query = sql_query(sql).into_boxed::(); - for bind in binds { - query = match bind { - FtsBind::Text(value) => query.bind::(value), - FtsBind::BigInt(value) => query.bind::(value), - }; - } - query -} - -#[derive(QueryableByName)] -struct CountRow { - #[diesel(sql_type = BigInt)] - count: i64, -} - -/// FTS hits across all owners, before any filter; reads only the index. -fn count_hits( - connection: &mut SqliteConnection, - query: &ValidatedSearchQuery, -) -> Result { - Ok(bound_query( - "SELECT COUNT(*) AS count FROM usage_events_fts WHERE usage_events_fts MATCH ?".to_owned(), - vec![FtsBind::Text(fts_query(&query.text))], - ) - .get_result::(connection)? - .count) -} - -#[derive(QueryableByName)] -struct SessionPageRow { - #[diesel(sql_type = Text)] - session_pk: String, - #[diesel(sql_type = BigInt)] - match_count: i64, - #[diesel(sql_type = BigInt)] - total_sessions: i64, -} - -/// One page of sessions with their match counts, ordered by each session's -/// strongest event, and the number of matching sessions. -fn load_session_page( - connection: &mut SqliteConnection, - query: &ValidatedSearchQuery, - source: &MatchSource, -) -> Result<(Vec<(Uuid, u64)>, u64), AppError> { - let (sql, mut binds) = source.sql_for(&[]); - let page_sql = format!( - "{sql} - SELECT session_pk, match_count, COUNT(*) OVER () AS total_sessions - FROM ranked - WHERE position = 1 - ORDER BY rank ASC, observed_at DESC, session_pk ASC - LIMIT ? OFFSET ?" - ); - binds.push(FtsBind::BigInt(i64::from(query.page_size))); - binds.push(FtsBind::BigInt(i64::from(query.offset))); - let rows = bound_query(page_sql, binds).load::(connection)?; - let total_sessions = match rows.first() { - Some(row) => row.total_sessions, - // A page past the end has no rows to carry the total. - None if query.offset > 0 => { - let (sql, binds) = source.sql_for(&[]); - bound_query( - format!("{sql} SELECT COUNT(*) AS count FROM ranked WHERE position = 1"), - binds, - ) - .get_result::(connection)? - .count - } - None => 0, - }; - let page = rows - .into_iter() - .map(|row| { - Ok(( - parse_uuid(&row.session_pk, "FTS session id")?, - u64::try_from(row.match_count.max(0)).map_err(|error| { - AppError::internal(format!("invalid FTS match count: {error}")) - })?, - )) - }) - .collect::, AppError>>()?; - let total_sessions = u64::try_from(total_sessions.max(0)) - .map_err(|error| AppError::internal(format!("invalid FTS session count: {error}")))?; - Ok((page, total_sessions)) -} - -/// The strongest matches of each session, with highlighted fragments built -/// only for the events returned. -fn load_matches( - connection: &mut SqliteConnection, - session_ids: &[Uuid], - source: &MatchSource, -) -> Result>, AppError> { - if session_ids.is_empty() { - return Ok(HashMap::new()); - } - let (sql, mut binds) = source.sql_for(session_ids); - let matches_sql = format!( - "{sql}, - top AS MATERIALIZED ( - SELECT * FROM ranked WHERE position <= {MATCHES_PER_SESSION} - ) - SELECT - top.event_pk AS event_pk, - top.session_pk AS session_pk, - e.turn_pk AS turn_pk, - e.turn_index AS turn_index, - e.event_type AS event_type, - e.llm_provider AS llm_provider, - e.llm_model AS llm_model, - top.observed_at AS observed_at, - snippet(usage_events_fts, 4, '{HIGHLIGHT_START}', '{HIGHLIGHT_END}', '…', 32) - AS content_fragment, - snippet(usage_events_fts, 5, '{HIGHLIGHT_START}', '{HIGHLIGHT_END}', '…', 32) - AS tool_name_fragment, - snippet(usage_events_fts, 6, '{HIGHLIGHT_START}', '{HIGHLIGHT_END}', '…', 32) - AS tool_content_fragment - FROM top - INNER JOIN usage_events_fts ON usage_events_fts.rowid = top.fts_rowid - INNER JOIN llm_usage_events e ON e.id = top.event_pk - WHERE usage_events_fts MATCH ? - ORDER BY top.session_pk, top.position" - ); - binds.push(source.binds[0].clone()); - let rows = bound_query(matches_sql, binds).load::(connection)?; - let mut by_session: HashMap> = HashMap::new(); - for row in rows { - let matched = row.into_matched()?; - by_session - .entry(matched.session_pk) - .or_default() - .push(matched); - } - Ok(by_session) -} - -#[derive(QueryableByName)] -struct MatchedEventRow { - #[diesel(sql_type = Text)] - event_pk: String, - #[diesel(sql_type = Text)] - session_pk: String, - #[diesel(sql_type = Text)] - turn_pk: String, - #[diesel(sql_type = Integer)] - turn_index: i32, - #[diesel(sql_type = Text)] - event_type: String, - #[diesel(sql_type = Text)] - llm_provider: String, - #[diesel(sql_type = Text)] - llm_model: String, - #[diesel(sql_type = BigInt)] - observed_at: i64, - #[diesel(sql_type = Nullable)] - content_fragment: Option, - #[diesel(sql_type = Nullable)] - tool_name_fragment: Option, - #[diesel(sql_type = Nullable)] - tool_content_fragment: Option, -} - -impl MatchedEventRow { - fn into_matched(self) -> Result { - let fragments = [ - self.content_fragment, - self.tool_name_fragment, - self.tool_content_fragment, - ] - .into_iter() - .flatten() - .filter(|fragment| fragment.contains(HIGHLIGHT_START)) - .take(2) - .collect(); - Ok(MatchedEvent { - event_pk: parse_uuid(&self.event_pk, "FTS event id")?, - session_pk: parse_uuid(&self.session_pk, "FTS session id")?, - turn_pk: parse_uuid(&self.turn_pk, "FTS turn id")?, - turn_index: self.turn_index, - event_type: self.event_type, - llm_provider: self.llm_provider, - llm_model: self.llm_model, - observed_at: timestamp_from_micros(self.observed_at)?, - fragments, - }) - } -} - -struct MatchedEvent { - event_pk: Uuid, - session_pk: Uuid, - turn_pk: Uuid, - turn_index: i32, - event_type: String, - llm_provider: String, - llm_model: String, - observed_at: DateTime, - fragments: Vec, -} - -impl MatchedEvent { - fn into_response(self) -> SessionSearchMatch { - SessionSearchMatch { - event_pk: self.event_pk, - turn_pk: self.turn_pk, - turn_index: self.turn_index, - event_type: self.event_type, - llm_provider: self.llm_provider, - llm_model: self.llm_model, - observed_at: self.observed_at, - fragments: self - .fragments - .into_iter() - .map(|fragment| SearchFragment::parse(&fragment)) - .collect(), - } - } -} - -struct SessionDetails { - session_id: String, - agent_name: String, - agent_version: Option, - host_name: String, - platform: String, - started_at: DateTime, - ended_at: Option>, -} - -fn load_session_details( - connection: &mut SqliteConnection, - user_id: Uuid, - session_ids: &[Uuid], -) -> Result, AppError> { - if session_ids.is_empty() { - return Ok(HashMap::new()); - } - let ids = session_ids.iter().map(Uuid::to_string).collect::>(); - let rows = agent_sessions::table - .inner_join(devices::table) - .filter(agent_sessions::user_id.eq(user_id.to_string())) - .filter(agent_sessions::id.eq_any(ids)) - .select(( - agent_sessions::id, - agent_sessions::session_id, - agent_sessions::agent_name, - agent_sessions::agent_version, - devices::host_name, - devices::platform, - agent_sessions::started_at, - agent_sessions::ended_at, - )) - .load::<( - String, - String, - String, - Option, - String, - String, - i64, - Option, - )>(connection)?; - rows.into_iter() - .map( - |( - session_pk, - session_id, - agent_name, - agent_version, - host_name, - platform, - started_at, - ended_at, - )| { - Ok(( - parse_uuid(&session_pk, "session id")?, - SessionDetails { - session_id, - agent_name, - agent_version, - host_name, - platform, - started_at: timestamp_from_micros(started_at)?, - ended_at: ended_at.map(timestamp_from_micros).transpose()?, - }, - )) - }, - ) - .collect() -} - -fn fts_query(value: &str) -> String { - value - .split_whitespace() - .map(|term| format!("\"{}\"", term.replace('"', "\"\""))) - .collect::>() - .join(" AND ") -} - -#[cfg(test)] -mod tests { - use diesel::{Connection, SqliteConnection}; - use serde_json::{Value, json}; - use uuid::Uuid; - - use crate::{ - search::{SessionSearchQuery, SessionSearchResponse}, - usage::IngestEventsRequest, - }; - - use super::{ - super::{migrations, repository::ingest_events}, - ROW_LOOKUP_MAX_HITS, fts_query, search, - }; - - const OWNER_ID: Uuid = Uuid::from_u128(1); - - #[test] - fn fts_query_uses_literal_and_terms() { - assert_eq!( - fts_query("blackbox response"), - r#""blackbox" AND "response""# - ); - assert_eq!(fts_query(r#"quoted"value"#), r#""quoted""value""#); - } - - fn event(id: &str, session: &str, observed_at: &str, text: &str) -> Value { - json!({ - "event_id": id, - "observed_at": observed_at, - "device": {"host_name": "alice-mbp", "platform": "macos"}, - "agent": {"name": "codex"}, - "session_id": session, - "turn_index": 1, - "llm": {"provider": "openai", "model": "gpt-5"}, - "event_type": "response", - "text": text, - "token_usage": {"input_tokens": 0_i64, "output_tokens": 1_i64, "cache_read_tokens": 0_i64, "cache_write_tokens": 0_i64, "reasoning_tokens": 0_i64, "total_tokens": 1_i64}, - "metadata": {} - }) - } - - /// Three sessions: `alpha` has five equal hits, `beta` one strong hit, and - /// `gamma` one hit diluted by other words. - fn seeded() -> SqliteConnection { - let mut database = SqliteConnection::establish(":memory:").expect("SQLite should open"); - migrations::run(&mut database).expect("migration should succeed"); - let mut events = (1_u32..=5_u32) - .map(|minute| { - event( - &format!("a{minute}"), - "alpha", - &format!("2026-09-01T00:0{minute}:00Z"), - "needle note", - ) - }) - .collect::>(); - events.push(event("b1", "beta", "2026-09-02T00:00:00Z", "needle")); - events.push(event( - "g1", - "gamma", - "2026-09-03T00:00:00Z", - "needle among many other unrelated words in a much longer reply", - )); - events.push(event("z1", "zeta", "2026-09-04T00:00:00Z", "nothing here")); - let request: IngestEventsRequest = - serde_json::from_value(json!({ "events": events })).expect("request should parse"); - ingest_events(&mut database, &request, OWNER_ID, 100).expect("ingest should succeed"); - database - } - - fn run(database: &mut SqliteConnection, query: Value, hits: i64) -> SessionSearchResponse { - let query: SessionSearchQuery = - serde_json::from_value(query).expect("query should deserialize"); - search( - database, - OWNER_ID, - query.validate().expect("query should validate"), - hits, - ) - .expect("search should succeed") - } - - fn summary(response: &SessionSearchResponse) -> Vec<(String, u64, usize)> { - response - .items - .iter() - .map(|item| { - ( - item.session_id.clone(), - item.match_count, - item.matches.len(), - ) - }) - .collect() - } - - #[test] - fn sessions_rank_by_their_strongest_event_and_page_in_order() { - let mut database = seeded(); - - let first = run(&mut database, json!({"q": "needle", "page_size": 2_u32}), 0); - let second = run( - &mut database, - json!({"q": "needle", "page": 2_u32, "page_size": 2_u32}), - 0, - ); - let past_end = run( - &mut database, - json!({"q": "needle", "page": 5_u32, "page_size": 2_u32}), - 0, - ); - - assert_eq!(first.total_sessions, 3); - assert!(first.has_more); - assert_eq!( - summary(&first), - [("beta".to_owned(), 1, 1), ("alpha".to_owned(), 5, 3)] - ); - assert_eq!(summary(&second), [("gamma".to_owned(), 1, 1)]); - assert!(!second.has_more); - assert_eq!(past_end.total_sessions, 3); - assert!(past_end.items.is_empty()); - } - - #[test] - fn a_session_returns_its_three_latest_equal_matches_and_counts_all() { - let mut database = seeded(); - - let response = run(&mut database, json!({"q": "needle note"}), 0); - - assert_eq!(summary(&response), [("alpha".to_owned(), 5, 3)]); - let minutes = response.items[0] - .matches - .iter() - .map(|matched| matched.observed_at.format("%M").to_string()) - .collect::>(); - // Equal scores fall back to the latest observation first. - assert_eq!(minutes, ["05", "04", "03"]); - assert!( - response.items[0].matches[0].fragments[0] - .segments - .iter() - .any(|segment| segment.highlighted) - ); - } - - #[test] - fn both_time_lookups_return_the_same_results_with_filters() { - let mut database = seeded(); - for query in [ - json!({"q": "needle"}), - json!({"q": "needle", "from": "2026-09-01T00:03:00Z", "to": "2026-09-02T12:00:00Z"}), - json!({"q": "needle", "agent_name": ["codex"], "event_type": "response", "page_size": 1_u32}), - json!({"q": "needle", "llm_model": ["absent"]}), - ] { - let by_rows = run(&mut database, query.clone(), 0); - let by_index = run(&mut database, query.clone(), ROW_LOOKUP_MAX_HITS + 1); - assert_eq!( - serde_json::to_value(&by_rows).expect("serializes"), - serde_json::to_value(&by_index).expect("serializes"), - "{query}" - ); - } - let windowed = run( - &mut database, - json!({"q": "needle", "from": "2026-09-01T00:03:00Z", "to": "2026-09-02T12:00:00Z"}), - ROW_LOOKUP_MAX_HITS + 1, - ); - assert_eq!( - summary(&windowed), - [("beta".to_owned(), 1, 1), ("alpha".to_owned(), 3, 3)] - ); - } -} diff --git a/crates/abyss-backend/src/storage/sqlite/search/index.rs b/crates/abyss-backend/src/storage/sqlite/search/index.rs new file mode 100644 index 0000000..332cd7b --- /dev/null +++ b/crates/abyss-backend/src/storage/sqlite/search/index.rs @@ -0,0 +1,80 @@ +//! The search index of one event, written in its ingest transaction: its +//! distinct texts once per owner with their folded trigrams, and one +//! reference per piece. + +use diesel::{ + ExpressionMethods, OptionalExtension, QueryDsl, RunQueryDsl, SqliteConnection, sql_query, + sql_types::{BigInt, Text}, +}; +use serde_json::Value; +use sha2::{Digest as _, Sha256}; + +use crate::{ + error::AppError, + search::matching::{fold::fold, pieces::pieces}, +}; + +use super::super::{ + models::{NewSearchRef, NewSearchText}, + schema::{search_refs, search_texts}, +}; + +/// References written per insert statement, well below SQLite's bound +/// parameter limit. +const REFS_PER_INSERT: usize = 1_000; + +/// Indexes every piece of one newly stored event of `owner`. Only this +/// event is read: a text is found again by its hash, never by comparing it +/// with other events. +pub(in crate::storage::sqlite) fn index_event( + connection: &mut SqliteConnection, + owner: &str, + event_pk: &str, + conversation: Option<&Value>, +) -> Result<(), AppError> { + let mut refs = Vec::new(); + for piece in pieces(conversation) { + refs.push(NewSearchRef { + event_pk, + entry_index: i32::try_from(piece.entry_index).map_err(|error| { + AppError::validation(format!("conversation has too many entries: {error}")) + })?, + field: piece.field.as_str(), + text_id: text_id(connection, owner, piece.text)?, + }); + } + for chunk in refs.chunks(REFS_PER_INSERT) { + diesel::insert_into(search_refs::table) + .values(chunk) + .execute(connection)?; + } + Ok(()) +} + +/// The id of `text` among the owner's texts. A new text is inserted, and +/// its folded form enters the trigram index under the same id. +fn text_id(connection: &mut SqliteConnection, owner: &str, text: &str) -> Result { + let sha256 = Sha256::digest(text.as_bytes()); + let inserted = diesel::insert_or_ignore_into(search_texts::table) + .values(NewSearchText { + user_id: owner, + sha256: sha256.as_slice(), + text, + }) + .returning(search_texts::id) + .get_result::(connection) + .optional()?; + if let Some(id) = inserted { + // The trigram index is a virtual table outside Diesel's schema. + sql_query("INSERT INTO search_texts_trigram (rowid, folded) VALUES (?, ?)") + .bind::(id) + .bind::(fold(text)) + .execute(connection)?; + return Ok(id); + } + Ok(search_texts::table + .filter(search_texts::user_id.eq(owner)) + .filter(search_texts::sha256.eq(sha256.as_slice())) + .select(search_texts::id) + .first::(connection)?) +} diff --git a/crates/abyss-backend/src/storage/sqlite/search/mod.rs b/crates/abyss-backend/src/storage/sqlite/search/mod.rs new file mode 100644 index 0000000..21401c7 --- /dev/null +++ b/crates/abyss-backend/src/storage/sqlite/search/mod.rs @@ -0,0 +1,527 @@ +//! SQLite search over the distinct texts of entries: the trigram index +//! narrows the candidate texts, the shared matcher verifies each of them, +//! and references expand the matching texts into hits of events. + +mod index; +#[cfg(test)] +mod tests; + +use std::collections::HashMap; + +use diesel::{ + Connection, ExpressionMethods, OptionalExtension, QueryDsl, QueryableByName, RunQueryDsl, + SqliteConnection, + connection::DefaultLoadingMode, + query_builder::{BoxedSqlQuery, SqlQuery}, + sql_query, + sql_types::{BigInt, Integer, Nullable, Text}, + sqlite::Sqlite, +}; +use uuid::Uuid; + +use crate::{ + error::AppError, + search::{ + InSessionHit, InSessionSearchQuery, InSessionSearchResponse, PieceField, + SessionSearchQuery, SessionSearchResponse, ValidatedSearchQuery, + matching::{ + fragment::fragment, + in_session::MAX_IN_SESSION_HITS, + pattern::{Pattern, utf16_ranges}, + ranking::{MatchedEvent, SearchHit, SessionDetails, is_new, rank_sessions}, + }, + }, + storage::stream_cache::StreamCache, + usage::UsageEventType, +}; + +use super::{ + models::{parse_uuid, timestamp_from_micros}, + repository::{session_structure, timestamp_to_micros}, + schema::{agent_sessions, devices, search_texts}, + session_titles, +}; + +pub(super) use self::index::index_event; + +pub(super) fn session_search( + connection: &mut SqliteConnection, + user_id: Uuid, + query: SessionSearchQuery, +) -> Result { + let query = query.validate()?; + let pattern = Pattern::new(&query.text); + let owner = user_id.to_string(); + let offset = usize::try_from(query.offset) + .map_err(|error| AppError::internal(format!("search offset: {error}")))?; + let page_size = usize::try_from(query.page_size) + .map_err(|error| AppError::internal(format!("search page size: {error}")))?; + // One read transaction keeps the hits, the session details, and the + // fragments on the same snapshot. + connection.transaction(|connection| { + let mut matching = Vec::new(); + for text in candidates(&owner, &pattern, None) + .load_iter::(connection)? + { + let text = text?; + if pattern.matches(&text.text) { + matching.push(text.id); + } + } + let hits = session_hits(connection, &owner, &matching, &query)?; + let page = rank_sessions(hits, offset, page_size); + let details = session_details(connection, &owner, &page.session_pks())?; + let events = matched_events(connection, &owner, &page.event_pks())?; + let fragments = search_texts::table + .filter(search_texts::user_id.eq(&owner)) + .filter(search_texts::id.eq_any(page.fragment_texts())) + .select((search_texts::id, search_texts::text)) + .load::<(i64, String)>(connection)? + .into_iter() + .filter_map(|(id, text)| Some((id, fragment(&text, &pattern.occurrences(&text))?))) + .collect(); + let mut response = page.into_response(query, details, &events, &fragments); + let sessions = response + .items + .iter() + .map(|item| item.session_pk) + .collect::>(); + let mut titles = session_titles::titles(connection, user_id, &sessions)?; + for item in &mut response.items { + item.title = titles.remove(&item.session_pk); + } + Ok(response) + }) +} + +/// Every hit in one session, each with the head of the segment holding its +/// event, derived from the session's stream facts as its streams are. +pub(super) fn search_in_session( + connection: &mut SqliteConnection, + user_id: Uuid, + session_pk: Uuid, + query: &InSessionSearchQuery, + cache: &StreamCache, +) -> Result { + let text = query.validate()?; + let pattern = Pattern::new(&text); + let owner = user_id.to_string(); + let session = session_pk.to_string(); + connection.transaction(|connection| { + let exists = agent_sessions::table + .filter(agent_sessions::user_id.eq(&owner)) + .filter(agent_sessions::id.eq(&session)) + .select(agent_sessions::id) + .first::(connection) + .optional()? + .is_some(); + if !exists { + return Err(AppError::not_found(format!( + "session not found: {session_pk}" + ))); + } + let mut ranges = HashMap::new(); + for candidate in candidates(&owner, &pattern, Some(&session)) + .load_iter::(connection)? + { + let candidate = candidate?; + let occurrences = pattern.occurrences(&candidate.text); + if !occurrences.is_empty() { + ranges.insert(candidate.id, utf16_ranges(&candidate.text, &occurrences)); + } + } + let matching = ranges.keys().copied().collect::>(); + let mut hits = hits_in_session(connection, &owner, &session, &matching)?; + if !hits.is_empty() { + let structure = session_structure(connection, user_id, session_pk, cache)?; + for hit in &mut hits { + hit.segment = structure.event_segment(hit.event_pk).map(str::to_owned); + } + } + Ok(InSessionSearchResponse::from_hits(text, hits, ranges)) + }) +} + +#[derive(Clone)] +enum Bind { + Text(String), + BigInt(i64), +} + +fn bound_query(sql: String, binds: Vec) -> BoxedSqlQuery<'static, Sqlite, SqlQuery> { + let mut query = sql_query(sql).into_boxed::(); + for bind in binds { + query = match bind { + Bind::Text(value) => query.bind::(value), + Bind::BigInt(value) => query.bind::(value), + }; + } + query +} + +#[derive(QueryableByName)] +struct TextRow { + #[diesel(sql_type = BigInt)] + id: i64, + #[diesel(sql_type = Text)] + text: String, +} + +/// Every text of the owner `?`. +const OWNER_TEXTS_SQL: &str = + "SELECT t.id AS id, t.text AS text FROM search_texts AS t WHERE t.user_id = ?"; + +/// The texts matching the trigram expression `?1` of the owner `?2`. +const TRIGRAM_TEXTS_SQL: &str = "\ + SELECT t.id AS id, t.text AS text \ + FROM search_texts_trigram AS f \ + CROSS JOIN search_texts AS t ON t.id = f.rowid \ + WHERE search_texts_trigram MATCH ? AND t.user_id = ?"; + +/// The owner's texts that may hold the query: for three characters or +/// more, those whose folded form has every trigram of the folded query, +/// else all of them. With a session key, only texts its events refer to. +fn candidates( + owner: &str, + pattern: &Pattern, + session_pk: Option<&str>, +) -> BoxedSqlQuery<'static, Sqlite, SqlQuery> { + // The trigram index, when it applies, drives the lookup; the owner + // filter narrows little, since most texts belong to one owner. + let (mut sql, mut binds) = pattern.trigram_expression().map_or_else( + || { + ( + String::from(OWNER_TEXTS_SQL), + vec![Bind::Text(owner.to_owned())], + ) + }, + |expression| { + ( + String::from(TRIGRAM_TEXTS_SQL), + vec![Bind::Text(expression), Bind::Text(owner.to_owned())], + ) + }, + ); + if let Some(session_pk) = session_pk { + sql.push_str( + " AND t.id IN (SELECT r.text_id FROM llm_usage_events AS e \ + CROSS JOIN search_refs AS r ON r.event_pk = e.id \ + WHERE e.user_id = ? AND e.session_pk = ?)", + ); + binds.push(Bind::Text(owner.to_owned())); + binds.push(Bind::Text(session_pk.to_owned())); + } + bound_query(sql, binds) +} + +#[derive(QueryableByName)] +struct HitRow { + #[diesel(sql_type = Text)] + event_pk: String, + #[diesel(sql_type = Text)] + session_pk: String, + #[diesel(sql_type = Text)] + session_id: String, + #[diesel(sql_type = BigInt)] + observed_at: i64, + #[diesel(sql_type = Text)] + event_type: String, + #[diesel(sql_type = Nullable)] + new_from: Option, + #[diesel(sql_type = Integer)] + entry_index: i32, + #[diesel(sql_type = Text)] + field: String, + #[diesel(sql_type = BigInt)] + text_id: i64, +} + +/// The pieces of the owner's events that refer to `text_ids`, within the +/// query's filters. The ids travel as one JSON array. A request's `new_from` +/// is in its stream facts; an event stored before they were read has none. +fn session_hits( + connection: &mut SqliteConnection, + owner: &str, + text_ids: &[i64], + query: &ValidatedSearchQuery, +) -> Result, AppError> { + if text_ids.is_empty() { + return Ok(Vec::new()); + } + let mut sql = String::from( + "SELECT r.event_pk AS event_pk, e.session_pk AS session_pk, \ + e.session_id AS session_id, e.observed_at AS observed_at, \ + e.event_type AS event_type, f.new_from AS new_from, \ + r.entry_index AS entry_index, r.field AS field, r.text_id AS text_id \ + FROM json_each(?) AS m \ + CROSS JOIN search_refs AS r ON r.text_id = m.value \ + CROSS JOIN llm_usage_events AS e ON e.id = r.event_pk \ + LEFT JOIN llm_usage_event_stream_facts AS f ON f.event_pk = e.id \ + WHERE e.user_id = ?", + ); + let mut binds = vec![ + Bind::Text(json_ids(text_ids)?), + Bind::Text(owner.to_owned()), + ]; + if let Some(from) = query.from { + sql.push_str(" AND e.observed_at >= ?"); + binds.push(Bind::BigInt(timestamp_to_micros(from))); + } + if let Some(to) = query.to { + sql.push_str(" AND e.observed_at < ?"); + binds.push(Bind::BigInt(timestamp_to_micros(to))); + } + push_any_filter(&mut sql, &mut binds, "e.agent_name", &query.agent_names); + push_any_filter(&mut sql, &mut binds, "e.llm_provider", &query.llm_providers); + push_any_filter(&mut sql, &mut binds, "e.llm_model", &query.llm_models); + if let Some(event_type) = &query.event_type { + sql.push_str(" AND e.event_type = ?"); + binds.push(Bind::Text(event_type.clone())); + } + bound_query(sql, binds) + .load::(connection)? + .into_iter() + .map(|row| { + let entry_index = entry_index(row.entry_index)?; + Ok(SearchHit { + new: is_new( + event_type(&row.event_type)?, + new_from(row.new_from)?, + entry_index, + ), + event_pk: parse_uuid(&row.event_pk, "search event id")?, + session_pk: parse_uuid(&row.session_pk, "search session id")?, + session_id: row.session_id, + observed_at: timestamp_from_micros(row.observed_at)?, + entry_index, + field: field(&row.field)?, + text_id: row.text_id, + }) + }) + .collect() +} + +/// Matches any of `values`; an empty list adds no filter. +fn push_any_filter(sql: &mut String, binds: &mut Vec, column: &str, values: &[String]) { + if values.is_empty() { + return; + } + sql.push_str(" AND "); + sql.push_str(column); + sql.push_str(" IN ("); + sql.push_str(&vec!["?"; values.len()].join(", ")); + sql.push(')'); + binds.extend(values.iter().cloned().map(Bind::Text)); +} + +#[derive(QueryableByName)] +struct SessionHitRow { + #[diesel(sql_type = Text)] + event_pk: String, + #[diesel(sql_type = Text)] + event_type: String, + #[diesel(sql_type = Nullable)] + new_from: Option, + #[diesel(sql_type = Integer)] + entry_index: i32, + #[diesel(sql_type = Text)] + field: String, + #[diesel(sql_type = BigInt)] + text_id: i64, +} + +/// The pieces of one session's events that refer to `text_ids`, by +/// observation, event, entry, and `name` before `text`; one more than the +/// cap, so the caller can tell the hits were cut. Segments are filled in +/// afterwards. +fn hits_in_session( + connection: &mut SqliteConnection, + owner: &str, + session_pk: &str, + text_ids: &[i64], +) -> Result, AppError> { + if text_ids.is_empty() { + return Ok(Vec::new()); + } + let limit = i64::try_from(MAX_IN_SESSION_HITS.saturating_add(1)) + .map_err(|error| AppError::internal(format!("hit limit: {error}")))?; + // The session's references are read by event, and the unary `+` keeps + // the text ids a filter on them, instead of a lookup of every text id + // for every event of the session. + sql_query( + "WITH events AS MATERIALIZED ( \ + SELECT id, observed_at, event_type FROM llm_usage_events \ + WHERE user_id = ?1 AND session_pk = ?2 \ + ) \ + SELECT r.event_pk AS event_pk, e.event_type AS event_type, f.new_from AS new_from, \ + r.entry_index AS entry_index, r.field AS field, r.text_id AS text_id \ + FROM events AS e \ + CROSS JOIN search_refs AS r ON r.event_pk = e.id \ + LEFT JOIN llm_usage_event_stream_facts AS f ON f.event_pk = e.id \ + WHERE +r.text_id IN (SELECT value FROM json_each(?3)) \ + ORDER BY e.observed_at, e.id, r.entry_index, r.field <> 'name' \ + LIMIT ?4", + ) + .bind::(owner) + .bind::(session_pk) + .bind::(json_ids(text_ids)?) + .bind::(limit) + .load::(connection)? + .into_iter() + .map(|row| { + let entry_index = entry_index(row.entry_index)?; + Ok(InSessionHit { + new: is_new( + event_type(&row.event_type)?, + new_from(row.new_from)?, + entry_index, + ), + event_pk: parse_uuid(&row.event_pk, "search event id")?, + entry_index, + field: field(&row.field)?, + text_id: row.text_id, + segment: None, + }) + }) + .collect() +} + +/// Event columns of the shown matches. +fn matched_events( + connection: &mut SqliteConnection, + owner: &str, + event_pks: &[Uuid], +) -> Result, AppError> { + let event_pks = serde_json::to_string(event_pks) + .map_err(|error| AppError::internal(format!("serialize search event ids: {error}")))?; + // The ids find the events, and the unary `+` keeps the owner a filter on + // them: the owner holds every event, so looking it up reads them all. + sql_query( + "SELECT id, event_type, llm_provider, llm_model FROM llm_usage_events \ + WHERE id IN (SELECT value FROM json_each(?1)) AND +user_id = ?2", + ) + .bind::(event_pks) + .bind::(owner) + .load::(connection)? + .into_iter() + .map(|row| { + Ok(( + parse_uuid(&row.id, "search event id")?, + MatchedEvent { + event_type: row.event_type, + llm_provider: row.llm_provider, + llm_model: row.llm_model, + }, + )) + }) + .collect() +} + +#[derive(QueryableByName)] +struct MatchedEventRow { + #[diesel(sql_type = Text)] + id: String, + #[diesel(sql_type = Text)] + event_type: String, + #[diesel(sql_type = Text)] + llm_provider: String, + #[diesel(sql_type = Text)] + llm_model: String, +} + +/// Session and device columns of the page's sessions. +fn session_details( + connection: &mut SqliteConnection, + owner: &str, + session_pks: &[Uuid], +) -> Result, AppError> { + if session_pks.is_empty() { + return Ok(HashMap::new()); + } + agent_sessions::table + .inner_join(devices::table) + .filter(agent_sessions::user_id.eq(owner)) + .filter(agent_sessions::id.eq_any(session_pks.iter().map(Uuid::to_string))) + .select(( + agent_sessions::id, + agent_sessions::session_id, + agent_sessions::agent_name, + agent_sessions::agent_version, + devices::host_name, + devices::platform, + agent_sessions::started_at, + agent_sessions::ended_at, + )) + .load::(connection)? + .into_iter() + .map( + |( + session_pk, + session_id, + agent_name, + agent_version, + host_name, + platform, + started_at, + ended_at, + )| { + Ok(( + parse_uuid(&session_pk, "session id")?, + SessionDetails { + session_id, + agent_name, + agent_version, + host_name, + platform, + started_at: timestamp_from_micros(started_at)?, + ended_at: ended_at.map(timestamp_from_micros).transpose()?, + }, + )) + }, + ) + .collect() +} + +/// Columns of one session row and its device. +type SessionDetailColumns = ( + String, + String, + String, + Option, + String, + String, + i64, + Option, +); + +fn json_ids(ids: &[i64]) -> Result { + serde_json::to_string(ids) + .map_err(|error| AppError::internal(format!("serialize text ids: {error}"))) +} + +fn event_type(value: &str) -> Result { + match value { + "request" => Ok(UsageEventType::Request), + "response" => Ok(UsageEventType::Response), + _ => Err(AppError::internal(format!( + "stored event type is invalid: {value}" + ))), + } +} + +fn entry_index(value: i32) -> Result { + u32::try_from(value) + .map_err(|error| AppError::internal(format!("stored entry index is invalid: {error}"))) +} + +fn new_from(value: Option) -> Result, AppError> { + value + .map(u32::try_from) + .transpose() + .map_err(|error| AppError::internal(format!("stored new_from is invalid: {error}"))) +} + +fn field(value: &str) -> Result { + PieceField::parse(value) + .ok_or_else(|| AppError::internal(format!("stored search field is invalid: {value}"))) +} diff --git a/crates/abyss-backend/src/storage/sqlite/search/tests.rs b/crates/abyss-backend/src/storage/sqlite/search/tests.rs new file mode 100644 index 0000000..d0cb45c --- /dev/null +++ b/crates/abyss-backend/src/storage/sqlite/search/tests.rs @@ -0,0 +1,839 @@ +//! SQLite search against stored events, and against a naive grep of every +//! piece of every event. + +use std::collections::{BTreeMap, BTreeSet, HashMap}; + +use diesel::{ + Connection, ExpressionMethods, QueryDsl, RunQueryDsl, SqliteConnection, + connection::SimpleConnection, dsl::count_star, +}; +use serde_json::{Value, json}; +use uuid::Uuid; + +use crate::{ + search::{ + InSessionSearchQuery, InSessionSearchResponse, SessionSearchQuery, SessionSearchResponse, + }, + usage::{IngestEventsRequest, segments::SegmentQuery}, +}; + +use crate::storage::stream_cache::StreamCache; + +use super::{ + super::{ + migrations, + repository::{ingest_events, session_segment, session_streams}, + schema::{agent_sessions, search_refs, search_texts}, + }, + search_in_session, session_search, +}; + +const OWNER_ID: Uuid = Uuid::from_u128(1); +const OTHER_OWNER_ID: Uuid = Uuid::from_u128(2); + +/// One event; a request `x` and its response `xr` share a response id, so +/// they form one exchange. +fn event(id: &str, session: &str, second: u32, event_type: &str, entries: &Value) -> Value { + json!({ + "event_id": id, + "observed_at": format!("2026-09-01T00:00:{second:02}Z"), + "device": {"host_name": "alice-mbp", "platform": "macos"}, + "agent": {"name": "codex_exec", "version": "0.9"}, + "session_id": session, + "response_id": format!("resp-{}", id.strip_suffix('r').unwrap_or(id)), + "llm": {"provider": "openai", "model": "gpt-5"}, + "event_type": event_type, + "token_usage": {"input_tokens": 1_i64, "output_tokens": 1_i64}, + "conversation": {"entries": entries, "completeness": "complete"}, + "metadata": { + "request_headers": [{"name": "x-api-key", "value": "test-secret"}] + } + }) +} + +fn database() -> SqliteConnection { + let mut database = SqliteConnection::establish(":memory:").expect("SQLite should open"); + migrations::run(&mut database).expect("migration should succeed"); + database + .batch_execute(&format!( + "INSERT INTO app_users (id, email, created_at) \ + VALUES ('{OTHER_OWNER_ID}', 'other@abyss.local', 0)" + )) + .expect("a second owner is created"); + database +} + +fn ingest(database: &mut SqliteConnection, owner: Uuid, events: &[Value]) { + let request: IngestEventsRequest = + serde_json::from_value(json!({ "events": events })).expect("request should parse"); + ingest_events(database, &request, owner, 1_000).expect("ingest should succeed"); +} + +fn search(database: &mut SqliteConnection, query: Value) -> SessionSearchResponse { + session_search( + database, + OWNER_ID, + serde_json::from_value::(query).expect("query should deserialize"), + ) + .expect("search should succeed") +} + +/// The key of the owner's session with the declared id. +fn session_pk(database: &mut SqliteConnection, owner: Uuid, session: &str) -> Uuid { + let key = agent_sessions::table + .filter(agent_sessions::user_id.eq(owner.to_string())) + .filter(agent_sessions::session_id.eq(session)) + .select(agent_sessions::id) + .first::(database) + .expect("the session is stored"); + Uuid::parse_str(&key).expect("a session key is a UUID") +} + +fn search_session( + database: &mut SqliteConnection, + session: &str, + q: &str, +) -> InSessionSearchResponse { + let session_pk = session_pk(database, OWNER_ID, session); + search_in_session( + database, + OWNER_ID, + session_pk, + &InSessionSearchQuery { q: q.to_owned() }, + &StreamCache::default(), + ) + .expect("search should succeed") +} + +/// The head of the segment holding each event of a session, read through +/// its streams and segments as the session page reads them. +fn segment_heads(database: &mut SqliteConnection, session: &str) -> HashMap { + let session_pk = session_pk(database, OWNER_ID, session); + let streams = session_streams(database, OWNER_ID, session_pk, &StreamCache::default()) + .expect("the session has streams"); + let query = SegmentQuery::default().validate().expect("a valid query"); + let mut heads = HashMap::new(); + for segment in streams.segments { + let run = session_segment( + database, + OWNER_ID, + session_pk, + &segment.head, + &query, + &StreamCache::default(), + ) + .expect("the segment exists"); + for exchange in run.exchanges { + for pk in [exchange.exchange.request_pk, exchange.exchange.response_pk] + .into_iter() + .flatten() + { + heads.insert(pk, segment.head.clone()); + } + } + } + heads +} + +fn user(text: &str) -> Value { + json!({"role": "user", "text": text}) +} + +fn assistant(text: &str) -> Value { + json!({"role": "assistant", "text": text}) +} + +fn call(id: &str, name: &str, text: &str) -> Value { + json!({"role": "tool_call", "text": text, "call_id": id, "name": name}) +} + +fn result(id: &str, text: &str) -> Value { + json!({"role": "tool_result", "text": text, "call_id": id}) +} + +/// Requests that replay their history, responses, and another owner's copy +/// of the same text. +fn fixture() -> Vec { + let system = json!({"role": "system", "text": "You are Abyss. 先搜索，再回答。Use NOTES.md."}); + let ask = user("find nimbus.example.org in NOTES.md"); + let grep = call("c1", "Grep", "{\"pattern\":\"nimbus\"}"); + let found = result("c1", "NOTES.md: nimbus.example.org 😀 NIMBUS.EXAMPLE.ORG"); + let answer = assistant("开始新的任务 for the branch; run cargo test."); + let mut events = vec![ + event("a1", "alpha", 1, "request", &json!([system, ask])), + event("a1r", "alpha", 2, "response", &json!([grep])), + event( + "a2", + "alpha", + 3, + "request", + &json!([system, ask, grep, found]), + ), + event("a2r", "alpha", 4, "response", &json!([answer])), + event( + "a3", + "alpha", + 5, + "request", + &json!([ + system, + ask, + grep, + found, + answer, + user("cargo test 搜索 again") + ]), + ), + event("a3r", "alpha", 6, "response", &json!([assistant("ok")])), + event( + "b1", + "beta", + 7, + "request", + &json!([ + json!({"role": "developer", "text": "ÄBC ΣΟΦΙΑ Straße İstanbul ＡＢＣ file"}), + user("ababa aaaaa \u{212a}elvin a\0b\0c abc"), + ]), + ), + event( + "b1r", + "beta", + 7, + "response", + &json!([call("c2", "Read", "NOTES.md")]), + ), + event( + "b2", + "beta", + 8, + "request", + &json!([ + json!({"role": "developer", "text": "ÄBC ΣΟΦΙΑ Straße İstanbul ＡＢＣ file"}), + user("ababa aaaaa \u{212a}elvin a\0b\0c abc"), + call("c2", "Read", "NOTES.md"), + result("c2", "# Notes\nbranches/search 开始新的任务"), + json!({"role": "reasoning", "text": ""}), + ]), + ), + event( + "g1", + "gamma", + 9, + "response", + &json!([assistant("搜索 cargo test nimbus")]), + ), + ]; + events.push(event("e1", "empty", 11, "request", &json!([]))); + events +} + +fn seeded() -> SqliteConnection { + let mut database = database(); + ingest(&mut database, OWNER_ID, &fixture()); + let other = fixture() + .into_iter() + .map(|mut event| { + event["event_id"] = json!(format!( + "other-{}", + event["event_id"].as_str().unwrap_or_default() + )); + event["session_id"] = json!("other-session"); + event + }) + .collect::>(); + ingest(&mut database, OTHER_OWNER_ID, &other); + database +} + +fn count(value: usize) -> u64 { + u64::try_from(value).expect("a count fits") +} + +#[derive(Clone, Copy)] +enum Table { + Texts, + Refs, +} + +fn stored(database: &mut SqliteConnection, table: Table) -> u64 { + let rows = match table { + Table::Texts => search_texts::table + .select(count_star()) + .first::(database), + Table::Refs => search_refs::table + .select(count_star()) + .first::(database), + }; + u64::try_from(rows.expect("rows count")).expect("a count fits") +} + +/// Every stored event of the owner as the naive grep reads it. +struct StoredEvent { + pk: Uuid, + event_id: String, + session_id: String, + observed_at: i64, + response: bool, + entries: Vec, +} + +fn stored_events(database: &mut SqliteConnection, owner: Uuid) -> Vec { + #[derive(diesel::QueryableByName)] + struct Row { + #[diesel(sql_type = diesel::sql_types::Text)] + id: String, + #[diesel(sql_type = diesel::sql_types::Text)] + event_id: String, + #[diesel(sql_type = diesel::sql_types::Text)] + session_id: String, + #[diesel(sql_type = diesel::sql_types::BigInt)] + observed_at: i64, + #[diesel(sql_type = diesel::sql_types::Text)] + event_type: String, + #[diesel(sql_type = diesel::sql_types::Nullable)] + conversation: Option, + } + diesel::sql_query( + "SELECT id, event_id, session_id, observed_at, event_type, conversation \ + FROM llm_usage_events WHERE user_id = ?", + ) + .bind::(owner.to_string()) + .load::(database) + .expect("events should load") + .into_iter() + .map(|row| { + let conversation = row.conversation.map_or(Value::Null, |conversation| { + serde_json::from_str::(&conversation).expect("conversation is JSON") + }); + StoredEvent { + pk: Uuid::parse_str(&row.id).expect("event id is a UUID"), + event_id: row.event_id, + session_id: row.session_id, + observed_at: row.observed_at, + response: row.event_type == "response", + entries: conversation["entries"] + .as_array() + .cloned() + .unwrap_or_default(), + } + }) + .collect() +} + +/// The reference fold, written without the shared implementation. +fn naive_fold(text: &str) -> Vec { + text.chars() + .map(|character| { + let lower = character.to_lowercase().collect::>(); + match lower.as_slice() { + [folded] => *folded, + _ => character, + } + }) + .collect() +} + +/// Non-overlapping occurrences as UTF-16 ranges, by sliding over characters. +fn naive_ranges(text: &str, query: &str) -> Vec<[usize; 2]> { + let haystack = naive_fold(text); + let needle = naive_fold(query.trim()); + let originals = text.chars().collect::>(); + let utf16 = |end: usize| { + originals[..end] + .iter() + .map(|character| character.len_utf16()) + .sum::() + }; + let mut ranges = Vec::new(); + let mut start = 0_usize; + while let Some(window) = haystack.get(start..start.saturating_add(needle.len())) { + if window.len() < needle.len() { + break; + } + if window == needle.as_slice() { + let end = start.saturating_add(needle.len()); + ranges.push([utf16(start), utf16(end)]); + start = end; + } else { + start = start.saturating_add(1); + } + } + ranges +} + +/// A request's new part starts after its last model-side entry. +fn naive_new_from(entries: &[Value]) -> usize { + entries + .iter() + .rposition(|entry| { + matches!( + entry["role"].as_str(), + Some("assistant" | "reasoning" | "tool_call") + ) + }) + .map_or(0, |last| last.saturating_add(1)) +} + +/// (event, entry, field) → (new, ranges) for every piece holding the query. +type NaiveHits = BTreeMap<(Uuid, usize, &'static str), (bool, Vec<[usize; 2]>)>; + +fn naive_hits(events: &[StoredEvent], query: &str, session: Option<&str>) -> NaiveHits { + let mut hits = BTreeMap::new(); + for event in events { + if session.is_some_and(|session| session != event.session_id) { + continue; + } + let new_from = naive_new_from(&event.entries); + for (index, entry) in event.entries.iter().enumerate() { + for field in ["name", "text"] { + let Some(text) = entry[field].as_str().filter(|text| !text.is_empty()) else { + continue; + }; + let ranges = naive_ranges(text, query); + if !ranges.is_empty() { + hits.insert( + (event.pk, index, field), + (event.response || index >= new_from, ranges), + ); + } + } + } + } + hits +} + +const QUERIES: [&str; 22] = [ + "搜索", + "nimbus.example.org", + "NIMBUS", + "branch", + "开始新的任务", + "cargo test", + "NOTES.md", + "a", + "s", + "ab", + "aa", + "😀", + "äbc", + "σοφ", + "grep", + "\"pattern\"", + "ss", + "ß", + "i̇", + "İ", + "kelvin", + "bc", +]; + +#[test] +fn session_search_returns_the_naive_greps_hit_events() { + let mut database = seeded(); + let events = stored_events(&mut database, OWNER_ID); + let times = events + .iter() + .map(|event| (event.pk, event.observed_at)) + .collect::>(); + for query in QUERIES.into_iter().chain(["nothing-matches-this", "a\0b"]) { + let naive = naive_hits(&events, query, None); + // Per session: hit events, and those with a new hit. + let mut expected = BTreeMap::, BTreeSet)>::new(); + for ((pk, _, _), (new, _)) in &naive { + let session = events + .iter() + .find(|event| event.pk == *pk) + .map(|event| event.session_id.clone()) + .expect("hits come from stored events"); + let entry = expected.entry(session).or_default(); + entry.0.insert(*pk); + if *new { + entry.1.insert(*pk); + } + } + let mut order = expected + .iter() + .map(|(session, (hit_events, _))| { + let latest = hit_events + .iter() + .map(|pk| times[pk]) + .max() + .expect("a session has hits"); + (std::cmp::Reverse(latest), session.clone()) + }) + .collect::>(); + order.sort(); + + let mut found = Vec::new(); + for page in 1_u32.. { + let response = search( + &mut database, + json!({"q": query, "page": page, "page_size": 2_u32}), + ); + assert_eq!(response.total_sessions, count(expected.len()), "{query:?}"); + let done = !response.has_more; + found.extend(response.items); + if done { + break; + } + } + assert_eq!( + found + .iter() + .map(|item| item.session_id.clone()) + .collect::>(), + order + .iter() + .map(|(_, session)| session.clone()) + .collect::>(), + "sessions order by their latest hit event for {query:?}" + ); + for item in &found { + let (hit_events, new_events) = &expected[&item.session_id]; + assert_eq!(item.match_count, count(hit_events.len()), "{query:?}"); + assert_eq!(item.new_match_count, count(new_events.len()), "{query:?}"); + let mut by_time = hit_events.iter().copied().collect::>(); + by_time.sort_by_key(|pk| (times[pk], *pk)); + let shown = by_time + .iter() + .filter(|pk| new_events.contains(*pk)) + .chain(by_time.iter().filter(|pk| !new_events.contains(*pk))) + .take(3) + .map(|pk| (*pk, new_events.contains(pk))) + .collect::>(); + assert_eq!( + item.matches + .iter() + .map(|matched| (matched.event_pk, matched.new)) + .collect::>(), + shown, + "shown matches for {query:?}" + ); + for matched in &item.matches { + assert!( + !matched.fragments.is_empty() && matched.fragments.len() <= 2, + "{query:?}" + ); + for fragment in &matched.fragments { + let highlighted = fragment + .segments + .iter() + .filter(|segment| segment.highlighted) + .map(|segment| naive_fold(&segment.text)) + .collect::>(); + assert!(!highlighted.is_empty(), "{query:?}"); + assert!( + highlighted.iter().all(|text| *text == naive_fold(query)), + "fragments highlight occurrences of {query:?}: {fragment:?}" + ); + } + } + } + } +} + +#[test] +fn search_in_session_returns_the_naive_greps_hits_and_ranges() { + let mut database = seeded(); + let events = stored_events(&mut database, OWNER_ID); + let times = events + .iter() + .map(|event| (event.pk, event.observed_at)) + .collect::>(); + for session in ["alpha", "beta", "gamma", "empty"] { + let heads = segment_heads(&mut database, session); + for query in QUERIES.into_iter().chain(["nothing-matches-this"]) { + let naive = naive_hits(&events, query, Some(session)); + let response = search_session(&mut database, session, query); + assert!(!response.truncated); + for hit in &response.hits { + assert_eq!( + hit.segment.as_ref(), + heads.get(&hit.event_pk), + "a hit names the segment the session page puts its event in" + ); + } + let ranges = response + .texts + .iter() + .map(|text| (text.id, text.ranges.clone())) + .collect::>(); + let got = response + .hits + .iter() + .map(|hit| { + ( + ( + hit.event_pk, + usize::try_from(hit.entry_index).expect("an index fits"), + hit.field.as_str(), + ), + (hit.new, ranges[&hit.text_id].clone()), + ) + }) + .collect::(); + assert_eq!(got, naive, "{session} {query:?}"); + assert_eq!(got.len(), response.hits.len(), "a piece is one hit"); + assert_eq!( + ranges.len(), + response + .hits + .iter() + .map(|hit| hit.text_id) + .collect::>() + .len(), + "texts are those the hits name" + ); + let mut ordered = response.hits.clone(); + ordered.sort_by_key(|hit| { + ( + times[&hit.event_pk], + hit.event_pk, + hit.entry_index, + hit.field, + ) + }); + assert_eq!( + ordered, response.hits, + "hits are in session order for {query:?}" + ); + } + } +} + +#[test] +fn ingest_stores_each_distinct_text_once_per_owner() { + let mut database = database(); + ingest(&mut database, OWNER_ID, &fixture()); + let texts = stored(&mut database, Table::Texts); + let refs = stored(&mut database, Table::Refs); + let events = stored_events(&mut database, OWNER_ID); + let pieces = events + .iter() + .flat_map(|event| event.entries.iter()) + .flat_map(|entry| { + ["name", "text"].map(|field| { + entry[field] + .as_str() + .filter(|text| !text.is_empty()) + .map(str::to_owned) + }) + }) + .flatten() + .collect::>(); + assert_eq!(count(pieces.len()), refs, "one reference per piece"); + assert_eq!( + count(pieces.iter().collect::>().len()), + texts, + "one text per distinct piece" + ); + + let other = fixture() + .into_iter() + .map(|mut event| { + event["event_id"] = json!(format!( + "other-{}", + event["event_id"].as_str().unwrap_or_default() + )); + event + }) + .collect::>(); + ingest(&mut database, OTHER_OWNER_ID, &other); + assert_eq!( + stored(&mut database, Table::Texts), + texts.saturating_mul(2), + "texts are not shared across owners" + ); + // A duplicate event adds nothing. + ingest(&mut database, OWNER_ID, &fixture()); + assert_eq!(stored(&mut database, Table::Refs), refs.saturating_mul(2)); +} + +#[test] +fn searches_are_scoped_to_the_owner() { + let mut database = seeded(); + let found = search(&mut database, json!({"q": "nimbus"})); + assert_eq!( + found + .items + .iter() + .map(|item| item.session_id.as_str()) + .collect::>(), + ["gamma", "alpha"], + "the other owner's session is not found" + ); + let other_session = session_pk(&mut database, OTHER_OWNER_ID, "other-session"); + let other = search_in_session( + &mut database, + OWNER_ID, + other_session, + &InSessionSearchQuery { + q: "nimbus".to_owned(), + }, + &StreamCache::default(), + ); + assert!( + matches!(other, Err(crate::error::AppError::NotFound(_))), + "another owner's session is not found" + ); + let unknown = search_in_session( + &mut database, + OWNER_ID, + Uuid::from_u128(99), + &InSessionSearchQuery { + q: "nimbus".to_owned(), + }, + &StreamCache::default(), + ); + assert!(matches!(unknown, Err(crate::error::AppError::NotFound(_)))); + let empty = search_session(&mut database, "empty", "nimbus"); + assert!(empty.hits.is_empty() && empty.texts.is_empty() && !empty.truncated); +} + +#[test] +fn hits_mark_new_and_inherited_pieces() { + let mut database = seeded(); + let response = search_session(&mut database, "alpha", "NOTES.md"); + let events = stored_events(&mut database, OWNER_ID); + // a2 replays [system, ask, grep, found]; only `found`, after the call, is new. + let a2 = events + .iter() + .find(|event| event.event_id == "a2") + .map(|event| event.pk) + .expect("a2 is stored"); + let a2_hits = response + .hits + .iter() + .filter(|hit| hit.event_pk == a2) + .map(|hit| (hit.entry_index, hit.new)) + .collect::>(); + assert_eq!(a2_hits, [(0, false), (1, false), (3, true)]); + assert!( + response.hits.iter().all(|hit| hit.segment.is_some()), + "every alpha event is in an exchange" + ); + + let found = search(&mut database, json!({"q": "NOTES.md"})); + let alpha = found + .items + .iter() + .find(|item| item.session_id == "alpha") + .expect("alpha matches"); + // a1 (all new), a2 (found is new), a3 (all inherited) hit; responses do not. + assert_eq!((alpha.match_count, alpha.new_match_count), (3, 2)); + assert_eq!( + alpha + .matches + .iter() + .map(|matched| matched.new) + .collect::>(), + [true, true, false] + ); +} + +#[test] +fn a_tool_name_is_its_own_piece_and_orders_before_the_text() { + let mut database = seeded(); + let response = search_session(&mut database, "beta", "read"); + let fields = response + .hits + .iter() + .map(|hit| (hit.entry_index, hit.field.as_str())) + .collect::>(); + // b1r's call is entry 0; b2 replays it as entry 2. + assert_eq!(fields, [(0, "name"), (2, "name")]); + let across = search_session(&mut database, "alpha", "grep{"); + assert!(across.hits.is_empty(), "a name and a text are never joined"); +} + +#[test] +fn filters_narrow_the_matching_events() { + let mut database = seeded(); + let sessions = |database: &mut SqliteConnection, query: Value| { + search(database, query) + .items + .into_iter() + .map(|item| (item.session_id, item.match_count)) + .collect::>() + }; + assert_eq!( + sessions(&mut database, json!({"q": "搜索"})), + [("gamma".to_owned(), 1), ("alpha".to_owned(), 3)] + ); + assert_eq!( + sessions( + &mut database, + json!({"q": "搜索", "event_type": "response"}) + ), + [("gamma".to_owned(), 1)] + ); + assert_eq!( + sessions( + &mut database, + json!({"q": "搜索", "from": "2026-09-01T00:00:03Z", "to": "2026-09-01T00:00:09Z"}) + ), + [("alpha".to_owned(), 2)] + ); + assert_eq!( + sessions( + &mut database, + json!({"q": "搜索", "agent_name": ["claude-cli"]}) + ), + [] + ); + assert_eq!( + sessions(&mut database, json!({"q": "搜索", "agent_name": ["claude-cli", "codex_exec"], "llm_provider": ["openai"], "llm_model": ["gpt-5"]})).len(), + 2 + ); + assert_eq!( + sessions(&mut database, json!({"q": "搜索", "llm_model": ["absent"]})), + [] + ); +} + +#[test] +fn a_page_past_the_end_still_counts_the_sessions() { + let mut database = seeded(); + let past_end = search( + &mut database, + json!({"q": "a", "page": 9_u32, "page_size": 1_u32}), + ); + assert_eq!(past_end.total_sessions, 3); + assert!(past_end.items.is_empty() && !past_end.has_more); +} + +#[test] +fn a_nul_character_does_not_hide_later_trigrams() { + let mut database = seeded(); + // b1 holds "a\0b\0c abc": the trigram index must still find "c abc". + let hits = search_session(&mut database, "beta", "c abc").hits; + assert_eq!(hits.len(), 2, "b1 and its replay in b2"); + let nul = search_session(&mut database, "beta", "a\0b\0c"); + assert_eq!(nul.hits.len(), 2, "a query with NUL matches literally"); +} + +#[test] +fn fragments_highlight_every_occurrence_in_their_window() { + let mut database = seeded(); + let found = search(&mut database, json!({"q": "nimbus.example.org"})); + let alpha = found + .items + .iter() + .find(|item| item.session_id == "alpha") + .expect("alpha matches"); + let rendered = alpha.matches[1].fragments[0] + .segments + .iter() + .map(|segment| { + if segment.highlighted { + format!("[{}]", segment.text) + } else { + segment.text.clone() + } + }) + .collect::(); + assert_eq!( + rendered, + "NOTES.md: [nimbus.example.org] 😀 [NIMBUS.EXAMPLE.ORG]" + ); +} diff --git a/crates/abyss-backend/src/storage/sqlite/session_list.rs b/crates/abyss-backend/src/storage/sqlite/session_list.rs index 74390a9..138166d 100644 --- a/crates/abyss-backend/src/storage/sqlite/session_list.rs +++ b/crates/abyss-backend/src/storage/sqlite/session_list.rs @@ -2,13 +2,16 @@ use diesel::{ QueryableByName, RunQueryDsl, SqliteConnection, sql_query, - sql_types::{BigInt, Bool, Nullable, Text}, + sql_types::{BigInt, Bool, Integer, Nullable, Text}, }; use uuid::Uuid; use crate::{ error::AppError, - usage::session_list::{SessionListItem, SessionListResponse, ValidatedSessionListQuery}, + usage::{ + session_list::{SessionListItem, SessionListResponse, ValidatedSessionListQuery}, + streams::EntryLabel, + }, }; use super::{ @@ -34,14 +37,29 @@ const ACTIVITY_CTE: &str = "\ s.started_at AS started_at, \ s.ended_at AS ended_at, \ COALESCE(MAX(e.observed_at), s.started_at) AS last_activity_at, \ - (SELECT COUNT(*) FROM agent_turns t WHERE t.session_pk = s.id) AS turns, \ + (SELECT CASE WHEN COUNT(tf.turn_id) > 0 THEN COUNT(DISTINCT tf.turn_id) \ + ELSE COUNT(*) FILTER (WHERE tf.starts_turn = 1 AND tf.continues_context = 1) \ + + COALESCE((SELECT 1 - ff.continues_context \ + FROM llm_usage_events fe \ + JOIN llm_usage_event_stream_facts ff ON ff.event_pk = fe.id \ + WHERE fe.session_pk = s.id AND ff.starts_turn = 1 \ + ORDER BY fe.observed_at, fe.id LIMIT 1), 0) END \ + FROM llm_usage_events te \ + JOIN llm_usage_event_stream_facts tf ON tf.event_pk = te.id \ + WHERE te.session_pk = s.id) AS turns, \ COUNT(e.session_pk) AS events, \ COALESCE(SUM(e.total_tokens), 0) AS total_tokens, \ - st.title AS title, \ + title_event.id AS title_event_pk, \ + title_facts.label_index AS title_index, \ + title_facts.label_text AS title_text, \ + title_facts.label_length AS title_length, \ st.purpose AS purpose \ FROM agent_sessions s \ JOIN devices d ON d.id = s.device_context_id \ - LEFT JOIN session_titles st ON st.session_pk = s.id \ + LEFT JOIN session_titles st ON st.session_pk = s.id AND st.user_id = s.user_id \ + LEFT JOIN llm_usage_events title_event ON title_event.id = st.title_event_pk \ + AND title_event.session_pk = s.id AND title_event.event_type = 'request' \ + LEFT JOIN llm_usage_event_stream_facts title_facts ON title_facts.event_pk = title_event.id \ LEFT JOIN llm_usage_events e ON e.session_pk = s.id \ WHERE s.user_id = ?1 \ AND (?2 = 0 OR s.agent_name IN (SELECT value FROM json_each(?3))) \ @@ -88,7 +106,13 @@ struct SessionListRow { #[diesel(sql_type = BigInt)] total_tokens: i64, #[diesel(sql_type = Nullable)] - title: Option, + title_event_pk: Option, + #[diesel(sql_type = Nullable)] + title_index: Option, + #[diesel(sql_type = Nullable)] + title_text: Option, + #[diesel(sql_type = Nullable)] + title_length: Option, #[diesel(sql_type = Nullable)] purpose: Option, } @@ -163,7 +187,15 @@ pub(super) fn list_sessions( turns: row.turns, events: row.events, total_tokens: row.total_tokens, - title: row.title, + title: EntryLabel::from_columns( + row.title_event_pk + .as_deref() + .map(|pk| parse_uuid(pk, "title event id")) + .transpose()?, + row.title_index, + row.title_text, + row.title_length, + ), purpose: row.purpose, }) }) @@ -213,12 +245,17 @@ mod tests { "device": {"host_name": "alice-mbp", "platform": "macos"}, "agent": {"name": agent}, "session_id": session, - "turn_id": format!("{session}-{turn}"), "llm": {"provider": "OpenAI", "model": model}, - "event_type": "response", - "text": "reply", + "event_type": "request", "token_usage": {"input_tokens": 0_i64, "output_tokens": 10_i64, "cache_read_tokens": 0_i64, "cache_write_tokens": 0_i64, "reasoning_tokens": 0_i64, "total_tokens": 10_i64}, - "metadata": {} + // Each event asks in words after an answer, so each opens a turn. + "conversation": { + "entries": [ + {"role": "assistant", "text": "earlier"}, + {"role": "user", "text": format!("ask {turn}")} + ], + "completeness": "complete" + } }) } diff --git a/crates/abyss-backend/src/storage/sqlite/session_titles.rs b/crates/abyss-backend/src/storage/sqlite/session_titles.rs index 86ec7e8..500eed7 100644 --- a/crates/abyss-backend/src/storage/sqlite/session_titles.rs +++ b/crates/abyss-backend/src/storage/sqlite/session_titles.rs @@ -1,168 +1,165 @@ -//! SQLite persistence of session names reported by the collector. +//! SQLite registration of each session's earliest request and its stored label. + +use std::collections::HashMap; use chrono::Utc; -use diesel::{ExpressionMethods, OptionalExtension, QueryDsl, RunQueryDsl, SqliteConnection}; -use serde_json::Value; +use diesel::{ + ExpressionMethods, OptionalExtension, QueryDsl, QueryableByName, RunQueryDsl, SqliteConnection, + sql_query, + sql_types::{Integer, Nullable, Text}, +}; use uuid::Uuid; -use super::{ - models::parse_uuid, - schema::{llm_usage_events, session_titles}, -}; -use crate::{ - error::AppError, - usage::{ - SessionInfo, - session_title::{TITLE_GENERATION_PURPOSE, TitleContribution, default_title}, - }, -}; +use super::{models::parse_uuid, schema::session_titles}; +use crate::{error::AppError, usage::streams::EntryLabel}; -/// Stored title, purpose, and matched task session of one session. -type StoredTitle = (Option, Option, Option); +/// First-request label, purpose, and matched task session. +type StoredTitle = (Option, Option, Option); -/// One stored event considered for title capture. +/// One stored event considered for the first-request registration. pub(super) struct CapturedEvent<'a> { pub(super) user_id: Uuid, pub(super) session_pk: Uuid, pub(super) event_pk: Uuid, pub(super) event_type: &'a str, - pub(super) metadata: &'a Value, - pub(super) naming_thread: bool, pub(super) observed_at: i64, } -/// Records what a newly stored event says about its session's name. -/// -/// An Agent title replaces a default title and an older Agent title. A default -/// title keeps the earliest user input, even after late delivery, and never -/// replaces an Agent title. +/// Keeps the earliest request even when it has no label, including late delivery. +/// Returns whether the registered request changed. pub(super) fn capture( connection: &mut SqliteConnection, event: &CapturedEvent<'_>, -) -> Result<(), AppError> { +) -> Result { + if event.event_type != "request" { + return Ok(false); + } let session = event.session_pk.to_string(); - let agent_title = - match TitleContribution::of(event.event_type, event.metadata, event.naming_thread) { - None => return Ok(()), - Some(TitleContribution::TitleThread) => { - return mark_title_thread(connection, &session, event.user_id); - } - Some(TitleContribution::Agent(title)) => Some(title), - Some(TitleContribution::UserInput) => None, - }; + let event_pk = event.event_pk.to_string(); let stored = session_titles::table .find(&session) .select(( - session_titles::title_event_pk, session_titles::title_observed_at, + session_titles::title_event_pk, )) - .first::<(Option, Option)>(connection) + .first::<(Option, Option)>(connection) .optional()? - .and_then(|(event_pk, observed_at)| event_pk.zip(observed_at)); - let agent = agent_title.is_some(); - let replaces = match stored { - None => true, - Some((_, stored_at)) if agent && stored_at < event.observed_at => true, - Some((_, stored_at)) if !agent && stored_at <= event.observed_at => false, - // Otherwise only a stored default title gives way. - Some((event_pk, _)) => !is_agent_title(connection, &event_pk)?, - }; - if !replaces { - return Ok(()); + .and_then(|(at, pk)| at.zip(pk)); + if stored.is_some_and(|(at, pk)| (at, pk) <= (event.observed_at, event_pk.clone())) { + return Ok(false); } - // Read the user input only once it could name the session. - let Some(title) = agent_title - .map(str::to_owned) - .or_else(|| default_title(event.metadata)) - else { - return Ok(()); - }; let now = Utc::now().timestamp_micros(); diesel::insert_into(session_titles::table) .values(( session_titles::session_pk.eq(&session), session_titles::user_id.eq(event.user_id.to_string()), - session_titles::title.eq(&title), - session_titles::title_event_pk.eq(event.event_pk.to_string()), + session_titles::title.eq(None::), + session_titles::title_event_pk.eq(&event_pk), session_titles::title_observed_at.eq(event.observed_at), session_titles::updated_at.eq(now), )) .on_conflict(session_titles::session_pk) .do_update() .set(( - session_titles::title.eq(&title), - session_titles::title_event_pk.eq(event.event_pk.to_string()), + session_titles::title.eq(None::), + session_titles::title_event_pk.eq(&event_pk), session_titles::title_observed_at.eq(event.observed_at), session_titles::updated_at.eq(now), )) .execute(connection)?; - Ok(()) + Ok(true) } -/// Agent titles are read from responses and default titles from requests, so -/// the side of the event that gave the stored title tells them apart. -fn is_agent_title(connection: &mut SqliteConnection, event_pk: &str) -> Result { - Ok(llm_usage_events::table - .find(event_pk) - .select(llm_usage_events::event_type) - .first::(connection) - .optional()? - .is_some_and(|event_type| event_type == "response")) +const TITLE_SELECT: &str = "\ + SELECT s.id AS session_pk, e.id AS event_pk, \ + f.label_index, f.label_text, f.label_length, st.purpose, st.parent_session_pk \ + FROM session_titles st \ + JOIN agent_sessions s ON s.id = st.session_pk AND s.user_id = st.user_id \ + LEFT JOIN llm_usage_events e ON e.id = st.title_event_pk \ + AND e.session_pk = s.id AND e.event_type = 'request' \ + LEFT JOIN llm_usage_event_stream_facts f ON f.event_pk = e.id"; + +#[derive(QueryableByName)] +struct TitleRow { + #[diesel(sql_type = Text)] + session_pk: String, + #[diesel(sql_type = Nullable)] + event_pk: Option, + #[diesel(sql_type = Nullable)] + label_index: Option, + #[diesel(sql_type = Nullable)] + label_text: Option, + #[diesel(sql_type = Nullable)] + label_length: Option, + #[diesel(sql_type = Nullable)] + purpose: Option, + #[diesel(sql_type = Nullable)] + parent_session_pk: Option, } -/// Marks a session as a thread that names another session. -fn mark_title_thread( +impl TitleRow { + fn into_title(self) -> Result { + let event_pk = self + .event_pk + .as_deref() + .map(|pk| parse_uuid(pk, "title event id")) + .transpose()?; + let label = EntryLabel::from_columns( + event_pk, + self.label_index, + self.label_text, + self.label_length, + ); + let parent = self + .parent_session_pk + .as_deref() + .map(|pk| parse_uuid(pk, "parent session id")) + .transpose()?; + Ok((label, self.purpose, parent)) + } +} + +/// Labels of an authorized result page, fetched in one query. +pub(super) fn titles( connection: &mut SqliteConnection, - session: &str, - user_id: Uuid, -) -> Result<(), AppError> { - let now = Utc::now().timestamp_micros(); - diesel::insert_into(session_titles::table) - .values(( - session_titles::session_pk.eq(session), - session_titles::user_id.eq(user_id.to_string()), - session_titles::purpose.eq(TITLE_GENERATION_PURPOSE), - session_titles::updated_at.eq(now), - )) - .on_conflict(session_titles::session_pk) - .do_update() - .set(( - session_titles::purpose.eq(TITLE_GENERATION_PURPOSE), - session_titles::updated_at.eq(now), - )) - .execute(connection)?; - Ok(()) + owner: Uuid, + sessions: &[Uuid], +) -> Result, AppError> { + if sessions.is_empty() { + return Ok(HashMap::new()); + } + let sessions = serde_json::to_string(sessions) + .map_err(|error| AppError::internal(format!("serialize title session ids: {error}")))?; + let rows = sql_query(format!( + "{TITLE_SELECT} WHERE s.user_id = ? AND s.id IN (SELECT value FROM json_each(?))" + )) + .bind::(owner.to_string()) + .bind::(sessions) + .load::(connection)?; + let mut titles = HashMap::new(); + for row in rows { + let session = parse_uuid(&row.session_pk, "session id")?; + if let Some(label) = row.into_title()?.0 { + titles.insert(session, label); + } + } + Ok(titles) } -/// Adds the stored title, purpose, and matched task session to a session. -pub(super) fn fill( +/// The registered request's label, shared with Turn labels. +pub(super) fn title( connection: &mut SqliteConnection, - session: &mut SessionInfo, -) -> Result<(), AppError> { - (session.title, session.purpose, session.parent_session_pk) = - load(connection, session.session_pk)?; - Ok(()) + session_pk: Uuid, +) -> Result, AppError> { + Ok(load(connection, session_pk)?.0) } -/// Returns the stored title and purpose of one session, and for a naming -/// thread the task session it was matched to. fn load(connection: &mut SqliteConnection, session_pk: Uuid) -> Result { - let Some((title, purpose, parent)) = session_titles::table - .find(session_pk.to_string()) - .select(( - session_titles::title, - session_titles::purpose, - session_titles::parent_session_pk, - )) - .first::<(Option, Option, Option)>(connection) + sql_query(format!("{TITLE_SELECT} WHERE s.id = ?")) + .bind::(session_pk.to_string()) + .get_result::(connection) .optional()? - else { - return Ok((None, None, None)); - }; - let parent = parent - .map(|parent| parse_uuid(&parent, "parent session id")) - .transpose()?; - Ok((title, purpose, parent)) + .map_or(Ok((None, None, None)), TitleRow::into_title) } #[cfg(test)] mod tests { @@ -184,35 +181,18 @@ mod tests { database } - /// A request whose collector conversation holds one user message. + /// A request whose collector conversation holds one user entry. fn request(agent: &str, session: &str, id: &str, at: &str, input: &str) -> Value { - event( - agent, - session, - id, - at, - "request", - json!({"conversation": { - "messages": [{"role": "user", "text": input}], - "completeness": "complete" - }}), - ) - } - - /// The response of a naming call, carrying the collector's title. - fn titled_response(session: &str, id: &str, at: &str, title: &str) -> Value { - event( - "claude-code", - session, - id, - at, - "response", - json!({"purpose": "session_title", "session_title": title}), - ) + let mut request = event(agent, session, id, at, "request"); + request["conversation"] = json!({ + "entries": [{"role": "user", "text": input}], + "completeness": "complete" + }); + request } - fn event(agent: &str, session: &str, id: &str, at: &str, kind: &str, metadata: Value) -> Value { - let mut event = json!({ + fn event(agent: &str, session: &str, id: &str, at: &str, kind: &str) -> Value { + json!({ "event_id": id, "observed_at": at, "device": {"host_name": "local", "platform": "macos"}, @@ -221,9 +201,7 @@ mod tests { "llm": {"provider": "anthropic", "model": "claude-haiku"}, "event_type": kind, "token_usage": {"input_tokens": 0_i64, "output_tokens": 1_i64, "cache_read_tokens": 0_i64, "cache_write_tokens": 0_i64, "reasoning_tokens": 0_i64, "total_tokens": 1_i64} - }); - event["metadata"] = metadata; - event + }) } fn ingest(database: &mut SqliteConnection, events: &[Value]) { @@ -243,101 +221,29 @@ mod tests { .expect("session should exist"); let (title, purpose, _) = load(database, Uuid::parse_str(&pk).expect("uuid")).expect("title should load"); - (title, purpose) + (title.map(|label| label.text), purpose) } #[test] - fn names_a_session_from_the_collector_title() { + fn responses_without_user_requests_have_no_title() { let mut database = database(); - let mut naming = request( + let mut reply = event( "claude-code", "cc", - "r2", - "2026-09-30T00:00:02Z", - "fix the build", - ); - naming["metadata"]["purpose"] = json!("session_title"); - ingest( - &mut database, - &[ - request( - "claude-code", - "cc", - "r1", - "2026-09-30T00:00:00Z", - "fix the build", - ), - naming, - titled_response("cc", "s2", "2026-09-30T00:00:03Z", "Build fix"), - ], - ); - - assert_eq!( - session_title(&mut database, "cc"), - (Some("Build fix".to_owned()), None) - ); - } - - #[test] - fn agent_titles_outrank_default_titles_in_any_arrival_order() { - for title_first in [false, true] { - let mut database = database(); - let mut events = vec![ - request( - "claude-code", - "cc", - "r1", - "2026-09-30T00:00:00Z", - "fix the build", - ), - titled_response("cc", "s1", "2026-09-30T00:00:03Z", "Build fix"), - ]; - // An earlier user request delivered late does not replace the title. - if title_first { - events.reverse(); - } - for value in events { - ingest(&mut database, &[value]); - } - assert_eq!( - session_title(&mut database, "cc").0.as_deref(), - Some("Build fix"), - "title first: {title_first}" - ); - } - } - - #[test] - fn keeps_the_newest_agent_title() { - let mut database = database(); - ingest( - &mut database, - &[ - titled_response("cc", "s2", "2026-09-30T02:00:01Z", "Newer"), - titled_response("cc", "s1", "2026-09-30T01:00:01Z", "Older"), - ], - ); - assert_eq!( - session_title(&mut database, "cc").0.as_deref(), - Some("Newer") - ); - ingest( - &mut database, - &[titled_response( - "cc", - "s3", - "2026-09-30T03:00:01Z", - "Newest", - )], - ); - assert_eq!( - session_title(&mut database, "cc").0.as_deref(), - Some("Newest") + "s1", + "2026-09-30T01:00:01Z", + "response", ); + reply["conversation"] = json!({ + "entries": [{"role": "assistant", "text": "Generated title"}], + "completeness": "complete" + }); + ingest(&mut database, &[reply]); + assert_eq!(session_title(&mut database, "cc"), (None, None)); } #[test] - fn title_like_replies_without_the_collector_field_are_ordinary_output() { + fn title_like_replies_do_not_name_a_session() { let mut database = database(); let mut reply = event( "claude-code", @@ -345,12 +251,11 @@ mod tests { "s1", "2026-09-30T00:00:01Z", "response", - json!({"conversation": { - "messages": [{"role": "assistant", "text": "{\"title\":\"Quoted\"}"}], - "completeness": "complete" - }}), ); - reply["text"] = json!("{\"title\":\"Quoted\"}"); + reply["conversation"] = json!({ + "entries": [{"role": "assistant", "text": "{\"title\":\"Quoted\"}"}], + "completeness": "complete" + }); ingest( &mut database, &[ @@ -365,71 +270,47 @@ mod tests { } #[test] - fn default_titles_keep_the_earliest_user_request_even_when_delivered_late() { + fn titles_keep_the_earliest_request_even_when_delivered_late() { let mut database = database(); for (id, at, kind, input, expected) in [ - ("r2", "2026-09-30T00:00:02Z", "request", "Second", "Second"), - ("r3", "2026-09-30T00:00:03Z", "request", "Third", "Second"), - ("r1", "2026-09-30T00:00:01Z", "request", "First", "First"), + ( + "r2", + "2026-09-30T00:00:02Z", + "request", + "Second", + Some("Second"), + ), + ( + "r3", + "2026-09-30T00:00:03Z", + "request", + "Third", + Some("Second"), + ), + ( + "r1", + "2026-09-30T00:00:01Z", + "request", + "First", + Some("First"), + ), ( "s0", "2026-09-30T00:00:00Z", "response", "Response", - "First", + Some("First"), ), + ("r0", "2026-09-30T00:00:00Z", "request", "", None), + ("r4", "2026-09-30T00:00:04Z", "request", "Fourth", None), ] { let mut value = request("cursor", "default", id, at, input); value["event_type"] = json!(kind); ingest(&mut database, &[value]); assert_eq!( session_title(&mut database, "default").0.as_deref(), - Some(expected) + expected ); } } - - #[test] - fn codex_task_sessions_keep_their_first_input_and_naming_threads_are_marked() { - let mut database = database(); - let mut naming = request( - "codex", - "title-thread", - "n1", - "2026-09-30T00:00:01Z", - "Generate a concise, single-line task title.\nUser prompt:\n画一只鹈鹕", - ); - naming["naming_thread"] = json!(true); - let task = request( - "codex", - "task", - "t1", - "2026-09-30T00:00:00Z", - "…", - ); - ingest( - &mut database, - &[ - task, - naming, - event( - "codex", - "title-thread", - "n2", - "2026-09-30T00:00:02Z", - "response", - json!({}), - ), - ], - ); - - assert_eq!( - session_title(&mut database, "task").0.as_deref(), - Some("…") - ); - assert_eq!( - session_title(&mut database, "title-thread"), - (None, Some("title_generation".to_owned())) - ); - } } diff --git a/crates/abyss-backend/src/storage/stream_cache.rs b/crates/abyss-backend/src/storage/stream_cache.rs new file mode 100644 index 0000000..cd90f8c --- /dev/null +++ b/crates/abyss-backend/src/storage/stream_cache.rs @@ -0,0 +1,313 @@ +//! Shared session structures for synchronous storage reads. The caller reads +//! the revision and event facts from one database snapshot. Per-session locks +//! merge concurrent misses without holding the cache's map lock during I/O. + +use std::{ + collections::HashMap, + sync::{Arc, Mutex, MutexGuard, PoisonError}, +}; + +use uuid::Uuid; + +use crate::{ + error::AppError, + usage::streams::{SessionStructure, StreamEvent}, +}; + +const DEFAULT_CAPACITY: usize = 128; + +type SessionKey = (Uuid, Uuid, i64); +type SessionSlot = Arc>>>; + +struct Entry { + last_used: u64, + slot: SessionSlot, +} + +#[derive(Default)] +struct State { + clock: u64, + entries: HashMap, +} + +/// At most `capacity` idle session revisions are retained. Active loaders and +/// their waiters pin their slots until they finish; pruning then restores the +/// limit. Structures already returned to callers remain valid after eviction. +pub struct StreamCache { + capacity: usize, + state: Mutex, +} + +impl Default for StreamCache { + fn default() -> Self { + Self::new(DEFAULT_CAPACITY) + } +} + +impl StreamCache { + /// With zero capacity, only overlapping calls can share a loaded result. + pub fn new(capacity: usize) -> Self { + Self { + capacity, + state: Mutex::new(State::default()), + } + } + + /// Reads and derives the structure only when this exact revision is absent. + /// A failed load is retryable. Revisions have separate slots, so older + /// snapshots share their own result without replacing newer snapshots. + pub fn get_or_load( + &self, + owner: Uuid, + session_pk: Uuid, + revision: i64, + load: F, + ) -> Result, AppError> + where + F: FnOnce() -> Result, AppError>, + { + let slot = self.slot((owner, session_pk, revision)); + let result = Self::load_slot(&slot, load); + drop(slot); + self.prune(); + result + } + + fn load_slot(slot: &SessionSlot, load: F) -> Result, AppError> + where + F: FnOnce() -> Result, AppError>, + { + let mut cached = lock(slot); + if let Some(cached) = cached.as_ref() { + return Ok(Arc::clone(cached)); + } + let structure = Arc::new(SessionStructure::derive(load()?)); + *cached = Some(Arc::clone(&structure)); + drop(cached); + Ok(structure) + } + + fn slot(&self, key: SessionKey) -> SessionSlot { + let mut state = lock(&self.state); + state.clock = state.clock.saturating_add(1); + let last_used = state.clock; + let entry = state.entries.entry(key).or_insert_with(|| Entry { + last_used, + slot: Arc::new(Mutex::new(None)), + }); + entry.last_used = last_used; + let slot = Arc::clone(&entry.slot); + let evicted = Self::trim(&mut state, self.capacity); + drop(state); + // A large evicted structure is freed outside the global map lock. + drop(evicted); + slot + } + + fn prune(&self) { + let evicted = Self::trim(&mut lock(&self.state), self.capacity); + drop(evicted); + } + + fn trim(state: &mut State, capacity: usize) -> Vec { + let mut evicted = Vec::new(); + while state.entries.len() > capacity { + let candidate = state + .entries + .iter() + .filter(|(_, entry)| Arc::strong_count(&entry.slot) == 1) + .min_by_key(|(_, entry)| entry.last_used) + .map(|(key, _)| *key); + let Some(key) = candidate else { + break; + }; + if let Some(entry) = state.entries.remove(&key) { + evicted.push(entry); + } + } + evicted + } +} + +// Only fully constructed structures are published; a loader panic leaves the +// previous value valid. Recover the lock so subsequent requests can retry. +fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(PoisonError::into_inner) +} + +#[cfg(test)] +mod tests { + use std::{ + sync::{ + Arc, Barrier, + atomic::{AtomicUsize, Ordering}, + mpsc, + }, + thread, + time::Duration, + }; + + use super::*; + + const OWNER: Uuid = Uuid::from_u128(1); + const SESSION: Uuid = Uuid::from_u128(2); + + #[test] + fn the_same_revision_reuses_one_structure_and_a_new_revision_rebuilds() { + let cache = StreamCache::default(); + let first = cache + .get_or_load(OWNER, SESSION, 1, || Ok(Vec::new())) + .expect("load"); + let repeated = cache + .get_or_load(OWNER, SESSION, 1, || panic!("cache hit must not load")) + .expect("hit"); + assert!(Arc::ptr_eq(&first, &repeated)); + let updated = cache + .get_or_load(OWNER, SESSION, 2, || Ok(Vec::new())) + .expect("reload"); + assert!(!Arc::ptr_eq(&first, &updated)); + let stale = cache + .get_or_load(OWNER, SESSION, 1, || Ok(Vec::new())) + .expect("old snapshot"); + assert!(!Arc::ptr_eq(&stale, &updated)); + let current = cache + .get_or_load(OWNER, SESSION, 2, || panic!("old snapshot replaced latest")) + .expect("latest hit"); + assert!(Arc::ptr_eq(&updated, ¤t)); + } + + #[test] + fn the_same_session_is_scoped_to_its_owner() { + let cache = StreamCache::default(); + let first = cache + .get_or_load(OWNER, SESSION, 1, || Ok(Vec::new())) + .expect("load"); + let other = cache + .get_or_load(Uuid::from_u128(3), SESSION, 1, || Ok(Vec::new())) + .expect("other owner"); + assert!(!Arc::ptr_eq(&first, &other)); + } + + #[test] + fn failed_loads_can_retry_without_discarding_an_older_revision() { + let cache = StreamCache::default(); + let first = cache + .get_or_load(OWNER, SESSION, 1, || Ok(Vec::new())) + .expect("load"); + assert!( + cache + .get_or_load(OWNER, SESSION, 2, || Err(AppError::internal( + "unavailable".to_owned() + ))) + .is_err() + ); + let previous = cache + .get_or_load(OWNER, SESSION, 1, || panic!("previous result was lost")) + .expect("previous hit"); + assert!(Arc::ptr_eq(&first, &previous)); + let retried = cache + .get_or_load(OWNER, SESSION, 2, || Ok(Vec::new())) + .expect("retry"); + assert!(!Arc::ptr_eq(&first, &retried)); + } + + #[test] + fn concurrent_misses_for_one_revision_load_only_once() { + let cache = StreamCache::default(); + let loads = AtomicUsize::new(0); + let barrier = Barrier::new(8); + let structures = thread::scope(|scope| { + let threads = std::iter::repeat_with(|| { + scope.spawn(|| { + barrier.wait(); + cache + .get_or_load(OWNER, SESSION, 1, || { + loads.fetch_add(1, Ordering::SeqCst); + Ok(Vec::new()) + }) + .expect("load shared session") + }) + }) + .take(8) + .collect::>(); + threads + .into_iter() + .map(|thread| thread.join().expect("worker")) + .collect::>() + }); + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert!( + structures + .iter() + .all(|value| Arc::ptr_eq(value, &structures[0])) + ); + } + + #[test] + fn independent_sessions_load_concurrently_and_are_pruned_afterwards() { + let cache = StreamCache::new(1); + let (started_tx, started_rx) = mpsc::channel(); + let (release_first_tx, release_first_rx) = mpsc::channel(); + let (release_second_tx, release_second_rx) = mpsc::channel(); + thread::scope(|scope| { + let cache = &cache; + let first_started = started_tx.clone(); + let first = scope.spawn(move || { + cache.get_or_load(OWNER, SESSION, 1, || { + first_started.send(()).expect("started"); + release_first_rx + .recv_timeout(Duration::from_secs(5)) + .expect("release first"); + Ok(Vec::new()) + }) + }); + let second = scope.spawn(move || { + cache.get_or_load(OWNER, Uuid::from_u128(3), 1, || { + started_tx.send(()).expect("started"); + release_second_rx + .recv_timeout(Duration::from_secs(5)) + .expect("release second"); + Ok(Vec::new()) + }) + }); + for _ in 0..2_usize { + started_rx + .recv_timeout(Duration::from_secs(2)) + .expect("both loads started"); + } + release_first_tx.send(()).expect("release first"); + release_second_tx.send(()).expect("release second"); + first.join().expect("first worker").expect("first load"); + second.join().expect("second worker").expect("second load"); + }); + assert_eq!(lock(&cache.state).entries.len(), 1); + } + + #[test] + fn capacity_evicts_the_least_recently_used_session() { + let cache = StreamCache::new(2); + let a = Uuid::from_u128(10); + let b = Uuid::from_u128(11); + let c = Uuid::from_u128(12); + let first = cache + .get_or_load(OWNER, a, 1, || Ok(Vec::new())) + .expect("a"); + let oldest = cache + .get_or_load(OWNER, b, 1, || Ok(Vec::new())) + .expect("b"); + let touched = cache + .get_or_load(OWNER, a, 1, || panic!("a must be cached")) + .expect("a hit"); + assert!(Arc::ptr_eq(&first, &touched)); + cache + .get_or_load(OWNER, c, 1, || Ok(Vec::new())) + .expect("c"); + assert_eq!(lock(&cache.state).entries.len(), 2); + assert!(lock(&cache.state).entries.contains_key(&(OWNER, a, 1))); + let reloaded = cache + .get_or_load(OWNER, b, 1, || Ok(Vec::new())) + .expect("b evicted"); + assert!(!Arc::ptr_eq(&oldest, &reloaded)); + assert_eq!(lock(&cache.state).entries.len(), 2); + } +} diff --git a/crates/abyss-backend/src/usage/conversation.rs b/crates/abyss-backend/src/usage/conversation.rs index 9814625..8104521 100644 --- a/crates/abyss-backend/src/usage/conversation.rs +++ b/crates/abyss-backend/src/usage/conversation.rs @@ -5,11 +5,8 @@ use serde_json::Value; #[derive(Deserialize)] pub struct Conversation { - pub messages: Vec, - /// Whether the first message is the request's own instructions, such as - /// a system prompt or OpenAI `instructions`; absent means none. - #[serde(default)] - pub instructions_present: bool, + /// One entry per protocol block or item, in provider order. + pub entries: Vec, completeness: Completeness, } @@ -21,21 +18,25 @@ enum Completeness { Unsupported, } -/// One message as the collector reports it; its role is kept as reported. +/// One entry as the collector reports it; its role is kept as reported. #[derive(Clone, Deserialize)] -pub struct Message { - #[serde(default)] - pub id: Option, +pub struct Entry { pub role: String, + /// Visible text; a tool call's input, a tool result's output. pub text: String, + /// Identifier pairing a tool call with its result; tool entries only. + #[serde(default)] + pub call_id: Option, + /// Identity of the item the entry belongs to, as the traffic carries it. + #[serde(default)] + pub item_id: Option, } impl Conversation { /// The completely decoded conversation of an event, as the `AgentEvent` /// contract defines it. - pub fn read(metadata: &Value) -> Option { - let value = metadata.get("conversation")?; - let conversation: Self = Self::deserialize(value).ok()?; + pub fn read(conversation: Option<&Value>) -> Option { + let conversation: Self = Self::deserialize(conversation?).ok()?; (conversation.completeness == Completeness::Complete).then_some(conversation) } } @@ -46,39 +47,38 @@ mod tests { #[test] fn roles_are_kept_as_reported() { - let metadata = serde_json::json!({"conversation": { - "messages": [ - {"role": "reasoning", "text": "plan"}, + let value = serde_json::json!({ + "entries": [ + {"role": "reasoning", "text": "plan", "item_id": "rs_1"}, {"role": "some_new_role", "text": "other"}, + {"role": "tool_call", "text": "{}", "call_id": "call_1", "name": "exec"}, {"role": "user", "text": "hello"} ], - "instructions_present": false, "completeness": "complete" - }}); - let conversation = Conversation::read(&metadata).expect("complete conversation"); + }); + let conversation = Conversation::read(Some(&value)).expect("complete conversation"); assert_eq!( conversation - .messages + .entries .iter() - .map(|message| message.role.as_str()) + .map(|entry| entry.role.as_str()) .collect::>(), - ["reasoning", "some_new_role", "user"] + ["reasoning", "some_new_role", "tool_call", "user"] ); + let call = &conversation.entries[2]; + assert_eq!(call.call_id.as_deref(), Some("call_1")); } #[test] fn conversations_outside_the_contract_are_not_read() { - let messages = serde_json::json!([{"role": "user", "text": "hello"}]); + let entries = serde_json::json!([{"role": "user", "text": "hello"}]); for conversation in [ - serde_json::json!({"messages": messages, "instructions_present": false}), - serde_json::json!({ - "messages": messages, - "instructions_present": false, - "completeness": "partial" - }), + serde_json::json!({"entries": entries}), + serde_json::json!({"entries": entries, "completeness": "partial"}), + serde_json::json!({"messages": entries, "completeness": "complete"}), ] { - let metadata = serde_json::json!({"conversation": conversation}); - assert!(Conversation::read(&metadata).is_none()); + assert!(Conversation::read(Some(&conversation)).is_none()); } + assert!(Conversation::read(None).is_none()); } } diff --git a/crates/abyss-backend/src/usage/event_order.rs b/crates/abyss-backend/src/usage/event_order.rs deleted file mode 100644 index 4d4421e..0000000 --- a/crates/abyss-backend/src/usage/event_order.rs +++ /dev/null @@ -1,189 +0,0 @@ -//! Deterministic ordering for normalized events inside Agent timeline turns. -//! -//! Events are ordered by turn, high-precision observation time, event side, -//! and stable event identity. - -use chrono::{DateTime, Utc}; - -/// Event fields required by deterministic timeline ordering. -pub trait TimelineEvent { - fn turn_index(&self) -> i32; - fn observed_at(&self) -> DateTime; - fn event_type(&self) -> &str; - fn event_id(&self) -> &str; -} - -#[cfg(feature = "postgres-es")] -impl TimelineEvent for crate::db::models::UsageEvent { - fn turn_index(&self) -> i32 { - self.turn_index - } - - fn observed_at(&self) -> DateTime { - self.observed_at - } - - fn event_type(&self) -> &str { - &self.event_type - } - - fn event_id(&self) -> &str { - &self.event_id - } -} - -impl TimelineEvent for super::UsageEventOutline { - fn turn_index(&self) -> i32 { - self.turn_index - } - - fn observed_at(&self) -> DateTime { - self.observed_at - } - - fn event_type(&self) -> &str { - &self.event_type - } - - fn event_id(&self) -> &str { - &self.event_id - } -} - -/// Applies the authoritative event order used by owned and shared timelines. -pub struct UsageEventTimelineOrder; - -impl UsageEventTimelineOrder { - /// Sorts events by turn, observation time, side, and ID. - /// - /// The final event identifier tie-breaker makes output deterministic even - /// for collectors with identical timestamps. - pub fn sort(events: &mut [T]) { - events.sort_by(|left, right| { - left.turn_index() - .cmp(&right.turn_index()) - .then_with(|| left.observed_at().cmp(&right.observed_at())) - .then_with(|| { - event_side_rank(left.event_type()).cmp(&event_side_rank(right.event_type())) - }) - .then_with(|| left.event_id().cmp(right.event_id())) - }); - } -} - -const fn event_side_rank(event_type: &str) -> u8 { - match event_type.as_bytes() { - b"request" => 0, - b"response" => 1, - _ => 2, - } -} - -#[cfg(test)] -mod tests { - use chrono::{DateTime, TimeZone as _, Timelike as _, Utc}; - - use super::{TimelineEvent, UsageEventTimelineOrder}; - - struct Sample { - turn_index: i32, - observed_at: DateTime, - event_type: &'static str, - event_id: &'static str, - } - - impl TimelineEvent for Sample { - fn turn_index(&self) -> i32 { - self.turn_index - } - - fn observed_at(&self) -> DateTime { - self.observed_at - } - - fn event_type(&self) -> &str { - self.event_type - } - - fn event_id(&self) -> &str { - self.event_id - } - } - - fn timestamp() -> DateTime { - Utc.with_ymd_and_hms(2026, 8, 4, 2, 12, 24) - .single() - .expect("timestamp should be valid") - } - - fn event( - turn_index: i32, - observed_at: DateTime, - event_type: &'static str, - event_id: &'static str, - ) -> Sample { - Sample { - turn_index, - observed_at, - event_type, - event_id, - } - } - - fn sorted_ids(mut events: Vec) -> Vec<&'static str> { - UsageEventTimelineOrder::sort(&mut events); - events.iter().map(|event| event.event_id).collect() - } - - #[test] - fn turn_precedes_observation_time() { - let first = timestamp(); - let second = first + chrono::TimeDelta::microseconds(1); - assert_eq!( - sorted_ids(vec![ - event(2, first, "request", "turn-2"), - event(1, second, "request", "turn-1"), - ]), - ["turn-1", "turn-2"] - ); - } - - #[test] - fn high_precision_timestamp_orders_events_within_a_turn() { - let first = timestamp() - .with_nanosecond(1_000) - .expect("microsecond should be valid"); - let second = first - .with_nanosecond(2_000) - .expect("microsecond should be valid"); - assert_eq!( - sorted_ids(vec![ - event(1, second, "request", "second"), - event(1, first, "response", "first"), - ]), - ["first", "second"] - ); - } - - #[test] - fn request_precedes_response_for_timestamp_ties() { - assert_eq!( - sorted_ids(vec![ - event(1, timestamp(), "response", "response"), - event(1, timestamp(), "request", "request"), - ]), - ["request", "response"] - ); - } - - #[test] - fn event_identity_breaks_remaining_ties() { - assert_eq!( - sorted_ids(vec![ - event(1, timestamp(), "response", "b"), - event(1, timestamp(), "response", "a"), - ]), - ["a", "b"] - ); - } -} diff --git a/crates/abyss-backend/src/usage/mod.rs b/crates/abyss-backend/src/usage/mod.rs index f75cd14..b214d42 100644 --- a/crates/abyss-backend/src/usage/mod.rs +++ b/crates/abyss-backend/src/usage/mod.rs @@ -1,7 +1,7 @@ //! API contracts and shared transformations for Agent usage collection. //! //! Ingest types preserve collector-provided observations, while response types -//! expose the normalized device/session/turn hierarchy returned by every storage +//! expose the normalized device/session hierarchy returned by every storage //! backend. Deterministic validation and metadata transformations live here; //! backend-specific persistence stays behind the storage compatibility layer. @@ -12,18 +12,27 @@ pub mod comparison; mod conversation; /// Opaque diagnostic-capture ingest contracts. pub mod diagnostics; -pub mod event_order; pub mod persistence; +/// Previews of an event's conversation entries. +pub mod preview; +/// One segment's exchanges for the session page. +pub mod segments; +/// The header of the session page. +pub mod session_header; /// Owner-scoped session index contracts. pub mod session_list; -pub mod session_title; -pub mod trajectory; +/// Streams of a session: stream facts, bases, and segments. +pub mod streams; +/// Local calendar days for summaries grouped by day. +pub mod time_zone; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; use uuid::Uuid; +use self::streams::EventLocation; + #[derive(Deserialize)] /// Atomic batch accepted by the event-ingestion endpoint. pub struct IngestEventsRequest { @@ -47,33 +56,37 @@ pub struct IngestUsageEvent { pub agent: AgentPayload, /// Agent-native session identifier. pub session_id: String, - /// The user turn this request works for, unique within its session; - /// absent when the collector cannot identify a particular turn. - pub turn_id: Option, - /// Whether the request comes from a thread the Agent opened only to name - /// another session. - #[serde(default)] - pub naming_thread: bool, /// The thread that sent this exchange's request, as the Agent declares /// it; requests with the same value continue one conversation. pub stream_id: Option, /// Provider identity of the response this exchange produced; both events - /// of the exchange carry it. Used to group a fallback turn when `turn_id` - /// is absent; one user turn can contain several such model calls. + /// of the exchange carry it. pub response_id: Option, /// The response this exchange's request continues after; the request then /// carries only the messages added since. pub previous_response_id: Option, + /// Id of the sub-agent this exchange belongs to, as the client declares + /// it; absent on the main line and its side requests. + pub subagent_id: Option, + /// Id of the sub-agent that started this one, as the client declares it; + /// absent when the main line started it. + pub parent_subagent_id: Option, + /// The turn this exchange works for: the user input that started it, as + /// the client declares it. A sub-agent and a side request carry the value + /// of the turn they work for. + pub turn_id: Option, /// LLM provider and model labels. pub llm: LlmPayload, /// Whether the event represents the request or response side of a call. pub event_type: UsageEventType, - /// Optional prompt or response text. - pub text: Option, + /// Conversation of the `AgentEvent`, stored exactly as received. + #[serde(default)] + pub conversation: Option, /// Token counters of the `AgentEvent`; absent counters are stored as zero. #[serde(default)] pub token_usage: TokenUsagePayload, - /// Collector metadata retained with the raw event. + /// Collector metadata retained with the raw event; the conversation is + /// not part of it. #[serde(default = "empty_metadata")] pub metadata: Value, /// Optional image metadata and content ordered within the event. @@ -180,6 +193,9 @@ pub struct SummaryQuery { pub to: Option>, /// Comma-separated grouping dimensions. pub group_by: Option, + /// IANA time zone whose calendar days the `day` dimension follows; UTC + /// when absent. + pub tz: Option, /// Selects the full or token-only response shape. pub fields: Option, /// User UUID, name, or email filter retained for the summary contract. @@ -265,7 +281,7 @@ pub struct SummaryRow { pub event_type: Option, /// Distinct sessions in this bucket. pub sessions: usize, - /// Distinct turns in this bucket. + /// Requests in this bucket whose stream facts say they open a turn. pub turns: usize, /// Request events in this bucket. pub requests: i64, @@ -363,8 +379,6 @@ pub struct RawEventsQuery { pub session_id: Vec, /// Backend session primary-key filter. pub session_pk: Option, - /// Normalized turn-number filter. - pub turn_index: Option, /// LLM provider filter; repeat the key to match any of several providers. #[serde(default)] pub llm_provider: Vec, @@ -379,10 +393,13 @@ pub struct RawEventsQuery { pub offset: Option, /// Selects complete events or the content-free list projection. pub fields: Option, + /// Includes each event's current segment and exchange when requested. + #[serde(default)] + pub locate: bool, } impl RawEventsQuery { - /// Whether events carry their text, metadata, and attachments. + /// Whether events carry their conversation, metadata, and attachments. pub fn includes_content(&self) -> bool { self.fields.unwrap_or_default() == RawEventFields::Full } @@ -392,11 +409,12 @@ impl RawEventsQuery { #[serde(rename_all = "snake_case")] /// Response projection requested from the raw-event endpoint. pub enum RawEventFields { - /// Return every event field, including text, metadata, and attachments. + /// Return every event field, including conversation, metadata, and attachments. #[default] Full, - /// Return what an event list shows. Text and metadata hold whole request - /// bodies, so storage does not read them and attachments are not loaded. + /// Return what an event list shows. Conversations hold whole request + /// bodies, so storage reads neither them nor metadata, and attachments are + /// not loaded. List, } @@ -434,7 +452,7 @@ impl RawEventListResponse { #[derive(Debug, Serialize)] /// One usage event as an event list shows it: identity, model, Agent, device, -/// and token counts, without text, metadata, or attachments. +/// and token counts, without conversation, metadata, or attachments. pub struct UsageEventListItem { /// Backend-generated event primary key. pub id: Uuid, @@ -452,16 +470,15 @@ pub struct UsageEventListItem { pub platform: Option, /// Backend session primary key. pub session_pk: Uuid, - /// Backend turn primary key. - pub turn_pk: Uuid, + /// Current stream position, when the request asks for event locations. + #[serde(skip_serializing_if = "Option::is_none")] + pub location: Option, /// Agent name as reported by the collector. pub agent_name: String, /// Agent version observed on this event. pub agent_version: Option, /// Agent-native session identifier. pub session_id: String, - /// Backend-normalized turn number. - pub turn_index: i32, /// LLM provider as reported by the collector. pub llm_provider: String, /// Provider-specific model identifier. @@ -494,11 +511,10 @@ impl From for UsageEventListItem { host_name: event.host_name, platform: event.platform, session_pk: event.session_pk, - turn_pk: event.turn_pk, + location: event.location, agent_name: event.agent_name, agent_version: event.agent_version, session_id: event.session_id, - turn_index: event.turn_index, llm_provider: event.llm_provider, llm_model: event.llm_model, event_type: event.event_type, @@ -532,26 +548,21 @@ pub struct UsageEventResponse { pub platform: Option, /// Backend session primary key. pub session_pk: Uuid, - /// Backend turn primary key. - pub turn_pk: Uuid, + /// Current stream position, when the request asks for event locations. + #[serde(skip_serializing_if = "Option::is_none")] + pub location: Option, /// Agent name as reported by the collector. pub agent_name: String, /// Agent version observed on this event. pub agent_version: Option, /// Agent-native session identifier. pub session_id: String, - /// Backend-normalized turn number. - pub turn_index: i32, /// LLM provider as reported by the collector. pub llm_provider: String, /// Provider-specific model identifier. pub llm_model: String, /// Canonical request or response value. pub event_type: String, - /// Optional prompt or response text. - pub text: Option, - /// Lowercase SHA-256 of text when text is present. - pub text_sha256: Option, /// Normalized input-token count. pub input_tokens: i64, /// Normalized output-token count. @@ -575,118 +586,15 @@ pub struct UsageEventResponse { /// Response whose conversation the exchange's request continues. #[serde(skip_serializing_if = "Option::is_none")] pub previous_response_id: Option, + /// Conversation of the event, as the collector reported it. + #[serde(skip_serializing_if = "Option::is_none")] + pub conversation: Option, /// Collector metadata retained without schema expansion. pub metadata: Value, /// Ordered attachment metadata associated with the event. pub attachments: Vec, } -#[derive(Debug, Serialize)] -/// Session metadata and its ordered turn timelines. -/// -/// `E` is the per-event representation: complete events for the full -/// timeline, or [`UsageEventOutline`] for the content-free outline. -pub struct SessionTimelineResponse { - /// Authoritative session information. - pub session: SessionInfo, - /// Turns ordered by normalized turn index. - pub turns: Vec>, -} - -/// Session timeline whose events omit text and metadata bodies. -pub type SessionOutlineResponse = SessionTimelineResponse; - -/// Maximum number of characters returned in an outline event preview. -pub const OUTLINE_PREVIEW_CHARS: i64 = 200; - -#[derive(Debug, Serialize)] -/// Content-free representation of one event, sized independently of its text. -/// -/// Full content is loaded separately by event primary key. -pub struct UsageEventOutline { - /// Backend-generated event primary key. - pub id: Uuid, - /// Collector-generated idempotency key. - pub event_id: String, - /// Owning turn, used only to group events. - #[serde(skip)] - pub turn_pk: Uuid, - /// Backend-normalized turn number. - pub turn_index: i32, - /// LLM provider as reported by the collector. - pub llm_provider: String, - /// Provider-specific model identifier. - pub llm_model: String, - /// Canonical request or response value. - pub event_type: String, - /// Normalized input-token count. - pub input_tokens: i64, - /// Normalized output-token count. - pub output_tokens: i64, - /// Normalized cache-read-token count. - pub cache_read_tokens: i64, - /// Normalized cache-write-token count. - pub cache_write_tokens: i64, - /// Normalized reasoning-token count. - pub reasoning_tokens: i64, - /// Provider-reported or derived total-token count. - pub total_tokens: i64, - /// Collector observation timestamp. - pub observed_at: DateTime, - /// Text length in characters; zero when text is absent. - pub text_chars: i64, - /// Up to [`OUTLINE_PREVIEW_CHARS`] characters of text: the tail of a - /// request, where the newest prompt is, or the head of a response. - pub preview: Option, - /// Number of image attachments on the event. - pub attachment_count: i64, -} - -#[derive(Debug, Serialize)] -/// API representation of an Agent session. -pub struct SessionInfo { - /// Backend session primary key. - pub session_pk: Uuid, - /// Authenticated owner identifier. - pub user_id: Uuid, - /// Most recently associated device context. - pub device_context_id: Uuid, - /// Agent name as reported by the collector. - pub agent_name: String, - /// Most recently known Agent version. - pub agent_version: Option, - /// Agent-native session identifier. - pub session_id: String, - /// Earliest event observation time in the session. - pub started_at: DateTime, - /// Latest event observation time in the session. - pub ended_at: Option>, - /// Stable, allowlisted session metadata. - pub metadata: Value, - /// Agent-generated title, or a first-user-input preview for other Agents. - pub title: Option, - /// Why the session exists when it only serves the Agent, such as - /// `title_generation`; `None` for ordinary sessions. - pub purpose: Option, - /// For a naming thread, the task session its title was matched to. - pub parent_session_pk: Option, -} - -#[derive(Debug, Serialize)] -/// One normalized turn and its deterministically ordered events. -pub struct TurnTimeline { - /// Backend turn primary key. - pub turn_pk: Uuid, - /// Backend-normalized turn number. - pub turn_index: i32, - /// Earliest event observation time in the turn. - pub started_at: DateTime, - /// Latest event observation time in the turn. - pub ended_at: Option>, - /// Request and response events in authoritative timeline order. - pub events: Vec, -} - /// Returns the JSON object used when optional metadata is absent. pub fn empty_metadata() -> Value { json!({}) @@ -717,14 +625,16 @@ mod tests { "version": "1.2.3" }, "session_id": "session-001", - "turn_index": 1_i32, "llm": { "provider": "OpenAI", "model": "gpt-5.5" }, "event_type": "request", - "text": "Explain this repository.", "token_usage": {"input_tokens": 10_i64, "output_tokens": 0_i64, "cache_read_tokens": 2_i64, "cache_write_tokens": 0_i64, "reasoning_tokens": 0_i64, "total_tokens": 12_i64}, + "conversation": { + "entries": [{"role": "user", "text": "hello"}], + "completeness": "complete" + }, "metadata": { "source": "unit-test" } @@ -757,6 +667,14 @@ mod tests { json!({ "source": "unit-test" }), "metadata should preserve arbitrary JSON object content" ); + assert_eq!( + event.conversation, + Some(json!({ + "entries": [{"role": "user", "text": "hello"}], + "completeness": "complete" + })), + "the conversation should be kept as received" + ); } #[test] @@ -797,16 +715,15 @@ mod tests { event.token_usage.total_tokens, None, "omitted token_usage should leave every counter absent" ); - assert_eq!(event.turn_id, None, "an omitted turn names no turn"); - assert!( - !event.naming_thread, - "an omitted flag is not a naming thread" - ); assert_eq!( event.metadata, json!({}), "omitted metadata should default to an empty object" ); + assert_eq!( + event.conversation, None, + "omitted conversation stays absent" + ); } #[test] diff --git a/crates/abyss-backend/src/usage/persistence.rs b/crates/abyss-backend/src/usage/persistence.rs index e8938f8..b2bbff5 100644 --- a/crates/abyss-backend/src/usage/persistence.rs +++ b/crates/abyss-backend/src/usage/persistence.rs @@ -4,7 +4,6 @@ use std::collections::HashSet; #[cfg(feature = "postgres-es")] use chrono::{DateTime, Utc}; -use sha2::{Digest, Sha256}; use crate::error::AppError; @@ -145,10 +144,6 @@ pub fn normalize_limit(limit: Option, fallback: i64, maximum: i64) -> i64 { limit.unwrap_or(fallback).clamp(1_i64, maximum) } -pub fn sha256_bytes(value: &str) -> Vec { - Sha256::digest(value.as_bytes()).to_vec() -} - #[cfg(test)] mod tests { use super::{escape_like_pattern, exact_values_filter, parse_group_by}; diff --git a/crates/abyss-backend/src/usage/preview.rs b/crates/abyss-backend/src/usage/preview.rs new file mode 100644 index 0000000..9bdcb25 --- /dev/null +++ b/crates/abyss-backend/src/usage/preview.rs @@ -0,0 +1,273 @@ +//! Previews of one event's conversation: the chosen runs of its entries, each +//! text cut to a length, with what was left out counted. +//! +//! Entry indexes are those of the event's full entry list, so a position found +//! elsewhere, such as a search hit, points into a preview unchanged. + +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use uuid::Uuid; + +use super::{UsageEventResponse, attachments::UsageEventAttachmentResponse}; +use crate::error::AppError; + +#[derive(Debug, Default, Deserialize)] +/// Which entries of a conversation to give, and how much of each text. +pub struct PreviewQuery { + /// Comma-separated runs of entries: `a-b` from index `a` to `b`, `-n` the + /// last `n`, `a` the one at `a`. All entries without it. + pub entries: Option, + /// Most UTF-16 code units of each entry text. Whole texts without it. + pub text_preview: Option, +} + +/// One run of entries as a query asks for it. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Run { + /// Entries `from` through `to`. + Span { from: usize, to: usize }, + /// The last `count` entries. + Last { count: usize }, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +/// A validated [`PreviewQuery`]. +pub struct Window { + runs: Option>, + text_preview: Option, +} + +impl Window { + /// Validates a query; a malformed run is a validation error. + pub fn parse(query: &PreviewQuery) -> Result { + let runs = query + .entries + .as_deref() + .map(|entries| entries.split(',').map(str::trim).map(parse_run).collect()) + .transpose()?; + Ok(Self { + runs, + text_preview: query.text_preview, + }) + } + + /// The runs over `total` entries as sorted, disjoint index ranges. + fn ranges(&self, total: usize) -> Vec> { + let Some(runs) = &self.runs else { + return std::iter::once(0..total).collect(); + }; + let mut ranges = runs + .iter() + .map(|run| match *run { + Run::Span { from, to } => from.min(total)..to.saturating_add(1).min(total), + Run::Last { count } => total.saturating_sub(count)..total, + }) + .filter(|range| !range.is_empty()) + .collect::>(); + ranges.sort_by_key(|range| range.start); + let mut merged: Vec> = Vec::new(); + for range in ranges { + match merged.last_mut() { + Some(last) if range.start <= last.end => last.end = last.end.max(range.end), + _ => merged.push(range), + } + } + merged + } +} + +fn parse_run(run: &str) -> Result { + let invalid = || AppError::validation(format!("invalid entries run: {run:?}")); + let index = |text: &str| text.parse::().map_err(|_| invalid()); + if let Some(count) = run.strip_prefix('-') { + return Ok(Run::Last { + count: index(count)?, + }); + } + let (from, to) = if let Some((from, to)) = run.split_once('-') { + (index(from)?, index(to)?) + } else { + let at = index(run)?; + (at, at) + }; + if from > to { + return Err(invalid()); + } + Ok(Run::Span { from, to }) +} + +#[derive(Debug, Serialize)] +/// The chosen part of one event's conversation. +pub struct ConversationPreview { + /// Event primary key. + pub event_pk: Uuid, + /// Completeness the collector reported for the conversation. + pub completeness: Option, + /// Number of entries in the full conversation. + pub entries_total: usize, + /// The chosen runs in order. + pub chunks: Vec, + /// Number of entries no chunk holds. + pub omitted_entries: usize, + /// UTF-8 bytes of the texts of those entries. + pub omitted_bytes: usize, + /// Image attachments of the event. + pub attachments: Vec, +} + +#[derive(Debug, Serialize)] +/// One run of consecutive entries. +pub struct Chunk { + /// Index of the first entry in the full conversation. + pub from: usize, + /// The entries as the collector reported them. A text cut to the + /// preview length carries `text_length`, its full length in UTF-16 code + /// units. + pub entries: Vec, +} + +/// Previews the conversation of `event` through `window`. +pub fn preview(event: UsageEventResponse, window: &Window) -> ConversationPreview { + let conversation = event.conversation.as_ref(); + let entries = conversation + .and_then(|conversation| conversation.get("entries")) + .and_then(Value::as_array) + .map(Vec::as_slice) + .unwrap_or_default(); + let ranges = window.ranges(entries.len()); + let shown = ranges.iter().map(std::ops::Range::len).sum::(); + let omitted_bytes = entries + .iter() + .enumerate() + .filter(|(index, _)| !ranges.iter().any(|range| range.contains(index))) + .map(|(_, entry)| text_of(entry).len()) + .sum(); + let completeness = conversation + .and_then(|conversation| conversation.get("completeness")) + .and_then(Value::as_str) + .map(str::to_owned); + let chunks = ranges + .into_iter() + .map(|range| Chunk { + from: range.start, + entries: entries[range] + .iter() + .map(|entry| cut(entry, window.text_preview)) + .collect(), + }) + .collect(); + let entries_total = entries.len(); + ConversationPreview { + event_pk: event.id, + completeness, + entries_total, + chunks, + omitted_entries: entries_total.saturating_sub(shown), + omitted_bytes, + attachments: event.attachments, + } +} + +fn text_of(entry: &Value) -> &str { + entry + .get("text") + .and_then(Value::as_str) + .unwrap_or_default() +} + +/// `entry` with its text cut to `limit` UTF-16 code units, never inside a +/// character. +fn cut(entry: &Value, limit: Option) -> Value { + let mut entry = entry.clone(); + let Some(limit) = limit else { + return entry; + }; + let text = text_of(&entry); + let length = text.encode_utf16().count(); + if length <= limit { + return entry; + } + let mut units = 0_usize; + let end = text + .char_indices() + .find(|(_, character)| { + units = units.saturating_add(character.len_utf16()); + units > limit + }) + .map_or(text.len(), |(at, _)| at); + let kept = text[..end].to_owned(); + if let Some(object) = entry.as_object_mut() { + object.insert("text".to_owned(), Value::String(kept)); + object.insert("text_length".to_owned(), Value::from(length)); + } + entry +} + +#[cfg(test)] +mod tests { + use serde_json::{Value, json}; + + use super::{PreviewQuery, Window, cut}; + + fn window(entries: Option<&str>, text_preview: Option) -> Window { + Window::parse(&PreviewQuery { + entries: entries.map(str::to_owned), + text_preview, + }) + .expect("a valid window") + } + + #[test] + fn runs_resolve_to_sorted_disjoint_ranges_within_the_entries() { + /// Entries run, entry count, and the expected `(start, end)` ranges. + type Case<'a> = (Option<&'a str>, usize, &'a [(usize, usize)]); + let cases: [Case<'_>; 6] = [ + (None, 4, &[(0, 4)]), + (Some("0-2,-5"), 10, &[(0, 3), (5, 10)]), + (Some("0-2,-5"), 6, &[(0, 6)]), + (Some("-2,3"), 10, &[(3, 4), (8, 10)]), + (Some("0-99"), 3, &[(0, 3)]), + (Some("7"), 3, &[]), + ]; + for (entries, total, expected) in cases { + let ranges = window(entries, None) + .ranges(total) + .into_iter() + .map(|range| (range.start, range.end)) + .collect::>(); + assert_eq!(ranges, expected, "{entries:?} of {total}"); + } + for invalid in ["", "a", "3-1", "1-", "-"] { + assert!( + Window::parse(&PreviewQuery { + entries: Some(invalid.to_owned()), + text_preview: None, + }) + .is_err(), + "{invalid:?}" + ); + } + } + + #[test] + fn long_texts_keep_their_start_and_full_length() { + let entry = json!({"role": "user", "text": "ab😀cd"}); + assert_eq!(cut(&entry, None), entry); + assert_eq!(cut(&entry, Some(6)), entry, "six UTF-16 code units fit"); + let cut_at = |limit| { + let cut = cut(&entry, Some(limit)); + (cut["text"].clone(), cut["text_length"].clone()) + }; + assert_eq!( + cut_at(3), + (json!("ab"), json!(6_usize)), + "never half a character" + ); + assert_eq!(cut_at(4), (json!("ab😀"), json!(6_usize))); + assert_eq!( + cut(&json!({"role": "user"}), Some(1)), + json!({"role": "user"}) + ); + assert_eq!(cut(&Value::Null, Some(1)), Value::Null); + } +} diff --git a/crates/abyss-backend/src/usage/segments.rs b/crates/abyss-backend/src/usage/segments.rs new file mode 100644 index 0000000..00db22e --- /dev/null +++ b/crates/abyss-backend/src/usage/segments.rs @@ -0,0 +1,147 @@ +//! The exchange layer of a session page: one segment's exchanges with their +//! times, models, token counts, and a preview of each request and response. + +use std::collections::HashMap; + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use super::{ + UsageEventResponse, + preview::{ConversationPreview, PreviewQuery, Window, preview}, + streams::{SegmentExchange, SegmentRun}, +}; +use crate::error::AppError; + +#[derive(Debug, Default, Deserialize)] +/// Query of a segment: how each exchange's preview is cut. +pub struct SegmentQuery { + /// Runs of entries each preview gives; see [`PreviewQuery`]. + pub entries: Option, + /// Most UTF-16 code units of each previewed entry text. + pub text_preview: Option, +} + +#[derive(Debug)] +/// A validated [`SegmentQuery`]. +pub struct ValidatedSegmentQuery { + pub window: Window, +} + +impl SegmentQuery { + pub fn validate(self) -> Result { + Ok(ValidatedSegmentQuery { + window: Window::parse(&PreviewQuery { + entries: self.entries, + text_preview: self.text_preview, + })?, + }) + } +} + +#[derive(Debug, Serialize)] +/// One segment's exchanges from the asked one to its end, all at once. +pub struct SegmentResponse { + /// Session primary key. + pub session_pk: Uuid, + /// Head of the segment that now holds the asked exchange. + pub segment: String, + /// The exchanges in order. + pub exchanges: Vec, +} + +#[derive(Debug, Serialize)] +/// One exchange with what its card shows. +pub struct SegmentExchangeResponse { + #[serde(flatten)] + pub exchange: SegmentExchange, + /// Observation time of the request event. + pub requested_at: Option>, + /// Observation time of the response event. + pub responded_at: Option>, + /// Model of the request, else of the response. + pub model: Option, + pub request_tokens: Option, + pub response_tokens: Option, + pub request: Option, + pub response: Option, +} + +#[derive(Debug, Default, PartialEq, Eq, Serialize)] +/// Token counts of one event as stored, or a sum of them. +pub struct TokenCounts { + pub input: i64, + pub output: i64, + pub cache_read: i64, + pub cache_write: i64, + pub reasoning: i64, + pub total: i64, +} + +impl TokenCounts { + /// Adds `other` to these counts. + pub const fn add(&mut self, other: &Self) { + self.input = self.input.saturating_add(other.input); + self.output = self.output.saturating_add(other.output); + self.cache_read = self.cache_read.saturating_add(other.cache_read); + self.cache_write = self.cache_write.saturating_add(other.cache_write); + self.reasoning = self.reasoning.saturating_add(other.reasoning); + self.total = self.total.saturating_add(other.total); + } + + const fn of(event: &UsageEventResponse) -> Self { + Self { + input: event.input_tokens, + output: event.output_tokens, + cache_read: event.cache_read_tokens, + cache_write: event.cache_write_tokens, + reasoning: event.reasoning_tokens, + total: event.total_tokens, + } + } +} + +/// The primary keys of the events a run's exchanges hold. +pub fn run_event_pks(run: &SegmentRun) -> Vec { + run.exchanges + .iter() + .flat_map(|exchange| [exchange.request_pk, exchange.response_pk]) + .flatten() + .collect() +} + +/// Fills a run with its events, keyed by primary key. +pub fn segment_response( + session_pk: Uuid, + run: SegmentRun, + mut events: HashMap, + window: &Window, +) -> SegmentResponse { + let exchanges = run + .exchanges + .into_iter() + .map(|exchange| { + let request = exchange.request_pk.and_then(|pk| events.remove(&pk)); + let response = exchange.response_pk.and_then(|pk| events.remove(&pk)); + SegmentExchangeResponse { + requested_at: request.as_ref().map(|event| event.observed_at), + responded_at: response.as_ref().map(|event| event.observed_at), + model: request + .as_ref() + .or(response.as_ref()) + .map(|event| event.llm_model.clone()), + request_tokens: request.as_ref().map(TokenCounts::of), + response_tokens: response.as_ref().map(TokenCounts::of), + request: request.map(|event| preview(event, window)), + response: response.map(|event| preview(event, window)), + exchange, + } + }) + .collect(); + SegmentResponse { + session_pk, + segment: run.segment, + exchanges, + } +} diff --git a/crates/abyss-backend/src/usage/session_header.rs b/crates/abyss-backend/src/usage/session_header.rs new file mode 100644 index 0000000..eb1b6dc --- /dev/null +++ b/crates/abyss-backend/src/usage/session_header.rs @@ -0,0 +1,99 @@ +//! The header of a session page: who took part, when, and how much. + +use std::collections::BTreeSet; + +use chrono::{DateTime, Utc}; +use serde::Serialize; +use uuid::Uuid; + +use super::{ + segments::TokenCounts, + streams::{EntryLabel, StreamSummary}, +}; + +/// The columns of one session event the header reads. +pub struct HeaderEvent { + pub session_id: String, + pub agent_name: String, + pub agent_version: Option, + pub provider: String, + pub model: String, + pub observed_at: DateTime, + pub tokens: TokenCounts, +} + +#[derive(Debug, Serialize)] +/// One session's agents, time span, counts, token totals, and models. +pub struct SessionHeaderResponse { + /// Session primary key. + pub session_pk: Uuid, + /// Session id the agent declared. + pub session_id: String, + /// The earliest request's new-part label, shared with every session view. + pub title: Option, + /// Agents that reported events, by name and version. + pub agents: Vec, + /// Earliest event time. + pub start: DateTime, + /// Latest event time. + pub end: DateTime, + pub event_count: usize, + pub exchange_count: usize, + /// Requests that open a turn. + pub turn_count: usize, + pub segment_count: usize, + /// Token counts summed over every event. + pub tokens: TokenCounts, + pub providers: Vec, + pub models: Vec, +} + +#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Serialize)] +pub struct HeaderAgent { + pub name: String, + pub version: Option, +} + +/// The header of a session with at least one event. +pub fn session_header( + session_pk: Uuid, + events: Vec, + summary: &StreamSummary, + title: Option, +) -> SessionHeaderResponse { + let start = events.iter().map(|event| event.observed_at).min(); + let end = events.iter().map(|event| event.observed_at).max(); + let event_count = events.len(); + let session_id = events + .first() + .map(|event| event.session_id.clone()) + .unwrap_or_default(); + let mut agents = BTreeSet::new(); + let mut providers = BTreeSet::new(); + let mut models = BTreeSet::new(); + let mut tokens = TokenCounts::default(); + for event in events { + tokens.add(&event.tokens); + agents.insert(HeaderAgent { + name: event.agent_name, + version: event.agent_version, + }); + providers.insert(event.provider); + models.insert(event.model); + } + SessionHeaderResponse { + session_pk, + session_id, + title, + agents: agents.into_iter().collect(), + start: start.unwrap_or_default(), + end: end.unwrap_or_default(), + event_count, + exchange_count: summary.exchanges, + turn_count: summary.turns, + segment_count: summary.segments, + tokens, + providers: providers.into_iter().collect(), + models: models.into_iter().collect(), + } +} diff --git a/crates/abyss-backend/src/usage/session_list.rs b/crates/abyss-backend/src/usage/session_list.rs index bd931fb..cdf334f 100644 --- a/crates/abyss-backend/src/usage/session_list.rs +++ b/crates/abyss-backend/src/usage/session_list.rs @@ -105,11 +105,12 @@ pub struct SessionListItem { pub ended_at: Option>, /// Latest event observation, or the session start when it has no events. pub last_activity_at: DateTime, + /// Requests whose stream facts say they open a turn. pub turns: i64, pub events: i64, pub total_tokens: i64, - /// Agent-generated title, or a first-user-input preview for other Agents. - pub title: Option, + /// The earliest request's new-part label, without display formatting. + pub title: Option, /// Why the session exists when it only serves the Agent, such as /// `title_generation`; `None` for ordinary sessions. pub purpose: Option, diff --git a/crates/abyss-backend/src/usage/session_title.rs b/crates/abyss-backend/src/usage/session_title.rs deleted file mode 100644 index 733062f..0000000 --- a/crates/abyss-backend/src/usage/session_title.rs +++ /dev/null @@ -1,190 +0,0 @@ -//! Session names: the title an Agent generated for its own session, else the -//! session's first user input. -//! -//! The collector reports a naming call as event fields: both sides carry a -//! `purpose`, and the response carries the `session_title` it produced for its -//! own session. Titles are therefore read without interpreting message text. -//! -//! An Agent may instead name a task from a separate thread whose requests do -//! not say which session they name. The collector marks those requests with -//! `naming_thread`; their sessions are only marked, and the task sessions keep -//! their first user input. - -use serde_json::Value; - -use super::conversation::Conversation; - -/// Session purpose stored for threads that name another session. -pub const TITLE_GENERATION_PURPOSE: &str = "title_generation"; -/// Longest default title, in characters. -const DEFAULT_TITLE_CHARS: usize = 80; - -/// What one stored event contributes to its session's name. -#[derive(Debug, PartialEq, Eq)] -pub enum TitleContribution<'a> { - /// Title the Agent generated for this event's session. - Agent(&'a str), - /// An ordinary request; its first user input may name the session. - UserInput, - /// The session is a thread that names another session. - TitleThread, -} - -impl<'a> TitleContribution<'a> { - /// Classifies an event from collector fields without reading its messages. - /// Requests the Agent made for its own purposes, such as naming calls, are - /// not user input. - pub fn of(event_type: &str, metadata: &'a Value, naming_thread: bool) -> Option { - if let Some(title) = metadata.get("session_title").and_then(Value::as_str) { - return Some(Self::Agent(title)); - } - if event_type != "request" || metadata.get("purpose").is_some() { - return None; - } - if naming_thread { - return Some(Self::TitleThread); - } - Some(Self::UserInput) - } -} - -/// Default name: the first non-empty user input, on one line, at most 80 chars. -pub fn default_title(metadata: &Value) -> Option { - let input = Conversation::read(metadata)? - .messages - .into_iter() - .find(|message| message.role == "user" && !message.text.trim().is_empty())? - .text; - let input = input.split_whitespace().collect::>().join(" "); - let mut chars = input.chars(); - let mut title: String = chars.by_ref().take(DEFAULT_TITLE_CHARS).collect(); - if chars.next().is_some() { - title.pop(); - title.push('…'); - } - (!title.is_empty()).then_some(title) -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use super::*; - - /// Metadata from the current collector, which always reports completeness. - fn request(conversation: &Value) -> Value { - let mut conversation = conversation.clone(); - if conversation.get("completeness").is_none() { - conversation["completeness"] = json!("complete"); - } - json!({"conversation": conversation}) - } - - fn default_of(messages: &Value) -> Option { - default_title(&request(&json!({"messages": messages}))) - } - - #[test] - fn default_title_reads_the_first_nonempty_user_input() { - assert_eq!( - default_of(&json!([ - {"role":"system","text":"Instructions"}, - {"role":"assistant","text":"Previous reply"}, - {"role":"tool","text":"Tool output"}, - {"role":"user","text":" "}, - {"role":"user","text":" 修复\n 登录问题 "}, - {"role":"user","text":"Second request"} - ])) - .as_deref(), - Some("修复 登录问题") - ); - } - - #[test] - fn default_title_truncates_unicode_and_leaves_short_input_intact() { - for (input, expected) in [ - ("修复登录".to_owned(), "修复登录".to_owned()), - ("修".repeat(80), "修".repeat(80)), - ("修".repeat(81), format!("{}…", "修".repeat(79))), - ] { - assert_eq!( - default_of(&json!([{"role":"user","text":input}])), - Some(expected) - ); - } - } - - #[test] - fn default_title_requires_readable_user_input() { - for messages in [ - json!([]), - json!([{"role":"system","text":"Instructions"}]), - json!([{"role":"user","text":" \n "}]), - json!([{"role":"tool","text":"Output"}]), - ] { - assert_eq!(default_of(&messages), None); - } - assert_eq!(default_title(&json!({"text":"No role"})), None); - } - - #[test] - fn agent_titles_come_only_from_the_collector_field() { - assert_eq!( - TitleContribution::of( - "response", - &json!({"purpose":"session_title","session_title":"登录修复"}), - false - ), - Some(TitleContribution::Agent("登录修复")) - ); - // A reply that merely looks like a title object is ordinary output. - assert_eq!( - TitleContribution::of( - "response", - &request(&json!({"messages":[ - {"role":"assistant","text":"{\"title\":\"Quoted\"}"} - ]})), - false - ), - None - ); - } - - #[test] - fn naming_requests_are_not_user_input() { - let mut metadata = request(&json!({"messages":[{"role":"user","text":"…"}]})); - metadata["purpose"] = json!("session_title"); - assert_eq!(TitleContribution::of("request", &metadata, false), None); - } - - #[test] - fn naming_threads_are_marked_from_the_collector_flag() { - let metadata = - request(&json!({"messages":[{"role":"user","text":"Generate a concise title"}]})); - assert_eq!( - TitleContribution::of("request", &metadata, true), - Some(TitleContribution::TitleThread) - ); - assert_eq!( - TitleContribution::of("request", &metadata, false), - Some(TitleContribution::UserInput) - ); - } - - #[test] - fn incomplete_conversations_give_no_default() { - for status in ["partial", "unsupported"] { - let metadata = request( - &json!({"completeness":status,"messages":[{"role":"user","text":"hello"}]}), - ); - assert_eq!(default_title(&metadata), None); - } - } - - #[test] - fn content_outside_collector_messages_is_never_interpreted() { - let metadata = json!({"raw_body":{"parts":[{"content_base64":"e30="}]}, - "conversation":{"messages":[], "completeness":"complete", "provider":{"input":"secret"}}}); - assert_eq!(default_title(&metadata), None); - } -} diff --git a/crates/abyss-backend/src/usage/streams.rs b/crates/abyss-backend/src/usage/streams.rs new file mode 100644 index 0000000..73f9453 --- /dev/null +++ b/crates/abyss-backend/src/usage/streams.rs @@ -0,0 +1,1016 @@ +//! Streams of one session: the facts ingest reads from each event, and the +//! exchanges, bases, and segments a query derives from them. +//! +//! A stream is one line of the model's context. Every request has at most one +//! base, the exchange whose output its context continues. Ingest reads the +//! roles, call ids, and item ids of one event's conversation entries, the +//! digest of model text that has neither id, plus whether a user entry has +//! text, and never looks at another event. A query +//! groups one session's events into exchanges and finds each base by key +//! lookups within the session, without comparing content. + +mod base; +mod facts; +#[cfg(test)] +mod fixtures; +mod segmentation; +mod timeline; +mod turns; + +use std::collections::HashMap; + +use chrono::{DateTime, Utc}; +use serde::Serialize; +use uuid::Uuid; + +pub use self::{ + base::{BaseRule, NoBaseReason}, + facts::{Label, RequestFacts, ResponseFacts}, + segmentation::SegmentKind, + timeline::SessionTurnsResponse, +}; +use self::{ + base::{Exchange, emitters, find_bases, group_exchanges}, + segmentation::{Segment, Segmentation}, + turns::Turns, +}; + +#[derive(Clone, Debug, PartialEq, Eq)] +/// Stored stream facts of one event, by its side of the exchange. +pub enum EventFacts { + /// The event is a request. + Request(RequestFacts), + /// The event is a response. + Response(ResponseFacts), +} + +#[derive(Clone, Debug)] +/// One session event with the stored columns streams are derived from. +pub struct StreamEvent { + /// Event primary key. + pub pk: Uuid, + /// Collector-generated idempotency key. + pub event_id: String, + /// Collector observation time. + pub observed_at: DateTime, + /// Shared by the request and response events of one exchange. + pub response_id: Option, + /// Response the request names as the one it continues. + pub previous_response_id: Option, + /// Sub-agent the client declared the event belongs to. + pub subagent_id: Option, + /// Sub-agent the client declared as the one that started it. + pub parent_subagent_id: Option, + /// Turn the client declared the event ultimately works for. + pub turn_id: Option, + /// Facts read at ingest; an event stored before they were read has the + /// default facts of its side. + pub facts: EventFacts, +} + +#[derive(Debug, Serialize)] +/// Segments of one session, without their exchanges or any content. +pub struct SessionStreamsResponse { + /// Session primary key. + pub session_pk: Uuid, + /// Segments by their first exchange. + pub segments: Vec, +} + +#[derive(Clone, Debug, Serialize)] +/// A run of exchanges along one stream. +pub struct StreamSegment { + /// Id of the segment's first exchange, which identifies the segment. + pub head: String, + /// Why the segment begins. + pub kind: SegmentKind, + /// Head of the segment holding the first exchange's base. + pub parent: Option, + /// Id of the first exchange's base. + pub forked_from: Option, + /// Earliest event time of the segment. + pub start: DateTime, + /// Latest event time of the segment. + pub end: DateTime, + /// Number of exchanges in the segment. + pub exchange_count: usize, + /// The new user words of the segment's first request. + pub label: Option, + /// Turn of the segment's first exchange: the declared value, or the id of + /// the exchange that opened it; `None` when it belongs to no turn. + pub turn_id: Option, + /// One-based number of that turn in the session; the segments of one + /// turn share it. + pub turn: Option, + /// Sub-agent the first exchange declares. + pub subagent_id: Option, + /// Always `None`: sub-agent segments are no longer attached to their + /// parent; the field stays for the response shape. + pub attached_to: Option, + /// Always `None`, as `attached_to`. + pub attached_by: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +#[expect( + dead_code, + reason = "attachment is gone; the type stays for `attached_by` in the response shape" +)] +/// How a sub-agent's root segment was attached to a segment of its parent +/// stream. Attachment only decides where the segment shows; it changes no +/// base and no segment. +pub enum AttachedBy { + /// The first parent segment holding a request with the sub-agent's + /// declared turn. + Declared, + /// The earliest parent segment whose time range holds the sub-agent's + /// start. + Time, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +/// Where a request's new user words are, and how they begin: the last +/// `user` entry with text in its new part. +pub struct EntryLabel { + /// The request event. + pub event_pk: Uuid, + /// Index of the entry among the event's entries. + pub entry_index: i32, + /// At most the first 50 UTF-16 code units of the entry's text. + pub text: String, + /// Length of the entry's whole text in UTF-16 code units. + pub length: i32, +} + +impl EntryLabel { + /// A session title or segment label from the same stored request facts. + pub fn from_columns( + event_pk: Option, + entry_index: Option, + text: Option, + length: Option, + ) -> Option { + let label = Label::from_columns(entry_index, text, length)?; + Some(Self { + event_pk: event_pk?, + entry_index: label.index, + text: label.text, + length: label.length, + }) + } +} + +/// Immutable structure of one session at one stored revision. Contains stream +/// facts and short labels, never event bodies. Share it through `Arc`. +pub struct SessionStructure { + exchanges: Vec, + bases: Vec>, + turns: Turns, + segmentation: Segmentation, + segments: Vec, + summary: StreamSummary, + exchange_by_id: HashMap, + event_to_exchange: HashMap, +} + +impl SessionStructure { + /// Derives the streams of one session from its events. + /// + /// `events` are the session's events in session order: observation time, + /// request before response, then primary key. + pub fn derive(events: Vec) -> Self { + let mut exchanges = group_exchanges(events); + let by_response_id = exchanges + .iter() + .enumerate() + .filter_map(|(index, exchange)| Some((exchange.response_id.as_deref()?, index))) + .collect::>(); + let by_id = emitters(&exchanges); + let bases = find_bases(&exchanges, &by_response_id, &by_id); + let turns = Turns::assign(&exchanges, &bases); + let segmentation = Segmentation::cut(&exchanges, &bases, &turns); + let summary = StreamSummary { + exchanges: exchanges.len(), + turns: turns.ids.len(), + segments: segmentation.segments.len(), + }; + let mut exchange_by_id = HashMap::with_capacity(exchanges.len()); + for (index, exchange) in exchanges.iter().enumerate() { + exchange_by_id.entry(exchange.id.clone()).or_insert(index); + } + let event_to_exchange = exchanges + .iter_mut() + .enumerate() + .flat_map(|(index, exchange)| { + std::mem::take(&mut exchange.event_pks) + .into_iter() + .map(move |pk| (pk, index)) + }) + .collect(); + let mut structure = Self { + exchanges, + bases, + turns, + segmentation, + segments: Vec::new(), + summary, + exchange_by_id, + event_to_exchange, + }; + structure.segments = structure + .segmentation + .segments + .iter() + .map(|segment| structure.segment(segment)) + .collect(); + structure + } + + /// The cached counts; no exchange traversal is needed. + pub const fn summary(&self) -> StreamSummary { + self.summary + } + + /// Copies only the segment metadata returned by the stream endpoint. + pub fn session_streams(&self, session_pk: Uuid) -> SessionStreamsResponse { + SessionStreamsResponse { + session_pk, + segments: self.segments.clone(), + } + } + + /// Exchanges from `from` to the end of the segment currently holding it. + pub fn segment_run(&self, from: &str) -> Option { + let from = *self.exchange_by_id.get(from)?; + let segment = self.segmentation.segment_of[from]; + let members = &self.segmentation.segments[segment].exchanges; + let at = members.iter().position(|&exchange| exchange == from)?; + Some(SegmentRun { + segment: self.segment_head(segment).to_owned(), + exchanges: members[at..] + .iter() + .map(|&exchange| self.exchange(exchange)) + .collect(), + }) + } + + /// Current segment and exchange of a stable event anchor. + pub fn event_position(&self, event_pk: Uuid) -> Option { + let exchange = *self.event_to_exchange.get(&event_pk)?; + let segment = self.segmentation.segment_of[exchange]; + Some(EventLocation { + segment: self.segment_head(segment).to_owned(), + turn: self.turn(exchange).1, + segment_kind: self.segments[segment].kind, + exchange: self.exchanges[exchange].id.clone(), + }) + } + + /// Location endpoint response, with the event and session identities. + pub fn event_location( + &self, + session_pk: Uuid, + event_pk: Uuid, + ) -> Option { + Some(EventLocationResponse { + session_pk, + event_pk, + location: self.event_position(event_pk)?, + }) + } + + /// Segment of one event, without building a map for unrelated events. + #[cfg(feature = "sqlite-fts")] + pub fn event_segment(&self, event_pk: Uuid) -> Option<&str> { + let exchange = *self.event_to_exchange.get(&event_pk)?; + Some(self.segment_head(self.segmentation.segment_of[exchange])) + } + + /// Segment heads of the events that participate in the derived exchanges. + #[cfg(test)] + pub fn event_segments(&self) -> HashMap { + self.event_to_exchange + .iter() + .map(|(&pk, &exchange)| { + ( + pk, + self.segment_head(self.segmentation.segment_of[exchange]) + .to_owned(), + ) + }) + .collect() + } + + /// The label of an exchange's request. + fn label(&self, exchange: usize) -> Option { + let request = self.exchanges.get(exchange)?.request.as_ref()?; + let label = request.facts.label.as_ref()?; + Some(EntryLabel { + event_pk: request.pk, + entry_index: label.index, + text: label.text.clone(), + length: label.length, + }) + } + + /// Identity and one-based number of an exchange's turn. + fn turn(&self, exchange: usize) -> (Option, Option) { + self.turns.of[exchange].map_or((None, None), |turn| { + ( + Some(self.turns.ids[turn].clone()), + Some(Turns::number(turn)), + ) + }) + } + + fn base(&self, exchange: usize) -> Option { + self.bases[exchange].ok().map(|(base, _)| base) + } + + fn segment(&self, segment: &Segment) -> StreamSegment { + let id = |exchange: usize| self.exchanges[exchange].id.clone(); + let first = segment.exchanges.first().copied().unwrap_or_default(); + let (turn_id, turn) = self.turn(first); + let members = segment + .exchanges + .iter() + .map(|&exchange| &self.exchanges[exchange]); + StreamSegment { + head: id(first), + kind: segment.kind, + parent: segment + .parent + .and_then(|parent| self.segmentation.segments[parent].exchanges.first()) + .map(|&head| id(head)), + forked_from: segment.parent.and_then(|_| self.base(first)).map(id), + start: members + .clone() + .map(|exchange| exchange.at) + .min() + .unwrap_or_default(), + end: members + .map(|exchange| exchange.end) + .max() + .unwrap_or_default(), + exchange_count: segment.exchanges.len(), + label: self.label(first), + turn_id, + turn, + subagent_id: self.exchanges[first].subagent_id.clone(), + attached_to: None, + attached_by: None, + } + } +} + +/// Derives the segments of one session from its events. +/// +/// `events` are the session's events in session order: observation time, +/// request before response, then primary key. +#[cfg(test)] +pub fn session_streams(session_pk: Uuid, events: Vec) -> SessionStreamsResponse { + SessionStructure::derive(events).session_streams(session_pk) +} + +/// How many exchanges, turns, and segments one session's events form. +#[derive(Clone, Copy, Debug)] +pub struct StreamSummary { + pub exchanges: usize, + /// Turns the exchanges belong to: the distinct declared turns when the + /// session declares any, else the requests that open a turn. + pub turns: usize, + pub segments: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +/// The current position of an event in its session's stream view. +pub struct EventLocation { + /// The current segment head; it can change when a later fork is observed. + pub segment: String, + /// One-based number of the turn of the event's exchange; `None` when the + /// exchange belongs to no turn. + pub turn: Option, + /// Why this segment begins. + pub segment_kind: SegmentKind, + /// The exchange containing the event. + pub exchange: String, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +/// The current location of an event in its owner's session structure. +pub struct EventLocationResponse { + pub session_pk: Uuid, + pub event_pk: Uuid, + #[serde(flatten)] + pub location: EventLocation, +} + +#[derive(Debug, Serialize)] +/// One exchange of a segment, as the segment layer gives it. +pub struct SegmentExchange { + /// Shared `response_id`, or the `event_id` of an event without one. + pub id: String, + /// Request event primary key. + pub request_pk: Option, + /// Response event primary key. + pub response_pk: Option, + /// Exchange the request's context continues. + pub base: Option, + /// Head of the segment holding the base. + pub base_segment: Option, + /// Rule that found the base. + pub rule: Option, + /// Why there is no base. + pub no_base_reason: Option, + /// Index of the request's entries where its new part begins. + pub new_from: Option, + /// Whether the request opens a turn. + pub starts_turn: bool, + /// Sub-agent the client declared: the request's, else the response's. + pub subagent_id: Option, + /// Parent sub-agent the client declared. + pub parent_subagent_id: Option, + /// Turn the exchange belongs to: the declared value, or the id of the + /// exchange that opened it; `None` when it belongs to no turn. + pub turn_id: Option, + /// One-based number of that turn in the session. + pub turn: Option, +} + +/// Exchanges of one segment from a given exchange on. +pub struct SegmentRun { + /// Head of the segment that now holds the asked exchange. + pub segment: String, + /// The exchanges in order. + pub exchanges: Vec, +} + +/// The exchanges of the segment holding exchange `from`, from it to the +/// segment's end. `None` when the session has no such exchange. +#[cfg(test)] +pub fn segment_run(events: Vec, from: &str) -> Option { + SessionStructure::derive(events).segment_run(from) +} + +/// The head of the segment holding each event's exchange, by event primary +/// key, including additional events of an exchange's already populated side. +#[cfg(test)] +pub fn event_segments(events: Vec) -> HashMap { + SessionStructure::derive(events).event_segments() +} + +impl SessionStructure { + /// Id of a segment's first exchange. + fn segment_head(&self, segment: usize) -> &str { + let head = self.segmentation.segments[segment] + .exchanges + .first() + .copied() + .unwrap_or_default(); + &self.exchanges[head].id + } + + fn exchange(&self, index: usize) -> SegmentExchange { + let exchange = &self.exchanges[index]; + let base = self.base(index); + let request = exchange.request.as_ref(); + let (turn_id, turn) = self.turn(index); + SegmentExchange { + id: exchange.id.clone(), + request_pk: request.map(|request| request.pk), + response_pk: exchange.response.as_ref().map(|response| response.pk), + base: base.map(|base| self.exchanges[base].id.clone()), + base_segment: base.map(|base| { + self.segment_head(self.segmentation.segment_of[base]) + .to_owned() + }), + rule: self.bases[index].ok().map(|(_, rule)| rule), + no_base_reason: self.bases[index].err(), + new_from: request.map(|request| request.facts.new_from), + starts_turn: exchange.starts_turn(), + subagent_id: exchange.subagent_id.clone(), + parent_subagent_id: exchange.parent_subagent_id.clone(), + turn_id, + turn, + } + } +} + +#[cfg(test)] +mod tests { + use chrono::TimeDelta; + use serde_json::json; + use uuid::Uuid; + + use std::collections::HashSet; + + use super::{ + AttachedBy, BaseRule, EventFacts, SegmentKind, SessionStreamsResponse, SessionStructure, + StreamEvent, event_segments, segment_run, session_streams, + }; + use crate::usage::streams::fixtures::{Session, assistant, call, item, result, user}; + + fn sorted_streams(session: &Session) -> SessionStreamsResponse { + let mut events = session.events.clone(); + events.sort_by_key(|event| event.observed_at); + session_streams(Uuid::nil(), events) + } + + /// `(head, subagent_id, attached_to, attached_by)` of one segment. + type Attachment<'a> = ( + &'a str, + Option<&'a str>, + Option<&'a str>, + Option, + ); + + fn attachments(streams: &SessionStreamsResponse) -> Vec> { + streams + .segments + .iter() + .map(|segment| { + ( + segment.head.as_str(), + segment.subagent_id.as_deref(), + segment.attached_to.as_deref(), + segment.attached_by, + ) + }) + .collect() + } + + #[test] + fn sub_agent_segments_are_grouped_by_their_declaration_and_never_attached() { + let mut session = Session::new(); + let main = (None, None, Some("p1")); + let a1 = (Some("a1"), None, Some("p1")); + session.declared( + "ask", + 0, + &json!([user("plan")]), + &json!([call("task")]), + main, + ); + session.declared( + "agent", + 5, + &json!([user("explore")]), + &json!([call("spawn")]), + a1, + ); + session.declared( + "agent-back", + 8, + &json!([user("explore"), call("spawn"), result("spawn")]), + &json!([assistant("explored")]), + a1, + ); + session.declared( + "back", + 30, + &json!([user("plan"), call("task"), result("task")]), + &json!([assistant("planned")]), + main, + ); + let streams = sorted_streams(&session); + assert_eq!( + attachments(&streams), + [("ask", None, None, None), ("agent", Some("a1"), None, None),] + ); + } + + #[test] + fn a_group_s_main_line_is_one_segment_and_each_side_exchange_its_own() { + let mut session = Session::new(); + let asked = json!([user("fix it")]); + let done = json!([user("fix it"), call("a"), result("a")]); + let answered = json!([ + user("fix it"), + call("a"), + result("a"), + item(assistant("done"), "msg_1"), + user("next") + ]); + let mut branch = answered.as_array().cloned().expect("an entry list"); + branch.extend([call("b"), result("b")]); + session.exchange("e1", 1, &asked, &json!([call("a")])); + session.exchange("e2", 2, &done, &json!([item(assistant("done"), "msg_1")])); + session.exchange("e3", 3, &answered, &json!([call("b")])); + // e3 is continued twice; only e4 is continued in turn, so e5, a + // recap whose answer nobody takes up, is a side exchange. + session.exchange("e4", 4, &json!(branch), &json!([call("c")])); + session.exchange("e5", 5, &json!(branch), &json!([assistant("two")])); + session.exchange( + "e6", + 6, + &json!([user("unrelated")]), + &json!([assistant("ok")]), + ); + let mut onward = branch.clone(); + onward.extend([call("c"), result("c")]); + session.exchange("e7", 7, &json!(onward), &json!([assistant("done")])); + + let streams = session_streams(Uuid::nil(), session.events); + let segments = streams + .segments + .iter() + .map(|segment| { + ( + segment.head.as_str(), + segment.kind, + segment.parent.as_deref(), + segment.forked_from.as_deref(), + segment.exchange_count, + ) + }) + .collect::>(); + assert_eq!( + segments, + [ + // Turn 1, then turn 2 opened by "next": e3, e4, e7. + ("e1", SegmentKind::Root, None, None, 2), + ("e3", SegmentKind::Turn, Some("e1"), Some("e2"), 3), + ("e5", SegmentKind::Branch, Some("e3"), Some("e3"), 1), + // A line of its own: nothing in it is continued. + ("e6", SegmentKind::Root, None, None, 1), + ] + ); + // A segment spans from its first request to its last response. + let first = &streams.segments[0]; + assert_eq!( + first.end.signed_duration_since(first.start), + TimeDelta::milliseconds(1_500) + ); + } + #[test] + fn a_segment_run_goes_from_its_exchange_to_the_segment_end() { + let mut session = Session::new(); + session.exchange("e1", 1, &json!([user("go")]), &json!([call("a")])); + session.exchange( + "e2", + 2, + &json!([user("go"), call("a"), result("a")]), + &json!([call("b")]), + ); + session.exchange( + "e3", + 3, + &json!([user("go"), call("a"), result("a"), call("b"), result("b")]), + &json!([assistant("done")]), + ); + let run = |from: &str| { + segment_run(session.events.clone(), from).map(|run| { + let ids = run + .exchanges + .iter() + .map(|exchange| exchange.id.as_str()) + .collect::>() + .join(","); + (run.segment, ids) + }) + }; + assert_eq!(run("e1"), Some(("e1".to_owned(), "e1,e2,e3".to_owned()))); + assert_eq!( + run("e2"), + Some(("e1".to_owned(), "e2,e3".to_owned())), + "an exchange inside a segment runs from itself and names the head" + ); + assert_eq!(run("zz"), None); + + let segment = segment_run(session.events.clone(), "e1").expect("a run"); + let e2 = &segment.exchanges[1]; + assert_eq!( + ( + e2.base.as_deref(), + e2.base_segment.as_deref(), + e2.rule, + e2.new_from + ), + ( + Some("e1"), + Some("e1"), + Some(BaseRule::LastOutput), + Some(2_i32) + ) + ); + assert!(e2.request_pk.is_some() && e2.response_pk.is_some()); + assert!(segment.exchanges[0].starts_turn && !e2.starts_turn); + } + + #[test] + fn every_event_of_an_exchange_maps_to_its_segment_head() { + let mut session = Session::new(); + session.exchange("e1", 1, &json!([user("go")]), &json!([call("a")])); + session.exchange( + "e2", + 2, + &json!([user("go"), call("a"), result("a")]), + &json!([assistant("done")]), + ); + session.exchange( + "e3", + 3, + &json!([user("unrelated")]), + &json!([assistant("ok")]), + ); + // A second request maps to e3 even though its content is not selected. + session.ask("e3", 4, None, &json!([user("again")])); + let segments = event_segments(session.events.clone()); + let heads = session + .events + .iter() + .map(|event| segments.get(&event.pk).map(String::as_str)) + .collect::>(); + assert_eq!( + heads, + [ + Some("e1"), + Some("e1"), + Some("e1"), + Some("e1"), + Some("e3"), + Some("e3"), + Some("e3") + ] + ); + } + + /// `(exchange, turn_id, turn)` of every exchange in session order. + fn exchange_turns(events: Vec) -> Vec<(String, Option, Option)> { + let structure = SessionStructure::derive(events); + (0..structure.exchanges.len()) + .map(|index| { + let exchange = structure.exchange(index); + (exchange.id, exchange.turn_id, exchange.turn) + }) + .collect() + } + + /// A session's turn count: its distinct declared turns when any event + /// declares one, else the requests that open a turn; read from events + /// alone. + fn counted_turns(events: &[StreamEvent]) -> usize { + let declared = events + .iter() + .filter_map(|event| event.turn_id.as_deref()) + .collect::>(); + if declared.is_empty() { + // The SQL formula: a turn-opening request that holds model output + // or names the response it continues, or the first turn-opening + // request of the session. + let mut opened = false; + events + .iter() + .filter(|event| match &event.facts { + EventFacts::Request(facts) if facts.starts_turn => { + let continues = + facts.new_from > 0_i32 || event.previous_response_id.is_some(); + let opens = continues || !opened; + opened = true; + opens + } + _ => false, + }) + .count() + } else { + declared.len() + } + } + + fn turn(id: &str, number: i64) -> (Option, Option) { + (Some(id.to_owned()), Some(number)) + } + + /// A Claude Code session: the main line, its sub-agent, and its side + /// requests declare the prompt they work for. + fn declared_session() -> Session { + let mut session = Session::new(); + let p1 = (None, None, Some("p1")); + // An auxiliary request before the first prompt. + session.exchange("pre", 0, &json!([user("quota")]), &json!([call("q")])); + session.declared( + "m1", + 1, + &json!([user("quota"), call("q"), result("q"), user("fix it")]), + &json!([call("c1")]), + p1, + ); + session.declared( + "sub", + 2, + &json!([user("explore")]), + &json!([assistant("found")]), + (Some("a1"), None, Some("p1")), + ); + // A title request declares nothing and has no base. + session.exchange( + "side", + 3, + &json!([user("name this")]), + &json!([assistant("Fix")]), + ); + session.declared( + "m2", + 4, + &json!([user("fix it"), call("c1"), result("c1")]), + &json!([call("c2")]), + p1, + ); + // The client resumes after hitting its output limit: user text + // without tool results, under the same prompt. + session.declared( + "resume", + 5, + &json!([ + user("fix it"), + call("c2"), + result("c2"), + user("Resume directly") + ]), + &json!([call("c3")]), + p1, + ); + // The user interjects: a new prompt that still carries tool results. + session.declared( + "interject", + 6, + &json!([call("c3"), result("c3"), user("also add a test")]), + &json!([call("c4")]), + (None, None, Some("p2")), + ); + // A compaction continues the interjection without a declaration. + session.exchange( + "compact", + 7, + &json!([call("c4"), result("c4"), user("summarize")]), + &json!([assistant("summary")]), + ); + // A recap starts afresh without a declaration. + session.exchange( + "recap", + 8, + &json!([user("recap")]), + &json!([assistant("so far")]), + ); + session + } + + #[test] + fn declared_turns_hold_their_sub_agents_and_side_requests() { + let session = declared_session(); + let turns = exchange_turns(session.events.clone()); + let expected = [ + ("pre", (None, None)), + ("m1", turn("p1", 1)), + ("sub", turn("p1", 1)), + ("side", turn("p1", 1)), + ("m2", turn("p1", 1)), + ("resume", turn("p1", 1)), + ("interject", turn("p2", 2)), + ("compact", turn("p2", 2)), + ("recap", turn("p2", 2)), + ]; + assert_eq!( + turns, + expected + .iter() + .map(|(id, (turn_id, number))| ((*id).to_owned(), turn_id.clone(), *number)) + .collect::>() + ); + let structure = SessionStructure::derive(session.events.clone()); + assert_eq!(structure.summary().turns, 2); + assert_eq!(counted_turns(&session.events), 2); + } + + #[test] + fn a_segment_opens_where_the_turn_differs_from_the_base_turn() { + let session = declared_session(); + let streams = session_streams(Uuid::nil(), session.events); + let segments = streams + .segments + .iter() + .map(|segment| { + ( + segment.head.as_str(), + segment.kind, + segment.exchange_count, + segment.turn_id.as_deref(), + segment.turn, + ) + }) + .collect::>(); + assert_eq!( + segments, + [ + ("pre", SegmentKind::Root, 1, None, None), + // Joining p1 from an exchange of no turn opens a segment. + ("m1", SegmentKind::Turn, 3, Some("p1"), Some(1)), + ("sub", SegmentKind::Root, 1, Some("p1"), Some(1)), + ("side", SegmentKind::Root, 1, Some("p1"), Some(1)), + // The resumption opens a turn by its request but stays in p1; + // the interjection carries tool results but declares p2. + ("interject", SegmentKind::Turn, 2, Some("p2"), Some(2)), + ("recap", SegmentKind::Root, 1, Some("p2"), Some(2)), + ] + ); + } + + #[test] + fn without_declarations_requests_that_open_a_turn_number_the_turns() { + let mut session = Session::new(); + // A response without its request, before any turn. + session.answer("lonely", 0, &json!([assistant("hi")])); + session.exchange("e1", 1, &json!([user("fix it")]), &json!([call("a")])); + session.exchange( + "e2", + 2, + &json!([user("fix it"), call("a"), result("a")]), + &json!([call("b")]), + ); + // A side request with user words but no model output joins the turn + // in progress: a line's first request opens a turn only before any. + session.exchange( + "title", + 3, + &json!([user("name this")]), + &json!([assistant("Fix")]), + ); + // Tool results without a known call and without user words: no base, + // so it joins the turn in progress. + session.exchange("stray", 4, &json!([result("zz")]), &json!([assistant("?")])); + // Continuing the main line by its base keeps the main line's turn. + session.exchange( + "e3", + 5, + &json!([ + user("fix it"), + call("a"), + result("a"), + call("b"), + result("b") + ]), + &json!([item(assistant("done"), "msg_1")]), + ); + session.exchange( + "e4", + 6, + &json!([item(assistant("done"), "msg_1"), user("next")]), + &json!([assistant("ok")]), + ); + // Only what is new, with a pointer to the answer it follows: the + // context continues, so the user's words open a turn. + session.ask("e5", 7, Some("e4"), &json!([user("and then")])); + session.answer("e5", 7, &json!([assistant("done")])); + let turns = exchange_turns(session.events.clone()); + assert_eq!( + turns + .iter() + .map(|(id, turn_id, number)| (id.as_str(), turn_id.as_deref(), *number)) + .collect::>(), + [ + ("lonely", None, None), + ("e1", Some("e1"), Some(1)), + ("e2", Some("e1"), Some(1)), + ("title", Some("e1"), Some(1)), + ("stray", Some("e1"), Some(1)), + ("e3", Some("e1"), Some(1)), + ("e4", Some("e4"), Some(2)), + ("e5", Some("e5"), Some(3)), + ] + ); + let structure = SessionStructure::derive(session.events.clone()); + assert_eq!(structure.summary().turns, 3); + assert_eq!(counted_turns(&session.events), 3); + let location = structure + .event_position(session.events[1].pk) + .expect("e1 is located"); + assert_eq!((location.segment.as_str(), location.turn), ("e1", Some(1))); + } + + #[test] + fn the_turn_count_matches_the_formula_on_every_fixture() { + let mut plain = Session::new(); + plain.exchange("e1", 1, &json!([user("go")]), &json!([call("a")])); + plain.exchange( + "e2", + 2, + &json!([user("go"), call("a"), result("a"), user("more")]), + &json!([assistant("ok")]), + ); + let mut events = declared_session().events; + events.sort_by_key(|event| event.observed_at); + for events in [events, plain.events] { + let structure = SessionStructure::derive(events.clone()); + let derived = structure + .segments + .iter() + .filter_map(|segment| segment.turn) + .max() + .and_then(|turn| usize::try_from(turn).ok()) + .unwrap_or_default(); + assert_eq!(structure.summary().turns, counted_turns(&events)); + assert_eq!(structure.summary().turns, derived); + } + } +} diff --git a/crates/abyss-backend/src/usage/streams/base.rs b/crates/abyss-backend/src/usage/streams/base.rs new file mode 100644 index 0000000..3ed3479 --- /dev/null +++ b/crates/abyss-backend/src/usage/streams/base.rs @@ -0,0 +1,674 @@ +//! Exchanges of one session and the base of each: the exchange whose output a +//! request's context continues, found by key lookups without comparing +//! content. + +use std::collections::HashMap; + +use chrono::{DateTime, Utc}; +use serde::Serialize; +use uuid::Uuid; + +use super::{ + EventFacts, StreamEvent, + facts::{MODEL_ROLES, RequestFacts, ResponseFacts, SkeletonEntry}, +}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +/// Rule that found a request's base. +pub enum BaseRule { + /// The request names a response of the session in `previous_response_id`. + Pointer, + /// The exchange that emitted the request's last identity said every + /// model-side entry that trails it. + LastOutput, + /// The request has no model output, and one exchange emitted every call + /// its tool results answer. + ResultsOnly, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +/// Why a request has no base. +pub enum NoBaseReason { + /// The last model output has no identity, or model output its emitter + /// did not say trails the last identity. + TextOnly, + /// The request has neither model output nor tool results. + NoOutput, + /// No other exchange of the session emitted the ids. + UnknownIds, + /// The ids come from more than one exchange. + MultipleResponses, + /// The exchange has no request event. + MissingRequest, +} + +/// One side of an exchange as streams use it. +pub(super) struct Side { + pub(super) pk: Uuid, + pub(super) previous_response_id: Option, + pub(super) facts: F, +} + +/// `subagent_id`, `parent_subagent_id`, and `turn_id` of one event. +type DeclaredIds = (Option, Option, Option); + +pub(super) struct Exchange { + pub(super) id: String, + pub(super) response_id: Option, + /// Every event grouped into the exchange, including duplicate sides. + /// Consumed when the session structure builds its event lookup. + pub(super) event_pks: Vec, + /// Request time, else response time. + pub(super) at: DateTime, + /// Latest event time. + pub(super) end: DateTime, + pub(super) request: Option>, + pub(super) response: Option>, + /// Ids the client declared: the request's, else the response's. + pub(super) subagent_id: Option, + pub(super) parent_subagent_id: Option, + pub(super) turn_id: Option, +} + +impl Exchange { + /// Adds `event` unless the exchange already has an event of its side. + fn add(&mut self, event: StreamEvent) { + let (pk, previous_response_id) = (event.pk, event.previous_response_id); + self.event_pks.push(pk); + self.end = self.end.max(event.observed_at); + let declared = (event.subagent_id, event.parent_subagent_id, event.turn_id); + match event.facts { + EventFacts::Request(facts) if self.request.is_none() => { + self.at = event.observed_at; + self.declare(declared); + self.request = Some(Side { + pk, + previous_response_id, + facts, + }); + } + EventFacts::Response(facts) if self.response.is_none() => { + if self.request.is_none() { + self.declare(declared); + } + self.response = Some(Side { + pk, + previous_response_id, + facts, + }); + } + EventFacts::Request(_) | EventFacts::Response(_) => {} + } + } + + fn declare(&mut self, (subagent, parent, turn): DeclaredIds) { + self.subagent_id = subagent; + self.parent_subagent_id = parent; + self.turn_id = turn; + } + + pub(super) fn starts_turn(&self) -> bool { + self.request + .as_ref() + .is_some_and(|request| request.facts.starts_turn) + } + + /// Whether the request continues a context: it holds model output (its + /// new part does not begin at its first entry), or it names the response + /// it continues, as a request that sends only what is new does. + pub(super) fn continues_context(&self) -> bool { + self.request.as_ref().is_some_and(|request| { + request.facts.new_from > 0_i32 || request.previous_response_id.is_some() + }) + } +} + +/// Pairs events into exchanges by `response_id`, ordered by request time; +/// an event without a `response_id` is its own exchange. +pub(super) fn group_exchanges(events: Vec) -> Vec { + let mut exchanges: Vec = Vec::new(); + let mut by_response_id: HashMap = HashMap::new(); + for event in events { + let known = event + .response_id + .as_ref() + .and_then(|response_id| by_response_id.get(response_id).copied()); + let index = known.unwrap_or_else(|| { + let id = event + .response_id + .clone() + .unwrap_or_else(|| event.event_id.clone()); + if let Some(response_id) = &event.response_id { + by_response_id.insert(response_id.clone(), exchanges.len()); + } + exchanges.push(Exchange { + id, + response_id: event.response_id.clone(), + event_pks: Vec::new(), + at: event.observed_at, + end: event.observed_at, + request: None, + response: None, + subagent_id: None, + parent_subagent_id: None, + turn_id: None, + }); + exchanges.len().saturating_sub(1) + }); + exchanges[index].add(event); + } + exchanges.sort_by_key(|exchange| exchange.at); + exchanges +} + +/// Exchanges whose response emitted each identity, in exchange order. +pub(super) fn emitters(exchanges: &[Exchange]) -> HashMap<&str, Vec> { + let mut by_id: HashMap<&str, Vec> = HashMap::new(); + for (index, exchange) in exchanges.iter().enumerate() { + let Some(response) = &exchange.response else { + continue; + }; + for id in &response.facts.emitted_ids { + let list = by_id.entry(id.as_str()).or_default(); + if list.last() != Some(&index) { + list.push(index); + } + } + } + by_id +} + +/// The one exchange other than `own` that emitted any of `ids`. +fn only_emitter<'a>( + ids: impl IntoIterator, + own: usize, + by_id: &HashMap<&str, Vec>, +) -> Result { + let mut found = None; + for id in ids { + for &emitter in by_id.get(id).into_iter().flatten() { + if emitter == own || found == Some(emitter) { + continue; + } + if found.is_some() { + return Err(NoBaseReason::MultipleResponses); + } + found = Some(emitter); + } + } + found.ok_or(NoBaseReason::UnknownIds) +} + +/// Whether `output` said every trailing role after its entry keyed by +/// `last`, in order: a client may drop entries of an output, but entries it +/// did not say come from a later reply. +fn said_after(output: &[SkeletonEntry], last: &str, trailing: &[String]) -> bool { + let is_model_side = |role: &str| MODEL_ROLES.contains(&role); + let Some(held) = output + .iter() + .rposition(|(role, key)| is_model_side(role) && key.as_deref() == Some(last)) + else { + return false; + }; + let mut said = output[held.saturating_add(1)..] + .iter() + .map(|(role, _)| role) + .filter(|role| is_model_side(role)); + trailing + .iter() + .all(|role| said.any(|spoken| spoken == role)) +} + +/// The base of every exchange and the rule that found it, or why it has +/// none: first each by [`find_base`]; then a request whose last identity +/// more than one other exchange emitted takes the emitter that saw what the +/// request has before that identity ([`tell_apart`]). +pub(super) fn find_bases( + exchanges: &[Exchange], + by_response_id: &HashMap<&str, usize>, + by_id: &HashMap<&str, Vec>, +) -> Vec> { + let first = (0..exchanges.len()) + .map(|index| find_base(index, exchanges, by_response_id, by_id)) + .collect::>(); + first + .iter() + .enumerate() + .map(|(index, &found)| match found { + Err(NoBaseReason::MultipleResponses) => { + tell_apart(index, exchanges, by_id, &first).unwrap_or(found) + } + found => found, + }) + .collect() +} + +/// The base of exchange `index` and the rule that found it: the response +/// its pointer names; else the one other exchange that emitted its last +/// identity, when that exchange said everything trailing it; else, without +/// model output, the one exchange that emitted the calls its tool results +/// answer. +fn find_base( + index: usize, + exchanges: &[Exchange], + by_response_id: &HashMap<&str, usize>, + by_id: &HashMap<&str, Vec>, +) -> Result<(usize, BaseRule), NoBaseReason> { + let request = exchanges[index] + .request + .as_ref() + .ok_or(NoBaseReason::MissingRequest)?; + if let Some(&base) = request + .previous_response_id + .as_deref() + .and_then(|response_id| by_response_id.get(response_id)) + { + return Ok((base, BaseRule::Pointer)); + } + let facts = &request.facts; + if let Some(last) = facts.last_identity.as_deref() { + let emitter = only_emitter([last], index, by_id)?; + return continued(emitter, exchanges, last, &facts.trailing); + } + // The last run has model output without any identity. + if facts.new_from > 0_i32 { + return Err(NoBaseReason::TextOnly); + } + if facts.result_ids.is_empty() { + return Err(NoBaseReason::NoOutput); + } + only_emitter(facts.result_ids.iter().map(String::as_str), index, by_id) + .map(|base| (base, BaseRule::ResultsOnly)) +} + +/// `emitter` as the base of a request whose last identity it emitted, when +/// it said everything trailing that identity. +fn continued( + emitter: usize, + exchanges: &[Exchange], + last: &str, + trailing: &[String], +) -> Result<(usize, BaseRule), NoBaseReason> { + let output = exchanges[emitter] + .response + .as_ref() + .map(|response| response.facts.output_skeleton.as_slice()) + .unwrap_or_default(); + if said_after(output, last, trailing) { + Ok((emitter, BaseRule::LastOutput)) + } else { + Err(NoBaseReason::TextOnly) + } +} + +/// For request `index`, whose last identity more than one other exchange +/// emitted: the one emitter that, before saying it, saw what the request +/// has before it. The context is resent as it grew, so the identity before +/// the last one either came just before the emitter spoke or was said by +/// the emitter itself. When one other exchange P0 emitted that identity, +/// the emitter is P0 or an exchange whose base by [`find_base`] is P0; +/// when nothing before the last identity has one, the emitter started +/// without any context and said it first. None emits it: the client +/// rewrote the context, `unknown_ids`; more than one: `multiple_responses`. +/// `None` when the identity before the last one has no single emitter. +fn tell_apart( + index: usize, + exchanges: &[Exchange], + by_id: &HashMap<&str, Vec>, + first: &[Result<(usize, BaseRule), NoBaseReason>], +) -> Option> { + let facts = &exchanges[index].request.as_ref()?.facts; + let last = facts.last_identity.as_deref()?; + let before = match facts.previous_identity.as_deref() { + Some(previous) => Some(only_emitter([previous], index, by_id).ok()?), + None => None, + }; + // An emitter that had no context and spoke the identity first. + let said_first_without_context = |emitter: usize| { + first[emitter] == Err(NoBaseReason::NoOutput) + && exchanges[emitter] + .request + .as_ref() + .is_some_and(|request| request.previous_response_id.is_none()) + && exchanges[emitter] + .response + .as_ref() + .and_then(|response| response.facts.emitted_ids.first()) + .is_some_and(|said| said == last) + }; + let saw = |emitter: usize| { + before.map_or_else( + || said_first_without_context(emitter), + |before| emitter == before || matches!(first[emitter], Ok((base, _)) if base == before), + ) + }; + let mut kept = by_id + .get(last) + .into_iter() + .flatten() + .copied() + .filter(|&emitter| emitter != index && saw(emitter)); + Some(match (kept.next(), kept.next()) { + (Some(emitter), None) => continued(emitter, exchanges, last, &facts.trailing), + (None, _) => Err(NoBaseReason::UnknownIds), + (Some(_), Some(_)) => Err(NoBaseReason::MultipleResponses), + }) +} + +#[cfg(test)] +mod tests { + use serde_json::{Value, json}; + + use super::{BaseRule, NoBaseReason}; + use crate::usage::streams::fixtures::{ + Session, assistant, call, item, reasoning, result, system, thought, user, + }; + + fn linked(base: &str) -> (String, BaseRule) { + (base.to_owned(), BaseRule::LastOutput) + } + + /// The base of a request after replies p1, p2, … with `outputs`. + fn base_after(outputs: &[Value], request: &Value) -> Result<(String, BaseRule), NoBaseReason> { + let mut session = Session::new(); + let mut second = 0_i64; + for output in outputs { + second = second.saturating_add(1); + session.exchange(&format!("p{second}"), second, &json!([user("go")]), output); + } + session.ask("next", second.saturating_add(1), None, request); + session.base("next") + } + + #[test] + fn a_request_continues_the_emitter_of_its_last_identity_unless_more_trails_it() { + // A request's base is the exchange that emitted the last identity in + // its history, when exactly one exchange did and it said all the model + // output after that identity. Each case lists the outputs of the + // earlier replies p1, p2, …, then what the next request carries back. + let cases = [ + ( + "1: the call's result and new input do not count", + vec![json!([assistant("look"), call("c1")])], + json!([ + user("go"), + assistant("look"), + call("c1"), + result("c1"), + user("more") + ]), + Ok(linked("p1")), + ), + ( + "2: rewritten text", + vec![json!([assistant("check the log"), call("c1")])], + json!([ + user("go"), + assistant("see the log"), + call("c1"), + result("c1") + ]), + Ok(linked("p1")), + ), + ( + "3: a dropped entry", + vec![json!([assistant("look"), call("c1")])], + json!([user("go"), call("c1"), result("c1")]), + Ok(linked("p1")), + ), + ( + "4: dropped thinking", + vec![json!([reasoning(), call("c1")])], + json!([user("go"), call("c1"), result("c1")]), + Ok(linked("p1")), + ), + ( + "5: text no observed reply said", + vec![json!([call("c1")])], + json!([user("go"), call("c1"), result("c1"), assistant("fixed")]), + Err(NoBaseReason::UnknownIds), + ), + ( + "6: earlier calls of the run come from earlier replies", + vec![json!([call("a")]), json!([call("b")])], + json!([user("go"), call("a"), result("a"), call("b"), result("b")]), + Ok(linked("p2")), + ), + ( + "7: a dropped server tool result", + vec![json!([call("s1"), result("s1"), assistant("found")])], + json!([user("go"), call("s1"), assistant("found")]), + Ok(linked("p1")), + ), + ( + "8: text aligned on its item id", + vec![json!([assistant("look"), item(assistant("so"), "msg_2")])], + json!([user("go"), item(assistant("so"), "msg_2"), user("more")]), + Ok(linked("p1")), + ), + ( + "9: dropped text after the call", + vec![json!([call("c1"), assistant("done")])], + json!([user("go"), call("c1"), result("c1")]), + Ok(linked("p1")), + ), + ( + "blank thinking p1 did not say", + vec![json!([call("c1")])], + json!([user("go"), call("c1"), result("c1"), reasoning()]), + Err(NoBaseReason::TextOnly), + ), + ( + "more blank text than p1 said", + vec![json!([call("c1"), assistant(" ")])], + json!([user("go"), call("c1"), assistant(" "), assistant("\n")]), + Err(NoBaseReason::TextOnly), + ), + ]; + for (case, outputs, request, expected) in cases { + assert_eq!(base_after(&outputs, &request), expected, "{case}"); + } + } + + #[test] + fn model_text_without_ids_aligns_on_the_one_reply_that_said_it_exactly() { + let cases = [ + ( + "text a later reply said", + vec![json!([call("c1")]), json!([assistant("fixed")])], + json!([user("go"), call("c1"), result("c1"), assistant("fixed")]), + Ok(linked("p2")), + ), + ( + "blank thinking before the replayed answer", + vec![json!([call("c1"), assistant("done")])], + json!([ + user("go"), + call("c1"), + result("c1"), + reasoning(), + assistant("done") + ]), + Ok(linked("p1")), + ), + ( + "a rewritten answer", + vec![json!([assistant("done")])], + json!([user("go"), assistant("done."), user("more")]), + Err(NoBaseReason::UnknownIds), + ), + ( + "an answer two replies said, both without context", + vec![json!([assistant("ok")]), json!([assistant("ok")])], + json!([user("go"), assistant("ok"), user("more")]), + Err(NoBaseReason::MultipleResponses), + ), + ]; + for (case, outputs, request, expected) in cases { + assert_eq!(base_after(&outputs, &request), expected, "{case}"); + } + } + + #[test] + fn a_repeated_answer_is_told_apart_by_what_came_before_it() { + // /loop: every round runs the check and answers the same words. + let mut session = Session::new(); + let mut history = vec![user("check")]; + for (round, call_id) in ["1", "2"].into_iter().enumerate() { + let second = i64::try_from(round).expect("small").saturating_mul(10); + session.exchange( + &format!("run{call_id}"), + second.saturating_add(1), + &json!(history), + &json!([call(call_id)]), + ); + history.extend([call(call_id), result(call_id)]); + session.exchange( + &format!("said{call_id}"), + second.saturating_add(2), + &json!(history), + &json!([assistant("no change")]), + ); + history.extend([assistant("no change"), user("check")]); + } + session.ask("next", 30, None, &json!(history)); + // The second round began after the first answer, which a later + // round said as well. + assert_eq!(session.base("run2"), Ok(linked("said1"))); + assert_eq!(session.base("next"), Ok(linked("said2"))); + + // The same words after thinking only one reply signed. + let mut session = Session::new(); + for (second, signature) in [(1, "sig1"), (2, "sig2")] { + session.exchange( + signature, + second, + &json!([user("go")]), + &json!([item(reasoning(), signature), assistant("same")]), + ); + } + session.ask( + "next", + 3, + None, + &json!([ + user("go"), + item(reasoning(), "sig1"), + assistant("same"), + user("more") + ]), + ); + assert_eq!(session.base("next"), Ok(linked("sig1"))); + } + + #[test] + fn a_repeated_answer_with_nothing_before_it_was_said_without_context() { + let mut session = Session::new(); + session.exchange("fresh", 1, &json!([user("hi")]), &json!([assistant("ok")])); + session.exchange("look", 2, &json!([user("go")]), &json!([call("c1")])); + let looked = [user("go"), call("c1"), result("c1")]; + session.exchange("found", 3, &json!(looked), &json!([assistant("ok")])); + session.exchange("found_too", 4, &json!(looked), &json!([assistant("fine")])); + session.exchange("again", 5, &json!(looked), &json!([assistant("fine")])); + // Only the reply that spoke first without any context fits. + session.ask( + "next", + 6, + None, + &json!([user("hi"), assistant("ok"), user("more")]), + ); + // A summary followed by the last answer, as compaction rebuilds a + // context: no reply said it without context. + session.ask( + "compacted", + 7, + None, + &json!([user("summary"), assistant("fine"), user("more")]), + ); + assert_eq!(session.base("next"), Ok(linked("fresh"))); + assert_eq!(session.base("compacted"), Err(NoBaseReason::UnknownIds)); + } + + #[test] + fn pointers_and_tool_results_find_a_base_and_requests_without_one_say_why() { + let mut session = Session::new(); + session.exchange("p", 1, &json!([user("go")]), &json!([call("a"), call("b")])); + session.exchange("twin", 2, &json!([user("go")]), &json!([call("d")])); + session.exchange("other", 3, &json!([user("go")]), &json!([call("d")])); + session.ask("pointer", 4, Some("p"), &json!([user("next")])); + session.ask("results", 5, None, &json!([result("a"), result("b")])); + session.ask( + "blank_text", + 6, + None, + &json!([user("hi"), assistant(" "), user("so")]), + ); + session.ask("nothing", 7, None, &json!([user("hi")])); + session.ask( + "unknown", + 8, + None, + &json!([user("go"), call("zz"), result("zz")]), + ); + session.ask("ambiguous", 9, None, &json!([result("d")])); + session.answer("lonely", 10, &json!([assistant("hi")])); + for (name, expected) in [ + ("pointer", Ok(("p".to_owned(), BaseRule::Pointer))), + ("results", Ok(("p".to_owned(), BaseRule::ResultsOnly))), + ("blank_text", Err(NoBaseReason::TextOnly)), + ("nothing", Err(NoBaseReason::NoOutput)), + ("unknown", Err(NoBaseReason::UnknownIds)), + ("ambiguous", Err(NoBaseReason::MultipleResponses)), + ("lonely", Err(NoBaseReason::MissingRequest)), + ] { + assert_eq!(session.base(name), expected, "{name}"); + } + } + + #[test] + fn a_chat_completions_turn_continues_the_exchange_that_gave_its_answer() { + // OpenCode on DeepSeek: thinking and answers come back without ids, + // as `reasoning_content` and `content`. + let mut session = Session::new(); + let ask = [system("rules"), user("fix it")]; + session.exchange( + "look", + 1, + &json!(ask), + &json!([thought("read the log"), call("c1")]), + ); + let mut resent = ask.to_vec(); + resent.extend([thought("read the log"), call("c1"), result("c1")]); + session.exchange( + "answer", + 2, + &json!(resent), + &json!([thought("it is fixed"), assistant("Fixed the build.")]), + ); + // The next turn resends the answer; OpenCode drops the thinking. + let mut next = resent.clone(); + next.extend([assistant("Fixed the build."), user("now test it")]); + session.exchange("next", 3, &json!(next), &json!([call("c2")])); + // OpenCode puts blank thinking before an answer that had none. + let mut replay = next.clone(); + replay.extend([call("c2"), result("c2")]); + session.exchange( + "tested", + 4, + &json!(replay), + &json!([assistant("Tests pass.")]), + ); + replay.extend([thought(" "), assistant("Tests pass."), user("thanks")]); + session.ask("thanks", 5, None, &json!(replay)); + for (name, base) in [ + ("answer", "look"), + ("next", "answer"), + ("tested", "next"), + ("thanks", "tested"), + ] { + assert_eq!(session.base(name), Ok(linked(base)), "{name}"); + } + } +} diff --git a/crates/abyss-backend/src/usage/streams/facts.rs b/crates/abyss-backend/src/usage/streams/facts.rs new file mode 100644 index 0000000..62ed97a --- /dev/null +++ b/crates/abyss-backend/src/usage/streams/facts.rs @@ -0,0 +1,551 @@ +//! The stream facts ingest reads from one event's conversation entries alone: +//! roles, call ids, item ids, digests of model text without ids, and whether a +//! user entry has text. + +use serde_json::Value; +use sha2::{Digest as _, Sha256}; + +use crate::usage::conversation::{Conversation, Entry}; + +/// Roles of the entries the model produced. +pub(super) const MODEL_ROLES: [&str; 3] = ["assistant", "reasoning", "tool_call"]; + +/// Roles of the model-side entries whose text can be their identity. +const TEXT_ROLES: [&str; 2] = ["assistant", "reasoning"]; + +/// Prefix that keeps item ids apart from call ids among entry identities. +const ITEM_PREFIX: &str = "item:"; + +/// Prefix that keeps text identities apart from call ids and item ids. +const TEXT_PREFIX: &str = "text:"; + +impl Entry { + fn is_model_side(&self) -> bool { + MODEL_ROLES.contains(&self.role.as_str()) + } + + /// The item id of a model-side entry, with [`ITEM_PREFIX`]. + fn item(&self) -> Option { + self.item_id + .as_deref() + .filter(|_| self.is_model_side()) + .map(|item_id| format!("{ITEM_PREFIX}{item_id}")) + } + + /// The identity of an `assistant` or `reasoning` entry that has neither + /// a call id nor an item id, read from its text when that text is not + /// blank: `text::` and the lowercase hex SHA-256 of the exact text. + fn text_identity(&self) -> Option { + let speaks = TEXT_ROLES.contains(&self.role.as_str()); + (speaks && self.call_id.is_none() && self.item_id.is_none() && !self.text.trim().is_empty()) + .then(|| { + let digest = hex::encode(Sha256::digest(self.text.as_bytes())); + format!("{TEXT_PREFIX}{}:{digest}", self.role) + }) + } + + /// The ids of the entry: its item, then the call id of a `tool_call` + /// entry; else, for model text without either, its text identity. + fn identities(&self) -> impl Iterator { + let call = self.call_id.clone().filter(|_| self.role == "tool_call"); + self.item() + .into_iter() + .chain(call) + .chain(self.text_identity()) + } + + /// The entry's key in a skeleton: its call id, else its item, else its + /// text identity. + fn key(&self) -> Option { + self.call_id + .clone() + .or_else(|| self.item()) + .or_else(|| self.text_identity()) + } +} + +/// The entries of the event's conversation; none when it is not completely +/// decoded. +fn conversation_entries(conversation: Option<&Value>) -> Vec { + Conversation::read(conversation) + .map(|conversation| conversation.entries) + .unwrap_or_default() +} + +/// One skeleton entry: an entry's role and key, its call id, else its +/// `item:`-prefixed item id, else its `text:`-prefixed text identity. +pub type SkeletonEntry = (String, Option); + +/// The `(role, key)` of each entry in order. +fn skeleton(entries: &[Entry]) -> Vec { + entries + .iter() + .map(|entry| (entry.role.clone(), entry.key())) + .collect() +} + +/// Call ids of the `tool_result` entries, in order. +fn result_ids(entries: &[Entry]) -> Vec { + entries + .iter() + .filter(|entry| entry.role == "tool_result") + .filter_map(|entry| entry.call_id.clone()) + .collect() +} + +/// Identities of the entries, in order. +fn identities(entries: &[Entry]) -> Vec { + entries.iter().flat_map(Entry::identities).collect() +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +/// Facts of one request event, read from its conversation entries alone. +pub struct RequestFacts { + /// Index after the last model-side entry, or 0 without one; the entries + /// from here on are the request's new part. + pub new_from: i32, + /// The last identity in the last run, the run reaching from the last + /// model-side entry back through model-side and `tool_result` entries to + /// any other role or the start: the call id of a `tool_call` entry, else + /// the `item:`-prefixed item id of a model-side entry, else the + /// `text:`-prefixed text identity of an `assistant` or `reasoning` entry. + pub last_identity: Option, + /// The identity before `last_identity` among all of the request's + /// identities in order, repeats of `last_identity` skipped; none when + /// nothing before it has an identity or without `last_identity`. It tells + /// apart the emitters of an identity that more than one response emitted. + pub previous_identity: Option, + /// Roles of the model-side entries after the one holding + /// `last_identity`, in order; none of them has an identity. Empty without + /// `last_identity`. + pub trailing: Vec, + /// Without model-side entries, the call ids of every `tool_result` + /// entry; otherwise empty. + pub result_ids: Vec, + /// Whether the new part holds a `user` entry with text and no + /// `tool_result` entry. + pub starts_turn: bool, + /// The last `user` entry with text in the new part. + pub label: Option