Skip to content

feat(agents): add an opt-in sub-agent token and priority class for orchestrating sessions #7

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
tee8z wants to merge 1 commit into from
Draft

feat(agents): add an opt-in sub-agent token and priority class for orchestrating sessions #7

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 81 additions & 5 deletions crates/buzz-backend-kubernetes/src/config.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,14 @@ pub struct PodOptions {
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
/// Expose the server-assigned Pod UID to a bounded owner setup gate.
pub setup_pending: bool,
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
/// Mount a short-lived ServiceAccount token so an orchestrating agent can
/// claim sub-agent Sandboxes (§Sub-agents). Off by default.
pub enable_subagents: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
/// Optional priority class (a DNS subdomain name), also usable by
/// namespace-scoped resource quotas.
pub priority_class_name: Option<String>,
}

impl PodOptions {
Expand Down Expand Up @@ -217,6 +225,20 @@ fn optional_u64(cfg: &serde_json::Value, field: &str) -> Result<Option<u64>, Str
}
}

/// An RFC 1123 subdomain (≤253 chars of dot-separated labels), the name
/// format of cluster objects such as PriorityClasses.
fn valid_dns_subdomain(name: &str) -> bool {
name.len() <= 253
&& name.split('.').all(|label| {
!label.is_empty()
&& label.starts_with(|c: char| c.is_ascii_lowercase() || c.is_ascii_digit())
&& label.ends_with(|c: char| c.is_ascii_lowercase() || c.is_ascii_digit())
&& label
.chars()
.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-')
})
}

/// A Kubernetes namespace name: RFC 1123 label, ≤63 chars. Validated here so a
/// typo fails with a named field instead of an apiserver rejection partway
/// through a deploy.
Expand Down Expand Up @@ -277,11 +299,21 @@ pub fn parse(cfg: &serde_json::Value) -> Result<ProviderConfig, String> {
Some(n) => Some(n),
};

let pod_options: PodOptions = match cfg.get("pod_options") {
let mut pod_options: PodOptions = match cfg.get("pod_options") {
None => PodOptions::default(),
Some(value) => serde_json::from_value(value.clone())
.map_err(|e| format!("provider_config.pod_options: {e}"))?,
};
if let Some(name) = pod_options.priority_class_name.as_mut() {
*name = name.trim().to_string();
if !valid_dns_subdomain(name) {
return Err(
"provider_config.pod_options.priority_class_name must be a PriorityClass name \
(DNS subdomain: lowercase letters, digits, '-' and '.')"
.into(),
);
}
}
if pod_options.active_deadline_seconds.is_some_and(|n| n <= 0) {
return Err("provider_config.pod_options.active_deadline_seconds must be positive".into());
}
Expand Down Expand Up @@ -320,13 +352,24 @@ pub fn parse(cfg: &serde_json::Value) -> Result<ProviderConfig, String> {
}
}

let service_account = optional_string(cfg, "service_account")?;
// Without a named ServiceAccount the token would be minted for the
// namespace's `default` account, which the deployment does not scope.
if pod_options.enable_subagents && service_account.is_none() {
return Err(
"provider_config.pod_options.enable_subagents requires provider_config.service_account: \
the sub-agent token is minted for that ServiceAccount, never `default`"
.into(),
);
}

