tokn_headers/
registry.rs

1//! Registry mapping `(provider, agent_id)` to the typed header schema pair.
2//!
3//! Each successful lookup yields a [`ResolvedSchema`] describing which header
4//! struct and (optional) overlay struct should be applied. Composition is
5//! overlay-wins via [`HeaderMap::merge_replacing`]: any name set by both the
6//! client-id headers and the overlay takes the overlay's value.
7//!
8//! Unknown agents for a known provider fall back to [`AgentKind::Opencode`]
9//! as a sensible default base. Unknown providers return [`None`].
10
11use crate::agent::AgentKind;
12use crate::map::HeaderMap;
13
14/// Closed enum of provider transport overlays.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
16pub enum OverlayKind {
17  Copilot,
18  Codex,
19}
20
21/// The schema pair selected for a given `(provider, persona)`.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub struct ResolvedSchema {
24  pub agent: AgentKind,
25  pub overlay: Option<OverlayKind>,
26}
27
28impl ResolvedSchema {
29  /// Compose persona-built and overlay-built maps into a single map. Overlay
30  /// values win on name conflict; otherwise persona ordering is preserved.
31  pub fn compose(persona_map: HeaderMap, overlay_map: Option<HeaderMap>) -> HeaderMap {
32    let mut out = persona_map;
33    if let Some(o) = overlay_map {
34      out.merge_replacing(o);
35    }
36    out
37  }
38}
39
40/// Look up the schema pair for a given `(provider, agent_id)`. Returns `None`
41/// for unknown providers; for known providers, falls back to
42/// [`AgentKind::Opencode`] as the base agent when `agent_id` is unknown.
43pub fn lookup(provider: &str, agent_id: &str) -> Option<ResolvedSchema> {
44  let base = AgentKind::from_agent_id(agent_id).unwrap_or(AgentKind::Opencode);
45  match provider {
46    "openai" => Some(ResolvedSchema {
47      agent: base,
48      overlay: matches!(base, AgentKind::CodexCli).then_some(OverlayKind::Codex),
49    }),
50    "copilot" | "github-copilot" => Some(ResolvedSchema {
51      agent: base,
52      overlay: Some(OverlayKind::Copilot),
53    }),
54    "deepseek" | "zai" | "zai-coding-plan" | "zhipuai" | "zhipuai-coding-plan" | "codex" => Some(ResolvedSchema {
55      agent: base,
56      overlay: matches!(provider, "codex").then_some(OverlayKind::Codex),
57    }),
58    _ => None,
59  }
60}
61
62#[cfg(test)]
63mod tests {
64  use super::*;
65  use crate::keys;
66  use crate::name::HeaderName;
67  use crate::value::HeaderValue;
68
69  #[test]
70  fn lookup_known_pairs() {
71    assert_eq!(
72      lookup("openai", "codex-cli"),
73      Some(ResolvedSchema {
74        agent: AgentKind::CodexCli,
75        overlay: Some(OverlayKind::Codex)
76      })
77    );
78    assert_eq!(
79      lookup("copilot", "opencode"),
80      Some(ResolvedSchema {
81        agent: AgentKind::Opencode,
82        overlay: Some(OverlayKind::Copilot)
83      })
84    );
85    assert_eq!(
86      lookup("deepseek", "claude-code"),
87      Some(ResolvedSchema {
88        agent: AgentKind::ClaudeCode,
89        overlay: None
90      })
91    );
92  }
93
94  #[test]
95  fn unknown_agent_falls_back_to_opencode() {
96    let r = lookup("copilot", "my-tool").unwrap();
97    assert_eq!(r.agent, AgentKind::Opencode);
98    assert_eq!(r.overlay, Some(OverlayKind::Copilot));
99  }
100
101  #[test]
102  fn unknown_provider_returns_none() {
103    assert!(lookup("nonesuch", "opencode").is_none());
104  }
105
106  #[test]
107  fn openai_with_non_codex_persona_has_no_overlay() {
108    let r = lookup("openai", "opencode").unwrap();
109    assert!(r.overlay.is_none());
110  }
111
112  #[test]
113  fn copilot_cli_resolves_with_copilot_overlay() {
114    let r = lookup("copilot", "copilot-cli").unwrap();
115    assert_eq!(r.agent, AgentKind::CopilotCli);
116    assert_eq!(r.overlay, Some(OverlayKind::Copilot));
117  }
118
119  #[test]
120  fn compose_overlay_wins_on_conflict() {
121    let mut persona_map = HeaderMap::new();
122    persona_map.insert(
123      HeaderName::new("X-Session-Id"),
124      HeaderValue::from_string("from-persona".into()),
125    );
126    persona_map.insert("X-Persona-Only", HeaderValue::from_string("p".into()));
127    let mut overlay_map = HeaderMap::new();
128    overlay_map.insert(
129      HeaderName::new("x-session-id"),
130      HeaderValue::from_string("from-overlay".into()),
131    );
132    overlay_map.insert("X-Overlay-Only", HeaderValue::from_string("o".into()));
133
134    let composed = ResolvedSchema::compose(persona_map, Some(overlay_map));
135    assert_eq!(composed.get(&keys::X_SESSION_ID).unwrap().as_str(), "from-overlay");
136    assert_eq!(composed.get("X-Persona-Only").unwrap().as_str(), "p");
137    assert_eq!(composed.get("X-Overlay-Only").unwrap().as_str(), "o");
138  }
139
140  #[test]
141  fn compose_without_overlay_is_identity() {
142    let mut persona_map = HeaderMap::new();
143    persona_map.insert("A", HeaderValue::from_string("1".into()));
144    let composed = ResolvedSchema::compose(persona_map.clone(), None);
145    assert_eq!(composed, persona_map);
146  }
147}
tokn_headers/registry.rs

tokn_headers/
registry.rs
