Line data Source code
1 : //! Checks the messages of the agent bridge against its JSON Schema, for the tests of the server and the bridge.
2 : //! It knows the keywords that proto/json/dashql/agent_bridge.json uses.
3 : //! The bridge tests include the file as their own module.
4 :
5 : #![allow(dead_code)]
6 :
7 : use serde_json::{Map, Value};
8 : use std::sync::OnceLock;
9 :
10 : /// The definitions of proto/json/dashql/agent_bridge.json
11 : pub struct BridgeSchema {
12 : definitions: Map<String, Value>,
13 : }
14 :
15 : impl BridgeSchema {
16 : /// A definition of the schema
17 1 : pub fn definition(&self, name: &str) -> Option<&Value> {
18 1 : self.definitions.get(name)
19 1 : }
20 :
21 107 : pub fn get() -> &'static BridgeSchema {
22 : static SCHEMA: OnceLock<BridgeSchema> = OnceLock::new();
23 107 : SCHEMA.get_or_init(|| {
24 1 : let text = include_str!("../../../../proto/json/dashql/agent_bridge.json");
25 1 : let schema: Value = serde_json::from_str(text).expect("the schema of the agent bridge is JSON");
26 1 : let definitions = schema["definitions"].as_object().expect("the schema has definitions").clone();
27 1 : BridgeSchema { definitions }
28 1 : })
29 107 : }
30 :
31 : /// Check a value against a definition
32 114 : pub fn check(&self, definition: &str, value: &Value) -> Result<(), String> {
33 114 : let schema = self.definitions.get(definition).ok_or(format!("unknown definition {definition}"))?;
34 114 : self.validate(schema, value, definition)
35 114 : }
36 :
37 : /// Check a message that a client receives on the WebSocket.
38 : /// `request` is the method of the client's request that a response answers.
39 101 : pub fn check_from_server(&self, message: &Value, request: Option<&str>) -> Result<(), String> {
40 101 : self.check("AgentRpcMessage", message)?;
41 101 : match (message["method"].as_str(), request) {
42 38 : (Some(method), _) => {
43 38 : let params = self.method("AgentClientMethods", method, "params")?;
44 37 : self.validate(params, &message["params"], method)
45 : }
46 63 : (None, _) if message.get("error").is_some() => Ok(()),
47 54 : (None, Some(method)) => {
48 54 : let result = self.method("AgentServerMethods", method, "result")?;
49 54 : self.validate(result, &message["result"], method)
50 : }
51 0 : (None, None) => Err(format!("a response to no request: {message}")),
52 : }
53 101 : }
54 :
55 : /// Check a message that the server receives on the WebSocket.
56 : /// `request` is the method of the server's request that a response answers.
57 1 : pub fn check_from_client(&self, message: &Value, request: Option<&str>) -> Result<(), String> {
58 1 : self.check("AgentRpcMessage", message)?;
59 1 : match (message["method"].as_str(), request) {
60 0 : (Some(method), _) => {
61 0 : let params = self.method("AgentServerMethods", method, "params")?;
62 0 : self.validate(params, &message["params"], method)
63 : }
64 1 : (None, _) if message.get("error").is_some() => Ok(()),
65 1 : (None, Some(method)) => {
66 1 : let result = self.method("AgentClientMethods", method, "result")?;
67 1 : self.validate(result, &message["result"], method)
68 : }
69 0 : (None, None) => Err(format!("a response to no request: {message}")),
70 : }
71 1 : }
72 :
73 : /// The schema of the params or the result of a method
74 93 : fn method(&self, methods: &str, method: &str, part: &str) -> Result<&Value, String> {
75 93 : let unknown = || format!("{methods} has no {part} for {method}");
76 93 : let definition = self.definitions.get(methods).ok_or_else(unknown)?;
77 93 : definition["properties"].get(method).and_then(|m| m["properties"].get(part)).ok_or_else(unknown)
78 93 : }
79 :
80 1142 : fn validate(&self, schema: &Value, value: &Value, path: &str) -> Result<(), String> {
81 1142 : let Some(schema) = schema.as_object() else { return Ok(()) };
82 1142 : if let Some(reference) = schema.get("$ref").and_then(Value::as_str) {
83 174 : let name = reference.strip_prefix("#/definitions/").ok_or(format!("unsupported reference {reference}"))?;
84 174 : let target = self.definitions.get(name).ok_or(format!("unknown definition {name}"))?;
85 174 : return self.validate(target, value, path);
86 968 : }
87 968 : if let Some(expected) = schema.get("const") {
88 11 : if value != expected {
89 0 : return Err(format!("{path}: expected {expected}, got {value}"));
90 11 : }
91 957 : }
92 968 : if let Some(Value::Array(values)) = schema.get("enum") {
93 36 : if !values.contains(value) {
94 1 : return Err(format!("{path}: {value} is none of {}", Value::Array(values.clone())));
95 35 : }
96 932 : }
97 967 : if let Some(Value::Array(options)) = schema.get("oneOf") {
98 24 : let matching = options.iter().filter(|option| self.validate(option, value, path).is_ok()).count();
99 12 : if matching != 1 {
100 1 : return Err(format!("{path}: {value} matches {matching} schemas of oneOf"));
101 11 : }
102 955 : }
103 966 : if let Some(types) = schema.get("type") {
104 848 : let types: Vec<&str> = match types {
105 94 : Value::Array(types) => types.iter().filter_map(Value::as_str).collect(),
106 754 : types => types.as_str().into_iter().collect(),
107 : };
108 928 : if !types.iter().any(|t| has_type(value, t)) {
109 0 : return Err(format!("{path}: {value} is no {}", types.join(" or ")));
110 848 : }
111 118 : }
112 966 : if let Value::Object(object) = value {
113 386 : for key in schema.get("required").and_then(Value::as_array).into_iter().flatten().filter_map(Value::as_str) {
114 386 : if !object.contains_key(key) {
115 14 : return Err(format!("{path}: {key} is missing"));
116 372 : }
117 : }
118 347 : let properties = schema.get("properties").and_then(Value::as_object);
119 813 : for (key, item) in object {
120 813 : let path = format!("{path}.{key}");
121 813 : match (properties.and_then(|p| p.get(key)), schema.get("additionalProperties")) {
122 567 : (Some(property), _) => self.validate(property, item, &path)?,
123 1 : (None, Some(Value::Bool(false))) => return Err(format!("{path} is not allowed")),
124 2 : (None, Some(additional)) => self.validate(additional, item, &path)?,
125 243 : (None, None) => {}
126 : }
127 : }
128 605 : }
129 950 : if let (Value::Array(items), Some(schema)) = (value, schema.get("items")) {
130 169 : for (i, item) in items.iter().enumerate() {
131 169 : self.validate(schema, item, &format!("{path}[{i}]"))?;
132 : }
133 908 : }
134 950 : Ok(())
135 1142 : }
136 : }
137 :
138 928 : fn has_type(value: &Value, name: &str) -> bool {
139 928 : match name {
140 928 : "object" => value.is_object(),
141 671 : "array" => value.is_array(),
142 625 : "string" => value.is_string(),
143 91 : "boolean" => value.is_boolean(),
144 83 : "null" => value.is_null(),
145 66 : "number" => value.is_number(),
146 3 : "integer" => value.is_i64() || value.is_u64(),
147 0 : _ => false,
148 : }
149 928 : }
150 :
151 : #[cfg(test)]
152 : mod tests {
153 : use super::*;
154 : use serde_json::json;
155 :
156 : #[test]
157 1 : fn defines_the_params_and_results_of_the_methods() {
158 1 : let schema = BridgeSchema::get();
159 2 : for methods in ["AgentServerMethods", "AgentClientMethods"] {
160 11 : for (method, parts) in schema.definitions[methods]["properties"].as_object().unwrap() {
161 19 : for part in parts["properties"].as_object().unwrap().values() {
162 19 : let reference = part["$ref"].as_str().unwrap_or_else(|| panic!("{method} refers to a definition"));
163 19 : let name = reference.strip_prefix("#/definitions/").unwrap();
164 19 : assert!(schema.definitions.contains_key(name), "{method} refers to the unknown {name}");
165 : }
166 : }
167 : }
168 1 : }
169 :
170 : #[test]
171 1 : fn checks_messages() {
172 1 : let schema = BridgeSchema::get();
173 1 : let status = json!({ "method": "agent/status", "params": { "session": "s1", "turnId": "t1", "state": "running", "text": "" } });
174 1 : assert_eq!(schema.check_from_server(&status, None), Ok(()));
175 1 : let done = json!({ "id": 1, "result": { "turnId": "t1", "harness": "codex" } });
176 1 : assert_eq!(schema.check_from_server(&done, Some("agent/request")), Ok(()));
177 1 : let error = json!({ "id": 1, "error": { "code": "usage", "message": "no" } });
178 1 : assert_eq!(schema.check_from_server(&error, Some("agent/request")), Ok(()));
179 1 : let reply = json!({ "id": "c1", "result": { "text": "ok", "data": [1] } });
180 1 : assert_eq!(schema.check_from_client(&reply, Some("command")), Ok(()));
181 :
182 1 : let unknown = json!({ "method": "agent/other", "params": {} });
183 1 : assert!(schema.check_from_server(&unknown, None).is_err());
184 1 : let extra = json!({ "method": "agent/done", "params": { "session": "s1", "turnId": "t1", "message": "", "x": 1 } });
185 1 : assert_eq!(schema.check_from_server(&extra, None), Err("agent/done.x is not allowed".into()));
186 1 : let state = json!({ "method": "agent/status", "params": { "session": "s1", "turnId": "t1", "state": "done", "text": "" } });
187 1 : assert!(schema.check_from_server(&state, None).unwrap_err().contains("is none of"));
188 1 : let missing = json!({ "id": 1, "result": { "turnId": "t1" } });
189 1 : assert_eq!(schema.check_from_server(&missing, Some("agent/request")), Err("agent/request: harness is missing".into()));
190 1 : let failed = json!({ "ok": false, "error": { "code": "busy", "message": "" } });
191 1 : assert_eq!(schema.check("AgentCallResponse", &failed), Ok(()));
192 1 : assert!(schema.check("AgentCallResponse", &json!({ "ok": true })).is_err());
193 1 : }
194 : }
|