Ok(ProviderConfig {
context: optional_string(cfg, "context")?,
namespace,
image,
resources,
inactivity_seconds,
service_account: optional_string(cfg, "service_account")?,
service_account,
pod_options,
})
}
Expand Down Expand Up @@ -410,7 +453,9 @@ pub fn config_schema() -> serde_json::Value {
"ephemeral_storage_request": {"type": "string", "minLength": 1},
"ephemeral_storage_limit": {"type": "string", "minLength": 1},
"active_deadline_seconds": {"type": "integer", "minimum": 1},
"setup_pending": {"type": "boolean", "default": false}
"setup_pending": {"type": "boolean", "default": false},
"enable_subagents": {"type": "boolean", "default": false},
"priority_class_name": {"type": "string", "minLength": 1}
}
},
"service_account": {
Expand Down Expand Up @@ -640,14 +685,45 @@ mod tests {
serde_json::json!({"workspace_size_limit": " "}),
serde_json::json!({"unknown": true}),
serde_json::json!({"node_selector": "amd64"}),
serde_json::json!({"priority_class_name": ""}),
serde_json::json!({"priority_class_name": " "}),
serde_json::json!({"priority_class_name": "Agent-Session"}),
serde_json::json!({"priority_class_name": "agent_session"}),
serde_json::json!({"priority_class_name": "agent..session"}),
serde_json::json!({"priority_class_name": "-agent"}),
serde_json::json!({"priority_class_name": "a".repeat(254)}),
serde_json::json!({"priority_class_name": 5}),
// No service account: the token would belong to `default`.
serde_json::json!({"enable_subagents": true}),
] {
let mut cfg = minimal();
cfg["pod_options"] = bad;
assert!(parse(&cfg).unwrap_err().contains("pod_options"));
cfg["pod_options"] = bad.clone();
assert!(parse(&cfg).unwrap_err().contains("pod_options"), "{bad}");
}
assert!(parse(&minimal()).unwrap().pod_options.is_default());
}

#[test]
fn accepts_subagents_with_a_service_account_and_a_trimmed_priority_class() {
let mut cfg = minimal();
cfg["service_account"] = " agent-sa ".into();
cfg["pod_options"] = serde_json::json!({
"enable_subagents": true,
"priority_class_name": " buzz-agent.session-1 ",
});
let parsed = parse(&cfg).unwrap();
assert_eq!(parsed.service_account.as_deref(), Some("agent-sa"));
assert!(parsed.pod_options.enable_subagents);
assert_eq!(
parsed.pod_options.priority_class_name.as_deref(),
Some("buzz-agent.session-1")
);
cfg["pod_options"]["priority_class_name"] = "a".repeat(253).into();
assert!(parse(&cfg).is_ok());
cfg["service_account"] = " ".into();
assert!(parse(&cfg).unwrap_err().contains("enable_subagents"));
}

