Publish harness and TUI open-source
initial sync from the monorepo
This commit is contained in:
commit
c68e39f604
2734 changed files with 1437016 additions and 0 deletions
545
crates/codegen/xai-hunk-tracker/src/loc/mod.rs
Normal file
545
crates/codegen/xai-hunk-tracker/src/loc/mod.rs
Normal file
|
|
@ -0,0 +1,545 @@
|
|||
//! LOC (Lines of Code) tracking — hunk-level attribution records.
|
||||
//!
|
||||
//! This module provides:
|
||||
//! - [`HunkRecord`]: a serializable attribution record derived from a [`Hunk`].
|
||||
//! - [`HunkRecordWriter`] / [`JsonlHunkRecordWriter`]: append-only JSONL persistence.
|
||||
//! - [`run_loc_sink`]: an async task that consumes [`HunkEvent`]s and writes records.
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use crate::events::{HunkEvent, HunkRemovalReason};
|
||||
use crate::types::{Hunk, HunkId, HunkSource};
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Enums
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Who authored a change.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub enum AuthorType {
|
||||
Agent,
|
||||
Human,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for AuthorType {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Agent => f.write_str("agent"),
|
||||
Self::Human => f.write_str("human"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Mirror of [`HunkSource`] for serialization.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub enum SourceType {
|
||||
AgentEdit,
|
||||
ExternalEditOnAgentFile,
|
||||
External,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for SourceType {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::AgentEdit => f.write_str("agent_edit"),
|
||||
Self::ExternalEditOnAgentFile => f.write_str("external_edit_on_agent_file"),
|
||||
Self::External => f.write_str("external"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether a record represents a new hunk or an in-place update.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub enum EventType {
|
||||
/// A new hunk was created.
|
||||
Added,
|
||||
/// An existing hunk's content changed in place.
|
||||
Updated,
|
||||
/// A hunk was removed. `lines_added` / `lines_removed` are negated
|
||||
/// so that `SUM` zeroes out the hunk's accumulated contribution.
|
||||
Removed,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for EventType {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Added => f.write_str("added"),
|
||||
Self::Updated => f.write_str("updated"),
|
||||
Self::Removed => f.write_str("removed"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// HunkRecord
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// A single LOC attribution record derived from a [`Hunk`].
|
||||
///
|
||||
/// Each record captures who authored a hunk (agent vs human), along with
|
||||
/// enough context (session, file, line range) for downstream analytics.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct HunkRecord {
|
||||
/// Stable hunk identifier (UUID).
|
||||
pub hunk_id: HunkId,
|
||||
/// Absolute file path.
|
||||
pub file_path: PathBuf,
|
||||
/// Start line of the hunk in the new file (1-indexed).
|
||||
pub hunk_start: usize,
|
||||
/// End line of the hunk in the new file (inclusive).
|
||||
pub hunk_end: usize,
|
||||
/// Lines added. For [`EventType::Added`] this is the full count (≥ 0).
|
||||
/// For [`EventType::Updated`] this is the delta from the previous state
|
||||
/// and may be negative (hunk shrank).
|
||||
pub lines_added: i64,
|
||||
/// Lines removed. For [`EventType::Added`] this is the full count (≥ 0).
|
||||
/// For [`EventType::Updated`] this is the delta and may be negative.
|
||||
pub lines_removed: i64,
|
||||
/// Who authored this change. `None` for [`EventType::Removed`] records.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub author_type: Option<AuthorType>,
|
||||
/// For agent edits: the agent id. For human edits: the user id (if known).
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub author_id: Option<String>,
|
||||
/// Machine-level agent identifier.
|
||||
pub agent_id: String,
|
||||
/// Session that produced this hunk.
|
||||
pub session_id: String,
|
||||
/// When the hunk was first detected.
|
||||
pub timestamp: DateTime<Utc>,
|
||||
/// Prompt index for agent edits, `None` for human edits.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub prompt_index: Option<usize>,
|
||||
/// Which [`HunkSource`] variant produced this change. `None` for [`EventType::Removed`] records.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub source_type: Option<SourceType>,
|
||||
/// Whether this is a new hunk or an in-place update.
|
||||
pub event_type: EventType,
|
||||
/// Why the hunk was removed. Only set for [`EventType::Removed`] records.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub removal_reason: Option<HunkRemovalReason>,
|
||||
}
|
||||
|
||||
impl HunkRecord {
|
||||
/// Derive a [`HunkRecord`] from a [`Hunk`].
|
||||
///
|
||||
/// `agent_id` is the stable machine-level identifier.
|
||||
/// `user_id` is the authenticated user id (used for human-attributed hunks).
|
||||
/// `event_type` distinguishes new hunks from in-place updates.
|
||||
///
|
||||
/// `attribution_source` controls which [`HunkSource`] is used for author
|
||||
/// attribution. For `HunkAdded` events this is `hunk.source`. For
|
||||
/// `HunkContentChanged` events this should be the *trigger* source
|
||||
/// (the source of the edit that caused the change), not the hunk's
|
||||
/// preserved source, since source-preservation logic may have kept the
|
||||
/// original agent attribution even though a human made the edit.
|
||||
pub fn from_hunk(
|
||||
hunk: &Hunk,
|
||||
session_id: &str,
|
||||
agent_id: &str,
|
||||
user_id: Option<&str>,
|
||||
event_type: EventType,
|
||||
attribution_source: &HunkSource,
|
||||
) -> Self {
|
||||
let (author_type, author_id, prompt_index, source_type) = match *attribution_source {
|
||||
HunkSource::AgentEdit { prompt_index } => (
|
||||
AuthorType::Agent,
|
||||
Some(agent_id.to_owned()),
|
||||
Some(prompt_index),
|
||||
SourceType::AgentEdit,
|
||||
),
|
||||
HunkSource::ExternalEditOnAgentFile => (
|
||||
AuthorType::Human,
|
||||
user_id.map(str::to_owned),
|
||||
None,
|
||||
SourceType::ExternalEditOnAgentFile,
|
||||
),
|
||||
HunkSource::External => (
|
||||
AuthorType::Human,
|
||||
user_id.map(str::to_owned),
|
||||
None,
|
||||
SourceType::External,
|
||||
),
|
||||
};
|
||||
|
||||
// For pure deletions (new_count == 0) use old_start/old_count.
|
||||
let (start, count) = if hunk.line_info.new_count == 0 {
|
||||
(hunk.line_info.old_start, hunk.line_info.old_count)
|
||||
} else {
|
||||
(hunk.line_info.new_start, hunk.line_info.new_count)
|
||||
};
|
||||
let end = if count == 0 { start } else { start + count - 1 };
|
||||
|
||||
Self {
|
||||
hunk_id: hunk.id.clone(),
|
||||
file_path: hunk.path.clone(),
|
||||
hunk_start: start,
|
||||
hunk_end: end,
|
||||
lines_added: hunk.line_info.new_count as i64,
|
||||
lines_removed: hunk.line_info.old_count as i64,
|
||||
author_type: Some(author_type),
|
||||
author_id,
|
||||
agent_id: agent_id.to_owned(),
|
||||
session_id: session_id.to_owned(),
|
||||
timestamp: hunk.created_at,
|
||||
prompt_index,
|
||||
source_type: Some(source_type),
|
||||
event_type,
|
||||
removal_reason: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// HunkRecordWriter
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Trait for persisting [`HunkRecord`]s.
|
||||
///
|
||||
/// Implementations may write to JSONL files, databases, etc.
|
||||
///
|
||||
/// All returned futures must be `Send` because `run_loc_sink` is spawned
|
||||
/// via `tokio::spawn` (which may run the task on any thread in the pool).
|
||||
pub trait HunkRecordWriter: Send {
|
||||
/// Write a single record. Errors are non-fatal; callers log and continue.
|
||||
fn write(
|
||||
&mut self,
|
||||
record: &HunkRecord,
|
||||
) -> impl std::future::Future<Output = std::io::Result<()>> + Send;
|
||||
|
||||
/// Flush any buffered data. Called during shutdown.
|
||||
fn flush(&mut self) -> impl std::future::Future<Output = std::io::Result<()>> + Send;
|
||||
}
|
||||
|
||||
/// Append-only JSONL writer for [`HunkRecord`]s.
|
||||
///
|
||||
/// The file is opened lazily on the first write so that sessions that produce
|
||||
/// no hunk events never create an empty file on disk.
|
||||
pub struct JsonlHunkRecordWriter {
|
||||
path: PathBuf,
|
||||
file: Option<tokio::fs::File>,
|
||||
}
|
||||
|
||||
impl JsonlHunkRecordWriter {
|
||||
/// Create a writer that will append to the given path.
|
||||
///
|
||||
/// The parent directory is created on the first write if it does not exist.
|
||||
pub fn new(path: PathBuf) -> Self {
|
||||
Self { path, file: None }
|
||||
}
|
||||
|
||||
/// Lazily open (or create) the file in append mode.
|
||||
async fn ensure_open(&mut self) -> std::io::Result<&mut tokio::fs::File> {
|
||||
if self.file.is_none() {
|
||||
if let Some(parent) = self.path.parent() {
|
||||
tokio::fs::create_dir_all(parent).await?;
|
||||
}
|
||||
let file = tokio::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&self.path)
|
||||
.await?;
|
||||
self.file = Some(file);
|
||||
}
|
||||
// The `if` block above guarantees `self.file` is `Some` at this point.
|
||||
Ok(self.file.as_mut().unwrap())
|
||||
}
|
||||
}
|
||||
|
||||
impl HunkRecordWriter for JsonlHunkRecordWriter {
|
||||
async fn write(&mut self, record: &HunkRecord) -> std::io::Result<()> {
|
||||
use tokio::io::AsyncWriteExt;
|
||||
|
||||
let file = self.ensure_open().await?;
|
||||
let mut line = serde_json::to_string(record)
|
||||
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
|
||||
line.push('\n');
|
||||
file.write_all(line.as_bytes()).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn flush(&mut self) -> std::io::Result<()> {
|
||||
use tokio::io::AsyncWriteExt;
|
||||
|
||||
if let Some(file) = self.file.as_mut() {
|
||||
file.flush().await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// LocAggregate (channel-based bridge to signals)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Lightweight aggregate update emitted by the LOC sink for consumption by
|
||||
/// an external bridge (e.g., the signals system in `xai-grok-shell`).
|
||||
///
|
||||
/// The sink sends one of these per processed `HunkEvent` that affects LOC.
|
||||
/// The bridge task translates them into `SignalEvent` variants.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum LocAggregate {
|
||||
/// Lines were added or changed (from HunkAdded or HunkContentChanged).
|
||||
LinesChanged {
|
||||
author_type: AuthorType,
|
||||
lines_added: i64,
|
||||
lines_removed: i64,
|
||||
file_path: PathBuf,
|
||||
},
|
||||
/// A hunk was reverted (rejected or superseded). The values are the
|
||||
/// accumulated totals that were zeroed out — always non-negative.
|
||||
LinesReverted {
|
||||
lines_added_reverted: i64,
|
||||
lines_removed_reverted: i64,
|
||||
},
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Sink configuration
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Context passed to the LOC sink at spawn time.
|
||||
pub struct LocSinkContext {
|
||||
/// Session identifier.
|
||||
pub session_id: String,
|
||||
/// Stable machine-level agent identifier.
|
||||
pub agent_id: String,
|
||||
/// Authenticated user id (if available). Used for human-attributed records.
|
||||
pub user_id: Option<String>,
|
||||
/// Optional channel for emitting LOC aggregates to an external consumer
|
||||
/// (e.g., the session signals system). When `None`, only JSONL is written.
|
||||
pub aggregate_tx: Option<mpsc::UnboundedSender<LocAggregate>>,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// run_loc_sink
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Consume [`HunkEvent`]s and write LOC attribution records.
|
||||
///
|
||||
/// This is the main entry point for the LOC tracking pipeline. It runs as a
|
||||
/// long-lived async task and should be spawned via `tokio::spawn`.
|
||||
///
|
||||
/// The sink maintains a `HashMap<HunkId, (i64, i64)>` tracking accumulated
|
||||
/// `(lines_added, lines_removed)` per hunk. When a `HunkRemoved` event
|
||||
/// arrives, the accumulated total is negated and written as a `Removed`
|
||||
/// record, zeroing out the hunk's contribution in SUM-based totals.
|
||||
///
|
||||
/// On cancellation the task drains any remaining events from the channel so
|
||||
/// that no in-flight records are lost.
|
||||
pub async fn run_loc_sink(
|
||||
mut event_rx: mpsc::UnboundedReceiver<HunkEvent>,
|
||||
mut writer: impl HunkRecordWriter,
|
||||
ctx: LocSinkContext,
|
||||
cancellation_token: tokio_util::sync::CancellationToken,
|
||||
) {
|
||||
// Accumulated (lines_added, lines_removed) per hunk_id.
|
||||
// Used to emit negating records when hunks are rejected/superseded.
|
||||
let mut acc: HashMap<HunkId, (i64, i64)> = HashMap::new();
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = cancellation_token.cancelled() => {
|
||||
tracing::debug!("LOC sink: cancellation received, draining remaining events");
|
||||
drain_remaining(&mut event_rx, &mut writer, &ctx, &mut acc).await;
|
||||
break;
|
||||
}
|
||||
event = event_rx.recv() => {
|
||||
let Some(event) = event else {
|
||||
// Channel closed — sender dropped.
|
||||
tracing::debug!("LOC sink: event channel closed");
|
||||
break;
|
||||
};
|
||||
handle_event(event, &mut writer, &ctx, &mut acc).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Err(e) = writer.flush().await {
|
||||
tracing::warn!(error = %e, "LOC sink: failed to flush writer on shutdown");
|
||||
}
|
||||
tracing::debug!("LOC sink: exiting");
|
||||
}
|
||||
|
||||
/// Process a single [`HunkEvent`].
|
||||
async fn handle_event(
|
||||
event: HunkEvent,
|
||||
writer: &mut impl HunkRecordWriter,
|
||||
ctx: &LocSinkContext,
|
||||
acc: &mut HashMap<HunkId, (i64, i64)>,
|
||||
) {
|
||||
match event {
|
||||
HunkEvent::HunkAdded { path: _, ref hunk } => {
|
||||
// For new hunks, the hunk's own source is the correct attribution.
|
||||
// lines_added/lines_removed are the full counts (no prior state).
|
||||
let record = HunkRecord::from_hunk(
|
||||
hunk,
|
||||
&ctx.session_id,
|
||||
&ctx.agent_id,
|
||||
ctx.user_id.as_deref(),
|
||||
EventType::Added,
|
||||
&hunk.source,
|
||||
);
|
||||
let entry = acc.entry(hunk.id.clone()).or_insert((0, 0));
|
||||
entry.0 += record.lines_added;
|
||||
entry.1 += record.lines_removed;
|
||||
// Emit aggregate for signals bridge
|
||||
if let Some(tx) = &ctx.aggregate_tx {
|
||||
let _ = tx.send(LocAggregate::LinesChanged {
|
||||
author_type: record.author_type.unwrap_or(AuthorType::Agent),
|
||||
lines_added: record.lines_added,
|
||||
lines_removed: record.lines_removed,
|
||||
file_path: record.file_path.clone(),
|
||||
});
|
||||
}
|
||||
write_record(&record, writer).await;
|
||||
}
|
||||
HunkEvent::HunkContentChanged {
|
||||
path: _,
|
||||
ref hunk,
|
||||
trigger_source,
|
||||
prev_lines_added,
|
||||
prev_lines_removed,
|
||||
} => {
|
||||
// For in-place changes, use the trigger source for attribution
|
||||
// and record only the delta (new - prev) so LOC totals can be
|
||||
// computed with a simple SUM grouped by author_type.
|
||||
let mut record = HunkRecord::from_hunk(
|
||||
hunk,
|
||||
&ctx.session_id,
|
||||
&ctx.agent_id,
|
||||
ctx.user_id.as_deref(),
|
||||
EventType::Updated,
|
||||
&trigger_source,
|
||||
);
|
||||
// Replace full counts with signed deltas so shrinking hunks
|
||||
// (e.g., human deletes 3 of 10 agent lines) produce negative
|
||||
// values that correctly reduce the total on SUM.
|
||||
record.lines_added = hunk.line_info.new_count as i64 - prev_lines_added as i64;
|
||||
record.lines_removed = hunk.line_info.old_count as i64 - prev_lines_removed as i64;
|
||||
let entry = acc.entry(hunk.id.clone()).or_insert((0, 0));
|
||||
entry.0 += record.lines_added;
|
||||
entry.1 += record.lines_removed;
|
||||
// Emit aggregate for signals bridge
|
||||
if let Some(tx) = &ctx.aggregate_tx {
|
||||
let _ = tx.send(LocAggregate::LinesChanged {
|
||||
author_type: record.author_type.unwrap_or(AuthorType::Human),
|
||||
lines_added: record.lines_added,
|
||||
lines_removed: record.lines_removed,
|
||||
file_path: record.file_path.clone(),
|
||||
});
|
||||
}
|
||||
write_record(&record, writer).await;
|
||||
}
|
||||
HunkEvent::HunkRemoved {
|
||||
path,
|
||||
hunk_id,
|
||||
reason,
|
||||
} => {
|
||||
match reason {
|
||||
HunkRemovalReason::Accepted => {
|
||||
// Accepted hunks keep their LOC contribution — just
|
||||
// clear the accumulated state without writing a
|
||||
// negating record.
|
||||
acc.remove(&hunk_id);
|
||||
}
|
||||
HunkRemovalReason::Rejected | HunkRemovalReason::Superseded => {
|
||||
// Rejected/superseded hunks lose their LOC — negate
|
||||
// the accumulated totals so SUM zeroes them out.
|
||||
if let Some((total_added, total_removed)) = acc.remove(&hunk_id)
|
||||
&& (total_added != 0 || total_removed != 0)
|
||||
{
|
||||
// Emit revert aggregate for signals bridge
|
||||
if let Some(tx) = &ctx.aggregate_tx {
|
||||
let _ = tx.send(LocAggregate::LinesReverted {
|
||||
lines_added_reverted: total_added.max(0),
|
||||
lines_removed_reverted: total_removed.max(0),
|
||||
});
|
||||
}
|
||||
let record = HunkRecord {
|
||||
hunk_id: hunk_id.clone(),
|
||||
file_path: path,
|
||||
hunk_start: 0,
|
||||
hunk_end: 0,
|
||||
lines_added: -total_added,
|
||||
lines_removed: -total_removed,
|
||||
author_type: None,
|
||||
author_id: None,
|
||||
agent_id: ctx.agent_id.clone(),
|
||||
session_id: ctx.session_id.clone(),
|
||||
timestamp: Utc::now(),
|
||||
prompt_index: None,
|
||||
source_type: None,
|
||||
event_type: EventType::Removed,
|
||||
removal_reason: Some(reason),
|
||||
};
|
||||
write_record(&record, writer).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
HunkEvent::HunkMoved { .. } => {
|
||||
tracing::trace!("LOC sink: ignoring HunkMoved event");
|
||||
}
|
||||
HunkEvent::FileAdded { .. } => {
|
||||
tracing::trace!("LOC sink: ignoring FileAdded event");
|
||||
}
|
||||
HunkEvent::FileRemoved { .. } => {
|
||||
tracing::trace!("LOC sink: ignoring FileRemoved event");
|
||||
}
|
||||
HunkEvent::BaselineUpdated { .. } => {
|
||||
tracing::trace!("LOC sink: ignoring BaselineUpdated event");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Write a [`HunkRecord`], logging on success/failure.
|
||||
async fn write_record(record: &HunkRecord, writer: &mut impl HunkRecordWriter) {
|
||||
if let Err(e) = writer.write(record).await {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
hunk_id = %record.hunk_id,
|
||||
file_path = %record.file_path.display(),
|
||||
"LOC sink: failed to write hunk record, dropping"
|
||||
);
|
||||
} else {
|
||||
tracing::debug!(
|
||||
hunk_id = %record.hunk_id,
|
||||
file_path = %record.file_path.display(),
|
||||
author_type = ?record.author_type,
|
||||
"LOC sink: wrote hunk record"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Drain remaining events after cancellation.
|
||||
async fn drain_remaining(
|
||||
event_rx: &mut mpsc::UnboundedReceiver<HunkEvent>,
|
||||
writer: &mut impl HunkRecordWriter,
|
||||
ctx: &LocSinkContext,
|
||||
acc: &mut HashMap<HunkId, (i64, i64)>,
|
||||
) {
|
||||
let mut count = 0usize;
|
||||
while let Ok(event) = event_rx.try_recv() {
|
||||
handle_event(event, writer, ctx, acc).await;
|
||||
count += 1;
|
||||
}
|
||||
if count > 0 {
|
||||
tracing::debug!(
|
||||
count,
|
||||
"LOC sink: drained remaining events after cancellation"
|
||||
);
|
||||
}
|
||||
}
|
||||
782
crates/codegen/xai-hunk-tracker/src/loc/tests.rs
Normal file
782
crates/codegen/xai-hunk-tracker/src/loc/tests.rs
Normal file
|
|
@ -0,0 +1,782 @@
|
|||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::Utc;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::events::HunkEvent;
|
||||
use crate::types::{Hunk, HunkId, HunkLineInfo, HunkSource};
|
||||
|
||||
use super::*;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn sample_agent_hunk() -> Hunk {
|
||||
Hunk {
|
||||
id: HunkId::from_string("test-hunk-001".into()),
|
||||
path: PathBuf::from("/tmp/foo.rs"),
|
||||
line_info: HunkLineInfo {
|
||||
old_start: 10,
|
||||
old_count: 3,
|
||||
new_start: 10,
|
||||
new_count: 5,
|
||||
},
|
||||
source: HunkSource::AgentEdit { prompt_index: 2 },
|
||||
old_text: Some("old\nlines\nhere".into()),
|
||||
new_text: "new\nlines\nhere\nplus\nmore".into(),
|
||||
patch: None,
|
||||
created_at: Utc::now(),
|
||||
selected: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_external_hunk() -> Hunk {
|
||||
Hunk {
|
||||
id: HunkId::from_string("test-hunk-002".into()),
|
||||
path: PathBuf::from("/tmp/bar.rs"),
|
||||
line_info: HunkLineInfo {
|
||||
old_start: 1,
|
||||
old_count: 0,
|
||||
new_start: 1,
|
||||
new_count: 4,
|
||||
},
|
||||
source: HunkSource::External,
|
||||
old_text: None,
|
||||
new_text: "line1\nline2\nline3\nline4".into(),
|
||||
patch: None,
|
||||
created_at: Utc::now(),
|
||||
selected: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_deletion_hunk() -> Hunk {
|
||||
Hunk {
|
||||
id: HunkId::from_string("test-hunk-003".into()),
|
||||
path: PathBuf::from("/tmp/del.rs"),
|
||||
line_info: HunkLineInfo {
|
||||
old_start: 5,
|
||||
old_count: 3,
|
||||
new_start: 0,
|
||||
new_count: 0,
|
||||
},
|
||||
source: HunkSource::AgentEdit { prompt_index: 1 },
|
||||
old_text: Some("deleted\nlines\nhere".into()),
|
||||
new_text: String::new(),
|
||||
patch: None,
|
||||
created_at: Utc::now(),
|
||||
selected: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// In-memory writer for testing.
|
||||
struct VecWriter {
|
||||
records: Vec<HunkRecord>,
|
||||
flush_count: usize,
|
||||
}
|
||||
|
||||
impl VecWriter {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
records: Vec::new(),
|
||||
flush_count: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HunkRecordWriter for VecWriter {
|
||||
async fn write(&mut self, record: &HunkRecord) -> std::io::Result<()> {
|
||||
self.records.push(record.clone());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn flush(&mut self) -> std::io::Result<()> {
|
||||
self.flush_count += 1;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Thread-safe wrapper around `VecWriter` for use with `run_loc_sink`
|
||||
/// (which takes ownership of the writer).
|
||||
struct SharedWriter(std::sync::Arc<std::sync::Mutex<VecWriter>>);
|
||||
|
||||
impl SharedWriter {
|
||||
fn new() -> (Self, std::sync::Arc<std::sync::Mutex<VecWriter>>) {
|
||||
let inner = std::sync::Arc::new(std::sync::Mutex::new(VecWriter::new()));
|
||||
(Self(inner.clone()), inner)
|
||||
}
|
||||
}
|
||||
|
||||
impl HunkRecordWriter for SharedWriter {
|
||||
async fn write(&mut self, record: &HunkRecord) -> std::io::Result<()> {
|
||||
// Do the work inside the lock synchronously — don't hold MutexGuard across .await
|
||||
self.0.lock().unwrap().records.push(record.clone());
|
||||
Ok(())
|
||||
}
|
||||
async fn flush(&mut self) -> std::io::Result<()> {
|
||||
self.0.lock().unwrap().flush_count += 1;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn make_ctx() -> LocSinkContext {
|
||||
LocSinkContext {
|
||||
session_id: "sess-001".into(),
|
||||
agent_id: "agent-abc".into(),
|
||||
user_id: Some("user-xyz".into()),
|
||||
aggregate_tx: None,
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Unit tests: HunkRecord::from_hunk
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn from_hunk_agent_edit() {
|
||||
let hunk = sample_agent_hunk();
|
||||
let record = HunkRecord::from_hunk(
|
||||
&hunk,
|
||||
"sess-1",
|
||||
"agent-1",
|
||||
Some("user-1"),
|
||||
EventType::Added,
|
||||
&hunk.source,
|
||||
);
|
||||
|
||||
assert_eq!(record.hunk_id, HunkId::from_string("test-hunk-001".into()));
|
||||
assert_eq!(record.file_path, PathBuf::from("/tmp/foo.rs"));
|
||||
assert_eq!(record.hunk_start, 10);
|
||||
assert_eq!(record.hunk_end, 14); // 10 + 5 - 1
|
||||
assert_eq!(record.lines_added, 5);
|
||||
assert_eq!(record.lines_removed, 3);
|
||||
assert_eq!(record.author_type, Some(AuthorType::Agent));
|
||||
assert_eq!(record.author_id, Some("agent-1".into()));
|
||||
assert_eq!(record.agent_id, "agent-1");
|
||||
assert_eq!(record.session_id, "sess-1");
|
||||
assert_eq!(record.prompt_index, Some(2));
|
||||
assert_eq!(record.source_type, Some(SourceType::AgentEdit));
|
||||
assert_eq!(record.event_type, EventType::Added);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_hunk_event_type_updated() {
|
||||
let hunk = sample_agent_hunk();
|
||||
let record = HunkRecord::from_hunk(
|
||||
&hunk,
|
||||
"sess-1",
|
||||
"agent-1",
|
||||
None,
|
||||
EventType::Updated,
|
||||
&hunk.source,
|
||||
);
|
||||
assert_eq!(record.event_type, EventType::Updated);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_hunk_external() {
|
||||
let hunk = sample_external_hunk();
|
||||
let record = HunkRecord::from_hunk(
|
||||
&hunk,
|
||||
"sess-1",
|
||||
"agent-1",
|
||||
Some("user-1"),
|
||||
EventType::Added,
|
||||
&hunk.source,
|
||||
);
|
||||
|
||||
assert_eq!(record.author_type, Some(AuthorType::Human));
|
||||
assert_eq!(record.author_id, Some("user-1".into()));
|
||||
assert_eq!(record.prompt_index, None);
|
||||
assert_eq!(record.source_type, Some(SourceType::External));
|
||||
assert_eq!(record.hunk_start, 1);
|
||||
assert_eq!(record.hunk_end, 4); // 1 + 4 - 1
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_hunk_external_no_user_id() {
|
||||
let hunk = sample_external_hunk();
|
||||
let record = HunkRecord::from_hunk(
|
||||
&hunk,
|
||||
"sess-1",
|
||||
"agent-1",
|
||||
None,
|
||||
EventType::Added,
|
||||
&hunk.source,
|
||||
);
|
||||
|
||||
assert_eq!(record.author_type, Some(AuthorType::Human));
|
||||
assert_eq!(record.author_id, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_hunk_external_edit_on_agent_file() {
|
||||
let mut hunk = sample_external_hunk();
|
||||
hunk.source = HunkSource::ExternalEditOnAgentFile;
|
||||
let record = HunkRecord::from_hunk(
|
||||
&hunk,
|
||||
"sess-1",
|
||||
"agent-1",
|
||||
Some("user-1"),
|
||||
EventType::Added,
|
||||
&hunk.source,
|
||||
);
|
||||
|
||||
assert_eq!(record.author_type, Some(AuthorType::Human));
|
||||
assert_eq!(
|
||||
record.source_type,
|
||||
Some(SourceType::ExternalEditOnAgentFile)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_hunk_pure_deletion() {
|
||||
let hunk = sample_deletion_hunk();
|
||||
let record = HunkRecord::from_hunk(
|
||||
&hunk,
|
||||
"sess-1",
|
||||
"agent-1",
|
||||
None,
|
||||
EventType::Added,
|
||||
&hunk.source,
|
||||
);
|
||||
|
||||
// Pure deletion: new_count == 0, so uses old_start/old_count
|
||||
assert_eq!(record.hunk_start, 5);
|
||||
assert_eq!(record.hunk_end, 7); // 5 + 3 - 1
|
||||
assert_eq!(record.lines_added, 0i64);
|
||||
assert_eq!(record.lines_removed, 3i64);
|
||||
}
|
||||
|
||||
/// Verify that attribution_source overrides the hunk's preserved source.
|
||||
#[test]
|
||||
fn from_hunk_trigger_source_overrides_preserved_source() {
|
||||
let hunk = sample_agent_hunk(); // hunk.source = AgentEdit
|
||||
let trigger = HunkSource::ExternalEditOnAgentFile;
|
||||
let record = HunkRecord::from_hunk(
|
||||
&hunk,
|
||||
"sess-1",
|
||||
"agent-1",
|
||||
Some("user-1"),
|
||||
EventType::Updated,
|
||||
&trigger,
|
||||
);
|
||||
|
||||
assert_eq!(record.author_type, Some(AuthorType::Human));
|
||||
assert_eq!(
|
||||
record.source_type,
|
||||
Some(SourceType::ExternalEditOnAgentFile)
|
||||
);
|
||||
assert_eq!(record.author_id, Some("user-1".into()));
|
||||
assert_eq!(record.prompt_index, None);
|
||||
assert_eq!(record.event_type, EventType::Updated);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Sink tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn sink_processes_added_and_content_changed() {
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
let ctx = make_ctx();
|
||||
let cancel = tokio_util::sync::CancellationToken::new();
|
||||
|
||||
let hunk = sample_agent_hunk();
|
||||
let mut updated_hunk = sample_agent_hunk();
|
||||
updated_hunk.line_info.new_count = 8; // grew from 5 to 8 lines
|
||||
|
||||
// Send a mix of events — only HunkAdded and HunkContentChanged should produce records
|
||||
tx.send(HunkEvent::FileAdded {
|
||||
path: PathBuf::from("/tmp/foo.rs"),
|
||||
is_agent_file: true,
|
||||
})
|
||||
.unwrap();
|
||||
tx.send(HunkEvent::HunkAdded {
|
||||
path: PathBuf::from("/tmp/foo.rs"),
|
||||
hunk: Arc::new(hunk),
|
||||
})
|
||||
.unwrap();
|
||||
tx.send(HunkEvent::HunkContentChanged {
|
||||
path: PathBuf::from("/tmp/foo.rs"),
|
||||
hunk: Arc::new(updated_hunk),
|
||||
trigger_source: HunkSource::AgentEdit { prompt_index: 2 },
|
||||
prev_lines_added: 5, // original hunk had 5 lines added
|
||||
prev_lines_removed: 3, // original hunk had 3 lines removed
|
||||
})
|
||||
.unwrap();
|
||||
tx.send(HunkEvent::HunkMoved {
|
||||
path: PathBuf::from("/tmp/foo.rs"),
|
||||
hunk_id: HunkId::from_string("test-hunk-001".into()),
|
||||
new_line_info: HunkLineInfo {
|
||||
old_start: 10,
|
||||
old_count: 3,
|
||||
new_start: 12,
|
||||
new_count: 5,
|
||||
},
|
||||
})
|
||||
.unwrap();
|
||||
tx.send(HunkEvent::HunkRemoved {
|
||||
path: PathBuf::from("/tmp/foo.rs"),
|
||||
hunk_id: HunkId::from_string("test-hunk-001".into()),
|
||||
reason: crate::events::HunkRemovalReason::Superseded,
|
||||
})
|
||||
.unwrap();
|
||||
tx.send(HunkEvent::FileRemoved {
|
||||
path: PathBuf::from("/tmp/foo.rs"),
|
||||
})
|
||||
.unwrap();
|
||||
tx.send(HunkEvent::BaselineUpdated {
|
||||
path: PathBuf::from("/tmp/foo.rs"),
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
// Drop sender to close the channel
|
||||
drop(tx);
|
||||
|
||||
let (shared_writer, shared) = SharedWriter::new();
|
||||
run_loc_sink(rx, shared_writer, ctx, cancel).await;
|
||||
|
||||
let w = shared.lock().unwrap();
|
||||
assert_eq!(
|
||||
w.records.len(),
|
||||
3,
|
||||
"HunkAdded + HunkContentChanged + HunkRemoved should produce 3 records"
|
||||
);
|
||||
|
||||
// First record: added (full counts)
|
||||
assert_eq!(
|
||||
w.records[0].hunk_id,
|
||||
HunkId::from_string("test-hunk-001".into())
|
||||
);
|
||||
assert_eq!(w.records[0].event_type, EventType::Added);
|
||||
assert_eq!(w.records[0].lines_added, 5);
|
||||
assert_eq!(w.records[0].lines_removed, 3);
|
||||
|
||||
// Second record: updated (delta: 8-5=3 added, 3-3=0 removed)
|
||||
assert_eq!(w.records[1].event_type, EventType::Updated);
|
||||
assert_eq!(w.records[1].lines_added, 3i64);
|
||||
assert_eq!(w.records[1].lines_removed, 0i64);
|
||||
|
||||
// Third record: removed (negates accumulated: -(5+3)=-8, -(3+0)=-3)
|
||||
assert_eq!(w.records[2].event_type, EventType::Removed);
|
||||
assert_eq!(w.records[2].lines_added, -8i64);
|
||||
assert_eq!(w.records[2].lines_removed, -3i64);
|
||||
|
||||
// SUM should be zero
|
||||
let total: i64 = w.records.iter().map(|r| r.lines_added).sum();
|
||||
assert_eq!(total, 0);
|
||||
}
|
||||
|
||||
/// When a hunk is removed, the sink must emit a negating record so that
|
||||
/// SUM-based totals zero out the hunk's contribution.
|
||||
#[tokio::test]
|
||||
async fn sink_removed_hunk_zeroes_out_accumulated_total() {
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
let ctx = make_ctx();
|
||||
let cancel = tokio_util::sync::CancellationToken::new();
|
||||
|
||||
let hunk = sample_agent_hunk(); // lines_added=5, lines_removed=3
|
||||
let hunk_id = hunk.id.clone();
|
||||
let path = hunk.path.clone();
|
||||
|
||||
// Add, then remove
|
||||
tx.send(HunkEvent::HunkAdded {
|
||||
path: path.clone(),
|
||||
hunk: Arc::new(hunk),
|
||||
})
|
||||
.unwrap();
|
||||
tx.send(HunkEvent::HunkRemoved {
|
||||
path: path.clone(),
|
||||
hunk_id: hunk_id.clone(),
|
||||
reason: crate::events::HunkRemovalReason::Rejected,
|
||||
})
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let (shared_writer, shared) = SharedWriter::new();
|
||||
run_loc_sink(rx, shared_writer, ctx, cancel).await;
|
||||
|
||||
let w = shared.lock().unwrap();
|
||||
assert_eq!(w.records.len(), 2, "Should have added + removed records");
|
||||
|
||||
// First: added
|
||||
assert_eq!(w.records[0].event_type, EventType::Added);
|
||||
assert_eq!(w.records[0].lines_added, 5);
|
||||
assert_eq!(w.records[0].lines_removed, 3);
|
||||
|
||||
// Second: removed (negated)
|
||||
assert_eq!(w.records[1].event_type, EventType::Removed);
|
||||
assert_eq!(w.records[1].lines_added, -5);
|
||||
assert_eq!(w.records[1].lines_removed, -3);
|
||||
|
||||
// SUM should be zero
|
||||
let total_added: i64 = w.records.iter().map(|r| r.lines_added).sum();
|
||||
let total_removed: i64 = w.records.iter().map(|r| r.lines_removed).sum();
|
||||
assert_eq!(total_added, 0, "Removed hunk should zero out lines_added");
|
||||
assert_eq!(
|
||||
total_removed, 0,
|
||||
"Removed hunk should zero out lines_removed"
|
||||
);
|
||||
}
|
||||
|
||||
/// Full scenario: agent adds, human expands, then hunk is removed.
|
||||
/// The negating record must cancel the entire accumulated total.
|
||||
#[tokio::test]
|
||||
async fn sink_removed_hunk_after_updates_zeroes_correctly() {
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
let ctx = make_ctx();
|
||||
let cancel = tokio_util::sync::CancellationToken::new();
|
||||
|
||||
let hunk = sample_agent_hunk(); // lines_added=5, lines_removed=3
|
||||
let hunk_id = hunk.id.clone();
|
||||
let path = hunk.path.clone();
|
||||
|
||||
let mut updated = sample_agent_hunk();
|
||||
updated.line_info.new_count = 8; // grew from 5 → 8
|
||||
|
||||
// Add → update → remove
|
||||
tx.send(HunkEvent::HunkAdded {
|
||||
path: path.clone(),
|
||||
hunk: Arc::new(hunk),
|
||||
})
|
||||
.unwrap();
|
||||
tx.send(HunkEvent::HunkContentChanged {
|
||||
path: path.clone(),
|
||||
hunk: Arc::new(updated),
|
||||
trigger_source: HunkSource::ExternalEditOnAgentFile,
|
||||
prev_lines_added: 5,
|
||||
prev_lines_removed: 3,
|
||||
})
|
||||
.unwrap();
|
||||
tx.send(HunkEvent::HunkRemoved {
|
||||
path: path.clone(),
|
||||
hunk_id: hunk_id.clone(),
|
||||
reason: crate::events::HunkRemovalReason::Superseded,
|
||||
})
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let (shared_writer, shared) = SharedWriter::new();
|
||||
run_loc_sink(rx, shared_writer, ctx, cancel).await;
|
||||
|
||||
let w = shared.lock().unwrap();
|
||||
assert_eq!(w.records.len(), 3, "added + updated + removed");
|
||||
|
||||
// SUM should be zero: the hunk was fully removed
|
||||
let total_added: i64 = w.records.iter().map(|r| r.lines_added).sum();
|
||||
let total_removed: i64 = w.records.iter().map(|r| r.lines_removed).sum();
|
||||
assert_eq!(total_added, 0);
|
||||
assert_eq!(total_removed, 0);
|
||||
}
|
||||
|
||||
/// Accepted hunks keep their LOC contribution — no negating record is written.
|
||||
#[tokio::test]
|
||||
async fn sink_accepted_hunk_preserves_loc() {
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
let ctx = make_ctx();
|
||||
let cancel = tokio_util::sync::CancellationToken::new();
|
||||
|
||||
let hunk = sample_agent_hunk(); // lines_added=5, lines_removed=3
|
||||
let hunk_id = hunk.id.clone();
|
||||
let path = hunk.path.clone();
|
||||
|
||||
tx.send(HunkEvent::HunkAdded {
|
||||
path: path.clone(),
|
||||
hunk: Arc::new(hunk),
|
||||
})
|
||||
.unwrap();
|
||||
tx.send(HunkEvent::HunkRemoved {
|
||||
path: path.clone(),
|
||||
hunk_id: hunk_id.clone(),
|
||||
reason: crate::events::HunkRemovalReason::Accepted,
|
||||
})
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let (shared_writer, shared) = SharedWriter::new();
|
||||
run_loc_sink(rx, shared_writer, ctx, cancel).await;
|
||||
|
||||
let w = shared.lock().unwrap();
|
||||
// Only the Added record — no Removed record for accepted hunks
|
||||
assert_eq!(
|
||||
w.records.len(),
|
||||
1,
|
||||
"Accepted hunk should NOT produce a Removed record"
|
||||
);
|
||||
assert_eq!(w.records[0].event_type, EventType::Added);
|
||||
|
||||
// LOC is preserved
|
||||
let total_added: i64 = w.records.iter().map(|r| r.lines_added).sum();
|
||||
assert_eq!(total_added, 5, "Accepted hunk's LOC should be preserved");
|
||||
}
|
||||
|
||||
/// When a hunk *shrinks* (e.g., human deletes 3 of 10 agent lines), the
|
||||
/// delta must be negative so SUM-based LOC totals stay accurate.
|
||||
#[tokio::test]
|
||||
async fn sink_shrinking_hunk_produces_negative_delta() {
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
let ctx = make_ctx();
|
||||
let cancel = tokio_util::sync::CancellationToken::new();
|
||||
|
||||
// Agent adds 10 lines
|
||||
let mut hunk = sample_agent_hunk();
|
||||
hunk.line_info.new_count = 10;
|
||||
hunk.line_info.old_count = 0;
|
||||
|
||||
// Human deletes 3 → hunk shrinks to 7
|
||||
let mut shrunk = sample_agent_hunk();
|
||||
shrunk.line_info.new_count = 7;
|
||||
shrunk.line_info.old_count = 0;
|
||||
|
||||
tx.send(HunkEvent::HunkAdded {
|
||||
path: hunk.path.clone(),
|
||||
hunk: Arc::new(hunk),
|
||||
})
|
||||
.unwrap();
|
||||
tx.send(HunkEvent::HunkContentChanged {
|
||||
path: shrunk.path.clone(),
|
||||
hunk: Arc::new(shrunk),
|
||||
trigger_source: HunkSource::ExternalEditOnAgentFile,
|
||||
prev_lines_added: 10,
|
||||
prev_lines_removed: 0,
|
||||
})
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let (shared_writer, shared) = SharedWriter::new();
|
||||
run_loc_sink(rx, shared_writer, ctx, cancel).await;
|
||||
|
||||
let w = shared.lock().unwrap();
|
||||
assert_eq!(w.records.len(), 2);
|
||||
|
||||
// First: agent added 10 lines
|
||||
assert_eq!(w.records[0].author_type, Some(AuthorType::Agent));
|
||||
assert_eq!(w.records[0].lines_added, 10i64);
|
||||
|
||||
// Second: human shrunk the hunk by 3 → negative delta
|
||||
assert_eq!(w.records[1].author_type, Some(AuthorType::Human));
|
||||
assert_eq!(w.records[1].event_type, EventType::Updated);
|
||||
assert_eq!(w.records[1].lines_added, -3i64);
|
||||
assert_eq!(w.records[1].lines_removed, 0i64);
|
||||
|
||||
// SUM(lines_added) by author: agent=10, human=-3, net=7 ✅
|
||||
let agent_total: i64 = w
|
||||
.records
|
||||
.iter()
|
||||
.filter(|r| r.author_type == Some(AuthorType::Agent))
|
||||
.map(|r| r.lines_added)
|
||||
.sum();
|
||||
let human_total: i64 = w
|
||||
.records
|
||||
.iter()
|
||||
.filter(|r| r.author_type == Some(AuthorType::Human))
|
||||
.map(|r| r.lines_added)
|
||||
.sum();
|
||||
assert_eq!(agent_total, 10);
|
||||
assert_eq!(human_total, -3);
|
||||
assert_eq!(agent_total + human_total, 7); // net lines in file
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sink_drains_on_cancellation() {
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
let ctx = make_ctx();
|
||||
let cancel = tokio_util::sync::CancellationToken::new();
|
||||
|
||||
let hunk1 = sample_agent_hunk();
|
||||
let hunk2 = sample_external_hunk();
|
||||
|
||||
// Send events before cancellation
|
||||
tx.send(HunkEvent::HunkAdded {
|
||||
path: hunk1.path.clone(),
|
||||
hunk: Arc::new(hunk1),
|
||||
})
|
||||
.unwrap();
|
||||
tx.send(HunkEvent::HunkAdded {
|
||||
path: hunk2.path.clone(),
|
||||
hunk: Arc::new(hunk2),
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
// Cancel immediately
|
||||
cancel.cancel();
|
||||
|
||||
let (shared_writer, shared) = SharedWriter::new();
|
||||
run_loc_sink(rx, shared_writer, ctx, cancel).await;
|
||||
|
||||
let w = shared.lock().unwrap();
|
||||
assert_eq!(
|
||||
w.records.len(),
|
||||
2,
|
||||
"Both events should be drained on cancellation"
|
||||
);
|
||||
assert!(w.flush_count > 0, "Writer should be flushed on shutdown");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// JSONL round-trip test
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn jsonl_round_trip() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("hunk_records.jsonl");
|
||||
let mut writer = JsonlHunkRecordWriter::new(path.clone());
|
||||
|
||||
let hunk = sample_agent_hunk();
|
||||
let record = HunkRecord::from_hunk(
|
||||
&hunk,
|
||||
"sess-rt",
|
||||
"agent-rt",
|
||||
Some("user-rt"),
|
||||
EventType::Added,
|
||||
&hunk.source,
|
||||
);
|
||||
|
||||
writer.write(&record).await.unwrap();
|
||||
writer.flush().await.unwrap();
|
||||
|
||||
// Read back and deserialize
|
||||
let contents = tokio::fs::read_to_string(&path).await.unwrap();
|
||||
let lines: Vec<&str> = contents.trim().lines().collect();
|
||||
assert_eq!(lines.len(), 1);
|
||||
|
||||
let deserialized: HunkRecord = serde_json::from_str(lines[0]).unwrap();
|
||||
assert_eq!(deserialized, record);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn jsonl_writer_creates_parent_dirs() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("nested").join("deep").join("records.jsonl");
|
||||
let mut writer = JsonlHunkRecordWriter::new(path.clone());
|
||||
|
||||
let hunk = sample_agent_hunk();
|
||||
let record = HunkRecord::from_hunk(
|
||||
&hunk,
|
||||
"sess-1",
|
||||
"agent-1",
|
||||
None,
|
||||
EventType::Added,
|
||||
&hunk.source,
|
||||
);
|
||||
|
||||
writer.write(&record).await.unwrap();
|
||||
assert!(path.exists());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn jsonl_writer_appends() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("records.jsonl");
|
||||
let mut writer = JsonlHunkRecordWriter::new(path.clone());
|
||||
|
||||
let hunk1 = sample_agent_hunk();
|
||||
let hunk2 = sample_external_hunk();
|
||||
let r1 = HunkRecord::from_hunk(&hunk1, "s", "a", None, EventType::Added, &hunk1.source);
|
||||
let r2 = HunkRecord::from_hunk(&hunk2, "s", "a", None, EventType::Added, &hunk2.source);
|
||||
|
||||
writer.write(&r1).await.unwrap();
|
||||
writer.write(&r2).await.unwrap();
|
||||
writer.flush().await.unwrap();
|
||||
|
||||
let contents = tokio::fs::read_to_string(&path).await.unwrap();
|
||||
let lines: Vec<&str> = contents.trim().lines().collect();
|
||||
assert_eq!(lines.len(), 2);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Deserialization validation
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Invalid enum values must be rejected during deserialization.
|
||||
/// This validates that the serde enum gate works — a typo like "foo"
|
||||
/// in the JSONL can't silently sneak past.
|
||||
#[test]
|
||||
fn deserialize_rejects_invalid_author_type() {
|
||||
let hunk = sample_agent_hunk();
|
||||
let record = HunkRecord::from_hunk(&hunk, "s", "a", None, EventType::Added, &hunk.source);
|
||||
let mut json = serde_json::to_string(&record).unwrap();
|
||||
|
||||
// Replace valid "agent" with invalid "foo"
|
||||
json = json.replacen("\"agent\"", "\"foo\"", 1);
|
||||
let result = serde_json::from_str::<HunkRecord>(&json);
|
||||
assert!(result.is_err(), "Should reject invalid author_type");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deserialize_rejects_invalid_event_type() {
|
||||
let hunk = sample_agent_hunk();
|
||||
let record = HunkRecord::from_hunk(&hunk, "s", "a", None, EventType::Added, &hunk.source);
|
||||
let mut json = serde_json::to_string(&record).unwrap();
|
||||
|
||||
json = json.replacen("\"added\"", "\"foo\"", 1);
|
||||
let result = serde_json::from_str::<HunkRecord>(&json);
|
||||
assert!(result.is_err(), "Should reject invalid event_type");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deserialize_rejects_invalid_source_type() {
|
||||
let hunk = sample_agent_hunk();
|
||||
let record = HunkRecord::from_hunk(&hunk, "s", "a", None, EventType::Added, &hunk.source);
|
||||
let mut json = serde_json::to_string(&record).unwrap();
|
||||
|
||||
json = json.replacen("\"agentEdit\"", "\"foo\"", 1);
|
||||
let result = serde_json::from_str::<HunkRecord>(&json);
|
||||
assert!(result.is_err(), "Should reject invalid source_type");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Writer failure resilience
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// The sink must continue processing events even when the writer fails.
|
||||
/// This validates the "log warning and drop the record" error policy.
|
||||
#[tokio::test]
|
||||
async fn sink_continues_after_writer_failure() {
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
let ctx = make_ctx();
|
||||
let cancel = tokio_util::sync::CancellationToken::new();
|
||||
|
||||
let hunk1 = sample_agent_hunk();
|
||||
let hunk2 = sample_external_hunk();
|
||||
|
||||
tx.send(HunkEvent::HunkAdded {
|
||||
path: hunk1.path.clone(),
|
||||
hunk: Arc::new(hunk1),
|
||||
})
|
||||
.unwrap();
|
||||
tx.send(HunkEvent::HunkAdded {
|
||||
path: hunk2.path.clone(),
|
||||
hunk: Arc::new(hunk2),
|
||||
})
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
/// Writer that always fails on write but tracks flush calls.
|
||||
struct FailingWriter(std::sync::Arc<std::sync::Mutex<bool>>);
|
||||
|
||||
impl HunkRecordWriter for FailingWriter {
|
||||
async fn write(&mut self, _record: &HunkRecord) -> std::io::Result<()> {
|
||||
Err(std::io::Error::other("disk full"))
|
||||
}
|
||||
async fn flush(&mut self) -> std::io::Result<()> {
|
||||
*self.0.lock().unwrap() = true;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
let flush_called = std::sync::Arc::new(std::sync::Mutex::new(false));
|
||||
let writer = FailingWriter(flush_called.clone());
|
||||
|
||||
// This must not panic — the sink should log warnings and continue.
|
||||
run_loc_sink(rx, writer, ctx, cancel).await;
|
||||
|
||||
// Flush must still be called on shutdown (sink didn't abort early).
|
||||
assert!(
|
||||
*flush_called.lock().unwrap(),
|
||||
"Sink should flush on shutdown even after write failures"
|
||||
);
|
||||
}
|
||||
Loading…
Reference in a new issue