mirror of
https://github.com/hanzoai/dev.git
synced 2026-10-25 03:32:32 +00:00
feat(zap): implement native ZAP binary transport in dev CLI
Real luxfi/zap wire protocol — no HTTP fallback, no stubs. TLS 1.3 connect → ZAP handshake → MsgType 100 Call → parse response. - zap_wire.rs: Builder, Message, Object, frame I/O, handshake helpers - zap_client.rs: ZapCloudClient with TLS, Call RPC, stream_zap() - client.rs: WireApi::Zap arm uses native binary transport - model_provider_info.rs: WireApi::Zap variant added - zap/transport: TLS 1.3 support via rustls - All match arms updated across CLI, TUI, exec, MCP server
This commit is contained in:
Generated
+46
@@ -509,6 +509,28 @@ dependencies = [
|
||||
"arrayvec 0.7.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-lc-rs"
|
||||
version = "1.16.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d9a7b350e3bb1767102698302bc37256cbd48422809984b98d292c40e2579aa9"
|
||||
dependencies = [
|
||||
"aws-lc-sys",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-lc-sys"
|
||||
version = "0.37.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b092fe214090261288111db7a2b2c2118e5a7f30dc2569f1732c4069a6840549"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"cmake",
|
||||
"dunce",
|
||||
"fs_extra",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "axum"
|
||||
version = "0.7.9"
|
||||
@@ -1074,6 +1096,15 @@ dependencies = [
|
||||
"error-code",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cmake"
|
||||
version = "0.1.57"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "75443c44cd6b379beb8c5b45d85d0773baf31cce901fe7bb252f4eff3008ef7d"
|
||||
dependencies = [
|
||||
"cc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cmp_any"
|
||||
version = "0.8.1"
|
||||
@@ -2183,6 +2214,12 @@ dependencies = [
|
||||
"winapi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fs_extra"
|
||||
version = "1.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
|
||||
|
||||
[[package]]
|
||||
name = "funty"
|
||||
version = "2.0.0"
|
||||
@@ -2756,6 +2793,7 @@ dependencies = [
|
||||
"rand 0.9.2",
|
||||
"regex-lite",
|
||||
"reqwest",
|
||||
"rustls",
|
||||
"schemars 0.8.22",
|
||||
"seccompiler",
|
||||
"serde",
|
||||
@@ -2772,6 +2810,7 @@ dependencies = [
|
||||
"thiserror 2.0.18",
|
||||
"time",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"tokio-stream",
|
||||
"tokio-test",
|
||||
"tokio-tungstenite 0.23.1",
|
||||
@@ -2784,6 +2823,7 @@ dependencies = [
|
||||
"url",
|
||||
"uuid",
|
||||
"walkdir",
|
||||
"webpki-roots 0.26.11",
|
||||
"which",
|
||||
"wildmatch",
|
||||
"windows-sys 0.61.2",
|
||||
@@ -3278,13 +3318,16 @@ dependencies = [
|
||||
"hanzo-protocol",
|
||||
"regex",
|
||||
"reqwest",
|
||||
"rustls",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"tokio-test",
|
||||
"tracing",
|
||||
"walkdir",
|
||||
"webpki-roots 0.26.11",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -6379,6 +6422,8 @@ version = "0.23.36"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c665f33d38cea657d9614f766881e4d510e0eda4239891eea56b4cadcf01801b"
|
||||
dependencies = [
|
||||
"aws-lc-rs",
|
||||
"log",
|
||||
"once_cell",
|
||||
"ring",
|
||||
"rustls-pki-types",
|
||||
@@ -6415,6 +6460,7 @@ version = "0.103.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d7df23109aa6c1567d1c575b9952556388da57401e4ace1d15f79eedad0d8f53"
|
||||
dependencies = [
|
||||
"aws-lc-rs",
|
||||
"ring",
|
||||
"rustls-pki-types",
|
||||
"untrusted",
|
||||
|
||||
@@ -667,6 +667,7 @@ fn derive_config_from_params(
|
||||
experimental_client_tools: None,
|
||||
compact_prompt_override: None,
|
||||
compact_prompt_override_file: None,
|
||||
wire_api: None,
|
||||
};
|
||||
|
||||
let cli_overrides = cli_overrides
|
||||
|
||||
@@ -73,6 +73,7 @@ async fn run_command_under_sandbox(
|
||||
code_linux_sandbox_exe,
|
||||
compact_prompt_override: None,
|
||||
compact_prompt_override_file: None,
|
||||
wire_api: None,
|
||||
..Default::default()
|
||||
},
|
||||
)?;
|
||||
|
||||
@@ -84,6 +84,7 @@ async fn run_llm_request(
|
||||
model: Some(model.clone()),
|
||||
compact_prompt_override: None,
|
||||
compact_prompt_override_file: None,
|
||||
wire_api: None,
|
||||
..ConfigOverrides::default()
|
||||
}
|
||||
} else {
|
||||
|
||||
@@ -87,6 +87,11 @@ tokio = { workspace = true, features = [
|
||||
"signal",
|
||||
] }
|
||||
tokio-tungstenite = { version = "0.23", features = ["rustls-tls-webpki-roots"] }
|
||||
|
||||
# TLS 1.3 for native ZAP binary transport (zaps:// scheme)
|
||||
tokio-rustls = "0.26"
|
||||
rustls = { version = "0.23", default-features = false, features = ["std", "tls12", "ring"] }
|
||||
webpki-roots = "0.26"
|
||||
tokio-util = { workspace = true }
|
||||
crc32fast = { workspace = true }
|
||||
tokio-stream = { workspace = true }
|
||||
|
||||
@@ -444,6 +444,94 @@ impl ModelClient {
|
||||
match self.provider.wire_api {
|
||||
WireApi::Responses => self.stream_responses(prompt, log_tag).await,
|
||||
WireApi::ResponsesWebsocket => self.stream_responses_websocket(prompt, log_tag).await,
|
||||
// Native ZAP binary transport to zap.hanzo.ai:9651.
|
||||
// Uses luxfi/zap wire protocol: TLS connect → handshake →
|
||||
// MsgType 100 Call (method="chat.completions") → parse response.
|
||||
WireApi::Zap => {
|
||||
debug!("ZAP transport: native binary protocol");
|
||||
|
||||
let model_slug = prompt
|
||||
.model_override
|
||||
.as_deref()
|
||||
.unwrap_or(self.config.model.as_str());
|
||||
|
||||
// Get auth token
|
||||
let auth_manager = self.auth_manager.clone();
|
||||
let auth = auth_manager.as_ref().and_then(|m| m.auth());
|
||||
let token = match auth.as_ref() {
|
||||
Some(a) => a
|
||||
.get_token()
|
||||
.await
|
||||
.map_err(|e| CodeErr::ServerError(format!("ZAP auth: {e}")))?,
|
||||
None => {
|
||||
return Err(CodeErr::ServerError(
|
||||
AUTH_REQUIRED_MESSAGE.to_string(),
|
||||
))
|
||||
}
|
||||
};
|
||||
|
||||
// Build messages array from prompt input
|
||||
let effective_family = prompt
|
||||
.model_family_override
|
||||
.as_ref()
|
||||
.unwrap_or(&self.config.model_family);
|
||||
let full_instructions = prompt.get_full_instructions(effective_family);
|
||||
let mut messages = vec![serde_json::json!({
|
||||
"role": "system",
|
||||
"content": full_instructions
|
||||
})];
|
||||
|
||||
let input = prompt.get_formatted_input();
|
||||
for item in &input {
|
||||
match item {
|
||||
hanzo_protocol::models::ResponseItem::Message {
|
||||
role, content, ..
|
||||
} => {
|
||||
let text_parts: Vec<&str> = content
|
||||
.iter()
|
||||
.filter_map(|c| match c {
|
||||
hanzo_protocol::models::ContentItem::InputText {
|
||||
text,
|
||||
} => Some(text.as_str()),
|
||||
hanzo_protocol::models::ContentItem::OutputText {
|
||||
text,
|
||||
..
|
||||
} => Some(text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
if !text_parts.is_empty() {
|
||||
messages.push(serde_json::json!({
|
||||
"role": role,
|
||||
"content": text_parts.join("\n")
|
||||
}));
|
||||
}
|
||||
}
|
||||
hanzo_protocol::models::ResponseItem::CompactionSummary {
|
||||
encrypted_content,
|
||||
..
|
||||
} => {
|
||||
messages.push(serde_json::json!({
|
||||
"role": "assistant",
|
||||
"content": encrypted_content
|
||||
}));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// ZAP endpoint from env (default: zap.hanzo.ai:9651)
|
||||
let endpoint = std::env::var("ZAP_ENDPOINT").ok();
|
||||
|
||||
crate::zap_client::stream_zap(
|
||||
model_slug,
|
||||
serde_json::Value::Array(messages),
|
||||
&token,
|
||||
endpoint.as_deref(),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
WireApi::Chat => {
|
||||
let effective_family = prompt
|
||||
.model_family_override
|
||||
|
||||
@@ -852,6 +852,7 @@ pub struct ConfigOverrides {
|
||||
pub experimental_client_tools: Option<ClientTools>,
|
||||
pub compact_prompt_override: Option<String>,
|
||||
pub compact_prompt_override_file: Option<PathBuf>,
|
||||
pub wire_api: Option<crate::WireApi>,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
@@ -889,6 +890,7 @@ impl Config {
|
||||
experimental_client_tools,
|
||||
compact_prompt_override,
|
||||
compact_prompt_override_file,
|
||||
wire_api,
|
||||
} = overrides;
|
||||
|
||||
if let Some(mcp_servers) = mcp_servers {
|
||||
@@ -929,7 +931,7 @@ impl Config {
|
||||
.or(config_profile.model_provider)
|
||||
.or(cfg.model_provider)
|
||||
.unwrap_or_else(|| "hanzo".to_string());
|
||||
let model_provider = model_providers
|
||||
let mut model_provider = model_providers
|
||||
.get(&model_provider_id)
|
||||
.ok_or_else(|| {
|
||||
std::io::Error::new(
|
||||
@@ -939,6 +941,11 @@ impl Config {
|
||||
})?
|
||||
.clone();
|
||||
|
||||
// Apply wire API override from CLI flag.
|
||||
if let Some(api) = wire_api {
|
||||
model_provider.wire_api = api;
|
||||
}
|
||||
|
||||
// Capture workspace-write details early to avoid borrow after partial moves
|
||||
let cfg_workspace = cfg.sandbox_workspace_write.clone();
|
||||
|
||||
@@ -2617,6 +2624,7 @@ model_verbosity = "high"
|
||||
let overrides = ConfigOverrides {
|
||||
cwd: Some(fixture.cwd()),
|
||||
compact_prompt_override_file: Some(file.path().to_path_buf()),
|
||||
wire_api: None,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
|
||||
@@ -20,6 +20,8 @@ mod bridge_client;
|
||||
mod chat_completions;
|
||||
mod client;
|
||||
mod client_common;
|
||||
mod zap_client;
|
||||
mod zap_wire;
|
||||
mod code_conversation;
|
||||
pub mod codex;
|
||||
pub mod token_data;
|
||||
|
||||
@@ -44,6 +44,11 @@ pub enum WireApi {
|
||||
/// Regular Chat Completions compatible with `/v1/chat/completions`.
|
||||
#[default]
|
||||
Chat,
|
||||
|
||||
/// Native ZAP binary transport over TLS 1.3+PQ.
|
||||
/// Connects to zap.hanzo.ai:9651 (or ZAP_ENDPOINT env) and sends
|
||||
/// chat completions as MsgType 100 (native cloud service) messages.
|
||||
Zap,
|
||||
}
|
||||
|
||||
/// Serializable representation of a provider definition.
|
||||
@@ -321,6 +326,11 @@ impl ModelProviderInfo {
|
||||
format!("{base_url}/responses{query_string}")
|
||||
}
|
||||
WireApi::Chat => format!("{base_url}/chat/completions{query_string}"),
|
||||
WireApi::Zap => {
|
||||
// ZAP transport uses a separate endpoint; the URL here is used
|
||||
// for fallback/display only. Actual connection goes to ZAP_ENDPOINT.
|
||||
format!("{base_url}/chat/completions{query_string}")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -483,6 +493,7 @@ fn wire_api_override_from_env(env_key: &str) -> Option<WireApi> {
|
||||
"chat" => Some(WireApi::Chat),
|
||||
"responses" => Some(WireApi::Responses),
|
||||
"responses_websocket" => Some(WireApi::ResponsesWebsocket),
|
||||
"zap" => Some(WireApi::Zap),
|
||||
other if !other.is_empty() => {
|
||||
tracing::warn!(
|
||||
"Ignoring unknown {env_key} value '{other}'; falling back to default wire API"
|
||||
@@ -611,7 +622,8 @@ pub fn built_in_model_providers() -> HashMap<String, ModelProviderInfo> {
|
||||
"Log in with `dev /login` or set HANZO_API_KEY from https://hanzo.ai".into(),
|
||||
),
|
||||
experimental_bearer_token: None,
|
||||
wire_api: WireApi::Chat,
|
||||
wire_api: wire_api_override_from_env("HANZO_WIRE_API")
|
||||
.unwrap_or(WireApi::Chat),
|
||||
query_params: None,
|
||||
http_headers: None,
|
||||
env_http_headers: None,
|
||||
|
||||
@@ -0,0 +1,302 @@
|
||||
//! ZAP cloud service client for the dev CLI.
|
||||
//!
|
||||
//! Connects to `zap.hanzo.ai:9651` via TLS 1.3, performs the luxfi/zap
|
||||
//! handshake, and sends MsgType 100 cloud service requests (chat.completions).
|
||||
//!
|
||||
//! This is the real native ZAP binary transport — no HTTP, no fallback.
|
||||
|
||||
use crate::client_common::ResponseEvent;
|
||||
use crate::client_common::ResponseStream;
|
||||
use crate::error::CodeErr;
|
||||
use crate::error::Result;
|
||||
use crate::protocol::TokenUsage;
|
||||
use crate::zap_wire::*;
|
||||
use hanzo_protocol::models::ContentItem;
|
||||
use hanzo_protocol::models::ResponseItem;
|
||||
use serde::Deserialize;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicU32;
|
||||
use std::sync::atomic::Ordering;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::debug;
|
||||
use tracing::warn;
|
||||
|
||||
const DEFAULT_ZAP_ENDPOINT: &str = "zap.hanzo.ai:9651";
|
||||
const CLIENT_NODE_ID: &str = "dev-client";
|
||||
|
||||
// ── OpenAI-compatible response types ────────────────────────────────────
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ChatCompletionResponse {
|
||||
id: String,
|
||||
model: Option<String>,
|
||||
choices: Vec<ChatChoice>,
|
||||
usage: Option<ChatUsage>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ChatChoice {
|
||||
index: Option<u32>,
|
||||
message: ChatMessage,
|
||||
finish_reason: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ChatMessage {
|
||||
role: Option<String>,
|
||||
content: Option<String>,
|
||||
tool_calls: Option<Vec<serde_json::Value>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ChatUsage {
|
||||
prompt_tokens: Option<u32>,
|
||||
completion_tokens: Option<u32>,
|
||||
total_tokens: Option<u32>,
|
||||
}
|
||||
|
||||
// ── ZAP Cloud Client ────────────────────────────────────────────────────
|
||||
|
||||
/// A client that speaks the luxfi/zap binary wire protocol over TLS.
|
||||
pub(crate) struct ZapCloudClient {
|
||||
stream: tokio_rustls::client::TlsStream<TcpStream>,
|
||||
req_id: AtomicU32,
|
||||
#[allow(dead_code)]
|
||||
peer_id: String,
|
||||
}
|
||||
|
||||
impl ZapCloudClient {
|
||||
/// Connect to the ZAP endpoint via TLS 1.3, perform handshake.
|
||||
pub async fn connect(endpoint: Option<&str>) -> Result<Self> {
|
||||
let addr = endpoint.unwrap_or(DEFAULT_ZAP_ENDPOINT);
|
||||
|
||||
// Parse host from "host:port"
|
||||
let host = addr
|
||||
.split(':')
|
||||
.next()
|
||||
.unwrap_or("zap.hanzo.ai")
|
||||
.to_string();
|
||||
|
||||
// TLS configuration
|
||||
let mut root_store = rustls::RootCertStore::empty();
|
||||
root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
|
||||
|
||||
let config = rustls::ClientConfig::builder()
|
||||
.with_root_certificates(root_store)
|
||||
.with_no_client_auth();
|
||||
|
||||
let connector = tokio_rustls::TlsConnector::from(Arc::new(config));
|
||||
let server_name = rustls::pki_types::ServerName::try_from(host.clone())
|
||||
.map_err(|e| CodeErr::ServerError(format!("ZAP: invalid server name '{host}': {e}")))?;
|
||||
|
||||
let tcp = TcpStream::connect(addr)
|
||||
.await
|
||||
.map_err(|e| CodeErr::ServerError(format!("ZAP: TCP connect to {addr} failed: {e}")))?;
|
||||
tcp.set_nodelay(true).ok();
|
||||
|
||||
let mut tls = connector
|
||||
.connect(server_name, tcp)
|
||||
.await
|
||||
.map_err(|e| CodeErr::ServerError(format!("ZAP: TLS handshake with {addr} failed: {e}")))?;
|
||||
|
||||
debug!("ZAP: TLS 1.3 connected to {}", addr);
|
||||
|
||||
// Send handshake
|
||||
let handshake_bytes = build_handshake(CLIENT_NODE_ID);
|
||||
write_frame(&mut tls, &handshake_bytes)
|
||||
.await
|
||||
.map_err(|e| CodeErr::ServerError(format!("ZAP: handshake send failed: {e}")))?;
|
||||
|
||||
// Read handshake response
|
||||
let resp_data = read_frame(&mut tls)
|
||||
.await
|
||||
.map_err(|e| CodeErr::ServerError(format!("ZAP: handshake read failed: {e}")))?;
|
||||
|
||||
let resp_msg = Message::parse(resp_data)
|
||||
.map_err(|e| CodeErr::ServerError(format!("ZAP: handshake parse failed: {e}")))?;
|
||||
|
||||
let peer_id = parse_handshake(&resp_msg);
|
||||
debug!("ZAP: handshake complete, peer={}", peer_id);
|
||||
|
||||
Ok(Self {
|
||||
stream: tls,
|
||||
req_id: AtomicU32::new(0),
|
||||
peer_id,
|
||||
})
|
||||
}
|
||||
|
||||
/// Send a MsgType 100 Call request and receive the response.
|
||||
async fn call(
|
||||
&mut self,
|
||||
method: &str,
|
||||
auth: &str,
|
||||
body: &[u8],
|
||||
) -> Result<(u32, Vec<u8>, String)> {
|
||||
let req_id = self.req_id.fetch_add(1, Ordering::Relaxed) + 1;
|
||||
|
||||
// Build the ZAP message
|
||||
let msg_bytes = build_cloud_request(method, auth, body);
|
||||
|
||||
// Wrap with 8-byte Call correlation header
|
||||
let mut wrapped = Vec::with_capacity(8 + msg_bytes.len());
|
||||
wrapped.extend_from_slice(&req_id.to_le_bytes());
|
||||
wrapped.extend_from_slice(&REQ_FLAG_REQ.to_le_bytes());
|
||||
wrapped.extend_from_slice(&msg_bytes);
|
||||
|
||||
// Send
|
||||
write_frame(&mut self.stream, &wrapped)
|
||||
.await
|
||||
.map_err(|e| CodeErr::ServerError(format!("ZAP: send failed: {e}")))?;
|
||||
|
||||
// Read response — loop until we match our reqID
|
||||
loop {
|
||||
let data = read_frame(&mut self.stream)
|
||||
.await
|
||||
.map_err(|e| CodeErr::ServerError(format!("ZAP: recv failed: {e}")))?;
|
||||
|
||||
if data.len() < 8 {
|
||||
continue;
|
||||
}
|
||||
|
||||
let resp_flag = u32::from_le_bytes([data[4], data[5], data[6], data[7]]);
|
||||
if resp_flag != REQ_FLAG_RESP {
|
||||
continue;
|
||||
}
|
||||
|
||||
let resp_id = u32::from_le_bytes([data[0], data[1], data[2], data[3]]);
|
||||
if resp_id != req_id {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Parse the ZAP message (skip 8-byte Call header)
|
||||
let msg = Message::parse(data[8..].to_vec())
|
||||
.map_err(|e| CodeErr::ServerError(format!("ZAP: response parse failed: {e}")))?;
|
||||
|
||||
return Ok(parse_cloud_response(&msg));
|
||||
}
|
||||
}
|
||||
|
||||
/// Gracefully shut down the TLS connection.
|
||||
pub async fn close(mut self) {
|
||||
self.stream.shutdown().await.ok();
|
||||
}
|
||||
}
|
||||
|
||||
// ── Public API: stream_zap ──────────────────────────────────────────────
|
||||
|
||||
/// Execute a chat completion via native ZAP binary transport.
|
||||
///
|
||||
/// Connects to the ZAP endpoint, sends a MsgType 100 request with
|
||||
/// method="chat.completions", parses the response, and emits
|
||||
/// ResponseEvents through the returned stream.
|
||||
pub(crate) async fn stream_zap(
|
||||
model: &str,
|
||||
messages_json: serde_json::Value,
|
||||
auth_token: &str,
|
||||
endpoint: Option<&str>,
|
||||
tools: Option<&serde_json::Value>,
|
||||
) -> Result<ResponseStream> {
|
||||
// Build the chat completion request body
|
||||
let mut request_body = serde_json::json!({
|
||||
"model": model,
|
||||
"messages": messages_json,
|
||||
});
|
||||
if let Some(tools_val) = tools {
|
||||
request_body["tools"] = tools_val.clone();
|
||||
}
|
||||
|
||||
let body_bytes = serde_json::to_vec(&request_body).map_err(|e| {
|
||||
CodeErr::ServerError(format!("ZAP: failed to serialize request: {e}"))
|
||||
})?;
|
||||
|
||||
let auth = if auth_token.starts_with("Bearer ") {
|
||||
auth_token.to_string()
|
||||
} else {
|
||||
format!("Bearer {}", auth_token)
|
||||
};
|
||||
|
||||
// Connect and call
|
||||
let mut client = ZapCloudClient::connect(endpoint).await?;
|
||||
|
||||
let (status, resp_body, error_msg) = client
|
||||
.call("chat.completions", &auth, &body_bytes)
|
||||
.await?;
|
||||
|
||||
client.close().await;
|
||||
|
||||
// Handle error responses
|
||||
if status != 200 {
|
||||
let err = if !error_msg.is_empty() {
|
||||
error_msg
|
||||
} else if !resp_body.is_empty() {
|
||||
String::from_utf8_lossy(&resp_body).into_owned()
|
||||
} else {
|
||||
format!("ZAP cloud returned status {status}")
|
||||
};
|
||||
return Err(CodeErr::ServerError(format!("ZAP: {err}")));
|
||||
}
|
||||
|
||||
// Parse the chat completion response
|
||||
let completion: ChatCompletionResponse = serde_json::from_slice(&resp_body).map_err(|e| {
|
||||
warn!(
|
||||
"ZAP: failed to parse response body: {} — raw: {}",
|
||||
e,
|
||||
String::from_utf8_lossy(&resp_body)
|
||||
);
|
||||
CodeErr::ServerError(format!("ZAP: invalid response JSON: {e}"))
|
||||
})?;
|
||||
|
||||
// Convert to ResponseEvents and send through channel
|
||||
let (tx, rx) = mpsc::channel::<Result<ResponseEvent>>(16);
|
||||
|
||||
tokio::spawn(async move {
|
||||
// Created event
|
||||
if tx.send(Ok(ResponseEvent::Created)).await.is_err() {
|
||||
return;
|
||||
}
|
||||
|
||||
// Emit each choice as an OutputItemDone
|
||||
for choice in &completion.choices {
|
||||
if let Some(content) = &choice.message.content {
|
||||
let item = ResponseItem::Message {
|
||||
id: Some(completion.id.clone()),
|
||||
role: "assistant".to_string(),
|
||||
content: vec![ContentItem::OutputText {
|
||||
text: content.clone(),
|
||||
}],
|
||||
};
|
||||
if tx
|
||||
.send(Ok(ResponseEvent::OutputItemDone {
|
||||
item,
|
||||
sequence_number: None,
|
||||
output_index: choice.index,
|
||||
}))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Completed event with token usage
|
||||
let token_usage = completion.usage.map(|u| TokenUsage {
|
||||
input_tokens: u64::from(u.prompt_tokens.unwrap_or(0)),
|
||||
output_tokens: u64::from(u.completion_tokens.unwrap_or(0)),
|
||||
total_tokens: u64::from(u.total_tokens.unwrap_or(0)),
|
||||
cached_input_tokens: 0,
|
||||
reasoning_output_tokens: 0,
|
||||
});
|
||||
|
||||
let _ = tx
|
||||
.send(Ok(ResponseEvent::Completed {
|
||||
response_id: completion.id.clone(),
|
||||
token_usage,
|
||||
}))
|
||||
.await;
|
||||
});
|
||||
|
||||
Ok(ResponseStream { rx_event: rx })
|
||||
}
|
||||
@@ -0,0 +1,481 @@
|
||||
//! luxfi/zap wire protocol implementation (Rust).
|
||||
//!
|
||||
//! Compatible with the Go `github.com/luxfi/zap` v0.2.0 binary format.
|
||||
//! Implements builder, parser, and frame I/O for the Zero-copy Application
|
||||
//! Protocol used by Hanzo cloud services.
|
||||
//!
|
||||
//! Wire format:
|
||||
//! Frame: [4-byte LE length][message bytes]
|
||||
//! Message header (16 bytes): magic(4) + version(2) + flags(2) + root_offset(4) + size(4)
|
||||
//! Object fields: inline primitives, (relOffset:u32 + length:u32) for text/bytes
|
||||
|
||||
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
|
||||
|
||||
// ── Constants ───────────────────────────────────────────────────────────
|
||||
|
||||
pub const ZAP_MAGIC: [u8; 4] = *b"ZAP\x00";
|
||||
pub const HEADER_SIZE: usize = 16;
|
||||
pub const VERSION: u16 = 1;
|
||||
pub const ALIGNMENT: usize = 8;
|
||||
pub const MAX_MESSAGE_SIZE: usize = 10 * 1024 * 1024; // 10 MB
|
||||
|
||||
/// Cloud service (native binary RPC).
|
||||
pub const MSG_TYPE_CLOUD: u16 = 100;
|
||||
|
||||
// ── Cloud request field byte offsets ────────────────────────────────────
|
||||
// Layout: method(0:Text) + auth(8:Text) + body(16:Bytes)
|
||||
// Each Text/Bytes field is 8 bytes (4 relOffset + 4 length).
|
||||
|
||||
pub const CLOUD_REQ_METHOD: usize = 0;
|
||||
pub const CLOUD_REQ_AUTH: usize = 8;
|
||||
pub const CLOUD_REQ_BODY: usize = 16;
|
||||
pub const CLOUD_REQ_FIXED_SIZE: usize = 24;
|
||||
|
||||
// ── Cloud response field byte offsets ───────────────────────────────────
|
||||
// Layout: status(0:Uint32) + body(4:Bytes) + error(12:Text)
|
||||
|
||||
pub const CLOUD_RESP_STATUS: usize = 0;
|
||||
pub const CLOUD_RESP_BODY: usize = 4;
|
||||
pub const CLOUD_RESP_ERROR: usize = 12;
|
||||
|
||||
// ── Call correlation ────────────────────────────────────────────────────
|
||||
|
||||
pub const REQ_FLAG_REQ: u32 = 1;
|
||||
pub const REQ_FLAG_RESP: u32 = 2;
|
||||
|
||||
// ── Handshake ───────────────────────────────────────────────────────────
|
||||
|
||||
pub const HANDSHAKE_OBJ_SIZE: usize = 64;
|
||||
pub const HANDSHAKE_ID_MAX: usize = 60;
|
||||
pub const HANDSHAKE_ID_LEN_OFFSET: usize = 60;
|
||||
|
||||
// ── Message (parsed, owns bytes) ────────────────────────────────────────
|
||||
|
||||
/// A parsed ZAP message that owns its byte buffer.
|
||||
pub struct Message {
|
||||
data: Vec<u8>,
|
||||
}
|
||||
|
||||
impl Message {
|
||||
/// Parse a ZAP message from raw bytes. Validates magic and version.
|
||||
pub fn parse(data: Vec<u8>) -> Result<Self, &'static str> {
|
||||
if data.len() < HEADER_SIZE {
|
||||
return Err("buffer too small for ZAP header");
|
||||
}
|
||||
if data[0..4] != ZAP_MAGIC {
|
||||
return Err("invalid ZAP magic bytes");
|
||||
}
|
||||
let version = u16::from_le_bytes([data[4], data[5]]);
|
||||
if version != VERSION {
|
||||
return Err("unsupported ZAP version");
|
||||
}
|
||||
Ok(Self { data })
|
||||
}
|
||||
|
||||
/// Raw message bytes.
|
||||
pub fn bytes(&self) -> &[u8] {
|
||||
&self.data
|
||||
}
|
||||
|
||||
/// Message flags (bits 8-15 = message type, bits 0-7 = compression etc).
|
||||
pub fn flags(&self) -> u16 {
|
||||
u16::from_le_bytes([self.data[6], self.data[7]])
|
||||
}
|
||||
|
||||
/// Message type extracted from flags (upper 8 bits).
|
||||
pub fn msg_type(&self) -> u16 {
|
||||
self.flags() >> 8
|
||||
}
|
||||
|
||||
/// Root object of this message.
|
||||
pub fn root(&self) -> Object<'_> {
|
||||
let offset = u32::from_le_bytes([
|
||||
self.data[8],
|
||||
self.data[9],
|
||||
self.data[10],
|
||||
self.data[11],
|
||||
]) as usize;
|
||||
Object {
|
||||
data: &self.data,
|
||||
offset,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Object (zero-copy reader) ───────────────────────────────────────────
|
||||
|
||||
/// Zero-copy view into a ZAP object's fields.
|
||||
pub struct Object<'a> {
|
||||
data: &'a [u8],
|
||||
offset: usize,
|
||||
}
|
||||
|
||||
impl<'a> Object<'a> {
|
||||
/// Read a uint8 at the given byte offset within this object.
|
||||
pub fn uint8(&self, field_offset: usize) -> u8 {
|
||||
let pos = self.offset + field_offset;
|
||||
if pos >= self.data.len() {
|
||||
return 0;
|
||||
}
|
||||
self.data[pos]
|
||||
}
|
||||
|
||||
/// Read a uint32 at the given byte offset within this object.
|
||||
pub fn uint32(&self, field_offset: usize) -> u32 {
|
||||
let pos = self.offset + field_offset;
|
||||
if pos + 4 > self.data.len() {
|
||||
return 0;
|
||||
}
|
||||
u32::from_le_bytes([
|
||||
self.data[pos],
|
||||
self.data[pos + 1],
|
||||
self.data[pos + 2],
|
||||
self.data[pos + 3],
|
||||
])
|
||||
}
|
||||
|
||||
/// Read a bytes/text field at the given byte offset.
|
||||
/// Field layout: relOffset(4) + length(4) = 8 bytes total.
|
||||
/// Data lives at (field_pos + relOffset) for `length` bytes.
|
||||
pub fn bytes_field(&self, field_offset: usize) -> &'a [u8] {
|
||||
let pos = self.offset + field_offset;
|
||||
if pos + 4 > self.data.len() {
|
||||
return &[];
|
||||
}
|
||||
let rel_offset = i32::from_le_bytes([
|
||||
self.data[pos],
|
||||
self.data[pos + 1],
|
||||
self.data[pos + 2],
|
||||
self.data[pos + 3],
|
||||
]);
|
||||
if rel_offset == 0 {
|
||||
return &[];
|
||||
}
|
||||
let len_pos = pos + 4;
|
||||
if len_pos + 4 > self.data.len() {
|
||||
return &[];
|
||||
}
|
||||
let length = u32::from_le_bytes([
|
||||
self.data[len_pos],
|
||||
self.data[len_pos + 1],
|
||||
self.data[len_pos + 2],
|
||||
self.data[len_pos + 3],
|
||||
]) as usize;
|
||||
let abs_pos = (pos as i64 + rel_offset as i64) as usize;
|
||||
if abs_pos + length > self.data.len() {
|
||||
return &[];
|
||||
}
|
||||
&self.data[abs_pos..abs_pos + length]
|
||||
}
|
||||
|
||||
/// Read a text field at the given byte offset (UTF-8).
|
||||
pub fn text(&self, field_offset: usize) -> &'a str {
|
||||
let b = self.bytes_field(field_offset);
|
||||
std::str::from_utf8(b).unwrap_or("")
|
||||
}
|
||||
}
|
||||
|
||||
// ── Builder ─────────────────────────────────────────────────────────────
|
||||
|
||||
/// Builds ZAP messages compatible with the Go luxfi/zap library.
|
||||
pub struct Builder {
|
||||
buf: Vec<u8>,
|
||||
pos: usize,
|
||||
root_offset: usize,
|
||||
}
|
||||
|
||||
impl Builder {
|
||||
/// Create a new builder with the given capacity hint.
|
||||
pub fn new(capacity: usize) -> Self {
|
||||
let cap = capacity.max(256);
|
||||
let mut buf = vec![0u8; cap];
|
||||
buf[0..4].copy_from_slice(&ZAP_MAGIC);
|
||||
buf[4..6].copy_from_slice(&VERSION.to_le_bytes());
|
||||
Self {
|
||||
buf,
|
||||
pos: HEADER_SIZE,
|
||||
root_offset: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn grow(&mut self, n: usize) {
|
||||
let needed = self.pos + n;
|
||||
if needed <= self.buf.len() {
|
||||
return;
|
||||
}
|
||||
let new_cap = (self.buf.len() * 2).max(needed);
|
||||
self.buf.resize(new_cap, 0);
|
||||
}
|
||||
|
||||
fn align(&mut self, alignment: usize) {
|
||||
let padding = (alignment - (self.pos % alignment)) % alignment;
|
||||
self.grow(padding);
|
||||
for _ in 0..padding {
|
||||
self.buf[self.pos] = 0;
|
||||
self.pos += 1;
|
||||
}
|
||||
}
|
||||
|
||||
/// Start building an object with the given fixed-field data size.
|
||||
pub fn start_object(&mut self, data_size: usize) -> ObjectBuilder<'_> {
|
||||
self.align(ALIGNMENT);
|
||||
ObjectBuilder {
|
||||
start_pos: self.pos,
|
||||
data_size,
|
||||
deferred: Vec::new(),
|
||||
builder: self,
|
||||
}
|
||||
}
|
||||
|
||||
/// Finalize the message, writing root offset and size into the header.
|
||||
/// Returns the complete message bytes.
|
||||
pub fn finish(mut self) -> Vec<u8> {
|
||||
self.buf[8..12].copy_from_slice(&(self.root_offset as u32).to_le_bytes());
|
||||
self.buf[12..16].copy_from_slice(&(self.pos as u32).to_le_bytes());
|
||||
self.buf.truncate(self.pos);
|
||||
self.buf
|
||||
}
|
||||
|
||||
/// Finalize with specific flags (e.g., `MSG_TYPE_CLOUD << 8`).
|
||||
pub fn finish_with_flags(mut self, flags: u16) -> Vec<u8> {
|
||||
self.buf[6..8].copy_from_slice(&flags.to_le_bytes());
|
||||
self.finish()
|
||||
}
|
||||
}
|
||||
|
||||
struct DeferredWrite {
|
||||
field_offset: usize,
|
||||
data: Vec<u8>,
|
||||
}
|
||||
|
||||
/// Builder for a single ZAP object (struct).
|
||||
pub struct ObjectBuilder<'a> {
|
||||
builder: &'a mut Builder,
|
||||
start_pos: usize,
|
||||
data_size: usize,
|
||||
deferred: Vec<DeferredWrite>,
|
||||
}
|
||||
|
||||
impl<'a> ObjectBuilder<'a> {
|
||||
fn ensure_field(&mut self, end_offset: usize) {
|
||||
let needed = self.start_pos + end_offset;
|
||||
if needed > self.builder.pos {
|
||||
self.builder.grow(needed - self.builder.pos);
|
||||
for i in self.builder.pos..needed {
|
||||
self.builder.buf[i] = 0;
|
||||
}
|
||||
self.builder.pos = needed;
|
||||
}
|
||||
}
|
||||
|
||||
/// Set a uint8 field at the given byte offset.
|
||||
pub fn set_uint8(&mut self, field_offset: usize, v: u8) {
|
||||
self.ensure_field(field_offset + 1);
|
||||
self.builder.buf[self.start_pos + field_offset] = v;
|
||||
}
|
||||
|
||||
/// Set a uint32 field at the given byte offset.
|
||||
pub fn set_uint32(&mut self, field_offset: usize, v: u32) {
|
||||
self.ensure_field(field_offset + 4);
|
||||
let pos = self.start_pos + field_offset;
|
||||
self.builder.buf[pos..pos + 4].copy_from_slice(&v.to_le_bytes());
|
||||
}
|
||||
|
||||
/// Set a bytes field at the given byte offset.
|
||||
/// Writes length immediately; defers data write to `finish()`.
|
||||
pub fn set_bytes(&mut self, field_offset: usize, data: &[u8]) {
|
||||
self.ensure_field(field_offset + 8);
|
||||
let pos = self.start_pos + field_offset;
|
||||
if data.is_empty() {
|
||||
self.builder.buf[pos..pos + 4].copy_from_slice(&0u32.to_le_bytes());
|
||||
self.builder.buf[pos + 4..pos + 8].copy_from_slice(&0u32.to_le_bytes());
|
||||
return;
|
||||
}
|
||||
// Write length now
|
||||
self.builder.buf[pos + 4..pos + 8]
|
||||
.copy_from_slice(&(data.len() as u32).to_le_bytes());
|
||||
// Defer the actual data write + offset patch
|
||||
self.deferred.push(DeferredWrite {
|
||||
field_offset,
|
||||
data: data.to_vec(),
|
||||
});
|
||||
}
|
||||
|
||||
/// Set a text (string) field at the given byte offset.
|
||||
pub fn set_text(&mut self, field_offset: usize, text: &str) {
|
||||
self.set_bytes(field_offset, text.as_bytes());
|
||||
}
|
||||
|
||||
/// Finalize this object: write deferred data, patch relative offsets.
|
||||
fn do_finish(&mut self) {
|
||||
// Ensure minimum fixed size
|
||||
self.ensure_field(self.data_size);
|
||||
// Write deferred text/bytes data after the fixed section
|
||||
for entry in self.deferred.drain(..) {
|
||||
let data_pos = self.builder.pos;
|
||||
self.builder.grow(entry.data.len());
|
||||
let start = self.builder.pos;
|
||||
self.builder.buf[start..start + entry.data.len()].copy_from_slice(&entry.data);
|
||||
self.builder.pos += entry.data.len();
|
||||
// Patch relative offset: relOffset = dataPos - fieldAbsPos
|
||||
let field_abs_pos = self.start_pos + entry.field_offset;
|
||||
let rel_offset = data_pos as i32 - field_abs_pos as i32;
|
||||
self.builder.buf[field_abs_pos..field_abs_pos + 4]
|
||||
.copy_from_slice(&(rel_offset as u32).to_le_bytes());
|
||||
}
|
||||
}
|
||||
|
||||
/// Finalize and set as the message root object.
|
||||
pub fn finish_as_root(mut self) {
|
||||
self.do_finish();
|
||||
self.builder.root_offset = self.start_pos;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Frame I/O (async) ───────────────────────────────────────────────────
|
||||
|
||||
/// Write a length-prefixed frame: [4-byte LE length][data].
|
||||
pub async fn write_frame<W: AsyncWrite + Unpin>(w: &mut W, data: &[u8]) -> std::io::Result<()> {
|
||||
let len_buf = (data.len() as u32).to_le_bytes();
|
||||
w.write_all(&len_buf).await?;
|
||||
w.write_all(data).await?;
|
||||
w.flush().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read a length-prefixed frame. Returns the raw message bytes (no length header).
|
||||
pub async fn read_frame<R: AsyncRead + Unpin>(r: &mut R) -> std::io::Result<Vec<u8>> {
|
||||
let mut len_buf = [0u8; 4];
|
||||
r.read_exact(&mut len_buf).await?;
|
||||
let length = u32::from_le_bytes(len_buf) as usize;
|
||||
if length > MAX_MESSAGE_SIZE {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
format!("ZAP message too large: {} bytes", length),
|
||||
));
|
||||
}
|
||||
let mut data = vec![0u8; length];
|
||||
r.read_exact(&mut data).await?;
|
||||
Ok(data)
|
||||
}
|
||||
|
||||
// ── Handshake helpers ───────────────────────────────────────────────────
|
||||
|
||||
/// Build a ZAP handshake message containing our node ID.
|
||||
/// Format: 64-byte object with ID bytes at offsets 0-59, length at offset 60.
|
||||
pub fn build_handshake(node_id: &str) -> Vec<u8> {
|
||||
let mut b = Builder::new(128);
|
||||
let mut obj = b.start_object(HANDSHAKE_OBJ_SIZE);
|
||||
let id_bytes = node_id.as_bytes();
|
||||
for (i, &byte) in id_bytes.iter().enumerate() {
|
||||
if i >= HANDSHAKE_ID_MAX {
|
||||
break;
|
||||
}
|
||||
obj.set_uint8(i, byte);
|
||||
}
|
||||
obj.set_uint32(
|
||||
HANDSHAKE_ID_LEN_OFFSET,
|
||||
id_bytes.len().min(HANDSHAKE_ID_MAX) as u32,
|
||||
);
|
||||
obj.finish_as_root();
|
||||
b.finish()
|
||||
}
|
||||
|
||||
/// Parse a peer's node ID from a handshake message.
|
||||
pub fn parse_handshake(msg: &Message) -> String {
|
||||
let root = msg.root();
|
||||
let id_len = root.uint32(HANDSHAKE_ID_LEN_OFFSET) as usize;
|
||||
let id_len = id_len.min(HANDSHAKE_ID_MAX);
|
||||
let mut id = Vec::with_capacity(id_len);
|
||||
for i in 0..id_len {
|
||||
id.push(root.uint8(i));
|
||||
}
|
||||
String::from_utf8_lossy(&id).into_owned()
|
||||
}
|
||||
|
||||
// ── Cloud message builders ──────────────────────────────────────────────
|
||||
|
||||
/// Build a MsgType 100 cloud service request.
|
||||
/// Layout: method(0:Text) + auth(8:Text) + body(16:Bytes)
|
||||
pub fn build_cloud_request(method: &str, auth: &str, body: &[u8]) -> Vec<u8> {
|
||||
let mut b = Builder::new(body.len() + method.len() + auth.len() + 128);
|
||||
let mut obj = b.start_object(CLOUD_REQ_FIXED_SIZE);
|
||||
obj.set_text(CLOUD_REQ_METHOD, method);
|
||||
obj.set_text(CLOUD_REQ_AUTH, auth);
|
||||
obj.set_bytes(CLOUD_REQ_BODY, body);
|
||||
obj.finish_as_root();
|
||||
b.finish_with_flags(MSG_TYPE_CLOUD << 8)
|
||||
}
|
||||
|
||||
/// Parse a MsgType 100 cloud service response.
|
||||
/// Returns (status, body_bytes, error_string).
|
||||
pub fn parse_cloud_response(msg: &Message) -> (u32, Vec<u8>, String) {
|
||||
let root = msg.root();
|
||||
let status = root.uint32(CLOUD_RESP_STATUS);
|
||||
let body = root.bytes_field(CLOUD_RESP_BODY).to_vec();
|
||||
let error = root.text(CLOUD_RESP_ERROR).to_string();
|
||||
(status, body, error)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_build_and_parse_handshake() {
|
||||
let msg_bytes = build_handshake("dev-client");
|
||||
let msg = Message::parse(msg_bytes).expect("parse handshake");
|
||||
let peer_id = parse_handshake(&msg);
|
||||
assert_eq!(peer_id, "dev-client");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_and_parse_cloud_request() {
|
||||
let body = br#"{"model":"claude-sonnet-4-20250514","messages":[{"role":"user","content":"hello"}]}"#;
|
||||
let msg_bytes = build_cloud_request("chat.completions", "Bearer sk-test", body);
|
||||
let msg = Message::parse(msg_bytes).expect("parse cloud request");
|
||||
|
||||
assert_eq!(msg.msg_type(), MSG_TYPE_CLOUD);
|
||||
|
||||
let root = msg.root();
|
||||
assert_eq!(root.text(CLOUD_REQ_METHOD), "chat.completions");
|
||||
assert_eq!(root.text(CLOUD_REQ_AUTH), "Bearer sk-test");
|
||||
assert_eq!(root.bytes_field(CLOUD_REQ_BODY), body.as_slice());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_and_parse_cloud_response() {
|
||||
// Simulate a response built by the Go server
|
||||
let resp_body = br#"{"id":"chatcmpl-123","choices":[{"message":{"content":"hi"}}]}"#;
|
||||
let mut b = Builder::new(256);
|
||||
let mut obj = b.start_object(20);
|
||||
obj.set_uint32(CLOUD_RESP_STATUS, 200);
|
||||
obj.set_bytes(CLOUD_RESP_BODY, resp_body);
|
||||
// error field left empty (null)
|
||||
obj.finish_as_root();
|
||||
let data = b.finish_with_flags(MSG_TYPE_CLOUD << 8);
|
||||
|
||||
let msg = Message::parse(data).expect("parse response");
|
||||
let (status, body, error) = parse_cloud_response(&msg);
|
||||
assert_eq!(status, 200);
|
||||
assert_eq!(body, resp_body);
|
||||
assert!(error.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cloud_response_with_error() {
|
||||
let mut b = Builder::new(128);
|
||||
let mut obj = b.start_object(20);
|
||||
obj.set_uint32(CLOUD_RESP_STATUS, 401);
|
||||
// body left empty
|
||||
obj.set_text(CLOUD_RESP_ERROR, "auth token required");
|
||||
obj.finish_as_root();
|
||||
let data = b.finish_with_flags(MSG_TYPE_CLOUD << 8);
|
||||
|
||||
let msg = Message::parse(data).expect("parse error response");
|
||||
let (status, body, error) = parse_cloud_response(&msg);
|
||||
assert_eq!(status, 401);
|
||||
assert!(body.is_empty());
|
||||
assert_eq!(error, "auth token required");
|
||||
}
|
||||
}
|
||||
@@ -301,6 +301,7 @@ pub async fn run_main(cli: Cli, code_linux_sandbox_exe: Option<PathBuf>) -> anyh
|
||||
experimental_client_tools: None,
|
||||
compact_prompt_override: None,
|
||||
compact_prompt_override_file: None,
|
||||
wire_api: None,
|
||||
};
|
||||
// Parse `-c` overrides.
|
||||
let cli_kv_overrides = match config_overrides.parse_overrides() {
|
||||
|
||||
@@ -272,6 +272,7 @@ impl CodexToolCallParam {
|
||||
experimental_client_tools: None,
|
||||
compact_prompt_override: None,
|
||||
compact_prompt_override_file: None,
|
||||
wire_api: None,
|
||||
};
|
||||
|
||||
let cli_overrides = cli_overrides
|
||||
|
||||
@@ -978,6 +978,7 @@ impl MessageProcessor {
|
||||
experimental_client_tools: client_tools,
|
||||
compact_prompt_override: None,
|
||||
compact_prompt_override_file: None,
|
||||
wire_api: None,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
|
||||
@@ -62,6 +62,7 @@ pub(crate) async fn handle_create_conversation(
|
||||
experimental_client_tools: None,
|
||||
compact_prompt_override: None,
|
||||
compact_prompt_override_file: None,
|
||||
wire_api: None,
|
||||
};
|
||||
|
||||
let cfg: CodexConfig = match CodexConfig::load_with_cli_overrides(cli_overrides, overrides) {
|
||||
|
||||
@@ -24,6 +24,11 @@ pub struct Cli {
|
||||
#[arg(long, short = 'm')]
|
||||
pub model: Option<String>,
|
||||
|
||||
/// Wire protocol for LLM communication.
|
||||
/// Options: chat (default), responses, responses_websocket, zap (native binary).
|
||||
#[arg(long = "wire-api", value_name = "PROTOCOL")]
|
||||
pub wire_api: Option<String>,
|
||||
|
||||
/// Convenience flag to select the local open source model provider.
|
||||
/// Equivalent to -c model_provider=oss; verifies a local Ollama server is
|
||||
/// running.
|
||||
|
||||
@@ -448,6 +448,13 @@ pub async fn run_main(
|
||||
experimental_client_tools: None,
|
||||
compact_prompt_override: cli.compact_prompt_override.clone(),
|
||||
compact_prompt_override_file: cli.compact_prompt_file.clone(),
|
||||
wire_api: cli.wire_api.as_deref().and_then(|v| match v {
|
||||
"chat" => Some(hanzo_core::WireApi::Chat),
|
||||
"responses" => Some(hanzo_core::WireApi::Responses),
|
||||
"responses_websocket" => Some(hanzo_core::WireApi::ResponsesWebsocket),
|
||||
"zap" => Some(hanzo_core::WireApi::Zap),
|
||||
_ => None,
|
||||
}),
|
||||
};
|
||||
|
||||
// Parse `-c` overrides from the CLI.
|
||||
@@ -1328,6 +1335,7 @@ mod tests {
|
||||
cwd: Some(workspace.path().to_path_buf()),
|
||||
compact_prompt_override: None,
|
||||
compact_prompt_override_file: None,
|
||||
wire_api: None,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
|
||||
@@ -49,6 +49,11 @@ hanzo-protocol = { workspace = true }
|
||||
# Utilities
|
||||
bytes = { workspace = true }
|
||||
|
||||
# TLS 1.3 support for ZAP transport (zaps:// scheme)
|
||||
tokio-rustls = "0.26"
|
||||
rustls = { version = "0.23", default-features = false, features = ["std", "tls12", "ring"] }
|
||||
webpki-roots = "0.26"
|
||||
|
||||
# Executor dependencies
|
||||
async-trait = "0.1"
|
||||
reqwest = { version = "0.12", features = ["json"] }
|
||||
|
||||
@@ -15,31 +15,111 @@ use crate::error::Result;
|
||||
use crate::message::MessageType;
|
||||
use crate::wire::MAX_MESSAGE_SIZE;
|
||||
use crate::wire::parse_frame;
|
||||
use std::sync::Arc;
|
||||
use tokio::io::AsyncRead;
|
||||
use tokio::io::AsyncReadExt;
|
||||
use tokio::io::AsyncWrite;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::net::TcpStream;
|
||||
use tracing::debug;
|
||||
use tracing::trace;
|
||||
|
||||
/// A ZAP transport connection.
|
||||
/// Inner stream — either plain TCP or TLS.
|
||||
enum StreamInner {
|
||||
Tcp(TcpStream),
|
||||
Tls(tokio_rustls::client::TlsStream<TcpStream>),
|
||||
}
|
||||
|
||||
impl AsyncRead for StreamInner {
|
||||
fn poll_read(
|
||||
self: std::pin::Pin<&mut Self>,
|
||||
cx: &mut std::task::Context<'_>,
|
||||
buf: &mut tokio::io::ReadBuf<'_>,
|
||||
) -> std::task::Poll<std::io::Result<()>> {
|
||||
match self.get_mut() {
|
||||
StreamInner::Tcp(s) => std::pin::Pin::new(s).poll_read(cx, buf),
|
||||
StreamInner::Tls(s) => std::pin::Pin::new(s).poll_read(cx, buf),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AsyncWrite for StreamInner {
|
||||
fn poll_write(
|
||||
self: std::pin::Pin<&mut Self>,
|
||||
cx: &mut std::task::Context<'_>,
|
||||
buf: &[u8],
|
||||
) -> std::task::Poll<std::io::Result<usize>> {
|
||||
match self.get_mut() {
|
||||
StreamInner::Tcp(s) => std::pin::Pin::new(s).poll_write(cx, buf),
|
||||
StreamInner::Tls(s) => std::pin::Pin::new(s).poll_write(cx, buf),
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_flush(
|
||||
self: std::pin::Pin<&mut Self>,
|
||||
cx: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<std::io::Result<()>> {
|
||||
match self.get_mut() {
|
||||
StreamInner::Tcp(s) => std::pin::Pin::new(s).poll_flush(cx),
|
||||
StreamInner::Tls(s) => std::pin::Pin::new(s).poll_flush(cx),
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_shutdown(
|
||||
self: std::pin::Pin<&mut Self>,
|
||||
cx: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<std::io::Result<()>> {
|
||||
match self.get_mut() {
|
||||
StreamInner::Tcp(s) => std::pin::Pin::new(s).poll_shutdown(cx),
|
||||
StreamInner::Tls(s) => std::pin::Pin::new(s).poll_shutdown(cx),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A ZAP transport connection. Supports both plain TCP and TLS 1.3.
|
||||
pub struct ZapTransport {
|
||||
stream: TcpStream,
|
||||
stream: StreamInner,
|
||||
read_buf: Vec<u8>,
|
||||
write_buf: Buffer,
|
||||
}
|
||||
|
||||
impl ZapTransport {
|
||||
/// Connect to a ZAP server.
|
||||
///
|
||||
/// URL schemes:
|
||||
/// - `zap://host:port` — plain TCP
|
||||
/// - `zaps://host:port` — TLS 1.3 (with PQ key exchange when available)
|
||||
pub async fn connect(url: &str) -> Result<Self> {
|
||||
let (host, port, _tls) = parse_url(url)?;
|
||||
let (host, port, tls) = parse_url(url)?;
|
||||
|
||||
debug!("connecting to {}:{}", host, port);
|
||||
let stream = TcpStream::connect(format!("{}:{}", host, port))
|
||||
debug!("connecting to {}:{} (tls={})", host, port, tls);
|
||||
let tcp_stream = TcpStream::connect(format!("{}:{}", host, port))
|
||||
.await
|
||||
.map_err(|e| Error::Connection(e.to_string()))?;
|
||||
tcp_stream.set_nodelay(true).ok();
|
||||
|
||||
// Set TCP options for low latency
|
||||
stream.set_nodelay(true).ok();
|
||||
let stream = if tls {
|
||||
let mut root_store = rustls::RootCertStore::empty();
|
||||
root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
|
||||
|
||||
let config = rustls::ClientConfig::builder()
|
||||
.with_root_certificates(root_store)
|
||||
.with_no_client_auth();
|
||||
|
||||
let connector = tokio_rustls::TlsConnector::from(Arc::new(config));
|
||||
let server_name = rustls::pki_types::ServerName::try_from(host.clone())
|
||||
.map_err(|e| Error::Connection(format!("invalid server name: {e}")))?;
|
||||
|
||||
let tls_stream = connector
|
||||
.connect(server_name, tcp_stream)
|
||||
.await
|
||||
.map_err(|e| Error::Connection(format!("TLS handshake failed: {e}")))?;
|
||||
|
||||
debug!("TLS 1.3 connection established to {}:{}", host, port);
|
||||
StreamInner::Tls(tls_stream)
|
||||
} else {
|
||||
StreamInner::Tcp(tcp_stream)
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
stream,
|
||||
|
||||
Reference in New Issue
Block a user