Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
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
8 changes: 6 additions & 2 deletions crates/aura-cli/src/backend/direct.rs
Original file line number Diff line number Diff line change
Expand Up @@ -340,8 +340,10 @@ impl DirectBackend {
let selected = get_selected_model();
let mut req = Self::build_chat_request(messages, tools, selected);

let session =
aura::SessionId::new(session_id).map_err(|e| anyhow::anyhow!("session id: {e}"))?;
let setup =
handlers::prepare_request(&self.app_state, &mut req, session_id, &self.extra_headers)
handlers::prepare_request(&self.app_state, &mut req, &session, &self.extra_headers)
.await
.map_err(|e| anyhow::anyhow!("{e}"))?;

Expand Down Expand Up @@ -406,8 +408,10 @@ impl DirectBackend {
tools: None,
};

let session =
aura::SessionId::new(session_id).map_err(|e| anyhow::anyhow!("session id: {e}"))?;
let setup =
handlers::prepare_request(&self.app_state, &mut req, session_id, &self.extra_headers)
handlers::prepare_request(&self.app_state, &mut req, &session, &self.extra_headers)
.await
.map_err(|e| anyhow::anyhow!("{e}"))?;

Expand Down
8 changes: 7 additions & 1 deletion crates/aura-events/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,12 @@ macro_rules! nonempty_string_newtype {
}
}

/// The identifier `value` names, or `None` for the empty string:
/// for a boundary where a blank and an absent value mean the same.
pub fn non_empty(value: impl Into<String>) -> Option<Self> {
Self::new(value).ok()
}

