use contextgraph_types::{
Capabilities, ContextQuery, ContextQueryResult, ErrorCode, ProviderInfo, VerifyRequest,
VerifyResponse,
};
use serde::{Deserialize, Serialize};
use crate::error::HostError;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum Envelope {
Handshake { protocol_version: String },
HandshakeAck {
protocol_version: String,
provider: ProviderInfo,
capabilities: Capabilities,
},
Query {
#[serde(default, skip_serializing_if = "Option::is_none")]
id: Option<String>,
query: ContextQuery,
},
Frames {
#[serde(default, skip_serializing_if = "Option::is_none")]
id: Option<String>,
result: ContextQueryResult,
},
Verify { request: VerifyRequest },
Verified { response: VerifyResponse },
Shutdown,
Error {
#[serde(default, skip_serializing_if = "Option::is_none")]
id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
code: Option<ErrorCode>,
message: String,
},
}
impl Envelope {
pub fn correlation_id(&self) -> Option<&str> {
match self {
Envelope::Query { id, .. }
| Envelope::Frames { id, .. }
| Envelope::Error { id, .. } => id.as_deref(),
_ => None,
}
}
pub fn error_code(&self) -> Option<ErrorCode> {
match self {
Envelope::Error { code, .. } => Some(code.clone().unwrap_or(ErrorCode::Internal)),
_ => None,
}
}
}
pub fn next_correlation_id() -> String {
use std::sync::atomic::{AtomicU64, Ordering};
static COUNTER: AtomicU64 = AtomicU64::new(1);
format!("q{}", COUNTER.fetch_add(1, Ordering::Relaxed))
}
pub fn verify_correlation(
provider_id: &str,
sent: Option<&str>,
echoed: Option<&str>,
) -> Result<(), HostError> {
let Some(expected) = sent else {
return Ok(());
};
match echoed {
Some(got) if got == expected => Ok(()),
other => Err(HostError::CorrelationMismatch {
id: provider_id.to_string(),
expected: expected.to_string(),
got: other.unwrap_or("<absent>").to_string(),
}),
}
}
pub fn envelope_kind(env: &Envelope) -> &'static str {
match env {
Envelope::Handshake { .. } => "handshake",
Envelope::HandshakeAck { .. } => "handshake_ack",
Envelope::Query { .. } => "query",
Envelope::Frames { .. } => "frames",
Envelope::Verify { .. } => "verify",
Envelope::Verified { .. } => "verified",
Envelope::Shutdown => "shutdown",
Envelope::Error { .. } => "error",
}
}
pub fn encode_line(env: &Envelope) -> Result<String, HostError> {
let mut line = serde_json::to_string(env).map_err(|e| HostError::Wire(e.to_string()))?;
line.push('\n');
Ok(line)
}
pub fn decode_line(line: &str) -> Result<Envelope, HostError> {
serde_json::from_str(line.trim_end()).map_err(|e| HostError::Wire(e.to_string()))
}
pub fn versions_compatible(a: &str, b: &str) -> bool {
protocol_family(a) == protocol_family(b)
}
fn protocol_family(version: &str) -> &str {
match version.split_once('.') {
Some((family, _)) => family,
None => version,
}
}
#[cfg(test)]
mod tests {
use super::*;
use contextgraph_types::capability::QueryCapability;
use contextgraph_types::{DataFlow, FrameKind, PROTOCOL_VERSION};
fn sample_ack() -> Envelope {
Envelope::HandshakeAck {
protocol_version: PROTOCOL_VERSION.to_string(),
provider: ProviderInfo {
name: "contextgraph-docs".into(),
version: "0.1.0".into(),
data_flow: DataFlow {
reads: true,
writes: false,
egress: false,
egress_scopes: vec![],
},
},
capabilities: Capabilities {
query: QueryCapability {
kinds: vec!["doc".into()],
},
..Capabilities::default()
},
}
}
#[test]
fn envelope_kind_matches_the_serialized_type_tag_for_every_variant() {
let variants: Vec<Envelope> = vec![
Envelope::Handshake {
protocol_version: PROTOCOL_VERSION.to_string(),
},
sample_ack(),
Envelope::Query {
id: None,
query: ContextQuery {
goal: "g".into(),
query_text: None,
embedding: None,
kinds: vec![],
anchors: vec![],
max_frames: 1,
max_tokens: 1,
as_of: None,
representation_preferences: vec![],
},
},
Envelope::Frames {
id: None,
result: ContextQueryResult {
frames: vec![],
truncated: false,
dropped_estimate: None,
},
},
Envelope::Shutdown,
Envelope::Error {
id: None,
code: None,
message: "m".into(),
},
];
for env in &variants {
let value: serde_json::Value =
serde_json::from_str(encode_line(env).unwrap().trim_end()).unwrap();
assert_eq!(
value["type"].as_str(),
Some(envelope_kind(env)),
"envelope_kind drifted from the serde tag for {env:?}"
);
}
}
#[test]
fn envelope_roundtrips_through_a_single_ndjson_line() {
let env = sample_ack();
let line = encode_line(&env).unwrap();
assert!(line.ends_with('\n'), "a frame is exactly one line");
assert_eq!(line.matches('\n').count(), 1, "no embedded newlines");
let back = decode_line(&line).unwrap();
assert_eq!(back, env);
}
#[test]
fn query_envelope_carries_contextgraph_types_shapes_verbatim() {
let query = ContextQuery {
goal: "fix the failing test".into(),
query_text: None,
embedding: None,
kinds: vec![FrameKind::Doc],
anchors: vec![],
max_frames: 5,
max_tokens: 2000,
as_of: None,
representation_preferences: vec![],
};
let env = Envelope::Query {
id: None,
query: query.clone(),
};
let line = encode_line(&env).unwrap();
match decode_line(&line).unwrap() {
Envelope::Query { query: back, .. } => assert_eq!(back, query),
other => panic!("expected query, got {}", envelope_kind(&other)),
}
}
#[test]
fn a_garbage_line_is_a_clean_wire_error_never_a_panic() {
let err = decode_line("this is not json {{{").unwrap_err();
assert!(matches!(err, HostError::Wire(_)));
}
#[test]
fn version_families_interoperate_within_a_major_but_not_across() {
assert!(versions_compatible(
"contextgraph/1.0-draft",
"contextgraph/1.0"
));
assert!(versions_compatible(
"contextgraph/1.0-draft",
"contextgraph/1.0-draft"
));
assert!(versions_compatible(PROTOCOL_VERSION, "contextgraph/1.9"));
assert!(!versions_compatible(
"contextgraph/1.0-draft",
"contextgraph/2.0"
));
assert!(!versions_compatible("contextgraph/1.0", "mcp/1.0"));
}
}