#[test]
fn applies_spec_defaults() {
let c = parse(&minimal()).unwrap();
Expand Down
160 changes: 145 additions & 15 deletions crates/buzz-backend-kubernetes/src/pod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,11 @@ use crate::naming::{
AgentIdentity, ANNOTATION_CREATE_INTENT, ANNOTATION_IMAGE, ANNOTATION_PUBKEY_FULL,
};
use k8s_openapi::api::core::v1::{
Capabilities, Container, EmptyDirVolumeSource, EnvFromSource, EnvVar, EnvVarSource,
ObjectFieldSelector, Pod, PodSecurityContext, PodSpec, ResourceRequirements, SeccompProfile,
Secret, SecretEnvSource, SecurityContext, Volume, VolumeMount,
Capabilities, ConfigMapProjection, Container, DownwardAPIProjection, DownwardAPIVolumeFile,
EmptyDirVolumeSource, EnvFromSource, EnvVar, EnvVarSource, KeyToPath, ObjectFieldSelector, Pod,
PodSecurityContext, PodSpec, ProjectedVolumeSource, ResourceRequirements, SeccompProfile,
Secret, SecretEnvSource, SecurityContext, ServiceAccountTokenProjection, Volume, VolumeMount,
VolumeProjection,
};
use k8s_openapi::apimachinery::pkg::api::resource::Quantity;
use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta;
Expand All @@ -26,6 +28,61 @@ const WORKSPACE_VOLUME: &str = "workspace";
/// The container name. Fixed: log and exec tooling addresses it by name.
const CONTAINER_NAME: &str = "agent";

/// Where an orchestrating agent finds its sub-agent credentials (§Sub-agents):
/// `token`, `ca.crt`, and `namespace`. Deliberately not the default
/// ServiceAccount path, so tools do not pick the token up implicitly.
pub const SUBAGENT_TOKEN_PATH: &str = "/var/run/secrets/buzz-subagents";
/// Lifetime of the projected token; the kubelet rotates it before expiry.
pub const SUBAGENT_TOKEN_SECONDS: i64 = 600;
const SUBAGENT_TOKEN_VOLUME: &str = "subagent-token";

/// Short-lived API credentials for claiming sub-agent Sandboxes. RBAC in the
/// developer namespace limits the token to SandboxClaims there.
fn subagent_token_volume() -> Volume {
Volume {
name: SUBAGENT_TOKEN_VOLUME.to_string(),
projected: Some(ProjectedVolumeSource {
sources: Some(vec![
VolumeProjection {
service_account_token: Some(ServiceAccountTokenProjection {
path: "token".to_string(),
expiration_seconds: Some(SUBAGENT_TOKEN_SECONDS),
..Default::default()
}),
..Default::default()
},
VolumeProjection {
config_map: Some(ConfigMapProjection {
name: "kube-root-ca.crt".to_string(),
items: Some(vec![KeyToPath {
key: "ca.crt".to_string(),
path: "ca.crt".to_string(),
..Default::default()
}]),
..Default::default()
}),
..Default::default()
},
VolumeProjection {
downward_api: Some(DownwardAPIProjection {
items: Some(vec![DownwardAPIVolumeFile {
path: "namespace".to_string(),
field_ref: Some(ObjectFieldSelector {
field_path: "metadata.namespace".to_string(),
..Default::default()
}),
..Default::default()
}]),
}),
..Default::default()
},
]),
..Default::default()
}),
..Default::default()
}
}

/// Build the per-attempt Secret holding the resolved environment.
///
/// `immutable: true` — the Secret is written once per attempt and never
Expand Down Expand Up @@ -144,11 +201,22 @@ pub fn build_pod(
limits: Some(limits),
..Default::default()
}),
volume_mounts: Some(vec![VolumeMount {
name: WORKSPACE_VOLUME.to_string(),
mount_path: WORKSPACE_PATH.to_string(),
..Default::default()
}]),
volume_mounts: Some({
let mut mounts = vec![VolumeMount {
name: WORKSPACE_VOLUME.to_string(),
mount_path: WORKSPACE_PATH.to_string(),
..Default::default()
}];
if cfg.pod_options.enable_subagents {
mounts.push(VolumeMount {
name: SUBAGENT_TOKEN_VOLUME.to_string(),
mount_path: SUBAGENT_TOKEN_PATH.to_string(),
read_only: Some(true),
..Default::default()
});
}
mounts
}),
security_context: Some(SecurityContext {
allow_privilege_escalation: Some(false),
capabilities: Some(Capabilities {
Expand Down Expand Up @@ -177,6 +245,7 @@ pub fn build_pod(
tolerations: (!cfg.pod_options.tolerations.is_empty())
.then(|| cfg.pod_options.tolerations.iter().map(Into::into).collect()),
active_deadline_seconds: cfg.pod_options.active_deadline_seconds,
priority_class_name: cfg.pod_options.priority_class_name.clone(),
restart_policy: Some(RESTART_POLICY.to_string()),
termination_grace_period_seconds: Some(TERMINATION_GRACE_SECONDS),
// The agent runs prompted, untrusted code while holding an nsec;
Expand All @@ -197,14 +266,20 @@ pub fn build_pod(
}),
..Default::default()
}),
volumes: Some(vec![Volume {
name: WORKSPACE_VOLUME.to_string(),
empty_dir: Some(EmptyDirVolumeSource {
size_limit: cfg.pod_options.workspace_size_limit.clone().map(Quantity),
volumes: Some({
let mut volumes = vec![Volume {
name: WORKSPACE_VOLUME.to_string(),
empty_dir: Some(EmptyDirVolumeSource {
size_limit: cfg.pod_options.workspace_size_limit.clone().map(Quantity),
..Default::default()
}),
..Default::default()
}),
..Default::default()
}]),
}];
if cfg.pod_options.enable_subagents {
volumes.push(subagent_token_volume());
}
volumes
}),
..Default::default()
}),
..Default::default()
Expand Down Expand Up @@ -376,6 +451,42 @@ mod tests {
);
}

