use std::collections::HashMap;
use std::sync::{Arc, Mutex, RwLock};
use std::time::{SystemTime, UNIX_EPOCH};
use crate::trace::{Origin, TraceEnvelope, TraceEvent, TraceEventContext, TraceSink};
use crate::usage::{Capability, IntentGraph, Observation};
struct Pending {
query: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[non_exhaustive]
pub enum OriginFilter {
#[default]
Any,
Exactly(Origin),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[non_exhaustive]
pub enum Provenance {
#[default]
Live,
Seeded,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[non_exhaustive]
pub struct ObservationPolicy {
pub origins: OriginFilter,
pub provenance: Provenance,
}
impl ObservationPolicy {
pub fn with_origins(mut self, origins: OriginFilter) -> Self {
self.origins = origins;
self
}
pub fn with_provenance(mut self, provenance: Provenance) -> Self {
self.provenance = provenance;
self
}
}
#[derive(Debug, PartialEq)]
pub(crate) enum Step<'a> {
Remember(&'a str),
Confirm(Capability, &'a str),
Ignore,
}
pub(crate) fn classify(event: &TraceEvent, policy: ObservationPolicy) -> Step<'_> {
match event {
TraceEvent::Search { query, origin, .. }
| TraceEvent::SkillSearch { query, origin, .. }
if accepts(policy, *origin) =>
{
Step::Remember(query)
}
TraceEvent::InvokeStart { tool_id, .. } => Step::Confirm(Capability::Tool, tool_id),
TraceEvent::SkillInvoke { skill_id, .. } => Step::Confirm(Capability::Skill, skill_id),
_ => Step::Ignore,
}
}
pub(crate) fn replay_log_into(
graph: &mut IntentGraph,
envelopes: &[TraceEnvelope],
policy: ObservationPolicy,
embeddings: &HashMap<String, Vec<f32>>,
fingerprint: Option<&str>,
) {
let mut pending: HashMap<&str, (&str, bool)> = HashMap::new();
for env in envelopes {
let session = env.session_id.as_str();
let (kind, capability_id) = match classify(&env.event, policy) {
Step::Remember(query) => {
pending.insert(session, (query, false));
continue;
}
Step::Confirm(kind, id) => (kind, id),
Step::Ignore => continue,
};
let Some(entry) = pending.get_mut(session) else {
continue; };
let query = entry.0;
let first_confirmation = !entry.1;
entry.1 = true;
if let (Some(vector), Some(fp)) = (embeddings.get(query), fingerprint) {
graph.note_query_vector(query, vector, fp);
}
graph.observe(Observation {
query,
kind,
capability_id,
ts_ms: env.ts,
first_confirmation,
seeded: policy.provenance == Provenance::Seeded,
});
}
}
pub(crate) fn queries_to_embed(
envelopes: &[TraceEnvelope],
policy: ObservationPolicy,
) -> Vec<String> {
let mut seen = std::collections::HashSet::new();
let mut out = Vec::new();
for env in envelopes {
if let TraceEvent::Search { query, origin, .. }
| TraceEvent::SkillSearch { query, origin, .. } = &env.event
&& accepts(policy, *origin)
&& seen.insert(query.as_str())
{
out.push(query.clone());
}
}
out
}
fn accepts(policy: ObservationPolicy, origin: Origin) -> bool {
match policy.origins {
OriginFilter::Any => true,
OriginFilter::Exactly(wanted) => origin == wanted,
}
}
pub struct UsageLearner {
inner: Arc<dyn TraceSink>,
graph: Arc<RwLock<IntentGraph>>,
pending: Mutex<Option<Pending>>,
policy: ObservationPolicy,
}
impl UsageLearner {
pub fn new(graph: Arc<RwLock<IntentGraph>>, inner: Arc<dyn TraceSink>) -> Self {
Self::with_policy(graph, inner, ObservationPolicy::default())
}
pub fn with_policy(
graph: Arc<RwLock<IntentGraph>>,
inner: Arc<dyn TraceSink>,
policy: ObservationPolicy,
) -> Self {
Self {
inner,
graph,
pending: Mutex::new(None),
policy,
}
}
pub fn policy(&self) -> ObservationPolicy {
self.policy
}
pub fn graph(&self) -> Arc<RwLock<IntentGraph>> {
self.graph.clone()
}
fn remember_query(&self, query: &str) {
if let Ok(mut pending) = self.pending.lock() {
*pending = Some(Pending {
query: query.to_string(),
});
}
if let Ok(graph) = self.graph.read() {
graph.arm_credit(query);
}
}
fn confirm(&self, kind: Capability, capability_id: &str, ts_ms: u64) {
let Ok(pending) = self.pending.lock() else {
return;
};
let Some(query) = pending.as_ref().map(|p| p.query.clone()) else {
return; };
drop(pending);
if let Ok(mut graph) = self.graph.write() {
let first_confirmation = graph.claim_credit(&query);
graph.observe(Observation {
query: &query,
kind,
capability_id,
ts_ms,
first_confirmation,
seeded: self.policy.provenance == Provenance::Seeded,
});
}
}
pub fn replay(&self, envelope: &TraceEnvelope) {
self.learn_from(&envelope.event, envelope.ts);
}
fn learn_from(&self, event: &TraceEvent, ts_ms: u64) {
match classify(event, self.policy) {
Step::Remember(query) => self.remember_query(query),
Step::Confirm(kind, capability_id) => self.confirm(kind, capability_id, ts_ms),
Step::Ignore => {}
}
}
}
impl TraceSink for UsageLearner {
fn record(&self, event: TraceEvent) {
self.learn_from(&event, now_ms());
self.inner.record(event);
}
fn record_with_context(&self, event: TraceEvent, context: TraceEventContext) {
self.learn_from(&event, now_ms());
self.inner.record_with_context(event, context);
}
fn record_envelope(&self, envelope: TraceEnvelope) {
self.learn_from(&envelope.event, envelope.ts);
self.inner.record_envelope(envelope);
}
fn sample_rate(&self) -> f64 {
self.inner.sample_rate()
}
}
fn now_ms() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_millis() as u64)
.unwrap_or(0)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::trace::{MemorySink, NoopSink, Origin};
fn learner() -> (Arc<UsageLearner>, Arc<RwLock<IntentGraph>>) {
let graph = Arc::new(RwLock::new(IntentGraph::empty()));
let l = Arc::new(UsageLearner::new(graph.clone(), Arc::new(NoopSink)));
(l, graph)
}
fn search(query: &str) -> TraceEvent {
TraceEvent::Search {
query: query.into(),
origin: Origin::Agent,
top_k: 5,
hits: Vec::new(),
stages: Vec::new(),
took_ms: 0,
}
}
fn invoke(tool_id: &str) -> TraceEvent {
TraceEvent::InvokeStart {
tool_id: tool_id.into(),
args_size_bytes: 0,
}
}
#[test]
fn a_search_then_invoke_becomes_one_observation() {
let (l, graph) = learner();
l.record(search("why is the build broken"));
l.record(invoke("gh_run_list"));
let g = graph.read().unwrap();
assert_eq!(g.len(), 1);
assert_eq!(g.intents[0].support, 1);
assert_eq!(g.intents[0].tools.get("gh_run_list"), Some(&1.0));
}
#[test]
fn a_search_nobody_acts_on_teaches_nothing() {
let (l, graph) = learner();
l.record(search("why is the build broken"));
assert!(graph.read().unwrap().is_empty());
}
#[test]
fn an_invoke_with_no_preceding_search_teaches_nothing() {
let (l, graph) = learner();
l.record(invoke("gh_run_list"));
assert!(graph.read().unwrap().is_empty());
}
#[test]
fn what_retrieval_returned_never_becomes_an_edge() {
let (l, graph) = learner();
l.record(TraceEvent::Search {
query: "why is the build broken".into(),
origin: Origin::Agent,
top_k: 5,
hits: vec![crate::trace::SearchHitTrace {
tool_id: "docker_build".into(),
score: 9.9,
}],
stages: Vec::new(),
took_ms: 0,
});
l.record(invoke("gh_run_list"));
let g = graph.read().unwrap();
assert_eq!(
g.intents[0].tools.keys().collect::<Vec<_>>(),
vec!["gh_run_list"]
);
}
#[test]
fn several_invokes_after_one_search_all_count_as_capabilities() {
let (l, graph) = learner();
l.record(search("why is the build broken"));
l.record(invoke("gh_run_list"));
l.record(invoke("gh_run_view"));
l.record(invoke("read_file"));
let g = graph.read().unwrap();
assert_eq!(g.len(), 1);
assert_eq!(g.intents[0].tools.len(), 3, "three capabilities were used");
assert_eq!(g.intents[0].support, 1, "but only one question was asked");
for (id, w) in &g.intents[0].tools {
assert_eq!(*w, 1.0, "{id} was used once");
}
}
#[test]
fn the_same_question_asked_twice_counts_twice() {
let (l, graph) = learner();
l.record(search("why is the build broken"));
l.record(invoke("gh_run_list"));
l.record(search("why is the build broken"));
l.record(invoke("gh_run_list"));
let g = graph.read().unwrap();
assert_eq!(g.intents[0].support, 2);
assert_eq!(g.intents[0].tools["gh_run_list"], 2.0);
}
#[test]
fn separate_searches_each_count() {
let (l, graph) = learner();
l.record(search("why is the build broken"));
l.record(invoke("gh_run_list"));
l.record(search("is the build broken again"));
l.record(invoke("gh_run_list"));
let g = graph.read().unwrap();
assert_eq!(g.intents[0].support, 2, "two questions, two observations");
}
#[test]
fn a_capability_search_across_both_registries_counts_once() {
let (l, graph) = learner();
l.record(search("why is the build broken"));
l.record(TraceEvent::SkillSearch {
query: "why is the build broken".into(),
origin: Origin::Agent,
top_k: 5,
hits: Vec::new(),
stages: Vec::new(),
took_ms: 0,
});
l.record(invoke("gh_run_list"));
l.record(TraceEvent::SkillInvoke {
skill_id: "ci-triage".into(),
took_ms: 1,
});
let g = graph.read().unwrap();
assert_eq!(g.len(), 1);
assert_eq!(
g.intents[0].support, 1,
"one question, however many catalogs it hit"
);
assert_eq!(g.intents[0].tools.len(), 1);
assert_eq!(g.intents[0].skills.len(), 1);
}
#[test]
fn two_learners_sharing_a_graph_count_a_capability_search_once() {
let graph = Arc::new(RwLock::new(IntentGraph::empty()));
let tools = Arc::new(UsageLearner::new(graph.clone(), Arc::new(NoopSink)));
let skills = Arc::new(UsageLearner::new(graph.clone(), Arc::new(NoopSink)));
tools.record(search("why is the build broken"));
skills.record(TraceEvent::SkillSearch {
query: "why is the build broken".into(),
origin: Origin::Agent,
top_k: 5,
hits: Vec::new(),
stages: Vec::new(),
took_ms: 0,
});
tools.record(invoke("gh_run_list"));
skills.record(TraceEvent::SkillInvoke {
skill_id: "ci-triage".into(),
took_ms: 1,
});
let g = graph.read().unwrap();
assert_eq!(g.len(), 1);
assert_eq!(
g.intents[0].support, 1,
"one question, even across two per-catalog learners"
);
assert_eq!(g.intents[0].tools.get("gh_run_list"), Some(&1.0));
assert_eq!(g.intents[0].skills.get("ci-triage"), Some(&1.0));
}
#[test]
fn a_new_search_replaces_the_pending_query() {
let (l, graph) = learner();
l.record(search("why is the build broken"));
l.record(search("rotate the signing key"));
l.record(invoke("vault_rotate"));
let g = graph.read().unwrap();
assert_eq!(g.len(), 1, "only the later query should have been credited");
assert!(
g.intents[0]
.members
.contains(&"rotate the signing key".to_string())
);
}
#[test]
fn skill_searches_and_skill_invokes_pair_on_the_skill_edges() {
let (l, graph) = learner();
l.record(TraceEvent::SkillSearch {
query: "why is the build broken".into(),
origin: Origin::Agent,
top_k: 5,
hits: Vec::new(),
stages: Vec::new(),
took_ms: 0,
});
l.record(TraceEvent::SkillInvoke {
skill_id: "ci-triage".into(),
took_ms: 1,
});
let g = graph.read().unwrap();
assert_eq!(g.intents[0].skills.get("ci-triage"), Some(&1.0));
assert!(g.intents[0].tools.is_empty());
}
#[test]
fn a_rejected_search_is_ignored_not_a_boundary() {
let policy =
ObservationPolicy::default().with_origins(OriginFilter::Exactly(Origin::Baseline));
assert_eq!(
classify(&search_from("q", Origin::Direct), policy),
Step::Ignore
);
assert_eq!(
classify(&search_from("q", Origin::Baseline), policy),
Step::Remember("q")
);
}
#[test]
fn only_the_attempt_confirms_an_observation() {
let policy = ObservationPolicy::default();
assert_eq!(
classify(&invoke("t"), policy),
Step::Confirm(Capability::Tool, "t")
);
assert_eq!(classify(&invoke_end("t"), policy), Step::Ignore);
assert_eq!(classify(&invoke_error("t"), policy), Step::Ignore);
}
#[test]
fn an_unrelated_event_is_never_evidence() {
assert_eq!(
classify(
&TraceEvent::AuthNeeds {
upstream: "gh".into()
},
ObservationPolicy::default()
),
Step::Ignore
);
}
fn search_from(query: &str, origin: Origin) -> TraceEvent {
TraceEvent::Search {
query: query.into(),
origin,
top_k: 5,
hits: Vec::new(),
stages: Vec::new(),
took_ms: 0,
}
}
fn invoke_end(tool_id: &str) -> TraceEvent {
TraceEvent::InvokeEnd {
tool_id: tool_id.into(),
took_ms: 1,
}
}
fn invoke_error(tool_id: &str) -> TraceEvent {
TraceEvent::InvokeError {
tool_id: tool_id.into(),
took_ms: 1,
error: "bad args".into(),
}
}
fn policy_learner(policy: ObservationPolicy) -> (Arc<UsageLearner>, Arc<RwLock<IntentGraph>>) {
let graph = Arc::new(RwLock::new(IntentGraph::empty()));
let l = Arc::new(UsageLearner::with_policy(
graph.clone(),
Arc::new(NoopSink),
policy,
));
(l, graph)
}
#[test]
fn the_default_policy_reproduces_todays_pairing_exactly() {
let events = || {
vec![
search("why is the build broken"),
invoke("gh_run_list"),
invoke("gh_run_view"),
search("rotate the signing key"),
invoke("vault_rotate"),
]
};
let (old, old_graph) = learner();
for e in events() {
old.record(e);
}
let (new, new_graph) = policy_learner(ObservationPolicy::default());
for e in events() {
new.record(e);
}
let old_g = old_graph.read().unwrap();
let new_g = new_graph.read().unwrap();
assert_eq!(old_g.intents, new_g.intents);
assert_eq!(old_g.rev(), new_g.rev());
for it in &new_g.intents {
assert_eq!(it.seeded_support, 0, "the default policy is live");
}
}
#[test]
fn only_the_required_origin_opens_an_observation_window() {
let (l, graph) = policy_learner(
ObservationPolicy::default().with_origins(OriginFilter::Exactly(Origin::Baseline)),
);
l.record(search_from("why is the build broken", Origin::Agent));
l.record(invoke("gh_run_list"));
assert!(
graph.read().unwrap().is_empty(),
"an agent search must not teach a baseline-only learner"
);
l.record(search_from("why is the build broken", Origin::Baseline));
l.record(invoke("gh_run_list"));
assert_eq!(graph.read().unwrap().len(), 1);
}
#[test]
fn a_filtered_out_search_leaves_the_pending_query_intact() {
let (l, graph) = policy_learner(
ObservationPolicy::default().with_origins(OriginFilter::Exactly(Origin::Baseline)),
);
l.record(search_from("why is the build broken", Origin::Baseline));
l.record(search_from("some internal probe", Origin::Direct));
l.record(invoke("gh_run_list"));
let g = graph.read().unwrap();
assert_eq!(g.len(), 1);
assert!(
g.intents[0]
.members
.contains(&"why is the build broken".to_string()),
"the baseline query still owns the invoke, got {:?}",
g.intents[0].members
);
}
#[test]
fn the_default_policy_still_pairs_on_the_attempt() {
let (l, graph) = learner();
l.record(search("why is the build broken"));
l.record(invoke("gh_run_list"));
assert_eq!(graph.read().unwrap().len(), 1);
}
#[test]
fn a_seeded_policy_stamps_provenance_on_what_it_credits() {
let (l, graph) = policy_learner(
ObservationPolicy::default()
.with_origins(OriginFilter::Exactly(Origin::Baseline))
.with_provenance(Provenance::Seeded),
);
l.record(search_from("why is the build broken", Origin::Baseline));
l.record(invoke("gh_run_list"));
l.record(invoke("gh_run_view"));
let g = graph.read().unwrap();
assert_eq!(g.intents[0].support, 1, "one question");
assert_eq!(g.intents[0].seeded_support, 1, "and it was seeded");
assert_eq!(g.intents[0].tools.len(), 2, "two capabilities");
}
#[test]
fn a_skill_invoke_confirms_like_a_tool_invoke() {
let (l, graph) = policy_learner(ObservationPolicy::default());
l.record(TraceEvent::SkillSearch {
query: "why is the build broken".into(),
origin: Origin::Agent,
top_k: 5,
hits: Vec::new(),
stages: Vec::new(),
took_ms: 0,
});
l.record(TraceEvent::SkillInvoke {
skill_id: "ci-triage".into(),
took_ms: 1,
});
assert_eq!(
graph.read().unwrap().intents[0].skills.get("ci-triage"),
Some(&1.0)
);
}
#[test]
fn every_event_is_forwarded_to_the_inner_sink() {
let inner = Arc::new(MemorySink::new("s"));
let graph = Arc::new(RwLock::new(IntentGraph::empty()));
let l = UsageLearner::new(graph, inner.clone());
l.record(search("why is the build broken"));
l.record(invoke("gh_run_list"));
l.record(TraceEvent::AuthNeeds {
upstream: "gh".into(),
});
assert_eq!(inner.snapshot().len(), 3);
}
#[test]
fn unrelated_events_are_forwarded_without_learning() {
let (l, graph) = learner();
l.record(TraceEvent::AuthNeeds {
upstream: "gh".into(),
});
assert!(graph.read().unwrap().is_empty());
}
}