From 5630203735fd3b95f4981a40fc024cc3f22d9ee6 Mon Sep 17 00:00:00 2001 From: ForeverAngry Date: Mon, 1 Jun 2026 16:53:29 -0400 Subject: [PATCH] feat(projection): use stabilized typed ContextProvenance vocabulary Projects core fields into `ContextProvenance` and stores remaining source-specific fields in `ContextItem::metadata`, satisfying the downstream half of rig-compose's Context packing stabilization. --- src/projection.rs | 138 ++++++++++++++++++++------------- src/security/finding.rs | 76 ++++++++++-------- tests/projection_provenance.rs | 94 +++++++++++++++------- 3 files changed, 195 insertions(+), 113 deletions(-) diff --git a/src/projection.rs b/src/projection.rs index a2ccd99..e74dd28 100644 --- a/src/projection.rs +++ b/src/projection.rs @@ -1,7 +1,8 @@ //! Projection helpers for `rig-compose` context packing. use rig_compose::{ - ContextItem, ContextPack, ContextPackConfig, ContextSourceKind, Evidence, InvestigationContext, + ContextItem, ContextPack, ContextPackConfig, ContextProjectionState, ContextProvenance, + ContextSourceKind, Evidence, InvestigationContext, }; use serde_json::{Value, json}; @@ -10,10 +11,6 @@ use crate::{BehaviorPattern, EntityBaseline, MemoryLookupHit}; #[cfg(feature = "graph")] use crate::Subgraph; -const STATE_CANDIDATE: &str = "candidate"; -#[cfg(feature = "graph")] -const STATE_EXPANDED: &str = "expanded"; - /// Convert resource-native records into [`ContextItem`] values. pub trait IntoContextItem { /// Project this resource record into a prompt-ready context item. @@ -30,12 +27,15 @@ impl IntoContextItem for BehaviorPattern { }; ContextItem::new(ContextSourceKind::Resource, source_id, text) .with_score(f64::from(self.confidence_delta)) - .with_provenance(json!({ + .with_context_provenance( + ContextProvenance::new() + .with_source_uri(format!("behavior-pattern://{}@v{}", self.id, self.version)) + .with_confidence(f64::from(self.confidence_delta)) + .with_version_key(&self.id) + .with_projection_state(ContextProjectionState::Candidate), + ) + .with_metadata(json!({ "resource": "behavior_pattern", - "source_uri": format!("behavior-pattern://{}@v{}", self.id, self.version), - "confidence": self.confidence_delta, - "version_key": self.id, - "projection_state": STATE_CANDIDATE, "id": self.id, "version": self.version, "required": self.rule.required, @@ -57,12 +57,15 @@ impl IntoContextItem for EntityBaseline { ), ) .with_score(self.samples as f64) - .with_provenance(json!({ + .with_context_provenance( + ContextProvenance::new() + .with_source_uri(format!("baseline://{}/{}", self.entity, self.metric)) + .with_principal(&self.entity) + .with_confidence(self.samples as f64) + .with_projection_state(ContextProjectionState::Candidate), + ) + .with_metadata(json!({ "resource": "baseline", - "source_uri": format!("baseline://{}/{}", self.entity, self.metric), - "principal": self.entity, - "confidence": self.samples, - "projection_state": STATE_CANDIDATE, "entity": self.entity, "metric": self.metric, "mean": self.mean, @@ -92,19 +95,32 @@ pub fn memory_hit_to_context_item(hit: &MemoryLookupHit, rank: usize) -> Context .key .clone() .unwrap_or_else(|| format!("memory.hit/{rank}")); + let mut prov = ContextProvenance::new() + .with_projection_state(ContextProjectionState::Candidate) + .with_confidence(f64::from(hit.score)); + if let Some(uri) = &hit.source_uri { + prov = prov.with_source_uri(uri); + } + if let Some(principal) = &hit.principal { + prov = prov.with_principal(principal); + } + if let Some(scope) = &hit.scope { + prov = prov.with_scope(scope); + } + if let Some(ms) = hit.recorded_at_millis { + prov = prov.with_recorded_at_millis(ms); + } + if let Some(key) = &hit.key { + prov = prov.with_source_frame_id(key); + } + ContextItem::new(ContextSourceKind::Memory, source_id, hit.summary.clone()) .with_rank(rank) .with_score(f64::from(hit.score)) - .with_provenance(json!({ + .with_context_provenance(prov) + .with_metadata(json!({ "resource": "memory.lookup", "key": hit.key, - "source_uri": hit.source_uri, - "principal": hit.principal, - "scope": hit.scope, - "recorded_at_millis": hit.recorded_at_millis, - "confidence": hit.score, - "source_frame_id": hit.key, - "projection_state": STATE_CANDIDATE, "score": hit.score, "metadata": hit.metadata, })) @@ -146,12 +162,15 @@ pub fn subgraph_to_context_item(subgraph: &Subgraph, rank: usize) -> ContextItem ) .with_rank(rank) .with_score(node_count.saturating_add(edge_count) as f64) - .with_provenance(json!({ + .with_context_provenance( + ContextProvenance::new() + .with_source_uri(format!("graph://{}", subgraph.seed)) + .with_principal(&subgraph.seed) + .with_projection_state(ContextProjectionState::Expanded) + .with_reason("graph_expansion"), + ) + .with_metadata(json!({ "resource": "graph.subgraph", - "source_uri": format!("graph://{}", subgraph.seed), - "principal": subgraph.seed, - "projection_state": STATE_EXPANDED, - "reason": "graph_expansion", "seed": subgraph.seed, "nodes": subgraph.nodes, "edges": subgraph.edges, @@ -172,11 +191,17 @@ pub fn evidence_to_context_item(evidence: &Evidence, rank: usize) -> ContextItem ContextItem::new(source, source_id, evidence_text(evidence)) .with_rank(rank) .with_score(evidence_score(&evidence.detail)) - .with_provenance(json!({ + .with_context_provenance( + ContextProvenance::new() + .with_source_uri(format!( + "evidence://{}/{}", + evidence.source_skill, evidence.label + )) + .with_confidence(evidence_score(&evidence.detail)) + .with_projection_state(ContextProjectionState::Candidate), + ) + .with_metadata(json!({ "resource": "investigation.evidence", - "source_uri": format!("evidence://{}/{}", evidence.source_skill, evidence.label), - "confidence": evidence_score(&evidence.detail), - "projection_state": STATE_CANDIDATE, "source_skill": evidence.source_skill, "label": evidence.label, "detail": evidence.detail, @@ -233,10 +258,14 @@ mod tests { assert_eq!(item.source_id, "behavior_pattern/spray@v2"); assert_eq!(item.text, "password spray around one host"); assert!((item.score - 0.25).abs() < 1e-9); - assert_eq!(item.provenance["resource"], "behavior_pattern"); - assert_eq!(item.provenance["source_uri"], "behavior-pattern://spray@v2"); - assert_eq!(item.provenance["projection_state"], "candidate"); - assert_eq!(item.provenance["required"][0], "auth.failure.burst"); + assert_eq!(item.metadata["resource"], "behavior_pattern"); + let prov = item.context_provenance().unwrap(); + assert_eq!(prov.source_uri.unwrap(), "behavior-pattern://spray@v2"); + assert_eq!( + prov.projection_state.unwrap(), + ContextProjectionState::Candidate + ); + assert_eq!(item.metadata["required"][0], "auth.failure.burst"); } #[test] @@ -257,19 +286,20 @@ mod tests { assert_eq!(items[0].source, ContextSourceKind::Memory); assert_eq!(items[0].source_id, "episode-1"); assert_eq!(items[0].rank, 0); - assert_eq!(items[0].provenance["source_uri"], "memory://episode/1"); - assert_eq!(items[0].provenance["principal"], "alice"); - assert_eq!(items[0].provenance["scope"], "workspace"); - assert_eq!( - items[0].provenance["recorded_at_millis"], - 1_700_000_000_000_i64 - ); - let confidence = items[0].provenance["confidence"] - .as_f64() + let prov = items[0].context_provenance().unwrap(); + assert_eq!(prov.source_uri.unwrap(), "memory://episode/1"); + assert_eq!(prov.principal.unwrap(), "alice"); + assert_eq!(prov.scope.unwrap(), "workspace"); + assert_eq!(prov.recorded_at_millis.unwrap(), 1_700_000_000_000_i64); + let confidence = prov + .confidence .expect("confidence should serialize as a number"); assert!((confidence - 0.9).abs() < 1e-6); - assert_eq!(items[0].provenance["source_frame_id"], "episode-1"); - assert_eq!(items[0].provenance["projection_state"], "candidate"); + assert_eq!(prov.source_frame_id.unwrap(), "episode-1"); + assert_eq!( + prov.projection_state.unwrap(), + ContextProjectionState::Candidate + ); assert_eq!(items[1].source_id, "memory.hit/1"); assert_eq!(items[1].rank, 1); } @@ -313,10 +343,14 @@ mod tests { assert_eq!(item.source_id, "graph/host-1"); assert_eq!(item.rank, 3); assert_eq!(item.score, 3.0); - assert_eq!(item.provenance["resource"], "graph.subgraph"); - assert_eq!(item.provenance["source_uri"], "graph://host-1"); - assert_eq!(item.provenance["projection_state"], "expanded"); - assert_eq!(item.provenance["reason"], "graph_expansion"); - assert_eq!(item.provenance["seed"], "host-1"); + assert_eq!(item.metadata["resource"], "graph.subgraph"); + let prov = item.context_provenance().unwrap(); + assert_eq!(prov.source_uri.unwrap(), "graph://host-1"); + assert_eq!( + prov.projection_state.unwrap(), + ContextProjectionState::Expanded + ); + assert_eq!(prov.reason.unwrap(), "graph_expansion"); + assert_eq!(item.metadata["seed"], "host-1"); } } diff --git a/src/security/finding.rs b/src/security/finding.rs index cd78449..dce6662 100644 --- a/src/security/finding.rs +++ b/src/security/finding.rs @@ -26,18 +26,17 @@ //! .with_technique_id("T1110.003"); //! //! let item = security_finding_to_context_item(&finding, 0); -//! assert_eq!(item.provenance["finding_id"], "credential.password_spray"); -//! assert_eq!(item.provenance["severity"], "high"); +//! assert_eq!(item.metadata["finding_id"], "credential.password_spray"); +//! assert_eq!(item.metadata["severity"], "high"); //! ``` -use rig_compose::{ContextItem, ContextSourceKind}; +use rig_compose::{ContextItem, ContextProjectionState, ContextProvenance, ContextSourceKind}; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; use crate::projection::IntoContextItem; use crate::trace::ResourceTraceEnvelope; -const STATE_CANDIDATE: &str = "candidate"; const TRACE_RESOURCE: &str = "security"; const TRACE_OPERATION: &str = "finding"; const TRACE_KIND: &str = "security_finding"; @@ -275,6 +274,21 @@ pub fn security_finding_to_context_item(finding: &SecurityFinding, rank: usize) let source_uri = finding.resolved_source_uri(); let confidence = finding.effective_confidence(); + let mut prov = ContextProvenance::new() + .with_projection_state(ContextProjectionState::Candidate) + .with_confidence(confidence); + + prov = prov.with_source_uri(source_uri); + if let Some(principal) = &finding.principal { + prov = prov.with_principal(principal); + } + if let Some(scope) = &finding.scope { + prov = prov.with_scope(scope); + } + if let Some(ms) = finding.recorded_at_millis { + prov = prov.with_recorded_at_millis(ms); + } + ContextItem::new( ContextSourceKind::Resource, source_id, @@ -282,14 +296,9 @@ pub fn security_finding_to_context_item(finding: &SecurityFinding, rank: usize) ) .with_rank(rank) .with_score(confidence) - .with_provenance(json!({ + .with_context_provenance(prov) + .with_metadata(json!({ "resource": "security.finding", - "source_uri": source_uri, - "principal": finding.principal, - "scope": finding.scope, - "recorded_at_millis": finding.recorded_at_millis, - "confidence": confidence, - "projection_state": STATE_CANDIDATE, "finding_id": finding.id, "severity": finding.severity.as_str(), "technique_id": finding.technique_id, @@ -402,24 +411,27 @@ mod tests { assert_eq!(item.rank, 3); assert!((item.score - FindingSeverity::High.confidence_weight()).abs() < 1e-9); - let p = &item.provenance; - // Shared vocabulary - assert_eq!(p["resource"], "security.finding"); - assert_eq!(p["source_uri"], "siem://event/42"); - assert_eq!(p["principal"], "host-1"); - assert_eq!(p["scope"], "workspace"); - assert_eq!(p["recorded_at_millis"], 1_700_000_000_000_i64); - assert_eq!(p["projection_state"], "candidate"); - let confidence = p["confidence"].as_f64().unwrap(); + let m = &item.metadata; + assert_eq!(m["resource"], "security.finding"); + assert_eq!(m["finding_id"], "credential.password_spray"); + assert_eq!(m["severity"], "high"); + assert_eq!(m["technique_id"], "T1110.003"); + assert_eq!(m["tactic"], "credential-access"); + assert_eq!(m["source_skill"], "credential.password_spray"); + assert_eq!(m["signals"][0], "auth.failure.burst"); + assert_eq!(m["detail"]["distinct_accounts"], 17); + + let prov = item.context_provenance().unwrap(); + assert_eq!(prov.source_uri.unwrap(), "siem://event/42"); + assert_eq!(prov.principal.unwrap(), "host-1"); + assert_eq!(prov.scope.unwrap(), "workspace"); + assert_eq!(prov.recorded_at_millis.unwrap(), 1_700_000_000_000_i64); + assert_eq!( + prov.projection_state.unwrap(), + ContextProjectionState::Candidate + ); + let confidence = prov.confidence.unwrap(); assert!((confidence - FindingSeverity::High.confidence_weight()).abs() < 1e-9); - // Security-specific - assert_eq!(p["finding_id"], "credential.password_spray"); - assert_eq!(p["severity"], "high"); - assert_eq!(p["technique_id"], "T1110.003"); - assert_eq!(p["tactic"], "credential-access"); - assert_eq!(p["source_skill"], "credential.password_spray"); - assert_eq!(p["signals"][0], "auth.failure.burst"); - assert_eq!(p["detail"]["distinct_accounts"], 17); } #[test] @@ -429,10 +441,10 @@ mod tests { assert_eq!(item.source_id, "security.finding/recon.high_fanout"); assert_eq!( - item.provenance["source_uri"], + item.context_provenance().unwrap().source_uri.unwrap(), "security-finding://recon.high_fanout" ); - assert!(item.provenance["principal"].is_null()); + assert!(item.context_provenance().unwrap().principal.is_none()); } #[test] @@ -440,7 +452,7 @@ mod tests { let finding = SecurityFinding::new("exfil.slow_beacon", FindingSeverity::Low, "beacon") .with_confidence(0.99); let item = security_finding_to_context_item(&finding, 0); - let confidence = item.provenance["confidence"].as_f64().unwrap(); + let confidence = item.context_provenance().unwrap().confidence.unwrap(); assert!((confidence - 0.99).abs() < 1e-9); assert!((item.score - 0.99).abs() < 1e-9); } @@ -467,7 +479,7 @@ mod tests { "chain", ); let item: ContextItem = finding.to_context_item(); - assert_eq!(item.provenance["severity"], "critical"); + assert_eq!(item.metadata["severity"], "critical"); } #[test] diff --git a/tests/projection_provenance.rs b/tests/projection_provenance.rs index 462fbb2..54800e1 100644 --- a/tests/projection_provenance.rs +++ b/tests/projection_provenance.rs @@ -1,17 +1,21 @@ #![allow(clippy::unwrap_used, clippy::panic, clippy::indexing_slicing)] -use rig_compose::{ContextSourceKind, Evidence, InvestigationContext}; +use rig_compose::{ContextProjectionState, ContextSourceKind, Evidence, InvestigationContext}; use rig_resources::{ BehaviorPattern, EntityBaseline, IntoContextItem, MemoryLookupHit, PatternRule, evidence_to_context_item, evidence_to_context_items, memory_hit_to_context_item, }; use serde_json::json; -fn assert_shared_candidate_keys(provenance: &serde_json::Value) { - assert!(provenance.get("resource").is_some()); - assert!(provenance.get("source_uri").is_some()); - assert_eq!(provenance["projection_state"], "candidate"); - assert!(provenance.get("confidence").is_some()); +fn assert_shared_candidate_keys(item: &rig_compose::ContextItem) { + assert!(item.metadata.get("resource").is_some()); + let prov = item.context_provenance().unwrap(); + assert!(prov.source_uri.is_some()); + assert_eq!( + prov.projection_state.unwrap(), + ContextProjectionState::Candidate + ); + assert!(prov.confidence.is_some()); } #[test] @@ -28,9 +32,17 @@ fn default_resource_projections_carry_shared_provenance_keys() { .with_description("password spray"); let pattern_item = pattern.to_context_item(); assert_eq!(pattern_item.source, ContextSourceKind::Resource); - assert_shared_candidate_keys(&pattern_item.provenance); - assert_eq!(pattern_item.provenance["resource"], "behavior_pattern"); - assert_eq!(pattern_item.provenance["version_key"], "spray"); + assert_shared_candidate_keys(&pattern_item); + assert_eq!(pattern_item.metadata["resource"], "behavior_pattern"); + assert_eq!( + pattern_item + .context_provenance() + .unwrap() + .version_key + .as_deref() + .unwrap(), + "spray" + ); let baseline = EntityBaseline { entity: "host-1".into(), @@ -41,8 +53,16 @@ fn default_resource_projections_carry_shared_provenance_keys() { }; let baseline_item = baseline.to_context_item(); assert_eq!(baseline_item.source, ContextSourceKind::Resource); - assert_shared_candidate_keys(&baseline_item.provenance); - assert_eq!(baseline_item.provenance["principal"], "host-1"); + assert_shared_candidate_keys(&baseline_item); + assert_eq!( + baseline_item + .context_provenance() + .unwrap() + .principal + .as_deref() + .unwrap(), + "host-1" + ); let memory = MemoryLookupHit::new(0.8, "prior incident") .with_key("frame-1") @@ -52,9 +72,18 @@ fn default_resource_projections_carry_shared_provenance_keys() { .with_recorded_at_millis(1_700_000_000_000); let memory_item = memory_hit_to_context_item(&memory, 3); assert_eq!(memory_item.source, ContextSourceKind::Memory); - assert_shared_candidate_keys(&memory_item.provenance); - assert_eq!(memory_item.provenance["scope"], "tenant-a"); - assert_eq!(memory_item.provenance["source_frame_id"], "frame-1"); + assert_shared_candidate_keys(&memory_item); + let memory_prov = memory_item.context_provenance().unwrap(); + assert_eq!(memory_prov.scope.as_deref().unwrap(), "tenant-a"); + assert_eq!( + memory_prov + .source_frame_id + .as_ref() + .unwrap() + .as_str() + .unwrap(), + "frame-1" + ); let evidence = Evidence::new("detector", "finding").with_detail(json!({ "summary": "high fan-out", @@ -62,8 +91,8 @@ fn default_resource_projections_carry_shared_provenance_keys() { })); let evidence_item = evidence_to_context_item(&evidence, 1); assert_eq!(evidence_item.source, ContextSourceKind::Resource); - assert_shared_candidate_keys(&evidence_item.provenance); - assert_eq!(evidence_item.provenance["source_skill"], "detector"); + assert_shared_candidate_keys(&evidence_item); + assert_eq!(evidence_item.metadata["source_skill"], "detector"); let mut ctx = InvestigationContext::new("host-1", "tenant-a"); ctx.evidence.push(evidence); @@ -83,11 +112,15 @@ fn graph_projection_uses_expanded_state_and_reason() { let item = subgraph_to_context_item(&subgraph, 0); assert_eq!(item.source, ContextSourceKind::Resource); - assert_eq!(item.provenance["resource"], "graph.subgraph"); - assert_eq!(item.provenance["source_uri"], "graph://host-1"); - assert_eq!(item.provenance["principal"], "host-1"); - assert_eq!(item.provenance["projection_state"], "expanded"); - assert_eq!(item.provenance["reason"], "graph_expansion"); + assert_eq!(item.metadata["resource"], "graph.subgraph"); + let prov = item.context_provenance().unwrap(); + assert_eq!(prov.source_uri.as_deref().unwrap(), "graph://host-1"); + assert_eq!(prov.principal.as_deref().unwrap(), "host-1"); + assert_eq!( + prov.projection_state.unwrap(), + ContextProjectionState::Expanded + ); + assert_eq!(prov.reason.as_deref().unwrap(), "graph_expansion"); } #[cfg(feature = "security")] @@ -112,12 +145,15 @@ fn security_projection_carries_shared_and_security_specific_keys() { let item = security_finding_to_context_item(&finding, 0); assert_eq!(item.source, ContextSourceKind::Resource); - assert_shared_candidate_keys(&item.provenance); - assert_eq!(item.provenance["resource"], "security.finding"); - assert_eq!(item.provenance["principal"], "host-1"); - assert_eq!(item.provenance["scope"], "tenant-a"); - assert_eq!(item.provenance["finding_id"], "credential.password_spray"); - assert_eq!(item.provenance["severity"], "high"); - assert_eq!(item.provenance["technique_id"], "T1110.003"); - assert_eq!(item.provenance["signals"][0], "auth.failure.burst"); + assert_shared_candidate_keys(&item); + assert_eq!(item.metadata["resource"], "security.finding"); + + let prov = item.context_provenance().unwrap(); + assert_eq!(prov.principal.as_deref().unwrap(), "host-1"); + assert_eq!(prov.scope.as_deref().unwrap(), "tenant-a"); + + assert_eq!(item.metadata["finding_id"], "credential.password_spray"); + assert_eq!(item.metadata["severity"], "high"); + assert_eq!(item.metadata["technique_id"], "T1110.003"); + assert_eq!(item.metadata["signals"][0], "auth.failure.burst"); }