agentic_server/handler/http/
messages.rs

1use std::sync::Arc;
2
3use axum::extract::{Request, State};
4use axum::response::{IntoResponse, Response};
5use bytes::Bytes;
6use http::HeaderMap;
7use tracing::debug;
8
9use agentic_core::executor::{
10    ExecutorError, MessagesUpstream, normalize_native_web_search_for_upstream, run_messages_loop, run_messages_stream,
11    validate_native_web_search_request,
12};
13use agentic_core::proxy::{
14    ProxyAuth, ProxyBody, ProxyRequest, ProxyResponse, error_response_for_auth, proxy_request_with_path,
15    upstream_request_headers,
16};
17use agentic_core::tool::ToolRegistry;
18use agentic_core::types::messages::{MessagesRequest, has_gateway_tool, registry_tools};
19
20use super::super::common::{convert_response, read_bytes_with_auth, sse_response_with_headers};
21use crate::app::AppState;
22
23async fn proxy_messages(
24    state: &AppState,
25    parts: axum::http::request::Parts,
26    body: Bytes,
27    path: &'static str,
28) -> Response {
29    convert_response(
30        proxy_request_with_path(
31            ProxyRequest {
32                headers: parts.headers,
33                body,
34                query: parts.uri.query().map(str::to_owned),
35            },
36            path,
37            ProxyAuth::Anthropic,
38            &state.proxy_state,
39        )
40        .await,
41    )
42}
43
44/// Preserve upstream Messages errors verbatim; render local executor failures
45/// as an Anthropic error envelope, consistent with the proxy path (E14).
46fn messages_error_response(err: ExecutorError) -> Response {
47    if let ExecutorError::LLMRequest {
48        status,
49        body,
50        mut headers,
51    } = err
52    {
53        headers
54            .entry(http::header::CONTENT_TYPE)
55            .or_insert(http::HeaderValue::from_static("application/json"));
56        return convert_response(ProxyResponse {
57            status,
58            headers,
59            body: ProxyBody::Full(Bytes::from(body)),
60        });
61    }
62    convert_response(error_response_for_auth(
63        err.http_status(),
64        err.error_code(),
65        &err.to_string(),
66        ProxyAuth::Anthropic,
67    ))
68}
69
70/// Drive the Messages-native gateway tool loop (non-streaming or streaming) for
71/// a request that declares a gateway-owned tool.
72async fn execute_messages(
73    state: &AppState,
74    headers: &HeaderMap,
75    query: Option<&str>,
76    req: &MessagesRequest,
77    body: &Bytes,
78) -> Response {
79    // Build the request-scoped registry from the declared tools (M6). Gateway
80    // ownership (incl. configured aliases like Claude Code's `WebSearch`) is
81    // resolved against the operator-configured map.
82    let gateway_map = &state.exec_ctx.messages_gateway_tools;
83    let mut tools = registry_tools(req.tools.as_ref(), gateway_map);
84    let mut executors = state.exec_ctx.gateway_executors.clone();
85    let registry = match ToolRegistry::build_with_handlers(&mut tools, &mut executors).await {
86        Ok(r) => r,
87        Err(e) => return messages_error_response(ExecutorError::from(e)),
88    };
89
90    // Parse the raw body to a JSON Value the loop forwards upstream untouched —
91    // preserving every Anthropic field (tool_choice, stop_sequences, …).
92    let request_json: serde_json::Value = match serde_json::from_slice(body) {
93        Ok(v) => v,
94        Err(e) => return messages_error_response(ExecutorError::from(e)),
95    };
96    if let Err(error) = validate_native_web_search_request(&request_json) {
97        return messages_error_response(error);
98    }
99
100    let upstream = MessagesUpstream::new(
101        &state.exec_ctx.llm_base_url,
102        query,
103        upstream_request_headers(headers, &state.proxy_state.config, ProxyAuth::Anthropic),
104    );
105    if req.stream {
106        match run_messages_stream(request_json, Arc::new(registry), Arc::clone(&state.exec_ctx), upstream).await {
107            Ok(response) => sse_response_with_headers(response.body, response.headers),
108            Err(e) => messages_error_response(e),
109        }
110    } else {
111        match run_messages_loop(request_json, &registry, &state.exec_ctx, &upstream).await {
112            Ok(message) => {
113                let mut response = axum::Json(message.body).into_response();
114                response.headers_mut().extend(message.headers);
115                response.headers_mut().insert(
116                    http::header::CONTENT_TYPE,
117                    http::HeaderValue::from_static("application/json"),
118                );
119                response
120            }
121            Err(e) => messages_error_response(e),
122        }
123    }
124}
125
126pub async fn messages(State(state): State<AppState>, request: Request) -> Response {
127    let (parts, body) = request.into_parts();
128    let bytes: Bytes = match read_bytes_with_auth(body, ProxyAuth::Anthropic).await {
129        Ok(bytes) => bytes,
130        Err(response) => return response,
131    };
132
133    // Route to the loop only when a gateway-owned tool is declared; everything
134    // else keeps the transparent proxy path.
135    if let Ok(req) = serde_json::from_slice::<MessagesRequest>(&bytes) {
136        let route_to_loop = has_gateway_tool(req.tools.as_ref(), &state.exec_ctx.messages_gateway_tools);
137        debug!(
138            route = if route_to_loop { "messages_loop" } else { "proxy" },
139            stream = req.stream,
140            tools = req.tools.as_ref().map_or(0, Vec::len),
141            "routing HTTP messages request"
142        );
143        if route_to_loop {
144            return execute_messages(&state, &parts.headers, parts.uri.query(), &req, &bytes).await;
145        }
146    }
147
148    proxy_messages(&state, parts, bytes, "/v1/messages").await
149}
150
151pub async fn count_tokens(State(state): State<AppState>, request: Request) -> Response {
152    let (parts, body) = request.into_parts();
153    let mut bytes: Bytes = match read_bytes_with_auth(body, ProxyAuth::Anthropic).await {
154        Ok(bytes) => bytes,
155        Err(response) => return response,
156    };
157    if let Ok(mut request_json) = serde_json::from_slice::<serde_json::Value>(&bytes) {
158        match normalize_native_web_search_for_upstream(&mut request_json) {
159            Ok(true) => match serde_json::to_vec(&request_json) {
160                Ok(body) => bytes = Bytes::from(body),
161                Err(error) => return messages_error_response(ExecutorError::from(error)),
162            },
163            Ok(false) => {}
164            Err(error) => return messages_error_response(error),
165        }
166    }
167    proxy_messages(&state, parts, bytes, "/v1/messages/count_tokens").await
168}
agentic_server/handler/http/messages.rs

agentic_server/handler/http/
messages.rs