pub fn as_str(&self) -> &str {
&self.0
}
Expand Down Expand Up @@ -236,7 +242,7 @@ impl std::fmt::Display for EmptyId {

impl std::error::Error for EmptyId {}

string_newtype! {
nonempty_string_newtype! {
/// Identifier for a session: an agent's identity over time, and the key
/// its [`SessionEvent`](run::SessionEvent) stream is ordered under.
SessionId
Expand Down
34 changes: 32 additions & 2 deletions crates/aura-events/src/run.rs
Original file line number Diff line number Diff line change
Expand Up @@ -296,7 +296,7 @@ mod tests {

fn envelope(seq: u64, run: Option<RunId>, payload: SessionEventPayload) -> SessionEvent {
SessionEvent {
session_id: SessionId::new("sess_1"),
session_id: SessionId::new("sess_1").unwrap(),
run_id: run,
seq: SequenceNumber::try_from(seq).expect("a sequence starts at 1"),
at: Timestamp::from_unix_millis(1_700_000_000_000),
Expand Down Expand Up @@ -631,7 +631,7 @@ mod tests {
fn ids_serialize_as_bare_strings() {
assert_eq!(serde_json::to_value(run_id()).unwrap(), json!(RUN));
assert_eq!(
serde_json::to_value(SessionId::new("sess_1")).unwrap(),
serde_json::to_value(SessionId::new("sess_1").unwrap()).unwrap(),
json!("sess_1")
);
assert_eq!(
Expand Down Expand Up @@ -688,6 +688,36 @@ mod tests {
);
}

/// A session named by the empty string names nothing, like an observer
/// or a checkpoint, so neither the constructor nor deserialization will
/// build one, and a stream keyed by it cannot exist.
#[test]
fn an_empty_session_id_is_refused() {
let err = SessionId::new("").unwrap_err();
assert_eq!(err.to_string(), "SessionId cannot be empty");
assert!(SessionId::try_from(String::new()).is_err());
assert!(serde_json::from_value::<SessionId>(json!("")).is_err());
assert_eq!(
serde_json::from_value::<SessionId>(json!("sess_1")).unwrap(),
"sess_1"
);
}

#[test]
fn a_blank_names_no_session() {
assert_eq!(SessionId::non_empty(""), None);
assert_eq!(SessionId::non_empty(String::new()), None);
assert_eq!(
SessionId::non_empty("sess_1"),
Some(SessionId::new("sess_1").unwrap())
);
assert_eq!(
Some("sess_1".to_string()).and_then(SessionId::non_empty),
Some(SessionId::new("sess_1").unwrap())
);
assert_eq!(Some(String::new()).and_then(SessionId::non_empty), None);
}

/// The refusal reaches the envelope: an attach naming an empty observer
/// does not parse into an event.
#[test]
Expand Down
34 changes: 20 additions & 14 deletions crates/aura-web-server/src/handlers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -231,7 +231,7 @@ pub struct RequestSetup {
pub async fn prepare_request(
data: &AppState,
req: &mut ChatCompletionRequest,
chat_session_id: &str,
chat_session_id: &aura::SessionId,
req_headers_map: &HashMap<String, String>,
) -> Result<RequestSetup, PrepareError> {
// Client-side tools are gated per-agent by `[agent].enable_client_tools`
Expand Down Expand Up @@ -298,10 +298,8 @@ pub async fn prepare_request(
// server-generated session id is echoed back via `X-Chat-Session-Id`, so
// the client may adopt it on its next request, and an id that is never
// reused just leaves TTL-bounded orphan records.
let skill_log = aura::session_store::SkillLogKey::new(
aura::SessionId::new(chat_session_id),
config.agent_id(),
);
let skill_log =
aura::session_store::SkillLogKey::new(chat_session_id.clone(), config.agent_id());
let skill_recorder = (!config.agent.skills.local.is_empty()).then(|| {
Arc::new(aura::skill_tool::SkillInvocationRecorder::new(
data.session_store.skills(),
Expand Down Expand Up @@ -1341,8 +1339,15 @@ pub async fn resolve_approval(
/// A blank value counts as none given. Taken as given, every client that sends
/// one would share a single session, and with it each other's recorded skill
/// invocations; orchestration persistence would refuse it outright.
fn chat_session_id(metadata: Option<&HashMap<String, String>>, headers: &HeaderMap) -> String {
let given = |value: &str| (!value.trim().is_empty()).then(|| value.to_owned());
fn chat_session_id(
metadata: Option<&HashMap<String, String>>,
headers: &HeaderMap,
) -> aura::SessionId {
let given = |value: &str| {
(!value.trim().is_empty())
.then(|| value.to_owned())
.and_then(aura::SessionId::non_empty)
};
let header = |name: &str| {
headers
.get(name)
Expand All @@ -1358,8 +1363,9 @@ fn chat_session_id(metadata: Option<&HashMap<String, String>>, headers: &HeaderM
}

/// Generate a chat session ID (simple GUID)
fn generate_chat_session_id() -> String {
format!("cs_{}", Uuid::new_v4().simple())
fn generate_chat_session_id() -> aura::SessionId {
aura::SessionId::new(format!("cs_{}", Uuid::new_v4().simple()))
.expect("a generated id carries its prefix, so it is never empty")
}

/// The HMAC-ON ingress path for `resolve_approval`: buffer the raw bytes
Expand Down Expand Up @@ -1526,7 +1532,7 @@ mod tests {
let first = chat_session_id(Some(&metadata("")), &blank);
let second = chat_session_id(Some(&metadata("")), &blank);

assert!(first.starts_with("cs_"), "got: {first}");
assert!(first.as_str().starts_with("cs_"), "got: {first}");
assert_ne!(first, second);
}
}
Expand Down Expand Up @@ -3141,13 +3147,13 @@ source = '{}'
skill_agent(dir.path(), "agent-a", "# Agent A instructions"),
skill_agent(dir.path(), "agent-b", "# Agent B instructions"),
]);
let session = "sess-switch";
let session = aura::SessionId::new("sess-switch").unwrap();
// Agent A loaded `shared` during the session's first turn.
state
.session_store
.skills()
.record(
&SkillLogKey::new(aura::SessionId::new(session), "agent-a"),
&SkillLogKey::new(session.clone(), "agent-a"),
SkillInvocationRecord {
version: SKILL_INVOCATION_RECORD_VERSION,
invocation: SkillInvocation::LoadSkill {
Expand All @@ -3165,7 +3171,7 @@ source = '{}'
// The client keeps the session id and switches to agent B, which
// configures a skill of the same name it never loaded.
let setup =
prepare_request(&state, &mut follow_up("agent-b"), session, &HashMap::new())
prepare_request(&state, &mut follow_up("agent-b"), &session, &HashMap::new())
.await
.unwrap();
assert!(
Expand All @@ -3177,7 +3183,7 @@ source = '{}'

// Agent A, back on the same session, still gets its own.
let setup =
prepare_request(&state, &mut follow_up("agent-a"), session, &HashMap::new())
prepare_request(&state, &mut follow_up("agent-a"), &session, &HashMap::new())
.await
.unwrap();
assert_eq!(setup.rehydrated_skills, ["shared"]);
Expand Down
4 changes: 2 additions & 2 deletions crates/aura-web-server/tests/file_session_store_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -525,7 +525,7 @@ fn skill_record(name: &str, anchor: u32, seq: u32) -> SkillInvocationRecord {
}

fn skill_log(session_id: &str) -> SkillLogKey {
SkillLogKey::new(SessionId::new(session_id), "agent")
SkillLogKey::new(SessionId::new(session_id).unwrap(), "agent")
}

fn skill_store(dir: &tempfile::TempDir, ttl_secs: Option<u64>) -> FileSkillInvocationStore {
Expand Down Expand Up @@ -633,7 +633,7 @@ async fn skill_sessions_are_isolated_and_a_hostile_id_stays_inside_the_root() {
async fn skill_agents_sharing_a_session_are_isolated() {
let dir = tempfile::tempdir().unwrap();
let store = skill_store(&dir, Some(60));
let session = SessionId::new("sess-switch");
let session = SessionId::new("sess-switch").unwrap();
let first = SkillLogKey::new(session.clone(), "agent-a");
let second = SkillLogKey::new(session, "agent-b");

Expand Down
4 changes: 2 additions & 2 deletions crates/aura-web-server/tests/redis_session_store_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1481,7 +1481,7 @@ async fn corrupt_record_is_skipped_from_list() {
// ---------------------------------------------------------------------------

fn skill_log(session_id: &str) -> SkillLogKey {
SkillLogKey::new(SessionId::new(session_id), "agent")
SkillLogKey::new(SessionId::new(session_id).unwrap(), "agent")
}

/// The Redis key holding `log`, per the store's key schema.
Expand Down Expand Up @@ -1614,7 +1614,7 @@ async fn skill_sessions_are_isolated_by_prefix_and_session() {
#[tokio::test]
async fn skill_agents_sharing_a_session_are_isolated() {
let skills = connect(&test_config(60)).await.skills();
let session = SessionId::new("sess-switch");
let session = SessionId::new("sess-switch").unwrap();
let first = SkillLogKey::new(session.clone(), "agent-a");
let second = SkillLogKey::new(session, "agent-b");

Expand Down
2 changes: 1 addition & 1 deletion crates/aura/src/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -392,7 +392,7 @@ impl PreparedAgent {
session_id: config_owned
.session_id
.clone()
.map(crate::config::SessionId::new),
.and_then(crate::config::SessionId::non_empty),
};
let wrapper = Arc::new(crate::hitl::HitlApprovalWrapper::new(
hitl.patterns.clone(),
Expand Down
2 changes: 1 addition & 1 deletion crates/aura/src/hitl/route.rs
Original file line number Diff line number Diff line change
Expand Up @@ -867,7 +867,7 @@ mod tests {
scope: AgentScope::Worker {
run_id,
task: crate::orchestration::TaskIdentity::new(2, Some("k8s-agent".to_string())),
session_id: Some(crate::config::SessionId::new("sess-abc".to_string())),
session_id: Some(crate::config::SessionId::new("sess-abc").unwrap()),
},
origin: ApprovalOrigin::AgentRequested {
reason: "deleting prod ns".to_string(),
Expand Down
4 changes: 2 additions & 2 deletions crates/aura/src/orchestration/orchestrator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -962,7 +962,7 @@ impl Orchestrator {
let scope = crate::hitl::AgentScope::Worker {
run_id,
task: super::TaskIdentity::new(task_id, worker_name.map(String::from)),
session_id: session_id_owned.map(crate::config::SessionId::new),
session_id: session_id_owned.and_then(crate::config::SessionId::non_empty),
};
let mut gate = crate::hitl::HitlApprovalWrapper::new(
hitl.patterns.clone(),
Expand Down Expand Up @@ -1207,7 +1207,7 @@ impl Orchestrator {
.map(|run_id| crate::hitl::AgentScope::Worker {
run_id,
task: super::TaskIdentity::new(task_id, worker_name.map(String::from)),
session_id: session_id.map(crate::config::SessionId::new),
session_id: session_id.and_then(crate::config::SessionId::non_empty),
})
}

Expand Down
2 changes: 1 addition & 1 deletion crates/aura/src/session_store/memory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -534,7 +534,7 @@ mod tests {
}

fn skill_log(session_id: impl Into<String>, agent_id: &str) -> SkillLogKey {
SkillLogKey::new(SessionId::new(session_id), agent_id)
SkillLogKey::new(SessionId::new(session_id).unwrap(), agent_id)
}

#[tokio::test]
Expand Down
30 changes: 27 additions & 3 deletions crates/aura/src/session_store/record.rs
Original file line number Diff line number Diff line change
Expand Up @@ -179,7 +179,7 @@ impl TryFrom<ScopeRecord> for AgentScope {
fn try_from(record: ScopeRecord) -> Result<Self, Self::Error> {
Ok(match record {
ScopeRecord::Single { session_id } => AgentScope::Single {
session_id: session_id.map(SessionId::new),
session_id: session_id.and_then(SessionId::non_empty),
},
ScopeRecord::Worker {
run_id,
Expand All @@ -189,7 +189,7 @@ impl TryFrom<ScopeRecord> for AgentScope {
} => AgentScope::Worker {
run_id: parse_run_id(&run_id)?,
task: TaskIdentity::new(task_id, worker),
session_id: session_id.map(SessionId::new),
session_id: session_id.and_then(SessionId::non_empty),
},
ScopeRecord::Coordinator { run_id } => AgentScope::Coordinator {
run_id: parse_run_id(&run_id)?,
Expand Down Expand Up @@ -281,7 +281,7 @@ mod tests {
fn single_scope_round_trips() {
assert_round_trip(parked(
AgentScope::Single {
session_id: Some(SessionId::new("sess-9")),
session_id: Some(SessionId::new("sess-9").unwrap()),
},
ApprovalOrigin::ConfigGate {
matched_pattern: "kubectl_*".to_string(),
Expand Down Expand Up @@ -332,6 +332,30 @@ mod tests {
assert_eq!(json["origin"]["kind"], "agent_requested");
}

/// A stored scope carrying an empty session id restores as a scope with
/// no session rather than failing to decode.
#[test]
fn a_blank_session_id_in_a_stored_scope_restores_as_none() {
let single = ScopeRecord::Single {
session_id: Some(String::new()),
};
assert_eq!(
AgentScope::try_from(single).unwrap(),
AgentScope::Single { session_id: None }
);

let worker = ScopeRecord::Worker {
run_id: "0191e8c0-1111-7000-8000-000000000000".to_string(),
task_id: 1,
worker: None,
session_id: Some(String::new()),
};
let AgentScope::Worker { session_id, .. } = AgentScope::try_from(worker).unwrap() else {
panic!("a worker scope restores as one");
};
assert_eq!(session_id, None);
}

#[test]
fn malformed_run_id_is_an_invalid_record() {
let scope = ScopeRecord::Coordinator {
Expand Down
2 changes: 1 addition & 1 deletion crates/aura/src/skill_rehydration.rs
Original file line number Diff line number Diff line change
Expand Up @@ -406,7 +406,7 @@ mod tests {
let dir = TempDir::new().unwrap();
let skills = vec![make_skill(dir.path(), "alpha", "# Alpha instructions")];
let store = Arc::new(InMemorySkillInvocationStore::new());
let log = SkillLogKey::new(crate::config::SessionId::new("sess-loop"), "agent");
let log = SkillLogKey::new(crate::config::SessionId::new("sess-loop").unwrap(), "agent");

// Turn N: history [user], anchor = 0 + 1. The LLM calls load_skill.
let recorder = Arc::new(SkillInvocationRecorder::new(store.clone(), log.clone(), 1));
Expand Down
5 changes: 4 additions & 1 deletion crates/aura/src/skill_tool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -870,7 +870,10 @@ mod tests {
let dir = TempDir::new().unwrap();
let configs = make_skill_configs(dir.path());
let store = Arc::new(InMemorySkillInvocationStore::new());
let log = SkillLogKey::new(crate::config::SessionId::new("sess-reuse"), "agent");
let log = SkillLogKey::new(
crate::config::SessionId::new("sess-reuse").unwrap(),
"agent",
);
let turn = |id: &str, anchor| {
let recorder = Arc::new(SkillInvocationRecorder::new(
store.clone(),
Expand Down
Loading