lda/wakey
1
0
Fork 0
Code Issues 5 Pull Requests Actions Packages Projects Releases 18 Wiki Activity

advertise and enforce per-agent terminal session limits

Browse Source
This commit is contained in:
lda
2026-07-16 19:26:12 +07:00 Verified
parent 9a14410bed
commit 0e99a584a3
10 changed files with 174 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ui/src/api.ts
+3
View File
@@ -3,6 +3,9 @@ export type Agent = {
connected: boolean;
nickname?: string | null;
capabilities: "terminal"[];
capability_options?: {
terminal?: { max_sessions: number };
};
};
export type TerminalSession = {
ui/src/pages/TerminalPage.tsx
+5 -2
View File
@@ -156,7 +156,10 @@ export function TerminalPage({
const selectedAgentSessionCount = sessions.filter(
(item) => item.agent_id === selectedAgentId,
).length;
const agentAtSessionLimit = selectedAgentSessionCount >= 2;
const selectedAgentSessionLimit =
selectedAgent?.capability_options?.terminal?.max_sessions ?? 2;
const agentAtSessionLimit =
selectedAgentSessionCount >= selectedAgentSessionLimit;
const canRequestStart =
selectedAgent?.connected &&
selectedAgent.capabilities.includes("terminal") &&
@@ -744,7 +747,7 @@ export function TerminalPage({
</TooltipTrigger>
<TooltipContent>
{agentAtSessionLimit
? "Two-session limit reached"
? `${selectedAgentSessionLimit}-session limit reached`
: "New session"}
</TooltipContent>
</Tooltip>
wakey-agent/src/config.rs
+22 -1
View File
@@ -4,6 +4,8 @@ use std::fmt;
use std::path::{Path, PathBuf};
use std::sync::LazyLock;
use crate::protocol::DEFAULT_TERMINAL_MAX_SESSIONS;
pub const DEFAULT_CONFIG_PATH: &str = "/etc/wakey-agent/config.toml";
pub const DEFAULT_PID_FILE: &str = "/var/run/wakey-agent.pid";
const WAKEY_DHCP_LEASES_ENV: &str = "WAKEY_DHCP_LEASES";
@@ -40,6 +42,7 @@ pub struct AgentConfig {
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct TerminalConfig {
#[serde(default)]
pub enabled: bool,
@@ -142,7 +145,7 @@ fn default_terminal_shell() -> PathBuf {
}
const fn default_terminal_max_sessions() -> usize {
2
DEFAULT_TERMINAL_MAX_SESSIONS
}
impl AgentConfig {
@@ -312,4 +315,22 @@ max_sessions = 2
assert!(config.terminal.args.is_empty());
assert!(config.terminal.current_dir.is_none());
}
#[test]
fn terminal_config_rejects_misspelled_fields() {
let error = toml::from_str::<AgentConfig>(
r#"
server_url = "https://example.com"
agent_id = "agent-1"
agent_token = "secret"
[terminal]
enabled = true
max_session = 67
"#,
)
.expect_err("unknown terminal fields must not silently use defaults");
assert!(error.to_string().contains("max_session"));
}
}
wakey-agent/src/protocol.rs
+40
View File
@@ -65,6 +65,30 @@ pub enum AgentCapability {
Terminal,
}
pub const DEFAULT_TERMINAL_MAX_SESSIONS: usize = 2;
/// Optional parameters attached to advertised agent capabilities.
///
/// Keep this separate from `AgentCapability`: the capability list remains a
/// compact, backward-compatible feature check, while this object can grow as
/// individual capabilities gain configurable limits or modes.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct AgentCapabilityOptions {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub terminal: Option<TerminalCapabilityOptions>,
}
impl AgentCapabilityOptions {
fn is_empty(&self) -> bool {
self.terminal.is_none()
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct TerminalCapabilityOptions {
pub max_sessions: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AgentTerminalSession {
pub terminal_id: TerminalId,
@@ -238,6 +262,8 @@ pub enum ClientMessage {
agent_id: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
capabilities: Vec<AgentCapability>,
#[serde(default, skip_serializing_if = "AgentCapabilityOptions::is_empty")]
capability_options: AgentCapabilityOptions,
},
Auth {
agent_id: String,
@@ -397,4 +423,18 @@ mod tests {
let json = serde_json::to_string(&resume).expect("serialize terminal resume");
assert!(json.contains("\"type\":\"resume_terminal\""));
}
#[test]
fn hello_serializes_typed_capability_options() {
let message = ClientMessage::Hello {
agent_id: "router".into(),
capabilities: vec![AgentCapability::Terminal],
capability_options: AgentCapabilityOptions {
terminal: Some(TerminalCapabilityOptions { max_sessions: 3 }),
},
};
let value = serde_json::to_value(message).expect("serialize hello");
assert_eq!(value["capability_options"]["terminal"]["max_sessions"], 3);
}
}
wakey-agent/src/session.rs
+12 -1
View File
@@ -8,7 +8,10 @@ use tracing::{debug, error, info, info_span, warn};
use crate::config::AgentConfig;
use crate::dispatch::{dispatch_command, inventory_for_config};
use crate::protocol::{AgentCapability, AgentCommand, ClientMessage, ErrorPayload, ServerMessage};
use crate::protocol::{
AgentCapability, AgentCapabilityOptions, AgentCommand, ClientMessage, ErrorPayload,
ServerMessage, TerminalCapabilityOptions,
};
use crate::terminal::TerminalManager;
pub async fn run(config: AgentConfig) -> Result<()> {
@@ -82,6 +85,14 @@ async fn run_once(
&ClientMessage::Hello {
agent_id: config.agent_id.clone(),
capabilities: agent_capabilities(config),
capability_options: AgentCapabilityOptions {
terminal: config
.terminal
.enabled
.then_some(TerminalCapabilityOptions {
max_sessions: config.terminal.max_sessions.max(1),
}),
},
},
)
.await?;
wakey-control-plane/src/api/commands.rs
+8 -2
View File
@@ -6,8 +6,9 @@ use serde::{Deserialize, Serialize};
use std::time::Instant;
use tracing::{info, info_span, warn};
use uuid::Uuid;
use wakey_agent::protocol::AgentCapability;
use wakey_agent::protocol::{AgentCommand, ErrorPayload, RequestId, ServerMessage};
use wakey_agent::protocol::{
AgentCapability, AgentCapabilityOptions, AgentCommand, ErrorPayload, RequestId, ServerMessage,
};
use crate::api::ApiError;
use crate::runtime::{AgentReply, AppState, SessionEvent};
@@ -19,6 +20,7 @@ pub struct AgentStatus {
pub connected: bool,
pub nickname: Option<String>,
pub capabilities: Vec<AgentCapability>,
pub capability_options: AgentCapabilityOptions,
}
#[derive(Debug, Deserialize)]
@@ -49,6 +51,10 @@ pub async fn list_agents(State(state): State<AppState>) -> Result<impl IntoRespo
.get(&agent_id)
.map(|session| session.capabilities.clone())
.unwrap_or_default(),
capability_options: sessions
.get(&agent_id)
.map(|session| session.capability_options.clone())
.unwrap_or_default(),
agent_id,
nickname,
})
wakey-control-plane/src/api/terminals.rs
+12 -5
View File
@@ -7,8 +7,8 @@ use futures_util::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use tracing::{info, warn};
use wakey_agent::protocol::{
AgentCapability, ServerMessage, TerminalAgentHandshake, TerminalControl, TerminalId,
TerminalOperatorHandshake,
AgentCapability, DEFAULT_TERMINAL_MAX_SESSIONS, ServerMessage, TerminalAgentHandshake,
TerminalControl, TerminalId, TerminalOperatorHandshake,
};
use crate::api::ApiError;
@@ -52,7 +52,7 @@ pub async fn create_terminal(
Json(request): Json<CreateTerminalRequest>,
) -> Result<(StatusCode, Json<TerminalSessionResponse>), ApiError> {
validate_size(request.rows, request.cols)?;
let agent_tx = {
let (agent_tx, max_sessions) = {
let sessions = state.sessions.read().await;
let session = sessions.get(&request.agent_id).ok_or_else(|| {
ApiError::new(
@@ -68,12 +68,19 @@ pub async fn create_terminal(
"agent has not advertised terminal capability",
));
}
session.tx.clone()
let max_sessions = session
.capability_options
.terminal
.as_ref()
.map(|terminal| terminal.max_sessions)
.unwrap_or(DEFAULT_TERMINAL_MAX_SESSIONS)
.max(1);
(session.tx.clone(), max_sessions)
};
let created = state
.terminals
.create(request.agent_id.clone())
.create_with_limit(request.agent_id.clone(), max_sessions)
.await
.map_err(registry_error)?;
let terminal_id = TerminalId::new(created.terminal_id.clone()).map_err(|message| {
wakey-control-plane/src/runtime/mod.rs
+2 -1
View File
@@ -16,7 +16,7 @@ use tower_http::services::ServeFile;
use tracing::info;
#[cfg(unix)]
use tracing::warn;
use wakey_agent::protocol::{AgentCapability, ErrorPayload, ServerMessage};
use wakey_agent::protocol::{AgentCapability, AgentCapabilityOptions, ErrorPayload, ServerMessage};
use crate::api;
use crate::config;
@@ -50,6 +50,7 @@ pub struct AgentSession {
pub connection_id: String,
pub tx: mpsc::UnboundedSender<SessionEvent>,
pub capabilities: Vec<AgentCapability>,
pub capability_options: AgentCapabilityOptions,
}
#[derive(Clone)]
pub enum SessionEvent {
wakey-control-plane/src/runtime/terminals.rs
+31 -3
View File
@@ -9,7 +9,6 @@ use wakey_agent::protocol::{AgentTerminalSession, TerminalId};
pub const TERMINAL_RELAY_QUEUE: usize = 32;
pub const TERMINAL_MAX_FRAME_BYTES: usize = 64 * 1024;
pub const TERMINAL_PENDING_AGENT_BYTES: usize = 256 * 1024;
pub const TERMINAL_MAX_SESSIONS_PER_AGENT: usize = 2;
pub const TERMINAL_ATTACH_TIMEOUT: Duration = Duration::from_secs(10);
pub const TERMINAL_ABSOLUTE_TIMEOUT: Duration = Duration::from_secs(12 * 60 * 60);
const TERMINAL_TOMBSTONE_TTL: Duration = Duration::from_secs(5 * 60);
@@ -75,14 +74,30 @@ impl TerminalRegistry {
}
}
pub async fn create(&self, agent_id: String) -> Result<CreatedTerminal, &'static str> {
#[cfg(test)]
async fn create(&self, agent_id: String) -> Result<CreatedTerminal, &'static str> {
self.create_with_limit(
agent_id,
wakey_agent::protocol::DEFAULT_TERMINAL_MAX_SESSIONS,
)
.await
}
/// Creates a session using the limit advertised by the connected agent.
/// The caller snapshots the limit with the agent connection so UI hints and
/// server-side enforcement use the same value.
pub async fn create_with_limit(
&self,
agent_id: String,
max_sessions: usize,
) -> Result<CreatedTerminal, &'static str> {
let mut sessions = self.inner.lock().await;
prune_expired_sessions(&mut sessions);
if sessions
.values()
.filter(|session| session.agent_id == agent_id)
.count()
>= TERMINAL_MAX_SESSIONS_PER_AGENT
>= max_sessions.max(1)
{
return Err("agent_terminal_limit_reached");
}
@@ -643,6 +658,19 @@ mod tests {
assert_eq!(third.err(), Some("agent_terminal_limit_reached"));
}
#[tokio::test]
async fn advertised_limit_controls_session_capacity() {
let registry = TerminalRegistry::new();
for _ in 0..3 {
registry
.create_with_limit("router".into(), 3)
.await
.expect("within advertised limit");
}
let fourth = registry.create_with_limit("router".into(), 3).await;
assert_eq!(fourth.err(), Some("agent_terminal_limit_reached"));
}
#[tokio::test]
async fn removed_session_is_distinct_from_unknown_session() {
let registry = TerminalRegistry::new();
wakey-control-plane/src/ws.rs
+39 -3
View File
@@ -9,8 +9,8 @@ use tokio::sync::mpsc;
use tracing::{debug, info, info_span, warn};
use uuid::Uuid;
use wakey_agent::protocol::{
AgentCapability, AgentTerminalSession, ErrorPayload, RequestId, ServerMessage, TerminalControl,
TerminalId,
AgentCapability, AgentCapabilityOptions, AgentTerminalSession, DEFAULT_TERMINAL_MAX_SESSIONS,
ErrorPayload, RequestId, ServerMessage, TerminalCapabilityOptions, TerminalControl, TerminalId,
};
use wakey_core::Device;
@@ -24,6 +24,8 @@ enum IncomingClientMessage {
agent_id: String,
#[serde(default)]
capabilities: Vec<AgentCapability>,
#[serde(default)]
capability_options: AgentCapabilityOptions,
},
Auth {
agent_id: String,
@@ -59,6 +61,7 @@ struct AgentConnectionState {
hello_agent_id: Option<String>,
hello_at: Option<Instant>,
capabilities: Vec<AgentCapability>,
capability_options: AgentCapabilityOptions,
}
pub async fn agent_ws(ws: WebSocketUpgrade, State(state): State<AppState>) -> impl IntoResponse {
@@ -193,6 +196,7 @@ async fn process_agent_text(
IncomingClientMessage::Hello {
agent_id,
capabilities,
capability_options,
} => {
if connection
.hello_agent_id
@@ -208,6 +212,17 @@ async fn process_agent_text(
let connect_to_hello_ms = connected_at.elapsed().as_millis() as u64;
info!(agent_id = %agent_id, connect_to_hello_ms, "agent hello received");
connection.hello_agent_id = Some(agent_id);
connection.capability_options = AgentCapabilityOptions {
terminal: capabilities.contains(&AgentCapability::Terminal).then_some(
TerminalCapabilityOptions {
max_sessions: capability_options
.terminal
.map(|terminal| terminal.max_sessions)
.unwrap_or(DEFAULT_TERMINAL_MAX_SESSIONS)
.max(1),
},
),
};
connection.capabilities = capabilities;
}
IncomingClientMessage::Auth {
@@ -253,6 +268,7 @@ async fn process_agent_text(
connection_id: connection_id.to_string(),
tx: tx.clone(),
capabilities: connection.capabilities.clone(),
capability_options: connection.capability_options.clone(),
},
);
connection.authed_agent_id = Some(agent_id.clone());
@@ -424,10 +440,13 @@ mod tests {
use std::collections::HashMap;
use tokio::sync::mpsc;
use wakey_agent::protocol::AgentCapabilityOptions;
use crate::runtime::AgentSession;
use super::{AgentConnectionState, is_current_session, validate_auth_identity};
use super::{
AgentConnectionState, IncomingClientMessage, is_current_session, validate_auth_identity,
};
#[test]
fn current_session_check_rejects_stale_connection_ids() {
@@ -439,6 +458,7 @@ mod tests {
connection_id: "conn-new".to_string(),
tx,
capabilities: Vec::new(),
capability_options: AgentCapabilityOptions::default(),
},
);
@@ -456,4 +476,20 @@ mod tests {
assert!(validate_auth_identity(&connection, "agent-a").is_ok());
assert!(validate_auth_identity(&connection, "agent-b").is_err());
}
#[test]
fn legacy_hello_without_capability_options_deserializes() {
let message: IncomingClientMessage = serde_json::from_str(
r#"{"type":"hello","agent_id":"agent-a","capabilities":["terminal"]}"#,
)
.expect("deserialize legacy hello");
let IncomingClientMessage::Hello {
capability_options, ..
} = message
else {
panic!("expected hello");
};
assert_eq!(capability_options, AgentCapabilityOptions::default());
}
}