1use crate::llm::actions::{
2 client_trait::{Client, ClientActionResult},
3 protocol_trait::Protocol,
4 ActionDefinition, ParameterDefinition, StartupExamples,
5};
6use crate::protocol::{ConnectContext, EventType};
7use crate::server::zipkin::{
8 actions::{action, parameter, SPAN_NOTE},
9 wire,
10};
11use crate::state::app_state::AppState;
12use anyhow::{bail, ensure, Context, Result};
13use serde_json::{json, Value};
14use std::sync::LazyLock;
15
16#[derive(Default)]
17pub struct ZipkinClientProtocol;
18impl ZipkinClientProtocol {
19 pub fn new() -> Self {
20 Self
21 }
22}
23
24fn report_action() -> ActionDefinition {
25 action(
26 "zipkin_report",
27 "POST spans to the collector's /api/v2/spans; the outcome arrives as zipkin_report_result.",
28 vec![
29 parameter("spans", "array", SPAN_NOTE, true),
30 parameter(
31 "gzip",
32 "boolean",
33 "Compress the body with gzip (default false)",
34 false,
35 ),
36 ],
37 json!({"type":"zipkin_report","spans":[{"traceId":"5af7183fb1d4cf5f","id":"352bff9a74ca9ad2","name":"get /cart","kind":"SERVER","timestamp":1700000000000000u64,"duration":2500,"localEndpoint":{"serviceName":"frontend"},"tags":{"http.method":"GET"}}]}),
38 "-> Zipkin report {preview(spans,100)}",
39 )
40}
41
42fn query_action() -> ActionDefinition {
43 action(
44 "zipkin_query",
45 "GET a read-API endpoint; the answer arrives as zipkin_query_result.",
46 vec![
47 parameter(
48 "endpoint",
49 "string",
50 "services, spans, remoteServices, traces, trace, traceMany, dependencies, autocompleteKeys or autocompleteValues",
51 true,
52 ),
53 parameter(
54 "trace_id",
55 "string",
56 "Required for endpoint trace: the trace id to fetch",
57 false,
58 ),
59 parameter(
60 "query",
61 "object",
62 "Query parameters the endpoint takes, e.g. {\"serviceName\": \"frontend\"}",
63 false,
64 ),
65 ],
66 json!({"type":"zipkin_query","endpoint":"trace","trace_id":"5af7183fb1d4cf5f"}),
67 "-> Zipkin query {endpoint}",
68 )
69}
70
71fn disconnect_action() -> ActionDefinition {
72 action(
73 "disconnect",
74 "End this Zipkin reporter session; no connection stays open between requests.",
75 vec![],
76 json!({"type":"disconnect"}),
77 "-> Zipkin disconnect",
78 )
79}
80
81pub fn actions() -> Vec<ActionDefinition> {
82 vec![report_action(), query_action(), disconnect_action()]
83}
84
85pub static CONNECTED_EVENT: LazyLock<EventType> = LazyLock::new(|| {
86 EventType::new(
87 "zipkin_connected",
88 "The reporter is ready; each report or query opens its own HTTP connection.",
89 report_action().example.clone(),
90 )
91 .with_parameters(vec![parameter(
92 "remote_addr",
93 "string",
94 "The collector's HTTP origin",
95 true,
96 )])
97 .with_actions(actions())
98});
99
100pub static REPORT_EVENT: LazyLock<EventType> = LazyLock::new(|| {
101 EventType::new(
102 "zipkin_report_result",
103 "The collector answered a report: 202 means accepted.",
104 query_action().example.clone(),
105 )
106 .with_parameters(vec![
107 parameter("status", "number", "HTTP status of the report", true),
108 parameter("accepted", "boolean", "True when the status is 2xx", true),
109 parameter("span_count", "number", "Spans the report carried", true),
110 parameter(
111 "message",
112 "string",
113 "The collector's error text, empty on success",
114 true,
115 ),
116 ])
117 .with_actions(actions())
118});
119
120pub static QUERY_EVENT: LazyLock<EventType> = LazyLock::new(|| {
121 EventType::new(
122 "zipkin_query_result",
123 "The collector answered a query. result holds the decoded, validated JSON on 200.",
124 disconnect_action().example.clone(),
125 )
126 .with_parameters(vec![
127 parameter("endpoint", "string", "The endpoint that was queried", true),
128 parameter("status", "number", "HTTP status of the query", true),
129 parameter(
130 "result",
131 "array|null",
132 "Names, spans, traces or links by endpoint",
133 true,
134 ),
135 parameter(
136 "message",
137 "string",
138 "Error text when status is not 200",
139 true,
140 ),
141 ])
142 .with_actions(actions())
143});
144
145pub fn query_path(v: &Value) -> Result<(String, wire::Shape)> {
147 let endpoint = v["endpoint"].as_str().context("endpoint required")?;
148 let (allowed, shape) = wire::endpoint_shape(endpoint).context("unknown Zipkin endpoint")?;
149 let mut path = format!("/api/v2/{endpoint}");
150 if endpoint == "trace" {
151 let id = wire::trace_id(v.get("trace_id").context("trace_id required")?)?;
152 path.push('/');
153 path.push_str(&id);
154 }
155 if let Some(q) = v.get("query").filter(|q| !q.is_null()) {
156 let q = q.as_object().context("query must be an object")?;
157 let mut sep = '?';
158 for (k, value) in q {
159 ensure!(
160 allowed.contains(&k.as_str()),
161 "{endpoint} does not take {k}"
162 );
163 let value = match value {
164 Value::String(s) => s.clone(),
165 Value::Number(n) => n.to_string(),
166 _ => bail!("query values must be strings or numbers"),
167 };
168 path.push(sep);
169 path.push_str(&wire::escape(k));
170 path.push('=');
171 path.push_str(&wire::escape(&value));
172 sep = '&';
173 }
174 }
175 Ok((path, shape))
176}
177
178impl Protocol for ZipkinClientProtocol {
179 fn protocol_name(&self) -> &'static str {
180 "Zipkin"
181 }
182 fn stack_name(&self) -> &'static str {
183 "ETH>IP>TCP>HTTP>Zipkin"
184 }
185 fn keywords(&self) -> Vec<&'static str> {
186 vec!["zipkin", "zipkin reporter", "tracing", "spans"]
187 }
188 fn description(&self) -> &'static str {
189 "Zipkin reporter: POSTs v2 spans to a collector and queries its read API"
190 }
191 fn get_async_actions(&self, _: &AppState) -> Vec<ActionDefinition> {
192 actions()
193 }
194 fn get_sync_actions(&self) -> Vec<ActionDefinition> {
195 vec![]
196 }
197 fn get_event_types(&self) -> Vec<EventType> {
198 vec![
199 CONNECTED_EVENT.clone(),
200 REPORT_EVENT.clone(),
201 QUERY_EVENT.clone(),
202 ]
203 }
204 fn get_startup_parameters(&self) -> Vec<ParameterDefinition> {
205 vec![]
206 }
207 fn metadata(&self) -> crate::protocol::metadata::ProtocolMetadataV2 {
208 use crate::protocol::metadata::*;
209 ProtocolMetadataV2::builder()
210 .state(DevelopmentState::Experimental)
211 .privilege_requirement(PrivilegeRequirement::None)
212 .implementation("HTTP/1.1 via hyper, one connection per request: POST /api/v2/spans (JSON v2, optional gzip) and GET on the read API, with spans validated before sending and after receiving")
213 .llm_control("Which spans to report and which queries to run, and what to do with each answer")
214 .e2e_testing("tests/client/zipkin: the official Zipkin server (zipkin-server 3.5.1, in-memory storage) and Jaeger all-in-one's Zipkin collector, each read back through its own query API")
215 .notes("JSON v2 only, cleartext HTTP. Request and response bodies are capped at 1 MiB, reports at 1000 spans. A handler chain stops after 8 follow-ups.")
216 .build()
217 }
218 fn example_prompt(&self) -> &'static str {
219 "Report a two-span trace to the Zipkin collector at localhost:9411, then fetch it back"
220 }
221 fn get_startup_examples(&self) -> StartupExamples {
222 let llm = json!({"type":"open_client","protocol":"zipkin","remote_addr":"127.0.0.1:9411","instruction":"Report one SERVER span for service frontend, then list the services"});
223 let mut scripted = llm.clone();
224 scripted["event_handlers"] = json!([{"event_pattern":"zipkin_connected","handler":{"type":"script","language":"python","code":"import json\nprint(json.dumps({'actions':[{'type':'zipkin_query','endpoint':'services'}]}))"}}]);
225 let mut static_example = llm.clone();
226 static_example["event_handlers"] = json!([{"event_pattern":"zipkin_connected","handler":{"type":"static","actions":[{"type":"zipkin_query","endpoint":"services"}]}}]);
227 StartupExamples::new(llm, scripted, static_example)
228 }
229 fn group_name(&self) -> &'static str {
230 "Network Services"
231 }
232}
233
234impl Client for ZipkinClientProtocol {
235 fn connect(
236 &self,
237 ctx: ConnectContext,
238 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<std::net::SocketAddr>> + Send>>
239 {
240 Box::pin(super::connect(ctx))
241 }
242 fn execute_action(&self, v: Value) -> Result<ClientActionResult> {
243 let name = v["type"].as_str().unwrap_or_default().to_string();
244 match name.as_str() {
245 "disconnect" => return Ok(ClientActionResult::Disconnect),
246 "zipkin_report" => {
247 let spans = wire::spans(v.get("spans").context("spans required")?)?;
248 ensure!(!spans.is_empty(), "spans must not be empty");
249 ensure!(
250 v.get("gzip").is_none_or(|g| g.is_null() || g.is_boolean()),
251 "gzip must be a boolean"
252 );
253 }
254 "zipkin_query" => {
255 query_path(&v)?;
256 }
257 _ => bail!("Unknown Zipkin client action"),
258 }
259 Ok(ClientActionResult::Custom { name, data: v })
260 }
261}