Skip to main content

netget/client/otlp/
mod.rs

1//! Reusable OTLP connections with owned exports, response handlers and socket shutdown.
2pub mod actions;
3pub mod wire;
4use crate::server::otlp::codec::{self, Encoding, Signal};
5use crate::{
6    client::{command_support, llm_budget::call_llm_for_client},
7    llm::actions::client_trait::{Client, ClientActionResult},
8    protocol::{ConnectContext, Event},
9    state::{
10        client_handles::{ClientCommand, ClientSendOutcome},
11        AccessLogOwner, ClientStatus,
12    },
13};
14use actions::{OtlpClientProtocol, CONNECTED, ERROR, RESULT};
15use anyhow::{ensure, Context, Result};
16use bytes::Bytes;
17use futures::{future::BoxFuture, stream::FuturesUnordered, FutureExt, StreamExt};
18use http_body_util::{BodyExt, Full};
19use hyper::{client::conn::http1, Request};
20use hyper_util::rt::TokioIo;
21use opentelemetry_proto::tonic::collector::{
22    logs::v1 as logs, metrics::v1 as metrics, trace::v1 as traces,
23};
24use prost::Message;
25use serde_json::{json, Value};
26use std::{
27    net::{Shutdown, SocketAddr},
28    sync::Arc,
29    time::Duration,
30};
31use tonic::{
32    codec::CompressionEncoding,
33    transport::{Channel, ClientTlsConfig, Endpoint},
34};
35pub const DEFAULT_TRANSPORT: &str = "grpc";
36pub const DEFAULT_TLS: bool = true;
37pub const DEFAULT_GZIP: bool = false;
38pub const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
39pub const EXPORT_TIMEOUT: Duration = Duration::from_secs(30);
40pub const IDLE_TIMEOUT: Duration = Duration::from_secs(120);
41pub const MAX_OPERATIONS: usize = 16;
42const MAX_DEPTH: u8 = 4;
43const CONNECT_DEPTH: u8 = u8::MAX;
44type Handler = BoxFuture<'static, (u8, Result<Vec<Value>>)>;
45type Operation = BoxFuture<'static, (u8, Event)>;
46trait Io: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Unpin {}
47impl<T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Unpin> Io for T {}
48struct SocketGuard(std::net::TcpStream);
49impl Drop for SocketGuard {
50    fn drop(&mut self) {
51        let _ = self.0.shutdown(Shutdown::Both);
52    }
53}
54#[derive(Clone)]
55enum Transport {
56    Grpc(Channel),
57    Http(Arc<tokio::sync::Mutex<http1::SendRequest<Full<Bytes>>>>),
58}
59#[derive(Clone)]
60struct Session {
61    transport: Transport,
62    host: String,
63    gzip: bool,
64    deadline: Duration,
65}
66fn seconds(
67    params: Option<&crate::protocol::StartupParams>,
68    key: &str,
69    default: Duration,
70    max: u64,
71) -> Result<Duration> {
72    let n = params
73        .map(|p| p.get_optional_u64(key))
74        .transpose()?
75        .flatten()
76        .unwrap_or(default.as_secs());
77    ensure!(n > 0 && n <= max, "{key} must be 1..{max}");
78    Ok(Duration::from_secs(n))
79}
80pub struct OtlpClient;
81impl OtlpClient {
82    pub async fn connect(ctx: ConnectContext) -> Result<SocketAddr> {
83        let p = ctx.startup_params.as_ref();
84        let connect = seconds(p, "connect_timeout_secs", CONNECT_TIMEOUT, 60)?;
85        let deadline = seconds(p, "export_timeout_secs", EXPORT_TIMEOUT, 300)?;
86        let idle = seconds(p, "idle_timeout_secs", IDLE_TIMEOUT, 3600)?;
87        let transport = p
88            .map(|p| p.get_optional_string("transport"))
89            .transpose()?
90            .flatten()
91            .unwrap_or_else(|| DEFAULT_TRANSPORT.into());
92        ensure!(
93            matches!(transport.as_str(), "grpc" | "http"),
94            "transport must be grpc or http"
95        );
96        let tls = p
97            .map(|p| p.get_optional_bool("tls"))
98            .transpose()?
99            .flatten()
100            .unwrap_or(DEFAULT_TLS);
101        let gzip = p
102            .map(|p| p.get_optional_bool("gzip"))
103            .transpose()?
104            .flatten()
105            .unwrap_or(DEFAULT_GZIP);
106        let ca = p
107            .map(|p| p.get_optional_string("ca_cert_path"))
108            .transpose()?
109            .flatten();
110        let name = p
111            .map(|p| p.get_optional_string("server_name"))
112            .transpose()?
113            .flatten();
114        ensure!(
115            tls || (ca.is_none() && name.is_none()),
116            "TLS trust/name parameters require tls=true"
117        );
118        let uri: hyper::Uri = format!(
119            "{}://{}",
120            if tls { "https" } else { "http" },
121            ctx.remote_addr
122        )
123        .parse()?;
124        let authority = uri.authority().context("remote_addr must be host:port")?;
125        ensure!(
126            authority.port_u16().is_some()
127                && authority.port_u16() != Some(0)
128                && !authority.as_str().contains('@')
129                && uri.path() == "/",
130            "remote_addr must be host:port"
131        );
132        let name = name.unwrap_or_else(|| uri.host().unwrap().trim_matches(['[', ']']).to_owned());
133        let pem = if let Some(path) = ca {
134            Some(
135                tokio::task::spawn_blocking(move || -> Result<Vec<u8>> {
136                    use std::io::Read;
137                    ensure!(
138                        std::fs::metadata(&path)?.is_file(),
139                        "CA path must be a regular file"
140                    );
141                    let file = std::fs::File::open(path)?;
142                    ensure!(file.metadata()?.is_file(), "CA path must be a regular file");
143                    let mut data = Vec::new();
144                    file.take(1048577).read_to_end(&mut data)?;
145                    ensure!(
146                        !data.is_empty() && data.len() <= 1048576,
147                        "CA file must be 1..1048576 bytes"
148                    );
149                    Ok(data)
150                })
151                .await??,
152            )
153        } else {
154            None
155        };
156        let _ = rustls::crypto::ring::default_provider().install_default();
157        let (session, socket, local, remote, driver) = tokio::time::timeout(connect, async {
158            let stream = tokio::net::TcpStream::connect(&ctx.remote_addr).await?;
159            let local = stream.local_addr()?;
160            let remote = stream.peer_addr()?;
161            let raw = stream.into_std()?;
162            let socket = SocketGuard(raw.try_clone()?);
163            let stream = tokio::net::TcpStream::from_std(raw)?;
164            let (transport, driver): (Transport, BoxFuture<'static, Result<(), hyper::Error>>) =
165                if transport == "grpc" {
166                    let mut endpoint = Endpoint::from(uri.clone())
167                        .connect_timeout(connect)
168                        .concurrency_limit(MAX_OPERATIONS)
169                        .buffer_size(MAX_OPERATIONS)
170                        .initial_stream_window_size(65536)
171                        .initial_connection_window_size(1048576)
172                        .http2_max_header_list_size(32768);
173                    if tls {
174                        let mut config = ClientTlsConfig::new()
175                            .with_webpki_roots()
176                            .domain_name(name.clone());
177                        if let Some(pem) = pem.clone() {
178                            config =
179                                config.ca_certificate(tonic::transport::Certificate::from_pem(pem));
180                        }
181                        endpoint = endpoint.tls_config(config)?;
182                    }
183                    // Reconnection cannot create an unowned socket. The sole connector
184                    // consumes the socket whose shutdown guard belongs to this client.
185                    let mut stream = Some(stream);
186                    let connector = tower::service_fn(move |_: hyper::Uri| {
187                        let stream = stream.take();
188                        async move {
189                            stream.map(TokioIo::new).ok_or_else(|| {
190                                std::io::Error::new(
191                                    std::io::ErrorKind::NotConnected,
192                                    "OTLP automatic reconnect is disabled",
193                                )
194                            })
195                        }
196                    });
197                    (
198                        Transport::Grpc(endpoint.connect_with_connector(connector).await?),
199                        futures::future::pending().boxed(),
200                    )
201                } else {
202                    let io: Box<dyn Io> = if tls {
203                        let mut roots = rustls::RootCertStore::empty();
204                        roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
205                        if let Some(pem) = pem.clone() {
206                            let certs = rustls_pemfile::certs(&mut pem.as_slice())
207                                .collect::<std::result::Result<Vec<_>, _>>()?;
208                            ensure!(!certs.is_empty(), "CA file contains no certificates");
209                            for cert in certs {
210                                roots.add(cert)?;
211                            }
212                        }
213                        let mut config = rustls::ClientConfig::builder()
214                            .with_root_certificates(roots)
215                            .with_no_client_auth();
216                        config.alpn_protocols = vec![b"http/1.1".to_vec()];
217                        Box::new(
218                            tokio_rustls::TlsConnector::from(Arc::new(config))
219                                .connect(
220                                    rustls::pki_types::ServerName::try_from(name.clone())?,
221                                    stream,
222                                )
223                                .await?,
224                        )
225                    } else {
226                        Box::new(stream)
227                    };
228                    let (sender, driver) = http1::Builder::new()
229                        .max_headers(128)
230                        .max_buf_size(32768)
231                        .handshake(TokioIo::new(io))
232                        .await?;
233                    (
234                        Transport::Http(Arc::new(tokio::sync::Mutex::new(sender))),
235                        driver.boxed(),
236                    )
237                };
238            Ok::<_, anyhow::Error>((
239                Session {
240                    transport,
241                    host: ctx.remote_addr.clone(),
242                    gzip,
243                    deadline,
244                },
245                socket,
246                local,
247                remote,
248                driver,
249            ))
250        })
251        .await
252        .context("OTLP connect deadline exceeded")??;
253        let now = crate::utils::clock::Instant::now();
254        ctx.state
255            .with_client_mut(ctx.client_id, |client| {
256                client.connection = Some(crate::state::ClientConnectionState {
257                    id: ctx.client_id,
258                    remote_addr: ctx.remote_addr.clone(),
259                    connected_addr: Some(remote),
260                    local_addr: Some(local),
261                    bytes_sent: 0,
262                    bytes_received: 0,
263                    packets_sent: 0,
264                    packets_received: 0,
265                    last_activity: now,
266                    status: ClientStatus::Connected,
267                    status_changed_at: now,
268                    protocol_info: crate::state::server::ProtocolConnectionInfo::new(
269                        json!({"transport":transport,"tls_verified":tls}),
270                    ),
271                });
272            })
273            .await;
274        ctx.state
275            .update_client_status(ctx.client_id, ClientStatus::Connected)
276            .await;
277        let commands = command_support::register_command_channel(&ctx.state, ctx.client_id).await;
278        let event = Event::new(
279            &CONNECTED,
280            json!({"remote_addr":ctx.remote_addr,"transport":transport,"tls_verified":tls,"server_name":if tls {Some(name)} else {None}}),
281        );
282        let state = ctx.state.clone();
283        let id = ctx.client_id;
284        state
285            .spawn_client_task(
286                id,
287                Self::run(ctx, session, socket, commands, event, driver, idle),
288            )
289            .await;
290        Ok(local)
291    }
292    #[allow(clippy::too_many_arguments)]
293    async fn run(
294        ctx: ConnectContext,
295        session: Session,
296        _socket: SocketGuard,
297        mut commands: tokio::sync::mpsc::Receiver<ClientCommand>,
298        event: Event,
299        mut driver: BoxFuture<'static, Result<(), hyper::Error>>,
300        idle: Duration,
301    ) {
302        let mut operations: FuturesUnordered<Operation> = FuturesUnordered::new();
303        let mut handlers: FuturesUnordered<Handler> = FuturesUnordered::new();
304        handlers.push(handler(ctx.clone(), event, CONNECT_DEPTH));
305        let mut activity = tokio::time::Instant::now();
306        loop {
307            tokio::select! {
308                _ = &mut driver => break,
309                command = commands.recv() => {
310                    let Some(command) = command else { break };
311                    activity = tokio::time::Instant::now();
312                    if enqueue(
313                        &ctx, &session, &mut operations, handlers.len(),
314                        command.action.clone(), Some(command), 0,
315                    ).await {
316                        break;
317                    }
318                }
319                Some((depth, event)) = operations.next(), if !operations.is_empty() => {
320                    activity = tokio::time::Instant::now();
321                    handlers.push(handler(ctx.clone(), event, depth));
322                }
323                Some((depth, result)) = handlers.next(), if !handlers.is_empty() => {
324                    activity = tokio::time::Instant::now();
325                    match result {
326                        Ok(actions) if actions.len() <= MAX_OPERATIONS => {
327                            let mut stop = false;
328                            for action in actions {
329                                if depth != CONNECT_DEPTH && depth >= MAX_DEPTH
330                                    && action["type"] != "disconnect"
331                                {
332                                    continue;
333                                }
334                                let next = if depth == CONNECT_DEPTH { 0 } else { depth + 1 };
335                                if enqueue(
336                                    &ctx, &session, &mut operations, handlers.len(), action, None, next,
337                                ).await {
338                                    stop = true;
339                                    break;
340                                }
341                            }
342                            if stop { break; }
343                        }
344                        Ok(_) => crate::logging::emit::Log::new(Some(&ctx.status_tx))
345                            .warn("OTLP handler exceeded 16 actions"),
346                        Err(e) => crate::logging::emit::Log::new(Some(&ctx.status_tx))
347                            .warn(format!("OTLP handler failed: {e}")),
348                    }
349                }
350                _ = tokio::time::sleep_until(activity + idle),
351                    if operations.is_empty() && handlers.is_empty() => break,
352            }
353        }
354        drop(operations);
355        drop(handlers);
356        drop(session);
357        drop(driver);
358        drop(_socket);
359        ctx.state.remove_client_handle(ctx.client_id).await;
360        ctx.state
361            .with_client_mut(ctx.client_id, |client| {
362                if let Some(c) = &mut client.connection {
363                    c.status = ClientStatus::Disconnected;
364                    c.status_changed_at = crate::utils::clock::Instant::now();
365                }
366            })
367            .await;
368        ctx.state
369            .update_client_status(ctx.client_id, ClientStatus::Disconnected)
370            .await;
371        let _ = ctx.status_tx.send("__UPDATE_UI__".into());
372    }
373}
374fn handler(ctx: ConnectContext, event: Event, depth: u8) -> Handler {
375    async move {
376        let result = async {
377            let instruction = ctx
378                .state
379                .get_instruction_for_client(ctx.client_id)
380                .await
381                .unwrap_or_default();
382            let memory = ctx
383                .state
384                .get_memory_for_client(ctx.client_id)
385                .await
386                .unwrap_or_default();
387            let result = call_llm_for_client(
388                &ctx.llm_client,
389                &ctx.state,
390                ctx.client_id.to_string(),
391                &instruction,
392                &memory,
393                Some(&event),
394                &OtlpClientProtocol,
395                &ctx.status_tx,
396            )
397            .await?;
398            if let Some(memory) = result.memory_updates {
399                ctx.state.set_memory_for_client(ctx.client_id, memory).await;
400            }
401            Ok(result.actions)
402        }
403        .await;
404        (depth, result)
405    }
406    .boxed()
407}
408async fn record(ctx: &ConnectContext, action: &Value, outcome: &Result<ClientSendOutcome>) {
409    ctx.state
410        .record_access_log(
411            AccessLogOwner::Client(ctx.client_id.as_u32()),
412            "OTLP",
413            None,
414            "injected_action",
415            action.clone(),
416            vec![match outcome {
417                Ok(v) => serde_json::to_value(v).unwrap_or(Value::Null),
418                Err(e) => json!({"error":e.to_string()}),
419            }],
420        )
421        .await;
422}
423async fn enqueue(
424    ctx: &ConnectContext,
425    session: &Session,
426    operations: &mut FuturesUnordered<Operation>,
427    handlers: usize,
428    action: Value,
429    command: Option<ClientCommand>,
430    depth: u8,
431) -> bool {
432    let outcome = match OtlpClientProtocol.execute_action(action.clone()) {
433        Err(e) => Ok(ClientSendOutcome::Rejected {
434            error: e.to_string(),
435        }),
436        Ok(ClientActionResult::Disconnect) => Ok(ClientSendOutcome::Disconnected),
437        Ok(ClientActionResult::WaitForMore | ClientActionResult::NoAction) => {
438            Ok(ClientSendOutcome::Executed {
439                detail: "wait_for_more".into(),
440            })
441        }
442        Ok(ClientActionResult::Custom { name, .. }) if name == "otlp_export" => {
443            if operations.len() + handlers >= MAX_OPERATIONS {
444                Err(anyhow::anyhow!("OTLP client busy: 16 exports or handlers"))
445            } else {
446                let ctx = ctx.clone();
447                let session = session.clone();
448                operations.push(
449                    async move {
450                        let result = tokio::time::timeout(session.deadline, export(&session, &action))
451                            .await
452                            .context("OTLP export deadline exceeded")
453                            .and_then(|r| r);
454                        let acknowledged = result.is_ok();
455                        let (outcome, event) = match result {
456                            Ok(data) => (
457                                Ok(ClientSendOutcome::Executed {
458                                    detail: format!(
459                                        "OTLP {}: {}",
460                                        data["signal"].as_str().unwrap_or(""),
461                                        data["result"].as_str().unwrap_or(""),
462                                    ),
463                                }),
464                                Event::new(&RESULT, data),
465                            ),
466                            Err(e) => (
467                                Err(anyhow::anyhow!(e.to_string())),
468                                Event::new(
469                                    &ERROR,
470                                    json!({"action_type":action["type"],"error":diagnostic(&e.to_string())}),
471                                ),
472                            ),
473                        };
474                        ctx.state
475                            .with_client_mut(ctx.client_id, |client| {
476                                if let Some(c) = &mut client.connection {
477                                    c.last_activity = crate::utils::clock::Instant::now();
478                                    if acknowledged {
479                                        c.packets_sent += 1;
480                                        c.packets_received += 1;
481                                    }
482                                }
483                            })
484                            .await;
485                        record(&ctx, &action, &outcome).await;
486                        if let Some(command) = command {
487                            command_support::reply(command, outcome);
488                        }
489                        (depth, event)
490                    }
491                    .boxed(),
492                );
493                return false;
494            }
495        }
496        Ok(_) => Ok(ClientSendOutcome::Rejected {
497            error: "unsupported OTLP action".into(),
498        }),
499    };
500    let stop = matches!(outcome, Ok(ClientSendOutcome::Disconnected));
501    record(ctx, &action, &outcome).await;
502    if let Some(command) = command {
503        command_support::reply(command, outcome);
504    }
505    stop
506}
507fn partial(signal: Signal, body: &[u8]) -> Result<(i64, String)> {
508    Ok(match signal {
509        Signal::Traces => traces::ExportTraceServiceResponse::decode(body)?
510            .partial_success
511            .map(|v| (v.rejected_spans, v.error_message)),
512        Signal::Metrics => metrics::ExportMetricsServiceResponse::decode(body)?
513            .partial_success
514            .map(|v| (v.rejected_data_points, v.error_message)),
515        Signal::Logs => logs::ExportLogsServiceResponse::decode(body)?
516            .partial_success
517            .map(|v| (v.rejected_log_records, v.error_message)),
518    }
519    .unwrap_or_default())
520}
521async fn export(session: &Session, action: &Value) -> Result<Value> {
522    let export = wire::build(action)?;
523    let signal = export.signal();
524    let body = export.encode();
525    let summary =
526        codec::summarize(signal, Encoding::Protobuf, &body).map_err(anyhow::Error::msg)?;
527    let mut data = json!({"signal":signal.as_str(),"items":summary.item_count,"service_name":action["service_name"],"retryable":false,"retry_after_secs":null,"rejected":0,"message":""});
528    let reply = match &session.transport {
529        Transport::Grpc(channel) => {
530            data["transport"] = json!("grpc");
531            macro_rules! call {
532                ($client:ty,$message:expr) => {{
533                    let mut client = <$client>::new(channel.clone())
534                        .accept_compressed(CompressionEncoding::Gzip)
535                        .max_decoding_message_size(codec::MAX_BODY_BYTES)
536                        .max_encoding_message_size(wire::MAX_EXPORT_BYTES);
537                    if session.gzip {
538                        client = client.send_compressed(CompressionEncoding::Gzip);
539                    }
540                    let mut request = tonic::Request::new($message);
541                    request.set_timeout(session.deadline);
542                    client
543                        .export(request)
544                        .await
545                        .map(|r| r.into_inner().encode_to_vec())
546                }};
547            }
548            let result = match export {
549                wire::Export::Traces(v) => {
550                    call!(traces::trace_service_client::TraceServiceClient<Channel>, v)
551                }
552                wire::Export::Metrics(v) => call!(
553                    metrics::metrics_service_client::MetricsServiceClient<Channel>,
554                    v
555                ),
556                wire::Export::Logs(v) => {
557                    call!(logs::logs_service_client::LogsServiceClient<Channel>, v)
558                }
559            };
560            match result {
561                Ok(body) => {
562                    data["grpc_code"] = json!(0);
563                    body
564                }
565                Err(status) => {
566                    if std::error::Error::source(&status).is_some() {
567                        return Err(anyhow::anyhow!(status));
568                    }
569                    data["grpc_code"] = json!(status.code() as i32);
570                    data["result"] = json!("rejected");
571                    data["message"] = json!(diagnostic(status.message()));
572                    let retry = retry_info(status.details());
573                    data["retry_after_secs"] = json!(retry);
574                    data["retryable"] = json!(
575                        matches!(
576                            status.code(),
577                            tonic::Code::Cancelled
578                                | tonic::Code::Unavailable
579                                | tonic::Code::DeadlineExceeded
580                                | tonic::Code::Aborted
581                                | tonic::Code::OutOfRange
582                                | tonic::Code::DataLoss
583                        ) || (status.code() == tonic::Code::ResourceExhausted && retry.is_some())
584                    );
585                    return Ok(data);
586                }
587            }
588        }
589        Transport::Http(sender) => {
590            data["transport"] = json!("http");
591            let mut sender = sender.lock().await;
592            sender.ready().await?;
593            let body = if session.gzip {
594                use std::io::Write;
595                let mut encoder =
596                    flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
597                encoder.write_all(&body)?;
598                encoder.finish()?
599            } else {
600                body
601            };
602            let mut request = Request::builder()
603                .method("POST")
604                .uri(format!("/v1/{}", signal.as_str()))
605                .header("host", &session.host)
606                .header("content-type", "application/x-protobuf")
607                .body(Full::new(Bytes::from(body)))?;
608            if session.gzip {
609                request.headers_mut().insert(
610                    "content-encoding",
611                    hyper::header::HeaderValue::from_static("gzip"),
612                );
613            }
614            let reply = sender.send_request(request).await?;
615            let (parts, body) = reply.into_parts();
616            let body = http_body_util::Limited::new(body, codec::MAX_BODY_BYTES)
617                .collect()
618                .await
619                .map_err(|e| anyhow::anyhow!(e.to_string()))?
620                .to_bytes();
621            let body = if parts
622                .headers
623                .get("content-encoding")
624                .is_some_and(|v| v == "gzip")
625            {
626                codec::gunzip_bounded(&body, codec::MAX_BODY_BYTES)
627                    .map_err(|e| anyhow::anyhow!("invalid bounded gzip response: {e:?}"))?
628            } else {
629                body.to_vec()
630            };
631            data["http_status"] = json!(parts.status.as_u16());
632            if parts.status.as_u16() != 200 {
633                data["result"] = json!("rejected");
634                data["retryable"] = json!(codec::retryable(parts.status.as_u16()));
635                data["retry_after_secs"] = json!(parts
636                    .headers
637                    .get("retry-after")
638                    .and_then(|v| v.to_str().ok())
639                    .and_then(|v| v.parse::<u64>().ok()));
640                let message = codec::RpcStatus::decode(body.as_slice())
641                    .map(|v| v.message)
642                    .unwrap_or_else(|_| format!("HTTP {}", parts.status.as_u16()));
643                data["message"] = json!(diagnostic(&message));
644                return Ok(data);
645            }
646            ensure!(
647                parts
648                    .headers
649                    .get("content-type")
650                    .and_then(|v| v.to_str().ok())
651                    .and_then(Encoding::from_content_type)
652                    == Some(Encoding::Protobuf),
653                "OTLP response must be protobuf"
654            );
655            body
656        }
657    };
658    let (rejected, message) = partial(signal, &reply)?;
659    ensure!(
660        rejected >= 0 && rejected as usize <= summary.item_count,
661        "invalid partial-success count"
662    );
663    data["result"] = json!(if rejected > 0 || !message.is_empty() {
664        "partial_success"
665    } else {
666        "accepted"
667    });
668    data["rejected"] = json!(rejected);
669    data["message"] = json!(diagnostic(&message));
670    Ok(data)
671}
672#[derive(Clone, PartialEq, Message)]
673struct DetailedStatus {
674    #[prost(message, repeated, tag = "3")]
675    details: Vec<prost_types::Any>,
676}
677#[derive(Clone, PartialEq, Message)]
678struct RetryInfo {
679    #[prost(message, optional, tag = "1")]
680    retry_delay: Option<prost_types::Duration>,
681}
682fn retry_info(bytes: &[u8]) -> Option<u64> {
683    if bytes.len() > 32768 {
684        return None;
685    }
686    let status = DetailedStatus::decode(bytes).ok()?;
687    for detail in status.details {
688        if detail.type_url == "type.googleapis.com/google.rpc.RetryInfo" {
689            let delay = RetryInfo::decode(detail.value.as_slice())
690                .ok()?
691                .retry_delay?;
692            if delay.seconds >= 0 && (0..1_000_000_000).contains(&delay.nanos) {
693                return (delay.seconds as u64).checked_add(u64::from(delay.nanos > 0));
694            }
695        }
696    }
697    None
698}
699
700fn diagnostic(text: &str) -> String {
701    if text.len() > 512 {
702        crate::utils::truncate_for_llm(text, 512 - crate::utils::truncate::TRUNCATION_MARKER.len())
703    } else {
704        text.to_owned()
705    }
706}