M0: compilable skeleton — Kigi 0.1.0 fork surgery
Hard fork of xai-org/grok-build (Apache-2.0) re-targeted as Kigi, an
unofficial Kimi Code CLI community build.
Rename & identity
- 72 xai-*/xai-grok-* crates -> kigi-* (explicit: xai-grok-pager-bin ->
kigi-bin [binary `kigi`], xai-grok-pager -> kigi-tui; rest mechanical);
ptyctl, ptyctl-cli, third_party/ unchanged; proto package
xai.grok.tools.v1 -> kigi.tools.v1
- Config home ~/.kigi (KIGI_SHARE_DIR override), env prefix GROK_* ->
KIGI_*, `kigi --version` carries the unofficial-community-build notice
- clap identity, help text, startup banner, prompt templates rebranded
(templates re-encrypted)
Deletions (PRD removal list #5/#6/#7/#9/#10)
- voice input (xai-grok-voice) and all TUI wiring
- telemetry: Mixpanel client, external OTel stream, Sentry, OTLP layers,
trace/GCS/S3 upload queues (kigi-file-utils halved), workspace upload
module & dc_log, heap-profile uploader, auth-diagnostics uploader,
session-analytics halves of feedback; local zero-egress observability
preserved in new kigi-log crate (unified log, --debug firehose,
subsystem file logs, opt-in instrumentation)
- announcements (crate, remote-settings fields, TUI surfaces)
- plugin marketplace (crate, sources/browse/CTA/extensions-modal tab);
direct plugin install/uninstall/update via kigi-agent git_install kept
- relay/gateway/assets endpoints and features (agent relay, headless
relay transport, gateway bridge, LeaderEnvUrls); leader IPC socket now
~/.kigi/leader.sock + KIGI_LEADER_SOCKET, no ws-url derivation
- functional types rehomed instead of deleted: PermissionMode ->
kigi-config-types, McpInitStrategy -> kigi-mcp, PrCreationSource ->
session signals, TerminalDiagnostics -> kigi-pager-render, agent_id ->
shell util
Endpoints
- kigi-env rewritten: single production KigiEndpoints {coding_api_base_url
https://api.kimi.com/coding/v1 (KIGI_CODE_BASE_URL), oauth_host
https://auth.kimi.com (KIGI_OAUTH_HOST), update_base_url (GitHub
Releases API), upgrade_page_url}; GrokBuildEnvironment enum deleted
Toolchain & workspace hygiene
- Rust 1.97.0 pinned; edition 2024; full cargo update; git2 hoisted to
workspace at 0.21 (Option->Result API migration), quick-xml 0.41
- Root Cargo.toml hand-maintained (PRD §8.1): version 0.1.0 inherited by
all members, members sorted, unused deps pruned
- cargo-deny advisories gate (deny.toml with documented transitive
exceptions); CI workflow (check/clippy/fmt/deny/test, macOS+Linux)
- cross-crate test seams re-gated behind `test-support` cargo feature;
insta snapshot baselines renamed to the kigi_tui prefix
- clippy --workspace --all-targets: zero warnings; fmt clean
Fixes surfaced by the port
- updater probe/installer divergence (bin/kigi vs bin/grok symlink set)
- idle model-metadata refresh dead under KIGI_CODE_BASE_URL override
(new is_effective_coding_endpoint_url, loopback+override aware)
- macOS symlinked-TMPDIR fixture canonicalization (foreign_sessions,
fast-worktree); RSS measurement tests serialized via serial_test
Docs & legal (Apache §4)
- NOTICE added (upstream attribution + change statement); THIRD-PARTY
notices sustained; kigi-tools ported-code notices extended; README,
CONTRIBUTING, SECURITY, AGENTS.md rewritten
Out of scope for M0 (tracked): Kimi auth/inference (M1), search/fetch,
command parity, config import (M2), Computer Hub excision & final
brand-token sweep (M2), distribution & self-update rewrite (M3).
This commit is contained in:
@@ -0,0 +1,487 @@
|
||||
//! WebSocket server for remote agent connections.
|
||||
//!
|
||||
//! This module provides a WebSocket server that allows remote TUI clients to
|
||||
//! connect to a grok agent running on a different machine.
|
||||
//!
|
||||
//! The agent persists across WebSocket reconnections: a single MvpAgent instance
|
||||
//! is created on first connection and reused for all subsequent connections. This
|
||||
//! ensures that session actors (and any in-flight prompts) survive client
|
||||
//! disconnects — when a client reconnects and loads an existing session, ongoing
|
||||
//! work continues to stream to the new connection.
|
||||
|
||||
use std::cell::RefCell;
|
||||
use std::net::SocketAddr;
|
||||
use std::rc::Rc;
|
||||
use std::sync::Arc;
|
||||
use std::thread;
|
||||
|
||||
use axum::{
|
||||
Router,
|
||||
extract::{
|
||||
ConnectInfo, Query, State,
|
||||
ws::{Message, WebSocket, WebSocketUpgrade},
|
||||
},
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::{IntoResponse, Response},
|
||||
routing::get,
|
||||
};
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, simplex};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::time::Duration;
|
||||
use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
|
||||
use tracing::{info, warn};
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use kigi_acp_lib::{
|
||||
AcpAgentGatewayReceiver as GatewayReceiver, AcpAgentGatewaySender as GatewaySender,
|
||||
AcpClientMessage, LineBufferedRead,
|
||||
};
|
||||
|
||||
use crate::agent::config::{Config as AgentConfig, ModelEntry};
|
||||
use crate::agent::models::{ModelFetchAuth, prefetch_models_blocking};
|
||||
use crate::agent::mvp_agent::MvpAgent;
|
||||
|
||||
use indexmap::IndexMap;
|
||||
|
||||
/// Swappable destination for the relay task.
|
||||
///
|
||||
/// Points at the current ACP connection's gateway sender. When no client is
|
||||
/// connected, the value is `None` and outbound messages are silently dropped
|
||||
/// (matching the old behaviour where the gateway channel's receiver was simply
|
||||
/// gone).
|
||||
type RelayDest = Rc<RefCell<Option<mpsc::UnboundedSender<AcpClientMessage>>>>;
|
||||
|
||||
const MAX_BUFFER_SIZE: usize = 8 * 1024 * 1024;
|
||||
const KEEPALIVE_INTERVAL_SECS: u64 = 15;
|
||||
|
||||
/// Configuration for the agent WebSocket server.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ServerConfig {
|
||||
/// Address to bind the server to
|
||||
pub bind_addr: SocketAddr,
|
||||
/// Secret token for client authentication (required)
|
||||
pub secret: String,
|
||||
}
|
||||
|
||||
/// Shared state for the WebSocket server.
|
||||
struct ServerState {
|
||||
agent_config: AgentConfig,
|
||||
secret: String,
|
||||
/// Channel to send new WebSocket connections to the persistent agent thread.
|
||||
/// Lazily initialised on first connection; protected by a tokio Mutex so the
|
||||
/// axum handler (which is `Send`) can acquire it.
|
||||
agent_conn_tx: tokio::sync::Mutex<Option<mpsc::UnboundedSender<NewConnectionChannels>>>,
|
||||
}
|
||||
|
||||
/// Channels bridging a single WebSocket connection to the agent thread.
|
||||
struct NewConnectionChannels {
|
||||
from_ws_rx: mpsc::UnboundedReceiver<String>,
|
||||
to_ws_tx: mpsc::UnboundedSender<String>,
|
||||
}
|
||||
|
||||
/// Query parameters for WebSocket connection.
|
||||
#[derive(Debug, serde::Deserialize, Default)]
|
||||
pub struct WsQueryParams {
|
||||
#[serde(rename = "server-key")]
|
||||
pub server_key: Option<String>,
|
||||
}
|
||||
|
||||
/// Validate the bearer token from request headers or query parameters.
|
||||
fn validate_auth(headers: &HeaderMap, query: &WsQueryParams, expected_secret: &str) -> bool {
|
||||
// Try Authorization header
|
||||
if let Some(token) = headers
|
||||
.get("authorization")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|v| v.strip_prefix("Bearer "))
|
||||
{
|
||||
return token == expected_secret;
|
||||
}
|
||||
|
||||
// Fall back to query parameter for browser connections
|
||||
if let Some(ref key) = query.server_key {
|
||||
return key == expected_secret;
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
|
||||
/// WebSocket upgrade handler with authentication.
|
||||
async fn ws_handler(
|
||||
ws: WebSocketUpgrade,
|
||||
State(state): State<Arc<ServerState>>,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
headers: HeaderMap,
|
||||
Query(query): Query<WsQueryParams>,
|
||||
) -> Response {
|
||||
// Validate secret token from header or query param
|
||||
if !validate_auth(&headers, &query, &state.secret) {
|
||||
warn!("Unauthorized connection attempt from {}", addr);
|
||||
return (
|
||||
StatusCode::UNAUTHORIZED,
|
||||
"Invalid or missing authorization token",
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
info!("Authenticated WebSocket connection from {}", addr);
|
||||
ws.on_upgrade(move |socket| handle_connection(socket, state, addr))
|
||||
}
|
||||
|
||||
/// Handle an authenticated WebSocket connection.
|
||||
///
|
||||
/// On first connection, spawns a persistent agent thread that owns the MvpAgent.
|
||||
/// On subsequent connections (reconnects), sends new WS channels to the existing
|
||||
/// agent thread so that session actors can continue streaming to the new client.
|
||||
async fn handle_connection(ws: WebSocket, state: Arc<ServerState>, peer_addr: SocketAddr) {
|
||||
info!("New WebSocket connection from {}", peer_addr);
|
||||
|
||||
let (mut ws_write, mut ws_read) = ws.split();
|
||||
|
||||
// Channels for bridging WS <-> Agent thread
|
||||
let (to_agent_tx, to_agent_rx) = mpsc::unbounded_channel::<String>();
|
||||
let (from_agent_tx, mut from_agent_rx) = mpsc::unbounded_channel::<String>();
|
||||
|
||||
// Ensure the persistent agent thread is running (lazy init on first connection).
|
||||
// If the previous agent thread died (panic, etc.), clear the stale sender so we
|
||||
// respawn a fresh one.
|
||||
{
|
||||
let mut agent_tx_guard = state.agent_conn_tx.lock().await;
|
||||
|
||||
// Check if existing sender is still alive (receiver not dropped)
|
||||
if let Some(ref tx) = *agent_tx_guard
|
||||
&& tx.is_closed()
|
||||
{
|
||||
warn!("Persistent agent thread died — will respawn");
|
||||
*agent_tx_guard = None;
|
||||
}
|
||||
|
||||
if agent_tx_guard.is_none() {
|
||||
let (conn_tx, conn_rx) = mpsc::unbounded_channel();
|
||||
|
||||
let agent_config = state.agent_config.clone();
|
||||
let _agent_thread = thread::Builder::new()
|
||||
.name("agent-persistent".to_string())
|
||||
.spawn(move || {
|
||||
// Prefetch models before creating the runtime (blocking is OK here)
|
||||
let auth = agent_config.create_auth_manager().current();
|
||||
let fetch_auth =
|
||||
ModelFetchAuth::resolve(&agent_config.endpoints, auth.is_some());
|
||||
let prefetched_models = if auth.is_some()
|
||||
|| agent_config.endpoints.has_custom_endpoint()
|
||||
|| fetch_auth != ModelFetchAuth::Session
|
||||
{
|
||||
prefetch_models_blocking(&agent_config.endpoints, auth.as_ref(), fetch_auth)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
info!("Prefetched models: {:?}", prefetched_models);
|
||||
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("Failed to create runtime for agent");
|
||||
|
||||
let local_set = tokio::task::LocalSet::new();
|
||||
local_set.block_on(&rt, async move {
|
||||
run_persistent_agent(agent_config, conn_rx, prefetched_models).await
|
||||
});
|
||||
|
||||
warn!("Persistent agent thread exiting");
|
||||
});
|
||||
|
||||
*agent_tx_guard = Some(conn_tx);
|
||||
info!("Persistent agent thread spawned");
|
||||
}
|
||||
|
||||
// Send new WS channels to the agent thread
|
||||
if let Some(ref tx) = *agent_tx_guard
|
||||
&& tx
|
||||
.send(NewConnectionChannels {
|
||||
from_ws_rx: to_agent_rx,
|
||||
to_ws_tx: from_agent_tx,
|
||||
})
|
||||
.is_err()
|
||||
{
|
||||
warn!("Failed to send connection channels to agent thread");
|
||||
}
|
||||
}
|
||||
|
||||
// Task: Read from WS, send to agent thread
|
||||
let read_task = tokio::spawn(async move {
|
||||
while let Some(msg) = ws_read.next().await {
|
||||
match msg {
|
||||
Ok(Message::Text(text)) => {
|
||||
let text_str: &str = text.as_ref();
|
||||
let trimmed = text_str.trim_end_matches(['\r', '\n']);
|
||||
// Skip browser keepalive pings (non-JSON text)
|
||||
if trimmed == "ping" || trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if to_agent_tx.send(trimmed.to_string()).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(Message::Binary(bin)) => {
|
||||
if let Ok(s) = std::str::from_utf8(&bin) {
|
||||
let trimmed = s.trim_end_matches(['\r', '\n']);
|
||||
if trimmed == "ping" || trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if to_agent_tx.send(trimmed.to_string()).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(Message::Close(frame)) => {
|
||||
if let Some(f) = frame {
|
||||
info!(
|
||||
"WebSocket close from {}: {} {}",
|
||||
peer_addr, f.code, f.reason
|
||||
);
|
||||
}
|
||||
break;
|
||||
}
|
||||
Ok(Message::Ping(_)) | Ok(Message::Pong(_)) => {}
|
||||
Err(e) => {
|
||||
warn!("WebSocket read error from {}: {:?}", peer_addr, e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Task: Read from agent thread, send to WS (with keepalive)
|
||||
let write_task = tokio::spawn(async move {
|
||||
let mut keepalive = tokio::time::interval(Duration::from_secs(KEEPALIVE_INTERVAL_SECS));
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
Some(msg) = from_agent_rx.recv() => {
|
||||
if ws_write.send(Message::Text(msg.into())).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
_ = keepalive.tick() => {
|
||||
if ws_write.send(Message::Ping(vec![].into())).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
else => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Wait for either task to complete
|
||||
tokio::select! {
|
||||
_ = read_task => {}
|
||||
_ = write_task => {}
|
||||
}
|
||||
|
||||
info!("WebSocket connection ended for {}", peer_addr);
|
||||
}
|
||||
|
||||
/// Run the persistent agent on a dedicated thread with LocalSet.
|
||||
///
|
||||
/// The MvpAgent is created **once** and reused across WebSocket reconnections.
|
||||
/// A persistent gateway channel ensures that session actors (which hold cloned
|
||||
/// `GatewaySender` handles) can always send notifications. A relay task forwards
|
||||
/// messages from the persistent channel to the *current* ACP connection's channel,
|
||||
/// so notifications reach whichever client is currently connected.
|
||||
async fn run_persistent_agent(
|
||||
agent_config: AgentConfig,
|
||||
mut connection_rx: mpsc::UnboundedReceiver<NewConnectionChannels>,
|
||||
prefetched_models: Option<IndexMap<String, ModelEntry>>,
|
||||
) {
|
||||
// Persistent gateway channel — the MvpAgent and all session actors hold
|
||||
// clones of `gw_tx`. This channel survives across reconnections.
|
||||
let (gw_tx, mut gw_rx) = tokio::sync::mpsc::unbounded_channel::<AcpClientMessage>();
|
||||
let gateway = GatewaySender::new(gw_tx);
|
||||
|
||||
// Create MvpAgent ONCE -- it persists for the lifetime of the server.
|
||||
let auth_manager = Arc::new(agent_config.create_auth_manager());
|
||||
// Proactive token refresh; runs until process exit.
|
||||
auth_manager.start_proactive_refresh(tokio_util::sync::CancellationToken::new());
|
||||
// Restore managed policy right before bootstrap reads it — the agent is created lazily here,
|
||||
// so an earlier restore could go stale before the gate.
|
||||
crate::managed_config::ensure_managed_policy_present(&auth_manager).await;
|
||||
let agent = Rc::new(
|
||||
MvpAgent::new(gateway, &agent_config, auth_manager, prefetched_models)
|
||||
.unwrap_or_else(crate::agent::init::exit_on_config_error),
|
||||
);
|
||||
|
||||
let relay_dest: RelayDest = Rc::new(RefCell::new(None));
|
||||
|
||||
// Relay task: reads from the persistent gateway channel and forwards to
|
||||
// whichever ACP connection is currently active.
|
||||
let relay_dest_for_task = relay_dest.clone();
|
||||
tokio::task::spawn_local(async move {
|
||||
while let Some(msg) = gw_rx.recv().await {
|
||||
let maybe_tx = relay_dest_for_task.borrow().clone();
|
||||
if let Some(tx) = maybe_tx
|
||||
&& tx.send(msg).is_err()
|
||||
{
|
||||
// Connection's gateway receiver was dropped — clear it.
|
||||
*relay_dest_for_task.borrow_mut() = None;
|
||||
}
|
||||
// If no connection, the message (and its response_tx) is dropped.
|
||||
// The caller (session actor) gets a send error which is already
|
||||
// handled with `let _ = ...`.
|
||||
}
|
||||
});
|
||||
|
||||
// Accept new connections in a loop
|
||||
while let Some(channels) = connection_rx.recv().await {
|
||||
info!("Agent thread: setting up new ACP connection (reconnect)");
|
||||
setup_acp_connection(agent.clone(), channels, relay_dest.clone());
|
||||
}
|
||||
|
||||
info!("Agent thread: connection channel closed, exiting");
|
||||
}
|
||||
|
||||
/// Set up a new ACP connection for a WebSocket connection, reusing the existing
|
||||
/// MvpAgent. The relay destination is updated so that session actor notifications
|
||||
/// flow to the new client.
|
||||
fn setup_acp_connection(
|
||||
agent: Rc<MvpAgent>,
|
||||
channels: NewConnectionChannels,
|
||||
relay_dest: RelayDest,
|
||||
) {
|
||||
let NewConnectionChannels {
|
||||
mut from_ws_rx,
|
||||
to_ws_tx,
|
||||
} = channels;
|
||||
|
||||
// Create new simplex IO streams for this ACP connection
|
||||
let (agent_read_rx, mut agent_read_tx) = simplex(MAX_BUFFER_SIZE);
|
||||
let (agent_write_rx, agent_write_tx) = simplex(MAX_BUFFER_SIZE);
|
||||
|
||||
let incoming = agent_read_rx.compat();
|
||||
let outgoing = agent_write_tx.compat_write();
|
||||
|
||||
// Create a per-connection gateway channel for the GatewayReceiver.
|
||||
// The relay task will forward persistent-channel messages here.
|
||||
let (conn_gw_tx, conn_gw_rx) = tokio::sync::mpsc::unbounded_channel::<AcpClientMessage>();
|
||||
|
||||
// Point the relay at this new connection's channel
|
||||
*relay_dest.borrow_mut() = Some(conn_gw_tx);
|
||||
|
||||
// Create new ACP connection reusing the same MvpAgent (via Rc clone).
|
||||
// `Agent` is implemented for `Rc<T: Agent>` so this works.
|
||||
let incoming = LineBufferedRead::spawn_local(incoming);
|
||||
let (conn, handle_io) = acp::AgentSideConnection::new(agent, outgoing, incoming, |fut| {
|
||||
tokio::task::spawn_local(fut);
|
||||
});
|
||||
tokio::task::spawn_local(
|
||||
GatewayReceiver::new(conn_gw_rx, conn)
|
||||
.with_on_meta(kigi_file_utils::trace_context::span_from_meta_traceparent)
|
||||
.run(),
|
||||
);
|
||||
|
||||
// Task: Forward WS messages → agent (incoming ACP bytes)
|
||||
tokio::task::spawn_local(async move {
|
||||
while let Some(msg) = from_ws_rx.recv().await {
|
||||
// Log messages that lack both `id` and `method` — the ACP layer
|
||||
// only prints "received message with neither id nor method" without
|
||||
// the payload, making debugging impossible.
|
||||
if let Ok(v) = serde_json::from_str::<serde_json::Value>(&msg)
|
||||
&& v.get("id").is_none()
|
||||
&& v.get("method").is_none()
|
||||
{
|
||||
warn!(
|
||||
len = msg.len(),
|
||||
"incoming WS message has neither id nor method"
|
||||
);
|
||||
}
|
||||
if agent_read_tx.write_all(msg.as_bytes()).await.is_err() {
|
||||
break;
|
||||
}
|
||||
if agent_read_tx.write_all(b"\n").await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
// WS disconnected — the simplex writer is dropped, causing `handle_io`
|
||||
// to complete. The GatewayReceiver for this connection will also stop.
|
||||
// But the MvpAgent and session actors stay alive, ready for the next
|
||||
// connection.
|
||||
});
|
||||
|
||||
// Task: Forward agent messages → WS (outgoing ACP bytes)
|
||||
tokio::task::spawn_local(async move {
|
||||
let mut reader = BufReader::new(agent_write_rx);
|
||||
let mut line = String::new();
|
||||
|
||||
loop {
|
||||
line.clear();
|
||||
match reader.read_line(&mut line).await {
|
||||
Ok(0) => break,
|
||||
Ok(_) => {
|
||||
let msg = line.trim_end_matches(['\r', '\n']);
|
||||
if !msg.is_empty() && to_ws_tx.send(msg.to_string()).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Run the ACP IO handler — fire-and-forget since we don't block the
|
||||
// connection loop. It completes when the WS disconnects.
|
||||
tokio::task::spawn_local(async move {
|
||||
let _ = handle_io.await;
|
||||
info!("ACP connection IO handler completed");
|
||||
});
|
||||
}
|
||||
|
||||
/// Run the agent WebSocket server.
|
||||
///
|
||||
/// This starts a WebSocket server that accepts authenticated connections from
|
||||
/// remote TUI clients. A single agent instance is shared across all connections
|
||||
/// (persisted across reconnections) so that in-flight session work survives
|
||||
/// client disconnects.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `config` - Server configuration (bind address and secret)
|
||||
/// * `agent_config` - Agent configuration to use for each connection
|
||||
///
|
||||
/// # Example
|
||||
/// ```ignore
|
||||
/// let server_config = ServerConfig {
|
||||
/// bind_addr: "0.0.0.0:9000".parse().unwrap(),
|
||||
/// secret: "my-secret-token".to_string(),
|
||||
/// };
|
||||
/// run_agent_server(server_config, agent_config).await?;
|
||||
/// ```
|
||||
pub async fn run_agent_server(
|
||||
config: ServerConfig,
|
||||
agent_config: AgentConfig,
|
||||
) -> anyhow::Result<()> {
|
||||
let state = Arc::new(ServerState {
|
||||
agent_config,
|
||||
secret: config.secret,
|
||||
agent_conn_tx: tokio::sync::Mutex::new(None),
|
||||
});
|
||||
|
||||
let app = Router::new()
|
||||
.route("/ws", get(ws_handler))
|
||||
.with_state(state);
|
||||
|
||||
let listener = TcpListener::bind(config.bind_addr).await?;
|
||||
info!("Agent server listening on ws://{}/ws", config.bind_addr);
|
||||
info!(
|
||||
"Clients should connect with: --remote ws://{}:{}/ws --secret <token>",
|
||||
config.bind_addr.ip(),
|
||||
config.bind_addr.port()
|
||||
);
|
||||
|
||||
axum::serve(
|
||||
listener,
|
||||
app.into_make_service_with_connect_info::<SocketAddr>(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user