Skip to main content

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, normalize_native_web_search_for_upstream, run_messages_loop, run_messages_stream,
11    validate_native_web_search_request,
12};
13use agentic_core::proxy::{ProxyAuth, ProxyRequest, error_response_for_auth, proxy_request_with_path};
14use agentic_core::tool::ToolRegistry;
15use agentic_core::types::messages::{MessagesRequest, has_gateway_tool, registry_tools};
16
17use super::super::common::{convert_response, read_bytes_with_auth, sse_response};
18use crate::app::AppState;
19
20async fn proxy_messages(
21    state: &AppState,
22    parts: axum::http::request::Parts,
23    body: Bytes,
24    path: &'static str,
25) -> Response {
26    convert_response(
27        proxy_request_with_path(
28            ProxyRequest {
29                headers: parts.headers,
30                body,
31                query: parts.uri.query().map(str::to_owned),
32            },
33            path,
34            ProxyAuth::Anthropic,
35            &state.proxy_state,
36        )
37        .await,
38    )
39}
40
41/// Extract the client's Anthropic credential — `x-api-key` (Anthropic-native) or
42/// an `Authorization: Bearer` — falling back to the server's configured key.
43/// Consistent with the proxy path forwarding the client's `x-api-key` (E15).
44fn extract_client_key(headers: &HeaderMap, config_key: Option<&str>) -> Option<String> {
45    headers
46        .get("x-api-key")
47        .and_then(|v| v.to_str().ok())
48        .filter(|s| !s.is_empty())
49        .map(str::to_owned)
50        .or_else(|| {
51            headers
52                .get("authorization")
53                .and_then(|v| v.to_str().ok())
54                .and_then(|v| v.strip_prefix("Bearer "))
55                .filter(|s| !s.is_empty())
56                .map(str::to_owned)
57        })
58        .or_else(|| config_key.filter(|s| !s.is_empty()).map(str::to_owned))
59}
60
61/// Render an executor error as the Anthropic error envelope
62/// (`{"type":"error","error":{"type","message"}}`), consistent with the proxy
63/// path (E14).
64fn messages_error_response(err: &ExecutorError) -> Response {
65    convert_response(error_response_for_auth(
66        err.http_status(),
67        err.error_code(),
68        &err.to_string(),
69        ProxyAuth::Anthropic,
70    ))
71}
72
73/// Drive the Messages-native gateway tool loop (non-streaming or streaming) for
74/// a request that declares a gateway-owned tool.
75async fn execute_messages(state: &AppState, headers: &HeaderMap, req: &MessagesRequest, body: &Bytes) -> Response {
76    let auth = extract_client_key(headers, state.openai_api_key.as_deref());
77
78    // Build the request-scoped registry from the declared tools (M6). Gateway
79    // ownership (incl. configured aliases like Claude Code's `WebSearch`) is
80    // resolved against the operator-configured map.
81    let gateway_map = &state.exec_ctx.messages_gateway_tools;
82    let mut tools = registry_tools(req.tools.as_ref(), gateway_map);
83    let mut executors = state.exec_ctx.gateway_executors.clone();
84    let registry = match ToolRegistry::build_with_handlers(&mut tools, &mut executors).await {
85        Ok(r) => r,
86        Err(e) => return messages_error_response(&ExecutorError::from(e)),
87    };
88
89    // Parse the raw body to a JSON Value so unmodeled Anthropic fields remain
90    // intact. Native web-search declarations are normalized before upstream use.
91    let request_json: serde_json::Value = match serde_json::from_slice(body) {
92        Ok(v) => v,
93        Err(e) => return messages_error_response(&ExecutorError::from(e)),
94    };
95    if let Err(error) = validate_native_web_search_request(&request_json) {
96        return messages_error_response(&error);
97    }
98
99    if req.stream {
100        let stream = run_messages_stream(request_json, Arc::new(registry), Arc::clone(&state.exec_ctx), auth);
101        sse_response(stream)
102    } else {
103        match run_messages_loop(request_json, &registry, &state.exec_ctx, auth.as_deref()).await {
104            Ok(message) => axum::Json(message).into_response(),
105            Err(e) => messages_error_response(&e),
106        }
107    }
108}
109
110pub async fn messages(State(state): State<AppState>, request: Request) -> Response {
111    let (parts, body) = request.into_parts();
112    let bytes: Bytes = match read_bytes_with_auth(body, ProxyAuth::Anthropic).await {
113        Ok(bytes) => bytes,
114        Err(response) => return response,
115    };
116
117    // Route to the loop only when a gateway-owned tool is declared; everything
118    // else keeps the transparent proxy path.
119    if let Ok(req) = serde_json::from_slice::<MessagesRequest>(&bytes) {
120        let route_to_loop = has_gateway_tool(req.tools.as_ref(), &state.exec_ctx.messages_gateway_tools);
121        debug!(
122            route = if route_to_loop { "messages_loop" } else { "proxy" },
123            stream = req.stream,
124            tools = req.tools.as_ref().map_or(0, Vec::len),
125            "routing HTTP messages request"
126        );
127        if route_to_loop {
128            return execute_messages(&state, &parts.headers, &req, &bytes).await;
129        }
130    }
131
132    proxy_messages(&state, parts, bytes, "/v1/messages").await
133}
134
135pub async fn count_tokens(State(state): State<AppState>, request: Request) -> Response {
136    let (parts, body) = request.into_parts();
137    let mut bytes: Bytes = match read_bytes_with_auth(body, ProxyAuth::Anthropic).await {
138        Ok(bytes) => bytes,
139        Err(response) => return response,
140    };
141    if let Ok(mut request_json) = serde_json::from_slice::<serde_json::Value>(&bytes) {
142        match normalize_native_web_search_for_upstream(&mut request_json) {
143            Ok(true) => match serde_json::to_vec(&request_json) {
144                Ok(body) => bytes = Bytes::from(body),
145                Err(error) => return messages_error_response(&ExecutorError::from(error)),
146            },
147            Ok(false) => {}
148            Err(error) => return messages_error_response(&error),
149        }
150    }
151    proxy_messages(&state, parts, bytes, "/v1/messages/count_tokens").await
152}