phi_ext/phi/
schema.rs

1//! Typed JSON Schema for tool parameters.
2//!
3//! Codex generates these via `schemars` from Rust argument structs. This crate
4//! stays dependency-free, so authors build an equivalent schema with the
5//! builders below; the wire still carries opaque JSON Schema bytes (same as
6//! Go's `Parameters map[string]any` after `json.Marshal`).
7
8use std::collections::BTreeMap;
9
10/// JSON Schema body for an LLM tool's parameters
11/// (`type` / `properties` / `required` / …).
12#[derive(Debug, Clone)]
13pub struct Schema {
14    inner: SchemaInner,
15}
16
17#[derive(Debug, Clone)]
18enum SchemaInner {
19    Built(Node),
20    /// Escape hatch for hand-written JSON Schema bytes.
21    Raw(Vec<u8>),
22}
23
24#[derive(Debug, Clone)]
25struct Node {
26    kind: Kind,
27    description: Option<String>,
28    properties: BTreeMap<String, Node>,
29    required: Vec<String>,
30    additional_properties: Option<bool>,
31    enum_values: Vec<String>,
32    items: Option<Box<Node>>,
33}
34
35#[derive(Debug, Clone, Copy, PartialEq, Eq)]
36enum Kind {
37    Object,
38    String,
39    Number,
40    Integer,
41    Boolean,
42    Array,
43}
44
45impl Default for Node {
46    fn default() -> Self {
47        Self {
48            kind: Kind::Object,
49            description: None,
50            properties: BTreeMap::new(),
51            required: Vec::new(),
52            additional_properties: None,
53            enum_values: Vec::new(),
54            items: None,
55        }
56    }
57}
58
59impl Schema {
60    fn from_node(node: Node) -> Self {
61        Self {
62            inner: SchemaInner::Built(node),
63        }
64    }
65
66    /// Object schema (`{"type":"object",…}`). Default for tool parameters.
67    pub fn object() -> Self {
68        Self::from_node(Node {
69            kind: Kind::Object,
70            ..Node::default()
71        })
72    }
73
74    pub fn string() -> Self {
75        Self::from_node(Node {
76            kind: Kind::String,
77            ..Node::default()
78        })
79    }
80
81    pub fn number() -> Self {
82        Self::from_node(Node {
83            kind: Kind::Number,
84            ..Node::default()
85        })
86    }
87
88    pub fn integer() -> Self {
89        Self::from_node(Node {
90            kind: Kind::Integer,
91            ..Node::default()
92        })
93    }
94
95    pub fn boolean() -> Self {
96        Self::from_node(Node {
97            kind: Kind::Boolean,
98            ..Node::default()
99        })
100    }
101
102    /// Array schema. `items` must be a builder schema (not [`Schema::raw`]).
103    pub fn array(items: Schema) -> Self {
104        let SchemaInner::Built(items_node) = items.inner else {
105            panic!("Schema::array requires a builder schema, not Schema::raw");
106        };
107        Self::from_node(Node {
108            kind: Kind::Array,
109            items: Some(Box::new(items_node)),
110            ..Node::default()
111        })
112    }
113
114    /// Opaque JSON Schema bytes (the previous `Vec<u8>` API).
115    pub fn raw(json: impl Into<Vec<u8>>) -> Self {
116        Self {
117            inner: SchemaInner::Raw(json.into()),
118        }
119    }
120
121    pub fn description(mut self, d: impl Into<String>) -> Self {
122        if let SchemaInner::Built(n) = &mut self.inner {
123            n.description = Some(d.into());
124        }
125        self
126    }
127
128    /// Add an object property. No-op on non-object / raw schemas.
129    pub fn property(mut self, name: impl Into<String>, schema: Schema) -> Self {
130        if let SchemaInner::Built(n) = &mut self.inner {
131            if n.kind == Kind::Object {
132                if let SchemaInner::Built(child) = schema.inner {
133                    n.properties.insert(name.into(), child);
134                }
135            }
136        }
137        self
138    }
139
140    /// Mark object property names as required.
141    pub fn required<I, S>(mut self, names: I) -> Self
142    where
143        I: IntoIterator<Item = S>,
144        S: Into<String>,
145    {
146        if let SchemaInner::Built(n) = &mut self.inner {
147            if n.kind == Kind::Object {
148                n.required.extend(names.into_iter().map(Into::into));
149            }
150        }
151        self
152    }
153
154    pub fn additional_properties(mut self, allow: bool) -> Self {
155        if let SchemaInner::Built(n) = &mut self.inner {
156            if n.kind == Kind::Object {
157                n.additional_properties = Some(allow);
158            }
159        }
160        self
161    }
162
163    /// Restrict a string schema to an enum (Codex-style compact enums).
164    pub fn enum_values<I, S>(mut self, values: I) -> Self
165    where
166        I: IntoIterator<Item = S>,
167        S: Into<String>,
168    {
169        if let SchemaInner::Built(n) = &mut self.inner {
170            if n.kind == Kind::String {
171                n.enum_values.extend(values.into_iter().map(Into::into));
172            }
173        }
174        self
175    }
176
177    /// Serialize to JSON Schema bytes for `RegisterTool`.
178    pub fn to_json_bytes(&self) -> Vec<u8> {
179        match &self.inner {
180            SchemaInner::Raw(b) => b.clone(),
181            SchemaInner::Built(n) => {
182                let mut s = String::new();
183                write_node(&mut s, n);
184                s.into_bytes()
185            }
186        }
187    }
188}
189
190impl From<Vec<u8>> for Schema {
191    fn from(json: Vec<u8>) -> Self {
192        Schema::raw(json)
193    }
194}
195
196impl From<&[u8]> for Schema {
197    fn from(json: &[u8]) -> Self {
198        Schema::raw(json.to_vec())
199    }
200}
201
202fn write_node(out: &mut String, n: &Node) {
203    out.push('{');
204    let mut first = true;
205    push_key(out, &mut first, "type");
206    push_json_string(out, kind_str(n.kind));
207
208    if let Some(d) = &n.description {
209        push_key(out, &mut first, "description");
210        push_json_string(out, d);
211    }
212
213    if n.kind == Kind::Object {
214        push_key(out, &mut first, "properties");
215        out.push('{');
216        let mut pfirst = true;
217        for (name, child) in &n.properties {
218            if !pfirst {
219                out.push(',');
220            }
221            pfirst = false;
222            push_json_string(out, name);
223            out.push(':');
224            write_node(out, child);
225        }
226        out.push('}');
227
228        if !n.required.is_empty() {
229            push_key(out, &mut first, "required");
230            out.push('[');
231            for (i, name) in n.required.iter().enumerate() {
232                if i > 0 {
233                    out.push(',');
234                }
235                push_json_string(out, name);
236            }
237            out.push(']');
238        }
239
240        if let Some(allow) = n.additional_properties {
241            push_key(out, &mut first, "additionalProperties");
242            out.push_str(if allow { "true" } else { "false" });
243        }
244    }
245
246    if n.kind == Kind::String && !n.enum_values.is_empty() {
247        push_key(out, &mut first, "enum");
248        out.push('[');
249        for (i, v) in n.enum_values.iter().enumerate() {
250            if i > 0 {
251                out.push(',');
252            }
253            push_json_string(out, v);
254        }
255        out.push(']');
256    }
257
258    if n.kind == Kind::Array {
259        if let Some(items) = &n.items {
260            push_key(out, &mut first, "items");
261            write_node(out, items);
262        }
263    }
264
265    out.push('}');
266}
267
268fn kind_str(k: Kind) -> &'static str {
269    match k {
270        Kind::Object => "object",
271        Kind::String => "string",
272        Kind::Number => "number",
273        Kind::Integer => "integer",
274        Kind::Boolean => "boolean",
275        Kind::Array => "array",
276    }
277}
278
279fn push_key(out: &mut String, first: &mut bool, key: &str) {
280    if !*first {
281        out.push(',');
282    }
283    *first = false;
284    push_json_string(out, key);
285    out.push(':');
286}
287
288fn push_json_string(out: &mut String, s: &str) {
289    out.push('"');
290    for c in s.chars() {
291        match c {
292            '"' => out.push_str("\\\""),
293            '\\' => out.push_str("\\\\"),
294            '\n' => out.push_str("\\n"),
295            '\r' => out.push_str("\\r"),
296            '\t' => out.push_str("\\t"),
297            c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)),
298            c => out.push(c),
299        }
300    }
301    out.push('"');
302}
303
304#[cfg(test)]
305mod tests {
306    use super::*;
307
308    #[test]
309    fn object_with_required_string_prop() {
310        let s = Schema::object()
311            .property("text", Schema::string().description("input text"))
312            .required(["text"]);
313        assert_eq!(
314            String::from_utf8(s.to_json_bytes()).unwrap(),
315            r#"{"type":"object","properties":{"text":{"type":"string","description":"input text"}},"required":["text"]}"#
316        );
317    }
318
319    #[test]
320    fn string_enum_and_additional_properties() {
321        let s = Schema::object()
322            .property(
323                "mode",
324                Schema::string().enum_values(["read-only", "workspace-write"]),
325            )
326            .additional_properties(false);
327        let json = String::from_utf8(s.to_json_bytes()).unwrap();
328        assert!(json.contains(r#""enum":["read-only","workspace-write"]"#));
329        assert!(json.contains(r#""additionalProperties":false"#));
330    }
331
332    #[test]
333    fn array_of_strings() {
334        let s = Schema::object().property("tags", Schema::array(Schema::string()));
335        assert_eq!(
336            String::from_utf8(s.to_json_bytes()).unwrap(),
337            r#"{"type":"object","properties":{"tags":{"type":"array","items":{"type":"string"}}}}"#
338        );
339    }
340
341    #[test]
342    fn raw_passthrough() {
343        let raw = br#"{"type":"object"}"#;
344        assert_eq!(Schema::raw(raw.to_vec()).to_json_bytes(), raw);
345    }
346}
phi_ext/phi/schema.rs

phi_ext/phi/
schema.rs