#[test]
fn subagent_token_is_opt_in_short_lived_and_off_the_default_path() {
let mut cfg = provider_config();
let fp = intent_template(&cfg, ["A".into()]).fingerprint();
let plain = build_pod(&identity(), &cfg, "g", &fp);
assert_eq!(spec(&plain).volumes.as_ref().unwrap().len(), 1);
assert_eq!(
spec(&plain).containers[0]
.volume_mounts
.as_ref()
.unwrap()
.len(),
1
);

cfg.pod_options.enable_subagents = true;
let changed = intent_template(&cfg, ["A".into()]).fingerprint();
let pod = build_pod(&identity(), &cfg, "g", &changed);
let spec = spec(&pod);
// Still no ambient token at the default ServiceAccount path.
assert_eq!(spec.automount_service_account_token, Some(false));
let volume = &spec.volumes.as_ref().unwrap()[1];
let sources = volume.projected.as_ref().unwrap().sources.as_ref().unwrap();
let token = sources[0].service_account_token.as_ref().unwrap();
assert_eq!(token.expiration_seconds, Some(SUBAGENT_TOKEN_SECONDS));
assert!(token.audience.is_none());
let mount = &spec.containers[0].volume_mounts.as_ref().unwrap()[1];
assert_eq!(mount.name, volume.name);
assert_eq!(mount.mount_path, SUBAGENT_TOKEN_PATH);
assert_eq!(mount.read_only, Some(true));
assert_ne!(
SUBAGENT_TOKEN_PATH,
"/var/run/secrets/kubernetes.io/serviceaccount"
);
}

#[test]
fn each_optional_pod_control_changes_intent() {
let cfg = provider_config();
Expand All @@ -388,6 +499,8 @@ mod tests {
serde_json::json!({"ephemeral_storage_limit": "4Gi"}),
serde_json::json!({"active_deadline_seconds": 28800}),
serde_json::json!({"setup_pending": true}),
serde_json::json!({"enable_subagents": true}),
serde_json::json!({"priority_class_name": "agent-session"}),
] {
let mut changed = cfg.clone();
changed.pod_options = serde_json::from_value(value).unwrap();
Expand All @@ -400,6 +513,23 @@ mod tests {
assert!(original.get("pod_options").is_none());
}

#[test]
fn priority_class_reaches_the_pod_without_changing_default_intent() {
let mut cfg = provider_config();
assert!(serde_json::to_value(&cfg.pod_options)
.unwrap()
.get("priority_class_name")
.is_none());
cfg.pod_options.priority_class_name = Some("agent-session".into());
let identity = identity();
let fingerprint = intent_template(&cfg, ["A".into()]).fingerprint();
let pod = build_pod(&identity, &cfg, "generation", &fingerprint);
assert_eq!(
pod.spec.unwrap().priority_class_name.as_deref(),
Some("agent-session")
);
}

#[test]
fn optional_placement_storage_and_lifetime_reach_pod_and_intent() {
let mut cfg = provider_config();
Expand Down
18 changes: 18 additions & 0 deletions desktop/src-tauri/src/managed_agents/backend_tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -521,6 +521,24 @@ esac"#,
assert!(error.to_string().contains("[REDACTED]"), "{error}");
}

#[test]
fn subagent_capability_is_public_but_credentials_still_fail() {
for enabled in [true, false] {
assert!(
validate_provider_config(&serde_json::json!({"pod_options": {
"enable_subagents": enabled, "priority_class_name": "agent-session"
}}))
.is_ok()
);
}
for credential in [
serde_json::json!({"access_token":"secret"}),
serde_json::json!({"api_key":"secret"}),
] {
assert!(validate_provider_config(&serde_json::json!({"pod_options":credential})).is_err());
}
}

#[test]
fn validate_provider_config_rejects_secret_key() {
let cfg = serde_json::json!({"api_key": "val"});
Expand Down
Loading
Loading